学习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。因此,可以通过这么一个简洁的函数优雅 的完成在两个函数之间进行流量转发的功能。这段代码确实值得学习。


更多文章
  • C, Go, Python的错误处理和异常机制杂谈
  • 好的命名是最好的文档
  • 读《系统之美:决策者的系统思考》
  • Linux高分屏支持
  • GCC默认的头文件搜索路径
  • 读《远见-如何规划职业生涯3大阶段》
  • 后端工程师学前端(五): SASS
  • 后端工程师学前端(四): CSS进阶(盒子模型)
  • 读《投资中最简单的事》
  • 后端工程师学前端(三): CSS进阶(特指度、单位和字体族)
  • 后端工程师学前端(二): CSS基础知识(规则与选择器)
  • Swift语法笔记
  • 后端工程师学前端(一): HTML
  • 读《管理的实践》
  • frp 源码阅读与分析(二):TCP内网穿透的实现