<commit_before>use std::{
    collections::VecDeque,
    fmt, io,
    sync::{
        atomic::{AtomicU64, AtomicUsize, Ordering},
        Arc, Mutex,
    },
    time::{Duration, Instant},
};

use crate::{
    config::ServerConfig,
    context::SharedContext,
    relay::tcprelay::EncryptedWrite,
    relay::{loadbalancing::server::LoadBalancer, socks5::Address, tcprelay::client::ServerClient},
};

use futures::{Future, Stream};
use log::{debug, error, info};
use tokio::{
    self,
    io::read_exact,
    timer::{Interval, Timeout},
};

struct Server {
    config: Arc<ServerConfig>,
    elapsed: AtomicU64,
}

impl Server {
    fn new(config: ServerConfig) -> Server {
        Server {
            config: Arc::new(config),
            elapsed: AtomicU64::new(0),
        }
    }

    fn delay(&self) -> u64 {
        self.elapsed.load(Ordering::Acquire)
    }
}

const MAX_LATENCY_QUEUE_SIZE: usize = 17;

struct ServerLatencyInner {
    latency_queue: VecDeque<u64>,
}

impl ServerLatencyInner {
    fn new() -> ServerLatencyInner {
        ServerLatencyInner {
            latency_queue: VecDeque::with_capacity(MAX_LATENCY_QUEUE_SIZE),
        }
    }

    fn push(&mut self, lat: u64) -> u64 {
        self.latency_queue.push_back(lat);
        if self.latency_queue.len() > MAX_LATENCY_QUEUE_SIZE {
            self.latency_queue.pop_front();
        }

        self.representation()
    }

    fn representation(&self) -> u64 {
        if self.latency_queue.is_empty() {
            return 0;
        }

        let sum: u64 = self.latency_queue.iter().sum();
        let avg = sum as f64 / self.latency_queue.len() as f64;
        let s_sum: u64 = self.latency_queue.iter().map(|n| *n * *n).sum();
        let dev = (s_sum as f64 / self.latency_queue.len() as f64 - avg * avg).sqrt();

        let mut v = self
            .latency_queue
            .iter()
            .cloned()
            .filter(|n| *n as f64 >= (avg - dev) && *n as f64 <= (avg + dev))
            .collect::<Vec<u64>>();
        v.sort();

        debug!(
            "Latency queue {:?}, avg {}, dev {}, sorted {:?}",
            self.latency_queue, avg, dev, v
        );

        if v.is_empty() {
            return 0;
        }

        let mid = v.len() / 2;
        if (v.len() & 1) == 0 {
            (v[mid - 1] + v[mid]) / 2
        } else {
            v[mid]
        }
    }
}

#[derive(Clone)]
struct ServerLatency {
    inner: Arc<Mutex<ServerLatencyInner>>,
}

impl ServerLatency {
    fn new() -> ServerLatency {
        ServerLatency {
            inner: Arc::new(Mutex::new(ServerLatencyInner::new())),
        }
    }

    fn push(&self, lat: u64) -> u64 {
        let mut inner = self.inner.lock().unwrap();
        inner.push(lat)
    }
}

impl fmt::Debug for ServerLatency {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let inner = self.inner.lock().unwrap();
        write!(f, "ServerLatency {{ latency: {:?} }}", inner.latency_queue)
    }
}

const DEFAULT_CHECK_INTERVAL_SEC: u64 = 10;
const DEFAULT_CHECK_TIMEOUT_SEC: u64 = 5;

struct Inner {
    servers: Vec<Arc<Server>>,
    best_idx: AtomicUsize,
}

impl Inner {
    fn new(context: SharedContext) -> Inner {
        let config = context.config();

        if config.server.is_empty() {
            panic!("no servers");
        }

        let check_required = config.server.len() > 1;

        let mut servers = Vec::new();
        for svr in &config.server {
            let sc = Arc::new(Server::new(svr.clone()));
            servers.push(sc.clone());

            if !check_required {
                continue;
            }

            let context = context.clone();
            let latency = ServerLatency::new();

            tokio::spawn(
                // Check every 10 seconds
                Interval::new(
                    Instant::now() + Duration::from_secs(1),
                    Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC),
                )
                .for_each(move |_| {
                    let sc = sc.clone();
                    let lat = latency.clone();

                    Inner::check_delay(sc.clone(), context.clone()).then(move |res| {
                        let md = match res {
                            Ok(d) => lat.push(d),
                            Err(..) => lat.push(DEFAULT_CHECK_TIMEOUT_SEC * 2 * 1000), // Penalty
                        };
                        debug!(
                            "updated remote server {} (median: {} ms) {:?}",
                            sc.config.addr(),
                            md,
                            lat
                        );
                        sc.elapsed.store(md, Ordering::Release);

                        Ok(())
                    })
                })
                .map_err(|_| ()),
            );
        }

        Inner {
            servers,
            best_idx: AtomicUsize::new(0),
        }
    }

    fn check_delay(sc: Arc<Server>, context: SharedContext) -> impl Future<Item = u64, Error = io::Error> {
        static GET_BODY: &[u8] = b"GET / HTTP/1.1\r\nHost: www.example.com\r\nConnection: close\r\nAccept: */*\r\n\r\n";

        let addr = Address::DomainNameAddress("www.example.com".to_owned(), 80);

        let start = Instant::now();

        // Send HTTP GET and read the first byte
        let fut = ServerClient::connect(context, addr, sc.config.clone()).and_then(|ServerClient { r, w }| {
            w.write_all(GET_BODY)
                .and_then(|(mut w, _)| w.flush())
                .and_then(move |_| read_exact(r, [0u8, 1]))
        });

        Timeout::new(fut, Duration::from_secs(DEFAULT_CHECK_TIMEOUT_SEC)).then(move |res| {
            let elapsed = Instant::now() - start;
            let elapsed = elapsed.as_secs() * 1000 + u64::from(elapsed.subsec_millis()); // Converted to ms
            match res {
                Ok(..) => {
                    // Got the result ... record its time
                    debug!("checked remote server {} latency with {} ms", sc.config.addr(), elapsed);
                    Ok(elapsed)
                }
                Err(err) => {
                    match err.into_inner() {
                        Some(err) => {
                            error!("failed to check server {}, error: {}", sc.config.addr(), err);

                            // NOTE: connection / handshake error, server is down
                            Err(err)
                        }
                        None => {
                            // Timeout
                            debug!(
                                "checked remote server {} latency timeout, elapsed {} ms",
                                sc.config.addr(),
                                elapsed
                            );

                            // NOTE: timeout is still available, but server is too slow
                            Ok(elapsed)
                        }
                    }
                }
            }
        })
    }

    fn best_idx(&self) -> usize {
        self.best_idx.load(Ordering::Acquire)
    }

    fn set_best_idx(&self, idx: usize) {
        self.best_idx.store(idx, Ordering::Release)
    }

    fn best_server(&self) -> &Arc<Server> {
        &self.servers[self.best_idx()]
    }
}

#[derive(Clone)]
pub struct PingBalancer {
    inner: Arc<Inner>,
}

impl PingBalancer {
    pub fn new(context: SharedContext) -> PingBalancer {
        let inner = Arc::new(Inner::new(context));

        let inner_cloned = inner.clone();
        tokio::spawn(
            Interval::new(
                Instant::now() + Duration::from_secs(2),
                Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC),
            )
            .for_each(move |_| {
                let inner = inner_cloned.clone();

                if inner.servers.is_empty() {
                    panic!("No server");
                }

                // Choose the best one
                let mut svr_idx = 0;

                for (idx, svr) in inner.servers.iter().enumerate() {
                    let choosen_svr = &inner.servers[svr_idx];
                    if svr.delay() < choosen_svr.delay() {
                        svr_idx = idx;
                    }
                }

                let choosen_svr = &inner.servers[svr_idx];
                debug!(
                    "chosen the best server {} (delay: {} ms)",
                    choosen_svr.config.addr(),
                    choosen_svr.delay()
                );

                if inner.best_idx() != svr_idx {
                    info!(
                        "switched server from {} (latency: {} ms) to {} (latency: {} ms)",
                        inner.best_server().config.addr(),
                        inner.best_server().delay(),
                        choosen_svr.config.addr(),
                        choosen_svr.delay(),
                    );
                }

                inner.set_best_idx(svr_idx);

                Ok(())
            })
            .map_err(|_| ()),
        );

        PingBalancer { inner }
    }
}

impl LoadBalancer for PingBalancer {
    fn pick_server(&mut self) -> Arc<ServerConfig> {
        if self.inner.servers.is_empty() {
            panic!("no server in configuration");
        }
        self.inner.best_server().config.clone()
    }

    fn total(&self) -> usize {
        self.inner.servers.len()
    }
}
<commit_msg>Skipped server delay that are not initialized yet<commit_after>use std::{
    collections::VecDeque,
    fmt, io,
    sync::{
        atomic::{AtomicU64, AtomicUsize, Ordering},
        Arc, Mutex,
    },
    time::{Duration, Instant},
};

use crate::{
    config::ServerConfig,
    context::SharedContext,
    relay::tcprelay::EncryptedWrite,
    relay::{loadbalancing::server::LoadBalancer, socks5::Address, tcprelay::client::ServerClient},
};

use futures::{Future, Stream};
use log::{debug, error, info};
use tokio::{
    self,
    io::read_exact,
    timer::{Interval, Timeout},
};

struct Server {
    config: Arc<ServerConfig>,
    elapsed: AtomicU64,
}

impl Server {
    fn new(config: ServerConfig) -> Server {
        Server {
            config: Arc::new(config),
            elapsed: AtomicU64::new(0),
        }
    }

    fn delay(&self) -> u64 {
        self.elapsed.load(Ordering::Acquire)
    }
}

const MAX_LATENCY_QUEUE_SIZE: usize = 17;

struct ServerLatencyInner {
    latency_queue: VecDeque<u64>,
}

impl ServerLatencyInner {
    fn new() -> ServerLatencyInner {
        ServerLatencyInner {
            latency_queue: VecDeque::with_capacity(MAX_LATENCY_QUEUE_SIZE),
        }
    }

    fn push(&mut self, lat: u64) -> u64 {
        self.latency_queue.push_back(lat);
        if self.latency_queue.len() > MAX_LATENCY_QUEUE_SIZE {
            self.latency_queue.pop_front();
        }

        self.representation()
    }

    fn representation(&self) -> u64 {
        if self.latency_queue.is_empty() {
            return 0;
        }

        let sum: u64 = self.latency_queue.iter().sum();
        let avg = sum as f64 / self.latency_queue.len() as f64;
        let s_sum: u64 = self.latency_queue.iter().map(|n| *n * *n).sum();
        let dev = (s_sum as f64 / self.latency_queue.len() as f64 - avg * avg).sqrt();

        let mut v = self
            .latency_queue
            .iter()
            .cloned()
            .filter(|n| *n as f64 >= (avg - dev) && *n as f64 <= (avg + dev))
            .collect::<Vec<u64>>();
        v.sort();

        debug!(
            "Latency queue {:?}, avg {}, dev {}, sorted {:?}",
            self.latency_queue, avg, dev, v
        );

        if v.is_empty() {
            return 0;
        }

        let mid = v.len() / 2;
        if (v.len() & 1) == 0 {
            (v[mid - 1] + v[mid]) / 2
        } else {
            v[mid]
        }
    }
}

#[derive(Clone)]
struct ServerLatency {
    inner: Arc<Mutex<ServerLatencyInner>>,
}

impl ServerLatency {
    fn new() -> ServerLatency {
        ServerLatency {
            inner: Arc::new(Mutex::new(ServerLatencyInner::new())),
        }
    }

