Merge ~ines-almeida/launchpad:fetch-service-update-end-session into launchpad:master

[Ines Almeida]

Update how to keep track of session_id within a builder session

[Ines Almeida (ines-almeida)]

This has not been cleaned up (and mypy issues not fixed). Opening it up for an early review.

[Ines Almeida (ines-almeida)]

Fixed the mypy issues and updated the logic to use the `use_fetch_service` flag directly from the builder proxy info. I still need to re-work a test, but this is now ready for review.

[Jürgen Gmach (jugmac00)]

review: Approve

diff --git a/lib/lp/buildmaster/builderproxy.py b/lib/lp/buildmaster/builderproxy.py
index 1fc0d9b..a52ce06 100644
--- a/lib/lp/buildmaster/builderproxy.py
+++ b/lib/lp/buildmaster/builderproxy.py
@@ -17,8 +17,9 @@ __all__ = [
 
 import base64
 import os
+import re
 import time
-from typing import Dict, Generator, Type
+from typing import Optional
 
 from twisted.internet import defer
 
@@ -37,6 +38,10 @@ from lp.services.config import config
 BUILD_METADATA_FILENAME_FORMAT = "{build_id}_metadata.json"
 
 
+class ProxyServiceException(Exception):
+    pass
+
+
 def _get_proxy_config(name):
     """Get a config item from builddmaster (current) or snappy (deprecated)."""
     return getattr(config.builddmaster, name) or getattr(config.snappy, name)
@@ -56,28 +61,23 @@ class BuilderProxyMixin:
     build: BuildFarmJob
     _worker: BuilderWorker
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self._use_fetch_service = False
-        self._proxy_service = None
-
     @defer.inlineCallbacks
     def startProxySession(
         self,
         args: BuildArgs,
         allow_internet: bool = True,
         use_fetch_service: bool = False,
-    ) -> Generator[None, Dict[str, str], None]:
-
-        self._use_fetch_service = use_fetch_service
+    ):
 
         if not allow_internet:
             return
 
+        proxy_service: IProxyService
+
         if not use_fetch_service and _get_proxy_config("builder_proxy_host"):
-            proxy_service: Type[IProxyService] = BuilderProxy
+            proxy_service = BuilderProxy(worker=self._worker)
         elif use_fetch_service and _get_proxy_config("fetch_service_host"):
-            proxy_service = FetchService
+            proxy_service = FetchService(worker=self._worker)
 
             # Append the fetch-service certificate to BuildArgs secrets.
             if "secrets" not in args:
@@ -91,10 +91,9 @@ class BuilderProxyMixin:
             # non-production environments.
             return
 
