<commit_before>//! Store buffer when failing to send events.

use std::fs::File;
use std::path::PathBuf;
use std::io;
use std::io::Write;
use std::fmt::Debug;
use retry_conf::RetryConf;
use error::FluentError;
use forwardable::forward::Forward;
use std::fs::OpenOptions;
use record::Record;
use event_record::EventRecord;
use serde::ser::Serialize;

/// Create file with write, create, append, and open option.
fn ensure_file_with_wca(path: PathBuf) -> Result<File, io::Error> {
    let file = OpenOptions::new().write(true).create(true).append(true).open(path)?;
    Ok(file)
}

/// Write event buffer into file with TSV format.
pub fn maybe_write_record<T>(conf: &RetryConf,
                             record: Record<T>,
                             err: FluentError)
                             -> Result<(), FluentError>
    where T: Serialize + Debug
{
    let store_needed = conf.clone().need_to_store();
    let store_path = conf.clone().store_path();
    if store_needed {
        match ensure_file_with_wca(store_path.clone().unwrap()) {
            Ok(mut f) => {
                let mut w = Vec::new();
                write!(&mut w, "{}", record.dump()).unwrap();
                f.write_all(&w)?;
                f.sync_data()?;
                Err(FluentError::FileStored(format!("stored buffer in specified file: \
                                                     {:?}",
                                                    store_path.unwrap())))
            }
            Err(e) => Err(From::from(e)),
        }
    } else {
        Err(err)
    }
}

/// Write event buffer into file with TSV format.
pub fn maybe_write_event_record<T>(conf: &RetryConf,
                             record: EventRecord<T>,
                             err: FluentError)
                             -> Result<(), FluentError>
    where T: Serialize + Debug
{
    let store_needed = conf.clone().need_to_store();
    let store_path = conf.clone().store_path();
    if store_needed {
        match ensure_file_with_wca(store_path.clone().unwrap()) {
            Ok(mut f) => {
                let mut w = Vec::new();
                write!(&mut w, "{}", record.dump()).unwrap();
                f.write_all(&w)?;
                f.sync_data()?;
                Err(FluentError::FileStored(format!("stored buffer in specified file: \
                                                     {:?}",
                                                    store_path.unwrap())))
            }
            Err(e) => Err(From::from(e)),
        }
    } else {
        Err(err)
    }
}

/// Write events buffer into file with TSV format.
pub fn maybe_write_records<T>(conf: &RetryConf,
                              forward: Forward<T>,
                              err: FluentError)
                              -> Result<(), FluentError>
    where T: Serialize + Debug
{
    let store_needed = conf.clone().need_to_store();
    let store_path = conf.clone().store_path();
    if store_needed {
        match ensure_file_with_wca(store_path.clone().unwrap()) {
            Ok(mut f) => {
                let mut w = Vec::new();
                write!(&mut w, "{}", forward.dump()).unwrap();
                f.write_all(&w)?;
                f.sync_data()?;
                Err(FluentError::FileStored(format!("stored buffer in specified file: \
                                                     {:?}",
                                                    store_path.unwrap())))
            }
            Err(e) => Err(From::from(e)),
        }
    } else {
        Err(err)
    }
}

#[cfg(test)]
mod tests {
    extern crate tempdir;
    use super::*;
    use time;
    use std::collections::HashMap;
    use self::tempdir::TempDir;
    use record::Record;
    use retry_conf::RetryConf;
    use error::FluentError;
    use forwardable::forward::Forward;
    use event_time::EventTime;

    #[test]
    fn test_write_record() {
        let tag = "fruently".to_string();
        let time = time::now();
        let mut obj: HashMap<String, String> = HashMap::new();
        obj.insert("name".to_string(), "fruently".to_string());
        let record = Record::new(tag.clone(), time, obj.clone());
        let tmp = TempDir::new("fruently").unwrap().into_path().join("buffer");
        let conf = RetryConf::new().store_file(tmp.clone());
        assert!(maybe_write_record(&conf, record, FluentError::Dummy("dummy".to_string()))
            .is_err());
        assert!(tmp.exists())
    }

    #[test]
    fn test_write_record_2_times() {
        let tag = "fruently".to_string();
        let time = time::now();
        let mut obj: HashMap<String, String> = HashMap::new();
        obj.insert("name".to_string(), "fruently".to_string());
        let record = Record::new(tag.clone(), time, obj.clone());
        let tmp = TempDir::new("fruently").unwrap().into_path().join("buffer");
        let conf = RetryConf::new().store_file(tmp.clone());
        assert!(maybe_write_record(&conf, record, FluentError::Dummy("dummy".to_string()))
            .is_err());
        assert!(tmp.exists());
        let mut obj2: HashMap<String, String> = HashMap::new();
        obj2.insert("name2".to_string(), "fruently2".to_string());
        let record2 = Record::new(tag.clone(), time, obj2.clone());
        let conf2 = RetryConf::new().store_file(tmp.clone());
        assert!(maybe_write_record(&conf2, record2, FluentError::Dummy("dummy".to_string()))
            .is_err());
        assert!(tmp.exists())
    }

    #[test]
    fn test_write_records() {
        let tag = "fruently".to_string();
        let mut obj1: HashMap<String, String> = HashMap::new();
        obj1.insert("hey".to_string(), "Rust with forward mode!".to_string());
        let mut obj2: HashMap<String, String> = HashMap::new();
        obj2.insert("yeah".to_string(), "Also sent together!".to_string());
        let time = time::now();
        let entry = (EventTime::new(time), obj1);
        let entry2 = (EventTime::new(time), obj2);
        let entries = vec![(entry), (entry2)];
        let forward = Forward::new(tag, entries);
        let tmp = TempDir::new("fruently").unwrap().into_path().join("buffer");
        let conf = RetryConf::new().store_file(tmp.clone());
        assert!(maybe_write_records(&conf, forward, FluentError::Dummy("dummy".to_string()))
            .is_err());
        assert!(tmp.exists())
    }
}
<commit_msg>Add test case for EventRecord<commit_after>//! Store buffer when failing to send events.

use std::fs::File;
use std::path::PathBuf;
use std::io;
use std::io::Write;
use std::fmt::Debug;
use retry_conf::RetryConf;
use error::FluentError;
use forwardable::forward::Forward;
use std::fs::OpenOptions;
use record::Record;
use event_record::EventRecord;
use serde::ser::Serialize;

/// Create file with write, create, append, and open option.
fn ensure_file_with_wca(path: PathBuf) -> Result<File, io::Error> {
    let file = OpenOptions::new().write(true).create(true).append(true).open(path)?;
    Ok(file)
}

/// Write event buffer into file with TSV format.
pub fn maybe_write_record<T>(conf: &RetryConf,
                             record: Record<T>,
                             err: FluentError)
                             -> Result<(), FluentError>
    where T: Serialize + Debug
{
    let store_needed = conf.clone().need_to_store();
    let store_path = conf.clone().store_path();
    if store_needed {
        match ensure_file_with_wca(store_path.clone().unwrap()) {
            Ok(mut f) => {
                let mut w = Vec::new();
                write!(&mut w, "{}", record.dump()).unwrap();
                f.write_all(&w)?;
                f.sync_data()?;
                Err(FluentError::FileStored(format!("stored buffer in specified file: \
                                                     {:?}",
                                                    store_path.unwrap())))
            }
            Err(e) => Err(From::from(e)),
        }
    } else {
        Err(err)
    }
}

/// Write event buffer into file with TSV format.
pub fn maybe_write_event_record<T>(conf: &RetryConf,
                             record: EventRecord<T>,
                             err: FluentError)
                             -> Result<(), FluentError>
    where T: Serialize + Debug
{
    let store_needed = conf.clone().need_to_store();
    let store_path = conf.clone().store_path();
    if store_needed {
        match ensure_file_with_wca(store_path.clone().unwrap()) {
            Ok(mut f) => {
                let mut w = Vec::new();
                write!(&mut w, "{}", record.dump()).unwrap();
                f.write_all(&w)?;
                f.sync_data()?;
                Err(FluentError::FileStored(format!("stored buffer in specified file: \
                                                     {:?}",
                                                    store_path.unwrap())))
            }
            Err(e) => Err(From::from(e)),
        }
    } else {
        Err(err)
    }
}

/// Write events buffer into file with TSV format.
pub fn maybe_write_records<T>(conf: &RetryConf,
                              forward: Forward<T>,
                              err: FluentError)
                              -> Result<(), FluentError>
    where T: Serialize + Debug
{
    let store_needed = conf.clone().need_to_store();
    let store_path = conf.clone().store_path();
    if store_needed {
        match ensure_file_with_wca(store_path.clone().unwrap()) {
            Ok(mut f) => {
                let mut w = Vec::new();
                write!(&mut w, "{}", forward.dump()).unwrap();
                f.write_all(&w)?;
                f.sync_data()?;
                Err(FluentError::FileStored(format!("stored buffer in specified file: \
                                                     {:?}",
                                                    store_path.unwrap())))
            }
            Err(e) => Err(From::from(e)),
        }
    } else {
        Err(err)
    }
}

#[cfg(test)]
mod tests {
    extern crate tempdir;
    use super::*;
    use time;
    use std::collections::HashMap;
    use self::tempdir::TempDir;
    use record::Record;
    use retry_conf::RetryConf;
    use error::FluentError;
    use forwardable::forward::Forward;
    use event_time::EventTime;

    #[test]
    fn test_write_record() {
        let tag = "fruently".to_string();
        let time = time::now();
        let mut obj: HashMap<String, String> = HashMap::new();
        obj.insert("name".to_string(), "fruently".to_string());
        let record = Record::new(tag.clone(), time, obj.clone());
        let tmp = TempDir::new("fruently").unwrap().into_path().join("buffer");
        let conf = RetryConf::new().store_file(tmp.clone());
        assert!(maybe_write_record(&conf, record, FluentError::Dummy("dummy".to_string()))
            .is_err());
        assert!(tmp.exists())
    }

