package server

import (
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/mem9-ai/drive9/pkg/innerauth"
	"github.com/mem9-ai/drive9/pkg/logger"
	"github.com/mem9-ai/drive9/pkg/tenant"
	"go.uber.org/zap"
	"go.uber.org/zap/zaptest/observer"
)

func TestInnerControlPlaneRouteMappings(t *testing.T) {
	const tokenID = "4f74bb7f-e1bc-4472-b7c1-50782c669fa9"
	tests := []struct {
		method   string
		path     string
		resource string
	}{
		{http.MethodPost, "/v1/provision", "/drive9/orgs/{org_id}/admin/tenants/modify"},
		{http.MethodPost, "/v1/admin/tenants", "/drive9/orgs/{org_id}/admin/tenants/modify"},
		{http.MethodDelete, "/v1/admin/tenants/tenant-1", "/drive9/orgs/{org_id}/admin/tenants/modify"},
		{http.MethodDelete, "/v1/admin/tenants/tenant-1/", "/drive9/orgs/{org_id}/admin/tenants/modify"},
		{http.MethodGet, "/v1/admin/tenants", "/drive9/orgs/{org_id}/admin/tenants/view"},
		{http.MethodGet, "/v1/admin/tenants/tenant-1", "/drive9/orgs/{org_id}/admin/tenants/view"},
		{http.MethodPost, "/v1/admin/tenant-pool", "/drive9/orgs/{org_id}/admin/tenant-pool/modify"},
		{http.MethodPatch, "/v1/admin/tenant-pool", "/drive9/orgs/{org_id}/admin/tenant-pool/modify"},
		{http.MethodDelete, "/v1/admin/tenant-pool", "/drive9/orgs/{org_id}/admin/tenant-pool/modify"},
		{http.MethodGet, "/v1/admin/tenant-pool", "/drive9/orgs/{org_id}/admin/tenant-pool/view"},
		{http.MethodGet, "/v1/quota", "/drive9/orgs/{org_id}/quota/view"},
		{http.MethodPost, "/v1/quota", "/drive9/orgs/{org_id}/quota/modify"},
		{http.MethodPost, "/v1/admin/tenants/tenant-1/quota", "/drive9/orgs/{org_id}/quota/modify"},
		{http.MethodPost, "/v1/admin/tenants/tenant-1/quota/", "/drive9/orgs/{org_id}/quota/modify"},
		{http.MethodGet, "/v1/tokens", "/drive9/orgs/{org_id}/tokens/view"},
		{http.MethodPost, "/v1/tokens/generate", "/drive9/orgs/{org_id}/tokens/modify"},
		{http.MethodPost, "/v1/tokens/" + tokenID + "/activate", "/drive9/orgs/{org_id}/tokens/modify"},
		{http.MethodPost, "/v1/tokens/" + tokenID + "/deactivate", "/drive9/orgs/{org_id}/tokens/modify"},
		{http.MethodDelete, "/v1/tokens/" + tokenID, "/drive9/orgs/{org_id}/tokens/modify"},
		{http.MethodPost, "/v1/tokens/revoke", "/drive9/orgs/{org_id}/tokens/modify"},
		{http.MethodGet, "/v1/admin/object-backends", "/drive9/orgs/{org_id}/object-backend/view"},
		{http.MethodPost, "/v1/admin/object-backends", "/drive9/orgs/{org_id}/object-backend/modify"},
		{http.MethodGet, "/v1/admin/object-backends/obb_1", "/drive9/orgs/{org_id}/object-backend/view"},
		{http.MethodPatch, "/v1/admin/object-backends/obb_1", "/drive9/orgs/{org_id}/object-backend/modify"},
		{http.MethodDelete, "/v1/admin/object-backends/obb_1", "/drive9/orgs/{org_id}/object-backend/modify"},
		{http.MethodGet, "/v1/admin/tenants/tenant-1/object-namespace", "/drive9/orgs/{org_id}/object-namespace/view"},
		{http.MethodPut, "/v1/admin/tenants/tenant-1/object-namespace", "/drive9/orgs/{org_id}/object-namespace/modify"},
		{http.MethodDelete, "/v1/admin/tenants/tenant-1/object-namespace", "/drive9/orgs/{org_id}/object-namespace/modify"},
	}

	for _, tt := range tests {
		t.Run(tt.method+" "+tt.path, func(t *testing.T) {
			got, ok := matchInnerControlPlaneRoute(tt.method, tt.path)
			if !ok {
				t.Fatal("route did not match")
			}
			if got != tt.resource {
				t.Fatalf("resource = %q, want %q", got, tt.resource)
			}
		})
	}
}

func TestInnerAuthURIsUseHighestProjectScopeForAPIKey(t *testing.T) {
	subject := innerauth.ControlPlaneSubject{
		Kind:           innerauth.CredentialAPIKey,
		OrganizationID: 60205,
		ResourceInfos: []innerauth.APIKeyResourceInfo{
			{ScopeType: "ORG", ResourceID: 60205, Role: "org:member"},
			{ScopeType: "PROJECT", ResourceID: 103547, Role: "project:readonly"},
			{ScopeType: "PROJECT", ResourceID: 103548, Role: "project:owner"},
			{ScopeType: "PROJECT", ResourceID: 103549, Role: "project:dev"},
		},
	}

	requestURI, resourceURI := innerAuthURIs(subject, "/drive9/orgs/{org_id}/quota/view")
	if requestURI != "/drive9/orgs/60205/projects/103548/quota/view" {
		t.Fatalf("request URI = %q, want highest project owner URI", requestURI)
	}
	if resourceURI != "/drive9/orgs/{org_id}/projects/{project_id}/quota/view" {
		t.Fatalf("resource URI = %q, want project-scoped URI", resourceURI)
	}
}

