fix(remote): retry once on a proxy 502/503/504 page and keep server errors from ending a relicense

tests/remote/unit/test_remote_client.py

@@ -85,6 +85,42 @@ def test_4xx_raises_systemexit_with_logged_error(client: RemoteClient, caplog: p
     assert "no such session" in caplog.text
 
 
+PROXY_502 = "<html><head><title>unshackle.dev | 502: Bad gateway</title></head></html>"
+
+
+@responses.activate
+def test_proxy_502_page_is_retried_once(client: RemoteClient, monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr("unshackle.core.remote_service.time.sleep", lambda _s: None)
+    url = "http://srv:8786/api/session/abc/license"
+    responses.add(responses.POST, url, body=PROXY_502, status=502, content_type="text/html")
+    responses.add(responses.POST, url, json={"keys": {}}, status=200)
+    assert client.post("/api/session/abc/license", {}) == {"keys": {}}
+    assert len(responses.calls) == 2
+
+
+@responses.activate
+def test_second_proxy_502_page_exits(client: RemoteClient, monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr("unshackle.core.remote_service.time.sleep", lambda _s: None)
+    url = "http://srv:8786/api/session/abc/license"
+    responses.add(responses.POST, url, body=PROXY_502, status=502, content_type="text/html")
+    with pytest.raises(SystemExit):
+        client.post("/api/session/abc/license", {})
+    assert len(responses.calls) == 2
+
+
+@responses.activate
+def test_server_json_503_is_not_retried(client: RemoteClient) -> None:
+    responses.add(
+        responses.GET,
+        "http://srv:8786/api/services",
+        json={"error_code": "SERVICE_UNAVAILABLE", "message": "busy"},
+        status=503,
+    )
+    with pytest.raises(SystemExit):
+        client.get("/api/services")
+    assert len(responses.calls) == 1
+
+
 @responses.activate
 def test_401_falls_back_to_next_auth_header(client: RemoteClient) -> None:
     seen: list[tuple[str | None, str | None]] = []

unshackle/commands/dl.py

@@ -4877,6 +4877,8 @@ class dl:
                     return not drm.content_keys or bool(track_kid and track_kid not in drm.content_keys)
 
                 if missing_track_key():
+                    from unshackle.core.remote_service import RemoteError
+
                     try:
                         licence(
                             drm_system="playready" if drm.__class__.__name__ == "PlayReady" else "widevine",
@@ -4884,7 +4886,7 @@ class dl:
                         )
                     except click.ClickException:
                         raise
-                    except Exception as e:
+                    except (Exception, RemoteError) as e:
                         self.log.debug(f"Server CDM licence with an empty challenge failed: {e!r}")
 
                 server_vault_keys = {} if cdm_only else getattr(svc_for_cdm, "server_vault_keys", {})

unshackle/core/remote_service.py

@@ -144,6 +144,7 @@ class RemoteClient:
         payload: Dict[str, Any] = (
             {"json": data} if body is None else {"data": body, "headers": {"Content-Type": "application/octet-stream"}}
         )
+        gateway_retried = False
         while True:
             try:
                 with self.in_flight(endpoint):
@@ -160,6 +161,16 @@ class RemoteClient:
             # servers differ on which header carries the key, so retry the rest before giving up
             if resp.status_code == 401 and self.next_auth_header():
                 continue
+            # a proxy HTML page, not the server's JSON, means a dropped tunnel stream, which one retry gets past
+            if (
+                resp.status_code in (502, 503, 504)
+                and not gateway_retried
+                and "json" not in resp.headers.get("Content-Type", "")
+            ):
+                gateway_retried = True
+                log.debug(f"HTTP {resp.status_code} from the proxy in front of the server on {endpoint}, retrying once")
+                time.sleep(1)
+                continue
             break
         if resp.status_code < 400 and "octet-stream" in resp.headers.get("Content-Type", ""):
             return {"body": resp.content}
@@ -1870,7 +1881,7 @@ class RemoteService:
                                 for kid_hex, key_hex in keys.items():
                                     drm_obj.content_keys[UUID(hex=kid_hex)] = key_hex
                         return challenge
-                except Exception as e:
+                except (Exception, RemoteError) as e:
                     self.log.warning("server_cdm license failed: %s", e)
             else:
                 self.log.warning(f"Track {track.id} has no PSSH to send to the server CDM")