<commit_before>#![feature(libc)]

extern crate libc;
use std::ffi::CString;
use std::slice;
use std::ptr;

mod ffi;

#[derive(Debug, Copy, Clone)]
pub struct Env {
    env: ffi::Voidptr
}

#[derive(Debug, Copy, Clone)]
pub struct Db {
    db: ffi::Voidptr
}

// XXX: Use Arc
unsafe impl Sync for Db {}
unsafe impl Send for Db {}
unsafe impl Sync for Env {}
unsafe impl Send for Env {}

use std::marker::PhantomData;

#[derive(Debug)]
pub struct DbObject<'a> {
    o: ffi::Voidptr,
    phantom: PhantomData<&'a ()>
}

impl<'a> DbObject<'a> {
    // We name it attr(), to not mix it up with the set() method.
    pub fn attr<'b>(&'b mut self, path: &str, data: &'b[u8]) {
        let path = CString::new(path).unwrap();
        unsafe {ffi::setstring(self.o, path.as_bytes(), data); }
    }

    // Shortcut for attr("key", ...)
    pub fn key<'b>(&'b mut self, data: &'b[u8]) {
        unsafe {ffi::setstring(self.o, "key\0".as_bytes(), data); }
    }

    // Shortcut for attr("value", ...)
    pub fn value<'b>(&'b mut self, data: &'b[u8]) {
        unsafe {ffi::setstring(self.o, "value\0".as_bytes(), data); }
    }
}

impl<'a> Drop for DbObject<'a> {
    fn drop(&mut self) {
        if !self.o.is_null() {
            unsafe {ffi::sp_destroy(self.o)};
        }
    }
}

#[derive(Debug)]
pub struct Transaction<'a> {
    tx: ffi::Voidptr,
    phantom: PhantomData<&'a ()>
}

impl<'a> Drop for Transaction<'a> {
    // Unless commit() is called, this will roll back the transaction
    fn drop(&mut self) {
        if !self.tx.is_null() {
            unsafe {ffi::sp_destroy(self.tx);}
        }
    }
}

impl<'a> Transaction<'a> {
    fn commit(mut self) -> i32 {
        let rc = unsafe{ffi::sp_commit(self.tx)};
        self.tx = ptr::null_mut();
        rc as i32
    }

    // Consumes the DbObject
    pub fn put<'b>(&mut self, obj: DbObject<'b>) {
        let mut obj = obj;
        unsafe {ffi::sp_set(self.tx, obj.o)}; // XXX: Check error code
        obj.o = ptr::null_mut(); // sp_set drops it
    }
}

impl Env {
    pub fn new() -> Env {
        Env {env: unsafe{ffi::sp_env()}}
    }

    pub fn destroy(self) {
        unsafe{ffi::sp_destroy(self.env)};
    }

    pub fn db(&mut self, dbname: &str) {
        // XXX: is dbname copied?
        let dbname = CString::new(dbname).unwrap();
        unsafe {ffi::setstring(self.env, "db\0".as_bytes(), dbname.as_bytes()) };
    }

    pub fn open(&mut self) {
        unsafe {ffi::sp_open(self.env)};
    }

    pub fn begin<'a>(&'a self) -> Transaction<'a> {
        let tx = unsafe{ffi::sp_begin(self.env)};
        assert!(!tx.is_null());
        Transaction {
            tx: tx,
            phantom: PhantomData
        }
    }

    /*
    fn get_object(&mut self, attr: &str) -> Voidptr {
        let attr = CString::new(attr).unwrap();
        unsafe {ffi::sp_getobject(self.env, attr.as_ptr() as *const u8)}
    }
    */

    pub fn get_db(&self, dbname: &str) -> Option<Db> {
        let dbstr = "db.".to_string() + dbname;
        let attr = CString::new(&dbstr[..]).unwrap();
        unsafe {
            let obj = ffi::sp_getobject(self.env, attr.as_ptr() as *const u8);
            if obj.is_null() {
                None
            }
            else {
                Some(Db {db: obj})
            }
        }
    }

    pub fn setattr(&mut self, key: &str, val: &str) {
        let key = CString::new(key).unwrap();
        let val = CString::new(val).unwrap();
        unsafe {ffi::setstring(self.env, key.as_bytes(), val.as_bytes()) };
    }

    pub fn setintattr(&mut self, key: &str, val: i64) {
        let key = CString::new(key).unwrap();
        unsafe {ffi::sp_setint(self.env, key.as_ptr(), val) };
    }
}

pub struct Object<'a> {
    obj: ffi::Voidptr,
    key: &'a [u8]
}

impl<'a> Object<'a> {
    pub fn get_value<'b>(&'b mut self) -> Option<&'b[u8]> {
        unsafe {
            let mut sz = 0;
            let valptr = ffi::sp_getstring(self.obj, "value\0".as_ptr(), &mut sz);

            // XXX: what if we stored an empty value?
            if valptr.is_null() {
                return None;
            }

            let s = slice::from_raw_parts(valptr as *const u8, sz as usize);
            return Some(s);
        }
    }
}

impl<'a> Drop for Object<'a> {
    fn drop(&mut self) {
        unsafe {ffi::sp_destroy(self.obj);}
    }
}

pub struct Cursor {
    obj: ffi::Voidptr,
}

pub struct CursorObject<'a> {
    obj: ffi::Voidptr,
    cur: &'a Cursor
}

impl<'a> CursorObject<'a> {
    pub fn get_key<'b>(&'b mut self) -> Option<&'b[u8]> {
        unsafe {
            let mut sz = 0;
            let keyptr = ffi::sp_getstring(self.obj, "key\0".as_ptr(), &mut sz);

            if keyptr.is_null() {
                return None;
            }

            let s = slice::from_raw_parts(keyptr as *const u8, sz as usize);
            return Some(s);
        }
    }

    pub fn get_value<'b>(&'b mut self) -> Option<&'b[u8]> {
        unsafe {
            let mut sz = 0;
            let valptr = ffi::sp_getstring(self.obj, "value\0".as_ptr(), &mut sz);

            // XXX: what if we stored an empty value?
            if valptr.is_null() {
                return None;
            }

            let s = slice::from_raw_parts(valptr as *const u8, sz as usize);
            return Some(s);
        }
    }
}

impl<'a> Drop for CursorObject<'a> {
    fn drop(&mut self) {
        unsafe {ffi::sp_destroy(self.obj);}
    }
}

impl Drop for Cursor {
    fn drop(&mut self) {
        unsafe {ffi::sp_destroy(self.obj);}
    }
}

impl Cursor {
    pub fn next<'a>(&'a self) -> Option<CursorObject<'a>> {
        let res = unsafe { ffi::sp_get(self.obj, ptr::null_mut()) };
        if res.is_null() {
            None
        }
        else {
            Some(CursorObject{obj: res, cur: &self})
        }
    }
}

impl Db {
    pub fn obj<'a>(&'a self) -> DbObject<'a> {
        let obj = unsafe { ffi::sp_object(self.db) };
        assert!(!obj.is_null());
        DbObject{o: obj, phantom: PhantomData}
    }

    // Consumes the DbObject
    pub fn put<'a>(&mut self, obj: DbObject<'a>) {
        let mut obj = obj;
        unsafe {ffi::sp_set(self.db, obj.o)}; // XXX: Check error code
        obj.o = ptr::null_mut(); // sp_set drops it
    }

    // XXX: Make sure self.db == dbobject.db
    pub fn get_<'a>(&'a self, pattern: DbObject<'a>) -> Option<DbObject<'a>> {
        unsafe {
            let res = ffi::sp_get(self.db, pattern.o);
            if res.is_null() {
                return None;
            }
            let mut pattern = pattern;
            pattern.o = res;
            return Some(pattern);
        }
    }

    pub fn set(&mut self, key: &[u8], value: &[u8]) {
        unsafe {
            //ffi::sp_set_kv(self.db, key.as_ptr(), key.len() as i32, value.as_ptr(), value.len() as i32);
            let obj = ffi::sp_object(self.db);
            assert!(!obj.is_null());
            ffi::setkey(obj, key);
            ffi::setvalue(obj, value);
            ffi::sp_set(self.db, obj);
            //ffi::sp_destroy(obj);
        }
    }

    pub fn iter_all(&mut self) -> Cursor {
        unsafe {
            let obj = ffi::sp_object(self.db);
            assert!(!obj.is_null());
            ffi::setstring(obj, "order\0".as_bytes(), ">=".as_bytes());
            let cursor = ffi::sp_cursor(self.db, obj);
            assert!(!cursor.is_null());
            Cursor {obj: cursor}
        }
    }

    pub fn get<'a>(&self, key: &'a [u8]) -> Option<Object<'a>> {
        assert!(key.len() > 0);
        unsafe {
            let pattern = ffi::sp_object(self.db);
            if pattern.is_null() {
                // XXX: Return Err
                return None;
            }
            ffi::setkey(pattern, key);
            let res = ffi::sp_get(self.db, pattern);
            if res.is_null() {
                ffi::sp_destroy(pattern);
                return None;
            }
            return Some(Object{obj: res, key: key});
        }
    }

    // XXX: Cannot distinguish between key not found and no value.
    pub fn get_value(&self, key: &[u8]) -> Option<Vec<u8>> {
        match self.get(key) {
            None => None,
            Some(mut valobj) => {
                match valobj.get_value() {
                    Some(val) => Some(val.to_vec()),
                    None => None
                }
            }
        }
    }
}
<commit_msg>Transaction#get_<commit_after>#![feature(libc)]

extern crate libc;
use std::ffi::CString;
use std::slice;
use std::ptr;

mod ffi;

#[derive(Debug, Copy, Clone)]
pub struct Env {
    env: ffi::Voidptr
}

#[derive(Debug, Copy, Clone)]
pub struct Db {
    db: ffi::Voidptr
}

// XXX: Use Arc
unsafe impl Sync for Db {}
unsafe impl Send for Db {}
unsafe impl Sync for Env {}
unsafe impl Send for Env {}

use std::marker::PhantomData;

#[derive(Debug)]
pub struct DbObject<'a> {
    o: ffi::Voidptr,
    phantom: PhantomData<&'a ()>
}

impl<'a> DbObject<'a> {
    // We name it attr(), to not mix it up with the set() method.
    pub fn attr<'b>(&'b mut self, path: &str, data: &'b[u8]) {
        let path = CString::new(path).unwrap();
        unsafe {ffi::setstring(self.o, path.as_bytes(), data); }
    }

    // Shortcut for attr("key", ...)
    pub fn key<'b>(&'b mut self, data: &'b[u8]) {
        unsafe {ffi::setstring(self.o, "key\0".as_bytes(), data); }
    }

    // Shortcut for attr("value", ...)
    pub fn value<'b>(&'b mut self, data: &'b[u8]) {
        unsafe {ffi::setstring(self.o, "value\0".as_bytes(), data); }
    }
}

