From 0aa3b36b35f5924bc98cac13f192b576960fa501 Mon Sep 17 00:00:00 2001
From: Karan Sharma <hello@mrkaran.dev>
Date: Sat, 12 Dec 2020 13:16:45 +0530
Subject: [PATCH] feat: add support for tcp, dot

---
 TODO.md                  |  6 +++++-
 cmd/cli.go               | 18 ++++++++++++++----
 cmd/hub.go               |  2 +-
 cmd/resolver.go          | 10 ++++------
 pkg/resolvers/classic.go | 33 ++++++++++++++++++++++++---------
 pkg/resolvers/dot.go     |  1 -
 6 files changed, 48 insertions(+), 22 deletions(-)
 delete mode 100644 pkg/resolvers/dot.go

diff --git a/TODO.md b/TODO.md
index b1be4b7..bd46882 100644
--- a/TODO.md
+++ b/TODO.md
@@ -2,10 +2,14 @@
 
 ## Resolver
 - [x] Create a DNS Resolver struct
-- [ ] Add methods to initialise the config, set defaults
+- [x]] Add methods to initialise the config, set defaults
 - [x] Add a resolve method
 - [x] Make it separate from Hub
 - [ ] Parse output into separate fields
+- [ ] Test UDP6
+- [x] Add DOH support
+- [x] Add DOT support
+- [x] Add DNS protocol on TCP mode support.
 
 ## CLI Features
 - [ ] `digfile`
diff --git a/cmd/cli.go b/cmd/cli.go
index 0945a9c..760295a 100644
--- a/cmd/cli.go
+++ b/cmd/cli.go
@@ -67,18 +67,28 @@ func main() {
 			Destination: hub.QueryFlags.QClasses,
 		},
 		&cli.BoolFlag{
-			Name:  "udp",
-			Usage: "Use the DNS protocol over UDP",
+			Name:    "udp",
+			Usage:   "Use the DNS protocol over UDP",
+			Aliases: []string{"U"},
 		},
 		&cli.BoolFlag{
-			Name:  "tcp",
-			Usage: "Use the DNS protocol over TCP",
+			Name:        "tcp",
+			Usage:       "Use the DNS protocol over TCP",
+			Aliases:     []string{"T"},
+			Destination: &hub.QueryFlags.UseTCP,
 		},
 		&cli.BoolFlag{
 			Name:        "https",
 			Usage:       "Use the DNS-over-HTTPS protocol",
+			Aliases:     []string{"H"},
 			Destination: &hub.QueryFlags.IsDOH,
 		},