    fn push(&self, lat: u64) -> u64 {
        let mut inner = self.inner.lock().unwrap();
        inner.push(lat)
    }
}

impl fmt::Debug for ServerLatency {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let inner = self.inner.lock().unwrap();
        write!(f, "ServerLatency {{ latency: {:?} }}", inner.latency_queue)
    }
}

const DEFAULT_CHECK_INTERVAL_SEC: u64 = 10;
const DEFAULT_CHECK_TIMEOUT_SEC: u64 = 5;

struct Inner {
    servers: Vec<Arc<Server>>,
    best_idx: AtomicUsize,
}

impl Inner {
    fn new(context: SharedContext) -> Inner {
        let config = context.config();

        if config.server.is_empty() {
            panic!("no servers");
        }

        let check_required = config.server.len() > 1;

        let mut servers = Vec::new();
        for svr in &config.server {
            let sc = Arc::new(Server::new(svr.clone()));
            servers.push(sc.clone());

            if !check_required {
                continue;
            }

            let context = context.clone();
            let latency = ServerLatency::new();

            tokio::spawn(
                // Check every 10 seconds
                Interval::new(
                    Instant::now() + Duration::from_secs(1),
                    Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC),
                )
                .for_each(move |_| {
                    let sc = sc.clone();
                    let lat = latency.clone();

                    Inner::check_delay(sc.clone(), context.clone()).then(move |res| {
                        let md = match res {
                            Ok(d) => lat.push(d),
                            Err(..) => lat.push(DEFAULT_CHECK_TIMEOUT_SEC * 2 * 1000), // Penalty
                        };
                        debug!(
                            "updated remote server {} (median: {} ms) {:?}",
                            sc.config.addr(),
                            md,
                            lat
                        );
                        sc.elapsed.store(md, Ordering::Release);

                        Ok(())
                    })
                })
                .map_err(|_| ()),
            );
        }

        Inner {
            servers,
            best_idx: AtomicUsize::new(0),
        }
    }

    fn check_delay(sc: Arc<Server>, context: SharedContext) -> impl Future<Item = u64, Error = io::Error> {
        static GET_BODY: &[u8] = b"GET / HTTP/1.1\r\nHost: www.example.com\r\nConnection: close\r\nAccept: */*\r\n\r\n";

        let addr = Address::DomainNameAddress("www.example.com".to_owned(), 80);

        let start = Instant::now();

        // Send HTTP GET and read the first byte
        let fut = ServerClient::connect(context, addr, sc.config.clone()).and_then(|ServerClient { r, w }| {
            w.write_all(GET_BODY)
                .and_then(|(mut w, _)| w.flush())
                .and_then(move |_| read_exact(r, [0u8, 1]))
        });

        Timeout::new(fut, Duration::from_secs(DEFAULT_CHECK_TIMEOUT_SEC)).then(move |res| {
            let elapsed = Instant::now() - start;
            let elapsed = elapsed.as_secs() * 1000 + u64::from(elapsed.subsec_millis()); // Converted to ms
            match res {
                Ok(..) => {
                    // Got the result ... record its time
                    debug!("checked remote server {} latency with {} ms", sc.config.addr(), elapsed);
                    Ok(elapsed)
                }
                Err(err) => {
                    match err.into_inner() {
                        Some(err) => {
                            error!("failed to check server {}, error: {}", sc.config.addr(), err);

                            // NOTE: connection / handshake error, server is down
                            Err(err)
                        }
                        None => {
                            // Timeout
                            debug!(
                                "checked remote server {} latency timeout, elapsed {} ms",
                                sc.config.addr(),
                                elapsed
                            );

                            // NOTE: timeout is still available, but server is too slow
                            Ok(elapsed)
                        }
                    }
                }
            }
        })
    }

    fn best_idx(&self) -> usize {
        self.best_idx.load(Ordering::Acquire)
    }

    fn set_best_idx(&self, idx: usize) {
        self.best_idx.store(idx, Ordering::Release)
    }

    fn best_server(&self) -> &Arc<Server> {
        &self.servers[self.best_idx()]
    }
}

#[derive(Clone)]
pub struct PingBalancer {
    inner: Arc<Inner>,
}

impl PingBalancer {
    pub fn new(context: SharedContext) -> PingBalancer {
        let inner = Arc::new(Inner::new(context));

        let inner_cloned = inner.clone();
        tokio::spawn(
            Interval::new(
                Instant::now() + Duration::from_secs(2),
                Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC),
            )
            .for_each(move |_| {
                let inner = inner_cloned.clone();

                if inner.servers.is_empty() {
                    panic!("No server");
                }

                // Choose the best one
                let mut svr_idx = 0;

                for (idx, svr) in inner.servers.iter().enumerate() {
                    let choosen_svr = &inner.servers[svr_idx];
                    if choosen_svr.delay() == 0 {
                        continue;
                    }

                    if svr.delay() < choosen_svr.delay() {
                        svr_idx = idx;
                    }
                }

                let choosen_svr = &inner.servers[svr_idx];
                debug!(
                    "chosen the best server {} (delay: {} ms)",
                    choosen_svr.config.addr(),
                    choosen_svr.delay()
                );

                if inner.best_idx() != svr_idx {
                    info!(
                        "switched server from {} (latency: {} ms) to {} (latency: {} ms)",
                        inner.best_server().config.addr(),
                        inner.best_server().delay(),
                        choosen_svr.config.addr(),
                        choosen_svr.delay(),
                    );
                }

                inner.set_best_idx(svr_idx);

                Ok(())
            })
            .map_err(|_| ()),
        );

        PingBalancer { inner }
    }
}

impl LoadBalancer for PingBalancer {
    fn pick_server(&mut self) -> Arc<ServerConfig> {
        if self.inner.servers.is_empty() {
            panic!("no server in configuration");
        }
        self.inner.best_server().config.clone()
    }

    fn total(&self) -> usize {
        self.inner.servers.len()
    }
}
<|endoftext|>


<commit_before>//!
//! The `secretbox` module symmetrically encrypts given plaintext and
//! then uses a one time authenticator to ensure tamper-resistance.
//!
//! In other words, this uses an encrypt-then-MAC scheme.
//!
//! ## Usage
//!
//! ```rust{.example}
//! use tweetnacl::secretbox::{SecretKey, SecretMsg};
//!
//! let key = SecretKey::from_str("some secret key");
//! let enc: SecretMsg = key.encrypt("my secret msg".as_bytes());
//!
//! // ...
//!
//! let decr_opt = key.decrypt(&enc);
//! println!("decrypted: {}", decr_opt.unwrap());
//! ```
//!
//! TODO: document me.

use bindings::*;
use std::slice::bytes::copy_memory;

/// Size of shared secret key used for symmetric encryption.
pub static KEY_BYTES: uint = 32;
/// Size of the nonce value.
pub static NONCE_BYTES: uint = 24;
/// Size of the zero padding applied to each message.
pub static ZERO_BYTES: uint = 32;

/// Encapsulates both the nonce value and cipher text returned by `encrypt`.
pub struct SecretMsg {
    /// Nonce value used for this ciphertext.
    pub nonce: [u8, ..NONCE_BYTES],
    pub cipher: Vec<u8>
}

/// Shared secret key. Must be `<= KEY_BYTES` bytes long.
///
/// This struct wraps access to encrypting and decrypting messages.
pub struct SecretKey ([u8, ..KEY_BYTES]);

impl SecretKey {
    /// Generate a secret key from the bytes of a given string.
    pub fn from_str(str: &str) -> SecretKey {
        SecretKey::from_slice(str.as_bytes())
    }

    /// Generate a secret key from a slice (turn a slice into a sized slice).
    pub fn from_slice(slice: &[u8]) -> SecretKey {
        assert!(slice.len() <= KEY_BYTES);

        let mut sized = [0u8, ..KEY_BYTES];
        copy_memory(sized, slice);

        SecretKey(sized)
    }

    /// Using this secret key, symmetrically encrypt the given message.
    ///
    /// A random nonce value will be securely generated and returned
    /// as part of the response.
    pub fn encrypt(&self, msg: &[u8]) -> SecretMsg {
        let mut stretched: Vec<u8> = Vec::from_elem(ZERO_BYTES, 0u8);
        stretched.push_all(msg);

        let mut nonce = [0u8, ..NONCE_BYTES];
        let mut cipher = Vec::from_elem(stretched.len(), 0u8);
        let &SecretKey(sk) = self;

        unsafe {
            randombytes(nonce.as_mut_ptr(), NONCE_BYTES as u64);

            // TODO: Better error handling
            match crypto_secretbox(cipher.as_mut_ptr(),
                                   stretched.as_ptr(),
                                   stretched.len() as u64,
                                   nonce.as_ptr(),
                                   sk.as_ptr()) {
                0 => SecretMsg { nonce: nonce, cipher: cipher },
                _ => fail!("crypto_secretbox failed")
            }
        }
    }

    /// Using this box's secret key, decrypt the given ciphertext into
    /// plain text.
    ///
    /// If the cipher text fails to conform to the MAC (message was
    /// tampered with or corrupted), then None will be returned instead.
    pub fn decrypt(&self, msg: &SecretMsg) -> Option<Vec<u8>> {
        let &SecretKey(sk) = self;
        let mut plaintext = Vec::from_elem(msg.cipher.len(), 0u8);

        unsafe {
            match crypto_secretbox_open(plaintext.as_mut_ptr(),
                                        msg.cipher.as_ptr(),
                                        msg.cipher.len() as u64,
                                        msg.nonce.as_ptr(),
                                        sk.as_ptr()) {
                0 => Some(Vec::from_slice(plaintext.slice(ZERO_BYTES, plaintext.len()))),
                -2 => None,
                _ => fail!("crypto_secretbox_open failed")
            }
        }
    }
}

#[test]
fn test_secretbox_sanity() {
    for i in range(0 as uint, 16) {
        let msg = Vec::from_elem(i * 4, i as u8);

        let key = SecretKey::from_str("some secret key");
        let SecretMsg { nonce, cipher } = key.encrypt(msg.as_slice());

        println!("enc:\t{}\nnonce:\t{}", cipher, Vec::from_slice(nonce));

        let decr_opt = key.decrypt(&SecretMsg { nonce: nonce, cipher: cipher });

        assert!(decr_opt.is_some());

        let decr = decr_opt.unwrap();
        println!("dec:\t{}", decr);

        assert!(msg == decr);
    }
}

#[test]
fn test_secretbox_uniqueness() {
    let msg = Vec::from_elem(128, 0x53u8);

    let key1 = SecretKey::from_str("1");
    let key2 = SecretKey::from_str("");

    let SecretMsg { nonce: n1, cipher: c1 } = key1.encrypt(msg.as_slice());
    let SecretMsg { nonce: n2, cipher: c2 } = key2.encrypt(msg.as_slice());

    assert!(n1 != n2);
    assert!(c1 != c2);
    assert!(c1 != msg);
    assert!(c2 != msg);
}

#[test]
fn test_secretbox_mac_sanity() {

    let msg = Vec::from_elem(0xff, 0xff as u8);

    let key = SecretKey::from_str("some secret key");

    let SecretMsg { nonce, cipher } = key.encrypt(msg.as_slice());

    let mut ciphers = [cipher.clone(), cipher.clone(), cipher.clone()];

    // tamper with the cipher text in various ways
    ciphers[0].push(0u8);
    ciphers[1].pop();

    let last = ciphers[2].pop().unwrap();
    ciphers[2].push(last + 1);

    for c in ciphers.iter() {
        let decr = key.decrypt(&SecretMsg { nonce: nonce, cipher: c.clone() });

        println!("cipher:\t{}\ndecr:\t{}", c, decr);
        assert!(decr.is_none());
    }

}
<commit_msg>Add SecretMsg from_bytes/as_bytes<commit_after>//!
//! The `secretbox` module symmetrically encrypts given plaintext and
//! then uses a one time authenticator to ensure tamper-resistance.
//!
//! In other words, this uses an encrypt-then-MAC scheme.
//!
//! ## Usage
//!
//! ```rust{.example}
//! use tweetnacl::secretbox::{SecretKey, SecretMsg};
//!
//! let key = SecretKey::from_str("some secret key");
//! let enc: SecretMsg = key.encrypt("my secret msg".as_bytes());
//!
//! // ...
//!
//! let decr_opt = key.decrypt(&enc);
//! println!("decrypted: {}", decr_opt.unwrap());
//! ```
//!
//! TODO: document me.