impl<'a> Drop for DbObject<'a> {
    fn drop(&mut self) {
        if !self.o.is_null() {
            unsafe {ffi::sp_destroy(self.o)};
        }
    }
}

#[derive(Debug)]
pub struct Transaction<'a> {
    tx: ffi::Voidptr,
    phantom: PhantomData<&'a ()>
}

impl<'a> Drop for Transaction<'a> {
    // Unless commit() is called, this will roll back the transaction
    fn drop(&mut self) {
        if !self.tx.is_null() {
            unsafe {ffi::sp_destroy(self.tx);}
        }
    }
}

impl<'a> Transaction<'a> {
    fn commit(mut self) -> i32 {
        let rc = unsafe{ffi::sp_commit(self.tx)};
        self.tx = ptr::null_mut();
        rc as i32
    }

    // Consumes the DbObject
    pub fn put<'b>(&mut self, obj: DbObject<'b>) {
        let mut obj = obj;
        unsafe {ffi::sp_set(self.tx, obj.o)}; // XXX: Check error code
        obj.o = ptr::null_mut(); // sp_set drops it
    }

    // XXX: Make sure self.db == dbobject.db
    pub fn get_<'b>(&'b self, pattern: DbObject<'b>) -> Option<DbObject<'b>> {
        unsafe {
            let res = ffi::sp_get(self.tx, pattern.o);
            if res.is_null() {
                return None;
            }
            let mut pattern = pattern;
            pattern.o = res;
            return Some(pattern);
        }
    }
}

impl Env {
    pub fn new() -> Env {
        Env {env: unsafe{ffi::sp_env()}}
    }

    pub fn destroy(self) {
        unsafe{ffi::sp_destroy(self.env)};
    }

    pub fn db(&mut self, dbname: &str) {
        // XXX: is dbname copied?
        let dbname = CString::new(dbname).unwrap();
        unsafe {ffi::setstring(self.env, "db\0".as_bytes(), dbname.as_bytes()) };
    }

    pub fn open(&mut self) {
        unsafe {ffi::sp_open(self.env)};
    }

    pub fn begin<'a>(&'a self) -> Transaction<'a> {
        let tx = unsafe{ffi::sp_begin(self.env)};
        assert!(!tx.is_null());
        Transaction {
            tx: tx,
            phantom: PhantomData
        }
    }

    /*
    fn get_object(&mut self, attr: &str) -> Voidptr {
        let attr = CString::new(attr).unwrap();
        unsafe {ffi::sp_getobject(self.env, attr.as_ptr() as *const u8)}
    }
    */

    pub fn get_db(&self, dbname: &str) -> Option<Db> {
        let dbstr = "db.".to_string() + dbname;
        let attr = CString::new(&dbstr[..]).unwrap();
        unsafe {
            let obj = ffi::sp_getobject(self.env, attr.as_ptr() as *const u8);
            if obj.is_null() {
                None
            }
            else {
                Some(Db {db: obj})
            }
        }
    }

    pub fn setattr(&mut self, key: &str, val: &str) {
        let key = CString::new(key).unwrap();
        let val = CString::new(val).unwrap();
        unsafe {ffi::setstring(self.env, key.as_bytes(), val.as_bytes()) };
    }

    pub fn setintattr(&mut self, key: &str, val: i64) {
        let key = CString::new(key).unwrap();
        unsafe {ffi::sp_setint(self.env, key.as_ptr(), val) };
    }
}

pub struct Object<'a> {
    obj: ffi::Voidptr,
    key: &'a [u8]
}

impl<'a> Object<'a> {
    pub fn get_value<'b>(&'b mut self) -> Option<&'b[u8]> {
        unsafe {
            let mut sz = 0;
            let valptr = ffi::sp_getstring(self.obj, "value\0".as_ptr(), &mut sz);

            // XXX: what if we stored an empty value?
            if valptr.is_null() {
                return None;
            }

            let s = slice::from_raw_parts(valptr as *const u8, sz as usize);
            return Some(s);
        }
    }
}

impl<'a> Drop for Object<'a> {
    fn drop(&mut self) {
        unsafe {ffi::sp_destroy(self.obj);}
    }
}

pub struct Cursor {
    obj: ffi::Voidptr,
}

pub struct CursorObject<'a> {
    obj: ffi::Voidptr,
    cur: &'a Cursor
}

impl<'a> CursorObject<'a> {
    pub fn get_key<'b>(&'b mut self) -> Option<&'b[u8]> {
        unsafe {
            let mut sz = 0;
            let keyptr = ffi::sp_getstring(self.obj, "key\0".as_ptr(), &mut sz);

            if keyptr.is_null() {
                return None;
            }

            let s = slice::from_raw_parts(keyptr as *const u8, sz as usize);
            return Some(s);
        }
    }

    pub fn get_value<'b>(&'b mut self) -> Option<&'b[u8]> {
        unsafe {
            let mut sz = 0;
            let valptr = ffi::sp_getstring(self.obj, "value\0".as_ptr(), &mut sz);

            // XXX: what if we stored an empty value?
            if valptr.is_null() {
                return None;
            }

            let s = slice::from_raw_parts(valptr as *const u8, sz as usize);
            return Some(s);
        }
    }
}

impl<'a> Drop for CursorObject<'a> {
    fn drop(&mut self) {
        unsafe {ffi::sp_destroy(self.obj);}
    }
}

impl Drop for Cursor {
    fn drop(&mut self) {
        unsafe {ffi::sp_destroy(self.obj);}
    }
}

impl Cursor {
    pub fn next<'a>(&'a self) -> Option<CursorObject<'a>> {
        let res = unsafe { ffi::sp_get(self.obj, ptr::null_mut()) };
        if res.is_null() {
            None
        }
        else {
            Some(CursorObject{obj: res, cur: &self})
        }
    }
}

