From a627494e056e9dbab17cf888914778e819f2078a Mon Sep 17 00:00:00 2001
From: Joseph Schorr <josephschorr@users.noreply.github.com>
Date: Thu, 17 Dec 2015 13:39:01 -0500
Subject: [PATCH] Ensure the squashed estimated size is an int

Also adds a test to verify
---
 formats/squashed.py    |  4 ++++
 test/registry_tests.py | 39 +++++++++++++++++++++++++++++++++++++--
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/formats/squashed.py b/formats/squashed.py
index 13b7dabcb..2dc5675f2 100644
--- a/formats/squashed.py
+++ b/formats/squashed.py
@@ -5,6 +5,7 @@ from formats.tarimageformatter import TarImageFormatter
 
 import copy
 import json
+import math
 
 class FileEstimationException(Exception):
   """ Exception raised by build_docker_load_stream if the estimated size of the layer TAR
@@ -68,6 +69,9 @@ class SquashedDockerImage(TarImageFormatter):
         image_json = get_image_json(image)
         estimated_file_size += image_json.get('Size', 0) * SquashedDockerImage.SIZE_MULTIPLIER
 
+    # Make sure the estimated file size is an integer number of bytes.
+    estimated_file_size = int(math.ceil(estimated_file_size))
+
     yield self.tar_file_header(synthetic_image_id + '/layer.tar', estimated_file_size)
 
     # Yield the contents of the merged layer.
diff --git a/test/registry_tests.py b/test/registry_tests.py
index fcda370aa..cf1414c35 100644
--- a/test/registry_tests.py
+++ b/test/registry_tests.py
@@ -60,6 +60,7 @@ logger = logging.getLogger(__name__)
 def generate_csrf():
   return generate_csrf_token()
 
+
 @testbp.route('/feature/<feature_name>', methods=['POST'])
 def set_feature(feature_name):
   import features
@@ -67,6 +68,15 @@ def set_feature(feature_name):
   features._FEATURES[feature_name].value = request.get_json()['value']
   return jsonify({'old_value': old_value})
 
+
+@testbp.route('/removeuncompressed/<image_id>', methods=['POST'])
+def removeuncompressed(image_id):
+  image = model.image.get_image_by_id('devtable', 'newrepo', image_id)
+  image.storage.uncompressed_size = None
+  image.storage.save()
+  return 'OK'
+
+
 @testbp.route('/addtoken', methods=['POST'])
 def addtoken():
   another_token = model.token.create_delegate_token('devtable', 'newrepo', 'my-new-token', 'write')
@@ -263,8 +273,12 @@ class V1RegistryPushMixin(V1RegistryMixin):
       last_image_id = image_id
 
       # PUT /v1/images/{imageID}/json
+      image_json_data = {'id': image_id}
+      if 'size' in image_data:
+        image_json_data['Size'] = image_data['size']
+
       self.conduct('PUT', '/v1/images/%s/json' % image_id,
-                   data=json.dumps({'id': image_id}), auth='sig')
+                   data=json.dumps(image_json_data), auth='sig')
 
       # PUT /v1/images/{imageID}/layer
       tar_file_info = tarfile.TarInfo(name='image_name')
@@ -284,7 +298,7 @@ class V1RegistryPushMixin(V1RegistryMixin):
                    data=StringIO(layer_bytes), auth='sig')
 
       # PUT /v1/images/{imageID}/checksum
-      checksum = compute_simple(StringIO(layer_bytes), json.dumps({'id': image_id}))
+      checksum = compute_simple(StringIO(layer_bytes), json.dumps(image_json_data))
       self.conduct('PUT', '/v1/images/%s/checksum' % image_id,
                    headers={'X-Docker-Checksum-Payload': checksum},
                    auth='sig')
@@ -1112,6 +1126,27 @@ class VerbTests(RegistryTestCaseMixin, V1RegistryPushMixin, LiveServerTestCase):
     self.assertTrue(updated_image_id in tar.getnames())
 
 
+  def test_estimated_squashing(self):
+    initial_images = [
+      {
+        'id': 'initialid',
+        'contents': 'the initial image',
+        'size': 2002,
+      },
+    ]
+
+    # Create the repo.
+    self.do_push('devtable', 'newrepo', 'devtable', 'password', images=initial_images)
+
+    # NULL out the uncompressed size to force estimation.
+    self.conduct('POST', '/__test/removeuncompressed/initialid')
+
+    # Pull the squashed version of the tag.
+    initial_image_id = '91081df45b58dc62dd207441785eef2b895f0383fbe601c99a3cf643c79957dc'
+    tar = self.get_squashed_image()
+    self.assertTrue(initial_image_id in tar.getnames())
+
+
   def test_multilayer_squashing(self):
     images = [
       {