Made clients for each protocol to reuse connections

internal/protocols/do53/do53.go

@@ -4,60 +4,62 @@ import (
 	"fmt"
 	"net"
 
-	"golang.org/x/net/dns/dnsmessage"
+	"github.com/miekg/dns"
 )
 
-func Run(domain, queryType, dest string, dnssec bool) error {
+type Do53Client struct {
+	udpConn *net.UDPConn
+}
 
-	message, err := MakeDNSMessage(domain, queryType)
-	if err != nil {
-		return err
-	}
+func New(dest string) (*Do53Client, error) {
 
 	udpAddr, err := net.ResolveUDPAddr("udp", dest)
 	if err != nil {
-		return fmt.Errorf("failed to resolve UDP address: %v", err)
+		return nil, fmt.Errorf("failed to resolve UDP address: %v", err)
 	}
 
 	udpConn, err := net.DialUDP("udp", nil, udpAddr)
 	if err != nil {
-		return fmt.Errorf("failed to dial UDP connection: %v", err)
+		return nil, fmt.Errorf("failed to dial UDP connection: %v", err)
 	}
-	defer udpConn.Close()
+	return &Do53Client{udpConn: udpConn}, nil
+}
 
-	_, err = udpConn.Write(message)
+func (c *Do53Client) Close() {
+	if c.udpConn != nil {
+		c.udpConn.Close()
+	}
+}
+
+func (c *Do53Client) Query(domain, queryType, dest string, dnssec bool) error {
+
+	message, err := NewDNSMessage(domain, queryType)
+	if err != nil {
+		return err
+	}
+
+	_, err = c.udpConn.Write(message)
 	if err != nil {
 		return fmt.Errorf("failed to send DNS query: %v", err)
 	}
 
 	buf := make([]byte, 4096)
-	n, err := udpConn.Read(buf)
+	n, err := c.udpConn.Read(buf)
 	if err != nil {
 		return fmt.Errorf("failed to read DNS response: %v", err)
 	}
 
-	var parser dnsmessage.Parser
-	_, err = parser.Start(buf[:n])
+	recvMsg := new(dns.Msg)
+	err = recvMsg.Unpack(buf[:n])
 	if err != nil {
 		return fmt.Errorf("failed to parse DNS response: %v", err)
 	}
 
 	// TODO: Check if the response had no errors or TD bit set
 
-	err = parser.SkipAllQuestions()
-	if err != nil {
-		return fmt.Errorf("failed to skip questions: %v", err)
-	}
-
-	answers, err := parser.AllAnswers()
-	if err != nil {
-		return err
-	}
-
-	for _, answer := range answers {
-		fmt.Println(answer.GoString())
+	for _, answer := range recvMsg.Answer {
+		fmt.Println(answer.String())
 	}
 
 	return nil
 }
-

internal/protocols/do53/packet.go

@@ -1,60 +1,40 @@
 package do53
 
 import (
-	"fmt"
-
-	"golang.org/x/net/dns/dnsmessage"
+	"github.com/miekg/dns"
 )
 
-func MakeDNSMessage(domain string, queryType string) ([]byte, error) {
-	messageHeader := dnsmessage.Header{
-		ID:               1234, // FIX: Use a random ID
-		Response:         false,
-		OpCode:           dnsmessage.OpCode(0),
-		RecursionDesired: true,
-	}
+func NewDNSMessage(domain string, queryType string) ([]byte, error) {
 
-	messageBuilder := dnsmessage.NewBuilder(nil, messageHeader)
-	queryName, err := dnsmessage.NewName(domain)
-	if err != nil {
-		return nil, fmt.Errorf("failed to create query name: %v", err)
-	}
-
-	// Determine query type
-	var queryTypeValue dnsmessage.Type
+	// TODO: Move this somewhere else and receive the type already parsed
+	var queryTypeValue uint16
 	switch queryType {
 	case "A":
-		queryTypeValue = dnsmessage.TypeA
+		queryTypeValue = dns.TypeA
 	case "AAAA":
-		queryTypeValue = dnsmessage.TypeAAAA
+		queryTypeValue = dns.TypeAAAA
 	case "MX":
-		queryTypeValue = dnsmessage.TypeMX
+		queryTypeValue = dns.TypeMX
 	case "CNAME":
-		queryTypeValue = dnsmessage.TypeCNAME
+		queryTypeValue = dns.TypeCNAME
 	case "TXT":
-		queryTypeValue = dnsmessage.TypeTXT
+		queryTypeValue = dns.TypeTXT
 	default:
-		queryTypeValue = dnsmessage.TypeA
+		queryTypeValue = dns.TypeA
 	}
 
-	messageQuestion := dnsmessage.Question{
-		Name:  queryName,
-		Type:  queryTypeValue,
-		Class: dnsmessage.ClassINET,
-	}
+	message := new(dns.Msg)
 
-	err = messageBuilder.StartQuestions()
+	message.Id = dns.Id()
+	message.Response = false
+	message.Opcode = dns.OpcodeQuery
+	message.Question = make([]dns.Question, 1)
+	message.Question[0] = dns.Question{Name: domain, Qtype: uint16(queryTypeValue), Qclass: dns.ClassINET}
+	message.Compress = true
+	wireMsg, err := message.Pack()
 	if err != nil {
 		return nil, err
 	}
-	err = messageBuilder.Question(messageQuestion)
-	if err != nil {
-		return nil, fmt.Errorf("failed to add question: %v", err)
-	}
 
-	message, err := messageBuilder.Finish()
-	if err != nil {
-		return nil, fmt.Errorf("failed to build message: %v", err)
-	}
-	return message, nil
+	return wireMsg, nil
 }