impl Db {
    pub fn obj<'a>(&'a self) -> DbObject<'a> {
        let obj = unsafe { ffi::sp_object(self.db) };
        assert!(!obj.is_null());
        DbObject{o: obj, phantom: PhantomData}
    }

    // Consumes the DbObject
    pub fn put<'a>(&mut self, obj: DbObject<'a>) {
        let mut obj = obj;
        unsafe {ffi::sp_set(self.db, obj.o)}; // XXX: Check error code
        obj.o = ptr::null_mut(); // sp_set drops it
    }

    // XXX: Make sure self.db == dbobject.db
    pub fn get_<'a>(&'a self, pattern: DbObject<'a>) -> Option<DbObject<'a>> {
        unsafe {
            let res = ffi::sp_get(self.db, pattern.o);
            if res.is_null() {
                return None;
            }
            let mut pattern = pattern;
            pattern.o = res;
            return Some(pattern);
        }
    }

    pub fn set(&mut self, key: &[u8], value: &[u8]) {
        unsafe {
            //ffi::sp_set_kv(self.db, key.as_ptr(), key.len() as i32, value.as_ptr(), value.len() as i32);
            let obj = ffi::sp_object(self.db);
            assert!(!obj.is_null());
            ffi::setkey(obj, key);
            ffi::setvalue(obj, value);
            ffi::sp_set(self.db, obj);
            //ffi::sp_destroy(obj);
        }
    }

    pub fn iter_all(&mut self) -> Cursor {
        unsafe {
            let obj = ffi::sp_object(self.db);
            assert!(!obj.is_null());
            ffi::setstring(obj, "order\0".as_bytes(), ">=".as_bytes());
            let cursor = ffi::sp_cursor(self.db, obj);
            assert!(!cursor.is_null());
            Cursor {obj: cursor}
        }
    }

    pub fn get<'a>(&self, key: &'a [u8]) -> Option<Object<'a>> {
        assert!(key.len() > 0);
        unsafe {
            let pattern = ffi::sp_object(self.db);
            if pattern.is_null() {
                // XXX: Return Err
                return None;
            }
            ffi::setkey(pattern, key);
            let res = ffi::sp_get(self.db, pattern);
            if res.is_null() {
                ffi::sp_destroy(pattern);
                return None;
            }
            return Some(Object{obj: res, key: key});
        }
    }

    // XXX: Cannot distinguish between key not found and no value.
    pub fn get_value(&self, key: &[u8]) -> Option<Vec<u8>> {
        match self.get(key) {
            None => None,
            Some(mut valobj) => {
                match valobj.get_value() {
                    Some(val) => Some(val.to_vec()),
                    None => None
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>pub trait Histogram {

    //TODO should be default impl of this trait
    fn record_single_value(&self, value: i64) -> Result<(), String>;

    //TODO should be default impl of this trait
    fn counts_array_index(&self, value: i64) -> Result<i32, String>;

    //TODO should be default impl of this trait
    fn get_bucket_index(&self, value: i64) -> i32;

    //TODO should be default impl of this trait
    fn get_sub_bucket_index(&self, value: i64, bucket_index: i32) -> i32;

}

pub struct SimpleHdrHistogram {
    pub leading_zeros_count_base: i32,
    pub sub_bucket_mask: i64,
    pub unit_magnitude: i32,
}

impl Default for SimpleHdrHistogram {
    fn default () -> SimpleHdrHistogram {
        SimpleHdrHistogram {
            leading_zeros_count_base: 0,
            sub_bucket_mask: 0,
            unit_magnitude: 0
        }
    }
}

impl Histogram for SimpleHdrHistogram {

    fn get_sub_bucket_index(&self, value: i64, bucket_index: i32) -> i32 {
        let sum = bucket_index + self.unit_magnitude;
        value.rotate_right(sum as u32) as i32
    }

    fn get_bucket_index(&self, value: i64) -> i32 {
        let valueOrred = value | self.sub_bucket_mask;
        self.leading_zeros_count_base - valueOrred.leading_zeros() as i32
    }

    fn record_single_value(&self, value: i64) -> Result<(), String> {

        let counts_index = self.counts_array_index(value);

        if true {
            Ok(())
        } else {
            Err(String::from("Could not record single value"))
        }

    }

    fn counts_array_index(&self, value: i64) -> Result<i32, String> {

        if (value < 0) {
            Err(String::from("Histogram recorded values cannot be negative."))
        } else {
            Ok(0)
        }
    }
}

#[test]
fn can_record_single_value() {
    let the_hist = SimpleHdrHistogram { ..Default::default() };
    let result = the_hist.record_single_value(99);

    match result {
        Ok(_) => (),
        Err(err) => panic!(format!("could not add single record to histogram because error: {}", err))
    }
}

#[test]
fn can_compute_counts_array_index() {
    let the_hist = SimpleHdrHistogram { ..Default::default() };
    let result = the_hist.counts_array_index(99);

    match result {
        Ok(_) => (),
        Err(err) => panic!(format!("could not compute counts array index because error: {}", err))
    }
}

#[test]
fn can_get_bucket_index() {
    let the_hist = SimpleHdrHistogram { ..Default::default() };
    let result = the_hist.get_bucket_index(99);
    assert_eq!(result, 0)
}

#[test]
fn can_get_sub_bucket_index() {
    let the_hist = SimpleHdrHistogram { ..Default::default() };
    let result = the_hist.get_sub_bucket_index(99, 1);
    assert_eq!(result, 0)
}<commit_msg>wip stuff<commit_after>pub struct SimpleHdrHistogram {
    pub leading_zeros_count_base: i32,
    pub sub_bucket_mask: i64,
    pub unit_magnitude: i32,
    pub sub_bucket_count: i32,
    pub sub_bucket_half_count: i32,
    pub sub_bucket_half_count_magnitude: i32,
}

impl Default for SimpleHdrHistogram {
    fn default () -> SimpleHdrHistogram {
        SimpleHdrHistogram {
            leading_zeros_count_base: 0,
            sub_bucket_mask: 0,
            unit_magnitude: 0,
            sub_bucket_count: 0,
            sub_bucket_half_count: 0,
            sub_bucket_half_count_magnitude: 0,
        }
    }
}

pub trait Histogram {

    //TODO should be default impl of this trait
    fn record_single_value(&self, value: i64) -> Result<(), String>;

    //TODO should be default impl of this trait
    fn counts_array_index(&self, value: i64) -> Result<i32, String>;
    fn counts_array_index_sub(&self, bucket_index: i32, sub_bucket_index: i32) -> i32;

    //TODO should be default impl of this trait
    fn get_bucket_index(&self, value: i64) -> i32;

    //TODO should be default impl of this trait
    fn get_sub_bucket_index(&self, value: i64, bucket_index: i32) -> i32;

}

impl Histogram for SimpleHdrHistogram {

    fn get_sub_bucket_index(&self, value: i64, bucket_index: i32) -> i32 {
        let sum = bucket_index + self.unit_magnitude;
        value.rotate_right(sum as u32) as i32
    }

    fn get_bucket_index(&self, value: i64) -> i32 {
        let valueOrred = value | self.sub_bucket_mask;
        self.leading_zeros_count_base - valueOrred.leading_zeros() as i32
    }

    fn record_single_value(&self, value: i64) -> Result<(), String> {
        let counts_index = self.counts_array_index(value);
        if true {
            Ok(())
        } else {
            Err(String::from("Could not record single value"))
        }
    }

    fn counts_array_index(&self, value: i64) -> Result<i32, String> {
        if (value < 0) {
            Err(String::from("Histogram recorded values cannot be negative."))
        } else {
            let bucket_index = self.get_bucket_index(value);
            let sub_bucket_index = self.get_sub_bucket_index(value, bucket_index);
            let result = self.counts_array_index_sub(bucket_index, sub_bucket_index);
            Ok(result)
        }
    }

    fn counts_array_index_sub(&self, bucket_index: i32, sub_bucket_index: i32) -> i32 {
        assert!(sub_bucket_index < self.sub_bucket_count);
        assert!(bucket_index == 0 || (sub_bucket_index >= self.sub_bucket_half_count));

        let bucket_base_index = (bucket_index + 1) << self.sub_bucket_half_count_magnitude;
        let offset_in_bucket = sub_bucket_index - self.sub_bucket_half_count;

        bucket_base_index + offset_in_bucket
    }
}

#[test]
fn can_record_single_value() {
    let the_hist = SimpleHdrHistogram { ..Default::default() };
    let result = the_hist.record_single_value(99);

    match result {
        Ok(_) => (),
        Err(err) => panic!(format!("could not add single record to histogram because error: {}", err))
    }
}

#[test]
fn can_compute_counts_array_index() {
    let the_hist = SimpleHdrHistogram { ..Default::default() };
    let result = the_hist.counts_array_index(99);

    match result {
        Ok(_) => (),
        Err(err) => panic!(format!("could not compute counts array index because error: {}", err))
    }
}

#[test]
fn can_get_bucket_index() {
    let the_hist = SimpleHdrHistogram { ..Default::default() };
    let result = the_hist.get_bucket_index(99);
    assert_eq!(result, 0)
}

#[test]
fn can_get_sub_bucket_index() {
    let the_hist = SimpleHdrHistogram { ..Default::default() };
    let result = the_hist.get_sub_bucket_index(99, 1);
    assert_eq!(result, 0)
}<|endoftext|>


<commit_before>use super::utils::{LspPositionConversion, TupleToPosition};
use crate::{
    compiler::{
        cst::{Cst, CstKind, UnwrapWhitespaceAndComment},
        rcst_to_cst::RcstToCst,
    },
    input::Input,
};
use lsp_types::{FoldingRange, FoldingRangeKind};

#[salsa::query_group(FoldingRangeDbStorage)]
pub trait FoldingRangeDb: LspPositionConversion + RcstToCst {
    fn folding_ranges(&self, input: Input) -> Vec<FoldingRange>;
}

fn folding_ranges(db: &dyn FoldingRangeDb, input: Input) -> Vec<FoldingRange> {
    let mut context = Context::new(db, input.clone());
    let cst = db.cst(input).unwrap();
    context.visit_csts(&cst);
    context.ranges
}

struct Context<'a> {
    db: &'a dyn FoldingRangeDb,
    input: Input,
    ranges: Vec<FoldingRange>,
}
impl<'a> Context<'a> {
    fn new(db: &'a dyn FoldingRangeDb, input: Input) -> Self {
        Context {
            db,
            input,
            ranges: vec![],
        }
    }

    fn visit_csts(&mut self, csts: &[Cst]) {
        for cst in csts {
            self.visit_cst(cst);
        }
    }
    fn visit_cst(&mut self, cst: &Cst) {
        match &cst.kind {
            CstKind::EqualsSign => {}
            CstKind::Comma => {}
            CstKind::Dot => {}
            CstKind::Colon => {}
            CstKind::ColonEqualsSign => {}
            CstKind::OpeningParenthesis => {}
            CstKind::ClosingParenthesis => {}
            CstKind::OpeningBracket => {}
            CstKind::ClosingBracket => {}
            CstKind::OpeningCurlyBrace => {}
            CstKind::ClosingCurlyBrace => {}
            CstKind::Arrow => {}
            CstKind::DoubleQuote => {}
            CstKind::Octothorpe => {}
            CstKind::Whitespace(_) => {}
            CstKind::Newline(_) => {}
            // TODO: support folding ranges for comments
            CstKind::Comment { .. } => {}
            CstKind::TrailingWhitespace { child, .. } => self.visit_cst(child),
            CstKind::Identifier(_) => {}
            CstKind::Symbol(_) => {}
            CstKind::Int { .. } => {}
            // TODO: support folding ranges for multiline texts
            CstKind::Text { .. } => {}
            CstKind::TextPart(_) => {}
            CstKind::Parenthesized { inner, .. } => self.visit_cst(inner),
            CstKind::Call { name, arguments } => {
                if !arguments.is_empty() {
                    let name = name.unwrap_whitespace_and_comment();
                    let last_argument = arguments.last().unwrap().unwrap_whitespace_and_comment();
                    self.push(
                        name.span.end,
                        last_argument.span.end,
                        FoldingRangeKind::Region,
                    );
                }

                self.visit_cst(name);
                self.visit_csts(&arguments);
            }
            // TODO: support folding ranges for structs
            CstKind::Struct { fields, .. } => self.visit_csts(fields),
            CstKind::StructField { key, value, .. } => {
                self.visit_cst(key);
                self.visit_cst(value);
            }
            CstKind::StructAccess { struct_, dot, key } => {
                self.visit_cst(struct_);
                self.visit_cst(dot);
                self.visit_cst(key);
            }
            CstKind::Lambda {
                opening_curly_brace,
                parameters_and_arrow,
                body,
                closing_curly_brace,
            } => {
                let opening_curly_brace = opening_curly_brace.unwrap_whitespace_and_comment();
                assert!(matches!(
                    opening_curly_brace.kind,
                    CstKind::OpeningCurlyBrace { .. }
                ));

                let closing_curly_brace = closing_curly_brace.unwrap_whitespace_and_comment();
                assert!(matches!(
                    closing_curly_brace.kind,
                    CstKind::ClosingCurlyBrace { .. }
                ));

                self.push(
                    opening_curly_brace.span.end,
                    closing_curly_brace.span.start,
                    FoldingRangeKind::Region,
                );
                if let Some((parameters, _)) = parameters_and_arrow {
                    self.visit_csts(&parameters);
                }
                self.visit_csts(&body);
            }
            CstKind::Assignment {
                name,
                assignment_sign,
                parameters,
                body,
            } => {
                if !body.is_empty() {
                    let equals_sign = assignment_sign.unwrap_whitespace_and_comment();
                    assert!(matches!(equals_sign.kind, CstKind::EqualsSign { .. }));

                    let last_expression = body.last().unwrap().unwrap_whitespace_and_comment();

                    self.push(
                        equals_sign.span.end,
                        last_expression.span.end,
                        FoldingRangeKind::Region,
                    );
                }

                self.visit_cst(name);
                self.visit_csts(&parameters);
                self.visit_csts(&body);
            }
            CstKind::Error { .. } => {}
        }
    }

    fn push(&mut self, start: usize, end: usize, kind: FoldingRangeKind) {
        let start = self
            .db
            .offset_to_lsp(self.input.clone(), start)
            .to_position();
        let end = self.db.offset_to_lsp(self.input.clone(), end).to_position();

        self.ranges.push(FoldingRange {
            start_line: start.line,
            start_character: Some(start.character),
            end_line: end.line,
            end_character: Some(end.character),
            kind: Some(kind),
        });
    }
}
<commit_msg>Fix folding ranges<commit_after>use super::utils::{LspPositionConversion, TupleToPosition};
use crate::{
    compiler::{
        cst::{Cst, CstKind, UnwrapWhitespaceAndComment},
        rcst_to_cst::RcstToCst,
    },
    input::Input,
};
use lsp_types::{FoldingRange, FoldingRangeKind};

#[salsa::query_group(FoldingRangeDbStorage)]
pub trait FoldingRangeDb: LspPositionConversion + RcstToCst {
    fn folding_ranges(&self, input: Input) -> Vec<FoldingRange>;
}

fn folding_ranges(db: &dyn FoldingRangeDb, input: Input) -> Vec<FoldingRange> {
    let mut context = Context::new(db, input.clone());
    let cst = db.cst(input).unwrap();
    context.visit_csts(&cst);
    context.ranges
}

struct Context<'a> {
    db: &'a dyn FoldingRangeDb,
    input: Input,
    ranges: Vec<FoldingRange>,
}
impl<'a> Context<'a> {
    fn new(db: &'a dyn FoldingRangeDb, input: Input) -> Self {
        Context {
            db,
            input,
            ranges: vec![],
        }
    }

    fn visit_csts(&mut self, csts: &[Cst]) {
        for cst in csts {
            self.visit_cst(cst);
        }
    }
    fn visit_cst(&mut self, cst: &Cst) {
        match &cst.kind {
            CstKind::EqualsSign => {}
            CstKind::Comma => {}
            CstKind::Dot => {}
            CstKind::Colon => {}
            CstKind::ColonEqualsSign => {}
            CstKind::OpeningParenthesis => {}
            CstKind::ClosingParenthesis => {}
            CstKind::OpeningBracket => {}
            CstKind::ClosingBracket => {}
            CstKind::OpeningCurlyBrace => {}
            CstKind::ClosingCurlyBrace => {}
            CstKind::Arrow => {}
            CstKind::DoubleQuote => {}
            CstKind::Octothorpe => {}
            CstKind::Whitespace(_) => {}
            CstKind::Newline(_) => {}
            // TODO: support folding ranges for comments
            CstKind::Comment { .. } => {}
            CstKind::TrailingWhitespace { child, .. } => self.visit_cst(child),
            CstKind::Identifier(_) => {}
            CstKind::Symbol(_) => {}
            CstKind::Int { .. } => {}
            // TODO: support folding ranges for multiline texts
            CstKind::Text { .. } => {}
            CstKind::TextPart(_) => {}
            CstKind::Parenthesized { inner, .. } => self.visit_cst(inner),
            CstKind::Call { name, arguments } => {
                if !arguments.is_empty() {
                    let name = name.unwrap_whitespace_and_comment();
                    let last_argument = arguments.last().unwrap().unwrap_whitespace_and_comment();
                    self.push(
                        name.span.end,
                        last_argument.span.end,
                        FoldingRangeKind::Region,
                    );
                }

                self.visit_cst(name);
                self.visit_csts(&arguments);
            }
            // TODO: support folding ranges for structs
            CstKind::Struct { fields, .. } => self.visit_csts(fields),
            CstKind::StructField { key, value, .. } => {
                self.visit_cst(key);
                self.visit_cst(value);
            }
            CstKind::StructAccess { struct_, dot, key } => {
                self.visit_cst(struct_);
                self.visit_cst(dot);
                self.visit_cst(key);
            }
            CstKind::Lambda {
                opening_curly_brace,
                parameters_and_arrow,
                body,
                closing_curly_brace,
            } => {
                let opening_curly_brace = opening_curly_brace.unwrap_whitespace_and_comment();
                assert!(matches!(
                    opening_curly_brace.kind,
                    CstKind::OpeningCurlyBrace { .. }
                ));

                let closing_curly_brace = closing_curly_brace.unwrap_whitespace_and_comment();
                assert!(matches!(
                    closing_curly_brace.kind,
                    CstKind::ClosingCurlyBrace { .. }
                ));

                self.push(
                    opening_curly_brace.span.end,
                    closing_curly_brace.span.start,
                    FoldingRangeKind::Region,
                );
                if let Some((parameters, _)) = parameters_and_arrow {
                    self.visit_csts(&parameters);
                }
                self.visit_csts(&body);
            }
            CstKind::Assignment {
                name,
                assignment_sign,
                parameters,
                body,
            } => {
                if !body.is_empty() {
                    let assignment_sign = assignment_sign.unwrap_whitespace_and_comment();
                    let last_expression = body.last().unwrap().unwrap_whitespace_and_comment();

                    self.push(
                        assignment_sign.span.end,
                        last_expression.span.end,
                        FoldingRangeKind::Region,
                    );
                }

                self.visit_cst(name);
                self.visit_csts(&parameters);
                self.visit_csts(&body);
            }
            CstKind::Error { .. } => {}
        }
    }

    fn push(&mut self, start: usize, end: usize, kind: FoldingRangeKind) {
        let start = self
            .db
            .offset_to_lsp(self.input.clone(), start)
            .to_position();
        let end = self.db.offset_to_lsp(self.input.clone(), end).to_position();

        self.ranges.push(FoldingRange {
            start_line: start.line,
            start_character: Some(start.character),
            end_line: end.line,
            end_character: Some(end.character),
            kind: Some(kind),
        });
    }
}
<|endoftext|>


<commit_before>use clap::{Args, IntoApp, Parser, Subcommand};

#[test]
fn arg_help_heading_applied() {
    #[derive(Debug, Clone, Parser)]
    struct CliOptions {
        #[clap(long)]
        #[clap(help_heading = Some("HEADING A"))]
        should_be_in_section_a: Option<u32>,

        #[clap(long)]
        no_section: Option<u32>,
    }

    let app = CliOptions::into_app();

    let should_be_in_section_a = app
        .get_arguments()
        .find(|a| a.get_name() == "should-be-in-section-a")
        .unwrap();
    assert_eq!(should_be_in_section_a.get_help_heading(), Some("HEADING A"));

    let should_be_in_section_b = app
        .get_arguments()
        .find(|a| a.get_name() == "no-section")
        .unwrap();
    assert_eq!(should_be_in_section_b.get_help_heading(), None);
}

#[test]
fn app_help_heading_applied() {
    #[derive(Debug, Clone, Parser)]
    #[clap(help_heading = "DEFAULT")]
    struct CliOptions {
        #[clap(long)]
        #[clap(help_heading = Some("HEADING A"))]
        should_be_in_section_a: Option<u32>,

        #[clap(long)]
        should_be_in_default_section: Option<u32>,
    }

    let app = CliOptions::into_app();

    let should_be_in_section_a = app
        .get_arguments()
        .find(|a| a.get_name() == "should-be-in-section-a")
        .unwrap();
    assert_eq!(should_be_in_section_a.get_help_heading(), Some("HEADING A"));

    let should_be_in_default_section = app
        .get_arguments()
        .find(|a| a.get_name() == "should-be-in-default-section")
        .unwrap();
    assert_eq!(
        should_be_in_default_section.get_help_heading(),
        Some("DEFAULT")
    );
}

#[test]
fn app_help_heading_flattened() {
    #[derive(Debug, Clone, Parser)]
    struct CliOptions {
        #[clap(flatten)]
        options_a: OptionsA,

        #[clap(flatten)]
        options_b: OptionsB,

        #[clap(subcommand)]
        sub_a: SubA,

        #[clap(long)]
        should_be_in_default_section: Option<u32>,
    }

    #[derive(Debug, Clone, Args)]
    #[clap(help_heading = "HEADING A")]
    struct OptionsA {
        #[clap(long)]
        should_be_in_section_a: Option<u32>,
    }

    #[derive(Debug, Clone, Args)]
    #[clap(help_heading = "HEADING B")]
    struct OptionsB {
        #[clap(long)]
        should_be_in_section_b: Option<u32>,
    }

    #[derive(Debug, Clone, Subcommand)]
    enum SubA {
        #[clap(flatten)]
        SubB(SubB),
        #[clap(subcommand)]
        SubC(SubC),
        SubAOne,
        #[clap(help_heading = "SUB A")]
        SubATwo {
            should_be_in_sub_a: Option<u32>,
        },
    }

    #[derive(Debug, Clone, Subcommand)]
    enum SubB {
        #[clap(help_heading = "SUB B")]
        SubBOne { should_be_in_sub_b: Option<u32> },
    }

    #[derive(Debug, Clone, Subcommand)]
    enum SubC {
        #[clap(help_heading = "SUB C")]
        SubCOne { should_be_in_sub_c: Option<u32> },
    }

    let app = CliOptions::into_app();

    let should_be_in_section_a = app
        .get_arguments()
        .find(|a| a.get_name() == "should-be-in-section-a")
        .unwrap();
    assert_eq!(should_be_in_section_a.get_help_heading(), Some("HEADING A"));

    let should_be_in_section_b = app
        .get_arguments()
        .find(|a| a.get_name() == "should-be-in-section-b")
        .unwrap();
    assert_eq!(should_be_in_section_b.get_help_heading(), Some("HEADING B"));

    let should_be_in_default_section = app
        .get_arguments()
        .find(|a| a.get_name() == "should-be-in-default-section")
        .unwrap();
    assert_eq!(should_be_in_default_section.get_help_heading(), None);

    let sub_a_two = app.find_subcommand("sub-a-two").unwrap();

    let should_be_in_sub_a = sub_a_two
        .get_arguments()
        .find(|a| a.get_name() == "should-be-in-sub-a")
        .unwrap();
    assert_eq!(should_be_in_sub_a.get_help_heading(), Some("SUB A"));

    let sub_b_one = app.find_subcommand("sub-b-one").unwrap();

    let should_be_in_sub_b = sub_b_one
        .get_arguments()
        .find(|a| a.get_name() == "should-be-in-sub-b")
        .unwrap();
    assert_eq!(should_be_in_sub_b.get_help_heading(), Some("SUB B"));

    let sub_c = app.find_subcommand("sub-c").unwrap();
    let sub_c_one = sub_c.find_subcommand("sub-c-one").unwrap();

    let should_be_in_sub_c = sub_c_one
        .get_arguments()
        .find(|a| a.get_name() == "should-be-in-sub-c")
        .unwrap();
    assert_eq!(should_be_in_sub_c.get_help_heading(), Some("SUB C"));
}

#[test]
fn flatten_field_with_help_heading() {
    #[derive(Debug, Clone, Parser)]
    struct CliOptions {
        #[clap(flatten)]
        #[clap(help_heading = "HEADING A")]
        options_a: OptionsA,
    }

    #[derive(Debug, Clone, Args)]
    struct OptionsA {
        #[clap(long)]
        should_be_in_section_a: Option<u32>,
    }

    let app = CliOptions::into_app();

    let should_be_in_section_a = app
        .get_arguments()
        .find(|a| a.get_name() == "should-be-in-section-a")
        .unwrap();
    assert_eq!(should_be_in_section_a.get_help_heading(), Some("HEADING A"));
}
<commit_msg>test(derive): Verify derive-genned errors are formatted<commit_after>use clap::{AppSettings, Args, ColorChoice, IntoApp, Parser, Subcommand};

#[test]
fn arg_help_heading_applied() {
    #[derive(Debug, Clone, Parser)]
    struct CliOptions {
        #[clap(long)]
        #[clap(help_heading = Some("HEADING A"))]
        should_be_in_section_a: Option<u32>,

        #[clap(long)]
        no_section: Option<u32>,
    }

    let app = CliOptions::into_app();

    let should_be_in_section_a = app
        .get_arguments()
        .find(|a| a.get_name() == "should-be-in-section-a")
        .unwrap();
    assert_eq!(should_be_in_section_a.get_help_heading(), Some("HEADING A"));

    let should_be_in_section_b = app
        .get_arguments()
        .find(|a| a.get_name() == "no-section")
        .unwrap();
    assert_eq!(should_be_in_section_b.get_help_heading(), None);
}

#[test]
fn app_help_heading_applied() {
    #[derive(Debug, Clone, Parser)]
    #[clap(help_heading = "DEFAULT")]
    struct CliOptions {
        #[clap(long)]
        #[clap(help_heading = Some("HEADING A"))]
        should_be_in_section_a: Option<u32>,

        #[clap(long)]
        should_be_in_default_section: Option<u32>,
    }

    let app = CliOptions::into_app();

    let should_be_in_section_a = app
        .get_arguments()
        .find(|a| a.get_name() == "should-be-in-section-a")
        .unwrap();
    assert_eq!(should_be_in_section_a.get_help_heading(), Some("HEADING A"));

    let should_be_in_default_section = app
        .get_arguments()
        .find(|a| a.get_name() == "should-be-in-default-section")
        .unwrap();
    assert_eq!(
        should_be_in_default_section.get_help_heading(),
        Some("DEFAULT")
    );
}

#[test]
fn app_help_heading_flattened() {
    #[derive(Debug, Clone, Parser)]
    struct CliOptions {
        #[clap(flatten)]
        options_a: OptionsA,

        #[clap(flatten)]
        options_b: OptionsB,

        #[clap(subcommand)]
        sub_a: SubA,

        #[clap(long)]
        should_be_in_default_section: Option<u32>,
    }

    #[derive(Debug, Clone, Args)]
    #[clap(help_heading = "HEADING A")]
    struct OptionsA {
        #[clap(long)]
        should_be_in_section_a: Option<u32>,
    }

    #[derive(Debug, Clone, Args)]
    #[clap(help_heading = "HEADING B")]
    struct OptionsB {
        #[clap(long)]
        should_be_in_section_b: Option<u32>,
    }

    #[derive(Debug, Clone, Subcommand)]
    enum SubA {
        #[clap(flatten)]
        SubB(SubB),
        #[clap(subcommand)]
        SubC(SubC),
        SubAOne,
        #[clap(help_heading = "SUB A")]
        SubATwo {
            should_be_in_sub_a: Option<u32>,
        },
    }

    #[derive(Debug, Clone, Subcommand)]
    enum SubB {
        #[clap(help_heading = "SUB B")]
        SubBOne { should_be_in_sub_b: Option<u32> },
    }

    #[derive(Debug, Clone, Subcommand)]
    enum SubC {
        #[clap(help_heading = "SUB C")]
        SubCOne { should_be_in_sub_c: Option<u32> },
    }

    let app = CliOptions::into_app();

    let should_be_in_section_a = app
        .get_arguments()
        .find(|a| a.get_name() == "should-be-in-section-a")
        .unwrap();
    assert_eq!(should_be_in_section_a.get_help_heading(), Some("HEADING A"));

    let should_be_in_section_b = app
        .get_arguments()
        .find(|a| a.get_name() == "should-be-in-section-b")
        .unwrap();
    assert_eq!(should_be_in_section_b.get_help_heading(), Some("HEADING B"));

    let should_be_in_default_section = app
        .get_arguments()
        .find(|a| a.get_name() == "should-be-in-default-section")
        .unwrap();
    assert_eq!(should_be_in_default_section.get_help_heading(), None);

    let sub_a_two = app.find_subcommand("sub-a-two").unwrap();

    let should_be_in_sub_a = sub_a_two
        .get_arguments()
        .find(|a| a.get_name() == "should-be-in-sub-a")
        .unwrap();
    assert_eq!(should_be_in_sub_a.get_help_heading(), Some("SUB A"));

    let sub_b_one = app.find_subcommand("sub-b-one").unwrap();

    let should_be_in_sub_b = sub_b_one
        .get_arguments()
        .find(|a| a.get_name() == "should-be-in-sub-b")
        .unwrap();
    assert_eq!(should_be_in_sub_b.get_help_heading(), Some("SUB B"));

    let sub_c = app.find_subcommand("sub-c").unwrap();
    let sub_c_one = sub_c.find_subcommand("sub-c-one").unwrap();

    let should_be_in_sub_c = sub_c_one
        .get_arguments()
        .find(|a| a.get_name() == "should-be-in-sub-c")
        .unwrap();
    assert_eq!(should_be_in_sub_c.get_help_heading(), Some("SUB C"));
}

#[test]
fn flatten_field_with_help_heading() {
    #[derive(Debug, Clone, Parser)]
    struct CliOptions {
        #[clap(flatten)]
        #[clap(help_heading = "HEADING A")]
        options_a: OptionsA,
    }

    #[derive(Debug, Clone, Args)]
    struct OptionsA {
        #[clap(long)]
        should_be_in_section_a: Option<u32>,
    }

    let app = CliOptions::into_app();

    let should_be_in_section_a = app
        .get_arguments()
        .find(|a| a.get_name() == "should-be-in-section-a")
        .unwrap();
    assert_eq!(should_be_in_section_a.get_help_heading(), Some("HEADING A"));
}

// The challenge with this test is creating an error situation not caught by `clap`'s error checking
// but by the code that `clap_derive` generates.
//
// Ultimately, the easiest way to confirm is to put a debug statement in the desired error path.
#[test]
fn derive_generated_error_has_full_context() {
    #[derive(Debug, Parser)]
    #[clap(setting(AppSettings::SubcommandsNegateReqs), color = ColorChoice::Never)]
    struct Opts {
        #[clap(long)]
        req_str: String,

        #[clap(subcommand)]
        cmd: Option<SubCommands>,
    }

    #[derive(Debug, Parser)]
    enum SubCommands {
        Sub {
            #[clap(short, long, parse(from_occurrences))]
            verbose: u8,
        },
    }

    let result = Opts::try_parse_from(&["test", "sub"]);
    assert!(
        result.is_err(),
        "`SubcommandsNegateReqs` with non-optional `req_str` should fail: {:?}",
        result.unwrap()
    );

    let expected = r#"error: The following required argument was not provided: req-str

USAGE:
    clap_derive --req-str <REQ_STR>
    clap_derive <SUBCOMMAND>

For more information try --help
"#;
    assert_eq!(result.unwrap_err().to_string(), expected);
}
<|endoftext|>


<commit_before>extern mod std;
extern mod zmq;

use std::comm::{stream, Chan};
use std::io;
use std::str;
use std::task;

fn new_server(ctx: zmq::Context, ch: &Chan<()>) {
    let socket = ctx.socket(zmq::REP).unwrap();
    socket.bind("tcp://127.0.0.1:3456").get();

    let msg = socket.recv_str(0).unwrap();
    io::println(fmt!("server received %?", msg));

    let msg = fmt!("hello %?", msg);
    io::println(fmt!("server sending %?", msg));

    match socket.send_str(msg, 0) {
        Ok(()) => { },
        Err(e) => fail!(e.to_str())
    }

    // Let the main thread know we're done.
    ch.send(());
}

fn new_client(ctx: zmq::Context) {
    io::println("starting client");

    let socket = ctx.socket(zmq::REQ).unwrap();

    socket.set_sndhwm(10).get();
    socket.set_rcvhwm(10).get();
    io::println(fmt!("rcvhwm: %?", socket.get_rcvhwm().get()));
    io::println(fmt!("sndhwm: %?", socket.get_sndhwm().get()));

    socket.set_identity(str::to_bytes("identity")).get();

    let identity = socket.get_identity().unwrap();
    io::println(fmt!("identity: %?", str::from_bytes(identity)));

    io::println("client connecting to server");

    socket.connect("tcp://127.0.0.1:3456").get();

    let msg = "foo";
    io::println(fmt!("client sending %?", msg));
    socket.send_str(msg, 0).get();

    let msg = socket.recv_str(0).unwrap();
    io::println(fmt!("client recieving %?", msg));
}

fn main() {
    let (major, minor, patch) = zmq::version();

    io::println(fmt!("version: %d %d %d", major, minor, patch));

    let ctx = zmq::init(1).unwrap();

    // We need to start the server in a separate scheduler as it blocks.
    let (po, ch) = stream();
    do task::spawn_sched(task::SingleThreaded) { new_server(ctx, &ch) }

    new_client(ctx);

    // Wait for the server to shut down.
    po.recv();

    ctx.term().get();
}
<commit_msg>Fix other example<commit_after>extern mod std;
extern mod zmq;

use std::comm::{stream, Chan};
use std::io;
use std::str;
use std::task;

fn new_server(ctx: zmq::Context, ch: &Chan<()>) {
    let socket = ctx.socket(zmq::REP).unwrap();
    socket.bind("tcp://127.0.0.1:3456").get();

    let msg = socket.recv_str(0).unwrap();
    io::println(fmt!("server received %?", msg));

    let msg = fmt!("hello %?", msg);
    io::println(fmt!("server sending %?", msg));

    match socket.send_str(msg, 0) {
        Ok(()) => { },
        Err(e) => fail!(e.to_str())
    }

    // Let the main thread know we're done.
    ch.send(());
}

fn new_client(ctx: zmq::Context) {
    io::println("starting client");

    let socket = ctx.socket(zmq::REQ).unwrap();

    socket.set_sndhwm(10).get();
    socket.set_rcvhwm(10).get();
    io::println(fmt!("rcvhwm: %?", socket.get_rcvhwm().get()));
    io::println(fmt!("sndhwm: %?", socket.get_sndhwm().get()));

    socket.set_identity("identity".as_bytes()).get();

    let identity = socket.get_identity().unwrap();
    io::println(fmt!("identity: %?", str::from_bytes(identity)));

    io::println("client connecting to server");

    socket.connect("tcp://127.0.0.1:3456").get();

    let msg = "foo";
    io::println(fmt!("client sending %?", msg));
    socket.send_str(msg, 0).get();

    let msg = socket.recv_str(0).unwrap();
    io::println(fmt!("client recieving %?", msg));
}

fn main() {
    let (major, minor, patch) = zmq::version();

    io::println(fmt!("version: %d %d %d", major, minor, patch));

    let ctx = zmq::init(1).unwrap();

    // We need to start the server in a separate scheduler as it blocks.
    let (po, ch) = stream();
    do task::spawn_sched(task::SingleThreaded) { new_server(ctx, &ch) }

    new_client(ctx);

    // Wait for the server to shut down.
    po.recv();

    ctx.term().get();
}
<|endoftext|>


<commit_before>use std::{
    io,
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use tokio::sync::mpsc;

use log::{debug, error, info, warn};
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, Query},
    rr::{DNSClass, Name, RData, RecordType},
};

use crate::{
    acl::AccessControl,
    config::ConfigType,
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

mod upstream;

fn should_forward_by_ptr_name(acl: &AccessControl, name: &Name) -> bool {
    let mut iter = name.iter().rev();
    let mut next = || std::str::from_utf8(iter.next().unwrap_or(&[48])).unwrap_or("*");
    if !"arpa".eq_ignore_ascii_case(next()) {
        return acl.is_default_in_proxy_list();
    }
    match &next().to_ascii_lowercase()[..] {
        "in-addr" => {
            let mut octets: [u8; 4] = [0; 4];
            for octet in octets.iter_mut() {
                match next().parse() {
                    Ok(result) => *octet = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V4(Ipv4Addr::new(octets[0], octets[1], octets[2], octets[3])))
        }
        "ip6" => {
            let mut segments: [u16; 8] = [0; 8];
            for segment in segments.iter_mut() {
                match u16::from_str_radix(&[next(), next(), next(), next()].concat(), 16) {
                    Ok(result) => *segment = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V6(Ipv6Addr::new(
                segments[0], segments[1], segments[2], segments[3], segments[4], segments[5], segments[6], segments[7]
            )))
        }
        _ => acl.is_default_in_proxy_list(),
    }
}

/// given the query, determine whether remote/local query should be used, or inconclusive
fn should_forward_by_query(acl: &Option<AccessControl>, query: &Query) -> Option<bool> {
    if let Some(acl) = acl {
        if query.query_class() != DNSClass::IN {
            // unconditionally use default for all non-IN queries
            Some(acl.is_default_in_proxy_list())
        } else if query.query_type() == RecordType::PTR {
            Some(should_forward_by_ptr_name(acl, query.name()))
        } else {
            let result = acl.check_name_in_proxy_list(query.name());
            if result == None && match query.query_type() {
                RecordType::A => acl.is_ipv4_empty(),
                RecordType::AAAA => acl.is_ipv6_empty(),
                RecordType::ANY => acl.is_ipv4_empty() && acl.is_ipv6_empty(),
                RecordType::PTR => unreachable!(),
                _ => true,
            } {
                Some(acl.is_default_in_proxy_list())
            } else {
                result
            }
        }
    } else {
        Some(true)
    }
}

/// given the local response, determine whether remote response should be used instead
fn should_forward_by_response(
    acl: &Option<AccessControl>,
    local_response: &io::Result<Message>,
    query: &Query,
) -> bool {
    if let Some(acl) = acl {
        macro_rules! examine_record {
            ($rec:ident, $is_answer:expr) => {
                if let RData::CNAME(ref name) = $rec.rdata() {
                    match acl.check_name_in_proxy_list(name) {
                        Some(value) => return value,
                        None => continue,
                    }
                } else if $is_answer && !query.query_type().is_any() && $rec.record_type() != query.query_type() {
                    warn!("local DNS response has inconsistent answer type {} for query {}", $rec.record_type(), query);
                    return true;
                }
                let forward = match $rec.rdata() {
                    RData::A(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V4(*ip)),
                    RData::AAAA(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V6(*ip)),
                    RData::PTR(_) => unreachable!(),
                    _ => acl.is_default_in_proxy_list(),
                };
                if !forward {
                    return false;
                }
            };
        }
        if let Ok(ref local_response) = local_response {
            for rec in local_response.answers() {
                examine_record!(rec, true);
            }
            for rec in local_response.additionals() {
                examine_record!(rec, false);
            }
        }
        true
    } else {
        unreachable!()
    }
}

async fn acl_lookup<Local, Remote>(
    acl: &Option<AccessControl>,
    local: Arc<Local>,
    remote: Arc<Remote>,
    query: &Query
) -> (io::Result<Message>, bool)
    where
        Local: upstream::Upstream,
        Remote: upstream::Upstream,
{
    // Start querying name servers
    debug!(
        "attempting lookup of {:?} {} with ns {:?} and {:?}",
        query.query_type(), query.name(), local, remote
    );

    let remote_response_fut = try_timeout(remote.lookup(query), Some(Duration::new(3, 0)));
    let local_response_fut = try_timeout(local.lookup(query), Some(Duration::new(3, 0)));

    match should_forward_by_query(acl, query) {
        Some(true) => {
            let remote_response = remote_response_fut.await;
            debug!("pick remote response (query): {:?}", remote_response);
            return (remote_response, true);
        }
        Some(false) => {
            let local_response = local_response_fut.await;
            debug!("pick local response (query): {:?}", local_response);
            return (local_response, false);
        }
        None => (),
    }

    let decider = async {
        let local_response = local_response_fut.await;
        if should_forward_by_response(acl, &local_response, query) {
            None
        } else {
            Some(local_response)
        }
    };
    tokio::pin!(remote_response_fut, decider);
    let mut use_remote = false;
    let mut remote_response = None;
    loop {
        tokio::select! {
            response = &mut remote_response_fut => {
                if use_remote {
                    debug!("pick remote response (response): {:?}", response);
                    return (response, true);
                } else {
                    remote_response = Some(response);
                }
            }
            decision = &mut decider => {
                if let Some(local_response) = decision {
                    debug!("pick local response (response): {:?}", local_response);
                    return (local_response, false);
                } else if let Some(remote_response) = remote_response {
                    debug!("pick remote response (response): {:?}", remote_response);
                    return (remote_response, true);
                } else {
                    use_remote = true;
                }
            }
        }
    }
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = match context.config().config_type {
        ConfigType::DnsLocal => {
            // Standalone server
            context.config().local_addr.as_ref().expect("local config")
        }
        c if c.is_local() => {
            // Integrated mode
            context.config().dns_local_addr.as_ref().expect("dns relay addr")
        }
        _ => {
            panic!("ConfigType must be DnsLocal");
        }
    };

    let bind_addr = local_addr.bind_addr(&context).await?;

    let socket = create_udp_socket(&bind_addr).await?;

    let actual_local_addr = socket.local_addr()?;
    info!("shadowsocks DNS relay listening on {}", actual_local_addr);

    let (mut rx, mut tx) = socket.split();
    let (qtx, mut qrx) = mpsc::channel::<(SocketAddr, Vec<u8>)>(1024);

    tokio::spawn(async move {
        while let Some((src, pkt)) = qrx.recv().await {
            if let Err(err) = tx.send_to(&pkt, &src).await {
                error!("failed to send packet {} bytes to {}, error: {}", pkt.len(), src, err);
            }
        }
    });

    let config = context.config();
    #[cfg(target_os = "android")]
        let local_upstream = Arc::new(upstream::UnixSocketUpstream {
        path: config.local_dns_path.clone().expect("local query DNS path"),
    });
    #[cfg(not(target_os = "android"))]
        let local_upstream = Arc::new(upstream::UdpUpstream {
        server: config.local_dns_addr.clone().expect("local query DNS address"),
    });
    // FIXME: We use TCP to send remote queries by default, which should be configuable.
    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await;
    let remote_upstream = Arc::new(upstream::ProxyTcpUpstream {
        context: context.clone(),
        svr_cfg: move || balancer.pick_server().server_config().clone(),
        ns: config.remote_dns_addr.clone().expect("remote query DNS address"),
    });

    loop {
        let mut req_buffer: [u8; 512] = [0; 512];
        let (_, src) = match rx.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                error!("DNS relay read from UDP socket error: {}", e);
                continue;
            }
        };

        let request = match Message::from_vec(&req_buffer) {
            Ok(x) => x,
            Err(e) => {
                error!("failed to parse UDP query message, error: {:?}", e);
                continue;
            }
        };

        debug!("received src: {}, query: {:?}", src, request);

        let context = context.clone();
        let local_upstream = Arc::clone(&local_upstream);
        let remote_upstream = Arc::clone(&remote_upstream);
        let mut qtx = qtx.clone();

        tokio::spawn(async move {
            let mut message = Message::new();
            message.set_id(request.id());
            message.set_recursion_desired(true);
            message.set_recursion_available(true);
            message.set_message_type(MessageType::Response);

            if !request.recursion_desired() {
                message.set_recursion_desired(false);
                message.set_response_code(ResponseCode::NotImp);
            } else if request.query_count() > 0 {
                let question = &request.queries()[0];
                let (r, forward) = acl_lookup(context.acl(), local_upstream, remote_upstream, question).await;

                if let Ok(result) = r {
                    for rec in result.answers() {
                        debug!("dns answer: {:?}", rec);

                        match rec.rdata() {
                            RData::A(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::V4(*ip), forward),
                            RData::AAAA(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::V6(*ip), forward),
                            _ => (),
                        }
                    }

                    message = result;
                    message.set_id(request.id());
                } else {
                    message.set_response_code(ResponseCode::ServFail);
                }
            }

            debug!("DNS src: {}, final response: {:?}", src, message);

            match message.to_vec() {
                Err(err) => {
                    error!("failed to serialize message, error: {}", err);
                }
                Ok(res_buffer) => {
                    if let Err(..) = qtx.send((src, res_buffer)).await {
                        error!("DNS send back queue is closed unexpectly");
                    }
                }
            }
        });
    }
}
<commit_msg>Prevent resuming completed Futures<commit_after>use std::{
    io,
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use tokio::sync::mpsc;

use log::{debug, error, info, warn};
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, Query},
    rr::{DNSClass, Name, RData, RecordType},
};

use crate::{
    acl::AccessControl,
    config::ConfigType,
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

mod upstream;

fn should_forward_by_ptr_name(acl: &AccessControl, name: &Name) -> bool {
    let mut iter = name.iter().rev();
    let mut next = || std::str::from_utf8(iter.next().unwrap_or(&[48])).unwrap_or("*");
    if !"arpa".eq_ignore_ascii_case(next()) {
        return acl.is_default_in_proxy_list();
    }
    match &next().to_ascii_lowercase()[..] {
        "in-addr" => {
            let mut octets: [u8; 4] = [0; 4];
            for octet in octets.iter_mut() {
                match next().parse() {
                    Ok(result) => *octet = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V4(Ipv4Addr::new(octets[0], octets[1], octets[2], octets[3])))
        }
        "ip6" => {
            let mut segments: [u16; 8] = [0; 8];
            for segment in segments.iter_mut() {
                match u16::from_str_radix(&[next(), next(), next(), next()].concat(), 16) {
                    Ok(result) => *segment = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V6(Ipv6Addr::new(
                segments[0], segments[1], segments[2], segments[3], segments[4], segments[5], segments[6], segments[7]
            )))
        }
        _ => acl.is_default_in_proxy_list(),
    }
}

/// given the query, determine whether remote/local query should be used, or inconclusive
fn should_forward_by_query(acl: &Option<AccessControl>, query: &Query) -> Option<bool> {
    if let Some(acl) = acl {
        if query.query_class() != DNSClass::IN {
            // unconditionally use default for all non-IN queries
            Some(acl.is_default_in_proxy_list())
        } else if query.query_type() == RecordType::PTR {
            Some(should_forward_by_ptr_name(acl, query.name()))
        } else {
            let result = acl.check_name_in_proxy_list(query.name());
            if result == None && match query.query_type() {
                RecordType::A => acl.is_ipv4_empty(),
                RecordType::AAAA => acl.is_ipv6_empty(),
                RecordType::ANY => acl.is_ipv4_empty() && acl.is_ipv6_empty(),
                RecordType::PTR => unreachable!(),
                _ => true,
            } {
                Some(acl.is_default_in_proxy_list())
            } else {
                result
            }
        }
    } else {
        Some(true)
    }
}

/// given the local response, determine whether remote response should be used instead
fn should_forward_by_response(
    acl: &Option<AccessControl>,
    local_response: &io::Result<Message>,
    query: &Query,
) -> bool {
    if let Some(acl) = acl {
        macro_rules! examine_record {
            ($rec:ident, $is_answer:expr) => {
                if let RData::CNAME(ref name) = $rec.rdata() {
                    match acl.check_name_in_proxy_list(name) {
                        Some(value) => return value,
                        None => continue,
                    }
                } else if $is_answer && !query.query_type().is_any() && $rec.record_type() != query.query_type() {
                    warn!("local DNS response has inconsistent answer type {} for query {}", $rec.record_type(), query);
                    return true;
                }
                let forward = match $rec.rdata() {
                    RData::A(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V4(*ip)),
                    RData::AAAA(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V6(*ip)),
                    RData::PTR(_) => unreachable!(),
                    _ => acl.is_default_in_proxy_list(),
                };
                if !forward {
                    return false;
                }
            };
        }
        if let Ok(ref local_response) = local_response {
            for rec in local_response.answers() {
                examine_record!(rec, true);
            }
            for rec in local_response.additionals() {
                examine_record!(rec, false);
            }
        }
        true
    } else {
        unreachable!()
    }
}

async fn acl_lookup<Local, Remote>(
    acl: &Option<AccessControl>,
    local: Arc<Local>,
    remote: Arc<Remote>,
    query: &Query
) -> (io::Result<Message>, bool)
    where
        Local: upstream::Upstream,
        Remote: upstream::Upstream,
{
    // Start querying name servers
    debug!(
        "attempting lookup of {:?} {} with ns {:?} and {:?}",
        query.query_type(), query.name(), local, remote
    );

    let remote_response_fut = try_timeout(remote.lookup(query), Some(Duration::new(3, 0)));
    let local_response_fut = try_timeout(local.lookup(query), Some(Duration::new(3, 0)));

    match should_forward_by_query(acl, query) {
        Some(true) => {
            let remote_response = remote_response_fut.await;
            debug!("pick remote response (query): {:?}", remote_response);
            return (remote_response, true);
        }
        Some(false) => {
            let local_response = local_response_fut.await;
            debug!("pick local response (query): {:?}", local_response);
            return (local_response, false);
        }
        None => (),
    }

    let decider = async {
        let local_response = local_response_fut.await;
        if should_forward_by_response(acl, &local_response, query) {
            None
        } else {
            Some(local_response)
        }
    };
    tokio::pin!(remote_response_fut, decider);
    let mut use_remote = false;
    let mut remote_response = None;
    loop {
        tokio::select! {
            response = &mut remote_response_fut, if remote_response.is_none() => {
                if use_remote {
                    debug!("pick remote response (response): {:?}", response);
                    return (response, true);
                } else {
                    remote_response = Some(response);
                }
            }
            decision = &mut decider, if !use_remote => {
                if let Some(local_response) = decision {
                    debug!("pick local response (response): {:?}", local_response);
                    return (local_response, false);
                } else if let Some(remote_response) = remote_response {
                    debug!("pick remote response (response): {:?}", remote_response);
                    return (remote_response, true);
                } else {
                    use_remote = true;
                }
            }
        }
    }
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = match context.config().config_type {
        ConfigType::DnsLocal => {
            // Standalone server
            context.config().local_addr.as_ref().expect("local config")
        }
        c if c.is_local() => {
            // Integrated mode
            context.config().dns_local_addr.as_ref().expect("dns relay addr")
        }
        _ => {
            panic!("ConfigType must be DnsLocal");
        }
    };

    let bind_addr = local_addr.bind_addr(&context).await?;

    let socket = create_udp_socket(&bind_addr).await?;

    let actual_local_addr = socket.local_addr()?;
    info!("shadowsocks DNS relay listening on {}", actual_local_addr);

    let (mut rx, mut tx) = socket.split();
    let (qtx, mut qrx) = mpsc::channel::<(SocketAddr, Vec<u8>)>(1024);

    tokio::spawn(async move {
        while let Some((src, pkt)) = qrx.recv().await {
            if let Err(err) = tx.send_to(&pkt, &src).await {
                error!("failed to send packet {} bytes to {}, error: {}", pkt.len(), src, err);
            }
        }
    });

    let config = context.config();
    #[cfg(target_os = "android")]
        let local_upstream = Arc::new(upstream::UnixSocketUpstream {
        path: config.local_dns_path.clone().expect("local query DNS path"),
    });
    #[cfg(not(target_os = "android"))]
        let local_upstream = Arc::new(upstream::UdpUpstream {
        server: config.local_dns_addr.clone().expect("local query DNS address"),
    });
    // FIXME: We use TCP to send remote queries by default, which should be configuable.
    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await;
    let remote_upstream = Arc::new(upstream::ProxyTcpUpstream {
        context: context.clone(),
        svr_cfg: move || balancer.pick_server().server_config().clone(),
        ns: config.remote_dns_addr.clone().expect("remote query DNS address"),
    });

    loop {
        let mut req_buffer: [u8; 512] = [0; 512];
        let (_, src) = match rx.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                error!("DNS relay read from UDP socket error: {}", e);
                continue;
            }
        };

        let request = match Message::from_vec(&req_buffer) {
            Ok(x) => x,
            Err(e) => {
                error!("failed to parse UDP query message, error: {:?}", e);
                continue;
            }
        };

        debug!("received src: {}, query: {:?}", src, request);

        let context = context.clone();
        let local_upstream = Arc::clone(&local_upstream);
        let remote_upstream = Arc::clone(&remote_upstream);
        let mut qtx = qtx.clone();

        tokio::spawn(async move {
            let mut message = Message::new();
            message.set_id(request.id());
            message.set_recursion_desired(true);
            message.set_recursion_available(true);
            message.set_message_type(MessageType::Response);

            if !request.recursion_desired() {
                message.set_recursion_desired(false);
                message.set_response_code(ResponseCode::NotImp);
            } else if request.query_count() > 0 {
                let question = &request.queries()[0];
                let (r, forward) = acl_lookup(context.acl(), local_upstream, remote_upstream, question).await;

                if let Ok(result) = r {
                    for rec in result.answers() {
                        debug!("dns answer: {:?}", rec);

                        match rec.rdata() {
                            RData::A(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::V4(*ip), forward),
                            RData::AAAA(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::V6(*ip), forward),
                            _ => (),
                        }
                    }

                    message = result;
                    message.set_id(request.id());
                } else {
                    message.set_response_code(ResponseCode::ServFail);
                }
            }

            debug!("DNS src: {}, final response: {:?}", src, message);

            match message.to_vec() {
                Err(err) => {
                    error!("failed to serialize message, error: {}", err);
                }
                Ok(res_buffer) => {
                    if let Err(..) = qtx.send((src, res_buffer)).await {
                        error!("DNS send back queue is closed unexpectly");
                    }
                }
            }
        });
    }
}
<|endoftext|>


<commit_before>//! Android specific features

use std::{
    convert::TryInto,
    io::{self, Error, ErrorKind, Read, Write},
    mem::{self, MaybeUninit},
    net::Shutdown,
    os::unix::io::{AsRawFd, RawFd},
    path::Path,
    pin::Pin,
    ptr,
    slice,
    task::{Context, Poll},
};

use futures::{future, ready};
use mio::net::UnixStream as MioUnixStream;
use tokio::io::{AsyncRead, AsyncWrite, PollEvented};

/// A UnixStream supports transferring FDs between processes
pub struct UnixStream {
    io: PollEvented<MioUnixStream>,
}

impl UnixStream {
    /// Connects to the socket named by `path`.
    pub async fn connect<P: AsRef<Path>>(path: P) -> io::Result<UnixStream> {
        let uds = MioUnixStream::connect(path)?;
        let io = PollEvented::new(uds)?;

        future::poll_fn(|cx| io.poll_write_ready(cx)).await?;
        Ok(UnixStream { io })
    }

    fn poll_send_with_fd(&self, cx: &mut Context, buf: &[u8], fds: &[RawFd]) -> Poll<io::Result<usize>> {
        ready!(self.io.poll_write_ready(cx))?;

        let fd = self.io.get_ref().as_raw_fd();
        match send_with_fd(fd, buf, fds) {
            Err(ref err) if err.kind() == ErrorKind::WouldBlock => Poll::Pending,
            x => Poll::Ready(x),
        }
    }

    /// Send data with file descriptors
    pub async fn send_with_fd(&mut self, buf: &[u8], fds: &[RawFd]) -> io::Result<usize> {
        future::poll_fn(|cx| self.poll_send_with_fd(cx, buf, fds)).await
    }

    /// Shuts down the read, write, or both halves of this connection.
    ///
    /// This function will cause all pending and future I/O calls on the
    /// specified portions to immediately return with an appropriate value
    /// (see the documentation of `Shutdown`).
    pub fn shutdown(&self, how: Shutdown) -> io::Result<()> {
        self.io.get_ref().shutdown(how)
    }
}

impl AsyncRead for UnixStream {
    unsafe fn prepare_uninitialized_buffer(&self, _: &mut [MaybeUninit<u8>]) -> bool {
        false
    }

    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
        self.poll_read_priv(cx, buf)
    }
}

impl AsyncWrite for UnixStream {
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        self.poll_write_priv(cx, buf)
    }

    fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
        Poll::Ready(Ok(()))
    }

    fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
        self.shutdown(Shutdown::Write)?;
        Poll::Ready(Ok(()))
    }
}

