+89 -23

domain_test.go

reviewed

··· 7 7 ) 8 8 9 9 func TestDecodeDomain(t *testing.T) { 10 10 - buf := []byte{ 11 11 - 0x03, 0x63, 0x6f, 0x6d, 0x00, 10 10 + tests := []struct { 11 11 + name string 12 12 + offset int 13 13 + input []byte 14 14 + expectedDomain string 15 15 + expectedOffset int 16 16 + expectedError error 17 17 + }{ 18 18 + { 19 19 + name: "Simple domain", 20 20 + input: []byte{3, 'w', 'w', 'w', 7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0}, 21 21 + expectedDomain: "www.example.com", 22 22 + expectedOffset: 17, 23 23 + expectedError: nil, 24 24 + }, 25 25 + { 26 26 + name: "Domain with compression", 27 27 + offset: 17, 28 28 + 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}, 29 29 + expectedDomain: "foo.www.example.com", 30 30 + expectedOffset: 23, 31 31 + expectedError: nil, 32 32 + }, 33 33 + { 34 34 + name: "Invalid label length", 35 35 + input: []byte{64, 'x'}, 36 36 + expectedDomain: "", 37 37 + expectedOffset: 2, 38 38 + expectedError: &InvalidLabelError{Length: 64}, 39 39 + }, 40 40 + { 41 41 + name: "Compression loop", 42 42 + input: []byte{0xC0, 0, 0xC0, 0}, 43 43 + expectedDomain: "", 44 44 + expectedOffset: 4, 45 45 + expectedError: &DomainCompressionError{}, 46 46 + }, 47 47 + { 48 48 + name: "Truncated input", 49 49 + input: []byte{3, 'w', 'w'}, 50 50 + expectedDomain: "", 51 51 + expectedOffset: 3, 52 52 + expectedError: &BufferOverflowError{Length: 3, Offset: 4}, 53 53 + }, 12 54 } 13 55 14 14 - domain, offset, err := decode_domain(buf, 0) 15 15 - assert.Equal(t, "com", domain) 16 16 - assert.Equal(t, 5, offset) 17 17 - assert.NoError(t, err) 18 18 - } 56 56 + for _, tt := range tests { 57 57 + t.Run(tt.name, func(t *testing.T) { 58 58 + domain, offset, err := decode_domain(tt.input, tt.offset) 19 59 20 20 - func TestDecodeDomainWithCompression(t *testing.T) { 21 21 - buf := []byte{ 22 22 - 0x03, 0x63, 0x6f, 0x6d, 0x00, 0x01, 0x63, 0xC0, 0x00, 60 60 + t.Log(tt.name) 61 61 + assert.Equal(t, tt.expectedError, err) 62 62 + assert.Equal(t, tt.expectedDomain, domain) 63 63 + assert.Equal(t, tt.expectedOffset, offset) 64 64 + }) 23 65 } 24 24 - 25 25 - domain, offset, err := decode_domain(buf, 5) 26 26 - assert.Equal(t, "c.com", domain) 27 27 - assert.Equal(t, 9, offset) 28 28 - assert.NoError(t, err) 29 66 } 30 67 31 31 - func TestDecodeDomainWithCompressionLoop(t *testing.T) { 32 32 - buf := []byte{ 33 33 - 0x03, 0x63, 0x6f, 0x6d, 0xC0, 0x00, 68 68 + func TestEncodeDomain(t *testing.T) { 69 69 + tests := []struct { 70 70 + name string 71 71 + input string 72 72 + offsets map[string]uint16 73 73 + expected []byte 74 74 + newOffsets map[string]uint16 75 75 + }{ 76 76 + { 77 77 + name: "Simple domain", 78 78 + input: "example.com", 79 79 + offsets: make(map[string]uint16), 80 80 + expected: []byte{7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0}, 81 81 + newOffsets: map[string]uint16{"example.com": 0, "com": 8}, 82 82 + }, 83 83 + { 84 84 + name: "Domain with existing offset", 85 85 + input: "test.example.com", 86 86 + offsets: map[string]uint16{"example.com": 10}, 87 87 + expected: []byte{4, 't', 'e', 's', 't', 0xC0, 0x0A}, 88 88 + newOffsets: map[string]uint16{"test.example.com": 0, "example.com": 10}, 89 89 + }, 90 90 + { 91 91 + name: "Multiple subdomains", 92 92 + input: "a.b.c.d", 93 93 + offsets: make(map[string]uint16), 94 94 + expected: []byte{1, 'a', 1, 'b', 1, 'c', 1, 'd', 0}, 95 95 + newOffsets: map[string]uint16{"a.b.c.d": 0, "b.c.d": 2, "c.d": 4, "d": 6}, 96 96 + }, 34 97 } 35 98 36 36 - domain, offset, err := decode_domain(buf, 0) 37 37 - assert.Equal(t, "", domain) 38 38 - assert.Equal(t, 6, offset) 39 39 - assert.Error(t, err) 99 99 + for _, tt := range tests { 100 100 + t.Run(tt.name, func(t *testing.T) { 101 101 + result := encode_domain([]byte{}, tt.input, &tt.offsets) 102 102 + assert.Equal(t, tt.expected, result, "Encoded domain does not match expected output") 103 103 + assert.Equal(t, tt.newOffsets, tt.offsets, "Offsets map does not match expected state") 104 104 + }) 105 105 + } 40 106 } 41 107 42 108 func FuzzDecodeDomain(f *testing.F) { ··· 52 118 }, 53 119 } 54 120 for _, tc := range testcases { 55 55 - f.Add(tc) // Use f.Add to provide a seed corpus 121 121 + f.Add(tc) 56 122 } 57 123 f.Fuzz(func(t *testing.T, msg []byte) { 58 124 decode_domain(msg, 0)

+71

errors_test.go

reviewed

··· 1 1 + package magna 2 2 + 3 3 + import ( 4 4 + "testing" 5 5 + 6 6 + "github.com/stretchr/testify/assert" 7 7 + ) 8 8 + 9 9 + func TestBufferOverflowError(t *testing.T) { 10 10 + tests := []struct { 11 11 + name string 12 12 + length int 13 13 + offset int 14 14 + expected string 15 15 + }{ 16 16 + {"basic overflow", 10, 15, "magna: offset 15 is past the buffer length 10"}, 17 17 + {"zero length", 0, 5, "magna: offset 5 is past the buffer length 0"}, 18 18 + {"negative offset", 10, -1, "magna: offset -1 is past the buffer length 10"}, 19 19 + } 20 20 + 21 21 + for _, tt := range tests { 22 22 + t.Run(tt.name, func(t *testing.T) { 23 23 + err := &BufferOverflowError{Length: tt.length, Offset: tt.offset} 24 24 + assert.Equal(t, tt.expected, err.Error()) 25 25 + }) 26 26 + } 27 27 + } 28 28 + 29 29 + func TestInvalidLabelError(t *testing.T) { 30 30 + tests := []struct { 31 31 + name string 32 32 + length int 33 33 + expected string 34 34 + }{ 35 35 + {"zero length", 0, "magna: received invalid label length 0"}, 36 36 + {"negative length", -1, "magna: received invalid label length -1"}, 37 37 + {"large length", 1000, "magna: received invalid label length 1000"}, 38 38 + } 39 39 + 40 40 + for _, tt := range tests { 41 41 + t.Run(tt.name, func(t *testing.T) { 42 42 + err := &InvalidLabelError{Length: tt.length} 43 43 + assert.Equal(t, tt.expected, err.Error()) 44 44 + }) 45 45 + } 46 46 + } 47 47 + 48 48 + func TestDomainCompressionError(t *testing.T) { 49 49 + err := &DomainCompressionError{} 50 50 + expected := "magna: loop detected in domain compression" 51 51 + assert.Equal(t, expected, err.Error()) 52 52 + } 53 53 + 54 54 + func TestMagnaError(t *testing.T) { 55 55 + tests := []struct { 56 56 + name string 57 57 + message string 58 58 + expected string 59 59 + }{ 60 60 + {"empty message", "", "magna: "}, 61 61 + {"basic message", "test error", "magna: test error"}, 62 62 + {"message with punctuation", "error: invalid input!", "magna: error: invalid input!"}, 63 63 + } 64 64 + 65 65 + for _, tt := range tests { 66 66 + t.Run(tt.name, func(t *testing.T) { 67 67 + err := &MagnaError{Message: tt.message} 68 68 + assert.Equal(t, tt.expected, err.Error()) 69 69 + }) 70 70 + } 71 71 + }

+348 -49

header_test.go

reviewed

··· 1 1 package magna 2 2 3 3 import ( 4 4 + "encoding/binary" 4 5 "testing" 5 6 6 7 "github.com/stretchr/testify/assert" 7 8 ) 8 9 9 10 func TestHeaderDecode(t *testing.T) { 10 10 - bytes := []byte{ 11 11 - 0x01, 0x02, // ID 12 12 - 0xaa, 0xaa, // QR, Opcode, AA, TC, RD, RA, Z, RCODE 13 13 - 0x00, 0x01, // QDCOUNT 14 14 - 0x00, 0x02, // ANCOUNT 15 15 - 0x00, 0x03, // NSCOUNT 16 16 - 0x00, 0x04, // ARCOUNT 11 11 + tests := []struct { 12 12 + name string 13 13 + input []byte 14 14 + expectedHeader Header 15 15 + expectedOffset int 16 16 + expectedErr error 17 17 + }{ 18 18 + { 19 19 + name: "Valid header", 20 20 + input: []byte{0x12, 0x34, 0x81, 0x80, 0x00, 0x01, 0x00, 0x02, 0x00, 0x03, 0x00, 0x04}, 21 21 + expectedHeader: Header{ 22 22 + ID: 0x1234, 23 23 + QR: true, 24 24 + OPCode: OPCode(0), 25 25 + AA: false, 26 26 + TC: false, 27 27 + RD: true, 28 28 + RA: true, 29 29 + Z: 0, 30 30 + RCode: RCode(0), 31 31 + QDCount: 1, 32 32 + ANCount: 2, 33 33 + NSCount: 3, 34 34 + ARCount: 4, 35 35 + }, 36 36 + expectedOffset: 12, 37 37 + expectedErr: nil, 38 38 + }, 39 39 + { 40 40 + name: "Insufficient buffer length", 41 41 + input: []byte{0x12, 0x34, 0x81}, 42 42 + expectedHeader: Header{}, 43 43 + expectedOffset: 3, 44 44 + expectedErr: &BufferOverflowError{Length: 3, Offset: 3}, 45 45 + }, 46 46 + { 47 47 + name: "Invalid ID", 48 48 + input: []byte{0x12}, 49 49 + expectedHeader: Header{}, 50 50 + expectedOffset: 1, 51 51 + expectedErr: &BufferOverflowError{Length: 1, Offset: 1}, 52 52 + }, 53 53 + { 54 54 + name: "Missing QDCount", 55 55 + input: []byte{0x12, 0x34, 0x81, 0x80, 0x00}, 56 56 + expectedHeader: Header{}, 57 57 + expectedOffset: 5, 58 58 + expectedErr: &BufferOverflowError{}, 59 59 + }, 60 60 + { 61 61 + name: "Missing ANCount", 62 62 + input: []byte{0x12, 0x34, 0x81, 0x80, 0x00, 0x01}, 63 63 + expectedHeader: Header{}, 64 64 + expectedOffset: 6, 65 65 + expectedErr: &BufferOverflowError{}, 66 66 + }, 67 67 + { 68 68 + name: "Missing NSCount", 69 69 + input: []byte{0x12, 0x34, 0x81, 0x80, 0x00, 0x01, 0x00, 0x02}, 70 70 + expectedHeader: Header{}, 71 71 + expectedOffset: 8, 72 72 + expectedErr: &BufferOverflowError{}, 73 73 + }, 74 74 + { 75 75 + name: "Missing ARCount", 76 76 + input: []byte{0x12, 0x34, 0x81, 0x80, 0x00, 0x01, 0x00, 0x02, 0x00, 0x03}, 77 77 + expectedHeader: Header{}, 78 78 + expectedOffset: 10, 79 79 + expectedErr: &BufferOverflowError{}, 80 80 + }, 17 81 } 18 82 19 19 - var header Header 20 20 - offset, err := header.Decode(bytes, 0) 21 21 - if err != nil { 22 22 - t.Errorf("error should be nil\n") 83 83 + for _, tt := range tests { 84 84 + t.Run(tt.name, func(t *testing.T) { 85 85 + h := &Header{} 86 86 + offset, err := h.Decode(tt.input, 0) 87 87 + 88 88 + if tt.expectedErr != nil { 89 89 + assert.Error(t, err) 90 90 + assert.IsType(t, tt.expectedErr, err) 91 91 + } else { 92 92 + assert.NoError(t, err) 93 93 + assert.Equal(t, tt.expectedHeader, *h) 94 94 + } 95 95 + 96 96 + assert.Equal(t, tt.expectedOffset, offset) 97 97 + }) 23 98 } 99 99 + } 24 100 25 25 - if offset != 12 { 26 26 - t.Errorf("offset should be 12 not %v\n", offset) 101 101 + func TestHeaderDecodeFlags(t *testing.T) { 102 102 + tests := []struct { 103 103 + name string 104 104 + flags uint16 105 105 + expected Header 106 106 + }{ 107 107 + { 108 108 + name: "All flags set", 109 109 + flags: 0xFFFF, 110 110 + expected: Header{ 111 111 + QR: true, 112 112 + OPCode: OPCode(15), 113 113 + AA: true, 114 114 + TC: true, 115 115 + RD: true, 116 116 + RA: true, 117 117 + Z: 7, 118 118 + RCode: RCode(15), 119 119 + }, 120 120 + }, 121 121 + { 122 122 + name: "No flags set", 123 123 + flags: 0x0000, 124 124 + expected: Header{ 125 125 + QR: false, 126 126 + OPCode: OPCode(0), 127 127 + AA: false, 128 128 + TC: false, 129 129 + RD: false, 130 130 + RA: false, 131 131 + Z: 0, 132 132 + RCode: RCode(0), 133 133 + }, 134 134 + }, 135 135 + { 136 136 + name: "Mixed flags", 137 137 + flags: 0x8510, 138 138 + expected: Header{ 139 139 + QR: true, 140 140 + OPCode: OPCode(0), 141 141 + AA: true, 142 142 + TC: false, 143 143 + RD: true, 144 144 + RA: false, 145 145 + Z: 1, 146 146 + RCode: RCode(0), 147 147 + }, 148 148 + }, 27 149 } 28 150 29 29 - assert.Equal(t, header.ID, uint16(258)) 30 30 - assert.Equal(t, header.QR, true) 31 31 - assert.Equal(t, header.OPCode, OPCode(5)) 32 32 - assert.Equal(t, header.AA, false) 33 33 - assert.Equal(t, header.TC, true) 34 34 - assert.Equal(t, header.RD, false) 35 35 - assert.Equal(t, header.RA, true) 36 36 - assert.Equal(t, header.Z, uint8(0b010)) 37 37 - assert.Equal(t, header.RCode, RCode(0b1010)) 38 38 - assert.Equal(t, header.QDCount, uint16(1)) 39 39 - assert.Equal(t, header.ANCount, uint16(2)) 40 40 - assert.Equal(t, header.NSCount, uint16(3)) 41 41 - assert.Equal(t, header.ARCount, uint16(4)) 151 151 + for _, tt := range tests { 152 152 + t.Run(tt.name, func(t *testing.T) { 153 153 + input := []byte{ 154 154 + 0x00, 0x00, 155 155 + byte(tt.flags >> 8), byte(tt.flags), 156 156 + 0x00, 0x00, 157 157 + 0x00, 0x00, 158 158 + 0x00, 0x00, 159 159 + 0x00, 0x00, 160 160 + } 161 161 + 162 162 + h := &Header{} 163 163 + _, err := h.Decode(input, 0) 164 164 + 165 165 + assert.NoError(t, err) 166 166 + assert.Equal(t, tt.expected.QR, h.QR) 167 167 + assert.Equal(t, tt.expected.OPCode, h.OPCode) 168 168 + assert.Equal(t, tt.expected.AA, h.AA) 169 169 + assert.Equal(t, tt.expected.TC, h.TC) 170 170 + assert.Equal(t, tt.expected.RD, h.RD) 171 171 + assert.Equal(t, tt.expected.RA, h.RA) 172 172 + assert.Equal(t, tt.expected.Z, h.Z) 173 173 + assert.Equal(t, tt.expected.RCode, h.RCode) 174 174 + }) 175 175 + } 42 176 } 43 177 44 178 func TestHeaderEncode(t *testing.T) { 45 45 - bytes := []byte{ 46 46 - 0x01, 0x02, // ID 47 47 - 0xaa, 0xaa, // QR, Opcode, AA, TC, RD, RA, Z, RCODE 48 48 - 0x00, 0x01, // QDCOUNT 49 49 - 0x00, 0x02, // ANCOUNT 50 50 - 0x00, 0x03, // NSCOUNT 51 51 - 0x00, 0x04, // ARCOUNT 179 179 + tests := []struct { 180 180 + name string 181 181 + header Header 182 182 + expected []byte 183 183 + }{ 184 184 + { 185 185 + name: "All fields set", 186 186 + header: Header{ 187 187 + ID: 0x1234, 188 188 + QR: true, 189 189 + OPCode: OPCode(1), 190 190 + AA: true, 191 191 + TC: true, 192 192 + RD: true, 193 193 + RA: true, 194 194 + Z: 5, 195 195 + RCode: RCode(3), 196 196 + QDCount: 1, 197 197 + ANCount: 2, 198 198 + NSCount: 3, 199 199 + ARCount: 4, 200 200 + }, 201 201 + expected: []byte{ 202 202 + 0x12, 0x34, 203 203 + 0x8f, 0xd3, 204 204 + 0x00, 0x01, 205 205 + 0x00, 0x02, 206 206 + 0x00, 0x03, 207 207 + 0x00, 0x04, 208 208 + }, 209 209 + }, 210 210 + { 211 211 + name: "No flags set", 212 212 + header: Header{ 213 213 + ID: 0x5678, 214 214 + QR: false, 215 215 + OPCode: OPCode(0), 216 216 + AA: false, 217 217 + TC: false, 218 218 + RD: false, 219 219 + RA: false, 220 220 + Z: 0, 221 221 + RCode: RCode(0), 222 222 + QDCount: 5, 223 223 + ANCount: 6, 224 224 + NSCount: 7, 225 225 + ARCount: 8, 226 226 + }, 227 227 + expected: []byte{ 228 228 + 0x56, 0x78, 229 229 + 0x00, 0x00, 230 230 + 0x00, 0x05, 231 231 + 0x00, 0x06, 232 232 + 0x00, 0x07, 233 233 + 0x00, 0x08, 234 234 + }, 235 235 + }, 236 236 + { 237 237 + name: "Mixed flags", 238 238 + header: Header{ 239 239 + ID: 0x9abc, 240 240 + QR: true, 241 241 + OPCode: OPCode(2), 242 242 + AA: false, 243 243 + TC: true, 244 244 + RD: false, 245 245 + RA: true, 246 246 + Z: 3, 247 247 + RCode: RCode(4), 248 248 + QDCount: 9, 249 249 + ANCount: 10, 250 250 + NSCount: 11, 251 251 + ARCount: 12, 252 252 + }, 253 253 + expected: []byte{ 254 254 + 0x9a, 0xbc, 255 255 + 0x92, 0xb4, 256 256 + 0x00, 0x09, 257 257 + 0x00, 0x0a, 258 258 + 0x00, 0x0b, 259 259 + 0x00, 0x0c, 260 260 + }, 261 261 + }, 262 262 + } 263 263 + 264 264 + for _, tt := range tests { 265 265 + t.Run(tt.name, func(t *testing.T) { 266 266 + encoded := tt.header.Encode() 267 267 + assert.Equal(t, tt.expected, encoded) 268 268 + }) 52 269 } 270 270 + } 53 271 54 54 - var header Header 55 55 - _, err := header.Decode(bytes, 0) 272 272 + func TestHeaderEncodeDecodeRoundTrip(t *testing.T) { 273 273 + originalHeader := Header{ 274 274 + ID: 0xdcba, 275 275 + QR: true, 276 276 + OPCode: OPCode(3), 277 277 + AA: true, 278 278 + TC: false, 279 279 + RD: true, 280 280 + RA: false, 281 281 + Z: 6, 282 282 + RCode: RCode(2), 283 283 + QDCount: 13, 284 284 + ANCount: 14, 285 285 + NSCount: 15, 286 286 + ARCount: 16, 287 287 + } 288 288 + 289 289 + encoded := originalHeader.Encode() 290 290 + 291 291 + decodedHeader := &Header{} 292 292 + offset, err := decodedHeader.Decode(encoded, 0) 293 293 + 56 294 assert.NoError(t, err) 295 295 + assert.Equal(t, len(encoded), offset) 296 296 + assert.Equal(t, originalHeader, *decodedHeader) 297 297 + } 57 298 58 58 - actual := header.Encode() 59 59 - assert.Equal(t, bytes, actual) 299 299 + func TestHeaderEncodeFlagCombinations(t *testing.T) { 300 300 + testCases := []struct { 301 301 + name string 302 302 + header Header 303 303 + expected uint16 304 304 + }{ 305 305 + {"QR flag", Header{QR: true}, 0x8000}, 306 306 + {"OPCode", Header{OPCode: OPCode(5)}, 0x2800}, 307 307 + {"AA flag", Header{AA: true}, 0x0400}, 308 308 + {"TC flag", Header{TC: true}, 0x0200}, 309 309 + {"RD flag", Header{RD: true}, 0x0100}, 310 310 + {"RA flag", Header{RA: true}, 0x0080}, 311 311 + {"Z value", Header{Z: 5}, 0x0050}, 312 312 + {"RCode", Header{RCode: RCode(7)}, 0x0007}, 313 313 + {"All flags set", Header{QR: true, OPCode: OPCode(15), AA: true, TC: true, RD: true, RA: true, Z: 7, RCode: RCode(15)}, 0xffff}, 314 314 + } 315 315 + 316 316 + for _, tc := range testCases { 317 317 + t.Run(tc.name, func(t *testing.T) { 318 318 + encoded := tc.header.Encode() 319 319 + flags := binary.BigEndian.Uint16(encoded[2:4]) 320 320 + assert.Equal(t, tc.expected, flags) 321 321 + }) 322 322 + } 60 323 } 61 324 62 325 func FuzzDecodeHeader(f *testing.F) { 63 326 testcases := [][]byte{ 64 64 - { 65 65 - 0x01, 0x02, 66 66 - 0xaa, 0xaa, 67 67 - 0x00, 0x01, 68 68 - 0x00, 0x02, 69 69 - 0x00, 0x03, 70 70 - 0x00, 0x04, 71 71 - }, 327 327 + {0x12, 0x34, 0x81, 0x80, 0x00, 0x01, 0x00, 0x02, 0x00, 0x03, 0x00, 0x04}, 328 328 + {0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}, 329 329 + {0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff}, 72 330 } 331 331 + 73 332 for _, tc := range testcases { 74 74 - f.Add(tc) // Use f.Add to provide a seed corpus 333 333 + f.Add(tc) 75 334 } 76 76 - var header Header 77 77 - f.Fuzz(func(t *testing.T, msg []byte) { 78 78 - header.Decode(msg, 0) 335 335 + 336 336 + f.Fuzz(func(t *testing.T, data []byte) { 337 337 + // limit to only 12 bytes 338 338 + if len(data) > 12 { 339 339 + data = data[0:12] 340 340 + } 341 341 + 342 342 + h := &Header{} 343 343 + offset, err := h.Decode(data, 0) 344 344 + if err != nil { 345 345 + switch err.(type) { 346 346 + case *BufferOverflowError, *InvalidLabelError: 347 347 + // these are expected error types 348 348 + default: 349 349 + t.Errorf("unexpected error type: %T", err) 350 350 + } 351 351 + return 352 352 + } 353 353 + 354 354 + if offset != len(data) { 355 355 + t.Errorf("offset (%d) does not match data length (%d)", offset, len(data)) 356 356 + } 357 357 + 358 358 + if h.OPCode > 15 { 359 359 + t.Errorf("invalid OPCode: %d", h.OPCode) 360 360 + } 361 361 + 362 362 + if h.Z > 7 { 363 363 + t.Errorf("invalid Z value: %d", h.Z) 364 364 + } 365 365 + 366 366 + if h.RCode > 15 { 367 367 + t.Errorf("invalid RCode: %d", h.RCode) 368 368 + } 369 369 + 370 370 + encoded := h.Encode() 371 371 + if len(encoded) != len(data) { 372 372 + t.Errorf("encoded length (%d) does not match input length (%d)", len(encoded), len(data)) 373 373 + 374 374 + for i := 0; i < len(data); i++ { 375 375 + t.Errorf("mismatch at position: %d: encoded %02x, input: %02x", i, encoded[i], data[i]) 376 376 + } 377 377 + } 79 378 }) 80 379 }

+49

message.go

reviewed

··· 1 1 package magna 2 2 3 3 + import ( 4 4 + "math/rand" 5 5 + ) 6 6 + 3 7 // Decode decodes a DNS packet. 4 8 func (m *Message) Decode(buf []byte) (err error) { 5 9 offset, err := m.Header.Decode(buf, 0) ··· 74 78 75 79 return bytes 76 80 } 81 81 + 82 82 + func CreateRequest(op OPCode, rd bool) *Message { 83 83 + return &Message{ 84 84 + Header: Header{ 85 85 + ID: uint16(rand.Intn(65534) + 1), 86 86 + QR: false, 87 87 + OPCode: op, 88 88 + AA: false, 89 89 + TC: false, 90 90 + RD: rd, 91 91 + RA: false, 92 92 + Z: 0, 93 93 + RCode: NOERROR, 94 94 + QDCount: 0, 95 95 + ARCount: 0, 96 96 + NSCount: 0, 97 97 + ANCount: 0, 98 98 + }, 99 99 + Question: make([]Question, 0), 100 100 + Answer: make([]ResourceRecord, 0), 101 101 + Additional: make([]ResourceRecord, 0), 102 102 + Authority: make([]ResourceRecord, 0), 103 103 + } 104 104 + } 105 105 + 106 106 + func (m *Message) CreateReply(req *Message) *Message { 107 107 + m.Header.ID = req.Header.ID 108 108 + m.Header.QR = true 109 109 + m.Header.OPCode = req.Header.OPCode 110 110 + 111 111 + return m 112 112 + } 113 113 + 114 114 + func (m *Message) AddQuestion(q Question) *Message { 115 115 + m.Header.QDCount += 1 116 116 + m.Question = append(m.Question, q) 117 117 + 118 118 + return m 119 119 + } 120 120 + 121 121 + func (m *Message) SetRCode(rc RCode) *Message { 122 122 + m.Header.RCode = rc 123 123 + 124 124 + return m 125 125 + }

+105 -176

message_test.go

reviewed

··· 1 1 package magna 2 2 3 3 import ( 4 4 + "bytes" 5 5 + "encoding/binary" 6 6 + "net" 4 7 "testing" 5 8 6 9 "github.com/stretchr/testify/assert" 7 10 ) 8 11 9 12 func TestMessageDecode(t *testing.T) { 10 10 - bytes := []byte{ 11 11 - 0x8e, 0x19, 0x81, 0x80, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x04, 0x6e, 0x65, 12 12 - 0x77, 0x73, 0x0b, 0x79, 0x63, 0x6f, 0x6d, 0x62, 0x69, 0x6e, 0x61, 0x74, 0x6f, 0x72, 0x03, 13 13 - 0x63, 0x6f, 0x6d, 0x00, 0x00, 0x01, 0x00, 0x01, 0xc0, 0x0c, 0x00, 0x01, 0x00, 0x01, 0x00, 14 14 - 0x00, 0x00, 0x01, 0x00, 0x04, 0xd1, 0xd8, 0xe6, 0xcf, 13 13 + tests := []struct { 14 14 + name string 15 15 + input []byte 16 16 + expected Message 17 17 + wantErr bool 18 18 + }{ 19 19 + { 20 20 + name: "Valid DNS message with one question", 21 21 + input: func() []byte { 22 22 + buf := new(bytes.Buffer) 23 23 + binary.Write(buf, binary.BigEndian, uint16(1234)) 24 24 + binary.Write(buf, binary.BigEndian, uint16(0x0100)) 25 25 + binary.Write(buf, binary.BigEndian, uint16(1)) 26 26 + binary.Write(buf, binary.BigEndian, uint16(0)) 27 27 + binary.Write(buf, binary.BigEndian, uint16(0)) 28 28 + binary.Write(buf, binary.BigEndian, uint16(0)) 29 29 + buf.Write([]byte{3, 'w', 'w', 'w', 7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0}) 30 30 + binary.Write(buf, binary.BigEndian, uint16(1)) 31 31 + binary.Write(buf, binary.BigEndian, uint16(1)) 32 32 + return buf.Bytes() 33 33 + }(), 34 34 + expected: Message{ 35 35 + Header: Header{ 36 36 + ID: 1234, 37 37 + QR: false, 38 38 + RD: true, 39 39 + OPCode: 0, 40 40 + QDCount: 1, 41 41 + }, 42 42 + Question: []Question{ 43 43 + { 44 44 + QName: "www.example.com", 45 45 + QType: 1, 46 46 + QClass: 1, 47 47 + }, 48 48 + }, 49 49 + }, 50 50 + wantErr: false, 51 51 + }, 52 52 + { 53 53 + name: "Valid DNS message with one answer", 54 54 + input: func() []byte { 55 55 + buf := new(bytes.Buffer) 56 56 + binary.Write(buf, binary.BigEndian, uint16(5678)) 57 57 + binary.Write(buf, binary.BigEndian, uint16(0x8180)) 58 58 + binary.Write(buf, binary.BigEndian, uint16(0)) 59 59 + binary.Write(buf, binary.BigEndian, uint16(1)) 60 60 + binary.Write(buf, binary.BigEndian, uint16(0)) 61 61 + binary.Write(buf, binary.BigEndian, uint16(0)) 62 62 + buf.Write([]byte{3, 'w', 'w', 'w', 7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0}) 63 63 + binary.Write(buf, binary.BigEndian, uint16(1)) 64 64 + binary.Write(buf, binary.BigEndian, uint16(1)) 65 65 + binary.Write(buf, binary.BigEndian, uint32(3600)) 66 66 + binary.Write(buf, binary.BigEndian, uint16(4)) 67 67 + binary.Write(buf, binary.BigEndian, uint32(0x0A000001)) 68 68 + return buf.Bytes() 69 69 + }(), 70 70 + expected: Message{ 71 71 + Header: Header{ 72 72 + ID: 5678, 73 73 + QR: true, 74 74 + OPCode: 0, 75 75 + AA: false, 76 76 + RD: true, 77 77 + RA: true, 78 78 + RCode: 0, 79 79 + ANCount: 1, 80 80 + }, 81 81 + Answer: []ResourceRecord{ 82 82 + { 83 83 + Name: "www.example.com", 84 84 + RType: 1, 85 85 + RClass: 1, 86 86 + TTL: 3600, 87 87 + RDLength: 4, 88 88 + RData: &A{net.IP([]byte{10, 0, 0, 1})}, 89 89 + }, 90 90 + }, 91 91 + }, 92 92 + wantErr: false, 93 93 + }, 94 94 + { 95 95 + name: "Invalid input - empty buffer", 96 96 + input: []byte{}, 97 97 + wantErr: true, 98 98 + }, 15 99 } 16 100 17 17 - var msg Message 18 18 - msg.Decode(bytes) 19 19 - assert.Equal(t, uint16(0x8e19), msg.Header.ID) 20 20 - assert.Equal(t, true, msg.Header.QR) 21 21 - assert.Equal(t, OPCode(0), msg.Header.OPCode) 22 22 - assert.Equal(t, false, msg.Header.AA) 23 23 - assert.Equal(t, false, msg.Header.TC) 24 24 - assert.Equal(t, true, msg.Header.RD) 25 25 - assert.Equal(t, true, msg.Header.RA) 26 26 - assert.Equal(t, uint8(0), msg.Header.Z) 27 27 - assert.Equal(t, RCode(0), msg.Header.RCode) 28 28 - assert.Equal(t, uint16(1), msg.Header.QDCount) 29 29 - assert.Equal(t, uint16(1), msg.Header.ANCount) 30 30 - assert.Equal(t, uint16(0), msg.Header.NSCount) 31 31 - assert.Equal(t, uint16(0), msg.Header.ARCount) 101 101 + for _, tt := range tests { 102 102 + t.Run(tt.name, func(t *testing.T) { 103 103 + m := &Message{} 104 104 + err := m.Decode(tt.input) 32 105 33 33 - assert.Equal(t, 1, len(msg.Question)) 34 34 - 35 35 - question := msg.Question[0] 36 36 - assert.Equal(t, "news.ycombinator.com", question.QName) 37 37 - assert.Equal(t, DNSType(1), question.QType) 38 38 - assert.Equal(t, DNSClass(1), question.QClass) 39 39 - 40 40 - assert.Equal(t, 1, len(msg.Answer)) 41 41 - answer := msg.Answer[0] 42 42 - assert.Equal(t, answer.Name, "news.ycombinator.com") 43 43 - assert.Equal(t, DNSType(1), answer.RType) 44 44 - assert.Equal(t, DNSClass(1), answer.RClass) 45 45 - assert.Equal(t, uint32(1), answer.TTL) 46 46 - assert.Equal(t, uint16(4), answer.RDLength) 47 47 - } 48 48 - 49 49 - func TestMessageDecodeWithU14Offset(t *testing.T) { 50 50 - bytes := []byte{ 51 51 - 0x0e, 0xc3, 0x80, 0x00, 0x00, 0x01, 0x00, 0x00, 52 52 - 0x00, 0x0d, 0x00, 0x0e, 0x06, 0x6e, 0x73, 0x2d, 53 53 - 0x33, 0x37, 0x32, 0x09, 0x61, 0x77, 0x73, 0x64, 54 54 - 0x6e, 0x73, 0x2d, 0x34, 0x36, 0x03, 0x63, 0x6f, 55 55 - 0x6d, 0x00, 0x00, 0x01, 0x00, 0x01, 0xc0, 0x1d, 56 56 - 0x00, 0x02, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 57 57 - 0x00, 0x14, 0x01, 0x61, 0x0c, 0x67, 0x74, 0x6c, 58 58 - 0x64, 0x2d, 0x73, 0x65, 0x72, 0x76, 0x65, 0x72, 59 59 - 0x73, 0x03, 0x6e, 0x65, 0x74, 0x00, 0xc0, 0x1d, 60 60 - 0x00, 0x02, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 61 61 - 0x00, 0x04, 0x01, 0x62, 0xc0, 0x34, 0xc0, 0x1d, 62 62 - 0x00, 0x02, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 63 63 - 0x00, 0x04, 0x01, 0x63, 0xc0, 0x34, 0xc0, 0x1d, 64 64 - 0x00, 0x02, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 65 65 - 0x00, 0x04, 0x01, 0x64, 0xc0, 0x34, 0xc0, 0x1d, 66 66 - 0x00, 0x02, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 67 67 - 0x00, 0x04, 0x01, 0x65, 0xc0, 0x34, 0xc0, 0x1d, 68 68 - 0x00, 0x02, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 69 69 - 0x00, 0x04, 0x01, 0x66, 0xc0, 0x34, 0xc0, 0x1d, 70 70 - 0x00, 0x02, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 71 71 - 0x00, 0x04, 0x01, 0x67, 0xc0, 0x34, 0xc0, 0x1d, 72 72 - 0x00, 0x02, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 73 73 - 0x00, 0x04, 0x01, 0x68, 0xc0, 0x34, 0xc0, 0x1d, 74 74 - 0x00, 0x02, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 75 75 - 0x00, 0x04, 0x01, 0x69, 0xc0, 0x34, 0xc0, 0x1d, 76 76 - 0x00, 0x02, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 77 77 - 0x00, 0x04, 0x01, 0x6a, 0xc0, 0x34, 0xc0, 0x1d, 78 78 - 0x00, 0x02, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 79 79 - 0x00, 0x04, 0x01, 0x6b, 0xc0, 0x34, 0xc0, 0x1d, 80 80 - 0x00, 0x02, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 81 81 - 0x00, 0x04, 0x01, 0x6c, 0xc0, 0x34, 0xc0, 0x1d, 82 82 - 0x00, 0x02, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 83 83 - 0x00, 0x04, 0x01, 0x6d, 0xc0, 0x34, 0xc0, 0x32, 84 84 - 0x00, 0x01, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 85 85 - 0x00, 0x04, 0xc0, 0x05, 0x06, 0x1e, 0xc0, 0x52, 86 86 - 0x00, 0x01, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 87 87 - 0x00, 0x04, 0xc0, 0x21, 0x0e, 0x1e, 0xc0, 0x62, 88 88 - 0x00, 0x01, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 89 89 - 0x00, 0x04, 0xc0, 0x1a, 0x5c, 0x1e, 0xc0, 0x72, 90 90 - 0x00, 0x01, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 91 91 - 0x00, 0x04, 0xc0, 0x1f, 0x50, 0x1e, 0xc0, 0x82, 92 92 - 0x00, 0x01, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 93 93 - 0x00, 0x04, 0xc0, 0x0c, 0x5e, 0x1e, 0xc0, 0x92, 94 94 - 0x00, 0x01, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 95 95 - 0x00, 0x04, 0xc0, 0x23, 0x33, 0x1e, 0xc0, 0xa2, 96 96 - 0x00, 0x01, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 97 97 - 0x00, 0x04, 0xc0, 0x2a, 0x5d, 0x1e, 0xc0, 0xb2, 98 98 - 0x00, 0x01, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 99 99 - 0x00, 0x04, 0xc0, 0x36, 0x70, 0x1e, 0xc0, 0xc2, 100 100 - 0x00, 0x01, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 101 101 - 0x00, 0x04, 0xc0, 0x2b, 0xac, 0x1e, 0xc0, 0xd2, 102 102 - 0x00, 0x01, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 103 103 - 0x00, 0x04, 0xc0, 0x30, 0x4f, 0x1e, 0xc0, 0xe2, 104 104 - 0x00, 0x01, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 105 105 - 0x00, 0x04, 0xc0, 0x34, 0xb2, 0x1e, 0xc0, 0xf2, 106 106 - 0x00, 0x01, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 107 107 - 0x00, 0x04, 0xc0, 0x29, 0xa2, 0x1e, 0xc1, 0x02, 108 108 - 0x00, 0x01, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 109 109 - 0x00, 0x04, 0xc0, 0x37, 0x53, 0x1e, 0xc0, 0x32, 110 110 - 0x00, 0x1c, 0x00, 0x01, 0x00, 0x02, 0xa3, 0x00, 111 111 - 0x00, 0x10, 0x20, 0x01, 0x05, 0x03, 0xa8, 0x3e, 112 112 - 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 113 113 - 0x00, 0x30, 106 106 + if tt.wantErr { 107 107 + assert.Error(t, err) 108 108 + } else { 109 109 + assert.NoError(t, err) 110 110 + assert.Equal(t, tt.expected.Header, m.Header) 111 111 + assert.Equal(t, tt.expected.Question, m.Question) 112 112 + assert.Equal(t, tt.expected.Answer, m.Answer) 113 113 + assert.Equal(t, tt.expected.Authority, m.Authority) 114 114 + assert.Equal(t, tt.expected.Additional, m.Additional) 115 115 + } 116 116 + }) 114 117 } 115 115 - 116 116 - var msg Message 117 117 - _ = msg.Decode(bytes) 118 118 - // assert_no_error(t, err) 119 119 - 120 120 - // Header Section 121 121 - assert.Equal(t, uint16(0x0ec3), msg.Header.ID) 122 122 - assert.Equal(t, true, msg.Header.QR) 123 123 - assert.Equal(t, QUERY, msg.Header.OPCode) 124 124 - assert.Equal(t, false, msg.Header.AA) 125 125 - assert.Equal(t, false, msg.Header.TC) 126 126 - assert.Equal(t, false, msg.Header.RD) 127 127 - assert.Equal(t, false, msg.Header.RA) 128 128 - assert.Equal(t, uint8(0), msg.Header.Z) 129 129 - assert.Equal(t, NOERROR, msg.Header.RCode) 130 130 - assert.Equal(t, uint16(1), msg.Header.QDCount) 131 131 - assert.Equal(t, uint16(0), msg.Header.ANCount) 132 132 - assert.Equal(t, uint16(13), msg.Header.NSCount) 133 133 - assert.Equal(t, uint16(14), msg.Header.ARCount) 134 134 - 135 135 - // Query Section 136 136 - assert.Equal(t, 1, len(msg.Question)) 137 137 - question := msg.Question[0] 138 138 - assert.Equal(t, "ns-372.awsdns-46.com", question.QName) 139 139 - assert.Equal(t, AType, question.QType) 140 140 - assert.Equal(t, IN, question.QClass) 141 141 - 142 142 - assert.Equal(t, 13, len(msg.Authority)) 143 143 - assert.Equal(t, 14, len(msg.Additional)) 144 144 - } 145 145 - 146 146 - func TestMessageEncode(t *testing.T) { 147 147 - bytes := []byte{ 148 148 - 0x8e, 0x19, 0x81, 0x80, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x04, 0x6e, 0x65, 149 149 - 0x77, 0x73, 0x0b, 0x79, 0x63, 0x6f, 0x6d, 0x62, 0x69, 0x6e, 0x61, 0x74, 0x6f, 0x72, 0x03, 150 150 - 0x63, 0x6f, 0x6d, 0x00, 0x00, 0x01, 0x00, 0x01, 0xc0, 0x0c, 0x00, 0x01, 0x00, 0x01, 0x00, 151 151 - 0x00, 0x00, 0x01, 0x00, 0x04, 0xd1, 0xd8, 0xe6, 0xcf, 152 152 - } 153 153 - 154 154 - var msg Message 155 155 - err := msg.Decode(bytes) 156 156 - assert.NoError(t, err) 157 157 - 158 158 - actual := msg.Encode() 159 159 - assert.Equal(t, bytes, actual) 160 160 - } 161 161 - 162 162 - func TestMessageEncode2(t *testing.T) { 163 163 - bytes := []byte{ 164 164 - 0xfc, 0xa9, 0x81, 0x80, 0x00, 0x01, 0x00, 0x05, 165 165 - 0x00, 0x00, 0x00, 0x00, 0x03, 0x6f, 0x6c, 0x64, 166 166 - 0x06, 0x72, 0x65, 0x64, 0x64, 0x69, 0x74, 0x03, 167 167 - 0x63, 0x6f, 0x6d, 0x00, 0x00, 0x01, 0x00, 0x01, 168 168 - 0xc0, 0x0c, 0x00, 0x05, 0x00, 0x01, 0x00, 0x00, 169 169 - 0x2a, 0x17, 0x00, 0x17, 0x06, 0x72, 0x65, 0x64, 170 170 - 0x64, 0x69, 0x74, 0x03, 0x6d, 0x61, 0x70, 0x06, 171 171 - 0x66, 0x61, 0x73, 0x74, 0x6c, 0x79, 0x03, 0x6e, 172 172 - 0x65, 0x74, 0x00, 0xc0, 0x2c, 0x00, 0x01, 0x00, 173 173 - 0x01, 0x00, 0x00, 0x00, 0x23, 0x00, 0x04, 0x97, 174 174 - 0x65, 0x01, 0x8c, 0xc0, 0x2c, 0x00, 0x01, 0x00, 175 175 - 0x01, 0x00, 0x00, 0x00, 0x23, 0x00, 0x04, 0x97, 176 176 - 0x65, 0xc1, 0x8c, 0xc0, 0x2c, 0x00, 0x01, 0x00, 177 177 - 0x01, 0x00, 0x00, 0x00, 0x23, 0x00, 0x04, 0x97, 178 178 - 0x65, 0x41, 0x8c, 0xc0, 0x2c, 0x00, 0x01, 0x00, 179 179 - 0x01, 0x00, 0x00, 0x00, 0x23, 0x00, 0x04, 0x97, 180 180 - 0x65, 0x81, 0x8c, 181 181 - } 182 182 - 183 183 - var msg Message 184 184 - err := msg.Decode(bytes) 185 185 - assert.NoError(t, err) 186 186 - 187 187 - actual := msg.Encode() 188 188 - assert.Equal(t, bytes, actual) 189 118 } 190 119 191 120 func FuzzDecodeMessage(f *testing.F) { ··· 198 127 }, 199 128 } 200 129 for _, tc := range testcases { 201 201 - f.Add(tc) // Use f.Add to provide a seed corpus 130 130 + f.Add(tc) 202 131 } 203 132 f.Fuzz(func(t *testing.T, msg []byte) { 204 133 var m Message

+240

question_test.go

reviewed

··· 1 1 + package magna 2 2 + 3 3 + import ( 4 4 + "testing" 5 5 + 6 6 + "github.com/stretchr/testify/assert" 7 7 + ) 8 8 + 9 9 + func TestQuestionDecode(t *testing.T) { 10 10 + tests := []struct { 11 11 + name string 12 12 + input []byte 13 13 + expectedOffset int 14 14 + expected Question 15 15 + expectedErr error 16 16 + }{ 17 17 + { 18 18 + name: "Valid question - example.com A IN", 19 19 + input: []byte{7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0, 0, 1, 0, 1}, 20 20 + expectedOffset: 17, 21 21 + expected: Question{ 22 22 + QName: "example.com", 23 23 + QType: DNSType(1), 24 24 + QClass: DNSClass(1), 25 25 + }, 26 26 + expectedErr: nil, 27 27 + }, 28 28 + { 29 29 + name: "Valid question - example.com MX CH", 30 30 + input: []byte{7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0, 0, 15, 0, 3}, 31 31 + expectedOffset: 17, 32 32 + expected: Question{ 33 33 + QName: "example.com", 34 34 + QType: DNSType(15), 35 35 + QClass: DNSClass(3), 36 36 + }, 37 37 + expectedErr: nil, 38 38 + }, 39 39 + { 40 40 + name: "Invalid domain name", 41 41 + input: []byte{255, 'i', 'n', 'v', 'a', 'l', 'i', 'd', 0, 0, 1, 0, 1}, 42 42 + expectedOffset: 13, 43 43 + expected: Question{}, 44 44 + expectedErr: &BufferOverflowError{}, 45 45 + }, 46 46 + { 47 47 + name: "Insufficient buffer for QType", 48 48 + input: []byte{7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0, 0}, 49 49 + expectedOffset: 14, 50 50 + expected: Question{QName: "example.com"}, 51 51 + expectedErr: &BufferOverflowError{}, 52 52 + }, 53 53 + { 54 54 + name: "Insufficient buffer for QClass", 55 55 + input: []byte{7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0, 0, 1, 0}, 56 56 + expectedOffset: 16, 57 57 + expected: Question{QName: "example.com", QType: DNSType(1)}, 58 58 + expectedErr: &BufferOverflowError{}, 59 59 + }, 60 60 + } 61 61 + 62 62 + for _, tt := range tests { 63 63 + t.Run(tt.name, func(t *testing.T) { 64 64 + q := &Question{} 65 65 + offset, err := q.Decode(tt.input, 0) 66 66 + 67 67 + assert.Equal(t, tt.expectedOffset, offset) 68 68 + 69 69 + if tt.expectedErr != nil { 70 70 + assert.Error(t, err) 71 71 + assert.IsType(t, tt.expectedErr, err) 72 72 + } else { 73 73 + assert.NoError(t, err) 74 74 + assert.Equal(t, tt.expected, *q) 75 75 + } 76 76 + }) 77 77 + } 78 78 + } 79 79 + 80 80 + func TestQuestionEncode(t *testing.T) { 81 81 + tests := []struct { 82 82 + name string 83 83 + question Question 84 84 + offsets map[string]uint16 85 85 + expected []byte 86 86 + }{ 87 87 + { 88 88 + name: "Simple domain - example.com A IN", 89 89 + question: Question{ 90 90 + QName: "example.com", 91 91 + QType: DNSType(1), 92 92 + QClass: DNSClass(1), 93 93 + }, 94 94 + offsets: make(map[string]uint16), 95 95 + expected: []byte{7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0, 0, 1, 0, 1}, 96 96 + }, 97 97 + { 98 98 + name: "Subdomain - subdomain.example.com AAAA IN", 99 99 + question: Question{ 100 100 + QName: "subdomain.example.com", 101 101 + QType: DNSType(28), 102 102 + QClass: DNSClass(1), 103 103 + }, 104 104 + offsets: make(map[string]uint16), 105 105 + expected: []byte{9, 's', 'u', 'b', 'd', 'o', 'm', 'a', 'i', 'n', 7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0, 0, 28, 0, 1}, 106 106 + }, 107 107 + { 108 108 + name: "Different class - example.com MX CH", 109 109 + question: Question{ 110 110 + QName: "example.com", 111 111 + QType: DNSType(15), 112 112 + QClass: DNSClass(3), 113 113 + }, 114 114 + offsets: make(map[string]uint16), 115 115 + expected: []byte{7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0, 0, 15, 0, 3}, 116 116 + }, 117 117 + { 118 118 + name: "Domain compression - example.com after subdomain.example.com", 119 119 + question: Question{ 120 120 + QName: "example.com", 121 121 + QType: DNSType(1), 122 122 + QClass: DNSClass(1), 123 123 + }, 124 124 + offsets: map[string]uint16{ 125 125 + "com": 22, 126 126 + "example.com": 19, 127 127 + }, 128 128 + expected: []byte{0xC0, 0x13, 0x00, 0x01, 0x00, 0x01}, 129 129 + }, 130 130 + } 131 131 + 132 132 + for _, tt := range tests { 133 133 + t.Run(tt.name, func(t *testing.T) { 134 134 + result := tt.question.Encode(nil, &tt.offsets) 135 135 + assert.Equal(t, tt.expected, result) 136 136 + 137 137 + if len(tt.offsets) == 0 { 138 138 + expectedOffsets := map[string]uint16{ 139 139 + tt.question.QName: 0, 140 140 + } 141 141 + for i := 0; i < len(tt.question.QName); i++ { 142 142 + if tt.question.QName[i] == '.' { 143 143 + expectedOffsets[tt.question.QName[i+1:]] = uint16(i + 1) 144 144 + } 145 145 + } 146 146 + assert.Equal(t, expectedOffsets, tt.offsets) 147 147 + } 148 148 + }) 149 149 + } 150 150 + } 151 151 + 152 152 + func TestQuestionEncodeDecodeRoundTrip(t *testing.T) { 153 153 + tests := []struct { 154 154 + name string 155 155 + question Question 156 156 + }{ 157 157 + { 158 158 + name: "Simple domain - example.com A IN", 159 159 + question: Question{ 160 160 + QName: "example.com", 161 161 + QType: DNSType(1), 162 162 + QClass: DNSClass(1), 163 163 + }, 164 164 + }, 165 165 + { 166 166 + name: "Subdomain - subdomain.example.com AAAA IN", 167 167 + question: Question{ 168 168 + QName: "subdomain.example.com", 169 169 + QType: DNSType(28), 170 170 + QClass: DNSClass(1), 171 171 + }, 172 172 + }, 173 173 + { 174 174 + name: "Different class - example.com MX CH", 175 175 + question: Question{ 176 176 + QName: "example.com", 177 177 + QType: DNSType(15), 178 178 + QClass: DNSClass(3), 179 179 + }, 180 180 + }, 181 181 + } 182 182 + 183 183 + for _, tt := range tests { 184 184 + t.Run(tt.name, func(t *testing.T) { 185 185 + offsets := make(map[string]uint16) 186 186 + encoded := tt.question.Encode(nil, &offsets) 187 187 + 188 188 + decodedQuestion := &Question{} 189 189 + _, err := decodedQuestion.Decode(encoded, 0) 190 190 + 191 191 + assert.NoError(t, err) 192 192 + assert.Equal(t, tt.question, *decodedQuestion) 193 193 + }) 194 194 + } 195 195 + } 196 196 + 197 197 + func TestQuestionEncodeWithExistingBuffer(t *testing.T) { 198 198 + question := Question{ 199 199 + QName: "example.com", 200 200 + QType: DNSType(1), 201 201 + QClass: DNSClass(1), 202 202 + } 203 203 + 204 204 + existingBuffer := []byte{0xFF, 0xFF, 0xFF, 0xFF} 205 205 + offsets := make(map[string]uint16) 206 206 + 207 207 + result := question.Encode(existingBuffer, &offsets) 208 208 + 209 209 + expected := append( 210 210 + existingBuffer, 211 211 + []byte{7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'c', 'o', 'm', 0, 0, 1, 0, 1}..., 212 212 + ) 213 213 + 214 214 + assert.Equal(t, expected, result) 215 215 + } 216 216 + 217 217 + func TestQuestionEncodeLongDomainName(t *testing.T) { 218 218 + longLabel := make([]byte, 63) 219 219 + for i := range longLabel { 220 220 + longLabel[i] = 'a' 221 221 + } 222 222 + longDomainName := string(longLabel) + "." + string(longLabel) + "." + string(longLabel) + "." + string(longLabel[:61]) 223 223 + 224 224 + question := Question{ 225 225 + QName: longDomainName, 226 226 + QType: DNSType(1), 227 227 + QClass: DNSClass(1), 228 228 + } 229 229 + 230 230 + offsets := make(map[string]uint16) 231 231 + encoded := question.Encode(nil, &offsets) 232 232 + 233 233 + assert.Equal(t, 259, len(encoded)) 234 234 + 235 235 + decodedQuestion := &Question{} 236 236 + _, err := decodedQuestion.Decode(encoded, 0) 237 237 + 238 238 + assert.NoError(t, err) 239 239 + assert.Equal(t, question, *decodedQuestion) 240 240 + }