invoices: refine update handling

diff --git a/invoices/invoiceregistry_test.go b/invoices/invoiceregistry_test.go
index cad3e27..f3026b2 100644
--- a/invoices/invoiceregistry_test.go
+++ b/invoices/invoiceregistry_test.go
@@ -98,6 +98,10 @@ func TestInvoiceRegistry(t *testing.T) {
 			name: "AMPWithoutMPPPayload",
 			test: testAMPWithoutMPPPayload,
 		},
+		{
+			name: "AMPWithoutMPPExistingInvoice",
+			test: testAMPWithoutMPPExistingInvoice,
+		},
 		{
 			name: "SpontaneousAmpPayment",
 			test: testSpontaneousAmpPayment,
@@ -1877,6 +1881,46 @@ func testAMPWithoutMPPPayload(t *testing.T,
 	checkFailResolution(t, resolution, invpkg.ResultAmpError)
 }
 
+// testAMPWithoutMPPExistingInvoice checks AMP handling for an existing invoice
+// when spontaneous AMP payments are disabled.
+func testAMPWithoutMPPExistingInvoice(t *testing.T,
+	makeDB func(t *testing.T) (invpkg.InvoiceDB, *clock.TestClock)) {
+
+	t.Parallel()
+	defer timeout()()
+
+	cfg := defaultRegistryConfig()
+	cfg.AcceptAMP = false
+	ctx := newTestContext(t, &cfg, makeDB)
+	ctxb := t.Context()
+
+	invoice := newInvoice(t, false, true)
+	_, err := ctx.registry.AddInvoice(
+		ctxb, invoice, testInvoicePaymentHash,
+	)
+	require.NoError(t, err)
+
+	payload := &mockPayload{
+		amp: record.NewAMP([32]byte{}, [32]byte{}, 0),
+	}
+
+	hodlChan := make(chan interface{}, 1)
+	resolution, err := ctx.registry.NotifyExitHopHtlc(
+		testInvoicePaymentHash, invoice.Terms.Value, testHtlcExpiry,
+		testCurrentHeight, getCircuitKey(10), hodlChan, nil, payload,
+	)
+	require.NoError(t, err)
+	require.NotNil(t, resolution)
+	checkFailResolution(t, resolution, invpkg.ResultAmpError)
+
+	storedInvoice, err := ctx.registry.LookupInvoice(
+		ctxb, testInvoicePaymentHash,
+	)
+	require.NoError(t, err)
+	require.Equal(t, invpkg.ContractOpen, storedInvoice.State)
+	require.Empty(t, storedInvoice.Htlcs)
+}
+
 // testSpontaneousAmpPayment tests receiving a spontaneous AMP payment with both
 // valid and invalid reconstructions.
 func testSpontaneousAmpPayment(t *testing.T,
diff --git a/invoices/update.go b/invoices/update.go
index 7db86c0..937bff2 100644
--- a/invoices/update.go
+++ b/invoices/update.go
@@ -128,16 +128,36 @@ func resolveReplayedHtlc(ctx *invoiceUpdateCtx, inv *Invoice) (bool,
 		return true, ctx.acceptRes(resultReplayToAccepted), nil
 
 	case HtlcStateSettled:
-		pre := inv.Terms.PaymentPreimage
+		var preimage *lntypes.Preimage
+		switch {
+		// AMP invoices store a separate preimage on each HTLC.
+		case inv.IsAMP():
+			if htlc.AMP == nil || htlc.AMP.Preimage == nil {
+				return true, nil, ErrHTLCPreimageMissing
+			}
+
+			preimage = htlc.AMP.Preimage
+			if htlc.AMP.Hash != ctx.hash ||
+				!preimage.Matches(htlc.AMP.Hash) {
+
+				return true, nil, ErrHTLCPreimageMismatch
+			}
+
+		// Regular invoices store their preimage at the invoice level.
+		case inv.Terms.PaymentPreimage == nil:
+			return true, nil, errors.New(
+				"settled invoice missing payment preimage",
+			)
 
-		// Terms.PaymentPreimage will be nil for AMP invoices.
-		// Set it to the HTLCs AMP Preimage instead.
-		if pre == nil {
-			pre = htlc.AMP.Preimage
+		default:
+			preimage = inv.Terms.PaymentPreimage
+			if !preimage.Matches(ctx.hash) {
+				return true, nil, ErrInvoicePreimageMismatch
+			}
 		}
 
 		return true, ctx.settleRes(
-			*pre,
+			*preimage,
 			ResultReplayToSettled,
 		), nil
 
@@ -155,6 +175,12 @@ func resolveReplayedHtlc(ctx *invoiceUpdateCtx, inv *Invoice) (bool,
 func updateInvoice(ctx *invoiceUpdateCtx, inv *Invoice) (
 	*InvoiceUpdateDesc, HtlcResolution, error) {
 
+	// AMP records are processed together with their corresponding MPP
+	// payload.
+	if ctx.amp != nil && ctx.mpp == nil {
+		return nil, ctx.failRes(ResultAmpError), nil
+	}
+
 	// If no MPP payload was provided, then we expect this to be a keysend,
 	// or a payment to an invoice created before we started to require the
 	// MPP payload.
@@ -414,6 +440,12 @@ func reconstructAMPPreimages(ctx *invoiceUpdateCtx,
 func updateLegacy(ctx *invoiceUpdateCtx,
 	inv *Invoice) (*InvoiceUpdateDesc, HtlcResolution, error) {
 
+	// AMP invoices use the MPP update path, where each HTLC's AMP data is
+	// available for processing.
+	if inv.IsAMP() {
+		return nil, ctx.failRes(ResultHtlcInvoiceTypeMismatch), nil
+	}
+
 	// If the invoice is already canceled, there is no further
 	// checking to do.
 	if inv.State == ContractCanceled {
@@ -432,12 +464,11 @@ func updateLegacy(ctx *invoiceUpdateCtx,
 	// if we're in this method it means that the remote party didn't supply
 	// the expected payload. However if this is a keysend payment, then
 	// we'll permit it to pass.
-	_, isKeySend := ctx.customRecords[record.KeySendType]
 	invoiceFeatures := inv.Terms.Features
 	paymentAddrRequired := invoiceFeatures.RequiresFeature(
 		lnwire.PaymentAddrRequired,
 	)
-	if !isKeySend && paymentAddrRequired {
+	if !isValidKeySend(ctx) && paymentAddrRequired {
 		log.Warnf("Payment to pay_hash=%v doesn't include MPP "+
 			"payload, rejecting", ctx.hash)
 		return nil, ctx.failRes(ResultAddressMismatch), nil
@@ -489,8 +520,15 @@ func updateLegacy(ctx *invoiceUpdateCtx,
 		return &update, ctx.acceptRes(resultDuplicateToAccepted), nil
 
 	case ContractSettled:
+		// Legacy settlement uses the invoice-level payment preimage.
+		preimage := inv.Terms.PaymentPreimage
+		if preimage == nil {
+			return nil, ctx.failRes(ResultHtlcInvoiceTypeMismatch),
+				nil
+		}
+
 		return &update, ctx.settleRes(
-			*inv.Terms.PaymentPreimage, ResultDuplicateToSettled,
+			*preimage, ResultDuplicateToSettled,
 		), nil
 	}
 
@@ -504,12 +542,35 @@ func updateLegacy(ctx *invoiceUpdateCtx,
 		return &update, ctx.acceptRes(resultAccepted), nil
 	}
 
+	// A legacy invoice provides its settlement preimage at the invoice
+	// level.
+	preimage := inv.Terms.PaymentPreimage
+	if preimage == nil {
+		return nil, ctx.failRes(ResultHtlcInvoiceTypeMismatch), nil
+	}
+
 	update.State = &InvoiceStateUpdateDesc{
 		NewState: ContractSettled,
-		Preimage: inv.Terms.PaymentPreimage,
+		Preimage: preimage,
 	}
 
 	return &update, ctx.settleRes(
-		*inv.Terms.PaymentPreimage, ResultSettled,
+		*preimage, ResultSettled,
 	), nil
 }
+
+// isValidKeySend reports whether the custom records contain a keysend
+// preimage whose hash matches the payment hash.
+func isValidKeySend(ctx *invoiceUpdateCtx) bool {
+	preimageBytes, ok := ctx.customRecords[record.KeySendType]
+	if !ok {
+		return false
+	}
+
+	preimage, err := lntypes.MakePreimage(preimageBytes)
+	if err != nil {
+		return false
+	}
+
+	return preimage.Hash() == ctx.hash
+}
diff --git a/invoices/update_invoice_test.go b/invoices/update_invoice_test.go
index d6e4ed9..74f4b9a 100644
--- a/invoices/update_invoice_test.go
+++ b/invoices/update_invoice_test.go
@@ -763,3 +763,363 @@ func testUpdateHTLC(t *testing.T, test updateHTLCTest, now time.Time) {
 	require.Equal(t, test.expErr, err)
 	require.Equal(t, test.output, *htlc)
 }
+
+// TestResolveReplayedHtlcSettled checks preimage selection for settled HTLC
+// replays.
+func TestResolveReplayedHtlcSettled(t *testing.T) {
+	t.Parallel()
+
+	const missingPreimageErr = "settled invoice missing payment preimage"
+
+	validPreimage := lntypes.Preimage{1}
+	otherPreimage := lntypes.Preimage{2}
+	validHash := validPreimage.Hash()
+	otherHash := otherPreimage.Hash()
+	setID := [32]byte{3}
+	ampRecord := record.NewAMP([32]byte{4}, setID, 5)
+	ampFeatures := lnwire.NewFeatureVector(
+		lnwire.NewRawFeatureVector(lnwire.AMPRequired),
+		lnwire.Features,
+	)
+
+	tests := []struct {
+		name             string
+		invoicePreimage  *lntypes.Preimage
+		invoiceFeatures  *lnwire.FeatureVector
+		htlcAMP          *InvoiceHtlcAMPData
+		paymentHash      lntypes.Hash
+		expectedPreimage *lntypes.Preimage
+		expectedErr      error
+		expectedErrText  string
+	}{
+		{
+			name:             "regular invoice",
+			invoicePreimage:  &validPreimage,
+			paymentHash:      validHash,
+			expectedPreimage: &validPreimage,
+		},
+		{
+			name:            "regular invoice missing preimage",
+			paymentHash:     validHash,
+			expectedErrText: missingPreimageErr,
+		},
+		{
+			name:            "regular invoice preimage mismatch",
+			invoicePreimage: &otherPreimage,
+			paymentHash:     validHash,
+			expectedErr:     ErrInvoicePreimageMismatch,
+		},
+		{
+			name:            "AMP invoice",
+			invoiceFeatures: ampFeatures,
+			htlcAMP: &InvoiceHtlcAMPData{
+				Record:   *ampRecord,
+				Hash:     validHash,
+				Preimage: &validPreimage,
+			},
+			paymentHash:      validHash,
+			expectedPreimage: &validPreimage,
+		},
+		{
+			name:            "AMP invoice missing HTLC data",
+			invoiceFeatures: ampFeatures,
+			paymentHash:     validHash,
+			expectedErr:     ErrHTLCPreimageMissing,
+		},
+		{
+			name:            "AMP invoice missing preimage",
+			invoiceFeatures: ampFeatures,
+			htlcAMP: &InvoiceHtlcAMPData{
+				Record: *ampRecord,
+				Hash:   validHash,
+			},
+			paymentHash: validHash,
+			expectedErr: ErrHTLCPreimageMissing,
+		},
+		{
+			name:            "AMP invoice preimage mismatch",
+			invoiceFeatures: ampFeatures,
+			htlcAMP: &InvoiceHtlcAMPData{
+				Record:   *ampRecord,
+				Hash:     validHash,
+				Preimage: &otherPreimage,
+			},
+			paymentHash: validHash,
+			expectedErr: ErrHTLCPreimageMismatch,
+		},
+		{
+			name:            "AMP invoice hash mismatch",
+			invoiceFeatures: ampFeatures,
+			htlcAMP: &InvoiceHtlcAMPData{
+				Record:   *ampRecord,
+				Hash:     otherHash,
+				Preimage: &otherPreimage,
+			},
+			paymentHash: validHash,
+			expectedErr: ErrHTLCPreimageMismatch,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+
+			circuitKey := CircuitKey{HtlcID: 1}
+			ctx := &invoiceUpdateCtx{
+				hash:       test.paymentHash,
+				circuitKey: circuitKey,
+			}
+			invoice := &Invoice{
+				Terms: ContractTerm{
+					PaymentPreimage: test.invoicePreimage,
+					Features:        test.invoiceFeatures,
+				},
+				Htlcs: map[CircuitKey]*InvoiceHTLC{
+					circuitKey: {
+						State: HtlcStateSettled,
+						AMP:   test.htlcAMP,
+					},
+				},
+			}
+
+			replayed, resolution, err := resolveReplayedHtlc(
+				ctx, invoice,
+			)
+			require.True(t, replayed)
+
+			switch {
+			case test.expectedErr != nil:
+				require.ErrorIs(t, err, test.expectedErr)
+				require.Nil(t, resolution)
+
+			case test.expectedErrText != "":
+				require.EqualError(t, err, test.expectedErrText)
+				require.Nil(t, resolution)
+
+			default:
+				require.NoError(t, err)
+				requireSettleResolution(
+					t, resolution, ResultReplayToSettled,
+				)
+				settleResolution, ok :=
+					resolution.(*HtlcSettleResolution)
+				require.True(t, ok)
+				require.Equal(
+					t, *test.expectedPreimage,
+					settleResolution.Preimage,
+				)
+			}
+		})
+	}
+}
+
+// TestUpdateInvoiceRejectsAmpWithoutMPP checks that AMP records follow the MPP
+// update path.
+func TestUpdateInvoiceRejectsAmpWithoutMPP(t *testing.T) {
+	t.Parallel()
+
+	ctx, invoice := newLegacyUpdateTestContext(t, ContractOpen)
+	ctx.amp = record.NewAMP([32]byte{1}, [32]byte{2}, 3)
+
+	update, resolution, err := updateInvoice(ctx, invoice)
+	require.NoError(t, err)
+	require.Nil(t, update)
+	requireFailResolution(t, resolution, ResultAmpError)
+}
+
+// TestUpdateInvoiceRejectsAmpInvoiceInLegacyPath checks that AMP invoices are
+// handled by the MPP update path.
+func TestUpdateInvoiceRejectsAmpInvoiceInLegacyPath(t *testing.T) {
+	t.Parallel()
+
+	ctx, invoice := newLegacyUpdateTestContext(t, ContractOpen)
+	invoice.Terms.PaymentPreimage = nil
+	invoice.Terms.Features = lnwire.NewFeatureVector(
+		lnwire.NewRawFeatureVector(
+			lnwire.TLVOnionPayloadOptional,
+			lnwire.PaymentAddrOptional,
+			lnwire.AMPRequired,
+		),
+		lnwire.Features,
+	)
+
+	update, resolution, err := updateInvoice(ctx, invoice)
+	require.NoError(t, err)
+	require.Nil(t, update)
+	requireFailResolution(t, resolution, ResultHtlcInvoiceTypeMismatch)
+}
+
+// TestUpdateLegacyRejectsNilPreimageSettle checks the outcome when a legacy
+// settlement has no invoice-level preimage.
+func TestUpdateLegacyRejectsNilPreimageSettle(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name  string
+		state ContractState
+	}{
+		{
+			name:  "new settle",
+			state: ContractOpen,
+		},
+		{
+			name:  "duplicate settled",
+			state: ContractSettled,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+
+			ctx, invoice := newLegacyUpdateTestContext(
+				t, test.state,
+			)
+			invoice.Terms.PaymentPreimage = nil
+
+			update, resolution, err := updateLegacy(ctx, invoice)
+			require.NoError(t, err)
+			require.Nil(t, update)
+			requireFailResolution(
+				t, resolution, ResultHtlcInvoiceTypeMismatch,
+			)
+		})
+	}
+}
+
+// TestUpdateLegacyValidatesKeysendRecord checks that the keysend record is
+// well-formed and corresponds to the payment hash.
+func TestUpdateLegacyValidatesKeysendRecord(t *testing.T) {
+	t.Parallel()
+
+	validPreimage := lntypes.Preimage{1}
+	invalidPreimage := lntypes.Preimage{2}
+
+	tests := []struct {
+		name           string
+		keysendRecord  []byte
+		expectFail     bool
+		expectedResult FailResolutionResult
+	}{
+		{
+			name:           "missing keysend",
+			expectFail:     true,
+			expectedResult: ResultAddressMismatch,
+		},
+		{
+			name:           "invalid keysend length",
+			keysendRecord:  []byte{1, 2, 3},
+			expectFail:     true,
+			expectedResult: ResultAddressMismatch,
+		},
+		{
+			name:           "wrong keysend preimage",
+			keysendRecord:  invalidPreimage[:],
+			expectFail:     true,
+			expectedResult: ResultAddressMismatch,
+		},
+		{
+			name:          "valid keysend",
+			keysendRecord: validPreimage[:],
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+
+			ctx, invoice := newLegacyUpdateTestContext(
+				t, ContractOpen,
+			)
+			ctx.hash = validPreimage.Hash()
+			ctx.customRecords = make(record.CustomSet)
+			invoice.Terms.PaymentPreimage = &validPreimage
+			invoice.Terms.Features = lnwire.NewFeatureVector(
+				lnwire.NewRawFeatureVector(
+					lnwire.TLVOnionPayloadRequired,
+					lnwire.PaymentAddrRequired,
+				),
+				lnwire.Features,
+			)
+
+			if test.keysendRecord != nil {
+				ctx.customRecords[record.KeySendType] =
+					test.keysendRecord
+			}
+
+			update, resolution, err := updateLegacy(ctx, invoice)
+			require.NoError(t, err)
+
+			if test.expectFail {
+				require.Nil(t, update)
+				requireFailResolution(
+					t, resolution, test.expectedResult,
+				)
+
+				return
+			}
+
+			require.NotNil(t, update)
+			requireSettleResolution(t, resolution, ResultSettled)
+		})
+	}
+}
+
+// newLegacyUpdateTestContext creates a minimal legacy invoice and update
+// context for exercising update selection and settlement outcomes.
+func newLegacyUpdateTestContext(t *testing.T,
+	state ContractState) (*invoiceUpdateCtx, *Invoice) {
+
+	t.Helper()
+
+	preimage := lntypes.Preimage{1}
+	payHash := preimage.Hash()
+
+	ctx := &invoiceUpdateCtx{
+		hash:                 payHash,
+		circuitKey:           CircuitKey{HtlcID: 1},
+		amtPaid:              lnwire.MilliSatoshi(1000),
+		expiry:               40,
+		currentHeight:        10,
+		finalCltvRejectDelta: 10,
+		customRecords:        make(record.CustomSet),
+		wireCustomRecords:    make(lnwire.CustomRecords),
+	}
+
+	invoice := &Invoice{
+		State: state,
+		Terms: ContractTerm{
+			FinalCltvDelta:  10,
+			PaymentPreimage: &preimage,
+			Value:           1000,
+			Features: lnwire.NewFeatureVector(
+				nil, lnwire.Features,
+			),
+		},
+		Htlcs: make(map[CircuitKey]*InvoiceHTLC),
+	}
+
+	return ctx, invoice
+}
+
+// requireFailResolution checks the resolution type and its reported outcome.
+func requireFailResolution(t *testing.T, resolution HtlcResolution,
+	expected FailResolutionResult) {
+
+	t.Helper()
+
+	failResolution, ok := resolution.(*HtlcFailResolution)
+	require.True(t, ok)
+	require.Equal(t, expected, failResolution.Outcome)
+}
+
+// requireSettleResolution checks the resolution type and its reported outcome.
+func requireSettleResolution(t *testing.T, resolution HtlcResolution,
+	expected SettleResolutionResult) {
+
+	t.Helper()
+
+	settleResolution, ok := resolution.(*HtlcSettleResolution)
+	require.True(t, ok)
+	require.Equal(t, expected, settleResolution.Outcome)
+}