|
4 | 4 | import binascii |
5 | 5 | import pathlib |
6 | 6 | import shelve |
| 7 | +import threading |
| 8 | +from contextlib import contextmanager |
7 | 9 |
|
8 | 10 | from cryptography.hazmat.primitives import serialization |
9 | 11 | from cryptography.hazmat.primitives.asymmetric import x25519 |
@@ -77,6 +79,7 @@ def __init__( |
77 | 79 | self._local_root: x25519.X25519PrivateKey | None = None |
78 | 80 | self._local_root_public: bytes | None = None |
79 | 81 | self._cache_path = cache_path |
| 82 | + self._cache_lock = threading.Lock() |
80 | 83 | if local_root: |
81 | 84 | with local_root.open() as f: |
82 | 85 | private_key = serialization.load_pem_private_key(f.read().encode("utf-8"), password=None) |
@@ -130,18 +133,23 @@ def observe_secondary_remote_public_key(self, infuse_id: int, secondary_pub_key: |
130 | 133 | dev.secondary_device_key_id = binascii.crc32(self._local_root_public + dev.device_public_key) & 0xFFFFFF |
131 | 134 | dev.local_shared_key = self._local_root.exchange(device_public_key) |
132 | 135 |
|
| 136 | + @contextmanager |
| 137 | + def _with_cache(self): |
| 138 | + with self._cache_lock as _lock, shelve.open(str(self._cache_path)) as cache: |
| 139 | + yield cache |
| 140 | + |
133 | 141 | def _update_cache(self, infuse_id: int, device_pub_key: bytes, shared_key: bytes): |
134 | 142 | if self._cache_path is None: |
135 | 143 | return |
136 | 144 | infuse_id_str = f"{infuse_id:016x}" |
137 | | - with shelve.open(str(self._cache_path)) as cache: |
| 145 | + with self._with_cache() as cache: |
138 | 146 | cache[infuse_id_str] = {"public_key": device_pub_key, "shared_key": shared_key} |
139 | 147 |
|
140 | 148 | def _from_cache(self, infuse_id: int, device_pub_key: bytes) -> bytes | None: |
141 | 149 | if self._cache_path is None: |
142 | 150 | return None |
143 | 151 | infuse_id_str = f"{infuse_id:016x}" |
144 | | - with shelve.open(str(self._cache_path)) as cache: |
| 152 | + with self._with_cache() as cache: |
145 | 153 | state = cache.get(infuse_id_str, None) |
146 | 154 | if state is None or state["public_key"] != device_pub_key: |
147 | 155 | return None |
|
0 commit comments