From 30f072aeff832d21e2b2c13f2ce241800ec900ba Mon Sep 17 00:00:00 2001
From: Joseph Schorr <joseph.schorr@coreos.com>
Date: Mon, 12 Nov 2018 23:27:49 +0200
Subject: [PATCH] Add support for creating schema 2 manifests and manifest
 lists via the OCI model

---
 data/model/oci/manifest.py                    | 173 +++++++++++++---
 data/model/oci/test/test_oci_manifest.py      | 194 ++++++++++++++++--
 data/registry_model/datatypes.py              |   5 +-
 data/registry_model/interface.py              |   5 +-
 data/registry_model/manifestbuilder.py        |  13 +-
 data/registry_model/registry_oci_model.py     |  57 ++---
 data/registry_model/registry_pre_oci_model.py |   5 +-
 data/registry_model/test/test_interface.py    |  10 +-
 .../test/test_manifestbuilder.py              |   9 +-
 endpoints/api/tag.py                          |   6 +-
 endpoints/v1/index.py                         |   6 +-
 endpoints/v1/registry.py                      |   6 +-
 endpoints/v1/tag.py                           |   5 +-
 endpoints/v2/manifest.py                      |   5 +-
 initdb.py                                     |   4 +
 test/test_api_usage.py                        |   5 +-
 16 files changed, 398 insertions(+), 110 deletions(-)

diff --git a/data/model/oci/manifest.py b/data/model/oci/manifest.py
index a677df0ab..2cc177ecb 100644
--- a/data/model/oci/manifest.py
+++ b/data/model/oci/manifest.py
@@ -1,15 +1,28 @@
 import logging
 
+from collections import namedtuple
+
 from peewee import IntegrityError
 
-from data.database import Tag, Manifest, ManifestBlob, ManifestLegacyImage, db_transaction
+from data.database import (Tag, Manifest, ManifestBlob, ManifestLegacyImage, ManifestChild,
+                           db_transaction)
+from data.model import BlobDoesNotExist
 from data.model.oci.tag import filter_to_alive_tags
-from data.model.storage import lookup_repo_storages_by_content_checksum
+from data.model.oci.label import create_manifest_label
+from data.model.storage import (lookup_repo_storages_by_content_checksum, get_storage_locations,
+                                get_layer_path)
+from data.model.blob import get_repository_blob_by_digest
 from data.model.image import lookup_repository_images, get_image, synthesize_v1_image
-from image.docker.schema1 import DockerSchema1Manifest, ManifestException
+from image.docker.schema1 import ManifestException
+from image.docker.schema2.list import MalformedSchema2ManifestList
+from util.validation import is_json
+
 
 logger = logging.getLogger(__name__)
 
+CreatedManifest = namedtuple('CreatedManifest', ['manifest', 'newly_created', 'labels_to_apply'])
+
+
 def lookup_manifest(repository_id, manifest_digest, allow_dead=False):
   """ Returns the manifest with the specified digest under the specified repository
       or None if none. If allow_dead is True, then manifests referenced by only
@@ -29,43 +42,97 @@ def lookup_manifest(repository_id, manifest_digest, allow_dead=False):
     return None
 
 
-def get_or_create_manifest(repository_id, manifest_interface_instance):
-  """ Returns a tuple of the manifest in the specified repository with the matching digest
-      (if it already exists) or, if not yet created, creates and returns the manifest, as well as
-      if the manifest was created. Returns (None, None) if there was an error creating the manifest.
+def get_or_create_manifest(repository_id, manifest_interface_instance, storage):
+  """ Returns a CreatedManifest for the manifest in the specified repository with the matching
+      digest (if it already exists) or, if not yet created, creates and returns the manifest.
+      Returns None if there was an error creating the manifest.
       Note that *all* blobs referenced by the manifest must exist already in the repository or this
-      method will fail with a (None, None).
+      method will fail with a None.
   """
   existing = lookup_manifest(repository_id, manifest_interface_instance.digest, allow_dead=True)
   if existing is not None:
-    return existing, False
+    return CreatedManifest(manifest=existing, newly_created=False, labels_to_apply=None)
 
-  assert len(list(manifest_interface_instance.layers)) > 0
+  return _create_manifest(repository_id, manifest_interface_instance, storage)
 
-  # TODO(jschorr): Switch this to supporting schema2 once we're ready.
-  assert isinstance(manifest_interface_instance, DockerSchema1Manifest)
+
+def _create_manifest(repository_id, manifest_interface_instance, storage):
+  digests = set(manifest_interface_instance.blob_digests)
+
+  def _lookup_digest(digest):
+    return _retrieve_bytes_in_storage(repository_id, digest, storage)
+
+  # Retrieve the child manifests, if any. If we do retrieve a child manifest, we also remove its
+  # blob from the list of blobs for this manifest, as the blob isn't really a "blob".
+  child_manifest_refs = manifest_interface_instance.child_manifests(_lookup_digest)
+  child_manifest_rows = []
+  child_manifest_label_dicts = []
+
+  if child_manifest_refs is not None:
+    for child_manifest_ref in child_manifest_refs:
+      # Load and parse the child manifest.
+      try:
+        child_manifest = child_manifest_ref.manifest_obj
+      except ManifestException:
+        logger.exception('Could not load manifest list for manifest `%s`',
+                         manifest_interface_instance.digest)
+        return None
+      except MalformedSchema2ManifestList:
+        logger.exception('Could not load manifest list for manifest `%s`',
+                         manifest_interface_instance.digest)
+        return None
+      except BlobDoesNotExist:
+        logger.exception('Could not load manifest list for manifest `%s`',
+                         manifest_interface_instance.digest)
+        return None
+      except IOError:
+        logger.exception('Could not load manifest list for manifest `%s`',
+                         manifest_interface_instance.digest)
+        return None
+
+      # Retrieve its labels.
+      labels = child_manifest.get_manifest_labels(_lookup_digest)
+      if labels is None:
+        logger.exception('Could not load manifest labels for child manifest')
+        return None
+
+      # Get/create the child manifest in the database.
+      assert list(child_manifest.layers)
+      child_manifest_info = get_or_create_manifest(repository_id, child_manifest, storage)
+      if child_manifest_info is None:
+        logger.error('Could not get/create child manifest')
+        return None
+
+      child_manifest_rows.append(child_manifest_info.manifest)
+      child_manifest_label_dicts.append(labels)
+      digests.remove(child_manifest.digest)
 
   # Ensure all the blobs in the manifest exist.
-  digests = manifest_interface_instance.checksums
-  query = lookup_repo_storages_by_content_checksum(repository_id, digests)
-  blob_map = {s.content_checksum: s for s in query}
-  for digest_str in manifest_interface_instance.blob_digests:
-    if digest_str not in blob_map:
-      logger.warning('Unknown blob `%s` under manifest `%s` for repository `%s`', digest_str,
-                     manifest_interface_instance.digest, repository_id)
-      return None, None
+  blob_map = {}
+  if digests:
+    query = lookup_repo_storages_by_content_checksum(repository_id, digests)
+    blob_map = {s.content_checksum: s for s in query}
+    for digest_str in digests:
+      if digest_str not in blob_map:
+        logger.warning('Unknown blob `%s` under manifest `%s` for repository `%s`', digest_str,
+                       manifest_interface_instance.digest, repository_id)
+        return None
 
-  # Determine and populate the legacy image if necessary.
-  legacy_image_id = _populate_legacy_image(repository_id, manifest_interface_instance, blob_map)
-  if legacy_image_id is None:
-    return None, None
+  # Determine and populate the legacy image if necessary. Manifest lists will not have a legacy
+  # image.
+  legacy_image = None
+  if manifest_interface_instance.leaf_layer_v1_image_id is not None:
+    legacy_image_id = _populate_legacy_image(repository_id, manifest_interface_instance, blob_map,
+                                             storage)
+    if legacy_image_id is None:
+      return None
 
-  legacy_image = get_image(repository_id, legacy_image_id)
-  if legacy_image is None:
-    return None, None
+    legacy_image = get_image(repository_id, legacy_image_id)
+    if legacy_image is None:
+      return None
 
   # Create the manifest and its blobs.
-  media_type = Manifest.media_type.get_id(manifest_interface_instance.content_type)
+  media_type = Manifest.media_type.get_id(manifest_interface_instance.media_type)
   storage_ids = {storage.id for storage in blob_map.values()}
 
   with db_transaction():
@@ -77,7 +144,7 @@ def get_or_create_manifest(repository_id, manifest_interface_instance):
                                  manifest_bytes=manifest_interface_instance.bytes)
     except IntegrityError:
       manifest = Manifest.get(repository=repository_id, digest=manifest_interface_instance.digest)
-      return manifest, False
+      return CreatedManifest(manifest=manifest, newly_created=False, labels_to_apply=None)
 
     # Insert the blobs.
     blobs_to_insert = [dict(manifest=manifest, repository=repository_id,
@@ -86,12 +153,42 @@ def get_or_create_manifest(repository_id, manifest_interface_instance):
       ManifestBlob.insert_many(blobs_to_insert).execute()
 
     # Set the legacy image (if applicable).
-    ManifestLegacyImage.create(repository=repository_id, image=legacy_image, manifest=manifest)
+    if legacy_image is not None:
+      ManifestLegacyImage.create(repository=repository_id, image=legacy_image, manifest=manifest)
 
-    return manifest, True
+    # Insert the manifest child rows (if applicable).
+    if child_manifest_rows:
+      children_to_insert = [dict(manifest=manifest, child_manifest=child_manifest,
+                                 repository=repository_id)
+                            for child_manifest in child_manifest_rows]
+      ManifestChild.insert_many(children_to_insert).execute()
+
+  # Define the labels for the manifest (if any).
+  labels = manifest_interface_instance.get_manifest_labels(_lookup_digest)
+  if labels:
+    for key, value in labels.iteritems():
+      media_type = 'application/json' if is_json(value) else 'text/plain'
+      create_manifest_label(manifest, key, value, 'manifest', media_type)
+
+  # Return the dictionary of labels to apply. We only return those labels either defined on
+  # the manifest or shared amongst all the child manifest.
+  labels_to_apply = labels or {}
+  if child_manifest_label_dicts:
+    labels_to_apply = child_manifest_label_dicts[0].viewitems()
+    for child_manifest_label_dict in child_manifest_label_dicts[1:]:
+      # Intersect the key+values of the labels to ensure we get the exact same result
+      # for all the child manifests.
+      labels_to_apply = labels_to_apply & child_manifest_label_dict.viewitems()
+
+    labels_to_apply = dict(labels_to_apply)
+
+  return CreatedManifest(manifest=manifest, newly_created=True, labels_to_apply=labels_to_apply)
 
 
-def _populate_legacy_image(repository_id, manifest_interface_instance, blob_map):
+def _populate_legacy_image(repository_id, manifest_interface_instance, blob_map, storage):
+  def _lookup_digest(digest):
+    return _retrieve_bytes_in_storage(repository_id, digest, storage)
+
   # Lookup all the images and their parent images (if any) inside the manifest.
   # This will let us know which v1 images we need to synthesize and which ones are invalid.
   docker_image_ids = list(manifest_interface_instance.legacy_image_ids)
@@ -100,7 +197,8 @@ def _populate_legacy_image(repository_id, manifest_interface_instance, blob_map)
 
   # Rewrite any v1 image IDs that do not match the checksum in the database.
   try:
-    rewritten_images = manifest_interface_instance.rewrite_invalid_image_ids(image_storage_map)
+    rewritten_images = manifest_interface_instance.generate_legacy_layers(image_storage_map,
+                                                                          _lookup_digest)
     rewritten_images = list(rewritten_images)
     parent_image_map = {}
 
@@ -132,3 +230,12 @@ def _populate_legacy_image(repository_id, manifest_interface_instance, blob_map)
     return None
 
   return rewritten_images[-1].image_id
+
+
+def _retrieve_bytes_in_storage(repository_id, digest, storage):
+  blob = get_repository_blob_by_digest(repository_id, digest)
+  if blob is None:
+    return None
+
+  placements = list(get_storage_locations(blob.uuid))
+  return storage.get_content(placements, get_layer_path(blob))
diff --git a/data/model/oci/test/test_oci_manifest.py b/data/model/oci/test/test_oci_manifest.py
index 70c758c88..8d3b77b2f 100644
--- a/data/model/oci/test/test_oci_manifest.py
+++ b/data/model/oci/test/test_oci_manifest.py
@@ -1,13 +1,22 @@
+import json
+
 from playhouse.test_utils import assert_query_count
 
-from app import docker_v2_signing_key
+from app import docker_v2_signing_key, storage
 
-from data.database import Tag, ManifestBlob, get_epoch_timestamp_ms
+from digest.digest_tools import sha256_digest
+from data.database import Tag, ManifestBlob, ImageStorageLocation, ManifestChild, get_epoch_timestamp_ms
 from data.model.oci.manifest import lookup_manifest, get_or_create_manifest
 from data.model.oci.tag import filter_to_alive_tags, get_tag
 from data.model.oci.shared import get_legacy_image_for_manifest
-from data.model.repository import get_repository
+from data.model.oci.label import list_manifest_labels
+from data.model.repository import get_repository, create_repository
+from data.model.image import find_create_or_link_image
+from data.model.blob import store_blob_record_and_temp_link
+from data.model.storage import get_layer_path
 from image.docker.schema1 import DockerSchema1ManifestBuilder, DockerSchema1Manifest
+from image.docker.schema2.manifest import DockerSchema2ManifestBuilder
+from image.docker.schema2.list import DockerSchema2ManifestListBuilder
 
 from test.fixtures import *
 
@@ -38,35 +47,104 @@ def test_lookup_manifest_dead_tag(initialized_db):
           dead_tag.manifest)
 
 
-def test_get_or_create_manifest(initialized_db):
-  repository = get_repository('devtable', 'simple')
+def _populate_blob(content):
+  digest = str(sha256_digest(content))
+  location = ImageStorageLocation.get(name='local_us')
+  blob = store_blob_record_and_temp_link('devtable', 'newrepo', digest, location,  len(content), 120)
+  storage.put_content(['local_us'], get_layer_path(blob), content)
+  return blob, digest
 
-  latest_tag = get_tag(repository, 'latest')
-  legacy_image = get_legacy_image_for_manifest(latest_tag.manifest)
-  parsed = DockerSchema1Manifest(latest_tag.manifest.manifest_bytes, validate=False)
 
-  builder = DockerSchema1ManifestBuilder('devtable', 'simple', 'anothertag')
-  builder.add_layer(parsed.blob_digests[0], '{"id": "%s"}' % legacy_image.docker_image_id)
-  sample_manifest_instance = builder.build(docker_v2_signing_key)
+@pytest.mark.parametrize('schema_version', [
+  1,
+  2,
+])
+def test_get_or_create_manifest(schema_version, initialized_db):
+  repository = create_repository('devtable', 'newrepo', None)
+
+  expected_labels = {
+    'Foo': 'Bar',
+    'Baz': 'Meh',
+  }
+
+  layer_json = json.dumps({
+    'id': 'somelegacyid',
+    'config': {
+      'Labels': expected_labels,
+    },
+    "rootfs": {
+      "type": "layers",
+      "diff_ids": []
+    },
+    "history": [
+      {
+        "created": "2018-04-03T18:37:09.284840891Z",
+        "created_by": "do something",
+      },
+    ],
+  })
+
+  # Create a legacy image.
+  find_create_or_link_image('somelegacyid', repository, 'devtable', {}, 'local_us')
+
+  # Add a blob containing the config.
+  _, config_digest = _populate_blob(layer_json)
+
+  # Add a blob of random data.
+  random_data = 'hello world'
+  _, random_digest = _populate_blob(random_data)
+
+  # Build the manifest.
+  if schema_version == 1:
+    builder = DockerSchema1ManifestBuilder('devtable', 'simple', 'anothertag')
+    builder.add_layer(random_digest, layer_json)
+    sample_manifest_instance = builder.build(docker_v2_signing_key)
+  elif schema_version == 2:
+    builder = DockerSchema2ManifestBuilder()
+    builder.set_config_digest(config_digest, len(layer_json))
+    builder.add_layer(random_digest, len(random_data))
+    sample_manifest_instance = builder.build()
 
   # Create a new manifest.
-  created, newly_created = get_or_create_manifest(repository, sample_manifest_instance)
+  created_manifest = get_or_create_manifest(repository, sample_manifest_instance, storage)
+  created = created_manifest.manifest
+  newly_created = created_manifest.newly_created
+
   assert newly_created
   assert created is not None
+  assert created.media_type.name == sample_manifest_instance.media_type
   assert created.digest == sample_manifest_instance.digest
   assert created.manifest_bytes == sample_manifest_instance.bytes
+  assert created_manifest.labels_to_apply == expected_labels
 
-  assert get_legacy_image_for_manifest(created) is not None
+  # Verify the legacy image.
+  legacy_image = get_legacy_image_for_manifest(created)
+  assert legacy_image is not None
+  assert legacy_image.storage.content_checksum == random_digest
 
+  # Verify the linked blobs.
   blob_digests = [mb.blob.content_checksum for mb
                   in ManifestBlob.select().where(ManifestBlob.manifest == created)]
-  assert parsed.blob_digests[0] in blob_digests
+
+  assert random_digest in blob_digests
+  if schema_version == 2:
+    assert config_digest in blob_digests
 
   # Retrieve it again and ensure it is the same manifest.
-  created2, newly_created2 = get_or_create_manifest(repository, sample_manifest_instance)
+  created_manifest2 = get_or_create_manifest(repository, sample_manifest_instance, storage)
+  created2 = created_manifest2.manifest
+  newly_created2 = created_manifest2.newly_created
+
   assert not newly_created2
   assert created2 == created
 
+  # Ensure the labels were added.
+  labels = list(list_manifest_labels(created))
+  assert len(labels) == 2
+
+  labels_dict = {label.key: label.value for label in labels}
+  assert labels_dict == expected_labels
+
 
 def test_get_or_create_manifest_invalid_image(initialized_db):
   repository = get_repository('devtable', 'simple')
@@ -78,6 +156,86 @@ def test_get_or_create_manifest_invalid_image(initialized_db):
   builder.add_layer(parsed.blob_digests[0], '{"id": "foo", "parent": "someinvalidimageid"}')
   sample_manifest_instance = builder.build(docker_v2_signing_key)
 
-  created, newly_created = get_or_create_manifest(repository, sample_manifest_instance)
-  assert created is None
-  assert newly_created is None
+  created_manifest = get_or_create_manifest(repository, sample_manifest_instance, storage)
+  assert created_manifest is None
+
+
+def test_get_or_create_manifest_list(initialized_db):
+  repository = create_repository('devtable', 'newrepo', None)
+
+  expected_labels = {
+    'Foo': 'Bar',
+    'Baz': 'Meh',
+  }
+
+  layer_json = json.dumps({
+    'id': 'somelegacyid',
+    'config': {
+      'Labels': expected_labels,
+    },
+    "rootfs": {
+      "type": "layers",
+      "diff_ids": []
+    },
+    "history": [
+      {
+        "created": "2018-04-03T18:37:09.284840891Z",
+        "created_by": "do something",
+      },
+    ],
+  })
+
+  # Create a legacy image.
+  find_create_or_link_image('somelegacyid', repository, 'devtable', {}, 'local_us')
+
+  # Add a blob containing the config.
+  _, config_digest = _populate_blob(layer_json)
+
+  # Add a blob of random data.
+  random_data = 'hello world'
+  _, random_digest = _populate_blob(random_data)
+
+  # Build the manifests.
+  v1_builder = DockerSchema1ManifestBuilder('devtable', 'simple', 'anothertag')
+  v1_builder.add_layer(random_digest, layer_json)
+  v1_manifest = v1_builder.build(docker_v2_signing_key).unsigned()
+
+  v2_builder = DockerSchema2ManifestBuilder()
+  v2_builder.set_config_digest(config_digest, len(layer_json))
+  v2_builder.add_layer(random_digest, len(random_data))
+  v2_manifest = v2_builder.build()
+
+  # Write the manifests as blobs.
+  location = ImageStorageLocation.get(name='local_us')
+  blob = store_blob_record_and_temp_link('devtable', 'newrepo', v1_manifest.digest, location,
+                                         len(v1_manifest.bytes), 120)
+  storage.put_content(['local_us'], get_layer_path(blob), v1_manifest.bytes)
+
+  blob = store_blob_record_and_temp_link('devtable', 'newrepo', v2_manifest.digest, location,
+                                         len(v2_manifest.bytes), 120)
+  storage.put_content(['local_us'], get_layer_path(blob), v2_manifest.bytes)
+
+  # Build the manifest list.
+  list_builder = DockerSchema2ManifestListBuilder()
+  list_builder.add_manifest(v1_manifest, 'amd64', 'linux')
+  list_builder.add_manifest(v2_manifest, 'amd32', 'linux')
+  manifest_list = list_builder.build()
+
+  # Write the manifest list, which should also write the manifests themselves.
+  created_tuple = get_or_create_manifest(repository, manifest_list, storage)
+  assert created_tuple is not None
+
+  created_list = created_tuple.manifest
+  assert created_list
+  assert created_list.media_type.name == manifest_list.media_type
+  assert created_list.digest == manifest_list.digest
+
+  # Ensure the child manifest links exist.
+  child_manifests = {cm.child_manifest.digest: cm.child_manifest
+                     for cm in ManifestChild.select().where(ManifestChild.manifest == created_list)}
+  assert len(child_manifests) == 2
+  assert v1_manifest.digest in child_manifests
+  assert v2_manifest.digest in child_manifests
+
+  assert child_manifests[v1_manifest.digest].media_type.name == v1_manifest.media_type
+  assert child_manifests[v2_manifest.digest].media_type.name == v2_manifest.media_type
diff --git a/data/registry_model/datatypes.py b/data/registry_model/datatypes.py
index 798e4df3e..8c9fde328 100644
--- a/data/registry_model/datatypes.py
+++ b/data/registry_model/datatypes.py
@@ -7,7 +7,8 @@ from cachetools import lru_cache
 
 from data import model
 from data.registry_model.datatype import datatype, requiresinput, optionalinput
-from image.docker.schema1 import DockerSchema1Manifest, DOCKER_SCHEMA1_SIGNED_MANIFEST_CONTENT_TYPE
+from image.docker.schemas import parse_manifest_from_bytes
+from image.docker.schema1 import DOCKER_SCHEMA1_SIGNED_MANIFEST_CONTENT_TYPE
 
 
 class RepositoryReference(datatype('Repository', [])):
@@ -191,7 +192,7 @@ class Manifest(datatype('Manifest', ['digest', 'media_type', 'manifest_bytes']))
 
   def get_parsed_manifest(self, validate=True):
     """ Returns the parsed manifest for this manifest. """
-    return DockerSchema1Manifest(self.manifest_bytes, validate=validate)
+    return parse_manifest_from_bytes(self.manifest_bytes, self.media_type, validate=validate)
 
 
 class LegacyImage(datatype('LegacyImage', ['docker_image_id', 'created', 'comment', 'command',
diff --git a/data/registry_model/interface.py b/data/registry_model/interface.py
index 544462c84..b342314c8 100644
--- a/data/registry_model/interface.py
+++ b/data/registry_model/interface.py
@@ -36,7 +36,8 @@ class RegistryDataInterface(object):
         or None if none. """
 
   @abstractmethod
-  def create_manifest_and_retarget_tag(self, repository_ref, manifest_interface_instance, tag_name):
+  def create_manifest_and_retarget_tag(self, repository_ref, manifest_interface_instance, tag_name,
+                                       storage):
     """ Creates a manifest in a repository, adding all of the necessary data in the model.
 
         The `manifest_interface_instance` parameter must be an instance of the manifest
@@ -127,7 +128,7 @@ class RegistryDataInterface(object):
 
   @abstractmethod
   def retarget_tag(self, repository_ref, tag_name, manifest_or_legacy_image,
-                   is_reversion=False):
+                   storage, is_reversion=False):
     """
     Creates, updates or moves a tag to a new entry in history, pointing to the manifest or
     legacy image specified. If is_reversion is set to True, this operation is considered a
diff --git a/data/registry_model/manifestbuilder.py b/data/registry_model/manifestbuilder.py
index 7ab2e18be..ce8d6d95b 100644
--- a/data/registry_model/manifestbuilder.py
+++ b/data/registry_model/manifestbuilder.py
@@ -18,17 +18,17 @@ _BuilderState = namedtuple('_BuilderState', ['builder_id', 'images', 'tags', 'ch
 _SESSION_KEY = '__manifestbuilder'
 
 
-def create_manifest_builder(repository_ref):
+def create_manifest_builder(repository_ref, storage):
   """ Creates a new manifest builder for populating manifests under the specified repository
       and returns it. Returns None if the builder could not be constructed.
   """
   builder_id = str(uuid.uuid4())
-  builder = _ManifestBuilder(repository_ref, _BuilderState(builder_id, {}, {}, {}))
+  builder = _ManifestBuilder(repository_ref, _BuilderState(builder_id, {}, {}, {}), storage)
   builder._save_to_session()
   return builder
 
 
-def lookup_manifest_builder(repository_ref, builder_id):
+def lookup_manifest_builder(repository_ref, builder_id, storage):
   """ Looks up the manifest builder with the given ID under the specified repository and returns
       it or None if none.
   """
@@ -40,16 +40,17 @@ def lookup_manifest_builder(repository_ref, builder_id):
   if builder_state.builder_id != builder_id:
     return None
 
-  return _ManifestBuilder(repository_ref, builder_state)
+  return _ManifestBuilder(repository_ref, builder_state, storage)
 
 
 class _ManifestBuilder(object):
   """ Helper class which provides an interface for bookkeeping the layers and configuration of
       manifests being constructed.
   """
-  def __init__(self, repository_ref, builder_state):
+  def __init__(self, repository_ref, builder_state, storage):
     self._repository_ref = repository_ref
     self._builder_state = builder_state
+    self._storage = storage
 
   @property
   def builder_id(self):
@@ -183,7 +184,7 @@ class _ManifestBuilder(object):
     if legacy_image is None:
       return None
 
-    tag = registry_model.retarget_tag(self._repository_ref, tag_name, legacy_image)
+    tag = registry_model.retarget_tag(self._repository_ref, tag_name, legacy_image, self._storage)
     if tag is None:
       return None
 
diff --git a/data/registry_model/registry_oci_model.py b/data/registry_model/registry_oci_model.py
index 455cff48f..0f93b1716 100644
--- a/data/registry_model/registry_oci_model.py
+++ b/data/registry_model/registry_oci_model.py
@@ -11,7 +11,6 @@ from data.registry_model.interface import RegistryDataInterface
 from data.registry_model.datatypes import Tag, Manifest, LegacyImage, Label, SecurityScanStatus
 from data.registry_model.shared import SharedModel
 from data.registry_model.label_handlers import apply_label_to_manifest
-from util.validation import is_json
 
 
 logger = logging.getLogger(__name__)
@@ -176,7 +175,8 @@ class OCIModel(SharedModel, RegistryDataInterface):
 
     return Tag.for_tag(tag, legacy_image=LegacyImage.for_image(legacy_image))
 
-  def create_manifest_and_retarget_tag(self, repository_ref, manifest_interface_instance, tag_name):
+  def create_manifest_and_retarget_tag(self, repository_ref, manifest_interface_instance, tag_name,
+                                       storage):
     """ Creates a manifest in a repository, adding all of the necessary data in the model.
 
         The `manifest_interface_instance` parameter must be an instance of the manifest
@@ -187,41 +187,47 @@ class OCIModel(SharedModel, RegistryDataInterface):
 
         Returns a reference to the (created manifest, tag) or (None, None) on error.
     """
+    def _retrieve_repo_blob(digest):
+      blob_found = self.get_repo_blob_by_digest(repository_ref, digest, include_placements=True)
+      if blob_found is None:
+        return None
+
+      try:
+        return storage.get_content(blob_found.placements, blob_found.storage_path)
+      except IOError:
+        logger.exception('Could not retrieve configuration blob `%s`', digest)
+        return None
+
     # Get or create the manifest itself.
-    manifest, newly_created = oci.manifest.get_or_create_manifest(repository_ref._db_id,
-                                                                  manifest_interface_instance)
-    if manifest is None:
+    created_manifest = oci.manifest.get_or_create_manifest(repository_ref._db_id,
+                                                           manifest_interface_instance,
+                                                           storage)
+    if created_manifest is None:
       return (None, None)
 
     # Re-target the tag to it.
-    tag = oci.tag.retarget_tag(tag_name, manifest)
+    tag = oci.tag.retarget_tag(tag_name, created_manifest.manifest)
     if tag is None:
       return (None, None)
 
-    legacy_image = oci.shared.get_legacy_image_for_manifest(manifest)
+    legacy_image = oci.shared.get_legacy_image_for_manifest(created_manifest.manifest)
     if legacy_image is None:
       return (None, None)
 
-    # Save the labels on the manifest. Note that order is important here: This must come after the
-    # tag has been changed.
-    # TODO(jschorr): Support schema2 here when we're ready.
-    if newly_created:
-      has_labels = False
+    li = LegacyImage.for_image(legacy_image)
+    wrapped_manifest = Manifest.for_manifest(created_manifest.manifest, li)
 
-      with self.batch_create_manifest_labels(Manifest.for_manifest(manifest, None)) as add_label:
-        for key, value in manifest_interface_instance.layers[-1].v1_metadata.labels.iteritems():
-          media_type = 'application/json' if is_json(value) else 'text/plain'
-          add_label(key, value, 'manifest', media_type)
-          has_labels = True
+    # Apply any labels that should modify the created tag.
+    if created_manifest.labels_to_apply:
+      for key, value in created_manifest.labels_to_apply.iteritems():
+        apply_label_to_manifest(dict(key=key, value=value), wrapped_manifest, self)
 
       # Reload the tag in case any updates were applied.
-      if has_labels:
-        tag = database.Tag.get(id=tag.id)
+      tag = database.Tag.get(id=tag.id)
 
-    li = LegacyImage.for_image(legacy_image)
-    return (Manifest.for_manifest(manifest, li), Tag.for_tag(tag, li))
+    return (wrapped_manifest, Tag.for_tag(tag, li))
 
-  def retarget_tag(self, repository_ref, tag_name, manifest_or_legacy_image,
+  def retarget_tag(self, repository_ref, tag_name, manifest_or_legacy_image, storage,
                    is_reversion=False):
     """
     Creates, updates or moves a tag to a new entry in history, pointing to the manifest or
@@ -240,11 +246,12 @@ class OCIModel(SharedModel, RegistryDataInterface):
       if manifest_instance is None:
         return None
 
-      manifest, _ = oci.manifest.get_or_create_manifest(repository_ref._db_id, manifest_instance)
-      if manifest is None:
+      created = oci.manifest.get_or_create_manifest(repository_ref._db_id, manifest_instance,
+                                                    storage)
+      if created is None:
         return None
 
-      manifest_id = manifest.id
+      manifest_id = created.manifest.id
 
     tag = oci.tag.retarget_tag(tag_name, manifest_id, is_reversion=is_reversion)
     legacy_image = LegacyImage.for_image(oci.shared.get_legacy_image_for_manifest(manifest_id))
diff --git a/data/registry_model/registry_pre_oci_model.py b/data/registry_model/registry_pre_oci_model.py
index b70045b5f..1b6517ed4 100644
--- a/data/registry_model/registry_pre_oci_model.py
+++ b/data/registry_model/registry_pre_oci_model.py
@@ -79,7 +79,8 @@ class PreOCIModel(SharedModel, RegistryDataInterface):
 
     return Manifest.for_tag_manifest(tag_manifest, legacy_image)
 
-  def create_manifest_and_retarget_tag(self, repository_ref, manifest_interface_instance, tag_name):
+  def create_manifest_and_retarget_tag(self, repository_ref, manifest_interface_instance, tag_name,
+                                       storage):
     """ Creates a manifest in a repository, adding all of the necessary data in the model.
 
         The `manifest_interface_instance` parameter must be an instance of the manifest
@@ -298,7 +299,7 @@ class PreOCIModel(SharedModel, RegistryDataInterface):
     manifest_digest = tag_manifest.digest if tag_manifest else None
     return Tag.for_repository_tag(tag, legacy_image=legacy_image, manifest_digest=manifest_digest)
 
-  def retarget_tag(self, repository_ref, tag_name, manifest_or_legacy_image,
+  def retarget_tag(self, repository_ref, tag_name, manifest_or_legacy_image, storage,
                    is_reversion=False):
     """
     Creates, updates or moves a tag to a new entry in history, pointing to the manifest or
diff --git a/data/registry_model/test/test_interface.py b/data/registry_model/test/test_interface.py
index 0c75d06b7..c0ebe625f 100644
--- a/data/registry_model/test/test_interface.py
+++ b/data/registry_model/test/test_interface.py
@@ -9,7 +9,7 @@ import pytest
 from mock import patch
 from playhouse.test_utils import assert_query_count
 
-from app import docker_v2_signing_key
+from app import docker_v2_signing_key, storage
 from data import model
 from data.database import (TagManifestLabelMap, TagManifestToManifest, Manifest, ManifestBlob,
                            ManifestLegacyImage, ManifestLabel, TagManifest, RepositoryTag, Image,
@@ -306,7 +306,7 @@ def test_retarget_tag_history(use_manifest, registry_model):
   # Retarget the tag.
   assert manifest_or_legacy_image
   updated_tag = registry_model.retarget_tag(repository_ref, 'latest', manifest_or_legacy_image,
-                                            is_reversion=True)
+                                            storage, is_reversion=True)
 
   # Ensure the tag has changed targets.
   if use_manifest:
@@ -698,7 +698,8 @@ def test_create_manifest_and_retarget_tag(registry_model):
 
   another_manifest, tag = registry_model.create_manifest_and_retarget_tag(repository_ref,
                                                                           sample_manifest,
-                                                                          'anothertag')
+                                                                          'anothertag',
+                                                                          storage)
   assert another_manifest is not None
   assert tag is not None
 
@@ -730,7 +731,8 @@ def test_create_manifest_and_retarget_tag_with_labels(registry_model):
 
   another_manifest, tag = registry_model.create_manifest_and_retarget_tag(repository_ref,
                                                                           sample_manifest,
-                                                                          'anothertag')
+                                                                          'anothertag',
+                                                                          storage)
   assert another_manifest is not None
   assert tag is not None
 
diff --git a/data/registry_model/test/test_manifestbuilder.py b/data/registry_model/test/test_manifestbuilder.py
index 381ac9ea3..89838f525 100644
--- a/data/registry_model/test/test_manifestbuilder.py
+++ b/data/registry_model/test/test_manifestbuilder.py
@@ -40,9 +40,9 @@ def test_build_manifest(layers, fake_session, registry_model):
   settings = BlobUploadSettings('2M', 512 * 1024, 3600)
   app_config = {'TESTING': True}
 
-  builder = create_manifest_builder(repository_ref)
-  assert lookup_manifest_builder(repository_ref, 'anotherid') is None
-  assert lookup_manifest_builder(repository_ref, builder.builder_id) is not None
+  builder = create_manifest_builder(repository_ref, storage)
+  assert lookup_manifest_builder(repository_ref, 'anotherid', storage) is None
+  assert lookup_manifest_builder(repository_ref, builder.builder_id, storage) is not None
 
   blobs_by_layer = {}
   for layer_id, parent_id, layer_bytes in layers:
@@ -89,8 +89,9 @@ def test_build_manifest(layers, fake_session, registry_model):
 
 
 def test_build_manifest_missing_parent(fake_session, registry_model):
+  storage = DistributedStorage({'local_us': FakeStorage(None)}, ['local_us'])
   repository_ref = registry_model.lookup_repository('devtable', 'complex')
-  builder = create_manifest_builder(repository_ref)
+  builder = create_manifest_builder(repository_ref, storage)
 
   assert builder.start_layer('somelayer', json.dumps({'id': 'somelayer', 'parent': 'someparent'}),
                              'local_us', None, 60) is None
diff --git a/endpoints/api/tag.py b/endpoints/api/tag.py
index 203aabe6c..413491093 100644
--- a/endpoints/api/tag.py
+++ b/endpoints/api/tag.py
@@ -3,6 +3,7 @@
 from datetime import datetime
 from flask import request, abort
 
+from app import storage
 from auth.auth_context import get_authenticated_user
 from data.registry_model import registry_model
 from endpoints.api import (resource, nickname, require_repo_read, require_repo_write,
@@ -154,7 +155,7 @@ class RepositoryTag(RepositoryParamResource):
       if image is None:
         raise NotFound()
 
-      if not registry_model.retarget_tag(repo_ref, tag, image):
+      if not registry_model.retarget_tag(repo_ref, tag, image, storage):
         raise InvalidRequest('Could not move tag')
 
       username = get_authenticated_user().username
@@ -287,7 +288,8 @@ class RestoreTag(RepositoryParamResource):
     if manifest_or_legacy_image is None:
       raise NotFound()
 
-    if not registry_model.retarget_tag(repo_ref, tag, manifest_or_legacy_image, is_reversion=True):
+    if not registry_model.retarget_tag(repo_ref, tag, manifest_or_legacy_image, storage,
+                                       is_reversion=True):
       raise InvalidRequest('Could not restore tag')
 
     log_action('revert_tag', namespace, log_data, repo_name=repository)
diff --git a/endpoints/v1/index.py b/endpoints/v1/index.py
index 6df84d0d2..0c0d52945 100644
--- a/endpoints/v1/index.py
+++ b/endpoints/v1/index.py
@@ -6,7 +6,7 @@ from functools import wraps
 
 from flask import request, make_response, jsonify, session
 
-from app import userevents, metric_queue
+from app import userevents, metric_queue, storage
 from auth.auth_context import get_authenticated_context, get_authenticated_user
 from auth.credentials import validate_credentials, CredentialKind
 from auth.decorators import process_auth
@@ -217,7 +217,7 @@ def create_repository(namespace_name, repo_name):
 
   # Start a new builder for the repository and save its ID in the session.
   assert repository_ref
-  builder = create_manifest_builder(repository_ref)
+  builder = create_manifest_builder(repository_ref, storage)
   logger.debug('Started repo push with manifest builder %s', builder)
   if builder is None:
     abort(404, message='Unknown repository', issue='unknown-repo')
@@ -243,7 +243,7 @@ def update_images(namespace_name, repo_name):
       # Make sure the repo actually exists.
       abort(404, message='Unknown repository', issue='unknown-repo')
 
-    builder = lookup_manifest_builder(repository_ref, session.get('manifest_builder'))
+    builder = lookup_manifest_builder(repository_ref, session.get('manifest_builder'), storage)
     if builder is None:
       abort(400)
 
diff --git a/endpoints/v1/registry.py b/endpoints/v1/registry.py
index 2cdd5b504..08d99313d 100644
--- a/endpoints/v1/registry.py
+++ b/endpoints/v1/registry.py
@@ -166,7 +166,7 @@ def put_image_layer(namespace, repository, image_id):
     exact_abort(409, 'Image already exists')
 
   logger.debug('Checking for image in manifest builder')
-  builder = lookup_manifest_builder(repository_ref, session.get('manifest_builder'))
+  builder = lookup_manifest_builder(repository_ref, session.get('manifest_builder'), store)
   if builder is None:
     abort(400)
 
@@ -268,7 +268,7 @@ def put_image_checksum(namespace, repository, image_id):
           image_id=image_id)
 
   logger.debug('Checking for image in manifest builder')
-  builder = lookup_manifest_builder(repository_ref, session.get('manifest_builder'))
+  builder = lookup_manifest_builder(repository_ref, session.get('manifest_builder'), store)
   if builder is None:
     abort(400)
 
@@ -361,7 +361,7 @@ def put_image_json(namespace, repository, image_id):
   if repository_ref is None:
     abort(403)
 
-  builder = lookup_manifest_builder(repository_ref, session.get('manifest_builder'))
+  builder = lookup_manifest_builder(repository_ref, session.get('manifest_builder'), store)
   if builder is None:
     abort(400)
 
diff --git a/endpoints/v1/tag.py b/endpoints/v1/tag.py
index 620d0df8e..2b4da596d 100644
--- a/endpoints/v1/tag.py
+++ b/endpoints/v1/tag.py
@@ -3,6 +3,7 @@ import json
 
 from flask import abort, request, jsonify, make_response, session
 
+from app import storage
 from auth.decorators import process_auth
 from auth.permissions import (ReadRepositoryPermission, ModifyRepositoryPermission)
 from data.registry_model import registry_model
@@ -70,7 +71,7 @@ def put_tag(namespace_name, repo_name, tag):
     image_id = json.loads(request.data)
 
     # Check for the image ID first in a builder (for an in-progress push).
-    builder = lookup_manifest_builder(repository_ref, session.get('manifest_builder'))
+    builder = lookup_manifest_builder(repository_ref, session.get('manifest_builder'), storage)
     if builder is not None:
       layer = builder.lookup_layer(image_id)
       if layer is not None:
@@ -86,7 +87,7 @@ def put_tag(namespace_name, repo_name, tag):
     if legacy_image is None:
       abort(400)
 
-    if registry_model.retarget_tag(repository_ref, tag, legacy_image) is None:
+    if registry_model.retarget_tag(repository_ref, tag, legacy_image, storage) is None:
       abort(400)
 
     return make_response('Created', 200)
diff --git a/endpoints/v2/manifest.py b/endpoints/v2/manifest.py
index f1c897868..31cdf7e7e 100644
--- a/endpoints/v2/manifest.py
+++ b/endpoints/v2/manifest.py
@@ -6,7 +6,7 @@ from flask import request, url_for, Response
 
 import features
 
-from app import app, metric_queue
+from app import app, metric_queue, storage
 from auth.registry_jwt_auth import process_registry_jwt_auth
 from digest import digest_tools
 from data.registry_model import registry_model
@@ -227,7 +227,8 @@ def _write_manifest(namespace_name, repo_name, manifest_impl):
     raise NameUnknown()
 
   manifest, tag = registry_model.create_manifest_and_retarget_tag(repository_ref, manifest_impl,
-                                                                  manifest_impl.tag)
+                                                                  manifest_impl.tag,
+                                                                  storage)
   if manifest is None:
     raise ManifestInvalid()
 
