ax/src/serve/oidc.go

package serve

import (
	"context"
	"net/http"
	"strings"
)

type contextKey string

const userContextKey contextKey = "ax_user"

// withSessionAuth wraps a handler with ax session token authentication.
// Auth endpoints (/auth/*) are passed through without a token check.
// All other requests must supply Authorization: Bearer <server_token>.
// If the token is not a valid OIDC session, agentLookup is tried as a fallback.
func withSessionAuth(ah *authHandler, agentLookup func(string) string, next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasPrefix(r.URL.Path, "/auth/") {
			next.ServeHTTP(w, r)
			return
		}
		auth := r.Header.Get("Authorization")
		if !strings.HasPrefix(auth, "Bearer ") {
			writeError(w, http.StatusUnauthorized, "Bearer token required")
			return
		}
		token := strings.TrimPrefix(auth, "Bearer ")
		username := ah.lookupSession(token)
		if username == "" && agentLookup != nil {
			username = agentLookup(token)
		}
		if username == "" {
			writeError(w, http.StatusUnauthorized, "invalid or expired session; run 'ax login'")
			return
		}
		ctx := context.WithValue(r.Context(), userContextKey, username)
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

func userFromContext(r *http.Request) string {
	v, _ := r.Context().Value(userContextKey).(string)
	return v
}