Golang net/rpc中获取客户端真实IP及当前net.Conn的方法

需求说明 #

• RPC服务器需获取客户端的真实IP（NAT转发场景下也能正确获取）
• 需要拿到当前RPC调用对应的net.Conn，用于后续向特定IP客户端发起双向RPC请求（如给Agent下发脚本执行指令）

原有实现局限 #

原代码在Accept连接时通过conn.RemoteAddr()获取IP，但NAT场景下只能拿到NAT网关的IP，无法得到客户端真实IP；且RPC方法内部无法直接获取当前连接的net.Conn。示例代码如下：

// example_server.go
type MathService struct{}
type Args struct {
    A, B int
}
type Reply struct {
    Result int
}
func (m *MathService) Multiply(args *Args, reply *Reply) error {
    reply.Result = args.A * args.B
    return nil
}
func main() {
    mathService := new(MathService)
    rpc.Register(mathService)
    listener, err := net.Listen("tcp", "0.0.0.0:1234")
    if err != nil {
        fmt.Println("Error starting server:", err)
        return
    }
    defer listener.Close()
    fmt.Println("Server listening on :1234")
    for {
        conn, err := listener.Accept()
        if err != nil {
            fmt.Println("Error accepting connection:", err)
            continue
        }
        fmt.Printf("remote ip addr: %s\n", conn.RemoteAddr().String())
        go rpc.ServeConn(conn)
    }
}

// example_client.go
type Args struct {
    A, B int
}
type Reply struct {
    Result int
}
func main() {
    client, err := rpc.Dial("tcp", "127.0.0.1:1234")
    if err != nil {
        fmt.Println("Error connecting to server:", err)
        return
    }
    defer client.Close()
    args := &Args{A: 5, B: 3}
    var reply Reply
    err = client.Call("MathService.Multiply", args, &reply)
    if err != nil {
        fmt.Println("Error calling Multiply:", err)
        return
    }
    fmt.Printf("Result: %d\n", reply.Result)
}

期望方案痛点 #

希望客户端通过调用ReportRealIP上报真实IP，但RPC方法内部无法获取当前调用对应的net.Conn，无法将真实IP与连接关联存储。伪代码如下：

// client.go
rpc.Call("Server.ReportRealIP","1,2,3,4", resp)

// server.go
var globalMap map[string]net.Conn
func (m *Server) ReportRealIP(args *string, reply *Reply) error {
    clientRealIP := args
    // 无法获取当前调用的net.Conn
    globalMap[clientRealIP] = current net.Conn
    return nil
}

解决方案：自定义Codec注入连接信息 #

通过自定义ServerCodec和连接绑定的服务实例，实现RPC方法内获取net.Conn，并关联客户端上报的真实IP。

1. 定义带连接信息的基础服务 #

创建基础结构体存储当前连接，所有需要访问连接的服务嵌入该结构体：

import (
    "net"
    "net/rpc"
    "sync"
    "fmt"
)

type BaseService struct {
    Conn net.Conn
}

var globalConnMap = make(map[string]net.Conn)
var mu sync.Mutex // 保护全局连接map的并发读写

2. 实现自定义ServerCodec #

包装标准gobCodec，确保连接信息能传递到服务实例：

type customCodec struct {
    rpc.ServerCodec
    conn net.Conn
}

func NewCustomCodec(conn net.Conn) rpc.ServerCodec {
    return &customCodec{
        ServerCodec: rpc.NewServerCodec(conn),
        conn:        conn,
    }
}

// 复用标准Codec的读写逻辑
func (c *customCodec) ReadRequestHeader(r *rpc.Request) error {
    return c.ServerCodec.ReadRequestHeader(r)
}
func (c *customCodec) ReadRequestBody(body interface{}) error {
    return c.ServerCodec.ReadRequestBody(body)
}
func (c *customCodec) WriteResponse(r *rpc.Response, body interface{}) error {
    return c.ServerCodec.WriteResponse(r, body)
}
func (c *customCodec) Close() error {
    return c.ServerCodec.Close()
}

3. 实现带连接绑定的服务及上报逻辑 #

