Commit 1f336e7d by Guba Sándor

Merge branch 'max-files-size' into 'master'

Max files size

fixes #3

See merge request !7
parents 2d0631fb 205e1a73
...@@ -19,11 +19,17 @@ re_qemu_img = re.compile(r'(file format: (?P<format>(qcow2|raw))|' ...@@ -19,11 +19,17 @@ re_qemu_img = re.compile(r'(file format: (?P<format>(qcow2|raw))|'
r'virtual size: \w+ \((?P<size>[0-9]+) bytes\)|' r'virtual size: \w+ \((?P<size>[0-9]+) bytes\)|'
r'backing file: \S+ \(actual path: (?P<base>\S+)\))$') r'backing file: \S+ \(actual path: (?P<base>\S+)\))$')
maximum_size = float(os.getenv("DOWNLOAD_MAX_SIZE", 1024*1024*1024*10))
class AbortException(Exception): class AbortException(Exception):
pass pass
class FileTooBig(Exception):
pass
class Disk(object): class Disk(object):
''' Storage driver DISK object. ''' Storage driver DISK object.
...@@ -204,7 +210,9 @@ class Disk(object): ...@@ -204,7 +210,9 @@ class Disk(object):
# undocumented zlib feature http://stackoverflow.com/a/2424549 # undocumented zlib feature http://stackoverflow.com/a/2424549
elif ext == 'bz2': elif ext == 'bz2':
decompressor = BZ2Decompressor() decompressor = BZ2Decompressor()
clen = max(int(r.headers.get('content-length', 700000000)), 1) clen = int(r.headers.get('content-length', maximum_size))
if clen > maximum_size:
raise FileTooBig()
percent = 0 percent = 0
try: try:
with open(disk_path, 'wb') as f: with open(disk_path, 'wb') as f:
...@@ -213,6 +221,8 @@ class Disk(object): ...@@ -213,6 +221,8 @@ class Disk(object):
chunk = decompressor.decompress(chunk) chunk = decompressor.decompress(chunk)
f.write(chunk) f.write(chunk)
actsize = f.tell() actsize = f.tell()
if actsize > maximum_size:
raise FileTooBig()
new_percent = min(100, round(actsize * 100.0 / clen)) new_percent = min(100, round(actsize * 100.0 / clen))
if new_percent > percent: if new_percent > percent:
percent = new_percent percent = new_percent
...@@ -234,6 +244,10 @@ class Disk(object): ...@@ -234,6 +244,10 @@ class Disk(object):
os.unlink(disk_path) os.unlink(disk_path)
logger.info("Download %s aborted %s removed.", logger.info("Download %s aborted %s removed.",
url, disk_path) url, disk_path)
except FileTooBig:
os.unlink(disk_path)
raise Exception("%s file is too big. Maximum size "
"is %s" % url, maximum_size)
except: except:
os.unlink(disk_path) os.unlink(disk_path)
logger.error("Download %s failed, %s removed.", logger.error("Download %s failed, %s removed.",
...@@ -309,7 +323,7 @@ class Disk(object): ...@@ -309,7 +323,7 @@ class Disk(object):
percent = 0 percent = 0
diff_disk = Disk.get(self.dir, self.name) diff_disk = Disk.get(self.dir, self.name)
base_disk = Disk.get(self.dir, self.base_name) base_disk = Disk.get(self.dir, self.base_name)
clen = max(base_disk.actual_size + diff_disk.actual_size, clen = min(base_disk.actual_size + diff_disk.actual_size,
diff_disk.size) diff_disk.size)
output = new_disk.get_path() output = new_disk.get_path()
proc = subprocess.Popen(cmdline) proc = subprocess.Popen(cmdline)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment