fix websocket disconnecting after one minute of inactivity

Signed-off-by: Jianhui Zhao <[email protected]>

client/client.go

@@ -2,7 +2,7 @@ package client
 
 // Client abstract device and user
 type Client interface {
-	WriteMsg(typ int, data []byte)
+	WriteMsg(typ int, data []byte) error
 
 	// For users, return the device ID that the user wants to access
 	// For devices, return the ID of the device

device.go

@@ -112,12 +112,14 @@ func (dev *device) Close() {
 	}
 }
 
-func (dev *device) WriteMsg(typ int, data []byte) {
+func (dev *device) WriteMsg(typ int, data []byte) error {
 	b := []byte{byte(typ), 0, 0}
 
 	binary.BigEndian.PutUint16(b[1:], uint16(len(data)))
 
-	dev.conn.Write(append(b, data...))
+	_, err := dev.conn.Write(append(b, data...))
+
+	return err
 }
 
 func parseDeviceInfo(b []byte) (string, string, string) {

user.go

@@ -3,6 +3,7 @@ package main
 import (
 	"fmt"
 	"sync"
+	"time"
 
 	"github.com/gin-gonic/gin"
 	"github.com/gorilla/websocket"
@@ -47,8 +48,8 @@ func (u *user) DeviceID() string {
 	return u.devid
 }
 
-func (u *user) WriteMsg(typ int, data []byte) {
-	u.conn.WriteMessage(typ, data)
+func (u *user) WriteMsg(typ int, data []byte) error {
+	return u.conn.WriteMessage(typ, data)
 }
 
 func (u *user) Close() {
@@ -68,6 +69,17 @@ func userLoginAck(code int, c client.Client) {
 	c.WriteMsg(websocket.TextMessage, []byte(msg))
 }
 
+func (u *user) keepAlive() {
+	for {
+		err := u.WriteMsg(websocket.PingMessage, []byte{})
+		if err != nil {
+			return
+		}
+
+		time.Sleep(time.Second * 5)
+	}
+}
+
 func (u *user) readLoop() {
 	defer u.Close()
 
@@ -104,6 +116,7 @@ func serveUser(br *broker, c *gin.Context) {
 		devid: devid,
 	}
 
+	go u.keepAlive()
 	go u.readLoop()
 
 	br.register <- u