proxy-pool/internal/controller/distribution/handler_test.go

package distribution

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"slices"
	"strings"
	"testing"
	"time"

	controllerExtraction "proxy-pool/internal/controller/extraction"
	"proxy-pool/internal/domain/authorization"
	"proxy-pool/internal/domain/clientpolicy"
	domainExtraction "proxy-pool/internal/domain/extraction"
	"proxy-pool/internal/platform/httpapi"
	"proxy-pool/internal/platform/httpsecurity"
)

func TestNewHandlerValidatesDependenciesAndBodyLimit(t *testing.T) {
	t.Parallel()

	identity := fakeIdentityResolver{identity: Identity{ClientID: "client-1", SourceIP: "198.51.100.8"}}
	readiness := fakeReadinessChecker{}
	extractor := &fakeExtractor{}

	tests := []struct {
		name    string
		config  Config
		deps    Dependencies
		wantErr string
	}{
		{
			name:    "missing extractor",
			config:  Config{BodyLimitBytes: 1024},
			deps:    Dependencies{Identity: identity, Readiness: readiness},
			wantErr: "extractor",
		},
		{
			name:    "missing identity",
			config:  Config{BodyLimitBytes: 1024},
			deps:    Dependencies{Extractor: extractor, Readiness: readiness},
			wantErr: "identity",
		},
		{
			name:    "missing readiness",
			config:  Config{BodyLimitBytes: 1024},
			deps:    Dependencies{Extractor: extractor, Identity: identity},
			wantErr: "readiness",
		},
		{
			name:    "non-positive body limit",
			config:  Config{BodyLimitBytes: 0},
			deps:    Dependencies{Extractor: extractor, Identity: identity, Readiness: readiness},
			wantErr: "body limit",
		},
	}

	for _, test := range tests {
		test := test
		t.Run(test.name, func(t *testing.T) {
			t.Parallel()
			handler, err := NewHandler(test.config, test.deps)
			if err == nil || !strings.Contains(strings.ToLower(err.Error()), test.wantErr) {
				t.Fatalf("NewHandler() = (%v, %v), want error containing %q", handler, err, test.wantErr)
			}
		})
	}
}

func TestHandlerExtractSuccessMapsOpenAPIDTOAndReturnsRequestID(t *testing.T) {
	t.Parallel()

	extractor := &fakeExtractor{response: controllerExtraction.Response{
		RequestID: "unexpected-service-request-id",
		Requested: 2,
		Returned:  1,
		Proxies: []controllerExtraction.ExtractedProxy{{
			ID:                  "px-1",
			Protocol:            "http",
			Host:                "192.0.2.10",
			Port:                8080,
			Username:            "user",
			Password:            "pass",
			URL:                 "http://user:pass@192.0.2.10:8080",
			Region:              "shanghai",
			Carrier:             "ct",
			Upstream:            "provider-a",
			ExpiresAt:           time.Date(2026, 7, 28, 12, 5, 0, 0, time.UTC),
			RemainingTTLSeconds: 300,
			ExtractedAt:         time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC),
		}},
	}}
	handler := mustNewHandler(t, Config{BodyLimitBytes: 4096}, Dependencies{
		Extractor: extractor,
		Identity:  fakeIdentityResolver{identity: Identity{ClientID: "tenant-a", SourceIP: "198.51.100.8"}},
		Readiness: fakeReadinessChecker{},
	})
	request := httptest.NewRequest(http.MethodPost, "/api/v1/proxies/extract", strings.NewReader(`{
		"count": 2,
		"fulfillment": "partial",
		"filters": {
			"protocols": ["http"],
			"regions": ["shanghai"],
			"carriers": ["ct"],
			"allowedUpstreams": ["provider-a"]
		}
	}`))
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set(httpapi.HeaderRequestID, "req-caller")
	request.Header.Set("Idempotency-Key", "idem-12345678")

	response := httptest.NewRecorder()
	handler.ServeHTTP(response, request)

	if response.Code != http.StatusOK {
		t.Fatalf("status = %d, want 200", response.Code)
	}
	if got := response.Header().Get(httpapi.HeaderRequestID); got != "req-caller" {
		t.Fatalf("X-Request-ID = %q, want req-caller", got)
	}
	if got := response.Header().Get("Content-Type"); got != httpapi.JSONContentType {
		t.Fatalf("Content-Type = %q, want %q", got, httpapi.JSONContentType)
	}
	if extractor.calls != 1 {
		t.Fatalf("extractor calls = %d, want 1", extractor.calls)
	}
	if extractor.request.ClientID != "tenant-a" || extractor.request.SourceIP != "198.51.100.8" {
		t.Fatalf("identity request = %+v", extractor.request)
	}
	if extractor.request.IdempotencyKey != "idem-12345678" {
		t.Fatalf("idempotency key = %q", extractor.request.IdempotencyKey)
	}
	if got := extractor.request.Filters.Upstreams; len(got) != 1 || got[0] != "provider-a" {
		t.Fatalf("allowedUpstreams mapping = %v", got)
	}

	var payload struct {
		RequestID string `json:"requestId"`
		Requested int    `json:"requested"`
		Returned  int    `json:"returned"`
		Proxies   []struct {
			ID       string `json:"id"`
			Password string `json:"password"`
			URL      string `json:"url"`
			Upstream string `json:"upstream"`
		} `json:"proxies"`
	}
	if err := json.Unmarshal(response.Body.Bytes(), &payload); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if payload.RequestID != "req-caller" || payload.Requested != 2 || payload.Returned != 1 {
		t.Fatalf("payload = %+v", payload)
	}
	if len(payload.Proxies) != 1 || payload.Proxies[0].Password != "pass" || payload.Proxies[0].URL == "" {
		t.Fatalf("proxies payload = %+v", payload.Proxies)
	}
}

