add better generate

Signed-off-by: Jess Frazelle <acidburn@microsoft.com>

vendor/github.com/docker/go-connections/proxy/tcp_proxy.go

@@ -0,0 +1,105 @@
+package proxy
+
+import (
+	"io"
+	"net"
+	"syscall"
+)
+
+// TCPProxy is a proxy for TCP connections. It implements the Proxy interface to
+// handle TCP traffic forwarding between the frontend and backend addresses.
+type TCPProxy struct {
+	Logger       logger
+	listener     *net.TCPListener
+	frontendAddr *net.TCPAddr
+	backendAddr  *net.TCPAddr
+}
+
+// NewTCPProxy creates a new TCPProxy.
+func NewTCPProxy(frontendAddr, backendAddr *net.TCPAddr, ops ...func(*TCPProxy)) (*TCPProxy, error) {
+	listener, err := net.ListenTCP("tcp", frontendAddr)
+	if err != nil {
+		return nil, err
+	}
+	// If the port in frontendAddr was 0 then ListenTCP will have a picked
+	// a port to listen on, hence the call to Addr to get that actual port:
+	proxy := &TCPProxy{
+		listener:     listener,
+		frontendAddr: listener.Addr().(*net.TCPAddr),
+		backendAddr:  backendAddr,
+		Logger:       &noopLogger{},
+	}
+
+	for _, op := range ops {
+		op(proxy)
+	}
+
+	return proxy, nil
+}
+
+func (proxy *TCPProxy) clientLoop(client *net.TCPConn, quit chan bool) {
+	backend, err := net.DialTCP("tcp", nil, proxy.backendAddr)
+	if err != nil {
+		proxy.Logger.Printf("Can't forward traffic to backend tcp/%v: %s\n", proxy.backendAddr, err)
+		client.Close()
+		return
+	}
+
+	event := make(chan int64)
+	var broker = func(to, from *net.TCPConn) {
+		written, err := io.Copy(to, from)
+		if err != nil {
+			// If the socket we are writing to is shutdown with
+			// SHUT_WR, forward it to the other end of the pipe:
+			if err, ok := err.(*net.OpError); ok && err.Err == syscall.EPIPE {
+				from.CloseWrite()
+			}
+		}
+		to.CloseRead()
+		event <- written
+	}
+
+	go broker(client, backend)
+	go broker(backend, client)
+
+	var transferred int64
+	for i := 0; i < 2; i++ {
+		select {
+		case written := <-event:
+			transferred += written
+		case <-quit:
+			// Interrupt the two brokers and "join" them.
+			client.Close()
+			backend.Close()
+			for ; i < 2; i++ {
+				transferred += <-event
+			}
+			return
+		}
+	}
+	client.Close()
+	backend.Close()
+}
+
+// Run starts forwarding the traffic using TCP.
+func (proxy *TCPProxy) Run() {
+	quit := make(chan bool)
+	defer close(quit)
+	for {
+		client, err := proxy.listener.Accept()
+		if err != nil {
+			proxy.Logger.Printf("Stopping proxy on tcp/%v for tcp/%v (%s)", proxy.frontendAddr, proxy.backendAddr, err)
+			return
+		}
+		go proxy.clientLoop(client.(*net.TCPConn), quit)
+	}
+}
+
+// Close stops forwarding the traffic.
+func (proxy *TCPProxy) Close() { proxy.listener.Close() }
+
+// FrontendAddr returns the TCP address on which the proxy is listening.
+func (proxy *TCPProxy) FrontendAddr() net.Addr { return proxy.frontendAddr }
+
+// BackendAddr returns the TCP proxied address.
+func (proxy *TCPProxy) BackendAddr() net.Addr { return proxy.backendAddr }