fix an issue where aggressive NATs or proxy connections pooled for too long could hang the tunneling of public connections

src/ngrok/server/control.go

@@ -16,6 +16,8 @@ const (
 	pingTimeoutInterval = 30 * time.Second
 	connReapInterval    = 10 * time.Second
 	controlWriteTimeout = 10 * time.Second
+	proxyStaleDuration  = 60 * time.Second
+	proxyMaxPoolSize    = 10
 )
 
 type Control struct {
@@ -104,7 +106,7 @@ func NewControl(ctlConn conn.Conn, authMsg *msg.Auth) {
 		replaced.shutdown.WaitComplete()
 	}
 
-	// start the writer first so that the follow messages get sent
+	// start the writer first so that the following messages get sent
 	go c.writer()
 
 	// Respond to authentication
@@ -292,6 +294,7 @@ func (c *Control) stopper() {
 func (c *Control) RegisterProxy(conn conn.Conn) {
 	conn.AddLogPrefix(c.id)
 
+	conn.SetDeadline(time.Now().Add(proxyStaleDuration))
 	select {
 	case c.proxies <- conn:
 		conn.Info("Registered")
@@ -307,35 +310,34 @@ func (c *Control) RegisterProxy(conn conn.Conn) {
 // Returns an error if we couldn't get a proxy because it took too long
 // or the tunnel is closing
 func (c *Control) GetProxy() (proxyConn conn.Conn, err error) {
-	// initial timeout is zero to try to get a proxy connection without asking for one
-	timeout := time.NewTimer(0)
+	var ok bool
+
+	// get a proxy connection from the pool
+	select {
+	case proxyConn, ok = <-c.proxies:
+		if !ok {
+			err = fmt.Errorf("No proxy connections available, control is closing")
+			return
+		}
+	default:
+		// no proxy available in the pool, ask for one over the control channel
+		c.conn.Debug("No proxy in pool, requesting proxy from control . . .")
+		if err = util.PanicToError(func() { c.out <- &msg.ReqProxy{} }); err != nil {
+			return
+		}
 
-	// get a proxy connection. if we timeout, request one over the control channel
-	for proxyConn == nil {
-		var ok bool
 		select {
 		case proxyConn, ok = <-c.proxies:
 			if !ok {
 				err = fmt.Errorf("No proxy connections available, control is closing")
 				return
 			}
-			continue
-		case <-timeout.C:
-			c.conn.Debug("Requesting new proxy connection")
-			// request a proxy connection
-			if err = util.PanicToError(func() { c.out <- &msg.ReqProxy{} }); err != nil {
-				return
-			}
 
-			// timeout after 1 second if we don't get one
-			timeout.Reset(1 * time.Second)
+		case <-time.After(pingTimeoutInterval):
+			err = fmt.Errorf("Timeout trying to get proxy connection")
+			return
 		}
 	}
-
-	// To try to reduce latency hanndling tunnel connections, we employ
-	// the following curde heuristic:
-	// Whenever we take a proxy connection from the pool, replace it with a new one
-	err = util.PanicToError(func() { c.out <- &msg.ReqProxy{} })
 	return
 }
 

src/ngrok/server/metrics.go

@@ -156,7 +156,7 @@ func (m *LocalMetrics) Report() {
 
 type KeenIoMetric struct {
 	Collection string
-	Event interface{}
+	Event      interface{}
 }
 
 type KeenIoMetrics struct {

src/ngrok/server/tunnel.go

@@ -8,6 +8,7 @@ import (
 	"ngrok/conn"
 	"ngrok/log"
 	"ngrok/msg"
+	"ngrok/util"
 	"os"
 	"strconv"
 	"strings"
@@ -253,9 +254,8 @@ func (t *Tunnel) HandlePublicConnection(publicConn conn.Conn) {
 	metrics.OpenConnection(t, publicConn)
 
 	var proxyConn conn.Conn
-	var attempts int
 	var err error
-	for {
+	for i := 0; i < (2 * proxyMaxPoolSize); i++ {
 		// get a proxy connection
 		if proxyConn, err = t.ctl.GetProxy(); err != nil {
 			t.Warn("Failed to get proxy connection: %v", err)
@@ -270,20 +270,29 @@ func (t *Tunnel) HandlePublicConnection(publicConn conn.Conn) {
 			Url:        t.url,
 			ClientAddr: publicConn.RemoteAddr().String(),
 		}
+
 		if err = msg.WriteMsg(proxyConn, startPxyMsg); err != nil {
-			attempts += 1
-			proxyConn.Warn("Failed to write StartProxyMessage: %v, attempt %d", err, attempts)
-			if attempts > 3 {
-				// give up
-				publicConn.Error("Too many failures starting proxy connection")
-				return
-			}
+			proxyConn.Warn("Failed to write StartProxyMessage: %v, attempt %d", err, i)
+			proxyConn.Close()
 		} else {
 			// success
 			break
 		}
 	}
 
+	if err != nil {
+		// give up
+		publicConn.Error("Too many failures starting proxy connection")
+		return
+	}
+
+	// To reduce latency handling tunnel connections, we employ the following curde heuristic:
+	// Whenever we take a proxy connection from the pool, replace it with a new one
+	util.PanicToError(func() { t.ctl.out <- &msg.ReqProxy{} })
+
+	// no timeouts while connections are joined
+	proxyConn.SetDeadline(time.Time{})
+
 	// join the public and proxy connections
 	bytesIn, bytesOut := conn.Join(publicConn, proxyConn)
 	metrics.CloseConnection(t, publicConn, startTime, bytesIn, bytesOut)