func TestHandlerRejectsIdentityWithoutExtractPermission(t *testing.T) {
	t.Parallel()
	extractor := &fakeExtractor{}
	handler := mustNewHandler(t, Config{BodyLimitBytes: 1024}, Dependencies{
		Extractor: extractor,
		Identity: fakeIdentityResolver{identity: Identity{
			ClientID: "tenant-read-only", SourceIP: "198.51.100.8", Permissions: []string{authorization.AdminRead},
		}},
		Readiness: fakeReadinessChecker{},
	})
	request := httptest.NewRequest(http.MethodPost, "/api/v1/proxies/extract", strings.NewReader(`{"count":1}`))
	request.Header.Set("Content-Type", httpapi.JSONContentType)
	request.Header.Set(headerIdempotencyKey, "idem-12345678")
	recorder := httptest.NewRecorder()

	handler.ServeHTTP(recorder, request)

	if recorder.Code != http.StatusForbidden || extractor.calls != 0 || !strings.Contains(recorder.Body.String(), `"code":"FORBIDDEN"`) {
		t.Fatalf("response = status %d calls=%d body=%s", recorder.Code, extractor.calls, recorder.Body.String())
	}
}

func TestHandlerEnforcesCredentialExtractionPolicy(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name              string
		body              string
		wantStatus        int
		wantUpstreams     []string
		wantRegions       []string
		wantExtractorCall int
	}{
		{
			name:              "omitted filters are bound to credential policy",
			body:              "{\"count\":2}",
			wantStatus:        http.StatusOK,
			wantUpstreams:     []string{"provider-a"},
			wantRegions:       []string{"shanghai"},
			wantExtractorCall: 1,
		},
		{
			name:              "count exceeds credential maximum",
			body:              "{\"count\":3}",
			wantStatus:        http.StatusForbidden,
			wantExtractorCall: 0,
		},
		{
			name:              "upstream exceeds credential boundary",
			body:              "{\"count\":1,\"filters\":{\"allowedUpstreams\":[\"provider-b\"]}}",
			wantStatus:        http.StatusForbidden,
			wantExtractorCall: 0,
		},
		{
			name:              "region exceeds credential boundary",
			body:              "{\"count\":1,\"filters\":{\"regions\":[\"beijing\"]}}",
			wantStatus:        http.StatusForbidden,
			wantExtractorCall: 0,
		},
	}

	for _, test := range tests {
		test := test
		t.Run(test.name, func(t *testing.T) {
			t.Parallel()
			extractor := &fakeExtractor{}
			handler := mustNewHandler(t, Config{BodyLimitBytes: 1024}, Dependencies{
				Extractor: extractor,
				Identity: fakeIdentityResolver{identity: Identity{
					ClientID:    "tenant-restricted",
					Permissions: []string{authorization.DistributionExtract},
					ClientPolicy: clientpolicy.Policy{
						MaxExtractCount:  2,
						AllowedUpstreams: []string{"provider-a"},
						AllowedRegions:   []string{"shanghai"},
					},
				}},
				Readiness: fakeReadinessChecker{},
			})
			request := httptest.NewRequest(http.MethodPost, pathExtract, strings.NewReader(test.body))
			request.Header.Set("Content-Type", httpapi.JSONContentType)
			response := httptest.NewRecorder()

			handler.ServeHTTP(response, request)

			if response.Code != test.wantStatus || extractor.calls != test.wantExtractorCall {
				t.Fatalf("response = status %d calls=%d, want status %d calls=%d",
					response.Code, extractor.calls, test.wantStatus, test.wantExtractorCall)
			}
			if response.Code == http.StatusForbidden && !strings.Contains(response.Body.String(), "\"code\":\"FORBIDDEN\"") {
				t.Fatalf("forbidden response = %s", response.Body.String())
			}
			if got := extractor.request.Filters.Upstreams; !slices.Equal(got, test.wantUpstreams) {
				t.Fatalf("upstreams = %v, want %v", got, test.wantUpstreams)
			}
			if got := extractor.request.Filters.Regions; !slices.Equal(got, test.wantRegions) {
				t.Fatalf("regions = %v, want %v", got, test.wantRegions)
			}
		})
	}
}

