codecrafters submit [skip ci]HEADmaster

3 files changed, 115 insertions, 4 deletions

diff --git a/app/forwarder.go b/app/forwarder.go
new file mode 100644
index 0000000..c4eac58
--- /dev/null
+++ b/app/forwarder.go
@@ -0,0 +1,60 @@
+package main
+
+import (
+	"fmt"
+	"net"
+)
+
+func ForwardQuestions(addr string, questions []DNSQuestion, id uint16) []DNSAnswer {
+	answers := []DNSAnswer{}
+
+	udpConn, err := net.Dial("udp", addr)
+	if err != nil {
+		fmt.Println("Failed to connect to address:", err)
+		return answers
+	}
+
+	defer udpConn.Close()
+
+	for _, question := range questions {
+		questionBytes := question.Bytes()
+
+		header := DNSHeader{
+			ID:      id,
+			QR:      0,
+			OPCODE:  0,
+			AA:      0,
+			TC:      0,
+			RD:      1,
+			RA:      0,
+			Z:       0,
+			RCODE:   0,
+			QDCOUNT: 1,
+			ANCOUNT: 0,
+			NSCOUNT: 0,
+			ARCOUNT: 0,
+		}
+
+		headerBytes := header.Bytes()
+		buf := append(headerBytes, questionBytes...)
+
+		_, err = udpConn.Write(buf)
+		if err != nil {
+			fmt.Println("Failed to send data:", err)
+			return answers
+		}
+
+		responseBuf := make([]byte, 512)
+
+		n, err := udpConn.Read(responseBuf)
+		if err != nil {
+			fmt.Println("Failed to receive data:", err)
+			return answers
+		}
+
+		responseHeader := ParseHeader(responseBuf[:n])
+		answers = append(answers, ParseAnswers(responseBuf[:n], responseHeader.QDCOUNT, responseHeader.ANCOUNT)...)
+	}
+
+	return answers
+}
diff --git a/app/main.go b/app/main.go
index b02972d..023988c 100644
--- a/app/main.go
+++ b/app/main.go
@@ -1,11 +1,16 @@
 package main
 
 import (
+	"flag"
 	"fmt"
 	"net"
 )
 
 func main() {
+	var resolver string
+	flag.StringVar(&resolver, "resolver", "1.1.1.1:53", "DNS resolver to forward queries to")
+	flag.Parse()
+
 	udpAddr, err := net.ResolveUDPAddr("udp", "127.0.0.1:2053")
 	if err != nil {
 		fmt.Println("Failed to resol[]byteve UDP address:", err)
@@ -31,7 +36,9 @@ func main() {
 		receivedData := string(buf[:size])
 		fmt.Printf("Received %d bytes from %s: %s\n", size, source, receivedData)
 		requestHeader := ParseHeader(buf[:size])
-		questions := ParseQuestions(buf[:size], requestHeader.QDCOUNT)
+		questions, _ := ParseQuestions(buf[:size], requestHeader.QDCOUNT)
+
+		answers := ForwardQuestions(resolver, questions, requestHeader.ID)
 
 		rcode := uint8(4)
 		if requestHeader.OPCODE == 0 {
@@ -58,7 +65,9 @@ func main() {
 
 		for _, question := range questions {
 			response.AddQuestion(question)
-			answer := MakeAnswer(question.Name, []byte("\x08\x08\x08\x08"))
+		}
+
+		for _, answer := range answers {
 			response.AddAnswer(answer)
 		}
 
diff --git a/app/message.go b/app/message.go
index 3d523ab..40c720d 100644
--- a/app/message.go
+++ b/app/message.go
@@ -136,7 +136,7 @@ func MakeQuestion(name []byte) DNSQuestion {
 	return DNSQuestion{Name: name, Type: 1, Class: 1}
 }
 
-func ParseQuestions(buf []byte, questionCount uint16) []DNSQuestion {
+func ParseQuestions(buf []byte, questionCount uint16) ([]DNSQuestion, int) {
 	questions := []DNSQuestion{}
 	offset := 12
 
@@ -148,7 +148,7 @@ func ParseQuestions(buf []byte, questionCount uint16) []DNSQuestion {
 		offset += len + 1 + 4
 	}
 
-	return questions
+	return questions, offset
 }
 
 func ParseDomain(data []byte, source []byte) []byte {
@@ -190,3 +190,45 @@ func decodeDNSPacket(packet []byte, source []byte) string {
 
 	return strings.Join(labels, ".")
 }
+
+func ParseAnswers(buf []byte, questionCount uint16, answerCount uint16) []DNSAnswer {
+	answers := []DNSAnswer{}
+	_, offset := ParseQuestions(buf, questionCount)
+
+	for i := 0; i < int(answerCount); i++ {
+		len := bytes.Index(buf[offset:], []byte{0})
+		label := ParseDomain(buf[offset:offset+len+1], buf)
+		offset += len + 1
+
+		answer := DNSAnswer{
+			Name: label,
+		}
+
+		answer.Type = extractUint16(buf, &offset)
+		answer.Class = extractUint16(buf, &offset)
+		answer.TTL = extractUint32(buf, &offset)
+		answer.RDLength = extractUint16(buf, &offset)
+		answer.RData = extractBytes(buf, &offset, int(answer.RDLength))
+		answers = append(answers, answer)
+	}
+
+	return answers
+}
+
+func extractBytes(src []byte, offset *int, length int) []byte {
+	result := src[*offset : *offset+length]
+	*offset += length
+	return result
+}
+
+func extractUint16(src []byte, offset *int) uint16 {
+	result := []byte{src[*offset], src[*offset+1]}
+	*offset += 2
+	return binary.BigEndian.Uint16(result[:])
+}
+
+func extractUint32(src []byte, offset *int) uint32 {
+	result := [4]byte{src[*offset], src[*offset+1], src[*offset+2], src[*offset+3]}
+	*offset += 4
+	return binary.BigEndian.Uint32(result[:])
+}