build source

vendor/github.com/jackc/pgx/v5/pgconn/auth_oauth.go

@@ -0,0 +1,67 @@
+package pgconn
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+
+	"github.com/jackc/pgx/v5/pgproto3"
+)
+
+func (c *PgConn) oauthAuth(ctx context.Context) error {
+	if c.config.OAuthTokenProvider == nil {
+		return errors.New("OAuth authentication required but no token provider configured")
+	}
+
+	token, err := c.config.OAuthTokenProvider(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to obtain OAuth token: %w", err)
+	}
+
+	// https://www.rfc-editor.org/rfc/rfc7628.html#section-3.1
+	initialResponse := []byte("n,,\x01auth=Bearer " + token + "\x01\x01")
+
+	saslInitialResponse := &pgproto3.SASLInitialResponse{
+		AuthMechanism: "OAUTHBEARER",
+		Data:          initialResponse,
+	}
+	c.frontend.Send(saslInitialResponse)
+	err = c.flushWithPotentialWriteReadDeadlock()
+	if err != nil {
+		return err
+	}
+
+	msg, err := c.receiveMessage()
+	if err != nil {
+		return err
+	}
+
+	switch m := msg.(type) {
+	case *pgproto3.AuthenticationOk:
+		return nil
+	case *pgproto3.AuthenticationSASLContinue:
+		// Server sent error response in SASL continue
+		// https://www.rfc-editor.org/rfc/rfc7628.html#section-3.2.2
+		// https://www.rfc-editor.org/rfc/rfc7628.html#section-3.2.3
+		errResponse := struct {
+			Status              string `json:"status"`
+			Scope               string `json:"scope"`
+			OpenIDConfiguration string `json:"openid-configuration"`
+		}{}
+		err := json.Unmarshal(m.Data, &errResponse)
+		if err != nil {
+			return fmt.Errorf("invalid OAuth error response from server: %w", err)
+		}
+
+		// Per RFC 7628 section 3.2.3, we should send a SASLResponse which only contains \x01.
+		// However, since the connection will be closed anyway, we can skip this
+		return fmt.Errorf("OAuth authentication failed: %s", errResponse.Status)
+
+	case *pgproto3.ErrorResponse:
+		return ErrorResponseToPgError(m)
+
+	default:
+		return fmt.Errorf("unexpected message type during OAuth auth: %T", msg)
+	}
+}