From 0a26c228474a7c3f0b8b8c9cb0d29e2920732624 Mon Sep 17 00:00:00 2001 From: Anton Kovalenko Date: Mon, 13 May 2024 02:11:34 +0300 Subject: [PATCH] Initial commit --- README.md | 3 + go.mod | 10 +++ go.sum | 17 +++++ main.go | 218 +++++++++++++++++++++++++++++++++++++++++++++++++++++ peers.json | 18 +++++ 5 files changed, 266 insertions(+) create mode 100644 README.md create mode 100644 go.mod create mode 100644 go.sum create mode 100644 main.go create mode 100644 peers.json diff --git a/README.md b/README.md new file mode 100644 index 0000000..fa3dc39 --- /dev/null +++ b/README.md @@ -0,0 +1,3 @@ +# Reimplementation of cisoakk's websocket relay example in GO + +https://github.com/cisoakk/wsServer/tree/master/examples/relay diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..ff10f03 --- /dev/null +++ b/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 diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..9356ab7 --- /dev/null +++ b/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= diff --git a/main.go b/main.go new file mode 100644 index 0000000..0c1e7bc --- /dev/null +++ b/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)) +} diff --git a/peers.json b/peers.json new file mode 100644 index 0000000..237dfee --- /dev/null +++ b/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" + } +]