Add piece hashing to verbs generated image storages

endpoints/v1/registry.py

@@ -239,7 +239,7 @@ def put_image_layer(namespace, repository, image_id):
   updated_storage = model.storage.set_image_storage_metadata(image_id, namespace, repository,
                                                              size_info.compressed_size,
                                                              size_info.uncompressed_size)
-  pieces_bytes = piece_hasher.piece_hashes + piece_hasher.hash_fragment.digest()
+  pieces_bytes = piece_hasher.final_piece_hashes()
   model.storage.save_torrent_info(updated_storage, app.config['TORRENT_PIECE_SIZE'], pieces_bytes)
 
   # Append the computed checksum.

endpoints/verbs.py

@@ -15,7 +15,10 @@ from storage import Storage
 
 from util.registry.queuefile import QueueFile
 from util.registry.queueprocess import QueueProcess
-from util.registry.torrent import make_torrent, per_user_torrent_filename, public_torrent_filename
+from util.registry.torrent import (make_torrent, per_user_torrent_filename, public_torrent_filename,
+                                   PieceHasher)
+from util.registry.filelike import wrap_with_handler
+from util.registry.gzipstream import SizeInfo
 from formats.squashed import SquashedDockerImage
 from formats.aci import ACIImage
 from endpoints.v2.blob import BLOB_DIGEST_ROUTE
@@ -25,7 +28,8 @@ verbs = Blueprint('verbs', __name__)
 logger = logging.getLogger(__name__)
 
 
-def _open_stream(formatter, namespace, repository, tag, synthetic_image_id, image_json, repo_image):
+def _open_stream(formatter, namespace, repository, tag, synthetic_image_id, image_json, repo_image,
+                 handlers):
   store = Storage(app)
 
   # For performance reasons, we load the full image list here, cache it, then disconnect from
@@ -55,6 +59,9 @@ def _open_stream(formatter, namespace, repository, tag, synthetic_image_id, imag
   stream = formatter.build_stream(namespace, repository, tag, synthetic_image_id, image_json,
                                   get_next_image, get_next_layer, get_image_json)
 
+  for handler_fn in handlers:
+    stream = wrap_with_handler(stream, handler_fn)
+
   return stream.read
 
 
@@ -200,12 +207,20 @@ def _repo_verb(namespace, repository, tag, verb, formatter, sign=False, checker=
     # Close any existing DB connection once the process has exited.
     database.close_db_filter(None)
 
+  hasher = PieceHasher(app.config['TORRENT_PIECE_SIZE'])
+  def _store_metadata_and_cleanup():
+    with database.UseThenDisconnect(app.config):
+      model.storage.save_torrent_info(derived, app.config['TORRENT_PIECE_SIZE'],
+                                      hasher.final_piece_hashes())
+
   # Create a queue process to generate the data. The queue files will read from the process
   # and send the results to the client and storage.
-  args = (formatter, namespace, repository, tag, synthetic_image_id, image_json, repo_image)
+  handlers = [hasher.update]
+  args = (formatter, namespace, repository, tag, synthetic_image_id, image_json, repo_image,
+          handlers)
   queue_process = QueueProcess(_open_stream,
                                8 * 1024, 10 * 1024 * 1024, # 8K/10M chunk/max
-                               args, finished=_cleanup)
+                               args, finished=_store_metadata_and_cleanup)
 
   client_queue_file = QueueFile(queue_process.create_queue(), 'client')
   storage_queue_file = QueueFile(queue_process.create_queue(), 'storage')

util/registry/torrent.py

@@ -80,3 +80,6 @@ class PieceHasher(object):
   @property
   def hash_fragment(self):
     return self._hash_fragment
+
+  def final_piece_hashes(self):
+    return self._piece_hashes + self._hash_fragment.digest()