<commit_before>use std::collections::HashMap;
use std::time::Duration;

use anyhow::bail;
use async_recursion::async_recursion;
use serenity::http::{Http, HttpError};
use serenity::model::application::interaction::application_command::{
    ApplicationCommandInteraction, CommandData, CommandDataOptionValue,
};
use serenity::model::application::interaction::{autocomplete, InteractionResponseType};
use serenity::model::application::interaction::MessageFlags;
use serenity::model::channel::{Message, PartialChannel};
use serenity::model::event::MessageUpdateEvent;
use serenity::model::guild::{Member, PartialMember, Role};
use serenity::model::id::{ChannelId, GuildId, MessageId, RoleId, UserId};
use serenity::model::user::User;
use serenity::model::Permissions;
use serenity::Error;
use serenity::model::prelude::interaction::autocomplete::AutocompleteData;
use tracing::{error, warn};
use typemap_rev::TypeMap;

use crate::database::settings::{
    Setting, ACTION_BAN, ACTION_HARDBAN, ACTION_KICK, ACTION_MUTE, ACTION_SOFTBAN, ACTION_WARN,
};
use crate::service::ban::BanService;
use crate::service::guild::CachedRole;
use crate::service::hardban::HardbanService;
use crate::service::kick::KickService;
use crate::service::mute::MuteService;
use crate::service::softban::SoftbanService;
use crate::service::warn::WarnService;
use crate::util::now;

pub mod mod_log;
pub mod user_dm;

pub trait AutocompleteDataExt {
    fn option(&self, option_name: &str) -> Option<&autocomplete::CommandDataOptionValue>;
    fn string(&self, option_name: &str) -> Option<&str>;
}

impl AutocompleteDataExt for AutocompleteData {
    fn option(&self, option_name: &str) -> Option<&autocomplete::CommandDataOptionValue> {
        self.options
            .iter()
            .find(|option| option.name == option_name)
            .map(|option| &option.value)
    }

    fn string(&self, option_name: &str) -> Option<&str> {
        if let Some(autocomplete::CommandDataOptionValue::String(str)) = self.option(option_name) {
            Some(str.as_str())
        } else {
            None
        }
    }
}

pub trait CommandDataExt {
    fn option(&self, option_name: &str) -> Option<&CommandDataOptionValue>;
    fn string(&self, option_name: &str) -> Option<&str>;
    fn integer(&self, option_name: &str) -> Option<i64>;
    fn boolean(&self, option_name: &str) -> Option<bool>;
    fn user(&self, option_name: &str) -> Option<(&User, Option<&PartialMember>)>;
    fn channel(&self, option_name: &str) -> Option<&PartialChannel>;
    fn role(&self, option_name: &str) -> Option<&Role>;
    fn number(&self, option_name: &str) -> Option<f64>;
}

impl CommandDataExt for CommandData {
    fn option(&self, option_name: &str) -> Option<&CommandDataOptionValue> {
        self.options
            .iter()
            .find(|option| option.name == option_name)
            .and_then(|option| option.resolved.as_ref())
    }

    fn string(&self, option_name: &str) -> Option<&str> {
        if let Some(CommandDataOptionValue::String(str)) = self.option(option_name) {
            Some(str.as_str())
        } else {
            None
        }
    }

    fn integer(&self, option_name: &str) -> Option<i64> {
        if let Some(CommandDataOptionValue::Integer(i)) = self.option(option_name) {
            Some(*i)
        } else {
            None
        }
    }

    fn boolean(&self, option_name: &str) -> Option<bool> {
        if let Some(CommandDataOptionValue::Boolean(b)) = self.option(option_name) {
            Some(*b)
        } else {
            None
        }
    }

    fn user(&self, option_name: &str) -> Option<(&User, Option<&PartialMember>)> {
        if let Some(CommandDataOptionValue::User(user, member)) = self.option(option_name) {
            Some((user, member.as_ref()))
        } else {
            None
        }
    }

    fn channel(&self, option_name: &str) -> Option<&PartialChannel> {
        if let Some(CommandDataOptionValue::Channel(channel)) = self.option(option_name) {
            Some(channel)
        } else {
            None
        }
    }

    fn role(&self, option_name: &str) -> Option<&Role> {
        if let Some(CommandDataOptionValue::Role(role)) = self.option(option_name) {
            Some(role)
        } else {
            None
        }
    }

    fn number(&self, option_name: &str) -> Option<f64> {
        if let Some(CommandDataOptionValue::Number(n)) = self.option(option_name) {
            Some(*n)
        } else {
            None
        }
    }
}

pub fn is_staff(permissions: Permissions) -> bool {
    permissions.administrator()
        || permissions.ban_members()
        || permissions.kick_members()
        || permissions.manage_roles()
        || permissions.manage_messages()
}

pub fn get_permissions(
    member_id: UserId,
    member_roles: &[RoleId],
    guild_id: GuildId,
    guild_roles: &HashMap<RoleId, CachedRole>,
    guild_owner_id: UserId,
) -> Option<Permissions> {
    if member_id == guild_owner_id {
        return Some(Permissions::ADMINISTRATOR);
    }

    let everyone = match guild_roles.get(&RoleId(guild_id.0)) {
        Some(everyone) => everyone,
        None => {
            return None;
        }
    };

    let mut permissions = everyone.permissions;

    for role in member_roles {
        if let Some(role) = guild_roles.get(role) {
            if role.permissions.contains(Permissions::ADMINISTRATOR) {
                return Some(Permissions::ADMINISTRATOR);
            }

            permissions |= role.permissions;
        } else {
            return None;
        }
    }

    Some(permissions)
}

pub trait UserExt {
    fn tag_and_id(&self) -> String;
}

impl UserExt for User {
    fn tag_and_id(&self) -> String {
        format!("{} ({})", self.tag(), self.id.0)
    }
}

pub struct GuildSlashCommandInteraction<'a> {
    pub guild_id: GuildId,
    pub member: &'a Member,
    pub permissions: Permissions,
}

pub fn verify_guild_slash_command(
    interaction: &ApplicationCommandInteraction,
) -> anyhow::Result<GuildSlashCommandInteraction<'_>> {
    let member = if let Some(member) = &interaction.member {
        member
    } else {
        bail!("interaction has no member field");
    };

    let permissions = if let Some(permissions) = member.permissions {
        permissions
    } else {
        bail!("interaction's member has no permissions field");
    };

    let guild_id = if let Some(guild_id) = interaction.guild_id {
        guild_id
    } else {
        bail!("interaction has missing guild id");
    };

    Ok(GuildSlashCommandInteraction {
        guild_id,
        member,
        permissions,
    })
}

pub struct GuildMessageCreated<'a> {
    pub guild_id: GuildId,
    pub member: &'a PartialMember,
}

pub fn verify_guild_message_create(message: &Message) -> Option<GuildMessageCreated> {
    Some(GuildMessageCreated {
        guild_id: message.guild_id?,
        member: message.member.as_ref()?,
    })
}

pub struct GuildMessageUpdated<'a> {
    pub guild_id: GuildId,
    pub content: &'a str,
}

pub fn verify_guild_message_update(message: &MessageUpdateEvent) -> Option<GuildMessageUpdated> {
    Some(GuildMessageUpdated {
        guild_id: message.guild_id?,
        content: message.content.as_ref()?,
    })
}

pub async fn reply_with_str(
    http: &Http,
    interaction: &ApplicationCommandInteraction,
    content: &str,
) {
    let _ = interaction
        .create_interaction_response(http, |response| {
            response
                .kind(InteractionResponseType::ChannelMessageWithSource)
                .interaction_response_data(|message| message.content(content))
        })
        .await
        .map_err(|err| {
            error!("failed to reply to interaction {:?} {}", interaction, err);
            err
        });
}

async fn invisible_reply_with_str(
    http: &Http,
    interaction: &ApplicationCommandInteraction,
    content: &str,
) {
    let _ = interaction
        .create_interaction_response(http, |response| {
            response
                .kind(InteractionResponseType::ChannelMessageWithSource)
                .interaction_response_data(|message| {
                    message.content(content).flags(MessageFlags::EPHEMERAL)
                })
        })
        .await
        .map_err(|err| {
            error!("failed to reply to interaction {:?} {}", interaction, err);
            err
        });
}

pub async fn edit_interaction_response(
    http: &Http,
    interaction: &ApplicationCommandInteraction,
    content: &str,
) {
    let _ = interaction
        .edit_original_interaction_response(http, |response| response.content(content))
        .await
        .map_err(|err| {
            error!("failed to edit interaction reply {:?} {}", interaction, err);
            err
        });
}

pub async fn unauthorized_reply(
    http: &Http,
    interaction: &ApplicationCommandInteraction,
    required_permission: Permissions,
) {
    invisible_failure_reply(
        http,
        interaction,
        &format!(
            "You don't have enough permissions to execute this command! Required permission: {}",
            required_permission
        ),
    )
    .await;
}

pub async fn invisible_success_reply(
    http: &Http,
    interaction: &ApplicationCommandInteraction,
    failure_message: &str,
) {
    invisible_reply_with_str(http, interaction, failure_message).await;
}

pub async fn invisible_failure_reply(
    http: &Http,
    interaction: &ApplicationCommandInteraction,
    failure_message: &str,
) {
    invisible_reply_with_str(http, interaction, failure_message).await;
}

pub enum CleanMessagesFailure {
    Unauthorized,
    Other,
}

pub async fn clean_messages(
    http: &Http,
    channel: ChannelId,
    messages: Vec<MessageId>,
) -> Result<(), CleanMessagesFailure> {
    let mut old_messages = vec![];
    let mut new_messages = vec![];
    let now = now();

    for message in messages.into_iter() {
        let timestamp = message.created_at().unix_timestamp() as u64;
        if (now - timestamp) <= 60 * 60 * 24 * 12 {
            new_messages.push(message);
        } else {
            old_messages.push(message);
        }
    }

    let result = if new_messages.len() >= 2 {
        channel.delete_messages(http, new_messages).await
    } else if new_messages.len() == 1 {
        channel.delete_message(http, new_messages[0]).await
    } else {
        Ok(())
    };

    if let Err(err) = result {
        return match err.discord_error_code() {
            Some(50013) => Err(CleanMessagesFailure::Unauthorized),
            _ => {
                error!("failed to delete messages in bulk {}", err);
                Err(CleanMessagesFailure::Other)
            }
        };
    }

    let mut result = Ok(());
    for message_id in old_messages {
        if let Err(err) = channel.delete_message(http, message_id).await {
            result = Err(err);
            break;
        }
    }

    if let Err(err) = result {
        return match err.discord_error_code() {
            Some(50013) => Err(CleanMessagesFailure::Unauthorized),
            _ => {
                error!("failed to delete messages one by one {}", err);
                Err(CleanMessagesFailure::Other)
            }
        };
    }

    Ok(())
}

