ModelRouter/internal/gateway/apiv1_test.go

package gateway

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"llmsproxy/internal/config"
)

// v1Gateway builds a gateway with one source and two keys, and returns the
// secrets so a test can assert they never appear in a response.
func v1Gateway(t *testing.T) (*Gateway, string, string) {
	t.Helper()
	g := newTestGateway(t, config.Source{
		Name:    "up",
		BaseURL: "http://up.example/v1",
		APIKey:  "sk-up-secret",
		Adapter: "openai",
		Models:  []config.Model{{ID: "m1", Kind: "chat"}},
	})
	// The shared helper only seeds one admin key; add a user key through the
	// core so role-gated endpoints have something to reject.
	if _, err := g.core.CreateKey("agent", "user", nil, "test agent key"); err != nil {
		t.Fatalf("CreateKey: %v", err)
	}
	rec := doReq(t, g, http.MethodGet, "/api/keys", "")
	var doc struct {
		Keys []config.GWKey `json:"keys"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &doc); err != nil {
		t.Fatalf("decode keys: %v", err)
	}
	userKey := ""
	for _, k := range doc.Keys {
		if k.Role == "user" {
			userKey = k.Key
		}
	}
	if userKey == "" {
		t.Fatal("no user key was created")
	}
	return g, "sk-test", userKey
}

func decodeJSON(t *testing.T, body string, v interface{}) {
	t.Helper()
	if err := json.Unmarshal([]byte(body), v); err != nil {
		t.Fatalf("decode %s: %v", body, err)
	}
}

func TestAPIV1IndexIsDiscoverable(t *testing.T) {
	g, _, _ := v1Gateway(t)
	rec := doReq(t, g, http.MethodGet, "/api/v1", "")
	if rec.Code != http.StatusOK {
		t.Fatalf("status = %d, want 200: %s", rec.Code, rec.Body.String())
	}
	var doc struct {
		APIVersion  string `json:"api_version"`
		Description string `json:"description"`
		Discovery   []struct {
			Method  string `json:"method"`
			Path    string `json:"path"`
			Auth    string `json:"auth"`
			Summary string `json:"summary"`
		} `json:"discovery"`
		Conventions map[string]interface{} `json:"conventions"`
	}
	decodeJSON(t, rec.Body.String(), &doc)
	if doc.APIVersion != "v1" {
		t.Errorf("api_version = %q", doc.APIVersion)
	}
	if doc.Description == "" {
		t.Error("index should carry a usage description")
	}
	if len(doc.Discovery) == 0 {
		t.Fatal("discovery list is empty")
	}
	want := map[string]bool{
		"/api/v1/overview": false, "/api/v1/sources": false,
		"/api/v1/auto": false, "/api/v1/models": false,
		"/api/v1/health": false, "/api/v1/keys": false,
		"/api/sources": false, "/api/auto": false, "/v1/chat/completions": false,
	}
	for _, e := range doc.Discovery {
		if _, ok := want[e.Path]; ok {
			want[e.Path] = true
		}
		if e.Method == "" || e.Path == "" || e.Auth == "" || e.Summary == "" {
			t.Errorf("incomplete discovery entry: %+v", e)
		}
	}
	for p, found := range want {
		if !found {
			t.Errorf("discovery is missing %s", p)
		}
	}
	if doc.Conventions["errors"] == nil {
		t.Error("conventions should document the error shape")
	}
}

func TestAPIV1OverviewSummarisesState(t *testing.T) {
	g, _, _ := v1Gateway(t)
	rec := doReq(t, g, http.MethodGet, "/api/v1/overview", "")
	if rec.Code != http.StatusOK {
		t.Fatalf("status = %d: %s", rec.Code, rec.Body.String())
	}
	var doc struct {
		Summary struct {
			SourceCount     int `json:"source_count"`
			ModelCount      int `json:"model_count"`
			GatewayKeyCount int `json:"gateway_key_count"`
		} `json:"summary"`
		Sources []map[string]interface{} `json:"sources"`
		Auto    []map[string]interface{} `json:"auto"`
		Health  []map[string]interface{} `json:"health"`
		Caller  struct {
			Role string `json:"role"`
		} `json:"caller"`
	}
	decodeJSON(t, rec.Body.String(), &doc)
	if doc.Summary.SourceCount != 1 {
		t.Errorf("source_count = %d, want 1", doc.Summary.SourceCount)
	}
	if doc.Summary.ModelCount != 1 {
		t.Errorf("model_count = %d, want 1", doc.Summary.ModelCount)
	}
	if doc.Summary.GatewayKeyCount < 2 {
		t.Errorf("gateway_key_count = %d, want >= 2", doc.Summary.GatewayKeyCount)
	}
	if len(doc.Sources) != 1 || doc.Sources[0]["name"] != "up" {
		t.Errorf("sources = %+v", doc.Sources)
	}
	if doc.Caller.Role != "admin" {
		t.Errorf("caller.role = %q, want admin", doc.Caller.Role)
	}
}

func TestAPIV1NeverEchoesSecrets(t *testing.T) {
	g, admin, user := v1Gateway(t)
	for _, path := range []string{"/api/v1", "/api/v1/overview", "/api/v1/sources", "/api/v1/sources/up", "/api/v1/keys", "/api/v1/models", "/api/v1/health", "/api/v1/auto"} {
		rec := doReq(t, g, http.MethodGet, path, "")
		if rec.Code != http.StatusOK {
			t.Fatalf("%s status = %d: %s", path, rec.Code, rec.Body.String())
		}
		if strings.Contains(rec.Body.String(), "sk-up-secret") {
			t.Errorf("%s leaked the source api_key", path)
		}
		if admin != "" && strings.Contains(rec.Body.String(), admin) {
			t.Errorf("%s leaked the admin gateway key", path)
		}
		if strings.Contains(rec.Body.String(), user) {
			t.Errorf("%s leaked the user gateway key", path)
		}
	}
	// Masking must still be useful: it says a key is configured.
	rec := doReq(t, g, http.MethodGet, "/api/v1/sources/up", "")
	if !strings.Contains(rec.Body.String(), `"api_key_set":true`) {
		t.Error("masked source should report that a key is configured")
	}
}

func TestAPIV1KeysRequiresAdminAndMasks(t *testing.T) {
	g, _, user := v1Gateway(t)
	// Authenticate as the user key by swapping the shared helper's header.
	req := newAuthedRequest(t, http.MethodGet, "/api/v1/keys", user)
	rec := serveViaHandler(t, g, req)
	if rec.Code != http.StatusForbidden {
		t.Errorf("user role on /api/v1/keys = %d, want 403: %s", rec.Code, rec.Body.String())
	}
}

func TestAPIV1SourcesNotFound(t *testing.T) {
	g, _, _ := v1Gateway(t)
	rec := doReq(t, g, http.MethodGet, "/api/v1/sources/nope", "")
	if rec.Code != http.StatusNotFound {
		t.Errorf("status = %d, want 404: %s", rec.Code, rec.Body.String())
	}
}

func TestAPIV1RejectsWrites(t *testing.T) {
	g, _, _ := v1Gateway(t)
	for _, m := range []string{http.MethodPost, http.MethodPut, http.MethodDelete, http.MethodPatch} {
		rec := doReq(t, g, m, "/api/v1/sources", `{}`)
		if rec.Code == http.StatusOK {
			t.Errorf("%s /api/v1/sources returned 200; the v1 read facade must not mutate", m)
		}
	}
	// And the source must be untouched.
	rec := doReq(t, g, http.MethodGet, "/api/v1/sources", "")
	if !strings.Contains(rec.Body.String(), `"name":"up"`) {
		t.Error("the source disappeared after a rejected write")
	}
}

func TestAPIV1ModelsGroupsBySource(t *testing.T) {
	g, _, _ := v1Gateway(t)
	rec := doReq(t, g, http.MethodGet, "/api/v1/models", "")
	if rec.Code != http.StatusOK {
		t.Fatalf("status = %d", rec.Code)
	}
	var doc struct {
		Count    int                 `json:"count"`
		Models   []string            `json:"models"`
		BySource map[string][]string `json:"by_source"`
	}
	decodeJSON(t, rec.Body.String(), &doc)
	if doc.Count != 1 || doc.Models[0] != "m1" {
		t.Errorf("models = %+v", doc.Models)
	}
	if len(doc.BySource["up"]) != 1 || doc.BySource["up"][0] != "m1" {
		t.Errorf("by_source = %+v", doc.BySource)
	}
}

func TestAPIV1RequiresAuthentication(t *testing.T) {
	g, _, _ := v1Gateway(t)
	for _, path := range []string{"/api/v1", "/api/v1/overview", "/api/v1/sources", "/api/v1/health", "/api/v1/models"} {
		req := newRequest(http.MethodGet, path, "")
		rec := serveViaHandler(t, g, req)
		if rec.Code != http.StatusUnauthorized {
			t.Errorf("%s without a key = %d, want 401", path, rec.Code)
		}
	}
}

func TestMaskKeyHidesTheSecret(t *testing.T) {
	cases := map[string]string{
		"":                   "",
		"ab":                 "ab…",
		"sk-gw-abcdefghijkl": "sk-gw-…ijkl",
		"enc:v1:AAAA":        "(sealed)",
	}
	for in, want := range cases {
		if got := maskKey(in); got != want {
			t.Errorf("maskKey(%q) = %q, want %q", in, got, want)
		}
	}
	// The masked form must not reveal the middle of the secret.
	full := maskKey("sk-gw-1234567890abcdef")
	if strings.Contains(full, "4567890abcdef") {
		t.Errorf("maskKey leaked the middle: %q", full)
	}
}

// newRequest builds an unauthenticated request.
func newRequest(method, path, body string) *http.Request {
	req, _ := http.NewRequest(method, path, strings.NewReader(body))
	if body != "" {
		req.Header.Set("Content-Type", "application/json")
	}
	return req
}

// newAuthedRequest builds a request carrying the given bearer key.
func newAuthedRequest(t *testing.T, method, path, key string) *http.Request {
	t.Helper()
	req := newRequest(method, path, "")
	req.Header.Set("Authorization", "Bearer "+key)
	return req
}

// serveViaHandler pushes a request through the full handler chain (auth included).
func serveViaHandler(t *testing.T, g *Gateway, req *http.Request) *httptest.ResponseRecorder {
	t.Helper()
	rec := httptest.NewRecorder()
	g.Handler().ServeHTTP(rec, req)
	return rec
}

// TestAPIV1RevealRequiresAdmin: the only way to read a key in the clear is the
// explicit admin-only reveal. Without the role check this endpoint would hand
// every source credential to any valid (even user-scoped) key.
func TestAPIV1RevealRequiresAdmin(t *testing.T) {
	g, _, user := v1Gateway(t)
	rec := serveViaHandler(t, g, newAuthedRequest(t, http.MethodGet, "/api/v1/sources/up?reveal=credentials", user))
	if rec.Code != http.StatusForbidden {
		t.Errorf("user reveal = %d, want 403: %s", rec.Code, rec.Body.String())
	}
	if strings.Contains(rec.Body.String(), "sk-up-secret") {
		t.Error("a forbidden reveal still leaked the key")
	}
	rec = doReq(t, g, http.MethodGet, "/api/v1/sources/up?reveal=credentials", "")
	if rec.Code != http.StatusOK {
		t.Fatalf("admin reveal = %d: %s", rec.Code, rec.Body.String())
	}
	if !strings.Contains(rec.Body.String(), "sk-up-secret") {
		t.Error("admin reveal should return the key in the clear (the edit form needs it)")
	}
	// The masked default view must still hide it on the same path.
	rec = doReq(t, g, http.MethodGet, "/api/v1/sources/up", "")
	if strings.Contains(rec.Body.String(), "sk-up-secret") {
		t.Error("the default view leaked the key")
	}
}