Repository: DataBiosphere/toil

Issue:
Toil AWS JobStore does not protect workflow from truncated downloads
I just ran a workflow that failed when one of my jobs encountered a partially-downloaded file.

I imported this from S3:

```
2018-09-03 13:59:35 9786900533 platinum-WG38_NA12878_haplo_thread_1.xg
```

And I got this in the job store bucket:

```
2018-09-20 17:21:21 9786900533 f506b22a-c4b8-4376-835d-85e822b8e7e3
```

But after downloading it with `job.fileStore.readGlobalFile` the file was truncated. This caused the `vg` command my job was running on the file to crash, and when my code dumped out the file it as it saw it, I got this:

```
2018-09-20 19:14:44 9786896570 index.xg
```

This file is 3963 bytes shorter than, and a prefix of, the file I tried to download from the file store:

```
cmp index.xg platinum-WG38_NA12878_haplo_thread_1.xg
cmp: EOF on index.xg
```

I believe this is related to [this sort of behavior from boto](https://forums.aws.amazon.com/thread.jspa?messageID=701669); downloads from S3 can appear to succeed, with a 200 HTTP code, but with only part of the file contents.

The Toil code to download the file just runs Boto until it gets no error:
https://github.com/DataBiosphere/toil/blob/5c7b39501338fc0d01a600846d04f6a98b259cce/src/toil/jobStores/aws/jobStore.py#L1148-L1152

It needs to check Boto's work to verify that the entire contents of the file were obtained, checking the bytes downloaded against the expected file length. If the file data was truncated, Toil needs to transparently resume the download where Boto left off, to hide this partial download behavior from application code.

Toil *especially* needs to make sure that partially-downloaded files don't make it into the cache.

┆Issue is synchronized with this [Jira Story](https://ucsc-cgl.atlassian.net/browse/TOIL-279)
┆Issue Number: TOIL-279



Discussion:
Here's the log from my workflow where the truncation happened. It looks like S3 might have been having some cloud weather since a bunch of the initial imports had to be retried, but the file in question seems to have made it into the file store OK and then been truncated coming out.

[log.wg38.2018-09-20.txt](https://github.com/DataBiosphere/toil/files/2406551/log.wg38.2018-09-20.txt)

@arostamianfar added checksumming support in https://github.com/DataBiosphere/toil/pull/2889 that ought to have fixed this, but Jordan just reported to me that a toil-vg run failed, apparently with a truncated XG. I'm not sure if Jordan used an older Toil, or if toil-vg is reading in an unprotected way (streaming) or what.

Resolving patch:
diff --git a/src/toil/jobStores/aws/jobStore.py b/src/toil/jobStores/aws/jobStore.py
--- a/src/toil/jobStores/aws/jobStore.py
+++ b/src/toil/jobStores/aws/jobStore.py
@@ -33,11 +33,10 @@
 import base64
 import hashlib
 import itertools
+import reprlib
 import urllib.parse
 import urllib.request, urllib.parse, urllib.error
 from io import BytesIO
-# Python 3 compatibility imports
-from six.moves import StringIO, reprlib
 
 from toil.lib.memoize import strict_bool
 from toil.lib.exceptions import panic
@@ -50,7 +49,7 @@
 import botocore.session
 import botocore.credentials
 
-from toil.lib.compatibility import compat_bytes, compat_plain, USING_PYTHON2
+from toil.lib.compatibility import compat_bytes, compat_plain
 from toil.lib.misc import AtomicFileCreate
 from toil.fileStores import FileID
 from toil.jobStores.abstractJobStore import (AbstractJobStore,
@@ -69,7 +68,7 @@
                                       bucket_location_to_region,
                                       region_to_bucket_location, copyKeyMultipart,
                                       uploadFromPath, chunkedFileUpload, fileSizeAndTime)
-from toil.jobStores.utils import WritablePipe, ReadablePipe
+from toil.jobStores.utils import WritablePipe, ReadablePipe, ReadableTransformingPipe
 from toil.jobGraph import JobGraph
 import toil.lib.encryption as encryption
 
@@ -82,6 +81,11 @@
 s3_boto3_client = boto3_session.client('s3')
 log = logging.getLogger(__name__)
 
+class ChecksumError(Exception):
+    """
+    Raised when a download from AWS does not contain the correct data.
+    """
+    pass
 
 class AWSJobStore(AbstractJobStore):
     """
@@ -420,7 +424,7 @@ def _importFile(self, otherCls, url, sharedFileName=None, hardlink=False):
                 info.copyFrom(srcKey)
                 info.save()
             finally:
-                srcKey.bucket.connection.close()
+                srcKey.bucket.connection.close() 
             return FileID(info.fileID, size) if sharedFileName is None else None
         else:
             return super(AWSJobStore, self)._importFile(otherCls, url,
@@ -866,7 +870,8 @@ def __init__(self, fileID, ownerID, encrypted,
             :param numContentChunks: the number of SDB domain attributes occupied by this files
 
             :type checksum: str|None
-            :param checksum: the sha1 checksum of the file, if available.
+            :param checksum: the checksum of the file, if available. Formatted
+            as <algorithm>$<lowercase hex hash>.
 
             inlined content. Note that an inlined empty string still occupies one chunk.
             """
@@ -1078,6 +1083,8 @@ def upload(self, localFilePath, calculateChecksum=True):
             if file_size <= self.maxInlinedSize():
                 with open(localFilePath, 'rb') as f:
                     self.content = f.read()
+                # Clear out any old checksum in case of overwrite
+                self.checksum = ''
             else:
                 headers = self._s3EncryptionHeaders()
                 self.checksum = self._get_file_checksum(localFilePath) if calculateChecksum else None
@@ -1085,14 +1092,69 @@ def upload(self, localFilePath, calculateChecksum=True):
                                               bucket=self.outer.filesBucket, fileID=compat_bytes(self.fileID),
                                               headers=headers)
 
-        def _get_file_checksum(self, localFilePath):
+        def _start_checksum(self, to_match=None, algorithm='sha1'):
+            """
+            Get a hasher that can be used with _update_checksum and
+            _finish_checksum.
+            
+            If to_match is set, it is a precomputed checksum which we expect
+            the result to match.
+            
+            The right way to compare checksums is to feed in the checksum to be
+            matched, so we can see its algorithm, instead of getting a new one
+            and comparing. If a checksum to match is fed in, _finish_checksum()
+            will raise a ChecksumError if it isn't matched.
+            """
+            
+            # If we have an expexted result it will go here.
+            expected = None
+            
+            if to_match is not None:
+                parts = to_match.split('$')
+                algorithm = parts[0]
+                expected = parts[1]
+            
+            wrapped = getattr(hashlib, algorithm)()
+            
+            log.debug('Starting %s checksum to match %s', algorithm, expected)
+            
+            return (algorithm, wrapped, expected)
+            
+        def _update_checksum(self, checksum_in_progress, data):
+            """
+            Update a checksum in progress from _start_checksum with new data.
+            """
+            log.debug('Updating checksum with %d bytes', len(data))
+            checksum_in_progress[1].update(data)
+        
+        def _finish_checksum(self, checksum_in_progress):
+            """
+            Complete a checksum in progress from _start_checksum and return the
+            checksum result string.
+            """
+            
+            result_hash = checksum_in_progress[1].hexdigest()
+            
+            log.debug('Completed checksum with hash %s vs. expected %s', result_hash, checksum_in_progress[2])
+            
+            if checksum_in_progress[2] is not None:
+                # We expected a particular hash
+                if result_hash != checksum_in_progress[2]:
+                    raise ChecksumError('Checksum mismatch. Expected: %s Actual: %s' %
+                        (checksum_in_progress[2], result_hash))
+                    
+            return '$'.join([checksum_in_progress[0], result_hash])
+            
+            
+        
+        def _get_file_checksum(self, localFilePath, to_match=None):
             with open(localFilePath, 'rb') as f:
-                checksum = hashlib.sha1()
+                hasher = self._start_checksum(to_match=to_match)
                 contents = f.read(1024 * 1024)
                 while contents != b'':
-                    checksum.update(contents)
+                    self._update_checksum(hasher, contents)
                     contents = f.read(1024 * 1024)
-                return 'sha1$%s' % checksum.hexdigest()
+                return self._finish_checksum(hasher)
 
         @contextmanager
         def uploadStream(self, multipart=True, allowInlining=True):
@@ -1111,12 +1173,20 @@ def uploadStream(self, multipart=True, allowInlining=True):
 
             class MultiPartPipe(WritablePipe):
                 def readFrom(self, readable):
+                    # Get the first block of data we want to put
                     buf = readable.read(store.partSize)
                     assert isinstance(buf, bytes)
+                    
                     if allowInlining and len(buf) <= info.maxInlinedSize():
                         log.debug('Inlining content of %d bytes', len(buf))
                         info.content = buf
+                        # There will be no checksum
+                        info.checksum = ''
                     else:
+                        # We will compute a checksum
+                        hasher = info._start_checksum()
+                        info._update_checksum(hasher, buf)
+                    
                         headers = info._s3EncryptionHeaders()
                         for attempt in retry_s3():
                             with attempt:
@@ -1126,27 +1196,21 @@ def readFrom(self, readable):
                                     headers=headers)
                         try:
                             for part_num in itertools.count():
-                                # There must be at least one part, even if the file is empty.
-                                if len(buf) == 0 and part_num > 0:
-                                    break
                                 for attempt in retry_s3():
                                     with attempt:
                                         log.debug('Uploading part %d of %d bytes', part_num + 1, len(buf))
-                                        if USING_PYTHON2:
-                                            upload.upload_part_from_file(fp=StringIO(buf),
-                                                                         # part numbers are 1-based
-                                                                         part_num=part_num + 1,
-                                                                         headers=headers)
-                                        else:
-                                            upload.upload_part_from_file(fp=BytesIO(buf),
-                                                                         # part numbers are 1-based
-                                                                         part_num=part_num + 1,
-                                                                         headers=headers)
-
-                                if len(buf) == 0:
-                                    break
+                                        upload.upload_part_from_file(fp=BytesIO(buf),
+                                                                     # part numbers are 1-based
+                                                                     part_num=part_num + 1,
+                                                                     headers=headers)
+                                
+                                # Get the next block of data we want to put
                                 buf = readable.read(info.outer.partSize)
                                 assert isinstance(buf, bytes)
+                                if len(buf) == 0:
+                                    # Don't allow any part other than the very first to be empty.
+                                    break
+                                info._update_checksum(hasher, buf)
                         except:
                             with panic(log=log):
                                 for attempt in retry_s3():
@@ -1157,6 +1221,9 @@ def readFrom(self, readable):
                             while store._getBucketVersioning(store.filesBucket) != True:
                                 log.warning('Versioning does not appear to be enabled yet. Deferring multipart upload completion...')
                                 time.sleep(1)
+                                
+                            # Save the checksum
+                            info.checksum = info._finish_checksum(hasher)
 
                             for attempt in retry_s3():
                                 with attempt:
@@ -1186,7 +1253,14 @@ def readFrom(self, readable):
                     if allowInlining and dataLength <= info.maxInlinedSize():
                         log.debug('Inlining content of %d bytes', len(buf))
                         info.content = buf
+                        # There will be no checksum
+                        info.checksum = ''
                     else:
+                        # We will compute a checksum
+                        hasher = info._start_checksum()
+                        info._update_checksum(hasher, buf)
+                        info.checksum = info._finish_checksum(hasher)
+                        
                         key = store.filesBucket.new_key(key_name=compat_bytes(info.fileID))
                         buf = BytesIO(buf)
                         headers = info._s3EncryptionHeaders()
@@ -1194,6 +1268,8 @@ def readFrom(self, readable):
                         while store._getBucketVersioning(store.filesBucket) != True:
                             log.warning('Versioning does not appear to be enabled yet. Deferring single part upload...')
                             time.sleep(1)
+                            
+                        
 
                         for attempt in retry_s3():
                             with attempt:
@@ -1290,23 +1366,27 @@ def download(self, localFilePath, verifyChecksum=True):
             elif self.version:
                 headers = self._s3EncryptionHeaders()
                 key = self.outer.filesBucket.get_key(compat_bytes(self.fileID), validate=False)
-                for attempt in retry_s3():
+                for attempt in retry_s3(predicate=lambda e: retryable_s3_errors(e) or isinstance(e, ChecksumError)):
                     with attempt:
                         with AtomicFileCreate(localFilePath) as tmpPath:
                             key.get_contents_to_filename(tmpPath,
                                                          version_id=self.version,
                                                          headers=headers)
-                if verifyChecksum and self.checksum:
-                    downloadedChecksum = self._get_file_checksum(localFilePath)
-                    if self.checksum != downloadedChecksum:
-                        raise AssertionError(
-                            'Checksums do not match for file %s. Expected: %s Actual: %s' % (
-                                localFilePath, self.checksum, downloadedChecksum))
+                                                         
+                        if verifyChecksum and self.checksum:
+                            try:
+                                # This automatically compares the result and matches the algorithm.
+                                self._get_file_checksum(localFilePath, self.checksum)
+                            except ChecksumError as e:
+                                # Annotate checksum mismatches with file name
+                                raise ChecksumError('Checksums do not match for file %s.' % localFilePath) from e
+                                # The error will get caught and result in a retry of the download until we run out of retries.
+                                # TODO: handle obviously truncated downloads by resuming instead.
             else:
                 assert False
 
         @contextmanager
-        def downloadStream(self):
+        def downloadStream(self, verifyChecksum=True):
             info = self
 
             class DownloadPipe(ReadablePipe):
@@ -1323,9 +1403,41 @@ def writeTo(self, writable):
                                                          version_id=info.version)
                     else:
                         assert False
