Cleanup, add local client

7 files changed, 109 insertions, 68 deletions

diff --git a/Makefile b/Makefile
index 7a77930..3e599d0 100644
--- a/Makefile
+++ b/Makefile
@@ -1,7 +1,7 @@
 IMAGE="docker.crute.me/websocket-proxy"
 VERSION="0.1.0"
 
-websocket-proxy: main.go sockets.go server.go client.go
+websocket-proxy: main.go sockets.go server.go client.go localclient.go
         CGO_ENABLED=0 go build -ldflags "-X main.version=$(VERSION)" -o $@ $^
 
 .PHONY: docker
diff --git a/client.go b/client.go
index d0716b1..d9a0f25 100644
--- a/client.go
+++ b/client.go
@@ -26,7 +26,14 @@ func (h *ClientHandler) ServiceConnection(proxyconn net.Conn) {
 
         log.Println("Connected to server")
 
-        serviceBoth(wsconn, proxyconn, h.Context)
+        errc := make(chan error)
+        ws := &WebsocketReadWriter{wsconn}
+
+        go serviceBoth(proxyconn, ws, errc)
+        go serviceBoth(ws, proxyconn, errc)
+
+        <-errc
+        log.Println("Closing client connection")
 }
 
 func (h *ClientHandler) Run() {
diff --git a/go.mod b/go.mod
index 423ea5c..3f372f6 100644
--- a/go.mod
+++ b/go.mod
@@ -1,8 +1,13 @@
 module code.crute.us/mcrute/websocket_proxy
 
-go 1.13
+go 1.20
 
 require (
         github.com/gorilla/websocket v1.4.1
         github.com/spf13/cobra v0.0.5
 )
+
+require (
+        github.com/inconshreveable/mousetrap v1.0.0 // indirect
+        github.com/spf13/pflag v1.0.3 // indirect
+)
diff --git a/localclient.go b/localclient.go
new file mode 100644
index 0000000..c1f5c38
--- /dev/null
+++ b/localclient.go
@@ -0,0 +1,31 @@
+package main
+
+import (
+        "context"
+        "log"
+        "os"
+
+        "github.com/gorilla/websocket"
+)
+
+type LocalClientHandler struct {
+        WebsocketServer string
+        Context         context.Context
+}
+
+func (h *LocalClientHandler) Run() {
+        wsconn, _, err := websocket.DefaultDialer.Dial(h.WebsocketServer, nil)
+        if err != nil {
+                log.Println(err)
+                return
+        }
+        defer wsconn.Close()
+
+        errc := make(chan error)
+        ws := &WebsocketReadWriter{wsconn}
+
+        go serviceBoth(os.Stdout, ws, errc)
+        go serviceBoth(ws, os.Stdin, errc)
+
+        log.Printf("Closing client connection %s", <-errc)
+}
diff --git a/main.go b/main.go
index 00fe063..1ac3351 100644
--- a/main.go
+++ b/main.go
@@ -50,6 +50,32 @@ var clientCmd = &cobra.Command{
         },
 }
 
+var localClientCmd = &cobra.Command{
+        Use:   "localclient [server host]",
+        Short: "Act as a client for a websocket-proxy server",
+        Args: func(cmd *cobra.Command, args []string) error {
+                if len(args) != 1 || args[0] == "" {
+                        return errors.New("Server host is a required argument")
+                }
+                if !strings.HasPrefix(args[0], "ws://") && !strings.HasPrefix(args[0], "wss://") {
+                        return errors.New("Server host format is ws[s]://host[:port]/[path]")
+                }
+                return nil
+        },
+        Run: func(cmd *cobra.Command, args []string) {
+                // TODO: Handle signals
+                ctx, cancel := context.WithCancel(context.Background())
+                defer cancel()
+
+                h := &LocalClientHandler{
+                        WebsocketServer: args[0],
+                        Context:         ctx,
+                }
+
+                h.Run()
+        },
+}
+
 var serverCmd = &cobra.Command{
         Use:   "server [next-hop host]",
         Short: "Serve websocket proxy client",
@@ -69,7 +95,10 @@ var serverCmd = &cobra.Command{
 }
 
 func main() {
+        log.SetOutput(os.Stderr)
+
         rootCmd.AddCommand(clientCmd)
+        rootCmd.AddCommand(localClientCmd)
         rootCmd.AddCommand(serverCmd)
 
         clientCmd.Flags().StringP("listen", "l", ":9013", "[address]:port to bind for serving clients")
diff --git a/server.go b/server.go
index 1faf33d..9a9d897 100644
--- a/server.go
+++ b/server.go
@@ -42,5 +42,12 @@ func (h *ServerHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
         log.Println("Connected to SSH server")
 
-        serviceBoth(wsconn, proxyconn, r.Context())
+        errc := make(chan error)
+        ws := &WebsocketReadWriter{wsconn}
+
+        go serviceBoth(proxyconn, ws, errc)
+        go serviceBoth(ws, proxyconn, errc)
+
+        <-errc
+        log.Println("Closing client server")
 }
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
 }