<commit_before>pub use super::{Command, CommandGroup, CommandType};

use std::collections::HashMap;
use std::default::Default;
use std::sync::Arc;
use client::Context;
use model::{Message, Permissions};

pub struct CreateCommand(pub Command);

impl CreateCommand {
    /// Adds multiple aliases.
    pub fn batch_known_as(mut self, names: Vec<&str>) -> Self {
        self.0.aliases.extend(names.into_iter().map(|n| n.to_owned()));

        self
    }

    /// Adds a ratelimit bucket.
    pub fn bucket(mut self, bucket: &str) -> Self {
        self.0.bucket = Some(bucket.to_owned());

        self
    }

    /// Adds a "check" to a command, which checks whether or not the command's
    /// function should be called.
    ///
    /// # Examples
    ///
    /// Ensure that the user who created a message, calling a "ping" command,
    /// is the owner.
    ///
    /// ```rust,no_run
    /// # use serenity::prelude::EventHandler;
    /// # struct Handler;
    /// # impl EventHandler for Handler {}
    /// use serenity::client::{Client, Context};
    /// use serenity::framework::{Command, BuiltinFramework};
    /// use serenity::model::Message;
    /// use std::env;
    /// use std::sync::Arc;
    ///
    /// let mut client = Client::new(&env::var("DISCORD_TOKEN").unwrap(), Handler);
    ///
    /// client.with_framework(BuiltinFramework::new()
    ///     .configure(|c| c.prefix("~"))
    ///     .command("ping", |c| c
    ///         .check(owner_check)
    ///         .desc("Replies to a ping with a pong")
    ///         .exec(ping)));
    ///
    /// fn ping(_context: &mut Context, message: &Message, _args: Vec<String>) -> Result<(),
    /// String> {
    ///     let _ = message.channel_id.say("Pong!");
    ///
    ///     Ok(())
    /// }
    ///
    /// fn owner_check(_context: &mut Context, message: &Message, _: &[String], _: &Arc<Command>) -> bool {
    ///     // replace with your user ID
    ///     message.author.id == 7
    /// }
    /// ```
    pub fn check<F>(mut self, check: F) -> Self
        where F: Fn(&mut Context, &Message, &[String], &Arc<Command>) -> bool + Send + Sync + 'static {
        self.0.checks.push(Box::new(check));

        self
    }

    /// Description, used by other commands.
    pub fn desc(mut self, desc: &str) -> Self {
        self.0.desc = Some(desc.to_owned());

        self
    }

    /// Whether command can be used only privately or not.
    pub fn dm_only(mut self, dm_only: bool) -> Self {
        self.0.dm_only = dm_only;

        self
    }

    /// Example arguments, used by other commands.
    pub fn example(mut self, example: &str) -> Self {
        self.0.example = Some(example.to_owned());

        self
    }

    /// A function that can be called when a command is received.
    /// You can return `Err(string)` if there's an error.
    ///
    /// See [`exec_str`] if you _only_ need to return a string on command use.
    ///
    /// [`exec_str`]: #method.exec_str
    pub fn exec<F>(mut self, func: F) -> Self
        where F: Fn(&mut Context, &Message, Vec<String>, String) -> Result<(), String>
                     + Send
                     + Sync
                     + 'static {
        self.0.exec = CommandType::Basic(Box::new(func));

        self
    }

    /// Sets a function that's called when a command is called that can access
    /// the internal HashMap of commands, used specifically for creating a help
    /// command.
    ///
    /// You can return `Err(string)` if there's an error.
    pub fn exec_help<F>(mut self, f: F) -> Self
        where F: Fn(&mut Context, &Message, HashMap<String, Arc<CommandGroup>>, &[String])
                    -> Result<(), String>
                     + 'static {
        self.0.exec = CommandType::WithCommands(Box::new(f));

        self
    }

    /// Sets a string to be sent in the channel of context on command. This can
    /// be useful for an `about`, `invite`, `ping`, etc. command.
    ///
    /// # Examples
    ///
    /// Create a command named "ping" that returns "Pong!":
    ///
    /// ```rust,ignore
    /// client.with_framework(|f| f
    ///     .command("ping", |c| c.exec_str("Pong!")));
    /// ```
    pub fn exec_str(mut self, content: &str) -> Self {
        self.0.exec = CommandType::StringResponse(content.to_owned());

        self
    }

    /// Whether command can be used only in guilds or not.
    pub fn guild_only(mut self, guild_only: bool) -> Self {
        self.0.guild_only = guild_only;

        self
    }

    /// Whether command should be displayed in help list or not, used by other commands.
    pub fn help_available(mut self, help_available: bool) -> Self {
        self.0.help_available = help_available;

        self
    }

    /// Adds an alias, allowing users to use the command under a different name.
    pub fn known_as(mut self, name: &str) -> Self {
        self.0.aliases.push(name.to_owned());

        self
    }

    /// Maximum amount of arguments that can be passed.
    pub fn max_args(mut self, max_args: i32) -> Self {
        self.0.max_args = Some(max_args);

        self
    }

    /// Minumum amount of arguments that should be passed.
    pub fn min_args(mut self, min_args: i32) -> Self {
        self.0.min_args = Some(min_args);

        self
    }

    /// Whether command can be used only privately or not.
    pub fn owners_only(mut self, owners_only: bool) -> Self {
        self.0.owners_only = owners_only;

        self
    }

    /// The permissions that a user must have in the contextual channel in order
    /// for the command to be processed.
    pub fn required_permissions(mut self, permissions: Permissions) -> Self {
        self.0.required_permissions = permissions;

        self
    }

    /// Command usage schema, used by other commands.
    pub fn usage(mut self, usage: &str) -> Self {
        self.0.usage = Some(usage.to_owned());

        self
    }

    /// Whether or not arguments should be parsed using the quotation parser.
    ///
    /// Enabling this will parse `~command "this is arg 1" "this is arg 2"` as
    /// two arguments: `this is arg 1` and `this is arg 2`.
    ///
    /// Disabling this will parse `~command "this is arg 1" "this is arg 2"` as
    /// eight arguments: `"this`, `is`, `arg`, `1"`, `"this`, `is`, `arg`, `2"`.
    ///
    /// Refer to [`utils::parse_quotes`] for information on the parser.
    ///
    /// [`utils::parse_quotes`]: ../../utils/fn.parse_quotes.html
    pub fn use_quotes(mut self, use_quotes: bool) -> Self {
        self.0.use_quotes = use_quotes;

        self
    }
}

impl Default for Command {
    fn default() -> Command {
        Command {
            aliases: Vec::new(),
            checks: Vec::default(),
            exec: CommandType::Basic(Box::new(|_, _, _, _| Ok(()))),
            desc: None,
            usage: None,
            example: None,
            use_quotes: false,
            min_args: None,
            bucket: None,
            max_args: None,
            required_permissions: Permissions::empty(),
            dm_only: false,
            guild_only: false,
            help_available: true,
            owners_only: false,
        }
    }
}
<commit_msg>Fix a failing test (#140)<commit_after>pub use super::{Command, CommandGroup, CommandType};

use std::collections::HashMap;
use std::default::Default;
use std::sync::Arc;
use client::Context;
use model::{Message, Permissions};

pub struct CreateCommand(pub Command);

impl CreateCommand {
    /// Adds multiple aliases.
    pub fn batch_known_as(mut self, names: Vec<&str>) -> Self {
        self.0.aliases.extend(names.into_iter().map(|n| n.to_owned()));

        self
    }

    /// Adds a ratelimit bucket.
    pub fn bucket(mut self, bucket: &str) -> Self {
        self.0.bucket = Some(bucket.to_owned());

        self
    }

    /// Adds a "check" to a command, which checks whether or not the command's
    /// function should be called.
    ///
    /// # Examples
    ///
    /// Ensure that the user who created a message, calling a "ping" command,
    /// is the owner.
    ///
    /// ```rust,no_run
    /// # use serenity::prelude::EventHandler;
    /// # struct Handler;
    /// # impl EventHandler for Handler {}
    /// use serenity::client::{Client, Context};
    /// use serenity::framework::{Command, BuiltinFramework};
    /// use serenity::model::Message;
    /// use std::env;
    /// use std::sync::Arc;
    ///
    /// let mut client = Client::new(&env::var("DISCORD_TOKEN").unwrap(), Handler);
    ///
    /// client.with_framework(BuiltinFramework::new()
    ///     .configure(|c| c.prefix("~"))
    ///     .command("ping", |c| c
    ///         .check(owner_check)
    ///         .desc("Replies to a ping with a pong")
    ///         .exec(ping)));
    ///
    /// fn ping(_context: &mut Context, message: &Message, _args: Vec<String>, _original_msg: String) -> Result<(),
    /// String> {
    ///     let _ = message.channel_id.say("Pong!");
    ///
    ///     Ok(())
    /// }
    ///
    /// fn owner_check(_context: &mut Context, message: &Message, _: &[String], _: &Arc<Command>) -> bool {
    ///     // replace with your user ID
    ///     message.author.id == 7
    /// }
    /// ```
    pub fn check<F>(mut self, check: F) -> Self
        where F: Fn(&mut Context, &Message, &[String], &Arc<Command>) -> bool + Send + Sync + 'static {
        self.0.checks.push(Box::new(check));

        self
    }

    /// Description, used by other commands.
    pub fn desc(mut self, desc: &str) -> Self {
        self.0.desc = Some(desc.to_owned());

        self
    }

    /// Whether command can be used only privately or not.
    pub fn dm_only(mut self, dm_only: bool) -> Self {
        self.0.dm_only = dm_only;

        self
    }

    /// Example arguments, used by other commands.
    pub fn example(mut self, example: &str) -> Self {
        self.0.example = Some(example.to_owned());

        self
    }

    /// A function that can be called when a command is received.
    /// You can return `Err(string)` if there's an error.
    ///
    /// See [`exec_str`] if you _only_ need to return a string on command use.
    ///
    /// [`exec_str`]: #method.exec_str
    pub fn exec<F>(mut self, func: F) -> Self
        where F: Fn(&mut Context, &Message, Vec<String>, String) -> Result<(), String>
                     + Send
                     + Sync
                     + 'static {
        self.0.exec = CommandType::Basic(Box::new(func));

        self
    }

    /// Sets a function that's called when a command is called that can access
    /// the internal HashMap of commands, used specifically for creating a help
    /// command.
    ///
    /// You can return `Err(string)` if there's an error.
    pub fn exec_help<F>(mut self, f: F) -> Self
        where F: Fn(&mut Context, &Message, HashMap<String, Arc<CommandGroup>>, &[String])
                    -> Result<(), String>
                     + 'static {
        self.0.exec = CommandType::WithCommands(Box::new(f));

        self
    }

    /// Sets a string to be sent in the channel of context on command. This can
    /// be useful for an `about`, `invite`, `ping`, etc. command.
    ///
    /// # Examples
    ///
    /// Create a command named "ping" that returns "Pong!":
    ///
    /// ```rust,ignore
    /// client.with_framework(|f| f
    ///     .command("ping", |c| c.exec_str("Pong!")));
    /// ```
    pub fn exec_str(mut self, content: &str) -> Self {
        self.0.exec = CommandType::StringResponse(content.to_owned());

        self
    }

    /// Whether command can be used only in guilds or not.
    pub fn guild_only(mut self, guild_only: bool) -> Self {
        self.0.guild_only = guild_only;

        self
    }

    /// Whether command should be displayed in help list or not, used by other commands.
    pub fn help_available(mut self, help_available: bool) -> Self {
        self.0.help_available = help_available;

        self
    }

    /// Adds an alias, allowing users to use the command under a different name.
    pub fn known_as(mut self, name: &str) -> Self {
        self.0.aliases.push(name.to_owned());

        self
    }

    /// Maximum amount of arguments that can be passed.
    pub fn max_args(mut self, max_args: i32) -> Self {
        self.0.max_args = Some(max_args);

        self
    }

    /// Minumum amount of arguments that should be passed.
    pub fn min_args(mut self, min_args: i32) -> Self {
        self.0.min_args = Some(min_args);

        self
    }

    /// Whether command can be used only privately or not.
    pub fn owners_only(mut self, owners_only: bool) -> Self {
        self.0.owners_only = owners_only;

        self
    }

    /// The permissions that a user must have in the contextual channel in order
    /// for the command to be processed.
    pub fn required_permissions(mut self, permissions: Permissions) -> Self {
        self.0.required_permissions = permissions;

        self
    }

    /// Command usage schema, used by other commands.
    pub fn usage(mut self, usage: &str) -> Self {
        self.0.usage = Some(usage.to_owned());

        self
    }

    /// Whether or not arguments should be parsed using the quotation parser.
    ///
    /// Enabling this will parse `~command "this is arg 1" "this is arg 2"` as
    /// two arguments: `this is arg 1` and `this is arg 2`.
    ///
    /// Disabling this will parse `~command "this is arg 1" "this is arg 2"` as
    /// eight arguments: `"this`, `is`, `arg`, `1"`, `"this`, `is`, `arg`, `2"`.
    ///
    /// Refer to [`utils::parse_quotes`] for information on the parser.
    ///
    /// [`utils::parse_quotes`]: ../../utils/fn.parse_quotes.html
    pub fn use_quotes(mut self, use_quotes: bool) -> Self {
        self.0.use_quotes = use_quotes;

        self
    }
}

