Fix API rate limit bug

Add `onmicrosoft.com` DKIM selector fallback (#35)

Author: wolveix

cmd/dss/scan.go

@@ -16,7 +16,7 @@ func newScanCMD() *cobra.Command {
 		Example: "  dss scan <STDIN>\n  dss scan globalcyberalliance.org gcaaide.org google.com\n  dss scan -z < zonefile",
 		Short:   "Scan DNS records for one or multiple domains.",
 		Long:    "Scan DNS records for one or multiple domains.\nBy default, the command will listen on STDIN, allowing you to type or pipe multiple domains.",
-		Run: func(_ *cobra.Command, args []string) {
+		Run: func(cmd *cobra.Command, args []string) {
 			sc, err := scanner.New(log, timeout, getScannerOpts()...)
 			if err != nil {
 				log.Fatal().Err(err).Msg("An unexpected error occurred.")
@@ -56,7 +56,7 @@ func newScanCMD() *cobra.Command {
 				for lineScanner.Scan() {
 					domain := lineScanner.Text()
 
-					results, err = sc.Scan(domain)
+					results, err = sc.Scan(cmd.Context(), domain)
 					if err != nil {
 						log.Fatal().Err(err).Msg("An unexpected error occurred.")
 					}
@@ -74,7 +74,7 @@ func newScanCMD() *cobra.Command {
 					printToConsole(resultsWithAdvice)
 				}
 			} else {
-				results, err = sc.Scan(args...)
+				results, err = sc.Scan(cmd.Context(), args...)
 				if err != nil {
 					log.Fatal().Err(err).Msg("An unexpected error occurred.")
 				}

pkg/http/scan.go

@@ -26,7 +26,7 @@ func (s *Server) registerScanRoutes() {
 		Method:      http.MethodGet,
 		Path:        s.apiPath + "/scan/{domain}",
 		Tags:        []string{"Scan Domains"},
-	}, func(_ context.Context, input *ScanSingleDomainRequest) (*ScanSingleDomainResponse, error) {
+	}, func(ctx context.Context, input *ScanSingleDomainRequest) (*ScanSingleDomainResponse, error) {
 		resp := ScanSingleDomainResponse{}
 
 		if len(input.DKIMSelectors) > 0 {
@@ -35,7 +35,7 @@ func (s *Server) registerScanRoutes() {
 			}
 		}
 
-		results, err := s.Scanner.Scan(input.Domain)
+		results, err := s.Scanner.Scan(ctx, input.Domain)
 		if err != nil {
 			return nil, huma.Error500InternalServerError(err.Error())
 		}
@@ -80,10 +80,10 @@ func (s *Server) registerScanRoutes() {
 		Method:      http.MethodPost,
 		Path:        s.apiPath + "/scan",
 		Tags:        []string{"Scan Domains"},
-	}, func(_ context.Context, input *ScanBulkDomainsRequest) (*ScanBulkDomainResponse, error) {
+	}, func(ctx context.Context, input *ScanBulkDomainsRequest) (*ScanBulkDomainResponse, error) {
 		resp := ScanBulkDomainResponse{}
 
-		results, err := s.Scanner.Scan(input.Body.Domains...)
+		results, err := s.Scanner.Scan(ctx, input.Body.Domains...)
 		if err != nil {
 			return nil, huma.Error500InternalServerError(err.Error())
 		}

pkg/http/server.go

@@ -65,6 +65,7 @@ func NewServer(logger zerolog.Logger, timeout time.Duration, rateLimit int, vers
 		MaxAge:           300, // The maximum value not ignored by any of the major browsers.
 	}))
 	mux.Use(httprate.Limit(rateLimit, 3*time.Second,
+		httprate.WithKeyFuncs(httprate.KeyByRealIP),
 		httprate.WithLimitHandler(func(w http.ResponseWriter, _ *http.Request) {
 			response, err := json.Marshal(huma.Error429TooManyRequests("try again later"))
 			if err != nil {

pkg/mail/server.go

@@ -1,6 +1,7 @@
 package mail
 
 import (
+	"context"
 	"fmt"
 	htmlTmpl "html/template"
 	textTmpl "text/template"
@@ -100,7 +101,7 @@ func (s *Server) handler() error {
 				}
 			}
 
-			results, err := s.Scanner.Scan(domainList...)
+			results, err := s.Scanner.Scan(context.Background(), domainList...)
 			if err != nil {
 				s.logger.Error().Err(err).Msg("An error occurred while scanning domains")
 				continue

pkg/scanner/requests.go

@@ -19,8 +19,8 @@ var (
 	DMARCPrefix = regexp.MustCompile(`^v\s*=\s*DMARC1`) // Matches v=DMARC1 with whitespace (RFC7489).
 	SPFPrefix   = regexp.MustCompile(`^v=(?i)spf1`)
 
-	// knownDkimSelectors is a list of known DKIM selectors.
-	knownDkimSelectors = []string{
+	// knownDKIMSelectors is a list of known DKIM selectors.
+	knownDKIMSelectors = []string{
 		"x",             // Generic
 		"google",        // Google
 		"selector1",     // Microsoft
@@ -94,7 +94,7 @@ func (s *Scanner) getDNSAnswers(domain string, recordType uint16) ([]dns.RR, err
 	}
 
 	if in.Rcode != dns.RcodeSuccess {
-		// disregard NXDOMAIN errors
+		// Disregard NXDOMAIN errors.
 		if in.Rcode == dns.RcodeNameError {
 			return nil, nil
 		}
@@ -140,7 +140,11 @@ func (s *Scanner) getTypeBIMI(domain string) (string, error) {
 // getTypeDKIM queries the DNS server for DKIM records of a domain.
 // It returns a string (DKIM record) and an error if any occurred.
 func (s *Scanner) getTypeDKIM(domain string) (string, error) {
-	selectors := append(s.dkimSelectors, knownDkimSelectors...)
+	selectors := append(s.dkimSelectors, knownDKIMSelectors...)
+
+	// Some Microsoft domains (e.g. "onmicrosoft.com") use dynamic selectors.
+	dashedDomain := strings.ReplaceAll(strings.ToLower(domain), ".", "-")
+	selectors = append(selectors, "selector1-"+dashedDomain, "selector2-"+dashedDomain)
 
 	for _, selector := range selectors {
 		records, err := s.getDNSRecords(selector+"._domainkey."+domain, dns.TypeTXT)
@@ -150,7 +154,7 @@ func (s *Scanner) getTypeDKIM(domain string) (string, error) {
 
 		for index, record := range records {
 			if strings.HasPrefix(record, DKIMPrefix) {
-				// TXT records can be split across multiple strings, so we need to join them
+				// TXT records can be split across multiple strings, we need to join them.
 				return strings.Join(records[index:], ""), nil
 			}
 		}
@@ -173,7 +177,7 @@ func (s *Scanner) getTypeDMARC(domain string) (string, error) {
 
 		for index, record := range records {
 			if DMARCPrefix.MatchString(record) {
-				// TXT records can be split across multiple strings, so we need to join them
+				// TXT records can be split across multiple strings, we need to join them.
 				return strings.Join(records[index:], ""), nil
 			}
 		}

pkg/scanner/scanner.go

@@ -1,6 +1,7 @@
 package scanner
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"runtime"
@@ -113,7 +114,7 @@ func New(logger zerolog.Logger, timeout time.Duration, opts ...Option) (*Scanner
 }
 
 // Scan scans a list of domains and returns the results.
-func (s *Scanner) Scan(domains ...string) ([]*Result, error) {
+func (s *Scanner) Scan(ctx context.Context, domains ...string) ([]*Result, error) {
 	if s.pool == nil {
 		return nil, errors.New("scanner is closed")
 	}
@@ -140,6 +141,12 @@ func (s *Scanner) Scan(domains ...string) ([]*Result, error) {
 				wg.Done()
 			}()
 
+			select {
+			case <-ctx.Done():
+				return
+			default:
+			}
+
 			var err error
 			result := &Result{
 				Domain: domainToScan,
@@ -247,6 +254,10 @@ func (s *Scanner) Scan(domains ...string) ([]*Result, error) {
 
 	wg.Wait()
 
+	if ctx.Err() != nil {
+		return nil, ctx.Err()
+	}
+
 	return results, nil
 }
 
@@ -274,7 +285,7 @@ func (s *Scanner) ScanZone(zone io.Reader) ([]*Result, error) {
 		domains = append(domains, domain)
 	}
 
-	return s.Scan(domains...)
+	return s.Scan(context.Background(), domains...)
 }
 
 // Close closes the scanner.