<commit_before>use structs::{CliError, Config, UserConfig, Member, Team, TeamGroup, PublicKey};
use std;
use std::fs::File;
use std::io::prelude::*;
use std::path::Path;
use std::collections::HashMap;

use glob::glob;
use reqwest;
use reqwest::header::Authorization;
use serde_json;

pub struct GithubClient {
    client: reqwest::Client,
    conf: Config,
}

impl GithubClient {
    pub fn new(config: &Config) -> Result<GithubClient, CliError> {
        if std::env::var("SSL_CERT_FILE").is_err() {
            std::env::set_var("SSL_CERT_FILE", &config.cert_path);
        }
        let client = reqwest::Client::new()?;
        Ok(GithubClient { client: client, conf: config.clone() })
    }

    fn get_cache_path(url: &str) -> std::path::PathBuf {
        let mut path = std::env::temp_dir();
        path.push("ghteam-auth-cache");
        path.push(url);
        path
    }

    fn load_content_from_cache(&self, url: &str) -> Result<(std::fs::Metadata, String), CliError> {
        let path = Self::get_cache_path(url);
        let metadata = std::fs::metadata(path.to_str().unwrap())?;
        let mut f = File::open(path.to_str().unwrap())?;
        let mut content = String::new();
        f.read_to_string(&mut content)?;
        Ok((metadata, content))
    }

    fn store_content_to_cache(&self, url: &str, content: &str) -> Result<(), CliError> {
        let path = Self::get_cache_path(url);
        std::fs::create_dir_all(path.parent().unwrap())?;
        let mut f = File::create(path.to_str().unwrap())?;
        f.write(content.as_bytes())?;
        Ok(())
    }

    fn get_content_from_url(&self, url: &str) -> Result<String, CliError> {
        let token = String::from("token ") + &self.conf.token;
        let res = self.client.get(url).header(Authorization(token)).send();
        let mut content = String::new();
        res?.read_to_string(&mut content)?;
        self.store_content_to_cache(url, &content)?;
        Ok(content)
    }

    fn get_content(&self, url: &str) -> Result<String, CliError> {
        match self.load_content_from_cache(url) {
            Ok((metadata, cache_content)) => {
                match std::time::SystemTime::now().duration_since(metadata.modified()?) {
                    Ok(caching_duration) => {
                        if caching_duration.as_secs() > self.conf.cache_duration {
                            match self.get_content_from_url(url) {
                                Ok(content_from_url) => Ok(content_from_url),
                                Err(_) => Ok(cache_content),
                            }
                        } else {
                            Ok(cache_content)
                        }
                    }
                    Err(_) => Ok(cache_content),
                }
            }
            Err(_) => self.get_content_from_url(url),
        }
    }

    #[allow(dead_code)]
    pub fn print_user_public_key(&self, user: &str) -> Result<(), CliError> {
        let keys = self.get_user_public_key(user)?;
        println!("{}", keys);
        Ok(())
    }

    fn get_user_public_key(&self, user: &str) -> Result<String, CliError> {
        let url = format!("{}/users/{}/keys", self.conf.endpoint, user);
        let content = self.get_content(&url)?;
        let keys = serde_json::from_str::<Vec<PublicKey>>(&content)?;
        Ok(keys.iter().map(|k| k.key.clone()).collect::<Vec<String>>().join("\n"))
    }

    #[allow(dead_code)]
    pub fn check_pam(&self, user: &str) -> Result<bool, CliError> {
        let teams = self.get_teams();
        Ok(teams.iter().any(|team| team.members.contains_key(user)))
    }

    pub fn get_teams(&self) -> Vec<TeamGroup> { self.get_teams_result().unwrap_or(Vec::new()) }

    fn get_teams_result(&self) -> Result<Vec<TeamGroup>, CliError> {
        let ghteams = self.get_ghteam_map()?;
        let mut teams = Vec::new();
        for team_conf in &self.conf.team {
            if let &Some(ghteam) = &ghteams.get(&team_conf.name) {
                teams.push(TeamGroup { team: ghteam.clone(),
                                       gid: team_conf.gid.clone(),
                                       group: team_conf.group.clone(),
                                       members: self.get_members(ghteam.id)?, });
            }
        }
        Ok(teams)
    }

    fn get_ghteam_map(&self) -> Result<HashMap<String, Team>, CliError> {
        let url = format!("{}/orgs/{}/teams", self.conf.endpoint, self.conf.org);
        let content = self.get_content(&url)?;
        let teams = serde_json::from_str::<Vec<Team>>(&content)?;
        Ok(teams.iter().map(|t| (t.name.clone(), t.clone())).collect())
    }

    fn get_members(&self, mid: u64) -> Result<HashMap<String, Member>, CliError> {
        let url = format!("{}/teams/{}/members", self.conf.endpoint, mid);
        let content = self.get_content(&url)?;
        let members = serde_json::from_str::<Vec<Member>>(&content)?;
        Ok(members.iter().map(|m| (m.login.clone(), m.clone())).collect())
    }

    fn get_user_conf(&self, user: &str) -> Result<UserConfig, CliError> {
        let home = self.conf.home.replace("{}", user);
        UserConfig::new(&Path::new(&home).join(&self.conf.user_conf_path))
    }

    #[allow(dead_code)]
    pub fn clear_all_caches(&self) -> Result<(), CliError> {
        let mut path = std::env::temp_dir();
        path.push("ghteam-auth-cache");
        path.push("**");
        path.push("*");
        for entry in glob(&path.to_str().unwrap()).unwrap() {
            match entry {
                Ok(path) => {
                    if path.is_file() {
                        std::fs::remove_file(path)?
                    }
                }
                Err(e) => println!("{:?}", e),
            }
        }
        Ok(())
    }
}
<commit_msg>Remove an unused method<commit_after>use structs::{CliError, Config, Member, Team, TeamGroup, PublicKey};
use std;
use std::fs::File;
use std::io::prelude::*;
use std::collections::HashMap;

use glob::glob;
use reqwest;
use reqwest::header::Authorization;
use serde_json;

pub struct GithubClient {
    client: reqwest::Client,
    conf: Config,
}

impl GithubClient {
    pub fn new(config: &Config) -> Result<GithubClient, CliError> {
        if std::env::var("SSL_CERT_FILE").is_err() {
            std::env::set_var("SSL_CERT_FILE", &config.cert_path);
        }
        let client = reqwest::Client::new()?;
        Ok(GithubClient { client: client, conf: config.clone() })
    }

    fn get_cache_path(url: &str) -> std::path::PathBuf {
        let mut path = std::env::temp_dir();
        path.push("ghteam-auth-cache");
        path.push(url);
        path
    }

    fn load_content_from_cache(&self, url: &str) -> Result<(std::fs::Metadata, String), CliError> {
        let path = Self::get_cache_path(url);
        let metadata = std::fs::metadata(path.to_str().unwrap())?;
        let mut f = File::open(path.to_str().unwrap())?;
        let mut content = String::new();
        f.read_to_string(&mut content)?;
        Ok((metadata, content))
    }

    fn store_content_to_cache(&self, url: &str, content: &str) -> Result<(), CliError> {
        let path = Self::get_cache_path(url);
        std::fs::create_dir_all(path.parent().unwrap())?;
        let mut f = File::create(path.to_str().unwrap())?;
        f.write(content.as_bytes())?;
        Ok(())
    }

