amcs/internal/app/oauth.go

package app

import (
	"crypto/sha256"
	"crypto/subtle"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"html"
	"log/slog"
	"net/http"
	"net/url"
	"strings"
	"time"

	"git.warky.dev/wdevs/amcs/internal/auth"
)

// --- JSON types ---

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

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

type oauthServerMetadata struct {
	Issuer                            string   `json:"issuer"`
	AuthorizationEndpoint             string   `json:"authorization_endpoint"`
	TokenEndpoint                     string   `json:"token_endpoint"`
	RegistrationEndpoint              string   `json:"registration_endpoint"`
	ScopesSupported                   []string `json:"scopes_supported"`
	ResponseTypesSupported            []string `json:"response_types_supported"`
	GrantTypesSupported               []string `json:"grant_types_supported"`
	TokenEndpointAuthMethodsSupported []string `json:"token_endpoint_auth_methods_supported"`
	CodeChallengeMethodsSupported     []string `json:"code_challenge_methods_supported"`
}

type registerRequest struct {
	ClientName              string   `json:"client_name"`
	RedirectURIs            []string `json:"redirect_uris"`
	GrantTypes              []string `json:"grant_types"`
	ResponseTypes           []string `json:"response_types"`
	TokenEndpointAuthMethod string   `json:"token_endpoint_auth_method"`
}

type registerResponse struct {
	ClientID                string   `json:"client_id"`
	ClientName              string   `json:"client_name"`
	RedirectURIs            []string `json:"redirect_uris"`
	GrantTypes              []string `json:"grant_types"`
	ResponseTypes           []string `json:"response_types"`
	TokenEndpointAuthMethod string   `json:"token_endpoint_auth_method"`
}

// --- Handlers ---

// oauthMetadataHandler serves GET /.well-known/oauth-authorization-server
// per RFC 8414 for OAuth 2.0 server metadata discovery.
func oauthMetadataHandler() http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		base := serverBaseURL(r)
		meta := oauthServerMetadata{
			Issuer:                            base,
			AuthorizationEndpoint:             base + "/authorize",
			TokenEndpoint:                     base + "/oauth/token",
			RegistrationEndpoint:              base + "/oauth/register",
			ScopesSupported:                   []string{"mcp"},
			ResponseTypesSupported:            []string{"code"},
			GrantTypesSupported:               []string{"authorization_code", "client_credentials"},
			TokenEndpointAuthMethodsSupported: []string{"client_secret_basic", "none"},
			CodeChallengeMethodsSupported:     []string{"S256"},
		}
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(meta)
	}
}

// oauthRegisterHandler serves POST /oauth/register per RFC 7591
// (OAuth 2.0 Dynamic Client Registration).
func oauthRegisterHandler(dynClients *auth.DynamicClientStore, log *slog.Logger) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			w.Header().Set("Allow", http.MethodPost)
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
			return
		}

		var req registerRequest
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			http.Error(w, "invalid request body", http.StatusBadRequest)
			return
		}
		if len(req.RedirectURIs) == 0 {
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(http.StatusBadRequest)
			_, _ = w.Write([]byte(`{"error":"invalid_client_metadata","error_description":"redirect_uris is required"}`))
			return
		}

		client, err := dynClients.Register(req.ClientName, req.RedirectURIs)
		if err != nil {
			log.Error("oauth register: failed", slog.String("error", err.Error()))
			http.Error(w, "internal server error", http.StatusInternalServerError)
			return
		}

		log.Info("oauth register: new client",
			slog.String("client_id", client.ClientID),
			slog.String("client_name", client.ClientName),
		)

		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusCreated)
		_ = json.NewEncoder(w).Encode(registerResponse{
			ClientID:                client.ClientID,
			ClientName:              client.ClientName,
			RedirectURIs:            client.RedirectURIs,
			GrantTypes:              []string{"authorization_code"},
			ResponseTypes:           []string{"code"},
			TokenEndpointAuthMethod: "none",
		})
	}
}

// oauthAuthorizeHandler serves GET and POST /oauth/authorize.
// GET shows an approval page; POST processes the user's approve/deny action.
func oauthAuthorizeHandler(dynClients *auth.DynamicClientStore, authCodes *auth.AuthCodeStore, log *slog.Logger) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		switch r.Method {
		case http.MethodGet:
			handleAuthorizeGET(w, r, dynClients)
		case http.MethodPost:
			handleAuthorizePOST(w, r, dynClients, authCodes, log)
		default:
			w.Header().Set("Allow", "GET, POST")
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		}
	}
}

