added websocket ping/pong mechanism (as described in rfc6455)

src/api/api.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"net/http"
 	"strconv"
+	"time"
 
 	"github.com/gabriel-vasile/mimetype"
 	"github.com/gin-gonic/gin"
@@ -15,6 +16,17 @@ import (
 	utils "github.com/bbernhard/signal-cli-rest-api/utils"
 )
 
+const (
+	// Time allowed to write the file to the client.
+	writeWait = 10 * time.Second
+
+	// Time allowed to read the next pong message from the client.
+	pongWait = 60 * time.Second
+
+	// Send pings to client with this period. Must be less than pongWait.
+	pingPeriod = (pongWait * 9) / 10
+)
+
 type GroupPermissions struct {
 	AddMembers string `json:"add_members" enums:"only-admins,every-member"`
 	EditGroup string `json:"edit_group" enums:"only-admins,every-member"`
@@ -267,26 +279,8 @@ func (a *Api) SendV2(c *gin.Context) {
 	c.JSON(201, SendMessageResponse{Timestamp: strconv.FormatInt((*timestamps)[0].Timestamp, 10)})
 }
 
-// @Summary Receive Signal Messages.
-// @Tags Messages
-// @Description Receives Signal Messages from the Signal Network.
-// @Accept  json
-// @Produce  json
-// @Success 200 {object} []string
-// @Failure 400 {object} Error
-// @Param number path string true "Registered Phone Number"
-// @Param timeout query string false "Receive timeout in seconds (default: 1)"
-// @Router /v1/receive/{number} [get]
-func (a *Api) Receive(c *gin.Context) {
-	number := c.Param("number")
 
-	if a.signalClient.GetSignalCliMode() == client.JsonRpc {
+func (a *Api) handleSignalReceive(ws *websocket.Conn, number string) {
-		ws, err := connectionUpgrader.Upgrade(c.Writer, c.Request, nil)
-		if err != nil {
-			c.JSON(400, Error{Msg: err.Error()})
-			return
-		}
-		defer ws.Close()
 	for {
 		data, err := a.signalClient.Receive(number, 0)
 		if err == nil {
@@ -309,6 +303,56 @@ func (a *Api) Receive(c *gin.Context) {
 			}
 		}
 	}
+}
+
+func wsPong(ws *websocket.Conn) {
+	ws.SetReadLimit(512)
+	ws.SetReadDeadline(time.Now().Add(pongWait))
+	ws.SetPongHandler(func(string) error { ws.SetReadDeadline(time.Now().Add(pongWait)); return nil })
+	for {
+		_, _, err := ws.ReadMessage()
+		if err != nil {
+			break
+		}
+	}
+}
+
+func wsPing(ws *websocket.Conn) {
+	pingTicker := time.NewTicker(pingPeriod)
+	for {
+		select {
+		case <-pingTicker.C:
+			ws.SetWriteDeadline(time.Now().Add(writeWait))
+			if err := ws.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
+				return
+			}
+		}
+	}
+}
+
+// @Summary Receive Signal Messages.
+// @Tags Messages
+// @Description Receives Signal Messages from the Signal Network.
+// @Accept  json
+// @Produce  json
+// @Success 200 {object} []string
+// @Failure 400 {object} Error
+// @Param number path string true "Registered Phone Number"
+// @Param timeout query string false "Receive timeout in seconds (default: 1)"
+// @Router /v1/receive/{number} [get]
+func (a *Api) Receive(c *gin.Context) {
+	number := c.Param("number")
+
+	if a.signalClient.GetSignalCliMode() == client.JsonRpc {
+		ws, err := connectionUpgrader.Upgrade(c.Writer, c.Request, nil)
+		if err != nil {
+			c.JSON(400, Error{Msg: err.Error()})
+			return
+		}
+		defer ws.Close()
+		go a.handleSignalReceive(ws, number)
+		go wsPing(ws)
+		wsPong(ws)
 	} else {
 		timeout := c.DefaultQuery("timeout", "1")
 		timeoutInt, err := strconv.ParseInt(timeout, 10, 32)