Browse Source

add support for specifying a remote tcp port

Alan Shreve 11 years ago
parent
commit
98902cff73
4 changed files with 57 additions and 42 deletions
  1. 5 4
      src/ngrok/client/config.go
  2. 6 5
      src/ngrok/client/model.go
  3. 7 2
      src/ngrok/msg/msg.go
  4. 39 31
      src/ngrok/server/tunnel.go

+ 5 - 4
src/ngrok/client/config.go

@@ -27,10 +27,11 @@ type Configuration struct {
 }
 
 type TunnelConfiguration struct {
-	Subdomain string            `yaml:"subdomain,omitempty"`
-	Hostname  string            `yaml:"hostname,omitempty"`
-	Protocols map[string]string `yaml:"proto,omitempty"`
-	HttpAuth  string            `yaml:"auth,omitempty"`
+	Subdomain  string            `yaml:"subdomain,omitempty"`
+	Hostname   string            `yaml:"hostname,omitempty"`
+	Protocols  map[string]string `yaml:"proto,omitempty"`
+	HttpAuth   string            `yaml:"auth,omitempty"`
+	RemotePort uint16            `yaml:"remote_port,omitempty"`
 }
 
 func LoadConfiguration(opts *Options) (config *Configuration, err error) {

+ 6 - 5
src/ngrok/client/model.go

@@ -258,11 +258,12 @@ func (c *ClientModel) control() {
 		}
 
 		reqTunnel := &msg.ReqTunnel{
-			ReqId:     util.RandId(8),
-			Protocol:  strings.Join(protocols, "+"),
-			Hostname:  config.Hostname,
-			Subdomain: config.Subdomain,
-			HttpAuth:  config.HttpAuth,
+			ReqId:      util.RandId(8),
+			Protocol:   strings.Join(protocols, "+"),
+			Hostname:   config.Hostname,
+			Subdomain:  config.Subdomain,
+			HttpAuth:   config.HttpAuth,
+			RemotePort: config.RemotePort,
 		}
 
 		// send the tunnel request

+ 7 - 2
src/ngrok/msg/msg.go

@@ -63,11 +63,16 @@ type AuthResp struct {
 // ReqId is a random number set by the client that it can pull
 // from future NewTunnel's to correlate then to the requesting ReqTunnel.
 type ReqTunnel struct {
-	ReqId     string
-	Protocol  string
+	ReqId    string
+	Protocol string
+
+	// http only
 	Hostname  string
 	Subdomain string
 	HttpAuth  string
+
+	// tcp only
+	RemotePort uint16
 }
 
 // When the server opens a new tunnel on behalf of

+ 39 - 31
src/ngrok/server/tunnel.go

@@ -104,50 +104,58 @@ func NewTunnel(m *msg.ReqTunnel, ctl *Control) (t *Tunnel, err error) {
 	proto := t.req.Protocol
 	switch proto {
 	case "tcp":
-		var port int = 0
+		bindTcp := func(port int) error {
+			if t.listener, err = net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP("0.0.0.0"), Port: port}); err != nil {
+				err = t.ctl.conn.Error("Error binding TCP listener: %v", err)
+				return err
+			}
+
+			// create the url
+			addr := t.listener.Addr().(*net.TCPAddr)
+			t.url = fmt.Sprintf("tcp://%s:%d", opts.domain, addr.Port)
+
+			// register it
+			if err = tunnelRegistry.RegisterAndCache(t.url, t); err != nil {
+				// This should never be possible because the OS will
+				// only assign available ports to us.
+				t.listener.Close()
+				err = fmt.Errorf("TCP listener bound, but failed to register %s", t.url)
+				return err
+			}
+
+			go t.listenTcp(t.listener)
+			return nil
+		}
+
+		// use the custom remote port you asked for
+		if t.req.RemotePort != 0 {
+			bindTcp(int(t.req.RemotePort))
+			return
+		}
 
 		// try to return to you the same port you had before
 		cachedUrl := tunnelRegistry.GetCachedRegistration(t)
 		if cachedUrl != "" {
+			var port int
 			parts := strings.Split(cachedUrl, ":")
 			portPart := parts[len(parts)-1]
 			port, err = strconv.Atoi(portPart)
 			if err != nil {
 				t.ctl.conn.Error("Failed to parse cached url port as integer: %s", portPart)
-				// continue with zero
-				port = 0
+			} else {
+				// we have a valid, cached port, let's try to bind with it
+				if bindTcp(port) != nil {
+					t.ctl.conn.Warn("Failed to get custom port %d: %v, trying a random one", port, err)
+				} else {
+					// success, we're done
+					return
+				}
 			}
 		}
 
 		// Bind for TCP connections
-		t.listener, err = net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP("0.0.0.0"), Port: port})
-
-		// If we failed with a custom port, try with a random one
-		if err != nil && port != 0 {
-			t.ctl.conn.Warn("Failed to get custom port %d: %v, trying a random one", port, err)
-			t.listener, err = net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP("0.0.0.0"), Port: 0})
-		}
-
-		// we tried to bind with a random port and failed (no more ports available?)
-		if err != nil {
-			err = t.ctl.conn.Error("Error binding TCP listener: %v", err)
-			return
-		}
-
-		// create the url
-		addr := t.listener.Addr().(*net.TCPAddr)
-		t.url = fmt.Sprintf("tcp://%s:%d", opts.domain, addr.Port)
-
-		// register it
-		if err = tunnelRegistry.RegisterAndCache(t.url, t); err != nil {
-			// This should never be possible because the OS will
-			// only assign available ports to us.
-			t.listener.Close()
-			err = fmt.Errorf("TCP listener bound, but failed to register %s", t.url)
-			return
-		}
-
-		go t.listenTcp(t.listener)
+		bindTcp(0)
+		return
 
 	case "http", "https":
 		l, ok := listeners[proto]