From 59cd69214140e4aa5356b6fc8480ba0e3cad1d5d Mon Sep 17 00:00:00 2001
From: Stefan Prodan <stefan.prodan@gmail.com>
Date: Tue, 11 Sep 2018 22:13:54 +0300
Subject: [PATCH] Add websocket echo handler

---
 README.md         |  3 +-
 pkg/api/echows.go | 82 +++++++++++++++++++++++++++++++++++++++++++++++
 pkg/api/server.go |  1 +
 3 files changed, 85 insertions(+), 1 deletion(-)
 create mode 100644 pkg/api/echows.go

diff --git a/README.md b/README.md
index 2389c91..566be0c 100644
--- a/README.md
+++ b/README.md
@@ -33,10 +33,11 @@ Web API:
 * `GET /headers` returns a JSON with the request HTTP headers
 * `GET /delay/{seconds}` waits for the specified period
 * `POST /token` issues a JWT token valid for one minute `JWT=$(curl -sd 'anon' podinfo:9898/token | jq -r .token)`
-* `GET /token/validate` validates the JWT token `curl -H "Authorization: Bearer ${JWT}" podinfo:9898/token/validate`
+* `GET /token/validate` validates the JWT token `curl -H "Authorization: Bearer $JWT" podinfo:9898/token/validate`
 * `GET /configs` returns a JSON with configmaps and/or secrets mounted in the `config` volume
 * `POST /write` writes the posted content to disk at /data/hash and returns the SHA1 hash of the content
 * `GET /read/{hash}` returns the content of the file /data/hash if exists
+* `GET /ws/echo` echos content via websockets `podcli ws ws://localhost:9898/ws/echo`
 
 ### Guides
 
diff --git a/pkg/api/echows.go b/pkg/api/echows.go
new file mode 100644
index 0000000..7302c0e
--- /dev/null
+++ b/pkg/api/echows.go
@@ -0,0 +1,82 @@
+package api
+
+import (
+	"net/http"
+	"strings"
+	"time"
+
+	"github.com/gorilla/websocket"
+	"go.uber.org/zap"
+)
+
+var wsCon = websocket.Upgrader{}
+
+// Test: go run ./cmd/podcli/* ws localhost:9898/ws/echo
+func (s *Server) echoWsHandler(w http.ResponseWriter, r *http.Request) {
+	c, err := wsCon.Upgrade(w, r, nil)
+	if err != nil {
+		if err != nil {
+			s.logger.Warn("websocket upgrade error", zap.Error(err))
+			return
+		}
+	}
+	defer c.Close()
+	done := make(chan struct{})
+	defer close(done)
+	in := make(chan interface{})
+	defer close(in)
+	go s.writeWs(c, in)
+	go s.sendHostWs(c, in, done)
+	for {
+		_, message, err := c.ReadMessage()
+		if err != nil {
+			if !strings.Contains(err.Error(), "close") {
+				s.logger.Warn("websocket read error", zap.Error(err))
+			}
+			break
+		}
+		var response = struct {
+			Time    time.Time `json:"ts"`
+			Message string    `json:"msg"`
+		}{
+			Time:    time.Now(),
+			Message: string(message),
+		}
+		in <- response
+	}
+}
+
+func (s *Server) sendHostWs(ws *websocket.Conn, in chan interface{}, done chan struct{}) {
+	ticker := time.NewTicker(5 * time.Second)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
+			var status = struct {
+				Time time.Time `json:"ts"`
+				Host string    `json:"server"`
+			}{
+				Time: time.Now(),
+				Host: s.config.Hostname,
+			}
+			in <- status
+		case <-done:
+			s.logger.Debug("websocket exit")
+			return
+		}
+	}
+}
+
+func (s *Server) writeWs(ws *websocket.Conn, in chan interface{}) {
+	for {
+		select {
+		case msg := <-in:
+			if err := ws.WriteJSON(msg); err != nil {
+				if !strings.Contains(err.Error(), "close") {
+					s.logger.Warn("websocket write error", zap.Error(err))
+				}
+				return
+			}
+		}
+	}
+}
diff --git a/pkg/api/server.go b/pkg/api/server.go
index b722eb9..554c593 100644
--- a/pkg/api/server.go
+++ b/pkg/api/server.go
@@ -78,6 +78,7 @@ func (s *Server) registerHandlers() {
 	s.router.HandleFunc("/token/validate", s.tokenValidateHandler).Methods("GET")
 	s.router.HandleFunc("/api/info", s.infoHandler).Methods("GET")
 	s.router.HandleFunc("/api/echo", s.echoHandler).Methods("POST")
+	s.router.HandleFunc("/ws/echo", s.echoWsHandler)
 }
 
 func (s *Server) registerMiddlewares() {