Skip to content
Merged
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions common/utils.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import (
"math/big"
"math/rand"
"net"
"net/url"
"os"
"os/exec"
"runtime"
Expand Down Expand Up @@ -284,3 +285,20 @@ func GetAudioDuration(ctx context.Context, filename string, ext string) (float64
}
return strconv.ParseFloat(durationStr, 64)
}

// BuildURL concatenates base and endpoint, returns the complete url string
func BuildURL(base string, endpoint string) string {
u, err := url.Parse(base)
if err != nil {
return base + endpoint
}
end := endpoint
if end == "" {
end = "/"
}
ref, err := url.Parse(end)
if err != nil {
return base + endpoint
}
return u.ResolveReference(ref).String()
}
24 changes: 24 additions & 0 deletions controller/ratio_config.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
package controller

import (
"net/http"
"one-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": ratio_setting.GetExposedData(),
})
}
293 changes: 293 additions & 0 deletions controller/ratio_sync.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,293 @@
package controller

import (
"encoding/json"
"net/http"
"one-api/model"
"one-api/setting/ratio_setting"
"one-api/dto"
"sync"
"time"

"github.com/gin-gonic/gin"
)

type upstreamResult struct {
Name string `json:"name"`
Data map[string]any `json:"data,omitempty"`
Err string `json:"err,omitempty"`
}

type TestResult struct {
Name string `json:"name"`
Status string `json:"status"`
Error string `json:"error,omitempty"`
}

type DifferenceItem struct {
Current interface{} `json:"current"` // 当前本地值,可能为null
Upstreams map[string]interface{} `json:"upstreams"` // 上游值:具体值/"same"/null
}

type SyncableChannel struct {
ID int `json:"id"`
Name string `json:"name"`
BaseURL string `json:"base_url"`
Status int `json:"status"`
}

func FetchUpstreamRatios(c *gin.Context) {
var req dto.UpstreamRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()})
return
}

if req.Timeout <= 0 {
req.Timeout = 10
}

var upstreams []dto.UpstreamDTO
if len(req.ChannelIDs) > 0 {
intIds := make([]int, 0, len(req.ChannelIDs))
for _, id64 := range req.ChannelIDs {
intIds = append(intIds, int(id64))
}
dbChannels, _ := model.GetChannelsByIds(intIds)
for _, ch := range dbChannels {
upstreams = append(upstreams, dto.UpstreamDTO{
Name: ch.Name,
BaseURL: ch.GetBaseURL(),
Endpoint: "",
})
}
}

var wg sync.WaitGroup
ch := make(chan upstreamResult, len(upstreams))

for _, chn := range upstreams {
wg.Add(1)
go func(chItem dto.UpstreamDTO) {
defer wg.Done()
endpoint := chItem.Endpoint
if endpoint == "" {
endpoint = "/api/ratio_config"
}
url := chItem.BaseURL + endpoint
client := http.Client{Timeout: time.Duration(req.Timeout) * time.Second}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Set reasonable timeout limits.

The timeout is user-controlled without upper bounds, which could lead to resource exhaustion.

if req.Timeout <= 0 {
    req.Timeout = 10
}
+if req.Timeout > 60 {
+    req.Timeout = 60  // Maximum 60 seconds
+}

Committable suggestion skipped: line range outside the PR's diff.

🤖 Prompt for AI Agents
In controller/ratio_sync.go at line 78, the HTTP client timeout is set directly
from a user-controlled value without any upper limit, which risks resource
exhaustion. Add a maximum timeout limit by checking if req.Timeout exceeds a
predefined safe threshold and, if so, cap it to that maximum value before
setting the client timeout.

resp, err := client.Get(url)
if err != nil {
ch <- upstreamResult{Name: chItem.Name, Err: err.Error()}
return
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
ch <- upstreamResult{Name: chItem.Name, Err: resp.Status}
return
}
var body struct {
Success bool `json:"success"`
Data map[string]any `json:"data"`
Message string `json:"message"`
}
if err := json.NewDecoder(resp.Body).Decode(&body); err != nil {
ch <- upstreamResult{Name: chItem.Name, Err: err.Error()}
return
}
if !body.Success {
ch <- upstreamResult{Name: chItem.Name, Err: body.Message}
return
}
ch <- upstreamResult{Name: chItem.Name, Data: body.Data}
}(chn)
}
Comment on lines +69 to +134

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Address security vulnerabilities in HTTP requests.

The code constructs URLs directly from user input and makes HTTP requests without proper validation, creating potential SSRF (Server-Side Request Forgery) vulnerabilities.

+import (
+    "net/url"
+    "strings"
+)

