<commit_before>use std::cell::Cell;
use std::mem;
use std::ptr;
use std::sync::atomic::{AtomicPtr, AtomicUsize, Ordering, fence};

use alloc;
use consume::CONSUME;
use countedindex::{CountedIndex, Transaction};
use maybe_acquire::{MAYBE_ACQUIRE, maybe_acquire_fence};
use memory::MemoryManager;

#[derive(Clone, Copy)]
enum ReaderState {
    Single,
    Multi,
}

pub struct ReaderMeta {
    state: Cell<ReaderState>,
    num_consumers: AtomicUsize,
}

pub struct Reader {
    pos_data: CountedIndex,
    meta: *const ReaderMeta,
}

/// This represents the reader attempt at loading a transaction
/// It behaves similarly to a Transaction but has logic for single/multi
/// readers
struct ReadAttempt<'a> {
    linked: Transaction<'a>,
    reader: &'a Reader,
    state: ReaderState,
}

/// This holds the set of readers currently active.
/// This struct is held out of line from the cursor so it's easy to atomically replace it
struct ReaderGroup {
    readers: Vec<*const Reader>,
}

#[repr(C)]
pub struct ReadCursor {
    readers: AtomicPtr<ReaderGroup>,
}

impl<'a> ReadAttempt<'a> {
    #[inline(always)]
    pub fn get(&self) -> (isize, usize) {
        self.linked.get()
    }

    #[inline(always)]
    pub fn commit_attempt(self, by: u32, ord: Ordering) -> Option<ReadAttempt<'a>> {
        match self.state {
            ReaderState::Single => {
                self.linked.commit_direct(by, ord);
                None
            }
            ReaderState::Multi => {
                if unsafe { (*self.reader.meta).num_consumers.load(Ordering::Relaxed) } == 1 {
                    fence(Ordering::Acquire);
                    unsafe {
                        (*self.reader.meta).state.set(ReaderState::Single);
                    }
                    self.linked.commit_direct(by, ord);
                    None
                } else {
                    match self.linked.commit(by, ord) {
                        Some(transaction) => {
                            Some(ReadAttempt {
                                linked: transaction,
                                reader: self.reader,
                                state: ReaderState::Multi,
                            })
                        }
                        None => None,
                    }
                }
            }
        }
    }
}

impl Reader {
    #[inline(always)]
    pub fn load_attempt(&self, ord: Ordering) -> ReadAttempt {
        ReadAttempt {
            linked: self.pos_data.load_transaction(ord),
            reader: self,
            state: unsafe { (*self.meta).state.get() },
        }
    }

    pub fn dup_consumer(&self) {
        unsafe {
            if (*self.meta).num_consumers.fetch_add(1, Ordering::SeqCst) == 1 {
                (*self.meta).state.set(ReaderState::Multi);
            }
        }
    }

    pub fn remove_consumer(&self) -> usize {
        unsafe { (*self.meta).num_consumers.fetch_sub(1, Ordering::SeqCst) }
    }

    pub fn get_consumers(&self) -> usize {
        unsafe { (*self.meta).num_consumers.load(Ordering::Relaxed) }
    }
}

impl ReaderGroup {
    pub fn new() -> ReaderGroup {
        ReaderGroup { readers: Vec::new() }
    }

    /// Only safe to call from a consumer of the queue!
    pub unsafe fn add_reader(&self, raw: usize, wrap: u32) -> (*mut ReaderGroup, *mut Reader) {
        let new_reader = alloc::allocate(1);
        let new_meta = alloc::allocate(1);
        let new_group = alloc::allocate(1);
        ptr::write(new_meta,
                   ReaderMeta {
                       state: Cell::new(ReaderState::Single),
                       num_consumers: AtomicUsize::new(1),
                   });
        ptr::write(new_reader,
                   Reader {
                       pos_data: CountedIndex::from_usize(raw, wrap),
                       meta: new_meta as *const ReaderMeta,
                   });
        let mut new_readers = self.readers.clone();
        new_readers.push(new_reader as *const Reader);
        ptr::write(new_group, ReaderGroup { readers: new_readers });
        (new_group, new_reader)
    }

    pub unsafe fn remove_reader(&self, reader: *const Reader) -> *mut ReaderGroup {
        let new_group = alloc::allocate(1);
        let mut new_readers = self.readers.clone();
        new_readers.retain(|pt| *pt != reader);
        ptr::write(new_group, ReaderGroup { readers: new_readers });
        new_group
    }

    pub fn get_max_diff(&self, cur_writer: usize) -> Option<u32> {
        let mut max_diff: usize = 0;
        unsafe {
            for reader_ptr in &self.readers {
                // If a reader has passed the writer during this function call
                // then what must have happened is that somebody else has completed this
                // written to the queue, and a reader has bypassed it. We should retry
                let rpos = (**reader_ptr).pos_data.load_count(MAYBE_ACQUIRE);
                let diff = cur_writer.wrapping_sub(rpos);
                if diff > (1 << 30) {
                    return None;
                }
                max_diff = if diff > max_diff { diff } else { max_diff };
            }
        }
        maybe_acquire_fence();

        assert!(max_diff <= (1 << 30));
        Some(max_diff as u32)
    }
}

impl ReadCursor {
    pub fn new(wrap: u32) -> (ReadCursor, *const Reader) {
        let rg = ReaderGroup::new();
        unsafe {
            let (real_group, reader) = rg.add_reader(0, wrap);
            (ReadCursor { readers: AtomicPtr::new(real_group) }, reader)
        }
    }

    #[inline(always)]
    pub fn prefetch_metadata(&self) {
        unsafe {
            let rg = &*self.readers.load(CONSUME);
            let dummy_ptr: *const usize = mem::transmute(&rg.readers);
            ptr::read_volatile(dummy_ptr);
        }
    }

    pub fn get_max_diff(&self, cur_writer: usize) -> Option<u32> {
        loop {
            unsafe {
                let first_ptr = self.readers.load(CONSUME);
                let rg = &*first_ptr;
                let rval = rg.get_max_diff(cur_writer);
                // This check ensures that the pointer hasn't changed
                // We must first read the diff, *and then* check the pointer
                // for changes. This is extremely similar to the seqlock except
                // with the pointer as the sequence lock instead of a different int
                //
                // We don't need another acquire fence here since the
                // relevant loads in get_max_diff are going to ordered
                // before all loads after the function exit, and also
                // ordered after the original pointer load
                let second_ptr = self.readers.load(MAYBE_ACQUIRE);
                if second_ptr == first_ptr {
                    maybe_acquire_fence();
                    return rval;
                }
            }
        }
    }

    // There's no fundamental reason these need to leak,
    // I just haven't implemented the memory management yet.
    // It's not too hard since we can track readers and writers active
    pub fn add_reader(&self, reader: &Reader, manager: &MemoryManager) -> *const Reader {
        let mut current_ptr = self.readers.load(CONSUME);
        loop {
            unsafe {
                let current_group = &*current_ptr;
                let raw = reader.pos_data.load_raw(Ordering::Relaxed);
                let wrap = reader.pos_data.wrap_at();
                let (new_group, new_reader) = current_group.add_reader(raw, wrap);
                match self.readers
                    .compare_exchange(current_ptr, new_group, Ordering::SeqCst, Ordering::SeqCst) {
                    Ok(_) => {
                        fence(Ordering::SeqCst);
                        manager.free(current_ptr, 1);
                        return new_reader as *const Reader;
                    }
                    Err(val) => {
                        current_ptr = val;
                        ptr::read(new_group);
                        alloc::deallocate((*new_reader).meta as *mut ReaderMeta, 1);
                        alloc::deallocate(new_group, 1);
                        alloc::deallocate(new_reader, 1);
                    }
                }
            }
        }
    }

    pub fn remove_reader(&self, reader: &Reader, mem: &MemoryManager) {
        let mut current_group = self.readers.load(CONSUME);
        loop {
            unsafe {
                let new_group = (*current_group).remove_reader(reader);
                match self.readers
                    .compare_exchange(current_group,
                                      new_group,
                                      Ordering::SeqCst,
                                      Ordering::SeqCst) {
                    Ok(_) => {
                        fence(Ordering::SeqCst);
                        mem.free(current_group, 1);
                        break;
                    }
                    Err(val) => {
                        current_group = val;
                        ptr::read(new_group);
                        alloc::deallocate(reader.meta as *mut ReaderMeta, 1);
                        alloc::deallocate(new_group, 1);
                    }
                }
            }
        }
    }
}
<commit_msg>Pos now a ptr<commit_after>use std::cell::Cell;
use std::mem;
use std::ptr;
use std::sync::atomic::{AtomicPtr, AtomicUsize, Ordering, fence};

use alloc;
use consume::CONSUME;
use countedindex::{CountedIndex, Transaction};
use maybe_acquire::{MAYBE_ACQUIRE, maybe_acquire_fence};
use memory::MemoryManager;

#[derive(Clone, Copy)]
enum ReaderState {
    Single,
    Multi,
}

pub struct ReaderPos {
    pos_data: CountedIndex,
}

pub struct ReaderMeta {
    state: Cell<ReaderState>,
    num_consumers: AtomicUsize,
}

pub struct Reader {
    pos: *const ReaderPos,
    meta: *const ReaderMeta,
}

/// This represents the reader attempt at loading a transaction
/// It behaves similarly to a Transaction but has logic for single/multi
/// readers
struct ReadAttempt<'a> {
    linked: Transaction<'a>,
    reader: &'a Reader,
    state: ReaderState,
}

/// This holds the set of readers currently active.
/// This struct is held out of line from the cursor so it's easy to atomically replace it
struct ReaderGroup {
    readers: Vec<*const Reader>,
}

#[repr(C)]
pub struct ReadCursor {
    readers: AtomicPtr<ReaderGroup>,
}