-        self._proxy_service = proxy_service(
-            build_id=self.build.build_cookie, worker=self._worker
+        session_data = yield proxy_service.startSession(
+            build_id=self.build.build_cookie
         )
-        session_data = yield self._proxy_service.startSession()
 
         args["proxy_url"] = session_data["proxy_url"]
         args["revocation_endpoint"] = session_data["revocation_endpoint"]
@@ -115,34 +114,55 @@ class BuilderProxyMixin:
         Sessions will be closed automatically within the Fetch Service after
         a certain amount of time configured by its charm (default 6 hours).
         """
-        if not self._use_fetch_service:
-            # No cleanup needed for the builder proxy
-            return
 
-        if self._proxy_service is None:
-            # A session was never started. This can happen if the proxy configs
-            # are not set (see `startProxySession`)
+        proxy_info = yield self._worker.proxy_info()
+        use_fetch_service = proxy_info.get("use_fetch_service")
+
+        if not use_fetch_service:
+            # No cleanup needed when not using the fetch service
+            # This is true both when we use the builder proxy and when we don't
+            # allow internet access to the builds.
             return
 
+        proxy_service = FetchService(worker=self._worker)
+
+        # XXX ines-almeida 2024-04-30: Getting the session_id from the
+        # revocation_endpoint feels a little like going back and forwards
+        # given the revocation_endpoint is created on `startProxySession()`.
+        # Ideally, we would update `buildd` and `buildd-manager` to senf and
+        # retrieve the session ID directly (instead of having to parse it).
+        revocation_endpoint = proxy_info.get("revocation_endpoint")
+        session_id = proxy_service.extractSessionIDFromRevocationEndpoint(
+            revocation_endpoint
+        )
+
+        if session_id is None:
+            raise ProxyServiceException(
+                "Could not parse the revocation endpoint fetched from buildd "
+                f"('{revocation_endpoint}') to get the fetch service "
+                "`session_id` used within the build."
+            )
+
         metadata_file_name = BUILD_METADATA_FILENAME_FORMAT.format(
             build_id=self.build.build_cookie
         )
         file_path = os.path.join(upload_path, metadata_file_name)
-        yield self._proxy_service.retrieveMetadataFromSession(
-            save_content_to=file_path
+        yield proxy_service.retrieveMetadataFromSession(
+            session_id=session_id,
+            save_content_to=file_path,
         )
 
-        yield self._proxy_service.endSession()
+        yield proxy_service.endSession(session_id=session_id)
 
 
 class IProxyService:
     """Interface for Proxy Services - either FetchService or BuilderProxy."""
 
-    def __init__(self, build_id: str, worker: BuilderWorker):
+    def __init__(self, worker: BuilderWorker):
         pass
 
     @defer.inlineCallbacks
-    def startSession(self):
+    def startSession(self, build_id: str):
         """Start a proxy session and request required
 
         :returns: dictionary with an authenticated `proxy_url` for the builder
@@ -159,7 +179,7 @@ class BuilderProxy(IProxyService):
     making API requests directly to the builder proxy control endpoint.
     """
 
-    def __init__(self, build_id: str, worker: BuilderWorker):
+    def __init__(self, worker: BuilderWorker):
         self.control_endpoint = _get_value_from_config(
             "builder_proxy_auth_api_endpoint"
         )
@@ -168,8 +188,6 @@ class BuilderProxy(IProxyService):
             port=_get_value_from_config("builder_proxy_port"),
         )
         self.auth_header = self._getAuthHeader()
-
-        self.build_id = build_id
         self.worker = worker
 
     @staticmethod
@@ -187,14 +205,14 @@ class BuilderProxy(IProxyService):
         return b"Basic " + base64.b64encode(auth_string.encode("ASCII"))
 
     @defer.inlineCallbacks
-    def startSession(self):
+    def startSession(self, build_id: str):
         """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_id, timestamp=timestamp
+            build_id=build_id, timestamp=timestamp
         )
 
         token = yield self.worker.process_pool.doWork(
@@ -233,7 +251,7 @@ class FetchService(IProxyService):
     RETRIEVE_METADATA_ENDPOINT = "{control_endpoint}/session/{session_id}"
     END_SESSION_ENDPOINT = "{control_endpoint}/session/{session_id}"
 
-    def __init__(self, build_id: str, worker: BuilderWorker):
+    def __init__(self, worker: BuilderWorker):
         self.control_endpoint = _get_value_from_config(
             "fetch_service_control_endpoint"
         )
@@ -242,10 +260,7 @@ class FetchService(IProxyService):
             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():
@@ -259,7 +274,7 @@ class FetchService(IProxyService):
         return b"Basic " + base64.b64encode(auth_string.encode("ASCII"))
 
     @defer.inlineCallbacks
-    def startSession(self):
+    def startSession(self, build_id: str):
         """Requests a fetch service session and returns session information.
 
         See IProxyService.
@@ -290,8 +305,28 @@ class FetchService(IProxyService):
             "revocation_endpoint": revocation_endpoint,
         }
 
+    def extractSessionIDFromRevocationEndpoint(
+        self, revocation_endpoint: str
+    ) -> Optional[str]:
+        """Helper method to get the session_id out of the revocation
+        endpoint.
+
+        Example `revocation_endpoint`:
+        http://fetch-service.net:9999/session/913890402914ce3be9ffd4d0/token
+
+        Would return: 913890402914ce3be9ffd4d0
+        """
+        re_pattern = self.TOKEN_REVOCATION.format(
+            control_endpoint=self.control_endpoint,
+            session_id="(?P<session_id>.*)",
+        )
+        match = re.match(re_pattern, revocation_endpoint)
+        return match["session_id"] if match else None
+
     @defer.inlineCallbacks
-    def retrieveMetadataFromSession(self, save_content_to: str):
+    def retrieveMetadataFromSession(
+        self, session_id: str, save_content_to: str
+    ):
         """Make request to retrieve metadata from the current session.
 
         Data is stored directly into a file whose path is `save_content_to`
@@ -300,7 +335,7 @@ class FetchService(IProxyService):
         """
         url = self.RETRIEVE_METADATA_ENDPOINT.format(
             control_endpoint=self.control_endpoint,
-            session_id=self.session_id,
+            session_id=session_id,
         )
         yield self.worker.process_pool.doWork(
             RetrieveFetchServiceSessionCommand,
@@ -310,14 +345,14 @@ class FetchService(IProxyService):
         )
 
     @defer.inlineCallbacks
-    def endSession(self):
+    def endSession(self, session_id: str):
         """End the proxy session and do any cleanup needed.
 
         :raises: RequestException if request to Fetch Service fails
         """
         url = self.END_SESSION_ENDPOINT.format(
             control_endpoint=self.control_endpoint,
-            session_id=self.session_id,
+            session_id=session_id,
         )
         yield self.worker.process_pool.doWork(
             EndFetchServiceSessionCommand,
diff --git a/lib/lp/buildmaster/interactor.py b/lib/lp/buildmaster/interactor.py
index dfcc86b..9bf1a71 100644
--- a/lib/lp/buildmaster/interactor.py
+++ b/lib/lp/buildmaster/interactor.py
@@ -209,6 +209,10 @@ class BuilderWorker:
         """Echo the arguments back."""
         return self._with_timeout(self._server.callRemote("echo", *args))
 
+    def proxy_info(self):
+        """Return the details for the proxy used by the manager."""
+        return self._with_timeout(self._server.callRemote("proxy_info"))
+
     def info(self):
         """Return the protocol version and the builder methods supported."""
         return self._with_timeout(self._server.callRemote("info"))
diff --git a/lib/lp/buildmaster/tests/mock_workers.py b/lib/lp/buildmaster/tests/mock_workers.py
index a880a9a..5f5563b 100644
--- a/lib/lp/buildmaster/tests/mock_workers.py
+++ b/lib/lp/buildmaster/tests/mock_workers.py
@@ -151,6 +151,18 @@ class OkWorker:
         self.call_log.append("info")
         return defer.succeed(("1.0", self.arch_tag, "binarypackage"))
 
+    def proxy_info(self):
+        self.call_log.append("proxy_info")
+        return defer.succeed(
+            {
+                "revocation_endpoint": (
+                    "http://fetch-service.test:9999/session/"
+                    "9138904ce3be9ffd4d0/token"
+                ),
+                "use_fetch_service": False,
+            }
+        )
+
     def resume(self):
         self.call_log.append("resume")
         return defer.succeed(("", "", 0))
diff --git a/lib/lp/snappy/tests/test_snapbuildbehaviour.py b/lib/lp/snappy/tests/test_snapbuildbehaviour.py
index 81d6926..cb2269b 100644
--- a/lib/lp/snappy/tests/test_snapbuildbehaviour.py
+++ b/lib/lp/snappy/tests/test_snapbuildbehaviour.py
@@ -10,6 +10,7 @@ import time
 import uuid
 from datetime import datetime
 from textwrap import dedent
+from unittest.mock import MagicMock
 from urllib.parse import urlsplit
 
 import fixtures
@@ -40,7 +41,6 @@ from lp.app.enums import InformationType
 from lp.archivepublisher.interfaces.archivegpgsigningkey import (
     IArchiveGPGSigningKey,
 )
-from lp.buildmaster.builderproxy import BuilderProxy
 from lp.buildmaster.enums import BuildBaseImageType, BuildStatus
 from lp.buildmaster.interactor import shut_down_default_process_pool
 from lp.buildmaster.interfaces.builder import CannotBuild
@@ -431,6 +431,20 @@ class TestAsyncSnapBuildBehaviourFetchService(
         snap = self.factory.makeSnap(use_fetch_service=True)
         request = self.factory.makeSnapBuildRequest(snap=snap)
         job = self.makeJob(snap=snap, build_request=request)
+
+        host = config.builddmaster.fetch_service_host
+        port = config.builddmaster.fetch_service_port
+        session_id = self.fetch_service_api.sessions.session_id
+        revocation_endpoint = (
+            f"http://{host}:{port}/session/{session_id}/token"
+        )
+
+        job._worker.proxy_info = MagicMock(
+            return_value={
+                "revocation_endpoint": revocation_endpoint,
+                "use_fetch_service": True,
+            }
+        )
         yield job.extraBuildArgs()
 
         # End the session
@@ -443,7 +457,6 @@ class TestAsyncSnapBuildBehaviourFetchService(
         start_session_request = self.fetch_service_api.sessions.requests[0]
         self.assertEqual(b"POST", start_session_request["method"])
         self.assertEqual(b"/session", start_session_request["uri"])
-        session_id = self.fetch_service_api.sessions.responses[0]["id"]
 
         # Request retrieve metadata
         retrieve_metadata_request = self.fetch_service_api.sessions.requests[1]
@@ -478,6 +491,14 @@ class TestAsyncSnapBuildBehaviourFetchService(
         snap = self.factory.makeSnap(use_fetch_service=False)
         request = self.factory.makeSnapBuildRequest(snap=snap)
         job = self.makeJob(snap=snap, build_request=request)
+
+        job._worker.proxy_info = MagicMock(
+            return_value={
+                "revocation_endpoint": "https://builder-proxy.test/revoke",
+                "use_fetch_service": False,
+            }
+        )
+
         yield job.extraBuildArgs()
         yield job.endProxySession(upload_path="test_path")
 
@@ -485,23 +506,33 @@ class TestAsyncSnapBuildBehaviourFetchService(
         self.assertEqual(0, len(self.fetch_service_api.sessions.requests))
 
     @defer.inlineCallbacks
-    def test_endProxySession_no_proxy_service(self):
-        """When the `fetch_service_host` is not set, the calls to the fetch
-        service don't go through."""
+    def test_endProxySession_fetch_Service_allow_internet_false(self):
+        """When `allow_internet` is False, we don't send proxy variables to the
+        buildd, and ending the session does not make calls to the fetch
+        service."""
         self.useFixture(
             FeatureFixture({SNAP_USE_FETCH_SERVICE_FEATURE_FLAG: "on"})
         )
-        self.useFixture(FeatureFixture({"fetch_service_host": None}))
-
-        snap = self.factory.makeSnap(use_fetch_service=True)
+        snap = self.factory.makeSnap(allow_internet=False)
         request = self.factory.makeSnapBuildRequest(snap=snap)
         job = self.makeJob(snap=snap, build_request=request)
-        yield job.extraBuildArgs()
+        args = yield job.extraBuildArgs()
+
+        # Scenario when we don't allow internet
+        job._worker.proxy_info = MagicMock(
+            return_value={
+                "revocation_endpoint": None,
+                "use_fetch_service": None,
+            }
+        )
+
         yield job.endProxySession(upload_path="test_path")
 
+        # No proxy config sent to buildd
+        self.assertIsNone(args.get("use_fetch_service"))
+        self.assertIsNone(args.get("revocation_endpoint"))
         # No calls go through to the fetch service
         self.assertEqual(0, len(self.fetch_service_api.sessions.requests))
-        self.assertEqual(None, job._proxy_service)
 
 
 class TestAsyncSnapBuildBehaviourBuilderProxy(
@@ -1470,8 +1501,6 @@ class TestAsyncSnapBuildBehaviourBuilderProxy(
 
         # End the session
         yield job.endProxySession(upload_path="test_path")
-        self.assertFalse(job.use_fetch_service)
-        self.assertTrue(isinstance(job.proxy_service, BuilderProxy))
 
     @defer.inlineCallbacks
     def test_composeBuildRequest_proxy_url_set(self):