    #[test]
    fn test_write_event_record() {
        let tag = "fruently".to_string();
        let time = time::now();
        let mut obj: HashMap<String, String> = HashMap::new();
        obj.insert("name".to_string(), "fruently".to_string());
        let record = EventRecord::new(tag.clone(), time, obj.clone());
        let tmp = TempDir::new("fruently").unwrap().into_path().join("buffer");
        let conf = RetryConf::new().store_file(tmp.clone());
        assert!(maybe_write_event_record(&conf, record, FluentError::Dummy("dummy".to_string()))
            .is_err());
        assert!(tmp.exists())
    }

    #[test]
    fn test_write_record_2_times() {
        let tag = "fruently".to_string();
        let time = time::now();
        let mut obj: HashMap<String, String> = HashMap::new();
        obj.insert("name".to_string(), "fruently".to_string());
        let record = Record::new(tag.clone(), time, obj.clone());
        let tmp = TempDir::new("fruently").unwrap().into_path().join("buffer");
        let conf = RetryConf::new().store_file(tmp.clone());
        assert!(maybe_write_record(&conf, record, FluentError::Dummy("dummy".to_string()))
            .is_err());
        assert!(tmp.exists());
        let mut obj2: HashMap<String, String> = HashMap::new();
        obj2.insert("name2".to_string(), "fruently2".to_string());
        let record2 = Record::new(tag.clone(), time, obj2.clone());
        let conf2 = RetryConf::new().store_file(tmp.clone());
        assert!(maybe_write_record(&conf2, record2, FluentError::Dummy("dummy".to_string()))
            .is_err());
        assert!(tmp.exists())
    }

    #[test]
    fn test_write_records() {
        let tag = "fruently".to_string();
        let mut obj1: HashMap<String, String> = HashMap::new();
        obj1.insert("hey".to_string(), "Rust with forward mode!".to_string());
        let mut obj2: HashMap<String, String> = HashMap::new();
        obj2.insert("yeah".to_string(), "Also sent together!".to_string());
        let time = time::now();
        let entry = (EventTime::new(time), obj1);
        let entry2 = (EventTime::new(time), obj2);
        let entries = vec![(entry), (entry2)];
        let forward = Forward::new(tag, entries);
        let tmp = TempDir::new("fruently").unwrap().into_path().join("buffer");
        let conf = RetryConf::new().store_file(tmp.clone());
        assert!(maybe_write_records(&conf, forward, FluentError::Dummy("dummy".to_string()))
            .is_err());
        assert!(tmp.exists())
    }
}
<|endoftext|>


