killbill-memoizeit

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/PaymentProcessor.java b/payment/src/main/java/org/killbill/billing/payment/core/PaymentProcessor.java
index 7885616..77dac47 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/PaymentProcessor.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/PaymentProcessor.java
@@ -370,7 +370,7 @@ public class PaymentProcessor extends ProcessorBase {
             if (paymentStateContext.getTransactionId() != null || paymentStateContext.getPaymentTransactionExternalKey() != null) {
                 // If a transaction id or key is passed, we are maybe completing an existing transaction (unless a new key was provided)
                 final List<PaymentTransactionModelDao> paymentTransactionsForCurrentPayment = daoHelper.getPaymentDao().getTransactionsForPayment(paymentStateContext.getPaymentId(), paymentStateContext.getInternalCallContext());
-                PaymentTransactionModelDao transactionToComplete = findTransactionToCompleteAndRunSanityChecks(paymentModelDao, paymentTransactionsForCurrentPayment, paymentStateContext);
+                PaymentTransactionModelDao transactionToComplete = findTransactionToCompleteAndRunSanityChecks(paymentModelDao, paymentTransactionsForCurrentPayment, paymentStateContext, internalCallContext);
 
                 if (transactionToComplete != null) {
                     // For completion calls, always invoke the Janitor first to get the latest state. The state machine will then
@@ -415,7 +415,8 @@ public class PaymentProcessor extends ProcessorBase {
 
     private PaymentTransactionModelDao findTransactionToCompleteAndRunSanityChecks(final PaymentModelDao paymentModelDao,
                                                                                    final Iterable<PaymentTransactionModelDao> paymentTransactionsForCurrentPayment,
-                                                                                   final PaymentStateContext paymentStateContext) throws PaymentApiException {
+                                                                                   final PaymentStateContext paymentStateContext,
+                                                                                   final InternalCallContext internalCallContext) throws PaymentApiException {
         final Collection<PaymentTransactionModelDao> completionCandidates = new LinkedList<PaymentTransactionModelDao>();
         for (final PaymentTransactionModelDao paymentTransactionModelDao : paymentTransactionsForCurrentPayment) {
             // Check if we already have a transaction for that id or key
@@ -434,7 +435,20 @@ public class PaymentProcessor extends ProcessorBase {
 
             // Sanity: if we already have a transaction for that id or key, the transaction type must match
             if (paymentTransactionModelDao.getTransactionType() != paymentStateContext.getTransactionType()) {
-                throw new PaymentApiException(ErrorCode.PAYMENT_INVALID_OPERATION, paymentStateContext.getTransactionType(), paymentModelDao.getStateName());
+                throw new PaymentApiException(ErrorCode.PAYMENT_INVALID_PARAMETER, "transactionType", String.format("%s doesn't match existing transaction type %s", paymentStateContext.getTransactionType(), paymentTransactionModelDao.getTransactionType()));
+            }
+
+            // Sanity: verify we don't already have a successful transaction for that key (chargeback reversals are a bit special, it's the only transaction type we can revert)
+            if (paymentTransactionModelDao.getTransactionExternalKey().equals(paymentStateContext.getPaymentTransactionExternalKey()) &&
+                paymentTransactionModelDao.getTransactionStatus() == TransactionStatus.SUCCESS &&
+                paymentTransactionModelDao.getTransactionType() != TransactionType.CHARGEBACK) {
+                throw new PaymentApiException(ErrorCode.PAYMENT_ACTIVE_TRANSACTION_KEY_EXISTS, paymentStateContext.getPaymentTransactionExternalKey());
+            }
+
+            // Sanity: don't share keys across accounts
+            if (paymentTransactionModelDao.getTransactionExternalKey().equals(paymentStateContext.getPaymentTransactionExternalKey()) &&
+                !paymentTransactionModelDao.getAccountRecordId().equals(internalCallContext.getAccountRecordId())) {
+                throw new PaymentApiException(ErrorCode.PAYMENT_ACTIVE_TRANSACTION_KEY_EXISTS, paymentStateContext.getPaymentTransactionExternalKey());
             }
 
             // UNKNOWN transactions are potential candidates, we'll invoke the Janitor first though

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentAutomatonRunner.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentAutomatonRunner.java
index 0610adf..f3e464e 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentAutomatonRunner.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentAutomatonRunner.java
@@ -18,6 +18,7 @@
 package org.killbill.billing.payment.core.sm;
 
 import java.math.BigDecimal;
+import java.util.List;
 import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 
@@ -42,6 +43,7 @@ import org.killbill.billing.catalog.api.Currency;
 import org.killbill.billing.osgi.api.OSGIServiceRegistration;
 import org.killbill.billing.payment.api.PaymentApiException;
 import org.killbill.billing.payment.api.PluginProperty;
+import org.killbill.billing.payment.api.TransactionStatus;
 import org.killbill.billing.payment.api.TransactionType;
 import org.killbill.billing.payment.core.PaymentExecutors;
 import org.killbill.billing.payment.core.sm.payments.AuthorizeCompleted;
@@ -67,6 +69,7 @@ import org.killbill.billing.payment.core.sm.payments.VoidInitiated;
 import org.killbill.billing.payment.core.sm.payments.VoidOperation;
 import org.killbill.billing.payment.dao.PaymentDao;
 import org.killbill.billing.payment.dao.PaymentModelDao;
+import org.killbill.billing.payment.dao.PaymentTransactionModelDao;
 import org.killbill.billing.payment.dispatcher.PluginDispatcher;
 import org.killbill.billing.payment.plugin.api.PaymentPluginApi;
 import org.killbill.billing.util.callcontext.CallContext;
@@ -126,7 +129,7 @@ public class PaymentAutomatonRunner {
                                                         final CallContext callContext,
                                                         final InternalCallContext internalCallContext) throws PaymentApiException {
         // Retrieve the payment id from the payment external key if needed
-        final UUID effectivePaymentId = paymentId != null ? paymentId : retrievePaymentId(paymentExternalKey, internalCallContext);
+        final UUID effectivePaymentId = paymentId != null ? paymentId : retrievePaymentId(paymentExternalKey, paymentTransactionExternalKey, internalCallContext);
 
         return new PaymentStateContext(isApiPayment,
                                        effectivePaymentId,
@@ -245,12 +248,41 @@ public class PaymentAutomatonRunner {
         }
     }
 
-    private UUID retrievePaymentId(@Nullable final String paymentExternalKey, final InternalCallContext internalCallContext) {
-        if (paymentExternalKey == null) {
+    // TODO Could we cache these to avoid extra queries in PaymentAutomatonDAOHelper?
+    private UUID retrievePaymentId(@Nullable final String paymentExternalKey, @Nullable final String paymentTransactionExternalKey, final InternalCallContext internalCallContext) {
+        if (paymentExternalKey != null) {
+            final PaymentModelDao payment = paymentDao.getPaymentByExternalKey(paymentExternalKey, internalCallContext);
+            if (payment != null) {
+                return payment.getId();
+            }
+        }
+
+        if (paymentTransactionExternalKey == null) {
             return null;
         }
 
-        final PaymentModelDao payment = paymentDao.getPaymentByExternalKey(paymentExternalKey, internalCallContext);
-        return payment == null ? null : payment.getId();
+        final List<PaymentTransactionModelDao> paymentTransactionModelDaos = paymentDao.getPaymentTransactionsByExternalKey(paymentTransactionExternalKey, internalCallContext);
+        for (final PaymentTransactionModelDao paymentTransactionModelDao : paymentTransactionModelDaos) {
+            if (paymentTransactionModelDao.getTransactionStatus() == TransactionStatus.SUCCESS ||
+                paymentTransactionModelDao.getTransactionStatus() == TransactionStatus.PENDING ||
+                paymentTransactionModelDao.getTransactionStatus() == TransactionStatus.UNKNOWN) {
+                return paymentTransactionModelDao.getPaymentId();
+            }
+        }
+
+        UUID paymentIdCandidate = null;
+        for (final PaymentTransactionModelDao paymentTransactionModelDao : paymentTransactionModelDaos) {
+            if (paymentTransactionModelDao.getTransactionStatus() == TransactionStatus.PAYMENT_FAILURE ||
+                paymentTransactionModelDao.getTransactionStatus() == TransactionStatus.PLUGIN_FAILURE) {
+                if (paymentIdCandidate == null) {
+                    paymentIdCandidate = paymentTransactionModelDao.getPaymentId();
+                } else if (!paymentIdCandidate.equals(paymentTransactionModelDao.getPaymentId())) {
+                    // Multiple failed payments sharing the key
+                    return null;
+                }
+            }
+        }
+
+        return paymentIdCandidate;
     }
 }

diff --git a/payment/src/test/java/org/killbill/billing/payment/api/TestPaymentApi.java b/payment/src/test/java/org/killbill/billing/payment/api/TestPaymentApi.java
index 8b1fee4..8c8ff6e 100644
--- a/payment/src/test/java/org/killbill/billing/payment/api/TestPaymentApi.java
+++ b/payment/src/test/java/org/killbill/billing/payment/api/TestPaymentApi.java
@@ -1739,7 +1739,7 @@ public class TestPaymentApi extends PaymentTestSuiteWithEmbeddedDB {
             createPayment(TransactionType.PURCHASE, null, paymentExternalKey, transactionExternalKey, requestedAmount, PaymentPluginStatus.PENDING);
             Assert.fail("PURCHASE transaction with same key should have failed");
         } catch (final PaymentApiException expected) {
-            Assert.assertEquals(expected.getCode(), ErrorCode.PAYMENT_INVALID_OPERATION.getCode());
+            Assert.assertEquals(expected.getCode(), ErrorCode.PAYMENT_INVALID_PARAMETER.getCode());
         }
     }
 
@@ -1770,7 +1770,7 @@ public class TestPaymentApi extends PaymentTestSuiteWithEmbeddedDB {
                 createPayment(transactionType, processedPayment.getId(), paymentExternalKey, keyA, requestedAmount, PaymentPluginStatus.PROCESSED);
                 Assert.fail("Retrying initial successful transaction (AUTHORIZE, PURCHASE, CREDIT) with same transaction key should fail");
             } catch (final PaymentApiException e) {
-                Assert.assertEquals(e.getCode(), ErrorCode.PAYMENT_INVALID_OPERATION.getCode());
+                Assert.assertEquals(e.getCode(), ErrorCode.PAYMENT_ACTIVE_TRANSACTION_KEY_EXISTS.getCode());
             }
         }
     }
@@ -1847,6 +1847,156 @@ public class TestPaymentApi extends PaymentTestSuiteWithEmbeddedDB {
         }
     }
 
+    @Test(groups = "slow")
+    public void testKeysSanityOnPending() throws Exception {
+        final String authKey = UUID.randomUUID().toString();
+        final Payment pendingAuthorization = createPayment(TransactionType.AUTHORIZE, null, null, authKey, BigDecimal.TEN, PaymentPluginStatus.PENDING);
+        assertNotNull(pendingAuthorization);
+        Assert.assertEquals(pendingAuthorization.getTransactions().size(), 1);
+        Assert.assertEquals(pendingAuthorization.getTransactions().get(0).getTransactionStatus(), TransactionStatus.PENDING);
+
+        try {
+            // Capture with the same external key should fail
+            createPayment(TransactionType.CAPTURE, pendingAuthorization.getId(), null, authKey, BigDecimal.TEN, PaymentPluginStatus.PROCESSED);
+            Assert.fail();
+        } catch (final PaymentApiException e) {
+            Assert.assertEquals(e.getCode(), ErrorCode.PAYMENT_INVALID_PARAMETER.getCode());
+        }
+
+        final Account account1 = testHelper.createTestAccount("bobo2@gmail.com", true);
+        try {
+            // Different auth with the same external key on a different account should fail
+            createPayment(account1, TransactionType.AUTHORIZE, null, null, authKey, BigDecimal.TEN, PaymentPluginStatus.PROCESSED);
+            Assert.fail();
+        } catch (final PaymentApiException e) {
+            Assert.assertEquals(e.getCode(), ErrorCode.PAYMENT_ACTIVE_TRANSACTION_KEY_EXISTS.getCode());
+        }
+
+        // Auth with the same external key should go through (completion)
+        final Payment authorization = createPayment(TransactionType.AUTHORIZE, null, null, authKey, BigDecimal.TEN, PaymentPluginStatus.PROCESSED);
+        assertNotNull(authorization);
+        Assert.assertEquals(authorization.getTransactions().size(), 1);
+        Assert.assertEquals(authorization.getTransactions().get(0).getTransactionStatus(), TransactionStatus.SUCCESS);
+
+        try {
+            // Different auth with the same external key on a different account should still fail
+            createPayment(account1, TransactionType.AUTHORIZE, null, null, authKey, BigDecimal.TEN, PaymentPluginStatus.PROCESSED);
+            Assert.fail();
+        } catch (final PaymentApiException e) {
+            Assert.assertEquals(e.getCode(), ErrorCode.PAYMENT_ACTIVE_TRANSACTION_KEY_EXISTS.getCode());
+        }
+
+        // Capture with a different key should go through
+        final String captureKey = UUID.randomUUID().toString();
+        final Payment pendingCapture = createPayment(TransactionType.CAPTURE, authorization.getId(), null, captureKey, BigDecimal.ONE, PaymentPluginStatus.PENDING);
+        Assert.assertEquals(pendingCapture.getTransactions().size(), 2);
+        Assert.assertEquals(pendingCapture.getTransactions().get(0).getTransactionStatus(), TransactionStatus.SUCCESS);
+        Assert.assertEquals(pendingCapture.getTransactions().get(1).getTransactionStatus(), TransactionStatus.PENDING);
+
+        try {
+            // Different auth with the same external key should fail
+            createPayment(TransactionType.AUTHORIZE, null, null, captureKey, BigDecimal.TEN, PaymentPluginStatus.PROCESSED);
+            Assert.fail();
+        } catch (final PaymentApiException e) {
+            Assert.assertEquals(e.getCode(), ErrorCode.PAYMENT_INVALID_PARAMETER.getCode());
+        }
+
+        try {
+            // Different auth with the same external key on a different account should fail
+            createPayment(account1, TransactionType.AUTHORIZE, null, null, captureKey, BigDecimal.TEN, PaymentPluginStatus.PROCESSED);
+            Assert.fail();
+        } catch (final PaymentApiException e) {
+            Assert.assertEquals(e.getCode(), ErrorCode.PAYMENT_INVALID_PARAMETER.getCode());
+        }
+
+        // Second capture with the same external key should go through (completion)
+        final Payment capturedPayment = createPayment(TransactionType.CAPTURE, authorization.getId(), null, captureKey, BigDecimal.ONE, PaymentPluginStatus.PROCESSED);
+        Assert.assertEquals(capturedPayment.getTransactions().size(), 2);
+        Assert.assertEquals(capturedPayment.getTransactions().get(0).getTransactionStatus(), TransactionStatus.SUCCESS);
+        Assert.assertEquals(capturedPayment.getTransactions().get(1).getTransactionStatus(), TransactionStatus.SUCCESS);
+
+        // Second capture with a different key should go through
+        final String captureKey2 = UUID.randomUUID().toString();
+        final Payment capturedPayment2 = createPayment(TransactionType.CAPTURE, authorization.getId(), null, captureKey2, BigDecimal.ONE, PaymentPluginStatus.PROCESSED);
+        Assert.assertEquals(capturedPayment2.getTransactions().size(), 3);
+        Assert.assertEquals(capturedPayment2.getTransactions().get(0).getTransactionStatus(), TransactionStatus.SUCCESS);
+        Assert.assertEquals(capturedPayment2.getTransactions().get(1).getTransactionStatus(), TransactionStatus.SUCCESS);
+        Assert.assertEquals(capturedPayment2.getTransactions().get(2).getTransactionStatus(), TransactionStatus.SUCCESS);
+    }
+
+    @Test(groups = "slow")
+    public void testKeysSanityOnSuccess() throws Exception {
+        final String authKey = UUID.randomUUID().toString();
+        final Payment authorization = createPayment(TransactionType.AUTHORIZE, null, null, authKey, BigDecimal.TEN, PaymentPluginStatus.PROCESSED);
+        assertNotNull(authorization);
+        Assert.assertEquals(authorization.getTransactions().size(), 1);
+        Assert.assertEquals(authorization.getTransactions().get(0).getTransactionStatus(), TransactionStatus.SUCCESS);
+
+        try {
+            // Capture with the same external key should fail
+            createPayment(TransactionType.CAPTURE, authorization.getId(), null, authKey, BigDecimal.TEN, PaymentPluginStatus.PROCESSED);
+            Assert.fail();
+        } catch (final PaymentApiException e) {
+            Assert.assertEquals(e.getCode(), ErrorCode.PAYMENT_INVALID_PARAMETER.getCode());
+        }
+
+        try {
+            // Different auth with the same external key should fail
+            createPayment(TransactionType.AUTHORIZE, null, null, authKey, BigDecimal.TEN, PaymentPluginStatus.PROCESSED);
+            Assert.fail();
+        } catch (final PaymentApiException e) {
+            Assert.assertEquals(e.getCode(), ErrorCode.PAYMENT_ACTIVE_TRANSACTION_KEY_EXISTS.getCode());
+        }
+
+        final Account account1 = testHelper.createTestAccount("bobo2@gmail.com", true);
+        try {
+            // Different auth with the same external key on a different account should fail
+            createPayment(account1, TransactionType.AUTHORIZE, null, null, authKey, BigDecimal.TEN, PaymentPluginStatus.PROCESSED);
+            Assert.fail();
+        } catch (final PaymentApiException e) {
+            Assert.assertEquals(e.getCode(), ErrorCode.PAYMENT_ACTIVE_TRANSACTION_KEY_EXISTS.getCode());
+        }
+
+        // Capture with a different key should go through
+        final String captureKey = UUID.randomUUID().toString();
+        final Payment capturedPayment = createPayment(TransactionType.CAPTURE, authorization.getId(), null, captureKey, BigDecimal.ONE, PaymentPluginStatus.PROCESSED);
+        Assert.assertEquals(capturedPayment.getTransactions().size(), 2);
+        Assert.assertEquals(capturedPayment.getTransactions().get(0).getTransactionStatus(), TransactionStatus.SUCCESS);
+        Assert.assertEquals(capturedPayment.getTransactions().get(1).getTransactionStatus(), TransactionStatus.SUCCESS);
+
+        try {
+            // Second capture with the same external key should fail
+            createPayment(TransactionType.CAPTURE, authorization.getId(), null, captureKey, BigDecimal.ONE, PaymentPluginStatus.PROCESSED);
+            Assert.fail();
+        } catch (final PaymentApiException e) {
+            Assert.assertEquals(e.getCode(), ErrorCode.PAYMENT_ACTIVE_TRANSACTION_KEY_EXISTS.getCode());
+        }
+
+        try {
+            // Different auth with the same external key should fail
+            createPayment(TransactionType.AUTHORIZE, null, null, captureKey, BigDecimal.TEN, PaymentPluginStatus.PROCESSED);
+            Assert.fail();
+        } catch (final PaymentApiException e) {
+            Assert.assertEquals(e.getCode(), ErrorCode.PAYMENT_INVALID_PARAMETER.getCode());
+        }
+
+        try {
+            // Different auth with the same external key on a different account should fail
+            createPayment(account1, TransactionType.AUTHORIZE, null, null, captureKey, BigDecimal.TEN, PaymentPluginStatus.PROCESSED);
+            Assert.fail();
+        } catch (final PaymentApiException e) {
+            Assert.assertEquals(e.getCode(), ErrorCode.PAYMENT_INVALID_PARAMETER.getCode());
+        }
+
+        // Second capture with a different key should go through
+        final String captureKey2 = UUID.randomUUID().toString();
+        final Payment capturedPayment2 = createPayment(TransactionType.CAPTURE, authorization.getId(), null, captureKey2, BigDecimal.ONE, PaymentPluginStatus.PROCESSED);
+        Assert.assertEquals(capturedPayment2.getTransactions().size(), 3);
+        Assert.assertEquals(capturedPayment2.getTransactions().get(0).getTransactionStatus(), TransactionStatus.SUCCESS);
+        Assert.assertEquals(capturedPayment2.getTransactions().get(1).getTransactionStatus(), TransactionStatus.SUCCESS);
+        Assert.assertEquals(capturedPayment2.getTransactions().get(2).getTransactionStatus(), TransactionStatus.SUCCESS);
+    }
+
     private void verifyRefund(final Payment refund, final String paymentExternalKey, final String paymentTransactionExternalKey, final String refundTransactionExternalKey, final BigDecimal requestedAmount, final BigDecimal refundAmount, final TransactionStatus transactionStatus) {
         Assert.assertEquals(refund.getExternalKey(), paymentExternalKey);
         Assert.assertEquals(refund.getTransactions().size(), 2);
@@ -1939,6 +2089,16 @@ public class TestPaymentApi extends PaymentTestSuiteWithEmbeddedDB {
                                   @Nullable final String paymentTransactionExternalKey,
                                   @Nullable final BigDecimal amount,
                                   final PaymentPluginStatus paymentPluginStatus) throws PaymentApiException {
+        return createPayment(account, transactionType, paymentId, paymentExternalKey, paymentTransactionExternalKey, amount, paymentPluginStatus);
+    }
+
+    private Payment createPayment(final Account account,
+                                  final TransactionType transactionType,
+                                  @Nullable final UUID paymentId,
+                                  @Nullable final String paymentExternalKey,
+                                  @Nullable final String paymentTransactionExternalKey,
+                                  @Nullable final BigDecimal amount,
+                                  final PaymentPluginStatus paymentPluginStatus) throws PaymentApiException {
         final Iterable<PluginProperty> pluginProperties = ImmutableList.<PluginProperty>of(new PluginProperty(MockPaymentProviderPlugin.PLUGIN_PROPERTY_PAYMENT_PLUGIN_STATUS_OVERRIDE, paymentPluginStatus.toString(), false));
         switch (transactionType) {
             case AUTHORIZE: