Make the AuthMechanism interface rather a lot less stinky.

Author: Simon MacMullen

src/com/rabbitmq/client/AuthMechanism.java

@@ -1,23 +1,22 @@
 package com.rabbitmq.client;
 
-import com.rabbitmq.client.impl.AMQChannel;
+import com.rabbitmq.client.impl.LongString;
 
 import java.io.IOException;
 
 /**
- * A pluggable authentication mechanism
+ * A pluggable authentication mechanism. Should be stateless.
  */
 public interface AuthMechanism {
     /**
-     * Send and receive start-ok / secure / secure-ok until a connection is
-     * established or an exception thrown.
-     * 
-     * @param channel to send methods on
+     * Handle one round of challenge-response
+     * @param round zero-based counter of the current round
+     * @param challenge the challenge this round, or null on round 0.
      * @param factory for reference to e.g. username and password.
-     * @return the Connection.Tune method sent by the server after authentication
-     * @throws IOException if the authentication failed or something else went wrong
+     * @return response
+     * @throws IOException
      */
-    AMQP.Connection.Tune doLogin(AMQChannel channel, ConnectionFactory factory) throws IOException;
+    LongString handleChallenge(int round, LongString challenge, ConnectionFactory factory);
 
     /**
      * The name of the authentication mechanism, as negotiated on the wire

src/com/rabbitmq/client/impl/AMQConnection.java

@@ -283,7 +283,30 @@ public void start()
             throw new IOException("No compatible authentication mechanism found - " +
                     "server offered [" + connStart.getMechanisms() + "]");
         }
-        AMQP.Connection.Tune connTune = mechanism.doLogin(_channel0, _factory);
+
+        int round = 0;
+        LongString challenge = null;
+        AMQP.Connection.Tune connTune = null;
+        do {
+            LongString response = mechanism.handleChallenge(round, challenge, _factory);
+            Method method = (round == 0)
+                ? new AMQImpl.Connection.StartOk(_clientProperties,
+                                                 mechanism.getName(), response,
+                                                 "en_US")
+                : new AMQImpl.Connection.SecureOk(response);
+
+            try {
+                Method serverResponse = _channel0.rpc(method).getMethod();
+                if (serverResponse instanceof AMQP.Connection.Tune) {
+                    connTune = (AMQP.Connection.Tune) serverResponse;
+                } else {
+                    challenge = ((AMQP.Connection.Secure) serverResponse).getChallenge();
+                    round++;
+                }
+            } catch (ShutdownSignalException e) {
+                throw AMQChannel.wrap(e, "Possibly caused by authentication failure");
+            }
+        } while (connTune == null);
 
         int channelMax =
             negotiatedMaxValue(_factory.getRequestedChannelMax(),

src/com/rabbitmq/client/impl/PlainMechanism.java

@@ -1,29 +1,16 @@
 package com.rabbitmq.client.impl;
 
-import com.rabbitmq.client.AMQP;
 import com.rabbitmq.client.AuthMechanism;
 import com.rabbitmq.client.ConnectionFactory;
-import com.rabbitmq.client.ShutdownSignalException;
-
-import java.io.IOException;
 
 /**
  * The PLAIN auth mechanism
  */
 public class PlainMechanism implements AuthMechanism {
-    public AMQP.Connection.Tune doLogin(AMQChannel channel,
-                                        ConnectionFactory factory) throws IOException {
-        LongString saslResponse = LongStringHelper.asLongString(
-                "\0" + factory.getUsername() + "\0" + factory.getPassword());
-        AMQImpl.Connection.StartOk startOk =
-            new AMQImpl.Connection.StartOk(factory.getClientProperties(), getName(),
-                                           saslResponse, "en_US");
-
-        try {
-            return (AMQP.Connection.Tune) channel.rpc(startOk).getMethod();
-        } catch (ShutdownSignalException e) {
-            throw AMQChannel.wrap(e, "Possibly caused by authentication failure");
-        }
+    public LongString handleChallenge(int round, LongString challenge,
+                                      ConnectionFactory factory) {
+        return LongStringHelper.asLongString("\0" + factory.getUsername() +
+                                             "\0" + factory.getPassword());
     }
 
     public String getName() {

src/com/rabbitmq/client/impl/ScramMD5Mechanism.java

@@ -1,11 +1,9 @@
 package com.rabbitmq.client.impl;
 
-import com.rabbitmq.client.AMQP;
 import com.rabbitmq.client.AuthMechanism;
 import com.rabbitmq.client.ConnectionFactory;
-import com.rabbitmq.client.ShutdownSignalException;
 
-import java.io.IOException;
+import java.io.UnsupportedEncodingException;
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
 import java.util.Arrays;
@@ -25,39 +23,40 @@ This is only somewhat improved security over PLAIN (if you can
 */
 
 public class ScramMD5Mechanism implements AuthMechanism {
-    public AMQP.Connection.Tune doLogin(AMQChannel channel,
-                                        ConnectionFactory factory) throws IOException {
-        try {
-            LongString resp1 = LongStringHelper.asLongString(factory.getUsername());
-            AMQImpl.Connection.StartOk startOk =
-                new AMQImpl.Connection.StartOk(factory.getClientProperties(), getName(),
-                                               resp1, "en_US");
-            AMQP.Connection.Secure secure =
-                    (AMQP.Connection.Secure) channel.rpc(startOk).getMethod();
-            byte[] challenge = secure.getChallenge().getBytes();
-            byte[] salt1 = Arrays.copyOfRange(challenge, 0, 4);
-            byte[] salt2 = Arrays.copyOfRange(challenge, 4, 8);
-
-            MessageDigest digest1 = MessageDigest.getInstance("MD5");
-            MessageDigest digest2 = MessageDigest.getInstance("MD5");
-            byte[] d1 = digest1.digest(concat(salt1, factory.getPassword().getBytes("utf-8")));
-            byte[] d2 = digest2.digest(concat(salt2, d1));
+    public LongString handleChallenge(int round, LongString challengeStr,
+                                      ConnectionFactory factory) {
+        if (round == 0) {
+            return LongStringHelper.asLongString(factory.getUsername());
+        } else {
+            try {
+                byte[] challenge = challengeStr.getBytes();
+                byte[] salt1 = Arrays.copyOfRange(challenge, 0, 4);
+                byte[] salt2 = Arrays.copyOfRange(challenge, 4, 8);
 
-            AMQImpl.Connection.SecureOk secureOk =
-                new AMQImpl.Connection.SecureOk(LongStringHelper.asLongString(d2));
+                byte[] pw = factory.getPassword().getBytes("utf-8");
+                byte[] d = digest(salt2, digest(salt1, pw));
 
-            return (AMQP.Connection.Tune) channel.rpc(secureOk).getMethod();
-        } catch (NoSuchAlgorithmException e) {
-            throw new RuntimeException(e);
-        } catch (ShutdownSignalException e) {
-            throw AMQChannel.wrap(e, "Possibly caused by authentication failure");
+                return LongStringHelper.asLongString(d);
+            } catch (UnsupportedEncodingException e) {
+                throw new RuntimeException(e);
+            }
         }
     }
 
     public String getName() {
         return "RABBIT-SCRAM-MD5";
     }
 
+    private static byte[] digest(byte[] arr1, byte[] arr2) {
+        try {
+            MessageDigest digest = MessageDigest.getInstance("MD5");
+            return digest.digest(concat(arr1, arr2));
+
+        } catch (NoSuchAlgorithmException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     private static byte[] concat(byte[] first, byte[] second) {
         byte[] result = Arrays.copyOf(first, first.length + second.length);
         System.arraycopy(second, 0, result, first.length, second.length);