From 8ea4a7245b96ff24e7296c8aa0525dae7eb0bc2e Mon Sep 17 00:00:00 2001
From: Tomas Kopecek <tkopecek@redhat.com>
Date: Jun 05 2020 08:58:38 +0000
Subject: lib: retry openRemoteFile and check size


---

diff --git a/koji/__init__.py b/koji/__init__.py
index 4dc8c7d..bf39ea9 100644
--- a/koji/__init__.py
+++ b/koji/__init__.py
@@ -55,6 +55,8 @@ import six
 import six.moves.configparser
 import six.moves.http_client
 import six.moves.urllib
+from requests.adapters import HTTPAdapter
+from requests.packages.urllib3.util.retry import Retry
 from six.moves import range, zip
 
 from koji.xmlrpcplus import Fault, dumps, getparser, loads, xmlrpc_client
@@ -1748,6 +1750,18 @@ def format_exc_plus():
     return rv
 
 
+def request_with_retry(retries=3, backoff_factor=0.3,
+                        status_forcelist=(500, 502, 504, 408, 429), session=None):
+    # stolen from https://www.peterbe.com/plog/best-practice-with-retries-with-requests
+    session = session or requests.Session()
+    retry = Retry(total=retries, read=retries, connect=retries,
+                backoff_factor=backoff_factor,
+                status_forcelist=status_forcelist)
+    adapter = HTTPAdapter(max_retries=retry)
+    session.mount('http://', adapter)
+    session.mount('https://', adapter)
+    return session
+
 def openRemoteFile(relpath, topurl=None, topdir=None, tempdir=None):
     """Open a file on the main server (read-only)
 
@@ -1756,12 +1770,15 @@ def openRemoteFile(relpath, topurl=None, topdir=None, tempdir=None):
     if topurl:
         url = "%s/%s" % (topurl, relpath)
         fo = tempfile.TemporaryFile(dir=tempdir)
-        resp = requests.get(url, stream=True)
+        resp = request_with_retry().get(url, stream=True)
         try:
             for chunk in resp.iter_content(chunk_size=8192):
                 fo.write(chunk)
         finally:
             resp.close()
+        if resp.headers.get('Content-Length') and fo.tell() != int(resp.headers['Content-Length']):
+            raise GenericError("Downloaded file %s doesn't match expected size (%s vs %s)" %
+                               (url, fo.tell(), resp.headers['Content-Length']))
         fo.seek(0)
     elif topdir:
         fn = "%s/%s" % (topdir, relpath)
diff --git a/tests/test_lib/test_utils.py b/tests/test_lib/test_utils.py
index c2b2e21..7ec1b22 100644
--- a/tests/test_lib/test_utils.py
+++ b/tests/test_lib/test_utils.py
@@ -161,6 +161,38 @@ class MiscFunctionTestCase(unittest.TestCase):
             for m in mocks:
                 m.assert_not_called()
 
+        for m in mocks:
+            m.reset_mock()
+
+        # downloaded size is larger than content-length
+        with requests_mock.Mocker() as m_requests:
+            text = 'random content'
+            m_requests.register_uri('GET', url, text=text,
+                    headers = {'Content-Length': "3"})
+            m_TemporaryFile.return_value.tell.return_value = len(text)
+            # using neither
+            with self.assertRaises(koji.GenericError):
+                koji.openRemoteFile(path, topurl=topurl)
+            m_TemporaryFile.assert_called_once()
+            m_TemporaryFile.return_value.tell.assert_called()
+            m_open.assert_not_called()
+
+        for m in mocks:
+            m.reset_mock()
+
+        # downloaded size is shorter than content-length
+        with requests_mock.Mocker() as m_requests:
+            text = 'random content'
+            m_requests.register_uri('GET', url, text=text,
+                    headers = {'Content-Length': "100"})
+            m_TemporaryFile.return_value.tell.return_value = len(text)
+            # using neither
+            with self.assertRaises(koji.GenericError):
+                koji.openRemoteFile(path, topurl=topurl)
+            m_TemporaryFile.assert_called_once()
+            m_TemporaryFile.return_value.tell.assert_called()
+            m_open.assert_not_called()
+
     def test_joinpath_bad(self):
         bad_joins = [
             ['/foo', '../bar'],