<commit_before>use serde::de::{Deserialize, Deserializer, Error};

use types::*;

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub struct Message {
    pub id: Integer,
    pub from: Option<User>,
    pub date: Integer,
    pub chat: Chat,
    pub forward: Option<Forward>,
    pub reply_to_message: Option<Box<Message>>,
    pub edit_date: Option<Integer>,
    pub caption: Option<String>,
    pub kind: MessageKind,
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub struct Forward {
    pub date: Integer,
    pub from: ForwardFrom,
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub enum ForwardFrom {
    User {
        user: User,
    },
    Channel {
        channel: Channel,
        message_id: Integer
    }
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub enum MessageKind {
    Text {
        data: String,
        entities: Vec<MessageEntity>,
    },
    Audio {
        data: Audio,
    },
    Document {
        data: Document,
    },
    Photo {
        data: Vec<PhotoSize>,
    },
    Sticker {
        data: Sticker,
    },
    Video {
        data: Video,
    },
    Voice {
        data: Voice,
    },
    Contact {
        data: Contact,
    },
    Location {
        data: Location,
    },
    Venue {
        data: Venue,
    },
    NewChatMember {
        data: User,
    },
    LeftChatMember {
        data: User,
    },
    NewChatTitle {
        data: String,
    },
    NewChatPhoto {
        data: PhotoSize,
    },
    DeleteChatPhoto,
    GroupChatCreated,
    SupergroupChatCreated,
    ChannelChatCreated,
    MigrateToChatId {
        data: Integer,
    },
    MigrateFromChatId {
        data: Integer,
    },
    PinnedMessage {
        data: Box<Message>,
    },
    #[doc(hidden)]
    Unknown {
        raw: RawMessage
    }
}

impl Deserialize for Message {  // TODO(knsd): Remove .clone()
    fn deserialize<D>(deserializer: D) -> Result<Message, D::Error> where D: Deserializer {
        let raw: RawMessage = Deserialize::deserialize(deserializer)?;

        let id = raw.message_id;
        let from = raw.from.clone();
        let date = raw.date;
        let chat = raw.chat.clone();
        let reply_to_message = raw.reply_to_message.clone();
        let edit_date = raw.edit_date;
        let caption = raw.caption.clone();

        let forward = match (raw.forward_date, &raw.forward_from,
                             &raw.forward_from_chat, raw.forward_from_message_id) {
            (None, &None, &None, None) => None,
            (Some(date), &Some(ref from), &None, None) => {
                Some(Forward {
                    date: date,
                    from: ForwardFrom::User {
                        user: from.clone(),
                    },
                })
            },
            (Some(date), &None, &Some(Chat::Channel(ref channel)), Some(message_id)) => {
                Some(Forward {
                    date: date,
                    from: ForwardFrom::Channel {
                        channel: channel.clone(),
                        message_id: message_id,
                    },
                })
            },
            _ => {
                return Err(D::Error::custom("invalid forward fields combination"))
            }
        };

        let make_message = |kind| {
            Ok(Message {
                id: id,
                from: from,
                date: date,
                chat: chat,
                forward: forward,
                reply_to_message: reply_to_message,
                edit_date: edit_date,
                caption: caption,
                kind: kind,
            })
        };

        macro_rules! maybe_field {
            ($name:ident, $variant:ident) => {{
                if let Some(val) = raw.$name {
                    return make_message(MessageKind::$variant {
                        data: val
                    })
                }
            }}
        }

        macro_rules! maybe_true_field {
            ($name:ident, $variant:ident) => {{
                if let Some(True) = raw.$name {
                    return make_message(MessageKind::$variant)
                }
            }}
        }

        if let Some(text) = raw.text {
            let entities = raw.entities.unwrap_or_else(Vec::new);
            return make_message(MessageKind::Text {
                data: text,
                entities: entities,
            })
        }

        maybe_field!(audio, Audio);
        maybe_field!(document, Document);
        maybe_field!(photo, Photo);
        maybe_field!(sticker, Sticker);
        maybe_field!(video, Video);
        maybe_field!(voice, Voice);
        maybe_field!(contact, Contact);
        maybe_field!(location, Location);
        maybe_field!(venue, Venue);
        maybe_field!(new_chat_member, NewChatMember);
        maybe_field!(left_chat_member, LeftChatMember);
        maybe_field!(new_chat_title, NewChatTitle);
        maybe_field!(new_chat_photo, NewChatPhoto);
        maybe_true_field!(delete_chat_photo, DeleteChatPhoto);
        maybe_true_field!(delete_chat_photo, DeleteChatPhoto);
        maybe_true_field!(group_chat_created, GroupChatCreated);
        maybe_true_field!(supergroup_chat_created, SupergroupChatCreated);
        maybe_true_field!(channel_chat_created, ChannelChatCreated);
        maybe_field!(migrate_to_chat_id, MigrateToChatId);
        maybe_field!(migrate_from_chat_id, MigrateFromChatId);
        maybe_field!(pinned_message, PinnedMessage);

        make_message(MessageKind::Unknown {
            raw: raw,
        })
    }
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct RawMessage {
    pub message_id: Integer,
    pub from: Option<User>,
    pub date: Integer,
    pub chat: Chat,
    pub forward_from: Option<User>,
    pub forward_from_chat: Option<Chat>,
    pub forward_from_message_id: Option<Integer>,
    pub forward_date: Option<Integer>,
    pub reply_to_message: Option<Box<Message>>,
    pub edit_date: Option<Integer>,
    pub text: Option<String>,
    pub entities: Option<Vec<MessageEntity>>,
    pub audio: Option<Audio>,
    pub document: Option<Document>,
//    pub game: Option<Game>,
    pub photo: Option<Vec<PhotoSize>>,
    pub sticker: Option<Sticker>,
    pub video: Option<Video>,
    pub voice: Option<Voice>,
    pub caption: Option<String>,
    pub contact: Option<Contact>,
    pub location: Option<Location>,
    pub venue: Option<Venue>,
    pub new_chat_member: Option<User>,
    pub left_chat_member: Option<User>,
    pub new_chat_title: Option<String>,
    pub new_chat_photo: Option<PhotoSize>,
    pub delete_chat_photo: Option<True>,
    pub group_chat_created: Option<True>,
    pub supergroup_chat_created: Option<True>,
    pub channel_chat_created: Option<True>,
    pub migrate_to_chat_id: Option<Integer>,
    pub migrate_from_chat_id: Option<Integer>,
    pub pinned_message: Option<Box<Message>>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub struct MessageEntity {
    offset: Integer,
    length: Integer,
    kind: MessageEntityKind,
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub enum MessageEntityKind {
    Mention,
    Hashtag,
    BotCommand,
    Url,
    Email,
    Bold,
    Italic,
    Code,
    Pre,
    TextLink(String), // TODO(knsd) URL?
    TextMention(User),
    #[doc(hidden)]
    Unknown(RawMessageEntity),
}

impl Deserialize for MessageEntity {
    fn deserialize<D>(deserializer: D) -> Result<MessageEntity, D::Error> where D: Deserializer {
        use self::MessageEntityKind::*;

        let raw: RawMessageEntity = Deserialize::deserialize(deserializer)?;

        let offset = raw.offset;
        let length = raw.length;

        macro_rules! required_field {
            ($name:ident) => {{
                match raw.$name {
                    Some(val) => val,
                    None => return Err(D::Error::missing_field(stringify!($name)))
                }
            }}
        }

        let kind = match raw.type_.as_str() {
            "mention" => Mention,
            "hashtag" => Hashtag,
            "bot_command" => BotCommand,
            "url" => Url,
            "email" => Email,
            "bold" => Bold,
            "italic" => Italic,
            "code" => Code,
            "pre" => Pre,
            "text_link" => TextLink(required_field!(url)),
            "text_mention" => TextMention(required_field!(user)),
            _ => Unknown(raw),
        };

        Ok(MessageEntity {
            offset: offset,
            length: length,
            kind: kind,
        })
    }
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct RawMessageEntity {
    #[serde(rename="type")]
    pub type_: String,
    pub offset: Integer,
    pub length: Integer,
    pub url: Option<String>,
    pub user: Option<User>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct PhotoSize {
    pub file_id: String,
    pub width: Integer,
    pub height: Integer,
    pub file_size: Option<Integer>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Audio {
    pub file_id: String,
    pub duration: Integer,
    pub performer: Option<String>,
    pub title: Option<String>,
    pub mime_type: Option<String>,
    pub file_size: Option<Integer>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Document {
    pub file_id: String,
    pub thumb: Option<PhotoSize>,
    pub file_name: Option<String>,
    pub mime_type: Option<String>,
    pub file_size: Option<Integer>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Sticker {
    pub file_id: String,
    pub width: Integer,
    pub height: Integer,
    pub thumb: Option<PhotoSize>,
    pub emoji: Option<String>,
    pub file_size: Option<Integer>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Video {
    pub file_id: String,
    pub width: Integer,
    pub height: Integer,
    pub duration: Integer,
    pub thumb: Option<PhotoSize>,
    pub mime_type: Option<String>,
    pub file_size: Option<Integer>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Voice {
    pub file_id: String,
    pub duration: Integer,
    pub mime_type: Option<String>,
    pub file_size: Option<Integer>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Contact {
    pub phone_number: String,
    pub first_name: String,
    pub last_name: Option<String>,
    pub user_id: Option<Integer>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Location {
    pub longitude: Float,
    pub latitude: Float,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Venue {
    pub location: Location,
    pub title: String,
    pub address: String,
    pub foursquare_id: Option<String>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct UserProfilePhotos {
    pub total_count: Integer,
    pub photos: Vec<Vec<PhotoSize>>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct File {
    pub file_id: String,
    pub file_size: Option<Integer>,
    pub file_path: Option<String>,
}<commit_msg>docs: document some message types<commit_after>use serde::de::{Deserialize, Deserializer, Error};

use types::*;

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub struct Message {
    pub id: Integer,
    pub from: Option<User>,
    pub date: Integer,
    pub chat: Chat,
    pub forward: Option<Forward>,
    pub reply_to_message: Option<Box<Message>>,
    pub edit_date: Option<Integer>,
    pub caption: Option<String>,
    pub kind: MessageKind,
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub struct Forward {
    pub date: Integer,
    pub from: ForwardFrom,
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub enum ForwardFrom {
    User {
        user: User,
    },
    Channel {
        channel: Channel,
        message_id: Integer
    }
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub enum MessageKind {
    Text {
        data: String,
        entities: Vec<MessageEntity>,
    },
    Audio {
        data: Audio,
    },
    Document {
        data: Document,
    },
    Photo {
        data: Vec<PhotoSize>,
    },
    Sticker {
        data: Sticker,
    },
    Video {
        data: Video,
    },
    Voice {
        data: Voice,
    },
    Contact {
        data: Contact,
    },
    Location {
        data: Location,
    },
    Venue {
        data: Venue,
    },
    NewChatMember {
        data: User,
    },
    LeftChatMember {
        data: User,
    },
    NewChatTitle {
        data: String,
    },
    NewChatPhoto {
        data: PhotoSize,
    },
    DeleteChatPhoto,
    GroupChatCreated,
    SupergroupChatCreated,
    ChannelChatCreated,
    MigrateToChatId {
        data: Integer,
    },
    MigrateFromChatId {
        data: Integer,
    },
    PinnedMessage {
        data: Box<Message>,
    },
    #[doc(hidden)]
    Unknown {
        raw: RawMessage
    }
}

impl Deserialize for Message {  // TODO(knsd): Remove .clone()
    fn deserialize<D>(deserializer: D) -> Result<Message, D::Error> where D: Deserializer {
        let raw: RawMessage = Deserialize::deserialize(deserializer)?;

        let id = raw.message_id;
        let from = raw.from.clone();
        let date = raw.date;
        let chat = raw.chat.clone();
        let reply_to_message = raw.reply_to_message.clone();
        let edit_date = raw.edit_date;
        let caption = raw.caption.clone();

        let forward = match (raw.forward_date, &raw.forward_from,
                             &raw.forward_from_chat, raw.forward_from_message_id) {
            (None, &None, &None, None) => None,
            (Some(date), &Some(ref from), &None, None) => {
                Some(Forward {
                    date: date,
                    from: ForwardFrom::User {
                        user: from.clone(),
                    },
                })
            },
            (Some(date), &None, &Some(Chat::Channel(ref channel)), Some(message_id)) => {
                Some(Forward {
                    date: date,
                    from: ForwardFrom::Channel {
                        channel: channel.clone(),
                        message_id: message_id,
                    },
                })
            },
            _ => {
                return Err(D::Error::custom("invalid forward fields combination"))
            }
        };

        let make_message = |kind| {
            Ok(Message {
                id: id,
                from: from,
                date: date,
                chat: chat,
                forward: forward,
                reply_to_message: reply_to_message,
                edit_date: edit_date,
                caption: caption,
                kind: kind,
            })
        };

        macro_rules! maybe_field {
            ($name:ident, $variant:ident) => {{
                if let Some(val) = raw.$name {
                    return make_message(MessageKind::$variant {
                        data: val
                    })
                }
            }}
        }

        macro_rules! maybe_true_field {
            ($name:ident, $variant:ident) => {{
                if let Some(True) = raw.$name {
                    return make_message(MessageKind::$variant)
                }
            }}
        }

        if let Some(text) = raw.text {
            let entities = raw.entities.unwrap_or_else(Vec::new);
            return make_message(MessageKind::Text {
                data: text,
                entities: entities,
            })
        }

        maybe_field!(audio, Audio);
        maybe_field!(document, Document);
        maybe_field!(photo, Photo);
        maybe_field!(sticker, Sticker);
        maybe_field!(video, Video);
        maybe_field!(voice, Voice);
        maybe_field!(contact, Contact);
        maybe_field!(location, Location);
        maybe_field!(venue, Venue);
        maybe_field!(new_chat_member, NewChatMember);
        maybe_field!(left_chat_member, LeftChatMember);
        maybe_field!(new_chat_title, NewChatTitle);
        maybe_field!(new_chat_photo, NewChatPhoto);
        maybe_true_field!(delete_chat_photo, DeleteChatPhoto);
        maybe_true_field!(delete_chat_photo, DeleteChatPhoto);
        maybe_true_field!(group_chat_created, GroupChatCreated);
        maybe_true_field!(supergroup_chat_created, SupergroupChatCreated);
        maybe_true_field!(channel_chat_created, ChannelChatCreated);
        maybe_field!(migrate_to_chat_id, MigrateToChatId);
        maybe_field!(migrate_from_chat_id, MigrateFromChatId);
        maybe_field!(pinned_message, PinnedMessage);

        make_message(MessageKind::Unknown {
            raw: raw,
        })
    }
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct RawMessage {
    pub message_id: Integer,
    pub from: Option<User>,
    pub date: Integer,
    pub chat: Chat,
    pub forward_from: Option<User>,
    pub forward_from_chat: Option<Chat>,
    pub forward_from_message_id: Option<Integer>,
    pub forward_date: Option<Integer>,
    pub reply_to_message: Option<Box<Message>>,
    pub edit_date: Option<Integer>,
    pub text: Option<String>,
    pub entities: Option<Vec<MessageEntity>>,
    pub audio: Option<Audio>,
    pub document: Option<Document>,
//    pub game: Option<Game>,
    pub photo: Option<Vec<PhotoSize>>,
    pub sticker: Option<Sticker>,
    pub video: Option<Video>,
    pub voice: Option<Voice>,
    pub caption: Option<String>,
    pub contact: Option<Contact>,
    pub location: Option<Location>,
    pub venue: Option<Venue>,
    pub new_chat_member: Option<User>,
    pub left_chat_member: Option<User>,
    pub new_chat_title: Option<String>,
    pub new_chat_photo: Option<PhotoSize>,
    pub delete_chat_photo: Option<True>,
    pub group_chat_created: Option<True>,
    pub supergroup_chat_created: Option<True>,
    pub channel_chat_created: Option<True>,
    pub migrate_to_chat_id: Option<Integer>,
    pub migrate_from_chat_id: Option<Integer>,
    pub pinned_message: Option<Box<Message>>,
}

/// This object represents one special entity in a text message.
/// For example, hashtags, usernames, URLs, etc.
#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub struct MessageEntity {
    /// Offset in UTF-16 code units to the start of the entity
    offset: Integer,
    /// Length of the entity in UTF-16 code units
    length: Integer,
    /// Kind of the entity.
    kind: MessageEntityKind,
}

/// Kind of the entity.
#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub enum MessageEntityKind {
    Mention,
    Hashtag,
    BotCommand,
    Url,
    Email,
    Bold,
    Italic,
    Code,
    Pre,
    TextLink(String), // TODO(knsd) URL?
    TextMention(User),
    #[doc(hidden)]
    Unknown(RawMessageEntity),
}

impl Deserialize for MessageEntity {
    fn deserialize<D>(deserializer: D) -> Result<MessageEntity, D::Error> where D: Deserializer {
        use self::MessageEntityKind::*;

        let raw: RawMessageEntity = Deserialize::deserialize(deserializer)?;

        let offset = raw.offset;
        let length = raw.length;

        macro_rules! required_field {
            ($name:ident) => {{
                match raw.$name {
                    Some(val) => val,
                    None => return Err(D::Error::missing_field(stringify!($name)))
                }
            }}
        }

        let kind = match raw.type_.as_str() {
            "mention" => Mention,
            "hashtag" => Hashtag,
            "bot_command" => BotCommand,
            "url" => Url,
            "email" => Email,
            "bold" => Bold,
            "italic" => Italic,
            "code" => Code,
            "pre" => Pre,
            "text_link" => TextLink(required_field!(url)),
            "text_mention" => TextMention(required_field!(user)),
            _ => Unknown(raw),
        };

        Ok(MessageEntity {
            offset: offset,
            length: length,
            kind: kind,
        })
    }
}

/// This object represents one special entity in a text message.
/// For example, hashtags, usernames, URLs, etc. Directly mapped.
#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct RawMessageEntity {
    /// Type of the entity. Can be mention (@username), hashtag, bot_command, url, email,
    /// bold (bold text), italic (italic text), code (monowidth string), pre (monowidth block),
    /// text_link (for clickable text URLs), text_mention (for users without usernames).
    #[serde(rename="type")]
    pub type_: String,
    /// Offset in UTF-16 code units to the start of the entity.
    pub offset: Integer,
    /// Length of the entity in UTF-16 code units.
    pub length: Integer,
    /// For “text_link” only, url that will be opened after user taps on the text.
    pub url: Option<String>,
    /// For “text_mention” only, the mentioned user.
    pub user: Option<User>,
}

/// This object represents one size of a photo or a file / sticker thumbnail.
#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct PhotoSize {
    /// Unique identifier for this file.
    pub file_id: String,
    /// Photo width.
    pub width: Integer,
    /// Photo height.
    pub height: Integer,
    /// File size.
    pub file_size: Option<Integer>,
}

/// This object represents an audio file to be treated as music by the Telegram clients.
#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Audio {
    /// Unique identifier for this file.
    pub file_id: String,
    /// Duration of the audio in seconds as defined by sender.
    pub duration: Integer,
    /// Performer of the audio as defined by sender or by audio tags.
    pub performer: Option<String>,
    /// Title of the audio as defined by sender or by audio tags.
    pub title: Option<String>,
    /// MIME type of the file as defined by sender.
    pub mime_type: Option<String>,
    /// File size.
    pub file_size: Option<Integer>,
}

/// This object represents a general file (as opposed to photos, voice messages and audio files).
#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Document {
    /// Unique file identifier.
    pub file_id: String,
    /// Document thumbnail as defined by sender.
    pub thumb: Option<PhotoSize>,
    /// Original filename as defined by sender.
    pub file_name: Option<String>,
    /// MIME type of the file as defined by sender.
    pub mime_type: Option<String>,
    /// File size.
    pub file_size: Option<Integer>,
}

/// This object represents a sticker.
#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Sticker {
    /// Unique identifier for this file.
    pub file_id: String,
    /// Sticker width.
    pub width: Integer,
    /// Sticker height.
    pub height: Integer,
    /// Sticker thumbnail in .webp or .jpg format.
    pub thumb: Option<PhotoSize>,
    /// Emoji associated with the sticker.
    pub emoji: Option<String>,
    /// File size.
    pub file_size: Option<Integer>,
}

/// This object represents a video file.
#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Video {
    /// Unique identifier for this file.
    pub file_id: String,
    /// Video width as defined by sender.
    pub width: Integer,
    /// Video height as defined by sender.
    pub height: Integer,
    /// Duration of the video in seconds as defined by sender.
    pub duration: Integer,
    /// Video thumbnail.
    pub thumb: Option<PhotoSize>,
    /// Mime type of a file as defined by sender.
    pub mime_type: Option<String>,
    /// File size.
    pub file_size: Option<Integer>,
}

/// This object represents a voice note.
#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Voice {
    /// Unique identifier for this file.
    pub file_id: String,
    /// Duration of the audio in seconds as defined by sender.
    pub duration: Integer,
    /// MIME type of the file as defined by sender.
    pub mime_type: Option<String>,
    /// File size.
    pub file_size: Option<Integer>,
}

/// This object represents a phone contact.
#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Contact {
    /// Contact's phone number.
    pub phone_number: String,
    /// Contact's first name.
    pub first_name: String,
    /// Contact's last name.
    pub last_name: Option<String>,
    /// Contact's user identifier in Telegram.
    pub user_id: Option<Integer>,
}

/// This object represents a point on the map.
#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Location {
    /// Longitude as defined by sender.
    pub longitude: Float,
    /// Latitude as defined by sender.
    pub latitude: Float,
}

/// This object represents a venue.
#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Venue {
    /// Venue location.
    pub location: Location,
    /// Name of the venue.
    pub title: String,
    /// Address of the venue.
    pub address: String,
    /// Foursquare identifier of the venue.
    pub foursquare_id: Option<String>,
}

/// This object represent a user's profile pictures.
#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct UserProfilePhotos {
    /// Total number of profile pictures the target user has.
    pub total_count: Integer,
    /// Requested profile pictures (in up to 4 sizes each).
    pub photos: Vec<Vec<PhotoSize>>,
}

/// This object represents a file ready to be downloaded.
/// The file can be downloaded via the link https://api.telegram.org/file/bot<token>/<file_path>.
/// It is guaranteed that the link will be valid for at least 1 hour.
#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct File {
    /// Unique identifier for this file.
    pub file_id: String,
    /// File size, if known.
    pub file_size: Option<Integer>,
    /// File path. Use https://api.telegram.org/file/bot<token>/<file_path> to get the file.
    pub file_path: Option<String>,
}
<|endoftext|>


<commit_before>use crate::codec::*;
use crate::endpoint::Endpoint;
use crate::error::ZmqResult;
use crate::message::*;
use crate::transport::AcceptStopHandle;
use crate::util::PeerIdentity;
use crate::{async_rt, CaptureSocket, SocketOptions};
use crate::{
    MultiPeerBackend, Socket, SocketBackend, SocketEvent, SocketSend, SocketType, ZmqError,
};

use async_trait::async_trait;
use dashmap::DashMap;
use futures::channel::{mpsc, oneshot};
use futures::FutureExt;
use parking_lot::Mutex;
use std::collections::HashMap;
use std::io::ErrorKind;
use std::pin::Pin;
use std::sync::Arc;

pub(crate) struct Subscriber {
    pub(crate) subscriptions: Vec<Vec<u8>>,
    pub(crate) send_queue: Pin<Box<ZmqFramedWrite>>,
    _subscription_coro_stop: oneshot::Sender<()>,
}

pub(crate) struct PubSocketBackend {
    subscribers: DashMap<PeerIdentity, Subscriber>,
    socket_monitor: Mutex<Option<mpsc::Sender<SocketEvent>>>,
    socket_options: SocketOptions,
}

impl PubSocketBackend {
    fn message_received(&self, peer_id: &PeerIdentity, message: Message) {
        let message = match message {
            Message::Message(m) => m,
            _ => return,
        };
        assert_eq!(message.len(), 1);
        let data: Vec<u8> = message.into_vec().pop().unwrap().to_vec();
        if data.is_empty() {
            return;
        }
        match data[0] {
            1 => {
                // Subscribe
                self.subscribers
                    .get_mut(&peer_id)
                    .unwrap()
                    .subscriptions
                    .push(Vec::from(&data[1..]));
            }
            0 => {
                // Unsubscribe
                let mut del_index = None;
                let sub = Vec::from(&data[1..]);
                for (idx, subscription) in self
                    .subscribers
                    .get(&peer_id)
                    .unwrap()
                    .subscriptions
                    .iter()
                    .enumerate()
                {
                    if &sub == subscription {
                        del_index = Some(idx);
                        break;
                    }
                }
                if let Some(index) = del_index {
                    self.subscribers
                        .get_mut(&peer_id)
                        .unwrap()
                        .subscriptions
                        .remove(index);
                }
            }
            _ => (),
        }
    }
}

impl SocketBackend for PubSocketBackend {
    fn socket_type(&self) -> SocketType {
        SocketType::PUB
    }

    fn socket_options(&self) -> &SocketOptions {
        &self.socket_options
    }

    fn shutdown(&self) {
        self.subscribers.clear();
    }

    fn monitor(&self) -> &Mutex<Option<mpsc::Sender<SocketEvent>>> {
        &self.socket_monitor
    }
}

impl MultiPeerBackend for PubSocketBackend {
    fn peer_connected(self: Arc<Self>, peer_id: &PeerIdentity, io: FramedIo) {
        let (mut recv_queue, send_queue) = io.into_parts();
        // TODO provide handling for recv_queue
        let (sender, stop_receiver) = oneshot::channel();
        self.subscribers.insert(
            peer_id.clone(),
            Subscriber {
                subscriptions: vec![],
                send_queue: Box::pin(send_queue),
                _subscription_coro_stop: sender,
            },
        );
        let backend = self;
        let peer_id = peer_id.clone();
        async_rt::task::spawn(async move {
            use futures::StreamExt;
            let mut stop_receiver = stop_receiver.fuse();
            loop {
                futures::select! {
                     _ = stop_receiver => {
                         break;
                     },
                     message = recv_queue.next().fuse() => {
                        match message {
                            Some(Ok(m)) => backend.message_received(&peer_id, m),
                            Some(Err(e)) => {
                                dbg!(e);
                                backend.peer_disconnected(&peer_id);
                                break;
                            }
                            None => {
                                backend.peer_disconnected(&peer_id);
                                break
                            }
                        }

                     }
                }
            }
        });
    }

    fn peer_disconnected(&self, peer_id: &PeerIdentity) {
        log::info!("Client disconnected {:?}", peer_id);
        self.subscribers.remove(peer_id);
    }
}

pub struct PubSocket {
    pub(crate) backend: Arc<PubSocketBackend>,
    binds: HashMap<Endpoint, AcceptStopHandle>,
}

impl Drop for PubSocket {
    fn drop(&mut self) {
        self.backend.shutdown();
    }
}

#[async_trait]
impl SocketSend for PubSocket {
    async fn send(&mut self, message: ZmqMessage) -> ZmqResult<()> {
        let mut dead_peers = Vec::new();
        for mut subscriber in self.backend.subscribers.iter_mut() {
            for sub_filter in &subscriber.subscriptions {
                if sub_filter.len() <= message.get(0).unwrap().len(){
                if sub_filter.as_slice() == &message.get(0).unwrap()[0..sub_filter.len()] {
                    let res = subscriber
                        .send_queue
                        .as_mut()
                        .try_send(Message::Message(message.clone()));
                    match res {
                        Ok(()) => {}
                        Err(ZmqError::Codec(CodecError::Io(e))) => {
                            if e.kind() == ErrorKind::BrokenPipe {
                                dead_peers.push(subscriber.key().clone());
                            } else {
                                dbg!(e);
                            }
                        }
                        Err(e) => {
                            dbg!(e);
                            todo!()
                        }
                    }
                    break;
                }
                }
            }
        }
        for peer in dead_peers {
            self.backend.peer_disconnected(&peer);
        }
        Ok(())
    }
}

impl CaptureSocket for PubSocket {}

#[async_trait]
impl Socket for PubSocket {
    fn with_options(options: SocketOptions) -> Self {
        Self {
            backend: Arc::new(PubSocketBackend {
                subscribers: DashMap::new(),
                socket_monitor: Mutex::new(None),
                socket_options: options,
            }),
            binds: HashMap::new(),
        }
    }

    fn backend(&self) -> Arc<dyn MultiPeerBackend> {
        self.backend.clone()
    }

    fn binds(&mut self) -> &mut HashMap<Endpoint, AcceptStopHandle> {
        &mut self.binds
    }

    fn monitor(&mut self) -> mpsc::Receiver<SocketEvent> {
        let (sender, receiver) = mpsc::channel(1024);
        self.backend.socket_monitor.lock().replace(sender);
        receiver
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::util::tests::{
        test_bind_to_any_port_helper, test_bind_to_unspecified_interface_helper,
    };
    use crate::ZmqResult;
    use std::net::IpAddr;

    #[async_rt::test]
    async fn test_bind_to_any_port() -> ZmqResult<()> {
        let s = PubSocket::new();
        test_bind_to_any_port_helper(s).await
    }

    #[async_rt::test]
    async fn test_bind_to_any_ipv4_interface() -> ZmqResult<()> {
        let any_ipv4: IpAddr = "0.0.0.0".parse().unwrap();
        let s = PubSocket::new();
        test_bind_to_unspecified_interface_helper(any_ipv4, s, 4000).await
    }

    #[async_rt::test]
    async fn test_bind_to_any_ipv6_interface() -> ZmqResult<()> {
        let any_ipv6: IpAddr = "::".parse().unwrap();
        let s = PubSocket::new();
        test_bind_to_unspecified_interface_helper(any_ipv6, s, 4010).await
    }
}
<commit_msg>Clippy fix<commit_after>use crate::codec::*;
use crate::endpoint::Endpoint;
use crate::error::ZmqResult;
use crate::message::*;
use crate::transport::AcceptStopHandle;
use crate::util::PeerIdentity;
use crate::{async_rt, CaptureSocket, SocketOptions};
use crate::{
    MultiPeerBackend, Socket, SocketBackend, SocketEvent, SocketSend, SocketType, ZmqError,
};

use async_trait::async_trait;
use dashmap::DashMap;
use futures::channel::{mpsc, oneshot};
use futures::FutureExt;
use parking_lot::Mutex;
use std::collections::HashMap;
use std::io::ErrorKind;
use std::pin::Pin;
use std::sync::Arc;

pub(crate) struct Subscriber {
    pub(crate) subscriptions: Vec<Vec<u8>>,
    pub(crate) send_queue: Pin<Box<ZmqFramedWrite>>,
    _subscription_coro_stop: oneshot::Sender<()>,
}

pub(crate) struct PubSocketBackend {
    subscribers: DashMap<PeerIdentity, Subscriber>,
    socket_monitor: Mutex<Option<mpsc::Sender<SocketEvent>>>,
    socket_options: SocketOptions,
}

impl PubSocketBackend {
    fn message_received(&self, peer_id: &PeerIdentity, message: Message) {
        let message = match message {
            Message::Message(m) => m,
            _ => return,
        };
        assert_eq!(message.len(), 1);
        let data: Vec<u8> = message.into_vec().pop().unwrap().to_vec();
        if data.is_empty() {
            return;
        }
        match data[0] {
            1 => {
                // Subscribe
                self.subscribers
                    .get_mut(peer_id)
                    .unwrap()
                    .subscriptions
                    .push(Vec::from(&data[1..]));
            }
            0 => {
                // Unsubscribe
                let mut del_index = None;
                let sub = Vec::from(&data[1..]);
                for (idx, subscription) in self
                    .subscribers
                    .get(peer_id)
                    .unwrap()
                    .subscriptions
                    .iter()
                    .enumerate()
                {
                    if &sub == subscription {
                        del_index = Some(idx);
                        break;
                    }
                }
                if let Some(index) = del_index {
                    self.subscribers
                        .get_mut(peer_id)
                        .unwrap()
                        .subscriptions
                        .remove(index);
                }
            }
            _ => (),
        }
    }
}

impl SocketBackend for PubSocketBackend {
    fn socket_type(&self) -> SocketType {
        SocketType::PUB
    }

    fn socket_options(&self) -> &SocketOptions {
        &self.socket_options
    }

    fn shutdown(&self) {
        self.subscribers.clear();
    }

    fn monitor(&self) -> &Mutex<Option<mpsc::Sender<SocketEvent>>> {
        &self.socket_monitor
    }
}

impl MultiPeerBackend for PubSocketBackend {
    fn peer_connected(self: Arc<Self>, peer_id: &PeerIdentity, io: FramedIo) {
        let (mut recv_queue, send_queue) = io.into_parts();
        // TODO provide handling for recv_queue
        let (sender, stop_receiver) = oneshot::channel();
        self.subscribers.insert(
            peer_id.clone(),
            Subscriber {
                subscriptions: vec![],
                send_queue: Box::pin(send_queue),
                _subscription_coro_stop: sender,
            },
        );
        let backend = self;
        let peer_id = peer_id.clone();
        async_rt::task::spawn(async move {
            use futures::StreamExt;
            let mut stop_receiver = stop_receiver.fuse();
            loop {
                futures::select! {
                     _ = stop_receiver => {
                         break;
                     },
                     message = recv_queue.next().fuse() => {
                        match message {
                            Some(Ok(m)) => backend.message_received(&peer_id, m),
                            Some(Err(e)) => {
                                dbg!(e);
                                backend.peer_disconnected(&peer_id);
                                break;
                            }
                            None => {
                                backend.peer_disconnected(&peer_id);
                                break
                            }
                        }

                     }
                }
            }
        });
    }

    fn peer_disconnected(&self, peer_id: &PeerIdentity) {
        log::info!("Client disconnected {:?}", peer_id);
        self.subscribers.remove(peer_id);
    }
}

pub struct PubSocket {
    pub(crate) backend: Arc<PubSocketBackend>,
    binds: HashMap<Endpoint, AcceptStopHandle>,
}

impl Drop for PubSocket {
    fn drop(&mut self) {
        self.backend.shutdown();
    }
}

#[async_trait]
impl SocketSend for PubSocket {
    async fn send(&mut self, message: ZmqMessage) -> ZmqResult<()> {
        let mut dead_peers = Vec::new();
        for mut subscriber in self.backend.subscribers.iter_mut() {
            for sub_filter in &subscriber.subscriptions {
                if sub_filter.len() <= message.get(0).unwrap().len()
                    && sub_filter.as_slice() == &message.get(0).unwrap()[0..sub_filter.len()]
                {
                    let res = subscriber
                        .send_queue
                        .as_mut()
                        .try_send(Message::Message(message.clone()));
                    match res {
                        Ok(()) => {}
                        Err(ZmqError::Codec(CodecError::Io(e))) => {
                            if e.kind() == ErrorKind::BrokenPipe {
                                dead_peers.push(subscriber.key().clone());
                            } else {
                                dbg!(e);
                            }
                        }
                        Err(e) => {
                            dbg!(e);
                            todo!()
                        }
                    }
                    break;
                }
            }
        }
        for peer in dead_peers {
            self.backend.peer_disconnected(&peer);
        }
        Ok(())
    }
}

impl CaptureSocket for PubSocket {}

#[async_trait]
impl Socket for PubSocket {
    fn with_options(options: SocketOptions) -> Self {
        Self {
            backend: Arc::new(PubSocketBackend {
                subscribers: DashMap::new(),
                socket_monitor: Mutex::new(None),
                socket_options: options,
            }),
            binds: HashMap::new(),
        }
    }

    fn backend(&self) -> Arc<dyn MultiPeerBackend> {
        self.backend.clone()
    }

    fn binds(&mut self) -> &mut HashMap<Endpoint, AcceptStopHandle> {
        &mut self.binds
    }

    fn monitor(&mut self) -> mpsc::Receiver<SocketEvent> {
        let (sender, receiver) = mpsc::channel(1024);
        self.backend.socket_monitor.lock().replace(sender);
        receiver
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::util::tests::{
        test_bind_to_any_port_helper, test_bind_to_unspecified_interface_helper,
    };
    use crate::ZmqResult;
    use std::net::IpAddr;

    #[async_rt::test]
    async fn test_bind_to_any_port() -> ZmqResult<()> {
        let s = PubSocket::new();
        test_bind_to_any_port_helper(s).await
    }

    #[async_rt::test]
    async fn test_bind_to_any_ipv4_interface() -> ZmqResult<()> {
        let any_ipv4: IpAddr = "0.0.0.0".parse().unwrap();
        let s = PubSocket::new();
        test_bind_to_unspecified_interface_helper(any_ipv4, s, 4000).await
    }

    #[async_rt::test]
    async fn test_bind_to_any_ipv6_interface() -> ZmqResult<()> {
        let any_ipv6: IpAddr = "::".parse().unwrap();
        let s = PubSocket::new();
        test_bind_to_unspecified_interface_helper(any_ipv6, s, 4010).await
    }
}
<|endoftext|>


<commit_before>// Copyright 2014-2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use core::ptr::PtrExt;

// FIXME: #13996: mark the `allocate` and `reallocate` return value as `noalias`

/// Return a pointer to `size` bytes of memory aligned to `align`.
///
/// On failure, return a null pointer.
///
/// Behavior is undefined if the requested size is 0 or the alignment is not a
/// power of 2. The alignment must be no larger than the largest supported page
/// size on the platform.
#[inline]
pub unsafe fn allocate(size: uint, align: uint) -> *mut u8 {
    imp::allocate(size, align)
}

/// Resize the allocation referenced by `ptr` to `size` bytes.
///
/// On failure, return a null pointer and leave the original allocation intact.
///
/// Behavior is undefined if the requested size is 0 or the alignment is not a
/// power of 2. The alignment must be no larger than the largest supported page
/// size on the platform.
///
/// The `old_size` and `align` parameters are the parameters that were used to
/// create the allocation referenced by `ptr`. The `old_size` parameter may be
/// any value in range_inclusive(requested_size, usable_size).
#[inline]
pub unsafe fn reallocate(ptr: *mut u8, old_size: uint, size: uint, align: uint) -> *mut u8 {
    imp::reallocate(ptr, old_size, size, align)
}

/// Resize the allocation referenced by `ptr` to `size` bytes.
///
/// If the operation succeeds, it returns `usable_size(size, align)` and if it
/// fails (or is a no-op) it returns `usable_size(old_size, align)`.
///
/// Behavior is undefined if the requested size is 0 or the alignment is not a
/// power of 2. The alignment must be no larger than the largest supported page
/// size on the platform.
///
/// The `old_size` and `align` parameters are the parameters that were used to
/// create the allocation referenced by `ptr`. The `old_size` parameter may be
/// any value in range_inclusive(requested_size, usable_size).
#[inline]
pub unsafe fn reallocate_inplace(ptr: *mut u8, old_size: uint, size: uint, align: uint) -> uint {
    imp::reallocate_inplace(ptr, old_size, size, align)
}

/// Deallocates the memory referenced by `ptr`.
///
/// The `ptr` parameter must not be null.
///
/// The `old_size` and `align` parameters are the parameters that were used to
/// create the allocation referenced by `ptr`. The `old_size` parameter may be
/// any value in range_inclusive(requested_size, usable_size).
#[inline]
pub unsafe fn deallocate(ptr: *mut u8, old_size: uint, align: uint) {
    imp::deallocate(ptr, old_size, align)
}

/// Returns the usable size of an allocation created with the specified the
/// `size` and `align`.
#[inline]
pub fn usable_size(size: uint, align: uint) -> uint {
    imp::usable_size(size, align)
}

/// Prints implementation-defined allocator statistics.
///
/// These statistics may be inconsistent if other threads use the allocator
/// during the call.
#[unstable]
pub fn stats_print() {
    imp::stats_print();
}

/// An arbitrary non-null address to represent zero-size allocations.
///
/// This preserves the non-null invariant for types like `Box<T>`. The address may overlap with
/// non-zero-size memory allocations.
pub const EMPTY: *mut () = 0x1 as *mut ();

/// The allocator for unique pointers.
#[cfg(not(test))]
#[lang="exchange_malloc"]
#[inline]
unsafe fn exchange_malloc(size: uint, align: uint) -> *mut u8 {
    if size == 0 {
        EMPTY as *mut u8
    } else {
        let ptr = allocate(size, align);
        if ptr.is_null() { ::oom() }
        ptr
    }
}

#[cfg(not(test))]
#[lang="exchange_free"]
#[inline]
unsafe fn exchange_free(ptr: *mut u8, old_size: uint, align: uint) {
    deallocate(ptr, old_size, align);
}

// The minimum alignment guaranteed by the architecture. This value is used to
// add fast paths for low alignment values. In practice, the alignment is a
// constant at the call site and the branch will be optimized out.
#[cfg(all(not(feature = "external_funcs"),
          not(feature = "external_crate"),
          any(target_arch = "arm",
              target_arch = "mips",
              target_arch = "mipsel")))]
const MIN_ALIGN: uint = 8;
#[cfg(all(not(feature = "external_funcs"),
          not(feature = "external_crate"),
          any(target_arch = "x86",
              target_arch = "x86_64",
              target_arch = "aarch64")))]
const MIN_ALIGN: uint = 16;

#[cfg(feature = "external_funcs")]
mod imp {
    extern {
        fn rust_allocate(size: uint, align: uint) -> *mut u8;
        fn rust_deallocate(ptr: *mut u8, old_size: uint, align: uint);
        fn rust_reallocate(ptr: *mut u8, old_size: uint, size: uint, align: uint) -> *mut u8;
        fn rust_reallocate_inplace(ptr: *mut u8, old_size: uint, size: uint,
                                   align: uint) -> uint;
        fn rust_usable_size(size: uint, align: uint) -> uint;
        fn rust_stats_print();
    }

    #[inline]
    pub unsafe fn allocate(size: uint, align: uint) -> *mut u8 {
        rust_allocate(size, align)
    }

    #[inline]
    pub unsafe fn deallocate(ptr: *mut u8, old_size: uint, align: uint) {
        rust_deallocate(ptr, old_size, align)
    }

    #[inline]
    pub unsafe fn reallocate(ptr: *mut u8, old_size: uint, size: uint, align: uint) -> *mut u8 {
        rust_reallocate(ptr, old_size, size, align)
    }

    #[inline]
    pub unsafe fn reallocate_inplace(ptr: *mut u8, old_size: uint, size: uint,
                                     align: uint) -> uint {
        rust_reallocate_inplace(ptr, old_size, size, align)
    }

    #[inline]
    pub fn usable_size(size: uint, align: uint) -> uint {
        unsafe { rust_usable_size(size, align) }
    }

    #[inline]
    pub fn stats_print() {
        unsafe { rust_stats_print() }
    }
}

#[cfg(feature = "external_crate")]
mod imp {
    extern crate external;
    pub use self::external::{allocate, deallocate, reallocate_inplace, reallocate};
    pub use self::external::{usable_size, stats_print};
}

#[cfg(all(not(feature = "external_funcs"),
          not(feature = "external_crate"),
          jemalloc))]
mod imp {
    use core::option::Option;
    use core::option::Option::None;
    use core::ptr::{null_mut, null};
    use core::num::Int;
    use libc::{c_char, c_int, c_void, size_t};
    use super::MIN_ALIGN;

    #[link(name = "jemalloc", kind = "static")]
    #[cfg(not(test))]
    extern {}

    extern {
        fn je_mallocx(size: size_t, flags: c_int) -> *mut c_void;
        fn je_rallocx(ptr: *mut c_void, size: size_t, flags: c_int) -> *mut c_void;
        fn je_xallocx(ptr: *mut c_void, size: size_t, extra: size_t, flags: c_int) -> size_t;
        fn je_sdallocx(ptr: *mut c_void, size: size_t, flags: c_int);
        fn je_nallocx(size: size_t, flags: c_int) -> size_t;
        fn je_malloc_stats_print(write_cb: Option<extern "C" fn(cbopaque: *mut c_void,
                                                                *const c_char)>,
                                 cbopaque: *mut c_void,
                                 opts: *const c_char);
    }

    // -lpthread needs to occur after -ljemalloc, the earlier argument isn't enough
    #[cfg(all(not(windows), not(target_os = "android")))]
    #[link(name = "pthread")]
    extern {}

    // MALLOCX_ALIGN(a) macro
    #[inline(always)]
    fn mallocx_align(a: uint) -> c_int { a.trailing_zeros() as c_int }

    #[inline(always)]
    fn align_to_flags(align: uint) -> c_int {
        if align <= MIN_ALIGN { 0 } else { mallocx_align(align) }
    }

    #[inline]
    pub unsafe fn allocate(size: uint, align: uint) -> *mut u8 {
        let flags = align_to_flags(align);
        je_mallocx(size as size_t, flags) as *mut u8
    }

    #[inline]
    pub unsafe fn reallocate(ptr: *mut u8, _old_size: uint, size: uint, align: uint) -> *mut u8 {
        let flags = align_to_flags(align);
        je_rallocx(ptr as *mut c_void, size as size_t, flags) as *mut u8
    }

    #[inline]
    pub unsafe fn reallocate_inplace(ptr: *mut u8, _old_size: uint, size: uint,
                                     align: uint) -> uint {
        let flags = align_to_flags(align);
        je_xallocx(ptr as *mut c_void, size as size_t, 0, flags) as uint
    }

    #[inline]
    pub unsafe fn deallocate(ptr: *mut u8, old_size: uint, align: uint) {
        let flags = align_to_flags(align);
        je_sdallocx(ptr as *mut c_void, old_size as size_t, flags)
    }

    #[inline]
    pub fn usable_size(size: uint, align: uint) -> uint {
        let flags = align_to_flags(align);
        unsafe { je_nallocx(size as size_t, flags) as uint }
    }

    pub fn stats_print() {
        unsafe {
            je_malloc_stats_print(None, null_mut(), null())
        }
    }
}

#[cfg(all(not(feature = "external_funcs"),
          not(feature = "external_crate"),
          not(jemalloc),
          unix))]
