Enhancement: support querying multiple servers at once

I made an attempt at implementing this, and it's relatively straightforward until we get to the output, since we don't store the origination server in the replies, and do not output the replying-server by default. Changing this is possible, but awkward due to #72. Here's the diff in case someone wants to pick up from where I left off; all but the final section should be relevant:

```diff diff --git a/cli/flags.go b/cli/flags.go index 5e396ae..536b5eb 100644 --- a/cli/flags.go +++ b/cli/flags.go @@ -4,7 +4,7 @@ import "time" type Flags struct { Name string `short:"q" long:"qname" description:"Query name"` - Server string `short:"s" long:"server" description:"DNS server"` + Server []string `short:"s" long:"server" description:"DNS server"` Types []string `short:"t" long:"type" description:"RR type (e.g. A, AAAA, MX, etc.) or type integer"` Reverse bool `short:"x" long:"reverse" description:"Reverse lookup"` DNSSEC bool `short:"d" long:"dnssec" description:"Set the DO (DNSSEC OK) bit in the OPT record"` diff --git a/main.go b/main.go index 3076d9e..0071d22 100644 --- a/main.go +++ b/main.go @@ -228,7 +228,7 @@ func parseServer(s string) (string, transport.Type, error) { // driver is the "main" function for this program that accepts a flag slice for testing func driver(args []string, out io.Writer) error { parser := flags.NewParser(&opts, flags.Default) - parser.Usage = `[OPTIONS] [@server] [type...] [name] + parser.Usage = `[OPTIONS] [@server...] [type...] [name] All long form (--) flags can be toggled with the dig-standard +[no]flag notation.` _, err := parser.ParseArgs(args) @@ -279,9 +279,9 @@ All long form (--) flags can be toggled with the dig-standard +[no]flag notation // Add non-flag RR types for _, arg := range args { - // Find a server by @ symbol if it isn't set by flag - if opts.Server == "" && strings.HasPrefix(arg, "@") { - opts.Server = strings.TrimPrefix(arg, "@") + // @ servers added to server list + if strings.HasPrefix(arg, "@") { + opts.Server = append(opts.Server, strings.TrimPrefix(arg, "@")) } // Parse chaos class @@ -335,23 +335,24 @@ All long form (--) flags can be toggled with the dig-standard +[no]flag notation log.Debugf("RR types: %+v", rrTypeStrings) } - // Set default DNS server - if opts.Server == "" { + // Set default DNS server if none were set explicitly + if len(opts.Server) == 0 { + opts.Server = make([]string, 1) if os.Getenv(defaultServerVar) != "" { - opts.Server = os.Getenv(defaultServerVar) + opts.Server[0] = os.Getenv(defaultServerVar) log.Debugf("Using %s from %s environment variable", opts.Server, defaultServerVar) } else { log.Debugf("No server specified or %s set, using /etc/resolv.conf", defaultServerVar) conf, err := dns.ClientConfigFromFile("/etc/resolv.conf") if err != nil { - opts.Server = "https://cloudflare-dns.com/dns-query" + opts.Server[0] = "https://cloudflare-dns.com/dns-query" log.Debugf("no server set, using %s", opts.Server) } else { if len(conf.Servers) == 0 { - opts.Server = "https://cloudflare-dns.com/dns-query" + opts.Server[0] = "https://cloudflare-dns.com/dns-query" log.Debugf("no server set, using %s", opts.Server) } else { - opts.Server = conf.Servers[0] + opts.Server[0] = conf.Servers[0] log.Debugf("found server %s from /etc/resolv.conf", opts.Server) } } @@ -363,8 +364,10 @@ All long form (--) flags can be toggled with the dig-standard +[no]flag notation if !strings.HasPrefix(opts.ODoHProxy, "https://") { return fmt.Errorf("ODoH proxy must use HTTPS") } - if !strings.HasPrefix(opts.Server, "https://") { - return fmt.Errorf("ODoH target must use HTTPS") + for _, v := range opts.Server { + if !strings.HasPrefix(v, "https://") { + return fmt.Errorf("ODoH target must use HTTPS") + } } } @@ -420,17 +423,20 @@ All long form (--) flags can be toggled with the dig-standard +[no]flag notation ) // Parse server address and transport type - server, transportType, err := parseServer(opts.Server) - if err != nil { - return err + type tserver struct { + server string + ttype transport.Type + trans *transport.Transport } - log.Debugf("Using server %s with transport %s", server, transportType) - - // QUIC specific overrides - if transportType == transport.TypeQUIC { - tlsConfig.NextProtos = opts.QUICALPNTokens - // Skip ID check if QUIC (https://datatracker.ietf.org/doc/html/rfc9250#section-4.2.1) - opts.NoIDCheck = true + servers := make([]*tserver, 0, len(opts.Server)) + for _, v := range opts.Server { + server, transportType, err := parseServer(v) + if err != nil { + log.Debugf("Skipping server %s due to error %s", v, err) + continue + } + log.Debugf("Adding server %s with transport %s", server, transportType) + servers = append(servers, &tserver{server, transportType, nil}) } // Recursive zone transfer @@ -438,28 +444,63 @@ All long form (--) flags can be toggled with the dig-standard +[no]flag notation if opts.Name == "" { return fmt.Errorf("no name specified for AXFR") } - _ = RecAXFR(opts.Name, server, out) + for _, v := range servers { + _ = RecAXFR(opts.Name, v.server, out) + } return nil } - // Create transport - txp, err := newTransport(server, transportType, tlsConfig) - if err != nil { - return err + // Create transports + for _, v := range servers { + tlsConfig := tlsConfig.Clone() + + // QUIC specific overrides + if v.ttype == transport.TypeQUIC { + tlsConfig.NextProtos = opts.QUICALPNTokens + // Skip ID check if QUIC (https://datatracker.ietf.org/doc/html/rfc9250#section-4.2.1) + opts.NoIDCheck = true // TODO: per-server overrides? + } + + // Create transport + txp, err := newTransport(v.server, v.ttype, tlsConfig) + if err != nil { + log.Debugf("Skipping server %s due to error %s", v.server, err) + continue + } + v.trans = txp } - startTime := time.Now() - var replies []*dns.Msg - for _, msg := range msgs { - reply, err := (*txp).Exchange(&msg) - if err != nil { - return err + // Filter failed servers, so we can preallocate responses + { + n := 0 + for _, v := range servers { + if v.trans == nil { + continue + } + servers[n] = v + n++ } + servers = servers[:n] + } - if !opts.NoIDCheck && reply.Id != msg.Id { - return fmt.Errorf("ID mismatch: expected %d, got %d", msg.Id, reply.Id) + // preallocate replies storage + replies := make([]*dns.Msg, 0, len(msgs) * len(servers)) + startTime := time.Now() + for _, v := range servers { + for _, msg := range msgs { + reply, err := (*v.trans).Exchange(&msg) + if err != nil { + log.Debugf("Skipping message %s with servers %s due to error %s", + &msg, v.server, err) + replies = append(replies, nil) // append nil so we can keep server sizes stable + continue + } + + if !opts.NoIDCheck && reply.Id != msg.Id { + return fmt.Errorf("ID mismatch: expected %d, got %d", msg.Id, reply.Id) + } + replies = append(replies, reply) } - replies = append(replies, reply) } queryTime := time.Since(startTime) @@ -474,27 +515,35 @@ All long form (--) flags can be toggled with the dig-standard +[no]flag notation output.PrettyPrintNSID(replies, out) } - printer := output.Printer{ - Server: server, - Out: out, - Opts: &opts, - QueryTime: queryTime, - NumReplies: len(replies), - Transport: txp, - } - if opts.Format == "column" { - printer.PrintColumn(replies) - } else { - for i, reply := range replies { - switch opts.Format { - case "pretty": - printer.PrintPretty(i, reply) - case "raw": - printer.PrintRaw(i, reply) - case "json", "yml", "yaml": - printer.PrintStructured(i, reply) - default: - return fmt.Errorf("invalid output format") + for i, v := range servers { + fmt.Println(v.server) + printer := output.Printer{ + Server: v.server, + Out: out, + Opts: &opts, + QueryTime: queryTime, + NumReplies: len(msgs), // TODO: filter nils + Transport: v.trans, + } + if opts.Format == "column" { + printer.PrintColumn(replies[i * len(msgs) : (i+1) * len(msgs)]) + } else { + for i, reply := range replies[i * len(msgs) : (i+1) * len(msgs)] { + if reply == nil { + continue + } + switch opts.Format { + case "pretty": + printer.PrintPretty(i, reply) + case "raw": + printer.PrintRaw(i, reply) + // TODO: jq and co can handle multipe separate json objects on stdin + // however, it would be nice to potentially return a [] instead + case "json", "yml", "yaml": + printer.PrintStructured(i, reply) + default: + return fmt.Errorf("invalid output format") + } } } } ```

natesales / q

Enhancement: support querying multiple servers at once #73