From ddd6634b38a8f5320b0fe1fdcd3fc329d9d271cd Mon Sep 17 00:00:00 2001
From: Andrew Boyle <pacejackson@users.noreply.github.com>
Date: Thu, 20 Feb 2025 09:12:22 -0600
Subject: [PATCH] always update thrift header write list in headerbp client
 middleware (#685)

This is really just to make the code cleaner since it technically works because thrift.GetWriteHeaderList returns the underlying list, not a copy, and slices.Delete deletes from the source slice. I think the code is a lot clearer and less likely to have a bug later if we always set the outgoing header list.

I verified this by adding a test case. The test passes without my changes, but it fails if I copy outgoing before deleting the headers from it.
---
 thriftbp/client_middlewares.go |  2 +-
 thriftbp/server_test.go        | 53 ++++++++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/thriftbp/client_middlewares.go b/thriftbp/client_middlewares.go
index d5a1dd5f1..68de2cb36 100644
--- a/thriftbp/client_middlewares.go
+++ b/thriftbp/client_middlewares.go
@@ -449,8 +449,8 @@ func ClientBaseplateHeadersMiddleware(service, client string) thrift.ClientMiddl
 					for k, v := range toAdd {
 						ctx = thrift.SetHeader(ctx, k, v)
 					}
-					ctx = thrift.SetWriteHeaderList(ctx, outgoing)
 				}
+				ctx = thrift.SetWriteHeaderList(ctx, outgoing)
 				return next.Call(ctx, method, args, result)
 			},
 		}
diff --git a/thriftbp/server_test.go b/thriftbp/server_test.go
index 731e1f655..564a035b8 100644
--- a/thriftbp/server_test.go
+++ b/thriftbp/server_test.go
@@ -141,6 +141,59 @@ func TestHeaderPropagation(t *testing.T) {
 	}
 }
 
+// verify that new headers are removed when there are no headers set to propagate
+func TestHeaderPropagation_removeOnly(t *testing.T) {
+	store := newSecretsStore(t)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+
+	ecImpl := ecinterface.Mock()
+
+	downstreamProcessor := baseplatethrift.NewBaseplateServiceV2Processor(&headerPropagationVerificationService{
+		wantUnset: []string{"x-bp-test"},
+	})
+	downstreamServer, err := thrifttest.NewBaseplateServer(thrifttest.ServerConfig{
+		Processor:       downstreamProcessor,
+		SecretStore:     store,
+		EdgeContextImpl: ecImpl,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	downstreamServer.Start(ctx)
+	time.Sleep(100 * time.Millisecond) // wait for the server to start
+
+	originProcessor := baseplatethrift.NewBaseplateServiceV2Processor(&headerPropagationVerificationService{
+		wantUnset: []string{"x-bp-test"},
+		client: func() baseplatethrift.BaseplateServiceV2 {
+			return baseplatethrift.NewBaseplateServiceV2Client(downstreamServer.ClientPool.TClient())
+		},
+	})
+	server, err := thrifttest.NewBaseplateServer(thrifttest.ServerConfig{
+		Processor:   originProcessor,
+		SecretStore: store,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	server.Start(ctx)
+	time.Sleep(100 * time.Millisecond) // wait for the server to start
+
+	client := baseplatethrift.NewBaseplateServiceV2Client(server.ClientPool.TClient())
+	got, err := client.IsHealthy(ctx, &baseplatethrift.IsHealthyRequest{
+		Probe: baseplatethrift.IsHealthyProbePtr(baseplatethrift.IsHealthyProbe_READINESS),
+	})
+
+	if err != nil {
+		t.Errorf("expected no error, got %v", err)
+	}
+
+	const want = true
+	if got != want {
+		t.Errorf("success mismatch, want %v, got %v", want, got)
+	}
+}
 func setHeader(ctx context.Context, key, value string) context.Context {
 	ctx = thrift.SetHeader(ctx, key, value)
 	return thrift.SetWriteHeaderList(ctx, append(thrift.GetWriteHeaderList(ctx), key))