1use 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
24const 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
31const 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 enable_load: false,
68 enable_load_reexport: false,
69 enable_top_level_stmt: true,
71 enable_f_strings: false,
72 _non_exhaustive: (),
75};
76
77impl DynamicCheckRunnerContext {
78 pub fn new(starlark_program: String) -> Result<Self, DynamicCheckRunnerError> {
88 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 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
210pub(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
235pub(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}