feat: publish v1.0.0-rc.36-custom.1
This commit is contained in:
Codex committed
commit
3a7e8f3b54
2381 files changed
+551494
No files matched your search
@@ -0,0 +1,113 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strconv"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/i18n"
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func GetAccessTokenStatus(c *gin.Context) {
|
||||
status, err := model.GetUserAccessTokenStatus(c.GetInt("id"))
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, status)
|
||||
}
|
||||
|
||||
func GenerateAccessToken(c *gin.Context) {
|
||||
if middleware.RequireSecurityProof(c, service.VerificationOperation{Scope: service.VerificationScopeAccessTokenGenerate}) == nil {
|
||||
return
|
||||
}
|
||||
id := c.GetInt("id")
|
||||
key, err := common.GenerateRandomKey(29 + common.GetRandomInt(4))
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
var existing int64
|
||||
if err := model.DB.Model(&model.User{}).Where("access_token = ?", key).Count(&existing).Error; err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
if existing != 0 {
|
||||
common.ApiErrorI18n(c, i18n.MsgUuidDuplicate)
|
||||
return
|
||||
}
|
||||
if err := model.UpdateUserAccessToken(id, key); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
recordUserSecurityAudit(c, id, "access_token.generate", map[string]any{"token_ref": model.AccessTokenFingerprint(key)})
|
||||
c.JSON(http.StatusOK, gin.H{"success": true, "message": "", "data": key})
|
||||
}
|
||||
|
||||
func RevokeAccessToken(c *gin.Context) {
|
||||
if middleware.RequireSecurityProof(c, service.VerificationOperation{Scope: service.VerificationScopeAccessTokenRevoke}) == nil {
|
||||
return
|
||||
}
|
||||
ref, err := model.RevokeUserAccessToken(c.GetInt("id"))
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
if ref != "" {
|
||||
recordUserSecurityAudit(c, c.GetInt("id"), "access_token.revoke", map[string]any{"token_ref": ref})
|
||||
}
|
||||
common.ApiSuccess(c, nil)
|
||||
}
|
||||
|
||||
func GetAuditLogs(c *gin.Context) {
|
||||
page := common.GetPageQuery(c)
|
||||
if page.Page < 1 || page.PageSize < 1 || page.Page > 100000000 {
|
||||
common.ApiErrorMsg(c, "Invalid audit pagination")
|
||||
return
|
||||
}
|
||||
filter := model.AuditLogFilter{Username: c.Query("username"), Category: c.Query("category"), TokenRef: c.Query("token_ref"), ExcludeTokenRef: c.Query("exclude_token_ref"), RequestId: c.Query("request_id")}
|
||||
viewerRole := c.GetInt("role")
|
||||
if c.FullPath() == "/api/audit/self" {
|
||||
filter.UserId = c.GetInt("id")
|
||||
filter.Username = ""
|
||||
filter.SelfView = true
|
||||
}
|
||||
if !model.ValidAuditCategory(filter.Category) || !model.ValidTokenFingerprint(filter.TokenRef) || !model.ValidTokenFingerprint(filter.ExcludeTokenRef) {
|
||||
common.ApiErrorMsg(c, "Invalid audit filters")
|
||||
return
|
||||
}
|
||||
for name, target := range map[string]*int64{"start_timestamp": &filter.StartTimestamp, "end_timestamp": &filter.EndTimestamp} {
|
||||
if raw := c.Query(name); raw != "" {
|
||||
parsed, err := strconv.ParseInt(raw, 10, 64)
|
||||
if err != nil || parsed < 0 {
|
||||
common.ApiErrorMsg(c, "Invalid audit time range")
|
||||
return
|
||||
}
|
||||
*target = parsed
|
||||
}
|
||||
}
|
||||
if filter.EndTimestamp > 0 && filter.EndTimestamp < filter.StartTimestamp {
|
||||
common.ApiErrorMsg(c, "Invalid audit time range")
|
||||
return
|
||||
}
|
||||
if raw := c.Query("success"); raw != "" {
|
||||
if raw != "true" && raw != "false" {
|
||||
common.ApiErrorMsg(c, "Invalid audit result")
|
||||
return
|
||||
}
|
||||
success := raw == "true"
|
||||
filter.Success = &success
|
||||
}
|
||||
logs, total, err := model.GetAuditLogs(filter, page.GetStartIdx(), page.GetPageSize(), viewerRole)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
page.SetItems(logs)
|
||||
page.SetTotal(int(total))
|
||||
common.ApiSuccess(c, page)
|
||||
}
|
||||
@@ -0,0 +1,851 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
sqlmysql "github.com/go-sql-driver/mysql"
|
||||
"gorm.io/driver/clickhouse"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/service/authz"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/driver/mysql"
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func setupAccessTokenAudit(t *testing.T) (*model.User, string) {
|
||||
t.Helper()
|
||||
previousDB, previousLogDB := model.DB, model.LOG_DB
|
||||
previousMain, previousLog := common.MainDatabaseType(), common.LogDatabaseType()
|
||||
previousRedis := common.RedisEnabled
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.AutoMigrate(&model.User{}, &model.UserSession{}, &model.Log{}, &model.AuditLog{}, &model.CasbinRule{}, &model.AuthzRole{}))
|
||||
model.DB, model.LOG_DB = db, db
|
||||
common.SetDatabaseTypes(common.DatabaseTypeSQLite, common.DatabaseTypeSQLite)
|
||||
common.RedisEnabled = false
|
||||
previousMaster := common.IsMasterNode
|
||||
common.IsMasterNode = true
|
||||
require.NoError(t, authz.Init(db))
|
||||
t.Cleanup(func() {
|
||||
model.DB, model.LOG_DB = previousDB, previousLogDB
|
||||
common.IsMasterNode = previousMaster
|
||||
common.SetDatabaseTypes(previousMain, previousLog)
|
||||
common.RedisEnabled = previousRedis
|
||||
})
|
||||
token := "legacy-opaque-token"
|
||||
user := &model.User{Username: "audit-owner", Password: "placeholder", Role: common.RoleAdminUser, Status: common.UserStatusEnabled, Group: "default", AccessToken: &token, AuthVersion: 1, AffCode: "audit-owner"}
|
||||
require.NoError(t, db.Create(user).Error)
|
||||
return user, token
|
||||
}
|
||||
|
||||
func auditRequest(router http.Handler, method, path, token string) *httptest.ResponseRecorder {
|
||||
request := httptest.NewRequest(method, path, strings.NewReader(`{"secret":"body-must-not-be-logged"}`))
|
||||
request.Header.Set("Authorization", "Bearer "+token)
|
||||
request.Header.Set("User-Agent", "audit-test-client")
|
||||
request.RemoteAddr = "192.0.2.8:4567"
|
||||
response := httptest.NewRecorder()
|
||||
router.ServeHTTP(response, request)
|
||||
return response
|
||||
}
|
||||
|
||||
func TestAccessTokenLifecycleAndLateRequests(t *testing.T) {
|
||||
user, old := setupAccessTokenAudit(t)
|
||||
router := gin.New()
|
||||
router.Use(middleware.RequestId(), middleware.AccessTokenAudit())
|
||||
router.GET("/api/user/token/status", middleware.UserAuth(), GetAccessTokenStatus)
|
||||
router.GET("/api/user/token", middleware.UserAuth(), GenerateAccessToken)
|
||||
router.POST("/api/user/token", middleware.UserAuth(), GenerateAccessToken)
|
||||
router.DELETE("/api/user/token", middleware.UserAuth(), RevokeAccessToken)
|
||||
// Rotate during the handler, after authentication has captured the old PAT.
|
||||
router.POST("/rotate-in-flight", middleware.UserAuth(), func(c *gin.Context) {
|
||||
require.NoError(t, model.UpdateUserAccessToken(user.Id, "new-token"))
|
||||
c.JSON(200, gin.H{"success": true})
|
||||
})
|
||||
status, err := model.GetUserAccessTokenStatus(user.Id)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, status.Exists)
|
||||
assert.Nil(t, status.CreatedAt)
|
||||
assert.Nil(t, status.LastUsedAt)
|
||||
assert.NotContains(t, auditRequest(router, "GET", "/api/user/token/status", old).Body.String(), old)
|
||||
require.Equal(t, 200, auditRequest(router, "POST", "/rotate-in-flight", old).Code)
|
||||
status, err = model.GetUserAccessTokenStatus(user.Id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, model.AccessTokenFingerprint("new-token"), status.TokenRef)
|
||||
assert.NotNil(t, status.CreatedAt)
|
||||
assert.Nil(t, status.LastUsedAt, "in-flight old requests must not mark the new generation as used")
|
||||
assert.Equal(t, 401, auditRequest(router, "GET", "/api/user/token/status", old).Code)
|
||||
for _, method := range []string{"POST", "GET", "DELETE"} {
|
||||
response := auditRequest(router, method, "/api/user/token", "new-token")
|
||||
assert.Equal(t, http.StatusForbidden, response.Code)
|
||||
assert.Contains(t, response.Body.String(), `"code":"SECURITY_PROOF_INVALID"`)
|
||||
stored, err := model.GetUserById(user.Id, true)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "new-token", stored.GetAccessToken(), "a PAT cannot manage itself without a dashboard verification")
|
||||
}
|
||||
_, err = model.RevokeUserAccessToken(user.Id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 401, auditRequest(router, "GET", "/api/user/token/status", "new-token").Code)
|
||||
ref, err := model.RevokeUserAccessToken(user.Id)
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, ref, "repeated revocation is idempotent")
|
||||
status, err = model.GetUserAccessTokenStatus(user.Id)
|
||||
require.NoError(t, err)
|
||||
assert.False(t, status.Exists)
|
||||
assert.Nil(t, status.CreatedAt)
|
||||
var history []model.AuditLog
|
||||
require.NoError(t, model.LOG_DB.Find(&history).Error)
|
||||
require.NotEmpty(t, history)
|
||||
encoded, err := common.Marshal(history)
|
||||
require.NoError(t, err)
|
||||
for _, secret := range []string{old, "new-token", "body-must-not-be-logged", "Authorization"} {
|
||||
assert.NotContains(t, string(encoded), secret)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAccessTokenAuditsResultsAndExcludesBrowserSessions(t *testing.T) {
|
||||
user, pat := setupAccessTokenAudit(t)
|
||||
router := gin.New()
|
||||
router.Use(middleware.RequestId(), middleware.AccessTokenAudit())
|
||||
router.NoRoute(func(c *gin.Context) { c.JSON(404, gin.H{"success": false}) })
|
||||
router.GET("/public", func(c *gin.Context) { c.JSON(200, gin.H{"success": true}) })
|
||||
router.GET("/rate-limited", func(c *gin.Context) { c.AbortWithStatusJSON(429, gin.H{"success": false}) }, middleware.UserAuth())
|
||||
router.GET("/read/:id", middleware.UserAuth(), func(c *gin.Context) { c.JSON(200, gin.H{"success": true}) })
|
||||
router.POST("/write", middleware.AdminAuth(), func(c *gin.Context) {
|
||||
recordManageAudit(c, "option.update", map[string]any{"key": "safe-setting"})
|
||||
c.JSON(200, gin.H{"success": true})
|
||||
})
|
||||
router.POST("/business-failure", middleware.AdminAuth(), func(c *gin.Context) { c.JSON(200, gin.H{"success": false, "message": "secret-response"}) })
|
||||
router.GET("/forbidden", middleware.RootAuth(), func(c *gin.Context) { c.Status(204) })
|
||||
cases := []struct {
|
||||
method, path string
|
||||
status int
|
||||
success bool
|
||||
}{{"GET", "/missing", 404, false}, {"GET", "/public", 200, true}, {"GET", "/rate-limited", 429, false}, {"GET", "/read/sensitive-id?password=secret-query", 200, true}, {"POST", "/write", 200, true}, {"POST", "/business-failure", 200, false}, {"GET", "/forbidden", 403, false}}
|
||||
for _, tc := range cases {
|
||||
response := auditRequest(router, tc.method, tc.path, pat)
|
||||
require.Equal(t, tc.status, response.Code)
|
||||
var access model.AuditLog
|
||||
require.NoError(t, model.LOG_DB.Where("request_id = ? AND category = ?", response.Header().Get(common.RequestIdKey), model.AuditCategoryAccessToken).First(&access).Error)
|
||||
assert.Equal(t, tc.success, access.Success)
|
||||
assert.Equal(t, tc.status, access.Status)
|
||||
assert.Equal(t, user.Id, access.UserId)
|
||||
assert.Equal(t, "192.0.2.8", access.Ip)
|
||||
assert.Equal(t, "audit-test-client", access.UserAgent)
|
||||
assert.NotContains(t, access.Route, "sensitive-id")
|
||||
}
|
||||
var operationCount int64
|
||||
require.NoError(t, model.LOG_DB.Model(&model.AuditLog{}).Where("category = ?", model.AuditCategoryOperation).Count(&operationCount).Error)
|
||||
assert.EqualValues(t, 2, operationCount, "manual operation must not be duplicated by fallback")
|
||||
now := time.Now().Unix()
|
||||
session := &model.UserSession{SID: "audit-session", UserID: user.Id, Version: 1, UserAuthVersion: 1, Status: model.UserSessionStatusActive, RefreshHash: "refresh-placeholder", LoginMethod: "password", LastActiveAt: now, ExpiresAt: now + 3600}
|
||||
require.NoError(t, model.CreateUserSession(session))
|
||||
jwt, _, err := service.IssueAccessToken(service.AuthIdentity{UserID: user.Id, SessionID: session.SID, UserAuthVersion: 1, SessionVersion: 1})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 200, auditRequest(router, "GET", "/read/123", jwt).Code)
|
||||
assert.Equal(t, 401, auditRequest(router, "GET", "/read/123", "unknown-token").Code)
|
||||
entries, total, err := model.GetAuditLogs(model.AuditLogFilter{Category: model.AuditCategoryAccessToken}, 0, 20, common.RoleRootUser)
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 7, total)
|
||||
encoded, err := common.Marshal(entries)
|
||||
require.NoError(t, err)
|
||||
for _, secret := range []string{pat, jwt, "secret-query", "secret-response", "body-must-not-be-logged", "sensitive-id"} {
|
||||
assert.NotContains(t, string(encoded), secret)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuditIsolationVisibilityAndFailureContracts(t *testing.T) {
|
||||
user, pat := setupAccessTokenAudit(t)
|
||||
metadata := model.AuditOther{
|
||||
Op: &model.AuditOperation{Action: "generic"},
|
||||
AdminInfo: &model.AuditAdminInfo{AdminID: 1},
|
||||
RootInfo: model.AuditFields{"private": "root-only"},
|
||||
}
|
||||
for _, owner := range []int{user.Id, user.Id + 1} {
|
||||
model.RecordAuditLog(nil, model.AuditLog{ActorRole: common.RoleAdminUser, UserId: owner, Username: fmt.Sprint(owner), Category: model.AuditCategorySecurity, Success: false, Other: metadata})
|
||||
}
|
||||
router := gin.New()
|
||||
router.Use(middleware.RequestId(), middleware.AccessTokenAudit())
|
||||
router.GET("/api/audit/self", middleware.UserAuth(), GetAuditLogs)
|
||||
router.GET("/api/audit", middleware.AdminAuth(), middleware.RequirePermission(authz.AuditRead), GetAuditLogs)
|
||||
response := auditRequest(router, "GET", "/api/audit/self?username=other&user_id=2&category=security&success=false&page_size=1", pat)
|
||||
var result struct {
|
||||
Success bool
|
||||
Data struct {
|
||||
Items []model.AuditLog
|
||||
Total int
|
||||
}
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.True(t, result.Success)
|
||||
require.Equal(t, 1, result.Data.Total)
|
||||
require.Len(t, result.Data.Items, 1)
|
||||
assert.Equal(t, user.Id, result.Data.Items[0].UserId)
|
||||
assert.NotContains(t, response.Body.String(), "admin_info")
|
||||
assert.NotContains(t, response.Body.String(), "root-only")
|
||||
var payload struct {
|
||||
Data struct {
|
||||
Items []struct {
|
||||
Other map[string]any `json:"other"`
|
||||
}
|
||||
}
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &payload), "audit metadata must be a JSON object, not an encoded string")
|
||||
require.Len(t, payload.Data.Items, 1)
|
||||
assert.Contains(t, payload.Data.Items[0].Other, "op")
|
||||
require.NoError(t, authz.SetUserPermissions(user.Id, authz.PermissionsMap{authz.ResourceAudit: {authz.ActionRead: true}}))
|
||||
response = auditRequest(router, "GET", "/api/audit?category=security", pat)
|
||||
assert.Contains(t, response.Body.String(), "admin_info")
|
||||
assert.NotContains(t, response.Body.String(), "root-only")
|
||||
for _, query := range []string{"success=bad", "category=bad", "token_ref=secret", "start_timestamp=-1", "start_timestamp=2&end_timestamp=1", "p=-1", "page_size=-1"} {
|
||||
assert.Contains(t, auditRequest(router, "GET", "/api/audit/self?"+query, pat).Body.String(), `"success":false`)
|
||||
}
|
||||
require.NoError(t, model.LOG_DB.Callback().Query().Before("gorm:query").Register("audit:fail", func(tx *gorm.DB) {
|
||||
if tx.Statement.Table == "audit_logs" {
|
||||
tx.AddError(errors.New("audit store unavailable"))
|
||||
}
|
||||
}))
|
||||
_, err := model.GetUserAccessTokenStatus(user.Id)
|
||||
require.Error(t, err, "audit query failure must not look like never used")
|
||||
model.LOG_DB.Callback().Query().Remove("audit:fail")
|
||||
require.NoError(t, model.DB.Callback().Update().Before("gorm:update").Register("audit:write-fail", func(tx *gorm.DB) { tx.AddError(errors.New("write unavailable")) }))
|
||||
require.Error(t, model.UpdateUserAccessToken(user.Id, "replacement"))
|
||||
_, err = model.RevokeUserAccessToken(user.Id)
|
||||
require.Error(t, err)
|
||||
model.DB.Callback().Update().Remove("audit:write-fail")
|
||||
status, err := model.GetUserAccessTokenStatus(user.Id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, model.AccessTokenFingerprint(pat), status.TokenRef)
|
||||
}
|
||||
|
||||
func TestAuditRoleVisibilityAndPermissions(t *testing.T) {
|
||||
admin, pat := setupAccessTokenAudit(t)
|
||||
rootToken := "root-audit-token"
|
||||
root := &model.User{Username: "root-audit", Role: common.RoleRootUser, Status: common.UserStatusEnabled, AuthVersion: 1, AccessToken: &rootToken, AffCode: "root-audit"}
|
||||
require.NoError(t, model.DB.Create(root).Error)
|
||||
metadata := model.AuditOther{
|
||||
AdminInfo: &model.AuditAdminInfo{AdminID: 1},
|
||||
RootInfo: model.AuditFields{"private": "root-only"},
|
||||
}
|
||||
for i, role := range []int{1, 10, 100, 0, -1, 99} {
|
||||
model.RecordAuditLog(nil, model.AuditLog{ActorRole: role, UserId: admin.Id, Username: admin.Username, Category: model.AuditCategorySecurity, RequestId: fmt.Sprintf("role-%d", role), CreatedAt: int64(100 + i), Other: metadata})
|
||||
}
|
||||
model.RecordAuditLog(nil, model.AuditLog{ActorRole: 100, UserId: root.Id, Username: root.Username, Category: model.AuditCategorySecurity, RequestId: "root-owned", Other: metadata})
|
||||
router := gin.New()
|
||||
router.Use(middleware.RequestId(), middleware.AccessTokenAudit())
|
||||
router.GET("/api/audit", middleware.AdminAuth(), middleware.RequirePermission(authz.AuditRead), GetAuditLogs)
|
||||
router.GET("/api/audit/self", middleware.UserAuth(), GetAuditLogs)
|
||||
now := time.Now().Unix()
|
||||
session := &model.UserSession{SID: "audit-permissions-session", UserID: admin.Id, Version: 1, UserAuthVersion: 1, Status: model.UserSessionStatusActive, RefreshHash: "placeholder", LoginMethod: "password", LastActiveAt: now, ExpiresAt: now + 3600}
|
||||
require.NoError(t, model.CreateUserSession(session))
|
||||
jwt, _, err := service.IssueAccessToken(service.AuthIdentity{UserID: admin.Id, SessionID: session.SID, UserAuthVersion: 1, SessionVersion: 1})
|
||||
require.NoError(t, err)
|
||||
for _, credential := range []string{pat, jwt} {
|
||||
assert.Equal(t, http.StatusForbidden, auditRequest(router, "GET", "/api/audit", credential).Code)
|
||||
}
|
||||
require.NoError(t, model.DB.Transaction(func(tx *gorm.DB) error {
|
||||
return authz.SetUserPermissionsInTx(tx, admin.Id, authz.PermissionsMap{authz.ResourceAudit: {authz.ActionRead: true}})
|
||||
}))
|
||||
require.NoError(t, authz.ReloadPolicy())
|
||||
for _, credential := range []string{pat, jwt} {
|
||||
for _, endpoint := range []string{"/api/audit", "/api/audit/self"} {
|
||||
response := auditRequest(router, "GET", endpoint+"?category=security&page_size=1", credential)
|
||||
assert.Equal(t, http.StatusOK, response.Code)
|
||||
var result struct {
|
||||
Data struct {
|
||||
Total int
|
||||
Items []model.AuditLog
|
||||
}
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
assert.Equal(t, 2, result.Data.Total)
|
||||
require.Len(t, result.Data.Items, 1)
|
||||
assert.Equal(t, 10, result.Data.Items[0].ActorRole)
|
||||
assert.NotContains(t, response.Body.String(), "root-only")
|
||||
for _, filter := range []string{"request_id=role-100", "username=" + root.Username, "request_id=role-0"} {
|
||||
filtered := auditRequest(router, "GET", endpoint+"?category=security&"+filter, credential)
|
||||
// self ignores supplied usernames, while still excluding every root row.
|
||||
if endpoint == "/api/audit/self" && strings.HasPrefix(filter, "username=") {
|
||||
continue
|
||||
}
|
||||
assert.Contains(t, filtered.Body.String(), `"total":0`)
|
||||
}
|
||||
}
|
||||
}
|
||||
rootSelf := auditRequest(router, "GET", "/api/audit/self?category=security", rootToken)
|
||||
assert.Contains(t, rootSelf.Body.String(), `"actor_role":100`)
|
||||
assert.NotContains(t, rootSelf.Body.String(), "admin_info")
|
||||
rootAll := auditRequest(router, "GET", "/api/audit?category=security", rootToken)
|
||||
assert.Contains(t, rootAll.Body.String(), `"total":7`)
|
||||
assert.Contains(t, rootAll.Body.String(), "root-only")
|
||||
require.NoError(t, authz.SetUserPermissions(admin.Id, authz.PermissionsMap{authz.ResourceAudit: {authz.ActionRead: false}}))
|
||||
for _, credential := range []string{pat, jwt} {
|
||||
assert.Equal(t, http.StatusForbidden, auditRequest(router, "GET", "/api/audit", credential).Code)
|
||||
assert.Equal(t, http.StatusOK, auditRequest(router, "GET", "/api/audit/self", credential).Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuditRoleSnapshotSurvivesActorChanges(t *testing.T) {
|
||||
user, pat := setupAccessTokenAudit(t)
|
||||
user.Role = common.RoleRootUser
|
||||
require.NoError(t, model.DB.Model(user).Update("role", user.Role).Error)
|
||||
router := gin.New()
|
||||
router.Use(middleware.RequestId(), middleware.AccessTokenAudit())
|
||||
router.POST("/change-role", middleware.UserAuth(), func(c *gin.Context) {
|
||||
recordLoginAudit(user, c)
|
||||
recordUserSecurityAudit(c, user.Id, "user.security_verify", nil)
|
||||
recordSubscriptionResetUserLogs(c, &model.SubscriptionResetResult{ResetCount: 1, PlanId: 1, PlanTitle: "Plan", AffectedUserIds: []int{999}}, auditOperatorInfo(c))
|
||||
require.NoError(t, model.DB.Model(user).Update("role", common.RoleAdminUser).Error)
|
||||
c.Status(200)
|
||||
})
|
||||
require.Equal(t, http.StatusOK, auditRequest(router, "POST", "/change-role", pat).Code)
|
||||
entries, total, err := model.GetAuditLogs(model.AuditLogFilter{}, 0, 20, common.RoleAdminUser)
|
||||
require.NoError(t, err)
|
||||
assert.Zero(t, total)
|
||||
assert.Empty(t, entries)
|
||||
status, err := model.GetUserAccessTokenStatus(user.Id)
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, status.LastUsedAt, "a root request must not leak through the last-use summary after demotion")
|
||||
require.NoError(t, model.DB.Unscoped().Delete(user).Error)
|
||||
entries, total, err = model.GetAuditLogs(model.AuditLogFilter{}, 0, 20, common.RoleRootUser)
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 4, total)
|
||||
for _, entry := range entries {
|
||||
assert.Equal(t, common.RoleRootUser, entry.ActorRole)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityAndOperationEventsUseAuditTable(t *testing.T) {
|
||||
user, _ := setupAccessTokenAudit(t)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest("POST", "/event", nil)
|
||||
c.Set("id", user.Id)
|
||||
c.Set("username", user.Username)
|
||||
c.Set("role", user.Role)
|
||||
c.Set(common.RequestIdKey, "correlated-request")
|
||||
recordLoginAudit(user, c)
|
||||
for _, action := range []string{"user.passkey_register", "user.passkey_delete", "user.2fa_setup", "user.2fa_enable", "user.2fa_disable_self", "user.2fa_backup_codes", "user.security_verify"} {
|
||||
recordUserSecurityAudit(c, user.Id, action, nil)
|
||||
}
|
||||
recordManageAudit(c, "option.update", map[string]any{"key": "safe"})
|
||||
recordSubscriptionResetUserLogs(c, &model.SubscriptionResetResult{ResetCount: 1, PlanId: 1, PlanTitle: "Plan", AffectedUserIds: []int{user.Id}}, &model.AuditAdminInfo{AdminID: user.Id})
|
||||
for _, typ := range []int{model.LogTypeTopup, model.LogTypeConsume, model.LogTypeRefund, model.LogTypeSystem} {
|
||||
model.RecordLog(user.Id, typ, "business entry")
|
||||
}
|
||||
var audits []model.AuditLog
|
||||
require.NoError(t, model.LOG_DB.Find(&audits).Error)
|
||||
assert.Len(t, audits, 10)
|
||||
for _, entry := range audits {
|
||||
assert.Equal(t, "correlated-request", entry.RequestId)
|
||||
assert.Equal(t, user.Role, entry.ActorRole)
|
||||
}
|
||||
var logs []model.Log
|
||||
require.NoError(t, model.LOG_DB.Find(&logs).Error)
|
||||
assert.Len(t, logs, 4)
|
||||
_, err := model.DeleteOldLogBatch(context.Background(), time.Now().Unix()+1, 100)
|
||||
require.NoError(t, err)
|
||||
var count int64
|
||||
require.NoError(t, model.LOG_DB.Model(&model.AuditLog{}).Count(&count).Error)
|
||||
assert.EqualValues(t, 10, count)
|
||||
}
|
||||
|
||||
// Released schemas copied from v1.0.0-rc.33; only the Go type names differ.
|
||||
|
||||
type releasedAuditUser struct {
|
||||
Id int `json:"id"`
|
||||
Username string `json:"username" gorm:"unique;index" validate:"max=20"`
|
||||
Password string `json:"password" gorm:"not null;" validate:"min=8,max=20"`
|
||||
OriginalPassword string `json:"original_password" gorm:"-:all"` // this field is only for Password change verification, don't save it to database!
|
||||
DisplayName string `json:"display_name" gorm:"index" validate:"max=20"`
|
||||
Role int `json:"role" gorm:"type:int;default:1"` // admin, common
|
||||
Status int `json:"status" gorm:"type:int;default:1"` // enabled, disabled
|
||||
Email string `json:"email" gorm:"index" validate:"max=50"`
|
||||
GitHubId string `json:"github_id" gorm:"column:github_id;index"`
|
||||
DiscordId string `json:"discord_id" gorm:"column:discord_id;index"`
|
||||
OidcId string `json:"oidc_id" gorm:"column:oidc_id;index"`
|
||||
WeChatId string `json:"wechat_id" gorm:"column:wechat_id;index"`
|
||||
TelegramId string `json:"telegram_id" gorm:"column:telegram_id;index"`
|
||||
VerificationCode string `json:"verification_code" gorm:"-:all"` // this field is only for Email verification, don't save it to database!
|
||||
AccessToken *string `json:"-" gorm:"type:char(32);column:access_token;uniqueIndex"` // this token is for system management
|
||||
Quota int `json:"quota" gorm:"type:int;default:0"`
|
||||
UsedQuota int `json:"used_quota" gorm:"type:int;default:0;column:used_quota"` // used quota
|
||||
RequestCount int `json:"request_count" gorm:"type:int;default:0;"` // request number
|
||||
Group string `json:"group" gorm:"type:varchar(64);default:'default'"`
|
||||
AffCode string `json:"aff_code" gorm:"type:varchar(32);column:aff_code;uniqueIndex"`
|
||||
AffCount int `json:"aff_count" gorm:"type:int;default:0;column:aff_count"`
|
||||
AffQuota int `json:"aff_quota" gorm:"type:int;default:0;column:aff_quota"` // 邀请剩余额度
|
||||
AffHistoryQuota int `json:"aff_history_quota" gorm:"type:int;default:0;column:aff_history"` // 邀请历史额度
|
||||
InviterId int `json:"inviter_id" gorm:"type:int;column:inviter_id;index"`
|
||||
DeletedAt gorm.DeletedAt `gorm:"index"`
|
||||
LinuxDOId string `json:"linux_do_id" gorm:"column:linux_do_id;index"`
|
||||
Setting string `json:"setting" gorm:"type:text;column:setting"`
|
||||
Remark string `json:"remark,omitempty" gorm:"type:varchar(255)" validate:"max=255"`
|
||||
StripeCustomer string `json:"stripe_customer" gorm:"type:varchar(64);column:stripe_customer;index"`
|
||||
CreatedAt int64 `json:"created_at" gorm:"autoCreateTime;column:created_at"`
|
||||
LastLoginAt int64 `json:"last_login_at" gorm:"default:0;column:last_login_at"`
|
||||
AuthVersion int64 `json:"-" gorm:"type:bigint;not null;default:1;column:auth_version"`
|
||||
AdminPermissions map[string]map[string]bool `json:"admin_permissions,omitempty" gorm:"-:all"`
|
||||
}
|
||||
|
||||
func (releasedAuditUser) TableName() string { return "users" }
|
||||
|
||||
type releasedAuditLog struct {
|
||||
Id int `json:"id" gorm:"index:idx_created_at_id,priority:2;index:idx_user_id_id,priority:2"`
|
||||
UserId int `json:"user_id" gorm:"index;index:idx_user_id_id,priority:1"`
|
||||
CreatedAt int64 `json:"created_at" gorm:"bigint;index:idx_created_at_id,priority:1;index:idx_created_at_type"`
|
||||
Type int `json:"type" gorm:"index:idx_created_at_type"`
|
||||
Content string `json:"content"`
|
||||
Username string `json:"username" gorm:"index;index:index_username_model_name,priority:2;default:''"`
|
||||
TokenName string `json:"token_name" gorm:"index;default:''"`
|
||||
ModelName string `json:"model_name" gorm:"index;index:index_username_model_name,priority:1;default:''"`
|
||||
Quota int `json:"quota" gorm:"default:0"`
|
||||
PromptTokens int `json:"prompt_tokens" gorm:"default:0"`
|
||||
CompletionTokens int `json:"completion_tokens" gorm:"default:0"`
|
||||
UseTime int `json:"use_time" gorm:"default:0"`
|
||||
IsStream bool `json:"is_stream"`
|
||||
ChannelId int `json:"channel" gorm:"index"`
|
||||
ChannelName string `json:"channel_name" gorm:"->"`
|
||||
TokenId int `json:"token_id" gorm:"default:0;index"`
|
||||
Group string `json:"group" gorm:"index"`
|
||||
Ip string `json:"ip" gorm:"index;default:''"`
|
||||
RequestId string `json:"request_id,omitempty" gorm:"type:varchar(64);index:idx_logs_request_id;default:''"`
|
||||
UpstreamRequestId string `json:"upstream_request_id,omitempty" gorm:"type:varchar(128);index:idx_logs_upstream_request_id;default:''"`
|
||||
Other string `json:"other"`
|
||||
}
|
||||
|
||||
func (releasedAuditLog) TableName() string { return "logs" }
|
||||
|
||||
// External tests create a new database per case on a loopback-only disposable
|
||||
// instance. They never drop databases or tables supplied through an environment variable.
|
||||
func newAuditTestDatabase(t *testing.T, kind, dsn string) (*gorm.DB, string) {
|
||||
t.Helper()
|
||||
if kind == "sqlite" {
|
||||
path := t.TempDir() + "/audit.db"
|
||||
db, err := gorm.Open(sqlite.Open(path), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
return db, path
|
||||
}
|
||||
require.NotEmpty(t, dsn)
|
||||
name := fmt.Sprintf("newapi_audit_%d", time.Now().UnixNano())
|
||||
var original, isolated gorm.Dialector
|
||||
var newDSN string
|
||||
if kind == "mysql" {
|
||||
config, err := sqlmysql.ParseDSN(dsn)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "tcp", config.Net)
|
||||
host, _, err := net.SplitHostPort(config.Addr)
|
||||
require.NoError(t, err)
|
||||
require.True(t, net.ParseIP(host).IsLoopback(), "database tests only permit loopback instances")
|
||||
original = mysql.Open(dsn)
|
||||
config.DBName = name
|
||||
newDSN = config.FormatDSN()
|
||||
isolated = mysql.Open(newDSN)
|
||||
} else {
|
||||
parsed, err := url.Parse(dsn)
|
||||
require.NoError(t, err)
|
||||
require.True(t, net.ParseIP(parsed.Hostname()).IsLoopback(), "database tests only permit loopback instances")
|
||||
parsed.Path = "/" + name
|
||||
newDSN = parsed.String()
|
||||
if kind == "clickhouse" {
|
||||
original = clickhouse.Open(dsn)
|
||||
isolated = clickhouse.Open(newDSN)
|
||||
} else {
|
||||
original = postgres.Open(dsn)
|
||||
isolated = postgres.Open(newDSN)
|
||||
}
|
||||
}
|
||||
admin, err := gorm.Open(original, &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
// No IF NOT EXISTS: a collision fails before any test data can be written.
|
||||
createSQL := "CREATE DATABASE " + name
|
||||
if kind == "mysql" {
|
||||
createSQL += " CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci"
|
||||
}
|
||||
require.NoError(t, admin.Exec(createSQL).Error)
|
||||
sqlDB, err := admin.DB()
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, sqlDB.Close())
|
||||
db, err := gorm.Open(isolated, &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
t.Logf("isolated database: %s (%s)", name, kind)
|
||||
t.Cleanup(func() {
|
||||
connection, err := db.DB()
|
||||
if err == nil {
|
||||
_ = connection.Close()
|
||||
}
|
||||
})
|
||||
return db, newDSN
|
||||
}
|
||||
|
||||
func verifyAuditRoleStorage(t *testing.T) {
|
||||
t.Helper()
|
||||
for i, role := range []int{1, 10, 100, 0, 99} {
|
||||
model.RecordAuditLog(nil, model.AuditLog{ActorRole: role, UserId: 1, Username: "role-owner", CreatedAt: int64(200 + i), Category: model.AuditCategoryOperation, RequestId: fmt.Sprintf("matrix-role-%d", role)})
|
||||
}
|
||||
filter := model.AuditLogFilter{Category: model.AuditCategoryOperation}
|
||||
visible, total, err := model.GetAuditLogs(filter, 0, 1, common.RoleAdminUser)
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 2, total)
|
||||
require.Len(t, visible, 1)
|
||||
assert.Equal(t, common.RoleAdminUser, visible[0].ActorRole)
|
||||
visible, total, err = model.GetAuditLogs(filter, 1, 1, common.RoleCommonUser)
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 2, total)
|
||||
require.Len(t, visible, 1)
|
||||
assert.Equal(t, common.RoleCommonUser, visible[0].ActorRole)
|
||||
filter.RequestId = "matrix-role-100"
|
||||
visible, total, err = model.GetAuditLogs(filter, 0, 20, common.RoleAdminUser)
|
||||
require.NoError(t, err)
|
||||
assert.Zero(t, total)
|
||||
assert.Empty(t, visible)
|
||||
visible, total, err = model.GetAuditLogs(filter, 0, 20, common.RoleRootUser)
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 1, total)
|
||||
require.Len(t, visible, 1)
|
||||
assert.Equal(t, common.RoleRootUser, visible[0].ActorRole)
|
||||
}
|
||||
|
||||
func verifyAuditJSONStorage(t *testing.T) {
|
||||
t.Helper()
|
||||
columns, err := model.LOG_DB.Migrator().ColumnTypes(&model.AuditLog{})
|
||||
require.NoError(t, err)
|
||||
var otherType string
|
||||
for _, column := range columns {
|
||||
if column.Name() == "other" {
|
||||
otherType = strings.ToLower(column.DatabaseTypeName())
|
||||
}
|
||||
}
|
||||
assert.Equal(t, "json", otherType)
|
||||
|
||||
metadata := model.AuditOther{
|
||||
Op: &model.AuditOperation{Action: "channel.update", Params: model.AuditFields{
|
||||
"id": 42, "name": "渠道", "changed_fields": []string{}, "large_id": uint64(9007199254740993),
|
||||
"extra": map[string]any{"attempts": 0, "permitted": false, "ratio": 1.25, "targets": []int{1, 2}},
|
||||
}},
|
||||
AdminInfo: &model.AuditAdminInfo{AdminID: 1},
|
||||
AuditInfo: &model.AuditRequestInfo{Method: "PUT", Route: "/api/channel/", Path: "/api/channel/", Status: 200, Success: false},
|
||||
RootInfo: model.AuditFields{"private": "root-only"},
|
||||
}
|
||||
model.RecordAuditLog(nil, model.AuditLog{ActorRole: common.RoleCommonUser, UserId: 1, Username: "json-owner", Category: model.AuditCategorySecurity, RequestId: "matrix-json", Other: metadata})
|
||||
filter := model.AuditLogFilter{RequestId: "matrix-json"}
|
||||
entries, total, err := model.GetAuditLogs(filter, 0, 20, common.RoleRootUser)
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, 1, total)
|
||||
require.Len(t, entries, 1)
|
||||
stored, err := common.Marshal(entries[0].Other)
|
||||
require.NoError(t, err)
|
||||
expected, err := common.Marshal(metadata)
|
||||
require.NoError(t, err)
|
||||
assert.JSONEq(t, string(expected), string(stored))
|
||||
require.NotNil(t, entries[0].Other.Op)
|
||||
assert.Equal(t, "channel.update", entries[0].Other.Op.Action)
|
||||
require.NotNil(t, entries[0].Other.AuditInfo)
|
||||
assert.False(t, entries[0].Other.AuditInfo.Success)
|
||||
var details struct {
|
||||
Op struct {
|
||||
Params struct {
|
||||
LargeId uint64 `json:"large_id"`
|
||||
}
|
||||
}
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(stored, &details))
|
||||
assert.EqualValues(t, 9007199254740993, details.Op.Params.LargeId, "JSON numbers must retain their type and precision")
|
||||
for _, role := range []int{common.RoleCommonUser, common.RoleAdminUser} {
|
||||
entries, _, err = model.GetAuditLogs(filter, 0, 20, role)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, entries, 1)
|
||||
assert.Nil(t, entries[0].Other.RootInfo)
|
||||
if role == common.RoleCommonUser {
|
||||
assert.Nil(t, entries[0].Other.AdminInfo)
|
||||
assert.Nil(t, entries[0].Other.AuditInfo)
|
||||
} else {
|
||||
assert.NotNil(t, entries[0].Other.AdminInfo)
|
||||
assert.NotNil(t, entries[0].Other.AuditInfo)
|
||||
}
|
||||
}
|
||||
model.RecordAuditLog(nil, model.AuditLog{ActorRole: common.RoleCommonUser, UserId: 1, Username: "json-owner", RequestId: "matrix-json-empty", Other: model.AuditOther{}})
|
||||
entries, _, err = model.GetAuditLogs(model.AuditLogFilter{RequestId: "matrix-json-empty"}, 0, 20, common.RoleRootUser)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, entries, 1)
|
||||
empty, err := common.Marshal(entries[0].Other)
|
||||
require.NoError(t, err)
|
||||
assert.JSONEq(t, `{}`, string(empty))
|
||||
for range 2 {
|
||||
require.NoError(t, model.InitLogDB())
|
||||
}
|
||||
entries, total, err = model.GetAuditLogs(filter, 0, 20, common.RoleRootUser)
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, 1, total)
|
||||
require.Len(t, entries, 1)
|
||||
stored, err = common.Marshal(entries[0].Other)
|
||||
require.NoError(t, err)
|
||||
assert.JSONEq(t, string(expected), string(stored), "repeated startup must retain structured audit metadata")
|
||||
}
|
||||
|
||||
func TestAuditOtherDatabaseEncoding(t *testing.T) {
|
||||
const payload = `{"op":{"action":"channel.update","params":{"id":9007199254740993,"nested":{"count":0,"enabled":false}}},"admin_info":{"admin_id":1,"admin_username":"root","admin_role":100,"auth_method":"session"},"audit_info":{"method":"PUT","route":"/api/channel/","path":"/api/channel/","status":200,"success":false},"root_info":{"generation":18446744073709551615}}`
|
||||
for _, input := range []any{payload, []byte(payload)} {
|
||||
var other model.AuditOther
|
||||
require.NoError(t, other.Scan(input))
|
||||
require.NotNil(t, other.Op)
|
||||
assert.Equal(t, "channel.update", other.Op.Action)
|
||||
require.NotNil(t, other.AdminInfo)
|
||||
assert.Equal(t, 100, other.AdminInfo.AdminRole)
|
||||
require.NotNil(t, other.AuditInfo)
|
||||
assert.False(t, other.AuditInfo.Success)
|
||||
encoded, err := other.Value()
|
||||
require.NoError(t, err)
|
||||
text, ok := encoded.(string)
|
||||
require.True(t, ok, "PostgreSQL simple protocol requires a string parameter")
|
||||
assert.JSONEq(t, payload, text)
|
||||
assert.Contains(t, text, "9007199254740993")
|
||||
assert.Contains(t, text, "18446744073709551615")
|
||||
for _, empty := range []any{nil, "", []byte(`null`), "{}"} {
|
||||
require.NoError(t, other.Scan(input))
|
||||
require.NoError(t, other.Scan(empty))
|
||||
encoded, err = other.Value()
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "{}", encoded, "empty input must clear fields from the previous row")
|
||||
}
|
||||
}
|
||||
var invalid model.AuditOther
|
||||
assert.Error(t, invalid.Scan(42))
|
||||
assert.Error(t, invalid.Scan([]byte(`{broken`)))
|
||||
}
|
||||
|
||||
func TestAuditDatabaseMatrix(t *testing.T) {
|
||||
previousDB, previousLogDB := model.DB, model.LOG_DB
|
||||
previousMain, previousLog := common.MainDatabaseType(), common.LogDatabaseType()
|
||||
previousRedis := common.RedisEnabled
|
||||
previousMaster, previousSQLite := common.IsMasterNode, common.SQLitePath
|
||||
common.IsMasterNode = true
|
||||
common.RedisEnabled = false
|
||||
t.Cleanup(func() {
|
||||
common.IsMasterNode, common.SQLitePath = previousMaster, previousSQLite
|
||||
model.DB, model.LOG_DB = previousDB, previousLogDB
|
||||
common.SetDatabaseTypes(previousMain, previousLog)
|
||||
common.RedisEnabled = previousRedis
|
||||
})
|
||||
cases := []struct {
|
||||
name, env string
|
||||
typ common.DatabaseType
|
||||
}{
|
||||
{"sqlite", "", common.DatabaseTypeSQLite}, {"mysql", "AUDIT_MYSQL_DSN", common.DatabaseTypeMySQL}, {"postgres", "AUDIT_POSTGRES_DSN", common.DatabaseTypePostgreSQL},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
dsn := os.Getenv(tc.env)
|
||||
if tc.env != "" && dsn == "" {
|
||||
t.Skip(tc.env + " is not configured")
|
||||
}
|
||||
for _, upgrade := range []bool{false, true} {
|
||||
t.Run(fmt.Sprintf("upgrade=%v", upgrade), func(t *testing.T) {
|
||||
db, isolatedDSN := newAuditTestDatabase(t, tc.name, dsn)
|
||||
t.Setenv("LOG_SQL_DSN", "")
|
||||
if tc.name == "sqlite" {
|
||||
common.SQLitePath = isolatedDSN
|
||||
t.Setenv("SQL_DSN", "local")
|
||||
} else {
|
||||
t.Setenv("SQL_DSN", isolatedDSN)
|
||||
}
|
||||
model.DB, model.LOG_DB = db, db
|
||||
common.SetDatabaseTypes(tc.typ, tc.typ)
|
||||
versionSQL := "SELECT version()"
|
||||
if tc.name == "sqlite" {
|
||||
versionSQL = "SELECT sqlite_version()"
|
||||
}
|
||||
var version string
|
||||
require.NoError(t, db.Raw(versionSQL).Scan(&version).Error)
|
||||
t.Logf("database version: %s", version)
|
||||
if upgrade {
|
||||
require.NoError(t, db.AutoMigrate(&releasedAuditUser{}, &releasedAuditLog{}))
|
||||
legacy := "released-token"
|
||||
require.NoError(t, db.Create(&releasedAuditUser{Username: "released-owner", Password: "placeholder", AccessToken: &legacy, AffCode: "released-aff", Quota: 1234}).Error)
|
||||
require.NoError(t, db.Create(&releasedAuditLog{UserId: 1, Type: model.LogTypeLogin, Content: "historical login", CreatedAt: 100, RequestId: "legacy-request"}).Error)
|
||||
}
|
||||
for range 2 {
|
||||
require.NoError(t, model.InitDB())
|
||||
require.NoError(t, model.InitLogDB())
|
||||
}
|
||||
if !upgrade {
|
||||
require.NoError(t, db.Create(&model.User{Username: "fresh-owner", Password: "placeholder", AffCode: "fresh-aff"}).Error)
|
||||
}
|
||||
status, err := model.GetUserAccessTokenStatus(1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, upgrade, status.Exists)
|
||||
assert.Nil(t, status.CreatedAt)
|
||||
if upgrade {
|
||||
assert.Equal(t, model.AccessTokenFingerprint("released-token"), status.TokenRef)
|
||||
legacyUser, validationErr := model.ValidateAccessToken("released-token")
|
||||
require.NoError(t, validationErr)
|
||||
require.NotNil(t, legacyUser)
|
||||
var user model.User
|
||||
require.NoError(t, db.First(&user, 1).Error)
|
||||
assert.Equal(t, 1234, user.Quota)
|
||||
var old model.Log
|
||||
require.NoError(t, db.First(&old).Error)
|
||||
assert.Equal(t, "historical login", old.Content)
|
||||
}
|
||||
require.NoError(t, model.UpdateUserAccessToken(1, "matrix-token"))
|
||||
require.Error(t, db.Create(&model.User{Username: "duplicate", Password: "placeholder", AffCode: "duplicate-aff", AccessToken: common.GetPointer("matrix-token")}).Error, "PAT uniqueness must survive upgrade")
|
||||
for _, timestamp := range []int64{101, 102, 103} {
|
||||
model.RecordAuditLog(nil, model.AuditLog{ActorRole: common.RoleAdminUser, UserId: 1, Username: "owner", CreatedAt: timestamp, Category: model.AuditCategoryAccessToken, TokenRef: model.AccessTokenFingerprint("matrix-token"), Ip: "192.0.2.1", Success: timestamp != 102})
|
||||
}
|
||||
first, total, err := model.GetAuditLogs(model.AuditLogFilter{UserId: 1}, 0, 2, common.RoleCommonUser)
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 3, total)
|
||||
require.Len(t, first, 2)
|
||||
assert.EqualValues(t, 103, first[0].CreatedAt)
|
||||
second, _, err := model.GetAuditLogs(model.AuditLogFilter{UserId: 1}, 2, 2, common.RoleCommonUser)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, second, 1)
|
||||
assert.EqualValues(t, 101, second[0].CreatedAt)
|
||||
failure := false
|
||||
failed, _, err := model.GetAuditLogs(model.AuditLogFilter{UserId: 1, Success: &failure}, 0, 10, 1)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, failed, 1)
|
||||
assert.EqualValues(t, 102, failed[0].CreatedAt)
|
||||
status, err = model.GetUserAccessTokenStatus(1)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, status.LastUsedAt)
|
||||
assert.EqualValues(t, 103, *status.LastUsedAt)
|
||||
_, err = model.RevokeUserAccessToken(1)
|
||||
require.NoError(t, err)
|
||||
_, err = model.RevokeUserAccessToken(1)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, model.MigrateAuditLogs())
|
||||
_, total, err = model.GetAuditLogs(model.AuditLogFilter{UserId: 1}, 0, 10, 1)
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 3, total)
|
||||
verifyAuditRoleStorage(t)
|
||||
verifyAuditJSONStorage(t)
|
||||
require.NoError(t, authz.Init(model.DB))
|
||||
assert.False(t, authz.Can(1, common.RoleAdminUser, authz.AuditRead))
|
||||
require.NoError(t, model.DB.Transaction(func(tx *gorm.DB) error {
|
||||
return authz.SetUserPermissionsInTx(tx, 1, authz.PermissionsMap{authz.ResourceAudit: {authz.ActionRead: true}})
|
||||
}))
|
||||
require.NoError(t, authz.ReloadPolicy())
|
||||
assert.True(t, authz.Can(1, common.RoleAdminUser, authz.AuditRead))
|
||||
require.NoError(t, authz.Init(model.DB))
|
||||
require.NoError(t, authz.Init(model.DB))
|
||||
assert.True(t, authz.Can(1, common.RoleAdminUser, authz.AuditRead))
|
||||
require.NoError(t, model.DB.Transaction(func(tx *gorm.DB) error {
|
||||
return authz.SetUserPermissionsInTx(tx, 1, authz.PermissionsMap{authz.ResourceAudit: {authz.ActionRead: false}})
|
||||
}))
|
||||
require.NoError(t, authz.ReloadPolicy())
|
||||
assert.False(t, authz.Can(1, common.RoleAdminUser, authz.AuditRead))
|
||||
assert.True(t, authz.Can(1, common.RoleRootUser, authz.AuditRead))
|
||||
})
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestIndependentAuditLogStores(t *testing.T) {
|
||||
_, _ = setupAccessTokenAudit(t)
|
||||
previousMaster := common.IsMasterNode
|
||||
common.IsMasterNode = true
|
||||
t.Cleanup(func() { common.IsMasterNode = previousMaster })
|
||||
for _, tc := range []struct{ kind, env string }{{"mysql", "AUDIT_MYSQL_DSN"}, {"postgres", "AUDIT_POSTGRES_DSN"}, {"clickhouse", "AUDIT_CLICKHOUSE_DSN"}} {
|
||||
for _, upgrade := range []bool{false, true} {
|
||||
t.Run(fmt.Sprintf("%s/upgrade=%v", tc.kind, upgrade), func(t *testing.T) {
|
||||
dsn := os.Getenv(tc.env)
|
||||
if dsn == "" {
|
||||
t.Skip(tc.env + " is not configured")
|
||||
}
|
||||
logDB, isolatedDSN := newAuditTestDatabase(t, tc.kind, dsn)
|
||||
if upgrade {
|
||||
if tc.kind == "clickhouse" {
|
||||
require.NoError(t, logDB.Exec(releasedClickHouseLogSchema).Error)
|
||||
} else {
|
||||
require.NoError(t, logDB.AutoMigrate(&releasedAuditLog{}))
|
||||
}
|
||||
require.NoError(t, logDB.Create(&releasedAuditLog{UserId: 1, Username: "legacy", CreatedAt: time.Now().Unix(), Type: model.LogTypeLogin, Content: "retained historical login", RequestId: "legacy-split-request"}).Error)
|
||||
}
|
||||
t.Setenv("LOG_SQL_DSN", isolatedDSN)
|
||||
t.Setenv("LOG_SQL_CLICKHOUSE_TTL_DAYS", "7")
|
||||
require.NoError(t, model.InitLogDB())
|
||||
require.NoError(t, model.InitLogDB())
|
||||
if upgrade {
|
||||
var old model.Log
|
||||
require.NoError(t, model.LOG_DB.Where("request_id = ?", "legacy-split-request").Take(&old).Error)
|
||||
assert.Equal(t, "retained historical login", old.Content)
|
||||
}
|
||||
require.NoError(t, model.UpdateUserAccessToken(1, "independent-pat"))
|
||||
model.RecordAuditLog(nil, model.AuditLog{ActorRole: common.RoleAdminUser, UserId: 1, Username: "independent", Category: model.AuditCategoryAccessToken, TokenRef: model.AccessTokenFingerprint("independent-pat"), Ip: "192.0.2.8", Success: false, Status: 403})
|
||||
entries, total, err := model.GetAuditLogs(model.AuditLogFilter{UserId: 1}, 0, 20, common.RoleAdminUser)
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 1, total)
|
||||
require.Len(t, entries, 1)
|
||||
assert.False(t, entries[0].Success)
|
||||
status, err := model.GetUserAccessTokenStatus(1)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, status.LastUsedAt)
|
||||
assert.Equal(t, "192.0.2.8", status.LastUsedIp)
|
||||
model.RecordLog(1, model.LogTypeTopup, "independent business")
|
||||
_, err = model.DeleteOldLogBatch(context.Background(), time.Now().Unix()+1, 100)
|
||||
require.NoError(t, err)
|
||||
_, total, err = model.GetAuditLogs(model.AuditLogFilter{UserId: 1}, 0, 20, 1)
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 1, total)
|
||||
var mainCount int64
|
||||
require.NoError(t, model.DB.Model(&model.AuditLog{}).Count(&mainCount).Error)
|
||||
assert.Zero(t, mainCount)
|
||||
verifyAuditRoleStorage(t)
|
||||
verifyAuditJSONStorage(t)
|
||||
if tc.kind == "clickhouse" {
|
||||
var create string
|
||||
require.NoError(t, model.LOG_DB.Raw("SHOW CREATE TABLE audit_logs").Scan(&create).Error)
|
||||
assert.NotContains(t, strings.ToUpper(create), "TTL")
|
||||
require.NoError(t, model.LOG_DB.Raw("SHOW CREATE TABLE logs").Scan(&create).Error)
|
||||
assert.Contains(t, strings.ToUpper(create), "TTL")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ClickHouse logs schema from v1.0.0-rc.33.
|
||||
const releasedClickHouseLogSchema = `
|
||||
CREATE TABLE IF NOT EXISTS logs (
|
||||
id Int64 DEFAULT 0,
|
||||
user_id Int32 DEFAULT 0,
|
||||
created_at Int64 DEFAULT 0,
|
||||
type Int32 DEFAULT 0,
|
||||
content String DEFAULT '',
|
||||
username String DEFAULT '',
|
||||
token_name String DEFAULT '',
|
||||
model_name String DEFAULT '',
|
||||
quota Int32 DEFAULT 0,
|
||||
prompt_tokens Int32 DEFAULT 0,
|
||||
completion_tokens Int32 DEFAULT 0,
|
||||
use_time Int32 DEFAULT 0,
|
||||
is_stream UInt8 DEFAULT 0,
|
||||
channel_id Int32 DEFAULT 0,
|
||||
token_id Int32 DEFAULT 0,
|
||||
` + "`group`" + ` String DEFAULT '',
|
||||
ip String DEFAULT '',
|
||||
request_id String DEFAULT '',
|
||||
upstream_request_id String DEFAULT '',
|
||||
other String DEFAULT ''
|
||||
)
|
||||
ENGINE = MergeTree()
|
||||
PARTITION BY toYYYYMM(toDateTime(created_at))
|
||||
ORDER BY (created_at, request_id)`
|
||||
@@ -0,0 +1,162 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// auditContentTemplates 将稳定的操作标识 action 映射为英文兜底模板，渲染后写入
|
||||
// Log.Content（供导出等非本地化消费者使用）。占位符为 ${name}，由该
|
||||
// action 的 params 填充。本地化展示文案在前端 i18n 模板中维护，本表是语言中立的
|
||||
// 英文基线——调用方因此无需在每个埋点处手写句子（避免与 params 重复书写同一份值）。
|
||||
var auditContentTemplates = map[string]string{
|
||||
"user.create": "Created user ${username} (role ${role})",
|
||||
"user.update": "Updated user ${username} (ID: ${id})",
|
||||
"user.delete": "Deleted user ${username} (ID: ${id})",
|
||||
"user.account_delete": "Account deletion",
|
||||
"user.manage": "Performed ${action} on user ${username} (ID: ${id})",
|
||||
"user.quota_add": "Increased user quota by ${quota}",
|
||||
"user.quota_subtract": "Decreased user quota by ${quota}",
|
||||
"user.quota_override": "Overrode user quota from ${from} to ${to}",
|
||||
"user.binding_clear": "Cleared ${bindingType} binding for user ${username}",
|
||||
"user.2fa_disable": "Force-disabled two-factor authentication for the user",
|
||||
"user.passkey_register": "Registered a passkey",
|
||||
"access_token.generate": "Generated a system access token",
|
||||
"access_token.revoke": "Revoked the system access token",
|
||||
"user.2fa_setup": "Started two-factor authentication setup",
|
||||
"user.2fa_enable": "Enabled two-factor authentication",
|
||||
"user.2fa_disable_self": "Disabled two-factor authentication",
|
||||
"user.2fa_backup_codes": "Regenerated two-factor backup codes",
|
||||
"user.security_verify": "Completed security verification",
|
||||
"user.password_change": "Account password change",
|
||||
"user.binding_start": "Account binding request",
|
||||
"user.binding_bind": "Account binding",
|
||||
"user.binding_unbind": "Account unlinking",
|
||||
"user.email_binding_resend": "Email confirmation code resend",
|
||||
"user.passkey_delete": "Deleted a passkey",
|
||||
"user.reset_passkey": "Reset the user passkey",
|
||||
"option.update": "Updated system setting ${key}",
|
||||
|
||||
"channel.create": "Created channel ${name} (type ${type}, count ${count})",
|
||||
"channel.update": "Updated channel ${name} (ID: ${id})",
|
||||
"channel.delete": "Deleted channel ${name} (ID: ${id})",
|
||||
"channel.delete_batch": "Batch deleted ${count} channels",
|
||||
"channel.delete_disabled": "Deleted all disabled channels (${count})",
|
||||
"channel.key_view": "Viewed channel key ${name} (ID: ${id})",
|
||||
"channel.tag_disable": "Disabled channels with tag ${tag}",
|
||||
"channel.tag_enable": "Enabled channels with tag ${tag}",
|
||||
"channel.tag_edit": "Edited channels with tag ${tag}",
|
||||
"channel.tag_batch_set": "Batch set tag for ${count} channels",
|
||||
"channel.copy": "Copied channel (source ID: ${sourceId}) to ${name} (new ID: ${id})",
|
||||
"channel.multi_key_manage": "Multi-key management ${action} on channel (ID: ${id})",
|
||||
"channel.upstream_apply": "Applied upstream model changes to channel (ID: ${id})",
|
||||
"channel.upstream_apply_all": "Applied upstream model changes to ${count} channels",
|
||||
|
||||
"redemption.create": "Created ${count} redemption codes named ${name} (${quota} each)",
|
||||
"redemption.delete_batch": "Batch deleted ${count} redemption codes",
|
||||
|
||||
"subscription.plan_reset": "Reset active subscriptions for plan ${plan_id}",
|
||||
"subscription.user_plan_reset": "Reset active plan ${plan_id} subscriptions for user ${target_user_id}",
|
||||
}
|
||||
|
||||
// auditContentEN 按 action 模板渲染英文兜底文本；未登记的 action 退回 action 本身。
|
||||
func auditContentEN(action string, params map[string]any) string {
|
||||
tmpl, ok := auditContentTemplates[action]
|
||||
if !ok {
|
||||
return action
|
||||
}
|
||||
return os.Expand(tmpl, func(key string) string {
|
||||
if v, ok := params[key]; ok {
|
||||
return fmt.Sprintf("%v", v)
|
||||
}
|
||||
return ""
|
||||
})
|
||||
}
|
||||
|
||||
// auditOperatorInfo 从上下文构建操作者身份信息（管理员 id/用户名/角色）。
|
||||
func auditOperatorInfo(c *gin.Context) *model.AuditAdminInfo {
|
||||
return &model.AuditAdminInfo{
|
||||
AdminID: c.GetInt("id"),
|
||||
AdminUsername: c.GetString("username"),
|
||||
AdminRole: c.GetInt("role"),
|
||||
AuthMethod: auditAuthMethod(c),
|
||||
}
|
||||
}
|
||||
|
||||
func auditAuthMethod(c *gin.Context) string {
|
||||
if c.GetBool("use_access_token") {
|
||||
return "access_token"
|
||||
}
|
||||
return "session"
|
||||
}
|
||||
|
||||
// markAuditLogged 标记当前请求已在 handler 内手动记录审计日志，
|
||||
// 使鉴权链路中的审计兜底（finishAdminAudit）跳过兜底记录，避免重复。
|
||||
func markAuditLogged(c *gin.Context) {
|
||||
common.SetContextKey(c, constant.ContextKeyAuditLogged, true)
|
||||
}
|
||||
|
||||
// recordManageAudit 记录一条由操作者本人归属的管理/高危审计日志（资源类操作：
|
||||
// 渠道 / 系统设置 / 兑换码等）。content 由 action+params 自动渲染。
|
||||
func recordManageAudit(c *gin.Context, action string, params map[string]any) {
|
||||
recordManageAuditFor(c, c.GetInt("id"), action, params)
|
||||
}
|
||||
|
||||
// recordManageAuditFor 记录一条管理审计日志，日志归属于操作者；targetUserId
|
||||
// 只表示被操作用户，用于在结构化参数中保留目标上下文。
|
||||
func recordManageAuditFor(c *gin.Context, targetUserId int, action string, params map[string]any) {
|
||||
if params == nil {
|
||||
params = map[string]any{}
|
||||
}
|
||||
operatorUserId := c.GetInt("id")
|
||||
if _, ok := params["target_user_id"]; !ok && targetUserId > 0 && targetUserId != operatorUserId {
|
||||
params["target_user_id"] = targetUserId
|
||||
}
|
||||
model.RecordOperationAuditLog(operatorUserId, c.GetInt("role"), auditContentEN(action, params), c.ClientIP(), action, params, auditOperatorInfo(c), nil, c)
|
||||
markAuditLogged(c)
|
||||
}
|
||||
|
||||
// recordUserSecurityAudit 记录普通用户自己的安全敏感操作（如 passkey 绑定/解绑）。
|
||||
// 这类日志没有管理员操作者，不写 admin_info；同时不依赖 AdminAuth/RootAuth 的兜底。
|
||||
func recordUserSecurityAudit(c *gin.Context, userId int, action string, params map[string]any) {
|
||||
if code := c.GetString("security_error_code"); code != "" {
|
||||
if params == nil {
|
||||
params = map[string]any{}
|
||||
}
|
||||
params["code"] = code
|
||||
}
|
||||
var auditInfo *model.AuditRequestInfo
|
||||
if success, ok := params["success"].(bool); ok {
|
||||
auditInfo = &model.AuditRequestInfo{
|
||||
Method: c.Request.Method, Route: c.FullPath(), Path: c.FullPath(),
|
||||
Status: c.Writer.Status(), Success: success,
|
||||
}
|
||||
}
|
||||
model.RecordOperationAuditLog(userId, c.GetInt("role"), auditContentEN(action, params), c.ClientIP(), action, params, nil, auditInfo, c)
|
||||
}
|
||||
|
||||
func tokenAuditParams(c *gin.Context) model.AuditFields {
|
||||
params, ok := common.GetContextKeyType[model.AuditFields](c, constant.ContextKeyTokenAuditParams)
|
||||
if !ok {
|
||||
params = model.AuditFields{}
|
||||
common.SetContextKey(c, constant.ContextKeyTokenAuditParams, params)
|
||||
}
|
||||
return params
|
||||
}
|
||||
|
||||
func tokenBatchAuditParams(c *gin.Context, ids []int) model.AuditFields {
|
||||
params := tokenAuditParams(c)
|
||||
params["total"] = len(ids)
|
||||
// Bound audit payloads without changing the batch operation's limits.
|
||||
params["requested_ids"] = append([]int{}, ids[:min(len(ids), 100)]...)
|
||||
if len(ids) > 100 {
|
||||
params["requested_ids_truncated"] = true
|
||||
}
|
||||
return params
|
||||
}
|
||||
@@ -0,0 +1,818 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/ecdsa"
|
||||
"crypto/elliptic"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/oauth"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
"github.com/fxamacker/cbor/v2"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/pquerna/otp/totp"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func newSecurityLoginPasskey(t *testing.T, userID int) *ecdsa.PrivateKey {
|
||||
t.Helper()
|
||||
key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
|
||||
require.NoError(t, err)
|
||||
credentialID := sha256.Sum256(elliptic.Marshal(key.Curve, key.X, key.Y))
|
||||
publicKey, err := cbor.Marshal(map[int]any{1: 2, 3: -7, -1: 1, -2: key.X.FillBytes(make([]byte, 32)), -3: key.Y.FillBytes(make([]byte, 32))})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, model.DB.Create(&model.PasskeyCredential{
|
||||
UserID: userID, CredentialID: base64.StdEncoding.EncodeToString(credentialID[:]), PublicKey: base64.StdEncoding.EncodeToString(publicKey),
|
||||
UserPresent: true, UserVerified: true,
|
||||
}).Error)
|
||||
return key
|
||||
}
|
||||
|
||||
func beginSecurityLoginPasskey(t *testing.T, parentToken string) (string, string) {
|
||||
t.Helper()
|
||||
body, err := common.Marshal(map[string]string{"flow_token": parentToken})
|
||||
require.NoError(t, err)
|
||||
response := securityEnrollmentRequest("POST", "/api/user/login/passkey/begin", string(body), "", service.AuthIdentity{}, LoginPasskeyBegin)
|
||||
var result struct {
|
||||
Success bool `json:"success"`
|
||||
Data struct {
|
||||
FlowToken string `json:"flow_token"`
|
||||
Options struct {
|
||||
PublicKey struct {
|
||||
Challenge string `json:"challenge"`
|
||||
UserVerification string `json:"userVerification"`
|
||||
} `json:"publicKey"`
|
||||
} `json:"options"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.True(t, result.Success, response.Body.String())
|
||||
require.Equal(t, "required", result.Data.Options.PublicKey.UserVerification)
|
||||
return result.Data.FlowToken, result.Data.Options.PublicKey.Challenge
|
||||
}
|
||||
|
||||
func TestSecurityLoginCodeCompletesOnce(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
path string
|
||||
backup bool
|
||||
}{{"/api/user/login/verify", false}, {"/api/user/login/2fa", false}, {"/api/user/login/verify", true}} {
|
||||
t.Run(fmt.Sprintf("%s/backup=%t", test.path, test.backup), func(t *testing.T) {
|
||||
user, _ := setupSecurityEnrollmentTest(t)
|
||||
factor := &model.TwoFA{UserId: user.Id, Secret: "JBSWY3DPEHPK3PXP", IsEnabled: true}
|
||||
require.NoError(t, model.DB.Create(factor).Error)
|
||||
challenge, err := service.StartLoginVerification(user, "password")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, challenge)
|
||||
code, err := totp.GenerateCode(factor.Secret, time.Now())
|
||||
require.NoError(t, err)
|
||||
if test.backup {
|
||||
code = "ABCD-1234"
|
||||
hash, err := common.HashBackupCode(code)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, model.DB.Create(&model.TwoFABackupCode{UserId: user.Id, CodeHash: hash}).Error)
|
||||
}
|
||||
body, err := common.Marshal(map[string]string{"flow_token": challenge.FlowToken, "code": code})
|
||||
require.NoError(t, err)
|
||||
router := gin.New()
|
||||
router.POST("/api/user/login/verify", VerifyLogin)
|
||||
router.POST("/api/user/login/2fa", Verify2FALogin)
|
||||
for attempt := range 2 {
|
||||
response := httptest.NewRecorder()
|
||||
router.ServeHTTP(response, httptest.NewRequest("POST", test.path, strings.NewReader(string(body))))
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
assert.Equal(t, attempt == 0, result.Success, response.Body.String())
|
||||
if attempt == 0 {
|
||||
var bundle service.AuthBundle
|
||||
require.NoError(t, common.Unmarshal(result.Data, &bundle))
|
||||
assert.NotEmpty(t, bundle.AccessToken)
|
||||
assert.Equal(t, "password", bundle.Session.LoginMethod)
|
||||
assert.NotEmpty(t, response.Header().Values("Set-Cookie"))
|
||||
} else {
|
||||
assert.Empty(t, response.Header().Values("Set-Cookie"))
|
||||
}
|
||||
}
|
||||
count, err := model.CountActiveUserSessions(user.Id, time.Now().Unix())
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 2, count)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityLoginRejectsChangedOrExpiredAuthorization(t *testing.T) {
|
||||
for _, change := range []string{"expired", "disabled user", "auth version", "factor removed", "factor locked", "legacy flow", "wrong method", "other purpose"} {
|
||||
t.Run(change, func(t *testing.T) {
|
||||
user, _ := setupSecurityEnrollmentTest(t)
|
||||
factor := &model.TwoFA{UserId: user.Id, Secret: "JBSWY3DPEHPK3PXP", IsEnabled: true}
|
||||
require.NoError(t, model.DB.Create(factor).Error)
|
||||
challenge, err := service.StartLoginVerification(user, "password")
|
||||
require.NoError(t, err)
|
||||
method := "2fa"
|
||||
switch change {
|
||||
case "expired":
|
||||
require.NoError(t, model.DB.Model(&model.AuthFlow{}).Where("purpose = ?", model.AuthFlowPurposeLoginVerification).Update("expires_at", time.Now().Add(-time.Minute)).Error)
|
||||
case "disabled user":
|
||||
require.NoError(t, model.DB.Model(user).Update("status", common.UserStatusDisabled).Error)
|
||||
case "auth version":
|
||||
require.NoError(t, model.DB.Model(user).Update("auth_version", user.AuthVersion+1).Error)
|
||||
case "factor removed":
|
||||
require.NoError(t, model.DB.Delete(factor).Error)
|
||||
case "factor locked":
|
||||
require.NoError(t, model.DB.Model(factor).Update("locked_until", time.Now().Add(time.Minute)).Error)
|
||||
case "legacy flow":
|
||||
require.NoError(t, model.DB.Model(&model.AuthFlow{}).Where("purpose = ?", model.AuthFlowPurposeLoginVerification).Update("purpose", model.AuthFlowPurposeTwoFALogin).Error)
|
||||
case "wrong method":
|
||||
method = "password"
|
||||
case "other purpose":
|
||||
require.NoError(t, model.DB.Model(&model.AuthFlow{}).Where("purpose = ?", model.AuthFlowPurposeLoginVerification).Update("purpose", model.AuthFlowPurposeSecurityProof).Error)
|
||||
}
|
||||
code, err := totp.GenerateCode(factor.Secret, time.Now())
|
||||
require.NoError(t, err)
|
||||
body, err := common.Marshal(map[string]string{"flow_token": challenge.FlowToken, "code": code, "method": method})
|
||||
require.NoError(t, err)
|
||||
response := securityEnrollmentRequest("POST", "/api/user/login/verify", string(body), "", service.AuthIdentity{}, VerifyLogin)
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
assert.False(t, result.Success, change)
|
||||
assert.Empty(t, response.Header().Values("Set-Cookie"))
|
||||
count, err := model.CountActiveUserSessions(user.Id, time.Now().Unix())
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 1, count)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityLoginPasskeyDoesNotRequireAdditionalTwoFA(t *testing.T) {
|
||||
for _, direct := range []bool{false, true} {
|
||||
t.Run(fmt.Sprintf("direct=%t", direct), func(t *testing.T) {
|
||||
user, _ := setupSecurityEnrollmentTest(t)
|
||||
key := newSecurityLoginPasskey(t, user.Id)
|
||||
require.NoError(t, model.DB.Create(&model.TwoFA{UserId: user.Id, Secret: "JBSWY3DPEHPK3PXP", IsEnabled: true}).Error)
|
||||
for _, verified := range []bool{false, true} {
|
||||
var flowToken, challenge, parentToken string
|
||||
if direct {
|
||||
response := securityEnrollmentRequest("POST", "/api/user/passkey/login/begin", "", "", service.AuthIdentity{}, PasskeyLoginBegin)
|
||||
var result struct {
|
||||
Data struct {
|
||||
FlowToken string `json:"flow_token"`
|
||||
Options struct {
|
||||
PublicKey struct {
|
||||
Challenge string `json:"challenge"`
|
||||
UserVerification string `json:"userVerification"`
|
||||
} `json:"publicKey"`
|
||||
} `json:"options"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.Equal(t, "required", result.Data.Options.PublicKey.UserVerification)
|
||||
flowToken, challenge = result.Data.FlowToken, result.Data.Options.PublicKey.Challenge
|
||||
} else {
|
||||
pending, err := service.StartLoginVerification(user, "oauth:github")
|
||||
require.NoError(t, err)
|
||||
parentToken = pending.FlowToken
|
||||
flowToken, challenge = beginSecurityLoginPasskey(t, parentToken)
|
||||
}
|
||||
var assertion map[string]any
|
||||
require.NoError(t, common.Unmarshal(securityPasskeyResponse(t, key, challenge, false, 0, verified), &assertion))
|
||||
if direct {
|
||||
response, ok := assertion["response"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
response["userHandle"] = base64.RawURLEncoding.EncodeToString([]byte(fmt.Sprint(user.Id)))
|
||||
}
|
||||
payload := map[string]any{"flow_token": flowToken, "credential": assertion}
|
||||
handler, path := PasskeyLoginFinish, "/api/user/passkey/login/finish"
|
||||
if !direct {
|
||||
payload["flow_token"], payload["passkey_flow_token"] = parentToken, flowToken
|
||||
handler, path = LoginPasskeyFinish, "/api/user/login/passkey/finish"
|
||||
}
|
||||
body, err := common.Marshal(payload)
|
||||
require.NoError(t, err)
|
||||
router := gin.New()
|
||||
router.POST(path, handler)
|
||||
response := httptest.NewRecorder()
|
||||
router.ServeHTTP(response, httptest.NewRequest("POST", path, strings.NewReader(string(body))))
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.Equal(t, verified, result.Success, response.Body.String())
|
||||
if verified {
|
||||
var bundle service.AuthBundle
|
||||
require.NoError(t, common.Unmarshal(result.Data, &bundle))
|
||||
assert.NotEmpty(t, bundle.AccessToken)
|
||||
} else {
|
||||
assert.Empty(t, response.Header().Values("Set-Cookie"))
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityLoginPasskeyConcurrentCompletionCreatesOneSession(t *testing.T) {
|
||||
user, _ := setupSecurityEnrollmentTest(t)
|
||||
key := newSecurityLoginPasskey(t, user.Id)
|
||||
pending, err := service.StartLoginVerification(user, "password")
|
||||
require.NoError(t, err)
|
||||
requests := make([]string, 2)
|
||||
for index := range requests {
|
||||
token, challenge := beginSecurityLoginPasskey(t, pending.FlowToken)
|
||||
body, err := common.Marshal(map[string]any{
|
||||
"flow_token": pending.FlowToken, "passkey_flow_token": token,
|
||||
"credential": securityPasskeyResponse(t, key, challenge, false, 0),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
requests[index] = string(body)
|
||||
}
|
||||
start := make(chan struct{})
|
||||
responses := make(chan *httptest.ResponseRecorder, 2)
|
||||
var workers sync.WaitGroup
|
||||
for _, body := range requests {
|
||||
workers.Go(func() {
|
||||
<-start
|
||||
responses <- securityEnrollmentRequest("POST", "/api/user/login/passkey/finish", body, "", service.AuthIdentity{}, LoginPasskeyFinish)
|
||||
})
|
||||
}
|
||||
close(start)
|
||||
workers.Wait()
|
||||
close(responses)
|
||||
successes := 0
|
||||
for response := range responses {
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
if result.Success {
|
||||
successes++
|
||||
}
|
||||
}
|
||||
assert.Equal(t, 1, successes)
|
||||
count, err := model.CountActiveUserSessions(user.Id, time.Now().Unix())
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 2, count)
|
||||
}
|
||||
|
||||
func TestSecurityLoginSessionFailureRollsBackChallengeConsumption(t *testing.T) {
|
||||
user, _ := setupSecurityEnrollmentTest(t)
|
||||
key := newSecurityLoginPasskey(t, user.Id)
|
||||
pending, err := service.StartLoginVerification(user, "password")
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, model.DB.Callback().Create().Before("gorm:create").Register("login_session_failure", func(tx *gorm.DB) {
|
||||
if tx.Statement.Table == "user_sessions" {
|
||||
tx.AddError(errors.New("private session creation failure"))
|
||||
}
|
||||
}))
|
||||
t.Cleanup(func() { _ = model.DB.Callback().Create().Remove("login_session_failure") })
|
||||
for _, fail := range []bool{true, false} {
|
||||
token, challenge := beginSecurityLoginPasskey(t, pending.FlowToken)
|
||||
body, err := common.Marshal(map[string]any{"flow_token": pending.FlowToken, "passkey_flow_token": token, "credential": securityPasskeyResponse(t, key, challenge, false, 0)})
|
||||
require.NoError(t, err)
|
||||
response := securityEnrollmentRequest("POST", "/api/user/login/passkey/finish", string(body), "", service.AuthIdentity{}, LoginPasskeyFinish)
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
assert.Equal(t, !fail, result.Success, response.Body.String())
|
||||
if fail {
|
||||
assert.Empty(t, response.Header().Values("Set-Cookie"))
|
||||
assert.NotContains(t, response.Body.String(), "private")
|
||||
_, err = model.GetAuthFlow(pending.FlowToken, model.AuthFlowMatch{Purpose: model.AuthFlowPurposeLoginVerification})
|
||||
require.NoError(t, err, "a failed session transaction must leave the parent challenge usable")
|
||||
require.NoError(t, model.DB.Callback().Create().Remove("login_session_failure"))
|
||||
}
|
||||
}
|
||||
count, err := model.CountActiveUserSessions(user.Id, time.Now().Unix())
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 2, count)
|
||||
}
|
||||
|
||||
func TestSecurityLoginFactorStateDoesNotAddPasswordLoginQueries(t *testing.T) {
|
||||
user, _ := setupSecurityEnrollmentTest(t)
|
||||
newSecurityLoginPasskey(t, user.Id)
|
||||
previousPasswordLogin := common.PasswordLoginEnabled
|
||||
common.PasswordLoginEnabled = true
|
||||
t.Cleanup(func() { common.PasswordLoginEnabled = previousPasswordLogin })
|
||||
queries := 0
|
||||
require.NoError(t, model.DB.Callback().Query().After("gorm:query").Register("login_query_count", func(tx *gorm.DB) {
|
||||
if !tx.DryRun {
|
||||
queries++
|
||||
}
|
||||
}))
|
||||
t.Cleanup(func() { _ = model.DB.Callback().Query().Remove("login_query_count") })
|
||||
state, err := model.GetUserVerificationState(user.Id)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, state.HasPassword)
|
||||
assert.True(t, state.HasPasskey)
|
||||
assert.False(t, state.HasTwoFA)
|
||||
assert.Equal(t, 1, queries, "factor availability must be one database round trip")
|
||||
queries = 0
|
||||
response := securityEnrollmentRequest("POST", "/api/user/login", `{"username":"enrollment-user","password":"enrollment-password"}`, "", service.AuthIdentity{}, Login)
|
||||
assert.Contains(t, response.Body.String(), `"require_verification":true`)
|
||||
assert.Equal(t, 2, queries, "only the existing credential lookup and the replacement factor-state lookup run before the challenge")
|
||||
}
|
||||
|
||||
type boundLoginOAuthProvider struct {
|
||||
authFlowTestOAuthProvider
|
||||
userID int
|
||||
}
|
||||
|
||||
func (*boundLoginOAuthProvider) IsUserIDTaken(string) bool { return true }
|
||||
func (*boundLoginOAuthProvider) ProviderUserIDColumn() string { return "github_id" }
|
||||
func (provider *boundLoginOAuthProvider) FillUserByProviderID(user *model.User, _ string) error {
|
||||
return model.DB.First(user, provider.userID).Error
|
||||
}
|
||||
|
||||
func TestSecurityLoginAllPrimaryTransportsRequireAdditionalVerification(t *testing.T) {
|
||||
for _, transport := range []string{"oauth", "custom oauth", "wechat", "telegram"} {
|
||||
t.Run(transport, func(t *testing.T) {
|
||||
var user *model.User
|
||||
var telegram *telegramOAuthFixture
|
||||
if transport == "telegram" {
|
||||
telegram = setupTelegramOAuthTest(t)
|
||||
user = telegram.user
|
||||
} else {
|
||||
user, _ = setupSecurityEnrollmentTest(t)
|
||||
}
|
||||
newSecurityLoginPasskey(t, user.Id)
|
||||
var response *httptest.ResponseRecorder
|
||||
switch transport {
|
||||
case "telegram":
|
||||
require.NoError(t, model.DB.Model(user).Update("telegram_id", "42").Error)
|
||||
state, code := telegram.authorization(t, "login", service.AuthIdentity{}, "", telegramIdentityClaims(42))
|
||||
response = telegramOAuthCallback(state, code, service.AuthIdentity{})
|
||||
case "wechat":
|
||||
require.NoError(t, model.DB.Model(user).Update("wechat_id", "bound-wechat").Error)
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
assert.Equal(t, "wechat-code", r.URL.Query().Get("code"))
|
||||
_, _ = w.Write([]byte(`{"success":true,"data":"bound-wechat"}`))
|
||||
}))
|
||||
t.Cleanup(upstream.Close)
|
||||
previousEnabled, previousAddress := common.WeChatAuthEnabled, common.WeChatServerAddress
|
||||
common.WeChatAuthEnabled, common.WeChatServerAddress = true, upstream.URL
|
||||
t.Cleanup(func() { common.WeChatAuthEnabled, common.WeChatServerAddress = previousEnabled, previousAddress })
|
||||
response = securityEnrollmentRequest("GET", "/api/oauth/wechat?code=wechat-code", "", "", service.AuthIdentity{}, WeChatAuth)
|
||||
default:
|
||||
const slug = "unified-login-test"
|
||||
if transport == "custom oauth" {
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if r.URL.Path == "/token" {
|
||||
_, _ = w.Write([]byte(`{"access_token":"provider-token","token_type":"Bearer"}`))
|
||||
return
|
||||
}
|
||||
_, _ = w.Write([]byte(`{"sub":"bound-custom","name":"User"}`))
|
||||
}))
|
||||
t.Cleanup(upstream.Close)
|
||||
oauth.RegisterCustom(slug, oauth.NewGenericOAuthProvider(&model.CustomOAuthProvider{
|
||||
Id: 42, Slug: slug, Name: "Custom login", Enabled: true, ClientId: "client", ClientSecret: "secret", UserIdField: "sub",
|
||||
TokenEndpoint: upstream.URL + "/token", UserInfoEndpoint: upstream.URL + "/userinfo",
|
||||
}))
|
||||
require.NoError(t, model.DB.Create(&model.UserOAuthBinding{UserId: user.Id, ProviderId: 42, ProviderUserId: "bound-custom"}).Error)
|
||||
} else {
|
||||
oauth.Register(slug, &boundLoginOAuthProvider{userID: user.Id})
|
||||
}
|
||||
t.Cleanup(func() { oauth.Unregister(slug) })
|
||||
token, _, err := model.CreateAuthFlow(model.AuthFlowCreate{Purpose: model.AuthFlowPurposeOAuth, Provider: slug, Intent: model.AuthFlowIntentLogin, Payload: `{}`, ExpiresAt: time.Now().Add(time.Minute)})
|
||||
require.NoError(t, err)
|
||||
router := gin.New()
|
||||
router.GET("/api/oauth/:provider", HandleOAuth)
|
||||
response = httptest.NewRecorder()
|
||||
router.ServeHTTP(response, httptest.NewRequest("GET", "/api/oauth/"+slug+"?state="+token+"&code=provider-code", nil))
|
||||
}
|
||||
var result struct {
|
||||
Success bool `json:"success"`
|
||||
Data service.LoginChallenge `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.True(t, result.Success, response.Body.String())
|
||||
assert.True(t, result.Data.RequireVerification)
|
||||
assert.Equal(t, []service.VerificationMethodOption{{Method: "passkey", Available: true}}, result.Data.Methods)
|
||||
assert.NotEmpty(t, result.Data.FlowToken)
|
||||
assert.Empty(t, response.Header().Values("Set-Cookie"))
|
||||
count, err := model.CountActiveUserSessions(user.Id, time.Now().Unix())
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 1, count)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityLoginPasskeyCannotCompleteAnotherChallenge(t *testing.T) {
|
||||
for _, otherUser := range []bool{false, true} {
|
||||
t.Run(fmt.Sprintf("other-user=%t", otherUser), func(t *testing.T) {
|
||||
user, _ := setupSecurityEnrollmentTest(t)
|
||||
key := newSecurityLoginPasskey(t, user.Id)
|
||||
first, err := service.StartLoginVerification(user, "password")
|
||||
require.NoError(t, err)
|
||||
passkeyToken, challenge := beginSecurityLoginPasskey(t, first.FlowToken)
|
||||
if otherUser {
|
||||
user = &model.User{Username: "other-login", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Group: "default", AffCode: "other-login", AuthVersion: 1}
|
||||
require.NoError(t, model.DB.Create(user).Error)
|
||||
newSecurityLoginPasskey(t, user.Id)
|
||||
}
|
||||
second, err := service.StartLoginVerification(user, "password")
|
||||
require.NoError(t, err)
|
||||
body, err := common.Marshal(map[string]any{"flow_token": second.FlowToken, "passkey_flow_token": passkeyToken, "credential": securityPasskeyResponse(t, key, challenge, false, 0)})
|
||||
require.NoError(t, err)
|
||||
response := securityEnrollmentRequest("POST", "/api/user/login/passkey/finish", string(body), "", service.AuthIdentity{}, LoginPasskeyFinish)
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
assert.False(t, result.Success)
|
||||
assert.Empty(t, response.Header().Values("Set-Cookie"))
|
||||
var count int64
|
||||
require.NoError(t, model.DB.Model(&model.UserSession{}).Count(&count).Error)
|
||||
assert.EqualValues(t, 1, count)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityLoginPasskeyCanManageTwoFAWithScopedProofs(t *testing.T) {
|
||||
for _, scope := range []string{service.VerificationScopeTwoFADisable, service.VerificationScopeTwoFABackupCodes} {
|
||||
t.Run(scope, func(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
key := newSecurityLoginPasskey(t, user.Id)
|
||||
require.NoError(t, model.DB.Create(&model.TwoFA{UserId: user.Id, Secret: "JBSWY3DPEHPK3PXP", IsEnabled: true}).Error)
|
||||
hash, err := common.HashBackupCode("ABCD-1234")
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, model.DB.Create(&model.TwoFABackupCode{UserId: user.Id, CodeHash: hash}).Error)
|
||||
if scope == service.VerificationScopeTwoFABackupCodes {
|
||||
_, err := service.VerifySecurityInput(identity, service.VerificationInput{Scope: scope, Method: "2fa", Code: "ABCD-1234"})
|
||||
assert.ErrorIs(t, err, service.ErrVerificationFailed)
|
||||
remaining, err := model.GetUnusedBackupCodeCount(user.Id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 1, remaining, "backup codes cannot authorize generating replacement backup codes")
|
||||
}
|
||||
beginBody, err := common.Marshal(map[string]string{"scope": scope})
|
||||
require.NoError(t, err)
|
||||
response := securityEnrollmentRequest("POST", "/api/user/passkey/verify/begin", string(beginBody), "", identity, PasskeyVerifyBegin)
|
||||
var started struct {
|
||||
Success bool `json:"success"`
|
||||
Data struct {
|
||||
FlowToken string `json:"flow_token"`
|
||||
Options struct {
|
||||
PublicKey struct {
|
||||
Challenge string `json:"challenge"`
|
||||
UserVerification string `json:"userVerification"`
|
||||
} `json:"publicKey"`
|
||||
} `json:"options"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &started))
|
||||
require.True(t, started.Success, response.Body.String())
|
||||
assert.Equal(t, "required", started.Data.Options.PublicKey.UserVerification)
|
||||
finish, err := common.Marshal(passkeyFinishRequest{FlowToken: started.Data.FlowToken, Credential: securityPasskeyResponse(t, key, started.Data.Options.PublicKey.Challenge, false, 0)})
|
||||
require.NoError(t, err)
|
||||
response = securityEnrollmentRequest("POST", "/api/user/passkey/verify/finish", string(finish), "", identity, PasskeyVerifyFinish)
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.True(t, result.Success, response.Body.String())
|
||||
var proof service.SecurityProof
|
||||
require.NoError(t, common.Unmarshal(result.Data, &proof))
|
||||
handler, path := Disable2FA, "/api/user/2fa/disable"
|
||||
if scope == service.VerificationScopeTwoFABackupCodes {
|
||||
handler, path = RegenerateBackupCodes, "/api/user/2fa/backup_codes"
|
||||
response = securityEnrollmentRequest("POST", "/api/user/2fa/disable", `{}`, proof.ProofToken, identity, Disable2FA)
|
||||
assert.Contains(t, response.Body.String(), "SECURITY_PROOF_SCOPE_MISMATCH")
|
||||
}
|
||||
response = securityEnrollmentRequest("POST", path, `{}`, proof.ProofToken, identity, handler)
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.True(t, result.Success, response.Body.String())
|
||||
stored, err := model.GetUserById(user.Id, false)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, identity.UserAuthVersion+1, stored.AuthVersion)
|
||||
_, err = model.GetPasskeyByUserID(user.Id)
|
||||
require.NoError(t, err, "managing 2FA must retain the alternative Passkey factor")
|
||||
factor, err := model.GetTwoFAByUserId(user.Id)
|
||||
require.NoError(t, err)
|
||||
if scope == service.VerificationScopeTwoFADisable {
|
||||
assert.Nil(t, factor)
|
||||
} else {
|
||||
require.NotNil(t, factor)
|
||||
assert.True(t, factor.IsEnabled)
|
||||
assert.Contains(t, string(result.Data), "backup_codes")
|
||||
}
|
||||
response = securityEnrollmentRequest("POST", path, `{}`, proof.ProofToken, identity, handler)
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
assert.False(t, result.Success, "the proof cannot authorize a second mutation")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityLoginRegisteredPasskeyRequiresUserVerification(t *testing.T) {
|
||||
for _, verified := range []bool{false, true} {
|
||||
t.Run(fmt.Sprintf("verified=%t", verified), func(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
|
||||
require.NoError(t, err)
|
||||
proof := issueSecurityEnrollmentProof(t, identity, service.VerificationOperation{Scope: service.VerificationScopePasskeyRegister}, service.VerificationMethodPassword)
|
||||
response := securityEnrollmentRequest("POST", "/api/user/passkey/register/begin", "", proof, identity, PasskeyRegisterBegin)
|
||||
var result struct {
|
||||
Success bool `json:"success"`
|
||||
Data struct {
|
||||
FlowToken string `json:"flow_token"`
|
||||
Options struct {
|
||||
PublicKey struct {
|
||||
Challenge string `json:"challenge"`
|
||||
AuthenticatorSelection struct {
|
||||
UserVerification string `json:"userVerification"`
|
||||
} `json:"authenticatorSelection"`
|
||||
} `json:"publicKey"`
|
||||
} `json:"options"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.True(t, result.Success, response.Body.String())
|
||||
assert.Equal(t, "required", result.Data.Options.PublicKey.AuthenticatorSelection.UserVerification)
|
||||
body, err := common.Marshal(passkeyFinishRequest{FlowToken: result.Data.FlowToken, Credential: securityPasskeyResponse(t, key, result.Data.Options.PublicKey.Challenge, true, 0, verified)})
|
||||
require.NoError(t, err)
|
||||
response = securityEnrollmentRequest("POST", "/api/user/passkey/register/finish", string(body), "", identity, PasskeyRegisterFinish)
|
||||
var finished securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &finished))
|
||||
assert.Equal(t, verified, finished.Success, response.Body.String())
|
||||
_, err = model.GetPasskeyByUserID(user.Id)
|
||||
if verified {
|
||||
require.NoError(t, err)
|
||||
} else {
|
||||
assert.ErrorIs(t, err, model.ErrPasskeyNotFound)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityLoginRequiresConfiguredFactors(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
name string
|
||||
twoFA, passkey bool
|
||||
locked, disabled bool
|
||||
methods []service.VerificationMethodOption
|
||||
unavailable bool
|
||||
}{
|
||||
{name: "password without additional factors"},
|
||||
{name: "passkey requires verification", passkey: true, methods: []service.VerificationMethodOption{{Method: "passkey", Available: true}}},
|
||||
{name: "twofa requires verification", twoFA: true, methods: []service.VerificationMethodOption{{Method: "2fa", Available: true}}},
|
||||
{name: "both factors are alternatives", twoFA: true, passkey: true, methods: []service.VerificationMethodOption{{Method: "2fa", Available: true}, {Method: "passkey", Available: true}}},
|
||||
{name: "locked twofa permits passkey", twoFA: true, passkey: true, locked: true, methods: []service.VerificationMethodOption{{Method: "2fa", Available: false, Reason: service.ErrVerificationLocked.Error()}, {Method: "passkey", Available: true}}},
|
||||
{name: "disabled passkey permits twofa", twoFA: true, passkey: true, disabled: true, methods: []service.VerificationMethodOption{{Method: "2fa", Available: true}, {Method: "passkey", Available: false, Reason: "Passkey authentication is disabled."}}},
|
||||
{name: "only passkey disabled blocks password", passkey: true, disabled: true, unavailable: true},
|
||||
{name: "both factors unavailable block password", passkey: true, twoFA: true, disabled: true, locked: true, unavailable: true},
|
||||
} {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
user, _ := setupSecurityEnrollmentTest(t)
|
||||
previousPasswordLogin := common.PasswordLoginEnabled
|
||||
common.PasswordLoginEnabled = true
|
||||
t.Cleanup(func() { common.PasswordLoginEnabled = previousPasswordLogin })
|
||||
if test.passkey {
|
||||
require.NoError(t, model.DB.Create(&model.PasskeyCredential{UserID: user.Id, CredentialID: "login-key", PublicKey: "public-key"}).Error)
|
||||
}
|
||||
if test.twoFA {
|
||||
factor := &model.TwoFA{UserId: user.Id, Secret: "JBSWY3DPEHPK3PXP", IsEnabled: true}
|
||||
if test.locked {
|
||||
until := time.Now().Add(time.Minute)
|
||||
factor.LockedUntil = &until
|
||||
}
|
||||
require.NoError(t, model.DB.Create(factor).Error)
|
||||
}
|
||||
system_setting.GetPasskeySettings().Enabled = !test.disabled
|
||||
before, err := model.CountActiveUserSessions(user.Id, time.Now().Unix())
|
||||
require.NoError(t, err)
|
||||
response := securityEnrollmentRequest(http.MethodPost, "/api/user/login", `{"username":"enrollment-user","password":"enrollment-password"}`, "", service.AuthIdentity{}, Login)
|
||||
var result struct {
|
||||
Success bool `json:"success"`
|
||||
Data struct {
|
||||
RequireVerification bool `json:"require_verification"`
|
||||
FlowToken string `json:"flow_token"`
|
||||
ExpiresAt int64 `json:"expires_at"`
|
||||
AccessToken string `json:"access_token"`
|
||||
Methods []service.VerificationMethodOption `json:"methods"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
after, err := model.CountActiveUserSessions(user.Id, time.Now().Unix())
|
||||
require.NoError(t, err)
|
||||
if test.unavailable {
|
||||
assert.False(t, result.Success)
|
||||
assert.Equal(t, before, after)
|
||||
assert.Empty(t, response.Header().Values("Set-Cookie"))
|
||||
return
|
||||
}
|
||||
require.True(t, result.Success, response.Body.String())
|
||||
if len(test.methods) == 0 {
|
||||
assert.False(t, result.Data.RequireVerification)
|
||||
assert.NotEmpty(t, result.Data.AccessToken)
|
||||
assert.Equal(t, before+1, after)
|
||||
return
|
||||
}
|
||||
assert.True(t, result.Data.RequireVerification)
|
||||
assert.NotEmpty(t, result.Data.FlowToken)
|
||||
assert.Greater(t, result.Data.ExpiresAt, time.Now().Unix())
|
||||
assert.LessOrEqual(t, result.Data.ExpiresAt, time.Now().Add(5*time.Minute).Unix())
|
||||
assert.Equal(t, test.methods, result.Data.Methods)
|
||||
assert.Empty(t, result.Data.AccessToken)
|
||||
assert.Empty(t, response.Header().Values("Set-Cookie"))
|
||||
assert.Equal(t, before, after, "a pending challenge must not create a session")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
type authFlowTestOAuthProvider struct {
|
||||
exchangeErr error
|
||||
userInfoErr error
|
||||
exchangeCalls int
|
||||
userInfoCalls int
|
||||
}
|
||||
|
||||
func (*authFlowTestOAuthProvider) GetName() string { return "Auth Flow Test" }
|
||||
func (*authFlowTestOAuthProvider) IsEnabled() bool { return true }
|
||||
func (provider *authFlowTestOAuthProvider) ExchangeToken(context.Context, string, *gin.Context) (*oauth.OAuthToken, error) {
|
||||
provider.exchangeCalls++
|
||||
if provider.exchangeErr != nil {
|
||||
return nil, provider.exchangeErr
|
||||
}
|
||||
return &oauth.OAuthToken{}, nil
|
||||
}
|
||||
func (provider *authFlowTestOAuthProvider) GetUserInfo(context.Context, *oauth.OAuthToken) (*oauth.OAuthUser, error) {
|
||||
provider.userInfoCalls++
|
||||
if provider.userInfoErr != nil {
|
||||
return nil, provider.userInfoErr
|
||||
}
|
||||
return &oauth.OAuthUser{ProviderUserID: "external-user"}, nil
|
||||
}
|
||||
func (*authFlowTestOAuthProvider) IsUserIDTaken(string) bool { return false }
|
||||
func (*authFlowTestOAuthProvider) FillUserByProviderID(*model.User, string) error { return nil }
|
||||
func (*authFlowTestOAuthProvider) SetProviderUserID(*model.User, string) {}
|
||||
func (*authFlowTestOAuthProvider) GetProviderPrefix() string { return "flow_" }
|
||||
func (*authFlowTestOAuthProvider) ProviderUserIDColumn() string { return "" }
|
||||
|
||||
func setupAuthFlowControllerTest(t *testing.T) *authFlowTestOAuthProvider {
|
||||
t.Helper()
|
||||
previousDB, previousLogDB := model.DB, model.LOG_DB
|
||||
previousRedis := common.RedisEnabled
|
||||
common.RedisEnabled = false
|
||||
previousType := common.MainDatabaseType()
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.AutoMigrate(&model.AuthFlow{}, &model.User{}, &model.UserSession{}, &model.AuditLog{}))
|
||||
model.DB, model.LOG_DB = db, db
|
||||
common.SetMainDatabaseType(common.DatabaseTypeSQLite)
|
||||
provider := &authFlowTestOAuthProvider{}
|
||||
oauth.Register("auth-flow-test", provider)
|
||||
t.Cleanup(func() {
|
||||
oauth.Unregister("auth-flow-test")
|
||||
model.DB, model.LOG_DB = previousDB, previousLogDB
|
||||
common.RedisEnabled = previousRedis
|
||||
common.SetMainDatabaseType(previousType)
|
||||
})
|
||||
return provider
|
||||
}
|
||||
|
||||
func TestGenerateOAuthCodeCarriesAffiliateInLoginFlow(t *testing.T) {
|
||||
setupAuthFlowControllerTest(t)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/api/oauth/state", strings.NewReader(`{"provider":"auth-flow-test","intent":"login","aff":"invite-code"}`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
GenerateOAuthCode(c)
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
var response struct {
|
||||
Success bool `json:"success"`
|
||||
Data struct {
|
||||
FlowToken string `json:"flow_token"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
|
||||
require.True(t, response.Success)
|
||||
flow, err := model.GetAuthFlow(response.Data.FlowToken, model.AuthFlowMatch{
|
||||
Purpose: model.AuthFlowPurposeOAuth, Provider: "auth-flow-test", Intent: model.AuthFlowIntentLogin,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
var payload oauthFlowPayload
|
||||
require.NoError(t, common.UnmarshalJsonStr(flow.Payload, &payload))
|
||||
assert.Equal(t, "invite-code", payload.AffiliateCode)
|
||||
assert.Zero(t, flow.UserId)
|
||||
assert.Empty(t, flow.SessionId)
|
||||
}
|
||||
|
||||
func TestGenerateOAuthCodeBindsFlowToAuthenticatedSession(t *testing.T) {
|
||||
_, identity := setupSecurityEnrollmentTest(t)
|
||||
oauth.Register("auth-flow-test", &authFlowTestOAuthProvider{})
|
||||
t.Cleanup(func() { oauth.Unregister("auth-flow-test") })
|
||||
proof := issueSecurityEnrollmentProof(t, identity, service.VerificationOperation{Scope: service.VerificationScopeAccountBind, Context: []byte(`{"provider":"auth-flow-test"}`)}, service.VerificationMethodPassword)
|
||||
response := securityEnrollmentRequest(http.MethodPost, "/api/oauth/state", `{"provider":"auth-flow-test","intent":"bind"}`, proof, identity, GenerateOAuthCode)
|
||||
var result struct {
|
||||
Success bool `json:"success"`
|
||||
Data struct {
|
||||
FlowToken string `json:"flow_token"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.True(t, result.Success, response.Body.String())
|
||||
flow, err := model.GetAuthFlow(result.Data.FlowToken, model.AuthFlowMatch{Purpose: model.AuthFlowPurposeOAuth, Provider: "auth-flow-test", Intent: model.AuthFlowIntentBind, UserId: identity.UserID, SessionId: identity.SessionID})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, identity.UserID, flow.UserId)
|
||||
assert.Equal(t, identity.SessionID, flow.SessionId)
|
||||
}
|
||||
|
||||
func TestOAuthLoginConsumesFlowOnlyAfterProviderIdentity(t *testing.T) {
|
||||
provider := setupAuthFlowControllerTest(t)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
exchangeErr error
|
||||
userInfoErr error
|
||||
}{
|
||||
{name: "exchange failure", exchangeErr: errors.New("exchange failed")},
|
||||
{name: "user info failure", userInfoErr: errors.New("user info failed")},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
provider.exchangeErr = test.exchangeErr
|
||||
provider.userInfoErr = test.userInfoErr
|
||||
token, _, err := model.CreateAuthFlow(model.AuthFlowCreate{
|
||||
Purpose: model.AuthFlowPurposeOAuth, Provider: "auth-flow-test", Intent: model.AuthFlowIntentLogin,
|
||||
Payload: `{}`, ExpiresAt: time.Now().Add(time.Minute),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
router := gin.New()
|
||||
router.GET("/api/oauth/:provider", HandleOAuth)
|
||||
request := httptest.NewRequest(http.MethodGet, "/api/oauth/auth-flow-test?state="+token+"&code=test", nil)
|
||||
response := httptest.NewRecorder()
|
||||
router.ServeHTTP(response, request)
|
||||
|
||||
flow, err := model.GetAuthFlow(token, model.AuthFlowMatch{
|
||||
Purpose: model.AuthFlowPurposeOAuth, Provider: "auth-flow-test", Intent: model.AuthFlowIntentLogin,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, flow.ConsumedAt)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOAuthLoginConsumesFlowAfterProviderIdentityAndOnProviderError(t *testing.T) {
|
||||
provider := setupAuthFlowControllerTest(t)
|
||||
|
||||
provider.exchangeErr = nil
|
||||
provider.userInfoErr = nil
|
||||
successToken, _, err := model.CreateAuthFlow(model.AuthFlowCreate{
|
||||
Purpose: model.AuthFlowPurposeOAuth, Provider: "auth-flow-test", Intent: model.AuthFlowIntentLogin,
|
||||
Payload: `{invalid`, ExpiresAt: time.Now().Add(time.Minute),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
router := gin.New()
|
||||
router.GET("/api/oauth/:provider", HandleOAuth)
|
||||
request := httptest.NewRequest(http.MethodGet, "/api/oauth/auth-flow-test?state="+successToken+"&code=test", nil)
|
||||
response := httptest.NewRecorder()
|
||||
router.ServeHTTP(response, request)
|
||||
_, err = model.GetAuthFlow(successToken, model.AuthFlowMatch{Purpose: model.AuthFlowPurposeOAuth})
|
||||
assert.ErrorIs(t, err, model.ErrAuthFlowConsumed)
|
||||
assert.Equal(t, 1, provider.exchangeCalls)
|
||||
assert.Equal(t, 1, provider.userInfoCalls)
|
||||
|
||||
providerErrorToken, _, err := model.CreateAuthFlow(model.AuthFlowCreate{
|
||||
Purpose: model.AuthFlowPurposeOAuth, Provider: "auth-flow-test", Intent: model.AuthFlowIntentLogin,
|
||||
Payload: `{}`, ExpiresAt: time.Now().Add(time.Minute),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
request = httptest.NewRequest(http.MethodGet, "/api/oauth/auth-flow-test?state="+providerErrorToken+"&error=access_denied", nil)
|
||||
response = httptest.NewRecorder()
|
||||
router.ServeHTTP(response, request)
|
||||
_, err = model.GetAuthFlow(providerErrorToken, model.AuthFlowMatch{Purpose: model.AuthFlowPurposeOAuth})
|
||||
assert.ErrorIs(t, err, model.ErrAuthFlowConsumed)
|
||||
assert.Equal(t, 1, provider.exchangeCalls)
|
||||
assert.Equal(t, 1, provider.userInfoCalls)
|
||||
}
|
||||
|
||||
func TestOAuthBindProviderErrorConsumesSessionBoundFlow(t *testing.T) {
|
||||
_, identity := setupSecurityEnrollmentTest(t)
|
||||
provider := &authFlowTestOAuthProvider{}
|
||||
oauth.Register("auth-flow-test", provider)
|
||||
t.Cleanup(func() { oauth.Unregister("auth-flow-test") })
|
||||
proof := issueSecurityEnrollmentProof(t, identity, service.VerificationOperation{Scope: service.VerificationScopeAccountBind, Context: []byte(`{"provider":"auth-flow-test"}`)}, service.VerificationMethodPassword)
|
||||
started := securityEnrollmentRequest(http.MethodPost, "/api/oauth/state", `{"provider":"auth-flow-test","intent":"bind"}`, proof, identity, GenerateOAuthCode)
|
||||
var result struct {
|
||||
Data struct {
|
||||
FlowToken string `json:"flow_token"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(started.Body.Bytes(), &result))
|
||||
require.NotEmpty(t, result.Data.FlowToken)
|
||||
response := securityEnrollmentRequest(http.MethodGet, "/api/oauth/auth-flow-test?state="+result.Data.FlowToken+"&error=access_denied&error_description=cancelled", "", "", identity, func(c *gin.Context) {
|
||||
c.Params = gin.Params{{Key: "provider", Value: "auth-flow-test"}}
|
||||
HandleOAuth(c)
|
||||
})
|
||||
assert.Equal(t, http.StatusOK, response.Code)
|
||||
_, err := model.GetAuthFlow(result.Data.FlowToken, model.AuthFlowMatch{Purpose: model.AuthFlowPurposeOAuth})
|
||||
assert.ErrorIs(t, err, model.ErrAuthFlowConsumed)
|
||||
assert.Zero(t, provider.exchangeCalls)
|
||||
assert.Zero(t, provider.userInfoCalls)
|
||||
}
|
||||
@@ -0,0 +1,197 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func RefreshAuth(c *gin.Context) {
|
||||
setAuthNoStore(c)
|
||||
rawRefreshToken, err := c.Cookie(service.RefreshCookieName)
|
||||
if err != nil || rawRefreshToken == "" {
|
||||
service.ClearRefreshCookie(c)
|
||||
writeAuthSessionError(c, service.ErrRefreshTokenInvalid)
|
||||
return
|
||||
}
|
||||
bundle, user, err := service.RefreshLoginSession(rawRefreshToken, c.GetHeader("X-Auth-Session"), c.ClientIP(), c.Request.UserAgent())
|
||||
if err != nil {
|
||||
if errors.Is(err, service.ErrRefreshTokenInvalid) || errors.Is(err, service.ErrLoginSessionRevoked) {
|
||||
service.ClearRefreshCookie(c)
|
||||
}
|
||||
writeAuthSessionError(c, err)
|
||||
return
|
||||
}
|
||||
service.WriteRefreshCookie(c, bundle.RefreshToken)
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": gin.H{
|
||||
"access_token": bundle.AccessToken,
|
||||
"token_type": bundle.TokenType,
|
||||
"access_expires_at": bundle.AccessExpiresAt,
|
||||
"user": buildSelfUserData(user),
|
||||
"session": bundle.Session,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func AuthLogout(c *gin.Context) {
|
||||
setAuthNoStore(c)
|
||||
expectedSID := strings.TrimSpace(c.GetHeader("X-Auth-Session"))
|
||||
rawRefreshToken, cookieErr := c.Cookie(service.RefreshCookieName)
|
||||
cookieSID, hasCookieSID := service.RefreshTokenSID(rawRefreshToken)
|
||||
if expectedSID != "" && cookieErr == nil && hasCookieSID && cookieSID != expectedSID {
|
||||
writeAuthSessionError(c, service.ErrLoginSessionMismatch)
|
||||
return
|
||||
}
|
||||
|
||||
if rawAccessToken, ok := dashboardBearer(c.GetHeader("Authorization")); ok {
|
||||
if identity, err := service.ParseAccessToken(rawAccessToken); err == nil {
|
||||
if expectedSID != "" && expectedSID != identity.SessionID {
|
||||
writeAuthSessionError(c, service.ErrLoginSessionMismatch)
|
||||
return
|
||||
}
|
||||
if _, err := model.RevokeUserSession(identity.UserID, identity.SessionID, "logout"); err != nil {
|
||||
writeAuthSessionError(c, err)
|
||||
return
|
||||
}
|
||||
cookieCleared := false
|
||||
if cookieErr == nil && hasCookieSID && cookieSID == identity.SessionID {
|
||||
if err := service.RevokeByRefreshToken(rawRefreshToken, identity.SessionID, "logout"); err != nil {
|
||||
writeAuthSessionError(c, err)
|
||||
return
|
||||
}
|
||||
service.ClearRefreshCookie(c)
|
||||
cookieCleared = true
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": gin.H{"revoked_sid": identity.SessionID, "cookie_cleared": cookieCleared},
|
||||
})
|
||||
return
|
||||
}
|
||||
}
|
||||
if cookieErr != nil || rawRefreshToken == "" {
|
||||
service.ClearRefreshCookie(c)
|
||||
c.JSON(http.StatusOK, gin.H{"success": true, "message": ""})
|
||||
return
|
||||
}
|
||||
if err := service.RevokeByRefreshToken(rawRefreshToken, expectedSID, "logout"); err != nil {
|
||||
writeAuthSessionError(c, err)
|
||||
return
|
||||
}
|
||||
service.ClearRefreshCookie(c)
|
||||
c.JSON(http.StatusOK, gin.H{"success": true, "message": ""})
|
||||
}
|
||||
|
||||
func GetLoginSessions(c *gin.Context) {
|
||||
identity, ok := requireBrowserSession(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
sessions, err := service.ListLoginSessions(identity.UserID, identity.SessionID)
|
||||
if err != nil {
|
||||
writeAuthSessionError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"success": true, "message": "", "data": sessions})
|
||||
}
|
||||
|
||||
func DeleteLoginSession(c *gin.Context) {
|
||||
identity, ok := requireBrowserSession(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
sid := strings.TrimSpace(c.Param("sid"))
|
||||
if sid == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"success": false, "code": "AUTH_SESSION_ID_REQUIRED", "message": "session id is required"})
|
||||
return
|
||||
}
|
||||
revoked, err := model.RevokeUserSession(identity.UserID, sid, "user_revoked")
|
||||
if err != nil {
|
||||
writeAuthSessionError(c, err)
|
||||
return
|
||||
}
|
||||
if !revoked {
|
||||
c.JSON(http.StatusNotFound, gin.H{"success": false, "code": "AUTH_SESSION_NOT_FOUND", "message": "session not found"})
|
||||
return
|
||||
}
|
||||
if rawRefreshToken, cookieErr := c.Cookie(service.RefreshCookieName); cookieErr == nil {
|
||||
cookieSID, ok := service.RefreshTokenSID(rawRefreshToken)
|
||||
if ok && cookieSID == sid {
|
||||
service.ClearRefreshCookie(c)
|
||||
}
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"success": true, "message": "", "data": gin.H{"revoked_sid": sid, "current": sid == identity.SessionID}})
|
||||
}
|
||||
|
||||
func RevokeOtherLoginSessions(c *gin.Context) {
|
||||
identity, ok := requireBrowserSession(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
count, err := model.RevokeOtherUserSessions(identity.UserID, identity.SessionID, "user_revoked_others")
|
||||
if err != nil {
|
||||
writeAuthSessionError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"success": true, "message": "", "data": gin.H{"revoked_count": count}})
|
||||
}
|
||||
|
||||
func requireBrowserSession(c *gin.Context) (service.AuthIdentity, bool) {
|
||||
identity, ok := middleware.GetSessionAuthIdentity(c)
|
||||
if !ok {
|
||||
c.JSON(http.StatusForbidden, gin.H{
|
||||
"success": false,
|
||||
"code": "AUTH_SESSION_REQUIRED",
|
||||
"message": "a dashboard login session is required",
|
||||
})
|
||||
return service.AuthIdentity{}, false
|
||||
}
|
||||
return identity, true
|
||||
}
|
||||
|
||||
func writeAuthSessionError(c *gin.Context, err error) {
|
||||
status, code := service.AuthSessionErrorCode(err)
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
status, code = http.StatusUnauthorized, "AUTH_UNAUTHORIZED"
|
||||
}
|
||||
if status == http.StatusInternalServerError {
|
||||
// The response body only carries the generic AUTH_INTERNAL_ERROR
|
||||
// code; without this log the underlying Redis/database/session
|
||||
// failure is indistinguishable from the client side.
|
||||
logger.LogError(c.Request.Context(), fmt.Sprintf("auth session internal error (%s %s): %v", c.Request.Method, c.Request.URL.Path, err))
|
||||
}
|
||||
c.JSON(status, gin.H{"success": false, "code": code, "message": http.StatusText(status)})
|
||||
}
|
||||
|
||||
func setAuthNoStore(c *gin.Context) {
|
||||
c.Header("Cache-Control", "no-store")
|
||||
}
|
||||
|
||||
func authRotationData(bundle *service.AuthBundle) gin.H {
|
||||
return gin.H{
|
||||
"access_token": bundle.AccessToken,
|
||||
"token_type": bundle.TokenType,
|
||||
"access_expires_at": bundle.AccessExpiresAt,
|
||||
"session": bundle.Session,
|
||||
}
|
||||
}
|
||||
|
||||
func dashboardBearer(header string) (string, bool) {
|
||||
parts := strings.Fields(header)
|
||||
if len(parts) != 2 || !strings.EqualFold(parts[0], "Bearer") || parts[1] == "" {
|
||||
return "", false
|
||||
}
|
||||
return parts[1], true
|
||||
}
|
||||
@@ -0,0 +1,155 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func TestAuthLogoutRejectsRefreshCookieSessionMismatch(t *testing.T) {
|
||||
previousDB := model.DB
|
||||
previousRedis := common.RedisEnabled
|
||||
previousSecret := common.SessionSecret
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.AutoMigrate(&model.User{}, &model.UserSession{}))
|
||||
model.DB = db
|
||||
common.RedisEnabled = false
|
||||
common.SessionSecret = "auth-logout-mismatch-test-secret"
|
||||
t.Cleanup(func() {
|
||||
model.DB = previousDB
|
||||
common.RedisEnabled = previousRedis
|
||||
common.SessionSecret = previousSecret
|
||||
})
|
||||
|
||||
user := &model.User{
|
||||
Username: "logout-mismatch-user", Password: "unused", Role: common.RoleCommonUser,
|
||||
Status: common.UserStatusEnabled, Group: "default", AuthVersion: 1,
|
||||
}
|
||||
require.NoError(t, db.Create(user).Error)
|
||||
sessionA, err := service.CreateLoginSession(user.Id, "password", "127.0.0.1", "agent-a")
|
||||
require.NoError(t, err)
|
||||
sessionB, err := service.CreateLoginSession(user.Id, "password", "127.0.0.1", "agent-b")
|
||||
require.NoError(t, err)
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/api/user/auth/logout", nil)
|
||||
c.Request.Header.Set("Authorization", "Bearer "+sessionA.AccessToken)
|
||||
c.Request.Header.Set("X-Auth-Session", sessionA.Session.SID)
|
||||
c.Request.AddCookie(&http.Cookie{Name: service.RefreshCookieName, Value: sessionB.RefreshToken})
|
||||
|
||||
AuthLogout(c)
|
||||
|
||||
assert.Equal(t, http.StatusConflict, recorder.Code)
|
||||
var response struct {
|
||||
Success bool `json:"success"`
|
||||
Code string `json:"code"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
|
||||
assert.False(t, response.Success)
|
||||
assert.Equal(t, "AUTH_SESSION_MISMATCH", response.Code)
|
||||
for _, sid := range []string{sessionA.Session.SID, sessionB.Session.SID} {
|
||||
stored, err := model.GetUserSessionBySID(sid)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, model.UserSessionStatusActive, stored.Status)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteAuthSessionErrorMapsSessionGrowthLimits(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
tests := []struct {
|
||||
name string
|
||||
err error
|
||||
expectedStatus int
|
||||
expectedCode string
|
||||
}{
|
||||
{
|
||||
name: "active session limit",
|
||||
err: model.ErrUserSessionLimit,
|
||||
expectedStatus: http.StatusConflict,
|
||||
expectedCode: "AUTH_SESSION_LIMIT",
|
||||
},
|
||||
{
|
||||
name: "issuance limit",
|
||||
err: model.ErrUserSessionIssuanceLimit,
|
||||
expectedStatus: http.StatusTooManyRequests,
|
||||
expectedCode: "AUTH_SESSION_ISSUANCE_LIMIT",
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
writeAuthSessionError(c, test.err)
|
||||
|
||||
assert.Equal(t, test.expectedStatus, recorder.Code)
|
||||
var response struct {
|
||||
Success bool `json:"success"`
|
||||
Code string `json:"code"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
|
||||
assert.False(t, response.Success)
|
||||
assert.Equal(t, test.expectedCode, response.Code)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionLimitDoesNotRecordRejectedLoginAsSuccessful(t *testing.T) {
|
||||
previousDB := model.DB
|
||||
previousRedis := common.RedisEnabled
|
||||
previousActiveLimit := common.UserSessionActiveLimit
|
||||
previousIssuanceLimit := common.UserSessionIssuanceLimit
|
||||
previousIssuanceWindow := common.UserSessionIssuanceWindowSeconds
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.AutoMigrate(&model.User{}, &model.UserSession{}, &model.TwoFA{}, &model.PasskeyCredential{}))
|
||||
model.DB = db
|
||||
common.RedisEnabled = false
|
||||
common.UserSessionActiveLimit = 1
|
||||
common.UserSessionIssuanceLimit = 100
|
||||
common.UserSessionIssuanceWindowSeconds = int64(common.DefaultUserSessionIssuanceWindowSeconds)
|
||||
t.Cleanup(func() {
|
||||
model.DB = previousDB
|
||||
common.RedisEnabled = previousRedis
|
||||
common.UserSessionActiveLimit = previousActiveLimit
|
||||
common.UserSessionIssuanceLimit = previousIssuanceLimit
|
||||
common.UserSessionIssuanceWindowSeconds = previousIssuanceWindow
|
||||
})
|
||||
|
||||
const previousLastLoginAt = int64(123)
|
||||
user := &model.User{
|
||||
Username: "rejected-login-audit-user", Password: "unused", Role: common.RoleCommonUser,
|
||||
Status: common.UserStatusEnabled, Group: "default", AuthVersion: 1, LastLoginAt: previousLastLoginAt,
|
||||
}
|
||||
require.NoError(t, db.Create(user).Error)
|
||||
now := time.Now().Unix()
|
||||
require.NoError(t, db.Create(&model.UserSession{
|
||||
SID: "existing-active-session", UserID: user.Id, Version: 1, UserAuthVersion: user.AuthVersion,
|
||||
Status: model.UserSessionStatusActive, RefreshHash: "hash", LoginMethod: "password",
|
||||
CreatedAt: now, LastActiveAt: now, ExpiresAt: now + 3600,
|
||||
}).Error)
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/api/user/login", nil)
|
||||
setupLogin(user, c)
|
||||
|
||||
assert.Equal(t, http.StatusConflict, recorder.Code)
|
||||
var stored model.User
|
||||
require.NoError(t, db.First(&stored, user.Id).Error)
|
||||
assert.Equal(t, previousLastLoginAt, stored.LastLoginAt)
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/QuantumNous/new-api/service/authz"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// GetPermissionCatalog returns the permission schema used by the client to
|
||||
// render the permission editor: the registry of resources with their actions
|
||||
// and display label keys, plus the roles with their baseline grant matrices.
|
||||
// Defining it in the authz package keeps the schema in a single place.
|
||||
func GetPermissionCatalog(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": gin.H{
|
||||
"resources": authz.Catalog(),
|
||||
"roles": authz.Roles(),
|
||||
},
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,108 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func GetSubscription(c *gin.Context) {
|
||||
var remainQuota int
|
||||
var usedQuota int
|
||||
var err error
|
||||
var token *model.Token
|
||||
var expiredTime int64
|
||||
if common.DisplayTokenStatEnabled {
|
||||
tokenId := c.GetInt("token_id")
|
||||
token, err = model.GetTokenById(tokenId)
|
||||
expiredTime = token.ExpiredTime
|
||||
remainQuota = token.RemainQuota
|
||||
usedQuota = token.UsedQuota
|
||||
} else {
|
||||
userId := c.GetInt("id")
|
||||
remainQuota, err = model.GetUserQuota(userId, false)
|
||||
usedQuota, err = model.GetUserUsedQuota(userId)
|
||||
}
|
||||
if expiredTime <= 0 {
|
||||
expiredTime = 0
|
||||
}
|
||||
if err != nil {
|
||||
openAIError := types.OpenAIError{
|
||||
Message: err.Error(),
|
||||
Type: "upstream_error",
|
||||
}
|
||||
c.JSON(200, gin.H{
|
||||
"error": openAIError,
|
||||
})
|
||||
return
|
||||
}
|
||||
quota := remainQuota + usedQuota
|
||||
amount := float64(quota)
|
||||
// OpenAI 兼容接口中的 *_USD 字段含义保持“额度单位”对应值：
|
||||
// 我们将其解释为以“站点展示类型”为准：
|
||||
// - USD: 直接除以 QuotaPerUnit
|
||||
// - CNY: 先转 USD 再乘汇率
|
||||
// - TOKENS: 直接使用 tokens 数量
|
||||
switch operation_setting.GetQuotaDisplayType() {
|
||||
case operation_setting.QuotaDisplayTypeCNY:
|
||||
amount = amount / common.QuotaPerUnit * operation_setting.USDExchangeRate
|
||||
case operation_setting.QuotaDisplayTypeTokens:
|
||||
// amount 保持 tokens 数值
|
||||
default:
|
||||
amount = amount / common.QuotaPerUnit
|
||||
}
|
||||
if token != nil && token.UnlimitedQuota {
|
||||
amount = 100000000
|
||||
}
|
||||
subscription := OpenAISubscriptionResponse{
|
||||
Object: "billing_subscription",
|
||||
HasPaymentMethod: true,
|
||||
SoftLimitUSD: amount,
|
||||
HardLimitUSD: amount,
|
||||
SystemHardLimitUSD: amount,
|
||||
AccessUntil: expiredTime,
|
||||
}
|
||||
c.JSON(200, subscription)
|
||||
return
|
||||
}
|
||||
|
||||
func GetUsage(c *gin.Context) {
|
||||
var quota int
|
||||
var err error
|
||||
var token *model.Token
|
||||
if common.DisplayTokenStatEnabled {
|
||||
tokenId := c.GetInt("token_id")
|
||||
token, err = model.GetTokenById(tokenId)
|
||||
quota = token.UsedQuota
|
||||
} else {
|
||||
userId := c.GetInt("id")
|
||||
quota, err = model.GetUserUsedQuota(userId)
|
||||
}
|
||||
if err != nil {
|
||||
openAIError := types.OpenAIError{
|
||||
Message: err.Error(),
|
||||
Type: "new_api_error",
|
||||
}
|
||||
c.JSON(200, gin.H{
|
||||
"error": openAIError,
|
||||
})
|
||||
return
|
||||
}
|
||||
amount := float64(quota)
|
||||
switch operation_setting.GetQuotaDisplayType() {
|
||||
case operation_setting.QuotaDisplayTypeCNY:
|
||||
amount = amount / common.QuotaPerUnit * operation_setting.USDExchangeRate
|
||||
case operation_setting.QuotaDisplayTypeTokens:
|
||||
// tokens 保持原值
|
||||
default:
|
||||
amount = amount / common.QuotaPerUnit
|
||||
}
|
||||
usage := OpenAIUsageResponse{
|
||||
Object: "list",
|
||||
TotalUsage: amount * 100,
|
||||
}
|
||||
c.JSON(200, usage)
|
||||
return
|
||||
}
|
||||
@@ -0,0 +1,203 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/setting/config"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func TestUpdateOptionRejectsInvalidTaskBillingExpressions(t *testing.T) {
|
||||
const pluginKey = "billing-save-probe"
|
||||
const modelName = "billing-save-model"
|
||||
source := `
|
||||
export const meta = {
|
||||
apiVersion: 1, key: "billing-save-probe", name: "Billing Save Probe", version: "1.0.0", author: {name: "Test"},
|
||||
models: ["billing-save-model"], fetchMode: "per_task",
|
||||
usageSchema: {seconds: {type: "number", unit: "second"}}
|
||||
};
|
||||
export function buildSubmitRequest() { return {}; }
|
||||
export function parseSubmitResponse() { return {}; }
|
||||
export function buildQueryRequest() { return {}; }
|
||||
export function parseTaskResult() { return {}; }
|
||||
`
|
||||
_, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(pluginKey) })
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
expression string
|
||||
errorText string
|
||||
}{
|
||||
{
|
||||
name: "invalid syntax",
|
||||
expression: `tier("base",`,
|
||||
errorText: "expr compile error",
|
||||
},
|
||||
{
|
||||
name: "undeclared usage key",
|
||||
expression: `tier("base", u("clips") * 0.1)`,
|
||||
errorText: `usage key \"clips\" is not declared`,
|
||||
},
|
||||
}
|
||||
|
||||
for _, testCase := range tests {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
expressions, marshalErr := common.Marshal(map[string]string{modelName: testCase.expression})
|
||||
require.NoError(t, marshalErr)
|
||||
body, marshalErr := common.Marshal(OptionUpdateRequest{
|
||||
Key: "billing_setting.billing_expr",
|
||||
Value: string(expressions),
|
||||
})
|
||||
require.NoError(t, marshalErr)
|
||||
recorder := httptest.NewRecorder()
|
||||
context, _ := gin.CreateTestContext(recorder)
|
||||
context.Request = httptest.NewRequest(http.MethodPut, "/api/option/", strings.NewReader(string(body)))
|
||||
|
||||
UpdateOption(context)
|
||||
|
||||
assert.Equal(t, http.StatusOK, recorder.Code)
|
||||
assert.Contains(t, recorder.Body.String(), `"success":false`)
|
||||
assert.Contains(t, recorder.Body.String(), modelName)
|
||||
assert.Contains(t, recorder.Body.String(), testCase.errorText)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpdateOptionRejectsUsageExpressionWithoutTaskPlugin(t *testing.T) {
|
||||
const modelName = "billing-save-model-without-plugin"
|
||||
expressions, err := common.Marshal(map[string]string{
|
||||
modelName: `u("mode") == "std" ? 1 : 2`,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
body, err := common.Marshal(OptionUpdateRequest{
|
||||
Key: "billing_setting.billing_expr",
|
||||
Value: string(expressions),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
recorder := httptest.NewRecorder()
|
||||
context, _ := gin.CreateTestContext(recorder)
|
||||
context.Request = httptest.NewRequest(
|
||||
http.MethodPut,
|
||||
"/api/option/",
|
||||
strings.NewReader(string(body)),
|
||||
)
|
||||
|
||||
UpdateOption(context)
|
||||
|
||||
assert.Equal(t, http.StatusOK, recorder.Code)
|
||||
assert.Contains(t, recorder.Body.String(), `"success":false`)
|
||||
assert.Contains(t, recorder.Body.String(), modelName)
|
||||
assert.Contains(t, recorder.Body.String(), "mode")
|
||||
assert.Contains(t, recorder.Body.String(), "no task plugin usage schema")
|
||||
}
|
||||
|
||||
func setupBillingAliasOptionDB(t *testing.T) {
|
||||
t.Helper()
|
||||
previousDB := model.DB
|
||||
previousLogDB := model.LOG_DB
|
||||
previousType := common.MainDatabaseType()
|
||||
previousCache := common.MemoryCacheEnabled
|
||||
previousMap := common.OptionMap
|
||||
previousRedis := common.RedisEnabled
|
||||
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, database.AutoMigrate(&model.Channel{}, &model.Option{}, &model.Log{}, &model.AuditLog{}, &model.User{}))
|
||||
model.DB = database
|
||||
model.LOG_DB = database
|
||||
common.SetMainDatabaseType(common.DatabaseTypeSQLite)
|
||||
common.MemoryCacheEnabled = false
|
||||
common.RedisEnabled = false
|
||||
common.OptionMap = map[string]string{}
|
||||
t.Cleanup(func() {
|
||||
model.DB = previousDB
|
||||
model.LOG_DB = previousLogDB
|
||||
common.SetMainDatabaseType(previousType)
|
||||
common.MemoryCacheEnabled = previousCache
|
||||
common.OptionMap = previousMap
|
||||
common.RedisEnabled = previousRedis
|
||||
model.InitChannelCache()
|
||||
})
|
||||
}
|
||||
|
||||
func TestUpdateOptionAliasBillingExprUsesPluginSchema(t *testing.T) {
|
||||
setupBillingAliasOptionDB(t)
|
||||
const pluginKey = "billing-alias-probe"
|
||||
source := `
|
||||
export const meta = {
|
||||
apiVersion: 1, key: "billing-alias-probe", name: "Billing Alias Probe", version: "1.0.0", author: {name: "Test"},
|
||||
models: ["declared-model"], fetchMode: "per_task",
|
||||
usageSchema: {seconds: {type: "number", unit: "second"}}
|
||||
};
|
||||
export function buildSubmitRequest() { return {}; }
|
||||
export function parseSubmitResponse() { return {}; }
|
||||
export function buildQueryRequest() { return {}; }
|
||||
export function parseTaskResult() { return {}; }
|
||||
`
|
||||
_, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(pluginKey) })
|
||||
|
||||
mapping := `{"alias-model":"declared-model"}`
|
||||
require.NoError(t, model.DB.Create(&model.Channel{
|
||||
Id: 1,
|
||||
Type: 54,
|
||||
Key: "key-1",
|
||||
Status: common.ChannelStatusEnabled,
|
||||
Name: "ch-1",
|
||||
Group: "default",
|
||||
Models: "alias-model,declared-model",
|
||||
ModelMapping: &mapping,
|
||||
}).Error)
|
||||
model.InitChannelCache()
|
||||
|
||||
saved := map[string]string{}
|
||||
require.NoError(t, config.GlobalConfig.SaveToDB(func(key, value string) error {
|
||||
saved[key] = value
|
||||
return nil
|
||||
}))
|
||||
t.Cleanup(func() {
|
||||
require.NoError(t, config.GlobalConfig.LoadFromDB(saved))
|
||||
})
|
||||
|
||||
putExpr := func(modelName, expression string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
expressions, marshalErr := common.Marshal(map[string]string{modelName: expression})
|
||||
require.NoError(t, marshalErr)
|
||||
body, marshalErr := common.Marshal(OptionUpdateRequest{
|
||||
Key: "billing_setting.billing_expr",
|
||||
Value: string(expressions),
|
||||
})
|
||||
require.NoError(t, marshalErr)
|
||||
recorder := httptest.NewRecorder()
|
||||
context, _ := gin.CreateTestContext(recorder)
|
||||
context.Request = httptest.NewRequest(http.MethodPut, "/api/option/", strings.NewReader(string(body)))
|
||||
UpdateOption(context)
|
||||
return recorder
|
||||
}
|
||||
|
||||
accepted := putExpr("alias-model", `u("seconds")`)
|
||||
assert.Equal(t, http.StatusOK, accepted.Code)
|
||||
assert.Contains(t, accepted.Body.String(), `"success":true`)
|
||||
|
||||
rejectedKey := putExpr("alias-model", `u("clips")`)
|
||||
assert.Equal(t, http.StatusOK, rejectedKey.Code)
|
||||
assert.Contains(t, rejectedKey.Body.String(), `"success":false`)
|
||||
assert.Contains(t, rejectedKey.Body.String(), `usage key \"clips\" is not declared`)
|
||||
|
||||
unresolvable := putExpr("unknown-alias-model", `u("seconds")`)
|
||||
assert.Equal(t, http.StatusOK, unresolvable.Code)
|
||||
assert.Contains(t, unresolvable.Body.String(), `"success":false`)
|
||||
assert.Contains(t, unresolvable.Body.String(), "no task plugin usage schema")
|
||||
}
|
||||
@@ -0,0 +1,620 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/relay/channel/advancedcustom"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
|
||||
"github.com/shopspring/decimal"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// https://github.com/songquanpeng/one-api/issues/79
|
||||
|
||||
type OpenAISubscriptionResponse struct {
|
||||
Object string `json:"object"`
|
||||
HasPaymentMethod bool `json:"has_payment_method"`
|
||||
SoftLimitUSD float64 `json:"soft_limit_usd"`
|
||||
HardLimitUSD float64 `json:"hard_limit_usd"`
|
||||
SystemHardLimitUSD float64 `json:"system_hard_limit_usd"`
|
||||
AccessUntil int64 `json:"access_until"`
|
||||
}
|
||||
|
||||
type OpenAIUsageDailyCost struct {
|
||||
Timestamp float64 `json:"timestamp"`
|
||||
LineItems []struct {
|
||||
Name string `json:"name"`
|
||||
Cost float64 `json:"cost"`
|
||||
}
|
||||
}
|
||||
|
||||
type OpenAICreditGrants struct {
|
||||
Object string `json:"object"`
|
||||
TotalGranted float64 `json:"total_granted"`
|
||||
TotalUsed float64 `json:"total_used"`
|
||||
TotalAvailable float64 `json:"total_available"`
|
||||
}
|
||||
|
||||
const maxAdvancedCustomBalanceResponseBytes = 256 << 10
|
||||
|
||||
type channelBalanceResult struct {
|
||||
Balance float64
|
||||
RawResponse string
|
||||
}
|
||||
|
||||
type OpenAIUsageResponse struct {
|
||||
Object string `json:"object"`
|
||||
//DailyCosts []OpenAIUsageDailyCost `json:"daily_costs"`
|
||||
TotalUsage float64 `json:"total_usage"` // unit: 0.01 dollar
|
||||
}
|
||||
|
||||
type OpenAISBUsageResponse struct {
|
||||
Msg string `json:"msg"`
|
||||
Data *struct {
|
||||
Credit string `json:"credit"`
|
||||
} `json:"data"`
|
||||
}
|
||||
|
||||
type AIProxyUserOverviewResponse struct {
|
||||
Success bool `json:"success"`
|
||||
Message string `json:"message"`
|
||||
ErrorCode int `json:"error_code"`
|
||||
Data struct {
|
||||
TotalPoints float64 `json:"totalPoints"`
|
||||
} `json:"data"`
|
||||
}
|
||||
|
||||
type API2GPTUsageResponse struct {
|
||||
Object string `json:"object"`
|
||||
TotalGranted float64 `json:"total_granted"`
|
||||
TotalUsed float64 `json:"total_used"`
|
||||
TotalRemaining float64 `json:"total_remaining"`
|
||||
}
|
||||
|
||||
type APGC2DGPTUsageResponse struct {
|
||||
//Grants interface{} `json:"grants"`
|
||||
Object string `json:"object"`
|
||||
TotalAvailable float64 `json:"total_available"`
|
||||
TotalGranted float64 `json:"total_granted"`
|
||||
TotalUsed float64 `json:"total_used"`
|
||||
}
|
||||
|
||||
type SiliconFlowUsageResponse struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
Status bool `json:"status"`
|
||||
Data struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Image string `json:"image"`
|
||||
Email string `json:"email"`
|
||||
IsAdmin bool `json:"isAdmin"`
|
||||
Balance string `json:"balance"`
|
||||
Status string `json:"status"`
|
||||
Introduction string `json:"introduction"`
|
||||
Role string `json:"role"`
|
||||
ChargeBalance string `json:"chargeBalance"`
|
||||
TotalBalance string `json:"totalBalance"`
|
||||
Category string `json:"category"`
|
||||
} `json:"data"`
|
||||
}
|
||||
|
||||
type DeepSeekUsageResponse struct {
|
||||
IsAvailable bool `json:"is_available"`
|
||||
BalanceInfos []struct {
|
||||
Currency string `json:"currency"`
|
||||
TotalBalance string `json:"total_balance"`
|
||||
GrantedBalance string `json:"granted_balance"`
|
||||
ToppedUpBalance string `json:"topped_up_balance"`
|
||||
} `json:"balance_infos"`
|
||||
}
|
||||
|
||||
type OpenRouterCreditResponse struct {
|
||||
Data struct {
|
||||
TotalCredits float64 `json:"total_credits"`
|
||||
TotalUsage float64 `json:"total_usage"`
|
||||
} `json:"data"`
|
||||
}
|
||||
|
||||
// GetAuthHeader get auth header
|
||||
func GetAuthHeader(token string) http.Header {
|
||||
h := http.Header{}
|
||||
h.Add("Authorization", fmt.Sprintf("Bearer %s", token))
|
||||
return h
|
||||
}
|
||||
|
||||
// GetClaudeAuthHeader get claude auth header
|
||||
func GetClaudeAuthHeader(token string) http.Header {
|
||||
h := http.Header{}
|
||||
h.Add("x-api-key", token)
|
||||
h.Add("anthropic-version", "2023-06-01")
|
||||
return h
|
||||
}
|
||||
|
||||
func GetResponseBody(method, url string, channel *model.Channel, headers http.Header) ([]byte, error) {
|
||||
req, err := http.NewRequest(method, url, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for k := range headers {
|
||||
req.Header.Add(k, headers.Get(k))
|
||||
}
|
||||
client, err := service.GetHttpClientWithProxy(channel.GetSetting().Proxy)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
res, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if res.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("status code: %d", res.StatusCode)
|
||||
}
|
||||
body, err := io.ReadAll(res.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
err = res.Body.Close()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return body, nil
|
||||
}
|
||||
|
||||
func updateChannelCloseAIBalance(channel *model.Channel) (float64, error) {
|
||||
url := fmt.Sprintf("%s/dashboard/billing/credit_grants", channel.GetBaseURL())
|
||||
body, err := GetResponseBody("GET", url, channel, GetAuthHeader(channel.Key))
|
||||
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
response := OpenAICreditGrants{}
|
||||
err = common.Unmarshal(body, &response)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
channel.UpdateBalance(response.TotalAvailable)
|
||||
return response.TotalAvailable, nil
|
||||
}
|
||||
|
||||
func updateChannelOpenAISBBalance(channel *model.Channel) (float64, error) {
|
||||
url := fmt.Sprintf("https://api.openai-sb.com/sb-api/user/status?api_key=%s", channel.Key)
|
||||
body, err := GetResponseBody("GET", url, channel, GetAuthHeader(channel.Key))
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
response := OpenAISBUsageResponse{}
|
||||
err = common.Unmarshal(body, &response)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if response.Data == nil {
|
||||
return 0, errors.New(response.Msg)
|
||||
}
|
||||
balance, err := strconv.ParseFloat(response.Data.Credit, 64)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
channel.UpdateBalance(balance)
|
||||
return balance, nil
|
||||
}
|
||||
|
||||
func updateChannelAIProxyBalance(channel *model.Channel) (float64, error) {
|
||||
url := "https://aiproxy.io/api/report/getUserOverview"
|
||||
headers := http.Header{}
|
||||
headers.Add("Api-Key", channel.Key)
|
||||
body, err := GetResponseBody("GET", url, channel, headers)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
response := AIProxyUserOverviewResponse{}
|
||||
err = common.Unmarshal(body, &response)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if !response.Success {
|
||||
return 0, fmt.Errorf("code: %d, message: %s", response.ErrorCode, response.Message)
|
||||
}
|
||||
channel.UpdateBalance(response.Data.TotalPoints)
|
||||
return response.Data.TotalPoints, nil
|
||||
}
|
||||
|
||||
func updateChannelAPI2GPTBalance(channel *model.Channel) (float64, error) {
|
||||
url := "https://api.api2gpt.com/dashboard/billing/credit_grants"
|
||||
body, err := GetResponseBody("GET", url, channel, GetAuthHeader(channel.Key))
|
||||
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
response := API2GPTUsageResponse{}
|
||||
err = common.Unmarshal(body, &response)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
channel.UpdateBalance(response.TotalRemaining)
|
||||
return response.TotalRemaining, nil
|
||||
}
|
||||
|
||||
func updateChannelSiliconFlowBalance(channel *model.Channel) (float64, error) {
|
||||
url := "https://api.siliconflow.cn/v1/user/info"
|
||||
body, err := GetResponseBody("GET", url, channel, GetAuthHeader(channel.Key))
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
response := SiliconFlowUsageResponse{}
|
||||
err = common.Unmarshal(body, &response)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if response.Code != 20000 {
|
||||
return 0, fmt.Errorf("code: %d, message: %s", response.Code, response.Message)
|
||||
}
|
||||
balance, err := strconv.ParseFloat(response.Data.TotalBalance, 64)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
channel.UpdateBalance(balance)
|
||||
return balance, nil
|
||||
}
|
||||
|
||||
func updateChannelDeepSeekBalance(channel *model.Channel) (float64, error) {
|
||||
url := "https://api.deepseek.com/user/balance"
|
||||
body, err := GetResponseBody("GET", url, channel, GetAuthHeader(channel.Key))
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
response := DeepSeekUsageResponse{}
|
||||
err = common.Unmarshal(body, &response)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
index := -1
|
||||
for i, balanceInfo := range response.BalanceInfos {
|
||||
if balanceInfo.Currency == "CNY" {
|
||||
index = i
|
||||
break
|
||||
}
|
||||
}
|
||||
if index == -1 {
|
||||
return 0, errors.New("currency CNY not found")
|
||||
}
|
||||
balance, err := strconv.ParseFloat(response.BalanceInfos[index].TotalBalance, 64)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
channel.UpdateBalance(balance)
|
||||
return balance, nil
|
||||
}
|
||||
|
||||
func updateChannelAIGC2DBalance(channel *model.Channel) (float64, error) {
|
||||
url := "https://api.aigc2d.com/dashboard/billing/credit_grants"
|
||||
body, err := GetResponseBody("GET", url, channel, GetAuthHeader(channel.Key))
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
response := APGC2DGPTUsageResponse{}
|
||||
err = common.Unmarshal(body, &response)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
channel.UpdateBalance(response.TotalAvailable)
|
||||
return response.TotalAvailable, nil
|
||||
}
|
||||
|
||||
func updateChannelOpenRouterBalance(channel *model.Channel) (float64, error) {
|
||||
url := "https://openrouter.ai/api/v1/credits"
|
||||
body, err := GetResponseBody("GET", url, channel, GetAuthHeader(channel.Key))
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
response := OpenRouterCreditResponse{}
|
||||
err = common.Unmarshal(body, &response)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
balance := response.Data.TotalCredits - response.Data.TotalUsage
|
||||
channel.UpdateBalance(balance)
|
||||
return balance, nil
|
||||
}
|
||||
|
||||
func updateChannelMoonshotBalance(channel *model.Channel) (float64, error) {
|
||||
url := "https://api.moonshot.cn/v1/users/me/balance"
|
||||
body, err := GetResponseBody("GET", url, channel, GetAuthHeader(channel.Key))
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
type MoonshotBalanceData struct {
|
||||
AvailableBalance float64 `json:"available_balance"`
|
||||
VoucherBalance float64 `json:"voucher_balance"`
|
||||
CashBalance float64 `json:"cash_balance"`
|
||||
}
|
||||
|
||||
type MoonshotBalanceResponse struct {
|
||||
Code int `json:"code"`
|
||||
Data MoonshotBalanceData `json:"data"`
|
||||
Scode string `json:"scode"`
|
||||
Status bool `json:"status"`
|
||||
}
|
||||
|
||||
response := MoonshotBalanceResponse{}
|
||||
err = common.Unmarshal(body, &response)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if !response.Status || response.Code != 0 {
|
||||
return 0, fmt.Errorf("failed to update moonshot balance, status: %v, code: %d, scode: %s", response.Status, response.Code, response.Scode)
|
||||
}
|
||||
availableBalanceCny := response.Data.AvailableBalance
|
||||
availableBalanceUsd := decimal.NewFromFloat(availableBalanceCny).Div(decimal.NewFromFloat(operation_setting.Price)).InexactFloat64()
|
||||
channel.UpdateBalance(availableBalanceUsd)
|
||||
return availableBalanceUsd, nil
|
||||
}
|
||||
|
||||
func fetchAdvancedCustomBalance(channel *model.Channel) (channelBalanceResult, error) {
|
||||
key := strings.TrimSpace(channel.Key)
|
||||
info := &relaycommon.RelayInfo{
|
||||
RelayFormat: types.RelayFormatOpenAI,
|
||||
RelayMode: relayconstant.RelayModeUnknown,
|
||||
RequestURLPath: dto.AdvancedCustomBalancePath,
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
ChannelType: constant.ChannelTypeAdvancedCustom,
|
||||
ChannelBaseUrl: channel.GetBaseURL(),
|
||||
ApiKey: key,
|
||||
ChannelOtherSettings: channel.GetOtherSettings(),
|
||||
},
|
||||
}
|
||||
requestURL, headers, err := (&advancedcustom.Adaptor{}).BuildBalanceRequest(info)
|
||||
if err != nil {
|
||||
return channelBalanceResult{}, sanitizeFetchModelsError(err, key)
|
||||
}
|
||||
if err := applyFetchModelsHeaderOverrides(channel, key, headers); err != nil {
|
||||
return channelBalanceResult{}, sanitizeFetchModelsError(err, key)
|
||||
}
|
||||
|
||||
request, err := http.NewRequest(http.MethodGet, requestURL, nil)
|
||||
if err != nil {
|
||||
return channelBalanceResult{}, sanitizeFetchModelsError(err, key)
|
||||
}
|
||||
for name, values := range headers {
|
||||
for _, value := range values {
|
||||
request.Header.Add(name, value)
|
||||
}
|
||||
if strings.EqualFold(name, "Host") {
|
||||
request.Host = headers.Get(name)
|
||||
}
|
||||
}
|
||||
client, err := service.GetHttpClientWithProxy(channel.GetSetting().Proxy)
|
||||
if err != nil {
|
||||
return channelBalanceResult{}, sanitizeFetchModelsError(err, key)
|
||||
}
|
||||
response, err := client.Do(request)
|
||||
if err != nil {
|
||||
return channelBalanceResult{}, sanitizeAdvancedCustomRequestError(err, key, requestURL)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
if response.StatusCode != http.StatusOK {
|
||||
return channelBalanceResult{}, fmt.Errorf("status code: %d", response.StatusCode)
|
||||
}
|
||||
body, err := io.ReadAll(io.LimitReader(response.Body, maxAdvancedCustomBalanceResponseBytes+1))
|
||||
if err != nil {
|
||||
return channelBalanceResult{}, sanitizeAdvancedCustomRequestError(err, key, requestURL)
|
||||
}
|
||||
if len(body) > maxAdvancedCustomBalanceResponseBytes {
|
||||
return channelBalanceResult{}, fmt.Errorf("balance response exceeds %d bytes", maxAdvancedCustomBalanceResponseBytes)
|
||||
}
|
||||
|
||||
var validated json.RawMessage
|
||||
if err := common.Unmarshal(body, &validated); err != nil {
|
||||
return channelBalanceResult{}, fmt.Errorf("invalid balance JSON response: %w", err)
|
||||
}
|
||||
if common.GetJsonType(validated) == "object" {
|
||||
var creditSummary struct {
|
||||
Object string `json:"object"`
|
||||
TotalAvailable json.RawMessage `json:"total_available"`
|
||||
}
|
||||
if err := common.Unmarshal(body, &creditSummary); err != nil {
|
||||
return channelBalanceResult{}, fmt.Errorf("invalid balance JSON response: %w", err)
|
||||
}
|
||||
if creditSummary.Object == "credit_summary" &&
|
||||
common.GetJsonType(creditSummary.TotalAvailable) == "number" {
|
||||
var balance float64
|
||||
if err := common.Unmarshal(creditSummary.TotalAvailable, &balance); err == nil &&
|
||||
balance >= 0 &&
|
||||
!math.IsNaN(balance) &&
|
||||
!math.IsInf(balance, 0) {
|
||||
channel.UpdateBalance(balance)
|
||||
return channelBalanceResult{Balance: balance}, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
formatted, err := common.IndentJson(body)
|
||||
if err != nil {
|
||||
return channelBalanceResult{}, fmt.Errorf("invalid balance JSON response: %w", err)
|
||||
}
|
||||
return channelBalanceResult{RawResponse: string(formatted)}, nil
|
||||
}
|
||||
|
||||
func updateChannelBalance(channel *model.Channel) (channelBalanceResult, error) {
|
||||
if channel.Type == constant.ChannelTypeAdvancedCustom {
|
||||
return fetchAdvancedCustomBalance(channel)
|
||||
}
|
||||
balance, err := updateStandardChannelBalance(channel)
|
||||
return channelBalanceResult{Balance: balance}, err
|
||||
}
|
||||
|
||||
func updateStandardChannelBalance(channel *model.Channel) (float64, error) {
|
||||
baseURL := constant.GetChannelBaseURL(channel.Type)
|
||||
if channel.GetBaseURL() == "" {
|
||||
channel.BaseURL = &baseURL
|
||||
}
|
||||
switch channel.Type {
|
||||
case constant.ChannelTypeOpenAI:
|
||||
if channel.GetBaseURL() != "" {
|
||||
baseURL = channel.GetBaseURL()
|
||||
}
|
||||
case constant.ChannelTypeAzure:
|
||||
return 0, errors.New("尚未实现")
|
||||
case constant.ChannelTypeCustom:
|
||||
baseURL = channel.GetBaseURL()
|
||||
//case common.ChannelTypeOpenAISB:
|
||||
// return updateChannelOpenAISBBalance(channel)
|
||||
case constant.ChannelTypeAIProxy:
|
||||
return updateChannelAIProxyBalance(channel)
|
||||
case constant.ChannelTypeAPI2GPT:
|
||||
return updateChannelAPI2GPTBalance(channel)
|
||||
case constant.ChannelTypeAIGC2D:
|
||||
return updateChannelAIGC2DBalance(channel)
|
||||
case constant.ChannelTypeSiliconFlow:
|
||||
return updateChannelSiliconFlowBalance(channel)
|
||||
case constant.ChannelTypeDeepSeek:
|
||||
return updateChannelDeepSeekBalance(channel)
|
||||
case constant.ChannelTypeOpenRouter:
|
||||
return updateChannelOpenRouterBalance(channel)
|
||||
case constant.ChannelTypeMoonshot:
|
||||
return updateChannelMoonshotBalance(channel)
|
||||
default:
|
||||
return 0, errors.New("尚未实现")
|
||||
}
|
||||
url := fmt.Sprintf("%s/v1/dashboard/billing/subscription", baseURL)
|
||||
|
||||
body, err := GetResponseBody("GET", url, channel, GetAuthHeader(channel.Key))
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
subscription := OpenAISubscriptionResponse{}
|
||||
err = common.Unmarshal(body, &subscription)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
now := time.Now()
|
||||
startDate := fmt.Sprintf("%s-01", now.Format("2006-01"))
|
||||
endDate := now.Format("2006-01-02")
|
||||
if !subscription.HasPaymentMethod {
|
||||
startDate = now.AddDate(0, 0, -100).Format("2006-01-02")
|
||||
}
|
||||
url = fmt.Sprintf("%s/v1/dashboard/billing/usage?start_date=%s&end_date=%s", baseURL, startDate, endDate)
|
||||
body, err = GetResponseBody("GET", url, channel, GetAuthHeader(channel.Key))
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
usage := OpenAIUsageResponse{}
|
||||
err = common.Unmarshal(body, &usage)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
balance := subscription.HardLimitUSD - usage.TotalUsage/100
|
||||
channel.UpdateBalance(balance)
|
||||
return balance, nil
|
||||
}
|
||||
|
||||
func UpdateChannelBalance(c *gin.Context) {
|
||||
id, err := strconv.Atoi(c.Param("id"))
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
channel, err := model.CacheGetChannel(id)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if channel.Type == constant.ChannelTypeTaskPlugin {
|
||||
c.JSON(http.StatusOK, gin.H{"success": false, "message": "Task Plugin channels do not support balance queries"})
|
||||
return
|
||||
}
|
||||
if channel.ChannelInfo.IsMultiKey {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "多密钥渠道不支持余额查询",
|
||||
})
|
||||
return
|
||||
}
|
||||
result, err := updateChannelBalance(channel)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
response := gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
}
|
||||
if result.RawResponse == "" {
|
||||
response["balance"] = result.Balance
|
||||
} else {
|
||||
response["raw_response"] = result.RawResponse
|
||||
}
|
||||
c.JSON(http.StatusOK, response)
|
||||
}
|
||||
|
||||
func updateAllChannelsBalance() error {
|
||||
channels, err := model.GetAllChannels(0, 0, true, false)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, channel := range channels {
|
||||
if channel.Status != common.ChannelStatusEnabled {
|
||||
continue
|
||||
}
|
||||
if channel.ChannelInfo.IsMultiKey {
|
||||
continue // skip multi-key channels
|
||||
}
|
||||
// TODO: support Azure
|
||||
//if channel.Type != common.ChannelTypeOpenAI && channel.Type != common.ChannelTypeCustom {
|
||||
// continue
|
||||
//}
|
||||
result, err := updateChannelBalance(channel)
|
||||
if err != nil {
|
||||
continue
|
||||
} else if result.RawResponse == "" {
|
||||
// err is nil & balance <= 0 means quota is used up
|
||||
if result.Balance <= 0 {
|
||||
service.DisableChannel(*types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, "", channel.GetAutoBan()), "余额不足")
|
||||
}
|
||||
}
|
||||
time.Sleep(common.RequestInterval)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func UpdateAllChannelsBalance(c *gin.Context) {
|
||||
// TODO: make it async
|
||||
err := updateAllChannelsBalance()
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
func AutomaticallyUpdateChannels(frequency int) {
|
||||
for {
|
||||
time.Sleep(time.Duration(frequency) * time.Minute)
|
||||
common.SysLog("updating all channels")
|
||||
_ = updateAllChannelsBalance()
|
||||
common.SysLog("channels update done")
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,88 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func GetChannelAffinityCacheStats(c *gin.Context) {
|
||||
stats := service.GetChannelAffinityCacheStats()
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": stats,
|
||||
})
|
||||
}
|
||||
|
||||
func ClearChannelAffinityCache(c *gin.Context) {
|
||||
all := strings.TrimSpace(c.Query("all"))
|
||||
ruleName := strings.TrimSpace(c.Query("rule_name"))
|
||||
|
||||
if all == "true" {
|
||||
deleted := service.ClearChannelAffinityCacheAll()
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": gin.H{
|
||||
"deleted": deleted,
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
if ruleName == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"success": false,
|
||||
"message": "缺少参数：rule_name，或使用 all=true 清空全部",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
deleted, err := service.ClearChannelAffinityCacheByRuleName(ruleName)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": gin.H{
|
||||
"deleted": deleted,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func GetChannelAffinityUsageCacheStats(c *gin.Context) {
|
||||
ruleName := strings.TrimSpace(c.Query("rule_name"))
|
||||
usingGroup := strings.TrimSpace(c.Query("using_group"))
|
||||
keyFp := strings.TrimSpace(c.Query("key_fp"))
|
||||
|
||||
if ruleName == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"success": false,
|
||||
"message": "missing param: rule_name",
|
||||
})
|
||||
return
|
||||
}
|
||||
if keyFp == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"success": false,
|
||||
"message": "missing param: key_fp",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
stats := service.GetChannelAffinityUsageCacheStats(ruleName, usingGroup, keyFp)
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": stats,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,136 @@
|
||||
package controller
|
||||
|
||||
import "github.com/QuantumNous/new-api/model"
|
||||
|
||||
func channelHasSensitiveChanges(channel *PatchChannel, origin *model.Channel, requestData map[string]any) bool {
|
||||
if _, ok := requestData["type"]; ok && channel.Type != origin.Type {
|
||||
return true
|
||||
}
|
||||
if _, ok := requestData["key"]; ok && channel.Key != "" && channel.Key != origin.Key {
|
||||
return true
|
||||
}
|
||||
if _, ok := requestData["base_url"]; ok && !equalStringPtr(channel.BaseURL, origin.BaseURL) {
|
||||
return true
|
||||
}
|
||||
if _, ok := requestData["openai_organization"]; ok && !equalStringPtr(channel.OpenAIOrganization, origin.OpenAIOrganization) {
|
||||
return true
|
||||
}
|
||||
if _, ok := requestData["header_override"]; ok && !equalStringPtr(channel.HeaderOverride, origin.HeaderOverride) {
|
||||
return true
|
||||
}
|
||||
if _, ok := requestData["param_override"]; ok && !equalStringPtr(channel.ParamOverride, origin.ParamOverride) {
|
||||
return true
|
||||
}
|
||||
if _, ok := requestData["setting"]; ok && !equalStringPtr(channel.Setting, origin.Setting) {
|
||||
return true
|
||||
}
|
||||
if _, ok := requestData["other"]; ok && channel.Other != origin.Other {
|
||||
return true
|
||||
}
|
||||
if _, ok := requestData["settings"]; ok && channel.OtherSettings != origin.OtherSettings {
|
||||
return true
|
||||
}
|
||||
if _, ok := requestData["key_mode"]; ok && channel.KeyMode != nil {
|
||||
return true
|
||||
}
|
||||
// Fail closed: any field present in the request that is neither a known
|
||||
// sensitive field (gated above) nor an explicitly classified non-sensitive
|
||||
// field must be treated as sensitive. This keeps a newly added channel field
|
||||
// from silently becoming editable by ChannelWrite-only admins until it is
|
||||
// consciously classified in channelNonSensitiveFields.
|
||||
for field := range requestData {
|
||||
if _, ok := channelSensitiveFields[field]; ok {
|
||||
continue
|
||||
}
|
||||
if _, ok := channelNonSensitiveFields[field]; ok {
|
||||
continue
|
||||
}
|
||||
if _, ok := channelOperationalFields[field]; ok {
|
||||
continue
|
||||
}
|
||||
if _, ok := channelReadOnlyFields[field]; ok {
|
||||
continue
|
||||
}
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// channelSensitiveFields lists the channel fields whose modification requires
|
||||
// ChannelSensitiveWrite. They are each checked individually in
|
||||
// channelHasSensitiveChanges with a precise old-vs-new comparison; this set is
|
||||
// used to exclude them from the fail-closed scan for unknown fields.
|
||||
var channelSensitiveFields = map[string]struct{}{
|
||||
"type": {},
|
||||
"key": {},
|
||||
"base_url": {},
|
||||
"openai_organization": {},
|
||||
"header_override": {},
|
||||
"param_override": {},
|
||||
"setting": {},
|
||||
"other": {},
|
||||
"settings": {},
|
||||
"key_mode": {},
|
||||
}
|
||||
|
||||
// channelOperationalFields lists fields managed by operation endpoints instead
|
||||
// of the general channel edit endpoint.
|
||||
var channelOperationalFields = map[string]struct{}{
|
||||
"status": {},
|
||||
}
|
||||
|
||||
// channelReadOnlyFields lists server-managed/accounting fields that the general
|
||||
// channel edit endpoint must ignore even if a client sends them.
|
||||
var channelReadOnlyFields = map[string]struct{}{
|
||||
"created_time": {},
|
||||
"test_time": {},
|
||||
"response_time": {},
|
||||
"balance": {},
|
||||
"balance_updated_time": {},
|
||||
"used_quota": {},
|
||||
}
|
||||
|
||||
func clearChannelReadOnlyFields(channel *PatchChannel, requestData map[string]any) {
|
||||
if _, ok := requestData["created_time"]; ok {
|
||||
channel.CreatedTime = 0
|
||||
}
|
||||
if _, ok := requestData["test_time"]; ok {
|
||||
channel.TestTime = 0
|
||||
}
|
||||
if _, ok := requestData["response_time"]; ok {
|
||||
channel.ResponseTime = 0
|
||||
}
|
||||
if _, ok := requestData["balance"]; ok {
|
||||
channel.Balance = 0
|
||||
}
|
||||
if _, ok := requestData["balance_updated_time"]; ok {
|
||||
channel.BalanceUpdatedTime = 0
|
||||
}
|
||||
if _, ok := requestData["used_quota"]; ok {
|
||||
channel.UsedQuota = 0
|
||||
}
|
||||
}
|
||||
|
||||
// channelNonSensitiveFields lists routing / server-managed channel
|
||||
// fields a ChannelWrite admin may edit without ChannelSensitiveWrite. When a new
|
||||
// field is added to model.Channel it must be added to either this set or
|
||||
// channelSensitiveFields or channelOperationalFields; otherwise it falls through
|
||||
// to the fail-closed branch and is treated as sensitive. The
|
||||
// TestChannelFieldsAreClassified guard test enforces this.
|
||||
var channelNonSensitiveFields = map[string]struct{}{
|
||||
"id": {},
|
||||
"test_model": {},
|
||||
"name": {},
|
||||
"weight": {},
|
||||
"models": {},
|
||||
"group": {},
|
||||
"model_mapping": {},
|
||||
"status_code_mapping": {},
|
||||
"priority": {},
|
||||
"auto_ban": {},
|
||||
"other_info": {},
|
||||
"tag": {},
|
||||
"remark": {},
|
||||
"channel_info": {},
|
||||
"multi_key_mode": {},
|
||||
}
|
||||
@@ -0,0 +1,204 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestChannelHasSensitiveChanges(t *testing.T) {
|
||||
baseURL := "https://api.example.com"
|
||||
headerOverride := `{"Authorization":"Bearer {api_key}"}`
|
||||
origin := &model.Channel{
|
||||
Type: 1,
|
||||
Key: "old-key",
|
||||
BaseURL: &baseURL,
|
||||
HeaderOverride: &headerOverride,
|
||||
Models: "gpt-4o",
|
||||
Group: "default",
|
||||
}
|
||||
|
||||
t.Run("non-sensitive routing fields", func(t *testing.T) {
|
||||
updated := PatchChannel{Channel: *origin}
|
||||
updated.Models = "gpt-4o,gpt-4o-mini"
|
||||
updated.Group = "vip"
|
||||
|
||||
assert.False(t, channelHasSensitiveChanges(&updated, origin, map[string]any{
|
||||
"models": updated.Models,
|
||||
"group": updated.Group,
|
||||
}))
|
||||
})
|
||||
|
||||
t.Run("key change", func(t *testing.T) {
|
||||
updated := PatchChannel{Channel: *origin}
|
||||
updated.Key = "new-key"
|
||||
|
||||
assert.True(t, channelHasSensitiveChanges(&updated, origin, map[string]any{"key": updated.Key}))
|
||||
})
|
||||
|
||||
t.Run("base url change", func(t *testing.T) {
|
||||
updated := PatchChannel{Channel: *origin}
|
||||
newBaseURL := "https://leak.example.com"
|
||||
updated.BaseURL = &newBaseURL
|
||||
|
||||
assert.True(t, channelHasSensitiveChanges(&updated, origin, map[string]any{"base_url": newBaseURL}))
|
||||
})
|
||||
|
||||
t.Run("header override change", func(t *testing.T) {
|
||||
updated := PatchChannel{Channel: *origin}
|
||||
newHeaderOverride := `{"X-Key":"{api_key}"}`
|
||||
updated.HeaderOverride = &newHeaderOverride
|
||||
|
||||
assert.True(t, channelHasSensitiveChanges(&updated, origin, map[string]any{"header_override": newHeaderOverride}))
|
||||
})
|
||||
|
||||
t.Run("omitted sensitive fields do not use zero values", func(t *testing.T) {
|
||||
updated := PatchChannel{}
|
||||
updated.Id = origin.Id
|
||||
updated.Priority = origin.Priority
|
||||
|
||||
assert.False(t, channelHasSensitiveChanges(&updated, origin, map[string]any{"priority": 10}))
|
||||
})
|
||||
|
||||
t.Run("unknown field fails closed", func(t *testing.T) {
|
||||
updated := PatchChannel{Channel: *origin}
|
||||
|
||||
assert.True(t, channelHasSensitiveChanges(&updated, origin, map[string]any{"future_secret_field": "x"}))
|
||||
})
|
||||
|
||||
t.Run("status is operational", func(t *testing.T) {
|
||||
updated := PatchChannel{Channel: *origin}
|
||||
updated.Status = common.ChannelStatusManuallyDisabled
|
||||
|
||||
assert.False(t, channelHasSensitiveChanges(&updated, origin, map[string]any{"status": updated.Status}))
|
||||
})
|
||||
|
||||
t.Run("read-only fields are ignored by sensitivity check", func(t *testing.T) {
|
||||
updated := PatchChannel{Channel: *origin}
|
||||
updated.Balance = 99
|
||||
updated.UsedQuota = 100
|
||||
updated.ResponseTime = 200
|
||||
|
||||
assert.False(t, channelHasSensitiveChanges(&updated, origin, map[string]any{
|
||||
"balance": updated.Balance,
|
||||
"used_quota": updated.UsedQuota,
|
||||
"response_time": updated.ResponseTime,
|
||||
}))
|
||||
})
|
||||
}
|
||||
|
||||
func TestClearChannelReadOnlyFields(t *testing.T) {
|
||||
channel := PatchChannel{Channel: model.Channel{
|
||||
CreatedTime: 11,
|
||||
TestTime: 22,
|
||||
ResponseTime: 33,
|
||||
Balance: 44.5,
|
||||
BalanceUpdatedTime: 55,
|
||||
UsedQuota: 66,
|
||||
Models: "gpt-4o",
|
||||
Group: "default",
|
||||
}}
|
||||
|
||||
clearChannelReadOnlyFields(&channel, map[string]any{
|
||||
"created_time": channel.CreatedTime,
|
||||
"test_time": channel.TestTime,
|
||||
"response_time": channel.ResponseTime,
|
||||
"balance": channel.Balance,
|
||||
"balance_updated_time": channel.BalanceUpdatedTime,
|
||||
"used_quota": channel.UsedQuota,
|
||||
"models": channel.Models,
|
||||
"group": channel.Group,
|
||||
})
|
||||
|
||||
assert.Zero(t, channel.CreatedTime)
|
||||
assert.Zero(t, channel.TestTime)
|
||||
assert.Zero(t, channel.ResponseTime)
|
||||
assert.Zero(t, channel.Balance)
|
||||
assert.Zero(t, channel.BalanceUpdatedTime)
|
||||
assert.Zero(t, channel.UsedQuota)
|
||||
assert.Equal(t, "gpt-4o", channel.Models)
|
||||
assert.Equal(t, "default", channel.Group)
|
||||
}
|
||||
|
||||
func TestUpdateChannelRejectsStatusField(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(
|
||||
http.MethodPut,
|
||||
"/api/channel/",
|
||||
bytes.NewBufferString(`{"id":1,"status":2}`),
|
||||
)
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
UpdateChannel(ctx)
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
var response struct {
|
||||
Success bool `json:"success"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
|
||||
assert.False(t, response.Success)
|
||||
}
|
||||
|
||||
func TestChannelStatusValidation(t *testing.T) {
|
||||
assert.True(t, isManageableChannelStatus(common.ChannelStatusEnabled))
|
||||
assert.True(t, isManageableChannelStatus(common.ChannelStatusManuallyDisabled))
|
||||
assert.False(t, isManageableChannelStatus(common.ChannelStatusAutoDisabled))
|
||||
assert.False(t, isManageableChannelStatus(0))
|
||||
}
|
||||
|
||||
// TestChannelFieldsAreClassified guards the fail-closed sensitivity check: every
|
||||
// JSON field of PatchChannel (including the embedded model.Channel) must be listed
|
||||
// in channelSensitiveFields, channelNonSensitiveFields, or
|
||||
// channelOperationalFields. A newly added field that is left unclassified will
|
||||
// fail this test, forcing a conscious permission decision instead of silently
|
||||
// defaulting either way.
|
||||
func TestChannelFieldsAreClassified(t *testing.T) {
|
||||
classified := func(name string) bool {
|
||||
if _, ok := channelSensitiveFields[name]; ok {
|
||||
return true
|
||||
}
|
||||
if _, ok := channelNonSensitiveFields[name]; ok {
|
||||
return true
|
||||
}
|
||||
if _, ok := channelOperationalFields[name]; ok {
|
||||
return true
|
||||
}
|
||||
_, ok := channelReadOnlyFields[name]
|
||||
return ok
|
||||
}
|
||||
|
||||
var collect func(rt reflect.Type) []string
|
||||
collect = func(rt reflect.Type) []string {
|
||||
var names []string
|
||||
for i := 0; i < rt.NumField(); i++ {
|
||||
field := rt.Field(i)
|
||||
if field.Anonymous && field.Type.Kind() == reflect.Struct {
|
||||
names = append(names, collect(field.Type)...)
|
||||
continue
|
||||
}
|
||||
name := strings.Split(field.Tag.Get("json"), ",")[0]
|
||||
if name == "" || name == "-" {
|
||||
continue
|
||||
}
|
||||
names = append(names, name)
|
||||
}
|
||||
return names
|
||||
}
|
||||
|
||||
for _, name := range collect(reflect.TypeFor[PatchChannel]()) {
|
||||
assert.Truef(t, classified(name),
|
||||
"channel field %q is not classified; add it to channelSensitiveFields, channelNonSensitiveFields, channelOperationalFields, or channelReadOnlyFields in channel_authz.go", name)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,95 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestShouldRetryHonorsPinRetryMode(t *testing.T) {
|
||||
openaiErr := types.NewOpenAIError(errors.New("upstream"), types.ErrorCodeBadResponseStatusCode, http.StatusInternalServerError)
|
||||
|
||||
c := newPinRetryContext()
|
||||
assert.True(t, shouldRetry(c, openaiErr, 1))
|
||||
|
||||
origin := newPinRetryContext()
|
||||
service.GetChannelConstraints(origin).AddPin(dto.ChannelPin{
|
||||
ChannelId: 2,
|
||||
Source: dto.PinSourceOriginTask,
|
||||
Rank: dto.PinRankOriginTask,
|
||||
RetryMode: dto.PinRetrySameChannel,
|
||||
})
|
||||
assert.True(t, shouldRetry(origin, openaiErr, 1), "origin pin retries on the same channel")
|
||||
|
||||
token := newPinRetryContext()
|
||||
service.GetChannelConstraints(token).AddPin(dto.ChannelPin{
|
||||
ChannelId: 1,
|
||||
Source: dto.PinSourceToken,
|
||||
Rank: dto.PinRankToken,
|
||||
RetryMode: dto.PinRetrySingleAttempt,
|
||||
})
|
||||
assert.False(t, shouldRetry(token, openaiErr, 1), "token pin suppresses retry")
|
||||
}
|
||||
|
||||
func TestShouldRetryTaskRelayHonorsPinRetryMode(t *testing.T) {
|
||||
taskErr := &dto.TaskError{StatusCode: http.StatusInternalServerError}
|
||||
|
||||
c := newPinRetryContext()
|
||||
assert.True(t, shouldRetryTaskRelay(c, 1, taskErr, 1))
|
||||
|
||||
origin := newPinRetryContext()
|
||||
service.GetChannelConstraints(origin).AddPin(dto.ChannelPin{
|
||||
ChannelId: 2,
|
||||
Source: dto.PinSourceOriginTask,
|
||||
Rank: dto.PinRankOriginTask,
|
||||
RetryMode: dto.PinRetrySameChannel,
|
||||
})
|
||||
assert.True(t, shouldRetryTaskRelay(origin, 2, taskErr, 1))
|
||||
|
||||
token := newPinRetryContext()
|
||||
service.GetChannelConstraints(token).AddPin(dto.ChannelPin{
|
||||
ChannelId: 1,
|
||||
Source: dto.PinSourceToken,
|
||||
Rank: dto.PinRankToken,
|
||||
RetryMode: dto.PinRetrySingleAttempt,
|
||||
})
|
||||
assert.False(t, shouldRetryTaskRelay(token, 1, taskErr, 1))
|
||||
}
|
||||
|
||||
func TestSameChannelPinsMergeToStricterRetryMode(t *testing.T) {
|
||||
c := newPinRetryContext()
|
||||
constraints := service.GetChannelConstraints(c)
|
||||
constraints.AddPin(dto.ChannelPin{
|
||||
ChannelId: 7,
|
||||
Source: dto.PinSourceOriginTask,
|
||||
Rank: dto.PinRankOriginTask,
|
||||
RetryMode: dto.PinRetrySameChannel,
|
||||
})
|
||||
constraints.AddPin(dto.ChannelPin{
|
||||
ChannelId: 7,
|
||||
Source: dto.PinSourceToken,
|
||||
Rank: dto.PinRankToken,
|
||||
RetryMode: dto.PinRetrySingleAttempt,
|
||||
})
|
||||
pin, found, overridden := constraints.ResolvedPin()
|
||||
require.True(t, found)
|
||||
assert.Equal(t, 7, pin.ChannelId)
|
||||
assert.Equal(t, dto.PinRetrySingleAttempt, pin.RetryMode)
|
||||
assert.Empty(t, overridden)
|
||||
assert.False(t, shouldRetry(c, types.NewOpenAIError(errors.New("upstream"), types.ErrorCodeBadResponseStatusCode, http.StatusInternalServerError), 1))
|
||||
}
|
||||
|
||||
func newPinRetryContext() *gin.Context {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
return c
|
||||
}
|
||||
@@ -0,0 +1,166 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/service/authz"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func setupTaskPluginBindChannelTest(t *testing.T) {
|
||||
t.Helper()
|
||||
wasMaster := common.IsMasterNode
|
||||
common.IsMasterNode = true
|
||||
previousRedisEnabled := common.RedisEnabled
|
||||
common.RedisEnabled = false
|
||||
originalDB, originalLogDB := model.DB, model.LOG_DB
|
||||
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := database.DB()
|
||||
require.NoError(t, err)
|
||||
sqlDB.SetMaxOpenConns(1)
|
||||
require.NoError(t, database.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.CasbinRule{}, &model.AuthzRole{}, &model.Log{}, &model.AuditLog{}, &model.User{}))
|
||||
model.DB = database
|
||||
model.LOG_DB = database
|
||||
require.NoError(t, authz.Init(database))
|
||||
t.Cleanup(func() {
|
||||
common.IsMasterNode = wasMaster
|
||||
common.RedisEnabled = previousRedisEnabled
|
||||
model.DB = originalDB
|
||||
model.LOG_DB = originalLogDB
|
||||
})
|
||||
}
|
||||
|
||||
func postAddChannel(t *testing.T, userID, role int, body string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
context, _ := gin.CreateTestContext(recorder)
|
||||
context.Set("id", userID)
|
||||
context.Set("role", role)
|
||||
context.Request = httptest.NewRequest(http.MethodPost, "/api/channel", strings.NewReader(body))
|
||||
context.Request.Header.Set("Content-Type", "application/json")
|
||||
AddChannel(context)
|
||||
return recorder
|
||||
}
|
||||
|
||||
func TestAddChannelTaskPluginRequiresBindPermission(t *testing.T) {
|
||||
setupTaskPluginBindChannelTest(t)
|
||||
const key = "channel-bind"
|
||||
source := `
|
||||
export const meta = {apiVersion: 1, key: "channel-bind", name: "Bind", version: "1.0.0", author: {name: "Test"}, models: ["doc"], fetchMode: "per_task"};
|
||||
export function buildSubmitRequest() { return {}; }
|
||||
export function parseSubmitResponse() { return {}; }
|
||||
export function buildQueryRequest() { return {}; }
|
||||
export function parseTaskResult() { return {}; }
|
||||
`
|
||||
_, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(key) })
|
||||
|
||||
taskPluginBody := `{"mode":"single","channel":{"type":61,"name":"plugin-channel","key":"sk","models":"doc","group":"default","base_url":"https://example.com","setting":"{\"task_plugin_key\":\"channel-bind\"}"}}`
|
||||
openaiBody := `{"mode":"single","channel":{"type":1,"name":"openai-channel","key":"sk","models":"gpt","group":"default"}}`
|
||||
|
||||
adminDenied := postAddChannel(t, 2, common.RoleAdminUser, taskPluginBody)
|
||||
assert.Contains(t, adminDenied.Body.String(), "task plugin channels require the task_plugin.bind permission")
|
||||
assert.Contains(t, adminDenied.Body.String(), `"success":false`)
|
||||
|
||||
rootAllowed := postAddChannel(t, 1, common.RoleRootUser, taskPluginBody)
|
||||
assert.Contains(t, rootAllowed.Body.String(), `"success":true`)
|
||||
assert.NotContains(t, rootAllowed.Body.String(), "task_plugin.bind")
|
||||
|
||||
adminOtherType := postAddChannel(t, 2, common.RoleAdminUser, openaiBody)
|
||||
assert.Contains(t, adminOtherType.Body.String(), `"success":true`)
|
||||
assert.NotContains(t, adminOtherType.Body.String(), "task_plugin.bind")
|
||||
}
|
||||
|
||||
func TestUpdateChannelTaskPluginRequiresBindPermission(t *testing.T) {
|
||||
setupTaskPluginBindChannelTest(t)
|
||||
const key = "channel-bind-update"
|
||||
source := `
|
||||
export const meta = {apiVersion: 1, key: "channel-bind-update", name: "Bind", version: "1.0.0", author: {name: "Test"}, models: ["doc"], fetchMode: "per_task"};
|
||||
export function buildSubmitRequest() { return {}; }
|
||||
export function parseSubmitResponse() { return {}; }
|
||||
export function buildQueryRequest() { return {}; }
|
||||
export function parseTaskResult() { return {}; }
|
||||
`
|
||||
_, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(key) })
|
||||
|
||||
baseURL := "https://example.com"
|
||||
setting := `{"task_plugin_key":"channel-bind-update"}`
|
||||
channel := model.Channel{
|
||||
Type: constant.ChannelTypeTaskPlugin,
|
||||
Status: common.ChannelStatusEnabled,
|
||||
Name: "existing-plugin",
|
||||
Models: "doc",
|
||||
Group: "default",
|
||||
Key: "sk",
|
||||
BaseURL: &baseURL,
|
||||
Setting: &setting,
|
||||
}
|
||||
require.NoError(t, channel.Insert())
|
||||
|
||||
payload := fmt.Sprintf(
|
||||
`{"id":%d,"type":61,"name":"existing-plugin","key":"sk","models":"doc","group":"default","base_url":"https://example.com","setting":"{\"task_plugin_key\":\"channel-bind-update\"}"}`,
|
||||
channel.Id,
|
||||
)
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
context, _ := gin.CreateTestContext(recorder)
|
||||
context.Set("id", 2)
|
||||
context.Set("role", common.RoleAdminUser)
|
||||
context.Request = httptest.NewRequest(http.MethodPut, "/api/channel", strings.NewReader(payload))
|
||||
context.Request.Header.Set("Content-Type", "application/json")
|
||||
UpdateChannel(context)
|
||||
assert.Contains(t, recorder.Body.String(), "task plugin channels require the task_plugin.bind permission")
|
||||
assert.Contains(t, recorder.Body.String(), `"success":false`)
|
||||
}
|
||||
|
||||
func TestAddChannelTaskPluginPersistsPluginDefaultBaseURLAndAuditsSource(t *testing.T) {
|
||||
setupTaskPluginBindChannelTest(t)
|
||||
for key, baseURLField := range map[string]string{"bind-default-url": `baseUrl: "http://10.0.0.5:8000/",`, "bind-no-default": ""} {
|
||||
source := fmt.Sprintf(`
|
||||
export const meta = {apiVersion: 1, key: %q, name: "Bind", version: "1.0.0", author: {name: "Test"}, %s models: ["doc"], fetchMode: "per_task"};
|
||||
export function buildSubmitRequest() { return {}; }
|
||||
export function parseSubmitResponse() { return {}; }
|
||||
export function buildQueryRequest() { return {}; }
|
||||
export function parseTaskResult() { return {}; }
|
||||
`, key, baseURLField)
|
||||
_, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(key) })
|
||||
}
|
||||
body := func(pluginKey string) string {
|
||||
return fmt.Sprintf(`{"mode":"single","channel":{"type":61,"name":"%s","key":"sk","models":"doc","group":"default","setting":"{\"task_plugin_key\":\"%s\"}"}}`, pluginKey, pluginKey)
|
||||
}
|
||||
|
||||
noDefault := postAddChannel(t, 1, common.RoleRootUser, body("bind-no-default"))
|
||||
assert.Contains(t, noDefault.Body.String(), "base URL is required for task plugin channels")
|
||||
|
||||
filled := postAddChannel(t, 1, common.RoleRootUser, body("bind-default-url"))
|
||||
require.Contains(t, filled.Body.String(), `"success":true`)
|
||||
var created model.Channel
|
||||
require.NoError(t, model.DB.Where("name = ?", "bind-default-url").First(&created).Error)
|
||||
require.NotNil(t, created.BaseURL)
|
||||
assert.Equal(t, "http://10.0.0.5:8000", *created.BaseURL, "the normalized plugin default is stored on the channel row")
|
||||
|
||||
var audits []model.AuditLog
|
||||
require.NoError(t, model.LOG_DB.Where("action = ?", "channel.create").Find(&audits).Error)
|
||||
encoded, err := common.Marshal(audits)
|
||||
require.NoError(t, err)
|
||||
assert.Contains(t, string(encoded), `"base_url_source":"plugin_default"`)
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestValidateTaskPluginChannel(t *testing.T) {
|
||||
source := `
|
||||
export const meta = {apiVersion: 1, key: "channel-validation", name: "Validation", version: "1.0.0", author: {name: "Test"}, models: ["doc"], fetchMode: "per_task"};
|
||||
export function buildSubmitRequest() { return {}; }
|
||||
export function parseSubmitResponse() { return {}; }
|
||||
export function buildQueryRequest() { return {}; }
|
||||
export function parseTaskResult() { return {}; }
|
||||
`
|
||||
_, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister("channel-validation") })
|
||||
baseURL := "https://example.com"
|
||||
|
||||
channel := &model.Channel{Type: constant.ChannelTypeTaskPlugin, BaseURL: &baseURL}
|
||||
require.ErrorContains(t, validateChannel(channel, false), "task plugin key is required")
|
||||
|
||||
missing := `{"task_plugin_key":"missing"}`
|
||||
channel.Setting = &missing
|
||||
require.ErrorContains(t, validateChannel(channel, false), "is not registered")
|
||||
|
||||
longKey := `{"task_plugin_key":"` + strings.Repeat("x", 31) + `"}`
|
||||
channel.Setting = &longKey
|
||||
require.ErrorContains(t, validateChannel(channel, false), "must not exceed 30")
|
||||
|
||||
valid := `{"task_plugin_key":"channel-validation"}`
|
||||
channel.Setting = &valid
|
||||
channel.BaseURL = nil
|
||||
require.ErrorContains(t, validateChannel(channel, false), "base URL is required")
|
||||
}
|
||||
|
||||
func TestValidateTaskPluginChannelFillsPluginDefaultBaseURL(t *testing.T) {
|
||||
source := `
|
||||
export const meta = {apiVersion: 1, key: "channel-default-url", name: "Default URL", version: "1.0.0", author: {name: "Test"}, baseUrl: "http://127.0.0.1:8000/", models: ["doc"], fetchMode: "per_task"};
|
||||
export function buildSubmitRequest() { return {}; }
|
||||
export function parseSubmitResponse() { return {}; }
|
||||
export function buildQueryRequest() { return {}; }
|
||||
export function parseTaskResult() { return {}; }
|
||||
`
|
||||
_, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister("channel-default-url") })
|
||||
bound := `{"task_plugin_key":"channel-default-url"}`
|
||||
|
||||
empty := " "
|
||||
for _, baseURL := range []*string{nil, &empty} {
|
||||
channel := &model.Channel{Type: constant.ChannelTypeTaskPlugin, Key: "sk", Setting: &bound, BaseURL: baseURL}
|
||||
require.NoError(t, validateChannel(channel, true))
|
||||
require.NotNil(t, channel.BaseURL)
|
||||
assert.Equal(t, "http://127.0.0.1:8000", *channel.BaseURL, "normalized plugin default is persisted onto the channel")
|
||||
}
|
||||
|
||||
explicit := "https://override.example.com"
|
||||
channel := &model.Channel{Type: constant.ChannelTypeTaskPlugin, Setting: &bound, BaseURL: &explicit}
|
||||
require.NoError(t, validateChannel(channel, false))
|
||||
assert.Equal(t, explicit, *channel.BaseURL, "an administrator value is never replaced by the plugin default")
|
||||
}
|
||||
@@ -0,0 +1,468 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
"github.com/QuantumNous/new-api/types"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestValidateChannelProxy(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
proxy string
|
||||
wantErr bool
|
||||
}{
|
||||
{name: "empty"},
|
||||
{name: "http", proxy: "http://proxy.example:8080"},
|
||||
{name: "https", proxy: "https://proxy.example:8443"},
|
||||
{name: "socks5", proxy: "socks5://proxy.example"},
|
||||
{name: "socks5h", proxy: "socks5h://proxy.example:1080/"},
|
||||
{name: "unsupported", proxy: "ftp://proxy.example", wantErr: true},
|
||||
{name: "path", proxy: "socks5://proxy.example:1080/path", wantErr: true},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
setting, err := common.Marshal(dto.ChannelSettings{Proxy: test.proxy})
|
||||
require.NoError(t, err)
|
||||
channel := &model.Channel{
|
||||
Type: constant.ChannelTypeOpenAI,
|
||||
Setting: common.GetPointer(string(setting)),
|
||||
}
|
||||
|
||||
err = validateChannel(channel, false)
|
||||
|
||||
if test.wantErr {
|
||||
require.ErrorContains(t, err, "invalid channel proxy")
|
||||
return
|
||||
}
|
||||
require.NoError(t, err)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateChannelRequiresNewAPIBaseURL(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
baseURL *string
|
||||
wantErr bool
|
||||
}{
|
||||
{name: "missing", wantErr: true},
|
||||
{name: "blank", baseURL: common.GetPointer(" "), wantErr: true},
|
||||
{name: "configured", baseURL: common.GetPointer("https://new-api.example")},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
channel := &model.Channel{
|
||||
Type: constant.ChannelTypeNewAPI,
|
||||
BaseURL: test.baseURL,
|
||||
}
|
||||
|
||||
err := validateChannel(channel, false)
|
||||
|
||||
if test.wantErr {
|
||||
require.ErrorContains(t, err, "New API channel base URL cannot be empty")
|
||||
return
|
||||
}
|
||||
require.NoError(t, err)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewAPIChannelRegistration(t *testing.T) {
|
||||
apiType, ok := common.ChannelType2APIType(constant.ChannelTypeNewAPI)
|
||||
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, constant.APITypeNewAPI, apiType)
|
||||
assert.Equal(t, "New API", constant.GetChannelTypeName(constant.ChannelTypeNewAPI))
|
||||
require.Greater(t, len(constant.ChannelBaseURLs), constant.ChannelTypeNewAPI)
|
||||
assert.Empty(t, constant.ChannelBaseURLs[constant.ChannelTypeNewAPI])
|
||||
}
|
||||
|
||||
func TestResponsesCompactChannelSupport(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
channelType int
|
||||
apiType int
|
||||
want bool
|
||||
}{
|
||||
{name: "OpenAI", channelType: constant.ChannelTypeOpenAI, apiType: constant.APITypeOpenAI, want: true},
|
||||
{name: "Azure", channelType: constant.ChannelTypeAzure, apiType: constant.APITypeOpenAI, want: true},
|
||||
{name: "Codex", channelType: constant.ChannelTypeCodex, apiType: constant.APITypeCodex, want: true},
|
||||
{name: "Advanced Custom", channelType: constant.ChannelTypeAdvancedCustom, apiType: constant.APITypeAdvancedCustom, want: true},
|
||||
{name: "Sub2API", channelType: constant.ChannelTypeSub2API, apiType: constant.APITypeSub2API, want: true},
|
||||
{name: "New API", channelType: constant.ChannelTypeNewAPI, apiType: constant.APITypeNewAPI, want: true},
|
||||
{name: "Anthropic", channelType: constant.ChannelTypeAnthropic, apiType: constant.APITypeAnthropic, want: false},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
assert.Equal(t, test.want, common.SupportsResponsesCompact(test.channelType, test.apiType))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMultiprotocolGatewayEndpointTypes(t *testing.T) {
|
||||
want := []constant.EndpointType{
|
||||
constant.EndpointTypeOpenAI,
|
||||
constant.EndpointTypeOpenAIResponse,
|
||||
constant.EndpointTypeOpenAIResponseCompact,
|
||||
constant.EndpointTypeAnthropic,
|
||||
constant.EndpointTypeGemini,
|
||||
constant.EndpointTypeOpenAIAlphaSearch,
|
||||
}
|
||||
|
||||
assert.Equal(t, want, common.GetEndpointTypesByChannelType(constant.ChannelTypeNewAPI, "gpt-5"))
|
||||
assert.Equal(t, want, common.GetEndpointTypesByChannelType(constant.ChannelTypeSub2API, "gpt-5"))
|
||||
}
|
||||
|
||||
func TestCopyChannelRejectsInvalidLegacyProxySettings(t *testing.T) {
|
||||
db := setupModelListControllerTestDB(t)
|
||||
settingBytes, err := common.Marshal(dto.ChannelSettings{
|
||||
Proxy: "socks5://proxy.example/legacy-path",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
setting := string(settingBytes)
|
||||
origin := &model.Channel{
|
||||
Type: constant.ChannelTypeOpenAI,
|
||||
Name: "legacy proxy channel",
|
||||
Key: "test-key",
|
||||
Models: "gpt-test",
|
||||
Group: "default",
|
||||
Setting: &setting,
|
||||
}
|
||||
require.NoError(t, db.Create(origin).Error)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Params = gin.Params{{Key: "id", Value: fmt.Sprintf("%d", origin.Id)}}
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/api/channel/copy", nil)
|
||||
|
||||
CopyChannel(ctx)
|
||||
|
||||
assert.Contains(t, recorder.Body.String(), "invalid channel settings")
|
||||
var channelCount int64
|
||||
require.NoError(t, db.Model(&model.Channel{}).Count(&channelCount).Error)
|
||||
assert.Equal(t, int64(1), channelCount)
|
||||
}
|
||||
|
||||
func TestDeleteChannelResetsProxyCacheWhenPreReadFails(t *testing.T) {
|
||||
db := setupModelListControllerTestDB(t)
|
||||
require.NoError(t, db.AutoMigrate(&model.Log{}, &model.AuditLog{}))
|
||||
service.ResetProxyClientCache()
|
||||
t.Cleanup(service.ResetProxyClientCache)
|
||||
|
||||
proxyURL := "http://proxy.example:8080"
|
||||
beforeDelete, err := service.GetHttpClientWithProxy(proxyURL)
|
||||
require.NoError(t, err)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Params = gin.Params{{Key: "id", Value: "999999"}}
|
||||
ctx.Request = httptest.NewRequest(http.MethodDelete, "/api/channel/999999", nil)
|
||||
|
||||
DeleteChannel(ctx)
|
||||
|
||||
assert.Contains(t, recorder.Body.String(), `"success":true`)
|
||||
afterDelete, err := service.GetHttpClientWithProxy(proxyURL)
|
||||
require.NoError(t, err)
|
||||
assert.NotSame(t, beforeDelete, afterDelete)
|
||||
}
|
||||
|
||||
func TestDeleteChannelBatchReportsAndAuditsActualDeletedCount(t *testing.T) {
|
||||
db := setupModelListControllerTestDB(t)
|
||||
require.NoError(t, db.AutoMigrate(&model.Log{}, &model.AuditLog{}))
|
||||
channel := &model.Channel{Name: "existing", Key: "test-key"}
|
||||
require.NoError(t, db.Create(channel).Error)
|
||||
|
||||
requestBody, err := common.Marshal(ChannelBatch{Ids: []int{channel.Id, 999999}})
|
||||
require.NoError(t, err)
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(http.MethodDelete, "/api/channel/batch", bytes.NewReader(requestBody))
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
DeleteChannelBatch(ctx)
|
||||
|
||||
var response struct {
|
||||
Success bool `json:"success"`
|
||||
Data int64 `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
|
||||
assert.True(t, response.Success)
|
||||
assert.Equal(t, int64(1), response.Data)
|
||||
|
||||
var auditLog model.AuditLog
|
||||
require.NoError(t, db.Order("id desc").First(&auditLog).Error)
|
||||
var auditData struct {
|
||||
Operation struct {
|
||||
Params map[string]any `json:"params"`
|
||||
} `json:"op"`
|
||||
}
|
||||
encodedAudit, err := common.Marshal(auditLog.Other)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, common.Unmarshal(encodedAudit, &auditData))
|
||||
assert.Equal(t, float64(1), auditData.Operation.Params["count"])
|
||||
}
|
||||
|
||||
func TestSettleTestQuotaUsesTieredBilling(t *testing.T) {
|
||||
info := &relaycommon.RelayInfo{
|
||||
TieredBillingSnapshot: &billingexpr.BillingSnapshot{
|
||||
BillingMode: "tiered_expr",
|
||||
ExprString: `param("stream") == true ? tier("stream", p * 3) : tier("base", p * 2)`,
|
||||
ExprHash: billingexpr.ExprHashString(`param("stream") == true ? tier("stream", p * 3) : tier("base", p * 2)`),
|
||||
GroupRatio: 1,
|
||||
EstimatedTier: "stream",
|
||||
QuotaPerUnit: common.QuotaPerUnit,
|
||||
ExprVersion: 1,
|
||||
},
|
||||
BillingRequestInput: &billingexpr.RequestInput{
|
||||
Body: []byte(`{"stream":true}`),
|
||||
},
|
||||
}
|
||||
|
||||
quota, result := settleTestQuota(info, types.PriceData{
|
||||
ModelRatio: 1,
|
||||
CompletionRatio: 2,
|
||||
}, &dto.Usage{
|
||||
PromptTokens: 1000,
|
||||
})
|
||||
|
||||
require.Equal(t, 1500, quota)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, "stream", result.MatchedTier)
|
||||
}
|
||||
|
||||
func TestBuildTestLogOtherInjectsTieredInfo(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
|
||||
info := &relaycommon.RelayInfo{
|
||||
TieredBillingSnapshot: &billingexpr.BillingSnapshot{
|
||||
BillingMode: "tiered_expr",
|
||||
ExprString: `tier("base", p * 2)`,
|
||||
},
|
||||
ChannelMeta: &relaycommon.ChannelMeta{},
|
||||
}
|
||||
priceData := types.PriceData{
|
||||
GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 1},
|
||||
}
|
||||
usage := &dto.Usage{
|
||||
PromptTokensDetails: dto.InputTokenDetails{
|
||||
CachedTokens: 12,
|
||||
},
|
||||
}
|
||||
|
||||
requestRules := []billingexpr.RequestRuleTrace{{
|
||||
Cond: `param("service_tier") == "fast"`,
|
||||
Multiplier: 2,
|
||||
Matched: true,
|
||||
}}
|
||||
other := buildTestLogOther(ctx, info, priceData, usage, &billingexpr.TieredResult{
|
||||
MatchedTier: "base",
|
||||
RequestRules: requestRules,
|
||||
})
|
||||
|
||||
fields := other.Snapshot()
|
||||
require.Equal(t, "tiered_expr", fields["billing_mode"])
|
||||
require.Equal(t, "base", fields["matched_tier"])
|
||||
require.Equal(t, requestRules, fields["request_rules"])
|
||||
require.NotEmpty(t, fields["expr_b64"])
|
||||
}
|
||||
|
||||
func TestResolveChannelTestUserIDUsesRequestUser(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
ctx.Set("id", 2)
|
||||
|
||||
userID, err := resolveChannelTestUserID(ctx)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 2, userID)
|
||||
}
|
||||
|
||||
func TestSelectChannelsForAutomaticTestPassiveRecoveryOnlyUsesAutoDisabled(t *testing.T) {
|
||||
channels := []*model.Channel{
|
||||
{Id: 1, Status: common.ChannelStatusEnabled},
|
||||
{Id: 2, Status: common.ChannelStatusAutoDisabled},
|
||||
{Id: 3, Status: common.ChannelStatusManuallyDisabled},
|
||||
}
|
||||
|
||||
selected := selectChannelsForAutomaticTest(channels, operation_setting.ChannelTestModePassiveRecovery)
|
||||
|
||||
require.Len(t, selected, 1)
|
||||
require.Equal(t, 2, selected[0].Id)
|
||||
}
|
||||
|
||||
func TestSelectChannelsForAutomaticTestScheduledSkipsManualDisabled(t *testing.T) {
|
||||
channels := []*model.Channel{
|
||||
{Id: 1, Status: common.ChannelStatusEnabled},
|
||||
{Id: 2, Status: common.ChannelStatusAutoDisabled},
|
||||
{Id: 3, Status: common.ChannelStatusManuallyDisabled},
|
||||
}
|
||||
|
||||
selected := selectChannelsForAutomaticTest(channels, operation_setting.ChannelTestModeScheduledAll)
|
||||
|
||||
require.Len(t, selected, 2)
|
||||
require.Equal(t, 1, selected[0].Id)
|
||||
require.Equal(t, 2, selected[1].Id)
|
||||
}
|
||||
|
||||
func TestSelectChannelsForAutomaticTestAutoBanOnlyUsesEligibleChannels(t *testing.T) {
|
||||
autoBanEnabled := 1
|
||||
autoBanDisabled := 0
|
||||
channels := []*model.Channel{
|
||||
{Id: 1, Status: common.ChannelStatusEnabled, AutoBan: &autoBanEnabled},
|
||||
{Id: 2, Status: common.ChannelStatusEnabled, AutoBan: &autoBanDisabled},
|
||||
{Id: 3, Status: common.ChannelStatusAutoDisabled, AutoBan: &autoBanEnabled},
|
||||
{Id: 4, Status: common.ChannelStatusManuallyDisabled, AutoBan: &autoBanEnabled},
|
||||
{Id: 5, Status: common.ChannelStatusEnabled},
|
||||
}
|
||||
|
||||
selected := selectChannelsForAutomaticTest(channels, operation_setting.ChannelTestModeAutoBanOnly)
|
||||
|
||||
require.Len(t, selected, 2)
|
||||
require.Equal(t, 1, selected[0].Id)
|
||||
require.Equal(t, 3, selected[1].Id)
|
||||
}
|
||||
|
||||
func TestRunChannelTestWorkersHonorsConfiguredConcurrency(t *testing.T) {
|
||||
originalInterval := common.RequestInterval
|
||||
common.RequestInterval = 0
|
||||
t.Cleanup(func() { common.RequestInterval = originalInterval })
|
||||
|
||||
channels := []*model.Channel{
|
||||
{Id: 1, Status: common.ChannelStatusEnabled},
|
||||
{Id: 2, Status: common.ChannelStatusEnabled},
|
||||
{Id: 3, Status: common.ChannelStatusEnabled},
|
||||
{Id: 4, Status: common.ChannelStatusEnabled},
|
||||
}
|
||||
started := make(chan struct{}, len(channels))
|
||||
release := make(chan struct{})
|
||||
var active atomic.Int32
|
||||
var maxActive atomic.Int32
|
||||
progress := make([]int, 0, len(channels)+1)
|
||||
summaryResult := make(chan channelTestSummary, 1)
|
||||
|
||||
go func() {
|
||||
summaryResult <- runChannelTestWorkers(
|
||||
context.Background(),
|
||||
channels,
|
||||
2,
|
||||
func(_ context.Context, _ *model.Channel) channelTestSummary {
|
||||
current := active.Add(1)
|
||||
defer active.Add(-1)
|
||||
for {
|
||||
observed := maxActive.Load()
|
||||
if current <= observed || maxActive.CompareAndSwap(observed, current) {
|
||||
break
|
||||
}
|
||||
}
|
||||
started <- struct{}{}
|
||||
<-release
|
||||
return channelTestSummary{Tested: 1, Succeeded: 1}
|
||||
},
|
||||
func(processed, _ int) {
|
||||
progress = append(progress, processed)
|
||||
},
|
||||
)
|
||||
}()
|
||||
|
||||
<-started
|
||||
<-started
|
||||
select {
|
||||
case <-started:
|
||||
t.Fatal("started more channel tests than the configured concurrency")
|
||||
default:
|
||||
}
|
||||
close(release)
|
||||
|
||||
summary := <-summaryResult
|
||||
|
||||
assert.Equal(t, int32(2), maxActive.Load())
|
||||
assert.Equal(t, channelTestSummary{Tested: 4, Succeeded: 4}, summary)
|
||||
assert.Equal(t, []int{0, 1, 2, 3, 4}, progress)
|
||||
}
|
||||
|
||||
func TestRunChannelTestWorkersStopsAfterCancellation(t *testing.T) {
|
||||
originalInterval := common.RequestInterval
|
||||
common.RequestInterval = 0
|
||||
t.Cleanup(func() { common.RequestInterval = originalInterval })
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
channels := []*model.Channel{
|
||||
{Id: 1, Status: common.ChannelStatusEnabled},
|
||||
{Id: 2, Status: common.ChannelStatusEnabled},
|
||||
{Id: 3, Status: common.ChannelStatusEnabled},
|
||||
{Id: 4, Status: common.ChannelStatusEnabled},
|
||||
}
|
||||
started := make(chan struct{}, len(channels))
|
||||
progress := make([]int, 0, 1)
|
||||
summaryResult := make(chan channelTestSummary, 1)
|
||||
|
||||
go func() {
|
||||
summaryResult <- runChannelTestWorkers(
|
||||
ctx,
|
||||
channels,
|
||||
2,
|
||||
func(ctx context.Context, _ *model.Channel) channelTestSummary {
|
||||
started <- struct{}{}
|
||||
<-ctx.Done()
|
||||
return channelTestSummary{Tested: 1, Succeeded: 1}
|
||||
},
|
||||
func(processed, _ int) {
|
||||
progress = append(progress, processed)
|
||||
},
|
||||
)
|
||||
}()
|
||||
|
||||
<-started
|
||||
<-started
|
||||
cancel()
|
||||
|
||||
summary := <-summaryResult
|
||||
|
||||
select {
|
||||
case <-started:
|
||||
t.Fatal("started another channel test after cancellation")
|
||||
default:
|
||||
}
|
||||
assert.Equal(t, channelTestSummary{Tested: 2, Succeeded: 2}, summary)
|
||||
assert.Equal(t, []int{0}, progress)
|
||||
}
|
||||
|
||||
func TestTestAllChannelsRejectsExistingActiveTask(t *testing.T) {
|
||||
db := setupModelListControllerTestDB(t)
|
||||
require.NoError(t, db.AutoMigrate(&model.SystemTask{}, &model.SystemTaskLock{}))
|
||||
|
||||
existing, err := model.CreateSystemTask(model.SystemTaskTypeChannelTest, nil, nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/api/channel/test", nil)
|
||||
|
||||
TestAllChannels(ctx)
|
||||
|
||||
require.Equal(t, http.StatusConflict, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), existing.TaskID)
|
||||
require.Contains(t, recorder.Body.String(), "已有通道测试任务正在运行或等待中")
|
||||
}
|
||||
@@ -0,0 +1,236 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/relay/channel/ali"
|
||||
"github.com/QuantumNous/new-api/relay/channel/openai"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relay/helper"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/setting/model_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/samber/lo"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func convertChatCompatibilityRequest(t *testing.T, request *dto.GeneralOpenAIRequest, channelType int, mapping map[string]string) []byte {
|
||||
t.Helper()
|
||||
oldMode := gin.Mode()
|
||||
gin.SetMode(gin.TestMode)
|
||||
t.Cleanup(func() { gin.SetMode(oldMode) })
|
||||
settings := model_setting.GetGlobalSettings()
|
||||
oldPassThrough, oldBlacklist := settings.PassThroughRequestEnabled, settings.ThinkingModelBlacklist
|
||||
oldEffortTailModels := settings.EffortTailModelIDs
|
||||
settings.PassThroughRequestEnabled = false
|
||||
settings.ThinkingModelBlacklist = nil
|
||||
settings.EffortTailModelIDs = []string{"gpt-5.1-codex-max"}
|
||||
t.Cleanup(func() {
|
||||
settings.PassThroughRequestEnabled = oldPassThrough
|
||||
settings.ThinkingModelBlacklist = oldBlacklist
|
||||
settings.EffortTailModelIDs = oldEffortTailModels
|
||||
})
|
||||
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest("POST", "/v1/chat/completions", nil)
|
||||
if mapping != nil {
|
||||
encoded, err := common.Marshal(mapping)
|
||||
require.NoError(t, err)
|
||||
c.Set("model_mapping", string(encoded))
|
||||
}
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: request.Model,
|
||||
Request: request,
|
||||
RelayFormat: types.RelayFormatOpenAI,
|
||||
IsStream: request.IsStream(nil),
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
ChannelType: channelType,
|
||||
UpstreamModelName: request.Model,
|
||||
SupportStreamOptions: true,
|
||||
},
|
||||
}
|
||||
require.NoError(t, helper.ModelMappedHelper(c, info, request))
|
||||
require.NoError(t, helper.ApplyReasoningModelSuffix(c, info, request))
|
||||
var converted any
|
||||
var err error
|
||||
if channelType == constant.ChannelTypeAli {
|
||||
converted, err = (&ali.Adaptor{}).ConvertOpenAIRequest(c, info, request)
|
||||
} else {
|
||||
converted, err = (&openai.Adaptor{}).ConvertOpenAIRequest(c, info, request)
|
||||
}
|
||||
require.NoError(t, err)
|
||||
encoded, err := common.Marshal(converted)
|
||||
require.NoError(t, err)
|
||||
return encoded
|
||||
}
|
||||
|
||||
func TestChannelTestOpenAIChatCompatibility(t *testing.T) {
|
||||
for _, tt := range []struct {
|
||||
name string
|
||||
model string
|
||||
upstream string
|
||||
endpoint string
|
||||
channelType int
|
||||
stream bool
|
||||
wantLimit string
|
||||
}{
|
||||
{name: "GPT6 automatic", model: "gpt-6-astra", upstream: "gpt-6-astra", channelType: constant.ChannelTypeOpenAI, wantLimit: "max_completion_tokens"},
|
||||
{name: "GPT6 explicit Azure stream", model: "gpt-6-astra", upstream: "gpt-6-astra", endpoint: string(constant.EndpointTypeOpenAI), channelType: constant.ChannelTypeAzure, stream: true, wantLimit: "max_completion_tokens"},
|
||||
{name: "alias maps to GPT6", model: "customer-model", upstream: "gpt-6-astra", channelType: constant.ChannelTypeOpenAI, wantLimit: "max_completion_tokens"},
|
||||
{name: "GPT5 alias maps to Qwen", model: "gpt-5.6-luna", upstream: "qwen-turbo", channelType: constant.ChannelTypeAli, wantLimit: "max_tokens"},
|
||||
{name: "GPT5 stream", model: "gpt-5.6-luna", upstream: "gpt-5.6-luna", channelType: constant.ChannelTypeOpenAI, stream: true, wantLimit: "max_completion_tokens"},
|
||||
{name: "GPT4 explicit", model: "gpt-4.1", upstream: "gpt-4.1", endpoint: string(constant.EndpointTypeOpenAI), channelType: constant.ChannelTypeOpenAI, wantLimit: "max_tokens"},
|
||||
{name: "o series", model: "o3-mini", upstream: "o3-mini", channelType: constant.ChannelTypeAzure, wantLimit: "max_completion_tokens"},
|
||||
} {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
request, ok := buildTestRequest(tt.model, tt.endpoint, &model.Channel{}, tt.stream).(*dto.GeneralOpenAIRequest)
|
||||
require.True(t, ok)
|
||||
encoded := convertChatCompatibilityRequest(t, request, tt.channelType, map[string]string{tt.model: tt.upstream})
|
||||
want := map[string]any{
|
||||
"model": tt.upstream,
|
||||
"messages": []dto.Message{{Role: "user", Content: "hi"}},
|
||||
"stream": tt.stream,
|
||||
tt.wantLimit: 16,
|
||||
}
|
||||
if tt.stream {
|
||||
want["stream_options"] = map[string]any{"include_usage": true}
|
||||
}
|
||||
wantJSON, err := common.Marshal(want)
|
||||
require.NoError(t, err)
|
||||
assert.JSONEq(t, string(wantJSON), string(encoded))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIChatSamplingCompatibility(t *testing.T) {
|
||||
const sampling = `{"temperature":0.2,"top_p":0.8,"logprobs":true,"top_logprobs":5}`
|
||||
for _, tt := range []struct {
|
||||
name string
|
||||
model string
|
||||
effort string
|
||||
reasoning string
|
||||
mapping map[string]string
|
||||
zeroValues bool
|
||||
wantModel string
|
||||
wantEffort string
|
||||
wantRole string
|
||||
wantParams string
|
||||
}{
|
||||
{name: "GPT5.1 explicit none", model: "gpt-5.1", effort: "none", wantEffort: "none", wantRole: "developer", wantParams: sampling},
|
||||
{name: "GPT5.2 default none", model: "gpt-5.2", wantRole: "developer", wantParams: sampling},
|
||||
{name: "GPT5.2 dated snapshot", model: "gpt-5.2-2025-12-11", wantRole: "developer", wantParams: sampling},
|
||||
{name: "GPT5.4 reasoning", model: "gpt-5.4", effort: "high", wantEffort: "high", wantRole: "developer", wantParams: `{}`},
|
||||
{name: "GPT5.4 snapshot none", model: "gpt-5.4-2026-03-05", effort: "none", wantEffort: "none", wantRole: "developer", wantParams: sampling},
|
||||
{name: "explicit zero values", model: "gpt-5.4", effort: "none", zeroValues: true, wantEffort: "none", wantRole: "developer", wantParams: `{"temperature":0,"top_p":0,"logprobs":false}`},
|
||||
{name: "GPT5 original", model: "gpt-5", wantRole: "developer", wantParams: `{}`},
|
||||
{name: "GPT5.6 existing policy", model: "gpt-5.6-luna", wantRole: "developer", wantParams: `{}`},
|
||||
{name: "pro variant", model: "gpt-5.2-pro-2025-12-11", wantRole: "developer", wantParams: `{}`},
|
||||
{name: "chat variant", model: "gpt-5.2-chat-latest", wantRole: "developer", wantParams: `{}`},
|
||||
{name: "codex model name keeps max", model: "gpt-5.1-codex-max", wantRole: "developer", wantParams: `{}`},
|
||||
{name: "GPT6", model: "gpt-6-astra", wantRole: "developer", wantParams: `{}`},
|
||||
{name: "GPT6 snapshot", model: "gpt-6-astra-2026-09-03", wantRole: "developer", wantParams: `{}`},
|
||||
{name: "GPT6 effort suffix", model: "gpt-6-astra-high", wantModel: "gpt-6-astra", wantEffort: "high", wantRole: "developer", wantParams: `{}`},
|
||||
{name: "none effort suffix", model: "gpt-5.2-none", wantModel: "gpt-5.2", wantEffort: "none", wantRole: "developer", wantParams: sampling},
|
||||
{name: "modifier overrides explicit effort", model: "gpt-5.2@thinking:off", effort: "high", wantModel: "gpt-5.2", wantEffort: "none", wantRole: "developer", wantParams: sampling},
|
||||
{name: "mapped modifier wins", model: "customer-model@thinking:off", mapping: map[string]string{"customer-model": "gpt-5.2@effort:high"}, wantModel: "gpt-5.2", wantEffort: "high", wantRole: "developer", wantParams: `{}`},
|
||||
{name: "nested reasoning disabled", model: "gpt-5.2", reasoning: `{"enabled":false}`, wantEffort: "none", wantRole: "developer", wantParams: sampling},
|
||||
{name: "o1 mini role exception", model: "o1-mini", wantRole: "system", wantParams: `{"top_p":0.8,"logprobs":true,"top_logprobs":5}`},
|
||||
{name: "GPT4 unchanged", model: "gpt-4.1", wantRole: "system", wantParams: sampling},
|
||||
{name: "future model unchanged", model: "gpt-7", wantRole: "system", wantParams: sampling},
|
||||
} {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
request := &dto.GeneralOpenAIRequest{
|
||||
Model: tt.model,
|
||||
Messages: []dto.Message{
|
||||
{Role: "system", Content: "first instruction"},
|
||||
{Role: "system", Content: "second instruction"},
|
||||
{Role: "user", Content: "hi"},
|
||||
},
|
||||
ReasoningEffort: tt.effort,
|
||||
}
|
||||
require.NoError(t, common.UnmarshalJsonStr(sampling, request))
|
||||
if tt.reasoning != "" {
|
||||
request.Reasoning = []byte(tt.reasoning)
|
||||
}
|
||||
if tt.zeroValues {
|
||||
request.Temperature = lo.ToPtr(0.0)
|
||||
request.TopP = lo.ToPtr(0.0)
|
||||
request.LogProbs = lo.ToPtr(false)
|
||||
request.TopLogProbs = nil
|
||||
}
|
||||
encoded := convertChatCompatibilityRequest(t, request, constant.ChannelTypeOpenAI, tt.mapping)
|
||||
var want map[string]any
|
||||
require.NoError(t, common.UnmarshalJsonStr(tt.wantParams, &want))
|
||||
want["model"] = tt.model
|
||||
if tt.wantModel != "" {
|
||||
want["model"] = tt.wantModel
|
||||
}
|
||||
want["messages"] = []dto.Message{
|
||||
{Role: tt.wantRole, Content: "first instruction"},
|
||||
{Role: "system", Content: "second instruction"},
|
||||
{Role: "user", Content: "hi"},
|
||||
}
|
||||
if tt.wantEffort != "" {
|
||||
want["reasoning_effort"] = tt.wantEffort
|
||||
}
|
||||
wantJSON, err := common.Marshal(want)
|
||||
require.NoError(t, err)
|
||||
assert.JSONEq(t, string(wantJSON), string(encoded))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIChatTokenLimitCompatibility(t *testing.T) {
|
||||
for _, modelName := range []string{"gpt-5", "o3-mini", "gpt-6-astra"} {
|
||||
for _, tt := range []struct {
|
||||
name string
|
||||
input string
|
||||
want string
|
||||
}{
|
||||
{name: "omitted", input: `{}`, want: `{}`},
|
||||
{name: "legacy only", input: `{"max_tokens":100}`, want: `{"max_completion_tokens":100}`},
|
||||
{name: "completion only", input: `{"max_completion_tokens":50}`, want: `{"max_completion_tokens":50}`},
|
||||
{name: "both positive stay present", input: `{"max_tokens":100,"max_completion_tokens":50}`, want: `{"max_tokens":100,"max_completion_tokens":50}`},
|
||||
{name: "zero completion falls back", input: `{"max_tokens":100,"max_completion_tokens":0}`, want: `{"max_completion_tokens":100}`},
|
||||
{name: "legacy zero stays present", input: `{"max_tokens":0}`, want: `{"max_tokens":0}`},
|
||||
{name: "completion zero stays present", input: `{"max_completion_tokens":0}`, want: `{"max_completion_tokens":0}`},
|
||||
{name: "both zero stay present", input: `{"max_tokens":0,"max_completion_tokens":0}`, want: `{"max_tokens":0,"max_completion_tokens":0}`},
|
||||
} {
|
||||
t.Run(modelName+"/"+tt.name, func(t *testing.T) {
|
||||
request := &dto.GeneralOpenAIRequest{Model: modelName, Messages: []dto.Message{{Role: "user", Content: "hi"}}}
|
||||
require.NoError(t, common.UnmarshalJsonStr(tt.input, request))
|
||||
encoded := convertChatCompatibilityRequest(t, request, constant.ChannelTypeOpenAI, nil)
|
||||
want := dto.GeneralOpenAIRequest{Model: modelName, Messages: []dto.Message{{Role: "user", Content: "hi"}}}
|
||||
require.NoError(t, common.UnmarshalJsonStr(tt.want, &want))
|
||||
wantJSON, err := common.Marshal(want)
|
||||
require.NoError(t, err)
|
||||
assert.JSONEq(t, string(wantJSON), string(encoded))
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDirectOpenAIResponsesKeepsExistingParameters(t *testing.T) {
|
||||
const body = `{"model":"gpt-6-astra","input":"hi","max_output_tokens":100,"temperature":0.2,"top_p":0.8,"top_logprobs":5,"include":["message.output_text.logprobs"],"reasoning":{"effort":"high"}}`
|
||||
var request dto.OpenAIResponsesRequest
|
||||
require.NoError(t, common.UnmarshalJsonStr(body, &request))
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: "gpt-6-astra",
|
||||
RelayFormat: types.RelayFormatOpenAIResponses,
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
ChannelType: constant.ChannelTypeOpenAI,
|
||||
UpstreamModelName: "gpt-6-astra",
|
||||
},
|
||||
}
|
||||
converted, err := (&openai.Adaptor{}).ConvertOpenAIResponsesRequest(nil, info, request)
|
||||
require.NoError(t, err)
|
||||
encoded, err := common.Marshal(converted)
|
||||
require.NoError(t, err)
|
||||
assert.JSONEq(t, body, string(encoded))
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,606 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func newAdvancedCustomModelListChannel(baseURL string, key string, upstreamPath string, auth *dto.AdvancedCustomRouteAuth) *model.Channel {
|
||||
config := &dto.AdvancedCustomConfig{
|
||||
Routes: []dto.AdvancedCustomRoute{
|
||||
{
|
||||
IncomingPath: dto.AdvancedCustomModelListPath,
|
||||
UpstreamPath: upstreamPath,
|
||||
Converter: "none",
|
||||
Auth: auth,
|
||||
},
|
||||
},
|
||||
}
|
||||
channel := &model.Channel{
|
||||
Type: constant.ChannelTypeAdvancedCustom,
|
||||
Key: key,
|
||||
BaseURL: &baseURL,
|
||||
}
|
||||
channel.SetOtherSettings(dto.ChannelOtherSettings{AdvancedCustom: config})
|
||||
return channel
|
||||
}
|
||||
|
||||
func TestParseOpenAIModelIDsStrictResponseContract(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
body string
|
||||
want []string
|
||||
wantError string
|
||||
}{
|
||||
{name: "malformed JSON", body: `{"data":`, wantError: "invalid OpenAI Models response"},
|
||||
{name: "missing data", body: `{"object":"list"}`, wantError: "data is required"},
|
||||
{name: "null data", body: `{"data":null}`, wantError: "data is required"},
|
||||
{name: "empty data", body: `{"data":[]}`, wantError: "no valid model IDs"},
|
||||
{name: "all IDs empty", body: `{"data":[{"id":""},{"id":" "}]}`, wantError: "no valid model IDs"},
|
||||
{
|
||||
name: "filters empty IDs and normalizes valid IDs",
|
||||
body: `{"data":[{"id":" gpt-4.1 "},{"id":""},{"id":"gpt-4.1"},{"id":"o3"}]}`,
|
||||
want: []string{"gpt-4.1", "o3"},
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
models, err := parseOpenAIModelIDs([]byte(test.body))
|
||||
if test.wantError != "" {
|
||||
require.ErrorContains(t, err, test.wantError)
|
||||
require.Nil(t, models)
|
||||
return
|
||||
}
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, test.want, models)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchAdvancedCustomModelsAppliesHeaderOverrideAfterRouteAuth(t *testing.T) {
|
||||
type receivedRequest struct {
|
||||
Headers http.Header
|
||||
Host string
|
||||
}
|
||||
received := make(chan receivedRequest, 1)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
received <- receivedRequest{Headers: r.Header.Clone(), Host: r.Host}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"data":[{"id":"gpt-4.1"}]}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
channel := newAdvancedCustomModelListChannel(server.URL, "secret-key", "/provider/models", &dto.AdvancedCustomRouteAuth{
|
||||
Type: dto.AdvancedCustomAuthTypeHeader,
|
||||
Name: "X-Route-Key",
|
||||
Value: "route-{api_key}",
|
||||
})
|
||||
headerOverride := `{
|
||||
"X-Route-Key":"global-{api_key}",
|
||||
"X-Static":"static-value",
|
||||
"X-Client":"{client_header:X-Client}",
|
||||
"Host":"models.example.test",
|
||||
"*":""
|
||||
}`
|
||||
channel.HeaderOverride = &headerOverride
|
||||
|
||||
models, err := fetchChannelUpstreamModelIDs(channel)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"gpt-4.1"}, models)
|
||||
|
||||
request := <-received
|
||||
require.Equal(t, "global-secret-key", request.Headers.Get("X-Route-Key"))
|
||||
require.Equal(t, "static-value", request.Headers.Get("X-Static"))
|
||||
require.Empty(t, request.Headers.Get("X-Client"))
|
||||
require.Equal(t, "models.example.test", request.Host)
|
||||
}
|
||||
|
||||
func TestFetchAdvancedCustomModelsUsesEnabledSavedMultiKey(t *testing.T) {
|
||||
authorization := make(chan string, 1)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
authorization <- r.Header.Get("Authorization")
|
||||
_, _ = w.Write([]byte(`{"data":[{"id":"gpt-4.1-mini"}]}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
channel := newAdvancedCustomModelListChannel(server.URL, "disabled-key\nenabled-key", "/v1/models", nil)
|
||||
channel.ChannelInfo = model.ChannelInfo{
|
||||
IsMultiKey: true,
|
||||
MultiKeyStatusList: map[int]int{
|
||||
0: common.ChannelStatusManuallyDisabled,
|
||||
1: common.ChannelStatusEnabled,
|
||||
},
|
||||
}
|
||||
|
||||
models, err := fetchChannelUpstreamModelIDs(channel)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"gpt-4.1-mini"}, models)
|
||||
require.Equal(t, "Bearer enabled-key", <-authorization)
|
||||
}
|
||||
|
||||
func TestFetchAdvancedCustomModelsRejectsNonOKResponse(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusBadGateway)
|
||||
_, _ = w.Write([]byte(`{"data":[{"id":"must-not-be-used"}]}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
channel := newAdvancedCustomModelListChannel(server.URL, "secret-key", "/v1/models", nil)
|
||||
models, err := fetchChannelUpstreamModelIDs(channel)
|
||||
require.ErrorContains(t, err, "status code: 502")
|
||||
require.Nil(t, models)
|
||||
}
|
||||
|
||||
func TestFetchAdvancedCustomModelsRedactsQueryKeyFromTransportErrors(t *testing.T) {
|
||||
const secret = "secret key/+"
|
||||
server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}))
|
||||
baseURL := server.URL
|
||||
server.Close()
|
||||
|
||||
channel := newAdvancedCustomModelListChannel(baseURL, secret, "/v1/models", &dto.AdvancedCustomRouteAuth{
|
||||
Type: dto.AdvancedCustomAuthTypeQuery,
|
||||
Name: "custom-token",
|
||||
Value: "prefix-{api_key}",
|
||||
})
|
||||
|
||||
_, err := fetchChannelUpstreamModelIDs(channel)
|
||||
require.Error(t, err)
|
||||
require.NotContains(t, err.Error(), secret)
|
||||
require.NotContains(t, err.Error(), "custom-token")
|
||||
require.NotContains(t, err.Error(), "prefix-")
|
||||
|
||||
direct := sanitizeFetchModelsError(&url.Error{
|
||||
Op: http.MethodGet,
|
||||
URL: baseURL + "/v1/models?custom-token=prefix-" + url.QueryEscape(secret),
|
||||
Err: errors.New("connection refused"),
|
||||
}, secret)
|
||||
require.EqualError(t, direct, "connection refused")
|
||||
|
||||
queryValue := "prefix-" + secret
|
||||
queryError := sanitizeAdvancedCustomRequestError(
|
||||
errors.New("dial "+queryValue+": connection refused"),
|
||||
queryValue,
|
||||
baseURL+"/v1/models?custom-token="+url.QueryEscape(queryValue),
|
||||
)
|
||||
require.NotContains(t, queryError.Error(), queryValue)
|
||||
require.EqualError(t, queryError, "dial [REDACTED]: connection refused")
|
||||
}
|
||||
|
||||
func TestFetchOrdinaryOpenAIModelsKeepsExistingEmptyDataBehavior(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
_, _ = w.Write([]byte(`{"object":"list"}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
baseURL := server.URL
|
||||
channel := &model.Channel{
|
||||
Type: constant.ChannelTypeOpenAI,
|
||||
Key: "ordinary-key",
|
||||
BaseURL: &baseURL,
|
||||
}
|
||||
models, err := fetchChannelUpstreamModelIDs(channel)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, models)
|
||||
}
|
||||
|
||||
func TestFetchModelsAdvancedCustomCreatePreview(t *testing.T) {
|
||||
receivedAuthorization := make(chan string, 1)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
receivedAuthorization <- r.Header.Get("Authorization")
|
||||
_, _ = w.Write([]byte(`{"data":[{"id":"preview-model"}]}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
config := dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{
|
||||
IncomingPath: dto.AdvancedCustomModelListPath,
|
||||
UpstreamPath: "/preview/models",
|
||||
Converter: "none",
|
||||
}}}
|
||||
configBytes, err := common.Marshal(config)
|
||||
require.NoError(t, err)
|
||||
rawConfig := string(configBytes)
|
||||
baseURL := server.URL
|
||||
emptyProxy := ""
|
||||
req := fetchModelsRequest{
|
||||
BaseURL: &baseURL,
|
||||
Type: constant.ChannelTypeAdvancedCustom,
|
||||
Key: "create-preview-key",
|
||||
AdvancedCustom: &rawConfig,
|
||||
Proxy: &emptyProxy,
|
||||
}
|
||||
body, err := common.Marshal(req)
|
||||
require.NoError(t, err)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/api/channel/fetch_models", bytes.NewReader(body))
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
FetchModels(ctx)
|
||||
|
||||
var response struct {
|
||||
Success bool `json:"success"`
|
||||
Message string `json:"message"`
|
||||
Data []string `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
|
||||
require.True(t, response.Success, response.Message)
|
||||
require.Equal(t, []string{"preview-model"}, response.Data)
|
||||
require.Equal(t, "Bearer create-preview-key", <-receivedAuthorization)
|
||||
}
|
||||
|
||||
func TestFetchModelsAdvancedCustomEditPreviewUsesSavedKeyAndExplicitClears(t *testing.T) {
|
||||
db := setupModelListControllerTestDB(t)
|
||||
receivedHeaders := make(chan http.Header, 1)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
receivedHeaders <- r.Header.Clone()
|
||||
_, _ = w.Write([]byte(`{"data":[{"id":"edited-preview-model"}]}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
savedChannel := newAdvancedCustomModelListChannel("http://127.0.0.1:1", "disabled-saved-key\nenabled-saved-key", "/saved/models", nil)
|
||||
savedChannel.Name = "saved advanced channel"
|
||||
savedChannel.Models = "old-model"
|
||||
savedChannel.ChannelInfo = model.ChannelInfo{
|
||||
IsMultiKey: true,
|
||||
MultiKeyStatusList: map[int]int{
|
||||
0: common.ChannelStatusManuallyDisabled,
|
||||
1: common.ChannelStatusEnabled,
|
||||
},
|
||||
}
|
||||
savedHeaderOverride := `{"X-Saved":"must-not-be-sent"}`
|
||||
savedChannel.HeaderOverride = &savedHeaderOverride
|
||||
savedChannel.SetSetting(dto.ChannelSettings{Proxy: "http://127.0.0.1:1"})
|
||||
require.NoError(t, db.Create(savedChannel).Error)
|
||||
|
||||
preserved, err := buildAdvancedCustomModelPreviewChannel(fetchModelsRequest{ChannelID: savedChannel.Id})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "http://127.0.0.1:1", preserved.GetBaseURL())
|
||||
require.Equal(t, savedHeaderOverride, *preserved.HeaderOverride)
|
||||
require.Equal(t, "http://127.0.0.1:1", preserved.GetSetting().Proxy)
|
||||
|
||||
previewConfig := dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{
|
||||
IncomingPath: dto.AdvancedCustomModelListPath,
|
||||
UpstreamPath: "/edited/models",
|
||||
Converter: "none",
|
||||
}}}
|
||||
configBytes, err := common.Marshal(previewConfig)
|
||||
require.NoError(t, err)
|
||||
rawConfig := string(configBytes)
|
||||
baseURL := server.URL
|
||||
explicitEmpty := ""
|
||||
req := fetchModelsRequest{
|
||||
ChannelID: savedChannel.Id,
|
||||
BaseURL: &baseURL,
|
||||
Type: constant.ChannelTypeAdvancedCustom,
|
||||
Key: "request-key-must-be-ignored",
|
||||
AdvancedCustom: &rawConfig,
|
||||
HeaderOverride: &explicitEmpty,
|
||||
Proxy: &explicitEmpty,
|
||||
}
|
||||
cleared, err := buildAdvancedCustomModelPreviewChannel(fetchModelsRequest{
|
||||
ChannelID: savedChannel.Id,
|
||||
BaseURL: &explicitEmpty,
|
||||
AdvancedCustom: &rawConfig,
|
||||
HeaderOverride: &explicitEmpty,
|
||||
Proxy: &explicitEmpty,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, cleared.BaseURL)
|
||||
require.Empty(t, *cleared.BaseURL)
|
||||
require.NotNil(t, cleared.HeaderOverride)
|
||||
require.Empty(t, *cleared.HeaderOverride)
|
||||
require.Empty(t, cleared.GetSetting().Proxy)
|
||||
|
||||
body, err := common.Marshal(req)
|
||||
require.NoError(t, err)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/api/channel/fetch_models", bytes.NewReader(body))
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
FetchModels(ctx)
|
||||
|
||||
var response struct {
|
||||
Success bool `json:"success"`
|
||||
Message string `json:"message"`
|
||||
Data []string `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
|
||||
require.True(t, response.Success, response.Message)
|
||||
require.Equal(t, []string{"edited-preview-model"}, response.Data)
|
||||
require.NotContains(t, recorder.Body.String(), "enabled-saved-key")
|
||||
require.NotContains(t, recorder.Body.String(), "request-key-must-be-ignored")
|
||||
|
||||
headers := <-receivedHeaders
|
||||
require.Equal(t, "Bearer enabled-saved-key", headers.Get("Authorization"))
|
||||
require.Empty(t, headers.Get("X-Saved"))
|
||||
}
|
||||
|
||||
func TestFailedAdvancedCustomDetectionDoesNotStageFullRemoval(t *testing.T) {
|
||||
db := setupModelListControllerTestDB(t)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
_, _ = w.Write([]byte(`{"data":[]}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
channel := newAdvancedCustomModelListChannel(server.URL, "secret-key", "/v1/models", nil)
|
||||
channel.Name = "empty discovery response"
|
||||
channel.Models = "gpt-4.1,o3"
|
||||
settings := channel.GetOtherSettings()
|
||||
settings.UpstreamModelUpdateCheckEnabled = true
|
||||
settings.UpstreamModelUpdateAutoSyncEnabled = true
|
||||
channel.SetOtherSettings(settings)
|
||||
require.NoError(t, db.Create(channel).Error)
|
||||
|
||||
modelsChanged, autoAdded, err := checkAndPersistChannelUpstreamModelUpdates(channel, &settings, true, true)
|
||||
require.ErrorContains(t, err, "no valid model IDs")
|
||||
require.False(t, modelsChanged)
|
||||
require.Zero(t, autoAdded)
|
||||
require.Empty(t, settings.UpstreamModelUpdateLastDetectedModels)
|
||||
require.Empty(t, settings.UpstreamModelUpdateLastRemovedModels)
|
||||
|
||||
reloaded, err := model.GetChannelById(channel.Id, true)
|
||||
require.NoError(t, err)
|
||||
persistedSettings := reloaded.GetOtherSettings()
|
||||
require.Empty(t, persistedSettings.UpstreamModelUpdateLastDetectedModels)
|
||||
require.Empty(t, persistedSettings.UpstreamModelUpdateLastRemovedModels)
|
||||
require.Equal(t, "gpt-4.1,o3", reloaded.Models)
|
||||
}
|
||||
|
||||
func TestFetchModelsUsesSharedChannelFetchBehavior(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/v1/models" {
|
||||
t.Errorf("unexpected path: %s", r.URL.Path)
|
||||
}
|
||||
if r.Header.Get("x-api-key") != "first-key" {
|
||||
t.Errorf("unexpected x-api-key header: %s", r.Header.Get("x-api-key"))
|
||||
}
|
||||
if r.Header.Get("Authorization") != "" {
|
||||
t.Errorf("unexpected Authorization header: %s", r.Header.Get("Authorization"))
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"data":[{"id":" claude-sonnet "},{"id":"claude-sonnet"}]}`))
|
||||
}))
|
||||
t.Cleanup(server.Close)
|
||||
|
||||
body, err := common.Marshal(map[string]any{
|
||||
"base_url": server.URL,
|
||||
"type": constant.ChannelTypeAnthropic,
|
||||
"key": "first-key\nsecond-key",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/api/channel/fetch_models", bytes.NewReader(body))
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
FetchModels(ctx)
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
require.JSONEq(t, `{"success":true,"message":"","data":["claude-sonnet"]}`, recorder.Body.String())
|
||||
}
|
||||
|
||||
func TestFetchNewAPIModelsUsesOpenAIContract(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
assert.Equal(t, "/v1/models", r.URL.Path)
|
||||
assert.Equal(t, "Bearer new-api-key", r.Header.Get("Authorization"))
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, err := w.Write([]byte(`{"data":[{"id":"gpt-5"},{"id":" gpt-5-mini "}]}`))
|
||||
assert.NoError(t, err)
|
||||
}))
|
||||
t.Cleanup(server.Close)
|
||||
|
||||
baseURL := server.URL
|
||||
channel := &model.Channel{
|
||||
Type: constant.ChannelTypeNewAPI,
|
||||
Key: "new-api-key",
|
||||
BaseURL: &baseURL,
|
||||
}
|
||||
|
||||
models, err := fetchChannelUpstreamModelIDs(channel)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"gpt-5", "gpt-5-mini"}, models)
|
||||
}
|
||||
|
||||
func TestNormalizeModelNames(t *testing.T) {
|
||||
result := normalizeModelNames([]string{
|
||||
" gpt-4o ",
|
||||
"",
|
||||
"gpt-4o",
|
||||
"gpt-4.1",
|
||||
" ",
|
||||
})
|
||||
|
||||
require.Equal(t, []string{"gpt-4o", "gpt-4.1"}, result)
|
||||
}
|
||||
|
||||
func TestMergeModelNames(t *testing.T) {
|
||||
result := mergeModelNames(
|
||||
[]string{"gpt-4o", "gpt-4.1"},
|
||||
[]string{"gpt-4.1", " gpt-4.1-mini ", "gpt-4o"},
|
||||
)
|
||||
|
||||
require.Equal(t, []string{"gpt-4o", "gpt-4.1", "gpt-4.1-mini"}, result)
|
||||
}
|
||||
|
||||
func TestSubtractModelNames(t *testing.T) {
|
||||
result := subtractModelNames(
|
||||
[]string{"gpt-4o", "gpt-4.1", "gpt-4.1-mini"},
|
||||
[]string{"gpt-4.1", "not-exists"},
|
||||
)
|
||||
|
||||
require.Equal(t, []string{"gpt-4o", "gpt-4.1-mini"}, result)
|
||||
}
|
||||
|
||||
func TestIntersectModelNames(t *testing.T) {
|
||||
result := intersectModelNames(
|
||||
[]string{"gpt-4o", "gpt-4.1", "gpt-4.1", "not-exists"},
|
||||
[]string{"gpt-4.1", "gpt-4o-mini", "gpt-4o"},
|
||||
)
|
||||
|
||||
require.Equal(t, []string{"gpt-4o", "gpt-4.1"}, result)
|
||||
}
|
||||
|
||||
func TestApplySelectedModelChanges(t *testing.T) {
|
||||
t.Run("add and remove together", func(t *testing.T) {
|
||||
result := applySelectedModelChanges(
|
||||
[]string{"gpt-4o", "gpt-4.1", "claude-3"},
|
||||
[]string{"gpt-4.1-mini"},
|
||||
[]string{"claude-3"},
|
||||
)
|
||||
|
||||
require.Equal(t, []string{"gpt-4o", "gpt-4.1", "gpt-4.1-mini"}, result)
|
||||
})
|
||||
|
||||
t.Run("add wins when conflict with remove", func(t *testing.T) {
|
||||
result := applySelectedModelChanges(
|
||||
[]string{"gpt-4o"},
|
||||
[]string{"gpt-4.1"},
|
||||
[]string{"gpt-4.1"},
|
||||
)
|
||||
|
||||
require.Equal(t, []string{"gpt-4o", "gpt-4.1"}, result)
|
||||
})
|
||||
}
|
||||
|
||||
func TestCollectPendingApplyUpstreamModelChanges(t *testing.T) {
|
||||
settings := dto.ChannelOtherSettings{
|
||||
UpstreamModelUpdateLastDetectedModels: []string{" gpt-4o ", "gpt-4o", "gpt-4.1"},
|
||||
UpstreamModelUpdateLastRemovedModels: []string{" old-model ", "", "old-model"},
|
||||
}
|
||||
|
||||
pendingAddModels, pendingRemoveModels := collectPendingApplyUpstreamModelChanges(settings)
|
||||
|
||||
require.Equal(t, []string{"gpt-4o", "gpt-4.1"}, pendingAddModels)
|
||||
require.Equal(t, []string{"old-model"}, pendingRemoveModels)
|
||||
}
|
||||
|
||||
func TestNormalizeChannelModelMapping(t *testing.T) {
|
||||
modelMapping := `{
|
||||
" alias-model ": " upstream-model ",
|
||||
"": "invalid",
|
||||
"invalid-target": ""
|
||||
}`
|
||||
channel := &model.Channel{
|
||||
ModelMapping: &modelMapping,
|
||||
}
|
||||
|
||||
result := normalizeChannelModelMapping(channel)
|
||||
require.Equal(t, map[string]string{
|
||||
"alias-model": "upstream-model",
|
||||
}, result)
|
||||
}
|
||||
|
||||
func TestCollectPendingUpstreamModelChangesFromModels_WithModelMapping(t *testing.T) {
|
||||
pendingAddModels, pendingRemoveModels := collectPendingUpstreamModelChangesFromModels(
|
||||
[]string{"alias-model", "gpt-4o", "stale-model"},
|
||||
[]string{"gpt-4o", "gpt-4.1", "mapped-target"},
|
||||
[]string{"gpt-4.1"},
|
||||
map[string]string{
|
||||
"alias-model": "mapped-target",
|
||||
},
|
||||
)
|
||||
|
||||
require.Equal(t, []string{}, pendingAddModels)
|
||||
require.Equal(t, []string{"stale-model"}, pendingRemoveModels)
|
||||
}
|
||||
|
||||
func TestCollectPendingUpstreamModelChangesFromModels_WithIgnoredRegexPatterns(t *testing.T) {
|
||||
pendingAddModels, pendingRemoveModels := collectPendingUpstreamModelChangesFromModels(
|
||||
[]string{"gpt-4o"},
|
||||
[]string{"gpt-4o", "claude-3-5-sonnet", "sora-video", "gpt-4.1"},
|
||||
[]string{"regex:^sora-.*$", "gpt-4.1"},
|
||||
nil,
|
||||
)
|
||||
|
||||
require.Equal(t, []string{"claude-3-5-sonnet"}, pendingAddModels)
|
||||
require.Equal(t, []string{}, pendingRemoveModels)
|
||||
}
|
||||
|
||||
func TestBuildUpstreamModelUpdateTaskNotificationContent_OmitOverflowDetails(t *testing.T) {
|
||||
channelSummaries := make([]upstreamModelUpdateChannelSummary, 0, 12)
|
||||
for i := range 12 {
|
||||
channelSummaries = append(channelSummaries, upstreamModelUpdateChannelSummary{
|
||||
ChannelName: "channel-" + string(rune('A'+i)),
|
||||
AddCount: i + 1,
|
||||
RemoveCount: i,
|
||||
})
|
||||
}
|
||||
|
||||
content := buildUpstreamModelUpdateTaskNotificationContent(
|
||||
24,
|
||||
12,
|
||||
56,
|
||||
21,
|
||||
9,
|
||||
[]int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12},
|
||||
channelSummaries,
|
||||
[]string{
|
||||
"gpt-4.1", "gpt-4.1-mini", "o3", "o4-mini", "gemini-2.5-pro", "claude-3.7-sonnet",
|
||||
"qwen-max", "deepseek-r1", "llama-3.3-70b", "mistral-large", "command-r-plus", "doubao-pro-32k",
|
||||
"hunyuan-large",
|
||||
},
|
||||
[]string{
|
||||
"gpt-3.5-turbo", "claude-2.1", "gemini-1.5-pro", "mixtral-8x7b", "qwen-plus", "glm-4",
|
||||
"yi-large", "moonshot-v1", "doubao-lite",
|
||||
},
|
||||
)
|
||||
|
||||
require.Contains(t, content, "其余 4 个渠道已省略")
|
||||
require.Contains(t, content, "其余 1 个已省略")
|
||||
require.Contains(t, content, "失败渠道 ID（展示 10/12）")
|
||||
require.Contains(t, content, "其余 2 个已省略")
|
||||
}
|
||||
|
||||
func TestShouldSendUpstreamModelUpdateNotification(t *testing.T) {
|
||||
channelUpstreamModelUpdateNotifyState.Lock()
|
||||
channelUpstreamModelUpdateNotifyState.lastNotifiedAt = 0
|
||||
channelUpstreamModelUpdateNotifyState.lastChangedChannels = 0
|
||||
channelUpstreamModelUpdateNotifyState.lastFailedChannels = 0
|
||||
channelUpstreamModelUpdateNotifyState.Unlock()
|
||||
|
||||
baseTime := int64(2000000)
|
||||
|
||||
require.True(t, shouldSendUpstreamModelUpdateNotification(baseTime, 6, 0))
|
||||
require.False(t, shouldSendUpstreamModelUpdateNotification(baseTime+3600, 6, 0))
|
||||
require.True(t, shouldSendUpstreamModelUpdateNotification(baseTime+3600, 7, 0))
|
||||
require.False(t, shouldSendUpstreamModelUpdateNotification(baseTime+7200, 7, 0))
|
||||
require.True(t, shouldSendUpstreamModelUpdateNotification(baseTime+8000, 0, 3))
|
||||
require.False(t, shouldSendUpstreamModelUpdateNotification(baseTime+9000, 0, 3))
|
||||
require.True(t, shouldSendUpstreamModelUpdateNotification(baseTime+10000, 0, 4))
|
||||
require.True(t, shouldSendUpstreamModelUpdateNotification(baseTime+90000, 7, 0))
|
||||
require.True(t, shouldSendUpstreamModelUpdateNotification(baseTime+90001, 0, 0))
|
||||
}
|
||||
|
||||
func TestDetectAllChannelUpstreamModelUpdatesRejectsExistingActiveTask(t *testing.T) {
|
||||
db := setupModelListControllerTestDB(t)
|
||||
require.NoError(t, db.AutoMigrate(&model.SystemTask{}, &model.SystemTaskLock{}))
|
||||
|
||||
existing, err := model.CreateSystemTask(model.SystemTaskTypeModelUpdate, nil, nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/api/channel/upstream-models/detect-all", nil)
|
||||
|
||||
DetectAllChannelUpstreamModelUpdates(ctx)
|
||||
|
||||
require.Equal(t, http.StatusConflict, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), existing.TaskID)
|
||||
require.Contains(t, recorder.Body.String(), "已有模型更新任务正在运行或等待中")
|
||||
}
|
||||
@@ -0,0 +1,72 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// GetCheckinStatus 获取用户签到状态和历史记录
|
||||
func GetCheckinStatus(c *gin.Context) {
|
||||
setting := operation_setting.GetCheckinSetting()
|
||||
if !setting.Enabled {
|
||||
common.ApiErrorMsg(c, "签到功能未启用")
|
||||
return
|
||||
}
|
||||
userId := c.GetInt("id")
|
||||
// 获取月份参数，默认为当前月份
|
||||
month := c.DefaultQuery("month", time.Now().Format("2006-01"))
|
||||
|
||||
stats, err := model.GetUserCheckinStats(userId, month)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"data": gin.H{
|
||||
"enabled": setting.Enabled,
|
||||
"min_quota": setting.MinQuota,
|
||||
"max_quota": setting.MaxQuota,
|
||||
"stats": stats,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// DoCheckin 执行用户签到
|
||||
func DoCheckin(c *gin.Context) {
|
||||
setting := operation_setting.GetCheckinSetting()
|
||||
if !setting.Enabled {
|
||||
common.ApiErrorMsg(c, "签到功能未启用")
|
||||
return
|
||||
}
|
||||
|
||||
userId := c.GetInt("id")
|
||||
|
||||
checkin, err := model.UserCheckin(userId)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
model.RecordLog(userId, model.LogTypeSystem, fmt.Sprintf("用户签到，获得额度 %s", logger.LogQuota(checkin.QuotaAwarded)))
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "签到成功",
|
||||
"data": gin.H{
|
||||
"quota_awarded": checkin.QuotaAwarded,
|
||||
"checkin_date": checkin.CheckinDate},
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,165 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/relay/channel/codex"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func GetCodexChannelUsage(c *gin.Context) {
|
||||
fetchCodexChannelWhamData(
|
||||
c,
|
||||
service.FetchCodexWhamUsage,
|
||||
"failed to fetch codex usage",
|
||||
"获取用量信息失败，请稍后重试",
|
||||
)
|
||||
}
|
||||
|
||||
func GetCodexChannelRateLimitResetCredits(c *gin.Context) {
|
||||
fetchCodexChannelWhamData(
|
||||
c,
|
||||
service.FetchCodexWhamRateLimitResetCredits,
|
||||
"failed to fetch codex reset credits",
|
||||
"获取重置次数详情失败，请稍后重试",
|
||||
)
|
||||
}
|
||||
|
||||
func ResetCodexChannelUsage(c *gin.Context) {
|
||||
fetchCodexChannelWhamData(
|
||||
c,
|
||||
service.ConsumeCodexWhamRateLimitResetCredit,
|
||||
"failed to reset codex usage",
|
||||
"重置用量失败，请稍后重试",
|
||||
)
|
||||
}
|
||||
|
||||
type codexWhamFetchFunc func(
|
||||
ctx context.Context,
|
||||
client *http.Client,
|
||||
baseURL string,
|
||||
accessToken string,
|
||||
accountID string,
|
||||
) (statusCode int, body []byte, err error)
|
||||
|
||||
func fetchCodexChannelWhamData(
|
||||
c *gin.Context,
|
||||
fetch codexWhamFetchFunc,
|
||||
logPrefix string,
|
||||
userMessage string,
|
||||
) {
|
||||
channelId, err := strconv.Atoi(c.Param("id"))
|
||||
if err != nil {
|
||||
common.ApiError(c, fmt.Errorf("invalid channel id: %w", err))
|
||||
return
|
||||
}
|
||||
|
||||
ch, err := model.GetChannelById(channelId, true)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if ch == nil {
|
||||
c.JSON(http.StatusOK, gin.H{"success": false, "message": "channel not found"})
|
||||
return
|
||||
}
|
||||
if ch.Type != constant.ChannelTypeCodex {
|
||||
c.JSON(http.StatusOK, gin.H{"success": false, "message": "channel type is not Codex"})
|
||||
return
|
||||
}
|
||||
if ch.ChannelInfo.IsMultiKey {
|
||||
c.JSON(http.StatusOK, gin.H{"success": false, "message": "multi-key channel is not supported"})
|
||||
return
|
||||
}
|
||||
|
||||
oauthKey, err := codex.ParseOAuthKey(strings.TrimSpace(ch.Key))
|
||||
if err != nil {
|
||||
common.SysError("failed to parse oauth key: " + err.Error())
|
||||
c.JSON(http.StatusOK, gin.H{"success": false, "message": "解析凭证失败，请检查渠道配置"})
|
||||
return
|
||||
}
|
||||
accessToken := strings.TrimSpace(oauthKey.AccessToken)
|
||||
accountID := strings.TrimSpace(oauthKey.AccountID)
|
||||
if accessToken == "" {
|
||||
c.JSON(http.StatusOK, gin.H{"success": false, "message": "codex channel: access_token is required"})
|
||||
return
|
||||
}
|
||||
if accountID == "" {
|
||||
c.JSON(http.StatusOK, gin.H{"success": false, "message": "codex channel: account_id is required"})
|
||||
return
|
||||
}
|
||||
|
||||
client, err := service.GetHttpClientWithProxy(ch.GetSetting().Proxy)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(c.Request.Context(), 15*time.Second)
|
||||
defer cancel()
|
||||
|
||||
statusCode, body, err := fetch(ctx, client, ch.GetBaseURL(), accessToken, accountID)
|
||||
if err != nil {
|
||||
common.SysError(logPrefix + ": " + err.Error())
|
||||
c.JSON(http.StatusOK, gin.H{"success": false, "message": userMessage})
|
||||
return
|
||||
}
|
||||
|
||||
if (statusCode == http.StatusUnauthorized || statusCode == http.StatusForbidden) && strings.TrimSpace(oauthKey.RefreshToken) != "" {
|
||||
refreshCtx, refreshCancel := context.WithTimeout(c.Request.Context(), 10*time.Second)
|
||||
defer refreshCancel()
|
||||
|
||||
res, refreshErr := service.RefreshCodexOAuthTokenWithProxy(refreshCtx, oauthKey.RefreshToken, ch.GetSetting().Proxy)
|
||||
if refreshErr == nil {
|
||||
oauthKey.AccessToken = res.AccessToken
|
||||
oauthKey.RefreshToken = res.RefreshToken
|
||||
oauthKey.LastRefresh = time.Now().Format(time.RFC3339)
|
||||
oauthKey.Expired = res.ExpiresAt.Format(time.RFC3339)
|
||||
if strings.TrimSpace(oauthKey.Type) == "" {
|
||||
oauthKey.Type = "codex"
|
||||
}
|
||||
|
||||
encoded, encErr := common.Marshal(oauthKey)
|
||||
if encErr == nil {
|
||||
_ = model.DB.Model(&model.Channel{}).Where("id = ?", ch.Id).Update("key", string(encoded)).Error
|
||||
model.InitChannelCache()
|
||||
}
|
||||
|
||||
ctx2, cancel2 := context.WithTimeout(c.Request.Context(), 15*time.Second)
|
||||
defer cancel2()
|
||||
statusCode, body, err = fetch(ctx2, client, ch.GetBaseURL(), oauthKey.AccessToken, accountID)
|
||||
if err != nil {
|
||||
common.SysError(logPrefix + " after refresh: " + err.Error())
|
||||
c.JSON(http.StatusOK, gin.H{"success": false, "message": userMessage})
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
var payload any
|
||||
if common.Unmarshal(body, &payload) != nil {
|
||||
payload = string(body)
|
||||
}
|
||||
|
||||
ok := statusCode >= 200 && statusCode < 300
|
||||
resp := gin.H{
|
||||
"success": ok,
|
||||
"message": "",
|
||||
"upstream_status": statusCode,
|
||||
"data": payload,
|
||||
}
|
||||
if !ok {
|
||||
resp["message"] = fmt.Sprintf("upstream status: %d", statusCode)
|
||||
}
|
||||
c.JSON(http.StatusOK, resp)
|
||||
}
|
||||
@@ -0,0 +1,606 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/oauth"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// CustomOAuthProviderResponse is the response structure for custom OAuth providers
|
||||
// It excludes sensitive fields like client_secret
|
||||
type CustomOAuthProviderResponse struct {
|
||||
Id int `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Slug string `json:"slug"`
|
||||
Icon string `json:"icon"`
|
||||
Enabled bool `json:"enabled"`
|
||||
ClientId string `json:"client_id"`
|
||||
AuthorizationEndpoint string `json:"authorization_endpoint"`
|
||||
TokenEndpoint string `json:"token_endpoint"`
|
||||
UserInfoEndpoint string `json:"user_info_endpoint"`
|
||||
Scopes string `json:"scopes"`
|
||||
UserIdField string `json:"user_id_field"`
|
||||
UsernameField string `json:"username_field"`
|
||||
DisplayNameField string `json:"display_name_field"`
|
||||
EmailField string `json:"email_field"`
|
||||
WellKnown string `json:"well_known"`
|
||||
AuthStyle int `json:"auth_style"`
|
||||
AccessPolicy string `json:"access_policy"`
|
||||
AccessDeniedMessage string `json:"access_denied_message"`
|
||||
}
|
||||
|
||||
type UserOAuthBindingResponse struct {
|
||||
ProviderId int `json:"provider_id"`
|
||||
ProviderName string `json:"provider_name"`
|
||||
ProviderSlug string `json:"provider_slug"`
|
||||
ProviderIcon string `json:"provider_icon"`
|
||||
ProviderUserId string `json:"provider_user_id"`
|
||||
}
|
||||
|
||||
func toCustomOAuthProviderResponse(p *model.CustomOAuthProvider) *CustomOAuthProviderResponse {
|
||||
return &CustomOAuthProviderResponse{
|
||||
Id: p.Id,
|
||||
Name: p.Name,
|
||||
Slug: p.Slug,
|
||||
Icon: p.Icon,
|
||||
Enabled: p.Enabled,
|
||||
ClientId: p.ClientId,
|
||||
AuthorizationEndpoint: p.AuthorizationEndpoint,
|
||||
TokenEndpoint: p.TokenEndpoint,
|
||||
UserInfoEndpoint: p.UserInfoEndpoint,
|
||||
Scopes: p.Scopes,
|
||||
UserIdField: p.UserIdField,
|
||||
UsernameField: p.UsernameField,
|
||||
DisplayNameField: p.DisplayNameField,
|
||||
EmailField: p.EmailField,
|
||||
WellKnown: p.WellKnown,
|
||||
AuthStyle: p.AuthStyle,
|
||||
AccessPolicy: p.AccessPolicy,
|
||||
AccessDeniedMessage: p.AccessDeniedMessage,
|
||||
}
|
||||
}
|
||||
|
||||
// GetCustomOAuthProviders returns all custom OAuth providers
|
||||
func GetCustomOAuthProviders(c *gin.Context) {
|
||||
providers, err := model.GetAllCustomOAuthProviders()
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
response := make([]*CustomOAuthProviderResponse, len(providers))
|
||||
for i, p := range providers {
|
||||
response[i] = toCustomOAuthProviderResponse(p)
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": response,
|
||||
})
|
||||
}
|
||||
|
||||
// GetCustomOAuthProvider returns a single custom OAuth provider by ID
|
||||
func GetCustomOAuthProvider(c *gin.Context) {
|
||||
idStr := c.Param("id")
|
||||
id, err := strconv.Atoi(idStr)
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, "无效的 ID")
|
||||
return
|
||||
}
|
||||
|
||||
provider, err := model.GetCustomOAuthProviderById(id)
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, "未找到该 OAuth 提供商")
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": toCustomOAuthProviderResponse(provider),
|
||||
})
|
||||
}
|
||||
|
||||
// CreateCustomOAuthProviderRequest is the request structure for creating a custom OAuth provider
|
||||
type CreateCustomOAuthProviderRequest struct {
|
||||
Name string `json:"name" binding:"required"`
|
||||
Slug string `json:"slug" binding:"required"`
|
||||
Icon string `json:"icon"`
|
||||
Enabled bool `json:"enabled"`
|
||||
ClientId string `json:"client_id" binding:"required"`
|
||||
ClientSecret string `json:"client_secret" binding:"required"`
|
||||
AuthorizationEndpoint string `json:"authorization_endpoint" binding:"required"`
|
||||
TokenEndpoint string `json:"token_endpoint" binding:"required"`
|
||||
UserInfoEndpoint string `json:"user_info_endpoint" binding:"required"`
|
||||
Scopes string `json:"scopes"`
|
||||
UserIdField string `json:"user_id_field"`
|
||||
UsernameField string `json:"username_field"`
|
||||
DisplayNameField string `json:"display_name_field"`
|
||||
EmailField string `json:"email_field"`
|
||||
WellKnown string `json:"well_known"`
|
||||
AuthStyle int `json:"auth_style"`
|
||||
AccessPolicy string `json:"access_policy"`
|
||||
AccessDeniedMessage string `json:"access_denied_message"`
|
||||
}
|
||||
|
||||
type FetchCustomOAuthDiscoveryRequest struct {
|
||||
WellKnownURL string `json:"well_known_url"`
|
||||
IssuerURL string `json:"issuer_url"`
|
||||
}
|
||||
|
||||
// FetchCustomOAuthDiscovery fetches OIDC discovery document via backend (root-only route)
|
||||
func FetchCustomOAuthDiscovery(c *gin.Context) {
|
||||
var req FetchCustomOAuthDiscoveryRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
common.ApiErrorMsg(c, "无效的请求参数: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
wellKnownURL := strings.TrimSpace(req.WellKnownURL)
|
||||
issuerURL := strings.TrimSpace(req.IssuerURL)
|
||||
|
||||
if wellKnownURL == "" && issuerURL == "" {
|
||||
common.ApiErrorMsg(c, "请先填写 Discovery URL 或 Issuer URL")
|
||||
return
|
||||
}
|
||||
|
||||
targetURL := wellKnownURL
|
||||
if targetURL == "" {
|
||||
targetURL = strings.TrimRight(issuerURL, "/") + "/.well-known/openid-configuration"
|
||||
}
|
||||
targetURL = strings.TrimSpace(targetURL)
|
||||
|
||||
parsedURL, err := url.Parse(targetURL)
|
||||
if err != nil || parsedURL.Host == "" || (parsedURL.Scheme != "http" && parsedURL.Scheme != "https") {
|
||||
common.ApiErrorMsg(c, "Discovery URL 无效，仅支持 http/https")
|
||||
return
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(c.Request.Context(), 20*time.Second)
|
||||
defer cancel()
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, targetURL, nil)
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, "创建 Discovery 请求失败: "+err.Error())
|
||||
return
|
||||
}
|
||||
httpReq.Header.Set("Accept", "application/json")
|
||||
|
||||
client := &http.Client{Timeout: 20 * time.Second}
|
||||
resp, err := client.Do(httpReq)
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, "获取 Discovery 配置失败: "+err.Error())
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
|
||||
message := strings.TrimSpace(string(body))
|
||||
if message == "" {
|
||||
message = resp.Status
|
||||
}
|
||||
common.ApiErrorMsg(c, "获取 Discovery 配置失败: "+message)
|
||||
return
|
||||
}
|
||||
|
||||
var discovery map[string]any
|
||||
if err = common.DecodeJson(resp.Body, &discovery); err != nil {
|
||||
common.ApiErrorMsg(c, "解析 Discovery 配置失败: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": gin.H{
|
||||
"well_known_url": targetURL,
|
||||
"discovery": discovery,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// CreateCustomOAuthProvider creates a new custom OAuth provider
|
||||
func CreateCustomOAuthProvider(c *gin.Context) {
|
||||
var req CreateCustomOAuthProviderRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
common.ApiErrorMsg(c, "无效的请求参数: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
// Check if slug is already taken
|
||||
if model.IsSlugTaken(req.Slug, 0) {
|
||||
common.ApiErrorMsg(c, "该 Slug 已被使用")
|
||||
return
|
||||
}
|
||||
|
||||
// Check if slug conflicts with built-in providers
|
||||
if oauth.IsProviderRegistered(req.Slug) && !oauth.IsCustomProvider(req.Slug) {
|
||||
common.ApiErrorMsg(c, "该 Slug 与内置 OAuth 提供商冲突")
|
||||
return
|
||||
}
|
||||
|
||||
provider := &model.CustomOAuthProvider{
|
||||
Name: req.Name,
|
||||
Slug: req.Slug,
|
||||
Icon: req.Icon,
|
||||
Enabled: req.Enabled,
|
||||
ClientId: req.ClientId,
|
||||
ClientSecret: req.ClientSecret,
|
||||
AuthorizationEndpoint: req.AuthorizationEndpoint,
|
||||
TokenEndpoint: req.TokenEndpoint,
|
||||
UserInfoEndpoint: req.UserInfoEndpoint,
|
||||
Scopes: req.Scopes,
|
||||
UserIdField: req.UserIdField,
|
||||
UsernameField: req.UsernameField,
|
||||
DisplayNameField: req.DisplayNameField,
|
||||
EmailField: req.EmailField,
|
||||
WellKnown: req.WellKnown,
|
||||
AuthStyle: req.AuthStyle,
|
||||
AccessPolicy: req.AccessPolicy,
|
||||
AccessDeniedMessage: req.AccessDeniedMessage,
|
||||
}
|
||||
|
||||
if err := model.CreateCustomOAuthProvider(provider); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
// Register the provider in the OAuth registry
|
||||
oauth.RegisterOrUpdateCustomProvider(provider)
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "创建成功",
|
||||
"data": toCustomOAuthProviderResponse(provider),
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateCustomOAuthProviderRequest is the request structure for updating a custom OAuth provider
|
||||
type UpdateCustomOAuthProviderRequest struct {
|
||||
Name string `json:"name"`
|
||||
Slug string `json:"slug"`
|
||||
Icon *string `json:"icon"` // Optional: if nil, keep existing
|
||||
Enabled *bool `json:"enabled"` // Optional: if nil, keep existing
|
||||
ClientId string `json:"client_id"`
|
||||
ClientSecret string `json:"client_secret"` // Optional: if empty, keep existing
|
||||
AuthorizationEndpoint string `json:"authorization_endpoint"`
|
||||
TokenEndpoint string `json:"token_endpoint"`
|
||||
UserInfoEndpoint string `json:"user_info_endpoint"`
|
||||
Scopes string `json:"scopes"`
|
||||
UserIdField string `json:"user_id_field"`
|
||||
UsernameField string `json:"username_field"`
|
||||
DisplayNameField string `json:"display_name_field"`
|
||||
EmailField string `json:"email_field"`
|
||||
WellKnown *string `json:"well_known"` // Optional: if nil, keep existing
|
||||
AuthStyle *int `json:"auth_style"` // Optional: if nil, keep existing
|
||||
AccessPolicy *string `json:"access_policy"` // Optional: if nil, keep existing
|
||||
AccessDeniedMessage *string `json:"access_denied_message"` // Optional: if nil, keep existing
|
||||
}
|
||||
|
||||
// UpdateCustomOAuthProvider updates an existing custom OAuth provider
|
||||
func UpdateCustomOAuthProvider(c *gin.Context) {
|
||||
idStr := c.Param("id")
|
||||
id, err := strconv.Atoi(idStr)
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, "无效的 ID")
|
||||
return
|
||||
}
|
||||
|
||||
var req UpdateCustomOAuthProviderRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
common.ApiErrorMsg(c, "无效的请求参数: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
// Get existing provider
|
||||
provider, err := model.GetCustomOAuthProviderById(id)
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, "未找到该 OAuth 提供商")
|
||||
return
|
||||
}
|
||||
|
||||
oldSlug := provider.Slug
|
||||
|
||||
// Check if new slug is taken by another provider
|
||||
if req.Slug != "" && req.Slug != provider.Slug {
|
||||
if model.IsSlugTaken(req.Slug, id) {
|
||||
common.ApiErrorMsg(c, "该 Slug 已被使用")
|
||||
return
|
||||
}
|
||||
// Check if slug conflicts with built-in providers
|
||||
if oauth.IsProviderRegistered(req.Slug) && !oauth.IsCustomProvider(req.Slug) {
|
||||
common.ApiErrorMsg(c, "该 Slug 与内置 OAuth 提供商冲突")
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Update fields
|
||||
if req.Name != "" {
|
||||
provider.Name = req.Name
|
||||
}
|
||||
if req.Slug != "" {
|
||||
provider.Slug = req.Slug
|
||||
}
|
||||
if req.Icon != nil {
|
||||
provider.Icon = *req.Icon
|
||||
}
|
||||
if req.Enabled != nil {
|
||||
provider.Enabled = *req.Enabled
|
||||
}
|
||||
if req.ClientId != "" {
|
||||
provider.ClientId = req.ClientId
|
||||
}
|
||||
if req.ClientSecret != "" {
|
||||
provider.ClientSecret = req.ClientSecret
|
||||
}
|
||||
if req.AuthorizationEndpoint != "" {
|
||||
provider.AuthorizationEndpoint = req.AuthorizationEndpoint
|
||||
}
|
||||
if req.TokenEndpoint != "" {
|
||||
provider.TokenEndpoint = req.TokenEndpoint
|
||||
}
|
||||
if req.UserInfoEndpoint != "" {
|
||||
provider.UserInfoEndpoint = req.UserInfoEndpoint
|
||||
}
|
||||
if req.Scopes != "" {
|
||||
provider.Scopes = req.Scopes
|
||||
}
|
||||
if req.UserIdField != "" {
|
||||
provider.UserIdField = req.UserIdField
|
||||
}
|
||||
if req.UsernameField != "" {
|
||||
provider.UsernameField = req.UsernameField
|
||||
}
|
||||
if req.DisplayNameField != "" {
|
||||
provider.DisplayNameField = req.DisplayNameField
|
||||
}
|
||||
if req.EmailField != "" {
|
||||
provider.EmailField = req.EmailField
|
||||
}
|
||||
if req.WellKnown != nil {
|
||||
provider.WellKnown = *req.WellKnown
|
||||
}
|
||||
if req.AuthStyle != nil {
|
||||
provider.AuthStyle = *req.AuthStyle
|
||||
}
|
||||
if req.AccessPolicy != nil {
|
||||
provider.AccessPolicy = *req.AccessPolicy
|
||||
}
|
||||
if req.AccessDeniedMessage != nil {
|
||||
provider.AccessDeniedMessage = *req.AccessDeniedMessage
|
||||
}
|
||||
|
||||
if err := model.UpdateCustomOAuthProvider(provider); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
// Update the provider in the OAuth registry
|
||||
if oldSlug != provider.Slug {
|
||||
oauth.UnregisterCustomProvider(oldSlug)
|
||||
}
|
||||
oauth.RegisterOrUpdateCustomProvider(provider)
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "更新成功",
|
||||
"data": toCustomOAuthProviderResponse(provider),
|
||||
})
|
||||
}
|
||||
|
||||
// DeleteCustomOAuthProvider deletes a custom OAuth provider
|
||||
func DeleteCustomOAuthProvider(c *gin.Context) {
|
||||
idStr := c.Param("id")
|
||||
id, err := strconv.Atoi(idStr)
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, "无效的 ID")
|
||||
return
|
||||
}
|
||||
|
||||
// Get existing provider to get slug
|
||||
provider, err := model.GetCustomOAuthProviderById(id)
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, "未找到该 OAuth 提供商")
|
||||
return
|
||||
}
|
||||
|
||||
// Check if there are any user bindings
|
||||
count, err := model.GetBindingCountByProviderId(id)
|
||||
if err != nil {
|
||||
common.SysError("Failed to get binding count for provider " + strconv.Itoa(id) + ": " + err.Error())
|
||||
common.ApiErrorMsg(c, "检查用户绑定时发生错误，请稍后重试")
|
||||
return
|
||||
}
|
||||
if count > 0 {
|
||||
common.ApiErrorMsg(c, "该 OAuth 提供商还有用户绑定，无法删除。请先解除所有用户绑定。")
|
||||
return
|
||||
}
|
||||
|
||||
if err := model.DeleteCustomOAuthProvider(id); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
// Unregister the provider from the OAuth registry
|
||||
oauth.UnregisterCustomProvider(provider.Slug)
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "删除成功",
|
||||
})
|
||||
}
|
||||
|
||||
func buildUserOAuthBindingsResponse(userId int) ([]UserOAuthBindingResponse, error) {
|
||||
bindings, err := model.GetUserOAuthBindingsByUserId(userId)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
response := make([]UserOAuthBindingResponse, 0, len(bindings))
|
||||
for _, binding := range bindings {
|
||||
provider, err := model.GetCustomOAuthProviderById(binding.ProviderId)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
response = append(response, UserOAuthBindingResponse{
|
||||
ProviderId: binding.ProviderId,
|
||||
ProviderName: provider.Name,
|
||||
ProviderSlug: provider.Slug,
|
||||
ProviderIcon: provider.Icon,
|
||||
ProviderUserId: binding.ProviderUserId,
|
||||
})
|
||||
}
|
||||
|
||||
return response, nil
|
||||
}
|
||||
|
||||
// GetUserOAuthBindings returns all OAuth bindings for the current user
|
||||
func GetUserOAuthBindings(c *gin.Context) {
|
||||
userId := c.GetInt("id")
|
||||
if userId == 0 {
|
||||
common.ApiErrorMsg(c, "未登录")
|
||||
return
|
||||
}
|
||||
|
||||
response, err := buildUserOAuthBindingsResponse(userId)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": response,
|
||||
})
|
||||
}
|
||||
|
||||
func GetUserOAuthBindingsByAdmin(c *gin.Context) {
|
||||
userIdStr := c.Param("id")
|
||||
userId, err := strconv.Atoi(userIdStr)
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, "invalid user id")
|
||||
return
|
||||
}
|
||||
|
||||
targetUser, err := model.GetUserById(userId, false)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
myRole := c.GetInt("role")
|
||||
if !canManageTargetRole(myRole, targetUser.Role) {
|
||||
common.ApiErrorMsg(c, "no permission")
|
||||
return
|
||||
}
|
||||
|
||||
response, err := buildUserOAuthBindingsResponse(userId)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": response,
|
||||
})
|
||||
}
|
||||
|
||||
// UnbindCustomOAuth unbinds a custom OAuth provider from the current user
|
||||
func UnbindCustomOAuth(c *gin.Context) {
|
||||
identity, ok := middleware.GetSessionAuthIdentity(c)
|
||||
if !ok {
|
||||
writeSecurityOperationError(c, service.ErrAuthTokenInvalid)
|
||||
return
|
||||
}
|
||||
|
||||
providerIdStr := c.Param("provider_id")
|
||||
providerId, err := strconv.Atoi(providerIdStr)
|
||||
if err != nil || providerId <= 0 {
|
||||
common.ApiErrorMsg(c, "无效的提供商 ID")
|
||||
return
|
||||
}
|
||||
|
||||
succeeded, notificationFailed := false, false
|
||||
defer func() {
|
||||
recordUserSecurityAudit(c, identity.UserID, "user.binding_unbind", map[string]any{"provider_id": providerId, "success": succeeded, "notification_failed": notificationFailed})
|
||||
}()
|
||||
context, err := common.Marshal(service.AccountUnbindingContext{ProviderID: providerId})
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
if middleware.RequireSecurityProof(c, service.VerificationOperation{Scope: service.VerificationScopeAccountUnbind, Context: context}) == nil {
|
||||
return
|
||||
}
|
||||
if err := service.UnbindAccountOAuth(identity, providerId); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
succeeded = true
|
||||
user, err := model.GetUserById(identity.UserID, false)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
notificationFailed = service.NotifyAccountSecurityChange(user.Email, "Login account unlinked") != nil
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "解绑成功",
|
||||
"data": gin.H{"notification_warning": notificationFailed},
|
||||
})
|
||||
}
|
||||
|
||||
func UnbindCustomOAuthByAdmin(c *gin.Context) {
|
||||
userIdStr := c.Param("id")
|
||||
userId, err := strconv.Atoi(userIdStr)
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, "invalid user id")
|
||||
return
|
||||
}
|
||||
|
||||
targetUser, err := model.GetUserById(userId, false)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
myRole := c.GetInt("role")
|
||||
if !canManageTargetRole(myRole, targetUser.Role) {
|
||||
common.ApiErrorMsg(c, "no permission")
|
||||
return
|
||||
}
|
||||
|
||||
providerIdStr := c.Param("provider_id")
|
||||
providerId, err := strconv.Atoi(providerIdStr)
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, "invalid provider id")
|
||||
return
|
||||
}
|
||||
|
||||
if err := model.DeleteUserOAuthBinding(userId, providerId); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "success",
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,807 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/pkg/ionet"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func getIoAPIKey(c *gin.Context) (string, bool) {
|
||||
common.OptionMapRWMutex.RLock()
|
||||
enabled := common.OptionMap["model_deployment.ionet.enabled"] == "true"
|
||||
apiKey := common.OptionMap["model_deployment.ionet.api_key"]
|
||||
common.OptionMapRWMutex.RUnlock()
|
||||
if !enabled || strings.TrimSpace(apiKey) == "" {
|
||||
common.ApiErrorMsg(c, "io.net model deployment is not enabled or api key missing")
|
||||
return "", false
|
||||
}
|
||||
return apiKey, true
|
||||
}
|
||||
|
||||
func GetModelDeploymentSettings(c *gin.Context) {
|
||||
common.OptionMapRWMutex.RLock()
|
||||
enabled := common.OptionMap["model_deployment.ionet.enabled"] == "true"
|
||||
hasAPIKey := strings.TrimSpace(common.OptionMap["model_deployment.ionet.api_key"]) != ""
|
||||
common.OptionMapRWMutex.RUnlock()
|
||||
|
||||
common.ApiSuccess(c, gin.H{
|
||||
"provider": "io.net",
|
||||
"enabled": enabled,
|
||||
"configured": hasAPIKey,
|
||||
"can_connect": enabled && hasAPIKey,
|
||||
})
|
||||
}
|
||||
|
||||
func getIoClient(c *gin.Context) (*ionet.Client, bool) {
|
||||
apiKey, ok := getIoAPIKey(c)
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
return ionet.NewClient(apiKey), true
|
||||
}
|
||||
|
||||
func getIoEnterpriseClient(c *gin.Context) (*ionet.Client, bool) {
|
||||
apiKey, ok := getIoAPIKey(c)
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
return ionet.NewEnterpriseClient(apiKey), true
|
||||
}
|
||||
|
||||
func TestIoNetConnection(c *gin.Context) {
|
||||
var req struct {
|
||||
APIKey string `json:"api_key"`
|
||||
}
|
||||
|
||||
rawBody, err := c.GetRawData()
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if len(bytes.TrimSpace(rawBody)) > 0 {
|
||||
if err := json.Unmarshal(rawBody, &req); err != nil {
|
||||
common.ApiErrorMsg(c, "invalid request payload")
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
apiKey := strings.TrimSpace(req.APIKey)
|
||||
if apiKey == "" {
|
||||
common.OptionMapRWMutex.RLock()
|
||||
storedKey := strings.TrimSpace(common.OptionMap["model_deployment.ionet.api_key"])
|
||||
common.OptionMapRWMutex.RUnlock()
|
||||
if storedKey == "" {
|
||||
common.ApiErrorMsg(c, "api_key is required")
|
||||
return
|
||||
}
|
||||
apiKey = storedKey
|
||||
}
|
||||
|
||||
client := ionet.NewEnterpriseClient(apiKey)
|
||||
result, err := client.GetMaxGPUsPerContainer()
|
||||
if err != nil {
|
||||
if apiErr, ok := err.(*ionet.APIError); ok {
|
||||
message := strings.TrimSpace(apiErr.Message)
|
||||
if message == "" {
|
||||
message = "failed to validate api key"
|
||||
}
|
||||
common.ApiErrorMsg(c, message)
|
||||
return
|
||||
}
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
totalHardware := 0
|
||||
totalAvailable := 0
|
||||
if result != nil {
|
||||
totalHardware = len(result.Hardware)
|
||||
totalAvailable = result.Total
|
||||
if totalAvailable == 0 {
|
||||
for _, hw := range result.Hardware {
|
||||
totalAvailable += hw.Available
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
common.ApiSuccess(c, gin.H{
|
||||
"hardware_count": totalHardware,
|
||||
"total_available": totalAvailable,
|
||||
})
|
||||
}
|
||||
|
||||
func requireDeploymentID(c *gin.Context) (string, bool) {
|
||||
deploymentID := strings.TrimSpace(c.Param("id"))
|
||||
if deploymentID == "" {
|
||||
common.ApiErrorMsg(c, "deployment ID is required")
|
||||
return "", false
|
||||
}
|
||||
return deploymentID, true
|
||||
}
|
||||
|
||||
func requireContainerID(c *gin.Context) (string, bool) {
|
||||
containerID := strings.TrimSpace(c.Param("container_id"))
|
||||
if containerID == "" {
|
||||
common.ApiErrorMsg(c, "container ID is required")
|
||||
return "", false
|
||||
}
|
||||
return containerID, true
|
||||
}
|
||||
|
||||
func mapIoNetDeployment(d ionet.Deployment) map[string]any {
|
||||
var created int64
|
||||
if d.CreatedAt.IsZero() {
|
||||
created = time.Now().Unix()
|
||||
} else {
|
||||
created = d.CreatedAt.Unix()
|
||||
}
|
||||
|
||||
timeRemainingHours := d.ComputeMinutesRemaining / 60
|
||||
timeRemainingMins := d.ComputeMinutesRemaining % 60
|
||||
var timeRemaining string
|
||||
if timeRemainingHours > 0 {
|
||||
timeRemaining = fmt.Sprintf("%d hour %d minutes", timeRemainingHours, timeRemainingMins)
|
||||
} else if timeRemainingMins > 0 {
|
||||
timeRemaining = fmt.Sprintf("%d minutes", timeRemainingMins)
|
||||
} else {
|
||||
timeRemaining = "completed"
|
||||
}
|
||||
|
||||
hardwareInfo := fmt.Sprintf("%s %s x%d", d.BrandName, d.HardwareName, d.HardwareQuantity)
|
||||
|
||||
return map[string]any{
|
||||
"id": d.ID,
|
||||
"deployment_name": d.Name,
|
||||
"container_name": d.Name,
|
||||
"status": strings.ToLower(d.Status),
|
||||
"type": "Container",
|
||||
"time_remaining": timeRemaining,
|
||||
"time_remaining_minutes": d.ComputeMinutesRemaining,
|
||||
"hardware_info": hardwareInfo,
|
||||
"hardware_name": d.HardwareName,
|
||||
"brand_name": d.BrandName,
|
||||
"hardware_quantity": d.HardwareQuantity,
|
||||
"completed_percent": d.CompletedPercent,
|
||||
"compute_minutes_served": d.ComputeMinutesServed,
|
||||
"compute_minutes_remaining": d.ComputeMinutesRemaining,
|
||||
"created_at": created,
|
||||
"updated_at": created,
|
||||
"model_name": "",
|
||||
"model_version": "",
|
||||
"instance_count": d.HardwareQuantity,
|
||||
"resource_config": map[string]any{
|
||||
"cpu": "",
|
||||
"memory": "",
|
||||
"gpu": strconv.Itoa(d.HardwareQuantity),
|
||||
},
|
||||
"description": "",
|
||||
"provider": "io.net",
|
||||
}
|
||||
}
|
||||
|
||||
func computeStatusCounts(total int, deployments []ionet.Deployment) map[string]int64 {
|
||||
counts := map[string]int64{
|
||||
"all": int64(total),
|
||||
}
|
||||
|
||||
for _, status := range []string{"running", "completed", "failed", "deployment requested", "termination requested", "destroyed"} {
|
||||
counts[status] = 0
|
||||
}
|
||||
|
||||
for _, d := range deployments {
|
||||
status := strings.ToLower(strings.TrimSpace(d.Status))
|
||||
counts[status] = counts[status] + 1
|
||||
}
|
||||
|
||||
return counts
|
||||
}
|
||||
|
||||
func GetAllDeployments(c *gin.Context) {
|
||||
pageInfo := common.GetPageQuery(c)
|
||||
client, ok := getIoEnterpriseClient(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
status := c.Query("status")
|
||||
opts := &ionet.ListDeploymentsOptions{
|
||||
Status: strings.ToLower(strings.TrimSpace(status)),
|
||||
Page: pageInfo.GetPage(),
|
||||
PageSize: pageInfo.GetPageSize(),
|
||||
SortBy: "created_at",
|
||||
SortOrder: "desc",
|
||||
}
|
||||
|
||||
dl, err := client.ListDeployments(opts)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
items := make([]map[string]any, 0, len(dl.Deployments))
|
||||
for _, d := range dl.Deployments {
|
||||
items = append(items, mapIoNetDeployment(d))
|
||||
}
|
||||
|
||||
data := gin.H{
|
||||
"page": pageInfo.GetPage(),
|
||||
"page_size": pageInfo.GetPageSize(),
|
||||
"total": dl.Total,
|
||||
"items": items,
|
||||
"status_counts": computeStatusCounts(dl.Total, dl.Deployments),
|
||||
}
|
||||
common.ApiSuccess(c, data)
|
||||
}
|
||||
|
||||
func SearchDeployments(c *gin.Context) {
|
||||
pageInfo := common.GetPageQuery(c)
|
||||
client, ok := getIoEnterpriseClient(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
status := strings.ToLower(strings.TrimSpace(c.Query("status")))
|
||||
keyword := strings.TrimSpace(c.Query("keyword"))
|
||||
|
||||
dl, err := client.ListDeployments(&ionet.ListDeploymentsOptions{
|
||||
Status: status,
|
||||
Page: pageInfo.GetPage(),
|
||||
PageSize: pageInfo.GetPageSize(),
|
||||
SortBy: "created_at",
|
||||
SortOrder: "desc",
|
||||
})
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
filtered := make([]ionet.Deployment, 0, len(dl.Deployments))
|
||||
if keyword == "" {
|
||||
filtered = dl.Deployments
|
||||
} else {
|
||||
kw := strings.ToLower(keyword)
|
||||
for _, d := range dl.Deployments {
|
||||
if strings.Contains(strings.ToLower(d.Name), kw) {
|
||||
filtered = append(filtered, d)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
items := make([]map[string]any, 0, len(filtered))
|
||||
for _, d := range filtered {
|
||||
items = append(items, mapIoNetDeployment(d))
|
||||
}
|
||||
|
||||
total := dl.Total
|
||||
if keyword != "" {
|
||||
total = len(filtered)
|
||||
}
|
||||
|
||||
data := gin.H{
|
||||
"page": pageInfo.GetPage(),
|
||||
"page_size": pageInfo.GetPageSize(),
|
||||
"total": total,
|
||||
"items": items,
|
||||
}
|
||||
common.ApiSuccess(c, data)
|
||||
}
|
||||
|
||||
func GetDeployment(c *gin.Context) {
|
||||
client, ok := getIoEnterpriseClient(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
deploymentID, ok := requireDeploymentID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
details, err := client.GetDeployment(deploymentID)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
data := map[string]any{
|
||||
"id": details.ID,
|
||||
"deployment_name": details.ID,
|
||||
"model_name": "",
|
||||
"model_version": "",
|
||||
"status": strings.ToLower(details.Status),
|
||||
"instance_count": details.TotalContainers,
|
||||
"hardware_id": details.HardwareID,
|
||||
"resource_config": map[string]any{
|
||||
"cpu": "",
|
||||
"memory": "",
|
||||
"gpu": strconv.Itoa(details.TotalGPUs),
|
||||
},
|
||||
"created_at": details.CreatedAt.Unix(),
|
||||
"updated_at": details.CreatedAt.Unix(),
|
||||
"description": "",
|
||||
"amount_paid": details.AmountPaid,
|
||||
"completed_percent": details.CompletedPercent,
|
||||
"gpus_per_container": details.GPUsPerContainer,
|
||||
"total_gpus": details.TotalGPUs,
|
||||
"total_containers": details.TotalContainers,
|
||||
"hardware_name": details.HardwareName,
|
||||
"brand_name": details.BrandName,
|
||||
"compute_minutes_served": details.ComputeMinutesServed,
|
||||
"compute_minutes_remaining": details.ComputeMinutesRemaining,
|
||||
"locations": details.Locations,
|
||||
"container_config": details.ContainerConfig,
|
||||
}
|
||||
|
||||
common.ApiSuccess(c, data)
|
||||
}
|
||||
|
||||
func UpdateDeploymentName(c *gin.Context) {
|
||||
client, ok := getIoEnterpriseClient(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
deploymentID, ok := requireDeploymentID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
var req struct {
|
||||
Name string `json:"name" binding:"required"`
|
||||
}
|
||||
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
updateReq := &ionet.UpdateClusterNameRequest{
|
||||
Name: strings.TrimSpace(req.Name),
|
||||
}
|
||||
|
||||
if updateReq.Name == "" {
|
||||
common.ApiErrorMsg(c, "deployment name cannot be empty")
|
||||
return
|
||||
}
|
||||
|
||||
available, err := client.CheckClusterNameAvailability(updateReq.Name)
|
||||
if err != nil {
|
||||
common.ApiError(c, fmt.Errorf("failed to check name availability: %w", err))
|
||||
return
|
||||
}
|
||||
|
||||
if !available {
|
||||
common.ApiErrorMsg(c, "deployment name is not available, please choose a different name")
|
||||
return
|
||||
}
|
||||
|
||||
resp, err := client.UpdateClusterName(deploymentID, updateReq)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
data := gin.H{
|
||||
"status": resp.Status,
|
||||
"message": resp.Message,
|
||||
"id": deploymentID,
|
||||
"name": updateReq.Name,
|
||||
}
|
||||
common.ApiSuccess(c, data)
|
||||
}
|
||||
|
||||
func UpdateDeployment(c *gin.Context) {
|
||||
client, ok := getIoEnterpriseClient(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
deploymentID, ok := requireDeploymentID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
var req ionet.UpdateDeploymentRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
resp, err := client.UpdateDeployment(deploymentID, &req)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
data := gin.H{
|
||||
"status": resp.Status,
|
||||
"deployment_id": resp.DeploymentID,
|
||||
}
|
||||
common.ApiSuccess(c, data)
|
||||
}
|
||||
|
||||
func ExtendDeployment(c *gin.Context) {
|
||||
client, ok := getIoEnterpriseClient(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
deploymentID, ok := requireDeploymentID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
var req ionet.ExtendDurationRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
details, err := client.ExtendDeployment(deploymentID, &req)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
data := mapIoNetDeployment(ionet.Deployment{
|
||||
ID: details.ID,
|
||||
Status: details.Status,
|
||||
Name: deploymentID,
|
||||
CompletedPercent: float64(details.CompletedPercent),
|
||||
HardwareQuantity: details.TotalGPUs,
|
||||
BrandName: details.BrandName,
|
||||
HardwareName: details.HardwareName,
|
||||
ComputeMinutesServed: details.ComputeMinutesServed,
|
||||
ComputeMinutesRemaining: details.ComputeMinutesRemaining,
|
||||
CreatedAt: details.CreatedAt,
|
||||
})
|
||||
|
||||
common.ApiSuccess(c, data)
|
||||
}
|
||||
|
||||
func DeleteDeployment(c *gin.Context) {
|
||||
client, ok := getIoEnterpriseClient(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
deploymentID, ok := requireDeploymentID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
resp, err := client.DeleteDeployment(deploymentID)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
data := gin.H{
|
||||
"status": resp.Status,
|
||||
"deployment_id": resp.DeploymentID,
|
||||
"message": "Deployment termination requested successfully",
|
||||
}
|
||||
common.ApiSuccess(c, data)
|
||||
}
|
||||
|
||||
func CreateDeployment(c *gin.Context) {
|
||||
client, ok := getIoEnterpriseClient(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
var req ionet.DeploymentRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
resp, err := client.DeployContainer(&req)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
data := gin.H{
|
||||
"deployment_id": resp.DeploymentID,
|
||||
"status": resp.Status,
|
||||
"message": "Deployment created successfully",
|
||||
}
|
||||
common.ApiSuccess(c, data)
|
||||
}
|
||||
|
||||
func GetHardwareTypes(c *gin.Context) {
|
||||
client, ok := getIoEnterpriseClient(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
hardwareTypes, totalAvailable, err := client.ListHardwareTypes()
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
data := gin.H{
|
||||
"hardware_types": hardwareTypes,
|
||||
"total": len(hardwareTypes),
|
||||
"total_available": totalAvailable,
|
||||
}
|
||||
common.ApiSuccess(c, data)
|
||||
}
|
||||
|
||||
func GetLocations(c *gin.Context) {
|
||||
client, ok := getIoClient(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
locationsResp, err := client.ListLocations()
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
total := locationsResp.Total
|
||||
if total == 0 {
|
||||
total = len(locationsResp.Locations)
|
||||
}
|
||||
|
||||
data := gin.H{
|
||||
"locations": locationsResp.Locations,
|
||||
"total": total,
|
||||
}
|
||||
common.ApiSuccess(c, data)
|
||||
}
|
||||
|
||||
func GetAvailableReplicas(c *gin.Context) {
|
||||
client, ok := getIoEnterpriseClient(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
hardwareIDStr := c.Query("hardware_id")
|
||||
gpuCountStr := c.Query("gpu_count")
|
||||
|
||||
if hardwareIDStr == "" {
|
||||
common.ApiErrorMsg(c, "hardware_id parameter is required")
|
||||
return
|
||||
}
|
||||
|
||||
hardwareID, err := strconv.Atoi(hardwareIDStr)
|
||||
if err != nil || hardwareID <= 0 {
|
||||
common.ApiErrorMsg(c, "invalid hardware_id parameter")
|
||||
return
|
||||
}
|
||||
|
||||
gpuCount := 1
|
||||
if gpuCountStr != "" {
|
||||
if parsed, err := strconv.Atoi(gpuCountStr); err == nil && parsed > 0 {
|
||||
gpuCount = parsed
|
||||
}
|
||||
}
|
||||
|
||||
replicas, err := client.GetAvailableReplicas(hardwareID, gpuCount)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
common.ApiSuccess(c, replicas)
|
||||
}
|
||||
|
||||
func GetPriceEstimation(c *gin.Context) {
|
||||
client, ok := getIoEnterpriseClient(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
var req ionet.PriceEstimationRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
priceResp, err := client.GetPriceEstimation(&req)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
common.ApiSuccess(c, priceResp)
|
||||
}
|
||||
|
||||
func CheckClusterNameAvailability(c *gin.Context) {
|
||||
client, ok := getIoEnterpriseClient(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
clusterName := strings.TrimSpace(c.Query("name"))
|
||||
if clusterName == "" {
|
||||
common.ApiErrorMsg(c, "name parameter is required")
|
||||
return
|
||||
}
|
||||
|
||||
available, err := client.CheckClusterNameAvailability(clusterName)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
data := gin.H{
|
||||
"available": available,
|
||||
"name": clusterName,
|
||||
}
|
||||
common.ApiSuccess(c, data)
|
||||
}
|
||||
|
||||
func GetDeploymentLogs(c *gin.Context) {
|
||||
client, ok := getIoClient(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
deploymentID, ok := requireDeploymentID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
containerID := c.Query("container_id")
|
||||
if containerID == "" {
|
||||
common.ApiErrorMsg(c, "container_id parameter is required")
|
||||
return
|
||||
}
|
||||
level := c.Query("level")
|
||||
stream := c.Query("stream")
|
||||
cursor := c.Query("cursor")
|
||||
limitStr := c.Query("limit")
|
||||
follow := c.Query("follow") == "true"
|
||||
|
||||
var limit int = 100
|
||||
if limitStr != "" {
|
||||
if parsedLimit, err := strconv.Atoi(limitStr); err == nil && parsedLimit > 0 {
|
||||
limit = min(parsedLimit, 1000)
|
||||
}
|
||||
}
|
||||
|
||||
opts := &ionet.GetLogsOptions{
|
||||
Level: level,
|
||||
Stream: stream,
|
||||
Limit: limit,
|
||||
Cursor: cursor,
|
||||
Follow: follow,
|
||||
}
|
||||
|
||||
if startTime := c.Query("start_time"); startTime != "" {
|
||||
if t, err := time.Parse(time.RFC3339, startTime); err == nil {
|
||||
opts.StartTime = &t
|
||||
}
|
||||
}
|
||||
if endTime := c.Query("end_time"); endTime != "" {
|
||||
if t, err := time.Parse(time.RFC3339, endTime); err == nil {
|
||||
opts.EndTime = &t
|
||||
}
|
||||
}
|
||||
|
||||
rawLogs, err := client.GetContainerLogsRaw(deploymentID, containerID, opts)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
common.ApiSuccess(c, rawLogs)
|
||||
}
|
||||
|
||||
func ListDeploymentContainers(c *gin.Context) {
|
||||
client, ok := getIoEnterpriseClient(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
deploymentID, ok := requireDeploymentID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
containers, err := client.ListContainers(deploymentID)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
items := make([]map[string]any, 0)
|
||||
if containers != nil {
|
||||
items = make([]map[string]any, 0, len(containers.Workers))
|
||||
for _, ctr := range containers.Workers {
|
||||
events := make([]map[string]any, 0, len(ctr.ContainerEvents))
|
||||
for _, event := range ctr.ContainerEvents {
|
||||
events = append(events, map[string]any{
|
||||
"time": event.Time.Unix(),
|
||||
"message": event.Message,
|
||||
})
|
||||
}
|
||||
|
||||
items = append(items, map[string]any{
|
||||
"container_id": ctr.ContainerID,
|
||||
"device_id": ctr.DeviceID,
|
||||
"status": strings.ToLower(strings.TrimSpace(ctr.Status)),
|
||||
"hardware": ctr.Hardware,
|
||||
"brand_name": ctr.BrandName,
|
||||
"created_at": ctr.CreatedAt.Unix(),
|
||||
"uptime_percent": ctr.UptimePercent,
|
||||
"gpus_per_container": ctr.GPUsPerContainer,
|
||||
"public_url": ctr.PublicURL,
|
||||
"events": events,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
response := gin.H{
|
||||
"total": 0,
|
||||
"containers": items,
|
||||
}
|
||||
if containers != nil {
|
||||
response["total"] = containers.Total
|
||||
}
|
||||
|
||||
common.ApiSuccess(c, response)
|
||||
}
|
||||
|
||||
func GetContainerDetails(c *gin.Context) {
|
||||
client, ok := getIoEnterpriseClient(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
deploymentID, ok := requireDeploymentID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
containerID, ok := requireContainerID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
details, err := client.GetContainerDetails(deploymentID, containerID)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if details == nil {
|
||||
common.ApiErrorMsg(c, "container details not found")
|
||||
return
|
||||
}
|
||||
|
||||
events := make([]map[string]any, 0, len(details.ContainerEvents))
|
||||
for _, event := range details.ContainerEvents {
|
||||
events = append(events, map[string]any{
|
||||
"time": event.Time.Unix(),
|
||||
"message": event.Message,
|
||||
})
|
||||
}
|
||||
|
||||
data := gin.H{
|
||||
"deployment_id": deploymentID,
|
||||
"container_id": details.ContainerID,
|
||||
"device_id": details.DeviceID,
|
||||
"status": strings.ToLower(strings.TrimSpace(details.Status)),
|
||||
"hardware": details.Hardware,
|
||||
"brand_name": details.BrandName,
|
||||
"created_at": details.CreatedAt.Unix(),
|
||||
"uptime_percent": details.UptimePercent,
|
||||
"gpus_per_container": details.GPUsPerContainer,
|
||||
"public_url": details.PublicURL,
|
||||
"events": events,
|
||||
}
|
||||
|
||||
common.ApiSuccess(c, data)
|
||||
}
|
||||
@@ -0,0 +1,112 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type emailBindRequest struct {
|
||||
Email string `json:"email"`
|
||||
FlowToken string `json:"flow_token"`
|
||||
NewCode string `json:"new_code"`
|
||||
OldCode string `json:"old_code"`
|
||||
}
|
||||
|
||||
func EmailBindStart(c *gin.Context) {
|
||||
identity, ok := middleware.GetSessionAuthIdentity(c)
|
||||
if !ok {
|
||||
writeSecurityOperationError(c, service.ErrAuthTokenInvalid)
|
||||
return
|
||||
}
|
||||
succeeded, notificationFailed := false, false
|
||||
defer func() {
|
||||
recordUserSecurityAudit(c, identity.UserID, "user.binding_start", map[string]any{"provider": "email", "success": succeeded, "notification_failed": notificationFailed})
|
||||
}()
|
||||
var request emailBindRequest
|
||||
if err := common.DecodeJson(c.Request.Body, &request); err != nil {
|
||||
writeSecurityOperationError(c, service.ErrVerificationContextInvalid)
|
||||
return
|
||||
}
|
||||
email, err := service.ValidateAccountEmail(request.Email)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
context, err := common.Marshal(service.AccountBindingContext{Provider: "email", Email: email})
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
authorization := middleware.RequireSecurityProof(c, service.VerificationOperation{Scope: service.VerificationScopeAccountBind, Context: context})
|
||||
if authorization == nil {
|
||||
return
|
||||
}
|
||||
data, err := service.StartEmailBinding(identity, authorization, email)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
succeeded, notificationFailed = true, data.NotificationWarning
|
||||
common.ApiSuccess(c, data)
|
||||
}
|
||||
|
||||
func EmailBindResend(c *gin.Context) {
|
||||
identity, ok := middleware.GetSessionAuthIdentity(c)
|
||||
if !ok {
|
||||
writeSecurityOperationError(c, service.ErrAuthTokenInvalid)
|
||||
return
|
||||
}
|
||||
succeeded := false
|
||||
defer func() {
|
||||
recordUserSecurityAudit(c, identity.UserID, "user.email_binding_resend", map[string]any{"success": succeeded})
|
||||
}()
|
||||
var request emailBindRequest
|
||||
if err := common.DecodeJson(c.Request.Body, &request); err != nil || request.FlowToken == "" {
|
||||
writeSecurityOperationError(c, model.ErrAuthFlowInvalid)
|
||||
return
|
||||
}
|
||||
data, err := service.ResendAccountEmailBinding(identity, request.FlowToken)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
succeeded = true
|
||||
common.ApiSuccess(c, data)
|
||||
}
|
||||
|
||||
func EmailBind(c *gin.Context) {
|
||||
identity, ok := middleware.GetSessionAuthIdentity(c)
|
||||
if !ok {
|
||||
writeSecurityOperationError(c, service.ErrAuthTokenInvalid)
|
||||
return
|
||||
}
|
||||
succeeded, notificationFailed := false, false
|
||||
defer func() {
|
||||
recordUserSecurityAudit(c, identity.UserID, "user.binding_bind", map[string]any{"provider": "email", "success": succeeded, "notification_failed": notificationFailed})
|
||||
}()
|
||||
var request emailBindRequest
|
||||
if err := common.DecodeJson(c.Request.Body, &request); err != nil || request.FlowToken == "" {
|
||||
writeSecurityOperationError(c, model.ErrAuthFlowInvalid)
|
||||
return
|
||||
}
|
||||
state, err := service.FinishEmailBinding(identity, request.FlowToken, request.NewCode, request.OldCode)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
succeeded = true
|
||||
notificationFailed = service.NotifyAccountSecurityChange(state.CurrentEmail, "Email address changed") != nil
|
||||
if err := service.NotifyAccountSecurityChange(state.Email, "Email address confirmed"); err != nil {
|
||||
notificationFailed = true
|
||||
}
|
||||
if err := model.PublishUserAuthCache(identity.UserID); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"success": true, "message": "", "data": gin.H{"notification_warning": notificationFailed}})
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func GetGroups(c *gin.Context) {
|
||||
groupNames := make([]string, 0)
|
||||
for groupName := range ratio_setting.GetGroupRatioCopy() {
|
||||
groupNames = append(groupNames, groupName)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": groupNames,
|
||||
})
|
||||
}
|
||||
|
||||
func GetUserGroups(c *gin.Context) {
|
||||
usableGroups := make(map[string]map[string]any)
|
||||
enabledGroups, err := model.GetEnabledGroupSet()
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
userGroup := ""
|
||||
userId := c.GetInt("id")
|
||||
userGroup, _ = model.GetUserGroup(userId, false)
|
||||
userUsableGroups := service.GetUserUsableGroups(userGroup)
|
||||
for groupName, _ := range ratio_setting.GetGroupRatioCopy() {
|
||||
if _, ok := enabledGroups[groupName]; !ok {
|
||||
continue
|
||||
}
|
||||
// UserUsableGroups contains the groups that the user can use
|
||||
if desc, ok := userUsableGroups[groupName]; ok {
|
||||
usableGroups[groupName] = map[string]any{
|
||||
"ratio": service.GetUserGroupRatio(userGroup, groupName),
|
||||
"desc": desc,
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(usableGroups) > 0 {
|
||||
if _, ok := userUsableGroups["auto"]; ok {
|
||||
usableGroups["auto"] = map[string]any{
|
||||
"ratio": "自动",
|
||||
"desc": setting.GetUsableGroupDescription("auto"),
|
||||
}
|
||||
}
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": usableGroups,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/setting"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestSelectableGroupsRequireEnabledChannel(t *testing.T) {
|
||||
originalUsableGroups := setting.UserUsableGroups2JSONString()
|
||||
originalRatios := ratio_setting.GroupRatio2JSONString()
|
||||
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"auto":"Auto","active":"Active","default":"Default","paused":"Paused","no-channel":"No channel"}`))
|
||||
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"active":0.2,"default":1,"paused":0.3,"no-channel":0.4}`))
|
||||
t.Cleanup(func() {
|
||||
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups))
|
||||
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalRatios))
|
||||
model.InvalidatePricingCache()
|
||||
})
|
||||
|
||||
db := setupModelListControllerTestDB(t)
|
||||
require.NoError(t, db.Create(&model.User{
|
||||
Id: 7001,
|
||||
Username: "group-availability-user",
|
||||
Password: "password",
|
||||
Group: "default",
|
||||
Status: common.UserStatusEnabled,
|
||||
}).Error)
|
||||
require.NoError(t, db.Create(&[]model.Channel{
|
||||
{Id: 11, Type: constant.ChannelTypeOpenAI, Key: "active-key", Name: "active", Status: common.ChannelStatusEnabled},
|
||||
{Id: 12, Type: constant.ChannelTypeOpenAI, Key: "default-key", Name: "default", Status: common.ChannelStatusManuallyDisabled},
|
||||
{Id: 13, Type: constant.ChannelTypeOpenAI, Key: "paused-key", Name: "paused", Status: common.ChannelStatusAutoDisabled},
|
||||
}).Error)
|
||||
require.NoError(t, db.Create(&[]model.Ability{
|
||||
{Group: "active", Model: "gpt-4o", ChannelId: 11, Enabled: true},
|
||||
{Group: "default", Model: "gpt-4o", ChannelId: 12, Enabled: true},
|
||||
{Group: "paused", Model: "gpt-4o", ChannelId: 13, Enabled: true},
|
||||
}).Error)
|
||||
model.InvalidatePricingCache()
|
||||
|
||||
userRecorder := httptest.NewRecorder()
|
||||
userContext, _ := gin.CreateTestContext(userRecorder)
|
||||
userContext.Request = httptest.NewRequest(http.MethodGet, "/api/user/groups", nil)
|
||||
userContext.Set("id", 7001)
|
||||
GetUserGroups(userContext)
|
||||
|
||||
var userResponse struct {
|
||||
Success bool `json:"success"`
|
||||
Data map[string]map[string]interface{} `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(userRecorder.Body.Bytes(), &userResponse))
|
||||
require.True(t, userResponse.Success)
|
||||
assert.ElementsMatch(t, []string{"active", "auto"}, mapKeys(userResponse.Data))
|
||||
|
||||
pricingRecorder := httptest.NewRecorder()
|
||||
pricingContext, _ := gin.CreateTestContext(pricingRecorder)
|
||||
pricingContext.Request = httptest.NewRequest(http.MethodGet, "/api/pricing", nil)
|
||||
GetPricing(pricingContext)
|
||||
|
||||
var pricingResponse struct {
|
||||
Success bool `json:"success"`
|
||||
GroupRatio map[string]float64 `json:"group_ratio"`
|
||||
UsableGroup map[string]string `json:"usable_group"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(pricingRecorder.Body.Bytes(), &pricingResponse))
|
||||
require.True(t, pricingResponse.Success)
|
||||
assert.Equal(t, map[string]float64{"active": 0.2}, pricingResponse.GroupRatio)
|
||||
assert.Equal(t, map[string]string{"active": "Active", "auto": "Auto"}, pricingResponse.UsableGroup)
|
||||
}
|
||||
|
||||
func mapKeys[V any](values map[string]V) []string {
|
||||
keys := make([]string, 0, len(values))
|
||||
for key := range values {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
return keys
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func GetImage(c *gin.Context) {
|
||||
|
||||
}
|
||||
@@ -0,0 +1,156 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strconv"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func GetAllLogs(c *gin.Context) {
|
||||
pageInfo := common.GetPageQuery(c)
|
||||
logType, _ := strconv.Atoi(c.Query("type"))
|
||||
startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64)
|
||||
endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64)
|
||||
username := c.Query("username")
|
||||
tokenName := c.Query("token_name")
|
||||
modelName := c.Query("model_name")
|
||||
channel, _ := strconv.Atoi(c.Query("channel"))
|
||||
group := c.Query("group")
|
||||
requestId := c.Query("request_id")
|
||||
upstreamRequestId := c.Query("upstream_request_id")
|
||||
logs, total, err := model.GetAllLogs(logType, startTimestamp, endTimestamp, modelName, username, tokenName, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), channel, group, requestId, upstreamRequestId)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if c.GetInt("role") < common.RoleRootUser {
|
||||
model.FormatAdminLogs(logs)
|
||||
} else {
|
||||
model.FormatRootLogs(logs)
|
||||
}
|
||||
pageInfo.SetTotal(int(total))
|
||||
pageInfo.SetItems(logs)
|
||||
common.ApiSuccess(c, pageInfo)
|
||||
return
|
||||
}
|
||||
|
||||
func GetUserLogs(c *gin.Context) {
|
||||
pageInfo := common.GetPageQuery(c)
|
||||
userId := c.GetInt("id")
|
||||
logType, _ := strconv.Atoi(c.Query("type"))
|
||||
startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64)
|
||||
endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64)
|
||||
tokenName := c.Query("token_name")
|
||||
modelName := c.Query("model_name")
|
||||
group := c.Query("group")
|
||||
requestId := c.Query("request_id")
|
||||
upstreamRequestId := c.Query("upstream_request_id")
|
||||
logs, total, err := model.GetUserLogs(userId, logType, startTimestamp, endTimestamp, modelName, tokenName, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), group, requestId, upstreamRequestId)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
pageInfo.SetTotal(int(total))
|
||||
pageInfo.SetItems(logs)
|
||||
common.ApiSuccess(c, pageInfo)
|
||||
return
|
||||
}
|
||||
|
||||
// Deprecated: SearchAllLogs 已废弃，前端未使用该接口。
|
||||
func SearchAllLogs(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "该接口已废弃",
|
||||
})
|
||||
}
|
||||
|
||||
// Deprecated: SearchUserLogs 已废弃，前端未使用该接口。
|
||||
func SearchUserLogs(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "该接口已废弃",
|
||||
})
|
||||
}
|
||||
|
||||
func GetLogByKey(c *gin.Context) {
|
||||
tokenId := c.GetInt("token_id")
|
||||
if tokenId == 0 {
|
||||
c.JSON(200, gin.H{
|
||||
"success": false,
|
||||
"message": "无效的令牌",
|
||||
})
|
||||
return
|
||||
}
|
||||
logs, err := model.GetLogByTokenId(tokenId)
|
||||
if err != nil {
|
||||
c.JSON(200, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
c.JSON(200, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": logs,
|
||||
})
|
||||
}
|
||||
|
||||
func GetLogsStat(c *gin.Context) {
|
||||
logType, _ := strconv.Atoi(c.Query("type"))
|
||||
startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64)
|
||||
endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64)
|
||||
tokenName := c.Query("token_name")
|
||||
username := c.Query("username")
|
||||
modelName := c.Query("model_name")
|
||||
channel, _ := strconv.Atoi(c.Query("channel"))
|
||||
group := c.Query("group")
|
||||
stat, err := model.SumUsedQuota(logType, startTimestamp, endTimestamp, modelName, username, tokenName, channel, group)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
//tokenNum := model.SumUsedToken(logType, startTimestamp, endTimestamp, modelName, username, "")
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": gin.H{
|
||||
"quota": stat.Quota,
|
||||
"rpm": stat.Rpm,
|
||||
"tpm": stat.Tpm,
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
func GetLogsSelfStat(c *gin.Context) {
|
||||
username := c.GetString("username")
|
||||
logType, _ := strconv.Atoi(c.Query("type"))
|
||||
startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64)
|
||||
endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64)
|
||||
tokenName := c.Query("token_name")
|
||||
modelName := c.Query("model_name")
|
||||
channel, _ := strconv.Atoi(c.Query("channel"))
|
||||
group := c.Query("group")
|
||||
quotaNum, err := model.SumUsedQuota(logType, startTimestamp, endTimestamp, modelName, username, tokenName, channel, group)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
//tokenNum := model.SumUsedToken(logType, startTimestamp, endTimestamp, modelName, username, tokenName)
|
||||
c.JSON(200, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": gin.H{
|
||||
"quota": quotaNum.Quota,
|
||||
"rpm": quotaNum.Rpm,
|
||||
"tpm": quotaNum.Tpm,
|
||||
//"token": tokenNum,
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
@@ -0,0 +1,151 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
passkeysvc "github.com/QuantumNous/new-api/service/passkey"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/go-webauthn/webauthn/protocol"
|
||||
webauthnlib "github.com/go-webauthn/webauthn/webauthn"
|
||||
)
|
||||
|
||||
func VerifyLogin(c *gin.Context) {
|
||||
var request struct {
|
||||
FlowToken string `json:"flow_token"`
|
||||
Method string `json:"method"`
|
||||
Code string `json:"code"`
|
||||
}
|
||||
if common.DecodeJson(c.Request.Body, &request) != nil || request.FlowToken == "" || request.Code == "" {
|
||||
common.ApiErrorMsg(c, "参数错误")
|
||||
return
|
||||
}
|
||||
if request.Method == "" {
|
||||
request.Method = service.VerificationMethodTwoFA
|
||||
}
|
||||
if request.Method != service.VerificationMethodTwoFA {
|
||||
writeSecurityOperationError(c, service.ErrProofMethod)
|
||||
return
|
||||
}
|
||||
bundle, err := service.VerifyLoginCode(request.FlowToken, request.Code, c.ClientIP(), c.Request.UserAgent())
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
completeVerifiedLoginResponse(c, bundle, service.VerificationMethodTwoFA)
|
||||
}
|
||||
|
||||
func LoginPasskeyBegin(c *gin.Context) {
|
||||
var request struct {
|
||||
FlowToken string `json:"flow_token"`
|
||||
}
|
||||
if common.DecodeJson(c.Request.Body, &request) != nil || request.FlowToken == "" {
|
||||
common.ApiErrorMsg(c, "参数错误")
|
||||
return
|
||||
}
|
||||
verification, err := service.RequireLoginVerification(request.FlowToken, service.VerificationMethodPasskey)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
credential, err := model.GetPasskeyByUserID(verification.State.UserID)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
wa, err := passkeysvc.BuildWebAuthn(c.Request)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
user := &model.User{Id: verification.State.UserID}
|
||||
options, sessionData, err := wa.BeginLogin(passkeysvc.NewWebAuthnUser(user, credential), webauthnlib.WithUserVerification(protocol.VerificationRequired))
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
token, expiresAt, err := passkeysvc.CreateSessionDataFlow(model.AuthFlowPurposeLoginPasskey, passkeysvc.FlowSecurity{
|
||||
AuthSessionIdentity: model.AuthSessionIdentity{UserID: user.Id, UserAuthVersion: verification.State.AuthVersion},
|
||||
LoginFlowID: verification.Flow.Id, LoginExpiresAt: verification.Flow.ExpiresAt.Unix(),
|
||||
}, sessionData)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, gin.H{"flow_token": token, "expires_at": expiresAt, "options": options})
|
||||
}
|
||||
|
||||
func LoginPasskeyFinish(c *gin.Context) {
|
||||
var request struct {
|
||||
FlowToken string `json:"flow_token"`
|
||||
PasskeyFlowToken string `json:"passkey_flow_token"`
|
||||
Credential json.RawMessage `json:"credential"`
|
||||
}
|
||||
if common.DecodeJson(c.Request.Body, &request) != nil || request.FlowToken == "" || request.PasskeyFlowToken == "" || len(request.Credential) == 0 {
|
||||
common.ApiErrorMsg(c, "参数错误")
|
||||
return
|
||||
}
|
||||
verification, err := service.RequireLoginVerification(request.FlowToken, service.VerificationMethodPasskey)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
parsed, err := protocol.ParseCredentialRequestResponseBytes(request.Credential)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
identity := model.AuthSessionIdentity{UserID: verification.State.UserID, UserAuthVersion: verification.State.AuthVersion}
|
||||
sessionData, security, err := passkeysvc.PopSessionDataFlow(request.PasskeyFlowToken, model.AuthFlowPurposeLoginPasskey, identity)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
if security.LoginFlowID != verification.Flow.Id || sessionData.UserVerification != protocol.VerificationRequired {
|
||||
writeSecurityOperationError(c, model.ErrAuthFlowInvalid)
|
||||
return
|
||||
}
|
||||
credential, err := model.GetPasskeyByUserID(identity.UserID)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
wa, err := passkeysvc.BuildWebAuthn(c.Request)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
validated, err := wa.ValidateLogin(passkeysvc.NewWebAuthnUser(&model.User{Id: identity.UserID}, credential), *sessionData, parsed)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
if err := model.UpdatePasskeyAssertionState(identity.UserID, validated, time.Now()); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
bundle, err := service.CompleteLoginVerification(request.FlowToken, verification, service.VerificationMethodPasskey, c.ClientIP(), c.Request.UserAgent())
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
completeVerifiedLoginResponse(c, bundle, service.VerificationMethodPasskey)
|
||||
}
|
||||
|
||||
func completeVerifiedLoginResponse(c *gin.Context, bundle *service.AuthBundle, method string) {
|
||||
identity, err := service.ParseAccessToken(bundle.AccessToken)
|
||||
if err != nil {
|
||||
writeAuthSessionError(c, err)
|
||||
return
|
||||
}
|
||||
user, err := model.GetSelfUserById(identity.UserID)
|
||||
if err != nil {
|
||||
writeAuthSessionError(c, err)
|
||||
return
|
||||
}
|
||||
c.Set("login_verification_method", method)
|
||||
writeLoginResponse(c, user, bundle)
|
||||
}
|
||||
@@ -0,0 +1,329 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting"
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// midjourneyPollSummary is the result recorded on a midjourney_poll system task
|
||||
// row, summarizing one polling pass.
|
||||
type midjourneyPollSummary struct {
|
||||
UnfinishedTasks int `json:"unfinished_tasks"`
|
||||
ChannelsScanned int `json:"channels_scanned"`
|
||||
NullTasksFailed int `json:"null_tasks_failed"`
|
||||
}
|
||||
|
||||
// runMidjourneyTaskUpdateOnce performs one Midjourney polling pass synchronously.
|
||||
// It honors ctx cancellation (the system-task runner cancels it when the lease
|
||||
// is lost) and, when report is non-nil, reports progress as (processedChannels,
|
||||
// totalChannels) so the system task surfaces a percentage.
|
||||
func runMidjourneyTaskUpdateOnce(ctx context.Context, report func(processed, total int)) midjourneyPollSummary {
|
||||
summary := midjourneyPollSummary{}
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
|
||||
tasks := model.GetAllUnFinishTasks()
|
||||
if len(tasks) == 0 {
|
||||
return summary
|
||||
}
|
||||
summary.UnfinishedTasks = len(tasks)
|
||||
|
||||
logger.LogInfo(ctx, fmt.Sprintf("检测到未完成的任务数有: %v", len(tasks)))
|
||||
taskChannelM := make(map[int][]string)
|
||||
taskM := make(map[string]*model.Midjourney)
|
||||
nullTaskIds := make([]int, 0)
|
||||
for _, task := range tasks {
|
||||
if task.MjId == "" {
|
||||
// 统计失败的未完成任务
|
||||
nullTaskIds = append(nullTaskIds, task.Id)
|
||||
continue
|
||||
}
|
||||
taskM[task.MjId] = task
|
||||
taskChannelM[task.ChannelId] = append(taskChannelM[task.ChannelId], task.MjId)
|
||||
}
|
||||
if len(nullTaskIds) > 0 {
|
||||
summary.NullTasksFailed = len(nullTaskIds)
|
||||
err := model.MjBulkUpdateByTaskIds(nullTaskIds, map[string]any{
|
||||
"status": "FAILURE",
|
||||
"progress": "100%",
|
||||
})
|
||||
if err != nil {
|
||||
logger.LogError(ctx, fmt.Sprintf("Fix null mj_id task error: %v", err))
|
||||
} else {
|
||||
logger.LogInfo(ctx, fmt.Sprintf("Fix null mj_id task success: %v", nullTaskIds))
|
||||
}
|
||||
}
|
||||
if len(taskChannelM) == 0 {
|
||||
return summary
|
||||
}
|
||||
|
||||
totalChannels := len(taskChannelM)
|
||||
processedChannels := 0
|
||||
for channelId, taskIds := range taskChannelM {
|
||||
if ctx != nil && ctx.Err() != nil {
|
||||
break
|
||||
}
|
||||
if report != nil {
|
||||
report(processedChannels, totalChannels)
|
||||
}
|
||||
processedChannels++
|
||||
summary.ChannelsScanned++
|
||||
logger.LogInfo(ctx, fmt.Sprintf("渠道 #%d 未完成的任务有: %d", channelId, len(taskIds)))
|
||||
if len(taskIds) == 0 {
|
||||
continue
|
||||
}
|
||||
midjourneyChannel, err := model.CacheGetChannel(channelId)
|
||||
if err != nil {
|
||||
logger.LogError(ctx, fmt.Sprintf("CacheGetChannel: %v", err))
|
||||
err := model.MjBulkUpdate(taskIds, map[string]any{
|
||||
"fail_reason": fmt.Sprintf("获取渠道信息失败，请联系管理员，渠道ID：%d", channelId),
|
||||
"status": "FAILURE",
|
||||
"progress": "100%",
|
||||
})
|
||||
if err != nil {
|
||||
logger.LogInfo(ctx, fmt.Sprintf("UpdateMidjourneyTask error: %v", err))
|
||||
}
|
||||
continue
|
||||
}
|
||||
requestUrl := fmt.Sprintf("%s/mj/task/list-by-condition", *midjourneyChannel.BaseURL)
|
||||
|
||||
body, err := common.Marshal(map[string]any{
|
||||
"ids": taskIds,
|
||||
})
|
||||
if err != nil {
|
||||
logger.LogError(ctx, fmt.Sprintf("Get Task marshal body error: %v", err))
|
||||
continue
|
||||
}
|
||||
timeout := time.Second * 15
|
||||
requestCtx, cancel := context.WithTimeout(ctx, timeout)
|
||||
req, err := http.NewRequestWithContext(requestCtx, "POST", requestUrl, bytes.NewBuffer(body))
|
||||
if err != nil {
|
||||
cancel()
|
||||
logger.LogError(ctx, fmt.Sprintf("Get Task error: %v", err))
|
||||
continue
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("mj-api-secret", midjourneyChannel.Key)
|
||||
resp, err := service.GetHttpClient().Do(req)
|
||||
if err != nil {
|
||||
logger.LogError(ctx, fmt.Sprintf("Get Task Do req error: %v", err))
|
||||
cancel()
|
||||
continue
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
logger.LogError(ctx, fmt.Sprintf("Get Task status code: %d", resp.StatusCode))
|
||||
resp.Body.Close()
|
||||
cancel()
|
||||
continue
|
||||
}
|
||||
responseBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
logger.LogError(ctx, fmt.Sprintf("Get Mjp Task parse body error: %v", err))
|
||||
resp.Body.Close()
|
||||
cancel()
|
||||
continue
|
||||
}
|
||||
var responseItems []dto.MidjourneyDto
|
||||
err = common.Unmarshal(responseBody, &responseItems)
|
||||
if err != nil {
|
||||
logger.LogError(ctx, fmt.Sprintf("Get Mjp Task parse body error2: %v, body: %s", err, string(responseBody)))
|
||||
resp.Body.Close()
|
||||
cancel()
|
||||
continue
|
||||
}
|
||||
resp.Body.Close()
|
||||
req.Body.Close()
|
||||
cancel()
|
||||
|
||||
for _, responseItem := range responseItems {
|
||||
task := taskM[responseItem.MjId]
|
||||
if task == nil {
|
||||
logger.LogWarn(ctx, fmt.Sprintf("Midjourney task response ignored: unknown mj_id=%s", responseItem.MjId))
|
||||
continue
|
||||
}
|
||||
|
||||
useTime := (time.Now().UnixNano() / int64(time.Millisecond)) - task.SubmitTime
|
||||
// 如果时间超过一小时，且进度不是100%，则认为任务失败
|
||||
if useTime > 3600000 && task.Progress != "100%" {
|
||||
responseItem.FailReason = "上游任务超时（超过1小时）"
|
||||
responseItem.Status = "FAILURE"
|
||||
}
|
||||
if !checkMjTaskNeedUpdate(task, responseItem) {
|
||||
continue
|
||||
}
|
||||
preStatus := task.Status
|
||||
task.Code = 1
|
||||
task.Progress = responseItem.Progress
|
||||
task.PromptEn = responseItem.PromptEn
|
||||
task.State = responseItem.State
|
||||
task.SubmitTime = responseItem.SubmitTime
|
||||
task.StartTime = responseItem.StartTime
|
||||
task.FinishTime = responseItem.FinishTime
|
||||
task.ImageUrl = responseItem.ImageUrl
|
||||
task.Status = responseItem.Status
|
||||
task.FailReason = responseItem.FailReason
|
||||
if responseItem.Properties != nil {
|
||||
propertiesStr, _ := common.Marshal(responseItem.Properties)
|
||||
task.Properties = string(propertiesStr)
|
||||
}
|
||||
if responseItem.Buttons != nil {
|
||||
buttonStr, _ := common.Marshal(responseItem.Buttons)
|
||||
task.Buttons = string(buttonStr)
|
||||
}
|
||||
// 映射 VideoUrl
|
||||
task.VideoUrl = responseItem.VideoUrl
|
||||
|
||||
// 映射 VideoUrls - 将数组序列化为 JSON 字符串
|
||||
if responseItem.VideoUrls != nil && len(responseItem.VideoUrls) > 0 {
|
||||
videoUrlsStr, err := common.Marshal(responseItem.VideoUrls)
|
||||
if err != nil {
|
||||
logger.LogError(ctx, fmt.Sprintf("序列化 VideoUrls 失败: %v", err))
|
||||
task.VideoUrls = "[]" // 失败时设置为空数组
|
||||
} else {
|
||||
task.VideoUrls = string(videoUrlsStr)
|
||||
}
|
||||
} else {
|
||||
task.VideoUrls = "" // 空值时清空字段
|
||||
}
|
||||
|
||||
shouldReturnQuota := false
|
||||
if (task.Progress != "100%" && responseItem.FailReason != "") || (task.Progress == "100%" && task.Status == "FAILURE") {
|
||||
logger.LogInfo(ctx, task.MjId+" 构建失败，"+task.FailReason)
|
||||
task.Progress = "100%"
|
||||
if task.Quota != 0 {
|
||||
shouldReturnQuota = true
|
||||
}
|
||||
}
|
||||
won, err := task.UpdateWithStatus(preStatus)
|
||||
if err != nil {
|
||||
logger.LogError(ctx, "UpdateMidjourneyTask task error: "+err.Error())
|
||||
} else if won && shouldReturnQuota {
|
||||
service.RefundMidjourneyQuota(ctx, task, "构图失败")
|
||||
}
|
||||
}
|
||||
}
|
||||
if report != nil && (ctx == nil || ctx.Err() == nil) {
|
||||
report(totalChannels, totalChannels)
|
||||
}
|
||||
return summary
|
||||
}
|
||||
|
||||
func checkMjTaskNeedUpdate(oldTask *model.Midjourney, newTask dto.MidjourneyDto) bool {
|
||||
if oldTask.Code != 1 {
|
||||
return true
|
||||
}
|
||||
if oldTask.Progress != newTask.Progress {
|
||||
return true
|
||||
}
|
||||
if oldTask.PromptEn != newTask.PromptEn {
|
||||
return true
|
||||
}
|
||||
if oldTask.State != newTask.State {
|
||||
return true
|
||||
}
|
||||
if oldTask.SubmitTime != newTask.SubmitTime {
|
||||
return true
|
||||
}
|
||||
if oldTask.StartTime != newTask.StartTime {
|
||||
return true
|
||||
}
|
||||
if oldTask.FinishTime != newTask.FinishTime {
|
||||
return true
|
||||
}
|
||||
if oldTask.ImageUrl != newTask.ImageUrl {
|
||||
return true
|
||||
}
|
||||
if oldTask.Status != newTask.Status {
|
||||
return true
|
||||
}
|
||||
if oldTask.FailReason != newTask.FailReason {
|
||||
return true
|
||||
}
|
||||
if oldTask.FinishTime != newTask.FinishTime {
|
||||
return true
|
||||
}
|
||||
if oldTask.Progress != "100%" && newTask.FailReason != "" {
|
||||
return true
|
||||
}
|
||||
// 检查 VideoUrl 是否需要更新
|
||||
if oldTask.VideoUrl != newTask.VideoUrl {
|
||||
return true
|
||||
}
|
||||
// 检查 VideoUrls 是否需要更新
|
||||
if newTask.VideoUrls != nil && len(newTask.VideoUrls) > 0 {
|
||||
newVideoUrlsStr, _ := common.Marshal(newTask.VideoUrls)
|
||||
if oldTask.VideoUrls != string(newVideoUrlsStr) {
|
||||
return true
|
||||
}
|
||||
} else if oldTask.VideoUrls != "" {
|
||||
// 如果新数据没有 VideoUrls 但旧数据有，需要更新（清空）
|
||||
return true
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
func GetAllMidjourney(c *gin.Context) {
|
||||
pageInfo := common.GetPageQuery(c)
|
||||
|
||||
// 解析其他查询参数
|
||||
queryParams := model.TaskQueryParams{
|
||||
ChannelID: c.Query("channel_id"),
|
||||
MjID: c.Query("mj_id"),
|
||||
StartTimestamp: c.Query("start_timestamp"),
|
||||
EndTimestamp: c.Query("end_timestamp"),
|
||||
}
|
||||
|
||||
items := model.GetAllTasks(pageInfo.GetStartIdx(), pageInfo.GetPageSize(), queryParams)
|
||||
total := model.CountAllTasks(queryParams)
|
||||
|
||||
if setting.MjForwardUrlEnabled {
|
||||
for i, midjourney := range items {
|
||||
midjourney.ImageUrl = system_setting.ServerAddress + "/mj/image/" + midjourney.MjId
|
||||
items[i] = midjourney
|
||||
}
|
||||
}
|
||||
pageInfo.SetTotal(int(total))
|
||||
pageInfo.SetItems(items)
|
||||
common.ApiSuccess(c, pageInfo)
|
||||
}
|
||||
|
||||
func GetUserMidjourney(c *gin.Context) {
|
||||
pageInfo := common.GetPageQuery(c)
|
||||
|
||||
userId := c.GetInt("id")
|
||||
|
||||
queryParams := model.TaskQueryParams{
|
||||
MjID: c.Query("mj_id"),
|
||||
StartTimestamp: c.Query("start_timestamp"),
|
||||
EndTimestamp: c.Query("end_timestamp"),
|
||||
}
|
||||
|
||||
items := model.GetAllUserTask(userId, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), queryParams)
|
||||
total := model.CountAllUserTask(userId, queryParams)
|
||||
|
||||
if setting.MjForwardUrlEnabled {
|
||||
for i, midjourney := range items {
|
||||
midjourney.ImageUrl = system_setting.ServerAddress + "/mj/image/" + midjourney.MjId
|
||||
items[i] = midjourney
|
||||
}
|
||||
}
|
||||
pageInfo.SetTotal(int(total))
|
||||
pageInfo.SetItems(items)
|
||||
common.ApiSuccess(c, pageInfo)
|
||||
}
|
||||
@@ -0,0 +1,313 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/i18n"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/oauth"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting"
|
||||
"github.com/QuantumNous/new-api/setting/console_setting"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func TestStatus(c *gin.Context) {
|
||||
err := model.PingDB()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{
|
||||
"success": false,
|
||||
"message": "数据库连接失败",
|
||||
})
|
||||
return
|
||||
}
|
||||
// 获取HTTP统计信息
|
||||
httpStats := middleware.GetStats()
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "Server is running",
|
||||
"http_stats": httpStats,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
func GetStatus(c *gin.Context) {
|
||||
|
||||
cs := console_setting.GetConsoleSetting()
|
||||
common.OptionMapRWMutex.RLock()
|
||||
defer common.OptionMapRWMutex.RUnlock()
|
||||
|
||||
passkeySetting := system_setting.GetPasskeySettings()
|
||||
legalSetting := system_setting.GetLegalSettings()
|
||||
|
||||
data := gin.H{
|
||||
"version": common.Version,
|
||||
"start_time": common.StartTime,
|
||||
"email_verification": common.EmailVerificationEnabled,
|
||||
"github_oauth": common.GitHubOAuthEnabled,
|
||||
"github_client_id": common.GitHubClientId,
|
||||
"discord_oauth": system_setting.GetDiscordSettings().Enabled,
|
||||
"discord_client_id": system_setting.GetDiscordSettings().ClientId,
|
||||
"linuxdo_oauth": common.LinuxDOOAuthEnabled,
|
||||
"linuxdo_client_id": common.LinuxDOClientId,
|
||||
"linuxdo_minimum_trust_level": common.LinuxDOMinimumTrustLevel,
|
||||
"telegram_oauth": common.TelegramOAuthEnabled,
|
||||
"telegram_oauth_configured": oauth.TelegramConfigurationError() == nil,
|
||||
"telegram_bot_name": common.TelegramBotName,
|
||||
"theme": "default",
|
||||
"system_name": common.SystemName,
|
||||
"logo": common.Logo,
|
||||
"footer_html": common.Footer,
|
||||
"wechat_qrcode": common.WeChatAccountQRCodeImageURL,
|
||||
"wechat_login": common.WeChatAuthEnabled,
|
||||
"server_address": system_setting.ServerAddress,
|
||||
"turnstile_check": common.TurnstileCheckEnabled,
|
||||
"turnstile_site_key": common.TurnstileSiteKey,
|
||||
"docs_link": operation_setting.GetGeneralSetting().DocsLink,
|
||||
"quota_per_unit": common.QuotaPerUnit,
|
||||
// 兼容旧前端：保留 display_in_currency，同时提供新的 quota_display_type
|
||||
"display_in_currency": operation_setting.IsCurrencyDisplay(),
|
||||
"quota_display_type": operation_setting.GetQuotaDisplayType(),
|
||||
"custom_currency_symbol": operation_setting.GetGeneralSetting().CustomCurrencySymbol,
|
||||
"custom_currency_exchange_rate": operation_setting.GetGeneralSetting().CustomCurrencyExchangeRate,
|
||||
"enable_batch_update": common.BatchUpdateEnabled,
|
||||
"enable_drawing": common.DrawingEnabled,
|
||||
"enable_task": common.TaskEnabled,
|
||||
"enable_data_export": common.DataExportEnabled,
|
||||
"data_export_default_time": common.DataExportDefaultTime,
|
||||
"default_collapse_sidebar": common.DefaultCollapseSidebar,
|
||||
"mj_notify_enabled": setting.MjNotifyEnabled,
|
||||
"chats": setting.Chats,
|
||||
"demo_site_enabled": operation_setting.DemoSiteEnabled,
|
||||
"self_use_mode_enabled": operation_setting.SelfUseModeEnabled,
|
||||
"register_enabled": common.RegisterEnabled,
|
||||
"password_login_enabled": common.PasswordLoginEnabled,
|
||||
"password_register_enabled": common.PasswordRegisterEnabled,
|
||||
"default_use_auto_group": setting.DefaultUseAutoGroup,
|
||||
|
||||
"password_login_encryption_enabled": common.PasswordLoginEncryptionEnabled,
|
||||
|
||||
"usd_exchange_rate": operation_setting.USDExchangeRate,
|
||||
"price": operation_setting.Price,
|
||||
"stripe_unit_price": setting.StripeUnitPrice,
|
||||
|
||||
// 面板启用开关
|
||||
"api_info_enabled": cs.ApiInfoEnabled,
|
||||
"uptime_kuma_enabled": cs.UptimeKumaEnabled,
|
||||
"announcements_enabled": cs.AnnouncementsEnabled,
|
||||
"faq_enabled": cs.FAQEnabled,
|
||||
|
||||
// 模块管理配置
|
||||
"HeaderNavModules": common.OptionMap["HeaderNavModules"],
|
||||
"SidebarModulesAdmin": common.OptionMap["SidebarModulesAdmin"],
|
||||
|
||||
"oidc_enabled": system_setting.GetOIDCSettings().Enabled,
|
||||
"oidc_client_id": system_setting.GetOIDCSettings().ClientId,
|
||||
"oidc_authorization_endpoint": system_setting.GetOIDCSettings().AuthorizationEndpoint,
|
||||
"oidc_display_name": system_setting.GetOIDCSettings().GetEffectiveDisplayName(),
|
||||
"passkey_login": passkeySetting.Enabled,
|
||||
"passkey_display_name": passkeySetting.RPDisplayName,
|
||||
"passkey_rp_id": passkeySetting.RPID,
|
||||
"passkey_origins": passkeySetting.Origins,
|
||||
"passkey_allow_insecure": passkeySetting.AllowInsecureOrigin,
|
||||
"passkey_user_verification": passkeySetting.UserVerification,
|
||||
"passkey_attachment": passkeySetting.AttachmentPreference,
|
||||
"setup": constant.Setup,
|
||||
"user_agreement_enabled": legalSetting.UserAgreement != "",
|
||||
"privacy_policy_enabled": legalSetting.PrivacyPolicy != "",
|
||||
"checkin_enabled": operation_setting.GetCheckinSetting().Enabled,
|
||||
}
|
||||
|
||||
// 根据启用状态注入可选内容
|
||||
if cs.ApiInfoEnabled {
|
||||
data["api_info"] = console_setting.GetApiInfo()
|
||||
}
|
||||
if cs.AnnouncementsEnabled {
|
||||
data["announcements"] = console_setting.GetAnnouncements()
|
||||
}
|
||||
if cs.FAQEnabled {
|
||||
data["faq"] = console_setting.GetFAQ()
|
||||
}
|
||||
|
||||
// Add enabled custom OAuth providers
|
||||
customProviders := oauth.GetEnabledCustomProviders()
|
||||
if len(customProviders) > 0 {
|
||||
type CustomOAuthInfo struct {
|
||||
Id int `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Slug string `json:"slug"`
|
||||
Icon string `json:"icon"`
|
||||
ClientId string `json:"client_id"`
|
||||
AuthorizationEndpoint string `json:"authorization_endpoint"`
|
||||
Scopes string `json:"scopes"`
|
||||
}
|
||||
providersInfo := make([]CustomOAuthInfo, 0, len(customProviders))
|
||||
for _, p := range customProviders {
|
||||
config := p.GetConfig()
|
||||
providersInfo = append(providersInfo, CustomOAuthInfo{
|
||||
Id: config.Id,
|
||||
Name: config.Name,
|
||||
Slug: config.Slug,
|
||||
Icon: config.Icon,
|
||||
ClientId: config.ClientId,
|
||||
AuthorizationEndpoint: config.AuthorizationEndpoint,
|
||||
Scopes: config.Scopes,
|
||||
})
|
||||
}
|
||||
data["custom_oauth_providers"] = providersInfo
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": data,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
func GetNotice(c *gin.Context) {
|
||||
common.OptionMapRWMutex.RLock()
|
||||
notice := common.OptionMap["Notice"]
|
||||
common.OptionMapRWMutex.RUnlock()
|
||||
serveRevalidatedJSON(c, notice)
|
||||
}
|
||||
|
||||
func GetAbout(c *gin.Context) {
|
||||
common.OptionMapRWMutex.RLock()
|
||||
about := common.OptionMap["About"]
|
||||
common.OptionMapRWMutex.RUnlock()
|
||||
serveRevalidatedJSON(c, about)
|
||||
}
|
||||
|
||||
func GetUserAgreement(c *gin.Context) {
|
||||
serveRevalidatedJSON(c, system_setting.GetLegalSettings().UserAgreement)
|
||||
}
|
||||
|
||||
func GetPrivacyPolicy(c *gin.Context) {
|
||||
serveRevalidatedJSON(c, system_setting.GetLegalSettings().PrivacyPolicy)
|
||||
}
|
||||
|
||||
func GetMidjourney(c *gin.Context) {
|
||||
common.OptionMapRWMutex.RLock()
|
||||
defer common.OptionMapRWMutex.RUnlock()
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": common.OptionMap["Midjourney"],
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
func GetHomePageContent(c *gin.Context) {
|
||||
common.OptionMapRWMutex.RLock()
|
||||
homePageContent := common.OptionMap["HomePageContent"]
|
||||
common.OptionMapRWMutex.RUnlock()
|
||||
serveRevalidatedJSON(c, homePageContent)
|
||||
}
|
||||
|
||||
func SendEmailVerification(c *gin.Context) {
|
||||
email, err := service.ValidateAccountEmail(c.Query("email"))
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
if model.IsEmailAlreadyTaken(email) {
|
||||
common.ApiErrorI18n(c, i18n.MsgUserEmailAlreadyTaken)
|
||||
return
|
||||
}
|
||||
code := common.GenerateVerificationCode(6)
|
||||
common.RegisterVerificationCodeWithKey(email, code, common.EmailVerificationPurpose)
|
||||
subject := fmt.Sprintf("%s邮箱验证邮件", common.SystemName)
|
||||
content := fmt.Sprintf("<p>您好，你正在进行%s邮箱验证。</p>"+
|
||||
"<p>您的验证码为: <strong>%s</strong></p>"+
|
||||
"<p>验证码 %d 分钟内有效，如果不是本人操作，请忽略。</p>", common.SystemName, code, common.VerificationValidMinutes)
|
||||
err = common.SendEmail(subject, email, content)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
func SendPasswordResetEmail(c *gin.Context) {
|
||||
email := model.NormalizeEmail(c.Query("email"))
|
||||
if err := common.Validate.Var(email, "required,email"); err != nil {
|
||||
common.ApiErrorI18n(c, i18n.MsgInvalidParams)
|
||||
return
|
||||
}
|
||||
if _, err := model.GetUniqueUserByEmail(email); err == nil {
|
||||
code := common.GenerateVerificationCode(0)
|
||||
common.RegisterVerificationCodeWithKey(email, code, common.PasswordResetPurpose)
|
||||
link := fmt.Sprintf("%s/user/reset?email=%s&token=%s", system_setting.ServerAddress, email, code)
|
||||
subject := fmt.Sprintf("%s密码重置", common.SystemName)
|
||||
content := fmt.Sprintf("<p>您好，你正在进行%s密码重置。</p>"+
|
||||
"<p>点击 <a href='%s'>此处</a> 进行密码重置。</p>"+
|
||||
"<p>如果链接无法点击，请尝试点击下面的链接或将其复制到浏览器中打开：<br> %s </p>"+
|
||||
"<p>重置链接 %d 分钟内有效，如果不是本人操作，请忽略。</p>", common.SystemName, link, link, common.VerificationValidMinutes)
|
||||
err := common.SendEmail(subject, email, content)
|
||||
if err != nil {
|
||||
logger.LogError(c.Request.Context(), fmt.Sprintf("failed to send password reset email to %s: %s", email, err.Error()))
|
||||
}
|
||||
} else if err != nil && !errors.Is(err, model.ErrEmailNotFound) {
|
||||
logger.LogWarn(c.Request.Context(), fmt.Sprintf("skip password reset email for %s: %s", email, err.Error()))
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
})
|
||||
}
|
||||
|
||||
type PasswordResetRequest struct {
|
||||
Email string `json:"email"`
|
||||
Token string `json:"token"`
|
||||
}
|
||||
|
||||
func ResetPassword(c *gin.Context) {
|
||||
var req PasswordResetRequest
|
||||
err := json.NewDecoder(c.Request.Body).Decode(&req)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
req.Email = model.NormalizeEmail(req.Email)
|
||||
if req.Email == "" || req.Token == "" {
|
||||
common.ApiErrorI18n(c, i18n.MsgInvalidParams)
|
||||
return
|
||||
}
|
||||
if !common.VerifyCodeWithKey(req.Email, req.Token, common.PasswordResetPurpose) {
|
||||
common.ApiErrorI18n(c, i18n.MsgUserPasswordResetLinkInvalid)
|
||||
return
|
||||
}
|
||||
password := common.GenerateVerificationCode(12)
|
||||
err = model.ResetUserPasswordByEmail(req.Email, password)
|
||||
if err != nil {
|
||||
if errors.Is(err, model.ErrEmailNotFound) || errors.Is(err, model.ErrEmailAmbiguous) {
|
||||
common.ApiErrorI18n(c, i18n.MsgUserPasswordResetLinkInvalid)
|
||||
return
|
||||
}
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
common.DeleteKey(req.Email, common.PasswordResetPurpose)
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": password,
|
||||
})
|
||||
return
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGetStatusReturnsEffectiveOIDCDisplayName(t *testing.T) {
|
||||
settings := system_setting.GetOIDCSettings()
|
||||
originalDisplayName := settings.DisplayName
|
||||
originalOptionMap := common.OptionMap
|
||||
t.Cleanup(func() {
|
||||
settings.DisplayName = originalDisplayName
|
||||
common.OptionMap = originalOptionMap
|
||||
})
|
||||
common.OptionMap = map[string]string{}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
displayName string
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "custom name is trimmed",
|
||||
displayName: " Acme SSO ",
|
||||
want: "Acme SSO",
|
||||
},
|
||||
{
|
||||
name: "whitespace-only name falls back",
|
||||
displayName: " ",
|
||||
want: "OIDC",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
settings.DisplayName = tt.displayName
|
||||
response := httptest.NewRecorder()
|
||||
context, _ := gin.CreateTestContext(response)
|
||||
context.Request = httptest.NewRequest(http.MethodGet, "/api/status", nil)
|
||||
|
||||
GetStatus(context)
|
||||
|
||||
var payload struct {
|
||||
Success bool `json:"success"`
|
||||
Data map[string]any `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &payload))
|
||||
require.True(t, payload.Success)
|
||||
assert.Equal(t, tt.want, payload.Data["oidc_display_name"])
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// GetMissingModels returns the list of model names that are referenced by channels
|
||||
// but do not have corresponding records in the models meta table.
|
||||
// This helps administrators quickly discover models that need configuration.
|
||||
func GetMissingModels(c *gin.Context) {
|
||||
missing, err := model.GetMissingModels()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"data": missing,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,373 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/relay"
|
||||
"github.com/QuantumNous/new-api/relay/channel/ai360"
|
||||
"github.com/QuantumNous/new-api/relay/channel/lingyiwanwu"
|
||||
"github.com/QuantumNous/new-api/relay/channel/minimax"
|
||||
"github.com/QuantumNous/new-api/relay/channel/moonshot"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relay/helper"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/samber/lo"
|
||||
)
|
||||
|
||||
// https://platform.openai.com/docs/api-reference/models/list
|
||||
|
||||
var openAIModels []dto.OpenAIModels
|
||||
var openAIModelsMap map[string]dto.OpenAIModels
|
||||
var channelId2Models map[int][]string
|
||||
|
||||
func init() {
|
||||
// https://platform.openai.com/docs/models/model-endpoint-compatibility
|
||||
for i := range constant.APITypeDummy {
|
||||
if i == constant.APITypeAIProxyLibrary {
|
||||
continue
|
||||
}
|
||||
adaptor := relay.GetAdaptor(i)
|
||||
channelName := adaptor.GetChannelName()
|
||||
modelNames := adaptor.GetModelList()
|
||||
for _, modelName := range modelNames {
|
||||
openAIModels = append(openAIModels, dto.OpenAIModels{
|
||||
Id: modelName,
|
||||
Object: "model",
|
||||
Created: 1626777600,
|
||||
OwnedBy: channelName,
|
||||
})
|
||||
}
|
||||
}
|
||||
for _, modelName := range ai360.ModelList {
|
||||
openAIModels = append(openAIModels, dto.OpenAIModels{
|
||||
Id: modelName,
|
||||
Object: "model",
|
||||
Created: 1626777600,
|
||||
OwnedBy: ai360.ChannelName,
|
||||
})
|
||||
}
|
||||
for _, modelName := range moonshot.ModelList {
|
||||
openAIModels = append(openAIModels, dto.OpenAIModels{
|
||||
Id: modelName,
|
||||
Object: "model",
|
||||
Created: 1626777600,
|
||||
OwnedBy: moonshot.ChannelName,
|
||||
})
|
||||
}
|
||||
for _, modelName := range lingyiwanwu.ModelList {
|
||||
openAIModels = append(openAIModels, dto.OpenAIModels{
|
||||
Id: modelName,
|
||||
Object: "model",
|
||||
Created: 1626777600,
|
||||
OwnedBy: lingyiwanwu.ChannelName,
|
||||
})
|
||||
}
|
||||
for _, modelName := range minimax.ModelList {
|
||||
openAIModels = append(openAIModels, dto.OpenAIModels{
|
||||
Id: modelName,
|
||||
Object: "model",
|
||||
Created: 1626777600,
|
||||
OwnedBy: minimax.ChannelName,
|
||||
})
|
||||
}
|
||||
for modelName, _ := range constant.MidjourneyModel2Action {
|
||||
openAIModels = append(openAIModels, dto.OpenAIModels{
|
||||
Id: modelName,
|
||||
Object: "model",
|
||||
Created: 1626777600,
|
||||
OwnedBy: "midjourney",
|
||||
})
|
||||
}
|
||||
openAIModelsMap = make(map[string]dto.OpenAIModels)
|
||||
for _, aiModel := range openAIModels {
|
||||
openAIModelsMap[aiModel.Id] = aiModel
|
||||
}
|
||||
channelId2Models = make(map[int][]string)
|
||||
for i := 1; i <= constant.ChannelTypeDummy; i++ {
|
||||
apiType, success := common.ChannelType2APIType(i)
|
||||
if !success || apiType == constant.APITypeAIProxyLibrary {
|
||||
if plugin, ok := jsplugin.DefaultRegistry.GetByChannelType(i); ok {
|
||||
channelId2Models[i] = append([]string(nil), plugin.Meta.Models...)
|
||||
}
|
||||
continue
|
||||
}
|
||||
meta := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{
|
||||
ChannelType: i,
|
||||
}}
|
||||
adaptor := relay.GetAdaptor(apiType)
|
||||
adaptor.Init(meta)
|
||||
channelId2Models[i] = adaptor.GetModelList()
|
||||
if len(channelId2Models[i]) == 0 {
|
||||
if plugin, ok := jsplugin.DefaultRegistry.GetByChannelType(i); ok {
|
||||
channelId2Models[i] = append([]string(nil), plugin.Meta.Models...)
|
||||
}
|
||||
}
|
||||
}
|
||||
openAIModels = lo.UniqBy(openAIModels, func(m dto.OpenAIModels) string {
|
||||
return m.Id
|
||||
})
|
||||
}
|
||||
|
||||
func channelOwnerName(channelType int) string {
|
||||
apiType, success := common.ChannelType2APIType(channelType)
|
||||
if !success {
|
||||
return strings.ToLower(constant.GetChannelTypeName(channelType))
|
||||
}
|
||||
adaptor := relay.GetAdaptor(apiType)
|
||||
if adaptor == nil {
|
||||
return strings.ToLower(constant.GetChannelTypeName(channelType))
|
||||
}
|
||||
adaptor.Init(&relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{
|
||||
ChannelType: channelType,
|
||||
}})
|
||||
if name := strings.TrimSpace(adaptor.GetChannelName()); name != "" {
|
||||
return name
|
||||
}
|
||||
return strings.ToLower(constant.GetChannelTypeName(channelType))
|
||||
}
|
||||
|
||||
func getPreferredModelOwners(modelNames []string, groups []string) map[string]string {
|
||||
channelTypes, err := model.GetPreferredModelOwnerChannelTypes(modelNames, groups)
|
||||
if err != nil {
|
||||
common.SysLog(fmt.Sprintf("GetPreferredModelOwnerChannelTypes error: %v", err))
|
||||
return map[string]string{}
|
||||
}
|
||||
|
||||
ownerByChannelType := make(map[int]string)
|
||||
owners := make(map[string]string, len(channelTypes))
|
||||
for modelName, channelType := range channelTypes {
|
||||
owner, ok := ownerByChannelType[channelType]
|
||||
if !ok {
|
||||
owner = channelOwnerName(channelType)
|
||||
ownerByChannelType[channelType] = owner
|
||||
}
|
||||
if owner != "" {
|
||||
owners[modelName] = owner
|
||||
}
|
||||
}
|
||||
return owners
|
||||
}
|
||||
|
||||
func buildOpenAIModel(modelName string, ownerByModel map[string]string) dto.OpenAIModels {
|
||||
var oaiModel dto.OpenAIModels
|
||||
if staticModel, ok := openAIModelsMap[modelName]; ok {
|
||||
oaiModel = staticModel
|
||||
} else {
|
||||
oaiModel = dto.OpenAIModels{
|
||||
Id: modelName,
|
||||
Object: "model",
|
||||
Created: 1626777600,
|
||||
OwnedBy: "custom",
|
||||
}
|
||||
}
|
||||
if owner, ok := ownerByModel[modelName]; ok && owner != "" {
|
||||
oaiModel.OwnedBy = owner
|
||||
}
|
||||
oaiModel.SupportedEndpointTypes = model.GetModelSupportEndpointTypes(modelName)
|
||||
return oaiModel
|
||||
}
|
||||
|
||||
type modelListGroups struct {
|
||||
userGroup string
|
||||
tokenGroup string
|
||||
ownerGroups []string
|
||||
}
|
||||
|
||||
func getModelListGroups(c *gin.Context) (modelListGroups, error) {
|
||||
tokenGroup := common.GetContextKeyString(c, constant.ContextKeyTokenGroup)
|
||||
userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup)
|
||||
if userGroup == "" && (tokenGroup == "" || tokenGroup == "auto") {
|
||||
var err error
|
||||
userGroup, err = model.GetUserGroup(c.GetInt("id"), false)
|
||||
if err != nil {
|
||||
return modelListGroups{}, err
|
||||
}
|
||||
}
|
||||
|
||||
if tokenGroup == "auto" {
|
||||
return modelListGroups{
|
||||
userGroup: userGroup,
|
||||
tokenGroup: tokenGroup,
|
||||
ownerGroups: service.GetRequestAutoGroups(c, userGroup),
|
||||
}, nil
|
||||
}
|
||||
|
||||
group := userGroup
|
||||
if tokenGroup != "" {
|
||||
group = tokenGroup
|
||||
}
|
||||
return modelListGroups{
|
||||
userGroup: userGroup,
|
||||
tokenGroup: tokenGroup,
|
||||
ownerGroups: []string{group},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func ListModels(c *gin.Context, modelType int) {
|
||||
acceptUnsetRatioModel := operation_setting.SelfUseModeEnabled
|
||||
if !acceptUnsetRatioModel {
|
||||
userId := c.GetInt("id")
|
||||
if userId > 0 {
|
||||
userSettings, _ := model.GetUserSetting(userId, false)
|
||||
if userSettings.AcceptUnsetRatioModel {
|
||||
acceptUnsetRatioModel = true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
userModelNames := make([]string, 0)
|
||||
groups, err := getModelListGroups(c)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "get user group failed",
|
||||
})
|
||||
return
|
||||
}
|
||||
ownerGroups := groups.ownerGroups
|
||||
modelLimitEnable := common.GetContextKeyBool(c, constant.ContextKeyTokenModelLimitEnabled)
|
||||
var tokenModelLimit map[string]bool
|
||||
if modelLimitEnable {
|
||||
s, ok := common.GetContextKey(c, constant.ContextKeyTokenModelLimit)
|
||||
if ok {
|
||||
tokenModelLimit, _ = s.(map[string]bool)
|
||||
}
|
||||
if tokenModelLimit == nil {
|
||||
tokenModelLimit = map[string]bool{}
|
||||
}
|
||||
}
|
||||
models := service.GetGroupsEnabledModels(ownerGroups)
|
||||
for _, modelName := range models {
|
||||
if modelLimitEnable {
|
||||
matchingName := ratio_setting.RoutingMatchModelName(modelName)
|
||||
if !tokenModelLimit[modelName] && !tokenModelLimit[matchingName] {
|
||||
continue
|
||||
}
|
||||
}
|
||||
if !acceptUnsetRatioModel && !helper.HasModelBillingConfig(modelName) {
|
||||
continue
|
||||
}
|
||||
userModelNames = append(userModelNames, modelName)
|
||||
}
|
||||
|
||||
ownerByModel := map[string]string{}
|
||||
if len(ownerGroups) > 0 {
|
||||
ownerByModel = getPreferredModelOwners(userModelNames, ownerGroups)
|
||||
}
|
||||
userOpenAiModels := make([]dto.OpenAIModels, 0, len(userModelNames))
|
||||
for _, modelName := range userModelNames {
|
||||
userOpenAiModels = append(userOpenAiModels, buildOpenAIModel(modelName, ownerByModel))
|
||||
}
|
||||
|
||||
switch modelType {
|
||||
case constant.ChannelTypeAnthropic:
|
||||
useranthropicModels := make([]dto.AnthropicModel, len(userOpenAiModels))
|
||||
for i, model := range userOpenAiModels {
|
||||
useranthropicModels[i] = dto.AnthropicModel{
|
||||
ID: model.Id,
|
||||
CreatedAt: time.Unix(int64(model.Created), 0).UTC().Format(time.RFC3339),
|
||||
DisplayName: model.Id,
|
||||
Type: "model",
|
||||
}
|
||||
}
|
||||
firstID := ""
|
||||
lastID := ""
|
||||
if len(useranthropicModels) > 0 {
|
||||
firstID = useranthropicModels[0].ID
|
||||
lastID = useranthropicModels[len(useranthropicModels)-1].ID
|
||||
}
|
||||
c.JSON(200, gin.H{
|
||||
"data": useranthropicModels,
|
||||
"first_id": firstID,
|
||||
"has_more": false,
|
||||
"last_id": lastID,
|
||||
})
|
||||
case constant.ChannelTypeGemini:
|
||||
userGeminiModels := make([]dto.GeminiModel, len(userOpenAiModels))
|
||||
for i, model := range userOpenAiModels {
|
||||
userGeminiModels[i] = dto.GeminiModel{
|
||||
Name: model.Id,
|
||||
DisplayName: model.Id,
|
||||
}
|
||||
}
|
||||
c.JSON(200, gin.H{
|
||||
"models": userGeminiModels,
|
||||
"nextPageToken": nil,
|
||||
})
|
||||
default:
|
||||
c.JSON(200, gin.H{
|
||||
"success": true,
|
||||
"data": userOpenAiModels,
|
||||
"object": "list",
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func ChannelListModels(c *gin.Context) {
|
||||
c.JSON(200, gin.H{
|
||||
"success": true,
|
||||
"data": openAIModels,
|
||||
})
|
||||
}
|
||||
|
||||
func DashboardListModels(c *gin.Context) {
|
||||
modelsByChannel := make(map[int][]string, len(channelId2Models))
|
||||
for channelType, models := range channelId2Models {
|
||||
modelsByChannel[channelType] = append([]string(nil), models...)
|
||||
}
|
||||
for channelType := 1; channelType <= constant.ChannelTypeDummy; channelType++ {
|
||||
if plugin, ok := jsplugin.DefaultRegistry.GetByChannelType(channelType); ok {
|
||||
modelsByChannel[channelType] = append([]string(nil), plugin.Meta.Models...)
|
||||
}
|
||||
}
|
||||
c.JSON(200, gin.H{
|
||||
"success": true,
|
||||
"data": modelsByChannel,
|
||||
})
|
||||
}
|
||||
|
||||
func EnabledListModels(c *gin.Context) {
|
||||
c.JSON(200, gin.H{
|
||||
"success": true,
|
||||
"data": model.GetEnabledModels(),
|
||||
})
|
||||
}
|
||||
|
||||
func RetrieveModel(c *gin.Context, modelType int) {
|
||||
modelId := c.Param("model")
|
||||
if aiModel, ok := openAIModelsMap[modelId]; ok {
|
||||
switch modelType {
|
||||
case constant.ChannelTypeAnthropic:
|
||||
c.JSON(200, dto.AnthropicModel{
|
||||
ID: aiModel.Id,
|
||||
CreatedAt: time.Unix(int64(aiModel.Created), 0).UTC().Format(time.RFC3339),
|
||||
DisplayName: aiModel.Id,
|
||||
Type: "model",
|
||||
})
|
||||
default:
|
||||
c.JSON(200, aiModel)
|
||||
}
|
||||
} else {
|
||||
openAIError := types.OpenAIError{
|
||||
Message: fmt.Sprintf("The model '%s' does not exist", modelId),
|
||||
Type: "invalid_request_error",
|
||||
Param: "model",
|
||||
Code: "model_not_found",
|
||||
}
|
||||
c.JSON(200, gin.H{
|
||||
"error": openAIError,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,531 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/setting"
|
||||
"github.com/QuantumNous/new-api/setting/config"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type listModelsResponse struct {
|
||||
Success bool `json:"success"`
|
||||
Data []dto.OpenAIModels `json:"data"`
|
||||
Object string `json:"object"`
|
||||
}
|
||||
|
||||
type userModelsResponse struct {
|
||||
Success bool `json:"success"`
|
||||
Data []string `json:"data"`
|
||||
}
|
||||
|
||||
func setupModelListControllerTestDB(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
|
||||
initModelListColumnNames(t)
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
common.SetDatabaseTypes(common.DatabaseTypeSQLite, common.DatabaseTypeSQLite)
|
||||
common.RedisEnabled = false
|
||||
|
||||
dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_"))
|
||||
db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
model.DB = db
|
||||
model.LOG_DB = db
|
||||
|
||||
require.NoError(t, db.AutoMigrate(&model.User{}, &model.Channel{}, &model.Ability{}, &model.Model{}, &model.Vendor{}))
|
||||
|
||||
t.Cleanup(func() {
|
||||
sqlDB, err := db.DB()
|
||||
if err == nil {
|
||||
_ = sqlDB.Close()
|
||||
}
|
||||
})
|
||||
|
||||
return db
|
||||
}
|
||||
|
||||
func initModelListColumnNames(t *testing.T) {
|
||||
t.Helper()
|
||||
|
||||
originalIsMasterNode := common.IsMasterNode
|
||||
originalSQLitePath := common.SQLitePath
|
||||
originalMainDatabaseType := common.MainDatabaseType()
|
||||
originalLogDatabaseType := common.LogDatabaseType()
|
||||
originalSQLDSN, hadSQLDSN := os.LookupEnv("SQL_DSN")
|
||||
defer func() {
|
||||
common.IsMasterNode = originalIsMasterNode
|
||||
common.SQLitePath = originalSQLitePath
|
||||
common.SetDatabaseTypes(originalMainDatabaseType, originalLogDatabaseType)
|
||||
if hadSQLDSN {
|
||||
require.NoError(t, os.Setenv("SQL_DSN", originalSQLDSN))
|
||||
} else {
|
||||
require.NoError(t, os.Unsetenv("SQL_DSN"))
|
||||
}
|
||||
}()
|
||||
|
||||
common.IsMasterNode = false
|
||||
common.SQLitePath = fmt.Sprintf("file:%s_init?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_"))
|
||||
common.SetDatabaseTypes(common.DatabaseTypeSQLite, common.DatabaseTypeSQLite)
|
||||
require.NoError(t, os.Setenv("SQL_DSN", "local"))
|
||||
|
||||
require.NoError(t, model.InitDB())
|
||||
if model.DB != nil {
|
||||
sqlDB, err := model.DB.DB()
|
||||
if err == nil {
|
||||
_ = sqlDB.Close()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func withTieredBillingConfig(t *testing.T, modes map[string]string, exprs map[string]string) {
|
||||
t.Helper()
|
||||
|
||||
saved := map[string]string{}
|
||||
require.NoError(t, config.GlobalConfig.SaveToDB(func(key, value string) error {
|
||||
if strings.HasPrefix(key, "billing_setting.") {
|
||||
saved[key] = value
|
||||
}
|
||||
return nil
|
||||
}))
|
||||
t.Cleanup(func() {
|
||||
require.NoError(t, config.GlobalConfig.LoadFromDB(saved))
|
||||
model.InvalidatePricingCache()
|
||||
})
|
||||
|
||||
modeBytes, err := common.Marshal(modes)
|
||||
require.NoError(t, err)
|
||||
exprBytes, err := common.Marshal(exprs)
|
||||
require.NoError(t, err)
|
||||
|
||||
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
|
||||
"billing_setting.billing_mode": string(modeBytes),
|
||||
"billing_setting.billing_expr": string(exprBytes),
|
||||
}))
|
||||
model.InvalidatePricingCache()
|
||||
}
|
||||
|
||||
func withSelfUseModeDisabled(t *testing.T) {
|
||||
t.Helper()
|
||||
|
||||
original := operation_setting.SelfUseModeEnabled
|
||||
operation_setting.SelfUseModeEnabled = false
|
||||
t.Cleanup(func() {
|
||||
operation_setting.SelfUseModeEnabled = original
|
||||
})
|
||||
}
|
||||
|
||||
func withSelfUseModeEnabled(t *testing.T) {
|
||||
t.Helper()
|
||||
|
||||
original := operation_setting.SelfUseModeEnabled
|
||||
operation_setting.SelfUseModeEnabled = true
|
||||
t.Cleanup(func() {
|
||||
operation_setting.SelfUseModeEnabled = original
|
||||
})
|
||||
}
|
||||
|
||||
func decodeListModelsPayload(t *testing.T, recorder *httptest.ResponseRecorder) listModelsResponse {
|
||||
t.Helper()
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
var payload listModelsResponse
|
||||
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &payload))
|
||||
require.True(t, payload.Success)
|
||||
require.Equal(t, "list", payload.Object)
|
||||
return payload
|
||||
}
|
||||
|
||||
func decodeListModelsResponse(t *testing.T, recorder *httptest.ResponseRecorder) map[string]struct{} {
|
||||
t.Helper()
|
||||
|
||||
payload := decodeListModelsPayload(t, recorder)
|
||||
ids := make(map[string]struct{}, len(payload.Data))
|
||||
for _, item := range payload.Data {
|
||||
ids[item.Id] = struct{}{}
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
func pricingByModelName(pricings []model.Pricing) map[string]model.Pricing {
|
||||
byName := make(map[string]model.Pricing, len(pricings))
|
||||
for _, pricing := range pricings {
|
||||
byName[pricing.ModelName] = pricing
|
||||
}
|
||||
return byName
|
||||
}
|
||||
|
||||
func decodeUserModelsResponse(t *testing.T, recorder *httptest.ResponseRecorder) []string {
|
||||
t.Helper()
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
var payload userModelsResponse
|
||||
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &payload))
|
||||
require.True(t, payload.Success)
|
||||
return payload.Data
|
||||
}
|
||||
|
||||
func TestGetUserModelsFiltersByRequestedGroup(t *testing.T) {
|
||||
db := setupModelListControllerTestDB(t)
|
||||
require.NoError(t, db.Create(&model.User{
|
||||
Id: 1002,
|
||||
Username: "playground-model-user",
|
||||
Password: "password",
|
||||
Group: "default",
|
||||
Status: common.UserStatusEnabled,
|
||||
}).Error)
|
||||
require.NoError(t, db.Create(&[]model.Ability{
|
||||
{Group: "default", Model: "zz-default-only-model", ChannelId: 1, Enabled: true},
|
||||
{Group: "default", Model: "zz-disabled-model", ChannelId: 1, Enabled: false},
|
||||
}).Error)
|
||||
|
||||
defaultRecorder := httptest.NewRecorder()
|
||||
defaultContext, _ := gin.CreateTestContext(defaultRecorder)
|
||||
defaultContext.Request = httptest.NewRequest(http.MethodGet, "/api/user/models?group=default", nil)
|
||||
defaultContext.Set("id", 1002)
|
||||
|
||||
GetUserModels(defaultContext)
|
||||
|
||||
defaultModels := decodeUserModelsResponse(t, defaultRecorder)
|
||||
require.ElementsMatch(t, []string{"zz-default-only-model"}, defaultModels)
|
||||
|
||||
vipRecorder := httptest.NewRecorder()
|
||||
vipContext, _ := gin.CreateTestContext(vipRecorder)
|
||||
vipContext.Request = httptest.NewRequest(http.MethodGet, "/api/user/models?group=vip", nil)
|
||||
vipContext.Set("id", 1002)
|
||||
|
||||
GetUserModels(vipContext)
|
||||
|
||||
require.Empty(t, decodeUserModelsResponse(t, vipRecorder))
|
||||
}
|
||||
|
||||
func TestGetUserModelsExpandsAutoGroupsInConfiguredOrder(t *testing.T) {
|
||||
originalAutoGroups := setting.AutoGroups2JsonString()
|
||||
originalUsableGroups := setting.UserUsableGroups2JSONString()
|
||||
originalSpecialGroups := ratio_setting.GetGroupRatioSetting().GroupSpecialUsableGroup.ReadAll()
|
||||
t.Cleanup(func() {
|
||||
require.NoError(t, setting.UpdateAutoGroupsByJsonString(originalAutoGroups))
|
||||
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups))
|
||||
specialGroups := ratio_setting.GetGroupRatioSetting().GroupSpecialUsableGroup
|
||||
specialGroups.Clear()
|
||||
specialGroups.AddAll(originalSpecialGroups)
|
||||
})
|
||||
|
||||
require.NoError(t, setting.UpdateAutoGroupsByJsonString(`["vip","default","unavailable"]`))
|
||||
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"auto":"自动分组","default":"默认分组","unavailable":"不可用分组"}`))
|
||||
specialGroups := ratio_setting.GetGroupRatioSetting().GroupSpecialUsableGroup
|
||||
specialGroups.Clear()
|
||||
specialGroups.Set("default", map[string]string{
|
||||
"+:vip": "VIP 分组",
|
||||
"-:unavailable": "",
|
||||
})
|
||||
|
||||
db := setupModelListControllerTestDB(t)
|
||||
require.NoError(t, db.Create(&model.User{
|
||||
Id: 1003,
|
||||
Username: "playground-auto-model-user",
|
||||
Password: "password",
|
||||
Group: "default",
|
||||
Status: common.UserStatusEnabled,
|
||||
}).Error)
|
||||
require.NoError(t, db.Create(&[]model.Ability{
|
||||
{Group: "vip", Model: "zz-vip-model", ChannelId: 1, Enabled: true},
|
||||
{Group: "vip", Model: "zz-shared-model", ChannelId: 1, Enabled: true},
|
||||
{Group: "default", Model: "zz-default-model", ChannelId: 1, Enabled: true},
|
||||
{Group: "default", Model: "zz-shared-model", ChannelId: 2, Enabled: true},
|
||||
{Group: "unavailable", Model: "zz-unavailable-model", ChannelId: 1, Enabled: true},
|
||||
}).Error)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
context, _ := gin.CreateTestContext(recorder)
|
||||
context.Request = httptest.NewRequest(http.MethodGet, "/api/user/models?group=auto", nil)
|
||||
context.Set("id", 1003)
|
||||
|
||||
GetUserModels(context)
|
||||
|
||||
models := decodeUserModelsResponse(t, recorder)
|
||||
require.Len(t, models, 3)
|
||||
assert.ElementsMatch(t, []string{"zz-vip-model", "zz-shared-model"}, models[:2])
|
||||
assert.Equal(t, "zz-default-model", models[2])
|
||||
}
|
||||
|
||||
func TestListModelsIncludesTieredBillingModel(t *testing.T) {
|
||||
withSelfUseModeDisabled(t)
|
||||
withTieredBillingConfig(t, map[string]string{
|
||||
"zz-tiered-visible-model": "tiered_expr",
|
||||
"zz-tiered-empty-expr-model": "tiered_expr",
|
||||
"zz-tiered-missing-expr-model": "tiered_expr",
|
||||
}, map[string]string{
|
||||
"zz-tiered-visible-model": `tier("base", p * 1 + c * 2)`,
|
||||
"zz-tiered-empty-expr-model": " ",
|
||||
})
|
||||
|
||||
db := setupModelListControllerTestDB(t)
|
||||
require.NoError(t, db.Create(&model.User{
|
||||
Id: 1001,
|
||||
Username: "model-list-user",
|
||||
Password: "password",
|
||||
Group: "default",
|
||||
Status: common.UserStatusEnabled,
|
||||
}).Error)
|
||||
require.NoError(t, db.Create(&[]model.Ability{
|
||||
{Group: "default", Model: "zz-tiered-visible-model", ChannelId: 1, Enabled: true},
|
||||
{Group: "default", Model: "zz-tiered-empty-expr-model", ChannelId: 1, Enabled: true},
|
||||
{Group: "default", Model: "zz-tiered-missing-expr-model", ChannelId: 1, Enabled: true},
|
||||
{Group: "default", Model: "zz-unpriced-model", ChannelId: 1, Enabled: true},
|
||||
}).Error)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
ctx.Set("id", 1001)
|
||||
|
||||
ListModels(ctx, constant.ChannelTypeOpenAI)
|
||||
|
||||
ids := decodeListModelsResponse(t, recorder)
|
||||
require.Contains(t, ids, "zz-tiered-visible-model")
|
||||
require.NotContains(t, ids, "zz-tiered-empty-expr-model")
|
||||
require.NotContains(t, ids, "zz-tiered-missing-expr-model")
|
||||
require.NotContains(t, ids, "zz-unpriced-model")
|
||||
|
||||
pricingByName := pricingByModelName(model.GetPricing())
|
||||
visiblePricing, ok := pricingByName["zz-tiered-visible-model"]
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "tiered_expr", visiblePricing.BillingMode)
|
||||
require.NotEmpty(t, visiblePricing.BillingExpr)
|
||||
|
||||
emptyExprPricing, ok := pricingByName["zz-tiered-empty-expr-model"]
|
||||
require.True(t, ok)
|
||||
require.Empty(t, emptyExprPricing.BillingMode)
|
||||
require.Empty(t, emptyExprPricing.BillingExpr)
|
||||
|
||||
missingExprPricing, ok := pricingByName["zz-tiered-missing-expr-model"]
|
||||
require.True(t, ok)
|
||||
require.Empty(t, missingExprPricing.BillingMode)
|
||||
require.Empty(t, missingExprPricing.BillingExpr)
|
||||
}
|
||||
|
||||
func TestListModelsUsesAdvancedCustomEndpointTypesFromPricingCache(t *testing.T) {
|
||||
withSelfUseModeEnabled(t)
|
||||
db := setupModelListControllerTestDB(t)
|
||||
|
||||
originalMemoryCacheEnabled := common.MemoryCacheEnabled
|
||||
common.MemoryCacheEnabled = true
|
||||
t.Cleanup(func() {
|
||||
common.MemoryCacheEnabled = originalMemoryCacheEnabled
|
||||
model.InvalidatePricingCache()
|
||||
})
|
||||
|
||||
require.NoError(t, db.Create(&model.User{
|
||||
Id: 1003,
|
||||
Username: "advanced-custom-model-list-user",
|
||||
Password: "password",
|
||||
Group: "default",
|
||||
Status: common.UserStatusEnabled,
|
||||
}).Error)
|
||||
|
||||
channel := &model.Channel{
|
||||
Id: 701,
|
||||
Type: constant.ChannelTypeAdvancedCustom,
|
||||
Key: "advanced-custom-key",
|
||||
Status: common.ChannelStatusEnabled,
|
||||
Name: "advanced-custom-channel",
|
||||
Group: "default",
|
||||
Models: "gemini-3.5-flash",
|
||||
}
|
||||
channel.SetOtherSettings(dto.ChannelOtherSettings{
|
||||
AdvancedCustom: &dto.AdvancedCustomConfig{
|
||||
Routes: []dto.AdvancedCustomRoute{
|
||||
{
|
||||
IncomingPath: "/v1/chat/completions",
|
||||
UpstreamPath: "/v1/chat/completions",
|
||||
},
|
||||
{
|
||||
IncomingPath: "/v1/responses",
|
||||
UpstreamPath: "/v1beta/models/{model}:generateContent",
|
||||
Converter: "openai_responses_to_gemini_generate_content",
|
||||
Models: []string{"re:^gemini-"},
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
require.NoError(t, db.Create(channel).Error)
|
||||
require.NoError(t, db.Create(&model.Ability{
|
||||
Group: "default",
|
||||
Model: "gemini-3.5-flash",
|
||||
ChannelId: 701,
|
||||
Enabled: true,
|
||||
}).Error)
|
||||
|
||||
model.InitChannelCache()
|
||||
model.GetPricing()
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
ctx.Set("id", 1003)
|
||||
|
||||
ListModels(ctx, constant.ChannelTypeOpenAI)
|
||||
|
||||
payload := decodeListModelsPayload(t, recorder)
|
||||
require.Len(t, payload.Data, 1)
|
||||
require.Equal(t, "gemini-3.5-flash", payload.Data[0].Id)
|
||||
require.Equal(t, []constant.EndpointType{
|
||||
constant.EndpointTypeOpenAI,
|
||||
constant.EndpointTypeOpenAIResponse,
|
||||
}, payload.Data[0].SupportedEndpointTypes)
|
||||
}
|
||||
|
||||
func TestListModelsTokenLimitIncludesTieredBillingModel(t *testing.T) {
|
||||
withSelfUseModeDisabled(t)
|
||||
withTieredBillingConfig(t, map[string]string{
|
||||
"zz-token-tiered-visible-model": "tiered_expr",
|
||||
"zz-token-tiered-empty-expr-model": "tiered_expr",
|
||||
"zz-token-tiered-missing-expr-model": "tiered_expr",
|
||||
}, map[string]string{
|
||||
"zz-token-tiered-visible-model": `tier("base", p * 1 + c * 2)`,
|
||||
"zz-token-tiered-empty-expr-model": "",
|
||||
})
|
||||
db := setupModelListControllerTestDB(t)
|
||||
require.NoError(t, db.Create(&[]model.Ability{
|
||||
{Group: "default", Model: "zz-token-tiered-visible-model", ChannelId: 1, Enabled: true},
|
||||
{Group: "default", Model: "zz-token-tiered-empty-expr-model", ChannelId: 1, Enabled: true},
|
||||
{Group: "default", Model: "zz-token-tiered-missing-expr-model", ChannelId: 1, Enabled: true},
|
||||
{Group: "default", Model: "zz-token-unpriced-model", ChannelId: 1, Enabled: true},
|
||||
}).Error)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default")
|
||||
common.SetContextKey(ctx, constant.ContextKeyTokenModelLimitEnabled, true)
|
||||
common.SetContextKey(ctx, constant.ContextKeyTokenModelLimit, map[string]bool{
|
||||
"zz-token-tiered-visible-model": true,
|
||||
"zz-token-tiered-empty-expr-model": true,
|
||||
"zz-token-tiered-missing-expr-model": true,
|
||||
"zz-token-unpriced-model": true,
|
||||
})
|
||||
|
||||
ListModels(ctx, constant.ChannelTypeOpenAI)
|
||||
|
||||
ids := decodeListModelsResponse(t, recorder)
|
||||
require.Contains(t, ids, "zz-token-tiered-visible-model")
|
||||
require.NotContains(t, ids, "zz-token-tiered-empty-expr-model")
|
||||
require.NotContains(t, ids, "zz-token-tiered-missing-expr-model")
|
||||
require.NotContains(t, ids, "zz-token-unpriced-model")
|
||||
}
|
||||
|
||||
func TestListModelsTokenLimitUsesResolvedCustomAutoGroups(t *testing.T) {
|
||||
withSelfUseModeEnabled(t)
|
||||
originalMax := setting.GetMaxTokenAutoGroups()
|
||||
originalUsableGroups := setting.UserUsableGroups2JSONString()
|
||||
originalRatios := ratio_setting.GroupRatio2JSONString()
|
||||
require.NoError(t, setting.UpdateMaxTokenAutoGroups("5"))
|
||||
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`))
|
||||
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1}`))
|
||||
t.Cleanup(func() {
|
||||
require.NoError(t, setting.UpdateMaxTokenAutoGroups(fmt.Sprintf("%d", originalMax)))
|
||||
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups))
|
||||
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalRatios))
|
||||
})
|
||||
|
||||
db := setupModelListControllerTestDB(t)
|
||||
require.NoError(t, db.Create(&[]model.Ability{
|
||||
{Group: "vip", Model: "zz-vip-allowed", ChannelId: 1, Enabled: true},
|
||||
{Group: "vip", Model: "zz-vip-denied", ChannelId: 1, Enabled: true},
|
||||
{Group: "default", Model: "zz-default-outside-snapshot", ChannelId: 1, Enabled: true},
|
||||
}).Error)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default")
|
||||
common.SetContextKey(ctx, constant.ContextKeyTokenGroup, "auto")
|
||||
common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip"})
|
||||
common.SetContextKey(ctx, constant.ContextKeyTokenModelLimitEnabled, true)
|
||||
common.SetContextKey(ctx, constant.ContextKeyTokenModelLimit, map[string]bool{
|
||||
"zz-vip-allowed": true,
|
||||
"zz-default-outside-snapshot": true,
|
||||
"zz-not-enabled": true,
|
||||
})
|
||||
|
||||
ListModels(ctx, constant.ChannelTypeOpenAI)
|
||||
ids := decodeListModelsResponse(t, recorder)
|
||||
require.Equal(t, map[string]struct{}{"zz-vip-allowed": {}}, ids)
|
||||
|
||||
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default"}`))
|
||||
emptyRecorder := httptest.NewRecorder()
|
||||
emptyCtx, _ := gin.CreateTestContext(emptyRecorder)
|
||||
emptyCtx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
common.SetContextKey(emptyCtx, constant.ContextKeyUserGroup, "default")
|
||||
common.SetContextKey(emptyCtx, constant.ContextKeyTokenGroup, "auto")
|
||||
common.SetContextKey(emptyCtx, constant.ContextKeyTokenAutoGroups, []string{"vip"})
|
||||
common.SetContextKey(emptyCtx, constant.ContextKeyTokenModelLimitEnabled, true)
|
||||
common.SetContextKey(emptyCtx, constant.ContextKeyTokenModelLimit, map[string]bool{"zz-vip-allowed": true})
|
||||
|
||||
require.NotPanics(t, func() {
|
||||
ListModels(emptyCtx, constant.ChannelTypeAnthropic)
|
||||
})
|
||||
var anthropicResponse struct {
|
||||
Data []dto.AnthropicModel `json:"data"`
|
||||
FirstID string `json:"first_id"`
|
||||
LastID string `json:"last_id"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(emptyRecorder.Body.Bytes(), &anthropicResponse))
|
||||
require.Empty(t, anthropicResponse.Data)
|
||||
require.Empty(t, anthropicResponse.FirstID)
|
||||
require.Empty(t, anthropicResponse.LastID)
|
||||
}
|
||||
|
||||
func TestSetupLoginDoesNotTouchPasswordWhenPasswordFieldOmitted(t *testing.T) {
|
||||
db := setupModelListControllerTestDB(t)
|
||||
require.NoError(t, db.AutoMigrate(&model.Log{}, &model.AuditLog{}, &model.UserSession{}, &model.TwoFA{}, &model.PasskeyCredential{}))
|
||||
|
||||
hashedPassword, err := common.Password2Hash("CurrentPassword123")
|
||||
require.NoError(t, err)
|
||||
user := &model.User{
|
||||
Username: "twofa-user",
|
||||
Password: hashedPassword,
|
||||
Role: common.RoleCommonUser,
|
||||
Status: common.UserStatusEnabled,
|
||||
Group: "default",
|
||||
}
|
||||
require.NoError(t, db.Create(user).Error)
|
||||
|
||||
router := gin.New()
|
||||
router.GET("/", func(c *gin.Context) {
|
||||
setupLogin(&model.User{
|
||||
Id: user.Id,
|
||||
AuthVersion: user.AuthVersion,
|
||||
Username: user.Username,
|
||||
Role: user.Role,
|
||||
Status: user.Status,
|
||||
Group: user.Group,
|
||||
}, c)
|
||||
})
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
router.ServeHTTP(recorder, request)
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
var stored model.User
|
||||
require.NoError(t, db.First(&stored, user.Id).Error)
|
||||
assert.Equal(t, hashedPassword, stored.Password)
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,337 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// GetAllModelsMeta 获取模型列表（分页）
|
||||
func GetAllModelsMeta(c *gin.Context) {
|
||||
listModelsMeta(c, "", "")
|
||||
}
|
||||
|
||||
// SearchModelsMeta 搜索模型列表
|
||||
func SearchModelsMeta(c *gin.Context) {
|
||||
listModelsMeta(c, c.Query("keyword"), c.Query("vendor"))
|
||||
}
|
||||
|
||||
func listModelsMeta(c *gin.Context, keyword, vendor string) {
|
||||
squareState := model.ModelSquareState(c.Query("square_state"))
|
||||
switch squareState {
|
||||
case "", model.ModelSquareVisible, model.ModelSquareUnavailable, model.ModelSquareHidden, model.ModelSquarePartial:
|
||||
default:
|
||||
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "Invalid model square state"})
|
||||
return
|
||||
}
|
||||
|
||||
pageInfo := common.GetPageQuery(c)
|
||||
if squareState != "" && (pageInfo.GetPage() < 1 || pageInfo.GetPageSize() < 1) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "Invalid pagination"})
|
||||
return
|
||||
}
|
||||
offset, limit := pageInfo.GetStartIdx(), pageInfo.GetPageSize()
|
||||
if squareState != "" {
|
||||
// Visibility depends on live channels and metadata rules. Filter the
|
||||
// enriched candidate set before counting and paginating the results.
|
||||
offset, limit = 0, -1
|
||||
}
|
||||
search := model.SearchModels
|
||||
if c.Query("include_channel_models") == "true" {
|
||||
search = model.SearchModelsWithChannels
|
||||
}
|
||||
modelsMeta, total, err := search(keyword, vendor, c.Query("status"), c.Query("sync_official"), offset, limit)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if err := enrichModels(modelsMeta); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if squareState != "" {
|
||||
filtered := make([]*model.Model, 0, len(modelsMeta))
|
||||
for _, metadata := range modelsMeta {
|
||||
if metadata.SquareState == squareState {
|
||||
filtered = append(filtered, metadata)
|
||||
}
|
||||
}
|
||||
total = int64(len(filtered))
|
||||
start := len(filtered)
|
||||
if pageInfo.GetPage()-1 <= len(filtered)/pageInfo.GetPageSize() {
|
||||
start = (pageInfo.GetPage() - 1) * pageInfo.GetPageSize()
|
||||
}
|
||||
end := min(start+pageInfo.GetPageSize(), len(filtered))
|
||||
modelsMeta = filtered[start:end]
|
||||
}
|
||||
|
||||
vendorCounts, _ := model.GetVendorModelCounts()
|
||||
pageInfo.SetTotal(int(total))
|
||||
pageInfo.SetItems(modelsMeta)
|
||||
common.ApiSuccess(c, gin.H{
|
||||
"items": modelsMeta,
|
||||
"total": total,
|
||||
"page": pageInfo.GetPage(),
|
||||
"page_size": pageInfo.GetPageSize(),
|
||||
"vendor_counts": vendorCounts,
|
||||
})
|
||||
}
|
||||
|
||||
// GetModelMeta 根据 ID 获取单条模型信息
|
||||
func GetModelMeta(c *gin.Context) {
|
||||
idStr := c.Param("id")
|
||||
id, err := strconv.Atoi(idStr)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
var m model.Model
|
||||
if err := model.DB.First(&m, id).Error; err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if err := enrichModels([]*model.Model{&m}); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, &m)
|
||||
}
|
||||
|
||||
// CreateModelMeta 新建模型
|
||||
func CreateModelMeta(c *gin.Context) {
|
||||
var m model.Model
|
||||
if err := c.ShouldBindJSON(&m); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if m.ModelName == "" {
|
||||
common.ApiErrorMsg(c, "模型名称不能为空")
|
||||
return
|
||||
}
|
||||
if err := model.ValidateMetadataValues(model.MetadataValues{Endpoints: m.Endpoints, Status: m.Status, NameRule: m.NameRule}); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
// 名称冲突检查
|
||||
if dup, err := model.IsModelNameDuplicated(0, m.ModelName); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
} else if dup {
|
||||
common.ApiErrorMsg(c, "模型名称已存在")
|
||||
return
|
||||
}
|
||||
|
||||
if err := m.Insert(); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
model.RefreshPricing()
|
||||
m.HasMetadata = m.Id > 0
|
||||
common.ApiSuccess(c, &m)
|
||||
}
|
||||
|
||||
// UpdateModelMeta 更新模型
|
||||
func UpdateModelMeta(c *gin.Context) {
|
||||
statusOnly := c.Query("status_only") == "true"
|
||||
|
||||
var m model.Model
|
||||
if err := c.ShouldBindJSON(&m); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if m.Id == 0 {
|
||||
common.ApiErrorMsg(c, "缺少模型 ID")
|
||||
return
|
||||
}
|
||||
|
||||
if statusOnly {
|
||||
if m.Status != 0 && m.Status != 1 {
|
||||
common.ApiErrorMsg(c, "invalid catalog visibility")
|
||||
return
|
||||
}
|
||||
// 只更新状态，防止误清空其他字段
|
||||
if err := model.DB.Model(&model.Model{}).Where("id = ?", m.Id).Update("status", m.Status).Error; err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
} else {
|
||||
if strings.TrimSpace(m.ModelName) == "" {
|
||||
common.ApiErrorMsg(c, "模型名称不能为空")
|
||||
return
|
||||
}
|
||||
if err := model.ValidateMetadataValues(model.MetadataValues{Endpoints: m.Endpoints, Status: m.Status, NameRule: m.NameRule}); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
// 名称冲突检查
|
||||
if dup, err := model.IsModelNameDuplicated(m.Id, m.ModelName); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
} else if dup {
|
||||
common.ApiErrorMsg(c, "模型名称已存在")
|
||||
return
|
||||
}
|
||||
|
||||
if err := m.Update(); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
model.RefreshPricing()
|
||||
m.HasMetadata = m.Id > 0
|
||||
common.ApiSuccess(c, &m)
|
||||
}
|
||||
|
||||
// DeleteModelMeta 删除模型
|
||||
func DeleteModelMeta(c *gin.Context) {
|
||||
idStr := c.Param("id")
|
||||
id, err := strconv.Atoi(idStr)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
removeFromChannels, err := strconv.ParseBool(c.DefaultQuery("remove_from_channels", "false"))
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
removePricing, err := strconv.ParseBool(c.DefaultQuery("remove_pricing", "false"))
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if removePricing && c.GetInt("role") != common.RoleRootUser {
|
||||
c.JSON(http.StatusForbidden, gin.H{"success": false, "message": "Model pricing is managed by a super administrator."})
|
||||
return
|
||||
}
|
||||
result, err := model.DeleteModelMetadata([]int{id}, removeFromChannels, removePricing)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
recordManageAudit(c, "model.delete", map[string]any{"model_ids": []int{id}, "remove_from_channels": removeFromChannels, "remove_pricing": removePricing, "updated_channels": result.UpdatedChannels})
|
||||
common.ApiSuccess(c, result)
|
||||
}
|
||||
|
||||
func BatchDeleteModelMeta(c *gin.Context) {
|
||||
var request struct {
|
||||
ModelIDs []int `json:"model_ids"`
|
||||
RemoveFromChannels bool `json:"remove_from_channels"`
|
||||
RemovePricing bool `json:"remove_pricing"`
|
||||
}
|
||||
if err := common.DecodeJson(c.Request.Body, &request); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if request.RemovePricing && c.GetInt("role") != common.RoleRootUser {
|
||||
c.JSON(http.StatusForbidden, gin.H{"success": false, "message": "Model pricing is managed by a super administrator."})
|
||||
return
|
||||
}
|
||||
result, err := model.DeleteModelMetadata(request.ModelIDs, request.RemoveFromChannels, request.RemovePricing)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
recordManageAudit(c, "model.delete_batch", map[string]any{"model_ids": request.ModelIDs, "remove_from_channels": request.RemoveFromChannels, "remove_pricing": request.RemovePricing, "updated_channels": result.UpdatedChannels})
|
||||
common.ApiSuccess(c, result)
|
||||
}
|
||||
|
||||
// enrichModels keeps configured endpoints intact and derives connections from
|
||||
// enabled routes, including hidden or unpriced models absent from the catalog.
|
||||
func enrichModels(models []*model.Model) error {
|
||||
if len(models) == 0 {
|
||||
return nil
|
||||
}
|
||||
configured, err := model.GetConfiguredModelChannels()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, metadata := range models {
|
||||
if metadata == nil {
|
||||
continue
|
||||
}
|
||||
metadata.HasMetadata = metadata.Id > 0
|
||||
channelIDs := make(map[int]struct{})
|
||||
for name, ids := range configured {
|
||||
if metadata.MatchesName(name) {
|
||||
for _, id := range ids {
|
||||
channelIDs[id] = struct{}{}
|
||||
}
|
||||
}
|
||||
}
|
||||
metadata.ConfiguredChannelCount = len(channelIDs)
|
||||
}
|
||||
connections, err := model.GetModelConnections()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := model.FillModelSquareStates(models, configured, connections); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, metadata := range models {
|
||||
if metadata == nil {
|
||||
continue
|
||||
}
|
||||
channels := make(map[int]model.BoundChannel)
|
||||
groups := make(map[string]bool)
|
||||
names := make(map[string]bool)
|
||||
endpoints := make(map[string]bool)
|
||||
quotas := make(map[int]bool)
|
||||
for _, connection := range connections {
|
||||
name := connection.Model
|
||||
if !metadata.MatchesName(name) {
|
||||
continue
|
||||
}
|
||||
names[name] = true
|
||||
groups[connection.Group] = true
|
||||
channels[connection.ChannelId] = model.BoundChannel{Name: connection.ChannelName, Type: connection.ChannelType}
|
||||
for _, endpoint := range model.GetModelSupportEndpointTypes(name) {
|
||||
endpoints[string(endpoint)] = true
|
||||
}
|
||||
for _, quota := range model.GetModelQuotaTypes(name) {
|
||||
quotas[quota] = true
|
||||
}
|
||||
}
|
||||
metadata.BoundChannels = nil
|
||||
metadata.EnableGroups = nil
|
||||
metadata.SupportedEndpoints = nil
|
||||
metadata.QuotaTypes = nil
|
||||
metadata.MatchedModels = nil
|
||||
for _, channel := range channels {
|
||||
metadata.BoundChannels = append(metadata.BoundChannels, channel)
|
||||
}
|
||||
sort.Slice(metadata.BoundChannels, func(i, j int) bool {
|
||||
a, b := metadata.BoundChannels[i], metadata.BoundChannels[j]
|
||||
if a.Name == b.Name {
|
||||
return a.Type < b.Type
|
||||
}
|
||||
return a.Name < b.Name
|
||||
})
|
||||
for group := range groups {
|
||||
metadata.EnableGroups = append(metadata.EnableGroups, group)
|
||||
}
|
||||
for endpoint := range endpoints {
|
||||
metadata.SupportedEndpoints = append(metadata.SupportedEndpoints, endpoint)
|
||||
}
|
||||
for quota := range quotas {
|
||||
metadata.QuotaTypes = append(metadata.QuotaTypes, quota)
|
||||
}
|
||||
sort.Strings(metadata.EnableGroups)
|
||||
sort.Strings(metadata.SupportedEndpoints)
|
||||
sort.Ints(metadata.QuotaTypes)
|
||||
if metadata.NameRule != model.NameRuleExact {
|
||||
for name := range names {
|
||||
metadata.MatchedModels = append(metadata.MatchedModels, name)
|
||||
}
|
||||
sort.Strings(metadata.MatchedModels)
|
||||
metadata.MatchedCount = len(names)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,118 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/setting"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestChannelOwnerNameUsesAdaptorChannelName(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
channelType int
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "openai",
|
||||
channelType: constant.ChannelTypeOpenAI,
|
||||
expected: "openai",
|
||||
},
|
||||
{
|
||||
name: "codex",
|
||||
channelType: constant.ChannelTypeCodex,
|
||||
expected: "codex",
|
||||
},
|
||||
{
|
||||
name: "openrouter",
|
||||
channelType: constant.ChannelTypeOpenRouter,
|
||||
expected: "openrouter",
|
||||
},
|
||||
{
|
||||
name: "azure fallback",
|
||||
channelType: constant.ChannelTypeAzure,
|
||||
expected: "azure",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
require.Equal(t, tt.expected, channelOwnerName(tt.channelType))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildOpenAIModelOverridesOwnedBy(t *testing.T) {
|
||||
modelItem := buildOpenAIModel("gpt-5.4", map[string]string{"gpt-5.4": "openai"})
|
||||
require.Equal(t, "gpt-5.4", modelItem.Id)
|
||||
require.Equal(t, "openai", modelItem.OwnedBy)
|
||||
}
|
||||
|
||||
func TestBuildOpenAIModelFallsBackToCustomForUnknownModels(t *testing.T) {
|
||||
modelItem := buildOpenAIModel("custom-test-model", nil)
|
||||
require.Equal(t, "custom-test-model", modelItem.Id)
|
||||
require.Equal(t, "custom", modelItem.OwnedBy)
|
||||
}
|
||||
|
||||
func TestGetModelListGroupsUsesUserGroupWhenTokenGroupIsEmpty(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default")
|
||||
|
||||
groups, err := getModelListGroups(ctx)
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Equal(t, "default", groups.userGroup)
|
||||
require.Empty(t, groups.tokenGroup)
|
||||
require.Equal(t, []string{"default"}, groups.ownerGroups)
|
||||
}
|
||||
|
||||
func TestGetModelListGroupsUsesExplicitTokenGroup(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default")
|
||||
common.SetContextKey(ctx, constant.ContextKeyTokenGroup, "vip")
|
||||
|
||||
groups, err := getModelListGroups(ctx)
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Equal(t, "default", groups.userGroup)
|
||||
require.Equal(t, "vip", groups.tokenGroup)
|
||||
require.Equal(t, []string{"vip"}, groups.ownerGroups)
|
||||
}
|
||||
|
||||
func TestGetModelListGroupsUsesFilteredTokenAutoGroupsSnapshot(t *testing.T) {
|
||||
originalMax := setting.GetMaxTokenAutoGroups()
|
||||
originalUsableGroups := setting.UserUsableGroups2JSONString()
|
||||
originalRatios := ratio_setting.GroupRatio2JSONString()
|
||||
require.NoError(t, setting.UpdateMaxTokenAutoGroups("1"))
|
||||
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`))
|
||||
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1}`))
|
||||
t.Cleanup(func() {
|
||||
require.NoError(t, setting.UpdateMaxTokenAutoGroups(fmt.Sprintf("%d", originalMax)))
|
||||
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups))
|
||||
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalRatios))
|
||||
})
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default")
|
||||
common.SetContextKey(ctx, constant.ContextKeyTokenGroup, "auto")
|
||||
common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip", "default"})
|
||||
|
||||
groups, err := getModelListGroups(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"vip"}, groups.ownerGroups)
|
||||
|
||||
common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip"})
|
||||
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default"}`))
|
||||
groups, err = getModelListGroups(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, groups.ownerGroups)
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func GetModelPricingConfig(c *gin.Context) {
|
||||
snapshot, err := model.GetModelPricingSnapshot(c.QueryArray("model"))
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, snapshot)
|
||||
}
|
||||
|
||||
func UpdateModelPricingConfig(c *gin.Context) {
|
||||
var request struct {
|
||||
Changes []model.ModelPricingChange `json:"changes"`
|
||||
}
|
||||
if err := common.DecodeJson(c.Request.Body, &request); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()})
|
||||
return
|
||||
}
|
||||
if err := model.UpdateModelPricing(request.Changes); err != nil {
|
||||
status := http.StatusBadRequest
|
||||
if errors.Is(err, model.ErrModelPricingConflict) {
|
||||
status = http.StatusConflict
|
||||
}
|
||||
c.JSON(status, gin.H{"success": false, "message": err.Error()})
|
||||
return
|
||||
}
|
||||
names := make([]string, 0, len(request.Changes))
|
||||
for _, change := range request.Changes {
|
||||
names = append(names, change.ModelName)
|
||||
}
|
||||
recordManageAudit(c, "model.pricing.update", map[string]any{"models": names})
|
||||
common.ApiSuccess(c, gin.H{"updated_models": names})
|
||||
}
|
||||
@@ -0,0 +1,452 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"math/rand"
|
||||
"net"
|
||||
"net/http"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// 上游地址
|
||||
const (
|
||||
upstreamModelsURL = "https://basellm.github.io/llm-metadata/api/newapi/models.json"
|
||||
upstreamVendorsURL = "https://basellm.github.io/llm-metadata/api/newapi/vendors.json"
|
||||
)
|
||||
|
||||
func normalizeLocale(locale string) (string, bool) {
|
||||
switch strings.ToLower(strings.TrimSpace(locale)) {
|
||||
case "", "zh", "zh-cn":
|
||||
return "zh", true
|
||||
case "en":
|
||||
return "en", true
|
||||
case "ja":
|
||||
return "ja", true
|
||||
default:
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
|
||||
func getUpstreamBase() string {
|
||||
return common.GetEnvOrDefaultString("SYNC_UPSTREAM_BASE", "https://basellm.github.io/llm-metadata")
|
||||
}
|
||||
|
||||
func getUpstreamURLs(locale string) (modelsURL, vendorsURL string) {
|
||||
base := strings.TrimRight(getUpstreamBase(), "/")
|
||||
if l, ok := normalizeLocale(locale); ok && l != "" {
|
||||
return fmt.Sprintf("%s/api/i18n/%s/newapi/models.json", base, l),
|
||||
fmt.Sprintf("%s/api/i18n/%s/newapi/vendors.json", base, l)
|
||||
}
|
||||
return fmt.Sprintf("%s/api/newapi/models.json", base), fmt.Sprintf("%s/api/newapi/vendors.json", base)
|
||||
}
|
||||
|
||||
type upstreamEnvelope[T any] struct {
|
||||
Success bool `json:"success"`
|
||||
Message string `json:"message"`
|
||||
Data []T `json:"data"`
|
||||
}
|
||||
|
||||
type upstreamModel struct {
|
||||
Description string `json:"description"`
|
||||
Endpoints json.RawMessage `json:"endpoints"`
|
||||
Icon string `json:"icon"`
|
||||
ModelName string `json:"model_name"`
|
||||
NameRule int `json:"name_rule"`
|
||||
Status int `json:"status"`
|
||||
Tags string `json:"tags"`
|
||||
VendorName string `json:"vendor_name"`
|
||||
}
|
||||
|
||||
type upstreamVendor struct {
|
||||
Description string `json:"description"`
|
||||
Icon string `json:"icon"`
|
||||
Name string `json:"name"`
|
||||
Status int `json:"status"`
|
||||
}
|
||||
|
||||
var (
|
||||
etagCache = make(map[string]string)
|
||||
bodyCache = make(map[string][]byte)
|
||||
cacheMutex sync.RWMutex
|
||||
)
|
||||
|
||||
func newHTTPClient() *http.Client {
|
||||
timeoutSec := common.GetEnvOrDefault("SYNC_HTTP_TIMEOUT_SECONDS", 10)
|
||||
dialer := &net.Dialer{Timeout: time.Duration(timeoutSec) * time.Second}
|
||||
transport := &http.Transport{
|
||||
MaxIdleConns: 100,
|
||||
IdleConnTimeout: 90 * time.Second,
|
||||
TLSHandshakeTimeout: time.Duration(timeoutSec) * time.Second,
|
||||
ExpectContinueTimeout: 1 * time.Second,
|
||||
ResponseHeaderTimeout: time.Duration(timeoutSec) * time.Second,
|
||||
}
|
||||
if common.TLSInsecureSkipVerify {
|
||||
transport.TLSClientConfig = common.InsecureTLSConfig
|
||||
}
|
||||
transport.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
|
||||
host, _, err := net.SplitHostPort(addr)
|
||||
if err != nil {
|
||||
host = addr
|
||||
}
|
||||
if strings.HasSuffix(host, "github.io") {
|
||||
if conn, err := dialer.DialContext(ctx, "tcp4", addr); err == nil {
|
||||
return conn, nil
|
||||
}
|
||||
return dialer.DialContext(ctx, "tcp6", addr)
|
||||
}
|
||||
return dialer.DialContext(ctx, network, addr)
|
||||
}
|
||||
return &http.Client{Transport: transport}
|
||||
}
|
||||
|
||||
var (
|
||||
httpClientOnce sync.Once
|
||||
httpClient *http.Client
|
||||
)
|
||||
|
||||
func getHTTPClient() *http.Client {
|
||||
httpClientOnce.Do(func() {
|
||||
httpClient = newHTTPClient()
|
||||
})
|
||||
return httpClient
|
||||
}
|
||||
|
||||
func fetchJSON[T any](ctx context.Context, url string, out *upstreamEnvelope[T]) error {
|
||||
var lastErr error
|
||||
attempts := max(common.GetEnvOrDefault("SYNC_HTTP_RETRY", 3), 1)
|
||||
baseDelay := 200 * time.Millisecond
|
||||
maxMB := common.GetEnvOrDefault("SYNC_HTTP_MAX_MB", 10)
|
||||
maxBytes := int64(maxMB) << 20
|
||||
for attempt := 0; attempt < attempts; attempt++ {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// ETag conditional request
|
||||
cacheMutex.RLock()
|
||||
if et := etagCache[url]; et != "" {
|
||||
req.Header.Set("If-None-Match", et)
|
||||
}
|
||||
cacheMutex.RUnlock()
|
||||
|
||||
resp, err := getHTTPClient().Do(req)
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
// backoff with jitter
|
||||
sleep := baseDelay * time.Duration(1<<attempt)
|
||||
jitter := time.Duration(rand.Intn(150)) * time.Millisecond
|
||||
time.Sleep(sleep + jitter)
|
||||
continue
|
||||
}
|
||||
func() {
|
||||
defer resp.Body.Close()
|
||||
switch resp.StatusCode {
|
||||
case http.StatusOK:
|
||||
// read body into buffer for caching and flexible decode
|
||||
limited := io.LimitReader(resp.Body, maxBytes+1)
|
||||
buf, err := io.ReadAll(limited)
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
return
|
||||
}
|
||||
if int64(len(buf)) > maxBytes {
|
||||
lastErr = errors.New("upstream metadata exceeds size limit")
|
||||
return
|
||||
}
|
||||
// cache body and ETag
|
||||
cacheMutex.Lock()
|
||||
if et := resp.Header.Get("ETag"); et != "" {
|
||||
etagCache[url] = et
|
||||
}
|
||||
bodyCache[url] = buf
|
||||
cacheMutex.Unlock()
|
||||
|
||||
// Try decode as envelope first
|
||||
if err := common.Unmarshal(buf, out); err != nil {
|
||||
// Try decode as pure array
|
||||
var arr []T
|
||||
if err2 := common.Unmarshal(buf, &arr); err2 != nil {
|
||||
lastErr = err
|
||||
return
|
||||
}
|
||||
out.Success = true
|
||||
out.Data = arr
|
||||
out.Message = ""
|
||||
} else {
|
||||
if !out.Success && len(out.Data) == 0 && out.Message == "" {
|
||||
out.Success = true
|
||||
}
|
||||
}
|
||||
lastErr = nil
|
||||
case http.StatusNotModified:
|
||||
// use cache
|
||||
cacheMutex.RLock()
|
||||
buf := bodyCache[url]
|
||||
cacheMutex.RUnlock()
|
||||
if len(buf) == 0 {
|
||||
lastErr = errors.New("cache miss for 304 response")
|
||||
return
|
||||
}
|
||||
if err := common.Unmarshal(buf, out); err != nil {
|
||||
var arr []T
|
||||
if err2 := common.Unmarshal(buf, &arr); err2 != nil {
|
||||
lastErr = err
|
||||
return
|
||||
}
|
||||
out.Success = true
|
||||
out.Data = arr
|
||||
out.Message = ""
|
||||
} else {
|
||||
if !out.Success && len(out.Data) == 0 && out.Message == "" {
|
||||
out.Success = true
|
||||
}
|
||||
}
|
||||
lastErr = nil
|
||||
default:
|
||||
lastErr = errors.New(resp.Status)
|
||||
}
|
||||
}()
|
||||
if lastErr == nil {
|
||||
return nil
|
||||
}
|
||||
sleep := baseDelay * time.Duration(1<<attempt)
|
||||
jitter := time.Duration(rand.Intn(150)) * time.Millisecond
|
||||
time.Sleep(sleep + jitter)
|
||||
}
|
||||
return lastErr
|
||||
}
|
||||
|
||||
type metadataSyncSource struct {
|
||||
Locale string `json:"locale"`
|
||||
ModelsURL string `json:"models_url"`
|
||||
VendorsURL string `json:"vendors_url"`
|
||||
Version string `json:"version"`
|
||||
}
|
||||
|
||||
type metadataSyncField struct {
|
||||
Field string `json:"field"`
|
||||
Local any `json:"local"`
|
||||
Upstream any `json:"upstream"`
|
||||
}
|
||||
|
||||
type metadataSyncCandidate struct {
|
||||
ModelName string `json:"model_name"`
|
||||
Kind string `json:"kind"`
|
||||
Scope string `json:"scope"`
|
||||
RecordVersion string `json:"record_version"`
|
||||
Fields []metadataSyncField `json:"fields"`
|
||||
Upstream *model.MetadataValues `json:"upstream,omitempty"`
|
||||
VendorToCreate string `json:"vendor_to_create,omitempty"`
|
||||
}
|
||||
|
||||
func fetchMetadataCatalog(c *gin.Context, locale string) (metadataSyncSource, map[string]model.MetadataValues, map[string]model.Vendor, error) {
|
||||
resolved, valid := normalizeLocale(locale)
|
||||
if !valid {
|
||||
return metadataSyncSource{}, nil, nil, errors.New("unsupported metadata language")
|
||||
}
|
||||
modelsURL, vendorsURL := getUpstreamURLs(resolved)
|
||||
source := metadataSyncSource{Locale: resolved, ModelsURL: modelsURL, VendorsURL: vendorsURL}
|
||||
ctx, cancel := context.WithTimeout(c.Request.Context(), time.Duration(common.GetEnvOrDefault("SYNC_HTTP_TIMEOUT_SECONDS", 15))*time.Second)
|
||||
defer cancel()
|
||||
var modelsEnv upstreamEnvelope[upstreamModel]
|
||||
var vendorsEnv upstreamEnvelope[upstreamVendor]
|
||||
var modelsErr, vendorsErr error
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(2)
|
||||
go func() { defer wg.Done(); modelsErr = fetchJSON(ctx, modelsURL, &modelsEnv) }()
|
||||
go func() { defer wg.Done(); vendorsErr = fetchJSON(ctx, vendorsURL, &vendorsEnv) }()
|
||||
wg.Wait()
|
||||
if modelsErr != nil {
|
||||
return source, nil, nil, fmt.Errorf("fetch models (%s, %s): %w", resolved, modelsURL, modelsErr)
|
||||
}
|
||||
if vendorsErr != nil {
|
||||
return source, nil, nil, fmt.Errorf("fetch vendors (%s, %s): %w", resolved, vendorsURL, vendorsErr)
|
||||
}
|
||||
if !modelsEnv.Success || !vendorsEnv.Success {
|
||||
return source, nil, nil, errors.New("upstream metadata source reported failure")
|
||||
}
|
||||
models := make(map[string]model.MetadataValues)
|
||||
vendors := make(map[string]model.Vendor)
|
||||
for _, vendor := range vendorsEnv.Data {
|
||||
vendor.Name = strings.TrimSpace(vendor.Name)
|
||||
if vendor.Name == "" {
|
||||
continue
|
||||
}
|
||||
vendors[vendor.Name] = model.Vendor{Name: vendor.Name, Description: vendor.Description, Icon: vendor.Icon, Status: vendor.Status}
|
||||
}
|
||||
for _, item := range modelsEnv.Data {
|
||||
if strings.TrimSpace(item.ModelName) == "" {
|
||||
continue
|
||||
}
|
||||
endpoints := ""
|
||||
if len(item.Endpoints) > 0 && string(item.Endpoints) != "null" {
|
||||
if err := common.Unmarshal(item.Endpoints, &endpoints); err != nil {
|
||||
endpoints = string(item.Endpoints)
|
||||
}
|
||||
}
|
||||
values := model.MetadataValues{Description: item.Description, Icon: item.Icon, Tags: item.Tags, Vendor: strings.TrimSpace(item.VendorName), Endpoints: endpoints, NameRule: item.NameRule, Status: item.Status}
|
||||
if err := model.ValidateMetadataValues(values); err != nil {
|
||||
return source, nil, nil, fmt.Errorf("model %s: %w", item.ModelName, err)
|
||||
}
|
||||
if _, duplicate := models[item.ModelName]; duplicate {
|
||||
return source, nil, nil, fmt.Errorf("duplicate upstream model: %s", item.ModelName)
|
||||
}
|
||||
models[item.ModelName] = values
|
||||
}
|
||||
encoded, err := common.Marshal([]any{source.Locale, models, vendors})
|
||||
if err != nil {
|
||||
return source, nil, nil, err
|
||||
}
|
||||
source.Version = fmt.Sprintf("%x", sha256.Sum256(encoded))
|
||||
return source, models, vendors, nil
|
||||
}
|
||||
|
||||
func SyncUpstreamPreview(c *gin.Context) {
|
||||
source, upstream, upstreamVendors, err := fetchMetadataCatalog(c, c.Query("locale"))
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
locals, vendors, err := model.GetMetadataSyncState(model.DB)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
missing, err := model.GetMissingModels()
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
siteNames := make(map[string]bool)
|
||||
allNames := make(map[string]bool)
|
||||
for name := range locals {
|
||||
siteNames[name] = true
|
||||
allNames[name] = true
|
||||
}
|
||||
for _, name := range missing {
|
||||
siteNames[name] = true
|
||||
allNames[name] = true
|
||||
}
|
||||
for name := range upstream {
|
||||
allNames[name] = true
|
||||
}
|
||||
names := make([]string, 0, len(allNames))
|
||||
for name := range allNames {
|
||||
names = append(names, name)
|
||||
}
|
||||
sort.Strings(names)
|
||||
vendorByID := make(map[int]*model.Vendor)
|
||||
for _, vendor := range vendors {
|
||||
vendorByID[vendor.Id] = vendor
|
||||
}
|
||||
candidates := make([]metadataSyncCandidate, 0, len(names))
|
||||
for _, name := range names {
|
||||
candidate := metadataSyncCandidate{ModelName: name, Scope: "catalog", Kind: "create", Fields: []metadataSyncField{}}
|
||||
if siteNames[name] {
|
||||
candidate.Scope = "site"
|
||||
}
|
||||
local := locals[name]
|
||||
up, found := upstream[name]
|
||||
if !found {
|
||||
candidate.Kind = "missing_upstream"
|
||||
candidates = append(candidates, candidate)
|
||||
continue
|
||||
}
|
||||
candidate.Upstream = &up
|
||||
var localVendor *model.Vendor
|
||||
if local != nil {
|
||||
localVendor = vendorByID[local.VendorID]
|
||||
}
|
||||
candidate.RecordVersion = model.MetadataRecordVersion(local, localVendor, model.FindMetadataVendor(vendors, up.Vendor))
|
||||
if local != nil && local.SyncOfficial == 0 {
|
||||
candidate.Kind = "blocked"
|
||||
candidates = append(candidates, candidate)
|
||||
continue
|
||||
}
|
||||
if up.Vendor != "" && model.FindMetadataVendor(vendors, up.Vendor) == nil {
|
||||
if _, exists := upstreamVendors[up.Vendor]; !exists {
|
||||
candidate.Kind = "missing_vendor"
|
||||
candidates = append(candidates, candidate)
|
||||
continue
|
||||
}
|
||||
candidate.VendorToCreate = up.Vendor
|
||||
}
|
||||
localValues := model.MetadataValues{}
|
||||
if local != nil {
|
||||
candidate.Kind = "update"
|
||||
localValues = model.MetadataValues{Description: local.Description, Icon: local.Icon, Tags: local.Tags, Endpoints: local.Endpoints, NameRule: local.NameRule, Status: local.Status}
|
||||
if localVendor != nil {
|
||||
localValues.Vendor = localVendor.Name
|
||||
}
|
||||
}
|
||||
localRaw, _ := common.Marshal(localValues)
|
||||
upRaw, _ := common.Marshal(up)
|
||||
var localFields, upFields map[string]any
|
||||
_ = common.Unmarshal(localRaw, &localFields)
|
||||
_ = common.Unmarshal(upRaw, &upFields)
|
||||
for _, field := range model.MetadataSyncFields {
|
||||
if local == nil || localFields[field] != upFields[field] {
|
||||
candidate.Fields = append(candidate.Fields, metadataSyncField{Field: field, Local: localFields[field], Upstream: upFields[field]})
|
||||
}
|
||||
}
|
||||
if local != nil && len(candidate.Fields) == 0 {
|
||||
candidate.Kind = "unchanged"
|
||||
}
|
||||
candidates = append(candidates, candidate)
|
||||
}
|
||||
common.ApiSuccess(c, gin.H{"source": source, "candidates": candidates})
|
||||
}
|
||||
|
||||
func SyncUpstreamModels(c *gin.Context) {
|
||||
var request struct {
|
||||
Locale string `json:"locale"`
|
||||
SourceVersion string `json:"source_version"`
|
||||
Selections []model.MetadataSyncSelection `json:"selections"`
|
||||
}
|
||||
if err := common.DecodeJson(c.Request.Body, &request); err != nil || len(request.Selections) == 0 || request.SourceVersion == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "Preview and select metadata changes before applying"})
|
||||
return
|
||||
}
|
||||
source, upstream, vendors, err := fetchMetadataCatalog(c, request.Locale)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if source.Version != request.SourceVersion {
|
||||
c.JSON(http.StatusConflict, gin.H{"success": false, "message": "Upstream metadata changed; preview again"})
|
||||
return
|
||||
}
|
||||
updates := make([]model.MetadataSyncUpdate, 0, len(request.Selections))
|
||||
for _, selection := range request.Selections {
|
||||
values, exists := upstream[selection.ModelName]
|
||||
if !exists {
|
||||
c.JSON(http.StatusConflict, gin.H{"success": false, "message": "Selected upstream model is no longer available"})
|
||||
return
|
||||
}
|
||||
updates = append(updates, model.MetadataSyncUpdate{MetadataSyncSelection: selection, Values: values})
|
||||
}
|
||||
result, err := model.ApplyMetadataSync(updates, vendors)
|
||||
if err != nil {
|
||||
status := http.StatusBadRequest
|
||||
if errors.Is(err, model.ErrMetadataSyncConflict) {
|
||||
status = http.StatusConflict
|
||||
}
|
||||
c.JSON(status, gin.H{"success": false, "message": err.Error()})
|
||||
return
|
||||
}
|
||||
recordManageAudit(c, "model.metadata.sync", map[string]any{"created_models": result.CreatedModels, "updated_models": result.UpdatedModels, "created_vendors": result.CreatedVendors})
|
||||
common.ApiSuccess(c, result)
|
||||
}
|
||||
@@ -0,0 +1,589 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/i18n"
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/oauth"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
const oauthAuthFlowTTL = 10 * time.Minute
|
||||
|
||||
type oauthStateRequest struct {
|
||||
Provider string `json:"provider"`
|
||||
Intent string `json:"intent"`
|
||||
Aff string `json:"aff,omitempty"`
|
||||
Scope string `json:"scope,omitempty"`
|
||||
Context json.RawMessage `json:"context,omitempty"`
|
||||
}
|
||||
|
||||
type oauthFlowPayload struct {
|
||||
AffiliateCode string `json:"affiliate_code,omitempty"`
|
||||
Verification *service.OAuthVerificationFlow `json:"verification,omitempty"`
|
||||
Telegram *oauth.TelegramOAuthFlow `json:"telegram,omitempty"`
|
||||
SessionIdentity *service.AuthIdentity `json:"session_identity,omitempty"`
|
||||
Authorization *model.AuthFlowAuthorization `json:"authorization,omitempty"`
|
||||
}
|
||||
|
||||
// providerParams returns map with Provider key for i18n templates
|
||||
func providerParams(name string) map[string]any {
|
||||
return map[string]any{"Provider": name}
|
||||
}
|
||||
|
||||
// GenerateOAuthCode generates a state code for OAuth CSRF protection
|
||||
func GenerateOAuthCode(c *gin.Context) {
|
||||
var request oauthStateRequest
|
||||
if err := common.DecodeJson(c.Request.Body, &request); err != nil {
|
||||
common.ApiErrorI18n(c, i18n.MsgInvalidParams)
|
||||
return
|
||||
}
|
||||
request.Provider = strings.TrimSpace(request.Provider)
|
||||
request.Intent = strings.TrimSpace(request.Intent)
|
||||
request.Aff = strings.TrimSpace(request.Aff)
|
||||
if oauth.GetProvider(request.Provider) == nil ||
|
||||
(request.Intent != model.AuthFlowIntentLogin && request.Intent != model.AuthFlowIntentBind && request.Intent != model.AuthFlowIntentVerify) ||
|
||||
len(request.Aff) > 32 ||
|
||||
(request.Intent != model.AuthFlowIntentLogin && request.Aff != "") ||
|
||||
(request.Intent != model.AuthFlowIntentVerify && (request.Scope != "" || len(request.Context) != 0)) {
|
||||
common.ApiErrorI18n(c, i18n.MsgInvalidParams)
|
||||
return
|
||||
}
|
||||
userID := 0
|
||||
sessionID := ""
|
||||
flowPayload := oauthFlowPayload{AffiliateCode: request.Aff}
|
||||
bindingStarted := false
|
||||
if request.Provider == "telegram" {
|
||||
telegramFlow, err := oauth.NewTelegramOAuthFlow()
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
flowPayload.Telegram = telegramFlow
|
||||
}
|
||||
if request.Intent == model.AuthFlowIntentBind || request.Intent == model.AuthFlowIntentVerify {
|
||||
identity, ok := middleware.GetSessionAuthIdentity(c)
|
||||
if !ok {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"success": false, "message": "绑定操作需要登录"})
|
||||
return
|
||||
}
|
||||
userID = identity.UserID
|
||||
sessionID = identity.SessionID
|
||||
if request.Intent == model.AuthFlowIntentBind {
|
||||
defer func() {
|
||||
recordUserSecurityAudit(c, userID, "user.binding_start", map[string]any{"provider": request.Provider, "success": bindingStarted})
|
||||
}()
|
||||
context, err := common.Marshal(service.AccountBindingContext{Provider: request.Provider})
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
flowPayload.Authorization = middleware.RequireSecurityProof(c, service.VerificationOperation{Scope: service.VerificationScopeAccountBind, Context: context})
|
||||
if flowPayload.Authorization == nil {
|
||||
return
|
||||
}
|
||||
flowPayload.SessionIdentity = &identity
|
||||
}
|
||||
if flowPayload.Telegram != nil {
|
||||
if _, _, err := service.ValidateLoginSession(identity); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
flowPayload.SessionIdentity = &identity
|
||||
}
|
||||
if request.Intent == model.AuthFlowIntentVerify {
|
||||
verification, err := service.StartOAuthVerification(identity, service.VerificationOperation{Scope: request.Scope, Context: request.Context}, request.Provider)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
flowPayload.Verification = verification
|
||||
}
|
||||
}
|
||||
payload, err := common.Marshal(flowPayload)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
expiresAt := time.Now().Add(oauthAuthFlowTTL)
|
||||
state, _, err := model.CreateAuthFlow(model.AuthFlowCreate{
|
||||
Purpose: model.AuthFlowPurposeOAuth,
|
||||
Provider: request.Provider,
|
||||
Intent: request.Intent,
|
||||
UserId: userID,
|
||||
SessionId: sessionID,
|
||||
Payload: string(payload),
|
||||
ExpiresAt: expiresAt,
|
||||
})
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
bindingStarted = request.Intent == model.AuthFlowIntentBind
|
||||
data := gin.H{"flow_token": state, "expires_at": expiresAt.Unix()}
|
||||
if flowPayload.Telegram != nil {
|
||||
data["authorization_url"] = flowPayload.Telegram.AuthorizationURL(state)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": data,
|
||||
})
|
||||
}
|
||||
|
||||
// HandleOAuth handles OAuth callback for all standard OAuth providers
|
||||
func HandleOAuth(c *gin.Context) {
|
||||
providerName := c.Param("provider")
|
||||
provider := oauth.GetProvider(providerName)
|
||||
if provider == nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"success": false,
|
||||
"message": i18n.T(c, i18n.MsgOAuthUnknownProvider),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// 1. Validate state (CSRF protection)
|
||||
state := c.Query("state")
|
||||
pendingFlow, err := model.GetAuthFlow(state, model.AuthFlowMatch{
|
||||
Purpose: model.AuthFlowPurposeOAuth,
|
||||
Provider: providerName,
|
||||
})
|
||||
if err != nil {
|
||||
c.JSON(http.StatusForbidden, gin.H{
|
||||
"success": false,
|
||||
"message": i18n.T(c, i18n.MsgOAuthStateInvalid),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
consumeMatch := model.AuthFlowMatch{
|
||||
Purpose: model.AuthFlowPurposeOAuth,
|
||||
Provider: providerName,
|
||||
Intent: pendingFlow.Intent,
|
||||
}
|
||||
bindSucceeded, notificationFailed := false, false
|
||||
if pendingFlow.Intent == model.AuthFlowIntentBind {
|
||||
defer func() {
|
||||
recordUserSecurityAudit(c, pendingFlow.UserId, "user.binding_bind", map[string]any{"provider": providerName, "success": bindSucceeded, "notification_failed": notificationFailed})
|
||||
}()
|
||||
}
|
||||
// Bind and verification callbacks must use the dashboard session that started them.
|
||||
if pendingFlow.Intent == model.AuthFlowIntentBind || pendingFlow.Intent == model.AuthFlowIntentVerify {
|
||||
identity, ok := middleware.GetSessionAuthIdentity(c)
|
||||
if !ok || identity.UserID != pendingFlow.UserId || identity.SessionID != pendingFlow.SessionId {
|
||||
c.JSON(http.StatusForbidden, gin.H{
|
||||
"success": false,
|
||||
"message": i18n.T(c, i18n.MsgOAuthStateInvalid),
|
||||
})
|
||||
return
|
||||
}
|
||||
consumeMatch.UserId = identity.UserID
|
||||
consumeMatch.SessionId = identity.SessionID
|
||||
if pendingFlow.Intent == model.AuthFlowIntentBind {
|
||||
var payload oauthFlowPayload
|
||||
if err := common.UnmarshalJsonStr(pendingFlow.Payload, &payload); err != nil {
|
||||
writeSecurityOperationError(c, model.ErrAuthFlowInvalid)
|
||||
return
|
||||
}
|
||||
context, err := common.Marshal(service.AccountBindingContext{Provider: providerName})
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
if err := service.ValidateFlowAuthorization(identity, service.VerificationOperation{Scope: service.VerificationScopeAccountBind, Context: context}, payload.Authorization); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
} else if pendingFlow.Intent != model.AuthFlowIntentLogin {
|
||||
common.ApiErrorI18n(c, i18n.MsgInvalidParams)
|
||||
return
|
||||
}
|
||||
|
||||
// 3. Check if provider is enabled
|
||||
var telegramPayload oauthFlowPayload
|
||||
if providerName == "telegram" {
|
||||
if err := oauth.TelegramConfigurationError(); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
if err := common.UnmarshalJsonStr(pendingFlow.Payload, &telegramPayload); err != nil || telegramPayload.Telegram == nil {
|
||||
writeSecurityOperationError(c, model.ErrAuthFlowInvalid)
|
||||
return
|
||||
}
|
||||
if pendingFlow.Intent != model.AuthFlowIntentLogin {
|
||||
identity, _ := middleware.GetSessionAuthIdentity(c)
|
||||
if telegramPayload.SessionIdentity == nil || *telegramPayload.SessionIdentity != identity {
|
||||
writeSecurityOperationError(c, model.ErrAuthFlowInvalid)
|
||||
return
|
||||
}
|
||||
if _, _, err := service.ValidateLoginSession(identity); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
c.Set(oauth.TelegramOAuthFlowContextKey, telegramPayload.Telegram)
|
||||
}
|
||||
if !provider.IsEnabled() {
|
||||
common.ApiErrorI18n(c, i18n.MsgOAuthNotEnabled, providerParams(provider.GetName()))
|
||||
return
|
||||
}
|
||||
|
||||
// 4. Handle error from provider
|
||||
errorCode := c.Query("error")
|
||||
if errorCode != "" {
|
||||
if _, err := model.ConsumeAuthFlow(state, consumeMatch); err != nil {
|
||||
c.JSON(http.StatusForbidden, gin.H{"success": false, "message": i18n.T(c, i18n.MsgOAuthStateInvalid)})
|
||||
return
|
||||
}
|
||||
errorDescription := c.Query("error_description")
|
||||
if errorDescription == "" {
|
||||
errorDescription = errorCode
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": errorDescription,
|
||||
})
|
||||
return
|
||||
}
|
||||
// 5. Exchange code for token
|
||||
code := c.Query("code")
|
||||
token, err := provider.ExchangeToken(c.Request.Context(), code, c)
|
||||
if err != nil {
|
||||
if providerName == "telegram" {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
handleOAuthError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
// 6. Get user info
|
||||
oauthUser, err := provider.GetUserInfo(c.Request.Context(), token)
|
||||
if err != nil {
|
||||
if providerName == "telegram" {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
handleOAuthError(c, err)
|
||||
return
|
||||
}
|
||||
if pendingFlow.Intent == model.AuthFlowIntentBind {
|
||||
bindSucceeded, notificationFailed = handleOAuthBind(c, providerName, provider, oauthUser, pendingFlow, state, consumeMatch)
|
||||
return
|
||||
}
|
||||
flow, err := model.ConsumeAuthFlow(state, consumeMatch)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusForbidden, gin.H{"success": false, "message": i18n.T(c, i18n.MsgOAuthStateInvalid)})
|
||||
return
|
||||
}
|
||||
|
||||
switch flow.Intent {
|
||||
case model.AuthFlowIntentLogin:
|
||||
handleOAuthLogin(c, provider, oauthUser, flow)
|
||||
case model.AuthFlowIntentVerify:
|
||||
handleOAuthVerification(c, providerName, oauthUser, flow)
|
||||
}
|
||||
}
|
||||
|
||||
func handleOAuthVerification(c *gin.Context, provider string, oauthUser *oauth.OAuthUser, flow *model.AuthFlow) {
|
||||
var payload oauthFlowPayload
|
||||
if err := common.UnmarshalJsonStr(flow.Payload, &payload); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
identity, _ := middleware.GetSessionAuthIdentity(c)
|
||||
proof, err := service.FinishOAuthVerification(identity, provider, oauthUser.ProviderUserID, payload.Verification)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
recordUserSecurityAudit(c, identity.UserID, "user.security_verify", map[string]any{"method": proof.Method, "scope": proof.Scope, "provider": provider})
|
||||
common.ApiSuccess(c, proof)
|
||||
}
|
||||
|
||||
func handleOAuthLogin(c *gin.Context, provider oauth.Provider, oauthUser *oauth.OAuthUser, flow *model.AuthFlow) {
|
||||
// 7. Find or create user
|
||||
var payload oauthFlowPayload
|
||||
if err := common.UnmarshalJsonStr(flow.Payload, &payload); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
user, err := findOrCreateOAuthUser(c, provider, oauthUser, payload.AffiliateCode)
|
||||
if err != nil {
|
||||
if errors.Is(err, model.ErrEmailAlreadyTaken) {
|
||||
common.ApiErrorI18n(c, i18n.MsgUserEmailAlreadyTaken)
|
||||
return
|
||||
}
|
||||
switch err.(type) {
|
||||
case *OAuthUserDeletedError:
|
||||
common.ApiErrorI18n(c, i18n.MsgOAuthUserDeleted)
|
||||
case *OAuthRegistrationDisabledError:
|
||||
common.ApiErrorI18n(c, i18n.MsgUserRegisterDisabled)
|
||||
case *OAuthEmailAlreadyTakenError:
|
||||
common.ApiErrorI18n(c, i18n.MsgUserEmailAlreadyTaken)
|
||||
default:
|
||||
writeSecurityOperationError(c, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// 8. Check user status
|
||||
if user.Status != common.UserStatusEnabled {
|
||||
common.ApiErrorI18n(c, i18n.MsgOAuthUserBanned)
|
||||
return
|
||||
}
|
||||
|
||||
// 9. Setup login
|
||||
setupLogin(user, c)
|
||||
}
|
||||
|
||||
// handleOAuthBind handles binding OAuth account to existing user
|
||||
func handleOAuthBind(c *gin.Context, providerName string, provider oauth.Provider, oauthUser *oauth.OAuthUser, flow *model.AuthFlow, state string, match model.AuthFlowMatch) (bool, bool) {
|
||||
identity, ok := middleware.GetSessionAuthIdentity(c)
|
||||
if !ok {
|
||||
writeSecurityOperationError(c, service.ErrAuthTokenInvalid)
|
||||
return false, false
|
||||
}
|
||||
var payload oauthFlowPayload
|
||||
if err := common.UnmarshalJsonStr(flow.Payload, &payload); err != nil {
|
||||
writeSecurityOperationError(c, model.ErrAuthFlowInvalid)
|
||||
return false, false
|
||||
}
|
||||
context, err := common.Marshal(service.AccountBindingContext{Provider: providerName})
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return false, false
|
||||
}
|
||||
// Recheck after the external provider round trip, then validate the session
|
||||
// under the transaction's locks before consuming the flow and writing.
|
||||
if err := service.ValidateFlowAuthorization(identity, service.VerificationOperation{Scope: service.VerificationScopeAccountBind, Context: context}, payload.Authorization); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return false, false
|
||||
}
|
||||
if provider.IsUserIDTaken(oauthUser.ProviderUserID) {
|
||||
common.ApiErrorI18n(c, i18n.MsgOAuthAlreadyBound, providerParams(provider.GetName()))
|
||||
return false, false
|
||||
}
|
||||
if legacyID, ok := oauthUser.Extra["legacy_id"].(string); ok && legacyID != "" && provider.IsUserIDTaken(legacyID) {
|
||||
common.ApiErrorI18n(c, i18n.MsgOAuthAlreadyBound, providerParams(provider.GetName()))
|
||||
return false, false
|
||||
}
|
||||
_, err = model.ConsumeAuthFlowWithAction(state, match, func(tx *gorm.DB, _ *model.AuthFlow) error {
|
||||
if providerName == "telegram" {
|
||||
return model.BindTelegramForSessionWithTx(tx, identity, oauthUser.ProviderUserID)
|
||||
}
|
||||
if custom, ok := provider.(*oauth.GenericOAuthProvider); ok {
|
||||
return model.UpdateUserOAuthBindingForSessionWithTx(tx, identity, custom.GetProviderId(), oauthUser.ProviderUserID)
|
||||
}
|
||||
return model.UpdateUserBindColumnForSessionWithTx(tx, identity, provider.ProviderUserIDColumn(), oauthUser.ProviderUserID)
|
||||
})
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return false, false
|
||||
}
|
||||
user, err := model.GetUserById(identity.UserID, false)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return true, true
|
||||
}
|
||||
notificationFailed := service.NotifyAccountSecurityChange(user.Email, "Login account linked: "+provider.GetName()) != nil
|
||||
common.ApiSuccessI18n(c, i18n.MsgOAuthBindSuccess, gin.H{"action": "bind", "notification_warning": notificationFailed})
|
||||
return true, notificationFailed
|
||||
}
|
||||
|
||||
// findOrCreateOAuthUser finds existing user or creates new user
|
||||
func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *oauth.OAuthUser, affiliateCode string) (*model.User, error) {
|
||||
user := &model.User{}
|
||||
if provider.ProviderUserIDColumn() == "telegram_id" {
|
||||
err := provider.FillUserByProviderID(user, oauthUser.ProviderUserID)
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, oauth.ErrTelegramAccountNotBound
|
||||
}
|
||||
return user, err
|
||||
}
|
||||
|
||||
// Check if user already exists with new ID
|
||||
if provider.IsUserIDTaken(oauthUser.ProviderUserID) {
|
||||
err := provider.FillUserByProviderID(user, oauthUser.ProviderUserID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Check if user has been deleted
|
||||
if user.Id == 0 {
|
||||
return nil, &OAuthUserDeletedError{}
|
||||
}
|
||||
return user, nil
|
||||
}
|
||||
|
||||
// Try to find user with legacy ID (for GitHub migration from login to numeric ID)
|
||||
if legacyID, ok := oauthUser.Extra["legacy_id"].(string); ok && legacyID != "" {
|
||||
if provider.IsUserIDTaken(legacyID) {
|
||||
err := provider.FillUserByProviderID(user, legacyID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if user.Id != 0 {
|
||||
// Found user with legacy ID, migrate to new ID
|
||||
common.SysLog(fmt.Sprintf("[OAuth] Migrating user %d from legacy_id=%s to new_id=%s",
|
||||
user.Id, legacyID, oauthUser.ProviderUserID))
|
||||
if err := user.UpdateGitHubId(oauthUser.ProviderUserID); err != nil {
|
||||
common.SysError(fmt.Sprintf("[OAuth] Failed to migrate user %d: %s", user.Id, err.Error()))
|
||||
// Continue with login even if migration fails
|
||||
}
|
||||
return user, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// User doesn't exist, create new user if registration is enabled
|
||||
if !common.RegisterEnabled {
|
||||
return nil, &OAuthRegistrationDisabledError{}
|
||||
}
|
||||
|
||||
// Set up new user
|
||||
user.Username = provider.GetProviderPrefix() + strconv.Itoa(model.GetMaxUserId()+1)
|
||||
|
||||
if oauthUser.Username != "" {
|
||||
if exists, err := model.CheckUserExistOrDeleted(oauthUser.Username, ""); err == nil && !exists {
|
||||
// 防止索引退化
|
||||
if len(oauthUser.Username) <= model.UserNameMaxLength {
|
||||
user.Username = oauthUser.Username
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if oauthUser.DisplayName != "" {
|
||||
user.DisplayName = oauthUser.DisplayName
|
||||
} else if oauthUser.Username != "" {
|
||||
user.DisplayName = oauthUser.Username
|
||||
} else {
|
||||
user.DisplayName = provider.GetName() + " User"
|
||||
}
|
||||
if oauthUser.Email != "" {
|
||||
user.Email = model.NormalizeEmail(oauthUser.Email)
|
||||
if err := model.EnsureEmailAvailable(user.Email, 0); err != nil {
|
||||
if errors.Is(err, model.ErrEmailAlreadyTaken) {
|
||||
return nil, &OAuthEmailAlreadyTakenError{}
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
user.Role = common.RoleCommonUser
|
||||
user.Status = common.UserStatusEnabled
|
||||
|
||||
// Handle affiliate code
|
||||
inviterId := 0
|
||||
if affiliateCode != "" {
|
||||
inviterId, _ = model.GetUserIdByAffCode(affiliateCode)
|
||||
}
|
||||
|
||||
// Use transaction to ensure user creation and OAuth binding are atomic
|
||||
if genericProvider, ok := provider.(*oauth.GenericOAuthProvider); ok {
|
||||
// Custom provider: create user and binding in a transaction
|
||||
err := model.DB.Transaction(func(tx *gorm.DB) error {
|
||||
// Create user
|
||||
if err := user.InsertWithTx(tx, inviterId); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Create OAuth binding
|
||||
binding := &model.UserOAuthBinding{
|
||||
UserId: user.Id,
|
||||
ProviderId: genericProvider.GetProviderId(),
|
||||
ProviderUserId: oauthUser.ProviderUserID,
|
||||
}
|
||||
if err := model.CreateUserOAuthBindingWithTx(tx, binding); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Perform post-transaction tasks (logs, sidebar config, inviter rewards)
|
||||
user.FinalizeOAuthUserCreation(inviterId)
|
||||
} else {
|
||||
// Built-in provider: create user and update provider ID in a transaction
|
||||
err := model.DB.Transaction(func(tx *gorm.DB) error {
|
||||
// Create user
|
||||
if err := user.InsertWithTx(tx, inviterId); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Set the provider user ID on the user model and update
|
||||
provider.SetProviderUserID(user, oauthUser.ProviderUserID)
|
||||
if err := tx.Model(user).Updates(map[string]any{
|
||||
"github_id": user.GitHubId,
|
||||
"discord_id": user.DiscordId,
|
||||
"oidc_id": user.OidcId,
|
||||
"linux_do_id": user.LinuxDOId,
|
||||
"wechat_id": user.WeChatId,
|
||||
"telegram_id": user.TelegramId,
|
||||
}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Perform post-transaction tasks
|
||||
user.FinalizeOAuthUserCreation(inviterId)
|
||||
}
|
||||
|
||||
return user, nil
|
||||
}
|
||||
|
||||
// Error types for OAuth
|
||||
type OAuthUserDeletedError struct{}
|
||||
|
||||
func (e *OAuthUserDeletedError) Error() string {
|
||||
return "user has been deleted"
|
||||
}
|
||||
|
||||
type OAuthRegistrationDisabledError struct{}
|
||||
|
||||
func (e *OAuthRegistrationDisabledError) Error() string {
|
||||
return "registration is disabled"
|
||||
}
|
||||
|
||||
type OAuthEmailAlreadyTakenError struct{}
|
||||
|
||||
func (e *OAuthEmailAlreadyTakenError) Error() string {
|
||||
return "email is already in use"
|
||||
}
|
||||
|
||||
// handleOAuthError handles OAuth errors and returns translated message
|
||||
func handleOAuthError(c *gin.Context, err error) {
|
||||
switch e := err.(type) {
|
||||
case *oauth.OAuthError:
|
||||
if e.Params != nil {
|
||||
common.ApiErrorI18n(c, e.MsgKey, e.Params)
|
||||
} else {
|
||||
common.ApiErrorI18n(c, e.MsgKey)
|
||||
}
|
||||
case *oauth.AccessDeniedError:
|
||||
common.ApiErrorMsg(c, e.Message)
|
||||
case *oauth.TrustLevelError:
|
||||
common.ApiErrorI18n(c, i18n.MsgOAuthTrustLevelLow)
|
||||
default:
|
||||
writeSecurityOperationError(c, err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,431 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"slices"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/i18n"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting"
|
||||
"github.com/QuantumNous/new-api/setting/billing_setting"
|
||||
"github.com/QuantumNous/new-api/setting/console_setting"
|
||||
"github.com/QuantumNous/new-api/setting/model_setting"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
var completionRatioMetaOptionKeys = []string{
|
||||
"ModelPrice",
|
||||
"ModelRatio",
|
||||
"CompletionRatio",
|
||||
"CacheRatio",
|
||||
"CreateCacheRatio",
|
||||
"ImageRatio",
|
||||
"AudioRatio",
|
||||
"AudioCompletionRatio",
|
||||
}
|
||||
|
||||
func isPaymentComplianceOptionKey(key string) bool {
|
||||
return strings.HasPrefix(key, "payment_setting.compliance_")
|
||||
}
|
||||
|
||||
func isPositiveOptionValue(value string) bool {
|
||||
intValue, err := strconv.Atoi(strings.TrimSpace(value))
|
||||
if err == nil {
|
||||
return intValue > 0
|
||||
}
|
||||
floatValue, err := strconv.ParseFloat(strings.TrimSpace(value), 64)
|
||||
return err == nil && floatValue > 0
|
||||
}
|
||||
|
||||
func collectModelNamesFromOptionValue(raw string, modelNames map[string]struct{}) {
|
||||
if strings.TrimSpace(raw) == "" {
|
||||
return
|
||||
}
|
||||
|
||||
var parsed map[string]any
|
||||
if err := common.UnmarshalJsonStr(raw, &parsed); err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
for modelName := range parsed {
|
||||
modelNames[modelName] = struct{}{}
|
||||
}
|
||||
}
|
||||
|
||||
func buildCompletionRatioMetaValue(optionValues map[string]string) string {
|
||||
modelNames := make(map[string]struct{})
|
||||
for _, key := range completionRatioMetaOptionKeys {
|
||||
collectModelNamesFromOptionValue(optionValues[key], modelNames)
|
||||
}
|
||||
|
||||
meta := make(map[string]ratio_setting.CompletionRatioInfo, len(modelNames))
|
||||
for modelName := range modelNames {
|
||||
meta[modelName] = ratio_setting.GetCompletionRatioInfo(modelName)
|
||||
}
|
||||
|
||||
jsonBytes, err := common.Marshal(meta)
|
||||
if err != nil {
|
||||
return "{}"
|
||||
}
|
||||
return string(jsonBytes)
|
||||
}
|
||||
|
||||
func GetOptions(c *gin.Context) {
|
||||
var options []*model.Option
|
||||
optionValues := make(map[string]string)
|
||||
common.OptionMapRWMutex.Lock()
|
||||
for k, v := range common.OptionMap {
|
||||
if k == "theme.frontend" || k == "billing_setting.billing_mode" || k == "billing_setting.billing_expr" {
|
||||
continue
|
||||
}
|
||||
value := common.Interface2String(v)
|
||||
isSensitiveKey := strings.HasSuffix(k, "Token") ||
|
||||
strings.HasSuffix(k, "Secret") ||
|
||||
strings.HasSuffix(k, "Key") ||
|
||||
strings.HasSuffix(k, "secret") ||
|
||||
strings.HasSuffix(k, "api_key")
|
||||
if isSensitiveKey {
|
||||
continue
|
||||
}
|
||||
options = append(options, &model.Option{
|
||||
Key: k,
|
||||
Value: value,
|
||||
})
|
||||
if slices.Contains(completionRatioMetaOptionKeys, k) {
|
||||
optionValues[k] = value
|
||||
}
|
||||
}
|
||||
common.OptionMapRWMutex.Unlock()
|
||||
// Display the same effective expressions used by pricing and settlement,
|
||||
// including built-in defaults absent from persisted administrator options.
|
||||
for key, values := range map[string]map[string]string{
|
||||
"billing_setting.billing_mode": billing_setting.GetBillingModeCopy(),
|
||||
"billing_setting.billing_expr": billing_setting.GetBillingExprCopy(),
|
||||
} {
|
||||
encoded, err := common.Marshal(values)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"success": false, "message": err.Error()})
|
||||
return
|
||||
}
|
||||
options = append(options, &model.Option{Key: key, Value: string(encoded)})
|
||||
}
|
||||
options = append(options, &model.Option{
|
||||
Key: "CompletionRatioMeta",
|
||||
Value: buildCompletionRatioMetaValue(optionValues),
|
||||
})
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": options,
|
||||
})
|
||||
}
|
||||
|
||||
type OptionUpdateRequest struct {
|
||||
Key string `json:"key"`
|
||||
Value any `json:"value"`
|
||||
}
|
||||
|
||||
func UpdateOption(c *gin.Context) {
|
||||
var option OptionUpdateRequest
|
||||
err := common.DecodeJson(c.Request.Body, &option)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"success": false,
|
||||
"message": "无效的参数",
|
||||
})
|
||||
return
|
||||
}
|
||||
switch option.Value.(type) {
|
||||
case bool:
|
||||
option.Value = common.Interface2String(option.Value.(bool))
|
||||
case float64:
|
||||
option.Value = common.Interface2String(option.Value.(float64))
|
||||
case int:
|
||||
option.Value = common.Interface2String(option.Value.(int))
|
||||
default:
|
||||
option.Value = fmt.Sprintf("%v", option.Value)
|
||||
}
|
||||
switch option.Key {
|
||||
case "QuotaForInviter", "QuotaForInvitee":
|
||||
if isPositiveOptionValue(option.Value.(string)) && !operation_setting.IsPaymentComplianceConfirmed() {
|
||||
common.ApiErrorI18n(c, i18n.MsgPaymentComplianceRequired)
|
||||
return
|
||||
}
|
||||
default:
|
||||
if isPaymentComplianceOptionKey(option.Key) {
|
||||
common.ApiErrorMsg(c, "合规确认字段不允许通过通用设置接口修改")
|
||||
return
|
||||
}
|
||||
}
|
||||
if option.Key == "TaskPublicAddress" && option.Value.(string) != "" {
|
||||
if err := service.ValidateTaskArtifactBaseURL(option.Value.(string)); err != nil {
|
||||
common.ApiErrorMsg(c, err.Error())
|
||||
return
|
||||
}
|
||||
}
|
||||
switch option.Key {
|
||||
case "GitHubOAuthEnabled":
|
||||
if option.Value == "true" && common.GitHubClientId == "" {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "无法启用 GitHub OAuth，请先填入 GitHub Client Id 以及 GitHub Client Secret！",
|
||||
})
|
||||
return
|
||||
}
|
||||
case "discord.enabled":
|
||||
if option.Value == "true" && system_setting.GetDiscordSettings().ClientId == "" {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "无法启用 Discord OAuth，请先填入 Discord Client Id 以及 Discord Client Secret！",
|
||||
})
|
||||
return
|
||||
}
|
||||
case "oidc.enabled":
|
||||
if option.Value == "true" && system_setting.GetOIDCSettings().ClientId == "" {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "无法启用 OIDC 登录，请先填入 OIDC Client Id 以及 OIDC Client Secret！",
|
||||
})
|
||||
return
|
||||
}
|
||||
case "LinuxDOOAuthEnabled":
|
||||
if option.Value == "true" && common.LinuxDOClientId == "" {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "无法启用 LinuxDO OAuth，请先填入 LinuxDO Client Id 以及 LinuxDO Client Secret！",
|
||||
})
|
||||
return
|
||||
}
|
||||
case "EmailDomainRestrictionEnabled":
|
||||
if option.Value == "true" && len(common.EmailDomainWhitelist) == 0 {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "无法启用邮箱域名限制，请先填入限制的邮箱域名！",
|
||||
})
|
||||
return
|
||||
}
|
||||
case "WeChatAuthEnabled":
|
||||
if option.Value == "true" && common.WeChatServerAddress == "" {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "无法启用微信登录，请先填入微信登录相关配置信息！",
|
||||
})
|
||||
return
|
||||
}
|
||||
case "TurnstileCheckEnabled":
|
||||
if option.Value == "true" && common.TurnstileSiteKey == "" {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "无法启用 Turnstile 校验，请先填入 Turnstile 校验相关配置信息！",
|
||||
})
|
||||
|
||||
return
|
||||
}
|
||||
case "TelegramOAuthEnabled":
|
||||
if option.Value == "true" && !system_setting.GetTelegramSettings().IsConfigured() {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"code": "TELEGRAM_OAUTH_NOT_CONFIGURED",
|
||||
"message": "Telegram OAuth is not configured or enabled. Please contact your administrator.",
|
||||
})
|
||||
return
|
||||
}
|
||||
case "theme.frontend":
|
||||
if option.Value != "default" {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "Classic 前端已移除，主题只能设置为 default",
|
||||
})
|
||||
return
|
||||
}
|
||||
case "GroupRatio":
|
||||
err = ratio_setting.CheckGroupRatio(option.Value.(string))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
case "gemini.safety_settings":
|
||||
err = model_setting.ValidateGeminiSafetySettings(option.Value.(string))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
case "claude.default_max_tokens":
|
||||
err = model_setting.ValidateClaudeDefaultMaxTokens(option.Value.(string))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
case operation_setting.ToolPriceOptionKey:
|
||||
err = operation_setting.ValidateToolPricesJSON(option.Value.(string))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
case "ImageRatio":
|
||||
err = ratio_setting.UpdateImageRatioByJSONString(option.Value.(string))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "图片倍率设置失败: " + err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
case "AudioRatio":
|
||||
err = ratio_setting.UpdateAudioRatioByJSONString(option.Value.(string))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "音频倍率设置失败: " + err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
case "AudioCompletionRatio":
|
||||
err = ratio_setting.UpdateAudioCompletionRatioByJSONString(option.Value.(string))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "音频补全倍率设置失败: " + err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
case "CreateCacheRatio":
|
||||
err = ratio_setting.UpdateCreateCacheRatioByJSONString(option.Value.(string))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "缓存创建倍率设置失败: " + err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
case "ModelRequestRateLimitGroup":
|
||||
err = setting.CheckModelRequestRateLimitGroup(option.Value.(string))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
case "AutomaticDisableStatusCodes":
|
||||
_, err = operation_setting.ParseHTTPStatusCodeRanges(option.Value.(string))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
case "AutomaticRetryStatusCodes":
|
||||
_, err = operation_setting.ParseHTTPStatusCodeRanges(option.Value.(string))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
case "billing_setting.billing_expr":
|
||||
expressions := make(map[string]string)
|
||||
if err = common.UnmarshalJsonStr(option.Value.(string), &expressions); err != nil {
|
||||
common.ApiErrorMsg(c, "计费表达式配置必须是模型到表达式的 JSON 对象: "+err.Error())
|
||||
return
|
||||
}
|
||||
models := make([]string, 0, len(expressions))
|
||||
for modelName := range expressions {
|
||||
models = append(models, modelName)
|
||||
}
|
||||
sort.Strings(models)
|
||||
generation := jsplugin.DefaultRegistry.Generation()
|
||||
for _, modelName := range models {
|
||||
expression := expressions[modelName]
|
||||
if plugin, ok := generation.GetByModel(modelName); ok {
|
||||
err = billing_setting.SmokeTestTaskExpr(expression, plugin.Meta.UsageSchema)
|
||||
} else if target, resolved := model.ResolveTaskModelAlias(generation, modelName); resolved {
|
||||
if plugin, ok := generation.Get(target.PluginKey); ok {
|
||||
err = billing_setting.SmokeTestTaskExpr(expression, plugin.Meta.UsageSchema)
|
||||
} else {
|
||||
err = billing_setting.SmokeTestExpr(expression)
|
||||
}
|
||||
} else {
|
||||
err = billing_setting.SmokeTestExpr(expression)
|
||||
}
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, fmt.Sprintf("模型 %s 的计费表达式无效: %v", modelName, err))
|
||||
return
|
||||
}
|
||||
}
|
||||
case "console_setting.api_info":
|
||||
err = console_setting.ValidateConsoleSettings(option.Value.(string), "ApiInfo")
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
case "console_setting.announcements":
|
||||
err = console_setting.ValidateConsoleSettings(option.Value.(string), "Announcements")
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
case "console_setting.faq":
|
||||
err = console_setting.ValidateConsoleSettings(option.Value.(string), "FAQ")
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
case "console_setting.uptime_kuma_groups":
|
||||
err = console_setting.ValidateConsoleSettings(option.Value.(string), "UptimeKumaGroups")
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
}
|
||||
err = model.UpdateOption(option.Key, option.Value.(string))
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
// 出于安全考虑只记录被修改的配置项名称，不记录配置值（可能含密钥等敏感信息）。
|
||||
recordManageAudit(c, "option.update", map[string]any{
|
||||
"key": option.Key,
|
||||
})
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestUpdateOptionRejectsNegativeClaudeDefaultMaxTokens(t *testing.T) {
|
||||
response := httptest.NewRecorder()
|
||||
context, _ := gin.CreateTestContext(response)
|
||||
context.Request = httptest.NewRequest(
|
||||
http.MethodPut,
|
||||
"/api/option/",
|
||||
strings.NewReader(`{"key":"claude.default_max_tokens","value":"{\"default\":-1}"}`),
|
||||
)
|
||||
|
||||
UpdateOption(context)
|
||||
|
||||
assert.Equal(t, http.StatusOK, response.Code)
|
||||
var payload struct {
|
||||
Success bool `json:"success"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &payload))
|
||||
assert.False(t, payload.Success)
|
||||
assert.Contains(t, payload.Message, "-1")
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestUpdateOptionRejectsInvalidGeminiSafetyThreshold(t *testing.T) {
|
||||
response := httptest.NewRecorder()
|
||||
context, _ := gin.CreateTestContext(response)
|
||||
context.Request = httptest.NewRequest(
|
||||
http.MethodPut,
|
||||
"/api/option/",
|
||||
strings.NewReader(`{"key":"gemini.safety_settings","value":"{\"default\":\"BLOCK_SOME\"}"}`),
|
||||
)
|
||||
|
||||
UpdateOption(context)
|
||||
|
||||
assert.Equal(t, http.StatusOK, response.Code)
|
||||
var payload struct {
|
||||
Success bool `json:"success"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &payload))
|
||||
assert.False(t, payload.Success)
|
||||
assert.Contains(t, payload.Message, "BLOCK_SOME")
|
||||
}
|
||||
@@ -0,0 +1,648 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
passkeysvc "github.com/QuantumNous/new-api/service/passkey"
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/go-webauthn/webauthn/protocol"
|
||||
webauthnlib "github.com/go-webauthn/webauthn/webauthn"
|
||||
)
|
||||
|
||||
type passkeyFinishRequest struct {
|
||||
FlowToken string `json:"flow_token"`
|
||||
Credential json.RawMessage `json:"credential"`
|
||||
}
|
||||
|
||||
type passkeyVerifyBeginRequest struct {
|
||||
Scope string `json:"scope"`
|
||||
Context json.RawMessage `json:"context,omitempty"`
|
||||
}
|
||||
|
||||
func parsePasskeyFinishRequest(c *gin.Context) (*passkeyFinishRequest, error) {
|
||||
var request passkeyFinishRequest
|
||||
if err := common.DecodeJson(c.Request.Body, &request); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if request.FlowToken == "" || len(request.Credential) == 0 {
|
||||
return nil, errors.New("Passkey 流程参数不完整")
|
||||
}
|
||||
return &request, nil
|
||||
}
|
||||
|
||||
func PasskeyRegisterBegin(c *gin.Context) {
|
||||
if !system_setting.GetPasskeySettings().Enabled {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "管理员未启用 Passkey 登录",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
user, err := getAuthenticatedUser(c)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
authorization := middleware.RequireSecurityProof(c, service.VerificationOperation{Scope: service.VerificationScopePasskeyRegister})
|
||||
if authorization == nil {
|
||||
return
|
||||
}
|
||||
|
||||
credential, err := model.GetPasskeyByUserID(user.Id)
|
||||
if err != nil && !errors.Is(err, model.ErrPasskeyNotFound) {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
if errors.Is(err, model.ErrPasskeyNotFound) {
|
||||
credential = nil
|
||||
}
|
||||
|
||||
wa, err := passkeysvc.BuildWebAuthn(c.Request)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
waUser := passkeysvc.NewWebAuthnUser(user, credential)
|
||||
selection := wa.Config.AuthenticatorSelection
|
||||
selection.UserVerification = protocol.VerificationRequired
|
||||
options := []webauthnlib.RegistrationOption{webauthnlib.WithAuthenticatorSelection(selection)}
|
||||
if credential != nil {
|
||||
descriptor := credential.ToWebAuthnCredential().Descriptor()
|
||||
options = append(options, webauthnlib.WithExclusions([]protocol.CredentialDescriptor{descriptor}))
|
||||
}
|
||||
|
||||
creation, sessionData, err := wa.BeginRegistration(waUser, options...)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
identity, ok := middleware.GetSessionAuthIdentity(c)
|
||||
if !ok {
|
||||
common.ApiErrorMsg(c, "当前认证方式不支持安全验证")
|
||||
return
|
||||
}
|
||||
flowToken, expiresAt, err := passkeysvc.CreateSessionDataFlow(
|
||||
model.AuthFlowPurposePasskeyRegister,
|
||||
passkeysvc.FlowSecurity{AuthSessionIdentity: identity, Scope: authorization.Scope, ContextHash: authorization.ContextHash, Authorization: authorization},
|
||||
sessionData,
|
||||
)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": gin.H{
|
||||
"options": creation,
|
||||
"flow_token": flowToken,
|
||||
"expires_at": expiresAt,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func PasskeyRegisterFinish(c *gin.Context) {
|
||||
if !system_setting.GetPasskeySettings().Enabled {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "管理员未启用 Passkey 登录",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
user, err := getAuthenticatedUser(c)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
request, err := parsePasskeyFinishRequest(c)
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, "无效的 Passkey 验证请求")
|
||||
return
|
||||
}
|
||||
parsedCredential, err := protocol.ParseCredentialCreationResponseBytes(request.Credential)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
wa, err := passkeysvc.BuildWebAuthn(c.Request)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
credentialRecord, err := model.GetPasskeyByUserID(user.Id)
|
||||
if err != nil && !errors.Is(err, model.ErrPasskeyNotFound) {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
if errors.Is(err, model.ErrPasskeyNotFound) {
|
||||
credentialRecord = nil
|
||||
}
|
||||
|
||||
identity, ok := middleware.GetSessionAuthIdentity(c)
|
||||
if !ok {
|
||||
common.ApiErrorMsg(c, "当前认证方式不支持安全验证")
|
||||
return
|
||||
}
|
||||
sessionData, security, err := passkeysvc.PopSessionDataFlow(
|
||||
request.FlowToken,
|
||||
model.AuthFlowPurposePasskeyRegister,
|
||||
identity,
|
||||
)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
if sessionData.UserVerification != protocol.VerificationRequired {
|
||||
writeSecurityOperationError(c, model.ErrAuthFlowInvalid)
|
||||
return
|
||||
}
|
||||
if err := service.ValidateFlowAuthorization(identity, service.VerificationOperation{Scope: service.VerificationScopePasskeyRegister}, security.Authorization); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
waUser := passkeysvc.NewWebAuthnUser(user, credentialRecord)
|
||||
credential, err := wa.CreateCredential(waUser, *sessionData, parsedCredential)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
passkeyCredential := model.NewPasskeyCredentialFromWebAuthn(user.Id, credential)
|
||||
if passkeyCredential == nil {
|
||||
common.ApiErrorMsg(c, "无法创建 Passkey 凭证")
|
||||
return
|
||||
}
|
||||
|
||||
if err := model.RegisterPasskeyForSession(identity, passkeyCredential); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
bundle, err := service.AdvanceCurrentSessionToUserVersion(identity, "passkey_registered")
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
recordUserSecurityAudit(c, user.Id, "user.passkey_register", nil)
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "Passkey 注册成功",
|
||||
"data": authRotationData(bundle),
|
||||
})
|
||||
}
|
||||
|
||||
func PasskeyDelete(c *gin.Context) {
|
||||
user, err := getAuthenticatedUser(c)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
if middleware.RequireSecurityProof(c, service.VerificationOperation{Scope: service.VerificationScopePasskeyDelete}) == nil {
|
||||
return
|
||||
}
|
||||
|
||||
identity, ok := middleware.GetSessionAuthIdentity(c)
|
||||
if !ok {
|
||||
common.ApiErrorMsg(c, "当前认证方式不支持安全验证")
|
||||
return
|
||||
}
|
||||
if err := model.DeletePasskeyForSession(identity); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
bundle, err := service.AdvanceCurrentSessionToUserVersion(identity, "passkey_deleted")
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
recordUserSecurityAudit(c, user.Id, "user.passkey_delete", nil)
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "Passkey 已解绑",
|
||||
"data": authRotationData(bundle),
|
||||
})
|
||||
}
|
||||
|
||||
func PasskeyStatus(c *gin.Context) {
|
||||
user, err := getAuthenticatedUser(c)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
credential, err := model.GetPasskeyByUserID(user.Id)
|
||||
if errors.Is(err, model.ErrPasskeyNotFound) {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": gin.H{
|
||||
"enabled": false,
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
data := gin.H{
|
||||
"enabled": true,
|
||||
"last_used_at": credential.LastUsedAt,
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": data,
|
||||
})
|
||||
}
|
||||
|
||||
func PasskeyLoginBegin(c *gin.Context) {
|
||||
if !system_setting.GetPasskeySettings().Enabled {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "管理员未启用 Passkey 登录",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
wa, err := passkeysvc.BuildWebAuthn(c.Request)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
assertion, sessionData, err := wa.BeginDiscoverableLogin(webauthnlib.WithUserVerification(protocol.VerificationRequired))
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
flowToken, expiresAt, err := passkeysvc.CreateSessionDataFlow(
|
||||
model.AuthFlowPurposePasskeyLogin,
|
||||
passkeysvc.FlowSecurity{},
|
||||
sessionData,
|
||||
)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": gin.H{
|
||||
"options": assertion,
|
||||
"flow_token": flowToken,
|
||||
"expires_at": expiresAt,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func PasskeyLoginFinish(c *gin.Context) {
|
||||
if !system_setting.GetPasskeySettings().Enabled {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "管理员未启用 Passkey 登录",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
request, err := parsePasskeyFinishRequest(c)
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, "无效的 Passkey 验证请求")
|
||||
return
|
||||
}
|
||||
parsedCredential, err := protocol.ParseCredentialRequestResponseBytes(request.Credential)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
wa, err := passkeysvc.BuildWebAuthn(c.Request)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
sessionData, _, err := passkeysvc.PopSessionDataFlow(
|
||||
request.FlowToken,
|
||||
model.AuthFlowPurposePasskeyLogin,
|
||||
service.AuthIdentity{},
|
||||
)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
if sessionData.UserVerification != protocol.VerificationRequired {
|
||||
writeSecurityOperationError(c, model.ErrAuthFlowInvalid)
|
||||
return
|
||||
}
|
||||
|
||||
handler := func(rawID, userHandle []byte) (webauthnlib.User, error) {
|
||||
// 首先通过凭证ID查找用户
|
||||
credential, err := model.GetPasskeyByCredentialID(rawID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("未找到 Passkey 凭证: %w", err)
|
||||
}
|
||||
|
||||
// 通过凭证获取用户
|
||||
user := &model.User{Id: credential.UserID}
|
||||
if err := user.FillUserById(); err != nil {
|
||||
return nil, fmt.Errorf("用户信息获取失败: %w", err)
|
||||
}
|
||||
|
||||
if user.Status != common.UserStatusEnabled {
|
||||
return nil, model.ErrUserSessionInactive
|
||||
}
|
||||
|
||||
if len(userHandle) > 0 {
|
||||
userID, parseErr := strconv.Atoi(string(userHandle))
|
||||
if parseErr != nil {
|
||||
// 记录异常但继续验证，因为某些客户端可能使用非数字格式
|
||||
common.SysLog(fmt.Sprintf("PasskeyLogin: userHandle parse error for credential, length: %d", len(userHandle)))
|
||||
} else if userID != user.Id {
|
||||
return nil, errors.New("用户句柄与凭证不匹配")
|
||||
}
|
||||
}
|
||||
|
||||
return passkeysvc.NewWebAuthnUser(user, credential), nil
|
||||
}
|
||||
|
||||
waUser, credential, err := wa.ValidatePasskeyLogin(handler, *sessionData, parsedCredential)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
userWrapper, ok := waUser.(*passkeysvc.WebAuthnUser)
|
||||
if !ok {
|
||||
common.ApiErrorMsg(c, "Passkey 登录状态异常")
|
||||
return
|
||||
}
|
||||
|
||||
modelUser := userWrapper.ModelUser()
|
||||
if modelUser == nil {
|
||||
common.ApiErrorMsg(c, "Passkey 登录状态异常")
|
||||
return
|
||||
}
|
||||
|
||||
if modelUser.Status != common.UserStatusEnabled {
|
||||
common.ApiErrorMsg(c, "该用户已被禁用")
|
||||
return
|
||||
}
|
||||
|
||||
if err := model.UpdatePasskeyAssertionState(modelUser.Id, credential, time.Now()); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
c.Set("login_verification_method", service.VerificationMethodPasskey)
|
||||
setupLoginAtAuthVersion(modelUser, modelUser.AuthVersion, c)
|
||||
}
|
||||
|
||||
func AdminResetPasskey(c *gin.Context) {
|
||||
id, err := strconv.Atoi(c.Param("id"))
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, "无效的用户 ID")
|
||||
return
|
||||
}
|
||||
|
||||
user := &model.User{Id: id}
|
||||
if err := user.FillUserById(); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
myRole := c.GetInt("role")
|
||||
if !canManageTargetRole(myRole, user.Role) {
|
||||
common.ApiErrorMsg(c, "no permission")
|
||||
return
|
||||
}
|
||||
|
||||
if _, err := model.GetPasskeyByUserID(user.Id); err != nil {
|
||||
if errors.Is(err, model.ErrPasskeyNotFound) {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "该用户尚未绑定 Passkey",
|
||||
})
|
||||
return
|
||||
}
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err := model.DeletePasskeyByUserIDWithAuthVersion(user.Id); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
if _, err := model.RevokeAllUserSessions(user.Id, "admin_passkey_reset"); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
recordManageAuditFor(c, user.Id, "user.reset_passkey", map[string]any{
|
||||
"username": user.Username,
|
||||
"id": user.Id,
|
||||
})
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "Passkey 已重置",
|
||||
})
|
||||
}
|
||||
|
||||
func PasskeyVerifyBegin(c *gin.Context) {
|
||||
if !system_setting.GetPasskeySettings().Enabled {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "管理员未启用 Passkey 登录",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
user, err := getAuthenticatedUser(c)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
var request passkeyVerifyBeginRequest
|
||||
if err := common.DecodeJson(c.Request.Body, &request); err != nil {
|
||||
common.ApiErrorMsg(c, "无效的 Passkey 验证请求")
|
||||
return
|
||||
}
|
||||
binding, err := service.BindVerificationOperation(service.VerificationOperation{Scope: request.Scope, Context: request.Context})
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
identity, ok := middleware.GetSessionAuthIdentity(c)
|
||||
if !ok {
|
||||
writeSecurityOperationError(c, service.ErrAuthTokenInvalid)
|
||||
return
|
||||
}
|
||||
if _, err := service.RequireVerificationMethod(identity, request.Scope, service.VerificationMethodPasskey); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
credential, err := model.GetPasskeyByUserID(user.Id)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "该用户尚未绑定 Passkey",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
wa, err := passkeysvc.BuildWebAuthn(c.Request)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
waUser := passkeysvc.NewWebAuthnUser(user, credential)
|
||||
assertion, sessionData, err := wa.BeginLogin(waUser, webauthnlib.WithUserVerification(protocol.VerificationRequired))
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
flowToken, expiresAt, err := passkeysvc.CreateSessionDataFlow(
|
||||
model.AuthFlowPurposePasskeyStepUp,
|
||||
passkeysvc.FlowSecurity{AuthSessionIdentity: identity, Scope: binding.Scope, ContextHash: binding.ContextHash},
|
||||
sessionData,
|
||||
)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": gin.H{
|
||||
"options": assertion,
|
||||
"flow_token": flowToken,
|
||||
"expires_at": expiresAt,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func PasskeyVerifyFinish(c *gin.Context) {
|
||||
if !system_setting.GetPasskeySettings().Enabled {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "管理员未启用 Passkey 登录",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
user, err := getAuthenticatedUser(c)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
request, err := parsePasskeyFinishRequest(c)
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, "无效的 Passkey 验证请求")
|
||||
return
|
||||
}
|
||||
parsedCredential, err := protocol.ParseCredentialRequestResponseBytes(request.Credential)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
wa, err := passkeysvc.BuildWebAuthn(c.Request)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
credential, err := model.GetPasskeyByUserID(user.Id)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "该用户尚未绑定 Passkey",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
identity, ok := middleware.GetSessionAuthIdentity(c)
|
||||
if !ok {
|
||||
common.ApiErrorMsg(c, "当前认证方式不支持安全验证")
|
||||
return
|
||||
}
|
||||
sessionData, security, err := passkeysvc.PopSessionDataFlow(
|
||||
request.FlowToken,
|
||||
model.AuthFlowPurposePasskeyStepUp,
|
||||
identity,
|
||||
)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
if sessionData.UserVerification != protocol.VerificationRequired {
|
||||
writeSecurityOperationError(c, model.ErrAuthFlowInvalid)
|
||||
return
|
||||
}
|
||||
|
||||
waUser := passkeysvc.NewWebAuthnUser(user, credential)
|
||||
validatedCredential, err := wa.ValidateLogin(waUser, *sessionData, parsedCredential)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
if err := model.UpdatePasskeyAssertionState(user.Id, validatedCredential, time.Now()); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
proof, err := service.CompleteSecurityVerification(identity, service.VerificationBinding{Scope: security.Scope, ContextHash: security.ContextHash}, service.VerificationMethodPasskey)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
recordUserSecurityAudit(c, user.Id, "user.security_verify", map[string]any{"method": proof.Method, "scope": proof.Scope})
|
||||
common.ApiSuccess(c, proof)
|
||||
}
|
||||
|
||||
func getAuthenticatedUser(c *gin.Context) (*model.User, error) {
|
||||
id := c.GetInt("id")
|
||||
if id == 0 {
|
||||
return nil, service.ErrAuthTokenInvalid
|
||||
}
|
||||
user, err := model.GetUserById(id, false)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if user.Status != common.UserStatusEnabled {
|
||||
return nil, model.ErrUserSessionInactive
|
||||
}
|
||||
return user, nil
|
||||
}
|
||||
@@ -0,0 +1,72 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"crypto/ecdsa"
|
||||
"crypto/elliptic"
|
||||
"crypto/rand"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type passkeyTestBody struct {
|
||||
*strings.Reader
|
||||
}
|
||||
|
||||
func (*passkeyTestBody) Close() error { return nil }
|
||||
|
||||
func TestParsePasskeyFinishRequestDoesNotRewriteRequestBody(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
bodyText := `{"flow_token":"flow-1","credential":{"id":"credential-1"}}`
|
||||
body := &passkeyTestBody{Reader: strings.NewReader(bodyText)}
|
||||
request := httptest.NewRequest(http.MethodPost, "/api/user/passkey/register/finish", nil)
|
||||
request.Body = body
|
||||
request.ContentLength = int64(len(bodyText))
|
||||
context, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
context.Request = request
|
||||
|
||||
parsed, err := parsePasskeyFinishRequest(context)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "flow-1", parsed.FlowToken)
|
||||
assert.JSONEq(t, `{"id":"credential-1"}`, string(parsed.Credential))
|
||||
assert.Same(t, body, context.Request.Body)
|
||||
assert.Equal(t, int64(len(bodyText)), context.Request.ContentLength)
|
||||
}
|
||||
|
||||
func TestPasskeyRegisterFinishRejectsUnapprovedFlowWithoutConsumingIt(t *testing.T) {
|
||||
_, identity := setupSecurityEnrollmentTest(t)
|
||||
system_setting.GetPasskeySettings().UserVerification = "required"
|
||||
key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
|
||||
require.NoError(t, err)
|
||||
payload, err := common.Marshal(map[string]any{"scope": service.VerificationScopePasskeyRegister})
|
||||
require.NoError(t, err)
|
||||
token, _, err := model.CreateAuthFlow(model.AuthFlowCreate{
|
||||
Purpose: model.AuthFlowPurposePasskeyRegister, UserId: identity.UserID, SessionId: identity.SessionID,
|
||||
Payload: string(payload), ExpiresAt: time.Now().Add(time.Minute),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
body, err := common.Marshal(passkeyFinishRequest{
|
||||
FlowToken: token, Credential: securityPasskeyResponse(t, key, "test-challenge", true, 0),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
response := securityEnrollmentRequest("POST", "/api/user/passkey/register/finish", string(body), "", identity, PasskeyRegisterFinish)
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
assert.False(t, result.Success)
|
||||
assert.Equal(t, "AUTH_FLOW_INVALID", result.Code)
|
||||
_, err = model.GetPasskeyByUserID(identity.UserID)
|
||||
assert.ErrorIs(t, err, model.ErrPasskeyNotFound)
|
||||
flow, err := model.GetAuthFlow(token, model.AuthFlowMatch{Purpose: model.AuthFlowPurposePasskeyRegister})
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, flow.ConsumedAt)
|
||||
}
|
||||
@@ -0,0 +1,82 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/i18n"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type PaymentComplianceRequest struct {
|
||||
Confirmed bool `json:"confirmed"`
|
||||
}
|
||||
|
||||
func requirePaymentCompliance(c *gin.Context) bool {
|
||||
if !operation_setting.IsPaymentComplianceConfirmed() {
|
||||
common.ApiErrorI18n(c, i18n.MsgPaymentComplianceRequired)
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func ConfirmPaymentCompliance(c *gin.Context) {
|
||||
if c.GetBool("use_access_token") {
|
||||
c.JSON(http.StatusForbidden, gin.H{
|
||||
"success": false,
|
||||
"message": "This operation requires dashboard session authentication. API access token is not allowed.",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
var req PaymentComplianceRequest
|
||||
if err := common.DecodeJson(c.Request.Body, &req); err != nil {
|
||||
common.ApiErrorMsg(c, "参数错误")
|
||||
return
|
||||
}
|
||||
if !req.Confirmed {
|
||||
common.ApiErrorMsg(c, "请确认合规声明")
|
||||
return
|
||||
}
|
||||
|
||||
now := time.Now().Unix()
|
||||
userId := c.GetInt("id")
|
||||
clientIP := c.ClientIP()
|
||||
|
||||
updates := map[string]string{
|
||||
"payment_setting.compliance_confirmed": "true",
|
||||
"payment_setting.compliance_terms_version": operation_setting.CurrentComplianceTermsVersion,
|
||||
"payment_setting.compliance_confirmed_at": strconv.FormatInt(now, 10),
|
||||
"payment_setting.compliance_confirmed_by": strconv.Itoa(userId),
|
||||
"payment_setting.compliance_confirmed_ip": clientIP,
|
||||
}
|
||||
|
||||
for key, value := range updates {
|
||||
if err := model.UpdateOption(key, value); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
logger.LogInfo(c.Request.Context(), fmt.Sprintf(
|
||||
"payment compliance confirmed user_id=%d ip=%s terms_version=%s confirmed_at=%d",
|
||||
userId,
|
||||
clientIP,
|
||||
operation_setting.CurrentComplianceTermsVersion,
|
||||
now,
|
||||
))
|
||||
|
||||
common.ApiSuccess(c, gin.H{
|
||||
"confirmed": true,
|
||||
"terms_version": operation_setting.CurrentComplianceTermsVersion,
|
||||
"confirmed_at": now,
|
||||
"confirmed_by": userId,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,110 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/setting"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
)
|
||||
|
||||
func isPaymentComplianceConfirmed() bool {
|
||||
return operation_setting.IsPaymentComplianceConfirmed()
|
||||
}
|
||||
|
||||
func isStripeTopUpEnabled() bool {
|
||||
if !isPaymentComplianceConfirmed() {
|
||||
return false
|
||||
}
|
||||
return strings.TrimSpace(setting.StripeApiSecret) != "" &&
|
||||
strings.TrimSpace(setting.StripeWebhookSecret) != "" &&
|
||||
strings.TrimSpace(setting.StripePriceId) != ""
|
||||
}
|
||||
|
||||
func isStripeWebhookConfigured() bool {
|
||||
return strings.TrimSpace(setting.StripeWebhookSecret) != ""
|
||||
}
|
||||
|
||||
func isStripeWebhookEnabled() bool {
|
||||
return isStripeTopUpEnabled()
|
||||
}
|
||||
|
||||
func isCreemTopUpEnabled() bool {
|
||||
if !isPaymentComplianceConfirmed() {
|
||||
return false
|
||||
}
|
||||
products := strings.TrimSpace(setting.CreemProducts)
|
||||
return strings.TrimSpace(setting.CreemApiKey) != "" &&
|
||||
products != "" &&
|
||||
products != "[]"
|
||||
}
|
||||
|
||||
func isCreemWebhookConfigured() bool {
|
||||
return strings.TrimSpace(setting.CreemWebhookSecret) != ""
|
||||
}
|
||||
|
||||
func isCreemWebhookEnabled() bool {
|
||||
return isCreemTopUpEnabled() && isCreemWebhookConfigured()
|
||||
}
|
||||
|
||||
func isWaffoTopUpEnabled() bool {
|
||||
if !isPaymentComplianceConfirmed() {
|
||||
return false
|
||||
}
|
||||
if !setting.WaffoEnabled {
|
||||
return false
|
||||
}
|
||||
|
||||
return isWaffoWebhookConfigured()
|
||||
}
|
||||
|
||||
func isWaffoWebhookConfigured() bool {
|
||||
if setting.WaffoSandbox {
|
||||
return strings.TrimSpace(setting.WaffoSandboxApiKey) != "" &&
|
||||
strings.TrimSpace(setting.WaffoSandboxPrivateKey) != "" &&
|
||||
strings.TrimSpace(setting.WaffoSandboxPublicCert) != ""
|
||||
}
|
||||
|
||||
return strings.TrimSpace(setting.WaffoApiKey) != "" &&
|
||||
strings.TrimSpace(setting.WaffoPrivateKey) != "" &&
|
||||
strings.TrimSpace(setting.WaffoPublicCert) != ""
|
||||
}
|
||||
|
||||
func isWaffoWebhookEnabled() bool {
|
||||
return isWaffoTopUpEnabled()
|
||||
}
|
||||
|
||||
func isWaffoPancakeTopUpEnabled() bool {
|
||||
if !isPaymentComplianceConfirmed() {
|
||||
return false
|
||||
}
|
||||
// Presence-of-credentials = enabled. Webhook public keys ship inside
|
||||
// the SDK; mode (test/prod) is read from each event.
|
||||
return strings.TrimSpace(setting.WaffoPancakeMerchantID) != "" &&
|
||||
strings.TrimSpace(setting.WaffoPancakePrivateKey) != "" &&
|
||||
strings.TrimSpace(setting.WaffoPancakeProductID) != ""
|
||||
}
|
||||
|
||||
func isWaffoPancakeWebhookConfigured() bool {
|
||||
return isWaffoPancakeTopUpEnabled()
|
||||
}
|
||||
|
||||
func isWaffoPancakeWebhookEnabled() bool {
|
||||
return isWaffoPancakeTopUpEnabled()
|
||||
}
|
||||
|
||||
func isEpayTopUpEnabled() bool {
|
||||
if !isPaymentComplianceConfirmed() {
|
||||
return false
|
||||
}
|
||||
return isEpayWebhookConfigured() && len(operation_setting.PayMethods) > 0
|
||||
}
|
||||
|
||||
func isEpayWebhookConfigured() bool {
|
||||
return strings.TrimSpace(operation_setting.PayAddress) != "" &&
|
||||
strings.TrimSpace(operation_setting.EpayId) != "" &&
|
||||
strings.TrimSpace(operation_setting.EpayKey) != ""
|
||||
}
|
||||
|
||||
func isEpayWebhookEnabled() bool {
|
||||
return isEpayTopUpEnabled()
|
||||
}
|
||||
@@ -0,0 +1,169 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/setting"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func confirmPaymentComplianceForTest(t *testing.T) {
|
||||
t.Helper()
|
||||
paymentSetting := operation_setting.GetPaymentSetting()
|
||||
originalConfirmed := paymentSetting.ComplianceConfirmed
|
||||
originalTermsVersion := paymentSetting.ComplianceTermsVersion
|
||||
t.Cleanup(func() {
|
||||
paymentSetting.ComplianceConfirmed = originalConfirmed
|
||||
paymentSetting.ComplianceTermsVersion = originalTermsVersion
|
||||
})
|
||||
paymentSetting.ComplianceConfirmed = true
|
||||
paymentSetting.ComplianceTermsVersion = operation_setting.CurrentComplianceTermsVersion
|
||||
}
|
||||
|
||||
func TestStripeWebhookEnabledRequiresTopUpAndWebhookConfig(t *testing.T) {
|
||||
confirmPaymentComplianceForTest(t)
|
||||
originalAPISecret := setting.StripeApiSecret
|
||||
originalWebhookSecret := setting.StripeWebhookSecret
|
||||
originalPriceID := setting.StripePriceId
|
||||
t.Cleanup(func() {
|
||||
setting.StripeApiSecret = originalAPISecret
|
||||
setting.StripeWebhookSecret = originalWebhookSecret
|
||||
setting.StripePriceId = originalPriceID
|
||||
})
|
||||
|
||||
setting.StripeWebhookSecret = ""
|
||||
setting.StripeApiSecret = "sk_test_123"
|
||||
setting.StripePriceId = "price_123"
|
||||
require.False(t, isStripeWebhookEnabled())
|
||||
|
||||
setting.StripeWebhookSecret = "whsec_test"
|
||||
require.True(t, isStripeWebhookEnabled())
|
||||
|
||||
setting.StripePriceId = ""
|
||||
require.False(t, isStripeWebhookEnabled())
|
||||
}
|
||||
|
||||
func TestCreemWebhookEnabledRequiresTopUpAndWebhookConfig(t *testing.T) {
|
||||
confirmPaymentComplianceForTest(t)
|
||||
originalAPIKey := setting.CreemApiKey
|
||||
originalProducts := setting.CreemProducts
|
||||
originalWebhookSecret := setting.CreemWebhookSecret
|
||||
t.Cleanup(func() {
|
||||
setting.CreemApiKey = originalAPIKey
|
||||
setting.CreemProducts = originalProducts
|
||||
setting.CreemWebhookSecret = originalWebhookSecret
|
||||
})
|
||||
|
||||
setting.CreemWebhookSecret = ""
|
||||
setting.CreemApiKey = "creem_api_key"
|
||||
setting.CreemProducts = `[{"productId":"prod_123"}]`
|
||||
require.False(t, isCreemWebhookEnabled())
|
||||
|
||||
setting.CreemWebhookSecret = "creem_secret"
|
||||
require.True(t, isCreemWebhookEnabled())
|
||||
|
||||
setting.CreemProducts = "[]"
|
||||
require.False(t, isCreemWebhookEnabled())
|
||||
}
|
||||
|
||||
func TestWaffoWebhookEnabledRequiresTopUpAndWebhookConfig(t *testing.T) {
|
||||
confirmPaymentComplianceForTest(t)
|
||||
originalEnabled := setting.WaffoEnabled
|
||||
originalSandbox := setting.WaffoSandbox
|
||||
originalAPIKey := setting.WaffoApiKey
|
||||
originalPrivateKey := setting.WaffoPrivateKey
|
||||
originalPublicCert := setting.WaffoPublicCert
|
||||
originalSandboxAPIKey := setting.WaffoSandboxApiKey
|
||||
originalSandboxPrivateKey := setting.WaffoSandboxPrivateKey
|
||||
originalSandboxPublicCert := setting.WaffoSandboxPublicCert
|
||||
t.Cleanup(func() {
|
||||
setting.WaffoEnabled = originalEnabled
|
||||
setting.WaffoSandbox = originalSandbox
|
||||
setting.WaffoApiKey = originalAPIKey
|
||||
setting.WaffoPrivateKey = originalPrivateKey
|
||||
setting.WaffoPublicCert = originalPublicCert
|
||||
setting.WaffoSandboxApiKey = originalSandboxAPIKey
|
||||
setting.WaffoSandboxPrivateKey = originalSandboxPrivateKey
|
||||
setting.WaffoSandboxPublicCert = originalSandboxPublicCert
|
||||
})
|
||||
|
||||
setting.WaffoEnabled = true
|
||||
setting.WaffoSandbox = false
|
||||
setting.WaffoApiKey = ""
|
||||
setting.WaffoPrivateKey = "private"
|
||||
setting.WaffoPublicCert = "public"
|
||||
require.False(t, isWaffoWebhookEnabled())
|
||||
|
||||
setting.WaffoApiKey = "api"
|
||||
require.True(t, isWaffoWebhookEnabled())
|
||||
|
||||
setting.WaffoEnabled = false
|
||||
require.False(t, isWaffoWebhookEnabled())
|
||||
|
||||
setting.WaffoEnabled = true
|
||||
setting.WaffoSandbox = true
|
||||
setting.WaffoSandboxApiKey = ""
|
||||
setting.WaffoSandboxPrivateKey = "sandbox_private"
|
||||
setting.WaffoSandboxPublicCert = "sandbox_public"
|
||||
require.False(t, isWaffoWebhookEnabled())
|
||||
|
||||
setting.WaffoSandboxApiKey = "sandbox_api"
|
||||
require.True(t, isWaffoWebhookEnabled())
|
||||
}
|
||||
|
||||
func TestWaffoPancakeWebhookEnabledRequiresTopUpAndWebhookConfig(t *testing.T) {
|
||||
confirmPaymentComplianceForTest(t)
|
||||
originalMerchantID := setting.WaffoPancakeMerchantID
|
||||
originalPrivateKey := setting.WaffoPancakePrivateKey
|
||||
originalProductID := setting.WaffoPancakeProductID
|
||||
t.Cleanup(func() {
|
||||
setting.WaffoPancakeMerchantID = originalMerchantID
|
||||
setting.WaffoPancakePrivateKey = originalPrivateKey
|
||||
setting.WaffoPancakeProductID = originalProductID
|
||||
})
|
||||
|
||||
// Presence of all three credentials enables the gateway. Webhook public
|
||||
// keys are bundled in the SDK and there is no separate Enabled toggle —
|
||||
// clear any of the three fields to disable.
|
||||
setting.WaffoPancakeMerchantID = ""
|
||||
setting.WaffoPancakePrivateKey = "private"
|
||||
setting.WaffoPancakeProductID = "product"
|
||||
require.False(t, isWaffoPancakeWebhookEnabled())
|
||||
|
||||
setting.WaffoPancakeMerchantID = "merchant"
|
||||
require.True(t, isWaffoPancakeWebhookEnabled())
|
||||
|
||||
setting.WaffoPancakeProductID = ""
|
||||
require.False(t, isWaffoPancakeWebhookEnabled())
|
||||
|
||||
setting.WaffoPancakeProductID = "product"
|
||||
setting.WaffoPancakePrivateKey = ""
|
||||
require.False(t, isWaffoPancakeWebhookEnabled())
|
||||
}
|
||||
|
||||
func TestEpayWebhookEnabledRequiresTopUpAndWebhookConfig(t *testing.T) {
|
||||
confirmPaymentComplianceForTest(t)
|
||||
originalPayAddress := operation_setting.PayAddress
|
||||
originalEpayID := operation_setting.EpayId
|
||||
originalEpayKey := operation_setting.EpayKey
|
||||
originalPayMethods := operation_setting.PayMethods
|
||||
t.Cleanup(func() {
|
||||
operation_setting.PayAddress = originalPayAddress
|
||||
operation_setting.EpayId = originalEpayID
|
||||
operation_setting.EpayKey = originalEpayKey
|
||||
operation_setting.PayMethods = originalPayMethods
|
||||
})
|
||||
|
||||
operation_setting.PayAddress = "https://pay.example.com"
|
||||
operation_setting.EpayId = "epay_id"
|
||||
operation_setting.EpayKey = ""
|
||||
operation_setting.PayMethods = []map[string]string{{"type": "alipay"}}
|
||||
require.False(t, isEpayWebhookEnabled())
|
||||
|
||||
operation_setting.EpayKey = "epay_key"
|
||||
require.True(t, isEpayWebhookEnabled())
|
||||
|
||||
operation_setting.PayMethods = nil
|
||||
require.False(t, isEpayWebhookEnabled())
|
||||
}
|
||||
@@ -0,0 +1,82 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strconv"
|
||||
|
||||
perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/samber/lo"
|
||||
)
|
||||
|
||||
func GetPerfMetricsSummary(c *gin.Context) {
|
||||
hours := 24
|
||||
if rawHours := c.Query("hours"); rawHours != "" {
|
||||
if parsed, err := strconv.Atoi(rawHours); err == nil {
|
||||
hours = parsed
|
||||
}
|
||||
}
|
||||
|
||||
activeGroups := append(lo.Keys(ratio_setting.GetGroupRatioCopy()), "auto")
|
||||
result, err := perfmetrics.QuerySummaryAll(hours, activeGroups)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"data": result,
|
||||
})
|
||||
}
|
||||
|
||||
func GetPerfMetrics(c *gin.Context) {
|
||||
modelName := c.Query("model")
|
||||
if modelName == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"success": false,
|
||||
"message": "model is required",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
hours := 24
|
||||
if rawHours := c.Query("hours"); rawHours != "" {
|
||||
if parsed, err := strconv.Atoi(rawHours); err == nil {
|
||||
hours = parsed
|
||||
}
|
||||
}
|
||||
|
||||
result, err := perfmetrics.Query(perfmetrics.QueryParams{
|
||||
Model: modelName,
|
||||
Group: c.Query("group"),
|
||||
Hours: hours,
|
||||
})
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
result.Groups = filterActiveGroups(result.Groups)
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"data": result,
|
||||
})
|
||||
}
|
||||
|
||||
func filterActiveGroups(groups []perfmetrics.GroupResult) []perfmetrics.GroupResult {
|
||||
activeRatios := ratio_setting.GetGroupRatioCopy()
|
||||
return lo.Filter(groups, func(g perfmetrics.GroupResult, _ int) bool {
|
||||
_, ok := activeRatios[g.Group]
|
||||
return ok || g.Group == "auto"
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,385 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// PerformanceStats 性能统计信息
|
||||
type PerformanceStats struct {
|
||||
// 缓存统计
|
||||
CacheStats common.DiskCacheStats `json:"cache_stats"`
|
||||
// 系统内存统计
|
||||
MemoryStats MemoryStats `json:"memory_stats"`
|
||||
// 磁盘缓存目录信息
|
||||
DiskCacheInfo DiskCacheInfo `json:"disk_cache_info"`
|
||||
// 磁盘空间信息
|
||||
DiskSpaceInfo common.DiskSpaceInfo `json:"disk_space_info"`
|
||||
// 配置信息
|
||||
Config PerformanceConfig `json:"config"`
|
||||
}
|
||||
|
||||
// MemoryStats 内存统计
|
||||
type MemoryStats struct {
|
||||
// 已分配内存（字节）
|
||||
Alloc uint64 `json:"alloc"`
|
||||
// 总分配内存（字节）
|
||||
TotalAlloc uint64 `json:"total_alloc"`
|
||||
// 系统内存（字节）
|
||||
Sys uint64 `json:"sys"`
|
||||
// GC 次数
|
||||
NumGC uint32 `json:"num_gc"`
|
||||
// Goroutine 数量
|
||||
NumGoroutine int `json:"num_goroutine"`
|
||||
}
|
||||
|
||||
// DiskCacheInfo 磁盘缓存目录信息
|
||||
type DiskCacheInfo struct {
|
||||
// 缓存目录路径
|
||||
Path string `json:"path"`
|
||||
// 目录是否存在
|
||||
Exists bool `json:"exists"`
|
||||
// 文件数量
|
||||
FileCount int `json:"file_count"`
|
||||
// 总大小（字节）
|
||||
TotalSize int64 `json:"total_size"`
|
||||
}
|
||||
|
||||
// PerformanceConfig 性能配置
|
||||
type PerformanceConfig struct {
|
||||
// 是否启用磁盘缓存
|
||||
DiskCacheEnabled bool `json:"disk_cache_enabled"`
|
||||
// 磁盘缓存阈值（MB）
|
||||
DiskCacheThresholdMB int `json:"disk_cache_threshold_mb"`
|
||||
// 磁盘缓存最大大小（MB）
|
||||
DiskCacheMaxSizeMB int `json:"disk_cache_max_size_mb"`
|
||||
// 磁盘缓存路径
|
||||
DiskCachePath string `json:"disk_cache_path"`
|
||||
// 是否在容器中运行
|
||||
IsRunningInContainer bool `json:"is_running_in_container"`
|
||||
|
||||
// MonitorEnabled 是否启用性能监控
|
||||
MonitorEnabled bool `json:"monitor_enabled"`
|
||||
// MonitorCPUThreshold CPU 使用率阈值（%）
|
||||
MonitorCPUThreshold int `json:"monitor_cpu_threshold"`
|
||||
// MonitorMemoryThreshold 内存使用率阈值（%）
|
||||
MonitorMemoryThreshold int `json:"monitor_memory_threshold"`
|
||||
// MonitorDiskThreshold 磁盘使用率阈值（%）
|
||||
MonitorDiskThreshold int `json:"monitor_disk_threshold"`
|
||||
}
|
||||
|
||||
// GetPerformanceStats 获取性能统计信息
|
||||
func GetPerformanceStats(c *gin.Context) {
|
||||
// 不再每次获取统计都全量扫描磁盘，依赖原子计数器保证性能
|
||||
// 仅在系统启动或显式清理时同步
|
||||
cacheStats := common.GetDiskCacheStats()
|
||||
|
||||
// 获取内存统计
|
||||
var memStats runtime.MemStats
|
||||
runtime.ReadMemStats(&memStats)
|
||||
|
||||
// 获取磁盘缓存目录信息
|
||||
diskCacheInfo := getDiskCacheInfo()
|
||||
|
||||
// 获取配置信息
|
||||
diskConfig := common.GetDiskCacheConfig()
|
||||
monitorConfig := common.GetPerformanceMonitorConfig()
|
||||
config := PerformanceConfig{
|
||||
DiskCacheEnabled: diskConfig.Enabled,
|
||||
DiskCacheThresholdMB: diskConfig.ThresholdMB,
|
||||
DiskCacheMaxSizeMB: diskConfig.MaxSizeMB,
|
||||
DiskCachePath: diskConfig.Path,
|
||||
IsRunningInContainer: common.IsRunningInContainer(),
|
||||
MonitorEnabled: monitorConfig.Enabled,
|
||||
MonitorCPUThreshold: monitorConfig.CPUThreshold,
|
||||
MonitorMemoryThreshold: monitorConfig.MemoryThreshold,
|
||||
MonitorDiskThreshold: monitorConfig.DiskThreshold,
|
||||
}
|
||||
|
||||
// 获取磁盘空间信息
|
||||
// 使用缓存的系统状态，避免频繁调用系统 API
|
||||
systemStatus := common.GetSystemStatus()
|
||||
diskSpaceInfo := common.DiskSpaceInfo{
|
||||
UsedPercent: systemStatus.DiskUsage,
|
||||
}
|
||||
// 如果需要详细信息，可以按需获取，或者扩展 SystemStatus
|
||||
// 这里为了保持接口兼容性，我们仍然调用 GetDiskSpaceInfo，但注意这可能会有性能开销
|
||||
// 考虑到 GetPerformanceStats 是管理接口，频率较低，直接调用是可以接受的
|
||||
// 但为了一致性，我们也可以考虑从 SystemStatus 中获取部分信息
|
||||
diskSpaceInfo = common.GetDiskSpaceInfo()
|
||||
|
||||
stats := PerformanceStats{
|
||||
CacheStats: cacheStats,
|
||||
MemoryStats: MemoryStats{
|
||||
Alloc: memStats.Alloc,
|
||||
TotalAlloc: memStats.TotalAlloc,
|
||||
Sys: memStats.Sys,
|
||||
NumGC: memStats.NumGC,
|
||||
NumGoroutine: runtime.NumGoroutine(),
|
||||
},
|
||||
DiskCacheInfo: diskCacheInfo,
|
||||
DiskSpaceInfo: diskSpaceInfo,
|
||||
Config: config,
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"data": stats,
|
||||
})
|
||||
}
|
||||
|
||||
// ClearDiskCache 清理不活跃的磁盘缓存
|
||||
func ClearDiskCache(c *gin.Context) {
|
||||
// 清理超过 10 分钟未使用的缓存文件
|
||||
// 10 分钟是一个安全的阈值，确保正在进行的请求不会被误删
|
||||
err := common.CleanupOldDiskCacheFiles(10 * time.Minute)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "不活跃的磁盘缓存已清理",
|
||||
})
|
||||
}
|
||||
|
||||
// ResetPerformanceStats 重置性能统计
|
||||
func ResetPerformanceStats(c *gin.Context) {
|
||||
common.ResetDiskCacheStats()
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "统计信息已重置",
|
||||
})
|
||||
}
|
||||
|
||||
// ForceGC 强制执行 GC
|
||||
func ForceGC(c *gin.Context) {
|
||||
runtime.GC()
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "GC 已执行",
|
||||
})
|
||||
}
|
||||
|
||||
// LogFileInfo 日志文件信息
|
||||
type LogFileInfo struct {
|
||||
Name string `json:"name"`
|
||||
Size int64 `json:"size"`
|
||||
ModTime time.Time `json:"mod_time"`
|
||||
}
|
||||
|
||||
// LogFilesResponse 日志文件列表响应
|
||||
type LogFilesResponse struct {
|
||||
LogDir string `json:"log_dir"`
|
||||
Enabled bool `json:"enabled"`
|
||||
FileCount int `json:"file_count"`
|
||||
TotalSize int64 `json:"total_size"`
|
||||
OldestTime *time.Time `json:"oldest_time,omitempty"`
|
||||
NewestTime *time.Time `json:"newest_time,omitempty"`
|
||||
Files []LogFileInfo `json:"files"`
|
||||
}
|
||||
|
||||
// getLogFiles 读取日志目录中的日志文件列表
|
||||
func getLogFiles() ([]LogFileInfo, error) {
|
||||
if *common.LogDir == "" {
|
||||
return nil, nil
|
||||
}
|
||||
entries, err := os.ReadDir(*common.LogDir)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var files []LogFileInfo
|
||||
for _, entry := range entries {
|
||||
if entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
name := entry.Name()
|
||||
if !strings.HasPrefix(name, "oneapi-") || !strings.HasSuffix(name, ".log") {
|
||||
continue
|
||||
}
|
||||
info, err := entry.Info()
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
files = append(files, LogFileInfo{
|
||||
Name: name,
|
||||
Size: info.Size(),
|
||||
ModTime: info.ModTime(),
|
||||
})
|
||||
}
|
||||
// 按文件名降序排列（最新在前）
|
||||
sort.Slice(files, func(i, j int) bool {
|
||||
return files[i].Name > files[j].Name
|
||||
})
|
||||
return files, nil
|
||||
}
|
||||
|
||||
// GetLogFiles 获取日志文件列表
|
||||
func GetLogFiles(c *gin.Context) {
|
||||
if *common.LogDir == "" {
|
||||
common.ApiSuccess(c, LogFilesResponse{Enabled: false})
|
||||
return
|
||||
}
|
||||
files, err := getLogFiles()
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
var totalSize int64
|
||||
var oldest, newest time.Time
|
||||
for i, f := range files {
|
||||
totalSize += f.Size
|
||||
if i == 0 || f.ModTime.Before(oldest) {
|
||||
oldest = f.ModTime
|
||||
}
|
||||
if i == 0 || f.ModTime.After(newest) {
|
||||
newest = f.ModTime
|
||||
}
|
||||
}
|
||||
resp := LogFilesResponse{
|
||||
LogDir: *common.LogDir,
|
||||
Enabled: true,
|
||||
FileCount: len(files),
|
||||
TotalSize: totalSize,
|
||||
Files: files,
|
||||
}
|
||||
if len(files) > 0 {
|
||||
resp.OldestTime = &oldest
|
||||
resp.NewestTime = &newest
|
||||
}
|
||||
common.ApiSuccess(c, resp)
|
||||
}
|
||||
|
||||
// CleanupLogFiles 清理过期日志文件
|
||||
func CleanupLogFiles(c *gin.Context) {
|
||||
mode := c.Query("mode")
|
||||
valueStr := c.Query("value")
|
||||
if mode != "by_count" && mode != "by_days" {
|
||||
common.ApiErrorMsg(c, "invalid mode, must be by_count or by_days")
|
||||
return
|
||||
}
|
||||
value, err := strconv.Atoi(valueStr)
|
||||
if err != nil || value < 1 {
|
||||
common.ApiErrorMsg(c, "invalid value, must be a positive integer")
|
||||
return
|
||||
}
|
||||
if *common.LogDir == "" {
|
||||
common.ApiErrorMsg(c, "log directory not configured")
|
||||
return
|
||||
}
|
||||
|
||||
files, err := getLogFiles()
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
activeLogPath := logger.GetCurrentLogPath()
|
||||
var toDelete []LogFileInfo
|
||||
|
||||
switch mode {
|
||||
case "by_count":
|
||||
// files 已按名称降序（最新在前），保留前 value 个
|
||||
for i, f := range files {
|
||||
if i < value {
|
||||
continue
|
||||
}
|
||||
fullPath := filepath.Join(*common.LogDir, f.Name)
|
||||
if fullPath == activeLogPath {
|
||||
continue
|
||||
}
|
||||
toDelete = append(toDelete, f)
|
||||
}
|
||||
case "by_days":
|
||||
cutoff := time.Now().AddDate(0, 0, -value)
|
||||
for _, f := range files {
|
||||
if f.ModTime.Before(cutoff) {
|
||||
fullPath := filepath.Join(*common.LogDir, f.Name)
|
||||
if fullPath == activeLogPath {
|
||||
continue
|
||||
}
|
||||
toDelete = append(toDelete, f)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
var deletedCount int
|
||||
var freedBytes int64
|
||||
var failedFiles []string
|
||||
for _, f := range toDelete {
|
||||
fullPath := filepath.Join(*common.LogDir, f.Name)
|
||||
if err := os.Remove(fullPath); err != nil {
|
||||
failedFiles = append(failedFiles, f.Name)
|
||||
continue
|
||||
}
|
||||
deletedCount++
|
||||
freedBytes += f.Size
|
||||
}
|
||||
|
||||
result := gin.H{
|
||||
"deleted_count": deletedCount,
|
||||
"freed_bytes": freedBytes,
|
||||
"failed_files": failedFiles,
|
||||
}
|
||||
|
||||
if len(failedFiles) > 0 {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": fmt.Sprintf("部分文件删除失败（%d/%d）", len(failedFiles), len(toDelete)),
|
||||
"data": result,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": result,
|
||||
})
|
||||
}
|
||||
|
||||
// getDiskCacheInfo 获取磁盘缓存目录信息
|
||||
func getDiskCacheInfo() DiskCacheInfo {
|
||||
// 使用统一的缓存目录
|
||||
dir := common.GetDiskCacheDir()
|
||||
|
||||
info := DiskCacheInfo{
|
||||
Path: dir,
|
||||
Exists: false,
|
||||
}
|
||||
|
||||
entries, err := os.ReadDir(dir)
|
||||
if err != nil {
|
||||
return info
|
||||
}
|
||||
|
||||
info.Exists = true
|
||||
info.FileCount = 0
|
||||
info.TotalSize = 0
|
||||
|
||||
for _, entry := range entries {
|
||||
if entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
info.FileCount++
|
||||
if fileInfo, err := entry.Info(); err == nil {
|
||||
info.TotalSize += fileInfo.Size()
|
||||
}
|
||||
}
|
||||
|
||||
return info
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func Playground(c *gin.Context) {
|
||||
var newAPIError *types.NewAPIError
|
||||
|
||||
defer func() {
|
||||
if newAPIError != nil {
|
||||
c.JSON(newAPIError.StatusCode, gin.H{
|
||||
"error": newAPIError.ToOpenAIError(),
|
||||
})
|
||||
}
|
||||
}()
|
||||
|
||||
useAccessToken := c.GetBool("use_access_token")
|
||||
if useAccessToken {
|
||||
newAPIError = types.NewError(errors.New("暂不支持使用 access token"), types.ErrorCodeAccessDenied, types.ErrOptionWithSkipRetry())
|
||||
return
|
||||
}
|
||||
|
||||
relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatOpenAI, nil, nil)
|
||||
if err != nil {
|
||||
newAPIError = types.NewError(err, types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
|
||||
return
|
||||
}
|
||||
|
||||
userId := c.GetInt("id")
|
||||
|
||||
// Write user context to ensure acceptUnsetRatio is available
|
||||
userCache, err := model.GetUserCache(userId)
|
||||
if err != nil {
|
||||
newAPIError = types.NewError(err, types.ErrorCodeQueryDataError, types.ErrOptionWithSkipRetry())
|
||||
return
|
||||
}
|
||||
userCache.WriteContext(c)
|
||||
|
||||
tempToken := &model.Token{
|
||||
UserId: userId,
|
||||
Name: fmt.Sprintf("playground-%s", relayInfo.UsingGroup),
|
||||
Group: relayInfo.UsingGroup,
|
||||
}
|
||||
_ = middleware.SetupContextForToken(c, tempToken)
|
||||
|
||||
Relay(c, types.RelayFormatOpenAI)
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestRelayTaskPluginEndpointPreservesUnclaimedFallback(t *testing.T) {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
fallbackCalls := 0
|
||||
|
||||
RelayTaskPluginEndpoint(c, func(c *gin.Context) {
|
||||
fallbackCalls++
|
||||
c.Status(http.StatusNoContent)
|
||||
c.Writer.WriteHeaderNow()
|
||||
})
|
||||
|
||||
assert.Equal(t, 1, fallbackCalls)
|
||||
assert.Equal(t, http.StatusNoContent, recorder.Code)
|
||||
}
|
||||
|
||||
func TestRelayTaskPluginEndpointNeverEntersOrdinaryRelayWhenClaimed(t *testing.T) {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Set(jsplugin.ContextKeyPinnedEndpoint, jsplugin.PinnedEndpoint{
|
||||
Generation: &jsplugin.RoutingGeneration{},
|
||||
Plugin: &jsplugin.LoadedPlugin{},
|
||||
Protocol: "openai_responses",
|
||||
Operation: jsplugin.HostProtocolOperation{Name: "create"},
|
||||
})
|
||||
fallbackCalls := 0
|
||||
|
||||
RelayTaskPluginEndpoint(c, func(c *gin.Context) {
|
||||
fallbackCalls++
|
||||
c.Status(http.StatusNoContent)
|
||||
c.Writer.WriteHeaderNow()
|
||||
})
|
||||
|
||||
assert.Zero(t, fallbackCalls)
|
||||
assert.NotEqual(t, http.StatusNoContent, recorder.Code)
|
||||
}
|
||||
@@ -0,0 +1,244 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/relay"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type nativeRouteBilling struct {
|
||||
events []string
|
||||
preConsumed int
|
||||
userID int
|
||||
settled bool
|
||||
}
|
||||
|
||||
func (b *nativeRouteBilling) Settle(int) error {
|
||||
b.events = append(b.events, "settle")
|
||||
b.settled = true
|
||||
return nil
|
||||
}
|
||||
|
||||
func (b *nativeRouteBilling) Refund(*gin.Context) {
|
||||
b.events = append(b.events, "refund")
|
||||
if !b.settled && b.preConsumed > 0 {
|
||||
_ = model.IncreaseUserQuota(b.userID, b.preConsumed, true)
|
||||
b.preConsumed = 0
|
||||
}
|
||||
}
|
||||
|
||||
func (b *nativeRouteBilling) NeedsRefund() bool {
|
||||
return !b.settled && b.preConsumed > 0
|
||||
}
|
||||
|
||||
func (b *nativeRouteBilling) GetPreConsumedQuota() int {
|
||||
return b.preConsumed
|
||||
}
|
||||
|
||||
func (b *nativeRouteBilling) Reserve(quota int) error {
|
||||
b.events = append(b.events, "reserve")
|
||||
if err := model.DecreaseUserQuota(b.userID, quota, true); err != nil {
|
||||
return err
|
||||
}
|
||||
b.preConsumed = quota
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestKlingNativeRouteSubmitPollSettleAndQuery(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
service.InitHttpClient()
|
||||
|
||||
previousDB := model.DB
|
||||
previousLogDB := model.LOG_DB
|
||||
previousMemoryCache := common.MemoryCacheEnabled
|
||||
previousBatchUpdate := common.BatchUpdateEnabled
|
||||
previousLogConsume := common.LogConsumeEnabled
|
||||
previousRedisEnabled := common.RedisEnabled
|
||||
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, database.AutoMigrate(&model.User{}, &model.Channel{}, &model.Task{}, &model.Log{}))
|
||||
model.DB = database
|
||||
model.LOG_DB = database
|
||||
common.MemoryCacheEnabled = false
|
||||
common.BatchUpdateEnabled = false
|
||||
common.LogConsumeEnabled = false
|
||||
common.RedisEnabled = false
|
||||
previousModelRatios := ratio_setting.ModelRatio2JSONString()
|
||||
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"kling-v1":1}`))
|
||||
t.Cleanup(func() {
|
||||
model.DB = previousDB
|
||||
model.LOG_DB = previousLogDB
|
||||
common.MemoryCacheEnabled = previousMemoryCache
|
||||
common.BatchUpdateEnabled = previousBatchUpdate
|
||||
common.LogConsumeEnabled = previousLogConsume
|
||||
common.RedisEnabled = previousRedisEnabled
|
||||
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(previousModelRatios))
|
||||
})
|
||||
require.NoError(t, database.Create(&model.User{
|
||||
Id: 7,
|
||||
Username: "native-route-user",
|
||||
Group: "default",
|
||||
Quota: 1_000_000,
|
||||
}).Error)
|
||||
|
||||
var submitCalls atomic.Int32
|
||||
var queryCalls atomic.Int32
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
switch {
|
||||
case r.Method == http.MethodPost && r.URL.Path == "/kling/v1/videos/text2video":
|
||||
submitCalls.Add(1)
|
||||
body, readErr := io.ReadAll(r.Body)
|
||||
if !assert.NoError(t, readErr) {
|
||||
http.Error(w, "read request", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
assert.Contains(t, string(body), `"model_name":"kling-v1"`)
|
||||
_, _ = io.WriteString(w, `{"code":0,"message":"","data":{"task_id":"kling-private-1","task_status":"submitted"}}`)
|
||||
case r.Method == http.MethodGet && r.URL.Path == "/kling/v1/videos/text2video/kling-private-1":
|
||||
queryCalls.Add(1)
|
||||
_, _ = io.WriteString(w, `{"code":0,"message":"","data":{"task_id":"kling-private-1","task_status":"succeed","task_status_msg":"","task_result":{"videos":[{"id":"video-private","url":"https://cdn.example/video.mp4","duration":"5"}]},"final_unit_deduction":"1"}}`)
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
channel := model.Channel{
|
||||
Type: constant.ChannelTypeKling,
|
||||
Name: "kling-native-e2e",
|
||||
Key: "sk-test",
|
||||
BaseURL: &upstream.URL,
|
||||
Status: common.ChannelStatusEnabled,
|
||||
Models: "kling-v1",
|
||||
Group: "default",
|
||||
}
|
||||
require.NoError(t, database.Create(&channel).Error)
|
||||
|
||||
generation := pluginruntime.DefaultRegistry.Generation()
|
||||
require.NotNil(t, generation)
|
||||
submitBinding, found := generation.LookupDeclaredRoute(http.MethodPost, "/kling/v1/videos/text2video")
|
||||
require.True(t, found)
|
||||
require.Equal(t, "kling", submitBinding.Plugin.Meta.Key)
|
||||
|
||||
submitRecorder := httptest.NewRecorder()
|
||||
submitContext, _ := gin.CreateTestContext(submitRecorder)
|
||||
submitContext.Request = httptest.NewRequest(
|
||||
http.MethodPost,
|
||||
"/kling/v1/videos/text2video",
|
||||
bytes.NewBufferString(`{"model_name":"kling-v1","prompt":"a lighthouse"}`),
|
||||
)
|
||||
submitContext.Request.Header.Set("Content-Type", "application/json")
|
||||
submitContext.Set(pluginruntime.ContextKeyPinnedRoute, pluginruntime.PinnedRoute{
|
||||
Generation: generation,
|
||||
Plugin: submitBinding.Plugin,
|
||||
Route: submitBinding.Route,
|
||||
})
|
||||
common.SetContextKey(submitContext, constant.ContextKeyUserId, 7)
|
||||
common.SetContextKey(submitContext, constant.ContextKeyUserGroup, "default")
|
||||
common.SetContextKey(submitContext, constant.ContextKeyUsingGroup, "default")
|
||||
common.SetContextKey(submitContext, constant.ContextKeyTokenGroup, "default")
|
||||
common.SetContextKey(submitContext, constant.ContextKeyUserQuota, 1_000_000)
|
||||
|
||||
middleware.PrepareTaskPluginRoute()(submitContext)
|
||||
require.False(t, submitContext.IsAborted(), submitRecorder.Body.String())
|
||||
require.Equal(t, "kling-v1", submitContext.GetString("resolved_task_model"))
|
||||
require.Equal(t, "text_to_video", submitContext.GetString("task_action"))
|
||||
require.Nil(t, middleware.SetupContextForSelectedChannel(submitContext, &channel, "kling-v1"))
|
||||
|
||||
billing := &nativeRouteBilling{userID: 7}
|
||||
relayInfo := &relaycommon.RelayInfo{
|
||||
UserId: 7,
|
||||
UserGroup: "default",
|
||||
UsingGroup: "default",
|
||||
UserQuota: 1_000_000,
|
||||
TokenGroup: "default",
|
||||
OriginModelName: "kling-v1",
|
||||
Billing: billing,
|
||||
TaskRelayInfo: &relaycommon.TaskRelayInfo{
|
||||
Action: submitContext.GetString("task_action"),
|
||||
PublicTaskID: "task_kling_public",
|
||||
LockedChannel: &channel,
|
||||
},
|
||||
}
|
||||
|
||||
outcome, taskErr := executeTaskSubmissionWith(submitContext, relayInfo, relay.RelayTaskSubmit)
|
||||
require.Nil(t, taskErr)
|
||||
require.NotNil(t, outcome)
|
||||
require.Equal(t, []string{"reserve", "settle"}, billing.events)
|
||||
require.False(t, submitContext.Writer.Written())
|
||||
|
||||
presentTaskSubmission(submitContext, outcome)
|
||||
require.Equal(t, http.StatusOK, submitRecorder.Code)
|
||||
assert.Contains(t, submitRecorder.Body.String(), `"task_id":"task_kling_public"`)
|
||||
assert.NotContains(t, submitRecorder.Body.String(), "kling-private-1")
|
||||
assert.Equal(t, int32(1), submitCalls.Load())
|
||||
|
||||
var persisted model.Task
|
||||
require.NoError(t, database.Where("task_id = ?", "task_kling_public").First(&persisted).Error)
|
||||
assert.Equal(t, constant.TaskPlatform("kling"), persisted.Platform)
|
||||
assert.Equal(t, "kling-private-1", persisted.PrivateData.UpstreamTaskID)
|
||||
assert.Equal(t, model.TaskStatus(model.TaskStatusNotStart), persisted.Status)
|
||||
|
||||
previousAdaptorFactory := service.GetTaskAdaptorFunc
|
||||
service.GetTaskAdaptorFunc = func(platform constant.TaskPlatform) service.TaskPollingAdaptor {
|
||||
return relay.GetTaskAdaptor(platform)
|
||||
}
|
||||
t.Cleanup(func() { service.GetTaskAdaptorFunc = previousAdaptorFactory })
|
||||
service.DispatchPlatformUpdate(
|
||||
context.Background(),
|
||||
persisted.Platform,
|
||||
map[int][]string{channel.Id: {"kling-private-1"}},
|
||||
map[string]*model.Task{"kling-private-1": &persisted},
|
||||
)
|
||||
|
||||
require.NoError(t, database.Where("task_id = ?", "task_kling_public").First(&persisted).Error)
|
||||
assert.Equal(t, model.TaskStatus(model.TaskStatusSuccess), persisted.Status)
|
||||
assert.Equal(t, "100%", persisted.Progress)
|
||||
assert.Equal(t, 1, persisted.Quota)
|
||||
assert.Equal(t, int32(1), queryCalls.Load())
|
||||
var settledUser model.User
|
||||
require.NoError(t, database.First(&settledUser, 7).Error)
|
||||
assert.Equal(t, 999_999, settledUser.Quota)
|
||||
|
||||
queryBinding, found := generation.LookupDeclaredRoute(http.MethodGet, "/kling/v1/videos/text2video/:task_id")
|
||||
require.True(t, found)
|
||||
queryRecorder := httptest.NewRecorder()
|
||||
queryContext, _ := gin.CreateTestContext(queryRecorder)
|
||||
queryContext.Request = httptest.NewRequest(http.MethodGet, "/kling/v1/videos/text2video/task_kling_public", nil)
|
||||
queryContext.Params = gin.Params{{Key: "task_id", Value: "task_kling_public"}}
|
||||
queryContext.Set(pluginruntime.ContextKeyPinnedRoute, pluginruntime.PinnedRoute{
|
||||
Generation: generation,
|
||||
Plugin: queryBinding.Plugin,
|
||||
Route: queryBinding.Route,
|
||||
})
|
||||
common.SetContextKey(queryContext, constant.ContextKeyUserId, 7)
|
||||
|
||||
middleware.PrepareTaskPluginRoute()(queryContext)
|
||||
|
||||
require.True(t, queryContext.IsAborted())
|
||||
require.Equal(t, http.StatusOK, queryRecorder.Code)
|
||||
assert.Contains(t, queryRecorder.Body.String(), `"task_id":"task_kling_public"`)
|
||||
assert.Contains(t, queryRecorder.Body.String(), `"task_status":"succeed"`)
|
||||
assert.NotContains(t, queryRecorder.Body.String(), "kling-private-1")
|
||||
assert.NotContains(t, queryRecorder.Body.String(), upstream.URL)
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,157 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
)
|
||||
|
||||
var (
|
||||
errPluginProtocolObservationLimitExceeded = errors.New("plugin protocol observation limit exceeded")
|
||||
errInvalidPluginProtocolObservationIdentity = errors.New("invalid plugin protocol observation identity")
|
||||
)
|
||||
|
||||
type pluginProtocolObservationLimits struct {
|
||||
global int
|
||||
perPlugin int
|
||||
perUser int
|
||||
perToken int
|
||||
}
|
||||
|
||||
var defaultPluginProtocolObservationLimits = pluginProtocolObservationLimits{
|
||||
global: 128,
|
||||
perPlugin: 32,
|
||||
perUser: 4,
|
||||
perToken: 2,
|
||||
}
|
||||
|
||||
var pluginProtocolObservationAdmissions = newPluginProtocolObservationLimiter(
|
||||
defaultPluginProtocolObservationLimits,
|
||||
)
|
||||
|
||||
type pluginProtocolObservationLimitError struct {
|
||||
scope string
|
||||
limit int
|
||||
}
|
||||
|
||||
func (e *pluginProtocolObservationLimitError) Error() string {
|
||||
return fmt.Sprintf("%s: %s capacity is %d", errPluginProtocolObservationLimitExceeded, e.scope, e.limit)
|
||||
}
|
||||
|
||||
func (e *pluginProtocolObservationLimitError) Unwrap() error {
|
||||
return errPluginProtocolObservationLimitExceeded
|
||||
}
|
||||
|
||||
type pluginProtocolObservationLimiter struct {
|
||||
mu sync.Mutex
|
||||
|
||||
limits pluginProtocolObservationLimits
|
||||
global int
|
||||
plugin map[string]int
|
||||
user map[int]int
|
||||
token map[int]int
|
||||
}
|
||||
|
||||
func newPluginProtocolObservationLimiter(limits pluginProtocolObservationLimits) *pluginProtocolObservationLimiter {
|
||||
return &pluginProtocolObservationLimiter{
|
||||
limits: limits,
|
||||
plugin: make(map[string]int),
|
||||
user: make(map[int]int),
|
||||
token: make(map[int]int),
|
||||
}
|
||||
}
|
||||
|
||||
func (l *pluginProtocolObservationLimiter) acquire(
|
||||
pluginKey string,
|
||||
userID int,
|
||||
tokenID int,
|
||||
) (func(), error) {
|
||||
pluginKey = strings.TrimSpace(pluginKey)
|
||||
switch {
|
||||
case pluginKey == "":
|
||||
return nil, fmt.Errorf("%w: plugin key is required", errInvalidPluginProtocolObservationIdentity)
|
||||
case userID <= 0:
|
||||
return nil, fmt.Errorf("%w: user id must be positive", errInvalidPluginProtocolObservationIdentity)
|
||||
case tokenID <= 0:
|
||||
return nil, fmt.Errorf("%w: token id must be positive", errInvalidPluginProtocolObservationIdentity)
|
||||
}
|
||||
|
||||
l.mu.Lock()
|
||||
if l.global >= l.limits.global {
|
||||
l.mu.Unlock()
|
||||
return nil, &pluginProtocolObservationLimitError{
|
||||
scope: "global",
|
||||
limit: l.limits.global,
|
||||
}
|
||||
}
|
||||
l.global++
|
||||
|
||||
if l.plugin[pluginKey] >= l.limits.perPlugin {
|
||||
l.global--
|
||||
l.mu.Unlock()
|
||||
return nil, &pluginProtocolObservationLimitError{
|
||||
scope: "plugin",
|
||||
limit: l.limits.perPlugin,
|
||||
}
|
||||
}
|
||||
l.plugin[pluginKey]++
|
||||
|
||||
if l.user[userID] >= l.limits.perUser {
|
||||
l.global--
|
||||
l.plugin[pluginKey]--
|
||||
if l.plugin[pluginKey] == 0 {
|
||||
delete(l.plugin, pluginKey)
|
||||
}
|
||||
l.mu.Unlock()
|
||||
return nil, &pluginProtocolObservationLimitError{
|
||||
scope: "user",
|
||||
limit: l.limits.perUser,
|
||||
}
|
||||
}
|
||||
l.user[userID]++
|
||||
|
||||
if l.token[tokenID] >= l.limits.perToken {
|
||||
l.global--
|
||||
l.plugin[pluginKey]--
|
||||
if l.plugin[pluginKey] == 0 {
|
||||
delete(l.plugin, pluginKey)
|
||||
}
|
||||
l.user[userID]--
|
||||
if l.user[userID] == 0 {
|
||||
delete(l.user, userID)
|
||||
}
|
||||
l.mu.Unlock()
|
||||
return nil, &pluginProtocolObservationLimitError{
|
||||
scope: "token",
|
||||
limit: l.limits.perToken,
|
||||
}
|
||||
}
|
||||
l.token[tokenID]++
|
||||
l.mu.Unlock()
|
||||
|
||||
var once sync.Once
|
||||
return func() {
|
||||
once.Do(func() {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
|
||||
l.global--
|
||||
|
||||
l.plugin[pluginKey]--
|
||||
if l.plugin[pluginKey] == 0 {
|
||||
delete(l.plugin, pluginKey)
|
||||
}
|
||||
|
||||
l.user[userID]--
|
||||
if l.user[userID] == 0 {
|
||||
delete(l.user, userID)
|
||||
}
|
||||
|
||||
l.token[tokenID]--
|
||||
if l.token[tokenID] == 0 {
|
||||
delete(l.token, tokenID)
|
||||
}
|
||||
})
|
||||
}, nil
|
||||
}
|
||||
@@ -0,0 +1,248 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestPluginProtocolObservationLimiterCaps(t *testing.T) {
|
||||
t.Run("global", func(t *testing.T) {
|
||||
limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{
|
||||
global: 2,
|
||||
perPlugin: 2,
|
||||
perUser: 2,
|
||||
perToken: 2,
|
||||
})
|
||||
releaseFirst, err := limiter.acquire("first", 1, 1)
|
||||
require.NoError(t, err)
|
||||
defer releaseFirst()
|
||||
releaseSecond, err := limiter.acquire("second", 2, 2)
|
||||
require.NoError(t, err)
|
||||
defer releaseSecond()
|
||||
|
||||
release, err := limiter.acquire("third", 3, 3)
|
||||
assert.Nil(t, release)
|
||||
assertLimitError(t, err, "global", 2)
|
||||
})
|
||||
|
||||
t.Run("plugin", func(t *testing.T) {
|
||||
limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{
|
||||
global: 3,
|
||||
perPlugin: 1,
|
||||
perUser: 3,
|
||||
perToken: 3,
|
||||
})
|
||||
releaseFirst, err := limiter.acquire("shared", 1, 1)
|
||||
require.NoError(t, err)
|
||||
defer releaseFirst()
|
||||
|
||||
release, err := limiter.acquire("shared", 2, 2)
|
||||
assert.Nil(t, release)
|
||||
assertLimitError(t, err, "plugin", 1)
|
||||
|
||||
releaseOther, err := limiter.acquire("other", 2, 2)
|
||||
require.NoError(t, err)
|
||||
defer releaseOther()
|
||||
})
|
||||
|
||||
t.Run("user", func(t *testing.T) {
|
||||
limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{
|
||||
global: 3,
|
||||
perPlugin: 3,
|
||||
perUser: 1,
|
||||
perToken: 3,
|
||||
})
|
||||
releaseFirst, err := limiter.acquire("first", 1, 1)
|
||||
require.NoError(t, err)
|
||||
defer releaseFirst()
|
||||
|
||||
release, err := limiter.acquire("second", 1, 2)
|
||||
assert.Nil(t, release)
|
||||
assertLimitError(t, err, "user", 1)
|
||||
|
||||
releaseOther, err := limiter.acquire("second", 2, 2)
|
||||
require.NoError(t, err)
|
||||
defer releaseOther()
|
||||
})
|
||||
|
||||
t.Run("token", func(t *testing.T) {
|
||||
limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{
|
||||
global: 3,
|
||||
perPlugin: 3,
|
||||
perUser: 3,
|
||||
perToken: 1,
|
||||
})
|
||||
releaseFirst, err := limiter.acquire("first", 1, 1)
|
||||
require.NoError(t, err)
|
||||
defer releaseFirst()
|
||||
|
||||
release, err := limiter.acquire("second", 2, 1)
|
||||
assert.Nil(t, release)
|
||||
assertLimitError(t, err, "token", 1)
|
||||
|
||||
releaseOther, err := limiter.acquire("second", 2, 2)
|
||||
require.NoError(t, err)
|
||||
defer releaseOther()
|
||||
})
|
||||
}
|
||||
|
||||
func TestPluginProtocolObservationLimiterRollsBackFailedAdmission(t *testing.T) {
|
||||
t.Run("user failure", func(t *testing.T) {
|
||||
limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{
|
||||
global: 2,
|
||||
perPlugin: 2,
|
||||
perUser: 1,
|
||||
perToken: 2,
|
||||
})
|
||||
releaseHeld, err := limiter.acquire("first", 1, 1)
|
||||
require.NoError(t, err)
|
||||
defer releaseHeld()
|
||||
|
||||
release, err := limiter.acquire("second", 1, 2)
|
||||
assert.Nil(t, release)
|
||||
assertLimitError(t, err, "user", 1)
|
||||
|
||||
releaseReplacement, err := limiter.acquire("second", 2, 2)
|
||||
require.NoError(t, err)
|
||||
defer releaseReplacement()
|
||||
})
|
||||
|
||||
t.Run("token failure", func(t *testing.T) {
|
||||
limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{
|
||||
global: 2,
|
||||
perPlugin: 2,
|
||||
perUser: 2,
|
||||
perToken: 1,
|
||||
})
|
||||
releaseHeld, err := limiter.acquire("first", 1, 1)
|
||||
require.NoError(t, err)
|
||||
defer releaseHeld()
|
||||
|
||||
release, err := limiter.acquire("second", 2, 1)
|
||||
assert.Nil(t, release)
|
||||
assertLimitError(t, err, "token", 1)
|
||||
|
||||
releaseReplacement, err := limiter.acquire("second", 2, 2)
|
||||
require.NoError(t, err)
|
||||
defer releaseReplacement()
|
||||
})
|
||||
}
|
||||
|
||||
func TestPluginProtocolObservationLimiterReleaseIsIdempotent(t *testing.T) {
|
||||
limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{
|
||||
global: 1,
|
||||
perPlugin: 1,
|
||||
perUser: 1,
|
||||
perToken: 1,
|
||||
})
|
||||
release, err := limiter.acquire("plugin", 1, 1)
|
||||
require.NoError(t, err)
|
||||
|
||||
release()
|
||||
release()
|
||||
|
||||
releaseAgain, err := limiter.acquire("plugin", 1, 1)
|
||||
require.NoError(t, err)
|
||||
releaseAgain()
|
||||
}
|
||||
|
||||
func TestPluginProtocolObservationLimiterRejectsMissingIdentityWithoutConsumingCapacity(t *testing.T) {
|
||||
limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{
|
||||
global: 1,
|
||||
perPlugin: 1,
|
||||
perUser: 1,
|
||||
perToken: 1,
|
||||
})
|
||||
|
||||
for _, testCase := range []struct {
|
||||
name string
|
||||
pluginKey string
|
||||
userID int
|
||||
tokenID int
|
||||
}{
|
||||
{name: "empty plugin", userID: 1, tokenID: 1},
|
||||
{name: "blank plugin", pluginKey: " \t", userID: 1, tokenID: 1},
|
||||
{name: "zero user", pluginKey: "plugin", tokenID: 1},
|
||||
{name: "negative user", pluginKey: "plugin", userID: -1, tokenID: 1},
|
||||
{name: "zero token", pluginKey: "plugin", userID: 1},
|
||||
{name: "negative token", pluginKey: "plugin", userID: 1, tokenID: -1},
|
||||
} {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
release, err := limiter.acquire(testCase.pluginKey, testCase.userID, testCase.tokenID)
|
||||
assert.Nil(t, release)
|
||||
assert.ErrorIs(t, err, errInvalidPluginProtocolObservationIdentity)
|
||||
})
|
||||
}
|
||||
|
||||
release, err := limiter.acquire("plugin", 1, 1)
|
||||
require.NoError(t, err)
|
||||
release()
|
||||
}
|
||||
|
||||
func TestPluginProtocolObservationLimiterConcurrentAdmissionsRespectCap(t *testing.T) {
|
||||
const (
|
||||
workerCount = 8
|
||||
globalLimit = 3
|
||||
)
|
||||
limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{
|
||||
global: globalLimit,
|
||||
perPlugin: workerCount,
|
||||
perUser: workerCount,
|
||||
perToken: workerCount,
|
||||
})
|
||||
start := make(chan struct{})
|
||||
releases := make(chan func(), workerCount)
|
||||
errorsFound := make(chan error, workerCount)
|
||||
|
||||
var workers sync.WaitGroup
|
||||
workers.Add(workerCount)
|
||||
for worker := 1; worker <= workerCount; worker++ {
|
||||
go func(id int) {
|
||||
defer workers.Done()
|
||||
<-start
|
||||
release, err := limiter.acquire("plugin", id, id)
|
||||
if err != nil {
|
||||
errorsFound <- err
|
||||
return
|
||||
}
|
||||
releases <- release
|
||||
}(worker)
|
||||
}
|
||||
close(start)
|
||||
workers.Wait()
|
||||
close(releases)
|
||||
close(errorsFound)
|
||||
|
||||
assert.Len(t, releases, globalLimit)
|
||||
assert.Len(t, errorsFound, workerCount-globalLimit)
|
||||
for err := range errorsFound {
|
||||
assert.ErrorIs(t, err, errPluginProtocolObservationLimitExceeded)
|
||||
}
|
||||
for release := range releases {
|
||||
release()
|
||||
}
|
||||
|
||||
release, err := limiter.acquire("plugin", 1, 1)
|
||||
require.NoError(t, err)
|
||||
release()
|
||||
}
|
||||
|
||||
func assertLimitError(
|
||||
t *testing.T,
|
||||
err error,
|
||||
expectedScope string,
|
||||
expectedLimit int,
|
||||
) {
|
||||
t.Helper()
|
||||
require.Error(t, err)
|
||||
assert.ErrorIs(t, err, errPluginProtocolObservationLimitExceeded)
|
||||
|
||||
var limitError *pluginProtocolObservationLimitError
|
||||
require.True(t, errors.As(err, &limitError))
|
||||
assert.Equal(t, expectedScope, limitError.scope)
|
||||
assert.Equal(t, expectedLimit, limitError.limit)
|
||||
}
|
||||
@@ -0,0 +1,305 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
builtinplugins "github.com/QuantumNous/new-api/plugins"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/openai/openai-go"
|
||||
"github.com/openai/openai-go/option"
|
||||
"github.com/openai/openai-go/responses"
|
||||
"github.com/openai/openai-go/shared"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestTaskPluginResponsesNonStreamDecodesWithOfficialGoSDK(t *testing.T) {
|
||||
pinned := compilePluginProtocolTestEndpoint(t, "official-sdk-non-stream", `
|
||||
export const protocols = {openai_responses: {
|
||||
renderEvents: function() { throw new Error("non-stream called renderEvents"); },
|
||||
renderFinal: function() {
|
||||
return {
|
||||
output: [{
|
||||
type: "message",
|
||||
status: "completed",
|
||||
role: "assistant",
|
||||
content: [{
|
||||
type: "output_text",
|
||||
text: "official-sdk-final",
|
||||
annotations: [],
|
||||
logprobs: []
|
||||
}]
|
||||
}]
|
||||
};
|
||||
}
|
||||
}};
|
||||
`)
|
||||
deps := pluginProtocolTestDeps()
|
||||
deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) {
|
||||
return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_sdk_final", nil), nil
|
||||
}
|
||||
deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) {
|
||||
return &model.Task{
|
||||
TaskID: "task_sdk_final",
|
||||
UserId: 71,
|
||||
Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key),
|
||||
Status: model.TaskStatusSuccess,
|
||||
}, true, nil
|
||||
}
|
||||
server := newPluginProtocolSDKTestServer(t, pinned, deps)
|
||||
defer server.Close()
|
||||
client := openai.NewClient(
|
||||
option.WithAPIKey("test-key"),
|
||||
option.WithBaseURL(server.URL+"/v1/"),
|
||||
option.WithMaxRetries(0),
|
||||
)
|
||||
requestContext, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
response, err := client.Responses.New(requestContext, responses.ResponseNewParams{
|
||||
Model: shared.ResponsesModel("video-model"),
|
||||
Input: responses.ResponseNewParamsInputUnion{
|
||||
OfString: openai.String("create a video"),
|
||||
},
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "resp_sdk_final", response.ID)
|
||||
assert.Equal(t, responses.ResponseStatusCompleted, response.Status)
|
||||
assert.Equal(t, "video-model", response.Model)
|
||||
assert.Equal(t, "official-sdk-final", response.OutputText())
|
||||
assert.Equal(t, "task_sdk_final", response.Metadata["task_id"])
|
||||
}
|
||||
|
||||
func TestTaskPluginResponsesStreamDecodesWithOfficialGoSDK(t *testing.T) {
|
||||
pinned := compilePluginProtocolTestEndpoint(t, "official-sdk-stream", `
|
||||
export const protocols = {openai_responses: {
|
||||
renderEvents: function() {
|
||||
return {events: [{type: "output", data: "official-sdk-stream"}], done: true};
|
||||
},
|
||||
renderFinal: function() { throw new Error("stream called renderFinal"); }
|
||||
}};
|
||||
`)
|
||||
deps := pluginProtocolTestDeps()
|
||||
deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) {
|
||||
return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_sdk_stream", nil), nil
|
||||
}
|
||||
deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) {
|
||||
return &model.Task{
|
||||
TaskID: "task_sdk_stream",
|
||||
UserId: 71,
|
||||
Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key),
|
||||
Status: model.TaskStatusSuccess,
|
||||
}, true, nil
|
||||
}
|
||||
server := newPluginProtocolSDKTestServer(t, pinned, deps)
|
||||
defer server.Close()
|
||||
client := openai.NewClient(
|
||||
option.WithAPIKey("test-key"),
|
||||
option.WithBaseURL(server.URL+"/v1/"),
|
||||
option.WithMaxRetries(0),
|
||||
)
|
||||
requestContext, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
stream := client.Responses.NewStreaming(requestContext, responses.ResponseNewParams{
|
||||
Model: shared.ResponsesModel("video-model"),
|
||||
Input: responses.ResponseNewParamsInputUnion{
|
||||
OfString: openai.String("create a video"),
|
||||
},
|
||||
})
|
||||
defer stream.Close()
|
||||
eventTypes := make([]string, 0, 8)
|
||||
sequenceNumbers := make([]int64, 0, 8)
|
||||
var completedText string
|
||||
for stream.Next() {
|
||||
event := stream.Current()
|
||||
eventTypes = append(eventTypes, event.Type)
|
||||
sequenceNumbers = append(sequenceNumbers, event.SequenceNumber)
|
||||
if event.Type == "response.completed" {
|
||||
completedText = event.Response.OutputText()
|
||||
}
|
||||
}
|
||||
|
||||
require.NoError(t, stream.Err())
|
||||
assert.Equal(t, []string{
|
||||
"response.created",
|
||||
"response.output_item.added",
|
||||
"response.content_part.added",
|
||||
"response.output_text.delta",
|
||||
"response.output_text.done",
|
||||
"response.content_part.done",
|
||||
"response.output_item.done",
|
||||
"response.completed",
|
||||
}, eventTypes)
|
||||
assert.Equal(t, []int64{0, 1, 2, 3, 4, 5, 6, 7}, sequenceNumbers)
|
||||
assert.Equal(t, "official-sdk-stream", completedText)
|
||||
}
|
||||
|
||||
func TestBuiltInKlingResponsesNonStreamDecodesWithOfficialGoSDK(t *testing.T) {
|
||||
pinned, deps := builtInKlingProtocolSDKFixture(t, "task_kling_sdk_final")
|
||||
server := newPluginProtocolSDKTestServer(t, pinned, deps)
|
||||
defer server.Close()
|
||||
client := openai.NewClient(
|
||||
option.WithAPIKey("test-key"),
|
||||
option.WithBaseURL(server.URL+"/v1/"),
|
||||
option.WithMaxRetries(0),
|
||||
)
|
||||
requestContext, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
response, err := client.Responses.New(requestContext, responses.ResponseNewParams{
|
||||
Model: shared.ResponsesModel("kling-v2-master"),
|
||||
Input: responses.ResponseNewParamsInputUnion{
|
||||
OfString: openai.String("camera orbit"),
|
||||
},
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, responses.ResponseStatusCompleted, response.Status)
|
||||
assert.Equal(t, "kling-v2-master", response.Model)
|
||||
assert.Contains(t, response.OutputText(), "https://gateway.example/v1/tasks/task_kling_sdk_final/artifacts/video/content")
|
||||
assert.NotContains(t, response.OutputText(), "upstream.example")
|
||||
assert.Equal(t, "kling", response.Metadata["vendor"])
|
||||
assert.Equal(t, "task_kling_sdk_final", response.Metadata["task_id"])
|
||||
}
|
||||
|
||||
func TestBuiltInKlingResponsesStreamDecodesWithOfficialGoSDK(t *testing.T) {
|
||||
pinned, deps := builtInKlingProtocolSDKFixture(t, "task_kling_sdk_stream")
|
||||
server := newPluginProtocolSDKTestServer(t, pinned, deps)
|
||||
defer server.Close()
|
||||
client := openai.NewClient(
|
||||
option.WithAPIKey("test-key"),
|
||||
option.WithBaseURL(server.URL+"/v1/"),
|
||||
option.WithMaxRetries(0),
|
||||
)
|
||||
requestContext, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
stream := client.Responses.NewStreaming(requestContext, responses.ResponseNewParams{
|
||||
Model: shared.ResponsesModel("kling-v2-master"),
|
||||
Input: responses.ResponseNewParamsInputUnion{
|
||||
OfString: openai.String("camera orbit"),
|
||||
},
|
||||
})
|
||||
defer stream.Close()
|
||||
eventTypes := make([]string, 0, 8)
|
||||
sequenceNumbers := make([]int64, 0, 8)
|
||||
var completedText string
|
||||
for stream.Next() {
|
||||
event := stream.Current()
|
||||
eventTypes = append(eventTypes, event.Type)
|
||||
sequenceNumbers = append(sequenceNumbers, event.SequenceNumber)
|
||||
if event.Type == "response.completed" {
|
||||
completedText = event.Response.OutputText()
|
||||
}
|
||||
}
|
||||
|
||||
require.NoError(t, stream.Err())
|
||||
assert.Equal(t, []string{
|
||||
"response.created",
|
||||
"response.output_item.added",
|
||||
"response.content_part.added",
|
||||
"response.output_text.delta",
|
||||
"response.output_text.done",
|
||||
"response.content_part.done",
|
||||
"response.output_item.done",
|
||||
"response.completed",
|
||||
}, eventTypes)
|
||||
assert.Equal(t, []int64{0, 1, 2, 3, 4, 5, 6, 7}, sequenceNumbers)
|
||||
assert.Contains(t, completedText, "https://gateway.example/v1/tasks/task_kling_sdk_stream/artifacts/video/content")
|
||||
assert.NotContains(t, completedText, "upstream.example")
|
||||
}
|
||||
|
||||
func builtInKlingProtocolSDKFixture(t *testing.T, taskID string) (pluginruntime.PinnedEndpoint, pluginProtocolBridgeDeps) {
|
||||
t.Helper()
|
||||
source, err := builtinplugins.Source("kling")
|
||||
require.NoError(t, err)
|
||||
registry := pluginruntime.NewRegistry()
|
||||
plugin, err := registry.RegisterFactory(source, pluginruntime.Options{Key: "kling"})
|
||||
require.NoError(t, err)
|
||||
binding, found := registry.Generation().LookupEndpoint(http.MethodPost, "/v1/responses", "kling-v2-master")
|
||||
require.True(t, found)
|
||||
pinned := pluginruntime.PinnedEndpoint{
|
||||
Generation: registry.Generation(),
|
||||
Plugin: plugin,
|
||||
Protocol: binding.Protocol,
|
||||
Operation: binding.Operation,
|
||||
Model: binding.Model,
|
||||
Candidates: []pluginruntime.ProtocolBinding{binding},
|
||||
}
|
||||
deps := pluginProtocolTestDeps()
|
||||
deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) {
|
||||
return pluginProtocolTestOutcome(info, plugin.Meta.Key, taskID, nil), nil
|
||||
}
|
||||
deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) {
|
||||
task := &model.Task{
|
||||
TaskID: taskID, Platform: constant.TaskPlatform(plugin.Meta.Key), UserId: 71,
|
||||
Status: model.TaskStatusSuccess, Progress: "100%", CreatedAt: 1_710_000_000,
|
||||
}
|
||||
task.SetData(map[string]any{
|
||||
"code": 0,
|
||||
"data": map[string]any{
|
||||
"task_id": "upstream-private",
|
||||
"task_status": "succeed",
|
||||
"task_result": map[string]any{
|
||||
"videos": []any{map[string]any{"url": "https://upstream.example/private-video.mp4"}},
|
||||
},
|
||||
},
|
||||
})
|
||||
return task, true, nil
|
||||
}
|
||||
deps.artifactContentURL = func(publicTaskID, artifactKey string) (string, error) {
|
||||
return "https://gateway.example/v1/tasks/" + publicTaskID + "/artifacts/" + artifactKey + "/content", nil
|
||||
}
|
||||
return pinned, deps
|
||||
}
|
||||
|
||||
func newPluginProtocolSDKTestServer(
|
||||
t *testing.T,
|
||||
pinned pluginruntime.PinnedEndpoint,
|
||||
deps pluginProtocolBridgeDeps,
|
||||
) *httptest.Server {
|
||||
t.Helper()
|
||||
return httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
||||
if request.Method != http.MethodPost || request.URL.Path != "/v1/responses" {
|
||||
http.NotFound(writer, request)
|
||||
return
|
||||
}
|
||||
var requestBody map[string]any
|
||||
if err := common.DecodeJson(request.Body, &requestBody); err != nil {
|
||||
http.Error(writer, "invalid request", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
modelName, _ := requestBody["model"].(string)
|
||||
stream, _ := requestBody["stream"].(bool)
|
||||
c, _ := gin.CreateTestContext(writer)
|
||||
c.Request = request
|
||||
common.SetContextKey(c, constant.ContextKeyUserId, 71)
|
||||
common.SetContextKey(c, constant.ContextKeyTokenId, 81)
|
||||
common.SetContextKey(c, constant.ContextKeyUsingGroup, "default")
|
||||
c.Set("resolved_task_model", modelName)
|
||||
c.Set(pluginruntime.ContextKeyProtocolRequest, pluginruntime.ProtocolRequestContext{
|
||||
RouteRequestContext: pluginruntime.RouteRequestContext{
|
||||
Path: request.URL.Path,
|
||||
Method: request.Method,
|
||||
Params: map[string]string{},
|
||||
Query: request.URL.Query(),
|
||||
RequestBody: requestBody,
|
||||
},
|
||||
Protocol: pinned.Protocol,
|
||||
Stream: stream,
|
||||
})
|
||||
serveTaskPluginProtocol(c, pinned, deps)
|
||||
}))
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,90 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"strconv"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// GetPrefillGroups 获取预填组列表，可通过 ?type=xxx 过滤
|
||||
func GetPrefillGroups(c *gin.Context) {
|
||||
groupType := c.Query("type")
|
||||
groups, err := model.GetAllPrefillGroups(groupType)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, groups)
|
||||
}
|
||||
|
||||
// CreatePrefillGroup 创建新的预填组
|
||||
func CreatePrefillGroup(c *gin.Context) {
|
||||
var g model.PrefillGroup
|
||||
if err := c.ShouldBindJSON(&g); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if g.Name == "" || g.Type == "" {
|
||||
common.ApiErrorMsg(c, "组名称和类型不能为空")
|
||||
return
|
||||
}
|
||||
// 创建前检查名称
|
||||
if dup, err := model.IsPrefillGroupNameDuplicated(0, g.Name); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
} else if dup {
|
||||
common.ApiErrorMsg(c, "组名称已存在")
|
||||
return
|
||||
}
|
||||
|
||||
if err := g.Insert(); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, &g)
|
||||
}
|
||||
|
||||
// UpdatePrefillGroup 更新预填组
|
||||
func UpdatePrefillGroup(c *gin.Context) {
|
||||
var g model.PrefillGroup
|
||||
if err := c.ShouldBindJSON(&g); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if g.Id == 0 {
|
||||
common.ApiErrorMsg(c, "缺少组 ID")
|
||||
return
|
||||
}
|
||||
// 名称冲突检查
|
||||
if dup, err := model.IsPrefillGroupNameDuplicated(g.Id, g.Name); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
} else if dup {
|
||||
common.ApiErrorMsg(c, "组名称已存在")
|
||||
return
|
||||
}
|
||||
|
||||
if err := g.Update(); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, &g)
|
||||
}
|
||||
|
||||
// DeletePrefillGroup 删除预填组
|
||||
func DeletePrefillGroup(c *gin.Context) {
|
||||
idStr := c.Param("id")
|
||||
id, err := strconv.Atoi(idStr)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if err := model.DeletePrefillGroupByID(id); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, nil)
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"maps"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/billing_setting"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func filterPricingByUsableGroups(pricing []model.Pricing, usableGroup map[string]string) []model.Pricing {
|
||||
if len(pricing) == 0 {
|
||||
return pricing
|
||||
}
|
||||
if len(usableGroup) == 0 {
|
||||
return []model.Pricing{}
|
||||
}
|
||||
|
||||
filtered := make([]model.Pricing, 0, len(pricing))
|
||||
for _, item := range pricing {
|
||||
if common.StringsContains(item.EnableGroup, "all") {
|
||||
filtered = append(filtered, item)
|
||||
continue
|
||||
}
|
||||
for _, group := range item.EnableGroup {
|
||||
if _, ok := usableGroup[group]; ok {
|
||||
filtered = append(filtered, item)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
return filtered
|
||||
}
|
||||
|
||||
func GetPricing(c *gin.Context) {
|
||||
pricing := model.GetPricing()
|
||||
enabledGroups, err := model.GetEnabledGroupSet()
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
userId, exists := c.Get("id")
|
||||
usableGroup := map[string]string{}
|
||||
groupRatio := map[string]float64{}
|
||||
maps.Copy(groupRatio, ratio_setting.GetGroupRatioCopy())
|
||||
var group string
|
||||
if exists {
|
||||
user, err := model.GetUserCache(userId.(int))
|
||||
if err == nil {
|
||||
group = user.Group
|
||||
for g := range groupRatio {
|
||||
ratio, ok := ratio_setting.GetGroupGroupRatio(group, g)
|
||||
if ok {
|
||||
groupRatio[g] = ratio
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
usableGroup = service.GetUserUsableGroups(group)
|
||||
for group := range usableGroup {
|
||||
if group == "auto" {
|
||||
continue
|
||||
}
|
||||
if _, ok := enabledGroups[group]; !ok {
|
||||
delete(usableGroup, group)
|
||||
}
|
||||
}
|
||||
if len(enabledGroups) == 0 {
|
||||
delete(usableGroup, "auto")
|
||||
}
|
||||
pricing = filterPricingByUsableGroups(pricing, usableGroup)
|
||||
pricing = append([]model.Pricing(nil), pricing...)
|
||||
for i := range pricing {
|
||||
for _, group := range pricing[i].EnableGroup {
|
||||
if _, visible := usableGroup[group]; !visible {
|
||||
continue
|
||||
}
|
||||
if expr, ok := billing_setting.GetGroupBillingExpr(group, pricing[i].ModelName); ok {
|
||||
if pricing[i].GroupBillingExpr == nil {
|
||||
pricing[i].GroupBillingExpr = make(map[string]string)
|
||||
}
|
||||
pricing[i].GroupBillingExpr[group] = expr
|
||||
}
|
||||
}
|
||||
}
|
||||
// check groupRatio contains usableGroup
|
||||
for group := range ratio_setting.GetGroupRatioCopy() {
|
||||
_, usable := usableGroup[group]
|
||||
_, enabled := enabledGroups[group]
|
||||
if !usable || !enabled {
|
||||
delete(groupRatio, group)
|
||||
}
|
||||
}
|
||||
|
||||
c.JSON(200, gin.H{
|
||||
"success": true,
|
||||
"data": pricing,
|
||||
"vendors": model.GetVendors(),
|
||||
"group_ratio": groupRatio,
|
||||
"usable_group": usableGroup,
|
||||
"supported_endpoint": model.GetSupportedEndpointMap(),
|
||||
"auto_groups": service.GetUserAutoGroup(group),
|
||||
"pricing_version": "a42d372ccf0b5dd13ecf71203521f9d2",
|
||||
})
|
||||
}
|
||||
|
||||
func ResetModelRatio(c *gin.Context) {
|
||||
defaultStr := ratio_setting.DefaultModelRatio2JSONString()
|
||||
err := model.UpdateOption("ModelRatio", defaultStr)
|
||||
if err != nil {
|
||||
c.JSON(200, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
err = ratio_setting.UpdateModelRatioByJSONString(defaultStr)
|
||||
if err != nil {
|
||||
c.JSON(200, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
c.JSON(200, gin.H{
|
||||
"success": true,
|
||||
"message": "重置模型倍率成功",
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func GetRankings(c *gin.Context) {
|
||||
result, err := service.GetRankingsSnapshot(c.DefaultQuery("period", "week"))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"data": result,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/QuantumNous/new-api/setting/billing_setting"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func GetRatioConfig(c *gin.Context) {
|
||||
if !ratio_setting.IsExposeRatioEnabled() {
|
||||
c.JSON(http.StatusForbidden, gin.H{
|
||||
"success": false,
|
||||
"message": "倍率配置接口未启用",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": billing_setting.GetPricingSyncData(map[string]any(ratio_setting.GetExposedData())),
|
||||
})
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,142 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/setting/billing_setting"
|
||||
"github.com/QuantumNous/new-api/setting/config"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
)
|
||||
|
||||
func TestPricingSyncExpressionPriority(t *testing.T) {
|
||||
expression := `tier("base", p * 2 + c * 8 + cr * 0)`
|
||||
cases := []struct {
|
||||
name string
|
||||
local map[string]any
|
||||
source map[string]any
|
||||
wantFields []string
|
||||
}{
|
||||
{"equal expressions suppress stale ratios", map[string]any{"billing_mode": map[string]string{"m": "tiered_expr"}, "billing_expr": map[string]string{"m": expression}}, map[string]any{"billing_mode": map[string]string{"m": "tiered_expr"}, "billing_expr": map[string]string{"m": expression}, "model_ratio": map[string]float64{"m": 3}, "model_price": map[string]float64{"m": 2}}, nil},
|
||||
{"local expression excludes legacy-only source", map[string]any{"billing_mode": map[string]string{"m": "tiered_expr"}, "billing_expr": map[string]string{"m": expression}}, map[string]any{"model_ratio": map[string]float64{"m": 3}, "completion_ratio": map[string]float64{"m": 2}}, nil},
|
||||
{"expression imports without legacy conflicts", map[string]any{"model_ratio": map[string]float64{"m": 1}}, map[string]any{"billing_mode": map[string]string{"m": "tiered_expr"}, "billing_expr": map[string]string{"m": expression}, "model_ratio": map[string]float64{"m": 3}, "model_price": map[string]float64{"m": 2}}, []string{"billing_mode", "billing_expr"}},
|
||||
{"inactive expression follows explicit ratio mode", map[string]any{"model_ratio": map[string]float64{"m": 1}}, map[string]any{"billing_mode": map[string]string{"m": "ratio"}, "billing_expr": map[string]string{"m": expression}, "model_ratio": map[string]float64{"m": 3}}, []string{"model_ratio"}},
|
||||
{"empty active expression never imports a false free price", map[string]any{}, map[string]any{"billing_mode": map[string]string{"m": "tiered_expr"}, "billing_expr": map[string]string{"m": " "}, "model_ratio": map[string]float64{"m": 0}}, nil},
|
||||
}
|
||||
for _, tt := range cases {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
diff := buildDifferences(tt.local, []struct {
|
||||
name string
|
||||
data map[string]any
|
||||
}{{"source", tt.source}})
|
||||
fields := make([]string, 0, len(diff["m"]))
|
||||
for field := range diff["m"] {
|
||||
fields = append(fields, field)
|
||||
}
|
||||
assert.ElementsMatch(t, tt.wantFields, fields)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRatioConfigExportsEffectiveExpressions(t *testing.T) {
|
||||
before := config.GlobalConfig.ExportAllConfigs()
|
||||
expose := ratio_setting.IsExposeRatioEnabled()
|
||||
t.Cleanup(func() {
|
||||
config.UpdateConfigFromMap(config.GlobalConfig.Get("billing_setting"), map[string]string{"billing_mode": before["billing_setting.billing_mode"], "billing_expr": before["billing_setting.billing_expr"]})
|
||||
ratio_setting.SetExposeRatioEnabled(expose)
|
||||
})
|
||||
config.UpdateConfigFromMap(config.GlobalConfig.Get("billing_setting"), map[string]string{"billing_mode": `{"sync-export":"tiered_expr"}`, "billing_expr": `{"sync-export":"tier(\"base\", p * 2)"}`})
|
||||
ratio_setting.SetExposeRatioEnabled(true)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/api/ratio_config", nil)
|
||||
GetRatioConfig(c)
|
||||
var response struct {
|
||||
Success bool
|
||||
Data map[string]any
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
|
||||
require.True(t, response.Success)
|
||||
assert.Equal(t, billing_setting.BillingModeTieredExpr, valueMap(response.Data["billing_mode"])["sync-export"])
|
||||
assert.Equal(t, `tier("base", p * 2)`, valueMap(response.Data["billing_expr"])["sync-export"])
|
||||
}
|
||||
|
||||
func TestPricingSyncCompleteSourcesAndArrayFormats(t *testing.T) {
|
||||
before := config.GlobalConfig.ExportAllConfigs()
|
||||
oldRatios, oldCompletion := ratio_setting.ModelRatio2JSONString(), ratio_setting.CompletionRatio2JSONString()
|
||||
t.Cleanup(func() {
|
||||
config.UpdateConfigFromMap(config.GlobalConfig.Get("billing_setting"), map[string]string{"billing_mode": before["billing_setting.billing_mode"], "billing_expr": before["billing_setting.billing_expr"]})
|
||||
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(oldRatios))
|
||||
require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(oldCompletion))
|
||||
})
|
||||
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"sync-token":1}`))
|
||||
require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"sync-token":2}`))
|
||||
expression := `len <= 200000 ? tier("short", p * 2 + c * 8 + cr * 0) : tier("long", p * 4 + c * 12)`
|
||||
expressions, err := common.Marshal(map[string]string{"sync-already": expression})
|
||||
require.NoError(t, err)
|
||||
config.UpdateConfigFromMap(config.GlobalConfig.Get("billing_setting"), map[string]string{"billing_mode": `{"sync-already":"tiered_expr"}`, "billing_expr": string(expressions)})
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var data any
|
||||
if r.URL.Path == "/ratio_config" {
|
||||
data = map[string]any{
|
||||
"billing_mode": map[string]string{"sync-already": "tiered_expr", "sync-expression": "tiered_expr"},
|
||||
"billing_expr": map[string]string{"sync-already": expression, "sync-expression": expression},
|
||||
"model_ratio": map[string]float64{"sync-already": 9, "sync-expression": 9, "sync-token": 1},
|
||||
"model_price": map[string]float64{"sync-expression": 4},
|
||||
"completion_ratio": map[string]float64{"sync-token": 4},
|
||||
"cache_ratio": map[string]float64{"sync-token": 0},
|
||||
}
|
||||
} else {
|
||||
data = []map[string]any{
|
||||
{"model_name": "sync-already", "model_ratio": 5, "completion_ratio": 3},
|
||||
{"model_name": "sync-expression", "model_ratio": 2, "model_price": 1},
|
||||
{"model_name": "sync-array-expression", "billing_mode": "tiered_expr", "billing_expr": expression, "quota_type": 1, "model_price": 0},
|
||||
{"model_name": "sync-unpriced"},
|
||||
{"model_name": "sync-invalid-expression", "billing_mode": "tiered_expr", "billing_expr": "", "model_ratio": 0},
|
||||
{"model_name": "sync-free", "model_ratio": 0, "completion_ratio": 0},
|
||||
}
|
||||
}
|
||||
encoded, err := common.Marshal(map[string]any{"success": true, "data": data})
|
||||
require.NoError(t, err)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write(encoded)
|
||||
}))
|
||||
defer server.Close()
|
||||
var response struct {
|
||||
Success bool
|
||||
Data struct {
|
||||
Differences map[string]map[string]dto.DifferenceItem
|
||||
Prices map[string]struct {
|
||||
Current map[string]any
|
||||
Upstreams map[string]map[string]any
|
||||
}
|
||||
TestResults []dto.TestResult `json:"test_results"`
|
||||
}
|
||||
}
|
||||
body := map[string]any{"upstreams": []map[string]any{
|
||||
{"id": 1, "name": "Expressions", "base_url": server.URL, "endpoint": "/ratio_config"},
|
||||
{"id": 2, "name": "Legacy", "base_url": server.URL, "endpoint": "/pricing"},
|
||||
}}
|
||||
recorder := modelManagementRequest(t, FetchUpstreamRatios, http.MethodPost, "/api/channel/fetch_upstream_ratios", body, &response)
|
||||
require.True(t, response.Success, recorder.Body.String())
|
||||
require.Len(t, response.Data.TestResults, 2)
|
||||
for _, result := range response.Data.TestResults {
|
||||
require.Equal(t, "success", result.Status, result.Error)
|
||||
}
|
||||
assert.NotContains(t, response.Data.Differences, "sync-already")
|
||||
assert.NotContains(t, response.Data.Differences, "sync-unpriced")
|
||||
assert.NotContains(t, response.Data.Differences, "sync-invalid-expression")
|
||||
assert.Equal(t, map[string]any{"billing_mode": "tiered_expr", "billing_expr": expression}, response.Data.Prices["sync-expression"].Upstreams["Expressions(1)"])
|
||||
assert.NotContains(t, response.Data.Prices["sync-expression"].Upstreams, "Legacy(2)")
|
||||
assert.Equal(t, map[string]any{"billing_mode": "tiered_expr", "billing_expr": expression}, response.Data.Prices["sync-array-expression"].Upstreams["Legacy(2)"])
|
||||
assert.Equal(t, float64(1), response.Data.Prices["sync-token"].Upstreams["Expressions(1)"]["model_ratio"], "unchanged base prices are included for a complete price preview")
|
||||
assert.Equal(t, float64(4), response.Data.Prices["sync-token"].Upstreams["Expressions(1)"]["completion_ratio"])
|
||||
assert.Equal(t, float64(0), response.Data.Prices["sync-token"].Upstreams["Expressions(1)"]["cache_ratio"])
|
||||
assert.Equal(t, float64(0), response.Data.Prices["sync-free"].Upstreams["Legacy(2)"]["model_ratio"])
|
||||
}
|
||||
@@ -0,0 +1,238 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/i18n"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func GetAllRedemptions(c *gin.Context) {
|
||||
pageInfo := common.GetPageQuery(c)
|
||||
redemptions, total, err := model.GetAllRedemptions(pageInfo.GetStartIdx(), pageInfo.GetPageSize())
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
pageInfo.SetTotal(int(total))
|
||||
pageInfo.SetItems(redemptions)
|
||||
common.ApiSuccess(c, pageInfo)
|
||||
return
|
||||
}
|
||||
|
||||
func SearchRedemptions(c *gin.Context) {
|
||||
keyword := c.Query("keyword")
|
||||
status := c.Query("status")
|
||||
pageInfo := common.GetPageQuery(c)
|
||||
redemptions, total, err := model.SearchRedemptions(keyword, status, pageInfo.GetStartIdx(), pageInfo.GetPageSize())
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
pageInfo.SetTotal(int(total))
|
||||
pageInfo.SetItems(redemptions)
|
||||
common.ApiSuccess(c, pageInfo)
|
||||
return
|
||||
}
|
||||
|
||||
func GetRedemption(c *gin.Context) {
|
||||
id, err := strconv.Atoi(c.Param("id"))
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
redemption, err := model.GetRedemptionById(id)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": redemption,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
func AddRedemption(c *gin.Context) {
|
||||
if !operation_setting.IsPaymentComplianceConfirmed() {
|
||||
common.ApiErrorI18n(c, i18n.MsgPaymentComplianceRequired)
|
||||
return
|
||||
}
|
||||
|
||||
redemption := model.Redemption{}
|
||||
err := c.ShouldBindJSON(&redemption)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if utf8.RuneCountInString(redemption.Name) == 0 || utf8.RuneCountInString(redemption.Name) > 20 {
|
||||
common.ApiErrorI18n(c, i18n.MsgRedemptionNameLength)
|
||||
return
|
||||
}
|
||||
if redemption.Count <= 0 {
|
||||
common.ApiErrorI18n(c, i18n.MsgRedemptionCountPositive)
|
||||
return
|
||||
}
|
||||
if redemption.Count > 100 {
|
||||
common.ApiErrorI18n(c, i18n.MsgRedemptionCountMax)
|
||||
return
|
||||
}
|
||||
if redemption.Quota <= 0 {
|
||||
common.ApiError(c, errors.New("redemption quota must be positive"))
|
||||
return
|
||||
}
|
||||
if err := common.ValidateWalletQuota(redemption.Quota); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if valid, msg := validateExpiredTime(c, redemption.ExpiredTime); !valid {
|
||||
c.JSON(http.StatusOK, gin.H{"success": false, "message": msg})
|
||||
return
|
||||
}
|
||||
var keys []string
|
||||
for i := 0; i < redemption.Count; i++ {
|
||||
key := common.GetUUID()
|
||||
cleanRedemption := model.Redemption{
|
||||
UserId: c.GetInt("id"),
|
||||
Name: redemption.Name,
|
||||
Key: key,
|
||||
CreatedTime: common.GetTimestamp(),
|
||||
Quota: redemption.Quota,
|
||||
ExpiredTime: redemption.ExpiredTime,
|
||||
}
|
||||
err = cleanRedemption.Insert()
|
||||
if err != nil {
|
||||
common.SysError("failed to insert redemption: " + err.Error())
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": i18n.T(c, i18n.MsgRedemptionCreateFailed),
|
||||
"data": keys,
|
||||
})
|
||||
return
|
||||
}
|
||||
keys = append(keys, key)
|
||||
}
|
||||
recordManageAudit(c, "redemption.create", map[string]any{
|
||||
"name": redemption.Name,
|
||||
"count": redemption.Count,
|
||||
"quota": logger.LogQuota(redemption.Quota),
|
||||
})
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": keys,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
func DeleteRedemption(c *gin.Context) {
|
||||
id, _ := strconv.Atoi(c.Param("id"))
|
||||
err := model.DeleteRedemptionById(id)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
func UpdateRedemption(c *gin.Context) {
|
||||
statusOnly := c.Query("status_only")
|
||||
redemption := model.Redemption{}
|
||||
err := c.ShouldBindJSON(&redemption)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
cleanRedemption, err := model.GetRedemptionById(redemption.Id)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if statusOnly == "" {
|
||||
if redemption.Quota <= 0 {
|
||||
common.ApiError(c, errors.New("redemption quota must be positive"))
|
||||
return
|
||||
}
|
||||
if err := common.ValidateWalletQuota(redemption.Quota); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if valid, msg := validateExpiredTime(c, redemption.ExpiredTime); !valid {
|
||||
c.JSON(http.StatusOK, gin.H{"success": false, "message": msg})
|
||||
return
|
||||
}
|
||||
// If you add more fields, please also update redemption.Update()
|
||||
cleanRedemption.Name = redemption.Name
|
||||
cleanRedemption.Quota = redemption.Quota
|
||||
cleanRedemption.ExpiredTime = redemption.ExpiredTime
|
||||
}
|
||||
if statusOnly != "" {
|
||||
cleanRedemption.Status = redemption.Status
|
||||
}
|
||||
err = cleanRedemption.Update()
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": cleanRedemption,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
func DeleteInvalidRedemption(c *gin.Context) {
|
||||
rows, err := model.DeleteInvalidRedemptions()
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": rows,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
func validateExpiredTime(c *gin.Context, expired int64) (bool, string) {
|
||||
if expired != 0 && expired < common.GetTimestamp() {
|
||||
return false, i18n.T(c, i18n.MsgRedemptionExpireTimeInvalid)
|
||||
}
|
||||
return true, ""
|
||||
}
|
||||
|
||||
func DeleteRedemptionBatch(c *gin.Context) {
|
||||
var request struct {
|
||||
Ids []int `json:"ids" binding:"required,min=1,max=1000,dive,gt=0"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&request); err != nil {
|
||||
common.ApiErrorI18n(c, i18n.MsgInvalidParams)
|
||||
return
|
||||
}
|
||||
count, err := model.BatchDeleteRedemptions(request.Ids)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
recordManageAudit(c, "redemption.delete_batch", map[string]any{
|
||||
"count": count,
|
||||
"total": len(request.Ids),
|
||||
"requested_redemption_ids": request.Ids,
|
||||
})
|
||||
common.ApiSuccess(c, count)
|
||||
}
|
||||
@@ -0,0 +1,201 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/driver/mysql"
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func TestDeleteRedemptionBatch(t *testing.T) {
|
||||
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
|
||||
t.Run(dialect, func(t *testing.T) {
|
||||
var driver, logDriver gorm.Dialector
|
||||
dbType := common.DatabaseTypeSQLite
|
||||
switch dialect {
|
||||
case "sqlite":
|
||||
driver = sqlite.Open(":memory:")
|
||||
logDriver = sqlite.Open(":memory:")
|
||||
case "mysql":
|
||||
dsn := os.Getenv("TEST_MYSQL_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("TEST_MYSQL_DSN is not configured")
|
||||
}
|
||||
driver = mysql.Open(dsn)
|
||||
logDSN := os.Getenv("TEST_MYSQL_LOG_DSN")
|
||||
if logDSN == "" {
|
||||
logDSN = dsn
|
||||
}
|
||||
logDriver = mysql.Open(logDSN)
|
||||
dbType = common.DatabaseTypeMySQL
|
||||
case "postgres":
|
||||
dsn := os.Getenv("TEST_POSTGRES_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("TEST_POSTGRES_DSN is not configured")
|
||||
}
|
||||
driver = postgres.Open(dsn)
|
||||
logDSN := os.Getenv("TEST_POSTGRES_LOG_DSN")
|
||||
if logDSN == "" {
|
||||
logDSN = dsn
|
||||
}
|
||||
logDriver = postgres.Open(logDSN)
|
||||
dbType = common.DatabaseTypePostgreSQL
|
||||
}
|
||||
db, err := gorm.Open(driver, &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
sqlDB.SetMaxOpenConns(1)
|
||||
t.Cleanup(func() { require.NoError(t, sqlDB.Close()) })
|
||||
var version string
|
||||
query := "SELECT version()"
|
||||
if dialect == "sqlite" {
|
||||
query = "SELECT sqlite_version()"
|
||||
}
|
||||
require.NoError(t, db.Raw(query).Scan(&version).Error)
|
||||
t.Logf("database version: %s", version)
|
||||
|
||||
logDB, err := gorm.Open(logDriver, &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
logSQL, err := logDB.DB()
|
||||
require.NoError(t, err)
|
||||
logSQL.SetMaxOpenConns(1)
|
||||
t.Cleanup(func() { require.NoError(t, logSQL.Close()) })
|
||||
previousDB, previousLogDB := model.DB, model.LOG_DB
|
||||
previousMain, previousLog := common.MainDatabaseType(), common.LogDatabaseType()
|
||||
previousRedis := common.RedisEnabled
|
||||
model.DB, model.LOG_DB = db, logDB
|
||||
common.SetDatabaseTypes(dbType, dbType)
|
||||
common.RedisEnabled = false
|
||||
t.Cleanup(func() {
|
||||
model.DB, model.LOG_DB = previousDB, previousLogDB
|
||||
common.SetDatabaseTypes(previousMain, previousLog)
|
||||
common.RedisEnabled = previousRedis
|
||||
})
|
||||
for _, table := range []any{&model.User{}, &model.Redemption{}} {
|
||||
require.False(t, db.Migrator().HasTable(table), "use an empty test database")
|
||||
require.NoError(t, db.AutoMigrate(table))
|
||||
t.Cleanup(func() { require.NoError(t, db.Migrator().DropTable(table)) })
|
||||
}
|
||||
require.False(t, logDB.Migrator().HasTable(&model.AuditLog{}), "use an empty test log database")
|
||||
require.NoError(t, logDB.AutoMigrate(&model.AuditLog{}))
|
||||
t.Cleanup(func() { require.NoError(t, logDB.Migrator().DropTable(&model.AuditLog{})) })
|
||||
token := "redemption-audit-test-token"
|
||||
admin := model.User{Username: "redemption-audit-admin", Password: "unused", Role: common.RoleAdminUser, Status: common.UserStatusEnabled, Group: "default", AccessToken: &token}
|
||||
require.NoError(t, db.Create(&admin).Error)
|
||||
codes := make([]model.Redemption, 16)
|
||||
for index := range codes {
|
||||
codes[index] = model.Redemption{Name: "selected", Key: fmt.Sprintf("%032d", index+1), Quota: 100, Status: common.RedemptionCodeStatusEnabled}
|
||||
}
|
||||
codes[1].Status = common.RedemptionCodeStatusUsed
|
||||
codes[15].Name = "unselected"
|
||||
codes[15].Status = common.RedemptionCodeStatusDisabled
|
||||
require.NoError(t, model.DB.Create(&codes).Error)
|
||||
router := gin.New()
|
||||
router.Use(middleware.RequestId())
|
||||
router.POST("/api/redemption/batch", middleware.AdminAuth(), DeleteRedemptionBatch)
|
||||
|
||||
overLimit := make([]int, 1001)
|
||||
for index := range overLimit {
|
||||
overLimit[index] = codes[0].Id
|
||||
}
|
||||
oversized, err := common.Marshal(map[string]any{"ids": overLimit})
|
||||
require.NoError(t, err)
|
||||
for _, body := range []string{"{}", `{"ids":[]}`, `{"ids":null}`, `{"ids":[0]}`, `{"ids":[1,-1]}`, `{"ids":["1"]}`, "{", string(oversized)} {
|
||||
t.Run("invalid_"+body[:min(len(body), 30)], func(t *testing.T) {
|
||||
response := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(http.MethodPost, "/api/redemption/batch", bytes.NewBufferString(body))
|
||||
request.Header.Set("Authorization", "Bearer "+token)
|
||||
router.ServeHTTP(response, request)
|
||||
var result struct {
|
||||
Success bool `json:"success"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
assert.False(t, result.Success)
|
||||
var count int64
|
||||
require.NoError(t, model.DB.Model(&model.Redemption{}).Count(&count).Error)
|
||||
assert.EqualValues(t, 16, count)
|
||||
var events []model.AuditLog
|
||||
require.NoError(t, logDB.Where("request_id = ? AND category = ?", response.Header().Get(common.RequestIdKey), model.AuditCategoryOperation).Find(&events).Error)
|
||||
require.Len(t, events, 1)
|
||||
assert.False(t, events[0].Success)
|
||||
assert.Equal(t, "redemption.delete_batch", events[0].Action)
|
||||
})
|
||||
}
|
||||
_, err = model.BatchDeleteRedemptions(nil)
|
||||
require.Error(t, err)
|
||||
requestedIDs := make([]int, 0, 17)
|
||||
for _, code := range codes[:15] {
|
||||
requestedIDs = append(requestedIDs, code.Id)
|
||||
}
|
||||
requestedIDs = append(requestedIDs, codes[0].Id, 999999)
|
||||
payload, err := common.Marshal(map[string]any{"ids": requestedIDs})
|
||||
require.NoError(t, err)
|
||||
for _, expectedCount := range []int64{15, 0} {
|
||||
response := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(http.MethodPost, "/api/redemption/batch", bytes.NewReader(payload))
|
||||
request.Header.Set("Authorization", "Bearer "+token)
|
||||
router.ServeHTTP(response, request)
|
||||
assert.Equal(t, http.StatusOK, response.Code)
|
||||
var result struct {
|
||||
Success bool `json:"success"`
|
||||
Data int64 `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
assert.True(t, result.Success)
|
||||
assert.Equal(t, expectedCount, result.Data)
|
||||
var events []model.AuditLog
|
||||
require.NoError(t, logDB.Where("request_id = ? AND category = ?", response.Header().Get(common.RequestIdKey), model.AuditCategoryOperation).Find(&events).Error)
|
||||
require.Len(t, events, 1, "one operation event, without a duplicate single-delete fallback")
|
||||
event := events[0]
|
||||
assert.Equal(t, "redemption.delete_batch", event.Action)
|
||||
assert.Equal(t, fmt.Sprintf("Batch deleted %d redemption codes", expectedCount), event.Content)
|
||||
assert.True(t, event.Success)
|
||||
assert.Equal(t, admin.Id, event.UserId)
|
||||
assert.Equal(t, "/api/redemption/batch", event.Route)
|
||||
require.NotNil(t, event.Other.Op)
|
||||
encoded, err := common.Marshal(event.Other.Op.Params)
|
||||
require.NoError(t, err)
|
||||
var params struct {
|
||||
Count int64 `json:"count"`
|
||||
Total int `json:"total"`
|
||||
IDs []int `json:"requested_redemption_ids"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(encoded, ¶ms))
|
||||
assert.Equal(t, expectedCount, params.Count)
|
||||
assert.Equal(t, len(requestedIDs), params.Total)
|
||||
assert.Equal(t, requestedIDs, params.IDs)
|
||||
encoded, err = common.Marshal(event)
|
||||
require.NoError(t, err)
|
||||
assert.NotContains(t, string(encoded), token)
|
||||
for _, code := range codes {
|
||||
assert.NotContains(t, string(encoded), code.Key)
|
||||
}
|
||||
}
|
||||
var active []model.Redemption
|
||||
require.NoError(t, model.DB.Find(&active).Error)
|
||||
require.Len(t, active, 1)
|
||||
assert.Equal(t, codes[15], active[0])
|
||||
var all []model.Redemption
|
||||
require.NoError(t, model.DB.Unscoped().Order("id").Find(&all).Error)
|
||||
require.Len(t, all, 16)
|
||||
for _, code := range all[:15] {
|
||||
assert.True(t, code.DeletedAt.Valid)
|
||||
}
|
||||
assert.False(t, all[15].DeletedAt.Valid)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,934 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
taskdto "github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
|
||||
"github.com/QuantumNous/new-api/relay"
|
||||
"github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
||||
"github.com/QuantumNous/new-api/relay/helper"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
|
||||
"github.com/bytedance/gopkg/util/gopool"
|
||||
"github.com/samber/lo"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
func relayHandler(c *gin.Context, info *relaycommon.RelayInfo) *types.NewAPIError {
|
||||
var err *types.NewAPIError
|
||||
switch info.RelayMode {
|
||||
case relayconstant.RelayModeImagesGenerations, relayconstant.RelayModeImagesEdits:
|
||||
err = relay.ImageHelper(c, info)
|
||||
case relayconstant.RelayModeAudioSpeech:
|
||||
fallthrough
|
||||
case relayconstant.RelayModeAudioTranslation:
|
||||
fallthrough
|
||||
case relayconstant.RelayModeAudioTranscription:
|
||||
err = relay.AudioHelper(c, info)
|
||||
case relayconstant.RelayModeRerank:
|
||||
err = relay.RerankHelper(c, info)
|
||||
case relayconstant.RelayModeEmbeddings:
|
||||
err = relay.EmbeddingHelper(c, info)
|
||||
case relayconstant.RelayModeResponses, relayconstant.RelayModeResponsesCompact:
|
||||
err = relay.ResponsesHelper(c, info)
|
||||
case relayconstant.RelayModeAlphaSearch:
|
||||
err = relay.AlphaSearchHelper(c, info)
|
||||
default:
|
||||
err = relay.TextHelper(c, info)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func geminiRelayHandler(c *gin.Context, info *relaycommon.RelayInfo) *types.NewAPIError {
|
||||
var err *types.NewAPIError
|
||||
if strings.Contains(c.Request.URL.Path, "embed") {
|
||||
err = relay.GeminiEmbeddingHandler(c, info)
|
||||
} else {
|
||||
err = relay.GeminiHelper(c, info)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
const publicUpstreamErrorMessage = "The model service is temporarily unavailable. Please try again later."
|
||||
|
||||
func publicUpstreamErrorContent(statusCode int) string {
|
||||
return fmt.Sprintf("status_code=%d, %s", statusCode, publicUpstreamErrorMessage)
|
||||
}
|
||||
|
||||
func clientFacingOpenAIError(err *types.NewAPIError, upstreamFailure bool, requestId string) types.OpenAIError {
|
||||
if upstreamFailure {
|
||||
return types.OpenAIError{
|
||||
Message: common.MessageWithRequestId(publicUpstreamErrorMessage, requestId),
|
||||
Type: "api_error",
|
||||
Code: "model_service_error",
|
||||
}
|
||||
}
|
||||
result := err.ToOpenAIError()
|
||||
result.Message = common.MessageWithRequestId(result.Message, requestId)
|
||||
return result
|
||||
}
|
||||
|
||||
func clientFacingClaudeError(err *types.NewAPIError, upstreamFailure bool, requestId string) types.ClaudeError {
|
||||
if upstreamFailure {
|
||||
return types.ClaudeError{Type: "api_error", Message: common.MessageWithRequestId(publicUpstreamErrorMessage, requestId)}
|
||||
}
|
||||
result := err.ToClaudeError()
|
||||
result.Message = common.MessageWithRequestId(result.Message, requestId)
|
||||
return result
|
||||
}
|
||||
|
||||
func Relay(c *gin.Context, relayFormat types.RelayFormat) {
|
||||
|
||||
requestId := c.GetString(common.RequestIdKey)
|
||||
//group := common.GetContextKeyString(c, constant.ContextKeyUsingGroup)
|
||||
//originalModel := common.GetContextKeyString(c, constant.ContextKeyOriginalModel)
|
||||
|
||||
var (
|
||||
newAPIError *types.NewAPIError
|
||||
ws *websocket.Conn
|
||||
upstreamFailure bool
|
||||
)
|
||||
|
||||
if relayFormat == types.RelayFormatOpenAIRealtime {
|
||||
var err error
|
||||
ws, err = upgrader.Upgrade(c.Writer, c.Request, nil)
|
||||
if err != nil {
|
||||
helper.WssError(c, ws, types.NewError(err, types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()).ToOpenAIError())
|
||||
return
|
||||
}
|
||||
defer ws.Close()
|
||||
}
|
||||
|
||||
defer func() {
|
||||
if newAPIError != nil {
|
||||
logger.LogError(c, fmt.Sprintf("relay error: %s", common.LocalLogPreview(newAPIError.Error())))
|
||||
switch relayFormat {
|
||||
case types.RelayFormatOpenAIRealtime:
|
||||
helper.WssError(c, ws, clientFacingOpenAIError(newAPIError, upstreamFailure, requestId))
|
||||
case types.RelayFormatClaude:
|
||||
c.JSON(newAPIError.StatusCode, gin.H{
|
||||
"type": "error",
|
||||
"error": clientFacingClaudeError(newAPIError, upstreamFailure, requestId),
|
||||
})
|
||||
default:
|
||||
c.JSON(newAPIError.StatusCode, gin.H{
|
||||
"error": clientFacingOpenAIError(newAPIError, upstreamFailure, requestId),
|
||||
})
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
request, err := helper.GetAndValidateRequest(c, relayFormat)
|
||||
if err != nil {
|
||||
// Map "request body too large" to 413 so clients can handle it correctly
|
||||
if common.IsRequestBodyTooLargeError(err) || errors.Is(err, common.ErrRequestBodyTooLarge) {
|
||||
newAPIError = types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusRequestEntityTooLarge, types.ErrOptionWithSkipRetry())
|
||||
} else {
|
||||
newAPIError = types.NewError(err, types.ErrorCodeInvalidRequest, types.ErrOptionWithStatusCode(http.StatusBadRequest), types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
relayInfo, err := relaycommon.GenRelayInfo(c, relayFormat, request, ws)
|
||||
if err != nil {
|
||||
newAPIError = types.NewError(err, types.ErrorCodeGenRelayInfoFailed)
|
||||
return
|
||||
}
|
||||
|
||||
needSensitiveCheck := setting.ShouldCheckPromptSensitive()
|
||||
needCountToken := constant.CountToken
|
||||
// Avoid building huge CombineText (strings.Join) when token counting and sensitive check are both disabled.
|
||||
var meta *types.TokenCountMeta
|
||||
if needSensitiveCheck || needCountToken {
|
||||
meta = request.GetTokenCountMeta()
|
||||
} else {
|
||||
meta = fastTokenCountMetaForPricing(request)
|
||||
}
|
||||
|
||||
if needSensitiveCheck && meta != nil {
|
||||
contains, words := service.CheckSensitiveText(meta.CombineText)
|
||||
if contains {
|
||||
logger.LogWarn(c, fmt.Sprintf("user sensitive words detected: %s", strings.Join(words, ", ")))
|
||||
newAPIError = types.NewError(err, types.ErrorCodeSensitiveWordsDetected)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
tokens, err := service.EstimateRequestToken(c, meta, relayInfo)
|
||||
if err != nil {
|
||||
newAPIError = types.NewError(err, types.ErrorCodeCountTokenFailed)
|
||||
return
|
||||
}
|
||||
|
||||
relayInfo.SetEstimatePromptTokens(tokens)
|
||||
|
||||
priceData, err := helper.ModelPriceHelper(c, relayInfo, tokens, meta)
|
||||
if err != nil {
|
||||
newAPIError = types.NewError(err, types.ErrorCodeModelPriceError, types.ErrOptionWithStatusCode(http.StatusBadRequest))
|
||||
return
|
||||
}
|
||||
|
||||
// common.SetContextKey(c, constant.ContextKeyTokenCountMeta, meta)
|
||||
|
||||
if priceData.FreeModel {
|
||||
logger.LogInfo(c, fmt.Sprintf("模型 %s 免费，跳过预扣费", relayInfo.OriginModelName))
|
||||
} else {
|
||||
newAPIError = service.PreConsumeBilling(c, priceData.QuotaToPreConsume, relayInfo)
|
||||
if newAPIError != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
defer func() {
|
||||
// Only return quota if downstream failed and quota was actually pre-consumed
|
||||
if newAPIError != nil {
|
||||
newAPIError = service.NormalizeViolationFeeError(newAPIError)
|
||||
if relayInfo.Billing != nil {
|
||||
relayInfo.Billing.Refund(c)
|
||||
}
|
||||
service.ChargeViolationFeeIfNeeded(c, relayInfo, newAPIError)
|
||||
}
|
||||
}()
|
||||
|
||||
retryParam := &service.RetryParam{
|
||||
Ctx: c,
|
||||
TokenGroup: relayInfo.TokenGroup,
|
||||
ModelName: relayInfo.OriginModelName,
|
||||
RequestPath: c.Request.URL.Path,
|
||||
Retry: common.GetPointer(0),
|
||||
}
|
||||
relayInfo.RetryIndex = 0
|
||||
relayInfo.LastError = nil
|
||||
|
||||
for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() {
|
||||
relayInfo.RetryIndex = retryParam.GetRetry()
|
||||
channel, channelErr := getChannel(c, relayInfo, retryParam)
|
||||
if channelErr != nil {
|
||||
logger.LogError(c, channelErr.Error())
|
||||
newAPIError = channelErr
|
||||
break
|
||||
}
|
||||
addUsedChannel(c, channel.Id)
|
||||
if billingErr := service.PrepareTieredBillingForSelectedGroup(c, relayInfo); billingErr != nil {
|
||||
newAPIError = billingErr
|
||||
break
|
||||
}
|
||||
|
||||
bodyStorage, bodyErr := common.GetBodyStorage(c)
|
||||
if bodyErr != nil {
|
||||
// Ensure consistent 413 for oversized bodies even when error occurs later (e.g., retry path)
|
||||
if common.IsRequestBodyTooLargeError(bodyErr) || errors.Is(bodyErr, common.ErrRequestBodyTooLarge) {
|
||||
newAPIError = types.NewErrorWithStatusCode(bodyErr, types.ErrorCodeReadRequestBodyFailed, http.StatusRequestEntityTooLarge, types.ErrOptionWithSkipRetry())
|
||||
} else {
|
||||
newAPIError = types.NewErrorWithStatusCode(bodyErr, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
break
|
||||
}
|
||||
c.Request.Body = io.NopCloser(bodyStorage)
|
||||
|
||||
switch relayFormat {
|
||||
case types.RelayFormatOpenAIRealtime:
|
||||
newAPIError = relay.WssHelper(c, relayInfo)
|
||||
case types.RelayFormatClaude:
|
||||
newAPIError = relay.ClaudeHelper(c, relayInfo)
|
||||
case types.RelayFormatGemini:
|
||||
newAPIError = geminiRelayHandler(c, relayInfo)
|
||||
default:
|
||||
newAPIError = relayHandler(c, relayInfo)
|
||||
}
|
||||
|
||||
if newAPIError == nil {
|
||||
relayInfo.LastError = nil
|
||||
return
|
||||
}
|
||||
upstreamFailure = true
|
||||
|
||||
newAPIError = service.NormalizeViolationFeeError(newAPIError)
|
||||
relayInfo.LastError = newAPIError
|
||||
|
||||
processChannelError(c, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()), newAPIError, relayInfo)
|
||||
|
||||
if !shouldRetry(c, newAPIError, common.RetryTimes-retryParam.GetRetry()) {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
useChannel := c.GetStringSlice("use_channel")
|
||||
if len(useChannel) > 1 {
|
||||
retryLogStr := fmt.Sprintf("重试：%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]"))
|
||||
logger.LogInfo(c, retryLogStr)
|
||||
}
|
||||
if newAPIError != nil {
|
||||
gopool.Go(func() {
|
||||
perfmetrics.RecordRelaySample(relayInfo, false, 0)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// CountClaudeTokens implements Anthropic's token-counting utility endpoint.
|
||||
// It deliberately skips upstream generation and billing; callers use this
|
||||
// endpoint to size prompts before creating a Message.
|
||||
func CountClaudeTokens(c *gin.Context) {
|
||||
request, err := helper.GetAndValidateClaudeRequest(c)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"type": "error",
|
||||
"error": gin.H{
|
||||
"type": "invalid_request_error",
|
||||
"message": common.MessageWithRequestId(err.Error(), c.GetString(common.RequestIdKey)),
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
info := relaycommon.GenRelayInfoClaude(c, request)
|
||||
inputTokens, err := service.CountRequestToken(c, request.GetTokenCountMeta(), info)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{
|
||||
"type": "error",
|
||||
"error": gin.H{
|
||||
"type": "api_error",
|
||||
"message": common.MessageWithRequestId(err.Error(), c.GetString(common.RequestIdKey)),
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"input_tokens": inputTokens})
|
||||
}
|
||||
|
||||
var upgrader = websocket.Upgrader{
|
||||
Subprotocols: []string{"realtime"}, // WS 握手支持的协议，如果有使用 Sec-WebSocket-Protocol，则必须在此声明对应的 Protocol TODO add other protocol
|
||||
CheckOrigin: func(r *http.Request) bool {
|
||||
return true // 允许跨域
|
||||
},
|
||||
}
|
||||
|
||||
func addUsedChannel(c *gin.Context, channelId int) {
|
||||
useChannel := c.GetStringSlice("use_channel")
|
||||
useChannel = append(useChannel, fmt.Sprintf("%d", channelId))
|
||||
c.Set("use_channel", useChannel)
|
||||
}
|
||||
|
||||
func fastTokenCountMetaForPricing(request dto.Request) *types.TokenCountMeta {
|
||||
if request == nil {
|
||||
return &types.TokenCountMeta{}
|
||||
}
|
||||
meta := &types.TokenCountMeta{
|
||||
TokenType: types.TokenTypeTokenizer,
|
||||
}
|
||||
switch r := request.(type) {
|
||||
case *dto.GeneralOpenAIRequest:
|
||||
maxCompletionTokens := lo.FromPtrOr(r.MaxCompletionTokens, uint(0))
|
||||
maxTokens := lo.FromPtrOr(r.MaxTokens, uint(0))
|
||||
if maxCompletionTokens > maxTokens {
|
||||
meta.MaxTokens = int(maxCompletionTokens)
|
||||
} else {
|
||||
meta.MaxTokens = int(maxTokens)
|
||||
}
|
||||
case *dto.OpenAIResponsesRequest:
|
||||
meta.MaxTokens = int(lo.FromPtrOr(r.MaxOutputTokens, uint(0)))
|
||||
case *dto.ClaudeRequest:
|
||||
meta.MaxTokens = int(lo.FromPtr(r.MaxTokens))
|
||||
case *dto.ImageRequest:
|
||||
// Pricing for image requests depends on ImagePriceRatio; safe to compute even when CountToken is disabled.
|
||||
return r.GetTokenCountMeta()
|
||||
default:
|
||||
// Best-effort: leave CombineText empty to avoid large allocations.
|
||||
}
|
||||
return meta
|
||||
}
|
||||
|
||||
func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service.RetryParam) (*model.Channel, *types.NewAPIError) {
|
||||
if info.ChannelMeta == nil {
|
||||
autoBan := c.GetBool("auto_ban")
|
||||
autoBanInt := 1
|
||||
if !autoBan {
|
||||
autoBanInt = 0
|
||||
}
|
||||
return &model.Channel{
|
||||
Id: c.GetInt("channel_id"),
|
||||
Type: c.GetInt("channel_type"),
|
||||
Name: c.GetString("channel_name"),
|
||||
AutoBan: &autoBanInt,
|
||||
}, nil
|
||||
}
|
||||
channel, selectGroup, err := service.CacheGetRandomSatisfiedChannel(retryParam)
|
||||
if err != nil {
|
||||
return nil, types.NewError(fmt.Errorf("获取分组 %s 下模型 %s 的可用渠道失败（retry）: %s", selectGroup, info.OriginModelName, err.Error()), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
if channel == nil {
|
||||
return nil, types.NewError(fmt.Errorf("分组 %s 下模型 %s 的可用渠道不存在（retry）", selectGroup, info.OriginModelName), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
|
||||
info.PriceData.GroupRatioInfo = helper.HandleGroupRatio(c, info)
|
||||
|
||||
newAPIError := middleware.SetupContextForSelectedChannel(c, channel, info.OriginModelName)
|
||||
if newAPIError != nil {
|
||||
return nil, newAPIError
|
||||
}
|
||||
return channel, nil
|
||||
}
|
||||
|
||||
func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) bool {
|
||||
if openaiErr == nil {
|
||||
return false
|
||||
}
|
||||
if service.ShouldSkipRetryAfterChannelAffinityFailure(c) {
|
||||
return false
|
||||
}
|
||||
if types.IsChannelError(openaiErr) {
|
||||
return true
|
||||
}
|
||||
if types.IsSkipRetryError(openaiErr) {
|
||||
return false
|
||||
}
|
||||
if retryTimes <= 0 {
|
||||
return false
|
||||
}
|
||||
if service.GetChannelConstraints(c).SuppressesRetry() {
|
||||
return false
|
||||
}
|
||||
code := openaiErr.StatusCode
|
||||
if code >= 200 && code < 300 {
|
||||
return false
|
||||
}
|
||||
if code < 100 || code > 599 {
|
||||
return true
|
||||
}
|
||||
if operation_setting.IsAlwaysSkipRetryCode(openaiErr.GetErrorCode()) {
|
||||
return false
|
||||
}
|
||||
return operation_setting.ShouldRetryByStatusCode(code)
|
||||
}
|
||||
|
||||
func processChannelError(c *gin.Context, channelError types.ChannelError, err *types.NewAPIError, relayInfo *relaycommon.RelayInfo) {
|
||||
logger.LogError(c, fmt.Sprintf("channel error (channel #%d, status code: %d): %s", channelError.ChannelId, err.StatusCode, common.LocalLogPreview(err.Error())))
|
||||
// 不要使用context获取渠道信息，异步处理时可能会出现渠道信息不一致的情况
|
||||
// do not use context to get channel info, there may be inconsistent channel info when processing asynchronously
|
||||
if service.ShouldDisableChannel(err) && channelError.AutoBan {
|
||||
gopool.Go(func() {
|
||||
service.DisableChannel(channelError, err.ErrorWithStatusCode())
|
||||
})
|
||||
}
|
||||
|
||||
if constant.ErrorLogEnabled && types.IsRecordErrorLog(err) {
|
||||
// 保存错误日志到mysql中
|
||||
userId := c.GetInt("id")
|
||||
tokenName := c.GetString("token_name")
|
||||
modelName := c.GetString("original_model")
|
||||
tokenId := c.GetInt("token_id")
|
||||
userGroup := c.GetString("group")
|
||||
other := model.NewLogOther()
|
||||
if c.Request != nil && c.Request.URL != nil {
|
||||
other.SetPublic("request_path", c.Request.URL.Path)
|
||||
}
|
||||
other.SetPublic("error_type", err.GetErrorType())
|
||||
other.SetPublic("error_code", err.GetErrorCode())
|
||||
other.SetPublic("status_code", err.StatusCode)
|
||||
service.AppendRelayLogAdminInfo(c, relayInfo, other)
|
||||
service.AppendTaskPluginContextAuditInfo(c, other)
|
||||
startTime := common.GetContextKeyTime(c, constant.ContextKeyRequestStartTime)
|
||||
if startTime.IsZero() {
|
||||
startTime = time.Now()
|
||||
}
|
||||
useTimeSeconds := int(time.Since(startTime).Seconds())
|
||||
other.SetRoot("upstream_error", err.MaskSensitiveErrorWithStatusCode())
|
||||
model.RecordErrorLog(c, userId, channelError.ChannelId, modelName, tokenName, publicUpstreamErrorContent(err.StatusCode), tokenId, useTimeSeconds, common.GetContextKeyBool(c, constant.ContextKeyIsStream), userGroup, other)
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
func RelayMidjourney(c *gin.Context) {
|
||||
relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatMjProxy, nil, nil)
|
||||
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{
|
||||
"description": fmt.Sprintf("failed to generate relay info: %s", err.Error()),
|
||||
"type": "upstream_error",
|
||||
"code": 4,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
var mjErr *taskdto.MidjourneyResponse
|
||||
switch relayInfo.RelayMode {
|
||||
case relayconstant.RelayModeMidjourneyNotify:
|
||||
mjErr = relay.RelayMidjourneyNotify(c)
|
||||
case relayconstant.RelayModeMidjourneyTaskFetch, relayconstant.RelayModeMidjourneyTaskFetchByCondition:
|
||||
mjErr = relay.RelayMidjourneyTask(c, relayInfo.RelayMode)
|
||||
case relayconstant.RelayModeMidjourneyTaskImageSeed:
|
||||
mjErr = relay.RelayMidjourneyTaskImageSeed(c)
|
||||
case relayconstant.RelayModeSwapFace:
|
||||
mjErr = relay.RelaySwapFace(c, relayInfo)
|
||||
default:
|
||||
mjErr = relay.RelayMidjourneySubmit(c, relayInfo)
|
||||
}
|
||||
//err = relayMidjourneySubmit(c, relayMode)
|
||||
log.Println(mjErr)
|
||||
if mjErr != nil {
|
||||
statusCode := http.StatusBadRequest
|
||||
if mjErr.Code == 30 {
|
||||
mjErr.Result = "当前分组负载已饱和，请稍后再试，或升级账户以提升服务质量。"
|
||||
statusCode = http.StatusTooManyRequests
|
||||
}
|
||||
c.JSON(statusCode, gin.H{
|
||||
"description": fmt.Sprintf("%s %s", mjErr.Description, mjErr.Result),
|
||||
"type": "upstream_error",
|
||||
"code": mjErr.Code,
|
||||
})
|
||||
channelId := c.GetInt("channel_id")
|
||||
logger.LogError(c, fmt.Sprintf("relay error (channel #%d, status code %d): %s", channelId, statusCode, fmt.Sprintf("%s %s", mjErr.Description, mjErr.Result)))
|
||||
}
|
||||
}
|
||||
|
||||
func RelayNotImplemented(c *gin.Context) {
|
||||
err := types.OpenAIError{
|
||||
Message: "API not implemented",
|
||||
Type: "new_api_error",
|
||||
Param: "",
|
||||
Code: "api_not_implemented",
|
||||
}
|
||||
c.JSON(http.StatusNotImplemented, gin.H{
|
||||
"error": err,
|
||||
})
|
||||
}
|
||||
|
||||
func RelayNotFound(c *gin.Context) {
|
||||
// The web fallback may already have applied static-asset cache headers.
|
||||
// A missing API or asset can appear after an upgrade; never cache its 404.
|
||||
c.Header("Cache-Control", "no-store, no-cache, must-revalidate, private, max-age=0")
|
||||
c.Header("Pragma", "no-cache")
|
||||
c.Header("Expires", "0")
|
||||
err := types.OpenAIError{
|
||||
Message: fmt.Sprintf("Invalid URL (%s %s)", c.Request.Method, c.Request.URL.Path),
|
||||
Type: "invalid_request_error",
|
||||
Param: "",
|
||||
Code: "",
|
||||
}
|
||||
c.JSON(http.StatusNotFound, gin.H{
|
||||
"error": err,
|
||||
})
|
||||
}
|
||||
|
||||
// RelayTaskPluginEndpoint keeps unclaimed shared-endpoint traffic on its
|
||||
// existing handler while claimed requests enter the generation-pinned
|
||||
// host-owned protocol bridge.
|
||||
func RelayTaskPluginEndpoint(c *gin.Context, fallback gin.HandlerFunc) {
|
||||
pinnedValue, exists := c.Get(pluginruntime.ContextKeyPinnedEndpoint)
|
||||
if !exists {
|
||||
fallback(c)
|
||||
return
|
||||
}
|
||||
pinned, ok := pinnedValue.(pluginruntime.PinnedEndpoint)
|
||||
if !ok || pinned.Plugin == nil || pinned.Generation == nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{
|
||||
"error": gin.H{
|
||||
"message": "Task protocol request failed",
|
||||
"type": "new_api_error",
|
||||
"code": "task_protocol_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
if pinned.Protocol != "openai_responses" {
|
||||
fallback(c)
|
||||
return
|
||||
}
|
||||
serveTaskPluginProtocol(c, pinned, defaultPluginProtocolBridgeDeps())
|
||||
}
|
||||
|
||||
func RelayTaskFetch(c *gin.Context) {
|
||||
relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, &taskdto.TaskError{
|
||||
Code: "gen_relay_info_failed",
|
||||
Message: err.Error(),
|
||||
StatusCode: http.StatusInternalServerError,
|
||||
})
|
||||
return
|
||||
}
|
||||
if taskErr := relay.RelayTaskFetch(c, relayInfo.RelayMode); taskErr != nil {
|
||||
respondTaskError(c, taskErr)
|
||||
}
|
||||
}
|
||||
|
||||
type taskSubmissionOutcome struct {
|
||||
Result *relay.TaskSubmitResult
|
||||
Task *model.Task
|
||||
RelayInfo *relaycommon.RelayInfo
|
||||
}
|
||||
|
||||
func RelayTask(c *gin.Context) {
|
||||
relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil)
|
||||
if err != nil {
|
||||
respondTaskSubmissionError(c, &taskdto.TaskError{
|
||||
Code: "gen_relay_info_failed",
|
||||
Message: err.Error(),
|
||||
StatusCode: http.StatusInternalServerError,
|
||||
})
|
||||
return
|
||||
}
|
||||
if action := c.GetString("task_action"); action != "" {
|
||||
relayInfo.Action = action
|
||||
}
|
||||
|
||||
if taskErr := relay.ResolveOriginTask(c, relayInfo); taskErr != nil {
|
||||
respondTaskSubmissionError(c, taskErr)
|
||||
return
|
||||
}
|
||||
if taskErr := relay.ApplyOriginTaskAffinity(c, relayInfo); taskErr != nil {
|
||||
respondTaskSubmissionError(c, taskErr)
|
||||
return
|
||||
}
|
||||
|
||||
outcome, taskErr := executeTaskSubmission(c, relayInfo)
|
||||
if taskErr != nil {
|
||||
respondTaskSubmissionError(c, taskErr)
|
||||
return
|
||||
}
|
||||
presentTaskSubmission(c, outcome)
|
||||
}
|
||||
|
||||
// executeTaskSubmission owns the retry, billing, and persistence lifecycle.
|
||||
// It deliberately performs no client response writes so JSON and protocol
|
||||
// presenters share the same durable task barrier. Its cancellation semantics
|
||||
// come from c.Request.Context: native task endpoints use the client context,
|
||||
// while the Responses bridge supplies an independently bounded context.
|
||||
func executeTaskSubmission(c *gin.Context, relayInfo *relaycommon.RelayInfo) (*taskSubmissionOutcome, *taskdto.TaskError) {
|
||||
return executeTaskSubmissionWith(c, relayInfo, relay.RelayTaskSubmit)
|
||||
}
|
||||
|
||||
type taskSubmitAttempt func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *taskdto.TaskError)
|
||||
|
||||
func executeTaskSubmissionWith(
|
||||
c *gin.Context,
|
||||
relayInfo *relaycommon.RelayInfo,
|
||||
submit taskSubmitAttempt,
|
||||
) (*taskSubmissionOutcome, *taskdto.TaskError) {
|
||||
diagnostics := newTaskPluginSubmitDiagnostics(c)
|
||||
diagnostics.start(relayInfo)
|
||||
var result *relay.TaskSubmitResult
|
||||
var taskErr *taskdto.TaskError
|
||||
durable := false
|
||||
stage := "start"
|
||||
defer func() {
|
||||
if !durable && relayInfo.Billing != nil {
|
||||
diagnostics.refund(stage)
|
||||
relayInfo.Billing.Refund(c)
|
||||
}
|
||||
}()
|
||||
stage = "before_attempt"
|
||||
if requestErr := c.Request.Context().Err(); requestErr != nil {
|
||||
diagnostics.cancelled("before_attempt", 0)
|
||||
return nil, service.TaskErrorWrapperLocal(requestErr, "request_cancelled", http.StatusRequestTimeout)
|
||||
}
|
||||
|
||||
retryParam := &service.RetryParam{
|
||||
Ctx: c,
|
||||
TokenGroup: relayInfo.TokenGroup,
|
||||
ModelName: relayInfo.OriginModelName,
|
||||
RequestPath: c.Request.URL.Path,
|
||||
Retry: common.GetPointer(0),
|
||||
}
|
||||
|
||||
for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() {
|
||||
stage = "select_channel"
|
||||
if requestErr := c.Request.Context().Err(); requestErr != nil {
|
||||
diagnostics.cancelled("before_attempt", retryParam.GetRetry()+1)
|
||||
taskErr = service.TaskErrorWrapperLocal(requestErr, "request_cancelled", http.StatusRequestTimeout)
|
||||
break
|
||||
}
|
||||
var channel *model.Channel
|
||||
|
||||
if lockedCh, ok := relayInfo.LockedChannel.(*model.Channel); ok && lockedCh != nil {
|
||||
channel = lockedCh
|
||||
if retryParam.GetRetry() > 0 {
|
||||
if setupErr := middleware.SetupContextForSelectedChannel(c, channel, relayInfo.OriginModelName); setupErr != nil {
|
||||
taskErr = service.TaskErrorWrapperLocal(setupErr.Err, "setup_locked_channel_failed", http.StatusInternalServerError)
|
||||
break
|
||||
}
|
||||
}
|
||||
} else {
|
||||
var channelErr *types.NewAPIError
|
||||
channel, channelErr = getChannel(c, relayInfo, retryParam)
|
||||
if channelErr != nil {
|
||||
logger.LogError(c, channelErr.Error())
|
||||
taskErr = service.TaskErrorWrapperLocal(channelErr.Err, "get_channel_failed", http.StatusInternalServerError)
|
||||
break
|
||||
}
|
||||
}
|
||||
diagnostics.attempt(retryParam.GetRetry()+1, channel, relayInfo.LockedChannel != nil)
|
||||
|
||||
addUsedChannel(c, channel.Id)
|
||||
bodyStorage, bodyErr := common.GetBodyStorage(c)
|
||||
if bodyErr != nil {
|
||||
stage = "read_body"
|
||||
if common.IsRequestBodyTooLargeError(bodyErr) || errors.Is(bodyErr, common.ErrRequestBodyTooLarge) {
|
||||
taskErr = service.TaskErrorWrapperLocal(bodyErr, "read_request_body_failed", http.StatusRequestEntityTooLarge)
|
||||
} else {
|
||||
taskErr = service.TaskErrorWrapperLocal(bodyErr, "read_request_body_failed", http.StatusBadRequest)
|
||||
}
|
||||
break
|
||||
}
|
||||
c.Request.Body = io.NopCloser(bodyStorage)
|
||||
|
||||
stage = "submit"
|
||||
result, taskErr = submit(c, relayInfo)
|
||||
if requestErr := c.Request.Context().Err(); requestErr != nil {
|
||||
diagnostics.cancelled("after_submit", retryParam.GetRetry()+1)
|
||||
taskErr = service.TaskErrorWrapperLocal(requestErr, "request_cancelled", http.StatusRequestTimeout)
|
||||
break
|
||||
}
|
||||
if taskErr == nil {
|
||||
diagnostics.attemptSucceeded(retryParam.GetRetry()+1, result)
|
||||
break
|
||||
}
|
||||
|
||||
if !taskErr.LocalError {
|
||||
processChannelError(c,
|
||||
*types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey,
|
||||
common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()),
|
||||
types.NewOpenAIError(taskErr.Error, types.ErrorCodeBadResponseStatusCode, taskErr.StatusCode),
|
||||
relayInfo)
|
||||
}
|
||||
|
||||
willRetry := shouldRetryTaskRelay(c, channel.Id, taskErr, common.RetryTimes-retryParam.GetRetry())
|
||||
diagnostics.attemptFailed(retryParam.GetRetry()+1, channel, taskErr, willRetry)
|
||||
if !willRetry {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
useChannel := c.GetStringSlice("use_channel")
|
||||
if len(useChannel) > 1 {
|
||||
retryLogStr := fmt.Sprintf("重试：%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]"))
|
||||
logger.LogInfo(c, retryLogStr)
|
||||
}
|
||||
|
||||
if taskErr != nil {
|
||||
diagnostics.failed(stage, "task_error", taskErr, false)
|
||||
return nil, taskErr
|
||||
}
|
||||
if result == nil {
|
||||
taskErr = service.TaskErrorWrapperLocal(errors.New("task submission returned no result"), "task_submit_failed", http.StatusInternalServerError)
|
||||
diagnostics.failed("submit", "missing_result", taskErr, false)
|
||||
return nil, taskErr
|
||||
}
|
||||
if requestErr := c.Request.Context().Err(); requestErr != nil {
|
||||
diagnostics.cancelled("before_reserve", retryParam.GetRetry()+1)
|
||||
return nil, service.TaskErrorWrapperLocal(requestErr, "request_cancelled", http.StatusRequestTimeout)
|
||||
}
|
||||
|
||||
// Reserve any submit-time upward billing adjustment before persistence.
|
||||
// This keeps insertion failures fully refundable while ensuring settlement
|
||||
// after the barrier normally has a zero positive delta.
|
||||
if relayInfo.Billing != nil {
|
||||
stage = "reserve"
|
||||
diagnostics.reserve("reserve_start", result.Quota)
|
||||
if reserveErr := relayInfo.Billing.Reserve(result.Quota); reserveErr != nil {
|
||||
common.SysError("reserve adjusted task billing error: " + reserveErr.Error())
|
||||
taskErr = service.TaskErrorWrapperLocal(errors.New("insufficient quota for adjusted task cost"), string(types.ErrorCodeInsufficientUserQuota), http.StatusForbidden)
|
||||
diagnostics.failed("reserve", "insufficient_quota", taskErr, false)
|
||||
return nil, taskErr
|
||||
}
|
||||
diagnostics.reserve("reserve_complete", result.Quota)
|
||||
}
|
||||
if requestErr := c.Request.Context().Err(); requestErr != nil {
|
||||
diagnostics.cancelled("before_insert", retryParam.GetRetry()+1)
|
||||
return nil, service.TaskErrorWrapperLocal(requestErr, "request_cancelled", http.StatusRequestTimeout)
|
||||
}
|
||||
|
||||
stage = "insert"
|
||||
task := model.InitTask(result.Platform, relayInfo)
|
||||
task.PrivateData.Execution = service.TaskExecutionSnapshotFromContext(c)
|
||||
task.PrivateData.UpstreamTaskID = result.UpstreamTaskID
|
||||
task.PrivateData.BillingSource = relayInfo.BillingSource
|
||||
task.PrivateData.SubscriptionId = relayInfo.SubscriptionId
|
||||
task.PrivateData.TokenId = relayInfo.TokenId
|
||||
task.PrivateData.NodeName = common.NodeName
|
||||
task.PrivateData.BillingContext = &model.TaskBillingContext{
|
||||
ModelPrice: relayInfo.PriceData.ModelPrice,
|
||||
GroupRatio: relayInfo.PriceData.GroupRatioInfo.GroupRatio,
|
||||
ModelRatio: relayInfo.PriceData.ModelRatio,
|
||||
OtherRatios: relayInfo.PriceData.OtherRatios(),
|
||||
OriginModelName: relayInfo.OriginModelName,
|
||||
PerCallBilling: common.StringsContains(constant.TaskPricePatches, relayInfo.OriginModelName) || relayInfo.PriceData.UsePrice,
|
||||
TieredSnapshot: relayInfo.TieredBillingSnapshot,
|
||||
}
|
||||
task.Quota = result.Quota
|
||||
task.Data = result.TaskData
|
||||
if len(result.PluginState) > 0 {
|
||||
task.PrivateData.PluginState = result.PluginState
|
||||
}
|
||||
task.Action = relayInfo.Action
|
||||
if immediate := result.Immediate; immediate != nil {
|
||||
task.Status = model.TaskStatus(immediate.Status)
|
||||
task.Progress = immediate.Progress
|
||||
if immediate.Status == model.TaskStatusSuccess || immediate.Status == model.TaskStatusFailure {
|
||||
task.FinishTime = time.Now().Unix()
|
||||
}
|
||||
if immediate.Status == model.TaskStatusFailure {
|
||||
task.FailReason = immediate.Reason
|
||||
}
|
||||
if immediate.Url != "" {
|
||||
task.PrivateData.ResultURL = immediate.Url
|
||||
} else if immediate.Status == model.TaskStatusSuccess {
|
||||
task.PrivateData.ResultURL = taskcommon.BuildProxyURL(task.TaskID)
|
||||
}
|
||||
}
|
||||
diagnostics.insertStart(task)
|
||||
if insertErr := task.InsertWithContext(c.Request.Context()); insertErr != nil {
|
||||
common.SysError("insert task error: " + insertErr.Error())
|
||||
taskErr = service.TaskErrorWrapperLocal(errors.New("failed to persist task"), "task_insert_failed", http.StatusInternalServerError)
|
||||
diagnostics.failed("insert", "database_error", taskErr, false)
|
||||
return nil, taskErr
|
||||
}
|
||||
durable = true
|
||||
stage = "settle"
|
||||
diagnostics.durable(task)
|
||||
diagnostics.settleStart(task, result.Quota)
|
||||
|
||||
if settleErr := service.SettleBilling(c, relayInfo, result.Quota); settleErr != nil {
|
||||
common.SysError("settle task billing error: " + settleErr.Error())
|
||||
taskErr = service.TaskErrorWrapperLocal(errors.New("failed to settle task billing"), "task_billing_settlement_failed", http.StatusInternalServerError)
|
||||
diagnostics.failed("settle", "billing_error", taskErr, true)
|
||||
return nil, taskErr
|
||||
}
|
||||
service.LogTaskConsumption(c, relayInfo, task)
|
||||
diagnostics.complete(task, result.Quota)
|
||||
|
||||
return &taskSubmissionOutcome{Result: result, Task: task, RelayInfo: relayInfo}, nil
|
||||
}
|
||||
|
||||
func presentTaskSubmission(c *gin.Context, outcome *taskSubmissionOutcome) {
|
||||
diagnostics := newTaskPluginSubmitDiagnostics(c)
|
||||
otherRatios := outcome.RelayInfo.PriceData.OtherRatios()
|
||||
if otherRatios == nil {
|
||||
otherRatios = map[string]float64{}
|
||||
}
|
||||
if ratiosJSON, err := common.Marshal(otherRatios); err == nil {
|
||||
c.Header("X-New-Api-Other-Ratios", string(ratiosJSON))
|
||||
}
|
||||
if pinnedValue, exists := c.Get(pluginruntime.ContextKeyPinnedRoute); exists {
|
||||
if pinned, ok := pinnedValue.(pluginruntime.PinnedRoute); ok && pinned.Plugin != nil && pinned.Route.Render != "" {
|
||||
view, err := service.BuildTaskPluginView(outcome.Task)
|
||||
requestValue, _ := c.Get(pluginruntime.ContextKeyRouteRequest)
|
||||
requestContext, _ := requestValue.(pluginruntime.RouteRequestContext)
|
||||
if err == nil {
|
||||
viewValue, valueErr := taskPluginProtocolJSONValue(view)
|
||||
if valueErr == nil {
|
||||
if body, callErr := pinned.Plugin.Engine.CallPath(c.Request.Context(), "native", []string{pinned.Route.Render}, requestContext.JSValue(), viewValue); callErr == nil {
|
||||
diagnostics.present(outcome.Task, "native_presenter")
|
||||
c.JSON(http.StatusOK, body)
|
||||
return
|
||||
} else {
|
||||
logger.LogError(c, "task plugin native submit presenter failed: "+callErr.Error())
|
||||
}
|
||||
} else {
|
||||
logger.LogError(c, "encode task plugin native submit view failed: "+valueErr.Error())
|
||||
}
|
||||
} else {
|
||||
logger.LogError(c, "build task plugin native submit view failed: "+err.Error())
|
||||
}
|
||||
}
|
||||
}
|
||||
if pinnedValue, exists := c.Get(pluginruntime.ContextKeyPinnedEndpoint); exists {
|
||||
if pinned, ok := pinnedValue.(pluginruntime.PinnedEndpoint); ok && pinned.Protocol == "openai_video" && pinned.Operation.Name == "create" {
|
||||
diagnostics.present(outcome.Task, "openai_video_create")
|
||||
c.JSON(http.StatusOK, outcome.Task.ToOpenAIVideo())
|
||||
return
|
||||
}
|
||||
}
|
||||
createdAt := outcome.Task.CreatedAt
|
||||
if createdAt == 0 {
|
||||
createdAt = outcome.Task.SubmitTime
|
||||
}
|
||||
diagnostics.present(outcome.Task, "host_fallback")
|
||||
c.JSON(http.StatusOK, map[string]any{
|
||||
"id": outcome.Task.TaskID,
|
||||
"task_id": outcome.Task.TaskID,
|
||||
"status": "queued",
|
||||
"model": outcome.RelayInfo.OriginModelName,
|
||||
"created_at": createdAt,
|
||||
})
|
||||
}
|
||||
|
||||
func respondTaskSubmissionError(c *gin.Context, taskErr *taskdto.TaskError) {
|
||||
newTaskPluginSubmitDiagnostics(c).presentError(taskErr)
|
||||
if middleware.RespondTaskPluginError(c, taskErr) {
|
||||
return
|
||||
}
|
||||
respondTaskError(c, taskErr)
|
||||
}
|
||||
|
||||
// respondTaskError 统一输出 Task 错误响应（含 429 限流提示改写）
|
||||
func respondTaskError(c *gin.Context, taskErr *taskdto.TaskError) {
|
||||
if taskErr.StatusCode == http.StatusTooManyRequests {
|
||||
taskErr.Message = "当前分组上游负载已饱和，请稍后再试"
|
||||
}
|
||||
c.JSON(taskErr.StatusCode, taskErr)
|
||||
}
|
||||
|
||||
func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *taskdto.TaskError, retryTimes int) bool {
|
||||
if taskErr == nil {
|
||||
return false
|
||||
}
|
||||
if service.ShouldSkipRetryAfterChannelAffinityFailure(c) {
|
||||
return false
|
||||
}
|
||||
if retryTimes <= 0 {
|
||||
return false
|
||||
}
|
||||
if service.GetChannelConstraints(c).SuppressesRetry() {
|
||||
return false
|
||||
}
|
||||
if taskErr.StatusCode == http.StatusTooManyRequests {
|
||||
return true
|
||||
}
|
||||
if taskErr.StatusCode == 307 {
|
||||
return true
|
||||
}
|
||||
if taskErr.StatusCode/100 == 5 {
|
||||
// 超时不重试
|
||||
if operation_setting.IsAlwaysSkipRetryStatusCode(taskErr.StatusCode) {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
if taskErr.StatusCode == http.StatusBadRequest {
|
||||
return false
|
||||
}
|
||||
if taskErr.StatusCode == 408 {
|
||||
// azure处理超时不重试
|
||||
return false
|
||||
}
|
||||
if taskErr.LocalError {
|
||||
return false
|
||||
}
|
||||
if taskErr.StatusCode/100 == 2 {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
)
|
||||
|
||||
func TestClientFacingOpenAIErrorHidesUpstreamDetails(t *testing.T) {
|
||||
err := types.NewOpenAIError(errors.New("provider.example.invalid provider secret failure"), types.ErrorCodeBadResponseStatusCode, http.StatusBadGateway)
|
||||
got := clientFacingOpenAIError(err, true, "local-request-id")
|
||||
|
||||
if strings.Contains(got.Message, "example-provider") || strings.Contains(got.Message, "provider secret") {
|
||||
t.Fatalf("upstream details leaked: %q", got.Message)
|
||||
}
|
||||
if !strings.Contains(got.Message, "local-request-id") {
|
||||
t.Fatalf("local request id missing: %q", got.Message)
|
||||
}
|
||||
if got.Code != "model_service_error" {
|
||||
t.Fatalf("unexpected public code: %v", got.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientFacingClaudeErrorHidesUpstreamDetails(t *testing.T) {
|
||||
err := types.NewOpenAIError(errors.New("provider.example.invalid provider secret failure"), types.ErrorCodeBadResponseStatusCode, http.StatusBadGateway)
|
||||
got := clientFacingClaudeError(err, true, "local-request-id")
|
||||
|
||||
if strings.Contains(got.Message, "example-provider") || strings.Contains(got.Message, "provider secret") {
|
||||
t.Fatalf("upstream details leaked: %q", got.Message)
|
||||
}
|
||||
if !strings.Contains(got.Message, "local-request-id") {
|
||||
t.Fatalf("local request id missing: %q", got.Message)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientFacingErrorKeepsLocalValidationMessage(t *testing.T) {
|
||||
err := types.NewErrorWithStatusCode(errors.New("model is required"), types.ErrorCodeInvalidRequest, http.StatusBadRequest)
|
||||
got := clientFacingOpenAIError(err, false, "local-request-id")
|
||||
if !strings.Contains(got.Message, "model is required") {
|
||||
t.Fatalf("local validation message was hidden: %q", got.Message)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestCountClaudeTokensReturnsInputTokensWhenRelayCountingDisabled(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
originalCountToken := constant.CountToken
|
||||
constant.CountToken = false
|
||||
t.Cleanup(func() {
|
||||
constant.CountToken = originalCountToken
|
||||
})
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(
|
||||
http.MethodPost,
|
||||
"/v1/messages/count_tokens?beta=true",
|
||||
strings.NewReader(`{
|
||||
"model":"gemini-3.6-flash",
|
||||
"messages":[{"role":"user","content":"count this prompt"}],
|
||||
"tools":[{"name":"lookup","description":"Look up a value","input_schema":{"type":"object","properties":{"query":{"type":"string"}}}}]
|
||||
}`),
|
||||
)
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
common.SetContextKey(ctx, constant.ContextKeyOriginalModel, "gemini-3.6-flash")
|
||||
|
||||
CountClaudeTokens(ctx)
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
var response struct {
|
||||
InputTokens int `json:"input_tokens"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
|
||||
assert.Positive(t, response.InputTokens)
|
||||
}
|
||||
|
||||
func TestCountClaudeTokensRejectsMissingMessages(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(
|
||||
http.MethodPost,
|
||||
"/v1/messages/count_tokens",
|
||||
strings.NewReader(`{"model":"gemini-3.6-flash"}`),
|
||||
)
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
CountClaudeTokens(ctx)
|
||||
|
||||
require.Equal(t, http.StatusBadRequest, recorder.Code)
|
||||
var response struct {
|
||||
Type string `json:"type"`
|
||||
Error struct {
|
||||
Type string `json:"type"`
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
|
||||
assert.Equal(t, "error", response.Type)
|
||||
assert.Equal(t, "invalid_request_error", response.Error.Type)
|
||||
assert.Contains(t, response.Error.Message, "messages")
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func TestProcessChannelErrorUsesSnapshotWithoutLeakingChannelMetadata(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
previousDB, previousLogDB := model.DB, model.LOG_DB
|
||||
previousRedisEnabled := common.RedisEnabled
|
||||
previousMainDatabaseType := common.MainDatabaseType()
|
||||
previousLogDatabaseType := common.LogDatabaseType()
|
||||
previousErrorLogEnabled := constant.ErrorLogEnabled
|
||||
|
||||
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := database.DB()
|
||||
require.NoError(t, err)
|
||||
sqlDB.SetMaxOpenConns(1)
|
||||
require.NoError(t, database.AutoMigrate(&model.User{}, &model.Log{}))
|
||||
model.DB, model.LOG_DB = database, database
|
||||
common.RedisEnabled = false
|
||||
common.SetDatabaseTypes(common.DatabaseTypeSQLite, common.DatabaseTypeSQLite)
|
||||
constant.ErrorLogEnabled = true
|
||||
t.Cleanup(func() {
|
||||
model.DB, model.LOG_DB = previousDB, previousLogDB
|
||||
common.RedisEnabled = previousRedisEnabled
|
||||
common.SetDatabaseTypes(previousMainDatabaseType, previousLogDatabaseType)
|
||||
constant.ErrorLogEnabled = previousErrorLogEnabled
|
||||
require.NoError(t, sqlDB.Close())
|
||||
})
|
||||
|
||||
require.NoError(t, database.Create(&model.User{Id: 7, Username: "log-owner", Group: "default"}).Error)
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
ctx.Set("id", 7)
|
||||
ctx.Set("username", "log-owner")
|
||||
ctx.Set("token_name", "test-token")
|
||||
ctx.Set("token_id", 11)
|
||||
ctx.Set("original_model", "gpt-test")
|
||||
ctx.Set("group", "default")
|
||||
ctx.Set("channel_id", 202)
|
||||
ctx.Set("channel_name", "mutable-context-channel")
|
||||
ctx.Set("channel_type", 9)
|
||||
ctx.Set("use_channel", []string{"101"})
|
||||
common.SetContextKey(ctx, constant.ContextKeyRequestStartTime, time.Now().Add(-time.Second))
|
||||
|
||||
channelSnapshot := types.ChannelError{
|
||||
ChannelId: 101,
|
||||
ChannelType: 1,
|
||||
ChannelName: "snapshot-channel",
|
||||
AutoBan: false,
|
||||
}
|
||||
apiErr := types.NewOpenAIError(errors.New("upstream failed"), types.ErrorCodeBadResponseStatusCode, http.StatusBadGateway)
|
||||
|
||||
processChannelError(ctx, channelSnapshot, apiErr, nil)
|
||||
|
||||
var stored model.Log
|
||||
require.NoError(t, database.First(&stored).Error)
|
||||
assert.Equal(t, channelSnapshot.ChannelId, stored.ChannelId)
|
||||
storedOther, err := common.StrToMap(stored.Other)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, float64(http.StatusBadGateway), storedOther["status_code"])
|
||||
for _, key := range []string{"channel_id", "channel_name", "channel_type"} {
|
||||
assert.NotContains(t, storedOther, key)
|
||||
}
|
||||
adminInfo, ok := storedOther["admin_info"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, []any{"101"}, adminInfo["use_channel"])
|
||||
|
||||
logs, total, err := model.GetUserLogs(7, model.LogTypeError, 0, 0, "", "", 0, 10, "", "", "")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), total)
|
||||
require.Len(t, logs, 1)
|
||||
assert.Equal(t, channelSnapshot.ChannelId, logs[0].ChannelId)
|
||||
assert.Empty(t, logs[0].ChannelName)
|
||||
userOther, err := common.StrToMap(logs[0].Other)
|
||||
require.NoError(t, err)
|
||||
assert.NotContains(t, userOther, "admin_info")
|
||||
for _, key := range []string{"channel_id", "channel_name", "channel_type"} {
|
||||
assert.NotContains(t, userOther, key)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestPublicUpstreamErrorContent(t *testing.T) {
|
||||
got := publicUpstreamErrorContent(502)
|
||||
if strings.Contains(got, "example-provider") || strings.Contains(got, "provider") || strings.Contains(got, "upstream") {
|
||||
t.Fatalf("upstream details leaked: %q", got)
|
||||
}
|
||||
if !strings.Contains(got, "status_code=502") || !strings.Contains(got, publicUpstreamErrorMessage) {
|
||||
t.Fatalf("unexpected public error content: %q", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,389 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/relay"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/types"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type taskSubmissionTestBilling struct {
|
||||
events *[]string
|
||||
settleErr error
|
||||
onSettle func()
|
||||
refunds int
|
||||
}
|
||||
|
||||
func (b *taskSubmissionTestBilling) Settle(int) error {
|
||||
*b.events = append(*b.events, "settle")
|
||||
if b.onSettle != nil {
|
||||
b.onSettle()
|
||||
}
|
||||
return b.settleErr
|
||||
}
|
||||
|
||||
func (b *taskSubmissionTestBilling) Refund(*gin.Context) {
|
||||
*b.events = append(*b.events, "refund")
|
||||
b.refunds++
|
||||
}
|
||||
|
||||
func (b *taskSubmissionTestBilling) NeedsRefund() bool { return b.refunds == 0 }
|
||||
func (b *taskSubmissionTestBilling) GetPreConsumedQuota() int { return 0 }
|
||||
func (b *taskSubmissionTestBilling) Reserve(int) error {
|
||||
*b.events = append(*b.events, "reserve")
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestPresentTaskSubmissionUsesNativePresenterAfterPersistence(t *testing.T) {
|
||||
plugin, err := pluginruntime.CompilePlugin(`
|
||||
export const meta = {apiVersion:1,key:"presenter-test",name:"Presenter",version:"1.0.0",author:{name:"Test"},models:["model"],fetchMode:"per_task",routes:[{method:"POST",path:"/vendor/jobs",type:"submit",decode:"decode",render:"created"}]};
|
||||
export const native = {decode:function(ctx){return {kind:"submit",model:"model",requestBody:ctx.body.value};},created:function(ctx,task){return {data:{task_id:task.task_id},upstream:task.data};}};
|
||||
export function buildSubmitRequest(){return {}} export function parseSubmitResponse(){return {taskId:"upstream"}} export function buildQueryRequest(){return {}} export function parseTaskResult(){return {status:"SUCCESS"}}
|
||||
`, pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
priceData := types.PriceData{}
|
||||
priceData.AddOtherRatio("seconds", 5)
|
||||
task := &model.Task{TaskID: "task_public", SubmitTime: 123}
|
||||
task.SetData(map[string]any{"task_id": "upstream_private"})
|
||||
outcome := &taskSubmissionOutcome{
|
||||
Result: &relay.TaskSubmitResult{},
|
||||
Task: task,
|
||||
RelayInfo: &relaycommon.RelayInfo{PriceData: priceData},
|
||||
}
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{"model":"model"}`))
|
||||
c.Set(pluginruntime.ContextKeyPinnedRoute, pluginruntime.PinnedRoute{Plugin: plugin, Route: plugin.Meta.Routes[0]})
|
||||
c.Set(pluginruntime.ContextKeyRouteRequest, pluginruntime.RouteRequestContext{Path: "/vendor/jobs", Method: http.MethodPost, Body: map[string]any{"kind": "json", "value": map[string]any{"model": "model"}}})
|
||||
|
||||
presentTaskSubmission(c, outcome)
|
||||
|
||||
assert.JSONEq(t, `{
|
||||
"data":{"task_id":"task_public"},
|
||||
"upstream":{"task_id":"upstream_private"}
|
||||
}`, recorder.Body.String())
|
||||
assert.JSONEq(t, `{"seconds":5}`, recorder.Header().Get("X-New-Api-Other-Ratios"))
|
||||
}
|
||||
|
||||
func TestPresentTaskSubmissionFallbackUsesPersistedPublicID(t *testing.T) {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
outcome := &taskSubmissionOutcome{
|
||||
Result: &relay.TaskSubmitResult{},
|
||||
Task: &model.Task{TaskID: "task_persisted", SubmitTime: 456},
|
||||
RelayInfo: &relaycommon.RelayInfo{OriginModelName: "video-model"},
|
||||
}
|
||||
|
||||
presentTaskSubmission(c, outcome)
|
||||
|
||||
assert.JSONEq(t, `{
|
||||
"id":"task_persisted",
|
||||
"task_id":"task_persisted",
|
||||
"status":"queued",
|
||||
"model":"video-model",
|
||||
"created_at":456
|
||||
}`, recorder.Body.String())
|
||||
}
|
||||
|
||||
func TestPresentTaskSubmissionUsesHostOpenAIVideoCreateReceipt(t *testing.T) {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Set(pluginruntime.ContextKeyPinnedEndpoint, pluginruntime.PinnedEndpoint{
|
||||
Protocol: "openai_video",
|
||||
Operation: pluginruntime.HostProtocolOperation{Name: "create"},
|
||||
})
|
||||
task := &model.Task{
|
||||
TaskID: "task_public",
|
||||
Status: model.TaskStatusSubmitted,
|
||||
Progress: "0%",
|
||||
CreatedAt: 456,
|
||||
Properties: model.Properties{OriginModelName: "video-model"},
|
||||
}
|
||||
outcome := &taskSubmissionOutcome{Result: &relay.TaskSubmitResult{}, Task: task, RelayInfo: &relaycommon.RelayInfo{}}
|
||||
|
||||
presentTaskSubmission(c, outcome)
|
||||
|
||||
assert.JSONEq(t, `{"id":"task_public","object":"video","model":"video-model","status":"queued","progress":0,"created_at":456}`, recorder.Body.String())
|
||||
assert.NotContains(t, recorder.Body.String(), "task_id")
|
||||
}
|
||||
|
||||
func TestExecuteTaskSubmissionRefundsWhenInsertFails(t *testing.T) {
|
||||
events := make([]string, 0, 3)
|
||||
database := setupTaskSubmissionDatabase(t, false, &events)
|
||||
_ = database
|
||||
billing := &taskSubmissionTestBilling{events: &events}
|
||||
c := taskSubmissionTestContext()
|
||||
info := taskSubmissionRelayInfo(billing)
|
||||
|
||||
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
||||
return &relay.TaskSubmitResult{
|
||||
UpstreamTaskID: "upstream_private",
|
||||
Platform: constant.TaskPlatform("plugin"),
|
||||
}, nil
|
||||
})
|
||||
|
||||
assert.Nil(t, outcome)
|
||||
require.NotNil(t, taskErr)
|
||||
assert.Equal(t, "task_insert_failed", taskErr.Code)
|
||||
assert.Equal(t, []string{"reserve", "insert", "refund"}, events)
|
||||
assert.Equal(t, 1, billing.refunds)
|
||||
assert.False(t, c.Writer.Written())
|
||||
}
|
||||
|
||||
func TestExecuteTaskSubmissionSettlementFailureStaysDurableAndWritesNothing(t *testing.T) {
|
||||
events := make([]string, 0, 3)
|
||||
database := setupTaskSubmissionDatabase(t, true, &events)
|
||||
billing := &taskSubmissionTestBilling{events: &events, settleErr: errors.New("settlement failed")}
|
||||
c := taskSubmissionTestContext()
|
||||
info := taskSubmissionRelayInfo(billing)
|
||||
|
||||
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
||||
return &relay.TaskSubmitResult{
|
||||
UpstreamTaskID: "upstream_private",
|
||||
Platform: constant.TaskPlatform("plugin"),
|
||||
}, nil
|
||||
})
|
||||
|
||||
assert.Nil(t, outcome)
|
||||
require.NotNil(t, taskErr)
|
||||
assert.Equal(t, "task_billing_settlement_failed", taskErr.Code)
|
||||
assert.Equal(t, []string{"reserve", "insert", "settle"}, events)
|
||||
assert.Zero(t, billing.refunds)
|
||||
var count int64
|
||||
require.NoError(t, database.Model(&model.Task{}).Where("task_id = ?", "task_public").Count(&count).Error)
|
||||
assert.Equal(t, int64(1), count)
|
||||
assert.False(t, c.Writer.Written())
|
||||
}
|
||||
|
||||
func TestExecuteTaskSubmissionPersistsPinnedPluginProvenance(t *testing.T) {
|
||||
events := make([]string, 0, 3)
|
||||
database := setupTaskSubmissionDatabase(t, true, &events)
|
||||
previousLogConsumeEnabled := common.LogConsumeEnabled
|
||||
common.LogConsumeEnabled = false
|
||||
t.Cleanup(func() { common.LogConsumeEnabled = previousLogConsumeEnabled })
|
||||
|
||||
c := taskSubmissionTestContext()
|
||||
c.Set(common.RequestIdKey, "request-public")
|
||||
c.Set(pluginruntime.ContextKeyPinnedPlugin, pluginruntime.PinnedPlugin{
|
||||
Generation: &pluginruntime.RoutingGeneration{Number: 42},
|
||||
Plugin: &pluginruntime.LoadedPlugin{Meta: pluginruntime.Meta{
|
||||
Key: "document-parser",
|
||||
Name: "Document Parser",
|
||||
Version: "1.2.3",
|
||||
APIVersion: 1,
|
||||
Author: pluginruntime.AuthorMeta{
|
||||
Name: "Community Author",
|
||||
URL: "https://plugins.example/author",
|
||||
},
|
||||
}},
|
||||
})
|
||||
billing := &taskSubmissionTestBilling{events: &events}
|
||||
info := taskSubmissionRelayInfo(billing)
|
||||
|
||||
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
||||
return &relay.TaskSubmitResult{
|
||||
UpstreamTaskID: "upstream-private",
|
||||
Platform: constant.TaskPlatform("document-parser"),
|
||||
}, nil
|
||||
})
|
||||
|
||||
require.Nil(t, taskErr)
|
||||
require.NotNil(t, outcome)
|
||||
require.NotNil(t, outcome.Task.PrivateData.Execution)
|
||||
require.NotNil(t, outcome.Task.PrivateData.Execution.TaskPlugin)
|
||||
assert.Equal(t, "request-public", outcome.Task.PrivateData.Execution.RequestID)
|
||||
assert.Equal(t, "/plugin/submit", outcome.Task.PrivateData.Execution.RequestPath)
|
||||
assert.Equal(t, "1.2.3", outcome.Task.PrivateData.Execution.TaskPlugin.Version)
|
||||
assert.Equal(t, uint64(42), outcome.Task.PrivateData.Execution.TaskPlugin.Generation)
|
||||
require.NotNil(t, outcome.Task.PrivateData.Execution.TaskPlugin.Author)
|
||||
assert.Equal(t, "Community Author", outcome.Task.PrivateData.Execution.TaskPlugin.Author.Name)
|
||||
assert.Equal(t, "https://plugins.example/author", outcome.Task.PrivateData.Execution.TaskPlugin.Author.URL)
|
||||
|
||||
var stored model.Task
|
||||
require.NoError(t, database.Where("task_id = ?", "task_public").First(&stored).Error)
|
||||
require.NotNil(t, stored.PrivateData.Execution)
|
||||
require.NotNil(t, stored.PrivateData.Execution.TaskPlugin)
|
||||
assert.Equal(t, "document-parser", stored.PrivateData.Execution.TaskPlugin.Key)
|
||||
require.NotNil(t, stored.PrivateData.Execution.TaskPlugin.Author)
|
||||
assert.Equal(t, "Community Author", stored.PrivateData.Execution.TaskPlugin.Author.Name)
|
||||
assert.Equal(t, "upstream-private", stored.PrivateData.UpstreamTaskID)
|
||||
}
|
||||
|
||||
func TestExecuteTaskSubmissionRefundsCancellationBeforeDurableBarrier(t *testing.T) {
|
||||
events := make([]string, 0, 2)
|
||||
setupTaskSubmissionDatabase(t, true, &events)
|
||||
billing := &taskSubmissionTestBilling{events: &events}
|
||||
c := taskSubmissionTestContext()
|
||||
requestContext, cancel := context.WithCancel(c.Request.Context())
|
||||
c.Request = c.Request.WithContext(requestContext)
|
||||
info := taskSubmissionRelayInfo(billing)
|
||||
|
||||
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
||||
cancel()
|
||||
return &relay.TaskSubmitResult{
|
||||
UpstreamTaskID: "upstream_private",
|
||||
Platform: constant.TaskPlatform("plugin"),
|
||||
}, nil
|
||||
})
|
||||
|
||||
assert.Nil(t, outcome)
|
||||
require.NotNil(t, taskErr)
|
||||
assert.Equal(t, "request_cancelled", taskErr.Code)
|
||||
assert.Equal(t, []string{"refund"}, events)
|
||||
assert.Equal(t, 1, billing.refunds)
|
||||
assert.False(t, c.Writer.Written())
|
||||
}
|
||||
|
||||
func TestExecuteTaskSubmissionDisconnectBeforeUpstreamAcceptanceSkipsSubmitAndRefunds(t *testing.T) {
|
||||
events := make([]string, 0, 1)
|
||||
setupTaskSubmissionDatabase(t, true, &events)
|
||||
billing := &taskSubmissionTestBilling{events: &events}
|
||||
c := taskSubmissionTestContext()
|
||||
requestContext, cancel := context.WithCancel(c.Request.Context())
|
||||
cancel()
|
||||
c.Request = c.Request.WithContext(requestContext)
|
||||
info := taskSubmissionRelayInfo(billing)
|
||||
submitted := false
|
||||
|
||||
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
||||
submitted = true
|
||||
return nil, nil
|
||||
})
|
||||
|
||||
assert.Nil(t, outcome)
|
||||
require.NotNil(t, taskErr)
|
||||
assert.Equal(t, "request_cancelled", taskErr.Code)
|
||||
assert.False(t, submitted)
|
||||
assert.Equal(t, []string{"refund"}, events)
|
||||
assert.Equal(t, 1, billing.refunds)
|
||||
assert.False(t, c.Writer.Written())
|
||||
}
|
||||
|
||||
func TestExecuteTaskSubmissionCallerCancellationDuringSubmitRefundsBeforeDurableBarrier(t *testing.T) {
|
||||
events := make([]string, 0, 1)
|
||||
setupTaskSubmissionDatabase(t, true, &events)
|
||||
billing := &taskSubmissionTestBilling{events: &events}
|
||||
c := taskSubmissionTestContext()
|
||||
requestContext, cancel := context.WithCancel(c.Request.Context())
|
||||
c.Request = c.Request.WithContext(requestContext)
|
||||
info := taskSubmissionRelayInfo(billing)
|
||||
submitStarted := make(chan struct{})
|
||||
done := make(chan struct{})
|
||||
var outcome *taskSubmissionOutcome
|
||||
var taskErr *dto.TaskError
|
||||
|
||||
go func() {
|
||||
defer close(done)
|
||||
outcome, taskErr = executeTaskSubmissionWith(c, info, func(c *gin.Context, _ *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
||||
close(submitStarted)
|
||||
<-c.Request.Context().Done()
|
||||
return nil, service.TaskErrorWrapperLocal(c.Request.Context().Err(), "do_request_failed", http.StatusInternalServerError)
|
||||
})
|
||||
}()
|
||||
select {
|
||||
case <-submitStarted:
|
||||
case <-time.After(2 * time.Second):
|
||||
require.FailNow(t, "submission did not start")
|
||||
}
|
||||
cancel()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(2 * time.Second):
|
||||
require.FailNow(t, "submission did not stop after disconnect")
|
||||
}
|
||||
|
||||
assert.Nil(t, outcome)
|
||||
require.NotNil(t, taskErr)
|
||||
assert.Equal(t, "request_cancelled", taskErr.Code)
|
||||
assert.Equal(t, []string{"refund"}, events)
|
||||
assert.Equal(t, 1, billing.refunds)
|
||||
assert.False(t, c.Writer.Written())
|
||||
}
|
||||
|
||||
func TestExecuteTaskSubmissionDisconnectAfterDurableInsertDoesNotRefund(t *testing.T) {
|
||||
events := make([]string, 0, 3)
|
||||
database := setupTaskSubmissionDatabase(t, true, &events)
|
||||
previousLogConsumeEnabled := common.LogConsumeEnabled
|
||||
common.LogConsumeEnabled = false
|
||||
t.Cleanup(func() { common.LogConsumeEnabled = previousLogConsumeEnabled })
|
||||
c := taskSubmissionTestContext()
|
||||
requestContext, cancel := context.WithCancel(c.Request.Context())
|
||||
c.Request = c.Request.WithContext(requestContext)
|
||||
billing := &taskSubmissionTestBilling{
|
||||
events: &events,
|
||||
onSettle: cancel,
|
||||
}
|
||||
info := taskSubmissionRelayInfo(billing)
|
||||
|
||||
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
||||
return &relay.TaskSubmitResult{
|
||||
UpstreamTaskID: "upstream_private",
|
||||
Platform: constant.TaskPlatform("plugin"),
|
||||
}, nil
|
||||
})
|
||||
|
||||
require.Nil(t, taskErr)
|
||||
require.NotNil(t, outcome)
|
||||
assert.Equal(t, "task_public", outcome.Task.TaskID)
|
||||
assert.Equal(t, []string{"reserve", "insert", "settle"}, events)
|
||||
assert.Zero(t, billing.refunds)
|
||||
var count int64
|
||||
require.NoError(t, database.Model(&model.Task{}).Where("task_id = ?", "task_public").Count(&count).Error)
|
||||
assert.Equal(t, int64(1), count)
|
||||
assert.False(t, c.Writer.Written())
|
||||
}
|
||||
|
||||
func setupTaskSubmissionDatabase(t *testing.T, migrate bool, events *[]string) *gorm.DB {
|
||||
t.Helper()
|
||||
previousDB := model.DB
|
||||
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, database.Callback().Create().Before("gorm:create").Register("test:task-submit-order", func(*gorm.DB) {
|
||||
*events = append(*events, "insert")
|
||||
}))
|
||||
if migrate {
|
||||
require.NoError(t, database.AutoMigrate(&model.Task{}))
|
||||
}
|
||||
model.DB = database
|
||||
t.Cleanup(func() { model.DB = previousDB })
|
||||
return database
|
||||
}
|
||||
|
||||
func taskSubmissionTestContext() *gin.Context {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/plugin/submit", strings.NewReader(`{}`))
|
||||
return c
|
||||
}
|
||||
|
||||
func taskSubmissionRelayInfo(billing relaycommon.BillingSettler) *relaycommon.RelayInfo {
|
||||
return &relaycommon.RelayInfo{
|
||||
UserId: 1,
|
||||
UsingGroup: "default",
|
||||
OriginModelName: "plugin-model",
|
||||
Billing: billing,
|
||||
TaskRelayInfo: &relaycommon.TaskRelayInfo{
|
||||
PublicTaskID: "task_public",
|
||||
LockedChannel: &model.Channel{Id: 1, Type: constant.ChannelTypeTaskPlugin, Name: "plugin"},
|
||||
},
|
||||
ChannelMeta: &relaycommon.ChannelMeta{ChannelId: 1, ChannelType: constant.ChannelTypeTaskPlugin},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestRequestStatusReplacesReachableProbe(t *testing.T) {
|
||||
monitor := Monitor{Status: 1, Uptime: 1, LastChecked: "probe"}
|
||||
stats := perfmetrics.GroupResult{Group: "test", SuccessRate: 67.9, AvgTtftMs: 7420, Series: []perfmetrics.BucketPoint{{Ts: time.Now().Unix(), SuccessRate: 67.9, AvgTtftMs: 7420}}}
|
||||
applyRequestMetrics(&monitor, stats)
|
||||
require.Equal(t, 0, monitor.Status)
|
||||
require.InDelta(t, 0.679, monitor.Uptime, 0.00001)
|
||||
require.Equal(t, 7420, monitor.History[0].Ping)
|
||||
require.Equal(t, int64(7420), monitor.Metrics.AvgTtftMs)
|
||||
}
|
||||
|
||||
func TestRequestStatusNoSamplesAndDisabled(t *testing.T) {
|
||||
monitor := Monitor{Status: 1, Uptime: 1}
|
||||
applyRequestMetrics(&monitor, perfmetrics.GroupResult{})
|
||||
require.Equal(t, -1, monitor.Status)
|
||||
require.Empty(t, monitor.LastChecked)
|
||||
require.Empty(t, monitor.History)
|
||||
monitor.ChannelDisabled = true
|
||||
applyRequestMetrics(&monitor, perfmetrics.GroupResult{SuccessRate: 100, Series: []perfmetrics.BucketPoint{{Ts: time.Now().Unix(), SuccessRate: 100}}})
|
||||
require.Equal(t, 0, monitor.Status)
|
||||
}
|
||||
|
||||
func TestRequestStatusStaleSamplesStayUnknown(t *testing.T) {
|
||||
monitor := Monitor{}
|
||||
applyRequestMetrics(&monitor, perfmetrics.GroupResult{
|
||||
SuccessRate: 100,
|
||||
Series: []perfmetrics.BucketPoint{{
|
||||
Ts: time.Now().Add(-3 * time.Hour).Unix(),
|
||||
SuccessRate: 100,
|
||||
}},
|
||||
})
|
||||
require.Equal(t, -1, monitor.Status)
|
||||
require.NotEmpty(t, monitor.LastChecked)
|
||||
require.NotNil(t, monitor.Metrics)
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
)
|
||||
|
||||
func paymentReturnPath(suffix string) string {
|
||||
base := strings.TrimRight(system_setting.ServerAddress, "/")
|
||||
return base + suffix
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestPaymentReturnPathUsesDefaultDashboardRoutes(t *testing.T) {
|
||||
previousAddress := system_setting.ServerAddress
|
||||
system_setting.ServerAddress = "https://dashboard.example.com/"
|
||||
t.Cleanup(func() { system_setting.ServerAddress = previousAddress })
|
||||
|
||||
assert.Equal(
|
||||
t,
|
||||
"https://dashboard.example.com/wallet?pay=success",
|
||||
paymentReturnPath("/wallet?pay=success"),
|
||||
)
|
||||
assert.Equal(
|
||||
t,
|
||||
"https://dashboard.example.com/usage-logs",
|
||||
paymentReturnPath("/usage-logs"),
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// etagVersionPublicContent namespaces the public-content ETag; bump it when
|
||||
// the JSON envelope served by serveRevalidatedJSON changes shape.
|
||||
const etagVersionPublicContent = "public-content:v1"
|
||||
|
||||
type publicContentResponse struct {
|
||||
Success bool `json:"success"`
|
||||
Message string `json:"message"`
|
||||
Data string `json:"data"`
|
||||
}
|
||||
|
||||
// serveRevalidatedJSON writes public content as JSON with a weak
|
||||
// content-derived ETag and answers conditional requests with 304 Not
|
||||
// Modified. The ETag is a weak validator derived from the content, so it is
|
||||
// stable across replicas and JSON encodings, and a new one is issued when
|
||||
// the content changes. Cache-Control: no-cache forces revalidation before
|
||||
// reuse, so an admin edit takes effect on the next request; Vary:
|
||||
// Accept-Encoding keeps the gzip and identity encodings apart in shared
|
||||
// caches.
|
||||
func serveRevalidatedJSON(c *gin.Context, content string) {
|
||||
body, err := common.Marshal(publicContentResponse{
|
||||
Success: true,
|
||||
Message: "",
|
||||
Data: content,
|
||||
})
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
etag := common.ETagFor(etagVersionPublicContent, content)
|
||||
|
||||
c.Header("ETag", etag)
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Vary", "Accept-Encoding")
|
||||
|
||||
if common.ETagMatches(c.GetHeader("If-None-Match"), etag) {
|
||||
c.Status(http.StatusNotModified)
|
||||
return
|
||||
}
|
||||
|
||||
c.Data(http.StatusOK, "application/json; charset=utf-8", body)
|
||||
}
|
||||
@@ -0,0 +1,136 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/oauth"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/go-webauthn/webauthn/protocol"
|
||||
)
|
||||
|
||||
func GetVerificationMethods(c *gin.Context) {
|
||||
identity, ok := middleware.GetSessionAuthIdentity(c)
|
||||
if !ok {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"success": false, "message": "当前认证方式不支持安全验证"})
|
||||
return
|
||||
}
|
||||
requirements, err := service.GetVerificationRequirements(identity, c.Query("scope"))
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, requirements)
|
||||
}
|
||||
|
||||
// writeSecurityOperationError only exposes known, fixed business messages.
|
||||
// Unexpected errors retain their cause for the existing server-side auth logger.
|
||||
func writeSecurityOperationError(c *gin.Context, err error) {
|
||||
status := http.StatusOK
|
||||
var code, message string
|
||||
var protocolError *protocol.Error
|
||||
switch {
|
||||
case errors.Is(err, service.ErrAccountEmailInvalid), errors.Is(err, service.ErrAccountEmailRestricted):
|
||||
code, message = "EMAIL_ADDRESS_REJECTED", err.Error()
|
||||
case errors.Is(err, model.ErrEmailAlreadyTaken):
|
||||
code, message = "EMAIL_ALREADY_TAKEN", "This email address is already in use."
|
||||
case errors.Is(err, service.ErrEmailBindingDelivery):
|
||||
code, message = "EMAIL_BINDING_DELIVERY_FAILED", err.Error()
|
||||
case errors.Is(err, model.ErrEmailBindingCodeInvalid):
|
||||
code, message = "EMAIL_BINDING_CODE_INVALID", err.Error()
|
||||
case errors.Is(err, model.ErrEmailBindingLocked):
|
||||
code, message = "EMAIL_BINDING_LOCKED", err.Error()
|
||||
case errors.Is(err, model.ErrEmailBindingResendWait):
|
||||
status = http.StatusTooManyRequests
|
||||
code, message = "EMAIL_BINDING_RESEND_WAIT", err.Error()
|
||||
case errors.Is(err, common.ErrAccountPasswordLength), errors.Is(err, common.ErrAccountPasswordSame), errors.Is(err, common.ErrPasswordLegacyLimit):
|
||||
code, message = "PASSWORD_POLICY_REJECTED", err.Error()
|
||||
case errors.Is(err, model.ErrCurrentPasswordInvalid):
|
||||
code, message = "CURRENT_PASSWORD_INVALID", err.Error()
|
||||
case errors.Is(err, model.ErrAccountPasswordState), errors.Is(err, model.ErrAccountBindingChanged):
|
||||
status = http.StatusConflict
|
||||
code, message = "ACCOUNT_SECURITY_STATE_CHANGED", err.Error()
|
||||
case errors.Is(err, model.ErrLastLoginMethod):
|
||||
code, message = "LAST_LOGIN_METHOD", err.Error()
|
||||
case errors.Is(err, oauth.ErrTelegramOAuthNotConfigured):
|
||||
code, message = "TELEGRAM_OAUTH_NOT_CONFIGURED", oauth.ErrTelegramOAuthNotConfigured.Error()
|
||||
case errors.Is(err, oauth.ErrTelegramOAuthConflict):
|
||||
code, message = "TELEGRAM_OAUTH_CONFLICT", oauth.ErrTelegramOAuthConflict.Error()
|
||||
case errors.Is(err, oauth.ErrTelegramOAuthFailed):
|
||||
code, message = "TELEGRAM_OAUTH_FAILED", oauth.ErrTelegramOAuthFailed.Error()
|
||||
case errors.Is(err, oauth.ErrTelegramAccountNotBound):
|
||||
code, message = "TELEGRAM_ACCOUNT_NOT_BOUND", oauth.ErrTelegramAccountNotBound.Error()
|
||||
case errors.Is(err, model.ErrExternalIdentityAlreadyClaimed):
|
||||
code, message = "ACCOUNT_ALREADY_BOUND", "This external account is already bound."
|
||||
if c.Param("provider") == "telegram" {
|
||||
code, message = "TELEGRAM_BIND_ALREADY_BOUND", "This Telegram account is already bound."
|
||||
}
|
||||
case errors.Is(err, service.ErrVerificationContextInvalid):
|
||||
status = http.StatusBadRequest
|
||||
code, message = "SECURITY_CONTEXT_INVALID", service.ErrVerificationContextInvalid.Error()
|
||||
case errors.Is(err, service.ErrVerificationForbidden):
|
||||
status = http.StatusForbidden
|
||||
code, message = "SECURITY_ACTION_FORBIDDEN", service.ErrVerificationForbidden.Error()
|
||||
case errors.Is(err, service.ErrVerificationFailed), errors.As(err, &protocolError):
|
||||
code, message = "SECURITY_VERIFICATION_FAILED", service.ErrVerificationFailed.Error()
|
||||
case errors.Is(err, service.ErrVerificationLocked):
|
||||
code, message = "SECURITY_VERIFICATION_LOCKED", service.ErrVerificationLocked.Error()
|
||||
case errors.Is(err, service.ErrVerificationUnavailable):
|
||||
code, message = "SECURITY_METHOD_UNAVAILABLE", service.ErrVerificationUnavailable.Error()
|
||||
case errors.Is(err, service.ErrVerificationFlowRequired):
|
||||
status = http.StatusBadRequest
|
||||
code, message = "SECURITY_VERIFICATION_FLOW_REQUIRED", service.ErrVerificationFlowRequired.Error()
|
||||
case errors.Is(err, service.ErrProofMethod):
|
||||
code, message = "SECURITY_PROOF_METHOD_MISMATCH", "This verification method is not allowed for this action."
|
||||
case errors.Is(err, service.ErrProofScope):
|
||||
code, message = "SECURITY_PROOF_SCOPE_MISMATCH", "Verification does not match this action."
|
||||
case errors.Is(err, service.ErrOAuthAccountMismatch):
|
||||
code, message = "OAUTH_ACCOUNT_MISMATCH", service.ErrOAuthAccountMismatch.Error()
|
||||
case errors.Is(err, model.ErrTwoFASetupInvalid):
|
||||
status = http.StatusConflict
|
||||
code, message = "TWOFA_SETUP_INVALID", model.ErrTwoFASetupInvalid.Error()
|
||||
case errors.Is(err, model.ErrTwoFACodeInvalid):
|
||||
code, message = "TWOFA_CODE_INVALID", model.ErrTwoFACodeInvalid.Error()
|
||||
case errors.Is(err, model.ErrTwoFAAlreadyEnabled):
|
||||
code, message = "TWOFA_ALREADY_ENABLED", "Two-factor authentication is already enabled."
|
||||
case errors.Is(err, model.ErrTwoFANotEnabled):
|
||||
code, message = "TWOFA_NOT_ENABLED", "Two-factor authentication is not enabled."
|
||||
case errors.Is(err, model.ErrPasskeyNotFound):
|
||||
code, message = "PASSKEY_NOT_FOUND", "No Passkey is registered."
|
||||
case errors.Is(err, model.ErrAuthFlowInvalid), errors.Is(err, model.ErrAuthFlowExpired), errors.Is(err, model.ErrAuthFlowConsumed):
|
||||
code, message = "AUTH_FLOW_INVALID", "Verification flow expired"
|
||||
case errors.Is(err, model.ErrUserSessionInvalid), errors.Is(err, model.ErrUserSessionInactive):
|
||||
writeAuthSessionError(c, service.ErrAuthTokenInvalid)
|
||||
return
|
||||
default:
|
||||
c.Set("security_error_code", "AUTH_INTERNAL_ERROR")
|
||||
writeAuthSessionError(c, err)
|
||||
return
|
||||
}
|
||||
c.Set("security_error_code", code)
|
||||
c.JSON(status, gin.H{"success": false, "code": code, "message": message})
|
||||
}
|
||||
|
||||
func UniversalVerify(c *gin.Context) {
|
||||
identity, ok := middleware.GetSessionAuthIdentity(c)
|
||||
if !ok {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"success": false, "message": "当前认证方式不支持安全验证"})
|
||||
return
|
||||
}
|
||||
var request service.VerificationInput
|
||||
if err := common.DecodeJson(c.Request.Body, &request); err != nil {
|
||||
common.ApiErrorMsg(c, "参数错误")
|
||||
return
|
||||
}
|
||||
proof, err := service.VerifySecurityInput(identity, request)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
recordUserSecurityAudit(c, identity.UserID, "user.security_verify", map[string]any{"method": proof.Method, "scope": proof.Scope})
|
||||
common.ApiSuccess(c, proof)
|
||||
}
|
||||
@@ -0,0 +1,847 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"crypto/aes"
|
||||
"crypto/cipher"
|
||||
"crypto/rand"
|
||||
"crypto/rsa"
|
||||
"crypto/sha256"
|
||||
"crypto/x509"
|
||||
"encoding/base64"
|
||||
"encoding/pem"
|
||||
"errors"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/textproto"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/oauth"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/pquerna/otp/totp"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func TestSecurityAccountDeletionRequiresScopedProof(t *testing.T) {
|
||||
for _, scenario := range []string{"missing", "wrong scope", "expired", "consumed", "other session", "other account", "password disabled", "factor added"} {
|
||||
t.Run(scenario, func(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
originalIdentity := identity
|
||||
operation := service.VerificationOperation{Scope: service.VerificationScopeAccountDelete}
|
||||
proof := ""
|
||||
if scenario != "missing" {
|
||||
proof = issueSecurityEnrollmentProof(t, identity, operation, service.VerificationMethodPassword)
|
||||
}
|
||||
if scenario == "wrong scope" {
|
||||
proof = issueSecurityEnrollmentProof(t, identity, service.VerificationOperation{Scope: service.VerificationScopePasswordChange}, service.VerificationMethodPassword)
|
||||
}
|
||||
switch scenario {
|
||||
case "expired":
|
||||
require.NoError(t, model.DB.Model(&model.AuthFlow{}).Where("purpose = ?", model.AuthFlowPurposeSecurityProof).Update("expires_at", time.Now().Add(-time.Minute)).Error)
|
||||
case "consumed":
|
||||
_, err := service.ConsumeOperationProof(proof, identity, operation)
|
||||
require.NoError(t, err)
|
||||
case "other session", "other account":
|
||||
userID := user.Id
|
||||
if scenario == "other account" {
|
||||
other := &model.User{Username: "other", AffCode: "other", Group: "default", Password: user.Password, Status: common.UserStatusEnabled, Role: common.RoleCommonUser, AuthVersion: 1}
|
||||
require.NoError(t, model.DB.Create(other).Error)
|
||||
userID = other.Id
|
||||
}
|
||||
bundle, err := service.CreateLoginSession(userID, "password", "127.0.0.1", scenario)
|
||||
require.NoError(t, err)
|
||||
identity, err = service.ParseAccessToken(bundle.AccessToken)
|
||||
require.NoError(t, err)
|
||||
case "password disabled":
|
||||
previous := common.PasswordLoginEnabled
|
||||
common.PasswordLoginEnabled = false
|
||||
t.Cleanup(func() { common.PasswordLoginEnabled = previous })
|
||||
case "factor added":
|
||||
require.NoError(t, model.DB.Create(&model.TwoFA{UserId: user.Id, Secret: "JBSWY3DPEHPK3PXP", IsEnabled: true}).Error)
|
||||
}
|
||||
response := securityEnrollmentRequest("DELETE", "/api/user/self", "", proof, identity, DeleteSelf)
|
||||
assert.Equal(t, http.StatusForbidden, response.Code, response.Body.String())
|
||||
stored, err := model.GetUserById(user.Id, false)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, identity.UserAuthVersion, stored.AuthVersion)
|
||||
_, _, err = service.ValidateLoginSession(originalIdentity)
|
||||
assert.NoError(t, err)
|
||||
var audit model.AuditLog
|
||||
require.NoError(t, model.LOG_DB.Where("action = ?", "user.account_delete").Last(&audit).Error)
|
||||
assert.False(t, audit.Success)
|
||||
auditJSON, err := common.Marshal(audit)
|
||||
require.NoError(t, err)
|
||||
if proof != "" {
|
||||
assert.NotContains(t, string(auditJSON), proof)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityAccountDeletionAcceptsEitherFactorAndRevokesSessions(t *testing.T) {
|
||||
for _, method := range []string{"password", "oauth", "2fa", "passkey"} {
|
||||
t.Run(method, func(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
if method == "oauth" {
|
||||
require.NoError(t, model.DB.Model(user).Updates(map[string]any{"password": "", "github_id": "linked-user"}).Error)
|
||||
oauth.Register("account-delete-oauth", &enrollmentOAuthProvider{externalID: "linked-user"})
|
||||
t.Cleanup(func() { oauth.Unregister("account-delete-oauth") })
|
||||
}
|
||||
if method == "2fa" || method == "passkey" {
|
||||
newSecurityLoginPasskey(t, user.Id)
|
||||
factor := &model.TwoFA{UserId: user.Id, Secret: "JBSWY3DPEHPK3PXP", IsEnabled: true}
|
||||
if method == "passkey" {
|
||||
locked := time.Now().Add(time.Minute)
|
||||
factor.LockedUntil = &locked
|
||||
} else {
|
||||
system_setting.GetPasskeySettings().Enabled = false
|
||||
}
|
||||
require.NoError(t, model.DB.Create(factor).Error)
|
||||
}
|
||||
requirements, err := service.GetVerificationRequirements(identity, service.VerificationScopeAccountDelete)
|
||||
require.NoError(t, err)
|
||||
_, err = service.RequireVerificationMethod(identity, service.VerificationScopeAccountDelete, method)
|
||||
require.NoError(t, err)
|
||||
if method == "2fa" || method == "passkey" {
|
||||
assert.Len(t, requirements.Methods, 2)
|
||||
_, err = service.RequireVerificationMethod(identity, service.VerificationScopeAccountDelete, "password")
|
||||
assert.ErrorIs(t, err, service.ErrProofMethod)
|
||||
}
|
||||
proof := issueSecurityEnrollmentProof(t, identity, service.VerificationOperation{Scope: service.VerificationScopeAccountDelete}, method)
|
||||
if method == "password" || method == "2fa" {
|
||||
input := service.VerificationInput{Scope: service.VerificationScopeAccountDelete, Method: method, Password: "enrollment-password"}
|
||||
if method == "2fa" {
|
||||
input.Code, err = totp.GenerateCode("JBSWY3DPEHPK3PXP", time.Now())
|
||||
require.NoError(t, err)
|
||||
}
|
||||
verified, err := service.VerifySecurityInput(identity, input)
|
||||
require.NoError(t, err)
|
||||
proof = verified.ProofToken
|
||||
}
|
||||
otherSession, err := service.CreateLoginSession(user.Id, "password", "127.0.0.1", "second-session")
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, model.UpdateUserAccessToken(user.Id, "account-delete-access-token"))
|
||||
response := securityEnrollmentRequest("DELETE", "/api/user/self", "", proof, identity, DeleteSelf)
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.True(t, result.Success, response.Body.String())
|
||||
assert.Contains(t, response.Header().Get("Set-Cookie"), "Max-Age=0")
|
||||
assert.Contains(t, response.Header().Get("Cache-Control"), "no-store")
|
||||
var deleted model.User
|
||||
require.NoError(t, model.DB.Unscoped().First(&deleted, user.Id).Error)
|
||||
assert.True(t, deleted.DeletedAt.Valid)
|
||||
assert.Equal(t, user.AuthVersion+1, deleted.AuthVersion)
|
||||
_, _, err = service.ValidateLoginSession(identity)
|
||||
assert.Error(t, err)
|
||||
otherIdentity, err := service.ParseAccessToken(otherSession.AccessToken)
|
||||
require.NoError(t, err)
|
||||
_, _, err = service.ValidateLoginSession(otherIdentity)
|
||||
assert.Error(t, err)
|
||||
_, _, err = service.RefreshLoginSession(otherSession.RefreshToken, otherIdentity.SessionID, "127.0.0.1", "second-session")
|
||||
assert.Error(t, err)
|
||||
count, err := model.CountActiveUserSessions(user.Id, time.Now().Unix())
|
||||
require.NoError(t, err)
|
||||
assert.Zero(t, count)
|
||||
tokenUser, err := model.ValidateAccessToken("account-delete-access-token")
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, tokenUser)
|
||||
var audit model.AuditLog
|
||||
require.NoError(t, model.LOG_DB.Where("action = ?", "user.account_delete").Last(&audit).Error)
|
||||
assert.True(t, audit.Success)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityAccountDeletionRechecksTransactionAndConsumesFailedProof(t *testing.T) {
|
||||
for _, scenario := range []string{"revoked session", "auth version", "root", "write failure"} {
|
||||
t.Run(scenario, func(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
operation := service.VerificationOperation{Scope: service.VerificationScopeAccountDelete}
|
||||
proof := issueSecurityEnrollmentProof(t, identity, operation, "password")
|
||||
if scenario != "write failure" {
|
||||
_, err := service.ConsumeOperationProof(proof, identity, operation)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
switch scenario {
|
||||
case "revoked session":
|
||||
_, err := model.RevokeAllUserSessions(user.Id, "test")
|
||||
require.NoError(t, err)
|
||||
case "auth version":
|
||||
require.NoError(t, model.DB.Model(user).Update("auth_version", user.AuthVersion+1).Error)
|
||||
case "root":
|
||||
require.NoError(t, model.DB.Model(user).Update("role", common.RoleRootUser).Error)
|
||||
_, err := service.GetVerificationRequirements(identity, service.VerificationScopeAccountDelete)
|
||||
assert.ErrorIs(t, err, service.ErrVerificationForbidden)
|
||||
case "write failure":
|
||||
require.NoError(t, model.DB.Callback().Delete().Before("gorm:delete").Register("account-delete-failure", func(tx *gorm.DB) {
|
||||
if tx.Statement.Table == "users" {
|
||||
_ = tx.AddError(errors.New("injected deletion failure"))
|
||||
}
|
||||
}))
|
||||
response := securityEnrollmentRequest("DELETE", "/api/user/self", "", proof, identity, DeleteSelf)
|
||||
assert.Contains(t, response.Body.String(), `"success":false`)
|
||||
require.NoError(t, model.DB.Callback().Delete().Remove("account-delete-failure"))
|
||||
response = securityEnrollmentRequest("DELETE", "/api/user/self", "", proof, identity, DeleteSelf)
|
||||
assert.Contains(t, response.Body.String(), "SECURITY_PROOF_CONSUMED")
|
||||
}
|
||||
if scenario != "write failure" {
|
||||
assert.Error(t, model.DeleteUserForSession(identity))
|
||||
}
|
||||
_, err := model.GetUserById(user.Id, false)
|
||||
require.NoError(t, err)
|
||||
if scenario == "write failure" {
|
||||
_, _, err = service.ValidateLoginSession(identity)
|
||||
assert.NoError(t, err, "failed deletion must roll back the account version and preserve sessions")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityAccountDeletionConcurrentRequestsHaveOneWinner(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
proof := issueSecurityEnrollmentProof(t, identity, service.VerificationOperation{Scope: service.VerificationScopeAccountDelete}, "password")
|
||||
start := make(chan struct{})
|
||||
responses := make(chan string, 2)
|
||||
for range 2 {
|
||||
go func() {
|
||||
<-start
|
||||
response := securityEnrollmentRequest("DELETE", "/api/user/self", "", proof, identity, DeleteSelf)
|
||||
responses <- response.Body.String()
|
||||
}()
|
||||
}
|
||||
close(start)
|
||||
succeeded := 0
|
||||
for range 2 {
|
||||
var response securityEnrollmentResponse
|
||||
require.NoError(t, common.UnmarshalJsonStr(<-responses, &response))
|
||||
if response.Success {
|
||||
succeeded++
|
||||
}
|
||||
}
|
||||
assert.Equal(t, 1, succeeded)
|
||||
var deleted model.User
|
||||
require.NoError(t, model.DB.Unscoped().First(&deleted, user.Id).Error)
|
||||
assert.True(t, deleted.DeletedAt.Valid)
|
||||
assert.Equal(t, identity.UserAuthVersion+1, deleted.AuthVersion)
|
||||
}
|
||||
|
||||
type securityMailbox struct {
|
||||
mutex sync.Mutex
|
||||
mail map[string][]string
|
||||
}
|
||||
|
||||
// Use a real local SMTP boundary so delivery, failure and code handling are
|
||||
// exercised without exporting test-only mail hooks from production packages.
|
||||
func newSecurityMailbox(t *testing.T) *securityMailbox {
|
||||
t.Helper()
|
||||
listener, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
mailbox := &securityMailbox{mail: map[string][]string{}}
|
||||
previousServer, previousPort := common.SMTPServer, common.SMTPPort
|
||||
previousAccount, previousFrom, previousToken := common.SMTPAccount, common.SMTPFrom, common.SMTPToken
|
||||
previousSSL, previousTLS := common.SMTPSSLEnabled, common.SMTPStartTLSEnabled
|
||||
common.SMTPServer, common.SMTPPort = "127.0.0.1", listener.Addr().(*net.TCPAddr).Port
|
||||
common.SMTPAccount, common.SMTPFrom, common.SMTPToken = "", "sender@example.com", ""
|
||||
common.SMTPSSLEnabled, common.SMTPStartTLSEnabled = false, false
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
for {
|
||||
connection, err := listener.Accept()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
mailbox.receive(connection)
|
||||
}
|
||||
}()
|
||||
t.Cleanup(func() {
|
||||
_ = listener.Close()
|
||||
<-done
|
||||
common.SMTPServer, common.SMTPPort = previousServer, previousPort
|
||||
common.SMTPAccount, common.SMTPFrom, common.SMTPToken = previousAccount, previousFrom, previousToken
|
||||
common.SMTPSSLEnabled, common.SMTPStartTLSEnabled = previousSSL, previousTLS
|
||||
})
|
||||
return mailbox
|
||||
}
|
||||
|
||||
func (mailbox *securityMailbox) receive(connection net.Conn) {
|
||||
defer connection.Close()
|
||||
_ = connection.SetDeadline(time.Now().Add(10 * time.Second))
|
||||
client := textproto.NewConn(connection)
|
||||
if client.PrintfLine("220 localhost ESMTP") != nil {
|
||||
return
|
||||
}
|
||||
receiver := ""
|
||||
for {
|
||||
line, err := client.ReadLine()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
switch {
|
||||
case strings.HasPrefix(line, "RCPT TO:"):
|
||||
receiver = strings.TrimSuffix(strings.TrimPrefix(line, "RCPT TO:<"), ">")
|
||||
case line == "DATA":
|
||||
if client.PrintfLine("354 Send message") != nil {
|
||||
return
|
||||
}
|
||||
message, err := io.ReadAll(client.DotReader())
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
mailbox.mutex.Lock()
|
||||
mailbox.mail[receiver] = append(mailbox.mail[receiver], string(message))
|
||||
mailbox.mutex.Unlock()
|
||||
case line == "QUIT":
|
||||
_ = client.PrintfLine("221 Goodbye")
|
||||
return
|
||||
}
|
||||
if client.PrintfLine("250 OK") != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (mailbox *securityMailbox) code(t *testing.T, receiver string) string {
|
||||
t.Helper()
|
||||
mailbox.mutex.Lock()
|
||||
defer mailbox.mutex.Unlock()
|
||||
for index := len(mailbox.mail[receiver]) - 1; index >= 0; index-- {
|
||||
match := regexp.MustCompile(`<strong>([0-9]{6})</strong>`).FindStringSubmatch(mailbox.mail[receiver][index])
|
||||
if len(match) == 2 {
|
||||
return match[1]
|
||||
}
|
||||
}
|
||||
t.Fatalf("no verification email delivered to %s", receiver)
|
||||
return ""
|
||||
}
|
||||
|
||||
func startSecurityEmailBinding(t *testing.T, identity service.AuthIdentity, email, method string) service.EmailBindingData {
|
||||
t.Helper()
|
||||
context, err := common.Marshal(service.AccountBindingContext{Provider: "email", Email: email})
|
||||
require.NoError(t, err)
|
||||
proof := issueSecurityEnrollmentProof(t, identity, service.VerificationOperation{Scope: service.VerificationScopeAccountBind, Context: context}, method)
|
||||
request, err := common.Marshal(map[string]string{"email": email})
|
||||
require.NoError(t, err)
|
||||
response := securityEnrollmentRequest("POST", "/api/oauth/email/bind/start", string(request), proof, identity, EmailBindStart)
|
||||
var body securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &body))
|
||||
require.True(t, body.Success, response.Body.String())
|
||||
var flow service.EmailBindingData
|
||||
require.NoError(t, common.Unmarshal(body.Data, &flow))
|
||||
require.NotEmpty(t, flow.FlowToken)
|
||||
return flow
|
||||
}
|
||||
|
||||
func TestSecurityAccountRequiresProofBeforeMutation(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
name, path, body string
|
||||
handler gin.HandlerFunc
|
||||
}{
|
||||
{"password", "/api/user/self", `{"password":"account-password!42","original_password":"enrollment-password"}`, UpdateSelf},
|
||||
{"oauth binding", "/api/oauth/state", `{"provider":"security-account-test","intent":"bind"}`, GenerateOAuthCode},
|
||||
} {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
oauth.Register("security-account-test", &authFlowTestOAuthProvider{})
|
||||
t.Cleanup(func() { oauth.Unregister("security-account-test") })
|
||||
response := securityEnrollmentRequest(http.MethodPost, test.path, test.body, "", identity, test.handler)
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
assert.Equal(t, http.StatusForbidden, response.Code)
|
||||
assert.Equal(t, "SECURITY_PROOF_REQUIRED", result.Code)
|
||||
stored, err := model.GetUserById(user.Id, true)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, common.ValidatePasswordAndHash("enrollment-password", stored.Password))
|
||||
var flows int64
|
||||
require.NoError(t, model.DB.Model(&model.AuthFlow{}).Count(&flows).Error)
|
||||
assert.Zero(t, flows)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityAccountPasswordRequiresCurrentPassword(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
name, original string
|
||||
success bool
|
||||
}{
|
||||
{"missing", "", false},
|
||||
{"wrong", "wrong-password", false},
|
||||
{"correct", "enrollment-password", true},
|
||||
} {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
proof := issueSecurityEnrollmentProof(t, identity, service.VerificationOperation{Scope: service.VerificationScopePasswordChange}, service.VerificationMethodPassword)
|
||||
request, err := common.Marshal(map[string]string{"password": "password123", "original_password": test.original})
|
||||
require.NoError(t, err)
|
||||
response := securityEnrollmentRequest(http.MethodPut, "/api/user/self", string(request), proof, identity, UpdateSelf)
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
assert.Equal(t, test.success, result.Success, response.Body.String())
|
||||
stored, err := model.GetUserById(user.Id, true)
|
||||
require.NoError(t, err)
|
||||
if test.success {
|
||||
assert.True(t, common.ValidatePasswordAndHash("password123", stored.Password))
|
||||
assert.True(t, strings.HasPrefix(stored.Password, "$argon2id$"))
|
||||
} else {
|
||||
assert.Equal(t, "CURRENT_PASSWORD_INVALID", result.Code)
|
||||
assert.Equal(t, user.Password, stored.Password)
|
||||
repeated := securityEnrollmentRequest(http.MethodPut, "/api/user/self", string(request), proof, identity, UpdateSelf)
|
||||
require.NoError(t, common.Unmarshal(repeated.Body.Bytes(), &result))
|
||||
assert.Equal(t, "SECURITY_PROOF_CONSUMED", result.Code)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityAccountProfileReadsPasswordStatusInOneQuery(t *testing.T) {
|
||||
for _, hasPassword := range []bool{true, false} {
|
||||
name := "password set"
|
||||
if !hasPassword {
|
||||
name = "password unset"
|
||||
}
|
||||
t.Run(name, func(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
updates := map[string]any{"access_token": "private-profile-token", "remark": "private-profile-remark"}
|
||||
if !hasPassword {
|
||||
updates["password"] = ""
|
||||
}
|
||||
require.NoError(t, model.DB.Model(&model.User{}).Where("id = ?", user.Id).Updates(updates).Error)
|
||||
queries := 0
|
||||
require.NoError(t, model.DB.Callback().Query().Before("gorm:query").Register("profile_query_count", func(tx *gorm.DB) {
|
||||
queries++
|
||||
}))
|
||||
profile, err := model.GetSelfUserById(user.Id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 1, queries)
|
||||
assert.Equal(t, hasPassword, profile.HasPassword)
|
||||
assert.Empty(t, profile.Password)
|
||||
assert.Nil(t, profile.AccessToken)
|
||||
assert.Empty(t, profile.Remark)
|
||||
assert.Equal(t, user.AuthVersion, profile.AuthVersion)
|
||||
assert.Equal(t, user.Group, profile.Group)
|
||||
require.NoError(t, model.DB.Callback().Query().Remove("profile_query_count"))
|
||||
|
||||
response := securityEnrollmentRequest(http.MethodGet, "/api/user/self", "", "", identity, GetSelf)
|
||||
var result struct {
|
||||
Success bool `json:"success"`
|
||||
Data map[string]any `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.True(t, result.Success, response.Body.String())
|
||||
assert.Equal(t, hasPassword, result.Data["has_password"])
|
||||
assert.NotContains(t, result.Data, "password")
|
||||
assert.NotContains(t, result.Data, "access_token")
|
||||
assert.NotContains(t, result.Data, "remark")
|
||||
|
||||
bundle, err := service.CreateLoginSession(user.Id, "profile-test", "127.0.0.1", "profile-test")
|
||||
require.NoError(t, err)
|
||||
_, refreshed, err := service.RefreshLoginSession(bundle.RefreshToken, "", "127.0.0.1", "profile-test")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, hasPassword, refreshed.HasPassword)
|
||||
assert.Empty(t, refreshed.Password)
|
||||
assert.Nil(t, refreshed.AccessToken)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityAccountProfileUpdateDoesNotRequireProof(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
response := securityEnrollmentRequest(http.MethodPut, "/api/user/self", `{"display_name":"Updated"}`, "", identity, UpdateSelf)
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.True(t, result.Success, response.Body.String())
|
||||
stored, err := model.GetUserById(user.Id, true)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "Updated", stored.DisplayName)
|
||||
assert.Equal(t, user.Password, stored.Password)
|
||||
}
|
||||
|
||||
func TestSecurityAccountLongUnicodePasswordAndSessionRotation(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
other, err := service.CreateLoginSession(user.Id, "password", "127.0.0.1", "other-device")
|
||||
require.NoError(t, err)
|
||||
otherIdentity, err := service.ParseAccessToken(other.AccessToken)
|
||||
require.NoError(t, err)
|
||||
password := strings.Repeat("安全🔒 ", 24)
|
||||
proof := issueSecurityEnrollmentProof(t, identity, service.VerificationOperation{Scope: service.VerificationScopePasswordChange}, service.VerificationMethodPassword)
|
||||
body, err := common.Marshal(map[string]string{"password": password, "original_password": "enrollment-password"})
|
||||
require.NoError(t, err)
|
||||
response := securityEnrollmentRequest(http.MethodPut, "/api/user/self", string(body), proof, identity, UpdateSelf)
|
||||
var result struct {
|
||||
Success bool `json:"success"`
|
||||
Data struct {
|
||||
AccessToken string `json:"access_token"`
|
||||
HasPassword bool `json:"has_password"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.True(t, result.Success, response.Body.String())
|
||||
assert.True(t, result.Data.HasPassword)
|
||||
currentIdentity, err := service.ParseAccessToken(result.Data.AccessToken)
|
||||
require.NoError(t, err)
|
||||
_, _, err = service.ValidateLoginSession(currentIdentity)
|
||||
require.NoError(t, err)
|
||||
_, _, err = service.ValidateLoginSession(otherIdentity)
|
||||
require.Error(t, err)
|
||||
login := model.User{Username: user.Username, Password: password}
|
||||
require.NoError(t, login.ValidateAndFill())
|
||||
assert.False(t, common.ValidatePasswordAndHash(strings.TrimSpace(password), login.Password))
|
||||
assert.False(t, common.ValidatePasswordAndHash("enrollment-password", login.Password))
|
||||
}
|
||||
|
||||
func TestSecurityAccountPasswordPolicyAndHashCompatibility(t *testing.T) {
|
||||
t.Setenv("ACCOUNT_PASSWORD_HASH_ALGORITHM", "argon2id")
|
||||
for _, test := range []struct {
|
||||
name, password string
|
||||
valid bool
|
||||
}{
|
||||
{"common password accepted", "password123", true},
|
||||
{"short", "seven77", false},
|
||||
{"too long", strings.Repeat("界", 129), false},
|
||||
{"long Unicode", strings.Repeat("界🔒 ", 40), true},
|
||||
{"whitespace preserved", " phrase with whitespace ", true},
|
||||
} {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
hash, err := common.HashAccountPassword(test.password)
|
||||
if !test.valid {
|
||||
require.Error(t, err)
|
||||
assert.Empty(t, hash)
|
||||
return
|
||||
}
|
||||
require.NoError(t, err)
|
||||
assert.True(t, common.ValidatePasswordAndHash(test.password, hash))
|
||||
assert.False(t, common.ValidatePasswordAndHash(test.password+"x", hash))
|
||||
})
|
||||
}
|
||||
legacy, err := common.Password2Hash("123456")
|
||||
require.NoError(t, err)
|
||||
assert.True(t, common.ValidatePasswordAndHash("123456", legacy), "login preserves historical passwords without applying new policy")
|
||||
for _, invalid := range []string{"$argon2id$", "$argon2id$v=19$m=4294967295,t=2,p=1$bad$bad", "$argon2id$v=19$m=19456,t=2,p=1$bad$bad"} {
|
||||
assert.False(t, common.ValidatePasswordAndHash("example-password", invalid))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityAccountEncryptedLongPasswordLogin(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
password := strings.Repeat("🔒", 128)
|
||||
hash, err := common.HashAccountPassword(password)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, model.DB.Model(user).Update("password", hash).Error)
|
||||
keyID, publicPEM := common.PasswordEncryptionPublicKey()
|
||||
if keyID == "" {
|
||||
privatePEM, err := common.GeneratePasswordEncryptionPrivateKey()
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, common.LoadPasswordEncryptionPrivateKey(privatePEM))
|
||||
keyID, publicPEM = common.PasswordEncryptionPublicKey()
|
||||
}
|
||||
block, _ := pem.Decode([]byte(publicPEM))
|
||||
require.NotNil(t, block)
|
||||
parsed, err := x509.ParsePKIXPublicKey(block.Bytes)
|
||||
require.NoError(t, err)
|
||||
publicKey, ok := parsed.(*rsa.PublicKey)
|
||||
require.True(t, ok)
|
||||
key, nonce := make([]byte, 32), make([]byte, 12)
|
||||
_, err = rand.Read(key)
|
||||
require.NoError(t, err)
|
||||
_, err = rand.Read(nonce)
|
||||
require.NoError(t, err)
|
||||
wrappedKey, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, publicKey, key, []byte("password-v2"))
|
||||
require.NoError(t, err)
|
||||
aesBlock, err := aes.NewCipher(key)
|
||||
require.NoError(t, err)
|
||||
gcm, err := cipher.NewGCM(aesBlock)
|
||||
require.NoError(t, err)
|
||||
ciphertext := gcm.Seal(nil, nonce, []byte(password), []byte("password-v2:"+keyID))
|
||||
parts := []string{"v2", base64.StdEncoding.EncodeToString(wrappedKey), base64.StdEncoding.EncodeToString(nonce), base64.StdEncoding.EncodeToString(ciphertext)}
|
||||
encrypted := strings.Join(parts, ".")
|
||||
common.PasswordLoginEncryptionEnabled = true
|
||||
request, err := common.Marshal(LoginRequest{Username: user.Username, PasswordEncrypted: encrypted, EncryptionKeyID: keyID})
|
||||
require.NoError(t, err)
|
||||
response := securityEnrollmentRequest("POST", "/api/user/login", string(request), "", identity, Login)
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.True(t, result.Success, response.Body.String())
|
||||
assert.Contains(t, string(result.Data), `"has_password":true`)
|
||||
proof, err := service.VerifySecurityInput(identity, service.VerificationInput{Method: "password", Scope: service.VerificationScopePasswordChange, PasswordEncrypted: encrypted, EncryptionKeyID: keyID})
|
||||
require.NoError(t, err)
|
||||
assert.NotEmpty(t, proof.ProofToken)
|
||||
ciphertext[len(ciphertext)-1] ^= 1
|
||||
parts[3] = base64.StdEncoding.EncodeToString(ciphertext)
|
||||
for _, invalid := range []string{strings.Join(parts, "."), "v2.bad.bad.bad", "v2." + strings.Repeat("a", 4096)} {
|
||||
_, err := common.DecryptPassword(invalid, keyID)
|
||||
assert.ErrorIs(t, err, common.ErrPasswordEncryptionInvalid)
|
||||
}
|
||||
_, err = common.DecryptPassword(encrypted, "wrong-key-id")
|
||||
assert.ErrorIs(t, err, common.ErrPasswordEncryptionInvalid)
|
||||
}
|
||||
|
||||
func TestSecurityAccountEmailConfirmationAndAudit(t *testing.T) {
|
||||
for _, scenario := range []string{"first email", "replace without mfa", "replace with mfa"} {
|
||||
t.Run(scenario, func(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
mailbox := newSecurityMailbox(t)
|
||||
method, previous := service.VerificationMethodPassword, ""
|
||||
if scenario != "first email" {
|
||||
previous = "previous@example.com"
|
||||
require.NoError(t, model.DB.Model(user).Update("email", previous).Error)
|
||||
}
|
||||
if scenario == "replace with mfa" {
|
||||
require.NoError(t, model.DB.Create(&model.TwoFA{UserId: user.Id, Secret: "JBSWY3DPEHPK3PXP", IsEnabled: true}).Error)
|
||||
method = service.VerificationMethodTwoFA
|
||||
}
|
||||
flow := startSecurityEmailBinding(t, identity, "New@Example.com", method)
|
||||
assert.Equal(t, "new@example.com", flow.Email)
|
||||
assert.Equal(t, scenario == "replace without mfa", flow.OldEmailRequired)
|
||||
newCode, oldCode := mailbox.code(t, flow.Email), ""
|
||||
_, state, err := model.GetEmailBinding(identity, flow.FlowToken)
|
||||
require.NoError(t, err)
|
||||
assert.NotEqual(t, newCode, state.NewCodeHash)
|
||||
assert.True(t, common.ValidatePasswordAndHash(newCode, state.NewCodeHash))
|
||||
if flow.OldEmailRequired {
|
||||
oldCode = mailbox.code(t, previous)
|
||||
_, err := service.FinishEmailBinding(identity, flow.FlowToken, newCode, "")
|
||||
assert.ErrorIs(t, err, model.ErrEmailBindingCodeInvalid)
|
||||
stored, err := model.GetUserById(user.Id, false)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, previous, stored.Email)
|
||||
}
|
||||
request, err := common.Marshal(map[string]string{"flow_token": flow.FlowToken, "new_code": newCode, "old_code": oldCode})
|
||||
require.NoError(t, err)
|
||||
response := securityEnrollmentRequest("POST", "/api/oauth/email/bind", string(request), "", identity, EmailBind)
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.True(t, result.Success, response.Body.String())
|
||||
stored, err := model.GetUserById(user.Id, false)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, flow.Email, stored.Email)
|
||||
_, err = service.FinishEmailBinding(identity, flow.FlowToken, newCode, oldCode)
|
||||
assert.ErrorIs(t, err, model.ErrAuthFlowConsumed)
|
||||
legacy := securityEnrollmentRequest("POST", "/api/oauth/email/bind", `{"email":"other@example.com","code":"123456"}`, "", identity, EmailBind)
|
||||
assert.Contains(t, legacy.Body.String(), `"success":false`)
|
||||
var audits []model.AuditLog
|
||||
require.NoError(t, model.LOG_DB.Find(&audits).Error)
|
||||
require.NotEmpty(t, audits)
|
||||
assert.False(t, audits[len(audits)-1].Success, "a rejected legacy request must be recorded as failed")
|
||||
encoded, err := common.Marshal(audits)
|
||||
require.NoError(t, err)
|
||||
assert.Contains(t, string(encoded), "user.binding_bind")
|
||||
assert.NotContains(t, string(encoded), flow.FlowToken)
|
||||
assert.NotContains(t, string(encoded), newCode)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityAccountEmailResendAndAttemptLimit(t *testing.T) {
|
||||
_, identity := setupSecurityEnrollmentTest(t)
|
||||
mailbox := newSecurityMailbox(t)
|
||||
flow := startSecurityEmailBinding(t, identity, "new@example.com", service.VerificationMethodPassword)
|
||||
_, err := service.ResendAccountEmailBinding(identity, flow.FlowToken)
|
||||
assert.ErrorIs(t, err, model.ErrEmailBindingResendWait)
|
||||
_, err = service.FinishEmailBinding(identity, flow.FlowToken, "invalid", "")
|
||||
assert.ErrorIs(t, err, model.ErrEmailBindingCodeInvalid)
|
||||
stored, state, err := model.GetEmailBinding(identity, flow.FlowToken)
|
||||
require.NoError(t, err)
|
||||
oldHash := state.NewCodeHash
|
||||
state.ResendAt = time.Now().Add(-time.Second).Unix()
|
||||
payload, err := common.Marshal(state)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, model.DB.Model(stored).Update("payload", string(payload)).Error)
|
||||
replacement, err := service.ResendAccountEmailBinding(identity, flow.FlowToken)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, flow.ExpiresAt, replacement.ExpiresAt)
|
||||
_, state, err = model.GetEmailBinding(identity, flow.FlowToken)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 1, state.FailedAttempts)
|
||||
assert.NotEqual(t, oldHash, state.NewCodeHash)
|
||||
assert.True(t, common.ValidatePasswordAndHash(mailbox.code(t, flow.Email), state.NewCodeHash))
|
||||
for attempt := 2; attempt <= model.EmailBindingMaxAttempts; attempt++ {
|
||||
_, err = service.FinishEmailBinding(identity, flow.FlowToken, "invalid", "")
|
||||
if attempt < model.EmailBindingMaxAttempts {
|
||||
assert.ErrorIs(t, err, model.ErrEmailBindingCodeInvalid)
|
||||
} else {
|
||||
assert.ErrorIs(t, err, model.ErrEmailBindingLocked)
|
||||
}
|
||||
}
|
||||
_, err = service.FinishEmailBinding(identity, flow.FlowToken, mailbox.code(t, flow.Email), "")
|
||||
assert.ErrorIs(t, err, model.ErrEmailBindingLocked)
|
||||
}
|
||||
|
||||
func TestSecurityAccountEmailConcurrentClaimsHaveOneOwner(t *testing.T) {
|
||||
user, firstIdentity := setupSecurityEnrollmentTest(t)
|
||||
mailbox := newSecurityMailbox(t)
|
||||
other := &model.User{Username: "other-user", AffCode: "other-account", Group: "default", Password: user.Password, Status: common.UserStatusEnabled, Role: common.RoleCommonUser, AuthVersion: 1}
|
||||
require.NoError(t, model.DB.Create(other).Error)
|
||||
require.NoError(t, model.PublishUserAuthCache(other.Id))
|
||||
bundle, err := service.CreateLoginSession(other.Id, "password", "127.0.0.1", "other-user")
|
||||
require.NoError(t, err)
|
||||
secondIdentity, err := service.ParseAccessToken(bundle.AccessToken)
|
||||
require.NoError(t, err)
|
||||
start, results := make(chan struct{}), make(chan error, 2)
|
||||
for _, identity := range []service.AuthIdentity{firstIdentity, secondIdentity} {
|
||||
flow := startSecurityEmailBinding(t, identity, "shared@example.com", service.VerificationMethodPassword)
|
||||
code := mailbox.code(t, flow.Email)
|
||||
go func(identity service.AuthIdentity, token, code string) {
|
||||
<-start
|
||||
_, err := service.FinishEmailBinding(identity, token, code, "")
|
||||
results <- err
|
||||
}(identity, flow.FlowToken, code)
|
||||
}
|
||||
close(start)
|
||||
first, second := <-results, <-results
|
||||
assert.NotEqual(t, first == nil, second == nil, "only one account may claim the address: %v / %v", first, second)
|
||||
var owners int64
|
||||
require.NoError(t, model.DB.Model(&model.User{}).Where("email = ?", "shared@example.com").Count(&owners).Error)
|
||||
assert.EqualValues(t, 1, owners)
|
||||
}
|
||||
|
||||
func TestSecurityAccountEmailRejectsChangedAuthorization(t *testing.T) {
|
||||
for _, scenario := range []string{"other session", "expired", "email changed", "mfa enabled", "notification failure"} {
|
||||
t.Run(scenario, func(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
mailbox := newSecurityMailbox(t)
|
||||
flow := startSecurityEmailBinding(t, identity, "new@example.com", service.VerificationMethodPassword)
|
||||
code := mailbox.code(t, flow.Email)
|
||||
stored, _, err := model.GetEmailBinding(identity, flow.FlowToken)
|
||||
require.NoError(t, err)
|
||||
switch scenario {
|
||||
case "other session":
|
||||
identity.SessionID = "different-session"
|
||||
case "expired":
|
||||
require.NoError(t, model.DB.Model(stored).Update("expires_at", time.Now().Add(-time.Second)).Error)
|
||||
case "email changed":
|
||||
require.NoError(t, model.DB.Model(user).Update("email", "changed@example.com").Error)
|
||||
case "mfa enabled":
|
||||
require.NoError(t, model.DB.Create(&model.TwoFA{UserId: user.Id, Secret: "JBSWY3DPEHPK3PXP", IsEnabled: true}).Error)
|
||||
case "notification failure":
|
||||
common.SMTPServer = ""
|
||||
}
|
||||
request, err := common.Marshal(map[string]string{"flow_token": flow.FlowToken, "new_code": code})
|
||||
require.NoError(t, err)
|
||||
response := securityEnrollmentRequest("POST", "/api/oauth/email/bind", string(request), "", identity, EmailBind)
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
assert.Equal(t, scenario == "notification failure", result.Success, response.Body.String())
|
||||
if scenario == "notification failure" {
|
||||
assert.Contains(t, string(result.Data), `"notification_warning":true`)
|
||||
}
|
||||
current, err := model.GetUserById(user.Id, false)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, scenario == "notification failure", current.Email == flow.Email)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityAccountOAuthBindingRejectsInvalidFlow(t *testing.T) {
|
||||
for _, scenario := range []string{"success and replay", "other session", "wrong provider", "expired", "old flow", "mfa enabled"} {
|
||||
t.Run(scenario, func(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
oauth.Register("account-oauth", &enrollmentOAuthProvider{externalID: "new-binding"})
|
||||
oauth.Register("other-oauth", &enrollmentOAuthProvider{externalID: "wrong-binding"})
|
||||
t.Cleanup(func() { oauth.Unregister("account-oauth"); oauth.Unregister("other-oauth") })
|
||||
operation := service.VerificationOperation{Scope: service.VerificationScopeAccountBind, Context: []byte(`{"provider":"account-oauth"}`)}
|
||||
proof := issueSecurityEnrollmentProof(t, identity, operation, service.VerificationMethodPassword)
|
||||
response := securityEnrollmentRequest("POST", "/api/oauth/state", `{"provider":"account-oauth","intent":"bind"}`, proof, identity, GenerateOAuthCode)
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.True(t, result.Success, response.Body.String())
|
||||
var started struct {
|
||||
FlowToken string `json:"flow_token"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(result.Data, &started))
|
||||
flow, err := model.GetAuthFlow(started.FlowToken, model.AuthFlowMatch{Purpose: model.AuthFlowPurposeOAuth})
|
||||
require.NoError(t, err)
|
||||
provider := "account-oauth"
|
||||
switch scenario {
|
||||
case "other session":
|
||||
identity.SessionID = "different-session"
|
||||
case "wrong provider":
|
||||
provider = "other-oauth"
|
||||
case "expired":
|
||||
require.NoError(t, model.DB.Model(flow).Update("expires_at", time.Now().Add(-time.Second)).Error)
|
||||
case "old flow":
|
||||
require.NoError(t, model.DB.Model(flow).Update("payload", "{}").Error)
|
||||
case "mfa enabled":
|
||||
require.NoError(t, model.DB.Create(&model.TwoFA{UserId: user.Id, Secret: "JBSWY3DPEHPK3PXP", IsEnabled: true}).Error)
|
||||
}
|
||||
handler := func(c *gin.Context) { c.Params = gin.Params{{Key: "provider", Value: provider}}; HandleOAuth(c) }
|
||||
path := "/api/oauth/" + provider + "?state=" + started.FlowToken + "&code=provider-code"
|
||||
response = securityEnrollmentRequest("GET", path, "", "", identity, handler)
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
assert.Equal(t, scenario == "success and replay", result.Success, response.Body.String())
|
||||
current, err := model.GetUserById(user.Id, false)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, scenario == "success and replay", current.GitHubId == "new-binding")
|
||||
if result.Success {
|
||||
replay := securityEnrollmentRequest("GET", path, "", "", identity, handler)
|
||||
assert.Contains(t, replay.Body.String(), `"success":false`)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityAccountFirstPasswordRace(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
require.NoError(t, model.DB.Model(user).Update("password", "").Error)
|
||||
require.NoError(t, model.DB.Create(&model.PasskeyCredential{UserID: user.Id, CredentialID: "existing-key", PublicKey: "public-key"}).Error)
|
||||
proof := issueSecurityEnrollmentProof(t, identity, service.VerificationOperation{Scope: service.VerificationScopePasswordSet}, service.VerificationMethodPasskey)
|
||||
request := `{"password":"initial-account-password!42"}`
|
||||
response := securityEnrollmentRequest("PUT", "/api/user/self", request, proof, identity, UpdateSelf)
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.True(t, result.Success, response.Body.String())
|
||||
user, identity = setupSecurityEnrollmentTest(t)
|
||||
require.NoError(t, model.DB.Model(user).Update("password", "").Error)
|
||||
start, results := make(chan struct{}), make(chan error, 2)
|
||||
for _, password := range []string{"first-password!42", "second-password!42"} {
|
||||
go func(password string) {
|
||||
<-start
|
||||
results <- model.ChangeUserPassword(identity, &model.User{Id: user.Id, Password: password}, true)
|
||||
}(password)
|
||||
}
|
||||
close(start)
|
||||
first, second := <-results, <-results
|
||||
assert.NotEqual(t, first == nil, second == nil, "only one concurrent first password may succeed: %v / %v", first, second)
|
||||
stored, err := model.GetUserById(user.Id, true)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, common.ValidatePasswordAndHash("first-password!42", stored.Password) || common.ValidatePasswordAndHash("second-password!42", stored.Password))
|
||||
}
|
||||
|
||||
func TestSecurityAccountUnbindPreservesUsableLoginMethod(t *testing.T) {
|
||||
for _, scenario := range []string{"password", "passkey", "email only", "twofa only", "disabled password", "disabled passkey"} {
|
||||
t.Run(scenario, func(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
previousPasswordEnabled := common.PasswordLoginEnabled
|
||||
t.Cleanup(func() { common.PasswordLoginEnabled = previousPasswordEnabled })
|
||||
common.PasswordLoginEnabled = scenario != "disabled password"
|
||||
require.NoError(t, model.DB.Create(&model.UserOAuthBinding{UserId: user.Id, ProviderId: 31, ProviderUserId: "linked-subject"}).Error)
|
||||
if scenario != "password" && scenario != "disabled password" {
|
||||
require.NoError(t, model.DB.Model(user).Updates(map[string]any{"password": "", "email": "verified@example.com"}).Error)
|
||||
}
|
||||
if scenario == "passkey" || scenario == "disabled passkey" {
|
||||
require.NoError(t, model.DB.Create(&model.PasskeyCredential{UserID: user.Id, CredentialID: "existing-key", PublicKey: "public-key"}).Error)
|
||||
system_setting.GetPasskeySettings().Enabled = scenario == "passkey"
|
||||
}
|
||||
if scenario == "twofa only" {
|
||||
require.NoError(t, model.DB.Create(&model.TwoFA{UserId: user.Id, Secret: "JBSWY3DPEHPK3PXP", IsEnabled: true}).Error)
|
||||
}
|
||||
err := service.UnbindAccountOAuth(identity, 31)
|
||||
if scenario == "password" || scenario == "passkey" {
|
||||
require.NoError(t, err)
|
||||
} else {
|
||||
assert.ErrorIs(t, err, model.ErrLastLoginMethod)
|
||||
var binding model.UserOAuthBinding
|
||||
require.NoError(t, model.DB.Where("user_id = ? AND provider_id = ?", user.Id, 31).First(&binding).Error)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,174 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type Setup struct {
|
||||
Status bool `json:"status"`
|
||||
RootInit bool `json:"root_init"`
|
||||
DatabaseType string `json:"database_type"`
|
||||
}
|
||||
|
||||
type SetupRequest struct {
|
||||
Username string `json:"username"`
|
||||
Password string `json:"password"`
|
||||
ConfirmPassword string `json:"confirmPassword"`
|
||||
SelfUseModeEnabled bool `json:"SelfUseModeEnabled"`
|
||||
DemoSiteEnabled bool `json:"DemoSiteEnabled"`
|
||||
}
|
||||
|
||||
func GetSetup(c *gin.Context) {
|
||||
setup := Setup{
|
||||
Status: constant.Setup,
|
||||
}
|
||||
if constant.Setup {
|
||||
c.JSON(200, gin.H{
|
||||
"success": true,
|
||||
"data": setup,
|
||||
})
|
||||
return
|
||||
}
|
||||
setup.RootInit = model.RootUserExists()
|
||||
setup.DatabaseType = string(common.MainDatabaseType())
|
||||
c.JSON(200, gin.H{
|
||||
"success": true,
|
||||
"data": setup,
|
||||
})
|
||||
}
|
||||
|
||||
func PostSetup(c *gin.Context) {
|
||||
// Check if setup is already completed
|
||||
if constant.Setup {
|
||||
c.JSON(200, gin.H{
|
||||
"success": false,
|
||||
"message": "系统已经初始化完成",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Check if root user already exists
|
||||
rootExists := model.RootUserExists()
|
||||
|
||||
var req SetupRequest
|
||||
err := c.ShouldBindJSON(&req)
|
||||
if err != nil {
|
||||
c.JSON(200, gin.H{
|
||||
"success": false,
|
||||
"message": "请求参数有误",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// If root doesn't exist, validate and create admin account
|
||||
if !rootExists {
|
||||
// Validate username length: max 12 characters to align with model.User validation
|
||||
if len(req.Username) > 12 {
|
||||
c.JSON(200, gin.H{
|
||||
"success": false,
|
||||
"message": "用户名长度不能超过12个字符",
|
||||
})
|
||||
return
|
||||
}
|
||||
// Validate password
|
||||
if req.Password != req.ConfirmPassword {
|
||||
c.JSON(200, gin.H{
|
||||
"success": false,
|
||||
"message": "两次输入的密码不一致",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
if err := common.ValidateNewAccountPassword(req.Password); err != nil {
|
||||
c.JSON(200, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Create root user
|
||||
hashedPassword, err := common.HashAccountPassword(req.Password)
|
||||
if err != nil {
|
||||
c.JSON(200, gin.H{
|
||||
"success": false,
|
||||
"message": "系统错误: " + err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
rootUser := model.User{
|
||||
Username: req.Username,
|
||||
Password: hashedPassword,
|
||||
Role: common.RoleRootUser,
|
||||
Status: common.UserStatusEnabled,
|
||||
DisplayName: "Root User",
|
||||
AccessToken: nil,
|
||||
Quota: 100000000,
|
||||
}
|
||||
err = model.DB.Create(&rootUser).Error
|
||||
if err != nil {
|
||||
c.JSON(200, gin.H{
|
||||
"success": false,
|
||||
"message": "创建管理员账号失败: " + err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Set operation modes
|
||||
operation_setting.SelfUseModeEnabled = req.SelfUseModeEnabled
|
||||
operation_setting.DemoSiteEnabled = req.DemoSiteEnabled
|
||||
|
||||
// Save operation modes to database for persistence
|
||||
err = model.UpdateOption("SelfUseModeEnabled", boolToString(req.SelfUseModeEnabled))
|
||||
if err != nil {
|
||||
c.JSON(200, gin.H{
|
||||
"success": false,
|
||||
"message": "保存自用模式设置失败: " + err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
err = model.UpdateOption("DemoSiteEnabled", boolToString(req.DemoSiteEnabled))
|
||||
if err != nil {
|
||||
c.JSON(200, gin.H{
|
||||
"success": false,
|
||||
"message": "保存演示站点模式设置失败: " + err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Update setup status
|
||||
constant.Setup = true
|
||||
|
||||
setup := model.Setup{
|
||||
Version: common.Version,
|
||||
InitializedAt: time.Now().Unix(),
|
||||
}
|
||||
err = model.DB.Create(&setup).Error
|
||||
if err != nil {
|
||||
c.JSON(200, gin.H{
|
||||
"success": false,
|
||||
"message": "系统初始化失败: " + err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(200, gin.H{
|
||||
"success": true,
|
||||
"message": "系统初始化成功",
|
||||
})
|
||||
}
|
||||
|
||||
func boolToString(b bool) string {
|
||||
if b {
|
||||
return "true"
|
||||
}
|
||||
return "false"
|
||||
}
|
||||
@@ -0,0 +1,547 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// ---- Shared types ----
|
||||
|
||||
type SubscriptionPlanDTO struct {
|
||||
Plan model.SubscriptionPlan `json:"plan"`
|
||||
}
|
||||
|
||||
type BillingPreferenceRequest struct {
|
||||
BillingPreference string `json:"billing_preference"`
|
||||
}
|
||||
|
||||
type SubscriptionBalancePayRequest struct {
|
||||
PlanId int `json:"plan_id"`
|
||||
}
|
||||
|
||||
// ---- User APIs ----
|
||||
|
||||
func GetSubscriptionPlans(c *gin.Context) {
|
||||
if !operation_setting.IsPaymentComplianceConfirmed() {
|
||||
common.ApiSuccess(c, []SubscriptionPlanDTO{})
|
||||
return
|
||||
}
|
||||
|
||||
var plans []model.SubscriptionPlan
|
||||
if err := model.DB.Where("enabled = ?", true).Order("sort_order desc, id desc").Find(&plans).Error; err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
result := make([]SubscriptionPlanDTO, 0, len(plans))
|
||||
for _, p := range plans {
|
||||
p.NormalizeDefaults()
|
||||
result = append(result, SubscriptionPlanDTO{
|
||||
Plan: p,
|
||||
})
|
||||
}
|
||||
common.ApiSuccess(c, result)
|
||||
}
|
||||
|
||||
func GetSubscriptionSelf(c *gin.Context) {
|
||||
userId := c.GetInt("id")
|
||||
settingMap, _ := model.GetUserSetting(userId, false)
|
||||
pref := common.NormalizeBillingPreference(settingMap.BillingPreference)
|
||||
|
||||
// Get all subscriptions (including expired)
|
||||
allSubscriptions, err := model.GetAllUserSubscriptions(userId)
|
||||
if err != nil {
|
||||
allSubscriptions = []model.SubscriptionSummary{}
|
||||
}
|
||||
|
||||
// Get active subscriptions for backward compatibility
|
||||
activeSubscriptions, err := model.GetAllActiveUserSubscriptions(userId)
|
||||
if err != nil {
|
||||
activeSubscriptions = []model.SubscriptionSummary{}
|
||||
}
|
||||
|
||||
common.ApiSuccess(c, gin.H{
|
||||
"billing_preference": pref,
|
||||
"subscriptions": activeSubscriptions, // all active subscriptions
|
||||
"all_subscriptions": allSubscriptions, // all subscriptions including expired
|
||||
})
|
||||
}
|
||||
|
||||
func UpdateSubscriptionPreference(c *gin.Context) {
|
||||
userId := c.GetInt("id")
|
||||
var req BillingPreferenceRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
common.ApiErrorMsg(c, "参数错误")
|
||||
return
|
||||
}
|
||||
pref := common.NormalizeBillingPreference(req.BillingPreference)
|
||||
|
||||
user, err := model.GetUserById(userId, true)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
current := user.GetSetting()
|
||||
current.BillingPreference = pref
|
||||
if err := model.UpdateUserSetting(user.Id, current); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, gin.H{"billing_preference": pref})
|
||||
}
|
||||
|
||||
func SubscriptionRequestBalancePay(c *gin.Context) {
|
||||
if !requirePaymentCompliance(c) {
|
||||
return
|
||||
}
|
||||
|
||||
userId := c.GetInt("id")
|
||||
var req SubscriptionBalancePayRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil || req.PlanId <= 0 {
|
||||
common.ApiErrorMsg(c, "参数错误")
|
||||
return
|
||||
}
|
||||
|
||||
if err := model.PurchaseSubscriptionWithBalance(userId, req.PlanId); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, nil)
|
||||
}
|
||||
|
||||
// ---- Admin APIs ----
|
||||
|
||||
func AdminListSubscriptionPlans(c *gin.Context) {
|
||||
var plans []model.SubscriptionPlan
|
||||
if err := model.DB.Order("sort_order desc, id desc").Find(&plans).Error; err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
result := make([]SubscriptionPlanDTO, 0, len(plans))
|
||||
for _, p := range plans {
|
||||
p.NormalizeDefaults()
|
||||
result = append(result, SubscriptionPlanDTO{
|
||||
Plan: p,
|
||||
})
|
||||
}
|
||||
common.ApiSuccess(c, result)
|
||||
}
|
||||
|
||||
type AdminUpsertSubscriptionPlanRequest struct {
|
||||
Plan model.SubscriptionPlan `json:"plan"`
|
||||
}
|
||||
|
||||
func AdminCreateSubscriptionPlan(c *gin.Context) {
|
||||
if !requirePaymentCompliance(c) {
|
||||
return
|
||||
}
|
||||
|
||||
var req AdminUpsertSubscriptionPlanRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
common.ApiErrorMsg(c, "参数错误")
|
||||
return
|
||||
}
|
||||
req.Plan.Id = 0
|
||||
if strings.TrimSpace(req.Plan.Title) == "" {
|
||||
common.ApiErrorMsg(c, "套餐标题不能为空")
|
||||
return
|
||||
}
|
||||
if req.Plan.PriceAmount < 0 {
|
||||
common.ApiErrorMsg(c, "价格不能为负数")
|
||||
return
|
||||
}
|
||||
if req.Plan.PriceAmount > 9999 {
|
||||
common.ApiErrorMsg(c, "价格不能超过9999")
|
||||
return
|
||||
}
|
||||
if req.Plan.Currency == "" {
|
||||
req.Plan.Currency = "USD"
|
||||
}
|
||||
req.Plan.Currency = "USD"
|
||||
if req.Plan.AllowBalancePay == nil {
|
||||
req.Plan.AllowBalancePay = common.GetPointer(true)
|
||||
}
|
||||
if req.Plan.AllowWalletOverflow == nil {
|
||||
req.Plan.AllowWalletOverflow = common.GetPointer(true)
|
||||
}
|
||||
if req.Plan.DurationUnit == "" {
|
||||
req.Plan.DurationUnit = model.SubscriptionDurationMonth
|
||||
}
|
||||
if req.Plan.DurationValue <= 0 && req.Plan.DurationUnit != model.SubscriptionDurationCustom {
|
||||
req.Plan.DurationValue = 1
|
||||
}
|
||||
if req.Plan.MaxPurchasePerUser < 0 {
|
||||
common.ApiErrorMsg(c, "购买上限不能为负数")
|
||||
return
|
||||
}
|
||||
if req.Plan.TotalAmount < 0 {
|
||||
common.ApiErrorMsg(c, "总额度不能为负数")
|
||||
return
|
||||
}
|
||||
req.Plan.UpgradeGroup = strings.TrimSpace(req.Plan.UpgradeGroup)
|
||||
if req.Plan.UpgradeGroup != "" {
|
||||
if _, ok := ratio_setting.GetGroupRatioCopy()[req.Plan.UpgradeGroup]; !ok {
|
||||
common.ApiErrorMsg(c, "升级分组不存在")
|
||||
return
|
||||
}
|
||||
}
|
||||
req.Plan.DowngradeGroup = strings.TrimSpace(req.Plan.DowngradeGroup)
|
||||
if req.Plan.DowngradeGroup != "" {
|
||||
if _, ok := ratio_setting.GetGroupRatioCopy()[req.Plan.DowngradeGroup]; !ok {
|
||||
common.ApiErrorMsg(c, "降级分组不存在")
|
||||
return
|
||||
}
|
||||
}
|
||||
req.Plan.QuotaResetPeriod = model.NormalizeResetPeriod(req.Plan.QuotaResetPeriod)
|
||||
if req.Plan.QuotaResetPeriod == model.SubscriptionResetCustom && req.Plan.QuotaResetCustomSeconds <= 0 {
|
||||
common.ApiErrorMsg(c, "自定义重置周期需大于0秒")
|
||||
return
|
||||
}
|
||||
err := model.DB.Create(&req.Plan).Error
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
model.InvalidateSubscriptionPlanCache(req.Plan.Id)
|
||||
common.ApiSuccess(c, req.Plan)
|
||||
}
|
||||
|
||||
func AdminUpdateSubscriptionPlan(c *gin.Context) {
|
||||
if !requirePaymentCompliance(c) {
|
||||
return
|
||||
}
|
||||
|
||||
id, _ := strconv.Atoi(c.Param("id"))
|
||||
if id <= 0 {
|
||||
common.ApiErrorMsg(c, "无效的ID")
|
||||
return
|
||||
}
|
||||
var req AdminUpsertSubscriptionPlanRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
common.ApiErrorMsg(c, "参数错误")
|
||||
return
|
||||
}
|
||||
if strings.TrimSpace(req.Plan.Title) == "" {
|
||||
common.ApiErrorMsg(c, "套餐标题不能为空")
|
||||
return
|
||||
}
|
||||
if req.Plan.PriceAmount < 0 {
|
||||
common.ApiErrorMsg(c, "价格不能为负数")
|
||||
return
|
||||
}
|
||||
if req.Plan.PriceAmount > 9999 {
|
||||
common.ApiErrorMsg(c, "价格不能超过9999")
|
||||
return
|
||||
}
|
||||
req.Plan.Id = id
|
||||
if req.Plan.Currency == "" {
|
||||
req.Plan.Currency = "USD"
|
||||
}
|
||||
req.Plan.Currency = "USD"
|
||||
if req.Plan.DurationUnit == "" {
|
||||
req.Plan.DurationUnit = model.SubscriptionDurationMonth
|
||||
}
|
||||
if req.Plan.DurationValue <= 0 && req.Plan.DurationUnit != model.SubscriptionDurationCustom {
|
||||
req.Plan.DurationValue = 1
|
||||
}
|
||||
if req.Plan.MaxPurchasePerUser < 0 {
|
||||
common.ApiErrorMsg(c, "购买上限不能为负数")
|
||||
return
|
||||
}
|
||||
if req.Plan.TotalAmount < 0 {
|
||||
common.ApiErrorMsg(c, "总额度不能为负数")
|
||||
return
|
||||
}
|
||||
req.Plan.UpgradeGroup = strings.TrimSpace(req.Plan.UpgradeGroup)
|
||||
if req.Plan.UpgradeGroup != "" {
|
||||
if _, ok := ratio_setting.GetGroupRatioCopy()[req.Plan.UpgradeGroup]; !ok {
|
||||
common.ApiErrorMsg(c, "升级分组不存在")
|
||||
return
|
||||
}
|
||||
}
|
||||
req.Plan.DowngradeGroup = strings.TrimSpace(req.Plan.DowngradeGroup)
|
||||
if req.Plan.DowngradeGroup != "" {
|
||||
if _, ok := ratio_setting.GetGroupRatioCopy()[req.Plan.DowngradeGroup]; !ok {
|
||||
common.ApiErrorMsg(c, "降级分组不存在")
|
||||
return
|
||||
}
|
||||
}
|
||||
req.Plan.QuotaResetPeriod = model.NormalizeResetPeriod(req.Plan.QuotaResetPeriod)
|
||||
if req.Plan.QuotaResetPeriod == model.SubscriptionResetCustom && req.Plan.QuotaResetCustomSeconds <= 0 {
|
||||
common.ApiErrorMsg(c, "自定义重置周期需大于0秒")
|
||||
return
|
||||
}
|
||||
|
||||
err := model.DB.Transaction(func(tx *gorm.DB) error {
|
||||
// update plan (allow zero values updates with map)
|
||||
updateMap := map[string]any{
|
||||
"title": req.Plan.Title,
|
||||
"subtitle": req.Plan.Subtitle,
|
||||
"price_amount": req.Plan.PriceAmount,
|
||||
"currency": req.Plan.Currency,
|
||||
"duration_unit": req.Plan.DurationUnit,
|
||||
"duration_value": req.Plan.DurationValue,
|
||||
"custom_seconds": req.Plan.CustomSeconds,
|
||||
"enabled": req.Plan.Enabled,
|
||||
"sort_order": req.Plan.SortOrder,
|
||||
"stripe_price_id": req.Plan.StripePriceId,
|
||||
"creem_product_id": req.Plan.CreemProductId,
|
||||
"waffo_pancake_product_id": req.Plan.WaffoPancakeProductId,
|
||||
"max_purchase_per_user": req.Plan.MaxPurchasePerUser,
|
||||
"total_amount": req.Plan.TotalAmount,
|
||||
"upgrade_group": req.Plan.UpgradeGroup,
|
||||
"downgrade_group": req.Plan.DowngradeGroup,
|
||||
"quota_reset_period": req.Plan.QuotaResetPeriod,
|
||||
"quota_reset_custom_seconds": req.Plan.QuotaResetCustomSeconds,
|
||||
"updated_at": common.GetTimestamp(),
|
||||
}
|
||||
if req.Plan.AllowBalancePay != nil {
|
||||
updateMap["allow_balance_pay"] = *req.Plan.AllowBalancePay
|
||||
}
|
||||
if req.Plan.AllowWalletOverflow != nil {
|
||||
updateMap["allow_wallet_overflow"] = *req.Plan.AllowWalletOverflow
|
||||
}
|
||||
if err := tx.Model(&model.SubscriptionPlan{}).Where("id = ?", id).Updates(updateMap).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
model.InvalidateSubscriptionPlanCache(id)
|
||||
common.ApiSuccess(c, nil)
|
||||
}
|
||||
|
||||
type AdminUpdateSubscriptionPlanStatusRequest struct {
|
||||
Enabled *bool `json:"enabled"`
|
||||
}
|
||||
|
||||
func AdminUpdateSubscriptionPlanStatus(c *gin.Context) {
|
||||
if !requirePaymentCompliance(c) {
|
||||
return
|
||||
}
|
||||
|
||||
id, _ := strconv.Atoi(c.Param("id"))
|
||||
if id <= 0 {
|
||||
common.ApiErrorMsg(c, "无效的ID")
|
||||
return
|
||||
}
|
||||
var req AdminUpdateSubscriptionPlanStatusRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil || req.Enabled == nil {
|
||||
common.ApiErrorMsg(c, "参数错误")
|
||||
return
|
||||
}
|
||||
if err := model.DB.Model(&model.SubscriptionPlan{}).Where("id = ?", id).Update("enabled", *req.Enabled).Error; err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
model.InvalidateSubscriptionPlanCache(id)
|
||||
common.ApiSuccess(c, nil)
|
||||
}
|
||||
|
||||
type AdminBindSubscriptionRequest struct {
|
||||
UserId int `json:"user_id"`
|
||||
PlanId int `json:"plan_id"`
|
||||
}
|
||||
|
||||
func AdminBindSubscription(c *gin.Context) {
|
||||
if !requirePaymentCompliance(c) {
|
||||
return
|
||||
}
|
||||
|
||||
var req AdminBindSubscriptionRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil || req.UserId <= 0 || req.PlanId <= 0 {
|
||||
common.ApiErrorMsg(c, "参数错误")
|
||||
return
|
||||
}
|
||||
msg, err := model.AdminBindSubscription(req.UserId, req.PlanId, "")
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if msg != "" {
|
||||
common.ApiSuccess(c, gin.H{"message": msg})
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, nil)
|
||||
}
|
||||
|
||||
// ---- Admin: user subscription management ----
|
||||
|
||||
func AdminListUserSubscriptions(c *gin.Context) {
|
||||
userId, _ := strconv.Atoi(c.Param("id"))
|
||||
if userId <= 0 {
|
||||
common.ApiErrorMsg(c, "无效的用户ID")
|
||||
return
|
||||
}
|
||||
subs, err := model.GetAllUserSubscriptions(userId)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, subs)
|
||||
}
|
||||
|
||||
type AdminCreateUserSubscriptionRequest struct {
|
||||
PlanId int `json:"plan_id"`
|
||||
}
|
||||
|
||||
type AdminResetSubscriptionRequest struct {
|
||||
PlanId int `json:"plan_id"`
|
||||
AdvanceResetTime *bool `json:"advance_reset_time"`
|
||||
}
|
||||
|
||||
func resolveAdvanceResetTime(value *bool) bool {
|
||||
if value == nil {
|
||||
return true
|
||||
}
|
||||
return *value
|
||||
}
|
||||
|
||||
func recordSubscriptionResetUserLogs(c *gin.Context, result *model.SubscriptionResetResult, adminInfo *model.AuditAdminInfo) {
|
||||
if result == nil || result.ResetCount == 0 {
|
||||
return
|
||||
}
|
||||
content := fmt.Sprintf("管理员重置订阅套餐 %s（ID: %d）额度", result.PlanTitle, result.PlanId)
|
||||
for _, userId := range result.AffectedUserIds {
|
||||
model.RecordLogWithAdminInfo(userId, model.LogTypeManage, content, adminInfo, nil, c)
|
||||
}
|
||||
}
|
||||
|
||||
// AdminCreateUserSubscription creates a new user subscription from a plan (no payment).
|
||||
func AdminCreateUserSubscription(c *gin.Context) {
|
||||
if !requirePaymentCompliance(c) {
|
||||
return
|
||||
}
|
||||
|
||||
userId, _ := strconv.Atoi(c.Param("id"))
|
||||
if userId <= 0 {
|
||||
common.ApiErrorMsg(c, "无效的用户ID")
|
||||
return
|
||||
}
|
||||
var req AdminCreateUserSubscriptionRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil || req.PlanId <= 0 {
|
||||
common.ApiErrorMsg(c, "参数错误")
|
||||
return
|
||||
}
|
||||
msg, err := model.AdminBindSubscription(userId, req.PlanId, "")
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if msg != "" {
|
||||
common.ApiSuccess(c, gin.H{"message": msg})
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, nil)
|
||||
}
|
||||
|
||||
func AdminResetUserSubscriptionsByPlan(c *gin.Context) {
|
||||
userId, _ := strconv.Atoi(c.Param("id"))
|
||||
if userId <= 0 {
|
||||
common.ApiErrorMsg(c, "无效的用户ID")
|
||||
return
|
||||
}
|
||||
var req AdminResetSubscriptionRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
common.ApiErrorMsg(c, "参数错误")
|
||||
return
|
||||
}
|
||||
if req.PlanId <= 0 {
|
||||
common.ApiErrorMsg(c, "参数错误")
|
||||
return
|
||||
}
|
||||
advanceResetTime := resolveAdvanceResetTime(req.AdvanceResetTime)
|
||||
result, err := model.AdminResetUserSubscriptionsByPlan(userId, req.PlanId, advanceResetTime)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
recordSubscriptionResetUserLogs(c, result, auditOperatorInfo(c))
|
||||
recordManageAuditFor(c, userId, "subscription.user_plan_reset", map[string]any{
|
||||
"target_user_id": userId,
|
||||
"plan_id": result.PlanId,
|
||||
"plan_title": result.PlanTitle,
|
||||
"reset_count": result.ResetCount,
|
||||
"user_count": result.UserCount,
|
||||
"advance_reset_time": result.AdvanceResetTime,
|
||||
})
|
||||
common.ApiSuccess(c, result)
|
||||
}
|
||||
|
||||
func AdminResetPlanSubscriptions(c *gin.Context) {
|
||||
planId, _ := strconv.Atoi(c.Param("id"))
|
||||
if planId <= 0 {
|
||||
common.ApiErrorMsg(c, "无效的ID")
|
||||
return
|
||||
}
|
||||
var req AdminResetSubscriptionRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
common.ApiErrorMsg(c, "参数错误")
|
||||
return
|
||||
}
|
||||
advanceResetTime := resolveAdvanceResetTime(req.AdvanceResetTime)
|
||||
result, err := model.AdminResetPlanSubscriptions(planId, advanceResetTime)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
recordSubscriptionResetUserLogs(c, result, auditOperatorInfo(c))
|
||||
common.SysLog(fmt.Sprintf("admin reset subscription plan %d quota: reset_count=%d user_count=%d advance_reset_time=%t",
|
||||
result.PlanId, result.ResetCount, result.UserCount, result.AdvanceResetTime))
|
||||
recordManageAudit(c, "subscription.plan_reset", map[string]any{
|
||||
"plan_id": result.PlanId,
|
||||
"plan_title": result.PlanTitle,
|
||||
"reset_count": result.ResetCount,
|
||||
"user_count": result.UserCount,
|
||||
"advance_reset_time": result.AdvanceResetTime,
|
||||
})
|
||||
common.ApiSuccess(c, result)
|
||||
}
|
||||
|
||||
// AdminInvalidateUserSubscription cancels a user subscription immediately.
|
||||
func AdminInvalidateUserSubscription(c *gin.Context) {
|
||||
subId, _ := strconv.Atoi(c.Param("id"))
|
||||
if subId <= 0 {
|
||||
common.ApiErrorMsg(c, "无效的订阅ID")
|
||||
return
|
||||
}
|
||||
msg, err := model.AdminInvalidateUserSubscription(subId)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if msg != "" {
|
||||
common.ApiSuccess(c, gin.H{"message": msg})
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, nil)
|
||||
}
|
||||
|
||||
// AdminDeleteUserSubscription hard-deletes a user subscription.
|
||||
func AdminDeleteUserSubscription(c *gin.Context) {
|
||||
subId, _ := strconv.Atoi(c.Param("id"))
|
||||
if subId <= 0 {
|
||||
common.ApiErrorMsg(c, "无效的订阅ID")
|
||||
return
|
||||
}
|
||||
msg, err := model.AdminDeleteUserSubscription(subId)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if msg != "" {
|
||||
common.ApiSuccess(c, gin.H{"message": msg})
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, nil)
|
||||
}
|
||||
@@ -0,0 +1,136 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/setting"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/thanhpk/randstr"
|
||||
)
|
||||
|
||||
type SubscriptionCreemPayRequest struct {
|
||||
PlanId int `json:"plan_id"`
|
||||
}
|
||||
|
||||
func SubscriptionRequestCreemPay(c *gin.Context) {
|
||||
if !requirePaymentCompliance(c) {
|
||||
return
|
||||
}
|
||||
|
||||
var req SubscriptionCreemPayRequest
|
||||
|
||||
// Keep body for debugging consistency (like RequestCreemPay)
|
||||
bodyBytes, err := io.ReadAll(c.Request.Body)
|
||||
if err != nil {
|
||||
logger.LogError(c.Request.Context(), fmt.Sprintf("Creem 订阅支付请求读取失败 error=%q", err.Error()))
|
||||
c.JSON(http.StatusOK, gin.H{"message": "error", "data": "read query error"})
|
||||
return
|
||||
}
|
||||
c.Request.Body = io.NopCloser(bytes.NewReader(bodyBytes))
|
||||
|
||||
if err := c.ShouldBindJSON(&req); err != nil || req.PlanId <= 0 {
|
||||
c.JSON(http.StatusOK, gin.H{"message": "error", "data": "参数错误"})
|
||||
return
|
||||
}
|
||||
|
||||
plan, err := model.GetSubscriptionPlanById(req.PlanId)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if !plan.Enabled {
|
||||
common.ApiErrorMsg(c, "套餐未启用")
|
||||
return
|
||||
}
|
||||
if plan.CreemProductId == "" {
|
||||
common.ApiErrorMsg(c, "该套餐未配置 CreemProductId")
|
||||
return
|
||||
}
|
||||
if setting.CreemWebhookSecret == "" && !setting.CreemTestMode {
|
||||
common.ApiErrorMsg(c, "Creem Webhook 未配置")
|
||||
return
|
||||
}
|
||||
|
||||
userId := c.GetInt("id")
|
||||
user, err := model.GetUserById(userId, false)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if user == nil {
|
||||
common.ApiErrorMsg(c, "用户不存在")
|
||||
return
|
||||
}
|
||||
|
||||
if plan.MaxPurchasePerUser > 0 {
|
||||
count, err := model.CountUserSubscriptionsByPlan(userId, plan.Id)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if count >= int64(plan.MaxPurchasePerUser) {
|
||||
common.ApiErrorMsg(c, "已达到该套餐购买上限")
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
reference := "sub-creem-ref-" + randstr.String(6)
|
||||
referenceId := "sub_ref_" + common.Sha1([]byte(reference+time.Now().String()+user.Username))
|
||||
|
||||
// create pending order first
|
||||
order := &model.SubscriptionOrder{
|
||||
UserId: userId,
|
||||
PlanId: plan.Id,
|
||||
Money: plan.PriceAmount,
|
||||
TradeNo: referenceId,
|
||||
PaymentMethod: model.PaymentMethodCreem,
|
||||
PaymentProvider: model.PaymentProviderCreem,
|
||||
CreateTime: time.Now().Unix(),
|
||||
Status: common.TopUpStatusPending,
|
||||
}
|
||||
if err := order.Insert(); err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{"message": "error", "data": "创建订单失败"})
|
||||
return
|
||||
}
|
||||
|
||||
// Reuse Creem checkout generator by building a lightweight product reference.
|
||||
currency := "USD"
|
||||
switch operation_setting.GetGeneralSetting().QuotaDisplayType {
|
||||
case operation_setting.QuotaDisplayTypeCNY:
|
||||
currency = "CNY"
|
||||
case operation_setting.QuotaDisplayTypeUSD:
|
||||
currency = "USD"
|
||||
default:
|
||||
currency = "USD"
|
||||
}
|
||||
product := &CreemProduct{
|
||||
ProductId: plan.CreemProductId,
|
||||
Name: plan.Title,
|
||||
Price: plan.PriceAmount,
|
||||
Currency: currency,
|
||||
Quota: 0,
|
||||
}
|
||||
|
||||
checkoutUrl, err := genCreemLink(c.Request.Context(), referenceId, product, user.Email, user.Username)
|
||||
if err != nil {
|
||||
logger.LogError(c.Request.Context(), fmt.Sprintf("Creem 订阅支付链接创建失败 trade_no=%s product_id=%s error=%q", referenceId, product.ProductId, err.Error()))
|
||||
c.JSON(http.StatusOK, gin.H{"message": "error", "data": "拉起支付失败"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": "success",
|
||||
"data": gin.H{
|
||||
"checkout_url": checkoutUrl,
|
||||
"order_id": referenceId,
|
||||
},
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,220 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/Calcium-Ion/go-epay/epay"
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/samber/lo"
|
||||
)
|
||||
|
||||
type SubscriptionEpayPayRequest struct {
|
||||
PlanId int `json:"plan_id"`
|
||||
PaymentMethod string `json:"payment_method"`
|
||||
}
|
||||
|
||||
func SubscriptionRequestEpay(c *gin.Context) {
|
||||
if !requirePaymentCompliance(c) {
|
||||
return
|
||||
}
|
||||
|
||||
var req SubscriptionEpayPayRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil || req.PlanId <= 0 {
|
||||
common.ApiErrorMsg(c, "参数错误")
|
||||
return
|
||||
}
|
||||
|
||||
plan, err := model.GetSubscriptionPlanById(req.PlanId)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if !plan.Enabled {
|
||||
common.ApiErrorMsg(c, "套餐未启用")
|
||||
return
|
||||
}
|
||||
if plan.PriceAmount < 0.01 {
|
||||
common.ApiErrorMsg(c, "套餐金额过低")
|
||||
return
|
||||
}
|
||||
if !operation_setting.ContainsPayMethod(req.PaymentMethod) {
|
||||
common.ApiErrorMsg(c, "支付方式不存在")
|
||||
return
|
||||
}
|
||||
|
||||
userId := c.GetInt("id")
|
||||
if plan.MaxPurchasePerUser > 0 {
|
||||
count, err := model.CountUserSubscriptionsByPlan(userId, plan.Id)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if count >= int64(plan.MaxPurchasePerUser) {
|
||||
common.ApiErrorMsg(c, "已达到该套餐购买上限")
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
callBackAddress := service.GetCallbackAddress()
|
||||
returnUrl, err := url.Parse(callBackAddress + "/api/subscription/epay/return")
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, "回调地址配置错误")
|
||||
return
|
||||
}
|
||||
notifyUrl, err := url.Parse(callBackAddress + "/api/subscription/epay/notify")
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, "回调地址配置错误")
|
||||
return
|
||||
}
|
||||
|
||||
tradeNo := fmt.Sprintf("%s%d", common.GetRandomString(6), time.Now().Unix())
|
||||
tradeNo = fmt.Sprintf("SUBUSR%dNO%s", userId, tradeNo)
|
||||
|
||||
client := GetEpayClient()
|
||||
if client == nil {
|
||||
common.ApiErrorMsg(c, "当前管理员未配置支付信息")
|
||||
return
|
||||
}
|
||||
|
||||
order := &model.SubscriptionOrder{
|
||||
UserId: userId,
|
||||
PlanId: plan.Id,
|
||||
Money: plan.PriceAmount,
|
||||
TradeNo: tradeNo,
|
||||
PaymentMethod: req.PaymentMethod,
|
||||
PaymentProvider: model.PaymentProviderEpay,
|
||||
CreateTime: time.Now().Unix(),
|
||||
Status: common.TopUpStatusPending,
|
||||
}
|
||||
if err := order.Insert(); err != nil {
|
||||
common.ApiErrorMsg(c, "创建订单失败")
|
||||
return
|
||||
}
|
||||
uri, params, err := client.Purchase(&epay.PurchaseArgs{
|
||||
Type: req.PaymentMethod,
|
||||
ServiceTradeNo: tradeNo,
|
||||
Name: fmt.Sprintf("SUB:%s", plan.Title),
|
||||
Money: strconv.FormatFloat(plan.PriceAmount, 'f', 2, 64),
|
||||
Device: epay.PC,
|
||||
NotifyUrl: notifyUrl,
|
||||
ReturnUrl: returnUrl,
|
||||
})
|
||||
if err != nil {
|
||||
_ = model.ExpireSubscriptionOrder(tradeNo, model.PaymentProviderEpay)
|
||||
common.ApiErrorMsg(c, "拉起支付失败")
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "success", "data": params, "url": uri})
|
||||
}
|
||||
|
||||
func SubscriptionEpayNotify(c *gin.Context) {
|
||||
var params map[string]string
|
||||
|
||||
if c.Request.Method == "POST" {
|
||||
// POST 请求：从 POST body 解析参数
|
||||
if err := c.Request.ParseForm(); err != nil {
|
||||
_, _ = c.Writer.Write([]byte("fail"))
|
||||
return
|
||||
}
|
||||
params = lo.Reduce(lo.Keys(c.Request.PostForm), func(r map[string]string, t string, i int) map[string]string {
|
||||
r[t] = c.Request.PostForm.Get(t)
|
||||
return r
|
||||
}, map[string]string{})
|
||||
} else {
|
||||
// GET 请求：从 URL Query 解析参数
|
||||
params = lo.Reduce(lo.Keys(c.Request.URL.Query()), func(r map[string]string, t string, i int) map[string]string {
|
||||
r[t] = c.Request.URL.Query().Get(t)
|
||||
return r
|
||||
}, map[string]string{})
|
||||
}
|
||||
|
||||
if len(params) == 0 {
|
||||
_, _ = c.Writer.Write([]byte("fail"))
|
||||
return
|
||||
}
|
||||
|
||||
client := GetEpayClient()
|
||||
if client == nil {
|
||||
_, _ = c.Writer.Write([]byte("fail"))
|
||||
return
|
||||
}
|
||||
verifyInfo, err := client.Verify(params)
|
||||
if err != nil || !verifyInfo.VerifyStatus {
|
||||
_, _ = c.Writer.Write([]byte("fail"))
|
||||
return
|
||||
}
|
||||
|
||||
if verifyInfo.TradeStatus != epay.StatusTradeSuccess {
|
||||
_, _ = c.Writer.Write([]byte("fail"))
|
||||
return
|
||||
}
|
||||
|
||||
LockOrder(verifyInfo.ServiceTradeNo)
|
||||
defer UnlockOrder(verifyInfo.ServiceTradeNo)
|
||||
|
||||
if err := model.CompleteSubscriptionOrder(verifyInfo.ServiceTradeNo, common.GetJsonString(verifyInfo), model.PaymentProviderEpay, verifyInfo.Type); err != nil {
|
||||
_, _ = c.Writer.Write([]byte("fail"))
|
||||
return
|
||||
}
|
||||
|
||||
_, _ = c.Writer.Write([]byte("success"))
|
||||
}
|
||||
|
||||
// SubscriptionEpayReturn handles browser return after payment.
|
||||
// It verifies the payload and completes the order, then redirects to console.
|
||||
func SubscriptionEpayReturn(c *gin.Context) {
|
||||
var params map[string]string
|
||||
|
||||
if c.Request.Method == "POST" {
|
||||
// POST 请求：从 POST body 解析参数
|
||||
if err := c.Request.ParseForm(); err != nil {
|
||||
c.Redirect(http.StatusFound, paymentReturnPath("/wallet?pay=fail"))
|
||||
return
|
||||
}
|
||||
params = lo.Reduce(lo.Keys(c.Request.PostForm), func(r map[string]string, t string, i int) map[string]string {
|
||||
r[t] = c.Request.PostForm.Get(t)
|
||||
return r
|
||||
}, map[string]string{})
|
||||
} else {
|
||||
// GET 请求：从 URL Query 解析参数
|
||||
params = lo.Reduce(lo.Keys(c.Request.URL.Query()), func(r map[string]string, t string, i int) map[string]string {
|
||||
r[t] = c.Request.URL.Query().Get(t)
|
||||
return r
|
||||
}, map[string]string{})
|
||||
}
|
||||
|
||||
if len(params) == 0 {
|
||||
c.Redirect(http.StatusFound, paymentReturnPath("/wallet?pay=fail"))
|
||||
return
|
||||
}
|
||||
|
||||
client := GetEpayClient()
|
||||
if client == nil {
|
||||
c.Redirect(http.StatusFound, paymentReturnPath("/wallet?pay=fail"))
|
||||
return
|
||||
}
|
||||
verifyInfo, err := client.Verify(params)
|
||||
if err != nil || !verifyInfo.VerifyStatus {
|
||||
c.Redirect(http.StatusFound, paymentReturnPath("/wallet?pay=fail"))
|
||||
return
|
||||
}
|
||||
if verifyInfo.TradeStatus == epay.StatusTradeSuccess {
|
||||
LockOrder(verifyInfo.ServiceTradeNo)
|
||||
defer UnlockOrder(verifyInfo.ServiceTradeNo)
|
||||
if err := model.CompleteSubscriptionOrder(verifyInfo.ServiceTradeNo, common.GetJsonString(verifyInfo), model.PaymentProviderEpay, verifyInfo.Type); err != nil {
|
||||
c.Redirect(http.StatusFound, paymentReturnPath("/wallet?pay=fail"))
|
||||
return
|
||||
}
|
||||
c.Redirect(http.StatusFound, paymentReturnPath("/wallet?pay=success"))
|
||||
return
|
||||
}
|
||||
c.Redirect(http.StatusFound, paymentReturnPath("/wallet?pay=pending"))
|
||||
}
|
||||
@@ -0,0 +1,142 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stripe/stripe-go/v81"
|
||||
"github.com/stripe/stripe-go/v81/checkout/session"
|
||||
"github.com/thanhpk/randstr"
|
||||
)
|
||||
|
||||
type SubscriptionStripePayRequest struct {
|
||||
PlanId int `json:"plan_id"`
|
||||
}
|
||||
|
||||
func SubscriptionRequestStripePay(c *gin.Context) {
|
||||
if !requirePaymentCompliance(c) {
|
||||
return
|
||||
}
|
||||
|
||||
var req SubscriptionStripePayRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil || req.PlanId <= 0 {
|
||||
common.ApiErrorMsg(c, "参数错误")
|
||||
return
|
||||
}
|
||||
|
||||
plan, err := model.GetSubscriptionPlanById(req.PlanId)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if !plan.Enabled {
|
||||
common.ApiErrorMsg(c, "套餐未启用")
|
||||
return
|
||||
}
|
||||
if plan.StripePriceId == "" {
|
||||
common.ApiErrorMsg(c, "该套餐未配置 StripePriceId")
|
||||
return
|
||||
}
|
||||
if !strings.HasPrefix(setting.StripeApiSecret, "sk_") && !strings.HasPrefix(setting.StripeApiSecret, "rk_") {
|
||||
common.ApiErrorMsg(c, "Stripe 未配置或密钥无效")
|
||||
return
|
||||
}
|
||||
if setting.StripeWebhookSecret == "" {
|
||||
common.ApiErrorMsg(c, "Stripe Webhook 未配置")
|
||||
return
|
||||
}
|
||||
|
||||
userId := c.GetInt("id")
|
||||
user, err := model.GetUserById(userId, false)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if user == nil {
|
||||
common.ApiErrorMsg(c, "用户不存在")
|
||||
return
|
||||
}
|
||||
|
||||
if plan.MaxPurchasePerUser > 0 {
|
||||
count, err := model.CountUserSubscriptionsByPlan(userId, plan.Id)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if count >= int64(plan.MaxPurchasePerUser) {
|
||||
common.ApiErrorMsg(c, "已达到该套餐购买上限")
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
reference := fmt.Sprintf("sub-stripe-ref-%d-%d-%s", user.Id, time.Now().UnixMilli(), randstr.String(4))
|
||||
referenceId := "sub_ref_" + common.Sha1([]byte(reference))
|
||||
|
||||
payLink, err := genStripeSubscriptionLink(referenceId, user.StripeCustomer, user.Email, plan.StripePriceId)
|
||||
if err != nil {
|
||||
logger.LogError(c.Request.Context(), fmt.Sprintf("Stripe 订阅支付链接创建失败 trade_no=%s plan_id=%d error=%q", referenceId, plan.Id, err.Error()))
|
||||
c.JSON(http.StatusOK, gin.H{"message": "error", "data": "拉起支付失败"})
|
||||
return
|
||||
}
|
||||
|
||||
order := &model.SubscriptionOrder{
|
||||
UserId: userId,
|
||||
PlanId: plan.Id,
|
||||
Money: plan.PriceAmount,
|
||||
TradeNo: referenceId,
|
||||
PaymentMethod: model.PaymentMethodStripe,
|
||||
PaymentProvider: model.PaymentProviderStripe,
|
||||
CreateTime: time.Now().Unix(),
|
||||
Status: common.TopUpStatusPending,
|
||||
}
|
||||
if err := order.Insert(); err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{"message": "error", "data": "创建订单失败"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": "success",
|
||||
"data": gin.H{
|
||||
"pay_link": payLink,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func genStripeSubscriptionLink(referenceId string, customerId string, email string, priceId string) (string, error) {
|
||||
stripe.Key = setting.StripeApiSecret
|
||||
|
||||
params := &stripe.CheckoutSessionParams{
|
||||
ClientReferenceID: stripe.String(referenceId),
|
||||
SuccessURL: stripe.String(paymentReturnPath("/wallet")),
|
||||
CancelURL: stripe.String(paymentReturnPath("/wallet")),
|
||||
LineItems: []*stripe.CheckoutSessionLineItemParams{
|
||||
{
|
||||
Price: stripe.String(priceId),
|
||||
Quantity: stripe.Int64(1),
|
||||
},
|
||||
},
|
||||
Mode: stripe.String(string(stripe.CheckoutSessionModeSubscription)),
|
||||
}
|
||||
|
||||
if "" == customerId {
|
||||
if "" != email {
|
||||
params.CustomerEmail = stripe.String(email)
|
||||
}
|
||||
params.CustomerCreation = stripe.String(string(stripe.CheckoutSessionCustomerCreationAlways))
|
||||
} else {
|
||||
params.Customer = stripe.String(customerId)
|
||||
}
|
||||
|
||||
result, err := session.New(params)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return result.URL, nil
|
||||
}
|
||||
@@ -0,0 +1,130 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/shopspring/decimal"
|
||||
"github.com/thanhpk/randstr"
|
||||
)
|
||||
|
||||
type SubscriptionWaffoPancakePayRequest struct {
|
||||
PlanId int `json:"plan_id"`
|
||||
}
|
||||
|
||||
func SubscriptionRequestWaffoPancakePay(c *gin.Context) {
|
||||
if !requirePaymentCompliance(c) {
|
||||
return
|
||||
}
|
||||
|
||||
var req SubscriptionWaffoPancakePayRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil || req.PlanId <= 0 {
|
||||
common.ApiErrorMsg(c, "参数错误")
|
||||
return
|
||||
}
|
||||
|
||||
plan, err := model.GetSubscriptionPlanById(req.PlanId)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if !plan.Enabled {
|
||||
common.ApiErrorMsg(c, "套餐未启用")
|
||||
return
|
||||
}
|
||||
if strings.TrimSpace(plan.WaffoPancakeProductId) == "" {
|
||||
common.ApiErrorMsg(c, "该套餐未配置 WaffoPancakeProductId")
|
||||
return
|
||||
}
|
||||
// Plan targets its own Pancake product, so we only require credentials
|
||||
// here — not the gateway-level WaffoPancakeProductID.
|
||||
if strings.TrimSpace(setting.WaffoPancakeMerchantID) == "" ||
|
||||
strings.TrimSpace(setting.WaffoPancakePrivateKey) == "" {
|
||||
common.ApiErrorMsg(c, "Waffo Pancake 未配置或密钥无效")
|
||||
return
|
||||
}
|
||||
|
||||
userId := c.GetInt("id")
|
||||
user, err := model.GetUserById(userId, false)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if user == nil {
|
||||
common.ApiErrorMsg(c, "用户不存在")
|
||||
return
|
||||
}
|
||||
|
||||
if plan.MaxPurchasePerUser > 0 {
|
||||
count, err := model.CountUserSubscriptionsByPlan(userId, plan.Id)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if count >= int64(plan.MaxPurchasePerUser) {
|
||||
common.ApiErrorMsg(c, "已达到该套餐购买上限")
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// WAFFO_PANCAKE_SUB- prefix (vs. wallet's WAFFO_PANCAKE-) drives webhook
|
||||
// dispatch in WaffoPancakeWebhook.
|
||||
tradeNo := fmt.Sprintf("WAFFO_PANCAKE_SUB-%d-%d-%s", userId, time.Now().UnixMilli(), randstr.String(6))
|
||||
|
||||
order := &model.SubscriptionOrder{
|
||||
UserId: userId,
|
||||
PlanId: plan.Id,
|
||||
Money: plan.PriceAmount,
|
||||
TradeNo: tradeNo,
|
||||
PaymentMethod: model.PaymentMethodWaffoPancake,
|
||||
PaymentProvider: model.PaymentProviderWaffoPancake,
|
||||
CreateTime: time.Now().Unix(),
|
||||
Status: common.TopUpStatusPending,
|
||||
}
|
||||
if err := order.Insert(); err != nil {
|
||||
logger.LogError(c.Request.Context(), fmt.Sprintf("Waffo Pancake 订阅订单创建失败 user_id=%d plan_id=%d trade_no=%s error=%q", userId, plan.Id, tradeNo, err.Error()))
|
||||
c.JSON(http.StatusOK, gin.H{"message": "error", "data": "创建订单失败"})
|
||||
return
|
||||
}
|
||||
|
||||
expiresInSeconds := 45 * 60
|
||||
session, err := service.CreateWaffoPancakeCheckoutSession(c.Request.Context(), &service.WaffoPancakeCreateSessionParams{
|
||||
ProductID: plan.WaffoPancakeProductId,
|
||||
BuyerIdentity: service.WaffoPancakeBuyerIdentityFromUserID(user.Id),
|
||||
PriceSnapshot: &service.WaffoPancakePriceSnapshot{
|
||||
Amount: decimal.NewFromFloat(plan.PriceAmount).StringFixed(2),
|
||||
TaxCategory: "saas",
|
||||
},
|
||||
BuyerEmail: getWaffoPancakeBuyerEmail(user),
|
||||
ExpiresInSeconds: &expiresInSeconds,
|
||||
OrderMerchantExternalID: tradeNo,
|
||||
})
|
||||
if err != nil {
|
||||
logger.LogError(c.Request.Context(), fmt.Sprintf("Waffo Pancake 订阅结账会话创建失败 user_id=%d plan_id=%d trade_no=%s error=%q", userId, plan.Id, tradeNo, err.Error()))
|
||||
order.Status = common.TopUpStatusFailed
|
||||
_ = order.Update()
|
||||
c.JSON(http.StatusOK, gin.H{"message": "error", "data": "拉起支付失败"})
|
||||
return
|
||||
}
|
||||
logger.LogInfo(c.Request.Context(), fmt.Sprintf("Waffo Pancake 订阅订单创建成功 user_id=%d plan_id=%d trade_no=%s session_id=%s money=%.2f", userId, plan.Id, tradeNo, session.SessionID, plan.PriceAmount))
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": "success",
|
||||
"data": gin.H{
|
||||
"checkout_url": session.CheckoutURL,
|
||||
"session_id": session.SessionID,
|
||||
"expires_at": session.ExpiresAt,
|
||||
"order_id": tradeNo,
|
||||
"token": session.Token,
|
||||
"token_expires_at": session.TokenExpiresAt,
|
||||
},
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func ListSystemInstances(c *gin.Context) {
|
||||
instances, err := model.ListSystemInstances()
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
now := common.GetTimestamp()
|
||||
responses := make([]model.SystemInstanceResponse, 0, len(instances))
|
||||
for _, instance := range instances {
|
||||
responses = append(responses, instance.ToResponse(now))
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": responses,
|
||||
})
|
||||
}
|
||||
|
||||
func DeleteStaleSystemInstances(c *gin.Context) {
|
||||
deletedCount, err := model.DeleteStaleSystemInstances(common.GetTimestamp())
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
common.ApiSuccess(c, gin.H{
|
||||
"deleted_count": deletedCount,
|
||||
})
|
||||
}
|
||||
|
||||
func DeleteStaleSystemInstance(c *gin.Context) {
|
||||
nodeName := c.Param("node_name")
|
||||
if strings.TrimSpace(nodeName) == "" {
|
||||
common.ApiErrorMsg(c, "node name is required")
|
||||
return
|
||||
}
|
||||
|
||||
deleted, err := model.DeleteStaleSystemInstance(nodeName, common.GetTimestamp())
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if !deleted {
|
||||
common.ApiErrorMsg(c, "instance is not stale or no longer exists")
|
||||
return
|
||||
}
|
||||
|
||||
common.ApiSuccess(c, gin.H{
|
||||
"deleted_count": 1,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strconv"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func CreateLogCleanupSystemTask(c *gin.Context) {
|
||||
targetTimestamp, _ := strconv.ParseInt(c.Query("target_timestamp"), 10, 64)
|
||||
if targetTimestamp == 0 {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "target timestamp is required",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
task, err := service.StartLogCleanupTask(targetTimestamp)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": task.ToResponse(),
|
||||
})
|
||||
}
|
||||
|
||||
func GetCurrentSystemTask(c *gin.Context) {
|
||||
taskType := c.Query("type")
|
||||
if taskType == "" {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "type is required",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
task, err := model.GetActiveSystemTask(taskType)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if task == nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": nil,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": task.ToResponse(),
|
||||
})
|
||||
}
|
||||
|
||||
func ListSystemTasks(c *gin.Context) {
|
||||
limit, _ := strconv.Atoi(c.Query("limit"))
|
||||
|
||||
tasks, err := model.ListSystemTasks(limit)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
responses := make([]model.SystemTaskResponse, 0, len(tasks))
|
||||
for _, task := range tasks {
|
||||
responses = append(responses, task.ToResponse())
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": responses,
|
||||
})
|
||||
}
|
||||
|
||||
func GetSystemTask(c *gin.Context) {
|
||||
taskID := c.Param("task_id")
|
||||
if taskID == "" {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "task id is required",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
task, err := model.GetSystemTaskByTaskID(taskID)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if task == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{
|
||||
"success": false,
|
||||
"message": "task not found",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": task.ToResponse(),
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,163 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
)
|
||||
|
||||
// RegisterScheduledSystemTasks wires the periodic channel test, upstream model
|
||||
// update, and async task polling (Midjourney / Suno / video) jobs into the
|
||||
// system task framework so a DB lease dedups execution across multiple master
|
||||
// instances and each run is recorded as one task row. Call this before
|
||||
// service.StartSystemTaskRunner.
|
||||
func RegisterScheduledSystemTasks() {
|
||||
service.RegisterSystemTaskHandler(channelTestHandler{})
|
||||
service.RegisterSystemTaskHandler(modelUpdateHandler{})
|
||||
service.RegisterSystemTaskHandler(midjourneyPollHandler{})
|
||||
service.RegisterSystemTaskHandler(asyncTaskPollHandler{})
|
||||
}
|
||||
|
||||
// channelTestHandler runs the scheduled "test all channels" job. Enablement and
|
||||
// cadence still come from the monitor settings; only the execution path moved
|
||||
// into the system task runner.
|
||||
type channelTestHandler struct{}
|
||||
|
||||
func (channelTestHandler) Type() string { return model.SystemTaskTypeChannelTest }
|
||||
|
||||
func (channelTestHandler) Enabled() bool {
|
||||
return operation_setting.GetMonitorSetting().AutoTestChannelEnabled
|
||||
}
|
||||
|
||||
func (channelTestHandler) Interval() time.Duration {
|
||||
minutes := operation_setting.GetMonitorSetting().AutoTestChannelMinutes
|
||||
if minutes <= 0 {
|
||||
minutes = 10
|
||||
}
|
||||
return time.Duration(minutes * float64(time.Minute))
|
||||
}
|
||||
|
||||
func (channelTestHandler) NewPayload() any { return nil }
|
||||
|
||||
// channelTestTaskPayload controls one channel_test run. A nil/empty payload is a
|
||||
// scheduled run, which uses the configured monitor ChannelTestMode and does not
|
||||
// notify. A manual "test all channels" trigger sets Mode=scheduled_all and
|
||||
// Notify=true to reproduce the legacy manual behavior (test every channel and
|
||||
// notify root on completion).
|
||||
type channelTestTaskPayload struct {
|
||||
Mode string `json:"mode,omitempty"`
|
||||
Notify bool `json:"notify,omitempty"`
|
||||
}
|
||||
|
||||
func (channelTestHandler) Run(ctx context.Context, task *model.SystemTask, runnerID string) {
|
||||
payload := channelTestTaskPayload{}
|
||||
if err := task.DecodePayload(&payload); err != nil {
|
||||
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, err)
|
||||
return
|
||||
}
|
||||
summary, err := runChannelTestTask(ctx, payload.Mode, payload.Notify, service.NewSystemTaskProgressReporter(task, runnerID))
|
||||
if err != nil {
|
||||
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, err)
|
||||
return
|
||||
}
|
||||
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusSucceeded, summary, nil)
|
||||
}
|
||||
|
||||
// modelUpdateHandler runs the scheduled upstream model update detection job.
|
||||
type modelUpdateHandler struct{}
|
||||
|
||||
func (modelUpdateHandler) Type() string { return model.SystemTaskTypeModelUpdate }
|
||||
|
||||
func (modelUpdateHandler) Enabled() bool {
|
||||
return common.GetEnvOrDefaultBool("CHANNEL_UPSTREAM_MODEL_UPDATE_TASK_ENABLED", true)
|
||||
}
|
||||
|
||||
func (modelUpdateHandler) Interval() time.Duration {
|
||||
intervalMinutes := common.GetEnvOrDefault(
|
||||
"CHANNEL_UPSTREAM_MODEL_UPDATE_TASK_INTERVAL_MINUTES",
|
||||
channelUpstreamModelUpdateTaskDefaultIntervalMinutes,
|
||||
)
|
||||
if intervalMinutes < 1 {
|
||||
intervalMinutes = channelUpstreamModelUpdateTaskDefaultIntervalMinutes
|
||||
}
|
||||
return time.Duration(intervalMinutes) * time.Minute
|
||||
}
|
||||
|
||||
func (modelUpdateHandler) NewPayload() any { return nil }
|
||||
|
||||
// modelUpdateTaskPayload controls one model_update run. A scheduled run
|
||||
// (Manual=false) respects the per-channel minimum check interval and may
|
||||
// auto-apply detected models when a channel has auto-sync enabled. A manual
|
||||
// "detect all" trigger sets Manual=true to reproduce the legacy detect-all
|
||||
// semantics: force a re-check regardless of the interval and never auto-apply,
|
||||
// so the admin reviews and applies changes explicitly.
|
||||
type modelUpdateTaskPayload struct {
|
||||
Manual bool `json:"manual,omitempty"`
|
||||
}
|
||||
|
||||
func (modelUpdateHandler) Run(ctx context.Context, task *model.SystemTask, runnerID string) {
|
||||
payload := modelUpdateTaskPayload{}
|
||||
if err := task.DecodePayload(&payload); err != nil {
|
||||
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, err)
|
||||
return
|
||||
}
|
||||
summary := runChannelUpstreamModelUpdateTaskOnce(ctx, payload.Manual, !payload.Manual, service.NewSystemTaskProgressReporter(task, runnerID))
|
||||
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusSucceeded, summary, nil)
|
||||
}
|
||||
|
||||
// midjourneyPollHandler runs one Midjourney polling pass per scheduled run.
|
||||
// Enabled() folds the "are there unfinished tasks?" check into enablement so the
|
||||
// scheduler creates no row when the system is idle; only when at least one
|
||||
// Midjourney task is in progress does a row get scheduled.
|
||||
type midjourneyPollHandler struct{}
|
||||
|
||||
func (midjourneyPollHandler) Type() string { return model.SystemTaskTypeMidjourneyPoll }
|
||||
|
||||
func (midjourneyPollHandler) Enabled() bool {
|
||||
return constant.UpdateTask && model.HasUnfinishedMidjourneyTasks()
|
||||
}
|
||||
|
||||
func (midjourneyPollHandler) Interval() time.Duration { return 15 * time.Second }
|
||||
|
||||
func (midjourneyPollHandler) NewPayload() any { return nil }
|
||||
|
||||
func (midjourneyPollHandler) Run(ctx context.Context, task *model.SystemTask, runnerID string) {
|
||||
summary := runMidjourneyTaskUpdateOnce(ctx, service.NewSystemTaskProgressReporter(task, runnerID))
|
||||
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusSucceeded, summary, nil)
|
||||
}
|
||||
|
||||
// asyncTaskPollHandler runs one async-task (Suno/video) polling pass per
|
||||
// scheduled run. Like midjourneyPollHandler, Enabled() folds in the unfinished
|
||||
// task existence check so an idle system schedules no rows.
|
||||
type asyncTaskPollHandler struct{}
|
||||
|
||||
func (asyncTaskPollHandler) Type() string { return model.SystemTaskTypeAsyncTaskPoll }
|
||||
|
||||
func (asyncTaskPollHandler) Enabled() bool {
|
||||
return constant.UpdateTask && model.HasUnfinishedSyncTasks()
|
||||
}
|
||||
|
||||
func (asyncTaskPollHandler) Interval() time.Duration { return 15 * time.Second }
|
||||
|
||||
func (asyncTaskPollHandler) NewPayload() any { return nil }
|
||||
|
||||
func (asyncTaskPollHandler) Run(ctx context.Context, task *model.SystemTask, runnerID string) {
|
||||
summary := service.RunTaskPollingOnce(ctx, service.NewSystemTaskProgressReporter(task, runnerID))
|
||||
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusSucceeded, summary, nil)
|
||||
}
|
||||
|
||||
func finishSystemTaskHandler(task *model.SystemTask, runnerID string, status model.SystemTaskStatus, result any, runErr error) {
|
||||
errorMessage := ""
|
||||
if runErr != nil {
|
||||
errorMessage = runErr.Error()
|
||||
}
|
||||
if err := model.FinishSystemTask(task.TaskID, runnerID, status, result, errorMessage); err != nil {
|
||||
common.SysLog(fmt.Sprintf("system task %s failed to persist result: %v", task.TaskID, err))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,478 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/relay"
|
||||
relaychannel "github.com/QuantumNous/new-api/relay/channel"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/types"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type taskArtifactResponse struct {
|
||||
Key string `json:"key"`
|
||||
Type string `json:"type"`
|
||||
MimeType string `json:"mime_type,omitempty"`
|
||||
ContentURL string `json:"content_url"`
|
||||
}
|
||||
|
||||
var (
|
||||
taskArtifactKeyPattern = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._~-]{0,127}$`)
|
||||
errTaskArtifactPluginUnavailable = errors.New("task artifact plugin unavailable")
|
||||
errTaskArtifactPlugin = errors.New("task artifact plugin error")
|
||||
)
|
||||
|
||||
func GetTask(c *gin.Context) {
|
||||
task, exists, err := model.GetByTaskId(c.GetInt("id"), c.Param("key"))
|
||||
if err != nil {
|
||||
videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to query task")
|
||||
return
|
||||
}
|
||||
if !exists {
|
||||
videoProxyError(c, http.StatusNotFound, "invalid_request_error", "Task not found")
|
||||
return
|
||||
}
|
||||
createdAt := task.CreatedAt
|
||||
if createdAt == 0 {
|
||||
createdAt = task.SubmitTime
|
||||
}
|
||||
failReason := task.FailReason
|
||||
if task.Status == model.TaskStatusSuccess && taskFailReasonIsLegacyResultURL(task.FailReason) {
|
||||
failReason = ""
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"task_id": task.TaskID,
|
||||
"platform": task.Platform,
|
||||
"status": task.Status,
|
||||
"progress": task.Progress,
|
||||
"fail_reason": failReason,
|
||||
"created_at": createdAt,
|
||||
"finished_at": task.FinishTime,
|
||||
})
|
||||
}
|
||||
|
||||
func GetTaskArtifacts(c *gin.Context) {
|
||||
task, exists, err := model.GetByTaskId(c.GetInt("id"), c.Param("key"))
|
||||
if err != nil {
|
||||
writeTaskArtifactError(c, http.StatusInternalServerError, "artifact_internal_error", "Failed to query task")
|
||||
return
|
||||
}
|
||||
if !exists || task == nil {
|
||||
writeTaskArtifactError(c, http.StatusNotFound, "artifact_not_found", "Task or artifact not found")
|
||||
return
|
||||
}
|
||||
writeTaskArtifacts(c, task, false)
|
||||
}
|
||||
|
||||
func GetDashboardTaskArtifacts(c *gin.Context) {
|
||||
task, exists, err := getTaskForArtifactRequest(c, c.Param("task_id"))
|
||||
if err != nil {
|
||||
writeTaskArtifactError(c, http.StatusInternalServerError, "artifact_internal_error", "Failed to query task")
|
||||
return
|
||||
}
|
||||
if !exists || task == nil {
|
||||
writeTaskArtifactError(c, http.StatusNotFound, "artifact_not_found", "Task or artifact not found")
|
||||
return
|
||||
}
|
||||
writeTaskArtifacts(c, task, true)
|
||||
}
|
||||
|
||||
func writeTaskArtifacts(c *gin.Context, task *model.Task, dashboard bool) {
|
||||
c.Header("Cache-Control", "private, no-store")
|
||||
artifacts, err := projectTaskArtifacts(task)
|
||||
if err != nil {
|
||||
writeTaskArtifactProjectionError(c, err)
|
||||
return
|
||||
}
|
||||
items := make([]taskArtifactResponse, 0, len(artifacts))
|
||||
for _, artifact := range artifacts {
|
||||
contentURL, buildErr := service.BuildTaskArtifactContentURL(task.TaskID, artifact.Key)
|
||||
if buildErr != nil {
|
||||
writeTaskArtifactError(c, http.StatusInternalServerError, "artifact_url_error", "Failed to build artifact content URL")
|
||||
return
|
||||
}
|
||||
items = append(items, taskArtifactResponse{
|
||||
Key: artifact.Key,
|
||||
Type: artifact.Type,
|
||||
MimeType: artifact.MimeType,
|
||||
ContentURL: contentURL,
|
||||
})
|
||||
}
|
||||
response := gin.H{"task_id": task.TaskID, "artifacts": items}
|
||||
if legacyVideoAvailable(task) {
|
||||
legacyContentURL, buildErr := service.BuildTaskArtifactContentURL(task.TaskID, "video")
|
||||
if buildErr != nil {
|
||||
writeTaskArtifactError(c, http.StatusInternalServerError, "artifact_url_error", "Failed to build artifact content URL")
|
||||
return
|
||||
}
|
||||
response["legacy_content_url"] = legacyContentURL
|
||||
}
|
||||
if dashboard {
|
||||
common.ApiSuccess(c, response)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, response)
|
||||
}
|
||||
|
||||
func projectTaskArtifacts(task *model.Task) ([]relaychannel.TaskArtifact, error) {
|
||||
if task == nil || task.Status != model.TaskStatusSuccess || !taskHasPluginExecution(task) {
|
||||
return []relaychannel.TaskArtifact{}, nil
|
||||
}
|
||||
adaptor := relay.GetTaskAdaptor(task.Platform)
|
||||
if adaptor == nil {
|
||||
return nil, errTaskArtifactPluginUnavailable
|
||||
}
|
||||
provider, ok := adaptor.(relaychannel.TaskArtifactProvider)
|
||||
if !ok {
|
||||
return []relaychannel.TaskArtifact{}, nil
|
||||
}
|
||||
artifacts, err := provider.ListArtifacts(task)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%w: %v", errTaskArtifactPlugin, err)
|
||||
}
|
||||
return validateProjectedTaskArtifacts(artifacts)
|
||||
}
|
||||
|
||||
func validateProjectedTaskArtifacts(artifacts []relaychannel.TaskArtifact) ([]relaychannel.TaskArtifact, error) {
|
||||
if len(artifacts) > 64 {
|
||||
return nil, fmt.Errorf("%w: too many artifacts", errTaskArtifactPlugin)
|
||||
}
|
||||
seen := make(map[string]struct{}, len(artifacts))
|
||||
for i := range artifacts {
|
||||
if artifacts[i].Key != strings.TrimSpace(artifacts[i].Key) ||
|
||||
artifacts[i].Type != strings.TrimSpace(artifacts[i].Type) {
|
||||
return nil, fmt.Errorf("%w: invalid artifact identity", errTaskArtifactPlugin)
|
||||
}
|
||||
if !taskArtifactKeyPattern.MatchString(artifacts[i].Key) {
|
||||
return nil, fmt.Errorf("%w: invalid artifact key", errTaskArtifactPlugin)
|
||||
}
|
||||
if _, exists := seen[artifacts[i].Key]; exists {
|
||||
return nil, fmt.Errorf("%w: duplicate artifact key", errTaskArtifactPlugin)
|
||||
}
|
||||
seen[artifacts[i].Key] = struct{}{}
|
||||
switch artifacts[i].Type {
|
||||
case "video", "audio", "image", "file":
|
||||
default:
|
||||
return nil, fmt.Errorf("%w: invalid artifact type", errTaskArtifactPlugin)
|
||||
}
|
||||
if len(artifacts[i].MimeType) > 255 || strings.ContainsAny(artifacts[i].MimeType, "\r\n") {
|
||||
return nil, fmt.Errorf("%w: invalid artifact mime type", errTaskArtifactPlugin)
|
||||
}
|
||||
}
|
||||
return artifacts, nil
|
||||
}
|
||||
|
||||
func initTaskArtifactAdaptor(task *model.Task) (relaychannel.TaskAdaptor, error) {
|
||||
if task == nil || !taskHasPluginExecution(task) {
|
||||
return nil, errTaskArtifactPluginUnavailable
|
||||
}
|
||||
channelModel, err := model.CacheGetChannel(task.ChannelId)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%w: channel unavailable", errTaskArtifactPluginUnavailable)
|
||||
}
|
||||
adaptor := relay.GetTaskAdaptor(task.Platform)
|
||||
if adaptor == nil {
|
||||
return nil, errTaskArtifactPluginUnavailable
|
||||
}
|
||||
pluginKey := task.PrivateData.Key
|
||||
if pluginKey == "" {
|
||||
pluginKey = channelModel.Key
|
||||
}
|
||||
baseURL := channelModel.GetBaseURL()
|
||||
if baseURL == "" {
|
||||
baseURL = constant.GetChannelBaseURL(channelModel.Type)
|
||||
}
|
||||
adaptor.Init(&relaycommon.RelayInfo{
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
ChannelType: channelModel.Type,
|
||||
ChannelBaseUrl: baseURL,
|
||||
ApiKey: pluginKey,
|
||||
ChannelSetting: channelModel.GetSetting(),
|
||||
},
|
||||
})
|
||||
return adaptor, nil
|
||||
}
|
||||
|
||||
func taskHasPluginExecution(task *model.Task) bool {
|
||||
return task != nil &&
|
||||
task.PrivateData.Execution != nil &&
|
||||
task.PrivateData.Execution.TaskPlugin != nil &&
|
||||
strings.TrimSpace(task.PrivateData.Execution.TaskPlugin.Key) != ""
|
||||
}
|
||||
|
||||
func legacyVideoAvailable(task *model.Task) bool {
|
||||
if task == nil || task.Status != model.TaskStatusSuccess ||
|
||||
taskHasPluginExecution(task) || task.Platform == constant.TaskPlatformSuno ||
|
||||
strings.TrimSpace(task.GetResultURL()) == "" {
|
||||
return false
|
||||
}
|
||||
switch constant.NormalizeTaskAction(task.Action) {
|
||||
case constant.TaskActionImageToVideo,
|
||||
constant.TaskActionTextToVideo,
|
||||
constant.TaskActionFirstTailToVideo,
|
||||
constant.TaskActionReferenceToVideo,
|
||||
constant.TaskActionRemix:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func getTaskForArtifactRequest(c *gin.Context, taskID string) (*model.Task, bool, error) {
|
||||
if middleware.IsTaskArtifactAccess(c) {
|
||||
task, exists, err := model.GetUniqueByOnlyTaskId(taskID)
|
||||
if err != nil || !exists || task == nil {
|
||||
return task, exists, err
|
||||
}
|
||||
owner, err := model.GetUserCache(task.UserId)
|
||||
if err != nil || owner == nil || owner.Status != common.UserStatusEnabled {
|
||||
return nil, false, err
|
||||
}
|
||||
return task, true, nil
|
||||
}
|
||||
if c.GetInt("token_id") == 0 && c.GetInt("role") >= common.RoleAdminUser {
|
||||
return model.GetByOnlyTaskId(taskID)
|
||||
}
|
||||
return model.GetByTaskId(c.GetInt("id"), taskID)
|
||||
}
|
||||
|
||||
func writeTaskArtifactProjectionError(c *gin.Context, err error) {
|
||||
if errors.Is(err, errTaskArtifactPluginUnavailable) {
|
||||
writeTaskArtifactError(c, http.StatusServiceUnavailable, "artifact_plugin_unavailable", "Artifact preview plugin is unavailable")
|
||||
return
|
||||
}
|
||||
writeTaskArtifactError(c, http.StatusInternalServerError, "artifact_plugin_error", "Artifact preview plugin failed")
|
||||
}
|
||||
|
||||
func writeTaskArtifactError(c *gin.Context, status int, code, message string) {
|
||||
c.Header("Cache-Control", "private, no-store")
|
||||
if middleware.IsTaskArtifactAccess(c) {
|
||||
status = http.StatusNotFound
|
||||
code = "artifact_not_found"
|
||||
message = "Task or artifact not found"
|
||||
}
|
||||
if strings.HasPrefix(c.Request.URL.Path, "/api/") {
|
||||
c.JSON(status, gin.H{"success": false, "code": code, "message": message})
|
||||
return
|
||||
}
|
||||
c.JSON(status, gin.H{
|
||||
"error": gin.H{
|
||||
"message": message,
|
||||
"type": code,
|
||||
"code": code,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func TaskArtifactContent(c *gin.Context) {
|
||||
task, exists, err := getTaskForArtifactRequest(c, c.Param("key"))
|
||||
if err != nil {
|
||||
writeTaskArtifactError(c, http.StatusInternalServerError, "artifact_internal_error", "Failed to query task")
|
||||
return
|
||||
}
|
||||
if !exists || task == nil {
|
||||
writeTaskArtifactError(c, http.StatusNotFound, "artifact_not_found", "Task or artifact not found")
|
||||
return
|
||||
}
|
||||
artifactKey := strings.TrimSpace(c.Param("artifact_key"))
|
||||
if !taskArtifactKeyPattern.MatchString(artifactKey) {
|
||||
writeTaskArtifactError(c, http.StatusNotFound, "artifact_not_found", "Task or artifact not found")
|
||||
return
|
||||
}
|
||||
if task.Status != model.TaskStatusSuccess {
|
||||
writeTaskArtifactError(c, http.StatusConflict, "artifact_not_ready", "Task artifacts are not ready")
|
||||
return
|
||||
}
|
||||
if !taskHasPluginExecution(task) {
|
||||
if artifactKey != "video" || !legacyVideoAvailable(task) {
|
||||
writeTaskArtifactError(c, http.StatusNotFound, "artifact_not_found", "Task or artifact not found")
|
||||
return
|
||||
}
|
||||
descriptor := &relaychannel.TaskContentRequest{
|
||||
URL: task.GetResultURL(),
|
||||
Method: c.Request.Method,
|
||||
Credentialless: true,
|
||||
}
|
||||
if err := proxyTaskMedia(c, task, descriptor); err != nil {
|
||||
writeTaskMediaProxyError(c, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
artifacts, err := projectTaskArtifacts(task)
|
||||
if err != nil {
|
||||
writeTaskArtifactProjectionError(c, err)
|
||||
return
|
||||
}
|
||||
found := false
|
||||
for _, artifact := range artifacts {
|
||||
if artifact.Key == artifactKey {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
writeTaskArtifactError(c, http.StatusNotFound, "artifact_not_found", "Task or artifact not found")
|
||||
return
|
||||
}
|
||||
artifactStore := service.GetTaskArtifactStore()
|
||||
if ref, resolveErr := artifactStore.Resolve(task, artifactKey); resolveErr == nil && ref != nil {
|
||||
_ = artifactStore.Serve(c, task, ref)
|
||||
return
|
||||
}
|
||||
|
||||
adaptor, err := initTaskArtifactAdaptor(task)
|
||||
if err != nil {
|
||||
writeTaskArtifactProjectionError(c, err)
|
||||
return
|
||||
}
|
||||
provider, ok := adaptor.(relaychannel.TaskContentRequestProvider)
|
||||
if !ok {
|
||||
writeTaskArtifactError(c, http.StatusServiceUnavailable, "artifact_plugin_unavailable", "Artifact content plugin is unavailable")
|
||||
return
|
||||
}
|
||||
clientRequest := relaychannel.TaskArtifactClientRequest{
|
||||
Method: c.Request.Method,
|
||||
Headers: taskArtifactClientHeaders(c.Request.Header),
|
||||
}
|
||||
descriptor, err := provider.BuildContentRequest(task, artifactKey, clientRequest)
|
||||
if err != nil || descriptor == nil {
|
||||
writeTaskArtifactError(c, http.StatusInternalServerError, "artifact_plugin_error", "Artifact content plugin failed")
|
||||
return
|
||||
}
|
||||
if err := proxyTaskMedia(c, task, descriptor); err != nil {
|
||||
writeTaskMediaProxyError(c, err)
|
||||
}
|
||||
}
|
||||
|
||||
func taskArtifactClientHeaders(headers http.Header) map[string]string {
|
||||
result := make(map[string]string, 4)
|
||||
for _, name := range []string{"Range", "If-Range", "If-None-Match", "If-Modified-Since"} {
|
||||
if value := strings.TrimSpace(headers.Get(name)); value != "" {
|
||||
result[name] = value
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
/*
|
||||
The task list handlers below deliberately do not call projectTaskArtifacts.
|
||||
Artifact projection is confined to the explicit endpoints above.
|
||||
*/
|
||||
|
||||
func GetAllTask(c *gin.Context) {
|
||||
pageInfo := common.GetPageQuery(c)
|
||||
startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64)
|
||||
endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64)
|
||||
queryParams := model.SyncTaskQueryParams{Platform: constant.TaskPlatform(c.Query("platform")), TaskID: c.Query("task_id"), Status: c.Query("status"), Action: c.Query("action"), StartTimestamp: startTimestamp, EndTimestamp: endTimestamp, ChannelID: c.Query("channel_id")}
|
||||
items := model.TaskGetAllTasks(pageInfo.GetStartIdx(), pageInfo.GetPageSize(), queryParams)
|
||||
pageInfo.SetTotal(int(model.TaskCountAllTasks(queryParams)))
|
||||
pageInfo.SetItems(tasksToDto(items, true, c.GetInt("role")))
|
||||
common.ApiSuccess(c, pageInfo)
|
||||
}
|
||||
|
||||
func GetUserTask(c *gin.Context) {
|
||||
pageInfo := common.GetPageQuery(c)
|
||||
userID := c.GetInt("id")
|
||||
startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64)
|
||||
endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64)
|
||||
queryParams := model.SyncTaskQueryParams{Platform: constant.TaskPlatform(c.Query("platform")), TaskID: c.Query("task_id"), Status: c.Query("status"), Action: c.Query("action"), StartTimestamp: startTimestamp, EndTimestamp: endTimestamp}
|
||||
items := model.TaskGetAllUserTask(userID, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), queryParams)
|
||||
pageInfo.SetTotal(int(model.TaskCountAllUserTask(userID, queryParams)))
|
||||
pageInfo.SetItems(tasksToDto(items, false, common.RoleCommonUser))
|
||||
common.ApiSuccess(c, pageInfo)
|
||||
}
|
||||
|
||||
func tasksToDto(tasks []*model.Task, fillUser bool, viewerRole int) []*dto.TaskDto {
|
||||
var userIDMap map[int]*model.UserBase
|
||||
if fillUser {
|
||||
userIDMap = make(map[int]*model.UserBase)
|
||||
userIDs := types.NewSet[int]()
|
||||
for _, task := range tasks {
|
||||
userIDs.Add(task.UserId)
|
||||
}
|
||||
for _, userID := range userIDs.Items() {
|
||||
if cacheUser, err := model.GetUserCache(userID); err == nil {
|
||||
userIDMap[userID] = cacheUser
|
||||
}
|
||||
}
|
||||
}
|
||||
result := make([]*dto.TaskDto, len(tasks))
|
||||
for i, task := range tasks {
|
||||
if fillUser {
|
||||
if user, ok := userIDMap[task.UserId]; ok {
|
||||
task.Username = user.Username
|
||||
}
|
||||
}
|
||||
item := relay.TaskModel2Dto(task)
|
||||
item.LegacyVideoAvailable = legacyVideoAvailable(task)
|
||||
if task.Status == model.TaskStatusSuccess {
|
||||
item.ResultURL = ""
|
||||
if taskFailReasonIsLegacyResultURL(task.FailReason) {
|
||||
item.FailReason = ""
|
||||
}
|
||||
}
|
||||
if viewerRole >= common.RoleAdminUser {
|
||||
adminInfo := &dto.TaskAdminInfo{}
|
||||
if execution := task.PrivateData.Execution; execution != nil {
|
||||
adminInfo.RequestID = execution.RequestID
|
||||
adminInfo.RequestPath = execution.RequestPath
|
||||
if snapshot := execution.TaskPlugin; snapshot != nil {
|
||||
adminInfo.TaskPlugin = &dto.TaskPluginInfo{
|
||||
Key: snapshot.Key,
|
||||
Name: snapshot.Name,
|
||||
Version: snapshot.Version,
|
||||
}
|
||||
if snapshot.Author != nil {
|
||||
adminInfo.TaskPlugin.Author = &dto.TaskPluginAuthorInfo{
|
||||
Name: snapshot.Author.Name,
|
||||
URL: snapshot.Author.URL,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if adminInfo.RequestID != "" || adminInfo.RequestPath != "" || adminInfo.TaskPlugin != nil {
|
||||
item.AdminInfo = adminInfo
|
||||
}
|
||||
}
|
||||
if viewerRole >= common.RoleRootUser {
|
||||
rootInfo := &dto.TaskRootInfo{
|
||||
UpstreamTaskID: task.PrivateData.UpstreamTaskID,
|
||||
NodeName: task.PrivateData.NodeName,
|
||||
}
|
||||
if execution := task.PrivateData.Execution; execution != nil {
|
||||
if snapshot := execution.TaskPlugin; snapshot != nil {
|
||||
rootInfo.TaskPlugin = &dto.TaskPluginRuntimeInfo{
|
||||
Key: snapshot.Key,
|
||||
Version: snapshot.Version,
|
||||
APIVersion: snapshot.APIVersion,
|
||||
Generation: snapshot.Generation,
|
||||
}
|
||||
}
|
||||
}
|
||||
if rootInfo.TaskPlugin != nil || rootInfo.UpstreamTaskID != "" || rootInfo.NodeName != "" {
|
||||
item.RootInfo = rootInfo
|
||||
}
|
||||
}
|
||||
result[i] = item
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func taskFailReasonIsLegacyResultURL(value string) bool {
|
||||
value = strings.TrimSpace(value)
|
||||
return len(value) >= len("https://") && strings.EqualFold(value[:len("https://")], "https://") ||
|
||||
len(value) >= len("http://") && strings.EqualFold(value[:len("http://")], "http://") ||
|
||||
len(value) >= len("data:") && strings.EqualFold(value[:len("data:")], "data:")
|
||||
}
|
||||
@@ -0,0 +1,590 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
relaychannel "github.com/QuantumNous/new-api/relay/channel"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func setupGenericTaskTest(t *testing.T) *model.Task {
|
||||
t.Helper()
|
||||
originalDB := model.DB
|
||||
previousRedisEnabled := common.RedisEnabled
|
||||
common.RedisEnabled = false
|
||||
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, database.AutoMigrate(&model.Task{}, &model.Channel{}, &model.User{}))
|
||||
model.DB = database
|
||||
t.Cleanup(func() {
|
||||
model.DB = originalDB
|
||||
common.RedisEnabled = previousRedisEnabled
|
||||
})
|
||||
|
||||
require.NoError(t, database.Create(&model.User{
|
||||
Id: 7, Username: "artifact-owner", Status: common.UserStatusEnabled,
|
||||
Role: common.RoleCommonUser, Group: "default",
|
||||
}).Error)
|
||||
baseURL := "https://example.com"
|
||||
require.NoError(t, database.Create(&model.Channel{
|
||||
Id: 1, Name: "artifact", Key: "key", BaseURL: &baseURL, Status: common.ChannelStatusEnabled,
|
||||
}).Error)
|
||||
task := &model.Task{
|
||||
TaskID: "task_generic", Platform: "document", UserId: 7, ChannelId: 1,
|
||||
Status: model.TaskStatusSuccess, Progress: "100%", SubmitTime: 10, FinishTime: 20,
|
||||
}
|
||||
require.NoError(t, database.Create(task).Error)
|
||||
return task
|
||||
}
|
||||
|
||||
func allowPrivateTaskMediaTest(t *testing.T) {
|
||||
t.Helper()
|
||||
originalFetchSetting := *system_setting.GetFetchSetting()
|
||||
system_setting.GetFetchSetting().EnableSSRFProtection = true
|
||||
system_setting.GetFetchSetting().AllowPrivateIp = true
|
||||
system_setting.GetFetchSetting().AllowedPorts = []string{"1-65535"}
|
||||
t.Cleanup(func() { *system_setting.GetFetchSetting() = originalFetchSetting })
|
||||
service.InitHttpClient()
|
||||
}
|
||||
|
||||
func TestGetTaskDoesNotProjectArtifacts(t *testing.T) {
|
||||
task := setupGenericTaskTest(t)
|
||||
task.FailReason = "https://stale-upstream.invalid/video.mp4"
|
||||
task.PrivateData = model.TaskPrivateData{
|
||||
ResultURL: "https://private-upstream.invalid/video.mp4",
|
||||
Execution: &model.TaskExecutionSnapshot{
|
||||
TaskPlugin: &model.TaskPluginSnapshot{Key: "missing-plugin", Name: "Missing"},
|
||||
},
|
||||
}
|
||||
require.NoError(t, model.DB.Save(task).Error)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Set("id", 7)
|
||||
c.Params = gin.Params{{Key: "key", Value: task.TaskID}}
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/tasks/"+task.TaskID, nil)
|
||||
|
||||
GetTask(c)
|
||||
|
||||
assert.Equal(t, http.StatusOK, recorder.Code)
|
||||
var response map[string]any
|
||||
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
|
||||
assert.Equal(t, task.TaskID, response["task_id"])
|
||||
assert.NotContains(t, response, "artifacts")
|
||||
assert.NotContains(t, recorder.Body.String(), "upstream.invalid")
|
||||
}
|
||||
|
||||
func TestGetTaskArtifactsReturnsEmptyForLegacyTask(t *testing.T) {
|
||||
task := setupGenericTaskTest(t)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Set("id", task.UserId)
|
||||
c.Params = gin.Params{{Key: "key", Value: task.TaskID}}
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/tasks/"+task.TaskID+"/artifacts", nil)
|
||||
|
||||
GetTaskArtifacts(c)
|
||||
|
||||
assert.Equal(t, http.StatusOK, recorder.Code)
|
||||
assert.Equal(t, "private, no-store", recorder.Header().Get("Cache-Control"))
|
||||
var response struct {
|
||||
TaskID string `json:"task_id"`
|
||||
Artifacts []taskArtifactResponse `json:"artifacts"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
|
||||
assert.Equal(t, task.TaskID, response.TaskID)
|
||||
assert.Empty(t, response.Artifacts)
|
||||
}
|
||||
|
||||
func TestTaskArtifactAuthorizationKeepsForeignTasksHidden(t *testing.T) {
|
||||
task := setupGenericTaskTest(t)
|
||||
|
||||
commonUser, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
commonUser.Set("id", 8)
|
||||
commonUser.Set("role", common.RoleCommonUser)
|
||||
_, exists, err := getTaskForArtifactRequest(commonUser, task.TaskID)
|
||||
require.NoError(t, err)
|
||||
assert.False(t, exists)
|
||||
|
||||
admin, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
admin.Set("id", 8)
|
||||
admin.Set("role", common.RoleAdminUser)
|
||||
found, exists, err := getTaskForArtifactRequest(admin, task.TaskID)
|
||||
require.NoError(t, err)
|
||||
require.True(t, exists)
|
||||
assert.Equal(t, task.TaskID, found.TaskID)
|
||||
|
||||
apiToken, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
apiToken.Set("id", 8)
|
||||
apiToken.Set("role", common.RoleRootUser)
|
||||
apiToken.Set("token_id", 99)
|
||||
_, exists, err = getTaskForArtifactRequest(apiToken, task.TaskID)
|
||||
require.NoError(t, err)
|
||||
assert.False(t, exists)
|
||||
}
|
||||
|
||||
func TestDashboardTaskArtifactsReturnsLegacyCapabilityWithoutUpstreamURL(t *testing.T) {
|
||||
task := setupGenericTaskTest(t)
|
||||
previousSecret := common.CryptoSecret
|
||||
previousPublicAddress := system_setting.TaskPublicAddress
|
||||
common.CryptoSecret = "controller-task-artifact-access-secret"
|
||||
system_setting.TaskPublicAddress = "https://gateway.example/prefix"
|
||||
t.Cleanup(func() {
|
||||
common.CryptoSecret = previousSecret
|
||||
system_setting.TaskPublicAddress = previousPublicAddress
|
||||
})
|
||||
task.Action = constant.TaskActionTextToVideo
|
||||
task.FailReason = "https://upstream.invalid/private-video.mp4?signature=secret"
|
||||
require.NoError(t, model.DB.Save(task).Error)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Set("id", task.UserId)
|
||||
c.Set("role", common.RoleCommonUser)
|
||||
c.Params = gin.Params{{Key: "task_id", Value: task.TaskID}}
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/api/task/"+task.TaskID+"/artifacts", nil)
|
||||
|
||||
GetDashboardTaskArtifacts(c)
|
||||
|
||||
assert.Equal(t, http.StatusOK, recorder.Code)
|
||||
assert.Equal(t, "private, no-store", recorder.Header().Get("Cache-Control"))
|
||||
var response struct {
|
||||
Success bool `json:"success"`
|
||||
Data struct {
|
||||
Artifacts []taskArtifactResponse `json:"artifacts"`
|
||||
LegacyContentURL string `json:"legacy_content_url"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
|
||||
assert.True(t, response.Success)
|
||||
assert.Empty(t, response.Data.Artifacts)
|
||||
contentURL, err := url.Parse(response.Data.LegacyContentURL)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "/prefix/v1/tasks/"+task.TaskID+"/artifacts/video/content", contentURL.Path)
|
||||
assert.True(t, service.VerifyTaskArtifactAccess(
|
||||
contentURL.Query().Get(service.TaskArtifactAccessQueryParameter),
|
||||
task.TaskID,
|
||||
"video",
|
||||
))
|
||||
assert.NotContains(t, recorder.Body.String(), "upstream.invalid")
|
||||
assert.NotContains(t, recorder.Body.String(), "signature=secret")
|
||||
}
|
||||
|
||||
func TestTaskArtifactAccessRequiresActiveOwner(t *testing.T) {
|
||||
task := setupGenericTaskTest(t)
|
||||
task.Action = constant.TaskActionTextToVideo
|
||||
task.FailReason = "https://upstream.invalid/private-video.mp4"
|
||||
require.NoError(t, model.DB.Save(task).Error)
|
||||
require.NoError(t, model.DB.Model(&model.User{}).
|
||||
Where("id = ?", task.UserId).
|
||||
Update("status", common.UserStatusDisabled).Error)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Set(middleware.TaskArtifactAccessContextKey, true)
|
||||
c.Params = gin.Params{
|
||||
{Key: "key", Value: task.TaskID},
|
||||
{Key: "artifact_key", Value: "video"},
|
||||
}
|
||||
c.Request = httptest.NewRequest(
|
||||
http.MethodGet,
|
||||
"/v1/tasks/"+task.TaskID+"/artifacts/video/content",
|
||||
nil,
|
||||
)
|
||||
|
||||
TaskArtifactContent(c)
|
||||
|
||||
assert.Equal(t, http.StatusNotFound, recorder.Code)
|
||||
assert.Equal(t, "private, no-store", recorder.Header().Get("Cache-Control"))
|
||||
}
|
||||
|
||||
func TestTaskArtifactAccessRejectsAmbiguousHistoricalTaskID(t *testing.T) {
|
||||
task := setupGenericTaskTest(t)
|
||||
task.Action = constant.TaskActionTextToVideo
|
||||
task.FailReason = "https://first-upstream.invalid/video.mp4"
|
||||
require.NoError(t, model.DB.Save(task).Error)
|
||||
require.NoError(t, model.DB.Create(&model.User{
|
||||
Id: 8, Username: "other-artifact-owner", Status: common.UserStatusEnabled,
|
||||
Role: common.RoleCommonUser, Group: "default", AffCode: "artifact-owner-8",
|
||||
}).Error)
|
||||
require.NoError(t, model.DB.Create(&model.Task{
|
||||
TaskID: task.TaskID, Platform: task.Platform, UserId: 8, ChannelId: task.ChannelId,
|
||||
Action: constant.TaskActionTextToVideo, Status: model.TaskStatusSuccess,
|
||||
FailReason: "https://second-upstream.invalid/video.mp4",
|
||||
}).Error)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Set(middleware.TaskArtifactAccessContextKey, true)
|
||||
c.Params = gin.Params{
|
||||
{Key: "key", Value: task.TaskID},
|
||||
{Key: "artifact_key", Value: "video"},
|
||||
}
|
||||
c.Request = httptest.NewRequest(
|
||||
http.MethodGet,
|
||||
"/v1/tasks/"+task.TaskID+"/artifacts/video/content",
|
||||
nil,
|
||||
)
|
||||
|
||||
TaskArtifactContent(c)
|
||||
|
||||
assert.Equal(t, http.StatusNotFound, recorder.Code)
|
||||
assert.NotContains(t, recorder.Body.String(), "upstream.invalid")
|
||||
}
|
||||
|
||||
func TestLegacyVideoArtifactContentUsesGetResultURL(t *testing.T) {
|
||||
task := setupGenericTaskTest(t)
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
assert.Equal(t, "bytes=0-3", r.Header.Get("Range"))
|
||||
w.Header().Set("Content-Type", "video/mp4")
|
||||
w.Header().Set("Content-Range", "bytes 0-3/4")
|
||||
w.WriteHeader(http.StatusPartialContent)
|
||||
_, _ = w.Write([]byte("data"))
|
||||
}))
|
||||
defer upstream.Close()
|
||||
allowPrivateTaskMediaTest(t)
|
||||
|
||||
task.Action = constant.TaskActionTextToVideo
|
||||
task.PrivateData.ResultURL = upstream.URL
|
||||
task.FailReason = "https://stale.invalid/legacy-fallback.mp4"
|
||||
require.NoError(t, model.DB.Save(task).Error)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Set(middleware.TaskArtifactAccessContextKey, true)
|
||||
c.Params = gin.Params{
|
||||
{Key: "key", Value: task.TaskID},
|
||||
{Key: "artifact_key", Value: "video"},
|
||||
}
|
||||
c.Request = httptest.NewRequest(
|
||||
http.MethodGet,
|
||||
"/v1/tasks/"+task.TaskID+"/artifacts/video/content",
|
||||
nil,
|
||||
)
|
||||
c.Request.Header.Set("Range", "bytes=0-3")
|
||||
|
||||
TaskArtifactContent(c)
|
||||
|
||||
assert.Equal(t, http.StatusPartialContent, recorder.Code)
|
||||
assert.Equal(t, "data", recorder.Body.String())
|
||||
assert.Equal(t, "bytes 0-3/4", recorder.Header().Get("Content-Range"))
|
||||
}
|
||||
|
||||
func TestDisabledArtifactStorePreservesPluginUpstreamContent(t *testing.T) {
|
||||
task := setupGenericTaskTest(t)
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
assert.Equal(t, "provider-key", r.Header.Get("x-goog-api-key"))
|
||||
assert.Equal(t, "bytes=0-13", r.Header.Get("Range"))
|
||||
w.Header().Set("Content-Type", "video/mp4")
|
||||
w.Header().Set("Content-Range", "bytes 0-13/14")
|
||||
w.WriteHeader(http.StatusPartialContent)
|
||||
_, _ = w.Write([]byte("artifact-bytes"))
|
||||
}))
|
||||
defer upstream.Close()
|
||||
allowPrivateTaskMediaTest(t)
|
||||
previousMemoryCache := common.MemoryCacheEnabled
|
||||
common.MemoryCacheEnabled = false
|
||||
t.Cleanup(func() { common.MemoryCacheEnabled = previousMemoryCache })
|
||||
|
||||
require.NoError(t, model.DB.Model(&model.Channel{}).Where("id = ?", task.ChannelId).Updates(map[string]any{
|
||||
"type": constant.ChannelTypeGemini,
|
||||
"key": "provider-key",
|
||||
"base_url": upstream.URL,
|
||||
}).Error)
|
||||
task.Platform = constant.TaskPlatform("google")
|
||||
task.PrivateData.Execution = &model.TaskExecutionSnapshot{TaskPlugin: &model.TaskPluginSnapshot{
|
||||
Key: "google", Name: "Google Veo (Gemini API)", Version: "1.0.0", APIVersion: 1,
|
||||
}}
|
||||
task.SetData(map[string]any{"response": map[string]any{
|
||||
"generateVideoResponse": map[string]any{
|
||||
"generatedVideos": []any{map[string]any{"video": map[string]any{"uri": upstream.URL}}},
|
||||
},
|
||||
}})
|
||||
require.NoError(t, model.DB.Save(task).Error)
|
||||
require.False(t, service.GetTaskArtifactStore().Enabled())
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Set(middleware.TaskArtifactAccessContextKey, true)
|
||||
c.Params = gin.Params{
|
||||
{Key: "key", Value: task.TaskID},
|
||||
{Key: "artifact_key", Value: "video"},
|
||||
}
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/tasks/"+task.TaskID+"/artifacts/video/content", nil)
|
||||
c.Request.Header.Set("Range", "bytes=0-13")
|
||||
|
||||
TaskArtifactContent(c)
|
||||
|
||||
assert.Equal(t, http.StatusPartialContent, recorder.Code)
|
||||
assert.Equal(t, "artifact-bytes", recorder.Body.String())
|
||||
assert.Equal(t, "video/mp4", recorder.Header().Get("Content-Type"))
|
||||
assert.Equal(t, "bytes 0-13/14", recorder.Header().Get("Content-Range"))
|
||||
}
|
||||
|
||||
func TestProjectedTaskArtifactValidationRejectsAmbiguousIdentity(t *testing.T) {
|
||||
validated, err := validateProjectedTaskArtifacts([]relaychannel.TaskArtifact{
|
||||
{Key: "video-main", Type: "video", MimeType: "video/mp4"},
|
||||
{Key: "cover.main", Type: "image", MimeType: "image/png"},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, validated, 2)
|
||||
assert.Equal(t, "video-main", validated[0].Key)
|
||||
|
||||
for _, artifacts := range [][]relaychannel.TaskArtifact{
|
||||
{{Key: "../video", Type: "video"}},
|
||||
{{Key: "video/0", Type: "video"}},
|
||||
{{Key: "video-main", Type: "video"}, {Key: "video-main", Type: "image"}},
|
||||
{{Key: "video-main", Type: "unknown"}},
|
||||
{{Key: "video-main", Type: "video", MimeType: "video/mp4\r\nX-Test: injected"}},
|
||||
} {
|
||||
_, err := validateProjectedTaskArtifacts(artifacts)
|
||||
assert.ErrorIs(t, err, errTaskArtifactPlugin)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProxyTaskMediaForwardsRangeAndFiltersResponseHeaders(t *testing.T) {
|
||||
task := setupGenericTaskTest(t)
|
||||
var receivedRange, receivedAuthorization string
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
receivedRange = r.Header.Get("Range")
|
||||
receivedAuthorization = r.Header.Get("Authorization")
|
||||
w.Header().Set("Content-Type", "video/mp4")
|
||||
w.Header().Set("Content-Range", "bytes 0-3/10")
|
||||
w.Header().Set("Accept-Ranges", "bytes")
|
||||
w.Header().Set("Set-Cookie", "provider=secret")
|
||||
w.Header().Set("WWW-Authenticate", "Bearer provider")
|
||||
w.Header().Set("X-Provider-Secret", "hidden")
|
||||
w.Header().Set("Cache-Control", "public, max-age=86400")
|
||||
w.WriteHeader(http.StatusPartialContent)
|
||||
_, _ = w.Write([]byte("data"))
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
allowPrivateTaskMediaTest(t)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/tasks/task_generic/artifacts/video-main/content", nil)
|
||||
c.Request.Header.Set("Range", "bytes=0-3")
|
||||
|
||||
err := proxyTaskMedia(c, task, &relaychannel.TaskContentRequest{
|
||||
URL: upstream.URL,
|
||||
Method: http.MethodGet,
|
||||
Headers: map[string]string{"Authorization": "Bearer provider-secret"},
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, http.StatusPartialContent, recorder.Code)
|
||||
assert.Equal(t, "data", recorder.Body.String())
|
||||
assert.Equal(t, "bytes=0-3", receivedRange)
|
||||
assert.Equal(t, "Bearer provider-secret", receivedAuthorization)
|
||||
assert.Equal(t, "bytes 0-3/10", recorder.Header().Get("Content-Range"))
|
||||
assert.Equal(t, "bytes", recorder.Header().Get("Accept-Ranges"))
|
||||
assert.Equal(t, "private, no-store", recorder.Header().Get("Cache-Control"))
|
||||
assert.Equal(t, "sandbox; default-src 'none'", recorder.Header().Get("Content-Security-Policy"))
|
||||
assert.Equal(t, "no-referrer", recorder.Header().Get("Referrer-Policy"))
|
||||
assert.Equal(t, "nosniff", recorder.Header().Get("X-Content-Type-Options"))
|
||||
assert.Empty(t, recorder.Header().Get("Set-Cookie"))
|
||||
assert.Empty(t, recorder.Header().Get("WWW-Authenticate"))
|
||||
assert.Empty(t, recorder.Header().Get("X-Provider-Secret"))
|
||||
}
|
||||
|
||||
func TestProxyTaskMediaPassesThroughUnsatisfiedRange(t *testing.T) {
|
||||
task := setupGenericTaskTest(t)
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.Header().Set("Content-Range", "bytes */10")
|
||||
w.WriteHeader(http.StatusRequestedRangeNotSatisfiable)
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
allowPrivateTaskMediaTest(t)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/content", nil)
|
||||
|
||||
require.NoError(t, proxyTaskMedia(c, task, &relaychannel.TaskContentRequest{
|
||||
URL: upstream.URL, Method: http.MethodGet,
|
||||
}))
|
||||
assert.Equal(t, http.StatusRequestedRangeNotSatisfiable, recorder.Code)
|
||||
assert.Equal(t, "bytes */10", recorder.Header().Get("Content-Range"))
|
||||
assert.Equal(t, "private, no-store", recorder.Header().Get("Cache-Control"))
|
||||
}
|
||||
|
||||
func TestTaskMediaResponseHeaderTimeoutDoesNotTruncateBody(t *testing.T) {
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.Header().Set("Content-Type", "video/mp4")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.(http.Flusher).Flush()
|
||||
time.Sleep(75 * time.Millisecond)
|
||||
_, _ = w.Write([]byte("complete-body"))
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
request, err := http.NewRequest(http.MethodGet, upstream.URL, nil)
|
||||
require.NoError(t, err)
|
||||
response, err := doTaskMediaRequest(upstream.Client(), request, 20*time.Millisecond)
|
||||
require.NoError(t, err)
|
||||
defer response.Body.Close()
|
||||
body, err := io.ReadAll(response.Body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "complete-body", string(body))
|
||||
}
|
||||
|
||||
func TestTaskMediaResponseHeaderTimeoutCancelsBeforeHeaders(t *testing.T) {
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
time.Sleep(75 * time.Millisecond)
|
||||
_, _ = w.Write([]byte("late"))
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
request, err := http.NewRequest(http.MethodGet, upstream.URL, nil)
|
||||
require.NoError(t, err)
|
||||
_, err = doTaskMediaRequest(upstream.Client(), request, 10*time.Millisecond)
|
||||
assert.ErrorIs(t, err, context.DeadlineExceeded)
|
||||
}
|
||||
|
||||
func TestWriteVideoDataURLStreamsAndSupportsHead(t *testing.T) {
|
||||
const dataURL = "data:video/mp4;base64,Y29tcGxldGUtYm9keQ=="
|
||||
|
||||
getRecorder := httptest.NewRecorder()
|
||||
getContext, _ := gin.CreateTestContext(getRecorder)
|
||||
getContext.Request = httptest.NewRequest(http.MethodGet, "/content", nil)
|
||||
require.NoError(t, writeVideoDataURL(getContext, dataURL))
|
||||
assert.Equal(t, http.StatusOK, getRecorder.Code)
|
||||
assert.Equal(t, "complete-body", getRecorder.Body.String())
|
||||
assert.Equal(t, "13", getRecorder.Header().Get("Content-Length"))
|
||||
|
||||
headRecorder := httptest.NewRecorder()
|
||||
headContext, _ := gin.CreateTestContext(headRecorder)
|
||||
headContext.Request = httptest.NewRequest(http.MethodHead, "/content", nil)
|
||||
require.NoError(t, writeVideoDataURL(headContext, dataURL))
|
||||
assert.Equal(t, http.StatusOK, headRecorder.Code)
|
||||
assert.Empty(t, headRecorder.Body.String())
|
||||
assert.Equal(t, "13", headRecorder.Header().Get("Content-Length"))
|
||||
}
|
||||
|
||||
func TestWriteVideoDataURLRejectsOversizedPayloadBeforeDecode(t *testing.T) {
|
||||
previousLimit := taskMediaDataURLMaxEncodedBytes
|
||||
taskMediaDataURLMaxEncodedBytes = 32
|
||||
t.Cleanup(func() { taskMediaDataURLMaxEncodedBytes = previousLimit })
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/content", nil)
|
||||
|
||||
err := writeVideoDataURL(c, "data:video/mp4;base64,"+strings.Repeat("A", 64))
|
||||
|
||||
assert.ErrorIs(t, err, errTaskMediaRequestRejected)
|
||||
assert.Empty(t, recorder.Header().Get("Content-Type"))
|
||||
}
|
||||
|
||||
func TestProxyTaskMediaAllowsOnlyCredentiallessCrossOriginRedirect(t *testing.T) {
|
||||
task := setupGenericTaskTest(t)
|
||||
var destinationAuthorization, destinationRange string
|
||||
destination := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
destinationAuthorization = r.Header.Get("Authorization")
|
||||
destinationRange = r.Header.Get("Range")
|
||||
_, _ = w.Write([]byte("redirected"))
|
||||
}))
|
||||
defer destination.Close()
|
||||
source := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
http.Redirect(w, r, destination.URL, http.StatusFound)
|
||||
}))
|
||||
defer source.Close()
|
||||
allowPrivateTaskMediaTest(t)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/content", nil)
|
||||
c.Request.Header.Set("Range", "bytes=0-3")
|
||||
|
||||
err := proxyTaskMedia(c, task, &relaychannel.TaskContentRequest{
|
||||
URL: source.URL, Method: http.MethodGet, Credentialless: true,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, http.StatusOK, recorder.Code)
|
||||
assert.Equal(t, "redirected", recorder.Body.String())
|
||||
assert.Empty(t, destinationAuthorization)
|
||||
assert.Equal(t, "bytes=0-3", destinationRange)
|
||||
|
||||
destinationRange = ""
|
||||
rejectedRecorder := httptest.NewRecorder()
|
||||
rejectedContext, _ := gin.CreateTestContext(rejectedRecorder)
|
||||
rejectedContext.Request = httptest.NewRequest(http.MethodGet, "/content", nil)
|
||||
err = proxyTaskMedia(rejectedContext, task, &relaychannel.TaskContentRequest{
|
||||
URL: source.URL, Method: http.MethodGet,
|
||||
Headers: map[string]string{"Authorization": "Bearer provider-secret"},
|
||||
})
|
||||
var proxyErr *taskMediaProxyError
|
||||
require.ErrorAs(t, err, &proxyErr)
|
||||
assert.Equal(t, "artifact_request_rejected", proxyErr.code)
|
||||
assert.Empty(t, destinationRange)
|
||||
}
|
||||
|
||||
func TestTaskMediaRequestHeaderPolicy(t *testing.T) {
|
||||
header := http.Header{}
|
||||
require.NoError(t, applyTaskMediaRequestHeaders(header, map[string]string{
|
||||
"Authorization": "Bearer provider-secret",
|
||||
"X-Signature": "signed",
|
||||
}))
|
||||
assert.Equal(t, "Bearer provider-secret", header.Get("Authorization"))
|
||||
assert.Equal(t, "signed", header.Get("X-Signature"))
|
||||
|
||||
for _, name := range []string{"Host", "Content-Length", "Accept-Encoding", "Connection", "Proxy-Authorization", "Transfer-Encoding"} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
assert.ErrorIs(t, applyTaskMediaRequestHeaders(http.Header{}, map[string]string{name: "bad"}), errTaskMediaRequestRejected)
|
||||
})
|
||||
}
|
||||
assert.ErrorIs(t, applyTaskMediaRequestHeaders(http.Header{}, map[string]string{"X-Test": "bad\r\ninjected"}), errTaskMediaRequestRejected)
|
||||
}
|
||||
|
||||
func TestCredentiallessTaskMediaDescriptorRejectsCredentialsAndBody(t *testing.T) {
|
||||
task := setupGenericTaskTest(t)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/content", nil)
|
||||
|
||||
for _, descriptor := range []*relaychannel.TaskContentRequest{
|
||||
{URL: "https://example.com/video", Method: http.MethodPost, Credentialless: true},
|
||||
{URL: "https://example.com/video", Method: http.MethodGet, Body: []byte("secret"), Credentialless: true},
|
||||
{URL: "https://example.com/video", Method: http.MethodGet, Headers: map[string]string{"X-Key": "secret"}, Credentialless: true},
|
||||
} {
|
||||
err := proxyTaskMedia(c, task, descriptor)
|
||||
var proxyErr *taskMediaProxyError
|
||||
require.ErrorAs(t, err, &proxyErr)
|
||||
assert.Equal(t, "artifact_request_rejected", proxyErr.code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelfTaskMediaURLGuard(t *testing.T) {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "https://gateway.example/v1/videos/task-1/content", nil)
|
||||
c.Request.Host = "gateway.example"
|
||||
|
||||
selfURL, err := url.Parse("https://gateway.example/v1/videos/task-1/content")
|
||||
require.NoError(t, err)
|
||||
assert.True(t, isSelfTaskMediaURL(c, selfURL))
|
||||
|
||||
remoteURL, err := url.Parse("https://cdn.example/v1/videos/task-1/content")
|
||||
require.NoError(t, err)
|
||||
assert.False(t, isSelfTaskMediaURL(c, remoteURL))
|
||||
assert.True(t, isTaskMediaFallbackLoop(remoteURL.String(), "task-1"))
|
||||
assert.False(t, isTaskMediaFallbackLoop(remoteURL.String(), "task-2"))
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestTaskLogDTOSeparatesUserAdminAndRootDetails(t *testing.T) {
|
||||
task := &model.Task{
|
||||
TaskID: "task_public",
|
||||
Platform: "document-parser",
|
||||
PrivateData: model.TaskPrivateData{
|
||||
Key: "channel-secret-canary",
|
||||
UpstreamTaskID: "upstream-private",
|
||||
NodeName: "node-a",
|
||||
Execution: &model.TaskExecutionSnapshot{
|
||||
RequestID: "request-public",
|
||||
RequestPath: "/v1/documents",
|
||||
TaskPlugin: &model.TaskPluginSnapshot{
|
||||
Key: "document-parser",
|
||||
Name: "Document Parser",
|
||||
Version: "1.2.3",
|
||||
Author: &model.TaskPluginAuthorSnapshot{
|
||||
Name: "Community Author",
|
||||
URL: "https://plugins.example/author",
|
||||
},
|
||||
APIVersion: 1,
|
||||
Generation: 42,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
userView := tasksToDto([]*model.Task{task}, false, common.RoleCommonUser)[0]
|
||||
assert.Nil(t, userView.AdminInfo)
|
||||
assert.Nil(t, userView.RootInfo)
|
||||
|
||||
adminView := tasksToDto([]*model.Task{task}, false, common.RoleAdminUser)[0]
|
||||
require.NotNil(t, adminView.AdminInfo)
|
||||
require.NotNil(t, adminView.AdminInfo.TaskPlugin)
|
||||
assert.Equal(t, "document-parser", adminView.AdminInfo.TaskPlugin.Key)
|
||||
assert.Equal(t, "Document Parser", adminView.AdminInfo.TaskPlugin.Name)
|
||||
assert.Equal(t, "1.2.3", adminView.AdminInfo.TaskPlugin.Version)
|
||||
require.NotNil(t, adminView.AdminInfo.TaskPlugin.Author)
|
||||
assert.Equal(t, "Community Author", adminView.AdminInfo.TaskPlugin.Author.Name)
|
||||
assert.Equal(t, "https://plugins.example/author", adminView.AdminInfo.TaskPlugin.Author.URL)
|
||||
assert.Equal(t, "request-public", adminView.AdminInfo.RequestID)
|
||||
assert.Equal(t, "/v1/documents", adminView.AdminInfo.RequestPath)
|
||||
assert.Nil(t, adminView.RootInfo)
|
||||
|
||||
rootView := tasksToDto([]*model.Task{task}, false, common.RoleRootUser)[0]
|
||||
require.NotNil(t, rootView.AdminInfo)
|
||||
require.NotNil(t, rootView.RootInfo)
|
||||
require.NotNil(t, rootView.RootInfo.TaskPlugin)
|
||||
assert.Equal(t, 1, rootView.RootInfo.TaskPlugin.APIVersion)
|
||||
assert.Equal(t, uint64(42), rootView.RootInfo.TaskPlugin.Generation)
|
||||
assert.Equal(t, "upstream-private", rootView.RootInfo.UpstreamTaskID)
|
||||
assert.Equal(t, "node-a", rootView.RootInfo.NodeName)
|
||||
|
||||
adminJSON, err := common.Marshal(adminView)
|
||||
require.NoError(t, err)
|
||||
assert.NotContains(t, string(adminJSON), "channel-secret-canary")
|
||||
assert.NotContains(t, string(adminJSON), "upstream-private")
|
||||
|
||||
rootJSON, err := common.Marshal(rootView)
|
||||
require.NoError(t, err)
|
||||
assert.NotContains(t, string(rootJSON), "channel-secret-canary")
|
||||
assert.Contains(t, string(rootJSON), "upstream-private")
|
||||
}
|
||||
|
||||
func TestTaskLogDTODoesNotInventHistoricalPluginProvenance(t *testing.T) {
|
||||
task := &model.Task{
|
||||
TaskID: "task_without_snapshot",
|
||||
Platform: "document-parser",
|
||||
}
|
||||
|
||||
adminView := tasksToDto([]*model.Task{task}, false, common.RoleAdminUser)[0]
|
||||
|
||||
assert.Nil(t, adminView.AdminInfo)
|
||||
assert.Nil(t, adminView.RootInfo)
|
||||
}
|
||||
|
||||
func TestTaskLogDTOReplacesLegacyVideoURLWithAvailabilityFlag(t *testing.T) {
|
||||
task := &model.Task{
|
||||
TaskID: "task_legacy_video",
|
||||
Platform: "jimeng",
|
||||
Action: constant.TaskActionTextToVideo,
|
||||
Status: model.TaskStatusSuccess,
|
||||
FailReason: "https://private-upstream.invalid/video.mp4?signature=secret",
|
||||
}
|
||||
|
||||
view := tasksToDto([]*model.Task{task}, false, common.RoleCommonUser)[0]
|
||||
assert.True(t, view.LegacyVideoAvailable)
|
||||
assert.Empty(t, view.ResultURL)
|
||||
assert.Empty(t, view.FailReason)
|
||||
encoded, err := common.Marshal(view)
|
||||
require.NoError(t, err)
|
||||
assert.NotContains(t, string(encoded), "private-upstream.invalid")
|
||||
assert.NotContains(t, string(encoded), "result_url")
|
||||
assert.Contains(t, string(encoded), "legacy_video_available")
|
||||
}
|
||||
|
||||
func TestTaskLogDTOKeepsFailureReasonAndDoesNotMarkPluginTaskLegacy(t *testing.T) {
|
||||
failed := &model.Task{
|
||||
TaskID: "task_failed",
|
||||
Platform: "jimeng",
|
||||
Action: constant.TaskActionTextToVideo,
|
||||
Status: model.TaskStatusFailure,
|
||||
FailReason: "provider rejected the request",
|
||||
}
|
||||
failedView := tasksToDto([]*model.Task{failed}, false, common.RoleCommonUser)[0]
|
||||
assert.Equal(t, "provider rejected the request", failedView.FailReason)
|
||||
assert.False(t, failedView.LegacyVideoAvailable)
|
||||
|
||||
pluginTask := &model.Task{
|
||||
TaskID: "task_plugin_video",
|
||||
Platform: "community-video",
|
||||
Action: constant.TaskActionTextToVideo,
|
||||
Status: model.TaskStatusSuccess,
|
||||
FailReason: "https://stale-upstream.invalid/plugin-video.mp4",
|
||||
PrivateData: model.TaskPrivateData{
|
||||
ResultURL: "https://private-upstream.invalid/plugin-video.mp4",
|
||||
Execution: &model.TaskExecutionSnapshot{
|
||||
TaskPlugin: &model.TaskPluginSnapshot{Key: "community-video"},
|
||||
},
|
||||
},
|
||||
}
|
||||
pluginView := tasksToDto([]*model.Task{pluginTask}, false, common.RoleCommonUser)[0]
|
||||
assert.False(t, pluginView.LegacyVideoAvailable)
|
||||
assert.Empty(t, pluginView.ResultURL)
|
||||
assert.Empty(t, pluginView.FailReason)
|
||||
}
|
||||
@@ -0,0 +1,835 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"maps"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/plugins"
|
||||
"github.com/QuantumNous/new-api/setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
const maxTaskPluginSourceBytes = 1024 * 1024
|
||||
|
||||
type taskPluginUploadRequest struct {
|
||||
Source string `json:"source" binding:"required"`
|
||||
Enabled *bool `json:"enabled"`
|
||||
Remark string `json:"remark"`
|
||||
Force bool `json:"force"`
|
||||
SourceSha256 string `json:"sourceSha256"`
|
||||
// Icon carries the sidecar icon.svg / icon.png as a data URI. It is optional
|
||||
// and stored separately from the source so the JavaScript stays readable.
|
||||
Icon string `json:"icon"`
|
||||
}
|
||||
|
||||
func UploadTaskPlugin(c *gin.Context) {
|
||||
var request taskPluginUploadRequest
|
||||
if err := c.ShouldBindJSON(&request); err != nil {
|
||||
common.ApiErrorMsg(c, err.Error())
|
||||
return
|
||||
}
|
||||
if len(request.Source) > maxTaskPluginSourceBytes {
|
||||
common.ApiErrorMsg(c, "plugin source exceeds 1 MiB")
|
||||
return
|
||||
}
|
||||
if expected := strings.TrimSpace(request.SourceSha256); expected != "" {
|
||||
actual := fmt.Sprintf("%x", sha256.Sum256([]byte(request.Source)))
|
||||
if !strings.EqualFold(actual, expected) {
|
||||
common.ApiErrorMsg(c, "plugin source sha256 mismatch")
|
||||
return
|
||||
}
|
||||
}
|
||||
temporary := jsplugin.NewRegistry()
|
||||
loaded, err := temporary.Register(request.Source, jsplugin.Options{})
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, err.Error())
|
||||
return
|
||||
}
|
||||
if err = jsplugin.ValidateV1Meta(loaded.Meta); err != nil {
|
||||
common.ApiErrorMsg(c, err.Error())
|
||||
return
|
||||
}
|
||||
icon := strings.TrimSpace(request.Icon)
|
||||
if icon != "" {
|
||||
if _, _, err = jsplugin.DecodeIconDataURI(icon); err != nil {
|
||||
common.ApiErrorMsg(c, err.Error())
|
||||
return
|
||||
}
|
||||
}
|
||||
enabled := true
|
||||
if request.Enabled != nil {
|
||||
enabled = *request.Enabled
|
||||
}
|
||||
if enabled && !request.Force {
|
||||
if err = jsplugin.PreflightRoutingConflict(jsplugin.DefaultRegistry.Generation(), loaded); err != nil {
|
||||
common.ApiErrorMsg(c, err.Error())
|
||||
return
|
||||
}
|
||||
}
|
||||
plugin := model.TaskPlugin{
|
||||
Key: loaded.Meta.Key, APIVersion: loaded.Meta.APIVersion, Version: loaded.Meta.Version,
|
||||
Source: request.Source, SourceHash: fmt.Sprintf("%x", sha256.Sum256([]byte(request.Source))),
|
||||
Icon: icon, Enabled: enabled, Remark: request.Remark,
|
||||
}
|
||||
if err = model.SaveTaskPlugin(&plugin); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if err = syncTaskPluginsOnceContext(c.Request.Context()); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, taskPluginDetail{Plugin: &plugin, Meta: loaded.Meta, Source: plugin.Source, Layer: "override", HasIcon: plugin.HasIcon()})
|
||||
}
|
||||
|
||||
func GetTaskPluginVersions(c *gin.Context) {
|
||||
plugins, err := model.ListTaskPluginVersions(c.Param("key"))
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, plugins)
|
||||
}
|
||||
|
||||
type taskPluginListItem struct {
|
||||
Meta jsplugin.Meta `json:"meta"`
|
||||
Source string `json:"source"`
|
||||
Enabled bool `json:"enabled"`
|
||||
Active bool `json:"active"`
|
||||
SourceHash string `json:"source_hash"`
|
||||
HasIcon bool `json:"has_icon"`
|
||||
Remark string `json:"remark"`
|
||||
RuntimeStatus string `json:"runtime_status"`
|
||||
RuntimeError string `json:"runtime_error,omitempty"`
|
||||
FactoryMeta *jsplugin.Meta `json:"factory_meta,omitempty"`
|
||||
ChannelCount int `json:"channel_count"`
|
||||
InFlightCount int64 `json:"in_flight_count"`
|
||||
}
|
||||
|
||||
type taskPluginRebuildOutcome struct {
|
||||
Status string `json:"status"`
|
||||
AttemptedAt time.Time `json:"attempted_at"`
|
||||
Generation uint64 `json:"generation"`
|
||||
DatabaseRevision string `json:"database_revision,omitempty"`
|
||||
PluginErrorCount int `json:"plugin_error_count"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type taskPluginRuntimeStatus struct {
|
||||
CurrentGeneration uint64 `json:"current_generation"`
|
||||
GenerationPublishedAt time.Time `json:"generation_published_at"`
|
||||
DatabaseRevision string `json:"database_revision"`
|
||||
DatabaseError string `json:"database_error,omitempty"`
|
||||
LastRebuild taskPluginRebuildOutcome `json:"last_rebuild"`
|
||||
PluginErrors map[string]string `json:"plugin_errors"`
|
||||
}
|
||||
|
||||
func ListTaskPlugins(c *gin.Context) {
|
||||
databasePlugins, err := model.ListTaskPlugins()
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
snapshot := jsplugin.DefaultRegistry.Snapshot()
|
||||
factory := make(map[string]jsplugin.Meta, len(snapshot.Factory))
|
||||
override := make(map[string]jsplugin.Meta, len(snapshot.Override))
|
||||
for _, meta := range snapshot.Factory {
|
||||
factory[meta.Key] = meta
|
||||
}
|
||||
for _, meta := range snapshot.Override {
|
||||
override[meta.Key] = meta
|
||||
}
|
||||
activeRows := make(map[string]model.TaskPlugin)
|
||||
keys := make(map[string]struct{}, len(factory)+len(databasePlugins))
|
||||
for key := range factory {
|
||||
keys[key] = struct{}{}
|
||||
}
|
||||
for _, plugin := range databasePlugins {
|
||||
keys[plugin.Key] = struct{}{}
|
||||
if plugin.Active {
|
||||
activeRows[plugin.Key] = plugin
|
||||
}
|
||||
}
|
||||
|
||||
runtimeErrors := jsplugin.DefaultRegistry.RoutingErrors()
|
||||
taskPluginSyncState.Lock()
|
||||
maps.Copy(runtimeErrors, taskPluginSyncState.errors)
|
||||
taskPluginSyncState.Unlock()
|
||||
|
||||
items := make([]taskPluginListItem, 0, len(keys))
|
||||
for key := range keys {
|
||||
factoryMeta, hasFactory := factory[key]
|
||||
row, hasOverride := activeRows[key]
|
||||
item := taskPluginListItem{Enabled: true, Active: true, RuntimeStatus: "registered"}
|
||||
if hasOverride {
|
||||
item.Source = "override"
|
||||
if hasFactory {
|
||||
item.Source = "override_over_factory"
|
||||
factoryCopy := factoryMeta
|
||||
item.FactoryMeta = &factoryCopy
|
||||
}
|
||||
item.Meta = jsplugin.Meta{Key: row.Key, Version: row.Version, APIVersion: row.APIVersion}
|
||||
if compiled, compileErr := jsplugin.NewRegistry().Register(row.Source, jsplugin.Options{Key: row.Key, Version: row.Version}); compileErr == nil {
|
||||
item.Meta = compiled.Meta
|
||||
}
|
||||
item.Enabled = row.Enabled
|
||||
item.Active = row.Active
|
||||
item.SourceHash = row.SourceHash
|
||||
item.HasIcon = row.HasIcon()
|
||||
item.Remark = row.Remark
|
||||
if message := runtimeErrors[key]; message != "" {
|
||||
item.RuntimeStatus = "compile_failed"
|
||||
item.RuntimeError = message
|
||||
} else if runtimeMeta, ok := override[key]; ok {
|
||||
item.Meta = runtimeMeta
|
||||
} else if !row.Enabled {
|
||||
item.RuntimeStatus = "disabled_fallback"
|
||||
} else {
|
||||
item.RuntimeStatus = "not_registered"
|
||||
}
|
||||
// "disabled_fallback" promises that the built-in still serves. When
|
||||
// the factory layer is suppressed as well, nothing serves this key.
|
||||
if item.RuntimeStatus == "disabled_fallback" && hasFactory && setting.IsTaskPluginFactoryDisabled(key) {
|
||||
item.RuntimeStatus = "disabled"
|
||||
}
|
||||
} else {
|
||||
item.Source = "factory"
|
||||
item.Meta = factoryMeta
|
||||
item.Enabled = !setting.IsTaskPluginFactoryDisabled(key)
|
||||
_, _, item.HasIcon = plugins.Icon(key)
|
||||
source, sourceErr := plugins.Source(key)
|
||||
if sourceErr == nil {
|
||||
item.SourceHash = fmt.Sprintf("%x", sha256.Sum256([]byte(source)))
|
||||
}
|
||||
if !item.Enabled {
|
||||
item.RuntimeStatus = "disabled"
|
||||
} else if message := runtimeErrors[key]; message != "" {
|
||||
item.RuntimeStatus = "compile_failed"
|
||||
item.RuntimeError = message
|
||||
}
|
||||
}
|
||||
if !hasFactory {
|
||||
channels, inFlight, usageErr := model.GetTaskPluginUsage(key)
|
||||
if usageErr != nil {
|
||||
common.ApiError(c, usageErr)
|
||||
return
|
||||
}
|
||||
item.ChannelCount = len(channels)
|
||||
item.InFlightCount = inFlight
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
sort.Slice(items, func(i, j int) bool {
|
||||
if items[i].Meta.SortPriority != items[j].Meta.SortPriority {
|
||||
return items[i].Meta.SortPriority > items[j].Meta.SortPriority
|
||||
}
|
||||
return items[i].Meta.Key < items[j].Meta.Key
|
||||
})
|
||||
common.ApiSuccess(c, items)
|
||||
}
|
||||
|
||||
func GetTaskPluginRuntime(c *gin.Context) {
|
||||
routingStatus := jsplugin.DefaultRegistry.RoutingStatus()
|
||||
pluginErrors := routingStatus.Errors
|
||||
|
||||
taskPluginSyncState.Lock()
|
||||
maps.Copy(pluginErrors, taskPluginSyncState.errors)
|
||||
lastRebuild := taskPluginSyncState.lastRebuild
|
||||
lastDatabaseRevision := lastRebuild.DatabaseRevision
|
||||
taskPluginSyncState.Unlock()
|
||||
|
||||
registryRebuild := routingStatus.LastRebuild
|
||||
if lastRebuild.AttemptedAt.Before(registryRebuild.AttemptedAt) {
|
||||
lastRebuild = taskPluginRebuildOutcome{
|
||||
Status: registryRebuild.Status,
|
||||
AttemptedAt: registryRebuild.AttemptedAt,
|
||||
Generation: registryRebuild.Generation,
|
||||
Error: registryRebuild.Error,
|
||||
}
|
||||
}
|
||||
if lastRebuild.Status == "" {
|
||||
lastRebuild.Status = "never"
|
||||
}
|
||||
lastRebuild.PluginErrorCount = len(pluginErrors)
|
||||
if lastRebuild.Status == "success" && len(pluginErrors) > 0 {
|
||||
lastRebuild.Status = "partial"
|
||||
}
|
||||
|
||||
status := taskPluginRuntimeStatus{
|
||||
DatabaseRevision: lastDatabaseRevision,
|
||||
LastRebuild: lastRebuild,
|
||||
PluginErrors: pluginErrors,
|
||||
}
|
||||
databaseSnapshot, err := model.GetTaskPluginSyncSnapshot()
|
||||
if err != nil {
|
||||
status.DatabaseError = "database snapshot unavailable"
|
||||
} else {
|
||||
status.DatabaseRevision = databaseSnapshot.Revision
|
||||
}
|
||||
if routingStatus.Generation != nil {
|
||||
status.CurrentGeneration = routingStatus.Generation.Number
|
||||
status.GenerationPublishedAt = routingStatus.Generation.PublishedAt
|
||||
}
|
||||
common.ApiSuccess(c, status)
|
||||
}
|
||||
|
||||
type taskPluginDetail struct {
|
||||
Plugin *model.TaskPlugin `json:"plugin,omitempty"`
|
||||
Meta jsplugin.Meta `json:"meta"`
|
||||
Source string `json:"source"`
|
||||
Layer string `json:"layer"`
|
||||
HasIcon bool `json:"has_icon"`
|
||||
}
|
||||
|
||||
// GetTaskPluginIcon serves a plugin logo as an image. The active override wins
|
||||
// (or the requested ?version=), then the factory sidecar. Data icons are only
|
||||
// ever drawn through <img>, and the nosniff header keeps a browser from
|
||||
// treating an SVG response as a document.
|
||||
func GetTaskPluginIcon(c *gin.Context) {
|
||||
key := c.Param("key")
|
||||
icon := ""
|
||||
plugin, err := model.GetTaskPluginVersion(key, c.Query("version"))
|
||||
if err == nil {
|
||||
icon = plugin.Icon
|
||||
} else if !errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if icon == "" && c.Query("version") == "" {
|
||||
icon = plugins.IconDataURI(key)
|
||||
}
|
||||
if icon == "" {
|
||||
c.AbortWithStatus(http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
mediaType, data, err := jsplugin.DecodeIconDataURI(icon)
|
||||
if err != nil {
|
||||
c.AbortWithStatus(http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
c.Header("Cache-Control", "private, max-age=3600")
|
||||
c.Header("X-Content-Type-Options", "nosniff")
|
||||
c.Header("Content-Security-Policy", "default-src 'none'; style-src 'unsafe-inline'; sandbox")
|
||||
c.Data(http.StatusOK, mediaType, data)
|
||||
}
|
||||
|
||||
func GetTaskPlugin(c *gin.Context) {
|
||||
key := c.Param("key")
|
||||
version := c.Query("version")
|
||||
plugin, err := model.GetTaskPluginVersion(key, version)
|
||||
if err == nil {
|
||||
loaded, compileErr := jsplugin.NewRegistry().Register(plugin.Source, jsplugin.Options{Key: plugin.Key, Version: plugin.Version})
|
||||
if compileErr != nil {
|
||||
common.ApiErrorMsg(c, compileErr.Error())
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, taskPluginDetail{Plugin: plugin, Meta: loaded.Meta, Source: plugin.Source, Layer: "override", HasIcon: plugin.HasIcon()})
|
||||
return
|
||||
}
|
||||
if !errors.Is(err, gorm.ErrRecordNotFound) || version != "" {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
source, err := plugins.Source(key)
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, "task plugin not found")
|
||||
return
|
||||
}
|
||||
loaded, err := jsplugin.NewRegistry().RegisterFactory(source, jsplugin.Options{Key: key})
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
_, _, hasIcon := plugins.Icon(key)
|
||||
common.ApiSuccess(c, taskPluginDetail{Meta: loaded.Meta, Source: source, Layer: "factory", HasIcon: hasIcon})
|
||||
}
|
||||
|
||||
type taskPluginDryRunRequest struct {
|
||||
Hook string `json:"hook" binding:"required"`
|
||||
Member string `json:"member"`
|
||||
Args []json.RawMessage `json:"args"`
|
||||
}
|
||||
|
||||
func DryRunTaskPlugin(c *gin.Context) {
|
||||
var request taskPluginDryRunRequest
|
||||
if err := c.ShouldBindJSON(&request); err != nil {
|
||||
common.ApiErrorMsg(c, err.Error())
|
||||
return
|
||||
}
|
||||
detailSource := ""
|
||||
plugin, err := model.GetTaskPluginVersion(c.Param("key"), "")
|
||||
if err == nil {
|
||||
detailSource = plugin.Source
|
||||
} else if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
detailSource, err = plugins.Source(c.Param("key"))
|
||||
}
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, "task plugin not found")
|
||||
return
|
||||
}
|
||||
loaded, err := jsplugin.NewRegistry().Register(detailSource, jsplugin.Options{Key: c.Param("key")})
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, err.Error())
|
||||
return
|
||||
}
|
||||
args := make([]any, len(request.Args))
|
||||
for index, raw := range request.Args {
|
||||
if err = common.Unmarshal(raw, &args[index]); err != nil {
|
||||
common.ApiErrorMsg(c, fmt.Sprintf("invalid argument %d: %v", index+1, err))
|
||||
return
|
||||
}
|
||||
}
|
||||
var output any
|
||||
if request.Member == "" {
|
||||
output, err = loaded.Engine.Call(context.Background(), request.Hook, args...)
|
||||
} else {
|
||||
output, err = loaded.Engine.CallMember(context.Background(), request.Hook, request.Member, args...)
|
||||
}
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, err.Error())
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, output)
|
||||
}
|
||||
|
||||
func DeleteTaskPluginVersion(c *gin.Context) {
|
||||
key := c.Param("key")
|
||||
version := c.Param("version")
|
||||
plugin, lookupErr := model.GetTaskPluginVersion(key, version)
|
||||
if lookupErr != nil {
|
||||
if errors.Is(lookupErr, gorm.ErrRecordNotFound) {
|
||||
common.ApiErrorMsg(c, "override plugin version not found; factory plugins cannot be deleted")
|
||||
return
|
||||
}
|
||||
common.ApiError(c, lookupErr)
|
||||
return
|
||||
}
|
||||
if plugin.Active && !taskPluginHasFactory(key) {
|
||||
channels, inFlight, usageErr := model.GetTaskPluginUsage(key)
|
||||
if usageErr != nil {
|
||||
common.ApiError(c, usageErr)
|
||||
return
|
||||
}
|
||||
if (len(channels) > 0 || inFlight > 0) && c.Query("force") != "true" {
|
||||
c.JSON(200, gin.H{"success": false, "message": "task plugin is still in use", "data": gin.H{"channels": channels, "in_flight_count": inFlight}})
|
||||
return
|
||||
}
|
||||
}
|
||||
_, err := model.DeleteTaskPluginVersion(key, version)
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
common.ApiErrorMsg(c, "override plugin version not found; factory plugins cannot be deleted")
|
||||
return
|
||||
}
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if err = syncTaskPluginsOnceContext(c.Request.Context()); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, nil)
|
||||
}
|
||||
|
||||
type taskPluginActivateRequest struct {
|
||||
Version string `json:"version" binding:"required"`
|
||||
}
|
||||
|
||||
func ActivateTaskPlugin(c *gin.Context) {
|
||||
var request taskPluginActivateRequest
|
||||
if err := c.ShouldBindJSON(&request); err != nil {
|
||||
common.ApiErrorMsg(c, err.Error())
|
||||
return
|
||||
}
|
||||
versions, err := model.ListTaskPluginVersions(c.Param("key"))
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
var target *model.TaskPlugin
|
||||
for i := range versions {
|
||||
if versions[i].Version == request.Version {
|
||||
target = &versions[i]
|
||||
break
|
||||
}
|
||||
}
|
||||
if target == nil {
|
||||
common.ApiErrorMsg(c, "plugin version not found")
|
||||
return
|
||||
}
|
||||
if _, err = jsplugin.NewRegistry().Register(target.Source, jsplugin.Options{Key: target.Key, Version: target.Version}); err != nil {
|
||||
common.ApiErrorMsg(c, err.Error())
|
||||
return
|
||||
}
|
||||
if err = model.ActivateTaskPlugin(target.Key, target.Version); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if err = syncTaskPluginsOnceContext(c.Request.Context()); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, nil)
|
||||
}
|
||||
|
||||
type taskPluginStatusRequest struct {
|
||||
Enabled *bool `json:"enabled" binding:"required"`
|
||||
}
|
||||
|
||||
func SetTaskPluginStatus(c *gin.Context) {
|
||||
var request taskPluginStatusRequest
|
||||
if err := c.ShouldBindJSON(&request); err != nil || request.Enabled == nil {
|
||||
common.ApiErrorMsg(c, "enabled is required")
|
||||
return
|
||||
}
|
||||
key := c.Param("key")
|
||||
disabledChannels := 0
|
||||
if !*request.Enabled {
|
||||
channels, inFlight, usageErr := model.GetTaskPluginUsage(key)
|
||||
if usageErr != nil {
|
||||
common.ApiError(c, usageErr)
|
||||
return
|
||||
}
|
||||
cascade := c.Query("cascade") == "true"
|
||||
force := c.Query("force") == "true"
|
||||
if (len(channels) > 0 && !cascade) || (inFlight > 0 && !force) {
|
||||
c.JSON(200, gin.H{"success": false, "message": "task plugin is still in use", "data": gin.H{"channels": channels, "in_flight_count": inFlight}})
|
||||
return
|
||||
}
|
||||
if cascade {
|
||||
for _, channel := range channels {
|
||||
if model.UpdateChannelStatus(channel.Id, "", common.ChannelStatusManuallyDisabled, "task plugin disabled") {
|
||||
disabledChannels++
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
_, lookupErr := model.GetTaskPluginVersion(key, "")
|
||||
hasActiveOverride := lookupErr == nil
|
||||
if lookupErr != nil && !errors.Is(lookupErr, gorm.ErrRecordNotFound) {
|
||||
common.ApiError(c, lookupErr)
|
||||
return
|
||||
}
|
||||
// Switching a key off must silence every layer that can serve it. The
|
||||
// factory built-in goes into the disabled set even when an override row
|
||||
// exists; otherwise the built-in would keep routing the same models (and
|
||||
// blocking same-model uploads) right after the administrator disabled the
|
||||
// plugin. Switching on reverses both layers.
|
||||
if taskPluginHasFactory(key) {
|
||||
keys := setting.GetTaskPluginDisabledFactoryKeys()
|
||||
if *request.Enabled {
|
||||
next := make([]string, 0, len(keys))
|
||||
for _, item := range keys {
|
||||
if item != key {
|
||||
next = append(next, item)
|
||||
}
|
||||
}
|
||||
keys = next
|
||||
} else {
|
||||
keys = append(append([]string{}, keys...), key)
|
||||
}
|
||||
if err := setting.SetTaskPluginDisabledFactoryKeysOption(keys); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
encoded, err := common.Marshal(setting.GetTaskPluginDisabledFactoryKeys())
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if err = model.UpdateOption(setting.TaskPluginDisabledFactoryKeysKey, string(encoded)); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if !hasActiveOverride {
|
||||
common.ApiSuccess(c, gin.H{"plugin_enabled": *request.Enabled, "disabled_channels": disabledChannels})
|
||||
return
|
||||
}
|
||||
}
|
||||
if err := model.SetTaskPluginEnabled(key, *request.Enabled); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if err := syncTaskPluginsOnceContext(c.Request.Context()); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, gin.H{"plugin_enabled": *request.Enabled, "disabled_channels": disabledChannels})
|
||||
}
|
||||
|
||||
func taskPluginHasFactory(key string) bool {
|
||||
for _, meta := range jsplugin.DefaultRegistry.Snapshot().Factory {
|
||||
if meta.Key == key {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func GetTaskPluginMarketplaceSources(c *gin.Context) {
|
||||
common.ApiSuccess(c, setting.GetTaskPluginMarketplaceSources())
|
||||
}
|
||||
|
||||
func UpdateTaskPluginMarketplaceSources(c *gin.Context) {
|
||||
var sources []setting.TaskPluginMarketplaceSource
|
||||
if err := c.ShouldBindJSON(&sources); err != nil {
|
||||
common.ApiErrorMsg(c, err.Error())
|
||||
return
|
||||
}
|
||||
if sources == nil {
|
||||
sources = []setting.TaskPluginMarketplaceSource{}
|
||||
}
|
||||
for i := range sources {
|
||||
name := strings.TrimSpace(sources[i].Name)
|
||||
indexURL := strings.TrimSpace(sources[i].IndexURL)
|
||||
if name == "" {
|
||||
common.ApiErrorMsg(c, "marketplace source name is required")
|
||||
return
|
||||
}
|
||||
parsed, err := url.Parse(indexURL)
|
||||
if err != nil || !parsed.IsAbs() || parsed.Host == "" || (!strings.EqualFold(parsed.Scheme, "http") && !strings.EqualFold(parsed.Scheme, "https")) {
|
||||
common.ApiErrorMsg(c, "marketplace source index_url must be an absolute http(s) URL")
|
||||
return
|
||||
}
|
||||
sources[i].Name = name
|
||||
sources[i].IndexURL = indexURL
|
||||
}
|
||||
encoded, err := common.Marshal(sources)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if err = model.UpdateOption(setting.TaskPluginMarketplaceSourcesKey, string(encoded)); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, sources)
|
||||
}
|
||||
|
||||
func GetTaskPluginOptions(c *gin.Context) {
|
||||
snapshot := jsplugin.DefaultRegistry.Snapshot()
|
||||
seen := make(map[string]bool)
|
||||
options := make([]gin.H, 0, len(snapshot.Factory)+len(snapshot.Override))
|
||||
for layer, metas := range [][]jsplugin.Meta{snapshot.Override, snapshot.Factory} {
|
||||
for _, meta := range metas {
|
||||
if seen[meta.Key] {
|
||||
continue
|
||||
}
|
||||
// Disabled factory keys are omitted from bind options. The disabled
|
||||
// set suppresses only the factory fallback; an enabled override for
|
||||
// the same key is listed in the override pass and still appears.
|
||||
if layer == 1 && setting.IsTaskPluginFactoryDisabled(meta.Key) {
|
||||
continue
|
||||
}
|
||||
if _, ok := jsplugin.DefaultRegistry.Get(meta.Key); !ok {
|
||||
continue
|
||||
}
|
||||
seen[meta.Key] = true
|
||||
hasIcon := false
|
||||
if layer == 0 {
|
||||
if row, rowErr := model.GetTaskPluginVersion(meta.Key, ""); rowErr == nil {
|
||||
hasIcon = row.HasIcon()
|
||||
}
|
||||
} else {
|
||||
_, _, hasIcon = plugins.Icon(meta.Key)
|
||||
}
|
||||
options = append(options, gin.H{
|
||||
"key": meta.Key,
|
||||
"name": meta.Name,
|
||||
"icon": meta.Icon,
|
||||
"hasIcon": hasIcon,
|
||||
"baseUrl": meta.BaseURL,
|
||||
"sortPriority": meta.SortPriority,
|
||||
"website": meta.Website,
|
||||
"models": meta.Models,
|
||||
"usageSchema": meta.UsageSchema,
|
||||
})
|
||||
}
|
||||
}
|
||||
sort.Slice(options, func(i, j int) bool {
|
||||
left, right := options[i]["sortPriority"].(int), options[j]["sortPriority"].(int)
|
||||
if left != right {
|
||||
return left > right
|
||||
}
|
||||
return options[i]["key"].(string) < options[j]["key"].(string)
|
||||
})
|
||||
common.ApiSuccess(c, options)
|
||||
}
|
||||
|
||||
var taskPluginSyncState = struct {
|
||||
sync.Mutex
|
||||
hashes map[string]string
|
||||
errors map[string]string
|
||||
lastRebuild taskPluginRebuildOutcome
|
||||
}{hashes: map[string]string{}, errors: map[string]string{}}
|
||||
|
||||
func syncTaskPluginsOnce() error {
|
||||
return syncTaskPluginsOnceContext(context.Background())
|
||||
}
|
||||
|
||||
func syncTaskPluginsOnceContext(ctx context.Context) error {
|
||||
started := time.Now()
|
||||
taskPluginSyncState.Lock()
|
||||
defer taskPluginSyncState.Unlock()
|
||||
databaseSnapshot, err := model.GetTaskPluginSyncSnapshot()
|
||||
if err != nil {
|
||||
syncErr := fmt.Errorf("sync task plugins: %w", err)
|
||||
taskPluginSyncState.lastRebuild = taskPluginRebuildOutcome{
|
||||
Status: "failed",
|
||||
AttemptedAt: time.Now(),
|
||||
Generation: jsplugin.DefaultRegistry.Generation().Number,
|
||||
DatabaseRevision: taskPluginSyncState.lastRebuild.DatabaseRevision,
|
||||
Error: syncErr.Error(),
|
||||
}
|
||||
logger.LogDebug(
|
||||
ctx,
|
||||
"task_plugin subsystem=sync event=failed stage=database_snapshot retained_generation=%d elapsed_ms=%d",
|
||||
jsplugin.DefaultRegistry.Generation().Number,
|
||||
time.Since(started).Milliseconds(),
|
||||
)
|
||||
return syncErr
|
||||
}
|
||||
databasePlugins := databaseSnapshot.Plugins
|
||||
sort.Slice(databasePlugins, func(i, j int) bool { return databasePlugins[i].Key < databasePlugins[j].Key })
|
||||
currentOverrides := jsplugin.DefaultRegistry.OverridePlugins()
|
||||
generationBefore := jsplugin.DefaultRegistry.Generation().Number
|
||||
logger.LogDebug(
|
||||
ctx,
|
||||
"task_plugin subsystem=sync event=start database_revision=%q generation=%d desired_plugins=%d current_overrides=%d",
|
||||
databaseSnapshot.Revision,
|
||||
generationBefore,
|
||||
len(databasePlugins),
|
||||
len(currentOverrides),
|
||||
)
|
||||
nextOverrides := make([]*jsplugin.LoadedPlugin, 0, len(databasePlugins))
|
||||
nextHashes := make(map[string]string, len(databasePlugins))
|
||||
seen := make(map[string]bool, len(databasePlugins))
|
||||
for _, plugin := range databasePlugins {
|
||||
seen[plugin.Key] = true
|
||||
if current := currentOverrides[plugin.Key]; current != nil && taskPluginSyncState.hashes[plugin.Key] == plugin.SourceHash {
|
||||
nextOverrides = append(nextOverrides, current)
|
||||
nextHashes[plugin.Key] = plugin.SourceHash
|
||||
logger.LogDebug(
|
||||
ctx,
|
||||
"task_plugin subsystem=sync event=plugin plugin=%q version=%q action=reuse",
|
||||
plugin.Key,
|
||||
plugin.Version,
|
||||
)
|
||||
continue
|
||||
}
|
||||
logger.LogDebug(
|
||||
ctx,
|
||||
"task_plugin subsystem=sync event=plugin plugin=%q version=%q action=compile_start",
|
||||
plugin.Key,
|
||||
plugin.Version,
|
||||
)
|
||||
compiled, compileErr := jsplugin.CompilePlugin(plugin.Source, jsplugin.Options{Key: plugin.Key, Version: plugin.Version})
|
||||
if compileErr != nil {
|
||||
retainedIncumbent := false
|
||||
if current := currentOverrides[plugin.Key]; current != nil {
|
||||
nextOverrides = append(nextOverrides, current)
|
||||
retainedIncumbent = true
|
||||
if currentHash := taskPluginSyncState.hashes[plugin.Key]; currentHash != "" {
|
||||
nextHashes[plugin.Key] = currentHash
|
||||
}
|
||||
}
|
||||
taskPluginSyncState.errors[plugin.Key] = compileErr.Error()
|
||||
common.SysError(fmt.Sprintf("compile task plugin %s@%s: %v", plugin.Key, plugin.Version, compileErr))
|
||||
logger.LogDebug(
|
||||
ctx,
|
||||
"task_plugin subsystem=sync event=plugin plugin=%q version=%q action=compile_failed retained_incumbent=%t",
|
||||
plugin.Key,
|
||||
plugin.Version,
|
||||
retainedIncumbent,
|
||||
)
|
||||
continue
|
||||
}
|
||||
nextOverrides = append(nextOverrides, compiled)
|
||||
nextHashes[plugin.Key] = plugin.SourceHash
|
||||
delete(taskPluginSyncState.errors, plugin.Key)
|
||||
logger.LogDebug(
|
||||
ctx,
|
||||
"task_plugin subsystem=sync event=plugin plugin=%q version=%q action=compile_success",
|
||||
plugin.Key,
|
||||
plugin.Version,
|
||||
)
|
||||
}
|
||||
if err = jsplugin.DefaultRegistry.ReplaceOverrides(nextOverrides); err != nil {
|
||||
syncErr := fmt.Errorf("publish task plugin generation: %w", err)
|
||||
taskPluginSyncState.lastRebuild = taskPluginRebuildOutcome{
|
||||
Status: "failed",
|
||||
AttemptedAt: time.Now(),
|
||||
Generation: jsplugin.DefaultRegistry.Generation().Number,
|
||||
DatabaseRevision: databaseSnapshot.Revision,
|
||||
Error: syncErr.Error(),
|
||||
}
|
||||
logger.LogDebug(
|
||||
ctx,
|
||||
"task_plugin subsystem=sync event=failed stage=publish retained_generation=%d retained_generation_active=true database_revision=%q elapsed_ms=%d",
|
||||
jsplugin.DefaultRegistry.Generation().Number,
|
||||
databaseSnapshot.Revision,
|
||||
time.Since(started).Milliseconds(),
|
||||
)
|
||||
return syncErr
|
||||
}
|
||||
taskPluginSyncState.hashes = nextHashes
|
||||
for key := range taskPluginSyncState.errors {
|
||||
if !seen[key] {
|
||||
delete(taskPluginSyncState.errors, key)
|
||||
}
|
||||
}
|
||||
pluginErrors := jsplugin.DefaultRegistry.RoutingErrors()
|
||||
maps.Copy(pluginErrors, taskPluginSyncState.errors)
|
||||
pluginErrorCount := len(pluginErrors)
|
||||
status := "success"
|
||||
if pluginErrorCount > 0 {
|
||||
status = "partial"
|
||||
}
|
||||
taskPluginSyncState.lastRebuild = taskPluginRebuildOutcome{
|
||||
Status: status,
|
||||
AttemptedAt: time.Now(),
|
||||
Generation: jsplugin.DefaultRegistry.Generation().Number,
|
||||
DatabaseRevision: databaseSnapshot.Revision,
|
||||
PluginErrorCount: pluginErrorCount,
|
||||
}
|
||||
logger.LogDebug(
|
||||
ctx,
|
||||
"task_plugin subsystem=sync event=complete database_revision=%q previous_generation=%d generation=%d status=%q active_overrides=%d plugin_errors=%d elapsed_ms=%d",
|
||||
databaseSnapshot.Revision,
|
||||
generationBefore,
|
||||
jsplugin.DefaultRegistry.Generation().Number,
|
||||
status,
|
||||
len(jsplugin.DefaultRegistry.ActiveOverridePlugins()),
|
||||
pluginErrorCount,
|
||||
time.Since(started).Milliseconds(),
|
||||
)
|
||||
return nil
|
||||
}
|
||||
|
||||
func SyncTaskPluginsOnce() {
|
||||
if err := syncTaskPluginsOnce(); err != nil {
|
||||
common.SysError(err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func SyncTaskPlugins() {
|
||||
SyncTaskPluginsOnce()
|
||||
for range time.NewTicker(30 * time.Second).C {
|
||||
SyncTaskPluginsOnce()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,269 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/relay"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// taskPluginSubmitDiagnostics keeps plugin-only lifecycle logging out of the
|
||||
// ordinary task path. An empty plugin key makes every method a no-op.
|
||||
type taskPluginSubmitDiagnostics struct {
|
||||
context *gin.Context
|
||||
pluginKey string
|
||||
generation uint64
|
||||
}
|
||||
|
||||
func newTaskPluginSubmitDiagnostics(c *gin.Context) taskPluginSubmitDiagnostics {
|
||||
diagnostics := taskPluginSubmitDiagnostics{
|
||||
context: c,
|
||||
pluginKey: c.GetString("expected_task_plugin_key"),
|
||||
}
|
||||
if diagnostics.pluginKey == "" {
|
||||
return diagnostics
|
||||
}
|
||||
if pinnedValue, exists := c.Get(pluginruntime.ContextKeyPinnedPlugin); exists {
|
||||
if pinned, ok := pinnedValue.(pluginruntime.PinnedPlugin); ok && pinned.Generation != nil {
|
||||
diagnostics.generation = pinned.Generation.Number
|
||||
}
|
||||
}
|
||||
return diagnostics
|
||||
}
|
||||
|
||||
func (d taskPluginSubmitDiagnostics) start(info *relaycommon.RelayInfo) {
|
||||
if d.pluginKey == "" {
|
||||
return
|
||||
}
|
||||
logger.LogDebug(
|
||||
d.context,
|
||||
"task_plugin subsystem=submit event=start generation=%d plugin=%q model=%q action_present=%t",
|
||||
d.generation,
|
||||
d.pluginKey,
|
||||
info.OriginModelName,
|
||||
info.Action != "",
|
||||
)
|
||||
}
|
||||
|
||||
func (d taskPluginSubmitDiagnostics) refund(stage string) {
|
||||
if d.pluginKey == "" {
|
||||
return
|
||||
}
|
||||
logger.LogDebug(
|
||||
d.context,
|
||||
"task_plugin subsystem=submit event=refund_invoked generation=%d plugin=%q stage=%q durable=false",
|
||||
d.generation,
|
||||
d.pluginKey,
|
||||
stage,
|
||||
)
|
||||
}
|
||||
|
||||
func (d taskPluginSubmitDiagnostics) cancelled(stage string, attempt int) {
|
||||
if d.pluginKey == "" {
|
||||
return
|
||||
}
|
||||
logger.LogDebug(
|
||||
d.context,
|
||||
"task_plugin subsystem=submit event=cancelled generation=%d plugin=%q stage=%q attempt=%d",
|
||||
d.generation,
|
||||
d.pluginKey,
|
||||
stage,
|
||||
attempt,
|
||||
)
|
||||
}
|
||||
|
||||
func (d taskPluginSubmitDiagnostics) attempt(attempt int, channel *model.Channel, locked bool) {
|
||||
if d.pluginKey == "" || channel == nil {
|
||||
return
|
||||
}
|
||||
logger.LogDebug(
|
||||
d.context,
|
||||
"task_plugin subsystem=submit event=attempt generation=%d plugin=%q attempt=%d channel_id=%d channel_type=%d locked=%t",
|
||||
d.generation,
|
||||
d.pluginKey,
|
||||
attempt,
|
||||
channel.Id,
|
||||
channel.Type,
|
||||
locked,
|
||||
)
|
||||
}
|
||||
|
||||
func (d taskPluginSubmitDiagnostics) attemptSucceeded(attempt int, result *relay.TaskSubmitResult) {
|
||||
if d.pluginKey == "" || result == nil {
|
||||
return
|
||||
}
|
||||
logger.LogDebug(
|
||||
d.context,
|
||||
"task_plugin subsystem=submit event=attempt_succeeded generation=%d plugin=%q attempt=%d platform=%q quota=%d task_data_bytes=%d client_response=%t immediate=%t",
|
||||
d.generation,
|
||||
d.pluginKey,
|
||||
attempt,
|
||||
result.Platform,
|
||||
result.Quota,
|
||||
len(result.TaskData),
|
||||
result.ClientResponse != nil,
|
||||
result.Immediate != nil,
|
||||
)
|
||||
}
|
||||
|
||||
func (d taskPluginSubmitDiagnostics) attemptFailed(attempt int, channel *model.Channel, taskErr *dto.TaskError, willRetry bool) {
|
||||
if d.pluginKey == "" || channel == nil || taskErr == nil {
|
||||
return
|
||||
}
|
||||
logger.LogDebug(
|
||||
d.context,
|
||||
"task_plugin subsystem=submit event=attempt_failed generation=%d plugin=%q attempt=%d channel_id=%d channel_type=%d code=%q status=%d local=%t will_retry=%t",
|
||||
d.generation,
|
||||
d.pluginKey,
|
||||
attempt,
|
||||
channel.Id,
|
||||
channel.Type,
|
||||
taskErr.Code,
|
||||
taskErr.StatusCode,
|
||||
taskErr.LocalError,
|
||||
willRetry,
|
||||
)
|
||||
}
|
||||
|
||||
func (d taskPluginSubmitDiagnostics) failed(stage, reason string, taskErr *dto.TaskError, durable bool) {
|
||||
if d.pluginKey == "" {
|
||||
return
|
||||
}
|
||||
code := ""
|
||||
status := 0
|
||||
local := true
|
||||
if taskErr != nil {
|
||||
code = taskErr.Code
|
||||
status = taskErr.StatusCode
|
||||
local = taskErr.LocalError
|
||||
}
|
||||
logger.LogDebug(
|
||||
d.context,
|
||||
"task_plugin subsystem=submit event=failed generation=%d plugin=%q stage=%q reason=%q code=%q status=%d local=%t durable=%t",
|
||||
d.generation,
|
||||
d.pluginKey,
|
||||
stage,
|
||||
reason,
|
||||
code,
|
||||
status,
|
||||
local,
|
||||
durable,
|
||||
)
|
||||
}
|
||||
|
||||
func (d taskPluginSubmitDiagnostics) reserve(event string, quota int) {
|
||||
if d.pluginKey == "" {
|
||||
return
|
||||
}
|
||||
logger.LogDebug(
|
||||
d.context,
|
||||
"task_plugin subsystem=submit event=%s generation=%d plugin=%q quota=%d",
|
||||
event,
|
||||
d.generation,
|
||||
d.pluginKey,
|
||||
quota,
|
||||
)
|
||||
}
|
||||
|
||||
func (d taskPluginSubmitDiagnostics) insertStart(task *model.Task) {
|
||||
if d.pluginKey == "" || task == nil {
|
||||
return
|
||||
}
|
||||
logger.LogDebug(
|
||||
d.context,
|
||||
"task_plugin subsystem=submit event=insert_start generation=%d plugin=%q public_task_id=%q platform=%q channel_id=%d quota=%d",
|
||||
d.generation,
|
||||
d.pluginKey,
|
||||
task.TaskID,
|
||||
task.Platform,
|
||||
task.ChannelId,
|
||||
task.Quota,
|
||||
)
|
||||
}
|
||||
|
||||
func (d taskPluginSubmitDiagnostics) durable(task *model.Task) {
|
||||
if d.pluginKey == "" || task == nil {
|
||||
return
|
||||
}
|
||||
logger.LogDebug(
|
||||
d.context,
|
||||
"task_plugin subsystem=submit event=durable generation=%d plugin=%q public_task_id=%q status=%q durable=true",
|
||||
d.generation,
|
||||
d.pluginKey,
|
||||
task.TaskID,
|
||||
taskPluginDebugStatus(string(task.Status)),
|
||||
)
|
||||
}
|
||||
|
||||
func (d taskPluginSubmitDiagnostics) settleStart(task *model.Task, quota int) {
|
||||
if d.pluginKey == "" || task == nil {
|
||||
return
|
||||
}
|
||||
logger.LogDebug(
|
||||
d.context,
|
||||
"task_plugin subsystem=submit event=settle_start generation=%d plugin=%q public_task_id=%q quota=%d durable=true",
|
||||
d.generation,
|
||||
d.pluginKey,
|
||||
task.TaskID,
|
||||
quota,
|
||||
)
|
||||
}
|
||||
|
||||
func (d taskPluginSubmitDiagnostics) complete(task *model.Task, quota int) {
|
||||
if d.pluginKey == "" || task == nil {
|
||||
return
|
||||
}
|
||||
logger.LogDebug(
|
||||
d.context,
|
||||
"task_plugin subsystem=submit event=complete generation=%d plugin=%q public_task_id=%q quota=%d durable=true",
|
||||
d.generation,
|
||||
d.pluginKey,
|
||||
task.TaskID,
|
||||
quota,
|
||||
)
|
||||
}
|
||||
|
||||
func (d taskPluginSubmitDiagnostics) present(task *model.Task, presenter string) {
|
||||
if d.pluginKey == "" || task == nil {
|
||||
return
|
||||
}
|
||||
logger.LogDebug(
|
||||
d.context,
|
||||
"task_plugin subsystem=submit event=present generation=%d plugin=%q public_task_id=%q presenter=%q durable=true",
|
||||
d.generation,
|
||||
d.pluginKey,
|
||||
task.TaskID,
|
||||
presenter,
|
||||
)
|
||||
}
|
||||
|
||||
func (d taskPluginSubmitDiagnostics) presentError(taskErr *dto.TaskError) {
|
||||
if d.pluginKey == "" || taskErr == nil {
|
||||
return
|
||||
}
|
||||
logger.LogDebug(
|
||||
d.context,
|
||||
"task_plugin subsystem=submit event=present_error generation=%d plugin=%q code=%q status=%d local=%t",
|
||||
d.generation,
|
||||
d.pluginKey,
|
||||
taskErr.Code,
|
||||
taskErr.StatusCode,
|
||||
taskErr.LocalError,
|
||||
)
|
||||
}
|
||||
|
||||
func taskPluginDebugStatus(status string) string {
|
||||
switch model.TaskStatus(status) {
|
||||
case model.TaskStatusSubmitted,
|
||||
model.TaskStatusQueued,
|
||||
model.TaskStatusInProgress,
|
||||
model.TaskStatusSuccess,
|
||||
model.TaskStatusFailure:
|
||||
return status
|
||||
default:
|
||||
return "unknown"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/relay"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestTaskPluginSubmitDiagnosticsArePluginOnlyAndDoNotLogPayloads(t *testing.T) {
|
||||
previousDebug := common.DebugEnabled
|
||||
common.DebugEnabled = true
|
||||
t.Cleanup(func() { common.DebugEnabled = previousDebug })
|
||||
|
||||
var output bytes.Buffer
|
||||
common.LogWriterMu.Lock()
|
||||
previousWriter := gin.DefaultErrorWriter
|
||||
gin.DefaultErrorWriter = &output
|
||||
common.LogWriterMu.Unlock()
|
||||
t.Cleanup(func() {
|
||||
common.LogWriterMu.Lock()
|
||||
gin.DefaultErrorWriter = previousWriter
|
||||
common.LogWriterMu.Unlock()
|
||||
})
|
||||
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/tasks/debug-plugin", nil)
|
||||
c.Set(common.RequestIdKey, "plugin-submit-request")
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: "safe-model",
|
||||
TaskRelayInfo: &relaycommon.TaskRelayInfo{Action: "https://private-action.invalid/?key=hidden"},
|
||||
}
|
||||
|
||||
newTaskPluginSubmitDiagnostics(c).start(info)
|
||||
assert.Empty(t, output.String())
|
||||
|
||||
c.Set("expected_task_plugin_key", "debug-plugin")
|
||||
diagnostics := newTaskPluginSubmitDiagnostics(c)
|
||||
diagnostics.start(info)
|
||||
diagnostics.attemptSucceeded(1, &relay.TaskSubmitResult{
|
||||
UpstreamTaskID: "private-upstream-canary",
|
||||
TaskData: []byte("private-task-data-canary"),
|
||||
ClientResponse: map[string]any{"secret": "private-client-response-canary"},
|
||||
Platform: constant.TaskPlatform("debug-plugin"),
|
||||
Quota: 12,
|
||||
})
|
||||
task := &model.Task{
|
||||
TaskID: "public-task-id",
|
||||
Platform: constant.TaskPlatform("debug-plugin"),
|
||||
Status: model.TaskStatus("https://private-status.invalid/?key=hidden"),
|
||||
PrivateData: model.TaskPrivateData{
|
||||
UpstreamTaskID: "private-task-record-canary",
|
||||
ResultURL: "https://private-url.invalid/result",
|
||||
},
|
||||
}
|
||||
diagnostics.insertStart(task)
|
||||
diagnostics.durable(task)
|
||||
diagnostics.complete(task, 12)
|
||||
|
||||
logOutput := output.String()
|
||||
require.Contains(t, logOutput, "plugin-submit-request")
|
||||
assert.Contains(t, logOutput, `plugin="debug-plugin"`)
|
||||
assert.Contains(t, logOutput, `public_task_id="public-task-id"`)
|
||||
assert.Contains(t, logOutput, "task_data_bytes=24")
|
||||
assert.Contains(t, logOutput, "action_present=true")
|
||||
assert.Contains(t, logOutput, `status="unknown"`)
|
||||
for _, secret := range []string{
|
||||
"private-upstream-canary",
|
||||
"private-task-data-canary",
|
||||
"private-client-response-canary",
|
||||
"private-task-record-canary",
|
||||
"private-url.invalid",
|
||||
"private-action.invalid",
|
||||
"private-status.invalid",
|
||||
"key=hidden",
|
||||
} {
|
||||
assert.NotContains(t, logOutput, secret)
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,17 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// TelegramLegacyAuth retires the unsigned-flow widget endpoints. Existing
|
||||
// Telegram bindings are used by the unified OAuth provider instead.
|
||||
func TelegramLegacyAuth(c *gin.Context) {
|
||||
c.JSON(http.StatusGone, gin.H{
|
||||
"success": false,
|
||||
"code": "TELEGRAM_LEGACY_AUTH_REMOVED",
|
||||
"message": "Telegram login has changed. Reload the page and start Telegram OAuth again.",
|
||||
})
|
||||
}
|
||||
Loaded 100 of 2381 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user