<commit_before>use ast::punc::*;
use ast::token::*;
use std::io::{BufReader, BufferedReader, Buffer, IoError, IoResult, EndOfFile};
use std::char::from_u32;
use std::num::from_str_radix;
use std::from_str::FromStr;
/// A Javascript lexer
pub struct Lexer<B> {
    /// The list of tokens generated so far
    pub tokens : Vec<Token>,
    /// The current line number in the script
    line_number : uint,
    /// The current column number in the script
    column_number : uint,
    /// The reader
    buffer: B,
    /// The peeked character buffer
    peek_buffer: String
}
impl<B:Buffer> Lexer<B> {
    /// Creates a new lexer with empty buffers
    pub fn new(buffer: B) -> Lexer<B> {
        Lexer {
            tokens: Vec::new(),
            line_number: 1,
            column_number: 0,
            buffer: buffer,
            peek_buffer: String::new()
        }
    }
    #[inline(always)]
    fn push_token(&mut self, tk:TokenData) {
        self.tokens.push(Token::new(tk, self.line_number, self.column_number))
    }
    #[inline(always)]
    fn push_punc(&mut self, punc:Punctuator) {
        self.push_token(TPunctuator(punc));
    }
    /// Processes an input stream from a string into an array of tokens
    pub fn lex_str(script:&str) -> Vec<Token> {
        let script_bytes:&[u8] = script.as_bytes();
        let reader = BufReader::new(script_bytes);
        let buf_reader = BufferedReader::new(reader);
        let mut lexer = Lexer::new(buf_reader);
        lexer.lex().unwrap();
        lexer.tokens
    }
    fn next(&mut self) -> IoResult<char> {
        if self.peek_buffer.len() == 0 {
            self.buffer.read_char()
        } else {
            Ok(match self.peek_buffer.pop_char() {
                Some(v) => v,
                None => unreachable!()
            })
        }
    }
    fn next_is(&mut self, peek:char) -> bool {
        if self.peek_buffer.len() > 0 {
            let matched = self.peek_buffer.as_slice().char_at(self.peek_buffer.len() - 1) == peek;
            if matched {
                self.peek_buffer.pop_char();
            }
            matched
        } else {
            match self.buffer.read_char() {
                Ok(ch) if ch == peek => {
                    self.column_number += 1;
                    true
                },
                Ok(ch) if ch != peek => {
                    self.peek_buffer.push_char(ch);
                    false
                },
                _ => false
            }
        }
    }
    /// Processes an input stream from the `buffer` into a vector of tokens
    pub fn lex(&mut self) -> IoResult<()> {
        loop {
            let ch = match self.next() {
                Ok(ch) => ch,
                Err(ref err) if err.kind == EndOfFile => break,
                Err(err) => return Err(err)
            };
            self.column_number += 1;
            match ch {
                '"' | '\'' => {
                    let mut buf = String::new();
                    loop {
                        match try!(self.next()) {
                            '\'' if ch == '\'' => {
                                break;
                            },
                            '"' if ch == '"' => {
                                break;
                            },
                            '\\' => {
                                let escape = try!(self.next());
                                if escape != '\n' {
                                    let escaped_ch = match escape {
                                        'n' => '\n',
                                        'r' => '\r',
                                        't' => '\t',
                                        'b' => '\x08',
                                        'f' => '\x0c',
                                        '0' => '\0',
                                        'x' => {
                                            let mut nums = String::with_capacity(2);
                                            for _ in range(0, 2) {
                                                nums.push_char(try!(self.next()));
                                            }
                                            self.column_number += 2;
                                            let as_num = match from_str_radix(nums.as_slice(), 16) {
                                                Some(v) => v,
                                                None => 0
                                            };
                                            match from_u32(as_num) {
                                                Some(v) => v,
                                                None => fail!("{}:{}: {} is not a valid unicode scalar value", self.line_number, self.column_number, as_num)
                                            }
                                        },
                                        'u' => {
                                            let mut nums = String::new();
                                            for _ in range(0, 4) {
                                                nums.push_char(try!(self.next()));
                                            }
                                            self.column_number += 4;
                                            let as_num = match from_str_radix(nums.as_slice(), 16) {
                                                Some(v) => v,
                                                None => 0
                                            };
                                            match from_u32(as_num) {
                                                Some(v) => v,
                                                None => fail!("{}:{}: {} is not a valid unicode scalar value", self.line_number, self.column_number, as_num)
                                            }
                                        },
                                        '\'' | '"' => escape,
                                        _ => fail!("{}:{}: Invalid escape `{}`", self.line_number, self.column_number, ch)
                                    };
                                    buf.push_char(escaped_ch);
                                }
                            },
                            ch => buf.push_char(ch)
                        }
                    }
                    self.push_token(TStringLiteral(buf))
                },
                '0' if self.next_is('x') => {
                    let mut buf = String::new();
                    loop {
                        match try!(self.next()) {
                            ch if ch.is_digit_radix(16) => buf.push_char(ch),
                            ch => {
                                self.peek_buffer.push_char(ch);
                                break;
                            }
                        }
                    }
                    self.push_token(TNumericLiteral(from_str_radix(buf.as_slice(), 16).unwrap()));
                },
                '0' => {
                    let mut buf = "0".into_string();
                    let mut gone_decimal = false;
                    loop {
                        let ch = self.next();
                        match ch {
                            Ok(ch) if ch.is_digit_radix(8) =>
                                buf.push_char(ch),
                            Ok('8') | Ok('9') | Ok('.') => {
                                gone_decimal = true;
                                buf.push_char(ch.unwrap());
                            },
                            Ok(ch) => {
                                self.peek_buffer.push_char(ch);
                                break;
                            },
                            Err(IoError {kind: EndOfFile, ..}) =>
                                break,
                            Err(err) =>
                                return Err(err)
                        }
                    }
                    self.push_token(TNumericLiteral(if gone_decimal {
                        from_str(buf.as_slice())
                    } else {
                        from_str_radix(buf.as_slice(), 8)
                    }.unwrap()));
                },
                _ if ch.is_digit() => {
                    let mut buf = String::new();
                    buf.push_char(ch);
                    loop {
                        let ch = try!(self.next());
                        match ch {
                            '.' => buf.push_char(ch),
                            _ if ch.is_digit() => buf.push_char(ch),
                            _ => {
                                self.peek_buffer.push_char(ch);
                                break;
                            }
                        }
                    }
                    self.push_token(TNumericLiteral(from_str(buf.as_slice()).unwrap()));
                },
                _ if ch.is_alphabetic() || ch == '$' || ch == '_' => {
                    let mut buf = String::new();
                    buf.push_char(ch);
                    loop {
                        let ch = try!(self.next());
                        match ch {
                            _ if ch.is_alphabetic() || ch.is_digit() || ch == '_' => buf.push_char(ch),
                            _ => {
                                self.peek_buffer.push_char(ch);
                                break;
                            }
                        }
                    }
                    self.push_token(match buf.as_slice() {
                        "true" => TBooleanLiteral(true),
                        "false" => TBooleanLiteral(false),
                        "null" => TNullLiteral,
                        slice => match FromStr::from_str(slice) {
                            Some(keyword) => TKeyword(keyword),
                            None => TIdentifier(buf.clone())
                        }
                    });
                },
                ';' => self.push_punc(PSemicolon),
                ':' => self.push_punc(PColon),
                '.' => self.push_punc(PDot),
                '(' => self.push_punc(POpenParen),
                ')' => self.push_punc(PCloseParen),
                ',' => self.push_punc(PComma),
                '{' => self.push_punc(POpenBlock),
                '}' => self.push_punc(PCloseBlock),
                '[' => self.push_punc(POpenBracket),
                ']' => self.push_punc(PCloseBracket),
                '?' => self.push_punc(PQuestion),
                '/' if self.next_is('/') => {
                    let mut buf = String::new();
                    loop {
                        match try!(self.next()) {
                            '\n' => break,
                            ch => buf.push_char(ch)
                        }
                    }
                    self.push_token(TComment(buf));
                },
                '/' if self.next_is('*') => {
                    let mut buf = String::new();
                    loop {
                        match try!(self.next()) {
                            '\n' => break,
                            '*' if self.next_is('/') => break,
                            ch => buf.push_char(ch)
                        }
                    }
                    self.push_token(TComment(buf));
                },
                '/' if self.next_is('=') =>
                    self.push_punc(PAssignDiv),
                '/' => self.push_punc(PDiv),
                '*' if self.next_is('=') =>
                    self.push_punc(PAssignMul),
                '*' => self.push_punc(PMul),
                '+' if self.next_is('=') =>
                    self.push_punc(PAssignAdd),
                '+' if self.next_is('+') =>
                    self.push_punc(PInc),
                '+' => self.push_punc(PAdd),
                '-' if self.next_is('=') =>
                    self.push_punc(PAssignSub),
                '-' if self.next_is('-') =>
                    self.push_punc(PDec),
                '-' => self.push_punc(PSub),
                '%' if self.next_is('=') =>
                    self.push_punc(PAssignMod),
                '%' => self.push_punc(PMod),
                '|' if self.next_is('|') =>
                    self.push_punc(PBoolOr),
                '|' if self.next_is('=') =>
                    self.push_punc(PAssignOr),
                '|' => self.push_punc(POr),
                '&' if self.next_is('&') =>
                    self.push_punc(PBoolAnd),
                '&' if self.next_is('=') =>
                    self.push_punc(PAssignAnd),
                '&' => self.push_punc(PAnd),
                '^' if self.next_is('=') =>
                    self.push_punc(PAssignXor),
                '^' => self.push_punc(PXor),
                '=' if self.next_is('>') =>
                    self.push_punc(PArrow),
                '=' if self.next_is('=') => {
                    let punc = if self.next_is('=') {
                        PStrictEq
                    } else {
                        PEq
                    };
                    self.push_punc(punc)
                },
                '=' => self.push_punc(PAssign),
                '<' if self.next_is('=') =>
                    self.push_punc(PLessThanOrEq),
                '<' if self.next_is('<') => {
                    let punc = if self.next_is('=') {
                        PAssignLeftSh
                    } else {
                        PLeftSh
                    };
                    self.push_punc(punc)
                },
                '<' => self.push_punc(PLessThan),
                '>' if self.next_is('=') =>
                    self.push_punc(PGreaterThanOrEq),
                '>' if self.next_is('>') => {
                    let punc = if self.next_is('=') {
                        PAssignRightSh
                    } else if self.next_is('>') {
                        if self.next_is('=') {
                            PAssignURightSh
                        } else {
                            PURightSh
                        }
                    } else {
                        PRightSh
                    };
                    self.push_punc(punc)
                },
                '>' => self.push_punc(PGreaterThan),
                '!' if self.next_is('=') => {
                    let punc = if self.next_is('=') {
                        PStrictNotEq
                    } else {
                        PNotEq
                    };
                    self.push_punc(punc)
                },
                '!' => self.push_punc(PNot),
                '~' => self.push_punc(PNeg),
                '\n' | '\u2028'|'\u2029' => {
                    self.line_number += 1;
                    self.column_number = 0;
                },
                '\r' => {
                    self.column_number = 0;
                },
                ' ' => (),
                ch => fail!("{}:{}: Unexpected '{}'", self.line_number, self.column_number, ch)
            };
        };
        Ok(())
    }
}<commit_msg>Move check from guard into case<commit_after>use ast::punc::*;
use ast::token::*;
use std::io::{BufReader, BufferedReader, Buffer, IoError, IoResult, EndOfFile};
use std::char::from_u32;
use std::num::from_str_radix;
use std::from_str::FromStr;
/// A Javascript lexer
pub struct Lexer<B> {
    /// The list of tokens generated so far
    pub tokens : Vec<Token>,
    /// The current line number in the script
    line_number : uint,
    /// The current column number in the script
    column_number : uint,
    /// The reader
    buffer: B,
    /// The peeked character buffer
    peek_buffer: String
}
impl<B:Buffer> Lexer<B> {
    /// Creates a new lexer with empty buffers
    pub fn new(buffer: B) -> Lexer<B> {
        Lexer {
            tokens: Vec::new(),
            line_number: 1,
            column_number: 0,
            buffer: buffer,
            peek_buffer: String::new()
        }
    }
    #[inline(always)]
    fn push_token(&mut self, tk:TokenData) {
        self.tokens.push(Token::new(tk, self.line_number, self.column_number))
    }
    #[inline(always)]
    fn push_punc(&mut self, punc:Punctuator) {
        self.push_token(TPunctuator(punc));
    }
    /// Processes an input stream from a string into an array of tokens
    pub fn lex_str(script:&str) -> Vec<Token> {
        let script_bytes:&[u8] = script.as_bytes();
        let reader = BufReader::new(script_bytes);
        let buf_reader = BufferedReader::new(reader);
        let mut lexer = Lexer::new(buf_reader);
        lexer.lex().unwrap();
        lexer.tokens
    }
    fn next(&mut self) -> IoResult<char> {
        if self.peek_buffer.len() == 0 {
            self.buffer.read_char()
        } else {
            Ok(match self.peek_buffer.pop_char() {
                Some(v) => v,
                None => unreachable!()
            })
        }
    }
    fn next_is(&mut self, peek:char) -> bool {
        if self.peek_buffer.len() > 0 {
            let matched = self.peek_buffer.as_slice().char_at(self.peek_buffer.len() - 1) == peek;
            if matched {
                self.peek_buffer.pop_char();
            }
            matched
        } else {
            match self.buffer.read_char() {
                Ok(ch) if ch == peek => {
                    self.column_number += 1;
                    true
                },
                Ok(ch) if ch != peek => {
                    self.peek_buffer.push_char(ch);
                    false
                },
                _ => false
            }
        }
    }
    /// Processes an input stream from the `buffer` into a vector of tokens
    pub fn lex(&mut self) -> IoResult<()> {
        loop {
            let ch = match self.next() {
                Ok(ch) => ch,
                Err(IoError {kind: EndOfFile, ..}) => break,
                Err(err) => return Err(err)
            };
            self.column_number += 1;
            match ch {
                '"' | '\'' => {
                    let mut buf = String::new();
                    loop {
                        match try!(self.next()) {
                            '\'' if ch == '\'' => {
                                break;
                            },
                            '"' if ch == '"' => {
                                break;
                            },
                            '\\' => {
                                let escape = try!(self.next());
                                if escape != '\n' {
                                    let escaped_ch = match escape {
                                        'n' => '\n',
                                        'r' => '\r',
                                        't' => '\t',
                                        'b' => '\x08',
                                        'f' => '\x0c',
                                        '0' => '\0',
                                        'x' => {
                                            let mut nums = String::with_capacity(2);
                                            for _ in range(0, 2) {
                                                nums.push_char(try!(self.next()));
                                            }
                                            self.column_number += 2;
                                            let as_num = match from_str_radix(nums.as_slice(), 16) {
                                                Some(v) => v,
                                                None => 0
                                            };
                                            match from_u32(as_num) {
                                                Some(v) => v,
                                                None => fail!("{}:{}: {} is not a valid unicode scalar value", self.line_number, self.column_number, as_num)
                                            }
                                        },
                                        'u' => {
                                            let mut nums = String::new();
                                            for _ in range(0, 4) {
                                                nums.push_char(try!(self.next()));
                                            }
                                            self.column_number += 4;
                                            let as_num = match from_str_radix(nums.as_slice(), 16) {
                                                Some(v) => v,
                                                None => 0
                                            };
                                            match from_u32(as_num) {
                                                Some(v) => v,
                                                None => fail!("{}:{}: {} is not a valid unicode scalar value", self.line_number, self.column_number, as_num)
                                            }
                                        },
                                        '\'' | '"' => escape,
                                        _ => fail!("{}:{}: Invalid escape `{}`", self.line_number, self.column_number, ch)
                                    };
                                    buf.push_char(escaped_ch);
                                }
                            },
                            ch => buf.push_char(ch)
                        }
                    }
                    self.push_token(TStringLiteral(buf))
                },
                '0' if self.next_is('x') => {
                    let mut buf = String::new();
                    loop {
                        match try!(self.next()) {
                            ch if ch.is_digit_radix(16) => buf.push_char(ch),
                            ch => {
                                self.peek_buffer.push_char(ch);
                                break;
                            }
                        }
                    }
                    self.push_token(TNumericLiteral(from_str_radix(buf.as_slice(), 16).unwrap()));
                },
                '0' => {
                    let mut buf = "0".into_string();
                    let mut gone_decimal = false;
                    loop {
                        let ch = self.next();
                        match ch {
                            Ok(ch) if ch.is_digit_radix(8) =>
                                buf.push_char(ch),
                            Ok('8') | Ok('9') | Ok('.') => {
                                gone_decimal = true;
                                buf.push_char(ch.unwrap());
                            },
                            Ok(ch) => {
                                self.peek_buffer.push_char(ch);
                                break;
                            },
                            Err(IoError {kind: EndOfFile, ..}) =>
                                break,
                            Err(err) =>
                                return Err(err)
                        }
                    }
                    self.push_token(TNumericLiteral(if gone_decimal {
                        from_str(buf.as_slice())
                    } else {
                        from_str_radix(buf.as_slice(), 8)
                    }.unwrap()));
                },
                _ if ch.is_digit() => {
                    let mut buf = String::new();
                    buf.push_char(ch);
                    loop {
                        let ch = try!(self.next());
                        match ch {
                            '.' => buf.push_char(ch),
                            _ if ch.is_digit() => buf.push_char(ch),
                            _ => {
                                self.peek_buffer.push_char(ch);
                                break;
                            }
                        }
                    }
                    self.push_token(TNumericLiteral(from_str(buf.as_slice()).unwrap()));
                },
                _ if ch.is_alphabetic() || ch == '$' || ch == '_' => {
                    let mut buf = String::new();
                    buf.push_char(ch);
                    loop {
                        let ch = try!(self.next());
                        match ch {
                            _ if ch.is_alphabetic() || ch.is_digit() || ch == '_' => buf.push_char(ch),
                            _ => {
                                self.peek_buffer.push_char(ch);
                                break;
                            }
                        }
                    }
                    self.push_token(match buf.as_slice() {
                        "true" => TBooleanLiteral(true),
                        "false" => TBooleanLiteral(false),
                        "null" => TNullLiteral,
                        slice => match FromStr::from_str(slice) {
                            Some(keyword) => TKeyword(keyword),
                            None => TIdentifier(buf.clone())
                        }
                    });
                },
                ';' => self.push_punc(PSemicolon),
                ':' => self.push_punc(PColon),
                '.' => self.push_punc(PDot),
                '(' => self.push_punc(POpenParen),
                ')' => self.push_punc(PCloseParen),
                ',' => self.push_punc(PComma),
                '{' => self.push_punc(POpenBlock),
                '}' => self.push_punc(PCloseBlock),
                '[' => self.push_punc(POpenBracket),
                ']' => self.push_punc(PCloseBracket),
                '?' => self.push_punc(PQuestion),
                '/' if self.next_is('/') => {
                    let mut buf = String::new();
                    loop {
                        match try!(self.next()) {
                            '\n' => break,
                            ch => buf.push_char(ch)
                        }
                    }
                    self.push_token(TComment(buf));
                },
                '/' if self.next_is('*') => {
                    let mut buf = String::new();
                    loop {
                        match try!(self.next()) {
                            '\n' => break,
                            '*' if self.next_is('/') => break,
                            ch => buf.push_char(ch)
                        }
                    }
                    self.push_token(TComment(buf));
                },
                '/' if self.next_is('=') =>
                    self.push_punc(PAssignDiv),
                '/' => self.push_punc(PDiv),
                '*' if self.next_is('=') =>
                    self.push_punc(PAssignMul),
                '*' => self.push_punc(PMul),
                '+' if self.next_is('=') =>
                    self.push_punc(PAssignAdd),
                '+' if self.next_is('+') =>
                    self.push_punc(PInc),
                '+' => self.push_punc(PAdd),
                '-' if self.next_is('=') =>
                    self.push_punc(PAssignSub),
                '-' if self.next_is('-') =>
                    self.push_punc(PDec),
                '-' => self.push_punc(PSub),
                '%' if self.next_is('=') =>
                    self.push_punc(PAssignMod),
                '%' => self.push_punc(PMod),
                '|' if self.next_is('|') =>
                    self.push_punc(PBoolOr),
                '|' if self.next_is('=') =>
                    self.push_punc(PAssignOr),
                '|' => self.push_punc(POr),
                '&' if self.next_is('&') =>
                    self.push_punc(PBoolAnd),
                '&' if self.next_is('=') =>
                    self.push_punc(PAssignAnd),
                '&' => self.push_punc(PAnd),
                '^' if self.next_is('=') =>
                    self.push_punc(PAssignXor),
                '^' => self.push_punc(PXor),
                '=' if self.next_is('>') =>
                    self.push_punc(PArrow),
                '=' if self.next_is('=') => {
                    let punc = if self.next_is('=') {
                        PStrictEq
                    } else {
                        PEq
                    };
                    self.push_punc(punc)
                },
                '=' => self.push_punc(PAssign),
                '<' if self.next_is('=') =>
                    self.push_punc(PLessThanOrEq),
                '<' if self.next_is('<') => {
                    let punc = if self.next_is('=') {
                        PAssignLeftSh
                    } else {
                        PLeftSh
                    };
                    self.push_punc(punc)
                },
                '<' => self.push_punc(PLessThan),
                '>' if self.next_is('=') =>
                    self.push_punc(PGreaterThanOrEq),
                '>' if self.next_is('>') => {
                    let punc = if self.next_is('=') {
                        PAssignRightSh
                    } else if self.next_is('>') {
                        if self.next_is('=') {
                            PAssignURightSh
                        } else {
                            PURightSh
                        }
                    } else {
                        PRightSh
                    };
                    self.push_punc(punc)
                },
                '>' => self.push_punc(PGreaterThan),
                '!' if self.next_is('=') => {
                    let punc = if self.next_is('=') {
                        PStrictNotEq
                    } else {
                        PNotEq
                    };
                    self.push_punc(punc)
                },
                '!' => self.push_punc(PNot),
                '~' => self.push_punc(PNeg),
                '\n' | '\u2028'|'\u2029' => {
                    self.line_number += 1;
                    self.column_number = 0;
                },
                '\r' => {
                    self.column_number = 0;
                },
                ' ' => (),
                ch => fail!("{}:{}: Unexpected '{}'", self.line_number, self.column_number, ch)
            };
        };
        Ok(())
    }
}<|endoftext|>


