<commit_before>use crate::asm::program_parsers::program;
use crate::asm::Assembler;
use crate::compiler::Compiler;
use crate::repl::command_parser::CommandParser;
use crate::vm::VM;

use nom::types::CompleteStr;
use std::fs;
use std::io;
use std::io::Write;
use std::path::Path;

pub mod command_parser;

const INFO_TAG: char = '\u{29FB}';
const WARN_TAG: char = '\u{26A0}';
const ERROR_TAG: char = '\u{2620}';

#[derive(Debug)]
enum Mode {
    Assembly,
    Highlevel,
}

pub struct REPL {
    assembly_buffer: Vec<String>,
    command_buffer: Vec<String>,
    compiler: Compiler,
    asm: Assembler,
    vm: VM,
    mode: Mode,
}

impl REPL {
    pub fn new() -> REPL {
        REPL {
            assembly_buffer: vec![],
            command_buffer: vec![],
            compiler: Compiler::new(),
            asm: Assembler::new(),
            vm: VM::new(),
            mode: Mode::Highlevel,
        }
    }

    pub fn run(&mut self) {
        println!("{} What will you DO today?", INFO_TAG);

        loop {
            let mut buffer = String::new();

            print!("\u{21AA} ");
            io::stdout().flush().expect("Unable to flush stdout");

            io::stdin()
                .read_line(&mut buffer)
                .expect("Unable to read line from user");
            let buffer = buffer.trim();

            self.command_buffer.push(buffer.to_string());

            if buffer.starts_with(':') {
                self.execute_command(&buffer);
            } else {
                let assembly = match self.mode {
                    // TODO: figure out how to step enough in highlevel or halt if we run for
                    // assembly.
                    Mode::Assembly => vec![buffer.into()],
                    Mode::Highlevel => match self.compiler.compile_expr(buffer) {
                        Ok(compiled) => compiled.to_vec(),
                        Err(e) => {
                            println!("{} Unable to compile input: {}", ERROR_TAG, e);
                            continue;
                        }
                    },
                };
                self.assembly_buffer.append(&mut assembly.clone());
                let assembled = assembly.join("\n");
                let bytecode = match program(CompleteStr(&assembled)) {
                    Ok((_, prog)) => self.asm.process_second(&prog),
                    Err(e) => {
                        println!("{} Unable to parse input: {}", WARN_TAG, e);
                        continue;
                    }
                };
                if let Err(e) = bytecode {
                    println!("{} {}", ERROR_TAG, e);
                    continue;
                }
                self.vm.program.append(&mut bytecode.unwrap());
                for _i in 1..=assembly.len() {
                    if let Err(e) = self.vm.step() {
                        println!("{} {}", ERROR_TAG, e);
                        continue;
                    }
                }
            }
        }
    }

    fn execute_command(&mut self, input: &str) {
        let args = CommandParser::tokenize(input);
        match args[0] {
            ":c" => self.vm = VM::new(),
            ":h" => self.print_help(),
            ":history" => {
                for command in &self.command_buffer {
                    println!("  {}", command);
                }
            }
            ":list" => self.list_program(),
            ":load" => self.load_file(&args[1..]),
            ":m" => self.toggle_mode(),
            ":q" => {
                println!("{} Buh-bye!", INFO_TAG);
                std::process::exit(0);
            }
            ":r" => self.list_registers(),
            ":s" => {
                println!("{} Listing symbols:", INFO_TAG);
                println!("{:#?}", self.asm.symbols);
                println!("{} EOF", INFO_TAG);
            }
            _ => println!("{} Invalid command: '{}'", ERROR_TAG, input),
        }
    }

    fn print_help(&self) {
        println!(":c clears the VM");
        println!(":h prints this help");
        println!(":history shows the command buffer");
        println!(":list lists the current program");
        println!(":load loads a program from the given path");
        println!(":m toggles the mode between assembly and high-level language");
        println!(":q quits");
        println!(":r prints the registers");
        println!(":s lists the known symbols");
    }

    fn list_program(&self) {
        println!("{} Listing assembly:", INFO_TAG);
        for line in &self.assembly_buffer {
            println!("  {}", line);
        }
        println!("{} EOF", INFO_TAG);
        println!("{} Listing instructions:", INFO_TAG);
        for instr in self.vm.program.chunks(4) {
            println!("  {:?}", instr);
        }
        println!("{} EOF", INFO_TAG);
    }

    fn list_registers(&self) {
        println!("{} Listing integer registers:", INFO_TAG);
        for reg in self.vm.iregisters.chunks(4) {
            println!("  {}\t{}\t{}\t{}", reg[0], reg[1], reg[2], reg[3]);
        }
        println!("{} EOF", INFO_TAG);

        println!("{} Listing real registers:", INFO_TAG);
        for reg in self.vm.rregisters.chunks(4) {
            println!("  {}\t{}\t{}\t{}", reg[0], reg[1], reg[2], reg[3]);
        }
        println!("{} EOF", INFO_TAG);
    }

    fn load_file(&mut self, args: &[&str]) {
        // TODO: merge with `main.rs` file reading.
        let mut tmp = String::new();
        if args.is_empty() {
            print!("{} Please enter the path to the file: ", INFO_TAG);
            io::stdout().flush().expect("Failed to flush stdout");
            io::stdin()
                .read_line(&mut tmp)
                .expect("Failed to read line from user");
            tmp = tmp.trim().to_string();
        } else {
            tmp = args[1].to_string();
        }
        let path = Path::new(&tmp);
        println!(
            "{} Loading program from '{}'",
            INFO_TAG,
            path.to_str().unwrap()
        );
        let result = fs::read_to_string(path);
        let contents = match result {
            Ok(contents) => contents,
            Err(e) => {
                println!("{} Error loading that file: {:?}", ERROR_TAG, e);
                return;
            }
        };

        match self.asm.assemble(&contents) {
            Ok(assembled) => {
                println!("{} Sending program to VM", INFO_TAG);
                self.vm = VM::new();
                self.vm.set_bytecode(&assembled).unwrap();
                self.list_program();
                let result = self.vm.run();
                if result.is_err() {
                    println!("{} Runtime error: {}", ERROR_TAG, result.unwrap_err());
                }
            }
            Err(errors) => {
                println!("{} Errors encountered while assembling:", ERROR_TAG);
                for error in errors {
                    println!("  {}", error);
                }
            }
        }
    }

    fn toggle_mode(&mut self) {
        match self.mode {
            Mode::Highlevel => self.mode = Mode::Assembly,
            Mode::Assembly => self.mode = Mode::Highlevel,
        };
        println!("{} Mode is now {:?}", INFO_TAG, self.mode);
    }
}

impl Default for REPL {
    fn default() -> Self {
        REPL::new()
    }
}
<commit_msg>remove obsolete todo<commit_after>use crate::asm::program_parsers::program;
use crate::asm::Assembler;
use crate::compiler::Compiler;
use crate::repl::command_parser::CommandParser;
use crate::vm::VM;

use nom::types::CompleteStr;
use std::fs;
use std::io;
use std::io::Write;
use std::path::Path;

pub mod command_parser;

const INFO_TAG: char = '\u{29FB}';
const WARN_TAG: char = '\u{26A0}';
const ERROR_TAG: char = '\u{2620}';

#[derive(Debug)]
enum Mode {
    Assembly,
    Highlevel,
}

pub struct REPL {
    assembly_buffer: Vec<String>,
    command_buffer: Vec<String>,
    compiler: Compiler,
    asm: Assembler,
    vm: VM,
    mode: Mode,
}

impl REPL {
    pub fn new() -> REPL {
        REPL {
            assembly_buffer: vec![],
            command_buffer: vec![],
            compiler: Compiler::new(),
            asm: Assembler::new(),
            vm: VM::new(),
            mode: Mode::Highlevel,
        }
    }

    pub fn run(&mut self) {
        println!("{} What will you DO today?", INFO_TAG);

        loop {
            let mut buffer = String::new();

            print!("\u{21AA} ");
            io::stdout().flush().expect("Unable to flush stdout");

            io::stdin()
                .read_line(&mut buffer)
                .expect("Unable to read line from user");
            let buffer = buffer.trim();

            self.command_buffer.push(buffer.to_string());

            if buffer.starts_with(':') {
                self.execute_command(&buffer);
            } else {
                let assembly = match self.mode {
                    Mode::Assembly => vec![buffer.into()],
                    Mode::Highlevel => match self.compiler.compile_expr(buffer) {
                        Ok(compiled) => compiled.to_vec(),
                        Err(e) => {
                            println!("{} Unable to compile input: {}", ERROR_TAG, e);
                            continue;
                        }
                    },
                };
                self.assembly_buffer.append(&mut assembly.clone());
                let assembled = assembly.join("\n");
                let bytecode = match program(CompleteStr(&assembled)) {
                    Ok((_, prog)) => self.asm.process_second(&prog),
                    Err(e) => {
                        println!("{} Unable to parse input: {}", WARN_TAG, e);
                        continue;
                    }
                };
                if let Err(e) = bytecode {
                    println!("{} {}", ERROR_TAG, e);
                    continue;
                }
                self.vm.program.append(&mut bytecode.unwrap());
                for _i in 1..=assembly.len() {
                    if let Err(e) = self.vm.step() {
                        println!("{} {}", ERROR_TAG, e);
                        continue;
                    }
                }
            }
        }
    }

    fn execute_command(&mut self, input: &str) {
        let args = CommandParser::tokenize(input);
        match args[0] {
            ":c" => self.vm = VM::new(),
            ":h" => self.print_help(),
            ":history" => {
                for command in &self.command_buffer {
                    println!("  {}", command);
                }
            }
            ":list" => self.list_program(),
            ":load" => self.load_file(&args[1..]),
            ":m" => self.toggle_mode(),
            ":q" => {
                println!("{} Buh-bye!", INFO_TAG);
                std::process::exit(0);
            }
            ":r" => self.list_registers(),
            ":s" => {
                println!("{} Listing symbols:", INFO_TAG);
                println!("{:#?}", self.asm.symbols);
                println!("{} EOF", INFO_TAG);
            }
            _ => println!("{} Invalid command: '{}'", ERROR_TAG, input),
        }
    }

    fn print_help(&self) {
        println!(":c clears the VM");
        println!(":h prints this help");
        println!(":history shows the command buffer");
        println!(":list lists the current program");
        println!(":load loads a program from the given path");
        println!(":m toggles the mode between assembly and high-level language");
        println!(":q quits");
        println!(":r prints the registers");
        println!(":s lists the known symbols");
    }

