simploxide_client/bot/farm/
mod.rs

1//! Bot farm managing multiple bots on the same SimpleX instance
2
3use serde::Deserialize;
4use simploxide_api_types::{
5    NewUser, User,
6    client_api::ClientApi,
7    commands::{ApiDeleteUser, ApiSetActiveUser, CancelFile, ListUsers, ReceiveFile},
8    responses::{CancelFileResponse, ListUsersResponse, ReceiveFileResponse, UsersListResponse},
9};
10use tokio::sync::{
11    mpsc::{self, UnboundedReceiver, UnboundedSender},
12    oneshot,
13};
14
15use std::{
16    collections::{HashMap, hash_map::Entry},
17    sync::Arc,
18};
19
20use crate::{EventParser, EventStream, bot::BotSettings, ext::ClientApiExt as _, id::UserId};
21
22mod demux;
23mod mux;
24
25use demux::{BotMap, Channel};
26
27use super::Bot;
28
29#[cfg(feature = "xftp")]
30pub type FarmBot<C> = Bot<crate::xftp::XftpClient<DelegateClient<C>>>;
31
32#[cfg(not(feature = "xftp"))]
33pub type FarmBot<C> = Bot<DelegateClient<C>>;
34
35#[derive(Clone)]
36pub struct BotFarm<S> {
37    state: S,
38}
39
40impl<C: ClientApi, P: EventParser> BotFarm<Init<C, P>> {
41    /// The `farm_name` is the name of the special bot managing the bot farm, it cannot be accessed
42    /// directly. It is mostly used as an intermediary user deleting other users under the hood.
43    pub async fn init(
44        farm_name: String,
45        client: C,
46        events: EventStream<P>,
47    ) -> Result<Self, C::Error> {
48        let mut farm_id = BotId::anybot();
49        let mut active_name = String::new();
50        let mut cache = HashMap::new();
51        let bots = demux::FxDashMap::with_hasher(rustc_hash::FxBuildHasher);
52
53        let resp = client.list_users().await?;
54
55        for info in &resp.users {
56            let bot_id: BotId = UserId::from(info).into();
57
58            if info.user.active_user {
59                active_name = info.user.profile.display_name.clone();
60            }
61
62            if info.user.profile.display_name == farm_name {
63                farm_id = bot_id;
64                continue;
65            }
66
67            bots.insert(bot_id, Channel::Ghost);
68            cache.insert(info.user.profile.display_name.clone(), info.user.clone());
69        }
70
71        let farm_id = match farm_id.get() {
72            Some(user_id) => user_id,
73            None => {
74                let resp = client
75                    .create_active_user(NewUser {
76                        profile: Some(Bot::<C>::default_profile(farm_name.clone())),
77                        past_timestamp: false,
78                        user_chat_relay: false,
79                        undocumented: Default::default(),
80                    })
81                    .await?;
82
83                active_name = farm_name.clone();
84                UserId::from(&resp.user)
85            }
86        };
87
88        let state = Init {
89            client,
90            events,
91            farm_id,
92            farm_name,
93            active_name,
94            bots,
95            cache,
96        };
97
98        Ok(Self { state })
99    }
100
101    /// Total users count on the farm excluding the farm user
102    pub fn users_count(&self) -> usize {
103        self.state.cache.len()
104    }
105
106    /// Iterate over all users excluding the farm user
107    pub fn users(&self) -> impl Iterator<Item = &User> {
108        self.state.cache.values()
109    }
110
111    pub fn user(&self, name: &str) -> Option<&User> {
112        self.state.cache.get(name)
113    }
114
115    pub async fn remove(&mut self, user_id: UserId) -> Result<(), C::Error> {
116        self.state
117            .client
118            .api_set_active_user(ApiSetActiveUser::new(self.state.farm_id.raw()))
119            .await?;
120
121        self.state.active_name = self.state.farm_name.clone();
122
123        let resp = self
124            .state
125            .client
126            .api_delete_user(ApiDeleteUser {
127                user_id: user_id.raw(),
128                del_smp_queues: true,
129                view_pwd: None,
130            })
131            .await?;
132
133        let user = resp.user.as_ref().unwrap();
134
135        self.state.bots.remove(&user_id.into());
136        self.state.cache.remove(&user.profile.display_name);
137
138        Ok(())
139    }
140
141    pub async fn remove_by_name(&mut self, name: &str) -> Result<(), C::Error> {
142        let Some(user) = self.state.cache.remove(name) else {
143            return Ok(());
144        };
145
146        let result = self.remove(UserId::from(&user)).await;
147
148        if result.is_err() {
149            self.state.cache.insert(name.to_owned(), user);
150        }
151
152        result
153    }
154
155    /// Prepare a user with its own event stream. Use `take_bot` to extract the bot then.
156    pub async fn prepare_bot(
157        &mut self,
158        settings: BotSettings,
159    ) -> Result<UserId, CreateError<C::Error>>
160    where
161        C: Clone,
162    {
163        let user_id = self.prepare_inner(settings).await?;
164        self.state.bots.insert(user_id.into(), Channel::new_bot());
165
166        Ok(user_id)
167    }
168
169    /// Prepare a ghost user. Ghosts don't have their own event streams, all their events end up in
170    /// the general bot farm stream.
171    pub async fn prepare_ghost(
172        &mut self,
173        settings: BotSettings,
174    ) -> Result<UserId, CreateError<C::Error>>
175    where
176        C: Clone,
177    {
178        let user_id = self.prepare_inner(settings).await?;
179        self.state.bots.insert(user_id.into(), Channel::Ghost);
180
181        Ok(user_id)
182    }
183
184    /// Transition the farm to the running state by starting dispatching and routing events.
185    ///
186    /// Returns a running farm and a general [`EventStream`] that receives:
187    /// - events belonging to ghost users
188    /// - general events not addressed to a specific user(events without [`User`] struct)
189    ///
190    /// Farm user events are filtered out
191    ///
192    /// Handle events or [discard](EventStream::discard) the returned event stream to avoid memory leaks.
193    pub fn run(self) -> (BotFarm<Running<C, P>>, EventStream<P>)
194    where
195        C: 'static + Send,
196        C::Error: Send,
197        P: 'static + Send,
198    {
199        let (delegate_client, rx) = DelegateClient::new(self.state.farm_id.into());
200        mux::start(self.state.client, rx);
201
202        let bots = Arc::new(self.state.bots);
203        let (suspender, mut unmuxed_events) = demux::start(bots.clone(), self.state.events);
204
205        unmuxed_events.exclude_user(self.state.farm_id);
206
207        #[cfg(feature = "xftp")]
208        let (xftp_client, unmuxed_events) = unmuxed_events.hook_xftp(delegate_client.clone());
209
210        let state = Running {
211            farm_name: self.state.farm_name,
212            client: delegate_client,
213            suspender,
214            bots,
215            #[cfg(feature = "xftp")]
216            xftp: xftp_client.manager(),
217        };
218
219        (BotFarm { state }, unmuxed_events)
220    }
221
222    async fn prepare_inner(
223        &mut self,
224        settings: BotSettings,
225    ) -> Result<UserId, CreateError<C::Error>>
226    where
227        C: Clone,
228    {
229        if settings.display_name == self.state.farm_name {
230            return Err(CreateError::FarmUser);
231        }
232
233        match self.state.cache.entry(settings.display_name.clone()) {
234            Entry::Occupied(mut occupied) => {
235                let bot = Bot::<C>::init_existing(
236                    self.state.client.clone(),
237                    occupied.get_mut(),
238                    settings,
239                )
240                .await?;
241                let update = bot.info().await?;
242
243                *occupied.get_mut() = update.user.clone();
244                self.change_active_user(update.user.profile.display_name.clone());
245
246                Ok(bot.user_id())
247            }
248            Entry::Vacant(vacant) => {
249                let bot = Bot::<C>::init_new(self.state.client.clone(), settings).await?;
250                let update = bot.info().await?;
251
252                vacant.insert(update.user.clone());
253
254                self.change_active_user(update.user.profile.display_name.clone());
255                Ok(bot.user_id())
256            }
257        }
258    }
259
260    fn change_active_user(&mut self, new_active_username: String) {
261        if new_active_username == self.state.active_name {
262            return;
263        }
264
265        if let Some(user) = self.state.cache.get_mut(&self.state.active_name) {
266            user.active_user = false;
267        }
268
269        self.state.active_name = new_active_username;
270    }
271}
272
273impl<C: 'static + ClientApi, P: EventParser> BotFarm<Running<C, P>>
274where
275    C::Error: Send,
276{
277    /// Return a ghost handle for `user_id`, or `None` if the user does not exist or is a bot.
278    ///
279    /// Each call produces a new independent [`FarmBot`] handle. Multiple handles for the same
280    /// ghost share the underlying command channel and with the `xftp` feature enabled the same
281    /// download table, so concurrent `download_file` calls on different handles will work
282    /// correctly.
283    pub fn ghost(&self, user_id: UserId) -> Option<FarmBot<C>> {
284        let chan = self.state.bots.get(&user_id.into())?;
285
286        if let Channel::Ghost = chan.value() {
287            Some(self.make_ghost(user_id))
288        } else {
289            None
290        }
291    }
292
293    /// Take the bot handle and its [`EventStream`] out of the farm.
294    ///
295    /// This is a one-shot operation: the internal event receiver is consumed and cannot be taken
296    /// again. Panics if `user_id` is unknown, was registered as a ghost, or was already taken.
297    /// Use [`take_bot_checked`](Self::take_bot_checked) to avoid the panic.
298    pub fn take_bot(&self, user_id: UserId) -> (FarmBot<C>, EventStream<P>) {
299        let mut chan = self.state.bots.get_mut(&user_id.into()).unwrap();
300
301        if chan.is_ghost() {
302            panic!("The {user_id:?} was not initialized as bot");
303        }
304
305        let receiver = chan
306            .take_receiver()
307            .unwrap_or_else(|| panic!("The {user_id:?} was already taken"));
308
309        self.make_bot(user_id, receiver)
310    }
311
312    /// Non-panicking variant of [`take_bot`](Self::take_bot). Returns `None` if the user is
313    /// unknown, is a ghost, or was already taken.
314    pub fn take_bot_checked(&self, user_id: UserId) -> Option<(FarmBot<C>, EventStream<P>)> {
315        self.state
316            .bots
317            .get_mut(&user_id.into())
318            .and_then(|mut chan| chan.take_receiver())
319            .map(|receiver| self.make_bot(user_id, receiver))
320    }
321
322    /// Create a new SimpleX user, as a bot, and return its handle and event stream.
323    ///
324    /// Unlike `prepare_bot`, this is available at runtime after [`run`](crate::bot::BotFarm::run).
325    /// In order to route events correctly **all** event streams are paused and don't receive any
326    /// events during the bot creation process.
327    pub async fn create_bot(
328        &self,
329        settings: BotSettings,
330    ) -> Result<(FarmBot<C>, EventStream<P>), CreateError<C::Error>> {
331        let user_id = self.create_inner(settings, true).await?;
332        let (bot, stream) = self.take_bot(user_id);
333        Ok((bot, stream))
334    }
335
336    /// Return the existing bot if a user with the given display name is already known, otherwise
337    /// create one via [`create_bot`](Self::create_bot).
338    ///
339    /// # Eventual consistency
340    ///
341    /// This method is eventually consistent and may return [CreateError::Desync] if the same bot
342    /// is getting created/deleted from multiple threads. You're supposed to retry this call to get
343    /// the actual result on [CreateError::Desync]
344    pub async fn get_or_create_bot(
345        &self,
346        settings: BotSettings,
347    ) -> Result<(FarmBot<C>, EventStream<P>), CreateError<C::Error>> {
348        let resp = self.state.client.list_users().await?;
349
350        match resp.users.iter().find_map(|info| {
351            (info.user.profile.display_name == settings.display_name).then_some(UserId::from(info))
352        }) {
353            Some(user_id) => match self.state.bots.get_mut(&user_id.into()) {
354                Some(mut entry) => match entry.value_mut() {
355                    Channel::Bot(pipe) => {
356                        let receiver = pipe.take_receiver().ok_or(CreateError::BotAlreadyTaken)?;
357                        Ok(self.make_bot(user_id, receiver))
358                    }
359                    Channel::Ghost => Err(CreateError::BotIsGhost),
360                },
361                None => Err(CreateError::Desync),
362            },
363            None => self.create_bot(settings).await,
364        }
365    }
366
367    /// Create a new SimpleX user, register it as a ghost, and return a handle to it.
368    ///
369    /// The ghost's events are routed to the general [`EventStream`] returned when running a farm.
370    pub async fn create_ghost(
371        &self,
372        settings: BotSettings,
373    ) -> Result<FarmBot<C>, CreateError<C::Error>> {
374        let user_id = self.create_inner(settings, false).await?;
375        Ok(self.ghost(user_id).unwrap())
376    }
377
378    /// Return a ghost handle for the named user if it already exists, otherwise create one via
379    /// [`create_ghost`](Self::create_ghost).
380    ///
381    /// Same eventual consistency caveats as [`get_or_create_bot`](Self::get_or_create_bot).
382    pub async fn get_or_create_ghost(
383        &self,
384        settings: BotSettings,
385    ) -> Result<FarmBot<C>, CreateError<C::Error>> {
386        let resp = self.state.client.list_users().await?;
387
388        match resp.users.iter().find_map(|info| {
389            (info.user.profile.display_name == settings.display_name)
390                .then_some(UserId::from(&info.user))
391        }) {
392            Some(user_id) => match self.state.bots.get(&user_id.into()) {
393                Some(entry) => match entry.value() {
394                    Channel::Bot(_) => Err(CreateError::GhostIsBot),
395                    Channel::Ghost => Ok(self.make_ghost(user_id)),
396                },
397                None => Err(CreateError::Desync),
398            },
399            None => self.create_ghost(settings).await,
400        }
401    }
402
403    /// Permanently delete a user and remove it from the routing table.
404    ///
405    /// Any [`FarmBot`] handles that were already taken for this user remain alive but all
406    /// subsequent commands on them will fail with an API error.
407    pub async fn delete(&self, user_id: UserId) -> Result<(), C::Error> {
408        self.state
409            .client
410            .api_delete_user(ApiDeleteUser {
411                user_id: user_id.raw(),
412                del_smp_queues: true,
413                view_pwd: None,
414            })
415            .await?;
416
417        self.state.bots.remove(&user_id.into());
418        Ok(())
419    }
420
421    async fn create_inner(
422        &self,
423        settings: BotSettings,
424        is_bot: bool,
425    ) -> Result<UserId, CreateError<C::Error>> {
426        if settings.display_name == self.state.farm_name {
427            return Err(CreateError::FarmUser);
428        }
429
430        let (_guard, suspension) = oneshot::channel();
431        let _ = self.state.suspender.send(suspension);
432
433        let mut resp = self
434            .state
435            .client
436            .new_user(NewUser {
437                profile: Some(Bot::<C>::default_profile(&settings.display_name)),
438                past_timestamp: false,
439                user_chat_relay: false,
440                undocumented: Default::default(),
441            })
442            .await?;
443
444        if is_bot {
445            self.state
446                .bots
447                .insert(UserId::from(&resp.user).into(), Channel::new_bot());
448        } else {
449            self.state
450                .bots
451                .insert(UserId::from(&resp.user).into(), Channel::Ghost);
452        }
453
454        let resp = Arc::get_mut(&mut resp).unwrap();
455        let client = self.state.client.delegate_to(UserId::from(&resp.user));
456
457        match Bot::init_existing(client, &mut resp.user, settings).await {
458            Ok(bot) => Ok(bot.user_id()),
459            Err(e) => {
460                self.state
461                    .bots
462                    .insert(UserId::from(&resp.user).into(), Channel::Ghost);
463
464                Err(e.into())
465            }
466        }
467    }
468
469    fn make_bot(
470        &self,
471        user_id: UserId,
472        receiver: UnboundedReceiver<P>,
473    ) -> (FarmBot<C>, EventStream<P>) {
474        let bot_client = self.state.client.delegate_to(user_id);
475        let stream = EventStream::from(receiver);
476
477        #[cfg(feature = "xftp")]
478        let (bot_client, stream) = stream.hook_xftp(bot_client);
479
480        (Bot::new(bot_client, user_id), stream)
481    }
482
483    fn make_ghost(&self, user_id: UserId) -> FarmBot<C> {
484        let bot_client = self.state.client.delegate_to(user_id);
485
486        #[cfg(feature = "xftp")]
487        let bot_client = crate::xftp::XftpClient::new(bot_client, self.state.xftp.clone());
488
489        Bot::new(bot_client, user_id)
490    }
491}
492
493pub struct Init<C, P> {
494    client: C,
495    events: EventStream<P>,
496    farm_id: UserId,
497    farm_name: String,
498    active_name: String,
499    bots: BotMap<P>,
500    cache: HashMap<String, User>,
501}
502
503#[derive(Clone)]
504pub struct Running<C: ClientApi, P> {
505    farm_name: String,
506    client: DelegateClient<C>,
507    suspender: demux::Suspender,
508    bots: Arc<BotMap<P>>,
509    #[cfg(feature = "xftp")]
510    xftp: Arc<crate::xftp::XftpManager>,
511}
512
513pub struct DelegateClient<C: ClientApi> {
514    bot_id: BotId,
515    sender: DelegateSender<C>,
516}
517
518impl<C: ClientApi> Clone for DelegateClient<C> {
519    fn clone(&self) -> Self {
520        Self {
521            bot_id: self.bot_id,
522            sender: self.sender.clone(),
523        }
524    }
525}
526
527impl<C: ClientApi> DelegateClient<C> {
528    fn new(bot_id: BotId) -> (Self, DelegateReceiver<C>) {
529        let (sender, receiver) = mpsc::unbounded_channel();
530        (Self { bot_id, sender }, receiver)
531    }
532
533    fn delegate_to(&self, bot_id: impl Into<BotId>) -> Self {
534        Self {
535            bot_id: bot_id.into(),
536            sender: self.sender.clone(),
537        }
538    }
539}
540
541impl<C: ClientApi> ClientApi for DelegateClient<C>
542where
543    C::Error: Send,
544{
545    type ResponseShape<'de, T: 'de + Deserialize<'de>> = C::ResponseShape<'de, T>;
546    type Error = C::Error;
547
548    async fn send_raw(&self, cmd: String) -> Result<String, Self::Error> {
549        let (responder, response) = oneshot::channel();
550
551        let request = DelegateRequest {
552            bot_id: self.bot_id,
553            cmd,
554            responder,
555        };
556
557        self.sender
558            .send(request)
559            .expect("Delegate client cannot outlive background task");
560
561        response
562            .await
563            .expect("Delegate client cannot outlive background task")
564    }
565
566    async fn list_users(&self) -> Result<Arc<UsersListResponse>, Self::Error> {
567        let client = self.delegate_to(BotId::anybot());
568        let response: ListUsersResponse = client.send(ListUsers {}).await?;
569        Ok(response.into_inner())
570    }
571
572    async fn receive_file(&self, cmd: ReceiveFile) -> Result<ReceiveFileResponse, Self::Error> {
573        let client = self.delegate_to(BotId::anybot());
574        client.send(cmd).await
575    }
576
577    async fn cancel_file(&self, file_id: i64) -> Result<CancelFileResponse, Self::Error> {
578        let client = self.delegate_to(BotId::anybot());
579        client.send(CancelFile { file_id }).await
580    }
581}
582
583#[derive(Debug)]
584pub enum CreateError<E> {
585    /// Farm user cannot be interacted with directly
586    FarmUser,
587    /// Bot cannot be created because ghost with same name already exists
588    BotIsGhost,
589    /// Ghost cannot be craeted because bot with same name already exists
590    GhostIsBot,
591    /// The bot already exists and was already taken with the `take_bot`
592    BotAlreadyTaken,
593    /// The in memory state is not synced with the DB, retry later
594    Desync,
595    Api(E),
596}
597
598impl<E> From<E> for CreateError<E> {
599    fn from(value: E) -> Self {
600        Self::Api(value)
601    }
602}
603
604impl<E> std::fmt::Display for CreateError<E>
605where
606    E: std::fmt::Display,
607{
608    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
609        match self {
610            Self::FarmUser => write!(
611                f,
612                "Attempt to create a farm user. Farm user is special and cannot be interacted with directly"
613            ),
614            Self::BotIsGhost => write!(
615                f,
616                "Cannot create a bot because the ghost user with the same name already exists"
617            ),
618            Self::GhostIsBot => write!(
619                f,
620                "Cannot create a ghost because the bot user with the same name already exists"
621            ),
622            Self::BotAlreadyTaken => {
623                write!(
624                    f,
625                    "The bot already exists and has been taken from the farm. Cannot recreate operational bots"
626                )
627            }
628            Self::Desync => {
629                write!(
630                    f,
631                    "The DB state was not in sync with the memory state, try again"
632                )
633            }
634            Self::Api(e) => write!(f, "{e:#}"),
635        }
636    }
637}
638
639impl<E: 'static + std::error::Error> std::error::Error for CreateError<E> {
640    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
641        if let Self::Api(err) = self {
642            Some(err)
643        } else {
644            None
645        }
646    }
647}
648
649struct DelegateRequest<C: ClientApi> {
650    bot_id: BotId,
651    cmd: String,
652    responder: oneshot::Sender<Result<String, C::Error>>,
653}
654
655type DelegateSender<C> = UnboundedSender<DelegateRequest<C>>;
656type DelegateReceiver<C> = UnboundedReceiver<DelegateRequest<C>>;
657
658#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
659#[repr(transparent)]
660struct BotId(Option<UserId>);
661
662impl BotId {
663    /// Used as an optimization for commands that can execute from any active bot account
664    fn anybot() -> Self {
665        Self(None)
666    }
667
668    fn get(&self) -> Option<UserId> {
669        self.0
670    }
671}
672
673impl From<UserId> for BotId {
674    fn from(user_id: UserId) -> Self {
675        Self(Some(user_id))
676    }
677}
simploxide_client/bot/farm/mod.rs

simploxide_client/bot/farm/
mod.rs