func TestHandlerRejectsEmptyResolvedIdentity(t *testing.T) {
	t.Parallel()
	extractor := &fakeExtractor{}
	handler := mustNewHandler(t, Config{BodyLimitBytes: 1024}, Dependencies{
		Extractor: extractor,
		Identity:  fakeIdentityResolver{},
		Readiness: fakeReadinessChecker{},
	})
	request := httptest.NewRequest(http.MethodPost, "/api/v1/proxies/extract", strings.NewReader(`{"count":1}`))
	request.Header.Set("Content-Type", httpapi.JSONContentType)
	response := httptest.NewRecorder()

	handler.ServeHTTP(response, request)

	if response.Code != http.StatusBadRequest {
		t.Fatalf("status = %d, want %d; body=%s", response.Code, http.StatusBadRequest, response.Body.String())
	}
	if extractor.calls != 0 {
		t.Fatalf("extractor calls = %d, want 0", extractor.calls)
	}
}

func TestHandlerRejectsDuplicateIdempotencyHeader(t *testing.T) {
	t.Parallel()
	extractor := &fakeExtractor{}
	handler := mustNewHandler(t, Config{BodyLimitBytes: 1024}, Dependencies{
		Extractor: extractor,
		Identity:  fakeIdentityResolver{identity: Identity{ClientID: "client-1"}},
		Readiness: fakeReadinessChecker{},
	})
	request := httptest.NewRequest(http.MethodPost, "/api/v1/proxies/extract", strings.NewReader(`{"count":1}`))
	request.Header.Set("Content-Type", httpapi.JSONContentType)
	request.Header.Add("Idempotency-Key", "idempotency-one")
	request.Header.Add("Idempotency-Key", "idempotency-two")
	response := httptest.NewRecorder()

	handler.ServeHTTP(response, request)

	if response.Code != http.StatusBadRequest {
		t.Fatalf("status = %d, want %d; body=%s", response.Code, http.StatusBadRequest, response.Body.String())
	}
	if extractor.calls != 0 {
		t.Fatalf("extractor calls = %d, want 0", extractor.calls)
	}
}

