feat: publish sanitized custom.97 source

controller/chain_detection.go

@@ -0,0 +1,156 @@
+package controller
+
+import (
+	"context"
+	"crypto/sha256"
+	"crypto/subtle"
+	"encoding/hex"
+	"errors"
+	"net/http"
+	"os"
+	"slices"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/QuantumNous/new-api/common"
+	"github.com/QuantumNous/new-api/model"
+	fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
+	"github.com/QuantumNous/new-api/service"
+	"github.com/gin-gonic/gin"
+	"gorm.io/gorm"
+)
+
+var chainEnqueue sync.Mutex
+var chainKinds = []string{"fingerprint", "ciallo", "vision"}
+
+// All three child tasks are committed together, using the existing durable ledger.
+func enqueueChainTest(ctx context.Context, group string, channelID int) ([]model.DetectionRun, error) {
+	chainEnqueue.Lock()
+	defer chainEnqueue.Unlock()
+	var runs []model.DetectionRun
+	err := model.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		var active []model.DetectionRun
+		if err := tx.Where("group_name = ? AND request_key LIKE ? AND settled = ? AND deadline > ?", group, "chain:%", false, time.Now().Unix()).Order("created_at asc").Limit(1).Find(&active).Error; err != nil {
+			return err
+		}
+		if len(active) > 0 {
+			id, _, _ := strings.Cut(active[0].ID, "_")
+			return tx.Where("id IN ?", []string{id + "_fingerprint", id + "_ciallo", id + "_vision"}).Find(&runs).Error
+		}
+		id, err := common.GenerateRandomCharsKey(32)
+		if err != nil {
+			return err
+		}
+		now := time.Now().Unix()
+		for _, kind := range chainKinds {
+			runs = append(runs, model.DetectionRun{ID: id + "_" + kind, RequestKey: "chain:" + id + ":" + kind, ChannelID: channelID, GroupName: group, Kind: kind, Status: "pending", CreatedAt: now, Deadline: now + 600})
+		}
+		return tx.Create(&runs).Error
+	})
+	return runs, err
+}
+
+func BotChainTest(c *gin.Context) {
+	c.Header("Cache-Control", "no-store")
+	key := os.Getenv("BOT_BALANCE_KEY")
+	derived := sha256.Sum256([]byte("newapi-chaintest-v1:" + key))
+	raw, bearer := strings.CutPrefix(c.GetHeader("Authorization"), "Bearer ")
+	expected, supplied := sha256.Sum256([]byte(hex.EncodeToString(derived[:]))), sha256.Sum256([]byte(raw))
+	if len(key) < 32 || !bearer || subtle.ConstantTimeCompare(expected[:], supplied[:]) != 1 {
+		c.JSON(http.StatusUnauthorized, gin.H{"success": false, "code": "unauthorized"})
+		return
+	}
+	var runs []model.DetectionRun
+	if c.Request.Method == http.MethodPost {
+		c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, 1024)
+		var input struct {
+			Group string `json:"group"`
+		}
+		if c.ShouldBindJSON(&input) != nil || strings.TrimSpace(input.Group) == "" || len(input.Group) > 128 {
+			c.JSON(400, gin.H{"success": false, "code": "invalid_group"})
+			return
+		}
+		input.Group = strings.TrimSpace(input.Group)
+		channels, err := service.AstraSamplingChannels(c.Request.Context())
+		if err != nil {
+			c.JSON(503, gin.H{"success": false, "code": "unavailable"})
+			return
+		}
+		slices.SortFunc(channels, func(a, b model.Channel) int { return a.Id - b.Id })
+		selected := 0
+		for _, channel := range channels {
+			if slices.Contains(channel.GetGroups(), input.Group) {
+				selected = channel.Id
+				break
+			}
+		}
+		if selected == 0 {
+			c.JSON(404, gin.H{"success": false, "code": "group_unavailable"})
+			return
+		}
+		runs, err = enqueueChainTest(c.Request.Context(), input.Group, selected)
+		if err != nil {
+			c.JSON(503, gin.H{"success": false, "code": "unavailable"})
+			return
+		}
+		for _, run := range runs {
+			go processDetection(context.Background(), run)
+		}
+	} else {
+		id := c.Param("id")
+		if len(id) != 32 || strings.ContainsAny(id, "_%") {
+			c.JSON(404, gin.H{"success": false, "code": "not_found"})
+			return
+		}
+		err := model.DB.WithContext(c.Request.Context()).Where("id IN ? AND request_key LIKE ?", []string{id + "_fingerprint", id + "_ciallo", id + "_vision"}, "chain:"+id+":%").Find(&runs).Error
+		if err != nil {
+			c.JSON(503, gin.H{"success": false, "code": "unavailable"})
+			return
+		}
+	}
+	data, err := chainTestResult(runs, time.Now().Unix())
+	if err != nil {
+		c.JSON(404, gin.H{"success": false, "code": "not_found"})
+		return
+	}
+	c.JSON(200, gin.H{"success": true, "data": data})
+}
+
+func chainTestResult(runs []model.DetectionRun, now int64) (gin.H, error) {
+	if len(runs) != 3 {
+		return nil, errors.New("incomplete chain")
+	}
+	id, _, _ := strings.Cut(runs[0].ID, "_")
+	results := make(map[string]any, 3)
+	done := true
+	for _, run := range runs {
+		run = service.NormalizeVisionRun(run)
+		state, text, match := "running", "", ""
+		if run.Settled {
+			state = "failed"
+			var result fingerprint.Result
+			if run.Status == "succeeded" && common.UnmarshalJsonStr(run.Result, &result) == nil {
+				text = result.Closest
+				if run.Kind == "fingerprint" {
+					match = service.AstraSamplingResultStatus(result)
+					state = match
+					if state != "match" && state != "mismatch" && state != "inconclusive" {
+						state = "failed"
+					}
+				} else if strings.TrimSpace(text) != "" {
+					state = service.VisionQuality(text)
+					if run.Kind == "ciallo" {
+						state = service.CialloQuality(text)
+					}
+				}
+			}
+		} else if now >= run.Deadline {
+			state = "failed"
+		} else {
+			done = false
+		}
+		results[run.Kind] = gin.H{"status": state, "text": text}
+	}
+	return gin.H{"id": id, "group": runs[0].GroupName, "deadline": runs[0].Deadline, "done": done, "results": results}, nil
+}