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

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

参数解析日志异常也没有处理

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

数据转发

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

mainmain

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路由器下运行都没有问题

使用示例：

-h-v