mod imp {
    use core::cmp;
    use core::ptr;
    use libc;
    use super::MIN_ALIGN;

    extern {
        fn posix_memalign(memptr: *mut *mut libc::c_void,
                          align: libc::size_t,
                          size: libc::size_t) -> libc::c_int;
    }

    #[inline]
    pub unsafe fn allocate(size: uint, align: uint) -> *mut u8 {
        if align <= MIN_ALIGN {
            libc::malloc(size as libc::size_t) as *mut u8
        } else {
            let mut out = 0 as *mut libc::c_void;
            let ret = posix_memalign(&mut out,
                                     align as libc::size_t,
                                     size as libc::size_t);
            if ret != 0 {
                ptr::null_mut()
            } else {
                out as *mut u8
            }
        }
    }

    #[inline]
    pub unsafe fn reallocate(ptr: *mut u8, old_size: uint, size: uint, align: uint) -> *mut u8 {
        if align <= MIN_ALIGN {
            libc::realloc(ptr as *mut libc::c_void, size as libc::size_t) as *mut u8
        } else {
            let new_ptr = allocate(size, align);
            ptr::copy_memory(new_ptr, ptr as *const u8, cmp::min(size, old_size));
            deallocate(ptr, old_size, align);
            new_ptr
        }
    }

    #[inline]
    pub unsafe fn reallocate_inplace(_ptr: *mut u8, old_size: uint, _size: uint,
                                     _align: uint) -> uint {
        old_size
    }

