<commit_before>// SPDX-License-Identifier: MIT
//
// Copyright (C) 2020-2022 Jason Ish

#![allow(clippy::redundant_field_names)]

use clap::{Arg, Command};
use evebox::logger;
use evebox::prelude::*;
use evebox::version;

#[tokio::main]
async fn main() {
    if let Err(err) = _main().await {
        error!("{}", err);
        std::process::exit(1);
    }
}

async fn _main() -> Result<(), Box<dyn std::error::Error>> {
    let parser = clap::Command::new("EveBox")
        .version(std::env!("CARGO_PKG_VERSION"))
        .arg(
            Arg::new("verbose")
                .long("verbose")
                .short('v')
                .multiple_occurrences(true)
                .global(true)
                .help("Increase verbosity"),
        )
        .arg(
            Arg::new("data-directory")
                .long("data-directory")
                .short('D')
                .takes_value(true)
                .value_name("DIR")
                .help("Data directory")
                .global(true),
        )
        .subcommand(clap::Command::new("version").about("Display version"));

    let sqlite_import = Command::new("sqlite-import")
        .about("Import to SQLite")
        .arg(
            Arg::new("oneshot")
                .long("oneshot")
                .help("One shot mode (exit on EOF"),
        )
        .arg(Arg::new("end").long("end").help("Start at end of file"))
        .arg(Arg::new("INPUT").required(true).index(1));

    let server = clap::Command::new("server")
        .about("EveBox Server")
        .arg(
            clap::Arg::new("config")
                .long("config")
                .short('c')
                .takes_value(true)
                .value_name("FILE")
                .help("Configuration filename"),
        )
        .arg(
            clap::Arg::new("http.host")
                .long("host")
                .value_name("HOSTNAME")
                .takes_value(true)
                .default_value("127.0.0.1")
                .help("Hostname/IP address to bind to"),
        )
        .arg(
            clap::Arg::new("http.port")
                .short('p')
                .long("port")
                .value_name("PORT")
                .takes_value(true)
                .default_value("5636")
                .env("EVEBOX_HTTP_PORT")
                .help("Port to bind to"),
        )
        .arg(
            clap::Arg::new("database.elasticsearch.url")
                .short('e')
                .long("elasticsearch")
                .takes_value(true)
                .value_name("URL")
                .default_value("http://localhost:9200")
                .help("Elastic Search URL"),
        )
        .arg(
            clap::Arg::new("database.elasticsearch.index")
                .short('i')
                .long("index")
                .takes_value(true)
                .default_value("logstash")
                .value_name("INDEX")
                .help("Elastic Search index prefix"),
        )
        .arg(
            Arg::new("database.elasticsearch.no-index-suffix")
                .long("no-index-suffix")
                .takes_value(false)
                .help("Do not add a suffix to the index name"),
        )
        .arg(
            Arg::new("database.elasticsearch.ecs")
                .long("ecs")
                .help("Enable Elastic ECS support"),
        )
        .arg(
            Arg::new("database.type")
                .long("database")
                .aliases(&["datastore"])
                .takes_value(true)
                .value_name("DATABASE")
                .default_value("elasticsearch")
                .help("Database type"),
        )
        .arg(
            Arg::new("sqlite-filename")
                .long("sqlite-filename")
                .takes_value(true)
                .value_name("FILENAME")
                .default_value("events.sqlite")
                .help("SQLite events filename"),
        )
        .arg(
            clap::Arg::new("no-check-certificate")
                .short('k')
                .long("no-check-certificate")
                .help("Disable TLS certificate validation"),
        )
        .arg(
            Arg::new("http.request-logging")
                .long("http-request-logging")
                .env("EVEBOX_HTTP_REQUEST_LOGGING")
                .hide_env(true)
                .help("Enable HTTP access logging"),
        )
        .arg(
            Arg::new("http.tls.enabled")
                .long("tls")
                .help("Enable TLS")
                .env("EVEBOX_HTTP_TLS_ENABLED")
                .hide_env(true),
        )
        .arg(
            Arg::new("http.tls.certificate")
                .long("tls-cert")
                .takes_value(true)
                .value_name("FILENAME")
                .help("TLS certificate filename"),
        )
        .arg(
            Arg::new("http.tls.key")
                .long("tls-key")
                .takes_value(true)
                .value_name("FILENAME")
                .help("TLS key filename"),
        )
        .arg(
            Arg::new("input.filename")
                .long("input")
                .takes_value(true)
                .value_name("FILENAME")
                .help("Input Eve file to read"),
        )
        .arg(
            Arg::new("end")
                .long("end")
                .help("Read from end (tail) of input file"),
        )
        .arg(Arg::new("input-start").long("input-start").hide(true));

    let oneshot = Command::new("oneshot")
        .about("Import a single eve.json and review in EveBox")
        .arg(
            Arg::new("limit")
                .long("limit")
                .takes_value(true)
                .help("Limit the number of events read"),
        )
        .arg(
            Arg::new("no-open")
                .long("no-open")
                .help("Don't open browser"),
        )
        .arg(
            Arg::new("no-wait")
                .long("no-wait")
                .help("Don't wait for events to load"),
        )
        .arg(
            Arg::new("database-filename")
                .long("database-filename")
                .takes_value(true)
                .default_value("./oneshot.sqlite")
                .help("Database filename"),
        )
        // --host, but keep th name as http.host to be campatible with the
        // EVEBOX_HTTP_HOST environment variable.
        .arg(
            clap::Arg::new("http.host")
                .long("host")
                .value_name("HOSTNAME")
                .takes_value(true)
                .default_value("127.0.0.1")
                .help("Hostname/IP address to bind to"),
        )
        .arg(Arg::new("INPUT").required(true).index(1));

    let elastic_import = Command::new("elastic-import")
        .about("Import to Elastic Search")
        .arg(
            Arg::new("config")
                .short('c')
                .long("config")
                .takes_value(true)
                .help("Configuration file"),
        )
        .arg(
            Arg::new("oneshot")
                .long("oneshot")
                .help("One shot mode (exit on EOF)"),
        )
        .arg(Arg::new("end").long("end").help("Start at end of file"))
        .arg(
            clap::Arg::new("elasticsearch")
                .short('e')
                .long("elasticsearch")
                .takes_value(true)
                .default_value("http://localhost:9200")
                .hide_default_value(true)
                .help("Elastic Search URL"),
        )
        .arg(
            Arg::new("index")
                .long("index")
                .takes_value(true)
                .default_value("logstash")
                .hide_default_value(true)
                .help("Elastic Search index prefix"),
        )
        .arg(
            Arg::new("no-index-suffix")
                .long("no-index-suffix")
                .takes_value(false)
                .help("Do not add a suffix to the index name"),
        )
        .arg(
            Arg::new("bookmark")
                .long("bookmark")
                .help("Enable bookmarking"),
        )
        .arg(
            Arg::new("bookmark-filename")
                .long("bookmark-filename")
                .takes_value(true)
                .default_value("")
                .hide_default_value(true)
                .help("Bookmark filename"),
        )
        .arg(
            Arg::new("bookmark-dir")
                .long("bookmark-dir")
                .takes_value(true)
                .default_value(".")
                .hide_default_value(true)
                .help("Bookmark directory"),
        )
        .arg(
            Arg::new("stdout")
                .long("stdout")
                .help("Print events to stdout"),
        )
        .arg(
            Arg::new("username")
                .long("username")
                .short('u')
                .takes_value(true)
                .help("Elasticsearch username"),
        )
        .arg(
            Arg::new("password")
                .long("password")
                .short('p')
                .takes_value(true)
                .help("Elasticsearch password"),
        )
        .arg(
            clap::Arg::new(evebox::commands::elastic_import::NO_CHECK_CERTIFICATE)
                .short('k')
                .long("no-check-certificate")
                .help("Disable TLS certificate validation"),
        )
        .arg(
            Arg::new("geoip.disabled")
                .long("no-geoip")
                .help("Disable GeoIP"),
        )
        .arg(
            Arg::new("geoip.database-filename")
                .long("geoip-database")
                .takes_value(true)
                .value_name("filename")
                .help("GeoIP database filename"),
        )
        .arg(
            Arg::new("input")
                .multiple_values(true)
                .value_name("INPUT")
                .index(1),
        );

    let mut parser = parser
        .subcommand(server)
        .subcommand(elastic_import)
        .subcommand(oneshot)
        .subcommand(evebox::commands::agent::command())
        .subcommand(sqlite_import)
        .subcommand(evebox::commands::config::config_subcommand())
        .subcommand(evebox::commands::print::command())
        .subcommand(elastic_debug());

    let matches = parser.clone().get_matches();

    // Initialize logging.
    let log_level = {
        let verbosity = matches.occurrences_of("verbose");
        if verbosity > 1 {
            tracing::Level::TRACE
        } else if verbosity > 0 {
            tracing::Level::DEBUG
        } else {
            tracing::Level::INFO
        }
    };
    logger::init_logger(log_level);
    logger::init_stdlog();

    let rc: anyhow::Result<()> = match matches.subcommand() {
        Some(("server", args)) => {
            evebox::server::main(args).await?;
            Ok(())
        }
        Some(("version", _)) => {
            version::print_version();
            Ok(())
        }
        Some(("sqlite-import", args)) => evebox::commands::sqlite_import::main(args).await,
        Some(("elastic-import", args)) => {
            if let Err(err) = evebox::commands::elastic_import::main(args).await {
                error!("{}", err);
                std::process::exit(1);
            }
            Ok(())
        }
        Some(("elastic-debug", args)) => evebox::commands::elastic_debug::main(args).await,
        Some(("oneshot", args)) => evebox::commands::oneshot::main(args).await,
        Some(("agent", args)) => evebox::commands::agent::main(args).await,
        Some(("config", args)) => evebox::commands::config::main(args),
        Some(("print", args)) => evebox::commands::print::main(args),
        _ => {
            parser.print_help().ok();
            println!();
            std::process::exit(1);
        }
    };
    if let Err(err) = rc {
        error!("error: {}", err);
        std::process::exit(1);
    } else {
        Ok(())
    }
}

fn elastic_debug() -> clap::Command<'static> {
    Command::new("elastic-debug").arg(
        Arg::new("elasticsearch")
            .long("elasticsearch")
            .short('e')
            .takes_value(true)
            .default_value("127.0.0.1"),
    )
}
<commit_msg>server: fix some env vars<commit_after>// SPDX-License-Identifier: MIT
//
// Copyright (C) 2020-2022 Jason Ish

#![allow(clippy::redundant_field_names)]

use clap::{Arg, Command};
use evebox::logger;
use evebox::prelude::*;
use evebox::version;

#[tokio::main]
async fn main() {
    if let Err(err) = _main().await {
        error!("{}", err);
        std::process::exit(1);
    }
}