    fn list_program(&self) {
        println!("{} Listing assembly:", INFO_TAG);
        for line in &self.assembly_buffer {
            println!("  {}", line);
        }
        println!("{} EOF", INFO_TAG);
        println!("{} Listing instructions:", INFO_TAG);
        for instr in self.vm.program.chunks(4) {
            println!("  {:?}", instr);
        }
        println!("{} EOF", INFO_TAG);
    }

    fn list_registers(&self) {
        println!("{} Listing integer registers:", INFO_TAG);
        for reg in self.vm.iregisters.chunks(4) {
            println!("  {}\t{}\t{}\t{}", reg[0], reg[1], reg[2], reg[3]);
        }
        println!("{} EOF", INFO_TAG);

        println!("{} Listing real registers:", INFO_TAG);
        for reg in self.vm.rregisters.chunks(4) {
            println!("  {}\t{}\t{}\t{}", reg[0], reg[1], reg[2], reg[3]);
        }
        println!("{} EOF", INFO_TAG);
    }

    fn load_file(&mut self, args: &[&str]) {
        // TODO: merge with `main.rs` file reading.
        let mut tmp = String::new();
        if args.is_empty() {
            print!("{} Please enter the path to the file: ", INFO_TAG);
            io::stdout().flush().expect("Failed to flush stdout");
            io::stdin()
                .read_line(&mut tmp)
                .expect("Failed to read line from user");
            tmp = tmp.trim().to_string();
        } else {
            tmp = args[1].to_string();
        }
        let path = Path::new(&tmp);
        println!(
            "{} Loading program from '{}'",
            INFO_TAG,
            path.to_str().unwrap()
        );
        let result = fs::read_to_string(path);
        let contents = match result {
            Ok(contents) => contents,
            Err(e) => {
                println!("{} Error loading that file: {:?}", ERROR_TAG, e);
                return;
            }
        };

        match self.asm.assemble(&contents) {
            Ok(assembled) => {
                println!("{} Sending program to VM", INFO_TAG);
                self.vm = VM::new();
                self.vm.set_bytecode(&assembled).unwrap();
                self.list_program();
                let result = self.vm.run();
                if result.is_err() {
                    println!("{} Runtime error: {}", ERROR_TAG, result.unwrap_err());
                }
            }
            Err(errors) => {
                println!("{} Errors encountered while assembling:", ERROR_TAG);
                for error in errors {
                    println!("  {}", error);
                }
            }
        }
    }

    fn toggle_mode(&mut self) {
        match self.mode {
            Mode::Highlevel => self.mode = Mode::Assembly,
            Mode::Assembly => self.mode = Mode::Highlevel,
        };
        println!("{} Mode is now {:?}", INFO_TAG, self.mode);
    }
}

impl Default for REPL {
    fn default() -> Self {
        REPL::new()
    }
}
<|endoftext|>


<commit_before>use crate::publish_rate_limit::PublishRateLimit;
use crate::{env, env_optional, uploaders::Uploader, Env, Replica};

#[derive(Debug)]
pub struct Config {
    pub uploader: Uploader,
    pub session_key: String,
    pub gh_client_id: String,
    pub gh_client_secret: String,
    pub gh_base_url: String,
    pub db_primary_config: DbPoolConfig,
    pub db_replica_config: Option<DbPoolConfig>,
    pub env: Env,
    pub max_upload_size: u64,
    pub max_unpack_size: u64,
    pub mirror: Replica,
    pub api_protocol: String,
    pub publish_rate_limit: PublishRateLimit,
    pub blocked_traffic: Vec<(String, Vec<String>)>,
    pub max_allowed_page_offset: u32,
    pub page_offset_ua_blocklist: Vec<String>,
    pub domain_name: String,
    pub allowed_origins: Vec<String>,
    pub downloads_persist_interval_ms: usize,
    pub ownership_invitations_expiration_days: u64,
    pub metrics_authorization_token: Option<String>,
    pub use_test_database_pool: bool,
}

#[derive(Debug)]
pub struct DbPoolConfig {
    pub url: String,
    pub read_only_mode: bool,
}

impl Default for Config {
    /// Returns a default value for the application's config
    ///
    /// Sets the following default values:
    ///
    /// - `Config::max_upload_size`: 10MiB
    /// - `Config::api_protocol`: `https`
    /// - `Config::ownership_invitations_expiration_days`: 30
    ///
    /// Pulls values from the following environment variables:
    ///
    /// - `MIRROR`: Is this instance of cargo_registry a mirror of crates.io.
    /// - `HEROKU`: Is this instance of cargo_registry currently running on Heroku.
    /// - `S3_BUCKET`: The S3 bucket used to store crate files. If not present during development,
    ///    cargo_registry will fall back to a local uploader.
    /// - `S3_REGION`: The region in which the bucket was created. Optional if US standard.
    /// - `S3_ACCESS_KEY`: The access key to interact with S3. Optional if running a mirror.
    /// - `S3_SECRET_KEY`: The secret key to interact with S3. Optional if running a mirror.
    /// - `SESSION_KEY`: The key used to sign and encrypt session cookies.
    /// - `GH_CLIENT_ID`: The client ID of the associated GitHub application.
    /// - `GH_CLIENT_SECRET`: The client secret of the associated GitHub application.
    /// - `DATABASE_URL`: The URL of the postgres database to use.
    /// - `READ_ONLY_REPLICA_URL`: The URL of an optional postgres read-only replica database.
    /// - `BLOCKED_TRAFFIC`: A list of headers and environment variables to use for blocking
    ///   traffic. See the `block_traffic` module for more documentation.
    /// - `DOWNLOADS_PERSIST_INTERVAL_MS`: how frequent to persist download counts (in ms).
    /// - `METRICS_AUTHORIZATION_TOKEN`: authorization token needed to query metrics. If missing,
    ///   querying metrics will be completely disabled.
    /// - `DB_OFFLINE`: If set to `leader` then use the read-only follower as if it was the leader.
    ///   If set to `follower` then act as if `READ_ONLY_REPLICA_URL` was unset.
    /// - `READ_ONLY_MODE`: If defined (even as empty) then force all connections to be read-only.
    /// - `WEB_MAX_ALLOWED_PAGE_OFFSET`: Page offsets larger than this value are rejected. Defaults
    ///   to 200.
    /// - `WEB_PAGE_OFFSET_UA_BLOCKLIST`: A comma seperated list of user-agent substrings that will
    ///   be blocked if `WEB_MAX_ALLOWED_PAGE_OFFSET` is exceeded.
    ///
    /// # Panics
    ///
    /// This function panics if `DB_OFFLINE=leader` but `READ_ONLY_REPLICA_URL` is unset.
    fn default() -> Config {
        let api_protocol = String::from("https");
        let mirror = if dotenv::var("MIRROR").is_ok() {
            Replica::ReadOnlyMirror
        } else {
            Replica::Primary
        };
        let heroku = dotenv::var("HEROKU").is_ok();
        let cargo_env = if heroku {
            Env::Production
        } else {
            Env::Development
        };

        let leader_url = env("DATABASE_URL");
        let follower_url = dotenv::var("READ_ONLY_REPLICA_URL").ok();
        let read_only_mode = dotenv::var("READ_ONLY_MODE").is_ok();
        let (db_primary_config, db_replica_config) = match dotenv::var("DB_OFFLINE").as_deref() {
            // The actual leader is down, use the follower in read-only mode as the primary and
            // don't configure a replica.
            Ok("leader") => (
                DbPoolConfig {
                    url: follower_url
                        .expect("Must set `READ_ONLY_REPLICA_URL` when using `DB_OFFLINE=leader`."),
                    read_only_mode: true,
                },
                None,
            ),
            // The follower is down, don't configure the replica.
            Ok("follower") => (
                DbPoolConfig {
                    url: leader_url,
                    read_only_mode,
                },
                None,
            ),
            _ => (
                DbPoolConfig {
                    url: leader_url,
                    read_only_mode,
                },
                follower_url.map(|url| DbPoolConfig {
                    url,
                    // Always enable read-only mode for the follower. In staging, we attach the
                    // same leader database to both environment variables and this ensures the
                    // connection is opened read-only even when attached to a writeable database.
                    read_only_mode: true,
                }),
            ),
        };

        let uploader = match (cargo_env, mirror) {
            (Env::Production, Replica::Primary) => {
                // `env` panics if these vars are not set, and in production for a primary instance,
                // that's what we want since we don't want to be able to start the server if the
                // server doesn't know where to upload crates.
                Uploader::S3 {
                    bucket: s3::Bucket::new(
                        env("S3_BUCKET"),
                        dotenv::var("S3_REGION").ok(),
                        env("S3_ACCESS_KEY"),
                        env("S3_SECRET_KEY"),
                        &api_protocol,
                    ),
                    cdn: dotenv::var("S3_CDN").ok(),
                }
            }
            (Env::Production, Replica::ReadOnlyMirror) => {
                // Read-only mirrors don't need access key or secret key since by definition,
                // they'll only need to read from a bucket, not upload.
                //
                // Read-only mirrors might have access key or secret key, so use them if those
                // environment variables are set.
                //
                // Read-only mirrors definitely need bucket though, so that they know where
                // to serve crate files from.
                Uploader::S3 {
                    bucket: s3::Bucket::new(
                        env("S3_BUCKET"),
                        dotenv::var("S3_REGION").ok(),
                        dotenv::var("S3_ACCESS_KEY").unwrap_or_default(),
                        dotenv::var("S3_SECRET_KEY").unwrap_or_default(),
                        &api_protocol,
                    ),
                    cdn: dotenv::var("S3_CDN").ok(),
                }
            }
            // In Development mode, either running as a primary instance or a read-only mirror
            _ => {
                if dotenv::var("S3_BUCKET").is_ok() {
                    // If we've set the `S3_BUCKET` variable to any value, use all of the values
                    // for the related S3 environment variables and configure the app to upload to
                    // and read from S3 like production does. All values except for bucket are
                    // optional, like production read-only mirrors.
                    println!("Using S3 uploader");
                    Uploader::S3 {
                        bucket: s3::Bucket::new(
                            env("S3_BUCKET"),
                            dotenv::var("S3_REGION").ok(),
                            dotenv::var("S3_ACCESS_KEY").unwrap_or_default(),
                            dotenv::var("S3_SECRET_KEY").unwrap_or_default(),
                            &api_protocol,
                        ),
                        cdn: dotenv::var("S3_CDN").ok(),
                    }
                } else {
                    // If we don't set the `S3_BUCKET` variable, we'll use a development-only
                    // uploader that makes it possible to run and publish to a locally-running
                    // crates.io instance without needing to set up an account and a bucket in S3.
                    println!(
                        "Using local uploader, crate files will be in the local_uploads directory"
                    );
                    Uploader::Local
                }
            }
        };
        let allowed_origins = env("WEB_ALLOWED_ORIGINS")
            .split(',')
            .map(ToString::to_string)
            .collect();
        Config {
            uploader,
            session_key: env("SESSION_KEY"),
            gh_client_id: env("GH_CLIENT_ID"),
            gh_client_secret: env("GH_CLIENT_SECRET"),
            gh_base_url: "https://api.github.com".to_string(),
            db_primary_config,
            db_replica_config,
            env: cargo_env,
            max_upload_size: 10 * 1024 * 1024, // 10 MB default file upload size limit
            max_unpack_size: 512 * 1024 * 1024, // 512 MB max when decompressed
            mirror,
            api_protocol,
            publish_rate_limit: Default::default(),
            blocked_traffic: blocked_traffic(),
            max_allowed_page_offset: env_optional("WEB_MAX_ALLOWED_PAGE_OFFSET").unwrap_or(200),
            page_offset_ua_blocklist: env_optional::<String>("WEB_PAGE_OFFSET_UA_BLOCKLIST")
                .unwrap_or_default()
                .split(',')
                .map(String::from)
                .collect(),
            domain_name: domain_name(),
            allowed_origins,
            downloads_persist_interval_ms: dotenv::var("DOWNLOADS_PERSIST_INTERVAL_MS")
                .map(|interval| {
                    interval
                        .parse()
                        .expect("invalid DOWNLOADS_PERSIST_INTERVAL_MS")
                })
                .unwrap_or(60_000), // 1 minute
            ownership_invitations_expiration_days: 30,
            metrics_authorization_token: dotenv::var("METRICS_AUTHORIZATION_TOKEN").ok(),
            use_test_database_pool: false,
        }
    }
}

