Building an OAuth2 Authorization Server from Scratch

package authserver
 
import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
 
	"github.com/golang-jwt/jwt/v5"
)
 
// KeyManager handles RSA key generation and storage.
type KeyManager struct {
	privateKey *rsa.PrivateKey
	publicKey  *rsa.PublicKey
	keyID      string
}
 
// NewKeyManager generates a fresh 2048-bit RSA key pair.
func NewKeyManager() (*KeyManager, error) {
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, fmt.Errorf("key generation: %w", err)
	}
 
	return &KeyManager{
		privateKey: privateKey,
		publicKey:  &privateKey.PublicKey,
		keyID:      "default-key",
	}, nil
}
 
// ExportPublicKey returns the public key in PEM format for JWKS.
func (km *KeyManager) ExportPublicKey() (string, error) {
	pubKeyBytes, err := x509.MarshalPKIXPublicKey(km.publicKey)
	if err != nil {
		return "", fmt.Errorf("marshal public key: %w", err)
	}
 
	pubKeyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: pubKeyBytes,
	})
 
	return string(pubKeyPEM), nil
}
 
// IssueJWT signs claims with the server's RSA private key using RS256 and
// returns the compact JWT string. The key ID is set on the header so
// verifiers can pick the right JWK out of the JWKS response.
func (km *KeyManager) IssueJWT(claims jwt.MapClaims) (string, error) {
	token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
	token.Header["kid"] = km.keyID
 
	signed, err := token.SignedString(km.privateKey)
	if err != nil {
		return "", fmt.Errorf("sign jwt: %w", err)
	}
 
	return signed, nil
}

// AuthorizationRequest holds incoming authorization request parameters.
type AuthorizationRequest struct {
	ClientID            string
	RedirectURI         string
	CodeChallenge       string
	CodeChallengeMethod string // only "S256" is accepted below
	State               string // CSRF protection
}
 
// AuthorizationCode stores state needed for token exchange.
type AuthorizationCode struct {
	Code                string
	ClientID            string
	RedirectURI         string
	CodeChallenge       string
	CodeChallengeMethod string
	ExpiresAt           time.Time
	Used                bool
}
 
// AuthorizationEndpoint validates the request and generates a code.
func (as *AuthServer) AuthorizationEndpoint(w http.ResponseWriter, r *http.Request) {
	q := r.URL.Query()
	clientID := q.Get("client_id")
	redirectURI := q.Get("redirect_uri")
	// Exact match against the URIs registered for this client. On failure,
	// show an error page; never redirect to an unverified URI.
	if !as.clients.RedirectURIAllowed(r.Context(), clientID, redirectURI) {
		http.Error(w, "unknown client_id or unregistered redirect_uri", http.StatusBadRequest)
		return
	}
 
	// PKCE: code_challenge is required and only S256 is accepted. A missing
	// method means "plain", which the token endpoint below does not implement.
	codeChallenge := q.Get("code_challenge")
	if codeChallenge == "" || q.Get("code_challenge_method") != "S256" {
		http.Error(w, "PKCE code_challenge with method S256 required", http.StatusBadRequest)
		return
	}
 
	// Store the authorization code with PKCE challenge
	code, err := generateRandomCode(32)
	if err != nil {
		http.Error(w, "server error", http.StatusInternalServerError)
		return
	}
	authCode := AuthorizationCode{
		Code:                code,
		ClientID:            clientID,
		RedirectURI:         redirectURI,
		CodeChallenge:       codeChallenge,
		CodeChallengeMethod: "S256",
		ExpiresAt:           time.Now().Add(10 * time.Minute),
		Used:                false,
	}
	as.storeAuthCode(code, authCode)
 
	// Redirect back with code and state, keeping any query component the
	// registered URI already has instead of string-concatenating "?code=".
	u, err := url.Parse(redirectURI)
	if err != nil {
		http.Error(w, "server error", http.StatusInternalServerError)
		return
	}
	params := u.Query()
	params.Set("code", code)
	if state := q.Get("state"); state != "" {
		params.Set("state", state)
	}
	u.RawQuery = params.Encode()
	http.Redirect(w, r, u.String(), http.StatusFound)
}
 
func generateRandomCode(length int) (string, error) {
	b := make([]byte, length)
	// crypto/rand can fail; never ignore it — a short read would yield a
	// weak or all-zero token. Propagate the error so the caller aborts.
	if _, err := rand.Read(b); err != nil {
		return "", fmt.Errorf("generate random code: %w", err)
	}
	return base64.RawURLEncoding.EncodeToString(b), nil
}

