SASL SCRAM implementation started

Author: Marat Gainullin

src/main/java/com/github/pgasync/io/backend/AuthenticationDecoder.java

@@ -15,14 +15,15 @@
 package com.github.pgasync.io.backend;
 
 import com.github.pgasync.io.Decoder;
+import com.github.pgasync.io.IO;
 import com.github.pgasync.message.backend.Authentication;
 
 import java.nio.ByteBuffer;
 import java.nio.charset.Charset;
 
 /**
  * See <a href="www.postgresql.org/docs/9.3/static/protocol-message-formats.html">Postgres message formats</a>
- * 
+ *
  * <pre>
  * AuthenticationOk (B)
  *  Byte1('R')
@@ -31,7 +32,7 @@
  *      Length of message contents in bytes, including self.
  *  Int32(0)
  *      Specifies that the authentication was successful.
- *       
+ *
  * AuthenticationMD5Password (B)
  *  Byte1('R')
  *      Identifies the message as an authentication request.
@@ -42,14 +43,22 @@
  *  Byte4
  *      The salt to use when encrypting the password.
  * </pre>
- * 
+ *
  * @author Antti Laisi
  */
 public class AuthenticationDecoder implements Decoder<Authentication> {
 
     private static final int OK = 0;
+    private static final int SASL_FINAL = 12;
+    private static final int SASL_CONTINUE = 11;
+    private static final int SASL = 10;
+    private static final int AUTHENTICATION_SSPI = 9;
+    private static final int AUTHENTICATION_GSS = 7;
+    private static final int AUTHENTICATION_SCM_CREDENTIAL = 6;
     private static final int PASSWORD_MD5_CHALLENGE = 5;
     private static final int CLEARTEXT_PASSWORD = 3;
+    private static final int AUTHENTICATION_KERBEROS_V5 = 2;
+    public static final String AUTHENTICATION_IS_NOT_SUPPORTED = " authentication is not supported";
 
     @Override
     public byte getMessageId() {
@@ -60,14 +69,40 @@ public byte getMessageId() {
     public Authentication read(ByteBuffer buffer, Charset encoding) {
         int type = buffer.getInt();
         switch (type) {
-            case OK:
-                return new Authentication(true, null);
             case CLEARTEXT_PASSWORD:
-                return new Authentication(false, null);
+                return Authentication.CLEAR_TEXT;
             case PASSWORD_MD5_CHALLENGE:
                 byte[] salt = new byte[4];
                 buffer.get(salt);
-                return new Authentication(false, salt);
+                return new Authentication(false, false, salt);
+            case SASL:
+                boolean scramSha256Met = false;
+                String sasl = IO.getCString(buffer, encoding);
+                while (!sasl.isEmpty()) {
+                    if ("SCRAM-SHA-256".equals(sasl)) {
+                        scramSha256Met = true;
+                    }
+                    sasl = IO.getCString(buffer, encoding);
+                }
+                if (scramSha256Met) {
+                    return Authentication.SCRAM_SHA_256;
+                } else {
+                    throw new UnsupportedOperationException("The server doesn't support " + Authentication.SUPPORTED_SASL + " SASL mechanism");
+                }
+            case AUTHENTICATION_SSPI:
+                throw new UnsupportedOperationException(AUTHENTICATION_IS_NOT_SUPPORTED);
+            case AUTHENTICATION_GSS:
+                throw new UnsupportedOperationException("AuthenticationGss" + AUTHENTICATION_IS_NOT_SUPPORTED);
+            case AUTHENTICATION_SCM_CREDENTIAL:
+                throw new UnsupportedOperationException("AuthenticationScmCredential" + AUTHENTICATION_IS_NOT_SUPPORTED);
+            case AUTHENTICATION_KERBEROS_V5:
+                throw new UnsupportedOperationException("AuthenticationKerberosV5" + AUTHENTICATION_IS_NOT_SUPPORTED);
+            case OK:
+                return Authentication.OK;
+            case SASL_CONTINUE:
+                return Authentication.OK;
+            case SASL_FINAL:
+                return Authentication.OK;
             default:
                 throw new UnsupportedOperationException("Unsupported authentication type: " + type);
         }

src/main/java/com/github/pgasync/io/frontend/FIndicatorsEncoder.java

@@ -15,6 +15,8 @@
 package com.github.pgasync.io.frontend;
 
 import com.github.pgasync.io.Encoder;
+import com.github.pgasync.io.IO;
+import com.github.pgasync.message.backend.Authentication;
 import com.github.pgasync.message.frontend.FIndicators;
 
 import java.nio.ByteBuffer;
@@ -46,6 +48,9 @@ public void write(FIndicators msg, ByteBuffer buffer, Charset encoding) {
             case SYNC:
                 sync(buffer);
                 break;
+            case SASL_INITIAL:
+                saslInitial(buffer, encoding);
+                break;
             default:
                 throw new IllegalStateException(msg.name());
         }
@@ -56,4 +61,12 @@ void sync(ByteBuffer buffer) {
         buffer.putInt(4);
     }
 
+    void saslInitial(ByteBuffer buffer, Charset encoding) {
+        buffer.put((byte) 'p');
+        buffer.putInt(0);
+        IO.putCString(buffer, Authentication.SUPPORTED_SASL, encoding);
+        buffer.putInt(-1);
+        buffer.putInt(1, buffer.position() - 1);
+    }
+
 }

src/main/java/com/github/pgasync/io/frontend/SASLInitialResponseEncoder.java

@@ -0,0 +1,44 @@
+package com.github.pgasync.io.frontend;
+
+import com.github.pgasync.io.IO;
+import com.github.pgasync.message.frontend.SASLInitialResponse;
+
+import java.nio.ByteBuffer;
+import java.nio.charset.Charset;
+import java.nio.charset.StandardCharsets;
+import java.security.SecureRandom;
+
+public class SASLInitialResponseEncoder extends SkipableEncoder<SASLInitialResponse> {
+
+    private static String saslPrep(String value){
+        // TODO: Implement me
+        return value;
+    }
+
+    @Override
+    protected byte getMessageId() {
+        return 'p';
+    }
+
+    @Override
+    protected void writeBody(SASLInitialResponse msg, ByteBuffer buffer, Charset encoding) {
+        IO.putCString(buffer, msg.getSaslMechanism(), encoding);
+        StringBuilder clientFirstMessage = new StringBuilder();
+        if (msg.getChannelBindingType() != null && !msg.getChannelBindingType().isBlank()) {
+            clientFirstMessage.append("p=").append(msg.getChannelBindingType()).append(",");
+        } else {
+            clientFirstMessage.append("n,");
+        }
+        clientFirstMessage.append(",");// It could be 'clientFirstMessage.append("a=").append(msg.getUsername()).append(",");' but it is not needed unless we are using impersonate techniques.
+        clientFirstMessage.append("n=,"); // Postgres expects the username here to be an empty string because of the startup message.
+        clientFirstMessage.append("r=").append(msg.getNonce());
+        byte[] clientFirstMessageContent = clientFirstMessage.toString().getBytes(encoding); // RFC dictates us to use UTF-8 here
+        buffer.putInt(clientFirstMessageContent.length);
+        buffer.put(clientFirstMessageContent);
+    }
+
+    @Override
+    public Class<SASLInitialResponse> getMessageType() {
+        return SASLInitialResponse.class;
+    }
+}

src/main/java/com/github/pgasync/message/backend/Authentication.java

@@ -23,11 +23,18 @@
  */
 public class Authentication implements Message {
 
+    public static final Authentication OK = new Authentication(true, false, null);
+    public static final Authentication CLEAR_TEXT = new Authentication(false, false, null);
+    public static final Authentication SCRAM_SHA_256 = new Authentication(false, true, null);
+    public static final String SUPPORTED_SASL = "SCRAM-SHA-256";
+
     private final boolean success;
+    private final boolean scramSha256;
     private final byte[] md5salt;
 
-    public Authentication(boolean success, byte[] md5salt) {
+    public Authentication(boolean success, boolean scramSha256, byte[] md5salt) {
         this.success = success;
+        this.scramSha256 = scramSha256;
         this.md5salt = md5salt;
     }
 
@@ -39,8 +46,12 @@ public boolean isAuthenticationOk() {
         return success;
     }
 
+    public boolean isScramSha256() {
+        return scramSha256;
+    }
+
     @Override
     public String toString() {
-        return String.format("Authentication(success=%s, md5salt=%s)", success, md5salt != null ? printHexBinary(md5salt) : null);
+        return String.format("Authentication(success=%s, md5salt=%s, scramSha256=%s)", success, md5salt != null ? printHexBinary(md5salt) : null, scramSha256);
     }
 }

src/main/java/com/github/pgasync/message/backend/AuthenticationSaslContinue.java

@@ -0,0 +1,16 @@
+package com.github.pgasync.message.backend;
+
+import com.github.pgasync.message.Message;
+
+public class AuthenticationSaslContinue implements Message {
+
+    private final byte[] saslData;
+
+    public AuthenticationSaslContinue(byte[] saslData) {
+        this.saslData = saslData;
+    }
+
+    public byte[] getSaslData() {
+        return saslData;
+    }
+}

src/main/java/com/github/pgasync/message/backend/SSLHandshake.java renamed to src/main/java/com/github/pgasync/message/backend/SslHandshake.java

@@ -5,6 +5,6 @@
 /**
  * @author Marat Gainullin
  */
-public enum SSLHandshake implements Message {
+public enum SslHandshake implements Message {
     INSTANCE
 }

src/main/java/com/github/pgasync/message/frontend/FIndicators.java

@@ -20,5 +20,6 @@
  * @author Marat Gainullin
  */
 public enum FIndicators implements Message {
-    SYNC
+    SYNC,
+    SASL_INITIAL
 }

src/main/java/com/github/pgasync/message/frontend/SASLInitialResponse.java

@@ -0,0 +1,38 @@
+package com.github.pgasync.message.frontend;
+
+import com.github.pgasync.message.Message;
+
+/**
+ * it is a bit weired to have a username both here and in a {@link StartupMessage}.
+ * But according to the SCRAM specification in RFC5802 we have to have it :(
+ */
+public class SASLInitialResponse implements Message {
+
+    private final String saslMechanism;
+    private final String channelBindingType;
+    private final String username;
+    private final String nonce;
+
+    public SASLInitialResponse(String saslMechanism, String channelBindingType, String username, String nonce) {
+        this.saslMechanism = saslMechanism;
+        this.channelBindingType = channelBindingType;
+        this.username = username;
+        this.nonce = nonce;
+    }
+
+    public String getUsername() {
+        return username;
+    }
+
+    public String getChannelBindingType() {
+        return channelBindingType;
+    }
+
+    public String getSaslMechanism() {
+        return saslMechanism;
+    }
+
+    public String getNonce() {
+        return nonce;
+    }
+}

src/test/java/com/github/pgasync/netty/NettyMessageEncoder.java

@@ -15,17 +15,7 @@
 package com.github.pgasync.netty;
 
 import com.github.pgasync.io.Encoder;
-import com.github.pgasync.io.frontend.BindEncoder;
-import com.github.pgasync.io.frontend.CloseEncoder;
-import com.github.pgasync.io.frontend.DescribeEncoder;
-import com.github.pgasync.io.frontend.ExecuteEncoder;
-import com.github.pgasync.io.frontend.FIndicatorsEncoder;
-import com.github.pgasync.io.frontend.ParseEncoder;
-import com.github.pgasync.io.frontend.PasswordMessageEncoder;
-import com.github.pgasync.io.frontend.QueryEncoder;
-import com.github.pgasync.io.frontend.SSLRequestEncoder;
-import com.github.pgasync.io.frontend.StartupMessageEncoder;
-import com.github.pgasync.io.frontend.TerminateEncoder;
+import com.github.pgasync.io.frontend.*;
 import com.github.pgasync.message.Message;
 import io.netty.buffer.ByteBuf;
 import io.netty.channel.ChannelHandlerContext;
@@ -54,6 +44,7 @@ public class NettyMessageEncoder extends MessageToByteEncoder<Message> {
             new BindEncoder(),
             new DescribeEncoder(),
             new ExecuteEncoder(),
+            new SASLInitialResponseEncoder(),
             new CloseEncoder(),
             new FIndicatorsEncoder(),
             new TerminateEncoder()

src/test/java/com/github/pgasync/netty/NettyPgProtocolStream.java

@@ -16,7 +16,7 @@
 
 import com.github.pgasync.PgProtocolStream;
 import com.github.pgasync.message.Message;
-import com.github.pgasync.message.backend.SSLHandshake;
+import com.github.pgasync.message.backend.SslHandshake;
 import com.github.pgasync.message.frontend.SSLRequest;
 import com.github.pgasync.message.frontend.StartupMessage;
 import com.github.pgasync.message.frontend.Terminate;
@@ -76,7 +76,7 @@ public CompletableFuture<Message> connect(StartupMessage startup) {
                 .thenApply(this::send)
                 .thenCompose(Function.identity())
                 .thenApply(message -> {
-                    if (message == SSLHandshake.INSTANCE) {
+                    if (message == SslHandshake.INSTANCE) {
                         return send(startup);
                     } else {
                         return CompletableFuture.completedFuture(message);
@@ -171,7 +171,7 @@ public void channelActive(ChannelHandlerContext context) {
             @Override
             public void userEventTriggered(ChannelHandlerContext context, Object evt) {
                 if (evt instanceof SslHandshakeCompletionEvent && ((SslHandshakeCompletionEvent) evt).isSuccess()) {
-                    gotMessage(SSLHandshake.INSTANCE);
+                    gotMessage(SslHandshake.INSTANCE);
                 }
             }