+func isValidURL(rawURL string) bool {
+    u, err := url.Parse(rawURL)
+    if err != nil {
+        return false
+    }
+    // Only allow http/https schemes
+    if u.Scheme != "http" && u.Scheme != "https" {
+        return false
+    }
+    // Prevent access to local/private networks
+    if strings.Contains(u.Host, "localhost") || 
+       strings.Contains(u.Host, "127.0.0.1") ||
+       strings.Contains(u.Host, "0.0.0.0") {
+        return false
+    }
+    return true
+}

go func(chItem dto.UpstreamDTO) {
    defer wg.Done()
    endpoint := chItem.Endpoint
    if endpoint == "" {
        endpoint = "/api/ratio_config"
    }
-    url := chItem.BaseURL + endpoint
+    fullURL := chItem.BaseURL + endpoint
+    if !isValidURL(fullURL) {
+        ch <- upstreamResult{Name: chItem.Name, Err: "invalid or unsafe URL"}
+        return
+    }
    client := http.Client{Timeout: time.Duration(req.Timeout) * time.Second}
-    resp, err := client.Get(url)
+    resp, err := client.Get(fullURL)
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
for _, chn := range upstreams {
wg.Add(1)
go func(chItem dto.UpstreamDTO) {
defer wg.Done()
endpoint := chItem.Endpoint
if endpoint == "" {
endpoint = "/api/ratio_config"
}
url := chItem.BaseURL + endpoint
client := http.Client{Timeout: time.Duration(req.Timeout) * time.Second}
resp, err := client.Get(url)
if err != nil {
ch <- upstreamResult{Name: chItem.Name, Err: err.Error()}
return
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
ch <- upstreamResult{Name: chItem.Name, Err: resp.Status}
return
}
var body struct {
Success bool `json:"success"`
Data map[string]any `json:"data"`
Message string `json:"message"`
}
if err := json.NewDecoder(resp.Body).Decode(&body); err != nil {
ch <- upstreamResult{Name: chItem.Name, Err: err.Error()}
return
}
if !body.Success {
ch <- upstreamResult{Name: chItem.Name, Err: body.Message}
return
}
ch <- upstreamResult{Name: chItem.Name, Data: body.Data}
}(chn)
}
// At the top of controller/ratio_sync.go, add the new imports:
import (
// … other imports …
"net/url"
"strings"
)
// Add this helper right after your imports:
func isValidURL(rawURL string) bool {
u, err := url.Parse(rawURL)
if err != nil {
return false
}
// Only allow http/https schemes
if u.Scheme != "http" && u.Scheme != "https" {
return false
}
// Prevent access to local/private networks
if strings.Contains(u.Host, "localhost") ||
strings.Contains(u.Host, "127.0.0.1") ||
strings.Contains(u.Host, "0.0.0.0") {
return false
}
return true
}
// Then, in your existing loop (lines 69–104), update to:
for _, chn := range upstreams {
wg.Add(1)
go func(chItem dto.UpstreamDTO) {
defer wg.Done()
endpoint := chItem.Endpoint
if endpoint == "" {
endpoint = "/api/ratio_config"
}
fullURL := chItem.BaseURL + endpoint
if !isValidURL(fullURL) {
ch <- upstreamResult{Name: chItem.Name, Err: "invalid or unsafe URL"}
return
}
client := http.Client{Timeout: time.Duration(req.Timeout) * time.Second}
resp, err := client.Get(fullURL)
if err != nil {
ch <- upstreamResult{Name: chItem.Name, Err: err.Error()}
return
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
ch <- upstreamResult{Name: chItem.Name, Err: resp.Status}
return
}
var body struct {
Success bool `json:"success"`
Data map[string]any `json:"data"`
Message string `json:"message"`
}
if err := json.NewDecoder(resp.Body).Decode(&body); err != nil {
ch <- upstreamResult{Name: chItem.Name, Err: err.Error()}
return
}
if !body.Success {
ch <- upstreamResult{Name: chItem.Name, Err: body.Message}
return
}
ch <- upstreamResult{Name: chItem.Name, Data: body.Data}
}(chn)
}
🤖 Prompt for AI Agents
In controller/ratio_sync.go around lines 69 to 104, the code constructs URLs
directly from user input and makes HTTP requests without validating the URLs,
which can lead to SSRF vulnerabilities. To fix this, validate and sanitize the
BaseURL and endpoint before making the HTTP request. Use URL parsing to ensure
the URL is well-formed, restrict requests to allowed domains or IP ranges, and
avoid concatenating strings directly to form URLs. Implement checks to reject or
handle unsafe URLs before proceeding with the HTTP client call.


wg.Wait()
close(ch)

localData := ratio_setting.GetExposedData()

var testResults []dto.TestResult
var successfulChannels []struct {
name string
data map[string]any
}

for r := range ch {
if r.Err != "" {
testResults = append(testResults, dto.TestResult{
Name: r.Name,
Status: "error",
Error: r.Err,
})
} else {
testResults = append(testResults, dto.TestResult{
Name: r.Name,
Status: "success",
})
successfulChannels = append(successfulChannels, struct {
name string
data map[string]any
}{name: r.Name, data: r.Data})
}
}

differences := buildDifferences(localData, successfulChannels)

c.JSON(http.StatusOK, gin.H{
"success": true,
"data": gin.H{
"differences": differences,
"test_results": testResults,
},
})
}

