Add withdraw features

bifrurcated · Jan 18, 2024 · e9a7a7f · e9a7a7f
1 parent 03bddaa
commit e9a7a7f
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 11 deletions.
diff --git a/src/main/java/com/bifurcated/wallet/controller/WalletEntityController.java b/src/main/java/com/bifurcated/wallet/controller/WalletEntityController.java
@@ -1,15 +1,12 @@
 package com.bifurcated.wallet.controller;
 
-import com.bifurcated.wallet.errors.UnsupportedOperationTypeError;
 import com.bifurcated.wallet.operation.OperationType;
 import com.bifurcated.wallet.service.WalletEntityService;
-import com.bifurcated.wallet.service.WalletService;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.web.bind.annotation.*;
 
 import java.util.Map;
-import java.util.Optional;
 import java.util.UUID;
 
 @RestController
@@ -29,16 +26,20 @@ public WalletEntityController(WalletEntityService walletEntityService, Map<Strin
     public record WalletResponse(UUID id, Float amount){}
     public record WalletRequest(
             @JsonProperty("valletId") UUID walletId,
-            String operationType,
             Float amount
     ){}
-    @PostMapping("/wallet")
-    public WalletResponse wallet(@RequestBody WalletRequest request) {
-        var operationType = request.operationType().toLowerCase();
+    @PostMapping("/wallet/add")
+    public WalletResponse walletAdd(@RequestBody WalletRequest request) {
         var wallet = walletService.addAmount(request.walletId(), request.amount());
         return new WalletResponse(wallet.getId(), wallet.getAmount());
     }
 
+    @PostMapping("/wallet/reduce")
+    public WalletResponse walletReduce(@RequestBody WalletRequest request) {
+        var wallet = walletService.reduceAmount(request.walletId(), request.amount());
+        return new WalletResponse(wallet.getId(), wallet.getAmount());
+    }
+
     public record BalanceResponse(Float amount){}
     @GetMapping("/wallets/{WALLET_UUID}")
     public BalanceResponse balance(@PathVariable(value = "WALLET_UUID") UUID id) {

diff --git a/src/main/java/com/bifurcated/wallet/service/WalletEntityService.java b/src/main/java/com/bifurcated/wallet/service/WalletEntityService.java
@@ -1,6 +1,7 @@
 package com.bifurcated.wallet.service;
 
 import com.bifurcated.wallet.data.WalletEntity;
+import com.bifurcated.wallet.errors.NotEnoughMoneyError;
 import com.bifurcated.wallet.errors.WalletNotFoundError;
 import com.bifurcated.wallet.repository.WalletEntityRepository;
 import jakarta.persistence.EntityManager;
@@ -37,6 +38,18 @@ public WalletEntity addAmountUsingUpdate(UUID id, Float amount) {
         return repository.findById(id).orElseThrow(WalletNotFoundError::new);
     }
 
+    @Transactional
+    public WalletEntity reduceAmount(UUID id, Float amount) {
+        var entity = entityManager.find(WalletEntity.class, id, LockModeType.PESSIMISTIC_WRITE);
+        var wallet = Optional.ofNullable(entity).orElseThrow(WalletNotFoundError::new);
+        var reduce = wallet.getAmount() - amount;
+        if (reduce < 0) {
+            throw new NotEnoughMoneyError(wallet.getAmount(), amount);
+        }
+        wallet.setAmount(reduce);
+        return repository.save(wallet);
+    }
+
     public Float amount(UUID id) {
         return repository.findById(id).orElseThrow(WalletNotFoundError::new).getAmount();
     }

diff --git a/src/test/java/com/bifurcated/wallet/controller/WalletEntityControllerTest.java b/src/test/java/com/bifurcated/wallet/controller/WalletEntityControllerTest.java
@@ -1,9 +1,7 @@
 package com.bifurcated.wallet.controller;
 
-import com.bifurcated.wallet.data.Wallet;
 import com.bifurcated.wallet.data.WalletEntity;
 import com.bifurcated.wallet.repository.WalletEntityRepository;
-import com.bifurcated.wallet.repository.WalletRepo;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -19,7 +17,6 @@
 import org.springframework.test.web.servlet.ResultActions;
 
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.List;
 import java.util.UUID;
 import java.util.concurrent.Callable;
@@ -102,7 +99,7 @@ record BalanceResponse(Float amount){}
         List<Future<ResultActions>> futures = new ArrayList<>();
         for (int i = 0; i < 300; i++) {
             Callable<ResultActions> callable = () -> {
-                ResultActions perform = this.mockMvc.perform(post(END_POINT_PATH + "/wallet")
+                ResultActions perform = this.mockMvc.perform(post(END_POINT_PATH + "/wallet/add")
                         .contentType(MediaType.APPLICATION_JSON)
                         .content(requestBody));
                 return perform.andExpect(status().isOk());
@@ -123,4 +120,39 @@ record BalanceResponse(Float amount){}
                 .andExpect(content().json(response));
     }
 
+    @Test
+    public void testWalletWithdrawDDOSAttack() throws Exception {
+        record WalletRequest(UUID valletId, String operationType, Float amount){}
+        record BalanceResponse(Float amount){}
+        UUID id = UUID.fromString("9ebef4de-68e3-43ad-a812-42193919ff02");
+        WalletRequest walletRequest = new WalletRequest(
+                id, "WITHDRAW", 1000F);
+
+        String requestBody = objectMapper.writeValueAsString(walletRequest);
+
+        ExecutorService executorService = Executors.newFixedThreadPool(300);
+        List<Future<ResultActions>> futures = new ArrayList<>();
+        for (int i = 0; i < 300; i++) {
+            Callable<ResultActions> callable = () -> {
+                ResultActions perform = this.mockMvc.perform(post(END_POINT_PATH + "/wallet/reduce")
+                        .contentType(MediaType.APPLICATION_JSON)
+                        .content(requestBody));
+                return perform.andExpect(status().isOk());
+            };
+
+            Future<ResultActions> future = executorService.submit(callable);
+            futures.add(future);
+        }
+        for (Future<ResultActions> future : futures) {
+            future.get();
+        }
+
+        BalanceResponse balanceResponse = new BalanceResponse(2000F);
+        String response = objectMapper.writeValueAsString(balanceResponse);
+        this.mockMvc.perform(get(END_POINT_PATH+"/wallets/"+id))
+                .andDo(print())
+                .andExpect(status().isOk())
+                .andExpect(content().json(response));
+    }
+
 }