func TestInnerAuthURIsKeepOrgScopeForOrgOnlyAPIKey(t *testing.T) {
	subject := innerauth.ControlPlaneSubject{
		Kind:           innerauth.CredentialAPIKey,
		OrganizationID: 60205,
		ResourceInfos:  []innerauth.APIKeyResourceInfo{{ScopeType: "ORG", ResourceID: 60205, Role: "org:owner"}},
	}

	requestURI, resourceURI := innerAuthURIs(subject, "/drive9/orgs/{org_id}/quota/view")
	if requestURI != "/drive9/orgs/60205/quota/view" || resourceURI != "/drive9/orgs/{org_id}/quota/view" {
		t.Fatalf("URIs = (%q, %q), want org-scoped URIs", requestURI, resourceURI)
	}
}

func TestInnerAuthURIsKeepJWTOrgScopeEvenWithProjectMetadata(t *testing.T) {
	subject := innerauth.ControlPlaneSubject{
		Kind:           innerauth.CredentialJWT,
		OrganizationID: 60205,
		ResourceInfos:  []innerauth.APIKeyResourceInfo{{ScopeType: "PROJECT", ResourceID: 103548, Role: "project:owner"}},
	}

	requestURI, resourceURI := innerAuthURIs(subject, "/drive9/orgs/{org_id}/quota/view")
	if requestURI != "/drive9/orgs/60205/quota/view" || resourceURI != "/drive9/orgs/{org_id}/quota/view" {
		t.Fatalf("URIs = (%q, %q), want JWT org-scoped URIs", requestURI, resourceURI)
	}
}

func TestInnerAuthURIsKeepOrgScopeWhenResourceTemplateHasNoDrive9Prefix(t *testing.T) {
	subject := innerauth.ControlPlaneSubject{
		Kind:           innerauth.CredentialAPIKey,
		OrganizationID: 60205,
		ResourceInfos: []innerauth.APIKeyResourceInfo{
			{ScopeType: "PROJECT", ResourceID: 103548, Role: "project:owner"},
		},
	}

	requestURI, resourceURI := innerAuthURIs(subject, "/serverless/orgs/{org_id}/clusters/create")
	if requestURI != "/serverless/orgs/60205/clusters/create" || resourceURI != "/serverless/orgs/{org_id}/clusters/create" {
		t.Fatalf("URIs = (%q, %q), want unchanged org-scoped template", requestURI, resourceURI)
	}
}

func TestInnerControlPlaneRouteDoesNotMatchDataPlaneOrNearMisses(t *testing.T) {
	tests := []struct {
		method string
		path   string
	}{
		{http.MethodPost, "/v1/tokens"},
		{http.MethodPost, "/v1/tokens/refresh"},
		{http.MethodPost, "/v1/fork"},
		{http.MethodPost, "/v1/tenant"},
		{http.MethodGet, "/v1/provision"},
		{http.MethodDelete, "/v1/admin/tenants"},
		{http.MethodGet, "/v1/admin/tenants/tenant-1/quota"},
		{http.MethodPost, "/v1/admin/tenants/tenant-1/quota/extra"},
		{http.MethodPost, "/v1/tokens/not-a-uuid/activate"},
		{http.MethodPost, "/v1/tokens/not-a-uuid/deactivate"},
		{http.MethodDelete, "/v1/tokens/not-a-uuid"},
		{http.MethodPost, "/v1/tokens/token-1/activate/extra"},
		{http.MethodDelete, "/v1/tokens/token-1/extra"},
	}

	for _, tt := range tests {
		t.Run(tt.method+" "+tt.path, func(t *testing.T) {
			if resource, ok := matchInnerControlPlaneRoute(tt.method, tt.path); ok {
				t.Fatalf("unexpected match with resource %q", resource)
			}
		})
	}
}

func TestInnerControlPlaneRouteMatchesEveryAcceptedTokenMutationPath(t *testing.T) {
	const tokenID = "4f74bb7f-e1bc-4472-b7c1-50782c669fa9"
	tests := []struct {
		method string
		path   string
		id     string
	}{
		{method: http.MethodDelete, path: "/v1/tokens/" + tokenID, id: tokenIDFromPath("/v1/tokens/" + tokenID)},
		{method: http.MethodPost, path: "/v1/tokens/" + tokenID + "/activate", id: tokenIDFromStatusPath("/v1/tokens/"+tokenID+"/activate", "/activate")},
		{method: http.MethodPost, path: "/v1/tokens/" + tokenID + "/deactivate", id: tokenIDFromStatusPath("/v1/tokens/"+tokenID+"/deactivate", "/deactivate")},
	}

	for _, tt := range tests {
		t.Run(tt.method+" "+tt.path, func(t *testing.T) {
			if tt.id == "" {
				t.Fatal("handler rejected test path")
			}
			resource, ok := matchInnerControlPlaneRoute(tt.method, tt.path)
			if !ok || resource != "/drive9/orgs/{org_id}/tokens/modify" {
				t.Fatalf("route = (%q, %t), want tokens/modify", resource, ok)
			}
		})
	}
}

