Change state file to a state dir.

We may want to store multiple state files in the future.

Also, lock the state dir to ensure only one instance is running for a
state dir.

1 files changed, 83 insertions, 14 deletions

diff --git a/fshare.py b/fshare.py
index 56f554b..72a8d27 100755
--- a/fshare.py
+++ b/fshare.py
@@ -17,6 +17,7 @@
 # fshare.  If not, see <http://www.gnu.org/licenses/>.
 
 import argparse
+import fcntl
 import hmac
 import json
 import os
@@ -34,17 +35,62 @@ from urllib import parse as urlparse
 
 # TODO: detect and store mime types
 
+def _open_dirfd(path, mode = 'r', perms = 0o644, dir_fd = None, **kwargs):
+    """
+    Same as the built-in open(), but with support for file permissions
+    and operation with respect to a directory FD.
+    """
+    flags = 0
+
+    if '+' in mode:
+        flags |= os.O_RDWR
+    else:
+        if 'r' in mode:
+            flags |= os.O_RDONLY
+        elif 'w' in mode or 'a' in mode:
+            flags |= os.O_WRONLY
+
+    if 'w' in mode:
+        flags |= os.O_CREAT | os.O_TRUNC
+    elif 'a' in mode:
+        flags |= os.O_CREAT | os.O_TRUNC
+
+    fdopen_mode = mode
+    if 'x' in mode:
+        flags |= os.O_EXCL
+        fdopen_mode = fdopen_mode.replace('x', '')
+
+    fd = os.open(path, flags, mode = perms, dir_fd = dir_fd)
+    try:
+        return os.fdopen(fd, fdopen_mode, **kwargs)
+    except:
+        os.close(fd)
+        raise
+
+class StateCorruptError(Exception):
+    pass
+
 class PersistentState:
     key = None
 
-    def __init__(self, path):
-        if not os.path.exists(path):
-            data = { 'key' : secrets.token_hex(16) }
-            with open(path, 'x') as f:
+    _fname = 'state'
+
+    def __init__(self, state_dir_fd, logger):
+        try:
+            with _open_dirfd(self._fname, 'w+x', perms = 0o600,
+                             dir_fd = state_dir_fd) as f:
+                data = { 'key' : secrets.token_hex(16) }
                 json.dump(data, f, indent = 4)
+                logger.info('Generated a new state file')
+        except FileExistsError:
+            pass
+
+        with _open_dirfd(self._fname, 'r', dir_fd = state_dir_fd) as f:
+            try:
+                data = json.load(f)
+            except json.decoder.JSONDecodeError:
+                raise StateCorruptError
 
-        with open(path, 'r') as f:
-            data = json.load(f)
 
         self.key = bytes.fromhex(data['key'])
 
@@ -302,7 +348,7 @@ group.add_argument('-6', '--ipv6',    action = 'store_true')
 parser.add_argument('-l', '--loglevel', default = 'WARNING')
 parser.add_argument('-d', '--debug', action = 'store_true', help = 'log to stderr')
 
-parser.add_argument('state_file')
+parser.add_argument('state_dir')
 parser.add_argument('data_dir')
 
 args = parser.parse_args(sys.argv[1:])
@@ -329,10 +375,33 @@ def excepthook(t, v, tb, logger = logger):
     logger.error('Uncaught top-level exception', exc_info = (t, v, tb))
 sys.excepthook = excepthook
 
-# read the state file
-state = PersistentState(args.state_file)
-
-# launch the server
-server = FShareServer((args.address, args.port), args.ipv4, args.ipv6,
-                      state, args.data_dir, logger)
-server.serve_forever()
+# open the state dir
+try:
+    state_dir_fd = os.open(args.state_dir, os.O_RDONLY | os.O_DIRECTORY)
+except (FileNotFoundError, NotADirectoryError) as e:
+    logger.error('The state directory "%s" is not an existing directory: %s',
+                 args.state_dir, e)
+    sys.exit(1)
+
+# lock the state dir
+try:
+    fcntl.flock(state_dir_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+except BlockingIOError:
+    logger.error('The state directory is already locked by another process')
+    os.close(state_dir_fd)
+    sys.exit(1)
+
+try:
+    # read the state file
+    state = PersistentState(state_dir_fd, logger)
+
+    # launch the server
+    server = FShareServer((args.address, args.port), args.ipv4, args.ipv6,
+                          state, args.data_dir, logger)
+    server.serve_forever()
+except StateCorruptError:
+    logger.error('Corrupted state file')
+    sys.exit(1)
+finally:
+    fcntl.flock(state_dir_fd, fcntl.LOCK_UN)
+    os.close(state_dir_fd)