From a56d5e188cb45f687908aa2910f98bc969af3ec5 Mon Sep 17 00:00:00 2001
From: Vasiliy Tolstov <v.tolstov@selfip.ru>
Date: Mon, 26 Jun 2017 17:31:00 +0300
Subject: [PATCH] small fixups

Signed-off-by: Vasiliy Tolstov <v.tolstov@selfip.ru>
---
 client.go |  2 ++
 server.go | 10 ++++++++++
 2 files changed, 12 insertions(+)

diff --git a/client.go b/client.go
index 5938a3a..16e8f3d 100644
--- a/client.go
+++ b/client.go
@@ -39,7 +39,9 @@ func Connect(ctx context.Context, c net.Conn, cfg *ClientConfig) (*ClientConn, e
 	}
 
 	for _, h := range cfg.Handlers {
+		fmt.Printf("%#+v\n", h)
 		if err := h.Handle(conn); err != nil {
+			fmt.Printf("rrr %v\n", err)
 			conn.Close()
 			cfg.ErrorCh <- err
 			return nil, err
diff --git a/server.go b/server.go
index caeaa24..002a5d3 100644
--- a/server.go
+++ b/server.go
@@ -294,14 +294,19 @@ func (*DefaultServerMessageHandler) Handle(c Conn) error {
 	}
 	wg.Add(2)
 
+	quit := make(chan struct{})
+
 	// server
 	go func() {
 		defer wg.Done()
 		for {
 			select {
+			case <-quit:
+				return
 			case msg := <-cfg.ServerMessageCh:
 				if err = msg.Write(c); err != nil {
 					cfg.errorCh <- err
+					close(quit)
 					return
 				}
 			}
@@ -313,20 +318,25 @@ func (*DefaultServerMessageHandler) Handle(c Conn) error {
 		defer wg.Done()
 		for {
 			select {
+			case <-quit:
+				return
 			default:
 				var messageType ClientMessageType
 				if err := binary.Read(c, binary.BigEndian, &messageType); err != nil {
 					cfg.errorCh <- err
+					close(quit)
 					return
 				}
 				msg, ok := clientMessages[messageType]
 				if !ok {
 					cfg.errorCh <- fmt.Errorf("unsupported message-type: %v", messageType)
+					close(quit)
 					return
 				}
 				parsedMsg, err := msg.Read(c)
 				if err != nil {
 					cfg.errorCh <- err
+					close(quit)
 					return
 				}
 				cfg.ClientMessageCh <- parsedMsg