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

// Manages the SQL database that holds the list of registered domain names.
// Each records is made of the name, the private token and the Let's Encrypt
// challenge value.

use std::sync::mpsc::{channel, Receiver};
use std::thread;

use r2d2_sqlite::SqliteConnectionManager;
use r2d2;
use rusqlite::Row;
use rusqlite::Result as SqlResult;
use rusqlite::types::{ToSql, ToSqlOutput};
use std::time::{SystemTime, UNIX_EPOCH};

#[derive(Clone, Debug, PartialEq, Serialize)]
pub struct DomainRecord {
    pub token: String,
    pub local_name: String,
    pub remote_name: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub dns_challenge: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub local_ip: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub public_ip: Option<String>,
    pub description: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub email: Option<String>,
    pub timestamp: i64,
}

macro_rules! sqlstr {
    ($row:ident, $index:expr) => (
        {
            let raw = $row.get::<i32, String>($index);
            if raw.is_empty() {
                None
            } else {
                Some(raw)
            }
        }
    )
}

impl DomainRecord {
    fn from_sql(row: Row) -> Self {
        DomainRecord {
            token: row.get(0),
            local_name: row.get(1),
            remote_name: row.get(2),
            dns_challenge: sqlstr!(row, 3),
            local_ip: sqlstr!(row, 4),
            public_ip: sqlstr!(row, 5),
            description: row.get(6),
            email: sqlstr!(row, 7),
            timestamp: row.get(8),
        }
    }

    pub fn new(token: &str,
               local_name: &str,
               remote_name: &str,
               dns_challenge: Option<&str>,
               local_ip: Option<&str>,
               public_ip: Option<&str>,
               description: &str,
               email: Option<&str>,
               timestamp: i64)
               -> Self {
        macro_rules! str2sql {
            ($val:expr) => (
                if $val.is_some() {
                    Some($val.unwrap().to_owned())
                } else {
                    None
                }
            )
        }

        DomainRecord {
            local_name: local_name.to_owned(),
            remote_name: remote_name.to_owned(),
            token: token.to_owned(),
            dns_challenge: str2sql!(dns_challenge),
            local_ip: str2sql!(local_ip),
            public_ip: str2sql!(public_ip),
            description: description.to_owned(),
            email: str2sql!(email),
            timestamp: timestamp,
        }
    }
}

unsafe impl Send for DomainRecord {}
unsafe impl Sync for DomainRecord {}

#[derive(Debug)]
pub enum DomainError {
    DbUnavailable,
    SQLError(String),
    NoRecord,
}

#[derive(Clone)]
pub struct DomainDb {
    pool: r2d2::Pool<SqliteConnectionManager>,
}

// try! like macro that sends back the error of a SQLite operation
// over the mpsc channel if needed.
macro_rules! sqltry {
    ($sql:expr, $tx:ident, $err:expr) => (
        match $sql {
            Err(_) => {
                $tx.send(Err($err)).unwrap();
                return;
            }
            Ok(value) => value,
        }
    );

    ($sql:expr, $tx:ident) => (
        match $sql {
            Err(err) => {
                $tx.send(Err(DomainError::SQLError(format!("{}", err)))).unwrap();
                return;
            }
            Ok(value) => value,
        }
    );
}

#[derive(Clone)]
pub enum SqlParam {
    Text(String),
    Integer(i64),
}

impl ToSql for SqlParam {
    fn to_sql(&self) -> SqlResult<ToSqlOutput> {
        match *self {
            SqlParam::Text(ref text) => text.to_sql(),
            SqlParam::Integer(ref number) => number.to_sql(),
        }
    }
}

impl DomainDb {
    pub fn new(path: &str) -> Self {
        debug!("Opening database at {}", path);
        let config = r2d2::Config::default();
        let manager = SqliteConnectionManager::new(path);
        let pool = r2d2::Pool::new(config, manager).expect(&format!("Unable to open database at {}",
                                                                    path));

        // Create the database table if needed.
        let conn = pool.get().unwrap();
        conn.execute("CREATE TABLE IF NOT EXISTS domains (
                      token         TEXT NOT NULL PRIMARY KEY,
                      local_name    TEXT NOT NULL,
                      remote_name   TEXT NOT NULL,
                      dns_challenge TEXT NOT NULL,
                      local_ip      TEXT NOT NULL,
                      public_ip     TEXT NOT NULL,
                      description   TEXT NOT NULL,
                      email         TEXT NOT NULL,
                      timestamp     INTEGER)",
                     &[])
            .unwrap_or_else(|err| {
                                panic!("Unable to create the domains database: {}", err);
                            });

        conn.execute("CREATE UNIQUE INDEX IF NOT EXISTS domains_local ON domains(local_name)",
                     &[])
            .unwrap_or_else(|err| {
                                panic!("Unable to create the domains_local index: {}", err);
                            });

        conn.execute("CREATE UNIQUE INDEX IF NOT EXISTS domains_remote ON domains(remote_name)",
                     &[])
            .unwrap_or_else(|err| {
                                panic!("Unable to create the domains_remote index: {}", err);
                            });

        conn.execute("CREATE INDEX IF NOT EXISTS domains_timestamp ON domains(timestamp)",
                     &[])
            .unwrap_or_else(|err| {
                                panic!("Unable to create the domains_timestamp index: {}", err);
                            });

        conn.execute("CREATE INDEX IF NOT EXISTS domains_public_ip ON domains(public_ip)",
                     &[])
            .unwrap_or_else(|err| {
                                panic!("Unable to create the domains_public_ip index: {}", err);
                            });

        conn.execute("CREATE INDEX IF NOT EXISTS domains_email ON domains(email)",
                     &[])
            .unwrap_or_else(|err| {
                                panic!("Unable to create the domains_email index: {}", err);
                            });

        DomainDb { pool: pool }
    }

    fn select_record(&self,
                     request: &str,
                     value: &str)
                     -> Receiver<Result<DomainRecord, DomainError>> {
        let (tx, rx) = channel();

        // Run the sql command on a pooled thread.
        let pool = self.pool.clone();
        let value = value.to_owned();
        let request = request.to_owned();
        thread::spawn(move || {
            let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
            let mut stmt = sqltry!(conn.prepare(&request), tx);
            let mut rows = sqltry!(stmt.query(&[&value]), tx);
            if let Some(result_row) = rows.next() {
                let row = sqltry!(result_row, tx);
                tx.send(Ok(DomainRecord::from_sql(row))).unwrap();
            } else {
                tx.send(Err(DomainError::NoRecord)).unwrap();
            }
        });

        rx
    }

    fn select_records(&self,
                      request: &str,
                      value: &str)
                      -> Receiver<Result<Vec<DomainRecord>, DomainError>> {
        let (tx, rx) = channel();

        // Run the sql command on a pooled thread.
        let pool = self.pool.clone();
        let value = value.to_owned();
        let request = request.to_owned();
        thread::spawn(move || {
            let mut result = Vec::new();
            let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
            let mut stmt = sqltry!(conn.prepare(&request), tx);
            let mut rows = sqltry!(stmt.query(&[&value]), tx);
            while let Some(result_row) = rows.next() {
                let row = sqltry!(result_row, tx);
                result.push(DomainRecord::from_sql(row));
            }
            tx.send(Ok(result)).unwrap();
        });

        rx
    }

    pub fn get_records_by_public_ip(&self,
                                    public_ip: &str)
                                    -> Receiver<Result<Vec<DomainRecord>, DomainError>> {
        self.select_records("SELECT token, local_name, remote_name, dns_challenge, \
                            local_ip, public_ip, description, email, timestamp \
                            FROM domains WHERE public_ip=$1",
                            public_ip)
    }

    pub fn get_record_by_name(&self, name: &str) -> Receiver<Result<DomainRecord, DomainError>> {
        self.select_record("SELECT token, local_name, remote_name, dns_challenge, \
                            local_ip, public_ip, description, email, timestamp \
                            FROM domains WHERE local_name=$1 or remote_name=$1",
                           name)
    }

    pub fn get_record_by_token(&self, token: &str) -> Receiver<Result<DomainRecord, DomainError>> {
        self.select_record("SELECT token, local_name, remote_name, dns_challenge, \
                           local_ip, public_ip, description, email, timestamp \
                            FROM domains WHERE token=$1",
                           token)
    }

    pub fn add_record(&self, record: DomainRecord) -> Receiver<Result<(), DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        let record = record.clone();
        thread::spawn(move || {
            let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
            sqltry!(conn.execute("INSERT INTO domains VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)",
                                 &[&record.token,
                                   &record.local_name,
                                   &record.remote_name,
                                   &record.dns_challenge.unwrap_or("".to_owned()),
                                   &record.local_ip.unwrap_or("".to_owned()),
                                   &record.public_ip.unwrap_or("".to_owned()),
                                   &record.description,
                                   &record.email.unwrap_or("".to_owned()),
                                   &record.timestamp]),
                    tx);
            tx.send(Ok(())).unwrap();
        });

        rx
    }

    pub fn update_record(&self, record: DomainRecord) -> Receiver<Result<(), DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        let record = record.clone();
        thread::spawn(move || {
            let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);

            sqltry!(conn.execute("UPDATE domains SET dns_challenge=$1, local_ip=$2, \
                                  public_ip=$3, timestamp=$4 \
                                  WHERE (local_name=$5 OR remote_name=$6) AND token=$7",
                                 &[&record.dns_challenge.unwrap_or("".to_owned()),
                                   &record.local_ip.unwrap_or("".to_owned()),
                                   &record.public_ip.unwrap_or("".to_owned()),
                                   &record.timestamp,
                                   &record.local_name,
                                   &record.remote_name,
                                   &record.token]),
                    tx);
            tx.send(Ok(())).unwrap();
        });

        rx
    }

    // Evict records older than a given timestamp.
    // Returns the number of evicted records.
    pub fn evict_records(&self, timestamp: SqlParam) -> Receiver<Result<i32, DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        thread::spawn(move || {
            let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
            let now = SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_secs() as i64;
            let res = sqltry!(conn.execute("UPDATE domains SET local_ip=$1, \
                    public_ip=$2, timestamp=$3 where timestamp<$4",
                                           &[&"", &"", &now, &timestamp]),
                              tx);
            tx.send(Ok(res)).unwrap();
        });

        rx
    }

    // Returns the number of rows affected.
    pub fn execute_1param_sql(&self,
                              request: &str,
                              value: SqlParam)
                              -> Receiver<Result<i32, DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        let value = value.to_owned();
        let request = request.to_owned();
        thread::spawn(move || {
                          let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
                          let res = sqltry!(conn.execute(&request, &[&value]), tx);
                          tx.send(Ok(res)).unwrap();
                      });

        rx
    }

    pub fn delete_record_by_token(&self, token: &str) -> Receiver<Result<i32, DomainError>> {
        self.execute_1param_sql("DELETE FROM domains WHERE token=$1",
                                SqlParam::Text(token.to_owned()))
    }

    pub fn flush(&self) -> Receiver<Result<(), DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        thread::spawn(move || {
                          let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
                          sqltry!(conn.execute("DELETE FROM domains", &[]), tx);
                          tx.send(Ok(())).unwrap();
                      });
        rx
    }
}