async fn _main() -> Result<(), Box<dyn std::error::Error>> {
    let parser = clap::Command::new("EveBox")
        .version(std::env!("CARGO_PKG_VERSION"))
        .arg(
            Arg::new("verbose")
                .long("verbose")
                .short('v')
                .multiple_occurrences(true)
                .global(true)
                .help("Increase verbosity"),
        )
        .arg(
            Arg::new("data-directory")
                .long("data-directory")
                .short('D')
                .takes_value(true)
                .value_name("DIR")
                .help("Data directory")
                .global(true),
        )
        .subcommand(clap::Command::new("version").about("Display version"));

    let sqlite_import = Command::new("sqlite-import")
        .about("Import to SQLite")
        .arg(
            Arg::new("oneshot")
                .long("oneshot")
                .help("One shot mode (exit on EOF"),
        )
        .arg(Arg::new("end").long("end").help("Start at end of file"))
        .arg(Arg::new("INPUT").required(true).index(1));

    let server = clap::Command::new("server")
        .about("EveBox Server")
        .arg(
            clap::Arg::new("config")
                .long("config")
                .short('c')
                .takes_value(true)
                .value_name("FILE")
                .help("Configuration filename"),
        )
        .arg(
            clap::Arg::new("http.host")
                .long("host")
                .value_name("HOSTNAME")
                .takes_value(true)
                .default_value("127.0.0.1")
                .env("EVEBOX_HTTP_HOST")
                .help("Hostname/IP address to bind to"),
        )
        .arg(
            clap::Arg::new("http.port")
                .short('p')
                .long("port")
                .value_name("PORT")
                .takes_value(true)
                .default_value("5636")
                .env("EVEBOX_HTTP_PORT")
                .help("Port to bind to"),
        )
        .arg(
            clap::Arg::new("database.elasticsearch.url")
                .short('e')
                .long("elasticsearch")
                .takes_value(true)
                .value_name("URL")
                .default_value("http://localhost:9200")
                .env("ELASTICSEARCH_URL")
                .help("Elastic Search URL"),
        )
        .arg(
            clap::Arg::new("database.elasticsearch.index")
                .short('i')
                .long("index")
                .takes_value(true)
                .default_value("logstash")
                .value_name("INDEX")
                .help("Elastic Search index prefix"),
        )
        .arg(
            Arg::new("database.elasticsearch.no-index-suffix")
                .long("no-index-suffix")
                .takes_value(false)
                .help("Do not add a suffix to the index name"),
        )
        .arg(
            Arg::new("database.elasticsearch.ecs")
                .long("ecs")
                .help("Enable Elastic ECS support"),
        )
        .arg(
            Arg::new("database.type")
                .long("database")
                .aliases(&["datastore"])
                .takes_value(true)
                .value_name("DATABASE")
                .default_value("elasticsearch")
                .help("Database type"),
        )
        .arg(
            Arg::new("sqlite-filename")
                .long("sqlite-filename")
                .takes_value(true)
                .value_name("FILENAME")
                .default_value("events.sqlite")
                .help("SQLite events filename"),
        )
        .arg(
            clap::Arg::new("no-check-certificate")
                .short('k')
                .long("no-check-certificate")
                .help("Disable TLS certificate validation"),
        )
        .arg(
            Arg::new("http.request-logging")
                .long("http-request-logging")
                .env("EVEBOX_HTTP_REQUEST_LOGGING")
                .hide_env(true)
                .help("Enable HTTP access logging"),
        )
        .arg(
            Arg::new("http.tls.enabled")
                .long("tls")
                .help("Enable TLS")
                .env("EVEBOX_HTTP_TLS_ENABLED")
                .hide_env(true),
        )
        .arg(
            Arg::new("http.tls.certificate")
                .long("tls-cert")
                .takes_value(true)
                .value_name("FILENAME")
                .help("TLS certificate filename"),
        )
        .arg(
            Arg::new("http.tls.key")
                .long("tls-key")
                .takes_value(true)
                .value_name("FILENAME")
                .help("TLS key filename"),
        )
        .arg(
            Arg::new("input.filename")
                .long("input")
                .takes_value(true)
                .value_name("FILENAME")
                .help("Input Eve file to read"),
        )
        .arg(
            Arg::new("end")
                .long("end")
                .help("Read from end (tail) of input file"),
        )
        .arg(Arg::new("input-start").long("input-start").hide(true));

    let oneshot = Command::new("oneshot")
        .about("Import a single eve.json and review in EveBox")
        .arg(
            Arg::new("limit")
                .long("limit")
                .takes_value(true)
                .help("Limit the number of events read"),
        )
        .arg(
            Arg::new("no-open")
                .long("no-open")
                .help("Don't open browser"),
        )
        .arg(
            Arg::new("no-wait")
                .long("no-wait")
                .help("Don't wait for events to load"),
        )
        .arg(
            Arg::new("database-filename")
                .long("database-filename")
                .takes_value(true)
                .default_value("./oneshot.sqlite")
                .help("Database filename"),
        )
        // --host, but keep th name as http.host to be campatible with the
        // EVEBOX_HTTP_HOST environment variable.
        .arg(
            clap::Arg::new("http.host")
                .long("host")
                .value_name("HOSTNAME")
                .takes_value(true)
                .default_value("127.0.0.1")
                .help("Hostname/IP address to bind to"),
        )
        .arg(Arg::new("INPUT").required(true).index(1));

    let elastic_import = Command::new("elastic-import")
        .about("Import to Elastic Search")
        .arg(
            Arg::new("config")
                .short('c')
                .long("config")
                .takes_value(true)
                .help("Configuration file"),
        )
        .arg(
            Arg::new("oneshot")
                .long("oneshot")
                .help("One shot mode (exit on EOF)"),
        )
        .arg(Arg::new("end").long("end").help("Start at end of file"))
        .arg(
            clap::Arg::new("elasticsearch")
                .short('e')
                .long("elasticsearch")
                .takes_value(true)
                .default_value("http://localhost:9200")
                .hide_default_value(true)
                .help("Elastic Search URL"),
        )
        .arg(
            Arg::new("index")
                .long("index")
                .takes_value(true)
                .default_value("logstash")
                .hide_default_value(true)
                .help("Elastic Search index prefix"),
        )
        .arg(
            Arg::new("no-index-suffix")
                .long("no-index-suffix")
                .takes_value(false)
                .help("Do not add a suffix to the index name"),
        )
        .arg(
            Arg::new("bookmark")
                .long("bookmark")
                .help("Enable bookmarking"),
        )
        .arg(
            Arg::new("bookmark-filename")
                .long("bookmark-filename")
                .takes_value(true)
                .default_value("")
                .hide_default_value(true)
                .help("Bookmark filename"),
        )
        .arg(
            Arg::new("bookmark-dir")
                .long("bookmark-dir")
                .takes_value(true)
                .default_value(".")
                .hide_default_value(true)
                .help("Bookmark directory"),
        )
        .arg(
            Arg::new("stdout")
                .long("stdout")
                .help("Print events to stdout"),
        )
        .arg(
            Arg::new("username")
                .long("username")
                .short('u')
                .takes_value(true)
                .help("Elasticsearch username"),
        )
        .arg(
            Arg::new("password")
                .long("password")
                .short('p')
                .takes_value(true)
                .help("Elasticsearch password"),
        )
        .arg(
            clap::Arg::new(evebox::commands::elastic_import::NO_CHECK_CERTIFICATE)
                .short('k')
                .long("no-check-certificate")
                .help("Disable TLS certificate validation"),
        )
        .arg(
            Arg::new("geoip.disabled")
                .long("no-geoip")
                .help("Disable GeoIP"),
        )
        .arg(
            Arg::new("geoip.database-filename")
                .long("geoip-database")
                .takes_value(true)
                .value_name("filename")
                .help("GeoIP database filename"),
        )
        .arg(
            Arg::new("input")
                .multiple_values(true)
                .value_name("INPUT")
                .index(1),
        );

    let mut parser = parser
        .subcommand(server)
        .subcommand(elastic_import)
        .subcommand(oneshot)
        .subcommand(evebox::commands::agent::command())
        .subcommand(sqlite_import)
        .subcommand(evebox::commands::config::config_subcommand())
        .subcommand(evebox::commands::print::command())
        .subcommand(elastic_debug());

    let matches = parser.clone().get_matches();

    // Initialize logging.
    let log_level = {
        let verbosity = matches.occurrences_of("verbose");
        if verbosity > 1 {
            tracing::Level::TRACE
        } else if verbosity > 0 {
            tracing::Level::DEBUG
        } else {
            tracing::Level::INFO
        }
    };
    logger::init_logger(log_level);
    logger::init_stdlog();

    let rc: anyhow::Result<()> = match matches.subcommand() {
        Some(("server", args)) => {
            evebox::server::main(args).await?;
            Ok(())
        }
        Some(("version", _)) => {
            version::print_version();
            Ok(())
        }
        Some(("sqlite-import", args)) => evebox::commands::sqlite_import::main(args).await,
        Some(("elastic-import", args)) => {
            if let Err(err) = evebox::commands::elastic_import::main(args).await {
                error!("{}", err);
                std::process::exit(1);
            }
            Ok(())
        }
        Some(("elastic-debug", args)) => evebox::commands::elastic_debug::main(args).await,
        Some(("oneshot", args)) => evebox::commands::oneshot::main(args).await,
        Some(("agent", args)) => evebox::commands::agent::main(args).await,
        Some(("config", args)) => evebox::commands::config::main(args),
        Some(("print", args)) => evebox::commands::print::main(args),
        _ => {
            parser.print_help().ok();
            println!();
            std::process::exit(1);
        }
    };
    if let Err(err) = rc {
        error!("error: {}", err);
        std::process::exit(1);
    } else {
        Ok(())
    }
}

fn elastic_debug() -> clap::Command<'static> {
    Command::new("elastic-debug").arg(
        Arg::new("elasticsearch")
            .long("elasticsearch")
            .short('e')
            .takes_value(true)
            .default_value("127.0.0.1"),
    )
}
<|endoftext|>


<commit_before>#![allow(dead_code)]
use utils;
//use std::ops::{Index, Deref, DerefMut};
use std::cell::RefCell;
use std::rc::Rc;

pub type MemShared = Rc<RefCell<Memory>>;

pub enum MemType
{
    RAM,
    BASIC,
    CHARGEN,
    IO,
    KERNAL,
}

// specific memory bank - RAM, ROM, IO
struct MemBank
{
    bank_type: MemType, // what am I?
    read_only: bool,    // RAM or ROM?
    offset: u16,        // offset from start of address space
    data: Vec<u8>,
}

impl MemBank
{
    pub fn new(mem_type: MemType) -> MemBank
    {
        let mut mem_bank = MemBank
        {
            bank_type: mem_type,
            read_only: true,
            offset: 0x0000,
            data: Vec::<u8>::new(),
        };

        match mem_bank.bank_type
        {
            MemType::BASIC   => {
                mem_bank.data = utils::open_file("rom/basic.rom", 0);
                mem_bank.offset = 0xA000;
            },
            MemType::CHARGEN => {
                mem_bank.data = utils::open_file("rom/chargen.rom", 0);
                mem_bank.offset = 0xD000;
            },
            MemType::KERNAL  => {
                mem_bank.data = utils::open_file("rom/kernal.rom", 0);
                mem_bank.offset = 0xE000;
            },
            MemType::RAM => {
                mem_bank.data = Vec::<u8>::with_capacity(0x10000);
                for _ in 0..0x10000
                {
                    mem_bank.data.push(0);
                }

                mem_bank.read_only = false;
            },
            MemType::IO => {
                mem_bank.data = Vec::<u8>::with_capacity(0x1000);
                for _ in 0..0x1000
                {
                    mem_bank.data.push(0);
                }

                mem_bank.offset = 0xD000;
                mem_bank.read_only = false;
            }
        }

        mem_bank
    }

