From dd7f254f96d9aba974039f5871ed6a4e4b2d2db4 Mon Sep 17 00:00:00 2001
From: Joseph Schorr <josephschorr@users.noreply.github.com>
Date: Mon, 27 Feb 2017 13:32:09 -0500
Subject: [PATCH] Have blob uploads be checked against configurable max layer
 size

---
 endpoints/v2/blob.py   | 25 ++++++++++++++++++++++++-
 endpoints/v2/errors.py |  6 ++++++
 requirements-nover.txt |  1 +
 requirements.txt       |  1 +
 4 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/endpoints/v2/blob.py b/endpoints/v2/blob.py
index ba0acf9ad..f67e07fbf 100644
--- a/endpoints/v2/blob.py
+++ b/endpoints/v2/blob.py
@@ -4,6 +4,7 @@ import time
 
 from flask import url_for, request, redirect, Response, abort as flask_abort
 
+import bitmath
 import resumablehashlib
 
 from app import storage, app, get_app_url, metric_queue
@@ -14,7 +15,7 @@ from digest import digest_tools
 from endpoints.common import parse_repository_name
 from endpoints.v2 import v2_bp, require_repo_read, require_repo_write, get_input_stream
 from endpoints.v2.errors import (BlobUnknown, BlobUploadInvalid, BlobUploadUnknown, Unsupported,
-                                 NameUnknown)
+                                 NameUnknown, LayerTooLarge)
 from endpoints.decorators import anon_protect
 from util.cache import cache_control
 from util.registry.filelike import wrap_with_handler, StreamSlice
@@ -346,6 +347,8 @@ def _upload_chunk(blob_upload, range_header):
 
   Returns a BlobUpload object or None if there was a failure.
   """
+  max_layer_size = bitmath.parse_string_unsafe(app.config['MAXIMUM_LAYER_SIZE'])
+
   # Get the offset and length of the current chunk.
   start_offset, length = _start_offset_and_length(range_header)
   if blob_upload is None or None in {start_offset, length}:
@@ -356,6 +359,16 @@ def _upload_chunk(blob_upload, range_header):
     logger.error('start_offset provided to _upload_chunk greater than blob.upload.byte_count')
     return None
 
+  # Check if we should raise 413 before accepting the data.
+  uploaded = bitmath.Byte(length + start_offset)
+  if length > -1 and uploaded > max_layer_size:
+    detail = {
+      'reason': '%s is greater than maximum allowed size %s' % (uploaded, max_layer_size),
+      'max_allowed': max_layer_size.bytes,
+      'uploaded': uploaded.bytes,
+    }
+    raise LayerTooLarge(detail=detail)
+
   location_set = {blob_upload.location_name}
 
   upload_error = None
@@ -435,6 +448,16 @@ def _upload_chunk(blob_upload, range_header):
   blob_upload.byte_count += length_written
   blob_upload.chunk_count += 1
 
+  # Ensure we have not gone beyond the max layer size.
+  upload_size = bitmath.Byte(blob_upload.byte_count)
+  if upload_size > max_layer_size:
+    detail = {
+      'reason': '%s is greater than maximum allowed size %s' % (upload_size, max_layer_size),
+      'max_allowed': max_layer_size.bytes,
+      'uploaded': upload_size.bytes,
+    }
+    raise LayerTooLarge(detail=detail)
+
   return blob_upload
 
 
diff --git a/endpoints/v2/errors.py b/endpoints/v2/errors.py
index 0f8a5284e..9157ed056 100644
--- a/endpoints/v2/errors.py
+++ b/endpoints/v2/errors.py
@@ -112,6 +112,12 @@ class TagInvalid(V2RegistryException):
                                      'manifest tag did not match URI',
                                      detail)
 
+class LayerTooLarge(V2RegistryException):
+  def __init__(self, detail=None):
+    super(LayerTooLarge, self).__init__('BLOB_UPLOAD_INVALID',
+                                        'Uploaded layer is larger than allowed by this registry',
+                                        detail)
+
 
 class Unauthorized(V2RegistryException):
   def __init__(self, detail=None, repository=None, scopes=None):
diff --git a/requirements-nover.txt b/requirements-nover.txt
index 7fdef7c77..f03fe050d 100644
--- a/requirements-nover.txt
+++ b/requirements-nover.txt
@@ -20,6 +20,7 @@ autobahn==0.9.3-3
 beautifulsoup4
 bencode
 bintrees
+bitmath
 boto
 cachetools==1.1.6
 cryptography
diff --git a/requirements.txt b/requirements.txt
index eaf37efea..aed0f9659 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -8,6 +8,7 @@ Babel==2.3.4
 beautifulsoup4==4.5.1
 bencode==1.0
 bintrees==2.0.4
+bitmath==1.3.1.2
 blinker==1.4
 boto==2.43.0
 cachetools==1.1.6