// TokenRequest holds the code exchange request.
type TokenRequest struct {
	Code         string
	CodeVerifier string
	ClientID     string
}
 
// TokenResponse holds the issued tokens.
type TokenResponse struct {
	AccessToken  string `json:"access_token"`
	RefreshToken string `json:"refresh_token"`
	ExpiresIn    int    `json:"expires_in"`
	TokenType    string `json:"token_type"`
}
 
// TokenEndpoint exchanges an authorization code for tokens.
func (as *AuthServer) TokenEndpoint(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "only POST allowed", http.StatusMethodNotAllowed)
		return
	}
 
	code := r.FormValue("code")
	codeVerifier := r.FormValue("code_verifier")
 
	// Retrieve the authorization code
	authCode, exists := as.getAuthCode(code)
	if !exists || authCode.Used || authCode.ExpiresAt.Before(time.Now()) {
		http.Error(w, "invalid or expired code", http.StatusBadRequest)
		return
	}
	// The code is bound to the client and redirect_uri it was issued for.
	if r.FormValue("client_id") != authCode.ClientID ||
		r.FormValue("redirect_uri") != authCode.RedirectURI {
		http.Error(w, "invalid_grant", http.StatusBadRequest)
		return
	}
 
	// PKCE verification: verify SHA256(code_verifier) == code_challenge
	hash := sha256.Sum256([]byte(codeVerifier))
	challenge := base64.RawURLEncoding.EncodeToString(hash[:])
 
	if challenge != authCode.CodeChallenge {
		http.Error(w, "PKCE verification failed", http.StatusBadRequest)
		return
	}
 
	// Mark code as used (prevent replay attacks)
	authCode.Used = true
	as.storeAuthCode(code, authCode)
 
	// Issue JWT access token
	claims := jwt.MapClaims{
		"sub":   authCode.ClientID,
		"aud":   as.audience, // the API that will accept this token, not the client
		"iat":   time.Now().Unix(),
		"exp":   time.Now().Add(15 * time.Minute).Unix(),
		"scope": "read write",
	}
	accessToken, err := as.keyMgr.IssueJWT(claims)
	if err != nil {
		http.Error(w, "token issuance failed", http.StatusInternalServerError)
		return
	}
 
	// Issue refresh token
	refreshToken, err := generateRandomCode(64)
	if err != nil {
		http.Error(w, "token issuance failed", http.StatusInternalServerError)
		return
	}
	as.storeRefreshToken(refreshToken, authCode.ClientID)
 
	response := TokenResponse{
		AccessToken:  accessToken,
		RefreshToken: refreshToken,
		ExpiresIn:    900,
		TokenType:    "Bearer",
	}
 
	w.Header().Set("Content-Type", "application/json")
	// RFC 6749 §5.1: responses carrying tokens must never be cached.
	w.Header().Set("Cache-Control", "no-store")
	w.Header().Set("Pragma", "no-cache")
	if err := json.NewEncoder(w).Encode(response); err != nil {
		// Headers and status are already written; the only thing left
		// to do is record that the client got a truncated body.
		slog.Error("token endpoint: encode response", "error", err)
	}
}

// JWKS represents the JSON Web Key Set response.
type JWKS struct {
	Keys []JWK `json:"keys"`
}
 
// JWK represents a single JSON Web Key.
type JWK struct {
	Kty string `json:"kty"` // "RSA"
	Kid string `json:"kid"` // Key ID (matches JWT header "kid")
	Use string `json:"use"` // "sig" for signing
	N   string `json:"n"`   // RSA modulus (base64url)
	E   string `json:"e"`   // RSA exponent (base64url)
}
 
// JWKSEndpoint serves the public keys at /.well-known/jwks.json.
func (as *AuthServer) JWKSEndpoint(w http.ResponseWriter, r *http.Request) {
	pubKey := as.keyMgr.publicKey
	n := base64.RawURLEncoding.EncodeToString(pubKey.N.Bytes())
	e := base64.RawURLEncoding.EncodeToString(big.NewInt(int64(pubKey.E)).Bytes())
 
	jwk := JWK{
		Kty: "RSA",
		Kid: as.keyMgr.keyID,
		Use: "sig",
		N:   n,
		E:   e,
	}
 
	response := JWKS{Keys: []JWK{jwk}}
 
	w.Header().Set("Content-Type", "application/json")
	if err := json.NewEncoder(w).Encode(response); err != nil {
		slog.Error("jwks endpoint: encode response", "error", err)
	}
}