    pub fn write(&mut self, addr: u16, val: u8)
    {
        match self.bank_type
        {
            MemType::RAM => self.data[(addr - self.offset) as usize] = val,
            MemType::IO => {
                // TODO: IO access has specific behavior depending on address
                match addr
                {
                    0xD016          => self.data[(addr - self.offset) as usize] = 0xC0 | val,
                    0xD019          => self.data[(addr - self.offset) as usize] = 0x70 | val,
                    0xD01A          => self.data[(addr - self.offset) as usize] = 0xF0 | val,
                    //0xD01E...0xD01F => (),  // cannot be written on real C64 but allow the VIC to do it anyway
                    0xD020...0xD02E => self.data[(addr - self.offset) as usize] = 0xF0 | val,
                    0xD02F...0xD03F => (),                             // write ignored
                    0xD040...0xD3FF => self.write(0xD000 + (addr % 0x0040), val), // same as 0xD000-0xD03F
                    _ => self.data[(addr - self.offset) as usize] = val
                }

            },
            _ => panic!("Can't write to ROM!")
        }
    }

    pub fn read(&mut self, addr: u16) -> u8
    {
        match self.bank_type
        {
            MemType::IO => {
                // TODO: IO access has specific behavior depending on address
                match addr
                {
                    0xD016          => 0xC0 | self.data[(addr - self.offset) as usize],
                    0xD018          => 0x01 | self.data[(addr - self.offset) as usize],
                    0xD019          => 0x70 | self.data[(addr - self.offset) as usize],
                    0xD01A          => 0xF0 | self.data[(addr - self.offset) as usize],
                    0xD01E...0xD01F => {                                  // cannot be written, cleared on read
                        let value = self.data[(addr - self.offset) as usize];
                        self.data[(addr - self.offset) as usize] = 0;
                        value
                    },
                    0xD020...0xD02E => 0xF0 | self.data[(addr - self.offset) as usize],
                    0xD02F...0xD03F => 0xFF,                              // always returns 0xFF
                    0xD040...0xD3FF => self.read(0xD000 + (addr % 0x0040)),          // same as 0xD000-0xD03F
                    _ => self.data[(addr - self.offset) as usize]
                }
            },
            _ => self.data[(addr - self.offset) as usize]
        }
    }
}

/*impl Index<u16> for MemBank
{
    type Output = u8;

    fn index<'a>(&'a self, _index: u16) -> &'a u8
    {
        &self.data[_index as usize]
    }
}

impl Deref for MemBank
{
    type Target = Vec<u8>;

    fn deref(&self) -> &Vec<u8>
    {
        &self.data
    }
}

impl DerefMut for MemBank
{
    fn deref_mut(&mut self) -> &mut Vec<u8>
    {
        &mut self.data
    }
}*/

// collective memory storage with all the banks and bank switching support
pub struct Memory
{
    ram:     MemBank,
    basic:   MemBank,
    chargen: MemBank,
    io:      MemBank,
    kernal:  MemBank,

    // bank switching flags
    pub basic_on:   bool,
    pub chargen_on: bool,
    pub io_on:  bool,
    pub kernal_on:  bool,
    cart_lo_on: bool, // cart flag - unused for now
    cart_hi_on: bool  // cart flag - unused for now
}

impl Memory
{
    pub fn new_shared() -> MemShared
    {
        Rc::new(RefCell::new(Memory
        {
            ram:     MemBank::new(MemType::RAM),     // 64k
            basic:   MemBank::new(MemType::BASIC),   // 8k
            chargen: MemBank::new(MemType::CHARGEN), // 4k
            io:      MemBank::new(MemType::IO),      // 4k (VIC, SID, CIA, Color RAM)
            kernal:  MemBank::new(MemType::KERNAL),  // 8k
            basic_on:   false,
            chargen_on: false,
            io_on:      false,
            kernal_on:  false,
            cart_lo_on: false, // unused for now
            cart_hi_on: false, // unused for now
        }))
    }

    // returns memory bank for current latch setting and address
    pub fn get_bank(&mut self, addr: u16) -> (&mut MemBank)
    {
        match addr
        {
            0x0000...0x9FFF => &mut self.ram,
            0xA000...0xBFFF => if self.basic_on { &mut self.basic } else { &mut self.ram },
            0xC000...0xCFFF => &mut self.ram,
            0xD000...0xDFFF => {
                if self.chargen_on { return &mut self.chargen }
                if self.io_on      { return &mut self.io; }
                return &mut self.ram;
            },
            0xE000...0xFFFF => if self.kernal_on  { &mut self.kernal } else { &mut self.ram },
            _ => panic!("Address out of memory range")
        }
    }

    // returns specific modifiable memory bank
    pub fn get_ram_bank(&mut self, bank_type: MemType) -> (&mut MemBank)
    {
        match bank_type
        {
            MemType::RAM => &mut self.ram,
            MemType::IO  => &mut self.io,
            _            => panic!("Unrecognized RAM bank"),
        }
    }

    // returns specific non-modifiable memory bank
    pub fn get_rom_bank(&mut self, bank_type: MemType) -> (&mut MemBank)
    {
        match bank_type
        {
            MemType::BASIC   => &mut self.basic,
            MemType::CHARGEN => &mut self.chargen,
            MemType::KERNAL  => &mut self.kernal,
            _                => panic!("Unrecognized ROM Abank"),
        }
    }

    pub fn reset(&mut self)
    {
        self.write_byte(0x0000, 0xFF);
        self.write_byte(0x0001, 0x07); // enable kernal, chargen and basic ROMs
    }

    fn update_bank_flags(&mut self)
    {
        // latch state is determined by 3 least significant bits from this location
        let latch = self.ram.read(0x0001) & 0x07;

        self.chargen_on = ((latch & 0x04) == 0) && ((latch & 0x03) != 0); // %0xx except %000
        self.io_on      = ((latch & 0x04) != 0) && ((latch & 0x03) != 0); // %1xx except %100
        self.basic_on   = (latch & 0x03) == 0x03;
        self.kernal_on  = self.basic_on || ((latch & 0x03) == 0x02);
    }

    // Write a byte to memory - returns whether RAM was written (true) or RAM under ROM (false)
    pub fn write_byte(&mut self, addr: u16, value: u8) -> bool
    {
        // RAM under ROM written? Return false to let us know about it
        if self.get_bank(addr).read_only
        {
            self.ram.write(addr, value);
            return false;
        }
        else
        {
            self.get_bank(addr).write(addr, value);
        }

        // update the bank switching flags here, since they can only change on memory write
        // latch byte changed - update bank switching flags
        if addr == 0x0001 { self.update_bank_flags(); }

        return true;
    }

    // Read a byte from memory
    pub fn read_byte(&mut self, addr: u16) -> u8
    {
        self.get_bank(addr).read(addr)
    }

    // Read a word from memory (stored in little endian)
    pub fn read_word_le(&mut self, addr: u16) -> u16
    {
        let bank = self.get_bank(addr);
        let value_be: u16 = ((bank.read(addr) as u16) << 8 & 0xFF00) |
                            ((bank.read(addr + 0x0001) as u16) & 0x00FF);

        let value_le: u16 = ((value_be << 8) & 0xFF00) | ((value_be >> 8) & 0x00FF);
        value_le
    }

    // Read a word from memory (stored in big endian: swap low<->high)
    pub fn read_word_be(&mut self, addr: u16) -> u16
    {
        let bank = self.get_bank(addr);
        let value_le: u16 = ((bank.read(addr) as u16) << 8 & 0xFF00) |
                            ((bank.read(addr + 0x0001) as u16) & 0x00FF);
        value_le
    }

    // Write word in little endian format (low/high)
    pub fn write_word_le(&mut self, addr: u16, value: u16) -> bool
    {
        let value_le_lo: u8 = (((value << 8) & 0xFF00) >> 8 & 0xFF) as u8;
        let value_le_hi: u8 = ((value >> 8) & 0x00FF) as u8;

        let hi = self.write_byte(addr, value_le_lo);
        let lo = self.write_byte(addr + 0x0001, value_le_hi);

        return hi && lo;
    }

    // Write word in big endian format (high/low)
    pub fn write_word_be(&mut self, addr: u16, value: u16) -> bool
    {
        let value_le_lo: u8 = (((value << 8) & 0xFF00) >> 8 & 0xFF) as u8;
        let value_le_hi: u8 = ((value >> 8) & 0x00FF) as u8;

        let hi = self.write_byte(addr, value_le_hi);
        let lo = self.write_byte(addr + 0x0001, value_le_lo);

        return hi && lo;
    }
}

<commit_msg>fix problem with memory latch not being set correctly (inifinte loops in Fire Trap and 1942)<commit_after>#![allow(dead_code)]
use utils;
//use std::ops::{Index, Deref, DerefMut};
use std::cell::RefCell;
use std::rc::Rc;

pub type MemShared = Rc<RefCell<Memory>>;

pub enum MemType
{
    RAM,
    BASIC,
    CHARGEN,
    IO,
    KERNAL,
}

// specific memory bank - RAM, ROM, IO
struct MemBank
{
    bank_type: MemType, // what am I?
    read_only: bool,    // RAM or ROM?
    offset: u16,        // offset from start of address space
    data: Vec<u8>,
}

impl MemBank
{
    pub fn new(mem_type: MemType) -> MemBank
    {
        let mut mem_bank = MemBank
        {
            bank_type: mem_type,
            read_only: true,
            offset: 0x0000,
            data: Vec::<u8>::new(),
        };

        match mem_bank.bank_type
        {
            MemType::BASIC   => {
                mem_bank.data = utils::open_file("rom/basic.rom", 0);
                mem_bank.offset = 0xA000;
            },
            MemType::CHARGEN => {
                mem_bank.data = utils::open_file("rom/chargen.rom", 0);
                mem_bank.offset = 0xD000;
            },
            MemType::KERNAL  => {
                mem_bank.data = utils::open_file("rom/kernal.rom", 0);
                mem_bank.offset = 0xE000;
            },
            MemType::RAM => {
                mem_bank.data = Vec::<u8>::with_capacity(0x10000);
                for _ in 0..0x10000
                {
                    mem_bank.data.push(0);
                }

                mem_bank.read_only = false;
            },
            MemType::IO => {
                mem_bank.data = Vec::<u8>::with_capacity(0x1000);
                for _ in 0..0x1000
                {
                    mem_bank.data.push(0);
                }

                mem_bank.offset = 0xD000;
                mem_bank.read_only = false;
            }
        }

        mem_bank
    }

    pub fn write(&mut self, addr: u16, val: u8)
    {
        match self.bank_type
        {
            MemType::RAM => self.data[(addr - self.offset) as usize] = val,
            MemType::IO => {
                // TODO: IO access has specific behavior depending on address
                match addr
                {
                    0xD016          => self.data[(addr - self.offset) as usize] = 0xC0 | val,
                    0xD019          => self.data[(addr - self.offset) as usize] = 0x70 | val,
                    0xD01A          => self.data[(addr - self.offset) as usize] = 0xF0 | val,
                    //0xD01E...0xD01F => (),  // cannot be written on real C64 but allow the VIC to do it anyway
                    0xD020...0xD02E => self.data[(addr - self.offset) as usize] = 0xF0 | val,
                    0xD02F...0xD03F => (),                             // write ignored
                    0xD040...0xD3FF => self.write(0xD000 + (addr % 0x0040), val), // same as 0xD000-0xD03F
                    _ => self.data[(addr - self.offset) as usize] = val
                }

            },
            _ => panic!("Can't write to ROM!")
        }
    }

    pub fn read(&mut self, addr: u16) -> u8
    {
        match self.bank_type
        {
            MemType::IO => {
                // TODO: IO access has specific behavior depending on address
                match addr
                {
                    0xD016          => 0xC0 | self.data[(addr - self.offset) as usize],
                    0xD018          => 0x01 | self.data[(addr - self.offset) as usize],
                    0xD019          => 0x70 | self.data[(addr - self.offset) as usize],
                    0xD01A          => 0xF0 | self.data[(addr - self.offset) as usize],
                    0xD01E...0xD01F => {                                  // cannot be written, cleared on read
                        let value = self.data[(addr - self.offset) as usize];
                        self.data[(addr - self.offset) as usize] = 0;
                        value
                    },
                    0xD020...0xD02E => 0xF0 | self.data[(addr - self.offset) as usize],
                    0xD02F...0xD03F => 0xFF,                              // always returns 0xFF
                    0xD040...0xD3FF => self.read(0xD000 + (addr % 0x0040)),          // same as 0xD000-0xD03F
                    _ => self.data[(addr - self.offset) as usize]
                }
            },
            _ => self.data[(addr - self.offset) as usize]
        }
    }
}

/*impl Index<u16> for MemBank
{
    type Output = u8;

    fn index<'a>(&'a self, _index: u16) -> &'a u8
    {
        &self.data[_index as usize]
    }
}

impl Deref for MemBank
{
    type Target = Vec<u8>;

    fn deref(&self) -> &Vec<u8>
    {
        &self.data
    }
}

impl DerefMut for MemBank
{
    fn deref_mut(&mut self) -> &mut Vec<u8>
    {
        &mut self.data
    }
}*/

// collective memory storage with all the banks and bank switching support
pub struct Memory
{
    ram:     MemBank,
    basic:   MemBank,
    chargen: MemBank,
    io:      MemBank,
    kernal:  MemBank,

    // bank switching flags
    pub basic_on:   bool,
    pub chargen_on: bool,
    pub io_on:  bool,
    pub kernal_on:  bool,
    cart_lo_on: bool, // cart flag - unused for now
    cart_hi_on: bool  // cart flag - unused for now
}

impl Memory
{
    pub fn new_shared() -> MemShared
    {
        Rc::new(RefCell::new(Memory
        {
            ram:     MemBank::new(MemType::RAM),     // 64k
            basic:   MemBank::new(MemType::BASIC),   // 8k
            chargen: MemBank::new(MemType::CHARGEN), // 4k
            io:      MemBank::new(MemType::IO),      // 4k (VIC, SID, CIA, Color RAM)
            kernal:  MemBank::new(MemType::KERNAL),  // 8k
            basic_on:   false,
            chargen_on: false,
            io_on:      false,
            kernal_on:  false,
            cart_lo_on: false, // unused for now
            cart_hi_on: false, // unused for now
        }))
    }

    // returns memory bank for current latch setting and address
    pub fn get_bank(&mut self, addr: u16) -> (&mut MemBank)
    {
        match addr
        {
            0x0000...0x9FFF => &mut self.ram,
            0xA000...0xBFFF => if self.basic_on { &mut self.basic } else { &mut self.ram },
            0xC000...0xCFFF => &mut self.ram,
            0xD000...0xDFFF => {
                if self.chargen_on { return &mut self.chargen }
                if self.io_on      { return &mut self.io; }
                return &mut self.ram;
            },
            0xE000...0xFFFF => if self.kernal_on  { &mut self.kernal } else { &mut self.ram },
            _ => panic!("Address out of memory range")
        }
    }

    // returns specific modifiable memory bank
    pub fn get_ram_bank(&mut self, bank_type: MemType) -> (&mut MemBank)
    {
        match bank_type
        {
            MemType::RAM => &mut self.ram,
            MemType::IO  => &mut self.io,
            _            => panic!("Unrecognized RAM bank"),
        }
    }

    // returns specific non-modifiable memory bank
    pub fn get_rom_bank(&mut self, bank_type: MemType) -> (&mut MemBank)
    {
        match bank_type
        {
            MemType::BASIC   => &mut self.basic,
            MemType::CHARGEN => &mut self.chargen,
            MemType::KERNAL  => &mut self.kernal,
            _                => panic!("Unrecognized ROM Abank"),
        }
    }

    pub fn reset(&mut self)
    {
        self.write_byte(0x0000, 0xFF);
        self.write_byte(0x0001, 0x07); // enable kernal, chargen and basic ROMs
    }

    fn update_memory_latch(&mut self)
    {
        let ddr = self.ram.read(0x0000);
        let pr  = self.ram.read(0x0001);
        let latch = !ddr | pr;

        self.chargen_on = ((latch & 0x04) == 0) && ((latch & 0x03) != 0); // %0xx except %000
        self.io_on      = ((latch & 0x04) != 0) && ((latch & 0x03) != 0); // %1xx except %100
        self.basic_on   = (latch & 0x03) == 3;
        self.kernal_on  = (latch & 0x02) != 0;
    }

    // Write a byte to memory - returns whether RAM was written (true) or RAM under ROM (false)
    pub fn write_byte(&mut self, addr: u16, value: u8) -> bool
    {
        // RAM under ROM written? Return false to let us know about it
        if self.get_bank(addr).read_only
        {
            self.ram.write(addr, value);
            return false;
        }
        else
        {
            self.get_bank(addr).write(addr, value);
        }

        // update the bank switching flags here, since they can only change on memory write
        // latch byte changed - update bank switching flags
        if addr < 0x0002 { self.update_memory_latch(); }

        return true;
    }

    // Read a byte from memory
    pub fn read_byte(&mut self, addr: u16) -> u8
    {
        // special location: current memory latch settings
        if addr == 0x0001
        {
            let ddr = self.ram.read(0x0000);
            let pr  = self.ram.read(0x0001);
            return (ddr & pr) | (!ddr & 0x17);
        }

        self.get_bank(addr).read(addr)
    }

    // Read a word from memory (stored in little endian)
    pub fn read_word_le(&mut self, addr: u16) -> u16
    {
        let bank = self.get_bank(addr);
        let value_be: u16 = ((bank.read(addr) as u16) << 8 & 0xFF00) |
                            ((bank.read(addr + 0x0001) as u16) & 0x00FF);

        let value_le: u16 = ((value_be << 8) & 0xFF00) | ((value_be >> 8) & 0x00FF);
        value_le
    }

    // Read a word from memory (stored in big endian: swap low<->high)
    pub fn read_word_be(&mut self, addr: u16) -> u16
    {
        let bank = self.get_bank(addr);
        let value_le: u16 = ((bank.read(addr) as u16) << 8 & 0xFF00) |
                            ((bank.read(addr + 0x0001) as u16) & 0x00FF);
        value_le
    }

    // Write word in little endian format (low/high)
    pub fn write_word_le(&mut self, addr: u16, value: u16) -> bool
    {
        let value_le_lo: u8 = (((value << 8) & 0xFF00) >> 8 & 0xFF) as u8;
        let value_le_hi: u8 = ((value >> 8) & 0x00FF) as u8;

        let hi = self.write_byte(addr, value_le_lo);
        let lo = self.write_byte(addr + 0x0001, value_le_hi);

        return hi && lo;
    }

    // Write word in big endian format (high/low)
    pub fn write_word_be(&mut self, addr: u16, value: u16) -> bool
    {
        let value_le_lo: u8 = (((value << 8) & 0xFF00) >> 8 & 0xFF) as u8;
        let value_le_hi: u8 = ((value >> 8) & 0x00FF) as u8;

        let hi = self.write_byte(addr, value_le_hi);
        let lo = self.write_byte(addr + 0x0001, value_le_lo);

        return hi && lo;
    }
}

<|endoftext|>


<commit_before>use crate::{
    channel::Channel,
    channels::Channels,
    configuration::Configuration,
    connection_closer::ConnectionCloser,
    connection_properties::ConnectionProperties,
    connection_status::{ConnectionState, ConnectionStatus, ConnectionStep},
    executor::{within_executor, DefaultExecutor, Executor},
    frames::Frames,
    internal_rpc::{InternalRPC, InternalRPCHandle},
    io_loop::IoLoop,
    reactor::DefaultReactorBuilder,
    socket_state::{SocketState, SocketStateHandle},
    tcp::{AMQPUriTcpExt, HandshakeResult, OwnedTLSConfig},
    thread::ThreadHandle,
    types::ShortUInt,
    uri::AMQPUri,
    Error, Promise, Result,
};
use amq_protocol::frame::{AMQPFrame, ProtocolVersion};
use async_trait::async_trait;
use log::{log_enabled, Level::Trace};
use std::{fmt, io, sync::Arc};

pub struct Connection {
    configuration: Configuration,
    status: ConnectionStatus,
    channels: Channels,
    io_loop: ThreadHandle,
    closer: Arc<ConnectionCloser>,
}

impl Connection {
    fn new(
        waker: SocketStateHandle,
        internal_rpc: InternalRPCHandle,
        frames: Frames,
        executor: Arc<dyn Executor>,
    ) -> Self {
        let configuration = Configuration::default();
        let status = ConnectionStatus::default();
        let channels = Channels::new(
            configuration.clone(),
            status.clone(),
            waker,
            internal_rpc.clone(),
            frames,
            executor,
        );
        let closer = Arc::new(ConnectionCloser::new(status.clone(), internal_rpc));
        let connection = Self {
            configuration,
            status,
            channels,
            io_loop: ThreadHandle::default(),
            closer,
        };

        connection.channels.create_zero();
        connection
    }

    /// Connect to an AMQP Server
    pub async fn connect(uri: &str, options: ConnectionProperties) -> Result<Connection> {
        Connect::connect(uri, options, OwnedTLSConfig::default()).await
    }

    /// Connect to an AMQP Server
    pub async fn connect_with_config(
        uri: &str,
        options: ConnectionProperties,
        config: OwnedTLSConfig,
    ) -> Result<Connection> {
        Connect::connect(uri, options, config).await
    }

    /// Connect to an AMQP Server
    pub async fn connect_uri(uri: AMQPUri, options: ConnectionProperties) -> Result<Connection> {
        Connect::connect(uri, options, OwnedTLSConfig::default()).await
    }

    /// Connect to an AMQP Server
    pub async fn connect_uri_with_identity(
        uri: AMQPUri,
        options: ConnectionProperties,
        config: OwnedTLSConfig,
    ) -> Result<Connection> {
        Connect::connect(uri, options, config).await
    }

    pub async fn create_channel(&self) -> Result<Channel> {
        if !self.status.connected() {
            return Err(Error::InvalidConnectionState(self.status.state()));
        }
        let channel = self.channels.create(self.closer.clone())?;
        channel.clone().channel_open(channel).await
    }

    /// Block current thread while the connection is still active.
    /// This is useful when you only have a consumer and nothing else keeping your application
    /// "alive".
    pub fn run(self) -> Result<()> {
        let io_loop = self.io_loop.clone();
        drop(self);
        if within_executor() {
            // io_loop waits for the executor to shutdown, do not deadlock in the
            // unlikely case where run is called from within it
            Ok(())
        } else {
            io_loop.wait("io loop")
        }
    }

