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

     1  package bbolt
     2  
     3  /*
     4  import (
     5  	"bufio"
     6  	"bytes"
     7  	"os"
     8  	"strings"
     9  	"testing"
    10  
    11  	"github.com/Asutorufa/yuhaiin/internal/statistics"
    12  	"github.com/Asutorufa/yuhaiin/pkg/net/netapi"
    13  	"github.com/Asutorufa/yuhaiin/pkg/utils/assert"
    14  	"github.com/Asutorufa/yuhaiin/pkg/utils/syncmap"
    15  	bolt "go.etcd.io/bbolt"
    16  )
    17  
    18  func TestBBlotDomainMatcherSearch(t *testing.T) {
    19  	os.MkdirAll("/tmp/bbolt", os.ModePerm)
    20  	b, err := bolt.Open("/tmp/bbolt/test.db", os.ModePerm, &bolt.Options{})
    21  	assert.NoError(t, err)
    22  	defer b.Close()
    23  
    24  	root := NewBBoltDomainMapper(b)
    25  	insert := func(d, str string) {
    26  		root.Insert(d, []byte(str))
    27  	}
    28  
    29  	insert("*.baidu.com", "sub_baidu_test")
    30  	insert("www.baidu.com", "test_baidu")
    31  	insert("last.baidu.*", "test_last_baidu")
    32  	insert("*.baidu.*", "last_sub_baidu_test")
    33  	insert("spo.baidu.com", "test_no_sub_baidu")
    34  	insert("www.google.com", "test_google")
    35  	insert("music.111.com", "1111")
    36  	insert("163.com", "163")
    37  	insert("*.google.com", "google")
    38  	insert("*.dl.google.com", "google_dl")
    39  	insert("api.sec.miui.*", "ad_miui")
    40  	insert("*.miui.com", "miui")
    41  
    42  	search := func(s string) string {
    43  		res, _ := root.Search(netapi.ParseAddressSplit("", s, 0))
    44  		return string(res)
    45  	}
    46  	assert.Equal(t, "test_baidu", search("www.baidu.com"))
    47  	assert.Equal(t, "test_no_sub_baidu", search("spo.baidu.com"))
    48  	assert.Equal(t, "test_last_baidu", search("last.baidu.com.cn"))
    49  	assert.Equal(t, "sub_baidu_test", search("test.baidu.com"))
    50  	assert.Equal(t, "sub_baidu_test", search("test.test2.baidu.com"))
    51  	assert.Equal(t, "last_sub_baidu_test", search("www.baidu.cn"))
    52  	assert.Equal(t, "test_google", search("www.google.com"))
    53  	assert.Equal(t, "", search("www.google.cn"))
    54  	assert.Equal(t, "", search("music.163.com"))
    55  	assert.Equal(t, "163", search("163.com"))
    56  	assert.Equal(t, "google", search("www.x.google.com"))
    57  	assert.Equal(t, "google_dl", search("dl.google.com"))
    58  	assert.Equal(t, "ad_miui", search("api.sec.miui.com"))
    59  }
    60  
    61  func TestInsertBBolt(t *testing.T) {
    62  	rule := &rule{}
    63  	os.MkdirAll("/tmp/bbolt", os.ModePerm)
    64  	b, err := bolt.Open("/tmp/bbolt/test.db", os.ModePerm, &bolt.Options{})
    65  	assert.NoError(t, err)
    66  	defer b.Close()
    67  
    68  	root := NewBBoltDomainMapper(b)
    69  	insert := func(d string, str []byte) {
    70  		root.Insert(d, str)
    71  	}
    72  
    73  	f, err := os.Open("/mnt/share/Work/code/shell/ACL/yuhaiin/yuhaiin_my.conf")
    74  	assert.NoError(t, err)
    75  	defer f.Close()
    76  
    77  	br := bufio.NewScanner(f)
    78  	for {
    79  		if !br.Scan() {
    80  			break
    81  		}
    82  
    83  		a := br.Bytes()
    84  
    85  		i := bytes.IndexByte(a, '#')
    86  		if i != -1 {
    87  			a = a[:i]
    88  		}
    89  
    90  		i = bytes.IndexByte(a, ' ')
    91  		if i == -1 {
    92  			continue
    93  		}
    94  
    95  		c, b := a[:i], a[i+1:]
    96  
    97  		if len(c) != 0 && len(b) != 0 {
    98  			insert(string(c), []byte{byte(rule.GetID(string(b)))})
    99  		}
   100  	}
   101  }
   102  
   103  type rule struct {
   104  	id        statistics.IDGenerator
   105  	mapping   syncmap.SyncMap[string, uint16]
   106  	idMapping syncmap.SyncMap[uint16, string]
   107  }
   108  
   109  func (r *rule) GetID(s string) uint16 {
   110  	s = strings.ToUpper(s)
   111  	if v, ok := r.mapping.Load(s); ok {
   112  		return v
   113  	}
   114  	id := uint16(r.id.Generate())
   115  	r.mapping.Store(s, id)
   116  	r.idMapping.Store(id, s)
   117  	return id
   118  }
   119  
   120  func (r *rule) GetMode(id uint16) string {
   121  	if v, ok := r.idMapping.Load(id); ok {
   122  		return v
   123  	}
   124  	return ""
   125  }
   126  */