func serverError(w http.ResponseWriter, msg string, err error) {
    slog.Error(msg, "error", err)
    http.Error(w, "server_error", http.StatusInternalServerError)
}
 
// /token endpoint, grant_type=refresh_token branch
func handleRefresh(w http.ResponseWriter, r *http.Request, presented string) {
    rt, err := store.LookupRefreshToken(r.Context(), presented)
    if errors.Is(err, ErrNotFound) {
        http.Error(w, "invalid_grant", http.StatusBadRequest)
        return
    }
    if err != nil {
        // A database outage is not invalid_grant: that answer makes
        // clients discard their tokens and log every user out.
        serverError(w, "refresh: lookup", err)
        return
    }
 
    // CRITICAL: presented token has already been used (rotated)
    if rt.RotatedAt != nil {
        // Theft detection: a revoked token was replayed.
        // Revoke every token in this rotation family (one login);
        // the user's other sessions stay valid.
        if err := store.RevokeFamily(r.Context(), rt.FamilyID); err != nil {
            // The stolen chain may still be live; fail loudly.
            serverError(w, "refresh-token replay: family revocation failed", err)
            return
        }
        slog.Warn("refresh-token replay detected; family revoked",
            "family", rt.FamilyID, "user", rt.UserID, "client", rt.ClientID)
        http.Error(w, "invalid_grant", http.StatusBadRequest)
        return
    }
 
    // Issue a new access token + new refresh token
    newAT, err := signAccessToken(rt.UserID, rt.ClientID, rt.Scopes, 15*time.Minute)
    if err != nil {
        serverError(w, "refresh: sign access token", err)
        return
    }
    newRT, err := generateRandomCode(64)
    if err != nil {
        serverError(w, "refresh: generate refresh token", err)
        return
    }
 
    if err := store.RotateRefreshToken(r.Context(), rt.ID, newRT, 30*24*time.Hour); err != nil {
        serverError(w, "refresh: rotate", err)
        return
    }
 
    w.Header().Set("Content-Type", "application/json")
    w.Header().Set("Cache-Control", "no-store")
    w.Header().Set("Pragma", "no-cache")
    if err := json.NewEncoder(w).Encode(map[string]any{
        "access_token":  newAT,
        "refresh_token": newRT,
        "token_type":    "Bearer",
        "expires_in":    900,
    }); err != nil {
        slog.Error("refresh endpoint: encode response", "error", err)
    }
}

type SigningKeyStore struct {
    mu      sync.RWMutex
    current *KeyPair  // active for new tokens
    retired *KeyPair  // accepted on verify, not used for signing
}
 
// Rotate moves current → retired and brings in a new signing key.
// Retired key is dropped after the next rotation (TTL-based).
func (s *SigningKeyStore) Rotate(next *KeyPair) {
    s.mu.Lock()
    defer s.mu.Unlock()
    s.retired = s.current
    s.current = next
}
 
func (s *SigningKeyStore) JWKS() jose.JSONWebKeySet {
    s.mu.RLock()
    defer s.mu.RUnlock()
    keys := []jose.JSONWebKey{s.current.public()}
    if s.retired != nil {
        keys = append(keys, s.retired.public())
    }
    return jose.JSONWebKeySet{Keys: keys}
}

-- refresh_tokens table — every refresh keeps the family ID across rotations
CREATE TABLE refresh_tokens (
    id          TEXT PRIMARY KEY,
    family_id   TEXT NOT NULL,
    user_id     BIGINT NOT NULL,
    parent_id   TEXT,                 -- chain back to predecessor
    issued_at   TIMESTAMPTZ NOT NULL DEFAULT now(),
    expires_at  TIMESTAMPTZ NOT NULL,
    revoked_at  TIMESTAMPTZ
);
 
CREATE INDEX rt_family_idx ON refresh_tokens (family_id);

