From f9eddf1e6f23341718054d58cdcdb6b251779d5b Mon Sep 17 00:00:00 2001
From: Milos Gajdos <milosgajdos83@gmail.com>
Date: Wed, 14 Aug 2019 17:14:39 +0100
Subject: [PATCH] Fixing the tunnel loopback messaging

---
 tunnel/default.go | 27 ++++++++++++++++++---------
 tunnel/socket.go  | 11 ++++++++++-
 2 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/tunnel/default.go b/tunnel/default.go
index 90a641f0..03d7c503 100644
--- a/tunnel/default.go
+++ b/tunnel/default.go
@@ -41,7 +41,8 @@ type tun struct {
 
 type link struct {
 	transport.Socket
-	id string
+	id       string
+	loopback bool
 }
 
 // create new tunnel on top of a link
@@ -148,8 +149,11 @@ func (t *tun) process() {
 				log.Debugf("Zero links to send to")
 			}
 			for _, link := range t.links {
-				log.Debugf("Sending %+v to %s", newMsg, link.Remote())
 				// TODO: error check and reconnect
+				log.Debugf("Sending %+v to %s", newMsg, link.Remote())
+				if link.loopback && msg.outbound {
+					continue
+				}
 				link.Send(newMsg)
 			}
 			t.RUnlock()
@@ -160,7 +164,7 @@ func (t *tun) process() {
 }
 
 // process incoming messages
-func (t *tun) listen(link transport.Socket) {
+func (t *tun) listen(link *link) {
 	// loopback flag
 	var loopback bool
 
@@ -184,6 +188,7 @@ func (t *tun) listen(link transport.Socket) {
 			// are we connecting to ourselves?
 			if token == t.token {
 				loopback = true
+				link.loopback = true
 			}
 			continue
 		case "close":
@@ -298,10 +303,11 @@ func (t *tun) connect() error {
 			// save the link
 			id := uuid.New().String()
 			t.Lock()
-			t.links[id] = &link{
+			link := &link{
 				Socket: sock,
 				id:     id,
 			}
+			t.links[id] = link
 			t.Unlock()
 
 			// delete the link
@@ -313,7 +319,7 @@ func (t *tun) connect() error {
 			}()
 
 			// listen for inbound messages
-			t.listen(sock)
+			t.listen(link)
 		})
 
 		t.Lock()
@@ -349,15 +355,16 @@ func (t *tun) connect() error {
 			continue
 		}
 
-		// process incoming messages
-		go t.listen(c)
-
 		// save the link
 		id := uuid.New().String()
-		t.links[id] = &link{
+		link := &link{
 			Socket: c,
 			id:     id,
 		}
+		t.links[id] = link
+
+		// process incoming messages
+		go t.listen(link)
 	}
 
 	// process outbound messages to be sent
@@ -448,6 +455,8 @@ func (t *tun) Dial(addr string) (Conn, error) {
 	c.remote = addr
 	// set local
 	c.local = "local"
+	// outbound socket
+	c.outbound = true
 
 	return c, nil
 }
diff --git a/tunnel/socket.go b/tunnel/socket.go
index 8b288c40..2590a48e 100644
--- a/tunnel/socket.go
+++ b/tunnel/socket.go
@@ -25,6 +25,8 @@ type socket struct {
 	recv chan *message
 	// wait until we have a connection
 	wait chan bool
+	// outbound marks the socket as outbound
+	outbound bool
 }
 
 // message is sent over the send channel
@@ -33,6 +35,8 @@ type message struct {
 	id string
 	// the session id
 	session string
+	// outbound marks the message as outbound
+	outbound bool
 	// transport data
 	data *transport.Message
 }
@@ -72,7 +76,12 @@ func (s *socket) Send(m *transport.Message) error {
 	}
 
 	// append to backlog
-	msg := &message{id: s.id, session: s.session, data: data}
+	msg := &message{
+		id:       s.id,
+		session:  s.session,
+		outbound: s.outbound,
+		data:     data,
+	}
 	log.Debugf("Appending %+v to send backlog", msg)
 	s.send <- msg
 	return nil