Go语言中如何在流复制前读取TCP Conn的数据？

要实现读取p1数据同时不影响反向代理流复制，核心是通过io.TeeReader将p1的数据流分流：一份用于自定义处理，一份继续流向目标连接。你之前尝试失败的原因是仍在使用原p1连接作为复制源，而io.TeeReader已经读取了原连接的数据流，导致原连接后续读取直接返回EOF。

具体实现步骤 #

1. 定义数据处理器：实现io.Writer接口，在Write方法中处理从p1读取到的数据
2. 用io.TeeReader包装p1：将p1的数据流分流到处理器，同时保留完整的复制流
3. 替换原复制源：使用包装后的TeeReader作为p1到p2的复制源，而非原p1连接

修改后的代码 #

import (
    "io"
    "log"
    "net"
)

// 数据处理器，实现io.Writer接口处理p1的数据
type P1DataProcessor struct{}

func (p *P1DataProcessor) Write(b []byte) (int, error) {
    // 在这里编写你的数据处理逻辑，比如打印、解析、存储等
    log.Printf("读取到p1的数据：%s", string(b))
    // 必须返回完整的字节长度和nil，否则会中断流复制
    return len(b), nil
}

func handleClient(p2, p1 net.Conn) {
    streamCopy := func(dst io.Writer, src io.Reader) {
        if _, err := generic.Copy(dst, src); err != nil {
            log.Println(err, "in:", p1.RemoteAddr(), "out:", fmt.Sprint(p2.RemoteAddr(), "(", p2.ID(), ")"))
        }
        p1.Close()
        p2.Close()
    }

    // 初始化数据处理器，包装p1的数据流
    processor := &P1DataProcessor{}
    teeP1 := io.TeeReader(p1, processor)

    // p1 -> p2 的流使用包装后的TeeReader，同时完成数据读取和复制
    go streamCopy(p2, teeP1)
    // p2 -> p1 的流正常处理
    streamCopy(p1, p2)
}

// 你的generic.Copy函数保持不变
func Copy(dst io.Writer, src io.Reader) (written int64, err error) {
    const bufSize = 32 * 1024
    if wt, ok := src.(io.WriterTo); ok {
        log.Println("io.WriterTo")
        return wt.WriteTo(dst)
    }
    if rt, ok := dst.(io.ReaderFrom); ok {
        log.Println("io.ReaderFrom")
        return rt.ReadFrom(src)
    }

    buf := make([]byte, bufSize)
    return io.CopyBuffer(dst, src, buf)
}

关键说明 #

• io.TeeReader的作用是：每次从它读取数据时，会将读取到的内容同时写入绑定的Writer（即我们的处理器），且不会修改原数据流的内容和顺序
• 必须使用teeP1作为p1到p2的复制源，不能再直接使用原p1连接，否则会因数据流已被读取而导致复制失败
• 如果需要修改p1的数据而非仅读取，需自定义实现io.Reader接口，在Read方法中修改数据后返回，而非使用io.TeeReader

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