impl UnixStream {
    // == Poll IO functions that takes `&self` ==
    //
    // They are not public because (taken from the doc of `PollEvented`):
    //
    // While `PollEvented` is `Sync` (if the underlying I/O type is `Sync`), the
    // caller must ensure that there are at most two tasks that use a
    // `PollEvented` instance concurrently. One for reading and one for writing.
    // While violating this requirement is "safe" from a Rust memory model point
    // of view, it will result in unexpected behavior in the form of lost
    // notifications and tasks hanging.

    pub(crate) fn poll_read_priv(&self, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
        ready!(self.io.poll_read_ready(cx, mio::Ready::readable()))?;

        match self.io.get_ref().read(buf) {
            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                self.io.clear_read_ready(cx, mio::Ready::readable())?;
                Poll::Pending
            }
            x => Poll::Ready(x),
        }
    }

    pub(crate) fn poll_write_priv(&self, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        ready!(self.io.poll_write_ready(cx))?;

        match self.io.get_ref().write(buf) {
            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                self.io.clear_write_ready(cx)?;
                Poll::Pending
            }
            x => Poll::Ready(x),
        }
    }
}

/// A common implementation of `sendmsg` that sends provided bytes with ancillary file descriptors
/// over either a datagram or stream unix socket.
///
/// Borrowed from: https://github.com/Standard-Cognition/sendfd
fn send_with_fd(socket: RawFd, bs: &[u8], fds: &[RawFd]) -> io::Result<usize> {
    unsafe {
        let mut iov = libc::iovec {
            // NB: this casts *const to *mut, and in doing so we trust the OS to be a good citizen
            // and not mutate our buffer. This is the API we have to live with.
            iov_base: bs.as_ptr() as *const _ as *mut _,
            iov_len: bs.len(),
        };

        // Construct msghdr
        //
        // 1. Allocate memory for msg_control
        let cmsg_fd_len = fds.len() * mem::size_of::<RawFd>();
        let cmsg_buffer_len = libc::CMSG_SPACE(cmsg_fd_len as u32) as usize;
        let mut cmsg_buffer = Vec::with_capacity(cmsg_buffer_len);
        cmsg_buffer.set_len(cmsg_buffer_len);

        let mut msghdr = libc::msghdr {
            msg_name: ptr::null_mut(),
            msg_namelen: 0,
            msg_iov: &mut iov as *mut _,
            msg_iovlen: 1,
            msg_control: cmsg_buffer.as_mut_ptr(),
            msg_controllen: cmsg_buffer_len.try_into().unwrap(),
            ..mem::zeroed()
        };

        // Fill cmsg with the file descriptors we are sending.
        let cmsg_header = libc::CMSG_FIRSTHDR(&mut msghdr as *mut _);
        cmsg_header.write(libc::cmsghdr {
            cmsg_level: libc::SOL_SOCKET,
            cmsg_type: libc::SCM_RIGHTS,
            cmsg_len: libc::CMSG_LEN(cmsg_fd_len as u32).try_into().unwrap(),
        });

        let cmsg_data = libc::CMSG_DATA(cmsg_header);
        #[allow(clippy::cast_ptr_alignment)] // false positive
        let cmsg_data_slice = slice::from_raw_parts_mut(cmsg_data as *mut RawFd, fds.len());
        cmsg_data_slice.copy_from_slice(fds);

        let count = libc::sendmsg(socket, &msghdr as *const _, 0);
        if count < 0 {
            let err = Error::last_os_error();
            Err(err)
        } else {
            Ok(count as usize)
        }
    }
}
<commit_msg>Make uds compatible with tokio3<commit_after>//! Android specific features