diff --git a/initdb.py b/initdb.py
index 4ea4de50f..ed17005db 100644
--- a/initdb.py
+++ b/initdb.py
@@ -29,6 +29,7 @@ from data.registry_model.registry_pre_oci_model import pre_oci_model
 from app import app, storage as store, tf
 from storage.basestorage import StoragePaths
 from image.docker.schema1 import DOCKER_SCHEMA1_CONTENT_TYPES
+from image.docker.schema2 import DOCKER_SCHEMA2_CONTENT_TYPES
 
 
 from workers import repositoryactioncounter
@@ -435,6 +436,9 @@ def initialize_database():
   for media_type in DOCKER_SCHEMA1_CONTENT_TYPES:
     MediaType.create(name=media_type)
 
+  for media_type in DOCKER_SCHEMA2_CONTENT_TYPES:
+    MediaType.create(name=media_type)
+
   LabelSourceType.create(name='manifest')
   LabelSourceType.create(name='api', mutable=True)
   LabelSourceType.create(name='internal')
diff --git a/test/test_api_usage.py b/test/test_api_usage.py
index fab4026bf..808cb3406 100644
--- a/test/test_api_usage.py
+++ b/test/test_api_usage.py
@@ -21,7 +21,8 @@ from cryptography.hazmat.backends import default_backend
 from endpoints.api import api_bp, api
 from endpoints.building import PreparedBuild
 from endpoints.webhooks import webhooks
-from app import app, config_provider, all_queues, dockerfile_build_queue, notification_queue
+from app import (app, config_provider, all_queues, dockerfile_build_queue, notification_queue,
+                 storage)
 from buildtrigger.basehandler import BuildTriggerHandler
 from initdb import setup_database_for_testing, finished_database_for_testing
 from data import database, model, appr_model
@@ -2909,7 +2910,7 @@ class TestListAndDeleteTag(ApiTestCase):
     for i in xrange(1, 9):
       tag_name = "tag" + str(i)
       remaining_tags.add(tag_name)
-      assert registry_model.retarget_tag(repo_ref, tag_name, latest_tag.legacy_image)
+      assert registry_model.retarget_tag(repo_ref, tag_name, latest_tag.legacy_image, storage)
 
     # Make sure we can iterate over all of them.
     json = self.getJsonResponse(ListRepositoryTags, params=dict(