From 05a935aec6a9bc008f50206bb198c4220f423252 Mon Sep 17 00:00:00 2001
From: Valentin Fedoskin <valentin.fedoskin@gmail.com>
Date: Sat, 23 May 2020 17:31:34 +0200
Subject: [PATCH] handle websocket connection

---
 go.mod                    |  1 +
 go.sum                    |  2 ++
 pkg/notifier/listen.go    |  1 +
 pkg/notifier/websocket.go | 37 +++++++++++++++++++++++++++++++++++++
 4 files changed, 41 insertions(+)
 create mode 100644 pkg/notifier/websocket.go

diff --git a/go.mod b/go.mod
index ed16050..67b119e 100644
--- a/go.mod
+++ b/go.mod
@@ -3,6 +3,7 @@ module flux-notifier
 go 1.14
 
 require (
+	github.com/gorilla/websocket v1.4.2
 	github.com/prometheus/client_golang v1.6.0
 	github.com/sirupsen/logrus v1.6.0
 	github.com/spf13/cobra v1.0.0
diff --git a/go.sum b/go.sum
index ad1ff60..206aefd 100644
--- a/go.sum
+++ b/go.sum
@@ -57,6 +57,8 @@ github.com/google/go-cmp v0.4.0 h1:xsAVV57WRhGj6kEIi8ReJzQlHHqcBYCElAvkovg3B/4=
 github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
 github.com/gorilla/websocket v1.4.0/go.mod h1:E7qHFY5m1UJ88s3WnNqhKjPHQ0heANvMoAMk2YaljkQ=
+github.com/gorilla/websocket v1.4.2 h1:+/TMaTYc4QFitKJxsQ7Yye35DkWvkdLcvGKqM+x0Ufc=
+github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/grpc-ecosystem/go-grpc-middleware v1.0.0/go.mod h1:FiyG127CGDf3tlThmgyCl78X/SZQqEOJBCDaAfeWzPs=
 github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0/go.mod h1:8NvIoxWQoOIhqOTXgfV/d3M/q6VIi02HzZEHgUlZvzk=
 github.com/grpc-ecosystem/grpc-gateway v1.9.0/go.mod h1:vNeuVxBJEsws4ogUvrchl83t/GYV9WGTSLVdBhOQFDY=
diff --git a/pkg/notifier/listen.go b/pkg/notifier/listen.go
index e83bdb2..afd7d47 100644
--- a/pkg/notifier/listen.go
+++ b/pkg/notifier/listen.go
@@ -31,6 +31,7 @@ func serve(ctx context.Context, port int) (err error) {
 	mux.Handle("/health", http.HandlerFunc(handleHealthRequest))
 	mux.Handle("/metrics", promhttp.Handler())
 	mux.Handle("/v6/events", http.HandlerFunc(handleEventsRequest))
+	mux.Handle("/", http.HandlerFunc(handleWebsocketRequest))
 
 	srv := &http.Server{
 		Addr:    ":" + strconv.FormatInt(int64(port), 10),
diff --git a/pkg/notifier/websocket.go b/pkg/notifier/websocket.go
new file mode 100644
index 0000000..5219209
--- /dev/null
+++ b/pkg/notifier/websocket.go
@@ -0,0 +1,37 @@
+package notifier
+
+import (
+	"github.com/gorilla/websocket"
+	"github.com/sirupsen/logrus"
+	"net/http"
+)
+
+var upgrader = websocket.Upgrader{}
+
+func handleWebsocketRequest(w http.ResponseWriter, r *http.Request) {
+	conn, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		logrus.WithError(err).Error("failed to upgrade connection")
+		return
+	}
+	defer func() {
+		if err := conn.Close(); err != nil {
+			logrus.WithError(err).Error("failed to close connection")
+		}
+	}()
+
+	for {
+		mt, message, err := conn.ReadMessage()
+		if err != nil {
+			logrus.WithError(err).Error("failed to read message")
+			break
+		}
+
+		logrus.WithField("ws-message", string(message)).Info("received ws message")
+
+		if err := conn.WriteMessage(mt, message); err != nil {
+			logrus.WithError(err).Error("failed to write message")
+			break
+		}
+	}
+}