Skip to content

Commit 7ca6e6e

Browse files
darcciodarccio
committed
feat(tracer): RUM injection POC
1 parent 389686a commit 7ca6e6e

File tree

2 files changed

+394
-0
lines changed

2 files changed

+394
-0
lines changed

‎instrumentation/rum/inject.go‎

Lines changed: 215 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,215 @@
1+
// Unless explicitly stated otherwise all files in this repository are licensed
2+
// under the Apache License Version 2.0.
3+
// This product includes software developed at Datadog (https://www.datadoghq.com/).
4+
// Copyright 2025 Datadog, Inc.
5+
6+
package rum
7+
8+
import (
9+
"net/http"
10+
"unicode"
11+
)
12+
13+
type state int
14+
15+
const (
16+
sInit state = iota // looking for '<'
17+
sLt // saw '<', expect '/'
18+
sSlash // saw "</", allow spaces, expect 'h'
19+
sH // expect 'e' (no spaces allowed)
20+
sE // expect 'a' (no spaces allowed)
21+
sA // expect 'd' (no spaces allowed)
22+
sD // saw "...head", allow spaces, expect '>'
23+
sDone // "</head>" found
24+
)
25+
26+
var (
27+
snippet = []byte("<snippet>")
28+
)
29+
30+
// injector of a POC for RUM snippet injection.
31+
// It doesn't handle Content-Length manipulation.
32+
// It isn't concurrent safe.
33+
type injector struct {
34+
wrapped http.ResponseWriter
35+
st state
36+
lastSeen int
37+
seenInCurrent bool
38+
buf [][]byte
39+
}
40+
41+
// Header implements http.ResponseWriter.
42+
func (ij *injector) Header() http.Header {
43+
// TODO: this is a good place to inject Content-Length to the right
44+
// length, not the original one, if injection happened.
45+
return ij.wrapped.Header()
46+
}
47+
48+
// WriteHeader implements http.ResponseWriter.
49+
func (ij *injector) WriteHeader(statusCode int) {
50+
ij.wrapped.WriteHeader(statusCode)
51+
}
52+
53+
// Write implements http.ResponseWriter.
54+
// There are no guarantees that Write will be called with the whole payload.
55+
// We need to keep state of what we've written so far to find the pattern
56+
// "</head>" in all its variants.
57+
func (ij *injector) Write(chunk []byte) (int, error) {
58+
prev := ij.st
59+
// If we've already found the pattern, just write the chunk.
60+
if prev == sDone {
61+
return ij.wrapped.Write(chunk)
62+
}
63+
ij.match(chunk)
64+
if prev == sInit {
65+
// No partial or full match done so far.
66+
if ij.st == sInit {
67+
return ij.wrapped.Write(chunk)
68+
}
69+
// Full match done in the chunk.
70+
if ij.st == sDone {
71+
ij.st = sDone
72+
sz, err := multiWrite(ij.wrapped, chunk[:ij.lastSeen], snippet, chunk[ij.lastSeen:])
73+
if err != nil {
74+
return sz, err
75+
}
76+
return sz, nil
77+
}
78+
// Partial match in progress. We buffer the write.
79+
// ij.lastSeen should be the index of the first byte of the match
80+
// of the first chunk.
81+
ij.buf = append(ij.buf, chunk)
82+
return 0, nil
83+
}
84+
if ij.st != sDone {
85+
// Partial match in progress. We buffer the write.
86+
ij.buf = append(ij.buf, chunk)
87+
return 0, nil
88+
}
89+
// Partial match done.
90+
var (
91+
total int
92+
sz int
93+
err error
94+
)
95+
ij.buf = append(ij.buf, chunk)
96+
seenAt := 0
97+
if ij.seenInCurrent {
98+
seenAt = len(ij.buf) - 1
99+
}
100+
// Write the chunks before the chunk where the pattern starts.
101+
sz, err = multiWrite(ij.wrapped, ij.buf[:seenAt]...)
102+
if err != nil {
103+
return sz, err
104+
}
105+
total += sz
106+
// Write the snippet in the chunk where the pattern starts.
107+
head := ij.buf[seenAt]
108+
sz, err = multiWrite(ij.wrapped, head[:ij.lastSeen], snippet, head[ij.lastSeen:])
109+
if err != nil {
110+
return sz, err
111+
}
112+
total += sz
113+
// Write the rest of the buffered chunks.
114+
sz, err = multiWrite(ij.wrapped, ij.buf[seenAt+1:]...)
115+
if err != nil {
116+
return sz, err
117+
}
118+
total += sz
119+
// Reset the buffer.
120+
ij.buf = ij.buf[:0]
121+
return total, nil
122+
}
123+
124+
func multiWrite(w http.ResponseWriter, chunks ...[]byte) (int, error) {
125+
if len(chunks) == 0 {
126+
return 0, nil
127+
}
128+
sz := 0
129+
for _, chunk := range chunks {
130+
n, err := w.Write(chunk)
131+
if err != nil {
132+
return sz, err
133+
}
134+
sz += n
135+
}
136+
return sz, nil
137+
}
138+
139+
// match updates the state of the injector according on what step of
140+
// the pattern "</head>" have been found.
141+
func (ij *injector) match(p []byte) {
142+
if ij.st == sDone {
143+
return
144+
}
145+
ij.seenInCurrent = false
146+
for i := 0; i < len(p); i++ {
147+
c := unicode.ToLower(rune(p[i]))
148+
switch ij.st {
149+
case sInit:
150+
ij.transition('<', c, sLt, i)
151+
case sLt: // expect '/'
152+
ij.transition('/', c, sSlash, i)
153+
case sSlash: // expect 'h'
154+
if unicode.IsSpace(c) {
155+
continue
156+
}
157+
ij.transition('h', c, sH, i)
158+
case sH: // expect 'e'
159+
ij.transition('e', c, sE, i)
160+
case sE: // expect 'a'
161+
ij.transition('a', c, sA, i)
162+
case sA: // expect 'd'
163+
ij.transition('d', c, sD, i)
164+
case sD: // expect '>'
165+
if unicode.IsSpace(c) {
166+
continue
167+
}
168+
ij.transition('>', c, sDone, i)
169+
}
170+
}
171+
}
172+
173+
func (ij *injector) transition(expected, current rune, target state, pos int) {
174+
switch current {
175+
case expected:
176+
ij.st = target
177+
case '<':
178+
ij.st = sLt
179+
default:
180+
ij.st = sInit
181+
}
182+
if current == '<' {
183+
ij.lastSeen = pos
184+
ij.seenInCurrent = true
185+
}
186+
}
187+
188+
// Flush flushes the buffered chunks to the wrapped writer.
189+
func (ij *injector) Flush() (int, error) {
190+
if len(ij.buf) == 0 {
191+
return 0, nil
192+
}
193+
sz, err := multiWrite(ij.wrapped, ij.buf...)
194+
ij.buf = ij.buf[:0]
195+
return sz, err
196+
}
197+
198+
// Reset resets the state of the injector.
199+
func (i *injector) Reset() {
200+
i.st = sInit
201+
i.lastSeen = -1
202+
i.buf = i.buf[:0]
203+
}
204+
205+
func NewInjector(fn func(w http.ResponseWriter, r *http.Request)) http.Handler {
206+
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
207+
ij := &injector{
208+
wrapped: w,
209+
lastSeen: -1,
210+
buf: make([][]byte, 0, 10),
211+
}
212+
fn(ij, r)
213+
ij.Flush()
214+
})
215+
}