impl Default for Command {
    fn default() -> Command {
        Command {
            aliases: Vec::new(),
            checks: Vec::default(),
            exec: CommandType::Basic(Box::new(|_, _, _, _| Ok(()))),
            desc: None,
            usage: None,
            example: None,
            use_quotes: false,
            min_args: None,
            bucket: None,
            max_args: None,
            required_permissions: Permissions::empty(),
            dm_only: false,
            guild_only: false,
            help_available: true,
            owners_only: false,
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2012 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 cast;
use iter::{Iterator, range};
use libc;
use ops::Drop;
use option::{Option, Some, None};
use ptr::RawPtr;
use ptr;
use str::StrSlice;
use vec::{ImmutableVector, CopyableVector};
use container::Container;

/// Resolution options for the `null_byte` condition
pub enum NullByteResolution {
    /// Truncate at the null byte
    Truncate,
    /// Use a replacement byte
    ReplaceWith(libc::c_char)
}

condition! {
    // This should be &[u8] but there's a lifetime issue (#5370).
    // NOTE: this super::NullByteResolution should be NullByteResolution
    // Change this next time the snapshot is updated.
    pub null_byte: (~[u8]) -> super::NullByteResolution;
}

/// The representation of a C String.
///
/// This structure wraps a `*libc::c_char`, and will automatically free the
/// memory it is pointing to when it goes out of scope.
pub struct CString {
    priv buf: *libc::c_char,
    priv owns_buffer_: bool,
}

impl CString {
    /// Create a C String from a pointer.
    pub unsafe fn new(buf: *libc::c_char, owns_buffer: bool) -> CString {
        CString { buf: buf, owns_buffer_: owns_buffer }
    }

    /// Unwraps the wrapped `*libc::c_char` from the `CString` wrapper.
    /// Any ownership of the buffer by the `CString` wrapper is forgotten.
    pub unsafe fn unwrap(self) -> *libc::c_char {
        let mut c_str = self;
        c_str.owns_buffer_ = false;
        c_str.buf
    }

    /// Calls a closure with a reference to the underlying `*libc::c_char`.
    ///
    /// # Failure
    ///
    /// Fails if the CString is null.
    pub fn with_ref<T>(&self, f: &fn(*libc::c_char) -> T) -> T {
        if self.buf.is_null() { fail!("CString is null!"); }
        f(self.buf)
    }

    /// Calls a closure with a mutable reference to the underlying `*libc::c_char`.
    ///
    /// # Failure
    ///
    /// Fails if the CString is null.
    pub fn with_mut_ref<T>(&mut self, f: &fn(*mut libc::c_char) -> T) -> T {
        if self.buf.is_null() { fail!("CString is null!"); }
        f(unsafe { cast::transmute_mut_unsafe(self.buf) })
    }

    /// Returns true if the CString is a null.
    pub fn is_null(&self) -> bool {
        self.buf.is_null()
    }

    /// Returns true if the CString is not null.
    pub fn is_not_null(&self) -> bool {
        self.buf.is_not_null()
    }

    /// Returns whether or not the `CString` owns the buffer.
    pub fn owns_buffer(&self) -> bool {
        self.owns_buffer_
    }

    /// Converts the CString into a `&[u8]` without copying.
    ///
    /// # Failure
    ///
    /// Fails if the CString is null.
    pub fn as_bytes<'a>(&'a self) -> &'a [u8] {
        #[fixed_stack_segment]; #[inline(never)];
        if self.buf.is_null() { fail!("CString is null!"); }
        unsafe {
            let len = libc::strlen(self.buf) as uint;
            cast::transmute((self.buf, len + 1))
        }
    }

    /// Return a CString iterator.
    pub fn iter<'a>(&'a self) -> CStringIterator<'a> {
        CStringIterator {
            ptr: self.buf,
            lifetime: unsafe { cast::transmute(self.buf) },
        }
    }
}

impl Drop for CString {
    fn drop(&self) {
        #[fixed_stack_segment]; #[inline(never)];
        if self.owns_buffer_ {
            unsafe {
                libc::free(self.buf as *libc::c_void)
            }
        }
    }
}

/// A generic trait for converting a value to a CString.
pub trait ToCStr {
    /// Copy the receiver into a CString.
    ///
    /// # Failure
    ///
    /// Raises the `null_byte` condition if the receiver has an interior null.
    fn to_c_str(&self) -> CString;

    /// Unsafe variant of `to_c_str()` that doesn't check for nulls.
    unsafe fn to_c_str_unchecked(&self) -> CString;

    /// Work with a temporary CString constructed from the receiver.
    /// The provided `*libc::c_char` will be freed immediately upon return.
    ///
    /// # Example
    ///
    /// ~~~ {.rust}
    /// let s = "PATH".with_c_str(|path| libc::getenv(path))
    /// ~~~
    ///
    /// # Failure
    ///
    /// Raises the `null_byte` condition if the receiver has an interior null.
    #[inline]
    fn with_c_str<T>(&self, f: &fn(*libc::c_char) -> T) -> T {
        self.to_c_str().with_ref(f)
    }

    /// Unsafe variant of `with_c_str()` that doesn't check for nulls.
    #[inline]
    unsafe fn with_c_str_unchecked<T>(&self, f: &fn(*libc::c_char) -> T) -> T {
        self.to_c_str_unchecked().with_ref(f)
    }
}

impl<'self> ToCStr for &'self str {
    #[inline]
    fn to_c_str(&self) -> CString {
        self.as_bytes().to_c_str()
    }

    #[inline]
    unsafe fn to_c_str_unchecked(&self) -> CString {
        self.as_bytes().to_c_str_unchecked()
    }
}

impl<'self> ToCStr for &'self [u8] {
    fn to_c_str(&self) -> CString {
        #[fixed_stack_segment]; #[inline(never)];
        let mut cs = unsafe { self.to_c_str_unchecked() };
        do cs.with_mut_ref |buf| {
            for i in range(0, self.len()) {
                unsafe {
                    let p = buf.offset(i as int);
                    if *p == 0 {
                        match null_byte::cond.raise(self.to_owned()) {
                            Truncate => break,
                            ReplaceWith(c) => *p = c
                        }
                    }
                }
            }
        }
        cs
    }

    unsafe fn to_c_str_unchecked(&self) -> CString {
        #[fixed_stack_segment]; #[inline(never)];
        do self.as_imm_buf |self_buf, self_len| {
            let buf = libc::malloc(self_len as libc::size_t + 1) as *mut u8;
            if buf.is_null() {
                fail!("failed to allocate memory!");
            }

            ptr::copy_memory(buf, self_buf, self_len);
            *ptr::mut_offset(buf, self_len as int) = 0;

            CString::new(buf as *libc::c_char, true)
        }
    }
}

/// External iterator for a CString's bytes.
///
/// Use with the `std::iterator` module.
pub struct CStringIterator<'self> {
    priv ptr: *libc::c_char,
    priv lifetime: &'self libc::c_char, // FIXME: #5922
}

impl<'self> Iterator<libc::c_char> for CStringIterator<'self> {
    fn next(&mut self) -> Option<libc::c_char> {
        let ch = unsafe { *self.ptr };
        if ch == 0 {
            None
        } else {
            self.ptr = unsafe { ptr::offset(self.ptr, 1) };
            Some(ch)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use libc;
    use ptr;
    use option::{Some, None};

    #[test]
    fn test_to_c_str() {
        do "".to_c_str().with_ref |buf| {
            unsafe {
                assert_eq!(*ptr::offset(buf, 0), 0);
            }
        }

        do "hello".to_c_str().with_ref |buf| {
            unsafe {
                assert_eq!(*ptr::offset(buf, 0), 'h' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 1), 'e' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 2), 'l' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 3), 'l' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 4), 'o' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 5), 0);
            }
        }
    }

    #[test]
    fn test_is_null() {
        let c_str = unsafe { CString::new(ptr::null(), false) };
        assert!(c_str.is_null());
        assert!(!c_str.is_not_null());
    }

    #[test]
    fn test_unwrap() {
        #[fixed_stack_segment]; #[inline(never)];

        let c_str = "hello".to_c_str();
        unsafe { libc::free(c_str.unwrap() as *libc::c_void) }
    }

    #[test]
    fn test_with_ref() {
        #[fixed_stack_segment]; #[inline(never)];

        let c_str = "hello".to_c_str();
        let len = unsafe { c_str.with_ref(|buf| libc::strlen(buf)) };
        assert!(!c_str.is_null());
        assert!(c_str.is_not_null());
        assert_eq!(len, 5);
    }

    #[test]
    #[should_fail]
    fn test_with_ref_empty_fail() {
        let c_str = unsafe { CString::new(ptr::null(), false) };
        c_str.with_ref(|_| ());
    }

    #[test]
    fn test_iterator() {
        let c_str = "".to_c_str();
        let mut iter = c_str.iter();
        assert_eq!(iter.next(), None);

        let c_str = "hello".to_c_str();
        let mut iter = c_str.iter();
        assert_eq!(iter.next(), Some('h' as libc::c_char));
        assert_eq!(iter.next(), Some('e' as libc::c_char));
        assert_eq!(iter.next(), Some('l' as libc::c_char));
        assert_eq!(iter.next(), Some('l' as libc::c_char));
        assert_eq!(iter.next(), Some('o' as libc::c_char));
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_to_c_str_fail() {
        use c_str::null_byte::cond;

        let mut error_happened = false;
        do cond.trap(|err| {
            assert_eq!(err, bytes!("he", 0, "llo").to_owned())
            error_happened = true;
            Truncate
        }).inside {
            "he\x00llo".to_c_str()
        };
        assert!(error_happened);

        do cond.trap(|_| {
            ReplaceWith('?' as libc::c_char)
        }).inside(|| "he\x00llo".to_c_str()).with_ref |buf| {
            unsafe {
                assert_eq!(*buf.offset(0), 'h' as libc::c_char);
                assert_eq!(*buf.offset(1), 'e' as libc::c_char);
                assert_eq!(*buf.offset(2), '?' as libc::c_char);
                assert_eq!(*buf.offset(3), 'l' as libc::c_char);
                assert_eq!(*buf.offset(4), 'l' as libc::c_char);
                assert_eq!(*buf.offset(5), 'o' as libc::c_char);
                assert_eq!(*buf.offset(6), 0);
            }
        }
    }

    #[test]
    fn test_to_c_str_unchecked() {
        unsafe {
            do "he\x00llo".to_c_str_unchecked().with_ref |buf| {
                assert_eq!(*buf.offset(0), 'h' as libc::c_char);
                assert_eq!(*buf.offset(1), 'e' as libc::c_char);
                assert_eq!(*buf.offset(2), 0);
                assert_eq!(*buf.offset(3), 'l' as libc::c_char);
                assert_eq!(*buf.offset(4), 'l' as libc::c_char);
                assert_eq!(*buf.offset(5), 'o' as libc::c_char);
                assert_eq!(*buf.offset(6), 0);
            }
        }
    }
}
<commit_msg>auto merge of #9220 : kballard/rust/c_str-as_str, r=thestinger<commit_after>// Copyright 2012 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 cast;
use iter::{Iterator, range};
use libc;
use ops::Drop;
use option::{Option, Some, None};
use ptr::RawPtr;
use ptr;
use str;
use str::StrSlice;
use vec::{ImmutableVector, CopyableVector};
use container::Container;