use bindings::*;
use std::slice::bytes::copy_memory;

/// Size of shared secret key used for symmetric encryption.
pub static KEY_BYTES: uint = 32;
/// Size of the nonce value.
pub static NONCE_BYTES: uint = 24;
/// Size of the zero padding applied to each message.
pub static ZERO_BYTES: uint = 32;

/// Encapsulates both the nonce value and cipher text returned by `encrypt`.
pub struct SecretMsg {
    /// Nonce value used for this ciphertext.
    pub nonce: [u8, ..NONCE_BYTES],
    pub cipher: Vec<u8>
}

impl SecretMsg {
    pub fn from_bytes(bytes: &[u8]) -> Option<SecretMsg> {
        if bytes.len() <= NONCE_BYTES + ZERO_BYTES {
            return None
        }

        let mut nonce = [0u8, ..NONCE_BYTES];
        let cipher = bytes.slice_from(NONCE_BYTES);

        copy_memory(nonce, bytes.slice(0, NONCE_BYTES));

        Some(SecretMsg { nonce: nonce, cipher: Vec::from_slice(cipher) })
    }

    pub fn as_bytes(&self) -> Vec<u8> {
        let mut buf = Vec::from_slice(self.nonce);
        buf.push_all(self.cipher.as_slice());

        buf
    }
}

/// Shared secret key. Must be `<= KEY_BYTES` bytes long.
///
/// This struct wraps access to encrypting and decrypting messages.
pub struct SecretKey ([u8, ..KEY_BYTES]);

impl SecretKey {
    /// Generate a secret key from the bytes of a given string.
    pub fn from_str(str: &str) -> SecretKey {
        SecretKey::from_slice(str.as_bytes())
    }

    /// Generate a secret key from a slice (turn a slice into a sized slice).
    pub fn from_slice(slice: &[u8]) -> SecretKey {
        assert!(slice.len() <= KEY_BYTES);

        let mut sized = [0u8, ..KEY_BYTES];
        copy_memory(sized, slice);

        SecretKey(sized)
    }

    /// Using this secret key, symmetrically encrypt the given message.
    ///
    /// A random nonce value will be securely generated and returned
    /// as part of the response.
    pub fn encrypt(&self, msg: &[u8]) -> SecretMsg {
        let mut stretched: Vec<u8> = Vec::from_elem(ZERO_BYTES, 0u8);
        stretched.push_all(msg);

        let mut nonce = [0u8, ..NONCE_BYTES];
        let mut cipher = Vec::from_elem(stretched.len(), 0u8);
        let &SecretKey(sk) = self;

        unsafe {
            randombytes(nonce.as_mut_ptr(), NONCE_BYTES as u64);

            // TODO: Better error handling
            match crypto_secretbox(cipher.as_mut_ptr(),
                                   stretched.as_ptr(),
                                   stretched.len() as u64,
                                   nonce.as_ptr(),
                                   sk.as_ptr()) {
                0 => SecretMsg { nonce: nonce, cipher: cipher },
                _ => fail!("crypto_secretbox failed")
            }
        }
    }

    /// Using this box's secret key, decrypt the given ciphertext into
    /// plain text.
    ///
    /// If the cipher text fails to conform to the MAC (message was
    /// tampered with or corrupted), then None will be returned instead.
    pub fn decrypt(&self, msg: &SecretMsg) -> Option<Vec<u8>> {
        let &SecretKey(sk) = self;
        let mut plaintext = Vec::from_elem(msg.cipher.len(), 0u8);

        unsafe {
            match crypto_secretbox_open(plaintext.as_mut_ptr(),
                                        msg.cipher.as_ptr(),
                                        msg.cipher.len() as u64,
                                        msg.nonce.as_ptr(),
                                        sk.as_ptr()) {
                0 => Some(Vec::from_slice(plaintext.slice(ZERO_BYTES, plaintext.len()))),
                -2 => None,
                _ => fail!("crypto_secretbox_open failed")
            }
        }
    }
}

#[test]
fn test_secretbox_sanity() {
    for i in range(0 as uint, 16) {
        let msg = Vec::from_elem(i * 4, i as u8);

        let key = SecretKey::from_str("some secret key");
        let SecretMsg { nonce, cipher } = key.encrypt(msg.as_slice());

        println!("enc:\t{}\nnonce:\t{}", cipher, Vec::from_slice(nonce));

        let decr_opt = key.decrypt(&SecretMsg { nonce: nonce, cipher: cipher });

        assert!(decr_opt.is_some());

        let decr = decr_opt.unwrap();
        println!("dec:\t{}", decr);

        assert!(msg == decr);
    }
}

#[test]
fn test_secretbox_uniqueness() {
    let msg = Vec::from_elem(128, 0x53u8);

    let key1 = SecretKey::from_str("1");
    let key2 = SecretKey::from_str("");

    let SecretMsg { nonce: n1, cipher: c1 } = key1.encrypt(msg.as_slice());
    let SecretMsg { nonce: n2, cipher: c2 } = key2.encrypt(msg.as_slice());

    assert!(n1 != n2);
    assert!(c1 != c2);
    assert!(c1 != msg);
    assert!(c2 != msg);
}

#[test]
fn test_secretbox_mac_sanity() {

    let msg = Vec::from_elem(0xff, 0xff as u8);

    let key = SecretKey::from_str("some secret key");

    let SecretMsg { nonce, cipher } = key.encrypt(msg.as_slice());

    let mut ciphers = [cipher.clone(), cipher.clone(), cipher.clone()];

    // tamper with the cipher text in various ways
    ciphers[0].push(0u8);
    ciphers[1].pop();

    let last = ciphers[2].pop().unwrap();
    ciphers[2].push(last + 1);

    for c in ciphers.iter() {
        let decr = key.decrypt(&SecretMsg { nonce: nonce, cipher: c.clone() });

        println!("cipher:\t{}\ndecr:\t{}", c, decr);
        assert!(decr.is_none());
    }

}

#[test]
fn test_secretbox_secretmsg() {
    let msg = b"some message";
    let key = SecretKey::from_str("some secret key");
    let encr = key.encrypt(msg);

    let secretMsg = encr.as_bytes();
    let re_encr = SecretMsg::from_bytes(secretMsg.as_slice());

    assert!(re_encr.is_some());

    let decr_opt = key.decrypt(&re_encr.unwrap());

    assert!(decr_opt.is_some());
    assert!(decr_opt.unwrap().as_slice() == msg);
}
<|endoftext|>


<commit_before>#![deny(elided_lifetimes_in_paths)]
#![allow(unused)] // todo remove

mod deconstruct_pat;
// TODO: find a better place for this?
mod pat_util;
pub mod usefulness;

use hir_def::{body::Body, EnumVariantId, LocalFieldId, VariantId};
use la_arena::Idx;

use crate::{db::HirDatabase, AdtId, InferenceResult, Interner, Substitution, Ty, TyKind};

use self::{deconstruct_pat::ToDo, pat_util::EnumerateAndAdjustIterator};

pub type PatId = Idx<Pat>;

#[derive(Clone, Debug)]
pub(crate) enum PatternError {
    Unimplemented,
}

#[derive(Clone, Debug, PartialEq)]
pub struct FieldPat {
    pub field: LocalFieldId,
    pub pattern: Pat,
}

#[derive(Clone, Debug, PartialEq)]
pub struct Pat {
    pub ty: Ty,
    pub kind: Box<PatKind>,
}

impl Pat {
    pub(crate) fn wildcard_from_ty(ty: &Ty) -> Self {
        Pat { ty: ty.clone(), kind: Box::new(PatKind::Wild) }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum PatKind {
    Wild,

    /// `x`, `ref x`, `x @ P`, etc.
    Binding {
        subpattern: Option<Pat>,
        // todo: ToDo,
    },

    /// `Foo(...)` or `Foo{...}` or `Foo`, where `Foo` is a variant name from an ADT with
    /// multiple variants.
    Variant {
        substs: Substitution,
        enum_variant: EnumVariantId,
        subpatterns: Vec<FieldPat>,
    },

    /// `(...)`, `Foo(...)`, `Foo{...}`, or `Foo`, where `Foo` is a variant name from an ADT with
    /// a single variant.
    Leaf {
        subpatterns: Vec<FieldPat>,
    },

    /// `box P`, `&P`, `&mut P`, etc.
    Deref {
        subpattern: Pat,
    },

    /// An or-pattern, e.g. `p | q`.
    /// Invariant: `pats.len() >= 2`.
    Or {
        pats: Vec<Pat>,
    },
}

pub(crate) struct PatCtxt<'a> {
    db: &'a dyn HirDatabase,
    infer: &'a InferenceResult,
    body: &'a Body,
    pub(crate) errors: Vec<PatternError>,
}

impl<'a> PatCtxt<'a> {
    pub(crate) fn new(db: &'a dyn HirDatabase, infer: &'a InferenceResult, body: &'a Body) -> Self {
        Self { db, infer, body, errors: Vec::new() }
    }

    pub(crate) fn lower_pattern(&mut self, pat: hir_def::expr::PatId) -> Pat {
        // TODO: pattern adjustments (implicit dereference)
        // More info https://github.com/rust-lang/rust/issues/42640#issuecomment-313535089
        let unadjusted_pat = self.lower_pattern_unadjusted(pat);
        unadjusted_pat
    }

    fn lower_pattern_unadjusted(&mut self, pat: hir_def::expr::PatId) -> Pat {
        let ty = &self.infer[pat];

        let kind = match self.body[pat] {
            hir_def::expr::Pat::Wild => PatKind::Wild,

            hir_def::expr::Pat::Tuple { ref args, ellipsis } => {
                let arity = match *ty.kind(&Interner) {
                    TyKind::Tuple(arity, _) => arity,
                    _ => panic!("unexpected type for tuple pattern: {:?}", ty),
                };
                let subpatterns = self.lower_tuple_subpats(args, arity, ellipsis);
                PatKind::Leaf { subpatterns }
            }

            hir_def::expr::Pat::TupleStruct { ref args, ellipsis, .. } => {
                let variant_data = match self.infer.variant_resolution_for_pat(pat) {
                    Some(variant_id) => variant_id.variant_data(self.db.upcast()),
                    None => panic!("tuple struct pattern not applied to an ADT {:?}", ty),
                };
                let subpatterns =
                    self.lower_tuple_subpats(args, variant_data.fields().len(), ellipsis);
                self.lower_variant_or_leaf(pat, ty, subpatterns)
            }

            hir_def::expr::Pat::Record { ref args, .. } => {
                let variant_data = match self.infer.variant_resolution_for_pat(pat) {
                    Some(variant_id) => variant_id.variant_data(self.db.upcast()),
                    None => panic!("record pattern not applied to an ADT {:?}", ty),
                };
                let subpatterns = args
                    .iter()
                    .map(|field| FieldPat {
                        // XXX(iDawer): field lookup is inefficient
                        field: variant_data.field(&field.name).unwrap_or_else(|| todo!()),
                        pattern: self.lower_pattern(field.pat),
                    })
                    .collect();
                self.lower_variant_or_leaf(pat, ty, subpatterns)
            }

            hir_def::expr::Pat::Or(ref pats) => PatKind::Or { pats: self.lower_patterns(pats) },

            _ => {
                self.errors.push(PatternError::Unimplemented);
                PatKind::Wild
            }
        };

        Pat { ty: ty.clone(), kind: Box::new(kind) }
    }

