如何在 JSON-RPC 方法中获取 WebSocket 连接信息

在 go 的 `jsonrpc` + `websocket` 组合中，标准 rpc 方法无法直接访问底层 `*websocket.conn`。本文介绍一种基于 `context.context` 和自定义 `servercodec` 的安全、可扩展方案，使 rpc 处理函数能获取连接元数据（如远程地址、header 等）。

go 标准库的.net/rpc/jsonrpc 是面向无状态通信设计的，其 ServeConn 接口将 *websocket.Conn 封装后完全隔离，导致 RPC 方法（如 Multiply）无法感知调用来源。强行打破这一抽象虽可行，但会牺牲可维护性与兼容性。推荐采用 上下文注入（Context Injection） 方式，在不修改 RPC 协议语义的前提下，安全传递连接信息。

✅ 核心思路：通过自定义 ServerCodec 注入 context.Context

我们需实现一个符合 rpc.ServerCodec 接口的编码器，在反序列化请求参数时，自动将携带连接信息的 context.Context 注入到参数结构体的 Context 字段中。

1. 定义支持上下文的参数结构体

import (     "context"     "net/http"     "net/rpc"     "net/rpc/jsonrpc"     "reflect"     "sync" )  type Args struct {     A       int     B       int     Context context.Context // 新增字段，用于接收注入的上下文 }

2. 实现自定义 ServerCodec

以下是一个轻量级 WebSocketContextCodec，它包装原始 jsonrpc.ServerCodec，并在 ReadRequestBody 阶段注入 context.WithValue：

type WebSocketContextCodec struct {     codec rpc.ServerCodec     ctx   context.Context }  func (c *WebSocketContextCodec) ReadRequestHeader(r *rpc.Request) error {     return c.codec.ReadRequestHeader(r) }  func (c *WebSocketContextCodec) ReadRequestBody(x interface{}) error {     err := c.codec.ReadRequestBody(x)     if err != nil {         return err     }      // 使用反射向 x 中的 Context 字段注入上下文     v := reflect.ValueOf(x)     if v.Kind() == reflect.Ptr {         v = v.Elem()         if v.Kind() == reflect.Struct {             ctxField := v.FieldByName("Context")             if ctxField.IsValid() && ctxField.CanSet() && ctxField.Type() == reflect.TypeOf((*context.Context)(nil)).Elem().Elem() {                 ctxField.Set(reflect.ValueOf(c.ctx))             }         }     }     return nil }  func (c *WebSocketContextCodec) WriteResponse(r *rpc.Response, body interface{}) error {     return c.codec.WriteResponse(r, body) }  func (c *WebSocketContextCodec) Close() error {     return c.codec.Close() }

3. 在 WebSocket 处理中创建带上下文的 Codec

修改 serve 函数，为每个连接创建专属 context.Context（例如包含远程地址、握手 Header 等）：

import "code.google.com/p/go.net/websocket"  func serve(ws *websocket.Conn) {     // 构建连接上下文：可携带 ws.RemoteAddr(), ws.Config().Origin, 或自定义元数据     connCtx := context.WithValue(         context.Background(),         "websocket.conn",         ws,     )     connCtx = context.WithValue(connCtx, "remote_addr", ws.RemoteAddr().String())      // 创建自定义 codec     codec := &WebSocketContextCodec{         codec: jsonrpc.NewServerCodec(ws),         ctx:   connCtx,     }      // 使用自定义 codec 启动 RPC 服务     rpc.ServeCodec(codec) }

4. 在 RPC 方法中使用连接信息

现在 Multiply 可直接访问连接上下文：

func (t *Arith) Multiply(args *Args, reply *int) error {     *reply = args.A * args.B      // ✅ 安全获取 WebSocket 连接对象和元数据     if conn, ok := args.Context.Value("websocket.conn").(*websocket.Conn); ok {         // 例如：记录日志或做权限校验         println("RPC called from:", conn.RemoteAddr())     }      if addr, ok := args.Context.Value("remote_addr").(string); ok {         println("Client IP:", addr)     }      return nil }

⚠️ 注意事项与最佳实践

• 线程安全：context.Context 本身是并发安全的，但注入的值（如 *websocket.Conn）需确保调用方不执行阻塞/写操作；建议仅读取元数据（如 RemoteAddr, Config().Origin），避免在 RPC 中调用 ws.Write()。
• 性能影响：反射注入仅在每次请求时执行一次，开销极小；若追求极致性能，可预编译字段索引（如用 sync.Once 缓存 FieldByName 结果）。
• 兼容性：该方案完全兼容标准 jsonrpc 协议，客户端无需任何改动。
• 替代方案对比：不推荐使用全局 map + goroutine ID 模拟 Thread-local —— Go 不提供稳定 goroutine ID，且易引发内存泄漏；context 是 Go 官方推荐的跨 API 边界传递请求范围数据的标准方式。

通过此方法，你既保持了 RPC 层的清晰抽象，又获得了对底层连接的精细控制能力，是构建高可用、可观测 WebSocket-RPC 服务的关键一环。