/// Resolution options for the `null_byte` condition
pub enum NullByteResolution {
    /// Truncate at the null byte
    Truncate,
    /// Use a replacement byte
    ReplaceWith(libc::c_char)
}

condition! {
    // This should be &[u8] but there's a lifetime issue (#5370).
    // NOTE: this super::NullByteResolution should be NullByteResolution
    // Change this next time the snapshot is updated.
    pub null_byte: (~[u8]) -> super::NullByteResolution;
}

/// The representation of a C String.
///
/// This structure wraps a `*libc::c_char`, and will automatically free the
/// memory it is pointing to when it goes out of scope.
pub struct CString {
    priv buf: *libc::c_char,
    priv owns_buffer_: bool,
}

impl CString {
    /// Create a C String from a pointer.
    pub unsafe fn new(buf: *libc::c_char, owns_buffer: bool) -> CString {
        CString { buf: buf, owns_buffer_: owns_buffer }
    }

    /// Unwraps the wrapped `*libc::c_char` from the `CString` wrapper.
    /// Any ownership of the buffer by the `CString` wrapper is forgotten.
    pub unsafe fn unwrap(self) -> *libc::c_char {
        let mut c_str = self;
        c_str.owns_buffer_ = false;
        c_str.buf
    }

    /// Calls a closure with a reference to the underlying `*libc::c_char`.
    ///
    /// # Failure
    ///
    /// Fails if the CString is null.
    pub fn with_ref<T>(&self, f: &fn(*libc::c_char) -> T) -> T {
        if self.buf.is_null() { fail!("CString is null!"); }
        f(self.buf)
    }

    /// Calls a closure with a mutable reference to the underlying `*libc::c_char`.
    ///
    /// # Failure
    ///
    /// Fails if the CString is null.
    pub fn with_mut_ref<T>(&mut self, f: &fn(*mut libc::c_char) -> T) -> T {
        if self.buf.is_null() { fail!("CString is null!"); }
        f(unsafe { cast::transmute_mut_unsafe(self.buf) })
    }

    /// Returns true if the CString is a null.
    pub fn is_null(&self) -> bool {
        self.buf.is_null()
    }

    /// Returns true if the CString is not null.
    pub fn is_not_null(&self) -> bool {
        self.buf.is_not_null()
    }

    /// Returns whether or not the `CString` owns the buffer.
    pub fn owns_buffer(&self) -> bool {
        self.owns_buffer_
    }

    /// Converts the CString into a `&[u8]` without copying.
    ///
    /// # Failure
    ///
    /// Fails if the CString is null.
    #[inline]
    pub fn as_bytes<'a>(&'a self) -> &'a [u8] {
        if self.buf.is_null() { fail!("CString is null!"); }
        unsafe {
            let len = ptr::position(self.buf, |c| *c == 0);
            cast::transmute((self.buf, len + 1))
        }
    }

    /// Converts the CString into a `&str` without copying.
    /// Returns None if the CString is not UTF-8 or is null.
    #[inline]
    pub fn as_str<'a>(&'a self) -> Option<&'a str> {
        if self.buf.is_null() { return None; }
        let buf = self.as_bytes();
        let buf = buf.slice_to(buf.len()-1); // chop off the trailing NUL
        str::from_utf8_slice_opt(buf)
    }

    /// Return a CString iterator.
    pub fn iter<'a>(&'a self) -> CStringIterator<'a> {
        CStringIterator {
            ptr: self.buf,
            lifetime: unsafe { cast::transmute(self.buf) },
        }
    }
}

impl Drop for CString {
    fn drop(&self) {
        #[fixed_stack_segment]; #[inline(never)];
        if self.owns_buffer_ {
            unsafe {
                libc::free(self.buf as *libc::c_void)
            }
        }
    }
}

/// A generic trait for converting a value to a CString.
pub trait ToCStr {
    /// Copy the receiver into a CString.
    ///
    /// # Failure
    ///
    /// Raises the `null_byte` condition if the receiver has an interior null.
    fn to_c_str(&self) -> CString;

    /// Unsafe variant of `to_c_str()` that doesn't check for nulls.
    unsafe fn to_c_str_unchecked(&self) -> CString;

    /// Work with a temporary CString constructed from the receiver.
    /// The provided `*libc::c_char` will be freed immediately upon return.
    ///
    /// # Example
    ///
    /// ~~~ {.rust}
    /// let s = "PATH".with_c_str(|path| libc::getenv(path))
    /// ~~~
    ///
    /// # Failure
    ///
    /// Raises the `null_byte` condition if the receiver has an interior null.
    #[inline]
    fn with_c_str<T>(&self, f: &fn(*libc::c_char) -> T) -> T {
        self.to_c_str().with_ref(f)
    }

    /// Unsafe variant of `with_c_str()` that doesn't check for nulls.
    #[inline]
    unsafe fn with_c_str_unchecked<T>(&self, f: &fn(*libc::c_char) -> T) -> T {
        self.to_c_str_unchecked().with_ref(f)
    }
}

impl<'self> ToCStr for &'self str {
    #[inline]
    fn to_c_str(&self) -> CString {
        self.as_bytes().to_c_str()
    }

    #[inline]
    unsafe fn to_c_str_unchecked(&self) -> CString {
        self.as_bytes().to_c_str_unchecked()
    }
}

impl<'self> ToCStr for &'self [u8] {
    fn to_c_str(&self) -> CString {
        #[fixed_stack_segment]; #[inline(never)];
        let mut cs = unsafe { self.to_c_str_unchecked() };
        do cs.with_mut_ref |buf| {
            for i in range(0, self.len()) {
                unsafe {
                    let p = buf.offset(i as int);
                    if *p == 0 {
                        match null_byte::cond.raise(self.to_owned()) {
                            Truncate => break,
                            ReplaceWith(c) => *p = c
                        }
                    }
                }
            }
        }
        cs
    }

    unsafe fn to_c_str_unchecked(&self) -> CString {
        #[fixed_stack_segment]; #[inline(never)];
        do self.as_imm_buf |self_buf, self_len| {
            let buf = libc::malloc(self_len as libc::size_t + 1) as *mut u8;
            if buf.is_null() {
                fail!("failed to allocate memory!");
            }

            ptr::copy_memory(buf, self_buf, self_len);
            *ptr::mut_offset(buf, self_len as int) = 0;

            CString::new(buf as *libc::c_char, true)
        }
    }
}

/// External iterator for a CString's bytes.
///
/// Use with the `std::iterator` module.
pub struct CStringIterator<'self> {
    priv ptr: *libc::c_char,
    priv lifetime: &'self libc::c_char, // FIXME: #5922
}

impl<'self> Iterator<libc::c_char> for CStringIterator<'self> {
    fn next(&mut self) -> Option<libc::c_char> {
        let ch = unsafe { *self.ptr };
        if ch == 0 {
            None
        } else {
            self.ptr = unsafe { ptr::offset(self.ptr, 1) };
            Some(ch)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use libc;
    use ptr;
    use option::{Some, None};

    #[test]
    fn test_str_to_c_str() {
        do "".to_c_str().with_ref |buf| {
            unsafe {
                assert_eq!(*ptr::offset(buf, 0), 0);
            }
        }

        do "hello".to_c_str().with_ref |buf| {
            unsafe {
                assert_eq!(*ptr::offset(buf, 0), 'h' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 1), 'e' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 2), 'l' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 3), 'l' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 4), 'o' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 5), 0);
            }
        }
    }

    #[test]
    fn test_vec_to_c_str() {
        let b: &[u8] = [];
        do b.to_c_str().with_ref |buf| {
            unsafe {
                assert_eq!(*ptr::offset(buf, 0), 0);
            }
        }

        do bytes!("hello").to_c_str().with_ref |buf| {
            unsafe {
                assert_eq!(*ptr::offset(buf, 0), 'h' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 1), 'e' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 2), 'l' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 3), 'l' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 4), 'o' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 5), 0);
            }
        }

        do bytes!("foo", 0xff).to_c_str().with_ref |buf| {
            unsafe {
                assert_eq!(*ptr::offset(buf, 0), 'f' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 1), 'o' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 2), 'o' as libc::c_char);
                assert_eq!(*ptr::offset(buf, 3), 0xff);
                assert_eq!(*ptr::offset(buf, 4), 0);
            }
        }
    }

    #[test]
    fn test_is_null() {
        let c_str = unsafe { CString::new(ptr::null(), false) };
        assert!(c_str.is_null());
        assert!(!c_str.is_not_null());
    }

    #[test]
    fn test_unwrap() {
        #[fixed_stack_segment]; #[inline(never)];

        let c_str = "hello".to_c_str();
        unsafe { libc::free(c_str.unwrap() as *libc::c_void) }
    }

    #[test]
    fn test_with_ref() {
        #[fixed_stack_segment]; #[inline(never)];

        let c_str = "hello".to_c_str();
        let len = unsafe { c_str.with_ref(|buf| libc::strlen(buf)) };
        assert!(!c_str.is_null());
        assert!(c_str.is_not_null());
        assert_eq!(len, 5);
    }

    #[test]
    #[should_fail]
    fn test_with_ref_empty_fail() {
        let c_str = unsafe { CString::new(ptr::null(), false) };
        c_str.with_ref(|_| ());
    }

    #[test]
    fn test_iterator() {
        let c_str = "".to_c_str();
        let mut iter = c_str.iter();
        assert_eq!(iter.next(), None);

        let c_str = "hello".to_c_str();
        let mut iter = c_str.iter();
        assert_eq!(iter.next(), Some('h' as libc::c_char));
        assert_eq!(iter.next(), Some('e' as libc::c_char));
        assert_eq!(iter.next(), Some('l' as libc::c_char));
        assert_eq!(iter.next(), Some('l' as libc::c_char));
        assert_eq!(iter.next(), Some('o' as libc::c_char));
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_to_c_str_fail() {
        use c_str::null_byte::cond;

        let mut error_happened = false;
        do cond.trap(|err| {
            assert_eq!(err, bytes!("he", 0, "llo").to_owned())
            error_happened = true;
            Truncate
        }).inside {
            "he\x00llo".to_c_str()
        };
        assert!(error_happened);

        do cond.trap(|_| {
            ReplaceWith('?' as libc::c_char)
        }).inside(|| "he\x00llo".to_c_str()).with_ref |buf| {
            unsafe {
                assert_eq!(*buf.offset(0), 'h' as libc::c_char);
                assert_eq!(*buf.offset(1), 'e' as libc::c_char);
                assert_eq!(*buf.offset(2), '?' as libc::c_char);
                assert_eq!(*buf.offset(3), 'l' as libc::c_char);
                assert_eq!(*buf.offset(4), 'l' as libc::c_char);
                assert_eq!(*buf.offset(5), 'o' as libc::c_char);
                assert_eq!(*buf.offset(6), 0);
            }
        }
    }

    #[test]
    fn test_to_c_str_unchecked() {
        unsafe {
            do "he\x00llo".to_c_str_unchecked().with_ref |buf| {
                assert_eq!(*buf.offset(0), 'h' as libc::c_char);
                assert_eq!(*buf.offset(1), 'e' as libc::c_char);
                assert_eq!(*buf.offset(2), 0);
                assert_eq!(*buf.offset(3), 'l' as libc::c_char);
                assert_eq!(*buf.offset(4), 'l' as libc::c_char);
                assert_eq!(*buf.offset(5), 'o' as libc::c_char);
                assert_eq!(*buf.offset(6), 0);
            }
        }
    }

    #[test]
    fn test_as_bytes() {
        let c_str = "hello".to_c_str();
        assert_eq!(c_str.as_bytes(), bytes!("hello", 0));
        let c_str = "".to_c_str();
        assert_eq!(c_str.as_bytes(), bytes!(0));
        let c_str = bytes!("foo", 0xff).to_c_str();
        assert_eq!(c_str.as_bytes(), bytes!("foo", 0xff, 0));
    }

    #[test]
    #[should_fail]
    fn test_as_bytes_fail() {
        let c_str = unsafe { CString::new(ptr::null(), false) };
        c_str.as_bytes();
    }

    #[test]
    fn test_as_str() {
        let c_str = "hello".to_c_str();
        assert_eq!(c_str.as_str(), Some("hello"));
        let c_str = "".to_c_str();
        assert_eq!(c_str.as_str(), Some(""));
        let c_str = bytes!("foo", 0xff).to_c_str();
        assert_eq!(c_str.as_str(), None);
        let c_str = unsafe { CString::new(ptr::null(), false) };
        assert_eq!(c_str.as_str(), None);
    }
}
<|endoftext|>