    fn lower_tuple_subpats(
        &mut self,
        pats: &[hir_def::expr::PatId],
        expected_len: usize,
        ellipsis: Option<usize>,
    ) -> Vec<FieldPat> {
        pats.iter()
            .enumerate_and_adjust(expected_len, ellipsis)
            .map(|(i, &subpattern)| FieldPat {
                field: LocalFieldId::from_raw((i as u32).into()),
                pattern: self.lower_pattern(subpattern),
            })
            .collect()
    }

    fn lower_patterns(&mut self, pats: &[hir_def::expr::PatId]) -> Vec<Pat> {
        pats.iter().map(|&p| self.lower_pattern(p)).collect()
    }

    fn lower_variant_or_leaf(
        &mut self,
        pat: hir_def::expr::PatId,
        ty: &Ty,
        subpatterns: Vec<FieldPat>,
    ) -> PatKind {
        let kind = match self.infer.variant_resolution_for_pat(pat) {
            Some(variant_id) => {
                if let VariantId::EnumVariantId(enum_variant) = variant_id {
                    let substs = match ty.kind(&Interner) {
                        TyKind::Adt(_, substs) | TyKind::FnDef(_, substs) => substs.clone(),
                        TyKind::Error => {
                            return PatKind::Wild;
                        }
                        _ => panic!("inappropriate type for def: {:?}", ty),
                    };
                    PatKind::Variant { substs, enum_variant, subpatterns }
                } else {
                    PatKind::Leaf { subpatterns }
                }
            }
            None => {
                self.errors.push(PatternError::Unimplemented);
                PatKind::Wild
            }
        };
        // TODO: do we need PatKind::AscribeUserType ?
        kind
    }
}

#[cfg(test)]
mod tests {
    use crate::diagnostics::tests::check_diagnostics;

    use super::*;

    #[test]
    fn unit() {
        check_diagnostics(
            r#"
fn main() {
    match () { () => {} }
    match () {  _ => {} }
    match () {          }
        //^^ Missing match arm
}
"#,
        );
    }

    #[test]
    fn tuple_of_units() {
        check_diagnostics(
            r#"
fn main() {
    match ((), ()) { ((), ()) => {} }
    match ((), ()) {  ((), _) => {} }
    match ((), ()) {   (_, _) => {} }
    match ((), ()) {        _ => {} }
    match ((), ()) {                }
        //^^^^^^^^ Missing match arm
}
"#,
        );
    }

    #[test]
    fn tuple_with_ellipsis() {
        // TODO: test non-exhaustive match with ellipsis in the middle
        // of a pattern, check reported witness
        check_diagnostics(
            r#"
struct A; struct B;
fn main(v: (A, (), B)) {
    match v { (A, ..)    => {} }
    match v { (.., B)    => {} }
    match v { (A, .., B) => {} }
    match v { (..)       => {} }
    match v {                  }
        //^ Missing match arm
}
"#,
        );
    }

    #[test]
    fn strukt() {
        check_diagnostics(
            r#"
struct A; struct B;
struct S { a: A, b: B}
fn main(v: S) {
    match v { S { a, b }       => {} }
    match v { S { a: _, b: _ } => {} }
    match v { S { .. }         => {} }
    match v { _                => {} }
    match v {                        }
        //^ Missing match arm
}
"#,
        );
    }

    #[test]
    fn c_enum() {
        check_diagnostics(
            r#"
enum E { A, B }
fn main(v: E) {
    match v { E::A | E::B => {} }
    match v { _           => {} }
    match v { E::A        => {} }
        //^ Missing match arm
    match v {                   }
        //^ Missing match arm
}
"#,
        );
    }

    #[test]
    fn enum_() {
        check_diagnostics(
            r#"
struct A; struct B;
enum E { Tuple(A, B), Struct{ a: A, b: B } }
fn main(v: E) {
    match v {
        E::Tuple(a, b)    => {}
        E::Struct{ a, b } => {}
    }
    match v {
        E::Tuple(_, _) => {}
        E::Struct{..}  => {}
    }
    match v {
        E::Tuple(..) => {}
        _ => {}
    }
    match v { E::Tuple(..) => {} }
        //^ Missing match arm
    match v { }
        //^ Missing match arm
}
"#,
        );
    }

    #[test]
    fn boolean() {
        check_diagnostics(
            r#"
fn main() {
    match true {
        true  => {}
        false => {}
    }
    match true {
        true | false => {}
    }
    match true {
        true => {}
        _ => {}
    }
    match true {}
        //^^^^ Missing match arm
    match true { true => {} }
        //^^^^ Missing match arm

}
"#,
        );
    }
}
<commit_msg>Lower Pat::Path<commit_after>#![deny(elided_lifetimes_in_paths)]
#![allow(unused)] // todo remove

mod deconstruct_pat;
// TODO: find a better place for this?
mod pat_util;
pub mod usefulness;

use hir_def::{body::Body, EnumVariantId, LocalFieldId, VariantId};
use la_arena::Idx;

use crate::{db::HirDatabase, AdtId, InferenceResult, Interner, Substitution, Ty, TyKind};

use self::{deconstruct_pat::ToDo, pat_util::EnumerateAndAdjustIterator};

pub type PatId = Idx<Pat>;

#[derive(Clone, Debug)]
pub(crate) enum PatternError {
    Unimplemented,
}

#[derive(Clone, Debug, PartialEq)]
pub struct FieldPat {
    pub field: LocalFieldId,
    pub pattern: Pat,
}

#[derive(Clone, Debug, PartialEq)]
pub struct Pat {
    pub ty: Ty,
    pub kind: Box<PatKind>,
}

impl Pat {
    pub(crate) fn wildcard_from_ty(ty: &Ty) -> Self {
        Pat { ty: ty.clone(), kind: Box::new(PatKind::Wild) }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum PatKind {
    Wild,

    /// `x`, `ref x`, `x @ P`, etc.
    Binding {
        subpattern: Option<Pat>,
        // todo: ToDo,
    },

    /// `Foo(...)` or `Foo{...}` or `Foo`, where `Foo` is a variant name from an ADT with
    /// multiple variants.
    Variant {
        substs: Substitution,
        enum_variant: EnumVariantId,
        subpatterns: Vec<FieldPat>,
    },

    /// `(...)`, `Foo(...)`, `Foo{...}`, or `Foo`, where `Foo` is a variant name from an ADT with
    /// a single variant.
    Leaf {
        subpatterns: Vec<FieldPat>,
    },

    /// `box P`, `&P`, `&mut P`, etc.
    Deref {
        subpattern: Pat,
    },

    /// An or-pattern, e.g. `p | q`.
    /// Invariant: `pats.len() >= 2`.
    Or {
        pats: Vec<Pat>,
    },
}

pub(crate) struct PatCtxt<'a> {
    db: &'a dyn HirDatabase,
    infer: &'a InferenceResult,
    body: &'a Body,
    pub(crate) errors: Vec<PatternError>,
}

impl<'a> PatCtxt<'a> {
    pub(crate) fn new(db: &'a dyn HirDatabase, infer: &'a InferenceResult, body: &'a Body) -> Self {
        Self { db, infer, body, errors: Vec::new() }
    }

    pub(crate) fn lower_pattern(&mut self, pat: hir_def::expr::PatId) -> Pat {
        // TODO: pattern adjustments (implicit dereference)
        // More info https://github.com/rust-lang/rust/issues/42640#issuecomment-313535089
        let unadjusted_pat = self.lower_pattern_unadjusted(pat);
        unadjusted_pat
    }

    fn lower_pattern_unadjusted(&mut self, pat: hir_def::expr::PatId) -> Pat {
        let ty = &self.infer[pat];

        let kind = match self.body[pat] {
            hir_def::expr::Pat::Wild => PatKind::Wild,

            hir_def::expr::Pat::Path(ref path) => {
                return self.lower_path(pat, path);
            }

            hir_def::expr::Pat::Tuple { ref args, ellipsis } => {
                let arity = match *ty.kind(&Interner) {
                    TyKind::Tuple(arity, _) => arity,
                    _ => panic!("unexpected type for tuple pattern: {:?}", ty),
                };
                let subpatterns = self.lower_tuple_subpats(args, arity, ellipsis);
                PatKind::Leaf { subpatterns }
            }

            hir_def::expr::Pat::TupleStruct { ref args, ellipsis, .. } => {
                let variant_data = match self.infer.variant_resolution_for_pat(pat) {
                    Some(variant_id) => variant_id.variant_data(self.db.upcast()),
                    None => panic!("tuple struct pattern not applied to an ADT {:?}", ty),
                };
                let subpatterns =
                    self.lower_tuple_subpats(args, variant_data.fields().len(), ellipsis);
                self.lower_variant_or_leaf(pat, ty, subpatterns)
            }

            hir_def::expr::Pat::Record { ref args, .. } => {
                let variant_data = match self.infer.variant_resolution_for_pat(pat) {
                    Some(variant_id) => variant_id.variant_data(self.db.upcast()),
                    None => panic!("record pattern not applied to an ADT {:?}", ty),
                };
                let subpatterns = args
                    .iter()
                    .map(|field| FieldPat {
                        // XXX(iDawer): field lookup is inefficient
                        field: variant_data.field(&field.name).unwrap_or_else(|| todo!()),
                        pattern: self.lower_pattern(field.pat),
                    })
                    .collect();
                self.lower_variant_or_leaf(pat, ty, subpatterns)
            }

            hir_def::expr::Pat::Or(ref pats) => PatKind::Or { pats: self.lower_patterns(pats) },

            _ => {
                self.errors.push(PatternError::Unimplemented);
                PatKind::Wild
            }
        };

        Pat { ty: ty.clone(), kind: Box::new(kind) }
    }

    fn lower_tuple_subpats(
        &mut self,
        pats: &[hir_def::expr::PatId],
        expected_len: usize,
        ellipsis: Option<usize>,
    ) -> Vec<FieldPat> {
        pats.iter()
            .enumerate_and_adjust(expected_len, ellipsis)
            .map(|(i, &subpattern)| FieldPat {
                field: LocalFieldId::from_raw((i as u32).into()),
                pattern: self.lower_pattern(subpattern),
            })
            .collect()
    }

    fn lower_patterns(&mut self, pats: &[hir_def::expr::PatId]) -> Vec<Pat> {
        pats.iter().map(|&p| self.lower_pattern(p)).collect()
    }

    fn lower_variant_or_leaf(
        &mut self,
        pat: hir_def::expr::PatId,
        ty: &Ty,
        subpatterns: Vec<FieldPat>,
    ) -> PatKind {
        let kind = match self.infer.variant_resolution_for_pat(pat) {
            Some(variant_id) => {
                if let VariantId::EnumVariantId(enum_variant) = variant_id {
                    let substs = match ty.kind(&Interner) {
                        TyKind::Adt(_, substs) | TyKind::FnDef(_, substs) => substs.clone(),
                        TyKind::Error => {
                            return PatKind::Wild;
                        }
                        _ => panic!("inappropriate type for def: {:?}", ty),
                    };
                    PatKind::Variant { substs, enum_variant, subpatterns }
                } else {
                    PatKind::Leaf { subpatterns }
                }
            }
            None => {
                self.errors.push(PatternError::Unimplemented);
                PatKind::Wild
            }
        };
        // TODO: do we need PatKind::AscribeUserType ?
        kind
    }

