<commit_before>use std::{
    ffi::OsString,
    path::{Path, PathBuf},
};

use archive_utils::extract;
use http_utils::{download, download_temp, url};
use provider::{
    async_trait::async_trait,
    eyre::{bail, Result},
    once_cell::sync::Lazy,
    regex::Regex,
    run_schedule,
    stencila_schema::{Node, Thing},
    tokio::{self, sync::mpsc},
    tracing, ImportOptions, ParseItem, Provider, ProviderTrait, IMPORT,
};

pub struct HttpProvider;

#[async_trait]
impl ProviderTrait for HttpProvider {
    fn spec() -> Provider {
        Provider::new("http")
    }

    fn parse(string: &str) -> Vec<ParseItem> {
        static URL_REGEX: Lazy<Regex> =
            Lazy::new(|| Regex::new(r"https?://[^\s]+").expect("Unable to create regex"));

        URL_REGEX
            .captures_iter(string)
            .into_iter()
            .map(|captures| {
                let capture = captures.get(0).expect("Should be a match group");
                ParseItem {
                    begin: capture.start(),
                    end: capture.end(),
                    node: Node::Thing(Thing {
                        url: Some(Box::new(captures[0].to_string())),
                        ..Default::default()
                    }),
                }
            })
            .collect()
    }

    async fn import(node: &Node, dest: &Path, _options: Option<ImportOptions>) -> Result<bool> {
        let thing = match node {
            Node::Thing(thing) => thing,
            _ => return Ok(false),
        };
        let url = match &thing.url {
            Some(url) => url.as_ref(),
            None => return Ok(false),
        };

        let dest_ext = dest
            .extension()
            .map_or_else(OsString::new, |ext| ext.to_owned());
        let dest_ext = dest_ext.as_os_str();

        let src_url = url::Url::parse(url)?;
        let src_path = PathBuf::from(src_url.path());
        let src_filename = src_path.file_name();
        let src_ext = if src_path.ends_with("tar.gz") {
            "tar.gz".to_string()
        } else if src_path.ends_with("tar.xz") {
            "tar.xz".to_string()
        } else {
            src_path
                .extension()
                .map_or_else(String::new, |ext| ext.to_string_lossy().to_string())
        };

        tracing::info!("Downloading `{url}` to `{dest}`", dest = dest.display());
        if !dest_ext.is_empty() {
            // If the destination appears to be a file: just download it
            download(url, dest).await?
        } else if src_ext == "tar"
            || src_ext == "tgz"
            || src_ext == "tar.gz"
            || src_ext == "tar.xz"
            || src_ext == "zip"
        {
            // If the destination appears to be a folder and the source is an archive:
            // extract it into the folder
            let archive = download_temp(url, Some(&src_ext)).await?;
            extract(archive.path(), dest, 0, None)?;
        } else if let Some(filename) = src_filename {
            // If the destination appears to be a folder and the source has a filename:
            // download it into the folder using the filename
            download(url, &dest.join(filename)).await?;
        } else if dest_ext.is_empty() {
            // If the destination appears to be a folder and the source has no filename:
            // download it to a slugified name
            download(url, &dest.join(slug::slugify(url))).await?;
        } else {
            // Otherwise, just download to the destination (even though it has no extension e.g. `Makefile`)
            download(url, dest).await?;
        }

        Ok(true)
    }

    async fn cron(action: &str, schedule: &str, node: &Node, path: &Path) -> Result<bool> {
        let thing = match node {
            Node::Thing(thing) => thing,
            _ => return Ok(false),
        };
        let url = match &thing.url {
            Some(url) => url.as_ref(),
            None => return Ok(false),
        }
        .to_owned();

        if action != IMPORT {
            bail!("Only the import action is supported for `http` resources")
        }

        let (sender, mut receiver) = mpsc::channel(1);
        let node = node.to_owned();
        let path = path.to_owned();
        tokio::spawn(async move {
            while let Some(..) = receiver.recv().await {
                if let Err(error) = Self::import(&node, &path, None).await {
                    tracing::error!("While importing from `{}`: {}", url, error);
                }
            }
        });
        run_schedule(schedule, sender).await?;

        Ok(true)
    }
}

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

    #[test]
    fn parse() {
        assert_json_is!(
            HttpProvider::parse("http://example.com/file.csv")[0].node,
            {
                "type": "Thing",
                "url": "http://example.com/file.csv",
            }
        );

        assert_json_is!(
            HttpProvider::parse("https://example.com/anything")[0].node,
            {
                "type": "Thing",
                "url": "https://example.com/anything",
            }
        );

        let parse_items = HttpProvider::parse(
            "
            http://example.com/sub/file.csv som word to be ignored
            and then another url https://example.com/anything/archive.tar.gz
        ",
        );
        assert_eq!(parse_items.len(), 2);
        assert_json_is!(parse_items[0].node, {
            "type": "Thing",
            "url": "http://example.com/sub/file.csv",
        });
        assert_json_is!(parse_items[1].node, {
            "type": "Thing",
            "url": "https://example.com/anything/archive.tar.gz",
        });
    }
}
<commit_msg>fix(HTTP provider): Action defaults to "import" if empty<commit_after>use std::{
    ffi::OsString,
    path::{Path, PathBuf},
};

use archive_utils::extract;
use http_utils::{download, download_temp, url};
use provider::{
    async_trait::async_trait,
    eyre::{bail, Result},
    once_cell::sync::Lazy,
    regex::Regex,
    run_schedule,
    stencila_schema::{Node, Thing},
    tokio::{self, sync::mpsc},
    tracing, ImportOptions, ParseItem, Provider, ProviderTrait, IMPORT,
};

pub struct HttpProvider;

#[async_trait]
impl ProviderTrait for HttpProvider {
    fn spec() -> Provider {
        Provider::new("http")
    }

    fn parse(string: &str) -> Vec<ParseItem> {
        static URL_REGEX: Lazy<Regex> =
            Lazy::new(|| Regex::new(r"https?://[^\s]+").expect("Unable to create regex"));

        URL_REGEX
            .captures_iter(string)
            .into_iter()
            .map(|captures| {
                let capture = captures.get(0).expect("Should be a match group");
                ParseItem {
                    begin: capture.start(),
                    end: capture.end(),
                    node: Node::Thing(Thing {
                        url: Some(Box::new(captures[0].to_string())),
                        ..Default::default()
                    }),
                }
            })
            .collect()
    }

    async fn import(node: &Node, dest: &Path, _options: Option<ImportOptions>) -> Result<bool> {
        let thing = match node {
            Node::Thing(thing) => thing,
            _ => return Ok(false),
        };
        let url = match &thing.url {
            Some(url) => url.as_ref(),
            None => return Ok(false),
        };

        let dest_ext = dest
            .extension()
            .map_or_else(OsString::new, |ext| ext.to_owned());
        let dest_ext = dest_ext.as_os_str();

        let src_url = url::Url::parse(url)?;
        let src_path = PathBuf::from(src_url.path());
        let src_filename = src_path.file_name();
        let src_ext = if src_path.ends_with("tar.gz") {
            "tar.gz".to_string()
        } else if src_path.ends_with("tar.xz") {
            "tar.xz".to_string()
        } else {
            src_path
                .extension()
                .map_or_else(String::new, |ext| ext.to_string_lossy().to_string())
        };

        tracing::info!("Downloading `{url}` to `{dest}`", dest = dest.display());
        if !dest_ext.is_empty() {
            // If the destination appears to be a file: just download it
            download(url, dest).await?
        } else if src_ext == "tar"
            || src_ext == "tgz"
            || src_ext == "tar.gz"
            || src_ext == "tar.xz"
            || src_ext == "zip"
        {
            // If the destination appears to be a folder and the source is an archive:
            // extract it into the folder
            let archive = download_temp(url, Some(&src_ext)).await?;
            extract(archive.path(), dest, 0, None)?;
        } else if let Some(filename) = src_filename {
            // If the destination appears to be a folder and the source has a filename:
            // download it into the folder using the filename
            download(url, &dest.join(filename)).await?;
        } else if dest_ext.is_empty() {
            // If the destination appears to be a folder and the source has no filename:
            // download it to a slugified name
            download(url, &dest.join(slug::slugify(url))).await?;
        } else {
            // Otherwise, just download to the destination (even though it has no extension e.g. `Makefile`)
            download(url, dest).await?;
        }

        Ok(true)
    }

    async fn cron(action: &str, schedule: &str, node: &Node, path: &Path) -> Result<bool> {
        let thing = match node {
            Node::Thing(thing) => thing,
            _ => return Ok(false),
        };
        let url = match &thing.url {
            Some(url) => url.as_ref(),
            None => return Ok(false),
        }
        .to_owned();

        if !action.is_empty() && action != IMPORT {
            bail!("Only the import action is supported for `http` resources")
        }

        let (sender, mut receiver) = mpsc::channel(1);
        let node = node.to_owned();
        let path = path.to_owned();
        tokio::spawn(async move {
            while let Some(..) = receiver.recv().await {
                if let Err(error) = Self::import(&node, &path, None).await {
                    tracing::error!("While importing from `{}`: {}", url, error);
                }
            }
        });
        run_schedule(schedule, sender).await?;

        Ok(true)
    }
}

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

    #[test]
    fn parse() {
        assert_json_is!(
            HttpProvider::parse("http://example.com/file.csv")[0].node,
            {
                "type": "Thing",
                "url": "http://example.com/file.csv",
            }
        );

        assert_json_is!(
            HttpProvider::parse("https://example.com/anything")[0].node,
            {
                "type": "Thing",
                "url": "https://example.com/anything",
            }
        );

        let parse_items = HttpProvider::parse(
            "
            http://example.com/sub/file.csv som word to be ignored
            and then another url https://example.com/anything/archive.tar.gz
        ",
        );
        assert_eq!(parse_items.len(), 2);
        assert_json_is!(parse_items[0].node, {
            "type": "Thing",
            "url": "http://example.com/sub/file.csv",
        });
        assert_json_is!(parse_items[1].node, {
            "type": "Thing",
            "url": "https://example.com/anything/archive.tar.gz",
        });
    }
}
<|endoftext|>


<commit_before>fn main() {
    let age = 49u32;
    match age {
        1 => println!("got 1"),
        2 | 3 | 4 | 5 => println!("got something in between 2 and 5"),
        25 ... 50 => println!("got something in range of 25 and 50 (inclusive)"),
        _ => println!("no one matches")
    }
}<commit_msg>#1 Updated 10_match<commit_after>fn main() {
    simple_match();
    tuple_destructuring();
}

fn simple_match() {
    let age = 49u32;
    match age {
        1 => println!("got 1"),
        2 | 3 | 4 | 5 => println!("got something in between 2 and 5"),
        25 ... 50 => println!("got something in range of 25 and 50 (inclusive)"),
        _ => println!("no one matches")
    }
}

fn tuple_destructuring() {
    let my_tuple = (7, 8);
    match my_tuple {
        (7, num) => println!("{:?}", num),
        (num, 8) => println!("{:?}", num),
        _ => println!("no one matches")
    }
}<|endoftext|>


<commit_before>use flowgger::config::Config;
use flowgger::merger::Merger;
use openssl::bn::BigNum;
use openssl::dh::DH;
use openssl::ssl::*;
use openssl::x509::X509FileType;
use std::error::Error;
use std::io;
use std::io::{stderr, BufWriter, ErrorKind, Write};
use std::net::TcpStream;
use std::path::{Path, PathBuf};
use std::sync::mpsc::Receiver;
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;
use super::Output;