    fn get_content_from_url(&self, url: &str) -> Result<String, CliError> {
        let token = String::from("token ") + &self.conf.token;
        let res = self.client.get(url).header(Authorization(token)).send();
        let mut content = String::new();
        res?.read_to_string(&mut content)?;
        self.store_content_to_cache(url, &content)?;
        Ok(content)
    }

    fn get_content(&self, url: &str) -> Result<String, CliError> {
        match self.load_content_from_cache(url) {
            Ok((metadata, cache_content)) => {
                match std::time::SystemTime::now().duration_since(metadata.modified()?) {
                    Ok(caching_duration) => {
                        if caching_duration.as_secs() > self.conf.cache_duration {
                            match self.get_content_from_url(url) {
                                Ok(content_from_url) => Ok(content_from_url),
                                Err(_) => Ok(cache_content),
                            }
                        } else {
                            Ok(cache_content)
                        }
                    }
                    Err(_) => Ok(cache_content),
                }
            }
            Err(_) => self.get_content_from_url(url),
        }
    }

    #[allow(dead_code)]
    pub fn print_user_public_key(&self, user: &str) -> Result<(), CliError> {
        let keys = self.get_user_public_key(user)?;
        println!("{}", keys);
        Ok(())
    }

    fn get_user_public_key(&self, user: &str) -> Result<String, CliError> {
        let url = format!("{}/users/{}/keys", self.conf.endpoint, user);
        let content = self.get_content(&url)?;
        let keys = serde_json::from_str::<Vec<PublicKey>>(&content)?;
        Ok(keys.iter().map(|k| k.key.clone()).collect::<Vec<String>>().join("\n"))
    }

    #[allow(dead_code)]
    pub fn check_pam(&self, user: &str) -> Result<bool, CliError> {
        let teams = self.get_teams();
        Ok(teams.iter().any(|team| team.members.contains_key(user)))
    }

    pub fn get_teams(&self) -> Vec<TeamGroup> { self.get_teams_result().unwrap_or(Vec::new()) }

    fn get_teams_result(&self) -> Result<Vec<TeamGroup>, CliError> {
        let ghteams = self.get_ghteam_map()?;
        let mut teams = Vec::new();
        for team_conf in &self.conf.team {
            if let &Some(ghteam) = &ghteams.get(&team_conf.name) {
                teams.push(TeamGroup { team: ghteam.clone(),
                                       gid: team_conf.gid.clone(),
                                       group: team_conf.group.clone(),
                                       members: self.get_members(ghteam.id)?, });
            }
        }
        Ok(teams)
    }

    fn get_ghteam_map(&self) -> Result<HashMap<String, Team>, CliError> {
        let url = format!("{}/orgs/{}/teams", self.conf.endpoint, self.conf.org);
        let content = self.get_content(&url)?;
        let teams = serde_json::from_str::<Vec<Team>>(&content)?;
        Ok(teams.iter().map(|t| (t.name.clone(), t.clone())).collect())
    }

    fn get_members(&self, mid: u64) -> Result<HashMap<String, Member>, CliError> {
        let url = format!("{}/teams/{}/members", self.conf.endpoint, mid);
        let content = self.get_content(&url)?;
        let members = serde_json::from_str::<Vec<Member>>(&content)?;
        Ok(members.iter().map(|m| (m.login.clone(), m.clone())).collect())
    }