pub trait SerenityErrorExt {
    fn discord_error_code(&self) -> Option<isize>;
}

impl SerenityErrorExt for Error {
    fn discord_error_code(&self) -> Option<isize> {
        match self {
            Error::Http(HttpError::UnsuccessfulRequest(response)) => Some(response.error.code),
            _ => None,
        }
    }
}

#[async_recursion]
pub async fn execute_mod_action(
    action_kind: i32,
    http: &Http,
    guild_id: GuildId,
    guild_name: &str,
    setting: &Setting,
    services: &TypeMap,
    channel_id: Option<ChannelId>,
    mod_user_id: UserId,
    mod_user_tag_and_id: &str,
    target_user: &User,
    reason: String,
    duration: Option<Duration>,
    call_depth: i32,
) {
    if call_depth >= 3 {
        return;
    }

    let warn_service = if let Some(service) = services.get::<WarnService>() {
        service
    } else {
        error!("couldn't get warn service!");
        return;
    };

    let mute_service = if let Some(service) = services.get::<MuteService>() {
        service
    } else {
        error!("couldn't get mute service!");
        return;
    };

    let kick_service = if let Some(service) = services.get::<KickService>() {
        service
    } else {
        error!("couldn't get kick service!");
        return;
    };

    let ban_service = if let Some(service) = services.get::<BanService>() {
        service
    } else {
        error!("couldn't get ban service!");
        return;
    };

    let softban_service = if let Some(service) = services.get::<SoftbanService>() {
        service
    } else {
        error!("couldn't get softban service!");
        return;
    };

    let hardban_service = if let Some(service) = services.get::<HardbanService>() {
        service
    } else {
        error!("couldn't get hardban service!");
        return;
    };

    if action_kind == ACTION_WARN {
        let _ = warn_service
            .issue_warn(
                http,
                guild_id,
                guild_name,
                setting,
                services,
                channel_id,
                mod_user_id,
                mod_user_tag_and_id,
                target_user,
                reason,
                call_depth + 1,
            )
            .await;
    } else if action_kind == ACTION_MUTE {
        let _ = mute_service
            .issue_mute(
                http,
                guild_id,
                guild_name,
                setting,
                services,
                channel_id,
                mod_user_id,
                mod_user_tag_and_id,
                target_user,
                reason,
                duration,
                call_depth + 1,
            )
            .await;
    } else if action_kind == ACTION_KICK {
        let _ = kick_service
            .issue_kick(
                http,
                guild_id,
                guild_name,
                setting,
                services,
                channel_id,
                mod_user_id,
                mod_user_tag_and_id,
                target_user,
                reason,
                call_depth + 1,
            )
            .await;
    } else if action_kind == ACTION_BAN {
        let _ = ban_service
            .issue_ban(
                http,
                guild_id,
                guild_name,
                setting,
                channel_id,
                mod_user_id,
                mod_user_tag_and_id,
                target_user,
                reason,
                duration,
            )
            .await;
    } else if action_kind == ACTION_SOFTBAN {
        let _ = softban_service
            .issue_softban(
                http,
                guild_id,
                guild_name,
                setting,
                services,
                channel_id,
                mod_user_id,
                mod_user_tag_and_id,
                target_user,
                reason,
                1,
                call_depth + 1,
            )
            .await;
    } else if action_kind == ACTION_HARDBAN {
        let _ = hardban_service
            .issue_hardban(
                http,
                guild_id,
                guild_name,
                setting,
                channel_id,
                mod_user_id,
                mod_user_tag_and_id,
                target_user,
                reason,
            )
            .await;
    } else {
        warn!(
            action_kind = action_kind,
            guild_id = guild_id.0,
            "invalid state for mod action duration type!"
        );
    }
}
<commit_msg>Fix autocomplete not working for tag command<commit_after>use std::collections::HashMap;
use std::time::Duration;

use anyhow::bail;
use async_recursion::async_recursion;
use serenity::http::{Http, HttpError};
use serenity::model::application::interaction::application_command::{
    ApplicationCommandInteraction, CommandData, CommandDataOptionValue,
};
use serenity::model::application::interaction::{autocomplete, InteractionResponseType};
use serenity::model::application::interaction::MessageFlags;
use serenity::model::channel::{Message, PartialChannel};
use serenity::model::event::MessageUpdateEvent;
use serenity::model::guild::{Member, PartialMember, Role};
use serenity::model::id::{ChannelId, GuildId, MessageId, RoleId, UserId};
use serenity::model::user::User;
use serenity::model::Permissions;
use serenity::Error;
use serenity::model::prelude::interaction::autocomplete::AutocompleteData;
use tracing::{error, warn};
use typemap_rev::TypeMap;

use crate::database::settings::{
    Setting, ACTION_BAN, ACTION_HARDBAN, ACTION_KICK, ACTION_MUTE, ACTION_SOFTBAN, ACTION_WARN,
};
use crate::service::ban::BanService;
use crate::service::guild::CachedRole;
use crate::service::hardban::HardbanService;
use crate::service::kick::KickService;
use crate::service::mute::MuteService;
use crate::service::softban::SoftbanService;
use crate::service::warn::WarnService;
use crate::util::now;

pub mod mod_log;
pub mod user_dm;

pub trait AutocompleteDataExt {
    fn option(&self, option_name: &str) -> Option<&autocomplete::CommandDataOptionValue>;
    fn string(&self, option_name: &str) -> Option<&str>;
}

impl AutocompleteDataExt for AutocompleteData {
    fn option(&self, option_name: &str) -> Option<&autocomplete::CommandDataOptionValue> {
        self.options
            .iter()
            .find(|option| option.name == option_name)
            .map(|option| &option.value)
    }

    fn string(&self, option_name: &str) -> Option<&str> {
        if let Some(autocomplete::CommandDataOptionValue::Autocomplete {value, kind: _}) = self.option(option_name) {
            Some(value.as_str())
        } else {
            None
        }
    }
}

pub trait CommandDataExt {
    fn option(&self, option_name: &str) -> Option<&CommandDataOptionValue>;
    fn string(&self, option_name: &str) -> Option<&str>;
    fn integer(&self, option_name: &str) -> Option<i64>;
    fn boolean(&self, option_name: &str) -> Option<bool>;
    fn user(&self, option_name: &str) -> Option<(&User, Option<&PartialMember>)>;
    fn channel(&self, option_name: &str) -> Option<&PartialChannel>;
    fn role(&self, option_name: &str) -> Option<&Role>;
    fn number(&self, option_name: &str) -> Option<f64>;
}

impl CommandDataExt for CommandData {
    fn option(&self, option_name: &str) -> Option<&CommandDataOptionValue> {
        self.options
            .iter()
            .find(|option| option.name == option_name)
            .and_then(|option| option.resolved.as_ref())
    }

    fn string(&self, option_name: &str) -> Option<&str> {
        if let Some(CommandDataOptionValue::String(str)) = self.option(option_name) {
            Some(str.as_str())
        } else {
            None
        }
    }

    fn integer(&self, option_name: &str) -> Option<i64> {
        if let Some(CommandDataOptionValue::Integer(i)) = self.option(option_name) {
            Some(*i)
        } else {
            None
        }
    }

    fn boolean(&self, option_name: &str) -> Option<bool> {
        if let Some(CommandDataOptionValue::Boolean(b)) = self.option(option_name) {
            Some(*b)
        } else {
            None
        }
    }

    fn user(&self, option_name: &str) -> Option<(&User, Option<&PartialMember>)> {
        if let Some(CommandDataOptionValue::User(user, member)) = self.option(option_name) {
            Some((user, member.as_ref()))
        } else {
            None
        }
    }

    fn channel(&self, option_name: &str) -> Option<&PartialChannel> {
        if let Some(CommandDataOptionValue::Channel(channel)) = self.option(option_name) {
            Some(channel)
        } else {
            None
        }
    }

    fn role(&self, option_name: &str) -> Option<&Role> {
        if let Some(CommandDataOptionValue::Role(role)) = self.option(option_name) {
            Some(role)
        } else {
            None
        }
    }

    fn number(&self, option_name: &str) -> Option<f64> {
        if let Some(CommandDataOptionValue::Number(n)) = self.option(option_name) {
            Some(*n)
        } else {
            None
        }
    }
}

pub fn is_staff(permissions: Permissions) -> bool {
    permissions.administrator()
        || permissions.ban_members()
        || permissions.kick_members()
        || permissions.manage_roles()
        || permissions.manage_messages()
}

pub fn get_permissions(
    member_id: UserId,
    member_roles: &[RoleId],
    guild_id: GuildId,
    guild_roles: &HashMap<RoleId, CachedRole>,
    guild_owner_id: UserId,
) -> Option<Permissions> {
    if member_id == guild_owner_id {
        return Some(Permissions::ADMINISTRATOR);
    }

    let everyone = match guild_roles.get(&RoleId(guild_id.0)) {
        Some(everyone) => everyone,
        None => {
            return None;
        }
    };

    let mut permissions = everyone.permissions;

    for role in member_roles {
        if let Some(role) = guild_roles.get(role) {
            if role.permissions.contains(Permissions::ADMINISTRATOR) {
                return Some(Permissions::ADMINISTRATOR);
            }

            permissions |= role.permissions;
        } else {
            return None;
        }
    }

    Some(permissions)
}

pub trait UserExt {
    fn tag_and_id(&self) -> String;
}

impl UserExt for User {
    fn tag_and_id(&self) -> String {
        format!("{} ({})", self.tag(), self.id.0)
    }
}

pub struct GuildSlashCommandInteraction<'a> {
    pub guild_id: GuildId,
    pub member: &'a Member,
    pub permissions: Permissions,
}

pub fn verify_guild_slash_command(
    interaction: &ApplicationCommandInteraction,
) -> anyhow::Result<GuildSlashCommandInteraction<'_>> {
    let member = if let Some(member) = &interaction.member {
        member
    } else {
        bail!("interaction has no member field");
    };

    let permissions = if let Some(permissions) = member.permissions {
        permissions
    } else {
        bail!("interaction's member has no permissions field");
    };

    let guild_id = if let Some(guild_id) = interaction.guild_id {
        guild_id
    } else {
        bail!("interaction has missing guild id");
    };

    Ok(GuildSlashCommandInteraction {
        guild_id,
        member,
        permissions,
    })
}

pub struct GuildMessageCreated<'a> {
    pub guild_id: GuildId,
    pub member: &'a PartialMember,
}

pub fn verify_guild_message_create(message: &Message) -> Option<GuildMessageCreated> {
    Some(GuildMessageCreated {
        guild_id: message.guild_id?,
        member: message.member.as_ref()?,
    })
}

