119 lines
3.9 KiB
Go
119 lines
3.9 KiB
Go
package controller
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"crypto/subtle"
|
|
"encoding/hex"
|
|
"fmt"
|
|
"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"
|
|
)
|
|
|
|
var cialloEnqueue sync.Mutex
|
|
|
|
func scheduleCialloDetections(ctx context.Context) error {
|
|
cialloEnqueue.Lock()
|
|
defer cialloEnqueue.Unlock()
|
|
channels, schedules, err := service.AutomaticDetectionChannelsDue(ctx, "ciallo", time.Now())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for _, channel := range channels {
|
|
groups := channel.GetGroups()
|
|
if len(groups) == 0 {
|
|
continue
|
|
}
|
|
id, err := common.GenerateRandomCharsKey(32)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
run := model.DetectionRun{ID: id, RequestKey: fmt.Sprintf("ciallo:%d:%s", channel.Id, schedules[channel.Id].LastResultID), ChannelID: channel.Id, GroupName: groups[0], Kind: "ciallo", Status: "pending", CreatedAt: time.Now().Unix()}
|
|
if err := model.CreateDetection(&run); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// The derived capability only authorizes Ciallo probes, not balance queries.
|
|
func BotCiallo(c *gin.Context) {
|
|
c.Header("Cache-Control", "no-store")
|
|
key := os.Getenv("BOT_BALANCE_KEY")
|
|
derived := sha256.Sum256([]byte("newapi-ciallo-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 run model.DetectionRun
|
|
if c.Request.Method == http.MethodGet {
|
|
if model.DB.Where("id = ? AND kind = ?", c.Param("id"), "ciallo").First(&run).Error != nil {
|
|
c.JSON(http.StatusNotFound, gin.H{"success": false, "code": "not_found"})
|
|
return
|
|
}
|
|
} else {
|
|
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(http.StatusBadRequest, 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
|
|
}
|
|
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
|
|
}
|
|
cialloEnqueue.Lock()
|
|
var active []model.DetectionRun
|
|
err = model.DB.Where("channel_id = ? AND kind = ? AND settled = ?", selected, "ciallo", false).Order("created_at asc").Limit(1).Find(&active).Error
|
|
if err == nil && len(active) > 0 {
|
|
run = active[0]
|
|
} else if err == nil {
|
|
id, idErr := common.GenerateRandomCharsKey(32)
|
|
err = idErr
|
|
if err == nil {
|
|
run = model.DetectionRun{ID: id, RequestKey: "bot-ciallo:" + id, ChannelID: selected, GroupName: input.Group, Kind: "ciallo", Status: "pending", CreatedAt: time.Now().Unix(), Deadline: time.Now().Unix() + 200}
|
|
err = model.CreateDetection(&run)
|
|
}
|
|
}
|
|
cialloEnqueue.Unlock()
|
|
if err != nil {
|
|
c.JSON(503, gin.H{"success": false, "code": "unavailable"})
|
|
return
|
|
}
|
|
// Start immediately; the durable worker reconciles interruptions and settlement.
|
|
if len(active) == 0 {
|
|
go processDetection(context.Background(), run)
|
|
}
|
|
}
|
|
var result fingerprint.Result
|
|
_ = common.UnmarshalJsonStr(run.Result, &result)
|
|
c.JSON(http.StatusOK, gin.H{"success": true, "data": gin.H{"id": run.ID, "status": run.Status, "settled": run.Settled, "quality": run.Quality, "text": result.Closest}})
|
|
}
|