Merge "Provides proper error handling on builder unpickle"

5 files changed, 156 insertions, 8 deletions

diff --git a/swift/cli/ringbuilder.py b/swift/cli/ringbuilder.py
index 15de28fe7..280f5a72d 100755
--- a/swift/cli/ringbuilder.py
+++ b/swift/cli/ringbuilder.py
@@ -830,10 +830,18 @@ def main(arguments=None):
 
     builder_file, ring_file = parse_builder_ring_filename_args(argv)
 
-    if exists(builder_file):
+    try:
         builder = RingBuilder.load(builder_file)
-    elif len(argv) < 3 or argv[2] not in('create', 'write_builder'):
-        print 'Ring Builder file does not exist: %s' % argv[1]
+    except exceptions.UnPicklingError as e:
+        print e
+        exit(EXIT_ERROR)
+    except (exceptions.FileNotFoundError, exceptions.PermissionError) as e:
+        if len(argv) < 3 or argv[2] not in('create', 'write_builder'):
+            print e
+            exit(EXIT_ERROR)
+    except Exception as e:
+        print 'Problem occurred while reading builder file: %s. %s' % (
+            argv[1], e.message)
         exit(EXIT_ERROR)
 
     backup_dir = pathjoin(dirname(argv[1]), 'backups')
diff --git a/swift/common/exceptions.py b/swift/common/exceptions.py
index 856c48975..8cc5fd446 100644
--- a/swift/common/exceptions.py
+++ b/swift/common/exceptions.py
@@ -123,6 +123,18 @@ class DuplicateDeviceError(RingBuilderError):
     pass
 
 
+class UnPicklingError(SwiftException):
+    pass
+
+
+class FileNotFoundError(SwiftException):
+    pass
+
+
+class PermissionError(SwiftException):
+    pass
+
+
 class ListingIterError(SwiftException):
     pass
 
diff --git a/swift/common/ring/builder.py b/swift/common/ring/builder.py
index d75e64551..bad148247 100644
--- a/swift/common/ring/builder.py
+++ b/swift/common/ring/builder.py
@@ -15,6 +15,7 @@
 
 import bisect
 import copy
+import errno
 import itertools
 import math
 import random