    #[allow(dead_code)]
    pub fn clear_all_caches(&self) -> Result<(), CliError> {
        let mut path = std::env::temp_dir();
        path.push("ghteam-auth-cache");
        path.push("**");
        path.push("*");
        for entry in glob(&path.to_str().unwrap()).unwrap() {
            match entry {
                Ok(path) => {
                    if path.is_file() {
                        std::fs::remove_file(path)?
                    }
                }
                Err(e) => println!("{:?}", e),
            }
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>use resize::Pixel::Gray8;
use resize::Type::Triangle;
use resize::Type::Lanczos3;
use std::fs;
use std::path::Path;

fn load_png(mut data: &[u8]) -> (usize, usize, Vec<u8>) {
    let decoder = png::Decoder::new(&mut data);
    let (info, mut reader) = decoder.read_info().unwrap();
    let mut src = vec![0; info.buffer_size()];
    reader.next_frame(&mut src).unwrap();

    (info.width as usize, info.height as usize, src)
}

fn write_png(path: &Path, w2: usize, h2: usize, pixels: &[u8]) {
    let outfh = fs::File::create(path).unwrap();
    let encoder = png::Encoder::new(outfh, w2 as u32, h2 as u32);
    encoder.write_header().unwrap().write_image_data(pixels).unwrap();
}

fn img_diff(a: &[u8], b: &[u8]) -> f64 {
    assert_eq!(a.len(), b.len());
    let sum = a.iter().cloned().zip(b.iter().cloned()).map(|(a,b)| {
        (a as i32 - b as i32).pow(2) as u32
    }).sum::<u32>();
    sum as f64 / a.len() as f64
}

fn assert_equals(img: &[u8], w2: usize, h2: usize, expected_filename: &str) {
    let (_, _, expected) = load_png(&fs::read(&expected_filename).expect(expected_filename));

    let diff = img_diff(&img, &expected);
    if diff > 0.0004 {
        let bad_file = Path::new(expected_filename).with_extension("failed-test.png");
        write_png(&bad_file, w2, h2, img);
        panic!("Test failed: {} differs by {}; see {} ", expected_filename, diff, bad_file.display());
    }
}

#[test]
fn qa() {
    let tiger = &include_bytes!("../examples/tiger.png")[..];

    let (w1, h1, src) = load_png(&tiger);
    let mut res1 = vec![];
    let mut res2 = vec![];
    for w2 in [1, 2, 10, 100, 1000, 2000].iter().copied() {
        for h2 in [1, 2, 9, 99, 999, 1555].iter().copied() {
            res1.clear();
            res1.resize(w2 * h2, 0);

            resize::new(w1, h1, w2, h2, Gray8, Lanczos3).resize(&src, &mut res1);
            assert_equals(&res1, w2, h2, &format!("tests/t{}x{}.png", w2, h2));

            res2.clear();
            res2.resize(100 * 100, 255);

            resize::new(w2, h2, 100, 100, Gray8, Triangle).resize(&res1, &mut res2);
            assert_equals(&res2, 100, 100, &format!("tests/t{}x{}-100.png", w2, h2));
        }
    }
}
<commit_msg>Let cargo parallelize the tests<commit_after>use resize::Pixel::Gray8;
use resize::Type::Triangle;
use resize::Type::Lanczos3;
use std::fs;
use std::path::Path;

fn load_png(mut data: &[u8]) -> (usize, usize, Vec<u8>) {
    let decoder = png::Decoder::new(&mut data);
    let (info, mut reader) = decoder.read_info().unwrap();
    let mut src = vec![0; info.buffer_size()];
    reader.next_frame(&mut src).unwrap();

    (info.width as usize, info.height as usize, src)
}

fn write_png(path: &Path, w2: usize, h2: usize, pixels: &[u8]) {
    let outfh = fs::File::create(path).unwrap();
    let encoder = png::Encoder::new(outfh, w2 as u32, h2 as u32);
    encoder.write_header().unwrap().write_image_data(pixels).unwrap();
}

fn img_diff(a: &[u8], b: &[u8]) -> f64 {
    assert_eq!(a.len(), b.len());
    let sum = a.iter().cloned().zip(b.iter().cloned()).map(|(a,b)| {
        (a as i32 - b as i32).pow(2) as u32
    }).sum::<u32>();
    sum as f64 / a.len() as f64
}

fn assert_equals(img: &[u8], w2: usize, h2: usize, expected_filename: &str) {
    let (_, _, expected) = load_png(&fs::read(&expected_filename).expect(expected_filename));

    let diff = img_diff(&img, &expected);
    if diff > 0.0004 {
        let bad_file = Path::new(expected_filename).with_extension("failed-test.png");
        write_png(&bad_file, w2, h2, img);
        panic!("Test failed: {} differs by {}; see {} ", expected_filename, diff, bad_file.display());
    }
}

fn test_width(w2: usize) {
    let tiger = &include_bytes!("../examples/tiger.png")[..];
    let (w1, h1, src) = load_png(&tiger);
    let mut res1 = vec![];
    let mut res2 = vec![];

    for h2 in [1, 2, 9, 99, 999, 1555].iter().copied() {
        res1.clear();
        res1.resize(w2 * h2, 0);

        resize::new(w1, h1, w2, h2, Gray8, Lanczos3).resize(&src, &mut res1);
        assert_equals(&res1, w2, h2, &format!("tests/t{}x{}.png", w2, h2));

        res2.clear();
        res2.resize(100 * 100, 255);

        resize::new(w2, h2, 100, 100, Gray8, Triangle).resize(&res1, &mut res2);
        assert_equals(&res2, 100, 100, &format!("tests/t{}x{}-100.png", w2, h2));
    }
}

#[test]
fn test_w2000() {
    test_width(2000);
}
#[test]
fn test_w1000() {
    test_width(1000);
}
#[test]
fn test_w100() {
    test_width(100);
}
#[test]
fn test_w10() {
    test_width(10);
}
#[test]
fn test_w2() {
    test_width(2);
}
#[test]
fn test_w1() {
    test_width(1);
}
<|endoftext|>


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

use tokio::sync::mpsc;

#[cfg(not(target_os = "android"))]
use std::net::{Ipv4Addr, SocketAddrV4};
#[cfg(not(target_os = "android"))]
use tokio::net::UdpSocket;

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

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

mod upstream;

async fn acl_lookup<Local, Remote>(
    context: SharedContext,
    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 qname_in_proxy_list = context.check_query_in_proxy_list(query);
    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 qname_in_proxy_list {
        Some(true) => {
            let remote_response = remote_response_fut.await.unwrap_or_else(|_| Message::new());
            debug!("pick remote response (qname): {:?}", remote_response);
            return Ok((remote_response, true));
        }
        Some(false) => {
            let local_response = local_response_fut.await.unwrap_or_else(|_| Message::new());
            debug!("pick local response (qname): {:?}", local_response);
            return Ok((local_response, false));
        }
        None => (),
    }

    let local_response = local_response_fut.await.unwrap_or_else(|_| Message::new());
    for rec in local_response.answers() {
        if rec.record_type() != query.query_type() {
            warn!("local DNS response has inconsistent answer type {} for query {}", rec.record_type(), query);
            break
        }
        let forward = match rec.rdata() {
            RData::A(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            RData::AAAA(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            RData::PTR(_) => panic!("PTR records should not reach here"),
            _ => context.is_default_in_proxy_list(),
        };
        if !forward {
            debug!("pick local response (response): {:?}", local_response);
            return Ok((local_response, false));
        }
    }

    let remote_response = remote_response_fut.await.unwrap_or_else(|_| Message::new());
    debug!("pick remote response (response): {:?}", remote_response);
    Ok((remote_response, 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.queries().is_empty() {
                message.set_response_code(ResponseCode::FormErr);
            } else {
                let question = &request.queries()[0];
                let r = acl_lookup(context.clone(), local_upstream, remote_upstream, question).await;

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

                        match rec.rdata() {
                            RData::A(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*ip), forward),
                            RData::AAAA(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*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>Preserve errors in acl_lookup<commit_after>use std::{
    io,
    net::{IpAddr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use tokio::sync::mpsc;

#[cfg(not(target_os = "android"))]
use std::net::{Ipv4Addr, SocketAddrV4};
#[cfg(not(target_os = "android"))]
use tokio::net::UdpSocket;

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

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

mod upstream;

async fn acl_lookup<Local, Remote>(
    context: SharedContext,
    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 qname_in_proxy_list = context.check_query_in_proxy_list(query);
    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 qname_in_proxy_list {
        Some(true) => {
            let remote_response = remote_response_fut.await;
            debug!("pick remote response (qname): {:?}", remote_response);
            return (remote_response, true);
        }
        Some(false) => {
            let local_response = local_response_fut.await;
            debug!("pick local response (qname): {:?}", local_response);
            return (local_response, false);
        }
        None => (),
    }

    // FIXME: spawn(remote_response_fut)
    let local_response = local_response_fut.await;
    for rec in local_response.unwrap_or_else(|_| Message::new()).answers() {
        if rec.record_type() != query.query_type() {
            warn!("local DNS response has inconsistent answer type {} for query {}", rec.record_type(), query);
            break
        }
        let forward = match rec.rdata() {
            RData::A(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            RData::AAAA(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            RData::PTR(_) => panic!("PTR records should not reach here"),
            _ => context.is_default_in_proxy_list(),
        };
        if !forward {
            debug!("pick local response (response): {:?}", local_response);
            return (local_response, false);
        }
    }

    let remote_response = remote_response_fut.await;
    debug!("pick remote response (response): {:?}", remote_response);
    (remote_response, 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.queries().is_empty() {
                message.set_response_code(ResponseCode::FormErr);
            } else {
                let question = &request.queries()[0];
                let (r, forward) = acl_lookup(context.clone(), 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::from(*ip), forward),
                            RData::AAAA(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*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>//! User sessions.
//!
//! Actix provides a general solution for session management. Session
//! middlewares could provide different implementations which could
//! be accessed via general session api.
//!
//! By default, only cookie session backend is implemented. Other
//! backend implementations can be added.
//!
//! In general, you insert a *session* middleware and initialize it
//! , such as a `CookieSessionBackend`. To access session data,
//! [*Session*](struct.Session.html) extractor must be used. Session
//! extractor allows us to get or set session data.
//!
//! ```rust
//! use actix_web::{web, App, HttpServer, HttpResponse, Error};
//! use actix_session::{Session, CookieSession};
//!
//! fn index(session: Session) -> Result<&'static str, Error> {
//!     // access session data
//!     if let Some(count) = session.get::<i32>("counter")? {
//!         println!("SESSION value: {}", count);
//!         session.set("counter", count+1)?;
//!     } else {
//!         session.set("counter", 1)?;
//!     }
//!
//!     Ok("Welcome!")
//! }
//!
//! fn main() -> std::io::Result<()> {
//! # std::thread::spawn(||
//!     HttpServer::new(
//!         || App::new().wrap(
//!               CookieSession::signed(&[0; 32]) // <- create cookie based session middleware
//!                     .secure(false)
//!              )
//!             .service(web::resource("/").to(|| HttpResponse::Ok())))
//!         .bind("127.0.0.1:59880")?
//!         .run()
//! # );
//! # Ok(())
//! }
//! ```
use std::cell::RefCell;
use std::rc::Rc;

use actix_web::dev::{Extensions, Payload, ServiceRequest, ServiceResponse};
use actix_web::{Error, FromRequest, HttpMessage, HttpRequest};
use hashbrown::HashMap;
use serde::de::DeserializeOwned;
use serde::Serialize;
use serde_json;

#[cfg(feature = "cookie-session")]
mod cookie;
#[cfg(feature = "cookie-session")]
pub use crate::cookie::CookieSession;

/// The high-level interface you use to modify session data.
///
/// Session object could be obtained with
/// [`RequestSession::session`](trait.RequestSession.html#tymethod.session)
/// method. `RequestSession` trait is implemented for `HttpRequest`.
///
/// ```rust
/// use actix_session::Session;
/// use actix_web::*;
///
/// fn index(session: Session) -> Result<&'static str> {
///     // access session data
///     if let Some(count) = session.get::<i32>("counter")? {
///         session.set("counter", count + 1)?;
///     } else {
///         session.set("counter", 1)?;
///     }
///
///     Ok("Welcome!")
/// }
/// # fn main() {}
/// ```
pub struct Session(Rc<RefCell<SessionInner>>);

/// Helper trait that allows to get session
pub trait UserSession {
    fn get_session(&mut self) -> Session;
}

impl UserSession for HttpRequest {
    fn get_session(&mut self) -> Session {
        Session::get_session(&mut *self.extensions_mut())
    }
}

impl UserSession for ServiceRequest {
    fn get_session(&mut self) -> Session {
        Session::get_session(&mut *self.extensions_mut())
    }
}

#[derive(PartialEq, Clone, Debug)]
pub enum SessionStatus {
    Changed,
    Purged,
    Renewed,
    Unchanged
}
impl Default for SessionStatus {
    fn default() -> SessionStatus {
        SessionStatus::Unchanged
    }
}

#[derive(Default)]
struct SessionInner {
    state: HashMap<String, String>,
    pub status: SessionStatus,
}

impl Session {
    /// Get a `value` from the session.
    pub fn get<T: DeserializeOwned>(&self, key: &str) -> Result<Option<T>, Error> {
        if let Some(s) = self.0.borrow().state.get(key) {
            Ok(Some(serde_json::from_str(s)?))
        } else {
            Ok(None)
        }
    }

    /// Set a `value` from the session.
    pub fn set<T: Serialize>(&self, key: &str, value: T) -> Result<(), Error> {
        let mut inner = self.0.borrow_mut();
        if inner.status != SessionStatus::Purged {
            inner.status = SessionStatus::Changed;
            inner
                .state
                .insert(key.to_owned(), serde_json::to_string(&value)?);
        }
        Ok(())
    }

    /// Remove value from the session.
    pub fn remove(&self, key: &str) {
        let mut inner = self.0.borrow_mut();
        if inner.status != SessionStatus::Purged {
            inner.status = SessionStatus::Changed;
            inner.state.remove(key);
        }
    }

    /// Clear the session.
    pub fn clear(&self) {
        let mut inner = self.0.borrow_mut();
        if inner.status != SessionStatus::Purged {
            inner.status = SessionStatus::Changed;
            inner.state.clear()
        }
    }

    /// Removes session, both client and server side.
    pub fn purge(&self) {
        let mut inner = self.0.borrow_mut();
        inner.status = SessionStatus::Purged;
        inner.state.clear();
    }

    /// Renews the session key, assigning existing session state to new key.
    pub fn renew(&self) {
        let mut inner = self.0.borrow_mut();
        if inner.status != SessionStatus::Purged {
            inner.status = SessionStatus::Renewed;
        }
    }

    pub fn set_session(
        data: impl Iterator<Item = (String, String)>,
        req: &mut ServiceRequest,
    ) {
        let session = Session::get_session(&mut *req.extensions_mut());
        let mut inner = session.0.borrow_mut();
        inner.state.extend(data);
    }

    pub fn get_changes<B>(
        res: &mut ServiceResponse<B>,
    ) -> (SessionStatus, Option<impl Iterator<Item = (String, String)>>) {
        if let Some(s_impl) = res
            .request()
            .extensions()
            .get::<Rc<RefCell<SessionInner>>>()
        {
            let state =
                std::mem::replace(&mut s_impl.borrow_mut().state, HashMap::new());
            (s_impl.borrow().status.clone(), Some(state.into_iter()))
        } else {
            (SessionStatus::Unchanged, None)
        }
    }

    fn get_session(extensions: &mut Extensions) -> Session {
        if let Some(s_impl) = extensions.get::<Rc<RefCell<SessionInner>>>() {
            return Session(Rc::clone(&s_impl));
        }
        let inner = Rc::new(RefCell::new(SessionInner::default()));
        extensions.insert(inner.clone());
        Session(inner)
    }
}

/// Extractor implementation for Session type.
///
/// ```rust
/// # use actix_web::*;
/// use actix_session::Session;
///
/// fn index(session: Session) -> Result<&'static str> {
///     // access session data
///     if let Some(count) = session.get::<i32>("counter")? {
///         session.set("counter", count + 1)?;
///     } else {
///         session.set("counter", 1)?;
///     }
///
///     Ok("Welcome!")
/// }
/// # fn main() {}
/// ```
impl FromRequest for Session {
    type Error = Error;
    type Future = Result<Session, Error>;
    type Config = ();

    #[inline]
    fn from_request(req: &HttpRequest, _: &mut Payload) -> Self::Future {
        Ok(Session::get_session(&mut *req.extensions_mut()))
    }
}

#[cfg(test)]
mod tests {
    use actix_web::{test, HttpResponse};

    use super::*;

    #[test]
    fn session() {
        let mut req = test::TestRequest::default().to_srv_request();

        Session::set_session(
            vec![("key".to_string(), "\"value\"".to_string())].into_iter(),
            &mut req,
        );
        let session = Session::get_session(&mut *req.extensions_mut());
        let res = session.get::<String>("key").unwrap();
        assert_eq!(res, Some("value".to_string()));

        session.set("key2", "value2".to_string()).unwrap();
        session.remove("key");

        let mut res = req.into_response(HttpResponse::Ok().finish());
        let (_status, state) = Session::get_changes(&mut res);
        let changes: Vec<_> = state.unwrap().collect();
        assert_eq!(changes, [("key2".to_string(), "\"value2\"".to_string())]);
    }

    #[test]
    fn get_session() {
        let mut req = test::TestRequest::default().to_srv_request();

        Session::set_session(
            vec![("key".to_string(), "\"value\"".to_string())].into_iter(),
            &mut req,
        );

        let session = req.get_session();
        let res = session.get::<String>("key").unwrap();
        assert_eq!(res, Some("value".to_string()));
    }

    #[test]
    fn purge_session() {
        let mut req = test::TestRequest::default().to_srv_request();
        let session = Session::get_session(&mut *req.extensions_mut());
        assert_eq!(session.0.borrow().status, SessionStatus::Unchanged);
        session.purge();
        assert_eq!(session.0.borrow().status, SessionStatus::Purged);
    }

    #[test]
    fn renew_session() {
        let mut req = test::TestRequest::default().to_srv_request();
        let session = Session::get_session(&mut *req.extensions_mut());
        assert_eq!(session.0.borrow().status, SessionStatus::Unchanged);
        session.renew();
        assert_eq!(session.0.borrow().status, SessionStatus::Renewed);
    }
}
<commit_msg>made Session::get_session public<commit_after>//! User sessions.
//!
//! Actix provides a general solution for session management. Session
//! middlewares could provide different implementations which could
//! be accessed via general session api.
//!
//! By default, only cookie session backend is implemented. Other
//! backend implementations can be added.
//!
//! In general, you insert a *session* middleware and initialize it
//! , such as a `CookieSessionBackend`. To access session data,
//! [*Session*](struct.Session.html) extractor must be used. Session
//! extractor allows us to get or set session data.
//!
//! ```rust
//! use actix_web::{web, App, HttpServer, HttpResponse, Error};
//! use actix_session::{Session, CookieSession};
//!
//! fn index(session: Session) -> Result<&'static str, Error> {
//!     // access session data
//!     if let Some(count) = session.get::<i32>("counter")? {
//!         println!("SESSION value: {}", count);
//!         session.set("counter", count+1)?;
//!     } else {
//!         session.set("counter", 1)?;
//!     }
//!
//!     Ok("Welcome!")
//! }
//!
//! fn main() -> std::io::Result<()> {
//! # std::thread::spawn(||
//!     HttpServer::new(
//!         || App::new().wrap(
//!               CookieSession::signed(&[0; 32]) // <- create cookie based session middleware
//!                     .secure(false)
//!              )
//!             .service(web::resource("/").to(|| HttpResponse::Ok())))
//!         .bind("127.0.0.1:59880")?
//!         .run()
//! # );
//! # Ok(())
//! }
//! ```
use std::cell::RefCell;
use std::rc::Rc;

use actix_web::dev::{Extensions, Payload, ServiceRequest, ServiceResponse};
use actix_web::{Error, FromRequest, HttpMessage, HttpRequest};
use hashbrown::HashMap;
use serde::de::DeserializeOwned;
use serde::Serialize;
use serde_json;

#[cfg(feature = "cookie-session")]
mod cookie;
#[cfg(feature = "cookie-session")]
pub use crate::cookie::CookieSession;

/// The high-level interface you use to modify session data.
///
/// Session object could be obtained with
/// [`RequestSession::session`](trait.RequestSession.html#tymethod.session)
/// method. `RequestSession` trait is implemented for `HttpRequest`.
///
/// ```rust
/// use actix_session::Session;
/// use actix_web::*;
///
/// fn index(session: Session) -> Result<&'static str> {
///     // access session data
///     if let Some(count) = session.get::<i32>("counter")? {
///         session.set("counter", count + 1)?;
///     } else {
///         session.set("counter", 1)?;
///     }
///
///     Ok("Welcome!")
/// }
/// # fn main() {}
/// ```
pub struct Session(Rc<RefCell<SessionInner>>);

/// Helper trait that allows to get session
pub trait UserSession {
    fn get_session(&mut self) -> Session;
}

impl UserSession for HttpRequest {
    fn get_session(&mut self) -> Session {
        Session::get_session(&mut *self.extensions_mut())
    }
}

impl UserSession for ServiceRequest {
    fn get_session(&mut self) -> Session {
        Session::get_session(&mut *self.extensions_mut())
    }
}

#[derive(PartialEq, Clone, Debug)]
pub enum SessionStatus {
    Changed,
    Purged,
    Renewed,
    Unchanged
}
impl Default for SessionStatus {
    fn default() -> SessionStatus {
        SessionStatus::Unchanged
    }
}

#[derive(Default)]
struct SessionInner {
    state: HashMap<String, String>,
    pub status: SessionStatus,
}

impl Session {
    /// Get a `value` from the session.
    pub fn get<T: DeserializeOwned>(&self, key: &str) -> Result<Option<T>, Error> {
        if let Some(s) = self.0.borrow().state.get(key) {
            Ok(Some(serde_json::from_str(s)?))
        } else {
            Ok(None)
        }
    }

    /// Set a `value` from the session.
    pub fn set<T: Serialize>(&self, key: &str, value: T) -> Result<(), Error> {
        let mut inner = self.0.borrow_mut();
        if inner.status != SessionStatus::Purged {
            inner.status = SessionStatus::Changed;
            inner
                .state
                .insert(key.to_owned(), serde_json::to_string(&value)?);
        }
        Ok(())
    }

    /// Remove value from the session.
    pub fn remove(&self, key: &str) {
        let mut inner = self.0.borrow_mut();
        if inner.status != SessionStatus::Purged {
            inner.status = SessionStatus::Changed;
            inner.state.remove(key);
        }
    }

    /// Clear the session.
    pub fn clear(&self) {
        let mut inner = self.0.borrow_mut();
        if inner.status != SessionStatus::Purged {
            inner.status = SessionStatus::Changed;
            inner.state.clear()
        }
    }

    /// Removes session, both client and server side.
    pub fn purge(&self) {
        let mut inner = self.0.borrow_mut();
        inner.status = SessionStatus::Purged;
        inner.state.clear();
    }

    /// Renews the session key, assigning existing session state to new key.
    pub fn renew(&self) {
        let mut inner = self.0.borrow_mut();
        if inner.status != SessionStatus::Purged {
            inner.status = SessionStatus::Renewed;
        }
    }

    pub fn set_session(
        data: impl Iterator<Item = (String, String)>,
        req: &mut ServiceRequest,
    ) {
        let session = Session::get_session(&mut *req.extensions_mut());
        let mut inner = session.0.borrow_mut();
        inner.state.extend(data);
    }

    pub fn get_changes<B>(
        res: &mut ServiceResponse<B>,
    ) -> (SessionStatus, Option<impl Iterator<Item = (String, String)>>) {
        if let Some(s_impl) = res
            .request()
            .extensions()
            .get::<Rc<RefCell<SessionInner>>>()
        {
            let state =
                std::mem::replace(&mut s_impl.borrow_mut().state, HashMap::new());
            (s_impl.borrow().status.clone(), Some(state.into_iter()))
        } else {
            (SessionStatus::Unchanged, None)
        }
    }

    pub fn get_session(extensions: &mut Extensions) -> Session {
        if let Some(s_impl) = extensions.get::<Rc<RefCell<SessionInner>>>() {
            return Session(Rc::clone(&s_impl));
        }
        let inner = Rc::new(RefCell::new(SessionInner::default()));
        extensions.insert(inner.clone());
        Session(inner)
    }
}

/// Extractor implementation for Session type.
///
/// ```rust
/// # use actix_web::*;
/// use actix_session::Session;
///
/// fn index(session: Session) -> Result<&'static str> {
///     // access session data
///     if let Some(count) = session.get::<i32>("counter")? {
///         session.set("counter", count + 1)?;
///     } else {
///         session.set("counter", 1)?;
///     }
///
///     Ok("Welcome!")
/// }
/// # fn main() {}
/// ```
impl FromRequest for Session {
    type Error = Error;
    type Future = Result<Session, Error>;
    type Config = ();

    #[inline]
    fn from_request(req: &HttpRequest, _: &mut Payload) -> Self::Future {
        Ok(Session::get_session(&mut *req.extensions_mut()))
    }
}

#[cfg(test)]
mod tests {
    use actix_web::{test, HttpResponse};

    use super::*;

    #[test]
    fn session() {
        let mut req = test::TestRequest::default().to_srv_request();

        Session::set_session(
            vec![("key".to_string(), "\"value\"".to_string())].into_iter(),
            &mut req,
        );
        let session = Session::get_session(&mut *req.extensions_mut());
        let res = session.get::<String>("key").unwrap();
        assert_eq!(res, Some("value".to_string()));

        session.set("key2", "value2".to_string()).unwrap();
        session.remove("key");

        let mut res = req.into_response(HttpResponse::Ok().finish());
        let (_status, state) = Session::get_changes(&mut res);
        let changes: Vec<_> = state.unwrap().collect();
        assert_eq!(changes, [("key2".to_string(), "\"value2\"".to_string())]);
    }

    #[test]
    fn get_session() {
        let mut req = test::TestRequest::default().to_srv_request();

        Session::set_session(
            vec![("key".to_string(), "\"value\"".to_string())].into_iter(),
            &mut req,
        );

        let session = req.get_session();
        let res = session.get::<String>("key").unwrap();
        assert_eq!(res, Some("value".to_string()));
    }

    #[test]
    fn purge_session() {
        let mut req = test::TestRequest::default().to_srv_request();
        let session = Session::get_session(&mut *req.extensions_mut());
        assert_eq!(session.0.borrow().status, SessionStatus::Unchanged);
        session.purge();
        assert_eq!(session.0.borrow().status, SessionStatus::Purged);
    }

    #[test]
    fn renew_session() {
        let mut req = test::TestRequest::default().to_srv_request();
        let session = Session::get_session(&mut *req.extensions_mut());
        assert_eq!(session.0.borrow().status, SessionStatus::Unchanged);
        session.renew();
        assert_eq!(session.0.borrow().status, SessionStatus::Renewed);
    }
}
<|endoftext|>


<commit_before>
use super::common::*;

use std::iter;

use idents;
use methodinfo::Direction;
use model;

extern crate proc_macro;
use self::proc_macro::TokenStream;
use syn::*;

/// Expands the `com_impl` attribute.
///
/// The attribute expansion results in the following items:
///
/// - Implementation for the delegating methods when calling the Rust methods
///   from COM.
/// - Virtual table instance for the COM type.
pub fn expand_com_impl(
    _attr_tokens: &TokenStream,
    item_tokens: TokenStream,
) -> Result<TokenStream, model::ParseError>
{
    // Parse the attribute.
    let mut output = vec![] ;
    let imp = model::ComImpl::parse( &item_tokens.to_string() )?;
    let struct_ident = imp.struct_name();
    let itf_ident = imp.interface_name();
    let vtable_struct_ident = idents::vtable_struct( itf_ident );
    let vtable_instance_ident = idents::vtable_instance( struct_ident, itf_ident );
    let vtable_offset = idents::vtable_offset( struct_ident, itf_ident );

    /////////////////////
    // #itf::QueryInterface, AddRef & Release
    //
    // Note that the actual methods implementation for these bits differs from
    // the primary IUnknown methods. When the methods are being called through
    // this vtable, the self_vtable pointer will point to this vtable and not
    // the start of the CoClass instance.
    //
    // We can convert these to the ComBox references by offsetting the pointer
    // by the known vtable offset.

    // QueryInterface
    let calling_convetion = get_calling_convetion();
    let query_interface_ident = idents::method_impl(
            struct_ident, itf_ident, "query_interface" );
    output.push( quote!(
            #[allow(non_snake_case)]
            #[doc(hidden)]
            unsafe extern #calling_convetion fn #query_interface_ident(
                self_vtable : ::intercom::RawComPtr,
                riid : ::intercom::REFIID,
                out : *mut ::intercom::RawComPtr
            ) -> ::intercom::HRESULT
            {
                // Get the primary iunk interface by offsetting the current
                // self_vtable with the vtable offset. Once we have the primary
                // pointer we can delegate the call to the primary implementation.
                ::intercom::ComBox::< #struct_ident >::query_interface(
                        &mut *(( self_vtable as usize - #vtable_offset() ) as *mut _ ),
                        riid,
                        out )
            }
        ) );

    // AddRef
    let add_ref_ident = idents::method_impl(
            struct_ident, itf_ident, "add_ref" );
    output.push( quote!(
            #[allow(non_snake_case)]
            #[allow(dead_code)]
            #[doc(hidden)]
            unsafe extern #calling_convetion fn #add_ref_ident(
                self_vtable : ::intercom::RawComPtr
            ) -> u32 {
                ::intercom::ComBox::< #struct_ident >::add_ref(
                        &mut *(( self_vtable as usize - #vtable_offset() ) as *mut _ ) )
            }
        ) );

    // Release
    let release_ident = idents::method_impl(
            struct_ident, itf_ident, "release" );
    output.push( quote!(
            #[allow(non_snake_case)]
            #[allow(dead_code)]
            #[doc(hidden)]
            unsafe extern #calling_convetion fn #release_ident(
                self_vtable : ::intercom::RawComPtr
            ) -> u32 {
                ::intercom::ComBox::< #struct_ident >::release_ptr(
                        ( self_vtable as usize - #vtable_offset() ) as *mut _ )
            }
        ) );

    // Start the definition fo the vtable fields. The base interface is always
    // IUnknown at this point. We might support IDispatch later, but for now
    // we only support IUnknown.
    let mut vtable_fields = vec![
        quote!(
            __base : ::intercom::IUnknownVtbl {
                query_interface : #query_interface_ident,
                add_ref : #add_ref_ident,
                release : #release_ident,
            },
        ) ];

    // Process the impl items. This gathers all COM-visible methods and defines
    // delegating calls for them. These delegating calls are the ones that are
    // invoked by the clients. The calls then convert everything to the RUST
    // interface.
    //
    // The impl may have various kinds of items - we only support the ones that
    // seem okay. So in case we encounter any errors we'll just skip the method
    // silently. This is done by breaking out of the 'catch' before adding the
    // method to the vtable fields.
    for method_info in imp.methods() {

        let method_ident = &method_info.name;
        let method_impl_ident = idents::method_impl(
                struct_ident, itf_ident, method_ident.as_ref() );

        let in_out_args = method_info.raw_com_args()
                .into_iter()
                .map( |com_arg| {
                    let name = &com_arg.arg.name;
                    let com_ty = &com_arg.arg.handler.com_ty();
                    let dir = match com_arg.dir {
                        Direction::In => quote!(),
                        Direction::Out | Direction::Retval => quote!( *mut )
                    };
                    quote!( #name : #dir #com_ty )
                } );
        let self_arg = quote!( self_vtable : ::intercom::RawComPtr );
        let args = iter::once( self_arg ).chain( in_out_args );

        // Format stack variables if the conversion requires
        // temporary variable.
        let stack_variables = method_info.args
                .iter()
                .filter( |ca| ca.handler.com_to_rust( &ca.name ).stack.is_some() )
                .map( |ca| {
                    let stack_variable = ca.handler.com_to_rust( &ca.name ).stack.unwrap();
                    quote!( #stack_variable )
                } );

        // Format the in and out parameters for the Rust call.
        let in_params = method_info.args
                .iter()
                .map( |ca| {
                    let conversion = ca.handler.com_to_rust( &ca.name ).conversion;
                    quote!( #conversion )
                } );

        let return_ident = Ident::from( "__result" );
        let return_statement = method_info
                .returnhandler
                .rust_to_com_return( &return_ident );

        // Define the delegating method implementation.
        //
        // Note the self_vtable here will be a pointer to the start of the
        // vtable for the current interface. To get the coclass and thus
        // the actual 'data' struct, we'll need to offset the self_vtable
        // with the vtable offset.
        let ret_ty = method_info.returnhandler.com_ty();
        let self_struct_stmt = if method_info.is_const {
                quote!( let self_struct : &#itf_ident = &**self_combox )
            } else {
                quote!( let self_struct : &mut #itf_ident = &mut **self_combox )
            };

        output.push( quote!(
            #[allow(non_snake_case)]
            #[allow(dead_code)]
            #[doc(hidden)]
            unsafe extern #calling_convetion fn #method_impl_ident(
                #( #args ),*
            ) -> #ret_ty {
                // Acquire the reference to the ComBox. For this we need
                // to offset the current 'self_vtable' vtable pointer.
                let self_combox = ( self_vtable as usize - #vtable_offset() )
                        as *mut ::intercom::ComBox< #struct_ident >;

                // Store stack variables required by some conversions.
                // For example, to convert a COM type to "str" we may need to
                // convert the COM type into "String" first and then pass
                // this "String" to the Rust method as "str".
                #( #stack_variables )*

                #self_struct_stmt;
                let #return_ident = self_struct.#method_ident( #( #in_params ),* );

                #return_statement
            }
        ) );

        // Include the delegating method in the virtual table fields.
        vtable_fields.push( quote!( #method_ident : #method_impl_ident, ) );
    }

    // Now that we've gathered all the virtual table fields, we can finally
    // emit the virtual table instance.
    output.push( quote!(
            #[allow(non_upper_case_globals)]
            const #vtable_instance_ident : #vtable_struct_ident
                    = #vtable_struct_ident { #( #vtable_fields )* };
        ) );

    Ok( tokens_to_tokenstream( item_tokens, output ) )
}
<commit_msg>Improved code clarity: filter+map => filter_map<commit_after>
use super::common::*;

use std::iter;

use idents;
use methodinfo::Direction;
use model;

extern crate proc_macro;
use self::proc_macro::TokenStream;
use syn::*;

/// Expands the `com_impl` attribute.
///
/// The attribute expansion results in the following items:
///
/// - Implementation for the delegating methods when calling the Rust methods
///   from COM.
/// - Virtual table instance for the COM type.
pub fn expand_com_impl(
    _attr_tokens: &TokenStream,
    item_tokens: TokenStream,
) -> Result<TokenStream, model::ParseError>
{
    // Parse the attribute.
    let mut output = vec![] ;
    let imp = model::ComImpl::parse( &item_tokens.to_string() )?;
    let struct_ident = imp.struct_name();
    let itf_ident = imp.interface_name();
    let vtable_struct_ident = idents::vtable_struct( itf_ident );
    let vtable_instance_ident = idents::vtable_instance( struct_ident, itf_ident );
    let vtable_offset = idents::vtable_offset( struct_ident, itf_ident );

    /////////////////////
    // #itf::QueryInterface, AddRef & Release
    //
    // Note that the actual methods implementation for these bits differs from
    // the primary IUnknown methods. When the methods are being called through
    // this vtable, the self_vtable pointer will point to this vtable and not
    // the start of the CoClass instance.
    //
    // We can convert these to the ComBox references by offsetting the pointer
    // by the known vtable offset.

    // QueryInterface
    let calling_convetion = get_calling_convetion();
    let query_interface_ident = idents::method_impl(
            struct_ident, itf_ident, "query_interface" );
    output.push( quote!(
            #[allow(non_snake_case)]
            #[doc(hidden)]
            unsafe extern #calling_convetion fn #query_interface_ident(
                self_vtable : ::intercom::RawComPtr,
                riid : ::intercom::REFIID,
                out : *mut ::intercom::RawComPtr
            ) -> ::intercom::HRESULT
            {
                // Get the primary iunk interface by offsetting the current
                // self_vtable with the vtable offset. Once we have the primary
                // pointer we can delegate the call to the primary implementation.
                ::intercom::ComBox::< #struct_ident >::query_interface(
                        &mut *(( self_vtable as usize - #vtable_offset() ) as *mut _ ),
                        riid,
                        out )
            }
        ) );

    // AddRef
    let add_ref_ident = idents::method_impl(
            struct_ident, itf_ident, "add_ref" );
    output.push( quote!(
            #[allow(non_snake_case)]
            #[allow(dead_code)]
            #[doc(hidden)]
            unsafe extern #calling_convetion fn #add_ref_ident(
                self_vtable : ::intercom::RawComPtr
            ) -> u32 {
                ::intercom::ComBox::< #struct_ident >::add_ref(
                        &mut *(( self_vtable as usize - #vtable_offset() ) as *mut _ ) )
            }
        ) );

    // Release
    let release_ident = idents::method_impl(
            struct_ident, itf_ident, "release" );
    output.push( quote!(
            #[allow(non_snake_case)]
            #[allow(dead_code)]
            #[doc(hidden)]
            unsafe extern #calling_convetion fn #release_ident(
                self_vtable : ::intercom::RawComPtr
            ) -> u32 {
                ::intercom::ComBox::< #struct_ident >::release_ptr(
                        ( self_vtable as usize - #vtable_offset() ) as *mut _ )
            }
        ) );

    // Start the definition fo the vtable fields. The base interface is always
    // IUnknown at this point. We might support IDispatch later, but for now
    // we only support IUnknown.
    let mut vtable_fields = vec![
        quote!(
            __base : ::intercom::IUnknownVtbl {
                query_interface : #query_interface_ident,
                add_ref : #add_ref_ident,
                release : #release_ident,
            },
        ) ];

    // Process the impl items. This gathers all COM-visible methods and defines
    // delegating calls for them. These delegating calls are the ones that are
    // invoked by the clients. The calls then convert everything to the RUST
    // interface.
    //
    // The impl may have various kinds of items - we only support the ones that
    // seem okay. So in case we encounter any errors we'll just skip the method
    // silently. This is done by breaking out of the 'catch' before adding the
    // method to the vtable fields.
    for method_info in imp.methods() {

        let method_ident = &method_info.name;
        let method_impl_ident = idents::method_impl(
                struct_ident, itf_ident, method_ident.as_ref() );

        let in_out_args = method_info.raw_com_args()
                .into_iter()
                .map( |com_arg| {
                    let name = &com_arg.arg.name;
                    let com_ty = &com_arg.arg.handler.com_ty();
                    let dir = match com_arg.dir {
                        Direction::In => quote!(),
                        Direction::Out | Direction::Retval => quote!( *mut )
                    };
                    quote!( #name : #dir #com_ty )
                } );
        let self_arg = quote!( self_vtable : ::intercom::RawComPtr );
        let args = iter::once( self_arg ).chain( in_out_args );

        // Format stack variables if the conversion requires
        // temporary variable.
        let stack_variables = method_info.args.iter()
                .filter_map( |ca| match ca.handler.com_to_rust( &ca.name ).stack {
                    Some( stack ) => Some( quote!( #stack ) ),
                    None => None
                } );

        // Format the in and out parameters for the Rust call.
        let in_params = method_info.args
                .iter()
                .map( |ca| {
                    let conversion = ca.handler.com_to_rust( &ca.name ).conversion;
                    quote!( #conversion )
                } );

        let return_ident = Ident::from( "__result" );
        let return_statement = method_info
                .returnhandler
                .rust_to_com_return( &return_ident );

        // Define the delegating method implementation.
        //
        // Note the self_vtable here will be a pointer to the start of the
        // vtable for the current interface. To get the coclass and thus
        // the actual 'data' struct, we'll need to offset the self_vtable
        // with the vtable offset.
        let ret_ty = method_info.returnhandler.com_ty();
        let self_struct_stmt = if method_info.is_const {
                quote!( let self_struct : &#itf_ident = &**self_combox )
            } else {
                quote!( let self_struct : &mut #itf_ident = &mut **self_combox )
            };

        output.push( quote!(
            #[allow(non_snake_case)]
            #[allow(dead_code)]
            #[doc(hidden)]
            unsafe extern #calling_convetion fn #method_impl_ident(
                #( #args ),*
            ) -> #ret_ty {
                // Acquire the reference to the ComBox. For this we need
                // to offset the current 'self_vtable' vtable pointer.
                let self_combox = ( self_vtable as usize - #vtable_offset() )
                        as *mut ::intercom::ComBox< #struct_ident >;

                // Store stack variables required by some conversions.
                // For example, to convert a COM type to "str" we may need to
                // convert the COM type into "String" first and then pass
                // this "String" to the Rust method as "str".
                #( #stack_variables )*

                #self_struct_stmt;
                let #return_ident = self_struct.#method_ident( #( #in_params ),* );

                #return_statement
            }
        ) );

        // Include the delegating method in the virtual table fields.
        vtable_fields.push( quote!( #method_ident : #method_impl_ident, ) );
    }

    // Now that we've gathered all the virtual table fields, we can finally
    // emit the virtual table instance.
    output.push( quote!(
            #[allow(non_upper_case_globals)]
            const #vtable_instance_ident : #vtable_struct_ident
                    = #vtable_struct_ident { #( #vtable_fields )* };
        ) );

    Ok( tokens_to_tokenstream( item_tokens, output ) )
}
<|endoftext|>


<commit_before>use std::io::{self, Write};
use termcolor::{Color, ColorChoice, ColorSpec, StandardStream, WriteColor};

pub fn heading(content: &str) {
    let mut stdout = StandardStream::stdout(ColorChoice::Always);

    match stdout.set_color(ColorSpec::new().set_fg(Some(Color::Cyan))) {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    match writeln!(&mut stdout, "\n{}\n{}", content, underline(content, "=")) {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    match stdout.reset() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }
}

pub fn subheading(content: &str) {
    let mut stdout = StandardStream::stdout(ColorChoice::Always);

    match stdout.set_color(ColorSpec::new().set_fg(Some(Color::Cyan))) {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    match writeln!(&mut stdout, "\n{}\n{}", content, underline(content, "-")) {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    match stdout.reset() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }
}

pub fn progress(content: String) {
    let mut stdout = io::stdout();

    match writeln!(&mut stdout, "\n{}", content) {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }
}

pub fn info(content: String) {
    let mut stdout = io::stdout();

    match writeln!(&mut stdout, "\n{}", content) {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }
}

pub fn error(content: String) {
    let mut stdout = StandardStream::stdout(ColorChoice::Always);

    match stdout.set_color(ColorSpec::new().set_fg(Some(Color::Red))) {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    match writeln!(&mut stdout, "\n{}", content) {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    match stdout.reset() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }
}

pub fn list(items: Vec<&str>) {
    let mut stdout = io::stdout();

    match writeln!(&mut stdout, "") {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    items
        .iter()
        .for_each(|content| match writeln!(&mut stdout, "- {}", content) {
            Ok(_) => (),
            Err(error) => panic!("There was a problem: {:?}", error),
        })
}

pub fn pause() {
    let mut stdout = StandardStream::stdout(ColorChoice::Always);

    match stdout.set_color(ColorSpec::new().set_fg(Some(Color::Yellow))) {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    loop {
        let prompt = "Enter 'Done' when ready to continue:";

        match writeln!(&mut stdout, "\n{}", prompt) {
            Ok(_) => (),
            Err(error) => panic!("There was a problem: {:?}", error),
        }

        let mut input = String::new();

        match io::stdin().read_line(&mut input) {
            Ok(_) => (),
            Err(error) => panic!("There was a problem: {:?}", error),
        }

        if input.contains("Done") {
            match stdout.reset() {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }

            break;
        } else {
            continue;
        }
    }

    for x in (0..100).step_by(0) {
        println!("{}", x);
    };
}

fn underline(content: &str, symbol: &str) -> String {
                content.chars().map(|_| symbol).collect()
}

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

    #[test]
    fn test_underline() {
        assert_eq!(underline("Some line", "="), "=========");
        assert_eq!(underline("Another line", "-"), "------------");
    }
}
<commit_msg>Checking that linting works in CI/CD<commit_after>use std::io::{self, Write};
use termcolor::{Color, ColorChoice, ColorSpec, StandardStream, WriteColor};

pub fn heading(content: &str) {
    let mut stdout = StandardStream::stdout(ColorChoice::Always);

    match stdout.set_color(ColorSpec::new().set_fg(Some(Color::Cyan))) {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    match writeln!(&mut stdout, "\n{}\n{}", content, underline(content, "=")) {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    match stdout.reset() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }
}

pub fn subheading(content: &str) {
    let mut stdout = StandardStream::stdout(ColorChoice::Always);

    match stdout.set_color(ColorSpec::new().set_fg(Some(Color::Cyan))) {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    match writeln!(&mut stdout, "\n{}\n{}", content, underline(content, "-")) {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    match stdout.reset() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }
}

pub fn progress(content: String) {
    let mut stdout = io::stdout();

    match writeln!(&mut stdout, "\n{}", content) {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }
}

pub fn info(content: String) {
    let mut stdout = io::stdout();

    match writeln!(&mut stdout, "\n{}", content) {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }
}

pub fn error(content: String) {
    let mut stdout = StandardStream::stdout(ColorChoice::Always);

    match stdout.set_color(ColorSpec::new().set_fg(Some(Color::Red))) {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    match writeln!(&mut stdout, "\n{}", content) {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    match stdout.reset() {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }
}

pub fn list(items: Vec<&str>) {
    let mut stdout = io::stdout();

    match writeln!(&mut stdout, "") {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    items
        .iter()
        .for_each(|content| match writeln!(&mut stdout, "- {}", content) {
            Ok(_) => (),
            Err(error) => panic!("There was a problem: {:?}", error),
        })
}

pub fn pause() {
    let mut stdout = StandardStream::stdout(ColorChoice::Always);

    match stdout.set_color(ColorSpec::new().set_fg(Some(Color::Yellow))) {
        Ok(_) => (),
        Err(error) => panic!("There was a problem: {:?}", error),
    }

    loop {
        let prompt = "Enter 'Done' when ready to continue:";

        match writeln!(&mut stdout, "\n{}", prompt) {
            Ok(_) => (),
            Err(error) => panic!("There was a problem: {:?}", error),
        }

        let mut input = String::new();

        match io::stdin().read_line(&mut input) {
            Ok(_) => (),
            Err(error) => panic!("There was a problem: {:?}", error),
        }

        if input.contains("Done") {
            match stdout.reset() {
                Ok(_) => (),
                Err(error) => panic!("There was a problem: {:?}", error),
            }

            break;
        } else {
            continue;
        }
    }

    let x = 3.1415;
    let r = 8.0;
    println!("the area of the circle is {}", x * r * r);
}

fn underline(content: &str, symbol: &str) -> String {
                content.chars().map(|_| symbol).collect()
}

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

    #[test]
    fn test_underline() {
        assert_eq!(underline("Some line", "="), "=========");
        assert_eq!(underline("Another line", "-"), "------------");
    }
}
<|endoftext|>