1 files changed, 26 insertions, 64 deletions

diff --git a/sockets.go b/sockets.go
index 2ded258..319b548 100644
--- a/sockets.go
+++ b/sockets.go
@@ -1,79 +1,41 @@
 package main
 
 import (
-        "context"
+        "fmt"
         "io"
-        "log"
-        "net"
 
         "github.com/gorilla/websocket"
 )
 
-func wsReader(wsconn *websocket.Conn, out chan []byte, ctx context.Context) {
-        for {
-                messageType, p, err := wsconn.ReadMessage()
-                if err != nil {
-                        log.Printf("error: wsReader: %s", err)
-                        return
-                }
-                if messageType != websocket.BinaryMessage {
-                        log.Println("error: wsReader: only binary messages are supported")
-                        continue
-                }
-                select {
-                case out <- p:
-                        continue
-                case <-ctx.Done():
-                        return
-                }
-        }
+type WebsocketReadWriter struct {
+        W *websocket.Conn
 }
 
-func socketReader(proxyconn net.Conn, out chan []byte, ctx context.Context) {
-        for {
-                readBuffer := make([]byte, 2048)
-
-                i, err := proxyconn.Read(readBuffer)
-                if err != nil {
-                        if err == io.EOF {
-                                log.Println("info: socketReader: Disconnected")
-                        } else {
-                                log.Printf("error: socketReader: %s", err)
-                        }
-                        return
-                }
-
-                select {
-                case out <- readBuffer[:i]:
-                        continue
-                case <-ctx.Done():
-                        return
-                }
+func (w *WebsocketReadWriter) Read(p []byte) (int, error) {
+        t, m, err := w.W.ReadMessage()
+        if err != nil {
+                return 0, fmt.Errorf("error reading websocket: %w", err)
+        }
+        if t != websocket.BinaryMessage {
+                return 0, fmt.Errorf("error reading websocket: only binary messages are supported")
+        }
+        if cap(p) < len(m) {
+                return 0, fmt.Errorf("error reading websocket: outputbuffer too short for input")
+        }
+        if i := copy(p, m); i != len(m) {
+                return 0, fmt.Errorf("error reading websocket: copy failed %d of %d", i, len(m))
         }
+        return len(m), nil
 }
 
-func serviceBoth(wsconn *websocket.Conn, proxyconn net.Conn, ctx context.Context) {
-        sc := make(chan []byte)
-        wsc := make(chan []byte)
-
-        go socketReader(proxyconn, sc, ctx)
-        go wsReader(wsconn, wsc, ctx)
-
-        for {
-                select {
-                case sd := <-sc:
-                        if err := wsconn.WriteMessage(websocket.BinaryMessage, sd); err != nil {
-                                log.Printf("error: serviceBoth: %s", err)
-                                return
-                        }
-
-                case wsd := <-wsc:
-                        if _, err := proxyconn.Write(wsd); err != nil {
-                                log.Printf("error: serviceBoth: %s", err)
-                                return
-                        }
-                case <-ctx.Done():
-                        return
-                }
+func (w *WebsocketReadWriter) Write(p []byte) (int, error) {
+        if err := w.W.WriteMessage(websocket.BinaryMessage, p); err != nil {
+                return 0, fmt.Errorf("error writing websocket: %w", err)
         }
+        return len(p), nil
+}
+
+func serviceBoth(dst io.Writer, src io.Reader, errc chan<- error) {
+        _, err := io.Copy(dst, src)
+        errc <- err
 }