Merge branch 'mask-sensible-headers' of https://github.com/seonwoo960…

…000/armeria

core/build.gradle

@@ -281,3 +281,7 @@ class PublicSuffixesTask extends DefaultTask {
 }
 
 task publicSuffixes(type: PublicSuffixesTask)
+
+test {
+    testLogging.showStandardStreams = true
+}

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

@@ -0,0 +1,78 @@
+/*
+ * 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.BiFunction;
+
+import com.google.common.collect.ImmutableSet;
+
+import com.linecorp.armeria.common.annotation.Nullable;
+
+/**
+ * A skeletal builder implementation for {@link HeadersSanitizer}.
+ */
+abstract class AbstractHeadersSanitizerBuilder<T> {
+
+    @Nullable
+    private BiFunction<? super RequestContext, ? super HttpHeaders, ? extends T> headersSanitizer;
+
+    private Set<String> headersMask = ImmutableSet.of();
+
+    /**
+     * Sets the {@link BiFunction} to use to sanitize headers before logging. It is common to have the
+     * {@link BiFunction} that removes sensitive headers, like Cookie, before logging.
+     */
+    public AbstractHeadersSanitizerBuilder<T> headersSanitizer(
+            BiFunction<? super RequestContext, ? super HttpHeaders, ? extends T> headersSanitizer) {
+        this.headersSanitizer = requireNonNull(headersSanitizer, "headersSanitizer");
+        return this;
+    }
+
+    /**
+     * Returns the {@link BiFunction} to use to sanitize headers before logging.
+     */
+    @Nullable
+    final BiFunction<? super RequestContext, ? super HttpHeaders, ? extends T> headersSanitizer() {
+        return headersSanitizer;
+    }
+
+    /**
+     * Sets the {@link Set} to use to mask headers before logging.
+     */
+    public AbstractHeadersSanitizerBuilder<T> headersMask(String... headers) {
+        headersMask = ImmutableSet.copyOf(requireNonNull(headers, "headers"));
+        return this;
+    }
+
+    /**
+     * Sets the {@link Set} to use to mask headers before logging.
+     */
+    public AbstractHeadersSanitizerBuilder<T> headersMask(Iterable<String> headers) {
+        headersMask = ImmutableSet.copyOf(requireNonNull(headers, "headers"));
+        return this;
+    }
+
+    /**
+     * Returns the {@link Set} to use to mask headers before logging.
+     */
+    final Set<String> headersMask() {
+        return headersMask;
+    }
+}

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

@@ -0,0 +1,60 @@
+/*
+ * 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();
+    }
+
+    /**
+     * Returns the sanitized {@link HttpHeaders}.
+     */
+    T sanitizeHeaders(RequestContext requestContext, HttpHeaders headers);
+}

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

@@ -0,0 +1,61 @@
+/*
+ * 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.BiFunction;
+
+import com.fasterxml.jackson.databind.JsonNode;
+
+import com.linecorp.armeria.common.annotation.Nullable;
+
+/**
+ * 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 static final String MASK = "****";
+
+    private final BiFunction<? super RequestContext, ? super HttpHeaders,
+            ? extends @Nullable JsonNode> headersSanitizer;
+
+    private final Set<String> headersMask;
+
+    JsonHeadersSanitizer(
+            BiFunction<? super RequestContext, ? super HttpHeaders, ? extends @Nullable JsonNode>
+                    headersSanitizer,
+            Set<String> headersMask) {
+        this.headersSanitizer = headersSanitizer;
+        this.headersMask = headersMask;
+    }
+
+    @Override
+    public JsonNode sanitizeHeaders(RequestContext ctx, HttpHeaders headers) {
+        final HttpHeadersBuilder builder = headers.toBuilder();
+        headers.forEach(
+                (name, value) -> builder.set(name, headersMask.contains(name.toString()) ? MASK : value));
+
+        return headersSanitizer.apply(ctx, builder.build());
+    }
+
+    @Override
+    public JsonNode apply(RequestContext requestContext, HttpHeaders headers) {
+        return null;
+    }
+}

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

@@ -0,0 +1,87 @@
+/*
+ * 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.google.common.base.MoreObjects.firstNonNull;
+import static java.util.Objects.requireNonNull;
+
+import java.util.Set;
+import java.util.function.BiFunction;
+
+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 ObjectMapper} that will be used to convert an object into a JSON format message.
+     */
+    public JsonHeadersSanitizerBuilder objectMapper(ObjectMapper objectMapper) {
+        this.objectMapper = requireNonNull(objectMapper, "objectMapper");
+        return this;
+    }
+
+    /**
+     * Sets the {@link BiFunction} to use to sanitize headers before logging.
+     */
+    @Override
+    public JsonHeadersSanitizerBuilder headersSanitizer(
+            BiFunction<? super RequestContext, ? super HttpHeaders, ? extends @Nullable JsonNode>
+                    headersSanitizer) {
+        return (JsonHeadersSanitizerBuilder) super.headersSanitizer(headersSanitizer);
+    }
+
+    /**
+     * Sets the {@link Set} to use to mask headers before logging.
+     */
+    @Override
+    public JsonHeadersSanitizerBuilder headersMask(String... headers) {
+        return (JsonHeadersSanitizerBuilder) super.headersMask(headers);
+    }
+
+    /**
+     * Sets the {@link Set} to use to mask headers before logging.
+     */
+    @Override
+    public JsonHeadersSanitizerBuilder headersMask(Iterable<String> headers) {
+        return (JsonHeadersSanitizerBuilder) super.headersMask(headers);
+    }
+
+    /**
+     * 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(firstNonNull(headersSanitizer(), defaultSanitizer(objectMapper)),
+                                        headersMask());
+    }
+
+    private static <T> BiFunction<? super RequestContext, T, JsonNode>
+    defaultSanitizer(ObjectMapper objectMapper) {
+        return (requestContext, obj) -> objectMapper.valueToTree(obj);
+    }
+}

core/src/main/java/com/linecorp/armeria/common/TextHeadersSanitizer.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.Set;
+import java.util.function.BiFunction;
+
+import com.linecorp.armeria.common.annotation.Nullable;
+
+/**
+ * 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 static final String MASK = "****";
+
+    private final BiFunction<? super RequestContext, ? super HttpHeaders,
+            ? extends @Nullable String> headersSanitizer;
+
+    private final Set<String> headersMask;
+
+    TextHeadersSanitizer(
+            BiFunction<? super RequestContext, ? super HttpHeaders, ? extends @Nullable String>
+                    headersSanitizer,
+            Set<String> headersMask) {
+        this.headersSanitizer = headersSanitizer;
+        this.headersMask = headersMask;
+    }
+
+    @Override
+    public String apply(RequestContext ctx, HttpHeaders headers) {
+        final HttpHeadersBuilder builder = headers.toBuilder();
+        headers.forEach(
+                (name, value) -> builder.set(name, headersMask.contains(name.toString()) ? MASK : value)
+        );
+        return headersSanitizer.apply(ctx, builder.build());
+    }
+
+    @Override
+    public String sanitizeHeaders(RequestContext requestContext, HttpHeaders headers) {
+        return null;
+    }
+}