Easily mask sensitive headers from request logs by adding HeadersSanitizer

core/src/main/java/com/linecorp/armeria/common/logging/AbstractHeadersSanitizerBuilder.java

@@ -0,0 +1,106 @@
+/*
+ * Copyright 2023 LINE Corporation
+ *
+ * LINE Corporation licenses this file to you under the Apache License,
+ * version 2.0 (the "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at:
+ *
+ *   https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package com.linecorp.armeria.common.logging;
+
+import static java.util.Objects.requireNonNull;
+
+import java.util.HashSet;
+import java.util.Set;
+import java.util.function.Function;
+
+import com.google.common.collect.ImmutableSet;
+
+import com.linecorp.armeria.common.HttpHeaderNames;
+import com.linecorp.armeria.common.annotation.Nullable;
+
+import io.netty.util.AsciiString;
+
+/**
+ * A skeletal builder implementation for {@link HeadersSanitizer}.
+ */
+public abstract class AbstractHeadersSanitizerBuilder<T> {
+
+    // Referenced from:
+    // - https://docs.rs/tower-http/latest/tower_http/sensitive_headers/index.html
+    // - https://techdocs.akamai.com/edge-diagnostics/reference/sensitive-headers
+    // - https://cloud.spring.io/spring-cloud-netflix/multi/multi__router_and_filter_zuul.html#_cookies_and_sensitive_headers
+    private static final Set<AsciiString> DEFAULT_SENSITIVE_HEADERS =
+            ImmutableSet.of(HttpHeaderNames.AUTHORIZATION, HttpHeaderNames.COOKIE,
+                            HttpHeaderNames.SET_COOKIE, HttpHeaderNames.PROXY_AUTHORIZATION);
+
+    @Nullable
+    private Set<AsciiString> sensitiveHeaders;
+
+    private HeaderMaskingFunction maskingFunction = HeaderMaskingFunction.of();
+
+    AbstractHeadersSanitizerBuilder() {}
+
+    /**
+     * Adds the headers to mask before logging.
+     */
+    public AbstractHeadersSanitizerBuilder<T> sensitiveHeaders(CharSequence... headers) {
+        requireNonNull(headers, "headers");
+        return sensitiveHeaders(ImmutableSet.copyOf(headers));
+    }
+
+    /**
+     * Sets the headers to mask before logging.
+     */
+    public AbstractHeadersSanitizerBuilder<T> sensitiveHeaders(Iterable<? extends CharSequence> headers) {
+        requireNonNull(headers, "headers");
+        if (sensitiveHeaders == null) {
+            sensitiveHeaders = new HashSet<>();
+        }
+        headers.forEach(header -> sensitiveHeaders.add(AsciiString.of(header).toLowerCase()));
+        return this;
+    }
+
+    final Set<AsciiString> sensitiveHeaders() {
+        if (sensitiveHeaders != null) {
+            return ImmutableSet.copyOf(sensitiveHeaders);
+        }
+        return DEFAULT_SENSITIVE_HEADERS;
+    }
+
+    /**
+     * Sets the {@link Function} to use to maskFunction headers before logging.
+     * The default maskingFunction is {@link HeaderMaskingFunction#of()}
+     *
+     * <pre>{@code
+     * builder.maskingFunction((name, value) -> {
+     *   if (name.equals(HttpHeaderNames.AUTHORIZATION)) {
+     *      return "****";
+     *   } else if (name.equals(HttpHeaderNames.COOKIE)) {
+     *     return name.substring(0, 4) + "****";
+     *   } else {
+     *     return value;
+     *   }
+     * }
+     * }</pre>
+     */
+    public AbstractHeadersSanitizerBuilder<T> maskingFunction(HeaderMaskingFunction maskingFunction) {
+        this.maskingFunction = requireNonNull(maskingFunction, "maskingFunction");
+        return this;
+    }
+
+    /**
+     * Returns the {@link Function} to use to mask headers before logging.
+     */
+    final HeaderMaskingFunction maskingFunction() {
+        return maskingFunction;
+    }
+}

core/src/main/java/com/linecorp/armeria/common/logging/AbstractLogFormatterBuilder.java

