From 216bf5aa29ca1c162825c2c865977256bbe939a7 Mon Sep 17 00:00:00 2001
From: Alejandro Mery <amery@jpi.io>
Date: Mon, 28 Aug 2023 15:36:27 +0000
Subject: [PATCH] zones: WithLookuper()/WithResolver()

Signed-off-by: Alejandro Mery <amery@jpi.io>
---
 pkg/zones/options.go | 43 +++++++++++++++++++++++++++++++++++++------
 1 file changed, 37 insertions(+), 6 deletions(-)

diff --git a/pkg/zones/options.go b/pkg/zones/options.go
index 11b2afb..ae47e65 100644
--- a/pkg/zones/options.go
+++ b/pkg/zones/options.go
@@ -14,7 +14,41 @@ type ScanOption func(*Zones, *ScanOptions) error
 // ScanOptions contains flags used by the initial scan
 type ScanOptions struct{}
 
-func (*Zones) setDefaults(_ *ScanOptions) error {
+// WithLookuper specifies what resolver.Lookuper to use to
+// find public addresses
+func WithLookuper(h resolver.Lookuper) ScanOption {
+	return func(m *Zones, opt *ScanOptions) error {
+		if h == nil {
+			return fs.ErrInvalid
+		}
+		m.resolver = resolver.NewResolver(h)
+		return nil
+	}
+}
+
+// WithResolver specifies what resolver to use to find
+// public addresses. if nil is passed, the [net.Resolver] will be used.
+// The default is using Cloudflare's 1.1.1.1.
+func WithResolver(h resolver.Resolver) ScanOption {
+	return func(m *Zones, opt *ScanOptions) error {
+		if h == nil {
+			h = resolver.SystemResolver(true)
+		}
+
+		m.resolver = h
+		return nil
+	}
+}
+
+func (m *Zones) setDefaults(opt *ScanOptions) error {
+	if m.resolver == nil {
+		h := resolver.NewCloudflareLookuper()
+
+		if err := WithLookuper(h)(m, opt); err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
 
@@ -22,12 +56,9 @@ func (*Zones) setDefaults(_ *ScanOptions) error {
 func NewFS(dir fs.FS, domain string, opts ...ScanOption) (*Zones, error) {
 	var scanOptions ScanOptions
 
-	lockuper := resolver.NewCloudflareLookuper()
-
 	z := &Zones{
-		dir:      dir,
-		resolver: resolver.NewResolver(lockuper),
-		domain:   domain,
+		dir:    dir,
+		domain: domain,
 	}
 
 	for _, opt := range opts {