From 2549d64b8b9b6e57fbb436bee8a57ee05083bc6d Mon Sep 17 00:00:00 2001
From: "Jonathan Leibiusky @xetorthio" <ionathan@gmail.com>
Date: Mon, 21 Aug 2017 15:12:53 -0300
Subject: [PATCH] Avoid TIME_WAIT leaking connections

---
 docker/local_cached_factory.go |  3 ++-
 router/router.go               | 32 ++++++++++++++++++++++++--------
 2 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/docker/local_cached_factory.go b/docker/local_cached_factory.go
index 8812955..e9d2fc7 100644
--- a/docker/local_cached_factory.go
+++ b/docker/local_cached_factory.go
@@ -91,7 +91,8 @@ func (f *localCachedFactory) GetForInstance(instance *types.Instance) (DockerApi
 			Timeout:   1 * time.Second,
 			KeepAlive: 30 * time.Second,
 		}).DialContext,
-		Proxy: http.ProxyURL(proxyUrl),
+		MaxIdleConnsPerHost: 5,
+		Proxy:               http.ProxyURL(proxyUrl),
 	}
 	if tlsConfig != nil {
 		transport.TLSClientConfig = tlsConfig
diff --git a/router/router.go b/router/router.go
index b601470..60e6b1c 100644
--- a/router/router.go
+++ b/router/router.go
@@ -10,6 +10,7 @@ import (
 	"net/http"
 	"strings"
 	"sync"
+	"time"
 
 	"golang.org/x/crypto/ssh"
 
@@ -34,11 +35,12 @@ type proxyRouter struct {
 	keyPath      string
 	director     Director
 	closed       bool
-	httpListener net.Listener
+	httpListener *net.TCPListener
 	udpDnsServer *dns.Server
 	tcpDnsServer *dns.Server
 	sshListener  net.Listener
 	sshConfig    *ssh.ServerConfig
+	dialer       *net.Dialer
 }
 
 func (r *proxyRouter) Listen(httpAddr, dnsAddr, sshAddr string) {
@@ -52,8 +54,11 @@ func (r *proxyRouter) ListenAndWait(httpAddr, dnsAddr, sshAddr string) {
 }
 
 func (r *proxyRouter) listen(wg *sync.WaitGroup, httpAddr, dnsAddr, sshAddr string) {
-
-	l, err := net.Listen("tcp", httpAddr)
+	tcpAddr, err := net.ResolveTCPAddr("tcp", httpAddr)
+	if err != nil {
+		log.Fatal(err)
+	}
+	l, err := net.ListenTCP("tcp", tcpAddr)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -61,10 +66,12 @@ func (r *proxyRouter) listen(wg *sync.WaitGroup, httpAddr, dnsAddr, sshAddr stri
 	wg.Add(1)
 	go func() {
 		for !r.closed {
-			conn, err := r.httpListener.Accept()
+			conn, err := r.httpListener.AcceptTCP()
 			if err != nil {
 				continue
 			}
+			conn.SetKeepAlive(true)
+			conn.SetKeepAlivePeriod(3 * time.Minute)
 			go r.handleConnection(conn)
 		}
 		wg.Done()
@@ -333,7 +340,7 @@ func (r *proxyRouter) handleConnection(c net.Conn) {
 			log.Printf("Error directing request: %v\n", err)
 			return
 		}
-		d, err := net.Dial("tcp", dstHost.String())
+		d, err := r.dialer.Dial("tcp", dstHost.String())
 		if err != nil {
 			log.Printf("Error dialing backend %s: %v\n", dstHost.String(), err)
 			return
@@ -358,7 +365,7 @@ func (r *proxyRouter) handleConnection(c net.Conn) {
 			log.Printf("Error directing request: %v\n", err)
 			return
 		}
-		d, err := net.Dial("tcp", dstHost.String())
+		d, err := r.dialer.Dial("tcp", dstHost.String())
 		if err != nil {
 			log.Printf("Error dialing backend %s: %v\n", dstHost.String(), err)
 			return
@@ -423,10 +430,11 @@ func proxySsh(reqs1, reqs2 <-chan *ssh.Request, channel1, channel2 ssh.Channel)
 
 func proxyConn(src, dst net.Conn) {
 	errc := make(chan error, 2)
-	cp := func(dst io.Writer, src io.Reader) {
+	cp := func(dst net.Conn, src net.Conn) {
 		_, err := io.Copy(dst, src)
 		errc <- err
 	}
+
 	go cp(src, dst)
 	go cp(dst, src)
 	<-errc
@@ -450,5 +458,13 @@ func NewRouter(director Director, keyPath string) *proxyRouter {
 
 	sshConfig.AddHostKey(private)
 
-	return &proxyRouter{director: director, sshConfig: sshConfig}
+	return &proxyRouter{
+		director:  director,
+		sshConfig: sshConfig,
+		dialer: &net.Dialer{
+			Timeout:   30 * time.Second,
+			KeepAlive: 30 * time.Second,
+			DualStack: true,
+		},
+	}
 }