use std::{
    convert::TryInto,
    io::{self, Error, ErrorKind, Read, Write},
    mem,
    net::Shutdown,
    os::unix::io::{AsRawFd, RawFd},
    path::Path,
    pin::Pin,
    ptr,
    slice,
    task::{Context, Poll},
};

use futures::{future, ready};
use mio::net::UnixStream as MioUnixStream;
use tokio::io::{unix::AsyncFd, AsyncRead, AsyncWrite, ReadBuf};

/// A UnixStream supports transferring FDs between processes
pub struct UnixStream {
    io: AsyncFd<MioUnixStream>,
}

impl UnixStream {
    /// Connects to the socket named by `path`.
    pub async fn connect<P: AsRef<Path>>(path: P) -> io::Result<UnixStream> {
        let uds = MioUnixStream::connect(path)?;
        let io = AsyncFd::new(uds)?;

        future::poll_fn(|cx| io.poll_write_ready(cx)).await?;
        Ok(UnixStream { io })
    }

    fn poll_send_with_fd(&self, cx: &mut Context, buf: &[u8], fds: &[RawFd]) -> Poll<io::Result<usize>> {
        ready!(self.io.poll_write_ready(cx))?;

        let fd = self.io.get_ref().as_raw_fd();
        match send_with_fd(fd, buf, fds) {
            Err(ref err) if err.kind() == ErrorKind::WouldBlock => Poll::Pending,
            x => Poll::Ready(x),
        }
    }

