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

package app

import (
	"bytes"
	"context"
	"encoding/json"
	"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)
	}
}