const DEFAULT_CERT: &'static str = "flowgger.pem";
const DEFAULT_CIPHERS: &'static str = "ECDHE-ECDSA-AES128-GCM-SHA256:ECDHE-RSA-AES128-GCM-SHA256:ECDHE-ECDSA-CHACHA20-POLY1305:ECDHE-RSA-CHACHA20-POLY1305:ECDHE-ECDSA-AES128-SHA256:ECDHE-RSA-AES128-SHA256:ECDHE-ECDSA-AES128-SHA:ECDHE-RSA-AES128-SHA:ECDHE-ECDSA-AES256-GCM-SHA384:ECDHE-RSA-AES256-GCM-SHA384:ECDHE-ECDSA-AES256-SHA384:ECDHE-RSA-AES256-SHA384:ECDHE-ECDSA-AES256-SHA:ECDHE-RSA-AES256-SHA:AES128-GCM-SHA256:AES256-GCM-SHA384:AES128-SHA256:AES256-SHA256:AES128-SHA:AES256-SHA:ECDHE-ECDSA-DES-CBC3-SHA:ECDHE-RSA-DES-CBC3-SHA:DES-CBC3-SHA:!aNULL:!eNULL:!EXPORT:!DES:!RC4:!MD5:!PSK:!aECDH:!EDH-DSS-DES-CBC3-SHA:!EDH-RSA-DES-CBC3-SHA:!KRB5-DES-CBC3-SHA";
const DEFAULT_COMPRESSION: bool = false;
const DEFAULT_CONNECT: &'static str = "127.0.0.1:6514";
const DEFAULT_KEY: &'static str = "flowgger.pem";
const DEFAULT_SLEEP_AFTER_CONNECTION_FAILURE: u32 = 1000;
const DEFAULT_TIMEOUT: u64 = 3600;
const DEFAULT_TLS_METHOD: &'static str = "any";
const DEFAULT_VERIFY_PEER: bool = false;
const TLS_VERIFY_DEPTH: u32 = 6;
const TLS_DEFAULT_THREADS: u32 = 1;

pub struct TlsOutput {
    config: TlsConfig,
    threads: u32
}

#[derive(Clone)]
struct TlsConfig {
    timeout: Option<Duration>,
    connect: String,
    arc_ctx: Arc<SslContext>
}

struct TlsWorker {
    arx: Arc<Mutex<Receiver<Vec<u8>>>>,
    merger: Option<Box<Merger + Send>>,
    tls_config: TlsConfig
}

impl TlsWorker {
    fn new(arx: Arc<Mutex<Receiver<Vec<u8>>>>, merger: Option<Box<Merger + Send>>, tls_config: TlsConfig) -> TlsWorker {
        TlsWorker {
            arx: arx,
            merger: merger,
            tls_config: tls_config
        }
    }

    fn handle_connection(&self, connect: &str) -> io::Result<()> {
        let client = try!(new_tcp(connect));
        let _ = writeln!(stderr(), "Connected to {}", connect);
        let sslclient = match SslStream::connect(&*self.tls_config.arc_ctx, client) {
            Err(_) => return Err(io::Error::new(io::ErrorKind::ConnectionAborted,
                "SSL handshake aborted by the server")),
            Ok(sslclient) => sslclient
        };
        let _ = writeln!(stderr(), "Completed SSL handshake with {}", connect);
        let mut writer = BufWriter::new(sslclient);
        let merger = &self.merger;
        loop {
            let mut bytes = match { self.arx.lock().unwrap().recv() } {
                Ok(line) => line,
                Err(_) => return Err(io::Error::new(io::ErrorKind::Other, "Cannot read the message queue any more"))
            };
            if let Some(ref merger) = *merger {
                merger.frame(&mut bytes);
            }
            match writer.write_all(&bytes) {
                Ok(_) => { },
                Err(e) => match e.kind() {
                    ErrorKind::Interrupted => continue,
                    _ => return Err(e)
                }
            };
        }
    }

    fn run(self) {
        loop {
            let connect = &self.tls_config.connect;
            if let Err(e) = self.handle_connection(connect) {
                match e.kind() {
                    ErrorKind::ConnectionRefused => {
                        let _ = writeln!(stderr(), "Connection to {} refused", connect);
                    }
                    ErrorKind::ConnectionAborted | ErrorKind::ConnectionReset => {
                        let _ = writeln!(stderr(), "Connection to {} aborted by the server", connect);
                    }
                    _ => {
                        let _ = writeln!(stderr(), "Error while communicating with {} - {}", connect, e);
                    }
                }
            }
            thread::sleep_ms(DEFAULT_SLEEP_AFTER_CONNECTION_FAILURE);
            let _ = writeln!(stderr(), "Attempting to reconnect");
        }
    }
}

fn new_tcp(connect: &str) -> Result<TcpStream, io::Error> {
    loop {
        match TcpStream::connect(connect) {
            Ok(stream) => return Ok(stream),
            Err(e) => {
                return Err(e);
            }
        }
    }
}

impl TlsOutput {
    pub fn new(config: &Config) -> TlsOutput {
        let (tls_config, threads) = config_parse(&config);
        TlsOutput {
            config: tls_config,
            threads: threads
        }
    }
}

impl Output for TlsOutput {
    fn start(&self, arx: Arc<Mutex<Receiver<Vec<u8>>>>, merger: Option<Box<Merger>>) {
        for _ in 0..self.threads {
            let arx = arx.clone();
            let config = self.config.clone();
            let merger = match merger {
                Some(ref merger) => Some(merger.clone_boxed()) as Option<Box<Merger + Send>>,
                None => None
            };
            thread::spawn(move || {
                let worker = TlsWorker::new(arx, merger, config);
                worker.run();
            });
        }
    }
}

#[cfg(feature = "ecdh")]
fn set_ecdh(ctx: &mut SslContext) {
    ctx.set_ecdh_auto(true).unwrap();
}

#[cfg(not(feature = "ecdh"))]
fn set_ecdh(ctx: &mut SslContext) {
    let _ = ctx;
}

fn set_fs(ctx: &mut SslContext) {
    let p = BigNum::from_hex_str("87A8E61DB4B6663CFFBBD19C651959998CEEF608660DD0F25D2CEED4435E3B00E00DF8F1D61957D4FAF7DF4561B2AA3016C3D91134096FAA3BF4296D830E9A7C209E0C6497517ABD5A8A9D306BCF67ED91F9E6725B4758C022E0B1EF4275BF7B6C5BFC11D45F9088B941F54EB1E59BB8BC39A0BF12307F5C4FDB70C581B23F76B63ACAE1CAA6B7902D52526735488A0EF13C6D9A51BFA4AB3AD8347796524D8EF6A167B5A41825D967E144E5140564251CCACB83E6B486F6B3CA3F7971506026C0B857F689962856DED4010ABD0BE621C3A3960A54E710C375F26375D7014103A4B54330C198AF126116D2276E11715F693877FAD7EF09CADB094AE91E1A1597").unwrap();
    let g = BigNum::from_hex_str("3FB32C9B73134D0B2E77506660EDBD484CA7B18F21EF205407F4793A1A0BA12510DBC15077BE463FFF4FED4AAC0BB555BE3A6C1B0C6B47B1BC3773BF7E8C6F62901228F8C28CBB18A55AE31341000A650196F931C77A57F2DDF463E5E9EC144B777DE62AAAB8A8628AC376D282D6ED3864E67982428EBC831D14348F6F2F9193B5045AF2767164E1DFC967C1FB3F2E55A4BD1BFFE83B9C80D052B985D182EA0ADB2A3B7313D3FE14C8484B1E052588B9B7D2BBD2DF016199ECD06E1557CD0915B3353BBB64E0EC377FD028370DF92B52C7891428CDC67EB6184B523D1DB246C32F63078490F00EF8D647D148D47954515E2327CFEF98C582664B4C0F6CC41659").unwrap();
    let q = BigNum::from_hex_str("8CF83642A709A097B447997640129DA299B1A47D1EB3750BA308B0FE64F5FBD3").unwrap();
    let dh = DH::from_params(p, g, q).unwrap();
    ctx.set_tmp_dh(dh).unwrap();
    set_ecdh(ctx);
}