@@ -31,16 +31,16 @@
 abstract class AbstractLogFormatterBuilder<T> {
 
     @Nullable
-    private BiFunction<? super RequestContext, ? super HttpHeaders, ? extends T> requestHeadersSanitizer;
+    private HeadersSanitizer<T> requestHeadersSanitizer;
 
     @Nullable
-    private BiFunction<? super RequestContext, ? super HttpHeaders, ? extends T> responseHeadersSanitizer;
+    private HeadersSanitizer<T> responseHeadersSanitizer;
 
     @Nullable
-    private BiFunction<? super RequestContext, ? super HttpHeaders, ? extends T> requestTrailersSanitizer;
+    private HeadersSanitizer<T> requestTrailersSanitizer;
 
     @Nullable
-    private BiFunction<? super RequestContext, ? super HttpHeaders, ? extends T> responseTrailersSanitizer;
+    private HeadersSanitizer<T> responseTrailersSanitizer;
 
     @Nullable
     private BiFunction<? super RequestContext, Object, ? extends T> requestContentSanitizer;
@@ -52,73 +52,137 @@ abstract class AbstractLogFormatterBuilder<T> {
      * Sets the {@link BiFunction} to use to sanitize request headers before logging. It is common to have the
      * {@link BiFunction} that removes sensitive headers, like {@code Cookie}, before logging. If unset, will
      * not sanitize request headers.
+     *
+     * <pre>{@code
+     * HeadersSanitizer<String> headersSanitizer =
+     *   HeadersSanitizer
+     *     .builderForText()
+     *     .sensitiveHeaders("Authorization", "Cookie")
+     *     ...
+     *     .build();
+     *
+     *  LogFormatter
+     *    .builderForText()
+     *    .requestHeadersSanitizer(headersSanitizer)
+     *    ...
+     * }</pre>
      */
     public AbstractLogFormatterBuilder<T> requestHeadersSanitizer(
             BiFunction<? super RequestContext, ? super HttpHeaders, ? extends T> requestHeadersSanitizer) {
-        this.requestHeadersSanitizer = requireNonNull(requestHeadersSanitizer, "requestHeadersSanitizer");
+        requireNonNull(requestHeadersSanitizer, "requestHeadersSanitizer");
+        // TODO(ikhoon): Replace BiFunction with HeadersSanitizer in Armeria 2.0.
+        this.requestHeadersSanitizer = requestHeadersSanitizer::apply;
         return this;
     }
 
     /**
      * Returns the {@link BiFunction} to use to sanitize request headers before logging.
      */
     @Nullable
-    final BiFunction<? super RequestContext, ? super HttpHeaders, ? extends T> requestHeadersSanitizer() {
+    final HeadersSanitizer<T> requestHeadersSanitizer() {
         return requestHeadersSanitizer;
     }
 
     /**
      * Sets the {@link BiFunction} to use to sanitize response headers before logging. It is common to have the
      * {@link BiFunction} that removes sensitive headers, like {@code Set-Cookie}, before logging. If unset,
      * will not sanitize response headers.
+     *
+     * <pre>{@code
+     * HeadersSanitizer<String> headersSanitizer =
+     *   HeadersSanitizer
+     *     .builderForText()
+     *     .sensitiveHeaders("Set-Cookie")
+     *     ...
+     *     .build();
+     *
+     *  LogFormatter
+     *    .builderForText()
+     *    .responseHeadersSanitizer(headersSanitizer)
+     *    ...
+     * }</pre>
      */
     public AbstractLogFormatterBuilder<T> responseHeadersSanitizer(
             BiFunction<? super RequestContext, ? super HttpHeaders, ? extends T> responseHeadersSanitizer) {
-        this.responseHeadersSanitizer = requireNonNull(responseHeadersSanitizer, "responseHeadersSanitizer");
+        // TODO(ikhoon): Replace BiFunction with HeadersSanitizer in Armeria 2.0.
+        requireNonNull(responseHeadersSanitizer, "responseHeadersSanitizer");
+        this.responseHeadersSanitizer = responseHeadersSanitizer::apply;
         return this;
     }
 
     /**
      * Returns the {@link BiFunction} to use to sanitize response headers before logging.
      */
     @Nullable
-    final BiFunction<? super RequestContext, ? super HttpHeaders, ? extends T> responseHeadersSanitizer() {
+    final HeadersSanitizer<T> responseHeadersSanitizer() {
         return responseHeadersSanitizer;
     }
 
     /**
      * Sets the {@link BiFunction} to use to sanitize request trailers before logging. If unset,
      * will not sanitize request trailers.
+     *
+     * <pre>{@code
+     * HeadersSanitizer<String> headersSanitizer =
+     *   HeadersSanitizer
+     *     .builderForText()
+     *     .sensitiveHeaders("...")
+     *     ...
+     *     .build();
+     *
+     *  LogFormatter
+     *    .builderForText()
+     *    .requestTrailersSanitizer(headersSanitizer)
+     *    ...
+     * }</pre>
      */
     public AbstractLogFormatterBuilder<T> requestTrailersSanitizer(
             BiFunction<? super RequestContext, ? super HttpHeaders, ? extends T> requestTrailersSanitizer) {
-        this.requestTrailersSanitizer = requireNonNull(requestTrailersSanitizer, "requestTrailersSanitizer");
+        // TODO(ikhoon): Replace BiFunction with HeadersSanitizer in Armeria 2.0.
+        requireNonNull(requestTrailersSanitizer, "requestTrailersSanitizer");
+        this.requestTrailersSanitizer = requestTrailersSanitizer::apply;
         return this;
     }
 
     /**
      * Returns the {@link BiFunction} to use to sanitize request trailers before logging.
      */
     @Nullable
-    final BiFunction<? super RequestContext, ? super HttpHeaders, ? extends T> requestTrailersSanitizer() {
+    final HeadersSanitizer<T> requestTrailersSanitizer() {
         return requestTrailersSanitizer;
     }
 
     /**
      * Sets the {@link BiFunction} to use to sanitize response trailers before logging. If unset,
      * will not sanitize response trailers.
+     *
+     * <pre>{@code
+     * HeadersSanitizer<String> headersSanitizer =
+     *   HeadersSanitizer
+     *     .builderForText()
+     *     .sensitiveHeaders("...")
+     *     ...
+     *     .build();
+     *
+     *  LogFormatter
+     *    .builderForText()
+     *    .responseTrailersSanitizer(headersSanitizer)
+     *    ...
+     * }</pre>
      */
     public AbstractLogFormatterBuilder<T> responseTrailersSanitizer(
             BiFunction<? super RequestContext, ? super HttpHeaders, ? extends T> responseTrailersSanitizer) {
-        this.responseTrailersSanitizer = requireNonNull(responseTrailersSanitizer, "responseTrailersSanitizer");
+        // TODO(ikhoon): Replace BiFunction with HeadersSanitizer in Armeria 2.0.
+        requireNonNull(responseTrailersSanitizer, "responseTrailersSanitizer");
+        this.responseTrailersSanitizer = responseTrailersSanitizer::apply;
         return this;
     }
 
     /**
      * Returns the {@link Function} to use to sanitize response trailers before logging.
      */
     @Nullable
-    final BiFunction<? super RequestContext, ? super HttpHeaders, ? extends T> responseTrailersSanitizer() {
+    final HeadersSanitizer<T> responseTrailersSanitizer() {
         return responseTrailersSanitizer;
     }
 
@@ -127,6 +191,13 @@ public AbstractLogFormatterBuilder<T> responseTrailersSanitizer(
      * It is common to have the {@link BiFunction} that removes sensitive headers, like {@code "Cookie"} and
      * {@code "Set-Cookie"}, before logging. This method is a shortcut for:
      * <pre>{@code
+     * HeadersSanitizer<String> headersSanitizer =
+     *   HeadersSanitizer
+     *     .builderForText()
+     *     .sensitiveHeaders("...")
+     *     ...
+     *     .build();
+     *
      * builder.requestHeadersSanitizer(headersSanitizer);
      * builder.requestTrailersSanitizer(headersSanitizer);
      * builder.responseHeadersSanitizer(headersSanitizer);

core/src/main/java/com/linecorp/armeria/common/logging/HeaderMaskingFunction.java

@@ -0,0 +1,42 @@
+/*
+ *  Copyright 2024 LINE Corporation
+ *
+ *  LINE Corporation licenses this file to you under the Apache License,
+ *  version 2.0 (the "License"); you may not use this file except in compliance
+ *  with the License. You may obtain a copy of the License at:
+ *
+ *    https://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ *  License for the specific language governing permissions and limitations
+ *  under the License.
+ */
+
+package com.linecorp.armeria.common.logging;
+
+import com.linecorp.armeria.common.annotation.Nullable;
+
+import io.netty.util.AsciiString;
+
+/**
+ * A function that masks the specified header value.
+ */
+@FunctionalInterface
+public interface HeaderMaskingFunction {
+
+    /**
+     * Returns the default {@link HeaderMaskingFunction} that masks the given value with {@code ****}.
+     */
+    static HeaderMaskingFunction of() {
+        return (name, value) -> "****";
+    }
+
+    /**
+     * Masks the specified {@code value} of the specified {@code name}.
+     * If {@code null} is returned, the specified {@code value} will be removed from the log.
+     */
+    @Nullable
+    String mask(AsciiString name, String value);
+}