<commit_before>use std::f64::consts::PI;

use terrain::dem;

/// This file assumes that all coordinates are provided relative to the earth represented as a
/// perfect sphere. This isn't quite accurate: The coordinate system of the input datasets are
/// actually WGS84 or NAD83. However, for our purposes the difference should not be noticable.

/// The radius of the earth in meters.
const EARTH_RADIUS: f64 = 6371000.0;
const EARTH_CIRCUMFERENCE: f64 = 2.0 * PI * EARTH_RADIUS;

pub struct TerrainFile {
    width: usize,
    height: usize,
    cell_size: f32,

    elevations: Vec<f32>,
    slopes: Vec<(f32, f32)>,
}

impl TerrainFile {
    fn compute_slopes(width: usize,
                      height: usize,
                      cell_size: f32,
                      elevations: &[f32])
                      -> Vec<(f32, f32)> {
        let get_elevation = |x: usize, y: usize| {
            assert!(x < width);
            assert!(y < height);
            elevations[x + y * width]
        };

        let mut slopes = Vec::with_capacity(width * height);
        for y in 0..height {
            for x in 0..width {
                let xslope = if x == 0 {
                    (get_elevation(x + 1, y) - get_elevation(x, y)) / cell_size
                } else if x == width - 1 {
                    (get_elevation(x, y) - get_elevation(x - 1, y)) / cell_size
                } else {
                    (get_elevation(x + 1, y) - get_elevation(x - 1, y)) / (2.0 * cell_size)
                };

                let yslope = if y == 0 {
                    (get_elevation(x, y + 1) - get_elevation(x, y)) / cell_size
                } else if y == height - 1 {
                    (get_elevation(x, y) - get_elevation(x, y - 1)) / cell_size
                } else {
                    (get_elevation(x, y + 1) - get_elevation(x, y - 1)) / (2.0 * cell_size)
                };

                slopes.push((xslope, yslope));
            }
        }

        slopes
    }

    /// Construct a `TerrainFile` from a `DigitalElevationModel`.
    pub fn from_digital_elevation_model(dem: dem::DigitalElevationModel) -> Self {
        // Compute approximate cell size in meters.
        let ycenter = dem.yllcorner + dem.cell_size * dem.height as f64;
        let cell_size_x = (dem.cell_size / 360.0) * EARTH_CIRCUMFERENCE *
                          ycenter.to_radians().cos();
        let cell_size_y = (dem.cell_size / 360.0) * EARTH_CIRCUMFERENCE;
        let cell_size = 0.5 * (cell_size_x + cell_size_y) as f32;

        let slopes = Self::compute_slopes(dem.width, dem.height, cell_size, &dem.elevations[..]);

        TerrainFile {
            width: dem.width,
            height: dem.height,
            cell_size,
            elevations: dem.elevations,
            slopes,
        }
    }

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

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

    pub fn elevation(&self, x: usize, y: usize) -> f32 {
        assert!(x < self.width);
        assert!(y < self.height);
        self.elevations[x + y * self.width]
    }

    pub fn slope(&self, x: usize, y: usize) -> (f32, f32) {
        assert!(x < self.width);
        assert!(y < self.height);
        self.slopes[x + y * self.width]
    }

    pub fn elevations(&self) -> &[f32] {
        &self.elevations[..]
    }

    pub fn slopes(&self) -> &[(f32, f32)] {
        &self.slopes[..]
    }

    pub fn cell_size(&self) -> f32 {
        self.cell_size
    }
}
<commit_msg>Minimize stretching of heightmaps in the horizontal direction<commit_after>use std::f64::consts::PI;

use terrain::dem;

/// This file assumes that all coordinates are provided relative to the earth represented as a
/// perfect sphere. This isn't quite accurate: The coordinate system of the input datasets are
/// actually WGS84 or NAD83. However, for our purposes the difference should not be noticable.

/// The radius of the earth in meters.
const EARTH_RADIUS: f64 = 6371000.0;
const EARTH_CIRCUMFERENCE: f64 = 2.0 * PI * EARTH_RADIUS;

pub struct TerrainFile {
    width: usize,
    height: usize,
    cell_size: f32,

    elevations: Vec<f32>,
    slopes: Vec<(f32, f32)>,
}