pub(crate) fn domain_name() -> String {
    dotenv::var("DOMAIN_NAME").unwrap_or_else(|_| "crates.io".into())
}

fn blocked_traffic() -> Vec<(String, Vec<String>)> {
    let pattern_list = dotenv::var("BLOCKED_TRAFFIC").unwrap_or_default();
    parse_traffic_patterns(&pattern_list)
        .map(|(header, value_env_var)| {
            let value_list = dotenv::var(value_env_var).unwrap_or_default();
            let values = value_list.split(',').map(String::from).collect();
            (header.into(), values)
        })
        .collect()
}

fn parse_traffic_patterns(patterns: &str) -> impl Iterator<Item = (&str, &str)> {
    patterns.split_terminator(',').map(|pattern| {
        if let Some(idx) = pattern.find('=') {
            (&pattern[..idx], &pattern[(idx + 1)..])
        } else {
            panic!(
                "BLOCKED_TRAFFIC must be in the form HEADER=VALUE_ENV_VAR, \
                 got invalid pattern {}",
                pattern
            )
        }
    })
}

#[test]
fn parse_traffic_patterns_splits_on_comma_and_looks_for_equal_sign() {
    let pattern_string_1 = "Foo=BAR,Bar=BAZ";
    let pattern_string_2 = "Baz=QUX";
    let pattern_string_3 = "";

    let patterns_1 = parse_traffic_patterns(pattern_string_1).collect::<Vec<_>>();
    assert_eq!(vec![("Foo", "BAR"), ("Bar", "BAZ")], patterns_1);

    let patterns_2 = parse_traffic_patterns(pattern_string_2).collect::<Vec<_>>();
    assert_eq!(vec![("Baz", "QUX")], patterns_2);

    assert_none!(parse_traffic_patterns(pattern_string_3).next());
}
<commit_msg>Disable page offset blocking if UA blocklist is unset<commit_after>use crate::publish_rate_limit::PublishRateLimit;
use crate::{env, env_optional, uploaders::Uploader, Env, Replica};

#[derive(Debug)]
pub struct Config {
    pub uploader: Uploader,
    pub session_key: String,
    pub gh_client_id: String,
    pub gh_client_secret: String,
    pub gh_base_url: String,
    pub db_primary_config: DbPoolConfig,
    pub db_replica_config: Option<DbPoolConfig>,
    pub env: Env,
    pub max_upload_size: u64,
    pub max_unpack_size: u64,
    pub mirror: Replica,
    pub api_protocol: String,
    pub publish_rate_limit: PublishRateLimit,
    pub blocked_traffic: Vec<(String, Vec<String>)>,
    pub max_allowed_page_offset: u32,
    pub page_offset_ua_blocklist: Vec<String>,
    pub domain_name: String,
    pub allowed_origins: Vec<String>,
    pub downloads_persist_interval_ms: usize,
    pub ownership_invitations_expiration_days: u64,
    pub metrics_authorization_token: Option<String>,
    pub use_test_database_pool: bool,
}

#[derive(Debug)]
pub struct DbPoolConfig {
    pub url: String,
    pub read_only_mode: bool,
}

impl Default for Config {
    /// Returns a default value for the application's config
    ///
    /// Sets the following default values:
    ///
    /// - `Config::max_upload_size`: 10MiB
    /// - `Config::api_protocol`: `https`
    /// - `Config::ownership_invitations_expiration_days`: 30
    ///
    /// Pulls values from the following environment variables:
    ///
    /// - `MIRROR`: Is this instance of cargo_registry a mirror of crates.io.
    /// - `HEROKU`: Is this instance of cargo_registry currently running on Heroku.
    /// - `S3_BUCKET`: The S3 bucket used to store crate files. If not present during development,
    ///    cargo_registry will fall back to a local uploader.
    /// - `S3_REGION`: The region in which the bucket was created. Optional if US standard.
    /// - `S3_ACCESS_KEY`: The access key to interact with S3. Optional if running a mirror.
    /// - `S3_SECRET_KEY`: The secret key to interact with S3. Optional if running a mirror.
    /// - `SESSION_KEY`: The key used to sign and encrypt session cookies.
    /// - `GH_CLIENT_ID`: The client ID of the associated GitHub application.
    /// - `GH_CLIENT_SECRET`: The client secret of the associated GitHub application.
    /// - `DATABASE_URL`: The URL of the postgres database to use.
    /// - `READ_ONLY_REPLICA_URL`: The URL of an optional postgres read-only replica database.
    /// - `BLOCKED_TRAFFIC`: A list of headers and environment variables to use for blocking
    ///   traffic. See the `block_traffic` module for more documentation.
    /// - `DOWNLOADS_PERSIST_INTERVAL_MS`: how frequent to persist download counts (in ms).
    /// - `METRICS_AUTHORIZATION_TOKEN`: authorization token needed to query metrics. If missing,
    ///   querying metrics will be completely disabled.
    /// - `DB_OFFLINE`: If set to `leader` then use the read-only follower as if it was the leader.
    ///   If set to `follower` then act as if `READ_ONLY_REPLICA_URL` was unset.
    /// - `READ_ONLY_MODE`: If defined (even as empty) then force all connections to be read-only.
    /// - `WEB_MAX_ALLOWED_PAGE_OFFSET`: Page offsets larger than this value are rejected. Defaults
    ///   to 200.
    /// - `WEB_PAGE_OFFSET_UA_BLOCKLIST`: A comma seperated list of user-agent substrings that will
    ///   be blocked if `WEB_MAX_ALLOWED_PAGE_OFFSET` is exceeded. Including an empty string in the
    ///   list will block *all* user-agents exceeding the offset. If not set or empty, no blocking
    ///   will occur.
    ///
    /// # Panics
    ///
    /// This function panics if `DB_OFFLINE=leader` but `READ_ONLY_REPLICA_URL` is unset.
    fn default() -> Config {
        let api_protocol = String::from("https");
        let mirror = if dotenv::var("MIRROR").is_ok() {
            Replica::ReadOnlyMirror
        } else {
            Replica::Primary
        };
        let heroku = dotenv::var("HEROKU").is_ok();
        let cargo_env = if heroku {
            Env::Production
        } else {
            Env::Development
        };

        let leader_url = env("DATABASE_URL");
        let follower_url = dotenv::var("READ_ONLY_REPLICA_URL").ok();
        let read_only_mode = dotenv::var("READ_ONLY_MODE").is_ok();
        let (db_primary_config, db_replica_config) = match dotenv::var("DB_OFFLINE").as_deref() {
            // The actual leader is down, use the follower in read-only mode as the primary and
            // don't configure a replica.
            Ok("leader") => (
                DbPoolConfig {
                    url: follower_url
                        .expect("Must set `READ_ONLY_REPLICA_URL` when using `DB_OFFLINE=leader`."),
                    read_only_mode: true,
                },
                None,
            ),
            // The follower is down, don't configure the replica.
            Ok("follower") => (
                DbPoolConfig {
                    url: leader_url,
                    read_only_mode,
                },
                None,
            ),
            _ => (
                DbPoolConfig {
                    url: leader_url,
                    read_only_mode,
                },
                follower_url.map(|url| DbPoolConfig {
                    url,
                    // Always enable read-only mode for the follower. In staging, we attach the
                    // same leader database to both environment variables and this ensures the
                    // connection is opened read-only even when attached to a writeable database.
                    read_only_mode: true,
                }),
            ),
        };

        let uploader = match (cargo_env, mirror) {
            (Env::Production, Replica::Primary) => {
                // `env` panics if these vars are not set, and in production for a primary instance,
                // that's what we want since we don't want to be able to start the server if the
                // server doesn't know where to upload crates.
                Uploader::S3 {
                    bucket: s3::Bucket::new(
                        env("S3_BUCKET"),
                        dotenv::var("S3_REGION").ok(),
                        env("S3_ACCESS_KEY"),
                        env("S3_SECRET_KEY"),
                        &api_protocol,
                    ),
                    cdn: dotenv::var("S3_CDN").ok(),
                }
            }
            (Env::Production, Replica::ReadOnlyMirror) => {
                // Read-only mirrors don't need access key or secret key since by definition,
                // they'll only need to read from a bucket, not upload.
                //
                // Read-only mirrors might have access key or secret key, so use them if those
                // environment variables are set.
                //
                // Read-only mirrors definitely need bucket though, so that they know where
                // to serve crate files from.
                Uploader::S3 {
                    bucket: s3::Bucket::new(
                        env("S3_BUCKET"),
                        dotenv::var("S3_REGION").ok(),
                        dotenv::var("S3_ACCESS_KEY").unwrap_or_default(),
                        dotenv::var("S3_SECRET_KEY").unwrap_or_default(),
                        &api_protocol,
                    ),
                    cdn: dotenv::var("S3_CDN").ok(),
                }
            }
            // In Development mode, either running as a primary instance or a read-only mirror
            _ => {
                if dotenv::var("S3_BUCKET").is_ok() {
                    // If we've set the `S3_BUCKET` variable to any value, use all of the values
                    // for the related S3 environment variables and configure the app to upload to
                    // and read from S3 like production does. All values except for bucket are
                    // optional, like production read-only mirrors.
                    println!("Using S3 uploader");
                    Uploader::S3 {
                        bucket: s3::Bucket::new(
                            env("S3_BUCKET"),
                            dotenv::var("S3_REGION").ok(),
                            dotenv::var("S3_ACCESS_KEY").unwrap_or_default(),
                            dotenv::var("S3_SECRET_KEY").unwrap_or_default(),
                            &api_protocol,
                        ),
                        cdn: dotenv::var("S3_CDN").ok(),
                    }
                } else {
                    // If we don't set the `S3_BUCKET` variable, we'll use a development-only
                    // uploader that makes it possible to run and publish to a locally-running
                    // crates.io instance without needing to set up an account and a bucket in S3.
                    println!(
                        "Using local uploader, crate files will be in the local_uploads directory"
                    );
                    Uploader::Local
                }
            }
        };
        let allowed_origins = env("WEB_ALLOWED_ORIGINS")
            .split(',')
            .map(ToString::to_string)
            .collect();
        let page_offset_ua_blocklist = match env_optional::<String>("WEB_PAGE_OFFSET_UA_BLOCKLIST")
        {
            None => vec![],
            Some(s) if s.is_empty() => vec![],
            Some(s) => s.split(',').map(String::from).collect(),
        };
        Config {
            uploader,
            session_key: env("SESSION_KEY"),
            gh_client_id: env("GH_CLIENT_ID"),
            gh_client_secret: env("GH_CLIENT_SECRET"),
            gh_base_url: "https://api.github.com".to_string(),
            db_primary_config,
            db_replica_config,
            env: cargo_env,
            max_upload_size: 10 * 1024 * 1024, // 10 MB default file upload size limit
            max_unpack_size: 512 * 1024 * 1024, // 512 MB max when decompressed
            mirror,
            api_protocol,
            publish_rate_limit: Default::default(),
            blocked_traffic: blocked_traffic(),
            max_allowed_page_offset: env_optional("WEB_MAX_ALLOWED_PAGE_OFFSET").unwrap_or(200),
            page_offset_ua_blocklist,
            domain_name: domain_name(),
            allowed_origins,
            downloads_persist_interval_ms: dotenv::var("DOWNLOADS_PERSIST_INTERVAL_MS")
                .map(|interval| {
                    interval
                        .parse()
                        .expect("invalid DOWNLOADS_PERSIST_INTERVAL_MS")
                })
                .unwrap_or(60_000), // 1 minute
            ownership_invitations_expiration_days: 30,
            metrics_authorization_token: dotenv::var("METRICS_AUTHORIZATION_TOKEN").ok(),
            use_test_database_pool: false,
        }
    }
}

