Made clients for each protocol to reuse connections

cmd/resolver/main.go

@@ -5,6 +5,7 @@ import (
 
 	"github.com/afonsofrancof/sdns-perf/internal/protocols/do53"
 	"github.com/afonsofrancof/sdns-perf/internal/protocols/doh"
+	"github.com/afonsofrancof/sdns-perf/internal/protocols/doq"
 	"github.com/afonsofrancof/sdns-perf/internal/protocols/dot"
 	"github.com/alecthomas/kong"
 )
@@ -45,20 +46,35 @@ var cli struct {
 }
 
 func (c *Do53Cmd) Run() error {
-	return do53.Run(c.DomainName, c.QueryType, c.Server, c.DNSSEC)
+	do53client, err := do53.New(c.Server)
+	if err != nil {
+		return err
+	}
+	return do53client.Query(c.DomainName, c.QueryType, c.Server, c.DNSSEC)
 }
 
 func (c *DoHCmd) Run() error {
-	return doh.Run(c.DomainName, c.QueryType, c.Server, c.Path,c.Proxy, c.DNSSEC)
+	dohclient, err := doh.New(c.Server, c.Path, c.Proxy)
+	if err != nil {
+		return err
+	}
+	return dohclient.Query(c.DomainName, c.QueryType, c.DNSSEC)
 }
 
 func (c *DoTCmd) Run() error {
-	return dot.Run(c.DomainName, c.QueryType, c.Server, c.DNSSEC)
+	dotclient, err := dot.New(c.Server)
+	if err != nil {
+		return err
+	}
+	return dotclient.Query(c.DomainName, c.QueryType, c.Server, c.DNSSEC)
 }
 
 func (c *DoQCmd) Run() error {
-	// TODO: Implement DoQ query
+	doqclient, err := doq.New(c.Server)
-	return nil
+	if err != nil {
+		return err
+	}
+	return doqclient.Query(c.DomainName, c.QueryType, c.DNSSEC)
 }
 
 func main() {

go.mod

@@ -4,5 +4,20 @@ go 1.24.0
 
 require (
 	github.com/alecthomas/kong v1.8.1
-	golang.org/x/net v0.35.0
+	github.com/miekg/dns v1.1.63
+	github.com/quic-go/quic-go v0.50.0
+)
+
+require (
+	github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 // indirect
+	github.com/google/pprof v0.0.0-20210407192527-94a9f03dee38 // indirect
+	github.com/onsi/ginkgo/v2 v2.9.5 // indirect
+	go.uber.org/mock v0.5.0 // indirect
+	golang.org/x/crypto v0.33.0 // indirect
+	golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842 // indirect
+	golang.org/x/mod v0.18.0 // indirect
+	golang.org/x/net v0.35.0 // indirect
+	golang.org/x/sync v0.8.0 // indirect
+	golang.org/x/sys v0.30.0 // indirect
+	golang.org/x/tools v0.22.0 // indirect
 )

go.sum

@@ -4,7 +4,63 @@ github.com/alecthomas/kong v1.8.1 h1:6aamvWBE/REnR/BCq10EcozmcpUPc5aGI1lPAWdB0EE
 github.com/alecthomas/kong v1.8.1/go.mod h1:p2vqieVMeTAnaC83txKtXe8FLke2X07aruPWXyMPQrU=
 github.com/alecthomas/repr v0.4.0 h1:GhI2A8MACjfegCPVq9f1FLvIBS+DrQ2KQBFZP1iFzXc=
 github.com/alecthomas/repr v0.4.0/go.mod h1:Fr0507jx4eOXV7AlPV6AVZLYrLIuIeSOWtW57eE/O/4=
+github.com/chzyer/logex v1.1.10/go.mod h1:+Ywpsq7O8HXn0nuIou7OrIPyXbp3wmkHB+jjWRnGsAI=
+github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e/go.mod h1:nSuG5e5PlCu98SY8svDHJxuZscDgtXS6KTTbou5AhLI=
+github.com/chzyer/test v0.0.0-20180213035817-a1ea475d72b1/go.mod h1:Q3SI9o4m/ZMnBNeIyt5eFwwo7qiLfzFZmjNmxjkiQlU=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/go-logr/logr v1.2.4 h1:g01GSCwiDw2xSZfjJ2/T9M+S6pFdcNtFYsp+Y43HYDQ=
+github.com/go-logr/logr v1.2.4/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
+github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 h1:tfuBGBXKqDEevZMzYi5KSi8KkcZtzBcTgAUUtapy0OI=
+github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572/go.mod h1:9Pwr4B2jHnOSGXyyzV8ROjYa2ojvAY6HCGYYfMoC3Ls=
+github.com/golang/protobuf v1.5.3 h1:KhyjKVUg7Usr/dYsdSqoFveMYd5ko72D+zANwlG1mmg=
+github.com/golang/protobuf v1.5.3/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY=
+github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
+github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/google/pprof v0.0.0-20210407192527-94a9f03dee38 h1:yAJXTCF9TqKcTiHJAE8dj7HMvPfh66eeA2JYW7eFpSE=
+github.com/google/pprof v0.0.0-20210407192527-94a9f03dee38/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE=
 github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM=
 github.com/hexops/gotextdiff v1.0.3/go.mod h1:pSWU5MAI3yDq+fZBTazCSJysOMbxWL1BSow5/V2vxeg=
+github.com/ianlancetaylor/demangle v0.0.0-20200824232613-28f6c0f3b639/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
+github.com/miekg/dns v1.1.63 h1:8M5aAw6OMZfFXTT7K5V0Eu5YiiL8l7nUAkyN6C9YwaY=
+github.com/miekg/dns v1.1.63/go.mod h1:6NGHfjhpmr5lt3XPLuyfDJi5AXbNIPM9PY6H6sF1Nfs=
+github.com/onsi/ginkgo/v2 v2.9.5 h1:+6Hr4uxzP4XIUyAkg61dWBw8lb/gc4/X5luuxN/EC+Q=
+github.com/onsi/ginkgo/v2 v2.9.5/go.mod h1:tvAoo1QUJwNEU2ITftXTpR7R1RbCzoZUOs3RonqW57k=
+github.com/onsi/gomega v1.27.6 h1:ENqfyGeS5AX/rlXDd/ETokDz93u0YufY1Pgxuy/PvWE=
+github.com/onsi/gomega v1.27.6/go.mod h1:PIQNjfQwkP3aQAH7lf7j87O/5FiNr+ZR8+ipb+qQlhg=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/quic-go/quic-go v0.50.0 h1:3H/ld1pa3CYhkcc20TPIyG1bNsdhn9qZBGN3b9/UyUo=
+github.com/quic-go/quic-go v0.50.0/go.mod h1:Vim6OmUvlYdwBhXP9ZVrtGmCMWa3wEqhq3NgYrI8b4E=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+go.uber.org/mock v0.5.0 h1:KAMbZvZPyBPWgD14IrIQ38QCyjwpvVVV6K/bHl1IwQU=
+go.uber.org/mock v0.5.0/go.mod h1:ge71pBPLYDk7QIi1LupWxdAykm7KIEFchiOqd6z7qMM=
+golang.org/x/crypto v0.33.0 h1:IOBPskki6Lysi0lo9qQvbxiQ+FvsCC/YWOecCHAixus=
+golang.org/x/crypto v0.33.0/go.mod h1:bVdXmD7IV/4GdElGPozy6U7lWdRXA4qyRVGJV57uQ5M=
+golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842 h1:vr/HnozRka3pE4EsMEg1lgkXJkTFJCVUX+S/ZT6wYzM=
+golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842/go.mod h1:XtvwrStGgqGPLc4cjQfWqZHG1YFdYs6swckp8vpsjnc=
+golang.org/x/mod v0.18.0 h1:5+9lSbEzPSdWkH32vYPBwEpX8KwDbM52Ud9xBUvNlb0=
+golang.org/x/mod v0.18.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
 golang.org/x/net v0.35.0 h1:T5GQRQb2y08kTAByq9L4/bz8cipCdA8FbRTXewonqY8=
 golang.org/x/net v0.35.0/go.mod h1:EglIi67kWsHKlRzzVMUD93VMSWGFOMSZgxFjparz1Qk=
+golang.org/x/sync v0.8.0 h1:3NFvSEYkUoMifnESzZl15y791HH1qU2xm6eCJU5ZPXQ=
+golang.org/x/sync v0.8.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
+golang.org/x/sys v0.0.0-20191204072324-ce4227a45e2e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc=
+golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/text v0.22.0 h1:bofq7m3/HAFvbF51jz3Q9wLg3jkvSPuiZu/pD1XwgtM=
+golang.org/x/text v0.22.0/go.mod h1:YRoo4H8PVmsu+E3Ou7cqLVH8oXWIHVoX0jqUWALQhfY=
+golang.org/x/time v0.5.0 h1:o7cqy6amK/52YcAKIPlM3a+Fpj35zvRj2TP+e1xFSfk=
+golang.org/x/time v0.5.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
+golang.org/x/tools v0.22.0 h1:gqSGLZqv+AI9lIQzniJ0nZDRG5GBPsSi+DRNHWNz6yA=
+golang.org/x/tools v0.22.0/go.mod h1:aCwcsjqvq7Yqt6TNyX7QMU2enbQ/Gt0bo6krSeEri+c=
+google.golang.org/protobuf v1.33.0 h1:uNO2rsAINq/JlFpSdYEKIZ0uKD/R9cpdv0T+yoGwGmI=
+google.golang.org/protobuf v1.33.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

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)
+func New(dest string) (*Do53Client, error) {
-	if err != nil {
-		return err
-	}
 
 	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
+	recvMsg := new(dns.Msg)
-	_, err = parser.Start(buf[:n])
+	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()
+	for _, answer := range recvMsg.Answer {
-	if err != nil {
+		fmt.Println(answer.String())
-		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())
 	}
 
 	return nil
 }