每个连接创建独立的服务实例，绑定当前net.Conn，并实现ReportRealIP方法：

type ServerService struct {
    BaseService
}

type RealIPArgs struct {
    RealIP string
}

type EmptyReply struct{}

// 客户端调用此方法上报真实IP
func (s *ServerService) ReportRealIP(args *RealIPArgs, reply *EmptyReply) error {
    mu.Lock()
    defer mu.Unlock()
    // 将真实IP与当前连接关联存储
    globalConnMap[args.RealIP] = s.Conn
    fmt.Printf("绑定真实IP %s 到连接 %s\n", args.RealIP, s.Conn.RemoteAddr().String())
    return nil
}

// 保留原有业务方法示例
func (s *ServerService) Multiply(args *Args, reply *Reply) error {
    reply.Result = args.A * args.B
    return nil
}

4. 服务器主逻辑修改 #

每个连接创建独立的RPC服务器实例，绑定当前连接的服务：

func main() {
    listener, err := net.Listen("tcp", "0.0.0.0:1234")
    if err != nil {
        fmt.Println("Error starting server:", err)
        return
    }
    defer listener.Close()
    fmt.Println("Server listening on :1234")

    for {
        conn, err := listener.Accept()
        if err != nil {
            fmt.Println("Error accepting connection:", err)
            continue
        }
        // 为每个连接创建绑定了当前Conn的服务实例
        serverService := &ServerService{
            BaseService: BaseService{Conn: conn},
        }
        // 创建独立RPC服务器，避免全局注册冲突
        rpcServer := rpc.NewServer()
        rpcServer.Register(serverService)
        // 使用自定义Codec处理连接
        go rpcServer.ServeCodec(NewCustomCodec(conn))
    }
}

5. 客户端上报真实IP示例 #

func main() {
    client, err := rpc.Dial("tcp", "127.0.0.1:1234")
    if err != nil {
        fmt.Println("Error connecting to server:", err)
        return
    }
    defer client.Close()

    // 上报真实IP（实际场景中客户端可从HTTP头、配置等渠道获取）
    realIP := "192.168.1.100"
    args := &RealIPArgs{RealIP: realIP}
    var reply EmptyReply
    err = client.Call("ServerService.ReportRealIP", args, &reply)
    if err != nil {
        fmt.Println("Error reporting real IP:", err)
        return
    }
    fmt.Println("Real IP reported successfully")

    // 调用业务方法示例
    mathArgs := &Args{A: 5, B: 3}
    var mathReply Reply
    err = client.Call("ServerService.Multiply", mathArgs, &mathReply)
    if err != nil {
        fmt.Println("Error calling Multiply:", err)
        return
    }
    fmt.Printf("Multiply result: %d\n", mathReply.Result)
}

6. 双向RPC调用示例 #

通过真实IP获取连接，向客户端发起RPC请求：

// 假设客户端已注册ClientService服务
type ScriptArgs struct {
    Script string
}
type ScriptReply struct {
    Output string
}

func sendCommandToClient(targetIP string) error {
    mu.Lock()
    conn, ok := globalConnMap[targetIP]
    mu.Unlock()
    if !ok {
        return fmt.Errorf("no connection found for IP %s", targetIP)
    }
    // 复用已有连接创建RPC客户端
    client := rpc.NewClient(conn)
    defer client.Close()

    // 调用客户端的ExecuteScript方法
    args := &ScriptArgs{Script: "echo 'hello from server'"}
    var reply ScriptReply
    err := client.Call("ClientService.ExecuteScript", args, &reply)
    if err != nil {
        return err
    }
    fmt.Printf("Client %s executed script, output: %s\n", targetIP, reply.Output)
    return nil
}

关键说明 #

• 每个连接对应独立的RPC服务器和服务实例，避免全局注册冲突，确保服务实例能拿到当前连接的net.Conn
• 用互斥锁保护全局连接map，解决并发读写问题
• 客户端上报真实IP后，服务器将IP与连接绑定，后续可通过真实IP找到连接，实现双向RPC通信

内容的提问来源于stack exchange，提问作者moluzhui