server: Go 服务器（REST+WSS 网关+JWT+Argon2+配对+素材直链+任务状态机）；修复并发下线 send-on-closed-channel、下发查询 SQL 优先级、配对码原子占用、上传体积上限、JWT 默认密钥告警
This commit is contained in:
Qiufeng
@@ -0,0 +1,250 @@
|
||||
package ws
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/ed25519"
|
||||
"crypto/rand"
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"log"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/coder/websocket"
|
||||
"github.com/gin-gonic/gin"
|
||||
"gorm.io/gorm"
|
||||
|
||||
"everypublish/server/internal/config"
|
||||
"everypublish/server/internal/models"
|
||||
"everypublish/server/internal/task"
|
||||
"everypublish/shared/proto"
|
||||
)
|
||||
|
||||
// signPayload 签名原文：nonce + "|" + ts（两端一致）
|
||||
func signPayload(nonce string, ts int64) []byte {
|
||||
return []byte(nonce + "|" + strconv.FormatInt(ts, 10))
|
||||
}
|
||||
|
||||
// HandleAgent 处理 Agent WSS 连接：hello 验签 → 注册 → 消息循环
|
||||
func HandleAgent(cfg *config.Config, db *gorm.DB, hub *Hub, serverKey ed25519.PrivateKey) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
conn, err := websocket.Accept(c.Writer, c.Request, &websocket.AcceptOptions{InsecureSkipVerify: true})
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
// 首条消息必须是 hello（10s 超时）
|
||||
ctx, cancel := context.WithTimeout(c.Request.Context(), 10*time.Second)
|
||||
defer cancel()
|
||||
typ, raw, err := conn.Read(ctx)
|
||||
if err != nil {
|
||||
_ = conn.Close(websocket.StatusPolicyViolation, "hello timeout")
|
||||
return
|
||||
}
|
||||
if typ != websocket.MessageText {
|
||||
_ = conn.Close(websocket.StatusPolicyViolation, "expect text")
|
||||
return
|
||||
}
|
||||
var helloEnv proto.Envelope
|
||||
if err = json.Unmarshal(raw, &helloEnv); err != nil || helloEnv.Type != proto.TypeHello {
|
||||
_ = conn.Close(websocket.StatusPolicyViolation, "expect hello")
|
||||
return
|
||||
}
|
||||
var hello proto.DeviceHello
|
||||
if err = json.Unmarshal(helloEnv.Payload, &hello); err != nil {
|
||||
_ = conn.Close(websocket.StatusPolicyViolation, "bad hello")
|
||||
return
|
||||
}
|
||||
deviceID, _ := strconv.ParseUint(hello.DeviceID, 10, 64)
|
||||
var device models.AgentDevice
|
||||
if err = db.First(&device, deviceID).Error; err != nil {
|
||||
_ = conn.Close(websocket.StatusPolicyViolation, "unknown device")
|
||||
return
|
||||
}
|
||||
if device.Revoked {
|
||||
_ = conn.Close(websocket.StatusPolicyViolation, "device revoked")
|
||||
return
|
||||
}
|
||||
pubRaw, _ := base64.StdEncoding.DecodeString(device.PublicKey)
|
||||
if len(pubRaw) != ed25519.PublicKeySize {
|
||||
_ = conn.Close(websocket.StatusPolicyViolation, "bad key")
|
||||
return
|
||||
}
|
||||
sig, _ := base64.StdEncoding.DecodeString(hello.Sig)
|
||||
if !ed25519.Verify(ed25519.PublicKey(pubRaw), signPayload(hello.Nonce, hello.TS), sig) {
|
||||
_ = conn.Close(websocket.StatusPolicyViolation, "bad signature")
|
||||
return
|
||||
}
|
||||
// 注册 + 应答
|
||||
serverNonce := randHex(16)
|
||||
serverTS := time.Now().UnixMilli()
|
||||
ackSig := ed25519.Sign(serverKey, signPayload(serverNonce, serverTS))
|
||||
sess := &Conn{
|
||||
ID: hello.DeviceID,
|
||||
Kind: "agent",
|
||||
WSID: device.WorkspaceID,
|
||||
ws: conn,
|
||||
send: make(chan []byte, 64),
|
||||
}
|
||||
hub.RegisterAgent(sess)
|
||||
now := time.Now()
|
||||
_ = db.Model(&device).Updates(map[string]interface{}{"status": "online", "last_seen_at": &now, "ip": c.ClientIP(), "version": hello.Version}).Error
|
||||
hub.BroadcastWS(device.WorkspaceID, "agent.status", gin.H{"deviceId": device.ID, "online": true})
|
||||
ack := proto.NewEnvelope(randHex(8), proto.TypeHelloAck, proto.HelloAck{
|
||||
ServerNonce: serverNonce, SessionID: randHex(16), ServerTS: serverTS, Sig: base64.StdEncoding.EncodeToString(ackSig),
|
||||
})
|
||||
go sess.writeLoop()
|
||||
sess.sendEnvelope(ack)
|
||||
log.Printf("[ws] agent %d online (workspace %d)", device.ID, device.WorkspaceID)
|
||||
// 消息循环
|
||||
defer func() {
|
||||
hub.UnregisterAgent(sess)
|
||||
_ = db.Model(&device).Updates(map[string]interface{}{"status": "offline", "last_seen_at": time.Now()}).Error
|
||||
hub.BroadcastWS(device.WorkspaceID, "agent.status", gin.H{"deviceId": device.ID, "online": false})
|
||||
log.Printf("[ws] agent %d offline", device.ID)
|
||||
}()
|
||||
for {
|
||||
// 90s 无消息视为死连接（客户端心跳 10-30s）
|
||||
readCtx, readCancel := context.WithTimeout(c.Request.Context(), 90*time.Second)
|
||||
typ, raw, err = conn.Read(readCtx)
|
||||
readCancel()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if typ != websocket.MessageText {
|
||||
continue
|
||||
}
|
||||
var env proto.Envelope
|
||||
if err = json.Unmarshal(raw, &env); err != nil {
|
||||
continue
|
||||
}
|
||||
handleAgentMessage(db, hub, sess, device, env)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// handleAgentMessage 分发 Agent 消息（所有回写经 sess 统一通道，避免并发写）
|
||||
func handleAgentMessage(db *gorm.DB, hub *Hub, sess *Conn, device models.AgentDevice, env proto.Envelope) {
|
||||
var err error
|
||||
switch env.Type {
|
||||
case proto.TypeHeartbeat:
|
||||
var hb proto.Heartbeat
|
||||
if err = json.Unmarshal(env.Payload, &hb); err == nil {
|
||||
now := time.Now()
|
||||
_ = db.Model(&models.AgentDevice{}).Where("id = ?", device.ID).Update("last_seen_at", &now).Error
|
||||
ack := proto.NewEnvelope(env.ID, proto.TypeHeartbeatAck, proto.HeartbeatAck{ServerTS: now.UnixMilli()})
|
||||
sess.sendEnvelope(ack)
|
||||
}
|
||||
case proto.TypeTaskAck:
|
||||
var ack proto.TaskAck
|
||||
if err = json.Unmarshal(env.Payload, &ack); err == nil {
|
||||
applyTaskAck(db, hub, device, ack)
|
||||
}
|
||||
case proto.TypeTaskResult:
|
||||
var result proto.TaskResult
|
||||
if err = json.Unmarshal(env.Payload, &result); err == nil {
|
||||
log.Printf("[ws] task.result received: %+v", result)
|
||||
applyTaskResult(db, hub, device, result)
|
||||
} else {
|
||||
log.Printf("[ws] task.result unmarshal error: %v payload=%s", err, string(env.Payload))
|
||||
}
|
||||
case proto.TypeChallengeAck:
|
||||
var cack proto.ChallengeAck
|
||||
if err = json.Unmarshal(env.Payload, &cack); err == nil {
|
||||
log.Printf("[ws] challenge %s ack: %s", cack.ChallengeID, cack.Action)
|
||||
}
|
||||
case proto.TypeChallengeSolve:
|
||||
var cs proto.ChallengeSolve
|
||||
if err = json.Unmarshal(env.Payload, &cs); err == nil {
|
||||
applyChallengeSolve(db, hub, device, cs)
|
||||
}
|
||||
default:
|
||||
log.Printf("[ws] unknown message type: %s", env.Type)
|
||||
}
|
||||
}
|
||||
|
||||
// applyTaskAck 任务确认：dispatched → running
|
||||
func applyTaskAck(db *gorm.DB, hub *Hub, device models.AgentDevice, ack proto.TaskAck) {
|
||||
tid, _ := strconv.ParseUint(ack.TaskID, 10, 64)
|
||||
var row models.Task
|
||||
if err := db.Where("id = ? and workspace_id = ?", tid, device.WorkspaceID).First(&row).Error; err != nil {
|
||||
return
|
||||
}
|
||||
if !ack.Accept {
|
||||
_ = db.Model(&row).Updates(map[string]interface{}{"status": task.Failed, "error_message": ack.Reason}).Error
|
||||
hub.BroadcastWS(device.WorkspaceID, "task.status", taskEvent(row))
|
||||
return
|
||||
}
|
||||
if next, ok := task.Next(row.Status, task.ActionStart); ok {
|
||||
_ = db.Model(&row).Update("status", next).Error
|
||||
row.Status = next
|
||||
hub.BroadcastWS(device.WorkspaceID, "task.status", taskEvent(row))
|
||||
}
|
||||
}
|
||||
|
||||
// applyTaskResult 结果回传：running → success/failed
|
||||
func applyTaskResult(db *gorm.DB, hub *Hub, device models.AgentDevice, result proto.TaskResult) {
|
||||
tid, _ := strconv.ParseUint(result.TaskID, 10, 64)
|
||||
var row models.Task
|
||||
if err := db.Where("id = ? and workspace_id = ?", tid, device.WorkspaceID).First(&row).Error; err != nil {
|
||||
log.Printf("[ws] task.result: task %d not found: %v", tid, err)
|
||||
return
|
||||
}
|
||||
log.Printf("[ws] task.result: task %d current status %s, want %s", tid, row.Status, result.Status)
|
||||
if result.Status == "success" {
|
||||
if next, ok := task.Next(row.Status, task.ActionSuccess); ok {
|
||||
urls, _ := json.Marshal([]string{result.PublishedURL})
|
||||
receipts, _ := json.Marshal(result.Receipts)
|
||||
_ = db.Model(&row).Updates(map[string]interface{}{
|
||||
"status": next, "published_urls": string(urls), "receipts": string(receipts),
|
||||
"error_message": "", "published_at": time.Now(),
|
||||
}).Error
|
||||
row.Status = next
|
||||
hub.BroadcastWS(device.WorkspaceID, "task.status", taskEvent(row))
|
||||
log.Printf("[ws] task %d success in workspace %d", tid, device.WorkspaceID)
|
||||
}
|
||||
} else {
|
||||
if next, ok := task.Next(row.Status, task.ActionFail); ok {
|
||||
_ = db.Model(&row).Updates(map[string]interface{}{"status": next, "error_message": result.Error}).Error
|
||||
row.Status = next
|
||||
hub.BroadcastWS(device.WorkspaceID, "task.status", taskEvent(row))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// applyChallengeSolve 挑战完成：账号激活
|
||||
func applyChallengeSolve(db *gorm.DB, hub *Hub, device models.AgentDevice, cs proto.ChallengeSolve) {
|
||||
cid, _ := strconv.ParseUint(cs.ChallengeID, 10, 64)
|
||||
var ch models.Challenge
|
||||
if err := db.Where("id = ? and workspace_id = ?", cid, device.WorkspaceID).First(&ch).Error; err != nil {
|
||||
return
|
||||
}
|
||||
now := time.Now()
|
||||
_ = db.Model(&ch).Updates(map[string]interface{}{"status": "solved", "solved_at": &now, "qr_token": cs.Value}).Error
|
||||
_ = db.Model(&models.Account{}).Where("id = ?", ch.AccountID).Updates(map[string]interface{}{"status": "active", "agent_device_id": device.ID, "last_active_at": &now}).Error
|
||||
hub.BroadcastWS(device.WorkspaceID, "challenge.status", gin.H{"challengeId": ch.ID, "status": "solved", "accountId": ch.AccountID})
|
||||
log.Printf("[ws] challenge %d solved, account %d active", ch.ID, ch.AccountID)
|
||||
}
|
||||
|
||||
// taskEvent 任务状态事件
|
||||
func taskEvent(row models.Task) gin.H {
|
||||
return gin.H{"taskId": row.ID, "status": row.Status, "errorMessage": row.ErrorMessage, "publishedUrls": row.PublishedURLs}
|
||||
}
|
||||
|
||||
// writeLoop 写协程（唯一写出口）
|
||||
func (c *Conn) writeLoop() {
|
||||
for raw := range c.send {
|
||||
wctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
err := c.ws.Write(wctx, websocket.MessageText, raw)
|
||||
cancel()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func randHex(n int) string {
|
||||
b := make([]byte, n)
|
||||
_, _ = rand.Read(b)
|
||||
return hex.EncodeToString(b)
|
||||
}
|
||||
Reference in New Issue
Block a user