perf: Add AES key caching (#189)

bodgit · web-flow · commit 3d794c26c683 · 2024-04-02T00:11:14.000+01:00
* test: Add AES password benchmark

* perf: Add AES key caching

* style: Lint fixes
diff --git 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
@@ -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
@@ -4,12 +4,47 @@ import (
 	"bytes"
 	"crypto/sha256"
 	"encoding/binary"
+	"encoding/hex"
 
+	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
+}
+
+const cacheSize = 10
+
+//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](cacheSize)
+
+		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 +65,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
@@ -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
 	}
diff --git a/internal/zstd/reader.go b/internal/zstd/reader.go
@@ -52,6 +52,7 @@ func NewReader(_ []byte, _ uint64, readers []io.ReadCloser) (io.ReadCloser, erro
 		if r, err = zstd.NewReader(readers[0]); err != nil {
 			return nil, err
 		}
+
 		runtime.SetFinalizer(r, (*zstd.Decoder).Close)
 	}
 
diff --git a/reader.go b/reader.go
@@ -544,7 +544,7 @@ func toValidName(name string) string {
 	return p
 }
 
-//nolint:cyclop
+//nolint:cyclop,funlen
 func (z *Reader) initFileList() {
 	z.fileListOnce.Do(func() {
 		files := make(map[string]int)
@@ -583,12 +583,14 @@ func (z *Reader) initFileList() {
 				isDir: isDir,
 			}
 			z.fileList = append(z.fileList, entry)
+
 			if isDir {
 				knownDirs[name] = idx
 			} else {
 				files[name] = idx
 			}
 		}
+
 		for dir := range dirs {
 			if _, ok := knownDirs[dir]; !ok {
 				if idx, ok := files[dir]; ok {
diff --git a/reader_test.go b/reader_test.go
@@ -188,6 +188,7 @@ func TestOpenReader(t *testing.T) {
 
 		t.Run(table.name, func(t *testing.T) {
 			t.Parallel()
+
 			r, err := sevenzip.OpenReader(filepath.Join("testdata", table.file))
 			if err != nil {
 				assert.ErrorIs(t, err, table.err)
@@ -243,6 +244,7 @@ func TestOpenReaderWithPassword(t *testing.T) {
 
 		t.Run(table.name, func(t *testing.T) {
 			t.Parallel()
+
 			r, err := sevenzip.OpenReaderWithPassword(filepath.Join("testdata", table.file), table.password)
 			if err != nil {
 				t.Fatal(err)
@@ -362,13 +364,13 @@ func benchmarkArchiveNaiveParallel(b *testing.B, file string, workers int) {
 	}
 }
 
-func benchmarkArchive(b *testing.B, file string, optimised bool) {
+func benchmarkArchive(b *testing.B, file, password string, optimised bool) {
 	b.Helper()
 
 	h := crc32.NewIEEE()
 
 	for n := 0; n < b.N; n++ {
-		r, err := sevenzip.OpenReader(filepath.Join("testdata", file))
+		r, err := sevenzip.OpenReaderWithPassword(filepath.Join("testdata", file), password)
 		if err != nil {
 			b.Fatal(err)
 		}
@@ -382,56 +384,60 @@ func benchmarkArchive(b *testing.B, file string, optimised bool) {
 	}
 }
 
+func BenchmarkAES7z(b *testing.B) {
+	benchmarkArchive(b, "aes7z.7z", "password", true)
+}
+
 func BenchmarkBzip2(b *testing.B) {
-	benchmarkArchive(b, "bzip2.7z", true)
+	benchmarkArchive(b, "bzip2.7z", "", true)
 }
 
 func BenchmarkCopy(b *testing.B) {
-	benchmarkArchive(b, "copy.7z", true)
+	benchmarkArchive(b, "copy.7z", "", true)
 }
 
 func BenchmarkDeflate(b *testing.B) {
-	benchmarkArchive(b, "deflate.7z", true)
+	benchmarkArchive(b, "deflate.7z", "", true)
 }
 
 func BenchmarkDelta(b *testing.B) {
-	benchmarkArchive(b, "delta.7z", true)
+	benchmarkArchive(b, "delta.7z", "", true)
 }
 
 func BenchmarkLZMA(b *testing.B) {
-	benchmarkArchive(b, "lzma.7z", true)
+	benchmarkArchive(b, "lzma.7z", "", true)
 }
 
 func BenchmarkLZMA2(b *testing.B) {
-	benchmarkArchive(b, "lzma2.7z", true)
+	benchmarkArchive(b, "lzma2.7z", "", true)
 }
 
 func BenchmarkBCJ2(b *testing.B) {
-	benchmarkArchive(b, "bcj2.7z", true)
+	benchmarkArchive(b, "bcj2.7z", "", true)
 }
 
 func BenchmarkComplex(b *testing.B) {
-	benchmarkArchive(b, "lzma1900.7z", true)
+	benchmarkArchive(b, "lzma1900.7z", "", true)
 }
 
 func BenchmarkLZ4(b *testing.B) {
-	benchmarkArchive(b, "lz4.7z", true)
+	benchmarkArchive(b, "lz4.7z", "", true)
 }
 
 func BenchmarkBrotli(b *testing.B) {
-	benchmarkArchive(b, "brotli.7z", true)
+	benchmarkArchive(b, "brotli.7z", "", true)
 }
 
 func BenchmarkZstandard(b *testing.B) {
-	benchmarkArchive(b, "zstd.7z", true)
+	benchmarkArchive(b, "zstd.7z", "", true)
 }
 
 func BenchmarkNaiveReader(b *testing.B) {
-	benchmarkArchive(b, "lzma1900.7z", false)
+	benchmarkArchive(b, "lzma1900.7z", "", false)
 }
 
 func BenchmarkOptimisedReader(b *testing.B) {
-	benchmarkArchive(b, "lzma1900.7z", true)
+	benchmarkArchive(b, "lzma1900.7z", "", true)
 }
 
 func BenchmarkNaiveParallelReader(b *testing.B) {
@@ -447,17 +453,17 @@ func BenchmarkParallelReader(b *testing.B) {
 }
 
 func BenchmarkBCJ(b *testing.B) {
-	benchmarkArchive(b, "bcj.7z", true)
+	benchmarkArchive(b, "bcj.7z", "", true)
 }
 
 func BenchmarkPPC(b *testing.B) {
-	benchmarkArchive(b, "ppc.7z", true)
+	benchmarkArchive(b, "ppc.7z", "", true)
 }
 
 func BenchmarkARM(b *testing.B) {
-	benchmarkArchive(b, "arm.7z", true)
+	benchmarkArchive(b, "arm.7z", "", true)
 }
 
 func BenchmarkSPARC(b *testing.B) {
-	benchmarkArchive(b, "sparc.7z", true)
+	benchmarkArchive(b, "sparc.7z", "", true)
 }
diff --git a/testdata/aes7z.7z b/testdata/aes7z.7z

Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,12 @@ func (rc *readCloser) Close() error {`
`29`	`29`	`}`
`30`	`30`
`31`	`31`	`func (rc *readCloser) Password(p string) error {`
`32`		`- block, err := aes.NewCipher(calculateKey(p, rc.cycles, rc.salt))`
	`32`	`+ key, err := calculateKey(p, rc.cycles, rc.salt)`
	`33`	`+ if err != nil {`
	`34`	`+ return err`
	`35`	`+ }`
	`36`	`+`
	`37`	`+ block, err := aes.NewCipher(key)`
`33`	`38`	`if err != nil {`
`34`	`39`	`return err`
`35`	`40`	`}`
Original file line number	Diff line number	Diff line change
`@@ -52,6 +52,7 @@ func NewReader(_ []byte, _ uint64, readers []io.ReadCloser) (io.ReadCloser, erro`
`52`	`52`	`if r, err = zstd.NewReader(readers[0]); err != nil {`
`53`	`53`	`return nil, err`
`54`	`54`	`}`
	`55`	`+`
`55`	`56`	`runtime.SetFinalizer(r, (*zstd.Decoder).Close)`
`56`	`57`	`}`
`57`	`58`
Original file line number	Diff line number	Diff line change
`@@ -544,7 +544,7 @@ func toValidName(name string) string {`
`544`	`544`	`return p`
`545`	`545`	`}`
`546`	`546`
`547`		`-//nolint:cyclop`
	`547`	`+//nolint:cyclop,funlen`
`548`	`548`	`func (z *Reader) initFileList() {`
`549`	`549`	`z.fileListOnce.Do(func() {`
`550`	`550`	`files := make(map[string]int)`
`@@ -583,12 +583,14 @@ func (z *Reader) initFileList() {`
`583`	`583`	`isDir: isDir,`
`584`	`584`	`}`
`585`	`585`	`z.fileList = append(z.fileList, entry)`
	`586`	`+`
`586`	`587`	`if isDir {`
`587`	`588`	`knownDirs[name] = idx`
`588`	`589`	`} else {`
`589`	`590`	`files[name] = idx`
`590`	`591`	`}`
`591`	`592`	`}`
	`593`	`+`
`592`	`594`	`for dir := range dirs {`
`593`	`595`	`if _, ok := knownDirs[dir]; !ok {`
`594`	`596`	`if idx, ok := files[dir]; ok {`
Original file line number	Diff line number	Diff line change
`@@ -188,6 +188,7 @@ func TestOpenReader(t *testing.T) {`
`188`	`188`
`189`	`189`	`t.Run(table.name, func(t *testing.T) {`
`190`	`190`	`t.Parallel()`
	`191`	`+`
`191`	`192`	`r, err := sevenzip.OpenReader(filepath.Join("testdata", table.file))`
`192`	`193`	`if err != nil {`
`193`	`194`	`assert.ErrorIs(t, err, table.err)`
`@@ -243,6 +244,7 @@ func TestOpenReaderWithPassword(t *testing.T) {`
`243`	`244`
`244`	`245`	`t.Run(table.name, func(t *testing.T) {`
`245`	`246`	`t.Parallel()`
	`247`	`+`
`246`	`248`	`r, err := sevenzip.OpenReaderWithPassword(filepath.Join("testdata", table.file), table.password)`
`247`	`249`	`if err != nil {`
`248`	`250`	`t.Fatal(err)`
`@@ -362,13 +364,13 @@ func benchmarkArchiveNaiveParallel(b *testing.B, file string, workers int) {`
`362`	`364`	`}`
`363`	`365`	`}`
`364`	`366`
`365`		`-func benchmarkArchive(b *testing.B, file string, optimised bool) {`
	`367`	`+func benchmarkArchive(b *testing.B, file, password string, optimised bool) {`
`366`	`368`	`b.Helper()`
`367`	`369`
`368`	`370`	`h := crc32.NewIEEE()`
`369`	`371`
`370`	`372`	`for n := 0; n < b.N; n++ {`
`371`		`- r, err := sevenzip.OpenReader(filepath.Join("testdata", file))`
	`373`	`+ r, err := sevenzip.OpenReaderWithPassword(filepath.Join("testdata", file), password)`
`372`	`374`	`if err != nil {`
`373`	`375`	`b.Fatal(err)`
`374`	`376`	`}`
`@@ -382,56 +384,60 @@ func benchmarkArchive(b *testing.B, file string, optimised bool) {`
`382`	`384`	`}`
`383`	`385`	`}`
`384`	`386`
	`387`	`+func BenchmarkAES7z(b *testing.B) {`
	`388`	`+ benchmarkArchive(b, "aes7z.7z", "password", true)`
	`389`	`+}`
	`390`	`+`
`385`	`391`	`func BenchmarkBzip2(b *testing.B) {`
`386`		`- benchmarkArchive(b, "bzip2.7z", true)`
	`392`	`+ benchmarkArchive(b, "bzip2.7z", "", true)`
`387`	`393`	`}`
`388`	`394`
`389`	`395`	`func BenchmarkCopy(b *testing.B) {`
`390`		`- benchmarkArchive(b, "copy.7z", true)`
	`396`	`+ benchmarkArchive(b, "copy.7z", "", true)`
`391`	`397`	`}`
`392`	`398`
`393`	`399`	`func BenchmarkDeflate(b *testing.B) {`
`394`		`- benchmarkArchive(b, "deflate.7z", true)`
	`400`	`+ benchmarkArchive(b, "deflate.7z", "", true)`
`395`	`401`	`}`
`396`	`402`
`397`	`403`	`func BenchmarkDelta(b *testing.B) {`
`398`		`- benchmarkArchive(b, "delta.7z", true)`
	`404`	`+ benchmarkArchive(b, "delta.7z", "", true)`
`399`	`405`	`}`
`400`	`406`
`401`	`407`	`func BenchmarkLZMA(b *testing.B) {`
`402`		`- benchmarkArchive(b, "lzma.7z", true)`
	`408`	`+ benchmarkArchive(b, "lzma.7z", "", true)`
`403`	`409`	`}`
`404`	`410`
`405`	`411`	`func BenchmarkLZMA2(b *testing.B) {`
`406`		`- benchmarkArchive(b, "lzma2.7z", true)`
	`412`	`+ benchmarkArchive(b, "lzma2.7z", "", true)`
`407`	`413`	`}`
`408`	`414`
`409`	`415`	`func BenchmarkBCJ2(b *testing.B) {`
`410`		`- benchmarkArchive(b, "bcj2.7z", true)`
	`416`	`+ benchmarkArchive(b, "bcj2.7z", "", true)`
`411`	`417`	`}`
`412`	`418`
`413`	`419`	`func BenchmarkComplex(b *testing.B) {`
`414`		`- benchmarkArchive(b, "lzma1900.7z", true)`
	`420`	`+ benchmarkArchive(b, "lzma1900.7z", "", true)`
`415`	`421`	`}`
`416`	`422`
`417`	`423`	`func BenchmarkLZ4(b *testing.B) {`
`418`		`- benchmarkArchive(b, "lz4.7z", true)`
	`424`	`+ benchmarkArchive(b, "lz4.7z", "", true)`
`419`	`425`	`}`
`420`	`426`
`421`	`427`	`func BenchmarkBrotli(b *testing.B) {`
`422`		`- benchmarkArchive(b, "brotli.7z", true)`
	`428`	`+ benchmarkArchive(b, "brotli.7z", "", true)`
`423`	`429`	`}`
`424`	`430`
`425`	`431`	`func BenchmarkZstandard(b *testing.B) {`
`426`		`- benchmarkArchive(b, "zstd.7z", true)`
	`432`	`+ benchmarkArchive(b, "zstd.7z", "", true)`
`427`	`433`	`}`
`428`	`434`
`429`	`435`	`func BenchmarkNaiveReader(b *testing.B) {`
`430`		`- benchmarkArchive(b, "lzma1900.7z", false)`
	`436`	`+ benchmarkArchive(b, "lzma1900.7z", "", false)`
`431`	`437`	`}`
`432`	`438`
`433`	`439`	`func BenchmarkOptimisedReader(b *testing.B) {`
`434`		`- benchmarkArchive(b, "lzma1900.7z", true)`
	`440`	`+ benchmarkArchive(b, "lzma1900.7z", "", true)`
`435`	`441`	`}`
`436`	`442`
`437`	`443`	`func BenchmarkNaiveParallelReader(b *testing.B) {`
`@@ -447,17 +453,17 @@ func BenchmarkParallelReader(b *testing.B) {`
`447`	`453`	`}`
`448`	`454`
`449`	`455`	`func BenchmarkBCJ(b *testing.B) {`
`450`		`- benchmarkArchive(b, "bcj.7z", true)`
	`456`	`+ benchmarkArchive(b, "bcj.7z", "", true)`
`451`	`457`	`}`
`452`	`458`
`453`	`459`	`func BenchmarkPPC(b *testing.B) {`
`454`		`- benchmarkArchive(b, "ppc.7z", true)`
	`460`	`+ benchmarkArchive(b, "ppc.7z", "", true)`
`455`	`461`	`}`
`456`	`462`
`457`	`463`	`func BenchmarkARM(b *testing.B) {`
`458`		`- benchmarkArchive(b, "arm.7z", true)`
	`464`	`+ benchmarkArchive(b, "arm.7z", "", true)`
`459`	`465`	`}`
`460`	`466`
`461`	`467`	`func BenchmarkSPARC(b *testing.B) {`
`462`		`- benchmarkArchive(b, "sparc.7z", true)`
	`468`	`+ benchmarkArchive(b, "sparc.7z", "", true)`
`463`	`469`	`}`