Skip to main content

stepflow_flow/values/
redacted_value.rs

1// Copyright 2025 DataStax Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
4// in compliance with the License. You may obtain a copy of the License at
5//
6//     http://www.apache.org/licenses/LICENSE-2.0
7//
8// Unless required by applicable law or agreed to in writing, software distributed under the License
9// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
10// or implied. See the License for the specific language governing permissions and limitations under
11// the License.
12
13use std::{collections::HashMap, fmt};
14
15use crate::schema::SchemaRef;
16
17#[derive(Debug, PartialEq, Clone)]
18pub struct Secrets {
19    is_secret: bool,
20    fields: Option<HashMap<String, Secrets>>,
21}
22
23impl Default for Secrets {
24    fn default() -> Self {
25        Self::new()
26    }
27}
28
29impl Secrets {
30    const EMPTY: Secrets = Secrets {
31        is_secret: false,
32        fields: None,
33    };
34
35    pub const fn empty() -> &'static Self {
36        &Self::EMPTY
37    }
38
39    pub fn new() -> Self {
40        Self {
41            is_secret: false,
42            fields: None,
43        }
44    }
45
46    /// Return true if the current value is marked as a secret.
47    pub fn is_secret(&self) -> bool {
48        self.is_secret
49    }
50
51    /// Get secret information about the given field, if any.
52    pub fn field(&'_ self, field_name: &str) -> &'_ Self {
53        if let Some(fields) = &self.fields {
54            fields.get(field_name).unwrap_or(&Self::EMPTY)
55        } else {
56            &Self::EMPTY
57        }
58    }
59
60    fn reduce(mut self) -> Option<Self> {
61        // Assumption: All fields have already been reduced
62
63        // First, simplify the fields. We only need to keep things that either
64        // mark secrets or have nested secret fields.
65        self.fields = self.fields.and_then(|mut fields| {
66            fields.retain(|_, v| v.is_secret || v.fields.is_some());
67            if fields.is_empty() {
68                None
69            } else {
70                Some(fields)
71            }
72        });
73
74        // Second, simplify this node. We need it if it's marked as a secret
75        // or if it has any fields (which are reduced and must lead to a secret).
76        if self.is_secret || self.fields.is_some() {
77            Some(self)
78        } else {
79            None
80        }
81    }
82
83    pub fn add_field(&mut self, field_name: &str, secrets: Secrets) {
84        let Some(secrets) = secrets.reduce() else {
85            // No secrets in this subtree, nothing to add
86            return;
87        };
88
89        if self.fields.is_none() {
90            self.fields = Some(HashMap::new());
91        }
92        let fields = self.fields.as_mut().unwrap();
93        fields.insert(field_name.to_string(), secrets);
94    }
95
96    pub fn add_secret_field(&mut self, field_name: &str) {
97        self.add_field(
98            field_name,
99            Secrets {
100                is_secret: true,
101                fields: None,
102            },
103        );
104    }
105
106    pub fn from_schema(schema: &serde_json::Value) -> Secrets {
107        let schema = schema.as_object().expect("Schema must be an object");
108        let type_name = schema.get("type").and_then(|s| s.as_str());
109
110        let mut secrets = match type_name {
111            Some("object") => {
112                let mut secrets = Secrets::new();
113                if let Some(properties) = schema.get("properties").and_then(|s| s.as_object()) {
114                    for (prop_name, prop_schema) in properties {
115                        let mut prop_secrets = Self::from_schema(prop_schema);
116                        prop_secrets.is_secret = prop_schema
117                            .get("is_secret")
118                            .and_then(|s| s.as_bool())
119                            .unwrap_or(false);
120                        secrets.add_field(prop_name, prop_secrets);
121                    }
122                }
123                secrets
124            }
125            Some("array") => {
126                if let Some(items) = schema.get("items") {
127                    Secrets::from_schema(items)
128                } else {
129                    Secrets::new()
130                }
131            }
132            _ => {
133                // If it's not an object or array, it can't have nested secrets.
134                Secrets::new()
135            }
136        };
137
138        secrets.is_secret = schema
139            .get("is_secret")
140            .and_then(|s| s.as_bool())
141            .unwrap_or(false);
142
143        secrets
144    }
145
146    pub fn redacted<'a>(&'a self, value: &'a serde_json::Value) -> RedactedValue<'a> {
147        RedactedValue::new(value, self)
148    }
149}
150
151impl From<&SchemaRef> for Secrets {
152    fn from(value: &SchemaRef) -> Self {
153        let secrets = Secrets::from_schema(value.as_value());
154        secrets.reduce().unwrap_or_default()
155    }
156}
157
158/// A wrapper for ValueRef that implements safe Display formatting with secret redaction.
159///
160/// This ensures that values marked as secrets in the schema are redacted when printed,
161/// preventing accidental exposure of sensitive data in logs, error messages, and debug output.
162pub struct RedactedValue<'a>(&'a serde_json::Value, &'a Secrets);
163
164impl<'a> RedactedValue<'a> {
165    pub fn new(value: &'a serde_json::Value, secrets: &'a Secrets) -> Self {
166        Self(value, secrets)
167    }
168}
169
170impl<'a> fmt::Display for RedactedValue<'a> {
171    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
172        write!(f, "{:?}", self)
173    }
174}
175
176impl<'a> fmt::Debug for RedactedValue<'a> {
177    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
178        let value = self.0;
179        let secrets = self.1;
180
181        match value {
182            _ if secrets.is_secret => {
183                write!(f, "\"[REDACTED]\"")
184            }
185            serde_json::Value::Object(obj) => {
186                let mut f = f.debug_map();
187                for (key, val) in obj {
188                    let secrets = secrets.field(key);
189                    f.entry(key, &RedactedValue(val, secrets));
190                }
191                f.finish()
192            }
193            serde_json::Value::Array(arr) => {
194                let mut f = f.debug_list();
195                let secrets = secrets.field("$item");
196                for val in arr {
197                    f.entry(&RedactedValue(val, secrets));
198                }
199                f.finish()
200            }
201            _ => {
202                // For primitive values, just write them directly
203                write!(f, "{}", value)
204            }
205        }
206    }
207}
208
209#[cfg(test)]
210mod tests {
211    use crate::values::ValueRef;
212
213    use super::*;
214    use serde_json::json;
215
216    #[test]
217    fn test_secrets_from_schema() {
218        let secrets = Secrets::from_schema(&json!({
219            "type": "object",
220            "properties": {
221                "username": {
222                    "type": "string"
223                },
224                "password": {
225                    "type": "string",
226                    "is_secret": true
227                },
228                "api_key": {
229                    "type": "string",
230                    "is_secret": true
231                },
232                "config": {
233                    "type": "object",
234                    "properties": {
235                        "timeout": {
236                            "type": "number"
237                        },
238                        "secret_token": {
239                            "type": "string",
240                            "is_secret": true
241                        }
242                    }
243                },
244                "items": {
245                    "type": "array",
246                    "items": {
247                        "type": "object",
248                        "properties": {
249                            "id": {"type": "string"},
250                            "secret_data": {"type": "string", "is_secret": true}
251                        }
252                    }
253                }
254            }
255        }));
256
257        assert!(!secrets.is_secret);
258        assert!(!secrets.field("username").is_secret());
259        assert!(secrets.field("password").is_secret());
260        assert!(secrets.field("api_key").is_secret());
261
262        let config_secrets = secrets.field("config");
263        assert!(!config_secrets.is_secret());
264        assert!(!config_secrets.field("timeout").is_secret());
265        assert!(config_secrets.field("secret_token").is_secret());
266
267        let items_secrets = secrets.field("items");
268        assert!(!items_secrets.is_secret());
269        assert!(!items_secrets.field("id").is_secret());
270        assert!(items_secrets.field("secret_data").is_secret());
271    }
272
273    #[test]
274    fn test_secret_array_from_schema() {
275        let secrets = Secrets::from_schema(&json!({
276            "type": "object",
277            "properties": {
278                "items": {
279                    "type": "array",
280                    "items": {
281                        "type": "object",
282                        "properties": {
283                            "id": {"type": "string"},
284                            "secret_data": {"type": "string", "is_secret": true}
285                        }
286                    },
287                    "is_secret": true,
288                }
289            }
290        }));
291
292        assert!(!secrets.is_secret);
293        let items_secrets = secrets.field("items");
294        assert!(items_secrets.is_secret());
295        assert!(!items_secrets.field("id").is_secret());
296        assert!(items_secrets.field("secret_data").is_secret());
297    }
298
299    #[test]
300    fn test_secrets_from_primitive_secret() {
301        let secret_schema = SchemaRef::parse_json(
302            r#"{
303            "type": "string",
304            "is_secret": true
305        }"#,
306        )
307        .unwrap();
308
309        let secrets = Secrets::from(&secret_schema);
310        assert!(secrets.is_secret());
311        assert!(secrets.fields.is_none());
312    }
313
314    fn create_test_secrets() -> Secrets {
315        let mut secrets = Secrets::new();
316        secrets.add_secret_field("password");
317        secrets.add_secret_field("api_key");
318
319        let mut config_secrets = Secrets::new();
320        config_secrets.add_secret_field("secret_token");
321        secrets.add_field("config", config_secrets);
322
323        let mut item_secrets = Secrets::new();
324        item_secrets.add_secret_field("secret_data");
325        let mut items_secrets = Secrets::new();
326        items_secrets.add_field("$item", item_secrets);
327        secrets.add_field("items", items_secrets);
328
329        secrets
330    }
331
332    #[test]
333    fn test_redact_secrets_in_object() {
334        let value_json = json!({
335            "username": "alice",
336            "password": "secret123",
337            "api_key": "sk-abcd1234"
338        });
339
340        let value = ValueRef::new(value_json);
341
342        insta::assert_snapshot!(value.redacted(&create_test_secrets()), @r#"{"username": "alice", "password": "[REDACTED]", "api_key": "[REDACTED]"}"#);
343    }
344
345    #[test]
346    fn test_redacted_nested_secrets() {
347        let value_json = json!({
348            "username": "alice",
349            "config": {
350                "timeout": 30,
351                "secret_token": "token123"
352            }
353        });
354
355        let value = ValueRef::new(value_json);
356        insta::assert_snapshot!(value.redacted(&create_test_secrets()), @r#"{"username": "alice", "config": {"timeout": 30, "secret_token": "[REDACTED]"}}"#);
357    }
358
359    #[test]
360    fn test_redacted_array_with_secret_fields() {
361        let value_json = json!({
362            "items": [
363                {"id": "1", "secret_data": "data1"},
364                {"id": "2", "secret_data": "data2"}
365            ]
366        });
367
368        let value = ValueRef::new(value_json);
369        insta::assert_snapshot!(value.redacted(&create_test_secrets()), @r#"{"items": [{"id": "1", "secret_data": "[REDACTED]"}, {"id": "2", "secret_data": "[REDACTED]"}]}"#);
370    }
371
372    #[test]
373    fn test_redacted_secret_array() {
374        let mut secrets = Secrets::new();
375        secrets.add_secret_field("items");
376
377        let value_json = json!({
378            "items": [
379                {"id": "1", "secret_data": "data1"},
380                {"id": "2", "secret_data": "data2"}
381            ]
382        });
383
384        let value = ValueRef::new(value_json);
385        insta::assert_snapshot!(value.redacted(&secrets), @r#"{"items": "[REDACTED]"}"#);
386    }
387
388    #[test]
389    fn test_no_schema_shows_all_values() {
390        let value_json = json!({
391            "password": "secret123",
392            "username": "alice"
393        });
394
395        let value = ValueRef::new(value_json);
396        insta::assert_snapshot!(value.redacted(&Secrets::EMPTY), @r#"{"password": "secret123", "username": "alice"}"#);
397    }
398
399    #[test]
400    fn test_redacted_primitive_secret() {
401        // Test a primitive value that's marked as secret
402        let secrets = Secrets {
403            is_secret: true,
404            fields: None,
405        };
406
407        let value = ValueRef::new(json!("secret_token"));
408        insta::assert_snapshot!(value.redacted(&secrets), @r#""[REDACTED]""#);
409    }
410}