Building an OAuth2 Authorization Server from Scratch

package authserver
 
import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
)
 
// 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
}

// AuthorizationRequest holds incoming authorization request parameters.
type AuthorizationRequest struct {
	ClientID            string
	RedirectURI         string
	CodeChallenge       string
	CodeChallengeMethod string // "S256" (SHA256) or "plain"
	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) {
	clientID := r.URL.Query().Get("client_id")
	if clientID == "" {
		http.Error(w, "missing client_id", http.StatusBadRequest)
		return
	}
 
	redirectURI := r.URL.Query().Get("redirect_uri")
	if redirectURI == "" {
		http.Error(w, "missing redirect_uri", http.StatusBadRequest)
		return
	}
 
	// PKCE: code_challenge is required
	codeChallenge := r.URL.Query().Get("code_challenge")
	if codeChallenge == "" {
		http.Error(w, "PKCE code_challenge 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: r.URL.Query().Get("code_challenge_method"),
		ExpiresAt:           time.Now().Add(10 * time.Minute),
		Used:                false,
	}
	as.storeAuthCode(code, authCode)
 
	// Redirect back to the client with the code and state
	state := r.URL.Query().Get("state")
	redirectURL := fmt.Sprintf("%s?code=%s&state=%s", redirectURI, code, state)
	http.Redirect(w, r, redirectURL, 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
	}
 
	// 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":   authCode.RedirectURI,
		"iat":   time.Now().Unix(),
		"exp":   time.Now().Add(1 * time.Hour).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:    3600,
		TokenType:    "Bearer",
	}
 
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(response)
}

// 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")
	json.NewEncoder(w).Encode(response)
}

// /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 err != nil {
        http.Error(w, "invalid_grant", http.StatusBadRequest)
        return
    }
 
    // CRITICAL: presented token has already been used (rotated)
    if rt.RotatedAt != nil {
        // Theft detection: a revoked token was replayed.
        // Revoke EVERY token issued to this user from this client.
        _ = store.RevokeFamily(r.Context(), rt.UserID, rt.ClientID)
        slog.Warn("refresh-token replay detected; family revoked",
            "user", rt.UserID, "client", rt.ClientID)
        http.Error(w, "invalid_grant", http.StatusBadRequest)
        return
    }
 
    // Issue a new access token + new refresh token
    newAT := signAccessToken(rt.UserID, rt.ClientID, rt.Scopes, 15*time.Minute)
    newRT := generateRefreshToken()
 
    if err := store.RotateRefreshToken(r.Context(), rt.ID, newRT, 30*24*time.Hour); err != nil {
        http.Error(w, "server_error", http.StatusInternalServerError)
        return
    }
 
    json.NewEncoder(w).Encode(map[string]any{
        "access_token":  newAT,
        "refresh_token": newRT,
        "token_type":    "Bearer",
        "expires_in":    900,
    })
}

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 err != nil { return nil, ErrInvalidGrant }
 
    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.repo.RevokeFamily(ctx, rt.FamilyID)
        s.audit.Log(ctx, AuditEvent{
            Type:     "refresh_token_replay",
            FamilyID: rt.FamilyID,
            UserID:   rt.UserID,
        })
        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. Always returns 200 to avoid
// leaking whether the token exists.
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) {
        // 401 is the only error the spec permits — auth failure.
        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 {
        if rt.Kind == KindRefresh && rt.ClientID == clientID {
            _ = s.repo.RevokeFamily(r.Context(), rt.FamilyID)
        } else if rt.Kind == KindAccess && rt.ClientID == clientID {
            _ = s.repo.RevokeAccessToken(r.Context(), rt.ID)
        }
        s.audit.Log(r.Context(), AuditEvent{
            Type: "token_revoked", FamilyID: rt.FamilyID,
            ClientID: clientID, UserID: rt.UserID,
        })
    }
 
    // Always 200 — even if the token didn't exist or belonged to another client.
    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';