    #[inline]
    pub unsafe fn deallocate(ptr: *mut u8, _old_size: uint, _align: uint) {
        libc::free(ptr as *mut libc::c_void)
    }

    #[inline]
    pub fn usable_size(size: uint, _align: uint) -> uint {
        size
    }

    pub fn stats_print() {}
}

#[cfg(all(not(feature = "external_funcs"),
          not(feature = "external_crate"),
          not(jemalloc),
          windows))]
mod imp {
    use libc::{c_void, size_t};
    use libc;
    use super::MIN_ALIGN;

    extern {
        fn _aligned_malloc(size: size_t, align: size_t) -> *mut c_void;
        fn _aligned_realloc(block: *mut c_void, size: size_t,
                            align: size_t) -> *mut c_void;
        fn _aligned_free(ptr: *mut c_void);
    }

    #[inline]
    pub unsafe fn allocate(size: uint, align: uint) -> *mut u8 {
        if align <= MIN_ALIGN {
            libc::malloc(size as size_t) as *mut u8
        } else {
            _aligned_malloc(size as size_t, align as size_t) as *mut u8
        }
    }

    #[inline]
    pub unsafe fn reallocate(ptr: *mut u8, _old_size: uint, size: uint, align: uint) -> *mut u8 {
        if align <= MIN_ALIGN {
            libc::realloc(ptr as *mut c_void, size as size_t) as *mut u8
        } else {
            _aligned_realloc(ptr as *mut c_void, size as size_t, align as size_t) as *mut u8
        }
    }