    /// Send data with file descriptors
    pub async fn send_with_fd(&mut self, buf: &[u8], fds: &[RawFd]) -> io::Result<usize> {
        future::poll_fn(|cx| self.poll_send_with_fd(cx, buf, fds)).await
    }

    /// Shuts down the read, write, or both halves of this connection.
    ///
    /// This function will cause all pending and future I/O calls on the
    /// specified portions to immediately return with an appropriate value
    /// (see the documentation of `Shutdown`).
    pub fn shutdown(&self, how: Shutdown) -> io::Result<()> {
        self.io.get_ref().shutdown(how)
    }
}

impl AsyncRead for UnixStream {
    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        self.poll_read_priv(cx, buf)
    }
}

impl AsyncWrite for UnixStream {
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        self.poll_write_priv(cx, buf)
    }

    fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
        Poll::Ready(Ok(()))
    }

    fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
        self.shutdown(Shutdown::Write)?;
        Poll::Ready(Ok(()))
    }
}

impl UnixStream {
    // == Poll IO functions that takes `&self` ==
    //
    // They are not public because (taken from the doc of `PollEvented`):
    //
    // While `PollEvented` is `Sync` (if the underlying I/O type is `Sync`), the
    // caller must ensure that there are at most two tasks that use a
    // `PollEvented` instance concurrently. One for reading and one for writing.
    // While violating this requirement is "safe" from a Rust memory model point
    // of view, it will result in unexpected behavior in the form of lost
    // notifications and tasks hanging.

