From 5fcc5c72acb66c30747d73ab4b2ac6cdf3dec363 Mon Sep 17 00:00:00 2001
From: Bernie Hackett <bernie@10gen.com>
Date: Tue, 16 Jun 2015 16:43:15 -0700
Subject: [PATCH] PYTHON-951 - Raise CorruptGridFile for truncated chunks.

---
 gridfs/grid_file.py |  3 +++
 test/test_gridfs.py | 17 +++++++++++++++--
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/gridfs/grid_file.py b/gridfs/grid_file.py
index 441f12f2c..52505866c 100644
--- a/gridfs/grid_file.py
+++ b/gridfs/grid_file.py
@@ -434,6 +434,9 @@ class GridOut(object):
 
             chunk_data = chunk["data"][self.__position % chunk_size:]
 
+            if not chunk_data:
+                raise CorruptGridFile("truncated chunk")
+
         self.__position += len(chunk_data)
         self.__buffer = EMPTY
         return chunk_data
diff --git a/test/test_gridfs.py b/test/test_gridfs.py
index 65ec3483e..39ebc8308 100644
--- a/test/test_gridfs.py
+++ b/test/test_gridfs.py
@@ -24,12 +24,12 @@ import threading
 import time
 import gridfs
 
+from bson.binary import Binary
 from bson.py3compat import u, StringIO, string_type
 from pymongo.mongo_client import MongoClient
 from pymongo.errors import ConfigurationError, ConnectionFailure
 from pymongo.read_preferences import ReadPreference
-from gridfs.errors import (FileExists,
-                           NoFile)
+from gridfs.errors import CorruptGridFile, FileExists, NoFile
 from test.test_replica_set_client import TestReplicaSetClientBase
 from test import (client_context,
                   unittest,
@@ -159,6 +159,19 @@ class TestGridfs(IntegrationTest):
         self.assertEqual(255 * 1024, raw["chunkSize"])
         self.assertTrue(isinstance(raw["md5"], string_type))
 
+    def test_corrupt_chunk(self):
+        files_id = self.fs.put(b'foobar')
+        self.db.fs.chunks.update_one({'files_id': files_id},
+                                     {'$set': {'data': Binary(b'foo', 0)}})
+        try:
+            out = self.fs.get(files_id)
+            self.assertRaises(CorruptGridFile, out.read)
+
+            out = self.fs.get(files_id)
+            self.assertRaises(CorruptGridFile, out.readline)
+        finally:
+            self.fs.delete(files_id)
+
     def test_delete_ensures_index(self):
         # setUp has dropped collections.
         names = self.db.collection_names()