core: Add timeout to TSig (#330)

2 files changed, 49 insertions, 1 deletions

diff --git a/core/configuration-chain.go b/core/configuration-chain.go
index 03a4602..47a40b1 100644
--- a/core/configuration-chain.go
+++ b/core/configuration-chain.go
@@ -265,13 +265,24 @@ func (cc *configurationChain) runTSig(
 			}
 		}
 	}()
+	timeout := make(chan struct{}, 1)
+	go func() {
+		// TODO(jimmy-dexon): make timeout configurable.
+		time.Sleep(5 * time.Second)
+		timeout <- struct{}{}
+		cc.tsigReady.Broadcast()
+	}()
 	var signature crypto.Signature
 	var err error
 	for func() bool {
 		signature, err = cc.tsig[hash].signature()
+		select {
+		case <-timeout:
+			return false
+		default:
+		}
 		return err == ErrNotEnoughtPartialSignatures
 	}() {
-		// TODO(jimmy-dexon): add a timeout here.
 		cc.tsigReady.Wait()
 	}
 	delete(cc.tsig, hash)
diff --git a/core/configuration-chain_test.go b/core/configuration-chain_test.go
index 129d07c..281be06 100644
--- a/core/configuration-chain_test.go
+++ b/core/configuration-chain_test.go
@@ -330,6 +330,43 @@ func (s *ConfigurationChainTestSuite) TestMultipleTSig() {
 	}
 }
 
+func (s *ConfigurationChainTestSuite) TestTSigTimeout() {
+	k := 2
+	n := 7
+	round := uint64(0)
+	cfgChains := s.runDKG(k, n, round)
+	timeout := 6 * time.Second
+
+	hash := crypto.Keccak256Hash([]byte("🍯🍋"))
+
+	psigs := s.preparePartialSignature(hash, round, cfgChains)
+
+	errs := make(chan error, n)
+	qualify := 0
+	for nID, cc := range cfgChains {
+		if _, exist := cc.gpk[round].qualifyNodeIDs[nID]; !exist {
+			continue
+		}
+		qualify++
+		go func(cc *configurationChain) {
+			_, err := cc.runTSig(round, hash)
+			// Prevent racing by collecting errors and check in main thread.
+			errs <- err
+		}(cc)
+		// Only 1 partial signature is provided.
+		err := cc.processPartialSignature(psigs[0])
+		s.Require().NoError(err)
+	}
+	time.Sleep(timeout)
+	s.Require().Len(errs, qualify)
+	for nID, cc := range cfgChains {
+		if _, exist := cc.gpk[round].qualifyNodeIDs[nID]; !exist {
+			continue
+		}
+		s.Equal(<-errs, ErrNotEnoughtPartialSignatures)
+	}
+}
+
 func TestConfigurationChain(t *testing.T) {
 	suite.Run(t, new(ConfigurationChainTestSuite))
 }