diff --git a/controller/channel.go b/controller/channel.go index 0531927..7f9553b 100644 --- a/controller/channel.go +++ b/controller/channel.go @@ -19,6 +19,7 @@ import ( "github.com/QuantumNous/new-api/service" "github.com/gin-gonic/gin" + "gorm.io/gorm" ) type OpenAIModel struct { @@ -68,6 +69,30 @@ func clearChannelInfo(channel *model.Channel) { } } +func attachChannelAssetCredentialSummaries(channels []*model.Channel) error { + ids := make([]int, 0) + for _, channel := range channels { + if channel != nil && channel.Type == constant.ChannelTypeChinaMobileSeedance { + ids = append(ids, channel.Id) + } + } + summaries, err := model.GetChannelAssetCredentialSummaries(ids) + if err != nil { + return err + } + for _, channel := range channels { + if channel == nil || channel.Type != constant.ChannelTypeChinaMobileSeedance { + continue + } + summary, ok := summaries[channel.Id] + channel.AssetCredentialConfigured = ok + if ok { + channel.AssetCredentialPoolID = summary.PoolID + } + } + return nil +} + func GetAllChannels(c *gin.Context) { pageInfo := common.GetPageQuery(c) channelData := make([]*model.Channel, 0) @@ -144,6 +169,10 @@ func GetAllChannels(c *gin.Context) { } } + if err := attachChannelAssetCredentialSummaries(channelData); err != nil { + common.ApiError(c, err) + return + } for _, datum := range channelData { clearChannelInfo(datum) } @@ -482,6 +511,10 @@ func SearchChannels(c *gin.Context) { pagedData := channelData[startIdx:endIdx] + if err := attachChannelAssetCredentialSummaries(pagedData); err != nil { + common.ApiError(c, err) + return + } for _, datum := range pagedData { clearChannelInfo(datum) } @@ -510,6 +543,10 @@ func GetChannel(c *gin.Context) { return } if channel != nil { + if err := attachChannelAssetCredentialSummaries([]*model.Channel{channel}); err != nil { + common.ApiError(c, err) + return + } clearChannelInfo(channel) } c.JSON(http.StatusOK, gin.H{ @@ -675,10 +712,29 @@ func RefreshCodexChannelCredential(c *gin.Context) { } type AddChannelRequest struct { - Mode string `json:"mode"` - MultiKeyMode constant.MultiKeyMode `json:"multi_key_mode"` - BatchAddSetKeyPrefix2Name bool `json:"batch_add_set_key_prefix_2_name"` - Channel *model.Channel `json:"channel"` + Mode string `json:"mode"` + MultiKeyMode constant.MultiKeyMode `json:"multi_key_mode"` + BatchAddSetKeyPrefix2Name bool `json:"batch_add_set_key_prefix_2_name"` + Channel *model.Channel `json:"channel"` + AssetCredential *ChannelAssetCredentialInput `json:"asset_credential"` +} + +type ChannelAssetCredentialInput struct { + AccessKey string `json:"access_key"` + SecretKey string `json:"secret_key"` + PoolID string `json:"pool_id"` +} + +func channelAssetCredentialFromInput(channelType int, input *ChannelAssetCredentialInput) (*model.ChannelAssetCredential, error) { + if channelType != constant.ChannelTypeChinaMobileSeedance || input == nil { + return nil, nil + } + ak := strings.TrimSpace(input.AccessKey) + sk := strings.TrimSpace(input.SecretKey) + if ak == "" || sk == "" { + return nil, errors.New("移动云素材 AccessKey 和 SecretKey 必须同时填写") + } + return &model.ChannelAssetCredential{AccessKey: ak, SecretKey: sk, PoolID: strings.TrimSpace(input.PoolID)}, nil } func getVertexArrayKeys(keys string) ([]string, error) { @@ -729,6 +785,15 @@ func AddChannel(c *gin.Context) { }) return } + credential, err := channelAssetCredentialFromInput(addChannelRequest.Channel.Type, addChannelRequest.AssetCredential) + if err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()}) + return + } + if credential != nil && addChannelRequest.Mode == "batch" { + c.JSON(http.StatusOK, gin.H{"success": false, "message": "移动云素材凭证不支持批量添加渠道"}) + return + } addChannelRequest.Channel.CreatedTime = common.GetTimestamp() keys := make([]string, 0) @@ -800,7 +865,15 @@ func AddChannel(c *gin.Context) { } channels = append(channels, *localChannel) } - err = model.BatchInsertChannels(channels) + if credential != nil { + if len(channels) != 1 { + c.JSON(http.StatusOK, gin.H{"success": false, "message": "移动云素材凭证仅支持创建一个渠道"}) + return + } + err = model.InsertChannelWithAssetCredential(&channels[0], credential) + } else { + err = model.BatchInsertChannels(channels) + } if err != nil { common.ApiError(c, err) return @@ -985,8 +1058,9 @@ func DeleteChannelBatch(c *gin.Context) { type PatchChannel struct { model.Channel - MultiKeyMode *string `json:"multi_key_mode"` - KeyMode *string `json:"key_mode"` // 多key模式下密钥覆盖或者追加 + MultiKeyMode *string `json:"multi_key_mode"` + KeyMode *string `json:"key_mode"` // 多key模式下密钥覆盖或者追加 + AssetCredential *ChannelAssetCredentialInput `json:"asset_credential"` } func UpdateChannel(c *gin.Context) { @@ -1103,7 +1177,22 @@ func UpdateChannel(c *gin.Context) { // 覆盖模式：直接使用新密钥（默认行为，不需要特殊处理） } } - err = channel.Update() + credential, err := channelAssetCredentialFromInput(channel.Type, channel.AssetCredential) + if err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()}) + return + } + if credential != nil { + err = model.DB.Transaction(func(tx *gorm.DB) error { + if err := channel.UpdateWithTx(tx); err != nil { + return err + } + credential.ChannelId = channel.Id + return model.UpsertChannelAssetCredentialWithTx(tx, credential) + }) + } else { + err = channel.Update() + } if err != nil { common.ApiError(c, err) return @@ -1112,6 +1201,10 @@ func UpdateChannel(c *gin.Context) { service.ResetProxyClientCache() channel.Key = "" clearChannelInfo(&channel.Channel) + if err := attachChannelAssetCredentialSummaries([]*model.Channel{&channel.Channel}); err != nil { + common.ApiError(c, err) + return + } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", @@ -2105,4 +2198,3 @@ func OllamaVersion(c *gin.Context) { }, }) } - diff --git a/controller/channel_asset_credential_test.go b/controller/channel_asset_credential_test.go new file mode 100644 index 0000000..f87a56c --- /dev/null +++ b/controller/channel_asset_credential_test.go @@ -0,0 +1,50 @@ +package controller + +import ( + "testing" + + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" + "gorm.io/gorm/logger" +) + +func setupChannelAssetCredentialControllerDB(t *testing.T) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open("file:controller_channel_asset_credentials?mode=memory&cache=shared"), &gorm.Config{ + Logger: logger.Default.LogMode(logger.Silent), + }) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + originalDB := model.DB + model.DB = db + require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.ChannelAssetCredential{})) + t.Cleanup(func() { + model.DB = originalDB + require.NoError(t, sqlDB.Close()) + }) + return db +} + +func TestAttachChannelAssetCredentialSummariesDoesNotExposeSecrets(t *testing.T) { + db := setupChannelAssetCredentialControllerDB(t) + channel := &model.Channel{Id: 61, Type: constant.ChannelTypeChinaMobileSeedance, Key: "video-key", Name: "channel"} + require.NoError(t, db.Create(channel).Error) + require.NoError(t, model.UpsertChannelAssetCredential(&model.ChannelAssetCredential{ + ChannelId: 61, + AccessKey: "ak-secret", + SecretKey: "sk-secret", + PoolID: "pool-61", + })) + + require.NoError(t, attachChannelAssetCredentialSummaries([]*model.Channel{channel})) + assert.True(t, channel.AssetCredentialConfigured) + assert.Equal(t, "pool-61", channel.AssetCredentialPoolID) + assert.NotContains(t, channel.Key, "ak-secret") + assert.NotContains(t, channel.Key, "sk-secret") +} diff --git a/controller/doubao_asset_real_e2e_test.go b/controller/doubao_asset_real_e2e_test.go index 48bf769..58cd68c 100644 --- a/controller/doubao_asset_real_e2e_test.go +++ b/controller/doubao_asset_real_e2e_test.go @@ -32,6 +32,12 @@ func TestChinaMobileAssetControllerWithRealUpstream(t *testing.T) { } db := setupRealAssetE2EDB(t) createDoubaoAssetProxyChannel(t, db, 7101, constant.ChannelTypeChinaMobileSeedance, "default", "video-generation-key", common.ChannelStatusEnabled) + require.NoError(t, model.UpsertChannelAssetCredential(&model.ChannelAssetCredential{ + ChannelId: 7101, + AccessKey: ak, + SecretKey: sk, + PoolID: poolID, + })) router := setupRealAssetE2ERouter() suffix := time.Now().Format("20060102150405") groupName := "new-api-e2e-" + suffix @@ -119,7 +125,7 @@ func setupRealAssetE2EDB(t *testing.T) *gorm.DB { sqlDB, err := db.DB() require.NoError(t, err) sqlDB.SetMaxOpenConns(1) - require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.UserAssetChannel{})) + require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.UserAssetChannel{}, &model.ChannelAssetCredential{})) t.Cleanup(func() { _ = sqlDB.Close() diff --git a/model/channel.go b/model/channel.go index e16a544..bba7752 100644 --- a/model/channel.go +++ b/model/channel.go @@ -56,6 +56,10 @@ type Channel struct { // cache info Keys []string `json:"-" gorm:"-"` + + // Asset credential metadata is populated only for management API responses. + AssetCredentialConfigured bool `json:"asset_credential_configured,omitempty" gorm:"-"` + AssetCredentialPoolID string `json:"asset_credential_pool_id,omitempty" gorm:"-"` } type ChannelInfo struct { @@ -283,7 +287,6 @@ func GetAllChannels(startIdx int, num int, selectAll bool, idSort bool) ([]*Chan return channels, err } - func GetChannelsByTag(tag string, idSort bool, selectAll bool) ([]*Channel, error) { var channels []*Channel order := "priority desc" @@ -403,7 +406,7 @@ func BatchDeleteChannels(ids []int) error { return tx.Error } for _, chunk := range lo.Chunk(ids, 200) { - if err := tx.Where("id in (?)", chunk).Delete(&Channel{}).Error; err != nil { + if err := DeleteChannelAssetCredentialsWithTx(tx, chunk); err != nil { tx.Rollback() return err } @@ -411,6 +414,10 @@ func BatchDeleteChannels(ids []int) error { tx.Rollback() return err } + if err := tx.Where("id in (?)", chunk).Delete(&Channel{}).Error; err != nil { + tx.Rollback() + return err + } } return tx.Commit().Error } @@ -465,6 +472,12 @@ func (channel *Channel) Insert() error { } func (channel *Channel) Update() error { + return DB.Transaction(func(tx *gorm.DB) error { + return channel.UpdateWithTx(tx) + }) +} + +func (channel *Channel) UpdateWithTx(tx *gorm.DB) error { // If this is a multi-key channel, recalculate MultiKeySize based on the current key list to avoid inconsistency after editing keys if channel.ChannelInfo.IsMultiKey { var keyStr string @@ -472,7 +485,8 @@ func (channel *Channel) Update() error { keyStr = channel.Key } else { // If key is not provided, read the existing key from the database - if existing, err := GetChannelById(channel.Id, true); err == nil { + var existing Channel + if err := tx.First(&existing, "id = ?", channel.Id).Error; err == nil { keyStr = existing.Key } } @@ -503,16 +517,30 @@ func (channel *Channel) Update() error { } } } - var err error - err = DB.Model(channel).Updates(channel).Error + err := tx.Model(channel).Updates(channel).Error if err != nil { return err } - DB.Model(channel).First(channel, "id = ?", channel.Id) - err = channel.UpdateAbilities(nil) + if err = tx.First(channel, "id = ?", channel.Id).Error; err != nil { + return err + } + err = channel.UpdateAbilities(tx) return err } +func InsertChannelWithAssetCredential(channel *Channel, credential *ChannelAssetCredential) error { + return DB.Transaction(func(tx *gorm.DB) error { + if err := tx.Create(channel).Error; err != nil { + return err + } + if err := channel.AddAbilities(tx); err != nil { + return err + } + credential.ChannelId = channel.Id + return UpsertChannelAssetCredentialWithTx(tx, credential) + }) +} + func (channel *Channel) UpdateResponseTime(responseTime int64) { err := DB.Model(channel).Select("response_time", "test_time").Updates(Channel{ TestTime: common.GetTimestamp(), @@ -534,13 +562,19 @@ func (channel *Channel) UpdateBalance(balance float64) { } func (channel *Channel) Delete() error { - var err error - err = DB.Delete(channel).Error - if err != nil { + return DB.Transaction(func(tx *gorm.DB) error { + return channel.DeleteWithTx(tx) + }) +} + +func (channel *Channel) DeleteWithTx(tx *gorm.DB) error { + if err := DeleteChannelAssetCredentialWithTx(tx, channel.Id); err != nil { return err } - err = channel.DeleteAbilities() - return err + if err := tx.Where("channel_id = ?", channel.Id).Delete(&Ability{}).Error; err != nil { + return err + } + return tx.Delete(channel).Error } var channelStatusLock sync.Mutex @@ -778,13 +812,49 @@ func updateChannelUsedQuota(id int, quota int) { } func DeleteChannelByStatus(status int64) (int64, error) { - result := DB.Where("status = ?", status).Delete(&Channel{}) - return result.RowsAffected, result.Error + var ids []int + if err := DB.Model(&Channel{}).Where("status = ?", status).Pluck("id", &ids).Error; err != nil { + return 0, err + } + if len(ids) == 0 { + return 0, nil + } + var rows int64 + err := DB.Transaction(func(tx *gorm.DB) error { + if err := DeleteChannelAssetCredentialsWithTx(tx, ids); err != nil { + return err + } + if err := tx.Where("channel_id IN ?", ids).Delete(&Ability{}).Error; err != nil { + return err + } + result := tx.Where("id IN ?", ids).Delete(&Channel{}) + rows = result.RowsAffected + return result.Error + }) + return rows, err } func DeleteDisabledChannel() (int64, error) { - result := DB.Where("status = ? or status = ?", common.ChannelStatusAutoDisabled, common.ChannelStatusManuallyDisabled).Delete(&Channel{}) - return result.RowsAffected, result.Error + var ids []int + if err := DB.Model(&Channel{}).Where("status = ? or status = ?", common.ChannelStatusAutoDisabled, common.ChannelStatusManuallyDisabled).Pluck("id", &ids).Error; err != nil { + return 0, err + } + if len(ids) == 0 { + return 0, nil + } + var rows int64 + err := DB.Transaction(func(tx *gorm.DB) error { + if err := DeleteChannelAssetCredentialsWithTx(tx, ids); err != nil { + return err + } + if err := tx.Where("channel_id IN ?", ids).Delete(&Ability{}).Error; err != nil { + return err + } + result := tx.Where("id IN ?", ids).Delete(&Channel{}) + rows = result.RowsAffected + return result.Error + }) + return rows, err } func GetPaginatedTags(offset int, limit int) ([]*string, error) { diff --git a/model/channel_asset_credential.go b/model/channel_asset_credential.go new file mode 100644 index 0000000..d90801c --- /dev/null +++ b/model/channel_asset_credential.go @@ -0,0 +1,99 @@ +package model + +import ( + "github.com/QuantumNous/new-api/common" + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + +// ChannelAssetCredential stores the China Mobile asset credentials separately +// from the channel video-generation key. +type ChannelAssetCredential struct { + Id int `json:"id"` + ChannelId int `json:"channel_id" gorm:"uniqueIndex;not null"` + AccessKey string `json:"-" gorm:"not null;size:255"` + SecretKey string `json:"-" gorm:"not null;size:255"` + PoolID string `json:"pool_id" gorm:"size:255"` + CreatedAt int64 `json:"created_at" gorm:"bigint;not null"` + UpdatedAt int64 `json:"updated_at" gorm:"bigint;not null"` +} + +// ChannelAssetCredentialSummary is safe to expose in channel management APIs. +type ChannelAssetCredentialSummary struct { + ChannelId int + PoolID string +} + +func GetChannelAssetCredential(channelID int) (*ChannelAssetCredential, error) { + var credential ChannelAssetCredential + err := DB.Where("channel_id = ?", channelID).First(&credential).Error + if err == gorm.ErrRecordNotFound { + return nil, nil + } + if err != nil { + return nil, err + } + return &credential, nil +} + +func UpsertChannelAssetCredential(credential *ChannelAssetCredential) error { + return DB.Transaction(func(tx *gorm.DB) error { + return UpsertChannelAssetCredentialWithTx(tx, credential) + }) +} + +func UpsertChannelAssetCredentialWithTx(tx *gorm.DB, credential *ChannelAssetCredential) error { + now := common.GetTimestamp() + credential.CreatedAt = now + credential.UpdatedAt = now + return tx.Clauses(clause.OnConflict{ + Columns: []clause.Column{{Name: "channel_id"}}, + DoUpdates: clause.Assignments(map[string]any{ + "access_key": credential.AccessKey, + "secret_key": credential.SecretKey, + "pool_id": credential.PoolID, + "updated_at": credential.UpdatedAt, + }), + }).Create(credential).Error +} + +func DeleteChannelAssetCredentialWithTx(tx *gorm.DB, channelID int) error { + return tx.Where("channel_id = ?", channelID).Delete(&ChannelAssetCredential{}).Error +} + +func DeleteChannelAssetCredential(channelID int) error { + return DB.Transaction(func(tx *gorm.DB) error { + return DeleteChannelAssetCredentialWithTx(tx, channelID) + }) +} + +func DeleteChannelAssetCredentialsWithTx(tx *gorm.DB, channelIDs []int) error { + if len(channelIDs) == 0 { + return nil + } + return tx.Where("channel_id IN ?", channelIDs).Delete(&ChannelAssetCredential{}).Error +} + +func DeleteChannelAssetCredentials(channelIDs []int) error { + return DB.Transaction(func(tx *gorm.DB) error { + return DeleteChannelAssetCredentialsWithTx(tx, channelIDs) + }) +} + +func GetChannelAssetCredentialSummaries(channelIDs []int) (map[int]ChannelAssetCredentialSummary, error) { + summaries := make(map[int]ChannelAssetCredentialSummary) + if len(channelIDs) == 0 { + return summaries, nil + } + var rows []ChannelAssetCredentialSummary + if err := DB.Model(&ChannelAssetCredential{}). + Select("channel_id", "pool_id"). + Where("channel_id IN ?", channelIDs). + Find(&rows).Error; err != nil { + return nil, err + } + for _, row := range rows { + summaries[row.ChannelId] = row + } + return summaries, nil +} diff --git a/model/channel_asset_credential_test.go b/model/channel_asset_credential_test.go new file mode 100644 index 0000000..3614a16 --- /dev/null +++ b/model/channel_asset_credential_test.go @@ -0,0 +1,132 @@ +package model + +import ( + "testing" + + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" + "gorm.io/gorm/logger" +) + +func setupChannelAssetCredentialDB(t *testing.T) *gorm.DB { + t.Helper() + + db, err := gorm.Open(sqlite.Open("file:channel_asset_credentials?mode=memory&cache=shared"), &gorm.Config{ + Logger: logger.Default.LogMode(logger.Silent), + }) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + + originalDB := DB + DB = db + t.Cleanup(func() { + DB = originalDB + require.NoError(t, sqlDB.Close()) + }) + return db +} + +func TestChannelAssetCredentialUpsertKeepsOneCredentialPerChannel(t *testing.T) { + db := setupChannelAssetCredentialDB(t) + require.NoError(t, db.AutoMigrate(&ChannelAssetCredential{})) + + require.NoError(t, UpsertChannelAssetCredential(&ChannelAssetCredential{ + ChannelId: 61, + AccessKey: "ak-old", + SecretKey: "sk-old", + PoolID: "pool-old", + })) + require.NoError(t, UpsertChannelAssetCredential(&ChannelAssetCredential{ + ChannelId: 61, + AccessKey: "ak-new", + SecretKey: "sk-new", + PoolID: "pool-new", + })) + + credential, err := GetChannelAssetCredential(61) + require.NoError(t, err) + require.NotNil(t, credential) + assert.Equal(t, "ak-new", credential.AccessKey) + assert.Equal(t, "sk-new", credential.SecretKey) + assert.Equal(t, "pool-new", credential.PoolID) + + var count int64 + require.NoError(t, db.Model(&ChannelAssetCredential{}).Where("channel_id = ?", 61).Count(&count).Error) + assert.Equal(t, int64(1), count) +} + +func TestChannelAssetCredentialSummariesDoNotContainSecrets(t *testing.T) { + db := setupChannelAssetCredentialDB(t) + require.NoError(t, db.AutoMigrate(&ChannelAssetCredential{})) + require.NoError(t, UpsertChannelAssetCredential(&ChannelAssetCredential{ + ChannelId: 61, + AccessKey: "ak-secret", + SecretKey: "sk-secret", + PoolID: "pool-61", + })) + + summaries, err := GetChannelAssetCredentialSummaries([]int{61, 62}) + require.NoError(t, err) + require.Contains(t, summaries, 61) + assert.Equal(t, "pool-61", summaries[61].PoolID) + assert.NotContains(t, summaries, 62) +} + +func TestChannelDeleteRemovesOnlyItsAssetCredential(t *testing.T) { + db := setupChannelAssetCredentialDB(t) + require.NoError(t, db.AutoMigrate(&Channel{}, &Ability{}, &ChannelAssetCredential{})) + for _, id := range []int{61, 62} { + require.NoError(t, db.Create(&Channel{Id: id, Key: "key", Name: "channel", Group: "default", Models: "model"}).Error) + require.NoError(t, UpsertChannelAssetCredential(&ChannelAssetCredential{ChannelId: id, AccessKey: "ak", SecretKey: "sk"})) + } + + require.NoError(t, (&Channel{Id: 61}).Delete()) + credential, err := GetChannelAssetCredential(61) + require.NoError(t, err) + assert.Nil(t, credential) + credential, err = GetChannelAssetCredential(62) + require.NoError(t, err) + assert.NotNil(t, credential) +} + +func TestDeleteDisabledChannelRemovesAssetCredentials(t *testing.T) { + db := setupChannelAssetCredentialDB(t) + require.NoError(t, db.AutoMigrate(&Channel{}, &Ability{}, &ChannelAssetCredential{})) + require.NoError(t, db.Create(&Channel{Id: 61, Key: "key", Name: "disabled", Group: "default", Models: "model", Status: 2}).Error) + require.NoError(t, db.Create(&Channel{Id: 62, Key: "key", Name: "enabled", Group: "default", Models: "model", Status: 1}).Error) + for _, id := range []int{61, 62} { + require.NoError(t, UpsertChannelAssetCredential(&ChannelAssetCredential{ChannelId: id, AccessKey: "ak", SecretKey: "sk"})) + } + + _, err := DeleteDisabledChannel() + require.NoError(t, err) + credential, err := GetChannelAssetCredential(61) + require.NoError(t, err) + assert.Nil(t, credential) + credential, err = GetChannelAssetCredential(62) + require.NoError(t, err) + assert.NotNil(t, credential) +} + +func TestDeleteChannelByStatusRemovesAssetCredentials(t *testing.T) { + db := setupChannelAssetCredentialDB(t) + require.NoError(t, db.AutoMigrate(&Channel{}, &Ability{}, &ChannelAssetCredential{})) + require.NoError(t, db.Create(&Channel{Id: 61, Key: "key", Name: "disabled", Group: "default", Models: "model", Status: 2}).Error) + require.NoError(t, db.Create(&Channel{Id: 62, Key: "key", Name: "enabled", Group: "default", Models: "model", Status: 1}).Error) + for _, id := range []int{61, 62} { + require.NoError(t, UpsertChannelAssetCredential(&ChannelAssetCredential{ChannelId: id, AccessKey: "ak", SecretKey: "sk"})) + } + + _, err := DeleteChannelByStatus(2) + require.NoError(t, err) + credential, err := GetChannelAssetCredential(61) + require.NoError(t, err) + assert.Nil(t, credential) + credential, err = GetChannelAssetCredential(62) + require.NoError(t, err) + assert.NotNil(t, credential) +} diff --git a/model/main.go b/model/main.go index d50d47a..b60fa97 100644 --- a/model/main.go +++ b/model/main.go @@ -259,6 +259,7 @@ func migrateDB() error { err := DB.AutoMigrate( &Channel{}, + &ChannelAssetCredential{}, &Token{}, &User{}, &PasskeyCredential{}, @@ -329,6 +330,7 @@ func migrateDBFast() error { name string }{ {&Channel{}, "Channel"}, + {&ChannelAssetCredential{}, "ChannelAssetCredential"}, {&Token{}, "Token"}, {&User{}, "User"}, {&PasskeyCredential{}, "PasskeyCredential"}, diff --git a/service/asset_chinamobile.go b/service/asset_chinamobile.go index f3c0a92..ad359f5 100644 --- a/service/asset_chinamobile.go +++ b/service/asset_chinamobile.go @@ -62,7 +62,7 @@ func (a *ChinaMobileAssetAdapter) DoAssetRequest(ctx context.Context, channel *m return nil, err } - credential, err := chinaMobileAssetCredentialFromEnv() + credential, err := chinaMobileAssetCredentialFromChannel(channel) if err != nil { return nil, newAssetError(AssetErrorInvalidRequest, err.Error(), http.StatusBadRequest) } @@ -101,11 +101,30 @@ type chinaMobileAssetCredential struct { PoolID string `json:"pool_id"` } -func chinaMobileAssetCredentialFromEnv() (chinaMobileAssetCredential, error) { +func chinaMobileAssetCredentialFromChannel(channel *model.Channel) (chinaMobileAssetCredential, error) { + if channel == nil { + return chinaMobileAssetCredential{}, fmt.Errorf("China Mobile asset channel is required") + } + credential, err := model.GetChannelAssetCredential(channel.Id) + if err != nil { + return chinaMobileAssetCredential{}, err + } + if credential == nil { + legacyCredential, legacyErr := normalizeChinaMobileAssetCredential(chinaMobileAssetCredential{ + AK: os.Getenv(chinaMobileAssetAKEnv), + SK: os.Getenv(chinaMobileAssetSKEnv), + PoolID: os.Getenv(chinaMobileAssetPoolIDEnv), + }) + if legacyErr != nil { + return chinaMobileAssetCredential{}, fmt.Errorf("该移动云渠道未配置素材凭证") + } + common.SysLog(fmt.Sprintf("using legacy China Mobile asset credentials for channel %d; configure channel asset credentials before removing environment fallback", channel.Id)) + return legacyCredential, nil + } return normalizeChinaMobileAssetCredential(chinaMobileAssetCredential{ - AK: os.Getenv(chinaMobileAssetAKEnv), - SK: os.Getenv(chinaMobileAssetSKEnv), - PoolID: os.Getenv(chinaMobileAssetPoolIDEnv), + AK: credential.AccessKey, + SK: credential.SecretKey, + PoolID: credential.PoolID, }) } @@ -114,7 +133,7 @@ func normalizeChinaMobileAssetCredential(credential chinaMobileAssetCredential) credential.SK = strings.TrimSpace(credential.SK) credential.PoolID = strings.TrimSpace(credential.PoolID) if credential.AK == "" || credential.SK == "" { - return chinaMobileAssetCredential{}, fmt.Errorf("%s and %s are required for China Mobile asset library", chinaMobileAssetAKEnv, chinaMobileAssetSKEnv) + return chinaMobileAssetCredential{}, fmt.Errorf("China Mobile asset AccessKey and SecretKey are required") } if credential.PoolID == "" { credential.PoolID = defaultChinaMobileAssetPoolID diff --git a/service/asset_chinamobile_test.go b/service/asset_chinamobile_test.go index 23182d1..300830c 100644 --- a/service/asset_chinamobile_test.go +++ b/service/asset_chinamobile_test.go @@ -8,10 +8,13 @@ import ( "time" "github.com/QuantumNous/new-api/model" + "github.com/glebarez/sqlite" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" cmerrs "gitlab.ecloud.com/ecloud/ecloudsdkcore/errs" cmmodel "gitlab.ecloud.com/ecloud/ecloudsdkmaas/model" + "gorm.io/gorm" + "gorm.io/gorm/logger" ) func TestChinaMobileAssetAdapterCreateAssetMapsRequestAndResponse(t *testing.T) { @@ -49,10 +52,10 @@ func TestChinaMobileAssetAdapterCreateAssetMapsRequestAndResponse(t *testing.T) assert.Contains(t, string(resp.Body), `"Result":"asset-1"`) } -func TestChinaMobileAssetCredentialFromEnvDefaultsCenterPool(t *testing.T) { +func TestChinaMobileAssetCredentialFromChannelDefaultsCenterPool(t *testing.T) { setChinaMobileAssetTestEnv(t, "") - credential, err := chinaMobileAssetCredentialFromEnv() + credential, err := chinaMobileAssetCredentialFromChannel(&model.Channel{}) require.NoError(t, err) assert.Equal(t, "ak", credential.AK) @@ -60,10 +63,10 @@ func TestChinaMobileAssetCredentialFromEnvDefaultsCenterPool(t *testing.T) { assert.Equal(t, defaultChinaMobileAssetPoolID, credential.PoolID) } -func TestChinaMobileAssetCredentialFromEnvSupportsPoolID(t *testing.T) { +func TestChinaMobileAssetCredentialFromChannelSupportsPoolID(t *testing.T) { setChinaMobileAssetTestEnv(t, "CIDC-RP-29") - credential, err := chinaMobileAssetCredentialFromEnv() + credential, err := chinaMobileAssetCredentialFromChannel(&model.Channel{}) require.NoError(t, err) assert.Equal(t, "ak", credential.AK) @@ -71,16 +74,66 @@ func TestChinaMobileAssetCredentialFromEnvSupportsPoolID(t *testing.T) { assert.Equal(t, "CIDC-RP-29", credential.PoolID) } -func TestChinaMobileAssetCredentialFromEnvRequiresAKAndSK(t *testing.T) { +func TestChinaMobileAssetCredentialFromChannelRequiresCredential(t *testing.T) { + setupChinaMobileAssetTestDB(t) t.Setenv(chinaMobileAssetAKEnv, "") t.Setenv(chinaMobileAssetSKEnv, "") t.Setenv(chinaMobileAssetPoolIDEnv, "") - _, err := chinaMobileAssetCredentialFromEnv() + _, err := chinaMobileAssetCredentialFromChannel(&model.Channel{Id: 99}) require.Error(t, err) - assert.Contains(t, err.Error(), chinaMobileAssetAKEnv) - assert.Contains(t, err.Error(), chinaMobileAssetSKEnv) + assert.Contains(t, err.Error(), "未配置素材凭证") +} + +func TestChinaMobileAssetCredentialFromChannelFallsBackToLegacyEnvironment(t *testing.T) { + setupChinaMobileAssetTestDB(t) + t.Setenv(chinaMobileAssetAKEnv, "legacy-ak") + t.Setenv(chinaMobileAssetSKEnv, "legacy-sk") + t.Setenv(chinaMobileAssetPoolIDEnv, "legacy-pool") + + credential, err := chinaMobileAssetCredentialFromChannel(&model.Channel{Id: 99}) + + require.NoError(t, err) + assert.Equal(t, "legacy-ak", credential.AK) + assert.Equal(t, "legacy-sk", credential.SK) + assert.Equal(t, "legacy-pool", credential.PoolID) +} + +func TestChinaMobileAssetCredentialFromChannelPrefersChannelCredentialOverLegacyEnvironment(t *testing.T) { + setupChinaMobileAssetTestDB(t) + require.NoError(t, model.UpsertChannelAssetCredential(&model.ChannelAssetCredential{ + ChannelId: 99, + AccessKey: "channel-ak", + SecretKey: "channel-sk", + PoolID: "channel-pool", + })) + t.Setenv(chinaMobileAssetAKEnv, "legacy-ak") + t.Setenv(chinaMobileAssetSKEnv, "legacy-sk") + t.Setenv(chinaMobileAssetPoolIDEnv, "legacy-pool") + + credential, err := chinaMobileAssetCredentialFromChannel(&model.Channel{Id: 99}) + + require.NoError(t, err) + assert.Equal(t, "channel-ak", credential.AK) + assert.Equal(t, "channel-sk", credential.SK) + assert.Equal(t, "channel-pool", credential.PoolID) +} + +func TestChinaMobileAssetAdapterReturnsBadRequestWhenCredentialIsMissing(t *testing.T) { + setupChinaMobileAssetTestDB(t) + adapter := &ChinaMobileAssetAdapter{newClient: func(credential chinaMobileAssetCredential) chinaMobileAssetSDKClient { + return &fakeChinaMobileAssetSDKClient{} + }} + spec, ok := ParseAssetAction("CreateAsset") + require.True(t, ok) + + _, assetErr := adapter.DoAssetRequest(context.Background(), &model.Channel{Id: 61}, AssetRequest{Action: spec, Body: validChinaMobileCreateAssetBody()}) + + require.NotNil(t, assetErr) + assert.Equal(t, AssetErrorInvalidRequest, assetErr.Type) + assert.Equal(t, http.StatusBadRequest, assetErr.HTTPStatus) + assert.Contains(t, assetErr.Message, "未配置素材凭证") } func TestChinaMobileAssetAdapterAllOfficialActions(t *testing.T) { @@ -280,9 +333,32 @@ func TestCallChinaMobileAssetSDKReturnsWhenContextCancelled(t *testing.T) { func setChinaMobileAssetTestEnv(t *testing.T, poolID string) { t.Helper() - t.Setenv(chinaMobileAssetAKEnv, "ak") - t.Setenv(chinaMobileAssetSKEnv, "sk") - t.Setenv(chinaMobileAssetPoolIDEnv, poolID) + setupChinaMobileAssetTestDB(t) + require.NoError(t, model.UpsertChannelAssetCredential(&model.ChannelAssetCredential{ + ChannelId: 0, + AccessKey: "ak", + SecretKey: "sk", + PoolID: poolID, + })) +} + +func setupChinaMobileAssetTestDB(t *testing.T) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open("file:service_chinamobile_asset?mode=memory&cache=shared"), &gorm.Config{ + Logger: logger.Default.LogMode(logger.Silent), + }) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + originalDB := model.DB + model.DB = db + require.NoError(t, db.AutoMigrate(&model.ChannelAssetCredential{})) + t.Cleanup(func() { + model.DB = originalDB + require.NoError(t, sqlDB.Close()) + }) + return db } type blockingChinaMobileAssetSDKClient struct { diff --git a/setup_test.sql b/setup_test.sql index 79ffe20..5080e6a 100644 --- a/setup_test.sql +++ b/setup_test.sql @@ -1,5 +1,6 @@ DELETE FROM tokens WHERE user_id IN (SELECT id FROM users WHERE username LIKE 'asset-admin' OR username LIKE 'asset-user'); DELETE FROM user_asset_channels WHERE user_id IN (9001, 9002); +DELETE FROM channel_asset_credentials WHERE channel_id = 9001; DELETE FROM channels WHERE name = 'cm-asset-channel'; DELETE FROM users WHERE username LIKE 'asset-admin' OR username LIKE 'asset-user'; @@ -8,7 +9,9 @@ VALUES (9001, 'asset-admin', 'pbkdf2_sha256$600000$dGVzdA==$hash', '素材管理 INSERT INTO users (id, username, password, display_name, role, status, quota, "group") VALUES (9002, 'asset-user', 'pbkdf2_sha256$600000$dGVzdA==$hash', '素材用户', 0, 1, 1000000000, 'default'); INSERT INTO channels (id, type, key, status, name, "group", models, weight, priority, auto_ban, created_time) -VALUES (9001, 61, '__CHINAMOBILE_ASSET_CHANNEL_KEY__', 1, 'cm-asset-channel', 'default', 'seedance-2', 10, 1, 1, 1784103836); +VALUES (9001, 61, 'video-generation-key', 1, 'cm-asset-channel', 'default', 'seedance-2', 10, 1, 1, 1784103836); +INSERT INTO channel_asset_credentials (channel_id, access_key, secret_key, pool_id, created_at, updated_at) +VALUES (9001, '__CHINAMOBILE_ASSET_AK__', '__CHINAMOBILE_ASSET_SK__', '__CHINAMOBILE_ASSET_POOL_ID__', 1784103836, 1784103836); INSERT INTO tokens (id, user_id, key, status, name, "group", expired_time, unlimited_quota) VALUES (9001, 9001, 'assetadmintestkey00000000000000000000000000001', 1, 'asset-admin-token', 'default', -1, 1); INSERT INTO tokens (id, user_id, key, status, name, "group", expired_time, unlimited_quota) diff --git a/setup_test_env.ps1 b/setup_test_env.ps1 index a8d2e7c..775d690 100644 --- a/setup_test_env.ps1 +++ b/setup_test_env.ps1 @@ -17,8 +17,9 @@ if ($ak.Contains("'") -or $sk.Contains("'") -or $poolId.Contains("'")) { throw "China Mobile asset credentials must not contain single quotes" } -$channelKey = "$ak|$sk|$poolId" $sql = Get-Content -LiteralPath $TemplatePath -Raw -$sql = $sql.Replace("__CHINAMOBILE_ASSET_CHANNEL_KEY__", $channelKey) +$sql = $sql.Replace("__CHINAMOBILE_ASSET_AK__", $ak) +$sql = $sql.Replace("__CHINAMOBILE_ASSET_SK__", $sk) +$sql = $sql.Replace("__CHINAMOBILE_ASSET_POOL_ID__", $poolId) [System.IO.File]::WriteAllText($OutputPath, $sql, [System.Text.UTF8Encoding]::new($false)) Write-Output $OutputPath diff --git a/web/src/components/table/channels/modals/EditChannelModal.jsx b/web/src/components/table/channels/modals/EditChannelModal.jsx index 757f4ee..57a514c 100644 --- a/web/src/components/table/channels/modals/EditChannelModal.jsx +++ b/web/src/components/table/channels/modals/EditChannelModal.jsx @@ -183,6 +183,12 @@ const EditChannelModal = (props) => { allow_include_obfuscation: false, allow_inference_geo: false, claude_beta_query: false, + asset_credential: { + access_key: '', + secret_key: '', + pool_id: '', + }, + asset_credential_configured: false, }; const [batch, setBatch] = useState(false); const [multiToSingle, setMultiToSingle] = useState(false); @@ -1359,6 +1365,27 @@ const EditChannelModal = (props) => { if (isEdit && (!localInputs.key || localInputs.key.trim() === '')) { delete localInputs.key; } + if (localInputs.type === 61) { + const credential = localInputs.asset_credential || {}; + const accessKey = String(credential.access_key || '').trim(); + const secretKey = String(credential.secret_key || '').trim(); + const poolId = String(credential.pool_id || '').trim(); + if ((accessKey === '') !== (secretKey === '')) { + showInfo(t('请同时填写素材 AccessKey 和 SecretKey')); + return; + } + if (accessKey === '') { + delete localInputs.asset_credential; + } else { + localInputs.asset_credential = { + access_key: accessKey, + secret_key: secretKey, + pool_id: poolId, + }; + } + } else { + delete localInputs.asset_credential; + } delete localInputs.vertex_files; if (!isEdit && (!localInputs.name || !localInputs.key)) { @@ -1551,10 +1578,15 @@ const EditChannelModal = (props) => { key_mode: isMultiKeyChannel ? keyMode : undefined, // 只在多key模式下传递 }); } else { + const { + asset_credential: assetCredential, + ...channelInputs + } = localInputs; res = await API.post(`/api/channel/`, { mode: mode, multi_key_mode: mode === 'multi_to_single' ? multiKeyMode : undefined, - channel: localInputs, + channel: channelInputs, + asset_credential: assetCredential, }); } const { success, message } = res.data; @@ -2506,6 +2538,35 @@ const EditChannelModal = (props) => { )} + {inputs.type === 61 && ( + + {t('移动云素材凭证')} + + {inputs.asset_credential_configured + ? t('已配置素材凭证；留空不会覆盖现有凭证。') + : t('未配置素材凭证；素材上传接口将不可用。')} + + + + + + )} + {isEdit && isMultiKeyChannel && (