Skip to content

Commit 12629c8

Browse files
chadsec1chadsec1
committed
feat: Improve STRANDLOCK availability
1 parent aa09676 commit 12629c8

File tree

5 files changed

+135
-62
lines changed

5 files changed

+135
-62
lines changed

‎logic/background_worker.py‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
from core.requests import http_request
22
from logic.smp import smp_unanswered_questions, smp_data_handler
3-
from logic.pfs import pfs_data_handler, update_ephemeral_keys
3+
from logic.pfs import pfs_data_handler
44
from logic.message import messages_data_handler
55
from logic.user import validate_identifier
66
from core.constants import (
@@ -100,10 +100,6 @@ def background_worker(user_data, user_data_lock, ui_queue, stop_flag):
100100
logger.debug("Received data: %s", str(message)[:3000])
101101

102102
# Sanity check universal message fields
103-
if (not "sender" in message):
104-
logger.error("Impossible condition, either you have discovered a bug in Coldwire, or the server is attempting to denial-of-service you. Skipping data message with no sender...")
105-
continue
106-
107103
if not validate_identifier(message["sender"]):
108104
logger.error("Impossible condition, either you have discovered a bug in Coldwire, or the server is attempting to denial-of-service you. Skipping data message with malformed sender identifier (%s)...", message["sender"])
109105
continue
@@ -189,5 +185,5 @@ def background_worker(user_data, user_data_lock, ui_queue, stop_flag):
189185
# TODO: We need to keep the last used key and use it when decapsulation with new key gives invalid output
190186
# because it might actually take some time for our keys to be uploaded to server + other servers, and to the contact.
191187
#
192-
update_ephemeral_keys(user_data, user_data_lock)
188+
# update_ephemeral_keys(user_data, user_data_lock)
193189

‎logic/contacts.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,17 @@ def save_contact(user_data: dict, user_data_lock, contact_id: str) -> None:
8383
"public_key": None,
8484
"private_key": None,
8585
},
86+
},
87+
"staged_keys": {
88+
CLASSIC_MCELIECE_8_F_NAME: {
89+
"public_key": None,
90+
"private_key": None,
91+
},
92+
ML_KEM_1024_NAME: {
93+
"public_key": None,
94+
"private_key": None,
95+
},
96+
8697

8798
}
8899
},

‎logic/message.py‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -309,7 +309,7 @@ def messages_data_handler(user_data: dict, user_data_lock, user_data_copied: dic
309309

310310
user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][CLASSIC_MCELIECE_8_F_NAME]["rotation_counter"] += 1
311311

312-
new_ml_kem_keys = user_data["tmp"]["new_ml_kem_keys"]
312+
staged_kyber_private_key = bool(user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][ML_KEM_1024_NAME]["private_key"])
313313

314314
rotation_counter = user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][CLASSIC_MCELIECE_8_F_NAME]["rotation_counter"]
315315

