Add file locking

CHANGELOG.md

@@ -1,5 +1,16 @@
 # EasyWKS Changelog
 
+<!-- BEGIN RELEASE v0.1.7 -->
+## Version 0.1.7
+
+### Changes
+
+<!-- BEGIN CHANGES 0.1.7 -->
+- Add file locking in order to avoid races between LMTP/process and HTTP.
+<!-- END CHANGES 0.1.7 -->
+
+<!-- END RELEASE v0.1.7 -->
+
 <!-- BEGIN RELEASE v0.1.6 -->
 ## Version 0.1.6
 

ROADMAP.md

@@ -1,4 +0,0 @@
-# EasyWKS Roadmap
-
-- [ ] Figure out whether file locking in the working directory is necessary to avoid races.
-- [ ] Testing, testing, testing!

easywks/__init__.py

@@ -1,2 +1,2 @@
 
-__version__ = '0.1.6'
+__version__ = '0.1.7'

easywks/files.py

@@ -1,5 +1,6 @@
 
 import os
+import fcntl
 import stat
 from datetime import datetime, timedelta
 
@@ -10,6 +11,24 @@ from .crypto import create_pgp_key, privkey_to_pubkey
 from .util import hash_user_id
 
 
+def _locked_read(file: str, binary: bool = False):
+    with open(file, 'r' + 'b' * binary) as f:
+        fcntl.lockf(f, fcntl.LOCK_SH)
+        content = f.read()
+        fcntl.lockf(f, fcntl.LOCK_UN)
+    return content
+
+
+def _locked_write(file: str, content, binary: bool = False):
+    with open(file, 'a' + 'b' * binary) as f:
+        fcntl.lockf(f, fcntl.LOCK_EX)
+        f.seek(0)
+        f.truncate()
+        f.write(content)
+        fcntl.lockf(f, fcntl.LOCK_UN)
+    return content
+
+
 def make_submission_address_file(domain: str):
     return Config[domain].submission_address + '\n'
 
@@ -32,49 +51,44 @@ def init_working_directory():
         # Create necessary files and directories
         os.makedirs(os.path.join(wdir, domain, 'hu'), exist_ok=True)
         os.makedirs(os.path.join(wdir, domain, 'pending'), exist_ok=True)
-        with open(os.path.join(wdir, domain, 'submission-address'), 'w') as saf:
-            saf.write(make_submission_address_file(domain))
-        with open(os.path.join(wdir, domain, 'policy'), 'w') as polf:
-            polf.write(make_policy_file(domain))
+        _locked_write(os.path.join(wdir, domain, 'submission-address'), make_submission_address_file(domain))
+        _locked_write(os.path.join(wdir, domain, 'policy'), make_policy_file(domain))
         # Create PGP key if it doesn't exist yet
         create_pgp_key(domain)
         # Export submission key to hu dir
         key = privkey_to_pubkey(domain)
         uid = hash_user_id(Config[domain].submission_address)
-        with open(os.path.join(wdir, domain, 'hu', uid), 'wb') as hu:
-            hu.write(bytes(key))
+        _locked_write(os.path.join(wdir, domain, 'hu', uid), bytes(key), binary=True)
 
 
 def read_public_key(domain, user):
     hu = hash_user_id(user)
     keyfile = os.path.join(Config.working_directory, domain, 'hu', hu)
-    key, _ = PGPKey.from_file(keyfile)
+    key, _ = PGPKey.from_blob(_locked_read(keyfile, binary=True))
     return key
 
 
 def read_hashed_public_key(domain, hu):
     keyfile = os.path.join(Config.working_directory, domain, 'hu', hu)
-    key, _ = PGPKey.from_file(keyfile)
+    key, _ = PGPKey.from_blob(_locked_read(keyfile, binary=True))
     return key
 
 
 def write_public_key(domain, user, key):
     hu = hash_user_id(user)
     keyfile = os.path.join(Config.working_directory, domain, 'hu', hu)
-    with open(keyfile, 'wb') as f:
-        f.write(bytes(key))
+    _locked_write(keyfile, bytes(key), binary=True)
 
 
 def read_pending_key(domain, nonce):
     keyfile = os.path.join(Config.working_directory, domain, 'pending', nonce)
-    key, _ = PGPKey.from_file(keyfile)
+    key, _ = PGPKey.from_blob(_locked_read(keyfile))
     return key
 
 
 def write_pending_key(domain, nonce, key):
     keyfile = os.path.join(Config.working_directory, domain, 'pending', nonce)
-    with open(keyfile, 'w') as f:
-        f.write(str(key))
+    _locked_write(keyfile, str(key))
 
 
 def remove_pending_key(domain, nonce):