From d559ce9da2efaa493960fada7bf4a310a8b72eb8 Mon Sep 17 00:00:00 2001
From: Asim Aslam <asim@aslam.me>
Date: Wed, 4 Sep 2019 15:41:57 +0100
Subject: [PATCH] Provide Links() method in Tunnel

---
 tunnel/default.go     | 13 +++++++++++++
 tunnel/link.go        | 14 ++++++++++++--
 tunnel/tunnel.go      | 10 ++++++++++
 tunnel/tunnel_test.go | 23 +++++------------------
 4 files changed, 40 insertions(+), 20 deletions(-)

diff --git a/tunnel/default.go b/tunnel/default.go
index 0ab11ea2..d62dd66e 100644
--- a/tunnel/default.go
+++ b/tunnel/default.go
@@ -911,6 +911,19 @@ func (t *tun) Listen(channel string) (Listener, error) {
 	return tl, nil
 }
 
+func (t *tun) Links() []Link {
+	t.RLock()
+	defer t.RUnlock()
+
+	var links []Link
+
+	for _, link := range t.links {
+		links = append(links, link)
+	}
+
+	return links
+}
+
 func (t *tun) String() string {
 	return "mucp"
 }
diff --git a/tunnel/link.go b/tunnel/link.go
index 7e14e9d8..9470ef6a 100644
--- a/tunnel/link.go
+++ b/tunnel/link.go
@@ -81,14 +81,24 @@ func (l *link) run() {
 	}
 }
 
-func (l *link) Close() {
+func (l *link) Id() string {
+	l.RLock()
+	defer l.RUnlock()
+
+	return l.id
+}
+
+func (l *link) Close() error {
 	l.Lock()
 	defer l.Unlock()
 
 	select {
 	case <-l.closed:
-		return
+		return nil
 	default:
 		close(l.closed)
+		return l.Socket.Close()
 	}
+
+	return nil
 }
diff --git a/tunnel/tunnel.go b/tunnel/tunnel.go
index e76f8437..7c9e2afc 100644
--- a/tunnel/tunnel.go
+++ b/tunnel/tunnel.go
@@ -31,10 +31,20 @@ type Tunnel interface {
 	Dial(channel string, opts ...DialOption) (Session, error)
 	// Accept connections on a channel
 	Listen(channel string) (Listener, error)
+	// All the links the tunnel is connected to
+	Links() []Link
 	// Name of the tunnel implementation
 	String() string
 }
 
+// Link represents internal links to the tunnel
+type Link interface {
+	// The id of the link
+	Id() string
+	// honours transport socket
+	transport.Socket
+}
+
 // The listener provides similar constructs to the transport.Listener
 type Listener interface {
 	Accept() (Session, error)
diff --git a/tunnel/tunnel_test.go b/tunnel/tunnel_test.go
index 3fc84b6d..fc76421b 100644
--- a/tunnel/tunnel_test.go
+++ b/tunnel/tunnel_test.go
@@ -187,30 +187,15 @@ func testBrokenTunAccept(t *testing.T, tun Tunnel, wait chan bool, wg *sync.Wait
 	if err := c.Recv(m); err != nil {
 		t.Fatal(err)
 	}
-	tun.Close()
 
-	// re-start tunnel
-	err = tun.Connect()
-	if err != nil {
-		t.Fatal(err)
-	}
-	defer tun.Close()
-
-	// listen on some virtual address
-	tl, err = tun.Listen("test-tunnel")
-	if err != nil {
-		t.Fatal(err)
+	// close all the links
+	for _, link := range tun.Links() {
+		link.Close()
 	}
 
 	// receiver ready; notify sender
 	wait <- true
 
-	// accept a connection
-	c, err = tl.Accept()
-	if err != nil {
-		t.Fatal(err)
-	}
-
 	// accept the message
 	m = new(transport.Message)
 	if err := c.Recv(m); err != nil {
@@ -279,6 +264,7 @@ func TestReconnectTunnel(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	defer tunB.Close()
 
 	// we manually override the tunnel.ReconnectTime value here
 	// this is so that we make the reconnects faster than the default 5s
@@ -289,6 +275,7 @@ func TestReconnectTunnel(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	defer tunA.Close()
 
 	wait := make(chan bool)