Merge ~ines-almeida/launchpad:fetch-service-buildd-manager-refactor into launchpad:master

[Ines Almeida]

Refactor buildd-manager builder proxy logic

This separates the logic for making calls to the builder proxy and the fetch service out of the BuilderProxyMixin and into their own handlers.
This will make it easier to later add the logic to end sessions or, retrieve metadata.

All tests in `lp.snappy.tests` have be run and passed successfully.

[Ines Almeida (ines-almeida)]

The code itself should be tested from the existing unit tests, since the functionality was maintained. But I still need to have a look at adding new tests for the new classes.

[Jürgen Gmach (jugmac00)]

That is a much cleaner approach - I like it a lot.

It is a pity that we do not have a coverage report, as I would be absolutely happy with no further unit tests as the functionality "should" be already 100% covered.

mypy is not happy yet

[Ines Almeida (ines-almeida)]

Fixed all mypy issues (interestingly, some of them were issues that already existed in the code previously, but weren't an issue before), and added comments where it was relevant.

I will refrain from adding unit tests in this MP since the functionality seems to be highly tested between the multiple tests that contain `extraBuildArgs` and all the tests within `lp.snappy.tests.test_snapbuildbehaviour`.

[Jürgen Gmach (jugmac00)]

Thanks a lot for the MP!

- typo in the commit message
retreieve -> retrieve

This is a great improvement for readability and maintainability!

review: Approve

[Ines Almeida (ines-almeida)]

Thanks for the review! I'll update the commit when I squash them

9dc14ea... by Ines Almeida on 2024-04-17

Refactor buildd-manager builder proxy logic

This separates the logic for making calls to the builder proxy and the fetch service out of the BuilderProxyMixin and into their own handlers.
This will make it easier to later add the logic to end sessions or, retrieve metadata.

diff --git a/lib/lp/buildmaster/builderproxy.py b/lib/lp/buildmaster/builderproxy.py
index 5f3cf5b..a2b8f71 100644
--- a/lib/lp/buildmaster/builderproxy.py
+++ b/lib/lp/buildmaster/builderproxy.py
@@ -16,7 +16,7 @@ __all__ = [
 
 import base64
 import time
-from typing import Dict, Generator
+from typing import Dict, Generator, Type
 
 from twisted.internet import defer
 
@@ -24,8 +24,10 @@ from lp.buildmaster.downloader import (
     RequestFetchServiceSessionCommand,
     RequestProxyTokenCommand,
 )
+from lp.buildmaster.interactor import BuilderWorker
 from lp.buildmaster.interfaces.builder import CannotBuild
 from lp.buildmaster.interfaces.buildfarmjobbehaviour import BuildArgs
+from lp.buildmaster.model.buildfarmjob import BuildFarmJob
 from lp.services.config import config
 
 
@@ -44,43 +46,27 @@ def _get_value_from_config(key: str):
 class BuilderProxyMixin:
     """Methods for handling builds with the Snap Build Proxy enabled."""
 
+    # Variables that are set up by the class that uses the Mixin
+    build: BuildFarmJob
+    _worker: BuilderWorker
+
     @defer.inlineCallbacks
     def addProxyArgs(
         self,
         args: BuildArgs,
         allow_internet: bool = True,
-        fetch_service: bool = False,
+        use_fetch_service: bool = False,
     ) -> Generator[None, Dict[str, str], None]:
+
+        self.proxy_service = None
+
         if not allow_internet:
             return
-        if not fetch_service and _get_proxy_config("builder_proxy_host"):
-            token = yield self._requestProxyToken()
-            args["proxy_url"] = (
-                "http://{username}:{password}@{host}:{port}".format(
-                    username=token["username"],
-                    password=token["secret"],
-                    host=_get_proxy_config("builder_proxy_host"),
-                    port=_get_proxy_config("builder_proxy_port"),
-                )
-            )
-            args["revocation_endpoint"] = "{endpoint}/{token}".format(
-                endpoint=_get_proxy_config("builder_proxy_auth_api_endpoint"),
-                token=token["username"],
-            )
-        elif fetch_service and _get_proxy_config("fetch_service_host"):
-            session = yield self._requestFetchServiceSession()
-            args["proxy_url"] = (
-                "http://{session_id}:{token}@{host}:{port}".format(
-                    session_id=session["id"],
-                    token=session["token"],
-                    host=_get_proxy_config("fetch_service_host"),
-                    port=_get_proxy_config("fetch_service_port"),
-                )
-            )
-            args["revocation_endpoint"] = "{endpoint}/{session_id}".format(
-                endpoint=_get_proxy_config("fetch_service_control_endpoint"),
-                session_id=session["id"],
-            )
+
+        if not use_fetch_service and _get_proxy_config("builder_proxy_host"):
+            proxy_service: Type[IProxyService] = BuilderProxy
+        elif use_fetch_service and _get_proxy_config("fetch_service_host"):
+            proxy_service = FetchService
 
             # Append the fetch-service certificate to BuildArgs secrets.
             if "secrets" not in args:
@@ -88,47 +74,177 @@ class BuilderProxyMixin:
             args["secrets"]["fetch_service_mitm_certificate"] = (
                 _get_value_from_config("fetch_service_mitm_certificate")
             )
+        else:
+            # If the required config values aren't set, we skip adding proxy
+            # args or starting a proxy session. This might be relevant for
+            # non-production environments.
+            return
+
+        self.proxy_service = proxy_service(
+            build_id=self.build.build_cookie, worker=self._worker
+        )
+        new_session = yield self.proxy_service.startSession()
+        args["proxy_url"] = new_session["proxy_url"]
+        args["revocation_endpoint"] = new_session["revocation_endpoint"]
+
+
+class IProxyService:
+    """Interface for Proxy Services - either FetchService or BuilderProxy."""
+
+    def __init__(self, build_id: str, worker: BuilderWorker):
+        pass
 
     @defer.inlineCallbacks
-    def _requestProxyToken(self):
+    def startSession(self):
+        """Start a proxy session and request required
+
+        :returns: dictionary with an authenticated `proxy_url` for the builder
+        to use, and a `revocation_endpoint` to revoke the token when it's no
+        longer required.
+        """
+        pass
+
+
+class BuilderProxy(IProxyService):
+    """Handler for the Builder Proxy.
+
+    Handles the life-cycle of the proxy session used by a builder by
+    making API requests directly to the builder proxy control endpoint.
+    """
+
+    def __init__(self, build_id: str, worker: BuilderWorker):
+        self.control_endpoint = _get_value_from_config(
+            "builder_proxy_auth_api_endpoint"
+        )
+        self.proxy_endpoint = "{host}:{port}".format(
+            host=_get_value_from_config("builder_proxy_host"),
+            port=_get_value_from_config("builder_proxy_port"),
+        )
+        self.auth_header = self._getAuthHeader()
+
+        self.build_id = build_id
+        self.worker = worker
+
+    @staticmethod
+    def _getAuthHeader():
+        """Helper method to generate authentication needed to call the
+        builder proxy authentication endpoint."""
+
         admin_username = _get_value_from_config(
             "builder_proxy_auth_api_admin_username"
         )
-        secret = _get_value_from_config("builder_proxy_auth_api_admin_secret")
-        url = _get_value_from_config("builder_proxy_auth_api_endpoint")
+        admin_secret = _get_value_from_config(
+            "builder_proxy_auth_api_admin_secret"
+        )
+        auth_string = f"{admin_username}:{admin_secret}".strip()
+        return b"Basic " + base64.b64encode(auth_string.encode("ASCII"))
+
+    @defer.inlineCallbacks
+    def startSession(self):
+        """Request a token from the builder proxy to be used by the builders.
+
+        See IProxyService.
+        """
         timestamp = int(time.time())
         proxy_username = "{build_id}-{timestamp}".format(
-            build_id=self.build.build_cookie, timestamp=timestamp
+            build_id=self.build_id, timestamp=timestamp
         )
-        auth_string = f"{admin_username}:{secret}".strip()
-        auth_header = b"Basic " + base64.b64encode(auth_string.encode("ASCII"))
 
-        token = yield self._worker.process_pool.doWork(
+        token = yield self.worker.process_pool.doWork(
             RequestProxyTokenCommand,
-            url=url,
-            auth_header=auth_header,
+            url=self.control_endpoint,
+            auth_header=self.auth_header,
             proxy_username=proxy_username,
         )
-        return token
 
-    @defer.inlineCallbacks
-    def _requestFetchServiceSession(self):
+        proxy_url = "http://{username}:{password}@{proxy_endpoint}".format(
+            username=token["username"],
+            password=token["secret"],
+            proxy_endpoint=self.proxy_endpoint,
+        )
+        revocation_endpoint = "{endpoint}/{token}".format(
+            endpoint=self.control_endpoint,
+            token=token["username"],
+        )
+
+        return {
+            "proxy_url": proxy_url,
+            "revocation_endpoint": revocation_endpoint,
+        }
+
+
+class FetchService(IProxyService):
+    """Handler for the Fetch Service.
+
+    Handles the life-cycle of the fetch service session used by a builder by
+    making API requests directly to the fetch service control endpoint.
+    """
+
+    PROXY_URL = "http://{session_id}:{token}@{proxy_endpoint}"
+    START_SESSION_ENDPOINT = "{control_endpoint}"
+    TOKEN_REVOCATION = "{control_endpoint}/session/{session_id}/token"
+    RETRIEVE_METADATA_ENDPOINT = "{control_endpoint}/session/{session_id}"
+    END_SESSION_ENDPOINT = "{control_endpoint}/session/{session_id}"
+
+    def __init__(self, build_id: str, worker: BuilderWorker):
+        self.control_endpoint = _get_value_from_config(
+            "fetch_service_control_endpoint"
+        )
+        self.proxy_endpoint = "{host}:{port}".format(
+            host=_get_value_from_config("fetch_service_host"),
+            port=_get_value_from_config("fetch_service_port"),
+        )
+        self.auth_header = self._getAuthHeader()
+
+        self.build_id = build_id
+        self.worker = worker
+        self.session_id = None
+
+    @staticmethod
+    def _getAuthHeader():
+        """Helper method to generate authentication needed to call the
+        fetch service control endpoint."""
         admin_username = _get_value_from_config(
             "fetch_service_control_admin_username"
         )
         secret = _get_value_from_config("fetch_service_control_admin_secret")
-        url = _get_value_from_config("fetch_service_control_endpoint")
+        auth_string = f"{admin_username}:{secret}".strip()
+        return b"Basic " + base64.b64encode(auth_string.encode("ASCII"))
+
+    @defer.inlineCallbacks
+    def startSession(self):
+        """Requests a fetch service session and returns session information.
+
+        See IProxyService.
+        """
         timestamp = int(time.time())
         proxy_username = "{build_id}-{timestamp}".format(
-            build_id=self.build.build_cookie, timestamp=timestamp
+            build_id=self.build_id, timestamp=timestamp
         )
-        auth_string = f"{admin_username}:{secret}".strip()
-        auth_header = b"Basic " + base64.b64encode(auth_string.encode("ASCII"))
 
-        session = yield self._worker.process_pool.doWork(
+        session_data = yield self.worker.process_pool.doWork(
             RequestFetchServiceSessionCommand,
-            url=url,
-            auth_header=auth_header,
+            url=self.START_SESSION_ENDPOINT.format(
+                control_endpoint=self.control_endpoint
+            ),
+            auth_header=self.auth_header,
             proxy_username=proxy_username,
         )
-        return session
+
+        self.session_id = session_data["id"]
+        token = session_data["token"]
+
+        proxy_url = self.PROXY_URL.format(
+            session_id=self.session_id,
+            token=token,
+            proxy_endpoint=self.proxy_endpoint,
+        )
+        revocation_endpoint = self.TOKEN_REVOCATION.format(
+            control_endpoint=self.control_endpoint,
+            session_id=self.session_id,
+        )
+
+        return {
+            "proxy_url": proxy_url,
+            "revocation_endpoint": revocation_endpoint,
+        }