Skip to content

Commit 7d9336b

Browse files
jeromevdljeromevdl
committed
custom message resolver for json strings
1 parent 8a4f05e commit 7d9336b

File tree

6 files changed

+196
-8
lines changed

6 files changed

+196
-8
lines changed

‎powertools-logging/powertools-logging-log4j/pom.xml

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -53,17 +53,14 @@
5353
<dependency>
5454
<groupId>org.apache.logging.log4j</groupId>
5555
<artifactId>log4j-slf4j2-impl</artifactId>
56-
<scope>provided</scope>
5756
</dependency>
5857
<dependency>
5958
<groupId>org.apache.logging.log4j</groupId>
6059
<artifactId>log4j-core</artifactId>
61-
<scope>provided</scope>
6260
</dependency>
6361
<dependency>
6462
<groupId>org.apache.logging.log4j</groupId>
6563
<artifactId>log4j-layout-template-json</artifactId>
66-
<scope>provided</scope>
6764
</dependency>
6865

6966
<!-- Test dependencies -->

‎powertools-logging/powertools-logging-log4j/src/main/java/org/apache/logging/log4j/layout/template/json/resolver/PowertoolsResolver.java

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,11 +24,15 @@
2424
import static software.amazon.lambda.powertools.logging.internal.PowertoolsLoggedFields.SAMPLING_RATE;
2525
import static software.amazon.lambda.powertools.logging.internal.PowertoolsLoggedFields.SERVICE;
2626