@@ -319,8 +319,9 @@ def messages_data_handler(user_data: dict, user_data_lock, user_data_copied: dic
319319
logger.info("Saved contact (%s) new batch of One-Time-Pads, new strand key, and new hash chain seed", contact_id)
320320
save_account_data(user_data, user_data_lock)
321321

322-
323-
if contact_id not in new_ml_kem_keys:
322+
# Why ???????
323+
# Nvm, I know why, PFS.
324+
if not staged_kyber_private_key:
324325
logger.info("Rotating our ephemeral keys")
325326
send_new_ephemeral_keys(user_data, user_data_lock, contact_id, ui_queue)
326327
save_account_data(user_data, user_data_lock)

‎logic/pfs.py‎

Lines changed: 84 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,43 @@
4646
logger = logging.getLogger(__name__)
4747

4848

49+
def send_pfs_ack(user_data: dict, user_data_lock: threading.Lock, contact_id: str, ui_queue: queue.Queue) -> None:
50+
with user_data_lock:
51+
server_url = user_data["server_url"]
52+
auth_token = user_data["token"]
53+
session_headers = user_data["tmp"]["session_headers"]
54+
55+
56+
our_next_strand_nonce = user_data["contacts"][contact_id]["our_next_strand_nonce"]
57+
our_strand_key = user_data["contacts"][contact_id]["our_strand_key"]
58+
59+
60+
our_new_strand_nonce = sha3_512(secrets.token_bytes(XCHACHA20POLY1305_NONCE_LEN))[:XCHACHA20POLY1305_NONCE_LEN]
61+
_, ciphertext_blob = encrypt_xchacha20poly1305(
62+
our_strand_key,
63+
PFS_TYPE + b"\x01" + our_new_strand_nonce,
64+
nonce = our_next_strand_nonce
65+
)
66+
67+
try:
68+
http_request(f"{server_url}/data/send", "POST", metadata = {
69+
"recipient": contact_id
70+
},
71+
blob = ciphertext_blob,
72+
headers = session_headers,
73+
auth_token = auth_token
74+
)
75+
except Exception:
76+
ui_queue.put({"type": "showerror", "title": "Error", "message": "Failed to send our ephemeral keys to the server"})
77+
return
78+
79+
# We update at the very end to ensure if any of previous steps fail, we do not desync our state
80+
with user_data_lock:
81+
user_data["contacts"][contact_id]["our_next_strand_nonce"] = our_new_strand_nonce
82+
83+
84+
85+
4986
def send_new_ephemeral_keys(user_data: dict, user_data_lock: threading.Lock, contact_id: str, ui_queue: queue.Queue) -> None:
5087
"""
5188
Generate, encrypt, and send fresh ephemeral keys to a contact.
@@ -70,12 +107,12 @@ def send_new_ephemeral_keys(user_data: dict, user_data_lock: threading.Lock, con
70107
# we put here because it could've change between time copy finished copying.
71108

72109
our_next_strand_nonce = user_data["contacts"][contact_id]["our_next_strand_nonce"]
73-
110+
our_strand_key = user_data["contacts"][contact_id]["our_strand_key"]
111+
74112
server_url = user_data_copied["server_url"]
75113
auth_token = user_data_copied["token"]
76114
session_headers = user_data_copied["tmp"]["session_headers"]
77115

78-
our_strand_key = user_data_copied["contacts"][contact_id]["our_strand_key"]
79116

80117
rotation_counter = user_data_copied["contacts"][contact_id]["ephemeral_keys"]["our_keys"][CLASSIC_MCELIECE_8_F_NAME]["rotation_counter"]
81118
rotate_at = user_data_copied["contacts"][contact_id]["ephemeral_keys"]["our_keys"][CLASSIC_MCELIECE_8_F_NAME]["rotate_at"]
@@ -111,7 +148,7 @@ def send_new_ephemeral_keys(user_data: dict, user_data_lock: threading.Lock, con
111148
our_new_strand_nonce = sha3_512(secrets.token_bytes(XCHACHA20POLY1305_NONCE_LEN))[:XCHACHA20POLY1305_NONCE_LEN]
112149
_, ciphertext_blob = encrypt_xchacha20poly1305(
113150
our_strand_key,
114-
PFS_TYPE + our_new_strand_nonce + publickeys_hashchain_signature + publickeys_hashchain,
151+
PFS_TYPE + b"\x00" + our_new_strand_nonce + publickeys_hashchain_signature + publickeys_hashchain,
115152
nonce = our_next_strand_nonce,
116153
max_padding = 1024
117154
)
@@ -133,60 +170,24 @@ def send_new_ephemeral_keys(user_data: dict, user_data_lock: threading.Lock, con
133170
with user_data_lock:
134171
user_data["contacts"][contact_id]["our_next_strand_nonce"] = our_new_strand_nonce
135172

136-
user_data["tmp"]["new_ml_kem_keys"][contact_id] = {
137-
"private_key": kyber_private_key,
138-
"public_key": kyber_public_key
139-
}
140-
141-
if rotate_mceliece:
142-
user_data["tmp"]["new_code_kem_keys"][contact_id] = {
143-
"private_key": mceliece_private_key,
144-
"public_key": mceliece_public_key
145-
}
146-
147-
user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][CLASSIC_MCELIECE_8_F_NAME]["rotation_counter"] = 0
148-
user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][CLASSIC_MCELIECE_8_F_NAME]["rotate_at"] = CLASSIC_MCELIECE_8_F_ROTATE_AT
149173

174+
user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][ML_KEM_1024_NAME]["private_key"] = kyber_private_key
175+
user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][ML_KEM_1024_NAME]["public_key"] = kyber_public_key
150176

151177

178+
if rotate_mceliece:
179+
user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][CLASSIC_MCELIECE_8_F_NAME]["private_key"] = mceliece_private_key
180+
user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][CLASSIC_MCELIECE_8_F_NAME]["public_key"] = mceliece_public_key
152181

153-
user_data["contacts"][contact_id]["lt_sign_keys"]["our_hash_chain"] = our_hash_chain
154-
155-
156-
157-
def update_ephemeral_keys(user_data: dict, user_data_lock: threading.Lock) -> None:
158-
"""
159-
Commit newly generated ephemeral keys to permanent storage.
160-
161-
- Moves keys from `user_data["tmp"]` into `user_data["contacts"]`.
162-
- Clears temporary key buffers after update.
163-
- Saves account state to disk.
164182

165-
Args:
166-
user_data (dict): Shared user account state.
167-
user_data_lock (threading.Lock): Lock protecting shared state.
168-
"""
169-
with user_data_lock:
170-
new_ml_kem_keys = user_data["tmp"]["new_ml_kem_keys"]
171-
new_code_kem_keys = user_data["tmp"]["new_code_kem_keys"]
183+
user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][CLASSIC_MCELIECE_8_F_NAME]["rotation_counter"] = 0
184+
user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][CLASSIC_MCELIECE_8_F_NAME]["rotate_at"] = CLASSIC_MCELIECE_8_F_ROTATE_AT
172185

173-
for contact_id, v in new_ml_kem_keys.items():
174-
with user_data_lock:
175-
user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][ML_KEM_1024_NAME]["private_key"] = v["private_key"]
176-
user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][ML_KEM_1024_NAME]["public_key"] = v["public_key"]
177186

178-
for contact_id, v in new_code_kem_keys.items():
179-
with user_data_lock:
180-
user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][CLASSIC_MCELIECE_8_F_NAME]["private_key"] = v["private_key"]
181-
user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][CLASSIC_MCELIECE_8_F_NAME]["public_key"] = v["public_key"]
182187

183188

184-
with user_data_lock:
185-
user_data["tmp"]["new_ml_kem_keys"] = {}
186-
user_data["tmp"]["new_code_kem_keys"] = {}
189+
user_data["contacts"][contact_id]["lt_sign_keys"]["our_hash_chain"] = our_hash_chain
187190

188-
189-
save_account_data(user_data, user_data_lock)
190191

191192

192193

@@ -206,6 +207,7 @@ def pfs_data_handler(user_data: dict, user_data_lock: threading.Lock, user_data_
206207
user_data_lock (threading.Lock): Lock protecting shared state.
207208
user_data_copied (dict): A read-only copy of user_data for consistency.
208209
ui_queue (queue.Queue): UI queue for notifications/errors.
210+
contact_id (str): Sender ID.
209211
pfs_plaintext (bytes): Decrypted Incoming PFS plaintext from the server.
210212
211213
Returns:
@@ -232,6 +234,33 @@ def pfs_data_handler(user_data: dict, user_data_lock: threading.Lock, user_data_
232234
logger.error("Contact (%s) strand key key is missing! Skipping message...", contact_id)
233235
return
234236

237+
if pfs_plaintext[0] == 1:
238+
logger.info("Received acknowlegement of PFS keys from contact %s", contact_id)
239+
with user_data_lock:
240+
user_data["contacts"][contact_id]["contact_next_strand_nonce"] = pfs_plaintext[1:]
241+
242+
user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][ML_KEM_1024_NAME]["private_key"] = user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][ML_KEM_1024_NAME]["private_key"]
243+
user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][ML_KEM_1024_NAME]["public_key"] = user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][ML_KEM_1024_NAME]["public_key"]
244+
245+
if user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][ML_KEM_1024_NAME]["private_key"]:
246+
user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][CLASSIC_MCELIECE_8_F_NAME]["private_key"] = user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][CLASSIC_MCELIECE_8_F_NAME]["private_key"]
247+
user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][CLASSIC_MCELIECE_8_F_NAME]["public_key"] = user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][CLASSIC_MCELIECE_8_F_NAME]["public_key"]
248+
249+
user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][ML_KEM_1024_NAME]["private_key"] = None
250+
user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][ML_KEM_1024_NAME]["public_key"] = None
251+
252+
user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][CLASSIC_MCELIECE_8_F_NAME]["private_key"] = None
253+
user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][CLASSIC_MCELIECE_8_F_NAME]["public_key"] = None
254+
255+
save_account_data(user_data, user_data_lock)
256+
return
257+
258+
elif pfs_plaintext[0] == 0:
259+
pfs_plaintext = pfs_plaintext[1:]
260+
else:
261+
logger.error("Skipping unknown PFS of type (%d) from contact (%s)", pfs_plaintext[0], contact_id)
262+
return
263+
235264

236265
if (
237266
(len(pfs_plaintext) < ML_KEM_1024_PK_LEN + ML_DSA_87_SIGN_LEN + KEYS_HASH_CHAIN_LEN)
@@ -283,6 +312,10 @@ def pfs_data_handler(user_data: dict, user_data_lock: threading.Lock, user_data_
283312
logger.info("contact (%s) has rotated their Kyber keys", contact_id)
284313

285314

315+
logger.info("We are acknowledging contact's new PFS keys")
316+
send_pfs_ack(user_data, user_data_lock, contact_id, ui_queue)
317+
318+
286319
with user_data_lock:
287320
user_data["contacts"][contact_id]["contact_next_strand_nonce"] = contact_next_strand_nonce
288321

@@ -292,10 +325,11 @@ def pfs_data_handler(user_data: dict, user_data_lock: threading.Lock, user_data_
292325
our_kyber_private_key = user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][ML_KEM_1024_NAME]["private_key"]
293326
our_mceliece_private_key = user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][CLASSIC_MCELIECE_8_F_NAME]["private_key"]
294327

295-
new_ml_kem_keys = user_data["tmp"]["new_ml_kem_keys"]
296-
new_code_kem_keys = user_data["tmp"]["new_code_kem_keys"]
297328

298-
if (our_kyber_private_key is None or our_mceliece_private_key is None) and ((contact_id not in new_ml_kem_keys) and (contact_id not in new_code_kem_keys)):
329+
staged_kem_private_key = user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][ML_KEM_1024_NAME]["private_key"]
330+
staged_code_private_key = user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][CLASSIC_MCELIECE_8_F_NAME]["private_key"]
331+
332+
if (our_kyber_private_key is None or our_mceliece_private_key is None) and ((staged_kem_private_key is None) and (staged_code_private_key is None)):
299333
logger.info("We are sending the contact (%s) our ephemeral keys because we didnt do it before.", contact_id)
300334
send_new_ephemeral_keys(user_data, user_data_lock, contact_id, ui_queue)
301335

‎logic/storage.py‎

Lines changed: 34 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -44,9 +44,7 @@ def load_account_data(password = None) -> dict:
4444

4545
user_data["tmp"] = {
4646
"password": password,
47-
"session_headers": secrets.choice(list(browser_headers.values())),
48-
"new_ml_kem_keys": {},
49-
"new_code_kem_keys": {}
47+
"session_headers": secrets.choice(list(browser_headers.values()))
5048
}
5149

5250

@@ -66,6 +64,8 @@ def load_account_data(password = None) -> dict:
6664
except TypeError:
6765
pass
6866

67+
68+
# In-use
6969
try:
7070
user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][ML_KEM_1024_NAME]["private_key"] = b64decode(user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][ML_KEM_1024_NAME]["private_key"], validate=True)
7171
user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][ML_KEM_1024_NAME]["public_key"] = b64decode(user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][ML_KEM_1024_NAME]["public_key"], validate=True)
@@ -78,6 +78,22 @@ def load_account_data(password = None) -> dict:
7878
except TypeError:
7979
pass
8080

81+
# Staging
82+
try:
83+
user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][ML_KEM_1024_NAME]["private_key"] = b64decode(user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][ML_KEM_1024_NAME]["private_key"], validate=True)
84+
user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][ML_KEM_1024_NAME]["public_key"] = b64decode(user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][ML_KEM_1024_NAME]["public_key"], validate=True)
85+
except TypeError:
86+
pass
87+
88+
try:
89+
user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][CLASSIC_MCELIECE_8_F_NAME]["private_key"] = b64decode(user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][CLASSIC_MCELIECE_8_F_NAME]["private_key"], validate=True)
90+
user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][CLASSIC_MCELIECE_8_F_NAME]["public_key"] = b64decode(user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][CLASSIC_MCELIECE_8_F_NAME]["public_key"], validate=True)
91+
except TypeError:
92+
pass
93+
94+
95+
96+
8197
try:
8298
user_data["contacts"][contact_id]["lt_sign_keys"]["contact_public_key"] = b64decode(user_data["contacts"][contact_id]["lt_sign_keys"]["contact_public_key"], validate=True)
8399
except TypeError:
@@ -157,6 +173,7 @@ def save_account_data(user_data: dict, user_data_lock, password = None) -> None:
157173
pass
158174

159175

176+
# In-use
160177
try:
161178
user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][ML_KEM_1024_NAME]["private_key"] = b64encode(user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][ML_KEM_1024_NAME]["private_key"]).decode()
162179
user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][ML_KEM_1024_NAME]["public_key"] = b64encode(user_data["contacts"][contact_id]["ephemeral_keys"]["our_keys"][ML_KEM_1024_NAME]["public_key"]).decode()
@@ -169,6 +186,20 @@ def save_account_data(user_data: dict, user_data_lock, password = None) -> None:
169186
except TypeError:
170187
pass
171188

189+
# Staging
190+
try:
191+
user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][ML_KEM_1024_NAME]["private_key"] = b64encode(user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][ML_KEM_1024_NAME]["private_key"]).decode()
192+
user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][ML_KEM_1024_NAME]["public_key"] = b64encode(user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][ML_KEM_1024_NAME]["public_key"]).decode()
193+
except TypeError:
194+
pass
195+
196+
try:
197+
user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][CLASSIC_MCELIECE_8_F_NAME]["private_key"] = b64encode(user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][CLASSIC_MCELIECE_8_F_NAME]["private_key"]).decode()
198+
user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][CLASSIC_MCELIECE_8_F_NAME]["public_key"] = b64encode(user_data["contacts"][contact_id]["ephemeral_keys"]["staged_keys"][CLASSIC_MCELIECE_8_F_NAME]["public_key"]).decode()
199+
except TypeError:
200+
pass
201+
202+
172203

173204

174205
try:

0 commit comments

Comments
 (0)