Ensure server proof has been validated during SCRAM conversation

JAVA-3527

Author: jstewart148

driver-core/src/main/com/mongodb/internal/connection/SaslAuthenticator.java

@@ -66,6 +66,14 @@ public Void run() {
 
                         res = sendSaslContinue(conversationId, response, connection);
                     }
+                    if (!saslClient.isComplete()) {
+                        saslClient.evaluateChallenge((res.getBinary("payload")).getData());
+                        if (!saslClient.isComplete()) {
+                            throw new MongoSecurityException(getMongoCredential(),
+                                    "SASL protocol error: server completed challenges before client completed responses "
+                                            + getMongoCredential());
+                        }
+                    }
                 } catch (Exception e) {
                     throw wrapException(e);
                 } finally {
@@ -93,7 +101,7 @@ public void onResult(final BsonDocument result, final Throwable t) {
                                 if (t != null) {
                                     callback.onResult(null, wrapException(t));
                                 } else if (result.getBoolean("done").getValue()) {
-                                    callback.onResult(null, null);
+                                    verifySaslClientComplete(saslClient, result, callback);
                                 } else {
                                     new Continuator(saslClient, result, connection, callback).start();
                                 }
@@ -121,6 +129,26 @@ private void throwIfSaslClientIsNull(final SaslClient saslClient) {
         }
     }
 
+    private void verifySaslClientComplete(final SaslClient saslClient, final BsonDocument result,
+                                          final SingleResultCallback<Void> callback) {
+        if (saslClient.isComplete()) {
+            callback.onResult(null, null);
+        } else {
+            try {
+                saslClient.evaluateChallenge(result.getBinary("payload").getData());
+                if (saslClient.isComplete()) {
+                    callback.onResult(null, null);
+                } else {
+                    callback.onResult(null, new MongoSecurityException(getMongoCredential(),
+                            "SASL protocol error: server completed challenges before client completed responses "
+                                    + getMongoCredential()));
+                }
+            } catch (SaslException e) {
+                callback.onResult(null, wrapException(e));
+            }
+        }
+    }
+
     @Nullable
     private Subject getSubject() {
         return getMongoCredential().getMechanismProperty(JAVA_SUBJECT_KEY, null);
@@ -202,7 +230,7 @@ public void onResult(final BsonDocument result, final Throwable t) {
                 callback.onResult(null, wrapException(t));
                 disposeOfSaslClient(saslClient);
             } else if (result.getBoolean("done").getValue()) {
-                callback.onResult(null, null);
+                verifySaslClientComplete(saslClient, result, callback);
                 disposeOfSaslClient(saslClient);
             } else {
                 continueConversation(result);

driver-core/src/test/unit/com/mongodb/internal/connection/ScramShaAuthenticatorSpecification.groovy

@@ -360,13 +360,65 @@ class ScramShaAuthenticatorSpecification extends Specification {
         async << [true, false]
     }
 
-    def createConnection(List<String> serverResponses) {
+    def 'should complete authentication when done is set to true prematurely SHA-256'() {
+        given:
+        def serverResponses = createMessages('''
+            S: r=rOprNGfwEbeRWgbNEkqO%hvYDpWUa2RaTCAfuxFIlj)hNlF$k0,s=W22ZaJ0SNY7soEsUEjb6gQ==,i=4096
+            S: v=6rriTRBi23WpRR/wtup+mMhUZUn/dB5nLTJRsjl95G4=
+        ''').last()
+        def authenticator = new ScramShaAuthenticator(SHA256_CREDENTIAL, { 'rOprNGfwEbeRWgbNEkqO' }, { 'pencil' })
+
+        when:
+        // server sends done=true on first response, client is not complete after processing response
+        authenticate(createConnection(serverResponses, 0), authenticator, async)
+
+        then:
+        def e = thrown(MongoSecurityException)
+        e.getMessage().contains('server completed challenges before client completed responses')
+
+        when:
+        // server sends done=true on second response, client is complete after processing response
+        authenticate(createConnection(serverResponses, 1), authenticator, async)
+
+        then:
+        noExceptionThrown()
+
+        where:
+        async << [true, false]
+    }
+
+    def 'should throw exception when done is set to true prematurely and server response is invalid SHA-256'() {
+        given:
+        def serverResponses = createMessages('''
+            S: r=rOprNGfwEbeRWgbNEkqO%hvYDpWUa2RaTCAfuxFIlj)hNlF$k0,s=W22ZaJ0SNY7soEsUEjb6gQ==,i=4096
+            S: v=invalidResponse
+        ''').last()
+        def authenticator = new ScramShaAuthenticator(SHA256_CREDENTIAL, { 'rOprNGfwEbeRWgbNEkqO' }, { 'pencil' })
+
+        when:
+        // server sends done=true on second response, client throws exception on invalid server response
+        authenticate(createConnection(serverResponses, 1), authenticator, async)
+
+        then:
+        def e = thrown(MongoSecurityException)
+        e.getCause() instanceof SaslException
+        e.getCause().getMessage() == 'Server signature was invalid.'
+
+        where:
+        async << [true, false]
+    }
+
+    def createConnection(List<String> serverResponses, int responseWhereDoneIsTrue = -1) {
         TestInternalConnection connection = new TestInternalConnection(serverId)
-        serverResponses.each {
+        serverResponses.eachWithIndex { response, index ->
+            def isDone = (index == responseWhereDoneIsTrue).booleanValue()
             connection.enqueueReply(
-                    buildSuccessfulReply("{conversationId: 1, payload: BinData(0, '${encode64(it)}'), done: false, ok: 1}")
-            ) }
-        connection.enqueueReply(buildSuccessfulReply('{conversationId: 1, done: true, ok: 1}'))
+                    buildSuccessfulReply("{conversationId: 1, payload: BinData(0, '${encode64(response)}'), done: ${isDone}, ok: 1}")
+            )
+        }
+        if (responseWhereDoneIsTrue < 0) {
+            connection.enqueueReply(buildSuccessfulReply('{conversationId: 1, done: true, ok: 1}'))
+        }
         connection
     }