    pub fn on_error<E: FnMut(Error) + Send + 'static>(&self, handler: E) {
        self.channels.set_error_handler(handler);
    }

    pub fn configuration(&self) -> &Configuration {
        &self.configuration
    }

    pub fn status(&self) -> &ConnectionStatus {
        &self.status
    }

    pub async fn close(&self, reply_code: ShortUInt, reply_text: &str) -> Result<()> {
        if let Some(channel0) = self.channels.get(0) {
            channel0
                .connection_close(reply_code, reply_text, 0, 0)
                .await
        } else {
            Ok(())
        }
    }

    /// Block all consumers and publishers on this connection
    pub async fn block(&self, reason: &str) -> Result<()> {
        if let Some(channel0) = self.channels.get(0) {
            channel0.connection_blocked(reason).await
        } else {
            Err(Error::InvalidConnectionState(self.status.state()))
        }
    }

    /// Unblock all consumers and publishers on this connection
    pub async fn unblock(&self) -> Result<()> {
        if let Some(channel0) = self.channels.get(0) {
            channel0.connection_unblocked().await
        } else {
            Err(Error::InvalidConnectionState(self.status.state()))
        }
    }

    /// Update the secret used by some authentication module such as OAuth2
    pub async fn update_secret(&self, new_secret: &str, reason: &str) -> Result<()> {
        if let Some(channel0) = self.channels.get(0) {
            channel0.connection_update_secret(new_secret, reason).await
        } else {
            Err(Error::InvalidConnectionState(self.status.state()))
        }
    }

    pub async fn connector(
        uri: AMQPUri,
        connect: Box<dyn FnOnce(&AMQPUri) -> HandshakeResult + Send + Sync>,
        mut options: ConnectionProperties,
    ) -> Result<Connection> {
        let executor = options
            .executor
            .take()
            .map(Ok)
            .unwrap_or_else(|| DefaultExecutor::default())?;

        let (connect_promise, resolver) = pinky_swear::PinkySwear::<HandshakeResult>::new();
        let connect_uri = uri.clone();
        executor.spawn_blocking(Box::new(move || {
            resolver.swear(connect(&connect_uri));
        }))?;

        let reactor_builder = options
            .reactor_builder
            .take()
            .unwrap_or_else(|| Arc::new(DefaultReactorBuilder));
        let socket_state = SocketState::default();
        let waker = socket_state.handle();
        let internal_rpc = InternalRPC::new(executor.clone(), waker.clone());
        let frames = Frames::default();
        let conn = Connection::new(waker, internal_rpc.handle(), frames.clone(), executor);
        let status = conn.status.clone();
        let configuration = conn.configuration.clone();
        status.set_vhost(&uri.vhost);
        status.set_username(&uri.authority.userinfo.username);
        if let Some(frame_max) = uri.query.frame_max {
            configuration.set_frame_max(frame_max);
        }
        if let Some(channel_max) = uri.query.channel_max {
            configuration.set_channel_max(channel_max);
        }
        if let Some(heartbeat) = uri.query.heartbeat {
            configuration.set_heartbeat(heartbeat);
        }
        let (promise_out, resolver) = Promise::new();
        if log_enabled!(Trace) {
            promise_out.set_marker("ProtocolHeader".into());
        }
        let channels = conn.channels.clone();
        if let Some(channel0) = channels.get(0) {
            channel0.send_frame(
                AMQPFrame::ProtocolHeader(ProtocolVersion::amqp_0_9_1()),
                resolver,
                None,
            )
        };
        let (promise_in, resolver) = Promise::new();
        if log_enabled!(Trace) {
            promise_in.set_marker("ProtocolHeader.Ok".into());
        }
        let io_loop_handle = conn.io_loop.clone();
        status.set_state(ConnectionState::Connecting);
        status.set_connection_step(ConnectionStep::ProtocolHeader(
            resolver,
            conn,
            uri.authority.userinfo.into(),
            uri.query.auth_mechanism.unwrap_or_default(),
            options,
        ));
        let handshake_result = connect_promise.await;
        IoLoop::new(
            status,
            configuration,
            channels,
            internal_rpc,
            frames,
            socket_state,
            io_loop_handle,
            handshake_result,
            &*reactor_builder,
        )
        .and_then(IoLoop::start)?;
        promise_out.await?;
        promise_in.await
    }
}

impl fmt::Debug for Connection {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Connection")
            .field("configuration", &self.configuration)
            .field("status", &self.status)
            .field("channels", &self.channels)
            .finish()
    }
}

/// Trait providing a method to connect to an AMQP server
#[async_trait]
pub trait Connect {
    /// connect to an AMQP server
    async fn connect(
        self,
        options: ConnectionProperties,
        config: OwnedTLSConfig,
    ) -> Result<Connection>;
}

#[async_trait]
impl Connect for AMQPUri {
    async fn connect(
        self,
        options: ConnectionProperties,
        config: OwnedTLSConfig
    ) -> Result<Connection> {
        Connection::connector(
            self,
            Box::new(move |uri| AMQPUriTcpExt::connect_with_config(uri, config.as_ref())),
            options,
        )
        .await
    }
}

#[async_trait]
impl Connect for &str {
    async fn connect(
        self,
        options: ConnectionProperties,
        config: OwnedTLSConfig
    ) -> Result<Connection> {
        match self.parse::<AMQPUri>() {
            Ok(uri) => Connect::connect(uri, options, config).await,
            Err(err) => Err(io::Error::new(io::ErrorKind::Other, err).into()),
        }
    }
}

#[cfg(test)]
mod tests {
    use env_logger;

    use super::*;
    use crate::channel_receiver_state::ChannelReceiverState;
    use crate::channel_status::ChannelState;
    use crate::types::ShortString;
    use crate::BasicProperties;
    use amq_protocol::frame::AMQPContentHeader;
    use amq_protocol::protocol::{basic, AMQPClass};

    #[test]
    fn basic_consume_small_payload() {
        let _ = env_logger::try_init();

        use crate::consumer::Consumer;
        use crate::queue::{Queue, QueueState};

        // Bootstrap connection state to a consuming state
        let executor = DefaultExecutor::default().unwrap();
        let socket_state = SocketState::default();
        let waker = socket_state.handle();
        let internal_rpc = InternalRPC::new(executor.clone(), waker.clone());
        let conn = Connection::new(
            waker,
            internal_rpc.handle(),
            Frames::default(),
            executor.clone(),
        );
        conn.status.set_state(ConnectionState::Connected);
        conn.configuration.set_channel_max(2047);
        let channel = conn.channels.create(conn.closer.clone()).unwrap();
        channel.set_state(ChannelState::Connected);
        let queue_name = ShortString::from("consumed");
        let mut queue: QueueState = Queue::new(queue_name.clone(), 0, 0).into();
        let consumer_tag = ShortString::from("consumer-tag");
        let consumer = Consumer::new(consumer_tag.clone(), executor);
        queue.register_consumer(consumer_tag.clone(), consumer);
        conn.channels
            .get(channel.id())
            .map(|c| c.register_queue(queue));
        // Now test the state machine behaviour
        {
            let method = AMQPClass::Basic(basic::AMQPMethod::Deliver(basic::Deliver {
                consumer_tag: consumer_tag.clone(),
                delivery_tag: 1,
                redelivered: false,
                exchange: "".into(),
                routing_key: queue_name.clone(),
            }));
            let class_id = method.get_amqp_class_id();
            let deliver_frame = AMQPFrame::Method(channel.id(), method);
            conn.channels.handle_frame(deliver_frame).unwrap();
            let channel_state = channel.status().receiver_state();
            let expected_state = ChannelReceiverState::WillReceiveContent(
                class_id,
                Some(queue_name.clone()),
                Some(consumer_tag.clone()),
            );
            assert_eq!(channel_state, expected_state);
        }
        {
            let header_frame = AMQPFrame::Header(
                channel.id(),
                60,
                Box::new(AMQPContentHeader {
                    class_id: 60,
                    weight: 0,
                    body_size: 2,
                    properties: BasicProperties::default(),
                }),
            );
            conn.channels.handle_frame(header_frame).unwrap();
            let channel_state = channel.status().receiver_state();
            let expected_state = ChannelReceiverState::ReceivingContent(
                Some(queue_name.clone()),
                Some(consumer_tag.clone()),
                2,
            );
            assert_eq!(channel_state, expected_state);
        }
        {
            let body_frame = AMQPFrame::Body(channel.id(), "{}".as_bytes().to_vec());
            conn.channels.handle_frame(body_frame).unwrap();
            let channel_state = channel.status().state();
            let expected_state = ChannelState::Connected;
            assert_eq!(channel_state, expected_state);
        }
    }

    #[test]
    fn basic_consume_empty_payload() {
        let _ = env_logger::try_init();

        use crate::consumer::Consumer;
        use crate::queue::{Queue, QueueState};

        // Bootstrap connection state to a consuming state
        let socket_state = SocketState::default();
        let waker = socket_state.handle();
        let executor = DefaultExecutor::default().unwrap();
        let internal_rpc = InternalRPC::new(executor.clone(), waker.clone());
        let conn = Connection::new(
            waker,
            internal_rpc.handle(),
            Frames::default(),
            executor.clone(),
        );
        conn.status.set_state(ConnectionState::Connected);
        conn.configuration.set_channel_max(2047);
        let channel = conn.channels.create(conn.closer.clone()).unwrap();
        channel.set_state(ChannelState::Connected);
        let queue_name = ShortString::from("consumed");
        let mut queue: QueueState = Queue::new(queue_name.clone(), 0, 0).into();
        let consumer_tag = ShortString::from("consumer-tag");
        let consumer = Consumer::new(consumer_tag.clone(), executor);
        queue.register_consumer(consumer_tag.clone(), consumer);
        conn.channels
            .get(channel.id())
            .map(|c| c.register_queue(queue));
        // Now test the state machine behaviour
        {
            let method = AMQPClass::Basic(basic::AMQPMethod::Deliver(basic::Deliver {
                consumer_tag: consumer_tag.clone(),
                delivery_tag: 1,
                redelivered: false,
                exchange: "".into(),
                routing_key: queue_name.clone(),
            }));
            let class_id = method.get_amqp_class_id();
            let deliver_frame = AMQPFrame::Method(channel.id(), method);
            conn.channels.handle_frame(deliver_frame).unwrap();
            let channel_state = channel.status().receiver_state();
            let expected_state = ChannelReceiverState::WillReceiveContent(
                class_id,
                Some(queue_name.clone()),
                Some(consumer_tag.clone()),
            );
            assert_eq!(channel_state, expected_state);
        }
        {
            let header_frame = AMQPFrame::Header(
                channel.id(),
                60,
                Box::new(AMQPContentHeader {
                    class_id: 60,
                    weight: 0,
                    body_size: 0,
                    properties: BasicProperties::default(),
                }),
            );
            conn.channels.handle_frame(header_frame).unwrap();
            let channel_state = channel.status().state();
            let expected_state = ChannelState::Connected;
            assert_eq!(channel_state, expected_state);
        }
    }
}
<commit_msg>rustfmt<commit_after>use crate::{
    channel::Channel,
    channels::Channels,
    configuration::Configuration,
    connection_closer::ConnectionCloser,
    connection_properties::ConnectionProperties,
    connection_status::{ConnectionState, ConnectionStatus, ConnectionStep},
    executor::{within_executor, DefaultExecutor, Executor},
    frames::Frames,
    internal_rpc::{InternalRPC, InternalRPCHandle},
    io_loop::IoLoop,
    reactor::DefaultReactorBuilder,
    socket_state::{SocketState, SocketStateHandle},
    tcp::{AMQPUriTcpExt, HandshakeResult, OwnedTLSConfig},
    thread::ThreadHandle,
    types::ShortUInt,
    uri::AMQPUri,
    Error, Promise, Result,
};
use amq_protocol::frame::{AMQPFrame, ProtocolVersion};
use async_trait::async_trait;
use log::{log_enabled, Level::Trace};
use std::{fmt, io, sync::Arc};