fn config_parse(config: &Config) -> (TlsConfig, u32) {
    let threads = config.lookup("output.tls_threads").
        map_or(TLS_DEFAULT_THREADS, |x| x.as_integer().
            expect("output.tls_threads must be a 32-bit integer") as u32);
    let connect = config.lookup("output.connect").map_or(DEFAULT_CONNECT, |x|x.as_str().
        expect("output.connect must be an ip:port string")).to_owned();
    let cert = config.lookup("output.tls_cert").map_or(DEFAULT_CERT, |x| x.as_str().
        expect("output.tls_cert must be a path to a .pem file")).to_owned();
    let key = config.lookup("output.tls_key").map_or(DEFAULT_KEY, |x| x.as_str().
        expect("output.tls_key must be a path to a .pem file")).to_owned();
    let ciphers = config.lookup("output.tls_ciphers").map_or(DEFAULT_CIPHERS, |x| x.as_str().
        expect("output.tls_ciphers must be a string with a cipher suite")).to_owned();
    let tls_method = match config.lookup("output.tls_method").map_or(DEFAULT_TLS_METHOD, |x| x.as_str().
        expect("output.tls_method must be a string with the TLS method")).to_lowercase().as_ref() {
            "any" | "sslv23" => SslMethod::Sslv23,
            "tlsv1" | "tlsv1.0" => SslMethod::Tlsv1,
            "tlsv1.1" => SslMethod::Tlsv1_1,
            "tlsv1.2" => SslMethod::Tlsv1_2,
            _ => panic!(r#"TLS method must be "any", "TLSv1.0", "TLSv1.1" or "TLSv1.2""#)
    };
    let verify_peer = config.lookup("output.tls_verify_peer").map_or(DEFAULT_VERIFY_PEER, |x| x.as_bool().
        expect("output.tls_verify_peer must be a boolean"));
    let ca_file: Option<PathBuf> = config.lookup("output.tls_ca_file").map_or(None, |x|
        Some(PathBuf::from(x.as_str().expect("output.tls_ca_file must be a path to a file"))));
    let compression = config.lookup("output.tls_compression").map_or(DEFAULT_COMPRESSION, |x| x.as_bool().
        expect("output.tls_compression must be a boolean"));
    let timeout = config.lookup("output.timeout").map_or(DEFAULT_TIMEOUT, |x| x.as_integer().
        expect("output.timeout must be an integer") as u64);
    let mut ctx = SslContext::new(tls_method).unwrap();
    if verify_peer == false {
        ctx.set_verify(SSL_VERIFY_NONE, None);
    } else {
        ctx.set_verify_depth(TLS_VERIFY_DEPTH);
        ctx.set_verify(SSL_VERIFY_PEER | SSL_VERIFY_FAIL_IF_NO_PEER_CERT, None);
        if let Some(ca_file) = ca_file {
            if ctx.set_CA_file(&ca_file).is_err() {
                panic!("Unable to read the trusted CA file");
            }
        }
    }
    let mut opts = SSL_OP_CIPHER_SERVER_PREFERENCE | SSL_OP_NO_SESSION_RESUMPTION_ON_RENEGOTIATION;
    if compression == false {
        opts = opts | SSL_OP_NO_COMPRESSION;
    }
    ctx.set_options(opts);
    set_fs(&mut ctx);
    ctx.set_certificate_file(&Path::new(&cert), X509FileType::PEM).unwrap();
    ctx.set_private_key_file(&Path::new(&key), X509FileType::PEM).unwrap();
    ctx.set_cipher_list(&ciphers).unwrap();
    let arc_ctx = Arc::new(ctx);
    let tls_config = TlsConfig {
        connect: connect,
        timeout: Some(Duration::from_secs(timeout)),
        arc_ctx: arc_ctx
    };
    (tls_config, threads)
}
<commit_msg>Add a new `tls_async` output option to flush the writer when the buffer is full, or after every message<commit_after>use flowgger::config::Config;
use flowgger::merger::Merger;
use openssl::bn::BigNum;
use openssl::dh::DH;
use openssl::ssl::*;
use openssl::x509::X509FileType;
use std::error::Error;
use std::io;
use std::io::{stderr, BufWriter, ErrorKind, Write};
use std::net::TcpStream;
use std::path::{Path, PathBuf};
use std::sync::mpsc::Receiver;
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;
use super::Output;

const DEFAULT_CERT: &'static str = "flowgger.pem";
const DEFAULT_CIPHERS: &'static str = "ECDHE-ECDSA-AES128-GCM-SHA256:ECDHE-RSA-AES128-GCM-SHA256:ECDHE-ECDSA-CHACHA20-POLY1305:ECDHE-RSA-CHACHA20-POLY1305:ECDHE-ECDSA-AES128-SHA256:ECDHE-RSA-AES128-SHA256:ECDHE-ECDSA-AES128-SHA:ECDHE-RSA-AES128-SHA:ECDHE-ECDSA-AES256-GCM-SHA384:ECDHE-RSA-AES256-GCM-SHA384:ECDHE-ECDSA-AES256-SHA384:ECDHE-RSA-AES256-SHA384:ECDHE-ECDSA-AES256-SHA:ECDHE-RSA-AES256-SHA:AES128-GCM-SHA256:AES256-GCM-SHA384:AES128-SHA256:AES256-SHA256:AES128-SHA:AES256-SHA:ECDHE-ECDSA-DES-CBC3-SHA:ECDHE-RSA-DES-CBC3-SHA:DES-CBC3-SHA:!aNULL:!eNULL:!EXPORT:!DES:!RC4:!MD5:!PSK:!aECDH:!EDH-DSS-DES-CBC3-SHA:!EDH-RSA-DES-CBC3-SHA:!KRB5-DES-CBC3-SHA";
const DEFAULT_COMPRESSION: bool = false;
const DEFAULT_CONNECT: &'static str = "127.0.0.1:6514";
const DEFAULT_KEY: &'static str = "flowgger.pem";
const DEFAULT_SLEEP_AFTER_CONNECTION_FAILURE: u32 = 1000;
const DEFAULT_ASYNC: bool = false;
const DEFAULT_TIMEOUT: u64 = 3600;
const DEFAULT_TLS_METHOD: &'static str = "any";
const DEFAULT_VERIFY_PEER: bool = false;
const TLS_VERIFY_DEPTH: u32 = 6;
const TLS_DEFAULT_THREADS: u32 = 1;

pub struct TlsOutput {
    config: TlsConfig,
    threads: u32
}

#[derive(Clone)]
struct TlsConfig {
    timeout: Option<Duration>,
    connect: String,
    arc_ctx: Arc<SslContext>,
    async: bool
}

struct TlsWorker {
    arx: Arc<Mutex<Receiver<Vec<u8>>>>,
    merger: Option<Box<Merger + Send>>,
    tls_config: TlsConfig
}

impl TlsWorker {
    fn new(arx: Arc<Mutex<Receiver<Vec<u8>>>>, merger: Option<Box<Merger + Send>>, tls_config: TlsConfig) -> TlsWorker {
        TlsWorker {
            arx: arx,
            merger: merger,
            tls_config: tls_config
        }
    }

    fn handle_connection(&self, connect: &str) -> io::Result<()> {
        let client = try!(new_tcp(connect));
        let _ = writeln!(stderr(), "Connected to {}", connect);
        let sslclient = match SslStream::connect(&*self.tls_config.arc_ctx, client) {
            Err(_) => return Err(io::Error::new(io::ErrorKind::ConnectionAborted,
                "SSL handshake aborted by the server")),
            Ok(sslclient) => sslclient
        };
        let _ = writeln!(stderr(), "Completed SSL handshake with {}", connect);
        let mut writer = BufWriter::new(sslclient);
        let merger = &self.merger;
        loop {
            let mut bytes = match { self.arx.lock().unwrap().recv() } {
                Ok(line) => line,
                Err(_) => return Err(io::Error::new(io::ErrorKind::Other, "Cannot read the message queue any more"))
            };
            if let Some(ref merger) = *merger {
                merger.frame(&mut bytes);
            }
            match writer.write_all(&bytes) {
                Ok(_) => { },
                Err(e) => match e.kind() {
                    ErrorKind::Interrupted => continue,
                    _ => return Err(e)
                }
            };
            if self.tls_config.async == false {
                try!(writer.flush());
            }
        }
    }

    fn run(self) {
        loop {
            let connect = &self.tls_config.connect;
            if let Err(e) = self.handle_connection(connect) {
                match e.kind() {
                    ErrorKind::ConnectionRefused => {
                        let _ = writeln!(stderr(), "Connection to {} refused", connect);
                    }
                    ErrorKind::ConnectionAborted | ErrorKind::ConnectionReset => {
                        let _ = writeln!(stderr(), "Connection to {} aborted by the server", connect);
                    }
                    _ => {
                        let _ = writeln!(stderr(), "Error while communicating with {} - {}", connect, e);
                    }
                }
            }
            thread::sleep_ms(DEFAULT_SLEEP_AFTER_CONNECTION_FAILURE);
            let _ = writeln!(stderr(), "Attempting to reconnect");
        }
    }
}

fn new_tcp(connect: &str) -> Result<TcpStream, io::Error> {
    loop {
        match TcpStream::connect(connect) {
            Ok(stream) => return Ok(stream),
            Err(e) => {
                return Err(e);
            }
        }
    }
}

impl TlsOutput {
    pub fn new(config: &Config) -> TlsOutput {
        let (tls_config, threads) = config_parse(&config);
        TlsOutput {
            config: tls_config,
            threads: threads
        }
    }
}

impl Output for TlsOutput {
    fn start(&self, arx: Arc<Mutex<Receiver<Vec<u8>>>>, merger: Option<Box<Merger>>) {
        for _ in 0..self.threads {
            let arx = arx.clone();
            let config = self.config.clone();
            let merger = match merger {
                Some(ref merger) => Some(merger.clone_boxed()) as Option<Box<Merger + Send>>,
                None => None
            };
            thread::spawn(move || {
                let worker = TlsWorker::new(arx, merger, config);
                worker.run();
            });
        }
    }
}

#[cfg(feature = "ecdh")]
fn set_ecdh(ctx: &mut SslContext) {
    ctx.set_ecdh_auto(true).unwrap();
}

#[cfg(not(feature = "ecdh"))]
fn set_ecdh(ctx: &mut SslContext) {
    let _ = ctx;
}

fn set_fs(ctx: &mut SslContext) {
    let p = BigNum::from_hex_str("87A8E61DB4B6663CFFBBD19C651959998CEEF608660DD0F25D2CEED4435E3B00E00DF8F1D61957D4FAF7DF4561B2AA3016C3D91134096FAA3BF4296D830E9A7C209E0C6497517ABD5A8A9D306BCF67ED91F9E6725B4758C022E0B1EF4275BF7B6C5BFC11D45F9088B941F54EB1E59BB8BC39A0BF12307F5C4FDB70C581B23F76B63ACAE1CAA6B7902D52526735488A0EF13C6D9A51BFA4AB3AD8347796524D8EF6A167B5A41825D967E144E5140564251CCACB83E6B486F6B3CA3F7971506026C0B857F689962856DED4010ABD0BE621C3A3960A54E710C375F26375D7014103A4B54330C198AF126116D2276E11715F693877FAD7EF09CADB094AE91E1A1597").unwrap();
    let g = BigNum::from_hex_str("3FB32C9B73134D0B2E77506660EDBD484CA7B18F21EF205407F4793A1A0BA12510DBC15077BE463FFF4FED4AAC0BB555BE3A6C1B0C6B47B1BC3773BF7E8C6F62901228F8C28CBB18A55AE31341000A650196F931C77A57F2DDF463E5E9EC144B777DE62AAAB8A8628AC376D282D6ED3864E67982428EBC831D14348F6F2F9193B5045AF2767164E1DFC967C1FB3F2E55A4BD1BFFE83B9C80D052B985D182EA0ADB2A3B7313D3FE14C8484B1E052588B9B7D2BBD2DF016199ECD06E1557CD0915B3353BBB64E0EC377FD028370DF92B52C7891428CDC67EB6184B523D1DB246C32F63078490F00EF8D647D148D47954515E2327CFEF98C582664B4C0F6CC41659").unwrap();
    let q = BigNum::from_hex_str("8CF83642A709A097B447997640129DA299B1A47D1EB3750BA308B0FE64F5FBD3").unwrap();
    let dh = DH::from_params(p, g, q).unwrap();
    ctx.set_tmp_dh(dh).unwrap();
    set_ecdh(ctx);
}

fn config_parse(config: &Config) -> (TlsConfig, u32) {
    let threads = config.lookup("output.tls_threads").
        map_or(TLS_DEFAULT_THREADS, |x| x.as_integer().
            expect("output.tls_threads must be a 32-bit integer") as u32);
    let connect = config.lookup("output.connect").map_or(DEFAULT_CONNECT, |x|x.as_str().
        expect("output.connect must be an ip:port string")).to_owned();
    let cert = config.lookup("output.tls_cert").map_or(DEFAULT_CERT, |x| x.as_str().
        expect("output.tls_cert must be a path to a .pem file")).to_owned();
    let key = config.lookup("output.tls_key").map_or(DEFAULT_KEY, |x| x.as_str().
        expect("output.tls_key must be a path to a .pem file")).to_owned();
    let ciphers = config.lookup("output.tls_ciphers").map_or(DEFAULT_CIPHERS, |x| x.as_str().
        expect("output.tls_ciphers must be a string with a cipher suite")).to_owned();
    let tls_method = match config.lookup("output.tls_method").map_or(DEFAULT_TLS_METHOD, |x| x.as_str().
        expect("output.tls_method must be a string with the TLS method")).to_lowercase().as_ref() {
            "any" | "sslv23" => SslMethod::Sslv23,
            "tlsv1" | "tlsv1.0" => SslMethod::Tlsv1,
            "tlsv1.1" => SslMethod::Tlsv1_1,
            "tlsv1.2" => SslMethod::Tlsv1_2,
            _ => panic!(r#"TLS method must be "any", "TLSv1.0", "TLSv1.1" or "TLSv1.2""#)
    };
    let verify_peer = config.lookup("output.tls_verify_peer").map_or(DEFAULT_VERIFY_PEER, |x| x.as_bool().
        expect("output.tls_verify_peer must be a boolean"));
    let ca_file: Option<PathBuf> = config.lookup("output.tls_ca_file").map_or(None, |x|
        Some(PathBuf::from(x.as_str().expect("output.tls_ca_file must be a path to a file"))));
    let compression = config.lookup("output.tls_compression").map_or(DEFAULT_COMPRESSION, |x| x.as_bool().
        expect("output.tls_compression must be a boolean"));
    let timeout = config.lookup("output.timeout").map_or(DEFAULT_TIMEOUT, |x| x.as_integer().
        expect("output.timeout must be an integer") as u64);
    let async = config.lookup("output.tls_async").map_or(DEFAULT_ASYNC, |x| x.as_bool().
        expect("output.tls_async must be a boolean"));
    let mut ctx = SslContext::new(tls_method).unwrap();
    if verify_peer == false {
        ctx.set_verify(SSL_VERIFY_NONE, None);
    } else {
        ctx.set_verify_depth(TLS_VERIFY_DEPTH);
        ctx.set_verify(SSL_VERIFY_PEER | SSL_VERIFY_FAIL_IF_NO_PEER_CERT, None);
        if let Some(ca_file) = ca_file {
            if ctx.set_CA_file(&ca_file).is_err() {
                panic!("Unable to read the trusted CA file");
            }
        }
    }
    let mut opts = SSL_OP_CIPHER_SERVER_PREFERENCE | SSL_OP_NO_SESSION_RESUMPTION_ON_RENEGOTIATION;
    if compression == false {
        opts = opts | SSL_OP_NO_COMPRESSION;
    }
    ctx.set_options(opts);
    set_fs(&mut ctx);
    ctx.set_certificate_file(&Path::new(&cert), X509FileType::PEM).unwrap();
    ctx.set_private_key_file(&Path::new(&key), X509FileType::PEM).unwrap();
    ctx.set_cipher_list(&ciphers).unwrap();
    let arc_ctx = Arc::new(ctx);
    let tls_config = TlsConfig {
        connect: connect,
        timeout: Some(Duration::from_secs(timeout)),
        arc_ctx: arc_ctx,
        async: async
    };
    (tls_config, threads)
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use dom::attr::{Attr, AttrHelpers};
use dom::bindings::codegen::Bindings::DOMTokenListBinding;
use dom::bindings::codegen::Bindings::DOMTokenListBinding::DOMTokenListMethods;
use dom::bindings::error::{Fallible, InvalidCharacter, Syntax};
use dom::bindings::global::Window;
use dom::bindings::js::{JS, JSRef, Temporary, OptionalRootable};
use dom::bindings::utils::{Reflector, Reflectable, reflect_dom_object};
use dom::element::{Element, AttributeHandlers};
use dom::node::window_from_node;

use servo_util::str::{DOMString, HTML_SPACE_CHARACTERS};
use string_cache::Atom;

#[dom_struct]
pub struct DOMTokenList {
    reflector_: Reflector,
    element: JS<Element>,
    local_name: Atom,
}

impl DOMTokenList {
    pub fn new_inherited(element: JSRef<Element>, local_name: Atom) -> DOMTokenList {
        DOMTokenList {
            reflector_: Reflector::new(),
            element: JS::from_rooted(element),
            local_name: local_name,
        }
    }

    pub fn new(element: JSRef<Element>, local_name: &Atom) -> Temporary<DOMTokenList> {
        let window = window_from_node(element).root();
        reflect_dom_object(box DOMTokenList::new_inherited(element, local_name.clone()),
                           &Window(*window),
                           DOMTokenListBinding::Wrap)
    }
}

impl Reflectable for DOMTokenList {
    fn reflector<'a>(&'a self) -> &'a Reflector {
        &self.reflector_
    }
}

trait PrivateDOMTokenListHelpers {
    fn attribute(self) -> Option<Temporary<Attr>>;
    fn check_token_exceptions<'a>(self, token: &'a str) -> Fallible<&'a str>;
}

