basic stdin reader on loop to reuse connections

cmd/resolver/main.go

@@ -1,7 +1,10 @@
 package main
 
 import (
+	"bufio"
 	"log"
+	"os"
+	"strings"
 
 	"github.com/afonsofrancof/sdns-perf/internal/protocols/do53"
 	"github.com/afonsofrancof/sdns-perf/internal/protocols/doh"
@@ -36,6 +39,10 @@ type Do53Cmd struct {
 	CommonFlags
 }
 
+type Listen struct {
+
+}
+
 var cli struct {
 	Verbose bool `help:"Enable verbose logging" short:"v"`
 
@@ -43,6 +50,7 @@ var cli struct {
 	DoT  DoTCmd  `cmd:"dot" help:"Query using DNS-over-TLS" name:"dot"`
 	DoQ  DoQCmd  `cmd:"doq" help:"Query using DNS-over-QUIC" name:"doq"`
 	Do53 Do53Cmd `cmd:"doq" help:"Query using plain DNS over UDP" name:"do53"`
+	Listen Listen `cmd:"listen"`
 }
 
 func (c *Do53Cmd) Run() error {
@@ -50,6 +58,7 @@ func (c *Do53Cmd) Run() error {
 	if err != nil {
 		return err
 	}
+	defer do53client.Close()
 	return do53client.Query(c.DomainName, c.QueryType, c.Server, c.DNSSEC)
 }
 
@@ -58,6 +67,7 @@ func (c *DoHCmd) Run() error {
 	if err != nil {
 		return err
 	}
+	defer dohclient.Close()
 	return dohclient.Query(c.DomainName, c.QueryType, c.DNSSEC)
 }
 
@@ -66,6 +76,7 @@ func (c *DoTCmd) Run() error {
 	if err != nil {
 		return err
 	}
+	defer dotclient.Close()
 	return dotclient.Query(c.DomainName, c.QueryType, c.Server, c.DNSSEC)
 }
 
@@ -74,9 +85,172 @@ func (c *DoQCmd) Run() error {
 	if err != nil {
 		return err
 	}
+	defer doqclient.Close()
 	return doqclient.Query(c.DomainName, c.QueryType, c.DNSSEC)
 }
 
+func (l *Listen) Run() error {
+    // Maps to store clients for reuse
+    do53Clients := make(map[string]*do53.Do53Client)
+    dotClients := make(map[string]*dot.DoTClient)
+    doqClients := make(map[string]*doq.DoQClient)
+    dohClients := make(map[string]*doh.DoHClient) // Using server+path+proxy as key
+    
+    scanner := bufio.NewScanner(os.Stdin)
+    log.Println("Listening for input. Format: protocol domain server [options]")
+    
+    for scanner.Scan() {
+        line := scanner.Text()
+        fields := strings.Fields(line)
+        
+        if len(fields) < 3 {
+            log.Printf("Invalid input: %s. Format should be 'protocol domain server [options]'", line)
+            continue
+        }
+        
+        protocol := fields[0]
+        domain := fields[1]
+        server := fields[2]
+        
+        // Default query type and DNSSEC setting
+        queryType := "A"
+        dnssec := false
+        
+        switch protocol {
+        case "do53":
+            // Parse additional options
+            if len(fields) > 3 {
+                queryType = fields[3]
+            }
+            if len(fields) > 4 && fields[4] == "dnssec" {
+                dnssec = true
+            }
+            
+            // Check if client exists, if not create it
+            client, exists := do53Clients[server]
+            if !exists {
+                var err error
+                client, err = do53.New(server)
+                if err != nil {
+                    log.Printf("Error creating Do53 client: %v", err)
+                    continue
+                }
+                do53Clients[server] = client
+            }
+            
+            err := client.Query(domain, queryType, server, dnssec)
+            if err != nil {
+                log.Printf("Error querying with Do53: %v", err)
+            }
+            
+        case "dot":
+            // Parse additional options
+            if len(fields) > 3 {
+                queryType = fields[3]
+            }
+            if len(fields) > 4 && fields[4] == "dnssec" {
+                dnssec = true
+            }
+            
+            client, exists := dotClients[server]
+            if !exists {
+                var err error
+                client, err = dot.New(server)
+                if err != nil {
+                    log.Printf("Error creating DoT client: %v", err)
+                    continue
+                }
+                dotClients[server] = client
+            }
+            
+            err := client.Query(domain, queryType, server, dnssec)
+            if err != nil {
+                log.Printf("Error querying with DoT: %v", err)
+            }
+            
+        case "doq":
+            // Parse additional options
+            if len(fields) > 3 {
+                queryType = fields[3]
+            }
+            if len(fields) > 4 && fields[4] == "dnssec" {
+                dnssec = true
+            }
+            
+            client, exists := doqClients[server]
+            if !exists {
+                var err error
+                client, err = doq.New(server)
+                if err != nil {
+                    log.Printf("Error creating DoQ client: %v", err)
+                    continue
+                }
+                doqClients[server] = client
+            }
+            
+            err := client.Query(domain, queryType, dnssec)
+            if err != nil {
+                log.Printf("Error querying with DoQ: %v", err)
+            }
+            
+        case "doh":
+            // DoH requires path parameter
+            if len(fields) < 4 {
+                log.Printf("DoH requires a path parameter")
+                continue
+            }
+            
+            path := fields[3]
+            proxy := ""
+            
+            // Parse additional options
+            if len(fields) > 4 {
+                queryType = fields[4]
+            }
+            
+            if len(fields) > 5 {
+                if fields[5] == "dnssec" {
+                    dnssec = true
+                } else {
+                    proxy = fields[5]
+                }
+            }
+            
+            if len(fields) > 6 && fields[6] == "dnssec" {
+                dnssec = true
+            }
+            
+            // Create a composite key for DoH clients
+            key := server + ":" + path + ":" + proxy
+            client, exists := dohClients[key]
+            if !exists {
+                var err error
+                client, err = doh.New(server, path, proxy)
+                if err != nil {
+                    log.Printf("Error creating DoH client: %v", err)
+                    continue
+                }
+                dohClients[key] = client
+            }
+            
+            err := client.Query(domain, queryType, dnssec)
+            if err != nil {
+                log.Printf("Error querying with DoH: %v", err)
+            }
+            
+        default:
+            log.Printf("Unknown protocol: %s", protocol)
+        }
+    }
+    
+    if err := scanner.Err(); err != nil {
+        return err
+    }
+    
+    return nil
+}
+
+
 func main() {
 	ctx := kong.Parse(&cli,
 		kong.Name("dns-go"),