From 7e726b0afbea2fb1aafa406a1a5de96d538f9ea4 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Fri, 19 Apr 2019 23:04:09 +0100
Subject: [PATCH] Listener should clean up a bit more when closing

---
 src/yggdrasil/core.go     | 1 +
 src/yggdrasil/listener.go | 4 ++++
 2 files changed, 5 insertions(+)

diff --git a/src/yggdrasil/core.go b/src/yggdrasil/core.go
index dba1c64..b5d74e8 100644
--- a/src/yggdrasil/core.go
+++ b/src/yggdrasil/core.go
@@ -262,6 +262,7 @@ func (c *Core) ListenConn() (*Listener, error) {
 		return nil, errors.New("a listener already exists")
 	}
 	c.sessions.listener = &Listener{
+		core:  c,
 		conn:  make(chan *Conn),
 		close: make(chan interface{}),
 	}
diff --git a/src/yggdrasil/listener.go b/src/yggdrasil/listener.go
index 268d8b7..6222541 100644
--- a/src/yggdrasil/listener.go
+++ b/src/yggdrasil/listener.go
@@ -7,6 +7,7 @@ import (
 
 // Listener waits for incoming sessions
 type Listener struct {
+	core  *Core
 	conn  chan *Conn
 	close chan interface{}
 }
@@ -30,6 +31,9 @@ func (l *Listener) Close() (err error) {
 		recover()
 		err = errors.New("already closed")
 	}()
+	if l.core.sessions.listener == l {
+		l.core.sessions.listener = nil
+	}
 	close(l.close)
 	close(l.conn)
 	return nil