Update the pieces to use base64 encoded binary

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 ResumableSHA256Field, ResumableSHA1Field, JSONField
+from data.fields import ResumableSHA256Field, ResumableSHA1Field, JSONField, Base64BinaryField
 from sqlalchemy.engine.url import make_url
 from collections import defaultdict
 
@@ -810,7 +810,7 @@ class BlobUpload(BaseModel):
   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='')
+  piece_hashes = Base64BinaryField(default='')
 
   class Meta:
     database = db
@@ -851,7 +851,7 @@ class QuayRelease(BaseModel):
 class TorrentInfo(BaseModel):
   storage = ForeignKeyField(ImageStorage)
   piece_length = IntegerField()
-  pieces = TextField()
+  pieces = Base64BinaryField()
 
   indexes = (
     (('storage', 'piece_length'), True),

data/fields.py

@@ -5,7 +5,10 @@ import json
 from peewee import TextField
 
 
-class ResumableSHA256Field(TextField):
+class _ResumableSHAField(TextField):
+  def _create_sha(self):
+    raise NotImplementedError
+
   def db_value(self, value):
     sha_state = value.state()
 
@@ -16,7 +19,7 @@ class ResumableSHA256Field(TextField):
     return json.dumps(sha_state)
 
   def python_value(self, value):
-    to_resume = resumablehashlib.sha256()
+    to_resume = self._create_sha()
     if value is None:
       return to_resume
 
@@ -28,27 +31,14 @@ class ResumableSHA256Field(TextField):
     return to_resume
 
 
-class ResumableSHA1Field(TextField):
-  def db_value(self, value):
-    sha_state = value.state()
+class ResumableSHA256Field(_ResumableSHAField):
+  def _create_sha(self):
+    return resumablehashlib.sha256()
 
-    # 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 ResumableSHA1Field(_ResumableSHAField):
+  def _create_sha(self):
+    return resumablehashlib.sha1()
 
 
 class JSONField(TextField):
@@ -59,3 +49,15 @@ class JSONField(TextField):
     if value is None or value == "":
       return {}
     return json.loads(value)
+
+
+class Base64BinaryField(TextField):
+  def db_value(self, value):
+    if value is None:
+      return None
+    return base64.b64encode(value)
+
+  def python_value(self, value):
+    if value is None:
+      return None
+    return base64.b64decode(value)

endpoints/v2/blob.py

@@ -294,8 +294,8 @@ def _finish_upload(namespace, repo_name, upload_obj, expected_digest):
                                                             upload_obj.uncompressed_byte_count)
 
   if upload_obj.piece_sha_state is not None:
-    piece_string = upload_obj.piece_hashes + upload_obj.piece_sha_state.hexdigest()
-    model.storage.save_torrent_info(blob_storage, app.config['TORRENT_PIECE_SIZE'], piece_string)
+    piece_bytes = upload_obj.piece_hashes + upload_obj.piece_sha_state.digest()
+    model.storage.save_torrent_info(blob_storage, app.config['TORRENT_PIECE_SIZE'], piece_bytes)
 
   # Delete the upload tracking row.
   upload_obj.delete_instance()

util/registry/torrent.py

@@ -2,7 +2,7 @@ import resumablehashlib
 
 
 class PieceHasher(object):
-  def __init__(self, piece_size, starting_offset, starting_piece_hash_str, hash_fragment_to_resume):
+  def __init__(self, piece_size, starting_offset, starting_piece_hash_bytes, hash_fragment_to_resume):
     if not isinstance(starting_offset, (int, long)):
       raise TypeError('starting_offset must be an integer')
     elif not isinstance(piece_size, (int, long)):
@@ -11,7 +11,7 @@ class PieceHasher(object):
     self._current_offset = starting_offset
     self._piece_size = piece_size
     self._hash_fragment = hash_fragment_to_resume
-    self._piece_hashes = [starting_piece_hash_str]
+    self._piece_hashes = bytearray(starting_piece_hash_bytes)
 
   def update(self, buf):
     buf_offset = 0
@@ -21,7 +21,7 @@ class PieceHasher(object):
 
       if self._piece_offset() == 0 and to_hash_len > 0 and self._current_offset > 0:
         # We are opening a new piece
-        self._piece_hashes.append(self._hash_fragment.hexdigest())
+        self._piece_hashes.extend(self._hash_fragment.digest())
         self._hash_fragment = resumablehashlib.sha1()
 
       self._hash_fragment.update(buf_bytes_to_hash)
@@ -36,7 +36,7 @@ class PieceHasher(object):
 
   @property
   def piece_hashes(self):
-    return ''.join(self._piece_hashes)
+    return self._piece_hashes
 
   @property
   def hash_fragment(self):