fiuh

1 files changed, 138 insertions, 0 deletions

diff --git a/internal/dns/unpack.go b/internal/dns/unpack.go
new file mode 100644
index 0000000..48703b1
--- /dev/null
+++ b/internal/dns/unpack.go
@@ -0,0 +1,138 @@
+package dns
+
+import "fmt"
+
+type unpacker struct {
+	buf []byte
+	off int
+}
+
+func (u *unpacker) readUint16() (uint16, error) {
+	if u.off+2 > len(u.buf) {
+		return 0, fmt.Errorf("dns: short uint16 at %d", u.off)
+	}
+	v := uint16(u.buf[u.off])<<8 | uint16(u.buf[u.off+1])
+	u.off += 2
+	return v, nil
+}
+
+func (u *unpacker) readUint32() (uint32, error) {
+	if u.off+4 > len(u.buf) {
+		return 0, fmt.Errorf("dns: short uint32 at %d", u.off)
+	}
+	v := uint32(u.buf[u.off])<<24 | uint32(u.buf[u.off+1])<<16 |
+		uint32(u.buf[u.off+2])<<8 | uint32(u.buf[u.off+3])
+	u.off += 4
+	return v, nil
+}
+
+func (u *unpacker) readName() (Name, error) {
+	var name Name
+	start := u.off
+	jumped := false
+	jumps := 0
+
+	for {
+		if u.off >= len(u.buf) {
+			return Name{}, fmt.Errorf("dns: name overflow")
+		}
+		c := u.buf[u.off]
+
+		// Compression pointer
+		if c&0xC0 == 0xC0 {
+			if u.off+2 > len(u.buf) {
+				return Name{}, fmt.Errorf("dns: short pointer")
+			}
+			ptr := int(uint16(c&0x3F)<<8 | uint16(u.buf[u.off+1]))
+			if ptr >= len(u.buf) {
+				return Name{}, fmt.Errorf("dns: pointer out of range")
+			}
+			if !jumped {
+				start = u.off + 2
+			}
+			u.off = ptr
+			jumped = true
+			jumps++
+			if jumps > 10 {
+				return Name{}, fmt.Errorf("dns: too many compression jumps")
+			}
+			continue
+		}
+
+		// Reserved label type (0x40, 0x80)
+		if c&0xC0 != 0 {
+			return Name{}, fmt.Errorf("dns: reserved label type")
+		}
+
+		// Root label
+		if c == 0 {
+			name.Data = append(name.Data, 0)
+			u.off++
+			break
+		}
+
+		// Normal label
+		length := int(c)
+		if length > 63 {
+			return Name{}, fmt.Errorf("dns: label too long")
+		}
+		if u.off+1+length > len(u.buf) {
+			return Name{}, fmt.Errorf("dns: short label")
+		}
+		name.Data = append(name.Data, u.buf[u.off:u.off+1+length]...)
+		u.off += 1 + length
+	}
+
+	if jumped {
+		u.off = start
+	}
+	return name, nil
+}
+
+func Unpack(buf []byte) (*Msg, error) {
+	if len(buf) < 12 {
+		return nil, fmt.Errorf("dns: message too short")
+	}
+	u := &unpacker{buf: buf}
+	m := &Msg{}
+
+	m.Header.ID, _ = u.readUint16()
+	m.Header.Flags, _ = u.readUint16()
+	qdcount, _ := u.readUint16()
+	ancount, _ := u.readUint16()
+	nscount, _ := u.readUint16()
+	arcount, _ := u.readUint16()
+
+	m.Question = make([]Question, qdcount)
+	for i := range m.Question {
+		if err := m.Question[i].unpack(u); err != nil {
+			return nil, err
+		}
+	}
+	m.Answer = make([]RR, ancount)
+	for i := range m.Answer {
+		rr, err := unpackRR(u)
+		if err != nil {
+			return nil, err
+		}
+		m.Answer[i] = rr
+	}
+	m.Ns = make([]RR, nscount)
+	for i := range m.Ns {
+		rr, err := unpackRR(u)
+		if err != nil {
+			return nil, err
+		}
+		m.Ns[i] = rr
+	}
+	m.Extra = make([]RR, arcount)
+	for i := range m.Extra {
+		rr, err := unpackRR(u)
+		if err != nil {
+			return nil, err
+		}
+		m.Extra[i] = rr
+	}
+
+	return m, nil
+}