refactor DNS configuration and add NS query handling

app/dns.yml

@@ -1,8 +1,9 @@
-port: 53
+dnsServer:
-network: udp
-alternate_resolver: 
-  ip: 8.8.8.8
   port: 53
+  network: udp # protocol
+  alternate_resolver: 
+    ip: 8.8.8.8
+    port: 53
 records:
   type_a:
     test.com: 10.10.10.1

app/dns/dns.go

@@ -11,14 +11,13 @@ import (
 
 	miekgDNS "github.com/miekg/dns"
 	"github.com/spf13/viper"
-	"github.com/zeevdiukman/test/app/helper"
 )
 
 type ResponseWriter = miekgDNS.ResponseWriter
 type Msg = miekgDNS.Msg
 type HandlerFunc = miekgDNS.HandlerFunc
 
-// func(miekgDNS.ResponseWriter, *miekgDNS.Msg)
+// DNS struct represents the DNS server with its configuration, server, mux, records, and resolver.
 type DNS struct {
 	Config   Config
 	Server   Server
@@ -27,159 +26,110 @@ type DNS struct {
 	Resolver Resolver
 }
 
+// Config struct holds the configuration settings.
 type Config struct {
 	*viper.Viper
 }
 
+// Server struct wraps the miekgDNS.Server.
 type Server struct {
 	*miekgDNS.Server
 }
 
+// Mux struct wraps the miekgDNS.ServeMux.
 type Mux struct {
 	*miekgDNS.ServeMux
 }
 
+// Records struct holds the DNS records.
 type Records struct {
-	TypeA map[string]string
+	TypeA  map[string]string
+	TypeNS map[string]string // Add NS records
 }
 
+// Resolver struct wraps the net.Resolver.
 type Resolver struct {
 	*net.Resolver
 }
 
+// New initializes a new DNS instance with the given configuration file path.
 func New(filePath string) *DNS {
 	app := &DNS{}
-	app.ConfigInit(filePath)
+	app.ConfigInit(filePath)                                                              // Initialize configuration
-	app.Records.TypeA = app.Config.GetStringMapString("records.type_a")
+	app.Records.TypeA = app.Config.GetStringMapString("records.type_a")                   // Load Type A records
-	alternate_resolver_ip := app.Config.GetString("alternate_resolver.ip")
+	alternate_resolver_ip := app.Config.GetString("alternate_resolver.ip")                // Get alternate resolver IP
-	alternate_resolver_port := strconv.Itoa(app.Config.GetInt("alternate_resolver.port"))
+	alternate_resolver_port := strconv.Itoa(app.Config.GetInt("alternate_resolver.port")) // Get alternate resolver port
-	app.ServerInit()
+	app.ServerInit()                                                                      // Initialize server
-	app.MuxInit()
+	app.MuxInit()                                                                         // Initialize multiplexer
-	app.Resolver = NewResolver(alternate_resolver_ip + ":" + alternate_resolver_port)
+	app.Resolver = NewResolver(alternate_resolver_ip + ":" + alternate_resolver_port)     // Initialize resolver
-	app.Server.Handler = app.Mux
+	app.Server.Handler = app.Mux                                                          // Set handler
 
 	return app
 }
+
+// Run starts the DNS server.
 func (a *DNS) Run() {
-	err := a.Server.ListenAndServe()
+	err := a.Server.ListenAndServe() // Start server
 	if err != nil {
-		fmt.Println(err.Error())
+		fmt.Println(err.Error()) // Print error if any
 	}
 }
+
+// ConfigInit initializes the configuration from the given file path.
 func (a *DNS) ConfigInit(filePath string) {
+	dir, file := filepath.Split(filePath)  // Split file path
+	fe := filepath.Ext(file)               // Get file extension
+	ext, _ := strings.CutPrefix(fe, ".")   // Remove leading dot from extension
+	name, _ := strings.CutSuffix(file, fe) // Get file name without extension
 
-	dir, file := filepath.Split(filePath)
+	a.Config.Viper = viper.New() // Create new Viper instance
-	fe := filepath.Ext(file)
+	a.Config.AddConfigPath(dir)  // Add config path
-	ext, _ := strings.CutPrefix(fe, ".")
+	a.Config.SetConfigName(name) // Set config name
-	name, _ := strings.CutSuffix(file, fe)
+	a.Config.SetConfigType(ext)  // Set config type
-
+	a.Config.ReadInConfig()      // Read config
-	a.Config.Viper = viper.New()
-	a.Config.AddConfigPath(dir)
-	a.Config.SetConfigName(name)
-	a.Config.SetConfigType(ext)
-	a.Config.ReadInConfig()
-
 }
+
+// ServerInit initializes the DNS server with the configuration settings.
 func (a *DNS) ServerInit() {
-	port := a.Config.GetInt("port")
+	port := a.Config.GetInt("port") // Get port from config
 	a.Server.Server = &miekgDNS.Server{
-		Addr:    ":" + strconv.Itoa(port),
+		Addr:    ":" + strconv.Itoa(port),      // Set server address
-		Net:     a.Config.GetString("network"),
+		Net:     a.Config.GetString("network"), // Set network type
-		Handler: nil,
+		Handler: nil,                           // Handler will be set later
 	}
-	helper.P("DNS server started at port ", port)
-
 }
 
+// MuxInit initializes the DNS request multiplexer.
 func (a *DNS) MuxInit() {
-	a.Mux.ServeMux = miekgDNS.NewServeMux()
+	a.Mux.ServeMux = miekgDNS.NewServeMux() // Create new ServeMux
-	a.Mux.HandleFunc(".", a.HandleTypeA)
+	a.Mux.HandleFunc(".", a.HandleTypeA)    // Handle all requests with HandleTypeA
-
+	a.Mux.HandleFunc(".", a.HandleTypeNS)   // Handle all requests with HandleTypeNS
 }
 
+// NewResolver creates a new DNS resolver with the given DNS server address.
 func NewResolver(DNSserverAddr string) Resolver {
 	return Resolver{
 		&net.Resolver{
-			PreferGo: true,
+			PreferGo: true, // Prefer Go's resolver
 			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
-				d := net.Dialer{Timeout: 5 * time.Second}
+				d := net.Dialer{Timeout: 5 * time.Second}       // Set dial timeout
-				return d.DialContext(ctx, "udp", DNSserverAddr)
+				return d.DialContext(ctx, "udp", DNSserverAddr) // Dial DNS server
 			},
 		},
 	}
 }
 