<commit_before>extern crate futures;
extern crate tokio_core;
extern crate tokio_signal;

use futures::stream::Stream;
use tokio_core::reactor::Core;
use tokio_signal::unix::{Signal,SIGHUP};

fn main() {
    // set up a Tokio event loop
    let mut core = Core::new().unwrap();

    // on Unix, we can listen to whatever signal we want, in this case: SIGHUP
    let sighup = Signal::new(SIGHUP, &core.handle());
    let stream = core.run(sighup).unwrap();

    println!("Waiting for SIGHUPS (Ctrl+C to quit)");
    println!("  TIP: use `pkill -sighup sighup-example` from a second terminal to send a SIGHUP \
                to all processes named 'sighup-example' (i.e. this binary)");

    // for_each is a powerful primitive provided by the Futures crate
    // it turns a Stream into a Future that completes after all stream-items
    // have been completed.
    let future = stream.for_each(|the_signal| {
        println!("*Got a signal {}* I should probably reload my config or something", the_signal);
        Ok(())
    });

    // Up until now, we haven't really DONE anything, just prepared
    // now it's time to actually schedule, and thus execute, the stream
    // on our event loop, and loop forever
    core.run(future).unwrap();
}
<commit_msg>signal: Touch up the sighup-example slightly<commit_after>extern crate futures;
extern crate tokio_core;
extern crate tokio_signal;

use futures::{Stream, Future};
use tokio_core::reactor::Core;
use tokio_signal::unix::{Signal,SIGHUP};

fn main() {
    // set up a Tokio event loop
    let mut core = Core::new().unwrap();

    // on Unix, we can listen to whatever signal we want, in this case: SIGHUP
    let stream = Signal::new(SIGHUP, &core.handle()).flatten_stream();

    println!("Waiting for SIGHUPS (Ctrl+C to quit)");
    println!("  TIP: use `pkill -sighup sighup-example` from a second terminal \
                to send a SIGHUP to all processes named 'sighup-example' \
                (i.e. this binary)");

    // for_each is a powerful primitive provided by the Futures crate
    // it turns a Stream into a Future that completes after all stream-items
    // have been completed.
    let future = stream.for_each(|the_signal| {
        println!("*Got signal {:#x}* I should probably reload my config \
                  or something", the_signal);
        Ok(())
    });

    // Up until now, we haven't really DONE anything, just prepared
    // now it's time to actually schedule, and thus execute, the stream
    // on our event loop, and loop forever
    core.run(future).unwrap();
}
<|endoftext|>


<commit_before>// Copyright 2021 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::net::Ipv4Addr;
use std::str::FromStr;
use std::sync::Arc;
use std::thread;

use anyhow::{anyhow, bail, Context};
use base::{error, validate_raw_descriptor, warn, Event, RawDescriptor};
use cros_async::{EventAsync, Executor, IoSourceExt};
use data_model::DataInit;
use devices::virtio;
use devices::virtio::net::{
    build_config, process_rx, process_tx, validate_and_configure_tap,
    virtio_features_to_tap_offload, NetError,
};
use devices::ProtectionType;
use futures::future::{AbortHandle, Abortable};
use getopts::Options;
use net_util::{MacAddress, Tap, TapT};
use once_cell::sync::OnceCell;
use sync::Mutex;
use vhost_user_devices::{CallEvent, DeviceRequestHandler, VhostUserBackend};
use virtio_sys::virtio_net;
use vm_memory::GuestMemory;
use vmm_vhost::vhost_user::message::{VhostUserProtocolFeatures, VhostUserVirtioFeatures};

thread_local! {
    static NET_EXECUTOR: OnceCell<Executor> = OnceCell::new();
}

async fn run_tx_queue(
    mut queue: virtio::Queue,
    mem: GuestMemory,
    mut tap: Tap,
    call_evt: Arc<Mutex<CallEvent>>,
    kick_evt: EventAsync,
) {
    loop {
        if let Err(e) = kick_evt.next_val().await {
            error!("Failed to read kick event for tx queue: {}", e);
            break;
        }

        process_tx(&call_evt, &mut queue, &mem, &mut tap);
    }
}

async fn run_rx_queue(
    mut queue: virtio::Queue,
    mem: GuestMemory,
    mut tap: Box<dyn IoSourceExt<Tap>>,
    call_evt: Arc<Mutex<CallEvent>>,
    kick_evt: EventAsync,
) {
    loop {
        if let Err(e) = tap.wait_readable().await {
            error!("Failed to wait for tap device to become readable: {}", e);
            break;
        }

        match process_rx(&call_evt, &mut queue, &mem, tap.as_source_mut()) {
            Ok(()) => {}
            Err(NetError::RxDescriptorsExhausted) => {
                if let Err(e) = kick_evt.next_val().await {
                    error!("Failed to read kick event for rx queue: {}", e);
                    break;
                }
            }
            Err(e) => {
                error!("Failed to process rx queue: {}", e);
                break;
            }
        }
    }
}

struct TapConfig {
    host_ip: Ipv4Addr,
    netmask: Ipv4Addr,
    mac: MacAddress,
}

impl FromStr for TapConfig {
    type Err = anyhow::Error;

    fn from_str(arg: &str) -> Result<Self, Self::Err> {
        let args: Vec<&str> = arg.split(',').collect();
        if args.len() != 3 {
            bail!("TAP config must consist of 3 parts but {}", args.len());
        }

        let host_ip: Ipv4Addr = args[0]
            .parse()
            .map_err(|e| anyhow!("invalid IP address: {}", e))?;
        let netmask: Ipv4Addr = args[1]
            .parse()
            .map_err(|e| anyhow!("invalid net mask: {}", e))?;
        let mac: MacAddress = args[2]
            .parse()
            .map_err(|e| anyhow!("invalid MAC address: {}", e))?;

        Ok(Self {
            host_ip,
            netmask,
            mac,
        })
    }
}

struct NetBackend {
    tap: Tap,
    avail_features: u64,
    acked_features: u64,
    acked_protocol_features: VhostUserProtocolFeatures,
    workers: [Option<AbortHandle>; Self::MAX_QUEUE_NUM],
    mtu: u16,
}

impl NetBackend {
    pub fn new_from_config(config: &TapConfig) -> anyhow::Result<Self> {
        // Create a tap device.
        let tap = Tap::new(true /* vnet_hdr */, false /* multi_queue */)
            .context("failed to create tap device")?;
        tap.set_ip_addr(config.host_ip)
            .context("failed to set IP address")?;
        tap.set_netmask(config.netmask)
            .context("failed to set netmask")?;
        tap.set_mac_address(config.mac)
            .context("failed to set MAC address")?;

        Self::new(tap)
    }

    pub fn new_from_tap_fd(tap_fd: RawDescriptor) -> anyhow::Result<Self> {
        let tap_fd = validate_raw_descriptor(tap_fd).context("failed to validate tap fd")?;
        // Safe because we ensure that we get a unique handle to the fd.
        let tap =
            unsafe { Tap::from_raw_descriptor(tap_fd).context("failed to create tap device")? };

        Self::new(tap)
    }

    fn new(tap: Tap) -> anyhow::Result<Self> {
        let vq_pairs = Self::max_vq_pairs();
        tap.enable().context("failed to enable tap")?;
        validate_and_configure_tap(&tap, vq_pairs as u16)
            .context("failed to validate and configure tap")?;

        // TODO(keiichiw): Support CTRL_VQ and MQ.
        // Note that MQ cannot be enabled without CTRL_VQ.
        let avail_features = virtio::base_features(ProtectionType::Unprotected)
            | 1 << virtio_net::VIRTIO_NET_F_GUEST_CSUM
            | 1 << virtio_net::VIRTIO_NET_F_CSUM
            | 1 << virtio_net::VIRTIO_NET_F_CTRL_GUEST_OFFLOADS
            | 1 << virtio_net::VIRTIO_NET_F_GUEST_TSO4
            | 1 << virtio_net::VIRTIO_NET_F_GUEST_UFO
            | 1 << virtio_net::VIRTIO_NET_F_HOST_TSO4
            | 1 << virtio_net::VIRTIO_NET_F_HOST_UFO
            | VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits();

        let mtu = tap.mtu()?;

        Ok(Self {
            tap,
            avail_features,
            acked_features: 0,
            acked_protocol_features: VhostUserProtocolFeatures::empty(),
            workers: Default::default(),
            mtu,
        })
    }

    fn max_vq_pairs() -> usize {
        Self::MAX_QUEUE_NUM / 2
    }
}