func TestInnerControlPlaneAuthJWTAuthenticatesAuthorizesAndInjectsContext(t *testing.T) {
	account := &innerAuthAccountFake{userResult: innerauth.UserAuthnResult{UserID: 123}}
	defaults := tenant.CredentialProvisionRequest{PublicKey: "server-public", PrivateKey: "server-private"}
	dispatcher := innerauth.NewDispatcher(account, innerAuthDefaultsFake{credentials: defaults, ok: true}, time.Hour)
	wantBody := `{"max_file_count":10}`
	handlerCalled := false
	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		handlerCalled = true
		subject := InnerAuthSubjectFromContext(r.Context())
		if subject == nil || subject.Kind != innerauth.CredentialJWT || subject.OrganizationID != 60205 || subject.UserID != 123 {
			t.Errorf("subject = %#v", subject)
		}
		credentials := InnerAuthCredentialsFromContext(r.Context())
		if credentials == nil || *credentials != defaults {
			t.Errorf("execution credentials = %#v, want defaults", credentials)
		}
		body, err := io.ReadAll(r.Body)
		if err != nil {
			t.Errorf("read preserved body: %v", err)
		} else if string(body) != wantBody {
			t.Errorf("body = %q, want %q", body, wantBody)
		}
		w.WriteHeader(http.StatusNoContent)
	})

	req := httptest.NewRequest(http.MethodPost, "/v1/provision?orgId=60205", strings.NewReader(wantBody))
	req.Header.Set(innerauth.TiDBCloudAuthorizationHeader, "Bearer "+testInnerJWT())
	recorder := httptest.NewRecorder()
	innerControlPlaneAuthMiddleware(dispatcher, next).ServeHTTP(recorder, req)

	if recorder.Code != http.StatusNoContent || !handlerCalled {
		t.Fatalf("status = %d, handler called = %t", recorder.Code, handlerCalled)
	}
	if got := strings.Join(account.calls, ","); got != "user_authn,user_authz" {
		t.Fatalf("calls = %q, want authn then authz", got)
	}
	if account.requestURI != "/drive9/orgs/60205/admin/tenants/modify" || account.resourceURI != "/drive9/orgs/{org_id}/admin/tenants/modify" {
		t.Fatalf("authz URIs = (%q, %q)", account.requestURI, account.resourceURI)
	}
}

func TestInnerControlPlaneAuthAPIKeyFromBodyDoesNotRequireOrgID(t *testing.T) {
	account := &innerAuthAccountFake{apiKeyResult: innerauth.APIKeyAuthnResult{PublicKey: "caller-public", OrganizationID: 60205}}
	defaults := tenant.CredentialProvisionRequest{PublicKey: "shared-public", PrivateKey: "shared-private"}
	dispatcher := innerauth.NewDispatcher(account, innerAuthDefaultsFake{credentials: defaults, ok: true}, time.Hour)
	wantBody := `{"public_key":"caller-public","private_key":"caller-private","pool_size":3}`
	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		subject := InnerAuthSubjectFromContext(r.Context())
		if subject == nil || subject.Kind != innerauth.CredentialAPIKey || subject.OrganizationID != 60205 || subject.PublicKey != "caller-public" {
			t.Errorf("subject = %#v", subject)
		}
		credentials := InnerAuthCredentialsFromContext(r.Context())
		if credentials == nil || *credentials != defaults {
			t.Errorf("execution credentials = %#v", credentials)
		}
		body, _ := io.ReadAll(r.Body)
		if string(body) != wantBody {
			t.Errorf("body = %q, want %q", body, wantBody)
		}
		w.WriteHeader(http.StatusNoContent)
	})

	req := httptest.NewRequest(http.MethodPost, "/v1/admin/tenant-pool", strings.NewReader(wantBody))
	recorder := httptest.NewRecorder()
	innerControlPlaneAuthMiddleware(dispatcher, next).ServeHTTP(recorder, req)

	if recorder.Code != http.StatusNoContent {
		t.Fatalf("status = %d, body = %s", recorder.Code, recorder.Body.String())
	}
	if got := strings.Join(account.calls, ","); got != "api_key_authn,api_key_authz" {
		t.Fatalf("calls = %q, want authn then authz", got)
	}
	if account.requestURI != "/drive9/orgs/60205/admin/tenant-pool/modify" || account.resourceURI != "/drive9/orgs/{org_id}/admin/tenant-pool/modify" {
		t.Fatalf("authz URIs = (%q, %q)", account.requestURI, account.resourceURI)
	}
}

