diff --git a/loki-log4j2-appender/src/main/java/pl/tkowalcz/tjahzi/log4j2/LokiAppender.java b/loki-log4j2-appender/src/main/java/pl/tkowalcz/tjahzi/log4j2/LokiAppender.java index 3822229..7a82bdc 100644 --- a/loki-log4j2-appender/src/main/java/pl/tkowalcz/tjahzi/log4j2/LokiAppender.java +++ b/loki-log4j2-appender/src/main/java/pl/tkowalcz/tjahzi/log4j2/LokiAppender.java @@ -20,11 +20,13 @@ import pl.tkowalcz.tjahzi.http.HttpClientFactory; import pl.tkowalcz.tjahzi.http.NettyHttpClient; import java.io.Serializable; -import java.util.Arrays; import java.util.HashMap; import java.util.Map; import java.util.Objects; import java.util.concurrent.TimeUnit; +import java.util.stream.Stream; + +import static java.util.Arrays.stream; /** * Loki Appender. @@ -81,9 +83,16 @@ public class LokiAppender extends AbstractAppender { .withRequestTimeoutMillis(readTimeoutMillis) .build(); + String[] additionalHeaders = stream(headers) + .flatMap(header -> Stream.of(header.getName(), header.getValue())) + .toArray(String[]::new); + NettyHttpClient httpClient = HttpClientFactory .defaultFactory() - .getHttpClient(configurationBuilder); + .getHttpClient( + configurationBuilder, + additionalHeaders + ); LoggingSystem loggingSystem = new TjahziInitializer().createLoggingSystem( httpClient, @@ -92,17 +101,7 @@ public class LokiAppender extends AbstractAppender { ); HashMap lokiLabels = Maps.newHashMap(); - Arrays.stream(labels).forEach(property -> { - String value = getConfiguration() - .getStrSubstitutor() - .replace(property.getValue()); - - lokiLabels.put( - property.getName(), - value - ); - } - ); + stream(labels).forEach(__ -> lokiLabels.put(__.getName(), __.getValue())); return new LokiAppender( getName(), @@ -218,7 +217,10 @@ public class LokiAppender extends AbstractAppender { ); Objects.requireNonNull(layout, "layout"); - this.appenderLogic = new AppenderLogic(loggingSystem, lokiLabels); + this.appenderLogic = new AppenderLogic( + loggingSystem, + lokiLabels + ); } @Override diff --git a/tjahzi-core/pom.xml b/tjahzi-core/pom.xml index 9ba9453..7e1ca2b 100644 --- a/tjahzi-core/pom.xml +++ b/tjahzi-core/pom.xml @@ -48,6 +48,13 @@ test + + com.github.tomakehurst + wiremock-jre8 + 2.27.2 + test + + org.junit.jupiter diff --git a/tjahzi-core/src/main/java/pl/tkowalcz/tjahzi/http/HttpClientFactory.java b/tjahzi-core/src/main/java/pl/tkowalcz/tjahzi/http/HttpClientFactory.java index 5db0ad6..bb4b57a 100644 --- a/tjahzi-core/src/main/java/pl/tkowalcz/tjahzi/http/HttpClientFactory.java +++ b/tjahzi-core/src/main/java/pl/tkowalcz/tjahzi/http/HttpClientFactory.java @@ -4,8 +4,13 @@ public class HttpClientFactory { private static final HttpClientFactory INSTANCE = new HttpClientFactory(); - public NettyHttpClient getHttpClient(ClientConfiguration clientConfiguration) { - return new NettyHttpClient(clientConfiguration); + public NettyHttpClient getHttpClient( + ClientConfiguration clientConfiguration, + String[] additionalHeaders) { + return new NettyHttpClient( + clientConfiguration, + additionalHeaders + ); } public static HttpClientFactory defaultFactory() { diff --git a/tjahzi-core/src/main/java/pl/tkowalcz/tjahzi/http/NettyHttpClient.java b/tjahzi-core/src/main/java/pl/tkowalcz/tjahzi/http/NettyHttpClient.java index 0a44bbd..8a89ba3 100644 --- a/tjahzi-core/src/main/java/pl/tkowalcz/tjahzi/http/NettyHttpClient.java +++ b/tjahzi-core/src/main/java/pl/tkowalcz/tjahzi/http/NettyHttpClient.java @@ -15,15 +15,20 @@ import java.util.concurrent.ThreadFactory; public class NettyHttpClient implements AutoCloseable { + public static final String PROTOBUF_MIME_TYPE = "application/x-protobuf"; + private final EventLoopGroup group; private final String logEndpoint; private final Channel channel; private final String host; + private final HttpHeaders headers; private final Snappy snappy = new Snappy(); - public NettyHttpClient(ClientConfiguration clientConfiguration) { + public NettyHttpClient( + ClientConfiguration clientConfiguration, + String[] additionalHeaders) { host = clientConfiguration.getHost(); logEndpoint = clientConfiguration.getLogEndpoint(); @@ -46,6 +51,11 @@ public class NettyHttpClient implements AutoCloseable { .connect() .syncUninterruptibly() .channel(); + + headers = new ReadOnlyHttpHeaders( + true, + additionalHeaders + ); } public void log(ByteBuf dataBuffer) { @@ -59,9 +69,11 @@ public class NettyHttpClient implements AutoCloseable { output ); - request.headers().set(HttpHeaderNames.CONTENT_TYPE, "application/x-protobuf"); - request.headers().set(HttpHeaderNames.CONTENT_LENGTH, output.readableBytes()); - request.headers().set(HttpHeaderNames.HOST, host); + request.headers() + .add(headers) + .set(HttpHeaderNames.CONTENT_TYPE, PROTOBUF_MIME_TYPE) + .set(HttpHeaderNames.CONTENT_LENGTH, output.readableBytes()) + .set(HttpHeaderNames.HOST, host); channel.writeAndFlush(request); } diff --git a/tjahzi-core/src/test/java/pl/tkowalcz/thjazi/LoggingSystemSanityCheckTest.java b/tjahzi-core/src/test/java/pl/tkowalcz/thjazi/LoggingSystemSanityCheckTest.java index b575a81..9c181f3 100644 --- a/tjahzi-core/src/test/java/pl/tkowalcz/thjazi/LoggingSystemSanityCheckTest.java +++ b/tjahzi-core/src/test/java/pl/tkowalcz/thjazi/LoggingSystemSanityCheckTest.java @@ -54,7 +54,10 @@ class LoggingSystemSanityCheckTest { .build(); NettyHttpClient httpClient = HttpClientFactory.defaultFactory() - .getHttpClient(clientConfiguration); + .getHttpClient( + clientConfiguration, + new String[0] + ); TjahziInitializer initializer = new TjahziInitializer(); LoggingSystem loggingSystem = initializer.createLoggingSystem( diff --git a/tjahzi-core/src/test/java/pl/tkowalcz/tjahzi/HeadersTest.java b/tjahzi-core/src/test/java/pl/tkowalcz/tjahzi/HeadersTest.java new file mode 100644 index 0000000..85164e6 --- /dev/null +++ b/tjahzi-core/src/test/java/pl/tkowalcz/tjahzi/HeadersTest.java @@ -0,0 +1,198 @@ +package pl.tkowalcz.tjahzi; + +import com.github.tomakehurst.wiremock.WireMockServer; +import org.awaitility.Durations; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import pl.tkowalcz.tjahzi.http.ClientConfiguration; +import pl.tkowalcz.tjahzi.http.HttpClientFactory; +import pl.tkowalcz.tjahzi.http.NettyHttpClient; + +import java.util.Map; +import java.util.concurrent.TimeUnit; + +import static com.github.tomakehurst.wiremock.client.WireMock.*; +import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig; +import static org.awaitility.Awaitility.await; + +class HeadersTest { + + private WireMockServer wireMockServer; + + private TjahziInitializer initializer; + private LoggingSystem loggingSystem; + + @BeforeEach + void setUp() { + wireMockServer = new WireMockServer( + wireMockConfig() + .dynamicPort() + .dynamicHttpsPort() + ); + + wireMockServer.start(); + + initializer = new TjahziInitializer(); + } + + @AfterEach + void tearDown() { + wireMockServer.stop(); + + if (loggingSystem != null) { + loggingSystem.close( + (int) TimeUnit.SECONDS.toMillis(10), + System.out::println + ); + + loggingSystem = null; + } + } + + @Test + void shouldIncludeAdditionalHeaders() { + // Given + wireMockServer.stubFor( + get(urlEqualTo("/loki/api/v1/push")) + .willReturn( + aResponse().withStatus(200) + )); + + ClientConfiguration clientConfiguration = ClientConfiguration.builder() + .withConnectionTimeoutMillis(10_000) + .withHost("localhost") + .withPort(wireMockServer.port()) + .withMaxRetries(1) + .build(); + + NettyHttpClient httpClient = HttpClientFactory.defaultFactory() + .getHttpClient( + clientConfiguration, + new String[]{ + "X-Scope-OrgID", "Circus", + "C", "Control" + } + ); + + LoggingSystem loggingSystem = initializer.createLoggingSystem( + httpClient, + 1024 * 1024, + false + ); + + // WHen + TjahziLogger logger = loggingSystem.createLogger(); + logger.log( + System.currentTimeMillis(), + Map.of(), + "Test" + ); + + // Then + await() + .atMost(Durations.FIVE_SECONDS) + .untilAsserted(() -> + wireMockServer.verify( + postRequestedFor(urlMatching("/loki/api/v1/push")) + .withHeader("X-Scope-OrgID", equalTo("Circus")) + )); + } + + @Test + void shouldHandleCaseWithNoAdditionalHeaders() { + // Given + wireMockServer.stubFor( + get(urlEqualTo("/loki/api/v1/push")) + .willReturn( + aResponse().withStatus(200) + )); + + ClientConfiguration clientConfiguration = ClientConfiguration.builder() + .withConnectionTimeoutMillis(10_000) + .withHost("localhost") + .withPort(wireMockServer.port()) + .withMaxRetries(1) + .build(); + + NettyHttpClient httpClient = HttpClientFactory.defaultFactory() + .getHttpClient( + clientConfiguration, + new String[0] + ); + + LoggingSystem loggingSystem = initializer.createLoggingSystem( + httpClient, + 1024 * 1024, + false + ); + + // WHen + TjahziLogger logger = loggingSystem.createLogger(); + logger.log( + System.currentTimeMillis(), + Map.of(), + "Test" + ); + + // Then + await() + .atMost(Durations.FIVE_SECONDS) + .untilAsserted(() -> + wireMockServer.verify( + postRequestedFor(urlMatching("/loki/api/v1/push")) + )); + } + + @Test + void shouldNotOverrideCrucialHeaders() { + // Given + wireMockServer.stubFor( + get(urlEqualTo("/loki/api/v1/push")) + .willReturn( + aResponse().withStatus(200) + )); + + ClientConfiguration clientConfiguration = ClientConfiguration.builder() + .withConnectionTimeoutMillis(10_000) + .withHost("localhost") + .withPort(wireMockServer.port()) + .withMaxRetries(1) + .build(); + + NettyHttpClient httpClient = HttpClientFactory.defaultFactory() + .getHttpClient( + clientConfiguration, + new String[]{ + "content-type", "text/plain", + "content-length", "5232423423", + "host", "remote" + } + ); + + LoggingSystem loggingSystem = initializer.createLoggingSystem( + httpClient, + 1024 * 1024, + false + ); + + // When + TjahziLogger logger = loggingSystem.createLogger(); + logger.log( + System.currentTimeMillis(), + Map.of(), + "Test" + ); + + // Then + await() + .atMost(Durations.FIVE_SECONDS) + .untilAsserted(() -> + wireMockServer.verify( + postRequestedFor(urlMatching("/loki/api/v1/push")) + .withHeader("content-type", matching("application/x-protobuf")) + .withHeader("content-length", matching("29|30")) + .withHeader("host", matching("localhost")) + )); + } +} diff --git a/tjahzi-core/src/test/java/pl/tkowalcz/thjazi/LoggingSystemTest.java b/tjahzi-core/src/test/java/pl/tkowalcz/tjahzi/LoggingSystemTest.java similarity index 97% rename from tjahzi-core/src/test/java/pl/tkowalcz/thjazi/LoggingSystemTest.java rename to tjahzi-core/src/test/java/pl/tkowalcz/tjahzi/LoggingSystemTest.java index 31bb46d..9fb6c2f 100644 --- a/tjahzi-core/src/test/java/pl/tkowalcz/thjazi/LoggingSystemTest.java +++ b/tjahzi-core/src/test/java/pl/tkowalcz/tjahzi/LoggingSystemTest.java @@ -57,7 +57,10 @@ class LoggingSystemTest { .build(); NettyHttpClient httpClient = HttpClientFactory.defaultFactory() - .getHttpClient(clientConfiguration); + .getHttpClient( + clientConfiguration, + new String[0] + ); TjahziInitializer initializer = new TjahziInitializer(); LoggingSystem loggingSystem = initializer.createLoggingSystem(