a go dns packet parser
at main 8.1 kB view raw
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, &currentOffsets) 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}