Skip to content

Commit

Permalink
perf: Add AES key caching (#189)
Browse files Browse the repository at this point in the history
* test: Add AES password benchmark

* perf: Add AES key caching

* style: Lint fixes
  • Loading branch information
bodgit authored Apr 1, 2024
1 parent 31f77ef commit 3d794c2
Show file tree
Hide file tree
Showing 8 changed files with 77 additions and 23 deletions.
1 change: 1 addition & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down
41 changes: 39 additions & 2 deletions internal/aes7z/key.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
}
7 changes: 6 additions & 1 deletion internal/aes7z/reader.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down
1 change: 1 addition & 0 deletions internal/zstd/reader.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}

Expand Down
4 changes: 3 additions & 1 deletion reader.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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 {
Expand Down
44 changes: 25 additions & 19 deletions reader_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
}
Expand All @@ -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) {
Expand All @@ -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)
}
Binary file added testdata/aes7z.7z
Binary file not shown.

0 comments on commit 3d794c2

Please sign in to comment.