#[test]
fn test_domain_store() {
    let db = DomainDb::new("domain_db_test.sqlite");

    // Start with an empty db.
    db.flush().recv().unwrap().expect("Flushing the db");

    // Check that we don't find any record.
    match db.get_record_by_name("test.example.org")
              .recv()
              .unwrap() {
        Err(DomainError::NoRecord) => {}
        Err(err) => panic!("Should not find a record by name in an empty db. {:?}", err),
        _ => panic!("Should not find a record by name in an empty db."),
    }

    match db.get_record_by_token("test-token").recv().unwrap() {
        Err(DomainError::NoRecord) => {}
        _ => panic!("Should not find a record by token in an empty db."),
    }

    // Add a record without a dns challenge.
    let no_challenge_record = DomainRecord::new("test-token",
                                                "local.test.example.org",
                                                "test.example.org",
                                                None,
                                                None,
                                                None,
                                                "Test Server",
                                                None,
                                                0);
    db.add_record(no_challenge_record.clone())
        .recv()
        .unwrap()
        .expect("Adding the no_challenge record");

    // Check that we can find it and that it matches our record.
    match db.get_record_by_name("test.example.org")
              .recv()
              .unwrap() {
        Ok(record) => assert_eq!(record, no_challenge_record),
        Err(err) => panic!("Failed to find record by name: {:?}", err),
    }

    match db.get_record_by_token("test-token").recv().unwrap() {
        Ok(record) => assert_eq!(record, no_challenge_record),
        Err(err) => panic!("Failed to find record by token: {:?}", err),
    }

    // Update the record to have challenge.
    let challenge_record = DomainRecord::new("test-token",
                                             "local.test.example.org",
                                             "test.example.org",
                                             Some("dns-challenge"),
                                             None,
                                             None,
                                             "Test Server",
                                             None,
                                             0);
    db.update_record(challenge_record.clone())
        .recv()
        .unwrap()
        .expect("Updating the challenge record");

    // Check that we can find it and that it matches our record.
    match db.get_record_by_name("test.example.org")
              .recv()
              .unwrap() {
        Ok(record) => assert_eq!(record, challenge_record),
        Err(err) => panic!("Failed to find record by name: {:?}", err),
    }

    match db.get_record_by_token("test-token").recv().unwrap() {
        Ok(record) => assert_eq!(record, challenge_record),
        Err(err) => panic!("Failed to find record by token: {:?}", err),
    }

    // Remove by token.
    db.delete_record_by_token(&challenge_record.token)
        .recv()
        .unwrap()
        .expect("Should delete record by token");
    match db.get_record_by_name(&challenge_record.local_name)
              .recv()
              .unwrap() {
        Err(DomainError::NoRecord) => {}
        _ => panic!("Should not find this record anymore."),
    }
}
<commit_msg>macro-ify sql index declarations<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

// Manages the SQL database that holds the list of registered domain names.
// Each records is made of the name, the private token and the Let's Encrypt
// challenge value.

use std::sync::mpsc::{channel, Receiver};
use std::thread;

use r2d2_sqlite::SqliteConnectionManager;
use r2d2;
use rusqlite::Row;
use rusqlite::Result as SqlResult;
use rusqlite::types::{ToSql, ToSqlOutput};
use std::time::{SystemTime, UNIX_EPOCH};

#[derive(Clone, Debug, PartialEq, Serialize)]
pub struct DomainRecord {
    pub token: String,
    pub local_name: String,
    pub remote_name: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub dns_challenge: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub local_ip: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub public_ip: Option<String>,
    pub description: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub email: Option<String>,
    pub timestamp: i64,
}

macro_rules! sqlstr {
    ($row:ident, $index:expr) => (
        {
            let raw = $row.get::<i32, String>($index);
            if raw.is_empty() {
                None
            } else {
                Some(raw)
            }
        }
    )
}

impl DomainRecord {
    fn from_sql(row: Row) -> Self {
        DomainRecord {
            token: row.get(0),
            local_name: row.get(1),
            remote_name: row.get(2),
            dns_challenge: sqlstr!(row, 3),
            local_ip: sqlstr!(row, 4),
            public_ip: sqlstr!(row, 5),
            description: row.get(6),
            email: sqlstr!(row, 7),
            timestamp: row.get(8),
        }
    }

    pub fn new(token: &str,
               local_name: &str,
               remote_name: &str,
               dns_challenge: Option<&str>,
               local_ip: Option<&str>,
               public_ip: Option<&str>,
               description: &str,
               email: Option<&str>,
               timestamp: i64)
               -> Self {
        macro_rules! str2sql {
            ($val:expr) => (
                if $val.is_some() {
                    Some($val.unwrap().to_owned())
                } else {
                    None
                }
            )
        }

        DomainRecord {
            local_name: local_name.to_owned(),
            remote_name: remote_name.to_owned(),
            token: token.to_owned(),
            dns_challenge: str2sql!(dns_challenge),
            local_ip: str2sql!(local_ip),
            public_ip: str2sql!(public_ip),
            description: description.to_owned(),
            email: str2sql!(email),
            timestamp: timestamp,
        }
    }
}

unsafe impl Send for DomainRecord {}
unsafe impl Sync for DomainRecord {}

#[derive(Debug)]
pub enum DomainError {
    DbUnavailable,
    SQLError(String),
    NoRecord,
}

#[derive(Clone)]
pub struct DomainDb {
    pool: r2d2::Pool<SqliteConnectionManager>,
}

// try! like macro that sends back the error of a SQLite operation
// over the mpsc channel if needed.
macro_rules! sqltry {
    ($sql:expr, $tx:ident, $err:expr) => (
        match $sql {
            Err(_) => {
                $tx.send(Err($err)).unwrap();
                return;
            }
            Ok(value) => value,
        }
    );

    ($sql:expr, $tx:ident) => (
        match $sql {
            Err(err) => {
                $tx.send(Err(DomainError::SQLError(format!("{}", err)))).unwrap();
                return;
            }
            Ok(value) => value,
        }
    );
}

#[derive(Clone)]
pub enum SqlParam {
    Text(String),
    Integer(i64),
}

impl ToSql for SqlParam {
    fn to_sql(&self) -> SqlResult<ToSqlOutput> {
        match *self {
            SqlParam::Text(ref text) => text.to_sql(),
            SqlParam::Integer(ref number) => number.to_sql(),
        }
    }
}

impl DomainDb {
    pub fn new(path: &str) -> Self {
        debug!("Opening database at {}", path);
        let config = r2d2::Config::default();
        let manager = SqliteConnectionManager::new(path);
        let pool = r2d2::Pool::new(config, manager).expect(&format!("Unable to open database at {}",
                                                                    path));

        // Create the database table if needed.
        let conn = pool.get().unwrap();
        conn.execute("CREATE TABLE IF NOT EXISTS domains (
                      token         TEXT NOT NULL PRIMARY KEY,
                      local_name    TEXT NOT NULL,
                      remote_name   TEXT NOT NULL,
                      dns_challenge TEXT NOT NULL,
                      local_ip      TEXT NOT NULL,
                      public_ip     TEXT NOT NULL,
                      description   TEXT NOT NULL,
                      email         TEXT NOT NULL,
                      timestamp     INTEGER)",
                     &[])
            .unwrap_or_else(|err| {
                                panic!("Unable to create the domains database: {}", err);
                            });

        macro_rules! index {
            ($name:expr) => (
                conn.execute(&format!("CREATE UNIQUE INDEX IF NOT EXISTS domains_{} ON domains({})",
                                      $name, $name), &[]).unwrap_or_else(|err| {
                                panic!("Unable to create the domains_{} index: {}", $name, err);
                            });
            )
        }

        index!("local_name");
        index!("remote_name");
        index!("timestamp");
        index!("public_ip");
        index!("email");

        DomainDb { pool: pool }
    }

    fn select_record(&self,
                     request: &str,
                     value: &str)
                     -> Receiver<Result<DomainRecord, DomainError>> {
        let (tx, rx) = channel();

        // Run the sql command on a pooled thread.
        let pool = self.pool.clone();
        let value = value.to_owned();
        let request = request.to_owned();
        thread::spawn(move || {
            let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
            let mut stmt = sqltry!(conn.prepare(&request), tx);
            let mut rows = sqltry!(stmt.query(&[&value]), tx);
            if let Some(result_row) = rows.next() {
                let row = sqltry!(result_row, tx);
                tx.send(Ok(DomainRecord::from_sql(row))).unwrap();
            } else {
                tx.send(Err(DomainError::NoRecord)).unwrap();
            }
        });

        rx
    }

    fn select_records(&self,
                      request: &str,
                      value: &str)
                      -> Receiver<Result<Vec<DomainRecord>, DomainError>> {
        let (tx, rx) = channel();

        // Run the sql command on a pooled thread.
        let pool = self.pool.clone();
        let value = value.to_owned();
        let request = request.to_owned();
        thread::spawn(move || {
            let mut result = Vec::new();
            let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
            let mut stmt = sqltry!(conn.prepare(&request), tx);
            let mut rows = sqltry!(stmt.query(&[&value]), tx);
            while let Some(result_row) = rows.next() {
                let row = sqltry!(result_row, tx);
                result.push(DomainRecord::from_sql(row));
            }
            tx.send(Ok(result)).unwrap();
        });

        rx
    }

    pub fn get_records_by_public_ip(&self,
                                    public_ip: &str)
                                    -> Receiver<Result<Vec<DomainRecord>, DomainError>> {
        self.select_records("SELECT token, local_name, remote_name, dns_challenge, \
                            local_ip, public_ip, description, email, timestamp \
                            FROM domains WHERE public_ip=$1",
                            public_ip)
    }

    pub fn get_record_by_name(&self, name: &str) -> Receiver<Result<DomainRecord, DomainError>> {
        self.select_record("SELECT token, local_name, remote_name, dns_challenge, \
                            local_ip, public_ip, description, email, timestamp \
                            FROM domains WHERE local_name=$1 or remote_name=$1",
                           name)
    }

    pub fn get_record_by_token(&self, token: &str) -> Receiver<Result<DomainRecord, DomainError>> {
        self.select_record("SELECT token, local_name, remote_name, dns_challenge, \
                           local_ip, public_ip, description, email, timestamp \
                            FROM domains WHERE token=$1",
                           token)
    }

    pub fn add_record(&self, record: DomainRecord) -> Receiver<Result<(), DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        let record = record.clone();
        thread::spawn(move || {
            let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
            sqltry!(conn.execute("INSERT INTO domains VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)",
                                 &[&record.token,
                                   &record.local_name,
                                   &record.remote_name,
                                   &record.dns_challenge.unwrap_or("".to_owned()),
                                   &record.local_ip.unwrap_or("".to_owned()),
                                   &record.public_ip.unwrap_or("".to_owned()),
                                   &record.description,
                                   &record.email.unwrap_or("".to_owned()),
                                   &record.timestamp]),
                    tx);
            tx.send(Ok(())).unwrap();
        });

        rx
    }

    pub fn update_record(&self, record: DomainRecord) -> Receiver<Result<(), DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        let record = record.clone();
        thread::spawn(move || {
            let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);

            sqltry!(conn.execute("UPDATE domains SET dns_challenge=$1, local_ip=$2, \
                                  public_ip=$3, timestamp=$4 \
                                  WHERE (local_name=$5 OR remote_name=$6) AND token=$7",
                                 &[&record.dns_challenge.unwrap_or("".to_owned()),
                                   &record.local_ip.unwrap_or("".to_owned()),
                                   &record.public_ip.unwrap_or("".to_owned()),
                                   &record.timestamp,
                                   &record.local_name,
                                   &record.remote_name,
                                   &record.token]),
                    tx);
            tx.send(Ok(())).unwrap();
        });

        rx
    }

    // Evict records older than a given timestamp.
    // Returns the number of evicted records.
    pub fn evict_records(&self, timestamp: SqlParam) -> Receiver<Result<i32, DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        thread::spawn(move || {
            let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
            let now = SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_secs() as i64;
            let res = sqltry!(conn.execute("UPDATE domains SET local_ip=$1, \
                    public_ip=$2, timestamp=$3 where timestamp<$4",
                                           &[&"", &"", &now, &timestamp]),
                              tx);
            tx.send(Ok(res)).unwrap();
        });

        rx
    }

    // Returns the number of rows affected.
    pub fn execute_1param_sql(&self,
                              request: &str,
                              value: SqlParam)
                              -> Receiver<Result<i32, DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        let value = value.to_owned();
        let request = request.to_owned();
        thread::spawn(move || {
                          let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
                          let res = sqltry!(conn.execute(&request, &[&value]), tx);
                          tx.send(Ok(res)).unwrap();
                      });

        rx
    }

    pub fn delete_record_by_token(&self, token: &str) -> Receiver<Result<i32, DomainError>> {
        self.execute_1param_sql("DELETE FROM domains WHERE token=$1",
                                SqlParam::Text(token.to_owned()))
    }

    pub fn flush(&self) -> Receiver<Result<(), DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        thread::spawn(move || {
                          let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
                          sqltry!(conn.execute("DELETE FROM domains", &[]), tx);
                          tx.send(Ok(())).unwrap();
                      });
        rx
    }
}