pub(crate) fn domain_name() -> String {
    dotenv::var("DOMAIN_NAME").unwrap_or_else(|_| "crates.io".into())
}

fn blocked_traffic() -> Vec<(String, Vec<String>)> {
    let pattern_list = dotenv::var("BLOCKED_TRAFFIC").unwrap_or_default();
    parse_traffic_patterns(&pattern_list)
        .map(|(header, value_env_var)| {
            let value_list = dotenv::var(value_env_var).unwrap_or_default();
            let values = value_list.split(',').map(String::from).collect();
            (header.into(), values)
        })
        .collect()
}

fn parse_traffic_patterns(patterns: &str) -> impl Iterator<Item = (&str, &str)> {
    patterns.split_terminator(',').map(|pattern| {
        if let Some(idx) = pattern.find('=') {
            (&pattern[..idx], &pattern[(idx + 1)..])
        } else {
            panic!(
                "BLOCKED_TRAFFIC must be in the form HEADER=VALUE_ENV_VAR, \
                 got invalid pattern {}",
                pattern
            )
        }
    })
}

#[test]
fn parse_traffic_patterns_splits_on_comma_and_looks_for_equal_sign() {
    let pattern_string_1 = "Foo=BAR,Bar=BAZ";
    let pattern_string_2 = "Baz=QUX";
    let pattern_string_3 = "";

    let patterns_1 = parse_traffic_patterns(pattern_string_1).collect::<Vec<_>>();
    assert_eq!(vec![("Foo", "BAR"), ("Bar", "BAZ")], patterns_1);

    let patterns_2 = parse_traffic_patterns(pattern_string_2).collect::<Vec<_>>();
    assert_eq!(vec![("Baz", "QUX")], patterns_2);

    assert_none!(parse_traffic_patterns(pattern_string_3).next());
}
<|endoftext|>


<commit_before>//! Customize line editor
use std::default::Default;

/// User preferences
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Config {
    /// Maximum number of entries in History.
    max_history_size: usize,
    history_duplicates: HistoryDuplicates,
    history_ignore_space: bool,
    completion_type: CompletionType,
    /// When listing completion alternatives, only display
    /// one screen of possibilities at a time.
    completion_prompt_limit: usize,
    /// Duration (milliseconds) Rustyline will wait for a character when
    /// reading an ambiguous key sequence.
    keyseq_timeout: i32,
    /// Emacs or Vi mode
    edit_mode: EditMode,
    /// If true, each nonblank line returned by `readline` will be
    /// automatically added to the history.
    auto_add_history: bool,
}

impl Config {
    pub fn builder() -> Builder {
        Builder::new()
    }

    /// Tell the maximum length (i.e. number of entries) for the history.
    pub fn max_history_size(&self) -> usize {
        self.max_history_size
    }

    /// Tell if lines which match the previous history entry are saved or not
    /// in the history list.
    ///
    /// By default, they are ignored.
    pub fn history_duplicates(&self) -> HistoryDuplicates {
        self.history_duplicates
    }

    /// Tell if lines which begin with a space character are saved or not in
    /// the history list.
    ///
    /// By default, they are saved.
    pub fn history_ignore_space(&self) -> bool {
        self.history_ignore_space
    }

    pub fn completion_type(&self) -> CompletionType {
        self.completion_type
    }

    pub fn completion_prompt_limit(&self) -> usize {
        self.completion_prompt_limit
    }

    pub fn keyseq_timeout(&self) -> i32 {
        self.keyseq_timeout
    }

    pub fn edit_mode(&self) -> EditMode {
        self.edit_mode
    }

    /// Tell if lines are automatically added to the history.
    ///
    /// By default, they are not.
    pub fn auto_add_history(&self) -> bool {
        self.auto_add_history
    }
}

