From 0b4af7da4113fc95ef21e9e359d3629bbd858fd2 Mon Sep 17 00:00:00 2001
From: Hugefiver <i@iruri.moe>
Date: Wed, 10 Jul 2024 15:11:28 +0800
Subject: [PATCH] fix bugs && add rate limiter test cases

* fixed hashmap seed issue
* remove unneccesary `rand.Seed`
---
 .vscode/settings.json      |  2 +-
 main.go                    | 10 ++--
 rate.go                    |  6 ++-
 rate_test.go               | 94 ++++++++++++++++++++++++++++++++++++++
 utils.go => utils/utils.go |  2 +-
 5 files changed, 105 insertions(+), 9 deletions(-)
 create mode 100644 rate_test.go
 rename utils.go => utils/utils.go (97%)

diff --git a/.vscode/settings.json b/.vscode/settings.json
index b5d91c1..e18302b 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -1,6 +1,6 @@
 {
     "files.eol": "\n",
     "go.toolsEnvVars": {
-        "GOOS": "linux",
+        // "GOOS": "linux",
     }
 }
\ No newline at end of file
diff --git a/main.go b/main.go
index 92b2432..76da5f7 100644
--- a/main.go
+++ b/main.go
@@ -20,6 +20,7 @@ import (
 
 	"github.com/hugefiver/fakessh/conf"
 	"github.com/hugefiver/fakessh/third/ssh"
+	"github.com/hugefiver/fakessh/utils"
 
 	"go.uber.org/zap"
 )
@@ -32,7 +33,6 @@ const seedSize = 32 // 256 bits
 var seed []byte
 
 func init() {
-	rand.Seed(time.Now().UnixNano())
 	byteBuf := bytes.NewBuffer(make([]byte, 0, seedSize))
 	for i := 0; i < seedSize/8; i++ {
 		binary.Write(byteBuf, binary.BigEndian, rand.Uint64())
@@ -61,17 +61,17 @@ func main() {
 			signers = append(signers, signer)
 		}
 	} else {
-		var pairs []*KeyOption
+		var pairs []*utils.KeyOption
 		if used.Contains(conf.FlagKeyType) || args.GenKeyFile {
-			pairs = GetKeyOptionPairs(args.KeyType)
+			pairs = utils.GetKeyOptionPairs(args.KeyType)
 		} else {
-			pairs = GetKeyOptionPairs(sc.Key.KeyType)
+			pairs = utils.GetKeyOptionPairs(sc.Key.KeyType)
 		}
 
 		// only generate the first key option
 		if args.GenKeyFile {
 			// Marshal key
-			k := &KeyOption{}
+			k := &utils.KeyOption{}
 			if len(pairs) > 0 {
 				k = pairs[0]
 			}
diff --git a/rate.go b/rate.go
index 0075a65..7f2b455 100644
--- a/rate.go
+++ b/rate.go
@@ -3,6 +3,7 @@ package main
 import (
 	"hash/maphash"
 	"time"
+	"unsafe"
 
 	"github.com/cespare/xxhash/v2"
 	"github.com/hugefiver/fakessh/conf"
@@ -49,7 +50,7 @@ func NewRateLimiter(cs []*conf.RateLimitConfig) *RateLimiter {
 	rs := make([]*rate.Limiter, len(cs))
 
 	for i, c := range cs {
-		rs[i] = rate.NewLimiter(rate.Every(c.Interval.Duration()), c.Limit)
+		rs[i] = rate.NewLimiter(rate.Every(c.Interval.Duration()/time.Duration(c.Limit)), c.Limit)
 	}
 	return &RateLimiter{limiters: rs}
 }
@@ -90,7 +91,8 @@ type SSHRateLimiter struct {
 }
 
 func hashString(seed maphash.Seed, s string) uint64 {
-	h := xxhash.NewWithSeed(seedSize)
+	x := *(*uint64)(unsafe.Pointer(&seed))
+	h := xxhash.NewWithSeed(x)
 
 	_, _ = h.WriteString(s)
 	return h.Sum64()
diff --git a/rate_test.go b/rate_test.go
new file mode 100644
index 0000000..9755e1f
--- /dev/null
+++ b/rate_test.go
@@ -0,0 +1,94 @@
+package main
+
+import (
+	"testing"
+	"time"
+
+	"github.com/hugefiver/fakessh/conf"
+)
+
+func TestRateLimiter(t *testing.T) {
+	t.Parallel()
+
+	rl := NewRateLimiter([]*conf.RateLimitConfig{
+		{Limit: 3, Interval: conf.Duration(time.Second)},
+		{Limit: 5, Interval: conf.Duration(time.Second * 10)},
+	})
+
+	r1 := rl.Allow()
+	r2 := rl.Allow()
+	r3 := rl.Allow()
+
+	t.Logf("reservations: %#v, %#v, %#v", r1.reservations, r2.reservations, r3.reservations)
+	if !r1.OK() || !r2.OK() || !r3.OK() {
+		t.Errorf("r1: %v, r2: %v, r3: %v", r1.ok, r2.ok, r3.ok)
+	}
+
+	re := rl.Allow()
+	if re.OK() {
+		t.Errorf("re: %v", re)
+	}
+
+	time.Sleep(time.Second * 1)
+
+	r4 := rl.Allow()
+	r5 := rl.Allow()
+	r6 := rl.Allow()
+
+	t.Logf("reservations: %#v, %#v, %#v", r4.reservations, r5.reservations, r6.reservations)
+	if !r4.OK() || !r5.OK() || r6.OK() {
+		t.Errorf("r4: %v, r5: %v, r6: %v", r4.ok, r5.ok, r6.ok)
+	}
+}
+
+func TestSSHRateLimiter(t *testing.T) {
+	t.Parallel()
+
+	rl := NewSSHRateLimiter([]*conf.RateLimitConfig{
+		{Limit: 3, Interval: conf.Duration(time.Second)},
+		{Limit: 5, Interval: conf.Duration(time.Second * 10)},
+	}, []*conf.RateLimitConfig{
+		{Limit: 1, Interval: conf.Duration(time.Second), PerIP: true},
+	})
+
+	x1 := rl.Allow("1")
+	x2 := rl.Allow("1")
+
+	if !x1.OK() || x2.OK() {
+		t.Errorf("x1: %v, x2: %v", x1, x2)
+	}
+
+	y1 := rl.Allow("2")
+	y2 := rl.Allow("2")
+
+	if !y1.OK() || y2.OK() {
+		t.Errorf("y1: %v, y2: %v", y1, y2)
+	}
+
+	z1 := rl.Allow("3")
+	z2 := rl.Allow("3")
+
+	if !z1.OK() || z2.OK() {
+		t.Errorf("z1: %v, z2: %v", z1, z2)
+	}
+
+	time.Sleep(time.Second * 1)
+
+	x3 := rl.Allow("1")
+
+	if !x3.OK() {
+		t.Errorf("x3: %v", x3)
+	}
+
+	y3 := rl.Allow("2")
+
+	if !y3.OK() {
+		t.Errorf("y3: %v", y3)
+	}
+
+	z3 := rl.Allow("3")
+
+	if z3.OK() {
+		t.Errorf("z3: %v", z3)
+	}
+}
diff --git a/utils.go b/utils/utils.go
similarity index 97%
rename from utils.go
rename to utils/utils.go
index bf69fe8..8ed5f58 100644
--- a/utils.go
+++ b/utils/utils.go
@@ -1,4 +1,4 @@
-package main
+package utils
 
 import (
 	"fmt"