added SSL request and SASL auth messages

authentication.go

@@ -10,9 +10,14 @@ type AuthenticationMethod int
 
 // Available authentication methods
 const (
-	AuthenticationMethodOK        AuthenticationMethod = 0
-	AuthenticationMethodPlaintext AuthenticationMethod = 3
-	AuthenticationMethodMD5       AuthenticationMethod = 5
+	AuthenticationMethodOK           AuthenticationMethod = 0
+	AuthenticationMethodPlaintext    AuthenticationMethod = 3
+	AuthenticationMethodMD5          AuthenticationMethod = 5
+	AuthenticationMethodSASL         AuthenticationMethod = 10
+	AuthenticationMethodSASLContinue AuthenticationMethod = 11
+	AuthenticationMethodSASLFinal    AuthenticationMethod = 12
+	SASLMechanismScramSHA256                              = "SCRAM-SHA-256"
+	SASLMechanismScramSHA256Plus                          = "SCRAM-SHA-256-PLUS"
 )
 
 func (a AuthenticationMethod) String() string {
@@ -23,6 +28,8 @@ func (a AuthenticationMethod) String() string {
 		return "Plaintext"
 	case AuthenticationMethodMD5:
 		return "MD5"
+	case AuthenticationMethodSASL:
+		return "SASL"
 	}
 
 	return "Unknown"
@@ -31,8 +38,11 @@ func (a AuthenticationMethod) String() string {
 // AuthenticationRequest is a server response either asking the client to authenticate or
 // used to indicate that authentication was successful
 type AuthenticationRequest struct {
-	Method AuthenticationMethod
-	Salt   []byte
+	Method                   AuthenticationMethod
+	Salt                     []byte
+	SupportedScramSHA256     bool
+	SupportedScramSHA256Plus bool
+	Message                  []byte
 }
 
 func (a *AuthenticationRequest) server() {}
@@ -57,36 +67,80 @@ func ParseAuthenticationRequest(r io.Reader) (*AuthenticationRequest, error) {
 	if err != nil {
 		return nil, err
 	}
-	m := AuthenticationMethod(i)
-	if m != AuthenticationMethodOK && m != AuthenticationMethodPlaintext && m != AuthenticationMethodMD5 {
-		return nil, fmt.Errorf("received unknown authentication request method number %d", m)
-	}
-
-	a := &AuthenticationRequest{
-		Method: m,
-	}
-
-	if a.Method == AuthenticationMethodMD5 {
+	switch m := AuthenticationMethod(i); m {
+	case AuthenticationMethodOK, AuthenticationMethodPlaintext:
+		return &AuthenticationRequest{Method: m}, nil
+	case AuthenticationMethodMD5:
+		a := &AuthenticationRequest{
+			Method: m,
+		}
 		a.Salt, err = buf.ReadString(false)
 		if err != nil {
 			return nil, err
 		}
 		if len(a.Salt) != 4 {
 			return nil, fmt.Errorf("expected salt of length 4")
 		}
-	}
+		return a, nil
+	case AuthenticationMethodSASL:
+		a := &AuthenticationRequest{
+			Method: m,
+		}
 
-	return a, nil
+		for {
+			supportedSaslMechanism, err := buf.ReadString(true)
+			if err != nil {
+				return nil, err
+			}
+			switch string(supportedSaslMechanism) {
+			case SASLMechanismScramSHA256:
+				a.SupportedScramSHA256 = true
+			case SASLMechanismScramSHA256Plus:
+				a.SupportedScramSHA256Plus = true
+			case "":
+				return a, nil
+			default:
+				return nil, fmt.Errorf("server supports unknown SASL mechanism")
+			}
+		}
+	case AuthenticationMethodSASLContinue, AuthenticationMethodSASLFinal:
+
+		message, err := buf.ReadString(true)
+		if err != nil {
+			return nil, err
+		}
+		a := &AuthenticationRequest{
+			Method:  m,
+			Message: message,
+		}
+
+		return a, nil
+
+	default:
+		return nil, fmt.Errorf("received unknown authentication request method number %d", m)
+	}
 }
 
 // Encode will return the byte representation of this AuthenticationRequest message
 func (a *AuthenticationRequest) Encode() []byte {
 	// 'R' [int32 - length] [int32 - method] [other - optional]
 	w := newWriteBuffer()
 	w.WriteInt(int(a.Method))
-	if a.Method == AuthenticationMethodMD5 {
+	switch a.Method {
+	case AuthenticationMethodMD5:
 		w.WriteString(a.Salt, false)
+	case AuthenticationMethodSASL:
+		if a.SupportedScramSHA256 {
+			w.WriteString([]byte(SASLMechanismScramSHA256), true)
+		}
+		if a.SupportedScramSHA256Plus {
+			w.WriteString([]byte(SASLMechanismScramSHA256Plus), true)
+		}
+		w.WriteByte('\x00')
+	case AuthenticationMethodSASLContinue, AuthenticationMethodSASLFinal, AuthenticationMethodOK:
+		w.WriteBytes(a.Message)
 	}
+
 	w.Wrap('R')
 	return w.Bytes()
 }
@@ -98,14 +152,20 @@ func (a *AuthenticationRequest) Encode() []byte {
 //     "Payload": map[string]interface{}{
 //       "Method": <AuthenticationRequest.Method>,
 //       "Salt": <AuthenticationRequest.Salt>,
+//       "ScramSHA256": <AuthenticationRequest.ScramSHA256>,
+//       "ScramSHA256Plus": <AuthenticationRequest.ScramSHA256Plus>,
+//       "Message": <AuthenticationRequest.Message>,
 //     },
 //   }
 func (a *AuthenticationRequest) AsMap() map[string]interface{} {
 	return map[string]interface{}{
 		"Type": "AuthenticationRequest",
 		"Payload": map[string]interface{}{
-			"Method": int(a.Method),
-			"Salt":   a.Salt,
+			"Method":          int(a.Method),
+			"Salt":            a.Salt,
+			"ScramSHA256":     a.SupportedScramSHA256,
+			"ScramSHA256Plus": a.SupportedScramSHA256Plus,
+			"Message":         a.Message,
 		},
 	}
 }

error.go

@@ -124,7 +124,7 @@ func encodeError(e *Error, tag byte) []byte {
 func errorMap(e *Error, name string) map[string]interface{} {
 	return map[string]interface{}{
 		"Type": name,
-		"Payload": map[string]string{
+		"Payload": map[string]interface{}{
 			"Severity": string(e.Severity),
 			"Text":     string(e.Text),
 			"Code":     string(e.Code),

message.go

@@ -1,6 +1,7 @@
 package pgproto
 
 import (
+	"bytes"
 	"fmt"
 	"io"
 )
@@ -26,91 +27,70 @@ type ServerMessage interface {
 
 // ParseClientMessage will read the next ClientMessage from the provided io.Reader
 func ParseClientMessage(r io.Reader) (ClientMessage, error) {
-	// Create a buffer
-	buf := newReadBuffer(r)
 
-	// Look at the first byte to determine the type of message we have
-	start, err := buf.ReadByte()
+	msg, err := readRawMessage(r)
 	if err != nil {
 		return nil, err
 	}
+	start := msg[0]
+	msgReader := bytes.NewReader(msg)
 
 	// Startup message:
 	//   [int32 - length] [int32 - protocol] [[string]\0[string]\0] \0
 	// Regular message
 	//   [char - tag] [int32 - length] [payload]
 	switch start {
-	// TODO: We need to handle this case better, it might not always start with \x00
-	//       We could just make calling `ParseStartupMessage` explicit
 	case '\x00':
-		msgReader, err := readStartupMessage(start, buf)
-		if err != nil {
-			return nil, err
-		}
+		// TODO: We need to handle this case better, it might not always start with \x00
+		//       We could just make calling `ParseStartupMessage` explicit
 		return ParseStartupMessage(msgReader)
+	case 'p':
+		// Password message
+		return ParsePasswordMessage(msgReader)
+	case 'Q':
+		// Simple query
+		return ParseSimpleQuery(msgReader)
+	case 't':
+		// Parameter description
+		return ParseParameterDescription(msgReader)
+	case 'B':
+		// Binary parameters
+		return ParseBinaryParameters(msgReader)
+	case 'P':
+		// Parse
+		return ParseParse(msgReader)
+	case 'E':
+		// Execute
+		return ParseExecute(msgReader)
+	case 'H':
+		// Flush
+		return ParseFlush(msgReader)
+	case 'S':
+		// Sync
+		return ParseSync(msgReader)
+	case 'C':
+		// Close
+		return ParseClose(msgReader)
+	case 'D':
+		// Describe
+		return ParseDescribe(msgReader)
+	case 'X':
+		// Termination
+		return ParseTermination(msgReader)
 	default:
-		// Read the entire next message from the input reader
-		msgReader, err := readMessage(start, buf)
-		if err != nil {
-			return nil, err
-		}
-		switch start {
-		case 'p':
-			// Password message
-			return ParsePasswordMessage(msgReader)
-		case 'Q':
-			// Simple query
-			return ParseSimpleQuery(msgReader)
-		case 't':
-			// Parameter description
-			return ParseParameterDescription(msgReader)
-		case 'B':
-			// Binary parameters
-			return ParseBinaryParameters(msgReader)
-		case 'P':
-			// Parse
-			return ParseParse(msgReader)
-		case 'E':
-			// Execute
-			return ParseExecute(msgReader)
-		case 'H':
-			// Flush
-			return ParseFlush(msgReader)
-		case 'S':
-			// Sync
-			return ParseSync(msgReader)
-		case 'C':
-			// Close
-			return ParseClose(msgReader)
-		case 'D':
-			// Describe
-			return ParseDescribe(msgReader)
-		case 'X':
-			// Termination
-			return ParseTermination(msgReader)
-		default:
-			return nil, fmt.Errorf("unknown message tag '%c'", start)
-		}
+		return nil, fmt.Errorf("unknown message tag '%c'", start)
 	}
 }
 
 // ParseServerMessage will read the next ServerMessage from the provided io.Reader
 func ParseServerMessage(r io.Reader) (ServerMessage, error) {
-	// Create a buffer
-	buf := newReadBuffer(r)
 
-	// Look at the first byte to determine the type of message we have
-	start, err := buf.ReadByte()
+	msg, err := readRawMessage(r)
 	if err != nil {
 		return nil, err
 	}
-
-	// Read the entire next message from the input reader
-	msgReader, err := readMessage(start, buf)
-	if err != nil {
-		return nil, err
-	}
-
+	start := msg[0]
+	msgReader := bytes.NewReader(msg)
 	// Message
 	//   [char - tag] [int32 - length] [payload]
 	switch start {
@@ -182,54 +162,25 @@ func ParseServerMessage(r io.Reader) (ServerMessage, error) {
 	}
 }
 
-func readStartupMessage(start byte, buf *readBuffer) (io.Reader, error) {
-	// [int32 - length] [payload]
-	// StartupMessage
-	// Read the next 3 bytes, prepend with the 1 we already read to parse the length from this message
-	s := [4]byte{
-		start,
-	}
-	_, err := buf.Read(s[1:])
-	if err != nil {
-		return nil, err
-	}
-	l := bytesToInt(s[:])
+func readRawMessage(r io.Reader) (rawmsg []byte, err error) {
 
-	// Read the rest of the message into a []byte
-	// DEV: Subtract 4 to account for the length of the in32 we just read
-	b := make([]byte, l-4)
-	_, err = buf.Read(b)
+	startByte, err := ReadNBytes(r, 1)
 	if err != nil {
 		return nil, err
 	}
-
-	// Rebuild the message into a []byte
-	w := newWriteBuffer()
-	w.WriteInt(l)
-	w.WriteBytes(b)
-	return w.Reader(), nil
-}
-
-func readMessage(start byte, buf *readBuffer) (io.Reader, error) {
-	// [char tag] [int32 length] [payload]
-	// Parse length from the message
-	l, err := buf.ReadInt()
+	rawPkgLen, err := ReadNBytes(r, 4)
 	if err != nil {
 		return nil, err
 	}
-
-	// Read the rest of the message into a []byte
-	// DEV: Subtract 4 to account for the length of the int32 we just read
-	b := make([]byte, l-4)
-	_, err = buf.Read(b)
+	pkgLen := bytesToInt(rawPkgLen)
+	payload, err := ReadNBytes(r, pkgLen-4)
 	if err != nil {
 		return nil, err
 	}
+	rawMessage := make([]byte, pkgLen+1)
+	copy(rawMessage, startByte)
+	copy(rawMessage[1:], rawPkgLen)
+	copy(rawMessage[5:], payload)
+	return rawMessage, nil
 
-	// Rebuild the message into a []byte
-	w := newWriteBuffer()
-	w.WriteByte(start)
-	w.WriteInt(l)
-	w.WriteBytes(b)
-	return w.Reader(), nil
 }