sui_config/
dynamic_transaction_signing_checks.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use serde::{
5    Deserialize,
6    de::{self, Deserializer},
7    ser::Serializer,
8};
9use serde_json::Value as JsonValue;
10use starlark::{
11    ErrorKind,
12    environment::{Globals, Module},
13    eval::Evaluator,
14    syntax::{AstModule, Dialect, DialectTypes},
15    values::{AllocValue, Heap, Value, dict::AllocDict},
16};
17use sui_types::{
18    base_types::ObjectRef,
19    signature::GenericSignature,
20    transaction::{InputObjectKind, TransactionData},
21};
22use tracing::warn;
23
24/// The name of the input variables that the transaction will have in the Starlark program.
25const TX_DATA_NAME: &str = "tx_data";
26const TX_SIGNERS_NAME: &str = "tx_signers";
27const TX_INPUT_OBJECTS_NAME: &str = "tx_input_objects";
28const TX_RECEIVING_OBJECTS_NAME: &str = "tx_receiving_objects";
29const TX_DIGEST_NAME: &str = "tx_digest";
30
31/// The dummy name of the Starlark file being executed. We will just be passing the string of the
32/// program in directly so this is not important but may appear in error messages.
33const STAR_INPUT_FILE_NAME: &str = "dynamic_transaction_signing_checks.star";
34
35#[derive(Debug, thiserror::Error)]
36pub enum DynamicCheckRunnerError {
37    #[error("Failed to serialize transaction data to JSON: {0}")]
38    JSONSerializationError(String),
39    #[error("Failed to parse Starlark program value -- unsupported number type {0}")]
40    UnsupportedNumberFormat(String),
41    #[error(
42        "Failed to execute Starlark program -- invalid return type expected a bool but got {0}"
43    )]
44    InvalidReturnType(String),
45    #[error("Failed to execute Starlark program: {0}")]
46    ExecutionError(ErrorKind),
47    #[error("Failed to load Starlark program: {0}")]
48    LoadingError(ErrorKind),
49    #[error("Check failed -- transaction denied")]
50    CheckFailure,
51}
52
53#[derive(Debug, Clone)]
54pub struct DynamicCheckRunnerContext {
55    module: AstModule,
56    globals: Globals,
57    loaded_program: String,
58}
59
60const DIALECT: Dialect = Dialect {
61    enable_def: true,
62    enable_lambda: true,
63    enable_keyword_only_arguments: false,
64    enable_positional_only_arguments: false,
65    enable_types: DialectTypes::Disable,
66    // NB: set loader to false to prevent any external loading
67    enable_load: false,
68    enable_load_reexport: false,
69    // NB: Allow for top level statements to be used (e.g., top-level `for`, `if`, etc.)
70    enable_top_level_stmt: true,
71    enable_f_strings: false,
72    // NB: We explicitly fully initalize the struct to prevent any future changes to the dialect
73    // without us noticing and deciding whether or not the new feature should be enabled.
74    _non_exhaustive: (),
75};
76
77impl DynamicCheckRunnerContext {
78    /// Create a new `DynamicCheckRunnerContext` with the given Starlark program
79    /// `starlark_program` string. This will parse and validate the program is syntactically
80    /// correct and will set up shared (immutable) state that can be reused. This will not
81    /// run the program.
82    ///
83    /// The `starlark_program` string should be a valid Starlark program that returns a boolean
84    /// value when run -- `True` in the case that the transaction should be allowed, or `False` if
85    /// the transaction should be denied. Any other return value other than `True` (including
86    /// errors) should be considered a denial.
87    pub fn new(starlark_program: String) -> Result<Self, DynamicCheckRunnerError> {
88        // Adds global functions and variables to the dialect (e.g., True, False, Maps, Lists, etc.)
89        // The full spec of what exactly is added here can be found here:
90        // https://github.com/bazelbuild/starlark/blob/master/spec.md#built-in-constants-and-functions
91        let globals = Globals::standard();
92        warn!(
93            "Dynamic transaction checks are enabled. Make sure that you intend to be running \
94            dynamic checks on transactions."
95        );
96        let module = AstModule::parse(STAR_INPUT_FILE_NAME, starlark_program.clone(), &DIALECT)
97            .map_err(|e| DynamicCheckRunnerError::LoadingError(e.into_kind()))?;
98        Ok(Self {
99            module,
100            globals,
101            loaded_program: starlark_program,
102        })
103    }
104
105    /// Run the Starlark program in `self` with the given transaction data, signatures, input
106    /// object kinds, and receiving objects.
107    pub fn run_predicate(
108        &self,
109        tx_data: &TransactionData,
110        tx_signatures: &[GenericSignature],
111        input_object_kinds: &[InputObjectKind],
112        receiving_objects: &[ObjectRef],
113    ) -> Result<(), DynamicCheckRunnerError> {
114        let tx_data_json = serde_json::to_value(tx_data)
115            .map_err(|e| DynamicCheckRunnerError::JSONSerializationError(e.to_string()))?;
116        let tx_signatures_json = serde_json::to_value(tx_signatures)
117            .map_err(|e| DynamicCheckRunnerError::JSONSerializationError(e.to_string()))?;
118        let input_object_kinds_json = serde_json::to_value(input_object_kinds)
119            .map_err(|e| DynamicCheckRunnerError::JSONSerializationError(e.to_string()))?;
120        let receiving_objects_json = serde_json::to_value(receiving_objects)
121            .map_err(|e| DynamicCheckRunnerError::JSONSerializationError(e.to_string()))?;
122        let digest_json = serde_json::to_value(tx_data.digest())
123            .map_err(|e| DynamicCheckRunnerError::JSONSerializationError(e.to_string()))?;
124
125        self.run_starlark_predicate(
126            &tx_data_json,
127            &tx_signatures_json,
128            &input_object_kinds_json,
129            &receiving_objects_json,
130            &digest_json,
131        )
132    }
133
134    fn run_starlark_predicate(
135        &self,
136        tx_data: &JsonValue,
137        tx_signatures: &JsonValue,
138        tx_input_object_kinds: &JsonValue,
139        tx_receiving_objects: &JsonValue,
140        tx_digest: &JsonValue,
141    ) -> Result<(), DynamicCheckRunnerError> {
142        let heap = Heap::new();
143        let env = Module::new();
144
145        let tx_data_value = Self::json_to_starlark(tx_data, &heap)?;
146        let tx_signers_value = Self::json_to_starlark(tx_signatures, &heap)?;
147        let tx_input_object_kinds_value = Self::json_to_starlark(tx_input_object_kinds, &heap)?;
148        let tx_receiving_objects_value = Self::json_to_starlark(tx_receiving_objects, &heap)?;
149        let tx_digest_value = Self::json_to_starlark(tx_digest, &heap)?;
150
151        env.set(TX_DATA_NAME, tx_data_value);
152        env.set(TX_SIGNERS_NAME, tx_signers_value);
153        env.set(TX_INPUT_OBJECTS_NAME, tx_input_object_kinds_value);
154        env.set(TX_RECEIVING_OBJECTS_NAME, tx_receiving_objects_value);
155        env.set(TX_DIGEST_NAME, tx_digest_value);
156
157        let mut evaluator = Evaluator::new(&env);
158        let output_value = evaluator
159            .eval_module(self.module.clone(), &self.globals)
160            .map_err(|e| DynamicCheckRunnerError::ExecutionError(e.into_kind()))?;
161        let transaction_allowed = output_value
162            .unpack_bool()
163            .ok_or_else(|| DynamicCheckRunnerError::InvalidReturnType(output_value.to_repr()))?;
164        if transaction_allowed {
165            Ok(())
166        } else {
167            Err(DynamicCheckRunnerError::CheckFailure)
168        }
169    }
170
171    fn json_to_starlark<'v>(
172        value: &JsonValue,
173        heap: &'v Heap,
174    ) -> Result<Value<'v>, DynamicCheckRunnerError> {
175        Ok(match value {
176            JsonValue::Null => Value::new_none(),
177            JsonValue::Bool(b) => Value::new_bool(*b),
178            JsonValue::Number(n) => {
179                if let Some(i) = n.as_u64() {
180                    heap.alloc(i)
181                } else {
182                    return Err(DynamicCheckRunnerError::UnsupportedNumberFormat(
183                        n.to_string(),
184                    ));
185                }
186            }
187            JsonValue::String(s) => heap.alloc(s.as_str()),
188            JsonValue::Array(arr) => {
189                let list: Vec<_> = arr
190                    .iter()
191                    .map(|v| Self::json_to_starlark(v, heap))
192                    .collect::<Result<_, _>>()?;
193                list.alloc_value(heap)
194            }
195            JsonValue::Object(obj) => {
196                let kvs: Vec<_> = obj
197                    .iter()
198                    .map(|(k, v)| {
199                        let key = heap.alloc(k.as_str());
200                        let val = Self::json_to_starlark(v, heap)?;
201                        Ok((key, val))
202                    })
203                    .collect::<Result<_, _>>()?;
204                heap.alloc(AllocDict(kvs))
205            }
206        })
207    }
208}
209
210// Custom serialization/deserialization for the `DynamicCheckRunnerContext` struct. This allows us
211// to validate the program at the time that it is deserialized, rather than at the time that it is
212// first used. This provides better error message locality and allows us to fail fast if the
213// program is invalid. We keep the invariant here that `serialize(deserialize(program)) ==
214// program`.
215
216/// Deserialize a `DynamicCheckRunnerContext` from a string. This will parse the string as a
217/// Starlark program and validate that it is syntactically correct and setup the
218/// `DynamicCheckRunnerContext` for it. If the program is syntactically invalid, an error will be
219/// returned.
220pub(crate) fn deserialize_dynamic_transaction_checks<'de, D>(
221    deserializer: D,
222) -> Result<Option<DynamicCheckRunnerContext>, D::Error>
223where
224    D: Deserializer<'de>,
225{
226    let path_opt: Option<String> = Option::deserialize(deserializer)?;
227    match path_opt {
228        Some(p) => Ok(Some(
229            DynamicCheckRunnerContext::new(p).map_err(de::Error::custom)?,
230        )),
231        None => Ok(None),
232    }
233}
234
235/// Takes a `DynamicCheckRunnerContext` and serializes the original source program as the returned
236/// string. No parsed state or otherwise is serialized.
237pub(crate) fn serialize_dynamic_transaction_checks<S>(
238    value: &Option<DynamicCheckRunnerContext>,
239    serializer: S,
240) -> Result<S::Ok, S::Error>
241where
242    S: Serializer,
243{
244    match value {
245        Some(DynamicCheckRunnerContext { loaded_program, .. }) => {
246            serializer.serialize_some(&loaded_program)
247        }
248        None => serializer.serialize_none(),
249    }
250}
251
252#[cfg(test)]
253mod test {
254    #[test]
255    fn parse_on_load_invalid() {
256        let program = r#"
257            def main(): return 1
258        "#;
259        let result = super::DynamicCheckRunnerContext::new(program.to_string());
260        assert!(result.is_err());
261    }
262
263    #[test]
264    fn parse_on_load_valid() {
265        let program = r#"
266def main(): 
267    return 1
268        "#;
269        let result = super::DynamicCheckRunnerContext::new(program.to_string());
270        assert!(result.is_ok());
271    }
272}