pub struct Connection {
    configuration: Configuration,
    status: ConnectionStatus,
    channels: Channels,
    io_loop: ThreadHandle,
    closer: Arc<ConnectionCloser>,
}

impl Connection {
    fn new(
        waker: SocketStateHandle,
        internal_rpc: InternalRPCHandle,
        frames: Frames,
        executor: Arc<dyn Executor>,
    ) -> Self {
        let configuration = Configuration::default();
        let status = ConnectionStatus::default();
        let channels = Channels::new(
            configuration.clone(),
            status.clone(),
            waker,
            internal_rpc.clone(),
            frames,
            executor,
        );
        let closer = Arc::new(ConnectionCloser::new(status.clone(), internal_rpc));
        let connection = Self {
            configuration,
            status,
            channels,
            io_loop: ThreadHandle::default(),
            closer,
        };

        connection.channels.create_zero();
        connection
    }

    /// Connect to an AMQP Server
    pub async fn connect(uri: &str, options: ConnectionProperties) -> Result<Connection> {
        Connect::connect(uri, options, OwnedTLSConfig::default()).await
    }

    /// Connect to an AMQP Server
    pub async fn connect_with_config(
        uri: &str,
        options: ConnectionProperties,
        config: OwnedTLSConfig,
    ) -> Result<Connection> {
        Connect::connect(uri, options, config).await
    }

    /// Connect to an AMQP Server
    pub async fn connect_uri(uri: AMQPUri, options: ConnectionProperties) -> Result<Connection> {
        Connect::connect(uri, options, OwnedTLSConfig::default()).await
    }

    /// Connect to an AMQP Server
    pub async fn connect_uri_with_identity(
        uri: AMQPUri,
        options: ConnectionProperties,
        config: OwnedTLSConfig,
    ) -> Result<Connection> {
        Connect::connect(uri, options, config).await
    }

    pub async fn create_channel(&self) -> Result<Channel> {
        if !self.status.connected() {
            return Err(Error::InvalidConnectionState(self.status.state()));
        }
        let channel = self.channels.create(self.closer.clone())?;
        channel.clone().channel_open(channel).await
    }

    /// Block current thread while the connection is still active.
    /// This is useful when you only have a consumer and nothing else keeping your application
    /// "alive".
    pub fn run(self) -> Result<()> {
        let io_loop = self.io_loop.clone();
        drop(self);
        if within_executor() {
            // io_loop waits for the executor to shutdown, do not deadlock in the
            // unlikely case where run is called from within it
            Ok(())
        } else {
            io_loop.wait("io loop")
        }
    }

    pub fn on_error<E: FnMut(Error) + Send + 'static>(&self, handler: E) {
        self.channels.set_error_handler(handler);
    }

    pub fn configuration(&self) -> &Configuration {
        &self.configuration
    }

    pub fn status(&self) -> &ConnectionStatus {
        &self.status
    }

    pub async fn close(&self, reply_code: ShortUInt, reply_text: &str) -> Result<()> {
        if let Some(channel0) = self.channels.get(0) {
            channel0
                .connection_close(reply_code, reply_text, 0, 0)
                .await
        } else {
            Ok(())
        }
    }

    /// Block all consumers and publishers on this connection
    pub async fn block(&self, reason: &str) -> Result<()> {
        if let Some(channel0) = self.channels.get(0) {
            channel0.connection_blocked(reason).await
        } else {
            Err(Error::InvalidConnectionState(self.status.state()))
        }
    }

    /// Unblock all consumers and publishers on this connection
    pub async fn unblock(&self) -> Result<()> {
        if let Some(channel0) = self.channels.get(0) {
            channel0.connection_unblocked().await
        } else {
            Err(Error::InvalidConnectionState(self.status.state()))
        }
    }

    /// Update the secret used by some authentication module such as OAuth2
    pub async fn update_secret(&self, new_secret: &str, reason: &str) -> Result<()> {
        if let Some(channel0) = self.channels.get(0) {
            channel0.connection_update_secret(new_secret, reason).await
        } else {
            Err(Error::InvalidConnectionState(self.status.state()))
        }
    }

    pub async fn connector(
        uri: AMQPUri,
        connect: Box<dyn FnOnce(&AMQPUri) -> HandshakeResult + Send + Sync>,
        mut options: ConnectionProperties,
    ) -> Result<Connection> {
        let executor = options
            .executor
            .take()
            .map(Ok)
            .unwrap_or_else(|| DefaultExecutor::default())?;

        let (connect_promise, resolver) = pinky_swear::PinkySwear::<HandshakeResult>::new();
        let connect_uri = uri.clone();
        executor.spawn_blocking(Box::new(move || {
            resolver.swear(connect(&connect_uri));
        }))?;

        let reactor_builder = options
            .reactor_builder
            .take()
            .unwrap_or_else(|| Arc::new(DefaultReactorBuilder));
        let socket_state = SocketState::default();
        let waker = socket_state.handle();
        let internal_rpc = InternalRPC::new(executor.clone(), waker.clone());
        let frames = Frames::default();
        let conn = Connection::new(waker, internal_rpc.handle(), frames.clone(), executor);
        let status = conn.status.clone();
        let configuration = conn.configuration.clone();
        status.set_vhost(&uri.vhost);
        status.set_username(&uri.authority.userinfo.username);
        if let Some(frame_max) = uri.query.frame_max {
            configuration.set_frame_max(frame_max);
        }
        if let Some(channel_max) = uri.query.channel_max {
            configuration.set_channel_max(channel_max);
        }
        if let Some(heartbeat) = uri.query.heartbeat {
            configuration.set_heartbeat(heartbeat);
        }
        let (promise_out, resolver) = Promise::new();
        if log_enabled!(Trace) {
            promise_out.set_marker("ProtocolHeader".into());
        }
        let channels = conn.channels.clone();
        if let Some(channel0) = channels.get(0) {
            channel0.send_frame(
                AMQPFrame::ProtocolHeader(ProtocolVersion::amqp_0_9_1()),
                resolver,
                None,
            )
        };
        let (promise_in, resolver) = Promise::new();
        if log_enabled!(Trace) {
            promise_in.set_marker("ProtocolHeader.Ok".into());
        }
        let io_loop_handle = conn.io_loop.clone();
        status.set_state(ConnectionState::Connecting);
        status.set_connection_step(ConnectionStep::ProtocolHeader(
            resolver,
            conn,
            uri.authority.userinfo.into(),
            uri.query.auth_mechanism.unwrap_or_default(),
            options,
        ));
        let handshake_result = connect_promise.await;
        IoLoop::new(
            status,
            configuration,
            channels,
            internal_rpc,
            frames,
            socket_state,
            io_loop_handle,
            handshake_result,
            &*reactor_builder,
        )
        .and_then(IoLoop::start)?;
        promise_out.await?;
        promise_in.await
    }
}

impl fmt::Debug for Connection {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Connection")
            .field("configuration", &self.configuration)
            .field("status", &self.status)
            .field("channels", &self.channels)
            .finish()
    }
}

/// Trait providing a method to connect to an AMQP server
#[async_trait]
pub trait Connect {
    /// connect to an AMQP server
    async fn connect(
        self,
        options: ConnectionProperties,
        config: OwnedTLSConfig,
    ) -> Result<Connection>;
}

#[async_trait]
impl Connect for AMQPUri {
    async fn connect(
        self,
        options: ConnectionProperties,
        config: OwnedTLSConfig,
    ) -> Result<Connection> {
        Connection::connector(
            self,
            Box::new(move |uri| AMQPUriTcpExt::connect_with_config(uri, config.as_ref())),
            options,
        )
        .await
    }
}

#[async_trait]
impl Connect for &str {
    async fn connect(
        self,
        options: ConnectionProperties,
        config: OwnedTLSConfig,
    ) -> Result<Connection> {
        match self.parse::<AMQPUri>() {
            Ok(uri) => Connect::connect(uri, options, config).await,
            Err(err) => Err(io::Error::new(io::ErrorKind::Other, err).into()),
        }
    }
}

#[cfg(test)]
mod tests {
    use env_logger;

    use super::*;
    use crate::channel_receiver_state::ChannelReceiverState;
    use crate::channel_status::ChannelState;
    use crate::types::ShortString;
    use crate::BasicProperties;
    use amq_protocol::frame::AMQPContentHeader;
    use amq_protocol::protocol::{basic, AMQPClass};

    #[test]
    fn basic_consume_small_payload() {
        let _ = env_logger::try_init();

        use crate::consumer::Consumer;
        use crate::queue::{Queue, QueueState};

        // Bootstrap connection state to a consuming state
        let executor = DefaultExecutor::default().unwrap();
        let socket_state = SocketState::default();
        let waker = socket_state.handle();
        let internal_rpc = InternalRPC::new(executor.clone(), waker.clone());
        let conn = Connection::new(
            waker,
            internal_rpc.handle(),
            Frames::default(),
            executor.clone(),
        );
        conn.status.set_state(ConnectionState::Connected);
        conn.configuration.set_channel_max(2047);
        let channel = conn.channels.create(conn.closer.clone()).unwrap();
        channel.set_state(ChannelState::Connected);
        let queue_name = ShortString::from("consumed");
        let mut queue: QueueState = Queue::new(queue_name.clone(), 0, 0).into();
        let consumer_tag = ShortString::from("consumer-tag");
        let consumer = Consumer::new(consumer_tag.clone(), executor);
        queue.register_consumer(consumer_tag.clone(), consumer);
        conn.channels
            .get(channel.id())
            .map(|c| c.register_queue(queue));
        // Now test the state machine behaviour
        {
            let method = AMQPClass::Basic(basic::AMQPMethod::Deliver(basic::Deliver {
                consumer_tag: consumer_tag.clone(),
                delivery_tag: 1,
                redelivered: false,
                exchange: "".into(),
                routing_key: queue_name.clone(),
            }));
            let class_id = method.get_amqp_class_id();
            let deliver_frame = AMQPFrame::Method(channel.id(), method);
            conn.channels.handle_frame(deliver_frame).unwrap();
            let channel_state = channel.status().receiver_state();
            let expected_state = ChannelReceiverState::WillReceiveContent(
                class_id,
                Some(queue_name.clone()),
                Some(consumer_tag.clone()),
            );
            assert_eq!(channel_state, expected_state);
        }
        {
            let header_frame = AMQPFrame::Header(
                channel.id(),
                60,
                Box::new(AMQPContentHeader {
                    class_id: 60,
                    weight: 0,
                    body_size: 2,
                    properties: BasicProperties::default(),
                }),
            );
            conn.channels.handle_frame(header_frame).unwrap();
            let channel_state = channel.status().receiver_state();
            let expected_state = ChannelReceiverState::ReceivingContent(
                Some(queue_name.clone()),
                Some(consumer_tag.clone()),
                2,
            );
            assert_eq!(channel_state, expected_state);
        }
        {
            let body_frame = AMQPFrame::Body(channel.id(), "{}".as_bytes().to_vec());
            conn.channels.handle_frame(body_frame).unwrap();
            let channel_state = channel.status().state();
            let expected_state = ChannelState::Connected;
            assert_eq!(channel_state, expected_state);
        }
    }

