Better trie implementations

dataset/storage.go

@@ -1,21 +1,20 @@
 package dataset
 
 import (
-	"bufio"
 	"bytes"
 	"fmt"
 	"io"
 	"io/fs"
 	"net"
 	"net/http"
+	"net/netip"
 	"net/url"
 	"os"
 	"slices"
-	"strings"
 	"time"
 
+	"git.maze.io/maze/styx/dataset/parser"
 	_ "github.com/mattn/go-sqlite3" // SQLite3 driver
-	"github.com/miekg/dns"
 )
 
 type Storage interface {
@@ -27,11 +26,13 @@ type Storage interface {
 
 	Clients() (Clients, error)
 	ClientByID(int64) (Client, error)
-	ClientByIP(net.IP) (Client, error)
+	ClientByAddr(netip.Addr) (Client, error)
+	// ClientByIP(net.IP) (Client, error)
 	SaveClient(*Client) error
 	DeleteClient(Client) error
 
 	Lists() ([]List, error)
+	ListsByGroup(Group) ([]List, error)
 	ListByID(int64) (List, error)
 	SaveList(*List) error
 	DeleteList(List) error
@@ -44,7 +45,6 @@ type Group struct {
 	Description string    `json:"description"`
 	CreatedAt   time.Time `json:"created_at" bstore:"nonzero"`
 	UpdatedAt   time.Time `json:"updated_at" bstore:"nonzero"`
-	Storage     Storage   `json:"-" bstore:"-"`
 }
 
 type Client struct {
@@ -56,11 +56,6 @@ type Client struct {
 	Groups      []Group   `json:"groups,omitempty" bstore:"-"`
 	CreatedAt   time.Time `json:"created_at" bstore:"nonzero"`
 	UpdatedAt   time.Time `json:"updated_at" bstore:"nonzero"`
-	Storage     Storage   `json:"-" bstore:"-"`
-}
-
-type WithClient interface {
-	Client() (Client, error)
 }
 
 type ClientGroup struct {
@@ -80,6 +75,15 @@ func (c *Client) ContainsIP(ip net.IP) bool {
 	return ipnet.Contains(ip)
 }
 
+func (c *Client) ContainsAddr(ip netip.Addr) bool {
+	return c.Prefix().Contains(ip)
+}
+
+func (c Client) Prefix() netip.Prefix {
+	ip, _ := netip.ParseAddr(c.IP)
+	return netip.PrefixFrom(ip, c.Mask)
+}
+
 func (c *Client) String() string {
 	ipnet := &net.IPNet{
 		IP:   net.ParseIP(c.IP),
@@ -136,28 +140,29 @@ type List struct {
 	UpdatedAt    time.Time     `json:"updated_at"`
 }
 
-func (list *List) Domains() (*DomainTree, error) {
+func (list *List) Networks() (*NetworkTrie, error) {
+	if list.Type != ListTypeNetwork {
+		return nil, nil
+	}
+
+	prefixes, _, err := parser.ParseNetworks(bytes.NewReader(list.Cache))
+	if err != nil {
+		return nil, err
+	}
+	return NewNetworkTrie(prefixes...), nil
+}
+
+func (list *List) Domains() (*DomainTrie, error) {
 	if list.Type != ListTypeDomain {
 		return nil, nil
 	}
 
-	var (
-		tree = NewDomainList()
-		scan = bufio.NewScanner(bytes.NewReader(list.Cache))
-	)
-	for scan.Scan() {
-		line := strings.TrimSpace(scan.Text())
-		if line == "" || line[0] == '#' {
-			continue
-		}
-		if labels, ok := dns.IsDomainName(line); ok && labels >= 2 {
-			tree.Add(line)
-		}
-	}
-	if err := scan.Err(); err != nil {
+	domains, _, err := parser.ParseDomains(bytes.NewReader(list.Cache))
+	if err != nil {
 		return nil, err
 	}
-	return tree, nil
+
+	return NewDomainTrie(list.Permit, domains...)
 }
 
 func (list *List) Update() (updated bool, err error) {