    #[inline]
    pub unsafe fn reallocate_inplace(_ptr: *mut u8, old_size: uint, _size: uint,
                                     _align: uint) -> uint {
        old_size
    }

    #[inline]
    pub unsafe fn deallocate(ptr: *mut u8, _old_size: uint, align: uint) {
        if align <= MIN_ALIGN {
            libc::free(ptr as *mut libc::c_void)
        } else {
            _aligned_free(ptr as *mut c_void)
        }
    }

    #[inline]
    pub fn usable_size(size: uint, _align: uint) -> uint {
        size
    }

    pub fn stats_print() {}
}

#[cfg(test)]
mod test {
    extern crate test;
    use self::test::Bencher;
    use core::ptr::PtrExt;
    use heap;

    #[test]
    fn basic_reallocate_inplace_noop() {
        unsafe {
            let size = 4000;
            let ptr = heap::allocate(size, 8);
            if ptr.is_null() { ::oom() }
            let ret = heap::reallocate_inplace(ptr, size, size, 8);
            heap::deallocate(ptr, size, 8);
            assert_eq!(ret, heap::usable_size(size, 8));
        }
    }

    #[bench]
    fn alloc_owned_small(b: &mut Bencher) {
        b.iter(|| {
            box 10i
        })
    }
}
<commit_msg>powerpc: Tell liballoc about power alignment<commit_after>// Copyright 2014-2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use core::ptr::PtrExt;