func TestInnerControlPlaneAuthAPIKeyUsesHighestProjectScope(t *testing.T) {
	account := &innerAuthAccountFake{apiKeyResult: innerauth.APIKeyAuthnResult{
		PublicKey:      "caller-public",
		OrganizationID: 60205,
		ResourceInfos: []innerauth.APIKeyResourceInfo{
			{ScopeType: "ORG", ResourceID: 60205, Role: "org:member"},
			{ScopeType: "PROJECT", ResourceID: 103547, Role: "project:readonly"},
			{ScopeType: "PROJECT", ResourceID: 103548, Role: "project:owner"},
			{ScopeType: "PROJECT", ResourceID: 103549, Role: "project:dev"},
		},
	}}
	dispatcher := innerauth.NewDispatcher(account, innerAuthDefaultsFake{credentials: tenant.CredentialProvisionRequest{PublicKey: "shared-public", PrivateKey: "shared-private"}, ok: true}, time.Hour)
	next := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusNoContent)
	})
	req := httptest.NewRequest(http.MethodGet, "/v1/quota", nil)
	req.Header.Set(quotaPublicKeyHeader, "caller-public")
	req.Header.Set(quotaPrivateKeyHeader, "caller-private")
	recorder := httptest.NewRecorder()
	innerControlPlaneAuthMiddleware(dispatcher, next).ServeHTTP(recorder, req)

	if recorder.Code != http.StatusNoContent {
		t.Fatalf("status = %d, body = %s", recorder.Code, recorder.Body.String())
	}
	if account.requestURI != "/drive9/orgs/60205/projects/103548/quota/view" || account.resourceURI != "/drive9/orgs/{org_id}/projects/{project_id}/quota/view" {
		t.Fatalf("authz URIs = (%q, %q), want highest project owner URI", account.requestURI, account.resourceURI)
	}
}

func TestInnerControlPlaneAuthAPIKeyHeadersTakePrecedenceOverBody(t *testing.T) {
	account := &innerAuthAccountFake{apiKeyResult: innerauth.APIKeyAuthnResult{PublicKey: "header-public", OrganizationID: 60205}}
	dispatcher := innerauth.NewDispatcher(account, innerAuthDefaultsFake{credentials: tenant.CredentialProvisionRequest{PublicKey: "shared-public", PrivateKey: "shared-private"}, ok: true}, time.Hour)
	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		credentials := InnerAuthCredentialsFromContext(r.Context())
		if credentials == nil || credentials.PublicKey != "shared-public" || credentials.PrivateKey != "shared-private" {
			t.Errorf("execution credentials = %#v, want shared defaults", credentials)
		}
		w.WriteHeader(http.StatusNoContent)
	})
	req := httptest.NewRequest(http.MethodPost, "/v1/quota", strings.NewReader(`{"public_key":"body-public","private_key":"body-private"}`))
	req.Header.Set(quotaPublicKeyHeader, "header-public")
	req.Header.Set(quotaPrivateKeyHeader, "header-private")
	recorder := httptest.NewRecorder()
	innerControlPlaneAuthMiddleware(dispatcher, next).ServeHTTP(recorder, req)
	if recorder.Code != http.StatusNoContent {
		t.Fatalf("status = %d, body = %s", recorder.Code, recorder.Body.String())
	}
}

func TestInnerControlPlaneAuthCompleteAPIKeyHeadersDoNotReadBody(t *testing.T) {
	account := &innerAuthAccountFake{apiKeyResult: innerauth.APIKeyAuthnResult{PublicKey: "header-public", OrganizationID: 60205}}
	dispatcher := innerauth.NewDispatcher(account, innerAuthDefaultsFake{credentials: tenant.CredentialProvisionRequest{PublicKey: "shared-public", PrivateKey: "shared-private"}, ok: true}, time.Hour)
	body := &readTrackingBody{}
	req := httptest.NewRequest(http.MethodPost, "/v1/quota", nil)
	req.Body = body
	req.Header.Set(quotaPublicKeyHeader, "header-public")
	req.Header.Set(quotaPrivateKeyHeader, "header-private")
	recorder := httptest.NewRecorder()

	innerControlPlaneAuthMiddleware(dispatcher, http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusNoContent)
	})).ServeHTTP(recorder, req)

	if recorder.Code != http.StatusNoContent {
		t.Fatalf("status = %d, body = %s", recorder.Code, recorder.Body.String())
	}
	if body.read {
		t.Fatal("middleware read request body despite complete API key headers")
	}
}

func TestInnerControlPlaneAuthRejectsMalformedBodyOnlyCredentials(t *testing.T) {
	dispatcher := innerauth.NewDispatcher(&innerAuthAccountFake{}, innerAuthDefaultsFake{}, time.Hour)
	nextCalled := false
	next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
		nextCalled = true
	})
	req := httptest.NewRequest(http.MethodPost, "/v1/quota", strings.NewReader(`{"public_key":"public"`))
	recorder := httptest.NewRecorder()

	innerControlPlaneAuthMiddleware(dispatcher, next).ServeHTTP(recorder, req)

	if recorder.Code != http.StatusBadRequest {
		t.Fatalf("status = %d, want %d, body = %s", recorder.Code, http.StatusBadRequest, recorder.Body.String())
	}
	if nextCalled {
		t.Fatal("business handler called after malformed credential body")
	}
}

func TestInnerControlPlaneAuthEnabledPassesThroughDataPlaneRoute(t *testing.T) {
	dispatcher := innerauth.NewDispatcher(&innerAuthAccountFake{}, innerAuthDefaultsFake{}, time.Hour)
	called := false
	next := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		called = true
		w.WriteHeader(http.StatusNoContent)
	})
	req := httptest.NewRequest(http.MethodPost, "/v1/tokens", nil)
	req.Header.Set(innerauth.TiDBCloudAuthorizationHeader, "Bearer malformed")
	recorder := httptest.NewRecorder()
	innerControlPlaneAuthMiddleware(dispatcher, next).ServeHTTP(recorder, req)
	if recorder.Code != http.StatusNoContent || !called {
		t.Fatalf("status = %d, handler called = %t", recorder.Code, called)
	}
}

