From 02d06af3fc21193d3c6604f51e74c10ea70fc350 Mon Sep 17 00:00:00 2001
From: Ehren Kret <ehren@signal.org>
Date: Tue, 11 May 2021 14:08:20 -0500
Subject: [PATCH] Replace use of MDC with custom fields on the appender

---
 service/config/sample.yml                     |  2 --
 .../WhisperServerConfiguration.java           |  8 ------
 .../textsecuregcm/WhisperServerService.java   | 11 --------
 .../LogstashTcpSocketAppenderFactory.java     | 28 ++++++++++++++++---
 4 files changed, 24 insertions(+), 25 deletions(-)

diff --git a/service/config/sample.yml b/service/config/sample.yml
index 361270c78..dee867d48 100644
--- a/service/config/sample.yml
+++ b/service/config/sample.yml
@@ -1,5 +1,3 @@
-environment: # value
-
 twilio: # Twilio gateway configuration
   accountId: 
   accountToken: 
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerConfiguration.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerConfiguration.java
index 0d79d25bd..1551919dc 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerConfiguration.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerConfiguration.java
@@ -52,10 +52,6 @@ import org.whispersystems.websocket.configuration.WebSocketConfiguration;
 /** @noinspection MismatchedQueryAndUpdateOfCollection, WeakerAccess */
 public class WhisperServerConfiguration extends Configuration {
 
-  @NotEmpty
-  @JsonProperty
-  private String environment;
-
   @NotNull
   @Valid
   @JsonProperty
@@ -263,10 +259,6 @@ public class WhisperServerConfiguration extends Configuration {
 
   private Map<String, String> transparentDataIndex = new HashMap<>();
 
-  public String getEnvironment() {
-    return environment;
-  }
-
   public RecaptchaConfiguration getRecaptchaConfiguration() {
     return recaptcha;
   }
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
index 6672fb07f..d2b5f9ef9 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -43,8 +43,6 @@ import io.micrometer.core.instrument.Metrics;
 import io.micrometer.core.instrument.distribution.DistributionStatisticConfig;
 import io.micrometer.wavefront.WavefrontConfig;
 import io.micrometer.wavefront.WavefrontMeterRegistry;
-import java.net.InetAddress;
-import java.net.UnknownHostException;
 import java.net.http.HttpClient;
 import java.time.Duration;
 import java.util.ArrayList;
@@ -67,7 +65,6 @@ import org.jdbi.v3.core.Jdbi;
 import org.signal.zkgroup.ServerSecretParams;
 import org.signal.zkgroup.auth.ServerZkAuthOperations;
 import org.signal.zkgroup.profiles.ServerZkProfileOperations;
-import org.slf4j.MDC;
 import org.whispersystems.dispatch.DispatchManager;
 import org.whispersystems.textsecuregcm.auth.AccountAuthenticator;
 import org.whispersystems.textsecuregcm.auth.CertificateGenerator;
@@ -231,14 +228,6 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
   @Override
   public void run(WhisperServerConfiguration config, Environment environment)
       throws Exception {
-    try {
-      MDC.put("host", InetAddress.getLocalHost().getHostName());
-    } catch (UnknownHostException e) {
-      MDC.put("host", "unknown");
-    }
-    MDC.put("service", "chat");
-    MDC.put("ddsource", "logstash");
-    MDC.put("ddtags", "env:" + config.getEnvironment());
     SharedMetricRegistries.add(Constants.METRICS_NAME, environment.metrics());
 
     final WavefrontConfig wavefrontConfig = new WavefrontConfig() {
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/metrics/LogstashTcpSocketAppenderFactory.java b/service/src/main/java/org/whispersystems/textsecuregcm/metrics/LogstashTcpSocketAppenderFactory.java
index d56f65530..568fc05c1 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/metrics/LogstashTcpSocketAppenderFactory.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/metrics/LogstashTcpSocketAppenderFactory.java
@@ -13,10 +13,15 @@ import ch.qos.logback.core.encoder.LayoutWrappingEncoder;
 import ch.qos.logback.core.net.ssl.SSLConfiguration;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.annotation.JsonTypeName;
+import com.fasterxml.jackson.databind.node.JsonNodeFactory;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.fasterxml.jackson.databind.node.TextNode;
 import io.dropwizard.logging.AbstractAppenderFactory;
 import io.dropwizard.logging.async.AsyncAppenderFactory;
 import io.dropwizard.logging.filter.LevelFilterFactory;
 import io.dropwizard.logging.layout.LayoutFactory;
+import java.net.InetAddress;
+import java.net.UnknownHostException;
 import java.time.Duration;
 import javax.validation.constraints.NotEmpty;
 import net.logstash.logback.appender.LogstashTcpSocketAppender;
@@ -25,15 +30,13 @@ import net.logstash.logback.encoder.LogstashEncoder;
 @JsonTypeName("logstashtcpsocket")
 public class LogstashTcpSocketAppenderFactory extends AbstractAppenderFactory<ILoggingEvent> {
 
-  @NotEmpty
   private String destination;
-
   private Duration keepAlive = Duration.ofSeconds(20);
-
-  @NotEmpty
   private String apiKey;
+  private String environment;
 
   @JsonProperty
+  @NotEmpty
   public String getDestination() {
     return destination;
   }
@@ -44,10 +47,17 @@ public class LogstashTcpSocketAppenderFactory extends AbstractAppenderFactory<IL
   }
 
   @JsonProperty
+  @NotEmpty
   public String getApiKey() {
     return apiKey;
   }
 
+  @JsonProperty
+  @NotEmpty
+  public String getEnvironment() {
+    return environment;
+  }
+
   @Override
   public Appender<ILoggingEvent> build(
       final LoggerContext context,
@@ -65,6 +75,16 @@ public class LogstashTcpSocketAppenderFactory extends AbstractAppenderFactory<IL
     appender.setKeepAliveDuration(new ch.qos.logback.core.util.Duration(keepAlive.toMillis()));
 
     final LogstashEncoder encoder = new LogstashEncoder();
+    final ObjectNode customFieldsNode = new ObjectNode(JsonNodeFactory.instance);
+    try {
+      customFieldsNode.set("host", TextNode.valueOf(InetAddress.getLocalHost().getHostName()));
+    } catch (UnknownHostException e) {
+      customFieldsNode.set("host", TextNode.valueOf("unknown"));
+    }
+    customFieldsNode.set("service", TextNode.valueOf("chat"));
+    customFieldsNode.set("ddsource", TextNode.valueOf("logstash"));
+    customFieldsNode.set("ddtags", TextNode.valueOf("env:" + environment));
+    encoder.setCustomFields(customFieldsNode.toString());
     final LayoutWrappingEncoder<ILoggingEvent> prefix = new LayoutWrappingEncoder<>();
     final PatternLayout layout = new PatternLayout();
     layout.setPattern(String.format("%s ", apiKey));