Rework antireplay

antireplay/cache.go

@@ -1,31 +1,16 @@
 package antireplay
 
-import (
+import "github.com/allegro/bigcache"
-	"github.com/allegro/bigcache"
 
-	"github.com/9seconds/mtg/config"
+type cache struct {
-)
+	cache *bigcache.BigCache
-
-var cache *bigcache.BigCache
-
-func Add(data []byte) {
-	cache.Set(string(data), nil) // nolint: errcheck
 }
 
-func Has(data []byte) bool {
+func (c *cache) Add(data []byte) {
-	_, err := cache.Get(string(data))
+	c.cache.Set(string(data), nil) // nolint: errcheck
+}
+
+func (c *cache) Has(data []byte) bool {
+	_, err := c.cache.Get(string(data))
 	return err == nil
 }
-
-func Init() {
-	c, err := bigcache.NewBigCache(bigcache.Config{
-		Shards:           1024,
-		LifeWindow:       config.C.AntiReplay.EvictionTime,
-		Hasher:           hasher{},
-		HardMaxCacheSize: config.C.AntiReplay.MaxSize,
-	})
-	if err != nil {
-		panic(err)
-	}
-	cache = c
-}

antireplay/init.go

@@ -0,0 +1,31 @@
+package antireplay
+
+import (
+	"sync"
+
+	"github.com/9seconds/mtg/config"
+	"github.com/allegro/bigcache"
+)
+
+var (
+	Cache    *cache
+	initOnce sync.Once
+)
+
+func Init() {
+	initOnce.Do(func() {
+		c, err := bigcache.NewBigCache(bigcache.Config{
+			Shards:           1024,
+			LifeWindow:       config.C.AntiReplay.EvictionTime,
+			Hasher:           hasher{},
+			HardMaxCacheSize: config.C.AntiReplay.MaxSize,
+		})
+		if err != nil {
+			panic(err)
+		}
+
+		Cache = &cache{
+			cache: c,
+		}
+	})
+}

obfuscated2/client_protocol.go

@@ -80,10 +80,10 @@ func (c *ClientProtocol) Handshake(socket conntypes.StreamReadWriteCloser) (conn
 	}
 
 	antiReplayKey := decryptedFrame.Unique()
-	if antireplay.Has(antiReplayKey) {
+	if antireplay.Cache.Has(antiReplayKey) {
 		return nil, errors.New("Replay attack is detected")
 	}
-	antireplay.Add(antiReplayKey)
+	antireplay.Cache.Add(antiReplayKey)
 
 	return wrappers.NewObfuscated2(socket, encryptor, decryptor), nil
 }