128 lines
3.3 KiB
Go
128 lines
3.3 KiB
Go
package snowluma
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"strconv"
|
||
"sync"
|
||
"sync/atomic"
|
||
"time"
|
||
)
|
||
|
||
// Caller 是一条可执行 OneBot action 的通道（WS 会话或 HTTP API）。
|
||
type Caller interface {
|
||
Call(ctx context.Context, action string, params any, out any) error
|
||
}
|
||
|
||
// CallAPI 通过指定 Caller 调用一个 action 并把 data 解码为 T。
|
||
// ctx 为 nil 时使用 30 秒超时。
|
||
func CallAPI[T any](c Caller, ctx context.Context, action string, params any) (T, error) {
|
||
var zero T
|
||
if c == nil {
|
||
return zero, ErrNotConnected
|
||
}
|
||
if ctx == nil {
|
||
var cancel context.CancelFunc
|
||
ctx, cancel = context.WithTimeout(context.Background(), 30*time.Second)
|
||
defer cancel()
|
||
}
|
||
var out T
|
||
if err := c.Call(ctx, action, params, &out); err != nil {
|
||
return zero, err
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
// wsRequest 是 WS 上发出的 action 请求。
|
||
type wsRequest struct {
|
||
Action string `json:"action"`
|
||
Params any `json:"params,omitempty"`
|
||
Echo string `json:"echo"`
|
||
}
|
||
|
||
// wsCaller 基于一条 WebSocket 连接实现 Caller：
|
||
// 发出的请求带自增 echo，读循环把响应按 echo 路由回等待方。
|
||
type wsCaller struct {
|
||
wmu sync.Mutex
|
||
writer messageWriter
|
||
pending sync.Map // echo(string) -> chan Response
|
||
seq atomic.Int64
|
||
done <-chan struct{} // 连接关闭时被关闭
|
||
}
|
||
|
||
// messageWriter 抽象 gorilla/websocket 的写方法，便于单测替换。
|
||
type messageWriter interface {
|
||
WriteMessage(messageType int, data []byte) error
|
||
}
|
||
|
||
func newWSCaller(w messageWriter, done <-chan struct{}) *wsCaller {
|
||
return &wsCaller{writer: w, done: done}
|
||
}
|
||
|
||
func (w *wsCaller) Call(ctx context.Context, action string, params any, out any) error {
|
||
if ctx == nil {
|
||
var cancel context.CancelFunc
|
||
ctx, cancel = context.WithTimeout(context.Background(), 30*time.Second)
|
||
defer cancel()
|
||
}
|
||
echo := strconv.FormatInt(w.seq.Add(1), 10)
|
||
payload, err := json.Marshal(wsRequest{Action: action, Params: params, Echo: echo})
|
||
if err != nil {
|
||
return fmt.Errorf("snowluma: marshal action %s: %w", action, err)
|
||
}
|
||
ch := make(chan Response, 1)
|
||
w.pending.Store(echo, ch)
|
||
defer w.pending.Delete(echo)
|
||
|
||
w.wmu.Lock()
|
||
err = w.writer.WriteMessage(1 /* websocket.TextMessage */, payload)
|
||
w.wmu.Unlock()
|
||
if err != nil {
|
||
return fmt.Errorf("snowluma: send action %s: %w", action, err)
|
||
}
|
||
|
||
select {
|
||
case resp := <-ch:
|
||
if err := resp.Err(); err != nil {
|
||
return err
|
||
}
|
||
if err := resp.Decode(out); err != nil {
|
||
return fmt.Errorf("snowluma: decode data of %s: %w", action, err)
|
||
}
|
||
return nil
|
||
case <-ctx.Done():
|
||
return ctx.Err()
|
||
case <-w.done:
|
||
return ErrNotConnected
|
||
}
|
||
}
|
||
|
||
// routeResponse 由读循环调用：把带 echo 的响应分发给等待者，返回是否路由成功。
|
||
func (w *wsCaller) routeResponse(data []byte, log logger) bool {
|
||
var resp Response
|
||
if err := json.Unmarshal(data, &resp); err != nil {
|
||
log.Warn("snowluma: unmarshal response failed", "err", err)
|
||
return false
|
||
}
|
||
v, ok := w.pending.Load(resp.Echo)
|
||
if !ok {
|
||
log.Warn("snowluma: no pending action for echo", "echo", resp.Echo)
|
||
return false
|
||
}
|
||
ch := v.(chan Response)
|
||
select {
|
||
case ch <- resp:
|
||
default:
|
||
}
|
||
return true
|
||
}
|
||
|
||
// logger 是框架依赖的最小日志接口（*slog.Logger 满足它）。
|
||
type logger interface {
|
||
Warn(msg string, args ...any)
|
||
Info(msg string, args ...any)
|
||
Error(msg string, args ...any)
|
||
Debug(msg string, args ...any)
|
||
}
|