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
}

 

 

 

简单说下原理，proxy在收到下游的请求后，会建立一个conn1，此时proxy会解析conn1里面的目标地址，并与之建立TCP连接conn2， 接下来会开启两个协程，一个用来从conn1里读取请求报文，并将报文发给conn2，另一个协程从conn2读取响应数据，并写到conn1返回给上游。