impl<'a> PrivateDOMTokenListHelpers for JSRef<'a, DOMTokenList> {
    fn attribute(self) -> Option<Temporary<Attr>> {
        let element = self.element.root();
        element.get_attribute(ns!(""), &self.local_name)
    }

    fn check_token_exceptions<'a>(self, token: &'a str) -> Fallible<&'a str> {
        match token {
            "" => Err(Syntax),
            token if token.find(HTML_SPACE_CHARACTERS).is_some() => Err(InvalidCharacter),
            token => Ok(token)
        }
    }
}

// http://dom.spec.whatwg.org/#domtokenlist
impl<'a> DOMTokenListMethods for JSRef<'a, DOMTokenList> {
    // http://dom.spec.whatwg.org/#dom-domtokenlist-length
    fn Length(self) -> u32 {
        self.attribute().root().map(|attr| {
            attr.value().tokens().map(|tokens| tokens.len()).unwrap_or(0)
        }).unwrap_or(0) as u32
    }

    // http://dom.spec.whatwg.org/#dom-domtokenlist-item
    fn Item(self, index: u32) -> Option<DOMString> {
        self.attribute().root().and_then(|attr| attr.value().tokens().and_then(|tokens| {
            tokens.get(index as uint).map(|token| token.as_slice().to_string())
        }))
    }

    fn IndexedGetter(self, index: u32, found: &mut bool) -> Option<DOMString> {
        let item = self.Item(index);
        *found = item.is_some();
        item
    }

    // http://dom.spec.whatwg.org/#dom-domtokenlist-contains
    fn Contains(self, token: DOMString) -> Fallible<bool> {
        self.check_token_exceptions(token.as_slice()).map(|slice| {
            self.attribute().root().and_then(|attr| attr.value().tokens().map(|tokens| {
                let atom = Atom::from_slice(slice);
                tokens.iter().any(|token| *token == atom)
            })).unwrap_or(false)
        })
    }
}
<commit_msg>Panic if DOMTokenList#contains is called for an unparsed attribute.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use dom::attr::{Attr, AttrHelpers};
use dom::bindings::codegen::Bindings::DOMTokenListBinding;
use dom::bindings::codegen::Bindings::DOMTokenListBinding::DOMTokenListMethods;
use dom::bindings::error::{Fallible, InvalidCharacter, Syntax};
use dom::bindings::global::Window;
use dom::bindings::js::{JS, JSRef, Temporary, OptionalRootable};
use dom::bindings::utils::{Reflector, Reflectable, reflect_dom_object};
use dom::element::{Element, AttributeHandlers};
use dom::node::window_from_node;

use servo_util::str::{DOMString, HTML_SPACE_CHARACTERS};
use string_cache::Atom;

#[dom_struct]
pub struct DOMTokenList {
    reflector_: Reflector,
    element: JS<Element>,
    local_name: Atom,
}

impl DOMTokenList {
    pub fn new_inherited(element: JSRef<Element>, local_name: Atom) -> DOMTokenList {
        DOMTokenList {
            reflector_: Reflector::new(),
            element: JS::from_rooted(element),
            local_name: local_name,
        }
    }

    pub fn new(element: JSRef<Element>, local_name: &Atom) -> Temporary<DOMTokenList> {
        let window = window_from_node(element).root();
        reflect_dom_object(box DOMTokenList::new_inherited(element, local_name.clone()),
                           &Window(*window),
                           DOMTokenListBinding::Wrap)
    }
}

impl Reflectable for DOMTokenList {
    fn reflector<'a>(&'a self) -> &'a Reflector {
        &self.reflector_
    }
}

trait PrivateDOMTokenListHelpers {
    fn attribute(self) -> Option<Temporary<Attr>>;
    fn check_token_exceptions<'a>(self, token: &'a str) -> Fallible<&'a str>;
}

impl<'a> PrivateDOMTokenListHelpers for JSRef<'a, DOMTokenList> {
    fn attribute(self) -> Option<Temporary<Attr>> {
        let element = self.element.root();
        element.get_attribute(ns!(""), &self.local_name)
    }

    fn check_token_exceptions<'a>(self, token: &'a str) -> Fallible<&'a str> {
        match token {
            "" => Err(Syntax),
            token if token.find(HTML_SPACE_CHARACTERS).is_some() => Err(InvalidCharacter),
            token => Ok(token)
        }
    }
}

// http://dom.spec.whatwg.org/#domtokenlist
impl<'a> DOMTokenListMethods for JSRef<'a, DOMTokenList> {
    // http://dom.spec.whatwg.org/#dom-domtokenlist-length
    fn Length(self) -> u32 {
        self.attribute().root().map(|attr| {
            attr.value().tokens().map(|tokens| tokens.len()).unwrap_or(0)
        }).unwrap_or(0) as u32
    }

    // http://dom.spec.whatwg.org/#dom-domtokenlist-item
    fn Item(self, index: u32) -> Option<DOMString> {
        self.attribute().root().and_then(|attr| attr.value().tokens().and_then(|tokens| {
            tokens.get(index as uint).map(|token| token.as_slice().to_string())
        }))
    }

    fn IndexedGetter(self, index: u32, found: &mut bool) -> Option<DOMString> {
        let item = self.Item(index);
        *found = item.is_some();
        item
    }

    // http://dom.spec.whatwg.org/#dom-domtokenlist-contains
    fn Contains(self, token: DOMString) -> Fallible<bool> {
        self.check_token_exceptions(token.as_slice()).map(|slice| {
            self.attribute().root().map(|attr| {
                let value = attr.value();
                let tokens = value.tokens()
                                  .expect("Should have parsed this attribute");
                let atom = Atom::from_slice(slice);
                tokens.iter().any(|token| *token == atom)
            }).unwrap_or(false)
        })
    }
}
<|endoftext|>


<commit_before>//! A single threaded (st) runtime (rt) for Weld.

extern crate libc;
extern crate fnv;

use libc::{c_char, int32_t, int64_t, size_t};
use libc::{malloc, realloc, free};

use fnv::FnvHashMap;

use std::ffi::CStr;
use std::ptr;
use std::sync::{Once, ONCE_INIT};

use super::WeldRuntimeErrno;

pub type WeldRunRef = *mut WeldRun;

/// An opaque pointer.
type Pointer = *mut libc::c_void;

/// Initialize the Weld runtime only once.
static ONCE: Once = ONCE_INIT;
static mut INITIALIZE_FAILED: bool = false;

/// Maintains information about a single Weld run.
#[derive(Debug)]
pub struct WeldRun {
    allocations: FnvHashMap<Pointer, i64>,
    errno: WeldRuntimeErrno,
    result: Pointer,
    nworkers: i32,
    memlimit: i64,
    allocated: i64,
}

/// Private API used by the FFI.
impl WeldRun {
    fn new(nworkers: i32, memlimit: i64) -> WeldRun {
        WeldRun {
            allocations: FnvHashMap::default(),
            errno: WeldRuntimeErrno::Success,
            result: ptr::null_mut(),
            nworkers: nworkers,
            memlimit: memlimit,
            allocated: 0,
        }
    }

    unsafe fn malloc(&mut self, size: i64) -> Pointer {
        if self.allocated + size > self.memlimit {
            self.errno = WeldRuntimeErrno::OutOfMemory;
            panic!("Weld run ran out of memory (limit={}, attempted to allocate {}",
                   self.memlimit,
                   self.allocated + size);
        }
        let mem = malloc(size as size_t);
        assert!(mem as i64 != 0);
        self.allocated += size;
        self.allocations.insert(mem.clone(), size);
        mem
    }

    unsafe fn realloc(&mut self, pointer: Pointer, size: i64) -> Pointer {
        let allocated = *self.allocations.get(&pointer).unwrap();
        if self.allocated - allocated + size > self.memlimit {
            self.errno = WeldRuntimeErrno::OutOfMemory;
            panic!("Weld run ran out of memory (limit={}, attempted to allocate {}",
                   self.memlimit,
                   self.allocated - allocated + size);
        }

        let mem = realloc(pointer, size as size_t);
        assert!(mem as i64 != 0);
        self.allocated -= allocated;
        self.allocated += size;
        let _ = self.allocations.remove(&pointer);
        self.allocations.insert(mem.clone(), size);
        mem
    }

    unsafe fn free(&mut self, pointer: Pointer) {
        let allocated = *self.allocations.get(&pointer).unwrap();
        free(pointer);
        self.allocated -= allocated;
        self.allocations.remove(&pointer);
    }

    fn set_errno(&mut self, errno: WeldRuntimeErrno) {
        self.errno = errno;
        panic!("Weld runtime threw error: {}", self.errno)
    }

    fn set_result(&mut self, result: Pointer) {
        self.result = result;
    }

    fn result(&self) -> Pointer {
        self.result
    }

    fn errno(&self) -> WeldRuntimeErrno {
        self.errno
    }
}

/// Public read-only API used to query information about the run.
impl WeldRun {
    /// Returns the number of bytes allocated by this Weld run.
    pub fn memory_usage(&self) -> i64 {
        self.allocations.values().sum()
    }

    /// Returns a 64-bit ID identifying this run.
    pub fn run_id(&self) -> i64 {
        0
    }
}

impl Drop for WeldRun {
    fn drop(&mut self) {
        // Free memory allocated by the run.
        unsafe {
            for (pointer, _) in self.allocations.iter() {
                free(*pointer);
            }
        }
    }
}

unsafe fn initialize() {
    // Hack to prevent symbols from being compiled out in a Rust binary.
    let mut x =  weld_runst_init as i64;
    x += weld_runst_set_result as i64;
    x += weld_runst_get_result as i64;
    x += weld_runst_malloc as i64;
    x += weld_runst_realloc as i64;
    x += weld_runst_free as i64;
    x += weld_runst_set_errno as i64;

    use super::link::*;
    x += weld_st_dict_new as i64;
    x += weld_st_dict_get as i64;
    x += weld_st_dict_upsert as i64;
    x += weld_st_dict_keyexists as i64;
    x += weld_st_dict_size as i64;
    x += weld_st_dict_tovec as i64;
    x += weld_st_dict_serialize as i64;

    x += weld_st_gb_new as i64;
    x += weld_st_gb_merge as i64;
    x += weld_st_gb_result as i64;

    trace!("Runtime initialized with hashed values {}", x);
}