#[test]
fn test_domain_store() {
    let db = DomainDb::new("domain_db_test.sqlite");

    // Start with an empty db.
    db.flush().recv().unwrap().expect("Flushing the db");

    // Check that we don't find any record.
    match db.get_record_by_name("test.example.org")
              .recv()
              .unwrap() {
        Err(DomainError::NoRecord) => {}
        Err(err) => panic!("Should not find a record by name in an empty db. {:?}", err),
        _ => panic!("Should not find a record by name in an empty db."),
    }

    match db.get_record_by_token("test-token").recv().unwrap() {
        Err(DomainError::NoRecord) => {}
        _ => panic!("Should not find a record by token in an empty db."),
    }

    // Add a record without a dns challenge.
    let no_challenge_record = DomainRecord::new("test-token",
                                                "local.test.example.org",
                                                "test.example.org",
                                                None,
                                                None,
                                                None,
                                                "Test Server",
                                                None,
                                                0);
    db.add_record(no_challenge_record.clone())
        .recv()
        .unwrap()
        .expect("Adding the no_challenge record");

    // Check that we can find it and that it matches our record.
    match db.get_record_by_name("test.example.org")
              .recv()
              .unwrap() {
        Ok(record) => assert_eq!(record, no_challenge_record),
        Err(err) => panic!("Failed to find record by name: {:?}", err),
    }

    match db.get_record_by_token("test-token").recv().unwrap() {
        Ok(record) => assert_eq!(record, no_challenge_record),
        Err(err) => panic!("Failed to find record by token: {:?}", err),
    }

    // Update the record to have challenge.
    let challenge_record = DomainRecord::new("test-token",
                                             "local.test.example.org",
                                             "test.example.org",
                                             Some("dns-challenge"),
                                             None,
                                             None,
                                             "Test Server",
                                             None,
                                             0);
    db.update_record(challenge_record.clone())
        .recv()
        .unwrap()
        .expect("Updating the challenge record");

    // Check that we can find it and that it matches our record.
    match db.get_record_by_name("test.example.org")
              .recv()
              .unwrap() {
        Ok(record) => assert_eq!(record, challenge_record),
        Err(err) => panic!("Failed to find record by name: {:?}", err),
    }

    match db.get_record_by_token("test-token").recv().unwrap() {
        Ok(record) => assert_eq!(record, challenge_record),
        Err(err) => panic!("Failed to find record by token: {:?}", err),
    }

    // Remove by token.
    db.delete_record_by_token(&challenge_record.token)
        .recv()
        .unwrap()
        .expect("Should delete record by token");
    match db.get_record_by_name(&challenge_record.local_name)
              .recv()
              .unwrap() {
        Err(DomainError::NoRecord) => {}
        _ => panic!("Should not find this record anymore."),
    }
}
<|endoftext|>


<commit_before>// format with label break value.
fn main() {
    {}

    {
        do_thing();
        if condition_not_met() {
            break 'block;
        }
        do_next_thing();
        if condition_not_met() {
            break 'block;
        }
        do_last_thing();
    }

    let result = 'block: {
        if foo() {
            // comment
            break 'block 1;
        }
        if bar() {
            /* comment */
            break 'block 2;
        }
        3
    };
}
<commit_msg>update test<commit_after>// format with label break value.
fn main() {
    'emty_block: {}

    'block: {
        do_thing();
        if condition_not_met() {
            break 'block;
        }
        do_next_thing();
        if condition_not_met() {
            break 'block;
        }
        do_last_thing();
    }

    let result = 'block: {
        if foo() {
            // comment
            break 'block 1;
        }
        if bar() {
            /* comment */
            break 'block 2;
        }
        3
    };
}
<|endoftext|>


<commit_before>// This returns a miri pointer at type usize, if the argument is a proper pointer
fn transmute_ptr_to_int<T>(x: *const T) -> usize {
    unsafe { std::mem::transmute(x) }
}

fn cast() {
    // Some casting-to-int with arithmetic.
    let x = &42 as *const i32 as usize;
    let y = x * 2;
    assert_eq!(y, x + x);
    let z = y as u8 as usize;
    assert_eq!(z, y % 256);
}

fn format() {
    // Pointer string formatting! We can't check the output as it changes when libstd changes,
    // but we can make sure Miri does not error.
    format!("{:?}", &mut 13 as *mut _);
}

fn transmute() {
    // Check that intptrcast is triggered for explicit casts and that it is consistent with
    // transmuting.
    let a: *const i32 = &42;
    let b = transmute_ptr_to_int(a) as u8;
    let c = a as usize as u8;
    assert_eq!(b, c);
}

fn ptr_bitops1() {
    let bytes = [0i8, 1, 2, 3, 4, 5, 6, 7, 8, 9];
    let one = bytes.as_ptr().wrapping_offset(1);
    let three = bytes.as_ptr().wrapping_offset(3);
    let res = (one as usize) | (three as usize);
    format!("{}", res);
}

fn ptr_bitops2() {
    let val = 13usize;
    let addr = &val as *const _ as usize;
    let _val = addr & 13;
}

fn ptr_eq_dangling() {
    let b = Box::new(0);
    let x = &*b as *const i32; // soon-to-be dangling
    drop(b);
    let b = Box::new(0);
    let y = &*b as *const i32; // different allocation
    // We cannot compare these even though both are inbounds -- they *could* be
    // equal if memory was reused.
    assert!(x != y);
}

fn ptr_eq_out_of_bounds() {
    let b = Box::new(0);
    let x = (&*b as *const i32).wrapping_sub(0x800); // out-of-bounds
    let b = Box::new(0);
    let y = &*b as *const i32; // different allocation
    // We cannot compare these even though both allocations are live -- they *could* be
    // equal (with the right base addresses).
    assert!(x != y);
}

fn ptr_eq_out_of_bounds_null() {
    let b = Box::new(0);
    let x = (&*b as *const i32).wrapping_sub(0x800); // out-of-bounds
    // We cannot compare this with NULL. After all, this *could* be NULL (with the right base address).
    assert!(x != std::ptr::null());
}

fn ptr_eq_integer() {
    let b = Box::new(0);
    let x = &*b as *const i32;
    // We cannot compare this with a non-NULL integer. After all, these *could* be equal (with the right base address).
    assert!(x != 64 as *const i32);
}

fn main() {
    cast();
    format();
    transmute();
    ptr_bitops1();
    ptr_bitops2();
    ptr_eq_dangling();
    ptr_eq_out_of_bounds();
    ptr_eq_out_of_bounds_null();
    ptr_eq_integer();
}
<commit_msg>Auto merge of #1475 - RalfJung:dangcast, r=RalfJung<commit_after>// This returns a miri pointer at type usize, if the argument is a proper pointer
fn transmute_ptr_to_int<T>(x: *const T) -> usize {
    unsafe { std::mem::transmute(x) }
}

fn cast() {
    // Some casting-to-int with arithmetic.
    let x = &42 as *const i32 as usize;
    let y = x * 2;
    assert_eq!(y, x + x);
    let z = y as u8 as usize;
    assert_eq!(z, y % 256);
}

/// Test usize->ptr cast for dangling and OOB address.
/// That is safe, and thus has to work.
fn cast_dangling() {
    let b = Box::new(0);
    let x = &*b as *const i32 as usize;
    drop(b);
    let _val = x as *const i32;

    let b = Box::new(0);
    let mut x = &*b as *const i32 as usize;
    x += 0x100;
    let _val = x as *const i32;
}

fn format() {
    // Pointer string formatting! We can't check the output as it changes when libstd changes,
    // but we can make sure Miri does not error.
    format!("{:?}", &mut 13 as *mut _);
}

fn transmute() {
    // Check that intptrcast is triggered for explicit casts and that it is consistent with
    // transmuting.
    let a: *const i32 = &42;
    let b = transmute_ptr_to_int(a) as u8;
    let c = a as usize as u8;
    assert_eq!(b, c);
}

fn ptr_bitops1() {
    let bytes = [0i8, 1, 2, 3, 4, 5, 6, 7, 8, 9];
    let one = bytes.as_ptr().wrapping_offset(1);
    let three = bytes.as_ptr().wrapping_offset(3);
    let res = (one as usize) | (three as usize);
    format!("{}", res);
}

fn ptr_bitops2() {
    let val = 13usize;
    let addr = &val as *const _ as usize;
    let _val = addr & 13;
}

fn ptr_eq_dangling() {
    let b = Box::new(0);
    let x = &*b as *const i32; // soon-to-be dangling
    drop(b);
    let b = Box::new(0);
    let y = &*b as *const i32; // different allocation
    // They *could* be equal if memory was reused, but probably are not.
    assert!(x != y);
}

fn ptr_eq_out_of_bounds() {
    let b = Box::new(0);
    let x = (&*b as *const i32).wrapping_sub(0x800); // out-of-bounds
    let b = Box::new(0);
    let y = &*b as *const i32; // different allocation
    // They *could* be equal (with the right base addresses), but probably are not.
    assert!(x != y);
}

fn ptr_eq_out_of_bounds_null() {
    let b = Box::new(0);
    let x = (&*b as *const i32).wrapping_sub(0x800); // out-of-bounds
    // This *could* be NULL (with the right base address), but probably is not.
    assert!(x != std::ptr::null());
}

fn ptr_eq_integer() {
    let b = Box::new(0);
    let x = &*b as *const i32;
    // These *could* be equal (with the right base address), but probably are not.
    assert!(x != 64 as *const i32);
}