func handleAuthorizeGET(w http.ResponseWriter, r *http.Request, dynClients *auth.DynamicClientStore) {
	q := r.URL.Query()
	clientID := q.Get("client_id")
	redirectURI := q.Get("redirect_uri")
	responseType := q.Get("response_type")
	state := q.Get("state")
	codeChallenge := q.Get("code_challenge")
	codeChallengeMethod := q.Get("code_challenge_method")
	scope := q.Get("scope")

	// Validate client and redirect_uri before any redirect (prevents open redirect).
	client, ok := dynClients.Lookup(clientID)
	if !ok {
		http.Error(w, "unknown client_id", http.StatusBadRequest)
		return
	}
	if !client.HasRedirectURI(redirectURI) {
		http.Error(w, "redirect_uri not registered for this client", http.StatusBadRequest)
		return
	}

	// Errors from here can safely redirect back to the client.
	if responseType != "code" {
		oauthRedirectError(w, r, redirectURI, "unsupported_response_type", state)
		return
	}
	if codeChallenge == "" || codeChallengeMethod != "S256" {
		oauthRedirectError(w, r, redirectURI, "invalid_request", state)
		return
	}

	serveAuthorizePage(w, client.ClientName, clientID, redirectURI, state, codeChallenge, codeChallengeMethod, scope)
}

func handleAuthorizePOST(w http.ResponseWriter, r *http.Request, dynClients *auth.DynamicClientStore, authCodes *auth.AuthCodeStore, log *slog.Logger) {
	if err := r.ParseForm(); err != nil {
		http.Error(w, "invalid form", http.StatusBadRequest)
		return
	}

	clientID := r.FormValue("client_id")
	redirectURI := r.FormValue("redirect_uri")
	state := r.FormValue("state")
	codeChallenge := r.FormValue("code_challenge")
	codeChallengeMethod := r.FormValue("code_challenge_method")
	scope := r.FormValue("scope")
	action := r.FormValue("action")

	client, ok := dynClients.Lookup(clientID)
	if !ok {
		http.Error(w, "unknown client_id", http.StatusBadRequest)
		return
	}
	if !client.HasRedirectURI(redirectURI) {
		http.Error(w, "redirect_uri not registered for this client", http.StatusBadRequest)
		return
	}

	if action == "deny" {
		oauthRedirectError(w, r, redirectURI, "access_denied", state)
		return
	}

	code, err := authCodes.Issue(auth.AuthCode{
		ClientID:            clientID,
		RedirectURI:         redirectURI,
		Scope:               scope,
		CodeChallenge:       codeChallenge,
		CodeChallengeMethod: codeChallengeMethod,
		KeyID:               clientID,
	})
	if err != nil {
		log.Error("oauth authorize: failed to issue code", slog.String("error", err.Error()))
		oauthRedirectError(w, r, redirectURI, "server_error", state)
		return
	}

	target := redirectURI + "?code=" + url.QueryEscape(code)
	if state != "" {
		target += "&state=" + url.QueryEscape(state)
	}
	http.Redirect(w, r, target, http.StatusFound)
}

// oauthTokenHandler serves POST /oauth/token.
// Supports grant_type=client_credentials and grant_type=authorization_code.
func oauthTokenHandler(oauthRegistry *auth.OAuthRegistry, tokenStore *auth.TokenStore, authCodes *auth.AuthCodeStore, log *slog.Logger) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			w.Header().Set("Allow", http.MethodPost)
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
			return
		}
		if err := r.ParseForm(); err != nil {
			writeTokenError(w, "invalid_request", http.StatusBadRequest)
			return
		}

		switch r.FormValue("grant_type") {
		case "client_credentials":
			handleClientCredentials(w, r, oauthRegistry, tokenStore, log)
		case "authorization_code":
			handleAuthorizationCode(w, r, authCodes, tokenStore, log)
		default:
			writeTokenError(w, "unsupported_grant_type", http.StatusBadRequest)
		}
	}
}

func handleClientCredentials(w http.ResponseWriter, r *http.Request, oauthRegistry *auth.OAuthRegistry, tokenStore *auth.TokenStore, log *slog.Logger) {
	clientID, clientSecret := extractOAuthBasicOrBody(r)
	if clientID == "" || clientSecret == "" {
		w.Header().Set("WWW-Authenticate", `Basic realm="oauth"`)
		writeTokenError(w, "invalid_client", http.StatusUnauthorized)
		return
	}
	keyID, ok := oauthRegistry.Lookup(clientID, clientSecret)
	if !ok {
		log.Warn("oauth token: invalid client credentials", slog.String("remote_addr", r.RemoteAddr))
		w.Header().Set("WWW-Authenticate", `Basic realm="oauth"`)
		writeTokenError(w, "invalid_client", http.StatusUnauthorized)
		return
	}
	issueToken(w, keyID, tokenStore, log)
}