impl Default for Config {
    fn default() -> Config {
        Config {
            max_history_size: 100,
            history_duplicates: HistoryDuplicates::IgnoreConsecutive,
            history_ignore_space: false,
            completion_type: CompletionType::Circular, // TODO Validate
            completion_prompt_limit: 100,
            keyseq_timeout: 500,
            edit_mode: EditMode::Emacs,
            auto_add_history: false,
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum HistoryDuplicates {
    AlwaysAdd,
    /// a line will not be added to the history if it matches the previous entry
    IgnoreConsecutive,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CompletionType {
    /// Complete the next full match (like in Vim by default)
    Circular,
    /// Complete till longest match.
    /// When more than one match, list all matches
    /// (like in Bash/Readline).
    List,
}

/// Style of editing / Standard keymaps
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum EditMode {
    Emacs,
    Vi,
}

/// Configuration builder
#[derive(Debug, Default)]
pub struct Builder {
    p: Config,
}

impl Builder {
    pub fn new() -> Builder {
        Builder {
            p: Config::default(),
        }
    }

    /// Set the maximum length for the history.
    pub fn max_history_size(mut self, max_size: usize) -> Builder {
        self.p.max_history_size = max_size;
        self
    }

    /// Tell if lines which match the previous history entry are saved or not
    /// in the history list.
    ///
    /// By default, they are ignored.
    pub fn history_ignore_dups(mut self, yes: bool) -> Builder {
        self.p.history_duplicates = if yes {
            HistoryDuplicates::IgnoreConsecutive
        } else {
            HistoryDuplicates::AlwaysAdd
        };
        self
    }

    /// Tell if lines which begin with a space character are saved or not in
    /// the history list.
    ///
    /// By default, they are saved.
    pub fn history_ignore_space(mut self, yes: bool) -> Builder {
        self.p.history_ignore_space = yes;
        self
    }

    /// Set `completion_type`.
    pub fn completion_type(mut self, completion_type: CompletionType) -> Builder {
        self.p.completion_type = completion_type;
        self
    }

    /// The number of possible completions that determines when the user is
    /// asked whether the list of possibilities should be displayed.
    pub fn completion_prompt_limit(mut self, completion_prompt_limit: usize) -> Builder {
        self.p.completion_prompt_limit = completion_prompt_limit;
        self
    }

    /// Timeout for ambiguous key sequences in milliseconds.
    /// Currently, it is used only to distinguish a single ESC from an ESC
    /// sequence.
    /// After seeing an ESC key, wait at most `keyseq_timeout_ms` for another
    /// byte.
    pub fn keyseq_timeout(mut self, keyseq_timeout_ms: i32) -> Builder {
        self.p.keyseq_timeout = keyseq_timeout_ms;
        self
    }

    /// Choose between Emacs or Vi mode.
    pub fn edit_mode(mut self, edit_mode: EditMode) -> Builder {
        self.p.edit_mode = edit_mode;
        self
    }

    /// Tell if lines are automatically added to the history.
    ///
    /// By default, they are not.
    pub fn auto_add_history(mut self, yes: bool) -> Builder {
        self.p.auto_add_history = yes;
        self
    }

    pub fn build(self) -> Config {
        self.p
    }
}
<commit_msg>Fix timeout in emacs mode<commit_after>//! Customize line editor
use std::default::Default;

/// User preferences
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Config {
    /// Maximum number of entries in History.
    max_history_size: usize,
    history_duplicates: HistoryDuplicates,
    history_ignore_space: bool,
    completion_type: CompletionType,
    /// When listing completion alternatives, only display
    /// one screen of possibilities at a time.
    completion_prompt_limit: usize,
    /// Duration (milliseconds) Rustyline will wait for a character when
    /// reading an ambiguous key sequence.
    keyseq_timeout: i32,
    /// Emacs or Vi mode
    edit_mode: EditMode,
    /// If true, each nonblank line returned by `readline` will be
    /// automatically added to the history.
    auto_add_history: bool,
}

impl Config {
    pub fn builder() -> Builder {
        Builder::new()
    }

    /// Tell the maximum length (i.e. number of entries) for the history.
    pub fn max_history_size(&self) -> usize {
        self.max_history_size
    }

    /// Tell if lines which match the previous history entry are saved or not
    /// in the history list.
    ///
    /// By default, they are ignored.
    pub fn history_duplicates(&self) -> HistoryDuplicates {
        self.history_duplicates
    }

    /// Tell if lines which begin with a space character are saved or not in
    /// the history list.
    ///
    /// By default, they are saved.
    pub fn history_ignore_space(&self) -> bool {
        self.history_ignore_space
    }

    pub fn completion_type(&self) -> CompletionType {
        self.completion_type
    }

    pub fn completion_prompt_limit(&self) -> usize {
        self.completion_prompt_limit
    }

    pub fn keyseq_timeout(&self) -> i32 {
        self.keyseq_timeout
    }

    pub fn edit_mode(&self) -> EditMode {
        self.edit_mode
    }

    /// Tell if lines are automatically added to the history.
    ///
    /// By default, they are not.
    pub fn auto_add_history(&self) -> bool {
        self.auto_add_history
    }
}

impl Default for Config {
    fn default() -> Config {
        Config {
            max_history_size: 100,
            history_duplicates: HistoryDuplicates::IgnoreConsecutive,
            history_ignore_space: false,
            completion_type: CompletionType::Circular, // TODO Validate
            completion_prompt_limit: 100,
            keyseq_timeout: -1,
            edit_mode: EditMode::Emacs,
            auto_add_history: false,
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum HistoryDuplicates {
    AlwaysAdd,
    /// a line will not be added to the history if it matches the previous entry
    IgnoreConsecutive,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CompletionType {
    /// Complete the next full match (like in Vim by default)
    Circular,
    /// Complete till longest match.
    /// When more than one match, list all matches
    /// (like in Bash/Readline).
    List,
}

/// Style of editing / Standard keymaps
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum EditMode {
    Emacs,
    Vi,
}

/// Configuration builder
#[derive(Debug, Default)]
pub struct Builder {
    p: Config,
}

impl Builder {
    pub fn new() -> Builder {
        Builder {
            p: Config::default(),
        }
    }

    /// Set the maximum length for the history.
    pub fn max_history_size(mut self, max_size: usize) -> Builder {
        self.p.max_history_size = max_size;
        self
    }

    /// Tell if lines which match the previous history entry are saved or not
    /// in the history list.
    ///
    /// By default, they are ignored.
    pub fn history_ignore_dups(mut self, yes: bool) -> Builder {
        self.p.history_duplicates = if yes {
            HistoryDuplicates::IgnoreConsecutive
        } else {
            HistoryDuplicates::AlwaysAdd
        };
        self
    }

    /// Tell if lines which begin with a space character are saved or not in
    /// the history list.
    ///
    /// By default, they are saved.
    pub fn history_ignore_space(mut self, yes: bool) -> Builder {
        self.p.history_ignore_space = yes;
        self
    }

    /// Set `completion_type`.
    pub fn completion_type(mut self, completion_type: CompletionType) -> Builder {
        self.p.completion_type = completion_type;
        self
    }

    /// The number of possible completions that determines when the user is
    /// asked whether the list of possibilities should be displayed.
    pub fn completion_prompt_limit(mut self, completion_prompt_limit: usize) -> Builder {
        self.p.completion_prompt_limit = completion_prompt_limit;
        self
    }

    /// Timeout for ambiguous key sequences in milliseconds.
    /// Currently, it is used only to distinguish a single ESC from an ESC
    /// sequence.
    /// After seeing an ESC key, wait at most `keyseq_timeout_ms` for another
    /// byte.
    pub fn keyseq_timeout(mut self, keyseq_timeout_ms: i32) -> Builder {
        self.p.keyseq_timeout = keyseq_timeout_ms;
        self
    }

    /// Choose between Emacs or Vi mode.
    pub fn edit_mode(mut self, edit_mode: EditMode) -> Builder {
        self.p.edit_mode = edit_mode;
        match edit_mode {
            EditMode::Emacs => self.p.keyseq_timeout = -1, // no timeout
            EditMode::Vi => self.p.keyseq_timeout = 500,
        };
        self
    }

    /// Tell if lines are automatically added to the history.
    ///
    /// By default, they are not.
    pub fn auto_add_history(mut self, yes: bool) -> Builder {
        self.p.auto_add_history = yes;
        self
    }

    pub fn build(self) -> Config {
        self.p
    }
}
<|endoftext|>


<commit_before>extern crate libc;

/// Defines different types of ØMQ sockets.
pub enum SocketType {
    //PAIR = 0,
    //PUB = 1,
    //SUB = 2,

    /// A socket of type `zeromq::REQ` is used by a *client* to send requests to and receive
    /// replies from a *service*. This socket type allows only an alternating sequence of
    /// `send(request)` and subsequent `recv(reply)` calls. Each request sent is round-robined
    /// among all `services`, and each reply received is matched with the last issued request.
    ///
    /// If no services are available, then any send operation on the socket shall block until at
    /// least one service becomes available. The `zeromq::REQ` socket shall not discard messages.
    REQ = 3,

    /// A socket of type `zeromq::REP` is used by a *service* to receive requests from and send
    /// replies to a *client*. This socket type allows only an alternating sequence of
    /// `recv(request)` and subsequent `send(reply)` calls. Each request received is fair-queued
    /// from among all *clients*, and each reply sent is routed to the *client* that issued the
    /// last request. If the original requester does not exist any more the reply is silently
    /// discarded.
    REP = 4,

    //DEALER = 5,
    //ROUTER = 6,
    //PULL = 7,
    //PUSH = 8,
    //XPUB = 9,
    //XSUB = 10,
    //STREAM = 11,
}

/// Defines different options for a ØMQ socket.
///
/// Option value can be retrieved from a socket by
/// [`getsockopt`](trait.ZmqSocket.html#tymethod.getsockopt).
pub enum SocketOption {
    /// *(readonly)* The `zeromq::TYPE` option shall retrieve the socket type for the specified
    /// *socket*. The socket type is specified at socket creation time and cannot be modified
    /// afterwards.
    TYPE = 16,
}

/// A number random enough not to collide with different errno ranges on
/// different OSes. The assumption is that error_t is at least 32-bit type.
const HAUSNUMERO: int = 156384712;

/// ØMQ errors.
#[deriving(PartialEq, Show)]
pub enum ErrorCode {
    /// Invalid argument
    EINVAL = libc::EINVAL as int,

    /// Permission denied
    EACCES = libc::EACCES as int,

    /// Connection refused
    ECONNREFUSED = libc::ECONNREFUSED as int,

    /// Connection reset by peer
    ECONNRESET = libc::ECONNRESET as int,

    /// Software caused connection abort
    ECONNABORTED = libc::ECONNABORTED as int,

    /// Socket is not connected
    ENOTCONN = libc::ENOTCONN as int,

    /// Connection timed out
    ETIMEDOUT = libc::ETIMEDOUT as int,

    /// Protocol not supported
    EPROTONOSUPPORT = HAUSNUMERO + 2,

    /// Message too long
    EMSGSIZE = HAUSNUMERO + 10,

    /// Operation cannot be accomplished in current state
    EFSM = HAUSNUMERO + 51,

    /// Unknown I/O error
    EIOERROR = HAUSNUMERO - 1,
}
<commit_msg>Implement Copy for enums<commit_after>extern crate libc;

/// Defines different types of ØMQ sockets.
pub enum SocketType {
    //PAIR = 0,
    //PUB = 1,
    //SUB = 2,

    /// A socket of type `zeromq::REQ` is used by a *client* to send requests to and receive
    /// replies from a *service*. This socket type allows only an alternating sequence of
    /// `send(request)` and subsequent `recv(reply)` calls. Each request sent is round-robined
    /// among all `services`, and each reply received is matched with the last issued request.
    ///
    /// If no services are available, then any send operation on the socket shall block until at
    /// least one service becomes available. The `zeromq::REQ` socket shall not discard messages.
    REQ = 3,

    /// A socket of type `zeromq::REP` is used by a *service* to receive requests from and send
    /// replies to a *client*. This socket type allows only an alternating sequence of
    /// `recv(request)` and subsequent `send(reply)` calls. Each request received is fair-queued
    /// from among all *clients*, and each reply sent is routed to the *client* that issued the
    /// last request. If the original requester does not exist any more the reply is silently
    /// discarded.
    REP = 4,

    //DEALER = 5,
    //ROUTER = 6,
    //PULL = 7,
    //PUSH = 8,
    //XPUB = 9,
    //XSUB = 10,
    //STREAM = 11,
}

impl Copy for SocketType {}

/// Defines different options for a ØMQ socket.
///
/// Option value can be retrieved from a socket by
/// [`getsockopt`](trait.ZmqSocket.html#tymethod.getsockopt).
pub enum SocketOption {
    /// *(readonly)* The `zeromq::TYPE` option shall retrieve the socket type for the specified
    /// *socket*. The socket type is specified at socket creation time and cannot be modified
    /// afterwards.
    TYPE = 16,
}

impl Copy for SocketOption {}

/// A number random enough not to collide with different errno ranges on
/// different OSes. The assumption is that error_t is at least 32-bit type.
const HAUSNUMERO: int = 156384712;

/// ØMQ errors.
#[deriving(PartialEq, Show)]
pub enum ErrorCode {
    /// Invalid argument
    EINVAL = libc::EINVAL as int,

    /// Permission denied
    EACCES = libc::EACCES as int,

    /// Connection refused
    ECONNREFUSED = libc::ECONNREFUSED as int,

    /// Connection reset by peer
    ECONNRESET = libc::ECONNRESET as int,

    /// Software caused connection abort
    ECONNABORTED = libc::ECONNABORTED as int,

    /// Socket is not connected
    ENOTCONN = libc::ENOTCONN as int,

    /// Connection timed out
    ETIMEDOUT = libc::ETIMEDOUT as int,

    /// Protocol not supported
    EPROTONOSUPPORT = HAUSNUMERO + 2,

    /// Message too long
    EMSGSIZE = HAUSNUMERO + 10,

    /// Operation cannot be accomplished in current state
    EFSM = HAUSNUMERO + 51,

    /// Unknown I/O error
    EIOERROR = HAUSNUMERO - 1,
}

impl Copy for ErrorCode {}

<|endoftext|>


<commit_before>//!
//!	This module contains the code to sum up n-mer aligned read counts based on a
//! SeqTable instance.
//!
use htslib::bam;
use htslib::bam::Read;
use htslib::bam::Reader;
use htslib::bam::Records;
use std::fs::File;
use std::io::Read as ioRead;
use std::io::Seek;
use std::process::exit;
use std::iter::Peekable;
use seqtable::SeqTable;

fn process_bam_seq<R: ioRead+Seek>(counts: &mut Vec<(u64, u64, u64, u64)>, table: &mut SeqTable<R>, bamrecs: &mut Peekable<Records<Reader>>, tid: &mut i32, map: &Vec<usize>, rlen: usize, minqual: u8) -> bool {
    // skip unmapped sequences (tid = -1)
    if *tid < 0 {
        match bamrecs.next() {
            Some(Ok(ref rec)) => { *tid = rec.tid(); return true; },
            Some(Err(_)) => return false,
            None => return false,
        }
    }
    // TODO: store counts table in seqtable file ...

    let mut rdr = table.get_sequence_by_idx(map[*tid as usize]).ok().expect("read sequence");

    loop {
        // check if we changed sequence
        match bamrecs.peek() {
            Some(&Ok(ref rec)) => if rec.tid() != *tid { *tid = rec.tid(); return true; },
            Some(&Err(_)) => return false,
            None => return false,
        }

        // if not count position
        if let Some(Ok(record)) = bamrecs.next() {
            if !record.is_unmapped() && record.seq().len() == rlen && record.mapq() >= minqual {
                let pair = rdr.get(record.pos() as u32).unwrap();

                if record.is_reverse() {
                    counts[pair.1 as usize].3 += 1;
                } else {
                    counts[pair.0 as usize].2 += 1;
                }
            }
        }
    }
}

pub fn tabulate(seqfile: &str, bamfile: Option<String>, minqual: u8) -> Vec<(u64, u64, u64, u64)> {
    // read
    let file = File::open(seqfile).ok().expect("read file");
    let mut table = match SeqTable::open(file) {
        Ok(value) => value,
        Err(e) => {
            println!("Error:tabulate: {}", e.to_string());
            exit(1);
        },
    };

    // allocate counts table
    let mut counts: Vec<(u64, u64, u64, u64)> = Vec::new();
    let nmer_count = 4u64.pow(table.params.cut_length as u32) + 1;

    for _ in 0..nmer_count {
        counts.push((0,0,0,0));
    }

    let rlen = table.params.read_length as usize;
    let seqinfos = table.sequences();

    for idx in 0..seqinfos.len() {
        let length = seqinfos[idx].length;
        let mut rdr = table.get_sequence_by_idx(idx).ok().expect("read sequence");

        for i in 0..length {
            let pair = rdr.get(i).unwrap();

            counts[pair.0 as usize].0 += 1;
            counts[pair.1 as usize].1 += 1;
        }
    }

    // if we received a BAM file, parse it
    if let Some(bamfilename) = bamfile {
        let bam = bam::Reader::new(&bamfilename).ok().expect("Error opening bam.");
	    let names = bam.header.target_names();
        let mut cur_tid = 0;

        // map BAM tid's to SeqTable idx's
        let map : Vec<usize> = names.iter().map(|&id| {
            let chrom = String::from_utf8_lossy(id);
            for i in 0..seqinfos.len() {
                if seqinfos[i].name.eq(&chrom) {
                    return i;
                }
            }
            println!("Error: Unknown sequence name in BAM: {}", chrom);
            exit(1);
        }
        ).collect();

        // reads
        let mut iter = bam.records().peekable();

        while process_bam_seq(&mut counts, &mut table, &mut iter, &mut cur_tid, &map, rlen, minqual) {}
    }

    counts
}

pub fn print_counts(counts: &Vec<(u64, u64, u64, u64)>, with_bam: bool) {
    // TODO: tmp: output table to file
    if with_bam {
        for i in 1..counts.len() {
            let (plus, minus, bam_plus, bam_minus) = counts[i];
            println!("{}\t{}\t{}\t{}\t{}", i, plus, minus, bam_plus, bam_minus);
        }
    } else {
        for i in 1..counts.len() {
            let (plus, minus, _, _) = counts[i];
            println!("{}\t{}\t{}", i, plus, minus);
        }
    }
}<commit_msg>Output kmer sequence key along with table<commit_after>//!
//!	This module contains the code to sum up n-mer aligned read counts based on a
//! SeqTable instance.
//!
use htslib::bam;
use htslib::bam::Read;
use htslib::bam::Reader;
use htslib::bam::Records;
use std::fs::File;
use std::io::Read as ioRead;
use std::io::Seek;
use std::process::exit;
use std::iter::Peekable;
use seqtable::SeqTable;

struct KeyIter {
    kmer: Vec<u8>,
    alph: [char; 4],
}

impl KeyIter {
    fn new(k: u8) -> KeyIter {
        let mut kmer = vec![1; k as usize];
        kmer[(k - 1) as usize] = 0;
        KeyIter { kmer: kmer, alph: ['A','C','G','T'] }
    }

    fn key(&self) -> String {
        let mut key = String::new();

        for b in &self.kmer {
            key.push(self.alph[(b - 1) as usize]);
        }
        key
    }
}

impl Iterator for KeyIter {
    type Item = String;

    fn next(&mut self) -> Option<String> {
        let mut idx = self.kmer.len() - 1;

        loop {
            self.kmer[idx] += 1;
            if self.kmer[idx] > 4 {
                self.kmer[idx] = 1;

                if idx > 0 {
                  idx -= 1;
                } else {
                  return None;
                }
            } else {
                break;
            }
        }

        Some(self.key())
    }
}

fn process_bam_seq<R: ioRead+Seek>(counts: &mut Vec<(u64, u64, u64, u64)>, table: &mut SeqTable<R>, bamrecs: &mut Peekable<Records<Reader>>, tid: &mut i32, map: &Vec<usize>, rlen: usize, minqual: u8) -> bool {
    // skip unmapped sequences (tid = -1)
    if *tid < 0 {
        match bamrecs.next() {
            Some(Ok(ref rec)) => { *tid = rec.tid(); return true; },
            Some(Err(_)) => return false,
            None => return false,
        }
    }
    // TODO: store counts table in seqtable file ...

    let mut rdr = table.get_sequence_by_idx(map[*tid as usize]).ok().expect("read sequence");

    loop {
        // check if we changed sequence
        match bamrecs.peek() {
            Some(&Ok(ref rec)) => if rec.tid() != *tid { *tid = rec.tid(); return true; },
            Some(&Err(_)) => return false,
            None => return false,
        }

        // if not count position
        if let Some(Ok(record)) = bamrecs.next() {
            if !record.is_unmapped() && record.seq().len() == rlen && record.mapq() >= minqual {
                let pair = rdr.get(record.pos() as u32).unwrap();

                if record.is_reverse() {
                    counts[pair.1 as usize].3 += 1;
                } else {
                    counts[pair.0 as usize].2 += 1;
                }
            }
        }
    }
}

pub fn tabulate(seqfile: &str, bamfile: Option<String>, minqual: u8) -> Vec<(u64, u64, u64, u64)> {
    // read
    let file = File::open(seqfile).ok().expect("read file");
    let mut table = match SeqTable::open(file) {
        Ok(value) => value,
        Err(e) => {
            println!("Error:tabulate: {}", e.to_string());
            exit(1);
        },
    };

    // allocate counts table
    let mut counts: Vec<(u64, u64, u64, u64)> = Vec::new();
    let nmer_count = 4u64.pow(table.params.cut_length as u32) + 1;

    for _ in 0..nmer_count {
        counts.push((0,0,0,0));
    }

    let rlen = table.params.read_length as usize;
    let seqinfos = table.sequences();

    for idx in 0..seqinfos.len() {
        let length = seqinfos[idx].length;
        let mut rdr = table.get_sequence_by_idx(idx).ok().expect("read sequence");

        for i in 0..length {
            let pair = rdr.get(i).unwrap();

            counts[pair.0 as usize].0 += 1;
            counts[pair.1 as usize].1 += 1;
        }
    }

    // if we received a BAM file, parse it
    if let Some(bamfilename) = bamfile {
        let bam = bam::Reader::new(&bamfilename).ok().expect("Error opening bam.");
	    let names = bam.header.target_names();
        let mut cur_tid = 0;

        // map BAM tid's to SeqTable idx's
        let map : Vec<usize> = names.iter().map(|&id| {
            let chrom = String::from_utf8_lossy(id);
            for i in 0..seqinfos.len() {
                if seqinfos[i].name.eq(&chrom) {
                    return i;
                }
            }
            println!("Error: Unknown sequence name in BAM: {}", chrom);
            exit(1);
        }
        ).collect();

        // reads
        let mut iter = bam.records().peekable();

        while process_bam_seq(&mut counts, &mut table, &mut iter, &mut cur_tid, &map, rlen, minqual) {}
    }

    counts
}

pub fn print_counts(counts: &Vec<(u64, u64, u64, u64)>, with_bam: bool) {
    let k = (((counts.len() - 1) as f32).log2() / 2.0) as u8;
    let mut keys = KeyIter::new(k);

    // TODO: tmp: output table to file
    if with_bam {
        for i in 1..counts.len() {
            let (plus, minus, bam_plus, bam_minus) = counts[i];
            let key = keys.next().unwrap();
            println!("{}\t{}\t{}\t{}\t{}\t{}", i, key, plus, minus, bam_plus, bam_minus);
        }
    } else {
        for i in 1..counts.len() {
            let (plus, minus, _, _) = counts[i];
            let key = keys.next().unwrap();
            println!("{}\t{}\t{}\t{}", i, key, plus, minus);
        }
    }
}<|endoftext|>


<commit_before>use super::{ReadAt, Size, WriteAt};

use std::io;
use std::io::{Read, Seek, SeekFrom, Write};

/// Adapts a `ReadAt` or `WriteAt` into a `Read` or `Write`.
///
/// This wraps anything that read and write at offsets, turning into an object that can
/// read or write at a file position. This allows you to use those types with all the many
/// functions that expect a `Read` or `Write`.
///
/// Note that seeking on `Cursor` has limited functionality. We don't know how many bytes are
/// available, so we can't use SeekFrom::End. See [`SizeCursor`][SizeCursor] for another option.
///
/// [SizeCursor]: struct.SizeCursor.html
///
/// # Examples
///
/// ```no_run
/// # use std::io::{self, Result, Read};
/// # use std::fs::File;
/// use positioned_io::{ReadAt, Cursor};
///
/// struct NetworkStorage {
///     // A remote disk that supports random access.
/// }
/// # impl NetworkStorage {
/// #   fn new(i: i32) -> Self { NetworkStorage { } }
/// # }
/// impl ReadAt for NetworkStorage {
///     // ...
/// #   fn read_at(&self, pos: u64, buf: &mut [u8]) -> Result<usize> {
/// #       Ok(0)
/// #   }
/// }
///
/// # const SOME_LOCATION: i32 = 1;
/// # fn foo() -> Result<()> {
/// // Adapt our storage into a Read.
/// let storage = NetworkStorage::new(SOME_LOCATION);
/// let curs = Cursor::new_pos(storage, 1 << 30);
///
/// // Copy a segment to a file.
/// let mut input = curs.take(1 << 20);
/// let mut output = File::create("segment.out")?;
/// io::copy(&mut input, &mut output)?;
/// # Ok(())
/// # }
/// ```
#[derive(Debug, Clone)]
pub struct Cursor<I> {
    io: I,
    pos: u64,
}
impl<I> Cursor<I> {
    /// Create a new `Cursor` which starts reading at a specified offset.
    ///
    /// Pass in a `ReadAt` or `WriteAt` as `io`.
    pub fn new_pos(io: I, pos: u64) -> Self {
        Cursor { io, pos }
    }
    /// Create a new Cursor which starts reading at offset zero.
    ///
    /// Pass in a `ReadAt` or `WriteAt` as `io`.
    pub fn new(io: I) -> Self {
        Self::new_pos(io, 0)
    }

    /// Consume `self` and yield the inner `ReadAt` or `WriteAt`.
    pub fn into_inner(self) -> I {
        self.io
    }
    /// Borrow the inner `ReadAt` or `WriteAt`.
    pub fn get_ref(&self) -> &I {
        &self.io
    }
    /// Borrow the inner `ReadAt` or `WriteAt` mutably.
    pub fn get_mut(&mut self) -> &mut I {
        &mut self.io
    }

    /// Get the current read/write position.
    pub fn position(&self) -> u64 {
        self.pos
    }
    /// Set the current read/write position.
    pub fn set_position(&mut self, pos: u64) {
        self.pos = pos;
    }
}

impl<I> Seek for Cursor<I> {
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        match pos {
            SeekFrom::Start(p) => self.pos = p,
            SeekFrom::Current(p) => {
                let pos = self.pos as i64 + p;
                if pos < 0 {
                    return Err(io::Error::new(io::ErrorKind::InvalidInput, "seek to a negative position"));
                }
                self.pos = pos as u64;
            }
            SeekFrom::End(_) => {
                return Err(io::Error::new(io::ErrorKind::InvalidInput, "seek from unknown end"))
            }
        };
        Ok(self.pos)
    }
}

impl<I> Read for Cursor<I>
where
    I: ReadAt,
{
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let bytes = self.get_ref().read_at(self.pos, buf)?;
        self.pos += bytes as u64;
        Ok(bytes)
    }
}
impl<I> Write for Cursor<I>
where
    I: WriteAt,
{
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let pos = self.pos;
        let bytes = self.get_mut().write_at(pos, buf)?;
        self.pos += bytes as u64;
        Ok(bytes)
    }
    fn flush(&mut self) -> io::Result<()> {
        WriteAt::flush(self.get_mut())
    }
}

/// Adapts a `ReadAt` or `WriteAt` into a `Read` or `Write`, with better seeking.
///
/// This is just like [`Cursor`][Cursor], except that it requires an object that implements
/// [`Size`][Size], and that it can seek from the end of the I/O object.
///
/// Eventually it will be legal to specialize `Cursor` for types that implement `Size`, see
/// [RFC 1210][RFC].
///
/// [Cursor]: struct.Cursor.html
/// [Size]: trait.Size.html
/// [RFC]: https://github.com/rust-lang/rfcs/pull/1210
#[derive(Debug, Clone)]
pub struct SizeCursor<I: Size> {
    cursor: Cursor<I>,
}

impl<I: Size> SizeCursor<I> {
    /// Create a new `SizeCursor` which starts reading at a specified offset.
    ///
    /// Pass in a `ReadAt` or `WriteAt` as `io`.
    pub fn new_pos(io: I, pos: u64) -> Self {
        SizeCursor {
            cursor: Cursor::new_pos(io, pos)
        }
    }
    /// Create a new `SizeCursor` which starts reading at offset zero.
    ///
    /// Pass in a `ReadAt` or `WriteAt` as `io`.
    pub fn new(io: I) -> Self {
        SizeCursor {
            cursor: Cursor::new(io)
        }
    }

    pub fn as_cursor(&self) -> &Cursor<I> {
        &self.cursor
    }

    pub fn as_cursor_mut(&mut self) -> &mut Cursor<I> {
        &mut self.cursor
    }

    pub fn into_cursor(self) -> Cursor<I> {
        self.cursor
    }

    pub fn into_inner(self) -> I {
        self.cursor.io
    }

    pub fn get_ref(&self) -> &I {
        &self.cursor.io
    }

    pub fn get_mut(&mut self) -> &mut I {
        &mut self.cursor.io
    }

    pub fn position(&self) -> u64 {
        self.cursor.position()
    }

    pub fn set_position(&mut self, pos: u64) {
        self.cursor.set_position(pos)
    }
}

impl<I: Size + ReadAt> Read for SizeCursor<I> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.cursor.read(buf)
    }
}

impl<I: Size + WriteAt> Write for SizeCursor<I> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.cursor.write(buf)
    }
    fn flush(&mut self) -> io::Result<()> {
        self.cursor.flush()
    }
}

// We know how to seek from the end for SizeCursor.
impl<I> Seek for SizeCursor<I>
where
    I: Size,
{
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        let pos = match pos {
            SeekFrom::Start(p) => p as i64,
            SeekFrom::Current(p) => self.cursor.pos as i64 + p,
            SeekFrom::End(p) => {
                match self.get_ref().size() {
                    Err(e) => return Err(e),
                    Ok(None) => {
                        return Err(io::Error::new(io::ErrorKind::InvalidData, "seek from unknown end"))
                    }
                    Ok(Some(s)) => s as i64 + p,
                }
            }
        };
        self.cursor.pos = pos as u64;
        Ok(self.cursor.pos)
    }
}
<commit_msg>tweak cursor docs<commit_after>use super::{ReadAt, Size, WriteAt};

