From 72885aae835cf962286d3a87f88209f18177617b Mon Sep 17 00:00:00 2001 From: Alexander Chuprov Date: Fri, 12 Jan 2024 18:09:28 +0300 Subject: [PATCH] [#62] signature: Refactor BufferPool Signed-off-by: Alexander Chuprov --- util/pool/buffer.go | 54 +++++++++++++++++++++++++++++++++++++++ util/signature/buffer.go | 40 ----------------------------- util/signature/data.go | 21 +++++++-------- util/signature/options.go | 16 ++++++------ 4 files changed, 73 insertions(+), 58 deletions(-) create mode 100644 util/pool/buffer.go delete mode 100644 util/signature/buffer.go diff --git a/util/pool/buffer.go b/util/pool/buffer.go new file mode 100644 index 0000000..e0a7185 --- /dev/null +++ b/util/pool/buffer.go @@ -0,0 +1,54 @@ +package pool + +import ( + "sync" +) + +// Buffer contains a byte slice. +type Buffer struct { + Data []byte +} + +// BufferPool manages a pool of Buffers. +type BufferPool struct { + poolSliceSize uint32 // Size for the buffer slices in the pool. + buffersPool *sync.Pool +} + +// NewBufferPool creates a BufferPool with a specified size. +func NewBufferPool(poolSliceSize uint32) BufferPool { + pool := sync.Pool{ + New: func() any { + return new(Buffer) + }, + } + return BufferPool{poolSliceSize: poolSliceSize, buffersPool: &pool} +} + +// Get retrieves a Buffer from the pool or creates a new one if necessary. +// It ensures the buffer's capacity is at least the specified size. +func (pool BufferPool) Get(size uint32) *Buffer { + result := pool.buffersPool.Get().(*Buffer) + + if cap(result.Data) < int(size) { + result.Data = make([]byte, size) + } else { + result.Data = result.Data[:size] + } + return result +} + +// Put returns a Buffer to the pool if its capacity does not exceed poolSliceSize. +func (pool BufferPool) Put(buf *Buffer) { + if cap(buf.Data) > int(pool.poolSliceSize) { + return + } + + buf.Data = buf.Data[:0] + pool.buffersPool.Put(buf) +} + +// PoolSliceSize returns the size for buffer slices in the pool. +func (pool BufferPool) PoolSliceSize() uint32 { + return uint32(pool.poolSliceSize) +} diff --git a/util/signature/buffer.go b/util/signature/buffer.go deleted file mode 100644 index 8102495..0000000 --- a/util/signature/buffer.go +++ /dev/null @@ -1,40 +0,0 @@ -package signature - -import "sync" - -const poolSliceMaxSize = 128 * 1024 - -type buffer struct { - data []byte -} - -var buffersPool = sync.Pool{ - New: func() any { - return new(buffer) - }, -} - -func tryGetNewBufferFromPool(size int) (*buffer, bool) { - if size > poolSliceMaxSize { - return &buffer{}, false - } - return newBufferFromPool(size), true -} - -func newBufferFromPool(size int) *buffer { - result := buffersPool.Get().(*buffer) - if cap(result.data) < size { - result.data = make([]byte, size) - } else { - result.data = result.data[:size] - } - return result -} - -func returnBufferToPool(buf *buffer) { - if cap(buf.data) > poolSliceMaxSize { - return - } - buf.data = buf.data[:0] - buffersPool.Put(buf) -} diff --git a/util/signature/data.go b/util/signature/data.go index 4170e6c..5e7c310 100644 --- a/util/signature/data.go +++ b/util/signature/data.go @@ -4,9 +4,14 @@ import ( "crypto/ecdsa" "git.frostfs.info/TrueCloudLab/frostfs-api-go/v2/refs" + "git.frostfs.info/TrueCloudLab/frostfs-api-go/v2/util/pool" crypto "git.frostfs.info/TrueCloudLab/frostfs-crypto" ) +const poolSliceMaxSize = 128 * 1024 + +var buffersPool = pool.NewBufferPool(poolSliceMaxSize) + type DataSource interface { ReadSignedData([]byte) ([]byte, error) SignedDataSize() int @@ -35,12 +40,10 @@ func SignDataWithHandler(key *ecdsa.PrivateKey, src DataSource, handler KeySigna opts[i](cfg) } - buffer, ok := tryGetNewBufferFromPool(src.SignedDataSize()) - if ok { - defer returnBufferToPool(buffer) - } + buffer := buffersPool.Get(uint32(src.SignedDataSize())) + defer buffersPool.Put(buffer) - data, err := src.ReadSignedData(buffer.data) + data, err := src.ReadSignedData(buffer.Data) if err != nil { return err } @@ -66,12 +69,10 @@ func VerifyDataWithSource(dataSrc DataSource, sigSrc KeySignatureSource, opts .. opts[i](cfg) } - buffer, ok := tryGetNewBufferFromPool(dataSrc.SignedDataSize()) - if ok { - defer returnBufferToPool(buffer) - } + buffer := buffersPool.Get(uint32(dataSrc.SignedDataSize())) + defer buffersPool.Put(buffer) - data, err := dataSrc.ReadSignedData(buffer.data) + data, err := dataSrc.ReadSignedData(buffer.Data) if err != nil { return err } diff --git a/util/signature/options.go b/util/signature/options.go index b9bf68d..ca2a5e9 100644 --- a/util/signature/options.go +++ b/util/signature/options.go @@ -35,10 +35,10 @@ func verify(cfg *cfg, data []byte, sig *refs.Signature) error { case refs.ECDSA_RFC6979_SHA256: return crypto.VerifyRFC6979(pub, data, sig.GetSign()) case refs.ECDSA_RFC6979_SHA256_WALLET_CONNECT: - buffer := newBufferFromPool(base64.StdEncoding.EncodedLen(len(data))) - defer returnBufferToPool(buffer) - base64.StdEncoding.Encode(buffer.data, data) - if !walletconnect.Verify(pub, buffer.data, sig.GetSign()) { + buffer := buffersPool.Get(uint32(base64.StdEncoding.EncodedLen(len(data)))) + defer buffersPool.Put(buffer) + base64.StdEncoding.Encode(buffer.Data, data) + if !walletconnect.Verify(pub, buffer.Data, sig.GetSign()) { return crypto.ErrInvalidSignature } return nil @@ -54,10 +54,10 @@ func sign(cfg *cfg, key *ecdsa.PrivateKey, data []byte) ([]byte, error) { case refs.ECDSA_RFC6979_SHA256: return crypto.SignRFC6979(key, data) case refs.ECDSA_RFC6979_SHA256_WALLET_CONNECT: - buffer := newBufferFromPool(base64.StdEncoding.EncodedLen(len(data))) - defer returnBufferToPool(buffer) - base64.StdEncoding.Encode(buffer.data, data) - return walletconnect.Sign(key, buffer.data) + buffer := buffersPool.Get(uint32(base64.StdEncoding.EncodedLen(len(data)))) + defer buffersPool.Put(buffer) + base64.StdEncoding.Encode(buffer.Data, data) + return walletconnect.Sign(key, buffer.Data) default: panic(fmt.Sprintf("unsupported scheme %s", cfg.scheme)) }