transport/websocket: add heartbeat parameter for connection keep-alive

Signed-off-by: hr567 <[email protected]>

infra/conf/transport_internet.go

@@ -149,6 +149,7 @@ type WebSocketConfig struct {
 	Path                string            `json:"path"`
 	Headers             map[string]string `json:"headers"`
 	AcceptProxyProtocol bool              `json:"acceptProxyProtocol"`
+	Heartbeat           uint32            `json:"heartbeat"`
 }
 
 // Build implements Buildable.
@@ -178,6 +179,7 @@ func (c *WebSocketConfig) Build() (proto.Message, error) {
 		Header:              c.Headers,
 		AcceptProxyProtocol: c.AcceptProxyProtocol,
 		Ed:                  ed,
+		Heartbeat:           c.Heartbeat,
 	}
 	return config, nil
 }

transport/internet/splithttp/browser_client.go

@@ -25,7 +25,7 @@ func (c *BrowserDialerClient) OpenDownload(ctx context.Context, baseURL string)
 		return nil, dummyAddr, dummyAddr, err
 	}
 
-	return websocket.NewConnection(conn, dummyAddr, nil), conn.RemoteAddr(), conn.LocalAddr(), nil
+	return websocket.NewConnection(conn, dummyAddr, nil, 0), conn.RemoteAddr(), conn.LocalAddr(), nil
 }
 
 func (c *BrowserDialerClient) SendUploadRequest(ctx context.Context, url string, payload io.ReadWriteCloser, contentLength int64) error {

transport/internet/websocket/config.proto

@@ -12,4 +12,5 @@ message Config {
   map<string, string> header = 3;
   bool accept_proxy_protocol = 4;
   uint32 ed = 5;
+  uint32 heartbeat = 6;
 }

transport/internet/websocket/connection.go

@@ -22,7 +22,18 @@ type connection struct {
 	remoteAddr net.Addr
 }
 
-func NewConnection(conn *websocket.Conn, remoteAddr net.Addr, extraReader io.Reader) *connection {
+func NewConnection(conn *websocket.Conn, remoteAddr net.Addr, extraReader io.Reader, heartbeat uint32) *connection {
+	if heartbeat != 0 {
+		go func() {
+			for {
+				time.Sleep(time.Duration(heartbeat) * time.Second)
+				if err := conn.WriteControl(websocket.PingMessage, []byte(""), time.Time{}); err != nil {
+					break
+				}
+			}
+		}()
+	}
+
 	return &connection{
 		conn:       conn,
 		remoteAddr: remoteAddr,

transport/internet/websocket/dialer.go

@@ -99,7 +99,7 @@ func dialWebSocket(ctx context.Context, dest net.Destination, streamSettings *in
 			return nil, err
 		}
 
-		return NewConnection(conn, conn.RemoteAddr(), nil), nil
+		return NewConnection(conn, conn.RemoteAddr(), nil, wsSettings.Heartbeat), nil
 	}
 
 	header := wsSettings.GetRequestHeader()
@@ -117,7 +117,7 @@ func dialWebSocket(ctx context.Context, dest net.Destination, streamSettings *in
 		return nil, errors.New("failed to dial to (", uri, "): ", reason).Base(err)
 	}
 
-	return NewConnection(conn, conn.RemoteAddr(), nil), nil
+	return NewConnection(conn, conn.RemoteAddr(), nil, wsSettings.Heartbeat), nil
 }
 
 type delayDialConn struct {

transport/internet/websocket/hub.go

@@ -73,7 +73,7 @@ func (h *requestHandler) ServeHTTP(writer http.ResponseWriter, request *http.Req
 		}
 	}
 
-	h.ln.addConn(NewConnection(conn, remoteAddr, extraReader))
+	h.ln.addConn(NewConnection(conn, remoteAddr, extraReader, h.ln.config.Heartbeat))
 }
 
 type Listener struct {