Hash files while uploading them, and accept StringIO to avoid use of temporary files

When uploading big file, this avoids trashing the kernel cache once more.

slapos/libnetworkcache.py

@@ -17,6 +17,7 @@ import base64
 import hashlib
 import httplib
 import json
+import os
 import socket
 import subprocess
 import tempfile
@@ -40,6 +41,35 @@ def urljoin(a, b):
   return a + b
 
 
+class hashing_file(object):
+
+  def __init__(self, file):
+    self._f = file
+    self._h = hashlib.sha512()
+
+  def hexdigest(self):
+    assert not self._f.read(1)
+    return self._h.hexdigest()
+
+  def read(self, *args, **kw):
+    d = self._f.read(*args, **kw)
+    self._h.update(d)
+    return d
+
+  def __len__(self):
+    f = self._f
+    try:
+      fd = f.fileno()
+    except AttributeError:
+      pos = f.tell()
+      try:
+        f.seek(0, 2)
+        return f.tell()
+      finally:
+        f.seek(pos)
+    return os.fstat(fd).st_size
+
+
 class NetworkcacheClient(object):
   '''
     NetworkcacheClient is a wrapper for httplib.
@@ -140,17 +170,9 @@ class NetworkcacheClient(object):
     If urlmd5 is None it must only upload to SHACACHE.
     Otherwise, it must create a new entry on SHADIR.
     '''
-    sha512sum = hashlib.sha512()
     # do not trust, go to beginning of opened file
     file_descriptor.seek(0)
-    while True:
-      d = file_descriptor.read(sha512sum.block_size)
-      if not d:
-        break
-      sha512sum.update(d)
-    sha512sum = sha512sum.hexdigest()
-
-    file_descriptor.seek(0)
+    file_descriptor = hashing_file(file_descriptor)
     if self.shacache_scheme == 'https':
       shacache_connection = httplib.HTTPSConnection(self.shacache_host,
         self.shacache_port, key_file=self.shacache_key_file,
@@ -164,15 +186,15 @@ class NetworkcacheClient(object):
       print 'uploaded'
       result = shacache_connection.getresponse()
       print 'answered'
-      data = result.read()
+      sha512sum = result.read()
       print 'read'
     finally:
       shacache_connection.close()
 
-    if result.status != 201 or data != sha512sum:
-      raise UploadError('Failed to upload the file to SHACACHE Server.' \
-                        'URL: %s. Response code: %s. Response data: %s' % \
-                                   (self.shacache_host, result.status, data))
+    if result.status != 201 or sha512sum != file_descriptor.hexdigest():
+      raise UploadError('Failed to upload the file to SHACACHE Server.'
+                        'URL: %s. Response code: %s. Response data: %s'
+                        % (self.shacache_host, result.status, sha512sum))
 
     if key is not None:
       if file_name is None or urlmd5 is None:
@@ -222,15 +244,8 @@ class NetworkcacheClient(object):
     If key is None, it must only upload to SHACACHE.
     Otherwise, it must create a new entry on SHADIR.
     '''
-    sha512sum = hashlib.sha512()
-    file_descriptor.seek(0)
-    while True:
-      d = file_descriptor.read(sha512sum.block_size)
-      if not d:
-        break
-      sha512sum.update(d)
-    sha512sum = sha512sum.hexdigest()
     file_descriptor.seek(0)
+    file_descriptor = hashing_file(file_descriptor)
     if self.shacache_scheme == 'https':
       shacache_connection = httplib.HTTPSConnection(self.shacache_host,
         self.shacache_port, key_file = self.shacache_key_file,
@@ -242,16 +257,17 @@ class NetworkcacheClient(object):
       shacache_connection.request('POST', self.shacache_path, file_descriptor,
                                       self.shacache_header_dict)
       result = shacache_connection.getresponse()
-      data = result.read()
+      sha512sum = result.read()
     finally:
       shacache_connection.close()
-    if result.status != 201 or data != sha512sum:
-      raise UploadError('Failed to upload the file to SHACACHE Server.' \
-                        'URL: %s. Response code: %s. Response data: %s' % \
-                            (self.shacache_host, result.status, data))
+
+    if result.status != 201 or sha512sum != file_descriptor.hexdigest():
+      raise UploadError('Failed to upload the file to SHACACHE Server.'
+                        'URL: %s. Response code: %s. Response data: %s'
+                        % (self.shacache_host, result.status, sha512sum))
 
     if key is not None:
-      kw['sha512'] = data # always update sha512sum
+      kw['sha512'] = sha512sum # always update sha512sum
       sha_entry = json.dumps(kw)
       try:
         signature = self._getSignatureString(sha_entry)
@@ -283,6 +299,9 @@ class NetworkcacheClient(object):
   def download(self, sha512sum):
     ''' Download the file.
     It uses http GET request method.
+    # do not trust, go to beginning of opened file
+    file_descriptor.seek(0)
+    file_descriptor = hashing_file(file_descriptor)
     '''
     sha_cache_url = urljoin(self.shacache_url, sha512sum)
     request = urllib2.Request(url=sha_cache_url, data=None,

slapos/libnetworkcachetests.py

@@ -18,6 +18,7 @@ import unittest
 import slapos.libnetworkcache
 import slapos.signature
 import sys
+from cStringIO import StringIO
 
 
 class NCHandler(BaseHTTPServer.BaseHTTPRequestHandler):
@@ -150,7 +151,7 @@ class OfflineTest(unittest.TestCase):
   def test_upload_offline(self):
     nc = slapos.libnetworkcache.NetworkcacheClient(self.shacache_url,
       self.shadir_url)
-    self.assertRaises(IOError, nc.upload, tempfile.TemporaryFile())
+    self.assertRaises(IOError, nc.upload, StringIO())
 
   def test_init_method_normal_http_url(self):
     """
@@ -192,9 +193,8 @@ class OnlineMixin:
     if not 'TEST_SHA_CACHE' in os.environ and not 'TEST_SHA_DIR' in os.environ:
       self.tree = tempfile.mkdtemp()
       self.thread = Server.run(self.tree, (self.host, self.port), self.handler)
-    self.test_data = tempfile.TemporaryFile()
     self.test_string = str(random.random())
-    self.test_data.write(self.test_string)
+    self.test_data = StringIO(self.test_string)
     self.test_shasum = hashlib.sha512(self.test_string).hexdigest()
 
   def tearDown(self):