Move proxy into pkg

Docker-DCO-1.1-Signed-off-by: Michael Crosby <michael@crosbymichael.com> (github: crosbymichael)

proxy/tcp_proxy.go

@@ -0,0 +1,93 @@
+package proxy
+
+import (
+	"io"
+	"log"
+	"net"
+	"syscall"
+)
+
+type TCPProxy struct {
+	listener     *net.TCPListener
+	frontendAddr *net.TCPAddr
+	backendAddr  *net.TCPAddr
+}
+
+func NewTCPProxy(frontendAddr, backendAddr *net.TCPAddr) (*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:
+	return &TCPProxy{
+		listener:     listener,
+		frontendAddr: listener.Addr().(*net.TCPAddr),
+		backendAddr:  backendAddr,
+	}, nil
+}
+
+func (proxy *TCPProxy) clientLoop(client *net.TCPConn, quit chan bool) {
+	backend, err := net.DialTCP("tcp", nil, proxy.backendAddr)
+	if err != nil {
+		log.Printf("Can't forward traffic to backend tcp/%v: %v\n", proxy.backendAddr, err.Error())
+		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
+	}
+
+	log.Printf("Forwarding traffic between tcp/%v and tcp/%v", client.RemoteAddr(), backend.RemoteAddr())
+	go broker(client, backend)
+	go broker(backend, client)
+
+	var transferred int64 = 0
+	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
+			}
+			goto done
+		}
+	}
+	client.Close()
+	backend.Close()
+done:
+	log.Printf("%v bytes transferred between tcp/%v and tcp/%v", transferred, client.RemoteAddr(), backend.RemoteAddr())
+}
+
+func (proxy *TCPProxy) Run() {
+	quit := make(chan bool)
+	defer close(quit)
+	log.Printf("Starting proxy on tcp/%v for tcp/%v", proxy.frontendAddr, proxy.backendAddr)
+	for {
+		client, err := proxy.listener.Accept()
+		if err != nil {
+			log.Printf("Stopping proxy on tcp/%v for tcp/%v (%v)", proxy.frontendAddr, proxy.backendAddr, err.Error())
+			return
+		}
+		go proxy.clientLoop(client.(*net.TCPConn), quit)
+	}
+}
+
+func (proxy *TCPProxy) Close()                 { proxy.listener.Close() }
+func (proxy *TCPProxy) FrontendAddr() net.Addr { return proxy.frontendAddr }
+func (proxy *TCPProxy) BackendAddr() net.Addr  { return proxy.backendAddr }