fn main() {
    cast();
    cast_dangling();
    format();
    transmute();
    ptr_bitops1();
    ptr_bitops2();
    ptr_eq_dangling();
    ptr_eq_out_of_bounds();
    ptr_eq_out_of_bounds_null();
    ptr_eq_integer();
}
<|endoftext|>


<commit_before>use nom::{
  be_u8, be_u16,
  IResult,
  ErrorCode, Err,
};

use frame::{
  ChannelAssignment, NumberType,
  Frame,
  Header, Footer,
};

use utility::to_u32;

pub fn frame_parser(input: &[u8], channels: u8) -> IResult<&[u8], Frame> {
  chain!(input,
    frame_header: header ~
    frame_footer: footer,
    || {
      Frame {
        header: frame_header,
        footer: frame_footer,
      }
    }
  )
}

fn blocking_strategy(input: &[u8]) -> IResult<&[u8], bool> {
  match take!(input, 2) {
    IResult::Done(i, bytes)   => {
      let sync_code = ((bytes[0] as u16) << 6) +
                      (bytes[1] as u16) >> 2;
      let is_valid  = ((bytes[1] >> 1) & 0b01) == 0;

      if sync_code == 0b11111111111110 && is_valid {
        let is_variable_block_size = (bytes[1] & 0b01) == 1;

        IResult::Done(i, is_variable_block_size)
      } else {
        IResult::Error(Err::Position(ErrorCode::Digit as u32, input))
      }
    }
    IResult::Error(error)     => IResult::Error(error),
    IResult::Incomplete(need) => IResult::Incomplete(need),
  }
}

fn block_sample(input: &[u8]) -> IResult<&[u8], (u32, u32)> {
  match take!(input, 1) {
    IResult::Done(i, bytes)   => {
      let sample_rate = bytes[0] & 0x0f;

      if sample_rate != 0x0f {
        let block_size = bytes[0] >> 4;

        IResult::Done(i, (block_size as u32, sample_rate as u32))
      } else {
        IResult::Error(Err::Position(ErrorCode::Digit as u32, input))
      }
    }
    IResult::Error(error)     => IResult::Error(error),
    IResult::Incomplete(need) => IResult::Incomplete(need),
  }
}

fn channel_bits(input: &[u8]) -> IResult<&[u8], (ChannelAssignment, u8)> {
  match take!(input, 1) {
    IResult::Done(i, bytes)   => {
      let channel_assignment = match bytes[0] >> 4 {
        0b0000...0b0111 => ChannelAssignment::Independent,
        0b1000          => ChannelAssignment::LeftSide,
        0b1001          => ChannelAssignment::RightSide,
        0b1010          => ChannelAssignment::MiddleSide,
        _               => ChannelAssignment::Independent,
      };
      let size_byte = (bytes[0] >> 1) & 0b0111;
      let is_valid  = (bytes[0] & 0b01) == 0;

      if is_valid {
        IResult::Done(i, (channel_assignment, size_byte))
      } else {
        IResult::Error(Err::Position(ErrorCode::Digit as u32, input))
      }
    }
    IResult::Error(error)     => IResult::Error(error),
    IResult::Incomplete(need) => IResult::Incomplete(need),
  }
}

fn utf8_size(input: &[u8], is_u64: bool)
             -> IResult<&[u8], Option<(usize, u8)>> {
  map!(input, be_u8, |utf8_header| {
    match utf8_header {
      0b00000000...0b01111111 => Some((0, utf8_header)),
      0b11000000...0b11011111 => Some((1, utf8_header & 0b00011111)),
      0b11100000...0b11101111 => Some((2, utf8_header & 0b00001111)),
      0b11110000...0b11110111 => Some((3, utf8_header & 0b00000111)),
      0b11111000...0b11111011 => Some((4, utf8_header & 0b00000011)),
      0b11111100...0b11111101 => Some((5, utf8_header & 0b00000001)),
      0b11111110              => if is_u64 { Some((6, 0)) } else { None },
      _                       => None,
    }
  })
}

fn sample_or_frame_number(input: &[u8], is_sample: bool,
                          (size, value): (usize, u8))
                          -> IResult<&[u8], NumberType> {
  let mut result   = value as u64;
  let mut is_error = false;

  match take!(input, size) {
    IResult::Done(i, bytes)   => {
      for i in (0..size) {
        let byte = bytes[i] as u64;

        match byte {
          0b10000000...10111111 => {
            result = (result << 6) + (byte & 0b00111111);
          }
          _                     => {
            is_error = true;
            break;
          }
        }
      }

      if is_error {
        IResult::Error(Err::Position(ErrorCode::Digit as u32, input))
      } else if is_sample {
        IResult::Done(i, NumberType::Sample(result))
      } else {
        IResult::Done(i, NumberType::Frame(result as u32))
      }
    }
    IResult::Error(error)     => IResult::Error(error),
    IResult::Incomplete(need) => IResult::Incomplete(need),
  }
}

fn secondary_block_size(input: &[u8], block_byte: u32)
                        -> IResult<&[u8], Option<u32>> {
  match block_byte {
    0b0110 => opt!(input, map!(take!(1), to_u32)),
    0b0111 => opt!(input, map!(take!(2), to_u32)),
    _      => IResult::Done(input, None)
  }
}

fn secondary_sample_rate(input: &[u8], sample_byte: u32)
                        -> IResult<&[u8], Option<u32>> {
  match sample_byte {
    0b1100 => opt!(input, map!(take!(1), to_u32)),
    0b1101 => opt!(input, map!(take!(2), to_u32)),
    0b1110 => opt!(input, map!(take!(2), to_u32)),
    _      => IResult::Done(input, None)
  }
}

named!(header <&[u8], Header>,
  chain!(
    is_variable_block_size: blocking_strategy ~
    tuple0: block_sample ~
    tuple1: channel_bits ~
    number_opt: apply!(utf8_size, is_variable_block_size) ~
    number_length: expr_opt!(number_opt) ~
    number: apply!(sample_or_frame_number, is_variable_block_size,
                   number_length) ~
    alt_block_size: apply!(secondary_block_size, tuple0.0) ~
    alt_sample_rate: apply!(secondary_sample_rate, tuple0.1) ~
    crc: be_u8,
    || {
      let (block_byte, sample_byte)       = tuple0;
      let (channel_assignment, size_byte) = tuple1;

      let block_size = match block_byte {
        0b0000          => 0,
        0b0001          => 192,
        0b0010...0b0101 => 576 * 2_u32.pow(block_byte - 2),
        0b0110 | 0b0111 => alt_block_size.unwrap() + 1,
        0b1000...0b1111 => 256 * 2_u32.pow(block_byte - 8),
        _               => unreachable!(),
      };

      let sample_rate = match sample_byte {
        0b0000 => 0,
        0b0001 => 88200,
        0b0010 => 176400,
        0b0011 => 192000,
        0b0100 => 8000,
        0b0101 => 16000,
        0b0110 => 22050,
        0b0111 => 24000,
        0b1000 => 32000,
        0b1001 => 44100,
        0b1010 => 48000,
        0b1011 => 96000,
        0b1100 => alt_sample_rate.unwrap() * 1000,
        0b1101 => alt_sample_rate.unwrap(),
        0b1110 => alt_sample_rate.unwrap() * 10,
        0b1111 => 0,
        _      => unreachable!(),
      };

      let bits_per_sample = match size_byte {
        0b0000 => 0,
        0b0001 => 8,
        0b0010 => 12,
        0b0011 => 0,
        0b0100 => 16,
        0b0101 => 20,
        0b0110 => 24,
        0b0111 => 0,
        _      => unreachable!(),
      };

      Header {
        block_size: block_size,
        sample_rate: sample_rate,
        channel_assignment: channel_assignment,
        bits_per_sample: bits_per_sample,
        number: number,
        crc: crc,
      }
    }
  )
);

named!(footer <&[u8], Footer>, map!(be_u16, Footer));
<commit_msg>Change frame_parser and header to take StreamInfo<commit_after>use nom::{
  be_u8, be_u16,
  IResult,
  ErrorCode, Err,
};

use frame::{
  ChannelAssignment, NumberType,
  Frame,
  Header, Footer,
};

use metadata::StreamInfo;
use utility::to_u32;

pub fn frame_parser<'a>(input: &'a [u8], stream_info: &StreamInfo)
                        -> IResult<'a, &'a [u8], Frame> {
  chain!(input,
    frame_header: apply!(header, stream_info) ~
    frame_footer: footer,
    || {
      Frame {
        header: frame_header,
        footer: frame_footer,
      }
    }
  )
}

fn blocking_strategy(input: &[u8]) -> IResult<&[u8], bool> {
  match take!(input, 2) {
    IResult::Done(i, bytes)   => {
      let sync_code = ((bytes[0] as u16) << 6) +
                      (bytes[1] as u16) >> 2;
      let is_valid  = ((bytes[1] >> 1) & 0b01) == 0;

      if sync_code == 0b11111111111110 && is_valid {
        let is_variable_block_size = (bytes[1] & 0b01) == 1;

        IResult::Done(i, is_variable_block_size)
      } else {
        IResult::Error(Err::Position(ErrorCode::Digit as u32, input))
      }
    }
    IResult::Error(error)     => IResult::Error(error),
    IResult::Incomplete(need) => IResult::Incomplete(need),
  }
}

fn block_sample(input: &[u8]) -> IResult<&[u8], (u32, u32)> {
  match take!(input, 1) {
    IResult::Done(i, bytes)   => {
      let sample_rate = bytes[0] & 0x0f;

      if sample_rate != 0x0f {
        let block_size = bytes[0] >> 4;

        IResult::Done(i, (block_size as u32, sample_rate as u32))
      } else {
        IResult::Error(Err::Position(ErrorCode::Digit as u32, input))
      }
    }
    IResult::Error(error)     => IResult::Error(error),
    IResult::Incomplete(need) => IResult::Incomplete(need),
  }
}

fn channel_bits(input: &[u8]) -> IResult<&[u8], (ChannelAssignment, u8)> {
  match take!(input, 1) {
    IResult::Done(i, bytes)   => {
      let channel_assignment = match bytes[0] >> 4 {
        0b0000...0b0111 => ChannelAssignment::Independent,
        0b1000          => ChannelAssignment::LeftSide,
        0b1001          => ChannelAssignment::RightSide,
        0b1010          => ChannelAssignment::MiddleSide,
        _               => ChannelAssignment::Independent,
      };
      let size_byte = (bytes[0] >> 1) & 0b0111;
      let is_valid  = (bytes[0] & 0b01) == 0;

      if is_valid {
        IResult::Done(i, (channel_assignment, size_byte))
      } else {
        IResult::Error(Err::Position(ErrorCode::Digit as u32, input))
      }
    }
    IResult::Error(error)     => IResult::Error(error),
    IResult::Incomplete(need) => IResult::Incomplete(need),
  }
}

fn utf8_size(input: &[u8], is_u64: bool)
             -> IResult<&[u8], Option<(usize, u8)>> {
  map!(input, be_u8, |utf8_header| {
    match utf8_header {
      0b00000000...0b01111111 => Some((0, utf8_header)),
      0b11000000...0b11011111 => Some((1, utf8_header & 0b00011111)),
      0b11100000...0b11101111 => Some((2, utf8_header & 0b00001111)),
      0b11110000...0b11110111 => Some((3, utf8_header & 0b00000111)),
      0b11111000...0b11111011 => Some((4, utf8_header & 0b00000011)),
      0b11111100...0b11111101 => Some((5, utf8_header & 0b00000001)),
      0b11111110              => if is_u64 { Some((6, 0)) } else { None },
      _                       => None,
    }
  })
}