func handleAuthorizationCode(w http.ResponseWriter, r *http.Request, authCodes *auth.AuthCodeStore, tokenStore *auth.TokenStore, log *slog.Logger) {
	code := r.FormValue("code")
	redirectURI := r.FormValue("redirect_uri")
	clientID := r.FormValue("client_id")
	codeVerifier := r.FormValue("code_verifier")

	if code == "" || redirectURI == "" || codeVerifier == "" {
		writeTokenError(w, "invalid_request", http.StatusBadRequest)
		return
	}

	entry, ok := authCodes.Consume(code)
	if !ok {
		writeTokenError(w, "invalid_grant", http.StatusBadRequest)
		return
	}
	if entry.ClientID != clientID || entry.RedirectURI != redirectURI {
		writeTokenError(w, "invalid_grant", http.StatusBadRequest)
		return
	}
	if !verifyPKCE(codeVerifier, entry.CodeChallenge, entry.CodeChallengeMethod) {
		log.Warn("oauth token: PKCE verification failed", slog.String("remote_addr", r.RemoteAddr))
		writeTokenError(w, "invalid_grant", http.StatusBadRequest)
		return
	}

	issueToken(w, entry.KeyID, tokenStore, log)
}

func issueToken(w http.ResponseWriter, keyID string, tokenStore *auth.TokenStore, log *slog.Logger) {
	token, ttl, err := tokenStore.Issue(keyID)
	if err != nil {
		log.Error("oauth token: failed to issue token", slog.String("error", err.Error()))
		http.Error(w, "internal server error", http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Cache-Control", "no-store")
	w.Header().Set("Pragma", "no-cache")
	_ = json.NewEncoder(w).Encode(tokenResponse{
		AccessToken: token,
		TokenType:   "bearer",
		ExpiresIn:   int(ttl / time.Second),
	})
}

// --- Helpers ---

func serveAuthorizePage(w http.ResponseWriter, clientName, clientID, redirectURI, state, codeChallenge, codeChallengeMethod, scope string) {
	e := html.EscapeString
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	_, _ = fmt.Fprintf(w, `<!DOCTYPE html>
<html>
<head>
<meta charset=utf-8>
<title>Authorize — AMCS</title>
<style>
body{font-family:system-ui,sans-serif;max-width:480px;margin:80px auto;padding:0 1rem}
button{padding:.5rem 1.2rem;margin-right:.5rem;cursor:pointer;font-size:1rem}
</style>
</head>
<body>
<h2>Authorize Access</h2>
<p><strong>%s</strong> is requesting access to this AMCS server.</p>
<form method=POST action=/oauth/authorize>
<input type=hidden name=client_id value="%s">
<input type=hidden name=redirect_uri value="%s">
<input type=hidden name=state value="%s">
<input type=hidden name=code_challenge value="%s">
<input type=hidden name=code_challenge_method value="%s">
<input type=hidden name=scope value="%s">
<button type=submit name=action value=approve>Approve</button>
<button type=submit name=action value=deny>Deny</button>
</form>
</body>
</html>`,
		e(clientName), e(clientID), e(redirectURI), e(state),
		e(codeChallenge), e(codeChallengeMethod), e(scope))
}

func oauthRedirectError(w http.ResponseWriter, r *http.Request, redirectURI, errCode, state string) {
	target := redirectURI + "?error=" + url.QueryEscape(errCode)
	if state != "" {
		target += "&state=" + url.QueryEscape(state)
	}
	http.Redirect(w, r, target, http.StatusFound)
}

func verifyPKCE(verifier, challenge, method string) bool {
	if method != "S256" {
		return false
	}
	h := sha256.Sum256([]byte(verifier))
	got := base64.RawURLEncoding.EncodeToString(h[:])
	return subtle.ConstantTimeCompare([]byte(got), []byte(challenge)) == 1
}

func serverBaseURL(r *http.Request) string {
	scheme := "https"
	if proto := r.Header.Get("X-Forwarded-Proto"); proto != "" {
		scheme = strings.ToLower(proto)
	} else if r.TLS == nil {
		scheme = "http"
	}
	return scheme + "://" + r.Host
}

func extractOAuthBasicOrBody(r *http.Request) (string, string) {
	if id, secret, ok := r.BasicAuth(); ok {
		return id, secret
	}
	return strings.TrimSpace(r.FormValue("client_id")), strings.TrimSpace(r.FormValue("client_secret"))
}

func writeTokenError(w http.ResponseWriter, errCode string, status int) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(tokenErrorResponse{Error: errCode})
}