func TestInnerControlPlaneAuthPreservesDrive9BearerOnMatchedRoutes(t *testing.T) {
	dispatcher := innerauth.NewDispatcher(&innerAuthAccountFake{}, innerAuthDefaultsFake{}, time.Hour)
	tests := []struct {
		method string
		path   string
	}{
		{http.MethodGet, "/v1/tokens"},
		{http.MethodGet, "/v1/quota?tenant_id=tenant-1"},
	}
	for _, tt := range tests {
		t.Run(tt.method+" "+tt.path, func(t *testing.T) {
			called := false
			next := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
				called = true
				w.WriteHeader(http.StatusNoContent)
			})
			req := httptest.NewRequest(tt.method, tt.path, nil)
			req.Header.Set("Authorization", "Bearer drive9-api-key")
			recorder := httptest.NewRecorder()

			innerControlPlaneAuthMiddleware(dispatcher, next).ServeHTTP(recorder, req)

			if recorder.Code != http.StatusNoContent || !called {
				t.Fatalf("status = %d, handler called = %t, body = %s", recorder.Code, called, recorder.Body.String())
			}
		})
	}
}

func TestInnerControlPlaneAuthDrive9BearerDoesNotReadOrTruncateBody(t *testing.T) {
	dispatcher := innerauth.NewDispatcher(&innerAuthAccountFake{}, innerAuthDefaultsFake{}, time.Hour)
	wantBody := strings.Repeat("x", int(maxCredentialProvisionBodyBytes)+17)
	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotBody, err := io.ReadAll(r.Body)
		if err != nil {
			t.Errorf("read body: %v", err)
		} else if string(gotBody) != wantBody {
			t.Errorf("body length = %d, want %d", len(gotBody), len(wantBody))
		}
		w.WriteHeader(http.StatusNoContent)
	})
	req := httptest.NewRequest(http.MethodPost, "/v1/tokens/generate", strings.NewReader(wantBody))
	req.Header.Set("Authorization", "Bearer drive9-api-key")
	recorder := httptest.NewRecorder()

	innerControlPlaneAuthMiddleware(dispatcher, next).ServeHTTP(recorder, req)

	if recorder.Code != http.StatusNoContent {
		t.Fatalf("status = %d, body = %s", recorder.Code, recorder.Body.String())
	}
}

func TestInnerControlPlaneAuthDoesNotBypassRBACForUnsupportedDrive9BearerRoutes(t *testing.T) {
	dispatcher := innerauth.NewDispatcher(&innerAuthAccountFake{}, innerAuthDefaultsFake{}, time.Hour)
	routes := []struct {
		method string
		path   string
	}{
		{http.MethodPost, "/v1/provision"},
		{http.MethodPost, "/v1/admin/tenants"},
		{http.MethodPost, "/v1/admin/tenant-pool"},
		{http.MethodPost, "/v1/quota"},
	}
	credentialHeaders := []struct {
		name string
		set  func(*http.Request)
	}{
		{"Authorization", func(req *http.Request) { req.Header.Set("Authorization", "Bearer drive9-api-key") }},
		{"X-API-Key", func(req *http.Request) { req.Header.Set("X-API-Key", "drive9-api-key") }},
	}
	for _, route := range routes {
		for _, credentialHeader := range credentialHeaders {
			t.Run(route.method+" "+route.path+" "+credentialHeader.name, func(t *testing.T) {
				called := false
				next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
					called = true
				})
				req := httptest.NewRequest(route.method, route.path, nil)
				credentialHeader.set(req)
				recorder := httptest.NewRecorder()

				innerControlPlaneAuthMiddleware(dispatcher, next).ServeHTTP(recorder, req)

				if recorder.Code != http.StatusUnauthorized {
					t.Fatalf("status = %d, want %d, body = %s", recorder.Code, http.StatusUnauthorized, recorder.Body.String())
				}
				if called {
					t.Fatal("business handler called without inner authentication")
				}
			})
		}
	}
}

func TestInnerControlPlaneAuthRejectsUnsupportedDrive9BearerCombinedWithBodyAPIKey(t *testing.T) {
	dispatcher := innerauth.NewDispatcher(&innerAuthAccountFake{}, innerAuthDefaultsFake{}, time.Hour)
	called := false
	next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
		called = true
	})
	req := httptest.NewRequest(http.MethodPost, "/v1/provision", strings.NewReader(`{"public_key":"caller-public","private_key":"caller-private"}`))
	req.Header.Set("Authorization", "Bearer drive9-api-key")
	recorder := httptest.NewRecorder()

	innerControlPlaneAuthMiddleware(dispatcher, next).ServeHTTP(recorder, req)

	if recorder.Code != http.StatusBadRequest {
		t.Fatalf("status = %d, want %d, body = %s", recorder.Code, http.StatusBadRequest, recorder.Body.String())
	}
	if called {
		t.Fatal("business handler called with ambiguous credentials")
	}
}

