github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/net/trie/domain/domain.go (about)

     1  package domain
     2  
     3  import (
     4  	"encoding/json"
     5  	"sync"
     6  
     7  	"github.com/Asutorufa/yuhaiin/pkg/net/netapi"
     8  )
     9  
    10  type Fqdn[T any] struct {
    11  	Root         *trie[T] `json:"root"`          // for example.com, example.*
    12  	WildcardRoot *trie[T] `json:"wildcard_root"` // for *.example.com, *.example.*
    13  	mu           sync.Mutex
    14  }
    15  
    16  func (d *Fqdn[T]) Insert(domain string, mark T) {
    17  	d.mu.Lock()
    18  	defer d.mu.Unlock()
    19  
    20  	if len(domain) == 0 {
    21  		return
    22  	}
    23  
    24  	r := newReader(domain)
    25  	if domain[0] == '*' {
    26  		insert(d.WildcardRoot, r, mark)
    27  	} else {
    28  		insert(d.Root, r, mark)
    29  	}
    30  }
    31  
    32  func (d *Fqdn[T]) Search(domain netapi.Address) (mark T, ok bool) {
    33  	r := newReader(domain.Hostname())
    34  
    35  	mark, ok = search(d.Root, r)
    36  	if ok {
    37  		return
    38  	}
    39  
    40  	r.reset()
    41  
    42  	return search(d.WildcardRoot, r)
    43  }
    44  func (d *Fqdn[T]) SearchString(domain string) (mark T, ok bool) {
    45  	r := newReader(domain)
    46  
    47  	mark, ok = search(d.Root, r)
    48  	if ok {
    49  		return
    50  	}
    51  
    52  	r.reset()
    53  
    54  	return search(d.WildcardRoot, r)
    55  }
    56  
    57  func (d *Fqdn[T]) Remove(domain string) {
    58  	d.mu.Lock()
    59  	defer d.mu.Unlock()
    60  
    61  	if len(domain) == 0 {
    62  		return
    63  	}
    64  
    65  	r := newReader(domain)
    66  	if domain[0] == '*' {
    67  		remove(d.WildcardRoot, r)
    68  	} else {
    69  		remove(d.Root, r)
    70  	}
    71  }
    72  
    73  func (d *Fqdn[T]) Marshal() ([]byte, error) {
    74  	return json.MarshalIndent(d, "", "  ")
    75  }
    76  
    77  func (d *Fqdn[T]) Clear() error {
    78  	d.Root = &trie[T]{Child: map[string]*trie[T]{}}
    79  	d.WildcardRoot = &trie[T]{Child: map[string]*trie[T]{}}
    80  	return nil
    81  }
    82  
    83  func NewDomainMapper[T any]() *Fqdn[T] {
    84  	return &Fqdn[T]{
    85  		Root:         &trie[T]{Child: map[string]*trie[T]{}},
    86  		WildcardRoot: &trie[T]{Child: map[string]*trie[T]{}},
    87  	}
    88  }