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

JAVA-5139

Author: stIncMale

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

@@ -48,6 +48,7 @@
 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 +142,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 +157,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 +180,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 +199,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 +219,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 +234,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 +257,7 @@ BsonDocument rewrapManyDataKey(final BsonDocument filter, final RewrapManyDataKe
                 return executeStateMachine(rewrapManyDatakeyContext, null);
             }
         } catch (MongoCryptException e) {
-            throw wrapInClientException(e);
+            throw wrapInMongoException(e);
         }
     }
 
@@ -324,7 +325,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 +349,8 @@ private void decryptKeys(final MongoCryptContext cryptContext) {
             }
             cryptContext.completeKeyDecryptors();
         } catch (Throwable t) {
-            throw wrapInClientException(t);
+            throw translateInterruptedException(t, "Interrupted while doing IO")
+                    .orElseThrow(() -> wrapInMongoException(t));
         }
     }
 
@@ -366,7 +368,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

@@ -98,7 +98,7 @@ private void enableHostNameVerification(final SSLSocket socket) {
     private void closeSocket(final Socket socket) {
         try {
             socket.close();
-        } catch (IOException e) {
+        } catch (IOException | RuntimeException e) {
             // ignore
         }
     }