Skip to main content

buntrail_engine/environment/
container_picker.rs

1//! Picker work uses the connector's admission, worker retirement and authentication broker.
2use super::*;
3use buntrail_services::connections::{
4    container_engine::{PinnedEngine, StartOutcome},
5    container_identity::ContainerDescriptor,
6};
7
8#[derive(Clone)]
9pub struct ContainerInventory {
10    pub locator: EnvironmentLocator,
11    pub host_binding: Vec<u8>,
12    pub pinned: PinnedEngine,
13    pub containers: Vec<ContainerDescriptor>,
14    pub start_uncertain: bool,
15    pub started: bool,
16    pub message: String,
17}
18#[derive(Clone, PartialEq, Eq)]
19pub(super) struct SelectedContainer {
20    pub host_binding: Vec<u8>,
21    pub pinned: PinnedEngine,
22    pub descriptor: ContainerDescriptor,
23}
24impl SelectedContainer {
25    pub(super) fn matches_existing(
26        &self,
27        target: &TargetIdentity,
28        runtime: &RuntimeGeneration,
29    ) -> bool {
30        self.descriptor.state == ContainerState::Running
31            && self
32                .descriptor
33                .actual_uid
34                .is_none_or(|uid| target.effective_user == uid.to_be_bytes())
35            && target.security_binding == self.host_binding
36            && target.container.as_ref().is_some_and(|container| {
37                container.engine_endpoint == self.descriptor.engine_endpoint
38                    && container.immutable_id == self.descriptor.immutable_id.as_bytes()
39            })
40            && runtime.container_start == self.descriptor.start_generation
41    }
42}
43impl ContainerInventory {
44    pub(super) fn selected(&self, index: usize) -> Result<SelectedContainer, CoreError> {
45        if self.start_uncertain {
46            return Err(unavailable(
47                "Original Start result is uncertain; inspect this exact container ID before connecting",
48            ));
49        }
50        Ok(SelectedContainer {
51            host_binding: self.host_binding.clone(),
52            pinned: self.pinned.clone(),
53            descriptor: self
54                .containers
55                .get(index)
56                .ok_or_else(|| {
57                    CoreError::new(ErrorCode::InvalidArgument, "Choose an existing container")
58                })?
59                .clone(),
60        })
61    }
62}
63pub(super) struct QueryState {
64    id: u64,
65    cancel: Arc<AtomicBool>,
66    done: Arc<Done>,
67    result: Mutex<Option<Result<Arc<ContainerInventory>, CoreError>>>,
68    broker: Mutex<Option<Arc<AuthenticationBroker>>>,
69    prompt: Mutex<Option<AuthenticationRequest>>,
70}
71impl QueryState {
72    pub(super) fn cancel(&self) {
73        self.cancel.store(true, Ordering::SeqCst);
74        if let Some(b) = self.broker.lock().unwrap().as_ref() {
75            b.cancel();
76        }
77    }
78}
79pub struct ContainerQuery {
80    state: Arc<QueryState>,
81    _permit: Arc<ConnectionPermit>,
82}
83impl Drop for ContainerQuery {
84    fn drop(&mut self) {
85        self.cancel();
86    }
87}
88impl ContainerQuery {
89    pub fn cancel(&self) {
90        self.state.cancel();
91    }
92    pub fn result(&self) -> Option<Result<Arc<ContainerInventory>, CoreError>> {
93        self.state.result.lock().unwrap().clone()
94    }
95    pub fn finished(&self) -> bool {
96        self.state.done.finished.load(Ordering::SeqCst)
97    }
98    pub fn take_authentication(&self) -> Option<ConnectionPrompt> {
99        if self.state.cancel.load(Ordering::SeqCst) {
100            return None;
101        }
102        let mut prompt = self.state.prompt.lock().unwrap();
103        if prompt.is_some() {
104            return None;
105        }
106        let request = self.state.broker.lock().unwrap().as_ref()?.take_request()?;
107        let copied = AuthenticationRequest {
108            id: request.id,
109            kind: request.kind.clone(),
110            prompt: request.prompt.clone(),
111        };
112        *prompt = Some(request);
113        Some(ConnectionPrompt {
114            attempt: self.state.id,
115            request: copied,
116        })
117    }
118    pub fn respond(&self, attempt: u64, prompt: u64, response: AuthenticationResponse) -> bool {
119        let mut pending = self.state.prompt.lock().unwrap();
120        if self.state.cancel.load(Ordering::SeqCst)
121            || attempt != self.state.id
122            || pending.as_ref().is_none_or(|p| p.id != prompt)
123        {
124            return false;
125        }
126        let accepted = self
127            .state
128            .broker
129            .lock()
130            .unwrap()
131            .as_ref()
132            .is_some_and(|b| b.respond(prompt, response));
133        if accepted {
134            *pending = None;
135        }
136        accepted
137    }
138}
139impl EnvironmentConnector {
140    pub fn container_inventory(
141        &self,
142        locator: EnvironmentLocator,
143    ) -> Result<ContainerQuery, CoreError> {
144        self.container_query(locator, None)
145    }
146    pub fn start_container(
147        &self,
148        inventory: Arc<ContainerInventory>,
149        index: usize,
150    ) -> Result<ContainerQuery, CoreError> {
151        let selected = inventory.selected(index)?;
152        self.container_query(inventory.locator.clone(), Some(selected))
153    }
154    fn container_query(
155        &self,
156        locator: EnvironmentLocator,
157        start: Option<SelectedContainer>,
158    ) -> Result<ContainerQuery, CoreError> {
159        let valid = match locator.route {
160            RouteKind::LocalContainer => locator.ssh_profile.is_none(),
161            RouteKind::SshContainer => locator.ssh_profile.as_ref().is_some_and(|p| !p.is_empty()),
162            _ => false,
163        };
164        if !valid || locator.engine.is_none() {
165            return Err(CoreError::new(
166                ErrorCode::InvalidArgument,
167                "Choose a container engine and its local or saved SSH route",
168            ));
169        }
170        let permit = self.budget.acquire(false)?;
171        static NEXT: AtomicU64 = AtomicU64::new(1);
172        let state = Arc::new(QueryState {
173            id: NEXT.fetch_add(1, Ordering::SeqCst),
174            cancel: Arc::new(AtomicBool::new(false)),
175            done: Arc::new(Done::default()),
176            result: Mutex::new(None),
177            broker: Mutex::new(None),
178            prompt: Mutex::new(None),
179        });
180        let mut registry = self.state.lock().unwrap();
181        if registry.closed {
182            return Err(CoreError::new(
183                ErrorCode::Closed,
184                "Connector is shutting down",
185            ));
186        }
187        registry.queries.retain(|q| q.strong_count() != 0);
188        if registry.queries.len() >= MAX_CONNECTIONS {
189            return Err(CoreError::new(
190                ErrorCode::Busy,
191                "Container picker capacity reached",
192            ));
193        }
194        registry.queries.push(Arc::downgrade(&state));
195        drop(registry);
196        let configuration = self.configuration.clone();
197        let worker = state.clone();
198        let worker_permit = permit.clone();
199        spawn(&self.workers, move || {
200            let _permit = worker_permit;
201            let result = (|| {
202                let broker = if locator.route == RouteKind::SshContainer {
203                    Some(AuthenticationBroker::new(configuration.askpass.clone())?)
204                } else {
205                    None
206                };
207                *worker.broker.lock().unwrap() = broker.clone();
208                let outcome = (|| {
209                    let host = if locator.route == RouteKind::SshContainer {
210                        EngineHost::ssh(SshConnector::new(
211                            SshProfile {
212                                name: locator.ssh_profile.clone().unwrap(),
213                                config: configuration.config.clone(),
214                            },
215                            broker.clone(),
216                            worker.cancel.clone(),
217                        )?)?
218                    } else {
219                        EngineHost::local(worker.cancel.clone())?
220                    };
221                    let client = ContainerEngineClient::new(
222                        host,
223                        locator.engine.clone().unwrap(),
224                        locator.engine_context.clone(),
225                        locator.effective_user.clone(),
226                    )?;
227                    let mut inventory = ContainerInventory {
228                        locator: locator.clone(),
229                        host_binding: client.host.binding().to_vec(),
230                        pinned: client.pinned.clone(),
231                        containers: vec![],
232                        start_uncertain: false,
233                        started: false,
234                        message: String::new(),
235                    };
236                    if let Some(selected) = start {
237                        if inventory.host_binding != selected.host_binding
238                            || inventory.pinned != selected.pinned
239                        {
240                            return Err(CoreError::new(
241                                ErrorCode::TargetChanged,
242                                "Chosen engine endpoint or host changed before Start",
243                            ));
244                        }
245                        match client.start(&selected.descriptor) {
246                            StartOutcome::Started(v) => {
247                                inventory.started = true;
248                                inventory.containers.push(v);
249                                inventory.message =
250                                    "Container started; verify the environment and choose a folder"
251                                        .into();
252                            }
253                            StartOutcome::StartedUnsupported { selected, error } => {
254                                inventory.started = true;
255                                inventory.containers.push(selected);
256                                inventory.message = format!(
257                                    "Container started, but environment verification is unavailable: {}",
258                                    error.message
259                                );
260                            }
261                            StartOutcome::Failed(error) => return Err(error),
262                            StartOutcome::Uncertain { selected, error } => {
263                                inventory.containers.push(selected);
264                                inventory.start_uncertain = true;
265                                inventory.message = format!(
266                                    "Start result is uncertain: {}. Refresh and inspect this original ID; no automatic retry was issued.",
267                                    error.message
268                                );
269                            }
270                        }
271                    } else if let Some(selector) = locator.container_selector.as_ref() {
272                        inventory.containers.push(client.resolve(selector)?);
273                    } else {
274                        inventory.containers = client.list()?;
275                    }
276                    Ok(Arc::new(inventory))
277                })();
278                if let Some(broker) = broker {
279                    run(Box::pin(broker.close()));
280                }
281                *worker.broker.lock().unwrap() = None;
282                *worker.prompt.lock().unwrap() = None;
283                outcome
284            })();
285            *worker.result.lock().unwrap() = Some(result);
286            worker.done.finish();
287        })?;
288        Ok(ContainerQuery {
289            state,
290            _permit: permit,
291        })
292    }
293    pub fn connect_selected_container(
294        &self,
295        inventory: Arc<ContainerInventory>,
296        index: usize,
297    ) -> Result<ConnectRequest, CoreError> {
298        let selected = inventory.selected(index)?;
299        let mut locator = inventory.locator.clone();
300        locator.container_selector = Some(selected.descriptor.immutable_id.clone());
301        self.start_admitted(locator, None, None, None, false, None, Some(selected))
302    }
303}