Skip to main content

updatehub_sdk/
listener.rs

1// Copyright (C) 2018, 2019, 2020 O.S. Systems Sofware LTDA
2//
3// SPDX-License-Identifier: Apache-2.0
4
5//! Used to communicate with UpdateHub and listen to appropriate callbacks
6//! when the state changes.
7
8use crate::{Error, Result};
9use log::warn;
10use std::{
11    collections::HashMap, env, fs, future::Future, io, path::Path, pin::Pin, result, str::FromStr,
12    sync::Arc,
13};
14use tokio::{
15    io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
16    net::{UnixListener, UnixStream},
17    sync::Mutex,
18};
19
20const SDK_TRIGGER_FILENAME: &str =
21    "/usr/share/updatehub/state-change-callbacks.d/10-updatehub-sdk-statechange-trigger";
22const SOCKET_PATH: &str = "/run/updatehub-statechange.sock";
23
24type CallbackFn = dyn Fn(Handler) -> Pin<Box<dyn Future<Output = Result<()>>>>;
25
26/// Struct that store the callbacks for the states.
27#[derive(Default)]
28pub struct StateChange {
29    callbacks: HashMap<State, Vec<Box<CallbackFn>>>,
30}
31
32/// The state of the agent that can be handled.
33#[derive(Debug, PartialEq, Eq, Hash)]
34pub enum State {
35    Probe,
36    Download,
37    Install,
38    Reboot,
39    Error,
40}
41
42impl FromStr for State {
43    type Err = io::Error;
44
45    fn from_str(s: &str) -> result::Result<Self, Self::Err> {
46        match s {
47            "probe" => Ok(State::Probe),
48            "download" => Ok(State::Download),
49            "install" => Ok(State::Install),
50            "reboot" => Ok(State::Reboot),
51            "error" => Ok(State::Error),
52            _ => Err(io::Error::new(
53                io::ErrorKind::InvalidInput,
54                format!("the '{}' is not a valid state", s),
55            )),
56        }
57    }
58}
59
60/// Handler used to communicate with UpdateHub
61/// to call commands on the state callbacks.
62pub struct Handler {
63    stream: Arc<Mutex<UnixStream>>,
64}
65
66impl Handler {
67    /// Cancels the current action on the agent.
68    pub async fn cancel(&mut self) -> Result<()> {
69        self.stream.lock().await.write_all(b"cancel").await.map_err(Error::Io)
70    }
71
72    /// Tell the agent to proceed with the transition.
73    pub async fn proceed(&self) -> Result<()> {
74        // No message need to be sent to the connection in order to the
75        // agent to proceed handling the current state.
76        Ok(())
77    }
78}
79
80impl StateChange {
81    /// Creates a new `StateChange` struct.
82    #[inline]
83    pub fn new() -> Self {
84        StateChange::default()
85    }
86
87    /// Function that register callback(s) for the state(s).
88    /// # Example
89    ///
90    /// ```no_run
91    /// use updatehub_sdk::listener;
92    ///
93    /// # let mut listener = listener::StateChange::default();
94    /// listener.on_state(listener::State::Download, |mut handler| async move {
95    ///     println!("function called when starting the Download state");
96    ///
97    ///     // Cancels the current download.
98    ///     handler.cancel().await
99    /// });
100    /// ```
101    pub fn on_state<F, Fut>(&mut self, state: State, f: F)
102    where
103        F: Fn(Handler) -> Fut + 'static,
104        Fut: Future<Output = Result<()>> + 'static,
105    {
106        self.callbacks.entry(state).or_insert_with(Vec::new).push(Box::new(move |d| Box::pin(f(d))))
107    }
108
109    /// Start the agent to listen for messages on the socket.
110    pub async fn listen(&self) -> Result<()> {
111        let sdk_trigger = Path::new(SDK_TRIGGER_FILENAME);
112        if !sdk_trigger.exists() {
113            warn!("WARNING: updatehub-sdk-statechange-trigger not found on {:?}", sdk_trigger);
114        }
115
116        let socket_path = env::var("UH_LISTENER_TEST").unwrap_or_else(|_| SOCKET_PATH.to_string());
117        let socket_path = Path::new(&socket_path);
118        if socket_path.exists() {
119            fs::remove_file(socket_path)?;
120        }
121
122        let listener = UnixListener::bind(socket_path)?;
123        loop {
124            let (socket, ..) = listener.accept().await?;
125            self.handle_connection(socket).await?;
126        }
127    }
128
129    async fn handle_connection(&self, mut stream: UnixStream) -> Result<()> {
130        let mut line = String::new();
131
132        {
133            let mut reader = BufReader::new(&mut stream);
134            reader.read_line(&mut line).await?;
135        }
136
137        self.emit(stream, line.trim()).await
138    }
139
140    async fn emit(&self, stream: UnixStream, input: &str) -> Result<()> {
141        let state = State::from_str(input)?;
142        if let Some(callbacks) = self.callbacks.get(&state) {
143            // Since tokio::net::UnixStream doesn implement clone (or try_clone)
144            // we use an Arc + Mutex in order to be able to clone it into the handle (hance
145            // Arc) and to ensure we can get an mutable reference to it (hance the Mutex)
146            let stream = Arc::new(Mutex::new(stream));
147            for f in callbacks {
148                let stream = stream.clone();
149                f(Handler { stream }).await?;
150            }
151        }
152
153        Ok(())
154    }
155}