Merge pull request #1155 from coreos-inc/torrenttest

Add torrent tests

storage/fakestorage.py

@@ -15,6 +15,15 @@ class FakeStorage(BaseStorageV2):
   def _init_path(self, path=None, create=False):
     return path
 
+  def get_direct_download_url(self, path, expires_in=60, requires_cors=False):
+    try:
+      if self.get_content('supports_direct_download') == 'true':
+        return 'http://somefakeurl'
+    except:
+      pass
+
+    return None
+
   def get_content(self, path):
     if not path in _FAKE_STORAGE_MAP:
       raise IOError('Fake file %s not found' % path)

test/registry_tests.py

@@ -4,13 +4,15 @@ import os
 import math
 import random
 import string
+import resumablehashlib
+import binascii
 
 import Crypto.Random
 from flask import request, jsonify
 from flask.blueprints import Blueprint
 from flask.ext.testing import LiveServerTestCase
 
-from app import app
+from app import app, storage
 from data.database import close_db_filter, configure
 from data import model
 from endpoints.v1 import v1_bp
@@ -28,6 +30,7 @@ import json
 import features
 import hashlib
 import logging
+import bencode
 
 import tarfile
 import shutil
@@ -61,6 +64,12 @@ def generate_csrf():
   return generate_csrf_token()
 
 
+@testbp.route('/fakestoragedd/<enabled>', methods=['POST'])
+def set_fakestorage_directdownload(enabled):
+  storage.put_content(['local_us'], 'supports_direct_download', enabled)
+  return 'OK'
+
+
 @testbp.route('/feature/<feature_name>', methods=['POST'])
 def set_feature(feature_name):
   import features
@@ -129,6 +138,26 @@ _CLEAN_DATABASE_PATH = None
 _JWK = RSAKey(key=RSA.generate(2048))
 
 
+def get_full_contents(image_data):
+  if 'chunks' in image_data:
+    # Data is just for chunking; no need for a real TAR.
+    return image_data['contents']
+
+  layer_data = StringIO()
+
+  tar_file_info = tarfile.TarInfo(name='contents')
+  tar_file_info.type = tarfile.REGTYPE
+  tar_file_info.size = len(image_data['contents'])
+
+  tar_file = tarfile.open(fileobj=layer_data, mode='w|gz')
+  tar_file.addfile(tar_file_info, StringIO(image_data['contents']))
+  tar_file.close()
+
+  layer_bytes = layer_data.getvalue()
+  layer_data.close()
+  return layer_bytes
+
+
 def get_new_database_uri():
   # If a clean copy of the database has not yet been created, create one now.
   global _CLEAN_DATABASE_PATH
@@ -281,19 +310,7 @@ class V1RegistryPushMixin(V1RegistryMixin):
                    data=json.dumps(image_json_data), auth='sig')
 
       # PUT /v1/images/{imageID}/layer
-      tar_file_info = tarfile.TarInfo(name='image_name')
+      layer_bytes = get_full_contents(image_data)
-      tar_file_info.type = tarfile.REGTYPE
-      tar_file_info.size = len(image_id)
-
-      layer_data = StringIO()
-
-      tar_file = tarfile.open(fileobj=layer_data, mode='w|gz')
-      tar_file.addfile(tar_file_info, StringIO(image_id))
-      tar_file.close()
-
-      layer_bytes = layer_data.getvalue()
-      layer_data.close()
-
       self.conduct('PUT', '/v1/images/%s/layer' % image_id,
                    data=StringIO(layer_bytes), auth='sig')
 
@@ -436,8 +453,12 @@ class V2RegistryPushMixin(V2RegistryMixin):
     # Build a fake manifest.
     tag_name = tag_name or 'latest'
     builder = SignedManifestBuilder(namespace, repository, tag_name)
+    full_contents = {}
+
     for image_data in images:
-      checksum = 'sha256:' + hashlib.sha256(image_data['contents']).hexdigest()
+      full_contents[image_data['id']] = get_full_contents(image_data)
+
+      checksum = 'sha256:' + hashlib.sha256(full_contents[image_data['id']]).hexdigest()
       if invalid:
         checksum = 'sha256:' + hashlib.sha256('foobarbaz').hexdigest()
 
@@ -450,11 +471,11 @@ class V2RegistryPushMixin(V2RegistryMixin):
     checksums = {}
     for image_data in images:
       image_id = image_data['id']
-      full_contents = image_data['contents']
+      layer_bytes = full_contents[image_data['id']]
       chunks = image_data.get('chunks')
 
       # Layer data should not yet exist.
-      checksum = 'sha256:' + hashlib.sha256(full_contents).hexdigest()
+      checksum = 'sha256:' + hashlib.sha256(layer_bytes).hexdigest()
       self.conduct('HEAD', '/v2/%s/%s/blobs/%s' % (namespace, repository, checksum),
                    expected_code=404, auth='jwt')
 
@@ -467,7 +488,7 @@ class V2RegistryPushMixin(V2RegistryMixin):
 
       # PATCH the image data into the layer.
       if chunks is None:
-        self.conduct('PATCH', location, data=full_contents, expected_code=204, auth='jwt')
+        self.conduct('PATCH', location, data=layer_bytes, expected_code=204, auth='jwt')
       else:
         for chunk in chunks:
           if len(chunk) == 3:
@@ -476,7 +497,7 @@ class V2RegistryPushMixin(V2RegistryMixin):
             (start_byte, end_byte) = chunk
             expected_code = 204
 
-          contents_chunk = full_contents[start_byte:end_byte]
+          contents_chunk = layer_bytes[start_byte:end_byte]
           self.conduct('PATCH', location, data=contents_chunk, expected_code=expected_code, auth='jwt',
                        headers={'Range': 'bytes=%s-%s' % (start_byte, end_byte)})
 
@@ -511,7 +532,7 @@ class V2RegistryPushMixin(V2RegistryMixin):
       response = self.conduct('HEAD', '/v2/%s/%s/blobs/%s' % (namespace, repository, checksum),
                               expected_code=200, auth='jwt')
       self.assertEquals(response.headers['Docker-Content-Digest'], checksum)
-      self.assertEquals(response.headers['Content-Length'], str(len(full_contents)))
+      self.assertEquals(response.headers['Content-Length'], str(len(layer_bytes)))
 
     # Write the manifest.
     put_code = 404 if invalid else expected_manifest_code
@@ -555,7 +576,7 @@ class V2RegistryPullMixin(V2RegistryMixin):
       result = self.conduct('GET', '/v2/%s/%s/blobs/%s' % (namespace, repository, blob_id),
                             expected_code=200, auth='jwt')
 
-      blobs[blob_id] = result.text
+      blobs[blob_id] = result.content
 
     # Verify the V1 metadata is present for each expected image.
     found_v1_layers = set()
@@ -570,6 +591,39 @@ class V2RegistryPullMixin(V2RegistryMixin):
     return blobs
 
 
+class V1RegistryLoginMixin(object):
+  def do_login(self, username, password, scope, expect_success=True):
+    data = {
+      'username': username,
+      'password': password,
+    }
+
+    response = self.conduct('POST', '/v1/users/', json_data=data, expected_code=400)
+    if expect_success:
+      self.assertEquals(response.text, '"Username or email already exists"')
+    else:
+      self.assertNotEquals(response.text, '"Username or email already exists"')
+
+
+class V2RegistryLoginMixin(object):
+  def do_login(self, username, password, scope, expect_success=True, expected_code=None):
+    params = {
+      'account': username,
+      'scope': scope,
+      'service': app.config['SERVER_HOSTNAME'],
+    }
+
+    if expected_code is None:
+      if expect_success:
+        expected_code = 200
+      else:
+        expected_code = 403
+
+    response = self.conduct('GET', '/v2/auth', params=params, auth=(username, password),
+                            expected_code=expected_code)
+    return response
+
+
 class RegistryTestsMixin(object):
   def test_push_pull_logging(self):
     # Push a new repository.