#[no_mangle]
/// Initialize the runtime.
///
/// This call currently circumvents an awkward problem with binaries using Rust, where the FFI
/// below used by the Weld runtime is compiled out.
pub unsafe extern "C" fn weld_init() {
    ONCE.call_once(|| initialize());
    if INITIALIZE_FAILED {
        panic!("Weld runtime initialization failed")
    }
}

#[no_mangle]
/// Initialize the runtime if necessary, and then initialize run-specific information.
pub unsafe extern "C" fn weld_runst_init(nworkers: int32_t, memlimit: int64_t) -> WeldRunRef {
    Box::into_raw(Box::new(WeldRun::new(nworkers, memlimit)))
}

#[no_mangle]
/// Allocate memory using `libc` `malloc`, tracking memory usage information.
pub unsafe extern "C" fn weld_runst_malloc(run: WeldRunRef, size: int64_t) -> Pointer {
    let run = &mut *run;
    run.malloc(size)
}

#[no_mangle]
/// Allocate memory using `libc` `remalloc`, tracking memory usage information.
pub unsafe extern "C" fn weld_runst_realloc(run: WeldRunRef,
                                     ptr: Pointer,
                                     newsize: int64_t) -> Pointer {
    let run = &mut *run;
    run.realloc(ptr, newsize)
}

#[no_mangle]
/// Free memory using `libc` `free`, tracking memory usage information.
pub unsafe extern "C" fn weld_runst_free(run: WeldRunRef, ptr: Pointer) {
    let run = &mut *run;
    run.free(ptr)
}

#[no_mangle]
/// Set the result pointer.
pub unsafe extern "C" fn weld_runst_set_result(run: WeldRunRef, ptr: Pointer) {
    let run = &mut *run;
    run.set_result(ptr)
}

#[no_mangle]
/// Set the errno value.
pub unsafe extern "C" fn weld_runst_set_errno(run: WeldRunRef, errno: WeldRuntimeErrno) {
    let run = &mut *run;
    run.set_errno(errno)
}

#[no_mangle]
/// Get the errno value.
pub unsafe extern "C" fn weld_runst_get_errno(run: WeldRunRef) -> WeldRuntimeErrno {
    let run = &mut *run;
    run.errno()
}

#[no_mangle]
/// Get the result pointer.
pub unsafe extern "C" fn weld_runst_get_result(run: WeldRunRef) -> Pointer {
    let run = &mut *run;
    run.result()
}

#[no_mangle]
/// Print a value from generated code.
pub unsafe extern "C" fn weld_runst_print(_run: WeldRunRef, string: *const c_char) {
    let string = CStr::from_ptr(string).to_str().unwrap();
    // XXX We could switch to a custom printer here too.
    println!("{}", string);
}
<commit_msg>Use symbols to prevent them from being compiled out<commit_after>//! A single threaded (st) runtime (rt) for Weld.

extern crate libc;
extern crate fnv;

use libc::{c_char, int32_t, int64_t, size_t};
use libc::{malloc, realloc, free};

use fnv::FnvHashMap;

use std::ffi::CStr;
use std::ptr;
use std::sync::{Once, ONCE_INIT};

use super::WeldRuntimeErrno;

pub type WeldRunRef = *mut WeldRun;

/// An opaque pointer.
type Pointer = *mut libc::c_void;

/// Initialize the Weld runtime only once.
static ONCE: Once = ONCE_INIT;
static mut INITIALIZE_FAILED: bool = false;

/// Maintains information about a single Weld run.
#[derive(Debug)]
pub struct WeldRun {
    allocations: FnvHashMap<Pointer, i64>,
    errno: WeldRuntimeErrno,
    result: Pointer,
    nworkers: i32,
    memlimit: i64,
    allocated: i64,
}

/// Private API used by the FFI.
impl WeldRun {
    fn new(nworkers: i32, memlimit: i64) -> WeldRun {
        WeldRun {
            allocations: FnvHashMap::default(),
            errno: WeldRuntimeErrno::Success,
            result: ptr::null_mut(),
            nworkers: nworkers,
            memlimit: memlimit,
            allocated: 0,
        }
    }

    unsafe fn malloc(&mut self, size: i64) -> Pointer {
        if self.allocated + size > self.memlimit {
            self.errno = WeldRuntimeErrno::OutOfMemory;
            panic!("Weld run ran out of memory (limit={}, attempted to allocate {}",
                   self.memlimit,
                   self.allocated + size);
        }
        let mem = malloc(size as size_t);
        assert!(mem as i64 != 0);
        self.allocated += size;
        self.allocations.insert(mem.clone(), size);
        mem
    }

    unsafe fn realloc(&mut self, pointer: Pointer, size: i64) -> Pointer {
        let allocated = *self.allocations.get(&pointer).unwrap();
        if self.allocated - allocated + size > self.memlimit {
            self.errno = WeldRuntimeErrno::OutOfMemory;
            panic!("Weld run ran out of memory (limit={}, attempted to allocate {}",
                   self.memlimit,
                   self.allocated - allocated + size);
        }

        let mem = realloc(pointer, size as size_t);
        assert!(mem as i64 != 0);
        self.allocated -= allocated;
        self.allocated += size;
        let _ = self.allocations.remove(&pointer);
        self.allocations.insert(mem.clone(), size);
        mem
    }

    unsafe fn free(&mut self, pointer: Pointer) {
        let allocated = *self.allocations.get(&pointer).unwrap();
        free(pointer);
        self.allocated -= allocated;
        self.allocations.remove(&pointer);
    }

    fn set_errno(&mut self, errno: WeldRuntimeErrno) {
        self.errno = errno;
        panic!("Weld runtime threw error: {}", self.errno)
    }

    fn set_result(&mut self, result: Pointer) {
        self.result = result;
    }

    fn result(&self) -> Pointer {
        self.result
    }

    fn errno(&self) -> WeldRuntimeErrno {
        self.errno
    }
}

/// Public read-only API used to query information about the run.
impl WeldRun {
    /// Returns the number of bytes allocated by this Weld run.
    pub fn memory_usage(&self) -> i64 {
        self.allocations.values().sum()
    }

    /// Returns a 64-bit ID identifying this run.
    pub fn run_id(&self) -> i64 {
        0
    }
}

impl Drop for WeldRun {
    fn drop(&mut self) {
        // Free memory allocated by the run.
        unsafe {
            for (pointer, _) in self.allocations.iter() {
                free(*pointer);
            }
        }
    }
}

unsafe fn initialize() {
    // Hack to prevent symbols from being compiled out in a Rust binary.
    let mut x =  weld_runst_init as i64;
    x += weld_runst_set_result as i64;
    x += weld_runst_get_result as i64;
    x += weld_runst_malloc as i64;
    x += weld_runst_realloc as i64;
    x += weld_runst_free as i64;
    x += weld_runst_get_errno as i64;
    x += weld_runst_set_errno as i64;

    use super::link::*;
    x += weld_st_dict_new as i64;
    x += weld_st_dict_get as i64;
    x += weld_st_dict_upsert as i64;
    x += weld_st_dict_keyexists as i64;
    x += weld_st_dict_size as i64;
    x += weld_st_dict_tovec as i64;
    x += weld_st_dict_serialize as i64;

    x += weld_st_gb_new as i64;
    x += weld_st_gb_merge as i64;
    x += weld_st_gb_result as i64;

    // Weld Multi-threaded runtime.
    x += weld_runtime_init as i64;
    x += weld_rt_thread_id as i64;
    x += weld_rt_abort_thread as i64;
    x += weld_rt_get_nworkers as i64;
    x += weld_rt_get_run_id as i64;
    x += weld_rt_start_loop as i64;
    x += weld_rt_set_result as i64;
    x += weld_rt_new_vb as i64;
    x += weld_rt_new_vb_piece as i64;
    x += weld_rt_cur_vb_piece as i64;
    x += weld_rt_set_vb_offset_if_fixed as i64;
    x += weld_rt_result_vb as i64;
    x += weld_rt_new_merger as i64;
    x += weld_rt_get_merger_at_index as i64;
    x += weld_rt_free_merger as i64;
    x += weld_run_begin as i64;
    x += weld_run_get_result as i64;
    x += weld_run_dispose as i64;
    x += weld_run_malloc as i64;
    x += weld_run_realloc as i64;
    x += weld_run_free as i64;
    x += weld_run_memory_usage as i64;
    x += weld_run_get_errno as i64;
    x += weld_run_set_errno as i64;
    x += weld_rt_dict_new as i64;
    x += weld_rt_dict_new_finalized as i64;
    x += weld_rt_dict_lookup as i64;
    x += weld_rt_dict_merge as i64;
    x += weld_rt_dict_finalize as i64;
    x += weld_rt_dict_to_array as i64;
    x += weld_rt_dict_size as i64;
    x += weld_rt_dict_serialize as i64;
    x += weld_rt_dict_free as i64;
    x += weld_rt_gb_new as i64;
    x += weld_rt_gb_merge as i64;
    x += weld_rt_gb_result as i64;
    x += weld_rt_gb_free as i64;

    trace!("Runtime initialized with hashed values {}", x);
}

#[no_mangle]
/// Initialize the runtime.
///
/// This call currently circumvents an awkward problem with binaries using Rust, where the FFI
/// below used by the Weld runtime is compiled out.
pub unsafe extern "C" fn weld_init() {
    ONCE.call_once(|| initialize());
    if INITIALIZE_FAILED {
        panic!("Weld runtime initialization failed")
    }
}

#[no_mangle]
/// Initialize the runtime if necessary, and then initialize run-specific information.
pub unsafe extern "C" fn weld_runst_init(nworkers: int32_t, memlimit: int64_t) -> WeldRunRef {
    Box::into_raw(Box::new(WeldRun::new(nworkers, memlimit)))
}

#[no_mangle]
/// Allocate memory using `libc` `malloc`, tracking memory usage information.
pub unsafe extern "C" fn weld_runst_malloc(run: WeldRunRef, size: int64_t) -> Pointer {
    let run = &mut *run;
    run.malloc(size)
}

#[no_mangle]
/// Allocate memory using `libc` `remalloc`, tracking memory usage information.
pub unsafe extern "C" fn weld_runst_realloc(run: WeldRunRef,
                                     ptr: Pointer,
                                     newsize: int64_t) -> Pointer {
    let run = &mut *run;
    run.realloc(ptr, newsize)
}

#[no_mangle]
/// Free memory using `libc` `free`, tracking memory usage information.
pub unsafe extern "C" fn weld_runst_free(run: WeldRunRef, ptr: Pointer) {
    let run = &mut *run;
    run.free(ptr)
}

#[no_mangle]
/// Set the result pointer.
pub unsafe extern "C" fn weld_runst_set_result(run: WeldRunRef, ptr: Pointer) {
    let run = &mut *run;
    run.set_result(ptr)
}

#[no_mangle]
/// Set the errno value.
pub unsafe extern "C" fn weld_runst_set_errno(run: WeldRunRef, errno: WeldRuntimeErrno) {
    let run = &mut *run;
    run.set_errno(errno)
}

#[no_mangle]
/// Get the errno value.
pub unsafe extern "C" fn weld_runst_get_errno(run: WeldRunRef) -> WeldRuntimeErrno {
    let run = &mut *run;
    run.errno()
}

#[no_mangle]
/// Get the result pointer.
pub unsafe extern "C" fn weld_runst_get_result(run: WeldRunRef) -> Pointer {
    let run = &mut *run;
    run.result()
}

#[no_mangle]
/// Print a value from generated code.
pub unsafe extern "C" fn weld_runst_print(_run: WeldRunRef, string: *const c_char) {
    let string = CStr::from_ptr(string).to_str().unwrap();
    // XXX We could switch to a custom printer here too.
    println!("{}", string);
}
<|endoftext|>


<commit_before>use std::vec;

