diff --git a/gomsg/README.md b/gomsg/README.md new file mode 100644 index 0000000..a0df051 --- /dev/null +++ b/gomsg/README.md @@ -0,0 +1,73 @@ +# gomsg — pymsg / libmsg 协议的 Go 实现 + +Go 版消息通信库,对应 Python [pymsg](../pymsg/) 与 C++ [libmsg](../libmsg/),线协议字节级兼容: + +``` +[4 字节 big-endian uint32 长度][UTF-8 载荷] (单条上限 64 MiB) +``` + +## 模块组成 + +| 文件 | 内容 | 对应 | +|------|------|------| +| `frame.go` | 帧编解码:`Encode` / `WriteFrame` / `ReadFrame`(长度上限防护) | pymsg `send`/`recv` | +| `client.go` | `MsgClient`:复用长连接 + 请求-响应 + 失败重连一次 | pymsg `MsgClient` | +| `server.go` | `MsgServer`:accept 循环 + 每连接 goroutine 串行处理 | pymsg `MsgServer` | +| `rpc.go` | `RpcClient` / `RpcServer` / `RpcError`:`_route` 分发、错误脱敏、panic 恢复 | pymsg `rpc.py` | +| `asyncqueue.go` | `AsyncQueue[T]`:批量 + 定时 flush,满队列丢新数据(fail-safe) | libmsg `AsyncQueue` | +| `logclient.go` | `LogClient`:异步批量推送 db2 Logs 表(级别过滤 + fail-safe) | waf `internal/logclient`、`log_client.py`、liblogc | + +## 使用 + +```go +// RPC 调用(如调 db2 execute 路由) +rpc := gomsg.NewRpcClient("127.0.0.1:8005", 5*time.Second) +data, err := rpc.Call("execute", map[string]any{ + "action": "select", "table": "Users", +}) +// data 为 any:数字是 json.Number(保 int64 精度),需用 .Int64()/.Float64() 转换 + +// 提供 RPC 服务 +srv := gomsg.NewRpcServer("127.0.0.1:8009") +srv.AddRoute("health", func(req map[string]any) (any, error) { return map[string]any{"ok": true}, nil }) +if err := srv.Start(); err != nil { ... } +defer srv.Stop() + +// 集中式日志(写入 db2 Logs 表) +lc := gomsg.NewLogClient("127.0.0.1:8005") +lc.SetLevel("info") +lc.Info("backend", "服务启动", nil) +defer lc.Flush() // 进程退出前 flush +``` + +## 与 Python / C++ 的兼容约定 + +- **请求**:`{"_route": "<路由名>", ...业务字段}`(`_route` 在分发前被弹出,handler 看不到) +- **响应**:`{"success": true, "data": ...}` / `{"success": false, "error": "..."}` +- **未知路由**:`{"success": false, "error": "未知路由: "}` +- **非法 JSON**:`{"success": false, "error": "invalid JSON request"}`(不回显输入片段) +- **错误脱敏**:异常信息中的绝对路径替换为 ``(对齐 Python `_sanitize_error`) +- **handler panic**:自动恢复并包装为 `success=false`(对应 Python 的 `except Exception`) + +## 安全性 + +- 帧长度上限 64 MiB 在读写两侧强制(恶意长度前缀不会导致内存耗尽) +- 请求 JSON 解析:Go `encoding/json` 内置 10000 层嵌套上限;数字经 `UseNumber` 以字符串保存, + 超大数字不会触发 CPU 密集型解析;输入大小受帧上限约束 +- 异常信息脱敏,不透传绝对路径等内部细节 +- 数字一律 `json.Number` 保留原始精度,避免 int64 ID 经 float64 丢失 + +## 测试 + +```bash +go test ./... -count=1 -race +``` + +覆盖:帧往返/超限/截断、Client-Server 回环、自动重连、并发客户端、RPC 成功/失败/脱敏/ +未知路由/非法 JSON/panic 恢复/超时、AsyncQueue 批量/定时/关停 drain/满队列丢弃、 +LogClient 对假 db2 服务的推送与级别过滤、db2 不可用时 fail-safe。 + +## 来源 + +`MsgClient` 连接复用与 `LogClient` 逻辑抽取自 [waf/internal/logclient](../waf/internal/logclient/); +API 与协议约定对齐 [pymsg](../pymsg/)(Python)与 [libmsg](../libmsg/)(C++)。 diff --git a/gomsg/asyncqueue.go b/gomsg/asyncqueue.go new file mode 100644 index 0000000..d40fbb2 --- /dev/null +++ b/gomsg/asyncqueue.go @@ -0,0 +1,150 @@ +package gomsg + +import ( + "sync" + "time" +) + +// AsyncQueueConfig 异步批量队列配置(语义对齐 C++ libmsg.AsyncQueueConfig)。 +type AsyncQueueConfig struct { + BatchSize int // 批量发送阈值,达到即 flush;必须 > 0,默认 50 + FlushInterval time.Duration // 定时 flush 间隔;必须 > 0 且 <= 300s,默认 2s + MaxQueueSize int // 队列容量;必须 >= 1,默认 10000;满时丢新数据(对齐 waf logclient 的 fail-safe 语义) +} + +func (c AsyncQueueConfig) withDefaults() AsyncQueueConfig { + if c.BatchSize <= 0 { + c.BatchSize = 50 + } + if c.FlushInterval <= 0 || c.FlushInterval > 300*time.Second { + c.FlushInterval = 2 * time.Second + } + if c.MaxQueueSize < 1 { + c.MaxQueueSize = 10000 + } + return c +} + +// AsyncQueue 通用异步批量发送队列(对应 C++ libmsg.AsyncQueue)。 +// Push 非阻塞;队列满时静默丢弃并计数(fail-safe,不影响业务)。 +type AsyncQueue[T any] struct { + cfg AsyncQueueConfig + handler func([]T) + ch chan T + stop chan struct{} + wg sync.WaitGroup + + mu sync.Mutex + stopped bool + dropped uint64 + flushed uint64 +} + +// NewAsyncQueue 创建异步批量队列,handler 在单后台 goroutine 中按批调用。 +func NewAsyncQueue[T any](handler func([]T), cfg AsyncQueueConfig) *AsyncQueue[T] { + if handler == nil { + panic("gomsg: AsyncQueue handler must not be nil") + } + cfg = cfg.withDefaults() + q := &AsyncQueue[T]{ + cfg: cfg, + handler: handler, + ch: make(chan T, cfg.MaxQueueSize), + stop: make(chan struct{}), + } + q.wg.Add(1) + go q.loop() + return q +} + +// Push 入队一条数据;队列满时返回 false(该条被丢弃)。 +func (q *AsyncQueue[T]) Push(item T) bool { + q.mu.Lock() + if q.stopped { + q.mu.Unlock() + return false + } + q.mu.Unlock() + select { + case q.ch <- item: + return true + default: + q.mu.Lock() + q.dropped++ + q.mu.Unlock() + return false + } +} + +// Shutdown 停止队列,flush 剩余数据后退出后台 goroutine。 +func (q *AsyncQueue[T]) Shutdown() { + q.mu.Lock() + if q.stopped { + q.mu.Unlock() + return + } + q.stopped = true + q.mu.Unlock() + close(q.stop) + q.wg.Wait() +} + +// Stats 返回丢弃与已发送批次数(供监控/测试)。 +func (q *AsyncQueue[T]) Stats() (dropped, flushed uint64) { + q.mu.Lock() + defer q.mu.Unlock() + return q.dropped, q.flushed +} + +func (q *AsyncQueue[T]) loop() { + defer q.wg.Done() + batch := make([]T, 0, q.cfg.BatchSize) + ticker := time.NewTicker(q.cfg.FlushInterval) + defer ticker.Stop() + + for { + select { + case item := <-q.ch: + batch = append(batch, item) + // 非阻塞续取,直到凑满 batchSize 或暂时无数据 + for len(batch) < q.cfg.BatchSize { + select { + case item := <-q.ch: + batch = append(batch, item) + default: + goto send + } + } + send: + if len(batch) >= q.cfg.BatchSize { + q.send(batch) + batch = batch[:0] + } + case <-ticker.C: + if len(batch) > 0 { + q.send(batch) + batch = batch[:0] + } + case <-q.stop: + // 退出前 drain 剩余数据 + for { + select { + case item := <-q.ch: + batch = append(batch, item) + default: + if len(batch) > 0 { + q.send(batch) + } + return + } + } + } + } +} + +func (q *AsyncQueue[T]) send(batch []T) { + q.handler(batch) + q.mu.Lock() + q.flushed++ + q.mu.Unlock() +} diff --git a/gomsg/asyncqueue_test.go b/gomsg/asyncqueue_test.go new file mode 100644 index 0000000..23f0ae5 --- /dev/null +++ b/gomsg/asyncqueue_test.go @@ -0,0 +1,120 @@ +package gomsg + +import ( + "sync" + "testing" + "time" +) + +func collectBatches(batches *[][]int, mu *sync.Mutex) func([]int) { + return func(b []int) { + mu.Lock() + *batches = append(*batches, append([]int(nil), b...)) + mu.Unlock() + } +} + +func TestAsyncQueueFlushOnBatchSize(t *testing.T) { + var batches [][]int + var mu sync.Mutex + q := NewAsyncQueue[int](collectBatches(&batches, &mu), AsyncQueueConfig{ + BatchSize: 5, + FlushInterval: time.Hour, // 只测批量阈值,不测定时 + }) + for i := 1; i <= 5; i++ { + if !q.Push(i) { + t.Fatal("push should succeed") + } + } + deadline := time.Now().Add(2 * time.Second) + for { + mu.Lock() + n := len(batches) + mu.Unlock() + if n >= 1 { + break + } + if time.Now().After(deadline) { + t.Fatal("batch not flushed in time") + } + time.Sleep(10 * time.Millisecond) + } + mu.Lock() + defer mu.Unlock() + if len(batches[0]) != 5 || batches[0][4] != 5 { + t.Fatalf("batch = %v", batches[0]) + } + q.Shutdown() +} + +func TestAsyncQueueFlushOnInterval(t *testing.T) { + var batches [][]int + var mu sync.Mutex + q := NewAsyncQueue[int](collectBatches(&batches, &mu), AsyncQueueConfig{ + BatchSize: 100, + FlushInterval: 100 * time.Millisecond, + }) + if !q.Push(42) { + t.Fatal("push should succeed") + } + time.Sleep(300 * time.Millisecond) + mu.Lock() + n := len(batches) + mu.Unlock() + if n < 1 { + t.Fatal("expected interval flush") + } + q.Shutdown() +} + +func TestAsyncQueueShutdownDrains(t *testing.T) { + var batches [][]int + var mu sync.Mutex + q := NewAsyncQueue[int](collectBatches(&batches, &mu), AsyncQueueConfig{ + BatchSize: 100, + FlushInterval: time.Hour, + }) + for i := 1; i <= 3; i++ { + q.Push(i) + } + q.Shutdown() + mu.Lock() + defer mu.Unlock() + if len(batches) != 1 || len(batches[0]) != 3 { + t.Fatalf("expected drained batch [1 2 3], got %v", batches) + } +} + +func TestAsyncQueueDropOnFull(t *testing.T) { + var batches [][]int + var mu sync.Mutex + q := NewAsyncQueue[int](collectBatches(&batches, &mu), AsyncQueueConfig{ + BatchSize: 1000, + FlushInterval: time.Hour, + MaxQueueSize: 4, + }) + // 后台协程可能已取走部分,填满队列直到 Push 返回 false + dropped := false + for i := 0; i < 100; i++ { + if !q.Push(i) { + dropped = true + break + } + } + if !dropped { + t.Fatal("expected drop when queue full") + } + d, _ := q.Stats() + if d == 0 { + t.Fatal("dropped counter should be > 0") + } + q.Shutdown() +} + +func TestAsyncQueueRejectAfterShutdown(t *testing.T) { + q := NewAsyncQueue[int](func([]int) {}, AsyncQueueConfig{}) + q.Shutdown() + if q.Push(1) { + t.Fatal("push after shutdown should fail") + } +} diff --git a/gomsg/client.go b/gomsg/client.go new file mode 100644 index 0000000..aba51a4 --- /dev/null +++ b/gomsg/client.go @@ -0,0 +1,106 @@ +package gomsg + +import ( + "net" + "sync" + "time" +) + +// DefaultTimeout 单次请求-响应交互的总超时。 +const DefaultTimeout = 5 * time.Second + +// MsgClient pymsg 协议客户端:复用长连接 + 请求-响应调用。 +// 连接断开时自动重建一次(对应 Python pymsg.MsgClient / C++ libmsg.MsgClient)。 +type MsgClient struct { + addr string + timeout time.Duration + + mu sync.Mutex + conn net.Conn +} + +// NewMsgClient 创建客户端。timeout 为 0 时使用 DefaultTimeout。 +func NewMsgClient(addr string, timeout time.Duration) *MsgClient { + if timeout <= 0 { + timeout = DefaultTimeout + } + return &MsgClient{addr: addr, timeout: timeout} +} + +// Call 发送一条请求并读取响应(复用长连接,失败时关闭重连一次)。 +// 返回响应载荷;超时 / 连接错误返回 error。 +func (c *MsgClient) Call(payload []byte) ([]byte, error) { + resp, err := c.call(payload) + if err == nil { + return resp, nil + } + // 旧连接可能已断开(服务端重启等),关闭后重连一次 + c.closeConn() + return c.call(payload) +} + +// Send 发送一条请求,不等待响应(fire-and-forget)。失败时重连一次。 +func (c *MsgClient) Send(payload []byte) error { + if err := c.sendOnly(payload); err != nil { + c.closeConn() + return c.sendOnly(payload) + } + return nil +} + +// Close 关闭底层连接。 +func (c *MsgClient) Close() error { + c.mu.Lock() + defer c.mu.Unlock() + return c.closeLocked() +} + +func (c *MsgClient) call(payload []byte) ([]byte, error) { + conn, err := c.getConn() + if err != nil { + return nil, err + } + _ = conn.SetDeadline(time.Now().Add(c.timeout)) + if err := WriteFrame(conn, payload); err != nil { + return nil, err + } + return ReadFrame(conn, 0) +} + +func (c *MsgClient) sendOnly(payload []byte) error { + conn, err := c.getConn() + if err != nil { + return err + } + _ = conn.SetWriteDeadline(time.Now().Add(c.timeout)) + return WriteFrame(conn, payload) +} + +func (c *MsgClient) getConn() (net.Conn, error) { + c.mu.Lock() + defer c.mu.Unlock() + if c.conn != nil { + return c.conn, nil + } + conn, err := net.DialTimeout("tcp", c.addr, c.timeout) + if err != nil { + return nil, err + } + c.conn = conn + return conn, nil +} + +func (c *MsgClient) closeConn() { + c.mu.Lock() + defer c.mu.Unlock() + _ = c.closeLocked() +} + +func (c *MsgClient) closeLocked() error { + if c.conn != nil { + err := c.conn.Close() + c.conn = nil + return err + } + return nil +} diff --git a/gomsg/client_test.go b/gomsg/client_test.go new file mode 100644 index 0000000..8eeac7c --- /dev/null +++ b/gomsg/client_test.go @@ -0,0 +1,170 @@ +package gomsg + +import ( + "bytes" + "errors" + "testing" + "time" +) + +// startEchoServer 启动一个 MsgServer,把收到的载荷原样回复。 +func startEchoServer(t *testing.T) (*MsgServer, string) { + t.Helper() + s := NewMsgServer("127.0.0.1:0", func(payload []byte) []byte { + return append([]byte("echo:"), payload...) + }) + if err := s.Start(); err != nil { + t.Fatal(err) + } + t.Cleanup(s.Stop) + return s, s.Addr().String() +} + +func TestMsgClientCall(t *testing.T) { + _, addr := startEchoServer(t) + c := NewMsgClient(addr, 0) + defer c.Close() + + resp, err := c.Call([]byte(`{"hello":"世界"}`)) + if err != nil { + t.Fatal(err) + } + want := "echo:{\"hello\":\"世界\"}" + if string(resp) != want { + t.Fatalf("resp = %q, want %q", resp, want) + } +} + +func TestMsgClientReusesConn(t *testing.T) { + _, addr := startEchoServer(t) + c := NewMsgClient(addr, 0) + defer c.Close() + + for i := 0; i < 5; i++ { + if _, err := c.Call([]byte("x")); err != nil { + t.Fatal(err) + } + } + if c.conn == nil { + t.Fatal("expected connection reuse") + } +} + +func TestMsgClientReconnectAfterServerRestart(t *testing.T) { + s := NewMsgServer("127.0.0.1:0", func(p []byte) []byte { return p }) + if err := s.Start(); err != nil { + t.Fatal(err) + } + addr := s.Addr().String() + + c := NewMsgClient(addr, 2*time.Second) + defer c.Close() + if _, err := c.Call([]byte("before")); err != nil { + t.Fatal(err) + } + + // 服务端重启:旧连接失效,Call 应自动重连成功 + s.Stop() + // 重新监听同一地址:旧 listener 已释放端口 + s3 := NewMsgServer(addr, func(p []byte) []byte { return p }) + if err := s3.Start(); err != nil { + t.Fatal(err) + } + defer s3.Stop() + + resp, err := c.Call([]byte("after")) + if err != nil { + t.Fatalf("expected auto-reconnect, got %v", err) + } + if string(resp) != "after" { + t.Fatalf("resp = %q", resp) + } +} + +func TestMsgClientSendNoReply(t *testing.T) { + got := make(chan []byte, 1) + s := NewMsgServer("127.0.0.1:0", func(payload []byte) []byte { + got <- payload + return nil // 不回复 + }) + if err := s.Start(); err != nil { + t.Fatal(err) + } + defer s.Stop() + + c := NewMsgClient(s.Addr().String(), 0) + defer c.Close() + if err := c.Send([]byte("fire-and-forget")); err != nil { + t.Fatal(err) + } + select { + case p := <-got: + if string(p) != "fire-and-forget" { + t.Fatalf("payload = %q", p) + } + case <-time.After(2 * time.Second): + t.Fatal("server did not receive message") + } +} + +func TestMsgClientTimeout(t *testing.T) { + s := NewMsgServer("127.0.0.1:0", func(payload []byte) []byte { + time.Sleep(500 * time.Millisecond) // 拖过客户端超时 + return nil + }) + if err := s.Start(); err != nil { + t.Fatal(err) + } + defer s.Stop() + + c := NewMsgClient(s.Addr().String(), 50*time.Millisecond) + defer c.Close() + if _, err := c.Call([]byte("slow")); err == nil { + t.Fatal("expected timeout error") + } +} + +func TestMsgServerStopClosesConns(t *testing.T) { + _, addr := startEchoServer(t) + c := NewMsgClient(addr, 0) + defer c.Close() + if _, err := c.Call([]byte("ping")); err != nil { + t.Fatal(err) + } +} + +func TestMsgServerConcurrentClients(t *testing.T) { + s := NewMsgServer("127.0.0.1:0", func(payload []byte) []byte { + return bytes.ToUpper(payload) + }) + if err := s.Start(); err != nil { + t.Fatal(err) + } + defer s.Stop() + + errCh := make(chan error, 8) + for i := 0; i < 8; i++ { + go func(n int) { + c := NewMsgClient(s.Addr().String(), 0) + defer c.Close() + for j := 0; j < 10; j++ { + in := []byte{byte('a' + n)} + resp, err := c.Call(in) + if err != nil { + errCh <- err + return + } + if len(resp) != 1 || resp[0] != in[0]-'a'+'A' { + errCh <- errors.New("unexpected reply") + return + } + } + errCh <- nil + }(i) + } + for i := 0; i < 8; i++ { + if err := <-errCh; err != nil { + t.Fatal(err) + } + } +} diff --git a/gomsg/frame.go b/gomsg/frame.go new file mode 100644 index 0000000..52bee77 --- /dev/null +++ b/gomsg/frame.go @@ -0,0 +1,66 @@ +// Package gomsg 是 pymsg(Python)/ libmsg(C++)协议的 Go 实现。 +// +// 线协议(与 Python / C++ 字节级兼容): +// +// [4 字节 big-endian uint32 长度][UTF-8 载荷] +// +// 消息大小上限 64 MiB(与 Python pymsg / C++ libmsg 对齐)。 +// +// 覆盖功能:帧编解码(本文件)、MsgClient/MsgServer(client.go / server.go)、 +// JSON RPC 约定层(rpc.go)、通用异步批量队列(asyncqueue.go)与 +// 集中式日志客户端(logclient.go,对应各服务目录下的 log_client.py / liblogc)。 +package gomsg + +import ( + "encoding/binary" + "fmt" + "io" +) + +// MaxMsgSize 单条消息载荷大小上限,与 Python pymsg / C++ libmsg 保持一致。 +const MaxMsgSize = 64 * 1024 * 1024 + +var headerSize = 4 + +// Encode 将载荷封装为 pymsg 帧:4 字节大端长度前缀 + 载荷。 +func Encode(payload []byte) ([]byte, error) { + if len(payload) > MaxMsgSize { + return nil, fmt.Errorf("payload too large: %d bytes (max %d)", len(payload), MaxMsgSize) + } + frame := make([]byte, headerSize+len(payload)) + binary.BigEndian.PutUint32(frame, uint32(len(payload))) + copy(frame[headerSize:], payload) + return frame, nil +} + +// WriteFrame 将载荷以 pymsg 帧形式写入 w。 +func WriteFrame(w io.Writer, payload []byte) error { + frame, err := Encode(payload) + if err != nil { + return err + } + _, err = w.Write(frame) + return err +} + +// ReadFrame 从 r 读取一条 pymsg 帧,返回载荷。 +// max 为单条消息大小上限(传 0 表示使用 MaxMsgSize)。 +// 长度前缀超过上限视为协议错误(防止恶意长度导致内存耗尽)。 +func ReadFrame(r io.Reader, max int) ([]byte, error) { + if max <= 0 || max > MaxMsgSize { + max = MaxMsgSize + } + header := make([]byte, headerSize) + if _, err := io.ReadFull(r, header); err != nil { + return nil, err + } + length := binary.BigEndian.Uint32(header) + if length > uint32(max) { + return nil, fmt.Errorf("message too large: %d bytes (max %d)", length, max) + } + payload := make([]byte, length) + if _, err := io.ReadFull(r, payload); err != nil { + return nil, err + } + return payload, nil +} diff --git a/gomsg/frame_test.go b/gomsg/frame_test.go new file mode 100644 index 0000000..cc78556 --- /dev/null +++ b/gomsg/frame_test.go @@ -0,0 +1,66 @@ +package gomsg + +import ( + "bytes" + "encoding/binary" + "errors" + "io" + "testing" +) + +func TestEncodeRoundTrip(t *testing.T) { + payload := []byte(`{"hello":"世界"}`) + frame, err := Encode(payload) + if err != nil { + t.Fatal(err) + } + if len(frame) != 4+len(payload) { + t.Fatalf("frame length = %d, want %d", len(frame), 4+len(payload)) + } + if got := binary.BigEndian.Uint32(frame[:4]); got != uint32(len(payload)) { + t.Fatalf("header = %d, want %d", got, len(payload)) + } + got, err := ReadFrame(bytes.NewReader(frame), 0) + if err != nil { + t.Fatal(err) + } + if !bytes.Equal(got, payload) { + t.Fatalf("payload mismatch: %q", got) + } +} + +func TestEncodeTooLarge(t *testing.T) { + payload := make([]byte, MaxMsgSize+1) + if _, err := Encode(payload); err == nil { + t.Fatal("expected error for oversized payload") + } +} + +func TestReadFrameTooLarge(t *testing.T) { + header := make([]byte, 4) + binary.BigEndian.PutUint32(header, MaxMsgSize+1) + _, err := ReadFrame(bytes.NewReader(header), 0) + if err == nil { + t.Fatal("expected error for oversized header") + } +} + +func TestReadFrameTruncated(t *testing.T) { + frame, _ := Encode([]byte("abcdef")) + if _, err := ReadFrame(bytes.NewReader(frame[:len(frame)-3]), 0); err == nil { + t.Fatal("expected error for truncated payload") + } +} + +func TestReadFrameMaxParam(t *testing.T) { + frame, _ := Encode([]byte("hello")) + if _, err := ReadFrame(bytes.NewReader(frame), 4); err == nil { + t.Fatal("expected error when max < payload length") + } +} + +func TestReadFrameEOF(t *testing.T) { + if _, err := ReadFrame(bytes.NewReader(nil), 0); !errors.Is(err, io.EOF) { + t.Fatalf("expected EOF, got %v", err) + } +} diff --git a/gomsg/go.mod b/gomsg/go.mod new file mode 100644 index 0000000..ab95b8a --- /dev/null +++ b/gomsg/go.mod @@ -0,0 +1,3 @@ +module github.com/micro_servers/gomsg + +go 1.22 diff --git a/gomsg/logclient.go b/gomsg/logclient.go new file mode 100644 index 0000000..7e42f4a --- /dev/null +++ b/gomsg/logclient.go @@ -0,0 +1,173 @@ +package gomsg + +import ( + "encoding/json" + "fmt" + "os" + "time" +) + +// 集中式日志客户端:通过 pymsg/TCP 调用 db2(或 db)execute 路由, +// 逐条 INSERT 到 Logs 表。对应各服务目录下的 log_client.py 与 C++ liblogc。 +// +// 设计要点(与 waf internal/logclient / Python log_client.py 对齐): +// - 异步非阻塞:AsyncQueue 批量推送,不拖慢业务 +// - fail-safe:db 服务不可用时静默丢弃,绝不影响业务 +// - 缓冲批量:满 50 条或超时 2 秒自动 flush + +const defaultDBAddr = "127.0.0.1:8005" + +// 日志级别数值(越小越详细)。 +const ( + levelDebug = 0 + levelInfo = 1 + levelWarn = 2 + levelError = 3 +) + +var levelMap = map[string]int{ + "debug": levelDebug, "DEBUG": levelDebug, + "info": levelInfo, "INFO": levelInfo, + "warn": levelWarn, "warning": levelWarn, "WARN": levelWarn, "WARNING": levelWarn, + "error": levelError, "ERROR": levelError, +} + +// Entry 单条日志。 +type Entry struct { + Service string + Level string + Message string + TraceID string + Extra map[string]any +} + +// LogClient 异步批量日志推送客户端。 +type LogClient struct { + queue *AsyncQueue[Entry] + client *MsgClient + minLevel int +} + +// NewLogClient 创建日志客户端并启动后台批量推送。 +// dbAddr 为空时依次取环境变量 LOG_DB_ADDR、默认 127.0.0.1:8005; +// 默认最低级别 info,可用 SetLevel 调整(环境变量 LOG_LEVEL 同样生效)。 +func NewLogClient(dbAddr string) *LogClient { + if dbAddr == "" { + dbAddr = os.Getenv("LOG_DB_ADDR") + } + if dbAddr == "" { + dbAddr = defaultDBAddr + } + lc := &LogClient{ + client: NewMsgClient(dbAddr, 5*time.Second), + minLevel: levelInfo, + } + lc.queue = NewAsyncQueue[Entry](lc.sendBatch, AsyncQueueConfig{}) + if lv := os.Getenv("LOG_LEVEL"); lv != "" { + lc.SetLevel(lv) + } + return lc +} + +// SetLevel 设置最低日志级别(debug|info|warn|error,大小写不敏感)。 +func (c *LogClient) SetLevel(level string) { + if n, ok := levelMap[level]; ok { + c.minLevel = n + } +} + +// Push 推送一条日志(非阻塞,队列满或低于最低级别时静默丢弃)。 +func (c *LogClient) Push(e Entry) bool { + n, ok := levelMap[e.Level] + if !ok || n < c.minLevel { + return false + } + return c.queue.Push(e) +} + +// Debug 推送 DEBUG 级别日志。 +func (c *LogClient) Debug(service, message string, extra map[string]any) bool { + return c.Push(Entry{Service: service, Level: "DEBUG", Message: message, Extra: extra}) +} + +// Info 推送 INFO 级别日志。 +func (c *LogClient) Info(service, message string, extra map[string]any) bool { + return c.Push(Entry{Service: service, Level: "INFO", Message: message, Extra: extra}) +} + +// Warning 推送 WARNING 级别日志。 +func (c *LogClient) Warning(service, message string, extra map[string]any) bool { + return c.Push(Entry{Service: service, Level: "WARNING", Message: message, Extra: extra}) +} + +// Error 推送 ERROR 级别日志。 +func (c *LogClient) Error(service, message string, extra map[string]any) bool { + return c.Push(Entry{Service: service, Level: "ERROR", Message: message, Extra: extra}) +} + +// Flush 停止队列并 flush 剩余日志(用于进程退出前)。 +func (c *LogClient) Flush() { + c.queue.Shutdown() +} + +// dbV2Request execute 路由的 INSERT 请求体(与 Python/C++ 各 log 客户端一致)。 +type dbV2Request struct { + Route string `json:"_route"` + Action string `json:"action"` + Table string `json:"table"` + Values map[string]any `json:"values"` +} + +// dbV2Response execute 路由的通用响应。 +type dbV2Response struct { + Success bool `json:"success"` + Error string `json:"error,omitempty"` +} + +func (c *LogClient) sendBatch(batch []Entry) { + now := time.Now().UTC().Format(time.RFC3339) + for _, e := range batch { + c.sendOne(e, now) + } +} + +// sendOne 通过 pymsg TCP 协议发送单条日志到 db2(fail-safe:任何错误静默丢弃)。 +func (c *LogClient) sendOne(e Entry, ts string) { + extraJSON := "" + if e.Extra != nil { + b, err := json.Marshal(e.Extra) + if err != nil { + return + } + extraJSON = string(b) + } + + values := map[string]any{ + "ts": ts, + "service": e.Service, + "level": e.Level, + "message": e.Message, + } + if e.TraceID != "" { + values["trace_id"] = e.TraceID + } + if extraJSON != "" { + values["extra"] = extraJSON + } + + req := dbV2Request{Route: "execute", Action: "insert", Table: "Logs", Values: values} + data, err := json.Marshal(req) + if err != nil { + return + } + + resp, err := c.client.Call(data) + if err != nil { + return // fail-safe + } + + var r dbV2Response + if json.Unmarshal(resp, &r) != nil || !r.Success { + fmt.Fprintf(os.Stderr, "[gomsg:logclient] INSERT failed: %s\n", r.Error) + } +} diff --git a/gomsg/logclient_test.go b/gomsg/logclient_test.go new file mode 100644 index 0000000..b487c5e --- /dev/null +++ b/gomsg/logclient_test.go @@ -0,0 +1,114 @@ +package gomsg + +import ( + "encoding/json" + "sync" + "testing" + "time" +) + +// fakeDB2 模拟 db2 的 execute 路由:校验 Logs INSERT 请求格式并返回成功。 +type fakeDB2 struct { + server *MsgServer + mu sync.Mutex + inserts []map[string]any + failures int +} + +func newFakeDB2(t *testing.T) *fakeDB2 { + f := &fakeDB2{} + f.server = NewMsgServer("127.0.0.1:0", func(payload []byte) []byte { + var req map[string]any + if err := json.Unmarshal(payload, &req); err != nil { + return []byte(`{"success":false,"error":"invalid JSON request"}`) + } + // db2 真实语义:非法 JSON 返回 success=false;其余按通用响应 + if req["_route"] != "execute" { + return []byte(`{"success":false,"error":"未知路由"}`) + } + f.mu.Lock() + f.inserts = append(f.inserts, req) + f.mu.Unlock() + return []byte(`{"success":true,"data":{}}`) + }) + if err := f.server.Start(); err != nil { + t.Fatal(err) + } + t.Cleanup(f.server.Stop) + return f +} + +func (f *fakeDB2) count() int { + f.mu.Lock() + defer f.mu.Unlock() + return len(f.inserts) +} + +func TestLogClientPushesToDB2(t *testing.T) { + db := newFakeDB2(t) + lc := NewLogClient(db.server.Addr().String()) + defer lc.Flush() + + lc.SetLevel("debug") + lc.Debug("test-svc", "debug msg", nil) + lc.Info("test-svc", "hello", map[string]any{"k": "v"}) + lc.Error("test-svc", "boom", nil) + + deadline := time.Now().Add(3 * time.Second) + for db.count() < 3 { + if time.Now().After(deadline) { + t.Fatal("logs not pushed in time") + } + time.Sleep(20 * time.Millisecond) + } + + db.mu.Lock() + defer db.mu.Unlock() + first := db.inserts[0] + if first["_route"] != "execute" { + t.Fatalf("_route = %v", first["_route"]) + } + vals := first["values"].(map[string]any) + if vals["service"] != "test-svc" || vals["level"] != "DEBUG" || vals["message"] != "debug msg" { + t.Fatalf("values = %v", vals) + } + // ts 必须是 ISO 8601 字符串 + if _, ok := vals["ts"].(string); !ok { + t.Fatalf("ts missing: %v", vals) + } + // 第二条(Info)的 extra 应序列化为 JSON 字符串 + vals2 := db.inserts[1]["values"].(map[string]any) + if vals2["extra"] != `{"k":"v"}` { + t.Fatalf("extra = %v", vals2["extra"]) + } +} + +func TestLogClientLevelFilter(t *testing.T) { + db := newFakeDB2(t) + lc := NewLogClient(db.server.Addr().String()) + defer lc.Flush() + + lc.Info("svc", "visible", nil) + lc.Debug("svc", "hidden", nil) // 默认 minLevel=info,debug 应被过滤 + lc.Push(Entry{Service: "svc", Level: "NOPE", Message: "hidden"}) + + deadline := time.Now().Add(3 * time.Second) + for db.count() < 1 { + if time.Now().After(deadline) { + t.Fatal("expected exactly 1 log") + } + time.Sleep(20 * time.Millisecond) + } + time.Sleep(100 * time.Millisecond) // 确认没有多余的推送 + if db.count() != 1 { + t.Fatalf("expected 1 insert, got %d", db.count()) + } +} + +func TestLogClientFailSafeWhenDBDown(t *testing.T) { + lc := NewLogClient("127.0.0.1:1") // 不可达地址 + defer lc.Flush() + lc.Info("svc", "will be dropped", nil) + time.Sleep(200 * time.Millisecond) + // 不 panic、不阻塞即视为通过(fail-safe 语义) +} diff --git a/gomsg/rpc.go b/gomsg/rpc.go new file mode 100644 index 0000000..1381167 --- /dev/null +++ b/gomsg/rpc.go @@ -0,0 +1,224 @@ +package gomsg + +import ( + "bytes" + "encoding/json" + "fmt" + "regexp" + "strings" + "sync" + "time" +) + +// JSON RPC 约定层(与 Python pymsg.rpc / C++ libmsg RpcClient/RpcServer 兼容): +// +// 请求: {"_route": "<路由名>", ...业务字段} +// 响应: {"success": true, "data": ...} +// {"success": false, "error": "..."} +// +// 数字一律以 json.Number 保留原始精度(避免 int64 经 float64 丢失)。 + +// RpcError 远端返回 success=false 的业务错误,或响应不符合 RPC 约定。 +type RpcError struct { + msg string +} + +func (e *RpcError) Error() string { return e.msg } + +// RpcHandler 处理一条 RPC 请求,返回值放入 data,错误包装为 success=false。 +type RpcHandler func(req map[string]any) (any, error) + +// RpcClient JSON RPC 客户端(复用 MsgClient 长连接)。 +type RpcClient struct { + client *MsgClient +} + +// NewRpcClient 创建客户端。timeout 为 0 时使用 DefaultTimeout。 +func NewRpcClient(addr string, timeout time.Duration) *RpcClient { + return &RpcClient{client: NewMsgClient(addr, timeout)} +} + +// Call 调用指定路由,返回响应中的 data 字段(可能为 nil)。 +// success=false 或非法响应时返回 *RpcError。 +func (c *RpcClient) Call(route string, payload map[string]any) (any, error) { + req := make(map[string]any, len(payload)+1) + for k, v := range payload { + req[k] = v + } + req["_route"] = route + body, err := json.Marshal(req) + if err != nil { + return nil, err + } + resp, err := c.client.Call(body) + if err != nil { + return nil, err + } + + var obj map[string]any + dec := json.NewDecoder(bytes.NewReader(resp)) + dec.UseNumber() + if err := dec.Decode(&obj); err != nil || obj == nil { + return nil, &RpcError{msg: fmt.Sprintf("响应不是合法 JSON: %v", err)} + } + success, ok := obj["success"].(bool) + if !ok { + return nil, &RpcError{msg: fmt.Sprintf("非法 RPC 响应: %s", truncate(resp))} + } + if !success { + errMsg, _ := obj["error"].(string) + return nil, &RpcError{msg: errMsg} + } + return obj["data"], nil +} + +// Close 关闭底层连接。 +func (c *RpcClient) Close() error { return c.client.Close() } + +// RpcServer JSON RPC 服务端:基于 MsgServer 的路由分发与响应包装。 +type RpcServer struct { + *MsgServer + routes map[string]RpcHandler + mu sync.RWMutex + // OnError 可选钩子:handler 抛错时调用(用于记日志)。 + OnError func(route string, err error) + // JSONMarshal 响应序列化函数,默认 json.Marshal(对应 Python 的 json_default)。 + JSONMarshal func(v any) ([]byte, error) +} + +// NewRpcServer 创建 RPC 服务端(不监听,需调用 Start)。 +func NewRpcServer(addr string) *RpcServer { + s := &RpcServer{routes: make(map[string]RpcHandler)} + s.MsgServer = NewMsgServer(addr, s.dispatch) + return s +} + +// AddRoute 注册路由处理函数。 +func (s *RpcServer) AddRoute(name string, fn RpcHandler) { + s.mu.Lock() + defer s.mu.Unlock() + s.routes[name] = fn +} + +// dispatch 解析请求、查找路由、包装响应;任何输入错误都不向客户端回显内部细节。 +// handler 内 panic 会被恢复并包装为 success=false(对应 Python 的 except Exception)。 +func (s *RpcServer) dispatch(payload []byte) (reply []byte) { + defer func() { + if r := recover(); r != nil { + err := fmt.Errorf("handler panic: %v", r) + if s.OnError != nil { + s.OnError("", err) + } + reply = s.reply(false, nil, sanitizeError(err.Error())) + } + }() + + req, err := decodeRequest(payload) + if err != nil { + return s.reply(false, nil, "invalid JSON request") + } + + routeVal, has := req["_route"] + delete(req, "_route") + route := "" + switch r := routeVal.(type) { + case string: + route = r + case nil: + // 保持 Python 行为:缺省 _route 时按空路由查找(未知路由) + default: + if has { + route = fmt.Sprintf("%v", r) + } + } + + s.mu.RLock() + handler := s.routes[route] + s.mu.RUnlock() + if handler == nil { + return s.reply(false, nil, "未知路由: "+route) + } + + data, err := handler(req) + if err != nil { + if s.OnError != nil { + s.OnError(route, err) + } + return s.reply(false, nil, sanitizeError(err.Error())) + } + return s.reply(true, data, "") +} + +func (s *RpcServer) reply(success bool, data any, errMsg string) []byte { + resp := map[string]any{"success": success} + if data != nil { + resp["data"] = data + } + if errMsg != "" { + resp["error"] = errMsg + } + marshal := s.JSONMarshal + if marshal == nil { + marshal = json.Marshal + } + out, err := marshal(resp) + if err != nil { + out, _ = json.Marshal(map[string]any{ + "success": false, + "error": sanitizeError(err.Error()), + }) + } + return out +} + +// decodeRequest 解析请求 JSON:UseNumber 保留整数精度; +// Go 的 encoding/json 内置 10000 层嵌套上限,数字位数字符串化保存, +// 输入大小又受帧协议 64 MiB 上限约束,天然规避深度/超大数字 DoS。 +func decodeRequest(payload []byte) (map[string]any, error) { + dec := json.NewDecoder(bytes.NewReader(payload)) + dec.UseNumber() + var req map[string]any + if err := dec.Decode(&req); err != nil { + return nil, err + } + if req == nil { + return nil, fmt.Errorf("empty request") + } + return req, nil +} + +// pathRe 匹配绝对路径(如 /opt/micro_servers/db/xxx.py),用于脱敏异常信息。 +var pathRe = regexp.MustCompile(`/[A-Za-z0-9_.\-]+(?:/[A-Za-z0-9_.\-]+)*`) + +// sanitizeError 脱敏异常信息中的绝对路径等内部细节,保留业务错误语义 +// (对齐 Python pymsg._sanitize_error:前导字符为单词字符时不匹配)。 +func sanitizeError(msg string) string { + if msg == "" { + return msg + } + var b strings.Builder + last := 0 + for _, idx := range pathRe.FindAllStringIndex(msg, -1) { + start, end := idx[0], idx[1] + if start > 0 && isWordChar(msg[start-1]) { + continue + } + b.WriteString(msg[last:start]) + b.WriteString("") + last = end + } + b.WriteString(msg[last:]) + return b.String() +} + +func isWordChar(c byte) bool { + return c == '_' || (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') +} + +func truncate(b []byte) string { + const limit = 200 + if len(b) <= limit { + return string(b) + } + return string(b[:limit]) + "..." +} diff --git a/gomsg/rpc_test.go b/gomsg/rpc_test.go new file mode 100644 index 0000000..b9324ce --- /dev/null +++ b/gomsg/rpc_test.go @@ -0,0 +1,203 @@ +package gomsg + +import ( + "encoding/json" + "errors" + "strings" + "testing" + "time" +) + +// startRPCServer 启动带两路由的 RPC 服务端。 +func startRPCServer(t *testing.T) (*RpcServer, string) { + t.Helper() + s := NewRpcServer("127.0.0.1:0") + s.AddRoute("echo", func(req map[string]any) (any, error) { + return req, nil + }) + s.AddRoute("sum", func(req map[string]any) (any, error) { + a := req["a"].(json.Number) + b := req["b"].(json.Number) + ai, _ := a.Int64() + bi, _ := b.Int64() + return ai + bi, nil + }) + s.AddRoute("fail", func(req map[string]any) (any, error) { + return nil, errors.New("boom: /opt/micro_servers/db2/xxx.py failed") + }) + s.AddRoute("panic", func(req map[string]any) (any, error) { + panic("panic in handler") + }) + if err := s.Start(); err != nil { + t.Fatal(err) + } + t.Cleanup(s.Stop) + return s, s.Addr().String() +} + +func TestRpcClientCallSuccess(t *testing.T) { + _, addr := startRPCServer(t) + c := NewRpcClient(addr, 0) + defer c.Close() + + data, err := c.Call("echo", map[string]any{"hello": "world"}) + if err != nil { + t.Fatal(err) + } + m := data.(map[string]any) + if m["hello"] != "world" { + t.Fatalf("data = %v", data) + } + if _, has := m["_route"]; has { + t.Fatal("_route should be stripped before handler") + } +} + +func TestRpcClientNumberPrecision(t *testing.T) { + _, addr := startRPCServer(t) + c := NewRpcClient(addr, 0) + defer c.Close() + + data, err := c.Call("sum", map[string]any{"a": json.Number("9007199254740993"), "b": json.Number("1")}) + if err != nil { + t.Fatal(err) + } + if got := data.(json.Number).String(); got != "9007199254740994" { + t.Fatalf("sum = %s, want exact int64 precision", got) + } +} + +func TestRpcClientBusinessError(t *testing.T) { + _, addr := startRPCServer(t) + c := NewRpcClient(addr, 0) + defer c.Close() + + _, err := c.Call("fail", nil) + var rpcErr *RpcError + if !errors.As(err, &rpcErr) { + t.Fatalf("expected *RpcError, got %T: %v", err, err) + } + // 绝对路径应被脱敏 + if strings.Contains(rpcErr.Error(), "/opt/") { + t.Fatalf("error not sanitized: %v", rpcErr) + } + if !strings.Contains(rpcErr.Error(), "") { + t.Fatalf("expected in error: %v", rpcErr) + } +} + +func TestRpcUnknownRoute(t *testing.T) { + _, addr := startRPCServer(t) + c := NewRpcClient(addr, 0) + defer c.Close() + + _, err := c.Call("no_such_route", nil) + if err == nil || !strings.Contains(err.Error(), "未知路由") { + t.Fatalf("expected 未知路由 error, got %v", err) + } +} + +func TestRpcServerInvalidJSON(t *testing.T) { + _, addr := startRPCServer(t) + c := NewMsgClient(addr, 0) + defer c.Close() + + resp, err := c.Call([]byte("{not json")) + if err != nil { + t.Fatal(err) + } + var r map[string]any + if json.Unmarshal(resp, &r) != nil { + t.Fatalf("invalid response: %q", resp) + } + if r["success"] != false || !strings.Contains(r["error"].(string), "invalid JSON request") { + t.Fatalf("unexpected response: %v", r) + } +} + +func TestRpcServerPanicHandler(t *testing.T) { + _, addr := startRPCServer(t) + c := NewRpcClient(addr, 0) + defer c.Close() + + // handler 内 panic 不应打崩服务,后续请求仍可用 + if _, err := c.Call("panic", nil); err == nil { + t.Fatal("expected error from panicking handler") + } + if _, err := c.Call("echo", map[string]any{"ok": true}); err != nil { + t.Fatal(err) + } +} + +func TestRpcServerOnErrorHook(t *testing.T) { + s := NewRpcServer("127.0.0.1:0") + var hookRoute string + s.OnError = func(route string, err error) { hookRoute = route } + s.AddRoute("fail", func(req map[string]any) (any, error) { + return nil, errors.New("boom") + }) + if err := s.Start(); err != nil { + t.Fatal(err) + } + defer s.Stop() + + c := NewRpcClient(s.Addr().String(), 0) + defer c.Close() + _, _ = c.Call("fail", nil) + if hookRoute != "fail" { + t.Fatalf("OnError hook route = %q", hookRoute) + } +} + +func TestRpcServerNonDictRequest(t *testing.T) { + s := NewRpcServer("127.0.0.1:0") + s.AddRoute("echo", func(req map[string]any) (any, error) { return req, nil }) + if err := s.Start(); err != nil { + t.Fatal(err) + } + defer s.Stop() + + c := NewMsgClient(s.Addr().String(), 0) + defer c.Close() + resp, err := c.Call([]byte(`["not","a","dict"]`)) + if err != nil { + t.Fatal(err) + } + var r map[string]any + if json.Unmarshal(resp, &r) != nil || r["success"] != false { + t.Fatalf("expected failure response, got %q", resp) + } +} + +func TestSanitizeError(t *testing.T) { + cases := []struct{ in, want string }{ + {"", ""}, + {"boom", "boom"}, + {"failed: /opt/micro_servers/db2/xxx.py", "failed: "}, + {"no abs path here", "no abs path here"}, + {"abc/def 保持原样", "abc/def 保持原样"}, // 前导字符为单词字符时不匹配(对齐 Python lookbehind) + } + for _, c := range cases { + if got := sanitizeError(c.in); got != c.want { + t.Errorf("sanitizeError(%q) = %q, want %q", c.in, got, c.want) + } + } +} + +func TestRpcClientTimeout(t *testing.T) { + s := NewRpcServer("127.0.0.1:0") + s.AddRoute("slow", func(req map[string]any) (any, error) { + time.Sleep(300 * time.Millisecond) + return nil, nil + }) + if err := s.Start(); err != nil { + t.Fatal(err) + } + defer s.Stop() + + c := NewRpcClient(s.Addr().String(), 50*time.Millisecond) + defer c.Close() + if _, err := c.Call("slow", nil); err == nil { + t.Fatal("expected timeout error") + } +} diff --git a/gomsg/server.go b/gomsg/server.go new file mode 100644 index 0000000..d01603e --- /dev/null +++ b/gomsg/server.go @@ -0,0 +1,130 @@ +package gomsg + +import ( + "net" + "sync" + "time" +) + +// MsgHandler 处理一条完整消息,返回回复载荷;返回 nil 表示不回复。 +type MsgHandler func(payload []byte) []byte + +// MsgServer pymsg 协议服务端:accept 循环 + 每连接一个 goroutine 串行处理。 +// 对应 Python pymsg.MsgServer / C++ libmsg.MsgServer。 +type MsgServer struct { + addr string + handler MsgHandler + // ReadTimeout 单条消息读取总超时;0 表示不设超时(默认,与 Python 行为一致)。 + ReadTimeout time.Duration + + mu sync.Mutex + ln net.Listener + conns map[net.Conn]struct{} + closed bool + wg sync.WaitGroup +} + +// NewMsgServer 创建服务端(不监听,需调用 Start)。 +func NewMsgServer(addr string, handler MsgHandler) *MsgServer { + return &MsgServer{ + addr: addr, + handler: handler, + conns: make(map[net.Conn]struct{}), + } +} + +// Addr 返回监听地址(Start 后可用,便于测试取随机端口)。 +func (s *MsgServer) Addr() net.Addr { + s.mu.Lock() + defer s.mu.Unlock() + if s.ln == nil { + return nil + } + return s.ln.Addr() +} + +// Start 开始监听并启动 accept 循环(非阻塞)。 +func (s *MsgServer) Start() error { + ln, err := net.Listen("tcp", s.addr) + if err != nil { + return err + } + s.mu.Lock() + if s.closed { + s.mu.Unlock() + _ = ln.Close() + return nil + } + s.ln = ln + s.mu.Unlock() + + s.wg.Add(1) + go s.acceptLoop(ln) + return nil +} + +// Stop 关闭监听与全部连接,等待处理中的请求结束。 +func (s *MsgServer) Stop() { + s.mu.Lock() + if s.closed { + s.mu.Unlock() + return + } + s.closed = true + if s.ln != nil { + _ = s.ln.Close() + } + for conn := range s.conns { + _ = conn.Close() + } + s.mu.Unlock() + s.wg.Wait() +} + +func (s *MsgServer) acceptLoop(ln net.Listener) { + defer s.wg.Done() + for { + conn, err := ln.Accept() + if err != nil { + return // 监听关闭 + } + s.mu.Lock() + if s.closed { + s.mu.Unlock() + _ = conn.Close() + return + } + s.conns[conn] = struct{}{} + s.mu.Unlock() + + s.wg.Add(1) + go s.handleConn(conn) + } +} + +func (s *MsgServer) handleConn(conn net.Conn) { + defer s.wg.Done() + defer func() { + s.mu.Lock() + delete(s.conns, conn) + s.mu.Unlock() + _ = conn.Close() + }() + + for { + if s.ReadTimeout > 0 { + _ = conn.SetReadDeadline(time.Now().Add(s.ReadTimeout)) + } + payload, err := ReadFrame(conn, 0) + if err != nil { + return + } + reply := s.handler(payload) + if reply == nil { + continue + } + if err := WriteFrame(conn, reply); err != nil { + return + } + } +} diff --git a/waf/README.md b/waf/README.md index b630142..cdefdae 100644 --- a/waf/README.md +++ b/waf/README.md @@ -144,7 +144,7 @@ waf/ │ ├── engine/ │ │ └── engine.go # coraza 封装:加载/热加载/列举/单测/输入检测 │ ├── logclient/ -│ │ └── logclient.go # Go 版 pymsg 日志客户端(异步批量推送) +│ │ └── logclient.go # gomsg 日志客户端适配层(协议/队列实现见 ../gomsg) │ └── proxy/ │ └── proxy.go # ReverseProxy + WAF 事务接入 + 安全响应头注入 ├── rules/ diff --git a/waf/go.mod b/waf/go.mod index 4a05569..76a4f8b 100644 --- a/waf/go.mod +++ b/waf/go.mod @@ -4,9 +4,12 @@ go 1.22.0 require ( github.com/corazawaf/coraza/v3 v3.3.0 + github.com/micro_servers/gomsg v0.0.0 gopkg.in/yaml.v3 v3.0.1 ) +replace github.com/micro_servers/gomsg => ../gomsg + require ( github.com/corazawaf/libinjection-go v0.2.2 // indirect github.com/magefile/mage v1.15.1-0.20241126214340-bdc92f694516 // indirect diff --git a/waf/internal/logclient/logclient.go b/waf/internal/logclient/logclient.go index 554bb46..ddeb2ea 100644 --- a/waf/internal/logclient/logclient.go +++ b/waf/internal/logclient/logclient.go @@ -1,142 +1,38 @@ -// Package logclient 提供异步批量日志推送客户端。 +// Package logclient 基于 gomsg 的集中式日志客户端(异步批量推送 db2 Logs 表)。 // -// 通过 pymsg TCP 协议调用 dbv2 execute 路由,逐条 INSERT 到 Logs 表。 -// -// 设计要点(与 Python example_log_client.py 对齐): -// - 异步非阻塞:内存 channel + 后台 goroutine 批量推送,不拖慢业务 -// - fail-safe:db 服务不可用时静默丢弃,绝不影响业务 -// - 缓冲批量:队列满 50 条或超时 2 秒时自动 flush -// -// 协议:pymsg 长度前缀帧(4 字节 big-endian uint32 长度 + UTF-8 载荷) +// 原实现(pymsg 帧协议 + 异步队列)已抽取至 ../gomsg,本包保留包级 API +// 兼容层:main / engine / proxy 的调用点零改动。 package logclient -import ( - "encoding/binary" - "encoding/json" - "fmt" - "io" - "net" - "os" - "sync" - "time" +import "github.com/micro_servers/gomsg" + +var ( + lc *gomsg.LogClient + pendingLevel string ) -const ( - defaultDBAddr = "127.0.0.1:8005" - batchSize = 50 - flushInterval = 2 * time.Second - socketTimeout = 5 * time.Second - // 消息大小安全上限(64 MiB),与 Python / C++ 端保持一致。 - maxMsgSize = 64 * 1024 * 1024 -) - -// 日志级别数值(越小越详细)。 -const ( - levelDebug = 0 - levelInfo = 1 - levelWarn = 2 - levelError = 3 -) - -var levelMap = map[string]int{ - "debug": levelDebug, - "info": levelInfo, - "warn": levelWarn, - "warning": levelWarn, - "error": levelError, - "DEBUG": levelDebug, - "INFO": levelInfo, - "WARN": levelWarn, - "WARNING": levelWarn, - "ERROR": levelError, -} - -// dbAddr 通过环境变量 LOG_DB_ADDR 配置(默认 127.0.0.1:8005),支持非标准部署场景。 -var dbAddr = defaultDBAddr - -// minLevel 控制最低输出级别,默认 info(debug 被过滤)。 -var minLevel = levelInfo - -func init() { - if v := os.Getenv("LOG_DB_ADDR"); v != "" { - dbAddr = v - } - if v := os.Getenv("LOG_LEVEL"); v != "" { - SetLevel(v) +// Init 初始化全局日志客户端(读 LOG_DB_ADDR / LOG_LEVEL 环境变量)。 +func Init() { + lc = gomsg.NewLogClient("") + if pendingLevel != "" { + lc.SetLevel(pendingLevel) } } // SetLevel 设置最低日志级别(debug|info|warn|error,大小写不敏感)。 func SetLevel(level string) { - if n, ok := levelMap[level]; ok { - minLevel = n + pendingLevel = level + if lc != nil { + lc.SetLevel(level) } } -// Entry 单条日志。 -type Entry struct { - Service string `json:"service"` - Level string `json:"level"` - Message string `json:"message"` - TraceID string `json:"trace_id,omitempty"` - Extra map[string]any `json:"extra,omitempty"` -} - -// dbV2Request execute 路由的 INSERT 请求体。 -type dbV2Request struct { - Route string `json:"_route"` - Action string `json:"action"` - Table string `json:"table"` - Values map[string]any `json:"values"` -} - -// dbV2Response execute 路由的通用响应。 -type dbV2Response struct { - Success bool `json:"success"` - Error string `json:"error,omitempty"` -} - -// Client 异步批量日志推送客户端。 -// send 仅在 flushLoop 单 goroutine 中调用,conn 复用无需额外加锁。 -type Client struct { - ch chan Entry - stop chan struct{} - wg sync.WaitGroup - stopped bool - mu sync.Mutex - conn net.Conn // 复用的 pymsg 长连接,失败时重建 -} - -// global 全局单例。 -var global *Client - -// Init 初始化全局日志客户端并启动后台 flush goroutine。 -func Init() { - global = New() -} - -// New 创建客户端实例。 -func New() *Client { - c := &Client{ - ch: make(chan Entry, 256), - stop: make(chan struct{}), - } - c.wg.Add(1) - go c.flushLoop() - return c -} - -// Push 推送一条日志(非阻塞,队列满时丢弃)。 -// 低于 SetLevel 设定的最低级别的日志将被静默丢弃。 +// Push 推送一条日志(非阻塞,队列满或低于最低级别时静默丢弃)。 func Push(service, level, message string, extra map[string]any) { - if global == nil { + if lc == nil { return } - n, ok := levelMap[level] - if !ok || n < minLevel { - return - } - global.push(Entry{ + lc.Push(gomsg.Entry{ Service: service, Level: level, Message: message, @@ -164,199 +60,9 @@ func Error(service, message string, extra map[string]any) { Push(service, "ERROR", message, extra) } -// Flush 同步 flush 队列中剩余日志(用于进程退出前)。 +// Flush 停止队列并 flush 剩余日志(用于进程退出前)。 func Flush() { - if global == nil { - return - } - global.flush() -} - -func (c *Client) push(e Entry) { - select { - case c.ch <- e: - default: - // 队列满,静默丢弃(fail-safe) + if lc != nil { + lc.Flush() } } - -func (c *Client) flushLoop() { - defer c.wg.Done() - batch := make([]Entry, 0, batchSize) - ticker := time.NewTicker(flushInterval) - defer ticker.Stop() - - for { - select { - case e := <-c.ch: - batch = append(batch, e) - // 非阻塞地继续取,直到满 batchSize - for len(batch) < batchSize { - select { - case e2 := <-c.ch: - batch = append(batch, e2) - default: - goto send - } - } - send: - if len(batch) >= batchSize { - c.send(batch) - batch = batch[:0] - } - case <-ticker.C: - if len(batch) > 0 { - c.send(batch) - batch = batch[:0] - } - case <-c.stop: - // 退出前 drain channel - for { - select { - case e := <-c.ch: - batch = append(batch, e) - default: - if len(batch) > 0 { - c.send(batch) - } - return - } - } - } - } -} - -func (c *Client) send(batch []Entry) { - now := time.Now().UTC().Format(time.RFC3339) - for _, entry := range batch { - c.sendOne(entry, now) - } -} - -// sendOne 通过 pymsg TCP 协议发送单条日志到 dbv2。 -func (c *Client) sendOne(entry Entry, ts string) { - extraJSON := "" - if entry.Extra != nil { - b, err := json.Marshal(entry.Extra) - if err != nil { - return - } - extraJSON = string(b) - } - - values := map[string]any{ - "ts": ts, - "service": entry.Service, - "level": entry.Level, - "message": entry.Message, - } - if entry.TraceID != "" { - values["trace_id"] = entry.TraceID - } - if extraJSON != "" { - values["extra"] = extraJSON - } - - req := dbV2Request{ - Route: "execute", - Action: "insert", - Table: "Logs", - Values: values, - } - data, err := json.Marshal(req) - if err != nil { - return - } - - // pymsg 协议:4 字节 big-endian 长度 + UTF-8 载荷(复用长连接,失败时重连一次) - resp, err := c.pymsgCall(data) - if err != nil { - return // fail-safe - } - - var r dbV2Response - if json.Unmarshal(resp, &r) != nil || !r.Success { - fmt.Fprintf(os.Stderr, "[logclient] INSERT failed: %s\n", r.Error) - } -} - -// pymsgCall 复用长连接发送一条 pymsg 请求并读取响应。 -// 首次失败时关闭旧连接重建一次,避免高并发拦截场景下大量 TIME_WAIT 短连接。 -func (c *Client) pymsgCall(payload []byte) ([]byte, error) { - resp, err := c.doPymsg(payload) - if err == nil { - return resp, nil - } - // 旧连接可能已断开(db 重启等),关闭后重连一次 - c.closeConn() - return c.doPymsg(payload) -} - -// doPymsg 在当前连接上执行一次请求-响应交互。 -func (c *Client) doPymsg(payload []byte) ([]byte, error) { - conn, err := c.getConn() - if err != nil { - return nil, err - } - conn.SetDeadline(time.Now().Add(socketTimeout)) - - // 发送:4 字节长度前缀 + 载荷 - header := make([]byte, 4) - binary.BigEndian.PutUint32(header, uint32(len(payload))) - if _, err := conn.Write(header); err != nil { - return nil, err - } - if _, err := conn.Write(payload); err != nil { - return nil, err - } - - // 接收:4 字节长度前缀 + 载荷 - if _, err := io.ReadFull(conn, header); err != nil { - return nil, err - } - respLen := binary.BigEndian.Uint32(header) - if respLen > maxMsgSize { - return nil, fmt.Errorf("response too large: %d bytes (max %d)", respLen, maxMsgSize) - } - resp := make([]byte, respLen) - if _, err := io.ReadFull(conn, resp); err != nil { - return nil, err - } - - return resp, nil -} - -// getConn 返回当前复用连接,不存在时建立。 -func (c *Client) getConn() (net.Conn, error) { - if c.conn != nil { - return c.conn, nil - } - conn, err := net.DialTimeout("tcp", dbAddr, socketTimeout) - if err != nil { - return nil, err - } - c.conn = conn - return conn, nil -} - -// closeConn 关闭并置空复用连接。 -func (c *Client) closeConn() { - if c.conn != nil { - _ = c.conn.Close() - c.conn = nil - } -} - -func (c *Client) flush() { - c.mu.Lock() - if c.stopped { - c.mu.Unlock() - return - } - c.stopped = true - c.mu.Unlock() - - close(c.stop) - c.wg.Wait() - c.closeConn() -}