1  package htlcswitch
  2  
  3  import (
  4  	"context"
  5  	"crypto/sha256"
  6  	"testing"
  7  	"time"
  8  
  9  	"github.com/lightningnetwork/lnd/lntypes"
 10  	"github.com/lightningnetwork/lnd/lnwallet"
 11  	"github.com/lightningnetwork/lnd/lnwire"
 12  )
 13  
 14  type linkTestContext struct {
 15  	t *testing.T
 16  
 17  	aliceSwitch *Switch
 18  	aliceLink   ChannelLink
 19  	bobChannel  *lnwallet.LightningChannel
 20  	aliceMsgs   <-chan lnwire.Message
 21  }
 22  
 23  // sendHtlcBobToAlice sends an HTLC from Bob to Alice, that pays to a preimage
 24  // already in Alice's registry.
 25  func (l *linkTestContext) sendHtlcBobToAlice(htlc *lnwire.UpdateAddHTLC) {
 26  	l.t.Helper()
 27  
 28  	_, err := l.bobChannel.AddHTLC(htlc, nil)
 29  	if err != nil {
 30  		l.t.Fatalf("bob failed adding htlc: %v", err)
 31  	}
 32  
 33  	l.aliceLink.HandleChannelUpdate(htlc)
 34  }
 35  
 36  // sendHtlcAliceToBob sends an HTLC from Alice to Bob, by first committing the
 37  // HTLC in the circuit map, then delivering the outgoing packet to Alice's link.
 38  // The HTLC will be sent to Bob via Alice's message stream.
 39  func (l *linkTestContext) sendHtlcAliceToBob(htlcID int,
 40  	htlc *lnwire.UpdateAddHTLC) {
 41  
 42  	l.t.Helper()
 43  
 44  	circuitMap := l.aliceSwitch.circuits
 45  	fwdActions, err := circuitMap.CommitCircuits(
 46  		&PaymentCircuit{
 47  			Incoming: CircuitKey{
 48  				HtlcID: uint64(htlcID),
 49  			},
 50  			PaymentHash: htlc.PaymentHash,
 51  		},
 52  	)
 53  	if err != nil {
 54  		l.t.Fatalf("unable to commit circuit: %v", err)
 55  	}
 56  
 57  	if len(fwdActions.Adds) != 1 {
 58  		l.t.Fatalf("expected 1 adds, found %d", len(fwdActions.Adds))
 59  	}
 60  
 61  	err = l.aliceLink.handleSwitchPacket(&htlcPacket{
 62  		incomingHTLCID: uint64(htlcID),
 63  		htlc:           htlc,
 64  	})
 65  	if err != nil {
 66  		l.t.Fatal(err)
 67  	}
 68  }
 69  
 70  // receiveHtlcAliceToBob pulls the next message from Alice's message stream,
 71  // asserts that it is an UpdateAddHTLC, then applies it to Bob's state machine.
 72  func (l *linkTestContext) receiveHtlcAliceToBob() {
 73  	l.t.Helper()
 74  
 75  	var msg lnwire.Message
 76  	select {
 77  	case msg = <-l.aliceMsgs:
 78  	case <-time.After(15 * time.Second):
 79  		l.t.Fatalf("did not received htlc from alice")
 80  	}
 81  
 82  	htlcAdd, ok := msg.(*lnwire.UpdateAddHTLC)
 83  	if !ok {
 84  		l.t.Fatalf("expected UpdateAddHTLC, got %T", msg)
 85  	}
 86  
 87  	_, err := l.bobChannel.ReceiveHTLC(htlcAdd)
 88  	if err != nil {
 89  		l.t.Fatalf("bob failed receiving htlc: %v", err)
 90  	}
 91  }
 92  
 93  // sendCommitSigBobToAlice makes Bob sign a new commitment and send it to
 94  // Alice, asserting that it signs expHtlcs number of HTLCs.
 95  func (l *linkTestContext) sendCommitSigBobToAlice(expHtlcs int) {
 96  	l.t.Helper()
 97  
 98  	testQuit, testQuitFunc := context.WithCancel(l.t.Context())
 99  	defer testQuitFunc()
100  	sigs, err := l.bobChannel.SignNextCommitment(testQuit)
101  	if err != nil {
102  		l.t.Fatalf("error signing commitment: %v", err)
103  	}
104  
105  	commitSig := &lnwire.CommitSig{
106  		CommitSig: sigs.CommitSig,
107  		HtlcSigs:  sigs.HtlcSigs,
108  	}
109  
110  	if len(commitSig.HtlcSigs) != expHtlcs {
111  		l.t.Fatalf("Expected %d htlc sigs, got %d", expHtlcs,
112  			len(commitSig.HtlcSigs))
113  	}
114  
115  	l.aliceLink.HandleChannelUpdate(commitSig)
116  }
117  
118  // receiveRevAndAckAliceToBob waits for Alice to send a RevAndAck to Bob, then
119  // hands this to Bob.
120  func (l *linkTestContext) receiveRevAndAckAliceToBob() {
121  	l.t.Helper()
122  
123  	var msg lnwire.Message
124  	select {
125  	case msg = <-l.aliceMsgs:
126  	case <-time.After(15 * time.Second):
127  		l.t.Fatalf("did not receive message")
128  	}
129  
130  	rev, ok := msg.(*lnwire.RevokeAndAck)
131  	if !ok {
132  		l.t.Fatalf("expected RevokeAndAck, got %T", msg)
133  	}
134  
135  	_, _, err := l.bobChannel.ReceiveRevocation(rev)
136  	if err != nil {
137  		l.t.Fatalf("bob failed receiving revocation: %v", err)
138  	}
139  }
140  
141  // receiveCommitSigAliceToBob waits for Alice to send a CommitSig to Bob,
142  // signing expHtlcs numbers of HTLCs, then hands this to Bob.
143  func (l *linkTestContext) receiveCommitSigAliceToBob(expHtlcs int) {
144  	l.t.Helper()
145  
146  	comSig := l.receiveCommitSigAlice(expHtlcs)
147  
148  	err := l.bobChannel.ReceiveNewCommitment(&lnwallet.CommitSigs{
149  		CommitSig: comSig.CommitSig,
150  		HtlcSigs:  comSig.HtlcSigs,
151  	})
152  	if err != nil {
153  		l.t.Fatalf("bob failed receiving commitment: %v", err)
154  	}
155  }
156  
157  // receiveCommitSigAlice waits for Alice to send a CommitSig, signing expHtlcs
158  // numbers of HTLCs.
159  func (l *linkTestContext) receiveCommitSigAlice(expHtlcs int) *lnwire.CommitSig {
160  	l.t.Helper()
161  
162  	var msg lnwire.Message
163  	select {
164  	case msg = <-l.aliceMsgs:
165  	case <-time.After(15 * time.Second):
166  		l.t.Fatalf("did not receive message")
167  	}
168  
169  	comSig, ok := msg.(*lnwire.CommitSig)
170  	if !ok {
171  		l.t.Fatalf("expected CommitSig, got %T", msg)
172  	}
173  
174  	if len(comSig.HtlcSigs) != expHtlcs {
175  		l.t.Fatalf("expected %d htlc sigs, got %d", expHtlcs,
176  			len(comSig.HtlcSigs))
177  	}
178  
179  	return comSig
180  }
181  
182  // sendRevAndAckBobToAlice make Bob revoke his current commitment, then hand
183  // the RevokeAndAck to Alice.
184  func (l *linkTestContext) sendRevAndAckBobToAlice() {
185  	l.t.Helper()
186  
187  	rev, _, _, err := l.bobChannel.RevokeCurrentCommitment()
188  	if err != nil {
189  		l.t.Fatalf("unable to revoke commitment: %v", err)
190  	}
191  
192  	l.aliceLink.HandleChannelUpdate(rev)
193  }
194  
195  // receiveSettleAliceToBob waits for Alice to send a HTLC settle message to
196  // Bob, then hands this to Bob.
197  func (l *linkTestContext) receiveSettleAliceToBob() {
198  	l.t.Helper()
199  
200  	var msg lnwire.Message
201  	select {
202  	case msg = <-l.aliceMsgs:
203  	case <-time.After(15 * time.Second):
204  		l.t.Fatalf("did not receive message")
205  	}
206  
207  	settleMsg, ok := msg.(*lnwire.UpdateFulfillHTLC)
208  	if !ok {
209  		l.t.Fatalf("expected UpdateFulfillHTLC, got %T", msg)
210  	}
211  
212  	err := l.bobChannel.ReceiveHTLCSettle(settleMsg.PaymentPreimage,
213  		settleMsg.ID)
214  	if err != nil {
215  		l.t.Fatalf("failed settling htlc: %v", err)
216  	}
217  }
218  
219  // sendSettleBobToAlice settles an HTLC on Bob's state machine, then sends an
220  // UpdateFulfillHTLC message to Alice's upstream inbox.
221  func (l *linkTestContext) sendSettleBobToAlice(htlcID uint64,
222  	preimage lntypes.Preimage) {
223  
224  	l.t.Helper()
225  
226  	err := l.bobChannel.SettleHTLC(preimage, htlcID, nil, nil, nil)
227  	if err != nil {
228  		l.t.Fatalf("alice failed settling htlc id=%d hash=%x",
229  			htlcID, sha256.Sum256(preimage[:]))
230  	}
231  
232  	settle := &lnwire.UpdateFulfillHTLC{
233  		ID:              htlcID,
234  		PaymentPreimage: preimage,
235  	}
236  
237  	l.aliceLink.HandleChannelUpdate(settle)
238  }
239  
240  // receiveSettleAliceToBob waits for Alice to send a HTLC settle message to
241  // Bob, then hands this to Bob.
242  func (l *linkTestContext) receiveFailAliceToBob() {
243  	l.t.Helper()
244  
245  	var msg lnwire.Message
246  	select {
247  	case msg = <-l.aliceMsgs:
248  	case <-time.After(15 * time.Second):
249  		l.t.Fatalf("did not receive message")
250  	}
251  
252  	failMsg, ok := msg.(*lnwire.UpdateFailHTLC)
253  	if !ok {
254  		l.t.Fatalf("expected UpdateFailHTLC, got %T", msg)
255  	}
256  
257  	err := l.bobChannel.ReceiveFailHTLC(failMsg.ID, failMsg.Reason)
258  	if err != nil {
259  		l.t.Fatalf("unable to apply received fail htlc: %v", err)
260  	}
261  }
262  
263  // assertNoMsgFromAlice asserts that Alice hasn't sent a message. Before
264  // calling, make sure that Alice has had the opportunity to send the message.
265  func (l *linkTestContext) assertNoMsgFromAlice(timeout time.Duration) {
266  	l.t.Helper()
267  
268  	select {
269  	case msg := <-l.aliceMsgs:
270  		l.t.Fatalf("unexpected message from Alice: %v", msg)
271  	case <-time.After(timeout):
272  	}
273  }