1use crate::{SensitiveTextRedactor, truncate_progress_message, v1};
6use std::sync::{
7 Arc, Mutex,
8 atomic::{AtomicU32, Ordering},
9};
10use std::time::{Duration, Instant};
11
12pub const DEFAULT_PROGRESS_CHUNK_BYTES: usize = 32 * 1024;
15const TIMED_FLUSH_INTERVAL: Duration = Duration::from_secs(2);
16#[derive(Clone)]
19pub struct ProgressReporter {
20 milestone: Arc<dyn Fn(String) + Send + Sync>,
21 detail: Arc<dyn Fn(String) + Send + Sync>,
22}
23
24impl ProgressReporter {
25 pub fn new(
26 milestone: impl Fn(String) + Send + Sync + 'static,
27 detail: impl Fn(String) + Send + Sync + 'static,
28 ) -> Self {
29 Self { milestone: Arc::new(milestone), detail: Arc::new(detail) }
30 }
31
32 pub fn report(&self, message: impl Into<String>) {
33 (self.milestone)(message.into());
34 }
35
36 pub fn detail(&self, message: impl Into<String>) {
37 (self.detail)(message.into());
38 }
39}
40
41#[derive(Clone)]
42pub struct EventSequencer {
43 action_id: Arc<str>,
44 next: Arc<AtomicU32>,
45 redactor: Arc<Mutex<SensitiveTextRedactor>>,
46}
47
48impl EventSequencer {
49 pub fn new(action_id: impl Into<Arc<str>>) -> Self {
50 Self::with_sensitive_values(action_id, std::iter::empty())
51 }
52
53 pub fn with_sensitive_values<'a>(
54 action_id: impl Into<Arc<str>>,
55 values: impl IntoIterator<Item = &'a [u8]>,
56 ) -> Self {
57 Self {
58 action_id: action_id.into(),
59 next: Arc::new(AtomicU32::new(1)),
60 redactor: Arc::new(Mutex::new(SensitiveTextRedactor::with_sensitive_values(values))),
61 }
62 }
63
64 pub fn progress(&self, message: &str) -> v1::ApplyResponse {
65 self.event(v1::EventPhase::Progress, message, false, None)
66 }
67
68 pub fn detail(&self, message: &str) -> v1::ApplyResponse {
69 self.event(v1::EventPhase::Progress, message, true, None)
70 }
71
72 pub fn completed(&self, message: &str) -> v1::ApplyResponse {
73 self.event(v1::EventPhase::Completed, message, false, None)
74 }
75
76 pub fn failed(&self, message: &str, error: v1::ProviderError) -> v1::ApplyResponse {
77 self.event(v1::EventPhase::Failed, message, false, Some(error))
78 }
79
80 fn event(
81 &self,
82 phase: v1::EventPhase,
83 message: &str,
84 detail: bool,
85 error: Option<v1::ProviderError>,
86 ) -> v1::ApplyResponse {
87 let redacted = self.redactor.lock().expect("provider event redactor lock").redact(message);
88 let safe_message = truncate_progress_message(&redacted);
89 v1::ApplyResponse {
90 sequence: self.next.fetch_add(1, Ordering::Relaxed),
91 phase: phase.into(),
92 action_id: self.action_id.to_string(),
93 safe_message,
94 error,
95 detail,
96 confidential_outputs: Default::default(),
97 public_outputs: Default::default(),
98 }
99 }
100}
101
102#[derive(Clone)]
103pub struct ProgressBatcher {
104 label: Arc<str>,
105 reporter: ProgressReporter,
106 state: Arc<Mutex<BatchState>>,
107 chunk_bytes: usize,
108}
109
110struct BatchState {
111 text: String,
112 last_emit: Instant,
113 label_emitted: bool,
114 redactor: SensitiveTextRedactor,
115}
116
117impl Default for BatchState {
118 fn default() -> Self {
119 Self {
120 text: String::new(),
121 last_emit: Instant::now(),
122 label_emitted: false,
123 redactor: SensitiveTextRedactor::default(),
124 }
125 }
126}
127
128impl ProgressBatcher {
129 pub fn new(label: impl Into<Arc<str>>, reporter: ProgressReporter) -> Self {
130 Self {
131 label: label.into(),
132 reporter,
133 state: Arc::new(Mutex::new(BatchState::default())),
134 chunk_bytes: DEFAULT_PROGRESS_CHUNK_BYTES,
135 }
136 }
137
138 pub fn push_line(&self, line: &str) {
139 let mut state = self.state.lock().expect("provider progress batch lock");
140 let safe = state.redactor.redact(line.trim_end_matches(['\r', '\n']));
141 state.text.push_str(&safe);
142 state.text.push('\n');
143 let timed = state.last_emit.elapsed() >= TIMED_FLUSH_INTERVAL;
144 if state.text.len() >= self.chunk_bytes || timed {
145 state.last_emit = Instant::now();
146 self.emit(&mut state);
147 }
148 }
149
150 pub fn finish(&self) {
151 let mut state = self.state.lock().expect("provider progress batch lock");
152 self.emit(&mut state);
153 }
154
155 fn emit(&self, state: &mut BatchState) {
156 let chunk = std::mem::take(&mut state.text);
157 let detail = chunk.trim_end();
158 if !detail.is_empty() {
159 let message = if state.label_emitted {
160 detail.to_string()
161 } else {
162 state.label_emitted = true;
163 format!("{}:\n{detail}", self.label)
164 };
165 self.reporter.detail(truncate_progress_message(&message));
167 }
168 }
169}
170
171#[cfg(test)]
172mod tests {
173 use super::*;
174
175 #[test]
176 fn sequencer_orders_and_sanitizes_events() {
177 let events =
178 EventSequencer::with_sensitive_values("action/test", [b"opaque-provider-value".as_slice()]);
179 let progress = events.progress("token=not-safe opaque-provider-value");
180 let detail = events.detail("ordinary command detail");
181 let completed = events.completed("ordinary completion");
182 assert_eq!(progress.sequence, 1);
183 assert_eq!(progress.safe_message, "token=[REDACTED] [REDACTED]");
184 assert!(detail.detail);
185 assert!(!progress.detail);
186 assert_eq!(completed.sequence, 3);
187 assert_eq!(completed.action_id, "action/test");
188 }
189
190 #[test]
191 fn batcher_preserves_full_normal_output_and_redacts_values() {
192 let messages = Arc::new(Mutex::new(Vec::new()));
193 let captured = messages.clone();
194 let batcher = ProgressBatcher::new(
195 "Nix copy",
196 ProgressReporter::new(
197 |_| {},
198 move |message| {
199 captured.lock().unwrap().push(message);
200 },
201 ),
202 );
203 for index in 0..100 {
204 batcher.push_line(&format!("copying /nix/store/path-{index}"));
205 }
206 batcher.push_line("copying /nix/store/final from https://user:[email protected]");
207 batcher.push_line("token=not-safe");
208 batcher.finish();
209 let joined = messages.lock().unwrap().join("\n");
210 assert!(joined.contains("copying /nix/store/path-0"));
211 assert!(joined.contains("copying /nix/store/path-99"));
212 assert!(joined.contains("copying /nix/store/final"));
213 assert!(joined.contains("https://[REDACTED]@example.test"));
214 assert!(joined.contains("token=[REDACTED]"));
215 assert!(!joined.contains("user:pass"));
216 assert!(!joined.contains("not-safe"));
217 }
218
219 #[test]
220 fn batcher_labels_only_the_first_bounded_chunk() {
221 let messages = Arc::new(Mutex::new(Vec::new()));
222 let captured = messages.clone();
223 let mut batcher = ProgressBatcher::new(
224 "Nix activation",
225 ProgressReporter::new(|_| {}, move |message| captured.lock().unwrap().push(message)),
226 );
227 batcher.chunk_bytes = 8;
228 for line in ["first", "second", "third"] {
229 batcher.push_line(line);
230 }
231 batcher.finish();
232 let joined = messages.lock().unwrap().join("\n");
233 assert_eq!(joined.matches("Nix activation:").count(), 1);
234 for line in ["first", "second", "third"] {
235 assert!(joined.contains(line));
236 }
237 }
238}