    fn lower_path(&mut self, pat: hir_def::expr::PatId, path: &hir_def::path::Path) -> Pat {
        let ty = &self.infer[pat];

        let pat_from_kind = |kind| Pat { ty: ty.clone(), kind: Box::new(kind) };

        match self.infer.variant_resolution_for_pat(pat) {
            Some(_) => pat_from_kind(self.lower_variant_or_leaf(pat, ty, Vec::new())),
            None => {
                self.errors.push(PatternError::Unimplemented);
                pat_from_kind(PatKind::Wild)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::diagnostics::tests::check_diagnostics;

    use super::*;

    #[test]
    fn unit() {
        check_diagnostics(
            r#"
fn main() {
    match () { () => {} }
    match () {  _ => {} }
    match () {          }
        //^^ Missing match arm
}
"#,
        );
    }

    #[test]
    fn tuple_of_units() {
        check_diagnostics(
            r#"
fn main() {
    match ((), ()) { ((), ()) => {} }
    match ((), ()) {  ((), _) => {} }
    match ((), ()) {   (_, _) => {} }
    match ((), ()) {        _ => {} }
    match ((), ()) {                }
        //^^^^^^^^ Missing match arm
}
"#,
        );
    }

    #[test]
    fn tuple_with_ellipsis() {
        // TODO: test non-exhaustive match with ellipsis in the middle
        // of a pattern, check reported witness
        check_diagnostics(
            r#"
struct A; struct B;
fn main(v: (A, (), B)) {
    match v { (A, ..)    => {} }
    match v { (.., B)    => {} }
    match v { (A, .., B) => {} }
    match v { (..)       => {} }
    match v {                  }
        //^ Missing match arm
}
"#,
        );
    }

    #[test]
    fn strukt() {
        check_diagnostics(
            r#"
struct A; struct B;
struct S { a: A, b: B}
fn main(v: S) {
    match v { S { a, b }       => {} }
    match v { S { a: _, b: _ } => {} }
    match v { S { .. }         => {} }
    match v { _                => {} }
    match v {                        }
        //^ Missing match arm
}
"#,
        );
    }

    #[test]
    fn c_enum() {
        check_diagnostics(
            r#"
enum E { A, B }
fn main(v: E) {
    match v { E::A | E::B => {} }
    match v { _           => {} }
    match v { E::A        => {} }
        //^ Missing match arm
    match v {                   }
        //^ Missing match arm
}
"#,
        );
    }

    #[test]
    fn enum_() {
        check_diagnostics(
            r#"
struct A; struct B;
enum E { Tuple(A, B), Struct{ a: A, b: B } }
fn main(v: E) {
    match v {
        E::Tuple(a, b)    => {}
        E::Struct{ a, b } => {}
    }
    match v {
        E::Tuple(_, _) => {}
        E::Struct{..}  => {}
    }
    match v {
        E::Tuple(..) => {}
        _ => {}
    }
    match v { E::Tuple(..) => {} }
        //^ Missing match arm
    match v { }
        //^ Missing match arm
}
"#,
        );
    }

    #[test]
    fn boolean() {
        check_diagnostics(
            r#"
fn main() {
    match true {
        true  => {}
        false => {}
    }
    match true {
        true | false => {}
    }
    match true {
        true => {}
        _ => {}
    }
    match true {}
        //^^^^ Missing match arm
    match true { true => {} }
        //^^^^ Missing match arm

}
"#,
        );
    }
}
<|endoftext|>


<commit_before>use format_buf::format;
use hir::{db::HirDatabase, FromSource};
use join_to_string::join;
use ra_syntax::{
    ast::{
        self, AstNode, NameOwner, StructKind, TypeAscriptionOwner, TypeParamsOwner, VisibilityOwner,
    },
    TextUnit, T,
};
use std::fmt::Write;

use crate::{Assist, AssistCtx, AssistId};

// Assist: add_new
//
// Adds a new inherent impl for a type.
//
// ```
// struct Ctx<T: Clone> {
//      data: T,<|>
// }
// ```
// ->
// ```
// struct Ctx<T: Clone> {
//      data: T,
// }
//
// impl<T: Clone> Ctx<T> {
//     fn new(data: T) -> Self { Self { data } }
// }
//
// ```
pub(crate) fn add_new(ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let strukt = ctx.find_node_at_offset::<ast::StructDef>()?;

    // We want to only apply this to non-union structs with named fields
    let field_list = match (strukt.kind(), strukt.is_union()) {
        (StructKind::Named(named), false) => named,
        _ => return None,
    };

    // Return early if we've found an existing new fn
    let impl_block = find_struct_impl(&ctx, &strukt)?;

    ctx.add_assist(AssistId("add_new"), "add new fn", |edit| {
        edit.target(strukt.syntax().text_range());

        let mut buf = String::with_capacity(512);

        if impl_block.is_some() {
            buf.push('\n');
        }

        let vis = strukt.visibility().map(|v| format!("{} ", v.syntax()));
        let vis = vis.as_ref().map(String::as_str).unwrap_or("");
        write!(&mut buf, "    {}fn new(", vis).unwrap();

        join(field_list.fields().map(|f| {
            format!(
                "{}: {}",
                f.name().unwrap().syntax().text(),
                f.ascribed_type().unwrap().syntax().text()
            )
        }))
        .separator(", ")
        .to_buf(&mut buf);

        buf.push_str(") -> Self { Self {");

        join(field_list.fields().map(|f| f.name().unwrap().syntax().text()))
            .separator(", ")
            .surround_with(" ", " ")
            .to_buf(&mut buf);

        buf.push_str("} }");

        let (start_offset, end_offset) = if let Some(impl_block) = impl_block {
            buf.push('\n');
            let start = impl_block
                .syntax()
                .descendants_with_tokens()
                .find(|t| t.kind() == T!['{'])
                .unwrap()
                .text_range()
                .end();

            (start, TextUnit::from_usize(1))
        } else {
            buf = generate_impl_text(&strukt, &buf);
            let start = strukt.syntax().text_range().end();

            (start, TextUnit::from_usize(3))
        };

        edit.set_cursor(start_offset + TextUnit::of_str(&buf) - end_offset);
        edit.insert(start_offset, buf);
    })
}

// Generates the surrounding `impl Type { <code> }` including type and lifetime
// parameters
fn generate_impl_text(strukt: &ast::StructDef, code: &str) -> String {
    let type_params = strukt.type_param_list();
    let mut buf = String::with_capacity(code.len());
    buf.push_str("\n\nimpl");
    if let Some(type_params) = &type_params {
        format!(buf, "{}", type_params.syntax());
    }
    buf.push_str(" ");
    buf.push_str(strukt.name().unwrap().text().as_str());
    if let Some(type_params) = type_params {
        let lifetime_params = type_params
            .lifetime_params()
            .filter_map(|it| it.lifetime_token())
            .map(|it| it.text().clone());
        let type_params =
            type_params.type_params().filter_map(|it| it.name()).map(|it| it.text().clone());
        join(lifetime_params.chain(type_params)).surround_with("<", ">").to_buf(&mut buf);
    }

    format!(&mut buf, " {{\n{}\n}}\n", code);

    buf
}

// Uses a syntax-driven approach to find any impl blocks for the struct that
// exist within the module/file
//
// Returns `None` if we've found an existing `new` fn
//
// FIXME: change the new fn checking to a more semantic approach when that's more
// viable (e.g. we process proc macros, etc)
fn find_struct_impl(
    ctx: &AssistCtx<impl HirDatabase>,
    strukt: &ast::StructDef,
) -> Option<Option<ast::ImplBlock>> {
    let db = ctx.db;
    let module = strukt.syntax().ancestors().find(|node| {
        ast::Module::can_cast(node.kind()) || ast::SourceFile::can_cast(node.kind())
    })?;

    let struct_ty = {
        let src = hir::Source { file_id: ctx.frange.file_id.into(), ast: strukt.clone() };
        hir::Struct::from_source(db, src).unwrap().ty(db)
    };

    let mut found_new_fn = false;

    let block = module.descendants().filter_map(ast::ImplBlock::cast).find(|impl_blk| {
        if found_new_fn {
            return false;
        }

        let src = hir::Source { file_id: ctx.frange.file_id.into(), ast: impl_blk.clone() };
        let blk = hir::ImplBlock::from_source(db, src).unwrap();

        let same_ty = blk.target_ty(db) == struct_ty;
        let not_trait_impl = blk.target_trait(db).is_none();

        found_new_fn = has_new_fn(impl_blk);

        same_ty && not_trait_impl
    });

    if found_new_fn {
        None
    } else {
        Some(block)
    }
}

fn has_new_fn(imp: &ast::ImplBlock) -> bool {
    if let Some(il) = imp.item_list() {
        for item in il.impl_items() {
            if let ast::ImplItem::FnDef(f) = item {
                if f.name().unwrap().text().eq_ignore_ascii_case("new") {
                    return true;
                }
            }
        }
    }

    false
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::helpers::{check_assist, check_assist_not_applicable, check_assist_target};

    #[test]
    #[rustfmt::skip]
    fn test_add_new() {
        // Check output of generation
        check_assist(
            add_new,
"struct Foo {<|>}",
"struct Foo {}

impl Foo {
    fn new() -> Self { Self {  } }<|>
}
",
        );
        check_assist(
            add_new,
"struct Foo<T: Clone> {<|>}",
"struct Foo<T: Clone> {}

impl<T: Clone> Foo<T> {
    fn new() -> Self { Self {  } }<|>
}
",
        );
        check_assist(
            add_new,
"struct Foo<'a, T: Foo<'a>> {<|>}",
"struct Foo<'a, T: Foo<'a>> {}

impl<'a, T: Foo<'a>> Foo<'a, T> {
    fn new() -> Self { Self {  } }<|>
}
",
        );
        check_assist(
            add_new,
"struct Foo { baz: String <|>}",
"struct Foo { baz: String }

impl Foo {
    fn new(baz: String) -> Self { Self { baz } }<|>
}
",
        );
        check_assist(
            add_new,
"struct Foo { baz: String, qux: Vec<i32> <|>}",
"struct Foo { baz: String, qux: Vec<i32> }

impl Foo {
    fn new(baz: String, qux: Vec<i32>) -> Self { Self { baz, qux } }<|>
}
",
        );

        // Check that visibility modifiers don't get brought in for fields
        check_assist(
            add_new,
"struct Foo { pub baz: String, pub qux: Vec<i32> <|>}",
"struct Foo { pub baz: String, pub qux: Vec<i32> }

impl Foo {
    fn new(baz: String, qux: Vec<i32>) -> Self { Self { baz, qux } }<|>
}
",
        );

        // Check that it reuses existing impls
        check_assist(
            add_new,
"struct Foo {<|>}

impl Foo {}
",
"struct Foo {}

impl Foo {
    fn new() -> Self { Self {  } }<|>
}
",
        );
        check_assist(
            add_new,
"struct Foo {<|>}

impl Foo {
    fn qux(&self) {}
}
",
"struct Foo {}

impl Foo {
    fn new() -> Self { Self {  } }<|>

    fn qux(&self) {}
}
",
        );

        check_assist(
            add_new,
"struct Foo {<|>}

impl Foo {
    fn qux(&self) {}
    fn baz() -> i32 {
        5
    }
}
",
"struct Foo {}

impl Foo {
    fn new() -> Self { Self {  } }<|>

    fn qux(&self) {}
    fn baz() -> i32 {
        5
    }
}
",
        );

        // Check visibility of new fn based on struct
        check_assist(
            add_new,
"pub struct Foo {<|>}",
"pub struct Foo {}

impl Foo {
    pub fn new() -> Self { Self {  } }<|>
}
",
        );
        check_assist(
            add_new,
"pub(crate) struct Foo {<|>}",
"pub(crate) struct Foo {}

impl Foo {
    pub(crate) fn new() -> Self { Self {  } }<|>
}
",
        );
    }

    #[test]
    fn add_new_not_applicable_if_fn_exists() {
        check_assist_not_applicable(
            add_new,
            "
struct Foo {<|>}

impl Foo {
    fn new() -> Self {
        Self
    }
}",
        );

        check_assist_not_applicable(
            add_new,
            "
struct Foo {<|>}

impl Foo {
    fn New() -> Self {
        Self
    }
}",
        );
    }

    #[test]
    fn add_new_target() {
        check_assist_target(
            add_new,
            "
struct SomeThingIrrelevant;
/// Has a lifetime parameter
struct Foo<'a, T: Foo<'a>> {<|>}
struct EvenMoreIrrelevant;
",
            "/// Has a lifetime parameter
struct Foo<'a, T: Foo<'a>> {}",
        );
    }
}
<commit_msg>Fix add-new assist<commit_after>use format_buf::format;
use hir::{db::HirDatabase, FromSource};
use join_to_string::join;
use ra_syntax::{
    ast::{
        self, AstNode, NameOwner, StructKind, TypeAscriptionOwner, TypeParamsOwner, VisibilityOwner,
    },
    TextUnit, T,
};
use std::fmt::Write;

use crate::{Assist, AssistCtx, AssistId};

// Assist: add_new
//
// Adds a new inherent impl for a type.
//
// ```
// struct Ctx<T: Clone> {
//      data: T,<|>
// }
// ```
// ->
// ```
// struct Ctx<T: Clone> {
//      data: T,
// }
//
// impl<T: Clone> Ctx<T> {
//     fn new(data: T) -> Self { Self { data } }
// }
//
// ```
pub(crate) fn add_new(ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let strukt = ctx.find_node_at_offset::<ast::StructDef>()?;

    // We want to only apply this to non-union structs with named fields
    let field_list = match (strukt.kind(), strukt.is_union()) {
        (StructKind::Named(named), false) => named,
        _ => return None,
    };

    // Return early if we've found an existing new fn
    let impl_block = find_struct_impl(&ctx, &strukt)?;

    ctx.add_assist(AssistId("add_new"), "add new fn", |edit| {
        edit.target(strukt.syntax().text_range());

        let mut buf = String::with_capacity(512);

        if impl_block.is_some() {
            buf.push('\n');
        }

        let vis = strukt.visibility().map(|v| format!("{} ", v.syntax()));
        let vis = vis.as_ref().map(String::as_str).unwrap_or("");
        write!(&mut buf, "    {}fn new(", vis).unwrap();

        join(field_list.fields().map(|f| {
            format!(
                "{}: {}",
                f.name().unwrap().syntax().text(),
                f.ascribed_type().unwrap().syntax().text()
            )
        }))
        .separator(", ")
        .to_buf(&mut buf);

        buf.push_str(") -> Self { Self {");

        join(field_list.fields().map(|f| f.name().unwrap().syntax().text()))
            .separator(", ")
            .surround_with(" ", " ")
            .to_buf(&mut buf);

        buf.push_str("} }");

        let (start_offset, end_offset) = if let Some(impl_block) = impl_block {
            buf.push('\n');
            let start = impl_block
                .syntax()
                .descendants_with_tokens()
                .find(|t| t.kind() == T!['{'])
                .unwrap()
                .text_range()
                .end();

            (start, TextUnit::from_usize(1))
        } else {
            buf = generate_impl_text(&strukt, &buf);
            let start = strukt.syntax().text_range().end();

            (start, TextUnit::from_usize(3))
        };

        edit.set_cursor(start_offset + TextUnit::of_str(&buf) - end_offset);
        edit.insert(start_offset, buf);
    })
}

// Generates the surrounding `impl Type { <code> }` including type and lifetime
// parameters
fn generate_impl_text(strukt: &ast::StructDef, code: &str) -> String {
    let type_params = strukt.type_param_list();
    let mut buf = String::with_capacity(code.len());
    buf.push_str("\n\nimpl");
    if let Some(type_params) = &type_params {
        format!(buf, "{}", type_params.syntax());
    }
    buf.push_str(" ");
    buf.push_str(strukt.name().unwrap().text().as_str());
    if let Some(type_params) = type_params {
        let lifetime_params = type_params
            .lifetime_params()
            .filter_map(|it| it.lifetime_token())
            .map(|it| it.text().clone());
        let type_params =
            type_params.type_params().filter_map(|it| it.name()).map(|it| it.text().clone());
        join(lifetime_params.chain(type_params)).surround_with("<", ">").to_buf(&mut buf);
    }

    format!(&mut buf, " {{\n{}\n}}\n", code);

    buf
}

// Uses a syntax-driven approach to find any impl blocks for the struct that
// exist within the module/file
//
// Returns `None` if we've found an existing `new` fn
//
// FIXME: change the new fn checking to a more semantic approach when that's more
// viable (e.g. we process proc macros, etc)
fn find_struct_impl(
    ctx: &AssistCtx<impl HirDatabase>,
    strukt: &ast::StructDef,
) -> Option<Option<ast::ImplBlock>> {
    let db = ctx.db;
    let module = strukt.syntax().ancestors().find(|node| {
        ast::Module::can_cast(node.kind()) || ast::SourceFile::can_cast(node.kind())
    })?;

    let struct_ty = {
        let src = hir::Source { file_id: ctx.frange.file_id.into(), ast: strukt.clone() };
        hir::Struct::from_source(db, src).unwrap().ty(db)
    };

    let mut found_new_fn = false;

    let block = module.descendants().filter_map(ast::ImplBlock::cast).find(|impl_blk| {
        if found_new_fn {
            return false;
        }

        let src = hir::Source { file_id: ctx.frange.file_id.into(), ast: impl_blk.clone() };
        let blk = hir::ImplBlock::from_source(db, src).unwrap();

        let same_ty = blk.target_ty(db) == struct_ty;
        let not_trait_impl = blk.target_trait(db).is_none();

        if !(same_ty && not_trait_impl) {
            return false;
        }

        found_new_fn = has_new_fn(impl_blk);
        true
    });

    if found_new_fn {
        None
    } else {
        Some(block)
    }
}

fn has_new_fn(imp: &ast::ImplBlock) -> bool {
    if let Some(il) = imp.item_list() {
        for item in il.impl_items() {
            if let ast::ImplItem::FnDef(f) = item {
                if f.name().unwrap().text().eq_ignore_ascii_case("new") {
                    return true;
                }
            }
        }
    }

    false
}

#[cfg(test)]
mod tests {
    use crate::helpers::{check_assist, check_assist_not_applicable, check_assist_target};

    use super::*;

    #[test]
    #[rustfmt::skip]
    fn test_add_new() {
        // Check output of generation
        check_assist(
            add_new,
"struct Foo {<|>}",
"struct Foo {}

impl Foo {
    fn new() -> Self { Self {  } }<|>
}
",
        );
        check_assist(
            add_new,
"struct Foo<T: Clone> {<|>}",
"struct Foo<T: Clone> {}

impl<T: Clone> Foo<T> {
    fn new() -> Self { Self {  } }<|>
}
",
        );
        check_assist(
            add_new,
"struct Foo<'a, T: Foo<'a>> {<|>}",
"struct Foo<'a, T: Foo<'a>> {}

impl<'a, T: Foo<'a>> Foo<'a, T> {
    fn new() -> Self { Self {  } }<|>
}
",
        );
        check_assist(
            add_new,
"struct Foo { baz: String <|>}",
"struct Foo { baz: String }

impl Foo {
    fn new(baz: String) -> Self { Self { baz } }<|>
}
",
        );
        check_assist(
            add_new,
"struct Foo { baz: String, qux: Vec<i32> <|>}",
"struct Foo { baz: String, qux: Vec<i32> }

impl Foo {
    fn new(baz: String, qux: Vec<i32>) -> Self { Self { baz, qux } }<|>
}
",
        );

        // Check that visibility modifiers don't get brought in for fields
        check_assist(
            add_new,
"struct Foo { pub baz: String, pub qux: Vec<i32> <|>}",
"struct Foo { pub baz: String, pub qux: Vec<i32> }

impl Foo {
    fn new(baz: String, qux: Vec<i32>) -> Self { Self { baz, qux } }<|>
}
",
        );

        // Check that it reuses existing impls
        check_assist(
            add_new,
"struct Foo {<|>}

impl Foo {}
",
"struct Foo {}

impl Foo {
    fn new() -> Self { Self {  } }<|>
}
",
        );
        check_assist(
            add_new,
"struct Foo {<|>}

impl Foo {
    fn qux(&self) {}
}
",
"struct Foo {}

impl Foo {
    fn new() -> Self { Self {  } }<|>

    fn qux(&self) {}
}
",
        );

        check_assist(
            add_new,
"struct Foo {<|>}

impl Foo {
    fn qux(&self) {}
    fn baz() -> i32 {
        5
    }
}
",
"struct Foo {}

impl Foo {
    fn new() -> Self { Self {  } }<|>

    fn qux(&self) {}
    fn baz() -> i32 {
        5
    }
}
",
        );

        // Check visibility of new fn based on struct
        check_assist(
            add_new,
"pub struct Foo {<|>}",
"pub struct Foo {}

impl Foo {
    pub fn new() -> Self { Self {  } }<|>
}
",
        );
        check_assist(
            add_new,
"pub(crate) struct Foo {<|>}",
"pub(crate) struct Foo {}

impl Foo {
    pub(crate) fn new() -> Self { Self {  } }<|>
}
",
        );
    }

    #[test]
    fn add_new_not_applicable_if_fn_exists() {
        check_assist_not_applicable(
            add_new,
            "
struct Foo {<|>}

impl Foo {
    fn new() -> Self {
        Self
    }
}",
        );

        check_assist_not_applicable(
            add_new,
            "
struct Foo {<|>}

impl Foo {
    fn New() -> Self {
        Self
    }
}",
        );
    }

    #[test]
    fn add_new_target() {
        check_assist_target(
            add_new,
            "
struct SomeThingIrrelevant;
/// Has a lifetime parameter
struct Foo<'a, T: Foo<'a>> {<|>}
struct EvenMoreIrrelevant;
",
            "/// Has a lifetime parameter
struct Foo<'a, T: Foo<'a>> {}",
        );
    }

    #[test]
    fn test_unrelated_new() {
        check_assist(
            add_new,
            r##"
pub struct AstId<N: AstNode> {
    file_id: HirFileId,
    file_ast_id: FileAstId<N>,
}

impl<N: AstNode> AstId<N> {
    pub fn new(file_id: HirFileId, file_ast_id: FileAstId<N>) -> AstId<N> {
        AstId { file_id, file_ast_id }
    }
}

pub struct Source<T> {
    pub file_id: HirFileId,<|>
    pub ast: T,
}

impl<T> Source<T> {
    pub fn map<F: FnOnce(T) -> U, U>(self, f: F) -> Source<U> {
        Source { file_id: self.file_id, ast: f(self.ast) }
    }
}
"##,
            r##"
pub struct AstId<N: AstNode> {
    file_id: HirFileId,
    file_ast_id: FileAstId<N>,
}

impl<N: AstNode> AstId<N> {
    pub fn new(file_id: HirFileId, file_ast_id: FileAstId<N>) -> AstId<N> {
        AstId { file_id, file_ast_id }
    }
}

pub struct Source<T> {
    pub file_id: HirFileId,
    pub ast: T,
}

impl<T> Source<T> {
    pub fn new(file_id: HirFileId, ast: T) -> Self { Self { file_id, ast } }<|>

    pub fn map<F: FnOnce(T) -> U, U>(self, f: F) -> Source<U> {
        Source { file_id: self.file_id, ast: f(self.ast) }
    }
}
"##,
        );
    }
}
<|endoftext|>


<commit_before>type Asn1Len = u64;
type Asn1TagNum = u64;

type Asn1Data = Vec<u8>;
type Asn1Slice<'a> = &'a [u8];

enum Asn1Tag {
  // Universal types here ...
  UniversalTag(Asn1TagNum)
  AplicationTag(Asn1TagNum),
  PrivateTag(Asn1TagNum),
  ContextSpecificTag(Asn1TagNum),
}

impl Asn1Tag {
  /// Returns true when this is a structured type.
  fn is_structured(&self) -> bool {
    if let Asn1Tag::UniversalTag(tag) == *self {
      match tag {
        // SEQUENCE (OF)
        16 => true,
        // SET (OF)
        17 => true,
        _ => false,

      }
    } else {
      false
    }
  }
}

type Asn1Type = String;

trait Asn1Data {
  fn get_asn1_type() -> Asn1Type;

