feat: custom response format

cmd/lookup.go

@@ -1,18 +1,28 @@
 package main
 
 import (
+	"fmt"
 	"strings"
 
 	"github.com/miekg/dns"
 	"github.com/urfave/cli/v2"
 )
 
+// Lookup sends the DNS queries to the server.
 func (hub *Hub) Lookup(c *cli.Context) error {
 	hub.prepareQuestions()
-	err := hub.Resolver.Lookup(hub.Questions)
+	responses, err := hub.Resolver.Lookup(hub.Questions)
 	if err != nil {
 		hub.Logger.Error(err)
 	}
+	for _, r := range responses {
+		for _, a := range r.Message.Answer {
+			if t, ok := a.(*dns.A); ok {
+				fmt.Println(t.String())
+			}
+		}
+	}
+
 	return nil
 }
 

cmd/output.go

@@ -0,0 +1,12 @@
+package main
+
+// Output takes a list of `dns.Answers` and based
+// on the output format specified displays the information.
+// func (hub *Hub) Output(c *cli.Context) error {
+// 	hub.prepareQuestions()
+// 	answers, err := hub.Resolver.Lookup(hub.Questions)
+// 	if err != nil {
+// 		hub.Logger.Error(err)
+// 	}
+// 	return nil
+// }

cmd/parse.go

@@ -12,7 +12,7 @@ func (hub *Hub) loadQueryArgs(c *cli.Context) error {
 	if err != nil {
 		cli.Exit("Error parsing arguments", -1)
 	}
-	err = hub.loadResolver(c)
+	err = hub.initResolver(c)
 	if err != nil {
 		cli.Exit("Error parsing nameservers", -1)
 	}

cmd/resolver.go

@@ -7,9 +7,9 @@ import (
 	"github.com/urfave/cli/v2"
 )
 
-// loadResolver checks for various flags and initialises
+// initResolver checks for various flags and initialises
 // the correct resolver based on the config.
-func (hub *Hub) loadResolver(c *cli.Context) error {
+func (hub *Hub) initResolver(c *cli.Context) error {
 	// check if DOH flag is set.
 	if hub.QueryFlags.IsDOH {
 		rslvr, err := resolvers.NewDOHResolver(hub.QueryFlags.Nameservers.Value())

pkg/resolvers/classic.go

@@ -98,22 +98,26 @@ func NewResolverFromResolvFile(resolvFilePath string) (Resolver, error) {
 // Lookup prepare a list of DNS messages to be sent to the server.
 // It's possible to send multiple question in one message
 // but some nameservers are not able to
-func (c *ClassicResolver) Lookup(questions []dns.Question) error {
-	messages := prepareMessages(questions)
+func (c *ClassicResolver) Lookup(questions []dns.Question) ([]Response, error) {
+	var (
+		messages  = prepareMessages(questions)
+		responses []Response
+	)
 
 	for _, msg := range messages {
 		for _, srv := range c.servers {
 			in, rtt, err := c.client.Exchange(&msg, srv)
 			if err != nil {
-				return err
+				return nil, err
 			}
-			for _, ans := range in.Answer {
-				if t, ok := ans.(*dns.A); ok {
-					fmt.Println(t.String())
+			msg.Answer = in.Answer
+			rsp := Response{
+				Message:    msg,
+				RTT:        rtt,
+				Nameserver: srv,
+			}
+			responses = append(responses, rsp)
 		}
 	}
-			fmt.Println("rtt is", rtt)
-		}
-	}
-	return nil
+	return responses, nil
 }

pkg/resolvers/doh.go

@@ -2,7 +2,6 @@ package resolvers
 
 import (
 	"bytes"
-	"fmt"
 	"io/ioutil"
 	"net/http"
 	"time"
@@ -27,39 +26,46 @@ func NewDOHResolver(servers []string) (Resolver, error) {
 	}, nil
 }
 
-func (r *DOHResolver) Lookup(questions []dns.Question) error {
-	messages := prepareMessages(questions)
+func (r *DOHResolver) Lookup(questions []dns.Question) ([]Response, error) {
+	var (
+		messages  = prepareMessages(questions)
+		responses []Response
+	)
 
-	for _, m := range messages {
-		b, err := m.Pack()
+	for _, msg := range messages {
+		// get the DNS Message in wire format.
+		b, err := msg.Pack()
 		if err != nil {
-			return err
+			return nil, err
 		}
 		for _, srv := range r.servers {
+			now := time.Now()
+			// Make an HTTP POST request to the DNS server with the DNS message as wire format bytes in the body.
 			resp, err := r.client.Post(srv, "application/dns-message", bytes.NewBuffer(b))
 			if err != nil {
-				return err
+				return nil, err
 			}
 			if resp.StatusCode != http.StatusOK {
-				return err
+				return nil, err
 			}
-
+			rtt := time.Since(now)
+			// extract the binary response in DNS Message.
 			body, err := ioutil.ReadAll(resp.Body)
 			if err != nil {
-				return err
+				return nil, err
 			}
 
-			r := &dns.Msg{}
-			err = r.Unpack(body)
+			err = msg.Unpack(body)
 			if err != nil {
-				return err
+				return nil, err
 			}
-			for _, ans := range r.Answer {
-				if t, ok := ans.(*dns.A); ok {
-					fmt.Println(t.String())
+			rsp := Response{
+				Message:    msg,
+				RTT:        rtt,
+				Nameserver: srv,
+			}
+			responses = append(responses, rsp)
 		}
 	}
-		}
-	}
-	return nil
+	return responses, nil
 }

pkg/resolvers/resolver.go

@@ -1,10 +1,23 @@
 package resolvers
 
-import "github.com/miekg/dns"
+import (
+	"time"
+
+	"github.com/miekg/dns"
+)
 
 // Resolver implements the configuration for a DNS
 // Client. Different types of client like (UDP/TCP/DOH/DOT)
 // can be initialised.
 type Resolver interface {
-	Lookup([]dns.Question) error
+	Lookup([]dns.Question) ([]Response, error)
+}
+
+// Response represents a custom output format
+// which wraps certain metadata about the DNS query
+// and the DNS Answer as well.
+type Response struct {
+	Message    dns.Msg
+	RTT        time.Duration
+	Nameserver string
 }