From 8d87d4facdd2c7ec6be0310e1603f9d090d3c8d8 Mon Sep 17 00:00:00 2001
From: Devin Carr <dcarr@cloudflare.com>
Date: Thu, 6 Apr 2023 16:00:19 -0700
Subject: [PATCH] TUN-7351: Add streaming logs session ping and timeout

Sends a ping every 15 seconds to keep the session alive even if no
protocol messages are being propagated. Additionally, sets a hard
timeout of 5 minutes when not actively streaming logs to drop the
connection.
---
 management/service.go | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/management/service.go b/management/service.go
index 0514178a..daf5687f 100644
--- a/management/service.go
+++ b/management/service.go
@@ -5,6 +5,7 @@ import (
 	"net/http"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	"github.com/go-chi/chi/v5"
 	"github.com/rs/zerolog"
@@ -19,6 +20,9 @@ const (
 	// value will return this error to incoming requests.
 	StatusSessionLimitExceeded websocket.StatusCode = 4002
 	reasonSessionLimitExceeded                      = "limit exceeded for streaming sessions"
+
+	StatusIdleLimitExceeded websocket.StatusCode = 4003
+	reasonIdleLimitExceeded                      = "session was idle for too long"
 )
 
 type ManagementService struct {
@@ -147,10 +151,20 @@ func (m *ManagementService) logs(w http.ResponseWriter, r *http.Request) {
 	}
 	// Make sure the connection is closed if other go routines fail to close the connection after completing.
 	defer c.Close(websocket.StatusInternalError, "")
-	ctx := r.Context()
+	ctx, cancel := context.WithCancel(r.Context())
+	defer cancel()
 	events := make(chan *ClientEvent)
 	go m.readEvents(c, ctx, events)
 
+	// Send a heartbeat ping to hold the connection open even if not streaming.
+	ping := time.NewTicker(15 * time.Second)
+	defer ping.Stop()
+
+	// Close the connection if no operation has occurred after the idle timeout.
+	idleTimeout := 5 * time.Minute
+	idle := time.NewTimer(idleTimeout)
+	defer idle.Stop()
+
 	for {
 		select {
 		case <-ctx.Done():
@@ -160,9 +174,11 @@ func (m *ManagementService) logs(w http.ResponseWriter, r *http.Request) {
 		case event := <-events:
 			switch event.Type {
 			case StartStreaming:
+				idle.Stop()
 				m.startStreaming(c, ctx, event)
 				continue
 			case StopStreaming:
+				idle.Reset(idleTimeout)
 				// TODO: limit StopStreaming to only halt streaming for clients that are already streaming
 				m.streaming.Store(false)
 			case UnknownClientEventType:
@@ -176,6 +192,11 @@ func (m *ManagementService) logs(w http.ResponseWriter, r *http.Request) {
 				}
 				return
 			}
+		case <-ping.C:
+			go c.Ping(ctx)
+		case <-idle.C:
+			c.Close(StatusIdleLimitExceeded, reasonIdleLimitExceeded)
+			return
 		}
 	}
 }