func TestInnerControlPlaneAuthRejectsDrive9BearerCombinedWithInnerCredentials(t *testing.T) {
	tests := []struct {
		name    string
		prepare func(*http.Request)
	}{
		{
			name: "JWT",
			prepare: func(req *http.Request) {
				req.Header.Set(innerauth.TiDBCloudAuthorizationHeader, "Bearer "+testInnerJWT())
			},
		},
		{
			name: "TiDB Cloud API key",
			prepare: func(req *http.Request) {
				req.Header.Set(quotaPublicKeyHeader, "public")
				req.Header.Set(quotaPrivateKeyHeader, "private")
			},
		},
	}
	dispatcher := innerauth.NewDispatcher(&innerAuthAccountFake{}, innerAuthDefaultsFake{}, time.Hour)
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			called := false
			next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) { called = true })
			req := httptest.NewRequest(http.MethodGet, "/v1/tokens?orgId=60205", nil)
			req.Header.Set("Authorization", "Bearer drive9-api-key")
			tt.prepare(req)
			recorder := httptest.NewRecorder()

			innerControlPlaneAuthMiddleware(dispatcher, next).ServeHTTP(recorder, req)

			if recorder.Code != http.StatusBadRequest {
				t.Fatalf("status = %d, want %d, body = %s", recorder.Code, http.StatusBadRequest, recorder.Body.String())
			}
			if !strings.Contains(recorder.Body.String(), "ambiguous credentials") {
				t.Fatalf("body = %s, want ambiguous credentials", recorder.Body.String())
			}
			if called {
				t.Fatal("business handler called with ambiguous credential families")
			}
		})
	}
}

func TestInnerControlPlaneAuthErrorStatusMapping(t *testing.T) {
	tests := []struct {
		name       string
		request    func() *http.Request
		accountErr error
		wantStatus int
	}{
		{
			name: "malformed JWT",
			request: func() *http.Request {
				req := httptest.NewRequest(http.MethodGet, "/v1/quota?orgId=60205", nil)
				req.Header.Set(innerauth.TiDBCloudAuthorizationHeader, "Bearer malformed")
				return req
			},
			wantStatus: http.StatusBadRequest,
		},
		{
			name: "JWT missing organization",
			request: func() *http.Request {
				req := httptest.NewRequest(http.MethodGet, "/v1/quota", nil)
				req.Header.Set(innerauth.TiDBCloudAuthorizationHeader, "Bearer "+testInnerJWT())
				return req
			},
			wantStatus: http.StatusBadRequest,
		},
		{
			name: "ambiguous JWT and API key",
			request: func() *http.Request {
				req := httptest.NewRequest(http.MethodGet, "/v1/quota?orgId=60205", nil)
				req.Header.Set(innerauth.TiDBCloudAuthorizationHeader, "Bearer "+testInnerJWT())
				req.Header.Set(quotaPublicKeyHeader, "public")
				req.Header.Set(quotaPrivateKeyHeader, "private")
				return req
			},
			wantStatus: http.StatusBadRequest,
		},
		{
			name: "ambiguous takes precedence over missing organization",
			request: func() *http.Request {
				req := httptest.NewRequest(http.MethodGet, "/v1/quota", nil)
				req.Header.Set(innerauth.TiDBCloudAuthorizationHeader, "Bearer "+testInnerJWT())
				req.Header.Set(quotaPublicKeyHeader, "public")
				req.Header.Set(quotaPrivateKeyHeader, "private")
				return req
			},
			wantStatus: http.StatusBadRequest,
		},
		{
			name: "unauthenticated",
			request: func() *http.Request {
				return apiKeyInnerRequest()
			},
			accountErr: innerauth.ErrUnauthenticated,
			wantStatus: http.StatusUnauthorized,
		},
		{
			name: "unauthorized",
			request: func() *http.Request {
				return apiKeyInnerRequest()
			},
			accountErr: innerauth.ErrForbidden,
			wantStatus: http.StatusForbidden,
		},
		{
			name: "upstream protocol failure",
			request: func() *http.Request {
				return apiKeyInnerRequest()
			},
			accountErr: innerauth.ErrInvalidResponse,
			wantStatus: http.StatusBadGateway,
		},
		{
			name: "server credentials unavailable",
			request: func() *http.Request {
				req := httptest.NewRequest(http.MethodGet, "/v1/quota?orgId=60205", nil)
				req.Header.Set(innerauth.TiDBCloudAuthorizationHeader, "Bearer "+testInnerJWT())
				return req
			},
			wantStatus: http.StatusServiceUnavailable,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			account := &innerAuthAccountFake{
				apiKeyResult: innerauth.APIKeyAuthnResult{PublicKey: "public", OrganizationID: 60205},
				userResult:   innerauth.UserAuthnResult{UserID: 123},
			}
			if errors.Is(tt.accountErr, innerauth.ErrForbidden) {
				account.authzErr = tt.accountErr
			} else {
				account.authnErr = tt.accountErr
			}
			defaults := innerAuthDefaultsFake{}
			if tt.name != "server credentials unavailable" {
				defaults = innerAuthDefaultsFake{credentials: tenant.CredentialProvisionRequest{PublicKey: "server-public", PrivateKey: "server-private"}, ok: true}
			}
			dispatcher := innerauth.NewDispatcher(account, defaults, time.Hour)
			handlerCalled := false
			next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) { handlerCalled = true })
			recorder := httptest.NewRecorder()
			innerControlPlaneAuthMiddleware(dispatcher, next).ServeHTTP(recorder, tt.request())
			if recorder.Code != tt.wantStatus {
				t.Fatalf("status = %d, want %d, body = %s", recorder.Code, tt.wantStatus, recorder.Body.String())
			}
			if strings.Contains(tt.name, "ambiguous") && !strings.Contains(recorder.Body.String(), "ambiguous credentials") {
				t.Fatalf("body = %s, want ambiguous credentials", recorder.Body.String())
			}
			if tt.name == "server credentials unavailable" && !strings.Contains(recorder.Body.String(), "server credentials unavailable") {
				t.Fatalf("body = %s, want server credentials unavailable", recorder.Body.String())
			}
			if handlerCalled {
				t.Fatal("business handler called after authentication or authorization failure")
			}
		})
	}
}