pub struct GuildMessageUpdated<'a> {
    pub guild_id: GuildId,
    pub content: &'a str,
}

pub fn verify_guild_message_update(message: &MessageUpdateEvent) -> Option<GuildMessageUpdated> {
    Some(GuildMessageUpdated {
        guild_id: message.guild_id?,
        content: message.content.as_ref()?,
    })
}

pub async fn reply_with_str(
    http: &Http,
    interaction: &ApplicationCommandInteraction,
    content: &str,
) {
    let _ = interaction
        .create_interaction_response(http, |response| {
            response
                .kind(InteractionResponseType::ChannelMessageWithSource)
                .interaction_response_data(|message| message.content(content))
        })
        .await
        .map_err(|err| {
            error!("failed to reply to interaction {:?} {}", interaction, err);
            err
        });
}

async fn invisible_reply_with_str(
    http: &Http,
    interaction: &ApplicationCommandInteraction,
    content: &str,
) {
    let _ = interaction
        .create_interaction_response(http, |response| {
            response
                .kind(InteractionResponseType::ChannelMessageWithSource)
                .interaction_response_data(|message| {
                    message.content(content).flags(MessageFlags::EPHEMERAL)
                })
        })
        .await
        .map_err(|err| {
            error!("failed to reply to interaction {:?} {}", interaction, err);
            err
        });
}

pub async fn edit_interaction_response(
    http: &Http,
    interaction: &ApplicationCommandInteraction,
    content: &str,
) {
    let _ = interaction
        .edit_original_interaction_response(http, |response| response.content(content))
        .await
        .map_err(|err| {
            error!("failed to edit interaction reply {:?} {}", interaction, err);
            err
        });
}

pub async fn unauthorized_reply(
    http: &Http,
    interaction: &ApplicationCommandInteraction,
    required_permission: Permissions,
) {
    invisible_failure_reply(
        http,
        interaction,
        &format!(
            "You don't have enough permissions to execute this command! Required permission: {}",
            required_permission
        ),
    )
    .await;
}

pub async fn invisible_success_reply(
    http: &Http,
    interaction: &ApplicationCommandInteraction,
    failure_message: &str,
) {
    invisible_reply_with_str(http, interaction, failure_message).await;
}

pub async fn invisible_failure_reply(
    http: &Http,
    interaction: &ApplicationCommandInteraction,
    failure_message: &str,
) {
    invisible_reply_with_str(http, interaction, failure_message).await;
}

pub enum CleanMessagesFailure {
    Unauthorized,
    Other,
}

pub async fn clean_messages(
    http: &Http,
    channel: ChannelId,
    messages: Vec<MessageId>,
) -> Result<(), CleanMessagesFailure> {
    let mut old_messages = vec![];
    let mut new_messages = vec![];
    let now = now();

    for message in messages.into_iter() {
        let timestamp = message.created_at().unix_timestamp() as u64;
        if (now - timestamp) <= 60 * 60 * 24 * 12 {
            new_messages.push(message);
        } else {
            old_messages.push(message);
        }
    }

    let result = if new_messages.len() >= 2 {
        channel.delete_messages(http, new_messages).await
    } else if new_messages.len() == 1 {
        channel.delete_message(http, new_messages[0]).await
    } else {
        Ok(())
    };

    if let Err(err) = result {
        return match err.discord_error_code() {
            Some(50013) => Err(CleanMessagesFailure::Unauthorized),
            _ => {
                error!("failed to delete messages in bulk {}", err);
                Err(CleanMessagesFailure::Other)
            }
        };
    }

    let mut result = Ok(());
    for message_id in old_messages {
        if let Err(err) = channel.delete_message(http, message_id).await {
            result = Err(err);
            break;
        }
    }

    if let Err(err) = result {
        return match err.discord_error_code() {
            Some(50013) => Err(CleanMessagesFailure::Unauthorized),
            _ => {
                error!("failed to delete messages one by one {}", err);
                Err(CleanMessagesFailure::Other)
            }
        };
    }

    Ok(())
}

pub trait SerenityErrorExt {
    fn discord_error_code(&self) -> Option<isize>;
}

impl SerenityErrorExt for Error {
    fn discord_error_code(&self) -> Option<isize> {
        match self {
            Error::Http(HttpError::UnsuccessfulRequest(response)) => Some(response.error.code),
            _ => None,
        }
    }
}

#[async_recursion]
pub async fn execute_mod_action(
    action_kind: i32,
    http: &Http,
    guild_id: GuildId,
    guild_name: &str,
    setting: &Setting,
    services: &TypeMap,
    channel_id: Option<ChannelId>,
    mod_user_id: UserId,
    mod_user_tag_and_id: &str,
    target_user: &User,
    reason: String,
    duration: Option<Duration>,
    call_depth: i32,
) {
    if call_depth >= 3 {
        return;
    }

    let warn_service = if let Some(service) = services.get::<WarnService>() {
        service
    } else {
        error!("couldn't get warn service!");
        return;
    };

    let mute_service = if let Some(service) = services.get::<MuteService>() {
        service
    } else {
        error!("couldn't get mute service!");
        return;
    };

    let kick_service = if let Some(service) = services.get::<KickService>() {
        service
    } else {
        error!("couldn't get kick service!");
        return;
    };

    let ban_service = if let Some(service) = services.get::<BanService>() {
        service
    } else {
        error!("couldn't get ban service!");
        return;
    };

    let softban_service = if let Some(service) = services.get::<SoftbanService>() {
        service
    } else {
        error!("couldn't get softban service!");
        return;
    };

    let hardban_service = if let Some(service) = services.get::<HardbanService>() {
        service
    } else {
        error!("couldn't get hardban service!");
        return;
    };

    if action_kind == ACTION_WARN {
        let _ = warn_service
            .issue_warn(
                http,
                guild_id,
                guild_name,
                setting,
                services,
                channel_id,
                mod_user_id,
                mod_user_tag_and_id,
                target_user,
                reason,
                call_depth + 1,
            )
            .await;
    } else if action_kind == ACTION_MUTE {
        let _ = mute_service
            .issue_mute(
                http,
                guild_id,
                guild_name,
                setting,
                services,
                channel_id,
                mod_user_id,
                mod_user_tag_and_id,
                target_user,
                reason,
                duration,
                call_depth + 1,
            )
            .await;
    } else if action_kind == ACTION_KICK {
        let _ = kick_service
            .issue_kick(
                http,
                guild_id,
                guild_name,
                setting,
                services,
                channel_id,
                mod_user_id,
                mod_user_tag_and_id,
                target_user,
                reason,
                call_depth + 1,
            )
            .await;
    } else if action_kind == ACTION_BAN {
        let _ = ban_service
            .issue_ban(
                http,
                guild_id,
                guild_name,
                setting,
                channel_id,
                mod_user_id,
                mod_user_tag_and_id,
                target_user,
                reason,
                duration,
            )
            .await;
    } else if action_kind == ACTION_SOFTBAN {
        let _ = softban_service
            .issue_softban(
                http,
                guild_id,
                guild_name,
                setting,
                services,
                channel_id,
                mod_user_id,
                mod_user_tag_and_id,
                target_user,
                reason,
                1,
                call_depth + 1,
            )
            .await;
    } else if action_kind == ACTION_HARDBAN {
        let _ = hardban_service
            .issue_hardban(
                http,
                guild_id,
                guild_name,
                setting,
                channel_id,
                mod_user_id,
                mod_user_tag_and_id,
                target_user,
                reason,
            )
            .await;
    } else {
        warn!(
            action_kind = action_kind,
            guild_id = guild_id.0,
            "invalid state for mod action duration type!"
        );
    }
}
<|endoftext|>


<commit_before>#![allow(dead_code)]

use std::io::Write;
use std::io::Read;

use std::str;
use std::net;
use std::net::ToSocketAddrs;

use bytes::Bytes;

use httpbis;
use httpbis::message::SimpleHttpMessage;
use httpbis::solicit::StreamId;
use httpbis::error::ErrorCode;
use httpbis::solicit::header::*;
use httpbis::solicit::frame::FrameIR;
use httpbis::solicit::frame::settings::SettingsFrame;
use httpbis::solicit::frame::headers::HeadersFrame;
use httpbis::solicit::frame::headers::HeadersFlag;
use httpbis::solicit::frame::data::DataFrame;
use httpbis::solicit::frame::data::DataFlag;
use httpbis::solicit::frame::goaway::GoawayFrame;
use httpbis::solicit::frame::RawFrame;
use httpbis::solicit::frame::rst_stream::RstStreamFrame;
use httpbis::solicit::connection::HttpFrame;
use httpbis::solicit::connection::HttpConnection;

pub struct HttpServerTester(net::TcpListener);

impl HttpServerTester {
    pub fn on_port(port: u16) -> HttpServerTester {
        let socket = net::TcpListener::bind(("::1", port)).expect("bind");
        let server = HttpServerTester(socket);

        debug!("started HttpServerTester on {}", server.port());

        server
    }

    pub fn new() -> HttpServerTester {
        HttpServerTester::on_port(0)
    }

    pub fn port(&self) -> u16 {
        self.0.local_addr().unwrap().port()
    }

    pub fn accept(&self) -> HttpConnectionTester {
        debug!("accept connection...");
        let r = HttpConnectionTester {
            tcp: self.0.accept().unwrap().0,
            conn: HttpConnection::new(),
            waiting_settings_ack: true,
        };
        debug!("accept connection.");
        r
    }
}

static PREFACE: &'static [u8] = b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n";

pub struct HttpConnectionTester {
    tcp: net::TcpStream,
    conn: HttpConnection,
    waiting_settings_ack: bool,
}

impl HttpConnectionTester {
    pub fn connect(port: u16) -> HttpConnectionTester {
        HttpConnectionTester {
            tcp: net::TcpStream::connect(("::1", port).to_socket_addrs().unwrap().next().unwrap())
                .expect("connect"),
            conn: HttpConnection::new(),
            waiting_settings_ack: true,
        }
    }

    pub fn recv_preface(&mut self) {
        let mut preface = Vec::new();
        preface.resize(PREFACE.len(), 0);
        self.tcp.read_exact(&mut preface).unwrap();
        assert_eq!(PREFACE, &preface[..]);
    }

    pub fn recv_eof(&mut self) {
        assert_eq!(0, self.tcp.read(&mut [0]).expect("read"));
    }

    pub fn send_preface(&mut self) {
        self.tcp.write(PREFACE).expect("send");
    }

    pub fn send_frame<F : FrameIR>(&mut self, frame: F) {
        self.tcp.write(&frame.serialize_into_vec()).expect("send_frame");
    }

    pub fn send_goaway(&mut self, last_stream_id: StreamId) {
        self.send_frame(GoawayFrame::new(last_stream_id, ErrorCode::InadequateSecurity));
    }

