From 8be71be540cd7fad1efa8b5e5853e27fcf2fb8d6 Mon Sep 17 00:00:00 2001
From: skshetry <18718008+skshetry@users.noreply.github.com>
Date: Tue, 4 Apr 2023 13:15:08 +0545
Subject: [PATCH] utils.studio: disallow redirects, misc API refactoring
 (#9301)

---
 dvc/utils/studio.py             | 16 +++++++++-------
 tests/unit/utils/test_studio.py |  1 +
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/dvc/utils/studio.py b/dvc/utils/studio.py
index 8da724ad06..b98166dc6d 100644
--- a/dvc/utils/studio.py
+++ b/dvc/utils/studio.py
@@ -15,21 +15,23 @@
 
 
 def post(
-    endpoint: str,
+    url: str,
     token: str,
     data: Dict[str, Any],
-    url: Optional[str] = STUDIO_URL,
+    base_url: Optional[str] = STUDIO_URL,
     max_retries: int = 3,
     timeout: int = 5,
 ) -> "Response":
-    endpoint = urljoin(url or STUDIO_URL, endpoint)
+    url = urljoin(base_url or STUDIO_URL, url)
     session = requests.Session()
-    session.mount(endpoint, HTTPAdapter(max_retries=max_retries))
+    session.mount(url, HTTPAdapter(max_retries=max_retries))
 
-    logger.trace("Sending %s to %s", data, endpoint)  # type: ignore[attr-defined]
+    logger.trace("Sending %s to %s", data, url)  # type: ignore[attr-defined]
 
     headers = {"Authorization": f"token {token}"}
-    r = session.post(endpoint, json=data, headers=headers, timeout=timeout)
+    r = session.post(
+        url, json=data, headers=headers, timeout=timeout, allow_redirects=False
+    )
     r.raise_for_status()
     return r
 
@@ -55,7 +57,7 @@ def notify_refs(
     data = {"repo_url": repo_url, "client": "dvc", "refs": refs}
 
     try:
-        r = post("/webhook/dvc", token, data, url=studio_url)
+        r = post("/webhook/dvc", token, data, base_url=studio_url)
     except requests.RequestException as e:
         logger.trace("", exc_info=True)  # type: ignore[attr-defined]
 
diff --git a/tests/unit/utils/test_studio.py b/tests/unit/utils/test_studio.py
index e47c546e05..a10c88bb04 100644
--- a/tests/unit/utils/test_studio.py
+++ b/tests/unit/utils/test_studio.py
@@ -41,4 +41,5 @@ def test_notify_refs(mocker, status_code, side_effect):
         },
         headers={"Authorization": "token TOKEN"},
         timeout=5,
+        allow_redirects=False,
     )