    #[test]
    fn basic_consume_empty_payload() {
        let _ = env_logger::try_init();

        use crate::consumer::Consumer;
        use crate::queue::{Queue, QueueState};

        // Bootstrap connection state to a consuming state
        let socket_state = SocketState::default();
        let waker = socket_state.handle();
        let executor = DefaultExecutor::default().unwrap();
        let internal_rpc = InternalRPC::new(executor.clone(), waker.clone());
        let conn = Connection::new(
            waker,
            internal_rpc.handle(),
            Frames::default(),
            executor.clone(),
        );
        conn.status.set_state(ConnectionState::Connected);
        conn.configuration.set_channel_max(2047);
        let channel = conn.channels.create(conn.closer.clone()).unwrap();
        channel.set_state(ChannelState::Connected);
        let queue_name = ShortString::from("consumed");
        let mut queue: QueueState = Queue::new(queue_name.clone(), 0, 0).into();
        let consumer_tag = ShortString::from("consumer-tag");
        let consumer = Consumer::new(consumer_tag.clone(), executor);
        queue.register_consumer(consumer_tag.clone(), consumer);
        conn.channels
            .get(channel.id())
            .map(|c| c.register_queue(queue));
        // Now test the state machine behaviour
        {
            let method = AMQPClass::Basic(basic::AMQPMethod::Deliver(basic::Deliver {
                consumer_tag: consumer_tag.clone(),
                delivery_tag: 1,
                redelivered: false,
                exchange: "".into(),
                routing_key: queue_name.clone(),
            }));
            let class_id = method.get_amqp_class_id();
            let deliver_frame = AMQPFrame::Method(channel.id(), method);
            conn.channels.handle_frame(deliver_frame).unwrap();
            let channel_state = channel.status().receiver_state();
            let expected_state = ChannelReceiverState::WillReceiveContent(
                class_id,
                Some(queue_name.clone()),
                Some(consumer_tag.clone()),
            );
            assert_eq!(channel_state, expected_state);
        }
        {
            let header_frame = AMQPFrame::Header(
                channel.id(),
                60,
                Box::new(AMQPContentHeader {
                    class_id: 60,
                    weight: 0,
                    body_size: 0,
                    properties: BasicProperties::default(),
                }),
            );
            conn.channels.handle_frame(header_frame).unwrap();
            let channel_state = channel.status().state();
            let expected_state = ChannelState::Connected;
            assert_eq!(channel_state, expected_state);
        }
    }
}
<|endoftext|>


<commit_before>use crate::syntax::Type;
use proc_macro2::Ident;

#[derive(Copy, Clone, PartialEq)]
pub enum Atom {
    Bool,
    U8,
    U16,
    U32,
    U64,
    Usize,
    I8,
    I16,
    I32,
    I64,
    Isize,
    F32,
    F64,
    CxxString,
    RustString,
}

impl Atom {
    pub fn from(ident: &Ident) -> Option<Self> {
        use self::Atom::*;
        match ident.to_string().as_str() {
            "bool" => Some(Bool),
            "u8" => Some(U8),
            "u16" => Some(U16),
            "u32" => Some(U32),
            "u64" => Some(U64),
            "usize" => Some(Usize),
            "i8" => Some(I8),
            "i16" => Some(I16),
            "i32" => Some(I32),
            "i64" => Some(I64),
            "isize" => Some(Isize),
            "f32" => Some(F32),
            "f64" => Some(F64),
            "CxxString" => Some(CxxString),
            "String" => Some(RustString),
            _ => None,
        }
    }
}

impl AsRef<str> for Atom {
    fn as_ref(&self) -> &str {
        use self::Atom::*;
        match self {
            Bool => "bool",
            U8 => "u8",
            U16 => "u16",
            U32 => "u32",
            U64 => "u64",
            Usize => "usize",
            I8 => "i8",
            I16 => "i16",
            I32 => "i32",
            I64 => "i64",
            Isize => "isize",
            F32 => "f32",
            F64 => "f64",
            CxxString => "CxxString",
            RustString => "String",
        }
    }
}

impl PartialEq<Atom> for Type {
    fn eq(&self, atom: &Atom) -> bool {
        match self {
            Type::Ident(ident) => ident == atom,
            _ => false,
        }
    }
}

impl PartialEq<Atom> for &Ident {
    fn eq(&self, atom: &Atom) -> bool {
        *self == atom
    }
}

impl PartialEq<Atom> for &Type {
    fn eq(&self, atom: &Atom) -> bool {
        *self == atom
    }
}
<commit_msg>Display for Atom<commit_after>use crate::syntax::Type;
use proc_macro2::Ident;
use std::fmt::{self, Display};

#[derive(Copy, Clone, PartialEq)]
pub enum Atom {
    Bool,
    U8,
    U16,
    U32,
    U64,
    Usize,
    I8,
    I16,
    I32,
    I64,
    Isize,
    F32,
    F64,
    CxxString,
    RustString,
}

impl Atom {
    pub fn from(ident: &Ident) -> Option<Self> {
        use self::Atom::*;
        match ident.to_string().as_str() {
            "bool" => Some(Bool),
            "u8" => Some(U8),
            "u16" => Some(U16),
            "u32" => Some(U32),
            "u64" => Some(U64),
            "usize" => Some(Usize),
            "i8" => Some(I8),
            "i16" => Some(I16),
            "i32" => Some(I32),
            "i64" => Some(I64),
            "isize" => Some(Isize),
            "f32" => Some(F32),
            "f64" => Some(F64),
            "CxxString" => Some(CxxString),
            "String" => Some(RustString),
            _ => None,
        }
    }
}

impl Display for Atom {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str(self.as_ref())
    }
}

impl AsRef<str> for Atom {
    fn as_ref(&self) -> &str {
        use self::Atom::*;
        match self {
            Bool => "bool",
            U8 => "u8",
            U16 => "u16",
            U32 => "u32",
            U64 => "u64",
            Usize => "usize",
            I8 => "i8",
            I16 => "i16",
            I32 => "i32",
            I64 => "i64",
            Isize => "isize",
            F32 => "f32",
            F64 => "f64",
            CxxString => "CxxString",
            RustString => "String",
        }
    }
}

impl PartialEq<Atom> for Type {
    fn eq(&self, atom: &Atom) -> bool {
        match self {
            Type::Ident(ident) => ident == atom,
            _ => false,
        }
    }
}

impl PartialEq<Atom> for &Ident {
    fn eq(&self, atom: &Atom) -> bool {
        *self == atom
    }
}

impl PartialEq<Atom> for &Type {
    fn eq(&self, atom: &Atom) -> bool {
        *self == atom
    }
}
<|endoftext|>


<commit_before>pub mod der;
pub mod rsa;

/// Helper function to encrypt mysql password using a public key loaded from a server.
pub fn encrypt(pass: &[u8], key: &[u8]) -> Vec<u8> {
    let pub_key = self::rsa::PublicKey::from_pem(key);
    let rng = ::rand::OsRng::new().unwrap();
    let pad = self::rsa::Pkcs1OaepPadding::new(rng);
    pub_key.encrypt_block(pass, pad)
}
<commit_msg>Clarify MySql versions, supported by crypto::encrypt function<commit_after>pub mod der;
pub mod rsa;

/// Helper function to encrypt mysql password using a public key loaded from a server.
///
/// It will use OAEP padding, so MySql versions prior to 8.0.5 are not supported.
pub fn encrypt(pass: &[u8], key: &[u8]) -> Vec<u8> {
    let pub_key = self::rsa::PublicKey::from_pem(key);
    let rng = ::rand::OsRng::new().unwrap();
    let pad = self::rsa::Pkcs1OaepPadding::new(rng);
    pub_key.encrypt_block(pass, pad)
}
<|endoftext|>


<commit_before>extern crate libc;

use defines::AfError;
use error::HANDLE_ERROR;
use self::libc::{c_int, size_t, c_char, c_void};
use std::ffi::CString;

extern {
    fn af_get_version(major: *mut c_int, minor: *mut c_int, patch: *mut c_int) -> c_int;
    fn af_info() -> c_int;
    fn af_info_string(str: *mut *mut c_char, verbose: bool) -> c_int;
    fn af_init() -> c_int;
    fn af_get_device_count(nDevices: *mut c_int) -> c_int;
    fn af_get_dbl_support(available: *mut c_int, device: c_int) -> c_int;
    fn af_set_device(device: c_int) -> c_int;
    fn af_get_device(device: *mut c_int) -> c_int;
    fn af_device_mem_info(alloc_bytes: *mut size_t, alloc_buffers: *mut size_t,
                          lock_bytes: *mut size_t, lock_buffers: *mut size_t) -> c_int;
    fn af_print_mem_info(msg: *const c_char, device_id: c_int) -> c_int;
    fn af_set_mem_step_size(step_bytes: size_t) -> c_int;
    fn af_get_mem_step_size(step_bytes: *mut size_t) -> c_int;
    fn af_device_gc() -> c_int;
    fn af_sync(device: c_int) -> c_int;

    fn af_free_host (ptr: *mut c_void) -> c_int;
}

/// Get ArrayFire Version Number
///
/// # Return Values
/// A triplet of integers indicating major, minor & fix release version numbers.
pub fn get_version() -> (i32, i32, i32) {
    unsafe {
        let mut maj: i32 = 0;
        let mut min: i32 = 0;
        let mut pat: i32 = 0;
        let err_val = af_get_version(&mut maj as *mut c_int,
                                     &mut min as *mut c_int, &mut pat as *mut c_int);
        HANDLE_ERROR(AfError::from(err_val));
        (maj, min, pat)
    }
}

/// Print library meta-info
///
/// # Examples
///
/// An example output of `af::info` call looks like below
///
/// ```text
/// ArrayFire v3.0.0 (CUDA, 64-bit Mac OSX, build d8d4b38)
/// Platform: CUDA Toolkit 7, Driver: CUDA Driver Version: 7000
/// [0] GeForce GT 750M, 2048 MB, CUDA Compute 3.0
/// ```
pub fn info() {
    unsafe {
        let err_val = af_info();
        HANDLE_ERROR(AfError::from(err_val));
    }
}

/// Return library meta-info as `String`
///
/// # Examples
///
/// An example output of `af::info_string` call looks like below
///
/// ```text
/// ArrayFire v3.0.0 (CUDA, 64-bit Mac OSX, build d8d4b38)
/// Platform: CUDA Toolkit 7, Driver: CUDA Driver Version: 7000
/// [0] GeForce GT 750M, 2048 MB, CUDA Compute 3.0
/// ```
pub fn info_string(verbose: bool) -> String {
    let mut tmp: *mut c_char = 0 as *mut c_char;
    unsafe {
        let err_val = af_info_string(&mut tmp, verbose);
        HANDLE_ERROR(AfError::from(err_val));

        let result = CString::from_raw(tmp).into_string().unwrap();

        let err_val = af_free_host(tmp as *mut c_void);
        HANDLE_ERROR(AfError::from(err_val));
        result
    }
}

/// Initialize ArrayFire library
///
/// 0th device will be the default device unless init call
/// is followed by set_device
pub fn init() {
    unsafe {
        let err_val = af_init();
        HANDLE_ERROR(AfError::from(err_val));
    }
}

/// Get total number of available devices
pub fn device_count() -> i32 {
    unsafe {
        let mut temp: i32 = 0;
        let err_val = af_get_device_count(&mut temp as *mut c_int);
        HANDLE_ERROR(AfError::from(err_val));
        temp
    }
}

/// Check if a device has double support
///
/// # Parameters
///
/// - `device` is the device for which double support is checked for
///
/// # Return Values
///
/// `True` if `device` device has double support, `False` otherwise.
pub fn is_double_available(device: i32) -> bool {
    unsafe {
        let mut temp: i32 = 0;
        let err_val = af_get_dbl_support(&mut temp as *mut c_int, device as c_int);
        HANDLE_ERROR(AfError::from(err_val));
        temp > 0
    }
}