use naulang::interpreter::frame::{Frame, FrameInfo};
use naulang::objectspace::primitives::Object;
use naulang::interpreter::bytecode::Bytecode;

#[derive(Clone)]
pub struct MethodObject<'m> {
    pub literals: &'m vec::Vec<Object<'m>>,
    pub stack_depth: usize,
    bytecodes: vec::Vec<u32>,
    arg_count: usize,
    // TODO: SourceMaps
}

impl<'m> MethodObject<'m> {

    pub fn new(
        bytecodes: vec::Vec<u32>, literals: &'m vec::Vec<Object<'m>>,  arg_count: usize)
    -> MethodObject<'m> {

        let stack_depth = Bytecode::calculate_stack_depth(&bytecodes);
        MethodObject {
            bytecodes: bytecodes.clone(),
            literals: literals,
            arg_count: arg_count,
            stack_depth: stack_depth,
        }
    }

    pub fn get_bytecode(&self, at_point: usize) -> u32 {
        self.bytecodes[at_point]
    }

    fn create_new_frame<'a>(&'a self, mut previous_frame: Frame<'a>, is_async: bool) -> Box<Frame> {
        let mut new_frame = Frame::new(FrameInfo {
            stack_depth: self.stack_depth,
            local_count: 0,
            literal_count: 0,
            method: self,
        });

        for index in (self.arg_count - 1)..0 {
            match previous_frame.pop() {
                Some(object) => new_frame.set_local_at(index, object),
                None => ()
            }
        }

        return new_frame;
    }
}

#[cfg(test)]
mod tests {
    use super::MethodObject;
    use naulang::interpreter::bytecode::Bytecode;

    #[test]
    fn test_get_bytecode() {
        let literals = vec!();
        let method: MethodObject = MethodObject::new(vec![
            Bytecode::LOAD_CONST, 0,
        ], &literals, 0);

        assert!(method.get_bytecode(0) == Bytecode::LOAD_CONST);
    }
}
<commit_msg>Add get_literal_at method for MethodObject<commit_after>use std::vec;

use naulang::interpreter::frame::{Frame, FrameInfo};
use naulang::objectspace::primitives::Object;
use naulang::interpreter::bytecode::Bytecode;

#[derive(Clone)]
pub struct MethodObject<'m> {
    pub literals: &'m vec::Vec<Object<'m>>,
    pub stack_depth: usize,
    bytecodes: vec::Vec<u32>,
    arg_count: usize,
    // TODO: SourceMaps
}

impl<'m> MethodObject<'m> {

    pub fn new(
        bytecodes: vec::Vec<u32>, literals: &'m vec::Vec<Object<'m>>,  arg_count: usize)
    -> MethodObject<'m> {

        let stack_depth = Bytecode::calculate_stack_depth(&bytecodes);
        MethodObject {
            bytecodes: bytecodes.clone(),
            literals: literals,
            arg_count: arg_count,
            stack_depth: stack_depth,
        }
    }

    pub fn get_bytecode(&self, at_point: usize) -> u32 {
        self.bytecodes[at_point]
    }

    pub fn get_literal_at(&self, index: usize) -> Option<&Object> {
        self.literals.get(index)
    }

    fn create_new_frame<'a>(&'a self, mut previous_frame: Frame<'a>, is_async: bool) -> Box<Frame> {
        let mut new_frame = Frame::new(FrameInfo {
            stack_depth: self.stack_depth,
            local_count: 0,
            literal_count: 0,
            method: self,
        });

        for index in (self.arg_count - 1)..0 {
            match previous_frame.pop() {
                Some(object) => new_frame.set_local_at(index, object),
                None => ()
            }
        }

        return new_frame;
    }
}

#[cfg(test)]
mod tests {
    use super::MethodObject;
    use naulang::objectspace::primitives::{Object};
    use naulang::interpreter::bytecode::Bytecode;

    #[test]
    fn test_get_literal_at() {
        let literals = vec!(Object::None);
        let method = MethodObject::new(vec!(Bytecode::HALT), &literals, 0);

        assert!(match method.get_literal_at(0) {
            Some(obj) => match *obj {
                Object::None => true,
                _ => false
            },
            None => false
        });
    }

    #[test]
    fn test_get_bytecode() {
        let literals = vec!();
        let method: MethodObject = MethodObject::new(vec!(
            Bytecode::LOAD_CONST, 0,
        ), &literals, 0);

        assert!(method.get_bytecode(0) == Bytecode::LOAD_CONST);
    }
}
<|endoftext|>


<commit_before>// Copyright 2020 The Tink-Rust Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
////////////////////////////////////////////////////////////////////////////////

use tink::{keyset, subtle::random::get_random_bytes};

#[test]
fn test_validate_key_version() {
    assert!(keyset::validate_key_version(2, 1).is_err());
    assert!(keyset::validate_key_version(1, 1).is_ok());
    assert!(keyset::validate_key_version(1, 2).is_ok());
}

#[test]
fn test_validate() {
    // empty keyset
    let empty_keys = vec![];
    assert!(
        keyset::validate(&tink_testutil::new_keyset(1, empty_keys)).is_err(),
        "expect an error when keyset is empty"
    );
    // no primary key
    let keys = vec![tink_testutil::new_dummy_key(
        1,
        tink::proto::KeyStatusType::Enabled,
        tink::proto::OutputPrefixType::Tink,
    )];
    assert!(
        keyset::validate(&tink_testutil::new_keyset(2, keys)).is_err(),
        "expect an error when there is no primary key"
    );
    // primary key is disabled
    let keys = vec![
        tink_testutil::new_dummy_key(
            1,
            tink::proto::KeyStatusType::Enabled,
            tink::proto::OutputPrefixType::Tink,
        ),
        tink_testutil::new_dummy_key(
            2,
            tink::proto::KeyStatusType::Disabled,
            tink::proto::OutputPrefixType::Legacy,
        ),
    ];
    assert!(
        keyset::validate(&tink_testutil::new_keyset(2, keys)).is_err(),
        "expect an error when primary key is disabled"
    );
    // multiple primary keys
    let keys = vec![
        tink_testutil::new_dummy_key(
            1,
            tink::proto::KeyStatusType::Enabled,
            tink::proto::OutputPrefixType::Tink,
        ),
        tink_testutil::new_dummy_key(
            1,
            tink::proto::KeyStatusType::Enabled,
            tink::proto::OutputPrefixType::Legacy,
        ),
    ];
    assert!(
        keyset::validate(&tink_testutil::new_keyset(1, keys)).is_err(),
        "expect an error when there are multiple primary keys"
    );
    // invalid keys
    let invalid_keys = generate_invalid_keys();
    for (i, key) in invalid_keys.into_iter().enumerate() {
        assert!(
            keyset::validate(&tink_testutil::new_keyset(1, vec![key])).is_err(),
            "expect an error when validate invalid key {}",
            i
        );
    }
    // no primary keys
    let keys = vec![
        tink_testutil::new_dummy_key(
            1,
            tink::proto::KeyStatusType::Disabled,
            tink::proto::OutputPrefixType::Tink,
        ),
        tink_testutil::new_dummy_key(
            1,
            tink::proto::KeyStatusType::Disabled,
            tink::proto::OutputPrefixType::Legacy,
        ),
    ];
    assert!(
        keyset::validate(&tink_testutil::new_keyset(1, keys)).is_err(),
        "expect an error when there are no primary keys"
    );
    // public key only
    let keys = vec![tink_testutil::new_key(
        &tink_testutil::new_key_data(
            tink_testutil::ECIES_AEAD_HKDF_PUBLIC_KEY_TYPE_URL,
            &get_random_bytes(10),
            tink::proto::key_data::KeyMaterialType::AsymmetricPublic,
        ),
        tink::proto::KeyStatusType::Enabled,
        1,
        tink::proto::OutputPrefixType::Tink,
    )];
    assert!(
        keyset::validate(&tink_testutil::new_keyset(1, keys)).is_ok(),
        "valid test failed when using public key only"
    );
    // private key
    let keys = vec![
        tink_testutil::new_key(
            &tink_testutil::new_key_data(
                tink_testutil::ECIES_AEAD_HKDF_PUBLIC_KEY_TYPE_URL,
                &get_random_bytes(10),
                tink::proto::key_data::KeyMaterialType::AsymmetricPublic,
            ),
            tink::proto::KeyStatusType::Enabled,
            1,
            tink::proto::OutputPrefixType::Tink,
        ),
        tink_testutil::new_key(
            &tink_testutil::new_key_data(
                tink_testutil::ECIES_AEAD_HKDF_PRIVATE_KEY_TYPE_URL,
                &get_random_bytes(10),
                tink::proto::key_data::KeyMaterialType::AsymmetricPrivate,
            ),
            tink::proto::KeyStatusType::Enabled,
            1,
            tink::proto::OutputPrefixType::Tink,
        ),
    ];
    assert!(
        keyset::validate(&tink_testutil::new_keyset(1, keys)).is_err(),
        "expect an error when there are keydata other than public"
    );
}

fn generate_invalid_keys() -> Vec<tink::proto::keyset::Key> {
    vec![
        // unknown status
        tink_testutil::new_key(
            &tink::proto::KeyData::default(),
            tink::proto::KeyStatusType::UnknownStatus,
            1,
            tink::proto::OutputPrefixType::Tink,
        ),
        // unknown prefix
        tink_testutil::new_key(
            &tink::proto::KeyData::default(),
            tink::proto::KeyStatusType::Enabled,
            1,
            tink::proto::OutputPrefixType::UnknownPrefix,
        ),
    ]
}
<commit_msg>tink::keyset: more validate_key tests<commit_after>// Copyright 2020 The Tink-Rust Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
////////////////////////////////////////////////////////////////////////////////

use tink::{keyset, subtle::random::get_random_bytes};

#[test]
fn test_validate_key_version() {
    assert!(keyset::validate_key_version(2, 1).is_err());
    assert!(keyset::validate_key_version(1, 1).is_ok());
    assert!(keyset::validate_key_version(1, 2).is_ok());
}

