Multiple download protection

Author: codeSamuraii

lib/store.py

@@ -22,6 +22,7 @@ def __init__(self, transfer_id: str):
         self._k_queue = self.key('queue')
         self._k_meta = self.key('metadata')
         self._k_cleanup = f'cleanup:{transfer_id}'
+        self._k_receiver_connected = self.key('receiver_connected')
 
     @classmethod
     def get_redis(cls) -> redis.Redis:
@@ -121,6 +122,17 @@ async def get_metadata(self) -> str | None:
 
     ## Transfer state operations ##
 
+    async def set_receiver_connected(self) -> bool:
+        """
+        Mark that a receiver has connected for this transfer.
+        Returns True if the flag was set, False if it was already created.
+        """
+        return bool(await self.redis.set(self._k_receiver_connected, '1', ex=300, nx=True))
+
+    async def is_receiver_connected(self) -> bool:
+        """Check if a receiver has already connected."""
+        return await self.redis.exists(self._k_receiver_connected) > 0
+
     async def set_completed(self) -> None:
         """Mark the transfer as completed."""
         await self.redis.set(f'completed:{self.transfer_id}', '1', ex=300, nx=True)

lib/transfer.py

@@ -58,6 +58,12 @@ async def wait_for_client_connected(self):
         await self.wait_for_event('client_connected')
         self.debug(f"△ Received client connected notification.")
 
+    async def is_receiver_connected(self) -> bool:
+        return await self.store.is_receiver_connected()
+
+    async def set_receiver_connected(self) -> bool:
+        return await self.store.set_receiver_connected()
+
     async def is_interrupted(self) -> bool:
         return await self.store.is_interrupted()
 

static/download.html

@@ -38,5 +38,12 @@ <h2>Ready to download</h2>
             <strong>©</strong> <a href="https://github.com/codeSamuraii/">Rémi Héneault</a></p>
         </footer>
     </div>
+    <script>
+        document.getElementById('download-button').addEventListener('click', function() {{
+            this.classList.add('disabled');
+            this.textContent = 'Downloading...';
+            this.style.pointerEvents = 'none';
+        }});
+    </script>
 </body>
 </html>

tests/test_endpoints.py

@@ -34,6 +34,8 @@ def mock_redis():
     """Mock Redis client for testing."""
     with patch('lib.store.Store.get_redis') as mock:
         redis_mock = AsyncMock()
+        # Default for receiver connected check
+        redis_mock.exists.return_value = 0
         mock.return_value = redis_mock
         yield redis_mock
 
@@ -146,7 +148,9 @@ async def test_http_download_success(self, test_client, test_file, mock_redis):
 
         # Mock Redis operations for metadata retrieval
         mock_redis.get.return_value = file_metadata.to_json()
-        mock_redis.set.return_value = None
+        # Mock receiver not connected, then successfully set
+        mock_redis.exists.return_value = 0
+        mock_redis.set.return_value = True
         mock_redis.publish.return_value = None
         mock_redis.scan.return_value = (0, [])
         mock_redis.delete.return_value = 1
@@ -197,6 +201,7 @@ async def test_http_download_prefetch_protection(self, test_client, mock_redis):
         )
 
         mock_redis.get.return_value = file_metadata.to_json()
+        mock_redis.exists.return_value = 0  # Receiver not connected
 
         # Simulate WhatsApp prefetch request
         headers = {"user-agent": "WhatsApp/2.21.1"}
@@ -207,6 +212,24 @@ async def test_http_download_prefetch_protection(self, test_client, mock_redis):
         if response.status_code == 200:
             assert "text/html" in response.headers.get("content-type", "") or "test.txt" in response.text
 
+    @pytest.mark.asyncio
+    async def test_http_download_already_connected(self, test_client, mock_redis):
+        """Test that a second download attempt is rejected."""
+        uid = "test-already-connected"
+        file_metadata = FileMetadata(
+            name="test.txt",
+            size=1000,
+            content_type="text/plain"
+        )
+
+        mock_redis.get.return_value = file_metadata.to_json()
+        mock_redis.exists.return_value = 1  # Simulate receiver already connected
+
+        response = test_client.get(f"/{uid}?download=true")
+
+        assert response.status_code == 409
+        assert "A client is already downloading this file" in response.text
+
 
 class TestHTTPUpload:
     """Test HTTP upload endpoint."""
@@ -320,7 +343,7 @@ async def mock_brpop(keys, timeout=None):
             content_type="text/plain"
         ).to_json()
         mock_redis.llen.return_value = 0
-        mock_redis.exists.return_value = False
+        mock_redis.exists.return_value = 0
         mock_redis.publish.return_value = None
         mock_redis.scan.return_value = (0, [])
         mock_redis.delete.return_value = 0
@@ -351,6 +374,8 @@ def upload_file():
         # Step 2: Download via HTTP
         def download_file():
             client = TestClient(app)
+            # Mock set_receiver_connected to succeed
+            mock_redis.set.return_value = True
             response = client.get(f"/{uid}?download=true")
             assert response.status_code == 200
             return response.content

views/http.py

@@ -115,6 +115,9 @@ async def http_download(request: Request, uid: str):
     else:
         log.info(f"▼ HTTP download request for: {transfer.file}")
 
+    if await transfer.is_receiver_connected():
+        raise HTTPException(status_code=409, detail="A client is already downloading this file.")
+
     file_name, file_size, file_type = transfer.get_file_info()
     user_agent = request.headers.get('user-agent', '').lower()
     is_prefetcher = any(prefetch_ua in user_agent for prefetch_ua in PREFETCHER_USER_AGENTS)
@@ -135,6 +138,9 @@ async def http_download(request: Request, uid: str):
         )
         return HTMLResponse(content=html_download, status_code=200)
 
+    if not await transfer.set_receiver_connected():
+        raise HTTPException(status_code=409, detail="A client is already downloading this file.")
+
     await transfer.set_client_connected()
 
     transfer.info("▼ Starting download...")

views/websockets.py

@@ -78,6 +78,11 @@ async def websocket_download(background_tasks: BackgroundTasks, websocket: WebSo
         await websocket.send_text("File not found")
         return
 
+    if await transfer.is_receiver_connected():
+        log.warning("▼ A client is already downloading this file.")
+        await websocket.send_text("Error: A client is already downloading this file.")
+        return
+
     file_name, file_size, file_type = transfer.get_file_info()
     transfer.debug(f"▼ File: name={file_name}, size={file_size}, type={file_type}")
     await websocket.send_json({'file_name': file_name, 'file_size': file_size, 'file_type': file_type})
@@ -93,6 +98,11 @@ async def websocket_download(background_tasks: BackgroundTasks, websocket: WebSo
             transfer.warning("▼ Client disconnected while waiting for go-ahead")
             return
 
+    if not await transfer.set_receiver_connected():
+        log.warning("▼ A client is already downloading this file.")
+        await websocket.send_text("Error: A client is already downloading this file.")
+        return
+
     transfer.info("▼ Notifying client is connected.")
     await transfer.set_client_connected()
     background_tasks.add_task(transfer.finalize_download)