Handle interrupts: synchronous KeyManagementService uses Socket IO (open, read, write), which is interruptible in a virtual thread

driver-sync/src/main/com/mongodb/client/internal/Crypt.java

@@ -42,12 +42,14 @@
 import java.io.InputStream;
 import java.nio.ByteBuffer;
 import java.util.Map;
+import java.util.function.Function;
 import java.util.function.Supplier;
 
 import static com.mongodb.assertions.Assertions.assertNotNull;
 import static com.mongodb.assertions.Assertions.notNull;
 import static com.mongodb.crypt.capi.MongoCryptContext.State;
 import static com.mongodb.internal.client.vault.EncryptOptionsHelper.asMongoExplicitEncryptOptions;
+import static com.mongodb.internal.thread.InterruptionUtil.translateInterruptedException;
 
 /**
  * <p>This class is not part of the public API and may be removed or changed at any time</p>
@@ -141,7 +143,7 @@ RawBsonDocument encrypt(final String databaseName, final RawBsonDocument command
        try (MongoCryptContext encryptionContext = mongoCrypt.createEncryptionContext(databaseName, command)) {
            return executeStateMachine(encryptionContext, databaseName);
         } catch (MongoCryptException e) {
-            throw wrapInClientException(e);
+            throw wrapInMongoException(e);
         }
     }
 
@@ -156,7 +158,7 @@ RawBsonDocument decrypt(final RawBsonDocument commandResponse) {
         try (MongoCryptContext decryptionContext = mongoCrypt.createDecryptionContext(commandResponse)) {
             return executeStateMachine(decryptionContext, null);
         } catch (MongoCryptException e) {
-            throw wrapInClientException(e);
+            throw wrapInMongoException(e);
         }
     }
 
@@ -179,7 +181,7 @@ BsonDocument createDataKey(final String kmsProvider, final DataKeyOptions option
                         .build())) {
             return executeStateMachine(dataKeyCreationContext, null);
         } catch (MongoCryptException e) {
-            throw wrapInClientException(e);
+            throw wrapInMongoException(e);
         }
     }
 
@@ -198,7 +200,7 @@ BsonBinary encryptExplicitly(final BsonValue value, final EncryptOptions options
                 new BsonDocument("v", value), asMongoExplicitEncryptOptions(options))) {
             return executeStateMachine(encryptionContext, null).getBinary("v");
         } catch (MongoCryptException e) {
-            throw wrapInClientException(e);
+            throw wrapInMongoException(e);
         }
     }
 
@@ -218,7 +220,7 @@ BsonDocument encryptExpression(final BsonDocument expression, final EncryptOptio
                 new BsonDocument("v", expression), asMongoExplicitEncryptOptions(options))) {
             return executeStateMachine(encryptionContext, null).getDocument("v");
         } catch (MongoCryptException e) {
-            throw wrapInClientException(e);
+            throw wrapInMongoException(e);
         }
     }
 
@@ -233,7 +235,7 @@ BsonValue decryptExplicitly(final BsonBinary value) {
         try (MongoCryptContext decryptionContext = mongoCrypt.createExplicitDecryptionContext(new BsonDocument("v", value))) {
             return assertNotNull(executeStateMachine(decryptionContext, null).get("v"));
         } catch (MongoCryptException e) {
-            throw wrapInClientException(e);
+            throw wrapInMongoException(e);
         }
     }
 
@@ -256,7 +258,7 @@ BsonDocument rewrapManyDataKey(final BsonDocument filter, final RewrapManyDataKe
                 return executeStateMachine(rewrapManyDatakeyContext, null);
             }
         } catch (MongoCryptException e) {
-            throw wrapInClientException(e);
+            throw wrapInMongoException(e);
         }
     }
 
@@ -324,7 +326,7 @@ private void mark(final MongoCryptContext cryptContext, final String databaseNam
             cryptContext.addMongoOperationResult(markedCommand);
             cryptContext.completeMongoOperation();
         } catch (Throwable t) {
-            throw wrapInClientException(t);
+            throw wrapInMongoException(t);
         }
     }
 
@@ -348,7 +350,9 @@ private void decryptKeys(final MongoCryptContext cryptContext) {
             }
             cryptContext.completeKeyDecryptors();
         } catch (Throwable t) {
-            throw wrapInClientException(t);
+            throw translateInterruptedException(t, "Interrupted while reading")
+                    .<MongoException>map(Function.identity())
+                    .orElseGet(() -> wrapInMongoException(t));
         }
     }
 
@@ -366,7 +370,11 @@ private void decryptKey(final MongoKeyDecryptor keyDecryptor) throws IOException
         }
     }
 
-    private MongoClientException wrapInClientException(final Throwable t) {
-        return new MongoClientException("Exception in encryption library: " + t.getMessage(), t);
+    private MongoException wrapInMongoException(final Throwable t) {
+        if (t instanceof MongoException) {
+            return (MongoException) t;
+        } else {
+            return new MongoClientException("Exception in encryption library: " + t.getMessage(), t);
+        }
     }
 }

driver-sync/src/main/com/mongodb/client/internal/KeyManagementService.java

@@ -16,6 +16,7 @@
 
 package com.mongodb.client.internal;
 
+import com.mongodb.MongoInterruptedException;
 import com.mongodb.ServerAddress;
 import com.mongodb.internal.diagnostics.logging.Logger;
 import com.mongodb.internal.diagnostics.logging.Loggers;
@@ -34,8 +35,10 @@
 import java.net.Socket;
 import java.nio.ByteBuffer;
 import java.util.Map;
+import java.util.Optional;
 
 import static com.mongodb.assertions.Assertions.notNull;
+import static com.mongodb.internal.thread.InterruptionUtil.translateInterruptedException;
 
 class KeyManagementService {
     private static final Logger LOGGER = Loggers.getLogger("client");
@@ -63,7 +66,7 @@ public InputStream stream(final String kmsProvider, final String host, final Byt
             socket.connect(new InetSocketAddress(InetAddress.getByName(serverAddress.getHost()), serverAddress.getPort()), timeoutMillis);
         } catch (IOException e) {
             closeSocket(socket);
-            throw e;
+            throw handleInterruptAndThrow(e, "Interrupted while connecting");
         }
 
         try {
@@ -75,7 +78,7 @@ public InputStream stream(final String kmsProvider, final String host, final Byt
             outputStream.write(bytes);
         } catch (IOException e) {
             closeSocket(socket);
-            throw e;
+            throw handleInterruptAndThrow(e, "Interrupted while writing");
         }
 
         try {
@@ -98,8 +101,21 @@ private void enableHostNameVerification(final SSLSocket socket) {
     private void closeSocket(final Socket socket) {
         try {
             socket.close();
-        } catch (IOException e) {
+        } catch (IOException | RuntimeException e) {
             // ignore
         }
     }
+
+    /**
+     * @return Never.
+     */
+    private static RuntimeException handleInterruptAndThrow(final IOException e, final String message) throws IOException,
+            MongoInterruptedException {
+        Optional<MongoInterruptedException> interruptedException = translateInterruptedException(e, message);
+        if (interruptedException.isPresent()) {
+            throw interruptedException.get();
+        } else {
+            throw e;
+        }
+    }
 }