impl<'a> ReadAttempt<'a> {
    #[inline(always)]
    pub fn get(&self) -> (isize, usize) {
        self.linked.get()
    }

    #[inline(always)]
    pub fn commit_attempt(self, by: u32, ord: Ordering) -> Option<ReadAttempt<'a>> {
        match self.state {
            ReaderState::Single => {
                self.linked.commit_direct(by, ord);
                None
            }
            ReaderState::Multi => {
                if unsafe { (*self.reader.meta).num_consumers.load(Ordering::Relaxed) } == 1 {
                    fence(Ordering::Acquire);
                    unsafe {
                        (*self.reader.meta).state.set(ReaderState::Single);
                    }
                    self.linked.commit_direct(by, ord);
                    None
                } else {
                    match self.linked.commit(by, ord) {
                        Some(transaction) => {
                            Some(ReadAttempt {
                                linked: transaction,
                                reader: self.reader,
                                state: ReaderState::Multi,
                            })
                        }
                        None => None,
                    }
                }
            }
        }
    }
}

impl Reader {
    #[inline(always)]
    pub fn load_attempt(&self, ord: Ordering) -> ReadAttempt {
        unsafe {
            ReadAttempt {
                linked: (*self.pos).pos_data.load_transaction(ord),
                reader: self,
                state: (*self.meta).state.get(),
            }
        }
    }

    pub fn dup_consumer(&self) {
        unsafe {
            if (*self.meta).num_consumers.fetch_add(1, Ordering::SeqCst) == 1 {
                (*self.meta).state.set(ReaderState::Multi);
            }
        }
    }

    pub fn remove_consumer(&self) -> usize {
        unsafe { (*self.meta).num_consumers.fetch_sub(1, Ordering::SeqCst) }
    }

    pub fn get_consumers(&self) -> usize {
        unsafe { (*self.meta).num_consumers.load(Ordering::Relaxed) }
    }
}

impl ReaderGroup {
    pub fn new() -> ReaderGroup {
        ReaderGroup { readers: Vec::new() }
    }

    /// Only safe to call from a consumer of the queue!
    pub unsafe fn add_reader(&self, raw: usize, wrap: u32) -> (*mut ReaderGroup, *mut Reader) {
        let new_reader = alloc::allocate(1);
        let new_meta = alloc::allocate(1);
        let new_group = alloc::allocate(1);
        let new_pos = alloc::allocate(1);
        ptr::write(new_pos, ReaderPos { pos_data: CountedIndex::from_usize(raw, wrap) });
        ptr::write(new_meta,
                   ReaderMeta {
                       state: Cell::new(ReaderState::Single),
                       num_consumers: AtomicUsize::new(1),
                   });
        ptr::write(new_reader,
                   Reader {
                       pos: new_pos,
                       meta: new_meta as *const ReaderMeta,
                   });
        let mut new_readers = self.readers.clone();
        new_readers.push(new_reader as *const Reader);
        ptr::write(new_group, ReaderGroup { readers: new_readers });
        (new_group, new_reader)
    }

    pub unsafe fn remove_reader(&self, reader: *const Reader) -> *mut ReaderGroup {
        let new_group = alloc::allocate(1);
        let mut new_readers = self.readers.clone();
        new_readers.retain(|pt| *pt != reader);
        ptr::write(new_group, ReaderGroup { readers: new_readers });
        new_group
    }

    pub fn get_max_diff(&self, cur_writer: usize) -> Option<u32> {
        let mut max_diff: usize = 0;
        unsafe {
            for reader_ptr in &self.readers {
                // If a reader has passed the writer during this function call
                // then what must have happened is that somebody else has completed this
                // written to the queue, and a reader has bypassed it. We should retry
                let rpos = (*(**reader_ptr).pos).pos_data.load_count(MAYBE_ACQUIRE);
                let diff = cur_writer.wrapping_sub(rpos);
                if diff > (1 << 30) {
                    return None;
                }
                max_diff = if diff > max_diff { diff } else { max_diff };
            }
        }
        maybe_acquire_fence();

        assert!(max_diff <= (1 << 30));
        Some(max_diff as u32)
    }
}

impl ReadCursor {
    pub fn new(wrap: u32) -> (ReadCursor, *const Reader) {
        let rg = ReaderGroup::new();
        unsafe {
            let (real_group, reader) = rg.add_reader(0, wrap);
            (ReadCursor { readers: AtomicPtr::new(real_group) }, reader)
        }
    }

    #[inline(always)]
    pub fn prefetch_metadata(&self) {
        unsafe {
            let rg = &*self.readers.load(CONSUME);
            let dummy_ptr: *const usize = mem::transmute(&rg.readers);
            ptr::read_volatile(dummy_ptr);
        }
    }

    pub fn get_max_diff(&self, cur_writer: usize) -> Option<u32> {
        loop {
            unsafe {
                let first_ptr = self.readers.load(CONSUME);
                let rg = &*first_ptr;
                let rval = rg.get_max_diff(cur_writer);
                // This check ensures that the pointer hasn't changed
                // We must first read the diff, *and then* check the pointer
                // for changes. This is extremely similar to the seqlock except
                // with the pointer as the sequence lock instead of a different int
                //
                // We don't need another acquire fence here since the
                // relevant loads in get_max_diff are going to ordered
                // before all loads after the function exit, and also
                // ordered after the original pointer load
                let second_ptr = self.readers.load(MAYBE_ACQUIRE);
                if second_ptr == first_ptr {
                    maybe_acquire_fence();
                    return rval;
                }
            }
        }
    }

    // There's no fundamental reason these need to leak,
    // I just haven't implemented the memory management yet.
    // It's not too hard since we can track readers and writers active
    pub fn add_reader(&self, reader: &Reader, manager: &MemoryManager) -> *const Reader {
        let mut current_ptr = self.readers.load(CONSUME);
        loop {
            unsafe {
                let current_group = &*current_ptr;
                let raw = (*reader.pos).pos_data.load_raw(Ordering::Relaxed);
                let wrap = (*reader.pos).pos_data.wrap_at();
                let (new_group, new_reader) = current_group.add_reader(raw, wrap);
                match self.readers
                    .compare_exchange(current_ptr, new_group, Ordering::SeqCst, Ordering::SeqCst) {
                    Ok(_) => {
                        fence(Ordering::SeqCst);
                        manager.free(current_ptr, 1);
                        return new_reader as *const Reader;
                    }
                    Err(val) => {
                        current_ptr = val;
                        ptr::read(new_group);
                        alloc::deallocate((*new_reader).meta as *mut ReaderMeta, 1);
                        alloc::deallocate((*new_reader).pos as *mut ReaderPos, 1);
                        alloc::deallocate(new_group, 1);
                        alloc::deallocate(new_reader, 1);
                    }
                }
            }
        }
    }