// FIXME: #13996: mark the `allocate` and `reallocate` return value as `noalias`

/// Return a pointer to `size` bytes of memory aligned to `align`.
///
/// On failure, return a null pointer.
///
/// Behavior is undefined if the requested size is 0 or the alignment is not a
/// power of 2. The alignment must be no larger than the largest supported page
/// size on the platform.
#[inline]
pub unsafe fn allocate(size: uint, align: uint) -> *mut u8 {
    imp::allocate(size, align)
}

/// Resize the allocation referenced by `ptr` to `size` bytes.
///
/// On failure, return a null pointer and leave the original allocation intact.
///
/// Behavior is undefined if the requested size is 0 or the alignment is not a
/// power of 2. The alignment must be no larger than the largest supported page
/// size on the platform.
///
/// The `old_size` and `align` parameters are the parameters that were used to
/// create the allocation referenced by `ptr`. The `old_size` parameter may be
/// any value in range_inclusive(requested_size, usable_size).
#[inline]
pub unsafe fn reallocate(ptr: *mut u8, old_size: uint, size: uint, align: uint) -> *mut u8 {
    imp::reallocate(ptr, old_size, size, align)
}

/// Resize the allocation referenced by `ptr` to `size` bytes.
///
/// If the operation succeeds, it returns `usable_size(size, align)` and if it
/// fails (or is a no-op) it returns `usable_size(old_size, align)`.
///
/// Behavior is undefined if the requested size is 0 or the alignment is not a
/// power of 2. The alignment must be no larger than the largest supported page
/// size on the platform.
///
/// The `old_size` and `align` parameters are the parameters that were used to
/// create the allocation referenced by `ptr`. The `old_size` parameter may be
/// any value in range_inclusive(requested_size, usable_size).
#[inline]
pub unsafe fn reallocate_inplace(ptr: *mut u8, old_size: uint, size: uint, align: uint) -> uint {
    imp::reallocate_inplace(ptr, old_size, size, align)
}

/// Deallocates the memory referenced by `ptr`.
///
/// The `ptr` parameter must not be null.
///
/// The `old_size` and `align` parameters are the parameters that were used to
/// create the allocation referenced by `ptr`. The `old_size` parameter may be
/// any value in range_inclusive(requested_size, usable_size).
#[inline]
pub unsafe fn deallocate(ptr: *mut u8, old_size: uint, align: uint) {
    imp::deallocate(ptr, old_size, align)
}

/// Returns the usable size of an allocation created with the specified the
/// `size` and `align`.
#[inline]
pub fn usable_size(size: uint, align: uint) -> uint {
    imp::usable_size(size, align)
}

/// Prints implementation-defined allocator statistics.
///
/// These statistics may be inconsistent if other threads use the allocator
/// during the call.
#[unstable]
pub fn stats_print() {
    imp::stats_print();
}

/// An arbitrary non-null address to represent zero-size allocations.
///
/// This preserves the non-null invariant for types like `Box<T>`. The address may overlap with
/// non-zero-size memory allocations.
pub const EMPTY: *mut () = 0x1 as *mut ();

/// The allocator for unique pointers.
#[cfg(not(test))]
#[lang="exchange_malloc"]
#[inline]
unsafe fn exchange_malloc(size: uint, align: uint) -> *mut u8 {
    if size == 0 {
        EMPTY as *mut u8
    } else {
        let ptr = allocate(size, align);
        if ptr.is_null() { ::oom() }
        ptr
    }
}

#[cfg(not(test))]
#[lang="exchange_free"]
#[inline]
unsafe fn exchange_free(ptr: *mut u8, old_size: uint, align: uint) {
    deallocate(ptr, old_size, align);
}

// The minimum alignment guaranteed by the architecture. This value is used to
// add fast paths for low alignment values. In practice, the alignment is a
// constant at the call site and the branch will be optimized out.
#[cfg(all(not(feature = "external_funcs"),
          not(feature = "external_crate"),
          any(target_arch = "arm",
              target_arch = "mips",
              target_arch = "mipsel",
              target_arch = "powerpc")))]
const MIN_ALIGN: uint = 8;
#[cfg(all(not(feature = "external_funcs"),
          not(feature = "external_crate"),
          any(target_arch = "x86",
              target_arch = "x86_64",
              target_arch = "aarch64")))]
const MIN_ALIGN: uint = 16;

#[cfg(feature = "external_funcs")]
mod imp {
    extern {
        fn rust_allocate(size: uint, align: uint) -> *mut u8;
        fn rust_deallocate(ptr: *mut u8, old_size: uint, align: uint);
        fn rust_reallocate(ptr: *mut u8, old_size: uint, size: uint, align: uint) -> *mut u8;
        fn rust_reallocate_inplace(ptr: *mut u8, old_size: uint, size: uint,
                                   align: uint) -> uint;
        fn rust_usable_size(size: uint, align: uint) -> uint;
        fn rust_stats_print();
    }

    #[inline]
    pub unsafe fn allocate(size: uint, align: uint) -> *mut u8 {
        rust_allocate(size, align)
    }

    #[inline]
    pub unsafe fn deallocate(ptr: *mut u8, old_size: uint, align: uint) {
        rust_deallocate(ptr, old_size, align)
    }

    #[inline]
    pub unsafe fn reallocate(ptr: *mut u8, old_size: uint, size: uint, align: uint) -> *mut u8 {
        rust_reallocate(ptr, old_size, size, align)
    }

    #[inline]
    pub unsafe fn reallocate_inplace(ptr: *mut u8, old_size: uint, size: uint,
                                     align: uint) -> uint {
        rust_reallocate_inplace(ptr, old_size, size, align)
    }

    #[inline]
    pub fn usable_size(size: uint, align: uint) -> uint {
        unsafe { rust_usable_size(size, align) }
    }

    #[inline]
    pub fn stats_print() {
        unsafe { rust_stats_print() }
    }
}

#[cfg(feature = "external_crate")]
mod imp {
    extern crate external;
    pub use self::external::{allocate, deallocate, reallocate_inplace, reallocate};
    pub use self::external::{usable_size, stats_print};
}

#[cfg(all(not(feature = "external_funcs"),
          not(feature = "external_crate"),
          jemalloc))]
mod imp {
    use core::option::Option;
    use core::option::Option::None;
    use core::ptr::{null_mut, null};
    use core::num::Int;
    use libc::{c_char, c_int, c_void, size_t};
    use super::MIN_ALIGN;

    #[link(name = "jemalloc", kind = "static")]
    #[cfg(not(test))]
    extern {}