  /// Create ASN.1 data from this struct.
  // FIXME: Should this use &self?
  fn into_asn1(&self) -> Result<Asn1Data, Asn1Error>;

  /// Create this struct from ASN.1 data.
  fn from_asn1(slice: Asn1Slice) -> Result<Self, Asn1Error>;
}

/// A macro to generate a generic Asn1Data trait implementation for a struct.
macro_rules! asn1_impl {
  ( $impl_type:ident, $asn1_type:expr,
    $( $name:ident, $rusttype:ident, $asn1type:expr ),*
  ) =>
(

impl Asn1Data for $impl_type {
  fn get_asn1_type() -> Asn1Type {
    $asn1_type
  }

  fn into_asn1(&self) -> Result<Asn1Data, Asn1Error> {
    Err(Asn1Error::EncodingError)
  }

  fn from_asn1(slice: Asn1Slice) -> Result<Self, Asn1Error> {
    Err(Asn1Error::InvalidAsn1)
  }
}

)
}

/// A list of errors that can occur decoding or encoding Asn1 data.
enum Asn1Error {
  /// Invalid Asn1 data.
  InvalidAsn1
  /// An error occured while encoding Asn1 data.
  EncodingError
  /// An invalid tag was decoded
  InvalidTag(Asn1Tag)
}
<commit_msg>Redo types<commit_after>/// An enum representing the length of an ASN.1 element.
enum Asn1Len {
  /// A Definite length element.
  pub Definite(u64),
  /// An Indefinite length element, not known before decoding.
  pub Indefinite,
}

impl From<u64> for Asn1Len {
  fn from(len: u64) -> Self {
    match len {
      l => Asn1Len::Def(l),
      0 => Asn1Len::Indef,
    }
  }
}

impl From<Asn1Len> for u64 {
  fn from(len: Asn1Len) -> Self {
    match len {
      Asn1Len::Def(l) => l,
      Asn1Len::Indef => 0,
    }
  }
}

/// An ASN.1 tag number.
type Asn1TagNum = u64;

/// An ASN.1 Class.
enum Asn1Class {
  /// Universal class.
  Universal,
  /// Application class.
  Aplication,
  /// Private class.
  Private,
  /// Context-specific class.
  ContextSpecific,
}

/// A struct representing an ASN.1 element.
struct Asn1Tag {
  /// The class of the ASN.1 element.
  pub class: Asn1Class,
  /// The tag number of the ASN.1 element.
  pub tagnum: Asn1TagNum,
  /// The length of the ASN.1 element.
  pub len: Asn1Len,
  /// A flag indicating whether an element is constructed.
  pub constructed: bool,
}

impl Asn1Tag {
  /// Returns true when this is a structured type.
  fn is_structured(&self) -> bool {
    if let Asn1Class::Universal == *self.class {
      match self.tagnum {
        // SEQUENCE (OF)
        16 => true,
        // SET (OF)
        17 => true,
        _ => false,
      }
    } else {
      false
    }
  }
}

type Asn1Type = String;

trait Asn1Data {
  fn get_asn1_type() -> Asn1Type;