use std::io;
use std::io::{Read, Seek, SeekFrom, Write};

/// Adapts a `ReadAt` or `WriteAt` into a `Read` or `Write`.
///
/// This wraps anything that read and write at offsets, turning into an object
/// that can read or write at a file position. This allows you to use those
/// types with all the many functions that expect a `Read` or `Write`.
///
/// Note that seeking on `Cursor` has limited functionality. We don't know how
/// many bytes are available, so we can't use `SeekFrom::End`.
/// See [`SizeCursor`][SizeCursor] for another option.
///
/// [SizeCursor]: struct.SizeCursor.html
///
/// # Examples
///
/// ```no_run
/// # use std::io::{self, Result, Read};
/// # use std::fs::File;
/// use positioned_io::{ReadAt, Cursor};
///
/// struct NetworkStorage {
///     // A remote disk that supports random access.
/// }
/// # impl NetworkStorage {
/// #   fn new(i: i32) -> Self { NetworkStorage { } }
/// # }
///
/// impl ReadAt for NetworkStorage {
///     // ...
/// #   fn read_at(&self, pos: u64, buf: &mut [u8]) -> Result<usize> {
/// #       Ok(0)
/// #   }
/// }
///
/// # const SOME_LOCATION: i32 = 1;
/// # fn foo() -> Result<()> {
/// // Adapt our storage into a Read.
/// let storage = NetworkStorage::new(SOME_LOCATION);
/// let curs = Cursor::new_pos(storage, 1 << 30);
///
/// // Copy a segment to a file.
/// let mut input = curs.take(1 << 20);
/// let mut output = File::create("segment.out")?;
/// io::copy(&mut input, &mut output)?;
/// # Ok(())
/// # }
/// ```
#[derive(Debug, Clone)]
pub struct Cursor<I> {
    io: I,
    pos: u64,
}