    extern {
        fn je_mallocx(size: size_t, flags: c_int) -> *mut c_void;
        fn je_rallocx(ptr: *mut c_void, size: size_t, flags: c_int) -> *mut c_void;
        fn je_xallocx(ptr: *mut c_void, size: size_t, extra: size_t, flags: c_int) -> size_t;
        fn je_sdallocx(ptr: *mut c_void, size: size_t, flags: c_int);
        fn je_nallocx(size: size_t, flags: c_int) -> size_t;
        fn je_malloc_stats_print(write_cb: Option<extern "C" fn(cbopaque: *mut c_void,
                                                                *const c_char)>,
                                 cbopaque: *mut c_void,
                                 opts: *const c_char);
    }

    // -lpthread needs to occur after -ljemalloc, the earlier argument isn't enough
    #[cfg(all(not(windows), not(target_os = "android")))]
    #[link(name = "pthread")]
    extern {}

    // MALLOCX_ALIGN(a) macro
    #[inline(always)]
    fn mallocx_align(a: uint) -> c_int { a.trailing_zeros() as c_int }

    #[inline(always)]
    fn align_to_flags(align: uint) -> c_int {
        if align <= MIN_ALIGN { 0 } else { mallocx_align(align) }
    }

    #[inline]
    pub unsafe fn allocate(size: uint, align: uint) -> *mut u8 {
        let flags = align_to_flags(align);
        je_mallocx(size as size_t, flags) as *mut u8
    }

    #[inline]
    pub unsafe fn reallocate(ptr: *mut u8, _old_size: uint, size: uint, align: uint) -> *mut u8 {
        let flags = align_to_flags(align);
        je_rallocx(ptr as *mut c_void, size as size_t, flags) as *mut u8
    }

    #[inline]
    pub unsafe fn reallocate_inplace(ptr: *mut u8, _old_size: uint, size: uint,
                                     align: uint) -> uint {
        let flags = align_to_flags(align);
        je_xallocx(ptr as *mut c_void, size as size_t, 0, flags) as uint
    }

    #[inline]
    pub unsafe fn deallocate(ptr: *mut u8, old_size: uint, align: uint) {
        let flags = align_to_flags(align);
        je_sdallocx(ptr as *mut c_void, old_size as size_t, flags)
    }

    #[inline]
    pub fn usable_size(size: uint, align: uint) -> uint {
        let flags = align_to_flags(align);
        unsafe { je_nallocx(size as size_t, flags) as uint }
    }

    pub fn stats_print() {
        unsafe {
            je_malloc_stats_print(None, null_mut(), null())
        }
    }
}

#[cfg(all(not(feature = "external_funcs"),
          not(feature = "external_crate"),
          not(jemalloc),
          unix))]
mod imp {
    use core::cmp;
    use core::ptr;
    use libc;
    use super::MIN_ALIGN;

    extern {
        fn posix_memalign(memptr: *mut *mut libc::c_void,
                          align: libc::size_t,
                          size: libc::size_t) -> libc::c_int;
    }

    #[inline]
    pub unsafe fn allocate(size: uint, align: uint) -> *mut u8 {
        if align <= MIN_ALIGN {
            libc::malloc(size as libc::size_t) as *mut u8
        } else {
            let mut out = 0 as *mut libc::c_void;
            let ret = posix_memalign(&mut out,
                                     align as libc::size_t,
                                     size as libc::size_t);
            if ret != 0 {
                ptr::null_mut()
            } else {
                out as *mut u8
            }
        }
    }

    #[inline]
    pub unsafe fn reallocate(ptr: *mut u8, old_size: uint, size: uint, align: uint) -> *mut u8 {
        if align <= MIN_ALIGN {
            libc::realloc(ptr as *mut libc::c_void, size as libc::size_t) as *mut u8
        } else {
            let new_ptr = allocate(size, align);
            ptr::copy_memory(new_ptr, ptr as *const u8, cmp::min(size, old_size));
            deallocate(ptr, old_size, align);
            new_ptr
        }
    }

    #[inline]
    pub unsafe fn reallocate_inplace(_ptr: *mut u8, old_size: uint, _size: uint,
                                     _align: uint) -> uint {
        old_size
    }

    #[inline]
    pub unsafe fn deallocate(ptr: *mut u8, _old_size: uint, _align: uint) {
        libc::free(ptr as *mut libc::c_void)
    }

    #[inline]
    pub fn usable_size(size: uint, _align: uint) -> uint {
        size
    }

    pub fn stats_print() {}
}

#[cfg(all(not(feature = "external_funcs"),
          not(feature = "external_crate"),
          not(jemalloc),
          windows))]
mod imp {
    use libc::{c_void, size_t};
    use libc;
    use super::MIN_ALIGN;

    extern {
        fn _aligned_malloc(size: size_t, align: size_t) -> *mut c_void;
        fn _aligned_realloc(block: *mut c_void, size: size_t,
                            align: size_t) -> *mut c_void;
        fn _aligned_free(ptr: *mut c_void);
    }

    #[inline]
    pub unsafe fn allocate(size: uint, align: uint) -> *mut u8 {
        if align <= MIN_ALIGN {
            libc::malloc(size as size_t) as *mut u8
        } else {
            _aligned_malloc(size as size_t, align as size_t) as *mut u8
        }
    }

    #[inline]
    pub unsafe fn reallocate(ptr: *mut u8, _old_size: uint, size: uint, align: uint) -> *mut u8 {
        if align <= MIN_ALIGN {
            libc::realloc(ptr as *mut c_void, size as size_t) as *mut u8
        } else {
            _aligned_realloc(ptr as *mut c_void, size as size_t, align as size_t) as *mut u8
        }
    }

    #[inline]
    pub unsafe fn reallocate_inplace(_ptr: *mut u8, old_size: uint, _size: uint,
                                     _align: uint) -> uint {
        old_size
    }

    #[inline]
    pub unsafe fn deallocate(ptr: *mut u8, _old_size: uint, align: uint) {
        if align <= MIN_ALIGN {
            libc::free(ptr as *mut libc::c_void)
        } else {
            _aligned_free(ptr as *mut c_void)
        }
    }

    #[inline]
    pub fn usable_size(size: uint, _align: uint) -> uint {
        size
    }

    pub fn stats_print() {}
}

#[cfg(test)]
mod test {
    extern crate test;
    use self::test::Bencher;
    use core::ptr::PtrExt;
    use heap;

    #[test]
    fn basic_reallocate_inplace_noop() {
        unsafe {
            let size = 4000;
            let ptr = heap::allocate(size, 8);
            if ptr.is_null() { ::oom() }
            let ret = heap::reallocate_inplace(ptr, size, size, 8);
            heap::deallocate(ptr, size, 8);
            assert_eq!(ret, heap::usable_size(size, 8));
        }
    }

    #[bench]
    fn alloc_owned_small(b: &mut Bencher) {
        b.iter(|| {
            box 10i
        })
    }
}
<|endoftext|>


<commit_before>//! Parse vbaProject.bin file
//!
//! Retranscription from:
//! https://github.com/unixfreak0037/officeparser/blob/master/officeparser.py

use std::collections::HashMap;
use std::io::Read;
use std::path::PathBuf;

use byteorder::{LittleEndian, ReadBytesExt};
use log::{debug, log_enabled, warn, Level};

use crate::cfb::{Cfb, XlsEncoding};
use crate::utils::read_u16;

/// A VBA specific error enum
#[derive(Debug)]
pub enum VbaError {
    /// Error comes from a cfb parsing
    Cfb(crate::cfb::CfbError),
    /// Io error
    Io(std::io::Error),

    /// Cannot find module
    ModuleNotFound(String),
    /// Generic unknown u16 value
    Unknown {
        /// error type
        typ: &'static str,
        /// value found
        val: u16,
    },
    /// Invalid libid format
    LibId,
    /// Invalid record id
    InvalidRecordId {
        /// expected record id
        expected: u16,
        /// record if found
        found: u16,
    },
}

from_err!(crate::cfb::CfbError, VbaError, Cfb);
from_err!(std::io::Error, VbaError, Io);

impl std::fmt::Display for VbaError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            VbaError::Io(e) => write!(f, "I/O error: {}", e),
            VbaError::Cfb(e) => write!(f, "Cfb error: {}", e),

            VbaError::ModuleNotFound(e) => write!(f, "Cannot find module '{}'", e),
            VbaError::Unknown { typ, val } => write!(f, "Unknown {} '{:X}'", typ, val),
            VbaError::LibId => write!(f, "Unexpected libid format"),
            VbaError::InvalidRecordId { expected, found } => write!(
                f,
                "Invalid record id: expecting {:X} found {:X}",
                expected, found
            ),
        }
    }
}

impl std::error::Error for VbaError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            VbaError::Io(e) => Some(e),
            VbaError::Cfb(e) => Some(e),
            _ => None,
        }
    }
}

/// A struct for managing VBA reading
#[allow(dead_code)]
#[derive(Clone)]
pub struct VbaProject {
    references: Vec<Reference>,
    modules: HashMap<String, Vec<u8>>,
    encoding: XlsEncoding,
}

impl VbaProject {
    /// Create a new `VbaProject` out of the vbaProject.bin `ZipFile` or xls file
    ///
    /// Starts reading project metadata (header, directories, sectors and minisectors).
    pub fn new<R: Read>(r: &mut R, len: usize) -> Result<VbaProject, VbaError> {
        let mut cfb = Cfb::new(r, len)?;
        VbaProject::from_cfb(r, &mut cfb)
    }

    /// Creates a new `VbaProject` out of a Compound File Binary and the corresponding reader
    pub fn from_cfb<R: Read>(r: &mut R, cfb: &mut Cfb) -> Result<VbaProject, VbaError> {
        // dir stream
        let stream = cfb.get_stream("dir", r)?;
        let stream = crate::cfb::decompress_stream(&*stream)?;
        let stream = &mut &*stream;

        // read dir information record (not used)
        let encoding = read_dir_information(stream)?;

        // array of REFERENCE records
        let refs = Reference::from_stream(stream, &encoding)?;

        // modules
        let mods: Vec<Module> = read_modules(stream, &encoding)?;

        // read all modules
        let modules: HashMap<String, Vec<u8>> = mods
            .into_iter()
            .map(|m| {
                cfb.get_stream(&m.stream_name, r).and_then(|s| {
                    crate::cfb::decompress_stream(&s[m.text_offset..]).map(move |s| (m.name, s))
                })
            })
            .collect::<Result<HashMap<_, _>, _>>()?;

        Ok(VbaProject {
            references: refs,
            modules,
            encoding,
        })
    }

    /// Gets the list of `Reference`s
    pub fn get_references(&self) -> &[Reference] {
        &self.references
    }

    /// Gets the list of `Module` names
    pub fn get_module_names(&self) -> Vec<&str> {
        self.modules.keys().map(|k| &**k).collect()
    }

    /// Reads module content and tries to convert to utf8
    ///
    /// While it works most of the time, the modules are MBCS encoding and the conversion
    /// may fail. If this is the case you should revert to `read_module_raw` as there is
    /// no built in decoding provided in this crate
    ///
    /// # Examples
    /// ```
    /// use calamine::{Reader, open_workbook, Xlsx};
    ///
    /// # let path = format!("{}/tests/vba.xlsm", env!("CARGO_MANIFEST_DIR"));
    /// let mut xl: Xlsx<_> = open_workbook(path).expect("Cannot find excel file");
    /// if let Some(Ok(mut vba)) = xl.vba_project() {
    ///     let vba = vba.to_mut();
    ///     let modules = vba.get_module_names().into_iter()
    ///                      .map(|s| s.to_string()).collect::<Vec<_>>();
    ///     for m in modules {
    ///         println!("Module {}:", m);
    ///         println!("{}", vba.get_module(&m)
    ///                           .expect(&format!("cannot read {:?} module", m)));
    ///     }
    /// }
    /// ```
    pub fn get_module(&self, name: &str) -> Result<String, VbaError> {
        debug!("read module {}", name);
        let data = self.get_module_raw(name)?;
        Ok(self.encoding.decode_all(data))
    }

    /// Reads module content (MBCS encoded) and output it as-is (binary output)
    pub fn get_module_raw(&self, name: &str) -> Result<&[u8], VbaError> {
        match self.modules.get(name) {
            Some(m) => Ok(&**m),
            None => Err(VbaError::ModuleNotFound(name.into())),
        }
    }
}

/// A vba reference
#[derive(Debug, Clone, Hash, Eq, PartialEq)]
pub struct Reference {
    /// name
    pub name: String,
    /// description
    pub description: String,
    /// location of the reference
    pub path: PathBuf,
}

impl Reference {
    /// Check if the reference location is accessible
    pub fn is_missing(&self) -> bool {
        !self.path.exists()
    }

    /// Gets the list of references from the dir_stream relevant part
    fn from_stream(stream: &mut &[u8], encoding: &XlsEncoding) -> Result<Vec<Reference>, VbaError> {
        debug!("read all references metadata");

        let mut references = Vec::new();
        let mut reference = Reference {
            name: "".to_string(),
            description: "".to_string(),
            path: "".into(),
        };

        loop {
            let check = stream.read_u16::<LittleEndian>();
            match check? {
                0x000F => {
                    // termination of references array
                    if !reference.name.is_empty() {
                        references.push(reference);
                    }
                    break;
                }
                0x0016 => {
                    // REFERENCENAME
                    if !reference.name.is_empty() {
                        references.push(reference);
                    }
                    let name = read_variable_record(stream, 1)?;
                    let name = encoding.decode_all(name);
                    reference = Reference {
                        name: name.clone(),
                        description: name,
                        path: "".into(),
                    };
                    check_variable_record(0x003E, stream)?; // unicode
                }
                0x0033 => {
                    // REFERENCEORIGINAL (followed by REFERENCECONTROL)
                    reference.set_libid(stream, encoding)?;
                }
                0x002F => {
                    // REFERENCECONTROL
                    *stream = &stream[4..]; // SizeTwiddled: len of total ref control
                    reference.set_libid(stream, encoding)?;

                    *stream = &stream[6..];
                    match stream.read_u16::<LittleEndian>()? {
                        0x0016 => {
                            // optional name record extended
                            read_variable_record(stream, 1)?; // name extended
                            check_variable_record(0x003E, stream)?; // name extended unicode
                            check_record(0x0030, stream)?;
                        }
                        0x0030 => (),
                        e => {
                            return Err(VbaError::Unknown {
                                typ: "token in reference control",
                                val: e,
                            });
                        }
                    }
                    *stream = &stream[4..];
                    reference.set_libid(stream, encoding)?;
                    *stream = &stream[26..];
                }
                0x000D => {
                    // REFERENCEREGISTERED
                    *stream = &stream[4..];
                    reference.set_libid(stream, encoding)?;
                    *stream = &stream[6..];
                }
                0x000E => {
                    // REFERENCEPROJECT
                    *stream = &stream[4..];
                    let absolute = read_variable_record(stream, 1)?; // project libid absolute
                    {
                        let absolute = encoding.decode_all(absolute);
                        reference.path = absolute.strip_prefix("*\\C").unwrap_or(&absolute).into();
                    }
                    read_variable_record(stream, 1)?; // project libid relative
                    *stream = &stream[6..];
                }
                c => {
                    return Err(VbaError::Unknown {
                        typ: "check id",
                        val: c,
                    });
                }
            }
        }

        debug!("references: {:#?}", references);
        Ok(references)
    }

    fn set_libid(&mut self, stream: &mut &[u8], encoding: &XlsEncoding) -> Result<(), VbaError> {
        let libid = read_variable_record(stream, 1)?; //libid twiddled
        if libid.is_empty() || libid.ends_with(b"##") {
            return Ok(());
        }
        let libid = encoding.decode_all(libid);
        let mut parts = libid.rsplit('#');
        match (parts.next(), parts.next()) {
            (Some(desc), Some(path)) => {
                self.description = desc.into();
                // use original path if already set
                if !path.is_empty() && self.path.as_os_str().is_empty() {
                    self.path = path.into();
                }
                Ok(())
            }
            _ => Err(VbaError::LibId),
        }
    }
}

/// A vba module
#[derive(Debug, Clone, Default)]
struct Module {
    /// module name as it appears in vba project
    name: String,
    stream_name: String,
    text_offset: usize,
}

fn read_dir_information(stream: &mut &[u8]) -> Result<XlsEncoding, VbaError> {
    debug!("read dir header");

    // PROJECTSYSKIND, PROJECTLCID and PROJECTLCIDINVOKE Records
    *stream = &stream[30..];

    // PROJECT Codepage
    let encoding = XlsEncoding::from_codepage(read_u16(&stream[6..8]))?;
    *stream = &stream[8..];

    // PROJECTNAME Record
    check_variable_record(0x0004, stream)?;

    // PROJECTDOCSTRING Record
    check_variable_record(0x0005, stream)?;
    check_variable_record(0x0040, stream)?; // unicode

    // PROJECTHELPFILEPATH Record - MS-OVBA 2.3.4.2.1.7
    check_variable_record(0x0006, stream)?;
    check_variable_record(0x003D, stream)?;

    // PROJECTHELPCONTEXT PROJECTLIBFLAGS and PROJECTVERSION Records
    *stream = &stream[32..];

    // PROJECTCONSTANTS Record
    check_variable_record(0x000C, stream)?;
    check_variable_record(0x003C, stream)?; // unicode

    Ok(encoding)
}

fn read_modules(stream: &mut &[u8], encoding: &XlsEncoding) -> Result<Vec<Module>, VbaError> {
    debug!("read all modules metadata");
    *stream = &stream[4..];

    let module_len = stream.read_u16::<LittleEndian>()? as usize;

    *stream = &stream[8..]; // PROJECTCOOKIE record
    let mut modules = Vec::with_capacity(module_len);

    for _ in 0..module_len {
        // name
        let name = check_variable_record(0x0019, stream)?;
        let name = encoding.decode_all(name);

        check_variable_record(0x0047, stream)?; // unicode

        let stream_name = check_variable_record(0x001A, stream)?; // stream name
        let stream_name = encoding.decode_all(stream_name);

        check_variable_record(0x0032, stream)?; // stream name unicode
        check_variable_record(0x001C, stream)?; // doc string
        check_variable_record(0x0048, stream)?; // doc string unicode

        // offset
        check_record(0x0031, stream)?;
        *stream = &stream[4..];
        let offset = stream.read_u32::<LittleEndian>()? as usize;

        // help context
        check_record(0x001E, stream)?;
        *stream = &stream[8..];

        // cookie
        check_record(0x002C, stream)?;
        *stream = &stream[6..];

        match stream.read_u16::<LittleEndian>()? {
            0x0021 /* procedural module */ |
            0x0022 /* document, class or designer module */ => (),
            e => return Err(VbaError::Unknown { typ: "module typ", val: e }),
        }

        loop {
            *stream = &stream[4..]; // reserved
            match stream.read_u16::<LittleEndian>() {
                Ok(0x0025) /* readonly */ | Ok(0x0028) /* private */ => (),
                Ok(0x002B) => break,
                Ok(e) => return Err(VbaError::Unknown { typ: "record id", val: e }),
                Err(e) => return Err(VbaError::Io(e)),
            }
        }
        *stream = &stream[4..]; // reserved

        modules.push(Module {
            name,
            stream_name,
            text_offset: offset,
        });
    }

    Ok(modules)
}

/// Reads a variable length record
///
/// `mult` is a multiplier of the length (e.g 2 when parsing XLWideString)
fn read_variable_record<'a>(r: &mut &'a [u8], mult: usize) -> Result<&'a [u8], VbaError> {
    let len = r.read_u32::<LittleEndian>()? as usize * mult;
    let (read, next) = r.split_at(len);
    *r = next;
    Ok(read)
}

/// Check that next record matches `id` and returns a variable length record
fn check_variable_record<'a>(id: u16, r: &mut &'a [u8]) -> Result<&'a [u8], VbaError> {
    check_record(id, r)?;
    let record = read_variable_record(r, 1)?;
    if log_enabled!(Level::Warn) && record.len() > 100_000 {
        warn!(
            "record id {} as a suspicious huge length of {} (hex: {:x})",
            id,
            record.len(),
            record.len() as u32
        );
    }
    Ok(record)
}

/// Check that next record matches `id`
fn check_record(id: u16, r: &mut &[u8]) -> Result<(), VbaError> {
    debug!("check record {:x}", id);
    let record_id = r.read_u16::<LittleEndian>()?;
    if record_id != id {
        Err(VbaError::InvalidRecordId {
            expected: id,
            found: record_id,
        })
    } else {
        Ok(())
    }
}
<commit_msg>replace strip_prefix (uses unstable library feature str_strip)<commit_after>//! Parse vbaProject.bin file
//!
//! Retranscription from:
//! https://github.com/unixfreak0037/officeparser/blob/master/officeparser.py

use std::collections::HashMap;
use std::io::Read;
use std::path::PathBuf;

use byteorder::{LittleEndian, ReadBytesExt};
use log::{debug, log_enabled, warn, Level};

use crate::cfb::{Cfb, XlsEncoding};
use crate::utils::read_u16;

/// A VBA specific error enum
#[derive(Debug)]
pub enum VbaError {
    /// Error comes from a cfb parsing
    Cfb(crate::cfb::CfbError),
    /// Io error
    Io(std::io::Error),

    /// Cannot find module
    ModuleNotFound(String),
    /// Generic unknown u16 value
    Unknown {
        /// error type
        typ: &'static str,
        /// value found
        val: u16,
    },
    /// Invalid libid format
    LibId,
    /// Invalid record id
    InvalidRecordId {
        /// expected record id
        expected: u16,
        /// record if found
        found: u16,
    },
}

from_err!(crate::cfb::CfbError, VbaError, Cfb);
from_err!(std::io::Error, VbaError, Io);

impl std::fmt::Display for VbaError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            VbaError::Io(e) => write!(f, "I/O error: {}", e),
            VbaError::Cfb(e) => write!(f, "Cfb error: {}", e),

            VbaError::ModuleNotFound(e) => write!(f, "Cannot find module '{}'", e),
            VbaError::Unknown { typ, val } => write!(f, "Unknown {} '{:X}'", typ, val),
            VbaError::LibId => write!(f, "Unexpected libid format"),
            VbaError::InvalidRecordId { expected, found } => write!(
                f,
                "Invalid record id: expecting {:X} found {:X}",
                expected, found
            ),
        }
    }
}

impl std::error::Error for VbaError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            VbaError::Io(e) => Some(e),
            VbaError::Cfb(e) => Some(e),
            _ => None,
        }
    }
}

/// A struct for managing VBA reading
#[allow(dead_code)]
#[derive(Clone)]
pub struct VbaProject {
    references: Vec<Reference>,
    modules: HashMap<String, Vec<u8>>,
    encoding: XlsEncoding,
}

impl VbaProject {
    /// Create a new `VbaProject` out of the vbaProject.bin `ZipFile` or xls file
    ///
    /// Starts reading project metadata (header, directories, sectors and minisectors).
    pub fn new<R: Read>(r: &mut R, len: usize) -> Result<VbaProject, VbaError> {
        let mut cfb = Cfb::new(r, len)?;
        VbaProject::from_cfb(r, &mut cfb)
    }