fn sample_or_frame_number(input: &[u8], is_sample: bool,
                          (size, value): (usize, u8))
                          -> IResult<&[u8], NumberType> {
  let mut result   = value as u64;
  let mut is_error = false;

  match take!(input, size) {
    IResult::Done(i, bytes)   => {
      for i in (0..size) {
        let byte = bytes[i] as u64;

        match byte {
          0b10000000...10111111 => {
            result = (result << 6) + (byte & 0b00111111);
          }
          _                     => {
            is_error = true;
            break;
          }
        }
      }

      if is_error {
        IResult::Error(Err::Position(ErrorCode::Digit as u32, input))
      } else if is_sample {
        IResult::Done(i, NumberType::Sample(result))
      } else {
        IResult::Done(i, NumberType::Frame(result as u32))
      }
    }
    IResult::Error(error)     => IResult::Error(error),
    IResult::Incomplete(need) => IResult::Incomplete(need),
  }
}

fn secondary_block_size(input: &[u8], block_byte: u32)
                        -> IResult<&[u8], Option<u32>> {
  match block_byte {
    0b0110 => opt!(input, map!(take!(1), to_u32)),
    0b0111 => opt!(input, map!(take!(2), to_u32)),
    _      => IResult::Done(input, None)
  }
}

fn secondary_sample_rate(input: &[u8], sample_byte: u32)
                        -> IResult<&[u8], Option<u32>> {
  match sample_byte {
    0b1100 => opt!(input, map!(take!(1), to_u32)),
    0b1101 => opt!(input, map!(take!(2), to_u32)),
    0b1110 => opt!(input, map!(take!(2), to_u32)),
    _      => IResult::Done(input, None)
  }
}

fn header<'a>(input: &'a [u8], stream_info: &StreamInfo)
              -> IResult<'a, &'a [u8], Header> {
  chain!(input,
    is_variable_block_size: blocking_strategy ~
    tuple0: block_sample ~
    tuple1: channel_bits ~
    number_opt: apply!(utf8_size, is_variable_block_size) ~
    number_length: expr_opt!(number_opt) ~
    number: apply!(sample_or_frame_number, is_variable_block_size,
                   number_length) ~
    alt_block_size: apply!(secondary_block_size, tuple0.0) ~
    alt_sample_rate: apply!(secondary_sample_rate, tuple0.1) ~
    crc: be_u8,
    || {
      let (block_byte, sample_byte)       = tuple0;
      let (channel_assignment, size_byte) = tuple1;

      let block_size = match block_byte {
        0b0000          => 0,
        0b0001          => 192,
        0b0010...0b0101 => 576 * 2_u32.pow(block_byte - 2),
        0b0110 | 0b0111 => alt_block_size.unwrap() + 1,
        0b1000...0b1111 => 256 * 2_u32.pow(block_byte - 8),
        _               => unreachable!(),
      };

      let sample_rate = match sample_byte {
        0b0000 => 0,
        0b0001 => 88200,
        0b0010 => 176400,
        0b0011 => 192000,
        0b0100 => 8000,
        0b0101 => 16000,
        0b0110 => 22050,
        0b0111 => 24000,
        0b1000 => 32000,
        0b1001 => 44100,
        0b1010 => 48000,
        0b1011 => 96000,
        0b1100 => alt_sample_rate.unwrap() * 1000,
        0b1101 => alt_sample_rate.unwrap(),
        0b1110 => alt_sample_rate.unwrap() * 10,
        0b1111 => 0,
        _      => unreachable!(),
      };

      let bits_per_sample = match size_byte {
        0b0000 => 0,
        0b0001 => 8,
        0b0010 => 12,
        0b0011 => 0,
        0b0100 => 16,
        0b0101 => 20,
        0b0110 => 24,
        0b0111 => 0,
        _      => unreachable!(),
      };

      Header {
        block_size: block_size,
        sample_rate: sample_rate,
        channel_assignment: channel_assignment,
        bits_per_sample: bits_per_sample,
        number: number,
        crc: crc,
      }
    }
  )
}

named!(footer <&[u8], Footer>, map!(be_u16, Footer));
<|endoftext|>


<commit_before>use maplit::hashmap;

use crate::GpuState;
use std::{collections::HashMap, mem};

#[repr(C)]
#[derive(Copy, Clone)]
pub(crate) struct GenHeightmapsUniforms {
    pub position: [i32; 2],
    pub origin: [i32; 2],
    pub spacing: f32,
    pub in_slot: i32,
    pub out_slot: i32,
    pub level_resolution: i32,
    pub face: u32,
    pub _padding: [u32; 3],
}
unsafe impl bytemuck::Zeroable for GenHeightmapsUniforms {}
unsafe impl bytemuck::Pod for GenHeightmapsUniforms {}

#[repr(C)]
#[derive(Copy, Clone)]
pub(crate) struct GenDisplacementsUniforms {
    pub node_center: [f64; 3],
    pub padding0: f64,
    pub origin: [i32; 2],
    pub position: [i32; 2],
    pub stride: i32,
    pub heightmaps_slot: i32,
    pub displacements_slot: i32,
    pub face: i32,
    pub level_resolution: u32,
}
unsafe impl bytemuck::Zeroable for GenDisplacementsUniforms {}
unsafe impl bytemuck::Pod for GenDisplacementsUniforms {}

#[repr(C)]
#[derive(Copy, Clone)]
pub(crate) struct GenNormalsUniforms {
    pub heightmaps_origin: [i32; 2],
    pub heightmaps_slot: i32,
    pub normals_slot: i32,
    pub spacing: f32,
    pub padding: [f32; 3],
}
unsafe impl bytemuck::Zeroable for GenNormalsUniforms {}
unsafe impl bytemuck::Pod for GenNormalsUniforms {}

#[repr(C)]
#[derive(Copy, Clone)]
pub(crate) struct GenMaterialsUniforms {
    pub heightmaps_origin: [i32; 2],
    pub parent_origin: [u32; 2],
    pub heightmaps_slot: i32,
    pub normals_slot: i32,
    pub albedo_slot: i32,
    pub parent_slot: i32,
    pub spacing: f32,
    pub level: u32,
    pub position: [i32; 2],
    pub level_resolution: u32,
    pub _padding: [u32; 3],
}
unsafe impl bytemuck::Zeroable for GenMaterialsUniforms {}
unsafe impl bytemuck::Pod for GenMaterialsUniforms {}

pub(crate) struct ComputeShader<U> {
    shader: rshader::ShaderSet,
    bindgroup_pipeline: Option<(wgpu::BindGroup, wgpu::ComputePipeline)>,
    uniforms: Option<wgpu::Buffer>,
    name: String,
    _phantom: std::marker::PhantomData<U>,
}
#[allow(unused)]
impl<U: bytemuck::Pod> ComputeShader<U> {
    pub fn new(shader: rshader::ShaderSource, name: String) -> Self {
        Self {
            shader: rshader::ShaderSet::compute_only(shader).unwrap(),
            bindgroup_pipeline: None,
            uniforms: None,
            name,
            _phantom: std::marker::PhantomData,
        }
    }

    pub fn refresh(&mut self) -> bool {
        if self.shader.refresh() {
            self.bindgroup_pipeline = None;
            true
        } else {
            false
        }
    }

    pub fn run(
        &mut self,
        device: &wgpu::Device,
        encoder: &mut wgpu::CommandEncoder,
        state: &GpuState,
        dimensions: (u32, u32, u32),
        uniforms: &U,
    ) {
        if self.uniforms.is_none() {
            self.uniforms = Some(device.create_buffer(&wgpu::BufferDescriptor {
                size: mem::size_of::<U>() as u64,
                usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::UNIFORM,
                mapped_at_creation: false,
                label: Some(&format!("buffer.{}.uniforms", self.name)),
            }));
        }
        if self.bindgroup_pipeline.is_none() {
            let (bind_group, bind_group_layout) = state.bind_group_for_shader(
                device,
                &self.shader,
                hashmap!["ubo".into() => (false, wgpu::BindingResource::Buffer(wgpu::BufferBinding {
                    buffer: self.uniforms.as_ref().unwrap(),
                    offset: 0,
                    size: None,
                }))],
                HashMap::new(),
                &format!("bindgroup.{}", self.name),
            );
            self.bindgroup_pipeline = Some((
                bind_group,
                device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                    layout: Some(&device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                        bind_group_layouts: [&bind_group_layout][..].into(),
                        push_constant_ranges: &[],
                        label: Some(&format!("pipeline.{}.layout", self.name)),
                    })),
                    module: &device.create_shader_module(&wgpu::ShaderModuleDescriptor {
                        label: Some(&format!("shader.{}", self.name)),
                        source: self.shader.compute(),
                    }),
                    entry_point: "main",
                    label: Some(&format!("pipeline.{}", self.name)),
                }),
            ));
        }

        let staging = device.create_buffer(&wgpu::BufferDescriptor {
            size: mem::size_of::<U>() as u64,
            usage: wgpu::BufferUsages::COPY_SRC,
            label: Some(&format!("buffer.temporary.{}.upload", self.name)),
            mapped_at_creation: true,
        });
        let mut buffer_view = staging.slice(..).get_mapped_range_mut();
        buffer_view[..mem::size_of::<U>()].copy_from_slice(bytemuck::bytes_of(uniforms));
        drop(buffer_view);
        staging.unmap();

        encoder.copy_buffer_to_buffer(
            &staging,
            0,
            self.uniforms.as_ref().unwrap(),
            0,
            mem::size_of::<U>() as u64,
        );

        let mut cpass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { label: None });
        cpass.set_pipeline(&self.bindgroup_pipeline.as_ref().unwrap().1);
        cpass.set_bind_group(0, &self.bindgroup_pipeline.as_ref().unwrap().0, &[]);
        cpass.dispatch(dimensions.0, dimensions.1, dimensions.2);
    }
}
<commit_msg>Better support for ComputeShader<()><commit_after>use maplit::hashmap;

use crate::GpuState;
use std::{collections::HashMap, mem};

#[repr(C)]
#[derive(Copy, Clone)]
pub(crate) struct GenHeightmapsUniforms {
    pub position: [i32; 2],
    pub origin: [i32; 2],
    pub spacing: f32,
    pub in_slot: i32,
    pub out_slot: i32,
    pub level_resolution: i32,
    pub face: u32,
    pub _padding: [u32; 3],
}
unsafe impl bytemuck::Zeroable for GenHeightmapsUniforms {}
unsafe impl bytemuck::Pod for GenHeightmapsUniforms {}

#[repr(C)]
#[derive(Copy, Clone)]
pub(crate) struct GenDisplacementsUniforms {
    pub node_center: [f64; 3],
    pub padding0: f64,
    pub origin: [i32; 2],
    pub position: [i32; 2],
    pub stride: i32,
    pub heightmaps_slot: i32,
    pub displacements_slot: i32,
    pub face: i32,
    pub level_resolution: u32,
}
unsafe impl bytemuck::Zeroable for GenDisplacementsUniforms {}
unsafe impl bytemuck::Pod for GenDisplacementsUniforms {}

