feat(python): support encrypted CDN resolver state

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: nicklasl

openfeature-provider/python/src/confidence/local_resolver.py

@@ -30,6 +30,8 @@ class LocalResolver:
     and Go RecoveringResolver.
     """
 
+    EncryptedState = Tuple[bytes, bytes, object]
+
     def __init__(
         self,
         wasm_bytes: bytes,
@@ -46,6 +48,7 @@ def __init__(
         self._delegate_factory = delegate_factory or WasmResolver
         self._delegate: WasmResolver = self._delegate_factory(wasm_bytes)
         self._current_state: Optional[Tuple[bytes, str, object]] = None
+        self._current_encrypted_state: Optional["LocalResolver.EncryptedState"] = None
         self._buffered_logs: List[bytes] = []
 
     def _reload_instance(self, error: BaseException) -> None:
@@ -70,7 +73,15 @@ def _reload_instance(self, error: BaseException) -> None:
         self._delegate = self._delegate_factory(self._wasm_bytes)
 
         # Restore state if available
-        if self._current_state is not None:
+        if self._current_encrypted_state is not None:
+            encrypted_state, encryption_key, sdk = self._current_encrypted_state
+            try:
+                self._delegate.set_encrypted_resolver_state(
+                    encrypted_state, encryption_key, sdk
+                )
+            except Exception as e:
+                logger.error("Failed to restore encrypted state after reload: %s", e)
+        elif self._current_state is not None:
             state, account_id, sdk = self._current_state
             try:
                 self._delegate.set_resolver_state(state, account_id, sdk)
@@ -90,12 +101,29 @@ def set_resolver_state(
             sdk: Optional SDK identifier and version.
         """
         self._current_state = (state, account_id, sdk)
+        self._current_encrypted_state = None
         try:
             self._delegate.set_resolver_state(state, account_id, sdk)
         except WasmCrashError as error:
             self._reload_instance(error)
             raise
 
+    def set_encrypted_resolver_state(
+        self,
+        encrypted_state: bytes,
+        encryption_key: bytes,
+        sdk: Optional[object] = None,
+    ) -> None:
+        self._current_encrypted_state = (encrypted_state, encryption_key, sdk)
+        self._current_state = None
+        try:
+            self._delegate.set_encrypted_resolver_state(
+                encrypted_state, encryption_key, sdk
+            )
+        except WasmCrashError as error:
+            self._reload_instance(error)
+            raise
+
     def resolve_process(
         self, request: wasm_api_pb2.ResolveProcessRequest
     ) -> wasm_api_pb2.ResolveProcessResponse:

openfeature-provider/python/src/confidence/proto/confidence/wasm/messages_pb2.py