  /// Create ASN.1 data from this struct.
  // FIXME: Should this use &self?
  fn into_asn1(&self) -> Result<Asn1Data, Asn1Error>;

  /// Create this struct from ASN.1 data.
  fn from_asn1(slice: Asn1Slice) -> Result<Self, Asn1Error>;
}

/// A macro to generate a generic Asn1Data trait implementation for a struct.
macro_rules! asn1_impl {
  ( $impl_type:ident, $asn1_type:expr,
    $( $name:ident, $rusttype:ident, $asn1type:expr ),*
  ) =>
(

impl Asn1Data for $impl_type {
  fn get_asn1_type() -> Asn1Type {
    $asn1_type
  }

  fn into_asn1(&self) -> Result<Asn1Data, Asn1Error> {
    Err(Asn1Error::EncodingError)
  }

  fn from_asn1(slice: Asn1Slice) -> Result<Self, Asn1Error> {
    Err(Asn1Error::InvalidAsn1)
  }
}

)
}

/// A list of errors that can occur decoding or encoding Asn1 data.
enum Asn1Error {
  /// Invalid Asn1 data.
  InvalidAsn1
  /// An error occured while encoding Asn1 data.
  EncodingError
  /// An invalid tag was decoded
  InvalidTag(Asn1Tag)
}
<|endoftext|>


<commit_before>//!
//! Provides `udev` related functionality for automated device scanning.
//!
//! This module mainly provides the [`UdevBackend`](::backend::udev::UdevBackend), which
//! monitors available DRM devices and acts as an event source, generating events whenever these
//! devices change.
//!
//! *Note:* Once inserted into the event loop, the [`UdevBackend`](::backend::udev::UdevBackend) will
//! only notify you about *changes* in the device list. To get an initial snapshot of the state during
//! your initialization, you need to call its `device_list` method.
//!
//! ```no_run
//! use smithay::backend::udev::{UdevBackend, UdevEvent};
//!
//! let udev = UdevBackend::new("seat0", None).expect("Failed to monitor udev.");
//!
//! for (dev_id, node_path) in udev.device_list() {
//!     // process the initial list of devices
//! }
//!
//! # let event_loop = smithay::reexports::calloop::EventLoop::<()>::new().unwrap();
//! # let loop_handle = event_loop.handle();
//! // setup the event source for long-term monitoring
//! loop_handle.insert_source(udev, |event, _, _dispatch_data| match event {
//!     UdevEvent::Added { device_id, path } => {
//!         // a new device has been added
//!     },
//!     UdevEvent::Changed { device_id } => {
//!         // a device has been changed
//!     },
//!     UdevEvent::Removed { device_id } => {
//!         // a device has been removed
//!     }
//! }).expect("Failed to insert the udev source into the event loop");
//! ```
//!
//! Additionally this contains some utility functions related to scanning.
//!
//! See also `anvil/src/udev.rs` for pure hardware backed example of a compositor utilizing this
//! backend.

use nix::sys::stat::{dev_t, stat};
use std::{
    collections::HashMap,
    ffi::OsString,
    io::Result as IoResult,
    os::unix::io::{AsRawFd, RawFd},
    path::{Path, PathBuf},
};
use udev::{Enumerator, EventType, MonitorBuilder, MonitorSocket};

use calloop::{EventSource, Interest, Mode, Poll, Readiness, Token};

/// Backend to monitor available drm devices.
///
/// Provides a way to automatically scan for available gpus and notifies the
/// given handler of any changes. Can be used to provide hot-plug functionality for gpus and
/// attached monitors.
pub struct UdevBackend {
    devices: HashMap<dev_t, PathBuf>,
    monitor: MonitorSocket,
    logger: ::slog::Logger,
}

impl AsRawFd for UdevBackend {
    fn as_raw_fd(&self) -> RawFd {
        self.monitor.as_raw_fd()
    }
}

impl UdevBackend {
    /// Creates a new [`UdevBackend`]
    ///
    /// ## Arguments
    /// `seat`    - system seat which should be bound
    /// `logger`  - slog Logger to be used by the backend and its `DrmDevices`.
    pub fn new<L, S: AsRef<str>>(seat: S, logger: L) -> IoResult<UdevBackend>
    where
        L: Into<Option<::slog::Logger>>,
    {
        let log = crate::slog_or_stdlog(logger).new(o!("smithay_module" => "backend_udev"));

        let devices = all_gpus(seat)?
            .into_iter()
            // Create devices
            .flat_map(|path| match stat(&path) {
                Ok(stat) => Some((stat.st_rdev, path)),
                Err(err) => {
                    warn!(log, "Unable to get id of {:?}, Error: {:?}. Skipping", path, err);
                    None
                }
            })
            .collect();

        let monitor = MonitorBuilder::new()?.match_subsystem("drm")?.listen()?;

        Ok(UdevBackend {
            devices,
            monitor,
            logger: log,
        })
    }

    /// Get a list of DRM devices currently known to the backend
    ///
    /// You should call this once before inserting the event source into your
    /// event loop, to get an initial snapshot of the device state.
    pub fn device_list(&self) -> impl Iterator<Item = (dev_t, &Path)> {
        self.devices.iter().map(|(&id, path)| (id, path.as_ref()))
    }
}

impl EventSource for UdevBackend {
    type Event = UdevEvent;
    type Metadata = ();
    type Ret = ();

    fn process_events<F>(&mut self, _: Readiness, _: Token, mut callback: F) -> std::io::Result<()>
    where
        F: FnMut(UdevEvent, &mut ()),
    {
        let monitor = self.monitor.clone();
        for event in monitor {
            debug!(
                self.logger,
                "Udev event: type={}, devnum={:?} devnode={:?}",
                event.event_type(),
                event.devnum(),
                event.devnode()
            );
            match event.event_type() {
                // New device
                EventType::Add => {
                    if let (Some(path), Some(devnum)) = (event.devnode(), event.devnum()) {
                        info!(self.logger, "New device: #{} at {}", devnum, path.display());
                        if self.devices.insert(devnum, path.to_path_buf()).is_none() {
                            callback(
                                UdevEvent::Added {
                                    device_id: devnum,
                                    path: path.to_path_buf(),
                                },
                                &mut (),
                            );
                        }
                    }
                }
                // Device removed
                EventType::Remove => {
                    if let Some(devnum) = event.devnum() {
                        info!(self.logger, "Device removed: #{}", devnum);
                        if self.devices.remove(&devnum).is_some() {
                            callback(UdevEvent::Removed { device_id: devnum }, &mut ());
                        }
                    }
                }
                // New connector
                EventType::Change => {
                    if let Some(devnum) = event.devnum() {
                        info!(self.logger, "Device changed: #{}", devnum);
                        if self.devices.contains_key(&devnum) {
                            callback(UdevEvent::Changed { device_id: devnum }, &mut ());
                        }
                    }
                }
                _ => {}
            }
        }
        Ok(())
    }

    fn register(&mut self, poll: &mut Poll, token: Token) -> std::io::Result<()> {
        poll.register(self.as_raw_fd(), Interest::Readable, Mode::Level, token)
    }

    fn reregister(&mut self, poll: &mut Poll, token: Token) -> std::io::Result<()> {
        poll.reregister(self.as_raw_fd(), Interest::Readable, Mode::Level, token)
    }

