Initial import

8 files changed, 282 insertions, 0 deletions

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..4b301d2
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+/websocket-proxy
diff --git a/Makefile b/Makefile
new file mode 100644
index 0000000..17cd4ec
--- /dev/null
+++ b/Makefile
@@ -0,0 +1,7 @@
+websocket-proxy: main.go sockets.go server.go client.go
+        go build -o $@ $^
+
+
+.PHONY: clean
+clean:
+        rm -rf websocket-proxy
diff --git a/client.go b/client.go
new file mode 100644
index 0000000..00455ab
--- /dev/null
+++ b/client.go
@@ -0,0 +1,45 @@
+package main
+
+import (
+        "log"
+        "net"
+
+        "github.com/gorilla/websocket"
+)
+
+type ClientHandler struct {
+        SocketListenOn  string
+        WebsocketServer string
+}
+
+func (h *ClientHandler) ServiceConnection(proxyconn net.Conn) {
+        defer proxyconn.Close()
+
+        wsconn, _, err := websocket.DefaultDialer.Dial(h.WebsocketServer, nil)
+        if err != nil {
+                log.Println(err)
+                return
+        }
+        defer wsconn.Close()
+
+        log.Println("Connected to server")
+
+        serviceBoth(wsconn, proxyconn)
+}
+
+func (h *ClientHandler) Run() {
+        listener, err := net.Listen("tcp", h.SocketListenOn)
+        if err != nil {
+                log.Printf("error: Run: %s", err)
+                return
+        }
+
+        for {
+                conn, err := listener.Accept()
+                if err != nil {
+                        log.Printf("error: Run: %s", err)
+                        continue
+                }
+                go h.ServiceConnection(conn)
+        }
+}
diff --git a/go.mod b/go.mod
new file mode 100644
index 0000000..423ea5c
--- /dev/null
+++ b/go.mod
@@ -0,0 +1,8 @@
+module code.crute.us/mcrute/websocket_proxy
+
+go 1.13
+
+require (
+        github.com/gorilla/websocket v1.4.1
+        github.com/spf13/cobra v0.0.5
+)
diff --git a/go.sum b/go.sum
new file mode 100644
index 0000000..84ffc4f
--- /dev/null
+++ b/go.sum
@@ -0,0 +1,35 @@
+github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
+github.com/armon/consul-api v0.0.0-20180202201655-eb2c6b5be1b6/go.mod h1:grANhF5doyWs3UAsr3K4I6qtAmlQcZDesFNEHPZAzj8=
+github.com/coreos/etcd v3.3.10+incompatible/go.mod h1:uF7uidLiAD3TWHmW31ZFd/JWoc32PjwdhPthX9715RE=
+github.com/coreos/go-etcd v2.0.0+incompatible/go.mod h1:Jez6KQU2B/sWsbdaef3ED8NzMklzPG4d5KIOhIy30Tk=
+github.com/coreos/go-semver v0.2.0/go.mod h1:nnelYz7RCh+5ahJtPPxZlU+153eP4D4r3EedlOD2RNk=
+github.com/cpuguy83/go-md2man v1.0.10/go.mod h1:SmD6nW6nTyfqj6ABTjUi3V3JVMnlJmwcJI5acqYI6dE=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
+github.com/gorilla/websocket v1.4.1 h1:q7AeDBpnBk8AogcD4DSag/Ukw/KV+YhzLj2bP5HvKCM=
+github.com/gorilla/websocket v1.4.1/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
+github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
+github.com/inconshreveable/mousetrap v1.0.0 h1:Z8tu5sraLXCXIcARxBp/8cbvlwVa7Z1NHg9XEKhtSvM=
+github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8=
+github.com/magiconair/properties v1.8.0/go.mod h1:PppfXfuXeibc/6YijjN8zIbojt8czPbwD3XqdrwzmxQ=
+github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=
+github.com/mitchellh/mapstructure v1.1.2/go.mod h1:FVVH3fgwuzCH5S8UJGiWEs2h04kUh9fWfEaFds41c1Y=
+github.com/pelletier/go-toml v1.2.0/go.mod h1:5z9KED0ma1S8pY6P1sdut58dfprrGBbd/94hg7ilaic=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/russross/blackfriday v1.5.2/go.mod h1:JO/DiYxRf+HjHt06OyowR9PTA263kcR/rfWxYHBV53g=
+github.com/spf13/afero v1.1.2/go.mod h1:j4pytiNVoe2o6bmDsKpLACNPDBIoEAkihy7loJ1B0CQ=
+github.com/spf13/cast v1.3.0/go.mod h1:Qx5cxh0v+4UWYiBimWS+eyWzqEqokIECu5etghLkUJE=
+github.com/spf13/cobra v0.0.5 h1:f0B+LkLX6DtmRH1isoNA9VTtNUK9K8xYd28JNNfOv/s=
+github.com/spf13/cobra v0.0.5/go.mod h1:3K3wKZymM7VvHMDS9+Akkh4K60UwM26emMESw8tLCHU=
+github.com/spf13/jwalterweatherman v1.0.0/go.mod h1:cQK4TGJAtQXfYWX+Ddv3mKDzgVb68N+wFjFa4jdeBTo=
+github.com/spf13/pflag v1.0.3 h1:zPAT6CGy6wXeQ7NtTnaTerfKOsV6V6F8agHXFiazDkg=
+github.com/spf13/pflag v1.0.3/go.mod h1:DYY7MBk1bdzusC3SYhjObp+wFpr4gzcvqqNjLnInEg4=
+github.com/spf13/viper v1.3.2/go.mod h1:ZiWeW+zYFKm7srdB9IoDzzZXaJaI5eL9QjNiN/DMA2s=
+github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
+github.com/ugorji/go/codec v0.0.0-20181204163529-d75b2dcb6bc8/go.mod h1:VFNgLljTbGfSG7qAOspJ7OScBnGdDN/yBr0sguwnwf0=
+github.com/xordataexchange/crypt v0.0.3-0.20170626215501-b2862e3d0a77/go.mod h1:aYKd//L2LvnjZzWKhF00oedf4jCCReLcmhLdhm1A27Q=
+golang.org/x/crypto v0.0.0-20181203042331-505ab145d0a9/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4=
+golang.org/x/sys v0.0.0-20181205085412-a5c9d58dba9a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
+golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
diff --git a/main.go b/main.go
new file mode 100644
index 0000000..75555e6
--- /dev/null
+++ b/main.go
@@ -0,0 +1,74 @@
+package main
+
+import (
+        "errors"
+        "fmt"
+        "log"
+        "net/http"
+        "os"
+        "strings"
+
+        "github.com/spf13/cobra"
+)
+
+var rootCmd = &cobra.Command{
+        Use:     "websocket-proxy",
+        Version: "0.1.0",
+        Short:   "Proxy TCP connections over a websocket",
+}
+
+var clientCmd = &cobra.Command{
+        Use:   "client [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) {
+                listenOn := cmd.Flag("listen").Value.String()
+
+                h := &ClientHandler{
+                        SocketListenOn:  listenOn,
+                        WebsocketServer: args[0],
+                }
+
+                log.Printf("Serving on %s", listenOn)
+                h.Run()
+        },
+}
+
+var serverCmd = &cobra.Command{
+        Use:   "server [next-hop host]",
+        Short: "Serve websocket proxy client",
+        Args: func(cmd *cobra.Command, args []string) error {
+                if len(args) != 1 || args[0] == "" {
+                        return errors.New("Next-hop host is a required argument")
+                }
+                return nil
+        },
+        Run: func(cmd *cobra.Command, args []string) {
+                listenOn := cmd.Flag("listen").Value.String()
+                log.Printf("Serving on %s", listenOn)
+
+                http.Handle("/", NewServerHandler(args[0]))
+                log.Fatal(http.ListenAndServe(listenOn, nil))
+        },
+}
+
+func main() {
+        rootCmd.AddCommand(clientCmd)
+        rootCmd.AddCommand(serverCmd)
+
+        clientCmd.Flags().StringP("listen", "l", ":9013", "[address]:port to bind for serving clients")
+        serverCmd.Flags().StringP("listen", "l", ":9012", "[address]:port to bind for serving clients")
+
+        if err := rootCmd.Execute(); err != nil {
+                fmt.Println(err)
+                os.Exit(1)
+        }
+}
diff --git a/server.go b/server.go
new file mode 100644
index 0000000..47eb3e5
--- /dev/null
+++ b/server.go
@@ -0,0 +1,46 @@
+package main
+
+import (
+        "log"
+        "net"
+        "net/http"
+
+        "github.com/gorilla/websocket"
+)
+
+type ServerHandler struct {
+        ProxyToHost string
+        upgrader    websocket.Upgrader
+}
+
+func NewServerHandler(proxyToHost string) *ServerHandler {
+        return &ServerHandler{
+                ProxyToHost: proxyToHost,
+                upgrader: websocket.Upgrader{
+                        ReadBufferSize:  1024,
+                        WriteBufferSize: 1024,
+                },
+        }
+}
+
+func (h *ServerHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+        log.Println("Got new HTTP connection")
+
+        wsconn, err := h.upgrader.Upgrade(w, r, nil)
+        if err != nil {
+                log.Printf("error: ServeHTTP: %s", err)
+                return
+        }
+        defer wsconn.Close()
+
+        proxyconn, err := net.Dial("tcp", h.ProxyToHost)
+        if err != nil {
+                log.Printf("error: ServeHTTP: %s", err)
+                return
+        }
+        defer proxyconn.Close()
+
+        log.Println("Connected to SSH server")
+
+        serviceBoth(wsconn, proxyconn)
+}
diff --git a/sockets.go b/sockets.go
new file mode 100644
index 0000000..0ebbe43
--- /dev/null
+++ b/sockets.go
@@ -0,0 +1,66 @@
+package main
+
+import (
+        "io"
+        "log"
+        "net"
+
+        "github.com/gorilla/websocket"
+)
+
+func wsReader(wsconn *websocket.Conn, out chan []byte) {
+        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
+                }
+                out <- p
+        }
+}
+
+func socketReader(proxyconn net.Conn, out chan []byte) {
+        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
+                }
+
+                out <- readBuffer[:i]
+        }
+}
+
+func serviceBoth(wsconn *websocket.Conn, proxyconn net.Conn) {
+        sc := make(chan []byte)
+        wsc := make(chan []byte)
+
+        go socketReader(proxyconn, sc)
+        go wsReader(wsconn, wsc)
+
+        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
+                        }
+                }
+        }
+}