#[repr(C)]
#[derive(Copy, Clone)]
pub(crate) struct GenNormalsUniforms {
    pub heightmaps_origin: [i32; 2],
    pub heightmaps_slot: i32,
    pub normals_slot: i32,
    pub spacing: f32,
    pub padding: [f32; 3],
}
unsafe impl bytemuck::Zeroable for GenNormalsUniforms {}
unsafe impl bytemuck::Pod for GenNormalsUniforms {}

#[repr(C)]
#[derive(Copy, Clone)]
pub(crate) struct GenMaterialsUniforms {
    pub heightmaps_origin: [i32; 2],
    pub parent_origin: [u32; 2],
    pub heightmaps_slot: i32,
    pub normals_slot: i32,
    pub albedo_slot: i32,
    pub parent_slot: i32,
    pub spacing: f32,
    pub level: u32,
    pub position: [i32; 2],
    pub level_resolution: u32,
    pub _padding: [u32; 3],
}
unsafe impl bytemuck::Zeroable for GenMaterialsUniforms {}
unsafe impl bytemuck::Pod for GenMaterialsUniforms {}

pub(crate) struct ComputeShader<U> {
    shader: rshader::ShaderSet,
    bindgroup_pipeline: Option<(wgpu::BindGroup, wgpu::ComputePipeline)>,
    uniforms: Option<wgpu::Buffer>,
    name: String,
    _phantom: std::marker::PhantomData<U>,
}
#[allow(unused)]
impl<U: bytemuck::Pod> ComputeShader<U> {
    pub fn new(shader: rshader::ShaderSource, name: String) -> Self {
        Self {
            shader: rshader::ShaderSet::compute_only(shader).unwrap(),
            bindgroup_pipeline: None,
            uniforms: None,
            name,
            _phantom: std::marker::PhantomData,
        }
    }

    pub fn refresh(&mut self) -> bool {
        if self.shader.refresh() {
            self.bindgroup_pipeline = None;
            true
        } else {
            false
        }
    }

    pub fn run(
        &mut self,
        device: &wgpu::Device,
        encoder: &mut wgpu::CommandEncoder,
        state: &GpuState,
        dimensions: (u32, u32, u32),
        uniforms: &U,
    ) {
        if mem::size_of::<U>() > 0 && self.uniforms.is_none() {
            self.uniforms = Some(device.create_buffer(&wgpu::BufferDescriptor {
                size: mem::size_of::<U>() as u64,
                usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::UNIFORM,
                mapped_at_creation: false,
                label: Some(&format!("buffer.{}.uniforms", self.name)),
            }));
        }
        if self.bindgroup_pipeline.is_none() {
            let (bind_group, bind_group_layout) = state.bind_group_for_shader(
                device,
                &self.shader,
                if self.uniforms.is_some() {
                    hashmap!["ubo".into() => (false, wgpu::BindingResource::Buffer(wgpu::BufferBinding {
                    buffer: self.uniforms.as_ref().unwrap(),
                    offset: 0,
                    size: None,
                }))] } else {
                    HashMap::new()
                },
                HashMap::new(),
                &format!("bindgroup.{}", self.name),
            );
            self.bindgroup_pipeline = Some((
                bind_group,
                device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                    layout: Some(&device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                        bind_group_layouts: [&bind_group_layout][..].into(),
                        push_constant_ranges: &[],
                        label: Some(&format!("pipeline.{}.layout", self.name)),
                    })),
                    module: &device.create_shader_module(&wgpu::ShaderModuleDescriptor {
                        label: Some(&format!("shader.{}", self.name)),
                        source: self.shader.compute(),
                    }),
                    entry_point: "main",
                    label: Some(&format!("pipeline.{}", self.name)),
                }),
            ));
        }

        if self.uniforms.is_some() {
            let staging = device.create_buffer(&wgpu::BufferDescriptor {
                size: mem::size_of::<U>() as u64,
                usage: wgpu::BufferUsages::COPY_SRC,
                label: Some(&format!("buffer.temporary.{}.upload", self.name)),
                mapped_at_creation: true,
            });
            let mut buffer_view = staging.slice(..).get_mapped_range_mut();
            buffer_view[..mem::size_of::<U>()].copy_from_slice(bytemuck::bytes_of(uniforms));
            drop(buffer_view);
            staging.unmap();

            encoder.copy_buffer_to_buffer(
                &staging,
                0,
                self.uniforms.as_ref().unwrap(),
                0,
                mem::size_of::<U>() as u64,
            );
        }

        let mut cpass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { label: None });
        cpass.set_pipeline(&self.bindgroup_pipeline.as_ref().unwrap().1);
        cpass.set_bind_group(0, &self.bindgroup_pipeline.as_ref().unwrap().0, &[]);
        cpass.dispatch(dimensions.0, dimensions.1, dimensions.2);
    }
}
<|endoftext|>


<commit_before>//! Tests for public/private dependencies.

use cargo_test_support::registry::Package;
use cargo_test_support::{is_nightly, project};

#[cargo_test]
fn exported_priv_warning() {
    if !is_nightly() {
        return;
    }
    Package::new("priv_dep", "0.1.0")
        .file("src/lib.rs", "pub struct FromPriv;")
        .publish();

    let p = project()
        .file(
            "Cargo.toml",
            r#"
            cargo-features = ["public-dependency"]

            [package]
            name = "foo"
            version = "0.0.1"

            [dependencies]
            priv_dep = "0.1.0"
        "#,
        )
        .file(
            "src/lib.rs",
            "
            extern crate priv_dep;
            pub fn use_priv(_: priv_dep::FromPriv) {}
        ",
        )
        .build();

    p.cargo("build --message-format=short")
        .masquerade_as_nightly_cargo()
        .with_stderr_contains(
            "\
src/lib.rs:3:13: warning: type `priv_dep::FromPriv` from private dependency 'priv_dep' in public interface
"
        )
        .run()
}

#[cargo_test]
fn exported_pub_dep() {
    if !is_nightly() {
        return;
    }
    Package::new("pub_dep", "0.1.0")
        .file("src/lib.rs", "pub struct FromPub;")
        .publish();

    let p = project()
        .file(
            "Cargo.toml",
            r#"
            cargo-features = ["public-dependency"]

            [package]
            name = "foo"
            version = "0.0.1"

            [dependencies]
            pub_dep = {version = "0.1.0", public = true}
        "#,
        )
        .file(
            "src/lib.rs",
            "
            extern crate pub_dep;
            pub fn use_pub(_: pub_dep::FromPub) {}
        ",
        )
        .build();

    p.cargo("build --message-format=short")
        .masquerade_as_nightly_cargo()
        .with_stderr(
            "\
[UPDATING] `[..]` index
[DOWNLOADING] crates ...
[DOWNLOADED] pub_dep v0.1.0 ([..])
[COMPILING] pub_dep v0.1.0
[COMPILING] foo v0.0.1 ([CWD])
[FINISHED] dev [unoptimized + debuginfo] target(s) in [..]
",
        )
        .run()
}

#[cargo_test]
pub fn requires_nightly_cargo() {
    let p = project()
        .file(
            "Cargo.toml",
            r#"
            cargo-features = ["public-dependency"]
        "#,
        )
        .file("src/lib.rs", "")
        .build();

    p.cargo("build --message-format=short")
        .with_status(101)
        .with_stderr(
            "\
error: failed to parse manifest at `[..]`

Caused by:
  the cargo feature `public-dependency` requires a nightly version of Cargo, but this is the `stable` channel
  See https://doc.rust-lang.org/book/appendix-07-nightly-rust.html for more information about Rust release channels.
"
        )
        .run()
}

#[cargo_test]
fn requires_feature() {
    Package::new("pub_dep", "0.1.0")
        .file("src/lib.rs", "")
        .publish();

    let p = project()
        .file(
            "Cargo.toml",
            r#"

            [package]
            name = "foo"
            version = "0.0.1"

            [dependencies]
            pub_dep = { version = "0.1.0", public = true }
        "#,
        )
        .file("src/lib.rs", "")
        .build();

    p.cargo("build --message-format=short")
        .masquerade_as_nightly_cargo()
        .with_status(101)
        .with_stderr(
            "\
error: failed to parse manifest at `[..]`

Caused by:
  feature `public-dependency` is required

  consider adding `cargo-features = [\"public-dependency\"]` to the manifest
",
        )
        .run()
}

#[cargo_test]
fn pub_dev_dependency() {
    Package::new("pub_dep", "0.1.0")
        .file("src/lib.rs", "pub struct FromPub;")
        .publish();

    let p = project()
        .file(
            "Cargo.toml",
            r#"
            cargo-features = ["public-dependency"]

            [package]
            name = "foo"
            version = "0.0.1"

            [dev-dependencies]
            pub_dep = {version = "0.1.0", public = true}
        "#,
        )
        .file(
            "src/lib.rs",
            "
            extern crate pub_dep;
            pub fn use_pub(_: pub_dep::FromPub) {}
        ",
        )
        .build();

    p.cargo("build --message-format=short")
        .masquerade_as_nightly_cargo()
        .with_status(101)
        .with_stderr(
            "\
error: failed to parse manifest at `[..]`

Caused by:
  'public' specifier can only be used on regular dependencies, not Development dependencies
",
        )
        .run()
}
<commit_msg>Fix nightly exported_priv_warning test.<commit_after>//! Tests for public/private dependencies.

use cargo_test_support::registry::Package;
use cargo_test_support::{is_nightly, project};

#[cargo_test]
fn exported_priv_warning() {
    if !is_nightly() {
        return;
    }
    Package::new("priv_dep", "0.1.0")
        .file("src/lib.rs", "pub struct FromPriv;")
        .publish();

    let p = project()
        .file(
            "Cargo.toml",
            r#"
            cargo-features = ["public-dependency"]

            [package]
            name = "foo"
            version = "0.0.1"

            [dependencies]
            priv_dep = "0.1.0"
        "#,
        )
        .file(
            "src/lib.rs",
            "
            extern crate priv_dep;
            pub fn use_priv(_: priv_dep::FromPriv) {}
        ",
        )
        .build();

    p.cargo("build --message-format=short")
        .masquerade_as_nightly_cargo()
        .with_stderr_contains(
            "\
src/lib.rs:3:13: warning: type `[..]FromPriv` from private dependency 'priv_dep' in public interface
",
        )
        .run()
}

#[cargo_test]
fn exported_pub_dep() {
    if !is_nightly() {
        return;
    }
    Package::new("pub_dep", "0.1.0")
        .file("src/lib.rs", "pub struct FromPub;")
        .publish();

    let p = project()
        .file(
            "Cargo.toml",
            r#"
            cargo-features = ["public-dependency"]

            [package]
            name = "foo"
            version = "0.0.1"

            [dependencies]
            pub_dep = {version = "0.1.0", public = true}
        "#,
        )
        .file(
            "src/lib.rs",
            "
            extern crate pub_dep;
            pub fn use_pub(_: pub_dep::FromPub) {}
        ",
        )
        .build();

    p.cargo("build --message-format=short")
        .masquerade_as_nightly_cargo()
        .with_stderr(
            "\
[UPDATING] `[..]` index
[DOWNLOADING] crates ...
[DOWNLOADED] pub_dep v0.1.0 ([..])
[COMPILING] pub_dep v0.1.0
[COMPILING] foo v0.0.1 ([CWD])
[FINISHED] dev [unoptimized + debuginfo] target(s) in [..]
",
        )
        .run()
}

