<commit_before>//! A collection of newtypes defining type-strong IDs.

use std::fmt::{Display, Formatter, Result as FmtResult};

use chrono::{DateTime, TimeZone, Utc};

macro_rules! id_u64 {
    ($($name:ident;)*) => {
        $(
            impl $name {
                /// Retrieves the time that the Id was created at.
                pub fn created_at(&self) -> DateTime<Utc> {
                    const DISCORD_EPOCH: u64 = 1_420_070_400_000;
                    Utc.timestamp_millis(((self.0 >> 22) + DISCORD_EPOCH) as i64)
                }

                /// Immutably borrow inner Id.
                #[inline]
                pub fn as_u64(&self) -> &u64 {
                    &self.0
                }

                /// Mutably borrow inner Id.
                #[inline]
                pub fn as_mut_u64(&mut self) -> &mut u64 {
                    &mut self.0
                }
            }

            // This is a hack so functions can accept iterators that either:
            // 1. return the id itself (e.g: `MessageId`)
            // 2. return a reference to it (`&MessageId`).
            impl AsRef<$name> for $name {
                fn as_ref(&self) -> &Self {
                    self
                }
            }

            impl<'a> From<&'a $name> for $name {
                fn from(id: &'a $name) -> $name {
                    id.clone()
                }
            }

            impl From<u64> for $name {
                fn from(id_as_u64: u64) -> $name {
                    $name(id_as_u64)
                }
            }

            impl PartialEq<u64> for $name {
                fn eq(&self, u: &u64) -> bool {
                    self.0 == *u
                }
            }

            impl Display for $name {
                fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
                    Display::fmt(&self.0, f)
                }
            }

            impl From<$name> for u64 {
                fn from(id: $name) -> u64 {
                    id.0 as u64
                }
            }

            impl From<$name> for i64 {
                fn from(id: $name) -> i64 {
                    id.0 as i64
                }
            }
        )*
    }
}