@@ -123,6 +123,7 @@ class ConfidenceProvider(AbstractProvider):
     def __init__(
         self,
         client_secret: str,
+        encryption_key: Optional[str] = None,
         state_poll_interval: float = DEFAULT_STATE_POLL_INTERVAL,
         log_poll_interval: float = DEFAULT_LOG_POLL_INTERVAL,
         assign_poll_interval: float = DEFAULT_ASSIGN_POLL_INTERVAL,
@@ -150,6 +151,7 @@ def __init__(
             wasm_bytes: Optional WASM bytes for testing.
         """
         self._client_secret = client_secret
+        self._encryption_key = encryption_key
         self._state_poll_interval = state_poll_interval
         self._log_poll_interval = log_poll_interval
         self._assign_poll_interval = assign_poll_interval
@@ -229,6 +231,7 @@ def initialize(self, evaluation_context: EvaluationContext) -> None:
             self._state_fetcher = StateFetcher(
                 client_secret=self._client_secret,
                 http_client=self._http_client,
+                encryption_key=self._encryption_key,
             )
 
         # Create flag logger if not injected
@@ -243,12 +246,11 @@ def initialize(self, evaluation_context: EvaluationContext) -> None:
         # Fetch initial state - don't fail if this fails, background thread will retry
         try:
             state, account_id, _ = self._state_fetcher.fetch()
-            if account_id:
-                sdk = types_pb2.Sdk(
-                    id=types_pb2.SdkId.SDK_ID_PYTHON_PROVIDER,
-                    version=__version__,
-                )
-                self._resolver.set_resolver_state(state, account_id, sdk)
+            sdk = types_pb2.Sdk(
+                id=types_pb2.SdkId.SDK_ID_PYTHON_PROVIDER,
+                version=__version__,
+            )
+            if self._set_resolver_state(state, account_id, sdk):
                 self._status = ProviderStatus.READY
                 self.emit_provider_ready(ProviderEventDetails())
                 logger.info("ConfidenceProvider initialized successfully")
@@ -311,6 +313,18 @@ def shutdown(self) -> None:
 
         logger.info("ConfidenceProvider shutdown complete")
 
+    def _set_resolver_state(self, state: bytes, account_id: str, sdk: object) -> bool:
+        if self._encryption_key and self._state_fetcher is not None:
+            cdn_bytes = getattr(self._state_fetcher, "raw_cdn_bytes", None)
+            if cdn_bytes is not None:
+                key_bytes = bytes.fromhex(self._encryption_key)
+                self._resolver.set_encrypted_resolver_state(cdn_bytes, key_bytes, sdk)
+                return True
+        if account_id:
+            self._resolver.set_resolver_state(state, account_id, sdk)
+            return True
+        return False
+
     def _resolve_typed(
         self,
         flag_key: str,
@@ -786,17 +800,17 @@ def _state_poll_loop(self) -> None:
 
             try:
                 state, account_id, changed = self._state_fetcher.fetch()
-                if changed and account_id:
+                if changed:
                     sdk = types_pb2.Sdk(
                         id=types_pb2.SdkId.SDK_ID_PYTHON_PROVIDER,
                         version=__version__,
                     )
                     with self._resolver_lock:
-                        self._resolver.set_resolver_state(state, account_id, sdk)
+                        self._set_resolver_state(state, account_id, sdk)
                     logger.debug("Resolver state updated")
 
-                # If we were NOT_READY and now have valid state, transition to READY
-                if account_id and self._status == ProviderStatus.NOT_READY:
+                has_state = self._encryption_key is not None or bool(account_id)
+                if has_state and self._status == ProviderStatus.NOT_READY:
                     self._status = ProviderStatus.READY
                     self.emit_provider_ready(ProviderEventDetails())
                     logger.info("Provider recovered and is now READY")

openfeature-provider/python/src/confidence/provider.py

@@ -39,6 +39,7 @@ def __init__(
         self,
         client_secret: str,
         http_client: Optional[httpx.Client] = None,
+        encryption_key: Optional[str] = None,
     ) -> None:
         """Initialize the StateFetcher.
 
@@ -47,12 +48,14 @@ def __init__(
             http_client: Optional httpx.Client for custom HTTP configuration or testing.
         """
         self._client_secret = client_secret
+        self._encryption_key = encryption_key
         self._http_client = http_client
         self._owns_client = http_client is None
 
         # Cached state
         self._state: Optional[bytes] = None
         self._account_id: Optional[str] = None
+        self._raw_cdn_bytes: Optional[bytes] = None
         self._etag: Optional[str] = None
 
         # Build CDN URL from SHA256 hash of client secret
@@ -69,6 +72,11 @@ def account_id(self) -> Optional[str]:
         """Return the current cached account ID."""
         return self._account_id
 
+    @property
+    def raw_cdn_bytes(self) -> Optional[bytes]:
+        """Return raw encrypted CDN bytes, or None if unencrypted."""
+        return self._raw_cdn_bytes
+
     def _get_client(self) -> httpx.Client:
         """Get or create the HTTP client."""
         if self._http_client is not None:
@@ -114,22 +122,36 @@ def fetch(self) -> Tuple[bytes, str, bool]:
                     f"Failed to fetch state: HTTP {response.status_code}"
                 )
 
-            # Parse the protobuf response
-            state_request = SetResolverStateRequest()
-            state_request.ParseFromString(response.content)
-
-            # Cache the state, account ID, and ETag
-            self._state = state_request.state
-            self._account_id = state_request.account_id
+            encrypted = response.headers.get("x-amz-meta-encrypted") == "true"
             self._etag = response.headers.get("ETag")
 
-            logger.info(
-                "Loaded resolver state for account=%s, etag=%s",
-                self._account_id,
-                self._etag,
-            )
-
-            return self._state, self._account_id, True
+            if not encrypted:
+                try:
+                    state_request = SetResolverStateRequest()
+                    state_request.ParseFromString(response.content)
+                    self._state = state_request.state
+                    self._account_id = state_request.account_id
+                    self._raw_cdn_bytes = None
+                    logger.info(
+                        "Loaded resolver state for account=%s, etag=%s",
+                        self._account_id,
+                        self._etag,
+                    )
+                    return self._state, self._account_id, True
+                except Exception:
+                    logger.warning(
+                        "Protobuf decode failed, treating state as encrypted"
+                    )
+                    encrypted = True
+
+            if not self._encryption_key:
+                raise StateFetcherError(
+                    "Resolver state is encrypted but no encryption_key was provided. "
+                    "Set the encryption key for this client credential."
+                )
+            self._raw_cdn_bytes = response.content
+            logger.info("Loaded encrypted resolver state, etag=%s", self._etag)
+            return self._raw_cdn_bytes, "", True
 
         finally:
             if should_close:

openfeature-provider/python/src/confidence/state_fetcher.py

@@ -55,6 +55,9 @@ def __init__(self, wasm_bytes: bytes) -> None:
         self._wasm_msg_guest_set_resolver_state = exports[
             "wasm_msg_guest_set_resolver_state"
         ]
+        self._wasm_msg_guest_set_encrypted_resolver_state = exports[
+            "wasm_msg_guest_set_encrypted_resolver_state"
+        ]
         self._wasm_msg_guest_resolve_flags = exports["wasm_msg_guest_resolve_flags"]
         self._wasm_msg_guest_bounded_flush_logs = exports[
             "wasm_msg_guest_bounded_flush_logs"
@@ -119,6 +122,26 @@ def set_resolver_state(
         if resp_ptr != 0:
             self._consume_response(resp_ptr, None)
 
+    def set_encrypted_resolver_state(
+        self,
+        encrypted_state: bytes,
+        encryption_key: bytes,
+        sdk: Optional[protobuf_message.Message] = None,
+    ) -> None:
+        request = messages_pb2.SetEncryptedResolverStateRequest()
+        request.encrypted_state = encrypted_state
+        request.encryption_key = encryption_key
+        if sdk is not None:
+            request.sdk.CopyFrom(sdk)
+
+        req_ptr = self._transfer_request(request)
+        resp_ptr = self._wasm_msg_guest_set_encrypted_resolver_state(
+            self._store, req_ptr
+        )
+
+        if resp_ptr != 0:
+            self._consume_response(resp_ptr, None)
+
     def resolve_process(
         self, request: wasm_api_pb2.ResolveProcessRequest
     ) -> wasm_api_pb2.ResolveProcessResponse:

openfeature-provider/python/src/confidence/wasm_resolver.py

@@ -55,6 +55,24 @@ def test_account_id() -> str:
     return account_path.read_text().strip()
 
 
+@pytest.fixture
+def test_encrypted_state() -> bytes:
+    """Load encrypted test resolver state."""
+    path = get_data_dir() / "resolver_state_encrypted.pb"
+    if not path.exists():
+        pytest.skip(f"Encrypted test state not found at {path}")
+    return path.read_bytes()
+
+
+@pytest.fixture
+def test_encryption_key() -> bytes:
+    """Load test encryption key (raw bytes)."""
+    path = get_data_dir() / "encryption_key_test.hex"
+    if not path.exists():
+        pytest.skip(f"Test encryption key not found at {path}")
+    return bytes.fromhex(path.read_text().strip())
+
+
 @pytest.fixture
 def test_client_secret() -> str:
     """Test client secret for the demo account."""