-func (a *DNS) Handler(f func(a *DNS, w miekgDNS.ResponseWriter, r *miekgDNS.Msg)) miekgDNS.HandlerFunc {
+// Lookup performs a DNS lookup for the given address using the resolver.
-	return func(w miekgDNS.ResponseWriter, r *miekgDNS.Msg) {
-		f(a, w, r)
-	}
-}
-
 func (r *Resolver) Lookup(lookupAddr string) string {
 	var resp []string
 	var err error
 
-	ctx := context.Background()
+	ctx := context.Background()               // Create background context
-	resp, err = r.LookupHost(ctx, lookupAddr)
+	resp, err = r.LookupHost(ctx, lookupAddr) // Perform lookup
 
 	if err != nil {
-		fmt.Println(err.Error())
+		fmt.Println(err.Error()) // Print error if any
 	}
 
-	return resp[0]
+	return resp[0] // Return first result
-}
-
-func (a *DNS) HandleTypeA(w miekgDNS.ResponseWriter, r *miekgDNS.Msg) {
-	useAlternateResolver := false
-	t := time.Now()
-	msg := &miekgDNS.Msg{}
-	msg.SetReply(r)
-	q := r.Question[0]
-	domainName := helper.FtoD(q.Name)
-	ip := ""
-
-	if ipValue, ok := a.Records.TypeA[domainName]; ok {
-		helper.P("FOUND => ", domainName)
-		ip = ipValue
-	} else {
-		dSlices := strings.Split(domainName, ".")
-		//check if wild card
-		if len(dSlices) > 2 {
-
-			name := dSlices[len(dSlices)-2]
-			tld := dSlices[len(dSlices)-1]
-			cname := name + "." + tld
-			wildCard := "*." + cname
-			if ipValue, ok := a.Records.TypeA[wildCard]; ok {
-
-				ip = ipValue
-			} else {
-				useAlternateResolver = true
-			}
-		} else {
-			useAlternateResolver = true
-		}
-
-		if useAlternateResolver {
-			ip = a.Resolver.Lookup(domainName)
-		}
-	}
-	RR_Header := miekgDNS.RR_Header{
-		Name:   miekgDNS.Fqdn(domainName),
-		Rrtype: miekgDNS.TypeA,
-		Class:  miekgDNS.ClassINET,
-		Ttl:    3600,
-	}
-	answer_typeA := &miekgDNS.A{
-		Hdr: RR_Header,
-		A:   net.ParseIP(ip).To4(),
-	}
-	msg.Authoritative = true
-	msg.RecursionDesired = false
-	msg.SetRcode(r, miekgDNS.RcodeSuccess)
-	msg.Answer = append(msg.Answer, answer_typeA)
-	tt := time.Since(t)
-	helper.P(tt, " => ", domainName)
-
-	w.WriteMsg(msg)
 }

app/dns/handlers.go

