sweeptimelockmanual: add flags for num keys and num updates

cmd/chantools/sweeptimelockmanual.go

@@ -34,6 +34,9 @@ type sweepTimeLockManualCommand struct {
 	TimeLockAddr              string
 	RemoteRevocationBasePoint string
 
+	MaxNumChansTotal  uint16
+	MaxNumChanUpdates uint64
+
 	rootKey *rootKey
 	inputs  *inputFlags
 	cmd     *cobra.Command
@@ -79,6 +82,16 @@ address is always the one that's longer (because it's P2WSH and not P2PKH).`,
 		&cc.MaxCsvLimit, "maxcsvlimit", defaultCsvLimit, "maximum CSV "+
 			"limit to use",
 	)
+	cc.cmd.Flags().Uint16Var(
+		&cc.MaxNumChansTotal, "maxnumchanstotal", maxKeys, "maximum "+
+			"number of keys to try, set to maximum number of "+
+			"channels the local node potentially has or had",
+	)
+	cc.cmd.Flags().Uint64Var(
+		&cc.MaxNumChanUpdates, "maxnumchanupdates", maxPoints,
+		"maximum number of channel updates to try, set to maximum "+
+			"number of times the channel was used",
+	)
 	cc.cmd.Flags().Uint16Var(
 		&cc.FeeRate, "feerate", defaultFeeSatPerVByte, "fee rate to "+
 			"use for the sweep transaction in sat/vByte",
@@ -123,13 +136,15 @@ func (c *sweepTimeLockManualCommand) Execute(_ *cobra.Command, _ []string) error
 
 	return sweepTimeLockManual(
 		extendedKey, c.APIURL, c.SweepAddr, c.TimeLockAddr,
-		remoteRevPoint, c.MaxCsvLimit, c.Publish, c.FeeRate,
+		remoteRevPoint, c.MaxCsvLimit, c.MaxNumChansTotal,
+		c.MaxNumChanUpdates, c.Publish, c.FeeRate,
 	)
 }
 
 func sweepTimeLockManual(extendedKey *hdkeychain.ExtendedKey, apiURL string,
 	sweepAddr, timeLockAddr string, remoteRevPoint *btcec.PublicKey,
-	maxCsvTimeout uint16, publish bool, feeRate uint16) error {
+	maxCsvTimeout, maxNumChannels uint16, maxNumChanUpdates uint64,
+	publish bool, feeRate uint16) error {
 
 	// First of all, we need to parse the lock addr and make sure we can
 	// brute force the script with the information we have. If not, we can't
@@ -164,9 +179,10 @@ func sweepTimeLockManual(extendedKey *hdkeychain.ExtendedKey, apiURL string,
 		delayDesc   *keychain.KeyDescriptor
 		commitPoint *btcec.PublicKey
 	)
-	for i := uint32(0); i < maxKeys; i++ {
+	for i := uint16(0); i < maxNumChannels; i++ {
 		csvTimeout, script, scriptHash, commitPoint, delayDesc, err = tryKey(
-			baseKey, remoteRevPoint, maxCsvTimeout, lockScript, i,
+			baseKey, remoteRevPoint, maxCsvTimeout, lockScript,
+			uint32(i), maxNumChanUpdates,
 		)
 
 		if err == nil {
@@ -285,8 +301,9 @@ func sweepTimeLockManual(extendedKey *hdkeychain.ExtendedKey, apiURL string,
 }
 
 func tryKey(baseKey *hdkeychain.ExtendedKey, remoteRevPoint *btcec.PublicKey,
-	maxCsvTimeout uint16, lockScript []byte, idx uint32) (int32, []byte,
-	[]byte, *btcec.PublicKey, *keychain.KeyDescriptor, error) {
+	maxCsvTimeout uint16, lockScript []byte, idx uint32,
+	maxNumChanUpdates uint64) (int32, []byte, []byte, *btcec.PublicKey,
+	*keychain.KeyDescriptor, error) {
 
 	// The easy part first, let's derive the delay base point.
 	delayPath := []uint32{
@@ -316,7 +333,7 @@ func tryKey(baseKey *hdkeychain.ExtendedKey, remoteRevPoint *btcec.PublicKey,
 	// points and CSV values.
 	csvTimeout, script, scriptHash, commitPoint, err := bruteForceDelayPoint(
 		delayPrivKey.PubKey(), remoteRevPoint, revRoot, lockScript,
-		maxCsvTimeout,
+		maxCsvTimeout, maxNumChanUpdates,
 	)
 	if err == nil {
 		return csvTimeout, script, scriptHash, commitPoint,
@@ -346,7 +363,7 @@ func tryKey(baseKey *hdkeychain.ExtendedKey, remoteRevPoint *btcec.PublicKey,
 	// points and CSV values.
 	csvTimeout, script, scriptHash, commitPoint, err = bruteForceDelayPoint(
 		delayPrivKey.PubKey(), remoteRevPoint, revRoot2, lockScript,
-		maxCsvTimeout,
+		maxCsvTimeout, maxNumChanUpdates,
 	)
 	if err == nil {
 		return csvTimeout, script, scriptHash, commitPoint,
@@ -378,7 +395,7 @@ func tryKey(baseKey *hdkeychain.ExtendedKey, remoteRevPoint *btcec.PublicKey,
 
 	csvTimeout, script, scriptHash, commitPoint, err = bruteForceDelayPoint(
 		delayPrivKey.PubKey(), remoteRevPoint, revRoot3, lockScript,
-		maxCsvTimeout,
+		maxCsvTimeout, maxNumChanUpdates,
 	)
 	if err == nil {
 		return csvTimeout, script, scriptHash, commitPoint,
@@ -396,9 +413,10 @@ func tryKey(baseKey *hdkeychain.ExtendedKey, remoteRevPoint *btcec.PublicKey,
 
 func bruteForceDelayPoint(delayBase, revBase *btcec.PublicKey,
 	revRoot *shachain.RevocationProducer, lockScript []byte,
-	maxCsvTimeout uint16) (int32, []byte, []byte, *btcec.PublicKey, error) {
+	maxCsvTimeout uint16, maxChanUpdates uint64) (int32, []byte, []byte,
+	*btcec.PublicKey, error) {
 
-	for i := uint64(0); i < maxPoints; i++ {
+	for i := uint64(0); i < maxChanUpdates; i++ {
 		revPreimage, err := revRoot.AtIndex(i)
 		if err != nil {
 			return 0, nil, nil, nil, err