[#62] signature: Refactor BufferPool

Signed-off-by: Alexander Chuprov <a.chuprov@yadro.com>
Signed-off-by: Evgenii Stratonikov <e.stratonikov@yadro.com>

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)
+}

util/signature/buffer.go

@@ -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)
-}

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())
+	buffer := buffersPool.Get(uint32(src.SignedDataSize()))
-	if ok {
+	defer buffersPool.Put(buffer)
-		defer returnBufferToPool(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())
+	buffer := buffersPool.Get(uint32(dataSrc.SignedDataSize()))
-	if ok {
+	defer buffersPool.Put(buffer)
-		defer returnBufferToPool(buffer)
-	}
 
-	data, err := dataSrc.ReadSignedData(buffer.data)
+	data, err := dataSrc.ReadSignedData(buffer.Data)
 	if err != nil {
 		return err
 	}

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)))
+		buffer := buffersPool.Get(uint32(base64.StdEncoding.EncodedLen(len(data))))
-		defer returnBufferToPool(buffer)
+		defer buffersPool.Put(buffer)
-		base64.StdEncoding.Encode(buffer.data, data)
+		base64.StdEncoding.Encode(buffer.Data, data)
-		if !walletconnect.Verify(pub, buffer.data, sig.GetSign()) {
+		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)))
+		buffer := buffersPool.Get(uint32(base64.StdEncoding.EncodedLen(len(data))))
-		defer returnBufferToPool(buffer)
+		defer buffersPool.Put(buffer)
-		base64.StdEncoding.Encode(buffer.data, data)
+		base64.StdEncoding.Encode(buffer.Data, data)
-		return walletconnect.Sign(key, buffer.data)
+		return walletconnect.Sign(key, buffer.Data)
 	default:
 		panic(fmt.Sprintf("unsupported scheme %s", cfg.scheme))
 	}