From 9ab8459c0c68013680330318054cd0ab050fae64 Mon Sep 17 00:00:00 2001
From: Matt Dainty <matt@bodgit-n-scarper.com>
Date: Mon, 1 Apr 2024 23:35:16 +0100
Subject: [PATCH] perf: Add AES key caching

---
 go.mod                   |  1 +
 go.sum                   |  2 ++
 internal/aes7z/key.go    | 40 ++++++++++++++++++++++++++++++++++++++--
 internal/aes7z/reader.go |  7 ++++++-
 4 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/go.mod b/go.mod
index b6d447e..5895b4a 100644
--- a/go.mod
+++ b/go.mod
@@ -7,6 +7,7 @@ require (
 	github.com/bodgit/plumbing v1.3.0
 	github.com/bodgit/windows v1.0.1
 	github.com/hashicorp/go-multierror v1.1.1
+	github.com/hashicorp/golang-lru/v2 v2.0.7
 	github.com/klauspost/compress v1.17.7
 	github.com/pierrec/lz4/v4 v4.1.21
 	github.com/stretchr/testify v1.9.0
diff --git a/go.sum b/go.sum
index 07f3f9d..13955a0 100644
--- a/go.sum
+++ b/go.sum
@@ -66,6 +66,8 @@ github.com/hashicorp/go-multierror v1.1.1 h1:H5DkEtf6CXdFp0N0Em5UCwQpXMWke8IA0+l
 github.com/hashicorp/go-multierror v1.1.1/go.mod h1:iw975J/qwKPdAO1clOe2L8331t/9/fmwbPZ6JB6eMoM=
 github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
 github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
+github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
+github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
 github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
 github.com/jstemmer/go-junit-report v0.0.0-20190106144839-af01ea7f8024/go.mod h1:6v2b51hI/fHJwM22ozAgKL4VKDeJcHhJFhtBdhmNjmU=
 github.com/jstemmer/go-junit-report v0.9.1/go.mod h1:Brl9GWCQeLvo8nXZwPNNblvFj/XSXhF0NWZEnDohbsk=
diff --git a/internal/aes7z/key.go b/internal/aes7z/key.go
index 2ff37e7..228739c 100644
--- a/internal/aes7z/key.go
+++ b/internal/aes7z/key.go
@@ -4,12 +4,46 @@ import (
 	"bytes"
 	"crypto/sha256"
 	"encoding/binary"
+	"encoding/hex"
+	"runtime"
 
+	lru "github.com/hashicorp/golang-lru/v2"
+	"go4.org/syncutil"
 	"golang.org/x/text/encoding/unicode"
 	"golang.org/x/text/transform"
 )
 
-func calculateKey(password string, cycles int, salt []byte) []byte {
+type cacheKey struct {
+	password string
+	cycles   int
+	salt     string // []byte isn't comparable
+}
+
+//nolint:gochecknoglobals
+var (
+	once  syncutil.Once
+	cache *lru.Cache[cacheKey, []byte]
+)
+
+func calculateKey(password string, cycles int, salt []byte) ([]byte, error) {
+	if err := once.Do(func() (err error) {
+		cache, err = lru.New[cacheKey, []byte](runtime.NumCPU())
+
+		return
+	}); err != nil {
+		return nil, err
+	}
+
+	ck := cacheKey{
+		password: password,
+		cycles:   cycles,
+		salt:     hex.EncodeToString(salt),
+	}
+
+	if key, ok := cache.Get(ck); ok {
+		return key, nil
+	}
+
 	b := bytes.NewBuffer(salt)
 
 	// Convert password to UTF-16LE
@@ -30,5 +64,7 @@ func calculateKey(password string, cycles int, salt []byte) []byte {
 		copy(key, h.Sum(nil))
 	}
 
-	return key
+	_ = cache.Add(ck, key)
+
+	return key, nil
 }
diff --git a/internal/aes7z/reader.go b/internal/aes7z/reader.go
index 4d0d30c..760db29 100644
--- a/internal/aes7z/reader.go
+++ b/internal/aes7z/reader.go
@@ -29,7 +29,12 @@ func (rc *readCloser) Close() error {
 }
 
 func (rc *readCloser) Password(p string) error {
-	block, err := aes.NewCipher(calculateKey(p, rc.cycles, rc.salt))
+	key, err := calculateKey(p, rc.cycles, rc.salt)
+	if err != nil {
+		return err
+	}
+
+	block, err := aes.NewCipher(key)
 	if err != nil {
 		return err
 	}