/// An identifier for an Application.
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct ApplicationId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a Channel
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct ChannelId(#[serde(with = "snowflake")] pub u64);

/// An identifier for an Emoji
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct EmojiId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a Guild
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct GuildId(#[serde(with = "snowflake")] pub u64);

/// An identifier for an Integration
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct IntegrationId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a Message
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct MessageId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a Role
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct RoleId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a User
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct UserId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a [`Webhook`][super::webhook::Webhook]
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct WebhookId(#[serde(with = "snowflake")] pub u64);

/// An identifier for an audit log entry.
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct AuditLogEntryId(#[serde(with = "snowflake")] pub u64);

/// An identifier for an attachment.
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct AttachmentId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a sticker.
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct StickerId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a sticker pack.
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct StickerPackId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a sticker pack banner.
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct StickerPackBannerId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a SKU.
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct SkuId(#[serde(with = "snowflake")] pub u64);

/// An identifier for an interaction.
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct InteractionId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a slash command.
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct CommandId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a slash command permission Id. Can contain
/// a [`RoleId`] or [`UserId`].
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct CommandPermissionId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a slash command version Id.
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct CommandVersionId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a slash command target Id. Can contain
/// a [`UserId`] or [`MessageId`].
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct TargetId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a stage channel instance.
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct StageInstanceId(#[serde(with = "snowflake")] pub u64);

id_u64! {
    AttachmentId;
    ApplicationId;
    ChannelId;
    EmojiId;
    GuildId;
    IntegrationId;
    MessageId;
    RoleId;
    StickerId;
    StickerPackId;
    StickerPackBannerId;
    SkuId;
    UserId;
    WebhookId;
    AuditLogEntryId;
    InteractionId;
    CommandId;
    CommandPermissionId;
    CommandVersionId;
    TargetId;
    StageInstanceId;
}

/// Used with `#[serde(with|deserialize_with|serialize_with)]`
///
/// # Examples
///
/// ```rust,ignore
/// #[derive(Deserialize, Serialize)]
/// struct A {
///     #[serde(with = "snowflake")]
///     id: u64,
/// }
///
/// #[derive(Deserialize)]
/// struct B {
///     #[serde(deserialize_with = "snowflake::deserialize")]
///     id: u64,
/// }
///
/// #[derive(Serialize)]
/// struct C {
///     #[serde(serialize_with = "snowflake::serialize")]
///     id: u64,
/// }
/// ```
pub(crate) mod snowflake {
    use std::fmt;

    use serde::de::{Error, MapAccess, Visitor};
    use serde::{Deserialize, Deserializer, Serializer};

    pub fn deserialize<'de, D: Deserializer<'de>>(deserializer: D) -> Result<u64, D::Error> {
        deserializer.deserialize_any(SnowflakeVisitor)
    }

    #[allow(clippy::trivially_copy_pass_by_ref)]
    pub fn serialize<S: Serializer>(id: &u64, serializer: S) -> Result<S::Ok, S::Error> {
        serializer.collect_str(id)
    }

    struct SnowflakeVisitor;

    impl<'de> Visitor<'de> for SnowflakeVisitor {
        type Value = u64;

        fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
            formatter.write_str("string or integer snowflake")
        }

        fn visit_u64<E: Error>(self, value: u64) -> Result<Self::Value, E> {
            Ok(value)
        }

        fn visit_str<E: Error>(self, value: &str) -> Result<Self::Value, E> {
            value.parse().map_err(Error::custom)
        }

        // This is called when serde_json's `arbitrary_precision` feature is enabled.
        fn visit_map<A: MapAccess<'de>>(self, mut map: A) -> Result<Self::Value, A::Error> {
            let id = map.next_value::<Snowflake>()?;
            Ok(id.value)
        }
    }

    struct Snowflake {
        value: u64,
    }

    impl<'de> Deserialize<'de> for Snowflake {
        fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
            struct StrVisitor;

            impl<'de> Visitor<'de> for StrVisitor {
                type Value = Snowflake;

                fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
                    formatter.write_str("string containing a number")
                }

                fn visit_str<E: Error>(self, s: &str) -> Result<Snowflake, E> {
                    let value = s.parse().map_err(Error::custom)?;
                    Ok(Snowflake {
                        value,
                    })
                }
            }

            deserializer.deserialize_str(StrVisitor)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::GuildId;

    #[test]
    fn test_created_at() {
        // The id is from discord's snowflake docs
        assert_eq!(
            GuildId(175928847299117063).created_at().to_rfc3339(),
            "2016-04-30T11:18:25.796+00:00"
        );
    }

    #[test]
    fn test_id_serde() {
        use serde::{Deserialize, Serialize};
        use serde_test::{assert_de_tokens, assert_tokens, Token};

        use super::snowflake;

        let id = GuildId(17_5928_8472_9911_7063);
        assert_tokens(&id, &[
            Token::NewtypeStruct {
                name: "GuildId",
            },
            Token::Str("175928847299117063"),
        ]);
        assert_de_tokens(&id, &[
            Token::NewtypeStruct {
                name: "GuildId",
            },
            Token::U64(17_5928_8472_9911_7063),
        ]);

        #[derive(Debug, PartialEq, Deserialize, Serialize)]
        struct S {
            #[serde(with = "snowflake")]
            id: u64,
        }

        let s = S {
            id: 17_5928_8472_9911_7063,
        };
        assert_tokens(&s, &[
            Token::Struct {
                name: "S",
                len: 1,
            },
            Token::Str("id"),
            Token::Str("175928847299117063"),
            Token::StructEnd,
        ]);

        #[derive(Debug, PartialEq, Deserialize, Serialize)]
        struct Opt {
            id: Option<GuildId>,
        }

        let s = Opt {
            id: Some(GuildId(17_5928_8472_9911_7063)),
        };
        assert_tokens(&s, &[
            Token::Struct {
                name: "Opt",
                len: 1,
            },
            Token::Str("id"),
            Token::Some,
            Token::NewtypeStruct {
                name: "GuildId",
            },
            Token::Str("175928847299117063"),
            Token::StructEnd,
        ]);
    }
}
<commit_msg>Fix ID type deserialization from signed integers (#1668)<commit_after>//! A collection of newtypes defining type-strong IDs.

use std::fmt::{Display, Formatter, Result as FmtResult};

use chrono::{DateTime, TimeZone, Utc};

macro_rules! id_u64 {
    ($($name:ident;)*) => {
        $(
            impl $name {
                /// Retrieves the time that the Id was created at.
                pub fn created_at(&self) -> DateTime<Utc> {
                    const DISCORD_EPOCH: u64 = 1_420_070_400_000;
                    Utc.timestamp_millis(((self.0 >> 22) + DISCORD_EPOCH) as i64)
                }

                /// Immutably borrow inner Id.
                #[inline]
                pub fn as_u64(&self) -> &u64 {
                    &self.0
                }

                /// Mutably borrow inner Id.
                #[inline]
                pub fn as_mut_u64(&mut self) -> &mut u64 {
                    &mut self.0
                }
            }

            // This is a hack so functions can accept iterators that either:
            // 1. return the id itself (e.g: `MessageId`)
            // 2. return a reference to it (`&MessageId`).
            impl AsRef<$name> for $name {
                fn as_ref(&self) -> &Self {
                    self
                }
            }

            impl<'a> From<&'a $name> for $name {
                fn from(id: &'a $name) -> $name {
                    id.clone()
                }
            }

            impl From<u64> for $name {
                fn from(id_as_u64: u64) -> $name {
                    $name(id_as_u64)
                }
            }

            impl PartialEq<u64> for $name {
                fn eq(&self, u: &u64) -> bool {
                    self.0 == *u
                }
            }

            impl Display for $name {
                fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
                    Display::fmt(&self.0, f)
                }
            }

            impl From<$name> for u64 {
                fn from(id: $name) -> u64 {
                    id.0 as u64
                }
            }

            impl From<$name> for i64 {
                fn from(id: $name) -> i64 {
                    id.0 as i64
                }
            }
        )*
    }
}

/// An identifier for an Application.
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct ApplicationId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a Channel
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct ChannelId(#[serde(with = "snowflake")] pub u64);

/// An identifier for an Emoji
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct EmojiId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a Guild
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct GuildId(#[serde(with = "snowflake")] pub u64);

/// An identifier for an Integration
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct IntegrationId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a Message
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct MessageId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a Role
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct RoleId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a User
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct UserId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a [`Webhook`][super::webhook::Webhook]
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct WebhookId(#[serde(with = "snowflake")] pub u64);

/// An identifier for an audit log entry.
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct AuditLogEntryId(#[serde(with = "snowflake")] pub u64);

/// An identifier for an attachment.
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct AttachmentId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a sticker.
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct StickerId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a sticker pack.
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct StickerPackId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a sticker pack banner.
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct StickerPackBannerId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a SKU.
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct SkuId(#[serde(with = "snowflake")] pub u64);

/// An identifier for an interaction.
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct InteractionId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a slash command.
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct CommandId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a slash command permission Id. Can contain
/// a [`RoleId`] or [`UserId`].
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct CommandPermissionId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a slash command version Id.
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct CommandVersionId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a slash command target Id. Can contain
/// a [`UserId`] or [`MessageId`].
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct TargetId(#[serde(with = "snowflake")] pub u64);

/// An identifier for a stage channel instance.
#[derive(
    Copy, Clone, Default, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Deserialize, Serialize,
)]
pub struct StageInstanceId(#[serde(with = "snowflake")] pub u64);

id_u64! {
    AttachmentId;
    ApplicationId;
    ChannelId;
    EmojiId;
    GuildId;
    IntegrationId;
    MessageId;
    RoleId;
    StickerId;
    StickerPackId;
    StickerPackBannerId;
    SkuId;
    UserId;
    WebhookId;
    AuditLogEntryId;
    InteractionId;
    CommandId;
    CommandPermissionId;
    CommandVersionId;
    TargetId;
    StageInstanceId;
}

/// Used with `#[serde(with|deserialize_with|serialize_with)]`
///
/// # Examples
///
/// ```rust,ignore
/// #[derive(Deserialize, Serialize)]
/// struct A {
///     #[serde(with = "snowflake")]
///     id: u64,
/// }
///
/// #[derive(Deserialize)]
/// struct B {
///     #[serde(deserialize_with = "snowflake::deserialize")]
///     id: u64,
/// }
///
/// #[derive(Serialize)]
/// struct C {
///     #[serde(serialize_with = "snowflake::serialize")]
///     id: u64,
/// }
/// ```
pub(crate) mod snowflake {
    use std::convert::TryFrom;
    use std::fmt;

    use serde::de::{Error, MapAccess, Visitor};
    use serde::{Deserialize, Deserializer, Serializer};

    pub fn deserialize<'de, D: Deserializer<'de>>(deserializer: D) -> Result<u64, D::Error> {
        deserializer.deserialize_any(SnowflakeVisitor)
    }

    #[allow(clippy::trivially_copy_pass_by_ref)]
    pub fn serialize<S: Serializer>(id: &u64, serializer: S) -> Result<S::Ok, S::Error> {
        serializer.collect_str(id)
    }

    struct SnowflakeVisitor;

    impl<'de> Visitor<'de> for SnowflakeVisitor {
        type Value = u64;

        fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
            formatter.write_str("string or integer snowflake")
        }

        // Called by formats like TOML.
        fn visit_i64<E: Error>(self, value: i64) -> Result<Self::Value, E> {
            u64::try_from(value).map_err(Error::custom)
        }

        fn visit_u64<E: Error>(self, value: u64) -> Result<Self::Value, E> {
            Ok(value)
        }

        fn visit_str<E: Error>(self, value: &str) -> Result<Self::Value, E> {
            value.parse().map_err(Error::custom)
        }

        // This is called when serde_json's `arbitrary_precision` feature is enabled.
        fn visit_map<A: MapAccess<'de>>(self, mut map: A) -> Result<Self::Value, A::Error> {
            let id = map.next_value::<Snowflake>()?;
            Ok(id.value)
        }
    }

    struct Snowflake {
        value: u64,
    }

    impl<'de> Deserialize<'de> for Snowflake {
        fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
            struct StrVisitor;

            impl<'de> Visitor<'de> for StrVisitor {
                type Value = Snowflake;

                fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
                    formatter.write_str("string containing a number")
                }

                fn visit_str<E: Error>(self, s: &str) -> Result<Snowflake, E> {
                    let value = s.parse().map_err(Error::custom)?;
                    Ok(Snowflake {
                        value,
                    })
                }
            }

            deserializer.deserialize_str(StrVisitor)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::GuildId;

    #[test]
    fn test_created_at() {
        // The id is from discord's snowflake docs
        assert_eq!(
            GuildId(175928847299117063).created_at().to_rfc3339(),
            "2016-04-30T11:18:25.796+00:00"
        );
    }

    #[test]
    fn test_id_serde() {
        use serde::{Deserialize, Serialize};
        use serde_test::{assert_de_tokens, assert_tokens, Token};

        use super::snowflake;

        let id = GuildId(17_5928_8472_9911_7063);
        assert_tokens(&id, &[
            Token::NewtypeStruct {
                name: "GuildId",
            },
            Token::Str("175928847299117063"),
        ]);
        assert_de_tokens(&id, &[
            Token::NewtypeStruct {
                name: "GuildId",
            },
            Token::U64(17_5928_8472_9911_7063),
        ]);

        #[derive(Debug, PartialEq, Deserialize, Serialize)]
        struct S {
            #[serde(with = "snowflake")]
            id: u64,
        }

        let s = S {
            id: 17_5928_8472_9911_7063,
        };
        assert_tokens(&s, &[
            Token::Struct {
                name: "S",
                len: 1,
            },
            Token::Str("id"),
            Token::Str("175928847299117063"),
            Token::StructEnd,
        ]);

        #[derive(Debug, PartialEq, Deserialize, Serialize)]
        struct Opt {
            id: Option<GuildId>,
        }

        let s = Opt {
            id: Some(GuildId(17_5928_8472_9911_7063)),
        };
        assert_tokens(&s, &[
            Token::Struct {
                name: "Opt",
                len: 1,
            },
            Token::Str("id"),
            Token::Some,
            Token::NewtypeStruct {
                name: "GuildId",
            },
            Token::Str("175928847299117063"),
            Token::StructEnd,
        ]);
    }
}
<|endoftext|>


<commit_before>use std::thread;
use std::thread::JoinHandle;
use std::sync::mpsc::{TryRecvError, channel};
use std::time::Duration;
use std::io;
use std::sync::atomic::{AtomicUsize, Ordering};

use byteorder::{ByteOrder, LittleEndian};
use mio::tcp::{TcpListener, TcpStream};

use net::EventLoop;

/// The amount of time to wait while unit testing to make sure threads are syncronized when there is no other method.
const WAIT_TIME_MS: u64 = 250;

pub static CLIENT_SERVER_SEND_TEST_VAL: AtomicUsize = AtomicUsize::new(0);

pub static EVENT_LOOP_SEND_TEST_VAL: AtomicUsize = AtomicUsize::new(0);

#[derive(Serialize, Deserialize, Debug, Clone, Copy)]
pub enum TestValToModify {
    ClientServerSend,
    EventLoopImplSend,
}

#[test]
fn client_server_connect() {
    let (mut event_loop_ref_client, thread_client) = event_loop_helper();
    let (event_loop_ref_server, thread_server) = event_loop_helper();
    let listener = TcpListener::bind(&super::ip("127.0.0.1:25570")).unwrap();
    event_loop_ref_server.add_listener(listener);
    let client = super::client::Client::spawn_client(super::ip("127.0.0.1:25570"), &event_loop_ref_client).unwrap();
    client.shutdown(&mut event_loop_ref_client);
    event_loop_ref_client.shutdown();
    event_loop_ref_server.shutdown();
    let _event_loop_client = thread_client.join().unwrap();
    let _event_loop_server = thread_server.join().unwrap();
}

#[test]
fn client_server_send() {
    let (mut event_loop_ref_client, thread_client) = event_loop_helper();
    let (event_loop_ref_server, thread_server) = event_loop_helper();
    let listener = TcpListener::bind(&super::ip("127.0.0.1:25571")).unwrap();
    event_loop_ref_server.add_listener(listener);
    let client = super::client::Client::spawn_client(super::ip("127.0.0.1:25571"), &event_loop_ref_client).unwrap();
    let old_test_val = CLIENT_SERVER_SEND_TEST_VAL.load(Ordering::Relaxed);

    client.send(&event_loop_ref_client,
                super::NetworkPacket::Test(TestValToModify::ClientServerSend));
    thread::sleep(Duration::from_millis(WAIT_TIME_MS));

    client.shutdown(&mut event_loop_ref_client);
    event_loop_ref_client.shutdown();
    event_loop_ref_server.shutdown();
    let _event_loop_client = thread_client.join().unwrap();
    let _event_loop_server = thread_server.join().unwrap();
    let new_test_val = CLIENT_SERVER_SEND_TEST_VAL.load(Ordering::Relaxed);
    assert!(old_test_val < new_test_val,
            "old_test_val < new_test_val, old_test_val: {}, new_test_val: {}",
            old_test_val,
            new_test_val);
}

#[test]
fn event_loop_impl_add_listener() {
    let (event_loop_ref, thread) = event_loop_helper();
    let listener = TcpListener::bind(&super::ip("127.0.0.1:0")).unwrap();
    event_loop_ref.add_listener(listener);
    event_loop_ref.shutdown();
    let event_loop = thread.join().unwrap();
    assert_eq!(event_loop.handler.listeners.len(), 1);
}

#[test]
fn event_loop_impl_add_socket() {
    let (event_loop_ref, thread) = event_loop_helper();
    let listener = TcpListener::bind(&super::ip("127.0.0.1:25567")).unwrap();
    let _stream_local = TcpStream::connect(&super::ip("127.0.0.1:25567")).unwrap();
    let stream_remote;
    loop {
        match listener.accept().unwrap() {
            None => {}  // I think this is that there is no socket avalable to accept.
            Some((stream, _addr)) => {
                stream_remote = stream;
                break;
            }
        }
    }
    event_loop_ref.add_socket(stream_remote);
    event_loop_ref.shutdown();
    let event_loop = thread.join().unwrap();
    assert_eq!(event_loop.handler.connections.count(), 1);
}

fn event_loop_helper() -> (super::EventLoopImplRef, JoinHandle<super::EventLoopImpl>) {
    let mut event_loop = super::EventLoopImpl::new(super::MAX_CONNECTIONS).unwrap();
    let event_loop_ref: super::EventLoopImplRef = (&mut event_loop).into();
    (event_loop_ref,
     thread::spawn(move || {
        event_loop.run().unwrap();
        event_loop
    }))
}

#[test]
fn event_loop_impl_kill() {
    let (event_loop_ref, thread) = event_loop_helper();
    let listener = TcpListener::bind(&super::ip("127.0.0.1:25568")).unwrap();
    let stream_local = TcpStream::connect(&super::ip("127.0.0.1:25568")).unwrap();
    let _stream_remote;
    loop {
        match listener.accept().unwrap() {
            None => {}  // I think this is that there is no socket avalable to accept.
            Some((stream, _addr)) => {
                _stream_remote = stream;
                break;
            }
        }
    }
    let token = event_loop_ref.add_socket(stream_local);
    event_loop_ref.kill(token);
    event_loop_ref.shutdown();
    let event_loop = thread.join().unwrap();
    assert_eq!(event_loop.handler.connections.count(), 0);
    // TODO: Check if it actually closes the socket.
    // I can't think of a way to do this. It just isn't exposed in mio's API, io::Read's API, nowhere.
    // The closest I can get is reading 0 bytes, but even that is ambiguous.
}

#[test]
fn event_loop_impl_new() {
    super::EventLoopImpl::new(super::MAX_CONNECTIONS).unwrap();
}

#[test]
fn event_loop_impl_send() {
    let (event_loop_ref_local, thread_local) = event_loop_helper();
    let (event_loop_ref_remote, thread_remote) = event_loop_helper();
    let listener = TcpListener::bind(&super::ip("127.0.0.1:25569")).unwrap();
    let stream_local = TcpStream::connect(&super::ip("127.0.0.1:25569")).unwrap();
    let stream_remote;
    loop {
        match listener.accept().unwrap() {
            None => {}  // I think this is that there is no socket avalable to accept.
            Some((stream, _addr)) => {
                stream_remote = stream;
                break;
            }
        }
    }
    let _token_remote = event_loop_ref_local.add_socket(stream_local);
    let token_local = event_loop_ref_remote.add_socket(stream_remote);

    let old_test_val = EVENT_LOOP_SEND_TEST_VAL.load(Ordering::Relaxed);
    event_loop_ref_remote.send(token_local,
                               super::NetworkPacket::Test(TestValToModify::EventLoopImplSend));
    thread::sleep(Duration::from_millis(WAIT_TIME_MS));

    event_loop_ref_local.shutdown();
    event_loop_ref_remote.shutdown();
    let _event_loop_local = thread_local.join().unwrap();
    let _event_loop_remote = thread_remote.join().unwrap();
    let new_test_val = EVENT_LOOP_SEND_TEST_VAL.load(Ordering::Relaxed);
    assert!(old_test_val < new_test_val,
            "old_test_val < new_test_val, old_test_val: {}, new_test_val: {}",
            old_test_val,
            new_test_val);
}

#[test]
fn event_loop_impl_shutdown() {
    let (tx, rx) = channel::<Result<(), io::Error>>();
    let mut event_loop = super::EventLoopImpl::new(super::MAX_CONNECTIONS).unwrap();
    let event_loop_ref: super::EventLoopImplRef = (&mut event_loop).into();
    let thread = thread::spawn(move || tx.send(event_loop.run()).unwrap());
    event_loop_ref.shutdown();
    thread::sleep(Duration::from_millis(WAIT_TIME_MS));
    // This may, once in a very long time, fail. It really shouldn't, but it is possible.
    // Just raise the number or try again if it fails.
    match rx.try_recv() {
        Err(TryRecvError::Empty) => {
            panic!("EventLoopImpl did not stop after calling shutdown()! This fact is depended on \
                    by other unit tests, so ctrl+c here")
        }
        Err(TryRecvError::Disconnected) => {
            panic!("EventLoop somehow disconnected it's channel without stopping! This fact is \
                    depended on by other unit tests, so ctrl+c here")
        }
        Ok(res) => res,
    }
    .unwrap();
    thread.join().unwrap();
}

#[test]
fn get_packet_length_bad_magic_number() {
    let length: Option<u16> = super::get_packet_length([0xFE, 0xF0, 0xF6, 0xFD, 0, 10]);
    assert_eq!(length, None);
}

#[test]
fn get_packet_length_correct() {
    let mut m_number: [u8; 4] = [0; 4];
    LittleEndian::write_u32(&mut m_number, super::NET_MAGIC_NUMBER);
    let mut intended_length: [u8; 2] = [0; 2];
    LittleEndian::write_u16(&mut intended_length, 10);
    let length = super::get_packet_length([m_number[0], m_number[1], m_number[2], m_number[3], intended_length[0], intended_length[1]]).expect("get_packet_length wrongly checks NET_MAGIC_NUMBER");
    assert_eq!(length, 10);
}

#[test]
fn handler_new() {
    super::Handler::new(super::MAX_CONNECTIONS);
}
<commit_msg>Documents all the tests<commit_after>use std::thread;
use std::thread::JoinHandle;
use std::sync::mpsc::{TryRecvError, channel};
use std::time::Duration;
use std::io;
use std::sync::atomic::{AtomicUsize, Ordering};

use byteorder::{ByteOrder, LittleEndian};
use mio::tcp::{TcpListener, TcpStream};

use net::EventLoop;

/// The amount of time to wait while unit testing to make sure threads are syncronized when there is no other method.
const WAIT_TIME_MS: u64 = 250;

pub static CLIENT_SERVER_SEND_TEST_VAL: AtomicUsize = AtomicUsize::new(0);

pub static EVENT_LOOP_SEND_TEST_VAL: AtomicUsize = AtomicUsize::new(0);

#[derive(Serialize, Deserialize, Debug, Clone, Copy)]
pub enum TestValToModify {
    ClientServerSend,
    EventLoopImplSend,
}

/// Tests a client connecting to a server, going through the event loop.
#[test]
fn client_server_connect() {
    let (mut event_loop_ref_client, thread_client) = event_loop_helper();
    let (event_loop_ref_server, thread_server) = event_loop_helper();
    let listener = TcpListener::bind(&super::ip("127.0.0.1:25570")).unwrap();
    event_loop_ref_server.add_listener(listener);
    let client = super::client::Client::spawn_client(super::ip("127.0.0.1:25570"), &event_loop_ref_client).unwrap();
    client.shutdown(&mut event_loop_ref_client);
    event_loop_ref_client.shutdown();
    event_loop_ref_server.shutdown();
    let _event_loop_client = thread_client.join().unwrap();
    let _event_loop_server = thread_server.join().unwrap();
}

/// Tests the event loop sending and receving, using the client infrastructure.
#[test]
fn client_server_send() {
    let (mut event_loop_ref_client, thread_client) = event_loop_helper();
    let (event_loop_ref_server, thread_server) = event_loop_helper();
    let listener = TcpListener::bind(&super::ip("127.0.0.1:25571")).unwrap();
    event_loop_ref_server.add_listener(listener);
    let client = super::client::Client::spawn_client(super::ip("127.0.0.1:25571"), &event_loop_ref_client).unwrap();
    let old_test_val = CLIENT_SERVER_SEND_TEST_VAL.load(Ordering::Relaxed);

    client.send(&event_loop_ref_client,
                super::NetworkPacket::Test(TestValToModify::ClientServerSend));
    thread::sleep(Duration::from_millis(WAIT_TIME_MS));

    client.shutdown(&mut event_loop_ref_client);
    event_loop_ref_client.shutdown();
    event_loop_ref_server.shutdown();
    let _event_loop_client = thread_client.join().unwrap();
    let _event_loop_server = thread_server.join().unwrap();
    let new_test_val = CLIENT_SERVER_SEND_TEST_VAL.load(Ordering::Relaxed);
    assert!(old_test_val < new_test_val,
            "old_test_val < new_test_val, old_test_val: {}, new_test_val: {}",
            old_test_val,
            new_test_val);
}

/// Adds a listener to EventLoopImpl.
#[test]
fn event_loop_impl_add_listener() {
    let (event_loop_ref, thread) = event_loop_helper();
    let listener = TcpListener::bind(&super::ip("127.0.0.1:0")).unwrap();
    event_loop_ref.add_listener(listener);
    event_loop_ref.shutdown();
    let event_loop = thread.join().unwrap();
    assert_eq!(event_loop.handler.listeners.len(), 1);
}

/// Shuts down a EventLoopImpl.
///
/// If this test fails, all other tests will hang, due to the fact that all other tests block on shutting down the event loop.
#[test]
fn event_loop_impl_add_socket() {
    let (event_loop_ref, thread) = event_loop_helper();
    let listener = TcpListener::bind(&super::ip("127.0.0.1:25567")).unwrap();
    let _stream_local = TcpStream::connect(&super::ip("127.0.0.1:25567")).unwrap();
    let stream_remote;
    loop {
        match listener.accept().unwrap() {
            None => {}  // I think this is that there is no socket avalable to accept.
            Some((stream, _addr)) => {
                stream_remote = stream;
                break;
            }
        }
    }
    event_loop_ref.add_socket(stream_remote);
    event_loop_ref.shutdown();
    let event_loop = thread.join().unwrap();
    assert_eq!(event_loop.handler.connections.count(), 1);
}

/// Constructs an event loop on a new thread.
fn event_loop_helper() -> (super::EventLoopImplRef, JoinHandle<super::EventLoopImpl>) {
    let mut event_loop = super::EventLoopImpl::new(super::MAX_CONNECTIONS).unwrap();
    let event_loop_ref: super::EventLoopImplRef = (&mut event_loop).into();
    (event_loop_ref,
     thread::spawn(move || {
        event_loop.run().unwrap();
        event_loop
    }))
}

/// Kills a connection on an EventLoopImpl.
#[test]
fn event_loop_impl_kill() {
    let (event_loop_ref, thread) = event_loop_helper();
    let listener = TcpListener::bind(&super::ip("127.0.0.1:25568")).unwrap();
    let stream_local = TcpStream::connect(&super::ip("127.0.0.1:25568")).unwrap();
    let _stream_remote;
    loop {
        match listener.accept().unwrap() {
            None => {}  // I think this is that there is no socket avalable to accept.
            Some((stream, _addr)) => {
                _stream_remote = stream;
                break;
            }
        }
    }
    let token = event_loop_ref.add_socket(stream_local);
    event_loop_ref.kill(token);
    event_loop_ref.shutdown();
    let event_loop = thread.join().unwrap();
    assert_eq!(event_loop.handler.connections.count(), 0);
    // TODO: Check if it actually closes the socket.
    // I can't think of a way to do this. It just isn't exposed in mio's API, io::Read's API, nowhere.
    // The closest I can get is reading 0 bytes, but even that is ambiguous.
}

/// Constructs an EventLoopImpl.
#[test]
fn event_loop_impl_new() {
    super::EventLoopImpl::new(super::MAX_CONNECTIONS).unwrap();
}

/// Sends a packet using the event loop, but without the client helper struct.
#[test]
fn event_loop_impl_send() {
    let (event_loop_ref_local, thread_local) = event_loop_helper();
    let (event_loop_ref_remote, thread_remote) = event_loop_helper();
    let listener = TcpListener::bind(&super::ip("127.0.0.1:25569")).unwrap();
    let stream_local = TcpStream::connect(&super::ip("127.0.0.1:25569")).unwrap();
    let stream_remote;
    loop {
        match listener.accept().unwrap() {
            None => {}  // I think this is that there is no socket avalable to accept.
            Some((stream, _addr)) => {
                stream_remote = stream;
                break;
            }
        }
    }
    let _token_remote = event_loop_ref_local.add_socket(stream_local);
    let token_local = event_loop_ref_remote.add_socket(stream_remote);

    let old_test_val = EVENT_LOOP_SEND_TEST_VAL.load(Ordering::Relaxed);
    event_loop_ref_remote.send(token_local,
                               super::NetworkPacket::Test(TestValToModify::EventLoopImplSend));
    thread::sleep(Duration::from_millis(WAIT_TIME_MS));

    event_loop_ref_local.shutdown();
    event_loop_ref_remote.shutdown();
    let _event_loop_local = thread_local.join().unwrap();
    let _event_loop_remote = thread_remote.join().unwrap();
    let new_test_val = EVENT_LOOP_SEND_TEST_VAL.load(Ordering::Relaxed);
    assert!(old_test_val < new_test_val,
            "old_test_val < new_test_val, old_test_val: {}, new_test_val: {}",
            old_test_val,
            new_test_val);
}

/// Shuts down an EventLoopImpl.
#[test]
fn event_loop_impl_shutdown() {
    let (tx, rx) = channel::<Result<(), io::Error>>();
    let mut event_loop = super::EventLoopImpl::new(super::MAX_CONNECTIONS).unwrap();
    let event_loop_ref: super::EventLoopImplRef = (&mut event_loop).into();
    let thread = thread::spawn(move || tx.send(event_loop.run()).unwrap());
    event_loop_ref.shutdown();
    thread::sleep(Duration::from_millis(WAIT_TIME_MS));
    // This may, once in a very long time, fail. It really shouldn't, but it is possible.
    // Just raise the number or try again if it fails.
    match rx.try_recv() {
        Err(TryRecvError::Empty) => {
            panic!("EventLoopImpl did not stop after calling shutdown()! This fact is depended on \
                    by other unit tests, so ctrl+c here")
        }
        Err(TryRecvError::Disconnected) => {
            panic!("EventLoop somehow disconnected it's channel without stopping! This fact is \
                    depended on by other unit tests, so ctrl+c here")
        }
        Ok(res) => res,
    }
    .unwrap();
    thread.join().unwrap();
}

/// Tests get_packet_length with a bad magic number.
#[test]
fn get_packet_length_bad_magic_number() {
    let length: Option<u16> = super::get_packet_length([0xFE, 0xF0, 0xF6, 0xFD, 0, 10]);
    assert_eq!(length, None);
}

/// Tests get_packet_length with a correct magic number, and ensures the length is the same.
#[test]
fn get_packet_length_correct() {
    let mut m_number: [u8; 4] = [0; 4];
    LittleEndian::write_u32(&mut m_number, super::NET_MAGIC_NUMBER);
    let mut intended_length: [u8; 2] = [0; 2];
    LittleEndian::write_u16(&mut intended_length, 10);
    let length = super::get_packet_length([m_number[0], m_number[1], m_number[2], m_number[3], intended_length[0], intended_length[1]]).expect("get_packet_length wrongly checks NET_MAGIC_NUMBER");
    assert_eq!(length, 10);
}

/// Constructs a Handler.
#[test]
fn handler_new() {
    super::Handler::new(super::MAX_CONNECTIONS);
}
<|endoftext|>


<commit_before>use alloc::vec::Vec;
use core::convert::TryFrom;

type ByteArray = &'static [u8];

#[cfg_attr(test, derive(Clone, Debug))]
#[allow(non_camel_case_types)]
#[derive(PartialEq)]
pub enum ApduStatusCode {
    SW_SUCCESS,
    /// Command successfully executed; 'XX' bytes of data are
    /// available and can be requested using GET RESPONSE.
    SW_GET_RESPONSE,
    SW_WRONG_DATA,
    SW_WRONG_LENGTH,
    SW_COND_USE_NOT_SATISFIED,
    SW_FILE_NOT_FOUND,
    SW_INCORRECT_P1P2,
    /// Instruction code not supported or invalid
    SW_INS_INVALID,
    SW_CLA_INVALID,
    SW_INTERNAL_EXCEPTION,
}

impl From<ApduStatusCode> for ByteArray {
    fn from(status_code: ApduStatusCode) -> ByteArray {
        match status_code {
            ApduStatusCode::SW_SUCCESS => b"\x90\x00",
            ApduStatusCode::SW_GET_RESPONSE => b"\x61\x00",
            ApduStatusCode::SW_WRONG_DATA => b"\x6A\x80",
            ApduStatusCode::SW_WRONG_LENGTH => b"\x67\x00",
            ApduStatusCode::SW_COND_USE_NOT_SATISFIED => b"\x69\x85",
            ApduStatusCode::SW_FILE_NOT_FOUND => b"\x6a\x82",
            ApduStatusCode::SW_INCORRECT_P1P2 => b"\x6a\x86",
            ApduStatusCode::SW_INS_INVALID => b"\x6d\x00",
            ApduStatusCode::SW_CLA_INVALID => b"\x6e\x00",
            ApduStatusCode::SW_INTERNAL_EXCEPTION => b"\x6f\x00",
        }
    }
}

#[allow(non_camel_case_types, dead_code)]
pub enum ApduIns {
    SELECT = 0xA4,
    READ_BINARY = 0xB0,
    GET_RESPONSE = 0xC0,
}

#[cfg_attr(test, derive(Clone, Debug))]
#[allow(dead_code)]
#[derive(Default, PartialEq)]
pub struct ApduHeader {
    cla: u8,
    ins: u8,
    p1: u8,
    p2: u8,
}

impl From<&[u8]> for ApduHeader {
    fn from(header: &[u8]) -> Self {
        ApduHeader {
            cla: header[0],
            ins: header[1],
            p1: header[2],
            p2: header[3],
        }
    }
}

#[cfg_attr(test, derive(Clone, Debug))]
#[allow(dead_code)]
#[derive(Default, PartialEq)]
pub struct APDU {
    header: ApduHeader,
    lc: u16,
    data: Vec<u8>,
    le: u32,
    case_type: u8,
}

const APDU_HEADER_LEN: u8 = 4;

impl TryFrom<&[u8]> for APDU {
    type Error = ApduStatusCode;

    fn try_from(frame: &[u8]) -> Result<Self, ApduStatusCode> {
        if frame.len() < APDU_HEADER_LEN as usize {
            return Err(ApduStatusCode::SW_WRONG_DATA);
        }
        //        +-----+-----+----+----+
        // header | CLA | INS | P1 | P2 |
        //        +-----+-----+----+----+
        let (header, payload) = frame.split_at(APDU_HEADER_LEN as usize);

        let mut apdu = APDU {
            header: header.into(),
            lc: 0x00,
            data: Vec::new(),
            le: 0x00,
            case_type: 0x00,
        };

        // case 1
        if payload.is_empty() {
            apdu.case_type = 0x01;
        } else {
            let byte_0 = payload[0];
            // case 2S (Le)
            if payload.len() == 1 {
                apdu.case_type = 0x02;
                apdu.le = if byte_0 == 0x00 {
                    // Ne = 256
                    0x100
                } else {
                    byte_0.into()
                }
            }
            // case 3S (Lc + data)
            if payload.len() == (1 + byte_0) as usize && byte_0 != 0 {
                apdu.case_type = 0x03;
                apdu.lc = byte_0.into();
                apdu.data = payload[1..].to_vec();
            }
            // case 4S (Lc + data + Le)
            if payload.len() == (1 + byte_0 + 1) as usize && byte_0 != 0 {
                apdu.case_type = 0x04;
                apdu.lc = byte_0.into();
                apdu.data = payload[1..(payload.len() - 1)].to_vec();
                apdu.le = (*payload.last().unwrap()).into();
                if apdu.le == 0x00 {
                    apdu.le = 0x100;
                }
            }
        }
        // TODO: Add extended length cases
        if apdu.case_type == 0x00 {
            return Err(ApduStatusCode::SW_COND_USE_NOT_SATISFIED);
        }
        Ok(apdu)
    }
}

#[cfg(test)]
mod test {
    use super::*;

    fn pass_frame(frame: &[u8]) -> Result<APDU, ApduStatusCode> {
        APDU::try_from(frame)
    }

    #[test]
    fn test_case_type_1() {
        let frame: [u8; 4] = [0x00, 0x12, 0x00, 0x80];
        let response = pass_frame(&frame);
        assert!(response.is_ok());
        let expected = APDU {
            header: ApduHeader {
                cla: 0x00,
                ins: 0x12,
                p1: 0x00,
                p2: 0x80,
            },
            lc: 0x00,
            data: Vec::new(),
            le: 0x00,
            case_type: 0x01,
        };
        assert_eq!(expected, response.unwrap());
    }

    #[test]
    fn test_case_type_2_short() {
        let frame: [u8; 5] = [0x00, 0xb0, 0x00, 0x00, 0x0f];
        let response = pass_frame(&frame);
        assert!(response.is_ok());
        let expected = APDU {
            header: ApduHeader {
                cla: 0x00,
                ins: 0xb0,
                p1: 0x00,
                p2: 0x00,
            },
            lc: 0x00,
            data: Vec::new(),
            le: 0x0f,
            case_type: 0x02,
        };
        assert_eq!(expected, response.unwrap());
    }

    #[test]
    fn test_case_type_2_short_le() {
        let frame: [u8; 5] = [0x00, 0xb0, 0x00, 0x00, 0x00];
        let response = pass_frame(&frame);
        assert!(response.is_ok());
        let expected = APDU {
            header: ApduHeader {
                cla: 0x00,
                ins: 0xb0,
                p1: 0x00,
                p2: 0x00,
            },
            lc: 0x00,
            data: Vec::new(),
            le: 0x100,
            case_type: 0x02,
        };
        assert_eq!(expected, response.unwrap());
    }

    #[test]
    fn test_case_type_3_short() {
        let frame: [u8; 7] = [0x00, 0xa4, 0x00, 0x0c, 0x02, 0xe1, 0x04];
        let payload = [0xe1, 0x04];
        let response = pass_frame(&frame);
        assert!(response.is_ok());
        let expected = APDU {
            header: ApduHeader {
                cla: 0x00,
                ins: 0xa4,
                p1: 0x00,
                p2: 0x0c,
            },
            lc: 0x02,
            data: payload.to_vec(),
            le: 0x00,
            case_type: 0x03,
        };
        assert_eq!(expected, response.unwrap());
    }

    #[test]
    fn test_case_type_4_short() {
        let frame: [u8; 13] = [
            0x00, 0xa4, 0x04, 0x00, 0x07, 0xd2, 0x76, 0x00, 0x00, 0x85, 0x01, 0x01, 0xff,
        ];
        let payload = [0xd2, 0x76, 0x00, 0x00, 0x85, 0x01, 0x01];
        let response = pass_frame(&frame);
        assert!(response.is_ok());
        let expected = APDU {
            header: ApduHeader {
                cla: 0x00,
                ins: 0xa4,
                p1: 0x04,
                p2: 0x00,
            },
            lc: 0x07,
            data: payload.to_vec(),
            le: 0xff,
            case_type: 0x04,
        };
        assert_eq!(expected, response.unwrap());
    }

    #[test]
    fn test_case_type_4_short_le() {
        let frame: [u8; 13] = [
            0x00, 0xa4, 0x04, 0x00, 0x07, 0xd2, 0x76, 0x00, 0x00, 0x85, 0x01, 0x01, 0x00,
        ];
        let payload = [0xd2, 0x76, 0x00, 0x00, 0x85, 0x01, 0x01];
        let response = pass_frame(&frame);
        assert!(response.is_ok());
        let expected = APDU {
            header: ApduHeader {
                cla: 0x00,
                ins: 0xa4,
                p1: 0x04,
                p2: 0x00,
            },
            lc: 0x07,
            data: payload.to_vec(),
            le: 0x100,
            case_type: 0x04,
        };
        assert_eq!(expected, response.unwrap());
    }

    #[test]
    fn test_invalid_apdu_header_length() {
        let frame: [u8; 3] = [0x00, 0x12, 0x00];
        let response = pass_frame(&frame);
        assert!(response.is_err());
        assert_eq!(Some(ApduStatusCode::SW_WRONG_DATA), response.err());
    }

    #[test]
    fn test_unsupported_case_type() {
        let frame: [u8; 73] = [
            0x00, 0x01, 0x03, 0x00, 0x00, 0x00, 0x40, 0xe3, 0x8f, 0xde, 0x51, 0x3d, 0xac, 0x9d,
            0x1c, 0x6e, 0x86, 0x76, 0x31, 0x40, 0x25, 0x96, 0x86, 0x4d, 0x29, 0xe8, 0x07, 0xb3,
            0x56, 0x19, 0xdf, 0x4a, 0x00, 0x02, 0xae, 0x2a, 0x8c, 0x9d, 0x5a, 0xab, 0xc3, 0x4b,
            0x4e, 0xb9, 0x78, 0xb9, 0x11, 0xe5, 0x52, 0x40, 0xf3, 0x45, 0x64, 0x9c, 0xd3, 0xd7,
            0xe8, 0xb5, 0x83, 0xfb, 0xe0, 0x66, 0x98, 0x4d, 0x98, 0x81, 0xf7, 0xb5, 0x49, 0x4d,
            0xcb, 0x00, 0x00,
        ];
        let response = pass_frame(&frame);
        assert!(response.is_err());
        assert_eq!(
            Some(ApduStatusCode::SW_COND_USE_NOT_SATISFIED),
            response.err()
        );
    }
}
<commit_msg>Addressing some of the requested changes<commit_after>use alloc::vec::Vec;
use core::convert::TryFrom;

type ByteArray = &'static [u8];

const APDU_HEADER_LEN: usize = 4;

#[cfg_attr(test, derive(Clone, Debug))]
#[allow(non_camel_case_types)]
#[derive(PartialEq)]
pub enum ApduStatusCode {
    SW_SUCCESS,
    /// Command successfully executed; 'XX' bytes of data are
    /// available and can be requested using GET RESPONSE.
    SW_GET_RESPONSE,
    SW_WRONG_DATA,
    SW_WRONG_LENGTH,
    SW_COND_USE_NOT_SATISFIED,
    SW_FILE_NOT_FOUND,
    SW_INCORRECT_P1P2,
    /// Instruction code not supported or invalid
    SW_INS_INVALID,
    SW_CLA_INVALID,
    SW_INTERNAL_EXCEPTION,
}

impl From<ApduStatusCode> for ByteArray {
    fn from(status_code: ApduStatusCode) -> ByteArray {
        match status_code {
            ApduStatusCode::SW_SUCCESS => b"\x90\x00",
            ApduStatusCode::SW_GET_RESPONSE => b"\x61\x00",
            ApduStatusCode::SW_WRONG_DATA => b"\x6A\x80",
            ApduStatusCode::SW_WRONG_LENGTH => b"\x67\x00",
            ApduStatusCode::SW_COND_USE_NOT_SATISFIED => b"\x69\x85",
            ApduStatusCode::SW_FILE_NOT_FOUND => b"\x6a\x82",
            ApduStatusCode::SW_INCORRECT_P1P2 => b"\x6a\x86",
            ApduStatusCode::SW_INS_INVALID => b"\x6d\x00",
            ApduStatusCode::SW_CLA_INVALID => b"\x6e\x00",
            ApduStatusCode::SW_INTERNAL_EXCEPTION => b"\x6f\x00",
        }
    }
}

#[allow(dead_code)]
pub enum ApduInstructions {
    Select = 0xA4,
    ReadBinary = 0xB0,
    GetResponse = 0xC0,
}

#[cfg_attr(test, derive(Clone, Debug))]
#[allow(dead_code)]
#[derive(Default, PartialEq)]
pub struct ApduHeader {
    cla: u8,
    ins: u8,
    p1: u8,
    p2: u8,
}

impl From<&[u8]> for ApduHeader {
    fn from(header: &[u8]) -> Self {
        ApduHeader {
            cla: header[0],
            ins: header[1],
            p1: header[2],
            p2: header[3],
        }
    }
}

#[cfg_attr(test, derive(Clone, Debug))]
#[derive(PartialEq)]
/// The APDU cases
pub enum Case {
    Le,
    LcData,
    LcDataLe,
    // TODO: More cases to add as extended length APDUs
    // Le could be 2 or 3 Bytes
}

#[cfg_attr(test, derive(Clone, Debug))]
#[allow(dead_code)]
#[derive(PartialEq)]
pub enum ApduType {
    Instruction,
    Short(Case),
    Extended(Case),
    Unknown,
}

impl Default for ApduType {
    fn default() -> ApduType {
        ApduType::Unknown
    }
}

#[cfg_attr(test, derive(Clone, Debug))]
#[allow(dead_code)]
#[derive(Default, PartialEq)]
pub struct APDU {
    header: ApduHeader,
    lc: u16,
    data: Vec<u8>,
    le: u32,
    case_type: ApduType,
}

impl TryFrom<&[u8]> for APDU {
    type Error = ApduStatusCode;

    fn try_from(frame: &[u8]) -> Result<Self, ApduStatusCode> {
        if frame.len() < APDU_HEADER_LEN as usize {
            return Err(ApduStatusCode::SW_WRONG_DATA);
        }
        //        +-----+-----+----+----+
        // header | CLA | INS | P1 | P2 |
        //        +-----+-----+----+----+
        let (header, payload) = frame.split_at(APDU_HEADER_LEN);

        let mut apdu = APDU {
            header: header.into(),
            lc: 0x00,
            data: Vec::new(),
            le: 0x00,
            case_type: ApduType::default(),
        };

        // case 1
        if payload.is_empty() {
            apdu.case_type = ApduType::Instruction;
        } else {
            let byte_0 = payload[0];
            if payload.len() == 1 {
                apdu.case_type = ApduType::Short(Case::Le);
                apdu.le = if byte_0 == 0x00 {
                    // Ne = 256
                    0x100
                } else {
                    byte_0.into()
                }
            }
            if payload.len() == (1 + byte_0) as usize && byte_0 != 0 {
                apdu.case_type = ApduType::Short(Case::LcData);
                apdu.lc = byte_0.into();
                apdu.data = payload[1..].to_vec();
            }
            if payload.len() == (1 + byte_0 + 1) as usize && byte_0 != 0 {
                apdu.case_type = ApduType::Short(Case::LcDataLe);
                apdu.lc = byte_0.into();
                apdu.data = payload[1..(payload.len() - 1)].to_vec();
                apdu.le = (*payload.last().unwrap()).into();
                if apdu.le == 0x00 {
                    apdu.le = 0x100;
                }
            }
        }
        // TODO: Add extended length cases
        if apdu.case_type == ApduType::default() {
            return Err(ApduStatusCode::SW_COND_USE_NOT_SATISFIED);
        }
        Ok(apdu)
    }
}

#[cfg(test)]
mod test {
    use super::*;

    fn pass_frame(frame: &[u8]) -> Result<APDU, ApduStatusCode> {
        APDU::try_from(frame)
    }

    #[test]
    fn test_case_type_1() {
        let frame: [u8; 4] = [0x00, 0x12, 0x00, 0x80];
        let response = pass_frame(&frame);
        assert!(response.is_ok());
        let expected = APDU {
            header: ApduHeader {
                cla: 0x00,
                ins: 0x12,
                p1: 0x00,
                p2: 0x80,
            },
            lc: 0x00,
            data: Vec::new(),
            le: 0x00,
            case_type: ApduType::Instruction,
        };
        assert_eq!(Ok(expected), response);
    }

    #[test]
    fn test_case_type_2_short() {
        let frame: [u8; 5] = [0x00, 0xb0, 0x00, 0x00, 0x0f];
        let response = pass_frame(&frame);
        let expected = APDU {
            header: ApduHeader {
                cla: 0x00,
                ins: 0xb0,
                p1: 0x00,
                p2: 0x00,
            },
            lc: 0x00,
            data: Vec::new(),
            le: 0x0f,
            case_type: ApduType::Short(Case::Le),
        };
        assert_eq!(Ok(expected), response);
    }

    #[test]
    fn test_case_type_2_short_le() {
        let frame: [u8; 5] = [0x00, 0xb0, 0x00, 0x00, 0x00];
        let response = pass_frame(&frame);
        let expected = APDU {
            header: ApduHeader {
                cla: 0x00,
                ins: 0xb0,
                p1: 0x00,
                p2: 0x00,
            },
            lc: 0x00,
            data: Vec::new(),
            le: 0x100,
            case_type: ApduType::Short(Case::Le),
        };
        assert_eq!(Ok(expected), response);
    }

    #[test]
    fn test_case_type_3_short() {
        let frame: [u8; 7] = [0x00, 0xa4, 0x00, 0x0c, 0x02, 0xe1, 0x04];
        let payload = [0xe1, 0x04];
        let response = pass_frame(&frame);
        let expected = APDU {
            header: ApduHeader {
                cla: 0x00,
                ins: 0xa4,
                p1: 0x00,
                p2: 0x0c,
            },
            lc: 0x02,
            data: payload.to_vec(),
            le: 0x00,
            case_type: ApduType::Short(Case::LcData),
        };
        assert_eq!(Ok(expected), response);
    }

    #[test]
    fn test_case_type_4_short() {
        let frame: [u8; 13] = [
            0x00, 0xa4, 0x04, 0x00, 0x07, 0xd2, 0x76, 0x00, 0x00, 0x85, 0x01, 0x01, 0xff,
        ];
        let payload = [0xd2, 0x76, 0x00, 0x00, 0x85, 0x01, 0x01];
        let response = pass_frame(&frame);
        let expected = APDU {
            header: ApduHeader {
                cla: 0x00,
                ins: 0xa4,
                p1: 0x04,
                p2: 0x00,
            },
            lc: 0x07,
            data: payload.to_vec(),
            le: 0xff,
            case_type: ApduType::Short(Case::LcDataLe),
        };
        assert_eq!(Ok(expected), response);
    }

    #[test]
    fn test_case_type_4_short_le() {
        let frame: [u8; 13] = [
            0x00, 0xa4, 0x04, 0x00, 0x07, 0xd2, 0x76, 0x00, 0x00, 0x85, 0x01, 0x01, 0x00,
        ];
        let payload = [0xd2, 0x76, 0x00, 0x00, 0x85, 0x01, 0x01];
        let response = pass_frame(&frame);
        let expected = APDU {
            header: ApduHeader {
                cla: 0x00,
                ins: 0xa4,
                p1: 0x04,
                p2: 0x00,
            },
            lc: 0x07,
            data: payload.to_vec(),
            le: 0x100,
            case_type: ApduType::Short(Case::LcDataLe),
        };
        assert_eq!(Ok(expected), response);
    }

    #[test]
    fn test_invalid_apdu_header_length() {
        let frame: [u8; 3] = [0x00, 0x12, 0x00];
        let response = pass_frame(&frame);
        assert_eq!(Err(ApduStatusCode::SW_WRONG_DATA), response);
    }

    #[test]
    fn test_unsupported_case_type() {
        let frame: [u8; 73] = [
            0x00, 0x01, 0x03, 0x00, 0x00, 0x00, 0x40, 0xe3, 0x8f, 0xde, 0x51, 0x3d, 0xac, 0x9d,
            0x1c, 0x6e, 0x86, 0x76, 0x31, 0x40, 0x25, 0x96, 0x86, 0x4d, 0x29, 0xe8, 0x07, 0xb3,
            0x56, 0x19, 0xdf, 0x4a, 0x00, 0x02, 0xae, 0x2a, 0x8c, 0x9d, 0x5a, 0xab, 0xc3, 0x4b,
            0x4e, 0xb9, 0x78, 0xb9, 0x11, 0xe5, 0x52, 0x40, 0xf3, 0x45, 0x64, 0x9c, 0xd3, 0xd7,
            0xe8, 0xb5, 0x83, 0xfb, 0xe0, 0x66, 0x98, 0x4d, 0x98, 0x81, 0xf7, 0xb5, 0x49, 0x4d,
            0xcb, 0x00, 0x00,
        ];
        let response = pass_frame(&frame);
        assert_eq!(Err(ApduStatusCode::SW_COND_USE_NOT_SATISFIED), response);
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the terms of the
// Mozilla Public License, v. 2.0. If a copy of the MPL was not
// distributed with this file, You can obtain one at
// http://mozilla.org/MPL/2.0/.

use std::thread;
use std::net::{TcpListener, ToSocketAddrs};
use std::ops::DerefMut;
use std::os::unix::io::RawFd;
use std::collections::LinkedList;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{
    channel,
    Sender,
    Receiver
};

use epoll;
use epoll::util::*;
use epoll::EpollEvent;
use simple_stream::nbetstream::NbetStream;
use stats;
use types::*;
use socket::Socket;
use resources::ResourcePool;

// When added to epoll, these will be the conditions of kernel notification:
//
// EPOLLIN      - Available for read
// EPOLLRDHUP   - Connection has been closed
// EPOLLONESHOT - After an event has been received, and reported, do not track
//                further changes until explicitly told to do so.
const EVENTS: u32 = (
    event_type::EPOLLET | // Set fd to EdgeTrigger mode
    event_type::EPOLLIN |
    event_type::EPOLLRDHUP |
    event_type::EPOLLONESHOT |
    event_type::EPOLLERR |
    event_type::EPOLLHUP);

pub fn begin<T, K>(address: T, handler: Box<K>) where
    T: ToSocketAddrs + Send + 'static,
    K: EventHandler + Send + Sync + 'static {
    // Master socket list
    let sockets = Arc::new(Mutex::new(LinkedList::<Socket>::new()));

    // Epoll instance
    let result = epoll::create1(0);
    if result.is_err() {
        let err = result.unwrap_err();
        error!("Unable to create epoll instance: {}", err);
        panic!()
    }
    let epfd = result.unwrap();

    // EpollEvent handler thread
    let epfd1 = epfd.clone();
    let sockets1 = sockets.clone();
    let (tx_ep_event, rx_ep_event): (Sender<EpollEvent>, Receiver<EpollEvent>) = channel();
    let safe_handler = Arc::new(Mutex::new(*handler));
    thread::Builder::new()
        .name("EpollEvent Handler".to_string())
        .spawn(move || {
            epoll_event_handler(epfd1, rx_ep_event, sockets1, safe_handler);
        }).unwrap();

    // Epoll wait thread
    let epfd2 = epfd.clone();
    thread::Builder::new()
        .name("Epoll Wait".to_string())
        .spawn(move || {
            epoll_wait_loop(epfd2, tx_ep_event);
        }).unwrap();

    // New connection thread
    let epfd3 = epfd.clone();
    let sockets2 = sockets.clone();
    let prox = thread::Builder::new()
        .name("TCP Incoming Listener".to_string())
        .spawn(move || {
            listen(address, epfd3, sockets2);
        }).unwrap();

    // Stay alive forever
    let _ = prox.join();
}

fn listen<T: ToSocketAddrs>(address: T, epfd: RawFd, sockets: SocketList) {
    let listener = TcpListener::bind(address).unwrap();
    for connection in listener.incoming() {
        if connection.is_err() {
            warn!("Error encountered during TCP connection: {}", connection.unwrap_err());
            continue;
        }

        let result = NbetStream::new(connection.unwrap());
        if result.is_err() {
            warn!("Error creating nbstream: {:?}", result.unwrap_err());
            continue;
        }

        let socket_ptr;
        let socket = Socket::new(result.unwrap());
        { // Begin Mutex lock
            let mut guard = match sockets.lock() {
                Ok(guard) => guard,
                Err(poisoned) => {
                    warn!("SocketList Mutex failed, using anyway...");
                    poisoned.into_inner()
                }
            };
            let socket_list = guard.deref_mut();
            socket_list.push_back(socket.clone());
            socket_ptr = socket_list.back_mut().unwrap() as *mut Socket;
        } // End Mutex lock

        // Add to epoll
        let mut event = EpollEvent {
            data: (socket_ptr as usize) as u64,
            events: EVENTS
        };
        let _ = epoll::ctl(epfd, ctl_op::ADD, socket.raw_fd(), &mut event).map_err(|e| {
            error!("CtrlError during add: {}", e);
        });
    }
    drop(listener);
}

fn epoll_wait_loop(epfd: RawFd, tx_handler: Sender<EpollEvent>) {
    // Maximum number of events we want to be notified of at one time
    let mut events = Vec::<EpollEvent>::with_capacity(100);
    unsafe { events.set_len(100); }

    loop {
        match epoll::wait(epfd, &mut events[..], -1) {
            Ok(num_events) => {
                for x in 0..num_events as usize {
                    let _ = tx_handler.send(events[x]);
                }
            }
            Err(e) => {
                error!("Error on epoll::wait(): {}", e);
                panic!()
            }
        }
    }
}

fn epoll_event_handler(epfd: RawFd,
                       rx: Receiver<EpollEvent>,
                       sockets: SocketList,
                       handler: SafeHandler) {
    // Types of events we care about
    const READ_EVENT: u32 = event_type::EPOLLIN;
    const DROP_EVENT: u32 = event_type::EPOLLRDHUP | event_type::EPOLLHUP | event_type::EPOLLERR;

    // Thread pool
    let mut pool = ResourcePool::new();

    for event in rx.iter() {
        // Get a pointer to the socket in the list
        let socket_ptr;
        { // Begin Mutex lock
            let _ = match sockets.lock() {
                Ok(guard) => guard,
                Err(poisoned) => {
                    warn!("SocketList Mutex failed, using anyway...");
                    poisoned.into_inner()
                }
            };
            socket_ptr = event.data as *mut Socket;
        } // End Mutex lock

        // Move semantics are a little weird right here, so we must use clones of clones
        // in order to make the compiler happy
        if (event.events & DROP_EVENT) > 0 {
            let t_handler = handler.clone();
            let socket_list = sockets.clone();
            let socket = unsafe { (*socket_ptr).clone() };
            pool.run(move || {
                handle_drop_event(epfd, socket.clone(), socket_list.clone(), t_handler.clone());
            });
        } else if (event.events & READ_EVENT) > 0 {
            let t_handler = handler.clone();
            let socket_list = sockets.clone();
            let socket = unsafe { (*socket_ptr).clone() };
            let ptr_addr = socket_ptr as usize;
            pool.run(move || {
                handle_read_event(epfd, socket.clone(), socket_list.clone(),
                    t_handler.clone(), ptr_addr);
            });
        } else { warn!("Unknown epoll event received: {}", event.data); continue; }
    }
}

fn handle_read_event(epfd: RawFd,
                     socket: Socket,
                     sockets: SocketList,
                     handler: SafeHandler,
                     ptr_addr: usize) {
    let mut socket = socket;
    let result = socket.read();
    if result.is_err() {
        // There is no saveing grace here. If we can't read, only option
        // is to remove it from server
        handle_drop_event(epfd, socket, sockets, handler);
        return;
    }

    for msg in socket.buffer().iter() {
        // Pass the msg on to the handler
        { // EventHandler Mutex locked
            let mut guard = match handler.lock() {
                Ok(guard) => guard,
                Err(poisoned) => {
                    warn!("EventHandler Mutex failed, using anyway...");
                    poisoned.into_inner()
                }
            };
            let e_handler = guard.deref_mut();
            e_handler.on_data_received(socket.clone(), sockets.clone(), msg.clone());
        } // EventHandler Mutex unlocked

        // Collect some stats
        stats::msg_recv();
        stats::bytes_recv(msg.len());

        // Add socket's fd to epoll watch list
        //
        // std::net::TcpStream uses F_DUPFD flag when calling clone(). This means that
        // it is very likely that the read event we received for this socket may not actually
        // be the exact same file descriptor (RawFd) as the one that registered it. Therefore,
        // we have several steps to take in order to try and re-register the same one or add it
        // again if epoll tells us the fd does not exist.
        // This approach does not have any ill effects though, Rust happily sets the CLOEXEC flag
        // on all owned fds, so once any of our sockets go out of scope, the kernel will clean
        // them up for us. Epoll, as it happens, removes fds from the watch list also, once the
        // kernel cleans them up.
        // We start off with an add, because we're assuming the socket is a clone and fds are
        // different
        let mut event = EpollEvent {
            data: ptr_addr as u64,
            events: EVENTS
        };
        let _ = epoll::ctl(epfd, ctl_op::ADD, socket.raw_fd(), &mut event).map_err(|e| {
            if e == CtlError::EEXIST {
                // Fd does not currently exist in epoll, we'll add it then because it must have
                // been a F_DUPFD'd fd from one of the cloned sockets.
                let _ = epoll::ctl(epfd, ctl_op::MOD, socket.raw_fd(), &mut event).map_err(|e| {
                    error!("During epoll::ctl: {}", e);
                });
            } else {
                error!("During epoll::ctl: {}", e);
            }
        });
    }
}

fn handle_drop_event(epfd: RawFd, socket: Socket, sockets: SocketList, handler: SafeHandler) {
    epoll_remove_fd(epfd, socket.raw_fd());
    remove_socket_from_list(socket.id(), sockets);

    let mut guard = match handler.lock() {
        Ok(g) => g,
        Err(p) => {
            warn!("EventHandler Mutex failed, using anyway...");
            p.into_inner()
        }
    };
    let e_handler = guard.deref_mut();
    e_handler.on_socket_closed(socket.id());
}

fn epoll_remove_fd(epfd: RawFd, fd: RawFd) {
    debug!("remove_socket_from_epoll");

    // In kernel versions before 2.6.9, the EPOLL_CTL_DEL operation required
    // a non-null pointer in event, even though this argument is ignored.
    // Since Linux 2.6.9, event can be specified as NULL when using
    // EPOLL_CTL_DEL.  Applications that need to be portable to kernels
    // before 2.6.9 should specify a non-null pointer in event.
    let mut rm_event = EpollEvent {
        data: 0 as u64,
        events: 0 as u32
    };

    // Depending on how fds are duplicated with .clone(), this may fail.
    //
    // If the failure case is CtlError::ENOENT, we do not care, because
    // epoll will clean the up the descriptor after they are all dropped from
    // program memory
    match epoll::ctl(epfd, ctl_op::DEL, fd, &mut rm_event) {
        Ok(()) => trace!("Socket removed from epoll watch list"),
        Err(e) => {
            match e {
                CtlError::ENOENT => {
                    trace!("Fd not found in epoll, will remove when fd is syscall closed");
                }
                _ => warn!("Epoll CtrlError during del: {}", e)
            }
        }
    };
}

/// Removes socket_ids from master list
fn remove_socket_from_list(socket_id: u32, sockets: SocketList) {
    debug!("remove_socket_from_list");

    let mut s_guard = match sockets.lock() {
        Ok(guard) => guard,
        Err(poisoned) => {
            warn!("SocketList Mutex was poisoned, using anyway");
            poisoned.into_inner()
        }
    };
    let s_list = s_guard.deref_mut();

    let mut s_found = false;
    let mut index = 1usize;
    for s in s_list.iter() {
        if s.id() == socket_id {
            s_found = true;
            break;
        }
        index += 1;
    }

    if !s_found { return; }

    if index == 1 {
        s_list.pop_front();
    } else {
        let mut split = s_list.split_off(index - 1);
        split.pop_front();
        s_list.append(&mut split);
    }

    stats::conn_lost();

    trace!("socket removed from list");
}
<commit_msg>Docs added for registered epoll events<commit_after>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the terms of the
// Mozilla Public License, v. 2.0. If a copy of the MPL was not
// distributed with this file, You can obtain one at
// http://mozilla.org/MPL/2.0/.

use std::thread;
use std::net::{TcpListener, ToSocketAddrs};
use std::ops::DerefMut;
use std::os::unix::io::RawFd;
use std::collections::LinkedList;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{
    channel,
    Sender,
    Receiver
};

use epoll;
use epoll::util::*;
use epoll::EpollEvent;
use simple_stream::nbetstream::NbetStream;
use stats;
use types::*;
use socket::Socket;
use resources::ResourcePool;

// When added to epoll, these will be the conditions of kernel notification:
//
// EPOLLIN      - Available for read
// EPOLLRDHUP   - Connection has been closed
// EPOLLONESHOT - After an event has been received, and reported, do not track
//                further changes until explicitly told to do so.
// EPOLLERR     - Some error occurred
// EPOLLHUP     - Hang up happened
const EVENTS: u32 = (
    event_type::EPOLLET | // Set fd to EdgeTrigger mode
    event_type::EPOLLIN |
    event_type::EPOLLRDHUP |
    event_type::EPOLLONESHOT |
    event_type::EPOLLERR |
    event_type::EPOLLHUP);

pub fn begin<T, K>(address: T, handler: Box<K>) where
    T: ToSocketAddrs + Send + 'static,
    K: EventHandler + Send + Sync + 'static {
    // Master socket list
    let sockets = Arc::new(Mutex::new(LinkedList::<Socket>::new()));

    // Epoll instance
    let result = epoll::create1(0);
    if result.is_err() {
        let err = result.unwrap_err();
        error!("Unable to create epoll instance: {}", err);
        panic!()
    }
    let epfd = result.unwrap();

    // EpollEvent handler thread
    let epfd1 = epfd.clone();
    let sockets1 = sockets.clone();
    let (tx_ep_event, rx_ep_event): (Sender<EpollEvent>, Receiver<EpollEvent>) = channel();
    let safe_handler = Arc::new(Mutex::new(*handler));
    thread::Builder::new()
        .name("EpollEvent Handler".to_string())
        .spawn(move || {
            epoll_event_handler(epfd1, rx_ep_event, sockets1, safe_handler);
        }).unwrap();

    // Epoll wait thread
    let epfd2 = epfd.clone();
    thread::Builder::new()
        .name("Epoll Wait".to_string())
        .spawn(move || {
            epoll_wait_loop(epfd2, tx_ep_event);
        }).unwrap();

    // New connection thread
    let epfd3 = epfd.clone();
    let sockets2 = sockets.clone();
    let prox = thread::Builder::new()
        .name("TCP Incoming Listener".to_string())
        .spawn(move || {
            listen(address, epfd3, sockets2);
        }).unwrap();

    // Stay alive forever
    let _ = prox.join();
}

fn listen<T: ToSocketAddrs>(address: T, epfd: RawFd, sockets: SocketList) {
    let listener = TcpListener::bind(address).unwrap();
    for connection in listener.incoming() {
        if connection.is_err() {
            warn!("Error encountered during TCP connection: {}", connection.unwrap_err());
            continue;
        }

        let result = NbetStream::new(connection.unwrap());
        if result.is_err() {
            warn!("Error creating nbstream: {:?}", result.unwrap_err());
            continue;
        }

        let socket_ptr;
        let socket = Socket::new(result.unwrap());
        { // Begin Mutex lock
            let mut guard = match sockets.lock() {
                Ok(guard) => guard,
                Err(poisoned) => {
                    warn!("SocketList Mutex failed, using anyway...");
                    poisoned.into_inner()
                }
            };
            let socket_list = guard.deref_mut();
            socket_list.push_back(socket.clone());
            socket_ptr = socket_list.back_mut().unwrap() as *mut Socket;
        } // End Mutex lock

        // Add to epoll
        let mut event = EpollEvent {
            data: (socket_ptr as usize) as u64,
            events: EVENTS
        };
        let _ = epoll::ctl(epfd, ctl_op::ADD, socket.raw_fd(), &mut event).map_err(|e| {
            error!("CtrlError during add: {}", e);
        });
    }
    drop(listener);
}

fn epoll_wait_loop(epfd: RawFd, tx_handler: Sender<EpollEvent>) {
    // Maximum number of events we want to be notified of at one time
    let mut events = Vec::<EpollEvent>::with_capacity(100);
    unsafe { events.set_len(100); }

    loop {
        match epoll::wait(epfd, &mut events[..], -1) {
            Ok(num_events) => {
                for x in 0..num_events as usize {
                    let _ = tx_handler.send(events[x]);
                }
            }
            Err(e) => {
                error!("Error on epoll::wait(): {}", e);
                panic!()
            }
        }
    }
}

fn epoll_event_handler(epfd: RawFd,
                       rx: Receiver<EpollEvent>,
                       sockets: SocketList,
                       handler: SafeHandler) {
    // Types of events we care about
    const READ_EVENT: u32 = event_type::EPOLLIN;
    const DROP_EVENT: u32 = event_type::EPOLLRDHUP | event_type::EPOLLHUP | event_type::EPOLLERR;

    // Thread pool
    let mut pool = ResourcePool::new();

    for event in rx.iter() {
        // Get a pointer to the socket in the list
        let socket_ptr;
        { // Begin Mutex lock
            let _ = match sockets.lock() {
                Ok(guard) => guard,
                Err(poisoned) => {
                    warn!("SocketList Mutex failed, using anyway...");
                    poisoned.into_inner()
                }
            };
            socket_ptr = event.data as *mut Socket;
        } // End Mutex lock

        // Move semantics are a little weird right here, so we must use clones of clones
        // in order to make the compiler happy
        if (event.events & DROP_EVENT) > 0 {
            let t_handler = handler.clone();
            let socket_list = sockets.clone();
            let socket = unsafe { (*socket_ptr).clone() };
            pool.run(move || {
                handle_drop_event(epfd, socket.clone(), socket_list.clone(), t_handler.clone());
            });
        } else if (event.events & READ_EVENT) > 0 {
            let t_handler = handler.clone();
            let socket_list = sockets.clone();
            let socket = unsafe { (*socket_ptr).clone() };
            let ptr_addr = socket_ptr as usize;
            pool.run(move || {
                handle_read_event(epfd, socket.clone(), socket_list.clone(),
                    t_handler.clone(), ptr_addr);
            });
        } else { warn!("Unknown epoll event received: {}", event.data); continue; }
    }
}

fn handle_read_event(epfd: RawFd,
                     socket: Socket,
                     sockets: SocketList,
                     handler: SafeHandler,
                     ptr_addr: usize) {
    let mut socket = socket;
    let result = socket.read();
    if result.is_err() {
        // There is no saveing grace here. If we can't read, only option
        // is to remove it from server
        handle_drop_event(epfd, socket, sockets, handler);
        return;
    }

    for msg in socket.buffer().iter() {
        // Pass the msg on to the handler
        { // EventHandler Mutex locked
            let mut guard = match handler.lock() {
                Ok(guard) => guard,
                Err(poisoned) => {
                    warn!("EventHandler Mutex failed, using anyway...");
                    poisoned.into_inner()
                }
            };
            let e_handler = guard.deref_mut();
            e_handler.on_data_received(socket.clone(), sockets.clone(), msg.clone());
        } // EventHandler Mutex unlocked

        // Collect some stats
        stats::msg_recv();
        stats::bytes_recv(msg.len());

        // Add socket's fd to epoll watch list
        //
        // std::net::TcpStream uses F_DUPFD flag when calling clone(). This means that
        // it is very likely that the read event we received for this socket may not actually
        // be the exact same file descriptor (RawFd) as the one that registered it. Therefore,
        // we have several steps to take in order to try and re-register the same one or add it
        // again if epoll tells us the fd does not exist.
        // This approach does not have any ill effects though, Rust happily sets the CLOEXEC flag
        // on all owned fds, so once any of our sockets go out of scope, the kernel will clean
        // them up for us. Epoll, as it happens, removes fds from the watch list also, once the
        // kernel cleans them up.
        // We start off with an add, because we're assuming the socket is a clone and fds are
        // different
        let mut event = EpollEvent {
            data: ptr_addr as u64,
            events: EVENTS
        };
        let _ = epoll::ctl(epfd, ctl_op::ADD, socket.raw_fd(), &mut event).map_err(|e| {
            if e == CtlError::EEXIST {
                // Fd does not currently exist in epoll, we'll add it then because it must have
                // been a F_DUPFD'd fd from one of the cloned sockets.
                let _ = epoll::ctl(epfd, ctl_op::MOD, socket.raw_fd(), &mut event).map_err(|e| {
                    error!("During epoll::ctl: {}", e);
                });
            } else {
                error!("During epoll::ctl: {}", e);
            }
        });
    }
}

fn handle_drop_event(epfd: RawFd, socket: Socket, sockets: SocketList, handler: SafeHandler) {
    epoll_remove_fd(epfd, socket.raw_fd());
    remove_socket_from_list(socket.id(), sockets);

    let mut guard = match handler.lock() {
        Ok(g) => g,
        Err(p) => {
            warn!("EventHandler Mutex failed, using anyway...");
            p.into_inner()
        }
    };
    let e_handler = guard.deref_mut();
    e_handler.on_socket_closed(socket.id());
}

fn epoll_remove_fd(epfd: RawFd, fd: RawFd) {
    debug!("remove_socket_from_epoll");

    // In kernel versions before 2.6.9, the EPOLL_CTL_DEL operation required
    // a non-null pointer in event, even though this argument is ignored.
    // Since Linux 2.6.9, event can be specified as NULL when using
    // EPOLL_CTL_DEL.  Applications that need to be portable to kernels
    // before 2.6.9 should specify a non-null pointer in event.
    let mut rm_event = EpollEvent {
        data: 0 as u64,
        events: 0 as u32
    };

    // Depending on how fds are duplicated with .clone(), this may fail.
    //
    // If the failure case is CtlError::ENOENT, we do not care, because
    // epoll will clean the up the descriptor after they are all dropped from
    // program memory
    match epoll::ctl(epfd, ctl_op::DEL, fd, &mut rm_event) {
        Ok(()) => trace!("Socket removed from epoll watch list"),
        Err(e) => {
            match e {
                CtlError::ENOENT => {
                    trace!("Fd not found in epoll, will remove when fd is syscall closed");
                }
                _ => warn!("Epoll CtrlError during del: {}", e)
            }
        }
    };
}

/// Removes socket_ids from master list
fn remove_socket_from_list(socket_id: u32, sockets: SocketList) {
    debug!("remove_socket_from_list");

    let mut s_guard = match sockets.lock() {
        Ok(guard) => guard,
        Err(poisoned) => {
            warn!("SocketList Mutex was poisoned, using anyway");
            poisoned.into_inner()
        }
    };
    let s_list = s_guard.deref_mut();

    let mut s_found = false;
    let mut index = 1usize;
    for s in s_list.iter() {
        if s.id() == socket_id {
            s_found = true;
            break;
        }
        index += 1;
    }

    if !s_found { return; }

    if index == 1 {
        s_list.pop_front();
    } else {
        let mut split = s_list.split_off(index - 1);
        split.pop_front();
        s_list.append(&mut split);
    }

    stats::conn_lost();

    trace!("socket removed from list");
}
<|endoftext|>