@@ -1085,7 +1086,26 @@ class RingBuilder(object):
         :param builder_file: path to builder file to load
         :return: RingBuilder instance
         """
-        builder = pickle.load(open(builder_file, 'rb'))
+        try:
+            fp = open(builder_file, 'rb')
+        except IOError as e:
+            if e.errno == errno.ENOENT:
+                raise exceptions.FileNotFoundError(
+                    'Ring Builder file does not exist: %s' % builder_file)
+            elif e.errno in [errno.EPERM, errno.EACCES]:
+                raise exceptions.PermissionError(
+                    'Ring Builder file cannot be accessed: %s' % builder_file)
+            else:
+                raise
+        else:
+            with fp:
+                try:
+                    builder = pickle.load(fp)
+                except Exception:
+                    # raise error during unpickling as UnPicklingError
+                    raise exceptions.UnPicklingError(
+                        'Ring Builder file is invalid: %s' % builder_file)
+
         if not hasattr(builder, 'devs'):
             builder_dict = builder
             builder = RingBuilder(1, 1, 1)
diff --git a/test/unit/cli/test_ringbuilder.py b/test/unit/cli/test_ringbuilder.py
index eb2d3b672..61da4adb4 100644
--- a/test/unit/cli/test_ringbuilder.py
+++ b/test/unit/cli/test_ringbuilder.py
@@ -13,11 +13,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import mock
 import os
 import tempfile
 import unittest
+import uuid
 
 import swift.cli.ringbuilder
+from swift.common import exceptions
 from swift.common.ring import RingBuilder
 
 
@@ -199,6 +202,65 @@ class TestCommands(unittest.TestCase):
         ring = RingBuilder.load(self.tmpfile)
         self.assertEqual(ring.replicas, 3.14159265359)
 
+    def test_validate(self):
+        self.create_sample_ring()
+        ring = RingBuilder.load(self.tmpfile)
+        ring.rebalance()
+        ring.save(self.tmpfile)
+        argv = ["", self.tmpfile, "validate"]
+        self.assertRaises(SystemExit, swift.cli.ringbuilder.main, argv)
+
+    def test_validate_empty_file(self):
+        open(self.tmpfile, 'a').close
+        argv = ["", self.tmpfile, "validate"]
+        try:
+            swift.cli.ringbuilder.main(argv)
+        except SystemExit as e:
+            self.assertEquals(e.code, 2)
+
+    def test_validate_corrupted_file(self):
+        self.create_sample_ring()
+        ring = RingBuilder.load(self.tmpfile)
+        ring.rebalance()
+        self.assertTrue(ring.validate())  # ring is valid until now
+        ring.save(self.tmpfile)
+        argv = ["", self.tmpfile, "validate"]
+
+        # corrupt the file
+        with open(self.tmpfile, 'wb') as f:
+            f.write(os.urandom(1024))
+        try:
+            swift.cli.ringbuilder.main(argv)
+        except SystemExit as e:
+            self.assertEquals(e.code, 2)
+
+    def test_validate_non_existent_file(self):
+        rand_file = '%s/%s' % ('/tmp', str(uuid.uuid4()))
+        argv = ["", rand_file, "validate"]
+        try:
+            swift.cli.ringbuilder.main(argv)
+        except SystemExit as e:
+            self.assertEquals(e.code, 2)
+
+    def test_validate_non_accessible_file(self):
+        with mock.patch.object(
+                RingBuilder, 'load',
+                mock.Mock(side_effect=exceptions.PermissionError)):
+            argv = ["", self.tmpfile, "validate"]
+            try:
+                swift.cli.ringbuilder.main(argv)
+            except SystemExit as e:
+                self.assertEquals(e.code, 2)
+
+    def test_validate_generic_error(self):
+        with mock.patch.object(RingBuilder, 'load',
+                               mock.Mock(side_effect=
+                               IOError('Generic error occurred'))):
+            argv = ["", self.tmpfile, "validate"]
+            try:
+                swift.cli.ringbuilder.main(argv)
+            except SystemExit as e:
+                self.assertEquals(e.code, 2)
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/unit/common/ring/test_builder.py b/test/unit/common/ring/test_builder.py
index 155edd6f4..c8476278a 100644
--- a/test/unit/common/ring/test_builder.py
+++ b/test/unit/common/ring/test_builder.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import errno
 import mock
 import operator
 import os
@@ -913,17 +914,25 @@ class TestRingBuilder(unittest.TestCase):
         rb.rebalance()
 
         real_pickle = pickle.load
+        fake_open = mock.mock_open()
+
+        io_error_not_found = IOError()
+        io_error_not_found.errno = errno.ENOENT
+
+        io_error_no_perm = IOError()
+        io_error_no_perm.errno = errno.EPERM
+
+        io_error_generic = IOError()
+        io_error_generic.errno = errno.EOPNOTSUPP
         try:
             #test a legit builder
             fake_pickle = mock.Mock(return_value=rb)
-            fake_open = mock.Mock(return_value=None)
             pickle.load = fake_pickle
             builder = ring.RingBuilder.load('fake.builder', open=fake_open)
             self.assertEquals(fake_pickle.call_count, 1)
             fake_open.assert_has_calls([mock.call('fake.builder', 'rb')])
             self.assertEquals(builder, rb)
             fake_pickle.reset_mock()
-            fake_open.reset_mock()
 
             #test old style builder
             fake_pickle.return_value = rb.to_dict()
@@ -932,7 +941,6 @@ class TestRingBuilder(unittest.TestCase):
             fake_open.assert_has_calls([mock.call('fake.builder', 'rb')])
             self.assertEquals(builder.devs, rb.devs)
             fake_pickle.reset_mock()
-            fake_open.reset_mock()
 
             #test old devs but no meta
             no_meta_builder = rb
@@ -943,10 +951,48 @@ class TestRingBuilder(unittest.TestCase):
             builder = ring.RingBuilder.load('fake.builder', open=fake_open)
             fake_open.assert_has_calls([mock.call('fake.builder', 'rb')])
             self.assertEquals(builder.devs, rb.devs)
-            fake_pickle.reset_mock()
+
+            #test an empty builder
+            fake_pickle.side_effect = EOFError
+            pickle.load = fake_pickle
+            self.assertRaises(exceptions.UnPicklingError,
+                              ring.RingBuilder.load, 'fake.builder',
+                              open=fake_open)
+
+            #test a corrupted builder
+            fake_pickle.side_effect = pickle.UnpicklingError
+            pickle.load = fake_pickle
+            self.assertRaises(exceptions.UnPicklingError,
+                              ring.RingBuilder.load, 'fake.builder',
+                              open=fake_open)
+
+            #test some error
+            fake_pickle.side_effect = AttributeError
+            pickle.load = fake_pickle
+            self.assertRaises(exceptions.UnPicklingError,
+                              ring.RingBuilder.load, 'fake.builder',
+                              open=fake_open)
         finally:
             pickle.load = real_pickle
 
+        #test non existent builder file
+        fake_open.side_effect = io_error_not_found
+        self.assertRaises(exceptions.FileNotFoundError,
+                          ring.RingBuilder.load, 'fake.builder',
+                          open=fake_open)
+
+        #test non accessible builder file
+        fake_open.side_effect = io_error_no_perm
+        self.assertRaises(exceptions.PermissionError,
+                          ring.RingBuilder.load, 'fake.builder',
+                          open=fake_open)
+
+        #test an error other then ENOENT and ENOPERM
+        fake_open.side_effect = io_error_generic
+        self.assertRaises(IOError,
+                          ring.RingBuilder.load, 'fake.builder',
+                          open=fake_open)
+
     def test_save_load(self):
         rb = ring.RingBuilder(8, 3, 1)
         devs = [{'id': 0, 'region': 0, 'zone': 0, 'weight': 1,