impl VhostUserBackend for NetBackend {
    // TODO(keiichiw): Support multiple queue pairs.
    const MAX_QUEUE_NUM: usize = 2; /* 1 rx and 1 tx */
    const MAX_VRING_LEN: u16 = 256;

    type Error = anyhow::Error;

    fn features(&self) -> u64 {
        self.avail_features
    }

    fn ack_features(&mut self, value: u64) -> anyhow::Result<()> {
        let unrequested_features = value & !self.avail_features;
        if unrequested_features != 0 {
            bail!("invalid features are given: {:#x}", unrequested_features);
        }

        self.acked_features |= value;

        self.tap
            .set_offload(virtio_features_to_tap_offload(self.acked_features))
            .context("failed to set tap offload to match features")?;

        Ok(())
    }

    fn acked_features(&self) -> u64 {
        self.acked_features
    }

    fn protocol_features(&self) -> VhostUserProtocolFeatures {
        // TODO(keiichiw): Support MQ.
        VhostUserProtocolFeatures::CONFIG
    }

    fn ack_protocol_features(&mut self, features: u64) -> anyhow::Result<()> {
        let features = VhostUserProtocolFeatures::from_bits(features)
            .ok_or_else(|| anyhow!("invalid protocol features are given: {:#x}", features))?;
        let supported = self.protocol_features();
        self.acked_protocol_features = features & supported;
        Ok(())
    }

    fn acked_protocol_features(&self) -> u64 {
        self.acked_protocol_features.bits()
    }

    fn read_config(&self, offset: u64, data: &mut [u8]) {
        let config_space = build_config(Self::max_vq_pairs() as u16, self.mtu);
        virtio::copy_config(data, 0, config_space.as_slice(), offset);
    }

    fn reset(&mut self) {}

    fn start_queue(
        &mut self,
        idx: usize,
        queue: virtio::Queue,
        mem: GuestMemory,
        call_evt: Arc<Mutex<CallEvent>>,
        kick_evt: Event,
    ) -> anyhow::Result<()> {
        if let Some(handle) = self.workers.get_mut(idx).and_then(Option::take) {
            warn!("Starting new queue handler without stopping old handler");
            handle.abort();
        }

        NET_EXECUTOR.with(|ex| {
            // Safe because the executor is initialized in main() below.
            let ex = ex.get().expect("Executor not initialized");

            let kick_evt = EventAsync::new(kick_evt.0, ex)
                .context("failed to create EventAsync for kick_evt")?;
            let tap = self.tap.try_clone().context("failed to clone tap device")?;
            let (handle, registration) = AbortHandle::new_pair();
            match idx {
                0 => {
                    let tap = ex
                        .async_from(tap)
                        .context("failed to create async tap device")?;

                    ex.spawn_local(Abortable::new(
                        run_rx_queue(queue, mem, tap, call_evt, kick_evt),
                        registration,
                    ))
                    .detach();
                }
                1 => {
                    ex.spawn_local(Abortable::new(
                        run_tx_queue(queue, mem, tap, call_evt, kick_evt),
                        registration,
                    ))
                    .detach();
                }
                _ => bail!("attempted to start unknown queue: {}", idx),
            }

            self.workers[idx] = Some(handle);
            Ok(())
        })
    }

    fn stop_queue(&mut self, idx: usize) {
        if let Some(handle) = self.workers.get_mut(idx).and_then(Option::take) {
            handle.abort();
        }
    }
}

fn main() -> anyhow::Result<()> {
    let mut opts = Options::new();
    opts.optflag("h", "help", "print this help menu");
    // TODO(keiichiw): Support tap-fd option.
    opts.optmulti(
        "",
        "device",
        "TAP device config. (e.g. \"/path/to/sock,10.0.2.2,255.255.255.0,12:34:56:78:9a:bc\")",
        "SOCKET_PATH,IP_ADDR,NET_MASK,MAC_ADDR",
    );
    opts.optmulti(
        "",
        "tap-fd",
        "TAP FD with a socket path",
        "SOCKET_PATH,TAP_FD",
    );

    let mut args = std::env::args();
    let program_name = args.next().expect("args is empty");
    let matches = match opts.parse(args) {
        Ok(m) => m,
        Err(e) => {
            println!("{}", e);
            println!("{}", opts.short_usage(&program_name));
            return Ok(());
        }
    };

    if matches.opt_present("h") {
        println!("{}", opts.usage(&program_name));
        return Ok(());
    }

    base::syslog::init().context("failed to initialize syslog")?;

    let device_args = matches.opt_strs("device");
    let tap_fd_args = matches.opt_strs("tap-fd");
    let num_devices = device_args.len() + tap_fd_args.len();
    if num_devices == 0 {
        bail!("no device option was passed");
    }

    let mut devices: Vec<(String, NetBackend)> = Vec::with_capacity(num_devices);

    for arg in device_args {
        let pos = match arg.find(',') {
            Some(p) => p,
            None => {
                bail!("device must take comma-separated argument");
            }
        };
        let socket = &arg[0..pos];
        let cfg = &arg[pos + 1..]
            .parse::<TapConfig>()
            .context("failed to parse tap config")?;
        let backend = NetBackend::new_from_config(&cfg).context("failed to create NetBackend")?;
        devices.push((socket.to_string(), backend));
    }

    for arg in tap_fd_args {
        let pos = match arg.find(',') {
            Some(p) => p,
            None => {
                bail!("'tap-fd' flag must take comma-separated argument");
            }
        };
        let socket = &arg[0..pos];
        let tap_fd = &arg[pos + 1..]
            .parse::<i32>()
            .context("failed to parse tap-fd")?;
        let backend =
            NetBackend::new_from_tap_fd(*tap_fd).context("failed to create NetBackend")?;
        devices.push((socket.to_string(), backend));
    }

    let mut threads = Vec::with_capacity(num_devices);
    for (socket, backend) in devices {
        let handler = DeviceRequestHandler::new(backend);
        let ex = Executor::new().context("failed to create executor")?;

        threads.push(thread::spawn(move || {
            NET_EXECUTOR.with(|thread_ex| {
                let _ = thread_ex.set(ex.clone());
            });
            if let Err(e) = ex.run_until(handler.run(&socket, &ex)) {
                error!("error occurred: {}", e);
            }
        }));
    }

    threads
        .into_iter()
        .try_for_each(thread::JoinHandle::join)
        .map_err(|e| anyhow!("failed to join threads: {:?}", e))
}
<commit_msg>vhost_user_devices: net: Add MTU feature flag<commit_after>// Copyright 2021 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::net::Ipv4Addr;
use std::str::FromStr;
use std::sync::Arc;
use std::thread;

use anyhow::{anyhow, bail, Context};
use base::{error, validate_raw_descriptor, warn, Event, RawDescriptor};
use cros_async::{EventAsync, Executor, IoSourceExt};
use data_model::DataInit;
use devices::virtio;
use devices::virtio::net::{
    build_config, process_rx, process_tx, validate_and_configure_tap,
    virtio_features_to_tap_offload, NetError,
};
use devices::ProtectionType;
use futures::future::{AbortHandle, Abortable};
use getopts::Options;
use net_util::{MacAddress, Tap, TapT};
use once_cell::sync::OnceCell;
use sync::Mutex;
use vhost_user_devices::{CallEvent, DeviceRequestHandler, VhostUserBackend};
use virtio_sys::virtio_net;
use vm_memory::GuestMemory;
use vmm_vhost::vhost_user::message::{VhostUserProtocolFeatures, VhostUserVirtioFeatures};

thread_local! {
    static NET_EXECUTOR: OnceCell<Executor> = OnceCell::new();
}

async fn run_tx_queue(
    mut queue: virtio::Queue,
    mem: GuestMemory,
    mut tap: Tap,
    call_evt: Arc<Mutex<CallEvent>>,
    kick_evt: EventAsync,
) {
    loop {
        if let Err(e) = kick_evt.next_val().await {
            error!("Failed to read kick event for tx queue: {}", e);
            break;
        }

        process_tx(&call_evt, &mut queue, &mem, &mut tap);
    }
}

async fn run_rx_queue(
    mut queue: virtio::Queue,
    mem: GuestMemory,
    mut tap: Box<dyn IoSourceExt<Tap>>,
    call_evt: Arc<Mutex<CallEvent>>,
    kick_evt: EventAsync,
) {
    loop {
        if let Err(e) = tap.wait_readable().await {
            error!("Failed to wait for tap device to become readable: {}", e);
            break;
        }

        match process_rx(&call_evt, &mut queue, &mem, tap.as_source_mut()) {
            Ok(()) => {}
            Err(NetError::RxDescriptorsExhausted) => {
                if let Err(e) = kick_evt.next_val().await {
                    error!("Failed to read kick event for rx queue: {}", e);
                    break;
                }
            }
            Err(e) => {
                error!("Failed to process rx queue: {}", e);
                break;
            }
        }
    }
}

struct TapConfig {
    host_ip: Ipv4Addr,
    netmask: Ipv4Addr,
    mac: MacAddress,
}

impl FromStr for TapConfig {
    type Err = anyhow::Error;

    fn from_str(arg: &str) -> Result<Self, Self::Err> {
        let args: Vec<&str> = arg.split(',').collect();
        if args.len() != 3 {
            bail!("TAP config must consist of 3 parts but {}", args.len());
        }

        let host_ip: Ipv4Addr = args[0]
            .parse()
            .map_err(|e| anyhow!("invalid IP address: {}", e))?;
        let netmask: Ipv4Addr = args[1]
            .parse()
            .map_err(|e| anyhow!("invalid net mask: {}", e))?;
        let mac: MacAddress = args[2]
            .parse()
            .map_err(|e| anyhow!("invalid MAC address: {}", e))?;

        Ok(Self {
            host_ip,
            netmask,
            mac,
        })
    }
}

struct NetBackend {
    tap: Tap,
    avail_features: u64,
    acked_features: u64,
    acked_protocol_features: VhostUserProtocolFeatures,
    workers: [Option<AbortHandle>; Self::MAX_QUEUE_NUM],
    mtu: u16,
}

impl NetBackend {
    pub fn new_from_config(config: &TapConfig) -> anyhow::Result<Self> {
        // Create a tap device.
        let tap = Tap::new(true /* vnet_hdr */, false /* multi_queue */)
            .context("failed to create tap device")?;
        tap.set_ip_addr(config.host_ip)
            .context("failed to set IP address")?;
        tap.set_netmask(config.netmask)
            .context("failed to set netmask")?;
        tap.set_mac_address(config.mac)
            .context("failed to set MAC address")?;

        Self::new(tap)
    }

    pub fn new_from_tap_fd(tap_fd: RawDescriptor) -> anyhow::Result<Self> {
        let tap_fd = validate_raw_descriptor(tap_fd).context("failed to validate tap fd")?;
        // Safe because we ensure that we get a unique handle to the fd.
        let tap =
            unsafe { Tap::from_raw_descriptor(tap_fd).context("failed to create tap device")? };

        Self::new(tap)
    }

