From 8025296bc143d691f8e4d2e9979c4394615c4cf5 Mon Sep 17 00:00:00 2001
From: ffranr <robyn@lightning.engineering>
Date: Sat, 13 Apr 2024 11:39:51 +0100
Subject: [PATCH] itest: add itest for field modification HTLC interception
 response

Implement an integration test where an HTLC is intercepted and the
interception response modifies fields in the resultant p2p message.
---
 itest/list_on_test.go                 |   4 +
 itest/lnd_forward_interceptor_test.go | 135 ++++++++++++++++++++++++++
 2 files changed, 139 insertions(+)

diff --git a/itest/list_on_test.go b/itest/list_on_test.go
index b4f586fa2..ca197778d 100644
--- a/itest/list_on_test.go
+++ b/itest/list_on_test.go
@@ -422,6 +422,10 @@ var allTestCases = []*lntest.TestCase{
 		Name:     "forward interceptor",
 		TestFunc: testForwardInterceptorBasic,
 	},
+	{
+		Name:     "forward interceptor modified htlc",
+		TestFunc: testForwardInterceptorModifiedHtlc,
+	},
 	{
 		Name:     "zero conf channel open",
 		TestFunc: testZeroConfChannelOpen,
diff --git a/itest/lnd_forward_interceptor_test.go b/itest/lnd_forward_interceptor_test.go
index b5d08f5bf..213f777be 100644
--- a/itest/lnd_forward_interceptor_test.go
+++ b/itest/lnd_forward_interceptor_test.go
@@ -1,6 +1,7 @@
 package itest
 
 import (
+	"context"
 	"fmt"
 	"strings"
 	"time"
@@ -344,6 +345,140 @@ func testForwardInterceptorBasic(ht *lntest.HarnessTest) {
 	ht.CloseChannel(bob, cpBC)
 }
 
+// testForwardInterceptorModifiedHtlc tests that the interceptor can modify the
+// amount and custom records of an intercepted HTLC and resume it.
+func testForwardInterceptorModifiedHtlc(ht *lntest.HarnessTest) {
+	// Initialize the test context with 3 connected nodes.
+	ts := newInterceptorTestScenario(ht)
+
+	alice, bob, carol := ts.alice, ts.bob, ts.carol
+
+	// Open and wait for channels.
+	const chanAmt = btcutil.Amount(300000)
+	p := lntest.OpenChannelParams{Amt: chanAmt}
+	reqs := []*lntest.OpenChannelRequest{
+		{Local: alice, Remote: bob, Param: p},
+		{Local: bob, Remote: carol, Param: p},
+	}
+	resp := ht.OpenMultiChannelsAsync(reqs)
+	cpAB, cpBC := resp[0], resp[1]
+
+	// Make sure Alice is aware of channel Bob=>Carol.
+	ht.AssertTopologyChannelOpen(alice, cpBC)
+
+	// Connect an interceptor to Bob's node.
+	bobInterceptor, cancelBobInterceptor := bob.RPC.HtlcInterceptor()
+
+	// Prepare the test cases.
+	invoiceValueAmtMsat := int64(1000)
+	req := &lnrpc.Invoice{ValueMsat: invoiceValueAmtMsat}
+	addResponse := carol.RPC.AddInvoice(req)
+	invoice := carol.RPC.LookupInvoice(addResponse.RHash)
+	tc := &interceptorTestCase{
+		amountMsat: invoiceValueAmtMsat,
+		invoice:    invoice,
+		payAddr:    invoice.PaymentAddr,
+	}
+
+	// We initiate a payment from Alice.
+	done := make(chan struct{})
+	go func() {
+		// Signal that all the payments have been sent.
+		defer close(done)
+
+		ts.sendPaymentAndAssertAction(tc)
+	}()
+
+	// We start the htlc interceptor with a simple implementation that saves
+	// all intercepted packets. These packets are held to simulate a
+	// pending payment.
+	packet := ht.ReceiveHtlcInterceptor(bobInterceptor)
+
+	// Resume the intercepted HTLC with a modified amount and custom
+	// records.
+	if packet.CustomRecords == nil {
+		packet.CustomRecords = make(map[uint64][]byte)
+	}
+	customRecords := packet.CustomRecords
+
+	// Add custom records entry.
+	crKey := uint64(65537)
+	crValue := []byte("custom-records-test-value")
+	customRecords[crKey] = crValue
+
+	action := routerrpc.ResolveHoldForwardAction_RESUME_MODIFIED
+	newOutgoingAmountMsat := packet.OutgoingAmountMsat + 4000
+
+	err := bobInterceptor.Send(&routerrpc.ForwardHtlcInterceptResponse{
+		IncomingCircuitKey: packet.IncomingCircuitKey,
+		OutgoingAmountMsat: newOutgoingAmountMsat,
+		CustomRecords:      customRecords,
+		Action:             action,
+	})
+	require.NoError(ht, err, "failed to send request")
+
+	// Check that the modified UpdateAddHTLC message fields were reported in
+	// Carol's log.
+	targetLogPrefixStr := "Received UpdateAddHTLC("
+	targetOutgoingAmountMsatStr := fmt.Sprintf(
+		"amt=%d", newOutgoingAmountMsat,
+	)
+
+	// Formulate custom records target log string.
+	var asciiValues []string
+	for _, b := range crValue {
+		asciiValues = append(asciiValues, fmt.Sprintf("%d", b))
+	}
+
+	targetCustomRecordsStr := fmt.Sprintf(
+		"%d:[%s]", crKey, strings.Join(asciiValues, " "),
+	)
+
+	// logEntryCheck is a helper function that checks if the log entry
+	// contains the expected strings.
+	logEntryCheck := func(logEntry string) bool {
+		return strings.Contains(logEntry, targetLogPrefixStr) &&
+			strings.Contains(logEntry, targetCustomRecordsStr) &&
+			strings.Contains(logEntry, targetOutgoingAmountMsatStr)
+	}
+
+	// Wait for the log entry to appear in Carol's log.
+	require.Eventually(ht, func() bool {
+		ctx := context.Background()
+		dbgInfo, err := carol.RPC.LN.GetDebugInfo(
+			ctx, &lnrpc.GetDebugInfoRequest{},
+		)
+		require.NoError(ht, err, "failed to get Carol node debug info")
+
+		for _, logEntry := range dbgInfo.Log {
+			if logEntryCheck(logEntry) {
+				return true
+			}
+		}
+
+		return false
+	}, defaultTimeout, time.Second)
+
+	// Cancel the context, which will disconnect Bob's interceptor.
+	cancelBobInterceptor()
+
+	// Make sure all goroutines are finished.
+	select {
+	case <-done:
+	case <-time.After(defaultTimeout):
+		require.Fail(ht, "timeout waiting for sending payment")
+	}
+
+	// Assert that the payment was successful.
+	var preimage lntypes.Preimage
+	copy(preimage[:], invoice.RPreimage)
+	ht.AssertPaymentStatus(alice, preimage, lnrpc.Payment_SUCCEEDED)
+
+	// Finally, close channels.
+	ht.CloseChannel(alice, cpAB)
+	ht.CloseChannel(bob, cpBC)
+}
+
 // interceptorTestScenario is a helper struct to hold the test context and
 // provide the needed functionality.
 type interceptorTestScenario struct {