Merge pull request #1 from samvaity/it-works

Add record sanitizers to test proxy

sdk/core/azure-core-test/src/main/java/com/azure/core/test/InterceptorManager.java

@@ -9,8 +9,9 @@
 import com.azure.core.test.models.NetworkCallRecord;
 import com.azure.core.test.models.RecordedData;
 import com.azure.core.test.models.RecordingRedactor;
-import com.azure.core.test.policy.TestProxyRecordPolicy;
+import com.azure.core.test.models.TestProxySanitizer;
 import com.azure.core.test.policy.RecordNetworkCallPolicy;
+import com.azure.core.test.policy.TestProxyRecordPolicy;
 import com.azure.core.util.CoreUtils;
 import com.azure.core.util.logging.ClientLogger;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -72,6 +73,7 @@ public class InterceptorManager implements AutoCloseable {
     private TestProxyRecordPolicy testProxyRecordPolicy;
     private TestProxyPlaybackClient testProxyPlaybackClient;
     private final Queue<String> proxyVariableQueue = new LinkedList<>();
+    private List<TestProxySanitizer> recordSanitizers;
 
     /**
      * Creates a new InterceptorManager that either replays test-session records or saves them.
@@ -270,7 +272,7 @@ public Consumer<String> getProxyVariableConsumer() {
      */
     public HttpPipelinePolicy getRecordPolicy() {
         if (enableTestProxy) {
-            return startProxyRecording(Collections.emptyList());
+            return startProxyRecording();
         }
         return getRecordPolicy(Collections.emptyList());
     }
@@ -286,7 +288,7 @@ public HttpPipelinePolicy getRecordPolicy() {
     public HttpPipelinePolicy getRecordPolicy(List<Function<String, String>> recordingRedactors) {
         if (enableTestProxy) {
             proxyVariableQueue.clear();
-            return startProxyRecording(recordingRedactors);
+            return startProxyRecording();
         }
         return new RecordNetworkCallPolicy(recordedData, recordingRedactors);
     }
@@ -298,8 +300,8 @@ public HttpPipelinePolicy getRecordPolicy(List<Function<String, String>> recordi
      */
     public HttpClient getPlaybackClient() {
         if (enableTestProxy) {
-            testProxyPlaybackClient = new TestProxyPlaybackClient();
-            proxyVariableQueue.addAll(testProxyPlaybackClient.startPlayback(playbackRecordName, null));
+            testProxyPlaybackClient = new TestProxyPlaybackClient(this.recordSanitizers);
+            proxyVariableQueue.addAll(testProxyPlaybackClient.startPlayback(playbackRecordName));
             return testProxyPlaybackClient;
         } else {
             return new PlaybackClient(recordedData, textReplacementRules);
@@ -367,9 +369,9 @@ private static URI toURI(URL url, ClientLogger logger) {
         }
     }
 
-    private HttpPipelinePolicy startProxyRecording(List<Function<String, String>> recordingRedactors) {
-        this.testProxyRecordPolicy = new TestProxyRecordPolicy();
-        testProxyRecordPolicy.startRecording(playbackRecordName, recordingRedactors);
+    private HttpPipelinePolicy startProxyRecording() {
+        this.testProxyRecordPolicy = new TestProxyRecordPolicy(this.recordSanitizers);
+        testProxyRecordPolicy.startRecording(playbackRecordName);
         return testProxyRecordPolicy;
     }
 
@@ -426,4 +428,12 @@ private File createRecordFile(String testName) throws IOException {
     public void addTextReplacementRule(String regex, String replacement) {
         textReplacementRules.put(regex, replacement);
     }
+
+    /**
+     * Add text replacement rule (regex as key, the replacement text as value) into {@code recordSanitizers}
+     * @param recordSanitizers the list of replacement regex and rules.
+     */
+    public void addRecordSanitizers(List<TestProxySanitizer> recordSanitizers) {
+        this.recordSanitizers = recordSanitizers;
+    }
 }

sdk/core/azure-core-test/src/main/java/com/azure/core/test/http/TestProxyPlaybackClient.java

@@ -7,6 +7,7 @@
 import com.azure.core.http.HttpMethod;
 import com.azure.core.http.HttpRequest;
 import com.azure.core.http.HttpResponse;
+import com.azure.core.test.models.TestProxySanitizer;
 import com.azure.core.test.utils.HttpURLConnectionHttpClient;
 import com.azure.core.test.utils.TestProxyUtils;
 import com.azure.core.util.Context;
@@ -16,12 +17,18 @@
 import reactor.core.publisher.Mono;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.Comparator;
 import java.util.LinkedList;
+import java.util.List;
 import java.util.Map;
 import java.util.Queue;
 import java.util.stream.Collectors;
 
+import static com.azure.core.test.utils.TestProxyUtils.getRegexSanitizerRequests;
+import static com.azure.core.test.utils.TestProxyUtils.loadSanitizers;
+
 /**
  * A {@link HttpClient} that plays back test recordings from the external test proxy.
  */
@@ -31,19 +38,26 @@ public class TestProxyPlaybackClient implements HttpClient {
     private String xRecordingId;
     private static final SerializerAdapter SERIALIZER = new JacksonAdapter();
 
+    private static final List<TestProxySanitizer> DEFAULT_SANITIZERS = loadSanitizers();
+    private final List<TestProxySanitizer> sanitizers = new ArrayList<>();
+
+    public TestProxyPlaybackClient(List<TestProxySanitizer> customSanitizers) {
+        this.sanitizers.addAll(DEFAULT_SANITIZERS);
+        this.sanitizers.addAll(customSanitizers == null ? Collections.emptyList() : customSanitizers);
+    }
+
     /**
      * Starts playback of a test recording.
      * @param recordFile The name of the file to read.
-     * @param textReplacementRules Rules for replacing text in the playback.
      * @return A {@link Queue} representing the variables in the recording.
      * @throws RuntimeException if an {@link IOException} is thrown.
      */
-    public Queue<String> startPlayback(String recordFile, Map<String, String> textReplacementRules) {
-        // TODO: replacement rules
+    public Queue<String> startPlayback(String recordFile) {
         HttpRequest request = new HttpRequest(HttpMethod.POST, String.format("%s/playback/start", TestProxyUtils.getProxyUrl()))
             .setBody(String.format("{\"x-recording-file\": \"%s\"}", recordFile));
         try (HttpResponse response = client.sendSync(request, Context.NONE)) {
             xRecordingId = response.getHeaderValue("x-recording-id");
+            addProxySanitization();
             String body = response.getBodyAsString().block();
             // The test proxy stores variables in a map with no guaranteed order.
             // The Java implementation of recording did not use a map, but relied on the order
@@ -80,4 +94,12 @@ public Mono<HttpResponse> send(HttpRequest request) {
         TestProxyUtils.changeHeaders(request, xRecordingId, "playback");
         return client.send(request);
     }
+
+    private void addProxySanitization() {
+        getRegexSanitizerRequests(this.sanitizers)
+            .forEach(request -> {
+                request.setHeader("x-recording-id", xRecordingId);
+                client.sendSync(request, Context.NONE);
+            });
+    }
 }

sdk/core/azure-core-test/src/main/java/com/azure/core/test/http/TestProxyTestServer.java

@@ -17,6 +17,8 @@
 public class TestProxyTestServer implements Closeable {
     private final DisposableServer server;
 
+    private static final String TEST_RESPONSE_BODY = "{\"modelId\":\"0cd2728b-210e-4c05-b706-f70554276bcc\",\"createdDateTime\":\"2022-08-31T00:00:00Z\",\"apiVersion\":\"2022-08-31\"}";
+
     /**
      * Constructor for TestProxyTestServer
      */
@@ -32,6 +34,14 @@ public TestProxyTestServer() {
                         res.addHeader(requestHeader.getKey(), requestHeader.getValue());
                     }
                     return res.status(HttpResponseStatus.OK).sendString(Mono.just("echoheaders"));
+                })
+                .get("/fr/models", (req, res) -> {
+                    for (Map.Entry<String, String> requestHeader : req.requestHeaders()) {
+                        res.addHeader(requestHeader.getKey(), requestHeader.getValue());
+                    }
+                    return res.status(HttpResponseStatus.OK)
+                        .addHeader("Content-Type","application/json")
+                        .sendString(Mono.just(TEST_RESPONSE_BODY));
                 }))
             .bindNow();
     }

sdk/core/azure-core-test/src/main/java/com/azure/core/test/models/TestProxySanitizer.java

@@ -0,0 +1,49 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.core.test.models;
+
+/**
+ * Keeps track of different sanitizers that redact the sensitive information when recording
+ */
+public class TestProxySanitizer {
+    private final TestProxySanitizerType testProxySanitizerType;
+    private final String regexKey;
+    private final String redactedValue;
+
+    /**
+     * Creates an instance of TestProxySanitizer
+     * @param regexKey the regex key to lookup for redaction
+     * @param redactedValue the replacement for regex matched content
+     * @param testProxySanitizerType the type of sanitizer
+     */
+    public TestProxySanitizer(String regexKey, String redactedValue, TestProxySanitizerType testProxySanitizerType) {
+        this.testProxySanitizerType = testProxySanitizerType;
+        this.regexKey = regexKey;
+        this.redactedValue = redactedValue;
+    }
+
+    /**
+     * Get the type of proxy sanitizer
+     * @return the type of proxy sanitizer
+     */
+    public TestProxySanitizerType getType() {
+        return testProxySanitizerType;
+    }
+
+    /**
+     * Get the regex key to lookup for redaction
+     * @return the regex key to lookup for redaction
+     */
+    public String getRegex() {
+        return regexKey;
+    }
+
+    /**
+     * Get the  replacement for regex matched content
+     * @return the replacement for regex matched content
+     */
+    public String getRedactedValue() {
+        return redactedValue;
+    }
+}

sdk/core/azure-core-test/src/main/java/com/azure/core/test/models/TestProxySanitizerType.java

@@ -0,0 +1,28 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.core.test.models;
+
+/**
+ * The possible record sanitizer types.
+ * Each sanitizer is optionally prefaced with the specific part of the request/response pair that it applies to.
+ */
+public enum TestProxySanitizerType {
+    /**
+     * Sanitize the request url.
+     */
+    URL("UriRegexSanitizer"),
+    /**
+     * Sanitize the response body.
+     */
+    BODY("BodyKeySanitizer"),
+    /**
+     * Sanitize the request/response headers.
+     */
+    HEADER("HeaderRegexSanitizer");
+
+    public final String name;
+
+    TestProxySanitizerType(String name) {
+        this.name = name;
+    }
+}

sdk/core/azure-core-test/src/main/java/com/azure/core/test/policy/TestProxyRecordPolicy.java

@@ -6,9 +6,11 @@
 import com.azure.core.http.HttpMethod;
 import com.azure.core.http.HttpPipelineCallContext;
 import com.azure.core.http.HttpPipelineNextPolicy;
+import com.azure.core.http.HttpPipelineNextSyncPolicy;
 import com.azure.core.http.HttpRequest;
 import com.azure.core.http.HttpResponse;
 import com.azure.core.http.policy.HttpPipelinePolicy;
+import com.azure.core.test.models.TestProxySanitizer;
 import com.azure.core.test.utils.HttpURLConnectionHttpClient;
 import com.azure.core.test.utils.TestProxyUtils;
 import com.azure.core.util.Context;
@@ -18,34 +20,47 @@
 import reactor.core.publisher.Mono;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Queue;
 import java.util.concurrent.atomic.AtomicInteger;
-import java.util.function.Function;
 import java.util.stream.Collectors;
 
+import static com.azure.core.test.utils.TestProxyUtils.getRegexSanitizerRequests;
+import static com.azure.core.test.utils.TestProxyUtils.loadSanitizers;
+
+
 /**
  * A {@link HttpPipelinePolicy} for redirecting traffic through the test proxy for recording.
  */
 public class TestProxyRecordPolicy implements HttpPipelinePolicy {
     private static final SerializerAdapter SERIALIZER = new JacksonAdapter();
-
     private final HttpURLConnectionHttpClient client = new HttpURLConnectionHttpClient();
     private String xRecordingId;
+    private final List<TestProxySanitizer> sanitizers = new ArrayList<>();
+    private static final List<TestProxySanitizer> DEFAULT_SANITIZERS = loadSanitizers();
+
+    public TestProxyRecordPolicy(List<TestProxySanitizer> customSanitizers) {
+        this.sanitizers.addAll(DEFAULT_SANITIZERS);
+        this.sanitizers.addAll(customSanitizers == null ? Collections.emptyList() : customSanitizers);
+    }
 
     /**
      * Starts a recording of test traffic.
+     *
      * @param recordFile The name of the file to save the recording to.
-     * @param redactors The set of redactors to send to the test proxy.
      */
-    public void startRecording(String recordFile, List<Function<String, String>> redactors) {
-        // TODO: redactors
+    public void startRecording(String recordFile) {
         HttpRequest request = new HttpRequest(HttpMethod.POST, String.format("%s/record/start", TestProxyUtils.getProxyUrl()))
             .setBody(String.format("{\"x-recording-file\": \"%s\"}", recordFile));
+
         HttpResponse response = client.sendSync(request, Context.NONE);
 
-        xRecordingId = response.getHeaderValue("x-recording-id");
+        this.xRecordingId = response.getHeaderValue("x-recording-id");
+
+        addProxySanitization();
     }
 
     /**
@@ -78,11 +93,25 @@ private String serializeVariables(Queue<String> variables) {
         }
     }
 
+    @Override
+    public HttpResponse processSync(HttpPipelineCallContext context, HttpPipelineNextSyncPolicy next) {
+        TestProxyUtils.changeHeaders(context.getHttpRequest(), xRecordingId, "record");
+        return next.processSync();
+    }
+
     @Override
     public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) {
         HttpRequest request = context.getHttpRequest();
         TestProxyUtils.changeHeaders(request, xRecordingId, "record");
         return next.process();
     }
+
+    private void addProxySanitization() {
+        getRegexSanitizerRequests(this.sanitizers)
+            .forEach(request -> {
+                request.setHeader("x-recording-id", xRecordingId);
+                client.sendSync(request, Context.NONE);
+            });
+    }
 }