add lots of comments on the function of 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,14 +20,21 @@ 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 add_to_build_tracker(self) -> bool:
+    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

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
@@ -11,8 +13,8 @@ class InstalledDistribution(AbstractDistribution):
     """
 
     @property
-    def add_to_build_tracker(self) -> bool:
-        return False
+    def build_tracker_id(self) -> Optional[str]:
+        return None
 
     def get_metadata_distribution(self) -> BaseDistribution:
         dist = 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
@@ -19,8 +19,10 @@ class SourceDistribution(AbstractDistribution):
     """
 
     @property
-    def add_to_build_tracker(self) -> bool:
-        return True
+    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:
         assert (

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
@@ -16,8 +18,8 @@ class WheelDistribution(AbstractDistribution):
     """
 
     @property
-    def add_to_build_tracker(self) -> bool:
-        return True
+    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

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

@@ -51,10 +51,20 @@ 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 +79,15 @@ def __exit__(
     ) -> 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 +101,37 @@ def add(self, req: InstallRequirement) -> None:
             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

@@ -67,15 +67,12 @@ def _get_prepared_distribution(
 ) -> BaseDistribution:
     """Prepare a distribution for installation."""
     abstract_dist = make_distribution_for_install_requirement(req)
-    if abstract_dist.add_to_build_tracker:
-        with build_tracker.track(req):
+    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
             )
-    else:
-        abstract_dist.prepare_distribution_metadata(
-            finder, build_isolation, check_build_deps
-        )
     return abstract_dist.get_metadata_distribution()