Tweak content caching

Author: itssimon

src/main/java/io/apitally/common/ApitallyClient.java

@@ -66,6 +66,7 @@ public enum HubRequestStatus {
     private ScheduledFuture<?> syncTask;
     private StartupData startupData;
     private boolean startupDataSent = false;
+    private boolean enabled = true;
 
     public final RequestCounter requestCounter;
     public final RequestLogger requestLogger;
@@ -89,6 +90,10 @@ public ApitallyClient(String clientId, String env, RequestLoggingConfig requestL
         this.consumerRegistry = new ConsumerRegistry();
     }
 
+    public boolean isEnabled() {
+        return enabled;
+    }
+
     private HttpClient createHttpClient() {
         return HttpClient.newBuilder()
                 .version(HttpClient.Version.HTTP_1_1)
@@ -193,7 +198,10 @@ private void sendLogData() {
                         .POST(HttpRequest.BodyPublishers.ofInputStream(() -> inputStream))
                         .build();
                 HubRequestStatus status = sendHubRequest(request).join();
-                if (status == HubRequestStatus.RETRYABLE_ERROR) {
+                if (status == HubRequestStatus.PAYMENT_REQUIRED) {
+                    requestLogger.clear();
+                    requestLogger.setSuspendUntil(System.currentTimeMillis() + (3600 * 1000L));
+                } else if (status == HubRequestStatus.RETRYABLE_ERROR) {
                     requestLogger.retryFileLater(logFile);
                     break;
                 } else {
@@ -219,22 +227,11 @@ public CompletableFuture<HubRequestStatus> sendHubRequest(HttpRequest request) {
                         if (response.statusCode() >= 200 && response.statusCode() < 300) {
                             return HubRequestStatus.OK;
                         } else if (response.statusCode() == 402) {
-                            // 402 is only returned by the log endpoint, meaning request logging is not
-                            // allowed for this client
-                            Optional<String> retryAfter = response.headers().firstValue("Retry-After");
-                            if (retryAfter.isPresent()) {
-                                try {
-                                    int retryAfterSeconds = Integer.parseInt(retryAfter.get());
-                                    requestLogger
-                                            .setSuspendUntil(System.currentTimeMillis() + (retryAfterSeconds * 1000L));
-                                } catch (NumberFormatException e) {
-                                    // Ignore invalid Retry-After header
-                                }
-                            }
-                            requestLogger.clear();
                             return HubRequestStatus.PAYMENT_REQUIRED;
                         } else if (response.statusCode() == 404) {
+                            enabled = false;
                             stopSync();
+                            requestLogger.close();
                             logger.error("Invalid Apitally client ID: {}", clientId);
                             return HubRequestStatus.INVALID_CLIENT_ID;
                         } else if (response.statusCode() == 422) {
@@ -304,6 +301,7 @@ private void sync() {
 
     public void shutdown() {
         try {
+            enabled = false;
             boolean syncRunning = syncTask != null;
             stopSync();
             if (syncRunning) {

src/main/java/io/apitally/common/RequestLogger.java

@@ -21,11 +21,12 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+
 import io.apitally.common.dto.Header;
 import io.apitally.common.dto.Request;
 import io.apitally.common.dto.Response;
-import com.fasterxml.jackson.databind.ObjectMapper;
-import com.fasterxml.jackson.databind.node.ObjectNode;
 
 public class RequestLogger {
     private static final Logger logger = LoggerFactory.getLogger(RequestLogger.class);
@@ -37,7 +38,7 @@ public class RequestLogger {
     private static final byte[] BODY_TOO_LARGE = "<body too large>".getBytes(StandardCharsets.UTF_8);
     private static final byte[] BODY_MASKED = "<masked>".getBytes(StandardCharsets.UTF_8);
     private static final String MASKED = "******";
-    private static final List<String> ALLOWED_CONTENT_TYPES = Arrays.asList("application/json", "text/plain");
+    public static final List<String> ALLOWED_CONTENT_TYPES = Arrays.asList("application/json", "text/plain");
     private static final List<String> EXCLUDE_PATH_PATTERNS = Arrays.asList(
             "/_?healthz?$",
             "/_?health[_-]?checks?$",

src/main/java/io/apitally/spring/ApitallyFilter.java

@@ -15,6 +15,7 @@
 
 import io.apitally.common.ApitallyClient;
 import io.apitally.common.ConsumerRegistry;
+import io.apitally.common.RequestLogger;
 import io.apitally.common.dto.Consumer;
 import io.apitally.common.dto.Header;
 import io.apitally.common.dto.Request;
@@ -38,23 +39,33 @@ public ApitallyFilter(ApitallyClient client) {
     @Override
     protected void doFilterInternal(@NonNull HttpServletRequest request, @NonNull HttpServletResponse response,
             @NonNull FilterChain filterChain) throws ServletException, IOException {
-        final boolean isContentCachingEnabled = client.requestLogger.getConfig().isEnabled()
-                && (client.requestLogger.getConfig().isRequestBodyIncluded() || client.requestLogger.getConfig()
-                        .isResponseBodyIncluded());
+        if (!client.isEnabled()) {
+            filterChain.doFilter(request, response);
+            return;
+        }
+
+        String requestContentType = request.getContentType();
+        final boolean shouldCacheRequest = client.requestLogger.getConfig().isEnabled()
+                && client.requestLogger.getConfig().isRequestBodyIncluded()
+                && requestContentType != null
+                && RequestLogger.ALLOWED_CONTENT_TYPES.stream()
+                        .anyMatch(allowedContentType -> requestContentType.startsWith(allowedContentType));
+        final boolean shouldCacheResponse = client.requestLogger.getConfig().isEnabled()
+                && client.requestLogger.getConfig().isResponseBodyIncluded();
+        ContentCachingRequestWrapper cachingRequest = shouldCacheRequest
+                ? new ContentCachingRequestWrapper(request)
+                : null;
+        ContentCachingResponseWrapper cachingResponse = shouldCacheResponse
+                ? new ContentCachingResponseWrapper(response)
+                : null;
 
         Exception exception = null;
         final long startTime = System.currentTimeMillis();
-        ContentCachingRequestWrapper cachingRequest = null;
-        ContentCachingResponseWrapper cachingResponse = null;
 
         try {
-            if (isContentCachingEnabled) {
-                cachingRequest = new ContentCachingRequestWrapper(request);
-                cachingResponse = new ContentCachingResponseWrapper(response);
-                filterChain.doFilter(cachingRequest, cachingResponse);
-            } else {
-                filterChain.doFilter(request, response);
-            }
+            filterChain.doFilter(
+                    cachingRequest != null ? cachingRequest : request,
+                    cachingResponse != null ? cachingResponse : response);
         } catch (Exception e) {
             exception = e;
             throw e;