sui_config/
dynamic_transaction_signing_checks.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
// Copyright (c) Mysten Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

use serde::{
    de::{self, Deserializer},
    ser::Serializer,
    Deserialize,
};
use serde_json::Value as JsonValue;
use starlark::{
    environment::{Globals, Module},
    eval::Evaluator,
    syntax::{AstModule, Dialect, DialectTypes},
    values::{dict::AllocDict, AllocValue, Heap, Value},
    ErrorKind,
};
use sui_types::{
    base_types::ObjectRef,
    signature::GenericSignature,
    transaction::{InputObjectKind, TransactionData},
};
use tracing::warn;

/// The name of the input variables that the transaction will have in the Starlark program.
const TX_DATA_NAME: &str = "tx_data";
const TX_SIGNERS_NAME: &str = "tx_signers";
const TX_INPUT_OBJECTS_NAME: &str = "tx_input_objects";
const TX_RECEIVING_OBJECTS_NAME: &str = "tx_receiving_objects";
const TX_DIGEST_NAME: &str = "tx_digest";

/// The dummy name of the Starlark file being executed. We will just be passing the string of the
/// program in directly so this is not important but may appear in error messages.
const STAR_INPUT_FILE_NAME: &str = "dynamic_transaction_signing_checks.star";

#[derive(Debug, thiserror::Error)]
pub enum DynamicCheckRunnerError {
    #[error("Failed to serialize transaction data to JSON: {0}")]
    JSONSerializationError(String),
    #[error("Failed to parse Starlark program value -- unsupported number type {0}")]
    UnsupportedNumberFormat(String),
    #[error(
        "Failed to execute Starlark program -- invalid return type expected a bool but got {0}"
    )]
    InvalidReturnType(String),
    #[error("Failed to execute Starlark program: {0}")]
    ExecutionError(ErrorKind),
    #[error("Failed to load Starlark program: {0}")]
    LoadingError(ErrorKind),
    #[error("Check failed -- transaction denied")]
    CheckFailure,
}

#[derive(Debug, Clone)]
pub struct DynamicCheckRunnerContext {
    module: AstModule,
    globals: Globals,
    loaded_program: String,
}

const DIALECT: Dialect = Dialect {
    enable_def: true,
    enable_lambda: true,
    enable_keyword_only_arguments: false,
    enable_positional_only_arguments: false,
    enable_types: DialectTypes::Disable,
    // NB: set loader to false to prevent any external loading
    enable_load: false,
    enable_load_reexport: false,
    // NB: Allow for top level statements to be used (e.g., top-level `for`, `if`, etc.)
    enable_top_level_stmt: true,
    enable_f_strings: false,
    // NB: We explicitly fully initalize the struct to prevent any future changes to the dialect
    // without us noticing and deciding whether or not the new feature should be enabled.
    _non_exhaustive: (),
};

impl DynamicCheckRunnerContext {
    /// Create a new `DynamicCheckRunnerContext` with the given Starlark program
    /// `starlark_program` string. This will parse and validate the program is syntactically
    /// correct and will set up shared (immutable) state that can be reused. This will not
    /// run the program.
    ///
    /// The `starlark_program` string should be a valid Starlark program that returns a boolean
    /// value when run -- `True` in the case that the transaction should be allowed, or `False` if
    /// the transaction should be denied. Any other return value other than `True` (including
    /// errors) should be considered a denial.
    pub fn new(starlark_program: String) -> Result<Self, DynamicCheckRunnerError> {
        // Adds global functions and variables to the dialect (e.g., True, False, Maps, Lists, etc.)
        // The full spec of what exactly is added here can be found here:
        // https://github.com/bazelbuild/starlark/blob/master/spec.md#built-in-constants-and-functions
        let globals = Globals::standard();
        warn!(
            "Dynamic transaction checks are enabled. Make sure that you intend to be running \
            dynamic checks on transactions."
        );
        let module = AstModule::parse(STAR_INPUT_FILE_NAME, starlark_program.clone(), &DIALECT)
            .map_err(|e| DynamicCheckRunnerError::LoadingError(e.into_kind()))?;
        Ok(Self {
            module,
            globals,
            loaded_program: starlark_program,
        })
    }

    /// Run the Starlark program in `self` with the given transaction data, signatures, input
    /// object kinds, and receiving objects.
    pub fn run_predicate(
        &self,
        tx_data: &TransactionData,
        tx_signatures: &[GenericSignature],
        input_object_kinds: &[InputObjectKind],
        receiving_objects: &[ObjectRef],
    ) -> Result<(), DynamicCheckRunnerError> {
        let tx_data_json = serde_json::to_value(tx_data)
            .map_err(|e| DynamicCheckRunnerError::JSONSerializationError(e.to_string()))?;
        let tx_signatures_json = serde_json::to_value(tx_signatures)
            .map_err(|e| DynamicCheckRunnerError::JSONSerializationError(e.to_string()))?;
        let input_object_kinds_json = serde_json::to_value(input_object_kinds)
            .map_err(|e| DynamicCheckRunnerError::JSONSerializationError(e.to_string()))?;
        let receiving_objects_json = serde_json::to_value(receiving_objects)
            .map_err(|e| DynamicCheckRunnerError::JSONSerializationError(e.to_string()))?;
        let digest_json = serde_json::to_value(tx_data.digest())
            .map_err(|e| DynamicCheckRunnerError::JSONSerializationError(e.to_string()))?;

        self.run_starlark_predicate(
            &tx_data_json,
            &tx_signatures_json,
            &input_object_kinds_json,
            &receiving_objects_json,
            &digest_json,
        )
    }

