7

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

 3 years ago
source link: https://studygolang.com/articles/33976
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.

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

JoeLei · 大约7小时之前 · 7 次点击 · 预计阅读时间 6 分钟 · 大约8小时之前 开始浏览    

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

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

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

命令行解析

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

golang的命令行有很多强大的第三方库, 但是定位是小工具,编译的二进制越少约好,所有只用了官方的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是一个全局变量,可以在编译的使用flag指定编译版本

这个函数实现了参数定义,参数校验,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行

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

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

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

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

命令行解析

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

golang的命令行有很多强大的第三方库, 但是定位是小工具,编译的二进制越少约好,所有只用了官方的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是一个全局变量,可以在编译的使用flag指定编译版本

这个函数实现了参数定义,参数校验,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行

-w1118

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


有疑问加站长微信联系(非本文作者)

280

入群交流(和以上内容无关):加入Go大咖交流群,或添加微信:liuxiaoyan-s 备注:入群;或加QQ群:701969077


About Joyk


Aggregate valuable and interesting links.
Joyk means Joy of geeK