Skip to main content

pdk_data_storage_lib/
local.rs

1// Copyright (c) 2026, Salesforce, Inc.,
2// All rights reserved.
3// For full license text, see the LICENSE.txt file
4
5//! Local shared data storage primitives
6//!
7//! Provides a lightweight, in-memory shared data API used by policies within a
8//! single node. This module defines the low-level [`LocalStorage`] trait and the
9//! injectable [`SharedData`] implementation that integrates with the host
10//! runtime.
11
12use pdk_core::classy::extract::context::ConfigureContext;
13use pdk_core::classy::extract::{Extract, FromContext};
14use pdk_core::classy::proxy_wasm::types::Status;
15use std::convert::Infallible;
16use std::rc::Rc;
17use thiserror::Error;
18
19/// Defines the behavior for store operations.
20#[derive(PartialEq, Eq, Debug, Copy, Clone)]
21pub enum StoreMode {
22    /// Indicates that the store operation does not depend on any condition.
23    Always,
24    /// Indicates that the store operation should succeed only if no other value was present in the store.
25    Absent,
26    /// Indicates that the store operation should succeed only if the stored value matches the provided cas.
27    Cas(u32),
28}
29
30#[non_exhaustive]
31#[derive(Error, Debug)]
32/// Error type for local storage operations.
33pub enum LocalStorageError {
34    /// Indicates that the operation was aborted due to difference in the expected CAS value
35    #[error("Cas mismatch.")]
36    CasMismatch,
37    /// Unhandled error from the proxy wasm runtime.
38    #[error("Low level error.")]
39    ProxyWasm(Status),
40}
41
42impl From<Status> for LocalStorageError {
43    fn from(value: Status) -> Self {
44        Self::ProxyWasm(value)
45    }
46}
47
48/// The [`LocalStorage`] trait defines the methods for storing and retrieving data.
49pub trait LocalStorage {
50    /// Sets the value of a key in the store with the particular [`StoreMode`].
51    fn set(&self, key: &str, value: &[u8], mode: StoreMode) -> Result<(), LocalStorageError>;
52
53    /// Retrieves the value of a key from the store.
54    fn get(&self, key: &str) -> Result<Option<(Vec<u8>, u32)>, LocalStorageError>;
55
56    /// Deletes a key from the store.
57    // proxy_wasm_rust_sdk is panicking when a cas mismatch is returned, therefore we cannot
58    // take advantage of the cas for delete.
59    fn delete(&self, key: &str) -> Result<(), LocalStorageError>;
60
61    /// Returns all the keys in the store.
62    fn keys(&self) -> Vec<String>;
63}
64
65#[derive(Clone)]
66/// Implementation of [`LocalStorage`] trait for the shared data.
67pub struct SharedData {
68    shared_data: Rc<dyn pdk_core::classy::SharedData>,
69}
70
71impl SharedData {
72    fn store_mode_to_cas(
73        &self,
74        key: &str,
75        mode: StoreMode,
76    ) -> Result<Option<u32>, LocalStorageError> {
77        match mode {
78            StoreMode::Always => Ok(None),
79            StoreMode::Absent => {
80                match self.shared_data.shared_data_get(key) {
81                    (Some(data), cas) => {
82                        if data.is_empty() {
83                            Ok(Some(cas.unwrap_or(u32::MAX))) // Should never be undefined since it has data
84                        } else {
85                            Err(LocalStorageError::CasMismatch)
86                        }
87                    }
88                    (None, Some(cas)) => Ok(Some(cas)), // Scenario of store with empty value
89                    (None, None) => Ok(Some(u32::MAX)),
90                }
91            }
92            StoreMode::Cas(cas) => Ok(Some(cas)),
93        }
94    }
95}
96
97impl FromContext<ConfigureContext> for SharedData {
98    type Error = Infallible;
99
100    fn from_context(context: &ConfigureContext) -> Result<Self, Self::Error> {
101        let shared_data: Rc<dyn pdk_core::classy::SharedData> = context.extract()?;
102        Ok(Self { shared_data })
103    }
104}
105
106impl LocalStorage for SharedData {
107    fn set(&self, key: &str, value: &[u8], mode: StoreMode) -> Result<(), LocalStorageError> {
108        let cas = self.store_mode_to_cas(key, mode)?;
109
110        self.shared_data
111            .shared_data_set(key, value, cas)
112            .map_err(|e| match e {
113                Status::CasMismatch => LocalStorageError::CasMismatch,
114                v => LocalStorageError::ProxyWasm(v),
115            })
116    }
117
118    fn get(&self, key: &str) -> Result<Option<(Vec<u8>, u32)>, LocalStorageError> {
119        match self.shared_data.shared_data_get(key) {
120            (Some(data), Some(cas)) => Ok(Some((data, cas))),
121            _ => Ok(None),
122        }
123    }
124
125    fn delete(&self, key: &str) -> Result<(), LocalStorageError> {
126        let cas = self.store_mode_to_cas(key, StoreMode::Always)?;
127        let _ = self.shared_data.shared_data_remove(key, cas)?;
128        Ok(())
129    }
130
131    fn keys(&self) -> Vec<String> {
132        self.shared_data.shared_data_keys()
133    }
134}