    pub fn send_headers(&mut self, stream_id: StreamId, headers: Headers, end: bool) {
        let fragment = self.conn.encoder.encode(headers.0.iter().map(|h| (h.name(), h.value())));
        let mut headers_frame = HeadersFrame::new(fragment, stream_id);
        headers_frame.set_flag(HeadersFlag::EndHeaders);
        if end {
            headers_frame.set_flag(HeadersFlag::EndStream);
        }
        self.send_frame(headers_frame);
    }

    pub fn send_get(&mut self, stream_id: StreamId, path: &str) {
        let mut headers = Headers::new();
        headers.add(":method", "GET");
        headers.add(":path", path);
        self.send_headers(stream_id, headers, true);
    }

    pub fn send_data(&mut self, stream_id: StreamId, data: &[u8], end: bool) {
        let mut data_frame = DataFrame::new(stream_id);
        data_frame.data = Bytes::from(data);
        if end {
            data_frame.set_flag(DataFlag::EndStream);
        }
        self.send_frame(data_frame);
    }

    pub fn send_rst(&mut self, stream_id: StreamId, error_code: ErrorCode) {
        self.send_frame(RstStreamFrame::new(stream_id, error_code));
    }

    pub fn recv_raw_frame(&mut self) -> RawFrame {
        httpbis::solicit_async::recv_raw_frame_sync(
            &mut self.tcp,
            self.conn.peer_settings.max_frame_size)
                .expect("recv_raw_frame")
    }

    pub fn fn_recv_frame_no_check_ack(&mut self) -> HttpFrame {
        let raw_frame = self.recv_raw_frame();
        let frame = HttpFrame::from_raw(&raw_frame).expect("parse frame");
        debug!("received frame: {:?}", frame);
        frame
    }

    pub fn recv_frame(&mut self) -> HttpFrame {
        loop {
            let frame = self.fn_recv_frame_no_check_ack();
            if let HttpFrame::Settings(ref f) = frame {
                if self.waiting_settings_ack && f.is_ack() {
                    self.waiting_settings_ack = false;
                    continue;
                }
                continue;
            }
            return frame;
        }
    }

    pub fn recv_frame_settings(&mut self) -> SettingsFrame {
        match self.fn_recv_frame_no_check_ack() {
            HttpFrame::Settings(settings) => settings,
            f => panic!("expecting SETTINGS, got: {:?}", f),
        }
    }

    pub fn recv_frame_settings_set(&mut self) -> SettingsFrame {
        let settings = self.recv_frame_settings();
        assert!(!settings.is_ack());
        settings
    }

    pub fn recv_frame_settings_ack(&mut self) -> SettingsFrame {
        assert!(self.waiting_settings_ack);
        let settings = self.recv_frame_settings();
        assert!(settings.is_ack());
        self.waiting_settings_ack = false;
        settings
    }

    pub fn get(&mut self, stream_id: StreamId, path: &str) -> SimpleHttpMessage {
        self.send_get(stream_id, path);

        self.recv_message(stream_id)
    }

    // Perform handshape, but do not wait for ACK of my SETTINGS
    // Useful, because ACK may come e.g. after first request HEADERS
    pub fn settings_xchg_but_ack(&mut self) {
        self.send_frame(SettingsFrame::new());
        self.recv_frame_settings_set();
        self.send_frame(SettingsFrame::new_ack());
    }

    pub fn settings_xchg(&mut self) {
        self.settings_xchg_but_ack();
        self.recv_frame_settings_ack();
    }

    pub fn recv_rst_frame(&mut self) -> RstStreamFrame {
        match self.recv_frame() {
            HttpFrame::RstStream(rst) => rst,
            f => panic!("expecting RST, got: {:?}", f),
        }
    }

    pub fn recv_rst_frame_check(&mut self, stream_id: StreamId, error_code: ErrorCode) {
        let frame = self.recv_rst_frame();
        assert_eq!(stream_id, frame.stream_id);
        assert_eq!(error_code, frame.error_code());
    }

    pub fn recv_frame_headers(&mut self) -> HeadersFrame {
        match self.recv_frame() {
            HttpFrame::Headers(headers) => headers,
            f => panic!("expecting HEADERS, got: {:?}", f),
        }
    }

    pub fn recv_frame_data(&mut self) -> DataFrame {
        match self.recv_frame() {
            HttpFrame::Data(data) => data,
            f => panic!("expecting DATA, got: {:?}", f),
        }
    }

    pub fn recv_frame_headers_check(&mut self, stream_id: StreamId, end: bool) -> Headers {
        let headers = self.recv_frame_headers();
        assert_eq!(stream_id, headers.stream_id);
        assert_eq!(end, headers.is_end_of_stream());
        let headers = self.conn.decoder.decode(headers.header_fragment()).expect("decode");
        Headers(headers.into_iter().map(|(n, v)| Header::new(n, v)).collect())
    }

    pub fn recv_frame_data_check(&mut self, stream_id: StreamId, end: bool) -> Vec<u8> {
        let data = self.recv_frame_data();
        assert_eq!(stream_id, data.stream_id);
        assert_eq!(end, data.is_end_of_stream());
        (&data.data[..]).to_vec()
    }

    pub fn recv_frame_data_check_empty_end(&mut self, stream_id: StreamId) {
        let data = self.recv_frame_data_check(stream_id, true);
        assert!(data.is_empty());
    }

    pub fn recv_message(&mut self, stream_id: StreamId) -> SimpleHttpMessage {
        let mut r = SimpleHttpMessage::default();
        loop {
            let frame = self.recv_frame();
            assert_eq!(stream_id, frame.get_stream_id());
            let end_of_stream = match frame {
                HttpFrame::Headers(headers_frame) => {
                    let end_of_stream = headers_frame.is_end_of_stream();
                    let headers = self.conn.decoder.decode(headers_frame.header_fragment()).expect("decode");
                    let headers = Headers(headers.into_iter().map(|(n, v)| Header::new(n, v)).collect());
                    r.headers.extend(headers);
                    end_of_stream
                }
                HttpFrame::Data(data_frame) => {
                    let end_of_stream = data_frame.is_end_of_stream();
                    r.body.extend_from_slice(&data_frame.data);
                    end_of_stream
                }
                frame => panic!("expecting HEADERS or DATA, got: {:?}", frame),
            };
            if end_of_stream {
                return r;
            }
        }
    }
}
<commit_msg>Fix test on Linux<commit_after>#![allow(dead_code)]

use std::io;
use std::io::Write;
use std::io::Read;

use std::str;
use std::net;
use std::net::ToSocketAddrs;

use bytes::Bytes;

use httpbis;
use httpbis::message::SimpleHttpMessage;
use httpbis::solicit::StreamId;
use httpbis::error::ErrorCode;
use httpbis::solicit::header::*;
use httpbis::solicit::frame::FrameIR;
use httpbis::solicit::frame::settings::SettingsFrame;
use httpbis::solicit::frame::headers::HeadersFrame;
use httpbis::solicit::frame::headers::HeadersFlag;
use httpbis::solicit::frame::data::DataFrame;
use httpbis::solicit::frame::data::DataFlag;
use httpbis::solicit::frame::goaway::GoawayFrame;
use httpbis::solicit::frame::RawFrame;
use httpbis::solicit::frame::rst_stream::RstStreamFrame;
use httpbis::solicit::connection::HttpFrame;
use httpbis::solicit::connection::HttpConnection;

pub struct HttpServerTester(net::TcpListener);

impl HttpServerTester {
    pub fn on_port(port: u16) -> HttpServerTester {
        let socket = net::TcpListener::bind(("::1", port)).expect("bind");
        let server = HttpServerTester(socket);

        debug!("started HttpServerTester on {}", server.port());

        server
    }

    pub fn new() -> HttpServerTester {
        HttpServerTester::on_port(0)
    }

    pub fn port(&self) -> u16 {
        self.0.local_addr().unwrap().port()
    }

    pub fn accept(&self) -> HttpConnectionTester {
        debug!("accept connection...");
        let r = HttpConnectionTester {
            tcp: self.0.accept().unwrap().0,
            conn: HttpConnection::new(),
            waiting_settings_ack: true,
        };
        debug!("accept connection.");
        r
    }
}

static PREFACE: &'static [u8] = b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n";

pub struct HttpConnectionTester {
    tcp: net::TcpStream,
    conn: HttpConnection,
    waiting_settings_ack: bool,
}

impl HttpConnectionTester {
    pub fn connect(port: u16) -> HttpConnectionTester {
        HttpConnectionTester {
            tcp: net::TcpStream::connect(("::1", port).to_socket_addrs().unwrap().next().unwrap())
                .expect("connect"),
            conn: HttpConnection::new(),
            waiting_settings_ack: true,
        }
    }

    pub fn recv_preface(&mut self) {
        let mut preface = Vec::new();
        preface.resize(PREFACE.len(), 0);
        self.tcp.read_exact(&mut preface).unwrap();
        assert_eq!(PREFACE, &preface[..]);
    }

    pub fn recv_eof(&mut self) {
        let r = self.tcp.read(&mut [0]);
        match r {
            Ok(0) => {}
            Ok(_) => panic!("expecting EOF"),
            Err(e) => {
                // On Linux it returns ECONNRESET
                if e.kind() != io::ErrorKind::ConnectionReset {
                    panic!("bad error");
                }
            }
        }
    }

    pub fn send_preface(&mut self) {
        self.tcp.write(PREFACE).expect("send");
    }

    pub fn send_frame<F : FrameIR>(&mut self, frame: F) {
        self.tcp.write(&frame.serialize_into_vec()).expect("send_frame");
    }

    pub fn send_goaway(&mut self, last_stream_id: StreamId) {
        self.send_frame(GoawayFrame::new(last_stream_id, ErrorCode::InadequateSecurity));
    }

    pub fn send_headers(&mut self, stream_id: StreamId, headers: Headers, end: bool) {
        let fragment = self.conn.encoder.encode(headers.0.iter().map(|h| (h.name(), h.value())));
        let mut headers_frame = HeadersFrame::new(fragment, stream_id);
        headers_frame.set_flag(HeadersFlag::EndHeaders);
        if end {
            headers_frame.set_flag(HeadersFlag::EndStream);
        }
        self.send_frame(headers_frame);
    }

    pub fn send_get(&mut self, stream_id: StreamId, path: &str) {
        let mut headers = Headers::new();
        headers.add(":method", "GET");
        headers.add(":path", path);
        self.send_headers(stream_id, headers, true);
    }

    pub fn send_data(&mut self, stream_id: StreamId, data: &[u8], end: bool) {
        let mut data_frame = DataFrame::new(stream_id);
        data_frame.data = Bytes::from(data);
        if end {
            data_frame.set_flag(DataFlag::EndStream);
        }
        self.send_frame(data_frame);
    }