impl TerrainFile {
    fn compute_slopes(width: usize,
                      height: usize,
                      cell_size: f32,
                      elevations: &[f32])
                      -> Vec<(f32, f32)> {
        let get_elevation = |x: usize, y: usize| {
            assert!(x < width);
            assert!(y < height);
            elevations[x + y * width]
        };

        let mut slopes = Vec::with_capacity(width * height);
        for y in 0..height {
            for x in 0..width {
                let xslope = if x == 0 {
                    (get_elevation(x + 1, y) - get_elevation(x, y)) / cell_size
                } else if x == width - 1 {
                    (get_elevation(x, y) - get_elevation(x - 1, y)) / cell_size
                } else {
                    (get_elevation(x + 1, y) - get_elevation(x - 1, y)) / (2.0 * cell_size)
                };

                let yslope = if y == 0 {
                    (get_elevation(x, y + 1) - get_elevation(x, y)) / cell_size
                } else if y == height - 1 {
                    (get_elevation(x, y) - get_elevation(x, y - 1)) / cell_size
                } else {
                    (get_elevation(x, y + 1) - get_elevation(x, y - 1)) / (2.0 * cell_size)
                };

                slopes.push((xslope, yslope));
            }
        }

        slopes
    }

    /// Construct a `TerrainFile` from a `DigitalElevationModel`.
    pub fn from_digital_elevation_model(dem: dem::DigitalElevationModel) -> Self {
        // Compute approximate cell size in meters.
        let ycenter = dem.yllcorner + 0.5 * dem.cell_size * dem.height as f64;
        let cell_size_x = (dem.cell_size / 360.0) * EARTH_CIRCUMFERENCE *
                          ycenter.to_radians().cos();
        let cell_size_y = (dem.cell_size / 360.0) * EARTH_CIRCUMFERENCE;
        let cell_size = cell_size_y as f32;

        let cell_size_ratio = cell_size_y / cell_size_x;
        let width = (dem.width as f64 / cell_size_ratio).floor() as usize;
        let height = dem.height;
        let mut elevations = Vec::with_capacity(width * height);

        for y in 0..height {
            for x in 0..width {
                let fx = x as f64 * cell_size_ratio;
                let floor_fx = fx.floor() as usize;
                let ceil_fx = fx.ceil() as usize;
                assert!(floor_fx < dem.width);
                assert!(ceil_fx < dem.width);
                let t = fx.fract() as f32;
                let h = dem.elevations[floor_fx + y * dem.width] * (1.0 - t) + dem.elevations[ceil_fx + y * dem.width] * t;
                elevations.push(h)
            }
        }

        let slopes = Self::compute_slopes(width, height, cell_size, &elevations[..]);

        TerrainFile {
            width,
            height,
            cell_size,
            elevations,
            slopes,
        }
    }

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

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

    pub fn elevation(&self, x: usize, y: usize) -> f32 {
        assert!(x < self.width);
        assert!(y < self.height);
        self.elevations[x + y * self.width]
    }

    pub fn slope(&self, x: usize, y: usize) -> (f32, f32) {
        assert!(x < self.width);
        assert!(y < self.height);
        self.slopes[x + y * self.width]
    }

    pub fn elevations(&self) -> &[f32] {
        &self.elevations[..]
    }

    pub fn slopes(&self) -> &[(f32, f32)] {
        &self.slopes[..]
    }

    pub fn cell_size(&self) -> f32 {
        self.cell_size
    }
}
<|endoftext|>


<commit_before>use crate::agents::*;
use crate::config::LimitConfig;
use crate::crypto::SigningAlgorithm;
use crate::utils::{agent::*, unix_timestamp};
use ::rusqlite::{Connection, Error as SqlError, OptionalExtension, ToSql, NO_PARAMS};
use std::path::PathBuf;
use std::time::Duration;
use tokio::task::spawn_blocking;
use url::Url;

macro_rules! params {
    ($($list:expr),*) => (
        std::convert::identity::<&[&dyn ToSql]>(&[$($list),*])
    );
}

/// Database file `application_id` value. 'Prtr' in hex.
const APP_ID: u32 = 0x5072_7472;

/// Message sent at an interval to collect garbage.
struct Gc;
impl Message for Gc {
    type Reply = ();
}

/// Message used internally to save a cache entry.
struct SaveCache {
    url: Url,
    data: String,
    expires: i64,
}
impl Message for SaveCache {
    type Reply = Result<(), SqlError>;
}

/// Message used internally to save a key set.
struct SaveKeys(KeySet);
impl Message for SaveKeys {
    type Reply = Result<(), SqlError>;
}

/// Store implementation using `rusqlite`.
pub struct RusqliteStore {
    /// TTL of session keys
    expire_sessions: Duration,
    /// TTL of cache keys
    expire_cache: Duration,
    /// Rate limit configuration.
    limit_configs: Vec<LimitConfig>,
    /// SQLite connection.
    conn: Connection,
    /// The agent used for fetching on cache miss.
    fetcher: Addr<FetchAgent>,
    /// Key manager if rotating keys are enabled.
    key_manager: Option<Addr<RotatingKeys>>,
}

impl RusqliteStore {
    pub async fn new(
        sqlite_db: PathBuf,
        expire_sessions: Duration,
        expire_cache: Duration,
        limit_configs: Vec<LimitConfig>,
        fetcher: Addr<FetchAgent>,
    ) -> Result<Self, SqlError> {
        spawn_blocking(move || {
            let conn = Connection::open(&sqlite_db)?;
            conn.busy_timeout(Duration::from_millis(500))?;
            Self::verify_app_id(&conn)?;
            Self::verify_schema(&conn)?;
            log::warn!(
                "Storing sessions and keys in SQLite at: {}",
                sqlite_db.display()
            );
            log::warn!("Please always double check this directory has secure permissions!");
            log::warn!("(This warning can't be fixed; it's a friendly reminder.)");
            Ok(RusqliteStore {
                expire_sessions,
                expire_cache,
                limit_configs,
                conn,
                fetcher,
                key_manager: None,
            })
        })
        .await
        .unwrap()
    }

    fn verify_app_id(conn: &Connection) -> Result<(), SqlError> {
        // If this is 0, assume the file was just now created.
        let schema_version: u32 =
            conn.query_row("SELECT * FROM pragma_schema_version()", NO_PARAMS, |row| {
                row.get(0)
            })?;
        if schema_version == 0 {
            // Note: can't use parameter binding in pragma.
            conn.execute(&format!("PRAGMA application_id = {}", APP_ID), NO_PARAMS)?;
        } else {
            let app_id: u32 =
                conn.query_row("SELECT * FROM pragma_application_id()", NO_PARAMS, |row| {
                    row.get(0)
                })?;
            if app_id != APP_ID {
                panic!(
                    "The SQLite database has an invalid application ID: {}",
                    app_id
                );
            }
        }
        Ok(())
    }

    fn verify_schema(conn: &Connection) -> Result<(), SqlError> {
        let user_version: u32 =
            conn.query_row("SELECT * FROM pragma_user_version()", NO_PARAMS, |row| {
                row.get(0)
            })?;
        match user_version {
            0 => Self::init_schema(conn),
            1 => Ok(()),
            _ => panic!(
                "The SQLite database has an unknown version: {}",
                user_version
            ),
        }
    }

    fn init_schema(conn: &Connection) -> Result<(), SqlError> {
        conn.execute_batch(
            "
            BEGIN;

            CREATE TABLE sessions (
                id TEXT NOT NULL PRIMARY KEY,
                data TEXT NOT NULL,
                expires INTEGER NOT NULL
            );
            CREATE INDEX sessions_expires ON sessions (expires);

            CREATE TABLE cache_entries (
                url TEXT NOT NULL PRIMARY KEY,
                data TEXT NOT NULL,
                expires INTEGER NOT NULL
            );
            CREATE INDEX cache_entries_expires ON cache_entries (expires);

            CREATE TABLE rate_limits (
                id TEXT NOT NULL PRIMARY KEY,
                value INTEGER NOT NULL,
                expires INTEGER NOT NULL
            );
            CREATE INDEX rate_limits_expires ON rate_limits (expires);

            CREATE TABLE key_sets (
                signing_alg TEXT NOT NULL PRIMARY KEY,
                key_set TEXT NOT NULL
            );

            PRAGMA user_version = 1;
            COMMIT;
            ",
        )?;
        Ok(())
    }

    fn get_key_set(&mut self, signing_alg: SigningAlgorithm) -> KeySet {
        self.conn
            .query_row(
                "SELECT key_set FROM key_sets WHERE signing_alg = ?1 LIMIT 1",
                params![&signing_alg.as_str()],
                |row| row.get(0),
            )
            .optional()
            .expect("Could not fetch keys from SQLite")
            .map_or_else(
                || KeySet::empty(signing_alg),
                |data: String| serde_json::from_str(&data).expect("Invalid key set JSON in SQLite"),
            )
    }
}

impl Agent for RusqliteStore {
    fn started(&mut self, cx: Context<Self, AgentStarted>) {
        // Start the garbage collection loop.
        let addr = cx.addr().clone();
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_secs(60));
            loop {
                interval.tick().await;
                addr.send(Gc).await;
            }
        });
        cx.reply(());
    }
}

impl Handler<Gc> for RusqliteStore {
    fn handle(&mut self, _message: Gc, cx: Context<Self, Gc>) {
        let now = unix_timestamp() as i64;
        self.conn
            .execute("DELETE FROM sessions WHERE expires <= ?1", &[now])
            .expect("session cleanup failed");
        self.conn
            .execute("DELETE FROM cache_entries WHERE expires <= ?1", &[now])
            .expect("cache cleanup failed");
        self.conn
            .execute("DELETE FROM rate_limits WHERE expires <= ?1", &[now])
            .expect("rate limits cleanup failed");
        cx.reply(())
    }
}

impl Handler<SaveSession> for RusqliteStore {
    fn handle(&mut self, message: SaveSession, cx: Context<Self, SaveSession>) {
        cx.reply_with(move || {
            let expires = (unix_timestamp() + self.expire_sessions.as_secs()) as i64;
            let data = serde_json::to_string(&message.data)?;
            self.conn.execute(
                "REPLACE INTO sessions (id, data, expires) VALUES (?1, ?2, ?3)",
                params![&message.session_id, &data, &expires],
            )?;
            Ok(())
        });
    }
}

