From 1a0aa50471446dc181c80a47f939bb28d1bc31ef Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20Ru=CC=88th?= <rueth@comsys.rwth-aachen.de>
Date: Tue, 8 Jan 2019 10:43:20 +0100
Subject: [PATCH] delete unused connections from the multiplexer

---
 mock_multiplexer_test.go   | 12 ++++++++++++
 multiplexer.go             | 13 +++++++++++++
 packet_handler_map.go      |  2 +-
 packet_handler_map_test.go | 11 +++++++++++
 4 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/mock_multiplexer_test.go b/mock_multiplexer_test.go
index 2c1ac30d..921dfaae 100644
--- a/mock_multiplexer_test.go
+++ b/mock_multiplexer_test.go
@@ -46,3 +46,15 @@ func (m *MockMultiplexer) AddConn(arg0 net.PacketConn, arg1 int) (packetHandlerM
 func (mr *MockMultiplexerMockRecorder) AddConn(arg0, arg1 interface{}) *gomock.Call {
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AddConn", reflect.TypeOf((*MockMultiplexer)(nil).AddConn), arg0, arg1)
 }
+
+// RemoveConn mocks base method
+func (m *MockMultiplexer) RemoveConn(arg0 net.PacketConn) error {
+	ret := m.ctrl.Call(m, "RemoveConn", arg0)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// RemoveConn indicates an expected call of RemoveConn
+func (mr *MockMultiplexerMockRecorder) RemoveConn(arg0 interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RemoveConn", reflect.TypeOf((*MockMultiplexer)(nil).RemoveConn), arg0)
+}
diff --git a/multiplexer.go b/multiplexer.go
index c4482ac2..e8c3b7db 100644
--- a/multiplexer.go
+++ b/multiplexer.go
@@ -15,6 +15,7 @@ var (
 
 type multiplexer interface {
 	AddConn(net.PacketConn, int) (packetHandlerManager, error)
+	RemoveConn(net.PacketConn) error
 }
 
 type connManager struct {
@@ -61,3 +62,15 @@ func (m *connMultiplexer) AddConn(c net.PacketConn, connIDLen int) (packetHandle
 	}
 	return p.manager, nil
 }
+
+func (m *connMultiplexer) RemoveConn(c net.PacketConn) error {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+
+	if _, ok := m.conns[c]; !ok {
+		return fmt.Errorf("cannote remove connection, connection is unknown")
+	}
+
+	delete(m.conns, c)
+	return nil
+}
diff --git a/packet_handler_map.go b/packet_handler_map.go
index 3c0f5bdd..6ce4e7fd 100644
--- a/packet_handler_map.go
+++ b/packet_handler_map.go
@@ -139,7 +139,7 @@ func (h *packetHandlerMap) close(e error) error {
 	}
 	h.mutex.Unlock()
 	wg.Wait()
-	return nil
+	return getMultiplexer().RemoveConn(h.conn)
 }
 
 func (h *packetHandlerMap) listen() {
diff --git a/packet_handler_map_test.go b/packet_handler_map_test.go
index 41bd842d..6d9e83a8 100644
--- a/packet_handler_map_test.go
+++ b/packet_handler_map_test.go
@@ -45,6 +45,16 @@ var _ = Describe("Packet Handler Map", func() {
 	})
 
 	It("closes", func() {
+		getMultiplexer() // make the sync.Once execute
+		// replace the clientMuxer. getClientMultiplexer will now return the MockMultiplexer
+		mockMultiplexer := NewMockMultiplexer(mockCtrl)
+		origMultiplexer := connMuxer
+		connMuxer = mockMultiplexer
+
+		defer func() {
+			connMuxer = origMultiplexer
+		}()
+
 		testErr := errors.New("test error	")
 		sess1 := NewMockPacketHandler(mockCtrl)
 		sess1.EXPECT().destroy(testErr)
@@ -52,6 +62,7 @@ var _ = Describe("Packet Handler Map", func() {
 		sess2.EXPECT().destroy(testErr)
 		handler.Add(protocol.ConnectionID{1, 1, 1, 1}, sess1)
 		handler.Add(protocol.ConnectionID{2, 2, 2, 2}, sess2)
+		mockMultiplexer.EXPECT().RemoveConn(gomock.Any())
 		handler.close(testErr)
 	})