Don't attempt to cast to HashMap when Map can do, and also make sure runtimeContext isn't poisoned by a Segment. (#202)

Author: Anuraag Agrawal

aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/AWSXRayRecorder.java

@@ -475,7 +475,7 @@ private Segment beginSegment(Segment segment) {
                          + "named '" + current.getName() + "' to start new segment named '" + segment.getName() + "'.");
         }
 
-        segment.setAws(getAwsRuntimeContext());
+        segment.putAllAws(getAwsRuntimeContext());
         if (origin != null) {
             segment.setOrigin(origin);
         }

aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/SegmentImpl.java

@@ -149,7 +149,7 @@ public void setRuleName(String ruleName) {
         checkAlreadyEmitted();
         if (getAws().get("xray") instanceof Map) {
             @SuppressWarnings("unchecked")
-            Map<String, Object> a = (HashMap<String, Object>) getAws().get("xray");
+            Map<String, Object> a = (Map<String, Object>) getAws().get("xray");
             HashMap<String, Object> referA = new HashMap<>();
             if (a != null) {
                 referA.putAll(a);

aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/AWSXRayRecorderTest.java

@@ -334,6 +334,26 @@ public void testSubsegmentEmittedInLambdaContext() throws JSONException {
             captured.getEndTime()).toString(), captured.streamSerialize(), JSONCompareMode.NON_EXTENSIBLE);
     }
 
+    @Test
+    public void testBeginSegment_awsRuntimeContextNotMutableBySegment() {
+        Segment segment = AWSXRay.beginSegment("test");
+        segment.putAws("foo", "bar");
+        assertThat(segment.getAws().get("foo")).isEqualTo("bar");
+        AWSXRay.endSegment();
+
+        segment = AWSXRay.beginSegment("test");
+        assertThat(segment.getAws().get("foo")).isNull();
+        AWSXRay.endSegment();
+    }
+
+    @Test
+    public void testBeginSegment_canSetRuleName() {
+        Segment segment = AWSXRay.beginSegment("test");
+        segment.setRuleName("rule");
+        assertThat(segment.getAws().get("xray")).isInstanceOfSatisfying(
+            Map.class, xray -> assertThat(xray.get("rule_name")).isEqualTo("rule"));
+    }
+
     private ObjectNode expectedLambdaSubsegment(
         TraceID traceId, String segmentId, String subsegmentId, double startTime, double endTime) {
         ObjectNode expected = JsonNodeFactory.instance.objectNode();