From 8091baeb6ef2d4bd37d7845928bdaf516cc301d0 Mon Sep 17 00:00:00 2001
From: Timur Karimov <timur.karimow.95@gmail.com>
Date: Mon, 2 Dec 2024 16:56:35 +0100
Subject: [PATCH] Restrict Stripe Link to credit card payment method and
 improve cleanup (#9822)

Co-authored-by: Timur Karimov <timur@Timurs-MacBook-Pro.local>
---
 changelog/fix-stripe-link-button              |  4 +++
 client/checkout/blocks/payment-processor.js   | 22 ++++++++++++--
 client/checkout/stripe-link/index.js          | 19 ++++++++++--
 .../checkout/stripe-link/test/index.test.js   | 29 +++++++++++++++++++
 4 files changed, 69 insertions(+), 5 deletions(-)
 create mode 100644 changelog/fix-stripe-link-button

diff --git a/changelog/fix-stripe-link-button b/changelog/fix-stripe-link-button
new file mode 100644
index 00000000000..d8acf0626f1
--- /dev/null
+++ b/changelog/fix-stripe-link-button
@@ -0,0 +1,4 @@
+Significance: minor
+Type: fix
+
+Restrict Stripe Link to credit card payment method and improve cleanup.
diff --git a/client/checkout/blocks/payment-processor.js b/client/checkout/blocks/payment-processor.js
index ee7ef6af861..952470aa46b 100644
--- a/client/checkout/blocks/payment-processor.js
+++ b/client/checkout/blocks/payment-processor.js
@@ -28,7 +28,10 @@ import { useCustomerData } from './utils';
 import enableStripeLinkPaymentMethod from 'wcpay/checkout/stripe-link';
 import { getUPEConfig } from 'wcpay/utils/checkout';
 import { validateElements } from 'wcpay/checkout/classic/payment-processing';
-import { PAYMENT_METHOD_ERROR } from 'wcpay/checkout/constants';
+import {
+	PAYMENT_METHOD_ERROR,
+	PAYMENT_METHOD_NAME_CARD,
+} from 'wcpay/checkout/constants';
 
 const getBillingDetails = ( billingData ) => {
 	return {
@@ -70,6 +73,7 @@ const PaymentProcessor = ( {
 	const stripe = useStripe();
 	const elements = useElements();
 	const hasLoadErrorRef = useRef( false );
+	const linkCleanupRef = useRef( null );
 
 	const paymentMethodsConfig = getUPEConfig( 'paymentMethodsConfig' );
 	const isTestMode = getUPEConfig( 'testMode' );
@@ -81,7 +85,10 @@ const PaymentProcessor = ( {
 	} = useCustomerData();
 
 	useEffect( () => {
-		if ( isLinkEnabled( paymentMethodsConfig ) ) {
+		if (
+			activePaymentMethod === PAYMENT_METHOD_NAME_CARD &&
+			isLinkEnabled( paymentMethodsConfig )
+		) {
 			enableStripeLinkPaymentMethod( {
 				api: api,
 				elements: elements,
@@ -123,11 +130,22 @@ const PaymentProcessor = ( {
 					} );
 				},
 				onButtonShow: blocksShowLinkButtonHandler,
+			} ).then( ( cleanup ) => {
+				linkCleanupRef.current = cleanup;
 			} );
+
+			// Cleanup the Link button when the component unmounts
+			return () => {
+				if ( linkCleanupRef.current ) {
+					linkCleanupRef.current();
+					linkCleanupRef.current = null;
+				}
+			};
 		}
 	}, [
 		api,
 		elements,
+		activePaymentMethod,
 		paymentMethodsConfig,
 		setBillingAddress,
 		setShippingAddress,
diff --git a/client/checkout/stripe-link/index.js b/client/checkout/stripe-link/index.js
index 45f4feabf36..e44cf0d8236 100644
--- a/client/checkout/stripe-link/index.js
+++ b/client/checkout/stripe-link/index.js
@@ -4,6 +4,13 @@
 import { dispatchChangeEventFor } from '../utils/upe';
 
 export const switchToNewPaymentTokenElement = () => {
+	// Switch to card payment method before enabling new payment token element
+	document
+		.querySelector(
+			'input[name="payment_method"][value="woocommerce_payments"]'
+		)
+		?.click();
+
 	const newPaymentTokenElement = document.getElementById(
 		'wc-woocommerce_payments-payment-token-new'
 	);
@@ -44,16 +51,17 @@ const enableStripeLinkPaymentMethod = async ( options ) => {
 	const emailField = document.getElementById( options.emailId );
 
 	if ( ! emailField ) {
-		return;
+		return Promise.resolve( () => null );
 	}
 
 	const stripe = await options.api.getStripe();
 	// https://stripe.com/docs/payments/link/autofill-modal
 	const linkAutofill = stripe.linkAutofillModal( options.elements );
 
-	emailField.addEventListener( 'keyup', ( event ) => {
+	const handleKeyup = ( event ) => {
 		linkAutofill.launch( { email: event.target.value } );
-	} );
+	};
+	emailField.addEventListener( 'keyup', handleKeyup );
 
 	options.onButtonShow( linkAutofill );
 
@@ -65,6 +73,11 @@ const enableStripeLinkPaymentMethod = async ( options ) => {
 		);
 		switchToNewPaymentTokenElement();
 	} );
+
+	return () => {
+		emailField.removeEventListener( 'keyup', handleKeyup );
+		removeLinkButton();
+	};
 };
 
 export default enableStripeLinkPaymentMethod;
diff --git a/client/checkout/stripe-link/test/index.test.js b/client/checkout/stripe-link/test/index.test.js
index b8d907c3508..392a170b179 100644
--- a/client/checkout/stripe-link/test/index.test.js
+++ b/client/checkout/stripe-link/test/index.test.js
@@ -85,6 +85,35 @@ describe( 'Stripe Link elements behavior', () => {
 		expect( handleButtonShow ).toHaveBeenCalled();
 	} );
 
+	test( 'Should properly clean up when cleanup function is called', async () => {
+		createStripeLinkElements();
+		const billingEmailInput = document.getElementById( 'billing_email' );
+		const removeEventListenerSpy = jest.spyOn(
+			billingEmailInput,
+			'removeEventListener'
+		);
+		const removeLinkButtonSpy = jest.spyOn(
+			document.querySelector( '.wcpay-stripelink-modal-trigger' ),
+			'remove'
+		);
+
+		const cleanup = await enableStripeLinkPaymentMethod( {
+			api: WCPayAPI(),
+			emailId: 'billing_email',
+			onAutofill: () => null,
+			onButtonShow: () => null,
+		} );
+
+		// Call the cleanup function
+		cleanup();
+
+		expect( removeEventListenerSpy ).toHaveBeenCalledWith(
+			'keyup',
+			expect.any( Function )
+		);
+		expect( removeLinkButtonSpy ).toHaveBeenCalled();
+	} );
+
 	function createStripeLinkElements() {
 		// Create the input field
 		const billingEmailInput = document.createElement( 'input' );