@@ -0,0 +1,96 @@
+package dns
+
+import (
+	"net"
+	"strings"
+	"time"
+
+	"github.com/zeevdiukman/test/app/helper"
+
+	miekgDNS "github.com/miekg/dns"
+)
+
+// HandleTypeNS handles DNS NS queries.
+func (a *DNS) HandleTypeNS(w miekgDNS.ResponseWriter, r *miekgDNS.Msg) {
+	t := time.Now() // Start timer
+	msg := &miekgDNS.Msg{}
+	msg.SetReply(r)                   // Set reply message
+	q := r.Question[0]                // Get question
+	domainName := helper.FtoD(q.Name) // Convert FQDN to domain name
+	ns := ""
+	if nsValue, ok := a.Records.TypeNS[domainName]; ok {
+		ns = nsValue // Get NS from records
+	} else {
+		ns = "default-ns.example.com." // Default NS if not found
+	}
+	RR_Header := miekgDNS.RR_Header{
+		Name:   miekgDNS.Fqdn(domainName), // Set domain name
+		Rrtype: miekgDNS.TypeNS,           // Set record type
+		Class:  miekgDNS.ClassINET,        // Set class
+		Ttl:    3600,                      // Set TTL
+	}
+	answer_typeNS := &miekgDNS.NS{
+		Hdr: RR_Header, // Set header
+		Ns:  ns,        // Set NS
+	}
+	msg.Authoritative = true                       // Set authoritative flag
+	msg.RecursionDesired = false                   // Set recursion desired flag
+	msg.SetRcode(r, miekgDNS.RcodeSuccess)         // Set response code
+	msg.Answer = append(msg.Answer, answer_typeNS) // Add answer
+	tt := time.Since(t)                            // Calculate elapsed time
+	helper.P(tt, " => ", domainName)               // Print elapsed time and domain name
+
+	w.WriteMsg(msg) // Write response
+}
+
+// HandleTypeA handles DNS Type A queries.
+func (a *DNS) HandleTypeA(w miekgDNS.ResponseWriter, r *miekgDNS.Msg) {
+	useAlternateResolver := false
+	t := time.Now() // Start timer
+	msg := &miekgDNS.Msg{}
+	msg.SetReply(r)                   // Set reply message
+	q := r.Question[0]                // Get question
+	domainName := helper.FtoD(q.Name) // Convert FQDN to domain name
+	ip := ""
+	if ipValue, ok := a.Records.TypeA[domainName]; ok {
+		ip = ipValue // Get IP from records
+	} else {
+		dSlices := strings.Split(domainName, ".") // Split domain name
+		// Check if wildcard
+		if len(dSlices) > 2 {
+			name := dSlices[len(dSlices)-2]
+			tld := dSlices[len(dSlices)-1]
+			cname := name + "." + tld
+			wildCard := "*." + cname
+			if ipValue, ok := a.Records.TypeA[wildCard]; ok {
+				ip = ipValue // Get IP from wildcard record
+			} else {
+				useAlternateResolver = true // Use alternate resolver
+			}
+		} else {
+			useAlternateResolver = true // Use alternate resolver
+		}
+
+		if useAlternateResolver {
+			ip = a.Resolver.Lookup(domainName) // Lookup IP using resolver
+		}
+	}
+	RR_Header := miekgDNS.RR_Header{
+		Name:   miekgDNS.Fqdn(domainName), // Set domain name
+		Rrtype: miekgDNS.TypeA,            // Set record type
+		Class:  miekgDNS.ClassINET,        // Set class
+		Ttl:    3600,                      // Set TTL
+	}
+	answer_typeA := &miekgDNS.A{
+		Hdr: RR_Header,             // Set header
+		A:   net.ParseIP(ip).To4(), // Set IP address
+	}
+	msg.Authoritative = true                      // Set authoritative flag
+	msg.RecursionDesired = false                  // Set recursion desired flag
+	msg.SetRcode(r, miekgDNS.RcodeSuccess)        // Set response code
+	msg.Answer = append(msg.Answer, answer_typeA) // Add answer
+	tt := time.Since(t)                           // Calculate elapsed time
+	helper.P(tt, " => ", domainName)              // Print elapsed time and domain name
+
+	w.WriteMsg(msg) // Write response
+}

app/main.go

@@ -11,16 +11,19 @@ import (
 
 func main() {
 	// Clear the terminal screen
-
 	helper.ClearScreen()
+
+	// Initialize the DNS server with the configuration file
 	d := dns.New("./dns.yml")
+
 	// Run the HTTP and DNS servers concurrently
 	go runHTTP()
 
 	go func() {
+		helper.P("DNS server started at " + d.Server.Addr)
 		d.Run()
 	}()
-	// go runDNS()
+
 	// Wait for SIGINT (Ctrl+C) to gracefully shut down the server
 	stop := make(chan os.Signal, 1)
 	signal.Notify(stop, os.Interrupt)

app/tmp/build-errors.log

@@ -1 +1 @@
-exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1
+exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1exit status 1