From 694712252069d00c475014c959b4a03beba9d5b9 Mon Sep 17 00:00:00 2001
From: s3lph <account-gitlab-ideynizv@kernelpanic.lol>
Date: Sun, 3 Oct 2021 03:16:31 +0200
Subject: [PATCH] Add file locking

---
 CHANGELOG.md        | 11 +++++++++++
 ROADMAP.md          |  4 ----
 easywks/__init__.py |  2 +-
 easywks/files.py    | 40 +++++++++++++++++++++++++++-------------
 4 files changed, 39 insertions(+), 18 deletions(-)
 delete mode 100644 ROADMAP.md

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 940f508..cdd05b9 100644
--- a/CHANGELOG.md
+++ b/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
 
diff --git a/ROADMAP.md b/ROADMAP.md
deleted file mode 100644
index 62e2ff8..0000000
--- a/ROADMAP.md
+++ /dev/null
@@ -1,4 +0,0 @@
-# EasyWKS Roadmap
-
-- [ ] Figure out whether file locking in the working directory is necessary to avoid races.
-- [ ] Testing, testing, testing!
\ No newline at end of file
diff --git a/easywks/__init__.py b/easywks/__init__.py
index 82a4917..1c808d0 100644
--- a/easywks/__init__.py
+++ b/easywks/__init__.py
@@ -1,2 +1,2 @@
 
-__version__ = '0.1.6'
+__version__ = '0.1.7'
diff --git a/easywks/files.py b/easywks/files.py
index 9179964..8081ea9 100644
--- a/easywks/files.py
+++ b/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):