fix leaking connections when user client closes connection (#222)

Same as #220, but ported to this repo.

---

When the user client closes the connection, it is not propagated to the
backend. For simple "one shot" requests, there is no huge impact as the
response would just be dropped, however for long running requests
(chunked encoding or websockets), the backend doesn't know that the user
client has closed the connection and it keeps sending new updates to the
relay client, which also doesn't know and so keep forwarding that to the
relay server.

The fix applied is:

* Add `StopRelayRequest` to the broker to "forget" a relaying request.
This will cause the relay server to response with a permanent error the
next time the relay client tries to send a response for that id. This
will in turn cause the relay client to close the connection to the
backend.
* There is a bug in the relay client when checking for backoff permanent
errors. When the backoff operation encounters a permanent error, it
actually unwraps it and return the underlying error, but the client is
still checking for `backoff.Permanent`, resulting in the check to always
fail and never handling it.

---------

Signed-off-by: Teo Koon Peng <[email protected]>

Author: koonpeng

src/bootstrap/cloud/terraform/service-account.tf

@@ -60,7 +60,7 @@ resource "google_project_iam_member" "robot-service-roles" {
   project = data.google_project.project.project_id
   member  = "serviceAccount:${google_service_account.robot-service.email}"
   for_each = toset([
-    "roles/cloudtrace.agent", # Upload cloud traces
+    "roles/cloudtrace.agent",  # Upload cloud traces
     "roles/logging.logWriter", # Upload text logs to Cloud logging
     # Required to use robot-service@ for GKE clusters that simulate robots
     "roles/monitoring.viewer",

src/go/cmd/http-relay-client/client/client.go

@@ -91,6 +91,18 @@ type ClientConfig struct {
 	ForceHttp2   bool
 }
 
+type RelayServerError struct {
+	msg string
+}
+
+func NewRelayServerError(msg string) error {
+	return &RelayServerError{msg}
+}
+
+func (e *RelayServerError) Error() string {
+	return e.msg
+}
+
 func DefaultClientConfig() ClientConfig {
 	return ClientConfig{
 		RemoteRequestTimeout:   60 * time.Second,
@@ -384,7 +396,7 @@ func (c *Client) postResponse(remote *http.Client, br *pb.HttpResponse) error {
 		return fmt.Errorf("couldn't read relay server's response body: %v", err)
 	}
 	if resp.StatusCode != http.StatusOK {
-		err := fmt.Errorf("relay server responded %s: %s", http.StatusText(resp.StatusCode), body)
+		err := NewRelayServerError(fmt.Sprintf("relay server responded %s or the client cancelled the request: %s", http.StatusText(resp.StatusCode), body))
 		if resp.StatusCode == http.StatusBadRequest {
 			// http-relay-server may have restarted during the request.
 			return backoff.Permanent(err)
@@ -564,6 +576,7 @@ func (c *Client) handleRequest(remote *http.Client, local *http.Client, pbreq *p
 	defer span.End()
 
 	resp, hresp, err := makeBackendRequest(ctx, local, req, id)
+	defer hresp.Body.Close()
 	if err != nil {
 		// Even if we couldn't handle the backend request, send an
 		// answer to the relay that signals the error.
@@ -643,8 +656,11 @@ func (c *Client) handleRequest(remote *http.Client, local *http.Client, pbreq *p
 				log.Printf("[%s] Failed to post response to relay: %v", *resp.Id, err)
 			},
 		)
-		if _, ok := err.(*backoff.PermanentError); ok {
-			// A permanent error suggests the request should be aborted.
+		// Any error suggests the request should be aborted.
+		// A missing chunk will cause clients to receive corrupted data, in most cases it is better
+		// to close the connection to avoid that.
+		if err != nil {
+			log.Printf("[%s] Closing backend connection (%s)", *resp.Id, err)
 			break
 		}
 	}

src/go/cmd/http-relay-server/server/broker.go

@@ -171,6 +171,14 @@ func (r *broker) RelayRequest(server string, request *pb.HttpRequest) (<-chan *p
 	}
 }
 
+// StopRelayRequest forgets a relaying request, this causes the next chunk from the backend
+// with the relay id to not be recognized, resulting in the relay server returning an error.
+func (r *broker) StopRelayRequest(requestId string) {
+	r.m.Lock()
+	defer r.m.Unlock()
+	delete(r.resp, requestId)
+}
+
 // GetRequest obtains a client's request for the server identifier. It blocks
 // until a client makes a request.
 func (r *broker) GetRequest(ctx context.Context, server, path string) (*pb.HttpRequest, error) {

src/go/cmd/http-relay-server/server/server.go

@@ -257,8 +257,10 @@ func (s *Server) bidirectionalStream(backendCtx backendContext, w http.ResponseW
 
 	numBytes := 0
 	for responseChunk := range responseChunks {
-		// TODO(b/130706300): detect dropped connection and end request in broker
-		_, _ = bufrw.Write(responseChunk.Body)
+		if _, err = w.Write(responseChunk.Body); err != nil {
+			log.Printf("[%s] %s", backendCtx.Id, err)
+			return
+		}
 		bufrw.Flush()
 		numBytes += len(responseChunk.Body)
 	}
@@ -378,6 +380,7 @@ func (s *Server) userClientRequest(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		return
 	}
+	defer s.b.StopRelayRequest(backendCtx.Id)
 
 	header, responseChunksChan, done := s.waitForFirstResponseAndHandleSwitching(ctx, *backendCtx, w, backendRespChan)
 	if done {
@@ -392,8 +395,10 @@ func (s *Server) userClientRequest(w http.ResponseWriter, r *http.Request) {
 	// i.e. this will block until
 	numBytes := 0
 	for responseChunk := range responseChunksChan {
-		// TODO(b/130706300): detect dropped connection and end request in broker
-		_, _ = w.Write(responseChunk.Body)
+		if _, err = w.Write(responseChunk.Body); err != nil {
+			log.Printf("[%s] %s", backendCtx.Id, err)
+			return
+		}
 		if flush, ok := w.(http.Flusher); ok {
 			flush.Flush()
 		}

src/go/tests/relay/nok8s_relay_test.go

@@ -15,6 +15,7 @@
 package main
 
 import (
+	"bufio"
 	"bytes"
 	"context"
 	"fmt"
@@ -128,7 +129,7 @@ func (r *relay) stop() error {
 }
 
 // TestHttpRelay launches a local http relay (client + server) and connects a
-// test-hhtp-server as a backend. The test is then interacting with the backend
+// test-http-server as a backend. The test is then interacting with the backend
 // through the local relay.
 func TestHttpRelay(t *testing.T) {
 	tests := []struct {
@@ -201,6 +202,63 @@ func TestHttpRelay(t *testing.T) {
 	}
 }
 
+// TestDroppedUserClientFreesRelayChannel checks that when the user client closes a connection,
+// it is propagated to the relay server and client, closing the backend connection as well.
+func TestDroppedUserClientFreesRelayChannel(t *testing.T) {
+	// setup http test server
+	connClosed := make(chan error)
+	defer close(connClosed)
+	finishServer := make(chan bool)
+	defer close(finishServer)
+
+	// mock a long running backend that uses chunking to send periodic updates
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		for {
+			select {
+			case <-finishServer:
+				return
+			default:
+				if _, err := fmt.Fprintln(w, "DEADBEEF"); err != nil {
+					connClosed <- err
+					return
+				}
+				if flusher, ok := w.(http.Flusher); ok {
+					flusher.Flush()
+				} else {
+					t.Fatal("cannot flush")
+				}
+				time.Sleep(time.Second)
+			}
+		}
+	}))
+	defer ts.Close()
+
+	backendAddress := strings.TrimPrefix(ts.URL, "http://")
+	r := &relay{}
+	if err := r.start(backendAddress); err != nil {
+		t.Fatal("failed to start relay: ", err)
+	}
+	defer r.stop()
+	relayAddress := "http://127.0.0.1:" + r.rsPort
+
+	res, err := http.Get(relayAddress + "/client/remote1/")
+	if err != nil {
+		t.Fatal(err)
+	}
+	// receive the first chunk then terminates the connection
+	if _, err := bufio.NewReader(res.Body).ReadString('\n'); err != nil {
+		t.Fatal(err)
+	}
+	res.Body.Close()
+
+	// wait for up to 30s for backend connection to be closed
+	select {
+	case <-connClosed:
+	case <-time.After(30 * time.Second):
+		t.Error("Server did not close connection")
+	}
+}
+
 type testServer struct {
 	testpb.UnimplementedTestServiceServer
 	responsePayload []byte