    fn new(tap: Tap) -> anyhow::Result<Self> {
        let vq_pairs = Self::max_vq_pairs();
        tap.enable().context("failed to enable tap")?;
        validate_and_configure_tap(&tap, vq_pairs as u16)
            .context("failed to validate and configure tap")?;

        // TODO(keiichiw): Support CTRL_VQ and MQ.
        // Note that MQ cannot be enabled without CTRL_VQ.
        let avail_features = virtio::base_features(ProtectionType::Unprotected)
            | 1 << virtio_net::VIRTIO_NET_F_GUEST_CSUM
            | 1 << virtio_net::VIRTIO_NET_F_CSUM
            | 1 << virtio_net::VIRTIO_NET_F_CTRL_GUEST_OFFLOADS
            | 1 << virtio_net::VIRTIO_NET_F_GUEST_TSO4
            | 1 << virtio_net::VIRTIO_NET_F_GUEST_UFO
            | 1 << virtio_net::VIRTIO_NET_F_HOST_TSO4
            | 1 << virtio_net::VIRTIO_NET_F_HOST_UFO
            | 1 << virtio_net::VIRTIO_NET_F_MTU
            | VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits();

        let mtu = tap.mtu()?;

        Ok(Self {
            tap,
            avail_features,
            acked_features: 0,
            acked_protocol_features: VhostUserProtocolFeatures::empty(),
            workers: Default::default(),
            mtu,
        })
    }

    fn max_vq_pairs() -> usize {
        Self::MAX_QUEUE_NUM / 2
    }
}

impl VhostUserBackend for NetBackend {
    // TODO(keiichiw): Support multiple queue pairs.
    const MAX_QUEUE_NUM: usize = 2; /* 1 rx and 1 tx */
    const MAX_VRING_LEN: u16 = 256;

    type Error = anyhow::Error;

    fn features(&self) -> u64 {
        self.avail_features
    }

    fn ack_features(&mut self, value: u64) -> anyhow::Result<()> {
        let unrequested_features = value & !self.avail_features;
        if unrequested_features != 0 {
            bail!("invalid features are given: {:#x}", unrequested_features);
        }

        self.acked_features |= value;

        self.tap
            .set_offload(virtio_features_to_tap_offload(self.acked_features))
            .context("failed to set tap offload to match features")?;

        Ok(())
    }

    fn acked_features(&self) -> u64 {
        self.acked_features
    }

    fn protocol_features(&self) -> VhostUserProtocolFeatures {
        // TODO(keiichiw): Support MQ.
        VhostUserProtocolFeatures::CONFIG
    }

    fn ack_protocol_features(&mut self, features: u64) -> anyhow::Result<()> {
        let features = VhostUserProtocolFeatures::from_bits(features)
            .ok_or_else(|| anyhow!("invalid protocol features are given: {:#x}", features))?;
        let supported = self.protocol_features();
        self.acked_protocol_features = features & supported;
        Ok(())
    }

    fn acked_protocol_features(&self) -> u64 {
        self.acked_protocol_features.bits()
    }

    fn read_config(&self, offset: u64, data: &mut [u8]) {
        let config_space = build_config(Self::max_vq_pairs() as u16, self.mtu);
        virtio::copy_config(data, 0, config_space.as_slice(), offset);
    }

    fn reset(&mut self) {}

    fn start_queue(
        &mut self,
        idx: usize,
        queue: virtio::Queue,
        mem: GuestMemory,
        call_evt: Arc<Mutex<CallEvent>>,
        kick_evt: Event,
    ) -> anyhow::Result<()> {
        if let Some(handle) = self.workers.get_mut(idx).and_then(Option::take) {
            warn!("Starting new queue handler without stopping old handler");
            handle.abort();
        }

        NET_EXECUTOR.with(|ex| {
            // Safe because the executor is initialized in main() below.
            let ex = ex.get().expect("Executor not initialized");

            let kick_evt = EventAsync::new(kick_evt.0, ex)
                .context("failed to create EventAsync for kick_evt")?;
            let tap = self.tap.try_clone().context("failed to clone tap device")?;
            let (handle, registration) = AbortHandle::new_pair();
            match idx {
                0 => {
                    let tap = ex
                        .async_from(tap)
                        .context("failed to create async tap device")?;

                    ex.spawn_local(Abortable::new(
                        run_rx_queue(queue, mem, tap, call_evt, kick_evt),
                        registration,
                    ))
                    .detach();
                }
                1 => {
                    ex.spawn_local(Abortable::new(
                        run_tx_queue(queue, mem, tap, call_evt, kick_evt),
                        registration,
                    ))
                    .detach();
                }
                _ => bail!("attempted to start unknown queue: {}", idx),
            }

            self.workers[idx] = Some(handle);
            Ok(())
        })
    }

    fn stop_queue(&mut self, idx: usize) {
        if let Some(handle) = self.workers.get_mut(idx).and_then(Option::take) {
            handle.abort();
        }
    }
}

fn main() -> anyhow::Result<()> {
    let mut opts = Options::new();
    opts.optflag("h", "help", "print this help menu");
    // TODO(keiichiw): Support tap-fd option.
    opts.optmulti(
        "",
        "device",
        "TAP device config. (e.g. \"/path/to/sock,10.0.2.2,255.255.255.0,12:34:56:78:9a:bc\")",
        "SOCKET_PATH,IP_ADDR,NET_MASK,MAC_ADDR",
    );
    opts.optmulti(
        "",
        "tap-fd",
        "TAP FD with a socket path",
        "SOCKET_PATH,TAP_FD",
    );

    let mut args = std::env::args();
    let program_name = args.next().expect("args is empty");
    let matches = match opts.parse(args) {
        Ok(m) => m,
        Err(e) => {
            println!("{}", e);
            println!("{}", opts.short_usage(&program_name));
            return Ok(());
        }
    };

    if matches.opt_present("h") {
        println!("{}", opts.usage(&program_name));
        return Ok(());
    }

    base::syslog::init().context("failed to initialize syslog")?;

    let device_args = matches.opt_strs("device");
    let tap_fd_args = matches.opt_strs("tap-fd");
    let num_devices = device_args.len() + tap_fd_args.len();
    if num_devices == 0 {
        bail!("no device option was passed");
    }

    let mut devices: Vec<(String, NetBackend)> = Vec::with_capacity(num_devices);

    for arg in device_args {
        let pos = match arg.find(',') {
            Some(p) => p,
            None => {
                bail!("device must take comma-separated argument");
            }
        };
        let socket = &arg[0..pos];
        let cfg = &arg[pos + 1..]
            .parse::<TapConfig>()
            .context("failed to parse tap config")?;
        let backend = NetBackend::new_from_config(&cfg).context("failed to create NetBackend")?;
        devices.push((socket.to_string(), backend));
    }

    for arg in tap_fd_args {
        let pos = match arg.find(',') {
            Some(p) => p,
            None => {
                bail!("'tap-fd' flag must take comma-separated argument");
            }
        };
        let socket = &arg[0..pos];
        let tap_fd = &arg[pos + 1..]
            .parse::<i32>()
            .context("failed to parse tap-fd")?;
        let backend =
            NetBackend::new_from_tap_fd(*tap_fd).context("failed to create NetBackend")?;
        devices.push((socket.to_string(), backend));
    }

    let mut threads = Vec::with_capacity(num_devices);
    for (socket, backend) in devices {
        let handler = DeviceRequestHandler::new(backend);
        let ex = Executor::new().context("failed to create executor")?;

        threads.push(thread::spawn(move || {
            NET_EXECUTOR.with(|thread_ex| {
                let _ = thread_ex.set(ex.clone());
            });
            if let Err(e) = ex.run_until(handler.run(&socket, &ex)) {
                error!("error occurred: {}", e);
            }
        }));
    }

    threads
        .into_iter()
        .try_for_each(thread::JoinHandle::join)
        .map_err(|e| anyhow!("failed to join threads: {:?}", e))
}
<|endoftext|>


<commit_before>#[cfg(feature = "native_lib")]
use std::ffi::{CString, OsString};
#[cfg(feature = "native_lib")]
use std::os::unix::ffi::OsStringExt;
use std::sync::Arc;
use std::ops::Deref;

use Proxy;
use EventQueue;

#[cfg(feature = "native_lib")]
use wayland_sys::client::*;

/// Enum representing the possible reasons why connecting to the wayland server failed
#[derive(Debug)]
pub enum ConnectError {
    /// The library was compiled with the `dlopen` feature, and the `libwayland-client.so`
    /// library could not be found at runtime
    NoWaylandLib,
    /// Any needed library was found, but the listening socket of the server could not be
    /// found.
    ///
    /// Most of the time, this means that the program was not started from a wayland session.
    NoCompositorListening,
    InvalidName,
}

pub(crate) struct DisplayInner {
    proxy: Proxy<::protocol::wl_display::WlDisplay>,
}

impl DisplayInner {
    pub(crate) fn ptr(&self) -> *mut wl_display {
        self.proxy.c_ptr() as *mut _
    }
}

impl Drop for DisplayInner {
    fn drop(&mut self) {
        #[cfg(not(feature = "native_lib"))]
        {
            unimplemented!();
        }
        #[cfg(feature = "native_lib")]
        {
            unsafe {
                ffi_dispatch!(
                    WAYLAND_CLIENT_HANDLE,
                    wl_display_disconnect,
                    self.proxy.c_ptr() as *mut wl_display
                );
            }
        }
    }
}

pub struct Display {
    inner: Arc<DisplayInner>,
}

impl Display {
    #[cfg(feature = "native_lib")]
    unsafe fn make_display(ptr: *mut wl_display) -> Result<(Display, EventQueue), ConnectError> {
        if ptr.is_null() {
            return Err(ConnectError::NoCompositorListening);
        }

        let display = Display {
            inner: Arc::new(DisplayInner {
                proxy: Proxy::from_display(ptr),
            }),
        };

        let evq = EventQueue::new(display.inner.clone(), None);

        Ok((display, evq))
    }

    pub fn connect_to_env() -> Result<(Display, EventQueue), ConnectError> {
        #[cfg(not(feature = "native_lib"))]
        {
            unimplemented!()
        }
        #[cfg(feature = "native_lib")]
        {
            unsafe {
                let display_ptr = ffi_dispatch!(
                    WAYLAND_CLIENT_HANDLE,
                    wl_display_connect,
                    ::std::ptr::null()
                );

                Display::make_display(display_ptr)
            }
        }
    }

    pub fn connect_to_name<S: Into<OsString>>(name: S) -> Result<(Display, EventQueue), ConnectError> {
        #[cfg(not(feature = "native_lib"))]
        {
            unimplemented!()
        }
        #[cfg(feature = "native_lib")]
        {
            if !::wayland_sys::client::is_lib_available() {
                return Err(ConnectError::NoWaylandLib);
            }

            let name = CString::new(name.into().into_vec()).map_err(|_| ConnectError::InvalidName)?;

            unsafe {
                let display_ptr = ffi_dispatch!(WAYLAND_CLIENT_HANDLE, wl_display_connect, name.as_ptr());

                Display::make_display(display_ptr)
            }
        }
    }
}

impl Deref for Display {
    type Target = Proxy<::protocol::wl_display::WlDisplay>;
    fn deref(&self) -> &Proxy<::protocol::wl_display::WlDisplay> {
        &self.inner.proxy
    }
}
<commit_msg>client: create other event queues<commit_after>#[cfg(feature = "native_lib")]
use std::ffi::{CString, OsString};
#[cfg(feature = "native_lib")]
use std::os::unix::ffi::OsStringExt;
use std::sync::Arc;
use std::ops::Deref;

use Proxy;
use EventQueue;

#[cfg(feature = "native_lib")]
use wayland_sys::client::*;

