implements BloomRing and saltfilter

go.mod

@@ -4,6 +4,7 @@ go 1.12
 
 require (
 	github.com/aead/chacha20 v0.0.0-20180709150244-8b13a72661da
+	github.com/riobard/go-bloom v0.0.0-20200213042214-218e1707c495
 	golang.org/x/crypto v0.0.0-20190426145343-a29dc8fdc734
 )
 

go.sum

@@ -8,3 +8,7 @@ github.com/golang/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:5JyrLPvD/ZdaY
 github.com/golang/sys v0.0.0-20190412213103-97732733099d h1:blRtD+FQOxZ6P7jigy+HS0R8zyGOMOv8TET4wCpzVwM=
 github.com/golang/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 github.com/golang/text v0.3.0/go.mod h1:GUiq9pdJKRKKAZXiVgWFEvocYuREvC14NhI4OPgEjeE=
+github.com/riobard/go-bloom v0.0.0-20170218180955-2b113c64a69b h1:H9yjH/g5w8MOPjQR2zMSP/Md1kKtj/33fIht9ChC2OU=
+github.com/riobard/go-bloom v0.0.0-20170218180955-2b113c64a69b/go.mod h1:HgjTstvQsPGkxUsCd2KWxErBblirPizecHcpD3ffK+s=
+github.com/riobard/go-bloom v0.0.0-20200213042214-218e1707c495 h1:p7xbxYTzzfXghR1kpsJDeoVVRRWAotKc8u7FP/N48rU=
+github.com/riobard/go-bloom v0.0.0-20200213042214-218e1707c495/go.mod h1:HgjTstvQsPGkxUsCd2KWxErBblirPizecHcpD3ffK+s=

internal/bloomring.go

@@ -0,0 +1,67 @@
+package internal
+
+import (
+	"hash/fnv"
+	"sync"
+
+	"github.com/riobard/go-bloom"
+)
+
+// simply use Double FNV here as our Bloom Filter hash
+func doubleFNV(b []byte) (uint64, uint64) {
+	hx := fnv.New64()
+	hx.Write(b)
+	x := hx.Sum64()
+	hy := fnv.New64a()
+	hy.Write(b)
+	y := hy.Sum64()
+	return x, y
+}
+
+type BloomRing struct {
+	slotCapacity int
+	slotPosition int
+	slotCount    int
+	entryCounter int
+	slots        []bloom.Filter
+	mutex        sync.RWMutex
+}
+
+func NewBloomRing(slot, capacity int, falsePositiveRate float64) *BloomRing {
+	// Calculate entries for each slot
+	r := &BloomRing{
+		slotCapacity: capacity / slot,
+		slotCount:    slot,
+		slots:        make([]bloom.Filter, slot),
+	}
+	for i := 0; i < slot; i++ {
+		r.slots[i] = bloom.New(r.slotCapacity, falsePositiveRate, doubleFNV)
+	}
+	return r
+}
+
+func (r *BloomRing) Add(b []byte) {
+	r.mutex.Lock()
+	defer r.mutex.Unlock()
+	slot := r.slots[r.slotPosition]
+	if r.entryCounter > r.slotCapacity {
+		// Move to next slot and reset
+		r.slotPosition = (r.slotPosition + 1) % r.slotCount
+		slot = r.slots[r.slotPosition]
+		slot.Reset()
+		r.entryCounter = 0
+	}
+	r.entryCounter++
+	slot.Add(b)
+}
+
+func (r *BloomRing) Test(b []byte) bool {
+	r.mutex.RLock()
+	defer r.mutex.RUnlock()
+	for _, s := range r.slots {
+		if s.Test(b) {
+			return true
+		}
+	}
+	return false
+}

internal/bloomring_test.go