#[test]
fn test_validate() {
    // empty keyset
    let empty_keys = vec![];
    assert!(
        keyset::validate(&tink_testutil::new_keyset(1, empty_keys)).is_err(),
        "expect an error when keyset is empty"
    );
    // no primary key
    let keys = vec![tink_testutil::new_dummy_key(
        1,
        tink::proto::KeyStatusType::Enabled,
        tink::proto::OutputPrefixType::Tink,
    )];
    assert!(
        keyset::validate(&tink_testutil::new_keyset(2, keys)).is_err(),
        "expect an error when there is no primary key"
    );
    // primary key is disabled
    let keys = vec![
        tink_testutil::new_dummy_key(
            1,
            tink::proto::KeyStatusType::Enabled,
            tink::proto::OutputPrefixType::Tink,
        ),
        tink_testutil::new_dummy_key(
            2,
            tink::proto::KeyStatusType::Disabled,
            tink::proto::OutputPrefixType::Legacy,
        ),
    ];
    assert!(
        keyset::validate(&tink_testutil::new_keyset(2, keys)).is_err(),
        "expect an error when primary key is disabled"
    );
    // multiple primary keys
    let keys = vec![
        tink_testutil::new_dummy_key(
            1,
            tink::proto::KeyStatusType::Enabled,
            tink::proto::OutputPrefixType::Tink,
        ),
        tink_testutil::new_dummy_key(
            1,
            tink::proto::KeyStatusType::Enabled,
            tink::proto::OutputPrefixType::Legacy,
        ),
    ];
    assert!(
        keyset::validate(&tink_testutil::new_keyset(1, keys)).is_err(),
        "expect an error when there are multiple primary keys"
    );
    // invalid keys
    let invalid_keys = generate_invalid_keys();
    for (i, key) in invalid_keys.into_iter().enumerate() {
        assert!(
            keyset::validate(&tink_testutil::new_keyset(1, vec![key])).is_err(),
            "expect an error when validate invalid key {}",
            i
        );
    }
    // no primary keys
    let keys = vec![
        tink_testutil::new_dummy_key(
            1,
            tink::proto::KeyStatusType::Disabled,
            tink::proto::OutputPrefixType::Tink,
        ),
        tink_testutil::new_dummy_key(
            1,
            tink::proto::KeyStatusType::Disabled,
            tink::proto::OutputPrefixType::Legacy,
        ),
    ];
    assert!(
        keyset::validate(&tink_testutil::new_keyset(1, keys)).is_err(),
        "expect an error when there are no primary keys"
    );
    // public key only
    let keys = vec![tink_testutil::new_key(
        &tink_testutil::new_key_data(
            tink_testutil::ECIES_AEAD_HKDF_PUBLIC_KEY_TYPE_URL,
            &get_random_bytes(10),
            tink::proto::key_data::KeyMaterialType::AsymmetricPublic,
        ),
        tink::proto::KeyStatusType::Enabled,
        1,
        tink::proto::OutputPrefixType::Tink,
    )];
    assert!(
        keyset::validate(&tink_testutil::new_keyset(1, keys)).is_ok(),
        "valid test failed when using public key only"
    );
    // private key
    let keys = vec![
        tink_testutil::new_key(
            &tink_testutil::new_key_data(
                tink_testutil::ECIES_AEAD_HKDF_PUBLIC_KEY_TYPE_URL,
                &get_random_bytes(10),
                tink::proto::key_data::KeyMaterialType::AsymmetricPublic,
            ),
            tink::proto::KeyStatusType::Enabled,
            1,
            tink::proto::OutputPrefixType::Tink,
        ),
        tink_testutil::new_key(
            &tink_testutil::new_key_data(
                tink_testutil::ECIES_AEAD_HKDF_PRIVATE_KEY_TYPE_URL,
                &get_random_bytes(10),
                tink::proto::key_data::KeyMaterialType::AsymmetricPrivate,
            ),
            tink::proto::KeyStatusType::Enabled,
            1,
            tink::proto::OutputPrefixType::Tink,
        ),
    ];
    assert!(
        keyset::validate(&tink_testutil::new_keyset(1, keys)).is_err(),
        "expect an error when there are keydata other than public"
    );
}

fn generate_invalid_keys() -> Vec<tink::proto::keyset::Key> {
    vec![
        // unknown status
        tink_testutil::new_key(
            &tink::proto::KeyData::default(),
            tink::proto::KeyStatusType::UnknownStatus,
            1,
            tink::proto::OutputPrefixType::Tink,
        ),
        // unknown prefix
        tink_testutil::new_key(
            &tink::proto::KeyData::default(),
            tink::proto::KeyStatusType::Enabled,
            1,
            tink::proto::OutputPrefixType::UnknownPrefix,
        ),
        // zero key id
        tink_testutil::new_key(
            &tink::proto::KeyData::default(),
            tink::proto::KeyStatusType::Enabled,
            0,
            tink::proto::OutputPrefixType::Tink,
        ),
        // no key_data
        tink::proto::keyset::Key {
            key_data: None,
            status: tink::proto::KeyStatusType::Enabled as i32,
            key_id: 1,
            output_prefix_type: tink::proto::OutputPrefixType::Tink as i32,
        },
    ]
}
<|endoftext|>


<commit_before>use std::collections::VecDeque;

use codegen::Instructions;
use codegen::Instruction::*;

use super::{Optimize, OptimizationLevel};

impl Optimize for Instructions {
    fn optimize(mut self: Instructions, level: OptimizationLevel) -> Instructions {
        // Optimizers should be grouped by the number of assumptions
        // the optimizer makes about the code
        // The fewer assumptions, the lower the applicable optimization level
        // Optimizers should be ordered appropriately so that they
        // do not conflict or contradict each other
        let optimizers: Vec<fn(&mut Instructions)> = match level {
            OptimizationLevel::Off => vec![],
            OptimizationLevel::L1 => vec![
                remove_opposites,
            ],
            OptimizationLevel::L2 => vec![
                truncate_no_side_effects,
                remove_opposites,
            ],
        };

        for optimize in optimizers {
            optimize(&mut self);
        }

        self
    }
}

/// Removes instructions from the end of the given instructions which do not have any side effects.
///
/// The idea here is that lots of instructions like +, -, >, <, and any loop consisting of only
/// those instructions do not do anything to the output of the program. Thus, for our purposes at
/// least, it is safe to remove those instructions completely and avoid running them at all.
///
/// Often, the reason these instructions are there in the first place is because the compiler
/// ensures that certain memory cells are "freed" or "zeroed" after their use. There is no reason
/// to do this at the end of a program, so let's get rid of it entirely.
fn truncate_no_side_effects(instructions: &mut Instructions) {
    // In this algorithm, we're going to search for the last instruction that has a side effect,
    // and remove (truncate) everything past it
    let mut last_side_effect = None;
    // If we find a side effect inside a loop, we want to leave the loop intact
    let mut jump_stack = VecDeque::new();

    for (i, &instr) in instructions.iter().enumerate().rev() {
        match instr {
            // no side effects
            Left | Right | Increment | Decrement => {},
            // side effects
            Read | Write => {
                last_side_effect = Some(i);
                break;
            },
            JumpBackwardUnlessZero => jump_stack.push_back(i),
            JumpForwardIfZero => {jump_stack.pop_back().unwrap();},
        }
    }

    if let Some(last_side_effect) = last_side_effect {
        if jump_stack.is_empty() {
            instructions.truncate(last_side_effect + 1);
        }
        else {
            instructions.truncate(jump_stack.front().unwrap() + 1);
        }
    }
    else {
        // no side effects, so the instructions may as well not be there at all
        //TODO: Determine if we should actually be doing this?
        //TODO: Maybe it would be better to just leave them as is if there are no side effects
        //TODO: Even if that means that nothing can possibly be read or written at runtime
        instructions.clear();
    }
}

fn remove_opposites(instructions: &mut Instructions) {
    let mut i = 1;
    while i < instructions.len() {
        let prev = instructions[i - 1];
        let current = instructions[i];

        // Cancel out opposites
        match (prev, current) {
            (Left, Right) | (Right, Left) |
            (Increment, Decrement) | (Decrement, Increment) => {
                instructions.remove(i);
                instructions.remove(i - 1);
                i -= 1;
            },

            // Otherwise just move on
            _ => i += 1,
        }
    }
}
<commit_msg>Fixed a bug that was causing the optimizer to panic if it cancelled out opposites at the beginning of the program<commit_after>use std::collections::VecDeque;

use codegen::Instructions;
use codegen::Instruction::*;

use super::{Optimize, OptimizationLevel};

impl Optimize for Instructions {
    fn optimize(mut self: Instructions, level: OptimizationLevel) -> Instructions {
        // Optimizers should be grouped by the number of assumptions
        // the optimizer makes about the code
        // The fewer assumptions, the lower the applicable optimization level
        // Optimizers should be ordered appropriately so that they
        // do not conflict or contradict each other
        let optimizers: Vec<fn(&mut Instructions)> = match level {
            OptimizationLevel::Off => vec![],
            OptimizationLevel::L1 => vec![
                remove_opposites,
            ],
            OptimizationLevel::L2 => vec![
                truncate_no_side_effects,
                remove_opposites,
            ],
        };

        for optimize in optimizers {
            optimize(&mut self);
        }

        self
    }
}

/// Removes instructions from the end of the given instructions which do not have any side effects.
///
/// The idea here is that lots of instructions like +, -, >, <, and any loop consisting of only
/// those instructions do not do anything to the output of the program. Thus, for our purposes at
/// least, it is safe to remove those instructions completely and avoid running them at all.
///
/// Often, the reason these instructions are there in the first place is because the compiler
/// ensures that certain memory cells are "freed" or "zeroed" after their use. There is no reason
/// to do this at the end of a program, so let's get rid of it entirely.
fn truncate_no_side_effects(instructions: &mut Instructions) {
    // In this algorithm, we're going to search for the last instruction that has a side effect,
    // and remove (truncate) everything past it
    let mut last_side_effect = None;
    // If we find a side effect inside a loop, we want to leave the loop intact
    let mut jump_stack = VecDeque::new();

    for (i, &instr) in instructions.iter().enumerate().rev() {
        match instr {
            // no side effects
            Left | Right | Increment | Decrement => {},
            // side effects
            Read | Write => {
                last_side_effect = Some(i);
                break;
            },
            JumpBackwardUnlessZero => jump_stack.push_back(i),
            JumpForwardIfZero => {jump_stack.pop_back().unwrap();},
        }
    }

    if let Some(last_side_effect) = last_side_effect {
        if jump_stack.is_empty() {
            instructions.truncate(last_side_effect + 1);
        }
        else {
            instructions.truncate(jump_stack.front().unwrap() + 1);
        }
    }
    else {
        // no side effects, so the instructions may as well not be there at all
        //TODO: Determine if we should actually be doing this?
        //TODO: Maybe it would be better to just leave them as is if there are no side effects
        //TODO: Even if that means that nothing can possibly be read or written at runtime
        instructions.clear();
    }
}

