simplify protocol. switch on actual types of messages instead of a string. properly display error message when quitting, being moving to a model where the client pings the server. wire up version number to the terminal UI. fail to connect clients with a different protocol version

src/ngrok/client/main.go

@@ -108,6 +108,7 @@ func control(s *State, ctl *ui.Controller) {
 		Hostname:  s.opts.hostname,
 		Subdomain: s.opts.subdomain,
 		ClientId:  s.id,
+		Version:   msg.Version,
 	})
 
 	if err != nil {
@@ -131,6 +132,7 @@ func control(s *State, ctl *ui.Controller) {
 	//state.version = regAck.Version
 	s.publicUrl = regAck.Url
 	s.status = "online"
+	s.serverVersion = regAck.Version
 	ctl.Update(s)
 
 	// main control loop
@@ -140,12 +142,13 @@ func control(s *State, ctl *ui.Controller) {
 			panic(err)
 		}
 
-		switch m.GetType() {
-		case "ReqProxyMsg":
+		switch m.(type) {
+		case *msg.ReqProxyMsg:
 			go proxy(regAck.ProxyAddr, s, ctl)
 
-		case "PingMsg":
-			msg.WriteMsg(conn, &msg.PongMsg{})
+		case *msg.PongMsg:
+			//msg.WriteMsg(conn, &msg.PongMsg{})
+			// XXX: update our live status
 		}
 	}
 }
@@ -194,6 +197,7 @@ func Main() {
 			case cmd := <-ctl.Cmds:
 				switch cmd.Code {
 				case ui.QUIT:
+					quitMessage = cmd.Payload.(string)
 					ctl.DoShutdown()
 					return
 				case ui.REPLAY:

src/ngrok/client/state.go

@@ -2,23 +2,26 @@ package client
 
 import (
 	metrics "github.com/inconshreveable/go-metrics"
+	"ngrok/msg"
 	"ngrok/proto"
 )
 
 // client state
 type State struct {
-	id        string
-	publicUrl string
-	protocol  proto.Protocol
-	opts      *Options
-	metrics   *ClientMetrics
+	id            string
+	publicUrl     string
+	serverVersion string
+	protocol      proto.Protocol
+	opts          *Options
+	metrics       *ClientMetrics
 
 	// just for UI purposes
 	status string
 }
 
 // implement client.ui.State
-func (s State) GetVersion() string          { return "" }
+func (s State) GetClientVersion() string    { return msg.Version }
+func (s State) GetServerVersion() string    { return msg.Version }
 func (s State) GetPublicUrl() string        { return s.publicUrl }
 func (s State) GetLocalAddr() string        { return s.opts.localaddr }
 func (s State) GetWebPort() int             { return s.opts.webport }

src/ngrok/client/ui/state.go

@@ -6,7 +6,8 @@ import (
 )
 
 type State interface {
-	GetVersion() string
+	GetClientVersion() string
+	GetServerVersion() string
 	GetPublicUrl() string
 	GetLocalAddr() string
 	GetStatus() string

src/ngrok/client/views/term/view.go

@@ -80,7 +80,7 @@ func (v *TermView) Render() {
 	status := v.state.GetStatus()
 	v.APrintf(colorForConn(status), 0, 2, "%-30s%s", "Tunnel Status", status)
 
-	v.Printf(0, 3, "%-30s%s", "Version", v.state.GetVersion())
+	v.Printf(0, 3, "%-30s%s/%s", "Version", v.state.GetClientVersion(), v.state.GetServerVersion())
 	v.Printf(0, 4, "%-30s%s", "Protocol", v.state.GetProtocol().GetName())
 	v.Printf(0, 5, "%-30s%s -> %s", "Forwarding", v.state.GetPublicUrl(), v.state.GetLocalAddr())
 	webAddr := fmt.Sprintf("http://localhost:%d", v.state.GetWebPort())

src/ngrok/msg/msg.go

@@ -27,23 +27,15 @@ func init() {
 	TypeMap["MetricsRespMsg"] = t((*MetricsRespMsg)(nil))
 }
 
-type Message interface {
-	GetType() string
-	SetType(string)
-}
+type Message interface{}
 
 type Envelope struct {
-	Version string
 	Type    string
 	Payload json.RawMessage
 }
 
-type TypeEmbed struct {
-	Type string
-}
-
 type RegMsg struct {
-	TypeEmbed
+	Version   string
 	Protocol  string
 	Hostname  string
 	Subdomain string
@@ -56,52 +48,35 @@ type RegMsg struct {
 }
 
 type RegAckMsg struct {
-	TypeEmbed
-	Type      string
+	Version   string
 	Url       string
 	ProxyAddr string
 	Error     string
 }
 
 type RegProxyMsg struct {
-	TypeEmbed
 	Url string
 }
 
 type ReqProxyMsg struct {
-	TypeEmbed
 }
 
 type PingMsg struct {
-	TypeEmbed
 }
 
 type PongMsg struct {
-	TypeEmbed
 }
 
 type VersionMsg struct {
-	TypeEmbed
 }
 
 type VersionRespMsg struct {
-	TypeEmbed
 	Version string
 }
 
 type MetricsMsg struct {
-	TypeEmbed
 }
 
 type MetricsRespMsg struct {
-	TypeEmbed
 	Metrics string
 }
-
-func (m *TypeEmbed) GetType() string {
-	return m.Type
-}
-
-func (m *TypeEmbed) SetType(typ string) {
-	m.Type = typ
-}

src/ngrok/msg/pack.go

@@ -23,7 +23,6 @@ func unpack(buffer []byte, msgIn Message) (msg Message, err error) {
 
 		// guess type
 		msg = reflect.New(t).Interface().(Message)
-		msg.SetType(env.Type)
 	} else {
 		msg = msgIn
 	}
@@ -43,11 +42,9 @@ func Unpack(buffer []byte) (msg Message, err error) {
 
 func Pack(payload interface{}) ([]byte, error) {
 	return json.Marshal(struct {
-		Version string
 		Type    string
 		Payload interface{}
 	}{
-		Version: Version,
 		Type:    reflect.TypeOf(payload).Elem().Name(),
 		Payload: payload,
 	})

src/ngrok/server/control.go

@@ -11,8 +11,8 @@ import (
 )
 
 const (
-	pingInterval     = 30 * time.Second
-	connReapInterval = pingInterval * 5
+	pingTimeoutInterval = 30 * time.Second
+	connReapInterval    = 10 * time.Second
 )
 
 type Control struct {
@@ -25,7 +25,7 @@ type Control struct {
 	stop chan (msg.Message)
 
 	// heartbeat
-	lastPong int64
+	lastPing int64
 
 	// tunnel
 	tun *Tunnel
@@ -37,7 +37,7 @@ func NewControl(tcpConn *net.TCPConn) {
 		out:      make(chan (interface{}), 1),
 		in:       make(chan (msg.Message), 1),
 		stop:     make(chan (msg.Message), 1),
-		lastPong: time.Now().Unix(),
+		lastPing: time.Now().Unix(),
 	}
 
 	go c.managerThread()
@@ -45,7 +45,6 @@ func NewControl(tcpConn *net.TCPConn) {
 }
 
 func (c *Control) managerThread() {
-	ping := time.NewTicker(pingInterval)
 	reap := time.NewTicker(connReapInterval)
 
 	// all shutdown functionality in here
@@ -53,7 +52,6 @@ func (c *Control) managerThread() {
 		if err := recover(); err != nil {
 			c.conn.Info("Control::managerThread failed with error %v: %s", err, debug.Stack())
 		}
-		ping.Stop()
 		reap.Stop()
 		c.conn.Close()
 
@@ -68,11 +66,9 @@ func (c *Control) managerThread() {
 		case m := <-c.out:
 			msg.WriteMsg(c.conn, m)
 
-		case <-ping.C:
-			msg.WriteMsg(c.conn, &msg.PingMsg{})
-
 		case <-reap.C:
-			if (time.Now().Unix() - c.lastPong) > 60 {
+			lastPing := time.Unix(c.lastPing, 0)
+			if time.Since(lastPing) > pingTimeoutInterval {
 				c.conn.Info("Lost heartbeat")
 				metrics.lostHeartbeatMeter.Mark(1)
 				return
@@ -84,16 +80,17 @@ func (c *Control) managerThread() {
 			}
 			return
 
-		case m := <-c.in:
-			switch m.GetType() {
-			case "RegMsg":
+		case mRaw := <-c.in:
+			switch m := mRaw.(type) {
+			case *msg.RegMsg:
 				c.conn.Info("Registering new tunnel")
-				c.tun = newTunnel(m.(*msg.RegMsg), c)
+				c.tun = newTunnel(m, c)
 
-			case "PongMsg":
-				atomic.StoreInt64(&c.lastPong, time.Now().Unix())
+			case *msg.PingMsg:
+				atomic.StoreInt64(&c.lastPing, time.Now().Unix())
+				c.out <- &msg.PongMsg{}
 
-			case "VersionReqMsg":
+			case *msg.VersionMsg:
 				c.out <- &msg.VersionRespMsg{Version: version}
 			}
 		}

src/ngrok/server/tunnel.go

@@ -60,6 +60,10 @@ func newTunnel(m *msg.RegMsg, ctl *Control) (t *Tunnel) {
 		return
 	}
 
+	if m.Version != msg.Version {
+		t.ctl.stop <- &msg.RegAckMsg{Error: fmt.Sprintf("Incompatible versions. Server %s, client %s.", msg.Version, m.Version)}
+	}
+
 	// pre-encode the http basic auth for fast comparisons later
 	if m.HttpAuth != "" {
 		m.HttpAuth = "Basic " + base64.StdEncoding.EncodeToString([]byte(m.HttpAuth))
@@ -68,7 +72,12 @@ func newTunnel(m *msg.RegMsg, ctl *Control) (t *Tunnel) {
 	t.ctl.conn.AddLogPrefix(t.Id())
 	t.AddLogPrefix(t.Id())
 	t.Info("Registered new tunnel")
-	t.ctl.out <- &msg.RegAckMsg{Url: t.url, ProxyAddr: fmt.Sprintf("%s", proxyAddr)}
+	t.ctl.out <- &msg.RegAckMsg{
+		Url:       t.url,
+		ProxyAddr: fmt.Sprintf("%s", proxyAddr),
+		Version:   msg.Version,
+	}
+
 	return
 }