Easily mask sensitive headers from request logs by adding HeadersSanitizer

core/src/main/java/com/linecorp/armeria/common/AbstractHeadersSanitizerBuilder.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 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 io.netty.util.AsciiString;
+
+/**
+ * A skeletal builder implementation for {@link HeadersSanitizer}.
+ */
+public abstract class AbstractHeadersSanitizerBuilder<T> {
+
+    private static final Set<AsciiString> DEFAULT_MASKING_HEADERS =
+            ImmutableSet.of(HttpHeaderNames.AUTHORIZATION, HttpHeaderNames.SET_COOKIE);
+
+    private final Set<AsciiString> maskingHeaders = new HashSet<>();
+
+    private Function<String, String> maskingFunction = header -> "****";
+
+    /**
+     * Sets the headers to mask before logging.
+     */
+    public AbstractHeadersSanitizerBuilder<T> maskingHeaders(CharSequence... headers) {
+        requireNonNull(headers, "headers");
+        return maskingHeaders(ImmutableSet.copyOf(headers));
+    }
+
+    /**
+     * Sets the headers to mask before logging.
+     */
+    public AbstractHeadersSanitizerBuilder<T> maskingHeaders(Iterable<? extends CharSequence> headers) {
+        requireNonNull(headers, "headers");
+        headers.forEach(header -> maskingHeaders.add(AsciiString.of(header).toLowerCase()));
+        return this;
+    }
+
+    final Set<AsciiString> maskingHeaders() {
+        if (!maskingHeaders.isEmpty()) {
+            return ImmutableSet.copyOf(maskingHeaders);
+        }
+        return DEFAULT_MASKING_HEADERS;
+    }
+
+    /**
+     * Sets the {@link Function} to use to maskFunction headers before logging.
+     */
+    public AbstractHeadersSanitizerBuilder<T> maskingFunction(Function<String, String> maskingFunction) {
+        this.maskingFunction = requireNonNull(maskingFunction, "maskingFunction");
+        return this;
+    }
+
+    /**
+     * Returns the {@link Function} to use to mask headers before logging.
+     */
+    final Function<String, String> maskingFunction() {
+        return maskingFunction;
+    }
+}

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

@@ -0,0 +1,55 @@
+/*
+ * 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}.
+ */
+@FunctionalInterface
+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,56 @@
+/*
+ * 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 com.linecorp.armeria.common.TextHeadersSanitizer.maskHeaders;
+
+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}.
+ */
+final class JsonHeadersSanitizer implements HeadersSanitizer<JsonNode> {
+
+    static final HeadersSanitizer<JsonNode> INSTANCE = new JsonHeadersSanitizerBuilder().build();
+    private final Set<AsciiString> maskingHeaders;
+    private final Function<String, String> maskingFunction;
+    private final ObjectMapper objectMapper;
+
+    JsonHeadersSanitizer(Set<AsciiString> 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();
+        maskHeaders(headers, maskingHeaders, maskingFunction,
+                    (header, values) -> result.put(header.toString(), values.size() > 1 ?
+                                                                      values.toString() : values.get(0)));
+
+        return result;
+    }
+}

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

@@ -0,0 +1,68 @@
+/*
+ * 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.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 JSON {@link HeadersSanitizer}.
+ */
+public final class JsonHeadersSanitizerBuilder extends AbstractHeadersSanitizerBuilder<JsonNode> {
+
+    @Nullable
+    private ObjectMapper objectMapper;
+
+    @Override
+    public JsonHeadersSanitizerBuilder maskingHeaders(CharSequence... headers) {
+        return (JsonHeadersSanitizerBuilder) super.maskingHeaders(headers);
+    }
+
+    @Override
+    public JsonHeadersSanitizerBuilder maskingHeaders(Iterable<? extends CharSequence> headers) {
+        return (JsonHeadersSanitizerBuilder) super.maskingHeaders(headers);
+    }
+
+    @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 HeadersSanitizer<JsonNode> build() {
+        final ObjectMapper objectMapper = this.objectMapper != null ?
+                                          this.objectMapper : JacksonUtil.newDefaultObjectMapper();
+        return new JsonHeadersSanitizer(maskingHeaders(), maskingFunction(), objectMapper);
+    }
+}

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

@@ -0,0 +1,85 @@
+/*
+ * 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.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.BiConsumer;
+import java.util.function.Function;
+
+import io.netty.util.AsciiString;
+
+/**
+ * A sanitizer that sanitizes {@link HttpHeaders} and returns {@link String}.
+ */
+final class TextHeadersSanitizer implements HeadersSanitizer<String> {
+
+    static final HeadersSanitizer<String> INSTANCE = new TextHeadersSanitizerBuilder().build();
+
+    private final Set<AsciiString> maskingHeaders;
+
+    private final Function<String, String> maskingFunction;
+
+    TextHeadersSanitizer(Set<AsciiString> 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('[');
+        }
+
+        maskHeaders(headers, maskingHeaders, maskingFunction,
+                    (header, values) -> 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);
+    }
+
+    static void maskHeaders(
+            HttpHeaders headers, Set<AsciiString> maskingHeaders, Function<String, String> maskingFunction,
+            final BiConsumer<AsciiString, List<String>> consumer) {
+        final Map<AsciiString, List<String>> headersWithValuesAsList = new LinkedHashMap<>();
+        for (Map.Entry<AsciiString, String> entry : headers) {
+            final AsciiString header = entry.getKey().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<AsciiString, List<String>>> entries = headersWithValuesAsList.entrySet();
+        for (Map.Entry<AsciiString, List<String>> entry : entries) {
+            final AsciiString header = entry.getKey();
+            final List<String> values = entry.getValue();
+            consumer.accept(header, values);
+        }
+    }
+}

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

@@ -0,0 +1,47 @@
+/*
+ * 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.Function;
+
+/**
+ * A builder implementation for Text {@link HeadersSanitizer}.
+ */
+public final class TextHeadersSanitizerBuilder extends AbstractHeadersSanitizerBuilder<String> {
+
+    @Override
+    public TextHeadersSanitizerBuilder maskingHeaders(CharSequence... headers) {
+        return (TextHeadersSanitizerBuilder) super.maskingHeaders(headers);
+    }
+
+    @Override
+    public TextHeadersSanitizerBuilder maskingHeaders(Iterable<? extends CharSequence> headers) {
+        return (TextHeadersSanitizerBuilder) super.maskingHeaders(headers);
+    }
+
+    @Override
+    public TextHeadersSanitizerBuilder maskingFunction(Function<String, String> maskingFunction) {
+        return (TextHeadersSanitizerBuilder) super.maskingFunction(maskingFunction);
+    }
+
+    /**
+     * Returns a newly created text {@link HeadersSanitizer} based on the properties of this builder.
+     */
+    public HeadersSanitizer<String> build() {
+        return new TextHeadersSanitizer(maskingHeaders(), maskingFunction());
+    }
+}