fix(httpreplay): add ignore-header flag, fix tests

httpreplay/cmd/httpr/httpr.go

@@ -29,20 +29,24 @@ import (
 	"net/http"
 	"os"
 	"os/signal"
+	"strings"
 
 	"cloud.google.com/go/httpreplay/internal/proxy"
 	"github.com/google/martian/v3/martianhttp"
 )
 
 var (
-	port         = flag.Int("port", 8080, "port of the proxy")
-	controlPort  = flag.Int("control-port", 8181, "port for controlling the proxy")
-	record       = flag.String("record", "", "record traffic and save to filename")
-	replay       = flag.String("replay", "", "read filename and replay traffic")
-	debugHeaders = flag.Bool("debug-headers", false, "log header mismatches")
+	port          = flag.Int("port", 8080, "port of the proxy")
+	controlPort   = flag.Int("control-port", 8181, "port for controlling the proxy")
+	record        = flag.String("record", "", "record traffic and save to filename")
+	replay        = flag.String("replay", "", "read filename and replay traffic")
+	debugHeaders  = flag.Bool("debug-headers", false, "log header mismatches")
+	ignoreHeaders repeatedString
 )
 
 func main() {
+	flag.Var(&ignoreHeaders, "ignore-header", "header key(s) to ignore when matching")
+
 	flag.Parse()
 	if *record == "" && *replay == "" {
 		log.Fatal("provide either -record or -replay")
@@ -63,6 +67,9 @@ func main() {
 		log.Fatal(err)
 	}
 	proxy.DebugHeaders = *debugHeaders
+	for _, key := range ignoreHeaders {
+		pr.IgnoreHeader(key)
+	}
 
 	// Expose handlers on the control port.
 	mux := http.NewServeMux()
@@ -108,3 +115,18 @@ func handleInitial(pr *proxy.Proxy) http.HandlerFunc {
 		}
 	}
 }
+
+type repeatedString []string
+
+func (i *repeatedString) String() string {
+	v := make([]string, 0)
+	if i != nil {
+		v = *i
+	}
+	return strings.Join(v, ",")
+}
+
+func (i *repeatedString) Set(value string) error {
+	*i = append(*i, value)
+	return nil
+}

httpreplay/cmd/httpr/integration_test.go

@@ -143,7 +143,15 @@ func start(modeFlag, filename string) (*exec.Cmd, *http.Transport, string, error
 	if err != nil {
 		return nil, nil, "", err
 	}
-	cmd := exec.Command("./httpr", "-port", pport, "-control-port", cport, modeFlag, filename, "-debug-headers")
+	cmd := exec.Command("./httpr",
+		"-port", pport,
+		"-control-port", cport,
+		modeFlag,
+		filename,
+		"-debug-headers",
+		"-ignore-header", "X-Goog-Api-Client",
+		"-ignore-header", "X-Goog-Gcs-Idempotency-Token",
+	)
 	if err := cmd.Start(); err != nil {
 		return nil, nil, "", err
 	}
@@ -204,6 +212,9 @@ func proxyTransport(pport, cport string) (*http.Transport, error) {
 }
 
 func getBucketInfo(ctx context.Context, hc *http.Client) (string, error) {
+	ctx, cc := context.WithTimeout(ctx, 10*time.Second)
+	defer cc()
+
 	client, err := storage.NewClient(ctx, option.WithHTTPClient(hc))
 	if err != nil {
 		return "", err

httpreplay/httpreplay_test.go

@@ -68,6 +68,9 @@ func TestIntegration_RecordAndReplay(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	rec.RemoveRequestHeaders("X-Goog-Api-Client")
+	rec.RemoveRequestHeaders("X-Goog-Gcs-Idempotency-Token")
+
 	hc, err := rec.Client(ctx, option.WithTokenSource(
 		testutil.TokenSource(ctx, storage.ScopeFullControl)))
 	if err != nil {
@@ -84,6 +87,8 @@ func TestIntegration_RecordAndReplay(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	rep.IgnoreHeader("X-Goog-Api-Client")
+	rep.IgnoreHeader("X-Goog-Gcs-Idempotency-Token")
 	defer rep.Close()
 	hc, err = rep.Client(ctx)
 	if err != nil {
@@ -153,7 +158,9 @@ func setup(ctx context.Context) (cleanup func(), err error) {
 // TODO(jba): test errors
 
 func run(t *testing.T, hc *http.Client) (*storage.BucketAttrs, []byte) {
-	ctx := context.Background()
+	ctx, cc := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cc()
+
 	client, err := storage.NewClient(ctx, option.WithHTTPClient(hc))
 	if err != nil {
 		t.Fatal(err)
@@ -188,7 +195,9 @@ func run(t *testing.T, hc *http.Client) (*storage.BucketAttrs, []byte) {
 }
 
 func testReadCRC(t *testing.T, hc *http.Client, mode string) {
-	ctx := context.Background()
+	ctx, cc := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cc()
+
 	client, err := storage.NewClient(ctx, option.WithHTTPClient(hc))
 	if err != nil {
 		t.Fatalf("%s: %v", mode, err)