Repository: app-sre/sretoolbox

Issue:
Cache registry token
We acquire a new token on every request:

https://github.com/app-sre/sretoolbox/blob/master/sretoolbox/container/image.py#L355

We should probably make a more efficient use of it.

Discussion:

Patch:
diff --git a/sretoolbox/container/image.py b/sretoolbox/container/image.py
index 92bedf5..0721ab5 100644
--- a/sretoolbox/container/image.py
+++ b/sretoolbox/container/image.py
@@ -41,6 +41,7 @@ class NoTagForImageByDigest(Exception):
     Raised when the Image was constructed with a by-digest URL and an
     operation is attempted that requires a tag.
     """
+
     def __init__(self, image):
         super().__init__(
             f"Can't determine a unique tag for Image: {str(image)}")
@@ -65,7 +66,7 @@ class Image:
     MAX_CACHE_ITEM_SIZE = 50*1024
 
     def __init__(self, url, tag_override=None, username=None, password=None,
-                 auth_server=None, response_cache=None):
+                 auth_server=None, response_cache=None, auth_token=None):
         image_data = self._parse_image_url(url)
         self.scheme = image_data['scheme']
         self.registry = image_data['registry']
@@ -73,6 +74,7 @@ class Image:
         self.image = image_data['image']
         self.response_cache = response_cache
 
+        self.auth_token = auth_token
         if tag_override is None:
             self.tag = image_data['tag']
         else:
@@ -361,17 +363,23 @@ class Image:
 
     @retry(exceptions=(HTTPError, requests.ConnectionError), max_attempts=5)
     def _request_get(self, url, method=requests.get):
-        # Try first without 'Authorization' header
+        # Use any cached tokens, they may still be valid
         headers = {
             'Accept':
-                'application/vnd.docker.distribution.manifest.v1+json,'
-                'application/vnd.docker.distribution.manifest.v2+json,'
-                'application/vnd.docker.distribution.manifest.v1+prettyjws,'
+            'application/vnd.docker.distribution.manifest.v1+json,'
+            'application/vnd.docker.distribution.manifest.v2+json,'
+            'application/vnd.docker.distribution.manifest.v1+prettyjws,'
         }
 
-        response = method(url, headers=headers, auth=self.auth)
+        if self.auth_token:
+            headers['Authorization'] = self.auth_token
+            auth = None
+        else:
+            auth = self.auth
+
+        response = method(url, headers=headers, auth=auth)
 
-        # Unauthorized, meaning we have to acquire a token
+        # Unauthorized, meaning we have to acquire a new token
         if response.status_code == 401:
             auth_specs = response.headers.get('Www-Authenticate')
             if auth_specs is None:
@@ -379,8 +387,9 @@ class Image:
 
             www_auth = self._parse_www_auth(auth_specs)
 
-            # Try again, this time with the Authorization header
-            headers['Authorization'] = self._get_auth(www_auth)
+            # Try again, with the new Authorization header
+            self.auth_token = self._get_auth(www_auth)
+            headers['Authorization'] = self.auth_token
             response = method(url, headers=headers)
 
         self._raise_for_status(response)
@@ -467,7 +476,8 @@ class Image:
         return Image(url=str(self), tag_override=str(item),
                      username=self.username, password=self.password,
                      auth_server=self.auth_server,
-                     response_cache=self.response_cache)
+                     response_cache=self.response_cache,
+                     auth_token=self.auth_token)
 
     def __iter__(self):
         for tag in self.tags:


Test patch:
diff --git a/tests/test_container.py b/tests/test_container.py
index 542d716..cc92830 100644
--- a/tests/test_container.py
+++ b/tests/test_container.py
@@ -15,7 +15,7 @@ import requests
 
 import pytest
 
-from unittest.mock import patch
+from unittest.mock import patch, MagicMock
 
 from sretoolbox.container import Image
 
@@ -158,11 +158,12 @@ class TestContainer:
         assert e.typename == 'NoTagForImageByDigest'
 
     def test_getitem(self):
-        image = Image("quay.io/foo/bar:latest", response_cache={})
+        image = Image("quay.io/foo/bar:latest", response_cache={},
+                      auth_token="atoken")
         other = image['current']
         assert image.response_cache is other.response_cache
-
-
+        assert other.auth_token is image.auth_token
+        assert other.tag == 'current'
 
 
 @patch.object(Image, '_request_get', spec=Image)
@@ -229,3 +230,56 @@ class TestGetManifest:
             "https://quay.io/v2/foo/bar/manifests/latest"
         )
         should_cache.assert_not_called()
+
+
+@patch.object(Image, '_parse_www_auth')
+@patch.object(Image, '_get_auth')
+class TestRequestGet:
+    def test_username_and_password_ok(self, getauth, parseauth):
+        r = requests.Response()
+        r.status_code = 200
+        method = MagicMock(return_value=r)
+        i = Image("quay.io/foo/bar:latest", username="user", password="pass")
+        i._request_get.__wrapped__(i, "http://www.google.com", method=method)
+        method.assert_called_once()
+        c = method.call_args_list[0]
+
+        assert c[0] == ('http://www.google.com', )
+        assert 'Authorization' not in c[1]['headers']
+        assert c[1]['auth'] == i.auth
+        getauth.assert_not_called()
+        parseauth.assert_not_called()
+
+    def test_username_and_password_reauthenticate(self, getauth, parseauth):
+        r = requests.Response()
+        r.status_code = 401
+        r.headers['Www-Authenticate'] = 'something something'
+        gets = [r]
+        r = requests.Response()
+        r.status_code = 200
+        gets.append(r)
+        method = MagicMock(side_effect=gets)
+        r = requests.Response()
+        r.status_code = 200
+        i = Image("quay.io/foo/bar:latest", username="user", password="pass")
+        getauth.return_value = "anauthtoken"
+        parseauth.return_value = "aparsedauth"
+        i._request_get.__wrapped__(i, "http://www.google.com", method=method)
+        parseauth.assert_called_once_with('something something')
+        assert method.call_count == 2
+        assert i.auth_token == 'anauthtoken'
+
+    def test_persistent_failure(self, getauth, parseauth):
+        r = requests.Response()
+        r.status_code = 401
+        r.headers['Www-Authenticate'] = 'something something'
+        method = MagicMock(return_value=r)
+        r = requests.Response()
+        r.status_code = 200
+        i = Image("quay.io/foo/bar:latest", username="user", password="pass")
+        getauth.return_value = "anauthtoken"
+        parseauth.return_value = "aparsedauth"
+        with pytest.raises(requests.exceptions.HTTPError):
+            i._request_get.__wrapped__(i, "http://www.google.com", method=method)
+            getauth.assert_called_once()
+            parseauth.assert_called_once()