From 767ced2b6a4ca72d1296c497c54d0c902e987b14 Mon Sep 17 00:00:00 2001
From: Kavi Gupta <kavi@berkeley.edu>
Date: Wed, 15 Jul 2020 00:51:30 -0700
Subject: [PATCH 1/4] fix cross-platform db issue

---
 client/utils/storage.py | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/client/utils/storage.py b/client/utils/storage.py
index 92dd71ee..8de357cf 100644
--- a/client/utils/storage.py
+++ b/client/utils/storage.py
@@ -1,3 +1,4 @@
+import os
 import ctypes
 import shelve # persistance
 import hmac # security
@@ -30,14 +31,22 @@ def mac(value):
     mac.update(repr(value).encode('utf-8'))
     return mac.hexdigest()
 
+def open_db():
+    try:
+        return shelve.open(SHELVE_FILE)
+    except:
+        os.unlink(SHELVE_FILE)
+        return shelve.open(SHELVE_FILE)
+
+
 def contains(root, key):
     key = '{}-{}'.format(root, key)
-    with shelve.open(SHELVE_FILE) as db:
+    with open_db() as db:
         return key in db
 
 def store(root, key, value):
     key = '{}-{}'.format(root, key)
-    with shelve.open(SHELVE_FILE) as db:
+    with open_db() as db:
         db[key] = {'value': value, 'mac': mac(value)}
     return value
 
@@ -45,7 +54,7 @@ def get(root, key, default=None):
     if not contains(root, key):
         return default
     key = '{}-{}'.format(root, key)
-    with shelve.open(SHELVE_FILE) as db:
+    with open_db() as db:
         data = db[key]
         if not hmac.compare_digest(data['mac'], mac(data['value'])):
             raise ProtocolException('{} was tampered.  Reverse changes, or redownload assignment'.format(SHELVE_FILE))

From 9aa997e6139857a3930d96903b83c0f903d4e107 Mon Sep 17 00:00:00 2001
From: Kavi Gupta <kavi@berkeley.edu>
Date: Wed, 15 Jul 2020 17:19:35 -0700
Subject: [PATCH 2/4] fix

---
 client/utils/storage.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/client/utils/storage.py b/client/utils/storage.py
index 8de357cf..38113dba 100644
--- a/client/utils/storage.py
+++ b/client/utils/storage.py
@@ -35,7 +35,11 @@ def open_db():
     try:
         return shelve.open(SHELVE_FILE)
     except:
-        os.unlink(SHELVE_FILE)
+        # just catch everything because there are a variety of errors that a corrupt db can cause
+        # if there is some other error that will happen on the retry hopefully
+        for name in os.listdir("."):
+            if name.startswith(SHELVE_FILE):
+                os.rename(name, name + ".corrupt")
         return shelve.open(SHELVE_FILE)
 
 

From 3e74bbc7994e4d0dac31dbc99ea6cf7253e0a29d Mon Sep 17 00:00:00 2001
From: Kavi Gupta <kavi@berkeley.edu>
Date: Wed, 15 Jul 2020 17:19:39 -0700
Subject: [PATCH 3/4] add test

---
 tests/end_to_end/storage_corrupt_test.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)
 create mode 100644 tests/end_to_end/storage_corrupt_test.py

diff --git a/tests/end_to_end/storage_corrupt_test.py b/tests/end_to_end/storage_corrupt_test.py
new file mode 100644
index 00000000..2205dfc5
--- /dev/null
+++ b/tests/end_to_end/storage_corrupt_test.py
@@ -0,0 +1,17 @@
+import json
+import tempfile
+
+from tests.end_to_end.end_to_end_test import EndToEndTest
+
+
+class EncryptionTest(EndToEndTest):
+    def testEncrypt(self):
+        self.copy_examples()
+        stdout, stderr = self.run_ok("-q", "q1")
+        self.assertEqual(stderr, "")
+        # mess up the shelve
+        self.add_file(".ok_storage.dir", "\0")
+        self.add_file(".ok_storage.bak", "\0")
+        self.add_file(".ok_storage.dat", "\0")
+        stdout, stderr = self.run_ok("-q", "q1")
+        self.assertEqual(stderr, "")

From f0c3faa89727b7100c5795d668bbcef2f6c370f4 Mon Sep 17 00:00:00 2001
From: Kavi Gupta <kavi@berkeley.edu>
Date: Wed, 15 Jul 2020 18:05:56 -0700
Subject: [PATCH 4/4] fix test

---
 tests/end_to_end/storage_corrupt_test.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/end_to_end/storage_corrupt_test.py b/tests/end_to_end/storage_corrupt_test.py
index 2205dfc5..ec54623a 100644
--- a/tests/end_to_end/storage_corrupt_test.py
+++ b/tests/end_to_end/storage_corrupt_test.py
@@ -8,10 +8,10 @@ class EncryptionTest(EndToEndTest):
     def testEncrypt(self):
         self.copy_examples()
         stdout, stderr = self.run_ok("-q", "q1")
-        self.assertEqual(stderr, "")
+        self.assertOnlyInvalidGrant(stderr)
         # mess up the shelve
         self.add_file(".ok_storage.dir", "\0")
         self.add_file(".ok_storage.bak", "\0")
         self.add_file(".ok_storage.dat", "\0")
         stdout, stderr = self.run_ok("-q", "q1")
-        self.assertEqual(stderr, "")
+        self.assertOnlyInvalidGrant(stderr)