codecrafters submit [skip ci]

2 files changed, 78 insertions, 8 deletions

diff --git a/app/main.go b/app/main.go
index b9be1e5..012d188 100644
--- a/app/main.go
+++ b/app/main.go
@@ -31,17 +31,28 @@ func main() {
 		receivedData := string(buf[:size])
 		fmt.Printf("Received %d bytes from %s: %s\n", size, source, receivedData)
 
-		response := make([]byte, 12)
-		copy(response, buf[:12])
-		response[2] = flipIndicator(response[2])
+		header := DNSHeader{
+			ID:      1234,
+			QR:      1,
+			OPCODE:  1,
+			AA:      0,
+			TC:      0,
+			RD:      0,
+			RA:      0,
+			Z:       0,
+			RCODE:   0,
+			QDCOUNT: 1,
+			ANCOUNT: 1,
+			NSCOUNT: 0,
+			ARCOUNT: 0,
+		}
+
+		response := MakeMessage(header)
+		response.AddQuestion([]byte("\x0ccodecrafters\x02io\x00"))
 
-		_, err = udpConn.WriteToUDP(response, source)
+		_, err = udpConn.WriteToUDP(response.bytes(), source)
 		if err != nil {
 			fmt.Println("Failed to send response:", err)
 		}
 	}
 }
-
-func flipIndicator(b byte) byte {
-	return b | 0b10000000
-}
diff --git a/app/message.go b/app/message.go
new file mode 100644
index 0000000..368eac0
--- /dev/null
+++ b/app/message.go
@@ -0,0 +1,59 @@
+package main
+
+import "encoding/binary"
+
+type DNSHeader struct {
+	ID      uint16
+	QR      uint8
+	OPCODE  uint8
+	AA      uint8
+	TC      uint8
+	RD      uint8
+	RA      uint8
+	Z       uint8
+	RCODE   uint8
+	QDCOUNT uint16
+	ANCOUNT uint16
+	NSCOUNT uint16
+	ARCOUNT uint16
+}
+
+type DNSMessage struct {
+	Header   DNSHeader
+	Question []byte
+}
+
+func (m *DNSHeader) Bytes() []byte {
+	bytes := make([]byte, 12)
+	binary.BigEndian.PutUint16(bytes[0:], m.ID)
+	binary.BigEndian.PutUint16(bytes[2:4], m.combineFlags())
+	binary.BigEndian.PutUint16(bytes[4:6], m.QDCOUNT)
+	binary.BigEndian.PutUint16(bytes[6:8], m.ANCOUNT)
+	binary.BigEndian.PutUint16(bytes[8:10], m.NSCOUNT)
+	binary.BigEndian.PutUint16(bytes[10:12], m.ARCOUNT)
+	return bytes
+}
+
+func (m *DNSHeader) combineFlags() uint16 {
+	return uint16(m.QR)<<15 | uint16(m.OPCODE)<<11 | uint16(m.AA)<<10 |
+		uint16(m.TC)<<9 | uint16(m.RD)<<8 | uint16(m.RA)<<7 | uint16(m.Z)<<4 |
+		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) bytes() []byte {
+	headerBytes := m.Header.Bytes()
+	bytes := make([]byte, len(headerBytes)+len(m.Question))
+	copy(bytes, headerBytes)
+	copy(bytes[len(headerBytes):], m.Question)
+	return bytes
+}
+
+func MakeMessage(header DNSHeader) DNSMessage {
+	return DNSMessage{Header: header, Question: []byte{}}
+}