119 lines
2.9 KiB
Go
119 lines
2.9 KiB
Go
package engine
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"git.difuse.io/Difuse/Mellaris/analyzer"
|
|
"git.difuse.io/Difuse/Mellaris/ruleset"
|
|
|
|
"github.com/bwmarrin/snowflake"
|
|
)
|
|
|
|
type legacyUDPStreamValue struct {
|
|
Tuple udpTupleKey
|
|
}
|
|
|
|
type emptyRuleset struct{}
|
|
|
|
func (emptyRuleset) Analyzers(ruleset.StreamInfo) []analyzer.Analyzer { return nil }
|
|
func (emptyRuleset) Match(ruleset.StreamInfo) ruleset.MatchResult {
|
|
return ruleset.MatchResult{Action: ruleset.ActionMaybe}
|
|
}
|
|
|
|
func benchmarkUDPManager(b *testing.B, churn bool) {
|
|
node, err := snowflake.NewNode(0)
|
|
if err != nil {
|
|
b.Fatalf("create node: %v", err)
|
|
}
|
|
factory := &udpStreamFactory{WorkerID: 0, Logger: noopTestLogger{}, Node: node, Ruleset: emptyRuleset{}}
|
|
mgr, err := newUDPStreamManager(factory, 200000, &statsCounters{})
|
|
if err != nil {
|
|
b.Fatalf("new manager: %v", err)
|
|
}
|
|
|
|
const flowCount = 20000
|
|
tuples := make([]udpTupleKey, flowCount)
|
|
payloads := make([][]byte, flowCount)
|
|
for i := 0; i < flowCount; i++ {
|
|
a := byte(i >> 8)
|
|
c := byte(i)
|
|
var t udpTupleKey
|
|
t.AIP = [16]byte{10, a, 0, c}
|
|
t.ALen = 4
|
|
t.BIP = [16]byte{172, 16, a, c}
|
|
t.BLen = 4
|
|
t.APort = 1024 + uint16(i%20000)
|
|
t.BPort = 20000 + uint16((i*7)%20000)
|
|
tuples[i] = t
|
|
payloads[i] = []byte{0x01, 0x00, 0x00, 0x00}
|
|
}
|
|
|
|
ctx := &udpContext{Verdict: udpVerdictAccept}
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i++ {
|
|
idx := i % flowCount
|
|
streamID := uint32(idx + 1)
|
|
if churn {
|
|
streamID = uint32((i % flowCount) + 1 + ((i / flowCount) * flowCount))
|
|
}
|
|
ctx.Verdict = udpVerdictAccept
|
|
ctx.Packet = nil
|
|
mgr.MatchWithContext(streamID, tuples[idx], false, payloads[idx], ctx)
|
|
}
|
|
}
|
|
|
|
func BenchmarkUDPManagerMatchStableStreamID(b *testing.B) {
|
|
benchmarkUDPManager(b, false)
|
|
}
|
|
|
|
func BenchmarkUDPManagerMatchStreamIDChurn(b *testing.B) {
|
|
benchmarkUDPManager(b, true)
|
|
}
|
|
|
|
func BenchmarkLegacyUDPFallbackScanChurn(b *testing.B) {
|
|
const flowCount = 5000
|
|
tuples := make([]udpTupleKey, flowCount)
|
|
for i := 0; i < flowCount; i++ {
|
|
a := byte(i >> 8)
|
|
c := byte(i)
|
|
var t udpTupleKey
|
|
t.AIP = [16]byte{10, a, 0, c}
|
|
t.ALen = 4
|
|
t.BIP = [16]byte{172, 16, a, c}
|
|
t.BLen = 4
|
|
t.APort = 1024 + uint16(i%20000)
|
|
t.BPort = 20000 + uint16((i*7)%20000)
|
|
tuples[i] = t
|
|
}
|
|
|
|
streams := make(map[uint32]*legacyUDPStreamValue, flowCount)
|
|
keys := make([]uint32, 0, flowCount)
|
|
for i := 0; i < flowCount; i++ {
|
|
streamID := uint32(i + 1)
|
|
streams[streamID] = &legacyUDPStreamValue{Tuple: tuples[i]}
|
|
keys = append(keys, streamID)
|
|
}
|
|
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i++ {
|
|
idx := i % flowCount
|
|
streamID := uint32((i % flowCount) + 1 + ((i / flowCount) * flowCount))
|
|
if _, ok := streams[streamID]; ok {
|
|
continue
|
|
}
|
|
tuple := tuples[idx]
|
|
revTuple := reverseTuple(tuple)
|
|
for _, k := range keys {
|
|
v, ok := streams[k]
|
|
if !ok || v == nil {
|
|
continue
|
|
}
|
|
if v.Tuple == tuple || v.Tuple == revTuple {
|
|
delete(streams, k)
|
|
streams[streamID] = v
|
|
break
|
|
}
|
|
}
|
|
}
|
|
}
|