/// Enum representing the possible reasons why connecting to the wayland server failed
#[derive(Debug)]
pub enum ConnectError {
    /// The library was compiled with the `dlopen` feature, and the `libwayland-client.so`
    /// library could not be found at runtime
    NoWaylandLib,
    /// Any needed library was found, but the listening socket of the server could not be
    /// found.
    ///
    /// Most of the time, this means that the program was not started from a wayland session.
    NoCompositorListening,
    InvalidName,
}

pub(crate) struct DisplayInner {
    proxy: Proxy<::protocol::wl_display::WlDisplay>,
}

impl DisplayInner {
    #[cfg(feature = "native_lib")]
    pub(crate) fn ptr(&self) -> *mut wl_display {
        self.proxy.c_ptr() as *mut _
    }
}

impl Drop for DisplayInner {
    fn drop(&mut self) {
        #[cfg(not(feature = "native_lib"))]
        {
            unimplemented!();
        }
        #[cfg(feature = "native_lib")]
        {
            unsafe {
                ffi_dispatch!(
                    WAYLAND_CLIENT_HANDLE,
                    wl_display_disconnect,
                    self.proxy.c_ptr() as *mut wl_display
                );
            }
        }
    }
}

pub struct Display {
    inner: Arc<DisplayInner>,
}

impl Display {
    #[cfg(feature = "native_lib")]
    unsafe fn make_display(ptr: *mut wl_display) -> Result<(Display, EventQueue), ConnectError> {
        if ptr.is_null() {
            return Err(ConnectError::NoCompositorListening);
        }

        let display = Display {
            inner: Arc::new(DisplayInner {
                proxy: Proxy::from_display(ptr),
            }),
        };

        let evq = EventQueue::new(display.inner.clone(), None);

        Ok((display, evq))
    }

    pub fn connect_to_env() -> Result<(Display, EventQueue), ConnectError> {
        #[cfg(not(feature = "native_lib"))]
        {
            unimplemented!()
        }
        #[cfg(feature = "native_lib")]
        {
            unsafe {
                let display_ptr = ffi_dispatch!(
                    WAYLAND_CLIENT_HANDLE,
                    wl_display_connect,
                    ::std::ptr::null()
                );

                Display::make_display(display_ptr)
            }
        }
    }

    pub fn connect_to_name<S: Into<OsString>>(name: S) -> Result<(Display, EventQueue), ConnectError> {
        #[cfg(not(feature = "native_lib"))]
        {
            unimplemented!()
        }
        #[cfg(feature = "native_lib")]
        {
            if !::wayland_sys::client::is_lib_available() {
                return Err(ConnectError::NoWaylandLib);
            }

            let name = CString::new(name.into().into_vec()).map_err(|_| ConnectError::InvalidName)?;

            unsafe {
                let display_ptr = ffi_dispatch!(WAYLAND_CLIENT_HANDLE, wl_display_connect, name.as_ptr());

                Display::make_display(display_ptr)
            }
        }
    }

    pub fn create_event_queue(&self) -> EventQueue {
        #[cfg(not(feature = "native_lib"))]
        {
            unimplemented!()
        }
        #[cfg(feature = "native_lib")]
        unsafe {
            let ptr = ffi_dispatch!(
                WAYLAND_CLIENT_HANDLE,
                wl_display_create_queue,
                self.inner.ptr()
            );

            EventQueue::new(self.inner.clone(), Some(ptr))
        }
    }
}

