AI 网关切换 OpenAI 兼容面并移除 chat 端点,新增模型黑白名单

2026-07-12 17:48:28 +08:00
parent 7706f59549
commit 489cb49cb3
34 changed files with 2602 additions and 2901 deletions
@@ -0,0 +1,206 @@
+package api
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"strconv"
+	"strings"
+	"testing"
+	"time"
+
+	"gorm.io/gorm"
+
+	"oci-portal/internal/model"
+)
+
+// seedGatewayModel 直插启用渠道与模型缓存,绕过云同步。
+func seedGatewayModel(t *testing.T, db *gorm.DB, name string) {
+	t.Helper()
+	ch := &model.AiChannel{Name: "t-" + name, OciConfigID: 1, Region: "r-" + name, Enabled: true, Priority: 1, Weight: 1}
+	if err := db.Create(ch).Error; err != nil {
+		t.Fatalf("seed channel: %v", err)
+	}
+	cache := &model.AiModelCache{ChannelID: ch.ID, ModelOcid: "ocid1.." + name, Name: name, Vendor: "v", SyncedAt: time.Now()}
+	if err := db.Create(cache).Error; err != nil {
+		t.Fatalf("seed model cache: %v", err)
+	}
+}
+
+func TestAiGatewayKeyModelRestrict(t *testing.T) {
+	r, auth, _, db := newTestRouterDB(t)
+	token, _, err := auth.Login(context.Background(), "admin", "pass123", "127.0.0.1", "")
+	if err != nil {
+		t.Fatalf("login: %v", err)
+	}
+
+	// 受限密钥:白名单含脏输入,入库应规范化为 2 项(cohere + ghost)
+	w := doRequest(t, r, http.MethodPost, "/api/v1/ai-keys", token,
+		`{"name":"limited","value":"limited-key-1234","models":[" cohere.command-a-03-2025 ","cohere.command-a-03-2025","","ghost-model"]}`)
+	if w.Code != http.StatusCreated {
+		t.Fatalf("创建受限密钥 status = %d, body %s", w.Code, w.Body.String())
+	}
+	var created struct {
+		Item model.AiKey `json:"item"`
+	}
+	if err := json.Unmarshal(w.Body.Bytes(), &created); err != nil {
+		t.Fatalf("decode: %v", err)
+	}
+	if len(created.Item.Models) != 2 || created.Item.Models[0] != "cohere.command-a-03-2025" {
+		t.Fatalf("创建后 models 未规范化: %v", created.Item.Models)
+	}
+	// 不限密钥对照
+	w = doRequest(t, r, http.MethodPost, "/api/v1/ai-keys", token, `{"name":"open","value":"open-key-12345"}`)
+	if w.Code != http.StatusCreated {
+		t.Fatalf("创建不限密钥 status = %d, body %s", w.Code, w.Body.String())
+	}
+
+	seedGatewayModel(t, db, "cohere.command-a-03-2025")
+	seedGatewayModel(t, db, "meta.llama-3.3-70b-instruct")
+
+	deny := []string{"model_not_found", "无权访问"}
+	pass := []string{"未知模型"} // 穿过白名单闸门,由编排层以「未知模型」拒绝(池内无 ghost-model)
+	tests := []struct {
+		name     string
+		key      string
+		path     string
+		body     string
+		wantCode int
+		wantSub  []string
+	}{
+		{"responses 流式同样拦截", "limited-key-1234", "/ai/v1/responses",
+			`{"model":"meta.llama-3.3-70b-instruct","stream":true,"input":"hi"}`, 404, deny},
+		{"responses 白名单内穿透", "limited-key-1234", "/ai/v1/responses",
+			`{"model":"ghost-model","input":"hi"}`, 404, pass},
+		{"embeddings 白名单外拦截", "limited-key-1234", "/ai/v1/embeddings",
+			`{"model":"meta.llama-3.3-70b-instruct","input":["hi"]}`, 404, deny},
+		{"messages 白名单外拦截为 Anthropic 体", "limited-key-1234", "/ai/v1/messages",
+			`{"model":"meta.llama-3.3-70b-instruct","max_tokens":16,"messages":[{"role":"user","content":"hi"}]}`, 404,
+			[]string{`"type":"error"`, "model_not_found", "无权访问"}},
+		{"responses 白名单外拦截", "limited-key-1234", "/ai/v1/responses",
+			`{"model":"meta.llama-3.3-70b-instruct","input":"hi"}`, 404, deny},
+		{"不限密钥穿透闸门", "open-key-12345", "/ai/v1/responses",
+			`{"model":"ghost-model","input":"hi"}`, 404, pass},
+		{"chat completions 端点已删除", "open-key-12345", "/ai/v1/chat/completions",
+			`{"model":"ghost-model","messages":[{"role":"user","content":"hi"}]}`, 404, []string{"not found"}},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			w := doRequest(t, r, http.MethodPost, tt.path, tt.key, tt.body)
+			if w.Code != tt.wantCode {
+				t.Fatalf("status = %d, want %d, body %s", w.Code, tt.wantCode, w.Body.String())
+			}
+			for _, sub := range tt.wantSub {
+				if !strings.Contains(w.Body.String(), sub) {
+					t.Errorf("body 缺少 %q: %s", sub, w.Body.String())
+				}
+			}
+		})
+	}
+}
+
+func TestAiGatewayModelsFilteredByKey(t *testing.T) {
+	r, auth, _, db := newTestRouterDB(t)
+	token, _, err := auth.Login(context.Background(), "admin", "pass123", "127.0.0.1", "")
+	if err != nil {
+		t.Fatalf("login: %v", err)
+	}
+	seedGatewayModel(t, db, "cohere.command-a-03-2025")
+	seedGatewayModel(t, db, "meta.llama-3.3-70b-instruct")
+	w := doRequest(t, r, http.MethodPost, "/api/v1/ai-keys", token,
+		`{"name":"limited","value":"limited-key-1234","models":["cohere.command-a-03-2025"]}`)
+	if w.Code != http.StatusCreated {
+		t.Fatalf("创建密钥 status = %d", w.Code)
+	}
+	w = doRequest(t, r, http.MethodPost, "/api/v1/ai-keys", token, `{"name":"open","value":"open-key-12345"}`)
+	if w.Code != http.StatusCreated {
+		t.Fatalf("创建密钥 status = %d", w.Code)
+	}
+
+	tests := []struct {
+		name    string
+		key     string
+		want    []string
+		notWant []string
+	}{
+		{"受限密钥仅见交集", "limited-key-1234", []string{"cohere.command-a-03-2025"}, []string{"meta.llama-3.3-70b-instruct"}},
+		{"不限密钥全量可见", "open-key-12345", []string{"cohere.command-a-03-2025", "meta.llama-3.3-70b-instruct"}, nil},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			w := doRequest(t, r, http.MethodGet, "/ai/v1/models", tt.key, "")
+			if w.Code != http.StatusOK {
+				t.Fatalf("status = %d, body %s", w.Code, w.Body.String())
+			}
+			for _, sub := range tt.want {
+				if !strings.Contains(w.Body.String(), sub) {
+					t.Errorf("应包含 %q: %s", sub, w.Body.String())
+				}
+			}
+			for _, sub := range tt.notWant {
+				if strings.Contains(w.Body.String(), sub) {
+					t.Errorf("不应包含 %q: %s", sub, w.Body.String())
+				}
+			}
+		})
+	}
+}
+
+func TestAiKeyModelsUpdateRoundTrip(t *testing.T) {
+	r, auth, _, _ := newTestRouterDB(t)
+	token, _, err := auth.Login(context.Background(), "admin", "pass123", "127.0.0.1", "")
+	if err != nil {
+		t.Fatalf("login: %v", err)
+	}
+	w := doRequest(t, r, http.MethodPost, "/api/v1/ai-keys", token, `{"name":"k","value":"round-key-1234"}`)
+	if w.Code != http.StatusCreated {
+		t.Fatalf("创建 status = %d", w.Code)
+	}
+	var created struct {
+		Item model.AiKey `json:"item"`
+	}
+	_ = json.Unmarshal(w.Body.Bytes(), &created)
+	id := created.Item.ID
+
+	listModels := func() []string {
+		w := doRequest(t, r, http.MethodGet, "/api/v1/ai-keys", token, "")
+		var resp struct {
+			Items []model.AiKey `json:"items"`
+		}
+		_ = json.Unmarshal(w.Body.Bytes(), &resp)
+		for _, k := range resp.Items {
+			if k.ID == id {
+				return k.Models
+			}
+		}
+		t.Fatalf("密钥 %d 不在列表中", id)
+		return nil
+	}
+
+	tests := []struct {
+		name string
+		body string
+		want []string
+	}{
+		{"设置白名单", `{"models":[" a-model ","a-model","b-model"]}`, []string{"a-model", "b-model"}},
+		{"不传 models 保持不变", `{"name":"k2"}`, []string{"a-model", "b-model"}},
+		{"空数组清空恢复不限", `{"models":[]}`, nil},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			w := doRequest(t, r, http.MethodPut, "/api/v1/ai-keys/"+strconv.Itoa(int(id)), token, tt.body)
+			if w.Code != http.StatusNoContent {
+				t.Fatalf("update status = %d, body %s", w.Code, w.Body.String())
+			}
+			got := listModels()
+			if len(got) != len(tt.want) {
+				t.Fatalf("models = %v, want %v", got, tt.want)
+			}
+			for i := range tt.want {
+				if got[i] != tt.want[i] {
+					t.Fatalf("models = %v, want %v", got, tt.want)
+				}
+			}
+		})
+	}
+}