feat: send first sasl message

Author: ViniciusCestarii

internal/protocol/auth.go

@@ -2,12 +2,15 @@ package protocol
 
 import (
 	"crypto/md5"
+	"crypto/rand"
+	"encoding/base64"
 	"encoding/binary"
 	"encoding/hex"
 	"fmt"
 	"postgres-protocol-go/internal/pool"
 	"postgres-protocol-go/internal/protocol/messages"
 	"postgres-protocol-go/pkg/utils"
+	"strings"
 )
 
 func ProcessAuth(pgConnection PgConnection) error {
@@ -25,7 +28,6 @@ func ProcessAuth(pgConnection PgConnection) error {
 
 	authType := parseAuthType(answer)
 
-	// todo: implement SCRAM-SHA-256
 	switch authType {
 	case authenticationOk:
 		if pgConnection.isVerbose() {
@@ -49,6 +51,49 @@ func ProcessAuth(pgConnection PgConnection) error {
 				}
 			}
 		}
+	case authenticationSASL:
+		pgConnection.saslMethod = strings.Trim(string(answer[9:]), "\x00 \n\r")
+		switch pgConnection.saslMethod {
+		// https://datatracker.ietf.org/doc/html/rfc7677
+		case "SCRAM-SHA-256":
+			initialResponse, err := buildSCRAMInitialResponse(pgConnection.connConfig.User)
+			if err != nil {
+				return err
+			}
+
+			initialResponseBytes := []byte(initialResponse)
+
+			lengthBytes := make([]byte, 4)
+			binary.BigEndian.PutUint32(lengthBytes, uint32(len(initialResponseBytes)))
+			fmt.Println(lengthBytes)
+
+			buff := pool.NewWriteBuffer(1024)
+			buff.StartMessage(messages.SASLInitial)
+			buff.WriteString(pgConnection.saslMethod)
+			buff.WriteInt32(int32(len(initialResponseBytes)))
+			_, err = buff.Write(initialResponseBytes)
+			if err != nil {
+				return err
+			}
+			buff.FinishMessage()
+
+			err = pgConnection.sendMessage(buff)
+			if err != nil {
+				return err
+			}
+
+			return ProcessAuth(pgConnection)
+		default:
+			return fmt.Errorf("SASL authentication method %s is not supported", pgConnection.saslMethod)
+		}
+	case authenticationSASLContinue:
+		fmt.Println("SASLContinue", pgConnection.saslMethod)
+		switch pgConnection.saslMethod {
+		case "SCRAM-SHA-256":
+			return fmt.Errorf("not implemented")
+		default:
+			return fmt.Errorf("SASL authentication method %s is not supported", pgConnection.saslMethod)
+		}
 	case authenticationMD5Password:
 		if pgConnection.connConfig.Password == nil {
 			return fmt.Errorf("password is required for MD5 authentication")
@@ -119,3 +164,21 @@ func parseAuthType(message []byte) uint32 {
 func parseSalt(message []byte) string {
 	return string(message[9:13])
 }
+
+func buildSCRAMInitialResponse(username string) (string, error) {
+	nonce, err := generateNonce()
+	if err != nil {
+		return "", fmt.Errorf("failed to generate nonce: %v", err)
+	}
+	// Format: "n,,n=<username>,r=<nonce>"
+	initialResponse := fmt.Sprintf("n,,n=%s,r=%s", username, nonce)
+	return initialResponse, nil
+}
+
+func generateNonce() (string, error) {
+	nonceBytes := make([]byte, 16) // 16 bytes = 128 bits of randomness
+	if _, err := rand.Read(nonceBytes); err != nil {
+		return "", err
+	}
+	return base64.StdEncoding.EncodeToString(nonceBytes), nil
+}

internal/protocol/messages/messages.go

@@ -7,6 +7,8 @@ const (
 	Startup         = 0 // No identifier
 	SSL             = 0 // No identifier
 	Auth            = 'R'
+	SASLInitial     = 'p'
+	SASLResponse    = 'p'
 	Password        = 'p'
 	Error           = 'E'
 	SimpleQuery     = 'Q'

internal/protocol/pg_connection.go

@@ -14,6 +14,7 @@ import (
 
 type PgConnection struct {
 	conn        net.Conn
+	saslMethod  string
 	connConfig  models.ConnConfig
 	driveConfig models.DriveConfig
 }