cleanup: use go-set in place of custom set (#10)

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/coredns/caddy v1.1.1
 	github.com/coredns/coredns v1.8.5
 	github.com/hashicorp/go-cleanhttp v0.5.2
+	github.com/hashicorp/go-set v0.1.6
 	github.com/miekg/dns v1.1.43
 	github.com/shoenig/extractors v0.3.0
 	github.com/shoenig/ignore v0.4.0

go.sum

@@ -253,6 +253,8 @@ github.com/grpc-ecosystem/grpc-opentracing v0.0.0-20180507213350-8e809c8a8645 h1
 github.com/grpc-ecosystem/grpc-opentracing v0.0.0-20180507213350-8e809c8a8645/go.mod h1:6iZfnjpejD4L/4DwD7NryNaJyCQdzwWwH2MWhCA90Kw=
 github.com/hashicorp/go-cleanhttp v0.5.2 h1:035FKYIWjmULyFRBKPs8TBQoi0x6d9G4xc9neXJWAZQ=
 github.com/hashicorp/go-cleanhttp v0.5.2/go.mod h1:kO/YDlP8L1346E6Sodw+PrpBSV4/SoxCXGY6BqNFT48=
+github.com/hashicorp/go-set v0.1.6 h1:fj/JG5B97sAOd9OQN4GL880yCE384fz3asNDpGbxkPo=
+github.com/hashicorp/go-set v0.1.6/go.mod h1:ELvMcE+1mRHYPVgTFSQiecObIwZRxY5Q11EhbtmM5KQ=
 github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
 github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
 github.com/hexdigest/gowrap v1.1.7/go.mod h1:Z+nBFUDLa01iaNM+/jzoOA1JJ7sm51rnYFauKFUB5fs=

plugins/donutdns/plugin.go

@@ -3,11 +3,12 @@ package donutdns
 import (
 	"context"
 	"net"
+	"strings"
 
 	"github.com/coredns/coredns/plugin"
 	"github.com/coredns/coredns/request"
+	"github.com/hashicorp/go-set"
 	"github.com/miekg/dns"
-	"github.com/shoenig/donutdns/sources/set"
 )
 
 const (
@@ -19,20 +20,20 @@ type DonutDNS struct {
 	Next plugin.Handler
 
 	defaultLists bool
-	block        *set.Set
-	allow        *set.Set
+	block        *set.Set[string]
+	allow        *set.Set[string]
 }
 
 func (dd DonutDNS) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
 	state := request.Request{W: w, Req: r}
-	query := state.Name()
+	query := strings.TrimSuffix(state.Name(), ".")
 
-	if dd.allow.Has(query) {
+	if dd.allow.Contains(query) {
 		pLog.Infof("query for %s is explicitly allowed", query)
 		return plugin.NextOrFailure(dd.Name(), dd.Next, ctx, w, r)
 	}
 
-	if !dd.block.Has(query) {
+	if !dd.block.Contains(query) {
 		pLog.Debugf("query for %s is implicitly allowed", query)
 		return plugin.NextOrFailure(dd.Name(), dd.Next, ctx, w, r)
 	}

plugins/donutdns/setup.go

@@ -7,10 +7,10 @@ import (
 	"github.com/coredns/coredns/core/dnsserver"
 	"github.com/coredns/coredns/plugin"
 	"github.com/coredns/coredns/plugin/pkg/log"
+	"github.com/hashicorp/go-set"
 	"github.com/shoenig/donutdns/sources"
 	"github.com/shoenig/donutdns/sources/extract"
 	"github.com/shoenig/donutdns/sources/fetch"
-	"github.com/shoenig/donutdns/sources/set"
 	"github.com/shoenig/ignore"
 )
 
@@ -25,8 +25,8 @@ func setup(c *caddy.Controller) error {
 
 	dd := DonutDNS{
 		defaultLists: true,
-		block:        set.New(),
-		allow:        set.New(),
+		block:        set.New[string](100),
+		allow:        set.New[string](100),
 	}
 
 	for c.Next() {
@@ -62,19 +62,19 @@ func setup(c *caddy.Controller) error {
 				if !c.NextArg() {
 					return c.ArgErr()
 				}
-				dd.block.Add(c.Val())
+				dd.block.Insert(c.Val())
 
 			case "allow":
 				if !c.NextArg() {
 					return c.ArgErr()
 				}
-				dd.allow.Add(c.Val())
+				dd.allow.Insert(c.Val())
 			}
 		}
 	}
 
-	pLog.Infof("domains on custom allow-list: %d", dd.allow.Len())
-	pLog.Infof("domains on custom block-list: %d", dd.block.Len())
+	pLog.Infof("domains on custom allow-list: %d", dd.allow.Size())
+	pLog.Infof("domains on custom block-list: %d", dd.block.Size())
 
 	// Add the Plugin to CoreDNS, so Servers can use it in their plugin chain.
 	dnsserver.GetConfig(c).AddPlugin(func(next plugin.Handler) plugin.Handler {
@@ -86,16 +86,16 @@ func setup(c *caddy.Controller) error {
 	return nil
 }
 
-func defaults(set *set.Set) {
+func defaults(set *set.Set[string]) {
 	downloader := fetch.NewDownloader(pLog)
 	s, err := downloader.Download(sources.Defaults())
 	if err != nil {
 		panic(err)
 	}
-	set.Union(s)
+	set.InsertSet(s)
 }
 
-func custom(filename string, set *set.Set) {
+func custom(filename string, set *set.Set[string]) {
 	// for now, everything uses the generic domain extractor
 	ex := extract.New(extract.Generic)
 	f, err := os.Open(filename)
@@ -107,5 +107,5 @@ func custom(filename string, set *set.Set) {
 	if err != nil {
 		panic(err)
 	}
-	set.Union(s)
+	set.InsertSet(s)
 }

sources/extract/extract.go

@@ -5,7 +5,7 @@ import (
 	"io"
 	"regexp"
 
-	"github.com/shoenig/donutdns/sources/set"
+	"github.com/hashicorp/go-set"
 )
 
 const (
@@ -18,7 +18,7 @@ const (
 
 // An Extractor reads content from an io.Reader and extracts domains into a Set.
 type Extractor interface {
-	Extract(io.Reader) (*set.Set, error)
+	Extract(io.Reader) (*set.Set[string], error)
 }
 
 type extractor struct {
@@ -32,13 +32,11 @@ func New(re string) Extractor {
 	}
 }
 
-func (e *extractor) Extract(r io.Reader) (*set.Set, error) {
+func (e *extractor) Extract(r io.Reader) (*set.Set[string], error) {
 	scanner := bufio.NewScanner(r)
-	s := set.New()
+	s := set.New[string](100)
 	for scanner.Scan() {
-		line := scanner.Text()
-		domain := e.parse(line)
-		s.Add(domain)
+		e.insert(s, scanner.Text())
 	}
 	if err := scanner.Err(); err != nil {
 		return nil, err
@@ -54,6 +52,12 @@ func (e *extractor) parse(line string) string {
 	case line[0] == '#':
 		return ""
 	}
-
 	return e.re.FindString(line)
 }
+
+func (e *extractor) insert(s *set.Set[string], line string) {
+	domain := e.parse(line)
+	if domain != "" {
+		s.Insert(domain)
+	}
+}

sources/extract/extract_test.go

@@ -4,35 +4,42 @@ import (
 	"bytes"
 	"fmt"
 	"io"
-	"io/ioutil"
+	"os"
 	"testing"
 
 	"github.com/shoenig/test/must"
 )
 
 func openSample(t *testing.T, filename string) io.Reader {
-	b, err := ioutil.ReadFile(fmt.Sprintf("samples/%s", filename))
+	b, err := os.ReadFile(fmt.Sprintf("samples/%s", filename))
 	must.NoError(t, err)
 	return bytes.NewBuffer(b)
 }
 
 func TestExtractor_Extract(t *testing.T) {
-
-	try := func(filename, re string, exp int) {
-		ex := New(Generic)
-		result, err := ex.Extract(openSample(t, filename))
-		must.NoError(t, err)
-		must.EqOp(t, exp, result.Len())
+	cases := []struct {
+		file string
+		mode string
+		exp  int
+	}{
+		{"KADhosts.txt", Generic, 6},
+		{"w3kbl.txt", Generic, 7},
+		{"adaway.txt", Generic, 3},
+		{"Admiral.txt", Generic, 1},
+		{"adservers.txt", Generic, 2},
+		{"hostsVN.txt", Generic, 3},
+		{"AntiMalwareHosts.txt", Generic, 6}, // matches some ipv4 addresses
+		{"Prigent-Crypto.txt", Generic, 4},
+		{"notrack-malware.txt", Generic, 4},
+		{"abuse.txt", Generic, 5},
 	}
 
-	try("KADhosts.txt", Generic, 6)
-	try("w3kbl.txt", Generic, 7)
-	try("adaway.txt", Generic, 3)
-	try("Admiral.txt", Generic, 1)
-	try("adservers.txt", Generic, 2)
-	try("hostsVN.txt", Generic, 3)
-	try("AntiMalwareHosts.txt", Generic, 6) // matches some ipv4 addresses
-	try("Prigent-Crypto.txt", Generic, 4)
-	try("notrack-malware.txt", Generic, 4)
-	try("abuse.txt", Generic, 5)
+	for _, tc := range cases {
+		t.Run(tc.file, func(t *testing.T) {
+			ex := New(tc.mode)
+			result, err := ex.Extract(openSample(t, tc.file))
+			must.NoError(t, err)
+			must.Eq(t, tc.exp, result.Size(), must.Sprintf("result: %v", result))
+		})
+	}
 }

sources/fetch/fetch.go

@@ -5,16 +5,16 @@ import (
 	"net/http"
 
 	"github.com/coredns/coredns/plugin/pkg/log"
+	"github.com/hashicorp/go-set"
 	"github.com/shoenig/donutdns/sources"
 	"github.com/shoenig/donutdns/sources/extract"
-	"github.com/shoenig/donutdns/sources/set"
 	"github.com/shoenig/ignore"
 )
 
 // A Downloader is used to download a set of source lists.
 type Downloader interface {
 	// Download all sources in Lists.
-	Download(*sources.Lists) (*set.Set, error)
+	Download(*sources.Lists) (*set.Set[string], error)
 }
 
 type downloader struct {
@@ -28,24 +28,24 @@ func NewDownloader(pLog log.P) Downloader {
 	}
 }
 
-func (d *downloader) Download(lists *sources.Lists) (*set.Set, error) {
+func (d *downloader) Download(lists *sources.Lists) (*set.Set[string], error) {
 	g := NewGetter(d.pLog, extract.New(extract.Generic))
-	combo := set.New()
+	combo := set.New[string](100)
 	for _, source := range lists.All() {
 		single, err := g.Get(source)
 		if err != nil {
 			d.pLog.Errorf("failed to fetch source %q, skip: %s", source, err)
 			continue
 		}
-		combo.Union(single)
+		combo.InsertSet(single)
 	}
 	return combo, nil
 }
 
 // A Getter is used to download a single source list.
 type Getter interface {
 	// Get source and extract its domains into a Set.
-	Get(source string) (*set.Set, error)
+	Get(source string) (*set.Set[string], error)
 }
 
 type getter struct {
@@ -66,7 +66,7 @@ func NewGetter(pLog log.P, ex extract.Extractor) Getter {
 	}
 }
 
-func (g *getter) Get(source string) (*set.Set, error) {
+func (g *getter) Get(source string) (*set.Set[string], error) {
 	request, err := http.NewRequest(http.MethodGet, source, nil)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create request: %w", err)
@@ -88,7 +88,7 @@ func (g *getter) Get(source string) (*set.Set, error) {
 		return nil, fmt.Errorf("failed to extract sources: %w", err)
 	}
 
-	g.plog.Infof("got %d domains from %q", single.Len(), source)
+	g.plog.Infof("got %d domains from %q", single.Size(), source)
 
 	return single, nil
 }

sources/fetch/fetch_test.go

@@ -30,7 +30,7 @@ func Test_Get(t *testing.T) {
 	g := NewGetter(pLog, ex)
 	s, err := g.Get(ts.URL)
 	must.NoError(t, err)
-	must.EqOp(t, 3, s.Len())
+	must.EqOp(t, 3, s.Size())
 }
 
 func Test_Download(t *testing.T) {
@@ -52,6 +52,6 @@ func Test_Download(t *testing.T) {
 	d := NewDownloader(pLog)
 	s, err := d.Download(lists)
 	must.NoError(t, err)
-	must.EqOp(t, 3, s.Len())
+	must.EqOp(t, 3, s.Size())
 	must.EqOp(t, 5, hit)
 }