#[cargo_test]
pub fn requires_nightly_cargo() {
    let p = project()
        .file(
            "Cargo.toml",
            r#"
            cargo-features = ["public-dependency"]
        "#,
        )
        .file("src/lib.rs", "")
        .build();

    p.cargo("build --message-format=short")
        .with_status(101)
        .with_stderr(
            "\
error: failed to parse manifest at `[..]`

Caused by:
  the cargo feature `public-dependency` requires a nightly version of Cargo, but this is the `stable` channel
  See https://doc.rust-lang.org/book/appendix-07-nightly-rust.html for more information about Rust release channels.
"
        )
        .run()
}

#[cargo_test]
fn requires_feature() {
    Package::new("pub_dep", "0.1.0")
        .file("src/lib.rs", "")
        .publish();

    let p = project()
        .file(
            "Cargo.toml",
            r#"

            [package]
            name = "foo"
            version = "0.0.1"

            [dependencies]
            pub_dep = { version = "0.1.0", public = true }
        "#,
        )
        .file("src/lib.rs", "")
        .build();

    p.cargo("build --message-format=short")
        .masquerade_as_nightly_cargo()
        .with_status(101)
        .with_stderr(
            "\
error: failed to parse manifest at `[..]`

Caused by:
  feature `public-dependency` is required

  consider adding `cargo-features = [\"public-dependency\"]` to the manifest
",
        )
        .run()
}

#[cargo_test]
fn pub_dev_dependency() {
    Package::new("pub_dep", "0.1.0")
        .file("src/lib.rs", "pub struct FromPub;")
        .publish();

    let p = project()
        .file(
            "Cargo.toml",
            r#"
            cargo-features = ["public-dependency"]

            [package]
            name = "foo"
            version = "0.0.1"

            [dev-dependencies]
            pub_dep = {version = "0.1.0", public = true}
        "#,
        )
        .file(
            "src/lib.rs",
            "
            extern crate pub_dep;
            pub fn use_pub(_: pub_dep::FromPub) {}
        ",
        )
        .build();

    p.cargo("build --message-format=short")
        .masquerade_as_nightly_cargo()
        .with_status(101)
        .with_stderr(
            "\
error: failed to parse manifest at `[..]`

Caused by:
  'public' specifier can only be used on regular dependencies, not Development dependencies
",
        )
        .run()
}
<|endoftext|>


<commit_before>/// Used by `SQLBindParameter`.
#[repr(u16)]
#[allow(non_camel_case_types)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum InputOutput {
    SQL_PARAM_TYPE_UNKNOWN = 0,
    SQL_PARAM_INPUT = 1,
    SQL_PARAM_INPUT_OUTPUT = 2,
    SQL_RESULT_COL = 3,
    SQL_PARAM_OUTPUT = 4,
    SQL_RETURN_VALUE = 5,
    #[cfg(feature = "odbc_version_3_80")]
    SQL_PARAM_INPUT_OUTPUT_STREAM = 8,
    #[cfg(feature = "odbc_version_3_80")]
    SQL_PARAM_OUTPUT_STREAM = 16,
}
pub use InputOutput::*;
<commit_msg>change InputOutput representation type to i16 (SQLSMALLINT)<commit_after>/// Used by `SQLBindParameter`.
#[repr(i16)]
#[allow(non_camel_case_types)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum InputOutput {
    SQL_PARAM_TYPE_UNKNOWN = 0,
    SQL_PARAM_INPUT = 1,
    SQL_PARAM_INPUT_OUTPUT = 2,
    SQL_RESULT_COL = 3,
    SQL_PARAM_OUTPUT = 4,
    SQL_RETURN_VALUE = 5,
    #[cfg(feature = "odbc_version_3_80")]
    SQL_PARAM_INPUT_OUTPUT_STREAM = 8,
    #[cfg(feature = "odbc_version_3_80")]
    SQL_PARAM_OUTPUT_STREAM = 16,
}
pub use InputOutput::*;
<|endoftext|>


<commit_before>extern crate lz4;
extern crate filetime;

use std::env;
use std::fmt::{Display, Formatter};
use std::ffi::OsString;
use std::fs;
use std::fs::{File, PathExt, OpenOptions};
use std::io::{Error, ErrorKind, Read, Write, Seek, SeekFrom};
use std::hash::{Hasher, SipHasher};
use std::path::{Path, PathBuf};
use std::sync::RwLock;

use self::filetime::FileTime;

use super::super::cache::FileHasher;
use super::super::compiler::OutputInfo;
use super::super::utils::DEFAULT_BUF_SIZE;
use super::binary::*;
use super::statistic::Statistic;
use super::counter::Counter;

const HEADER: &'static [u8] = b"OBCF\x00\x02";
const FOOTER: &'static [u8] = b"END\x00";
const SUFFIX: &'static str = ".lz4";

#[derive(Debug)]
pub enum CacheError {
	InvalidHeader(PathBuf),
	InvalidFooter(PathBuf),
	PackedFilesMismatch(PathBuf),
	MutexError(String),
}

impl Display for CacheError {
	fn fmt(&self, f: &mut Formatter) -> Result<(), ::std::fmt::Error> {
		match self {
			&CacheError::InvalidHeader(ref path) => write!(f, "invalid cache file header: {}", path.display()),
			&CacheError::InvalidFooter(ref path) => write!(f, "invalid cache file footer: {}", path.display()),
			&CacheError::PackedFilesMismatch(ref path) => write!(f, "unexpected count of packed cached files: {}", path.display()),
			&CacheError::MutexError(ref message) => write!(f, "mutex error: {}", message),
		}
	}
}

impl ::std::error::Error for CacheError {
	fn description(&self) -> &str {
		match self {
			&CacheError::InvalidHeader(_) => "invalid cache file header",
			&CacheError::InvalidFooter(_) => "invalid cache file footer",
			&CacheError::PackedFilesMismatch(_) => "unexpected count of packed cached files",
			&CacheError::MutexError(_) => "mutex error",
		}
	}

	fn cause(&self) -> Option<&::std::error::Error> {
		None
	}
}

#[derive(Clone)]
pub struct FileCache {
	cache_dir: PathBuf,
}

struct CacheFile {
	path: PathBuf,
	size: u64,
	accessed: FileTime,
}

impl FileCache {
	pub fn new() -> Self {
		let cache_dir = match env::var("OCTOBUILD_CACHE") {
			Ok(value) => Path::new(&value).to_path_buf(),
			Err(_) => env::home_dir().unwrap().join(".octobuild").join("cache")
		};
		FileCache {
			cache_dir: cache_dir,
		}
	}

	pub fn run_cached<F: Fn()->Result<OutputInfo, Error>, C: Fn()->bool>(&self, file_hasher: &FileHasher, statistic: &RwLock<Statistic>, params: u64, inputs: &Vec<PathBuf>, outputs: &Vec<PathBuf>, worker: F, checker: C) -> Result<OutputInfo, Error> {
		let hash = try! (self.generate_hash(file_hasher, params, inputs));
		let path = self.cache_dir.join(&hash[0..2]).join(&(hash[2..].to_string() + SUFFIX));
		// Try to read data from cache.
		match read_cache(statistic, &path, outputs) {
			Ok(output) => {return Ok(output)}
			Err(_) => {}
		}
		// Run task and save result to cache.
		let output = try !(worker());
		if checker() {
			try !(write_cache(statistic, &path, outputs, &output));
		}
		Ok(output)
	}

	pub fn cleanup(&self, max_cache_size: u64) -> Result<(), Error> {
		let mut files = try! (find_cache_files(&self.cache_dir, Vec::new()));
		files.sort_by(|a, b| b.accessed.cmp(&a.accessed));

		let mut cache_size: u64 = 0;
		for item in files.into_iter() {
			cache_size += item.size;
			if cache_size > max_cache_size {
				let _ = try!(fs::remove_file(&item.path));
			}
		}
		Ok(())
	}

	fn generate_hash(&self, file_hasher: &FileHasher, params: u64, inputs: &Vec<PathBuf>) -> Result<String, Error> {
		let mut hash = SipHasher::new();
		// params
		hash.write_u64(params);
		// inputs
		for input in inputs.iter() {
			let file_hash = try! (file_hasher.file_hash(input));
			hash.write(file_hash.as_bytes());
		}
		Ok(format!("{:016x}", hash.finish()))
	}
}

fn find_cache_files(dir: &Path, mut files: Vec<CacheFile>) -> Result<Vec<CacheFile>, Error> {
	for entry in try!(fs::read_dir(dir)) {
		let entry = try!(entry);
		let path = entry.path();
		let stat = try! (fs::metadata(&path));
		if stat.is_dir() {
			let r = find_cache_files(&path, files);
			files = try! (r);
		} else {
			files.push(CacheFile {
				path: path,
				size: stat.len(),
				accessed: FileTime::from_last_modification_time(&stat),
			});
		}
	}
	Ok(files)
}

fn write_cached_file<W: Write>(stream: &mut W, path: &PathBuf) -> Result<(), Error> {
	let mut buf: [u8; DEFAULT_BUF_SIZE] = [0; DEFAULT_BUF_SIZE];
	let mut file = try! (File::open(path));
	loop {
		let size = try! (file.read(&mut buf));
		if size <= 0 {
			break;
		}
		try! (write_usize(stream, size));
		try! (stream.write_all(&buf[0..size]));
	}
	try! (write_usize(stream, 0));
	Ok(())
}

fn write_cache(statistic: &RwLock<Statistic>, path: &Path, paths: &Vec<PathBuf>, output: &OutputInfo) -> Result<(), Error> {
	if !output.success() {
		return Ok(());
	}
	match path.parent() {
		Some(parent) => try! (fs::create_dir_all(&parent)),
		None => ()
	}
	let mut stream = try! (lz4::EncoderBuilder::new().level(1).build(Counter::writer(try! (File::create(path)))));
	try! (stream.write_all(HEADER));
	try! (write_usize(&mut stream, paths.len()));
	for path in paths.iter() {
		try! (write_cached_file(&mut stream, path));
	}
	try! (write_output(&mut stream, output));
	try! (stream.write_all(FOOTER));
	match stream.finish() {
		(writer, result) => {
			statistic.write().unwrap().add_miss(writer.len());
			result
		}
	}
}

fn read_cached_file<R: Read>(stream: &mut R, path: &PathBuf) -> Result<(), Error> {
	let mut file = try! (File::create(path));
	loop {
		let size = try! (read_usize(stream));
		if size == 0 {break;}
		let block = try! (read_exact(stream, size));
		try! (file.write_all(&block));
	}
	Ok(())
}

