[agent] Cleanup unused, shimmed websocket connections after a configurable timeout.

This PR adds a configurable value for how long shimmed websocket connections can remain idle before they are garbage collected.

Previously, the underlying connection from the proxy agent to the backend server was only cleaned up when there was either an error while sending or receiving messages, or if the agent received an explicit close message.

When a client of shim protocol stopped communicating without first sending a close message, this caused the underlying connection between the proxy agent and backend server to be retained indefinitely. That, in turn, can result in a resource leak within the agent (and possibly in the backend server as well).

To account for that scenario, we needed to find some way to identify that a client of the shim protocol had stopped communicating, so that we could clean up those lost connections.

The shim protocol does not include any sort of ping or pong messages, but it does have polling requests sent by the client which can be used as a dead-man switch; once a sufficient amount of time has elapsed since the last such request, we can assume that the client is gone and the underlying connection can be removed.

This PR adds a configuration option to the agent to specify what that amount of time should be (with a default of 1 hour), and implements the agent-side cleanup logic to identify and close expired connections.

Author: ojarjur

agent/agent.go

@@ -73,6 +73,7 @@ var (
 	injectBanner              = flag.String("inject-banner", "", "HTML snippet to inject in served webpages")
 	bannerHeight              = flag.String("banner-height", "40px", "Height of the injected banner. This is ignored if no banner is set.")
 	shimWebsockets            = flag.Bool("shim-websockets", false, "Whether or not to replace websockets with a shim")
+	websocketShimTimeout      = flag.Duration("websocket-shim-timeout", 60*time.Minute, "Timeout for websocket shim connections to expire due to inactivity.")
 	shimPath                  = flag.String("shim-path", "", "Path under which to handle websocket shim requests")
 	healthCheckPath           = flag.String("health-check-path", "/", "Path on backend host to issue health checks against.  Defaults to the root.")
 	healthCheckFreq           = flag.Int("health-check-interval-seconds", 0, "Wait time in seconds between health checks.  Set to zero to disable health checks.  Checks disabled by default.")
@@ -126,7 +127,8 @@ func hostProxy(ctx context.Context, host, shimPath string, injectShimCode, force
 		// restricted to a path prefix not equal to "/" will fail for websocket open requests. Passing in the
 		// sessionHandler twice allows the websocket handler to ensure that cookies are applied based on the
 		// correct, restored path.
-		h, err = websockets.Proxy(ctx, h, host, shimPath, *rewriteWebsocketHost, *enableWebsocketsInjection, sessionLRU.SessionHandler, metricHandler)
+		h, err = websockets.Proxy(ctx, h, host, shimPath, *rewriteWebsocketHost, *enableWebsocketsInjection, sessionLRU.SessionHandler,
+			metricHandler, *websocketShimTimeout)
 		if injectShimCode {
 			shimFunc, err := websockets.ShimBody(shimPath)
 			if err != nil {

agent/websockets/connection.go

@@ -17,16 +17,16 @@ limitations under the License.
 package websockets
 
 import (
+	"context"
 	"encoding/base64"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"log"
 	"net/http"
+	"sync"
 	"time"
 
-	"context"
-
 	"github.com/gorilla/websocket"
 )
 
@@ -57,12 +57,14 @@ func (m *message) Serialize(version int) interface{} {
 // and encapsulates it in an API that is a little more amenable to how the server side
 // of our websocket shim is implemented.
 type Connection struct {
-	done            func() <-chan struct{}
-	cancel          context.CancelFunc
-	clientMessages  chan *message
-	serverMessages  chan *message
-	protocolVersion int
-	subprotocol     string
+	done             func() <-chan struct{}
+	cancel           context.CancelFunc
+	clientMessages   chan *message
+	serverMessages   chan *message
+	protocolVersion  int
+	subprotocol      string
+	mu               sync.Mutex
+	lastActivityTime time.Time
 }
 
 // This map defines the set of headers that should be stripped from the WS request, as they
@@ -87,6 +89,20 @@ func stripWSHeader(header http.Header) http.Header {
 	return result
 }
 
+// updateActivity updates the last activity timestamp.
+func (conn *Connection) updateActivity() {
+	conn.mu.Lock()
+	defer conn.mu.Unlock()
+	conn.lastActivityTime = time.Now()
+}
+
+// lastActivity returns the last activity timestamp.
+func (conn *Connection) lastActivity() time.Time {
+	conn.mu.Lock()
+	defer conn.mu.Unlock()
+	return conn.lastActivityTime
+}
+
 // NewConnection creates and returns a new Connection.
 func NewConnection(ctx context.Context, targetURL string, header http.Header, errCallback func(err error)) (*Connection, error) {
 	ctx, cancel := context.WithCancel(ctx)
@@ -162,11 +178,12 @@ func NewConnection(ctx context.Context, targetURL string, header http.Header, er
 		}
 	}()
 	return &Connection{
-		done:           ctx.Done,
-		cancel:         cancel,
-		clientMessages: clientMessages,
-		serverMessages: serverMessages,
-		subprotocol: serverConn.Subprotocol(),
+		done:             ctx.Done,
+		cancel:           cancel,
+		clientMessages:   clientMessages,
+		serverMessages:   serverMessages,
+		subprotocol:      serverConn.Subprotocol(),
+		lastActivityTime: time.Now(),
 	}, nil
 }
 
@@ -184,6 +201,7 @@ func (conn *Connection) Close() {
 //
 // The returned error value is non-nill if the connection has been closed.
 func (conn *Connection) SendClientMessage(msg interface{}, injectionEnabled bool, injectedHeaders map[string]string) error {
+	conn.updateActivity()
 	var clientMessage *message
 	if textMsg, ok := msg.(string); ok {
 		clientMessage = &message{
@@ -236,7 +254,9 @@ func (conn *Connection) SendClientMessage(msg interface{}, injectionEnabled bool
 //
 // The returned []string value is nil if the error is non-nil, or if the method
 // times out while waiting for a server message.
-func (conn *Connection) ReadServerMessages() ([]interface{}, error) {
+func (conn *Connection) ReadServerMessages(readTimeout time.Duration) ([]interface{}, error) {
+	conn.updateActivity()
+	defer conn.updateActivity()
 	var msgs []interface{}
 	select {
 	case serverMsg, ok := <-conn.serverMessages:
@@ -257,7 +277,7 @@ func (conn *Connection) ReadServerMessages() ([]interface{}, error) {
 				return msgs, nil
 			}
 		}
-	case <-time.After(time.Second * 20):
+	case <-time.After(readTimeout):
 		return nil, nil
 	}
 }

agent/websockets/shim.go

@@ -18,6 +18,7 @@ package websockets
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -31,8 +32,8 @@ import (
 	"sync"
 	"sync/atomic"
 	"text/template"
+	"time"
 
-	"context"
 	"github.com/google/inverting-proxy/agent/metrics"
 )
 
@@ -320,9 +321,33 @@ func (c *connectionErrorHandler) ReportError(err error) {
 	}
 }
 
-func createShimChannel(ctx context.Context, host, shimPath string, rewriteHost bool, openWebsocketWrapper func(http.Handler, *metrics.MetricHandler) http.Handler, enableWebsocketInjection bool, metricHandler *metrics.MetricHandler) http.Handler {
+func createShimChannel(ctx context.Context, host, shimPath string, rewriteHost bool, openWebsocketWrapper func(http.Handler, *metrics.MetricHandler) http.Handler, enableWebsocketInjection bool, metricHandler *metrics.MetricHandler, timeout time.Duration) http.Handler {
 	var connections sync.Map
 	var sessionCount uint64
+
+	// Background goroutine to clean up inactive websocket shim connections.
+	go func() {
+		ticker := time.NewTicker(min(timeout, 30*time.Second))
+		defer ticker.Stop()
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+				connections.Range(func(key, value any) bool {
+					sessionID := key.(string)
+					conn := value.(*Connection)
+					if time.Since(conn.lastActivity()) > timeout {
+						log.Printf("Closing inactive websocket shim session %q after timeout", sessionID)
+						conn.Close()
+						connections.Delete(sessionID)
+					}
+					return true // Continue iteration
+				})
+			}
+		}
+	}()
+
 	mux := http.NewServeMux()
 	errorHandler := &connectionErrorHandler{}
 	openWebsocketHandler := openWebsocketWrapper(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -351,9 +376,9 @@ func createShimChannel(ctx context.Context, host, shimPath string, rewriteHost b
 			}
 		}
 		resp := &sessionMessage{
-			ID:      sessionID,
-			Message: targetURL.String(),
-			Version: conn.protocolVersion,
+			ID:          sessionID,
+			Message:     targetURL.String(),
+			Version:     conn.protocolVersion,
 			Subprotocol: conn.Subprotocol(),
 		}
 		respBytes, err := json.Marshal(resp)
@@ -512,7 +537,7 @@ func createShimChannel(ctx context.Context, host, shimPath string, rewriteHost b
 			metricHandler.WriteResponseCodeMetric(statusCode)
 			return
 		}
-		serverMsgs, err := conn.ReadServerMessages()
+		serverMsgs, err := conn.ReadServerMessages(min(20*time.Second, timeout/2))
 		if err != nil {
 			statusCode := http.StatusBadRequest
 			errorMessage := fmt.Sprintf("attempt to read data from a closed session: %q", msg.ID)
@@ -548,11 +573,11 @@ func createShimChannel(ctx context.Context, host, shimPath string, rewriteHost b
 // openWebsocketWrapper is a http.Handler wrapper function that is invoked on websocket open requests after the original
 // targetURL of the request is restored. It must call the wrapped http.Handler with which it is created after it
 // is finished processing the request.
-func Proxy(ctx context.Context, wrapped http.Handler, host, shimPath string, rewriteHost, enableWebsocketInjection bool, openWebsocketWrapper func(wrapped http.Handler, metricHandler *metrics.MetricHandler) http.Handler, metricHandler *metrics.MetricHandler) (http.Handler, error) {
+func Proxy(ctx context.Context, wrapped http.Handler, host, shimPath string, rewriteHost, enableWebsocketInjection bool, openWebsocketWrapper func(wrapped http.Handler, metricHandler *metrics.MetricHandler) http.Handler, metricHandler *metrics.MetricHandler, timeout time.Duration) (http.Handler, error) {
 	mux := http.NewServeMux()
 	if shimPath != "" {
 		shimPath = path.Clean("/"+shimPath) + "/"
-		shimServer := createShimChannel(ctx, host, shimPath, rewriteHost, openWebsocketWrapper, enableWebsocketInjection, metricHandler)
+		shimServer := createShimChannel(ctx, host, shimPath, rewriteHost, openWebsocketWrapper, enableWebsocketInjection, metricHandler, timeout)
 		mux.Handle(shimPath, shimServer)
 	}
 	mux.Handle("/", wrapped)

agent/websockets/websockets_test.go

@@ -23,13 +23,15 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
 	"path"
 	"strings"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
@@ -239,7 +241,7 @@ func TestShimHandlers(t *testing.T) {
 	openWrapper := func(h http.Handler, metricHandler *metrics.MetricHandler) http.Handler {
 		return h
 	}
-	p, err := Proxy(context.Background(), h, serverURL.Host, testShimPath, false, false, openWrapper, nil)
+	p, err := Proxy(context.Background(), h, serverURL.Host, testShimPath, false, false, openWrapper, nil, 60*time.Second)
 	if err != nil {
 		t.Fatalf("Failure creating the websocket shim proxy: %+v", err)
 	}
@@ -354,3 +356,107 @@ func TestShimHandlers(t *testing.T) {
 		}
 	}
 }
+
+func TestShimPolling(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	// Setup a fake backend that accepts websocket connections but sends no messages.
+	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		upgrader := websocket.Upgrader{}
+		conn, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			t.Logf("Failed to upgrade websocket: %v", err)
+			return
+		}
+		defer conn.Close()
+		// Keep connection open until client closes it.
+		for {
+			if _, _, err := conn.NextReader(); err != nil {
+				break
+			}
+		}
+	}))
+	defer backend.Close()
+
+	backendURL, err := url.Parse(backend.URL)
+	if err != nil {
+		t.Fatalf("Failed to parse backend URL: %v", err)
+	}
+
+	shimPath := "/shim/"
+	idleTimeout := 3 * time.Second
+	shim := createShimChannel(
+		ctx,
+		backendURL.Host,
+		shimPath,
+		false,
+		func(h http.Handler, m *metrics.MetricHandler) http.Handler { return h },
+		false,
+		nil,
+		idleTimeout,
+	)
+	shimServer := httptest.NewServer(shim)
+	defer shimServer.Close()
+
+	// 1. Open a websocket connection via the shim.
+	openURL := shimServer.URL + shimPath + "open"
+	resp, err := http.Post(openURL, "text/plain", strings.NewReader(backendURL.String()))
+	if err != nil {
+		t.Fatalf("Failed to open shim connection: %v", err)
+	}
+	if resp.StatusCode != http.StatusOK {
+		t.Fatalf("Failed to open shim connection, status: %d", resp.StatusCode)
+	}
+	body, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		t.Fatalf("Failed to read open response body: %v", err)
+	}
+	resp.Body.Close()
+	var openResp sessionMessage
+	if err := json.Unmarshal(body, &openResp); err != nil {
+		t.Fatalf("Failed to unmarshal open response: %v", err)
+	}
+	sessionID := openResp.ID
+	if sessionID == "" {
+		t.Fatal("No sessionID in open response")
+	}
+
+	// 2. Poll repeatedly without any messages being sent.
+	pollURL := shimServer.URL + shimPath + "poll"
+	pollReq := fmt.Sprintf(`{"id": %q}`, sessionID)
+	timeout := time.Now().Add(idleTimeout + 1*time.Second)
+	for time.Now().Before(timeout) {
+		resp, err := http.Post(pollURL, "application/json", strings.NewReader(pollReq))
+		if err != nil {
+			t.Fatalf("Failed to poll shim connection: %v", err)
+		}
+		resp.Body.Close()
+		if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusRequestTimeout {
+			t.Fatalf("Unexpected status code during polling: %d", resp.StatusCode)
+		}
+		// Sleep for half of read timeout to simulate polling faster than idle timeout.
+		time.Sleep(500 * time.Millisecond)
+	}
+
+	// 3. After idleTimeout + 1s, one more poll should succeed.
+	resp, err = http.Post(pollURL, "application/json", strings.NewReader(pollReq))
+	if err != nil {
+		t.Fatalf("Failed to poll shim connection: %v", err)
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusRequestTimeout {
+		t.Errorf("Polling after idle timeout got status %d, want %d", resp.StatusCode, http.StatusRequestTimeout)
+	}
+
+	// 4. Close connection.
+	closeURL := shimServer.URL + shimPath + "close"
+	closeReq := fmt.Sprintf(`{"id": %q}`, sessionID)
+	resp, err = http.Post(closeURL, "application/json", strings.NewReader(closeReq))
+	if err != nil {
+		t.Fatalf("Failed to close shim connection: %v", err)
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		t.Errorf("Close shim connection got status %d, want %d", resp.StatusCode, http.StatusOK)
+	}
+}

server/server.go

@@ -220,7 +220,7 @@ func (p *proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	p.Lock()
 	p.requests[id] = pending
 	p.Unlock()
-	defer func(){
+	defer func() {
 		p.Lock()
 		delete(p.requests, id)
 		p.Unlock()

testing/websockets/main.go

@@ -31,6 +31,7 @@ import (
 	"net/http"
 	"net/http/httputil"
 	"net/url"
+	"time"
 
 	"github.com/google/inverting-proxy/agent/metrics"
 	"github.com/google/inverting-proxy/agent/websockets"
@@ -48,6 +49,7 @@ var (
 	monitoringEndpoint       = flag.String("monitoring-endpoint", "staging-monitoring.sandbox.googleapis.com:443", "The endpoint to which to write metrics. Eg: monitoring.googleapis.com corresponds to Cloud Monarch.")
 	monitoringResourceType   = flag.String("monitoring-resource-type", "gce_instance", "The monitoring resource type. Eg: gce_instance")
 	monitoringResourceLabels = flag.String("monitoring-resource-labels", "instance-id=fake-instance-id,instance-zone=us-west1-a", "Comma separated key value pairs for the purpose of monitoring configuration. Eg: 'instance-id=my-instance-id,instance-zone=us-west1-a")
+	websocketShimTimeout     = flag.Duration("websocket-shim-timeout", 60*time.Minute, "Timeout for websocket shim connections to expire due to inactivity.")
 )
 
 func main() {
@@ -69,7 +71,7 @@ func main() {
 	}
 
 	backendProxy := httputil.NewSingleHostReverseProxy(backendURL)
-	shimmingProxy, err := websockets.Proxy(context.Background(), backendProxy, backendURL.Host, *shimPath, true, *enableWebsocketInjection, func(h http.Handler, metricHandler *metrics.MetricHandler) http.Handler { return h }, metricHandler)
+	shimmingProxy, err := websockets.Proxy(context.Background(), backendProxy, backendURL.Host, *shimPath, true, *enableWebsocketInjection, func(h http.Handler, metricHandler *metrics.MetricHandler) http.Handler { return h }, metricHandler, *websocketShimTimeout)
 	if err != nil {
 		log.Fatalf("Failure starting the websocket-shimming proxy: %v", err)
 	}