Fix #131: add explicit zeroization in MPC sessions

pkg/mpc/ecdsa_keygen_session.go

@@ -13,6 +13,7 @@ import (
 	"github.com/fystack/mpcium/pkg/kvstore"
 	"github.com/fystack/mpcium/pkg/logger"
 	"github.com/fystack/mpcium/pkg/messaging"
+	"github.com/fystack/mpcium/pkg/security"
 )
 
 type KeyGenSession interface {
@@ -98,11 +99,13 @@ func (s *ecdsaKeygenSession) GenerateKey(done func()) {
 		case msg := <-s.outCh:
 			s.handleTssMessage(msg)
 		case saveData := <-s.endCh:
+			defer security.ZeroEcdsaKeygenLocalPartySaveData(saveData)
 			keyBytes, err := json.Marshal(saveData)
 			if err != nil {
 				s.ErrCh <- err
 				return
 			}
+			defer security.ZeroBytes(keyBytes)
 
 			err = s.kvstore.Put(s.composeKey(walletIDWithVersion(s.walletID, s.GetVersion())), keyBytes)
 			if err != nil {
@@ -148,3 +151,4 @@ func (s *ecdsaKeygenSession) GenerateKey(done func()) {
 		}
 	}
 }
+

pkg/mpc/ecdsa_resharing_session.go

@@ -14,6 +14,7 @@ import (
 	"github.com/fystack/mpcium/pkg/kvstore"
 	"github.com/fystack/mpcium/pkg/logger"
 	"github.com/fystack/mpcium/pkg/messaging"
+	"github.com/fystack/mpcium/pkg/security"
 )
 
 type ReshareSession interface {
@@ -172,11 +173,14 @@ func (s *ecdsaReshareSession) Reshare(done func()) {
 			// skip for old committee
 			if saveData.ECDSAPub != nil {
 
+				defer security.ZeroEcdsaKeygenLocalPartySaveData(saveData)
+
 				keyBytes, err := json.Marshal(saveData)
 				if err != nil {
 					s.ErrCh <- err
 					return
 				}
+				defer security.ZeroBytes(keyBytes)
 
 				newVersion := s.GetVersion() + 1
 				key := s.composeKey(walletIDWithVersion(s.walletID, newVersion))

pkg/mpc/ecdsa_signing_session.go

@@ -17,6 +17,7 @@ import (
 	"github.com/fystack/mpcium/pkg/kvstore"
 	"github.com/fystack/mpcium/pkg/logger"
 	"github.com/fystack/mpcium/pkg/messaging"
+	"github.com/fystack/mpcium/pkg/security"
 	"github.com/samber/lo"
 )
 
@@ -129,12 +130,14 @@ func (s *ecdsaSigningSession) Init(tx *big.Int) error {
 	if err != nil {
 		return errors.Wrap(err, "Failed to get wallet data from KVStore")
 	}
+	defer security.ZeroBytes(keyData)
 	// Check if all the participants of the key are present
 	var data keygen.LocalPartySaveData
 	err = json.Unmarshal(keyData, &data)
 	if err != nil {
 		return errors.Wrap(err, "Failed to unmarshal wallet data")
 	}
+	defer security.ZeroEcdsaKeygenLocalPartySaveData(&data)
 
 	if len(s.derivationPath) > 0 {
 		il, extendedChildPk, errorDerivation := s.ckd.Derive(s.walletID, data.ECDSAPub, s.derivationPath, tss.S256())
@@ -168,7 +171,6 @@ func (s *ecdsaSigningSession) Sign(onSuccess func(data []byte)) {
 	}()
 
 	for {
-
 		select {
 		case msg := <-s.outCh:
 			s.handleTssMessage(msg)
@@ -207,7 +209,6 @@ func (s *ecdsaSigningSession) Sign(onSuccess func(data []byte)) {
 			})
 			if err != nil {
 				s.ErrCh <- errors.Wrap(err, "Failed to publish sign success message")
-
 				return
 			}
 
@@ -220,6 +221,36 @@ func (s *ecdsaSigningSession) Sign(onSuccess func(data []byte)) {
 			onSuccess(bytes)
 			return
 		}
+	}
+}
+
+func (s *ecdsaSigningSession) Close() error {
+	if s == nil {
+		return nil
+	}
+
+	if s.data != nil {
+		security.ZeroEcdsaKeygenLocalPartySaveData(s.data)
+		s.data = nil
+	}
+
+	if s.tx != nil {
+		s.tx.SetInt64(0)
+		s.tx = nil
+	}
 
+	if s.derivationPath != nil {
+		for i := range s.derivationPath {
+			s.derivationPath[i] = 0
+		}
+		s.derivationPath = nil
 	}
+
+	s.ckd = nil
+
+	// Avoid closing endCh here to prevent send-on-closed-channel panics.
+	// Let the producer side (tss-lib) own channel lifetime.
+	s.endCh = nil
+
+	return s.session.Close()
 }

pkg/mpc/eddsa_keygen_session.go

@@ -12,6 +12,7 @@ import (
 	"github.com/fystack/mpcium/pkg/kvstore"
 	"github.com/fystack/mpcium/pkg/logger"
 	"github.com/fystack/mpcium/pkg/messaging"
+	"github.com/fystack/mpcium/pkg/security"
 )
 
 type eddsaKeygenSession struct {
@@ -86,11 +87,13 @@ func (s *eddsaKeygenSession) GenerateKey(done func()) {
 		case msg := <-s.outCh:
 			s.handleTssMessage(msg)
 		case saveData := <-s.endCh:
+			defer security.ZeroEddsaKeygenLocalPartySaveData(saveData)
 			keyBytes, err := json.Marshal(saveData)
 			if err != nil {
 				s.ErrCh <- err
 				return
 			}
+			defer security.ZeroBytes(keyBytes)
 
 			err = s.kvstore.Put(s.composeKey(walletIDWithVersion(s.walletID, s.GetVersion())), keyBytes)
 			if err != nil {

pkg/mpc/eddsa_resharing_session.go

@@ -13,6 +13,7 @@ import (
 	"github.com/fystack/mpcium/pkg/kvstore"
 	"github.com/fystack/mpcium/pkg/logger"
 	"github.com/fystack/mpcium/pkg/messaging"
+	"github.com/fystack/mpcium/pkg/security"
 )
 
 type eddsaReshareSession struct {
@@ -157,11 +158,14 @@ func (s *eddsaReshareSession) Reshare(done func()) {
 		select {
 		case saveData := <-s.endCh:
 			if saveData.EDDSAPub != nil {
-				keyBytes, err := json.Marshal(saveData)
+				defer security.ZeroEddsaKeygenLocalPartySaveData(saveData)
+
+					keyBytes, err := json.Marshal(saveData)
 				if err != nil {
 					s.ErrCh <- err
 					return
 				}
+				defer security.ZeroBytes(keyBytes)
 
 				newVersion := s.GetVersion() + 1
 				key := s.composeKey(walletIDWithVersion(s.walletID, newVersion))

pkg/mpc/eddsa_signing_session.go

@@ -17,6 +17,7 @@ import (
 	"github.com/fystack/mpcium/pkg/kvstore"
 	"github.com/fystack/mpcium/pkg/logger"
 	"github.com/fystack/mpcium/pkg/messaging"
+	"github.com/fystack/mpcium/pkg/security"
 	"github.com/samber/lo"
 )
 
@@ -118,12 +119,14 @@ func (s *eddsaSigningSession) Init(tx *big.Int) error {
 	if err != nil {
 		return errors.Wrap(err, "Failed to get wallet data from KVStore")
 	}
+	defer security.ZeroBytes(keyData)
 	// Check if all the participants of the key are present
 	var data keygen.LocalPartySaveData
 	err = json.Unmarshal(keyData, &data)
 	if err != nil {
 		return errors.Wrap(err, "Failed to unmarshal wallet data")
 	}
+	defer security.ZeroEddsaKeygenLocalPartySaveData(&data)
 
 	if len(s.derivationPath) > 0 {
 		il, extendedChildPk, errorDerivation := s.ckd.Derive(s.walletID, data.EDDSAPub, s.derivationPath, tss.Edwards())
@@ -210,3 +213,39 @@ func (s *eddsaSigningSession) Sign(onSuccess func(data []byte)) {
 
 	}
 }
+// Close cleans up the EDDSA signing session by zeroing all sensitive data.
+func (s *eddsaSigningSession) Close() error {
+	if s == nil {
+		return nil
+	}
+
+	// Zero out sensitive data
+	if s.data != nil {
+		security.ZeroEddsaKeygenLocalPartySaveData(s.data)
+		s.data = nil
+	}
+
+	// Clear other sensitive fields
+	if s.tx != nil {
+		s.tx.SetInt64(0)
+		s.tx = nil
+	}
+
+	// Clear the derivation path
+	if s.derivationPath != nil {
+		for i := range s.derivationPath {
+			s.derivationPath[i] = 0
+		}
+		s.derivationPath = nil
+	}
+
+	// Clear CKD reference
+	s.ckd = nil
+
+	// Avoid closing endCh here to prevent send-on-closed-channel panics.
+	// Let the producer side (tss-lib) own channel lifetime.
+	s.endCh = nil
+
+	// Call parent's Close() to handle cleanup of subscriptions
+	return s.session.Close()
+}

pkg/security/zeroize.go

@@ -0,0 +1,50 @@
+package security
+
+import (
+	"math/big"
+
+	ecdsakeygen "github.com/bnb-chain/tss-lib/v2/ecdsa/keygen"
+	eddsakeygen "github.com/bnb-chain/tss-lib/v2/eddsa/keygen"
+)
+
+// ZeroEcdsaKeygenLocalPartySaveData securely zeros out sensitive fields in an ECDSA LocalPartySaveData struct.
+// It handles big.Ints and other sensitive data within the struct.
+func ZeroEcdsaKeygenLocalPartySaveData(data *ecdsakeygen.LocalPartySaveData) {
+	if data == nil {
+		return
+	}
+	if data.Xi != nil {
+		zeroBigInt(data.Xi)
+	}
+	if data.ShareID != nil {
+		zeroBigInt(data.ShareID)
+	}
+}
+
+// ZeroEddsaKeygenLocalPartySaveData securely zeros out sensitive fields in an EDDSA LocalPartySaveData struct.
+func ZeroEddsaKeygenLocalPartySaveData(data *eddsakeygen.LocalPartySaveData) {
+	if data == nil {
+		return
+	}
+	if data.Xi != nil {
+		zeroBigInt(data.Xi)
+	}
+	if data.ShareID != nil {
+		zeroBigInt(data.ShareID)
+	}
+}
+
+func zeroBigInt(x *big.Int) {
+	if x == nil {
+		return
+	}
+	words := x.Bits()
+	for i := range words {
+		words[i] = 0
+	}
+	x.SetInt64(0)
+}
+func ZeroBigIntForensic(x *big.Int) {
+	zeroBigInt(x)
+}
+

pkg/security/zeroize_test.go

@@ -0,0 +1,69 @@
+package security
+
+import (
+	"math/big"
+	"testing"
+
+	ecdsakeygen "github.com/bnb-chain/tss-lib/v2/ecdsa/keygen"
+	eddsakeygen "github.com/bnb-chain/tss-lib/v2/eddsa/keygen"
+)
+func TestZeroBigInt_OverwritesBackingWords(t *testing.T) {
+	x := new(big.Int).Lsh(big.NewInt(1), 2048)
+	wordsBefore := append([]big.Word(nil), x.Bits()...)
+	if len(wordsBefore) == 0 {
+		t.Fatal("expected non-empty big.Int words")
+	}
+
+	zeroBigInt(x)
+
+	for i, w := range wordsBefore {
+		if w == 0 {
+			continue
+		}
+		if i < len(x.Bits()) && x.Bits()[i] != 0 {
+			t.Fatalf("expected big.Int word %d to be overwritten", i)
+		}
+	}
+
+	if x.Sign() != 0 {
+		t.Fatalf("expected big.Int value to be zero, got %v", x)
+	}
+}
+
+func TestZeroEcdsaKeygenLocalPartySaveData_Nil(t *testing.T) {
+	ZeroEcdsaKeygenLocalPartySaveData(nil)
+}
+
+func TestZeroEcdsaKeygenLocalPartySaveData_ZeroesFields(t *testing.T) {
+	data := &ecdsakeygen.LocalPartySaveData{}
+	data.Xi = big.NewInt(123)
+	data.ShareID = big.NewInt(456)
+
+	ZeroEcdsaKeygenLocalPartySaveData(data)
+
+	if data.Xi == nil || data.Xi.Sign() != 0 {
+		t.Fatalf("expected Xi to be zero, got %v", data.Xi)
+	}
+	if data.ShareID == nil || data.ShareID.Sign() != 0 {
+		t.Fatalf("expected ShareID to be zero, got %v", data.ShareID)
+	}
+}
+
+func TestZeroEddsaKeygenLocalPartySaveData_Nil(t *testing.T) {
+	ZeroEddsaKeygenLocalPartySaveData(nil)
+}
+
+func TestZeroEddsaKeygenLocalPartySaveData_ZeroesFields(t *testing.T) {
+	data := &eddsakeygen.LocalPartySaveData{}
+	data.Xi = big.NewInt(789)
+	data.ShareID = big.NewInt(101112)
+
+	ZeroEddsaKeygenLocalPartySaveData(data)
+
+	if data.Xi == nil || data.Xi.Sign() != 0 {
+		t.Fatalf("expected Xi to be zero, got %v", data.Xi)
+	}
+	if data.ShareID == nil || data.ShareID.Sign() != 0 {
+		t.Fatalf("expected ShareID to be zero, got %v", data.ShareID)
+	}
+}