ENG-37132: Adding bodyAttribute functions to support already truncate…

…d bodies (#219) * ENG-37132: Adding bodyAttribute functions to support already truncated bodies * adding tests and refactoring
hypertrace · Nov 11, 2023 · 157e76c · 157e76c
1 parent de9ef45
commit 157e76c
Show file tree

Hide file tree

Showing 2 changed files with 182 additions and 9 deletions.
diff --git a/sdk/instrumentation/bodyattribute/bodyattribute.go b/sdk/instrumentation/bodyattribute/bodyattribute.go
@@ -17,30 +17,55 @@ func SetTruncatedBodyAttribute(attrName string, body []byte, bodyMaxSize int, sp
 	}
 
 	if bodyLen <= bodyMaxSize {
-		span.SetAttribute(attrName, string(body))
+		SetBodyAttribute(attrName, body, false, span)
 		return
 	}
 
-	truncatedBody := body[:bodyMaxSize]
-	span.SetAttribute(fmt.Sprintf("%s.truncated", attrName), true)
-	span.SetAttribute(attrName, string(truncatedBody))
+	SetBodyAttribute(attrName, body[:bodyMaxSize], true, span)
 }
 
 // SetTruncatedEncodedBodyAttribute is like SetTruncatedBodyAttribute above but also base64 encodes the
 // body. This is usually due to non utf8 bytes in the body eg. for multipart/form-data content type.
 // The body attribute name has a ".base64" suffix.
 func SetTruncatedEncodedBodyAttribute(attrName string, body []byte, bodyMaxSize int, span sdk.Span) {
 	bodyLen := len(body)
-	if bodyLen == 0 {
+	if len(body) == 0 {
 		return
 	}
 
 	if bodyLen <= bodyMaxSize {
-		span.SetAttribute(attrName+".base64", base64.RawStdEncoding.EncodeToString(body))
+		SetEncodedBodyAttribute(attrName, body, false, span)
+		return
+	}
+
+	SetEncodedBodyAttribute(attrName, body[:bodyMaxSize], true, span)
+}
+
+// SetBodyAttribute sets the body as a span attribute.
+// also sets truncated attribute if truncated is true
+func SetBodyAttribute(attrName string, body []byte, truncated bool, span sdk.Span) {
+	if len(body) == 0 {
 		return
 	}
 
-	truncatedBody := body[:bodyMaxSize]
-	span.SetAttribute(fmt.Sprintf("%s.truncated", attrName), true)
-	span.SetAttribute(attrName+".base64", base64.RawStdEncoding.EncodeToString(truncatedBody))
+	span.SetAttribute(attrName, string(body))
+	// if already truncated then set attribute
+	if truncated {
+		span.SetAttribute(fmt.Sprintf("%s.truncated", attrName), true)
+	}
+}
+
+// SetEncodedBodyAttribute is like SetBodyAttribute above but also base64 encodes the
+// body. This is usually due to non utf8 bytes in the body eg. for multipart/form-data content type.
+// The body attribute name has a ".base64" suffix.
+func SetEncodedBodyAttribute(attrName string, body []byte, truncated bool, span sdk.Span) {
+	if len(body) == 0 {
+		return
+	}
+
+	span.SetAttribute(attrName+".base64", base64.RawStdEncoding.EncodeToString(body))
+	// if already truncated then set attribute
+	if truncated {
+		span.SetAttribute(fmt.Sprintf("%s.truncated", attrName), true)
+	}
 }
diff --git a/sdk/instrumentation/bodyattribute/bodyattribute_test.go b/sdk/instrumentation/bodyattribute/bodyattribute_test.go
@@ -51,3 +51,151 @@ func TestSetTruncatedEncodedBodyAttributeEmptyBody(t *testing.T) {
 	assert.Nil(t, s.ReadAttribute("http.request.body.base64"))
 	assert.Zero(t, s.RemainingAttributes())
 }
+
+func TestSetBodyAttribute(t *testing.T) {
+	testBody := "test1test2"
+	type args struct {
+		attrName  string
+		body      []byte
+		truncated bool
+		span      *mock.Span
+	}
+	tests := []struct {
+		name               string
+		args               args
+		expectedAssertions func(t *testing.T, gotSpan *mock.Span)
+	}{
+		{
+			name: "empty body, truncated",
+			args: args{
+				attrName:  "http.request.body",
+				body:      []byte(""),
+				truncated: true,
+				span:      mock.NewSpan(),
+			},
+			expectedAssertions: func(t *testing.T, gotSpan *mock.Span) {
+				assert.Nil(t, gotSpan.ReadAttribute("http.request.body"))
+				assert.Zero(t, gotSpan.RemainingAttributes())
+			},
+		},
+		{
+			name: "empty body, not truncated",
+			args: args{
+				attrName:  "http.request.body",
+				body:      []byte(""),
+				truncated: false,
+				span:      mock.NewSpan(),
+			},
+			expectedAssertions: func(t *testing.T, gotSpan *mock.Span) {
+				assert.Nil(t, gotSpan.ReadAttribute("http.request.body"))
+				assert.Zero(t, gotSpan.RemainingAttributes())
+			},
+		},
+		{
+			name: "non empty body, not truncated",
+			args: args{
+				attrName:  "http.request.body",
+				body:      []byte(testBody),
+				truncated: false,
+				span:      mock.NewSpan(),
+			},
+			expectedAssertions: func(t *testing.T, gotSpan *mock.Span) {
+				assert.Equal(t, testBody, gotSpan.ReadAttribute("http.request.body"))
+				assert.Zero(t, gotSpan.RemainingAttributes())
+			},
+		},
+		{
+			name: "non empty body, truncated",
+			args: args{
+				attrName:  "http.request.body",
+				body:      []byte(testBody),
+				truncated: true,
+				span:      mock.NewSpan(),
+			},
+			expectedAssertions: func(t *testing.T, gotSpan *mock.Span) {
+				assert.Equal(t, testBody, gotSpan.ReadAttribute("http.request.body"))
+				assert.True(t, (gotSpan.ReadAttribute("http.request.body.truncated")).(bool))
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			SetBodyAttribute(tt.args.attrName, tt.args.body, tt.args.truncated, tt.args.span)
+			tt.expectedAssertions(t, tt.args.span)
+		})
+	}
+}
+
+func TestSetEncodedBodyAttribute(t *testing.T) {
+	testBody := "test1test2"
+	type args struct {
+		attrName  string
+		body      []byte
+		truncated bool
+		span      *mock.Span
+	}
+	tests := []struct {
+		name               string
+		args               args
+		expectedAssertions func(t *testing.T, gotSpan *mock.Span)
+	}{
+		{
+			name: "empty body, truncated",
+			args: args{
+				attrName:  "http.request.body",
+				body:      []byte(""),
+				truncated: true,
+				span:      mock.NewSpan(),
+			},
+			expectedAssertions: func(t *testing.T, gotSpan *mock.Span) {
+				assert.Nil(t, gotSpan.ReadAttribute("http.request.body.base64"))
+				assert.Zero(t, gotSpan.RemainingAttributes())
+			},
+		},
+		{
+			name: "empty body, not truncated",
+			args: args{
+				attrName:  "http.request.body",
+				body:      []byte(""),
+				truncated: false,
+				span:      mock.NewSpan(),
+			},
+			expectedAssertions: func(t *testing.T, gotSpan *mock.Span) {
+				assert.Nil(t, gotSpan.ReadAttribute("http.request.body.base64"))
+				assert.Zero(t, gotSpan.RemainingAttributes())
+			},
+		},
+		{
+			name: "non empty body, not truncated",
+			args: args{
+				attrName:  "http.request.body",
+				body:      []byte(testBody),
+				truncated: false,
+				span:      mock.NewSpan(),
+			},
+			expectedAssertions: func(t *testing.T, gotSpan *mock.Span) {
+				assert.Equal(t, base64.RawStdEncoding.EncodeToString([]byte(testBody)), gotSpan.ReadAttribute("http.request.body.base64"))
+				assert.Zero(t, gotSpan.RemainingAttributes())
+			},
+		},
+		{
+			name: "non empty body, truncated",
+			args: args{
+				attrName:  "http.request.body",
+				body:      []byte(testBody),
+				truncated: true,
+				span:      mock.NewSpan(),
+			},
+			expectedAssertions: func(t *testing.T, gotSpan *mock.Span) {
+				assert.Equal(t, base64.RawStdEncoding.EncodeToString([]byte(testBody)), gotSpan.ReadAttribute("http.request.body.base64"))
+				assert.True(t, (gotSpan.ReadAttribute("http.request.body.truncated")).(bool))
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			SetEncodedBodyAttribute(tt.args.attrName, tt.args.body, tt.args.truncated, tt.args.span)
+			tt.expectedAssertions(t, tt.args.span)
+		})
+	}
+}