-
+            
+            class HashingPipe(ReadableTransformingPipe):
+                """
+                Class which checksums all the data read through it. If it
+                reaches EOF and the checksum isn't correct, raises
+                ChecksumError.
+                
+                Assumes info actually has a checksum.
+                """
+            
+                def transform(self, readable, writable):
+                    hasher = info._start_checksum(to_match=info.checksum)
+                    contents = readable.read(1024 * 1024)
+                    while contents != b'':
+                        info._update_checksum(hasher, contents)
+                        try:
+                            writable.write(contents)
+                        except BrokenPipeError:
+                            # Read was stopped early by user code.
+                            # Can't check the checksum.
+                            return
+                        contents = readable.read(1024 * 1024)
+                    # We reached EOF in the input.
+                    # Finish checksumming and verify.
+                    info._finish_checksum(hasher)
+                    # Now stop so EOF happens in the output.
+                    
             with DownloadPipe() as readable:
-                yield readable
+                if verifyChecksum and self.checksum:
+                    # Interpose a pipe to check the hash
+                    with HashingPipe(readable) as verified:
+                        yield verified
+                else:
+                    # No true checksum available, so don't hash
+                    yield readable
 
         def delete(self):
             store = self.outer
diff --git a/src/toil/jobStores/utils.py b/src/toil/jobStores/utils.py
--- a/src/toil/jobStores/utils.py
+++ b/src/toil/jobStores/utils.py
@@ -44,7 +44,7 @@ class WritablePipe(with_metaclass(ABCMeta, object)):
 
     More complicated, less illustrative tests:
 