fn remove_opposites(instructions: &mut Instructions) {
    let mut i = 1;
    while i < instructions.len() {
        let prev = instructions[i - 1];
        let current = instructions[i];

        // Cancel out opposites
        match (prev, current) {
            (Left, Right) | (Right, Left) |
            (Increment, Decrement) | (Decrement, Increment) => {
                instructions.remove(i);
                instructions.remove(i - 1);
                i -= 1;
            },

            // Otherwise just move on
            _ => i += 1,
        }

        // Sometimes this can result in getting back to the start
        if i == 0 {
            i = 1;
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::{Bitv, VecMap};
use std::collections::bitv::Iter;
use std::mem::swap;

use alphabets::Alphabet;

/// Construct suffix array for given text.
///
/// # Arguments
///
/// * `text` - the text
///
/// ```rust
/// use bio::data_structures::suffix_array;
/// let text = b"GCCTTAACATTATTACGCCTA$";
///
/// ```
pub fn construct(text: &[u8]) -> Vec<usize> {
    let n = text.len();
    let ranks = Alphabet::new(text).get_ranks();

    let transformed_text: Vec<usize> = text.iter()
        .map(|&c| *ranks.get(&(c as usize)).unwrap())
        .collect();
    if transformed_text[n-1] != 0 {
        panic!("Expecting extra sentinel character being lexicographically smallest at the end of the text.");
    }

    let mut sais = SAIS::new(n);
    sais.construct(transformed_text.as_slice());

    sais.pos
}

struct SAIS {
    pos: Vec<usize>,
    lms_pos: Vec<usize>,
    bucket_sizes: VecMap<usize>,
    bucket_start: Vec<usize>,
    bucket_end: Vec<usize>
}

impl SAIS {
    fn new(n: usize) -> Self {
        SAIS {
            pos: Vec::with_capacity(n),
            lms_pos: Vec::with_capacity(n),
            bucket_sizes: VecMap::new(),
            bucket_start: Vec::with_capacity(n),
            bucket_end: Vec::with_capacity(n)
        }
    }

    fn init_bucket_start(&mut self, text: &[usize]) {
        self.bucket_sizes.clear();
        self.bucket_start.clear();

        for &c in text.iter() {
            if !self.bucket_sizes.contains_key(&c) {
                self.bucket_sizes.insert(c, 0);
            }
            *(self.bucket_sizes.get_mut(&c).unwrap()) += 1;
        }

        let mut sum = 0;
        for &size in self.bucket_sizes.values() {
            self.bucket_start.push(sum);
            sum += size;
        }
    }

    fn init_bucket_end(&mut self, text: &[usize]) {
        self.bucket_end.clear();
        for &r in self.bucket_start[1..].iter() {
            self.bucket_end.push(r);
        }
        self.bucket_end.push(text.len());
    }

    fn construct(&mut self, text: &[usize]) {
        let pos_types = PosTypes::new(text);
        self.calc_lms_pos(text, &pos_types);
        self.calc_pos(text, &pos_types);
    }

    /// Step 1 of the SAIS algorithm.
    fn calc_lms_pos(&mut self, text: &[usize], pos_types: &PosTypes) {
        let n = text.len();

        // collect LMS positions
        self.lms_pos.clear();
        for r in (0..n) {
            if pos_types.is_lms_pos(r) {
                self.lms_pos.push(r);
            }
        }

        // sort LMS substrings by applying step 2 with unsorted LMS positions
        self.calc_pos(text, pos_types);

        // obtain pre-sorted LMS positions from sorted LMS substrings
        self.lms_pos.clear();
        for &p in self.pos.iter() {
            if pos_types.is_lms_pos(p) {
                self.lms_pos.push(p);
            }
        }
        let lms_substring_count = self.lms_pos.len() - 1;

        // if less than 2 LMS substrings are present, no further sorting is needed
        if lms_substring_count > 1 {
            // sort LMS suffixes
            let mut reduced_text = Vec::with_capacity(self.lms_pos.len());
            let mut label = 0;
            reduced_text.push(label);
            for ((&l, &m), &r) in self.lms_pos.iter()
                .zip(self.lms_pos[1..].iter())
                .zip(self.lms_pos[2..].iter()) {
                // choose same label if substrings are equal
                if text[l..m+1] != text[m..r+1] {
                    label += 1;
                }
                reduced_text.push(label);
            }
            // if we have less labels than substrings, we have to sort by recursion
            // because two or more substrings are equal
            if label + 1 < lms_substring_count {
                // backup lms_pos
                let lms_pos = self.lms_pos.clone();
                // recurse SA construction for reduced text
                self.construct(reduced_text.as_slice());
                // obtain sorted lms suffixes
                self.lms_pos.clear();
                for &p in self.pos.iter() {
                    self.lms_pos.push(lms_pos[p]);
                }
            }
        }
    }

    /// Step 2 of the SAIS algorithm.
    fn calc_pos(&mut self, text: &[usize], pos_types: &PosTypes) {
        let n = text.len();
        self.pos.clear();

        self.init_bucket_start(text);
        self.init_bucket_end(text);

        // init all positions as unknown (n-1 is max position)
        for _ in text.iter() {
            self.pos.push(n);
        }

        // insert LMS positions to the end of their buckets
        for &p in self.lms_pos.iter().rev() {
            let c = text[p];
            self.pos[self.bucket_end[c]] = p;
            self.bucket_end[c] -= 1;
        }

        // reset bucket ends
        self.init_bucket_end(text);

        // insert L-positions into buckets
        for r in (0..n) {
            let p = self.pos[r];
            if p == n {
                continue;
            }
            let prev = p - 1;
            if pos_types.is_l_pos(prev) {
                let c = text[prev];
                self.pos[self.bucket_start[c]] = prev;
                self.bucket_start[c] += 1;
            }
        }

        // insert S-positions into buckets
        for r in (0..n).rev() {
            let prev = self.pos[r] - 1;
            if pos_types.is_s_pos(prev) {
                let c = text[prev];
                self.pos[self.bucket_end[c]] = prev;
                self.bucket_end[c] -= 1;
            }
        }
    }
}

struct PosTypes {
    pos_types: Bitv,
}

impl PosTypes {
    /// Calculate the text position type.
    /// S-type marks suffixes being lexicographically smaller than their successor,
    /// L-type marks those being larger.
    /// This function fills a Bitv, with 1-bits denoting S-type
    /// and 0-bits denoting L-type.
    ///
    /// # Arguments
    ///
    /// * `text` - the text, ending with a sentinel.
    fn new(text: &[usize]) -> Self {
        let n = text.len();
        let mut pos_types = Bitv::from_elem(n, false);
        pos_types.set(n-1, true);

        for p in (0..n-1).rev() {
            if text[p] == text[p + 1] {
                // if the characters are equal, the next position determines
                // the lexicographical order
                let v = pos_types.get(p + 1).unwrap();
                pos_types.set(p, v);
            }
            else {
                pos_types.set(p, text[p] < text[p + 1]);
            }
        }

        PosTypes { pos_types: pos_types }
    }

    fn is_s_pos(&self, p: usize) -> bool {
        self.pos_types.get(p).unwrap()
    }

    fn is_l_pos(&self, p: usize) -> bool {
        !self.pos_types.get(p).unwrap()
    }

    fn is_lms_pos(&self, p: usize) -> bool {
        self.is_s_pos(p) && self.is_l_pos(p)
    }
}

// ```rust
// use std::collections::Bitv;
// use bio::data_structures::suffix_array::get_pos_types;
// let text = b"GCCTTAACATTATTACGCCTA$";
// let pos_type = get_pos_types(text);
// let mut test = Bitv::from_bytes(&[0b01100110, 0b10010011,  0b01100100]);
// test.truncate(text.len());
// assert_eq!(pos_type, test);
// ```
<commit_msg>Minor.<commit_after>use std::collections::{Bitv, VecMap};
use std::collections::bitv::Iter;
use std::mem::swap;

use alphabets::Alphabet;

/// Construct suffix array for given text.
///
/// # Arguments
///
/// * `text` - the text
///
/// ```rust
/// use bio::data_structures::suffix_array;
/// let text = b"GCCTTAACATTATTACGCCTA$";
/// let pos = suffix_array::construct(text);
/// ```
pub fn construct(text: &[u8]) -> Vec<usize> {
    let n = text.len();
    let ranks = Alphabet::new(text).get_ranks();

    let transformed_text: Vec<usize> = text.iter()
        .map(|&c| *ranks.get(&(c as usize)).unwrap())
        .collect();
    if transformed_text[n-1] != 0 {
        panic!("Expecting extra sentinel character being lexicographically smallest at the end of the text.");
    }

    let mut sais = SAIS::new(n);
    sais.construct(transformed_text.as_slice());

    sais.pos
}

struct SAIS {
    pos: Vec<usize>,
    lms_pos: Vec<usize>,
    bucket_sizes: VecMap<usize>,
    bucket_start: Vec<usize>,
    bucket_end: Vec<usize>
}

impl SAIS {
    fn new(n: usize) -> Self {
        SAIS {
            pos: Vec::with_capacity(n),
            lms_pos: Vec::with_capacity(n),
            bucket_sizes: VecMap::new(),
            bucket_start: Vec::with_capacity(n),
            bucket_end: Vec::with_capacity(n)
        }
    }

    fn init_bucket_start(&mut self, text: &[usize]) {
        self.bucket_sizes.clear();
        self.bucket_start.clear();

        for &c in text.iter() {
            if !self.bucket_sizes.contains_key(&c) {
                self.bucket_sizes.insert(c, 0);
            }
            *(self.bucket_sizes.get_mut(&c).unwrap()) += 1;
        }

        let mut sum = 0;
        for &size in self.bucket_sizes.values() {
            self.bucket_start.push(sum);
            sum += size;
        }
    }

    fn init_bucket_end(&mut self, text: &[usize]) {
        self.bucket_end.clear();
        for &r in self.bucket_start[1..].iter() {
            self.bucket_end.push(r);
        }
        self.bucket_end.push(text.len());
    }

    fn construct(&mut self, text: &[usize]) {
        let pos_types = PosTypes::new(text);
        self.calc_lms_pos(text, &pos_types);
        self.calc_pos(text, &pos_types);
    }

    /// Step 1 of the SAIS algorithm.
    fn calc_lms_pos(&mut self, text: &[usize], pos_types: &PosTypes) {
        let n = text.len();

        // collect LMS positions
        self.lms_pos.clear();
        for r in (0..n) {
            if pos_types.is_lms_pos(r) {
                self.lms_pos.push(r);
            }
        }

        // sort LMS substrings by applying step 2 with unsorted LMS positions
        self.calc_pos(text, pos_types);

        // obtain pre-sorted LMS positions from sorted LMS substrings
        self.lms_pos.clear();
        for &p in self.pos.iter() {
            if pos_types.is_lms_pos(p) {
                self.lms_pos.push(p);
            }
        }
        let lms_substring_count = self.lms_pos.len() - 1;

        // if less than 2 LMS substrings are present, no further sorting is needed
        if lms_substring_count > 1 {
            // sort LMS suffixes
            let mut reduced_text = Vec::with_capacity(self.lms_pos.len());
            let mut label = 0;
            reduced_text.push(label);
            for ((&l, &m), &r) in self.lms_pos.iter()
                .zip(self.lms_pos[1..].iter())
                .zip(self.lms_pos[2..].iter()) {
                // choose same label if substrings are equal
                if text[l..m+1] != text[m..r+1] {
                    label += 1;
                }
                reduced_text.push(label);
            }
            // if we have less labels than substrings, we have to sort by recursion
            // because two or more substrings are equal
            if label + 1 < lms_substring_count {
                // backup lms_pos
                let lms_pos = self.lms_pos.clone();
                // recurse SA construction for reduced text
                self.construct(reduced_text.as_slice());
                // obtain sorted lms suffixes
                self.lms_pos.clear();
                for &p in self.pos.iter() {
                    self.lms_pos.push(lms_pos[p]);
                }
            }
        }
    }

    /// Step 2 of the SAIS algorithm.
    fn calc_pos(&mut self, text: &[usize], pos_types: &PosTypes) {
        let n = text.len();
        self.pos.clear();

        self.init_bucket_start(text);
        self.init_bucket_end(text);

        // init all positions as unknown (n-1 is max position)
        for _ in text.iter() {
            self.pos.push(n);
        }

        // insert LMS positions to the end of their buckets
        for &p in self.lms_pos.iter().rev() {
            let c = text[p];
            self.pos[self.bucket_end[c]] = p;
            self.bucket_end[c] -= 1;
        }

        // reset bucket ends
        self.init_bucket_end(text);

        // insert L-positions into buckets
        for r in (0..n) {
            let p = self.pos[r];
            if p == n {
                continue;
            }
            let prev = p - 1;
            if pos_types.is_l_pos(prev) {
                let c = text[prev];
                self.pos[self.bucket_start[c]] = prev;
                self.bucket_start[c] += 1;
            }
        }

        // insert S-positions into buckets
        for r in (0..n).rev() {
            let prev = self.pos[r] - 1;
            if pos_types.is_s_pos(prev) {
                let c = text[prev];
                self.pos[self.bucket_end[c]] = prev;
                self.bucket_end[c] -= 1;
            }
        }
    }
}

struct PosTypes {
    pos_types: Bitv,
}

impl PosTypes {
    /// Calculate the text position type.
    /// S-type marks suffixes being lexicographically smaller than their successor,
    /// L-type marks those being larger.
    /// This function fills a Bitv, with 1-bits denoting S-type
    /// and 0-bits denoting L-type.
    ///
    /// # Arguments
    ///
    /// * `text` - the text, ending with a sentinel.
    fn new(text: &[usize]) -> Self {
        let n = text.len();
        let mut pos_types = Bitv::from_elem(n, false);
        pos_types.set(n-1, true);

        for p in (0..n-1).rev() {
            if text[p] == text[p + 1] {
                // if the characters are equal, the next position determines
                // the lexicographical order
                let v = pos_types.get(p + 1).unwrap();
                pos_types.set(p, v);
            }
            else {
                pos_types.set(p, text[p] < text[p + 1]);
            }
        }

        PosTypes { pos_types: pos_types }
    }

    fn is_s_pos(&self, p: usize) -> bool {
        self.pos_types.get(p).unwrap()
    }

    fn is_l_pos(&self, p: usize) -> bool {
        !self.pos_types.get(p).unwrap()
    }

    fn is_lms_pos(&self, p: usize) -> bool {
        self.is_s_pos(p) && self.is_l_pos(p)
    }
}

// ```rust
// use std::collections::Bitv;
// use bio::data_structures::suffix_array::get_pos_types;
// let text = b"GCCTTAACATTATTACGCCTA$";
// let pos_type = get_pos_types(text);
// let mut test = Bitv::from_bytes(&[0b01100110, 0b10010011,  0b01100100]);
// test.truncate(text.len());
// assert_eq!(pos_type, test);
// ```
<|endoftext|>