Check for token in every received message.

tunnel/default.go

@@ -96,6 +96,7 @@ func (t *tun) newSocket(id, session string) (*socket, bool) {
 		t.Unlock()
 		return nil, false
 	}
+
 	t.sockets[id+session] = s
 	t.Unlock()
 
@@ -158,21 +159,26 @@ func (t *tun) listen(link transport.Socket) {
 			return
 		}
 
+		// loopback flag
 		var loopback bool
 
-		switch msg.Header["Micro-Tunnel"] {
+		// TODO: figure out the way how to populate Micro-Tunnel-Token for every message
-		case "connect":
+
-			// TODO: handle the connect message
 		// check the Micro-Tunnel-Token
 		token, ok := msg.Header["Micro-Tunnel-Token"]
-			if !ok {
+		if ok {
-				// no token found; bailing
-				continue
-			}
 			// are we connecting to ourselves?
 			if token == t.token {
 				loopback = true
 			}
+		}
+
+		switch msg.Header["Micro-Tunnel"] {
+		case "connect":
+			// connecting without token is not allowed
+			if token == "" {
+				continue
+			}
 		case "close":
 			// TODO: handle the close message
 			// maybe report io.EOF or kill the link
@@ -187,6 +193,8 @@ func (t *tun) listen(link transport.Socket) {
 		session := msg.Header["Micro-Tunnel-Session"]
 		delete(msg.Header, "Micro-Tunnel-Session")
 
+		// TODO: should we delete Micro-Tunnel-Token header, too?
+
 		// if the session id is blank there's nothing we can do
 		// TODO: check this is the case, is there any reason
 		// why we'd have a blank session? Is the tunnel
@@ -208,6 +216,12 @@ func (t *tun) listen(link transport.Socket) {
 			// this could be something we dialed in which case
 			// we have a session for it otherwise its a listener
 			s, exists = t.getSocket(id, session)
+			if !exists {
+				// try get it based on just the tunnel id
+				// the assumption here is that a listener
+				// has no session but its set a listener session
+				s, exists = t.getSocket(id, "listener")
+			}
 		}
 		// bail if no socket has been found
 		if !exists {

tunnel/tunnel_test.go

@@ -98,3 +98,35 @@ func TestTunnel(t *testing.T) {
 	// wait until done
 	wg.Wait()
 }
+
+//func TestLoopbackTunnel(t *testing.T) {
+//	// create a new tunnel client
+//	tun := NewTunnel(
+//		Address("127.0.0.1:9096"),
+//		Nodes("127.0.0.1:9096"),
+//	)
+//
+//	// start tunB
+//	err := tun.Connect()
+//	if err != nil {
+//		t.Fatal(err)
+//	}
+//	defer tun.Close()
+//
+//	time.Sleep(time.Millisecond * 50)
+//
+//	var wg sync.WaitGroup
+//
+//	// start accepting connections
+//	// on tunnel A
+//	wg.Add(1)
+//	go testAccept(t, tun, &wg)
+//
+//	time.Sleep(time.Millisecond * 50)
+//
+//	// dial and send via B
+//	testSend(t, tun)
+//
+//	// wait until done
+//	wg.Wait()
+//}