golang实现TCP代理

前言

TCP代理的使用场景有很多，比如Nginx的http代理，本质上也是TCP的转发，微服务网格istio的Envoy也是基于这种理念来实现的。

之所以写这个demo，是为了可以很方便地调试上下游服务在通信过程中到底传递了什么数据，对各种协议的通信方式可以很快的掌握它们原理，比如rpc通信、http通信、grpc通信等等。

代码路径：http://gitee.com/zqwlai/go-test/tcpProxy.go

package main

import (
    "bufio"
    "flag"
    "fmt"
    "github.com/rs/zerolog"
    "net"
    "os"
    "strconv"
    "strings"

    "time"
)

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

func main() {
    help := flag.Bool("help", false, "print usage")
    bind := flag.String("bind", "0.0.0.0:9528", "The address to bind to")
    backend := flag.String("backend", "127.0.0.1:30000", "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, 1)

        go Proxy(conn, target, closed)
        go Proxy2(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
        }
        fmt.Println(1111, n1)
        fmt.Println(222, string(buffer[:n1]))
        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
        }
    }

}


func Proxy2(from net.Conn, to net.Conn, closed chan bool) {
    time.Sleep(1*time.Second)
    buffer := make([]byte, 4096)
    for {
        n1, err := from.Read(buffer)
        if err != nil {
            closed <- true
            return
        }
        fmt.Println(3333, string(buffer[:n1]))
        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
        }
    }
}

//解析HTTP报文
func parseHttp(conn net.Conn)(string){

    var s = ""
    reader := bufio.NewReader(conn)
    firstLine,_,_ := reader.ReadLine()
    s += string(firstLine) + "\n"

    tempList := strings.Split(string(firstLine), " ")
    method := tempList[0]
    path := tempList[1]
    proto := tempList[2]
    fmt.Println(1111, method, path, proto)
    var headers = make(map[string]string)
    var contenLength int
    for {
        line,_,_ := reader.ReadLine()

        s += string(line) + "\n"
        if len(line) == 0{
            break
        }
        tempList = strings.Split(string(line), ":")
        key := tempList[0]
        value := strings.Trim(tempList[1], " ")
        headers[key] = value
        if (key == "Content-Length"){
            contenLength,_ = strconv.Atoi(value)
        }
    }

    fmt.Println(2222, headers)
    if contenLength == 0{
        return s
    }
    pack := make([]byte, contenLength)
    t,_ := reader.Read(pack)
    s += string(pack[:t])

    return s
}