fix bugs && add rate limiter test cases

* fixed hashmap seed issue
* remove unneccesary `rand.Seed`

.vscode/settings.json

@@ -1,6 +1,6 @@
 {
     "files.eol": "\n",
     "go.toolsEnvVars": {
-        "GOOS": "linux",
+        // "GOOS": "linux",
     }
 }

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]
 			}

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()

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)
+	}
+}

utils.go → utils/utils.go

@@ -1,4 +1,4 @@
-package main
+package utils
 
 import (
 	"fmt"