Fix withdraw increment (#26)

* Fix: made increment/is_spent more specific
* fix: make safe `increment_withdraw_link `
* fix: remove old fix
* chore: code format

---------

Co-authored-by: benarc <ben@arc.wales>

helpers.py

@@ -0,0 +1,49 @@
+from threading import Lock
+from typing import Dict
+
+
+class CounterLock:
+    def __init__(self):
+        self.counter = 0
+        self.lock = Lock()
+
+    def acquire(self) -> bool:
+        self.counter += 1
+        return self.lock.acquire()
+
+    def release(self) -> None:
+        self.counter -= 1
+        return self.lock.release()
+
+    @property
+    def no_more_waiters(self) -> bool:
+        return self.counter == 0
+
+
+class NamedLock:
+    _lock = Lock()
+    _locks: Dict[str, CounterLock] = {}
+
+    def acquire(self, name: str) -> bool:
+        self._lock.acquire()
+
+        if name not in self._locks:
+            self._locks[name] = CounterLock()
+
+        self._lock.release()
+
+        return self._locks[name].acquire()
+
+
+
+    def release(self, name: str):
+        self._lock.acquire()
+
+        if name not in self._locks:
+            return self._lock.release()
+
+        self._locks[name].release()
+        if self._locks[name].no_more_waiters:
+            del self._locks[name]
+
+        return self._lock.release()

lnurl.py

@@ -19,7 +19,9 @@ from .crud import (
     remove_unique_withdraw_link,
 )
 from .models import WithdrawLink
+from .helpers import NamedLock
 
+withdraw_lock = NamedLock()
 
 @withdraw_ext.get(
     "/api/v1/lnurl/{unique_hash}",
@@ -81,6 +83,31 @@ async def api_lnurl_callback(
     pr: str = Query(...),
     id_unique_hash=None,
 ):
+    link = await _check_withdraw_link_safe(unique_hash, k1, id_unique_hash)
+
+    try:
+        payment_hash = await pay_invoice(
+            wallet_id=link.wallet,
+            payment_request=pr,
+            max_sat=link.max_withdrawable,
+            extra={"tag": "withdraw", "withdrawal_link_id": link.id},
+        )
+        if link.webhook_url:
+            await dispatch_webhook(link, payment_hash, pr)
+        return {"status": "OK"}
+    except Exception as e:
+        raise HTTPException(
+            status_code=HTTPStatus.BAD_REQUEST, detail=f"withdraw not working. {str(e)}"
+        )
+
+async def _check_withdraw_link_safe(unique_hash, k1, id_unique_hash) -> WithdrawLink:
+    try:
+        withdraw_lock.acquire(unique_hash)
+        return await _check_withdraw_link(unique_hash, k1, id_unique_hash)
+    finally:
+        withdraw_lock.release(unique_hash)
+
+async def _check_withdraw_link(unique_hash, k1, id_unique_hash) -> WithdrawLink:
     link = await get_withdraw_link_by_hash(unique_hash)
     now = int(datetime.now().timestamp())
     if not link:
@@ -94,7 +121,9 @@ async def api_lnurl_callback(
         )
 
     if link.k1 != k1:
-        raise HTTPException(status_code=HTTPStatus.BAD_REQUEST, detail="k1 is wrong.")
+        raise HTTPException(
+            status_code=HTTPStatus.BAD_REQUEST, detail="k1 is wrong."
+        )
 
     if now < link.open_time:
         raise HTTPException(
@@ -109,22 +138,13 @@ async def api_lnurl_callback(
             raise HTTPException(
                 status_code=HTTPStatus.NOT_FOUND, detail="withdraw not found."
             )
-
-    try:
     await increment_withdraw_link(link)
-        payment_hash = await pay_invoice(
+
-            wallet_id=link.wallet,
+    return link
-            payment_request=pr,
+
-            max_sat=link.max_withdrawable,
+
-            extra={"tag": "withdraw", "withdrawal_link_id": link.id},
+
-        )
+
-        if link.webhook_url:
-            await dispatch_webhook(link, payment_hash, pr)
-        return {"status": "OK"}
-    except Exception as e:
-        raise HTTPException(
-            status_code=HTTPStatus.BAD_REQUEST, detail=f"withdraw not working. {str(e)}"
-        )
 
 
 def check_unique_link(link: WithdrawLink, unique_hash: str) -> bool: