Initial commit

README.md

@@ -0,0 +1,3 @@
+# Reimplementation of cisoakk's websocket relay example in GO
+
+https://github.com/cisoakk/wsServer/tree/master/examples/relay

go.mod

@@ -0,0 +1,10 @@
+module git.int.sw4me.com/akovalenko/wsrelay
+
+go 1.17
+
+require (
+	github.com/sirupsen/logrus v1.9.3
+	nhooyr.io/websocket v1.8.11
+)
+
+require golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8 // indirect

go.sum

@@ -0,0 +1,17 @@
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
+github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
+github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8 h1:0A+M6Uqn+Eje4kHMK80dtF3JCXC4ykBgQG4Fe06QRhQ=
+golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+nhooyr.io/websocket v1.8.11 h1:f/qXNc2/3DpoSZkHt1DQu6rj4zGC8JmkkLkWss0MgN0=
+nhooyr.io/websocket v1.8.11/go.mod h1:rN9OFWIUwuxg4fR5tELlYC04bXYowCP9GX47ivo2l+c=

main.go

@@ -0,0 +1,218 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"flag"
+	"fmt"
+	"net/http"
+	"os"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"nhooyr.io/websocket"
+)
+
+var (
+	Listen  = flag.String("listen", "0.0.0.0:8080", "listen interface:port")
+	Peers   = flag.String("peers", "peers.json", "peer definition file")
+	Verbose = flag.Bool("verbose", false, "verbose logging")
+)
+
+type Pair struct {
+	P string `json:"p"`
+	U string `json:"u"`
+}
+
+type Intro struct {
+	Location int
+	Conn     *websocket.Conn
+	Addr     string
+}
+
+type Msg struct {
+	Location int
+	Bytes    []byte
+	Type     websocket.MessageType
+}
+
+type Handler struct {
+	pairs     []Pair
+	locations map[string]int
+	conns     []*websocket.Conn
+	addrs     []string
+	closeReqs chan int
+	intros    chan Intro
+	messages  chan Msg
+}
+
+func NewHandler() *Handler {
+	h := &Handler{}
+	h.init()
+	return h
+}
+
+func (h *Handler) init() {
+	h.closeReqs = make(chan int)
+	h.intros = make(chan Intro)
+	h.messages = make(chan Msg)
+	go h.monitorChannels()
+}
+
+func (h *Handler) peer(location int) int {
+	if location%2 == 0 {
+		return location + 1
+	} else {
+		return location - 1
+	}
+}
+
+func (h *Handler) monitorChannels() {
+	for {
+		select {
+		case intro := <-h.intros:
+			old := h.conns[intro.Location]
+			if old != nil {
+				log.WithField("old", h.addrs[intro.Location]).
+					WithField("new", intro.Addr).
+					WithField("uuid", h.uuidOf(intro.Location)).Info("duplicate user")
+				intro.Conn.Close(websocket.StatusNormalClosure, "duplicate user")
+				break
+			}
+			h.conns[intro.Location] = intro.Conn
+			h.addrs[intro.Location] = intro.Addr
+			log.WithField("client", intro.Addr).
+				WithField("uuid", h.uuidOf(intro.Location)).
+				Debug("drained")
+			go h.relay(intro.Conn, intro.Location)
+		case msg := <-h.messages:
+			peerLocation := h.peer(msg.Location)
+			peer := h.conns[peerLocation]
+			log.WithField("client", h.addrs[msg.Location]).
+				WithField("uuid", h.uuidOf(msg.Location)).
+				WithField("message", string(msg.Bytes)).Debug("received")
+			if peer != nil {
+				log.WithField("client", h.addrs[peerLocation]).
+					WithField("uuid", h.uuidOf(peerLocation)).
+					Debug("sent")
+				err := peer.Write(context.Background(),
+					msg.Type, msg.Bytes)
+				if err != nil {
+					err := peer.Close(websocket.StatusNormalClosure, "write error")
+					if err != nil {
+						log.Warning(err)
+					}
+					h.conns[peerLocation] = nil
+					h.addrs[peerLocation] = ""
+				}
+			}
+		case closeLoc := <-h.closeReqs:
+			err := h.conns[closeLoc].
+				Close(websocket.StatusNormalClosure, "")
+			if err != nil {
+				log.Warning(err)
+			}
+
+			h.conns[closeLoc] = nil
+			h.addrs[closeLoc] = ""
+			log.Trace(h.conns)
+		}
+	}
+}
+
+func (h *Handler) uuidOf(location int) string {
+	if location%2 == 0 {
+		return h.pairs[location/2].P
+	} else {
+		return h.pairs[location/2].U
+	}
+}
+
+func (h *Handler) relay(from *websocket.Conn, location int) {
+	for {
+		msgType, msgBytes, err := from.Read(context.Background())
+		if err != nil {
+			log.Debug("sending close message ", location)
+			h.closeReqs <- location
+			log.Debug("returning from Relay")
+			return
+		}
+		h.messages <- Msg{
+			Location: location,
+			Type:     msgType,
+			Bytes:    msgBytes,
+		}
+	}
+}
+
+func (h *Handler) LoadPeers(filename string) error {
+	file, err := os.Open(filename)
+	if err != nil {
+		return err
+	}
+	decoder := json.NewDecoder(file)
+	err = decoder.Decode(&h.pairs)
+	if err != nil {
+		return err
+	}
+	h.locations = map[string]int{}
+	for i, pair := range h.pairs {
+		if _, ok := h.locations[pair.P]; ok {
+			return fmt.Errorf("duplicate uuid %v", pair.P)
+		}
+		if _, ok := h.locations[pair.U]; ok {
+			return fmt.Errorf("duplicate uuid %v", pair.U)
+		}
+		h.locations[pair.P] = i * 2
+		h.locations[pair.U] = i*2 + 1
+	}
+	h.conns = make([]*websocket.Conn, len(h.pairs)*2)
+	h.addrs = make([]string, len(h.pairs)*2)
+	return nil
+}
+
+func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	wsConn, err := websocket.Accept(w, r, nil)
+	if err != nil {
+		log.WithField("client", r.RemoteAddr).Error(err)
+		return
+	}
+	go h.ServeWebSocket(wsConn, r.RemoteAddr)
+}
+
+func (h *Handler) ServeWebSocket(wsConn *websocket.Conn, rAddr string) {
+	authLog := log.WithField("client", rAddr)
+	authLog.Trace("waiting for auth")
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	_, msgBytes, err := wsConn.Read(ctx)
+	if err != nil {
+		authLog.Error(err)
+		wsConn.Close(websocket.StatusNormalClosure, "")
+		return
+	}
+	uuid := string(msgBytes)
+	loc, ok := h.locations[uuid]
+	if !ok {
+		authLog.WithField("uuid", uuid).Error("unknown user")
+	} else {
+		authLog.WithField("uuid", uuid).Info("logged in")
+	}
+	h.intros <- Intro{loc, wsConn, rAddr}
+}
+
+func main() {
+	flag.Parse()
+	if *Verbose {
+		log.SetLevel(log.TraceLevel)
+		log.Trace("verbose logging")
+	}
+	h := NewHandler()
+	log.Info("Parsing peers from ", *Peers)
+	err := h.LoadPeers(*Peers)
+	if err != nil {
+		log.Fatal(err)
+	}
+	log.Info("Starting relay on ", *Listen)
+	log.Fatal(http.ListenAndServe(*Listen, h))
+}

peers.json

@@ -0,0 +1,18 @@
+[
+  {
+    "p": "1e6479a4-5060-414a-a03e-bc083b4d93cb",
+    "u": "ec6d6805-2dc8-4002-8f9d-e82015231a10"
+  },
+  {
+    "p": "7cebc48a-ca2b-489b-9c99-2f26780baf6f",
+    "u": "01120c32-29e2-41e2-9994-26881ca90f59"
+  },
+  {
+    "p": "e6508779-e20c-4426-9171-c9236dd1c32f",
+    "u": "709e2184-e728-4711-8d66-b918aae624e1"
+  },
+  {
+    "p": "f85c13b8-6914-497b-8b79-d0d3b2125e3a",
+    "u": "3f51940e-7b63-412f-b9c7-d0aadeeb913d"
+  }
+]