impl Handler<GetSession> for RusqliteStore {
    fn handle(&mut self, message: GetSession, cx: Context<Self, GetSession>) {
        cx.reply_with(move || {
            let now = unix_timestamp() as i64;
            let data: Option<String> = self
                .conn
                .query_row(
                    "SELECT data FROM sessions WHERE id = ?1 AND expires > ?2 LIMIT 1",
                    params![&message.session_id, &now],
                    |row| row.get(0),
                )
                .optional()?;
            if let Some(data) = data {
                let data = serde_json::from_str(&data)?;
                Ok(Some(data))
            } else {
                Ok(None)
            }
        });
    }
}

impl Handler<DeleteSession> for RusqliteStore {
    fn handle(&mut self, message: DeleteSession, cx: Context<Self, DeleteSession>) {
        cx.reply_with(move || {
            self.conn
                .execute("DELETE FROM sessions WHERE id = ?1", &[&message.session_id])?;
            Ok(())
        });
    }
}

impl Handler<FetchUrlCached> for RusqliteStore {
    fn handle(&mut self, message: FetchUrlCached, cx: Context<Self, FetchUrlCached>) {
        // TODO: Add locking to coordinate multiple fetches for the same resource.
        let now = unix_timestamp() as i64;
        let data: Result<Option<String>, SqlError> = self
            .conn
            .query_row(
                "SELECT data FROM cache_entries WHERE url = ?1 AND expires > ?2 LIMIT 1",
                params![&message.url.as_str(), &now],
                |row| row.get(0),
            )
            .optional();
        match data {
            Err(e) => return cx.reply(Err(e.into())),
            Ok(Some(data)) => return cx.reply(Ok(data)),
            Ok(None) => {}
        }
        let me = cx.addr().clone();
        let fetcher = self.fetcher.clone();
        let expire_cache = self.expire_cache;
        cx.reply_later(async move {
            let result = fetcher.send(FetchUrl::get(&message.url.clone())).await?;
            let ttl = std::cmp::max(expire_cache, result.max_age);
            me.send(SaveCache {
                url: message.url,
                data: result.data.clone(),
                expires: (unix_timestamp() + ttl.as_secs()) as i64,
            })
            .await?;
            Ok(result.data)
        });
    }
}

impl Handler<SaveCache> for RusqliteStore {
    fn handle(&mut self, message: SaveCache, cx: Context<Self, SaveCache>) {
        cx.reply_with(move || {
            self.conn.execute(
                "REPLACE INTO cache_entries (url, data, expires) VALUES (?1, ?2, ?3)",
                params![&message.url.as_str(), &message.data, &message.expires],
            )?;
            Ok(())
        });
    }
}

impl Handler<IncrAndTestLimits> for RusqliteStore {
    fn handle(&mut self, message: IncrAndTestLimits, cx: Context<Self, IncrAndTestLimits>) {
        cx.reply_with(move || {
            let mut ok = true;
            for config in &self.limit_configs {
                let id = message.input.build_key(&config, "", "|");
                let now = unix_timestamp() as i64;
                let window = config.window.as_secs() as i64;
                let tx = self.conn.transaction()?;
                tx.execute(
                    "DELETE FROM rate_limits WHERE id = ?1 AND expires <= ?2",
                    params![&id, &now],
                )?;
                if config.extend_window {
                    tx.execute(
                        "INSERT INTO rate_limits (id, value, expires) VALUES (?1, 1, ?2 + ?3)
                        ON CONFLICT(id) DO UPDATE SET value = value + 1, expires = ?2 + ?3",
                        params![&id, &now, &window],
                    )?;
                } else {
                    tx.execute(
                        "INSERT INTO rate_limits (id, value, expires) VALUES (?1, 1, ?2 + ?3)
                        ON CONFLICT(id) DO UPDATE SET value = value + 1",
                        params![&id, &now, &window],
                    )?;
                }
                let count: i64 = tx.query_row(
                    "SELECT value FROM rate_limits WHERE id = ?1 LIMIT 1",
                    params![&id],
                    |row| row.get(0),
                )?;
                tx.commit()?;
                ok = ok && count as usize <= config.max_count;
            }
            Ok(ok)
        })
    }
}

impl Handler<DecrLimits> for RusqliteStore {
    fn handle(&mut self, message: DecrLimits, cx: Context<Self, DecrLimits>) {
        cx.reply_with(move || {
            for config in &self.limit_configs {
                if !config.decr_complete {
                    continue;
                }
                let id = message.input.build_key(&config, "", "|");
                let expires = (unix_timestamp() + config.window.as_secs()) as i64;
                let tx = self.conn.transaction()?;
                tx.execute(
                    "DELETE FROM rate_limits WHERE id = ?1 AND (expires <= ?2 OR value <= 1)",
                    params![&id, &expires],
                )?;
                tx.execute(
                    "UPDATE rate_limits SET value = value - 1 WHERE id = ?1",
                    params![&id],
                )?;
                tx.commit()?;
            }
            Ok(())
        })
    }
}

impl Handler<EnableRotatingKeys> for RusqliteStore {
    fn handle(&mut self, message: EnableRotatingKeys, cx: Context<Self, EnableRotatingKeys>) {
        self.key_manager = Some(message.key_manager.clone());
        let mut update_msgs = Vec::with_capacity(message.signing_algs.len());
        for signing_alg in &message.signing_algs {
            let key_set = self.get_key_set(*signing_alg);
            update_msgs.push(UpdateKeys(key_set.clone()));
        }
        cx.reply_later(async move {
            for update_msg in update_msgs {
                message.key_manager.send(update_msg).await;
            }
        });
    }
}

impl Handler<RotateKeysLocked> for RusqliteStore {
    fn handle(&mut self, message: RotateKeysLocked, cx: Context<Self, RotateKeysLocked>) {
        let me = cx.addr().clone();
        let key_set = self.get_key_set(message.0);
        let key_manager = self.key_manager.as_ref().unwrap().clone();
        cx.reply_later(async move {
            if let Some(key_set) = key_manager.send(RotateKeys(key_set)).await {
                me.send(SaveKeys(key_set.clone()))
                    .await
                    .expect("Could not save keys to SQLite");
                key_manager.send(UpdateKeys(key_set)).await;
            }
        });
    }
}

impl Handler<ImportKeySet> for RusqliteStore {
    fn handle(&mut self, message: ImportKeySet, cx: Context<Self, ImportKeySet>) {
        let me = cx.addr().clone();
        cx.reply_later(async move {
            me.send(SaveKeys(message.0))
                .await
                .expect("Could not save keys to SQLite");
        });
    }
}

impl Handler<SaveKeys> for RusqliteStore {
    fn handle(&mut self, message: SaveKeys, cx: Context<Self, SaveKeys>) {
        let key_set = message.0;
        cx.reply_with(move || {
            let data = serde_json::to_string(&key_set).expect("Could not encode key set as JSON");
            self.conn.execute(
                "REPLACE INTO key_sets (signing_alg, key_set) VALUES (?1, ?2)",
                params![&key_set.signing_alg.as_str(), &data],
            )?;
            Ok(())
        });
    }
}

impl StoreSender for Addr<RusqliteStore> {}
<commit_msg>Fixes for rusqlite 0.25 upgrade<commit_after>use crate::agents::*;
use crate::config::LimitConfig;
use crate::crypto::SigningAlgorithm;
use crate::utils::{agent::*, unix_timestamp};
use ::rusqlite::{Connection, Error as SqlError, OptionalExtension, ToSql};
use std::path::PathBuf;
use std::time::Duration;
use tokio::task::spawn_blocking;
use url::Url;

macro_rules! params {
    ($($list:expr),*) => (
        std::convert::identity::<&[&dyn ToSql]>(&[$($list),*])
    );
}

/// Database file `application_id` value. 'Prtr' in hex.
const APP_ID: u32 = 0x5072_7472;

/// Message sent at an interval to collect garbage.
struct Gc;
impl Message for Gc {
    type Reply = ();
}

/// Message used internally to save a cache entry.
struct SaveCache {
    url: Url,
    data: String,
    expires: i64,
}
impl Message for SaveCache {
    type Reply = Result<(), SqlError>;
}

/// Message used internally to save a key set.
struct SaveKeys(KeySet);
impl Message for SaveKeys {
    type Reply = Result<(), SqlError>;
}

/// Store implementation using `rusqlite`.
pub struct RusqliteStore {
    /// TTL of session keys
    expire_sessions: Duration,
    /// TTL of cache keys
    expire_cache: Duration,
    /// Rate limit configuration.
    limit_configs: Vec<LimitConfig>,
    /// SQLite connection.
    conn: Connection,
    /// The agent used for fetching on cache miss.
    fetcher: Addr<FetchAgent>,
    /// Key manager if rotating keys are enabled.
    key_manager: Option<Addr<RotatingKeys>>,
}

impl RusqliteStore {
    pub async fn new(
        sqlite_db: PathBuf,
        expire_sessions: Duration,
        expire_cache: Duration,
        limit_configs: Vec<LimitConfig>,
        fetcher: Addr<FetchAgent>,
    ) -> Result<Self, SqlError> {
        spawn_blocking(move || {
            let conn = Connection::open(&sqlite_db)?;
            conn.busy_timeout(Duration::from_millis(500))?;
            Self::verify_app_id(&conn)?;
            Self::verify_schema(&conn)?;
            log::warn!(
                "Storing sessions and keys in SQLite at: {}",
                sqlite_db.display()
            );
            log::warn!("Please always double check this directory has secure permissions!");
            log::warn!("(This warning can't be fixed; it's a friendly reminder.)");
            Ok(RusqliteStore {
                expire_sessions,
                expire_cache,
                limit_configs,
                conn,
                fetcher,
                key_manager: None,
            })
        })
        .await
        .unwrap()
    }