func buildDifferences(localData map[string]any, successfulChannels []struct {
name string
data map[string]any
}) map[string]map[string]dto.DifferenceItem {
differences := make(map[string]map[string]dto.DifferenceItem)
ratioTypes := []string{"model_ratio", "completion_ratio", "cache_ratio", "model_price"}

allModels := make(map[string]struct{})

for _, ratioType := range ratioTypes {
if localRatioAny, ok := localData[ratioType]; ok {
if localRatio, ok := localRatioAny.(map[string]float64); ok {
for modelName := range localRatio {
allModels[modelName] = struct{}{}
}
}
}
}

for _, channel := range successfulChannels {
for _, ratioType := range ratioTypes {
if upstreamRatio, ok := channel.data[ratioType].(map[string]any); ok {
for modelName := range upstreamRatio {
allModels[modelName] = struct{}{}
}
}
}
}

for modelName := range allModels {
for _, ratioType := range ratioTypes {
var localValue interface{} = nil
if localRatioAny, ok := localData[ratioType]; ok {
if localRatio, ok := localRatioAny.(map[string]float64); ok {
if val, exists := localRatio[modelName]; exists {
localValue = val
}
}
}

upstreamValues := make(map[string]interface{})
hasUpstreamValue := false
hasDifference := false

for _, channel := range successfulChannels {
var upstreamValue interface{} = nil

if upstreamRatio, ok := channel.data[ratioType].(map[string]any); ok {
if val, exists := upstreamRatio[modelName]; exists {
upstreamValue = val
hasUpstreamValue = true

if localValue != nil && localValue != val {
hasDifference = true
} else if localValue == val {
upstreamValue = "same"
}
Comment on lines +228 to +232

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Fix potential type comparison issue.

The comparison localValue != val may not work correctly for floating-point numbers due to precision issues, and the type assertion could fail silently.

+import "math"

+func compareFloatValues(a, b interface{}) bool {
+    aFloat, aOk := a.(float64)
+    bFloat, bOk := b.(float64)
+    if aOk && bOk {
+        return math.Abs(aFloat-bFloat) < 1e-9  // Use epsilon for float comparison
+    }
+    return a == b  // Fallback to direct comparison
+}

-if localValue != nil && localValue != val {
+if localValue != nil && !compareFloatValues(localValue, val) {
    hasDifference = true
-} else if localValue == val {
+} else if compareFloatValues(localValue, val) {
    upstreamValue = "same"
}
🤖 Prompt for AI Agents
In controller/ratio_sync.go around lines 199 to 203, the code compares
localValue and val directly, which can cause issues with floating-point
precision and silent type assertion failures. To fix this, ensure both values
are asserted to the correct numeric type safely, then compare them using a
tolerance threshold for floating-point numbers instead of direct equality. This
will prevent incorrect difference detection due to minor precision errors.

}
}
if upstreamValue == nil && localValue == nil {
upstreamValue = "same"
}

if localValue == nil && upstreamValue != nil && upstreamValue != "same" {
hasDifference = true
}

upstreamValues[channel.name] = upstreamValue
}

shouldInclude := false

if localValue != nil {
if hasDifference {
shouldInclude = true
}
} else {
if hasUpstreamValue {
shouldInclude = true
}
}

if shouldInclude {
if differences[modelName] == nil {
differences[modelName] = make(map[string]dto.DifferenceItem)
}
differences[modelName][ratioType] = dto.DifferenceItem{
Current: localValue,
Upstreams: upstreamValues,
}
}
}
}

channelHasDiff := make(map[string]bool)
for _, ratioMap := range differences {
for _, item := range ratioMap {
for chName, val := range item.Upstreams {
if val != nil && val != "same" {
channelHasDiff[chName] = true
}
}
}
}

for modelName, ratioMap := range differences {
for ratioType, item := range ratioMap {
for chName := range item.Upstreams {
if !channelHasDiff[chName] {
delete(item.Upstreams, chName)
}
}
differences[modelName][ratioType] = item
}
}

return differences
}
Comment on lines +147 to +293

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Refactor complex buildDifferences function for maintainability.

The buildDifferences function is overly complex with deeply nested logic, making it difficult to understand, test, and maintain. Consider breaking it into smaller, focused functions.

+func collectAllModels(localData map[string]any, channels []struct {
+    name string
+    data map[string]any
+}, ratioTypes []string) map[string]struct{} {
+    allModels := make(map[string]struct{})
+    
+    // Collect from local data
+    for _, ratioType := range ratioTypes {
+        if localRatioAny, ok := localData[ratioType]; ok {
+            if localRatio, ok := localRatioAny.(map[string]float64); ok {
+                for modelName := range localRatio {
+                    allModels[modelName] = struct{}{}
+                }
+            }
+        }
+    }
+    
+    // Collect from upstream data
+    for _, channel := range channels {
+        for _, ratioType := range ratioTypes {
+            if upstreamRatio, ok := channel.data[ratioType].(map[string]any); ok {
+                for modelName := range upstreamRatio {
+                    allModels[modelName] = struct{}{}
+                }
+            }
+        }
+    }
+    
+    return allModels
+}

+func compareModelRatio(modelName, ratioType string, localData map[string]any, 
+    channels []struct {
+        name string
+        data map[string]any
+    }) (dto.DifferenceItem, bool) {
+    // Extract comparison logic here
+    // Return the difference item and whether it should be included
+}

func buildDifferences(localData map[string]any, successfulChannels []struct {
    name string
    data map[string]any
}) map[string]map[string]dto.DifferenceItem {
    differences := make(map[string]map[string]dto.DifferenceItem)
    ratioTypes := []string{"model_ratio", "completion_ratio", "cache_ratio", "model_price"}

-    allModels := make(map[string]struct{})
-    // ... existing complex logic
+    allModels := collectAllModels(localData, successfulChannels, ratioTypes)
    
    for modelName := range allModels {
        for _, ratioType := range ratioTypes {
-            // ... existing complex comparison logic
+            if item, shouldInclude := compareModelRatio(modelName, ratioType, localData, successfulChannels); shouldInclude {
+                if differences[modelName] == nil {
+                    differences[modelName] = make(map[string]dto.DifferenceItem)
+                }
+                differences[modelName][ratioType] = item
+            }
        }
    }

+    return filterEmptyChannels(differences)
+}

+func filterEmptyChannels(differences map[string]map[string]dto.DifferenceItem) map[string]map[string]dto.DifferenceItem {
+    // Extract the channel filtering logic here

Committable suggestion skipped: line range outside the PR's diff.

🤖 Prompt for AI Agents
In controller/ratio_sync.go from lines 147 to 264, the buildDifferences function
is too complex and deeply nested. Refactor by splitting it into smaller helper
functions that each handle a specific task, such as extracting all model names,
comparing local and upstream values, and filtering differences. This will
improve readability, testability, and maintainability.


func GetSyncableChannels(c *gin.Context) {
channels, err := model.GetAllChannels(0, 0, true, false)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": err.Error(),
})
return
}

