Merge pull request #57 from Mavulp/fix-subscribe

3 files changed, 102 insertions, 48 deletions

diff --git a/src/bot/master.rs b/src/bot/master.rs
index 7a1201f..4cdb490 100644
--- a/src/bot/master.rs
+++ b/src/bot/master.rs
@@ -100,21 +100,16 @@ impl MasterBot {
         (bot, msg_loop)
     }
 
-    fn build_bot_args_for(&self, id: ClientId) -> Option<MusicBotArgs> {
-        let channel = self
-            .teamspeak
-            .channel_of_user(id)
-            .expect("Can find poke sender");
+    fn build_bot_args_for(&self, id: ClientId) -> Result<MusicBotArgs, BotCreationError> {
+        let channel = match self.teamspeak.channel_of_user(id) {
+            Some(channel) => channel,
+            None => return Err(BotCreationError::UnfoundUser),
+        };
 
         if channel == self.teamspeak.my_channel() {
-            self.teamspeak.send_message_to_user(
-                id,
-                &format!(
-                    "Joining the channel of \"{}\" is not allowed",
-                    self.config.master_name
-                ),
-            );
-            return None;
+            return Err(BotCreationError::MasterChannel(
+                self.config.master_name.clone(),
+            ));
         }
 
         let MusicBots {
@@ -126,15 +121,7 @@ impl MasterBot {
 
         for (_, bot) in connected_bots {
             if bot.my_channel() == channel {
-                self.teamspeak.send_message_to_user(
-                    id,
-                    &format!(
-                        "\"{}\" is already in this channel. \
-                         Multiple bots in one channel are not allowed.",
-                        bot.name()
-                    ),
-                );
-                return None;
+                return Err(BotCreationError::MultipleBots(bot.name().to_owned()));
             }
         }
 
@@ -147,9 +134,7 @@ impl MasterBot {
         let name_index = match available_names.pop() {
             Some(v) => v,
             None => {
-                self.teamspeak
-                    .send_message_to_user(id, "Out of names. Too many bots are already connected!");
-                return None;
+                return Err(BotCreationError::OutOfNames);
             }
         };
         let name = self.config.names[name_index].clone();
@@ -158,11 +143,7 @@ impl MasterBot {
         let id_index = match available_ids.pop() {
             Some(v) => v,
             None => {
-                self.teamspeak.send_message_to_user(
-                    id,
-                    "Out of identities. Too many bots are already connected!",
-                );
-                return None;
+                return Err(BotCreationError::OutOfIdentities);
             }
         };
 
@@ -178,7 +159,7 @@ impl MasterBot {
 
         info!("Connecting to {} on {}", channel_path, self.config.address);
 
-        Some(MusicBotArgs {
+        Ok(MusicBotArgs {
             name,
             name_index,
             id_index,
@@ -192,22 +173,32 @@ impl MasterBot {
     }
 
     async fn spawn_bot_for(&self, id: ClientId) {
-        if let Some(bot_args) = self.build_bot_args_for(id) {
-            let (bot, fut) = MusicBot::new(bot_args).await;
-            tokio::spawn(fut.unit_error().boxed().compat().map(|_| ()));
-            let mut music_bots = self.music_bots.write().expect("RwLock was not poisoned");
-            music_bots
-                .connected_bots
-                .insert(bot.name().to_string(), bot);
+        match self.build_bot_args_for(id) {
+            Ok(bot_args) => {
+                let (bot, fut) = MusicBot::new(bot_args).await;
+                tokio::spawn(fut.unit_error().boxed().compat().map(|_| ()));
+                let mut music_bots = self.music_bots.write().expect("RwLock was not poisoned");
+                music_bots
+                    .connected_bots
+                    .insert(bot.name().to_string(), bot);
+            }
+            Err(e) => self.teamspeak.send_message_to_user(id, &e.to_string()),
         }
     }
 
     async fn on_message(&self, message: MusicBotMessage) -> Result<(), AudioPlayerError> {
-        if let MusicBotMessage::TextMessage(message) = message {
-            if let MessageTarget::Poke(who) = message.target {
-                info!("Poked by {}, creating bot for their channel", who);
-                self.spawn_bot_for(who).await;
+        match message {
+            MusicBotMessage::TextMessage(message) => {
+                if let MessageTarget::Poke(who) = message.target {
+                    info!("Poked by {}, creating bot for their channel", who);
+                    self.spawn_bot_for(who).await;
+                }
+            }
+            MusicBotMessage::ChannelCreated(_) => {
+                // TODO Only subscribe to one channel
+                self.teamspeak.subscribe_all();
             }
+            _ => (),
         }
 
         Ok(())
@@ -270,6 +261,37 @@ impl MasterBot {
     }
 }
 
+#[derive(Debug)]
+pub enum BotCreationError {
+    UnfoundUser,
+    MasterChannel(String),
+    MultipleBots(String),
+    OutOfNames,
+    OutOfIdentities,
+}
+
+impl std::fmt::Display for BotCreationError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        use BotCreationError::*;
+        match self {
+            UnfoundUser => write!(
+                f,
+                "I can't find you in the channel list, \
+                    either I am not subscribed to your channel or this is a bug.",
+            ),
+            MasterChannel(name) => write!(f, "Joining the channel of \"{}\" is not allowed", name),
+            MultipleBots(name) => write!(
+                f,
+                "\"{}\" is already in this channel. \
+                         Multiple bots in one channel are not allowed.",
+                name
+            ),
+            OutOfNames => write!(f, "Out of names. Too many bots are already connected!"),
+            OutOfIdentities => write!(f, "Out of identities. Too many bots are already connected!"),
+        }
+    }
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 pub struct MasterArgs {
     #[serde(default = "default_name")]
diff --git a/src/bot/music.rs b/src/bot/music.rs
index 96cc383..14a73f9 100644
--- a/src/bot/music.rs
+++ b/src/bot/music.rs
@@ -13,10 +13,10 @@ use tsclientlib::{data, ChannelId, ClientId, ConnectOptions, Identity, Invoker,
 
 use crate::audio_player::{AudioPlayer, AudioPlayerError, PollResult};
 use crate::command::Command;
+use crate::command::VolumeChange;
 use crate::playlist::Playlist;
 use crate::teamspeak as ts;
 use crate::youtube_dl::AudioMetadata;
-use crate::command::VolumeChange;
 use ts::TeamSpeakConnection;
 
 #[derive(Debug)]
@@ -53,6 +53,7 @@ pub enum MusicBotMessage {
         client: ClientId,
         old_channel: ChannelId,
     },
+    ChannelCreated(ChannelId),
     ClientDisconnected {
         id: ClientId,
         client: data::Client,
@@ -175,7 +176,11 @@ impl MusicBot {
             format!("")
         };
 
-        self.send_message(&format!("Playing {} {}", ts::underline(&metadata.title), duration));
+        self.send_message(&format!(
+            "Playing {} {}",
+            ts::underline(&metadata.title),
+            duration
+        ));
         self.set_description(&format!("Currently playing '{}'", metadata.title));
         self.player.reset().unwrap();
         self.player.set_metadata(metadata).unwrap();
@@ -273,6 +278,10 @@ impl MusicBot {
         self.with_teamspeak(|ts| ts.set_description(desc));
     }
 
+    fn subscribe_all(&self) {
+        self.with_teamspeak(|ts| ts.subscribe_all());
+    }
+
     async fn on_text(&self, message: Message) -> Result<(), AudioPlayerError> {
         let msg = message.text;
         if msg.starts_with("!") {
@@ -310,7 +319,8 @@ impl MusicBot {
                 self.add_audio(url.to_string(), invoker.name).await;
             }
             Command::Search { query } => {
-                self.add_audio(format!("ytsearch:{}", query.join(" ")), invoker.name).await;
+                self.add_audio(format!("ytsearch:{}", query.join(" ")), invoker.name)
+                    .await;
             }
             Command::Pause => {
                 self.player.pause()?;
@@ -415,6 +425,10 @@ impl MusicBot {
                 let old_channel = client.channel;
                 self.on_client_left_channel(old_channel);
             }
+            MusicBotMessage::ChannelCreated(_) => {
+                // TODO Only subscribe to one channel
+                self.subscribe_all();
+            }
             MusicBotMessage::StateChange(state) => {
                 self.on_state(state)?;
             }
diff --git a/src/teamspeak/mod.rs b/src/teamspeak/mod.rs
index 4d31f2e..b53d1d0 100644
--- a/src/teamspeak/mod.rs
+++ b/src/teamspeak/mod.rs
@@ -35,6 +35,15 @@ fn get_message<'a>(event: &Event) -> Option<MusicBotMessage> {
             invoker: sender.clone(),
             text: msg.clone(),
         })),
+        Event::PropertyAdded {
+            id: property,
+            invoker: _,
+        } => match property {
+            PropertyId::Channel(id) => {
+                Some(MusicBotMessage::ChannelCreated(*id))
+            }
+            _ => None,
+        },
         Event::PropertyChanged {
             id: property,
             old: from,
@@ -188,7 +197,7 @@ impl TeamSpeakConnection {
                 .lock()
                 .to_mut()
                 .get_client(&self.conn.lock().own_client)
-                .expect("Can get myself")
+                .expect("can get myself")
                 .set_description(desc)
                 .map_err(|e| error!("Failed to change description: {}", e)),
         );
@@ -204,16 +213,25 @@ impl TeamSpeakConnection {
         );
     }
 
-    pub fn send_message_to_user(&self, id: ClientId, text: &str) {
+    pub fn send_message_to_user(&self, client: ClientId, text: &str) {
         tokio::spawn(
             self.conn
                 .lock()
                 .to_mut()
-                .send_message(MessageTarget::Client(id), text)
+                .send_message(MessageTarget::Client(client), text)
                 .map_err(|e| error!("Failed to send message: {}", e)),
         );
     }
 
+    pub fn subscribe_all(&self) {
+        let packet = self.conn.lock().to_mut().server.set_subscribed(true);
+        tokio::spawn(
+            self.conn
+                .send_packet(packet)
+                .map_err(|e| error!("Failed to send subscribe packet: {}", e)),
+        );
+    }
+
     pub fn disconnect(&self, reason: &str) {
         let opt = DisconnectOptions::new()
             .reason(Reason::Clientdisconnect)