func TestInnerControlPlaneAuthDisabledPreservesLegacyHandler(t *testing.T) {
	for _, dispatcher := range []*innerauth.Dispatcher{nil, innerauth.NewDispatcher(nil, innerAuthDefaultsFake{}, time.Hour)} {
		called := false
		next := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
			called = true
			w.WriteHeader(http.StatusNoContent)
		})
		recorder := httptest.NewRecorder()
		innerControlPlaneAuthMiddleware(dispatcher, next).ServeHTTP(recorder, httptest.NewRequest(http.MethodPost, "/v1/provision", nil))
		if recorder.Code != http.StatusNoContent || !called {
			t.Fatalf("dispatcher = %#v, status = %d, handler called = %t", dispatcher, recorder.Code, called)
		}
	}
}

func TestInnerControlPlaneAuthJWTDenialUsesRequestedOrganizationAndSafeLogFields(t *testing.T) {
	core, recorded := observer.New(zap.InfoLevel)
	restoreLogger := logger.L()
	logger.Set(zap.New(core))
	t.Cleanup(func() { logger.Set(restoreLogger) })

	account := &innerAuthAccountFake{userResult: innerauth.UserAuthnResult{UserID: 123}, authzErr: innerauth.ErrForbidden}
	defaults := innerAuthDefaultsFake{credentials: tenant.CredentialProvisionRequest{PublicKey: "server-public-secret", PrivateKey: "server-private-secret"}, ok: true}
	dispatcher := innerauth.NewDispatcher(account, defaults, time.Hour)
	handlerCalled := false
	req := httptest.NewRequest(http.MethodGet, "/v1/quota?orgId=60206", nil)
	req.Header.Set(innerauth.TiDBCloudAuthorizationHeader, "Bearer "+testInnerJWT())
	recorder := httptest.NewRecorder()

	innerControlPlaneAuthMiddleware(dispatcher, http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
		handlerCalled = true
	})).ServeHTTP(recorder, req)

	if recorder.Code != http.StatusForbidden || handlerCalled {
		t.Fatalf("status = %d, handler called = %t", recorder.Code, handlerCalled)
	}
	if account.requestURI != "/drive9/orgs/60206/quota/view" || account.resourceURI != "/drive9/orgs/{org_id}/quota/view" {
		t.Fatalf("authz URIs = (%q, %q)", account.requestURI, account.resourceURI)
	}
	entries := recorded.FilterMessage("inner_control_plane_authorization_denied").AllUntimed()
	if len(entries) != 1 {
		t.Fatalf("denial logs = %d, want 1", len(entries))
	}
	fields := entries[0].ContextMap()
	for key, want := range map[string]any{
		"credential_kind": "jwt",
		"resource_uri":    "/drive9/orgs/{org_id}/quota/view",
		"organization_id": uint64(60206),
	} {
		if got := fields[key]; got != want {
			t.Errorf("log field %s = %#v, want %#v", key, got, want)
		}
	}
	encoded := entries[0].Message
	for key, value := range fields {
		encoded += key + "=" + fmt.Sprint(value)
	}
	for _, secret := range []string{testInnerJWT(), "server-public-secret", "server-private-secret"} {
		if strings.Contains(encoded, secret) {
			t.Fatalf("denial log leaked credential material")
		}
	}
}

func TestInnerControlPlaneAuthJWTAuthorizationUsesCallerSuppliedOrganizationInRequestURI(t *testing.T) {
	account := &innerAuthAccountFake{
		userResult: innerauth.UserAuthnResult{UserID: 123},
		userAuthz: func(requestURI, _ string) error {
			if strings.HasPrefix(requestURI, "/drive9/orgs/60205/") {
				return nil
			}
			return innerauth.ErrForbidden
		},
	}
	defaults := innerAuthDefaultsFake{credentials: tenant.CredentialProvisionRequest{PublicKey: "server-public", PrivateKey: "server-private"}, ok: true}
	dispatcher := innerauth.NewDispatcher(account, defaults, time.Hour)
	handlerCalls := 0
	handler := innerControlPlaneAuthMiddleware(dispatcher, http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		handlerCalls++
		w.WriteHeader(http.StatusNoContent)
	}))

	for _, tt := range []struct {
		organizationID string
		wantStatus     int
	}{
		{organizationID: "60205", wantStatus: http.StatusNoContent},
		{organizationID: "60206", wantStatus: http.StatusForbidden},
	} {
		req := httptest.NewRequest(http.MethodGet, "/v1/quota?orgId="+tt.organizationID, nil)
		req.Header.Set(innerauth.TiDBCloudAuthorizationHeader, "Bearer "+testInnerJWT())
		recorder := httptest.NewRecorder()
		handler.ServeHTTP(recorder, req)
		if recorder.Code != tt.wantStatus {
			t.Fatalf("org %s status = %d, want %d, body = %s", tt.organizationID, recorder.Code, tt.wantStatus, recorder.Body.String())
		}
	}
	if handlerCalls != 1 {
		t.Fatalf("business handler calls = %d, want only allowed organization", handlerCalls)
	}
}

