Add feature flag to force all direct download URLs to be proxied

Fixes #1667

test/test_storageproxy.py

@@ -0,0 +1,89 @@
+import unittest
+import requests
+import os
+
+from flask import Flask
+from flask.ext.testing import LiveServerTestCase
+from initdb import setup_database_for_testing, finished_database_for_testing
+from util.security.instancekeys import InstanceKeys
+from storage import Storage
+
+_PORT_NUMBER = 5001
+
+class TestStorageProxy(LiveServerTestCase):
+  def setUp(self):
+    setup_database_for_testing(self)
+
+  def tearDown(self):
+    finished_database_for_testing(self)
+
+  def create_app(self):
+    global _PORT_NUMBER
+    _PORT_NUMBER = _PORT_NUMBER + 1
+
+    self.test_app = Flask('teststorageproxy')
+    self.test_app.config['LIVESERVER_PORT'] = _PORT_NUMBER
+
+    if os.environ.get('DEBUG') == 'true':
+      self.test_app.config['DEBUG'] = True
+
+    self.test_app.config['TESTING'] = True
+    self.test_app.config['SERVER_HOSTNAME'] = 'localhost:%s' % _PORT_NUMBER
+
+    self.test_app.config['INSTANCE_SERVICE_KEY_KID_LOCATION'] = 'test/data/test.kid'
+    self.test_app.config['INSTANCE_SERVICE_KEY_LOCATION'] = 'test/data/test.pem'
+    self.test_app.config['INSTANCE_SERVICE_KEY_SERVICE'] = 'quay'
+
+    # UGH... Such a stupid hack!
+    self.test_app.config['FEATURE_PROXY_STORAGE'] = self.id().find('notinstalled') < 0
+
+    self.test_app.config['DISTRIBUTED_STORAGE_CONFIG'] = {
+      'test': ['FakeStorage', {}],
+    }
+
+    instance_keys = InstanceKeys(self.test_app)
+    self.storage = Storage(self.test_app, instance_keys=instance_keys)
+    self.test_app.config['DISTRIBUTED_STORAGE_PREFERENCE'] = ['test']
+    return self.test_app
+
+  def test_storage_proxy_auth_notinstalled(self):
+    # Active direct download on the fake storage.
+    self.storage.put_content(['test'], 'supports_direct_download', 'true')
+
+    # Get the unwrapped URL.
+    direct_download_url = self.storage.get_direct_download_url(['test'], 'somepath')
+    self.assertEquals(-1, direct_download_url.find('/_storage_proxy/'))
+
+    # Ensure that auth returns 404.
+    headers = {
+      'X-Original-URI': 'someurihere'
+    }
+
+    resp = requests.get('http://%s/_storage_proxy_auth' % self.test_app.config['SERVER_HOSTNAME'],
+                        headers=headers)
+    self.assertEquals(404, resp.status_code)
+
+
+  def test_storage_proxy_auth(self):
+    # Active direct download on the fake storage.
+    self.storage.put_content(['test'], 'supports_direct_download', 'true')
+
+    # Get the wrapped URL.
+    direct_download_url = self.storage.get_direct_download_url(['test'], 'somepath')
+
+    # Ensure it refers to the storage proxy.
+    proxy_index = direct_download_url.find('/_storage_proxy/')
+    self.assertTrue(proxy_index > 0)
+
+    # Ensure that auth returns 200 for the URL pieces.
+    headers = {
+      'X-Original-URI': direct_download_url[proxy_index:]
+    }
+
+    resp = requests.get('http://%s/_storage_proxy_auth' % self.test_app.config['SERVER_HOSTNAME'],
+                        headers=headers)
+    self.assertEquals(200, resp.status_code)
+
+
+if __name__ == '__main__':
+  unittest.main()