From 6ff219420b6638e627d8c4c9bef06939b734c339 Mon Sep 17 00:00:00 2001
From: Mark Beamer Jr <markbeamerjr@gmail.com>
Date: Wed, 13 Jun 2018 20:21:47 -0400
Subject: [PATCH] put back connection closing.

---
 peer/server.go      | 11 +++++++++--
 reflector/server.go | 11 +++++++++--
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/peer/server.go b/peer/server.go
index 5143c85..9cce245 100644
--- a/peer/server.go
+++ b/peer/server.go
@@ -84,13 +84,20 @@ func (s *Server) listenAndServe(listener net.Listener) {
 			log.Error(err)
 		} else {
 			s.stop.Add(1)
-			go s.handleConnection(conn)
+			go func() {
+				s.stop.Done()
+				s.handleConnection(conn)
+			}()
 		}
 	}
 }
 
 func (s *Server) handleConnection(conn net.Conn) {
-	defer s.stop.Done()
+	defer func() {
+		if err := conn.Close(); err != nil {
+			log.Error("error closing client connection for peer server - ", err)
+		}
+	}()
 	timeoutDuration := 5 * time.Second
 
 	for {
diff --git a/reflector/server.go b/reflector/server.go
index 48cdaa1..9f473d6 100644
--- a/reflector/server.go
+++ b/reflector/server.go
@@ -74,13 +74,20 @@ func (s *Server) listenAndServe(listener net.Listener) {
 			log.Error(err)
 		} else {
 			s.stop.Add(1)
-			go s.handleConn(conn)
+			go func() {
+				s.stop.Done()
+				s.handleConn(conn)
+			}()
 		}
 	}
 }
 
 func (s *Server) handleConn(conn net.Conn) {
-	defer s.stop.Done()
+	defer func() {
+		if err := conn.Close(); err != nil {
+			log.Error("error closing client connection for reflector server - ", err)
+		}
+	}()
 	// TODO: connection should time out eventually
 	err := s.doHandshake(conn)
 	if err != nil {