27+
import com.fasterxml.jackson.core.JacksonException;
28+
import com.fasterxml.jackson.databind.DeserializationFeature;
29+
import com.fasterxml.jackson.databind.ObjectMapper;
2730
import java.util.Map;
2831
import java.util.stream.Collectors;
2932
import java.util.stream.Stream;
3033
import org.apache.logging.log4j.core.LogEvent;
3134
import org.apache.logging.log4j.layout.template.json.util.JsonWriter;
35+
import org.apache.logging.log4j.message.Message;
3236
import org.apache.logging.log4j.util.ReadOnlyStringMap;
3337
import software.amazon.lambda.powertools.common.internal.LambdaConstants;
3438
import software.amazon.lambda.powertools.common.internal.SystemWrapper;
@@ -147,11 +151,14 @@ public void resolve(LogEvent logEvent, JsonWriter jsonWriter) {
147151
(final LogEvent logEvent, final JsonWriter jsonWriter) ->
148152
jsonWriter.writeString(SystemWrapper.getenv(LambdaConstants.AWS_REGION_ENV));
149153

154+
public static final String LAMBDA_ARN_REGEX =
155+
"^arn:(aws|aws-us-gov|aws-cn):lambda:[a-zA-Z0-9-]+:\\d{12}:function:[a-zA-Z0-9-_]+(:[a-zA-Z0-9-_]+)?$";
156+
150157
private static final EventResolver ACCOUNT_ID_RESOLVER = new EventResolver() {
151158
@Override
152159
public boolean isResolvable(LogEvent logEvent) {
153160
final String arn = logEvent.getContextData().getValue(PowertoolsLoggedFields.FUNCTION_ARN.getName());
154-
return null != arn && !arn.isEmpty();
161+
return null != arn && !arn.isEmpty() && arn.matches(LAMBDA_ARN_REGEX);
155162
}
156163

157164
@Override
@@ -161,6 +168,42 @@ public void resolve(LogEvent logEvent, JsonWriter jsonWriter) {
161168
}
162169
};
163170

171+
/**
172+
* Use a custom message resolver to permit to log json string in json format without escaped quotes.
173+
*/
174+
private static final EventResolver MESSAGE_RESOLVER = new EventResolver() {
175+
private final ObjectMapper mapper = new ObjectMapper()
176+
.enable(DeserializationFeature.FAIL_ON_TRAILING_TOKENS);
177+
178+
public boolean isValidJson(String json) {
179+
if (!(json.startsWith("{") || json.startsWith("["))) {
180+
return false;
181+
}
182+
try {
183+
mapper.readTree(json);
184+
} catch (JacksonException e) {
185+
return false;
186+
}
187+
return true;
188+
}
189+
190+
@Override
191+
public boolean isResolvable(LogEvent logEvent) {
192+
final Message msg = logEvent.getMessage();
193+
return null != msg && null != msg.getFormattedMessage();
194+
}
195+
196+
@Override
197+
public void resolve(LogEvent logEvent, JsonWriter jsonWriter) {
198+
String message = logEvent.getMessage().getFormattedMessage();
199+
if (isValidJson(message)) {
200+
jsonWriter.writeRawString(message);
201+
} else {
202+
jsonWriter.writeString(message);
203+
}
204+
}
205+
};
206+
164207
private static final EventResolver NON_POWERTOOLS_FIELD_RESOLVER =
165208
(LogEvent logEvent, JsonWriter jsonWriter) -> {
166209
StringBuilder stringBuilder = jsonWriter.getStringBuilder();
@@ -193,6 +236,7 @@ public void resolve(LogEvent logEvent, JsonWriter jsonWriter) {
193236
{ SAMPLING_RATE.getName(), SAMPLING_RATE_RESOLVER },
194237
{ "region", REGION_RESOLVER },
195238
{ "account_id", ACCOUNT_ID_RESOLVER },
239+
{ "message", MESSAGE_RESOLVER }
196240
}).collect(Collectors.toMap(data -> (String) data[0], data -> (EventResolver) data[1]));
197241

198242

‎powertools-logging/powertools-logging-log4j/src/main/resources/LambdaEcsLayout.json

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,8 @@
1212
"field": "name"
1313
},
1414
"message": {
15-
"$resolver": "message",
16-
"stringified": true
15+
"$resolver": "powertools",
16+
"field": "message"
1717
},
1818
"error.type": {
1919
"$resolver": "exception",

‎powertools-logging/powertools-logging-log4j/src/main/resources/LambdaJsonLayout.json

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,8 @@
44
"field": "name"
55
},
66
"message": {
7-
"$resolver": "message",
8-
"stringified": true
7+
"$resolver": "powertools",
8+
"field": "message"
99
},
1010
"error": {
1111
"message": {

‎powertools-logging/powertools-logging-log4j/src/test/java/org/apache/logging/log4j/layout/template/json/resolver/PowertoolsMessageResolverTest.java

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
/*
2+
* Copyright 2023 Amazon.com, Inc. or its affiliates.
3+
* Licensed under the Apache License, Version 2.0 (the
4+
* "License"); you may not use this file except in compliance
5+
* with the License. You may obtain a copy of the License at
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
* Unless required by applicable law or agreed to in writing, software
8+
* distributed under the License is distributed on an "AS IS" BASIS,
9+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10+
* See the License for the specific language governing permissions and
11+
* limitations under the License.
12+
*
13+
*/
14+
15+
package org.apache.logging.log4j.layout.template.json.resolver;
16+
17+
import static org.assertj.core.api.Assertions.assertThat;
18+
import static org.assertj.core.api.Assertions.contentOf;
19+
import static org.mockito.Mockito.when;
20+
import static org.mockito.MockitoAnnotations.openMocks;
21+
22+
import com.amazonaws.services.lambda.runtime.Context;
23+
import com.amazonaws.services.lambda.runtime.events.SQSEvent;
24+
import java.io.File;
25+
import java.io.IOException;
26+
import java.nio.channels.FileChannel;
27+
import java.nio.file.NoSuchFileException;
28+
import java.nio.file.Paths;
29+
import java.nio.file.StandardOpenOption;
30+
import java.util.Arrays;
31+
import java.util.Collections;
32+
import org.junit.jupiter.api.AfterEach;
33+
import org.junit.jupiter.api.BeforeEach;
34+
import org.junit.jupiter.api.Test;
35+
import org.mockito.Mock;
36+
import org.slf4j.MDC;
37+
import software.amazon.lambda.powertools.logging.internal.handler.PowertoolsJsonMessage;
38+
import software.amazon.lambda.powertools.logging.internal.handler.PowertoolsLogEnabled;
39+
40+
class PowertoolsMessageResolverTest {
41+
42+
@Mock
43+
private Context context;
44+
45+
@BeforeEach
46+
void setUp() throws IOException {
47+
openMocks(this);
48+
MDC.clear();
49+
setupContext();
50+
51+
try {
52+
FileChannel.open(Paths.get("target/logfile.json"), StandardOpenOption.WRITE).truncate(0).close();
53+
FileChannel.open(Paths.get("target/ecslogfile.json"), StandardOpenOption.WRITE).truncate(0).close();
54+
} catch (NoSuchFileException e) {
55+
// may not be there in the first run
56+
}
57+
}
58+
59+
@AfterEach
60+
void cleanUp() throws IOException {
61+
//Make sure file is cleaned up before running full stack logging regression
62+
FileChannel.open(Paths.get("target/logfile.json"), StandardOpenOption.WRITE).truncate(0).close();
63+
FileChannel.open(Paths.get("target/ecslogfile.json"), StandardOpenOption.WRITE).truncate(0).close();
64+
}
65+
66+
@Test
67+
void shouldLogJsonMessageWithoutEscapedStrings() {
68+
// GIVEN
69+
PowertoolsJsonMessage requestHandler = new PowertoolsJsonMessage();
70+
SQSEvent.SQSMessage msg = new SQSEvent.SQSMessage();
71+
msg.setMessageId("1212abcd");
72+
msg.setBody("plop");
73+
msg.setEventSource("eb");
74+
msg.setAwsRegion("eu-west-1");
75+
SQSEvent.MessageAttribute attribute = new SQSEvent.MessageAttribute();
76+
attribute.setStringListValues(Arrays.asList("val1", "val2", "val3"));
77+
msg.setMessageAttributes(Collections.singletonMap("keyAttribute", attribute));
78+
79+
// WHEN
80+
requestHandler.handleRequest(msg, context);
81+
82+
// THEN
83+
File logFile = new File("target/logfile.json");
84+
assertThat(contentOf(logFile)).contains("\"message\":{\"messageId\":\"1212abcd\",\"receiptHandle\":null,\"body\":\"plop\",\"md5OfBody\":null,\"md5OfMessageAttributes\":null,\"eventSourceArn\":null,\"eventSource\":\"eb\",\"awsRegion\":\"eu-west-1\",\"attributes\":null,\"messageAttributes\":{\"keyAttribute\":{\"stringValue\":null,\"binaryValue\":null,\"stringListValues\":[\"val1\",\"val2\",\"val3\"],\"binaryListValues\":null,\"dataType\":null}}}");
85+
}
86+
87+
88+
@Test
89+
void shouldLogStringMessageWhenNotJson() {
90+
// GIVEN
91+
PowertoolsLogEnabled requestHandler = new PowertoolsLogEnabled();
92+
93+
// WHEN
94+
requestHandler.handleRequest(null, context);
95+
96+
// THEN
97+
File logFile = new File("target/logfile.json");
98+
assertThat(contentOf(logFile)).contains("\"message\":\"Test debug event\"");
99+
}
100+
101+
private void setupContext() {
102+
when(context.getFunctionName()).thenReturn("testFunction");
103+
when(context.getInvokedFunctionArn()).thenReturn("testArn");
104+
when(context.getFunctionVersion()).thenReturn("1");
105+
when(context.getMemoryLimitInMB()).thenReturn(10);
106+
when(context.getAwsRequestId()).thenReturn("RequestId");
107+
}
108+
}

‎powertools-logging/powertools-logging-log4j/src/test/java/software/amazon/lambda/powertools/logging/internal/handler/PowertoolsJsonMessage.java

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
/*
2+
* Copyright 2023 Amazon.com, Inc. or its affiliates.
3+
* Licensed under the Apache License, Version 2.0 (the
4+
* "License"); you may not use this file except in compliance
5+
* with the License. You may obtain a copy of the License at
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
* Unless required by applicable law or agreed to in writing, software
8+
* distributed under the License is distributed on an "AS IS" BASIS,
9+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10+
* See the License for the specific language governing permissions and
11+
* limitations under the License.
12+
*
13+
*/
14+
15+
package software.amazon.lambda.powertools.logging.internal.handler;
16+
17+
import com.amazonaws.services.lambda.runtime.Context;
18+
import com.amazonaws.services.lambda.runtime.RequestHandler;
19+
import com.amazonaws.services.lambda.runtime.events.SQSEvent;
20+
import com.fasterxml.jackson.core.JsonProcessingException;
21+
import org.slf4j.Logger;
22+
import org.slf4j.LoggerFactory;
23+
import software.amazon.lambda.powertools.logging.Logging;
24+
import software.amazon.lambda.powertools.utilities.JsonConfig;
25+
26+
public class PowertoolsJsonMessage implements RequestHandler<SQSEvent.SQSMessage, String> {
27+
private final Logger LOG = LoggerFactory.getLogger(PowertoolsJsonMessage.class);
28+
29+
@Override
30+
@Logging(clearState = true)
31+
public String handleRequest(SQSEvent.SQSMessage input, Context context) {
32+
try {
33+
LOG.debug(JsonConfig.get().getObjectMapper().writeValueAsString(input));
34+
} catch (JsonProcessingException e) {
35+
throw new RuntimeException(e);
36+
}
37+
return input.getMessageId();
38+
}
39+
}

0 commit comments

Comments
 (0)