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

     1  package bbolt
     2  
     3  /*
     4  import (
     5  	"fmt"
     6  	"log"
     7  
     8  	"github.com/Asutorufa/yuhaiin/pkg/net/netapi"
     9  	bolt "go.etcd.io/bbolt"
    10  )
    11  
    12  func sbbolt(s *bolt.Bucket, domain string) (resp []byte, ok bool) {
    13  	z := newDomainStr(domain)
    14  	first, asterisk := true, false
    15  
    16  	for {
    17  		if !z.hasNext() || s == nil {
    18  			return
    19  		}
    20  
    21  		if r := s.Bucket([]byte(z.str())); r != nil {
    22  			if symbol := r.Get(bSymbolKey); symbol != nil {
    23  				// log.Println(symbol[0])
    24  				if symbol[0] == wildcard {
    25  					resp, ok = r.Get(bMarkKey), true
    26  				}
    27  
    28  				if symbol[0] == last && z.last() {
    29  					return r.Get(bMarkKey), true
    30  				}
    31  			}
    32  
    33  			s, first, _ = r, false, z.next()
    34  			continue
    35  		}
    36  
    37  		if !first {
    38  			return
    39  		}
    40  
    41  		if !asterisk {
    42  			s, asterisk = s.Bucket([]byte("*")), true
    43  			// log.Println("asterisk")
    44  		} else {
    45  			z.next()
    46  		}
    47  	}
    48  }
    49  
    50  var bSymbolKey = []byte{0x01}
    51  var bMarkKey = []byte{0x02}
    52  
    53  func storeBBoltMark(b *bolt.Bucket, symbol uint8, mark []byte) error {
    54  	if err := b.Put(bSymbolKey, []byte{symbol}); err != nil {
    55  		return err
    56  	}
    57  
    58  	if err := b.Put(bMarkKey, mark); err != nil {
    59  		return err
    60  	}
    61  
    62  	return nil
    63  }
    64  
    65  func insertBBolt(b *bolt.Bucket, domain string, mark []byte) error {
    66  	var err error
    67  	z := newDomainStr(domain)
    68  	for z.hasNext() {
    69  		if z.last() && domain[0] == '*' {
    70  			// b, err = b.CreateBucketIfNotExists([]byte("*"))
    71  			// if err != nil {
    72  			// return err
    73  			// }
    74  
    75  			if err := storeBBoltMark(b, wildcard, mark); err != nil {
    76  				return err
    77  			}
    78  			break
    79  		}
    80  
    81  		b, err = b.CreateBucketIfNotExists([]byte(z.str()))
    82  		if err != nil {
    83  			return fmt.Errorf("create bucket [%s](%s) failed: %w", domain, z.str(), err)
    84  		}
    85  
    86  		if z.last() {
    87  			if err = storeBBoltMark(b, last, mark); err != nil {
    88  				return err
    89  			}
    90  		}
    91  
    92  		z.next()
    93  	}
    94  
    95  	return nil
    96  }
    97  
    98  type bboltDomain struct{ db *bolt.DB }
    99  
   100  func (d *bboltDomain) Insert(domain string, mark []byte) {
   101  	if len(domain) == 0 {
   102  		return
   103  	}
   104  
   105  	err := d.db.Update(func(tx *bolt.Tx) (err error) {
   106  		var b *bolt.Bucket
   107  		if domain[0] == '*' {
   108  			b, err = tx.CreateBucketIfNotExists([]byte{'w', 'i', 'l', 'd', 'c', 'a', 'r', 'd', 'r', 'o', 'o', 't'})
   109  		} else {
   110  			b, err = tx.CreateBucketIfNotExists([]byte{'r', 'o', 'o', 't'})
   111  		}
   112  		if err != nil {
   113  			return err
   114  		}
   115  
   116  		if err := insertBBolt(b, domain, mark); err != nil {
   117  			return fmt.Errorf("insert bbolt failed: %w", err)
   118  		}
   119  
   120  		return nil
   121  	})
   122  	if err != nil {
   123  		log.Printf("insert %s failed: %w", domain, err)
   124  	}
   125  }
   126  
   127  func (d *bboltDomain) Search(domain netapi.Address) (mark []byte, ok bool) {
   128  	err := d.db.View(func(tx *bolt.Tx) error {
   129  		b := tx.Bucket([]byte{'r', 'o', 'o', 't'})
   130  		if b != nil {
   131  			mark, ok = sbbolt(b, domain.Hostname())
   132  			if ok {
   133  				return nil
   134  			}
   135  		}
   136  
   137  		b = tx.Bucket([]byte{'w', 'i', 'l', 'd', 'c', 'a', 'r', 'd', 'r', 'o', 'o', 't'})
   138  		if b != nil {
   139  			mark, ok = sbbolt(b, domain.Hostname())
   140  		}
   141  
   142  		return nil
   143  	})
   144  	if err != nil {
   145  		log.Println(err)
   146  	}
   147  
   148  	return
   149  }
   150  
   151  func (d *bboltDomain) Clear() error {
   152  	return d.db.Update(func(tx *bolt.Tx) error {
   153  		if err := tx.DeleteBucket([]byte{'r', 'o', 'o', 't'}); err != nil {
   154  			return err
   155  		}
   156  		return tx.DeleteBucket([]byte{'w', 'i', 'l', 'd', 'c', 'a', 'r', 'd', 'r', 'o', 'o', 't'})
   157  	})
   158  }
   159  
   160  func NewBBoltDomainMapper(db *bolt.DB) *bboltDomain {
   161  	return &bboltDomain{db}
   162  }
   163  */