From 3babf643bc8fb33fbc414fb708d45677d0b5c9a5 Mon Sep 17 00:00:00 2001
From: Patryk Kalinowski <pkal@horizon.io>
Date: Wed, 4 Sep 2024 15:20:55 +0200
Subject: [PATCH] Raise client timeout to 30s (#69)

* rpc: create a client with a 30 second timeout

* cmd: pass the transport instead of http client to rpc.New
---
 cmd/waas-auth/main.go | 17 +++++++----------
 rpc/auth_test.go      |  8 ++++----
 rpc/helpers_test.go   |  6 +++---
 rpc/rpc.go            |  7 ++++---
 4 files changed, 18 insertions(+), 20 deletions(-)

diff --git a/cmd/waas-auth/main.go b/cmd/waas-auth/main.go
index d8fdaaad..6a39b1ac 100644
--- a/cmd/waas-auth/main.go
+++ b/cmd/waas-auth/main.go
@@ -40,17 +40,14 @@ func main() {
 		}
 	}
 
-	// HTTP client to use for all outgoing connections out of the enclave
-	client := &http.Client{
-		Timeout: 10 * time.Second,
-		Transport: transport.Chain(
-			baseTransport,
-			transport.SetHeader("User-Agent", "waas-authenticator/"+waasauthenticator.VERSION),
-			traceid.Transport,
-		),
-	}
+	// HTTP transport chain to use for all outgoing connections out of the enclave
+	transportChain := transport.Chain(
+		baseTransport,
+		transport.SetHeader("User-Agent", "waas-authenticator/"+waasauthenticator.VERSION),
+		traceid.Transport,
+	)
 
-	s, err := rpc.New(cfg, client)
+	s, err := rpc.New(cfg, transportChain)
 	if err != nil {
 		panic(err)
 	}
diff --git a/rpc/auth_test.go b/rpc/auth_test.go
index d846765c..0437d504 100644
--- a/rpc/auth_test.go
+++ b/rpc/auth_test.go
@@ -341,14 +341,14 @@ func TestStytchAuth(t *testing.T) {
 		stytchServer, tok := issueAccessTokenAndRunStytchJwksServer(t, "project-123", tokBuilderFn)
 		defer stytchServer.Close()
 
-		svc := initRPCWithClient(t, &http.Client{Transport: testTransport{
+		svc := initRPCWithClient(t, testTransport{
 			RoundTripper: http.DefaultTransport,
 			modifyRequest: func(req *http.Request) {
 				if strings.Contains(req.URL.String(), "stytch.com") {
 					req.URL.Host = stytchServer.Listener.Addr().String()
 				}
 			},
-		}})
+		})
 		tenant, _ := newTenantWithAuthConfig(t, svc.Enclave, proto.AuthConfig{
 			Stytch: proto.AuthStytchConfig{
 				Enabled:   true,
@@ -442,14 +442,14 @@ func TestPlayFabAuth(t *testing.T) {
 			playfabAPI := httptest.NewServer(http.HandlerFunc(testCase.playfabHandler))
 			defer playfabAPI.Close()
 
-			svc := initRPCWithClient(t, &http.Client{Transport: testTransport{
+			svc := initRPCWithClient(t, testTransport{
 				RoundTripper: http.DefaultTransport,
 				modifyRequest: func(req *http.Request) {
 					if strings.Contains(req.URL.String(), "playfabapi.com") {
 						req.URL.Host = playfabAPI.Listener.Addr().String()
 					}
 				},
-			}})
+			})
 			tenant, _ := newTenantWithAuthConfig(t, svc.Enclave, proto.AuthConfig{
 				Playfab: proto.AuthPlayfabConfig{
 					Enabled: true,
diff --git a/rpc/helpers_test.go b/rpc/helpers_test.go
index 3f2310a7..c1e670f7 100644
--- a/rpc/helpers_test.go
+++ b/rpc/helpers_test.go
@@ -51,7 +51,7 @@ func initRPC(t *testing.T, options ...func(*config.Config)) *rpc.RPC {
 		opt(cfg)
 	}
 
-	svc, err := rpc.New(cfg, &http.Client{Transport: &testTransport{RoundTripper: http.DefaultTransport}})
+	svc, err := rpc.New(cfg, &testTransport{RoundTripper: http.DefaultTransport})
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -59,13 +59,13 @@ func initRPC(t *testing.T, options ...func(*config.Config)) *rpc.RPC {
 	return svc
 }
 
-func initRPCWithClient(t *testing.T, client *http.Client, options ...func(*config.Config)) *rpc.RPC {
+func initRPCWithClient(t *testing.T, transport http.RoundTripper, options ...func(*config.Config)) *rpc.RPC {
 	cfg := initConfig(t, awsEndpoint)
 	for _, opt := range options {
 		opt(cfg)
 	}
 
-	svc, err := rpc.New(cfg, client)
+	svc, err := rpc.New(cfg, transport)
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/rpc/rpc.go b/rpc/rpc.go
index dddde452..60b3221a 100644
--- a/rpc/rpc.go
+++ b/rpc/rpc.go
@@ -75,9 +75,10 @@ type RPC struct {
 	running      int32
 }
 
-func New(cfg *config.Config, client *http.Client) (*RPC, error) {
-	if client == nil {
-		client = &http.Client{Timeout: 10 * time.Second}
+func New(cfg *config.Config, transport http.RoundTripper) (*RPC, error) {
+	client := &http.Client{
+		Timeout:   30 * time.Second,
+		Transport: transport,
 	}
 	wrappedClient := tracing.WrapClient(client)