From 56cf90b13dc88852e1938ab60107eadfa36cdf4d Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Thu, 10 Oct 2019 21:07:55 +0300
Subject: [PATCH] Correct closing of connections

---
 hub/connection.go |  1 +
 proxy/direct.go   | 11 +++++++----
 proxy/middle.go   | 13 ++++++++-----
 3 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/hub/connection.go b/hub/connection.go
index ee6f4ec..4420a97 100644
--- a/hub/connection.go
+++ b/hub/connection.go
@@ -52,6 +52,7 @@ func (c *connection) write(packet conntypes.Packet) error {
 
 func (c *connection) shutdown() {
 	c.shutdownOnce.Do(func() {
+		c.conn.Close()
 		close(c.done)
 		c.hub.channelBrokenSockets <- c.id
 	})
diff --git a/proxy/direct.go b/proxy/direct.go
index 19e7f69..23266ae 100644
--- a/proxy/direct.go
+++ b/proxy/direct.go
@@ -27,14 +27,17 @@ func directConnection(request *protocol.TelegramRequest) error {
 	go directPipe(telegramConn, request.ClientConn, wg, request.Logger)
 	go directPipe(request.ClientConn, telegramConn, wg, request.Logger)
 
-	<-request.Ctx.Done()
 	wg.Wait()
 
-	return request.Ctx.Err()
+	return nil
 }
 
-func directPipe(dst io.Writer, src io.Reader, wg *sync.WaitGroup, logger *zap.SugaredLogger) {
-	defer wg.Done()
+func directPipe(dst io.WriteCloser, src io.ReadCloser, wg *sync.WaitGroup, logger *zap.SugaredLogger) {
+	defer func() {
+		dst.Close()
+		src.Close()
+		wg.Done()
+	}()
 
 	buf := make([]byte, directPipeBufferSize)
 	if _, err := io.CopyBuffer(dst, src, buf); err != nil {
diff --git a/proxy/middle.go b/proxy/middle.go
index d806885..c8c79e3 100644
--- a/proxy/middle.go
+++ b/proxy/middle.go
@@ -32,17 +32,20 @@ func middleConnection(request *protocol.TelegramRequest) error {
 	go middlePipe(telegramConn, clientConn, wg, request.Logger)
 	go middlePipe(clientConn, telegramConn, wg, request.Logger)
 
-	<-request.Ctx.Done()
 	wg.Wait()
 
-	return request.Ctx.Err()
+	return nil
 }
 
-func middlePipe(dst conntypes.PacketAckWriter,
-	src conntypes.PacketAckReader,
+func middlePipe(dst conntypes.PacketAckWriteCloser,
+	src conntypes.PacketAckReadCloser,
 	wg *sync.WaitGroup,
 	logger *zap.SugaredLogger) {
-	defer wg.Done()
+	defer func() {
+		dst.Close()
+		src.Close()
+		wg.Done()
+	}()
 
 	for {
 		acks := conntypes.ConnectionAcks{}