    fn run_starlark_predicate(
        &self,
        tx_data: &JsonValue,
        tx_signatures: &JsonValue,
        tx_input_object_kinds: &JsonValue,
        tx_receiving_objects: &JsonValue,
        tx_digest: &JsonValue,
    ) -> Result<(), DynamicCheckRunnerError> {
        let heap = Heap::new();
        let env = Module::new();

        let tx_data_value = Self::json_to_starlark(tx_data, &heap)?;
        let tx_signers_value = Self::json_to_starlark(tx_signatures, &heap)?;
        let tx_input_object_kinds_value = Self::json_to_starlark(tx_input_object_kinds, &heap)?;
        let tx_receiving_objects_value = Self::json_to_starlark(tx_receiving_objects, &heap)?;
        let tx_digest_value = Self::json_to_starlark(tx_digest, &heap)?;

        env.set(TX_DATA_NAME, tx_data_value);
        env.set(TX_SIGNERS_NAME, tx_signers_value);
        env.set(TX_INPUT_OBJECTS_NAME, tx_input_object_kinds_value);
        env.set(TX_RECEIVING_OBJECTS_NAME, tx_receiving_objects_value);
        env.set(TX_DIGEST_NAME, tx_digest_value);

        let mut evaluator = Evaluator::new(&env);
        let output_value = evaluator
            .eval_module(self.module.clone(), &self.globals)
            .map_err(|e| DynamicCheckRunnerError::ExecutionError(e.into_kind()))?;
        let transaction_allowed = output_value
            .unpack_bool()
            .ok_or_else(|| DynamicCheckRunnerError::InvalidReturnType(output_value.to_repr()))?;
        if transaction_allowed {
            Ok(())
        } else {
            Err(DynamicCheckRunnerError::CheckFailure)
        }
    }

    fn json_to_starlark<'v>(
        value: &JsonValue,
        heap: &'v Heap,
    ) -> Result<Value<'v>, DynamicCheckRunnerError> {
        Ok(match value {
            JsonValue::Null => Value::new_none(),
            JsonValue::Bool(b) => Value::new_bool(*b),
            JsonValue::Number(n) => {
                if let Some(i) = n.as_u64() {
                    heap.alloc(i)
                } else {
                    return Err(DynamicCheckRunnerError::UnsupportedNumberFormat(
                        n.to_string(),
                    ));
                }
            }
            JsonValue::String(s) => heap.alloc(s.as_str()),
            JsonValue::Array(arr) => {
                let list: Vec<_> = arr
                    .iter()
                    .map(|v| Self::json_to_starlark(v, heap))
                    .collect::<Result<_, _>>()?;
                list.alloc_value(heap)
            }
            JsonValue::Object(obj) => {
                let kvs: Vec<_> = obj
                    .iter()
                    .map(|(k, v)| {
                        let key = heap.alloc(k.as_str());
                        let val = Self::json_to_starlark(v, heap)?;
                        Ok((key, val))
                    })
                    .collect::<Result<_, _>>()?;
                heap.alloc(AllocDict(kvs))
            }
        })
    }
}

// Custom serialization/deserialization for the `DynamicCheckRunnerContext` struct. This allows us
// to validate the program at the time that it is deserialized, rather than at the time that it is
// first used. This provides better error message locality and allows us to fail fast if the
// program is invalid. We keep the invariant here that `serialize(deserialize(program)) ==
// program`.

/// Deserialize a `DynamicCheckRunnerContext` from a string. This will parse the string as a
/// Starlark program and validate that it is syntactically correct and setup the
/// `DynamicCheckRunnerContext` for it. If the program is syntactically invalid, an error will be
/// returned.
pub(crate) fn deserialize_dynamic_transaction_checks<'de, D>(
    deserializer: D,
) -> Result<Option<DynamicCheckRunnerContext>, D::Error>
where
    D: Deserializer<'de>,
{
    let path_opt: Option<String> = Option::deserialize(deserializer)?;
    match path_opt {
        Some(p) => Ok(Some(
            DynamicCheckRunnerContext::new(p).map_err(de::Error::custom)?,
        )),
        None => Ok(None),
    }
}

/// Takes a `DynamicCheckRunnerContext` and serializes the original source program as the returned
/// string. No parsed state or otherwise is serialized.
pub(crate) fn serialize_dynamic_transaction_checks<S>(
    value: &Option<DynamicCheckRunnerContext>,
    serializer: S,
) -> Result<S::Ok, S::Error>
where
    S: Serializer,
{
    match value {
        Some(DynamicCheckRunnerContext { loaded_program, .. }) => {
            serializer.serialize_some(&loaded_program)
        }
        None => serializer.serialize_none(),
    }
}

#[cfg(test)]
mod test {
    #[test]
    fn parse_on_load_invalid() {
        let program = r#"
            def main(): return 1
        "#;
        let result = super::DynamicCheckRunnerContext::new(program.to_string());
        assert!(result.is_err());
    }

    #[test]
    fn parse_on_load_valid() {
        let program = r#"
def main(): 
    return 1
        "#;
        let result = super::DynamicCheckRunnerContext::new(program.to_string());
        assert!(result.is_ok());
    }
}