+		&cli.BoolFlag{
+			Name:        "tls",
+			Usage:       "Use the DNS-over-TLS",
+			Aliases:     []string{"S"},
+			Destination: &hub.QueryFlags.IsDOT,
+		},
 		&cli.BoolFlag{
 			Name:        "ipv6",
 			Aliases:     []string{"6"},
diff --git a/cmd/hub.go b/cmd/hub.go
index 418e641..bd648f6 100644
--- a/cmd/hub.go
+++ b/cmd/hub.go
@@ -25,7 +25,7 @@ type QueryFlags struct {
 	IsDOH       bool
 	IsDOT       bool
 	IsUDP       bool
-	IsTLS       bool
+	UseTCP      bool
 	UseIPv4     bool
 	UseIPv6     bool
 }
diff --git a/cmd/resolver.go b/cmd/resolver.go
index 682386e..c300cd9 100644
--- a/cmd/resolver.go
+++ b/cmd/resolver.go
@@ -7,7 +7,8 @@ import (
 	"github.com/urfave/cli/v2"
 )
 
-// loadResolver checks
+// loadResolver checks for various flags and initialises
+// the correct resolver based on the config.
 func (hub *Hub) loadResolver(c *cli.Context) error {
 	// check if DOH flag is set.
 	if hub.QueryFlags.IsDOH {
@@ -18,11 +19,6 @@ func (hub *Hub) loadResolver(c *cli.Context) error {
 		hub.Resolver = rslvr
 		return nil
 	}
-	// check if DOT flag is set.
-
-	// check if TCP flag is set.
-
-	// fallback to good ol UDP.
 	if len(hub.QueryFlags.Nameservers.Value()) == 0 {
 		if runtime.GOOS == "windows" {
 			// TODO: Add a method for reading system default nameserver in windows.
@@ -38,6 +34,8 @@ func (hub *Hub) loadResolver(c *cli.Context) error {
 		rslvr, err := resolvers.NewClassicResolver(hub.QueryFlags.Nameservers.Value(), resolvers.ClassicResolverOpts{
 			UseIPv4: hub.QueryFlags.UseIPv4,
 			UseIPv6: hub.QueryFlags.UseIPv6,
+			UseTLS:  hub.QueryFlags.IsDOT,
+			UseTCP:  hub.QueryFlags.UseTCP,
 		})
 		if err != nil {
 			return err
diff --git a/pkg/resolvers/classic.go b/pkg/resolvers/classic.go
index 5ee1761..e6d6f54 100644
--- a/pkg/resolvers/classic.go
+++ b/pkg/resolvers/classic.go
@@ -7,6 +7,15 @@ import (
 	"github.com/miekg/dns"
 )
 
+const (
+	// DefaultUDPPort specifies the default port for a DNS server connecting over UDP
+	DefaultUDPPort = "53"
+	// DefaultTLSPort specifies the default port for a DNS server connecting over TCP over TLS
+	DefaultTLSPort = "853"
+	//DefaultResolvConfPath specifies path to default resolv config file on UNIX.
+	DefaultResolvConfPath = "/etc/resolv.conf"
+)
+
 // ClassicResolver represents the config options for setting up a Resolver.
 type ClassicResolver struct {
 	client  *dns.Client
@@ -18,24 +27,24 @@ type ClassicResolverOpts struct {
 	UseIPv4 bool
 	UseIPv6 bool
 	UseTCP  bool
+	UseTLS  bool
 }
 
-//DefaultResolvConfPath specifies path to default resolv config file on UNIX.
-const DefaultResolvConfPath = "/etc/resolv.conf"
-
 // NewClassicResolver accepts a list of nameservers and configures a DNS resolver.
 func NewClassicResolver(servers []string, opts ClassicResolverOpts) (Resolver, error) {
 	client := &dns.Client{}
 	var nameservers []string
 	for _, srv := range servers {
 		if i := net.ParseIP(srv); i != nil {
-			nameservers = append(nameservers, net.JoinHostPort(srv, "53"))
-		} else {
-			host, port, err := net.SplitHostPort(srv)
-			if err != nil {
-				return nil, err
+			// if no port specified in nameserver, append defaults.
+			if opts.UseTLS == true {
+				nameservers = append(nameservers, net.JoinHostPort(srv, DefaultTLSPort))
+			} else {
+				nameservers = append(nameservers, net.JoinHostPort(srv, DefaultUDPPort))
 			}
-			nameservers = append(nameservers, fmt.Sprintf("%s:%s", host, port))
+		} else {
+			// use the port user specified.
+			nameservers = append(nameservers, srv)
 		}
 	}
 	client.Net = "udp"
@@ -45,6 +54,12 @@ func NewClassicResolver(servers []string, opts ClassicResolverOpts) (Resolver, e
 	if opts.UseIPv6 {
 		client.Net = "udp6"
 	}
+	if opts.UseTCP {
+		client.Net = "tcp"
+	}
+	if opts.UseTLS {
+		client.Net = "tcp-tls"
+	}
 	return &ClassicResolver{
 		client:  client,
 		servers: nameservers,
diff --git a/pkg/resolvers/dot.go b/pkg/resolvers/dot.go
deleted file mode 100644
index 91fea07..0000000
--- a/pkg/resolvers/dot.go
+++ /dev/null
@@ -1 +0,0 @@
-package resolvers