    fn verify_app_id(conn: &Connection) -> Result<(), SqlError> {
        // If this is 0, assume the file was just now created.
        let schema_version: u32 =
            conn.query_row("SELECT * FROM pragma_schema_version()", [], |row| {
                row.get(0)
            })?;
        if schema_version == 0 {
            // Note: can't use parameter binding in pragma.
            conn.execute(&format!("PRAGMA application_id = {}", APP_ID), [])?;
        } else {
            let app_id: u32 =
                conn.query_row("SELECT * FROM pragma_application_id()", [], |row| {
                    row.get(0)
                })?;
            if app_id != APP_ID {
                panic!(
                    "The SQLite database has an invalid application ID: {}",
                    app_id
                );
            }
        }
        Ok(())
    }

    fn verify_schema(conn: &Connection) -> Result<(), SqlError> {
        let user_version: u32 =
            conn.query_row("SELECT * FROM pragma_user_version()", [], |row| row.get(0))?;
        match user_version {
            0 => Self::init_schema(conn),
            1 => Ok(()),
            _ => panic!(
                "The SQLite database has an unknown version: {}",
                user_version
            ),
        }
    }

    fn init_schema(conn: &Connection) -> Result<(), SqlError> {
        conn.execute_batch(
            "
            BEGIN;

            CREATE TABLE sessions (
                id TEXT NOT NULL PRIMARY KEY,
                data TEXT NOT NULL,
                expires INTEGER NOT NULL
            );
            CREATE INDEX sessions_expires ON sessions (expires);

            CREATE TABLE cache_entries (
                url TEXT NOT NULL PRIMARY KEY,
                data TEXT NOT NULL,
                expires INTEGER NOT NULL
            );
            CREATE INDEX cache_entries_expires ON cache_entries (expires);

            CREATE TABLE rate_limits (
                id TEXT NOT NULL PRIMARY KEY,
                value INTEGER NOT NULL,
                expires INTEGER NOT NULL
            );
            CREATE INDEX rate_limits_expires ON rate_limits (expires);

            CREATE TABLE key_sets (
                signing_alg TEXT NOT NULL PRIMARY KEY,
                key_set TEXT NOT NULL
            );

            PRAGMA user_version = 1;
            COMMIT;
            ",
        )?;
        Ok(())
    }

    fn get_key_set(&mut self, signing_alg: SigningAlgorithm) -> KeySet {
        self.conn
            .query_row(
                "SELECT key_set FROM key_sets WHERE signing_alg = ?1 LIMIT 1",
                params![&signing_alg.as_str()],
                |row| row.get(0),
            )
            .optional()
            .expect("Could not fetch keys from SQLite")
            .map_or_else(
                || KeySet::empty(signing_alg),
                |data: String| serde_json::from_str(&data).expect("Invalid key set JSON in SQLite"),
            )
    }
}

impl Agent for RusqliteStore {
    fn started(&mut self, cx: Context<Self, AgentStarted>) {
        // Start the garbage collection loop.
        let addr = cx.addr().clone();
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_secs(60));
            loop {
                interval.tick().await;
                addr.send(Gc).await;
            }
        });
        cx.reply(());
    }
}

impl Handler<Gc> for RusqliteStore {
    fn handle(&mut self, _message: Gc, cx: Context<Self, Gc>) {
        let now = unix_timestamp() as i64;
        self.conn
            .execute("DELETE FROM sessions WHERE expires <= ?1", [now])
            .expect("session cleanup failed");
        self.conn
            .execute("DELETE FROM cache_entries WHERE expires <= ?1", [now])
            .expect("cache cleanup failed");
        self.conn
            .execute("DELETE FROM rate_limits WHERE expires <= ?1", [now])
            .expect("rate limits cleanup failed");
        cx.reply(())
    }
}

impl Handler<SaveSession> for RusqliteStore {
    fn handle(&mut self, message: SaveSession, cx: Context<Self, SaveSession>) {
        cx.reply_with(move || {
            let expires = (unix_timestamp() + self.expire_sessions.as_secs()) as i64;
            let data = serde_json::to_string(&message.data)?;
            self.conn.execute(
                "REPLACE INTO sessions (id, data, expires) VALUES (?1, ?2, ?3)",
                params![&message.session_id, &data, &expires],
            )?;
            Ok(())
        });
    }
}

impl Handler<GetSession> for RusqliteStore {
    fn handle(&mut self, message: GetSession, cx: Context<Self, GetSession>) {
        cx.reply_with(move || {
            let now = unix_timestamp() as i64;
            let data: Option<String> = self
                .conn
                .query_row(
                    "SELECT data FROM sessions WHERE id = ?1 AND expires > ?2 LIMIT 1",
                    params![&message.session_id, &now],
                    |row| row.get(0),
                )
                .optional()?;
            if let Some(data) = data {
                let data = serde_json::from_str(&data)?;
                Ok(Some(data))
            } else {
                Ok(None)
            }
        });
    }
}

impl Handler<DeleteSession> for RusqliteStore {
    fn handle(&mut self, message: DeleteSession, cx: Context<Self, DeleteSession>) {
        cx.reply_with(move || {
            self.conn
                .execute("DELETE FROM sessions WHERE id = ?1", &[&message.session_id])?;
            Ok(())
        });
    }
}

impl Handler<FetchUrlCached> for RusqliteStore {
    fn handle(&mut self, message: FetchUrlCached, cx: Context<Self, FetchUrlCached>) {
        // TODO: Add locking to coordinate multiple fetches for the same resource.
        let now = unix_timestamp() as i64;
        let data: Result<Option<String>, SqlError> = self
            .conn
            .query_row(
                "SELECT data FROM cache_entries WHERE url = ?1 AND expires > ?2 LIMIT 1",
                params![&message.url.as_str(), &now],
                |row| row.get(0),
            )
            .optional();
        match data {
            Err(e) => return cx.reply(Err(e.into())),
            Ok(Some(data)) => return cx.reply(Ok(data)),
            Ok(None) => {}
        }
        let me = cx.addr().clone();
        let fetcher = self.fetcher.clone();
        let expire_cache = self.expire_cache;
        cx.reply_later(async move {
            let result = fetcher.send(FetchUrl::get(&message.url.clone())).await?;
            let ttl = std::cmp::max(expire_cache, result.max_age);
            me.send(SaveCache {
                url: message.url,
                data: result.data.clone(),
                expires: (unix_timestamp() + ttl.as_secs()) as i64,
            })
            .await?;
            Ok(result.data)
        });
    }
}

impl Handler<SaveCache> for RusqliteStore {
    fn handle(&mut self, message: SaveCache, cx: Context<Self, SaveCache>) {
        cx.reply_with(move || {
            self.conn.execute(
                "REPLACE INTO cache_entries (url, data, expires) VALUES (?1, ?2, ?3)",
                params![&message.url.as_str(), &message.data, &message.expires],
            )?;
            Ok(())
        });
    }
}

impl Handler<IncrAndTestLimits> for RusqliteStore {
    fn handle(&mut self, message: IncrAndTestLimits, cx: Context<Self, IncrAndTestLimits>) {
        cx.reply_with(move || {
            let mut ok = true;
            for config in &self.limit_configs {
                let id = message.input.build_key(&config, "", "|");
                let now = unix_timestamp() as i64;
                let window = config.window.as_secs() as i64;
                let tx = self.conn.transaction()?;
                tx.execute(
                    "DELETE FROM rate_limits WHERE id = ?1 AND expires <= ?2",
                    params![&id, &now],
                )?;
                if config.extend_window {
                    tx.execute(
                        "INSERT INTO rate_limits (id, value, expires) VALUES (?1, 1, ?2 + ?3)
                        ON CONFLICT(id) DO UPDATE SET value = value + 1, expires = ?2 + ?3",
                        params![&id, &now, &window],
                    )?;
                } else {
                    tx.execute(
                        "INSERT INTO rate_limits (id, value, expires) VALUES (?1, 1, ?2 + ?3)
                        ON CONFLICT(id) DO UPDATE SET value = value + 1",
                        params![&id, &now, &window],
                    )?;
                }
                let count: i64 = tx.query_row(
                    "SELECT value FROM rate_limits WHERE id = ?1 LIMIT 1",
                    params![&id],
                    |row| row.get(0),
                )?;
                tx.commit()?;
                ok = ok && count as usize <= config.max_count;
            }
            Ok(ok)
        })
    }
}

impl Handler<DecrLimits> for RusqliteStore {
    fn handle(&mut self, message: DecrLimits, cx: Context<Self, DecrLimits>) {
        cx.reply_with(move || {
            for config in &self.limit_configs {
                if !config.decr_complete {
                    continue;
                }
                let id = message.input.build_key(&config, "", "|");
                let expires = (unix_timestamp() + config.window.as_secs()) as i64;
                let tx = self.conn.transaction()?;
                tx.execute(
                    "DELETE FROM rate_limits WHERE id = ?1 AND (expires <= ?2 OR value <= 1)",
                    params![&id, &expires],
                )?;
                tx.execute(
                    "UPDATE rate_limits SET value = value - 1 WHERE id = ?1",
                    params![&id],
                )?;
                tx.commit()?;
            }
            Ok(())
        })
    }
}

impl Handler<EnableRotatingKeys> for RusqliteStore {
    fn handle(&mut self, message: EnableRotatingKeys, cx: Context<Self, EnableRotatingKeys>) {
        self.key_manager = Some(message.key_manager.clone());
        let mut update_msgs = Vec::with_capacity(message.signing_algs.len());
        for signing_alg in &message.signing_algs {
            let key_set = self.get_key_set(*signing_alg);
            update_msgs.push(UpdateKeys(key_set.clone()));
        }
        cx.reply_later(async move {
            for update_msg in update_msgs {
                message.key_manager.send(update_msg).await;
            }
        });
    }
}

impl Handler<RotateKeysLocked> for RusqliteStore {
    fn handle(&mut self, message: RotateKeysLocked, cx: Context<Self, RotateKeysLocked>) {
        let me = cx.addr().clone();
        let key_set = self.get_key_set(message.0);
        let key_manager = self.key_manager.as_ref().unwrap().clone();
        cx.reply_later(async move {
            if let Some(key_set) = key_manager.send(RotateKeys(key_set)).await {
                me.send(SaveKeys(key_set.clone()))
                    .await
                    .expect("Could not save keys to SQLite");
                key_manager.send(UpdateKeys(key_set)).await;
            }
        });
    }
}