/// Set active device
///
/// # Parameters
///
/// - `device` is the value of the device identifier which has to be set as active
pub fn set_device(device: i32) {
    unsafe {
        let err_val = af_set_device(device as c_int);
        HANDLE_ERROR(AfError::from(err_val));
    }
}

/// Get the current active device id
pub fn get_device() -> i32 {
    unsafe {
        let mut temp: i32 = 0;
        let err_val = af_get_device(&mut temp as *mut c_int);
        HANDLE_ERROR(AfError::from(err_val));
        temp
    }
}

/// Get memory information from the memory manager for the current active device
///
/// # Parameters
///
/// This function doesn't take any input parameters
///
/// # Return Values
///
/// A quadruple of values regarding the following information.
///
/// * Number of bytes allocated
/// * Number of buffers allocated
/// * Number of bytes locked
/// * Number of buffers locked
pub fn device_mem_info() -> (usize, usize, usize, usize) {
    unsafe {
        let mut o0: usize = 0;
        let mut o1: usize = 0;
        let mut o2: usize = 0;
        let mut o3: usize = 0;
        let err_val = af_device_mem_info(&mut o0 as *mut size_t,
                                         &mut o1 as *mut size_t,
                                         &mut o2 as *mut size_t,
                                         &mut o3 as *mut size_t);
        HANDLE_ERROR(AfError::from(err_val));
        (o0, o1, o2, o3)
    }
}

/// Print buffer details from the ArrayFire device manager
///
/// This information is printed in the form of a table.
///
/// # Parameters
///
/// - `msg` is a message to print before the table
/// - `device` is the id of the device for which buffer details are to be printed
///
/// # Return Values
///
/// None
pub fn print_mem_info(msg: String, device: i32) {
    unsafe {
        let cmsg = CString::new(msg.as_bytes());
        match cmsg {
            Ok(v) => {
                let err_val = af_print_mem_info(v.to_bytes_with_nul().as_ptr() as * const c_char,
                                                device as c_int);
                HANDLE_ERROR(AfError::from(err_val));
            },
            Err(_) => HANDLE_ERROR(AfError::ERR_INTERNAL),
        }
    }
}

/// Set the minimum memory chunk size
///
/// # Parameters
///
/// - `step_bytes` is the size of minimum memory chunk in bytes
///
/// # Return Values
///
/// None
pub fn set_mem_step_size(step_bytes: usize) {
    unsafe {
        let err_val = af_set_mem_step_size(step_bytes as size_t);
        HANDLE_ERROR(AfError::from(err_val));
    }
}

/// Get the minimum memory chunk size
///
/// # Parameters
///
/// None
///
/// # Return Values
///
/// Returns is the size of minimum memory chunk in bytes
pub fn get_mem_step_size() -> usize {
    unsafe {
        let mut temp: usize = 0;
        let err_val = af_get_mem_step_size(&mut temp as *mut size_t);
        HANDLE_ERROR(AfError::from(err_val));
        temp
    }
}

/// Call the garbage collection routine
pub fn device_gc() {
    unsafe {
        let err_val = af_device_gc();
        HANDLE_ERROR(AfError::from(err_val));
    }
}

/// Sync all operations on given device
///
/// # Parameters
///
/// - `device` on which the operations are to be synced
///
/// # Return Values
///
/// None
pub fn sync(device: i32) {
    unsafe {
        let err_val = af_sync(device as c_int);
        HANDLE_ERROR(AfError::from(err_val));
    }
}
<commit_msg>Added alloc_host and changed to use `to_str` -> `to_owned`<commit_after>extern crate libc;

use defines::{AfError, DType};
use error::HANDLE_ERROR;
use self::libc::{c_int, size_t, c_char, c_void};
use std::ffi::CString;

extern {
    fn af_get_version(major: *mut c_int, minor: *mut c_int, patch: *mut c_int) -> c_int;
    fn af_info() -> c_int;
    fn af_info_string(str: *mut *mut c_char, verbose: bool) -> c_int;
    fn af_init() -> c_int;
    fn af_get_device_count(nDevices: *mut c_int) -> c_int;
    fn af_get_dbl_support(available: *mut c_int, device: c_int) -> c_int;
    fn af_set_device(device: c_int) -> c_int;
    fn af_get_device(device: *mut c_int) -> c_int;
    fn af_device_mem_info(alloc_bytes: *mut size_t, alloc_buffers: *mut size_t,
                          lock_bytes: *mut size_t, lock_buffers: *mut size_t) -> c_int;
    fn af_print_mem_info(msg: *const c_char, device_id: c_int) -> c_int;
    fn af_set_mem_step_size(step_bytes: size_t) -> c_int;
    fn af_get_mem_step_size(step_bytes: *mut size_t) -> c_int;
    fn af_device_gc() -> c_int;
    fn af_sync(device: c_int) -> c_int;

    fn af_alloc_host(elements: size_t, _type: DType) -> *mut c_void;
    fn af_free_host(ptr: *mut c_void) -> c_int;
}

/// Get ArrayFire Version Number
///
/// # Return Values
/// A triplet of integers indicating major, minor & fix release version numbers.
pub fn get_version() -> (i32, i32, i32) {
    unsafe {
        let mut maj: i32 = 0;
        let mut min: i32 = 0;
        let mut pat: i32 = 0;
        let err_val = af_get_version(&mut maj as *mut c_int,
                                     &mut min as *mut c_int, &mut pat as *mut c_int);
        HANDLE_ERROR(AfError::from(err_val));
        (maj, min, pat)
    }
}

/// Print library meta-info
///
/// # Examples
///
/// An example output of `af::info` call looks like below
///
/// ```text
/// ArrayFire v3.0.0 (CUDA, 64-bit Mac OSX, build d8d4b38)
/// Platform: CUDA Toolkit 7, Driver: CUDA Driver Version: 7000
/// [0] GeForce GT 750M, 2048 MB, CUDA Compute 3.0
/// ```
pub fn info() {
    unsafe {
        let err_val = af_info();
        HANDLE_ERROR(AfError::from(err_val));
    }
}

/// Return library meta-info as `String`
///
/// # Examples
///
/// An example output of `af::info_string` call looks like below
///
/// ```text
/// ArrayFire v3.0.0 (CUDA, 64-bit Mac OSX, build d8d4b38)
/// Platform: CUDA Toolkit 7, Driver: CUDA Driver Version: 7000
/// [0] GeForce GT 750M, 2048 MB, CUDA Compute 3.0
/// ```
pub fn info_string(verbose: bool) -> String {
    let mut tmp: *mut c_char = 0 as *mut c_char;
    unsafe {
        let err_val = af_info_string(&mut tmp, verbose);
        HANDLE_ERROR(AfError::from(err_val));

        let result = (*CString::from_raw(tmp)).to_str().unwrap().to_owned();

        let err_val = af_free_host(tmp as *mut c_void);
        HANDLE_ERROR(AfError::from(err_val));
        result
    }
}

/// Initialize ArrayFire library
///
/// 0th device will be the default device unless init call
/// is followed by set_device
pub fn init() {
    unsafe {
        let err_val = af_init();
        HANDLE_ERROR(AfError::from(err_val));
    }
}

/// Get total number of available devices
pub fn device_count() -> i32 {
    unsafe {
        let mut temp: i32 = 0;
        let err_val = af_get_device_count(&mut temp as *mut c_int);
        HANDLE_ERROR(AfError::from(err_val));
        temp
    }
}

/// Check if a device has double support
///
/// # Parameters
///
/// - `device` is the device for which double support is checked for
///
/// # Return Values
///
/// `True` if `device` device has double support, `False` otherwise.
pub fn is_double_available(device: i32) -> bool {
    unsafe {
        let mut temp: i32 = 0;
        let err_val = af_get_dbl_support(&mut temp as *mut c_int, device as c_int);
        HANDLE_ERROR(AfError::from(err_val));
        temp > 0
    }
}

/// Set active device
///
/// # Parameters
///
/// - `device` is the value of the device identifier which has to be set as active
pub fn set_device(device: i32) {
    unsafe {
        let err_val = af_set_device(device as c_int);
        HANDLE_ERROR(AfError::from(err_val));
    }
}

/// Get the current active device id
pub fn get_device() -> i32 {
    unsafe {
        let mut temp: i32 = 0;
        let err_val = af_get_device(&mut temp as *mut c_int);
        HANDLE_ERROR(AfError::from(err_val));
        temp
    }
}

/// Get memory information from the memory manager for the current active device
///
/// # Parameters
///
/// This function doesn't take any input parameters
///
/// # Return Values
///
/// A quadruple of values regarding the following information.
///
/// * Number of bytes allocated
/// * Number of buffers allocated
/// * Number of bytes locked
/// * Number of buffers locked
pub fn device_mem_info() -> (usize, usize, usize, usize) {
    unsafe {
        let mut o0: usize = 0;
        let mut o1: usize = 0;
        let mut o2: usize = 0;
        let mut o3: usize = 0;
        let err_val = af_device_mem_info(&mut o0 as *mut size_t,
                                         &mut o1 as *mut size_t,
                                         &mut o2 as *mut size_t,
                                         &mut o3 as *mut size_t);
        HANDLE_ERROR(AfError::from(err_val));
        (o0, o1, o2, o3)
    }
}

/// Print buffer details from the ArrayFire device manager
///
/// This information is printed in the form of a table.
///
/// # Parameters
///
/// - `msg` is a message to print before the table
/// - `device` is the id of the device for which buffer details are to be printed
///
/// # Return Values
///
/// None
pub fn print_mem_info(msg: String, device: i32) {
    unsafe {
        let cmsg = CString::new(msg.as_bytes());
        match cmsg {
            Ok(v) => {
                let err_val = af_print_mem_info(v.to_bytes_with_nul().as_ptr() as * const c_char,
                                                device as c_int);
                HANDLE_ERROR(AfError::from(err_val));
            },
            Err(_) => HANDLE_ERROR(AfError::ERR_INTERNAL),
        }
    }
}

/// Set the minimum memory chunk size
///
/// # Parameters
///
/// - `step_bytes` is the size of minimum memory chunk in bytes
///
/// # Return Values
///
/// None
pub fn set_mem_step_size(step_bytes: usize) {
    unsafe {
        let err_val = af_set_mem_step_size(step_bytes as size_t);
        HANDLE_ERROR(AfError::from(err_val));
    }
}

/// Get the minimum memory chunk size
///
/// # Parameters
///
/// None
///
/// # Return Values
///
/// Returns is the size of minimum memory chunk in bytes
pub fn get_mem_step_size() -> usize {
    unsafe {
        let mut temp: usize = 0;
        let err_val = af_get_mem_step_size(&mut temp as *mut size_t);
        HANDLE_ERROR(AfError::from(err_val));
        temp
    }
}

/// Call the garbage collection routine
pub fn device_gc() {
    unsafe {
        let err_val = af_device_gc();
        HANDLE_ERROR(AfError::from(err_val));
    }
}

/// Sync all operations on given device
///
/// # Parameters
///
/// - `device` on which the operations are to be synced
///
/// # Return Values
///
/// None
pub fn sync(device: i32) {
    unsafe {
        let err_val = af_sync(device as c_int);
        HANDLE_ERROR(AfError::from(err_val));
    }
}
<|endoftext|>