Align Flight auth behavior

Author: bill-ph

AGENTS.md

@@ -8,4 +8,5 @@
 - When following a plan file, mark tasks upon completion
 - When creating new branch from origin/main, do not track origin/main. 
 - Always run lint before committing.
-- Parallelize using subagents when possible.
+- Parallelize using subagents when possible.
+- Prefer correctness, maintanability, robustness over shortcut implementations

server/flightsqlingress/ingress.go

@@ -211,13 +211,6 @@ func (h *ControlPlaneFlightSQLHandler) sessionFromContext(ctx context.Context) (
 		remoteAddr = p.Addr
 	}
 
-	releaseRateLimit, rejectReason := server.BeginRateLimitedAuthAttempt(h.rateLimiter, remoteAddr)
-	defer releaseRateLimit()
-	if rejectReason != "" {
-		slog.Warn("Flight auth rejected by rate limit policy.", "remote_addr", remoteAddr, "reason", rejectReason)
-		return nil, status.Error(codes.ResourceExhausted, "authentication rate limit exceeded")
-	}
-
 	md, ok := metadata.FromIncomingContext(ctx)
 	if !ok {
 		server.RecordFailedAuthAttempt(h.rateLimiter, remoteAddr)
@@ -228,26 +221,44 @@ func (h *ControlPlaneFlightSQLHandler) sessionFromContext(ctx context.Context) (
 		return nil, status.Error(codes.Unavailable, "session store is not configured")
 	}
 
-	username, err := h.authenticateBasicCredentials(md, remoteAddr)
-	if err != nil {
-		return nil, err
-	}
-
 	if sessionToken := incomingSessionToken(md); sessionToken != "" {
 		s, ok := h.sessions.GetByToken(sessionToken)
 		if !ok {
+			server.RecordFailedAuthAttempt(h.rateLimiter, remoteAddr)
 			return nil, status.Error(codes.Unauthenticated, "session not found")
 		}
 
-		if username != s.username {
-			return nil, status.Error(codes.PermissionDenied, "session token does not match authenticated user")
+		// When Basic auth is included alongside a bearer session token, enforce
+		// principal consistency. Token-only auth is allowed after bootstrap.
+		if hasAuthorizationHeader(md) {
+			username, err := h.authenticateBasicCredentials(md, remoteAddr)
+			if err != nil {
+				return nil, err
+			}
+			if username != s.username {
+				server.RecordFailedAuthAttempt(h.rateLimiter, remoteAddr)
+				return nil, status.Error(codes.PermissionDenied, "session token does not match authenticated user")
+			}
 		}
 
 		setSessionTokenMetadata(ctx, sessionToken)
 		s.touch()
 		return s, nil
 	}
 
+	// Bootstrap requires Basic auth and is subject to auth rate limiting.
+	releaseRateLimit, rejectReason := server.BeginRateLimitedAuthAttempt(h.rateLimiter, remoteAddr)
+	defer releaseRateLimit()
+	if rejectReason != "" {
+		slog.Warn("Flight auth rejected by rate limit policy.", "remote_addr", remoteAddr, "reason", rejectReason)
+		return nil, status.Error(codes.ResourceExhausted, "authentication rate limit exceeded")
+	}
+
+	username, err := h.authenticateBasicCredentials(md, remoteAddr)
+	if err != nil {
+		return nil, err
+	}
+
 	s, err := h.sessions.Create(ctx, username)
 	if err != nil {
 		return nil, status.Errorf(codes.Unavailable, "create bootstrap session: %v", err)
@@ -258,6 +269,10 @@ func (h *ControlPlaneFlightSQLHandler) sessionFromContext(ctx context.Context) (
 	return s, nil
 }
 
+func hasAuthorizationHeader(md metadata.MD) bool {
+	return len(md.Get("authorization")) > 0
+}
+
 func (h *ControlPlaneFlightSQLHandler) authenticateBasicCredentials(md metadata.MD, remoteAddr net.Addr) (string, error) {
 	authHeaders := md.Get("authorization")
 	if len(authHeaders) == 0 {

server/flightsqlingress/ingress_test.go

@@ -172,10 +172,12 @@ func TestFlightAuthSessionKeyDoesNotTrustMetadataClientOverride(t *testing.T) {
 	}
 }
 
-func TestSessionFromContextRejectsServerIssuedSessionTokenWithoutBasicAuth(t *testing.T) {
+func TestSessionFromContextAcceptsServerIssuedSessionTokenWithoutBasicAuth(t *testing.T) {
+	s := newFlightClientSession(1234, "postgres", nil)
+	s.token = "issued-token"
 	store := &flightAuthSessionStore{
 		sessions: map[string]*flightClientSession{
-			"issued-token": newFlightClientSession(1234, "postgres", nil),
+			"issued-token": s,
 		},
 	}
 
@@ -185,8 +187,15 @@ func TestSessionFromContextRejectsServerIssuedSessionTokenWithoutBasicAuth(t *te
 	}
 
 	ctx := metadata.NewIncomingContext(context.Background(), metadata.Pairs("x-duckgres-session", "issued-token"))
-	if _, err := h.sessionFromContext(ctx); err == nil {
-		t.Fatalf("expected token-only auth to be rejected")
+	got, err := h.sessionFromContext(ctx)
+	if err != nil {
+		t.Fatalf("expected token-only auth to succeed, got %v", err)
+	}
+	if got == nil {
+		t.Fatalf("expected non-nil session")
+	}
+	if got != s {
+		t.Fatalf("expected existing token session to be reused")
 	}
 }
 
@@ -253,6 +262,44 @@ func TestSessionFromContextAcceptsServerIssuedSessionTokenWithBasicAuth(t *testi
 	}
 }
 
+func TestSessionFromContextTokenPathDoesNotClearRateLimiterFailures(t *testing.T) {
+	addr := &net.TCPAddr{IP: net.ParseIP("203.0.113.47"), Port: 30004}
+	rateLimiter := server.NewRateLimiter(server.RateLimitConfig{
+		MaxFailedAttempts:   2,
+		FailedAttemptWindow: time.Minute,
+		BanDuration:         time.Hour,
+		MaxConnectionsPerIP: 100,
+	})
+	rateLimiter.RecordFailedAuth(addr)
+
+	s := newFlightClientSession(1234, "postgres", nil)
+	s.token = "issued-token"
+	store := &flightAuthSessionStore{
+		sessions: map[string]*flightClientSession{
+			"issued-token": s,
+		},
+	}
+	h, err := NewControlPlaneFlightSQLHandler(store, map[string]string{"postgres": "postgres"})
+	if err != nil {
+		t.Fatalf("failed to construct handler: %v", err)
+	}
+	h.rateLimiter = rateLimiter
+
+	base := peer.NewContext(context.Background(), &peer.Peer{Addr: addr})
+	ctx := metadata.NewIncomingContext(base, metadata.Pairs("x-duckgres-session", "issued-token"))
+	if _, err := h.sessionFromContext(ctx); err != nil {
+		t.Fatalf("token-only auth failed: %v", err)
+	}
+
+	_, err = h.sessionFromContext(authContextForPeer(addr, "postgres", "wrong"))
+	if status.Code(err) != codes.Unauthenticated {
+		t.Fatalf("expected unauthenticated error for bad password, got %v", err)
+	}
+	if !rateLimiter.IsBanned(addr) {
+		t.Fatalf("expected prior failure + new failure to ban; token-only path should not clear failures")
+	}
+}
+
 func TestSessionFromContextWithoutTokenCreatesDistinctSessions(t *testing.T) {
 	var createCalls atomic.Int32
 	store := &flightAuthSessionStore{

tests/controlplane/flight_ingress_test.go

@@ -28,11 +28,6 @@ func flightAuthContext(username, password string) context.Context {
 	return metadata.NewOutgoingContext(context.Background(), metadata.Pairs("authorization", "Basic "+token))
 }
 
-func basicAuthHeader(username, password string) string {
-	token := base64.StdEncoding.EncodeToString([]byte(username + ":" + password))
-	return "Basic " + token
-}
-
 func newFlightClient(t *testing.T, port int) *flightsql.Client {
 	t.Helper()
 	addr := fmt.Sprintf("127.0.0.1:%d", port)
@@ -150,12 +145,10 @@ func TestFlightIngressIncludeSchemaLowWorkerRegression(t *testing.T) {
 				errCh <- fmt.Errorf("worker %d bootstrap missing x-duckgres-session header", workerID)
 				return
 			}
-			authHeader := basicAuthHeader("testuser", "testpass")
 			token := strings.TrimSpace(sessionTokens[0])
 
 			for i := 0; i < iterationsPerGoroutine; i++ {
 				iterBaseCtx := metadata.NewOutgoingContext(context.Background(), metadata.Pairs(
-					"authorization", authHeader,
 					"x-duckgres-session", token,
 				))
 				iterCtx, iterCancel := context.WithTimeout(iterBaseCtx, 20*time.Second)
@@ -177,7 +170,7 @@ func TestFlightIngressIncludeSchemaLowWorkerRegression(t *testing.T) {
 	}
 }
 
-func TestFlightIngressServerIssuedSessionTokenRequiresBasicAuth(t *testing.T) {
+func TestFlightIngressServerIssuedSessionTokenAllowsTokenOnlyAuth(t *testing.T) {
 	h := startControlPlane(t, cpOpts{
 		flightPort: freePort(t),
 		maxWorkers: 1,
@@ -206,8 +199,8 @@ func TestFlightIngressServerIssuedSessionTokenRequiresBasicAuth(t *testing.T) {
 	ctx2, cancel2 := context.WithTimeout(tokenCtx, 20*time.Second)
 	defer cancel2()
 
-	if _, err := client2.GetTables(ctx2, &flightsql.GetTablesOpts{}); err == nil {
-		t.Fatalf("expected token-only GetTables to fail without basic auth")
+	if _, err := client2.GetTables(ctx2, &flightsql.GetTablesOpts{}); err != nil {
+		t.Fatalf("expected token-only GetTables to succeed, got %v", err)
 	}
 
 	tokenAndAuthCtx := metadata.NewOutgoingContext(context.Background(), metadata.Pairs(