@@ -815,7 +869,8 @@ class V1RegistryTests(V1RegistryPullMixin, V1RegistryPushMixin, RegistryTestsMix
   def test_push_reponame_with_slashes(self):
     # Attempt to add a repository name with slashes. This should fail as we do not support it.
     images = [{
-      'id': 'onlyimagehere'
+      'id': 'onlyimagehere',
+      'contents': 'somecontents',
     }]
     self.do_push('public', 'newrepo/somesubrepo', 'public', 'password', images, expected_code=400)
 
@@ -824,7 +879,8 @@ class V1RegistryTests(V1RegistryPullMixin, V1RegistryPushMixin, RegistryTestsMix
 
     images = [{
       'id': 'onlyimagehere',
-      'comment': 'Pawe\xc5\x82 Kami\xc5\x84ski <pawel.kaminski@codewise.com>'.decode('utf-8')
+      'comment': 'Pawe\xc5\x82 Kami\xc5\x84ski <pawel.kaminski@codewise.com>'.decode('utf-8'),
+      'contents': 'somecontents',
     }]
 
     self.do_push('devtable', 'unicodetest', 'devtable', 'password', images)
@@ -833,8 +889,10 @@ class V1RegistryTests(V1RegistryPullMixin, V1RegistryPushMixin, RegistryTestsMix
   def test_tag_validation(self):
     image_id = 'onlyimagehere'
     images = [{
-      'id': image_id
+      'id': image_id,
+      'contents': 'somecontents',
     }]
+
     self.do_push('public', 'newrepo', 'public', 'password', images)
     self.do_tag('public', 'newrepo', '1', image_id)
     self.do_tag('public', 'newrepo', 'x' * 128, image_id)
@@ -887,8 +945,10 @@ class V2RegistryTests(V2RegistryPullMixin, V2RegistryPushMixin, RegistryTestsMix
   def test_push_reponame_with_slashes(self):
     # Attempt to add a repository name with slashes. This should fail as we do not support it.
     images = [{
-      'id': 'onlyimagehere'
+      'id': 'onlyimagehere',
+      'contents': 'somecontents',
     }]
+
     self.do_push('public', 'newrepo/somesubrepo', 'devtable', 'password', images,
                  expected_auth_code=400)
 
@@ -1081,13 +1141,74 @@ class V1PushV2PullRegistryTests(V2RegistryPullMixin, V1RegistryPushMixin, Regist
                                 RegistryTestCaseMixin, LiveServerTestCase):
   """ Tests for V1 push, V2 pull registry. """
 
+
 class V1PullV2PushRegistryTests(V1RegistryPullMixin, V2RegistryPushMixin, RegistryTestsMixin,
                                 RegistryTestCaseMixin, LiveServerTestCase):
   """ Tests for V1 pull, V2 push registry. """
 
 
-class VerbTests(RegistryTestCaseMixin, V1RegistryPushMixin, LiveServerTestCase):
+class TorrentTestMixin(V2RegistryPullMixin):
-  """ Tests for registry verbs. """
+  """ Mixin of tests for torrent support. """
+  def get_torrent(self, blobsum):
+    # Enable direct download URLs in fake storage.
+    self.conduct('POST', '/__test/fakestoragedd/true')
+
+    response = self.conduct('GET', '/c1/torrent/devtable/newrepo/blobs/' + blobsum,
+                            auth=('devtable', 'password'))
+
+    # Disable direct download URLs in fake storage.
+    self.conduct('POST', '/__test/fakestoragedd/false')
+
+    return response.content
+
+  def test_get_basic_torrent(self):
+    initial_images = [
+      {
+        'id': 'initialid',
+        'contents': 'the initial image',
+      },
+    ]
+
+    # Create the repo.
+    self.do_push('devtable', 'newrepo', 'devtable', 'password', images=initial_images)
+
+    # Retrieve the manifest for the tag.
+    blobs = self.do_pull('devtable', 'newrepo', 'devtable', 'password', manifest_id='latest',
+                         images=initial_images)
+    self.assertEquals(1, len(list(blobs.keys())))
+    blobsum = list(blobs.keys())[0]
+
+    # Retrieve the torrent for the tag.
+    torrent = self.get_torrent(blobsum)
+    contents = bencode.bdecode(torrent)
+
+    # Ensure that there is a webseed.
+    self.assertEquals(contents['url-list'], 'http://somefakeurl')
+
+    # Ensure there is an announce and some pieces.
+    self.assertIsNotNone(contents.get('info', {}).get('pieces'))
+    self.assertIsNotNone(contents.get('announce'))
+
+    sha = resumablehashlib.sha1()
+    sha.update(blobs[blobsum])
+
+    expected = binascii.hexlify(sha.digest())
+    found = binascii.hexlify(contents['info']['pieces'])
+
+    self.assertEquals(expected, found)
+
+
+class TorrentV1PushTests(RegistryTestCaseMixin, TorrentTestMixin, V1RegistryPushMixin, LiveServerTestCase):
+  """ Torrent tests via V1 push. """
+  pass
+
+class TorrentV2PushTests(RegistryTestCaseMixin, TorrentTestMixin, V2RegistryPushMixin, LiveServerTestCase):
+  """ Torrent tests via V2 push. """
+  pass
+
+
+class SquashingTests(RegistryTestCaseMixin, V1RegistryPushMixin, LiveServerTestCase):
+  """ Tests for registry squashing. """
 
   def get_squashed_image(self):
     response = self.conduct('GET', '/c1/squash/devtable/newrepo/latest', auth='sig')
@@ -1184,39 +1305,8 @@ class VerbTests(RegistryTestCaseMixin, V1RegistryPushMixin, LiveServerTestCase):
 
     # Ensure that the "image_name" file refers to the latest image, as it is the top layer.
     layer_tar = tarfile.open(fileobj=tar.extractfile(tar.getmember('%s/layer.tar' % expected_image_id)))
-    image_name = layer_tar.extractfile(layer_tar.getmember('image_name')).read()
+    image_contents = layer_tar.extractfile(layer_tar.getmember('contents')).read()
-    self.assertEquals('latestid', image_name)
+    self.assertEquals('the latest image', image_contents)
-
-class V1RegistryLoginMixin(object):
-  def do_login(self, username, password, scope, expect_success=True):
-    data = {
-      'username': username,
-      'password': password,
-    }
-
-    response = self.conduct('POST', '/v1/users/', json_data=data, expected_code=400)
-    if expect_success:
-      self.assertEquals(response.text, '"Username or email already exists"')
-    else:
-      self.assertNotEquals(response.text, '"Username or email already exists"')
-
-class V2RegistryLoginMixin(object):
-  def do_login(self, username, password, scope, expect_success=True, expected_code=None):
-    params = {
-      'account': username,
-      'scope': scope,
-      'service': app.config['SERVER_HOSTNAME'],
-    }
-
-    if expected_code is None:
-      if expect_success:
-        expected_code = 200
-      else:
-        expected_code = 403
-
-    response = self.conduct('GET', '/v2/auth', params=params, auth=(username, password),
-                            expected_code=expected_code)
-    return response
 
 
 class LoginTests(object):