impl Deref for Display {
    type Target = Proxy<::protocol::wl_display::WlDisplay>;
    fn deref(&self) -> &Proxy<::protocol::wl_display::WlDisplay> {
        &self.inner.proxy
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::ffi::{CStr, CString, OsStr, OsString};
use std::io::{Error as IoError, ErrorKind, Result as IoResult};
use std::os::raw::c_void;
use std::os::unix::ffi::{OsStrExt, OsStringExt};
use std::os::unix::io::{IntoRawFd, RawFd};
use std::path::PathBuf;
use std::ptr;
use std::sync::Arc;

use wayland_commons::{Implementation, Interface};

use {Client, EventLoop, Global, LoopToken, NewResource};
use globals::global_bind;

#[cfg(feature = "native_lib")]
use wayland_sys::server::*;

pub(crate) struct DisplayInner {
    #[cfg(feature = "native_lib")]
    pub(crate) ptr: *mut wl_display,
}

impl Drop for DisplayInner {
    fn drop(&mut self) {
        #[cfg(not(feature = "native_lib"))]
        {
            unimplemented!();
        }
        #[cfg(feature = "native_lib")]
        {
            unsafe {
                ffi_dispatch!(WAYLAND_SERVER_HANDLE, wl_display_destroy, self.ptr);
            }
        }
    }
}

/// The wayland display
///
/// This is the core of your wayland server, this object must
/// be kept alive as long as your server is running. It allows
/// you to manage listening sockets and clients.
pub struct Display {
    inner: Arc<DisplayInner>,
}

impl Display {
    #[cfg(feature = "native_lib")]
    /// Create a new display
    ///
    /// This method provides you a `Display` as well as the main `EventLoop`
    /// which will host your clients' objects.
    ///
    /// Note that at this point, your server is not yet ready to receive connections,
    /// your need to add listening sockets using the `add_socket*` methods.
    pub fn new() -> (Display, EventLoop) {
        let ptr = unsafe { ffi_dispatch!(WAYLAND_SERVER_HANDLE, wl_display_create,) };

        let display = Display {
            inner: Arc::new(DisplayInner { ptr: ptr }),
        };

        // setup the client_created listener
        unsafe {
            let listener = signal::rust_listener_create(client_created);
            ffi_dispatch!(
                WAYLAND_SERVER_HANDLE,
                wl_display_add_client_created_listener,
                ptr,
                listener
            );
        }

        let evq_ptr = unsafe { ffi_dispatch!(WAYLAND_SERVER_HANDLE, wl_display_get_event_loop, ptr) };

        let evq = unsafe { EventLoop::display_new(display.inner.clone(), evq_ptr) };

        (display, evq)
    }

    /// Create a new global object
    ///
    /// This object will be advertized to all clients, and they will
    /// be able to instanciate it from their registries.
    ///
    /// Your implementation will be called whenever a client instanciates
    /// this global.
    ///
    /// The version specified is the **highest supported version**, you must
    /// be able to handle clients that choose to instanciate this global with
    /// a lower version number.
    pub fn create_global<I: Interface, Impl>(
        &mut self,
        token: &LoopToken,
        version: u32,
        implementation: Impl,
    ) -> Global<I>
    where
        Impl: Implementation<NewResource<I>, u32> + 'static,
    {
        let token_inner = token
            .inner
            .inner
            .as_ref()
            .expect("Display::create_global requires the token associated with the display event loop.");
        assert!(
            Arc::ptr_eq(&self.inner, token_inner),
            "Display::create_global requires the token associated with the display event loop."
        );

        let data = Box::new(Box::new(implementation) as Box<Implementation<NewResource<I>, u32>>);

        unsafe {
            let ptr = ffi_dispatch!(
                WAYLAND_SERVER_HANDLE,
                wl_global_create,
                self.inner.ptr,
                I::c_interface(),
                version as i32,
                &*data as *const Box<_> as *mut _,
                global_bind::<I>
            );

            Global::create(ptr, data)
        }
    }

    /// Flush events to the clients
    ///
    /// Will send as many pending events as possible to the respective sockets of the clients.
    /// Will not block, but might not send everything if the socket buffer fills up.
    pub fn flush_clients(&self) {
        #[cfg(not(feature = "native_lib"))]
        {
            unimplemented!()
        }
        #[cfg(feature = "native_lib")]
        {
            unsafe {
                ffi_dispatch!(
                    WAYLAND_SERVER_HANDLE,
                    wl_display_flush_clients,
                    self.inner.ptr
                )
            };
        }
    }
}

#[cfg(feature = "native_lib")]
unsafe extern "C" fn client_created(_listener: *mut wl_listener, data: *mut c_void) {
    // init the client
    let _client = Client::from_ptr(data as *mut wl_client);
}

impl Display {
    /// Add a listening socket to this display
    ///
    /// Wayland clients will be able to connect to your compositor from this socket.
    ///
    /// Socket will be created in the directory specified by the environment variable
    /// `XDG_RUNTIME_DIR`.
    ///
    /// If a name is provided, it is used. Otherwise, if `WAYLAND_DISPLAY` environment
    /// variable is set, its contents are used as socket name. Otherwise, `wayland-0` is used.
    ///
    /// Errors if `name` contains an interior null, or if `XDG_RUNTIME_DIR` is not set,
    /// or if specified could not be bound (either it is already used or the compositor
    /// does not have the rights to create it).
    pub fn add_socket<S>(&mut self, name: Option<S>) -> IoResult<()>
    where
        S: AsRef<OsStr>,
    {
        #[cfg(not(feature = "native_lib"))]
        {
            unimplemented!();
        }
        #[cfg(feature = "native_lib")]
        {
            let cname = match name.as_ref().map(|s| CString::new(s.as_ref().as_bytes())) {
                Some(Ok(n)) => Some(n),
                Some(Err(_)) => {
                    return Err(IoError::new(
                        ErrorKind::InvalidInput,
                        "nulls are not allowed in socket name",
                    ))
                }
                None => None,
            };
            let ret = unsafe {
                ffi_dispatch!(
                    WAYLAND_SERVER_HANDLE,
                    wl_display_add_socket,
                    self.inner.ptr,
                    cname.as_ref().map(|s| s.as_ptr()).unwrap_or(ptr::null())
                )
            };
            if ret == -1 {
                // lets try to be helpfull
                let mut socket_name = get_runtime_dir()?;
                match name {
                    Some(s) => socket_name.push(s.as_ref()),
                    None => socket_name.push("wayland-0"),
                }
                Err(IoError::new(
                    ErrorKind::PermissionDenied,
                    format!("could not bind socket {}", socket_name.to_string_lossy()),
                ))
            } else {
                Ok(())
            }
        }
    }

    /// Add an automatically named listening socket to this display
    ///
    /// Wayland clients will be able to connect to your compositor from this socket.
    ///
    /// Socket will be created in the directory specified by the environment variable
    /// `XDG_RUNTIME_DIR`. The directory is scanned for any name in the form `wayland-$d` with
    /// `0 <= $d < 32` and the first available one is used.
    ///
    /// Errors if `XDG_RUNTIME_DIR` is not set, or all 32 names are already in use.
    pub fn add_socket_auto(&mut self) -> IoResult<OsString> {
        #[cfg(not(feature = "native_lib"))]
        {
            unimplemented!()
        }
        #[cfg(feature = "native_lib")]
        {
            let ret = unsafe {
                ffi_dispatch!(
                    WAYLAND_SERVER_HANDLE,
                    wl_display_add_socket_auto,
                    self.inner.ptr
                )
            };
            if ret.is_null() {
                // try to be helpfull
                let socket_name = get_runtime_dir()?;
                Err(IoError::new(
                    ErrorKind::Other,
                    format!(
                        "no available wayland-* name in {}",
                        socket_name.to_string_lossy()
                    ),
                ))
            } else {
                let sockname = unsafe { CStr::from_ptr(ret) };
                Ok(<OsString as OsStringExt>::from_vec(
                    sockname.to_bytes().into(),
                ))
            }
        }
    }

    /// Add existing listening socket to this display
    ///
    /// Wayland clients will be able to connect to your compositor from this socket.
    ///
    /// The existing socket fd must already be created, opened, and locked.
    /// The fd must be properly set to CLOEXEC and bound to a socket file
    /// with both bind() and listen() already called. An error is returned
    /// otherwise.
    pub fn add_socket_from<T>(&mut self, socket: T) -> IoResult<()>
    where
        T: IntoRawFd,
    {
        unsafe { self.add_socket_fd(socket.into_raw_fd()) }
    }

    /// Add existing listening socket to this display from a raw file descriptor
    ///
    /// Wayland clients will be able to connect to your compositor from this socket.
    ///
    /// The library takes ownership of the provided socket if this method returns
    /// successfully.
    ///
    /// The existing socket fd must already be created, opened, and locked.
    /// The fd must be properly set to CLOEXEC and bound to a socket file
    /// with both bind() and listen() already called. An error is returned
    /// otherwise.
    pub unsafe fn add_socket_fd(&mut self, fd: RawFd) -> IoResult<()> {
        #[cfg(not(feature = "native_lib"))]
        {
            unimplemented!()
        }
        #[cfg(feature = "native_lib")]
        {
            let ret = ffi_dispatch!(
                WAYLAND_SERVER_HANDLE,
                wl_display_add_socket_fd,
                self.inner.ptr,
                fd
            );
            if ret == 0 {
                Ok(())
            } else {
                Err(IoError::new(ErrorKind::InvalidInput, "invalid socket fd"))
            }
        }
    }
}

fn get_runtime_dir() -> IoResult<PathBuf> {
    match env::var_os("XDG_RUNTIME_DIR") {
        Some(s) => Ok(s.into()),
        None => Err(IoError::new(
            ErrorKind::NotFound,
            "XDG_RUNTIME_DIR env variable is not set",
        )),
    }
}
<commit_msg>server: Display:c_ptr()<commit_after>use std::env;
use std::ffi::{CStr, CString, OsStr, OsString};
use std::io::{Error as IoError, ErrorKind, Result as IoResult};
use std::os::raw::c_void;
use std::os::unix::ffi::{OsStrExt, OsStringExt};
use std::os::unix::io::{IntoRawFd, RawFd};
use std::path::PathBuf;
use std::ptr;
use std::sync::Arc;

use wayland_commons::{Implementation, Interface};

use {Client, EventLoop, Global, LoopToken, NewResource};
use globals::global_bind;

#[cfg(feature = "native_lib")]
use wayland_sys::server::*;

pub(crate) struct DisplayInner {
    #[cfg(feature = "native_lib")]
    pub(crate) ptr: *mut wl_display,
}

impl Drop for DisplayInner {
    fn drop(&mut self) {
        #[cfg(not(feature = "native_lib"))]
        {
            unimplemented!();
        }
        #[cfg(feature = "native_lib")]
        {
            unsafe {
                ffi_dispatch!(WAYLAND_SERVER_HANDLE, wl_display_destroy, self.ptr);
            }
        }
    }
}

/// The wayland display
///
/// This is the core of your wayland server, this object must
/// be kept alive as long as your server is running. It allows
/// you to manage listening sockets and clients.
pub struct Display {
    inner: Arc<DisplayInner>,
}

impl Display {
    #[cfg(feature = "native_lib")]
    /// Create a new display
    ///
    /// This method provides you a `Display` as well as the main `EventLoop`
    /// which will host your clients' objects.
    ///
    /// Note that at this point, your server is not yet ready to receive connections,
    /// your need to add listening sockets using the `add_socket*` methods.
    pub fn new() -> (Display, EventLoop) {
        let ptr = unsafe { ffi_dispatch!(WAYLAND_SERVER_HANDLE, wl_display_create,) };

        let display = Display {
            inner: Arc::new(DisplayInner { ptr: ptr }),
        };

        // setup the client_created listener
        unsafe {
            let listener = signal::rust_listener_create(client_created);
            ffi_dispatch!(
                WAYLAND_SERVER_HANDLE,
                wl_display_add_client_created_listener,
                ptr,
                listener
            );
        }

        let evq_ptr = unsafe { ffi_dispatch!(WAYLAND_SERVER_HANDLE, wl_display_get_event_loop, ptr) };

        let evq = unsafe { EventLoop::display_new(display.inner.clone(), evq_ptr) };

        (display, evq)
    }

    /// Create a new global object
    ///
    /// This object will be advertized to all clients, and they will
    /// be able to instanciate it from their registries.
    ///
    /// Your implementation will be called whenever a client instanciates
    /// this global.
    ///
    /// The version specified is the **highest supported version**, you must
    /// be able to handle clients that choose to instanciate this global with
    /// a lower version number.
    pub fn create_global<I: Interface, Impl>(
        &mut self,
        token: &LoopToken,
        version: u32,
        implementation: Impl,
    ) -> Global<I>
    where
        Impl: Implementation<NewResource<I>, u32> + 'static,
    {
        let token_inner = token
            .inner
            .inner
            .as_ref()
            .expect("Display::create_global requires the token associated with the display event loop.");
        assert!(
            Arc::ptr_eq(&self.inner, token_inner),
            "Display::create_global requires the token associated with the display event loop."
        );

        let data = Box::new(Box::new(implementation) as Box<Implementation<NewResource<I>, u32>>);

        unsafe {
            let ptr = ffi_dispatch!(
                WAYLAND_SERVER_HANDLE,
                wl_global_create,
                self.inner.ptr,
                I::c_interface(),
                version as i32,
                &*data as *const Box<_> as *mut _,
                global_bind::<I>
            );

            Global::create(ptr, data)
        }
    }

    /// Flush events to the clients
    ///
    /// Will send as many pending events as possible to the respective sockets of the clients.
    /// Will not block, but might not send everything if the socket buffer fills up.
    pub fn flush_clients(&self) {
        #[cfg(not(feature = "native_lib"))]
        {
            unimplemented!()
        }
        #[cfg(feature = "native_lib")]
        {
            unsafe {
                ffi_dispatch!(
                    WAYLAND_SERVER_HANDLE,
                    wl_display_flush_clients,
                    self.inner.ptr
                )
            };
        }
    }
}

#[cfg(feature = "native_lib")]
unsafe extern "C" fn client_created(_listener: *mut wl_listener, data: *mut c_void) {
    // init the client
    let _client = Client::from_ptr(data as *mut wl_client);
}

impl Display {
    /// Add a listening socket to this display
    ///
    /// Wayland clients will be able to connect to your compositor from this socket.
    ///
    /// Socket will be created in the directory specified by the environment variable
    /// `XDG_RUNTIME_DIR`.
    ///
    /// If a name is provided, it is used. Otherwise, if `WAYLAND_DISPLAY` environment
    /// variable is set, its contents are used as socket name. Otherwise, `wayland-0` is used.
    ///
    /// Errors if `name` contains an interior null, or if `XDG_RUNTIME_DIR` is not set,
    /// or if specified could not be bound (either it is already used or the compositor
    /// does not have the rights to create it).
    pub fn add_socket<S>(&mut self, name: Option<S>) -> IoResult<()>
    where
        S: AsRef<OsStr>,
    {
        #[cfg(not(feature = "native_lib"))]
        {
            unimplemented!();
        }
        #[cfg(feature = "native_lib")]
        {
            let cname = match name.as_ref().map(|s| CString::new(s.as_ref().as_bytes())) {
                Some(Ok(n)) => Some(n),
                Some(Err(_)) => {
                    return Err(IoError::new(
                        ErrorKind::InvalidInput,
                        "nulls are not allowed in socket name",
                    ))
                }
                None => None,
            };
            let ret = unsafe {
                ffi_dispatch!(
                    WAYLAND_SERVER_HANDLE,
                    wl_display_add_socket,
                    self.inner.ptr,
                    cname.as_ref().map(|s| s.as_ptr()).unwrap_or(ptr::null())
                )
            };
            if ret == -1 {
                // lets try to be helpfull
                let mut socket_name = get_runtime_dir()?;
                match name {
                    Some(s) => socket_name.push(s.as_ref()),
                    None => socket_name.push("wayland-0"),
                }
                Err(IoError::new(
                    ErrorKind::PermissionDenied,
                    format!("could not bind socket {}", socket_name.to_string_lossy()),
                ))
            } else {
                Ok(())
            }
        }
    }

    /// Add an automatically named listening socket to this display
    ///
    /// Wayland clients will be able to connect to your compositor from this socket.
    ///
    /// Socket will be created in the directory specified by the environment variable
    /// `XDG_RUNTIME_DIR`. The directory is scanned for any name in the form `wayland-$d` with
    /// `0 <= $d < 32` and the first available one is used.
    ///
    /// Errors if `XDG_RUNTIME_DIR` is not set, or all 32 names are already in use.
    pub fn add_socket_auto(&mut self) -> IoResult<OsString> {
        #[cfg(not(feature = "native_lib"))]
        {
            unimplemented!()
        }
        #[cfg(feature = "native_lib")]
        {
            let ret = unsafe {
                ffi_dispatch!(
                    WAYLAND_SERVER_HANDLE,
                    wl_display_add_socket_auto,
                    self.inner.ptr
                )
            };
            if ret.is_null() {
                // try to be helpfull
                let socket_name = get_runtime_dir()?;
                Err(IoError::new(
                    ErrorKind::Other,
                    format!(
                        "no available wayland-* name in {}",
                        socket_name.to_string_lossy()
                    ),
                ))
            } else {
                let sockname = unsafe { CStr::from_ptr(ret) };
                Ok(<OsString as OsStringExt>::from_vec(
                    sockname.to_bytes().into(),
                ))
            }
        }
    }

    /// Add existing listening socket to this display
    ///
    /// Wayland clients will be able to connect to your compositor from this socket.
    ///
    /// The existing socket fd must already be created, opened, and locked.
    /// The fd must be properly set to CLOEXEC and bound to a socket file
    /// with both bind() and listen() already called. An error is returned
    /// otherwise.
    pub fn add_socket_from<T>(&mut self, socket: T) -> IoResult<()>
    where
        T: IntoRawFd,
    {
        unsafe { self.add_socket_fd(socket.into_raw_fd()) }
    }

    /// Add existing listening socket to this display from a raw file descriptor
    ///
    /// Wayland clients will be able to connect to your compositor from this socket.
    ///
    /// The library takes ownership of the provided socket if this method returns
    /// successfully.
    ///
    /// The existing socket fd must already be created, opened, and locked.
    /// The fd must be properly set to CLOEXEC and bound to a socket file
    /// with both bind() and listen() already called. An error is returned
    /// otherwise.
    pub unsafe fn add_socket_fd(&mut self, fd: RawFd) -> IoResult<()> {
        #[cfg(not(feature = "native_lib"))]
        {
            unimplemented!()
        }
        #[cfg(feature = "native_lib")]
        {
            let ret = ffi_dispatch!(
                WAYLAND_SERVER_HANDLE,
                wl_display_add_socket_fd,
                self.inner.ptr,
                fd
            );
            if ret == 0 {
                Ok(())
            } else {
                Err(IoError::new(ErrorKind::InvalidInput, "invalid socket fd"))
            }
        }
    }
}

#[cfg(feature = "native_lib")]
impl Display {
    /// Retrieve a pointer from the C lib to this `wl_display`
    pub fn c_ptr(&self) -> *mut wl_display {
        self.inner.ptr
    }
}

fn get_runtime_dir() -> IoResult<PathBuf> {
    match env::var_os("XDG_RUNTIME_DIR") {
        Some(s) => Ok(s.into()),
        None => Err(IoError::new(
            ErrorKind::NotFound,
            "XDG_RUNTIME_DIR env variable is not set",
        )),
    }
}
<|endoftext|>