From 2eba78b0db9f5b2a5af75bd9223f427665655bf1 Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Thu, 10 Oct 2019 16:42:29 +0300
Subject: [PATCH] Add unregistering on telegram conn close

---
 hub/hub.go                  | 3 +++
 wrappers/packetack/proxy.go | 8 +++++++-
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/hub/hub.go b/hub/hub.go
index 6e76b4d..09c76d1 100644
--- a/hub/hub.go
+++ b/hub/hub.go
@@ -54,6 +54,9 @@ func (h *hub) getHub(req *protocol.TelegramRequest) *connectionHub {
 
 		rv, ok = h.subs[key]
 		if !ok {
+			h.logger.Debugw("Create new connection hub",
+				"dc", req.ClientProtocol.DC(),
+				"protocol", req.ClientProtocol.ConnectionProtocol())
 			rv = newConnectionHub(h.logger.With(
 				"dc", req.ClientProtocol.DC(),
 				"protocol", req.ClientProtocol.ConnectionProtocol(),
diff --git a/wrappers/packetack/proxy.go b/wrappers/packetack/proxy.go
index bce700f..301035c 100644
--- a/wrappers/packetack/proxy.go
+++ b/wrappers/packetack/proxy.go
@@ -5,6 +5,7 @@ import (
 	"encoding/binary"
 	"fmt"
 	"net"
+	"sync"
 
 	"github.com/9seconds/mtg/config"
 	"github.com/9seconds/mtg/conntypes"
@@ -19,6 +20,7 @@ type wrapperProxy struct {
 	clientIPPort []byte
 	ourIPPort    []byte
 	channelRead  hub.ChannelReadCloser
+	closeOnce    sync.Once
 }
 
 func (w *wrapperProxy) Write(packet conntypes.Packet, acks *conntypes.ConnectionAcks) error {
@@ -61,7 +63,11 @@ func (w *wrapperProxy) Read(acks *conntypes.ConnectionAcks) (conntypes.Packet, e
 }
 
 func (w *wrapperProxy) Close() error {
-	return w.channelRead.Close()
+	w.closeOnce.Do(func() {
+		w.channelRead.Close()
+		hub.Registry.Unregister(w.request.ConnID)
+	})
+	return nil
 }
 
 func NewProxy(request *protocol.TelegramRequest) conntypes.PacketAckReadWriteCloser {