teeporxy.go

package main

import (

    "bytes"

    "crypto/tls"

    "flag"

    "fmt"

    "io"

    "io/ioutil"

    "math/rand"

    "net"

    "net/http"

    "net/http/httputil"

    "runtime"

    "time"

)

// Console flags

//参数解析

var (

    listen                = flag.String("l", ":8888", "port to accept requests")  //接收请求端口 默认渡口是8888

    targetProduction      = flag.String("a", "localhost:8080", "where production traffic goes. http://localhost:8080/production")  //a代表产品机器  默认端口是8080

    altTarget             = flag.String("b", "localhost:8081", "where testing traffic goes. response are skipped. http://localhost:8081/test")  //b 测试机器 端口是8081 

    debug                 = flag.Bool("debug", false, "more logging, showing ignored output")  //日志开关

    productionTimeout     = flag.Int("a.timeout", 3, "timeout in seconds for production traffic")// 生产机器请求超时时间

    alternateTimeout      = flag.Int("b.timeout", 1, "timeout in seconds for alternate site traffic")//测试机器清酒超时时间

    productionHostRewrite = flag.Bool("a.rewrite", false, "rewrite the host header when proxying production traffic") //生产机器是重定向开关  

    alternateHostRewrite  = flag.Bool("b.rewrite", false, "rewrite the host header when proxying alternate site traffic")//测试机器是否重定向开关

    percent               = flag.Float64("p", 100.0, "float64 percentage of traffic to send to testing")// 生产数据发给测试机器数据的百分比  流量分割

    tlsPrivateKey         = flag.String("key.file", "", "path to the TLS private key file") //TSL 私钥证书

    tlsCertificate        = flag.String("cert.file", "", "path to the TLS certificate file")//Tsl 龚玥证书

)

// handler contains the address of the main Target and the one for the Alternative target

//handler 包含连个地址  其中一个是生产服务器  另个一是测试服务器

type handler struct {

    Target      string

    Alternative string

    Randomizer  rand.Rand

}

// ServeHTTP duplicates the incoming request (req) and does the request to the Target and the Alternate target discading the Alternate response

//sereHttp 复制获取到的req 并且发送到生产服务器和测试服务器   测试服务器丢弃响应结果

func (h handler) ServeHTTP(w http.ResponseWriter, req *http.Request) {

    var productionRequest, alternativeRequest *http.Request

    if *percent == 100.0 || h.Randomizer.Float64()*100 < *percent {

        alternativeRequest, productionRequest = DuplicateRequest(req) //复制数据到生产和测试请求中

        go func() {

            defer func() {

                if r := recover(); r != nil && *debug {

                    fmt.Println("Recovered in f", r)

                }

            }()

            // Open new TCP connection to the server

                      //获取客户端连接 带有超时时间   

            clientTcpConn, err := net.DialTimeout("tcp", h.Alternative, time.Duration(time.Duration(*alternateTimeout)*time.Second))

            if err != nil {

                if *debug {

                    fmt.Printf("Failed to connect to %s\n", h.Alternative)

                }

                return

            }

            clientHttpConn := httputil.NewClientConn(clientTcpConn, nil) // Start a new HTTP connection on it

            defer clientHttpConn.Close()                                 // Close the connection to the server

            if *alternateHostRewrite {

                alternativeRequest.Host = h.Alternative

            }

            err = clientHttpConn.Write(alternativeRequest) // Pass on the request

            if err != nil {

                if *debug {

                    fmt.Printf("Failed to send to %s: %v\n", h.Alternative, err)

                }

                return

            }

            _, err = clientHttpConn.Read(alternativeRequest) // Read back the reply

            if err != nil {

                if *debug {

                    fmt.Printf("Failed to receive from %s: %v\n", h.Alternative, err)

                }

                return

            }

        }()

    } else {

        productionRequest = req

    }

    defer func() {

        if r := recover(); r != nil && *debug {

            fmt.Println("Recovered in f", r)

        }

    }()

    // Open new TCP connection to the server

       //生产服务器

    clientTcpConn, err := net.DialTimeout("tcp", h.Target, time.Duration(time.Duration(*productionTimeout)*time.Second))

    if err != nil {

        fmt.Printf("Failed to connect to %s\n", h.Target)

        return

    }

    clientHttpConn := httputil.NewClientConn(clientTcpConn, nil) // Start a new HTTP connection on it

    defer clientHttpConn.Close()                                 // Close the connection to the server

    if *productionHostRewrite {

        productionRequest.Host = h.Target

    }

    err = clientHttpConn.Write(productionRequest) // Pass on the request

    if err != nil {

        fmt.Printf("Failed to send to %s: %v\n", h.Target, err)

        return

    }

    resp, err := clientHttpConn.Read(productionRequest) // Read back the reply

    if err != nil {

        fmt.Printf("Failed to receive from %s: %v\n", h.Target, err)

        return

    }

    defer resp.Body.Close()

    for k, v := range resp.Header {

        w.Header()[k] = v

    }

    w.WriteHeader(resp.StatusCode)

    body, _ := ioutil.ReadAll(resp.Body)

    w.Write(body)

}

func main() {

    flag.Parse()

    runtime.GOMAXPROCS(runtime.NumCPU())

    var err error

    var listener net.Listener

    if len(*tlsPrivateKey) > 0 {

        cer, err := tls.LoadX509KeyPair(*tlsCertificate, *tlsPrivateKey)

        if err != nil {

            fmt.Printf("Failed to load certficate: %s and private key: %s", *tlsCertificate, *tlsPrivateKey)

            return

        }

        config := &tls.Config{Certificates: []tls.Certificate{cer}}

        listener, err = tls.Listen("tcp", *listen, config)

        if err != nil {

            fmt.Printf("Failed to listen to %s: %s\n", *listen, err)

            return

        }

    } else {

        listener, err = net.Listen("tcp", *listen)

        if err != nil {

            fmt.Printf("Failed to listen to %s: %s\n", *listen, err)

            return

        }

    }

    h := handler{

        Target:      *targetProduction,

        Alternative: *altTarget,

        Randomizer:  *rand.New(rand.NewSource(time.Now().UnixNano())),

    }

    http.Serve(listener, h)

}

type nopCloser struct {

    io.Reader

}

func (nopCloser) Close() error { return nil }

//复制req到生茶服务器和测试服务器

func DuplicateRequest(request *http.Request) (request1 *http.Request, request2 *http.Request) {

    b1 := new(bytes.Buffer)

    b2 := new(bytes.Buffer)

    w := io.MultiWriter(b1, b2)  //同时向多个对象中写入数据

    io.Copy(w, request.Body) //复制数据到  w中

    defer request.Body.Close()

    request1 = &http.Request{

        Method:        request.Method,

        URL:           request.URL,

        Proto:         request.Proto,

        ProtoMajor:    request.ProtoMajor,

        ProtoMinor:    request.ProtoMinor,

        Header:        request.Header,

        Body:          nopCloser{b1},

        Host:          request.Host,

        ContentLength: request.ContentLength,

    }

    request2 = &http.Request{

        Method:        request.Method,

        URL:           request.URL,

        Proto:         request.Proto,

        ProtoMajor:    request.ProtoMajor,

        ProtoMinor:    request.ProtoMinor,

        Header:        request.Header,

        Body:          nopCloser{b2},

        Host:          request.Host,

        ContentLength: request.ContentLength,

    }

    return

}