Make auth mechanisms pluggable in the Java client.

Author: Simon MacMullen

src/com/rabbitmq/client/AuthMechanism.java

@@ -0,0 +1,26 @@
+package com.rabbitmq.client;
+
+import com.rabbitmq.client.impl.AMQChannel;
+
+import java.io.IOException;
+
+/**
+ * A pluggable authentication mechanism
+ */
+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
+     * @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
+     */
+    AMQP.Connection.Tune doLogin(AMQChannel channel, ConnectionFactory factory) throws IOException;
+
+    /**
+     * The name of the authentication mechanism, as negotiated on the wire
+     */
+    String getName();
+}

src/com/rabbitmq/client/ConnectionFactory.java

@@ -33,7 +33,7 @@
 import java.io.IOException;
 import java.security.KeyManagementException;
 import java.security.NoSuchAlgorithmException;
-import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 import java.net.Socket;
@@ -46,6 +46,7 @@
 
 import com.rabbitmq.client.impl.AMQConnection;
 import com.rabbitmq.client.impl.FrameHandler;
+import com.rabbitmq.client.impl.PlainMechanism;
 import com.rabbitmq.client.impl.SocketFrameHandler;
 
 /**
@@ -85,6 +86,10 @@ public class ConnectionFactory implements Cloneable {
     /** The default port to use for AMQP connections when using SSL */
     public static final int DEFAULT_AMQP_OVER_SSL_PORT = 5671;
 
+    /** The default list of authentication mechanisms to use */
+    public static final AuthMechanism[] DEFAULT_AUTH_MECHANISMS =
+        new AuthMechanism[] { new PlainMechanism() };
+
     /**
      * The default SSL protocol (currently "SSLv3").
      */
@@ -98,6 +103,7 @@ public class ConnectionFactory implements Cloneable {
     private int requestedChannelMax               = DEFAULT_CHANNEL_MAX;
     private int requestedFrameMax                 = DEFAULT_FRAME_MAX;
     private int requestedHeartbeat                = DEFAULT_HEARTBEAT;
+    private AuthMechanism[] authMechanisms        = DEFAULT_AUTH_MECHANISMS;
     private Map<String, Object> _clientProperties = AMQConnection.defaultClientProperties();
     private SocketFactory factory                 = SocketFactory.getDefault();
 
@@ -261,6 +267,32 @@ public void setClientProperties(Map<String, Object> clientProperties) {
         _clientProperties = clientProperties;
     }
 
+    /**
+     * Given a list of mechanism names supported by the server, select a
+     * preferred mechanism, or null if we have none in common.
+     * @param serverMechanisms
+     * @return
+     */
+    public AuthMechanism getAuthMechanism(List<String> serverMechanisms) {
+        // Our list is in order of preference, the server one is not.
+        for (AuthMechanism mechanism : authMechanisms) {
+            if (serverMechanisms.contains(mechanism.getName())) {
+                return mechanism;
+            }
+        }
+
+        return null;
+    }
+
+    /**
+     * Set authentication mechanisms to use (in descending preference order)
+     * @param mechanisms
+     * @see #DEFAULT_AUTH_MECHANISMS
+     */
+    public void setAuthMechanisms(AuthMechanism[] mechanisms) {
+        this.authMechanisms = mechanisms;
+    }
+
     /**
      * Retrieve the socket factory used to make connections with.
      */

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

@@ -34,6 +34,8 @@
 import java.io.EOFException;
 import java.io.IOException;
 import java.net.SocketException;
+import java.util.Arrays;
+import java.util.List;
 import java.util.Map;
 import java.util.HashMap;
 import java.util.concurrent.TimeoutException;
@@ -44,6 +46,7 @@
 import com.rabbitmq.client.AMQP;
 import com.rabbitmq.client.Address;
 import com.rabbitmq.client.AlreadyClosedException;
+import com.rabbitmq.client.AuthMechanism;
 import com.rabbitmq.client.Channel;
 import com.rabbitmq.client.Command;
 import com.rabbitmq.client.Connection;
@@ -154,7 +157,7 @@ public void ensureIsOpen()
      */
     private int _heartbeat;
 
-    private final String _username, _password, _virtualHost;
+    private final String _virtualHost;
     private final int _requestedChannelMax, _requestedFrameMax, _requestedHeartbeat;
     private final Map<String, Object> _clientProperties;
 
@@ -202,8 +205,6 @@ public AMQConnection(ConnectionFactory factory,
     {
         checkPreconditions();
 
-        _username = factory.getUsername();
-        _password = factory.getPassword();
         _virtualHost = factory.getVirtualHost();
         _requestedChannelMax = factory.getRequestedChannelMax();
         _requestedFrameMax = factory.getRequestedFrameMax();
@@ -254,8 +255,9 @@ public void start()
         ml.setName("AMQP Connection " + getHost() + ":" + getPort());
         ml.start();
 
+        AMQP.Connection.Start connStart = null;
         try {
-            AMQP.Connection.Start connStart =
+            connStart =
                 (AMQP.Connection.Start) connStartBlocker.getReply().getMethod();
 
             _serverProperties = connStart.getServerProperties();
@@ -273,20 +275,15 @@ public void start()
         } catch (ShutdownSignalException sse) {
             throw AMQChannel.wrap(sse);
         }
-        
-        LongString saslResponse = LongStringHelper.asLongString("\0" + _username +
-                                                                "\0" + _password);
-        AMQImpl.Connection.StartOk startOk =
-            new AMQImpl.Connection.StartOk(_clientProperties, "PLAIN",
-                                           saslResponse, "en_US");
-        
-        AMQP.Connection.Tune connTune = null;
 
-        try {
-            connTune = (AMQP.Connection.Tune) _channel0.rpc(startOk).getMethod();
-        } catch (ShutdownSignalException e) {
-            throw AMQChannel.wrap(e, "Possibly caused by authentication failure");
+        List<String> mechanisms = Arrays.asList(
+                    connStart.getMechanisms().toString().split(" "));
+        AuthMechanism mechanism = _factory.getAuthMechanism(mechanisms);
+        if (mechanism == null) {
+            throw new IOException("No compatible authentication mechanism found - " +
+                    "server offered [" + connStart.getMechanisms() + "]");
         }
+        AMQP.Connection.Tune connTune = mechanism.doLogin(_channel0, _factory);
 
         int channelMax =
             negotiatedMaxValue(_factory.getRequestedChannelMax(),
@@ -714,6 +711,6 @@ public void close(int closeCode,
     }
 
     @Override public String toString() {
-        return "amqp://" + _username + "@" + getHost() + ":" + getPort() + _virtualHost;
+        return "amqp://" + _factory.getUsername() + "@" + getHost() + ":" + getPort() + _virtualHost;
     }
 }

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

@@ -0,0 +1,32 @@
+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 String getName() {
+        return "PLAIN";
+    }
+}