func (s *TokenServer) Refresh(ctx context.Context, presented string) (*TokenPair, error) {
    rt, err := s.repo.GetByID(ctx, presented)
    if errors.Is(err, ErrNotFound) {
        return nil, ErrInvalidGrant
    }
    if err != nil {
        return nil, fmt.Errorf("load refresh token: %w", err)
    }
 
    if rt.RevokedAt.Valid {
        // Replay of a revoked token. Nuke the entire family — the legitimate
        // owner will be forced to re-authenticate, which is the correct
        // outcome when a token is suspected stolen.
        s.audit.Log(ctx, AuditEvent{
            Type:     "refresh_token_replay",
            FamilyID: rt.FamilyID,
            UserID:   rt.UserID,
        })
        if err := s.repo.RevokeFamily(ctx, rt.FamilyID); err != nil {
            // Not ErrInvalidGrant: the stolen chain may still be live.
            return nil, fmt.Errorf("revoke family %s after replay: %w", rt.FamilyID, err)
        }
        return nil, ErrInvalidGrant
    }
    // Mark this rotation step revoked, mint successor in same family.
    return s.rotate(ctx, rt)
}

CREATE TABLE auth_audit (
    id              BIGSERIAL PRIMARY KEY,
    occurred_at     TIMESTAMPTZ NOT NULL DEFAULT now(),
    event_type      TEXT NOT NULL,
    actor_user_id   BIGINT,
    client_id       TEXT,
    ip              INET,
    user_agent      TEXT,
    context         JSONB NOT NULL,
    -- Each row's hash includes the previous row's hash; tamper-detection chain.
    row_hash        BYTEA NOT NULL,
    prev_row_hash   BYTEA NOT NULL
);
 
-- Single index covers the two queries auditors actually run.
CREATE INDEX auth_audit_actor_time ON auth_audit (actor_user_id, occurred_at DESC);

POST /oauth/revoke HTTP/1.1
Host: auth.example.com
Authorization: Basic czZCaGRSa3F0MzpnWDFmQmF0M2JW
Content-Type: application/x-www-form-urlencoded
 
token=2YotnFZFEjr1zCsicMWpAA&token_type_hint=refresh_token

// RevokeEndpoint implements RFC 7009.
func (s *AuthServer) RevokeEndpoint(w http.ResponseWriter, r *http.Request) {
    if r.Method != http.MethodPost {
        http.Error(w, "only POST allowed", http.StatusMethodNotAllowed)
        return
    }
 
    clientID, clientSecret, ok := r.BasicAuth()
    if !ok || !s.clients.Authenticate(r.Context(), clientID, clientSecret) {
        w.Header().Set("WWW-Authenticate", `Basic realm="oauth"`)
        http.Error(w, "invalid_client", http.StatusUnauthorized)
        return
    }
 
    presented := r.FormValue("token")
    hint := r.FormValue("token_type_hint")
 
    // Resolve in hint order, fall back to the other type. Spec requires
    // we accept either type regardless of hint.
    rt, found := s.lookupForRevoke(r.Context(), presented, hint)
    if !found {
        w.WriteHeader(http.StatusOK) // invalid token: nothing left to revoke
        return
    }
    if rt.ClientID != clientID {
        // RFC 7009 §2.1: refuse a token that was issued to another client.
        http.Error(w, "unauthorized_client", http.StatusBadRequest)
        return
    }
 
    var err error
    switch rt.Kind {
    case KindRefresh:
        err = s.repo.RevokeFamily(r.Context(), rt.FamilyID)
    case KindAccess:
        err = s.repo.RevokeAccessToken(r.Context(), rt.ID)
    }
    if err != nil {
        // A 200 here would tell the client the token is dead when it isn't.
        slog.Error("revoke: storage write failed", "client", clientID, "error", err)
        w.Header().Set("Retry-After", "5")
        http.Error(w, "revocation temporarily unavailable", http.StatusServiceUnavailable)
        return
    }
    s.audit.Log(r.Context(), AuditEvent{
        Type: "token_revoked", FamilyID: rt.FamilyID,
        ClientID: clientID, UserID: rt.UserID,
    })
    w.WriteHeader(http.StatusOK)
}

CREATE TABLE revoked_tokens (
    jti           TEXT PRIMARY KEY,         -- JWT ID claim
    family_id     TEXT,                     -- null for access tokens
    revoked_at    TIMESTAMPTZ NOT NULL DEFAULT now(),
    expires_at    TIMESTAMPTZ NOT NULL,     -- original token exp
    reason        TEXT NOT NULL             -- 'user_logout' | 'admin_disable' | 'leak_detected'
);
 
CREATE INDEX revoked_tokens_expiry ON revoked_tokens (expires_at);
 
-- Background job: evict rows whose underlying token has already expired.
DELETE FROM revoked_tokens WHERE expires_at < now() - INTERVAL '1 hour';