From 917ca6c1c58696a13c58b4dc44a66644d1165d53 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Fri, 8 Mar 2019 10:26:46 +0000
Subject: [PATCH] Make changes based on review comments

---
 src/util/util.go       |  2 +-
 src/yggdrasil/admin.go |  2 +-
 src/yggdrasil/awdl.go  |  3 +--
 src/yggdrasil/link.go  |  2 +-
 src/yggdrasil/tcp.go   | 38 ++++++++++++++++++++++++++------------
 5 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/src/util/util.go b/src/util/util.go
index d669fa5..49e0207 100644
--- a/src/util/util.go
+++ b/src/util/util.go
@@ -87,7 +87,7 @@ func Difference(a, b []string) []string {
 		mb[x] = true
 	}
 	for _, x := range a {
-		if _, ok := mb[x]; !ok {
+		if !mb[x] {
 			ab = append(ab, x)
 		}
 	}
diff --git a/src/yggdrasil/admin.go b/src/yggdrasil/admin.go
index 8cb195e..a0854f2 100644
--- a/src/yggdrasil/admin.go
+++ b/src/yggdrasil/admin.go
@@ -676,7 +676,7 @@ func (a *admin) getData_getPeers() []admin_nodeInfo {
 			{"bytes_sent", atomic.LoadUint64(&p.bytesSent)},
 			{"bytes_recvd", atomic.LoadUint64(&p.bytesRecvd)},
 			{"proto", p.intf.info.linkType},
-			{"endpoint", p.intf.info.remote},
+			{"endpoint", p.intf.name},
 			{"box_pub_key", hex.EncodeToString(p.box[:])},
 		}
 		peerInfos = append(peerInfos, info)
diff --git a/src/yggdrasil/awdl.go b/src/yggdrasil/awdl.go
index fe64e8b..5e8cce1 100644
--- a/src/yggdrasil/awdl.go
+++ b/src/yggdrasil/awdl.go
@@ -54,8 +54,7 @@ func (a *awdl) init(l *link) error {
 	a.mutex.Unlock()
 
 	go func() {
-		for {
-			e := <-a.reconfigure
+		for e := range a.reconfigure {
 			e <- nil
 		}
 	}()
diff --git a/src/yggdrasil/link.go b/src/yggdrasil/link.go
index a81b50d..67ce5c1 100644
--- a/src/yggdrasil/link.go
+++ b/src/yggdrasil/link.go
@@ -105,7 +105,7 @@ func (l *link) call(uri string, sintf string) error {
 	case "tcp":
 		l.tcp.call(u.Host, nil, sintf)
 	case "socks":
-		l.tcp.call(pathtokens[0], &u.Host, sintf)
+		l.tcp.call(pathtokens[0], u.Host, sintf)
 	default:
 		return errors.New("unknown call scheme: " + u.Scheme)
 	}
diff --git a/src/yggdrasil/tcp.go b/src/yggdrasil/tcp.go
index 0179c20..2a177c0 100644
--- a/src/yggdrasil/tcp.go
+++ b/src/yggdrasil/tcp.go
@@ -159,7 +159,13 @@ func (t *tcp) listener(l *tcpListener, listenaddr string) {
 	t.mutex.Unlock()
 	// And here we go!
 	accepted := make(chan bool)
-	defer l.listener.Close()
+	defer func() {
+		t.link.core.log.Infoln("Stopping TCP listener on:", l.listener.Addr().String())
+		l.listener.Close()
+		t.mutex.Lock()
+		delete(t.listeners, listenaddr)
+		t.mutex.Unlock()
+	}()
 	t.link.core.log.Infoln("Listening for TCP on:", l.listener.Addr().String())
 	for {
 		var sock net.Conn
@@ -178,13 +184,8 @@ func (t *tcp) listener(l *tcpListener, listenaddr string) {
 				t.link.core.log.Errorln("Failed to accept connection:", err)
 				return
 			}
-			go t.handler(sock, true)
+			go t.handler(sock, true, nil)
 		case <-l.stop:
-			t.link.core.log.Infoln("Stopping TCP listener on:", l.listener.Addr().String())
-			l.listener.Close()
-			t.mutex.Lock()
-			delete(t.listeners, listenaddr)
-			t.mutex.Unlock()
 			return
 		}
 	}
@@ -230,8 +231,12 @@ func (t *tcp) call(saddr string, options interface{}, sintf string) {
 			if sintf != "" {
 				return
 			}
+			dialerdst, er := net.ResolveTCPAddr("tcp", socksaddr)
+			if er != nil {
+				return
+			}
 			var dialer proxy.Dialer
-			dialer, err = proxy.SOCKS5("tcp", socksaddr, nil, proxy.Direct)
+			dialer, err = proxy.SOCKS5("tcp", dialerdst.String(), nil, proxy.Direct)
 			if err != nil {
 				return
 			}
@@ -246,6 +251,7 @@ func (t *tcp) call(saddr string, options interface{}, sintf string) {
 					addr:    saddr,
 				},
 			}
+			t.handler(conn, false, dialerdst.String())
 		} else {
 			dialer := net.Dialer{
 				Control: t.tcpContext,
@@ -302,12 +308,12 @@ func (t *tcp) call(saddr string, options interface{}, sintf string) {
 			if err != nil {
 				return
 			}
+			t.handler(conn, false, nil)
 		}
-		t.handler(conn, false)
 	}()
 }
 
-func (t *tcp) handler(sock net.Conn, incoming bool) {
+func (t *tcp) handler(sock net.Conn, incoming bool, options interface{}) {
 	defer sock.Close()
 	t.setExtraOptions(sock)
 	stream := stream{}
@@ -315,8 +321,16 @@ func (t *tcp) handler(sock net.Conn, incoming bool) {
 	local, _, _ := net.SplitHostPort(sock.LocalAddr().String())
 	remote, _, _ := net.SplitHostPort(sock.RemoteAddr().String())
 	remotelinklocal := net.ParseIP(remote).IsLinkLocalUnicast()
-	name := "tcp://" + sock.RemoteAddr().String()
-	link, err := t.link.core.link.create(&stream, name, "tcp", local, remote, incoming, remotelinklocal)
+	var name string
+	var proto string
+	if socksaddr, issocks := options.(string); issocks {
+		name = "socks://" + socksaddr + "/" + sock.RemoteAddr().String()
+		proto = "socks"
+	} else {
+		name = "tcp://" + sock.RemoteAddr().String()
+		proto = "tcp"
+	}
+	link, err := t.link.core.link.create(&stream, name, proto, local, remote, incoming, remotelinklocal)
 	if err != nil {
 		t.link.core.log.Println(err)
 		panic(err)