    pub fn remove_reader(&self, reader: &Reader, mem: &MemoryManager) {
        let mut current_group = self.readers.load(CONSUME);
        loop {
            unsafe {
                let new_group = (*current_group).remove_reader(reader);
                match self.readers
                    .compare_exchange(current_group,
                                      new_group,
                                      Ordering::SeqCst,
                                      Ordering::SeqCst) {
                    Ok(_) => {
                        fence(Ordering::SeqCst);
                        mem.free(current_group, 1);
                        mem.free(reader.pos as *mut ReaderPos, 1);
                        alloc::deallocate(reader.meta as *mut ReaderMeta, 1);
                        break;
                    }
                    Err(val) => {
                        current_group = val;
                        ptr::read(new_group);
                        alloc::deallocate(new_group, 1);
                    }
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>//! Types for the *m.room.member* event.

use std::collections::BTreeMap;

use ruma_events_macros::StateEventContent;
use ruma_identifiers::UserId;
use serde::{Deserialize, Serialize};
use strum::{Display, EnumString};

use crate::StateEvent;

/// The current membership state of a user in the room.
///
/// Adjusts the membership state for a user in a room. It is preferable to use the membership
/// APIs (`/rooms/<room id>/invite` etc) when performing membership actions rather than
/// adjusting the state directly as there are a restricted set of valid transformations. For
/// example, user A cannot force user B to join a room, and trying to force this state change
/// directly will fail.
///
/// The `third_party_invite` property will be set if this invite is an *invite* event and is the
/// successor of an *m.room.third_party_invite* event, and absent otherwise.
///
/// This event may also include an `invite_room_state` key inside the event's unsigned data. If
/// present, this contains an array of `StrippedState` events. These events provide information
/// on a subset of state events such as the room name. Note that ruma-events treats unsigned
/// data on events as arbitrary JSON values, and the ruma-events types for this event don't
/// provide direct access to `invite_room_state`. If you need this data, you must extract and
/// convert it from a `serde_json::Value` yourself.
///
/// The user for which a membership applies is represented by the `state_key`. Under some
/// conditions, the `sender` and `state_key` may not match - this may be interpreted as the
/// `sender` affecting the membership state of the `state_key` user.
///
/// The membership for a given user can change over time. Previous membership can be retrieved
/// from the `prev_content` object on an event. If not present, the user's previous membership
/// must be assumed as leave.
pub type MemberEvent = StateEvent<MemberEventContent>;

/// The payload for `MemberEvent`.
#[derive(Clone, Debug, Deserialize, Serialize, StateEventContent)]
#[ruma_event(type = "m.room.member")]
pub struct MemberEventContent {
    /// The avatar URL for this user, if any. This is added by the homeserver.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub avatar_url: Option<String>,

    /// The display name for this user, if any. This is added by the homeserver.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub displayname: Option<String>,

    /// Flag indicating if the room containing this event was created
    /// with the intention of being a direct chat.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub is_direct: Option<bool>,

    /// The membership state of this user.
    pub membership: MembershipState,

    /// If this member event is the successor to a third party invitation, this field will
    /// contain information about that invitation.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub third_party_invite: Option<ThirdPartyInvite>,
}

/// The membership state of a user.
#[derive(Clone, Copy, Debug, PartialEq, Display, EnumString, Deserialize, Serialize)]
#[serde(rename_all = "lowercase")]
#[strum(serialize_all = "lowercase")]
pub enum MembershipState {
    /// The user is banned.
    Ban,

    /// The user has been invited.
    Invite,

    /// The user has joined.
    Join,

    /// The user has requested to join.
    Knock,

    /// The user has left.
    Leave,
}

/// Information about a third party invitation.
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct ThirdPartyInvite {
    /// A name which can be displayed to represent the user instead of their third party
    /// identifier.
    pub display_name: String,

    /// A block of content which has been signed, which servers can use to verify the event.
    /// Clients should ignore this.
    pub signed: SignedContent,
}

/// A block of content which has been signed, which servers can use to verify a third party
/// invitation.
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct SignedContent {
    /// The invited Matrix user ID.
    ///
    /// Must be equal to the user_id property of the event.
    pub mxid: UserId,

    /// A single signature from the verifying server, in the format specified by the Signing Events
    /// section of the server-server API.
    pub signatures: BTreeMap<String, BTreeMap<String, String>>,

    /// The token property of the containing third_party_invite object.
    pub token: String,
}

/// Translation of the membership change in `m.room.member` event.
#[derive(Clone, Copy, Debug, Deserialize, Serialize)]
pub enum MembershipChange {
    /// No change.
    None,

    /// Must never happen.
    Error,

    /// User joined the room.
    Joined,

    /// User left the room.
    Left,

    /// User was banned.
    Banned,

    /// User was unbanned.
    Unbanned,

    /// User was kicked.
    Kicked,

    /// User was invited.
    Invited,

    /// User was kicked and banned.
    KickedAndBanned,

    /// User rejected the invite.
    InvitationRejected,

    /// User had their invite revoked.
    InvitationRevoked,

    /// `displayname` or `avatar_url` changed.
    ProfileChanged,

    /// Not implemented.
    NotImplemented,
}

impl MemberEvent {
    /// Helper function for membership change. Check [the specification][spec] for details.
    ///
    /// [spec]: https://matrix.org/docs/spec/client_server/latest#m-room-member
    pub fn membership_change(&self) -> MembershipChange {
        use MembershipState::*;
        let prev_membership = if let Some(prev_content) = &self.prev_content {
            prev_content.membership
        } else {
            Leave
        };
        match (prev_membership, &self.content.membership) {
            (Invite, Invite) | (Leave, Leave) | (Ban, Ban) => MembershipChange::None,
            (Invite, Join) | (Leave, Join) => MembershipChange::Joined,
            (Invite, Leave) => {
                if self.sender == self.state_key {
                    MembershipChange::InvitationRevoked
                } else {
                    MembershipChange::InvitationRejected
                }
            }
            (Invite, Ban) | (Leave, Ban) => MembershipChange::Banned,
            (Join, Invite) | (Ban, Invite) | (Ban, Join) => MembershipChange::Error,
            (Join, Join) => MembershipChange::ProfileChanged,
            (Join, Leave) => {
                if self.sender == self.state_key {
                    MembershipChange::Left
                } else {
                    MembershipChange::Kicked
                }
            }
            (Join, Ban) => MembershipChange::KickedAndBanned,
            (Leave, Invite) => MembershipChange::Invited,
            (Ban, Leave) => MembershipChange::Unbanned,
            (Knock, _) | (_, Knock) => MembershipChange::NotImplemented,
        }
    }
}

#[cfg(test)]
mod tests {
    use std::time::{Duration, UNIX_EPOCH};

    use maplit::btreemap;
    use matches::assert_matches;
    use serde_json::{from_value as from_json_value, json};

    use super::{MemberEventContent, MembershipState, SignedContent, ThirdPartyInvite};
    use crate::{EventJson, StateEvent};

    #[test]
    fn serde_with_no_prev_content() {
        let json = json!({
            "type": "m.room.member",
            "content": {
                "membership": "join"
            },
            "event_id": "$h29iv0s8:example.com",
            "origin_server_ts": 1,
            "room_id": "!n8f893n9:example.com",
            "sender": "@carl:example.com",
            "state_key": "example.com"
        });

        assert_matches!(
            from_json_value::<EventJson<StateEvent<MemberEventContent>>>(json)
                .unwrap()
                .deserialize()
                .unwrap(),
            StateEvent::<MemberEventContent> {
                content: MemberEventContent {
                    avatar_url: None,
                    displayname: None,
                    is_direct: None,
                    membership: MembershipState::Join,
                    third_party_invite: None,
                },
                event_id,
                origin_server_ts,
                room_id,
                sender,
                state_key,
                unsigned,
                prev_content: None,
            } if event_id == "$h29iv0s8:example.com"
                && origin_server_ts == UNIX_EPOCH + Duration::from_millis(1)
                && room_id == "!n8f893n9:example.com"
                && sender == "@carl:example.com"
                && state_key == "example.com"
                && unsigned.is_empty()
        );
    }

    #[test]
    fn serde_with_prev_content() {
        let json = json!({
            "type": "m.room.member",
            "content": {
                "membership": "join"
            },
            "event_id": "$h29iv0s8:example.com",
            "origin_server_ts": 1,
            "prev_content": {
                "membership": "join"
            },
            "room_id": "!n8f893n9:example.com",
            "sender": "@carl:example.com",
            "state_key": "example.com"
        });

        assert_matches!(
            from_json_value::<EventJson<StateEvent<MemberEventContent>>>(json)
                .unwrap()
                .deserialize()
                .unwrap(),
            StateEvent::<MemberEventContent> {
                content: MemberEventContent {
                    avatar_url: None,
                    displayname: None,
                    is_direct: None,
                    membership: MembershipState::Join,
                    third_party_invite: None,
                },
                event_id,
                origin_server_ts,
                room_id,
                sender,
                state_key,
                unsigned,
                prev_content: Some(MemberEventContent {
                    avatar_url: None,
                    displayname: None,
                    is_direct: None,
                    membership: MembershipState::Join,
                    third_party_invite: None,
                }),
            } if event_id == "$h29iv0s8:example.com"
                && origin_server_ts == UNIX_EPOCH + Duration::from_millis(1)
                && room_id == "!n8f893n9:example.com"
                && sender == "@carl:example.com"
                && state_key == "example.com"
                && unsigned.is_empty()
        );
    }

    #[test]
    fn serde_with_content_full() {
        let json = json!({
            "type": "m.room.member",
            "content": {
                "avatar_url": "mxc://example.org/SEsfnsuifSDFSSEF",
                "displayname": "Alice Margatroid",
                "is_direct": true,
                "membership": "invite",
                "third_party_invite": {
                    "display_name": "alice",
                    "signed": {
                        "mxid": "@alice:example.org",
                        "signatures": {
                            "magic.forest": {
                                "ed25519:3": "foobar"
                            }
                        },
                        "token": "abc123"
                    }
                }
            },
            "event_id": "$143273582443PhrSn:example.org",
            "origin_server_ts": 233,
            "room_id": "!jEsUZKDJdhlrceRyVU:example.org",
            "sender": "@alice:example.org",
            "state_key": "@alice:example.org"
        });

        assert_matches!(
            from_json_value::<EventJson<StateEvent<MemberEventContent>>>(json)
                .unwrap()
                .deserialize()
                .unwrap(),
            StateEvent::<MemberEventContent> {
                content: MemberEventContent {
                    avatar_url: Some(avatar_url),
                    displayname: Some(displayname),
                    is_direct: Some(true),
                    membership: MembershipState::Invite,
                    third_party_invite: Some(ThirdPartyInvite {
                        display_name: third_party_displayname,
                        signed: SignedContent { mxid, signatures, token },
                    }),
                },
                event_id,
                origin_server_ts,
                room_id,
                sender,
                state_key,
                unsigned,
                prev_content: None,
            } if avatar_url == "mxc://example.org/SEsfnsuifSDFSSEF"
                && displayname == "Alice Margatroid"
                && third_party_displayname == "alice"
                && mxid == "@alice:example.org"
                && signatures == btreemap! {
                    "magic.forest".to_owned() => btreemap! {
                        "ed25519:3".to_owned() => "foobar".to_owned()
                    }
                }
                && token == "abc123"
                && event_id == "$143273582443PhrSn:example.org"
                && origin_server_ts == UNIX_EPOCH + Duration::from_millis(233)
                && room_id == "!jEsUZKDJdhlrceRyVU:example.org"
                && sender == "@alice:example.org"
                && state_key == "@alice:example.org"
                && unsigned.is_empty()
        )
    }

    #[test]
    fn serde_with_prev_content_full() {
        let json = json!({
            "type": "m.room.member",
            "content": {
                "membership": "join"
            },
            "event_id": "$143273582443PhrSn:example.org",
            "origin_server_ts": 233,
            "prev_content": {
                "avatar_url": "mxc://example.org/SEsfnsuifSDFSSEF",
                "displayname": "Alice Margatroid",
                "is_direct": true,
                "membership": "invite",
                "third_party_invite": {
                    "display_name": "alice",
                    "signed": {
                        "mxid": "@alice:example.org",
                        "signatures": {
                            "magic.forest": {
                                "ed25519:3": "foobar"
                            }
                        },
                        "token": "abc123"
                    }
                }
            },
            "room_id": "!jEsUZKDJdhlrceRyVU:example.org",
            "sender": "@alice:example.org",
            "state_key": "@alice:example.org"
        });

        assert_matches!(
            from_json_value::<EventJson<StateEvent<MemberEventContent>>>(json)
                .unwrap()
                .deserialize()
                .unwrap(),
            StateEvent::<MemberEventContent> {
                content: MemberEventContent {
                    avatar_url: None,
                    displayname: None,
                    is_direct: None,
                    membership: MembershipState::Join,
                    third_party_invite: None,
                },
                event_id,
                origin_server_ts,
                room_id,
                sender,
                state_key,
                unsigned,
                prev_content: Some(MemberEventContent {
                    avatar_url: Some(avatar_url),
                    displayname: Some(displayname),
                    is_direct: Some(true),
                    membership: MembershipState::Invite,
                    third_party_invite: Some(ThirdPartyInvite {
                        display_name: third_party_displayname,
                        signed: SignedContent { mxid, signatures, token },
                    }),
                }),
            } if event_id == "$143273582443PhrSn:example.org"
                && origin_server_ts == UNIX_EPOCH + Duration::from_millis(233)
                && room_id == "!jEsUZKDJdhlrceRyVU:example.org"
                && sender == "@alice:example.org"
                && state_key == "@alice:example.org"
                && unsigned.is_empty()
                && avatar_url == "mxc://example.org/SEsfnsuifSDFSSEF"
                && displayname == "Alice Margatroid"
                && third_party_displayname == "alice"
                && mxid == "@alice:example.org"
                && signatures == btreemap! {
                    "magic.forest".to_owned() => btreemap! {
                        "ed25519:3".to_owned() => "foobar".to_owned()
                    }
                }
                && token == "abc123"
        );
    }
}
<commit_msg>Extend MembershipChange::ProfileChanged<commit_after>//! Types for the *m.room.member* event.

use std::collections::BTreeMap;

use ruma_events_macros::StateEventContent;
use ruma_identifiers::UserId;
use serde::{Deserialize, Serialize};
use strum::{Display, EnumString};

use crate::StateEvent;

/// The current membership state of a user in the room.
///
/// Adjusts the membership state for a user in a room. It is preferable to use the membership
/// APIs (`/rooms/<room id>/invite` etc) when performing membership actions rather than
/// adjusting the state directly as there are a restricted set of valid transformations. For
/// example, user A cannot force user B to join a room, and trying to force this state change
/// directly will fail.
///
/// The `third_party_invite` property will be set if this invite is an *invite* event and is the
/// successor of an *m.room.third_party_invite* event, and absent otherwise.
///
/// This event may also include an `invite_room_state` key inside the event's unsigned data. If
/// present, this contains an array of `StrippedState` events. These events provide information
/// on a subset of state events such as the room name. Note that ruma-events treats unsigned
/// data on events as arbitrary JSON values, and the ruma-events types for this event don't
/// provide direct access to `invite_room_state`. If you need this data, you must extract and
/// convert it from a `serde_json::Value` yourself.
///
/// The user for which a membership applies is represented by the `state_key`. Under some
/// conditions, the `sender` and `state_key` may not match - this may be interpreted as the
/// `sender` affecting the membership state of the `state_key` user.
///
/// The membership for a given user can change over time. Previous membership can be retrieved
/// from the `prev_content` object on an event. If not present, the user's previous membership
/// must be assumed as leave.
pub type MemberEvent = StateEvent<MemberEventContent>;

/// The payload for `MemberEvent`.
#[derive(Clone, Debug, Deserialize, Serialize, StateEventContent)]
#[ruma_event(type = "m.room.member")]
pub struct MemberEventContent {
    /// The avatar URL for this user, if any. This is added by the homeserver.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub avatar_url: Option<String>,

    /// The display name for this user, if any. This is added by the homeserver.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub displayname: Option<String>,

    /// Flag indicating if the room containing this event was created
    /// with the intention of being a direct chat.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub is_direct: Option<bool>,

    /// The membership state of this user.
    pub membership: MembershipState,

    /// If this member event is the successor to a third party invitation, this field will
    /// contain information about that invitation.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub third_party_invite: Option<ThirdPartyInvite>,
}

/// The membership state of a user.
#[derive(Clone, Copy, Debug, PartialEq, Display, EnumString, Deserialize, Serialize)]
#[serde(rename_all = "lowercase")]
#[strum(serialize_all = "lowercase")]
pub enum MembershipState {
    /// The user is banned.
    Ban,

    /// The user has been invited.
    Invite,

    /// The user has joined.
    Join,

    /// The user has requested to join.
    Knock,

    /// The user has left.
    Leave,
}

/// Information about a third party invitation.
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct ThirdPartyInvite {
    /// A name which can be displayed to represent the user instead of their third party
    /// identifier.
    pub display_name: String,

    /// A block of content which has been signed, which servers can use to verify the event.
    /// Clients should ignore this.
    pub signed: SignedContent,
}

/// A block of content which has been signed, which servers can use to verify a third party
/// invitation.
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct SignedContent {
    /// The invited Matrix user ID.
    ///
    /// Must be equal to the user_id property of the event.
    pub mxid: UserId,

    /// A single signature from the verifying server, in the format specified by the Signing Events
    /// section of the server-server API.
    pub signatures: BTreeMap<String, BTreeMap<String, String>>,

    /// The token property of the containing third_party_invite object.
    pub token: String,
}

/// Translation of the membership change in `m.room.member` event.
#[derive(Clone, Copy, Debug, Deserialize, Serialize)]
pub enum MembershipChange {
    /// No change.
    None,

    /// Must never happen.
    Error,

    /// User joined the room.
    Joined,

    /// User left the room.
    Left,

    /// User was banned.
    Banned,

    /// User was unbanned.
    Unbanned,

    /// User was kicked.
    Kicked,

    /// User was invited.
    Invited,

    /// User was kicked and banned.
    KickedAndBanned,

    /// User rejected the invite.
    InvitationRejected,

    /// User had their invite revoked.
    InvitationRevoked,

    /// `displayname` or `avatar_url` changed.
    ProfileChanged {
        /// Whether the `displayname` changed.
        displayname_changed: bool,
        /// Whether the `avatar_url` changed.
        avatar_url_changed: bool,
    },

    /// Not implemented.
    NotImplemented,
}

impl MemberEvent {
    /// Helper function for membership change. Check [the specification][spec] for details.
    ///
    /// [spec]: https://matrix.org/docs/spec/client_server/latest#m-room-member
    pub fn membership_change(&self) -> MembershipChange {
        use MembershipState::*;
        let prev_content = if let Some(prev_content) = &self.prev_content {
            prev_content
        } else {
            &MemberEventContent {
                avatar_url: None,
                displayname: None,
                is_direct: None,
                membership: Leave,
                third_party_invite: None,
            }
        };

        match (prev_content.membership, &self.content.membership) {
            (Invite, Invite) | (Leave, Leave) | (Ban, Ban) => MembershipChange::None,
            (Invite, Join) | (Leave, Join) => MembershipChange::Joined,
            (Invite, Leave) => {
                if self.sender == self.state_key {
                    MembershipChange::InvitationRevoked
                } else {
                    MembershipChange::InvitationRejected
                }
            }
            (Invite, Ban) | (Leave, Ban) => MembershipChange::Banned,
            (Join, Invite) | (Ban, Invite) | (Ban, Join) => MembershipChange::Error,
            (Join, Join) => MembershipChange::ProfileChanged {
                displayname_changed: prev_content.displayname != self.content.displayname,
                avatar_url_changed: prev_content.avatar_url != self.content.avatar_url,
            },
            (Join, Leave) => {
                if self.sender == self.state_key {
                    MembershipChange::Left
                } else {
                    MembershipChange::Kicked
                }
            }
            (Join, Ban) => MembershipChange::KickedAndBanned,
            (Leave, Invite) => MembershipChange::Invited,
            (Ban, Leave) => MembershipChange::Unbanned,
            (Knock, _) | (_, Knock) => MembershipChange::NotImplemented,
        }
    }
}

#[cfg(test)]
mod tests {
    use std::time::{Duration, UNIX_EPOCH};

    use maplit::btreemap;
    use matches::assert_matches;
    use serde_json::{from_value as from_json_value, json};

    use super::{MemberEventContent, MembershipState, SignedContent, ThirdPartyInvite};
    use crate::{EventJson, StateEvent};

    #[test]
    fn serde_with_no_prev_content() {
        let json = json!({
            "type": "m.room.member",
            "content": {
                "membership": "join"
            },
            "event_id": "$h29iv0s8:example.com",
            "origin_server_ts": 1,
            "room_id": "!n8f893n9:example.com",
            "sender": "@carl:example.com",
            "state_key": "example.com"
        });

        assert_matches!(
            from_json_value::<EventJson<StateEvent<MemberEventContent>>>(json)
                .unwrap()
                .deserialize()
                .unwrap(),
            StateEvent::<MemberEventContent> {
                content: MemberEventContent {
                    avatar_url: None,
                    displayname: None,
                    is_direct: None,
                    membership: MembershipState::Join,
                    third_party_invite: None,
                },
                event_id,
                origin_server_ts,
                room_id,
                sender,
                state_key,
                unsigned,
                prev_content: None,
            } if event_id == "$h29iv0s8:example.com"
                && origin_server_ts == UNIX_EPOCH + Duration::from_millis(1)
                && room_id == "!n8f893n9:example.com"
                && sender == "@carl:example.com"
                && state_key == "example.com"
                && unsigned.is_empty()
        );
    }

    #[test]
    fn serde_with_prev_content() {
        let json = json!({
            "type": "m.room.member",
            "content": {
                "membership": "join"
            },
            "event_id": "$h29iv0s8:example.com",
            "origin_server_ts": 1,
            "prev_content": {
                "membership": "join"
            },
            "room_id": "!n8f893n9:example.com",
            "sender": "@carl:example.com",
            "state_key": "example.com"
        });

        assert_matches!(
            from_json_value::<EventJson<StateEvent<MemberEventContent>>>(json)
                .unwrap()
                .deserialize()
                .unwrap(),
            StateEvent::<MemberEventContent> {
                content: MemberEventContent {
                    avatar_url: None,
                    displayname: None,
                    is_direct: None,
                    membership: MembershipState::Join,
                    third_party_invite: None,
                },
                event_id,
                origin_server_ts,
                room_id,
                sender,
                state_key,
                unsigned,
                prev_content: Some(MemberEventContent {
                    avatar_url: None,
                    displayname: None,
                    is_direct: None,
                    membership: MembershipState::Join,
                    third_party_invite: None,
                }),
            } if event_id == "$h29iv0s8:example.com"
                && origin_server_ts == UNIX_EPOCH + Duration::from_millis(1)
                && room_id == "!n8f893n9:example.com"
                && sender == "@carl:example.com"
                && state_key == "example.com"
                && unsigned.is_empty()
        );
    }

    #[test]
    fn serde_with_content_full() {
        let json = json!({
            "type": "m.room.member",
            "content": {
                "avatar_url": "mxc://example.org/SEsfnsuifSDFSSEF",
                "displayname": "Alice Margatroid",
                "is_direct": true,
                "membership": "invite",
                "third_party_invite": {
                    "display_name": "alice",
                    "signed": {
                        "mxid": "@alice:example.org",
                        "signatures": {
                            "magic.forest": {
                                "ed25519:3": "foobar"
                            }
                        },
                        "token": "abc123"
                    }
                }
            },
            "event_id": "$143273582443PhrSn:example.org",
            "origin_server_ts": 233,
            "room_id": "!jEsUZKDJdhlrceRyVU:example.org",
            "sender": "@alice:example.org",
            "state_key": "@alice:example.org"
        });

        assert_matches!(
            from_json_value::<EventJson<StateEvent<MemberEventContent>>>(json)
                .unwrap()
                .deserialize()
                .unwrap(),
            StateEvent::<MemberEventContent> {
                content: MemberEventContent {
                    avatar_url: Some(avatar_url),
                    displayname: Some(displayname),
                    is_direct: Some(true),
                    membership: MembershipState::Invite,
                    third_party_invite: Some(ThirdPartyInvite {
                        display_name: third_party_displayname,
                        signed: SignedContent { mxid, signatures, token },
                    }),
                },
                event_id,
                origin_server_ts,
                room_id,
                sender,
                state_key,
                unsigned,
                prev_content: None,
            } if avatar_url == "mxc://example.org/SEsfnsuifSDFSSEF"
                && displayname == "Alice Margatroid"
                && third_party_displayname == "alice"
                && mxid == "@alice:example.org"
                && signatures == btreemap! {
                    "magic.forest".to_owned() => btreemap! {
                        "ed25519:3".to_owned() => "foobar".to_owned()
                    }
                }
                && token == "abc123"
                && event_id == "$143273582443PhrSn:example.org"
                && origin_server_ts == UNIX_EPOCH + Duration::from_millis(233)
                && room_id == "!jEsUZKDJdhlrceRyVU:example.org"
                && sender == "@alice:example.org"
                && state_key == "@alice:example.org"
                && unsigned.is_empty()
        )
    }

    #[test]
    fn serde_with_prev_content_full() {
        let json = json!({
            "type": "m.room.member",
            "content": {
                "membership": "join"
            },
            "event_id": "$143273582443PhrSn:example.org",
            "origin_server_ts": 233,
            "prev_content": {
                "avatar_url": "mxc://example.org/SEsfnsuifSDFSSEF",
                "displayname": "Alice Margatroid",
                "is_direct": true,
                "membership": "invite",
                "third_party_invite": {
                    "display_name": "alice",
                    "signed": {
                        "mxid": "@alice:example.org",
                        "signatures": {
                            "magic.forest": {
                                "ed25519:3": "foobar"
                            }
                        },
                        "token": "abc123"
                    }
                }
            },
            "room_id": "!jEsUZKDJdhlrceRyVU:example.org",
            "sender": "@alice:example.org",
            "state_key": "@alice:example.org"
        });

        assert_matches!(
            from_json_value::<EventJson<StateEvent<MemberEventContent>>>(json)
                .unwrap()
                .deserialize()
                .unwrap(),
            StateEvent::<MemberEventContent> {
                content: MemberEventContent {
                    avatar_url: None,
                    displayname: None,
                    is_direct: None,
                    membership: MembershipState::Join,
                    third_party_invite: None,
                },
                event_id,
                origin_server_ts,
                room_id,
                sender,
                state_key,
                unsigned,
                prev_content: Some(MemberEventContent {
                    avatar_url: Some(avatar_url),
                    displayname: Some(displayname),
                    is_direct: Some(true),
                    membership: MembershipState::Invite,
                    third_party_invite: Some(ThirdPartyInvite {
                        display_name: third_party_displayname,
                        signed: SignedContent { mxid, signatures, token },
                    }),
                }),
            } if event_id == "$143273582443PhrSn:example.org"
                && origin_server_ts == UNIX_EPOCH + Duration::from_millis(233)
                && room_id == "!jEsUZKDJdhlrceRyVU:example.org"
                && sender == "@alice:example.org"
                && state_key == "@alice:example.org"
                && unsigned.is_empty()
                && avatar_url == "mxc://example.org/SEsfnsuifSDFSSEF"
                && displayname == "Alice Margatroid"
                && third_party_displayname == "alice"
                && mxid == "@alice:example.org"
                && signatures == btreemap! {
                    "magic.forest".to_owned() => btreemap! {
                        "ed25519:3".to_owned() => "foobar".to_owned()
                    }
                }
                && token == "abc123"
        );
    }
}
<|endoftext|>


<commit_before>#[doc = "Generate markdown from a document tree"];

import std::io;
import std::io::writer_util;

export mk_pass;

fn mk_pass(
    writer: fn~() -> io::writer
) -> pass {
    ret fn~(
        _srv: astsrv::srv,
        doc: doc::cratedoc
    ) -> doc::cratedoc {
        write_markdown(doc, writer());
        doc
    };
}

type ctxt = {
    w: io::writer,
    mutable depth: uint
};

fn write_markdown(
    doc: doc::cratedoc,
    writer: io::writer
) {
    let ctxt = {
        w: writer,
        mutable depth: 1u
    };

    write_crate(ctxt, doc);
}

fn write_header(ctxt: ctxt, title: str) {
    let hashes = str::from_chars(vec::init_elt('#', ctxt.depth));
    ctxt.w.write_line(#fmt("%s %s", hashes, title));
    ctxt.w.write_line("");
}

fn subsection(ctxt: ctxt, f: fn&()) {
    ctxt.depth += 1u;
    f();
    ctxt.depth -= 1u;
}

fn write_crate(
    ctxt: ctxt,
    doc: doc::cratedoc
) {
    write_header(ctxt, #fmt("Crate %s", doc.topmod.name));
    write_top_module(ctxt, doc.topmod);
}

fn write_top_module(
    ctxt: ctxt,
    moddoc: doc::moddoc
) {
    write_mod_contents(ctxt, moddoc);
}

fn write_mod(
    ctxt: ctxt,
    moddoc: doc::moddoc
) {
    write_header(ctxt, #fmt("Module `%s`", moddoc.name));
    write_mod_contents(ctxt, moddoc);
}

fn write_mod_contents(
    ctxt: ctxt,
    doc: doc::moddoc
) {
    write_brief(ctxt, doc.brief);
    write_desc(ctxt, doc.desc);

    for fndoc in *doc.fns {
        subsection(ctxt) {||
            write_fn(ctxt, fndoc);
        }
    }

    for moddoc in *doc.mods {
        subsection(ctxt) {||
            write_mod(ctxt, moddoc);
        }
    }
}

#[test]
fn should_write_crate_brief_description() {
    let markdown = test::render("#[doc(brief = \"this is the crate\")];");
    assert str::contains(markdown, "this is the crate");
}

#[test]
fn should_write_crate_description() {
    let markdown = test::render("#[doc = \"this is the crate\"];");
    assert str::contains(markdown, "this is the crate");
}

fn write_fn(
    ctxt: ctxt,
    doc: doc::fndoc
) {
    write_header(ctxt, #fmt("Function `%s`", doc.name));
    write_brief(ctxt, doc.brief);
    write_desc(ctxt, doc.desc);
    write_args(ctxt, doc.args);
    write_return(ctxt, doc.return);
}

fn write_brief(
    ctxt: ctxt,
    brief: option<str>
) {
    alt brief {
      some(brief) {
        ctxt.w.write_line(brief);
        ctxt.w.write_line("");
      }
      none { }
    }
}

fn write_desc(
    ctxt: ctxt,
    desc: option<str>
) {
    alt desc {
        some(desc) {
            ctxt.w.write_line(desc);
            ctxt.w.write_line("");
        }
        none { }
    }
}

fn write_args(
    ctxt: ctxt,
    args: [doc::argdoc]
) {
    if vec::is_not_empty(args) {
        ctxt.w.write_line("Arguments:");
        ctxt.w.write_line("");
        vec::iter(args) {|arg| write_arg(ctxt, arg) };
        ctxt.w.write_line("");
    }
}

fn write_arg(ctxt: ctxt, arg: doc::argdoc) {
    assert option::is_some(arg.ty);
    ctxt.w.write_str(#fmt(
        "* `%s`: `%s`",
        arg.name,
        option::get(arg.ty)
    ));
    alt arg.desc {
      some(desc) {
        ctxt.w.write_str(#fmt(" - %s", desc));
      }
      none. { }
    }
    ctxt.w.write_line("");
}

#[test]
fn should_write_argument_list() {
    let source = "fn a(b: int, c: int) { }";
    let markdown = test::render(source);
    assert str::contains(
        markdown,
        "Arguments:\n\
         \n\
         * `b`: `int`\n\
         * `c`: `int`\n\
         \n"
    );
}

#[test]
fn should_not_write_arguments_if_none() {
    let source = "fn a() { } fn b() { }";
    let markdown = test::render(source);
    assert !str::contains(markdown, "Arguments");
}

#[test]
fn should_write_argument_description() {
    let source = "#[doc(args(a = \"milk\"))] fn f(a: bool) { }";
    let markdown = test::render(source);
    assert str::contains(markdown, "`a`: `bool` - milk");
}

fn write_return(
    ctxt: ctxt,
    return: option<doc::retdoc>
) {
    alt return {
      some(doc) {
        alt doc.ty {
          some(ty) {
            ctxt.w.write_line(#fmt("Returns `%s`", ty));
            ctxt.w.write_line("");
            alt doc.desc {
              some(d) {
                ctxt.w.write_line(d);
              }
              none { }
            }
          }
          none { fail "unimplemented"; }
        }
      }
      none { }
    }
}

#[test]
fn should_write_return_type_on_new_line() {
    let markdown = test::render("fn a() -> int { }");
    assert str::contains(markdown, "\nReturns `int`");
}

#[test]
fn should_write_blank_line_between_return_type_and_next_header() {
    let markdown = test::render(
        "fn a() -> int { } \
         fn b() -> int { }"
    );
    assert str::contains(markdown, "Returns `int`\n\n##");
}

#[test]
fn should_not_write_return_type_when_there_is_none() {
    let markdown = test::render("fn a() { }");
    assert !str::contains(markdown, "Returns");
}

#[cfg(test)]
mod test {
    fn render(source: str) -> str {
        let srv = astsrv::mk_srv_from_str(source);
        let doc = extract::from_srv(srv, "");
        let doc = attr_pass::mk_pass()(srv, doc);
        let doc = tystr_pass::mk_pass()(srv, doc);
        let markdown = write_markdown_str(doc);
        #debug("markdown: %s", markdown);
        markdown
    }

    fn write_markdown_str(
        doc: doc::cratedoc
    ) -> str {
        let buffer = io::mk_mem_buffer();
        let writer = io::mem_buffer_writer(buffer);
        write_markdown(doc, writer);
        ret io::mem_buffer_str(buffer);
    }

    #[test]
    fn write_markdown_should_write_crate_header() {
        let srv = astsrv::mk_srv_from_str("");
        let doc = extract::from_srv(srv, "belch");
        let doc = attr_pass::mk_pass()(srv, doc);
        let markdown = write_markdown_str(doc);
        assert str::contains(markdown, "# Crate belch");
    }

    #[test]
    fn write_markdown_should_write_function_header() {
        let markdown = render("fn func() { }");
        assert str::contains(markdown, "## Function `func`");
    }

    #[test]
    fn write_markdown_should_write_mod_headers() {
        let markdown = render("mod moo { }");
        assert str::contains(markdown, "## Module `moo`");
    }

    #[test]
    fn should_leave_blank_line_after_header() {
        let markdown = render("mod morp { }");
        assert str::contains(markdown, "Module `morp`\n\n");
    }

    #[test]
    fn should_leave_blank_line_between_fn_header_and_brief() {
        let markdown = render("#[doc(brief = \"brief\")] fn a() { }");
        assert str::contains(markdown, "Function `a`\n\nbrief");
    }

    #[test]
    fn should_leave_blank_line_after_brief() {
        let markdown = render("#[doc(brief = \"brief\")] fn a() { }");
        assert str::contains(markdown, "brief\n\n");
    }

    #[test]
    fn should_leave_blank_line_between_brief_and_desc() {
        let markdown = render(
            "#[doc(brief = \"brief\", desc = \"desc\")] fn a() { }"
        );
        assert str::contains(markdown, "brief\n\ndesc");
    }

}<commit_msg>rustdoc: Change the header scheme so that all mods are h2, fns h3<commit_after>#[doc = "Generate markdown from a document tree"];

import std::io;
import std::io::writer_util;

export mk_pass;

fn mk_pass(
    writer: fn~() -> io::writer
) -> pass {
    ret fn~(
        _srv: astsrv::srv,
        doc: doc::cratedoc
    ) -> doc::cratedoc {
        write_markdown(doc, writer());
        doc
    };
}

type ctxt = {
    w: io::writer
};

fn write_markdown(
    doc: doc::cratedoc,
    writer: io::writer
) {
    let ctxt = {
        w: writer
    };

    write_crate(ctxt, doc);
}

tag hlvl {
    h1 = 1;
    h2 = 2;
    h3 = 3;
}

fn write_header(ctxt: ctxt, lvl: hlvl, title: str) {
    let hashes = str::from_chars(vec::init_elt('#', lvl as uint));
    ctxt.w.write_line(#fmt("%s %s", hashes, title));
    ctxt.w.write_line("");
}

fn write_crate(
    ctxt: ctxt,
    doc: doc::cratedoc
) {
    write_header(ctxt, h1, #fmt("Crate %s", doc.topmod.name));
    write_top_module(ctxt, doc.topmod);
}

fn write_top_module(
    ctxt: ctxt,
    moddoc: doc::moddoc
) {
    write_mod_contents(ctxt, moddoc);
}

fn write_mod(
    ctxt: ctxt,
    moddoc: doc::moddoc
) {
    write_header(ctxt, h2, #fmt("Module `%s`", moddoc.name));
    write_mod_contents(ctxt, moddoc);
}

fn write_mod_contents(
    ctxt: ctxt,
    doc: doc::moddoc
) {
    write_brief(ctxt, doc.brief);
    write_desc(ctxt, doc.desc);

    for fndoc in *doc.fns {
        write_fn(ctxt, fndoc);
    }

    for moddoc in *doc.mods {
        write_mod(ctxt, moddoc);
    }
}

#[test]
fn should_write_crate_brief_description() {
    let markdown = test::render("#[doc(brief = \"this is the crate\")];");
    assert str::contains(markdown, "this is the crate");
}

#[test]
fn should_write_crate_description() {
    let markdown = test::render("#[doc = \"this is the crate\"];");
    assert str::contains(markdown, "this is the crate");
}

fn write_fn(
    ctxt: ctxt,
    doc: doc::fndoc
) {
    write_header(ctxt, h3, #fmt("Function `%s`", doc.name));
    write_brief(ctxt, doc.brief);
    write_desc(ctxt, doc.desc);
    write_args(ctxt, doc.args);
    write_return(ctxt, doc.return);
}

fn write_brief(
    ctxt: ctxt,
    brief: option<str>
) {
    alt brief {
      some(brief) {
        ctxt.w.write_line(brief);
        ctxt.w.write_line("");
      }
      none { }
    }
}

fn write_desc(
    ctxt: ctxt,
    desc: option<str>
) {
    alt desc {
        some(desc) {
            ctxt.w.write_line(desc);
            ctxt.w.write_line("");
        }
        none { }
    }
}

fn write_args(
    ctxt: ctxt,
    args: [doc::argdoc]
) {
    if vec::is_not_empty(args) {
        ctxt.w.write_line("Arguments:");
        ctxt.w.write_line("");
        vec::iter(args) {|arg| write_arg(ctxt, arg) };
        ctxt.w.write_line("");
    }
}

fn write_arg(ctxt: ctxt, arg: doc::argdoc) {
    assert option::is_some(arg.ty);
    ctxt.w.write_str(#fmt(
        "* `%s`: `%s`",
        arg.name,
        option::get(arg.ty)
    ));
    alt arg.desc {
      some(desc) {
        ctxt.w.write_str(#fmt(" - %s", desc));
      }
      none. { }
    }
    ctxt.w.write_line("");
}

#[test]
fn should_write_argument_list() {
    let source = "fn a(b: int, c: int) { }";
    let markdown = test::render(source);
    assert str::contains(
        markdown,
        "Arguments:\n\
         \n\
         * `b`: `int`\n\
         * `c`: `int`\n\
         \n"
    );
}

#[test]
fn should_not_write_arguments_if_none() {
    let source = "fn a() { } fn b() { }";
    let markdown = test::render(source);
    assert !str::contains(markdown, "Arguments");
}

#[test]
fn should_write_argument_description() {
    let source = "#[doc(args(a = \"milk\"))] fn f(a: bool) { }";
    let markdown = test::render(source);
    assert str::contains(markdown, "`a`: `bool` - milk");
}

fn write_return(
    ctxt: ctxt,
    return: option<doc::retdoc>
) {
    alt return {
      some(doc) {
        alt doc.ty {
          some(ty) {
            ctxt.w.write_line(#fmt("Returns `%s`", ty));
            ctxt.w.write_line("");
            alt doc.desc {
              some(d) {
                ctxt.w.write_line(d);
              }
              none { }
            }
          }
          none { fail "unimplemented"; }
        }
      }
      none { }
    }
}

#[test]
fn should_write_return_type_on_new_line() {
    let markdown = test::render("fn a() -> int { }");
    assert str::contains(markdown, "\nReturns `int`");
}

#[test]
fn should_write_blank_line_between_return_type_and_next_header() {
    let markdown = test::render(
        "fn a() -> int { } \
         fn b() -> int { }"
    );
    assert str::contains(markdown, "Returns `int`\n\n##");
}

#[test]
fn should_not_write_return_type_when_there_is_none() {
    let markdown = test::render("fn a() { }");
    assert !str::contains(markdown, "Returns");
}

#[cfg(test)]
mod test {
    fn render(source: str) -> str {
        let srv = astsrv::mk_srv_from_str(source);
        let doc = extract::from_srv(srv, "");
        let doc = attr_pass::mk_pass()(srv, doc);
        let doc = tystr_pass::mk_pass()(srv, doc);
        let markdown = write_markdown_str(doc);
        #debug("markdown: %s", markdown);
        markdown
    }

    fn write_markdown_str(
        doc: doc::cratedoc
    ) -> str {
        let buffer = io::mk_mem_buffer();
        let writer = io::mem_buffer_writer(buffer);
        write_markdown(doc, writer);
        ret io::mem_buffer_str(buffer);
    }

    #[test]
    fn write_markdown_should_write_crate_header() {
        let srv = astsrv::mk_srv_from_str("");
        let doc = extract::from_srv(srv, "belch");
        let doc = attr_pass::mk_pass()(srv, doc);
        let markdown = write_markdown_str(doc);
        assert str::contains(markdown, "# Crate belch");
    }

    #[test]
    fn write_markdown_should_write_function_header() {
        let markdown = render("fn func() { }");
        assert str::contains(markdown, "## Function `func`");
    }

    #[test]
    fn write_markdown_should_write_mod_headers() {
        let markdown = render("mod moo { }");
        assert str::contains(markdown, "## Module `moo`");
    }

    #[test]
    fn should_leave_blank_line_after_header() {
        let markdown = render("mod morp { }");
        assert str::contains(markdown, "Module `morp`\n\n");
    }

    #[test]
    fn should_leave_blank_line_between_fn_header_and_brief() {
        let markdown = render("#[doc(brief = \"brief\")] fn a() { }");
        assert str::contains(markdown, "Function `a`\n\nbrief");
    }

    #[test]
    fn should_leave_blank_line_after_brief() {
        let markdown = render("#[doc(brief = \"brief\")] fn a() { }");
        assert str::contains(markdown, "brief\n\n");
    }

    #[test]
    fn should_leave_blank_line_between_brief_and_desc() {
        let markdown = render(
            "#[doc(brief = \"brief\", desc = \"desc\")] fn a() { }"
        );
        assert str::contains(markdown, "brief\n\ndesc");
    }

}<|endoftext|>


<commit_before>use std::convert::TryInto;
use std::process::{Child, Command};

use anyhow::{bail, Result};
use log::{debug, info, warn};
use nix::{
    sys::signal::{self, Signal},
    unistd::Pid,
};
use procfs::process::{all_processes, Process};

use crate::task_handler::ProcessAction;

pub fn compile_shell_command(command_string: &str) -> Command {
    let mut command = Command::new("sh");
    command.arg("-c").arg(command_string);

    command
}

/// Send a signal to one of Pueue's child process handles.
/// We need a special since there exists some inconsistent behavior.
///
/// In some circumstances and environments `sh -c $command` doesn't spawn a shell,
/// but rather spawns the `$command` directly.
///
/// This makes things a lot more complicated, since we need to either send signals
/// to the root process directly OR to all it's child processes.
/// This also affects the `--children` flag on all commands. We then have to either send the signal
/// to all direct children or to all of the childrens' children.
pub fn send_signal_to_child(
    child: &Child,
    action: &ProcessAction,
    send_to_children: bool,
) -> Result<bool> {
    let signal = get_signal_from_action(action);
    let pid: i32 = child.id().try_into().unwrap();
    // Check whether this process actually spawned a shell.
    let is_shell = if let Ok(is_shell) = did_process_spawn_shell(pid) {
        is_shell
    } else {
        return Ok(false);
    };

    if is_shell {
        // If it's a shell, we have to send the signal to the actual shell and to all it's children.
        // There might be multiple children, for instance, when users use the `&` operator.
        // If the `send_to_children` flag is given, the

        // Send the signal to the shell, don't propagate to it's children yet.
        send_signal_to_process(pid, action, false)?;

        // Now send the signal to the shells child processes and their respective
        // children if the user wants to do so.
        let shell_children = get_child_processes(pid);
        for shell_child in shell_children {
            send_signal_to_process(shell_child.pid(), action, send_to_children)?;
        }
    } else {
        // If it isn't a shell, send the signal directly to the process.
        // Handle children normally.
        send_signal_to_process(pid, action, send_to_children)?;
    }

    signal::kill(Pid::from_raw(pid), signal)?;
    Ok(true)
}

/// This is a helper function to safely kill a child process.
/// It's purpose is to properly kill all processes and prevent any dangling processes.
///
/// Sadly, this needs some extra handling. Check the docstring of `send_signal_to_child` for
/// additional information on why this has to be done.
pub fn kill_child(task_id: usize, child: &mut Child, kill_children: bool) -> bool {
    let pid: i32 = child.id().try_into().unwrap();

    // Check whether this process actually spawned a shell.
    let is_shell = if let Ok(is_shell) = did_process_spawn_shell(pid) {
        is_shell
    } else {
        return false;
    };

    // We have to kill the root process first, to prevent it from spawning new processes.
    // However, this prevents us from getting it's child processes afterwards.
    // That's why we have to get the list of child processes already now.
    let mut child_processes = None;
    if kill_children || is_shell {
        child_processes = Some(get_child_processes(pid));
    }

    // Kill the parent first
    match child.kill() {
        Err(_) => {
            info!("Task {} has already finished by itself", task_id);
            false
        }
        _ => {
            // Now kill all remaining children, after the parent has been killed.
            // If a shell is spawned, we have to manually send the kill signal to all children.
            // Otherwise only send a signal to all children if the `kill_children` flag is set.
            if let Some(child_processes) = child_processes {
                if is_shell {
                    for child_process in child_processes {
                        // Send the signal to each child process, show warning if this fails.
                        let process_pid = child_process.pid();
                        if let Err(error) =
                            send_signal_to_process(process_pid, &ProcessAction::Kill, kill_children)
                        {
                            warn!(
                                "Failed to send kill to pid {} with error {:?}",
                                process_pid, error
                            );
                        }
                    }
                } else if kill_children {
                    send_signal_to_processes(child_processes, &ProcessAction::Kill);
                }
            }

            true
        }
    }
}

/// Check whether a process's commandline string is actually a shell or not
fn did_process_spawn_shell(pid: i32) -> Result<bool> {
    // Get the /proc representation of the child, so we can do some checks
    let process = if let Ok(process) = Process::new(pid) {
        process
    } else {
        info!(
            "Process to kill has probably just gone away. Process {}",
            pid
        );
        bail!("Process has just gone away");
    };

    // Get the root command and, so we check whether it's actually a shell with `sh -c`.
    let mut cmdline = if let Ok(cmdline) = process.cmdline() {
        cmdline
    } else {
        info!(
            "Process to kill has probably just gone away. Process {}",
            pid
        );
        bail!("Process has just gone away");
    };

    if cmdline.len() < 3 {
        return Ok(false);
    }

    if cmdline.remove(0) != "sh" {
        return Ok(false);
    }

    if cmdline.remove(0) != "-c" {
        return Ok(false);
    }

    Ok(true)
}

/// Send a signal to a unix process.
fn send_signal_to_process(
    pid: i32,
    action: &ProcessAction,
    send_to_children: bool,
) -> Result<bool, nix::Error> {
    let signal = get_signal_from_action(action);
    debug!("Sending signal {} to {}", signal, pid);

    // Send the signal to all children, if that's what the user wants.
    if send_to_children {
        send_signal_to_children(pid, action);
    }

    signal::kill(Pid::from_raw(pid), signal)?;
    Ok(true)
}

/// A small helper that sends a signal to all children of a specific process by id.
fn send_signal_to_children(pid: i32, action: &ProcessAction) {
    send_signal_to_processes(get_child_processes(pid), action);
}

/// Send a signal to a list of processes
fn send_signal_to_processes(processes: Vec<Process>, action: &ProcessAction) {
    let signal = get_signal_from_action(action);
    for process in processes {
        // Process is no longer alive, skip this one.
        if !process.is_alive() {
            continue;
        }

        let pid = Pid::from_raw(process.pid);
        if let Err(error) = signal::kill(pid, signal) {
            warn!(
                "Failed send signal {:?} to Pid {}: {:?}",
                signal, process.pid, error
            );
        }
    }
}

/// Get all children of a specific process
fn get_child_processes(pid: i32) -> Vec<Process> {
    let all_processes = match all_processes() {
        Err(error) => {
            warn!("Failed to get full process list: {}", error);
            return Vec::new();
        }
        Ok(processes) => processes,
    };

    all_processes
        .into_iter()
        .filter(|process| process.stat.ppid == pid)
        .collect()
}

fn get_signal_from_action(action: &ProcessAction) -> Signal {
    match action {
        ProcessAction::Kill => Signal::SIGKILL,
        ProcessAction::Pause => Signal::SIGSTOP,
        ProcessAction::Resume => Signal::SIGCONT,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::thread::sleep;
    use std::time::Duration;

    /// Assert that certain process id no longer exists
    fn process_is_gone(pid: i32) -> bool {
        match Process::new(pid) {
            Ok(process) => !process.is_alive(),
            Err(_) => true,
        }
    }

    #[test]
    fn test_spawn_command() {
        let mut child = compile_shell_command("echo 'this is a test'")
            .spawn()
            .expect("Failed to spawn echo");

        let ecode = child.wait().expect("failed to wait on echo");

        assert!(ecode.success());
    }

    #[test]
    /// Ensure a `sh -c` command will be properly killed without detached processes.
    fn test_shell_command_is_killed() {
        let mut child = compile_shell_command("sleep 60 & sleep 60 && echo 'this is a test'")
            .spawn()
            .expect("Failed to spawn echo");
        let pid: i32 = child.id().try_into().unwrap();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Make sure the process indeed spawned a shell.
        assert!(did_process_spawn_shell(pid).unwrap());

        // Get all child processes, so we can make sure they no longer exist afterwards.
        let child_processes = get_child_processes(pid);
        assert_eq!(child_processes.len(), 2);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for child_process in child_processes {
            assert!(process_is_gone(child_process.stat.pid));
        }
    }

    #[test]
    /// Ensure that a `sh -c` process with a child process that has children of it's own
    /// will properly kill all processes and their children's children without detached processes.
    fn test_shell_command_children_are_killed() {
        let mut child = compile_shell_command("bash -c 'sleep 60 && sleep 60' && sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid: i32 = child.id().try_into().unwrap();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Make sure the process indeed spawned a shell.
        assert!(did_process_spawn_shell(pid).unwrap());

        // Get all child processes and all childrens children,
        // so we can make sure they no longer exist afterwards.
        let child_processes = get_child_processes(pid);
        assert_eq!(child_processes.len(), 1);
        let mut childrens_children = Vec::new();
        for child_process in &child_processes {
            childrens_children.extend(get_child_processes(child_process.stat.pid));
        }
        assert_eq!(childrens_children.len(), 1);

        // Kill the process and make sure its childen will be killed.
        assert!(kill_child(0, &mut child, true));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for child_process in child_processes {
            assert!(process_is_gone(child_process.stat.pid));
        }

        // Assert that all children's child processes have been killed.
        for child_process in childrens_children {
            assert!(process_is_gone(child_process.stat.pid));
        }
    }

    #[test]
    /// Ensure a normal command without `sh -c` will be killed.
    fn test_normal_command_is_killed() {
        let mut child = Command::new("sleep")
            .arg("60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid: i32 = child.id().try_into().unwrap();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Make sure the process did not spawn a shell.
        assert!(!did_process_spawn_shell(pid).unwrap());

        // No childprocesses exist
        let child_processes = get_child_processes(pid);
        assert_eq!(child_processes.len(), 0);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        assert!(process_is_gone(pid));
    }

    #[test]
    /// Ensure a normal command and all it's children will be
    /// properly killed without any detached processes.
    fn test_normal_command_children_are_killed() {
        let mut child = Command::new("bash")
            .arg("-c")
            .arg("sleep 60 & sleep 60 && sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid: i32 = child.id().try_into().unwrap();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Make sure the process indeed spawned a shell.
        assert!(!did_process_spawn_shell(pid).unwrap());

        // Get all child processes, so we can make sure they no longer exist afterwards.
        let child_processes = get_child_processes(pid);
        assert_eq!(child_processes.len(), 2);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, true));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for child_process in child_processes {
            assert!(process_is_gone(child_process.stat.pid));
        }
    }
}
<commit_msg>Disable cmdline check for macos. Always handle as if it's a sub shell process<commit_after>use std::convert::TryInto;
use std::process::{Child, Command};

use anyhow::{bail, Result};
use log::{debug, info, warn};
use nix::{
    sys::signal::{self, Signal},
    unistd::Pid,
};
use procfs::process::{all_processes, Process};

use crate::task_handler::ProcessAction;

pub fn compile_shell_command(command_string: &str) -> Command {
    let mut command = Command::new("sh");
    command.arg("-c").arg(command_string);

    command
}

/// Send a signal to one of Pueue's child process handles.
/// We need a special since there exists some inconsistent behavior.
pub fn send_signal_to_child(
    child: &Child,
    action: &ProcessAction,
    send_to_children: bool,
) -> Result<bool> {
    let signal = get_signal_from_action(action);
    let pid: i32 = child.id().try_into().unwrap();

    // There might be multiple children, for instance, when users use the `&` operator.
    // If the `send_to_children` flag is given, the

    // Send the signal to the shell, don't propagate to it's children yet.
    send_signal_to_process(pid, action, false)?;

    // Now send the signal to the shells child processes and their respective
    // children if the user wants to do so.
    let shell_children = get_child_processes(pid);
    for shell_child in shell_children {
        send_signal_to_process(shell_child.pid(), action, send_to_children)?;
    }

    signal::kill(Pid::from_raw(pid), signal)?;
    Ok(true)
}

/// This is a helper function to safely kill a child process.
/// It's purpose is to properly kill all processes and prevent any dangling processes.
///
/// Sadly, this needs some extra handling. Check the docstring of `send_signal_to_child` for
/// additional information on why this has to be done.
pub fn kill_child(task_id: usize, child: &mut Child, kill_children: bool) -> bool {
    let pid: i32 = child.id().try_into().unwrap();

    // We have to kill the root process first, to prevent it from spawning new processes.
    // However, this prevents us from getting it's child processes afterwards.
    // That's why we have to get the list of child processes already now.
    let child_processes = get_child_processes(pid);

    // Kill the parent first
    match child.kill() {
        Err(_) => {
            info!("Task {} has already finished by itself", task_id);
            false
        }
        _ => {
            // Now kill all remaining children, after the parent has been killed.
            for child_process in child_processes {
                // Send the signal to each child process, show warning if this fails.
                let process_pid = child_process.pid();
                if let Err(error) =
                    send_signal_to_process(process_pid, &ProcessAction::Kill, kill_children)
                {
                    warn!(
                        "Failed to send kill to pid {} with error {:?}",
                        process_pid, error
                    );
                }
            }

            true
        }
    }
}

/// Send a signal to a unix process.
fn send_signal_to_process(
    pid: i32,
    action: &ProcessAction,
    send_to_children: bool,
) -> Result<bool, nix::Error> {
    let signal = get_signal_from_action(action);
    debug!("Sending signal {} to {}", signal, pid);

    // Send the signal to all children, if that's what the user wants.
    if send_to_children {
        send_signal_to_children(pid, action);
    }

    signal::kill(Pid::from_raw(pid), signal)?;
    Ok(true)
}

/// A small helper that sends a signal to all children of a specific process by id.
fn send_signal_to_children(pid: i32, action: &ProcessAction) {
    send_signal_to_processes(get_child_processes(pid), action);
}

/// Send a signal to a list of processes
fn send_signal_to_processes(processes: Vec<Process>, action: &ProcessAction) {
    let signal = get_signal_from_action(action);
    for process in processes {
        // Process is no longer alive, skip this one.
        if !process.is_alive() {
            continue;
        }

        let pid = Pid::from_raw(process.pid);
        if let Err(error) = signal::kill(pid, signal) {
            warn!(
                "Failed send signal {:?} to Pid {}: {:?}",
                signal, process.pid, error
            );
        }
    }
}

/// Get all children of a specific process
fn get_child_processes(pid: i32) -> Vec<Process> {
    let all_processes = match all_processes() {
        Err(error) => {
            warn!("Failed to get full process list: {}", error);
            return Vec::new();
        }
        Ok(processes) => processes,
    };

    all_processes
        .into_iter()
        .filter(|process| process.stat.ppid == pid)
        .collect()
}

fn get_signal_from_action(action: &ProcessAction) -> Signal {
    match action {
        ProcessAction::Kill => Signal::SIGKILL,
        ProcessAction::Pause => Signal::SIGSTOP,
        ProcessAction::Resume => Signal::SIGCONT,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::thread::sleep;
    use std::time::Duration;

    /// Assert that certain process id no longer exists
    fn process_is_gone(pid: i32) -> bool {
        match Process::new(pid) {
            Ok(process) => !process.is_alive(),
            Err(_) => true,
        }
    }

    #[test]
    fn test_spawn_command() {
        let mut child = compile_shell_command("echo 'this is a test'")
            .spawn()
            .expect("Failed to spawn echo");

        let ecode = child.wait().expect("failed to wait on echo");

        assert!(ecode.success());
    }

    #[test]
    /// Ensure a `sh -c` command will be properly killed without detached processes.
    fn test_shell_command_is_killed() {
        let mut child = compile_shell_command("sleep 60 & sleep 60 && echo 'this is a test'")
            .spawn()
            .expect("Failed to spawn echo");
        let pid: i32 = child.id().try_into().unwrap();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Get all child processes, so we can make sure they no longer exist afterwards.
        let child_processes = get_child_processes(pid);
        assert_eq!(child_processes.len(), 2);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for child_process in child_processes {
            assert!(process_is_gone(child_process.stat.pid));
        }
    }

    #[test]
    /// Ensure that a `sh -c` process with a child process that has children of it's own
    /// will properly kill all processes and their children's children without detached processes.
    fn test_shell_command_children_are_killed() {
        let mut child = compile_shell_command("bash -c 'sleep 60 && sleep 60' && sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid: i32 = child.id().try_into().unwrap();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Get all child processes and all childrens children,
        // so we can make sure they no longer exist afterwards.
        let child_processes = get_child_processes(pid);
        assert_eq!(child_processes.len(), 1);
        let mut childrens_children = Vec::new();
        for child_process in &child_processes {
            childrens_children.extend(get_child_processes(child_process.stat.pid));
        }
        assert_eq!(childrens_children.len(), 1);

        // Kill the process and make sure its childen will be killed.
        assert!(kill_child(0, &mut child, true));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for child_process in child_processes {
            assert!(process_is_gone(child_process.stat.pid));
        }

        // Assert that all children's child processes have been killed.
        for child_process in childrens_children {
            assert!(process_is_gone(child_process.stat.pid));
        }
    }

    #[test]
    /// Ensure a normal command without `sh -c` will be killed.
    fn test_normal_command_is_killed() {
        let mut child = Command::new("sleep")
            .arg("60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid: i32 = child.id().try_into().unwrap();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // No childprocesses exist
        let child_processes = get_child_processes(pid);
        assert_eq!(child_processes.len(), 0);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        assert!(process_is_gone(pid));
    }

    #[test]
    /// Ensure a normal command and all it's children will be
    /// properly killed without any detached processes.
    fn test_normal_command_children_are_killed() {
        let mut child = Command::new("bash")
            .arg("-c")
            .arg("sleep 60 & sleep 60 && sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid: i32 = child.id().try_into().unwrap();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Get all child processes, so we can make sure they no longer exist afterwards.
        let child_processes = get_child_processes(pid);
        assert_eq!(child_processes.len(), 2);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, true));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for child_process in child_processes {
            assert!(process_is_gone(child_process.stat.pid));
        }
    }
}
<|endoftext|>