-

internal/protocols/do53/packet.go

@@ -1,60 +1,40 @@
 package do53
 
 import (
-	"fmt"
+	"github.com/miekg/dns"
-
-	"golang.org/x/net/dns/dnsmessage"
 )
 
-func MakeDNSMessage(domain string, queryType string) ([]byte, error) {
+func NewDNSMessage(domain string, queryType string) ([]byte, error) {
-	messageHeader := dnsmessage.Header{
-		ID:               1234, // FIX: Use a random ID
-		Response:         false,
-		OpCode:           dnsmessage.OpCode(0),
-		RecursionDesired: true,
-	}
 
-	messageBuilder := dnsmessage.NewBuilder(nil, messageHeader)
+	// TODO: Move this somewhere else and receive the type already parsed
-	queryName, err := dnsmessage.NewName(domain)
+	var queryTypeValue uint16
-	if err != nil {
-		return nil, fmt.Errorf("failed to create query name: %v", err)
-	}
-
-	// Determine query type
-	var queryTypeValue dnsmessage.Type
 	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{
+	message := new(dns.Msg)
-		Name:  queryName,
-		Type:  queryTypeValue,
-		Class: dnsmessage.ClassINET,
-	}
 
-	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()
+	return wireMsg, nil
-	if err != nil {
-		return nil, fmt.Errorf("failed to build message: %v", err)
-	}
-	return message, nil
 }

internal/protocols/doh/doh.go

@@ -11,37 +11,41 @@ import (
 	"os"
 
 	"github.com/afonsofrancof/sdns-perf/internal/protocols/do53"
-	"golang.org/x/net/dns/dnsmessage"
+	"github.com/miekg/dns"
 )
 
-func Run(domain, queryType, server, path, proxy string, dnssec bool) error {
+type DoHClient struct {
+	tcpConn    *net.TCPConn
+	tlsConn    *tls.Conn
+	keyLogFile *os.File
+	target     string
+	path       string
+	proxy      string
+}
 
-	DNSMessage, err := do53.MakeDNSMessage(domain, queryType)
+func New(target, path, proxy string) (*DoHClient, error) {
+
+	tcpAddr, err := net.ResolveTCPAddr("tcp", target)
 	if err != nil {
-		return err
+		return nil, fmt.Errorf("failed to resolve TCP address: %v", err)
 	}
 
-	// Step 1 - Establish a TCP Connection
+	tcpConn, err := net.DialTCP("tcp", nil, tcpAddr)
-	tcpConn, err := net.Dial("tcp", server)
 	if err != nil {
-		return fmt.Errorf("failed to establish TCP connection: %v", err)
+		return nil, fmt.Errorf("failed to establish TCP connection: %v", err)
 	}
-	defer tcpConn.Close()
 
-	// Step 2 - Upgrade it to a TLS Connection
-
-	// Temporary keylog file to allow traffic inspection
 	keyLogFile, err := os.OpenFile(
 		"tls-key-log.txt",
 		os.O_APPEND|os.O_CREATE|os.O_WRONLY,
 		0600,
 	)
 	if err != nil {
-		return fmt.Errorf("failed opening key log file: %v", err)
+		return nil, fmt.Errorf("failed opening key log file: %v", err)
 	}
-	defer keyLogFile.Close()
 
 	tlsConfig := &tls.Config{
+		// FIX: Actually check the domain name
 		InsecureSkipVerify: true,
 		MinVersion:         tls.VersionTLS12,
 		KeyLogWriter:       keyLogFile,
@@ -50,24 +54,45 @@ func Run(domain, queryType, server, path, proxy string, dnssec bool) error {
 	tlsConn := tls.Client(tcpConn, tlsConfig)
 	err = tlsConn.Handshake()
 	if err != nil {
-		return fmt.Errorf("failed to execute the TLS handshake: %v", err)
+		return nil, fmt.Errorf("failed to execute the TLS handshake: %v", err)
 	}
-	defer tlsConn.Close()
 
-	// Step 3 - Create an HTTP request with the do53 message in the body
+	return &DoHClient{tcpConn: tcpConn, keyLogFile: keyLogFile, tlsConn: tlsConn, target: target, path: path, proxy: proxy}, err
-	httpReq, err := http.NewRequest("POST", "https://"+server+"/"+path, bytes.NewBuffer(DNSMessage))
+
+}
+
+func (c *DoHClient) Close() {
+	if c.tcpConn != nil {
+		c.tcpConn.Close()
+	}
+	if c.keyLogFile != nil {
+		c.keyLogFile.Close()
+	}
+	if c.tlsConn != nil {
+		c.tlsConn.Close()
+	}
+}
+
+func (c *DoHClient) Query(domain, queryType string, dnssec bool) error {
+
+	DNSMessage, err := do53.NewDNSMessage(domain, queryType)
+	if err != nil {
+		return err
+	}
+
+	httpReq, err := http.NewRequest("POST", "https://"+c.target+"/"+c.path, bytes.NewBuffer(DNSMessage))
 	if err != nil {
 		return fmt.Errorf("failed to create HTTP request: %v", err)
 	}
 	httpReq.Header.Add("Content-Type", "application/dns-message")
 	httpReq.Header.Set("Accept", "application/dns-message")
 
-	err = httpReq.Write(tlsConn)
+	err = httpReq.Write(c.tlsConn)
 	if err != nil {
 		return fmt.Errorf("failed writing HTTP request: %v", err)
 	}
 
-	reader := bufio.NewReader(tlsConn)
+	reader := bufio.NewReader(c.tlsConn)
 	resp, err := http.ReadResponse(reader, httpReq)
 	if err != nil {
 		return fmt.Errorf("failed reading HTTP response: %v", err)
@@ -80,68 +105,16 @@ func Run(domain, queryType, server, path, proxy string, dnssec bool) error {
 		return fmt.Errorf("failed reading response body: %v", err)
 	}
 
-	// Parse the response
+	recvMsg := new(dns.Msg)
-	var parser dnsmessage.Parser
+	err = recvMsg.Unpack(responseBody[:n])
-	header, err := parser.Start(responseBody[:n])
 	if err != nil {
 		return fmt.Errorf("failed to parse DNS response: %v", err)
 	}
 
-	fmt.Printf("DNS Response Header:\n")
+	// TODO: Check if the response had no errors or TD bit set
-	fmt.Printf("  ID: %d\n", header.ID)
-	fmt.Printf("  Response: %v\n", header.Response)
-	fmt.Printf("  RCode: %v\n", header.RCode)
 
-	// Skip all questions before reading answers
+	for _, answer := range recvMsg.Answer {
-	err = parser.SkipAllQuestions()
+		fmt.Println(answer.String())
-	if err != nil {
-		return fmt.Errorf("failed to skip questions: %v", err)
-	}
-
-	// Parse answers
-	fmt.Printf("\nAnswers:\n")
-	answers, err := parser.AllAnswers()
-
-	for i, answer := range answers {
-
-		if err != nil {
-			return fmt.Errorf("failed to parse answer %d: %v", i, err)
-		}
-
-		fmt.Printf("  Answer %d:\n", i+1)
-		fmt.Printf("    Name: %v\n", answer.Header.Name)
-		fmt.Printf("    Type: %v\n", answer.Header.Type)
-		fmt.Printf("    TTL: %v seconds\n", answer.Header.TTL)
-
-		// Handle different record types
-		switch answer.Header.Type {
-		case dnsmessage.TypeA:
-			if r, ok := answer.Body.(*dnsmessage.AResource); ok {
-				fmt.Printf("    IPv4: %d.%d.%d.%d\n", r.A[0], r.A[1], r.A[2], r.A[3])
-			}
-		case dnsmessage.TypeAAAA:
-			if r, ok := answer.Body.(*dnsmessage.AAAAResource); ok {
-				ip := r.AAAA
-				fmt.Printf("    IPv6: %02x%02x:%02x%02x:%02x%02x:%02x%02x:%02x%02x:%02x%02x:%02x%02x:%02x%02x\n",
-					ip[0], ip[1], ip[2], ip[3], ip[4], ip[5], ip[6], ip[7],
-					ip[8], ip[9], ip[10], ip[11], ip[12], ip[13], ip[14], ip[15])
-			}
-		case dnsmessage.TypeCNAME:
-			if r, ok := answer.Body.(*dnsmessage.CNAMEResource); ok {
-				fmt.Printf("    CNAME: %v\n", r.CNAME)
-			}
-		case dnsmessage.TypeMX:
-			if r, ok := answer.Body.(*dnsmessage.MXResource); ok {
-				fmt.Printf("    Preference: %v\n", r.Pref)
-				fmt.Printf("    MX: %v\n", r.MX)
-			}
-		case dnsmessage.TypeTXT:
-			if r, ok := answer.Body.(*dnsmessage.TXTResource); ok {
-				fmt.Printf("    TXT: %v\n", r.TXT)
-			}
-		default:
-			fmt.Printf("    [Unsupported record type]\n")
-		}
 	}
 
 	return nil

internal/protocols/dot/dot.go

@@ -9,35 +9,35 @@ import (
 	"os"
 
 	"github.com/afonsofrancof/sdns-perf/internal/protocols/do53"
-	"golang.org/x/net/dns/dnsmessage"
+	"github.com/miekg/dns"
 )
 
-func Run(domain, queryType, server string, dnssec bool) error {
+type DoTClient struct {
+	tcpConn    *net.TCPConn
+	tlsConn    *tls.Conn
+	keyLogFile *os.File
+}
 
-	DNSMessage, err := do53.MakeDNSMessage(domain, queryType)
+func New(target string) (*DoTClient, error) {
+
+	tcpAddr, err := net.ResolveTCPAddr("tcp", target)
 	if err != nil {
-		return err
+		return nil, fmt.Errorf("failed to resolve TCP address: %v", err)
 	}
 
-	// Step 1 - Establish a TCP Connection
+	tcpConn, err := net.DialTCP("tcp", nil, tcpAddr)
-	tcpConn, err := net.Dial("tcp", server)
 	if err != nil {
-		return fmt.Errorf("failed to establish TCP connection: %v", err)
+		return nil, fmt.Errorf("failed to establish TCP connection: %v", err)
 	}
-	defer tcpConn.Close()
 
-	// Step 2 - Upgrade it to a TLS Connection
-
-	// Temporary keylog file to allow traffic inspection
 	keyLogFile, err := os.OpenFile(
 		"tls-key-log.txt",
 		os.O_APPEND|os.O_CREATE|os.O_WRONLY,
 		0600,
 	)
 	if err != nil {
-		return fmt.Errorf("failed opening key log file: %v", err)
+		return nil, fmt.Errorf("failed opening key log file: %v", err)
 	}
-	defer keyLogFile.Close()
 
 	tlsConfig := &tls.Config{
 		InsecureSkipVerify: true,
@@ -48,22 +48,43 @@ func Run(domain, queryType, server string, dnssec bool) error {
 	tlsConn := tls.Client(tcpConn, tlsConfig)
 	err = tlsConn.Handshake()
 	if err != nil {
-		return fmt.Errorf("failed to execute the TLS handshake: %v", err)
+		return nil, fmt.Errorf("failed to execute the TLS handshake: %v", err)
+	}
+
+	return &DoTClient{tcpConn: tcpConn, tlsConn: tlsConn, keyLogFile: keyLogFile}, nil
+}
+
+func (c *DoTClient) Close() {
+	if c.tcpConn != nil {
+		c.tcpConn.Close()
+	}
+	if c.tlsConn != nil {
+		c.tlsConn.Close()
+	}
+	if c.keyLogFile != nil {
+		c.keyLogFile.Close()
+	}
+}
+
+func (c *DoTClient) Query(domain, queryType, target string, dnssec bool) error {
+
+	DNSMessage, err := do53.NewDNSMessage(domain, queryType)
+	if err != nil {
+		return err
 	}
-	defer tlsConn.Close()
 
 	var lengthPrefixedMessage bytes.Buffer
 	binary.Write(&lengthPrefixedMessage, binary.BigEndian, uint16(len(DNSMessage)))
 	lengthPrefixedMessage.Write(DNSMessage)
 
-	_, err = tlsConn.Write(lengthPrefixedMessage.Bytes())
+	_, err = c.tlsConn.Write(lengthPrefixedMessage.Bytes())
 	if err != nil {
 		return fmt.Errorf("failed writing TLS request: %v", err)
 	}
 
 	// Read the 2-byte length prefix
 	lengthBuf := make([]byte, 2)
-	_, err = tlsConn.Read(lengthBuf)
+	_, err = c.tlsConn.Read(lengthBuf)
 	if err != nil {
 		return fmt.Errorf("failed reading response length: %v", err)
 	}
@@ -71,73 +92,21 @@ func Run(domain, queryType, server string, dnssec bool) error {
 	messageLength := binary.BigEndian.Uint16(lengthBuf)
 
 	responseBuf := make([]byte, messageLength)
-	n, err := tlsConn.Read(responseBuf)
+	n, err := c.tlsConn.Read(responseBuf)
 	if err != nil {
 		return fmt.Errorf("failed reading TLS response: %v", err)
 	}
 
-	// Parse the response
+	recvMsg := new(dns.Msg)
-	var parser dnsmessage.Parser
+	err = recvMsg.Unpack(responseBuf[:n])
-	header, err := parser.Start(responseBuf[:n])
 	if err != nil {
 		return fmt.Errorf("failed to parse DNS response: %v", err)
 	}
 
-	fmt.Printf("DNS Response Header:\n")
+	// TODO: Check if the response had no errors or TD bit set
-	fmt.Printf("  ID: %d\n", header.ID)
-	fmt.Printf("  Response: %v\n", header.Response)
-	fmt.Printf("  RCode: %v\n", header.RCode)
 
-	// Skip all questions before reading answers
+	for _, answer := range recvMsg.Answer {
-	err = parser.SkipAllQuestions()
+		fmt.Println(answer.String())
-	if err != nil {
-		return fmt.Errorf("failed to skip questions: %v", err)
-	}
-
-	// Parse answers
-	fmt.Printf("\nAnswers:\n")
-	answers, err := parser.AllAnswers()
-
-	for i, answer := range answers {
-
-		if err != nil {
-			return fmt.Errorf("failed to parse answer %d: %v", i, err)
-		}
-
-		fmt.Printf("  Answer %d:\n", i+1)
-		fmt.Printf("    Name: %v\n", answer.Header.Name)
-		fmt.Printf("    Type: %v\n", answer.Header.Type)
-		fmt.Printf("    TTL: %v seconds\n", answer.Header.TTL)
-
-		// Handle different record types
-		switch answer.Header.Type {
-		case dnsmessage.TypeA:
-			if r, ok := answer.Body.(*dnsmessage.AResource); ok {
-				fmt.Printf("    IPv4: %d.%d.%d.%d\n", r.A[0], r.A[1], r.A[2], r.A[3])
-			}
-		case dnsmessage.TypeAAAA:
-			if r, ok := answer.Body.(*dnsmessage.AAAAResource); ok {
-				ip := r.AAAA
-				fmt.Printf("    IPv6: %02x%02x:%02x%02x:%02x%02x:%02x%02x:%02x%02x:%02x%02x:%02x%02x:%02x%02x\n",
-					ip[0], ip[1], ip[2], ip[3], ip[4], ip[5], ip[6], ip[7],
-					ip[8], ip[9], ip[10], ip[11], ip[12], ip[13], ip[14], ip[15])
-			}
-		case dnsmessage.TypeCNAME:
-			if r, ok := answer.Body.(*dnsmessage.CNAMEResource); ok {
-				fmt.Printf("    CNAME: %v\n", r.CNAME)
-			}
-		case dnsmessage.TypeMX:
-			if r, ok := answer.Body.(*dnsmessage.MXResource); ok {
-				fmt.Printf("    Preference: %v\n", r.Pref)
-				fmt.Printf("    MX: %v\n", r.MX)
-			}
-		case dnsmessage.TypeTXT:
-			if r, ok := answer.Body.(*dnsmessage.TXTResource); ok {
-				fmt.Printf("    TXT: %v\n", r.TXT)
-			}
-		default:
-			fmt.Printf("    [Unsupported record type]\n")
-		}
 	}
 
 	return nil