sub2api-cn-relay-manager/internal/app/key_self_service_test.go

package app

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"sub2api-cn-relay-manager/internal/metrics"
	"sub2api-cn-relay-manager/internal/store/sqlite"
)

func makeCreateBody(groupID, displayName string, models []string) io.Reader {
	b, _ := json.Marshal(map[string]any{
		"logical_group_id": groupID,
		"display_name":     displayName,
		"allowed_models":   models,
	})
	return bytes.NewReader(b)
}

func makeCreateRequest(t *testing.T, method, path string, body io.Reader) *http.Request {
	t.Helper()
	req := httptest.NewRequest(method, path, body)
	req.Header.Set("Content-Type", "application/json")
	return req
}

func TestUserKeyAPIUsesPortalSubjectHeader(t *testing.T) {
	t.Parallel()
	store := openAppTestStore(t)
	defer closeAppTestStore(t, store)

	// Seed a logical group + route + host so resolveLogicalGroupHost succeeds
	_, _ = store.Hosts().Create(context.Background(), sqlite.Host{
		HostID:              "test-host",
		BaseURL:             "http://127.0.0.1:1",
		HostVersion:         "0.0.1",
		CapabilityProbeJSON: "{}",
		AuthType:            "apikey",
		AuthToken:           "test-token",
	})
	_, _ = store.LogicalGroups().Create(context.Background(), sqlite.LogicalGroup{
		LogicalGroupID: "gpt-shared",
		DisplayName:    "GPT Shared",
		Status:         "active",
	})
	_, _ = store.LogicalGroupRoutes().Create(context.Background(), sqlite.LogicalGroupRoute{
		RouteID:        "test-route",
		LogicalGroupID: "gpt-shared",
		Name:           "Test Route",
		Status:         "active",
		ShadowHostID:   "test-host",
		ShadowGroupID:  "999",
	})

	handler := NewAPIHandler("t", ActionSet{
		UserKeyHandler: buildUserKeyHandler(appTestDSN(t, store)),
	})

	req := makeCreateRequest(t, http.MethodPost, "/api/keys", makeCreateBody("gpt-shared", "portal key", []string{"gpt-5.4"}))
	req.Header.Set("X-Portal-Subject", "smoke-user")
	resp := httptestRecorder(handler, req)

	// We expect 500 because test host is unreachable (port 1), but the important
	// thing is the request decoded the subject header and reached the host resolution
	// step (not 401 "user credentials required")
	if resp.code == http.StatusUnauthorized || resp.code == http.StatusNotImplemented {
		t.Fatalf("status code = %d, expected to pass auth layer", resp.code)
	}

	var errResp struct {
		Error struct {
			Code    string `json:"code"`
			Message string `json:"message"`
		} `json:"error"`
	}
	if err := json.Unmarshal(resp.Body().Bytes(), &errResp); err == nil {
		if strings.Contains(errResp.Error.Message, "POST /api/v1/auth/login") ||
			strings.Contains(errResp.Error.Message, "no such host") ||
			strings.Contains(errResp.Error.Message, "connect: connection refused") ||
			strings.Contains(errResp.Error.Message, "dial tcp") {
			t.Logf("expected host-level error (not auth): code=%s msg=%s", errResp.Error.Code, errResp.Error.Message)
		} else {
			t.Logf("unexpected error shape: code=%s msg=%s", errResp.Error.Code, errResp.Error.Message)
		}
	}
}

func TestUserKeyCreateRejectsMissingSubject(t *testing.T) {
	t.Parallel()
	handler := NewAPIHandler("t", ActionSet{
		UserKeyHandler: buildUserKeyHandler(appTestDSN(t, openAppTestStore(t))),
	})
	req := makeCreateRequest(t, http.MethodPost, "/api/keys", makeCreateBody("gpt-shared", "portal key", nil))
	resp := httptestRecorder(handler, req)
	if resp.code != http.StatusUnauthorized {
		t.Fatalf("status code = %d, want 401", resp.code)
	}
}

