Merge pull request #155 from lightninglabs/reorg-notifications

chainnotifier: add re-org chan option to RegisterConfirmationsNtfn

chainnotifier_client.go

@@ -19,6 +19,14 @@ type notifierOptions struct {
 	// includeBlock if true, then the dispatched confirmation notification
 	// will include the block that mined the transaction.
 	includeBlock bool
+
+	// reOrgChan if set, will be sent on if the transaction is re-organized
+	// out of the chain. This channel being set will also imply that we
+	// don't cancel the notification listener after having received one
+	// confirmation event. That means the caller manually needs to cancel
+	// the passed in context to cancel being notified once the required
+	// number of confirmations have been reached.
+	reOrgChan chan struct{}
 }
 
 // defaultNotifierOptions returns the set of default options for the notifier.
@@ -38,15 +46,27 @@ func WithIncludeBlock() NotifierOption {
 	}
 }
 
+// WithReOrgChan configures a channel that will be sent on if the transaction is
+// re-organized out of the chain. This channel being set will also imply that we
+// don't cancel the notification listener after having received one confirmation
+// event. That means the caller manually needs to cancel the passed in context
+// to cancel being notified once the required number of confirmations have been
+// reached.
+func WithReOrgChan(reOrgChan chan struct{}) NotifierOption {
+	return func(o *notifierOptions) {
+		o.reOrgChan = reOrgChan
+	}
+}
+
 // ChainNotifierClient exposes base lightning functionality.
 type ChainNotifierClient interface {
 	RegisterBlockEpochNtfn(ctx context.Context) (
 		chan int32, chan error, error)
 
 	RegisterConfirmationsNtfn(ctx context.Context, txid *chainhash.Hash,
 		pkScript []byte, numConfs, heightHint int32,
-		opts ...NotifierOption) (
-		chan *chainntnfs.TxConfirmation, chan error, error)
+		opts ...NotifierOption) (chan *chainntnfs.TxConfirmation,
+		chan error, error)
 
 	RegisterSpendNtfn(ctx context.Context,
 		outpoint *wire.OutPoint, pkScript []byte, heightHint int32) (
@@ -153,8 +173,8 @@ func (s *chainNotifierClient) RegisterSpendNtfn(ctx context.Context,
 
 func (s *chainNotifierClient) RegisterConfirmationsNtfn(ctx context.Context,
 	txid *chainhash.Hash, pkScript []byte, numConfs, heightHint int32,
-	optFuncs ...NotifierOption) (
-	chan *chainntnfs.TxConfirmation, chan error, error) {
+	optFuncs ...NotifierOption) (chan *chainntnfs.TxConfirmation,
+	chan error, error) {
 
 	opts := defaultNotifierOptions()
 	for _, optFunc := range optFuncs {
@@ -166,8 +186,7 @@ func (s *chainNotifierClient) RegisterConfirmationsNtfn(ctx context.Context,
 		txidSlice = txid[:]
 	}
 	confStream, err := s.client.RegisterConfirmationsNtfn(
-		s.chainMac.WithMacaroonAuth(ctx),
-		&chainrpc.ConfRequest{
+		s.chainMac.WithMacaroonAuth(ctx), &chainrpc.ConfRequest{
 			Script:       pkScript,
 			NumConfs:     uint32(numConfs),
 			HeightHint:   uint32(heightHint),
@@ -195,7 +214,7 @@ func (s *chainNotifierClient) RegisterConfirmationsNtfn(ctx context.Context,
 			}
 
 			switch c := confEvent.Event.(type) {
-			// Script confirmed
+			// Script confirmed.
 			case *chainrpc.ConfEvent_Conf:
 				tx, err := decodeTx(c.Conf.RawTx)
 				if err != nil {
@@ -205,7 +224,9 @@ func (s *chainNotifierClient) RegisterConfirmationsNtfn(ctx context.Context,
 
 				var block *wire.MsgBlock
 				if opts.includeBlock {
-					block, err = decodeBlock(c.Conf.RawBlock)
+					block, err = decodeBlock(
+						c.Conf.RawBlock,
+					)
 					if err != nil {
 						errChan <- err
 						return
@@ -227,10 +248,26 @@ func (s *chainNotifierClient) RegisterConfirmationsNtfn(ctx context.Context,
 					TxIndex:     c.Conf.TxIndex,
 					Block:       block,
 				}
-				return
 
-			// Ignore reorg events, not supported.
+				// If we're running in re-org aware mode, then
+				// we don't return here, since we might want to
+				// be informed about the new block we got
+				// confirmed in after a re-org.
+				if opts.reOrgChan == nil {
+					return
+				}
+
+			// On a re-org, we just need to signal, we don't have
+			// any additional information. But we only signal if the
+			// caller requested to be notified about re-orgs.
 			case *chainrpc.ConfEvent_Reorg:
+				if opts.reOrgChan != nil {
+					select {
+					case opts.reOrgChan <- struct{}{}:
+					case <-ctx.Done():
+						return
+					}
+				}
 				continue
 
 			// Nil event, should never happen.
@@ -240,9 +277,8 @@ func (s *chainNotifierClient) RegisterConfirmationsNtfn(ctx context.Context,
 
 			// Unexpected type.
 			default:
-				errChan <- fmt.Errorf(
-					"conf event has unexpected type",
-				)
+				errChan <- fmt.Errorf("conf event has " +
+					"unexpected type")
 				return
 			}
 		}