    pub fn send_rst(&mut self, stream_id: StreamId, error_code: ErrorCode) {
        self.send_frame(RstStreamFrame::new(stream_id, error_code));
    }

    pub fn recv_raw_frame(&mut self) -> RawFrame {
        httpbis::solicit_async::recv_raw_frame_sync(
            &mut self.tcp,
            self.conn.peer_settings.max_frame_size)
                .expect("recv_raw_frame")
    }

    pub fn fn_recv_frame_no_check_ack(&mut self) -> HttpFrame {
        let raw_frame = self.recv_raw_frame();
        let frame = HttpFrame::from_raw(&raw_frame).expect("parse frame");
        debug!("received frame: {:?}", frame);
        frame
    }

    pub fn recv_frame(&mut self) -> HttpFrame {
        loop {
            let frame = self.fn_recv_frame_no_check_ack();
            if let HttpFrame::Settings(ref f) = frame {
                if self.waiting_settings_ack && f.is_ack() {
                    self.waiting_settings_ack = false;
                    continue;
                }
                continue;
            }
            return frame;
        }
    }

    pub fn recv_frame_settings(&mut self) -> SettingsFrame {
        match self.fn_recv_frame_no_check_ack() {
            HttpFrame::Settings(settings) => settings,
            f => panic!("expecting SETTINGS, got: {:?}", f),
        }
    }

    pub fn recv_frame_settings_set(&mut self) -> SettingsFrame {
        let settings = self.recv_frame_settings();
        assert!(!settings.is_ack());
        settings
    }

    pub fn recv_frame_settings_ack(&mut self) -> SettingsFrame {
        assert!(self.waiting_settings_ack);
        let settings = self.recv_frame_settings();
        assert!(settings.is_ack());
        self.waiting_settings_ack = false;
        settings
    }

    pub fn get(&mut self, stream_id: StreamId, path: &str) -> SimpleHttpMessage {
        self.send_get(stream_id, path);

        self.recv_message(stream_id)
    }

    // Perform handshape, but do not wait for ACK of my SETTINGS
    // Useful, because ACK may come e.g. after first request HEADERS
    pub fn settings_xchg_but_ack(&mut self) {
        self.send_frame(SettingsFrame::new());
        self.recv_frame_settings_set();
        self.send_frame(SettingsFrame::new_ack());
    }

    pub fn settings_xchg(&mut self) {
        self.settings_xchg_but_ack();
        self.recv_frame_settings_ack();
    }

    pub fn recv_rst_frame(&mut self) -> RstStreamFrame {
        match self.recv_frame() {
            HttpFrame::RstStream(rst) => rst,
            f => panic!("expecting RST, got: {:?}", f),
        }
    }

    pub fn recv_rst_frame_check(&mut self, stream_id: StreamId, error_code: ErrorCode) {
        let frame = self.recv_rst_frame();
        assert_eq!(stream_id, frame.stream_id);
        assert_eq!(error_code, frame.error_code());
    }

    pub fn recv_frame_headers(&mut self) -> HeadersFrame {
        match self.recv_frame() {
            HttpFrame::Headers(headers) => headers,
            f => panic!("expecting HEADERS, got: {:?}", f),
        }
    }

    pub fn recv_frame_data(&mut self) -> DataFrame {
        match self.recv_frame() {
            HttpFrame::Data(data) => data,
            f => panic!("expecting DATA, got: {:?}", f),
        }
    }

    pub fn recv_frame_headers_check(&mut self, stream_id: StreamId, end: bool) -> Headers {
        let headers = self.recv_frame_headers();
        assert_eq!(stream_id, headers.stream_id);
        assert_eq!(end, headers.is_end_of_stream());
        let headers = self.conn.decoder.decode(headers.header_fragment()).expect("decode");
        Headers(headers.into_iter().map(|(n, v)| Header::new(n, v)).collect())
    }

    pub fn recv_frame_data_check(&mut self, stream_id: StreamId, end: bool) -> Vec<u8> {
        let data = self.recv_frame_data();
        assert_eq!(stream_id, data.stream_id);
        assert_eq!(end, data.is_end_of_stream());
        (&data.data[..]).to_vec()
    }

    pub fn recv_frame_data_check_empty_end(&mut self, stream_id: StreamId) {
        let data = self.recv_frame_data_check(stream_id, true);
        assert!(data.is_empty());
    }

    pub fn recv_message(&mut self, stream_id: StreamId) -> SimpleHttpMessage {
        let mut r = SimpleHttpMessage::default();
        loop {
            let frame = self.recv_frame();
            assert_eq!(stream_id, frame.get_stream_id());
            let end_of_stream = match frame {
                HttpFrame::Headers(headers_frame) => {
                    let end_of_stream = headers_frame.is_end_of_stream();
                    let headers = self.conn.decoder.decode(headers_frame.header_fragment()).expect("decode");
                    let headers = Headers(headers.into_iter().map(|(n, v)| Header::new(n, v)).collect());
                    r.headers.extend(headers);
                    end_of_stream
                }
                HttpFrame::Data(data_frame) => {
                    let end_of_stream = data_frame.is_end_of_stream();
                    r.body.extend_from_slice(&data_frame.data);
                    end_of_stream
                }
                frame => panic!("expecting HEADERS or DATA, got: {:?}", frame),
            };
            if end_of_stream {
                return r;
            }
        }
    }
}
<|endoftext|>


<commit_before>//! Manages serial port state and communication with clients,
//! and handling requests / responses

use std::collections::HashSet;
use std::sync::mpsc::{Receiver, TryRecvError};
use std::thread;
use std::time::Duration;

use base64;

use errors::*;
use messages::*;
use sub_manager::*;
use port_manager::*;
use writelock_manager::*;
use common::*;

/// Serial port management module supporting one
/// writer and multiple readers
///
/// Clients can lock a port for writing, but
/// subscribe to data from multiple ports for reads
///
/// The Manager takes actions in response to
/// [SerialRequest::*](../messages/index.html) messages sent on
/// on its receiver
pub struct Manager {
  /// Manage write lock status
  writelock_manager: WriteLockManager,
  /// Manage ports
  port_manager: PortManager,
  /// Manage subscriptions
  sub_manager: SubscriptionManager,
  /// Receiver for serial requests
  receiver: Receiver<(String, SerialRequest)>,
  /// Receiver for response subscription requests
  subsc_receiver: SubscReceiver,
}

impl Manager {
  ///Constructor
  pub fn new(
    receiver: Receiver<(String, SerialRequest)>,
    subsc_receiver: SubscReceiver,
  ) -> Manager {
    Manager {
      writelock_manager: WriteLockManager::new(),
      port_manager: PortManager::new(),
      sub_manager: SubscriptionManager::new(),
      receiver: receiver,
      subsc_receiver: subsc_receiver,
    }
  }

  ///Spawn an instance in a new thread.
  pub fn spawn(
    receiver: Receiver<(String, SerialRequest)>,
    subsc_receiver: SubscReceiver,
  ) -> thread::JoinHandle<()> {
    thread::spawn(move || { Manager::new(receiver, subsc_receiver).run(); })
  }

  /// Main loop
  fn run(&mut self) {

    // Bad ports we couldn't read from
    // A set of SerialResponse::Errors built from
    // from the serial port read/write error responses
    let mut bad_ports = HashSet::<String>::new();
    // Shutdown flag
    let mut shutdown = false;
    // Check about 30 times a second
    let sleep_dur = Duration::from_millis(33);

    while !shutdown {
      // Sleep for a little bit to avoid pegging cpu
      // TODO Dynamic timing loop
      thread::sleep(sleep_dur);

      // Handle serial operation requests
      match self.receiver.try_recv() {
        Err(e) => {
          match e {
            TryRecvError::Empty => {
              // nothing to do
            }
            TryRecvError::Disconnected => {
              // Remote end hung up, time to shutdown
              info!("Shutting down SerialPortManager");
              shutdown = true;
            }
          }
        }
        Ok(req) => self.handle_serial_request(&req.0, req.1),
      }

      // Check for new data on each port
      for (port_name, result) in self.port_manager.read_all_ports() {
        match result {
          Ok(data) => {
            let response = match String::from_utf8(data) {
              // We need to send as binary
              Err(e) => {
                SerialResponse::Read {
                  port: port_name.to_string(),
                  data: base64::encode(&e.into_bytes()),
                  base64: Some(true),
                }
              }
              Ok(s) => {
                SerialResponse::Read {
                  port: port_name.to_string(),
                  data: s,
                  base64: Some(false),
                }
              }
            };
            self.broadcast_message_for_port(&port_name, response);
          }
          // Send data reads
          Err(e) => {
            warn!("Error reading port!");
            warn!("{}", e);
            bad_ports.insert(port_name);
          }
        }
      }

      //Handle write requests
      let mut recv_count = 0;
      let mut do_recv = true;
      while do_recv {
        recv_count += 1;
        match self.subsc_receiver.try_recv() {
          Ok(sub_request) => self.sub_manager.add_subscription(sub_request),
          Err(e) => {
            match e {
              TryRecvError::Disconnected => {
                // Does this mean all senders have disconnected?
                // Or just one?
                debug!("Got disconnected when trying to get serial request");
              }
              TryRecvError::Empty => do_recv = false,
            }
          }
        }

        // Don't let this loop get stuck here forever
        // starving the outer loop
        if recv_count == 50 {
          do_recv = false;
        }
      }

      // Cleanup bad serial ports that failed read or write
      // We remove them from everything before
      self.cleanup_bad_ports(&bad_ports);
      bad_ports.clear();
    }
  }

  /// Handles and dispatches SerialRequest sent by
  /// the channel
  fn handle_serial_request(&mut self, sub_id: &String, msg: SerialRequest) {
    let response = match msg {
      SerialRequest::Open { port } => self.handle_open_port(sub_id, port),
      SerialRequest::WriteLock { port } => self.handle_write_lock(sub_id, port),
      SerialRequest::ReleaseWriteLock { port } => self.handle_release_write_lock(sub_id, port),
      SerialRequest::Write { port, data, base64 } => {
        self.handle_write_port(sub_id, port, data, base64.unwrap_or(false))
      }
      SerialRequest::Close { port } => self.handle_close_port(sub_id, port),
      SerialRequest::List {} => self.handle_list_ports(sub_id),
    };
    match response {
      Ok(response) => {
        // Send acknowledge?
      }
      Err(e) => {
        warn!("Error '{}' occured handling serial request message", e);
        // Send error?
        self.send_message(&sub_id, to_serial_response_error(e));
      }
    }
  }

