fix: remove archive.download dest before downloading if it's a symlink (#191)

Author: joshuarli

devenv/lib/archive.py

@@ -38,51 +38,60 @@ def download(
         dest = f"{cache_root}/{sha256}"
         os.makedirs(cache_root, exist_ok=True)
 
-    if not os.path.exists(dest):
-        headers = {}
-        if url.startswith("https://ghcr.io/v2/homebrew"):
-            # downloading homebrew blobs requires auth
-            # you can get an anonymous token from https://ghcr.io/token?service=ghcr.io&scope=repository%3Ahomebrew/core/go%3Apull
-            # but there's also a special shortcut token QQ==
-            # https://github.com/Homebrew/brew/blob/2184406bd8444e4de2626f5b0c749d4d08cb1aed/Library/Homebrew/brew.sh#L993
-            headers["Authorization"] = "bearer QQ=="
-
-        req = urllib.request.Request(url, headers=headers)
-
-        retry_sleep = 1.0
-        while retries >= 0:
-            try:
-                resp = urllib.request.urlopen(req)
-                break
-            except HTTPError as e:
-                if retries == 0:
-                    raise RuntimeError(f"Error getting {url}: {e}")
-                print(f"Error getting {url} ({retries} retries left): {e}")
-
-            time.sleep(retry_sleep)
-            retries -= 1
-            retry_sleep *= retry_exp
-
-        dest_dir = os.path.dirname(dest)
-        os.makedirs(dest_dir, exist_ok=True)
-
-        with tempfile.NamedTemporaryFile(delete=False, dir=dest_dir) as tmpf:
-            shutil.copyfileobj(resp, tmpf)
-            tmpf.seek(0)
-            checksum = hashlib.sha256()
+    if os.path.islink(dest):
+        # there are cases where dest can be an existing symlink
+        # (docker desktop starts and puts symlinks into ~/.docker/cli-plugins)
+        # such symlinks should be removed otherwise callers to download
+        # usually try to chmod after and end up with PermissionError
+        os.remove(dest)
+
+    if os.path.exists(dest):
+        return dest
+
+    headers = {}
+    if url.startswith("https://ghcr.io/v2/homebrew"):
+        # downloading homebrew blobs requires auth
+        # you can get an anonymous token from https://ghcr.io/token?service=ghcr.io&scope=repository%3Ahomebrew/core/go%3Apull
+        # but there's also a special shortcut token QQ==
+        # https://github.com/Homebrew/brew/blob/2184406bd8444e4de2626f5b0c749d4d08cb1aed/Library/Homebrew/brew.sh#L993
+        headers["Authorization"] = "bearer QQ=="
+
+    req = urllib.request.Request(url, headers=headers)
+
+    retry_sleep = 1.0
+    while retries >= 0:
+        try:
+            resp = urllib.request.urlopen(req)
+            break
+        except HTTPError as e:
+            if retries == 0:
+                raise RuntimeError(f"Error getting {url}: {e}")
+            print(f"Error getting {url} ({retries} retries left): {e}")
+
+        time.sleep(retry_sleep)
+        retries -= 1
+        retry_sleep *= retry_exp
+
+    dest_dir = os.path.dirname(dest)
+    os.makedirs(dest_dir, exist_ok=True)
+
+    with tempfile.NamedTemporaryFile(delete=False, dir=dest_dir) as tmpf:
+        shutil.copyfileobj(resp, tmpf)
+        tmpf.seek(0)
+        checksum = hashlib.sha256()
+        buf = tmpf.read(4096)
+        while buf:
+            checksum.update(buf)
             buf = tmpf.read(4096)
-            while buf:
-                checksum.update(buf)
-                buf = tmpf.read(4096)
-
-            if not secrets.compare_digest(checksum.hexdigest(), sha256):
-                raise RuntimeError(
-                    f"checksum mismatch for {url}:\n"
-                    f"- got: {checksum.hexdigest()}\n"
-                    f"- expected: {sha256}\n"
-                )
-
-            atomic_replace(tmpf.name, dest)
+
+        if not secrets.compare_digest(checksum.hexdigest(), sha256):
+            raise RuntimeError(
+                f"checksum mismatch for {url}:\n"
+                f"- got: {checksum.hexdigest()}\n"
+                f"- expected: {sha256}\n"
+            )
+
+        atomic_replace(tmpf.name, dest)
 
     return dest
 

tests/lib/test_archive.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import io
+import os
 import pathlib
 import tarfile
 import time
@@ -169,6 +170,12 @@ def test_download(tmp_path: pathlib.Path, mock_sleep: mock.MagicMock) -> None:
     )
 
     dest = f"{tmp_path}/a"
+
+    # if dest is already a valid symlink it should be paved over
+    with open(f"{tmp_path}/hi", "wb"):
+        pass
+    os.symlink(f"{tmp_path}/hi", dest)
+
     with mock.patch.object(
         urllib.request,
         "urlopen",
@@ -186,7 +193,22 @@ def test_download(tmp_path: pathlib.Path, mock_sleep: mock.MagicMock) -> None:
     with open(dest, "rb") as f:
         assert f.read() == data
 
-    dest = f"{tmp_path}/b"
+
+def test_download_exceeded_retries(
+    tmp_path: pathlib.Path, mock_sleep: mock.MagicMock
+) -> None:
+    data_sha256 = (
+        "b5bb9d8014a0f9b1d61e21e796d78dccdf1352f23cd32812f4850b878ae4944c"
+    )
+
+    err = urllib.error.HTTPError(
+        "https://example.com/foo",
+        503,
+        "Service Unavailable",
+        "",  # type: ignore
+        io.BytesIO(b""),
+    )
+    dest = f"{tmp_path}/a"
     with pytest.raises(RuntimeError) as excinfo:
         with mock.patch.object(
             urllib.request,
@@ -202,7 +224,42 @@ def test_download(tmp_path: pathlib.Path, mock_sleep: mock.MagicMock) -> None:
         == "Error getting https://example.com/foo: HTTP Error 503: Service Unavailable"
     )
 
-    dest = f"{tmp_path}/b"
+
+def test_download_dest_is_broken_symlink(
+    tmp_path: pathlib.Path, mock_sleep: mock.MagicMock
+) -> None:
+    data = b"foo\n"
+    data_sha256 = (
+        "b5bb9d8014a0f9b1d61e21e796d78dccdf1352f23cd32812f4850b878ae4944c"
+    )
+
+    dest = f"{tmp_path}/a"
+
+    # if dest is already a dead symlink it should be paved over as well
+    os.symlink(f"{tmp_path}/does-not-exist", dest)
+
+    with mock.patch.object(
+        urllib.request,
+        "urlopen",
+        autospec=True,
+        side_effect=(io.BytesIO(data),),
+    ):
+        archive.download("https://example.com/foo", data_sha256, dest)
+
+    with open(dest, "rb") as f:
+        assert f.read() == data
+
+
+def test_download_wrong_sha(
+    tmp_path: pathlib.Path, mock_sleep: mock.MagicMock
+) -> None:
+    data = b"foo\n"
+    data_sha256 = (
+        "b5bb9d8014a0f9b1d61e21e796d78dccdf1352f23cd32812f4850b878ae4944c"
+    )
+
+    dest = f"{tmp_path}/a"
+
     with pytest.raises(RuntimeError) as excinfo:
         with mock.patch.object(
             urllib.request,