-    Same as above, but provving that handles are closed:
+    Same as above, but proving that handles are closed:
 
     >>> x = os.dup(0); os.close(x)
     >>> class MyPipe(WritablePipe):
@@ -164,7 +164,7 @@ class ReadablePipe(with_metaclass(ABCMeta, object)):
 
     More complicated, less illustrative tests:
 
-    Same as above, but provving that handles are closed:
+    Same as above, but proving that handles are closed:
 
     >>> x = os.dup(0); os.close(x)
     >>> class MyPipe(ReadablePipe):
@@ -240,3 +240,57 @@ def __exit__(self, exc_type, exc_val, exc_tb):
                 # Only raise the child exception if there wasn't
                 # already an exception in the main thread
                 raise
+                
+class ReadableTransformingPipe(ReadablePipe):
+    """
+    A pipe which is constructed around a readable stream, and which provides a
+    context manager that gives a readable stream.
+    
+    Useful as a base class for pipes which have to transform or otherwise visit
+    bytes that flow through them, instead of just consuming or producing data.
+    
+    Clients should subclass it and implement :meth:`.transform`, like so:
+    
+    >>> import sys, shutil
+    >>> class MyPipe(ReadableTransformingPipe):
+    ...     def transform(self, readable, writable):
+    ...         writable.write(readable.read().decode('utf-8').upper().encode('utf-8'))
+    >>> class SourcePipe(ReadablePipe):
+    ...     def writeTo(self, writable):
+    ...         writable.write('Hello, world!\\n'.encode('utf-8'))
+    >>> with SourcePipe() as source:
+    ...     with MyPipe(source) as transformed:
+    ...         shutil.copyfileobj(codecs.getreader('utf-8')(transformed), sys.stdout)
+    HELLO, WORLD!
+    
+    The :meth:`.transform` method runs in its own thread, and should move data
+    chunk by chunk instead of all at once. It should finish normally if it
+    encounters either an EOF on the readable, or a :class:`BrokenPipeError` on
+    the writable. This means tat it should make sure to actually catch a
+    :class:`BrokenPipeError` when writing.
+    
+    See also: :class:`toil.lib.misc.WriteWatchingStream`.
+    
+    """
+    
+    def __init__(self, source):
+        super(ReadableTransformingPipe, self).__init__()
+        self.source = source
+        
+    @abstractmethod
+    def transform(self, readable, writable):
+        """
+        Implement this method to ship data through the pipe.
+
+        :param file readable: the input stream file object to transform.
+
+        :param file writable: the file object representing the writable end of the pipe. Do not
+        explicitly invoke the close() method of the object, that will be done automatically.
+        """
+        raise NotImplementedError()
+    
+    def writeTo(self, writable):
+        self.transform(self.source, writable)
+    
+    
+


Test patch: