Easily mask sensitive headers from request logs by adding HeadersSanitizer

core/src/main/java/com/linecorp/armeria/common/AbstractHeadersSanitizerBuilder.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 static java.util.Objects.requireNonNull;
+
+import java.util.Set;
+import java.util.function.Function;
+
+import com.google.common.collect.ImmutableSet;
+
+/**
+ * A skeletal builder implementation for {@link HeadersSanitizer}.
+ */
+abstract class AbstractHeadersSanitizerBuilder<T> {
+
+    private Set<String> maskHeaders = ImmutableSet.of();
+
+    private Function<String, String> mask = (header) -> "****";
+
+    /**
+     * Sets the {@link Set} which includes headers to mask before logging.
+     */
+    public AbstractHeadersSanitizerBuilder<T> maskHeaders(String... headers) {
+        maskHeaders = ImmutableSet.copyOf(requireNonNull(headers, "headers"));
+        return this;
+    }
+
+    /**
+     * Sets the {@link Set} which includes headers to mask before logging.
+     */
+    public AbstractHeadersSanitizerBuilder<T> maskHeaders(Iterable<String> headers) {
+        maskHeaders = ImmutableSet.copyOf(requireNonNull(headers, "headers"));
+        return this;
+    }
+
+    /**
+     * Returns the {@link Set} which includes headers to mask before logging.
+     */
+    final Set<String> maskHeaders() {
+        return maskHeaders;
+    }
+
+    /**
+     * Sets the {@link Function} to use to mask headers before logging.
+     */
+    public AbstractHeadersSanitizerBuilder<T> mask(Function<String, String> mask) {
+        this.mask = requireNonNull(mask, "mask");
+        return this;
+    }
+
+    /**
+     * Returns the {@link Function} to use to mask headers before logging.
+     */
+    final Function<String, String> mask() {
+        return mask;
+    }
+}

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,59 @@
+/*
+ * 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.Map;
+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<String> maskHeaders;
+    private final Function<String, String> mask;
+    private final ObjectMapper objectMapper;
+
+    JsonHeadersSanitizer(Set<String> maskHeaders, Function<String, String> mask, ObjectMapper objectMapper) {
+        this.maskHeaders = maskHeaders;
+        this.mask = mask;
+        this.objectMapper = objectMapper;
+    }
+
+    @Override
+    public JsonNode apply(RequestContext requestContext, HttpHeaders headers) {
+        final ObjectNode result = objectMapper.createObjectNode();
+        for (Map.Entry<AsciiString, String> e : headers) {
+            final String header = e.getKey().toString();
+            if (maskHeaders.contains(header)) {
+                result.put(header, mask.apply(e.getValue()));
+            } else {
+                result.put(header, e.getValue());
+            }
+        }
+
+        return result;
+    }
+}

core/src/main/java/com/linecorp/armeria/common/JsonHeadersSanitizerBuilder.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.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 maskHeaders(String... headers) {
+        return (JsonHeadersSanitizerBuilder) super.maskHeaders(headers);
+    }
+
+    /**
+     * Sets the {@link Set} which includes headers to mask before logging.
+     */
+    @Override
+    public JsonHeadersSanitizerBuilder maskHeaders(Iterable<String> headers) {
+        return (JsonHeadersSanitizerBuilder) super.maskHeaders(headers);
+    }
+
+    /**
+     * Sets the {@link Function} to use to mask headers before logging.
+     */
+    @Override
+    public JsonHeadersSanitizerBuilder mask(Function<String, String> mask) {
+        return (JsonHeadersSanitizerBuilder) super.mask(mask);
+    }
+
+    /**
+     * 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();
+
+        return new JsonHeadersSanitizer(maskHeaders(), mask(), objectMapper);
+    }
+}

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

@@ -0,0 +1,66 @@
+/*
+ * 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.Map;
+import java.util.Set;
+import java.util.function.Function;
+
+import io.netty.util.AsciiString;
+
+/**
+ * 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<String> maskHeaders;
+
+    private final Function<String, String> mask;
+
+    TextHeadersSanitizer(Set<String> maskHeaders, Function<String, String> mask) {
+        this.maskHeaders = maskHeaders;
+        this.mask = mask;
+    }
+
+    @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('[');
+        }
+
+        for (Map.Entry<AsciiString, String> e : headers) {
+            final String header = e.getKey().toString();
+            if (maskHeaders.contains(header)) {
+                sb.append(header).append('=').append(mask.apply(e.getValue())).append(", ");
+            } else {
+                sb.append(header).append('=').append(e.getValue()).append(", ");
+            }
+        }
+
+        sb.setCharAt(sb.length() - 2, ']');
+        return sb.substring(0, sb.length() - 1);
+    }
+}

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

@@ -0,0 +1,57 @@
+/*
+ * 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.Set;
+import java.util.function.Function;
+
+/**
+ * A builder implementation for {@link TextHeadersSanitizer}.
+ */
+public final class TextHeadersSanitizerBuilder extends AbstractHeadersSanitizerBuilder<String> {
+
+    /**
+     * Sets the {@link Set} which includes headers to mask before logging.
+     */
+    @Override
+    public TextHeadersSanitizerBuilder maskHeaders(String... headers) {
+        return (TextHeadersSanitizerBuilder) super.maskHeaders(headers);
+    }
+
+    /**
+     * Sets the {@link Set} which includes headers to mask before logging.
+     */
+    @Override
+    public TextHeadersSanitizerBuilder maskHeaders(Iterable<String> headers) {
+        return (TextHeadersSanitizerBuilder) super.maskHeaders(headers);
+    }
+
+    /**
+     * Sets the {@link Function} to use to mask headers before logging.
+     */
+    @Override
+    public TextHeadersSanitizerBuilder mask(Function<String, String> mask) {
+        return (TextHeadersSanitizerBuilder) super.mask(mask);
+    }
+
+    /**
+     * Returns a newly created text {@link HeadersSanitizer} based on the properties of this builder.
+     */
+    public TextHeadersSanitizer build() {
+        return new TextHeadersSanitizer(maskHeaders(), mask());
+    }
+}