Authenticate your backend

package main

import (
	"context"
	"crypto"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
	"uuid"

	"github.com/lestrrat-go/jwx/v4/jwa"
	"github.com/lestrrat-go/jwx/v4/jwk"
	"github.com/lestrrat-go/jwx/v4/jws"
)

// Values that you get from Gnosis Pay. The issuer is the "issuer" field
// of GET <issuer>/.well-known/oauth-authorization-server.
const (
	issuer         = "https://core.sandbox.gnosispay.in/"
	clientID       = "<PARTNER_CLIENT_ID>"
	privateKeyPath = "private.pem"
)

const clientAssertionType = "urn:ietf:params:oauth:client-assertion-type:jwt-bearer"

type tokenResponse struct {
	AccessToken string `json:"access_token"`
	TokenType   string `json:"token_type"`
	ExpiresIn   int64  `json:"expires_in"`
	Scope       string `json:"scope"`
}

type tokenErrorResponse struct {
	Error            string `json:"error"`
	ErrorDescription string `json:"error_description"`
}

func main() {
	ctx := context.Background()

	key, err := loadPrivateKey(privateKeyPath)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	assertion, err := signClientAssertion(key)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	token, err := requestAccessToken(ctx, assertion)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	// Step 4: print the access token. Only the token goes to stdout, so that
	// PARTNER_API_TOKEN=$(go run .) captures it.
	fmt.Fprintf(os.Stderr, "token expires in %d seconds, scope %q\n", token.ExpiresIn, token.Scope)
	fmt.Println(token.AccessToken)
}

// Step 1: load the private key and set its kid. The kid is the RFC 7638 JWK
// thumbprint of the public key (SHA-256, base64url without padding). It is
// equal to the kid that Gnosis Pay returned when it registered the public key.
func loadPrivateKey(path string) (jwk.Key, error) {
	pemBytes, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("read private key: %w", err)
	}

	key, err := jwk.ParseKey(pemBytes, jwk.WithX509(true))
	if err != nil {
		return nil, fmt.Errorf("parse private key: %w", err)
	}

	publicKey, err := key.PublicKey()
	if err != nil {
		return nil, fmt.Errorf("derive public key: %w", err)
	}
	thumbprint, err := publicKey.Thumbprint(crypto.SHA256)
	if err != nil {
		return nil, fmt.Errorf("compute kid: %w", err)
	}
	kid := base64.RawURLEncoding.EncodeToString(thumbprint)
	if err := key.Set(jwk.KeyIDKey, kid); err != nil {
		return nil, fmt.Errorf("set kid: %w", err)
	}
	return key, nil
}

// Step 2: sign the client assertion. The claims are a JSON object, so that
// "aud" is a single string and not an array.
func signClientAssertion(key jwk.Key) (string, error) {
	now := time.Now()
	claims, err := json.Marshal(map[string]any{
		"iss": clientID,
		"sub": clientID,
		"aud": issuer,
		"jti": uuid.NewV7().String(),
		"iat": now.Unix(),
		"exp": now.Add(time.Minute).Unix(),
	})
	if err != nil {
		return "", fmt.Errorf("encode claims: %w", err)
	}

	kid, _ := key.KeyID()
	headers := jws.NewHeaders()
	if err := headers.Set(jws.KeyIDKey, kid); err != nil {
		return "", fmt.Errorf("set kid header: %w", err)
	}
	if err := headers.Set(jws.TypeKey, "client-authentication+jwt"); err != nil {
		return "", fmt.Errorf("set typ header: %w", err)
	}

	// The server accepts only alg "EdDSA" for an Ed25519 key, so use jwa.EdDSA()
	// and not jwa.EdDSAEd25519(), which writes alg "Ed25519". Use jwa.ES256()
	// for an EC P-256 key.
	signed, err := jws.Sign(claims, jws.WithKey(jwa.EdDSA(), key, jws.WithProtectedHeaders(headers)))
	if err != nil {
		return "", fmt.Errorf("sign client assertion: %w", err)
	}
	return string(signed), nil
}

// Step 3: exchange the client assertion for an access token.
func requestAccessToken(ctx context.Context, assertion string) (tokenResponse, error) {
	form := url.Values{
		"grant_type":            {"client_credentials"},
		"client_assertion_type": {clientAssertionType},
		"client_assertion":      {assertion},
	}

	tokenEndpoint := strings.TrimSuffix(issuer, "/") + "/oauth/token"
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenEndpoint, strings.NewReader(form.Encode()))
	if err != nil {
		return tokenResponse{}, fmt.Errorf("build token request: %w", err)
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return tokenResponse{}, fmt.Errorf("send token request: %w", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return tokenResponse{}, fmt.Errorf("read token response: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		var tokenErr tokenErrorResponse
		if err := json.Unmarshal(body, &tokenErr); err != nil {
			return tokenResponse{}, fmt.Errorf("token request failed with status %d", resp.StatusCode)
		}
		return tokenResponse{}, errors.New("token request failed: " + tokenErr.Error + ": " + tokenErr.ErrorDescription)
	}

	var token tokenResponse
	if err := json.Unmarshal(body, &token); err != nil {
		return tokenResponse{}, fmt.Errorf("decode token response: %w", err)
	}
	return token, nil
}