impl<I> Cursor<I> {
    /// Create a new `Cursor` which starts reading at a specified offset.
    ///
    /// Pass in a `ReadAt` or `WriteAt` as `io`.
    pub fn new_pos(io: I, pos: u64) -> Self {
        Cursor { io, pos }
    }

    /// Create a new Cursor which starts reading at offset zero.
    ///
    /// Pass in a `ReadAt` or `WriteAt` as `io`.
    pub fn new(io: I) -> Self {
        Self::new_pos(io, 0)
    }

    /// Consume `self` and yield the inner `ReadAt` or `WriteAt`.
    pub fn into_inner(self) -> I {
        self.io
    }

    /// Borrow the inner `ReadAt` or `WriteAt`.
    pub fn get_ref(&self) -> &I {
        &self.io
    }

    /// Borrow the inner `ReadAt` or `WriteAt` mutably.
    pub fn get_mut(&mut self) -> &mut I {
        &mut self.io
    }

    /// Get the current read/write position.
    pub fn position(&self) -> u64 {
        self.pos
    }

    /// Set the current read/write position.
    pub fn set_position(&mut self, pos: u64) {
        self.pos = pos;
    }
}

impl<I> Seek for Cursor<I> {
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        match pos {
            SeekFrom::Start(p) => self.pos = p,
            SeekFrom::Current(p) => {
                let pos = self.pos as i64 + p;
                if pos < 0 {
                    return Err(io::Error::new(io::ErrorKind::InvalidInput, "seek to a negative position"));
                }
                self.pos = pos as u64;
            }
            SeekFrom::End(_) => {
                return Err(io::Error::new(io::ErrorKind::InvalidInput, "seek from unknown end"))
            }
        }
        Ok(self.pos)
    }
}

impl<I: ReadAt> Read for Cursor<I> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let bytes = self.get_ref().read_at(self.pos, buf)?;
        self.pos += bytes as u64;
        Ok(bytes)
    }
}

impl<I: WriteAt> Write for Cursor<I> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let pos = self.pos;
        let bytes = self.get_mut().write_at(pos, buf)?;
        self.pos += bytes as u64;
        Ok(bytes)
    }

    fn flush(&mut self) -> io::Result<()> {
        WriteAt::flush(self.get_mut())
    }
}

/// Adapts a `ReadAt` or `WriteAt` into a `Read` or `Write`, with better
/// seeking.
///
/// This is just like [`Cursor`][Cursor], except that it requires an object
/// that implements [`Size`][Size], and that it can seek from the end of the
/// I/O object.
///
/// Eventually it will be legal to specialize `Cursor` for types that implement
/// `Size`, see [RFC 1210][RFC].
///
/// [Cursor]: struct.Cursor.html
/// [Size]: trait.Size.html
/// [RFC]: https://github.com/rust-lang/rfcs/pull/1210
#[derive(Debug, Clone)]
pub struct SizeCursor<I: Size> {
    cursor: Cursor<I>,
}

