1package magna
2
3import (
4 "errors"
5 "testing"
6
7 "github.com/stretchr/testify/assert"
8)
9
10func BenchmarkDecodeDomainSimple(b *testing.B) {
11 input := []byte{3, 'w', 'w', 'w', 7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0}
12 for i := 0; i < b.N; i++ {
13 _, _, _ = decodeDomain(input, 0)
14 }
15}
16
17func BenchmarkDecodeDomainCompressed(b *testing.B) {
18 input := []byte{
19 0x07, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 0x03, 'c', 'o', 'm', 0x00,
20 0x03, 'f', 'o', 'o', 0xc0, 0x00,
21 }
22 offset := 13
23 b.ResetTimer()
24 for i := 0; i < b.N; i++ {
25 _, _, _ = decodeDomain(input, offset)
26 }
27}
28
29func BenchmarkEncodeDomainSimple(b *testing.B) {
30 domain := "www.example.com"
31 offsets := make(map[string]uint16)
32 out := make([]byte, 0, 64)
33 b.ResetTimer()
34 for i := 0; i < b.N; i++ {
35 _, _ = encodeDomain(out[:0], domain, &offsets)
36 for k := range offsets {
37 delete(offsets, k)
38 }
39 }
40}
41
42func BenchmarkEncodeDomainWithCompression(b *testing.B) {
43 domain1 := "www.example.com"
44 domain2 := "mail.example.com"
45 offsets := make(map[string]uint16)
46 out := make([]byte, 0, 128)
47 b.ResetTimer()
48 for i := 0; i < b.N; i++ {
49 tempOut, _ := encodeDomain(out[:0], domain1, &offsets)
50 _, _ = encodeDomain(tempOut, domain2, &offsets)
51 for k := range offsets {
52 delete(offsets, k)
53 }
54 }
55}
56
57func TestDecodeDomain(t *testing.T) {
58 tests := []struct {
59 name string
60 offset int
61 input []byte
62 expectedDomain string
63 expectedOffset int
64 expectedError error
65 errorCheck func(t *testing.T, err error)
66 }{
67 {
68 name: "Simple domain",
69 input: []byte{3, 'w', 'w', 'w', 7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0},
70 expectedDomain: "www.example.com",
71 expectedOffset: 17,
72 expectedError: nil,
73 },
74 {
75 name: "Domain with compression",
76 offset: 17,
77 input: []byte{3, 'w', 'w', 'w', 7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0, 3, 'f', 'o', 'o', 0xC0, 0},
78 expectedDomain: "foo.www.example.com",
79 expectedOffset: 23,
80 expectedError: nil,
81 },
82 {
83 name: "Invalid label length",
84 input: []byte{64, 'x'},
85 expectedDomain: "",
86 expectedOffset: 2,
87 expectedError: &InvalidLabelError{Length: 64},
88 errorCheck: func(t *testing.T, err error) {
89 var target *InvalidLabelError
90 assert.True(t, errors.As(err, &target))
91 assert.Equal(t, 64, target.Length)
92 },
93 },
94 {
95 name: "Compression loop",
96 input: []byte{0xC0, 0, 0xC0, 0},
97 expectedDomain: "",
98 expectedOffset: 4,
99 expectedError: &DomainCompressionError{},
100 errorCheck: func(t *testing.T, err error) {
101 assert.IsType(t, &DomainCompressionError{}, err)
102 },
103 },
104 {
105 name: "Truncated input",
106 input: []byte{3, 'w', 'w'},
107 expectedDomain: "",
108 expectedOffset: 3,
109 expectedError: &BufferOverflowError{Length: 3, Offset: 4},
110 errorCheck: func(t *testing.T, err error) {
111 var target *BufferOverflowError
112 assert.True(t, errors.As(err, &target), "Expected BufferOverflowError")
113 if target != nil {
114 assert.Equal(t, 3, target.Length)
115 assert.Equal(t, 1+3, target.Offset)
116 }
117 assert.Contains(t, err.Error(), "failed to read domain label data")
118 },
119 },
120 }
121
122 for _, tt := range tests {
123 t.Run(tt.name, func(t *testing.T) {
124 domain, offset, err := decodeDomain(tt.input, tt.offset)
125
126 t.Logf("Test: %s, Input: %x, OffsetIn: %d => Domain: '%s', OffsetOut: %d, Err: %v", tt.name, tt.input, tt.offset, domain, offset, err)
127
128 if tt.expectedError != nil {
129 assert.Error(t, err, "Expected an error but got nil")
130 if tt.errorCheck != nil {
131 tt.errorCheck(t, err)
132 } else {
133 assert.IsType(t, tt.expectedError, err, "Error type mismatch")
134 }
135 } else {
136 assert.NoError(t, err, "Expected no error but got one")
137 }
138
139 assert.Equal(t, tt.expectedDomain, domain, "Domain mismatch")
140 if tt.expectedError == nil {
141 assert.Equal(t, tt.expectedOffset, offset, "Offset mismatch")
142 }
143 })
144 }
145}
146
147func TestEncodeDomain(t *testing.T) {
148 tests := []struct {
149 name string
150 input string
151 initialBuf []byte
152 offsets map[string]uint16
153 expected []byte
154 expectedErr error
155 newOffsets map[string]uint16
156 }{
157 {
158 name: "Simple domain",
159 input: "example.com",
160 initialBuf: []byte{},
161 offsets: make(map[string]uint16),
162 expected: []byte{7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0},
163 newOffsets: map[string]uint16{"example.com": 0, "com": 8},
164 },
165 {
166 name: "Domain with existing offset for compression",
167 input: "test.example.com",
168 initialBuf: []byte{},
169 offsets: map[string]uint16{"example.com": 10},
170 expected: []byte{4, 't', 'e', 's', 't', 0xC0, 0x0A},
171 newOffsets: map[string]uint16{"test.example.com": 0, "example.com": 10},
172 },
173 {
174 name: "Multiple subdomains",
175 input: "a.b.c.d",
176 initialBuf: []byte{},
177 offsets: make(map[string]uint16),
178 expected: []byte{1, 'a', 1, 'b', 1, 'c', 1, 'd', 0},
179 newOffsets: map[string]uint16{"a.b.c.d": 0, "b.c.d": 2, "c.d": 4, "d": 6},
180 },
181 {
182 name: "Root domain",
183 input: ".",
184 initialBuf: []byte{},
185 offsets: make(map[string]uint16),
186 expected: []byte{0},
187 newOffsets: map[string]uint16{},
188 },
189 {
190 name: "Empty domain",
191 input: "",
192 initialBuf: []byte{},
193 offsets: make(map[string]uint16),
194 expected: []byte{0},
195 newOffsets: map[string]uint16{},
196 },
197 {
198 name: "Label too long",
199 input: "labeltoolonglabeltoolonglabeltoolonglabeltoolonglabeltoolonglabeltoolong.com",
200 initialBuf: []byte{},
201 offsets: make(map[string]uint16),
202 expected: nil,
203 expectedErr: &InvalidLabelError{Length: 72},
204 newOffsets: map[string]uint16{},
205 },
206 {
207 name: "Empty label inside domain",
208 input: "example..com",
209 initialBuf: []byte{},
210 offsets: make(map[string]uint16),
211 expected: nil,
212 expectedErr: &InvalidLabelError{Length: 0},
213 newOffsets: map[string]uint16{},
214 },
215 {
216 name: "Append to existing buffer",
217 input: "example.com",
218 initialBuf: []byte{0xAA, 0xBB},
219 offsets: make(map[string]uint16),
220 expected: []byte{0xAA, 0xBB, 7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0},
221 newOffsets: map[string]uint16{"example.com": 2, "com": 10},
222 },
223 }
224
225 for _, tt := range tests {
226 t.Run(tt.name, func(t *testing.T) {
227 currentOffsets := make(map[string]uint16)
228 for k, v := range tt.offsets {
229 currentOffsets[k] = v
230 }
231
232 result, err := encodeDomain(tt.initialBuf, tt.input, ¤tOffsets)
233
234 if tt.expectedErr != nil {
235 assert.Error(t, err, "Expected an error but got nil")
236 assert.IsType(t, tt.expectedErr, err, "Error type mismatch")
237 if expectedILE, ok := tt.expectedErr.(*InvalidLabelError); ok {
238 actualILE := &InvalidLabelError{}
239 if assert.True(t, errors.As(err, &actualILE)) {
240 assert.Equal(t, expectedILE.Length, actualILE.Length)
241 }
242 }
243 } else {
244 assert.NoError(t, err, "Expected no error but got one")
245 assert.Equal(t, tt.expected, result, "Encoded domain does not match expected output")
246 assert.Equal(t, tt.newOffsets, currentOffsets, "Offsets map does not match expected state")
247 }
248 })
249 }
250}
251
252func FuzzDecodeDomain(f *testing.F) {
253 testcases := [][]byte{
254 {
255 0x03, 0x63, 0x6f, 0x6d, 0x00,
256 },
257 {
258 0x03, 0x63, 0x6f, 0x6d, 0x00, 0x01, 0x63, 0xc0, 0x00,
259 },
260 {
261 0x03, 0x63, 0x6f, 0x6d, 0xc0, 0x00,
262 },
263 {
264 0xc0, 0x00,
265 },
266 {
267 0xc0, 0xff,
268 },
269 {
270 0x40,
271 },
272 {
273 0x03, 0x63, 0x6f,
274 },
275 {
276 0xc0,
277 },
278 }
279 for _, tc := range testcases {
280 f.Add(tc)
281 }
282 f.Fuzz(func(t *testing.T, msg []byte) {
283 _, _, err := decodeDomain(msg, 0)
284 if err != nil {
285 var bufErr *BufferOverflowError
286 var labelErr *InvalidLabelError
287 var compErr *DomainCompressionError
288
289 if !(errors.As(err, &bufErr) || errors.As(err, &labelErr) || errors.As(err, &compErr)) {
290 t.Errorf("Fuzzing decodeDomain: unexpected error type %T: %v for input %x", err, err, msg)
291 }
292 }
293 })
294}