Easily mask sensitive headers from request logs by adding HeadersSanitizer

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

@@ -0,0 +1,81 @@
+/*
+ * 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;
+
+import static java.util.Objects.requireNonNull;
+
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.function.Function;
+
+/**
+ * A skeletal builder implementation for {@link HeadersSanitizer}.
+ */
+abstract class AbstractHeadersSanitizerBuilder<T> {
+
+    private final Set<CharSequence> maskingHeaders = new HashSet<>();
+
+    private Function<String, String> mask = (header) -> "****";
+
+    /**
+     * Sets the {@link Set} which includes headers to mask before logging.
+     */
+    public AbstractHeadersSanitizerBuilder<T> maskingHeaders(CharSequence... headers) {
+        requireNonNull(headers, "headers");
+        Arrays.stream(headers).map(header -> header.toString().toLowerCase()).forEach(maskingHeaders::add);
+        return this;
+    }
+
+    /**
+     * Sets the {@link Set} which includes headers to mask before logging.
+     */
+    public AbstractHeadersSanitizerBuilder<T> maskingHeaders(Iterable<? extends CharSequence> headers) {
+        requireNonNull(headers, "headers");
+        headers.forEach(header -> maskingHeaders.add(header.toString().toLowerCase()));
+        return this;
+    }
+
+    /**
+     * Returns the {@link Set} which includes headers to mask before logging.
+     */
+    final Set<CharSequence> maskingHeaders() {
+        return maskingHeaders;
+    }
+
+    /**
+     * Sets the {@link Function} to use to maskFunction headers before logging.
+     */
+    public AbstractHeadersSanitizerBuilder<T> maskingFunction(Function<String, String> maskingFunction) {
+        this.mask = requireNonNull(maskingFunction, "maskingFunction");
+        return this;
+    }
+
+    /**
+     * Returns the {@link Function} to use to mask headers before logging.
+     */
+    final Function<String, String> maskingFunction() {
+        return mask;
+    }
+
+    protected final Set<CharSequence> defaultMaskingHeaders() {
+        final HashSet<CharSequence> defaultMaskingHeaders = new HashSet<>();
+        defaultMaskingHeaders.add(HttpHeaderNames.AUTHORIZATION.toLowerCase().toString());
+        defaultMaskingHeaders.add(HttpHeaderNames.SET_COOKIE.toLowerCase().toString());
+        return defaultMaskingHeaders;
+    }
+}

core/src/main/java/com/linecorp/armeria/common/HeadersSanitizer.java

@@ -0,0 +1,54 @@
+/*
+ * 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;
+
+import java.util.function.BiFunction;
+
+import com.fasterxml.jackson.databind.JsonNode;
+
+/**
+ * A sanitizer that sanitizes {@link HttpHeaders}.
+ */
+public interface HeadersSanitizer<T> extends BiFunction<RequestContext, HttpHeaders, T> {
+    /**
+     * Returns the default text {@link HeadersSanitizer}.
+     */
+    static HeadersSanitizer<String> ofText() {
+        return TextHeadersSanitizer.INSTANCE;
+    }
+
+    /**
+     * Returns a newly created {@link TextHeadersSanitizerBuilder}.
+     */
+    static TextHeadersSanitizerBuilder builderForText() {
+        return new TextHeadersSanitizerBuilder();
+    }
+
+    /**
+     * Returns the default json {@link HeadersSanitizer}.
+     */
+    static HeadersSanitizer<JsonNode> ofJson() {
+        return JsonHeadersSanitizer.INSTANCE;
+    }
+
+    /**
+     * Returns a newly created {@link JsonHeadersSanitizerBuilder}.
+     */
+    static JsonHeadersSanitizerBuilder builderForJson() {
+        return new JsonHeadersSanitizerBuilder();
+    }
+}

core/src/main/java/com/linecorp/armeria/common/JsonHeadersSanitizer.java

@@ -0,0 +1,72 @@
+/*
+ * 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;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Set;
+import java.util.function.Function;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+
+import io.netty.util.AsciiString;
+
+/**
+ * A sanitizer that sanitizes {@link HttpHeaders} and returns {@link JsonNode}.
+ */
+public final class JsonHeadersSanitizer implements HeadersSanitizer<JsonNode> {
+
+    static final HeadersSanitizer<JsonNode> INSTANCE = new JsonHeadersSanitizerBuilder().build();
+    private final Set<CharSequence> maskingHeaders;
+    private final Function<String, String> maskingFunction;
+    private final ObjectMapper objectMapper;
+
+    JsonHeadersSanitizer(Set<CharSequence> maskingHeaders, Function<String, String> maskingFunction,
+                         ObjectMapper objectMapper) {
+        this.maskingHeaders = maskingHeaders;
+        this.maskingFunction = maskingFunction;
+        this.objectMapper = objectMapper;
+    }
+
+    @Override
+    public JsonNode apply(RequestContext requestContext, HttpHeaders headers) {
+        final ObjectNode result = objectMapper.createObjectNode();
+        final Map<String, List<String>> headersWithValuesAsList = new LinkedHashMap<>();
+        for (Map.Entry<AsciiString, String> entry : headers) {
+            final String header = entry.getKey().toString().toLowerCase();
+            final String value = maskingHeaders.contains(header) ? maskingFunction.apply(entry.getValue())
+                                                                 : entry.getValue();
+            headersWithValuesAsList.computeIfAbsent(header, k -> new ArrayList<>()).add(value);
+        }
+
+        final Set<Entry<String, List<String>>> entries = headersWithValuesAsList.entrySet();
+        for (Map.Entry<String, List<String>> entry : entries) {
+            final String header = entry.getKey();
+            final List<String> values = entry.getValue();
+
+            result.put(header, values.size() > 1 ? values.toString() : values.get(0));
+        }
+
+        return result;
+    }
+}

