1 files changed, 20 insertions, 7 deletions

diff --git a/sockets.go b/sockets.go
index 0ebbe43..2ded258 100644
--- a/sockets.go
+++ b/sockets.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+        "context"
         "io"
         "log"
         "net"
@@ -8,7 +9,7 @@ import (
         "github.com/gorilla/websocket"
 )
 
-func wsReader(wsconn *websocket.Conn, out chan []byte) {
+func wsReader(wsconn *websocket.Conn, out chan []byte, ctx context.Context) {
         for {
                 messageType, p, err := wsconn.ReadMessage()
                 if err != nil {
@@ -19,11 +20,16 @@ func wsReader(wsconn *websocket.Conn, out chan []byte) {
                         log.Println("error: wsReader: only binary messages are supported")
                         continue
                 }
-                out <- p
+                select {
+                case out <- p:
+                        continue
+                case <-ctx.Done():
+                        return
+                }
         }
 }
 
-func socketReader(proxyconn net.Conn, out chan []byte) {
+func socketReader(proxyconn net.Conn, out chan []byte, ctx context.Context) {
         for {
                 readBuffer := make([]byte, 2048)
 
@@ -37,16 +43,21 @@ func socketReader(proxyconn net.Conn, out chan []byte) {
                         return
                 }
 
-                out <- readBuffer[:i]
+                select {
+                case out <- readBuffer[:i]:
+                        continue
+                case <-ctx.Done():
+                        return
+                }
         }
 }
 
-func serviceBoth(wsconn *websocket.Conn, proxyconn net.Conn) {
+func serviceBoth(wsconn *websocket.Conn, proxyconn net.Conn, ctx context.Context) {
         sc := make(chan []byte)
         wsc := make(chan []byte)
 
-        go socketReader(proxyconn, sc)
-        go wsReader(wsconn, wsc)
+        go socketReader(proxyconn, sc, ctx)
+        go wsReader(wsconn, wsc, ctx)
 
         for {
                 select {
@@ -61,6 +72,8 @@ func serviceBoth(wsconn *websocket.Conn, proxyconn net.Conn) {
                                 log.Printf("error: serviceBoth: %s", err)
                                 return
                         }
+                case <-ctx.Done():
+                        return
                 }
         }
 }