    /// Creates a new `VbaProject` out of a Compound File Binary and the corresponding reader
    pub fn from_cfb<R: Read>(r: &mut R, cfb: &mut Cfb) -> Result<VbaProject, VbaError> {
        // dir stream
        let stream = cfb.get_stream("dir", r)?;
        let stream = crate::cfb::decompress_stream(&*stream)?;
        let stream = &mut &*stream;

        // read dir information record (not used)
        let encoding = read_dir_information(stream)?;

        // array of REFERENCE records
        let refs = Reference::from_stream(stream, &encoding)?;

        // modules
        let mods: Vec<Module> = read_modules(stream, &encoding)?;

        // read all modules
        let modules: HashMap<String, Vec<u8>> = mods
            .into_iter()
            .map(|m| {
                cfb.get_stream(&m.stream_name, r).and_then(|s| {
                    crate::cfb::decompress_stream(&s[m.text_offset..]).map(move |s| (m.name, s))
                })
            })
            .collect::<Result<HashMap<_, _>, _>>()?;

        Ok(VbaProject {
            references: refs,
            modules,
            encoding,
        })
    }

    /// Gets the list of `Reference`s
    pub fn get_references(&self) -> &[Reference] {
        &self.references
    }

    /// Gets the list of `Module` names
    pub fn get_module_names(&self) -> Vec<&str> {
        self.modules.keys().map(|k| &**k).collect()
    }

    /// Reads module content and tries to convert to utf8
    ///
    /// While it works most of the time, the modules are MBCS encoding and the conversion
    /// may fail. If this is the case you should revert to `read_module_raw` as there is
    /// no built in decoding provided in this crate
    ///
    /// # Examples
    /// ```
    /// use calamine::{Reader, open_workbook, Xlsx};
    ///
    /// # let path = format!("{}/tests/vba.xlsm", env!("CARGO_MANIFEST_DIR"));
    /// let mut xl: Xlsx<_> = open_workbook(path).expect("Cannot find excel file");
    /// if let Some(Ok(mut vba)) = xl.vba_project() {
    ///     let vba = vba.to_mut();
    ///     let modules = vba.get_module_names().into_iter()
    ///                      .map(|s| s.to_string()).collect::<Vec<_>>();
    ///     for m in modules {
    ///         println!("Module {}:", m);
    ///         println!("{}", vba.get_module(&m)
    ///                           .expect(&format!("cannot read {:?} module", m)));
    ///     }
    /// }
    /// ```
    pub fn get_module(&self, name: &str) -> Result<String, VbaError> {
        debug!("read module {}", name);
        let data = self.get_module_raw(name)?;
        Ok(self.encoding.decode_all(data))
    }

    /// Reads module content (MBCS encoded) and output it as-is (binary output)
    pub fn get_module_raw(&self, name: &str) -> Result<&[u8], VbaError> {
        match self.modules.get(name) {
            Some(m) => Ok(&**m),
            None => Err(VbaError::ModuleNotFound(name.into())),
        }
    }
}

/// A vba reference
#[derive(Debug, Clone, Hash, Eq, PartialEq)]
pub struct Reference {
    /// name
    pub name: String,
    /// description
    pub description: String,
    /// location of the reference
    pub path: PathBuf,
}

impl Reference {
    /// Check if the reference location is accessible
    pub fn is_missing(&self) -> bool {
        !self.path.exists()
    }

    /// Gets the list of references from the dir_stream relevant part
    fn from_stream(stream: &mut &[u8], encoding: &XlsEncoding) -> Result<Vec<Reference>, VbaError> {
        debug!("read all references metadata");

        let mut references = Vec::new();
        let mut reference = Reference {
            name: "".to_string(),
            description: "".to_string(),
            path: "".into(),
        };

        loop {
            let check = stream.read_u16::<LittleEndian>();
            match check? {
                0x000F => {
                    // termination of references array
                    if !reference.name.is_empty() {
                        references.push(reference);
                    }
                    break;
                }
                0x0016 => {
                    // REFERENCENAME
                    if !reference.name.is_empty() {
                        references.push(reference);
                    }
                    let name = read_variable_record(stream, 1)?;
                    let name = encoding.decode_all(name);
                    reference = Reference {
                        name: name.clone(),
                        description: name,
                        path: "".into(),
                    };
                    check_variable_record(0x003E, stream)?; // unicode
                }
                0x0033 => {
                    // REFERENCEORIGINAL (followed by REFERENCECONTROL)
                    reference.set_libid(stream, encoding)?;
                }
                0x002F => {
                    // REFERENCECONTROL
                    *stream = &stream[4..]; // SizeTwiddled: len of total ref control
                    reference.set_libid(stream, encoding)?;

                    *stream = &stream[6..];
                    match stream.read_u16::<LittleEndian>()? {
                        0x0016 => {
                            // optional name record extended
                            read_variable_record(stream, 1)?; // name extended
                            check_variable_record(0x003E, stream)?; // name extended unicode
                            check_record(0x0030, stream)?;
                        }
                        0x0030 => (),
                        e => {
                            return Err(VbaError::Unknown {
                                typ: "token in reference control",
                                val: e,
                            });
                        }
                    }
                    *stream = &stream[4..];
                    reference.set_libid(stream, encoding)?;
                    *stream = &stream[26..];
                }
                0x000D => {
                    // REFERENCEREGISTERED
                    *stream = &stream[4..];
                    reference.set_libid(stream, encoding)?;
                    *stream = &stream[6..];
                }
                0x000E => {
                    // REFERENCEPROJECT
                    *stream = &stream[4..];
                    let absolute = read_variable_record(stream, 1)?; // project libid absolute
                    {
                        let absolute = encoding.decode_all(absolute);
                        reference.path = if absolute.starts_with("*\\C") {
                            absolute[3..].into()
                        } else {
                            absolute.into()
                        };
                    }
                    read_variable_record(stream, 1)?; // project libid relative
                    *stream = &stream[6..];
                }
                c => {
                    return Err(VbaError::Unknown {
                        typ: "check id",
                        val: c,
                    });
                }
            }
        }

        debug!("references: {:#?}", references);
        Ok(references)
    }

    fn set_libid(&mut self, stream: &mut &[u8], encoding: &XlsEncoding) -> Result<(), VbaError> {
        let libid = read_variable_record(stream, 1)?; //libid twiddled
        if libid.is_empty() || libid.ends_with(b"##") {
            return Ok(());
        }
        let libid = encoding.decode_all(libid);
        let mut parts = libid.rsplit('#');
        match (parts.next(), parts.next()) {
            (Some(desc), Some(path)) => {
                self.description = desc.into();
                // use original path if already set
                if !path.is_empty() && self.path.as_os_str().is_empty() {
                    self.path = path.into();
                }
                Ok(())
            }
            _ => Err(VbaError::LibId),
        }
    }
}

/// A vba module
#[derive(Debug, Clone, Default)]
struct Module {
    /// module name as it appears in vba project
    name: String,
    stream_name: String,
    text_offset: usize,
}

fn read_dir_information(stream: &mut &[u8]) -> Result<XlsEncoding, VbaError> {
    debug!("read dir header");

    // PROJECTSYSKIND, PROJECTLCID and PROJECTLCIDINVOKE Records
    *stream = &stream[30..];

    // PROJECT Codepage
    let encoding = XlsEncoding::from_codepage(read_u16(&stream[6..8]))?;
    *stream = &stream[8..];

    // PROJECTNAME Record
    check_variable_record(0x0004, stream)?;

    // PROJECTDOCSTRING Record
    check_variable_record(0x0005, stream)?;
    check_variable_record(0x0040, stream)?; // unicode

    // PROJECTHELPFILEPATH Record - MS-OVBA 2.3.4.2.1.7
    check_variable_record(0x0006, stream)?;
    check_variable_record(0x003D, stream)?;

    // PROJECTHELPCONTEXT PROJECTLIBFLAGS and PROJECTVERSION Records
    *stream = &stream[32..];

    // PROJECTCONSTANTS Record
    check_variable_record(0x000C, stream)?;
    check_variable_record(0x003C, stream)?; // unicode

    Ok(encoding)
}

fn read_modules(stream: &mut &[u8], encoding: &XlsEncoding) -> Result<Vec<Module>, VbaError> {
    debug!("read all modules metadata");
    *stream = &stream[4..];

    let module_len = stream.read_u16::<LittleEndian>()? as usize;

    *stream = &stream[8..]; // PROJECTCOOKIE record
    let mut modules = Vec::with_capacity(module_len);

    for _ in 0..module_len {
        // name
        let name = check_variable_record(0x0019, stream)?;
        let name = encoding.decode_all(name);

        check_variable_record(0x0047, stream)?; // unicode

        let stream_name = check_variable_record(0x001A, stream)?; // stream name
        let stream_name = encoding.decode_all(stream_name);

        check_variable_record(0x0032, stream)?; // stream name unicode
        check_variable_record(0x001C, stream)?; // doc string
        check_variable_record(0x0048, stream)?; // doc string unicode

        // offset
        check_record(0x0031, stream)?;
        *stream = &stream[4..];
        let offset = stream.read_u32::<LittleEndian>()? as usize;

        // help context
        check_record(0x001E, stream)?;
        *stream = &stream[8..];

        // cookie
        check_record(0x002C, stream)?;
        *stream = &stream[6..];

        match stream.read_u16::<LittleEndian>()? {
            0x0021 /* procedural module */ |
            0x0022 /* document, class or designer module */ => (),
            e => return Err(VbaError::Unknown { typ: "module typ", val: e }),
        }

        loop {
            *stream = &stream[4..]; // reserved
            match stream.read_u16::<LittleEndian>() {
                Ok(0x0025) /* readonly */ | Ok(0x0028) /* private */ => (),
                Ok(0x002B) => break,
                Ok(e) => return Err(VbaError::Unknown { typ: "record id", val: e }),
                Err(e) => return Err(VbaError::Io(e)),
            }
        }
        *stream = &stream[4..]; // reserved

        modules.push(Module {
            name,
            stream_name,
            text_offset: offset,
        });
    }

    Ok(modules)
}

/// Reads a variable length record
///
/// `mult` is a multiplier of the length (e.g 2 when parsing XLWideString)
fn read_variable_record<'a>(r: &mut &'a [u8], mult: usize) -> Result<&'a [u8], VbaError> {
    let len = r.read_u32::<LittleEndian>()? as usize * mult;
    let (read, next) = r.split_at(len);
    *r = next;
    Ok(read)
}

/// Check that next record matches `id` and returns a variable length record
fn check_variable_record<'a>(id: u16, r: &mut &'a [u8]) -> Result<&'a [u8], VbaError> {
    check_record(id, r)?;
    let record = read_variable_record(r, 1)?;
    if log_enabled!(Level::Warn) && record.len() > 100_000 {
        warn!(
            "record id {} as a suspicious huge length of {} (hex: {:x})",
            id,
            record.len(),
            record.len() as u32
        );
    }
    Ok(record)
}

/// Check that next record matches `id`
fn check_record(id: u16, r: &mut &[u8]) -> Result<(), VbaError> {
    debug!("check record {:x}", id);
    let record_id = r.read_u16::<LittleEndian>()?;
    if record_id != id {
        Err(VbaError::InvalidRecordId {
            expected: id,
            found: record_id,
        })
    } else {
        Ok(())
    }
}
<|endoftext|>