From d109594d751502af3a6595f6aa6959479e766bed Mon Sep 17 00:00:00 2001
From: Leonard Lyubich <leonard@nspcc.ru>
Date: Wed, 2 Mar 2022 14:41:37 +0300
Subject: [PATCH] [#157] util/signature: Revive removed features

Signed-off-by: Leonard Lyubich <leonard@nspcc.ru>
---
 util/signature/data.go    | 59 +++++++++++++++++++++++++++++++--------
 util/signature/options.go | 12 ++++----
 2 files changed, 55 insertions(+), 16 deletions(-)

diff --git a/util/signature/data.go b/util/signature/data.go
index 1de11a8..e18960f 100644
--- a/util/signature/data.go
+++ b/util/signature/data.go
@@ -44,14 +44,42 @@ var (
 	ErrInvalidSignature = errors.New("invalid signature")
 )
 
-func SignData(key *ecdsa.PrivateKey, src DataSource, opts ...SignOption) (*signature.Signature, error) {
+func SignData(key *ecdsa.PrivateKey, src DataSource, opts ...SignOption) (res *signature.Signature, err error) {
+	err = signDataWithHandler(key, src, func(key, sig []byte, scheme signature.Scheme) {
+		res = new(signature.Signature)
+		res.SetKey(key)
+		res.SetSign(sig)
+		res.SetScheme(scheme)
+	}, opts...)
+
+	return
+}
+
+func VerifyData(dataSrc DataSource, sig *signature.Signature, opts ...SignOption) error {
+	return verifyDataWithSource(dataSrc, func() ([]byte, []byte, signature.Scheme) {
+		return sig.Key(), sig.Sign(), sig.Scheme()
+	}, opts...)
+}
+
+func SignDataWithHandler(key *ecdsa.PrivateKey, src DataSource, handler func(key, sig []byte)) error {
+	return signDataWithHandler(key, src, func(key, sig []byte, scheme signature.Scheme) {
+		handler(key, sig)
+	})
+}
+
+func signDataWithHandler(
+	key *ecdsa.PrivateKey,
+	src DataSource,
+	handler func(key, sig []byte, scheme signature.Scheme),
+	opts ...SignOption,
+) error {
 	if key == nil {
-		return nil, ErrEmptyPrivateKey
+		return ErrEmptyPrivateKey
 	}
 
 	data, err := dataForSignature(src)
 	if err != nil {
-		return nil, err
+		return err
 	}
 	defer bytesPool.Put(&data)
 
@@ -59,17 +87,26 @@ func SignData(key *ecdsa.PrivateKey, src DataSource, opts ...SignOption) (*signa
 
 	sigData, err := sign(cfg.scheme, key, data)
 	if err != nil {
-		return nil, err
+		return err
 	}
 
-	sig := signature.New()
-	sig.SetKey((*keys.PublicKey)(&key.PublicKey).Bytes())
-	sig.SetSign(sigData)
-	sig.SetScheme(cfg.scheme)
-	return sig, nil
+	handler((*keys.PublicKey)(&key.PublicKey).Bytes(), sigData, cfg.scheme)
+
+	return nil
 }
 
-func VerifyData(dataSrc DataSource, sig *signature.Signature, opts ...SignOption) error {
+func VerifyDataWithSource(dataSrc DataSource, sigSrc func() (key, sig []byte)) error {
+	return verifyDataWithSource(dataSrc, func() ([]byte, []byte, signature.Scheme) {
+		key, sign := sigSrc()
+		return key, sign, signature.ECDSAWithSHA512
+	})
+}
+
+func verifyDataWithSource(
+	dataSrc DataSource,
+	sigSrc func() (key, sig []byte, scheme signature.Scheme),
+	opts ...SignOption,
+) error {
 	data, err := dataForSignature(dataSrc)
 	if err != nil {
 		return err
@@ -78,5 +115,5 @@ func VerifyData(dataSrc DataSource, sig *signature.Signature, opts ...SignOption
 
 	cfg := getConfig(opts...)
 
-	return verify(cfg, data, sig)
+	return verify(cfg, data, sigSrc)
 }
diff --git a/util/signature/options.go b/util/signature/options.go
index f768351..a825b9a 100644
--- a/util/signature/options.go
+++ b/util/signature/options.go
@@ -49,27 +49,29 @@ func sign(scheme signature.Scheme, key *ecdsa.PrivateKey, msg []byte) ([]byte, e
 	}
 }
 
-func verify(cfg *cfg, msg []byte, sig *signature.Signature) error {
-	pub, err := keys.NewPublicKeyFromBytes(sig.Key(), elliptic.P256())
+func verify(cfg *cfg, msg []byte, f func() (key, sign []byte, scheme signature.Scheme)) error {
+	key, sign, scheme := f()
+
+	pub, err := keys.NewPublicKeyFromBytes(key, elliptic.P256())
 	if err != nil {
 		return fmt.Errorf("%w: %v", ErrInvalidPublicKey, err)
 	}
 
 	if !cfg.schemeFixed {
-		cfg.scheme = sig.Scheme()
+		cfg.scheme = scheme
 	}
 
 	switch cfg.scheme {
 	case signature.ECDSAWithSHA512:
 		h := sha512.Sum512(msg)
-		r, s := unmarshalXY(sig.Sign())
+		r, s := unmarshalXY(sign)
 		if r != nil && s != nil && ecdsa.Verify((*ecdsa.PublicKey)(pub), h[:], r, s) {
 			return nil
 		}
 		return ErrInvalidSignature
 	case signature.RFC6979WithSHA256:
 		h := sha256.Sum256(msg)
-		if pub.Verify(sig.Sign(), h[:]) {
+		if pub.Verify(sign, h[:]) {
 			return nil
 		}
 		return ErrInvalidSignature