func TestUserKeyCreateRejectsMissingGroup(t *testing.T) {
	t.Parallel()
	handler := NewAPIHandler("t", ActionSet{
		UserKeyHandler: buildUserKeyHandler(appTestDSN(t, openAppTestStore(t))),
	})
	body := bytes.NewReader([]byte(`{"display_name":"portal key"}`))
	req := makeCreateRequest(t, http.MethodPost, "/api/keys", body)
	req.Header.Set("X-Portal-Subject", "smoke-user")
	resp := httptestRecorder(handler, req)
	if resp.code != http.StatusBadRequest {
		t.Fatalf("status code = %d, want 400", resp.code)
	}
}

func TestUserKeyResetRejectsMissingSubject(t *testing.T) {
	t.Parallel()
	handler := NewAPIHandler("t", ActionSet{
		UserKeyHandler: buildUserKeyHandler(appTestDSN(t, openAppTestStore(t))),
	})
	req := httptest.NewRequest(http.MethodPost, "/api/keys/key_123/reset", nil)
	req.Header.Set("Content-Type", "application/json")
	resp := httptestRecorder(handler, req)
	if resp.code != http.StatusUnauthorized {
		t.Fatalf("status code = %d, want 401", resp.code)
	}
}

func TestUserKeyRateLimitNoDB(t *testing.T) {
	t.Parallel()
	store := openAppTestStore(t)
	defer closeAppTestStore(t, store)
	_, _ = store.LogicalGroups().Create(context.Background(), sqlite.LogicalGroup{
		LogicalGroupID: "gpt-shared",
		DisplayName:    "GPT Shared",
		Status:         "active",
	})

	handler := NewAPIHandler("t", ActionSet{
		UserKeyHandler: buildUserKeyHandler(appTestDSN(t, store)),
	})

	req := makeCreateRequest(t, http.MethodPost, "/api/keys", makeCreateBody("gpt-shared", "rate-test", nil))
	req.Header.Set("X-Portal-Subject", "rate-user")
	resp := httptestRecorder(handler, req)
	if resp.code == http.StatusUnauthorized || resp.code == http.StatusNotImplemented {
		t.Fatalf("status code = %d, expected to pass auth layer", resp.code)
	}
}

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

	store := openAppTestStore(t)
	defer closeAppTestStore(t, store)

	const logicalGroupID = "gpt-shared"
	const hostGroupID = "999"
	const subjectID = "portal-user:13"

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch {
		case r.Method == http.MethodGet && strings.HasPrefix(r.URL.RequestURI(), "/api/v1/admin/users?"):
			w.Write([]byte(`{"data":{"items":[]}}`))
		case r.Method == http.MethodPost && r.URL.Path == "/api/v1/admin/users":
			w.Write([]byte(`{"data":{"id":84,"email":"managed@sub2api.local"}}`))
		case r.Method == http.MethodPut && r.URL.Path == "/api/v1/admin/users/84":
			w.Write([]byte(`{"data":{"id":84}}`))
		case r.Method == http.MethodPost && r.URL.Path == "/api/v1/admin/users/84/balance":
			w.Write([]byte(`{"data":{"id":84}}`))
		case r.Method == http.MethodPost && r.URL.Path == "/api/v1/admin/subscriptions/assign":
			w.Write([]byte(`{"data":{"id":401}}`))
		case r.Method == http.MethodPost && r.URL.Path == "/api/v1/auth/login":
			var req map[string]any
			if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
				t.Fatalf("decode login request: %v", err)
			}
			expected := expectedManagedIdentity(subjectID, hostGroupID)
			if got := fmt.Sprint(req["email"]); got != expected.Email {
				t.Fatalf("login email = %q, want subject-scoped %q", got, expected.Email)
			}
			w.Write([]byte(`{"data":{"access_token":"user-jwt"}}`))
		case r.Method == http.MethodPost && r.URL.Path == "/api/v1/keys":
			var req map[string]any
			if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
				t.Fatalf("decode managed key request: %v", err)
			}
			expected := expectedManagedIdentity(subjectID, hostGroupID)
			if got := fmt.Sprint(req["custom_key"]); got != expected.CustomKey {
				t.Fatalf("custom_key = %q, want subject-scoped %q", got, expected.CustomKey)
			}
			w.Write([]byte(`{"data":{"id":501,"key":"placeholder-from-host","name":"managed-key"}}`))
		case r.Method == http.MethodPut && r.URL.Path == "/api/v1/admin/api-keys/501":
			w.Write([]byte(`{"data":{"api_key":{"id":501}}}`))
		default:
			w.WriteHeader(http.StatusNotFound)
		}
	}))
	defer server.Close()

	_, _ = store.Hosts().Create(context.Background(), sqlite.Host{
		HostID:              "test-host",
		BaseURL:             server.URL,
		HostVersion:         "0.0.1",
		CapabilityProbeJSON: "{}",
		AuthType:            "apikey",
		AuthToken:           "test-token",
	})
	_, _ = store.LogicalGroups().Create(context.Background(), sqlite.LogicalGroup{
		LogicalGroupID: logicalGroupID,
		DisplayName:    "GPT Shared",
		Status:         "active",
	})
	_, _ = store.LogicalGroupRoutes().Create(context.Background(), sqlite.LogicalGroupRoute{
		RouteID:        "test-route",
		LogicalGroupID: logicalGroupID,
		Name:           "Test Route",
		Status:         "active",
		ShadowHostID:   "test-host",
		ShadowGroupID:  hostGroupID,
	})

	handler := NewAPIHandler("t", ActionSet{
		UserKeyHandler: buildUserKeyHandler(appTestDSN(t, store)),
	})

	req := makeCreateRequest(t, http.MethodPost, "/api/keys", makeCreateBody(logicalGroupID, "portal key", []string{"gpt-5.4"}))
	req.Header.Set("X-Portal-Subject", subjectID)
	resp := httptestRecorder(handler, req)
	if resp.code != http.StatusCreated {
		t.Fatalf("status code = %d, want 201, body=%s", resp.code, resp.Body().String())
	}

	var createResp CreateUserKeyResponse
	if err := json.Unmarshal(resp.Body().Bytes(), &createResp); err != nil {
		t.Fatalf("decode create response: %v", err)
	}

	expected := expectedManagedIdentity(subjectID, hostGroupID)
	if createResp.PlaintextKey != expected.CustomKey {
		t.Fatalf("plaintext_key = %q, want subject-scoped %q", createResp.PlaintextKey, expected.CustomKey)
	}
	wantMasked := "sk-****" + expected.CustomKey[len(expected.CustomKey)-4:]
	if createResp.Key.MaskedPreview != wantMasked {
		t.Fatalf("masked_preview = %q, want %q", createResp.Key.MaskedPreview, wantMasked)
	}

	record, err := store.UserKeys().GetByID(context.Background(), createResp.Key.KeyID)
	if err != nil {
		t.Fatalf("UserKeys().GetByID() error = %v", err)
	}
	if record.KeyFingerprint != "sha256:"+sha256Hex(expected.CustomKey) {
		t.Fatalf("key_fingerprint = %q, want sha256 of returned plaintext key", record.KeyFingerprint)
	}
	if record.MaskedPreview != wantMasked {
		t.Fatalf("stored masked_preview = %q, want %q", record.MaskedPreview, wantMasked)
	}
}