impl<I: Size> SizeCursor<I> {
    /// Create a new `SizeCursor` which starts reading at a specified offset.
    ///
    /// Pass in a `ReadAt` or `WriteAt` as `io`.
    pub fn new_pos(io: I, pos: u64) -> Self {
        SizeCursor {
            cursor: Cursor::new_pos(io, pos)
        }
    }
    /// Create a new `SizeCursor` which starts reading at offset zero.
    ///
    /// Pass in a `ReadAt` or `WriteAt` as `io`.
    pub fn new(io: I) -> Self {
        SizeCursor {
            cursor: Cursor::new(io)
        }
    }

    pub fn as_cursor(&self) -> &Cursor<I> {
        &self.cursor
    }

    pub fn as_cursor_mut(&mut self) -> &mut Cursor<I> {
        &mut self.cursor
    }

    pub fn into_cursor(self) -> Cursor<I> {
        self.cursor
    }

    pub fn into_inner(self) -> I {
        self.cursor.io
    }

    pub fn get_ref(&self) -> &I {
        &self.cursor.io
    }

    pub fn get_mut(&mut self) -> &mut I {
        &mut self.cursor.io
    }

    pub fn position(&self) -> u64 {
        self.cursor.position()
    }

    pub fn set_position(&mut self, pos: u64) {
        self.cursor.set_position(pos)
    }
}

impl<I: Size + ReadAt> Read for SizeCursor<I> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.cursor.read(buf)
    }
}

impl<I: Size + WriteAt> Write for SizeCursor<I> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.cursor.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.cursor.flush()
    }
}

// We know how to seek from the end for SizeCursor.
impl<I: Size> Seek for SizeCursor<I> {
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        let pos = match pos {
            SeekFrom::Start(p) => p as i64,
            SeekFrom::Current(p) => self.cursor.pos as i64 + p,
            SeekFrom::End(p) => {
                match self.get_ref().size() {
                    Err(e) => return Err(e),
                    Ok(None) => {
                        return Err(io::Error::new(io::ErrorKind::InvalidData, "seek from unknown end"))
                    }
                    Ok(Some(s)) => s as i64 + p,
                }
            }
        };
        self.cursor.pos = pos as u64;
        Ok(self.cursor.pos)
    }
}
<|endoftext|>


<commit_before>use std::sync::{Arc, RwLock};
use std::collections::HashMap;
use multimap::MultiMap;
use rand::{Rng, thread_rng};
use dns_parser::{self, QueryClass, Name, RRData};

pub type AnswerBuilder = dns_parser::Builder<dns_parser::Answers>;

/// A collection of registered services is shared between threads.
pub type Services = Arc<RwLock<ServicesInner>>;

pub struct ServicesInner {
    hostname: Name<'static>,
    /// main index
    by_id: HashMap<usize, ServiceData>,
    /// maps to id
    by_type: MultiMap<Name<'static>, usize>,
    /// maps to id
    by_name: HashMap<Name<'static>, usize>
}

impl ServicesInner {
    pub fn new(hostname: String) -> Self {
        ServicesInner {
            hostname: Name::from_str(hostname).unwrap(),
            by_id: HashMap::new(),
            by_type: MultiMap::new(),
            by_name: HashMap::new(),
        }
    }

    pub fn get_hostname(&self) -> &Name<'static> {
        &self.hostname
    }

    pub fn find_by_name<'a>(&'a self, name: &'a Name<'a>) -> Option<&ServiceData> {
        self.by_name.get(name)
            .and_then(|id| self.by_id.get(id))
    }

    pub fn find_by_type<'a>(&'a self, ty: &'a Name<'a>) -> FindByType<'a> {
        let ids = self.by_type.get_vec(ty)
            .map(|ids| &ids[..])
            .unwrap_or(&[]);
        FindByType {
            services: self,
            ids: ids
        }
    }

    pub fn register(&mut self, svc: ServiceData) -> usize {
        let mut id = thread_rng().gen::<usize>();
        while self.by_id.contains_key(&id) {
            id = thread_rng().gen::<usize>();
        }

        self.by_type.insert(svc.typ.clone(), id);
        self.by_name.insert(svc.name.clone(), id);
        self.by_id.insert(id, svc);

        id
    }

    pub fn unregister(&mut self, id: usize) -> ServiceData {
        use std::collections::hash_map::Entry;

        let svc = self.by_id.remove(&id).expect("unknown service");

        if let Some(entries) = self.by_type.get_vec_mut(&svc.typ) {
            entries.retain(|&e| e == id);
        }

        match self.by_name.entry(svc.name.clone()) {
            Entry::Occupied(entry) => {
                assert_eq!(*entry.get(), id);
                entry.remove();
            }
            _ => {
                panic!("unknown/wrong service for id {}", id);
            }
        }

        svc
    }
}

/// Returned by [`ServicesInner.find_by_type`](struct.ServicesInner.html#method.find_by_type)
pub struct FindByType<'a> {
    services: &'a ServicesInner,
    ids: &'a [usize],
}

impl<'a> Iterator for FindByType<'a> {
    type Item = &'a ServiceData;

    fn next(&mut self) -> Option<Self::Item> {
        match self.ids.split_first() {
            Some((id, rest)) => {
                self.ids = rest;
                self.services.by_id.get(id)
            }
            None => None
        }
    }
}

#[derive(Clone)]
pub struct ServiceData {
    pub name: Name<'static>,
    pub typ: Name<'static>,
    pub port: u16,
    pub txt: Vec<u8>,
}

/// Packet building helpers for `fsm` to respond with `ServiceData`
impl ServiceData {
    pub fn add_ptr_rr(&self, builder: AnswerBuilder, ttl: u32) -> AnswerBuilder {
        builder.add_answer(&self.typ, QueryClass::IN, ttl, &RRData::PTR(self.name.clone()))
    }

    pub fn add_srv_rr(&self, hostname: &Name, builder: AnswerBuilder, ttl: u32) -> AnswerBuilder {
        builder.add_answer(&self.name, QueryClass::IN, ttl, &RRData::SRV {
            priority: 0,
            weight: 0,
            port: self.port,
            target: hostname.clone(),
        })
    }

    pub fn add_txt_rr(&self, builder: AnswerBuilder, ttl: u32) -> AnswerBuilder {
        builder.add_answer(&self.name, QueryClass::IN, ttl, &RRData::TXT(&self.txt))
    }
}
<commit_msg>Make the FindByType iterator lazier<commit_after>use std::sync::{Arc, RwLock};
use std::collections::HashMap;
use std::slice;
use multimap::MultiMap;
use rand::{Rng, thread_rng};
use dns_parser::{self, QueryClass, Name, RRData};

pub type AnswerBuilder = dns_parser::Builder<dns_parser::Answers>;

/// A collection of registered services is shared between threads.
pub type Services = Arc<RwLock<ServicesInner>>;

pub struct ServicesInner {
    hostname: Name<'static>,
    /// main index
    by_id: HashMap<usize, ServiceData>,
    /// maps to id
    by_type: MultiMap<Name<'static>, usize>,
    /// maps to id
    by_name: HashMap<Name<'static>, usize>
}

impl ServicesInner {
    pub fn new(hostname: String) -> Self {
        ServicesInner {
            hostname: Name::from_str(hostname).unwrap(),
            by_id: HashMap::new(),
            by_type: MultiMap::new(),
            by_name: HashMap::new(),
        }
    }

    pub fn get_hostname(&self) -> &Name<'static> {
        &self.hostname
    }

    pub fn find_by_name<'a>(&'a self, name: &'a Name<'a>) -> Option<&ServiceData> {
        self.by_name.get(name)
            .and_then(|id| self.by_id.get(id))
    }

    pub fn find_by_type<'a>(&'a self, ty: &'a Name<'a>) -> FindByType<'a> {
        let ids = self.by_type.get_vec(ty)
                              .map(|ids| ids.iter());

        FindByType {
            services: self,
            ids: ids,
        }
    }

    pub fn register(&mut self, svc: ServiceData) -> usize {
        let mut id = thread_rng().gen::<usize>();
        while self.by_id.contains_key(&id) {
            id = thread_rng().gen::<usize>();
        }

        self.by_type.insert(svc.typ.clone(), id);
        self.by_name.insert(svc.name.clone(), id);
        self.by_id.insert(id, svc);

        id
    }

    pub fn unregister(&mut self, id: usize) -> ServiceData {
        use std::collections::hash_map::Entry;

        let svc = self.by_id.remove(&id).expect("unknown service");

        if let Some(entries) = self.by_type.get_vec_mut(&svc.typ) {
            entries.retain(|&e| e == id);
        }

        match self.by_name.entry(svc.name.clone()) {
            Entry::Occupied(entry) => {
                assert_eq!(*entry.get(), id);
                entry.remove();
            }
            _ => {
                panic!("unknown/wrong service for id {}", id);
            }
        }

        svc
    }
}

/// Returned by [`ServicesInner.find_by_type`](struct.ServicesInner.html#method.find_by_type)
pub struct FindByType<'a> {
    services: &'a ServicesInner,
    ids: Option<slice::Iter<'a, usize>>,
}

impl<'a> Iterator for FindByType<'a> {
    type Item = &'a ServiceData;

    fn next(&mut self) -> Option<Self::Item> {
        self.ids.as_mut()
            .and_then(Iterator::next)
            .map(|id| {
                let svc = self.services.by_id.get(id);
                svc.expect("missing service")
            })
    }
}

#[derive(Clone)]
pub struct ServiceData {
    pub name: Name<'static>,
    pub typ: Name<'static>,
    pub port: u16,
    pub txt: Vec<u8>,
}

/// Packet building helpers for `fsm` to respond with `ServiceData`
impl ServiceData {
    pub fn add_ptr_rr(&self, builder: AnswerBuilder, ttl: u32) -> AnswerBuilder {
        builder.add_answer(&self.typ, QueryClass::IN, ttl, &RRData::PTR(self.name.clone()))
    }

    pub fn add_srv_rr(&self, hostname: &Name, builder: AnswerBuilder, ttl: u32) -> AnswerBuilder {
        builder.add_answer(&self.name, QueryClass::IN, ttl, &RRData::SRV {
            priority: 0,
            weight: 0,
            port: self.port,
            target: hostname.clone(),
        })
    }

    pub fn add_txt_rr(&self, builder: AnswerBuilder, ttl: u32) -> AnswerBuilder {
        builder.add_answer(&self.name, QueryClass::IN, ttl, &RRData::TXT(&self.txt))
    }
}
<|endoftext|>