Add composite MdcEntryWriter to filter MDC keys with regex (#974)

jug · web-flow · commit 94ccc40345d5 · 2024-07-27T12:25:29.000-07:00
diff --git a/README.md b/README.md
@@ -1033,7 +1033,7 @@ specify `<mdcKeyFieldName>mdcKeyName=fieldName</mdcKeyFieldName>`:
 You can also manipulate the MDC entry values written to the JSON output.
 By default, no manipulations are done and all MDC entry values are written as text.
 
-Currently, MDC entry writers for the following value types are supported:
+Currently, the following MDC entry writers are supported:
 
 ```xml
 <encoder class="net.logstash.logback.encoder.LogstashEncoder">
@@ -1057,6 +1057,23 @@ Currently, MDC entry writers for the following value types are supported:
         e.g. Writes true instead of "true"
     -->
     <mdcEntryWriter class="net.logstash.logback.composite.loggingevent.mdc.BooleanMdcEntryWriter"/>
+
+    <!--
+        Composite MDC entry writer that delegates writing MDC entries to a list of `MdcEntryWriter`
+        if the key of an MDC entry does match a given include pattern and does not match a given
+        exclude pattern.
+                
+        Omitting the 'includeMdcKeyPattern' means to include all MDC keys.
+        Omitting the 'excludeMdcKeyPattern' means to exclude no MDC keys.
+        
+        The elements with the key patterns are optional. If provided, use a regular expression
+        with the syntax of `java.util.regex.Pattern`.
+    -->
+    <mdcEntryWriter class="net.logstash.logback.composite.loggingevent.mdc.RegexFilteringMdcEntryWriter">
+        <includeMdcKeyPattern>keyPatternToInclude</includeMdcKeyPattern>
+        <excludeMdcKeyPattern>keyPatternToExclude</excludeMdcKeyPattern>
+        <mdcEntryWriter class="net.logstash.logback.composite.loggingevent.LongMdcEntryWriter"/>
+    </mdcEntryWriter>
 </encoder>
 ```
 
diff --git a/src/main/java/net/logstash/logback/composite/loggingevent/mdc/RegexFilteringMdcEntryWriter.java b/src/main/java/net/logstash/logback/composite/loggingevent/mdc/RegexFilteringMdcEntryWriter.java
@@ -0,0 +1,87 @@
+/*
+ * Copyright 2013-2023 the original author or authors.
+ *
+ * Licensed 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
+ *
+ *      http://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 net.logstash.logback.composite.loggingevent.mdc;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.regex.Pattern;
+
+import com.fasterxml.jackson.core.JsonGenerator;
+
+/**
+ * Writes MDC entries by delegating to other instances of {@link MdcEntryWriter} if MDC key matches the given include
+ * and exclude pattern.
+ *
+ * <ol>
+ * <li>An MDC entry is written if the MDC key does match the {@link #includeMdcKeyPattern} AND does not match
+ *     the {@link #excludeMdcKeyPattern}.</li>
+ * <li>Omitting a {@link #includeMdcKeyPattern} means to include all MDC keys.</li>
+ * <li>Omitting a {@link #excludeMdcKeyPattern} means to exclude no MDC keys.</li>
+ * </ol>
+ */
+public class RegexFilteringMdcEntryWriter implements MdcEntryWriter {
+
+    private Pattern includeMdcKeyPattern;
+    private Pattern excludeMdcKeyPattern;
+    private final List<MdcEntryWriter> mdcEntryWriters = new ArrayList<>();
+
+    @Override
+    public boolean writeMdcEntry(JsonGenerator generator, String fieldName, String mdcKey, String mdcValue) throws IOException {
+        if (shouldWrite(mdcKey)) {
+            for (MdcEntryWriter mdcEntryWriter : this.mdcEntryWriters) {
+                if (mdcEntryWriter.writeMdcEntry(generator, fieldName, mdcKey, mdcValue)) {
+                    return true;
+                }
+            }
+        }
+
+        return false;
+    }
+
+    public Pattern getIncludeMdcKeyPattern() {
+        return includeMdcKeyPattern;
+    }
+    public void setIncludeMdcKeyPattern(String includeMdcKeyPattern) {
+        this.includeMdcKeyPattern = Pattern.compile(includeMdcKeyPattern);
+    }
+
+    public Pattern getExcludeMdcKeyPattern() {
+        return excludeMdcKeyPattern;
+    }
+    public void setExcludeMdcKeyPattern(String excludeMdcKeyPattern) {
+        this.excludeMdcKeyPattern = Pattern.compile(excludeMdcKeyPattern);
+    }
+
+    public List<MdcEntryWriter> getMdcEntryWriters() {
+        return Collections.unmodifiableList(mdcEntryWriters);
+    }
+    public void addMdcEntryWriter(MdcEntryWriter mdcEntryWriter) {
+        this.mdcEntryWriters.add(mdcEntryWriter);
+    }
+
+    /** Returns true if passed MDC key should be written to the JSON output. */
+    private boolean shouldWrite(String key) {
+        if (this.mdcEntryWriters.isEmpty()) {
+            return false;
+        }
+        boolean includeKey = includeMdcKeyPattern == null || includeMdcKeyPattern.matcher(key).matches();
+        boolean excludeKey = excludeMdcKeyPattern != null && excludeMdcKeyPattern.matcher(key).matches();
+        return includeKey && !excludeKey;
+    }
+
+}
diff --git a/src/test/java/net/logstash/logback/ConfigurationTest.java b/src/test/java/net/logstash/logback/ConfigurationTest.java
@@ -52,7 +52,9 @@
 import net.logstash.logback.composite.loggingevent.SequenceJsonProvider;
 import net.logstash.logback.composite.loggingevent.StackTraceJsonProvider;
 import net.logstash.logback.composite.loggingevent.TagsJsonProvider;
+import net.logstash.logback.composite.loggingevent.mdc.BooleanMdcEntryWriter;
 import net.logstash.logback.composite.loggingevent.mdc.LongMdcEntryWriter;
+import net.logstash.logback.composite.loggingevent.mdc.RegexFilteringMdcEntryWriter;
 import net.logstash.logback.encoder.LoggingEventCompositeJsonEncoder;
 import net.logstash.logback.marker.Markers;
 import net.logstash.logback.stacktrace.ShortenedThrowableConverter;
@@ -180,8 +182,17 @@ private void verifyCommonProviders(List<JsonProvider<ILoggingEvent>> providers)
         assertThat(mdcJsonProvider).isNotNull();
         assertThat(mdcJsonProvider.getIncludeMdcKeyNames()).containsExactly("included");
         assertThat(mdcJsonProvider.getMdcKeyFieldNames()).containsOnly(entry("key", "renamedKey"));
-        assertThat(mdcJsonProvider.getMdcEntryWriters()).hasSize(1);
+        assertThat(mdcJsonProvider.getMdcEntryWriters()).hasSize(2);
         assertThat(mdcJsonProvider.getMdcEntryWriters()).element(0).isExactlyInstanceOf(LongMdcEntryWriter.class);
+        assertThat(mdcJsonProvider.getMdcEntryWriters()).element(1).isExactlyInstanceOf(RegexFilteringMdcEntryWriter.class);
+        RegexFilteringMdcEntryWriter regexFilteringMdcEntryWriter =
+                (RegexFilteringMdcEntryWriter) mdcJsonProvider.getMdcEntryWriters().get(1);
+        assertThat(regexFilteringMdcEntryWriter.getIncludeMdcKeyPattern()).isNotNull();
+        assertThat("include").matches(regexFilteringMdcEntryWriter.getIncludeMdcKeyPattern());
+        assertThat(regexFilteringMdcEntryWriter.getExcludeMdcKeyPattern()).isNotNull();
+        assertThat("exclude").matches(regexFilteringMdcEntryWriter.getExcludeMdcKeyPattern());
+        assertThat(regexFilteringMdcEntryWriter.getMdcEntryWriters()).hasSize(1);
+        assertThat(regexFilteringMdcEntryWriter.getMdcEntryWriters()).element(0).isExactlyInstanceOf(BooleanMdcEntryWriter.class);
 
         KeyValuePairsJsonProvider keyValuePairsJsonProvider = getInstance(providers, KeyValuePairsJsonProvider.class);
         assertThat(keyValuePairsJsonProvider).isNotNull();
diff --git a/src/test/java/net/logstash/logback/composite/loggingevent/mdc/RegexFilteringMdcEntryWriterTest.java b/src/test/java/net/logstash/logback/composite/loggingevent/mdc/RegexFilteringMdcEntryWriterTest.java
@@ -0,0 +1,92 @@
+/*
+ * Copyright 2013-2023 the original author or authors.
+ *
+ * Licensed 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
+ *
+ *      http://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 net.logstash.logback.composite.loggingevent.mdc;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.when;
+
+import java.io.IOException;
+
+import com.fasterxml.jackson.core.JsonGenerator;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+
+@ExtendWith(MockitoExtension.class)
+class RegexFilteringMdcEntryWriterTest {
+
+    private final RegexFilteringMdcEntryWriter mdcEntryWriter = new RegexFilteringMdcEntryWriter();
+
+    @Mock
+    private JsonGenerator generator;
+    @Mock
+    private MdcEntryWriter mockedMdcEntryWriter;
+
+    @Test
+    void noIncludeAndNoExcludePattern() throws IOException {
+        mockMdcEntryWriter();
+
+        boolean result = mdcEntryWriter.writeMdcEntry(generator, "field", "key", "value");
+
+        assertThat(result).isTrue();
+        verify(mockedMdcEntryWriter).writeMdcEntry(generator, "field", "key", "value");
+        verifyNoMoreInteractions(mockedMdcEntryWriter);
+    }
+
+    @Test
+    void noMdcEntryWriter() throws IOException {
+        boolean result = mdcEntryWriter.writeMdcEntry(generator, "field", "key", "value");
+
+        assertThat(result).isFalse();
+        verifyNoMoreInteractions(mockedMdcEntryWriter);
+    }
+
+    @ParameterizedTest
+    @CsvSource({
+            "include,excl.*,true",
+            "include,.*lude,false",
+            "exclude,excl.*,false",
+            "other,excl.*,false"
+    })
+    void includeAndExcludePattern(String mdcKey, String excludePattern, boolean entryWritten) throws IOException {
+        if (entryWritten) {
+            mockMdcEntryWriter();
+        }
+        mdcEntryWriter.setIncludeMdcKeyPattern("incl.*");
+        mdcEntryWriter.setExcludeMdcKeyPattern(excludePattern);
+
+        boolean result = mdcEntryWriter.writeMdcEntry(generator, "field", mdcKey, "value");
+
+        assertThat(result).isEqualTo(entryWritten);
+        if (entryWritten) {
+            verify(mockedMdcEntryWriter).writeMdcEntry(generator, "field", mdcKey, "value");
+        }
+        verifyNoMoreInteractions(mockedMdcEntryWriter);
+    }
+
+    private void mockMdcEntryWriter() throws IOException {
+        when(mockedMdcEntryWriter.writeMdcEntry(eq(generator), anyString(), anyString(), anyString())).thenReturn(true);
+        mdcEntryWriter.addMdcEntryWriter(mockedMdcEntryWriter);
+    }
+
+}
diff --git a/src/test/resources/logback-test.xml b/src/test/resources/logback-test.xml
@@ -51,6 +51,11 @@
             <includeMdcKeyName>included</includeMdcKeyName>
             <mdcKeyFieldName>key=renamedKey</mdcKeyFieldName>
             <mdcEntryWriter class="net.logstash.logback.composite.loggingevent.mdc.LongMdcEntryWriter"/>
+            <mdcEntryWriter class="net.logstash.logback.composite.loggingevent.mdc.RegexFilteringMdcEntryWriter">
+                <includeMdcKeyPattern>incl.*</includeMdcKeyPattern>
+                <excludeMdcKeyPattern>excl.*</excludeMdcKeyPattern>
+                <mdcEntryWriter class="net.logstash.logback.composite.loggingevent.mdc.BooleanMdcEntryWriter"/>
+            </mdcEntryWriter>
             <includeKeyValueKeyName>included</includeKeyValueKeyName>
             <keyValueKeyFieldName>key=renamedKey</keyValueKeyFieldName>
             <customFields>{"customName":"customValue"}</customFields>
@@ -142,6 +147,11 @@
                     <includeMdcKeyName>included</includeMdcKeyName>
                     <mdcKeyFieldName>key=renamedKey</mdcKeyFieldName>
                     <mdcEntryWriter class="net.logstash.logback.composite.loggingevent.mdc.LongMdcEntryWriter"/>
+                    <mdcEntryWriter class="net.logstash.logback.composite.loggingevent.mdc.RegexFilteringMdcEntryWriter">
+                        <includeMdcKeyPattern>incl.*</includeMdcKeyPattern>
+                        <excludeMdcKeyPattern>excl.*</excludeMdcKeyPattern>
+                        <mdcEntryWriter class="net.logstash.logback.composite.loggingevent.mdc.BooleanMdcEntryWriter"/>
+                    </mdcEntryWriter>
                 </mdc>
                 <keyValuePairs>
                     <includeKeyName>included</includeKeyName>
