使用场景就不说了,可以支持任意TCP网络数据转发

用 google 搜索,有很多这样的代码片段,但是作为一个小工具,都不完整,比如没有参数解析,打印都是print,而不是规范的日志,有些异常也没有处理

下面通过104行代码,完成一个麻雀虽小,但是五脏六腑都齐全的小工具

100行代码实现一个高性能网络转发小工具

使用场景就不说了,可以支持任意TCP网络数据转发

用 google 搜索,有很多这样的代码片段,但是作为一个小工具,都不完整,比如没有参数解析,打印都是print,而不是规范的日志,有些异常也没有处理

下面通过104行代码,完成一个麻雀虽小,但是五脏六腑都齐全的小工具

命令行解析

对于一个命令行工具,参数解析是第一步, 也可能是用户使用你的工具交互的首选途径,不管其他参数如何,总的需要一个-v,打印下程序的版本吧

golang的命令行有很多强大的第三方库, 比如 cobra , kingpin 等,但是既然定位是小工具,编译的二进制越少约好,所有只用了官方的flag实现

var (
	version string
)

func ParseArgs() (string, string) {
	listenAddr := flag.String("l", ":8080", "listen address")
	forwardAddr := flag.String("f", "", "forwarding address")
	flagVersion := flag.Bool("v", false, "print version")

	flag.Parse()

	if *flagVersion {
		fmt.Println("version:", version)
		os.Exit(0)
	}

	if *forwardAddr == "" {
		flag.Usage()
		os.Exit(0)
	}

	return *listenAddr, *forwardAddr
}

version 是一个全局变量,可以在编译时使用 -ldflags "-X main.version={{now | date "2006-01-02T15:04:05"}}" 指定工具版本

这个函数实现了参数定义,参数校验,Usage 打印等,基本满足小工具的使用了

TCP Serve

func ListenAndServe(listenAddr string, forwardAddr string) {
	ln, err := net.Listen("tcp", listenAddr)
	if err != nil {
		log.Fatalf("listen addr %s failed: %s", listenAddr, err.Error())
	}

	log.Printf("accept %s to %s\n", listenAddr, forwardAddr)

	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Printf("accept %s error: %s\n", listenAddr, err.Error())
		}

		go HandleRequest(conn, forwardAddr)
	}
}

对于接受网络请求,需要启动一个TCP服务,这里需要处理下端口冲突异常,启动日志等,最后通过一个死循环,对于每一个请求,启动一个goroute 处理

golang实现就是这么简单

请求处理

对于请求,HTTP 服务一般是对象 Request 做处理,返回一个 Response,这里实现也是类似

只是我们是网络转发,所有先通过net.Dialer拨号,连接到远程服务器再说

这里要注意下拨号超时,搜索了很多代码片段,清一色的net.Dial直接使用,对应远端是个错误地址也不知道,所以注意加下超时错误日志

func HandleRequest(conn net.Conn, forwardAddr string) {
	d := net.Dialer{Timeout: time.Second * 10}

	proxy, err := d.Dial("tcp", forwardAddr)
	if err != nil {
		log.Printf("try connect %s -> %s failed: %s\n", conn.RemoteAddr(), forwardAddr, err.Error())
		conn.Close()
		return
	}
	log.Printf("connected: %s -> %s\n", conn.RemoteAddr(), forwardAddr)

	Pipe(conn, proxy)
}

作为透明的数据转发,连接后,就需要转发数据,这里单独用一个函数Pipe处理io请求

数据转发

func Pipe(src net.Conn, dest net.Conn) {
	var (
		readBytes  int64
		writeBytes int64
	)
	ts := time.Now()

	wg := sync.WaitGroup{}
	wg.Add(1)

	closeFun := func(err error) {
		dest.Close()
		src.Close()
	}

	go func() {
		defer wg.Done()
		n, err := io.Copy(dest, src)
		readBytes += n
		closeFun(err)
	}()

	n, err := io.Copy(src, dest)
	writeBytes += n
	closeFun(err)

	wg.Wait()
	log.Printf("connection %s -> %s closed: readBytes %d, writeBytes %d, duration %s", src.RemoteAddr(), dest.RemoteAddr(), readBytes, writeBytes, time.Now().Sub(ts))
}

命令行入口 - main

没啥好说点, 固定的函数名main,固定的package main, 调用参数解析,直接启动服务即可

func main() {
	listenAddr, forwardAddr := ParseArgs()
	ListenAndServe(listenAddr, forwardAddr)
}

结束

包含import, 空号,一共104行, 完整代码 Gist 查看 -> 传送门

编译: Taskfile.yml 格式

build:
    desc: Build the go binary.
    cmds:
      - go build -tags netgo -ldflags "-X main.version={{now | date "2006-01-02T15:04:05"}}" -v -o build/thanos -i cmd/forwarding/forwarding.go

大小 ~ 3Mb,果然小巧, 而且自带跨平台天赋,在 mac, linus, windows, arm路由器下运行都没有问题

使用示例: -w1118

-h 打印命令行标准, -v 打印版本,可以看到转发请求和io统计,时间等