func TestHandlerMapsSecurityFailureBeforeParsingBody(t *testing.T) {
	t.Parallel()
	extractor := &fakeExtractor{}
	handler := mustNewHandler(t, Config{BodyLimitBytes: 1024}, Dependencies{
		Extractor: extractor,
		Identity: fakeIdentityResolver{err: &httpsecurity.HTTPError{
			StatusCode: http.StatusUnauthorized,
			Code:       "UNAUTHORIZED",
			Header:     http.Header{"WWW-Authenticate": []string{`Bearer realm="proxy-pool"`}},
			Cause:      errors.New("token=secret"),
		}},
		Readiness: fakeReadinessChecker{},
	})
	request := httptest.NewRequest(http.MethodPost, "/api/v1/proxies/extract", strings.NewReader(`not-json`))
	response := httptest.NewRecorder()

	handler.ServeHTTP(response, request)

	if response.Code != http.StatusUnauthorized {
		t.Fatalf("status = %d, want %d; body=%s", response.Code, http.StatusUnauthorized, response.Body.String())
	}
	if challenge := response.Header().Get("WWW-Authenticate"); challenge != `Bearer realm="proxy-pool"` {
		t.Fatalf("challenge = %q", challenge)
	}
	if strings.Contains(response.Body.String(), "secret") {
		t.Fatalf("security response leaked cause: %s", response.Body.String())
	}
	if extractor.calls != 0 {
		t.Fatalf("extractor calls = %d, want 0", extractor.calls)
	}
}

func TestHandlerMapsUnexpectedIdentityFailureToInternalError(t *testing.T) {
	t.Parallel()
	extractor := &fakeExtractor{}
	handler := mustNewHandler(t, Config{BodyLimitBytes: 1024}, Dependencies{
		Extractor: extractor,
		Identity:  fakeIdentityResolver{err: errors.New("credential store secret")},
		Readiness: fakeReadinessChecker{},
	})
	request := httptest.NewRequest(http.MethodPost, "/api/v1/proxies/extract", strings.NewReader(`{"count":1}`))
	response := httptest.NewRecorder()

	handler.ServeHTTP(response, request)

	if response.Code != http.StatusInternalServerError {
		t.Fatalf("status = %d, want %d; body=%s", response.Code, http.StatusInternalServerError, response.Body.String())
	}
	if !strings.Contains(response.Body.String(), `"code":"INTERNAL_ERROR"`) {
		t.Fatalf("body = %s, want INTERNAL_ERROR", response.Body.String())
	}
	if strings.Contains(response.Body.String(), "secret") {
		t.Fatalf("response leaked dependency error: %s", response.Body.String())
	}
	if extractor.calls != 0 {
		t.Fatalf("extractor calls = %d, want 0", extractor.calls)
	}
}

func TestHandlerExtractMapsErrorsToProblemResponsesWithoutSensitiveLeakage(t *testing.T) {
	t.Parallel()
	tooManyRegions := `{"count":1,"filters":{"regions":["` + strings.Join(makeUniqueValues(65), `","`) + `"]}}`

	tests := []struct {
		name           string
		requestBody    string
		contentType    string
		idempotencyKey string
		extractErr     error
		bodyLimit      int64
		wantStatus     int
		wantCode       string
	}{
		{
			name:        "invalid json unknown field",
			requestBody: `{"count":1,"unexpected":true}`,
			contentType: "application/json",
			wantStatus:  http.StatusBadRequest,
			wantCode:    "INVALID_JSON",
		},
		{
			name:        "unsupported media type",
			requestBody: `{"count":1}`,
			contentType: "text/plain",
			wantStatus:  http.StatusUnsupportedMediaType,
			wantCode:    "UNSUPPORTED_MEDIA_TYPE",
		},
		{
			name:        "request body too large",
			requestBody: strings.Repeat(" ", 300) + `{"count":1}`,
			contentType: "application/json",
			wantStatus:  http.StatusRequestEntityTooLarge,
			wantCode:    "REQUEST_BODY_TOO_LARGE",
		},
		{
			name:           "invalid idempotency key",
			requestBody:    `{"count":1}`,
			contentType:    "application/json",
			idempotencyKey: "short",
			wantStatus:     http.StatusBadRequest,
			wantCode:       "INVALID_HEADER",
		},
		{
			name:        "invalid dto",
			requestBody: `{"count":1001,"filters":{"protocols":["http","http"]}}`,
			contentType: "application/json",
			wantStatus:  http.StatusUnprocessableEntity,
			wantCode:    "INVALID_REQUEST",
		},
		{
			name:        "too many filter values",
			requestBody: tooManyRegions,
			contentType: "application/json",
			bodyLimit:   4096,
			wantStatus:  http.StatusUnprocessableEntity,
			wantCode:    "INVALID_REQUEST",
		},
		{
			name:        "insufficient proxies",
			requestBody: `{"count":1,"fulfillment":"allOrNothing"}`,
			contentType: "application/json",
			extractErr:  domainExtraction.ErrInsufficientProxies,
			wantStatus:  http.StatusConflict,
			wantCode:    "INSUFFICIENT_PROXIES",
		},
		{
			name:        "idempotency conflict",
			requestBody: `{"count":1}`,
			contentType: "application/json",
			extractErr:  domainExtraction.ErrIdempotencyConflict,
			wantStatus:  http.StatusConflict,
			wantCode:    "IDEMPOTENCY_CONFLICT",
		},
		{
			name:        "admission rejected",
			requestBody: `{"count":1}`,
			contentType: "application/json",
			extractErr:  controllerExtraction.ErrAdmissionRejected,
			wantStatus:  http.StatusTooManyRequests,
			wantCode:    "RATE_LIMITED",
		},
		{
			name:        "count exceeded",
			requestBody: `{"count":1}`,
			contentType: "application/json",
			extractErr:  controllerExtraction.ErrCountExceeded,
			wantStatus:  http.StatusUnprocessableEntity,
			wantCode:    "COUNT_EXCEEDED",
		},
		{
			name:        "unexpected error",
			requestBody: `{"count":1}`,
			contentType: "application/json",
			extractErr:  errors.New("backend secret password leaked"),
			wantStatus:  http.StatusInternalServerError,
			wantCode:    "INTERNAL_ERROR",
		},
		{
			name:        "extraction unavailable",
			requestBody: `{"count":1}`,
			contentType: "application/json",
			extractErr: errors.Join(
				controllerExtraction.ErrUnavailable,
				domainExtraction.ErrStoreUnavailable,
				errors.New("redis dial failed: TOKEN"),
			),
			wantStatus: http.StatusServiceUnavailable,
			wantCode:   "SERVICE_UNAVAILABLE",
		},
	}

	for _, test := range tests {
		test := test
		t.Run(test.name, func(t *testing.T) {
			t.Parallel()
			extractor := &fakeExtractor{err: test.extractErr}
			bodyLimit := test.bodyLimit
			if bodyLimit == 0 {
				bodyLimit = 256
			}
			handler := mustNewHandler(t, Config{BodyLimitBytes: bodyLimit}, Dependencies{
				Extractor: extractor,
				Identity:  fakeIdentityResolver{identity: Identity{ClientID: "tenant-a", SourceIP: "198.51.100.8"}},
				Readiness: fakeReadinessChecker{},
			})
			request := httptest.NewRequest(http.MethodPost, "/api/v1/proxies/extract", strings.NewReader(test.requestBody))
			request.Header.Set("Content-Type", test.contentType)
			request.Header.Set(httpapi.HeaderRequestID, "req-err")
			if test.idempotencyKey != "" {
				request.Header.Set("Idempotency-Key", test.idempotencyKey)
			}

			response := httptest.NewRecorder()
			handler.ServeHTTP(response, request)

			if response.Code != test.wantStatus {
				t.Fatalf("status = %d, want %d", response.Code, test.wantStatus)
			}
			if got := response.Header().Get(httpapi.HeaderRequestID); got == "" {
				t.Fatal("X-Request-ID header is empty")
			}
			if got := response.Header().Get("Content-Type"); got != httpapi.ProblemContentType {
				t.Fatalf("Content-Type = %q, want %q", got, httpapi.ProblemContentType)
			}
			var problem httpapi.Problem
			if err := json.Unmarshal(response.Body.Bytes(), &problem); err != nil {
				t.Fatalf("decode problem: %v", err)
			}
			if problem.Status != test.wantStatus || problem.Code != test.wantCode {
				t.Fatalf("problem = %+v", problem)
			}
			if body := response.Body.String(); strings.Contains(body, `"password"`) ||
				strings.Contains(body, "secret") || strings.Contains(body, "redis") || strings.Contains(body, "TOKEN") {
				t.Fatalf("error body leaked sensitive data: %s", body)
			}
		})
	}
}