‎instrumentation/rum/inject_test.go‎

Lines changed: 179 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,179 @@
1+
// Unless explicitly stated otherwise all files in this repository are licensed
2+
// under the Apache License Version 2.0.
3+
// This product includes software developed at Datadog (https://www.datadoghq.com/).
4+
// Copyright 2025 Datadog, Inc.
5+
6+
package rum
7+
8+
import (
9+
"bytes"
10+
"io"
11+
"net/http"
12+
"net/http/httptest"
13+
"strings"
14+
"testing"
15+
16+
"github.com/stretchr/testify/assert"
17+
)
18+
19+
func TestInjector(t *testing.T) {
20+
payload := []byte("Hello, world!")
21+
h := func(w http.ResponseWriter, r *http.Request) {
22+
w.Write(payload)
23+
}
24+
injector := NewInjector(h)
25+
server := httptest.NewServer(injector)
26+
defer server.Close()
27+
28+
resp, err := http.DefaultClient.Get(server.URL)
29+
assert.NoError(t, err)
30+
defer resp.Body.Close()
31+
32+
respBody, err := io.ReadAll(resp.Body)
33+
assert.NoError(t, err)
34+
assert.Equal(t, "Hello, world!", string(respBody))
35+
// TODO: when Content-Length is implemented, uncomment this.
36+
// assert.Equal(t, int64(len(payload) + len(snippet)), resp.ContentLength)
37+
}
38+
39+
func TestInjectorMatch(t *testing.T) {
40+
cases := []struct {
41+
in []byte
42+
want state
43+
}{
44+
{[]byte("hello </head> world"), sDone},
45+
{[]byte("noise </ head > tail"), sDone}, // spaces after '/' and before '>'
46+
{[]byte("nope < /head>"), sInit}, // space between '<' and '/'
47+
{[]byte("nope </ he ad >"), sInit}, // spaces inside "head"
48+
{[]byte("ok </\tHead\t\t >"), sDone}, // tabs after '/', spaces before '>'
49+
{[]byte("partial </hea>"), sInit}, // missing 'd'
50+
{[]byte("wrong </header>"), sInit}, // extra letters before '>'
51+
{[]byte("caps </HEAD>"), sDone}, // case-insensitive
52+
{[]byte("attr-like </head foo>"), sInit}, // rejected by our custom rule
53+
{[]byte("prefix << / h e a d >"), sInit}, // multiple violations
54+
}
55+
56+
for _, tc := range cases {
57+
t.Run(string(tc.in), func(t *testing.T) {
58+
i := &injector{}
59+
i.match(tc.in)
60+
got := i.st
61+
i.Reset()
62+
if got != tc.want {
63+
t.Fatalf("match(%q) = %v; want %v", tc.in, got, tc.want)
64+
}
65+
})
66+
}
67+
}
68+
69+
func TestInjectorWrite(t *testing.T) {
70+
cases := []struct {
71+
category string
72+
in string // comma separated chunks
73+
out string
74+
}{
75+
{"basic", "abc</head>def", "abc<snippet></head>def"},
76+
{"basic", "abc</he,ad>def", "abc<snippet></head>def"},
77+
{"basic", "abc,</head>def", "abc<snippet></head>def"},
78+
{"basic", "abc</head>,def", "abc<snippet></head>def"},
79+
{"basic", "abc</h,ea,d>def", "abc<snippet></head>def"},
80+
{"basic", "abc,</hea,d>def", "abc<snippet></head>def"},
81+
{"no-head", "abc", "abc"},
82+
{"no-head", "abc</hea", "abc</hea"},
83+
{"empty", "", ""},
84+
{"empty", ",", ""},
85+
{"incomplete", "abc</he</head>def", "abc</he<snippet></head>def"},
86+
{"incomplete", "abc</he,</head>def", "abc</he<snippet></head>def"},
87+
{"casing", "abc</HeAd>def", "abc<snippet></HeAd>def"},
88+
{"casing", "abc</HEAD>def", "abc<snippet></HEAD>def"},
89+
{"spaces", "abc </head>def", "abc <snippet></head>def"},
90+
{"spaces", "abc </hea,d>def", "abc <snippet></head>def"},
91+
{"spaces", "abc</ head>def", "abc<snippet></ head>def"},
92+
{"spaces", "abc</h ead>def", "abc</h ead>def"},
93+
{"spaces", "abc</he ad>def", "abc</he ad>def"},
94+
{"spaces", "abc</hea d>def", "abc</hea d>def"},
95+
{"spaces", "abc</head >def", "abc<snippet></head >def"},
96+
{"spaces", "abc</head> def", "abc<snippet></head> def"},
97+
}
98+
99+
for _, tc := range cases {
100+
t.Run(tc.category+":"+tc.in, func(t *testing.T) {
101+
total := 0
102+
recorder := httptest.NewRecorder()
103+
i := &injector{
104+
wrapped: recorder,
105+
}
106+
chunks := strings.Split(tc.in, ",")
107+
for _, chunk := range chunks {
108+
w, err := i.Write([]byte(chunk))
109+
assert.NoError(t, err)
110+
total += w
111+
}
112+
sz, err := i.Flush()
113+
assert.NoError(t, err)
114+
total += sz
115+
body := recorder.Body.String()
116+
assert.Equal(t, tc.out, body)
117+
assert.Equal(t, len(tc.out), total)
118+
})
119+
}
120+
}
121+
122+
type sinkResponseWriter struct {
123+
out []byte
124+
}
125+
126+
func (s *sinkResponseWriter) Header() http.Header {
127+
return http.Header{}
128+
}
129+
func (s *sinkResponseWriter) Write(p []byte) (int, error) {
130+
s.out = append(s.out, p...)
131+
return len(p), nil
132+
}
133+
func (s *sinkResponseWriter) WriteHeader(int) {}
134+
135+
func BenchmarkInjectorWrite(b *testing.B) {
136+
b.ReportAllocs()
137+
b.ResetTimer()
138+
sink := &sinkResponseWriter{}
139+
ij := &injector{
140+
wrapped: sink,
141+
}
142+
for i := 0; i < b.N; i++ {
143+
ij.Write([]byte("abc</head>def"))
144+
if !bytes.Equal(sink.out, []byte("abc<snippet></head>def")) {
145+
b.Fatalf("out is not as expected: %q", sink.out)
146+
}
147+
sink.out = sink.out[:0]
148+
ij.Reset()
149+
}
150+
}
151+
152+
func FuzzInjectorWrite(f *testing.F) {
153+
cases := []string{
154+
"abc</head>def",
155+
"abc",
156+
"abc</hea",
157+
"abc</he</head>def",
158+
"abc</HeAd>def",
159+
"abc</HEAD>def",
160+
"abc </head>def",
161+
"abc</ head>def",
162+
"abc</h ead>def",
163+
"abc</he ad>def",
164+
"abc</hea d>def",
165+
"abc</head >def",
166+
"abc</head> def",
167+
"",
168+
}
169+
for _, tc := range cases {
170+
f.Add([]byte(tc))
171+
}
172+
f.Fuzz(func(t *testing.T, in []byte) {
173+
sink := &sinkResponseWriter{}
174+
ij := &injector{
175+
wrapped: sink,
176+
}
177+
ij.Write(in)
178+
})
179+
}

0 commit comments

Comments
 (0)