Hash and track layer file chunks for torrenting

data/database.py

@@ -9,7 +9,7 @@ from random import SystemRandom
 from datetime import datetime
 from peewee import *
 from data.read_slave import ReadSlaveModel
-from data.fields import ResumableSHAField, JSONField
+from data.fields import ResumableSHA256Field, ResumableSHA1Field, JSONField
 from sqlalchemy.engine.url import make_url
 from collections import defaultdict
 
@@ -803,12 +803,14 @@ class BlobUpload(BaseModel):
   repository = ForeignKeyField(Repository, index=True)
   uuid = CharField(index=True, unique=True)
   byte_count = IntegerField(default=0)
-  sha_state = ResumableSHAField(null=True, default=resumablehashlib.sha256)
+  sha_state = ResumableSHA256Field(null=True, default=resumablehashlib.sha256)
   location = ForeignKeyField(ImageStorageLocation)
   storage_metadata = JSONField(null=True, default={})
   chunk_count = IntegerField(default=0)
   uncompressed_byte_count = IntegerField(null=True)
   created = DateTimeField(default=datetime.now, index=True)
+  piece_sha_state = ResumableSHA1Field(null=True, default=resumablehashlib.sha1)
+  piece_hashes = TextField(default='')
 
   class Meta:
     database = db
@@ -846,6 +848,16 @@ class QuayRelease(BaseModel):
     )
 
 
+class TorrentInfo(BaseModel):
+  storage = ForeignKeyField(ImageStorage)
+  piece_length = IntegerField()
+  pieces = TextField()
+
+  indexes = (
+    (('storage', 'piece_length'), True),
+  )
+
+
 all_models = [User, Repository, Image, AccessToken, Role, RepositoryPermission, Visibility,
               RepositoryTag, EmailConfirmation, FederatedLogin, LoginService, QueueItem,
               RepositoryBuild, Team, TeamMember, TeamRole, LogEntryKind, LogEntry,
@@ -856,4 +868,4 @@ all_models = [User, Repository, Image, AccessToken, Role, RepositoryPermission,
               RepositoryAuthorizedEmail, ImageStorageTransformation,
               TeamMemberInvite, ImageStorageSignature, ImageStorageSignatureKind,
               AccessTokenKind, Star, RepositoryActionCount, TagManifest, UserRegion,
-              QuayService, QuayRegion, QuayRelease, BlobUpload, DerivedStorageForImage]
+              QuayService, QuayRegion, QuayRelease, BlobUpload, DerivedStorageForImage, TorrentInfo]

data/fields.py

@@ -5,7 +5,7 @@ import json
 from peewee import TextField
 
 
-class ResumableSHAField(TextField):
+class ResumableSHA256Field(TextField):
   def db_value(self, value):
     sha_state = value.state()
 
@@ -28,6 +28,29 @@ class ResumableSHAField(TextField):
     return to_resume
 
 
+class ResumableSHA1Field(TextField):
+  def db_value(self, value):
+    sha_state = value.state()
+
+    # One of the fields is a byte string, let's base64 encode it to make sure
+    # we can store and fetch it regardless of default collocation.
+    sha_state[3] = base64.b64encode(sha_state[3])
+
+    return json.dumps(sha_state)
+
+  def python_value(self, value):
+    to_resume = resumablehashlib.sha1()
+    if value is None:
+      return to_resume
+
+    sha_state = json.loads(value)
+
+    # We need to base64 decode the data bytestring.
+    sha_state[3] = base64.b64decode(sha_state[3])
+    to_resume.set_state(sha_state)
+    return to_resume
+
+
 class JSONField(TextField):
   def db_value(self, value):
     return json.dumps(value)

data/model/blob.py

@@ -51,6 +51,8 @@ def store_blob_record_and_temp_link(namespace, repo_name, blob_digest, location_
     image = Image.create(storage=storage, docker_image_id=random_image_name, repository=repo)
     tag.create_temporary_hidden_tag(repo, image, link_expiration_s)
 
+  return storage
+
 
 def get_blob_upload(namespace, repo_name, upload_uuid):
   """ Load the upload which is already in progress.

data/model/storage.py

@@ -5,7 +5,7 @@ from peewee import JOIN_LEFT_OUTER, fn, SQL
 from data.model import config, db_transaction, InvalidImageException
 from data.database import (ImageStorage, Image, DerivedStorageForImage, ImageStoragePlacement,
                            ImageStorageLocation, ImageStorageTransformation, ImageStorageSignature,
-                           ImageStorageSignatureKind, Repository, Namespace)
+                           ImageStorageSignatureKind, Repository, Namespace, TorrentInfo)
 
 
 logger = logging.getLogger(__name__)
@@ -212,3 +212,7 @@ def get_storage_locations(uuid):
            .where(ImageStorage.uuid == uuid))
 
   return [location.location.name for location in query]
+
+
+def save_torrent_info(storage_object, piece_length, pieces):
+  TorrentInfo.create(storage=storage_object, piece_length=piece_length, pieces=pieces)