refactor: deduplicate task and security logic

This commit is contained in:
martin committed 2026-07-26 14:07:53 +08:00
1 parent 35a20cffeb
commit 92c2cc033b
11 files changed
+222 -364

No files matched your search

+4 -30
View File
@@ -2,9 +2,7 @@ package admin
import (
"context"
"encoding/json"
"fmt"
do "github.com/qianfree/team-api/internal/model/do"
"strings"
"time"
@@ -17,6 +15,7 @@ import (
"github.com/qianfree/team-api/internal/consts"
"github.com/qianfree/team-api/internal/dao"
"github.com/qianfree/team-api/internal/logic/common"
do "github.com/qianfree/team-api/internal/model/do"
"github.com/qianfree/team-api/internal/model/entity"
"github.com/qianfree/team-api/internal/utility/crypto"
)
@@ -153,7 +152,7 @@ func (s *sAdmin) Login(ctx context.Context, req *v1.AdminLoginReq) (*v1.AdminLog
}
refreshTokenHash := common.HashRefreshToken(refreshToken)
deviceInfo := extractDeviceInfo(ctx)
deviceInfo := common.ExtractDeviceInfo(ctx)
// Create session with jti
jti := common.GenerateJti()
@@ -269,7 +268,7 @@ func (s *sAdmin) Refresh(ctx context.Context, req *v1.AdminRefreshReq) (*v1.Admi
// Rotate session
ipAddress := g.RequestFromCtx(ctx).GetClientIp()
deviceInfo := extractDeviceInfo(ctx)
deviceInfo := common.ExtractDeviceInfo(ctx)
err = common.RefreshSession(ctx, session.Id, refreshTokenHash, newRefreshTokenHash, ipAddress, deviceInfo)
if err != nil {
return nil, err
@@ -301,14 +300,7 @@ func (s *sAdmin) Refresh(ctx context.Context, req *v1.AdminRefreshReq) (*v1.Admi
// ListSessions returns active sessions for the current admin user.
func (s *sAdmin) ListSessions(ctx context.Context, req *v1.AdminSessionListReq) (*v1.AdminSessionListRes, error) {
currentSessionID := common.GetCtxSessionID(ctx)
page := req.Page
pageSize := req.PageSize
if page < 1 {
page = 1
}
if pageSize < 1 || pageSize > 100 {
pageSize = 20
}
page, pageSize := common.NormalizePagination(req.Page, req.PageSize)
q := dao.SysSessions.Ctx(ctx).
Where("user_type", "admin").
@@ -455,21 +447,3 @@ func (s *sAdmin) ChangePassword(ctx context.Context, req *v1.AdminChangePassword
return nil, nil
}
// extractDeviceInfo extracts device information from the request and returns it as a JSON string.
// The sys_sessions.device_info column is JSONB, so this must be valid JSON.
func extractDeviceInfo(ctx context.Context) string {
r := g.RequestFromCtx(ctx)
if r == nil {
return `{"user_agent":"unknown"}`
}
ua := r.Header.Get("User-Agent")
if len(ua) > 500 {
ua = ua[:500]
}
if ua == "" {
ua = "unknown"
}
b, _ := json.Marshal(map[string]string{"user_agent": ua})
return string(b)
}
+3 -27
View File
@@ -61,7 +61,7 @@ func (s *sAdmin) Verify2FA(ctx context.Context, req *v1.Admin2FAVerifyReq) (*v1.
refreshTokenHash := common.HashRefreshToken(refreshToken)
ipAddress := g.RequestFromCtx(ctx).GetClientIp()
deviceInfo := extractDeviceInfo(ctx)
deviceInfo := common.ExtractDeviceInfo(ctx)
// Create session with jti
jti := common.GenerateJti()
@@ -171,24 +171,7 @@ func (s *sAdmin) RegenerateBackupCodes(ctx context.Context, req *v1.Admin2FARege
// ConfirmHighRisk generates a confirm token for high-risk operations.
func (s *sAdmin) ConfirmHighRisk(ctx context.Context, req *v1.Admin2FAConfirmReq) (*v1.Admin2FAConfirmRes, error) {
userID := common.GetCtxUserID(ctx)
enabled, err := common.Is2FAEnabled(ctx, "admin", userID)
if err != nil {
return nil, err
}
if !enabled {
return nil, common.NewBusinessError(consts.CodeTotpNotEnabled, consts.MsgTotpNotEnabled)
}
valid, err := common.Verify2FACode(ctx, "admin", userID, req.Code)
if err != nil {
return nil, err
}
if !valid {
return nil, common.NewBusinessError(consts.CodeTotpInvalid, consts.MsgTotpInvalid)
}
token, err := common.GenerateConfirmToken(ctx, userID, "admin")
token, err := common.ConfirmHighRisk(ctx, "admin", userID, req.Code)
if err != nil {
return nil, err
}
@@ -198,14 +181,7 @@ func (s *sAdmin) ConfirmHighRisk(ctx context.Context, req *v1.Admin2FAConfirmReq
// LoginHistory returns the login history for all admin users with search filters.
func (s *sAdmin) LoginHistory(ctx context.Context, req *v1.AdminLoginHistoryReq) (*v1.AdminLoginHistoryRes, error) {
page := req.Page
pageSize := req.PageSize
if page < 1 {
page = 1
}
if pageSize < 1 || pageSize > 100 {
pageSize = 20
}
page, pageSize := common.NormalizePagination(req.Page, req.PageSize)
q := common.AuditModelCtx(ctx, "aud_login_history").Where("user_type", "admin")
+1 -8
View File
@@ -70,14 +70,7 @@ func (s *sAdmin) UsageLogCleanupCreate(ctx context.Context, req *v1.UsageLogClea
}
func (s *sAdmin) UsageLogCleanupList(ctx context.Context, req *v1.UsageLogCleanupListReq) (*v1.UsageLogCleanupListRes, error) {
page := req.Page
pageSize := req.PageSize
if page < 1 {
page = 1
}
if pageSize < 1 || pageSize > 100 {
pageSize = 20
}
page, pageSize := common.NormalizePagination(req.Page, req.PageSize)
m := dao.TskTasks.Ctx(ctx).
Where("handler", "usage_log_cleanup").
+42
View File
@@ -14,6 +14,7 @@ import (
"github.com/gogf/gf/v2/frame/g"
"github.com/golang-jwt/jwt/v5"
"github.com/qianfree/team-api/internal/consts"
"github.com/qianfree/team-api/internal/dao"
do "github.com/qianfree/team-api/internal/model/do"
"github.com/qianfree/team-api/internal/model/entity"
@@ -139,6 +140,27 @@ func ParseConfirmToken(tokenStr string) (*ConfirmClaims, error) {
return claims, nil
}
// ConfirmHighRisk 校验高风险操作的二次验证码并签发短期确认令牌。
func ConfirmHighRisk(ctx context.Context, userType string, userID int64, code string) (string, error) {
enabled, err := Is2FAEnabled(ctx, userType, userID)
if err != nil {
return "", err
}
if !enabled {
return "", NewBusinessError(consts.CodeTotpNotEnabled, consts.MsgTotpNotEnabled)
}
valid, err := Verify2FACode(ctx, userType, userID, code)
if err != nil {
return "", err
}
if !valid {
return "", NewBusinessError(consts.CodeTotpInvalid, consts.MsgTotpInvalid)
}
return GenerateConfirmToken(ctx, userID, userType)
}
// Setup2FA generates a new TOTP secret for a user. Returns the secret and otpauth URI.
// Callers must still persist the secret after the user confirms with a valid code.
func Setup2FA(ctx context.Context, userType string, userID int64) (secret, uri string, err error) {
@@ -468,6 +490,26 @@ func DeviceFingerprint(ua, ip string) string {
return hash[:32]
}
// ExtractDeviceInfo 从请求上下文提取会话设备信息并编码为合法 JSON。
func ExtractDeviceInfo(ctx context.Context) string {
r := g.RequestFromCtx(ctx)
if r == nil {
return buildDeviceInfo("")
}
return buildDeviceInfo(r.Header.Get("User-Agent"))
}
func buildDeviceInfo(userAgent string) string {
if len(userAgent) > 500 {
userAgent = userAgent[:500]
}
if userAgent == "" {
userAgent = "unknown"
}
data, _ := json.Marshal(map[string]string{"user_agent": userAgent})
return string(data)
}
// verifyAndConsumeBackupCode checks if code matches any backup code and marks it consumed.
func verifyAndConsumeBackupCode(ctx context.Context, userType string, userID int64, code string) (bool, error) {
_, backupCodes, err := GetUserTOTPSecret(ctx, userType, userID)
+22
View File
@@ -1,6 +1,8 @@
package common
import (
"context"
"encoding/json"
"strings"
"testing"
"time"
@@ -249,6 +251,26 @@ func TestDeviceFingerprint_Empty(t *testing.T) {
}
}
func TestExtractDeviceInfoWithoutRequest(t *testing.T) {
var info map[string]string
if err := json.Unmarshal([]byte(ExtractDeviceInfo(context.Background())), &info); err != nil {
t.Fatalf("ExtractDeviceInfo returned invalid JSON: %v", err)
}
if info["user_agent"] != "unknown" {
t.Fatalf("user_agent = %q, want unknown", info["user_agent"])
}
}
func TestBuildDeviceInfoTruncatesUserAgent(t *testing.T) {
var info map[string]string
if err := json.Unmarshal([]byte(buildDeviceInfo(strings.Repeat("a", 501))), &info); err != nil {
t.Fatalf("buildDeviceInfo returned invalid JSON: %v", err)
}
if len(info["user_agent"]) != 500 {
t.Fatalf("user_agent length = %d, want 500", len(info["user_agent"]))
}
}
func TestBackupCodeSHA256(t *testing.T) {
const code = "ABCD-EFGH"
hashed := hashBackupCode(code)
+69 -239
View File
@@ -9,11 +9,11 @@ import (
"github.com/gogf/gf/v2/errors/gerror"
"github.com/gogf/gf/v2/frame/g"
"github.com/gogf/gf/v2/os/gtime"
"github.com/shopspring/decimal"
"github.com/qianfree/team-api/internal/dao"
"github.com/qianfree/team-api/internal/logic/relay"
do "github.com/qianfree/team-api/internal/model/do"
"github.com/qianfree/team-api/internal/model/entity"
"github.com/qianfree/team-api/internal/utility/crypto"
"github.com/qianfree/team-api/relay/common"
)
@@ -23,6 +23,64 @@ type AsyncProvider struct{}
var DefaultAsyncProvider = &AsyncProvider{}
func asyncTaskFromEntity(row *entity.TskModelTasks) *common.AsyncTask {
if row == nil {
return nil
}
return &common.AsyncTask{
ID: row.Id,
PublicTaskID: row.PublicTaskId,
RequestID: row.RequestId,
Platform: row.Platform,
Action: row.Action,
Status: row.Status,
Progress: row.Progress,
FailReason: row.FailReason,
TenantID: row.TenantId,
UserID: row.UserId,
ApiKeyID: row.ApiKeyId,
ChannelID: row.ChannelId,
ModelName: row.ModelName,
UpstreamModel: row.UpstreamModel,
PreDeductAmount: row.PreDeductAmount,
ActualCost: row.ActualCost,
BillingSettled: row.BillingSettled,
ResultURL: row.ResultUrl,
Data: json.RawMessage(row.Data),
PrivateData: json.RawMessage(row.PrivateData),
SubmitTime: timePtrFromGTime(row.SubmitTime),
StartTime: timePtrFromGTime(row.StartTime),
FinishTime: timePtrFromGTime(row.FinishTime),
CreatedAt: timeFromGTime(row.CreatedAt),
UpdatedAt: timeFromGTime(row.UpdatedAt),
}
}
func asyncTasksFromEntities(rows []*entity.TskModelTasks) []*common.AsyncTask {
tasks := make([]*common.AsyncTask, 0, len(rows))
for _, row := range rows {
if task := asyncTaskFromEntity(row); task != nil {
tasks = append(tasks, task)
}
}
return tasks
}
func timePtrFromGTime(value *gtime.Time) *time.Time {
if value == nil {
return nil
}
t := value.Time
return &t
}
func timeFromGTime(value *gtime.Time) time.Time {
if value == nil {
return time.Time{}
}
return value.Time
}
// CreateTask 创建异步任务记录
func (p *AsyncProvider) CreateTask(ctx context.Context, task *common.AsyncTask) error {
g.Log().Infof(ctx, "[AsyncProvider] CreateTask start: public_id=%s, tenant=%d, user=%d, model=%s, platform=%s",
@@ -131,100 +189,19 @@ func (p *AsyncProvider) UpdateTaskCAS(ctx context.Context, task *common.AsyncTas
// GetTaskByPublicID 根据公开任务 ID 查询
func (p *AsyncProvider) GetTaskByPublicID(ctx context.Context, publicTaskID string) (*common.AsyncTask, error) {
var row *struct {
ID int64 `json:"id"`
PublicTaskID string `json:"public_task_id"`
RequestId string `json:"request_id"`
Platform string `json:"platform"`
Action string `json:"action"`
Status string `json:"status"`
Progress string `json:"progress"`
FailReason string `json:"fail_reason"`
TenantID int64 `json:"tenant_id"`
UserID int64 `json:"user_id"`
ApiKeyID int64 `json:"api_key_id"`
ChannelID int64 `json:"channel_id"`
ModelName string `json:"model_name"`
UpstreamModel string `json:"upstream_model"`
PreDeductAmount decimal.Decimal `json:"pre_deduct_amount"`
ActualCost decimal.Decimal `json:"actual_cost"`
BillingSettled bool `json:"billing_settled"`
ResultURL string `json:"result_url"`
Data json.RawMessage `json:"data"`
PrivateData json.RawMessage `json:"private_data"`
SubmitTime *time.Time `json:"submit_time"`
StartTime *time.Time `json:"start_time"`
FinishTime *time.Time `json:"finish_time"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
var row *entity.TskModelTasks
err := dao.TskModelTasks.Ctx(ctx).
Where("public_task_id", publicTaskID).
Scan(&row)
if err != nil {
return nil, gerror.Wrapf(err, "query async task failed: public_id=%s", publicTaskID)
}
if row == nil {
return nil, nil
}
return &common.AsyncTask{
ID: row.ID,
PublicTaskID: row.PublicTaskID,
RequestID: row.RequestId,
Platform: row.Platform,
Action: row.Action,
Status: row.Status,
Progress: row.Progress,
FailReason: row.FailReason,
TenantID: row.TenantID,
UserID: row.UserID,
ApiKeyID: row.ApiKeyID,
ChannelID: row.ChannelID,
ModelName: row.ModelName,
UpstreamModel: row.UpstreamModel,
PreDeductAmount: row.PreDeductAmount,
ActualCost: row.ActualCost,
BillingSettled: row.BillingSettled,
ResultURL: row.ResultURL,
Data: row.Data,
PrivateData: row.PrivateData,
SubmitTime: row.SubmitTime,
StartTime: row.StartTime,
FinishTime: row.FinishTime,
CreatedAt: row.CreatedAt,
UpdatedAt: row.UpdatedAt,
}, nil
return asyncTaskFromEntity(row), nil
}
// GetTaskByPublicIDAndUser 根据公开任务 ID + 用户 ID 查询
func (p *AsyncProvider) GetTaskByPublicIDAndUser(ctx context.Context, publicTaskID string, userID int64, tenantID int64) (*common.AsyncTask, error) {
var row *struct {
ID int64 `json:"id"`
PublicTaskID string `json:"public_task_id"`
RequestId string `json:"request_id"`
Platform string `json:"platform"`
Action string `json:"action"`
Status string `json:"status"`
Progress string `json:"progress"`
FailReason string `json:"fail_reason"`
TenantID int64 `json:"tenant_id"`
UserID int64 `json:"user_id"`
ApiKeyID int64 `json:"api_key_id"`
ChannelID int64 `json:"channel_id"`
ModelName string `json:"model_name"`
UpstreamModel string `json:"upstream_model"`
PreDeductAmount decimal.Decimal `json:"pre_deduct_amount"`
ActualCost decimal.Decimal `json:"actual_cost"`
BillingSettled bool `json:"billing_settled"`
ResultURL string `json:"result_url"`
Data json.RawMessage `json:"data"`
PrivateData json.RawMessage `json:"private_data"`
SubmitTime *time.Time `json:"submit_time"`
StartTime *time.Time `json:"start_time"`
FinishTime *time.Time `json:"finish_time"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
var row *entity.TskModelTasks
err := dao.TskModelTasks.Ctx(ctx).
Where("public_task_id", publicTaskID).
Where("user_id", userID).
@@ -233,62 +210,12 @@ func (p *AsyncProvider) GetTaskByPublicIDAndUser(ctx context.Context, publicTask
if err != nil {
return nil, gerror.Wrapf(err, "query async task failed: public_id=%s user_id=%d", publicTaskID, userID)
}
if row == nil {
return nil, nil
}
return &common.AsyncTask{
ID: row.ID,
PublicTaskID: row.PublicTaskID,
RequestID: row.RequestId,
Platform: row.Platform,
Action: row.Action,
Status: row.Status,
Progress: row.Progress,
FailReason: row.FailReason,
TenantID: row.TenantID,
UserID: row.UserID,
ApiKeyID: row.ApiKeyID,
ChannelID: row.ChannelID,
ModelName: row.ModelName,
UpstreamModel: row.UpstreamModel,
PreDeductAmount: row.PreDeductAmount,
ActualCost: row.ActualCost,
BillingSettled: row.BillingSettled,
ResultURL: row.ResultURL,
Data: row.Data,
PrivateData: row.PrivateData,
SubmitTime: row.SubmitTime,
StartTime: row.StartTime,
FinishTime: row.FinishTime,
CreatedAt: row.CreatedAt,
UpdatedAt: row.UpdatedAt,
}, nil
return asyncTaskFromEntity(row), nil
}
// GetNonTerminalTasks 获取所有非终态任务
func (p *AsyncProvider) GetNonTerminalTasks(ctx context.Context, limit int) ([]*common.AsyncTask, error) {
var rows []struct {
ID int64 `json:"id"`
PublicTaskID string `json:"public_task_id"`
RequestId string `json:"request_id"`
Platform string `json:"platform"`
Action string `json:"action"`
Status string `json:"status"`
Progress string `json:"progress"`
TenantID int64 `json:"tenant_id"`
UserID int64 `json:"user_id"`
ApiKeyID int64 `json:"api_key_id"`
ChannelID int64 `json:"channel_id"`
ModelName string `json:"model_name"`
UpstreamModel string `json:"upstream_model"`
PreDeductAmount decimal.Decimal `json:"pre_deduct_amount"`
ActualCost decimal.Decimal `json:"actual_cost"`
BillingSettled bool `json:"billing_settled"`
Data json.RawMessage `json:"data"`
PrivateData json.RawMessage `json:"private_data"`
SubmitTime *time.Time `json:"submit_time"`
CreatedAt time.Time `json:"created_at"`
}
var rows []*entity.TskModelTasks
err := dao.TskModelTasks.Ctx(ctx).
Where("status NOT IN (?, ?)", "SUCCESS", "FAILURE").
Order("submit_time ASC").
@@ -298,52 +225,13 @@ func (p *AsyncProvider) GetNonTerminalTasks(ctx context.Context, limit int) ([]*
return nil, gerror.Wrapf(err, "query non-terminal tasks failed")
}
tasks := make([]*common.AsyncTask, 0, len(rows))
for _, r := range rows {
tasks = append(tasks, &common.AsyncTask{
ID: r.ID,
PublicTaskID: r.PublicTaskID,
RequestID: r.RequestId,
Platform: r.Platform,
Action: r.Action,
Status: r.Status,
Progress: r.Progress,
TenantID: r.TenantID,
UserID: r.UserID,
ApiKeyID: r.ApiKeyID,
ChannelID: r.ChannelID,
ModelName: r.ModelName,
UpstreamModel: r.UpstreamModel,
PreDeductAmount: r.PreDeductAmount,
ActualCost: r.ActualCost,
BillingSettled: r.BillingSettled,
Data: r.Data,
PrivateData: r.PrivateData,
SubmitTime: r.SubmitTime,
CreatedAt: r.CreatedAt,
})
}
return tasks, nil
return asyncTasksFromEntities(rows), nil
}
// GetTimedOutTasks 获取超时未完成任务
func (p *AsyncProvider) GetTimedOutTasks(ctx context.Context, cutoffUnix int64, limit int) ([]*common.AsyncTask, error) {
cutoffTime := time.Unix(cutoffUnix, 0)
var rows []struct {
ID int64 `json:"id"`
PublicTaskID string `json:"public_task_id"`
RequestId string `json:"request_id"`
Platform string `json:"platform"`
Status string `json:"status"`
TenantID int64 `json:"tenant_id"`
UserID int64 `json:"user_id"`
ApiKeyID int64 `json:"api_key_id"`
ChannelID int64 `json:"channel_id"`
ModelName string `json:"model_name"`
PreDeductAmount decimal.Decimal `json:"pre_deduct_amount"`
PrivateData json.RawMessage `json:"private_data"`
SubmitTime *time.Time `json:"submit_time"`
}
var rows []*entity.TskModelTasks
err := dao.TskModelTasks.Ctx(ctx).
Where("status NOT IN (?, ?)", "SUCCESS", "FAILURE").
Where("submit_time < ?", cutoffTime).
@@ -354,48 +242,12 @@ func (p *AsyncProvider) GetTimedOutTasks(ctx context.Context, cutoffUnix int64,
return nil, gerror.Wrapf(err, "query timed-out tasks failed")
}
tasks := make([]*common.AsyncTask, 0, len(rows))
for _, r := range rows {
tasks = append(tasks, &common.AsyncTask{
ID: r.ID,
PublicTaskID: r.PublicTaskID,
RequestID: r.RequestId,
Platform: r.Platform,
Status: r.Status,
TenantID: r.TenantID,
UserID: r.UserID,
ApiKeyID: r.ApiKeyID,
ChannelID: r.ChannelID,
ModelName: r.ModelName,
PreDeductAmount: r.PreDeductAmount,
PrivateData: r.PrivateData,
SubmitTime: r.SubmitTime,
})
}
return tasks, nil
return asyncTasksFromEntities(rows), nil
}
// GetUnsettledTasks 获取终态但未结算的任务(用于结算重试)
func (p *AsyncProvider) GetUnsettledTasks(ctx context.Context, limit int) ([]*common.AsyncTask, error) {
var rows []struct {
ID int64 `json:"id"`
PublicTaskID string `json:"public_task_id"`
RequestId string `json:"request_id"`
Platform string `json:"platform"`
Status string `json:"status"`
TenantID int64 `json:"tenant_id"`
UserID int64 `json:"user_id"`
ApiKeyID int64 `json:"api_key_id"`
ChannelID int64 `json:"channel_id"`
ModelName string `json:"model_name"`
PreDeductAmount decimal.Decimal `json:"pre_deduct_amount"`
ActualCost decimal.Decimal `json:"actual_cost"`
BillingSettled bool `json:"billing_settled"`
Data json.RawMessage `json:"data"`
PrivateData json.RawMessage `json:"private_data"`
SubmitTime *time.Time `json:"submit_time"`
CreatedAt time.Time `json:"created_at"`
}
var rows []*entity.TskModelTasks
err := dao.TskModelTasks.Ctx(ctx).
Where("status IN (?, ?)", "SUCCESS", "FAILURE").
Where("billing_settled = ?", false).
@@ -407,29 +259,7 @@ func (p *AsyncProvider) GetUnsettledTasks(ctx context.Context, limit int) ([]*co
return nil, gerror.Wrapf(err, "query unsettled tasks failed")
}
tasks := make([]*common.AsyncTask, 0, len(rows))
for _, r := range rows {
tasks = append(tasks, &common.AsyncTask{
ID: r.ID,
PublicTaskID: r.PublicTaskID,
RequestID: r.RequestId,
Platform: r.Platform,
Status: r.Status,
TenantID: r.TenantID,
UserID: r.UserID,
ApiKeyID: r.ApiKeyID,
ChannelID: r.ChannelID,
ModelName: r.ModelName,
PreDeductAmount: r.PreDeductAmount,
ActualCost: r.ActualCost,
BillingSettled: r.BillingSettled,
Data: r.Data,
PrivateData: r.PrivateData,
SubmitTime: r.SubmitTime,
CreatedAt: r.CreatedAt,
})
}
return tasks, nil
return asyncTasksFromEntities(rows), nil
}
// GetChannelByID 获取渠道基本信息(含从 chn_channel_keys 解密的 API Key)
@@ -0,0 +1,72 @@
package task
import (
"testing"
"time"
"github.com/gogf/gf/v2/os/gtime"
"github.com/shopspring/decimal"
"github.com/qianfree/team-api/internal/model/entity"
)
func TestAsyncTaskFromEntity(t *testing.T) {
createdAt := time.Date(2026, 7, 26, 10, 30, 0, 0, time.UTC)
updatedAt := createdAt.Add(time.Minute)
row := &entity.TskModelTasks{
Id: 1,
PublicTaskId: "task_public",
RequestId: "req_public",
Platform: "sora",
Action: "generate",
Status: "IN_PROGRESS",
Progress: "50%",
FailReason: "temporary",
TenantId: 2,
UserId: 3,
ApiKeyId: 4,
ChannelId: 5,
ModelName: "requested-model",
UpstreamModel: "upstream-model",
PreDeductAmount: decimal.RequireFromString("0.25"),
ActualCost: decimal.RequireFromString("0.10"),
BillingSettled: true,
ResultUrl: "https://example.com/result",
Data: `{"result":"ok"}`,
PrivateData: `{"upstream_task_id":"upstream-1"}`,
SubmitTime: gtime.NewFromTime(createdAt),
StartTime: gtime.NewFromTime(createdAt.Add(time.Second)),
FinishTime: gtime.NewFromTime(createdAt.Add(2 * time.Second)),
CreatedAt: gtime.NewFromTime(createdAt),
UpdatedAt: gtime.NewFromTime(updatedAt),
}
task := asyncTaskFromEntity(row)
if task == nil {
t.Fatal("asyncTaskFromEntity returned nil")
}
if task.ID != row.Id || task.PublicTaskID != row.PublicTaskId || task.RequestID != row.RequestId {
t.Fatalf("identity fields not mapped: %+v", task)
}
if task.Action != row.Action || task.FailReason != row.FailReason || task.ResultURL != row.ResultUrl {
t.Fatalf("task detail fields not mapped: %+v", task)
}
if !task.PreDeductAmount.Equal(row.PreDeductAmount) || !task.ActualCost.Equal(row.ActualCost) || !task.BillingSettled {
t.Fatalf("billing fields not mapped: %+v", task)
}
if string(task.Data) != row.Data || string(task.PrivateData) != row.PrivateData {
t.Fatalf("JSON fields not mapped: data=%s private_data=%s", task.Data, task.PrivateData)
}
if task.SubmitTime == nil || !task.SubmitTime.Equal(createdAt) || !task.CreatedAt.Equal(createdAt) || !task.UpdatedAt.Equal(updatedAt) {
t.Fatalf("time fields not mapped: %+v", task)
}
}
func TestAsyncTaskFromEntityNil(t *testing.T) {
if task := asyncTaskFromEntity(nil); task != nil {
t.Fatalf("asyncTaskFromEntity(nil) = %+v, want nil", task)
}
if tasks := asyncTasksFromEntities([]*entity.TskModelTasks{nil}); len(tasks) != 0 {
t.Fatalf("asyncTasksFromEntities retained nil rows: %+v", tasks)
}
}
+4 -24
View File
@@ -2,7 +2,6 @@ package tenant
import (
"context"
"encoding/json"
"fmt"
"math"
"strings"
@@ -18,10 +17,9 @@ import (
"github.com/qianfree/team-api/internal/dao"
"github.com/qianfree/team-api/internal/logic/common"
"github.com/qianfree/team-api/internal/middleware"
do "github.com/qianfree/team-api/internal/model/do"
"github.com/qianfree/team-api/internal/model/entity"
"github.com/qianfree/team-api/internal/utility/crypto"
do "github.com/qianfree/team-api/internal/model/do"
)
// Register handles tenant registration.
@@ -224,7 +222,7 @@ func (s *sTenant) Register(ctx context.Context, req *v1.TenantRegisterReq) (*v1.
}
refreshTokenHash := common.HashRefreshToken(refreshToken)
deviceInfo := extractTenantDeviceInfo(ctx)
deviceInfo := common.ExtractDeviceInfo(ctx)
jti := common.GenerateJti()
sessionID, err := common.CreateSession(ctx, "tenant", ownerUserID, tenantID, refreshTokenHash, ipAddress, deviceInfo, jti)
if err != nil {
@@ -469,7 +467,7 @@ func (s *sTenant) Login(ctx context.Context, req *v1.TenantLoginReq) (*v1.Tenant
refreshTokenHash := common.HashRefreshToken(refreshToken)
// ipAddress already declared above for 2FA check
deviceInfo := extractTenantDeviceInfo(ctx)
deviceInfo := common.ExtractDeviceInfo(ctx)
jti := common.GenerateJti()
sessionID, err := common.CreateSession(ctx, "tenant", user.Id, tenant.Id, refreshTokenHash, ipAddress, deviceInfo, jti)
if err != nil {
@@ -573,7 +571,7 @@ func (s *sTenant) Refresh(ctx context.Context, req *v1.TenantRefreshReq) (*v1.Te
newRefreshTokenHash := common.HashRefreshToken(newRefreshToken)
ipAddress := g.RequestFromCtx(ctx).GetClientIp()
deviceInfo := extractTenantDeviceInfo(ctx)
deviceInfo := common.ExtractDeviceInfo(ctx)
err = common.RefreshSession(ctx, session.Id, refreshTokenHash, newRefreshTokenHash, ipAddress, deviceInfo)
if err != nil {
return nil, err
@@ -693,21 +691,3 @@ func (s *sTenant) RevokeSession(ctx context.Context, req *v1.TenantRevokeSession
}
return nil, nil
}
// extractTenantDeviceInfo extracts device information from the request and returns it as a JSON string.
// The sys_sessions.device_info column is JSONB, so this must be valid JSON.
func extractTenantDeviceInfo(ctx context.Context) string {
r := g.RequestFromCtx(ctx)
if r == nil {
return `{"user_agent":"unknown"}`
}
ua := r.Header.Get("User-Agent")
if len(ua) > 500 {
ua = ua[:500]
}
if ua == "" {
ua = "unknown"
}
b, _ := json.Marshal(map[string]string{"user_agent": ua})
return string(b)
}
+1 -1
View File
@@ -355,7 +355,7 @@ func (s *sTenant) JoinByInvite(ctx context.Context, req *v1.TenantMemberJoinReq)
refreshTokenHash := common.HashRefreshToken(refreshToken)
ipAddress := g.RequestFromCtx(ctx).GetClientIp()
deviceInfo := extractTenantDeviceInfo(ctx)
deviceInfo := common.ExtractDeviceInfo(ctx)
jti := common.GenerateJti()
sessionID, err := common.CreateSession(ctx, "tenant", userID, tenantID, refreshTokenHash, ipAddress, deviceInfo, jti)
if err != nil {
+2 -16
View File
@@ -42,14 +42,7 @@ func (s *sTenant) OpenAppList(ctx context.Context, req *v1.OpenAppListReq) (*v1.
return nil, common.NewForbiddenError("需要 owner 或 admin 权限")
}
tenantID := middleware.GetTenantID(ctx)
page := req.Page
pageSize := req.PageSize
if page < 1 {
page = 1
}
if pageSize < 1 || pageSize > 100 {
pageSize = 20
}
page, pageSize := common.NormalizePagination(req.Page, req.PageSize)
m := dao.OpnApps.Ctx(ctx).Where("tenant_id", tenantID)
if req.Keyword != "" {
@@ -489,14 +482,7 @@ func (s *sTenant) WebhookDeliveryLogs(ctx context.Context, req *v1.WebhookDelive
return nil, common.NewForbiddenError("需要 owner 或 admin 权限")
}
tenantID := middleware.GetTenantID(ctx)
page := req.Page
pageSize := req.PageSize
if page < 1 {
page = 1
}
if pageSize < 1 || pageSize > 100 {
pageSize = 20
}
page, pageSize := common.NormalizePagination(req.Page, req.PageSize)
// 计数查询(不 JOIN)
countQuery := dao.OpnWebhookDeliveryLogs.Ctx(ctx).Where("tenant_id", tenantID).Where("webhook_config_id", req.Id)
+2 -19
View File
@@ -60,7 +60,7 @@ func (s *sTenant) Verify2FA(ctx context.Context, req *v1.Tenant2FAVerifyReq) (*v
refreshTokenHash := common.HashRefreshToken(refreshToken)
ipAddress := g.RequestFromCtx(ctx).GetClientIp()
deviceInfo := extractTenantDeviceInfo(ctx)
deviceInfo := common.ExtractDeviceInfo(ctx)
jti := common.GenerateJti()
sessionID, err := common.CreateSession(ctx, "tenant", user.Id, user.TenantId, refreshTokenHash, ipAddress, deviceInfo, jti)
@@ -168,24 +168,7 @@ func (s *sTenant) RegenerateBackupCodes(ctx context.Context, req *v1.Tenant2FARe
// ConfirmHighRisk generates a confirm token for high-risk operations.
func (s *sTenant) ConfirmHighRisk(ctx context.Context, req *v1.Tenant2FAConfirmReq) (*v1.Tenant2FAConfirmRes, error) {
userID := middleware.GetUserID(ctx)
enabled, err := common.Is2FAEnabled(ctx, "tenant", userID)
if err != nil {
return nil, err
}
if !enabled {
return nil, common.NewBusinessError(consts.CodeTotpNotEnabled, consts.MsgTotpNotEnabled)
}
valid, err := common.Verify2FACode(ctx, "tenant", userID, req.Code)
if err != nil {
return nil, err
}
if !valid {
return nil, common.NewBusinessError(consts.CodeTotpInvalid, consts.MsgTotpInvalid)
}
token, err := common.GenerateConfirmToken(ctx, userID, "tenant")
token, err := common.ConfirmHighRisk(ctx, "tenant", userID, req.Code)
if err != nil {
return nil, err
}