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/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
}