使用golang網(wǎng)絡(luò)編程實現(xiàn)一個簡單的TCP代理（不支持HTTP）

package main

import (
    "flag"
    "github.com/rs/zerolog"
    "net"
    "os"
)

var logger = zerolog.New(os.Stdout).With().Timestamp().Logger()

func main() {
    help := flag.Bool("help", false, "print usage")
    bind := flag.String("bind", "127.0.0.1:6000", "The address to bind to")
    backend := flag.String("backend", "", "The backend server address")
    flag.Parse()

    logger.Level(zerolog.DebugLevel)

    if *help {
        flag.Usage()
        return
    }

    if *backend == "" {
        flag.Usage()
        return
    }

    if *bind == "" {
        //use default bind
        logger.Info().Str("bind", *bind).Msg("use default bind")
    }

    success, err := RunProxy(*bind, *backend)
    if !success {
        logger.Error().Err(err).Send()
        os.Exit(1)
    }
}

func RunProxy(bind, backend string) (bool, error) {
    listener, err := net.Listen("tcp", bind)
    if err != nil {
        return false, err
    }
    defer listener.Close()
    logger.Info().Str("bind", bind).Str("backend", backend).Msg("tcp-proxy started.")
    for {
        conn, err := listener.Accept()
        if err != nil {
            logger.Error().Err(err).Send()
        } else {
            go ConnectionHandler(conn, backend)
        }
    }
}

func ConnectionHandler(conn net.Conn, backend string) {
    logger.Info().Str("conn", conn.RemoteAddr().String()).Msg("client connected.")
    target, err := net.Dial("tcp", backend)
    defer conn.Close()
    if err != nil {
        logger.Error().Err(err).Send()
    } else {
        defer target.Close()
        logger.Info().Str("conn", conn.RemoteAddr().String()).Str("backend", target.LocalAddr().String()).Msg("backend connected.")
        closed := make(chan bool, 2)
        go Proxy(conn, target, closed)
        go Proxy(target, conn, closed)
        <-closed
        logger.Info().Str("conn", conn.RemoteAddr().String()).Msg("Connection closed.")
    }
}

func Proxy(from net.Conn, to net.Conn, closed chan bool) {
    buffer := make([]byte, 4096)
    for {
        n1, err := from.Read(buffer)
        if err != nil {
            closed <- true
            return
        }
        n2, err := to.Write(buffer[:n1])
        logger.Debug().Str("from", from.RemoteAddr().String()).Int("recv", n1).Str("to", to.RemoteAddr().String()).Int("send", n2).Send()
        if err != nil {
            closed <- true
            return
        }
    }
}

使用方式

代理監(jiān)聽9000端口，代理后端服務(wù)的8000端口

tcp-proxy --bind 0.0.0.0:9000 --backend 127.0.0.1:8000

源碼解析

這個簡單的小程序主要由三個函數(shù)構(gòu)成
1.RunProxy 啟動代理服務(wù)，監(jiān)聽bind參數(shù)指定的端口，接收客戶端請求
2.ConnectionHandler 客戶端請求處理，連接backend服務(wù)
3.Proxy 數(shù)據(jù)傳輸代理，將客戶端數(shù)據(jù)發(fā)送到backend服務(wù)，將backend數(shù)據(jù)發(fā)送給客戶端