Merge pull request #12194 from cosmicexplorer/build-tracker-comments

Add lots of comments on the function of BuildTracker

news/12194.trivial.rst

@@ -0,0 +1 @@
+Add lots of comments to the ``BuildTracker``.

src/pip/_internal/distributions/base.py

@@ -1,4 +1,5 @@
 import abc
+from typing import Optional
 
 from pip._internal.index.package_finder import PackageFinder
 from pip._internal.metadata.base import BaseDistribution
@@ -19,12 +20,23 @@ class AbstractDistribution(metaclass=abc.ABCMeta):
 
      - we must be able to create a Distribution object exposing the
        above metadata.
+
+     - if we need to do work in the build tracker, we must be able to generate a unique
+       string to identify the requirement in the build tracker.
     """
 
     def __init__(self, req: InstallRequirement) -> None:
         super().__init__()
         self.req = req
 
+    @abc.abstractproperty
+    def build_tracker_id(self) -> Optional[str]:
+        """A string that uniquely identifies this requirement to the build tracker.
+
+        If None, then this dist has no work to do in the build tracker, and
+        ``.prepare_distribution_metadata()`` will not be called."""
+        raise NotImplementedError()
+
     @abc.abstractmethod
     def get_metadata_distribution(self) -> BaseDistribution:
         raise NotImplementedError()

src/pip/_internal/distributions/installed.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 from pip._internal.distributions.base import AbstractDistribution
 from pip._internal.index.package_finder import PackageFinder
 from pip._internal.metadata import BaseDistribution
@@ -10,6 +12,10 @@ class InstalledDistribution(AbstractDistribution):
     been computed.
     """
 
+    @property
+    def build_tracker_id(self) -> Optional[str]:
+        return None
+
     def get_metadata_distribution(self) -> BaseDistribution:
         assert self.req.satisfied_by is not None, "not actually installed"
         return self.req.satisfied_by

src/pip/_internal/distributions/sdist.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Iterable, Set, Tuple
+from typing import Iterable, Optional, Set, Tuple
 
 from pip._internal.build_env import BuildEnvironment
 from pip._internal.distributions.base import AbstractDistribution
@@ -18,6 +18,12 @@ class SourceDistribution(AbstractDistribution):
     generated, either using PEP 517 or using the legacy `setup.py egg_info`.
     """
 
+    @property
+    def build_tracker_id(self) -> Optional[str]:
+        """Identify this requirement uniquely by its link."""
+        assert self.req.link
+        return self.req.link.url_without_fragment
+
     def get_metadata_distribution(self) -> BaseDistribution:
         return self.req.get_dist()
 

src/pip/_internal/distributions/wheel.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 from pip._vendor.packaging.utils import canonicalize_name
 
 from pip._internal.distributions.base import AbstractDistribution
@@ -15,6 +17,10 @@ class WheelDistribution(AbstractDistribution):
     This does not need any preparation as wheels can be directly unpacked.
     """
 
+    @property
+    def build_tracker_id(self) -> Optional[str]:
+        return None
+
     def get_metadata_distribution(self) -> BaseDistribution:
         """Loads the metadata from the wheel file into memory and returns a
         Distribution that uses it, not relying on the wheel file or

src/pip/_internal/operations/build/build_tracker.py

@@ -51,10 +51,22 @@ def get_build_tracker() -> Generator["BuildTracker", None, None]:
             yield tracker
 
 
+class TrackerId(str):
+    """Uniquely identifying string provided to the build tracker."""
+
+
 class BuildTracker:
+    """Ensure that an sdist cannot request itself as a setup requirement.
+
+    When an sdist is prepared, it identifies its setup requirements in the
+    context of ``BuildTracker.track()``. If a requirement shows up recursively, this
+    raises an exception.
+
+    This stops fork bombs embedded in malicious packages."""
+
     def __init__(self, root: str) -> None:
         self._root = root
-        self._entries: Set[InstallRequirement] = set()
+        self._entries: Dict[TrackerId, InstallRequirement] = {}
         logger.debug("Created build tracker: %s", self._root)
 
     def __enter__(self) -> "BuildTracker":
@@ -69,16 +81,15 @@ class BuildTracker:
     ) -> None:
         self.cleanup()
 
-    def _entry_path(self, link: Link) -> str:
-        hashed = hashlib.sha224(link.url_without_fragment.encode()).hexdigest()
+    def _entry_path(self, key: TrackerId) -> str:
+        hashed = hashlib.sha224(key.encode()).hexdigest()
         return os.path.join(self._root, hashed)
 
-    def add(self, req: InstallRequirement) -> None:
+    def add(self, req: InstallRequirement, key: TrackerId) -> None:
         """Add an InstallRequirement to build tracking."""
 
-        assert req.link
         # Get the file to write information about this requirement.
-        entry_path = self._entry_path(req.link)
+        entry_path = self._entry_path(key)
 
         # Try reading from the file. If it exists and can be read from, a build
         # is already in progress, so a LookupError is raised.
@@ -92,33 +103,37 @@ class BuildTracker:
             raise LookupError(message)
 
         # If we're here, req should really not be building already.
-        assert req not in self._entries
+        assert key not in self._entries
 
         # Start tracking this requirement.
         with open(entry_path, "w", encoding="utf-8") as fp:
             fp.write(str(req))
-        self._entries.add(req)
+        self._entries[key] = req
 
         logger.debug("Added %s to build tracker %r", req, self._root)
 
-    def remove(self, req: InstallRequirement) -> None:
+    def remove(self, req: InstallRequirement, key: TrackerId) -> None:
         """Remove an InstallRequirement from build tracking."""
 
-        assert req.link
-        # Delete the created file and the corresponding entries.
-        os.unlink(self._entry_path(req.link))
-        self._entries.remove(req)
+        # Delete the created file and the corresponding entry.
+        os.unlink(self._entry_path(key))
+        del self._entries[key]
 
         logger.debug("Removed %s from build tracker %r", req, self._root)
 
     def cleanup(self) -> None:
-        for req in set(self._entries):
-            self.remove(req)
+        for key, req in list(self._entries.items()):
+            self.remove(req, key)
 
         logger.debug("Removed build tracker: %r", self._root)
 
     @contextlib.contextmanager
-    def track(self, req: InstallRequirement) -> Generator[None, None, None]:
-        self.add(req)
+    def track(self, req: InstallRequirement, key: str) -> Generator[None, None, None]:
+        """Ensure that `key` cannot install itself as a setup requirement.
+
+        :raises LookupError: If `key` was already provided in a parent invocation of
+                             the context introduced by this method."""
+        tracker_id = TrackerId(key)
+        self.add(req, tracker_id)
         yield
-        self.remove(req)
+        self.remove(req, tracker_id)

src/pip/_internal/operations/prepare.py

@@ -65,10 +65,12 @@ def _get_prepared_distribution(
 ) -> BaseDistribution:
     """Prepare a distribution for installation."""
     abstract_dist = make_distribution_for_install_requirement(req)
-    with build_tracker.track(req):
-        abstract_dist.prepare_distribution_metadata(
-            finder, build_isolation, check_build_deps
-        )
+    tracker_id = abstract_dist.build_tracker_id
+    if tracker_id is not None:
+        with build_tracker.track(req, tracker_id):
+            abstract_dist.prepare_distribution_metadata(
+                finder, build_isolation, check_build_deps
+            )
     return abstract_dist.get_metadata_distribution()