fixup! Don't send bearer tokens to DRS servers outside of Terra (#7660)

Author: nadove-ucsc

src/azul/http.py

@@ -1,16 +1,23 @@
+from collections.abc import (
+    Callable,
+)
 import logging
 import sys
 import time
 from typing import (
     Any,
     ClassVar,
     Self,
+    Sequence,
 )
 
 import certifi
 from furl import (
     furl,
 )
+from google.auth.transport.urllib3 import (
+    AuthorizedHttp,
+)
 import urllib3
 import urllib3.connection
 import urllib3.connectionpool
@@ -322,6 +329,28 @@ def urlopen(self, method, url, *args, **kwargs) -> urllib3.HTTPResponse:
             return response
 
 
+class ConditionalAuthHttpClient(HttpClientDecorator):
+    """
+    Must directly wrap an instance of AuthorizedHttp. An HTTP client that only
+    authenticates if a condition is met. The arguments to the predicate function
+    are the same as to :meth:`urlopen`.
+    """
+
+    def __init__(self,
+                 inner: AuthorizedHttp,
+                 condition: Callable[[str, str, Sequence[Any], dict[str, Any]], bool],
+                 headers: dict | None = None):
+        super().__init__(inner, headers)
+        self.condition = condition
+
+    def urlopen(self, method, url, *args, **kwargs) -> urllib3.HTTPResponse:
+        assert isinstance(self._inner, AuthorizedHttp), self._inner
+        if self.condition(method, url, *args, **kwargs):
+            return super().urlopen(method, url, *args, **kwargs)
+        else:
+            return self._inner.http.request(method, url, *args, **kwargs)
+
+
 class HasCachedHttpClient:
     """
     A convenience mixin that provides a cached instance property referring to an

src/azul/plugins/repository/tdr.py

@@ -23,7 +23,6 @@
 
 from azul import (
     cache_per_thread,
-    config,
     require,
 )
 from azul.auth import (
@@ -271,12 +270,6 @@ def update(self,
             assert self.location is None, self
             assert self.retry_after is None, self
         else:
-            # Compact DRS URIs cannot be parsed by furl
-            _, _, drs_host, *_ = self.file.drs_uri.split('/')
-            if drs_host != config.tdr_service_url.host:
-                # External DRS URis can point to servers outside of Terra, which
-                # don't accept our OAuth tokens.
-                authentication = None
             drs_client = plugin.drs_client(authentication)
             access = drs_client.get_object(self.file.drs_uri,
                                            access_method=AccessMethod.gs)

src/azul/terra.py

@@ -36,6 +36,9 @@
 from google.auth.transport.requests import (
     Request,
 )
+from google.auth.transport.urllib3 import (
+    AuthorizedHttp,
+)
 from google.cloud import (
     bigquery,
 )
@@ -79,6 +82,7 @@
     DRSClient,
 )
 from azul.http import (
+    ConditionalAuthHttpClient,
     LimitedRetryHttpClient,
     LimitedTimeoutException,
     Propagate429HttpClient,
@@ -289,9 +293,20 @@ class TerraClient(OAuth2Client):
     credentials_provider: TerraCredentialsProvider
 
     def _create_http_client(self) -> urllib3.request.RequestMethods:
+
+        authorized_http = super()._create_http_client()
+        assert isinstance(authorized_http, AuthorizedHttp), authorized_http
+
+        def should_authenticate(method: str, url: str, *args, **kwargs) -> bool:
+            url = furl(url)
+            return url.host in (config.tdr_service_url.host, config.sam_service_url.host)
+
         return Propagate429HttpClient(
             LimitedRetryHttpClient(
-                super()._create_http_client()
+                ConditionalAuthHttpClient(
+                    authorized_http,
+                    should_authenticate,
+                )
             )
         )