1 files changed, 103 insertions, 0 deletions
diff --git a/internal/dns/pack.go b/internal/dns/pack.go
new file mode 100644
index 0000000..dddc2dd
--- /dev/null
+++ b/internal/dns/pack.go
@@ -0,0 +1,103 @@
+package dns
+
+import (
+	"strings"
+)
+
+type packer struct {
+	buf   []byte
+	names map[string]int
+}
+
+func newPacker() *packer {
+	return &packer{
+		names: make(map[string]int),
+	}
+}
+
+func (p *packer) writeUint16(v uint16) {
+	p.buf = append(p.buf, byte(v>>8), byte(v))
+}
+
+func (p *packer) writeUint32(v uint32) {
+	p.buf = append(p.buf, byte(v>>24), byte(v>>16), byte(v>>8), byte(v))
+}
+
+func (p *packer) writeName(n Name) error {
+	labels := SplitDomainName(n.String())
+	for i := 0; i < len(labels); i++ {
+		suffix := Fqdn(joinLabels(labels[i:]))
+		if off, ok := p.names[suffix]; ok {
+			for _, label := range labels[:i] {
+				p.buf = append(p.buf, byte(len(label)))
+				p.buf = append(p.buf, label...)
+			}
+			p.writeUint16(uint16(off) | 0xC000)
+			return nil
+		}
+	}
+
+	off := len(p.buf)
+	p.buf = append(p.buf, n.Data...)
+
+	for i := 0; i < len(labels); i++ {
+		suffix := Fqdn(joinLabels(labels[i:]))
+		suffixOff := off
+		for j := 0; j < i; j++ {
+			suffixOff += 1 + len(labels[j])
+		}
+		if _, exists := p.names[suffix]; !exists {
+			p.names[suffix] = suffixOff
+		}
+	}
+	return nil
+}
+
+func joinLabels(labels []string) string {
+	if len(labels) == 0 {
+		return "."
+	}
+	return strings.Join(labels, ".") + "."
+}
+
+func Pack(m *Msg) ([]byte, error) {
+	p := newPacker()
+
+	headerOff := len(p.buf)
+	p.buf = append(p.buf, make([]byte, 12)...)
+
+	for _, q := range m.Question {
+		if err := q.pack(p); err != nil {
+			return nil, err
+		}
+	}
+	for _, rr := range m.Answer {
+		if err := rr.pack(p); err != nil {
+			return nil, err
+		}
+	}
+	for _, rr := range m.Ns {
+		if err := rr.pack(p); err != nil {
+			return nil, err
+		}
+	}
+	for _, rr := range m.Extra {
+		if err := rr.pack(p); err != nil {
+			return nil, err
+		}
+	}
+
+	putUint16(p.buf[headerOff:], m.Header.ID)
+	putUint16(p.buf[headerOff+2:], m.Header.Flags)
+	putUint16(p.buf[headerOff+4:], uint16(len(m.Question)))
+	putUint16(p.buf[headerOff+6:], uint16(len(m.Answer)))
+	putUint16(p.buf[headerOff+8:], uint16(len(m.Ns)))
+	putUint16(p.buf[headerOff+10:], uint16(len(m.Extra)))
+
+	return p.buf, nil
+}
+
+func putUint16(b []byte, v uint16) {
+	b[0] = byte(v >> 8)
+	b[1] = byte(v)
+}