From b6b5b6c0b80798147dd8b27120aa888880998ccb Mon Sep 17 00:00:00 2001
From: jet2tlf <jet2tlf@gmail.com>
Date: Mon, 3 Jun 2024 02:51:52 -0300
Subject: codecrafters submit [skip ci]

---
 app/message.go | 79 +++++++++++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 70 insertions(+), 9 deletions(-)

(limited to 'app/message.go')

diff --git a/app/message.go b/app/message.go
index 931a6f0..c00346e 100644
--- a/app/message.go
+++ b/app/message.go
@@ -1,6 +1,9 @@
 package main
 
-import "encoding/binary"
+import (
+	"encoding/binary"
+	"strings"
+)
 
 type DNSHeader struct {
 	ID      uint16
@@ -33,6 +36,12 @@ type DNSAnswer struct {
 	RData    []byte
 }
 
+type DNSQuestion struct {
+	Name  []byte
+	Type  uint16
+	Class uint16
+}
+
 func (m *DNSHeader) Bytes() []byte {
 	bytes := make([]byte, 12)
 	binary.BigEndian.PutUint16(bytes[0:], m.ID)
@@ -50,10 +59,8 @@ func (m *DNSHeader) combineFlags() uint16 {
 		uint16(m.RCODE)
 }
 
-func (m *DNSMessage) AddQuestion(q []byte) {
-	m.Question = q
-	m.Question = binary.BigEndian.AppendUint16(m.Question, 1)
-	m.Question = binary.BigEndian.AppendUint16(m.Question, 1)
+func (m *DNSMessage) AddQuestion(q DNSQuestion) {
+	m.Question = q.Bytes()
 }
 
 func (m *DNSMessage) Bytes() []byte {
@@ -77,15 +84,23 @@ func (a *DNSAnswer) Bytes() []byte {
 	return bytes
 }
 
-func MakeMessage(header DNSHeader) DNSMessage {
-	return DNSMessage{Header: header, Question: []byte{}, Answer: DNSAnswer{}}
-}
-
 func (m *DNSMessage) AddAnswer(a DNSAnswer) {
 	m.Answer = a
 	m.Header.ANCOUNT = 1
 }
 
+func (q *DNSQuestion) Bytes() []byte {
+	bytes := []byte{}
+	bytes = append(bytes, q.Name...)
+	bytes = binary.BigEndian.AppendUint16(bytes, q.Type)
+	bytes = binary.BigEndian.AppendUint16(bytes, q.Class)
+	return bytes
+}
+
+func MakeMessage(header DNSHeader) DNSMessage {
+	return DNSMessage{Header: header, Question: []byte{}, Answer: DNSAnswer{}}
+}
+
 func MakeAnswer(name []byte, rdata []byte) DNSAnswer {
 	return DNSAnswer{Name: name, Type: 1, Class: 1, TTL: 60, RDLength: uint16(len(rdata)), RData: rdata}
 }
@@ -107,3 +122,49 @@ func ParseHeader(buf []byte) DNSHeader {
 		ARCOUNT: binary.BigEndian.Uint16(buf[10:12]),
 	}
 }
+
+func MakeQuestion(name []byte) DNSQuestion {
+	return DNSQuestion{Name: name, Type: 1, Class: 1}
+}
+
+func ParseQuestion(buf []byte) DNSQuestion {
+	return MakeQuestion(ParseDomain(buf))
+}
+
+func ParseDomain(data []byte) []byte {
+	domainByte := data[12:]
+	domain := decodeDNSPacket(domainByte)
+	segments := strings.Split(domain, ".")
+	var encodedDomain []byte
+
+	for _, segment := range segments {
+		encodedDomain = append(encodedDomain, byte(len(segment)))
+		encodedDomain = append(encodedDomain, []byte(segment)...)
+	}
+
+	encodedDomain = append(encodedDomain, 0x00)
+	return encodedDomain
+}
+
+func decodeDNSPacket(packet []byte) string {
+	var domain string
+	i := 0
+
+	for i < len(packet) && packet[i] != 0 {
+		labelLength := int(packet[i])
+
+		i++
+		if i+labelLength > len(packet) {
+			break
+		}
+
+		domain += string(packet[i : i+labelLength])
+
+		i += labelLength
+		if i < len(packet) && packet[i] != 0 {
+			domain += "."
+		}
+	}
+
+	return domain
+}
-- 
cgit v1.2.3