type readTrackingBody struct {
	read bool
}

func (b *readTrackingBody) Read([]byte) (int, error) {
	b.read = true
	return 0, errors.New("unexpected body read")
}

func (*readTrackingBody) Close() error { return nil }

func TestNewWithConfigWiresInnerAuthToControlPlaneRoutes(t *testing.T) {
	dispatcher := innerauth.NewDispatcher(&innerAuthAccountFake{}, innerAuthDefaultsFake{}, time.Hour)
	server := NewWithConfig(Config{
		DefaultTenantProvider: tenant.ProviderTiDBCloudNativeShared,
		InnerAuth:             dispatcher,
	})
	t.Cleanup(server.Close)
	tests := []struct {
		method string
		path   string
	}{
		{http.MethodPost, "/v1/provision?orgId=60205"},
		{http.MethodGet, "/v1/quota?orgId=60205"},
		{http.MethodGet, "/v1/admin/tenant-pool?orgId=60205"},
		{http.MethodGet, "/v1/admin/tenants?orgId=60205"},
		{http.MethodGet, "/v1/admin/tenants/tenant-1?orgId=60205"},
		{http.MethodDelete, "/v1/admin/tenants/tenant-1/?orgId=60205"},
		{http.MethodPost, "/v1/admin/tenants/tenant-1/quota/?orgId=60205"},
		{http.MethodGet, "/v1/tokens?orgId=60205"},
	}
	for _, tt := range tests {
		t.Run(tt.method+" "+tt.path, func(t *testing.T) {
			req := httptest.NewRequest(tt.method, tt.path, nil)
			req.Header.Set(innerauth.TiDBCloudAuthorizationHeader, "Bearer malformed")
			recorder := httptest.NewRecorder()
			server.ServeHTTP(recorder, req)
			if recorder.Code != http.StatusBadRequest {
				t.Fatalf("status = %d, want inner-auth 400, body = %s", recorder.Code, recorder.Body.String())
			}
		})
	}
}

func TestTokensCredentialDispatcherAcceptsCompletedInnerAuth(t *testing.T) {
	called := false
	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		called = true
		scope := ControlPlaneScopeFromContext(r.Context())
		if scope == nil || scope.OrganizationID != "60205" || scope.Identity.OrganizationID != "60205" {
			t.Errorf("control-plane scope = %#v", scope)
		}
		w.WriteHeader(http.StatusNoContent)
	})
	subject := &innerauth.ControlPlaneSubject{OrganizationID: 60205, PublicKey: "public", Kind: innerauth.CredentialAPIKey}
	ctx := withInnerAuthSubject(context.Background(), subject)
	req := httptest.NewRequest(http.MethodGet, "/v1/tokens", nil).WithContext(ctx)
	recorder := httptest.NewRecorder()
	new(Server).tokensCredentialDispatcher(Config{}, next).ServeHTTP(recorder, req)
	if recorder.Code != http.StatusNoContent || !called {
		t.Fatalf("status = %d, handler called = %t, body = %s", recorder.Code, called, recorder.Body.String())
	}
}

func apiKeyInnerRequest() *http.Request {
	req := httptest.NewRequest(http.MethodGet, "/v1/quota", nil)
	req.Header.Set(quotaPublicKeyHeader, "public")
	req.Header.Set(quotaPrivateKeyHeader, "private")
	return req
}

func testInnerJWT() string {
	return base64.RawURLEncoding.EncodeToString([]byte(`{"alg":"none"}`)) + "." +
		base64.RawURLEncoding.EncodeToString([]byte(`{"sub":"user"}`)) + ".c2lnbmF0dXJl"
}

type innerAuthAccountFake struct {
	userResult   innerauth.UserAuthnResult
	apiKeyResult innerauth.APIKeyAuthnResult
	authnErr     error
	authzErr     error
	userAuthz    func(requestURI, resourceURI string) error
	calls        []string
	requestURI   string
	resourceURI  string
}

func (f *innerAuthAccountFake) UserAuthn(context.Context, string) (innerauth.UserAuthnResult, error) {
	f.calls = append(f.calls, "user_authn")
	return f.userResult, f.authnErr
}

func (f *innerAuthAccountFake) UserAuthz(_ context.Context, _ uint64, requestURI, resourceURI string) error {
	f.calls = append(f.calls, "user_authz")
	f.requestURI, f.resourceURI = requestURI, resourceURI
	if f.userAuthz != nil {
		return f.userAuthz(requestURI, resourceURI)
	}
	return f.authzErr
}

func (f *innerAuthAccountFake) APIKeyAuthn(context.Context, string) (innerauth.APIKeyAuthnResult, error) {
	f.calls = append(f.calls, "api_key_authn")
	return f.apiKeyResult, f.authnErr
}

func (f *innerAuthAccountFake) APIKeyAuthz(_ context.Context, _ string, requestURI, resourceURI string) error {
	f.calls = append(f.calls, "api_key_authz")
	f.requestURI, f.resourceURI = requestURI, resourceURI
	return f.authzErr
}

type innerAuthDefaultsFake struct {
	credentials tenant.CredentialProvisionRequest
	ok          bool
}

func (f innerAuthDefaultsFake) DefaultCredentials() (tenant.CredentialProvisionRequest, bool) {
	return f.credentials, f.ok
}