From 4af2b17a75aa6c1f9af80a3cc6280b256715bf61 Mon Sep 17 00:00:00 2001
From: WorldObservationLog <wolc@duck.com>
Date: Wed, 8 May 2024 19:23:25 +0800
Subject: [PATCH] feat: retry when raise DecryptException

---
 src/decrypt.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/src/decrypt.py b/src/decrypt.py
index cd765cb..bf5536f 100644
--- a/src/decrypt.py
+++ b/src/decrypt.py
@@ -1,6 +1,8 @@
 import asyncio
+import logging
 
 from loguru import logger
+from tenacity import retry, retry_if_exception_type, stop_after_attempt, before_sleep_log
 
 from src.adb import Device
 from src.exceptions import DecryptException
@@ -9,6 +11,8 @@ from src.mp4 import SongInfo, SampleInfo
 from src.types import defaultId, prefetchKey
 
 
+@retry(retry=retry_if_exception_type(DecryptException), stop=stop_after_attempt(3),
+       before_sleep=before_sleep_log(logger, logging.WARN))
 async def decrypt(info: SongInfo, keys: list[str], manifest: Datum, device: Device) -> bytes:
     async with device.decryptLock:
         logger.info(f"Decrypting song: {manifest.attributes.artistName} - {manifest.attributes.name}")
@@ -28,7 +32,12 @@ async def decrypt(info: SongInfo, keys: list[str], manifest: Datum, device: Devi
                 writer.write(bytes([len(key_uri)]))
                 writer.write(key_uri.encode("utf-8"))
             last_index = sample.descIndex
-            result = await decrypt_sample(writer, reader, sample)
+            try:
+                result = await decrypt_sample(writer, reader, sample)
+            except DecryptException as e:
+                writer.write(bytes([0, 0, 0, 0]))
+                writer.close()
+                raise e
             decrypted += result
         writer.write(bytes([0, 0, 0, 0]))
         writer.close()