func makeUniqueValues(count int) []string {
	values := make([]string, count)
	for index := range values {
		values[index] = fmt.Sprintf("region-%d", index)
	}
	return values
}

func TestHandlerHealthRoutesAndRoutingEdges(t *testing.T) {
	t.Parallel()

	handler := mustNewHandler(t, Config{BodyLimitBytes: 1024}, Dependencies{
		Extractor: &fakeExtractor{},
		Identity:  fakeIdentityResolver{identity: Identity{ClientID: "tenant-a", SourceIP: "198.51.100.8"}},
		Readiness: fakeReadinessChecker{},
	})

	tests := []struct {
		name       string
		method     string
		target     string
		wantStatus int
		wantCT     string
	}{
		{name: "live ok", method: http.MethodGet, target: "/health/live", wantStatus: http.StatusOK, wantCT: httpapi.JSONContentType},
		{name: "ready ok", method: http.MethodGet, target: "/health/ready", wantStatus: http.StatusOK, wantCT: httpapi.JSONContentType},
		{name: "extract wrong method", method: http.MethodGet, target: "/api/v1/proxies/extract", wantStatus: http.StatusMethodNotAllowed, wantCT: httpapi.ProblemContentType},
		{name: "unknown route", method: http.MethodGet, target: "/missing", wantStatus: http.StatusNotFound, wantCT: httpapi.ProblemContentType},
	}

	for _, test := range tests {
		test := test
		t.Run(test.name, func(t *testing.T) {
			t.Parallel()
			request := httptest.NewRequest(test.method, test.target, nil)
			response := httptest.NewRecorder()
			handler.ServeHTTP(response, request)
			if response.Code != test.wantStatus {
				t.Fatalf("status = %d, want %d", response.Code, test.wantStatus)
			}
			if got := response.Header().Get(httpapi.HeaderRequestID); got == "" {
				t.Fatal("X-Request-ID header is empty")
			}
			if got := response.Header().Get("Content-Type"); got != test.wantCT {
				t.Fatalf("Content-Type = %q, want %q", got, test.wantCT)
			}
		})
	}
}

func TestHandlerReadyMapsDependencyFailureTo503Problem(t *testing.T) {
	t.Parallel()

	handler := mustNewHandler(t, Config{BodyLimitBytes: 1024}, Dependencies{
		Extractor: &fakeExtractor{},
		Identity:  fakeIdentityResolver{identity: Identity{ClientID: "tenant-a", SourceIP: "198.51.100.8"}},
		Readiness: fakeReadinessChecker{err: errors.New("storage unavailable")},
	})
	request := httptest.NewRequest(http.MethodGet, "/health/ready", nil)
	response := httptest.NewRecorder()

	handler.ServeHTTP(response, request)

	if response.Code != http.StatusServiceUnavailable {
		t.Fatalf("status = %d, want 503", response.Code)
	}
	var problem httpapi.Problem
	if err := json.Unmarshal(response.Body.Bytes(), &problem); err != nil {
		t.Fatalf("decode problem: %v", err)
	}
	if problem.Code != "SERVICE_UNAVAILABLE" {
		t.Fatalf("problem code = %q, want SERVICE_UNAVAILABLE", problem.Code)
	}
}

func mustNewHandler(t *testing.T, config Config, deps Dependencies) *Handler {
	t.Helper()
	handler, err := NewHandler(config, deps)
	if err != nil {
		t.Fatalf("NewHandler() error = %v", err)
	}
	return handler
}

type fakeExtractor struct {
	request  controllerExtraction.Request
	response controllerExtraction.Response
	err      error
	calls    int
}

func (f *fakeExtractor) Extract(_ context.Context, request controllerExtraction.Request) (controllerExtraction.Response, error) {
	f.calls++
	f.request = request
	return f.response, f.err
}

type fakeIdentityResolver struct {
	identity Identity
	err      error
}

func (f fakeIdentityResolver) Resolve(*http.Request) (Identity, error) {
	return f.identity, f.err
}

type fakeReadinessChecker struct {
	err error
}

func (f fakeReadinessChecker) Ready(context.Context) error {
	return f.err
}