@@ -0,0 +1,68 @@
+package internal_test
+
+import (
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/shadowsocks/go-shadowsocks2/internal"
+)
+
+var (
+	bloomRingInstance *internal.BloomRing
+)
+
+func TestMain(m *testing.M) {
+	bloomRingInstance = internal.NewBloomRing(internal.DefaultSFSlot, int(internal.DefaultSFCapacity),
+		internal.DefaultSFFPR)
+	os.Exit(m.Run())
+}
+
+func TestBloomRing_Add(t *testing.T) {
+	defer func() {
+		if any := recover(); any != nil {
+			t.Fatalf("Should not got panic while adding item: %v", any)
+		}
+	}()
+	bloomRingInstance.Add(make([]byte, 16))
+}
+
+func TestBloomRing_Test(t *testing.T) {
+	buf := []byte("shadowsocks")
+	bloomRingInstance.Add(buf)
+	if !bloomRingInstance.Test(buf) {
+		t.Fatal("Test on filter missing")
+	}
+}
+
+func BenchmarkBloomRing(b *testing.B) {
+	// Generate test samples with different length
+	samples := make([][]byte, internal.DefaultSFCapacity-internal.DefaultSFSlot)
+	var checkPoints [][]byte
+	for i := 0; i < len(samples); i++ {
+		samples[i] = []byte(fmt.Sprint(i))
+		if i%1000 == 0 {
+			checkPoints = append(checkPoints, samples[i])
+		}
+	}
+	b.Logf("Generated %d samples and %d check points", len(samples), len(checkPoints))
+	for i := 1; i < 16; i++ {
+		b.Run(fmt.Sprintf("Slot%d", i), benchmarkBloomRing(samples, checkPoints, i))
+	}
+}
+
+func benchmarkBloomRing(samples, checkPoints [][]byte, slot int) func(*testing.B) {
+	filter := internal.NewBloomRing(slot, int(internal.DefaultSFCapacity), internal.DefaultSFFPR)
+	for _, sample := range samples {
+		filter.Add(sample)
+	}
+	return func(b *testing.B) {
+		b.ResetTimer()
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			for _, cp := range checkPoints {
+				filter.Test(cp)
+			}
+		}
+	}
+}

internal/saltfilter.go

@@ -0,0 +1,80 @@
+package internal
+
+import (
+	"fmt"
+	"os"
+	"strconv"
+)
+
+// Those suggest value are all set according to
+// https://github.com/shadowsocks/shadowsocks-org/issues/44#issuecomment-281021054
+// Due to this package contains various internal implementation so const named with DefaultBR prefix
+const (
+	DefaultSFCapacity = 1e6
+	// FalsePositiveRate
+	DefaultSFFPR  = 1e-6
+	DefaultSFSlot = 10
+)
+
+const EnvironmentPrefix = "SHADOWSOCKS_"
+
+// A shared instance used for checking salt repeat
+var saltfilter *BloomRing
+
+func init() {
+	var (
+		finalCapacity = DefaultSFCapacity
+		finalFPR      = DefaultSFFPR
+		finalSlot     = float64(DefaultSFSlot)
+	)
+	for _, opt := range []struct {
+		ENVName string
+		Target  *float64
+	}{
+		{
+			ENVName: "CAPACITY",
+			Target:  &finalCapacity,
+		},
+		{
+			ENVName: "FPR",
+			Target:  &finalFPR,
+		},
+		{
+			ENVName: "SLOT",
+			Target:  &finalSlot,
+		},
+	} {
+		envKey := EnvironmentPrefix + "SF_" + opt.ENVName
+		env := os.Getenv(envKey)
+		if env != "" {
+			p, err := strconv.ParseFloat(env, 64)
+			if err != nil {
+				panic(fmt.Sprintf("Invalid envrionment `%s` setting in saltfilter: %s", envKey, env))
+			}
+			*opt.Target = p
+		}
+	}
+	// Support disable saltfilter by given a negative capacity
+	if finalCapacity <= 0 {
+		return
+	}
+	saltfilter = NewBloomRing(int(finalSlot), int(finalCapacity), finalFPR)
+}
+
+// TestSalt returns true if salt is repeated
+func TestSalt(b []byte) bool {
+	// If nil means feature disabled, return false to bypass salt repeat detection
+	if saltfilter == nil {
+		return false
+	}
+	return saltfilter.Test(b)
+}
+
+// AddSalt salt to filter
+func AddSalt(b []byte) {
+	// If nil means feature disabled
+	if saltfilter == nil {
+		return
+	}
+	saltfilter.Add(b)
+}