diff --git a/heicode/controller/discord.go b/heicode/controller/discord.go index 46ac0e6..55f3108 100644 --- a/heicode/controller/discord.go +++ b/heicode/controller/discord.go @@ -101,7 +101,8 @@ func getDiscordUserInfoByCode(code string) (*DiscordUser, error) { func DiscordOAuth(c *gin.Context) { session := sessions.Default(c) state := c.Query("state") - if state == "" || session.Get("oauth_state") == nil || state != session.Get("oauth_state").(string) { + storedState, _ := session.Get("oauth_state").(string) + if state == "" || storedState == "" || state != storedState { c.JSON(http.StatusForbidden, gin.H{ "success": false, "message": "state is empty or not same", diff --git a/heicode/controller/github.go b/heicode/controller/github.go index 02fd26c..466a2d0 100644 --- a/heicode/controller/github.go +++ b/heicode/controller/github.go @@ -81,7 +81,8 @@ func getGitHubUserInfoByCode(code string) (*GitHubUser, error) { func GitHubOAuth(c *gin.Context) { session := sessions.Default(c) state := c.Query("state") - if state == "" || session.Get("oauth_state") == nil || state != session.Get("oauth_state").(string) { + storedState, _ := session.Get("oauth_state").(string) + if state == "" || storedState == "" || state != storedState { c.JSON(http.StatusForbidden, gin.H{ "success": false, "message": "state is empty or not same", @@ -140,10 +141,9 @@ func GitHubOAuth(c *gin.Context) { user.Email = githubUser.Email user.Role = common.RoleCommonUser user.Status = common.UserStatusEnabled - affCode := session.Get("aff") inviterId := 0 - if affCode != nil { - inviterId, _ = model.GetUserIdByAffCode(affCode.(string)) + if affCodeStr, ok := session.Get("aff").(string); ok && affCodeStr != "" { + inviterId, _ = model.GetUserIdByAffCode(affCodeStr) } if err := user.Insert(inviterId); err != nil { diff --git a/heicode/controller/linuxdo.go b/heicode/controller/linuxdo.go index 725bf68..84659ec 100644 --- a/heicode/controller/linuxdo.go +++ b/heicode/controller/linuxdo.go @@ -168,7 +168,8 @@ func LinuxdoOAuth(c *gin.Context) { } state := c.Query("state") - if state == "" || session.Get("oauth_state") == nil || state != session.Get("oauth_state").(string) { + storedState, _ := session.Get("oauth_state").(string) + if state == "" || storedState == "" || state != storedState { c.JSON(http.StatusForbidden, gin.H{ "success": false, "message": "state is empty or not same", @@ -226,10 +227,9 @@ func LinuxdoOAuth(c *gin.Context) { user.Role = common.RoleCommonUser user.Status = common.UserStatusEnabled - affCode := session.Get("aff") inviterId := 0 - if affCode != nil { - inviterId, _ = model.GetUserIdByAffCode(affCode.(string)) + if affCodeStr, ok := session.Get("aff").(string); ok && affCodeStr != "" { + inviterId, _ = model.GetUserIdByAffCode(affCodeStr) } if err := user.Insert(inviterId); err != nil { diff --git a/heicode/controller/oauth.go b/heicode/controller/oauth.go index 8d7f216..25963e2 100644 --- a/heicode/controller/oauth.go +++ b/heicode/controller/oauth.go @@ -56,7 +56,8 @@ func HandleOAuth(c *gin.Context) { // 1. Validate state (CSRF protection) state := c.Query("state") - if state == "" || session.Get("oauth_state") == nil || state != session.Get("oauth_state").(string) { + storedState, _ := session.Get("oauth_state").(string) + if state == "" || storedState == "" || state != storedState { c.JSON(http.StatusForbidden, gin.H{ "success": false, "message": i18n.T(c, i18n.MsgOAuthStateInvalid), @@ -263,10 +264,9 @@ func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *o user.Status = common.UserStatusEnabled // Handle affiliate code - affCode := session.Get("aff") inviterId := 0 - if affCode != nil { - inviterId, _ = model.GetUserIdByAffCode(affCode.(string)) + if affCodeStr, ok := session.Get("aff").(string); ok && affCodeStr != "" { + inviterId, _ = model.GetUserIdByAffCode(affCodeStr) } // Use transaction to ensure user creation and OAuth binding are atomic diff --git a/heicode/controller/oidc.go b/heicode/controller/oidc.go index f28788c..e35e83f 100644 --- a/heicode/controller/oidc.go +++ b/heicode/controller/oidc.go @@ -103,7 +103,8 @@ func getOidcUserInfoByCode(code string) (*OidcUser, error) { func OidcAuth(c *gin.Context) { session := sessions.Default(c) state := c.Query("state") - if state == "" || session.Get("oauth_state") == nil || state != session.Get("oauth_state").(string) { + storedState, _ := session.Get("oauth_state").(string) + if state == "" || storedState == "" || state != storedState { c.JSON(http.StatusForbidden, gin.H{ "success": false, "message": "state is empty or not same", diff --git a/heicode/middleware/auth.go b/heicode/middleware/auth.go index a4943f9..83759b3 100644 --- a/heicode/middleware/auth.go +++ b/heicode/middleware/auth.go @@ -112,7 +112,8 @@ func authHelper(c *gin.Context, minRole int) { return } - if id != apiUserId { + idVal, _ := id.(int) + if idVal != apiUserId { c.JSON(http.StatusUnauthorized, gin.H{ "success": false, "message": common.TranslateMessage(c, i18n.MsgAuthUserIdMismatch), diff --git a/heicode/middleware/distributor.go b/heicode/middleware/distributor.go index 2607882..597ef1d 100644 --- a/heicode/middleware/distributor.go +++ b/heicode/middleware/distributor.go @@ -37,7 +37,8 @@ func Distribute() func(c *gin.Context) { return } if ok { - id, err := strconv.Atoi(channelId.(string)) + channelIdStr, _ := channelId.(string) + id, err := strconv.Atoi(channelIdStr) if err != nil { abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidChannelId)) return diff --git a/heicode/middleware/turnstile-check.go b/heicode/middleware/turnstile-check.go index 1dc94c1..fe1336f 100644 --- a/heicode/middleware/turnstile-check.go +++ b/heicode/middleware/turnstile-check.go @@ -72,6 +72,7 @@ func TurnstileCheck() gin.HandlerFunc { "message": "无法保存会话信息，请重试", "success": false, }) + c.Abort() return } } diff --git a/heicode/model/channel.go b/heicode/model/channel.go index 7f10580..d196ccc 100644 --- a/heicode/model/channel.go +++ b/heicode/model/channel.go @@ -74,7 +74,18 @@ func (c ChannelInfo) Value() (driver.Value, error) { // Scan implements sql.Scanner interface func (c *ChannelInfo) Scan(value interface{}) error { - bytesValue, _ := value.([]byte) + if value == nil { + return nil + } + var bytesValue []byte + switch v := value.(type) { + case []byte: + bytesValue = v + case string: + bytesValue = []byte(v) + default: + return fmt.Errorf("cannot scan %T into ChannelInfo", value) + } return common.Unmarshal(bytesValue, c) } @@ -562,7 +573,10 @@ func CleanupChannelPollingLocks() { } channelPollingLocks.Range(func(key, value interface{}) bool { - channelId := key.(int) + channelId, ok := key.(int) + if !ok { + return true + } if !activeChannelSet[channelId] { channelPollingLocks.Delete(channelId) } diff --git a/heicode/model/task.go b/heicode/model/task.go index 8011c8b..2f1bbbe 100644 --- a/heicode/model/task.go +++ b/heicode/model/task.go @@ -4,6 +4,7 @@ import ( "bytes" "database/sql/driver" "encoding/json" + "fmt" "time" "github.com/heicode/manager/common" @@ -81,7 +82,19 @@ type Properties struct { } func (m *Properties) Scan(val interface{}) error { - bytesValue, _ := val.([]byte) + if val == nil { + *m = Properties{} + return nil + } + var bytesValue []byte + switch v := val.(type) { + case []byte: + bytesValue = v + case string: + bytesValue = []byte(v) + default: + return fmt.Errorf("cannot scan %T into Properties", val) + } if len(bytesValue) == 0 { *m = Properties{} return nil @@ -142,7 +155,18 @@ func GenerateTaskID() string { } func (p *TaskPrivateData) Scan(val interface{}) error { - bytesValue, _ := val.([]byte) + if val == nil { + return nil + } + var bytesValue []byte + switch v := val.(type) { + case []byte: + bytesValue = v + case string: + bytesValue = []byte(v) + default: + return fmt.Errorf("cannot scan %T into TaskPrivateData", val) + } if len(bytesValue) == 0 { return nil } diff --git a/heicode/relay/audio_handler.go b/heicode/relay/audio_handler.go index 29e370b..e90226c 100644 --- a/heicode/relay/audio_handler.go +++ b/heicode/relay/audio_handler.go @@ -67,10 +67,14 @@ func AudioHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type service.ResetStatusCode(newAPIError, statusCodeMappingStr) return newAPIError } - if usage.(*dto.Usage).CompletionTokenDetails.AudioTokens > 0 || usage.(*dto.Usage).PromptTokensDetails.AudioTokens > 0 { - service.PostAudioConsumeQuota(c, info, usage.(*dto.Usage), "") + usageData, ok := usage.(*dto.Usage) + if !ok { + return types.NewError(fmt.Errorf("invalid usage type"), types.ErrorCodeBadResponseBody) + } + if usageData.CompletionTokenDetails.AudioTokens > 0 || usageData.PromptTokensDetails.AudioTokens > 0 { + service.PostAudioConsumeQuota(c, info, usageData, "") } else { - service.PostTextConsumeQuota(c, info, usage.(*dto.Usage), nil) + service.PostTextConsumeQuota(c, info, usageData, nil) } return nil diff --git a/heicode/relay/channel/aws/relay-aws.go b/heicode/relay/channel/aws/relay-aws.go index 84936bc..ee6a2c0 100644 --- a/heicode/relay/channel/aws/relay-aws.go +++ b/heicode/relay/channel/aws/relay-aws.go @@ -225,7 +225,11 @@ func awsHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types ctx, cancel := newAwsInvokeContext() defer cancel() - awsResp, err := a.AwsClient.InvokeModel(ctx, a.AwsReq.(*bedrockruntime.InvokeModelInput)) + invokeInput, ok := a.AwsReq.(*bedrockruntime.InvokeModelInput) + if !ok { + return types.NewOpenAIError(errors.New("invalid AWS request type"), types.ErrorCodeAwsInvokeError, http.StatusInternalServerError), nil + } + awsResp, err := a.AwsClient.InvokeModel(ctx, invokeInput) if err != nil { statusCode := getAwsErrorStatusCode(err) return types.NewOpenAIError(errors.Wrap(err, "InvokeModel"), types.ErrorCodeAwsInvokeError, statusCode), nil @@ -255,7 +259,11 @@ func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) ( ctx, cancel := newAwsInvokeContext() defer cancel() - awsResp, err := a.AwsClient.InvokeModelWithResponseStream(ctx, a.AwsReq.(*bedrockruntime.InvokeModelWithResponseStreamInput)) + streamInput, ok := a.AwsReq.(*bedrockruntime.InvokeModelWithResponseStreamInput) + if !ok { + return types.NewOpenAIError(errors.New("invalid AWS stream request type"), types.ErrorCodeAwsInvokeError, http.StatusInternalServerError), nil + } + awsResp, err := a.AwsClient.InvokeModelWithResponseStream(ctx, streamInput) if err != nil { statusCode := getAwsErrorStatusCode(err) return types.NewOpenAIError(errors.Wrap(err, "InvokeModelWithResponseStream"), types.ErrorCodeAwsInvokeError, statusCode), nil @@ -298,7 +306,11 @@ func handleNovaRequest(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) ctx, cancel := newAwsInvokeContext() defer cancel() - awsResp, err := a.AwsClient.InvokeModel(ctx, a.AwsReq.(*bedrockruntime.InvokeModelInput)) + novaInput, ok := a.AwsReq.(*bedrockruntime.InvokeModelInput) + if !ok { + return types.NewOpenAIError(errors.New("invalid AWS Nova request type"), types.ErrorCodeAwsInvokeError, http.StatusInternalServerError), nil + } + awsResp, err := a.AwsClient.InvokeModel(ctx, novaInput) if err != nil { statusCode := getAwsErrorStatusCode(err) return types.NewOpenAIError(errors.Wrap(err, "InvokeModel"), types.ErrorCodeAwsInvokeError, statusCode), nil diff --git a/heicode/relay/channel/claude/relay-claude.go b/heicode/relay/channel/claude/relay-claude.go index 3e3d7fe..f7fb130 100644 --- a/heicode/relay/channel/claude/relay-claude.go +++ b/heicode/relay/channel/claude/relay-claude.go @@ -248,7 +248,9 @@ func RequestOpenAI2ClaudeMessage(c *gin.Context, textRequest dto.GeneralOpenAIRe case []interface{}: stopSequences := make([]string, 0) for _, stop := range textRequest.Stop.([]interface{}) { - stopSequences = append(stopSequences, stop.(string)) + if s, ok := stop.(string); ok { + stopSequences = append(stopSequences, s) + } } claudeRequest.StopSequences = stopSequences } diff --git a/heicode/relay/channel/ollama/adaptor.go b/heicode/relay/channel/ollama/adaptor.go index e158ad0..d8c9bc3 100644 --- a/heicode/relay/channel/ollama/adaptor.go +++ b/heicode/relay/channel/ollama/adaptor.go @@ -29,11 +29,14 @@ func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayIn if err != nil { return nil, err } - openaiRequest.(*dto.GeneralOpenAIRequest).StreamOptions = &dto.StreamOptions{ + oaiReq, ok := openaiRequest.(*dto.GeneralOpenAIRequest) + if !ok { + return nil, errors.New("unexpected request type from Claude conversion") + } + oaiReq.StreamOptions = &dto.StreamOptions{ IncludeUsage: true, } - // map to ollama chat request (Claude -> OpenAI -> Ollama chat) - return openAIChatToOllamaChat(c, openaiRequest.(*dto.GeneralOpenAIRequest)) + return openAIChatToOllamaChat(c, oaiReq) } func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) { diff --git a/heicode/relay/channel/zhipu/relay-zhipu.go b/heicode/relay/channel/zhipu/relay-zhipu.go index 54f72e7..a59770c 100644 --- a/heicode/relay/channel/zhipu/relay-zhipu.go +++ b/heicode/relay/channel/zhipu/relay-zhipu.go @@ -32,7 +32,7 @@ var expSeconds int64 = 24 * 3600 func getZhipuToken(apikey string) string { data, ok := zhipuTokens.Load(apikey) if ok { - tokenData := data.(zhipuTokenData) + tokenData, _ := data.(zhipuTokenData) if time.Now().Before(tokenData.ExpiryTime) { return tokenData.Token } diff --git a/heicode/service/convert.go b/heicode/service/convert.go index 4ce82ab..7a8c7f6 100644 --- a/heicode/service/convert.go +++ b/heicode/service/convert.go @@ -753,7 +753,11 @@ func GeminiToOpenAIRequest(geminiRequest *dto.GeminiChatRequest, info *relaycomm } // gemini stop sequences 最多 5 个，openai stop 最多 4 个 if len(geminiRequest.GenerationConfig.StopSequences) > 0 { - openaiRequest.Stop = geminiRequest.GenerationConfig.StopSequences[:4] + maxStop := 4 + if len(geminiRequest.GenerationConfig.StopSequences) < maxStop { + maxStop = len(geminiRequest.GenerationConfig.StopSequences) + } + openaiRequest.Stop = geminiRequest.GenerationConfig.StopSequences[:maxStop] } if geminiRequest.GenerationConfig.CandidateCount != nil && *geminiRequest.GenerationConfig.CandidateCount > 0 { openaiRequest.N = lo.ToPtr(*geminiRequest.GenerationConfig.CandidateCount) diff --git a/heicode/service/file_service.go b/heicode/service/file_service.go index 2674918..8bd7fb1 100644 --- a/heicode/service/file_service.go +++ b/heicode/service/file_service.go @@ -144,9 +144,11 @@ func registerSourceForCleanup(c *gin.Context, source types.FileSource) { func CleanupFileSources(c *gin.Context) { key := string(constant.ContextKeyFileSourcesToCleanup) if sources, exists := c.Get(key); exists { - for _, source := range sources.([]types.FileSource) { - if cache := source.GetCache(); cache != nil { - cache.Close() + if sourcesSlice, ok := sources.([]types.FileSource); ok { + for _, source := range sourcesSlice { + if cache := source.GetCache(); cache != nil { + cache.Close() + } } } c.Set(key, nil)