Merge pull request #3 from minvws/iv-generator

Added IV generator for counters

Author: jaytaph

.gitignore

@@ -1,3 +1,4 @@
 __pycache__
 .venv
 ssl/*
+secrets/*

app.conf

@@ -46,4 +46,9 @@ path=secrets/keystore.json
 [hsm_keystore]
 library = /usr/lib/softhsm/libsofthsm2.so
 slot = 1588871416
-slot_pin = 1234
+slot_pin = 1234
+
+[iv]
+path = secrets/iv.json
+block_size = 100
+start_value = 9223372036854775807

app.conf.example

@@ -21,7 +21,6 @@ ssl_key_file = server.key
 key_name = REK
 key_version = 1
 alg = AES-256-GCM
-max_encryptions = 9223372036854775807
 key_renewal_at = 1000
 iv_prefix = RIVF
 ; How old may a RID be before it is considered invalid for exchanging to a PDN/BSN via VAD?
@@ -40,6 +39,14 @@ host = localhost
 port = 6379
 db = 0
 
+[iv]
+; Path to IV counter
+path = secrets/iv.json
+; Block size to reserve. If the IV counter is below this value, a new block will be reserved and written to disks
+block_size = 100
+; Start value for the IV counter
+start_value = 9223372036854775807
+
 [hsm_keystore]
 library = /usr/local/lib/softhsm/libsofthsm2.so
 slot = 0

app/config.py

@@ -26,7 +26,6 @@ class ConfigRid(BaseModel):
     key_name : str= Field(default="REK")
     key_version : str= Field(default=1)
     alg : str= Field(default="AES-256-GCM")
-    max_encryptions : str= Field(default=2**63 - 1)
     key_renewal_at : str= Field(default=1000)
     iv_prefix : str = Field(default="RIVF")
     max_age_for_pdn_exchange_via_vad : int = Field(default=10)
@@ -41,6 +40,10 @@ class ConfigBpg(BaseModel):
     class Config:
         extra = "allow"
 
+class ConfigIv(BaseModel):
+    path: str = Field(default="iv.json")
+    block_size: int = Field(default=100)
+    start: int = Field(default=9223372036854775807) # 2^63 - 1
 
 class ConfigUvicorn(BaseModel):
     swagger_enabled: bool = Field(default=False)
@@ -79,6 +82,7 @@ class Config(BaseModel):
     rid: ConfigRid
     bpg: ConfigBpg
     redis: ConfigRedis
+    iv: ConfigIv
     json_keystore: ConfigJsonKeystore
     hsm_keystore: ConfigHsmKeystore
 

app/container.py

@@ -8,6 +8,7 @@
 from app.services.crypto.hsm_crypto_service import HsmCryptoService
 from app.services.crypto.json_keystore import JsonKeyStorage
 from app.services.crypto.memory_crypto_service import MemoryCryptoService
+from app.services.iv_service import IvService
 from app.services.pdn_service import PdnService
 from app.services.rid_cache import RidCache
 from app.services.rid_service import RidService
@@ -26,13 +27,18 @@ def container_config(binder: inject.Binder) -> None:
     else:
         raise ValueError(f"Unknown keystore type {config.app.keystore}")
 
+
+    iv_service = IvService(config.iv.path, config.iv.block_size)
+    iv_service.set_iv_counter(config.iv.start, if_not_exists=True)
+    binder.bind(IvService, iv_service)
+
     redis_service = redis.Redis(config.redis.host, config.redis.port, config.redis.db)
     binder.bind(redis.Redis, redis_service)
 
     rid_cache = RidCache(redis_service)
     binder.bind(RidCache, rid_cache)
 
-    rid_service = RidService(config.rid, crypto_service, rid_cache)
+    rid_service = RidService(config.rid, crypto_service, rid_cache, iv_service)
     binder.bind(RidService, rid_service)
 
     bpg_service = BpgService(config.bpg, crypto_service)

app/services/iv_service.py

@@ -0,0 +1,87 @@
+import json
+import os
+
+from filelock import FileLock
+
+"""
+The IV service is responsible for generating and storing an IV counter. It is used to generate unique IVs for each message that is sent.
+This means that the IV counter must return a unique number each time. We do this by "reserving" a block of counters that can be used.
+Once we hit the low water mark for that block, we reserve a new block and store the new low water mark. This way, we can never (?) use
+a duplicate counter, but we can miss counters in case of a crash. That is not a problem though.
+
+Note we don't case about thread safety here, as we assume that the IV service is only used by one thread at a time.
+
+
+Counter     Low Water Mark
+--------------------------
+  100         95
+  99          95
+  98          95
+  97          95
+  96          95
+
+ Here we hit the low water mark, so we reserve a new block and store the LWM to disk
+ 
+  95          90
+  94          90
+  ...
+"""
+
+class IvError(Exception):
+    pass
+
+class IvService:
+    def __init__(self, filename="iv.json", block_size = 100) -> None:
+        self.filename = filename
+        self.block_size = block_size
+
+        try:
+            self.remaining = self.low_watermark = self.load_counter()
+        except IvError:
+            self.remaining = self.low_watermark = 0
+
+    def set_iv_counter(self, counter: int, if_not_exists: bool = True) -> None:
+        """
+        Manually set the IV counter.
+        """
+        if if_not_exists and os.path.exists(self.filename):
+            return
+        self.store_counter(counter)
+        self.low_watermark = self.remaining = counter
+
+    def get_iv_counter(self) -> int:
+        """
+        Return the current IV counter
+        """
+        if self.remaining <= self.low_watermark:
+            self.low_watermark -= self.block_size
+            self.store_counter(self.low_watermark)
+
+        self.remaining -= 1
+        return self.remaining
+
+    def load_counter(self) -> int:
+        """
+        Load counter from disk
+        """
+        try:
+            lock = FileLock(self.filename + ".lock", timeout=1)
+            with lock:
+                with open(self.filename, 'r') as f:
+                    data = json.load(f)
+                    return data['remaining']
+        except Exception as e:
+            raise IvError(e)
+
+    def store_counter(self, counter: int) -> None:
+        """
+        Store counter to disk
+        """
+        try:
+            lock = FileLock(self.filename + ".lock", timeout=1)
+            with lock:
+                with open(self.filename, 'w') as f:
+                    json.dump({'remaining': counter}, f)
+        except Exception as e:
+            raise IvError(e)
+

app/services/rid_service.py

@@ -8,6 +8,7 @@
 from app.services.crypto.crypto_service import CryptoService
 from app.services.jwe import ALGORITHMS, JWEVerifyException
 from app.services.rid_cache import RidCache
+from app.services.iv_service import IvService
 from app.types import BasePseudonym, Rid
 
 class RidException(Exception):
@@ -43,18 +44,19 @@ class VerificationException(RidException):
 
 
 class RidService:
-    def __init__(self, config: ConfigRid, crypto_service: CryptoService, rid_cache: RidCache) -> None:
+    def __init__(self, config: ConfigRid, crypto_service: CryptoService, rid_cache: RidCache, iv_service: IvService) -> None:
         self.crypto_service = crypto_service
         self.config = config
         self.rid_cache = rid_cache
+        self.iv_service = iv_service
 
         if self.config.alg == "AES-256-GCM":
             self.jwe_alg = ALGORITHMS.A256GCM
         else:
             raise ValueError("Unsupported algorithm")
 
         self.rid_env_key_id = self.config.key_name + "-" + str(self.config.key_version)
-        self.rid_enc_remaining = int(self.config.max_encryptions)
+        self.rid_enc_remaining = self.iv_service.get_iv_counter()
 
     def is_healthy(self) -> bool:
         """
@@ -67,7 +69,7 @@ def generate_iv(self) -> bytes:
         Generate an IV for the encryption
         :return: IV bytes
         """
-        self.rid_enc_remaining -= 1
+        self.rid_enc_remaining = self.iv_service.get_iv_counter()
         if self.rid_enc_remaining == 0:
             raise EncryptionExhaustedException("No more encryption attempts left for key")
 

requirements.txt

@@ -16,4 +16,6 @@ pydantic_core~=2.23.3
 anyio~=4.4.0
 click~=8.1.7
 annotated-types~=0.7.0
-setuptools~=72.2.0
+setuptools~=72.2.0
+cryptography~=43.0.1
+cffi~=1.17.1