  /// Handle write port requests
  fn handle_write_port(
    &mut self,
    sub_id: &String,
    port_name: String,
    data: String,
    base_64: bool,
  ) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    self.check_owns_writelock(&port_name, &sub_id)?;
    match base_64 {
      true => {
        base64::decode(&data)
          .map_err(|e| ErrorKind::Base64(e).into())
          .and_then(|d| self.port_manager.write_port(&port_name, &d))
          .map(|_| self.send_message(&sub_id, SerialResponse::Wrote { port: port_name }),)
      }
      false => {
        self
          .port_manager
          .write_port(&port_name, data.as_bytes())
          .map(|_| self.send_message(&sub_id, SerialResponse::Wrote { port: port_name }),)
      }
    }
  }

  /// Handle write lock requests
  fn handle_write_lock(&mut self, sub_id: &String, port_name: String) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    self
      .writelock_manager
      .lock_port(&port_name, &sub_id)
      .map(|_| self.send_message(&sub_id, SerialResponse::WriteLocked { port: port_name }),)
  }

  /// Handle write requests
  fn handle_release_write_lock(
    &mut self,
    sub_id: &String,
    port_name: Option<String>,
  ) -> Result<()> {
    self.check_sub_id(sub_id)?;
    match port_name {
      None => {
        self.writelock_manager.unlock_all_ports_for_sub(&sub_id);
        Ok(
          self.send_message(
            &sub_id,
            SerialResponse::WriteLockReleased { port: port_name },
          ),
        )
      }
      Some(port_name) => {
        self
          .writelock_manager
          .unlock_port(&port_name, &sub_id)
          .map(
            |_| {
              self.send_message(
                &sub_id,
                SerialResponse::WriteLockReleased { port: Some(port_name) },
              )
            },
          )
      }
    }
  }

  /// Handle open port requests
  fn handle_open_port(&mut self, sub_id: &String, port_name: String) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    self.port_manager.open_port(&port_name)?;
    self
      .sub_manager
      .add_port(&sub_id, &port_name)
      .map(|_| self.send_message(&sub_id, SerialResponse::Opened { port: port_name }),)
  }

  /// Handle list ports request
  fn handle_list_ports(&mut self, sub_id: &String) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    let port_names: Result<Vec<String>> =
      self
        .port_manager
        .list_ports()
        .map(|v| v.iter().map(|v| v.port_name.clone()).collect());
    let resp = port_names.map(|pns| SerialResponse::List { ports: pns });
    self.send_message(
      &sub_id,
      resp.unwrap_or(
        SerialResponse::Error {
          display: "".to_string(),
          description: "".to_string(),
        },
      ),
    );
    Ok(())
  }

  /// Handle close port requests
  fn handle_close_port(&mut self, sub_id: &String, port_name: Option<String>) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    // Unsubscribe from ports
    match port_name {
      None => {
        self.sub_manager.clear_ports(Some(&sub_id));
        Ok(())
      }
      Some(ref pn) => self.sub_manager.remove_port(&sub_id, &pn),
    }?;
    // Remove write locks
    match port_name {
      None => {
        self.writelock_manager.unlock_all_ports_for_sub(sub_id);
        Ok(())
      }
      Some(ref pn) => self.writelock_manager.unlock_port(&pn, sub_id),
    }?;
    // TODO: Ensure close messages are sent, etc

    // Close ports with no subscribers
    let open_ports = self.port_manager.open_ports();
    let subscribed_ports = self.sub_manager.subscribed_ports();
    let ports_with_no_subs = open_ports.difference(&subscribed_ports);

    // For each open port that isn't subscribed,
    for port_to_close in ports_with_no_subs {
      // close it
      self.port_manager.close_port(&port_to_close);
      // remove the write lock
      self.writelock_manager.clear_lock(&port_to_close);
      // Let them know its closed
      let close_resp = SerialResponse::Closed { port: port_to_close.clone() };
      self.send_message(&sub_id, close_resp);
    }
    Ok(())
  }

  /// Cleanup any bad ports
  fn cleanup_bad_ports(&mut self, bad_ports: &HashSet<String>) {
    for port_name in bad_ports.iter() {
      // Tell everyone port is sick
      let err_resp = to_serial_response_error(ErrorKind::PortReadError(port_name.to_owned()).into(),);
      self.broadcast_message_for_port(port_name, err_resp);
      // Tell everyone the sick ports were closed
      let close_resp = SerialResponse::Closed { port: port_name.clone() };
      self.broadcast_message_for_port(port_name, close_resp);
      // Close bad ports
      self.port_manager.close_port(port_name);
      // Remove write locks on bad ports
      self.writelock_manager.clear_lock(port_name);
      // Remove bad ports from subscriptions
      self.sub_manager.remove_port_from_all(port_name);
    }
  }

  fn send_message(&mut self, sub_id: &String, msg: SerialResponse) {
    match self.sub_manager.send_message(sub_id, msg) {
      Err(e) => {
        warn!("Error sending serial response to sub_id '{}'", sub_id);
        let mut bad_subs = Vec::new();
        bad_subs.push(e);
        self.cleanup_bad_subs(bad_subs);
      }
      _ => {}
    }
  }

  /// Broadcast a message to all subscribers and
  /// then cleanup any subs that errored
  fn broadcast_message(&mut self, msg: SerialResponse) {
    let bad_subs = self.sub_manager.broadcast_message(msg);
    self.cleanup_bad_subs(bad_subs);
  }

  /// Broadcast a message to all subscribers interested in the
  /// given port and then cleanup any subs that errored
  fn broadcast_message_for_port(&mut self, port_name: &String, msg: SerialResponse) {
    let bad_subs = self
      .sub_manager
      .broadcast_message_for_port(port_name, msg);
    self.cleanup_bad_subs(bad_subs);
  }

  /// Cleanup any bad subs where a message send failed
  fn cleanup_bad_subs(&mut self, bad_subs: Vec<Error>) {
    for e in bad_subs {
      match e {
        Error(ErrorKind::SubscriberSendError(sub_id), _) => {
          // Remove subscriptions
          self.sub_manager.end_subscription(&sub_id);
          // Remove all write locks held by dead subscription
          self.writelock_manager.unlock_all_ports_for_sub(&sub_id);
        }
        _ => { /*nop*/ }
      }
    }
  }

  fn check_sub_id(&self, sub_id: &String) -> Result<()> {
    self.sub_manager.check_subscription_exists(sub_id)
  }

  fn check_owns_writelock(&self, port_name: &String, sub_id: &String) -> Result<()> {
    self
      .writelock_manager
      .check_owns_write_lock(port_name, sub_id)
  }
}
<commit_msg>Switch out some matches for more idiomatic if-let<commit_after>//! Manages serial port state and communication with clients,
//! and handling requests / responses

use std::collections::HashSet;
use std::sync::mpsc::{Receiver, TryRecvError};
use std::thread;
use std::time::Duration;

use base64;

use errors::*;
use messages::*;
use sub_manager::*;
use port_manager::*;
use writelock_manager::*;
use common::*;

/// Serial port management module supporting one
/// writer and multiple readers
///
/// Clients can lock a port for writing, but
/// subscribe to data from multiple ports for reads
///
/// The Manager takes actions in response to
/// [SerialRequest::*](../messages/index.html) messages sent on
/// on its receiver
pub struct Manager {
  /// Manage write lock status
  writelock_manager: WriteLockManager,
  /// Manage ports
  port_manager: PortManager,
  /// Manage subscriptions
  sub_manager: SubscriptionManager,
  /// Receiver for serial requests
  receiver: Receiver<(String, SerialRequest)>,
  /// Receiver for response subscription requests
  subsc_receiver: SubscReceiver,
}

impl Manager {
  ///Constructor
  pub fn new(
    receiver: Receiver<(String, SerialRequest)>,
    subsc_receiver: SubscReceiver,
  ) -> Manager {
    Manager {
      writelock_manager: WriteLockManager::new(),
      port_manager: PortManager::new(),
      sub_manager: SubscriptionManager::new(),
      receiver: receiver,
      subsc_receiver: subsc_receiver,
    }
  }

  ///Spawn an instance in a new thread.
  pub fn spawn(
    receiver: Receiver<(String, SerialRequest)>,
    subsc_receiver: SubscReceiver,
  ) -> thread::JoinHandle<()> {
    thread::spawn(move || { Manager::new(receiver, subsc_receiver).run(); })
  }

  /// Main loop
  fn run(&mut self) {

    // Bad ports we couldn't read from
    // A set of SerialResponse::Errors built from
    // from the serial port read/write error responses
    let mut bad_ports = HashSet::<String>::new();
    // Shutdown flag
    let mut shutdown = false;
    // Check about 30 times a second
    let sleep_dur = Duration::from_millis(33);

    while !shutdown {
      // Sleep for a little bit to avoid pegging cpu
      // TODO Dynamic timing loop
      thread::sleep(sleep_dur);

      // Handle serial operation requests
      match self.receiver.try_recv() {
        Err(e) => {
          match e {
            TryRecvError::Empty => {
              // nothing to do
            }
            TryRecvError::Disconnected => {
              // Remote end hung up, time to shutdown
              info!("Shutting down SerialPortManager");
              shutdown = true;
            }
          }
        }
        Ok(req) => self.handle_serial_request(&req.0, req.1),
      }

      // Check for new data on each port
      for (port_name, result) in self.port_manager.read_all_ports() {
        match result {
          Ok(data) => {
            let response = match String::from_utf8(data) {
              // We need to send as binary
              Err(e) => {
                SerialResponse::Read {
                  port: port_name.to_string(),
                  data: base64::encode(&e.into_bytes()),
                  base64: Some(true),
                }
              }
              Ok(s) => {
                SerialResponse::Read {
                  port: port_name.to_string(),
                  data: s,
                  base64: Some(false),
                }
              }
            };
            self.broadcast_message_for_port(&port_name, response);
          }
          // Send data reads
          Err(e) => {
            warn!("Error reading port!");
            warn!("{}", e);
            bad_ports.insert(port_name);
          }
        }
      }

      //Handle write requests
      let mut recv_count = 0;
      let mut do_recv = true;
      while do_recv {
        recv_count += 1;
        match self.subsc_receiver.try_recv() {
          Ok(sub_request) => self.sub_manager.add_subscription(sub_request),
          Err(e) => {
            match e {
              TryRecvError::Disconnected => {
                // Does this mean all senders have disconnected?
                // Or just one?
                debug!("Got disconnected when trying to get serial request");
              }
              TryRecvError::Empty => do_recv = false,
            }
          }
        }

        // Don't let this loop get stuck here forever
        // starving the outer loop
        if recv_count == 50 {
          do_recv = false;
        }
      }

      // Cleanup bad serial ports that failed read or write
      // We remove them from everything before
      self.cleanup_bad_ports(&bad_ports);
      bad_ports.clear();
    }
  }