var syncableChannels []dto.SyncableChannel
for _, channel := range channels {
if channel.GetBaseURL() != "" {
syncableChannels = append(syncableChannels, dto.SyncableChannel{
ID: channel.Id,
Name: channel.Name,
BaseURL: channel.GetBaseURL(),
Status: channel.Status,
})
}
}

c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "",
"data": syncableChannels,
})
}
Comment on lines +295 to +322

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Add input validation and improve error handling.

The GetSyncableChannels function lacks input validation and returns success even when database errors occur.

func GetSyncableChannels(c *gin.Context) {
    channels, err := model.GetAllChannels(0, 0, true, false)
    if err != nil {
-        c.JSON(http.StatusOK, gin.H{
+        c.JSON(http.StatusInternalServerError, gin.H{
            "success": false,
            "message": err.Error(),
        })
        return
    }

    var syncableChannels []dto.SyncableChannel
    for _, channel := range channels {
-        if channel.GetBaseURL() != "" {
+        baseURL := channel.GetBaseURL()
+        if baseURL != "" && isValidURL(baseURL) {  // Reuse the validation function
            syncableChannels = append(syncableChannels, dto.SyncableChannel{
                ID:      channel.Id,
                Name:    channel.Name,
-                BaseURL: channel.GetBaseURL(),
+                BaseURL: baseURL,
                Status:  channel.Status,
            })
        }
    }

    c.JSON(http.StatusOK, gin.H{
        "success": true,
-        "message": "",
+        "message": "Successfully retrieved syncable channels",
        "data":    syncableChannels,
    })
}
🤖 Prompt for AI Agents
In controller/ratio_sync.go around lines 266 to 293, the GetSyncableChannels
function lacks input validation and improperly returns success status even when
database errors occur. Add validation for any input parameters received from the
gin.Context before processing. Modify the error handling to return an
appropriate HTTP error status code (e.g., 500) instead of http.StatusOK when a
database error occurs, and set "success" to false in the JSON response. Ensure
the function only returns success true when no errors happen and valid data is
returned.

Loading