type managedIdentityExpectation struct {
	Email     string
	CustomKey string
}

func expectedManagedIdentity(selector, groupID string) managedIdentityExpectation {
	normalizedSelector := strings.TrimSpace(strings.ToLower(selector))
	sum := sha256Hex(normalizedSelector + "|" + strings.TrimSpace(groupID))
	prefix := expectedManagedPrefix(selector)
	shortHash := sum[:16]
	return managedIdentityExpectation{
		Email:     fmt.Sprintf("%s-%s@sub2api.local", prefix, shortHash),
		CustomKey: "sk-relay-" + sum[:32],
	}
}

func expectedManagedPrefix(value string) string {
	value = strings.ToLower(strings.TrimSpace(value))
	var b strings.Builder
	lastDash := false
	for _, r := range value {
		switch {
		case r >= 'a' && r <= 'z', r >= '0' && r <= '9':
			b.WriteRune(r)
			lastDash = false
		case !lastDash:
			b.WriteByte('-')
			lastDash = true
		}
	}
	prefix := strings.Trim(b.String(), "-")
	if prefix == "" {
		prefix = "relay-sub"
	}
	if len(prefix) > 24 {
		prefix = strings.Trim(prefix[:24], "-")
	}
	return prefix
}

func TestUserKeyAPIMetricsMiddlewareAndCreateMetric(t *testing.T) {
	t.Parallel()
	handler := NewAPIHandler("t", ActionSet{
		UserKeyHandler: buildUserKeyHandler(appTestDSN(t, openAppTestStore(t))),
	})
	req := makeCreateRequest(t, http.MethodPost, "/api/keys", makeCreateBody("", "portal key", nil))
	req.Header.Set("X-Portal-Subject", "smoke-user")
	_ = httptestRecorder(handler, req)

	metricsReq := httptest.NewRequest(http.MethodGet, "/metrics", nil)
	metricsResp := httptest.NewRecorder()
	metrics.Handler().ServeHTTP(metricsResp, metricsReq)
	body := metricsResp.Body.String()
	if !strings.Contains(body, "http_requests_total") {
		t.Fatal("expected metrics endpoint to expose http_requests_total after middleware-wrapped request")
	}
	if !strings.Contains(body, "user_key_operations_total") {
		t.Fatal("expected metrics endpoint to expose user_key_operations_total after create validation failure")
	}
}

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

	store := openAppTestStore(t)
	defer closeAppTestStore(t, store)

	const logicalGroupID = "gpt-shared"
	const hostGroupID = "999"
	const subjectID = "portal-user:lifecycle"
	const keyID = "key_lifecycle"

	var allowedGroupsUpdates [][]int64
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch {
		case r.Method == http.MethodGet && strings.HasPrefix(r.URL.RequestURI(), "/api/v1/admin/users?"):
			w.Write([]byte(`{"data":{"items":[{"id":84,"email":"` + expectedManagedIdentity(subjectID, hostGroupID).Email + `"}]}}`))
		case r.Method == http.MethodPut && r.URL.Path == "/api/v1/admin/users/84":
			var payload struct {
				AllowedGroups []int64 `json:"allowed_groups"`
			}
			if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
				t.Fatalf("decode update payload: %v", err)
			}
			allowedGroupsUpdates = append(allowedGroupsUpdates, append([]int64(nil), payload.AllowedGroups...))
			w.Write([]byte(`{"data":{"id":84}}`))
		default:
			w.WriteHeader(http.StatusNotFound)
		}
	}))
	defer server.Close()

	_, _ = store.Hosts().Create(context.Background(), sqlite.Host{
		HostID:              "test-host",
		BaseURL:             server.URL,
		HostVersion:         "0.0.1",
		CapabilityProbeJSON: "{}",
		AuthType:            "apikey",
		AuthToken:           "test-token",
	})
	_, _ = store.LogicalGroups().Create(context.Background(), sqlite.LogicalGroup{
		LogicalGroupID: logicalGroupID,
		DisplayName:    "GPT Shared",
		Status:         "active",
	})
	_, _ = store.LogicalGroupRoutes().Create(context.Background(), sqlite.LogicalGroupRoute{
		RouteID:        "test-route",
		LogicalGroupID: logicalGroupID,
		Name:           "Test Route",
		Status:         "active",
		ShadowHostID:   "test-host",
		ShadowGroupID:  hostGroupID,
	})
	if _, err := store.UserKeys().Create(context.Background(), sqlite.UserKeyRecord{
		KeyID:          keyID,
		OwnerSubjectID: subjectID,
		KeyFingerprint: "sha256:test",
		MaskedPreview:  "sk-****test",
		DisplayName:    "lifecycle key",
		LogicalGroupID: logicalGroupID,
		AllowedModels:  []string{"gpt-5.4"},
		AdminStatus:    "active",
		QuotaStatus:    "ok",
	}); err != nil {
		t.Fatalf("UserKeys().Create() error = %v", err)
	}

	handler := buildUserKeyHandler(appTestDSN(t, store))
	paused, err := handler.pauseFn(context.Background(), keyID, subjectID, "")
	if err != nil {
		t.Fatalf("pauseFn() error = %v", err)
	}
	if paused.AdminStatus != "paused" {
		t.Fatalf("pauseFn() admin_status = %q, want paused", paused.AdminStatus)
	}
	row, err := store.UserKeys().GetByID(context.Background(), keyID)
	if err != nil {
		t.Fatalf("UserKeys().GetByID() after pause error = %v", err)
	}
	if row.AdminStatus != "paused" {
		t.Fatalf("stored admin_status after pause = %q, want paused", row.AdminStatus)
	}

	resumed, err := handler.resumeFn(context.Background(), keyID, subjectID)
	if err != nil {
		t.Fatalf("resumeFn() error = %v", err)
	}
	if resumed.AdminStatus != "active" {
		t.Fatalf("resumeFn() admin_status = %q, want active", resumed.AdminStatus)
	}
	row, err = store.UserKeys().GetByID(context.Background(), keyID)
	if err != nil {
		t.Fatalf("UserKeys().GetByID() after resume error = %v", err)
	}
	if row.AdminStatus != "active" {
		t.Fatalf("stored admin_status after resume = %q, want active", row.AdminStatus)
	}

	if err := handler.deleteFn(context.Background(), keyID, subjectID); err != nil {
		t.Fatalf("deleteFn() error = %v", err)
	}
	row, err = store.UserKeys().GetByID(context.Background(), keyID)
	if err != nil {
		t.Fatalf("UserKeys().GetByID() after delete error = %v", err)
	}
	if row.AdminStatus != "retired" {
		t.Fatalf("stored admin_status after delete = %q, want retired", row.AdminStatus)
	}

	if len(allowedGroupsUpdates) != 2 {
		t.Fatalf("allowedGroupsUpdates len = %d, want 2", len(allowedGroupsUpdates))
	}
	if len(allowedGroupsUpdates[0]) != 0 {
		t.Fatalf("pause allowed_groups = %#v, want empty", allowedGroupsUpdates[0])
	}
	if len(allowedGroupsUpdates[1]) != 1 || allowedGroupsUpdates[1][0] != 999 {
		t.Fatalf("resume allowed_groups = %#v, want [999]", allowedGroupsUpdates[1])
	}
}

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

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/api/v1/admin/groups":
			w.Write([]byte(`{"data":{"items":[{"id":321,"name":"group-by-name"}]}}`))
		case "/api/v1/admin/channels", "/api/v1/admin/payment/plans", "/api/v1/admin/accounts":
			w.Write([]byte(`{"data":{"items":[]}}`))
		default:
			w.WriteHeader(http.StatusNotFound)
		}
	}))
	defer server.Close()

	client, err := newSub2APIClient(server.URL, CreateHostAuth{Type: "apikey", Token: "test-token"})
	if err != nil {
		t.Fatalf("newSub2APIClient() error = %v", err)
	}

	groupID, err := resolveShadowHostGroupID(context.Background(), client, sqlite.LogicalGroupRoute{ShadowGroupID: "group-by-name"})
	if err != nil {
		t.Fatalf("resolveShadowHostGroupID() error = %v", err)
	}
	if groupID != "321" {
		t.Fatalf("groupID = %q, want 321", groupID)
	}
}