  /// Handles and dispatches SerialRequest sent by
  /// the channel
  fn handle_serial_request(&mut self, sub_id: &String, msg: SerialRequest) {
    let response = match msg {
      SerialRequest::Open { port } => self.handle_open_port(sub_id, port),
      SerialRequest::WriteLock { port } => self.handle_write_lock(sub_id, port),
      SerialRequest::ReleaseWriteLock { port } => self.handle_release_write_lock(sub_id, port),
      SerialRequest::Write { port, data, base64 } => {
        self.handle_write_port(sub_id, port, data, base64.unwrap_or(false))
      }
      SerialRequest::Close { port } => self.handle_close_port(sub_id, port),
      SerialRequest::List {} => self.handle_list_ports(sub_id),
    };
    if let Err(e) = response {
      warn!("Error '{}' occured handling serial request message", e);
      // Send error?
      self.send_message(&sub_id, to_serial_response_error(e));

    }
  }

  /// Handle write port requests
  fn handle_write_port(
    &mut self,
    sub_id: &String,
    port_name: String,
    data: String,
    base_64: bool,
  ) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    self.check_owns_writelock(&port_name, &sub_id)?;
    match base_64 {
      true => {
        base64::decode(&data)
          .map_err(|e| ErrorKind::Base64(e).into())
          .and_then(|d| self.port_manager.write_port(&port_name, &d))
          .map(|_| self.send_message(&sub_id, SerialResponse::Wrote { port: port_name }),)
      }
      false => {
        self
          .port_manager
          .write_port(&port_name, data.as_bytes())
          .map(|_| self.send_message(&sub_id, SerialResponse::Wrote { port: port_name }),)
      }
    }
  }

  /// Handle write lock requests
  fn handle_write_lock(&mut self, sub_id: &String, port_name: String) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    self
      .writelock_manager
      .lock_port(&port_name, &sub_id)
      .map(|_| self.send_message(&sub_id, SerialResponse::WriteLocked { port: port_name }),)
  }

  /// Handle write requests
  fn handle_release_write_lock(
    &mut self,
    sub_id: &String,
    port_name: Option<String>,
  ) -> Result<()> {
    self.check_sub_id(sub_id)?;
    match port_name {
      None => {
        self.writelock_manager.unlock_all_ports_for_sub(&sub_id);
        Ok(
          self.send_message(
            &sub_id,
            SerialResponse::WriteLockReleased { port: port_name },
          ),
        )
      }
      Some(port_name) => {
        self
          .writelock_manager
          .unlock_port(&port_name, &sub_id)
          .map(
            |_| {
              self.send_message(
                &sub_id,
                SerialResponse::WriteLockReleased { port: Some(port_name) },
              )
            },
          )
      }
    }
  }

  /// Handle open port requests
  fn handle_open_port(&mut self, sub_id: &String, port_name: String) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    self.port_manager.open_port(&port_name)?;
    self
      .sub_manager
      .add_port(&sub_id, &port_name)
      .map(|_| self.send_message(&sub_id, SerialResponse::Opened { port: port_name }),)
  }

  /// Handle list ports request
  fn handle_list_ports(&mut self, sub_id: &String) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    let port_names: Result<Vec<String>> =
      self
        .port_manager
        .list_ports()
        .map(|v| v.iter().map(|v| v.port_name.clone()).collect());
    let resp = port_names.map(|pns| SerialResponse::List { ports: pns });
    self.send_message(
      &sub_id,
      resp.unwrap_or(
        SerialResponse::Error {
          display: "".to_string(),
          description: "".to_string(),
        },
      ),
    );
    Ok(())
  }

  /// Handle close port requests
  fn handle_close_port(&mut self, sub_id: &String, port_name: Option<String>) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    // Unsubscribe from ports
    match port_name {
      None => {
        self.sub_manager.clear_ports(Some(&sub_id));
        Ok(())
      }
      Some(ref pn) => self.sub_manager.remove_port(&sub_id, &pn),
    }?;
    // Remove write locks
    match port_name {
      None => {
        self.writelock_manager.unlock_all_ports_for_sub(sub_id);
        Ok(())
      }
      Some(ref pn) => self.writelock_manager.unlock_port(&pn, sub_id),
    }?;
    // TODO: Ensure close messages are sent, etc

    // Close ports with no subscribers
    let open_ports = self.port_manager.open_ports();
    let subscribed_ports = self.sub_manager.subscribed_ports();
    let ports_with_no_subs = open_ports.difference(&subscribed_ports);

    // For each open port that isn't subscribed,
    for port_to_close in ports_with_no_subs {
      // close it
      self.port_manager.close_port(&port_to_close);
      // remove the write lock
      self.writelock_manager.clear_lock(&port_to_close);
      // Let them know its closed
      let close_resp = SerialResponse::Closed { port: port_to_close.clone() };
      self.send_message(&sub_id, close_resp);
    }
    Ok(())
  }

  /// Cleanup any bad ports
  fn cleanup_bad_ports(&mut self, bad_ports: &HashSet<String>) {
    for port_name in bad_ports.iter() {
      // Tell everyone port is sick
      let err_resp = to_serial_response_error(ErrorKind::PortReadError(port_name.to_owned()).into(),);
      self.broadcast_message_for_port(port_name, err_resp);
      // Tell everyone the sick ports were closed
      let close_resp = SerialResponse::Closed { port: port_name.clone() };
      self.broadcast_message_for_port(port_name, close_resp);
      // Close bad ports
      self.port_manager.close_port(port_name);
      // Remove write locks on bad ports
      self.writelock_manager.clear_lock(port_name);
      // Remove bad ports from subscriptions
      self.sub_manager.remove_port_from_all(port_name);
    }
  }

  fn send_message(&mut self, sub_id: &String, msg: SerialResponse) {
    if let Err(e) = self.sub_manager.send_message(sub_id, msg) {
      warn!("Error sending serial response to sub_id '{}'", sub_id);
      let mut bad_subs = Vec::new();
      bad_subs.push(e);
      self.cleanup_bad_subs(bad_subs);
    }
  }

  /// Broadcast a message to all subscribers and
  /// then cleanup any subs that errored
  fn broadcast_message(&mut self, msg: SerialResponse) {
    let bad_subs = self.sub_manager.broadcast_message(msg);
    self.cleanup_bad_subs(bad_subs);
  }

  /// Broadcast a message to all subscribers interested in the
  /// given port and then cleanup any subs that errored
  fn broadcast_message_for_port(&mut self, port_name: &String, msg: SerialResponse) {
    let bad_subs = self
      .sub_manager
      .broadcast_message_for_port(port_name, msg);
    self.cleanup_bad_subs(bad_subs);
  }

  /// Cleanup any bad subs where a message send failed
  fn cleanup_bad_subs(&mut self, bad_subs: Vec<Error>) {
    for e in bad_subs {
      if let Error(ErrorKind::SubscriberSendError(sub_id), _) = e {
        // Remove subscriptions
        self.sub_manager.end_subscription(&sub_id);
        // Remove all write locks held by dead subscription
        self.writelock_manager.unlock_all_ports_for_sub(&sub_id);
      }
    }
  }

  fn check_sub_id(&self, sub_id: &String) -> Result<()> {
    self.sub_manager.check_subscription_exists(sub_id)
  }

  fn check_owns_writelock(&self, port_name: &String, sub_id: &String) -> Result<()> {
    self
      .writelock_manager
      .check_owns_write_lock(port_name, sub_id)
  }
}
<|endoftext|>


<commit_before>
macro_rules! return_if {
  ($e:expr, $r:expr) => {
    if $e {
      return $r;
    }
  }
}

pub struct Square {
  a: usize,
  data: Vec<usize>
}

pub enum Diagonal {
  First,
  Second
}

impl From<usize> for Diagonal {
  fn from(d: usize) -> Diagonal {
    match d {
      0 => Diagonal::First,
      1 => Diagonal::Second,
      _ => panic!()
    }
  }
}

pub enum SquareLine {
  Row(usize),
  Col(usize),
  Diag(Diagonal)
}

impl SquareLine {
  pub fn sum(&self, square: &Square) -> usize {
    square.sum_of(self)
  }
}

pub struct LineIter {
  a: usize,
  i: usize,
  is_col: bool
}

impl Iterator for LineIter {
  type Item = SquareLine;

  fn next(&mut self) -> Option<SquareLine> {
    if self.i >= self.a + 2 {
      return None;
    }

    if self.i >= self.a {
      let diag = self.i - self.a;
      self.i = self.i + 1;
      return Some(SquareLine::Diag(Diagonal::from(diag)));
    }

    if self.is_col {
      let result = SquareLine::Col(self.i);

      self.is_col = false;
      self.i = self.i + 1;

      return Some(result);
    } else {
      let result = SquareLine::Row(self.i);

      self.is_col = true;

      return Some(result);
    }
  }
}

impl<'a> From<&'a Square> for LineIter {
  fn from(square: &'a Square) -> LineIter {
    LineIter {
      a: square.a,
      i: 0,
      is_col: false
    }
  }
}

pub struct SquareIter<'a> {
  square: &'a Square,
  start: usize,
  jump: usize,
  i: usize
}

impl<'a> Iterator for SquareIter<'a> {
  type Item = &'a usize;

  fn next(&mut self) -> Option<&'a usize> {

    if self.i >= self.square.a {
      return None;
    }

    let value = &self.square.data[self.start + self.i * self.jump];
    self.i = self.i + 1;

    Some(value)
  }
}

impl Square {
  pub fn new(data: Vec<usize>) -> Square {
    let len = data.len();
    let a_candidate = (len as f64).sqrt().floor() as usize;

    assert_eq!(a_candidate * a_candidate, len);

    Square {
      a: a_candidate,
      data: data
    }
  }

  pub fn get_expected_sum(&self) -> usize {
    (self.a * ( (self.a * self.a) + 1)) / 2
  }

  pub fn is_magic(&self) -> bool {
    let expected_sum = self.get_expected_sum();

    for line in LineIter::from(self) {
      let sum = line.sum(self);
      return_if!(sum != expected_sum, false);
    }

    return true;
  }

  pub fn sum_of(&self, line: &SquareLine) -> usize {
    self.get_iter(line).fold(0, |acc, &x| acc + x)
  }