core/src/main/java/com/linecorp/armeria/common/JsonHeadersSanitizerBuilder.java

@@ -0,0 +1,82 @@
+/*
+ * 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;
+
+import static java.util.Objects.requireNonNull;
+
+import java.util.Set;
+import java.util.function.Function;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+
+import com.linecorp.armeria.common.annotation.Nullable;
+import com.linecorp.armeria.internal.common.JacksonUtil;
+
+/**
+ * A builder implementation for {@link JsonHeadersSanitizer}.
+ */
+public final class JsonHeadersSanitizerBuilder extends AbstractHeadersSanitizerBuilder<JsonNode> {
+
+    @Nullable
+    private ObjectMapper objectMapper;
+
+    /**
+     * Sets the {@link Set} which includes headers to mask before logging.
+     */
+    @Override
+    public JsonHeadersSanitizerBuilder maskingHeaders(CharSequence... headers) {
+        return (JsonHeadersSanitizerBuilder) super.maskingHeaders(headers);
+    }
+
+    /**
+     * Sets the {@link Set} which includes headers to mask before logging.
+     */
+    @Override
+    public JsonHeadersSanitizerBuilder maskingHeaders(Iterable<? extends CharSequence> headers) {
+        return (JsonHeadersSanitizerBuilder) super.maskingHeaders(headers);
+    }
+
+    /**
+     * Sets the {@link Function} to use to mask headers before logging.
+     */
+    @Override
+    public JsonHeadersSanitizerBuilder maskingFunction(Function<String, String> maskingFunction) {
+        return (JsonHeadersSanitizerBuilder) super.maskingFunction(maskingFunction);
+    }
+
+    /**
+     * Sets the {@link ObjectMapper} that will be used to convert headers into a {@link JsonNode}.
+     */
+    public JsonHeadersSanitizerBuilder objectMapper(ObjectMapper objectMapper) {
+        this.objectMapper = requireNonNull(objectMapper, "objectMapper");
+        return this;
+    }
+
+    /**
+     * Returns a newly created JSON {@link HeadersSanitizer} based on the properties of this builder.
+     */
+    public JsonHeadersSanitizer build() {
+        final ObjectMapper objectMapper = this.objectMapper != null ?
+                                          this.objectMapper : JacksonUtil.newDefaultObjectMapper();
+
+        final Set<CharSequence> maskingHeaders = maskingHeaders();
+        return new JsonHeadersSanitizer(
+                !maskingHeaders.isEmpty() ? maskingHeaders : defaultMaskingHeaders(), maskingFunction(),
+                objectMapper);
+    }
+}

core/src/main/java/com/linecorp/armeria/common/TextHeadersSanitizer.java

@@ -0,0 +1,79 @@
+/*
+ * 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;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.Function;
+
+import io.netty.util.AsciiString;
+import scala.collection.generic.Sorted;
+
+/**
+ * A sanitizer that sanitizes {@link HttpHeaders} and returns {@link String}.
+ */
+public final class TextHeadersSanitizer implements HeadersSanitizer<String> {
+
+    static final HeadersSanitizer<String> INSTANCE = new TextHeadersSanitizerBuilder().build();
+
+    private final Set<CharSequence> maskingHeaders;
+
+    private final Function<String, String> maskingFunction;
+
+    TextHeadersSanitizer(Set<CharSequence> maskingHeaders, Function<String, String> maskingFunction) {
+        this.maskingHeaders = maskingHeaders;
+        this.maskingFunction = maskingFunction;
+    }
+
+    @Override
+    public String apply(RequestContext ctx, HttpHeaders headers) {
+        if (headers.isEmpty()) {
+            return headers.isEndOfStream() ? "[EOS]" : "[]";
+        }
+
+        final StringBuilder sb = new StringBuilder();
+        if (headers.isEndOfStream()) {
+            sb.append("[EOS], ");
+        } else {
+            sb.append('[');
+        }
+
+        final Map<String, List<String>> headersWithValuesAsList = new LinkedHashMap<>();
+        for (Map.Entry<AsciiString, String> entry : headers) {
+            final String header = entry.getKey().toString().toLowerCase();
+            final String value = maskingHeaders.contains(header) ? maskingFunction.apply(entry.getValue())
+                                                                 : entry.getValue();
+            headersWithValuesAsList.computeIfAbsent(header, k -> new ArrayList<>()).add(value);
+        }
+
+        final Set<Map.Entry<String, List<String>>> entries = headersWithValuesAsList.entrySet();
+        for (Map.Entry<String, List<String>> entry : entries) {
+            final String header = entry.getKey();
+            final List<String> values = entry.getValue();
+
+            sb.append(header).append('=')
+              .append(values.size() > 1 ? values.toString() : values.get(0)).append(", ");
+        }
+
+        sb.setCharAt(sb.length() - 2, ']');
+        return sb.substring(0, sb.length() - 1);
+    }
+}