impl Handler<ImportKeySet> for RusqliteStore {
    fn handle(&mut self, message: ImportKeySet, cx: Context<Self, ImportKeySet>) {
        let me = cx.addr().clone();
        cx.reply_later(async move {
            me.send(SaveKeys(message.0))
                .await
                .expect("Could not save keys to SQLite");
        });
    }
}

impl Handler<SaveKeys> for RusqliteStore {
    fn handle(&mut self, message: SaveKeys, cx: Context<Self, SaveKeys>) {
        let key_set = message.0;
        cx.reply_with(move || {
            let data = serde_json::to_string(&key_set).expect("Could not encode key set as JSON");
            self.conn.execute(
                "REPLACE INTO key_sets (signing_alg, key_set) VALUES (?1, ?2)",
                params![&key_set.signing_alg.as_str(), &data],
            )?;
            Ok(())
        });
    }
}

impl StoreSender for Addr<RusqliteStore> {}
<|endoftext|>


<commit_before>//! RustCrypto implementations of signing and verifying algorithms.
//! According to that organization, only hmac is safely implemented at the
//! moment.

use base64;
use crypto_mac::Mac;
use digest::generic_array::ArrayLength;
use digest::{BlockInput, FixedOutput, Update, Reset};
use hmac::Hmac;
use sha2;

use crate::algorithm::{AlgorithmType, SigningAlgorithm, VerifyingAlgorithm};
use crate::error::Error;
use crate::SEPARATOR;
/// A trait used to make the implementation of `SigningAlgorithm` and
/// `VerifyingAlgorithm` easier.
/// RustCrypto crates tend to have algorithm types defined at the type level,
/// so they cannot accept a self argument.
pub trait TypeLevelAlgorithmType {
    fn algorithm_type() -> AlgorithmType;
}

macro_rules! type_level_algorithm_type {
    ($rust_crypto_type: ty, $algorithm_type: expr) => {
        impl TypeLevelAlgorithmType for $rust_crypto_type {
            fn algorithm_type() -> AlgorithmType {
                $algorithm_type
            }
        }
    };
}

type_level_algorithm_type!(sha2::Sha256, AlgorithmType::Hs256);
type_level_algorithm_type!(sha2::Sha384, AlgorithmType::Hs384);
type_level_algorithm_type!(sha2::Sha512, AlgorithmType::Hs512);

impl<D> SigningAlgorithm for Hmac<D>
where
    D: Update + BlockInput + FixedOutput + Reset + Default + Clone + TypeLevelAlgorithmType,
    D::BlockSize: ArrayLength<u8>,
    D::OutputSize: ArrayLength<u8>,
{
    fn algorithm_type(&self) -> AlgorithmType {
        D::algorithm_type()
    }

    fn sign(&self, header: &str, claims: &str) -> Result<String, Error> {
        let hmac = get_hmac_with_data(&self, header, claims);
        let mac_result = hmac.finalize();
        let code = mac_result.into_bytes();
        Ok(base64::encode_config(&code, base64::URL_SAFE_NO_PAD))
    }
}

impl<D> VerifyingAlgorithm for Hmac<D>
where
    D: Update + BlockInput + FixedOutput + Reset + Default + Clone + TypeLevelAlgorithmType,
    D::BlockSize: ArrayLength<u8>,
    D::OutputSize: ArrayLength<u8>,
{
    fn algorithm_type(&self) -> AlgorithmType {
        D::algorithm_type()
    }

    fn verify_bytes(&self, header: &str, claims: &str, signature: &[u8]) -> Result<bool, Error> {
        let hmac = get_hmac_with_data(self, header, claims);
        hmac.verify(&signature)?;
        Ok(true)
    }
}

fn get_hmac_with_data<D>(hmac: &Hmac<D>, header: &str, claims: &str) -> Hmac<D>
where
    D: Update + BlockInput + FixedOutput + Reset + Default + Clone + TypeLevelAlgorithmType,
    D::BlockSize: ArrayLength<u8>,
    D::OutputSize: ArrayLength<u8>,
{
    let mut hmac = hmac.clone();
    hmac.reset();
    hmac.update(header.as_bytes());
    hmac.update(SEPARATOR.as_bytes());
    hmac.update(claims.as_bytes());
    hmac
}

#[cfg(test)]
mod tests {
    use crate::algorithm::{SigningAlgorithm, VerifyingAlgorithm};
    use crate::error::Error;
    use crypto_mac::NewMac;
    use hmac::Hmac;
    use sha2::Sha256;

    #[test]
    pub fn sign() -> Result<(), Error> {
        let header = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9";
        let claims = "eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWV9";
        let expected_signature = "TJVA95OrM7E2cBab30RMHrHDcEfxjoYZgeFONFh7HgQ";

        let signer: Hmac<Sha256> = Hmac::new_varkey(b"secret")?;
        let computed_signature = SigningAlgorithm::sign(&signer, &header, &claims)?;

        assert_eq!(computed_signature, expected_signature);
        Ok(())
    }

    #[test]
    pub fn verify() -> Result<(), Error> {
        let header = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9";
        let claims = "eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWV9";
        let signature = "TJVA95OrM7E2cBab30RMHrHDcEfxjoYZgeFONFh7HgQ";

        let verifier: Hmac<Sha256> = Hmac::new_varkey(b"secret")?;
        assert!(VerifyingAlgorithm::verify(
            &verifier, &header, &claims, &signature
        )?);
        Ok(())
    }
}
<commit_msg>Format import<commit_after>//! RustCrypto implementations of signing and verifying algorithms.
//! According to that organization, only hmac is safely implemented at the
//! moment.

use base64;
use crypto_mac::Mac;
use digest::generic_array::ArrayLength;
use digest::{BlockInput, FixedOutput, Reset, Update};
use hmac::Hmac;
use sha2;

use crate::algorithm::{AlgorithmType, SigningAlgorithm, VerifyingAlgorithm};
use crate::error::Error;
use crate::SEPARATOR;
/// A trait used to make the implementation of `SigningAlgorithm` and
/// `VerifyingAlgorithm` easier.
/// RustCrypto crates tend to have algorithm types defined at the type level,
/// so they cannot accept a self argument.
pub trait TypeLevelAlgorithmType {
    fn algorithm_type() -> AlgorithmType;
}

macro_rules! type_level_algorithm_type {
    ($rust_crypto_type: ty, $algorithm_type: expr) => {
        impl TypeLevelAlgorithmType for $rust_crypto_type {
            fn algorithm_type() -> AlgorithmType {
                $algorithm_type
            }
        }
    };
}

type_level_algorithm_type!(sha2::Sha256, AlgorithmType::Hs256);
type_level_algorithm_type!(sha2::Sha384, AlgorithmType::Hs384);
type_level_algorithm_type!(sha2::Sha512, AlgorithmType::Hs512);

impl<D> SigningAlgorithm for Hmac<D>
where
    D: Update + BlockInput + FixedOutput + Reset + Default + Clone + TypeLevelAlgorithmType,
    D::BlockSize: ArrayLength<u8>,
    D::OutputSize: ArrayLength<u8>,
{
    fn algorithm_type(&self) -> AlgorithmType {
        D::algorithm_type()
    }

    fn sign(&self, header: &str, claims: &str) -> Result<String, Error> {
        let hmac = get_hmac_with_data(&self, header, claims);
        let mac_result = hmac.finalize();
        let code = mac_result.into_bytes();
        Ok(base64::encode_config(&code, base64::URL_SAFE_NO_PAD))
    }
}

impl<D> VerifyingAlgorithm for Hmac<D>
where
    D: Update + BlockInput + FixedOutput + Reset + Default + Clone + TypeLevelAlgorithmType,
    D::BlockSize: ArrayLength<u8>,
    D::OutputSize: ArrayLength<u8>,
{
    fn algorithm_type(&self) -> AlgorithmType {
        D::algorithm_type()
    }

    fn verify_bytes(&self, header: &str, claims: &str, signature: &[u8]) -> Result<bool, Error> {
        let hmac = get_hmac_with_data(self, header, claims);
        hmac.verify(&signature)?;
        Ok(true)
    }
}

fn get_hmac_with_data<D>(hmac: &Hmac<D>, header: &str, claims: &str) -> Hmac<D>
where
    D: Update + BlockInput + FixedOutput + Reset + Default + Clone + TypeLevelAlgorithmType,
    D::BlockSize: ArrayLength<u8>,
    D::OutputSize: ArrayLength<u8>,
{
    let mut hmac = hmac.clone();
    hmac.reset();
    hmac.update(header.as_bytes());
    hmac.update(SEPARATOR.as_bytes());
    hmac.update(claims.as_bytes());
    hmac
}

#[cfg(test)]
mod tests {
    use crate::algorithm::{SigningAlgorithm, VerifyingAlgorithm};
    use crate::error::Error;
    use crypto_mac::NewMac;
    use hmac::Hmac;
    use sha2::Sha256;

    #[test]
    pub fn sign() -> Result<(), Error> {
        let header = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9";
        let claims = "eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWV9";
        let expected_signature = "TJVA95OrM7E2cBab30RMHrHDcEfxjoYZgeFONFh7HgQ";

        let signer: Hmac<Sha256> = Hmac::new_varkey(b"secret")?;
        let computed_signature = SigningAlgorithm::sign(&signer, &header, &claims)?;

        assert_eq!(computed_signature, expected_signature);
        Ok(())
    }

    #[test]
    pub fn verify() -> Result<(), Error> {
        let header = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9";
        let claims = "eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWV9";
        let signature = "TJVA95OrM7E2cBab30RMHrHDcEfxjoYZgeFONFh7HgQ";

        let verifier: Hmac<Sha256> = Hmac::new_varkey(b"secret")?;
        assert!(VerifyingAlgorithm::verify(
            &verifier, &header, &claims, &signature
        )?);
        Ok(())
    }
}
<|endoftext|>