    pub(crate) fn poll_read_priv(&self, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        let mut read_guard = ready!(self.io.poll_read_ready(cx))?;

        let b = unsafe { &mut *(buf.unfilled_mut() as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]) };
        match self.io.get_ref().read(b) {
            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                read_guard.clear_ready();
                Poll::Pending
            }
            Ok(n) => {
                // Safety: We trust `UnixStream::read` to have filled up `n` bytes
                // in the buffer.
                unsafe {
                    buf.assume_init(n);
                }
                buf.advance(n);
                return Poll::Ready(Ok(()));
            }
            Err(e) => return Poll::Ready(Err(e)),
        }
    }

    pub(crate) fn poll_write_priv(&self, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        let mut write_guard = ready!(self.io.poll_write_ready(cx))?;

        match self.io.get_ref().write(buf) {
            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                write_guard.clear_ready();
                Poll::Pending
            }
            x => Poll::Ready(x),
        }
    }
}

/// A common implementation of `sendmsg` that sends provided bytes with ancillary file descriptors
/// over either a datagram or stream unix socket.
///
/// Borrowed from: https://github.com/Standard-Cognition/sendfd
fn send_with_fd(socket: RawFd, bs: &[u8], fds: &[RawFd]) -> io::Result<usize> {
    unsafe {
        let mut iov = libc::iovec {
            // NB: this casts *const to *mut, and in doing so we trust the OS to be a good citizen
            // and not mutate our buffer. This is the API we have to live with.
            iov_base: bs.as_ptr() as *const _ as *mut _,
            iov_len: bs.len(),
        };

        // Construct msghdr
        //
        // 1. Allocate memory for msg_control
        let cmsg_fd_len = fds.len() * mem::size_of::<RawFd>();
        let cmsg_buffer_len = libc::CMSG_SPACE(cmsg_fd_len as u32) as usize;
        let mut cmsg_buffer = Vec::with_capacity(cmsg_buffer_len);
        cmsg_buffer.set_len(cmsg_buffer_len);

        let mut msghdr = libc::msghdr {
            msg_name: ptr::null_mut(),
            msg_namelen: 0,
            msg_iov: &mut iov as *mut _,
            msg_iovlen: 1,
            msg_control: cmsg_buffer.as_mut_ptr(),
            msg_controllen: cmsg_buffer_len.try_into().unwrap(),
            ..mem::zeroed()
        };

        // Fill cmsg with the file descriptors we are sending.
        let cmsg_header = libc::CMSG_FIRSTHDR(&mut msghdr as *mut _);
        cmsg_header.write(libc::cmsghdr {
            cmsg_level: libc::SOL_SOCKET,
            cmsg_type: libc::SCM_RIGHTS,
            cmsg_len: libc::CMSG_LEN(cmsg_fd_len as u32).try_into().unwrap(),
        });

        let cmsg_data = libc::CMSG_DATA(cmsg_header);
        #[allow(clippy::cast_ptr_alignment)] // false positive
        let cmsg_data_slice = slice::from_raw_parts_mut(cmsg_data as *mut RawFd, fds.len());
        cmsg_data_slice.copy_from_slice(fds);

        let count = libc::sendmsg(socket, &msghdr as *const _, 0);
        if count < 0 {
            let err = Error::last_os_error();
            Err(err)
        } else {
            Ok(count as usize)
        }
    }
}
<|endoftext|>