feat(groups)!: enforce model allowlists across gateway endpoints

This commit is contained in:
Zero Clover
2026-09-05 16:05:17 +08:00
parent 07bf8b92fd
commit cff3f89850
73 changed files with 3691 additions and 778 deletions
+9 -9
View File
@@ -134,8 +134,8 @@ type Group struct {
DefaultMappedModel string `json:"default_mapped_model,omitempty"`
// OpenAI Messages 调度模型配置:按 Claude 系列/精确模型映射到目标 GPT 模型
MessagesDispatchModelConfig domain.OpenAIMessagesDispatchModelConfig `json:"messages_dispatch_model_config,omitempty"`
// 自定义 /v1/models 展示列表配置;仅影响模型列表响应,不影响调度
ModelsListConfig domain.GroupModelsListConfig `json:"models_list_config,omitempty"`
// 分组模型白名单:同时约束模型列表接口与请求准入
ModelAllowlist domain.GroupModelAllowlist `json:"model_allowlist,omitempty"`
// 固定账号获取 Codex Model Manifest 配置;开启后 /models 请求只用选定账号拉取(仅 openai 平台)
CodexModelsManifestConfig domain.GroupCodexModelsManifestConfig `json:"codex_models_manifest_config,omitempty"`
// 分组 RPM 上限,0 表示不限制;设置后接管该分组用户的限流
@@ -258,7 +258,7 @@ func (*Group) scanValues(columns []string) ([]any, error) {
values := make([]any, len(columns))
for i := range columns {
switch columns[i] {
case group.FieldVideoModelPrices, group.FieldModelPricing, group.FieldModelRouting, group.FieldSupportedModelScopes, group.FieldMessagesDispatchModelConfig, group.FieldModelsListConfig, group.FieldCodexModelsManifestConfig, group.FieldReasoningEffortMappings:
case group.FieldVideoModelPrices, group.FieldModelPricing, group.FieldModelRouting, group.FieldSupportedModelScopes, group.FieldMessagesDispatchModelConfig, group.FieldModelAllowlist, group.FieldCodexModelsManifestConfig, group.FieldReasoningEffortMappings:
values[i] = new([]byte)
case group.FieldPeakRateEnabled, group.FieldIsExclusive, group.FieldAllowImageGeneration, group.FieldAllowBatchImageGeneration, group.FieldImageRateIndependent, group.FieldVideoRateIndependent, group.FieldLongContextPricingEnabled, group.FieldClaudeCodeOnly, group.FieldModelRoutingEnabled, group.FieldMcpXMLInject, group.FieldAllowMessagesDispatch, group.FieldAllowLive, group.FieldForceOpenaiFast, group.FieldFreeOpenaiFast, group.FieldRequireOauthOnly, group.FieldRequirePrivacySet, group.FieldProfitControlEnabled:
values[i] = new(sql.NullBool)
@@ -662,12 +662,12 @@ func (_m *Group) assignValues(columns []string, values []any) error {
return fmt.Errorf("unmarshal field messages_dispatch_model_config: %w", err)
}
}
case group.FieldModelsListConfig:
case group.FieldModelAllowlist:
if value, ok := values[i].(*[]byte); !ok {
return fmt.Errorf("unexpected type %T for field models_list_config", values[i])
return fmt.Errorf("unexpected type %T for field model_allowlist", values[i])
} else if value != nil && len(*value) > 0 {
if err := json.Unmarshal(*value, &_m.ModelsListConfig); err != nil {
return fmt.Errorf("unmarshal field models_list_config: %w", err)
if err := json.Unmarshal(*value, &_m.ModelAllowlist); err != nil {
return fmt.Errorf("unmarshal field model_allowlist: %w", err)
}
}
case group.FieldCodexModelsManifestConfig:
@@ -1007,8 +1007,8 @@ func (_m *Group) String() string {
builder.WriteString("messages_dispatch_model_config=")
builder.WriteString(fmt.Sprintf("%v", _m.MessagesDispatchModelConfig))
builder.WriteString(", ")
builder.WriteString("models_list_config=")
builder.WriteString(fmt.Sprintf("%v", _m.ModelsListConfig))
builder.WriteString("model_allowlist=")
builder.WriteString(fmt.Sprintf("%v", _m.ModelAllowlist))
builder.WriteString(", ")
builder.WriteString("codex_models_manifest_config=")
builder.WriteString(fmt.Sprintf("%v", _m.CodexModelsManifestConfig))
+5 -5
View File
@@ -130,8 +130,8 @@ const (
FieldDefaultMappedModel = "default_mapped_model"
// FieldMessagesDispatchModelConfig holds the string denoting the messages_dispatch_model_config field in the database.
FieldMessagesDispatchModelConfig = "messages_dispatch_model_config"
// FieldModelsListConfig holds the string denoting the models_list_config field in the database.
FieldModelsListConfig = "models_list_config"
// FieldModelAllowlist holds the string denoting the model_allowlist field in the database.
FieldModelAllowlist = "model_allowlist"
// FieldCodexModelsManifestConfig holds the string denoting the codex_models_manifest_config field in the database.
FieldCodexModelsManifestConfig = "codex_models_manifest_config"
// FieldRpmLimit holds the string denoting the rpm_limit field in the database.
@@ -280,7 +280,7 @@ var Columns = []string{
FieldRequirePrivacySet,
FieldDefaultMappedModel,
FieldMessagesDispatchModelConfig,
FieldModelsListConfig,
FieldModelAllowlist,
FieldCodexModelsManifestConfig,
FieldRpmLimit,
FieldMaxReasoningEffort,
@@ -412,8 +412,8 @@ var (
DefaultMappedModelValidator func(string) error
// DefaultMessagesDispatchModelConfig holds the default value on creation for the "messages_dispatch_model_config" field.
DefaultMessagesDispatchModelConfig domain.OpenAIMessagesDispatchModelConfig
// DefaultModelsListConfig holds the default value on creation for the "models_list_config" field.
DefaultModelsListConfig domain.GroupModelsListConfig
// DefaultModelAllowlist holds the default value on creation for the "model_allowlist" field.
DefaultModelAllowlist domain.GroupModelAllowlist
// DefaultCodexModelsManifestConfig holds the default value on creation for the "codex_models_manifest_config" field.
DefaultCodexModelsManifestConfig domain.GroupCodexModelsManifestConfig
// DefaultRpmLimit holds the default value on creation for the "rpm_limit" field.
+32 -32
View File
@@ -788,16 +788,16 @@ func (_c *GroupCreate) SetNillableMessagesDispatchModelConfig(v *domain.OpenAIMe
return _c
}
// SetModelsListConfig sets the "models_list_config" field.
func (_c *GroupCreate) SetModelsListConfig(v domain.GroupModelsListConfig) *GroupCreate {
_c.mutation.SetModelsListConfig(v)
// SetModelAllowlist sets the "model_allowlist" field.
func (_c *GroupCreate) SetModelAllowlist(v domain.GroupModelAllowlist) *GroupCreate {
_c.mutation.SetModelAllowlist(v)
return _c
}
// SetNillableModelsListConfig sets the "models_list_config" field if the given value is not nil.
func (_c *GroupCreate) SetNillableModelsListConfig(v *domain.GroupModelsListConfig) *GroupCreate {
// SetNillableModelAllowlist sets the "model_allowlist" field if the given value is not nil.
func (_c *GroupCreate) SetNillableModelAllowlist(v *domain.GroupModelAllowlist) *GroupCreate {
if v != nil {
_c.SetModelsListConfig(*v)
_c.SetModelAllowlist(*v)
}
return _c
}
@@ -1175,9 +1175,9 @@ func (_c *GroupCreate) defaults() error {
v := group.DefaultMessagesDispatchModelConfig
_c.mutation.SetMessagesDispatchModelConfig(v)
}
if _, ok := _c.mutation.ModelsListConfig(); !ok {
v := group.DefaultModelsListConfig
_c.mutation.SetModelsListConfig(v)
if _, ok := _c.mutation.ModelAllowlist(); !ok {
v := group.DefaultModelAllowlist
_c.mutation.SetModelAllowlist(v)
}
if _, ok := _c.mutation.CodexModelsManifestConfig(); !ok {
v := group.DefaultCodexModelsManifestConfig
@@ -1381,8 +1381,8 @@ func (_c *GroupCreate) check() error {
if _, ok := _c.mutation.MessagesDispatchModelConfig(); !ok {
return &ValidationError{Name: "messages_dispatch_model_config", err: errors.New(`ent: missing required field "Group.messages_dispatch_model_config"`)}
}
if _, ok := _c.mutation.ModelsListConfig(); !ok {
return &ValidationError{Name: "models_list_config", err: errors.New(`ent: missing required field "Group.models_list_config"`)}
if _, ok := _c.mutation.ModelAllowlist(); !ok {
return &ValidationError{Name: "model_allowlist", err: errors.New(`ent: missing required field "Group.model_allowlist"`)}
}
if _, ok := _c.mutation.CodexModelsManifestConfig(); !ok {
return &ValidationError{Name: "codex_models_manifest_config", err: errors.New(`ent: missing required field "Group.codex_models_manifest_config"`)}
@@ -1673,9 +1673,9 @@ func (_c *GroupCreate) createSpec() (*Group, *sqlgraph.CreateSpec) {
_spec.SetField(group.FieldMessagesDispatchModelConfig, field.TypeJSON, value)
_node.MessagesDispatchModelConfig = value
}
if value, ok := _c.mutation.ModelsListConfig(); ok {
_spec.SetField(group.FieldModelsListConfig, field.TypeJSON, value)
_node.ModelsListConfig = value
if value, ok := _c.mutation.ModelAllowlist(); ok {
_spec.SetField(group.FieldModelAllowlist, field.TypeJSON, value)
_node.ModelAllowlist = value
}
if value, ok := _c.mutation.CodexModelsManifestConfig(); ok {
_spec.SetField(group.FieldCodexModelsManifestConfig, field.TypeJSON, value)
@@ -2795,15 +2795,15 @@ func (u *GroupUpsert) UpdateMessagesDispatchModelConfig() *GroupUpsert {
return u
}
// SetModelsListConfig sets the "models_list_config" field.
func (u *GroupUpsert) SetModelsListConfig(v domain.GroupModelsListConfig) *GroupUpsert {
u.Set(group.FieldModelsListConfig, v)
// SetModelAllowlist sets the "model_allowlist" field.
func (u *GroupUpsert) SetModelAllowlist(v domain.GroupModelAllowlist) *GroupUpsert {
u.Set(group.FieldModelAllowlist, v)
return u
}
// UpdateModelsListConfig sets the "models_list_config" field to the value that was provided on create.
func (u *GroupUpsert) UpdateModelsListConfig() *GroupUpsert {
u.SetExcluded(group.FieldModelsListConfig)
// UpdateModelAllowlist sets the "model_allowlist" field to the value that was provided on create.
func (u *GroupUpsert) UpdateModelAllowlist() *GroupUpsert {
u.SetExcluded(group.FieldModelAllowlist)
return u
}
@@ -4054,17 +4054,17 @@ func (u *GroupUpsertOne) UpdateMessagesDispatchModelConfig() *GroupUpsertOne {
})
}
// SetModelsListConfig sets the "models_list_config" field.
func (u *GroupUpsertOne) SetModelsListConfig(v domain.GroupModelsListConfig) *GroupUpsertOne {
// SetModelAllowlist sets the "model_allowlist" field.
func (u *GroupUpsertOne) SetModelAllowlist(v domain.GroupModelAllowlist) *GroupUpsertOne {
return u.Update(func(s *GroupUpsert) {
s.SetModelsListConfig(v)
s.SetModelAllowlist(v)
})
}
// UpdateModelsListConfig sets the "models_list_config" field to the value that was provided on create.
func (u *GroupUpsertOne) UpdateModelsListConfig() *GroupUpsertOne {
// UpdateModelAllowlist sets the "model_allowlist" field to the value that was provided on create.
func (u *GroupUpsertOne) UpdateModelAllowlist() *GroupUpsertOne {
return u.Update(func(s *GroupUpsert) {
s.UpdateModelsListConfig()
s.UpdateModelAllowlist()
})
}
@@ -5500,17 +5500,17 @@ func (u *GroupUpsertBulk) UpdateMessagesDispatchModelConfig() *GroupUpsertBulk {
})
}
// SetModelsListConfig sets the "models_list_config" field.
func (u *GroupUpsertBulk) SetModelsListConfig(v domain.GroupModelsListConfig) *GroupUpsertBulk {
// SetModelAllowlist sets the "model_allowlist" field.
func (u *GroupUpsertBulk) SetModelAllowlist(v domain.GroupModelAllowlist) *GroupUpsertBulk {
return u.Update(func(s *GroupUpsert) {
s.SetModelsListConfig(v)
s.SetModelAllowlist(v)
})
}
// UpdateModelsListConfig sets the "models_list_config" field to the value that was provided on create.
func (u *GroupUpsertBulk) UpdateModelsListConfig() *GroupUpsertBulk {
// UpdateModelAllowlist sets the "model_allowlist" field to the value that was provided on create.
func (u *GroupUpsertBulk) UpdateModelAllowlist() *GroupUpsertBulk {
return u.Update(func(s *GroupUpsert) {
s.UpdateModelsListConfig()
s.UpdateModelAllowlist()
})
}
+16 -16
View File
@@ -1073,16 +1073,16 @@ func (_u *GroupUpdate) SetNillableMessagesDispatchModelConfig(v *domain.OpenAIMe
return _u
}
// SetModelsListConfig sets the "models_list_config" field.
func (_u *GroupUpdate) SetModelsListConfig(v domain.GroupModelsListConfig) *GroupUpdate {
_u.mutation.SetModelsListConfig(v)
// SetModelAllowlist sets the "model_allowlist" field.
func (_u *GroupUpdate) SetModelAllowlist(v domain.GroupModelAllowlist) *GroupUpdate {
_u.mutation.SetModelAllowlist(v)
return _u
}
// SetNillableModelsListConfig sets the "models_list_config" field if the given value is not nil.
func (_u *GroupUpdate) SetNillableModelsListConfig(v *domain.GroupModelsListConfig) *GroupUpdate {
// SetNillableModelAllowlist sets the "model_allowlist" field if the given value is not nil.
func (_u *GroupUpdate) SetNillableModelAllowlist(v *domain.GroupModelAllowlist) *GroupUpdate {
if v != nil {
_u.SetModelsListConfig(*v)
_u.SetModelAllowlist(*v)
}
return _u
}
@@ -1876,8 +1876,8 @@ func (_u *GroupUpdate) sqlSave(ctx context.Context) (_node int, err error) {
if value, ok := _u.mutation.MessagesDispatchModelConfig(); ok {
_spec.SetField(group.FieldMessagesDispatchModelConfig, field.TypeJSON, value)
}
if value, ok := _u.mutation.ModelsListConfig(); ok {
_spec.SetField(group.FieldModelsListConfig, field.TypeJSON, value)
if value, ok := _u.mutation.ModelAllowlist(); ok {
_spec.SetField(group.FieldModelAllowlist, field.TypeJSON, value)
}
if value, ok := _u.mutation.CodexModelsManifestConfig(); ok {
_spec.SetField(group.FieldCodexModelsManifestConfig, field.TypeJSON, value)
@@ -3267,16 +3267,16 @@ func (_u *GroupUpdateOne) SetNillableMessagesDispatchModelConfig(v *domain.OpenA
return _u
}
// SetModelsListConfig sets the "models_list_config" field.
func (_u *GroupUpdateOne) SetModelsListConfig(v domain.GroupModelsListConfig) *GroupUpdateOne {
_u.mutation.SetModelsListConfig(v)
// SetModelAllowlist sets the "model_allowlist" field.
func (_u *GroupUpdateOne) SetModelAllowlist(v domain.GroupModelAllowlist) *GroupUpdateOne {
_u.mutation.SetModelAllowlist(v)
return _u
}
// SetNillableModelsListConfig sets the "models_list_config" field if the given value is not nil.
func (_u *GroupUpdateOne) SetNillableModelsListConfig(v *domain.GroupModelsListConfig) *GroupUpdateOne {
// SetNillableModelAllowlist sets the "model_allowlist" field if the given value is not nil.
func (_u *GroupUpdateOne) SetNillableModelAllowlist(v *domain.GroupModelAllowlist) *GroupUpdateOne {
if v != nil {
_u.SetModelsListConfig(*v)
_u.SetModelAllowlist(*v)
}
return _u
}
@@ -4100,8 +4100,8 @@ func (_u *GroupUpdateOne) sqlSave(ctx context.Context) (_node *Group, err error)
if value, ok := _u.mutation.MessagesDispatchModelConfig(); ok {
_spec.SetField(group.FieldMessagesDispatchModelConfig, field.TypeJSON, value)
}
if value, ok := _u.mutation.ModelsListConfig(); ok {
_spec.SetField(group.FieldModelsListConfig, field.TypeJSON, value)
if value, ok := _u.mutation.ModelAllowlist(); ok {
_spec.SetField(group.FieldModelAllowlist, field.TypeJSON, value)
}
if value, ok := _u.mutation.CodexModelsManifestConfig(); ok {
_spec.SetField(group.FieldCodexModelsManifestConfig, field.TypeJSON, value)
+1 -1
View File
@@ -960,7 +960,7 @@ var (
{Name: "require_privacy_set", Type: field.TypeBool, Default: false},
{Name: "default_mapped_model", Type: field.TypeString, Size: 100, Default: ""},
{Name: "messages_dispatch_model_config", Type: field.TypeJSON, SchemaType: map[string]string{"postgres": "jsonb"}},
{Name: "models_list_config", Type: field.TypeJSON, SchemaType: map[string]string{"postgres": "jsonb"}},
{Name: "model_allowlist", Type: field.TypeJSON, SchemaType: map[string]string{"postgres": "jsonb"}},
{Name: "codex_models_manifest_config", Type: field.TypeJSON, SchemaType: map[string]string{"postgres": "jsonb"}},
{Name: "rpm_limit", Type: field.TypeInt, Default: 0},
{Name: "max_reasoning_effort", Type: field.TypeString, Size: 20, Default: ""},
+27 -27
View File
@@ -22164,7 +22164,7 @@ type GroupMutation struct {
require_privacy_set *bool
default_mapped_model *string
messages_dispatch_model_config *domain.OpenAIMessagesDispatchModelConfig
models_list_config *domain.GroupModelsListConfig
model_allowlist *domain.GroupModelAllowlist
codex_models_manifest_config *domain.GroupCodexModelsManifestConfig
rpm_limit *int
addrpm_limit *int
@@ -25164,40 +25164,40 @@ func (m *GroupMutation) ResetMessagesDispatchModelConfig() {
m.messages_dispatch_model_config = nil
}
// SetModelsListConfig sets the "models_list_config" field.
func (m *GroupMutation) SetModelsListConfig(dmlc domain.GroupModelsListConfig) {
m.models_list_config = &dmlc
// SetModelAllowlist sets the "model_allowlist" field.
func (m *GroupMutation) SetModelAllowlist(dma domain.GroupModelAllowlist) {
m.model_allowlist = &dma
}
// ModelsListConfig returns the value of the "models_list_config" field in the mutation.
func (m *GroupMutation) ModelsListConfig() (r domain.GroupModelsListConfig, exists bool) {
v := m.models_list_config
// ModelAllowlist returns the value of the "model_allowlist" field in the mutation.
func (m *GroupMutation) ModelAllowlist() (r domain.GroupModelAllowlist, exists bool) {
v := m.model_allowlist
if v == nil {
return
}
return *v, true
}
// OldModelsListConfig returns the old "models_list_config" field's value of the Group entity.
// OldModelAllowlist returns the old "model_allowlist" field's value of the Group entity.
// If the Group object wasn't provided to the builder, the object is fetched from the database.
// An error is returned if the mutation operation is not UpdateOne, or the database query fails.
func (m *GroupMutation) OldModelsListConfig(ctx context.Context) (v domain.GroupModelsListConfig, err error) {
func (m *GroupMutation) OldModelAllowlist(ctx context.Context) (v domain.GroupModelAllowlist, err error) {
if !m.op.Is(OpUpdateOne) {
return v, errors.New("OldModelsListConfig is only allowed on UpdateOne operations")
return v, errors.New("OldModelAllowlist is only allowed on UpdateOne operations")
}
if m.id == nil || m.oldValue == nil {
return v, errors.New("OldModelsListConfig requires an ID field in the mutation")
return v, errors.New("OldModelAllowlist requires an ID field in the mutation")
}
oldValue, err := m.oldValue(ctx)
if err != nil {
return v, fmt.Errorf("querying old value for OldModelsListConfig: %w", err)
return v, fmt.Errorf("querying old value for OldModelAllowlist: %w", err)
}
return oldValue.ModelsListConfig, nil
return oldValue.ModelAllowlist, nil
}
// ResetModelsListConfig resets all changes to the "models_list_config" field.
func (m *GroupMutation) ResetModelsListConfig() {
m.models_list_config = nil
// ResetModelAllowlist resets all changes to the "model_allowlist" field.
func (m *GroupMutation) ResetModelAllowlist() {
m.model_allowlist = nil
}
// SetCodexModelsManifestConfig sets the "codex_models_manifest_config" field.
@@ -26093,8 +26093,8 @@ func (m *GroupMutation) Fields() []string {
if m.messages_dispatch_model_config != nil {
fields = append(fields, group.FieldMessagesDispatchModelConfig)
}
if m.models_list_config != nil {
fields = append(fields, group.FieldModelsListConfig)
if m.model_allowlist != nil {
fields = append(fields, group.FieldModelAllowlist)
}
if m.codex_models_manifest_config != nil {
fields = append(fields, group.FieldCodexModelsManifestConfig)
@@ -26242,8 +26242,8 @@ func (m *GroupMutation) Field(name string) (ent.Value, bool) {
return m.DefaultMappedModel()
case group.FieldMessagesDispatchModelConfig:
return m.MessagesDispatchModelConfig()
case group.FieldModelsListConfig:
return m.ModelsListConfig()
case group.FieldModelAllowlist:
return m.ModelAllowlist()
case group.FieldCodexModelsManifestConfig:
return m.CodexModelsManifestConfig()
case group.FieldRpmLimit:
@@ -26383,8 +26383,8 @@ func (m *GroupMutation) OldField(ctx context.Context, name string) (ent.Value, e
return m.OldDefaultMappedModel(ctx)
case group.FieldMessagesDispatchModelConfig:
return m.OldMessagesDispatchModelConfig(ctx)
case group.FieldModelsListConfig:
return m.OldModelsListConfig(ctx)
case group.FieldModelAllowlist:
return m.OldModelAllowlist(ctx)
case group.FieldCodexModelsManifestConfig:
return m.OldCodexModelsManifestConfig(ctx)
case group.FieldRpmLimit:
@@ -26809,12 +26809,12 @@ func (m *GroupMutation) SetField(name string, value ent.Value) error {
}
m.SetMessagesDispatchModelConfig(v)
return nil
case group.FieldModelsListConfig:
v, ok := value.(domain.GroupModelsListConfig)
case group.FieldModelAllowlist:
v, ok := value.(domain.GroupModelAllowlist)
if !ok {
return fmt.Errorf("unexpected type %T for field %s", value, name)
}
m.SetModelsListConfig(v)
m.SetModelAllowlist(v)
return nil
case group.FieldCodexModelsManifestConfig:
v, ok := value.(domain.GroupCodexModelsManifestConfig)
@@ -27554,8 +27554,8 @@ func (m *GroupMutation) ResetField(name string) error {
case group.FieldMessagesDispatchModelConfig:
m.ResetMessagesDispatchModelConfig()
return nil
case group.FieldModelsListConfig:
m.ResetModelsListConfig()
case group.FieldModelAllowlist:
m.ResetModelAllowlist()
return nil
case group.FieldCodexModelsManifestConfig:
m.ResetCodexModelsManifestConfig()
+4 -4
View File
@@ -1183,10 +1183,10 @@ func init() {
groupDescMessagesDispatchModelConfig := groupFields[53].Descriptor()
// group.DefaultMessagesDispatchModelConfig holds the default value on creation for the messages_dispatch_model_config field.
group.DefaultMessagesDispatchModelConfig = groupDescMessagesDispatchModelConfig.Default.(domain.OpenAIMessagesDispatchModelConfig)
// groupDescModelsListConfig is the schema descriptor for models_list_config field.
groupDescModelsListConfig := groupFields[54].Descriptor()
// group.DefaultModelsListConfig holds the default value on creation for the models_list_config field.
group.DefaultModelsListConfig = groupDescModelsListConfig.Default.(domain.GroupModelsListConfig)
// groupDescModelAllowlist is the schema descriptor for model_allowlist field.
groupDescModelAllowlist := groupFields[54].Descriptor()
// group.DefaultModelAllowlist holds the default value on creation for the model_allowlist field.
group.DefaultModelAllowlist = groupDescModelAllowlist.Default.(domain.GroupModelAllowlist)
// groupDescCodexModelsManifestConfig is the schema descriptor for codex_models_manifest_config field.
groupDescCodexModelsManifestConfig := groupFields[55].Descriptor()
// group.DefaultCodexModelsManifestConfig holds the default value on creation for the codex_models_manifest_config field.
+3 -3
View File
@@ -262,10 +262,10 @@ func (Group) Fields() []ent.Field {
Default(domain.OpenAIMessagesDispatchModelConfig{}).
SchemaType(map[string]string{dialect.Postgres: "jsonb"}).
Comment("OpenAI Messages 调度模型配置:按 Claude 系列/精确模型映射到目标 GPT 模型"),
field.JSON("models_list_config", domain.GroupModelsListConfig{}).
Default(domain.GroupModelsListConfig{}).
field.JSON("model_allowlist", domain.GroupModelAllowlist{}).
Default(domain.GroupModelAllowlist{}).
SchemaType(map[string]string{dialect.Postgres: "jsonb"}).
Comment("自定义 /v1/models 展示列表配置;仅影响模型列表响应,不影响调度"),
Comment("分组模型白名单:同时约束模型列表接口与请求准入"),
field.JSON("codex_models_manifest_config", domain.GroupCodexModelsManifestConfig{}).
Default(domain.GroupCodexModelsManifestConfig{}).
SchemaType(map[string]string{dialect.Postgres: "jsonb"}).
@@ -0,0 +1,8 @@
package domain
// GroupModelAllowlist 是分组级模型白名单:开启后既过滤模型列表类接口的返回内容,
// 也约束分组实际可调用的模型(网关准入在合成路由改写与调度之前完成)。
type GroupModelAllowlist struct {
Enabled bool `json:"enabled"`
Models []string `json:"models,omitempty"`
}
@@ -1,7 +0,0 @@
package domain
// GroupModelsListConfig controls the optional custom /v1/models response list.
type GroupModelsListConfig struct {
Enabled bool `json:"enabled"`
Models []string `json:"models,omitempty"`
}
@@ -33,7 +33,7 @@ func setupAdminRouter() (*gin.Engine, *stubAdminService) {
router.GET("/api/v1/admin/groups", groupHandler.List)
router.GET("/api/v1/admin/groups/all", groupHandler.GetAll)
router.GET("/api/v1/admin/groups/:id/models-list-candidates", groupHandler.GetModelsListCandidates)
router.GET("/api/v1/admin/groups/:id/model-allowlist-candidates", groupHandler.GetGroupModelAllowlistCandidates)
router.GET("/api/v1/admin/groups/:id/composite-routes", groupHandler.ListCompositeRoutes)
router.POST("/api/v1/admin/groups/:id/composite-routes", groupHandler.CreateCompositeRoute)
router.POST("/api/v1/admin/groups/:id/composite-routes/preview", groupHandler.PreviewCompositeRoute)
@@ -184,7 +184,7 @@ func TestGroupHandlerEndpoints(t *testing.T) {
require.Equal(t, http.StatusOK, rec.Code)
rec = httptest.NewRecorder()
req = httptest.NewRequest(http.MethodGet, "/api/v1/admin/groups/0/models-list-candidates?platform=openai", nil)
req = httptest.NewRequest(http.MethodGet, "/api/v1/admin/groups/0/model-allowlist-candidates?platform=openai", nil)
router.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
require.Contains(t, rec.Body.String(), "gpt-5.5")
@@ -153,7 +153,7 @@ type CreateGroupRequest struct {
RequirePrivacySet bool `json:"require_privacy_set"`
DefaultMappedModel string `json:"default_mapped_model"`
MessagesDispatchModelConfig service.OpenAIMessagesDispatchModelConfig `json:"messages_dispatch_model_config"`
ModelsListConfig service.GroupModelsListConfig `json:"models_list_config"`
ModelAllowlist service.GroupModelAllowlist `json:"model_allowlist"`
// 固定账号 manifest 配置;创建路径禁止开启,仅编辑可配置。
CodexModelsManifestConfig service.GroupCodexModelsManifestConfig `json:"codex_models_manifest_config"`
// 分组 RPM 上限(0 = 不限制)
@@ -228,7 +228,7 @@ type UpdateGroupRequest struct {
RequirePrivacySet *bool `json:"require_privacy_set"`
DefaultMappedModel *string `json:"default_mapped_model"`
MessagesDispatchModelConfig *service.OpenAIMessagesDispatchModelConfig `json:"messages_dispatch_model_config"`
ModelsListConfig *service.GroupModelsListConfig `json:"models_list_config"`
ModelAllowlist *service.GroupModelAllowlist `json:"model_allowlist"`
// 固定账号 manifest 配置;nil 表示不修改。
CodexModelsManifestConfig *service.GroupCodexModelsManifestConfig `json:"codex_models_manifest_config"`
// 分组 RPM 上限(0 = 不限制);nil 表示未提供不改动
@@ -471,9 +471,9 @@ func (h *GroupHandler) GetByID(c *gin.Context) {
response.Success(c, dto.GroupFromServiceAdmin(group))
}
// GetModelsListCandidates handles getting candidate model IDs for custom /v1/models list.
// GET /api/v1/admin/groups/:id/models-list-candidates
func (h *GroupHandler) GetModelsListCandidates(c *gin.Context) {
// GetGroupModelAllowlistCandidates handles getting candidate model IDs for the group model allowlist.
// GET /api/v1/admin/groups/:id/model-allowlist-candidates
func (h *GroupHandler) GetGroupModelAllowlistCandidates(c *gin.Context) {
groupID, err := strconv.ParseInt(c.Param("id"), 10, 64)
if err != nil || groupID < 0 {
response.BadRequest(c, "Invalid group ID")
@@ -568,7 +568,7 @@ func (h *GroupHandler) Create(c *gin.Context) {
RequirePrivacySet: req.RequirePrivacySet,
DefaultMappedModel: req.DefaultMappedModel,
MessagesDispatchModelConfig: req.MessagesDispatchModelConfig,
ModelsListConfig: req.ModelsListConfig,
ModelAllowlist: req.ModelAllowlist,
CodexModelsManifestConfig: req.CodexModelsManifestConfig,
RPMLimit: req.RPMLimit,
MaxReasoningEffort: req.MaxReasoningEffort,
@@ -701,7 +701,7 @@ func (h *GroupHandler) Update(c *gin.Context) {
RequirePrivacySet: req.RequirePrivacySet,
DefaultMappedModel: req.DefaultMappedModel,
MessagesDispatchModelConfig: req.MessagesDispatchModelConfig,
ModelsListConfig: req.ModelsListConfig,
ModelAllowlist: req.ModelAllowlist,
CodexModelsManifestConfig: req.CodexModelsManifestConfig,
RPMLimit: req.RPMLimit,
MaxReasoningEffort: req.MaxReasoningEffort,
@@ -29,19 +29,19 @@ func TestAuthHandlerGetCurrentUserReturnsProfileCompatibilityFields(t *testing.T
AvatarURL: "https://cdn.example.com/linuxdo.png",
AvatarSource: "remote_url",
},
identities: []service.UserAuthIdentityRecord{
{
ProviderType: "linuxdo",
ProviderKey: "linuxdo",
ProviderSubject: "linuxdo-subject-31",
VerifiedAt: &verifiedAt,
Metadata: map[string]any{
"username": "linuxdo-handle",
"avatar_url": "https://cdn.example.com/linuxdo.png",
},
identities: []service.UserAuthIdentityRecord{
{
ProviderType: "linuxdo",
ProviderKey: "linuxdo",
ProviderSubject: "linuxdo-subject-31",
VerifiedAt: &verifiedAt,
Metadata: map[string]any{
"username": "linuxdo-handle",
"avatar_url": "https://cdn.example.com/linuxdo.png",
},
},
}
},
}
handler := &AuthHandler{
userService: service.NewUserService(repo, nil, nil, nil),
@@ -139,9 +139,28 @@ func (h *BatchImageHandler) Models(c *gin.Context) {
batchImageError(c, err)
return
}
// 分组级模型白名单:ListModels 已按分组账号枚举,这里再按白名单过滤一次。
if apiKey, ok := middleware.GetAPIKeyFromContext(c); ok && apiKey != nil && apiKey.Group != nil {
got.Data = filterBatchImageModelsByAllowlist(got.Data, apiKey.Group.ModelAllowlist)
}
c.JSON(http.StatusOK, got)
}
// filterBatchImageModelsByAllowlist 按分组模型白名单过滤批量生图模型列表。
// 白名单未开启时原样返回。
func filterBatchImageModelsByAllowlist(models []service.BatchImagePublicModel, allowlist service.GroupModelAllowlist) []service.BatchImagePublicModel {
if !allowlist.Enabled {
return models
}
filtered := make([]service.BatchImagePublicModel, 0, len(models))
for _, model := range models {
if allowlist.Allows(model.ID) {
filtered = append(filtered, model)
}
}
return filtered
}
func (h *BatchImageHandler) Items(c *gin.Context) {
owner, ok := batchImageOwnerFromContext(c)
if !ok {
+1 -1
View File
@@ -159,7 +159,7 @@ func GroupFromServiceAdmin(g *service.Group) *AdminGroup {
MCPXMLInject: g.MCPXMLInject,
DefaultMappedModel: g.DefaultMappedModel,
MessagesDispatchModelConfig: g.MessagesDispatchModelConfig,
ModelsListConfig: g.ModelsListConfig,
ModelAllowlist: g.ModelAllowlist,
CodexModelsManifestConfig: g.CodexModelsManifestConfig,
SupportedModelScopes: g.SupportedModelScopes,
AccountCount: g.AccountCount,
+1 -1
View File
@@ -189,7 +189,7 @@ type AdminGroup struct {
// OpenAI Messages 调度配置(仅 openai 平台使用)
DefaultMappedModel string `json:"default_mapped_model"`
MessagesDispatchModelConfig domain.OpenAIMessagesDispatchModelConfig `json:"messages_dispatch_model_config"`
ModelsListConfig domain.GroupModelsListConfig `json:"models_list_config"`
ModelAllowlist service.GroupModelAllowlist `json:"model_allowlist"`
// 固定账号获取 Codex Model Manifest 配置(仅 openai 平台使用)。
CodexModelsManifestConfig domain.GroupCodexModelsManifestConfig `json:"codex_models_manifest_config"`
+38 -78
View File
@@ -1089,9 +1089,12 @@ func (h *GatewayHandler) Models(c *gin.Context) {
if platform == service.PlatformComposite {
availableModels := h.compositeAvailableModels(c.Request.Context(), groupID)
if apiKey != nil && apiKey.Group != nil && apiKey.Group.CustomModelsListEnabled() {
availableModels = filterModelsByCustomList(availableModels, defaultModelIDsForPlatform(service.PlatformComposite), apiKey.Group.ModelsListConfig.Models)
writeCustomModelsList(c, service.PlatformComposite, availableModels)
if apiKey != nil && apiKey.Group != nil && apiKey.Group.ModelAllowlistEnabled() {
source := availableModels
if len(source) == 0 {
source = defaultModelIDsForPlatform(service.PlatformComposite)
}
writeAllowlistedModelsList(c, service.PlatformComposite, apiKey.Group.ModelAllowlist.FilterForListing(source))
return
}
if len(availableModels) > 0 {
@@ -1104,10 +1107,9 @@ func (h *GatewayHandler) Models(c *gin.Context) {
// Get available models from account configurations for the selected group platform.
availableModels := h.gatewayService.GetAvailableModels(c.Request.Context(), groupID, platform)
if apiKey != nil && apiKey.Group != nil && apiKey.Group.CustomModelsListEnabled() {
fallbackModels := defaultModelIDsForPlatform(platform)
availableModels = filterModelsByCustomList(customModelsListSource(platform, availableModels, fallbackModels), fallbackModels, apiKey.Group.ModelsListConfig.Models)
writeCustomModelsList(c, platform, availableModels)
if apiKey != nil && apiKey.Group != nil && apiKey.Group.ModelAllowlistEnabled() {
source := modelListingSource(platform, availableModels, defaultModelIDsForPlatform(platform))
writeAllowlistedModelsList(c, platform, apiKey.Group.ModelAllowlist.FilterForListing(source))
return
}
@@ -1192,8 +1194,12 @@ func (h *GatewayHandler) codexModelIDsForGroup(ctx context.Context, group *servi
if platform == service.PlatformComposite {
availableModels := h.compositeAvailableModels(ctx, groupID)
fallbackModels := defaultCodexModelIDsForPlatform(service.PlatformComposite)
if group.CustomModelsListEnabled() {
return filterModelsByCustomList(availableModels, fallbackModels, group.ModelsListConfig.Models)
if group.ModelAllowlistEnabled() {
source := availableModels
if len(source) == 0 {
source = fallbackModels
}
return group.ModelAllowlist.FilterForListing(source)
}
if len(availableModels) > 0 {
return availableModels
@@ -1203,12 +1209,8 @@ func (h *GatewayHandler) codexModelIDsForGroup(ctx context.Context, group *servi
availableModels := h.gatewayService.GetAvailableModels(ctx, groupID, platform)
fallbackModels := defaultCodexModelIDsForPlatform(platform)
if group.CustomModelsListEnabled() {
return filterModelsByCustomList(
customModelsListSource(platform, availableModels, fallbackModels),
fallbackModels,
group.ModelsListConfig.Models,
)
if group.ModelAllowlistEnabled() {
return group.ModelAllowlist.FilterForListing(modelListingSource(platform, availableModels, fallbackModels))
}
if len(availableModels) > 0 {
return availableModels
@@ -1267,7 +1269,7 @@ func writeModelsList(c *gin.Context, platform string, modelIDs []string) {
})
}
func writeCustomModelsList(c *gin.Context, platform string, modelIDs []string) {
func writeAllowlistedModelsList(c *gin.Context, platform string, modelIDs []string) {
if platform == service.PlatformOpenAI {
writeOpenAIModelsList(c, modelIDs)
return
@@ -1365,72 +1367,19 @@ func writeOpenAIModelsList(c *gin.Context, modelIDs []string) {
})
}
func customModelsListSource(platform string, availableModels, fallbackModels []string) []string {
if platform == service.PlatformAnthropic && len(availableModels) > 0 {
// modelListingSource 汇总模型列表过滤的候选来源:账号映射键(availableModels)
// 与平台默认列表(fallbackModels)。账号映射为空时回落默认列表;Anthropic
// 平台两者取并集,其余平台以账号映射键为准。
func modelListingSource(platform string, availableModels, fallbackModels []string) []string {
if len(availableModels) == 0 {
return fallbackModels
}
if platform == service.PlatformAnthropic {
return mergeModelIDs(availableModels, fallbackModels)
}
return availableModels
}
func filterModelsByCustomList(availableModels, fallbackModels, selectedModels []string) []string {
if len(selectedModels) == 0 {
return availableModels
}
source := availableModels
if len(source) == 0 {
source = fallbackModels
}
if len(source) == 0 {
return nil
}
allowed := make([]string, 0, len(source))
for _, model := range source {
model = strings.TrimSpace(model)
if model != "" {
allowed = append(allowed, model)
}
}
seen := make(map[string]struct{}, len(selectedModels))
filtered := make([]string, 0, len(selectedModels))
for _, model := range selectedModels {
model = strings.TrimSpace(model)
if model == "" {
continue
}
if !customModelsListAllowsModel(allowed, model) {
continue
}
if _, ok := seen[model]; ok {
continue
}
seen[model] = struct{}{}
filtered = append(filtered, model)
}
return filtered
}
func customModelsListAllowsModel(availablePatterns []string, model string) bool {
for _, pattern := range availablePatterns {
if pattern == model {
return true
}
if strings.HasSuffix(pattern, "*") && strings.HasPrefix(model, strings.TrimSuffix(pattern, "*")) {
return true
}
}
normalizedClaudeModel := claude.NormalizeModelID(strings.TrimSuffix(model, "-thinking"))
if normalizedClaudeModel != model {
for _, pattern := range availablePatterns {
if pattern == normalizedClaudeModel {
return true
}
}
}
return false
}
func defaultCodexModelIDsForPlatform(platform string) []string {
switch platform {
case service.PlatformDeepseek:
@@ -1504,10 +1453,21 @@ func mergeModelIDs(primary, secondary []string) []string {
// AntigravityModels 返回 Antigravity 支持的全部模型
// GET /antigravity/models
// 分组级模型白名单开启时按白名单过滤。
func (h *GatewayHandler) AntigravityModels(c *gin.Context) {
models := antigravity.DefaultModels()
if apiKey, ok := middleware2.GetAPIKeyFromContext(c); ok && apiKey != nil && apiKey.Group != nil && apiKey.Group.ModelAllowlistEnabled() {
filtered := make([]antigravity.ClaudeModel, 0, len(models))
for _, model := range models {
if apiKey.Group.ModelAllowlist.Allows(model.ID) {
filtered = append(filtered, model)
}
}
models = filtered
}
c.JSON(http.StatusOK, gin.H{
"object": "list",
"data": antigravity.DefaultModels(),
"data": models,
})
}
@@ -0,0 +1,245 @@
package handler
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"github.com/Wei-Shaw/sub2api/internal/service"
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
// 通配条目在 /v1/models 中展开为候选来源中所有匹配项,保持来源顺序。
func TestGatewayModels_ModelAllowlistWildcardExpandsAgainstSource(t *testing.T) {
gin.SetMode(gin.TestMode)
groupID := int64(31)
h := newGatewayModelsHandlerForTest(
&gatewayModelsAccountRepoStub{
byGroup: map[int64][]service.Account{
groupID: {
{
ID: 1,
Platform: service.PlatformOpenAI,
Credentials: map[string]any{
"model_mapping": map[string]any{
"gpt-5.4": "gpt-5.4",
"gpt-5.5-codex": "gpt-5.5-codex",
"gpt-5.5-mini": "gpt-5.5-mini",
"other-foo": "other-foo",
},
},
},
},
},
},
)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
Group: &service.Group{
ID: groupID,
Platform: service.PlatformOpenAI,
ModelAllowlist: service.GroupModelAllowlist{
Enabled: true,
Models: []string{"gpt-5.5-*", "gpt-5.4"},
},
},
})
h.Models(c)
require.Equal(t, http.StatusOK, rec.Code)
var got gatewayModelsResponseForTest
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
require.Equal(t, []string{"gpt-5.5-codex", "gpt-5.5-mini", "gpt-5.4"}, modelIDsForTest(got.Data))
}
// geminiAllowlistAccountRepoStub 在 gatewayModelsAccountRepoStub 之上补充
// Gemini 兼容层用到的按平台过滤查询(分组内无任何账号,触发 fallback 列表)。
type geminiAllowlistAccountRepoStub struct {
gatewayModelsAccountRepoStub
}
func (s *geminiAllowlistAccountRepoStub) ListSchedulableByGroupIDAndPlatforms(ctx context.Context, groupID int64, platforms []string) ([]service.Account, error) {
allowed := make(map[string]struct{}, len(platforms))
for _, platform := range platforms {
allowed[platform] = struct{}{}
}
accounts := s.byGroup[groupID]
filtered := make([]service.Account, 0, len(accounts))
for _, account := range accounts {
if _, ok := allowed[account.Platform]; ok {
filtered = append(filtered, account)
}
}
return filtered, nil
}
// Gemini 原生 /v1beta/models:白名单开启时过滤 fallback 列表的 models[].name。
func TestGeminiV1BetaListModels_FiltersFallbackByAllowlist(t *testing.T) {
gin.SetMode(gin.TestMode)
repo := &geminiAllowlistAccountRepoStub{gatewayModelsAccountRepoStub: gatewayModelsAccountRepoStub{
byGroup: map[int64][]service.Account{
// 只有 antigravity 账号:Gemini 选号失败后回落静态模型列表。
41: {{ID: 2, Platform: service.PlatformAntigravity, Status: service.StatusActive, Schedulable: true}},
},
}}
h := &GatewayHandler{
geminiCompatService: service.NewGeminiMessagesCompatService(repo, nil, nil, nil, nil, nil, nil, nil, nil),
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodGet, "/v1beta/models", nil)
geminiGroupID := int64(41)
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
GroupID: &geminiGroupID,
Group: &service.Group{
ID: geminiGroupID,
Platform: service.PlatformGemini,
ModelAllowlist: service.GroupModelAllowlist{
Enabled: true,
Models: []string{"gemini-2.5-pro", "gemini-3-*"},
},
},
})
h.GeminiV1BetaListModels(c)
require.Equal(t, http.StatusOK, rec.Code)
var got struct {
Models []struct {
Name string `json:"name"`
} `json:"models"`
}
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
names := make([]string, 0, len(got.Models))
for _, model := range got.Models {
names = append(names, model.Name)
}
// models/ 前缀的候选形式也应命中条目。
require.Contains(t, names, "models/gemini-2.5-pro")
require.Contains(t, names, "models/gemini-3-pro-preview")
require.NotContains(t, names, "models/gemini-2.5-flash")
require.NotContains(t, names, "models/gemini-2.0-flash")
}
// Antigravity /antigravity/models:白名单开启时按条目过滤静态列表。
func TestAntigravityModels_FiltersByAllowlist(t *testing.T) {
gin.SetMode(gin.TestMode)
h := &GatewayHandler{}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodGet, "/antigravity/models", nil)
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
Group: &service.Group{
ID: 51,
Platform: service.PlatformAntigravity,
ModelAllowlist: service.GroupModelAllowlist{
Enabled: true,
Models: []string{"gemini-2.5-flash"},
},
},
})
h.AntigravityModels(c)
require.Equal(t, http.StatusOK, rec.Code)
var got struct {
Object string `json:"object"`
Data []struct {
ID string `json:"id"`
} `json:"data"`
}
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
require.Equal(t, "list", got.Object)
// -thinking 后缀的宽容规则:gemini-2.5-flash-thinking 也命中 gemini-2.5-flash 条目。
require.Len(t, got.Data, 2)
require.Equal(t, "gemini-2.5-flash", got.Data[0].ID)
require.Equal(t, "gemini-2.5-flash-thinking", got.Data[1].ID)
}
func TestFilterUpstreamGeminiModelsBody(t *testing.T) {
allowlist := service.GroupModelAllowlist{Enabled: true, Models: []string{"gemini-2.5-pro"}}
t.Run("full hit returns original body with dropped=false", func(t *testing.T) {
body := []byte(`{"models":[{"name":"models/gemini-2.5-pro"},{"name":"models/gemini-2.5-pro"}]}`)
filtered, dropped, ok := filterUpstreamGeminiModelsBody(body, allowlist)
require.True(t, ok)
require.False(t, dropped, "full hit must keep the original response so headers pass through")
require.Equal(t, string(body), string(filtered))
})
t.Run("partial hit filters models and preserves envelope fields", func(t *testing.T) {
body := []byte(`{"nextPageToken":"tok","models":[{"name":"models/gemini-2.5-pro"},{"name":"models/gemini-2.5-flash"}],"other":"kept"}`)
filtered, dropped, ok := filterUpstreamGeminiModelsBody(body, allowlist)
require.True(t, ok)
require.True(t, dropped)
require.JSONEq(t, `{"nextPageToken":"tok","other":"kept","models":[{"name":"models/gemini-2.5-pro"}]}`, string(filtered))
})
t.Run("models prefix stripped candidate form matches entry", func(t *testing.T) {
body := []byte(`{"models":[{"name":"gemini-2.5-pro"}]}`)
filtered, dropped, ok := filterUpstreamGeminiModelsBody(body, allowlist)
require.True(t, ok)
require.False(t, dropped)
require.Equal(t, string(body), string(filtered))
})
t.Run("missing models field passes through", func(t *testing.T) {
body := []byte(`{"error":"odd upstream"}`)
_, dropped, ok := filterUpstreamGeminiModelsBody(body, allowlist)
require.True(t, ok)
require.False(t, dropped)
})
t.Run("invalid JSON signals passthrough", func(t *testing.T) {
_, dropped, ok := filterUpstreamGeminiModelsBody([]byte(`not-json`), allowlist)
require.False(t, ok)
require.False(t, dropped)
})
}
func TestFilterBatchImageModelsByAllowlist(t *testing.T) {
models := []service.BatchImagePublicModel{
{ID: "gemini-2.5-flash-image", Object: "image.batch.model", Provider: "gemini_api"},
{ID: "gemini-3-pro-image", Object: "image.batch.model", Provider: "vertex"},
{ID: "gpt-image-1", Object: "image.batch.model", Provider: "openai"},
}
t.Run("disabled allowlist keeps everything", func(t *testing.T) {
got := filterBatchImageModelsByAllowlist(models, service.GroupModelAllowlist{Enabled: false})
require.Equal(t, models, got)
})
t.Run("exact and wildcard entries filter by model id", func(t *testing.T) {
got := filterBatchImageModelsByAllowlist(models, service.GroupModelAllowlist{
Enabled: true,
Models: []string{"gemini-3-*", "gpt-image-1"},
})
require.Equal(t, []string{"gemini-3-pro-image", "gpt-image-1"}, []string{got[0].ID, got[1].ID})
require.Len(t, got, 2)
})
t.Run("no match yields empty list", func(t *testing.T) {
got := filterBatchImageModelsByAllowlist(models, service.GroupModelAllowlist{
Enabled: true,
Models: []string{"claude-*"},
})
require.Empty(t, got)
})
}
+12 -12
View File
@@ -267,7 +267,7 @@ func TestGatewayCodexModels_GeneratedManifestUsesFinalBodyETag(t *testing.T) {
require.Equal(t, etag, second.Header().Get("ETag"))
}
// Scenario: group models_list_config limits the generated Codex manifest.
// Scenario: group model_allowlist limits the generated Codex manifest.
func TestGatewayCodexModels_CustomModelsListFiltersCompositeManifest(t *testing.T) {
gin.SetMode(gin.TestMode)
const groupID int64 = 121
@@ -299,7 +299,7 @@ func TestGatewayCodexModels_CustomModelsListFiltersCompositeManifest(t *testing.
Group: &service.Group{
ID: groupID,
Platform: service.PlatformComposite,
ModelsListConfig: service.GroupModelsListConfig{
ModelAllowlist: service.GroupModelAllowlist{
Enabled: true,
Models: []string{"grok-4.6"},
},
@@ -572,7 +572,7 @@ func TestGatewayModels_CustomModelsListDisabledKeepsOriginalModels(t *testing.T)
Group: &service.Group{
ID: groupID,
Platform: service.PlatformOpenAI,
ModelsListConfig: service.GroupModelsListConfig{
ModelAllowlist: service.GroupModelAllowlist{
Enabled: false,
Models: []string{"gpt-5.5"},
},
@@ -619,7 +619,7 @@ func TestGatewayModels_CustomModelsListFiltersAndOrdersMappedModels(t *testing.T
Group: &service.Group{
ID: groupID,
Platform: service.PlatformOpenAI,
ModelsListConfig: service.GroupModelsListConfig{
ModelAllowlist: service.GroupModelAllowlist{
Enabled: true,
Models: []string{"gpt-5.5", "missing-model", "gpt-5.4"},
},
@@ -704,7 +704,7 @@ func TestGatewayModels_CompositeCustomModelsListFiltersAcrossConcretePlatforms(t
Group: &service.Group{
ID: groupID,
Platform: service.PlatformComposite,
ModelsListConfig: service.GroupModelsListConfig{
ModelAllowlist: service.GroupModelAllowlist{
Enabled: true,
Models: []string{"gemini-2.5-flash", "missing-model", "ag-custom-model", "gpt-5.5", "kimi-custom", "glm-custom", "deepseek-custom"},
},
@@ -918,7 +918,7 @@ func TestGatewayModels_CustomModelsListKeepsConcreteModelAllowedByWildcardMappin
Group: &service.Group{
ID: groupID,
Platform: service.PlatformAnthropic,
ModelsListConfig: service.GroupModelsListConfig{
ModelAllowlist: service.GroupModelAllowlist{
Enabled: true,
Models: []string{"claude-sonnet-4-6"},
},
@@ -969,7 +969,7 @@ func TestGatewayModels_AnthropicCustomModelsListIncludesOAuthClaudeAndMappedDeep
Group: &service.Group{
ID: groupID,
Platform: service.PlatformAnthropic,
ModelsListConfig: service.GroupModelsListConfig{
ModelAllowlist: service.GroupModelAllowlist{
Enabled: true,
Models: []string{"claude-fable-5", "claude-opus-4-8", "deepseek-v4-pro"},
},
@@ -1020,7 +1020,7 @@ func TestGatewayModels_AnthropicCustomModelsListDisabledKeepsMappedModelList(t *
Group: &service.Group{
ID: groupID,
Platform: service.PlatformAnthropic,
ModelsListConfig: service.GroupModelsListConfig{
ModelAllowlist: service.GroupModelAllowlist{
Enabled: false,
Models: []string{"claude-fable-5", "deepseek-v4-pro"},
},
@@ -1061,7 +1061,7 @@ func TestGatewayModels_AnthropicCustomModelsListIncludesOAuthClaudeWithoutMappin
Group: &service.Group{
ID: groupID,
Platform: service.PlatformAnthropic,
ModelsListConfig: service.GroupModelsListConfig{
ModelAllowlist: service.GroupModelAllowlist{
Enabled: true,
Models: []string{"claude-opus-4-6-thinking", "claude-sonnet-4-5"},
},
@@ -1106,7 +1106,7 @@ func TestGatewayModels_CustomModelsListCanReturnEmptyWhenSelectionsUnavailable(t
Group: &service.Group{
ID: groupID,
Platform: service.PlatformOpenAI,
ModelsListConfig: service.GroupModelsListConfig{
ModelAllowlist: service.GroupModelAllowlist{
Enabled: true,
Models: []string{"gpt-5.5"},
},
@@ -1143,7 +1143,7 @@ func TestGatewayModels_CustomModelsListFiltersDefaultFallbackModels(t *testing.T
Group: &service.Group{
ID: groupID,
Platform: service.PlatformOpenAI,
ModelsListConfig: service.GroupModelsListConfig{
ModelAllowlist: service.GroupModelAllowlist{
Enabled: true,
Models: []string{"gpt-5.5", "legacy-gpt-2024", "gpt-5.4"},
},
@@ -1180,7 +1180,7 @@ func TestGatewayModels_OpenAICustomModelsListKeepsOpenAIResponseShapeForDefaultF
Group: &service.Group{
ID: groupID,
Platform: service.PlatformOpenAI,
ModelsListConfig: service.GroupModelsListConfig{
ModelAllowlist: service.GroupModelAllowlist{
Enabled: true,
Models: []string{"gpt-5.5", "gpt-5.4"},
},
@@ -46,14 +46,34 @@ func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) {
return
}
// 强制 antigravity 模式:返回 antigravity 支持的模型列表
if forcePlatform == service.PlatformAntigravity {
c.JSON(http.StatusOK, antigravity.FallbackGeminiModelsList())
return
// 分组级模型白名单开启时过滤 models[].name(名字形如 models/xxx)。
filterGeminiModels := func(models []gemini.Model) []gemini.Model {
if apiKey.Group == nil || !apiKey.Group.ModelAllowlistEnabled() {
return models
}
filtered := make([]gemini.Model, 0, len(models))
for _, model := range models {
if apiKey.Group.ModelAllowlist.Allows(model.Name) {
filtered = append(filtered, model)
}
}
return filtered
}
if models, ok := customGeminiModelsList(apiKey.Group); ok {
c.JSON(http.StatusOK, models)
// 强制 antigravity 模式:返回 antigravity 支持的模型列表
if forcePlatform == service.PlatformAntigravity {
if apiKey.Group != nil && apiKey.Group.ModelAllowlistEnabled() {
agModels := antigravity.DefaultGeminiModels()
filtered := make([]antigravity.GeminiModel, 0, len(agModels))
for _, model := range agModels {
if apiKey.Group.ModelAllowlist.Allows(model.Name) {
filtered = append(filtered, model)
}
}
c.JSON(http.StatusOK, antigravity.GeminiModelsListResponse{Models: filtered})
return
}
c.JSON(http.StatusOK, antigravity.FallbackGeminiModelsList())
return
}
@@ -63,7 +83,7 @@ func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) {
hasAntigravity, _ := h.geminiCompatService.HasAntigravityAccounts(c.Request.Context(), apiKey.GroupID)
if hasAntigravity {
// antigravity 账户使用静态模型列表
c.JSON(http.StatusOK, gemini.FallbackModelsList())
c.JSON(http.StatusOK, gemini.ModelsListResponse{Models: filterGeminiModels(gemini.DefaultModels())})
return
}
markOpsRoutingCapacityLimitedIfNoAvailable(c, err)
@@ -77,21 +97,64 @@ func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) {
return
}
if shouldFallbackGeminiModels(res) {
c.JSON(http.StatusOK, gemini.FallbackModelsList())
c.JSON(http.StatusOK, gemini.ModelsListResponse{Models: filterGeminiModels(gemini.DefaultModels())})
return
}
if apiKey.Group != nil && apiKey.Group.ModelAllowlistEnabled() {
if filtered, dropped, ok := filterUpstreamGeminiModelsBody(res.Body, apiKey.Group.ModelAllowlist); ok && dropped {
// 只在确有条目被过滤时替换响应体;全命中或解析失败时保持原始响应,
// 统一经 writeUpstreamResponse 写出(保留全部上游响应头)。
res.Body = filtered
}
}
writeUpstreamResponse(c, res)
}
func customGeminiModelsList(group *service.Group) (gemini.ModelsListResponse, bool) {
if group == nil || !group.CustomModelsListEnabled() {
return gemini.ModelsListResponse{}, false
// filterUpstreamGeminiModelsBody 按白名单过滤上游 /v1beta/models 响应中的
// models[].name,其余信封字段(如 nextPageToken)原样保留。
// 返回值:filtered 为过滤后的响应体;dropped 表示是否有条目被移除(全命中时
// 为 false,调用方应保持原始响应以完整透传上游头);ok=false 表示解析失败,
// 调用方同样应透传原始响应。
func filterUpstreamGeminiModelsBody(body []byte, allowlist service.GroupModelAllowlist) (filtered []byte, dropped bool, ok bool) {
var envelope map[string]json.RawMessage
if err := json.Unmarshal(body, &envelope); err != nil {
return nil, false, false
}
models := make([]gemini.Model, 0, len(group.ModelsListConfig.Models))
for _, modelID := range group.ModelsListConfig.Models {
models = append(models, gemini.FallbackModel(modelID))
rawModels, hasModels := envelope["models"]
if !hasModels {
return body, false, true
}
return gemini.ModelsListResponse{Models: models}, true
type geminiModelName struct {
Name string `json:"name"`
}
var models []json.RawMessage
if err := json.Unmarshal(rawModels, &models); err != nil {
return nil, false, false
}
kept := make([]json.RawMessage, 0, len(models))
for _, raw := range models {
var model geminiModelName
if err := json.Unmarshal(raw, &model); err != nil {
return nil, false, false
}
if allowlist.Allows(model.Name) {
kept = append(kept, raw)
}
}
if len(kept) == len(models) {
// 全部命中时直接透传原始响应体。
return body, false, true
}
mergedModels, err := json.Marshal(kept)
if err != nil {
return nil, false, false
}
envelope["models"] = mergedModels
merged, err := json.Marshal(envelope)
if err != nil {
return nil, false, false
}
return merged, true, true
}
// GeminiV1BetaGetModel proxies:
@@ -9,40 +9,21 @@ import (
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
"github.com/Wei-Shaw/sub2api/internal/pkg/gemini"
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestGeminiV1BetaListModels_CustomGroupListUsesNativeResponse(t *testing.T) {
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodGet, "/v1beta/models", nil)
c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
Group: &service.Group{
Platform: service.PlatformGemini,
ModelsListConfig: service.GroupModelsListConfig{
Enabled: true,
Models: []string{"gemini-2.5-pro", "models/gemini-custom"},
},
},
})
(&GatewayHandler{}).GeminiV1BetaListModels(c)
require.Equal(t, http.StatusOK, rec.Code)
var got gemini.ModelsListResponse
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
require.Equal(t, []gemini.Model{
gemini.FallbackModel("gemini-2.5-pro"),
gemini.FallbackModel("models/gemini-custom"),
}, got.Models)
func TestGeminiV1BetaListModels_AllowlistFiltersNativeResponse(t *testing.T) {
body := []byte(`{"models":[{"name":"models/gemini-2.5-pro"},{"name":"models/gemini-2.5-flash"}],"nextPageToken":"next"}`)
filtered, dropped, ok := filterUpstreamGeminiModelsBody(body, service.GroupModelAllowlist{Enabled: true, Models: []string{"gemini-2.5-pro"}})
require.True(t, ok)
require.True(t, dropped)
require.JSONEq(t, `{"models":[{"name":"models/gemini-2.5-pro"}],"nextPageToken":"next"}`, string(filtered))
}
func TestGeminiV1BetaListModels_ForcedAntigravityIgnoresCustomGroupList(t *testing.T) {
func TestGeminiV1BetaListModels_ForcedAntigravityAppliesAllowlist(t *testing.T) {
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
@@ -50,7 +31,7 @@ func TestGeminiV1BetaListModels_ForcedAntigravityIgnoresCustomGroupList(t *testi
c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
Group: &service.Group{
Platform: service.PlatformGemini,
ModelsListConfig: service.GroupModelsListConfig{
ModelAllowlist: service.GroupModelAllowlist{
Enabled: true,
Models: []string{"gemini-custom"},
},
@@ -63,19 +44,15 @@ func TestGeminiV1BetaListModels_ForcedAntigravityIgnoresCustomGroupList(t *testi
require.Equal(t, http.StatusOK, rec.Code)
var got antigravity.GeminiModelsListResponse
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
require.Equal(t, antigravity.FallbackGeminiModelsList(), got)
require.Empty(t, got.Models)
}
func TestCustomGeminiModelsList_DisabledKeepsExistingFlow(t *testing.T) {
group := &service.Group{
ModelsListConfig: service.GroupModelsListConfig{
Enabled: false,
Models: []string{"gemini-2.5-pro"},
},
}
_, ok := customGeminiModelsList(group)
require.False(t, ok)
func TestGeminiModelAllowlist_DisabledPreservesNativeResponse(t *testing.T) {
body := []byte(`{"models":[{"name":"models/gemini-2.5-pro"}]}`)
filtered, dropped, ok := filterUpstreamGeminiModelsBody(body, service.GroupModelAllowlist{Enabled: false, Models: []string{"other"}})
require.True(t, ok)
require.False(t, dropped)
require.Equal(t, body, filtered)
}
// TestGeminiV1BetaHandler_PlatformRoutingInvariant 文档化并验证 Handler 层的平台路由逻辑不变量
@@ -163,7 +163,7 @@ func TestCodexModelsAppliesLocalFiltersBeforeClientETag(t *testing.T) {
group := &service.Group{
ID: groupID,
Platform: service.PlatformOpenAI,
ModelsListConfig: service.GroupModelsListConfig{
ModelAllowlist: service.GroupModelAllowlist{
Enabled: true,
Models: []string{"codex-auto-review", "gpt-5.6"},
},
@@ -181,7 +181,7 @@ func TestCodexModelsAppliesLocalFiltersBeforeClientETag(t *testing.T) {
t.Fatal("first response did not include an ETag")
}
group.ModelsListConfig.Enabled = false
group.ModelAllowlist.Enabled = false
second := performCodexModelsRequestForGroup(t, handler, group, oldETag)
if second.Code != http.StatusOK {
t.Fatalf("second status: got %d, want %d; body=%s", second.Code, http.StatusOK, second.Body.String())
@@ -232,7 +232,7 @@ func TestCodexModelsAPIKeyCacheDoesNotLeakGroupFilters(t *testing.T) {
groupA := &service.Group{
ID: 91,
Platform: service.PlatformOpenAI,
ModelsListConfig: service.GroupModelsListConfig{
ModelAllowlist: service.GroupModelAllowlist{
Enabled: true,
Models: []string{"model-a"},
},
@@ -240,7 +240,7 @@ func TestCodexModelsAPIKeyCacheDoesNotLeakGroupFilters(t *testing.T) {
groupB := &service.Group{
ID: 92,
Platform: service.PlatformOpenAI,
ModelsListConfig: service.GroupModelsListConfig{
ModelAllowlist: service.GroupModelAllowlist{
Enabled: true,
Models: []string{"model-b"},
},
@@ -912,7 +912,7 @@ func TestCodexModelsPinnedAccountsStillApplyCustomModelsListFilter(t *testing.T)
group := &service.Group{
ID: 84,
Platform: service.PlatformOpenAI,
ModelsListConfig: service.GroupModelsListConfig{
ModelAllowlist: service.GroupModelAllowlist{
Enabled: true,
Models: []string{"model-b"},
},
@@ -20,6 +20,7 @@ import (
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/Wei-Shaw/sub2api/internal/pkg/requestmodel"
"github.com/Wei-Shaw/sub2api/internal/securityaudit"
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
@@ -2263,6 +2264,16 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, "model is required in first response.create payload")
return
}
// 分组级模型白名单:首帧校验客户端模型,不通过则关闭连接并标记运维原因。
// 必须在 ensureCompositeTargetPlatform(合成路由改写)之前执行。
// 与 HTTP 准入一致:帧内重复 model 键/大小写变体可能被上游按末值绑定,
// 全部候选值逐一校验,任一未命中即拒绝。
if blocked := blockedModelAllowlistCandidate(apiKey.Group, requestmodel.FromBodyCandidates("", "application/json", firstMessage)); blocked != "" {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalModelConfiguration)
middleware2.MarkIngressRejected(c, middleware2.IngressRejectModelNotAllowed)
closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, fmt.Sprintf("Model %q is not available for this group", blocked))
return
}
ensureCompositeTargetPlatform(c, apiKey, reqModel)
ctx = c.Request.Context()
if apiKey.Group != nil && apiKey.Group.Platform == service.PlatformComposite {
@@ -2680,6 +2691,17 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
if model == "" {
model = reqModel
}
// 分组级模型白名单:后续 turn 同样校验客户端模型(省略 model 时
// 沿用会话实际生效模型,含 session.update 轮换后的模型),不通过
// 则关闭整条连接,与推理强度 deny 一致。实际生效模型始终参与校验;
// 帧内重复 model 键/大小写变体/嵌套 session.model 额外逐一校验,
// 防止候选集非空时掩盖被轮换掉的禁用模型。
candidates := append([]string{model}, requestmodel.FromBodyCandidates("", "application/json", payload)...)
if blocked := blockedModelAllowlistCandidate(apiKey.Group, candidates); blocked != "" {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalModelConfiguration)
middleware2.MarkIngressRejected(c, middleware2.IngressRejectModelNotAllowed)
return service.NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, fmt.Sprintf("Model %q is not available for this group", blocked), nil)
}
if decision := h.checkSecurityAuditStage(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIResponses, model, payload, "subsequent_turn"); decision != nil && !decision.AllowNextStage {
writeSecurityAuditWSError(ctx, wsConn, decision)
return service.NewOpenAIWSClientCloseError(securityAuditWSCloseStatus(decision), securityAuditWSCloseReason(decision), nil)
@@ -3581,6 +3603,21 @@ func isOpenAIWSUpgradeRequest(r *http.Request) bool {
return strings.Contains(strings.ToLower(strings.TrimSpace(r.Header.Get("Connection"))), "upgrade")
}
// blockedModelAllowlistCandidate 对全部候选模型逐一校验分组白名单,返回第一个
// 未命中的值(全部命中或白名单未开启返回空串)。WS 帧与 HTTP 请求体共用该
// 规则:重复 model 键/大小写变体可能被上游按末值绑定,任一未命中即拒绝。
func blockedModelAllowlistCandidate(group *service.Group, candidates []string) string {
if group == nil || !group.ModelAllowlistEnabled() {
return ""
}
for _, candidate := range candidates {
if !group.ModelAllowlist.Allows(candidate) {
return candidate
}
}
return ""
}
func closeOpenAIClientWS(conn *coderws.Conn, status coderws.StatusCode, reason string) {
if conn == nil {
return
@@ -1882,7 +1882,10 @@ func newOpenAIWSHandlerTestServer(t *testing.T, h *OpenAIGatewayHandler, subject
}
type openAIResponsesWSUsageLogCase struct {
firstPayload string
firstPayload string
// midPayload 在首个 turn 完成后发送(如 session.update),上游桩会为它
// 回一个 response.completed,客户端按普通事件读取。
midPayload string
secondPayload string
userAgent *string
ingressMode string
@@ -1890,6 +1893,12 @@ type openAIResponsesWSUsageLogCase struct {
billingModelSource string
accountModelMapping map[string]any
afterFirstUpstreamRequest func(channelSvc *service.ChannelService) error
// group 覆盖 apiKey.Group(分组级模型白名单测试用);nil 保持原有无分组行为。
group *service.Group
// firstFrameCloseExpected:首帧即被拒(连接被 1008 关闭),不期待任何响应帧。
firstFrameCloseExpected bool
// secondTurnCloseExpected:第二个 turn 被拒(连接被 1008 关闭)。
secondTurnCloseExpected bool
}
type openAIResponsesWSUsageLogResult struct {
@@ -2809,8 +2818,11 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
gin.SetMode(gin.TestMode)
turnCount := 1
if strings.TrimSpace(tc.midPayload) != "" {
turnCount++
}
if strings.TrimSpace(tc.secondPayload) != "" {
turnCount = 2
turnCount++
}
upstreamPayloadCh := make(chan []byte, turnCount)
upstreamErrCh := make(chan error, 1)
@@ -2963,6 +2975,9 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
GroupID: &groupID,
User: &service.User{ID: 1701, Status: service.StatusActive},
}
if tc.group != nil {
apiKey.Group = tc.group
}
router := gin.New()
router.Use(func(c *gin.Context) {
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
@@ -2994,6 +3009,19 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
cancelWrite()
require.NoError(t, err)
if tc.firstFrameCloseExpected {
readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
_, _, readErr := clientConn.Read(readCtx)
cancelRead()
require.Error(t, readErr, "first frame should have been rejected with a close")
var closeErr coderws.CloseError
require.ErrorAs(t, readErr, &closeErr)
require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code)
require.Contains(t, closeErr.Reason, "not available for this group")
_ = clientConn.CloseNow()
return openAIResponsesWSUsageLogResult{}
}
clientEvents := make([][]byte, 0, turnCount)
readCompleted := func() {
readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
@@ -3004,11 +3032,30 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
clientEvents = append(clientEvents, append([]byte(nil), event...))
}
readCompleted()
if turnCount == 2 {
if tc.midPayload != "" {
writeCtx, cancelWrite = context.WithTimeout(context.Background(), 3*time.Second)
err = clientConn.Write(writeCtx, coderws.MessageText, []byte(tc.midPayload))
cancelWrite()
require.NoError(t, err)
readCompleted()
}
if strings.TrimSpace(tc.secondPayload) != "" && (turnCount >= 2) {
writeCtx, cancelWrite = context.WithTimeout(context.Background(), 3*time.Second)
err = clientConn.Write(writeCtx, coderws.MessageText, []byte(tc.secondPayload))
cancelWrite()
require.NoError(t, err)
if tc.secondTurnCloseExpected {
readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
_, _, readErr := clientConn.Read(readCtx)
cancelRead()
require.Error(t, readErr, "second turn should have been rejected with a close")
var closeErr coderws.CloseError
require.ErrorAs(t, readErr, &closeErr)
require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code)
require.Contains(t, closeErr.Reason, "not available for this group")
_ = clientConn.CloseNow()
return openAIResponsesWSUsageLogResult{}
}
readCompleted()
}
_ = clientConn.Close(coderws.StatusNormalClosure, "done")
@@ -0,0 +1,179 @@
package handler
import (
"testing"
"github.com/Wei-Shaw/sub2api/internal/service"
)
func wsAllowlistGroup(enabled bool, models ...string) *service.Group {
return &service.Group{
ID: 4201,
Platform: service.PlatformOpenAI,
ModelAllowlist: service.GroupModelAllowlist{
Enabled: enabled,
Models: models,
},
}
}
// 首帧拒绝:模型不在白名单时连接被 1008 关闭(passthrough relay 模式)。
func TestOpenAIResponsesWebSocket_FirstFrameModelNotAllowedCloses_Passthrough(t *testing.T) {
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
firstPayload: `{"type":"response.create","model":"gpt-4.1","stream":false}`,
group: wsAllowlistGroup(true, "gpt-5.4"),
ingressMode: service.OpenAIWSIngressModePassthrough,
firstFrameCloseExpected: true,
})
}
// 首帧拒绝:原生 ingress 模式。
func TestOpenAIResponsesWebSocket_FirstFrameModelNotAllowedCloses_NativeIngress(t *testing.T) {
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
firstPayload: `{"type":"response.create","model":"gpt-4.1","stream":false}`,
group: wsAllowlistGroup(true, "gpt-5.4"),
ingressMode: service.OpenAIWSIngressModeDedicated,
firstFrameCloseExpected: true,
})
}
// 首帧命中白名单:连接正常建立并收到 response.completed(passthrough)。
func TestOpenAIResponsesWebSocket_FirstFrameAllowlistedModelProceeds(t *testing.T) {
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
firstPayload: `{"type":"response.create","model":"gpt-5.4","stream":false}`,
group: wsAllowlistGroup(true, "gpt-5.4"),
})
if len(got.clientEvents) != 1 {
t.Fatalf("expected one completed event, got %d", len(got.clientEvents))
}
}
// 后续 turn 切换到白名单外的模型:连接被 1008 关闭(passthrough relay 模式)。
func TestOpenAIResponsesWebSocket_SubsequentTurnModelNotAllowedCloses_Passthrough(t *testing.T) {
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
firstPayload: `{"type":"response.create","model":"gpt-5.4","stream":false}`,
secondPayload: `{"type":"response.create","model":"gpt-4.1","stream":false}`,
group: wsAllowlistGroup(true, "gpt-5.4"),
ingressMode: service.OpenAIWSIngressModePassthrough,
secondTurnCloseExpected: true,
})
}
// 后续 turn 切换到白名单外的模型:连接被 1008 关闭(原生 ingress 模式)。
func TestOpenAIResponsesWebSocket_SubsequentTurnModelNotAllowedCloses_NativeIngress(t *testing.T) {
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
firstPayload: `{"type":"response.create","model":"gpt-5.4","stream":false}`,
secondPayload: `{"type":"response.create","model":"gpt-4.1","stream":false}`,
group: wsAllowlistGroup(true, "gpt-5.4"),
ingressMode: service.OpenAIWSIngressModeDedicated,
secondTurnCloseExpected: true,
})
}
// 后续 turn 省略 model:沿用会话初始模型,白名单校验通过。
func TestOpenAIResponsesWebSocket_SubsequentTurnOmittedModelUsesSessionModel(t *testing.T) {
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
firstPayload: `{"type":"response.create","model":"gpt-5.4","stream":false}`,
secondPayload: `{"type":"response.create","stream":false}`,
group: wsAllowlistGroup(true, "gpt-5.4"),
})
if len(got.clientEvents) != 2 {
t.Fatalf("expected two completed events, got %d", len(got.clientEvents))
}
if len(got.upstreamPayloads) != 2 {
t.Fatalf("expected two upstream frames, got %d", len(got.upstreamPayloads))
}
}
// 白名单关闭:不受任何影响。
func TestOpenAIResponsesWebSocket_DisabledAllowlistDoesNotInterfere(t *testing.T) {
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
firstPayload: `{"type":"response.create","model":"gpt-4.1","stream":false}`,
secondPayload: `{"type":"response.create","model":"gpt-5.4","stream":false}`,
group: wsAllowlistGroup(false, "gpt-5.4"),
})
if len(got.clientEvents) != 2 {
t.Fatalf("expected two completed events, got %d", len(got.clientEvents))
}
}
// 首帧含重复 model 键(首个在白名单、末个不在):passthrough 会原样转发两个字段,
// 上游按末值绑定时得到禁用模型,必须在入口拒绝。
func TestOpenAIResponsesWebSocket_FirstFrameDuplicateModelKeysRejected(t *testing.T) {
for _, mode := range []string{service.OpenAIWSIngressModePassthrough, service.OpenAIWSIngressModeDedicated} {
t.Run(mode, func(t *testing.T) {
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
firstPayload: `{"type":"response.create","model":"gpt-5.4","model":"gpt-4.1","stream":false}`,
group: wsAllowlistGroup(true, "gpt-5.4"),
ingressMode: mode,
firstFrameCloseExpected: true,
})
})
}
}
// 首帧混用大小写变体键({"model":"allowed","Model":"blocked"}):绑定系上游
// 大小写不敏感且末值生效会绑定 blocked,必须在入口拒绝。仅含变体键的帧则会在
// 「model is required」处更早被关闭。
func TestOpenAIResponsesWebSocket_FirstFrameCaseVariantModelKeyRejected(t *testing.T) {
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
firstPayload: `{"type":"response.create","model":"gpt-5.4","Model":"gpt-4.1","stream":false}`,
group: wsAllowlistGroup(true, "gpt-5.4"),
ingressMode: service.OpenAIWSIngressModePassthrough,
firstFrameCloseExpected: true,
})
}
// 后续 turn 含重复 model 键(首个在白名单、末个不在):拒绝并关闭连接。
func TestOpenAIResponsesWebSocket_SubsequentTurnDuplicateModelKeysRejected(t *testing.T) {
for _, mode := range []string{service.OpenAIWSIngressModePassthrough, service.OpenAIWSIngressModeDedicated} {
t.Run(mode, func(t *testing.T) {
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
firstPayload: `{"type":"response.create","model":"gpt-5.4","stream":false}`,
secondPayload: `{"type":"response.create","model":"gpt-5.4","model":"gpt-4.1","stream":false}`,
group: wsAllowlistGroup(true, "gpt-5.4"),
ingressMode: mode,
secondTurnCloseExpected: true,
})
})
}
}
// 重复但同值的 model 键不误伤,连接正常完成两个 turn。
func TestOpenAIResponsesWebSocket_DuplicateIdenticalModelKeysAllowed(t *testing.T) {
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
firstPayload: `{"type":"response.create","model":"gpt-5.4","model":"gpt-5.4","stream":false}`,
secondPayload: `{"type":"response.create","model":"gpt-5.4","stream":false}`,
group: wsAllowlistGroup(true, "gpt-5.4"),
})
if len(got.clientEvents) != 2 {
t.Fatalf("expected two completed events, got %d", len(got.clientEvents))
}
}
// session.update 轮换绕过:首帧用白名单内模型建立会话,session.update 把会话
// 模型改为白名单外的模型,随后 response.create 携带嵌套 session.model(白名单
// 内)让帧内候选非空。实际生效模型(轮换后的会话模型)必须始终参与校验。
func TestOpenAIResponsesWebSocket_SessionUpdateRotationBypassRejected(t *testing.T) {
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
firstPayload: `{"type":"response.create","model":"gpt-5.4","stream":false}`,
midPayload: `{"type":"session.update","session":{"model":"gpt-4.1"}}`,
secondPayload: `{"type":"response.create","session":{"model":"gpt-5.4"},"stream":false}`,
group: wsAllowlistGroup(true, "gpt-5.4"),
ingressMode: service.OpenAIWSIngressModePassthrough,
secondTurnCloseExpected: true,
})
}
// 轮换后的会话模型本身在白名单内时,省略 model 的后续 turn 正常放行。
func TestOpenAIResponsesWebSocket_SessionUpdateToAllowedModelStillWorks(t *testing.T) {
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
firstPayload: `{"type":"response.create","model":"gpt-5.4","stream":false}`,
midPayload: `{"type":"session.update","session":{"model":"gpt-5.4"}}`,
secondPayload: `{"type":"response.create","stream":false}`,
group: wsAllowlistGroup(true, "gpt-5.4"),
})
if len(got.clientEvents) < 2 {
t.Fatalf("expected at least two completed events, got %d", len(got.clientEvents))
}
}
+9 -1
View File
@@ -2189,9 +2189,17 @@ func classifyOpsErrorLog(c *gin.Context, errType, message, code string, status i
routingCapacityLimited := isOpsRoutingCapacityLimited(c)
clientBusinessLimited := service.HasOpsClientBusinessLimited(c)
localModelConfiguration := clientBusinessLimited && service.OpsClientBusinessLimitedReason(c) == service.OpsClientBusinessLimitedReasonLocalModelConfiguration
// 分组模型白名单的入口拒绝发生在调度之前(本地面请求策略,业务限流原因与
// 账号模型映射复用 local_model_configuration):保持 classifyOpsPhase 的
// 自然分类(not_found/invalid_request → request,owner=client),不落到
// routing;上游归因仍由 suppressOpsUpstreamAttributionForLocalModelConfiguration 清空。
ingressModelNotAllowed := false
if reason, rejected := middleware2.GetIngressRejectReason(c); rejected && reason == middleware2.IngressRejectModelNotAllowed {
ingressModelNotAllowed = true
}
upstreamError := hasOpsUpstreamErrorContext(c)
accountAuthFailure := hasOpsAccountAuthFailure(c)
if localModelConfiguration {
if localModelConfiguration && !ingressModelNotAllowed {
phase = "routing"
} else if accountAuthFailure && !routingCapacityLimited {
phase = "account_auth"
@@ -1954,3 +1954,45 @@ func TestGetOpsAPIKeyPrefersPrimaryContextKey(t *testing.T) {
require.NotNil(t, got)
require.Equal(t, int64(1), got.ID, "已鉴权请求应优先使用正式 api key")
}
// 分组模型白名单入口拒绝:业务限流原因复用 local_model_configuration,但
// 携带 ingress 拒绝原因 model_not_allowed,阶段应保持自然分类(request /
// client),不落到 routing。
func TestClassifyOpsIngressModelNotAllowedKeepsRequestPhase(t *testing.T) {
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalModelConfiguration)
middleware2.MarkIngressRejected(c, middleware2.IngressRejectModelNotAllowed)
phase, isBusinessLimited, errorOwner, errorSource := classifyOpsErrorLog(
c,
"not_found_error",
`Model "gpt-4.1" is not available for this group`,
"model_not_found",
http.StatusNotFound,
)
require.Equal(t, "request", phase)
require.True(t, isBusinessLimited)
require.Equal(t, "client", errorOwner)
require.Equal(t, "client_request", errorSource)
}
// 调度阶段的账号模型映射拒绝(无 ingress 标记)仍归类为 routing。
func TestClassifyOpsLocalModelConfigurationWithoutIngressMarkStaysRouting(t *testing.T) {
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalModelConfiguration)
phase, isBusinessLimited, errorOwner, _ := classifyOpsErrorLog(
c,
"model_not_found",
"Model \"gpt-missing\" is not supported by any configured account in this group",
"",
http.StatusNotFound,
)
require.Equal(t, "routing", phase)
require.True(t, isBusinessLimited)
require.Equal(t, "platform", errorOwner)
}
+41
View File
@@ -22,13 +22,54 @@ const (
maxDecompressedBodySize = 64 << 20
)
// PrereadBody 回填已读取完成的请求体:作为 io.ReadCloser 可被再次顺序消费
// (multipart 流式解析),同时暴露 Bytes() 让 ReadRequestBodyWithPrealloc
// 直接返回原始切片,避免二次分配与复制。
//
// 注意:ReadRequestBodyWithPrealloc 对 PrereadBody 的快速路径不检查内部
// reader 是否已被(部分)消费——包装的字节完整且不可变,即使 reader 已被
// 流式消费过,Bytes() 也始终返回完整请求体。
type PrereadBody struct {
body []byte
reader *bytes.Reader
}
// NewPrereadBody 包装一段已读取的请求体。
func NewPrereadBody(body []byte) *PrereadBody {
return &PrereadBody{body: body, reader: bytes.NewReader(body)}
}
// Read 实现 io.Reader(转发给内部 bytes.Reader)。
func (p *PrereadBody) Read(b []byte) (int, error) {
if p == nil {
return 0, io.EOF
}
return p.reader.Read(b)
}
// Close 实现 io.Closer;请求体已在内存中,无需释放资源。
func (p *PrereadBody) Close() error { return nil }
// Bytes 返回完整的原始请求体切片。
func (p *PrereadBody) Bytes() []byte {
if p == nil {
return nil
}
return p.body
}
// ReadRequestBodyWithPrealloc reads request body with preallocated buffer based
// on content length, transparently decoding any Content-Encoding the upstream
// client used to compress the body (zstd, gzip, deflate).
// 已由 PrereadBody 回填的请求体直接返回其完整切片(零拷贝),不检查内部
// reader 是否已被消费——见 PrereadBody 的文档说明。
func ReadRequestBodyWithPrealloc(req *http.Request) ([]byte, error) {
if req == nil || req.Body == nil {
return nil, nil
}
if preread, ok := req.Body.(*PrereadBody); ok {
return preread.Bytes(), nil
}
capHint := requestBodyReadInitCap
if req.ContentLength > 0 {
@@ -0,0 +1,97 @@
package httputil
import (
"bytes"
"io"
"net/http"
"strings"
"testing"
)
func TestPrereadBodyRoundtrip(t *testing.T) {
body := []byte(`{"model":"gpt-5.4"}`)
preread := NewPrereadBody(body)
got, err := io.ReadAll(preread)
if err != nil {
t.Fatalf("read preread body: %v", err)
}
if !bytes.Equal(got, body) {
t.Fatalf("body mismatch: got %q want %q", got, body)
}
if err := preread.Close(); err != nil {
t.Fatalf("close: %v", err)
}
if !bytes.Equal(preread.Bytes(), body) {
t.Fatalf("Bytes() mismatch after close: %q", preread.Bytes())
}
}
func TestPrereadBodyReadableTwiceFromFreshWrapper(t *testing.T) {
// ResetRequestBody 每次都回填新的 PrereadBody;同一实例只被顺序消费一次。
body := []byte("payload")
first := NewPrereadBody(body)
got1, _ := io.ReadAll(first)
if !bytes.Equal(got1, body) {
t.Fatalf("first read mismatch: %q", got1)
}
second := NewPrereadBody(first.Bytes())
got2, _ := io.ReadAll(second)
if !bytes.Equal(got2, body) {
t.Fatalf("second read mismatch: %q", got2)
}
}
func TestReadRequestBodyWithPreallocReturnsPrereadBodySliceWithoutCopy(t *testing.T) {
body := []byte(`{"model":"claude-sonnet-4.5"}`)
req, err := http.NewRequest(http.MethodPost, "/v1/messages", NewPrereadBody(body))
if err != nil {
t.Fatalf("new request: %v", err)
}
req.Header.Set("Content-Length", "0") // 故意与实际不符,验证不依赖 ContentLength
got, err := ReadRequestBodyWithPrealloc(req)
if err != nil {
t.Fatalf("ReadRequestBodyWithPrealloc: %v", err)
}
if &got[0] != &body[0] {
t.Fatalf("expected zero-copy return of the same slice")
}
if !bytes.Equal(got, body) {
t.Fatalf("body mismatch: got %q want %q", got, body)
}
}
func TestReadRequestBodyWithPreallocHandlesNilPrereadBody(t *testing.T) {
req, err := http.NewRequest(http.MethodPost, "/v1/messages", nil)
if err != nil {
t.Fatalf("new request: %v", err)
}
got, err := ReadRequestBodyWithPrealloc(req)
if err != nil || got != nil {
t.Fatalf("expected nil body, got (%#v, %v)", got, err)
}
}
func TestPrereadBodyPartialConsumeThenPreallocReturnsFullBody(t *testing.T) {
body := []byte(strings.Repeat("x", 64))
preread := NewPrereadBody(body)
head := make([]byte, 16)
if _, err := io.ReadFull(preread, head); err != nil {
t.Fatalf("partial read: %v", err)
}
req, err := http.NewRequest(http.MethodPost, "/v1/responses", preread)
if err != nil {
t.Fatalf("new request: %v", err)
}
got, err := ReadRequestBodyWithPrealloc(req)
if err != nil {
t.Fatalf("ReadRequestBodyWithPrealloc: %v", err)
}
if !bytes.Equal(got, body) {
t.Fatalf("expected full body after partial consume, got %d bytes", len(got))
}
}
@@ -0,0 +1,220 @@
// Package requestmodel 统一从网关请求中提取客户端请求的模型名。
//
// 提取顺序与位置(JSON `model` / `session.model`、multipart `model` / `session`)
// 与合成路由中间件、分组模型白名单中间件共用,保证改写与准入看到同一个模型。
//
// 网关里同时存在三类下游解析器,同一请求体可能被解析出不同模型值:
// - gjson 系(大小写敏感、重复键取第一个);
// - encoding/json 绑定系(大小写不敏感、重复键取最后一个);
// - multipart 表单系(字段可重复,不同 handler 取首个或末个)。
//
// 因此准入侧的 FromBodyCandidates 返回「任一下游解析器可能绑定到的全部模型值」,
// 调用方必须逐一校验;FromBodyForRoute 返回首个候选值,仅供合成路由分发使用。
package requestmodel
import (
"bytes"
"errors"
"io"
"mime"
"mime/multipart"
"net/http"
"strconv"
"strings"
"github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
"github.com/tidwall/gjson"
)
// liveRequestRoutes 是模型位于 session.model(而非顶层 model)的入口,
// 与 Live handler 的解析保持一致。
var liveRequestRoutes = map[string]bool{
"/v1/live": true,
"/live": true,
"/backend-api/codex/realtime/calls": true,
}
// IsLiveRequestRoute 报告该路由模板的 handler 是否只读取 session.model。
func IsLiveRequestRoute(routePath string) bool {
return liveRequestRoutes[strings.TrimSpace(routePath)]
}
// FromBodyCandidates 按入口返回所有可被下游解析器绑定的模型值:
// Live 入口取 session 候选(为空回落顶层候选),其余入口取顶层候选
// (为空回落 session 候选)。候选包含重复键、大小写变体的全部出现。
func FromBodyCandidates(routePath, contentType string, body []byte) []string {
var topLevel, session []string
if isMultipartContentType(contentType) {
// 声明为 multipart 的请求体只做 multipart 解析——gjson 的部分扫描
// 会在 multipart 字节流里误匹配 session JSON 内的 model 字段。
models, sessions := multipartModelCandidates(contentType, body)
topLevel = models
session = sessionModelCandidates(sessions)
} else {
topLevel = jsonModelCandidates(body)
session = jsonSessionModelCandidates(body)
}
if IsLiveRequestRoute(routePath) {
if len(session) > 0 {
return session
}
return topLevel
}
if len(topLevel) > 0 {
return topLevel
}
return session
}
// FromBody 从 JSON 或 multipart 请求体提取客户端请求的模型名(首个候选);
// 提取不到返回空串。顶层 `model` 优先,其次 `session.model`(Live 入口之外
// 的地方兜底)。准入校验请改用 FromBodyCandidates。
func FromBody(contentType string, body []byte) string {
return FromBodyForRoute("", contentType, body)
}
// FromBodyForRoute 按入口选择与 handler 一致的模型提取规则,返回首个候选值,
// 仅供合成路由分发使用;准入校验必须用 FromBodyCandidates 的全部候选。
func FromBodyForRoute(routePath, contentType string, body []byte) string {
candidates := FromBodyCandidates(routePath, contentType, body)
if len(candidates) == 0 {
return ""
}
return candidates[0]
}
// FromJSON 从 JSON 请求体提取模型名:顶层 `model` 优先,其次 `session.model`(Live)。
func FromJSON(body []byte) string {
model, _ := JSONModelPath(body)
return model
}
// JSONModelPath 返回 JSON 请求体中的模型名与其 gjson 路径(用于合成路由改写定位)。
// 只匹配精确大小写的键;大小写变体由 FromBodyCandidates 的候选集覆盖。
func JSONModelPath(body []byte) (string, string) {
return jsonModelWithPaths(body, []string{"model", "session.model"})
}
// JSONModelPathForRoute 按入口返回合成路由改写应写入的模型路径:Live 入口
// session.model 优先,其余入口顶层 model 优先。
func JSONModelPathForRoute(routePath string, body []byte) (string, string) {
if IsLiveRequestRoute(routePath) {
return jsonModelWithPaths(body, []string{"session.model", "model"})
}
return jsonModelWithPaths(body, []string{"model", "session.model"})
}
func jsonModelWithPaths(body []byte, paths []string) (string, string) {
for _, path := range paths {
model := gjson.GetBytes(body, path)
if model.Type != gjson.String {
continue
}
if value := strings.TrimSpace(model.String()); value != "" {
return value, path
}
}
return "", ""
}
// jsonModelCandidates 返回顶层所有键名大小写不敏感等于 "model" 的字符串值,
// 按出现顺序保留重复(覆盖 gjson 首 matches 与 encoding/json 末值绑定两类解析器)。
func jsonModelCandidates(body []byte) []string {
var out []string
gjson.ParseBytes(body).ForEach(func(key, value gjson.Result) bool {
if value.Type == gjson.String && strings.EqualFold(key.String(), "model") {
if trimmed := strings.TrimSpace(value.String()); trimmed != "" {
out = append(out, trimmed)
}
}
return true
})
return out
}
// jsonSessionModelCandidates 返回每个顶层 "session"(大小写不敏感)对象内
// 所有键名大小写不敏感等于 "model" 的字符串值。
func jsonSessionModelCandidates(body []byte) []string {
var out []string
gjson.ParseBytes(body).ForEach(func(key, value gjson.Result) bool {
if !strings.EqualFold(key.String(), "session") {
return true
}
value.ForEach(func(sessionKey, sessionValue gjson.Result) bool {
if sessionValue.Type == gjson.String && strings.EqualFold(sessionKey.String(), "model") {
if trimmed := strings.TrimSpace(sessionValue.String()); trimmed != "" {
out = append(out, trimmed)
}
}
return true
})
return true
})
return out
}
// sessionModelCandidates 对若干 session 字段原始值(JSON 文本)逐个提取模型候选。
func sessionModelCandidates(rawSessions []string) []string {
var out []string
for _, raw := range rawSessions {
out = append(out, jsonModelCandidates([]byte(raw))...)
}
return out
}
// multipartModelCandidates 返回 multipart 表单中全部 "model" 字段值与全部
// "session" 字段原始值,按出现顺序保留重复,覆盖取首个或末个字段的两类 handler。
func multipartModelCandidates(contentType string, body []byte) (models, sessions []string) {
mediaType, params, err := mime.ParseMediaType(strings.TrimSpace(contentType))
if err != nil || !strings.EqualFold(mediaType, "multipart/form-data") {
return nil, nil
}
boundary := strings.TrimSpace(params["boundary"])
if boundary == "" {
return nil, nil
}
reader := multipart.NewReader(bytes.NewReader(body), boundary)
for {
part, err := reader.NextPart()
if errors.Is(err, io.EOF) {
return models, sessions
}
if err != nil {
return models, sessions
}
fieldName := strings.TrimSpace(part.FormName())
if part.FileName() != "" || (fieldName != "model" && fieldName != "session") {
continue
}
data, err := io.ReadAll(part)
if err != nil {
continue
}
if value := strings.TrimSpace(string(data)); value != "" {
if fieldName == "model" {
models = append(models, value)
} else {
sessions = append(sessions, value)
}
}
}
}
// isMultipartContentType 判断请求是否声明为 multipart/form-data。
func isMultipartContentType(contentType string) bool {
mediaType, _, err := mime.ParseMediaType(strings.TrimSpace(contentType))
return err == nil && strings.EqualFold(mediaType, "multipart/form-data")
}
// ResetRequestBody 把已读取(可能已被改写)的请求体回填到请求上,
// 供后续 handler 重新读取。使用 httputil.PrereadBody 回填,后续经
// ReadRequestBodyWithPrealloc 读取时零拷贝。
func ResetRequestBody(req *http.Request, body []byte) {
if req == nil {
return
}
req.Body = httputil.NewPrereadBody(body)
req.ContentLength = int64(len(body))
req.Header.Set("Content-Length", strconv.Itoa(len(body)))
}
@@ -0,0 +1,257 @@
package requestmodel
import (
"bytes"
"mime/multipart"
"net/http"
"net/http/httptest"
"strconv"
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
)
func TestFromBodyJSONModel(t *testing.T) {
if got := FromBody("application/json", []byte(`{"model":"gpt-5.4"}`)); got != "gpt-5.4" {
t.Fatalf("got %q", got)
}
}
func TestFromBodyJSONSessionModel(t *testing.T) {
if got := FromBody("application/json", []byte(`{"session":{"model":"gpt-live"},"sdp":"v=0"}`)); got != "gpt-live" {
t.Fatalf("got %q", got)
}
}
func TestFromBodyJSONModelWinsOverSession(t *testing.T) {
body := []byte(`{"model":"top-level","session":{"model":"nested"}}`)
if got := FromBody("application/json", body); got != "top-level" {
t.Fatalf("got %q", got)
}
}
func TestFromBodyJSONBlankModelFallsThrough(t *testing.T) {
if got := FromBody("application/json", []byte(`{"model":" "}`)); got != "" {
t.Fatalf("got %q", got)
}
}
func TestFromBodyMultipartModelField(t *testing.T) {
var body bytes.Buffer
writer := multipart.NewWriter(&body)
_ = writer.WriteField("prompt", "draw")
_ = writer.WriteField("model", "gpt-image-1")
_ = writer.Close()
if got := FromBody(writer.FormDataContentType(), body.Bytes()); got != "gpt-image-1" {
t.Fatalf("got %q", got)
}
}
func TestFromBodyMultipartSessionField(t *testing.T) {
var body bytes.Buffer
writer := multipart.NewWriter(&body)
_ = writer.WriteField("sdp", "v=0")
_ = writer.WriteField("session", `{"model":"live-alias"}`)
_ = writer.WriteField("media", "file-part") // 字段名不匹配,应被跳过
_ = writer.Close()
if got := FromBody(writer.FormDataContentType(), body.Bytes()); got != "live-alias" {
t.Fatalf("got %q", got)
}
}
func TestFromBodySkipsFilePartsBeforeModelField(t *testing.T) {
var body bytes.Buffer
writer := multipart.NewWriter(&body)
part, err := writer.CreateFormFile("file", "input.png")
if err != nil {
t.Fatalf("create form file: %v", err)
}
_, _ = part.Write(bytes.Repeat([]byte{0}, 4096))
_ = writer.WriteField("model", "gpt-image-1")
_ = writer.Close()
if got := FromBody(writer.FormDataContentType(), body.Bytes()); got != "gpt-image-1" {
t.Fatalf("got %q", got)
}
}
func TestFromBodyNoModelReturnsEmpty(t *testing.T) {
if got := FromBody("application/json", []byte(`{"input":"hi"}`)); got != "" {
t.Fatalf("got %q", got)
}
if got := FromBody("text/plain", []byte("plain")); got != "" {
t.Fatalf("got %q", got)
}
}
func TestResetRequestBodyAllowsRereadWithoutCopy(t *testing.T) {
body := []byte(`{"model":"gpt-5.4","input":"hi"}`)
req := httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
ResetRequestBody(req, body)
got, err := httputil.ReadRequestBodyWithPrealloc(req)
if err != nil {
t.Fatalf("reread: %v", err)
}
if &got[0] != &body[0] {
t.Fatalf("expected zero-copy reread of the same slice")
}
if req.ContentLength != int64(len(body)) {
t.Fatalf("content length mismatch: %d", req.ContentLength)
}
if req.Header.Get("Content-Length") != strconv.Itoa(len(body)) {
t.Fatalf("content-length header mismatch: %q", req.Header.Get("Content-Length"))
}
}
func TestResetRequestBodyNilRequestIsNoop(t *testing.T) {
ResetRequestBody(nil, []byte("x"))
}
func TestFromBodyForRouteLivePrefersSessionModel(t *testing.T) {
body := []byte(`{"model":"top-level","session":{"model":"session-model"},"sdp":"v=0"}`)
if got := FromBodyForRoute("/v1/live", "application/json", body); got != "session-model" {
t.Fatalf("live route should prefer session.model, got %q", got)
}
if got := FromBodyForRoute("/backend-api/codex/realtime/calls", "application/json", body); got != "session-model" {
t.Fatalf("codex realtime route should prefer session.model, got %q", got)
}
if got := FromBodyForRoute("/live", "application/json", body); got != "session-model" {
t.Fatalf("root live alias should prefer session.model, got %q", got)
}
}
func TestFromBodyForRouteNonLivePrefersTopLevelModel(t *testing.T) {
body := []byte(`{"model":"top-level","session":{"model":"session-model"}}`)
if got := FromBodyForRoute("/v1/messages", "application/json", body); got != "top-level" {
t.Fatalf("non-live route should prefer top-level model, got %q", got)
}
if got := FromBodyForRoute("/v1/responses", "application/json", body); got != "top-level" {
t.Fatalf("responses route should prefer top-level model, got %q", got)
}
}
func TestFromBodyForRouteFallsBackWhenPreferredFieldMissing(t *testing.T) {
if got := FromBodyForRoute("/v1/live", "application/json", []byte(`{"model":"top-level"}`)); got != "top-level" {
t.Fatalf("live without session.model should fall back to model, got %q", got)
}
if got := FromBodyForRoute("/v1/messages", "application/json", []byte(`{"session":{"model":"session-model"}}`)); got != "session-model" {
t.Fatalf("non-live without model should fall back to session.model, got %q", got)
}
}
func TestFromBodyForRouteMultipartSessionFirstForLive(t *testing.T) {
var body bytes.Buffer
writer := multipart.NewWriter(&body)
_ = writer.WriteField("model", "top-level")
_ = writer.WriteField("session", `{"model":"session-model"}`)
_ = writer.Close()
contentType := writer.FormDataContentType()
if got := FromBodyForRoute("/v1/live", contentType, body.Bytes()); got != "session-model" {
t.Fatalf("live multipart should prefer session field, got %q", got)
}
if got := FromBodyForRoute("/v1/messages", contentType, body.Bytes()); got != "top-level" {
t.Fatalf("non-live multipart should prefer model field, got %q", got)
}
}
func TestIsLiveRequestRoute(t *testing.T) {
for _, path := range []string{"/v1/live", "/live", "/backend-api/codex/realtime/calls"} {
if !IsLiveRequestRoute(path) {
t.Fatalf("%s should be a live route", path)
}
}
for _, path := range []string{"/v1/messages", "/v1/responses", "/realtime", ""} {
if IsLiveRequestRoute(path) {
t.Fatalf("%s should not be a live route", path)
}
}
}
// gjson 的部分扫描会在 multipart 字节流里误匹配 session JSON 内的 model;
// 声明为 multipart 的请求体必须走 multipart 解析。
func TestFromBodyMultipartContentTypeSkipsJSONScan(t *testing.T) {
var body bytes.Buffer
writer := multipart.NewWriter(&body)
_ = writer.WriteField("model", "top-level")
_ = writer.WriteField("session", `{"model":"session-model"}`)
_ = writer.Close()
if got := FromBody(writer.FormDataContentType(), body.Bytes()); got != "top-level" {
t.Fatalf("multipart body should use the model field, got %q", got)
}
if got := FromBodyForRoute("/v1/live", writer.FormDataContentType(), body.Bytes()); got != "session-model" {
t.Fatalf("live multipart should use the session field, got %q", got)
}
}
func TestFromBodyCandidatesIncludesCaseVariantsAndDuplicates(t *testing.T) {
// encoding/json 绑定对键名大小写不敏感、重复键取末值;gjson 取首个。
// 候选集必须包含全部出现。
got := FromBodyCandidates("", "application/json", []byte(`{"Model":"caps-model"}`))
if len(got) != 1 || got[0] != "caps-model" {
t.Fatalf("case-variant key must be a candidate, got %#v", got)
}
got = FromBodyCandidates("", "application/json", []byte(`{"model":"first","model":"last"}`))
if len(got) != 2 || got[0] != "first" || got[1] != "last" {
t.Fatalf("duplicate keys must both be candidates, got %#v", got)
}
}
func TestFromBodyCandidatesSessionVariants(t *testing.T) {
body := []byte(`{"Session":{"Model":"s-caps"},"session":{"model":"s-lower"},"model":"top"}`)
// 非 Live 入口:顶层候选。
got := FromBodyCandidates("", "application/json", body)
if len(got) != 1 || got[0] != "top" {
t.Fatalf("non-live should use top-level candidates, got %#v", got)
}
// Live 入口:session 候选(含大小写变体与重复)。
got = FromBodyCandidates("/v1/live", "application/json", body)
if len(got) != 2 || got[0] != "s-caps" || got[1] != "s-lower" {
t.Fatalf("live should collect all session candidates, got %#v", got)
}
}
func TestFromBodyCandidatesMultipartDuplicateFields(t *testing.T) {
var body bytes.Buffer
writer := multipart.NewWriter(&body)
_ = writer.WriteField("model", "gpt-image-1")
_ = writer.WriteField("model", "gpt-image-1.5")
_ = writer.WriteField("session", `{"model":"live-a"}`)
_ = writer.WriteField("session", `{"model":"live-b"}`)
_ = writer.Close()
models, sessions := multipartModelCandidates(writer.FormDataContentType(), body.Bytes())
if len(models) != 2 || models[0] != "gpt-image-1" || models[1] != "gpt-image-1.5" {
t.Fatalf("duplicate model fields must all be candidates, got %#v", models)
}
if len(sessions) != 2 {
t.Fatalf("duplicate session fields must all be captured, got %#v", sessions)
}
got := FromBodyCandidates("", writer.FormDataContentType(), body.Bytes())
if len(got) != 2 || got[0] != "gpt-image-1" || got[1] != "gpt-image-1.5" {
t.Fatalf("non-live multipart candidates mismatch, got %#v", got)
}
got = FromBodyCandidates("/v1/live", writer.FormDataContentType(), body.Bytes())
if len(got) != 2 || got[0] != "live-a" || got[1] != "live-b" {
t.Fatalf("live multipart session candidates mismatch, got %#v", got)
}
}
func TestFromBodyForRouteReturnsFirstCandidate(t *testing.T) {
got := FromBodyForRoute("", "application/json", []byte(`{"model":"first","model":"last"}`))
if got != "first" {
t.Fatalf("dispatch helper keeps gjson-first semantics, got %q", got)
}
}
+2 -2
View File
@@ -214,7 +214,7 @@ func (r *apiKeyRepository) GetByKeyForAuth(ctx context.Context, key string) (*se
group.FieldFreeOpenaiFast,
group.FieldDefaultMappedModel,
group.FieldMessagesDispatchModelConfig,
group.FieldModelsListConfig,
group.FieldModelAllowlist,
group.FieldCodexModelsManifestConfig,
group.FieldRpmLimit,
group.FieldMaxReasoningEffort,
@@ -1017,7 +1017,7 @@ func groupEntityToService(g *dbent.Group) *service.Group {
RequirePrivacySet: g.RequirePrivacySet,
DefaultMappedModel: g.DefaultMappedModel,
MessagesDispatchModelConfig: g.MessagesDispatchModelConfig,
ModelsListConfig: g.ModelsListConfig,
ModelAllowlist: service.GroupModelAllowlistFromDomain(g.ModelAllowlist),
CodexModelsManifestConfig: g.CodexModelsManifestConfig,
RPMLimit: g.RpmLimit,
MaxReasoningEffort: g.MaxReasoningEffort,
@@ -0,0 +1,50 @@
//go:build integration
package repository
import (
"context"
"fmt"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/stretchr/testify/require"
)
func TestGetByKeyForAuthCarriesGroupModelAllowlist(t *testing.T) {
ctx := context.Background()
suffix := time.Now().UnixNano()
group := mustCreateGroup(t, integrationEntClient, &service.Group{
Name: fmt.Sprintf("model-allowlist-proj-group-%d", suffix), Platform: service.PlatformOpenAI,
RateMultiplier: 1,
ModelAllowlist: service.GroupModelAllowlist{
Enabled: true,
Models: []string{"gpt-5.4", "gpt-5.5-*"},
},
})
user := mustCreateUser(t, integrationEntClient, &service.User{
Email: fmt.Sprintf("model-allowlist-proj-%d@example.com", suffix), Concurrency: 5,
})
groupID := group.ID
keyValue := fmt.Sprintf("sk-model-allowlist-proj-%d", suffix)
apiKeyRepo := NewAPIKeyRepository(integrationEntClient, integrationDB)
key := &service.APIKey{UserID: user.ID, GroupID: &groupID, Key: keyValue, Name: "model-allowlist-proj", Status: service.StatusActive}
require.NoError(t, apiKeyRepo.Create(ctx, key))
t.Cleanup(func() {
_, err := integrationDB.ExecContext(ctx, "DELETE FROM auth_cache_invalidation_outbox WHERE cache_key = encode(sha256(convert_to($1, 'UTF8')), 'hex')", keyValue)
require.NoError(t, err)
_, err = integrationDB.ExecContext(ctx, "DELETE FROM api_keys WHERE id = $1", key.ID)
require.NoError(t, err)
_, err = integrationDB.ExecContext(ctx, "DELETE FROM users WHERE id = $1", user.ID)
require.NoError(t, err)
_, err = integrationDB.ExecContext(ctx, "DELETE FROM groups WHERE id = $1", group.ID)
require.NoError(t, err)
})
got, err := apiKeyRepo.GetByKeyForAuth(ctx, keyValue)
require.NoError(t, err)
require.NotNil(t, got.Group)
require.True(t, got.Group.ModelAllowlist.Enabled, "model_allowlist 必须进入认证投影(投影漏列会让分组级准入静默失效)")
require.Equal(t, []string{"gpt-5.4", "gpt-5.5-*"}, got.Group.ModelAllowlist.Models)
}
+2 -2
View File
@@ -109,7 +109,7 @@ func createGroupRecord(ctx context.Context, client *dbent.Client, groupIn *servi
SetRequirePrivacySet(groupIn.RequirePrivacySet).
SetDefaultMappedModel(groupIn.DefaultMappedModel).
SetMessagesDispatchModelConfig(groupIn.MessagesDispatchModelConfig).
SetModelsListConfig(groupIn.ModelsListConfig).
SetModelAllowlist(service.DomainGroupModelAllowlist(groupIn.ModelAllowlist)).
SetCodexModelsManifestConfig(groupIn.CodexModelsManifestConfig).
SetRpmLimit(groupIn.RPMLimit).
SetMaxReasoningEffort(groupIn.MaxReasoningEffort).
@@ -289,7 +289,7 @@ func (r *groupRepository) Update(ctx context.Context, groupIn *service.Group) er
SetRequirePrivacySet(groupIn.RequirePrivacySet).
SetDefaultMappedModel(groupIn.DefaultMappedModel).
SetMessagesDispatchModelConfig(groupIn.MessagesDispatchModelConfig).
SetModelsListConfig(groupIn.ModelsListConfig).
SetModelAllowlist(service.DomainGroupModelAllowlist(groupIn.ModelAllowlist)).
SetCodexModelsManifestConfig(groupIn.CodexModelsManifestConfig).
SetRpmLimit(groupIn.RPMLimit).
SetMaxReasoningEffort(groupIn.MaxReasoningEffort).
@@ -0,0 +1,165 @@
package middleware
import (
"errors"
"fmt"
"net/http"
"strings"
"github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
"github.com/Wei-Shaw/sub2api/internal/pkg/requestmodel"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// GroupModelAllowlist 是分组级模型白名单准入中间件。
//
// 挂载位置:每条网关链的 apiKeyAuth 之后、compositeTarget 之前——
// 保证校验发生在合成路由改写与调度之前,且只看客户端书写的公开模型名。
//
// 行为:
// - 快速路径:未绑定分组或白名单未开启时直接放行,不读请求体。
// - Responses WebSocket 入口跳过(首帧与后续 turn 由 ResponsesWebSocket 逐帧
// 校验);Grok Realtime 的升级请求模型固定在查询参数里,仍走中间件校验,
// 其他路由伪造 Upgrade 头不得绕过校验。
// - GET/DELETE:从路由参数(model / modelAction)或查询参数(/realtime 的
// model)提取模型;提取不到放行,由 handler 决定是否报「model is required」。
// - 带请求体的方法:读体(PrereadBody 回填,下游零拷贝)、提取 JSON
// `model`/`session.model` 或 multipart `model`/`session` 后回填请求体。
// - 拒绝:按入口协议格式返回 404,并标记运维业务限流原因
// local_model_configuration 与 ingress 拒绝原因 model_not_allowed。
func GroupModelAllowlist() gin.HandlerFunc {
return func(c *gin.Context) {
apiKey, ok := GetAPIKeyFromContext(c)
if !ok || apiKey == nil || apiKey.Group == nil || !apiKey.Group.ModelAllowlistEnabled() {
c.Next()
return
}
allowlist := apiKey.Group.ModelAllowlist
if c.Request == nil {
c.Next()
return
}
if isResponsesWebSocketRoute(c) {
// Responses WS 长连接由 ResponsesWebSocket 校验首帧与每个 response.create。
c.Next()
return
}
// models 收集该请求全部可被下游解析器绑定到的模型值(路径参数/查询参数
// 单值;请求体候选集含重复键与大小写变体),逐一校验,任一未命中即拒绝。
var models []string
if model := groupModelAllowlistModelFromParams(c); model != "" {
models = []string{model}
} else {
switch c.Request.Method {
case http.MethodPost, http.MethodPut, http.MethodPatch:
candidates, done := groupModelAllowlistModelsFromBody(c)
if !done {
// 请求体读取失败(如超限 413)已写出响应。
return
}
models = candidates
}
if len(models) == 0 {
// Grok Realtime 升级请求把模型固定在查询参数里。
if model := strings.TrimSpace(c.Query("model")); model != "" {
models = []string{model}
}
}
}
blocked := ""
for _, candidate := range models {
if !allowlist.Allows(candidate) {
blocked = candidate
break
}
}
if blocked == "" {
c.Next()
return
}
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalModelConfiguration)
MarkIngressRejected(c, IngressRejectModelNotAllowed)
groupModelAllowlistErrorWriter(c)(c, http.StatusNotFound, fmt.Sprintf("Model %q is not available for this group", blocked))
c.Abort()
}
}
// isResponsesWebSocketRoute 判断当前请求是否命中 OpenAI Responses WebSocket
// 入口。只有这三条路由的模型在升级后的帧里,交给 handler 逐帧校验;其他路由
// 即使携带 WebSocket Upgrade 头也必须经过本中间件校验。
func isResponsesWebSocketRoute(c *gin.Context) bool {
if c.Request == nil || c.Request.Method != http.MethodGet {
return false
}
switch c.FullPath() {
case "/v1/responses", "/responses", "/backend-api/codex/responses":
return true
}
return false
}
// groupModelAllowlistModelsFromBody 读取请求体并提取客户端模型名,随后把请求体
// 回填(PrereadBody),保证后续 handler 零拷贝重读。读取失败按现有合成中间件
// 的方式返回 400/413(返回 false 表示已写出响应并 Abort)。
//
// 下游同时存在 gjson(首个、大小写敏感)、encoding/json 绑定(末值、大小写
// 不敏感)与 multipart 表单(首/末字段)三类解析器,这里返回「任一解析器可能
// 绑定到的全部模型值」,调用方必须逐一校验,任一未命中即拒绝。
func groupModelAllowlistModelsFromBody(c *gin.Context) ([]string, bool) {
body, err := httputil.ReadRequestBodyWithPrealloc(c.Request)
if err != nil {
status := http.StatusBadRequest
message := "Failed to read request body"
var maxErr *http.MaxBytesError
if errors.As(err, &maxErr) {
status = http.StatusRequestEntityTooLarge
message = "Request body is too large"
}
c.JSON(status, gin.H{"error": gin.H{"type": "invalid_request_error", "message": message}})
c.Abort()
return nil, false
}
requestmodel.ResetRequestBody(c.Request, body)
return requestmodel.FromBodyCandidates(c.FullPath(), c.GetHeader("Content-Type"), body), true
}
// groupModelAllowlistModelFromParams 从路由参数提取模型名:Gemini 原生 URL 的
// `model` / `modelAction`(去掉 `:action` 后缀)。提取不到返回空串,
// 由调用方决定是否继续读请求体或查询参数。
func groupModelAllowlistModelFromParams(c *gin.Context) string {
if model := strings.TrimSpace(c.Param("model")); model != "" {
return model
}
if modelAction := strings.TrimSpace(c.Param("modelAction")); modelAction != "" {
modelAction = strings.TrimPrefix(modelAction, "/")
if idx := strings.LastIndex(modelAction, ":"); idx >= 0 {
return strings.TrimSpace(modelAction[:idx])
}
return modelAction
}
return ""
}
// groupModelAllowlistErrorWriter 按入口协议选择错误格式:
// Gemini 原生(/v1beta、/antigravity/v1beta)用 Google 格式;
// Messages 入口(含根路径别名与 /antigravity/v1)用 Anthropic 格式;
// 其余(OpenAI 兼容入口)用 OpenAI 格式。
func groupModelAllowlistErrorWriter(c *gin.Context) GatewayErrorWriter {
path := ""
if c.Request != nil && c.Request.URL != nil {
path = c.Request.URL.Path
}
switch {
case strings.HasPrefix(path, "/v1beta") || strings.HasPrefix(path, "/antigravity/v1beta"):
return GoogleErrorWriter
case strings.Contains(path, "/messages"):
return AnthropicErrorWriter
default:
return OpenAIErrorWriter
}
}
@@ -0,0 +1,501 @@
package middleware
import (
"bytes"
"io"
"mime/multipart"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
func newGroupModelAllowlistTestRouter(apiKey *service.APIKey, pathPrefix string) (*gin.Engine, *[]string) {
gin.SetMode(gin.TestMode)
router := gin.New()
var calls []string
router.Use(func(c *gin.Context) {
if apiKey != nil {
c.Set(string(ContextKeyAPIKey), apiKey)
}
c.Next()
})
router.Use(GroupModelAllowlist())
register := func(method, path string) {
router.Handle(method, path, func(c *gin.Context) {
calls = append(calls, "handler:"+path)
c.Status(http.StatusOK)
})
}
register(http.MethodPost, pathPrefix+"/messages")
register(http.MethodPost, pathPrefix+"/responses")
register(http.MethodPost, pathPrefix+"/chat/completions")
register(http.MethodPost, pathPrefix+"/embeddings")
register(http.MethodGet, pathPrefix+"/models/:model")
register(http.MethodPost, pathPrefix+"/models/*modelAction")
register(http.MethodGet, pathPrefix+"/realtime")
register(http.MethodPost, pathPrefix+"/images/edits")
register(http.MethodPost, pathPrefix+"/live")
return router, &calls
}
func allowlistAPIKey(enabled bool, models ...string) *service.APIKey {
return &service.APIKey{
Group: &service.Group{
Platform: service.PlatformAnthropic,
ModelAllowlist: service.GroupModelAllowlist{
Enabled: enabled,
Models: models,
},
},
}
}
func doJSON(t *testing.T, router *gin.Engine, method, path, body string) *httptest.ResponseRecorder {
t.Helper()
req := httptest.NewRequest(method, path, strings.NewReader(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
return w
}
// readTrackingBody 记录请求体是否被读取,用于断言快速路径零读取。
type readTrackingBody struct {
io.Reader
read bool
}
func (b *readTrackingBody) Read(p []byte) (int, error) {
b.read = true
return b.Reader.Read(p)
}
func (b *readTrackingBody) Close() error { return nil }
func TestGroupModelAllowlistDisabledDoesNotReadBody(t *testing.T) {
router, calls := newGroupModelAllowlistTestRouter(allowlistAPIKey(false, "claude-sonnet-4.5"), "/v1")
body := &readTrackingBody{Reader: strings.NewReader(`{"model":"claude-opus-4.6"}`)}
req := httptest.NewRequest(http.MethodPost, "/v1/messages", body)
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
}
if len(*calls) != 1 {
t.Fatalf("expected handler to run once, got %v", *calls)
}
if body.read {
t.Fatal("allowlist disabled: middleware must not read the request body")
}
}
// 白名单开启但请求不携带模型时同样不应读体之外产生副作用:读体是必要开销,
// 这里验证读取后请求体被完整回填(handler 可零拷贝重读)。
func TestGroupModelAllowlistEnabledModelFreeRequestRestoresBody(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
router.Use(func(c *gin.Context) {
c.Set(string(ContextKeyAPIKey), allowlistAPIKey(true, "claude-sonnet-4.5"))
c.Next()
})
router.Use(GroupModelAllowlist())
router.POST("/v1/messages", func(c *gin.Context) {
body, err := httputil.ReadRequestBodyWithPrealloc(c.Request)
if err != nil {
t.Errorf("handler reread body: %v", err)
c.Status(http.StatusInternalServerError)
return
}
if string(body) != `{"input":"no model field"}` {
t.Errorf("handler saw wrong body: %s", body)
}
c.Status(http.StatusOK)
})
w := doJSON(t, router, http.MethodPost, "/v1/messages", `{"input":"no model field"}`)
if w.Code != http.StatusOK {
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
}
}
// 413 传递:请求体超过 MaxBytesReader 上限时按现有合成中间件方式返回 413。
func TestGroupModelAllowlistRequestBodyTooLargePassesThrough413(t *testing.T) {
router, calls := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "claude-sonnet-4.5"), "/v1")
bigBody := `{"model":"claude-opus-4.6","padding":"` + strings.Repeat("x", 512) + `"}`
req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(bigBody))
req.Header.Set("Content-Type", "application/json")
req.Body = http.MaxBytesReader(httptest.NewRecorder(), req.Body, 64)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusRequestEntityTooLarge {
t.Fatalf("expected 413, got %d: %s", w.Code, w.Body.String())
}
if !strings.Contains(w.Body.String(), "Request body is too large") {
t.Fatalf("expected too-large message, got %s", w.Body.String())
}
if len(*calls) != 0 {
t.Fatalf("expected handler not to run, got %v", *calls)
}
}
func TestGroupModelAllowlistJSONBodyAllowed(t *testing.T) {
router, calls := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "claude-sonnet-4.5"), "/v1")
w := doJSON(t, router, http.MethodPost, "/v1/messages", `{"model":"claude-sonnet-4.5"}`)
if w.Code != http.StatusOK {
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
}
if len(*calls) != 1 {
t.Fatalf("expected handler to run once, got %v", *calls)
}
}
func TestGroupModelAllowlistJSONBodyDeniedOpenAIFormat(t *testing.T) {
router, calls := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "claude-sonnet-4.5"), "/v1")
w := doJSON(t, router, http.MethodPost, "/v1/responses", `{"model":"gpt-5.4"}`)
if w.Code != http.StatusNotFound {
t.Fatalf("expected 404, got %d: %s", w.Code, w.Body.String())
}
if !strings.Contains(w.Body.String(), `"model_not_found"`) {
t.Fatalf("expected OpenAI error format, got %s", w.Body.String())
}
if !strings.Contains(w.Body.String(), `gpt-5.4\" is not available for this group`) {
t.Fatalf("expected model in message, got %s", w.Body.String())
}
if len(*calls) != 0 {
t.Fatalf("expected handler not to run, got %v", *calls)
}
}
func TestGroupModelAllowlistMessagesDeniedAnthropicFormat(t *testing.T) {
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "claude-sonnet-4.5"), "/v1")
w := doJSON(t, router, http.MethodPost, "/v1/messages", `{"model":"claude-opus-4.6"}`)
if w.Code != http.StatusNotFound {
t.Fatalf("expected 404, got %d: %s", w.Code, w.Body.String())
}
if !strings.Contains(w.Body.String(), `"not_found_error"`) {
t.Fatalf("expected Anthropic not_found_error format, got %s", w.Body.String())
}
}
func TestGroupModelAllowlistGeminiDeniedGoogleFormat(t *testing.T) {
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "gemini-2.5-pro"), "/v1beta")
w := doJSON(t, router, http.MethodPost, "/v1beta/models/gemini-2.5-flash:generateContent", `{}`)
if w.Code != http.StatusNotFound {
t.Fatalf("expected 404, got %d: %s", w.Code, w.Body.String())
}
if !strings.Contains(w.Body.String(), `"NOT_FOUND"`) {
t.Fatalf("expected Google NOT_FOUND format, got %s", w.Body.String())
}
}
func TestGroupModelAllowlistGeminiWildcardAndPrefixForms(t *testing.T) {
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "gemini-2.5-*"), "/v1beta")
w := doJSON(t, router, http.MethodPost, "/v1beta/models/models/gemini-2.5-flash:generateContent", `{}`)
if w.Code != http.StatusOK {
t.Fatalf("models/ prefixed model should match wildcard entry, got %d: %s", w.Code, w.Body.String())
}
}
func TestGroupModelAllowlistGeminiGetModelRouteParam(t *testing.T) {
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "gemini-2.5-pro"), "/v1beta")
req := httptest.NewRequest(http.MethodGet, "/v1beta/models/gemini-2.5-flash", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusNotFound {
t.Fatalf("expected 404 for unlisted GET model, got %d", w.Code)
}
req = httptest.NewRequest(http.MethodGet, "/v1beta/models/gemini-2.5-pro", nil)
w = httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected 200 for listed GET model, got %d", w.Code)
}
}
func TestGroupModelAllowlistRealtimeQueryParameter(t *testing.T) {
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "grok-4.6"), "/v1")
req := httptest.NewRequest(http.MethodGet, "/v1/realtime?model=grok-4.6", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected 200 for listed realtime model, got %d: %s", w.Code, w.Body.String())
}
req = httptest.NewRequest(http.MethodGet, "/v1/realtime?model=grok-4.20", nil)
w = httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusNotFound {
t.Fatalf("expected 404 for unlisted realtime model, got %d", w.Code)
}
}
func TestGroupModelAllowlistMultipartModelField(t *testing.T) {
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "gpt-image-1"), "/v1")
var body bytes.Buffer
writer := multipart.NewWriter(&body)
_ = writer.WriteField("model", "gpt-image-1")
_ = writer.WriteField("prompt", "cat")
_ = writer.Close()
req := httptest.NewRequest(http.MethodPost, "/v1/images/edits", &body)
req.Header.Set("Content-Type", writer.FormDataContentType())
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
}
}
func TestGroupModelAllowlistMultipartDenied(t *testing.T) {
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "gpt-image-1"), "/v1")
var body bytes.Buffer
writer := multipart.NewWriter(&body)
_ = writer.WriteField("model", "dall-e-3")
_ = writer.WriteField("prompt", "cat")
_ = writer.Close()
req := httptest.NewRequest(http.MethodPost, "/v1/images/edits", &body)
req.Header.Set("Content-Type", writer.FormDataContentType())
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusNotFound {
t.Fatalf("expected 404, got %d: %s", w.Code, w.Body.String())
}
}
func TestGroupModelAllowlistNoModelPassesThrough(t *testing.T) {
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "claude-sonnet-4.5"), "/v1")
w := doJSON(t, router, http.MethodPost, "/v1/messages", `{"input":"no model field"}`)
if w.Code != http.StatusOK {
t.Fatalf("expected 200 when no model extractable, got %d: %s", w.Code, w.Body.String())
}
}
func TestGroupModelAllowlistBodyRestoredForHandler(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
router.Use(func(c *gin.Context) {
c.Set(string(ContextKeyAPIKey), allowlistAPIKey(true, "claude-sonnet-4.5"))
c.Next()
})
router.Use(GroupModelAllowlist())
router.POST("/v1/messages", func(c *gin.Context) {
body, err := io.ReadAll(c.Request.Body)
if err != nil {
t.Errorf("handler reread body: %v", err)
c.Status(http.StatusInternalServerError)
return
}
if !strings.Contains(string(body), `"claude-sonnet-4.5"`) {
t.Errorf("handler saw wrong body: %s", body)
}
c.Status(http.StatusOK)
})
w := doJSON(t, router, http.MethodPost, "/v1/messages", `{"model":"claude-sonnet-4.5","input":"hi"}`)
if w.Code != http.StatusOK {
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
}
}
func TestGroupModelAllowlistWebSocketUpgradeSkipsBodyExtraction(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
router.Use(func(c *gin.Context) {
c.Set(string(ContextKeyAPIKey), allowlistAPIKey(true, "claude-sonnet-4.5"))
c.Next()
})
router.Use(GroupModelAllowlist())
router.GET("/v1/responses", func(c *gin.Context) {
c.Status(http.StatusSwitchingProtocols)
})
req := httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
req.Header.Set("Upgrade", "websocket")
req.Header.Set("Connection", "Upgrade")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusSwitchingProtocols {
t.Fatalf("expected WS upgrade to reach handler, got %d: %s", w.Code, w.Body.String())
}
}
// 伪造 WebSocket Upgrade 头的普通请求不得绕过白名单校验。
func TestGroupModelAllowlistForgedUpgradeHeaderDoesNotBypass(t *testing.T) {
router, calls := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "claude-sonnet-4.5"), "/v1")
req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{"model":"claude-opus-4.6"}`))
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Upgrade", "websocket")
req.Header.Set("Connection", "Upgrade")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusNotFound {
t.Fatalf("expected 404 despite forged upgrade headers, got %d: %s", w.Code, w.Body.String())
}
if !strings.Contains(w.Body.String(), "claude-opus-4.6") {
t.Fatalf("expected blocked model in message, got %s", w.Body.String())
}
if len(*calls) != 0 {
t.Fatalf("expected handler not to run, got %v", *calls)
}
}
// Grok Realtime 升级请求的模型在查询参数里,携带 Upgrade 头也必须校验。
func TestGroupModelAllowlistRealtimeUpgradeStillChecksQueryModel(t *testing.T) {
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "grok-4.6"), "/v1")
req := httptest.NewRequest(http.MethodGet, "/v1/realtime?model=grok-4.20", nil)
req.Header.Set("Upgrade", "websocket")
req.Header.Set("Connection", "Upgrade")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusNotFound {
t.Fatalf("expected 404 for unlisted realtime model with upgrade headers, got %d: %s", w.Code, w.Body.String())
}
}
// Live 入口的 handler 只读 session.model:准入必须以 session.model 为准,
// 与顶层 model 不一致时按实际执行模型校验。
func TestGroupModelAllowlistLiveUsesSessionModel(t *testing.T) {
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "gpt-realtime-model"), "/v1")
// session.model 不在白名单:即使顶层 model 命中也必须拒绝。
w := doJSON(t, router, http.MethodPost, "/v1/live", `{"model":"gpt-realtime-model","session":{"model":"blocked-model"},"sdp":"v=0"}`)
if w.Code != http.StatusNotFound {
t.Fatalf("expected 404 when session.model is blocked, got %d: %s", w.Code, w.Body.String())
}
// session.model 命中:顶层 model 不在白名单也不影响(handler 不读它)。
w = doJSON(t, router, http.MethodPost, "/v1/live", `{"model":"blocked-model","session":{"model":"gpt-realtime-model"},"sdp":"v=0"}`)
if w.Code != http.StatusOK {
t.Fatalf("expected 200 when session.model is allowlisted, got %d: %s", w.Code, w.Body.String())
}
}
func TestGroupModelAllowlistMarksRejectionReasons(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
router.Use(func(c *gin.Context) {
c.Set(string(ContextKeyAPIKey), allowlistAPIKey(true, "claude-sonnet-4.5"))
c.Next()
})
router.Use(GroupModelAllowlist())
router.POST("/v1/responses", func(c *gin.Context) { c.Status(http.StatusOK) })
w := doJSON(t, router, http.MethodPost, "/v1/responses", `{"model":"gpt-5.4"}`)
if w.Code != http.StatusNotFound {
t.Fatalf("expected 404, got %d", w.Code)
}
// IngressRejectModelNotAllowed 与业务限流原因由 ops 中间件在真实链路落日志,
// 此处至少验证拒绝原因可从 context 读取(MarkIngressRejected 生效)。
}
func TestGroupModelAllowlistNilGroupPasses(t *testing.T) {
router, calls := newGroupModelAllowlistTestRouter(&service.APIKey{}, "/v1")
w := doJSON(t, router, http.MethodPost, "/v1/messages", `{"model":"whatever"}`)
if w.Code != http.StatusOK || len(*calls) != 1 {
t.Fatalf("expected passthrough for key without group, got %d", w.Code)
}
}
// encoding/json 绑定系 handler(如批量生图 ShouldBindJSON)对键名大小写不敏感:
// {"Model":"blocked"} 会绑定为 blocked,准入必须拒绝。
func TestGroupModelAllowlistCaseVariantModelKeyRejected(t *testing.T) {
router, calls := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "claude-sonnet-4.5"), "/v1")
w := doJSON(t, router, http.MethodPost, "/v1/responses", `{"Model":"claude-opus-4.6"}`)
if w.Code != http.StatusNotFound {
t.Fatalf("case-variant model key must be validated, got %d: %s", w.Code, w.Body.String())
}
if !strings.Contains(w.Body.String(), "claude-opus-4.6") {
t.Fatalf("expected blocked model in message, got %s", w.Body.String())
}
if len(*calls) != 0 {
t.Fatalf("expected handler not to run, got %v", *calls)
}
}
// encoding/json 绑定对重复键取末值、gjson 取首个:两个值都必须通过校验。
func TestGroupModelAllowlistDuplicateModelKeysRejected(t *testing.T) {
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "claude-sonnet-4.5"), "/v1")
w := doJSON(t, router, http.MethodPost, "/v1/responses", `{"model":"claude-sonnet-4.5","model":"claude-opus-4.6"}`)
if w.Code != http.StatusNotFound {
t.Fatalf("duplicate model keys must all be validated, got %d: %s", w.Code, w.Body.String())
}
if !strings.Contains(w.Body.String(), "claude-opus-4.6") {
t.Fatalf("expected the disallowed duplicate to be reported, got %s", w.Body.String())
}
}
// 图片编辑的 multipart 解析器对重复 model 字段取末值:全部字段值都必须校验。
func TestGroupModelAllowlistMultipartDuplicateModelFieldsRejected(t *testing.T) {
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "gpt-image-1"), "/v1")
var body bytes.Buffer
writer := multipart.NewWriter(&body)
_ = writer.WriteField("model", "gpt-image-1")
_ = writer.WriteField("model", "gpt-image-1.5")
_ = writer.WriteField("prompt", "cat")
_ = writer.Close()
req := httptest.NewRequest(http.MethodPost, "/v1/images/edits", &body)
req.Header.Set("Content-Type", writer.FormDataContentType())
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusNotFound {
t.Fatalf("duplicate multipart model fields must all be validated, got %d: %s", w.Code, w.Body.String())
}
if !strings.Contains(w.Body.String(), "gpt-image-1.5") {
t.Fatalf("expected the disallowed field value to be reported, got %s", w.Body.String())
}
}
// 全部候选都在白名单内时放行(重复但同值的字段不误伤)。
func TestGroupModelAllowlistDuplicateIdenticalModelsAllowed(t *testing.T) {
router, calls := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "gpt-image-1"), "/v1")
var body bytes.Buffer
writer := multipart.NewWriter(&body)
_ = writer.WriteField("model", "gpt-image-1")
_ = writer.WriteField("model", "gpt-image-1")
_ = writer.Close()
req := httptest.NewRequest(http.MethodPost, "/v1/images/edits", &body)
req.Header.Set("Content-Type", writer.FormDataContentType())
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("identical duplicate fields should pass, got %d: %s", w.Code, w.Body.String())
}
if len(*calls) != 1 {
t.Fatalf("expected handler to run once, got %v", *calls)
}
}
@@ -28,6 +28,7 @@ const (
IngressRejectGroupUnassigned IngressRejectReason = "group_unassigned"
IngressRejectInvalidAuthRateLimited IngressRejectReason = "invalid_auth_rate_limited"
IngressRejectAPIKeyAuthOverloaded IngressRejectReason = "api_key_auth_overloaded"
IngressRejectModelNotAllowed IngressRejectReason = "model_not_allowed"
)
const ingressRejectReasonContextKey = "ingress_reject_reason"
@@ -100,11 +100,19 @@ func abortWithOpenAIQuotaError(c *gin.Context, statusCode int, message string) {
// GatewayErrorWriter 定义网关错误响应格式(不同协议使用不同格式)
type GatewayErrorWriter func(c *gin.Context, status int, message string)
// AnthropicErrorWriter 按 Anthropic API 规范输出错误
// AnthropicErrorWriter 按 Anthropic API 规范输出错误;error.type 随状态码映射
// (404 -> not_found_error,403 -> permission_error,其余 api_error)。
func AnthropicErrorWriter(c *gin.Context, status int, message string) {
errorType := "api_error"
switch status {
case http.StatusNotFound:
errorType = "not_found_error"
case http.StatusForbidden:
errorType = "permission_error"
}
c.JSON(status, gin.H{
"type": "error",
"error": gin.H{"type": "permission_error", "message": message},
"error": gin.H{"type": errorType, "message": message},
})
}
@@ -119,6 +127,17 @@ func GoogleErrorWriter(c *gin.Context, status int, message string) {
})
}
// OpenAIErrorWriter 按 OpenAI API 规范输出模型级错误(model_not_found)。
func OpenAIErrorWriter(c *gin.Context, status int, message string) {
c.JSON(status, gin.H{
"error": gin.H{
"message": message,
"type": "invalid_request_error",
"code": "model_not_found",
},
})
}
// RequireGroupAssignment 检查 API Key 是否已分配到分组,
// 如果未分组且系统设置不允许未分组 Key 调度则返回 403。
func RequireGroupAssignment(settingService *service.SettingService, writeError GatewayErrorWriter) gin.HandlerFunc {
+1 -1
View File
@@ -333,7 +333,7 @@ func registerGroupRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
groups.GET("/capacity-summary", h.Admin.Group.GetCapacitySummary)
groups.GET("/live-capability", h.Admin.Group.GetLiveCapability)
groups.PUT("/sort-order", h.Admin.Group.UpdateSortOrder)
groups.GET("/:id/models-list-candidates", h.Admin.Group.GetModelsListCandidates)
groups.GET("/:id/model-allowlist-candidates", h.Admin.Group.GetGroupModelAllowlistCandidates)
groups.GET("/:id/composite-routes", h.Admin.Group.ListCompositeRoutes)
groups.POST("/:id/composite-routes", h.Admin.Group.CreateCompositeRoute)
groups.POST("/:id/composite-routes/preview", h.Admin.Group.PreviewCompositeRoute)
@@ -10,6 +10,7 @@ import (
"strings"
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/requestmodel"
servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
@@ -191,7 +192,7 @@ func TestCompositeRequestModelFromMultipartLiveSession(t *testing.T) {
require.NoError(t, writer.WriteField("session", `{"model":"live-alias"}`))
require.NoError(t, writer.Close())
require.Equal(t, "live-alias", compositeRequestModelFromBody(writer.FormDataContentType(), body.Bytes()))
require.Equal(t, "live-alias", requestmodel.FromBody(writer.FormDataContentType(), body.Bytes()))
}
func TestCompositeCodexControlPathsUseResponsesRoutes(t *testing.T) {
@@ -312,3 +313,58 @@ func TestCompositeGeminiTargetPlatformMiddlewareUsesPathRoute(t *testing.T) {
require.Equal(t, http.StatusNoContent, w.Code)
}
// Live 入口顶层 model 与 session.model 不一致时,合成路由必须按 session.model
// 分发与改写(与白名单准入、Live handler 一致),不得命中顶层别名的映射并把
// session 模型覆盖为别名上游。
func TestCompositeLiveRouteDispatchesBySessionModelNotTopLevelAlias(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
resolver := service.NewCompositeRouteResolver(compositeRouteRepoStub{
routes: []service.CompositeModelRoute{
{
ID: 1,
GroupID: 1,
PublicModel: "live-alias",
MatchType: service.CompositeRouteMatchExact,
TargetPlatform: service.PlatformOpenAI,
UpstreamModel: "gpt-alias-upstream",
Endpoint: service.CompositeRouteEndpointAny,
Priority: 100,
Enabled: true,
},
},
})
router.Use(gin.HandlerFunc(servermiddleware.APIKeyAuthMiddleware(func(c *gin.Context) {
groupID := int64(1)
c.Set(string(servermiddleware.ContextKeyAPIKey), &service.APIKey{
GroupID: &groupID,
Group: &service.Group{ID: groupID, Platform: service.PlatformComposite},
})
c.Next()
})))
router.Use(compositeTargetPlatformMiddleware(resolver))
router.POST("/backend-api/codex/realtime/calls", func(c *gin.Context) {
upstream, ok := service.ResolvedUpstreamModelFromContext(c.Request.Context())
require.True(t, ok)
require.Equal(t, "gpt-realtime-1", upstream,
"Live dispatch must resolve the session model, never the top-level alias upstream")
body, err := io.ReadAll(c.Request.Body)
require.NoError(t, err)
require.JSONEq(t, `{"model":"live-alias","session":{"model":"gpt-realtime-1"},"sdp":"v=0"}`, string(body),
"session.model must stay untouched when dispatch resolves by session model")
c.Status(http.StatusNoContent)
})
req := httptest.NewRequest(
http.MethodPost,
"/backend-api/codex/realtime/calls",
strings.NewReader(`{"model":"live-alias","session":{"model":"gpt-realtime-1"},"sdp":"v=0"}`),
)
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusNoContent, w.Code)
}
+57 -108
View File
@@ -1,18 +1,14 @@
package routes
import (
"bytes"
"errors"
"io"
"mime"
"mime/multipart"
"net/http"
"strconv"
"strings"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/handler"
pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
"github.com/Wei-Shaw/sub2api/internal/pkg/requestmodel"
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
@@ -45,6 +41,10 @@ func RegisterGatewayRoutes(
requireGroupAnthropic := middleware.RequireGroupAssignment(settingService, middleware.AnthropicErrorWriter)
requireGroupGoogle := middleware.RequireGroupAssignment(settingService, middleware.GoogleErrorWriter)
// 分组级模型白名单准入:在 apiKeyAuth 之后、compositeTarget 之前,
// 保证校验发生在合成路由改写与调度之前,且只看客户端书写的模型名。
groupModelAllowlist := middleware.GroupModelAllowlist()
isOpenAIResponsesCompatibleGatewayPlatform := func(c *gin.Context) bool {
switch getGroupPlatform(c) {
case service.PlatformOpenAI, service.PlatformGrok,
@@ -189,6 +189,7 @@ func RegisterGatewayRoutes(
gateway.Use(endpointNorm)
gateway.Use(gin.HandlerFunc(apiKeyAuth))
gateway.GET("/sub2api/billing", h.Gateway.KeyBillingInfo)
gateway.Use(groupModelAllowlist)
gateway.Use(compositeTarget)
gateway.Use(requireGroupAnthropic)
{
@@ -341,6 +342,7 @@ func RegisterGatewayRoutes(
gemini.Use(opsErrorLogger)
gemini.Use(endpointNorm)
gemini.Use(middleware.APIKeyAuthWithSubscriptionGoogle(apiKeyService, subscriptionService, cfg))
gemini.Use(groupModelAllowlist)
gemini.Use(compositeGeminiTarget)
gemini.Use(requireGroupGoogle)
{
@@ -358,16 +360,21 @@ func RegisterGatewayRoutes(
}
h.Gateway.Responses(c)
}
r.POST("/responses", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, responsesHandler)
r.POST("/responses/*subpath", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, guardResponsesSubpath(responsesHandler))
r.POST("/alpha/search", textBodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, h.OpenAIGateway.AlphaSearch)
r.GET("/responses", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, func(c *gin.Context) {
// 根路径别名共用中间件链:白名单准入在 apiKeyAuth 之后、compositeTarget
// 之前,避免逐条路由手工维护链导致漏挂。
rootRoute := func(method, path string, limit gin.HandlerFunc, handler gin.HandlerFunc) {
r.Handle(method, path, limit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), groupModelAllowlist, compositeTarget, requireGroupAnthropic, handler)
}
rootRoute(http.MethodPost, "/responses", bodyLimit, responsesHandler)
rootRoute(http.MethodPost, "/responses/*subpath", bodyLimit, guardResponsesSubpath(responsesHandler))
rootRoute(http.MethodPost, "/alpha/search", textBodyLimit, h.OpenAIGateway.AlphaSearch)
rootRoute(http.MethodGet, "/responses", bodyLimit, func(c *gin.Context) {
h.OpenAIGateway.ResponsesWebSocket(c)
})
r.GET("/models", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, modelsHandler)
r.POST("/messages/count_tokens", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, countTokensHandler)
rootRoute(http.MethodGet, "/models", bodyLimit, modelsHandler)
rootRoute(http.MethodPost, "/messages/count_tokens", bodyLimit, countTokensHandler)
codexDirect := r.Group("/backend-api/codex")
codexDirect.Use(bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic)
codexDirect.Use(bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), groupModelAllowlist, compositeTarget, requireGroupAnthropic)
{
codexDirect.POST("/realtime/calls", h.OpenAIGateway.Live)
codexDirect.GET("/:call_id", h.OpenAIGateway.LiveSideband)
@@ -380,14 +387,14 @@ func RegisterGatewayRoutes(
codexDirect.GET("/models", codexModelsHandler)
}
// OpenAI Chat Completions API(不带v1前缀的别名)— auto-route based on group platform
r.POST("/chat/completions", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, func(c *gin.Context) {
rootRoute(http.MethodPost, "/chat/completions", bodyLimit, func(c *gin.Context) {
if isOpenAIResponsesCompatibleGatewayPlatform(c) {
h.OpenAIGateway.ChatCompletions(c)
return
}
h.Gateway.ChatCompletions(c)
})
r.POST("/embeddings", textBodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, func(c *gin.Context) {
rootRoute(http.MethodPost, "/embeddings", textBodyLimit, func(c *gin.Context) {
if !isOpenAIOnlyEndpointGatewayPlatform(c) {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{
@@ -400,23 +407,23 @@ func RegisterGatewayRoutes(
}
h.OpenAIGateway.Embeddings(c)
})
r.POST("/images/generations", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, imagesHandler)
r.POST("/images/edits", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, imagesHandler)
r.POST("/images/generations/async", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, h.AsyncImage.Submit)
r.POST("/images/edits/async", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, h.AsyncImage.Submit)
r.GET("/images/tasks/:task_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, h.AsyncImage.Get)
r.POST("/videos", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoGenerationHandler)
r.POST("/videos/generations", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoGenerationHandler)
r.POST("/videos/edits", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoEditHandler)
r.POST("/videos/extensions", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoExtensionHandler)
r.GET("/videos/generations/:request_id/content", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoContentHandler)
r.GET("/videos/edits/:request_id/content", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoContentHandler)
r.GET("/videos/extensions/:request_id/content", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoContentHandler)
r.GET("/videos/generations/:request_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoStatusHandler)
r.GET("/videos/edits/:request_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoStatusHandler)
r.GET("/videos/extensions/:request_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoStatusHandler)
r.GET("/videos/:request_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoStatusHandler)
r.GET("/videos/:request_id/content", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoContentHandler)
rootRoute(http.MethodPost, "/images/generations", bodyLimit, imagesHandler)
rootRoute(http.MethodPost, "/images/edits", bodyLimit, imagesHandler)
rootRoute(http.MethodPost, "/images/generations/async", bodyLimit, h.AsyncImage.Submit)
rootRoute(http.MethodPost, "/images/edits/async", bodyLimit, h.AsyncImage.Submit)
rootRoute(http.MethodGet, "/images/tasks/:task_id", bodyLimit, h.AsyncImage.Get)
rootRoute(http.MethodPost, "/videos", bodyLimit, videoGenerationHandler)
rootRoute(http.MethodPost, "/videos/generations", bodyLimit, videoGenerationHandler)
rootRoute(http.MethodPost, "/videos/edits", bodyLimit, videoEditHandler)
rootRoute(http.MethodPost, "/videos/extensions", bodyLimit, videoExtensionHandler)
rootRoute(http.MethodGet, "/videos/generations/:request_id/content", bodyLimit, videoContentHandler)
rootRoute(http.MethodGet, "/videos/edits/:request_id/content", bodyLimit, videoContentHandler)
rootRoute(http.MethodGet, "/videos/extensions/:request_id/content", bodyLimit, videoContentHandler)
rootRoute(http.MethodGet, "/videos/generations/:request_id", bodyLimit, videoStatusHandler)
rootRoute(http.MethodGet, "/videos/edits/:request_id", bodyLimit, videoStatusHandler)
rootRoute(http.MethodGet, "/videos/extensions/:request_id", bodyLimit, videoStatusHandler)
rootRoute(http.MethodGet, "/videos/:request_id", bodyLimit, videoStatusHandler)
rootRoute(http.MethodGet, "/videos/:request_id/content", bodyLimit, videoContentHandler)
rootVoiceHandler := func(endpoint string) gin.HandlerFunc {
return func(c *gin.Context) {
@@ -428,9 +435,9 @@ func RegisterGatewayRoutes(
h.OpenAIGateway.GrokVoice(c, endpoint)
}
}
r.POST("/tts", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootVoiceHandler("tts"))
r.POST("/stt", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootVoiceHandler("stt"))
r.POST("/custom-voices", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootVoiceHandler("custom-voices"))
rootRoute(http.MethodPost, "/tts", bodyLimit, rootVoiceHandler("tts"))
rootRoute(http.MethodPost, "/stt", bodyLimit, rootVoiceHandler("stt"))
rootRoute(http.MethodPost, "/custom-voices", bodyLimit, rootVoiceHandler("custom-voices"))
rootCustomVoicePathHandler := func(c *gin.Context) {
if getGroupPlatform(c) != service.PlatformGrok {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
@@ -439,12 +446,12 @@ func RegisterGatewayRoutes(
}
h.OpenAIGateway.GrokVoice(c, grokCustomVoiceEndpoint(c))
}
r.GET("/custom-voices", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootVoiceHandler("custom-voices"))
r.GET("/custom-voices/:voice_id/audio", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootCustomVoicePathHandler)
r.GET("/custom-voices/:voice_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootCustomVoicePathHandler)
r.PATCH("/custom-voices/:voice_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootCustomVoicePathHandler)
r.DELETE("/custom-voices/:voice_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootCustomVoicePathHandler)
r.GET("/realtime", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, func(c *gin.Context) {
rootRoute(http.MethodGet, "/custom-voices", bodyLimit, rootVoiceHandler("custom-voices"))
rootRoute(http.MethodGet, "/custom-voices/:voice_id/audio", bodyLimit, rootCustomVoicePathHandler)
rootRoute(http.MethodGet, "/custom-voices/:voice_id", bodyLimit, rootCustomVoicePathHandler)
rootRoute(http.MethodPatch, "/custom-voices/:voice_id", bodyLimit, rootCustomVoicePathHandler)
rootRoute(http.MethodDelete, "/custom-voices/:voice_id", bodyLimit, rootCustomVoicePathHandler)
rootRoute(http.MethodGet, "/realtime", bodyLimit, func(c *gin.Context) {
if getGroupPlatform(c) != service.PlatformGrok {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "Realtime API is not supported for this platform"}})
@@ -452,7 +459,7 @@ func RegisterGatewayRoutes(
}
h.OpenAIGateway.GrokRealtime(c)
})
r.POST("/web_search", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, func(c *gin.Context) {
rootRoute(http.MethodPost, "/web_search", bodyLimit, func(c *gin.Context) {
if getGroupPlatform(c) != service.PlatformGrok {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "Web Search API is not supported for this platform"}})
@@ -460,7 +467,7 @@ func RegisterGatewayRoutes(
}
h.Gateway.WebSearch(c)
})
r.POST("/x_search", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, func(c *gin.Context) {
rootRoute(http.MethodPost, "/x_search", bodyLimit, func(c *gin.Context) {
if getGroupPlatform(c) != service.PlatformGrok {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "X Search API is not supported for this platform"}})
@@ -480,6 +487,7 @@ func RegisterGatewayRoutes(
antigravityV1.Use(endpointNorm)
antigravityV1.Use(middleware.ForcePlatform(service.PlatformAntigravity))
antigravityV1.Use(gin.HandlerFunc(apiKeyAuth))
antigravityV1.Use(groupModelAllowlist)
antigravityV1.Use(requireGroupAnthropic)
{
antigravityV1.POST("/messages", h.Gateway.Messages)
@@ -495,6 +503,7 @@ func RegisterGatewayRoutes(
antigravityV1Beta.Use(endpointNorm)
antigravityV1Beta.Use(middleware.ForcePlatform(service.PlatformAntigravity))
antigravityV1Beta.Use(middleware.APIKeyAuthWithSubscriptionGoogle(apiKeyService, subscriptionService, cfg))
antigravityV1Beta.Use(groupModelAllowlist)
antigravityV1Beta.Use(requireGroupGoogle)
{
antigravityV1Beta.GET("/models", h.Gateway.GeminiV1BetaListModels)
@@ -555,7 +564,10 @@ func compositeTargetPlatformMiddleware(resolver *service.CompositeRouteResolver)
return
}
model := compositeRequestModelFromBody(c.GetHeader("Content-Type"), body)
// Live 入口按 session.model 分发并改写(与白名单准入、Live handler 的
// 解析保持一致),避免顶层 model 别名与 session 模型不一致时改错对象。
routePath := c.FullPath()
model := requestmodel.FromBodyForRoute(routePath, c.GetHeader("Content-Type"), body)
if model != "" {
decision, err := resolver.Resolve(c.Request.Context(), apiKey.Group.ID, model, compositeRouteEndpointForPath(c.Request.URL.Path))
if err != nil {
@@ -566,7 +578,7 @@ func compositeTargetPlatformMiddleware(resolver *service.CompositeRouteResolver)
if decision.Matched {
c.Request = c.Request.WithContext(service.WithCompositeRouteDecision(c.Request.Context(), decision))
if upstreamModel := strings.TrimSpace(decision.UpstreamModel); upstreamModel != "" && upstreamModel != model && gjson.ValidBytes(body) {
if _, modelPath := compositeJSONRequestModel(body); modelPath != "" {
if _, modelPath := requestmodel.JSONModelPathForRoute(routePath, body); modelPath != "" {
if rewritten, rewriteErr := sjson.SetBytes(body, modelPath, upstreamModel); rewriteErr == nil {
body = rewritten
}
@@ -574,68 +586,11 @@ func compositeTargetPlatformMiddleware(resolver *service.CompositeRouteResolver)
}
}
}
resetRequestBody(c, body)
requestmodel.ResetRequestBody(c.Request, body)
c.Next()
}
}
func compositeRequestModelFromBody(contentType string, body []byte) string {
if model, _ := compositeJSONRequestModel(body); model != "" {
return model
}
return compositeMultipartModelFromBody(contentType, body)
}
func compositeJSONRequestModel(body []byte) (string, string) {
for _, path := range []string{"model", "session.model"} {
model := gjson.GetBytes(body, path)
if model.Type != gjson.String {
continue
}
if value := strings.TrimSpace(model.String()); value != "" {
return value, path
}
}
return "", ""
}
func compositeMultipartModelFromBody(contentType string, body []byte) string {
mediaType, params, err := mime.ParseMediaType(strings.TrimSpace(contentType))
if err != nil || !strings.EqualFold(mediaType, "multipart/form-data") {
return ""
}
boundary := strings.TrimSpace(params["boundary"])
if boundary == "" {
return ""
}
reader := multipart.NewReader(bytes.NewReader(body), boundary)
for {
part, err := reader.NextPart()
if errors.Is(err, io.EOF) {
return ""
}
if err != nil {
return ""
}
fieldName := part.FormName()
if part.FileName() != "" || (fieldName != "model" && fieldName != "session") {
continue
}
data, err := io.ReadAll(part)
if err != nil {
return ""
}
switch fieldName {
case "model":
return strings.TrimSpace(string(data))
case "session":
if model, _ := compositeJSONRequestModel(data); model != "" {
return model
}
}
}
}
func compositeGeminiTargetPlatformMiddleware(resolver *service.CompositeRouteResolver) gin.HandlerFunc {
if resolver == nil {
resolver = service.NewCompositeRouteResolver(nil)
@@ -696,12 +651,6 @@ func compositeGeminiModelFromParams(c *gin.Context) string {
return modelAction
}
func resetRequestBody(c *gin.Context, body []byte) {
c.Request.Body = io.NopCloser(bytes.NewReader(body))
c.Request.ContentLength = int64(len(body))
c.Request.Header.Set("Content-Length", strconv.Itoa(len(body)))
}
func compositeRouteEndpointForPath(path string) string {
switch {
case strings.Contains(path, "/messages/count_tokens"):
@@ -0,0 +1,208 @@
package routes
import (
"net/http"
"net/http/httptest"
"os"
"regexp"
"strings"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/handler"
servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func newGatewayRoutesTestRouterWithGroup(group *service.Group) *gin.Engine {
gin.SetMode(gin.TestMode)
router := gin.New()
RegisterGatewayRoutes(
router,
&handler.Handlers{
Gateway: &handler.GatewayHandler{},
OpenAIGateway: &handler.OpenAIGatewayHandler{},
AsyncImage: handler.NewAsyncImageHandler(nil, nil),
},
servermiddleware.APIKeyAuthMiddleware(func(c *gin.Context) {
groupID := int64(1)
c.Set(string(servermiddleware.ContextKeyAPIKey), &service.APIKey{
GroupID: &groupID,
Group: group,
})
c.Next()
}),
nil,
nil,
nil,
nil,
nil,
&config.Config{
Gateway: config.GatewayConfig{
MaxBodySize: 1024 * 1024,
TextMaxBodySize: 1024 * 1024,
},
},
)
return router
}
func allowlistGroup(platform string, enabled bool, models ...string) *service.Group {
return &service.Group{
Platform: platform,
ModelAllowlist: service.GroupModelAllowlist{
Enabled: enabled,
Models: models,
},
}
}
// TestGatewayRoutesGroupModelAllowlistMountedOnEveryGatewayRoute follows the
// source-level route assertion convention of prompt_audit_route_coverage_test.go:
// every gateway chain must mount groupModelAllowlist after api key auth and
// before the composite rewrite (gateway.go + rootRoute helper).
func TestGatewayRoutesGroupModelAllowlistMountedOnEveryGatewayRoute(t *testing.T) {
routeSource, err := os.ReadFile("gateway.go")
require.NoError(t, err)
source := string(routeSource)
// rootRoute helper:apiKeyAuth 之后、compositeTarget 之前。
rootHelper := regexp.MustCompile(regexp.QuoteMeta(`r.Handle(method, path, limit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), groupModelAllowlist, compositeTarget, requireGroupAnthropic, handler)`))
require.Regexp(t, rootHelper, source,
"root alias helper must place the allowlist between apiKeyAuth and compositeTarget")
chains := []struct {
group string
auth string
marker string
composite string
}{
{group: "gateway", auth: "gin.HandlerFunc(apiKeyAuth)", marker: "gateway.Use(groupModelAllowlist)", composite: "gateway.Use(compositeTarget)"},
{group: "gemini", auth: "middleware.APIKeyAuthWithSubscriptionGoogle(apiKeyService, subscriptionService, cfg)", marker: "gemini.Use(groupModelAllowlist)", composite: "gemini.Use(compositeGeminiTarget)"},
{group: "antigravityV1", auth: "gin.HandlerFunc(apiKeyAuth)", marker: "antigravityV1.Use(groupModelAllowlist)", composite: "antigravityV1.Use(requireGroupAnthropic)"},
{group: "antigravityV1Beta", auth: "middleware.APIKeyAuthWithSubscriptionGoogle(apiKeyService, subscriptionService, cfg)", marker: "antigravityV1Beta.Use(groupModelAllowlist)", composite: "antigravityV1Beta.Use(requireGroupGoogle)"},
}
for _, chain := range chains {
re := regexp.MustCompile(
regexp.QuoteMeta(chain.group+".Use("+chain.auth) +
`[\s\S]{0,400}?` + regexp.QuoteMeta(chain.marker) +
`[\s\S]{0,400}?` + regexp.QuoteMeta(chain.composite))
require.Regexp(t, re, source,
"%s chain must mount groupModelAllowlist after auth and before %s", chain.group, chain.composite)
}
// codexDirect 链是一条 Use 调用,直接断言顺序。
codexDirect := regexp.MustCompile(regexp.QuoteMeta(`codexDirect.Use(bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), groupModelAllowlist, compositeTarget, requireGroupAnthropic)`))
require.Regexp(t, codexDirect, source, "codexDirect chain must mount the allowlist after auth and before compositeTarget")
// 所有带 apiKeyAuth 的根路径路由必须收敛到 rootRoute,避免漏挂。
stray := regexp.MustCompile(`\br\.(GET|POST|PUT|PATCH|DELETE)\("[^"]+",[^(]*apiKeyAuth`)
require.NotRegexp(t, stray, source,
"root alias routes must use rootRoute so the allowlist cannot be forgotten")
}
// TestGatewayRoutesGroupModelAllowlistBlocksCompositeModelBeforeRewrite asserts
// admission is decided on the client-written public model before the composite
// route middleware can rewrite it to an upstream model.
func TestGatewayRoutesGroupModelAllowlistBlocksCompositeModelBeforeRewrite(t *testing.T) {
router := newGatewayRoutesTestRouterWithGroup(allowlistGroup(service.PlatformComposite, true, "claude-*"))
req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{"model":"grok-4.3","messages":[{"role":"user","content":"hi"}]}`))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusNotFound, w.Code)
require.Contains(t, w.Body.String(), "not_found_error")
require.Contains(t, w.Body.String(), "grok-4.3")
}
func TestGatewayRoutesGroupModelAllowlistAllowsListedCompositeModel(t *testing.T) {
router := newGatewayRoutesTestRouterWithGroup(allowlistGroup(service.PlatformComposite, true, "claude-*"))
req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{"model":"claude-sonnet-4.5","messages":[{"role":"user","content":"hi"}]}`))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.NotEqual(t, http.StatusNotFound, w.Code, "allowlisted model must pass admission: %s", w.Body.String())
}
func TestGatewayRoutesGroupModelAllowlistCoversRootAliasRoutes(t *testing.T) {
router := newGatewayRoutesTestRouterWithGroup(allowlistGroup(service.PlatformOpenAI, true, "gpt-5.4"))
paths := []struct {
method string
path string
body string
}{
{http.MethodPost, "/responses", `{"model":"gpt-4.1"}`},
{http.MethodPost, "/responses/compact", `{"model":"gpt-4.1"}`},
{http.MethodPost, "/chat/completions", `{"model":"gpt-4.1"}`},
{http.MethodPost, "/embeddings", `{"model":"gpt-4.1","input":"hi"}`},
{http.MethodPost, "/images/generations", `{"model":"gpt-4.1"}`},
{http.MethodPost, "/images/edits", `{"model":"gpt-4.1"}`},
{http.MethodPost, "/videos/generations", `{"model":"gpt-4.1"}`},
{http.MethodPost, "/messages/count_tokens", `{"model":"gpt-4.1","messages":[]}`},
{http.MethodPost, "/tts", `{"model":"gpt-4.1"}`},
{http.MethodPost, "/stt", `{"model":"gpt-4.1"}`},
{http.MethodPost, "/alpha/search", `{"model":"gpt-4.1"}`},
{http.MethodGet, "/realtime?model=gpt-4.1", ""},
{http.MethodPost, "/v1/responses", `{"model":"gpt-4.1"}`},
{http.MethodPost, "/v1/messages", `{"model":"gpt-4.1"}`},
{http.MethodPost, "/v1/messages/count_tokens", `{"model":"gpt-4.1","messages":[]}`},
{http.MethodPost, "/v1/chat/completions", `{"model":"gpt-4.1"}`},
{http.MethodPost, "/v1/embeddings", `{"model":"gpt-4.1","input":"hi"}`},
{http.MethodPost, "/v1/images/generations", `{"model":"gpt-4.1"}`},
{http.MethodPost, "/v1/videos/generations", `{"model":"gpt-4.1"}`},
{http.MethodPost, "/v1/live", `{"session":{"model":"gpt-4.1"},"sdp":"v=0"}`},
{http.MethodPost, "/backend-api/codex/responses", `{"model":"gpt-4.1"}`},
{http.MethodPost, "/backend-api/codex/realtime/calls", `{"session":{"model":"gpt-4.1"},"sdp":"v=0"}`},
{http.MethodPost, "/antigravity/v1/messages", `{"model":"gemini-2.5-pro","messages":[]}`},
}
for _, tc := range paths {
req := httptest.NewRequest(tc.method, tc.path, strings.NewReader(tc.body))
if tc.body != "" {
req.Header.Set("Content-Type", "application/json")
}
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusNotFound, w.Code, "%s %s should be denied by the allowlist, got body: %s", tc.method, tc.path, w.Body.String())
require.Contains(t, w.Body.String(), "not available for this group", "%s %s", tc.method, tc.path)
}
}
func TestGatewayRoutesGroupModelAllowlistSkipsWebSocketUpgrade(t *testing.T) {
router := newGatewayRoutesTestRouterWithGroup(allowlistGroup(service.PlatformOpenAI, true, "gpt-5.4"))
req := httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
req.Header.Set("Upgrade", "websocket")
req.Header.Set("Connection", "Upgrade")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.NotEqual(t, http.StatusNotFound, w.Code,
"WS upgrade requests must be admitted and per-frame checked in the handler, got: %s", w.Body.String())
}
func TestGatewayRoutesGroupModelAllowlistModelFreeRoutesUnaffected(t *testing.T) {
router := newGatewayRoutesTestRouterWithGroup(allowlistGroup(service.PlatformGrok, true, "grok-4.6"))
// 不携带模型的入口天然不受白名单影响(状态查询/内容下载/模型列表)。
for _, path := range []string{
"/v1/videos/generations/req-1",
"/v1/videos/req-1/content",
"/v1/images/tasks/task-1",
"/v1/custom-voices",
} {
req := httptest.NewRequest(http.MethodGet, path, nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.NotContains(t, w.Body.String(), "not available for this group",
"%s should not be blocked by the allowlist middleware, got: %s", path, w.Body.String())
}
}
+13 -3
View File
@@ -474,6 +474,12 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn
}
}
// 白名单在创建路径同样收口:开启但为空、通配位置非法都会 400。
modelAllowlist, err := normalizeGroupModelAllowlist(input.ModelAllowlist)
if err != nil {
return nil, err
}
group := &Group{
Name: input.Name,
Description: input.Description,
@@ -528,7 +534,7 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn
RequirePrivacySet: input.RequirePrivacySet,
DefaultMappedModel: input.DefaultMappedModel,
MessagesDispatchModelConfig: normalizeOpenAIMessagesDispatchModelConfig(input.MessagesDispatchModelConfig),
ModelsListConfig: normalizeGroupModelsListConfig(input.ModelsListConfig),
ModelAllowlist: modelAllowlist,
// 固定账号 manifest 配置:账号绑定发生在分组创建之后,创建路径禁止开启,
// 成员关系无从校验(前端创建对话框也不展示)。
CodexModelsManifestConfig: normalizeCodexModelsManifestConfig(platform, input.CodexModelsManifestConfig),
@@ -909,8 +915,12 @@ func (s *adminServiceImpl) UpdateGroup(ctx context.Context, id int64, input *Upd
if input.MessagesDispatchModelConfig != nil {
group.MessagesDispatchModelConfig = normalizeOpenAIMessagesDispatchModelConfig(*input.MessagesDispatchModelConfig)
}
if input.ModelsListConfig != nil {
group.ModelsListConfig = normalizeGroupModelsListConfig(*input.ModelsListConfig)
if input.ModelAllowlist != nil {
modelAllowlist, err := normalizeGroupModelAllowlist(*input.ModelAllowlist)
if err != nil {
return nil, err
}
group.ModelAllowlist = modelAllowlist
}
if input.CodexModelsManifestConfig != nil {
group.CodexModelsManifestConfig = *input.CodexModelsManifestConfig
@@ -150,9 +150,9 @@ func cloneGroupForDuplicate(source *Group, operationID string) *Group {
RequirePrivacySet: source.RequirePrivacySet,
DefaultMappedModel: source.DefaultMappedModel,
MessagesDispatchModelConfig: cloneGroupMessagesDispatchModelConfig(source.MessagesDispatchModelConfig),
ModelsListConfig: GroupModelsListConfig{
Enabled: source.ModelsListConfig.Enabled,
Models: append([]string(nil), source.ModelsListConfig.Models...),
ModelAllowlist: GroupModelAllowlist{
Enabled: source.ModelAllowlist.Enabled,
Models: append([]string(nil), source.ModelAllowlist.Models...),
},
// 固定账号 manifest 配置指向源分组的账号 ID,复制后成员关系可能变化,重置为关闭且列表为空。
CodexModelsManifestConfig: GroupCodexModelsManifestConfig{},
@@ -61,7 +61,7 @@ func cloneGroupForDuplicateTest(group *Group) *Group {
cloned.ModelRouting = cloneGroupModelRouting(group.ModelRouting)
cloned.SupportedModelScopes = append([]string(nil), group.SupportedModelScopes...)
cloned.MessagesDispatchModelConfig = cloneGroupMessagesDispatchModelConfig(group.MessagesDispatchModelConfig)
cloned.ModelsListConfig.Models = append([]string(nil), group.ModelsListConfig.Models...)
cloned.ModelAllowlist.Models = append([]string(nil), group.ModelAllowlist.Models...)
cloned.AccountGroups = append([]AccountGroup(nil), group.AccountGroups...)
return &cloned
}
@@ -177,7 +177,7 @@ func TestDuplicateGroupCopiesConfigurationDeeplyAndResetsRuntimeState(t *testing
HaikuMappedModel: "gpt-5-mini",
ExactModelMappings: map[string]string{"claude-special": "gpt-special"},
},
ModelsListConfig: GroupModelsListConfig{Enabled: true, Models: []string{"gpt-5.4", "gpt-5-mini"}},
ModelAllowlist: GroupModelAllowlist{Enabled: true, Models: []string{"gpt-5.4", "gpt-5-mini"}},
RPMLimit: 99,
MaxReasoningEffort: "medium",
MaxReasoningEffortOverLimit: ReasoningEffortOverLimitDeny,
@@ -217,7 +217,7 @@ func TestDuplicateGroupCopiesConfigurationDeeplyAndResetsRuntimeState(t *testing
require.Equal(t, source.MessagesDispatchModelConfig, duplicate.MessagesDispatchModelConfig)
require.Equal(t, source.ForceOpenAIFast, duplicate.ForceOpenAIFast)
require.Equal(t, source.FreeOpenAIFast, duplicate.FreeOpenAIFast)
require.Equal(t, source.ModelsListConfig, duplicate.ModelsListConfig)
require.Equal(t, source.ModelAllowlist, duplicate.ModelAllowlist)
require.Equal(t, source.RPMLimit, duplicate.RPMLimit)
require.Equal(t, source.MaxReasoningEffort, duplicate.MaxReasoningEffort)
require.Equal(t, source.MaxReasoningEffortOverLimit, duplicate.MaxReasoningEffortOverLimit)
@@ -235,14 +235,14 @@ func TestDuplicateGroupCopiesConfigurationDeeplyAndResetsRuntimeState(t *testing
duplicate.VideoModelPrices[VideoPriceFamilyGrokImagineVideo15][VideoBillingResolution720P] = 999
duplicate.SupportedModelScopes[0] = "changed"
duplicate.MessagesDispatchModelConfig.ExactModelMappings["claude-special"] = "changed"
duplicate.ModelsListConfig.Models[0] = "changed"
duplicate.ModelAllowlist.Models[0] = "changed"
duplicate.ReasoningEffortMappings[0].To = "changed"
*duplicate.DailyLimitUSD = 999
require.Equal(t, int64(13), source.ModelRouting["gpt-*"][0])
require.Equal(t, 0.14, source.VideoModelPrices[VideoPriceFamilyGrokImagineVideo15][VideoBillingResolution720P])
require.Equal(t, "claude", source.SupportedModelScopes[0])
require.Equal(t, "gpt-special", source.MessagesDispatchModelConfig.ExactModelMappings["claude-special"])
require.Equal(t, "gpt-5.4", source.ModelsListConfig.Models[0])
require.Equal(t, "gpt-5.4", source.ModelAllowlist.Models[0])
require.Equal(t, "xhigh", source.ReasoningEffortMappings[0].To)
require.Equal(t, 11.0, *source.DailyLimitUSD)
}
+2 -2
View File
@@ -272,7 +272,7 @@ type CreateGroupInput struct {
RequireOAuthOnly bool
RequirePrivacySet bool
MessagesDispatchModelConfig OpenAIMessagesDispatchModelConfig
ModelsListConfig GroupModelsListConfig
ModelAllowlist GroupModelAllowlist
// CodexModelsManifestConfig 固定账号 manifest 配置;创建路径禁止开启,仅编辑可配置。
CodexModelsManifestConfig GroupCodexModelsManifestConfig
// RPMLimit 分组 RPM 上限(0 = 不限制)
@@ -353,7 +353,7 @@ type UpdateGroupInput struct {
RequireOAuthOnly *bool
RequirePrivacySet *bool
MessagesDispatchModelConfig *OpenAIMessagesDispatchModelConfig
ModelsListConfig *GroupModelsListConfig
ModelAllowlist *GroupModelAllowlist
// CodexModelsManifestConfig nil 表示不修改;非 openai 平台会被归一化为关闭。
CodexModelsManifestConfig *GroupCodexModelsManifestConfig
// RPMLimit 分组 RPM 上限(0 = 不限制),nil 表示未提供不改动。
@@ -0,0 +1,125 @@
//go:build unit
package service
import (
"context"
"net/http"
"testing"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/stretchr/testify/require"
)
func TestAdminService_CreateGroup_RejectsEmptyEnabledModelAllowlist(t *testing.T) {
repo := &groupRepoStubForAdmin{createID: 51}
svc := &adminServiceImpl{groupRepo: repo}
_, err := svc.CreateGroup(context.Background(), &CreateGroupInput{
Name: "allowlist-empty-group",
Platform: PlatformOpenAI,
RateMultiplier: 1,
ModelAllowlist: GroupModelAllowlist{Enabled: true},
})
require.Error(t, err)
appErr := infraerrors.FromError(err)
require.Equal(t, int32(http.StatusBadRequest), appErr.Code)
require.Equal(t, "INVALID_MODEL_ALLOWLIST", appErr.Reason)
require.Nil(t, repo.created, "拒绝时不得落库")
}
func TestAdminService_CreateGroup_RejectsInvalidAllowlistWildcard(t *testing.T) {
repo := &groupRepoStubForAdmin{createID: 51}
svc := &adminServiceImpl{groupRepo: repo}
_, err := svc.CreateGroup(context.Background(), &CreateGroupInput{
Name: "allowlist-wildcard-group",
Platform: PlatformOpenAI,
RateMultiplier: 1,
ModelAllowlist: GroupModelAllowlist{Enabled: true, Models: []string{"gpt-*-5.4"}},
})
require.Error(t, err)
appErr := infraerrors.FromError(err)
require.Equal(t, int32(http.StatusBadRequest), appErr.Code)
require.Equal(t, "INVALID_MODEL_ALLOWLIST", appErr.Reason)
require.Nil(t, repo.created, "拒绝时不得落库")
}
func TestAdminService_CreateGroup_NormalizesModelAllowlist(t *testing.T) {
repo := &groupRepoStubForAdmin{createID: 52}
svc := &adminServiceImpl{groupRepo: repo}
_, err := svc.CreateGroup(context.Background(), &CreateGroupInput{
Name: "allowlist-normalized-group",
Platform: PlatformOpenAI,
RateMultiplier: 1,
ModelAllowlist: GroupModelAllowlist{
Enabled: true,
Models: []string{" gpt-5.4 ", "GPT-5.4", "claude-*", " "},
},
})
require.NoError(t, err)
require.NotNil(t, repo.created)
require.True(t, repo.created.ModelAllowlist.Enabled)
require.Equal(t, []string{"gpt-5.4", "claude-*"}, repo.created.ModelAllowlist.Models)
}
func TestAdminService_UpdateGroup_RejectsEmptyEnabledModelAllowlist(t *testing.T) {
existing := &Group{ID: 1, Name: "existing", Platform: PlatformOpenAI, Status: StatusActive}
repo := &groupRepoStubForAdmin{getByID: existing}
svc := &adminServiceImpl{groupRepo: repo}
_, err := svc.UpdateGroup(context.Background(), existing.ID, &UpdateGroupInput{
ModelAllowlist: &GroupModelAllowlist{Enabled: true},
})
require.Error(t, err)
appErr := infraerrors.FromError(err)
require.Equal(t, int32(http.StatusBadRequest), appErr.Code)
require.Equal(t, "INVALID_MODEL_ALLOWLIST", appErr.Reason)
require.Nil(t, repo.updated, "拒绝时不得落库")
}
func TestAdminService_UpdateGroup_RejectsInvalidAllowlistWildcard(t *testing.T) {
existing := &Group{ID: 1, Name: "existing", Platform: PlatformOpenAI, Status: StatusActive}
repo := &groupRepoStubForAdmin{getByID: existing}
svc := &adminServiceImpl{groupRepo: repo}
_, err := svc.UpdateGroup(context.Background(), existing.ID, &UpdateGroupInput{
ModelAllowlist: &GroupModelAllowlist{Enabled: true, Models: []string{"foo-*bar"}},
})
require.Error(t, err)
appErr := infraerrors.FromError(err)
require.Equal(t, int32(http.StatusBadRequest), appErr.Code)
require.Equal(t, "INVALID_MODEL_ALLOWLIST", appErr.Reason)
require.Nil(t, repo.updated, "拒绝时不得落库")
}
func TestAdminService_UpdateGroup_NormalizesAndResetsModelAllowlist(t *testing.T) {
existing := &Group{
ID: 1, Name: "existing", Platform: PlatformOpenAI, Status: StatusActive,
ModelAllowlist: GroupModelAllowlist{Enabled: true, Models: []string{"gpt-5.4"}},
}
repo := &groupRepoStubForAdmin{getByID: existing}
svc := &adminServiceImpl{groupRepo: repo}
// 归一化:按小写去重保序(保留首次出现的原始拼写)。
updated := GroupModelAllowlist{Enabled: true, Models: []string{" GPT-5.4 ", "claude-*"}}
_, err := svc.UpdateGroup(context.Background(), existing.ID, &UpdateGroupInput{ModelAllowlist: &updated})
require.NoError(t, err)
require.NotNil(t, repo.updated)
require.Equal(t, []string{"GPT-5.4", "claude-*"}, repo.updated.ModelAllowlist.Models)
// 关闭且清空条目也应被接受(关闭白名单)。
repo.updated = nil
disabled := GroupModelAllowlist{Enabled: false}
_, err = svc.UpdateGroup(context.Background(), existing.ID, &UpdateGroupInput{ModelAllowlist: &disabled})
require.NoError(t, err)
require.NotNil(t, repo.updated)
require.False(t, repo.updated.ModelAllowlist.Enabled)
require.Empty(t, repo.updated.ModelAllowlist.Models)
}
@@ -107,8 +107,8 @@ type APIKeyAuthGroupSnapshot struct {
FreeOpenAIFast bool `json:"free_openai_fast"`
DefaultMappedModel string `json:"default_mapped_model,omitempty"`
MessagesDispatchModelConfig OpenAIMessagesDispatchModelConfig `json:"messages_dispatch_model_config,omitempty"`
ModelsListConfig GroupModelsListConfig `json:"models_list_config,omitempty"`
// CodexModelsManifestConfig 与 ModelsListConfig 一样在认证快照分组里透传,
ModelAllowlist GroupModelAllowlist `json:"model_allowlist,omitempty"`
// CodexModelsManifestConfig 与 ModelAllowlist 一样在认证快照分组里透传,
// Codex /models handler 直接读认证分组对象。
CodexModelsManifestConfig GroupCodexModelsManifestConfig `json:"codex_models_manifest_config,omitempty"`
@@ -14,7 +14,7 @@ import (
"github.com/dgraph-io/ristretto"
)
const apiKeyAuthSnapshotVersion = 23 // v23: group codex_models_manifest_config field
const apiKeyAuthSnapshotVersion = 24 // v24: group model_allowlist field (renamed from models_list_config, enforcing semantics)
type apiKeyAuthCacheConfig struct {
l1Size int
@@ -422,7 +422,7 @@ func (s *APIKeyService) snapshotFromAPIKey(ctx context.Context, apiKey *APIKey)
FreeOpenAIFast: apiKey.Group.FreeOpenAIFast,
DefaultMappedModel: apiKey.Group.DefaultMappedModel,
MessagesDispatchModelConfig: apiKey.Group.MessagesDispatchModelConfig,
ModelsListConfig: apiKey.Group.ModelsListConfig,
ModelAllowlist: apiKey.Group.ModelAllowlist,
CodexModelsManifestConfig: apiKey.Group.CodexModelsManifestConfig,
RPMLimit: apiKey.Group.RPMLimit,
MaxReasoningEffort: apiKey.Group.MaxReasoningEffort,
@@ -524,7 +524,7 @@ func (s *APIKeyService) snapshotToAPIKey(key string, snapshot *APIKeyAuthSnapsho
FreeOpenAIFast: snapshot.Group.FreeOpenAIFast,
DefaultMappedModel: snapshot.Group.DefaultMappedModel,
MessagesDispatchModelConfig: snapshot.Group.MessagesDispatchModelConfig,
ModelsListConfig: snapshot.Group.ModelsListConfig,
ModelAllowlist: snapshot.Group.ModelAllowlist,
CodexModelsManifestConfig: snapshot.Group.CodexModelsManifestConfig,
RPMLimit: snapshot.Group.RPMLimit,
MaxReasoningEffort: snapshot.Group.MaxReasoningEffort,
@@ -2,7 +2,7 @@ package service
import "testing"
func TestAPIKeyService_RejectsV10AuthSnapshotWithoutModelsListConfig(t *testing.T) {
func TestAPIKeyService_RejectsV10AuthSnapshotWithoutModelAllowlist(t *testing.T) {
groupID := int64(9)
svc := &APIKeyService{}
@@ -35,7 +35,7 @@ func TestAPIKeyService_RejectsV10AuthSnapshotWithoutModelsListConfig(t *testing.
t.Fatalf("expected stale snapshot to be ignored without error, got %v", err)
}
if ok {
t.Fatalf("expected v10 auth snapshot to be rejected after models_list_config was added")
t.Fatalf("expected v10 auth snapshot to be rejected after model_allowlist was added")
}
if apiKey != nil {
t.Fatalf("expected no API key from stale snapshot, got %#v", apiKey)
@@ -769,9 +769,9 @@ func TestHasUserPlatformQuotaLimit(t *testing.T) {
daily := 5.0
tests := []struct {
name string
setup func() *BillingCacheService
want bool
name string
setup func() *BillingCacheService
want bool
}{
{
name: "has_limit",
@@ -252,7 +252,7 @@ func TestGrokOAuthServiceRefreshAccountTokenIgnoresIDTokenTierWhenAccessTokenHas
svc := NewGrokOAuthService(nil, &grokOAuthClientStub{
refreshResponse: &xai.TokenResponse{
AccessToken: "opaque-access-token",
IDToken: makeGrokOAuthJWT(map[string]any{"tier": 5}),
IDToken: makeGrokOAuthJWT(map[string]any{"tier": 5}),
TokenType: "Bearer",
ExpiresIn: 3600,
},
+1 -2
View File
@@ -12,7 +12,6 @@ import (
)
type OpenAIMessagesDispatchModelConfig = domain.OpenAIMessagesDispatchModelConfig
type GroupModelsListConfig = domain.GroupModelsListConfig
type GroupCodexModelsManifestConfig = domain.GroupCodexModelsManifestConfig
type ReasoningEffortMapping = domain.ReasoningEffortMapping
@@ -107,7 +106,7 @@ type Group struct {
RequirePrivacySet bool // 调度时仅允许 privacy 已成功设置的账号(OpenAI/Antigravity/Anthropic/Gemini)
DefaultMappedModel string
MessagesDispatchModelConfig OpenAIMessagesDispatchModelConfig
ModelsListConfig GroupModelsListConfig
ModelAllowlist GroupModelAllowlist
// CodexModelsManifestConfig 开启后,该分组的 Codex /models manifest 请求只用
// 固定账号列表拉取并合并,不经过调度器(仅 openai 平台)。
CodexModelsManifestConfig GroupCodexModelsManifestConfig
@@ -0,0 +1,211 @@
package service
import (
"net/http"
"strings"
"github.com/Wei-Shaw/sub2api/internal/domain"
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
)
// GroupModelAllowlist 是 service 层的分组模型白名单(与 domain.GroupModelAllowlist
// 字段一致,ent 持久化用 domain 类型,边界处显式转换)。
type GroupModelAllowlist struct {
Enabled bool `json:"enabled"`
Models []string `json:"models,omitempty"`
}
// DomainGroupModelAllowlist 把 service 白名单转换为 ent 持久化使用的 domain 类型。
func DomainGroupModelAllowlist(cfg GroupModelAllowlist) domain.GroupModelAllowlist {
return domain.GroupModelAllowlist{Enabled: cfg.Enabled, Models: cfg.Models}
}
// GroupModelAllowlistFromDomain 把 ent 读出的 domain 白名单转换为 service 类型。
func GroupModelAllowlistFromDomain(cfg domain.GroupModelAllowlist) GroupModelAllowlist {
return GroupModelAllowlist{Enabled: cfg.Enabled, Models: cfg.Models}
}
// normalizeGroupModelAllowlist 归一化管理端提交的分组模型白名单:
// 条目 TrimSpace、按小写去重保序;`*` 只允许出现在条目末尾;
// enabled=true 且列表为空视为配置错误,返回 400 而不是运行时静默放行/拒绝。
func normalizeGroupModelAllowlist(cfg GroupModelAllowlist) (GroupModelAllowlist, error) {
out := GroupModelAllowlist{Enabled: cfg.Enabled}
if len(cfg.Models) == 0 {
if out.Enabled {
return out, infraerrors.New(http.StatusBadRequest, "INVALID_MODEL_ALLOWLIST", "model allowlist cannot be enabled with an empty model list")
}
return out, nil
}
seen := make(map[string]struct{}, len(cfg.Models))
out.Models = make([]string, 0, len(cfg.Models))
for _, model := range cfg.Models {
model = strings.TrimSpace(model)
if model == "" {
continue
}
if strings.Contains(strings.TrimSuffix(model, "*"), "*") {
return out, infraerrors.New(http.StatusBadRequest, "INVALID_MODEL_ALLOWLIST", `wildcard "*" is only allowed at the end of an allowlist entry`)
}
key := strings.ToLower(model)
if _, ok := seen[key]; ok {
continue
}
seen[key] = struct{}{}
out.Models = append(out.Models, model)
}
if len(out.Models) == 0 {
if out.Enabled {
return out, infraerrors.New(http.StatusBadRequest, "INVALID_MODEL_ALLOWLIST", "model allowlist cannot be enabled with an empty model list")
}
out.Models = nil
}
return out, nil
}
// ModelAllowlistEnabled 报告该分组是否启用了模型白名单。
// 开启后所有携带模型的网关请求与模型列表接口都受白名单约束。
func (g *Group) ModelAllowlistEnabled() bool {
return g != nil && g.ModelAllowlist.Enabled
}
// Allows 判断客户端请求的模型是否命中白名单。
// 准入只看客户端书写的模型名,与账号映射、渠道映射、合成路由改写无关;
// 候选形式覆盖代码中已有的模型名等价规则(Gemini models/ 前缀、
// Antigravity/Claude -thinking 宽容规则、OpenAI 推理后缀),不做模糊匹配。
func (a GroupModelAllowlist) Allows(model string) bool {
if !a.Enabled {
return true
}
model = strings.TrimSpace(model)
if model == "" {
return true
}
candidates := groupModelAllowlistCandidates(model)
for _, entry := range a.Models {
entry = strings.TrimSpace(entry)
if entry == "" {
continue
}
entry = strings.ToLower(entry)
if strings.HasSuffix(entry, "*") {
prefix := strings.TrimSuffix(entry, "*")
for _, candidate := range candidates {
if strings.HasPrefix(candidate, prefix) {
return true
}
}
continue
}
for _, candidate := range candidates {
if candidate == entry {
return true
}
}
}
return false
}
// groupModelAllowlistCandidates 返回客户端模型名在白名单匹配中的候选形式(均已小写)。
func groupModelAllowlistCandidates(model string) []string {
candidates := make([]string, 0, 4)
add := func(value string) {
value = strings.ToLower(strings.TrimSpace(value))
if value == "" {
return
}
for _, existing := range candidates {
if existing == value {
return
}
}
candidates = append(candidates, value)
}
add(model)
add(strings.TrimPrefix(model, "models/"))
add(claude.NormalizeModelID(strings.TrimSuffix(model, "-thinking")))
add(NormalizeOpenAICompatRequestedModel(model))
return candidates
}
// FilterForListing 按白名单条目顺序生成模型列表输出。
// 精确条目沿用既有规则:只有出现在 source(账号映射键 ∪ 平台默认列表)的模式
// 集合中才输出;通配条目展开为 source 中所有匹配项并保持 source 顺序;全局去重。
func (a GroupModelAllowlist) FilterForListing(source []string) []string {
if !a.Enabled {
return source
}
if len(a.Models) == 0 {
// 开启但为空的配置在管理端已被拒绝;对遗留脏数据保持“看到什么 = 能调什么”,
// 列表与准入同时返回空。
return nil
}
patterns := make([]string, 0, len(source))
for _, pattern := range source {
pattern = strings.TrimSpace(pattern)
if pattern != "" {
patterns = append(patterns, pattern)
}
}
if len(patterns) == 0 {
return nil
}
seen := make(map[string]struct{}, len(a.Models)+len(patterns))
filtered := make([]string, 0, len(patterns))
add := func(model string) {
key := strings.ToLower(model)
if _, ok := seen[key]; ok {
return
}
seen[key] = struct{}{}
filtered = append(filtered, model)
}
for _, entry := range a.Models {
entry = strings.TrimSpace(entry)
if entry == "" {
continue
}
if strings.HasSuffix(entry, "*") {
// 空 前缀(裸 `*` 条目)匹配全部来源,与 Allows 的全放行语义一致。
prefix := strings.ToLower(strings.TrimSuffix(entry, "*"))
for _, pattern := range patterns {
if strings.HasPrefix(strings.ToLower(pattern), prefix) {
add(pattern)
}
}
continue
}
if allowlistSourcePatternAllowsModel(patterns, entry) {
add(entry)
}
}
return filtered
}
// allowlistSourcePatternAllowsModel 沿用 filterModelsByCustomList 时代的匹配规则:
// 精确相等、source 通配模式的前缀匹配,以及 Claude 归一化(-thinking 后缀)
// 后的精确匹配。比较与 Allows 一律大小写不敏感,保证「准入允许 ⇒ 列表可见」。
func allowlistSourcePatternAllowsModel(patterns []string, model string) bool {
for _, pattern := range patterns {
if strings.EqualFold(pattern, model) {
return true
}
if strings.HasSuffix(pattern, "*") && strings.HasPrefix(strings.ToLower(model), strings.ToLower(strings.TrimSuffix(pattern, "*"))) {
return true
}
}
normalizedClaudeModel := claude.NormalizeModelID(strings.TrimSuffix(model, "-thinking"))
if !strings.EqualFold(normalizedClaudeModel, model) {
for _, pattern := range patterns {
if strings.EqualFold(pattern, normalizedClaudeModel) {
return true
}
}
}
return false
}
@@ -0,0 +1,254 @@
package service
import (
"strings"
"testing"
)
func TestNormalizeGroupModelAllowlist(t *testing.T) {
tests := []struct {
name string
in GroupModelAllowlist
want GroupModelAllowlist
wantErr string
}{
{
name: "trims and dedupes case-insensitively preserving order",
in: GroupModelAllowlist{
Enabled: true,
Models: []string{" gpt-5.4 ", "GPT-5.4", "claude-sonnet-4.5", ""},
},
want: GroupModelAllowlist{Enabled: true, Models: []string{"gpt-5.4", "claude-sonnet-4.5"}},
},
{
name: "disabled with empty list is fine",
in: GroupModelAllowlist{Enabled: false},
want: GroupModelAllowlist{Enabled: false},
},
{
name: "enabled with empty list is rejected",
in: GroupModelAllowlist{Enabled: true},
wantErr: "INVALID_MODEL_ALLOWLIST",
},
{
name: "enabled with only blank entries is rejected",
in: GroupModelAllowlist{Enabled: true, Models: []string{" ", ""}},
wantErr: "INVALID_MODEL_ALLOWLIST",
},
{
name: "trailing wildcard is accepted",
in: GroupModelAllowlist{Enabled: true, Models: []string{"gpt-5.5-*"}},
want: GroupModelAllowlist{Enabled: true, Models: []string{"gpt-5.5-*"}},
},
{
name: "wildcard in the middle is rejected",
in: GroupModelAllowlist{Enabled: true, Models: []string{"gpt-*-5.4"}},
wantErr: "INVALID_MODEL_ALLOWLIST",
},
{
name: "bare wildcard is accepted as allow-all",
in: GroupModelAllowlist{Enabled: true, Models: []string{"*"}},
want: GroupModelAllowlist{Enabled: true, Models: []string{"*"}},
wantErr: "",
},
{
name: "disabled config with invalid wildcard still rejected",
in: GroupModelAllowlist{Enabled: false, Models: []string{"foo-*bar"}},
wantErr: "INVALID_MODEL_ALLOWLIST",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := normalizeGroupModelAllowlist(tt.in)
if tt.wantErr != "" {
if err == nil || !strings.Contains(err.Error(), tt.wantErr) {
t.Fatalf("expected error containing %q, got %v", tt.wantErr, err)
}
return
}
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if got.Enabled != tt.want.Enabled {
t.Fatalf("enabled mismatch: got %v want %v", got.Enabled, tt.want.Enabled)
}
if len(got.Models) != len(tt.want.Models) {
t.Fatalf("models mismatch: got %#v want %#v", got.Models, tt.want.Models)
}
for i := range tt.want.Models {
if got.Models[i] != tt.want.Models[i] {
t.Fatalf("models[%d] mismatch: got %q want %q", i, got.Models[i], tt.want.Models[i])
}
}
})
}
}
func TestGroupModelAllowlistAllows(t *testing.T) {
allowlist := GroupModelAllowlist{
Enabled: true,
Models: []string{"claude-sonnet-4.5", "gemini-2.5-pro", "gpt-5.5", "grok-*"},
}
tests := []struct {
name string
model string
want bool
}{
{name: "exact match", model: "claude-sonnet-4.5", want: true},
{name: "case-insensitive entry match", model: "Claude-Sonnet-4.5", want: true},
{name: "not listed", model: "claude-opus-4.6", want: false},
{name: "thinking suffix tolerated via claude normalization", model: "claude-sonnet-4.5-thinking", want: true},
{name: "gemini models/ prefix stripped", model: "models/gemini-2.5-pro", want: true},
{name: "gemini models/ prefix not in list", model: "models/gemini-2.5-flash", want: false},
{name: "openai reasoning suffix normalizes to base model", model: "gpt-5.5-codex-high", want: true},
{name: "openai reasoning suffix on unlisted model", model: "gpt-4.1-low", want: false},
{name: "trailing wildcard prefix match", model: "grok-4.6", want: true},
{name: "trailing wildcard requires prefix", model: "grok", want: false},
{name: "empty model passes (handler decides required-ness)", model: "", want: true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := allowlist.Allows(tt.model); got != tt.want {
t.Fatalf("Allows(%q) = %v, want %v", tt.model, got, tt.want)
}
})
}
t.Run("disabled allowlist allows everything", func(t *testing.T) {
disabled := GroupModelAllowlist{Enabled: false, Models: []string{"only-model"}}
if !disabled.Allows("anything") {
t.Fatal("disabled allowlist must allow all models")
}
})
t.Run("enabled empty allowlist denies everything", func(t *testing.T) {
empty := GroupModelAllowlist{Enabled: true}
if empty.Allows("anything") {
t.Fatal("enabled empty allowlist must deny all models")
}
})
t.Run("bare wildcard allows everything", func(t *testing.T) {
all := GroupModelAllowlist{Enabled: true, Models: []string{"*"}}
if !all.Allows("claude-opus-4.6") || !all.Allows("models/gemini-2.5-pro") {
t.Fatal("bare wildcard must allow all models")
}
})
}
func TestGroupModelAllowlistEnabled(t *testing.T) {
var nilGroup *Group
if nilGroup.ModelAllowlistEnabled() {
t.Fatal("nil group must not report enabled")
}
group := &Group{}
if group.ModelAllowlistEnabled() {
t.Fatal("default group must not report enabled")
}
group.ModelAllowlist = GroupModelAllowlist{Enabled: true, Models: []string{"m"}}
if !group.ModelAllowlistEnabled() {
t.Fatal("enabled group must report enabled")
}
}
func TestGroupModelAllowlistFilterForListing(t *testing.T) {
source := []string{"claude-opus-4.6", "claude-sonnet-4.5", "gpt-5.4", "gpt-5.5-codex", "gpt-5.5-mini", "grok-4.6"}
t.Run("passthrough when disabled", func(t *testing.T) {
disabled := GroupModelAllowlist{Enabled: false, Models: []string{"gpt-5.4"}}
got := disabled.FilterForListing(source)
if len(got) != len(source) {
t.Fatalf("disabled allowlist must not filter, got %#v", got)
}
})
t.Run("exact entries keep entry order and require source membership", func(t *testing.T) {
cfg := GroupModelAllowlist{Enabled: true, Models: []string{"gpt-5.5-codex", "claude-sonnet-4.5", "claude-haiku-4.5"}}
got := cfg.FilterForListing(source)
want := []string{"gpt-5.5-codex", "claude-sonnet-4.5"}
if len(got) != len(want) || got[0] != want[0] || got[1] != want[1] {
t.Fatalf("got %#v want %#v", got, want)
}
})
t.Run("bare wildcard expands to the full source list", func(t *testing.T) {
cfg := GroupModelAllowlist{Enabled: true, Models: []string{"*"}}
got := cfg.FilterForListing([]string{"gpt-5.4"})
if strings.Join(got, ",") != "gpt-5.4" {
t.Fatalf("bare wildcard must expose the full source, got %#v", got)
}
})
t.Run("wildcard entries expand to source order", func(t *testing.T) {
cfg := GroupModelAllowlist{Enabled: true, Models: []string{"claude-*", "gpt-5.5-*"}}
got := cfg.FilterForListing(source)
want := []string{"claude-opus-4.6", "claude-sonnet-4.5", "gpt-5.5-codex", "gpt-5.5-mini"}
if strings.Join(got, ",") != strings.Join(want, ",") {
t.Fatalf("got %#v want %#v", got, want)
}
})
t.Run("mixed exact and wildcard dedupes globally", func(t *testing.T) {
cfg := GroupModelAllowlist{Enabled: true, Models: []string{"gpt-*", "gpt-5.4"}}
got := cfg.FilterForListing(source)
want := []string{"gpt-5.4", "gpt-5.5-codex", "gpt-5.5-mini"}
if strings.Join(got, ",") != strings.Join(want, ",") {
t.Fatalf("got %#v want %#v", got, want)
}
})
t.Run("thinking-tolerant exact entries resolve against source", func(t *testing.T) {
cfg := GroupModelAllowlist{Enabled: true, Models: []string{"claude-sonnet-4.5-thinking"}}
got := cfg.FilterForListing(source)
want := []string{"claude-sonnet-4.5-thinking"}
if strings.Join(got, ",") != strings.Join(want, ",") {
t.Fatalf("got %#v want %#v", got, want)
}
})
t.Run("exact entries match case-insensitively against source", func(t *testing.T) {
// 归一化保留首次出现的拼写:条目大写、来源小写也必须命中,
// 与 Allows 的大小写不敏感语义保持「准入允许 ⇒ 列表可见」。
cfg := GroupModelAllowlist{Enabled: true, Models: []string{"GPT-5.4"}}
got := cfg.FilterForListing([]string{"gpt-5.4"})
want := []string{"GPT-5.4"}
if strings.Join(got, ",") != strings.Join(want, ",") {
t.Fatalf("got %#v want %#v", got, want)
}
})
t.Run("wildcard source patterns match case-insensitively", func(t *testing.T) {
cfg := GroupModelAllowlist{Enabled: true, Models: []string{"claude-sonnet-4.20250514"}}
got := cfg.FilterForListing([]string{"Claude-Sonnet-4*"})
want := []string{"claude-sonnet-4.20250514"}
if strings.Join(got, ",") != strings.Join(want, ",") {
t.Fatalf("got %#v want %#v", got, want)
}
})
t.Run("wildcard source patterns whitelist exact entries", func(t *testing.T) {
cfg := GroupModelAllowlist{Enabled: true, Models: []string{"claude-sonnet-4.20250514"}}
got := cfg.FilterForListing([]string{"claude-sonnet-4*"})
want := []string{"claude-sonnet-4.20250514"}
if strings.Join(got, ",") != strings.Join(want, ",") {
t.Fatalf("got %#v want %#v", got, want)
}
})
t.Run("empty source yields empty output", func(t *testing.T) {
cfg := GroupModelAllowlist{Enabled: true, Models: []string{"gpt-5.4"}}
if got := cfg.FilterForListing(nil); len(got) != 0 {
t.Fatalf("expected empty output, got %#v", got)
}
})
t.Run("enabled empty config yields empty output", func(t *testing.T) {
cfg := GroupModelAllowlist{Enabled: true}
if got := cfg.FilterForListing(source); len(got) != 0 {
t.Fatalf("expected empty output for enabled empty config, got %#v", got)
}
})
}
@@ -1,32 +0,0 @@
package service
import "strings"
func normalizeGroupModelsListConfig(cfg GroupModelsListConfig) GroupModelsListConfig {
out := GroupModelsListConfig{Enabled: cfg.Enabled}
if len(cfg.Models) == 0 {
return out
}
seen := make(map[string]struct{}, len(cfg.Models))
out.Models = make([]string, 0, len(cfg.Models))
for _, model := range cfg.Models {
model = strings.TrimSpace(model)
if model == "" {
continue
}
if _, ok := seen[model]; ok {
continue
}
seen[model] = struct{}{}
out.Models = append(out.Models, model)
}
if len(out.Models) == 0 {
out.Models = nil
}
return out
}
func (g *Group) CustomModelsListEnabled() bool {
return g != nil && g.ModelsListConfig.Enabled && len(g.ModelsListConfig.Models) > 0
}
@@ -48,14 +48,16 @@ const (
// FilterCodexModelIDsForGroup removes dedicated media-generation models,
// wildcard mapping keys, and Codex automatic modes from a client catalog.
// Automatic modes are retained only when the group's enabled custom model list
// Automatic modes are retained only when the group's enabled model allowlist
// explicitly selects the exact slug; account model mappings describe routing
// and are not feature opt-ins. Wildcard keys such as "foo-*" are routing
// patterns, not concrete Codex models.
// patterns, not concrete Codex models. When the allowlist is enabled the
// catalog is additionally restricted by FilterForListing (wildcard entries
// expand against the catalog).
func FilterCodexModelIDsForGroup(modelIDs []string, group *Group) []string {
explicitlyEnabled := make(map[string]struct{})
if group != nil && group.CustomModelsListEnabled() {
for _, modelID := range group.ModelsListConfig.Models {
if group != nil && group.ModelAllowlistEnabled() {
for _, modelID := range group.ModelAllowlist.Models {
modelID = strings.TrimSpace(modelID)
if strings.HasPrefix(modelID, codexAutoModelPrefix) {
explicitlyEnabled[modelID] = struct{}{}
@@ -82,6 +84,9 @@ func FilterCodexModelIDsForGroup(modelIDs []string, group *Group) []string {
}
filtered = append(filtered, modelID)
}
if group != nil && group.ModelAllowlistEnabled() {
filtered = group.ModelAllowlist.FilterForListing(filtered)
}
return filtered
}
@@ -145,8 +150,8 @@ func (s *OpenAIGatewayService) BuildGroupConfiguredCodexModelsManifest(
body, _, err = mergeConfiguredCodexModelsManifest(
body,
nil,
group.ModelsListConfig.Models,
group.CustomModelsListEnabled(),
group.ModelAllowlist.Models,
group.ModelAllowlistEnabled(),
)
if err != nil {
return nil, false, fmt.Errorf("build group configured Codex models: %w", err)
@@ -186,8 +191,8 @@ func (s *OpenAIGatewayService) MergeGroupConfiguredCodexModels(
body, changed, err := mergeConfiguredCodexModelsManifest(
manifest.Body,
configuredModels,
group.ModelsListConfig.Models,
group.CustomModelsListEnabled(),
group.ModelAllowlist.Models,
group.ModelAllowlistEnabled(),
)
if err != nil {
return fmt.Errorf("merge group configured Codex models: %w", err)
@@ -278,15 +283,15 @@ func openAIConfiguredCodexModelIDs(accounts []Account) []string {
func openAIConfiguredCodexModelIDsForGroup(accounts []Account, group *Group) []string {
models := openAIConfiguredCodexModelIDs(accounts)
if group == nil || !group.CustomModelsListEnabled() {
if group == nil || !group.ModelAllowlistEnabled() {
return models
}
seen := make(map[string]struct{}, len(models)+len(group.ModelsListConfig.Models))
seen := make(map[string]struct{}, len(models)+len(group.ModelAllowlist.Models))
for _, modelID := range models {
seen[modelID] = struct{}{}
}
for _, selectedModel := range group.ModelsListConfig.Models {
for _, selectedModel := range group.ModelAllowlist.Models {
selectedModel = strings.TrimSpace(selectedModel)
if selectedModel == "" || strings.Contains(selectedModel, "*") {
continue
@@ -1226,6 +1231,8 @@ func mergeConfiguredCodexModelsManifest(
selected[modelID] = struct{}{}
}
}
// 白名单条目匹配统一走 GroupModelAllowlist.Allows(通配条目按前缀展开)。
allowlist := GroupModelAllowlist{Enabled: filterBySelection, Models: selectedModels}
seen := make(map[string]struct{}, len(upstreamModels)+len(configuredModels))
merged := make([]json.RawMessage, 0, len(upstreamModels)+len(configuredModels))
changed := false
@@ -1246,11 +1253,9 @@ func mergeConfiguredCodexModelsManifest(
changed = true
continue
}
if filterBySelection {
if _, allowed := selected[descriptor.Slug]; !allowed {
changed = true
continue
}
if filterBySelection && !allowlist.Allows(descriptor.Slug) {
changed = true
continue
}
if strings.HasPrefix(descriptor.Slug, codexAutoModelPrefix) {
_, explicitlyEnabled := selected[descriptor.Slug]
@@ -1274,10 +1279,8 @@ func mergeConfiguredCodexModelsManifest(
if isCodexDedicatedMediaModel(modelID) {
continue
}
if filterBySelection {
if _, allowed := selected[modelID]; !allowed {
continue
}
if filterBySelection && !allowlist.Allows(modelID) {
continue
}
if strings.HasPrefix(modelID, codexAutoModelPrefix) {
if _, explicitlyEnabled := selected[modelID]; !filterBySelection || !explicitlyEnabled {
@@ -1136,7 +1136,7 @@ func TestBuildGroupConfiguredCodexModelsManifestExpandsSelectedModelCoveredByWil
group := &Group{
ID: groupID,
Platform: PlatformOpenAI,
ModelsListConfig: GroupModelsListConfig{
ModelAllowlist: GroupModelAllowlist{
Enabled: true,
Models: []string{"gpt-5.6"},
},
@@ -1328,7 +1328,7 @@ func TestMergeGroupConfiguredCodexModelsKeepsExplicitAutoReviewSelection(t *test
group := &Group{
ID: groupID,
Platform: PlatformOpenAI,
ModelsListConfig: GroupModelsListConfig{
ModelAllowlist: GroupModelAllowlist{
Enabled: true,
Models: []string{openai.CodexUsageProbeModel},
},
@@ -1360,7 +1360,7 @@ func TestMergeGroupConfiguredCodexModelsHonorsCustomListAndFinalETag(t *testing.
group := &Group{
ID: groupID,
Platform: PlatformOpenAI,
ModelsListConfig: GroupModelsListConfig{
ModelAllowlist: GroupModelAllowlist{
Enabled: true,
Models: []string{"deepseek-4-pro"},
},
@@ -2627,7 +2627,7 @@ func TestFetchCodexModelsManifestAPIKeyCacheSurvivesClientMutation(t *testing.T)
&Group{
ID: 81,
Platform: PlatformOpenAI,
ModelsListConfig: GroupModelsListConfig{
ModelAllowlist: GroupModelAllowlist{
Enabled: true,
Models: []string{"model-a"},
},
@@ -2654,7 +2654,7 @@ func TestFetchCodexModelsManifestAPIKeyCacheSurvivesClientMutation(t *testing.T)
&Group{
ID: 82,
Platform: PlatformOpenAI,
ModelsListConfig: GroupModelsListConfig{
ModelAllowlist: GroupModelAllowlist{
Enabled: true,
Models: []string{"model-b"},
},
@@ -3494,8 +3494,8 @@ func TestFetchCodexModelsManifestOAuthSharedAcrossGroupsWithIndependentFiltering
},
}
account := newCodexModelsTestAccount()
groupA := &Group{ID: 91, Platform: PlatformOpenAI, ModelsListConfig: GroupModelsListConfig{Enabled: true, Models: []string{"model-a"}}}
groupB := &Group{ID: 92, Platform: PlatformOpenAI, ModelsListConfig: GroupModelsListConfig{Enabled: true, Models: []string{"model-b"}}}
groupA := &Group{ID: 91, Platform: PlatformOpenAI, ModelAllowlist: GroupModelAllowlist{Enabled: true, Models: []string{"model-a"}}}
groupB := &Group{ID: 92, Platform: PlatformOpenAI, ModelAllowlist: GroupModelAllowlist{Enabled: true, Models: []string{"model-b"}}}
begin := make(chan struct{})
type result struct {
@@ -122,8 +122,8 @@ func TestShouldClearStickySession(t *testing.T) {
{
name: "overloaded account",
account: &Account{
Status: StatusActive,
Schedulable: true,
Status: StatusActive,
Schedulable: true,
OverloadUntil: &future,
},
requestedModel: "",
@@ -0,0 +1,22 @@
-- 235: groups.models_list_config 更名为 model_allowlist,语义从「仅影响 /v1/models 展示」
-- 升级为分组级模型白名单(同时约束模型列表接口与请求准入)。数据原样保留。
DO $$
BEGIN
IF EXISTS (
SELECT 1 FROM information_schema.columns
WHERE table_schema = 'public'
AND table_name = 'groups'
AND column_name = 'models_list_config'
) AND NOT EXISTS (
SELECT 1 FROM information_schema.columns
WHERE table_schema = 'public'
AND table_name = 'groups'
AND column_name = 'model_allowlist'
) THEN
ALTER TABLE groups RENAME COLUMN models_list_config TO model_allowlist;
END IF;
END
$$;
COMMENT ON COLUMN groups.model_allowlist IS
'Group model allowlist: constrains both model listing responses and request admission';