    fn unregister(&mut self, poll: &mut Poll) -> std::io::Result<()> {
        poll.unregister(self.as_raw_fd())
    }
}

/// Events generated by the [`UdevBackend`], notifying you of changes in system devices
pub enum UdevEvent {
    /// A new device has been detected
    Added {
        /// ID of the new device
        device_id: dev_t,
        /// Path of the new device
        path: PathBuf,
    },
    /// A device has changed
    Changed {
        /// ID of the changed device
        device_id: dev_t,
    },
    /// A device has been removed
    Removed {
        /// ID of the removed device
        device_id: dev_t,
    },
}

/// Returns the path of the primary GPU device if any
///
/// Might be used for filtering in [`UdevHandler::device_added`] or for manual
/// [`LegacyDrmDevice`](::backend::drm::legacy::LegacyDrmDevice) initialization.
pub fn primary_gpu<S: AsRef<str>>(seat: S) -> IoResult<Option<PathBuf>> {
    let mut enumerator = Enumerator::new()?;
    enumerator.match_subsystem("drm")?;
    enumerator.match_sysname("card[0-9]*")?;

    if let Some(path) = enumerator
        .scan_devices()?
        .filter(|device| {
            let seat_name = device
                .property_value("ID_SEAT")
                .map(|x| x.to_os_string())
                .unwrap_or_else(|| OsString::from("seat0"));
            if seat_name == *seat.as_ref() {
                if let Ok(Some(pci)) = device.parent_with_subsystem(Path::new("pci")) {
                    if let Some(id) = pci.attribute_value("boot_vga") {
                        return id == "1";
                    }
                }
            }
            false
        })
        .flat_map(|device| device.devnode().map(PathBuf::from))
        .next()
    {
        Ok(Some(path))
    } else {
        all_gpus(seat).map(|all| all.into_iter().next())
    }
}

/// Returns the paths of all available GPU devices
///
/// Might be used for manual  [`LegacyDrmDevice`](::backend::drm::legacy::LegacyDrmDevice)
/// initialization.
pub fn all_gpus<S: AsRef<str>>(seat: S) -> IoResult<Vec<PathBuf>> {
    let mut enumerator = Enumerator::new()?;
    enumerator.match_subsystem("drm")?;
    enumerator.match_sysname("card[0-9]*")?;
    Ok(enumerator
        .scan_devices()?
        .filter(|device| {
            device
                .property_value("ID_SEAT")
                .map(|x| x.to_os_string())
                .unwrap_or_else(|| OsString::from("seat0"))
                == *seat.as_ref()
        })
        .flat_map(|device| device.devnode().map(PathBuf::from))
        .collect())
}
<commit_msg>udev: Add function to query the loaded driver of a device<commit_after>//!
//! Provides `udev` related functionality for automated device scanning.
//!
//! This module mainly provides the [`UdevBackend`](::backend::udev::UdevBackend), which
//! monitors available DRM devices and acts as an event source, generating events whenever these
//! devices change.
//!
//! *Note:* Once inserted into the event loop, the [`UdevBackend`](::backend::udev::UdevBackend) will
//! only notify you about *changes* in the device list. To get an initial snapshot of the state during
//! your initialization, you need to call its `device_list` method.
//!
//! ```no_run
//! use smithay::backend::udev::{UdevBackend, UdevEvent};
//!
//! let udev = UdevBackend::new("seat0", None).expect("Failed to monitor udev.");
//!
//! for (dev_id, node_path) in udev.device_list() {
//!     // process the initial list of devices
//! }
//!
//! # let event_loop = smithay::reexports::calloop::EventLoop::<()>::new().unwrap();
//! # let loop_handle = event_loop.handle();
//! // setup the event source for long-term monitoring
//! loop_handle.insert_source(udev, |event, _, _dispatch_data| match event {
//!     UdevEvent::Added { device_id, path } => {
//!         // a new device has been added
//!     },
//!     UdevEvent::Changed { device_id } => {
//!         // a device has been changed
//!     },
//!     UdevEvent::Removed { device_id } => {
//!         // a device has been removed
//!     }
//! }).expect("Failed to insert the udev source into the event loop");
//! ```
//!
//! Additionally this contains some utility functions related to scanning.
//!
//! See also `anvil/src/udev.rs` for pure hardware backed example of a compositor utilizing this
//! backend.

use nix::sys::stat::{dev_t, stat};
use std::{
    collections::HashMap,
    ffi::OsString,
    io::Result as IoResult,
    os::unix::io::{AsRawFd, RawFd},
    path::{Path, PathBuf},
};
use udev::{Enumerator, EventType, MonitorBuilder, MonitorSocket};

use calloop::{EventSource, Interest, Mode, Poll, Readiness, Token};

/// Backend to monitor available drm devices.
///
/// Provides a way to automatically scan for available gpus and notifies the
/// given handler of any changes. Can be used to provide hot-plug functionality for gpus and
/// attached monitors.
pub struct UdevBackend {
    devices: HashMap<dev_t, PathBuf>,
    monitor: MonitorSocket,
    logger: ::slog::Logger,
}

impl AsRawFd for UdevBackend {
    fn as_raw_fd(&self) -> RawFd {
        self.monitor.as_raw_fd()
    }
}

impl UdevBackend {
    /// Creates a new [`UdevBackend`]
    ///
    /// ## Arguments
    /// `seat`    - system seat which should be bound
    /// `logger`  - slog Logger to be used by the backend and its `DrmDevices`.
    pub fn new<L, S: AsRef<str>>(seat: S, logger: L) -> IoResult<UdevBackend>
    where
        L: Into<Option<::slog::Logger>>,
    {
        let log = crate::slog_or_stdlog(logger).new(o!("smithay_module" => "backend_udev"));

        let devices = all_gpus(seat)?
            .into_iter()
            // Create devices
            .flat_map(|path| match stat(&path) {
                Ok(stat) => Some((stat.st_rdev, path)),
                Err(err) => {
                    warn!(log, "Unable to get id of {:?}, Error: {:?}. Skipping", path, err);
                    None
                }
            })
            .collect();

        let monitor = MonitorBuilder::new()?.match_subsystem("drm")?.listen()?;

        Ok(UdevBackend {
            devices,
            monitor,
            logger: log,
        })
    }

    /// Get a list of DRM devices currently known to the backend
    ///
    /// You should call this once before inserting the event source into your
    /// event loop, to get an initial snapshot of the device state.
    pub fn device_list(&self) -> impl Iterator<Item = (dev_t, &Path)> {
        self.devices.iter().map(|(&id, path)| (id, path.as_ref()))
    }
}

impl EventSource for UdevBackend {
    type Event = UdevEvent;
    type Metadata = ();
    type Ret = ();

    fn process_events<F>(&mut self, _: Readiness, _: Token, mut callback: F) -> std::io::Result<()>
    where
        F: FnMut(UdevEvent, &mut ()),
    {
        let monitor = self.monitor.clone();
        for event in monitor {
            debug!(
                self.logger,
                "Udev event: type={}, devnum={:?} devnode={:?}",
                event.event_type(),
                event.devnum(),
                event.devnode()
            );
            match event.event_type() {
                // New device
                EventType::Add => {
                    if let (Some(path), Some(devnum)) = (event.devnode(), event.devnum()) {
                        info!(self.logger, "New device: #{} at {}", devnum, path.display());
                        if self.devices.insert(devnum, path.to_path_buf()).is_none() {
                            callback(
                                UdevEvent::Added {
                                    device_id: devnum,
                                    path: path.to_path_buf(),
                                },
                                &mut (),
                            );
                        }
                    }
                }
                // Device removed
                EventType::Remove => {
                    if let Some(devnum) = event.devnum() {
                        info!(self.logger, "Device removed: #{}", devnum);
                        if self.devices.remove(&devnum).is_some() {
                            callback(UdevEvent::Removed { device_id: devnum }, &mut ());
                        }
                    }
                }
                // New connector
                EventType::Change => {
                    if let Some(devnum) = event.devnum() {
                        info!(self.logger, "Device changed: #{}", devnum);
                        if self.devices.contains_key(&devnum) {
                            callback(UdevEvent::Changed { device_id: devnum }, &mut ());
                        }
                    }
                }
                _ => {}
            }
        }
        Ok(())
    }

    fn register(&mut self, poll: &mut Poll, token: Token) -> std::io::Result<()> {
        poll.register(self.as_raw_fd(), Interest::Readable, Mode::Level, token)
    }

    fn reregister(&mut self, poll: &mut Poll, token: Token) -> std::io::Result<()> {
        poll.reregister(self.as_raw_fd(), Interest::Readable, Mode::Level, token)
    }

    fn unregister(&mut self, poll: &mut Poll) -> std::io::Result<()> {
        poll.unregister(self.as_raw_fd())
    }
}

/// Events generated by the [`UdevBackend`], notifying you of changes in system devices
pub enum UdevEvent {
    /// A new device has been detected
    Added {
        /// ID of the new device
        device_id: dev_t,
        /// Path of the new device
        path: PathBuf,
    },
    /// A device has changed
    Changed {
        /// ID of the changed device
        device_id: dev_t,
    },
    /// A device has been removed
    Removed {
        /// ID of the removed device
        device_id: dev_t,
    },
}

/// Returns the path of the primary GPU device if any
///
/// Might be used for filtering in [`UdevHandler::device_added`] or for manual
/// [`LegacyDrmDevice`](::backend::drm::legacy::LegacyDrmDevice) initialization.
pub fn primary_gpu<S: AsRef<str>>(seat: S) -> IoResult<Option<PathBuf>> {
    let mut enumerator = Enumerator::new()?;
    enumerator.match_subsystem("drm")?;
    enumerator.match_sysname("card[0-9]*")?;

    if let Some(path) = enumerator
        .scan_devices()?
        .filter(|device| {
            let seat_name = device
                .property_value("ID_SEAT")
                .map(|x| x.to_os_string())
                .unwrap_or_else(|| OsString::from("seat0"));
            if seat_name == *seat.as_ref() {
                if let Ok(Some(pci)) = device.parent_with_subsystem(Path::new("pci")) {
                    if let Some(id) = pci.attribute_value("boot_vga") {
                        return id == "1";
                    }
                }
            }
            false
        })
        .flat_map(|device| device.devnode().map(PathBuf::from))
        .next()
    {
        Ok(Some(path))
    } else {
        all_gpus(seat).map(|all| all.into_iter().next())
    }
}

/// Returns the paths of all available GPU devices
///
/// Might be used for manual  [`LegacyDrmDevice`](::backend::drm::legacy::LegacyDrmDevice)
/// initialization.
pub fn all_gpus<S: AsRef<str>>(seat: S) -> IoResult<Vec<PathBuf>> {
    let mut enumerator = Enumerator::new()?;
    enumerator.match_subsystem("drm")?;
    enumerator.match_sysname("card[0-9]*")?;
    Ok(enumerator
        .scan_devices()?
        .filter(|device| {
            device
                .property_value("ID_SEAT")
                .map(|x| x.to_os_string())
                .unwrap_or_else(|| OsString::from("seat0"))
                == *seat.as_ref()
        })
        .flat_map(|device| device.devnode().map(PathBuf::from))
        .collect())
}

/// Returns the loaded driver for a device named by it's [`dev_t`](::nix::sys::stat::dev_t).
pub fn driver(dev: dev_t) -> IoResult<Option<OsString>> {
    let mut enumerator = Enumerator::new()?;
    enumerator.match_subsystem("drm")?;
    enumerator.match_sysname("card[0-9]*")?;
    Ok(enumerator
        .scan_devices()?
        .filter(|device| device.devnum() == Some(dev))
        .flat_map(|dev| {
            let mut device = Some(dev);
            while let Some(dev) = device {
                if dev.driver().is_some() {
                    return dev.driver().map(std::ffi::OsStr::to_os_string);
                }
                device = dev.parent();
            }
            None
        })
        .next())
}
<|endoftext|>