feat: add agent user nodes with access token authentication

src/serve/oidc.go

@@ -13,7 +13,8 @@ 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>.
-func withSessionAuth(ah *authHandler, next http.Handler) http.Handler {
+// 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)
@@ -26,6 +27,9 @@ func withSessionAuth(ah *authHandler, next http.Handler) http.Handler {
 		}
 		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

src/serve/server.go

@@ -14,8 +14,8 @@ import (
 // When oidcCfg is non-nil, every request must carry a valid Bearer token;
 // the authenticated username is derived from the token claim configured in
 // OIDCConfig.UserClaim. Without OIDC, the X-Ax-User header is used instead.
-func New(newSvc func(user string) (service.NodeService, error), oidcCfg *store.OIDCConfig) (http.Handler, error) {
-	s := &server{newSvc: newSvc}
+func New(newSvc func(user string) (service.NodeService, error), oidcCfg *store.OIDCConfig, agentLookup func(string) string) (http.Handler, error) {
+	s := &server{newSvc: newSvc, agentLookup: agentLookup}
 	mux := http.NewServeMux()
 	mux.HandleFunc("GET /nodes", s.listNodes)
 	mux.HandleFunc("POST /nodes", s.addNode)
@@ -35,17 +35,23 @@ func New(newSvc func(user string) (service.NodeService, error), oidcCfg *store.O
 		mux.HandleFunc("POST /auth/device/start", ah.deviceStart)
 		mux.HandleFunc("GET /auth/callback", ah.callback)
 		mux.HandleFunc("GET /auth/poll", ah.poll)
-		return withRateLimit(rl, withSessionAuth(ah, mux)), nil
+		return withRateLimit(rl, withSessionAuth(ah, agentLookup, mux)), nil
 	}
 	return withRateLimit(rl, mux), nil
 }
 
 type server struct {
-	newSvc func(user string) (service.NodeService, error)
+	newSvc      func(user string) (service.NodeService, error)
+	agentLookup func(string) string
 }
 
 func (s *server) svc(w http.ResponseWriter, r *http.Request) (service.NodeService, bool) {
 	user := userFromContext(r)
+	if user == "" && s.agentLookup != nil {
+		if auth := r.Header.Get("Authorization"); strings.HasPrefix(auth, "Bearer ") {
+			user = s.agentLookup(strings.TrimPrefix(auth, "Bearer "))
+		}
+	}
 	if user == "" {
 		user = r.Header.Get("X-Ax-User")
 	}