  pub fn get_iter<'a, 'b>(&'a self, line: &'b SquareLine) -> SquareIter<'a> {
    match line {
      &SquareLine::Row(r) => self.row(r),
      &SquareLine::Col(c) => self.col(c),
      &SquareLine::Diag(ref d) => self.diag(d)
    }
  }

  pub fn row<'a>(&'a self, r: usize) -> SquareIter<'a> {
    assert!(r < self.a);

    SquareIter {
      square: self,
      start: r * self.a,
      jump: 1,
      i: 0
    }
  }

  pub fn col<'a>(&'a self, c: usize) -> SquareIter<'a> {
    assert!(c < self.a);

    SquareIter {
      square: self,
      start: c,
      jump: self.a,
      i: 0
    }
  }

  pub fn diag<'a>(&'a self, d: &Diagonal) -> SquareIter<'a> {
    match d {
      &Diagonal::First => SquareIter {
        square: self,
        start: 0,
        jump: self.a + 1,
        i: 0
      },

      &Diagonal::Second => SquareIter {
        square: self,
        start: self.a - 1,
        jump: self.a - 1,
        i: 0
      }
    }
  }
}

#[cfg(test)]
mod test {

  use super::*;

  fn test_square() -> Square {
    Square::new(vec![8, 1, 6, 3, 5, 7, 4, 9, 2])
  }

  #[test]
  #[should_panic]
  fn test_non_square() {
    Square::new(vec![1, 2]);
  }

  #[test]
  fn test_row_iterator() {
    let square = test_square();

    assert_eq!(vec![8, 1, 6], square.row(0).map(|v| v.clone()).collect::<Vec<usize>>());
  }

  #[test]
  fn test_col_iterator() {
    let square = test_square();

    assert_eq!(vec![8, 3, 4], square.col(0).map(|v| v.clone()).collect::<Vec<usize>>());
  }

  #[test]
  fn test_diag_iterator() {
    let square = test_square();

    assert_eq!(vec![8, 5, 2], square.diag(&Diagonal::First).map(|v| v.clone()).collect::<Vec<usize>>());
    assert_eq!(vec![6, 5, 4], square.diag(&Diagonal::Second).map(|v| v.clone()).collect::<Vec<usize>>());
  }

  #[test]
  fn test_first_example() {
    let square = Square::new(vec![8, 1, 6, 3, 5, 7, 4, 9, 2]);

    assert_eq!(true, square.is_magic());
  }

  #[test]
  fn test_second_example() {
    let square = Square::new(vec![2, 7, 6, 9, 5, 1, 4, 3, 8]);

    assert_eq!(true, square.is_magic());
  }

  #[test]
  fn test_third_example() {
    let square = Square::new(vec![3, 5, 7, 8, 1, 6, 4, 9, 2]);

    assert_eq!(false, square.is_magic());
  }

  #[test]
  fn test_fourth_example() {
    let square = Square::new(vec![8, 1, 6, 7, 5, 3, 4, 9, 2]);

    assert_eq!(false, square.is_magic());
  }

  #[test]
  fn test_big() {
    let square = Square::new(vec![25, 13, 1, 19, 7, 16, 9, 22, 15, 3, 12, 5,
            18, 6, 24, 8, 21, 14, 2, 20, 4, 17, 10, 23, 11]);

    assert_eq!(true, square.is_magic());
  }
}
<commit_msg>bonus 2 dp261<commit_after>
macro_rules! return_if {
  ($e:expr, $r:expr) => {
    if $e {
      return $r;
    }
  }
}

pub struct Square {
  a: usize,
  data: Vec<usize>
}

pub enum Diagonal {
  First,
  Second
}

impl From<usize> for Diagonal {
  fn from(d: usize) -> Diagonal {
    match d {
      0 => Diagonal::First,
      1 => Diagonal::Second,
      _ => panic!()
    }
  }
}

pub enum SquareLine {
  Row(usize),
  Col(usize),
  Diag(Diagonal)
}

impl SquareLine {
  pub fn sum(&self, square: &Square) -> usize {
    square.sum_of(self)
  }
}

pub struct LineIter {
  a: usize,
  i: usize,
  is_col: bool
}

impl Iterator for LineIter {
  type Item = SquareLine;

  fn next(&mut self) -> Option<SquareLine> {
    if self.i >= self.a + 2 {
      return None;
    }

    if self.i >= self.a {
      let diag = self.i - self.a;
      self.i = self.i + 1;
      return Some(SquareLine::Diag(Diagonal::from(diag)));
    }

    if self.is_col {
      let result = SquareLine::Col(self.i);

      self.is_col = false;
      self.i = self.i + 1;

      return Some(result);
    } else {
      let result = SquareLine::Row(self.i);

      self.is_col = true;

      return Some(result);
    }
  }
}

impl<'a> From<&'a Square> for LineIter {
  fn from(square: &'a Square) -> LineIter {
    LineIter {
      a: square.a,
      i: 0,
      is_col: false
    }
  }
}

pub struct SquareIter<'a> {
  square: &'a Square,
  start: usize,
  jump: usize,
  i: usize
}

impl<'a> Iterator for SquareIter<'a> {
  type Item = &'a usize;

  fn next(&mut self) -> Option<&'a usize> {

    if self.i >= self.square.a {
      return None;
    }

    let value = &self.square.data[self.start + self.i * self.jump];
    self.i = self.i + 1;

    Some(value)
  }
}

impl From<Vec<usize>> for Square {
  fn from(data: Vec<usize>) -> Square {
    let len = data.len();
    let a_candidate = (len as f64).sqrt().floor() as usize;

    assert_eq!(a_candidate * a_candidate, len);

    Square {
      a: a_candidate,
      data: data
    }
  }
}

impl Square {

  pub fn try_new(data: Vec<usize>) -> Option<Square> {
    let len = data.len();
    let s = ((4 * len + 1) as f64).sqrt() as usize;

    let a_candidate = (s+1)/2;

    assert_eq!(a_candidate * (a_candidate - 1), len);

    return Square::try_complete(&data, &a_candidate);
  }

  fn try_complete(incomplete_data: &Vec<usize>, a: &usize) -> Option<Square> {
    let mut try_data = incomplete_data.clone();

    let begin_len = try_data.len();

    let num_missing = (a*a) - begin_len;

    let mut work_data = vec![1; num_missing];

    assert_eq!(a * a, try_data.len() + work_data.len());

    let mut done = false;

    while !done {
      try_data.resize( begin_len, 0 );
      try_data.extend_from_slice( &work_data );

      assert_eq!(a * a, try_data.len());

      {
        let square = Square::from(try_data.clone());

        if square.is_magic() {
          return Some(square);
        }
      }

      let mut cur = 0;

      loop {
        if cur == work_data.len() {
          done = true;
          break;
        }

        work_data[cur] = work_data[cur] + 1;
        if work_data[cur] > 9 {
          work_data[cur] = 1;
          cur = cur + 1;
        } else {
          break;
        }
      }
    }

    None
  }

  pub fn get_expected_sum(&self) -> usize {
    (self.a * ( (self.a * self.a) + 1)) / 2
  }

  pub fn is_magic(&self) -> bool {
    let expected_sum = self.get_expected_sum();

    for line in LineIter::from(self) {
      let sum = line.sum(self);
      return_if!(sum != expected_sum, false);
    }

    return true;
  }

  pub fn sum_of(&self, line: &SquareLine) -> usize {
    self.get_iter(line).fold(0, |acc, &x| acc + x)
  }

  pub fn get_iter<'a, 'b>(&'a self, line: &'b SquareLine) -> SquareIter<'a> {
    match line {
      &SquareLine::Row(r) => self.row(r),
      &SquareLine::Col(c) => self.col(c),
      &SquareLine::Diag(ref d) => self.diag(d)
    }
  }

  pub fn row<'a>(&'a self, r: usize) -> SquareIter<'a> {
    assert!(r < self.a);

    SquareIter {
      square: self,
      start: r * self.a,
      jump: 1,
      i: 0
    }
  }

  pub fn col<'a>(&'a self, c: usize) -> SquareIter<'a> {
    assert!(c < self.a);

    SquareIter {
      square: self,
      start: c,
      jump: self.a,
      i: 0
    }
  }

  pub fn diag<'a>(&'a self, d: &Diagonal) -> SquareIter<'a> {
    match d {
      &Diagonal::First => SquareIter {
        square: self,
        start: 0,
        jump: self.a + 1,
        i: 0
      },

      &Diagonal::Second => SquareIter {
        square: self,
        start: self.a - 1,
        jump: self.a - 1,
        i: 0
      }
    }
  }

  pub fn render(&self) {
    println!("");
    for i in 0..self.a {
      println!("{:?}", self.row(i).collect::<Vec<&usize>>());
    }
  }

}

#[cfg(test)]
mod test {

  use super::*;

  fn test_square() -> Square {
    Square::from(vec![8, 1, 6, 3, 5, 7, 4, 9, 2])
  }

  #[test]
  #[should_panic]
  fn test_non_square() {
    Square::from(vec![1, 2]);
  }

  #[test]
  fn test_row_iterator() {
    let square = test_square();

    assert_eq!(vec![8, 1, 6], square.row(0).map(|v| v.clone()).collect::<Vec<usize>>());
  }

  #[test]
  fn test_col_iterator() {
    let square = test_square();

    assert_eq!(vec![8, 3, 4], square.col(0).map(|v| v.clone()).collect::<Vec<usize>>());
  }

  #[test]
  fn test_diag_iterator() {
    let square = test_square();

    assert_eq!(vec![8, 5, 2], square.diag(&Diagonal::First).map(|v| v.clone()).collect::<Vec<usize>>());
    assert_eq!(vec![6, 5, 4], square.diag(&Diagonal::Second).map(|v| v.clone()).collect::<Vec<usize>>());
  }

  #[test]
  fn test_first_example() {
    let square = Square::from(vec![8, 1, 6, 3, 5, 7, 4, 9, 2]);

    assert!(square.is_magic());
  }

  #[test]
  fn test_second_example() {
    let square = Square::from(vec![2, 7, 6, 9, 5, 1, 4, 3, 8]);

    assert!(square.is_magic());
  }

  #[test]
  fn test_third_example() {
    let square = Square::from(vec![3, 5, 7, 8, 1, 6, 4, 9, 2]);

    assert!(!square.is_magic());
  }

  #[test]
  fn test_fourth_example() {
    let square = Square::from(vec![8, 1, 6, 7, 5, 3, 4, 9, 2]);

    assert!(!square.is_magic());
  }

  #[test]
  fn test_big() {
    let square = Square::from(vec![25, 13, 1, 19, 7, 16, 9, 22, 15, 3, 12, 5,
            18, 6, 24, 8, 21, 14, 2, 20, 4, 17, 10, 23, 11]);

    assert!(square.is_magic());
  }

  #[test]
  fn test_complete_1() {
    let square = Square::try_new(vec![8, 1, 6, 3, 5, 7]).unwrap();
    assert!(square.is_magic());

    square.render();
  }

  #[test]
  #[should_panic]
  fn test_complete_2() {
    Square::try_new(vec![3, 5, 7, 8, 1, 6]).unwrap();
  }
}
<|endoftext|>