学习frp源码之简洁的在两个connection之间转发流量

在阅读frp的代码时,发现这样一个技巧,即使用 io.CopyBuffer 来把两个连接之间的流量互转。

// Join two io.ReadWriteCloser and do some operations.
func Join(c1 io.ReadWriteCloser, c2 io.ReadWriteCloser) (inCount int64, outCount int64) {
    var wait sync.WaitGroup
    pipe := func(to io.ReadWriteCloser, from io.ReadWriteCloser, count *int64) {
        defer to.Close()
        defer from.Close()
        defer wait.Done()

        buf := pool.GetBuf(16 * 1024)
        defer pool.PutBuf(buf)
        *count, _ = io.CopyBuffer(to, from, buf)
    }

    wait.Add(2)
    go pipe(c1, c2, &inCount)
    go pipe(c2, c1, &outCount)
    wait.Wait()
    return
}

可以看看 io.CopyBuffer 的源码:

// CopyBuffer is identical to Copy except that it stages through the
// provided buffer (if one is required) rather than allocating a
// temporary one. If buf is nil, one is allocated; otherwise if it has
// zero length, CopyBuffer panics.
func CopyBuffer(dst Writer, src Reader, buf []byte) (written int64, err error) {
    if buf != nil && len(buf) == 0 {
        panic("empty buffer in io.CopyBuffer")
    }
    return copyBuffer(dst, src, buf)
}

// copyBuffer is the actual implementation of Copy and CopyBuffer.
// if buf is nil, one is allocated.
func copyBuffer(dst Writer, src Reader, buf []byte) (written int64, err error) {
    // If the reader has a WriteTo method, use it to do the copy.
    // Avoids an allocation and a copy.
    if wt, ok := src.(WriterTo); ok {
        return wt.WriteTo(dst)
    }
    // Similarly, if the writer has a ReadFrom method, use it to do the copy.
    if rt, ok := dst.(ReaderFrom); ok {
        return rt.ReadFrom(src)
    }
    if buf == nil {
        size := 32 * 1024
        if l, ok := src.(*LimitedReader); ok && int64(size) > l.N {
            if l.N < 1 {
                size = 1
            } else {
                size = int(l.N)
            }
        }
        buf = make([]byte, size)
    }
    for {
        nr, er := src.Read(buf)
        if nr > 0 {
            nw, ew := dst.Write(buf[0:nr])
            if nw > 0 {
                written += int64(nw)
            }
            if ew != nil {
                err = ew
                break
            }
            if nr != nw {
                err = ErrShortWrite
                break
            }
        }
        if er != nil {
            if er != EOF {
                err = er
            }
            break
        }
    }
    return written, err
}

可以看到,在连接关闭之前,一直都在 for 循环。而连接关闭时,读取到的就是 io.EOF。因此,可以通过这么一个简洁的函数优雅 的完成在两个函数之间进行流量转发的功能。这段代码确实值得学习。


更多文章