fn read_cache(statistic: &RwLock<Statistic>, path: &Path, paths: &Vec<PathBuf>) -> Result<OutputInfo, Error> {
	let mut file = try! (OpenOptions::new().read(true).write(true).open(Path::new(path)));
	try! (file.write(&[4]));
	try! (file.seek(SeekFrom::Start(0)));
	let mut stream = try! (lz4::Decoder::new (Counter::reader(file)));
	if try! (read_exact(&mut stream, HEADER.len())) != HEADER {
		return Err(Error::new(ErrorKind::InvalidInput, CacheError::InvalidHeader(path.to_path_buf())));
	}
	if try! (read_usize(&mut stream)) != paths.len() {
		return Err(Error::new(ErrorKind::InvalidInput, CacheError::PackedFilesMismatch(path.to_path_buf())));
	}
	for path in paths.iter() {
		let mut temp_name = OsString::from("~tmp~");
		temp_name.push(path.file_name().unwrap());
		let temp = path.with_file_name(temp_name);
		drop(fs::remove_file(&path));
		match read_cached_file(&mut stream, &temp).and_then(|_| fs::rename(&temp, &path)) {
			Ok(_) => {
			}
			Err(e) => {
				drop(fs::remove_file(&temp));
				return Err(e);
			}
		};
	}
	let output = try! (read_output(&mut stream));
	if try! (read_exact(&mut stream, FOOTER.len())) != FOOTER {
		return Err(Error::new(ErrorKind::InvalidInput, CacheError::InvalidFooter(path.to_path_buf())));
	}
	let mut eof = [0];
	if try! (stream.read(&mut eof)) != 0 {
		return Err(Error::new(ErrorKind::InvalidInput, CacheError::InvalidFooter(path.to_path_buf())));
	}
	statistic.write().unwrap().add_hit(stream.finish().0.len());
	Ok(output)
}

fn write_blob(stream: &mut Write, blob: &[u8]) -> Result<(), Error> {
	try! (write_usize(stream, blob.len()));
	try! (stream.write_all(blob));
	Ok(())
}

fn read_blob(stream: &mut Read) -> Result<Vec<u8>, Error> {
	let size = try! (read_usize(stream));
	read_exact(stream, size)
}

fn write_output(stream: &mut Write, output: &OutputInfo) -> Result<(), Error> {
	try! (write_blob(stream, &output.stdout));
	try! (write_blob(stream, &output.stderr));
	Ok(())
}

fn read_output(stream: &mut Read) -> Result<OutputInfo, Error> {
	let stdout = try! (read_blob(stream));
	let stderr = try! (read_blob(stream));
	Ok(OutputInfo {
		status: Some(0),
		stdout: stdout,
		stderr: stderr,
	})
}
<commit_msg>Rustup: 1.5.0<commit_after>extern crate lz4;
extern crate filetime;

use std::env;
use std::fmt::{Display, Formatter};
use std::ffi::OsString;
use std::fs;
use std::fs::{File, OpenOptions};
use std::io::{Error, ErrorKind, Read, Write, Seek, SeekFrom};
use std::hash::{Hasher, SipHasher};
use std::path::{Path, PathBuf};
use std::sync::RwLock;

use self::filetime::FileTime;

use super::super::cache::FileHasher;
use super::super::compiler::OutputInfo;
use super::super::utils::DEFAULT_BUF_SIZE;
use super::binary::*;
use super::statistic::Statistic;
use super::counter::Counter;

const HEADER: &'static [u8] = b"OBCF\x00\x02";
const FOOTER: &'static [u8] = b"END\x00";
const SUFFIX: &'static str = ".lz4";

#[derive(Debug)]
pub enum CacheError {
	InvalidHeader(PathBuf),
	InvalidFooter(PathBuf),
	PackedFilesMismatch(PathBuf),
	MutexError(String),
}

impl Display for CacheError {
	fn fmt(&self, f: &mut Formatter) -> Result<(), ::std::fmt::Error> {
		match self {
			&CacheError::InvalidHeader(ref path) => write!(f, "invalid cache file header: {}", path.display()),
			&CacheError::InvalidFooter(ref path) => write!(f, "invalid cache file footer: {}", path.display()),
			&CacheError::PackedFilesMismatch(ref path) => write!(f, "unexpected count of packed cached files: {}", path.display()),
			&CacheError::MutexError(ref message) => write!(f, "mutex error: {}", message),
		}
	}
}

impl ::std::error::Error for CacheError {
	fn description(&self) -> &str {
		match self {
			&CacheError::InvalidHeader(_) => "invalid cache file header",
			&CacheError::InvalidFooter(_) => "invalid cache file footer",
			&CacheError::PackedFilesMismatch(_) => "unexpected count of packed cached files",
			&CacheError::MutexError(_) => "mutex error",
		}
	}

	fn cause(&self) -> Option<&::std::error::Error> {
		None
	}
}

#[derive(Clone)]
pub struct FileCache {
	cache_dir: PathBuf,
}

struct CacheFile {
	path: PathBuf,
	size: u64,
	accessed: FileTime,
}

impl FileCache {
	pub fn new() -> Self {
		let cache_dir = match env::var("OCTOBUILD_CACHE") {
			Ok(value) => Path::new(&value).to_path_buf(),
			Err(_) => env::home_dir().unwrap().join(".octobuild").join("cache")
		};
		FileCache {
			cache_dir: cache_dir,
		}
	}

	pub fn run_cached<F: Fn()->Result<OutputInfo, Error>, C: Fn()->bool>(&self, file_hasher: &FileHasher, statistic: &RwLock<Statistic>, params: u64, inputs: &Vec<PathBuf>, outputs: &Vec<PathBuf>, worker: F, checker: C) -> Result<OutputInfo, Error> {
		let hash = try! (self.generate_hash(file_hasher, params, inputs));
		let path = self.cache_dir.join(&hash[0..2]).join(&(hash[2..].to_string() + SUFFIX));
		// Try to read data from cache.
		match read_cache(statistic, &path, outputs) {
			Ok(output) => {return Ok(output)}
			Err(_) => {}
		}
		// Run task and save result to cache.
		let output = try !(worker());
		if checker() {
			try !(write_cache(statistic, &path, outputs, &output));
		}
		Ok(output)
	}

	pub fn cleanup(&self, max_cache_size: u64) -> Result<(), Error> {
		let mut files = try! (find_cache_files(&self.cache_dir, Vec::new()));
		files.sort_by(|a, b| b.accessed.cmp(&a.accessed));

		let mut cache_size: u64 = 0;
		for item in files.into_iter() {
			cache_size += item.size;
			if cache_size > max_cache_size {
				let _ = try!(fs::remove_file(&item.path));
			}
		}
		Ok(())
	}

	fn generate_hash(&self, file_hasher: &FileHasher, params: u64, inputs: &Vec<PathBuf>) -> Result<String, Error> {
		let mut hash = SipHasher::new();
		// params
		hash.write_u64(params);
		// inputs
		for input in inputs.iter() {
			let file_hash = try! (file_hasher.file_hash(input));
			hash.write(file_hash.as_bytes());
		}
		Ok(format!("{:016x}", hash.finish()))
	}
}

fn find_cache_files(dir: &Path, mut files: Vec<CacheFile>) -> Result<Vec<CacheFile>, Error> {
	for entry in try!(fs::read_dir(dir)) {
		let entry = try!(entry);
		let path = entry.path();
		let stat = try! (fs::metadata(&path));
		if stat.is_dir() {
			let r = find_cache_files(&path, files);
			files = try! (r);
		} else {
			files.push(CacheFile {
				path: path,
				size: stat.len(),
				accessed: FileTime::from_last_modification_time(&stat),
			});
		}
	}
	Ok(files)
}

fn write_cached_file<W: Write>(stream: &mut W, path: &PathBuf) -> Result<(), Error> {
	let mut buf: [u8; DEFAULT_BUF_SIZE] = [0; DEFAULT_BUF_SIZE];
	let mut file = try! (File::open(path));
	loop {
		let size = try! (file.read(&mut buf));
		if size <= 0 {
			break;
		}
		try! (write_usize(stream, size));
		try! (stream.write_all(&buf[0..size]));
	}
	try! (write_usize(stream, 0));
	Ok(())
}

fn write_cache(statistic: &RwLock<Statistic>, path: &Path, paths: &Vec<PathBuf>, output: &OutputInfo) -> Result<(), Error> {
	if !output.success() {
		return Ok(());
	}
	match path.parent() {
		Some(parent) => try! (fs::create_dir_all(&parent)),
		None => ()
	}
	let mut stream = try! (lz4::EncoderBuilder::new().level(1).build(Counter::writer(try! (File::create(path)))));
	try! (stream.write_all(HEADER));
	try! (write_usize(&mut stream, paths.len()));
	for path in paths.iter() {
		try! (write_cached_file(&mut stream, path));
	}
	try! (write_output(&mut stream, output));
	try! (stream.write_all(FOOTER));
	match stream.finish() {
		(writer, result) => {
			statistic.write().unwrap().add_miss(writer.len());
			result
		}
	}
}

fn read_cached_file<R: Read>(stream: &mut R, path: &PathBuf) -> Result<(), Error> {
	let mut file = try! (File::create(path));
	loop {
		let size = try! (read_usize(stream));
		if size == 0 {break;}
		let block = try! (read_exact(stream, size));
		try! (file.write_all(&block));
	}
	Ok(())
}

fn read_cache(statistic: &RwLock<Statistic>, path: &Path, paths: &Vec<PathBuf>) -> Result<OutputInfo, Error> {
	let mut file = try! (OpenOptions::new().read(true).write(true).open(Path::new(path)));
	try! (file.write(&[4]));
	try! (file.seek(SeekFrom::Start(0)));
	let mut stream = try! (lz4::Decoder::new (Counter::reader(file)));
	if try! (read_exact(&mut stream, HEADER.len())) != HEADER {
		return Err(Error::new(ErrorKind::InvalidInput, CacheError::InvalidHeader(path.to_path_buf())));
	}
	if try! (read_usize(&mut stream)) != paths.len() {
		return Err(Error::new(ErrorKind::InvalidInput, CacheError::PackedFilesMismatch(path.to_path_buf())));
	}
	for path in paths.iter() {
		let mut temp_name = OsString::from("~tmp~");
		temp_name.push(path.file_name().unwrap());
		let temp = path.with_file_name(temp_name);
		drop(fs::remove_file(&path));
		match read_cached_file(&mut stream, &temp).and_then(|_| fs::rename(&temp, &path)) {
			Ok(_) => {
			}
			Err(e) => {
				drop(fs::remove_file(&temp));
				return Err(e);
			}
		};
	}
	let output = try! (read_output(&mut stream));
	if try! (read_exact(&mut stream, FOOTER.len())) != FOOTER {
		return Err(Error::new(ErrorKind::InvalidInput, CacheError::InvalidFooter(path.to_path_buf())));
	}
	let mut eof = [0];
	if try! (stream.read(&mut eof)) != 0 {
		return Err(Error::new(ErrorKind::InvalidInput, CacheError::InvalidFooter(path.to_path_buf())));
	}
	statistic.write().unwrap().add_hit(stream.finish().0.len());
	Ok(output)
}

fn write_blob(stream: &mut Write, blob: &[u8]) -> Result<(), Error> {
	try! (write_usize(stream, blob.len()));
	try! (stream.write_all(blob));
	Ok(())
}

fn read_blob(stream: &mut Read) -> Result<Vec<u8>, Error> {
	let size = try! (read_usize(stream));
	read_exact(stream, size)
}

fn write_output(stream: &mut Write, output: &OutputInfo) -> Result<(), Error> {
	try! (write_blob(stream, &output.stdout));
	try! (write_blob(stream, &output.stderr));
	Ok(())
}

fn read_output(stream: &mut Read) -> Result<OutputInfo, Error> {
	let stdout = try! (read_blob(stream));
	let stderr = try! (read_blob(stream));
	Ok(OutputInfo {
		status: Some(0),
		stdout: stdout,
		stderr: stderr,
	})
}
<|endoftext|>