github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/net/proxy/shadowsocks/internal/bloomring.go (about)

     1  package internal
     2  
     3  import (
     4  	"hash/fnv"
     5  	"sync"
     6  
     7  	"github.com/Asutorufa/yuhaiin/pkg/net/proxy/shadowsocks/bloom"
     8  )
     9  
    10  // simply use Double FNV here as our Bloom Filter hash
    11  func doubleFNV(b []byte) (uint64, uint64) {
    12  	hx := fnv.New64()
    13  	hx.Write(b)
    14  	x := hx.Sum64()
    15  	hy := fnv.New64a()
    16  	hy.Write(b)
    17  	y := hy.Sum64()
    18  	return x, y
    19  }
    20  
    21  type BloomRing struct {
    22  	slotCapacity int
    23  	slotPosition int
    24  	slotCount    int
    25  	entryCounter int
    26  	slots        []bloom.Filter
    27  	mu           sync.RWMutex
    28  }
    29  
    30  func NewBloomRing(slot, capacity int, falsePositiveRate float64) *BloomRing {
    31  	// Calculate entries for each slot
    32  	r := &BloomRing{
    33  		slotCapacity: capacity / slot,
    34  		slotCount:    slot,
    35  		slots:        make([]bloom.Filter, slot),
    36  	}
    37  	for i := 0; i < slot; i++ {
    38  		r.slots[i] = bloom.New(r.slotCapacity, falsePositiveRate, doubleFNV)
    39  	}
    40  	return r
    41  }
    42  
    43  func (r *BloomRing) Add(b []byte) {
    44  	if r == nil {
    45  		return
    46  	}
    47  	r.mu.Lock()
    48  	defer r.mu.Unlock()
    49  	r.add(b)
    50  }
    51  
    52  func (r *BloomRing) add(b []byte) {
    53  	slot := r.slots[r.slotPosition]
    54  	if r.entryCounter > r.slotCapacity {
    55  		// Move to next slot and reset
    56  		r.slotPosition = (r.slotPosition + 1) % r.slotCount
    57  		slot = r.slots[r.slotPosition]
    58  		slot.Reset()
    59  		r.entryCounter = 0
    60  	}
    61  	r.entryCounter++
    62  	slot.Add(b)
    63  }
    64  
    65  func (r *BloomRing) Test(b []byte) bool {
    66  	if r == nil {
    67  		return false
    68  	}
    69  	r.mu.RLock()
    70  	defer r.mu.RUnlock()
    71  	test := r.test(b)
    72  	return test
    73  }
    74  
    75  func (r *BloomRing) test(b []byte) bool {
    76  	for _, s := range r.slots {
    77  		if s.Test(b) {
    78  			return true
    79  		}
    80  	}
    81  	return false
    82  }
    83  
    84  func (r *BloomRing) Check(b []byte) bool {
    85  	r.mu.Lock()
    86  	defer r.mu.Unlock()
    87  	if r.Test(b) {
    88  		return true
    89  	}
    90  	r.Add(b)
    91  	return false
    92  }