<commit_before>use crate::postings::compression::AlignedBuffer;

/// This modules define the logic used to search for a doc in a given
/// block. (at most 128 docs)
///
/// Searching within a block is a hotspot when running intersection.
/// so it was worth defining it in its own module.

#[cfg(target_arch = "x86_64")]
mod sse2 {
    use crate::postings::compression::{AlignedBuffer, COMPRESSION_BLOCK_SIZE};
    use std::arch::x86_64::__m128i as DataType;
    use std::arch::x86_64::_mm_add_epi32 as op_add;
    use std::arch::x86_64::_mm_cmplt_epi32 as op_lt;
    use std::arch::x86_64::_mm_load_si128 as op_load; // requires 128-bits alignment
    use std::arch::x86_64::_mm_set1_epi32 as set1;
    use std::arch::x86_64::_mm_setzero_si128 as set0;
    use std::arch::x86_64::_mm_sub_epi32 as op_sub;
    use std::arch::x86_64::{_mm_cvtsi128_si32, _mm_shuffle_epi32};

    const MASK1: i32 = 78;
    const MASK2: i32 = 177;

    /// Performs an exhaustive linear search over the
    ///
    /// There is no early exit here. We simply count the
    /// number of elements that are `< target`.
    pub(crate) fn linear_search_sse2_128(arr: &AlignedBuffer, target: u32) -> usize {
        unsafe {
            let ptr = arr as *const AlignedBuffer as *const DataType;
            let vkey = set1(target as i32);
            let mut cnt = set0();
            // We work over 4 `__m128i` at a time.
            // A single `__m128i` actual contains 4 `u32`.
            for i in 0..(COMPRESSION_BLOCK_SIZE as isize) / (4 * 4) {
                let cmp1 = op_lt(op_load(ptr.offset(i * 4)), vkey);
                let cmp2 = op_lt(op_load(ptr.offset(i * 4 + 1)), vkey);
                let cmp3 = op_lt(op_load(ptr.offset(i * 4 + 2)), vkey);
                let cmp4 = op_lt(op_load(ptr.offset(i * 4 + 3)), vkey);
                let sum = op_add(op_add(cmp1, cmp2), op_add(cmp3, cmp4));
                cnt = op_sub(cnt, sum);
            }
            cnt = op_add(cnt, _mm_shuffle_epi32(cnt, MASK1));
            cnt = op_add(cnt, _mm_shuffle_epi32(cnt, MASK2));
            _mm_cvtsi128_si32(cnt) as usize
        }
    }

    #[cfg(test)]
    mod test {
        use super::linear_search_sse2_128;
        use crate::postings::compression::{AlignedBuffer, COMPRESSION_BLOCK_SIZE};

        #[test]
        fn test_linear_search_sse2_128_u32() {
            let mut block = [0u32; COMPRESSION_BLOCK_SIZE];
            for el in 0u32..128u32 {
                block[el as usize] = el * 2 + 1 << 18;
            }
            let target = block[64] + 1;
            assert_eq!(linear_search_sse2_128(&AlignedBuffer(block), target), 65);
        }
    }
}

/// This `linear search` browser exhaustively through the array.
/// but the early exit is very difficult to predict.
///
/// Coupled with `exponential search` this function is likely
/// to be called with the same `len`
fn linear_search(arr: &[u32], target: u32) -> usize {
    arr.iter().map(|&el| if el < target { 1 } else { 0 }).sum()
}

fn exponential_search(arr: &[u32], target: u32) -> (usize, usize) {
    let end = arr.len();
    let mut begin = 0;
    for &pivot in &[1, 3, 7, 15, 31, 63] {
        if pivot >= end {
            break;
        }
        if arr[pivot] > target {
            return (begin, pivot);
        }
        begin = pivot;
    }
    (begin, end)
}

fn galloping(block_docs: &[u32], target: u32) -> usize {
    let (start, end) = exponential_search(&block_docs, target);
    start + linear_search(&block_docs[start..end], target)
}

/// Tantivy may rely on SIMD instructions to search for a specific document within
/// a given block.
#[derive(Clone, Copy, PartialEq)]
pub enum BlockSearcher {
    #[cfg(target_arch = "x86_64")]
    SSE2,
    Scalar,
}

impl BlockSearcher {
    /// Search the first index containing an element greater or equal to
    /// the target.
    ///
    /// The results should be equivalent to
    /// ```ignore
    /// block[..]
    //       .iter()
    //       .take_while(|&&val| val < target)
    //       .count()
    /// ```
    ///
    /// The `start` argument is just used to hint that the response is
    /// greater than beyond `start`. The implementation may or may not use
    /// it for optimization.
    ///
    /// # Assumption
    ///
    /// The array len is > start.
    /// The block is sorted
    /// The target is assumed greater or equal to the `arr[start]`.
    /// The target is assumed smaller or equal to the last element of the block.
    ///
    /// Currently the scalar implementation starts by an exponential search, and
    /// then operates a linear search in the result subarray.
    ///
    /// If SSE2 instructions are available in the `(platform, running CPU)`,
    /// then we use a different implementation that does an exhaustive linear search over
    /// the full block whenever the block is full (`len == 128`). It is surprisingly faster, most likely because of the lack
    /// of branch.
    pub(crate) fn search_in_block(
        self,
        block_docs: &AlignedBuffer,
        len: usize,
        start: usize,
        target: u32,
    ) -> usize {
        #[cfg(target_arch = "x86_64")]
        {
            use crate::postings::compression::COMPRESSION_BLOCK_SIZE;
            if self == BlockSearcher::SSE2 && len == COMPRESSION_BLOCK_SIZE {
                return sse2::linear_search_sse2_128(block_docs, target);
            }
        }
        start + galloping(&block_docs.0[start..len], target)
    }
}

impl Default for BlockSearcher {
    fn default() -> BlockSearcher {
        #[cfg(target_arch = "x86_64")]
        {
            if is_x86_feature_detected!("sse2") {
                return BlockSearcher::SSE2;
            }
        }
        BlockSearcher::Scalar
    }
}

#[cfg(test)]
mod tests {
    use super::exponential_search;
    use super::linear_search;
    use super::BlockSearcher;
    use crate::postings::compression::{AlignedBuffer, COMPRESSION_BLOCK_SIZE};

    #[test]
    fn test_linear_search() {
        let len: usize = 50;
        let arr: Vec<u32> = (0..len).map(|el| 1u32 + (el as u32) * 2).collect();
        for target in 1..*arr.last().unwrap() {
            let res = linear_search(&arr[..], target);
            if res > 0 {
                assert!(arr[res - 1] < target);
            }
            if res < len {
                assert!(arr[res] >= target);
            }
        }
    }

    #[test]
    fn test_exponentiel_search() {
        assert_eq!(exponential_search(&[1, 2], 0), (0, 1));
        assert_eq!(exponential_search(&[1, 2], 1), (0, 1));
        assert_eq!(
            exponential_search(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], 7),
            (3, 7)
        );
    }

    fn util_test_search_in_block(block_searcher: BlockSearcher, block: &[u32], target: u32) {
        let cursor = search_in_block_trivial_but_slow(block, target);
        assert!(block.len() < COMPRESSION_BLOCK_SIZE);
        let mut output_buffer = [u32::max_value(); COMPRESSION_BLOCK_SIZE];
        output_buffer[..block.len()].copy_from_slice(block);
        for i in 0..cursor {
            assert_eq!(
                block_searcher.search_in_block(
                    &AlignedBuffer(output_buffer),
                    block.len(),
                    i,
                    target
                ),
                cursor
            );
        }
    }

    fn util_test_search_in_block_all(block_searcher: BlockSearcher, block: &[u32]) {
        use std::collections::HashSet;
        let mut targets = HashSet::new();
        for (i, val) in block.iter().cloned().enumerate() {
            if i > 0 {
                targets.insert(val - 1);
            }
            targets.insert(val);
        }
        for target in targets {
            util_test_search_in_block(block_searcher, block, target);
        }
    }

    fn search_in_block_trivial_but_slow(block: &[u32], target: u32) -> usize {
        block.iter().take_while(|&&val| val < target).count()
    }

    fn test_search_in_block_util(block_searcher: BlockSearcher) {
        for len in 1u32..128u32 {
            let v: Vec<u32> = (0..len).map(|i| i * 2).collect();
            util_test_search_in_block_all(block_searcher, &v[..]);
        }
    }

    #[test]
    fn test_search_in_block_scalar() {
        test_search_in_block_util(BlockSearcher::Scalar);
    }

    #[cfg(target_arch = "x86_64")]
    #[test]
    fn test_search_in_block_sse2() {
        test_search_in_block_util(BlockSearcher::SSE2);
    }
}
<commit_msg>ignore -> compile_fail<commit_after>use crate::postings::compression::AlignedBuffer;

/// This modules define the logic used to search for a doc in a given
/// block. (at most 128 docs)
///
/// Searching within a block is a hotspot when running intersection.
/// so it was worth defining it in its own module.

#[cfg(target_arch = "x86_64")]
mod sse2 {
    use crate::postings::compression::{AlignedBuffer, COMPRESSION_BLOCK_SIZE};
    use std::arch::x86_64::__m128i as DataType;
    use std::arch::x86_64::_mm_add_epi32 as op_add;
    use std::arch::x86_64::_mm_cmplt_epi32 as op_lt;
    use std::arch::x86_64::_mm_load_si128 as op_load; // requires 128-bits alignment
    use std::arch::x86_64::_mm_set1_epi32 as set1;
    use std::arch::x86_64::_mm_setzero_si128 as set0;
    use std::arch::x86_64::_mm_sub_epi32 as op_sub;
    use std::arch::x86_64::{_mm_cvtsi128_si32, _mm_shuffle_epi32};

    const MASK1: i32 = 78;
    const MASK2: i32 = 177;

    /// Performs an exhaustive linear search over the
    ///
    /// There is no early exit here. We simply count the
    /// number of elements that are `< target`.
    pub(crate) fn linear_search_sse2_128(arr: &AlignedBuffer, target: u32) -> usize {
        unsafe {
            let ptr = arr as *const AlignedBuffer as *const DataType;
            let vkey = set1(target as i32);
            let mut cnt = set0();
            // We work over 4 `__m128i` at a time.
            // A single `__m128i` actual contains 4 `u32`.
            for i in 0..(COMPRESSION_BLOCK_SIZE as isize) / (4 * 4) {
                let cmp1 = op_lt(op_load(ptr.offset(i * 4)), vkey);
                let cmp2 = op_lt(op_load(ptr.offset(i * 4 + 1)), vkey);
                let cmp3 = op_lt(op_load(ptr.offset(i * 4 + 2)), vkey);
                let cmp4 = op_lt(op_load(ptr.offset(i * 4 + 3)), vkey);
                let sum = op_add(op_add(cmp1, cmp2), op_add(cmp3, cmp4));
                cnt = op_sub(cnt, sum);
            }
            cnt = op_add(cnt, _mm_shuffle_epi32(cnt, MASK1));
            cnt = op_add(cnt, _mm_shuffle_epi32(cnt, MASK2));
            _mm_cvtsi128_si32(cnt) as usize
        }
    }

    #[cfg(test)]
    mod test {
        use super::linear_search_sse2_128;
        use crate::postings::compression::{AlignedBuffer, COMPRESSION_BLOCK_SIZE};

        #[test]
        fn test_linear_search_sse2_128_u32() {
            let mut block = [0u32; COMPRESSION_BLOCK_SIZE];
            for el in 0u32..128u32 {
                block[el as usize] = el * 2 + 1 << 18;
            }
            let target = block[64] + 1;
            assert_eq!(linear_search_sse2_128(&AlignedBuffer(block), target), 65);
        }
    }
}

/// This `linear search` browser exhaustively through the array.
/// but the early exit is very difficult to predict.
///
/// Coupled with `exponential search` this function is likely
/// to be called with the same `len`
fn linear_search(arr: &[u32], target: u32) -> usize {
    arr.iter().map(|&el| if el < target { 1 } else { 0 }).sum()
}

fn exponential_search(arr: &[u32], target: u32) -> (usize, usize) {
    let end = arr.len();
    let mut begin = 0;
    for &pivot in &[1, 3, 7, 15, 31, 63] {
        if pivot >= end {
            break;
        }
        if arr[pivot] > target {
            return (begin, pivot);
        }
        begin = pivot;
    }
    (begin, end)
}

fn galloping(block_docs: &[u32], target: u32) -> usize {
    let (start, end) = exponential_search(&block_docs, target);
    start + linear_search(&block_docs[start..end], target)
}

/// Tantivy may rely on SIMD instructions to search for a specific document within
/// a given block.
#[derive(Clone, Copy, PartialEq)]
pub enum BlockSearcher {
    #[cfg(target_arch = "x86_64")]
    SSE2,
    Scalar,
}

impl BlockSearcher {
    /// Search the first index containing an element greater or equal to
    /// the target.
    ///
    /// The results should be equivalent to
    /// ```compile_fail
    /// block[..]
    //       .iter()
    //       .take_while(|&&val| val < target)
    //       .count()
    /// ```
    ///
    /// The `start` argument is just used to hint that the response is
    /// greater than beyond `start`. The implementation may or may not use
    /// it for optimization.
    ///
    /// # Assumption
    ///
    /// The array len is > start.
    /// The block is sorted
    /// The target is assumed greater or equal to the `arr[start]`.
    /// The target is assumed smaller or equal to the last element of the block.
    ///
    /// Currently the scalar implementation starts by an exponential search, and
    /// then operates a linear search in the result subarray.
    ///
    /// If SSE2 instructions are available in the `(platform, running CPU)`,
    /// then we use a different implementation that does an exhaustive linear search over
    /// the full block whenever the block is full (`len == 128`). It is surprisingly faster, most likely because of the lack
    /// of branch.
    pub(crate) fn search_in_block(
        self,
        block_docs: &AlignedBuffer,
        len: usize,
        start: usize,
        target: u32,
    ) -> usize {
        #[cfg(target_arch = "x86_64")]
        {
            use crate::postings::compression::COMPRESSION_BLOCK_SIZE;
            if self == BlockSearcher::SSE2 && len == COMPRESSION_BLOCK_SIZE {
                return sse2::linear_search_sse2_128(block_docs, target);
            }
        }
        start + galloping(&block_docs.0[start..len], target)
    }
}

impl Default for BlockSearcher {
    fn default() -> BlockSearcher {
        #[cfg(target_arch = "x86_64")]
        {
            if is_x86_feature_detected!("sse2") {
                return BlockSearcher::SSE2;
            }
        }
        BlockSearcher::Scalar
    }
}

#[cfg(test)]
mod tests {
    use super::exponential_search;
    use super::linear_search;
    use super::BlockSearcher;
    use crate::postings::compression::{AlignedBuffer, COMPRESSION_BLOCK_SIZE};

    #[test]
    fn test_linear_search() {
        let len: usize = 50;
        let arr: Vec<u32> = (0..len).map(|el| 1u32 + (el as u32) * 2).collect();
        for target in 1..*arr.last().unwrap() {
            let res = linear_search(&arr[..], target);
            if res > 0 {
                assert!(arr[res - 1] < target);
            }
            if res < len {
                assert!(arr[res] >= target);
            }
        }
    }

    #[test]
    fn test_exponentiel_search() {
        assert_eq!(exponential_search(&[1, 2], 0), (0, 1));
        assert_eq!(exponential_search(&[1, 2], 1), (0, 1));
        assert_eq!(
            exponential_search(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], 7),
            (3, 7)
        );
    }

    fn util_test_search_in_block(block_searcher: BlockSearcher, block: &[u32], target: u32) {
        let cursor = search_in_block_trivial_but_slow(block, target);
        assert!(block.len() < COMPRESSION_BLOCK_SIZE);
        let mut output_buffer = [u32::max_value(); COMPRESSION_BLOCK_SIZE];
        output_buffer[..block.len()].copy_from_slice(block);
        for i in 0..cursor {
            assert_eq!(
                block_searcher.search_in_block(
                    &AlignedBuffer(output_buffer),
                    block.len(),
                    i,
                    target
                ),
                cursor
            );
        }
    }

    fn util_test_search_in_block_all(block_searcher: BlockSearcher, block: &[u32]) {
        use std::collections::HashSet;
        let mut targets = HashSet::new();
        for (i, val) in block.iter().cloned().enumerate() {
            if i > 0 {
                targets.insert(val - 1);
            }
            targets.insert(val);
        }
        for target in targets {
            util_test_search_in_block(block_searcher, block, target);
        }
    }

    fn search_in_block_trivial_but_slow(block: &[u32], target: u32) -> usize {
        block.iter().take_while(|&&val| val < target).count()
    }

    fn test_search_in_block_util(block_searcher: BlockSearcher) {
        for len in 1u32..128u32 {
            let v: Vec<u32> = (0..len).map(|i| i * 2).collect();
            util_test_search_in_block_all(block_searcher, &v[..]);
        }
    }

    #[test]
    fn test_search_in_block_scalar() {
        test_search_in_block_util(BlockSearcher::Scalar);
    }

    #[cfg(target_arch = "x86_64")]
    #[test]
    fn test_search_in_block_sse2() {
        test_search_in_block_util(BlockSearcher::SSE2);
    }
}
<|endoftext|>


<commit_before>/*!
Migrant can be used as a library so you can embed the management of migrations
into your binary and don't need to use a secondary tool in production environments.

The majority of `src/main.rs` could be copied, or just select functionality.
*/
extern crate migrant;

use std::env;
use std::path::PathBuf;
use migrant::Config;

pub fn main() {
    let dir = env::current_dir().unwrap();
    let config_path = match migrant::search_for_config(&dir) {
        None => {
            migrant::Config::init(&dir).expect("failed to initialize project");
            return;
        }
        Some(p) => p,
    };
    let base_dir = config_path.parent()
        .map(PathBuf::from)
        .expect("failed to get parent dir");
    let config = Config::load(&config_path)
        .expect("failed to load config");

    // This will fail if no migration files are present
    //migrant::apply_migration(&base_dir, &config_path, &config, migrant::Direction::Up, false, false, true)
    //    .expect("failed to apply migrations")

    migrant::list(&config, &base_dir)
        .expect("failed to list migrations");
}
<commit_msg>update embedded example<commit_after>/*!
Migrant can be used as a library so you can embed the management of migrations
into your binary and don't need to use a secondary tool in production environments.

The majority of `src/main.rs` could be copied, or just select functionality.
*/
extern crate migrant;

use std::env;
use std::path::PathBuf;
use migrant::Config;

pub fn main() {
    let dir = env::current_dir().unwrap();
    let config_path = match migrant::search_for_config(&dir) {
        None => {
            migrant::Config::init(&dir).expect("failed to initialize project");
            return;
        }
        Some(p) => p,
    };
    let base_dir = config_path.parent()
        .map(PathBuf::from)
        .expect("failed to get parent dir");
    let config = Config::load(&config_path)
        .expect("failed to load config");

    // This will fail if no migration files are present!
    // Run all available `up` migrations
    // migrant::Migrator::with_config(&config, &config_path, &base_dir)
    //     .direction(migrant::Direction::Up)
    //     .all(true)
    //     .apply()
    //     .expect("failed to apply migrations")

    migrant::list(&config, &base_dir)
        .expect("failed to list migrations");
}
<|endoftext|>


<commit_before>use std::process::exit;

extern crate nix;
extern crate rotor;
extern crate rotor_stream;
extern crate argparse;
extern crate httparse;

use std::cmp::max;
use std::str::from_utf8;
use std::net::ToSocketAddrs;
use std::io::{stdout, stderr, Write};
use std::time::Duration;

use argparse::{ArgumentParser, Store};
use rotor::mio::tcp::{TcpStream};
use rotor_stream::{Stream, Transport, Protocol, Intent};
use rotor::{Scope};

struct Context;

enum Http {
    SendRequest(String),
    ReadHeaders,
    ReadBody,
}

impl<'a> Protocol for Http {
    type Context = Context;
    type Socket = TcpStream;
    type Seed = (String, String);
    fn create((host, path): Self::Seed, _sock: &mut TcpStream,
        scope: &mut Scope<Context>)
        -> Intent<Self>
    {
        // Wait socket to become writable
        let req = format!(concat!(
                "GET {} HTTP/1.1\r\n",
                "Host: {}\r\n",
                "\r\n"),
            path, host);
        println!("----- Request -----");
        print!("{}", req);
        Intent::of(Http::SendRequest(req)).expect_flush()
            .deadline(scope.now() + Duration::new(10, 0))
    }
    fn bytes_flushed(self, transport: &mut Transport<TcpStream>,
        scope: &mut Scope<Context>)
        -> Intent<Self>
    {
        match self {
            Http::SendRequest(val) => {
                transport.output().extend(val.as_bytes());
                Intent::of(Http::ReadHeaders)
                    .expect_delimiter(b"\r\n\r\n", 4096)
                    .deadline(scope.now() + Duration::new(10, 0))
            }
            _ => unreachable!(),
        }
    }
    fn bytes_read(self, transport: &mut Transport<TcpStream>,
                  end: usize, scope: &mut Scope<Context>)
        -> Intent<Self>
    {
        match self {
            Http::SendRequest(_) => unreachable!(),
            Http::ReadHeaders => {
                let mut headers = [httparse::EMPTY_HEADER; 16];
                {
                    let mut req = httparse::Response::new(&mut headers);
                    let buf = &transport.input()[..end+4];
                    req.parse(buf).unwrap();
                }
                let mut clen = 0;
                println!("----- Headers -----");
                for header in headers.iter() {
                    println!("{}: {}",
                        header.name,
                        from_utf8(header.value).unwrap());
                    let hlower = header.name.to_lowercase();
                    if &hlower[..] == "content-length" {
                        if let Some(val) = from_utf8(header.value).ok()
                                           .and_then(|x| x.parse().ok()) {
                            clen = val;
                            break;
                        }
                    }
                }
                println!("----- Body [{}] -----", clen);
                transport.input().consume(end + 4);
                Intent::of(Http::ReadBody).expect_bytes(clen)
                    .deadline(scope.now() +
                              Duration::new(max(10, clen as u64), 0))
            }
            Http::ReadBody => {
                let newline = {
                    let data = &transport.input()[..];
                    data.len() > 0 && &data[data.len()-1..] != b"\n"
                };
                transport.input().write_to(&mut stdout()).ok();
                if newline {
                    println!("");
                }
                println!("----- Done -----");
                Intent::done()
            }
        }
    }
    fn timeout(self, _transport: &mut Transport<TcpStream>,
        _scope: &mut Scope<Context>)
        -> Intent<Self>
    {
        writeln!(&mut stderr(), "Timeout reached").ok();
        Intent::done()
    }

    /// Message received (from the main loop)
    fn wakeup(self, _transport: &mut Transport<TcpStream>,
        _scope: &mut Scope<Context>)
        -> Intent<Self>
    {
        unreachable!();
    }
}

fn main() {
    let mut url = "".to_string();
    {
        let mut ap = ArgumentParser::new();
        ap.set_description("
            Fetches some url. Only http, and only fixed size pages are
            supported (i.e. where Content-Length is set)
            ");
        ap.refer(&mut url).add_argument("url", Store, "
            The url to fetch.");
        ap.parse_args_or_exit();
    }
    if !url.starts_with("http://") {
        println!("Url should start with http://");
        exit(1);
    }
    let (host, path) = if let Some(idx) = url[7..].find('/') {
        (&url[7..idx+7], &url[idx+7..])
    } else {
        (&url[7..], "/")
    };
    println!("Host: {} (port: 80), path: {}", host, path);

    let event_loop = rotor::Loop::new(&rotor::Config::new()).unwrap();

    let sock = TcpStream::connect(
        // Any better way for current stable rust?
        &(host, 80).to_socket_addrs().unwrap().next().unwrap()).unwrap();

    let mut loop_inst = event_loop.instantiate(Context);
    loop_inst.add_machine_with(|scope| {
        Stream::<Http>::new(
            sock, (host.to_string(), path.to_string()),
            scope)
    }).unwrap();
    loop_inst.run().unwrap();
}
<commit_msg>Better error reporting in fetch example<commit_after>use std::process::exit;

extern crate nix;
extern crate rotor;
extern crate rotor_stream;
extern crate argparse;
extern crate httparse;

use std::cmp::max;
use std::str::from_utf8;
use std::net::ToSocketAddrs;
use std::io::{stdout, stderr, Write};
use std::time::Duration;

use argparse::{ArgumentParser, Store};
use rotor::mio::tcp::{TcpStream};
use rotor_stream::{Stream, Transport, Protocol, Intent, Exception};
use rotor::{Scope};

struct Context;

enum Http {
    SendRequest(String),
    ReadHeaders,
    ReadBody,
}

impl<'a> Protocol for Http {
    type Context = Context;
    type Socket = TcpStream;
    type Seed = (String, String);
    fn create((host, path): Self::Seed, _sock: &mut TcpStream,
        scope: &mut Scope<Context>)
        -> Intent<Self>
    {
        // Wait socket to become writable
        let req = format!(concat!(
                "GET {} HTTP/1.1\r\n",
                "Host: {}\r\n",
                "\r\n"),
            path, host);
        println!("----- Request -----");
        print!("{}", req);
        Intent::of(Http::SendRequest(req)).expect_flush()
            .deadline(scope.now() + Duration::new(10, 0))
    }
    fn bytes_flushed(self, transport: &mut Transport<TcpStream>,
        scope: &mut Scope<Context>)
        -> Intent<Self>
    {
        match self {
            Http::SendRequest(val) => {
                transport.output().extend(val.as_bytes());
                Intent::of(Http::ReadHeaders)
                    .expect_delimiter(b"\r\n\r\n", 4096)
                    .deadline(scope.now() + Duration::new(10, 0))
            }
            _ => unreachable!(),
        }
    }
    fn bytes_read(self, transport: &mut Transport<TcpStream>,
                  end: usize, scope: &mut Scope<Context>)
        -> Intent<Self>
    {
        match self {
            Http::SendRequest(_) => unreachable!(),
            Http::ReadHeaders => {
                let mut headers = [httparse::EMPTY_HEADER; 16];
                {
                    let mut req = httparse::Response::new(&mut headers);
                    let buf = &transport.input()[..end+4];
                    req.parse(buf).unwrap();
                }
                let mut clen = 0;
                println!("----- Headers -----");
                for header in headers.iter() {
                    println!("{}: {}",
                        header.name,
                        from_utf8(header.value).unwrap());
                    let hlower = header.name.to_lowercase();
                    if &hlower[..] == "content-length" {
                        if let Some(val) = from_utf8(header.value).ok()
                                           .and_then(|x| x.parse().ok()) {
                            clen = val;
                            break;
                        }
                    }
                }
                println!("----- Body [{}] -----", clen);
                transport.input().consume(end + 4);
                Intent::of(Http::ReadBody).expect_bytes(clen)
                    .deadline(scope.now() +
                              Duration::new(max(10, clen as u64), 0))
            }
            Http::ReadBody => {
                let newline = {
                    let data = &transport.input()[..];
                    data.len() > 0 && &data[data.len()-1..] != b"\n"
                };
                transport.input().write_to(&mut stdout()).ok();
                if newline {
                    println!("");
                }
                println!("----- Done -----");
                Intent::done()
            }
        }
    }
    fn timeout(self, _transport: &mut Transport<TcpStream>,
        _scope: &mut Scope<Context>)
        -> Intent<Self>
    {
        writeln!(&mut stderr(), "Timeout reached").ok();
        Intent::done()
    }

    /// Message received (from the main loop)
    fn wakeup(self, _transport: &mut Transport<TcpStream>,
        _scope: &mut Scope<Context>)
        -> Intent<Self>
    {
        unreachable!();
    }

    fn exception(self, _transport: &mut Transport<Self::Socket>,
        reason: Exception, scope: &mut Scope<Self::Context>)
        -> Intent<Self>
    {
        writeln!(&mut stderr(), "Error when fetching data: {}", reason).ok();
        scope.shutdown_loop();
        Intent::error(Box::new(reason))
    }
}

fn main() {
    let mut url = "".to_string();
    {
        let mut ap = ArgumentParser::new();
        ap.set_description("
            Fetches some url. Only http, and only fixed size pages are
            supported (i.e. where Content-Length is set)
            ");
        ap.refer(&mut url).add_argument("url", Store, "
            The url to fetch.");
        ap.parse_args_or_exit();
    }
    if !url.starts_with("http://") {
        println!("Url should start with http://");
        exit(1);
    }
    let (host, path) = if let Some(idx) = url[7..].find('/') {
        (&url[7..idx+7], &url[idx+7..])
    } else {
        (&url[7..], "/")
    };
    println!("Host: {} (port: 80), path: {}", host, path);

    let event_loop = rotor::Loop::new(&rotor::Config::new()).unwrap();

    let sock = TcpStream::connect(
        // Any better way for current stable rust?
        &(host, 80).to_socket_addrs().unwrap().next().unwrap()).unwrap();

    let mut loop_inst = event_loop.instantiate(Context);
    loop_inst.add_machine_with(|scope| {
        Stream::<Http>::new(
            sock, (host.to_string(), path.to_string()),
            scope)
    }).unwrap();
    loop_inst.run().unwrap();
}
<|endoftext|>


<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![forbid(unsafe_code)]

//! This crate implements the storage [GRPC](http://grpc.io) service.
//!
//! The user of storage service is supposed to use it via client lib provided in
//! [`storage-client`](../storage-client/index.html) instead of via
//! [`StorageClient`](../storage-proto/proto/storage_grpc/struct.StorageClient.html) directly.

#[cfg(feature = "fuzzing")]
pub mod mocks;

use anyhow::Result;
use futures::channel::mpsc;
use futures::sink::SinkExt;
use libra_config::config::NodeConfig;
use libra_logger::prelude::*;
use libra_types::proto::types::{
    UpdateToLatestLedgerRequest, UpdateToLatestLedgerResponse, ValidatorChangeProof,
};
use libradb::LibraDB;
use std::{convert::TryFrom, net::ToSocketAddrs, ops::Deref, path::Path, sync::Arc};
use storage_proto::proto::storage::{
    storage_server::{Storage, StorageServer},
    BackupAccountStateRequest, BackupAccountStateResponse, GetAccountStateRangeProofRequest,
    GetAccountStateRangeProofResponse, GetAccountStateWithProofByVersionRequest,
    GetAccountStateWithProofByVersionResponse, GetEpochChangeLedgerInfosRequest,
    GetLatestAccountStateRequest, GetLatestAccountStateResponse, GetLatestStateRootRequest,
    GetLatestStateRootResponse, GetStartupInfoRequest, GetStartupInfoResponse,
    GetTransactionsRequest, GetTransactionsResponse, SaveTransactionsRequest,
    SaveTransactionsResponse,
};
use tokio::runtime::Runtime;

/// Starts storage service according to config.
pub fn start_storage_service(config: &NodeConfig) -> Runtime {
    let rt = tokio::runtime::Runtime::new().unwrap();

    let storage_service = StorageService::new(&config.storage.dir());

    let addr = format!("{}:{}", config.storage.address, config.storage.port)
        .to_socket_addrs()
        .unwrap()
        .next()
        .unwrap();
    rt.spawn(
        tonic::transport::Server::builder()
            .add_service(StorageServer::new(storage_service))
            .serve(addr),
    );
    rt
}

/// The implementation of the storage [GRPC](http://grpc.io) service.
///
/// It serves [`LibraDB`] APIs over the network. See API documentation in [`storage-proto`] and
/// [`LibraDB`].
#[derive(Clone)]
pub struct StorageService {
    db: Arc<LibraDBWrapper>,
}

struct LibraDBWrapper {
    db: Option<LibraDB>,
}

impl LibraDBWrapper {
    pub fn new<P: AsRef<Path>>(path: &P) -> Self {
        let db = LibraDB::new(path);
        Self { db: Some(db) }
    }
}

impl Deref for LibraDBWrapper {
    type Target = LibraDB;

    fn deref(&self) -> &Self::Target {
        self.db.as_ref().expect("LibraDB is dropped unexpectedly")
    }
}

impl StorageService {
    /// This opens a [`LibraDB`] at `path` and returns a [`StorageService`] instance serving it.
    pub fn new<P: AsRef<Path>>(path: &P) -> Self {
        let db_wrapper = LibraDBWrapper::new(path);
        Self {
            db: Arc::new(db_wrapper),
        }
    }
}

impl StorageService {
    fn update_to_latest_ledger_inner(
        &self,
        req: UpdateToLatestLedgerRequest,
    ) -> Result<UpdateToLatestLedgerResponse> {
        let rust_req = libra_types::get_with_proof::UpdateToLatestLedgerRequest::try_from(req)?;

        let (
            response_items,
            ledger_info_with_sigs,
            validator_change_proof,
            ledger_consistency_proof,
        ) = self
            .db
            .update_to_latest_ledger(rust_req.client_known_version, rust_req.requested_items)?;

        let rust_resp = libra_types::get_with_proof::UpdateToLatestLedgerResponse {
            response_items,
            ledger_info_with_sigs,
            validator_change_proof,
            ledger_consistency_proof,
        };

        Ok(rust_resp.into())
    }

    fn get_transactions_inner(
        &self,
        req: GetTransactionsRequest,
    ) -> Result<GetTransactionsResponse> {
        let rust_req = storage_proto::GetTransactionsRequest::try_from(req)?;

        let txn_list_with_proof = self.db.get_transactions(
            rust_req.start_version,
            rust_req.batch_size,
            rust_req.ledger_version,
            rust_req.fetch_events,
        )?;

        let rust_resp = storage_proto::GetTransactionsResponse::new(txn_list_with_proof);

        Ok(rust_resp.into())
    }

    fn get_latest_state_root_inner(
        &self,
        _req: GetLatestStateRootRequest,
    ) -> Result<GetLatestStateRootResponse> {
        let (version, state_root_hash) = self.db.get_latest_state_root()?;
        let rust_resp = storage_proto::GetLatestStateRootResponse::new(version, state_root_hash);
        Ok(rust_resp.into())
    }

    fn get_latest_account_state_inner(
        &self,
        req: GetLatestAccountStateRequest,
    ) -> Result<GetLatestAccountStateResponse> {
        let rust_req = storage_proto::GetLatestAccountStateRequest::try_from(req)?;
        let account_state_blob = self.db.get_latest_account_state(rust_req.address)?;
        let rust_resp = storage_proto::GetLatestAccountStateResponse::new(account_state_blob);
        Ok(rust_resp.into())
    }

    fn get_account_state_with_proof_by_version_inner(
        &self,
        req: GetAccountStateWithProofByVersionRequest,
    ) -> Result<GetAccountStateWithProofByVersionResponse> {
        let rust_req = storage_proto::GetAccountStateWithProofByVersionRequest::try_from(req)?;

        let (account_state_blob, sparse_merkle_proof) = self
            .db
            .get_account_state_with_proof_by_version(rust_req.address, rust_req.version)?;

        let rust_resp = storage_proto::GetAccountStateWithProofByVersionResponse {
            account_state_blob,
            sparse_merkle_proof,
        };

        Ok(rust_resp.into())
    }

    fn save_transactions_inner(
        &self,
        req: SaveTransactionsRequest,
    ) -> Result<SaveTransactionsResponse> {
        let rust_req = storage_proto::SaveTransactionsRequest::try_from(req)?;
        self.db.save_transactions(
            &rust_req.txns_to_commit,
            rust_req.first_version,
            rust_req.ledger_info_with_signatures.as_ref(),
        )?;
        Ok(SaveTransactionsResponse::default())
    }

    fn get_startup_info_inner(&self) -> Result<GetStartupInfoResponse> {
        let info = self.db.get_startup_info()?;
        let rust_resp = storage_proto::GetStartupInfoResponse { info };
        Ok(rust_resp.into())
    }

    fn get_epoch_change_ledger_infos_inner(
        &self,
        req: GetEpochChangeLedgerInfosRequest,
    ) -> Result<ValidatorChangeProof> {
        let rust_req = storage_proto::GetEpochChangeLedgerInfosRequest::try_from(req)?;
        let (ledger_infos, more) = self
            .db
            .get_epoch_change_ledger_infos(rust_req.start_epoch, rust_req.end_epoch)?;
        let rust_resp =
            libra_types::validator_change::ValidatorChangeProof::new(ledger_infos, more);
        Ok(rust_resp.into())
    }

    fn get_account_state_range_proof_inner(
        &self,
        req: GetAccountStateRangeProofRequest,
    ) -> Result<GetAccountStateRangeProofResponse> {
        let rust_req = storage_proto::GetAccountStateRangeProofRequest::try_from(req)?;
        let proof = self
            .db
            .get_account_state_range_proof(rust_req.rightmost_key, rust_req.version)?;
        let rust_resp = storage_proto::GetAccountStateRangeProofResponse::new(proof);
        Ok(rust_resp.into())
    }
}

#[tonic::async_trait]
impl Storage for StorageService {
    async fn save_transactions(
        &self,
        request: tonic::Request<SaveTransactionsRequest>,
    ) -> Result<tonic::Response<SaveTransactionsResponse>, tonic::Status> {
        debug!("[GRPC] Storage::save_transactions");
        let req = request.into_inner();
        let resp = self
            .save_transactions_inner(req)
            .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;
        Ok(tonic::Response::new(resp))
    }

    async fn update_to_latest_ledger(
        &self,
        request: tonic::Request<UpdateToLatestLedgerRequest>,
    ) -> Result<tonic::Response<UpdateToLatestLedgerResponse>, tonic::Status> {
        debug!("[GRPC] Storage::update_to_latest_ledger");
        let req = request.into_inner();
        let resp = self
            .update_to_latest_ledger_inner(req)
            .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;
        Ok(tonic::Response::new(resp))
    }

    async fn get_transactions(
        &self,
        request: tonic::Request<GetTransactionsRequest>,
    ) -> Result<tonic::Response<GetTransactionsResponse>, tonic::Status> {
        debug!("[GRPC] Storage::get_transactions");
        let req = request.into_inner();
        let resp = self
            .get_transactions_inner(req)
            .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;
        Ok(tonic::Response::new(resp))
    }

    async fn get_latest_state_root(
        &self,
        request: tonic::Request<GetLatestStateRootRequest>,
    ) -> Result<tonic::Response<GetLatestStateRootResponse>, tonic::Status> {
        debug!("[GRPC] Storage::get_latest_state_root");
        let req = request.into_inner();
        let resp = self
            .get_latest_state_root_inner(req)
            .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;
        Ok(tonic::Response::new(resp))
    }

    async fn get_latest_account_state(
        &self,
        request: tonic::Request<GetLatestAccountStateRequest>,
    ) -> Result<tonic::Response<GetLatestAccountStateResponse>, tonic::Status> {
        debug!("[GRPC] Storage::get_latest_account_state");
        let req = request.into_inner();
        let resp = self
            .get_latest_account_state_inner(req)
            .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;
        Ok(tonic::Response::new(resp))
    }

    async fn get_account_state_with_proof_by_version(
        &self,
        request: tonic::Request<GetAccountStateWithProofByVersionRequest>,
    ) -> Result<tonic::Response<GetAccountStateWithProofByVersionResponse>, tonic::Status> {
        debug!("[GRPC] Storage::get_account_state_with_proof_by_version");
        let req = request.into_inner();
        let resp = self
            .get_account_state_with_proof_by_version_inner(req)
            .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;
        Ok(tonic::Response::new(resp))
    }

    async fn get_startup_info(
        &self,
        _request: tonic::Request<GetStartupInfoRequest>,
    ) -> Result<tonic::Response<GetStartupInfoResponse>, tonic::Status> {
        debug!("[GRPC] Storage::get_startup_info");
        let resp = self
            .get_startup_info_inner()
            .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;
        Ok(tonic::Response::new(resp))
    }

    async fn get_epoch_change_ledger_infos(
        &self,
        request: tonic::Request<GetEpochChangeLedgerInfosRequest>,
    ) -> Result<tonic::Response<ValidatorChangeProof>, tonic::Status> {
        debug!("[GRPC] Storage::get_epoch_change_ledger_infos");
        let req = request.into_inner();
        let resp = self
            .get_epoch_change_ledger_infos_inner(req)
            .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;
        Ok(tonic::Response::new(resp))
    }

    type BackupAccountStateStream =
        mpsc::Receiver<Result<BackupAccountStateResponse, tonic::Status>>;

    async fn backup_account_state(
        &self,
        request: tonic::Request<BackupAccountStateRequest>,
    ) -> Result<tonic::Response<Self::BackupAccountStateStream>, tonic::Status> {
        debug!("[GRPC] Storage::backup_account_state");
        let req = request.into_inner();
        let iter = self
            .db
            .get_account_iter(req.version)
            .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;

        let iter = iter.map(|res| match res {
            Ok((hash, blob)) => {
                let resp: BackupAccountStateResponse =
                    storage_proto::BackupAccountStateResponse::new(hash, blob).into();
                Ok(resp)
            }
            Err(e) => Err(tonic::Status::new(tonic::Code::Internal, e.to_string())),
        });

        // Channel to buffer the stream
        let (mut tx, rx) = mpsc::channel(32);

        tokio::spawn(async move {
            for resp in iter {
                tx.send(resp).await.unwrap();
            }
        });

        Ok(tonic::Response::new(rx))
    }

    async fn get_account_state_range_proof(
        &self,
        request: tonic::Request<GetAccountStateRangeProofRequest>,
    ) -> Result<tonic::Response<GetAccountStateRangeProofResponse>, tonic::Status> {
        debug!("[GRPC] Storage::get_account_state_range_proof");
        let req = request.into_inner();
        let resp = self
            .get_account_state_range_proof_inner(req)
            .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;
        Ok(tonic::Response::new(resp))
    }
}

#[cfg(test)]
mod storage_service_test;
<commit_msg>[Storage] Remove LibraDBWrapper<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![forbid(unsafe_code)]

//! This crate implements the storage [GRPC](http://grpc.io) service.
//!
//! The user of storage service is supposed to use it via client lib provided in
//! [`storage-client`](../storage-client/index.html) instead of via
//! [`StorageClient`](../storage-proto/proto/storage_grpc/struct.StorageClient.html) directly.

#[cfg(feature = "fuzzing")]
pub mod mocks;

use anyhow::Result;
use futures::channel::mpsc;
use futures::sink::SinkExt;
use libra_config::config::NodeConfig;
use libra_logger::prelude::*;
use libra_types::proto::types::{
    UpdateToLatestLedgerRequest, UpdateToLatestLedgerResponse, ValidatorChangeProof,
};
use libradb::LibraDB;
use std::{convert::TryFrom, net::ToSocketAddrs, path::Path, sync::Arc};
use storage_proto::proto::storage::{
    storage_server::{Storage, StorageServer},
    BackupAccountStateRequest, BackupAccountStateResponse, GetAccountStateRangeProofRequest,
    GetAccountStateRangeProofResponse, GetAccountStateWithProofByVersionRequest,
    GetAccountStateWithProofByVersionResponse, GetEpochChangeLedgerInfosRequest,
    GetLatestAccountStateRequest, GetLatestAccountStateResponse, GetLatestStateRootRequest,
    GetLatestStateRootResponse, GetStartupInfoRequest, GetStartupInfoResponse,
    GetTransactionsRequest, GetTransactionsResponse, SaveTransactionsRequest,
    SaveTransactionsResponse,
};
use tokio::runtime::Runtime;

/// Starts storage service according to config.
pub fn start_storage_service(config: &NodeConfig) -> Runtime {
    let rt = tokio::runtime::Runtime::new().unwrap();

    let storage_service = StorageService::new(&config.storage.dir());

    let addr = format!("{}:{}", config.storage.address, config.storage.port)
        .to_socket_addrs()
        .unwrap()
        .next()
        .unwrap();
    rt.spawn(
        tonic::transport::Server::builder()
            .add_service(StorageServer::new(storage_service))
            .serve(addr),
    );
    rt
}

/// The implementation of the storage [GRPC](http://grpc.io) service.
///
/// It serves [`LibraDB`] APIs over the network. See API documentation in [`storage-proto`] and
/// [`LibraDB`].
#[derive(Clone)]
pub struct StorageService {
    db: Arc<LibraDB>,
}

impl StorageService {
    /// This opens a [`LibraDB`] at `path` and returns a [`StorageService`] instance serving it.
    pub fn new<P: AsRef<Path>>(path: &P) -> Self {
        let db = Arc::new(LibraDB::new(path));
        Self { db }
    }
}

impl StorageService {
    fn update_to_latest_ledger_inner(
        &self,
        req: UpdateToLatestLedgerRequest,
    ) -> Result<UpdateToLatestLedgerResponse> {
        let rust_req = libra_types::get_with_proof::UpdateToLatestLedgerRequest::try_from(req)?;

        let (
            response_items,
            ledger_info_with_sigs,
            validator_change_proof,
            ledger_consistency_proof,
        ) = self
            .db
            .update_to_latest_ledger(rust_req.client_known_version, rust_req.requested_items)?;

        let rust_resp = libra_types::get_with_proof::UpdateToLatestLedgerResponse {
            response_items,
            ledger_info_with_sigs,
            validator_change_proof,
            ledger_consistency_proof,
        };

        Ok(rust_resp.into())
    }

    fn get_transactions_inner(
        &self,
        req: GetTransactionsRequest,
    ) -> Result<GetTransactionsResponse> {
        let rust_req = storage_proto::GetTransactionsRequest::try_from(req)?;

        let txn_list_with_proof = self.db.get_transactions(
            rust_req.start_version,
            rust_req.batch_size,
            rust_req.ledger_version,
            rust_req.fetch_events,
        )?;

        let rust_resp = storage_proto::GetTransactionsResponse::new(txn_list_with_proof);

        Ok(rust_resp.into())
    }

    fn get_latest_state_root_inner(
        &self,
        _req: GetLatestStateRootRequest,
    ) -> Result<GetLatestStateRootResponse> {
        let (version, state_root_hash) = self.db.get_latest_state_root()?;
        let rust_resp = storage_proto::GetLatestStateRootResponse::new(version, state_root_hash);
        Ok(rust_resp.into())
    }

    fn get_latest_account_state_inner(
        &self,
        req: GetLatestAccountStateRequest,
    ) -> Result<GetLatestAccountStateResponse> {
        let rust_req = storage_proto::GetLatestAccountStateRequest::try_from(req)?;
        let account_state_blob = self.db.get_latest_account_state(rust_req.address)?;
        let rust_resp = storage_proto::GetLatestAccountStateResponse::new(account_state_blob);
        Ok(rust_resp.into())
    }

    fn get_account_state_with_proof_by_version_inner(
        &self,
        req: GetAccountStateWithProofByVersionRequest,
    ) -> Result<GetAccountStateWithProofByVersionResponse> {
        let rust_req = storage_proto::GetAccountStateWithProofByVersionRequest::try_from(req)?;

        let (account_state_blob, sparse_merkle_proof) = self
            .db
            .get_account_state_with_proof_by_version(rust_req.address, rust_req.version)?;

        let rust_resp = storage_proto::GetAccountStateWithProofByVersionResponse {
            account_state_blob,
            sparse_merkle_proof,
        };

        Ok(rust_resp.into())
    }

    fn save_transactions_inner(
        &self,
        req: SaveTransactionsRequest,
    ) -> Result<SaveTransactionsResponse> {
        let rust_req = storage_proto::SaveTransactionsRequest::try_from(req)?;
        self.db.save_transactions(
            &rust_req.txns_to_commit,
            rust_req.first_version,
            rust_req.ledger_info_with_signatures.as_ref(),
        )?;
        Ok(SaveTransactionsResponse::default())
    }

    fn get_startup_info_inner(&self) -> Result<GetStartupInfoResponse> {
        let info = self.db.get_startup_info()?;
        let rust_resp = storage_proto::GetStartupInfoResponse { info };
        Ok(rust_resp.into())
    }

    fn get_epoch_change_ledger_infos_inner(
        &self,
        req: GetEpochChangeLedgerInfosRequest,
    ) -> Result<ValidatorChangeProof> {
        let rust_req = storage_proto::GetEpochChangeLedgerInfosRequest::try_from(req)?;
        let (ledger_infos, more) = self
            .db
            .get_epoch_change_ledger_infos(rust_req.start_epoch, rust_req.end_epoch)?;
        let rust_resp =
            libra_types::validator_change::ValidatorChangeProof::new(ledger_infos, more);
        Ok(rust_resp.into())
    }

    fn get_account_state_range_proof_inner(
        &self,
        req: GetAccountStateRangeProofRequest,
    ) -> Result<GetAccountStateRangeProofResponse> {
        let rust_req = storage_proto::GetAccountStateRangeProofRequest::try_from(req)?;
        let proof = self
            .db
            .get_account_state_range_proof(rust_req.rightmost_key, rust_req.version)?;
        let rust_resp = storage_proto::GetAccountStateRangeProofResponse::new(proof);
        Ok(rust_resp.into())
    }
}

#[tonic::async_trait]
impl Storage for StorageService {
    async fn save_transactions(
        &self,
        request: tonic::Request<SaveTransactionsRequest>,
    ) -> Result<tonic::Response<SaveTransactionsResponse>, tonic::Status> {
        debug!("[GRPC] Storage::save_transactions");
        let req = request.into_inner();
        let resp = self
            .save_transactions_inner(req)
            .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;
        Ok(tonic::Response::new(resp))
    }

    async fn update_to_latest_ledger(
        &self,
        request: tonic::Request<UpdateToLatestLedgerRequest>,
    ) -> Result<tonic::Response<UpdateToLatestLedgerResponse>, tonic::Status> {
        debug!("[GRPC] Storage::update_to_latest_ledger");
        let req = request.into_inner();
        let resp = self
            .update_to_latest_ledger_inner(req)
            .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;
        Ok(tonic::Response::new(resp))
    }

    async fn get_transactions(
        &self,
        request: tonic::Request<GetTransactionsRequest>,
    ) -> Result<tonic::Response<GetTransactionsResponse>, tonic::Status> {
        debug!("[GRPC] Storage::get_transactions");
        let req = request.into_inner();
        let resp = self
            .get_transactions_inner(req)
            .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;
        Ok(tonic::Response::new(resp))
    }

    async fn get_latest_state_root(
        &self,
        request: tonic::Request<GetLatestStateRootRequest>,
    ) -> Result<tonic::Response<GetLatestStateRootResponse>, tonic::Status> {
        debug!("[GRPC] Storage::get_latest_state_root");
        let req = request.into_inner();
        let resp = self
            .get_latest_state_root_inner(req)
            .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;
        Ok(tonic::Response::new(resp))
    }

    async fn get_latest_account_state(
        &self,
        request: tonic::Request<GetLatestAccountStateRequest>,
    ) -> Result<tonic::Response<GetLatestAccountStateResponse>, tonic::Status> {
        debug!("[GRPC] Storage::get_latest_account_state");
        let req = request.into_inner();
        let resp = self
            .get_latest_account_state_inner(req)
            .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;
        Ok(tonic::Response::new(resp))
    }

    async fn get_account_state_with_proof_by_version(
        &self,
        request: tonic::Request<GetAccountStateWithProofByVersionRequest>,
    ) -> Result<tonic::Response<GetAccountStateWithProofByVersionResponse>, tonic::Status> {
        debug!("[GRPC] Storage::get_account_state_with_proof_by_version");
        let req = request.into_inner();
        let resp = self
            .get_account_state_with_proof_by_version_inner(req)
            .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;
        Ok(tonic::Response::new(resp))
    }

    async fn get_startup_info(
        &self,
        _request: tonic::Request<GetStartupInfoRequest>,
    ) -> Result<tonic::Response<GetStartupInfoResponse>, tonic::Status> {
        debug!("[GRPC] Storage::get_startup_info");
        let resp = self
            .get_startup_info_inner()
            .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;
        Ok(tonic::Response::new(resp))
    }

    async fn get_epoch_change_ledger_infos(
        &self,
        request: tonic::Request<GetEpochChangeLedgerInfosRequest>,
    ) -> Result<tonic::Response<ValidatorChangeProof>, tonic::Status> {
        debug!("[GRPC] Storage::get_epoch_change_ledger_infos");
        let req = request.into_inner();
        let resp = self
            .get_epoch_change_ledger_infos_inner(req)
            .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;
        Ok(tonic::Response::new(resp))
    }

    type BackupAccountStateStream =
        mpsc::Receiver<Result<BackupAccountStateResponse, tonic::Status>>;

    async fn backup_account_state(
        &self,
        request: tonic::Request<BackupAccountStateRequest>,
    ) -> Result<tonic::Response<Self::BackupAccountStateStream>, tonic::Status> {
        debug!("[GRPC] Storage::backup_account_state");
        let req = request.into_inner();
        let iter = self
            .db
            .get_account_iter(req.version)
            .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;

        let iter = iter.map(|res| match res {
            Ok((hash, blob)) => {
                let resp: BackupAccountStateResponse =
                    storage_proto::BackupAccountStateResponse::new(hash, blob).into();
                Ok(resp)
            }
            Err(e) => Err(tonic::Status::new(tonic::Code::Internal, e.to_string())),
        });

        // Channel to buffer the stream
        let (mut tx, rx) = mpsc::channel(32);

        tokio::spawn(async move {
            for resp in iter {
                tx.send(resp).await.unwrap();
            }
        });

        Ok(tonic::Response::new(rx))
    }

    async fn get_account_state_range_proof(
        &self,
        request: tonic::Request<GetAccountStateRangeProofRequest>,
    ) -> Result<tonic::Response<GetAccountStateRangeProofResponse>, tonic::Status> {
        debug!("[GRPC] Storage::get_account_state_range_proof");
        let req = request.into_inner();
        let resp = self
            .get_account_state_range_proof_inner(req)
            .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;
        Ok(tonic::Response::new(resp))
    }
}

#[cfg(test)]
mod storage_service_test;
<|endoftext|>


<commit_before>use conduit::{Request, Response};
use diesel::prelude::*;
use time::Timespec;
use serde_json;

use db::RequestTransaction;
use schema::{crate_owner_invitations, users, crates, crate_owners};
use user::RequestUser;
use util::errors::{CargoResult, human};
use util::RequestUtils;
use owner::{CrateOwner, OwnerKind};

/// The model representing a row in the `crate_owner_invitations` database table.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Identifiable, Queryable)]
#[primary_key(invited_user_id, crate_id)]
pub struct CrateOwnerInvitation {
    pub invited_user_id: i32,
    pub invited_by_user_id: i32,
    pub crate_id: i32,
    pub created_at: Timespec,
}

#[derive(Insertable, Clone, Copy, Debug)]
#[table_name = "crate_owner_invitations"]
pub struct NewCrateOwnerInvitation {
    pub invited_user_id: i32,
    pub invited_by_user_id: i32,
    pub crate_id: i32,
}

impl CrateOwnerInvitation {
    pub fn invited_by_username(&self, conn: &PgConnection) -> String {
        users::table
            .find(self.invited_by_user_id)
            .select(users::gh_login)
            .first(&*conn)
            .unwrap_or_else(|_| String::from("(unknown username)"))
    }

    pub fn crate_name(&self, conn: &PgConnection) -> String {
        crates::table
            .find(self.crate_id)
            .select(crates::name)
            .first(&*conn)
            .unwrap_or_else(|_| String::from("(unknown crate name)"))
    }

    pub fn encodable(self, conn: &PgConnection) -> EncodableCrateOwnerInvitation {
        EncodableCrateOwnerInvitation {
            invited_by_username: self.invited_by_username(conn),
            crate_name: self.crate_name(conn),
            crate_id: self.crate_id,
            created_at: ::encode_time(self.created_at),
        }
    }
}

/// The serialization format for the `CrateOwnerInvitation` model.
#[derive(Deserialize, Serialize, Debug)]
pub struct EncodableCrateOwnerInvitation {
    pub invited_by_username: String,
    pub crate_name: String,
    pub crate_id: i32,
    pub created_at: String,
}

/// Handles the `GET /me/crate_owner_invitations` route.
pub fn list(req: &mut Request) -> CargoResult<Response> {
    let conn = &*req.db_conn()?;
    let user_id = req.user()?.id;

    let crate_owner_invitations = crate_owner_invitations::table
        .filter(crate_owner_invitations::invited_user_id.eq(user_id))
        .load::<CrateOwnerInvitation>(&*conn)?
        .into_iter()
        .map(|i| i.encodable(conn))
        .collect();

    #[derive(Serialize)]
    struct R {
        crate_owner_invitations: Vec<EncodableCrateOwnerInvitation>,
    }
    Ok(req.json(&R { crate_owner_invitations }))
}

#[derive(Deserialize)]
struct OwnerInvitation {
    crate_owner_invite: InvitationResponse,
}

#[derive(Deserialize, Serialize, Debug, Copy, Clone)]
pub struct InvitationResponse {
    pub crate_id: i32,
    pub accepted: bool,
}

/// Handles the `PUT /me/crate_owner_invitations/:crate_id` route.
pub fn handle_invite(req: &mut Request) -> CargoResult<Response> {

    let conn = &*req.db_conn()?;

    let mut body = String::new();
    req.body().read_to_string(&mut body)?;

    let crate_invite: OwnerInvitation = serde_json::from_str(&body).map_err(|_| {
        human("invalid json request")
    })?;

    let crate_invite = crate_invite.crate_owner_invite;

    let response = if crate_invite.accepted {
        accept_invite(req, conn, crate_invite)
    } else {
        #[derive(Serialize)]
        struct R {
            crate_owner_invitation: InvitationResponse,
        }
        Ok((req.json(&R { crate_owner_invitation: crate_invite })))
    };

    response
}

fn accept_invite(
    req: &mut Request,
    conn: &PgConnection,
    crate_invite: InvitationResponse,
) -> CargoResult<Response> {
    let user_id = req.user()?.id;
    use diesel::{insert, delete};
    let pending_crate_owner = crate_owner_invitations::table
        .filter(crate_owner_invitations::crate_id.eq(crate_invite.crate_id))
        .filter(crate_owner_invitations::invited_user_id.eq(user_id))
        .first::<CrateOwnerInvitation>(&*conn)?;

    let owner = CrateOwner {
        crate_id: crate_invite.crate_id,
        owner_id: user_id,
        created_by: pending_crate_owner.invited_by_user_id,
        owner_kind: OwnerKind::User as i32,
    };

    conn.transaction(|| {
        insert(&owner).into(crate_owners::table).execute(conn)?;
        delete(
            crate_owner_invitations::table
                .filter(crate_owner_invitations::crate_id.eq(crate_invite.crate_id))
                .filter(crate_owner_invitations::invited_user_id.eq(user_id)),
        ).execute(conn)?;

        #[derive(Serialize)]
        struct R {
            crate_owner_invitation: InvitationResponse,
        }
        Ok(req.json(&R { crate_owner_invitation: crate_invite }))
    })
}
<commit_msg>clippy<commit_after>use conduit::{Request, Response};
use diesel::prelude::*;
use time::Timespec;
use serde_json;

use db::RequestTransaction;
use schema::{crate_owner_invitations, users, crates, crate_owners};
use user::RequestUser;
use util::errors::{CargoResult, human};
use util::RequestUtils;
use owner::{CrateOwner, OwnerKind};

/// The model representing a row in the `crate_owner_invitations` database table.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Identifiable, Queryable)]
#[primary_key(invited_user_id, crate_id)]
pub struct CrateOwnerInvitation {
    pub invited_user_id: i32,
    pub invited_by_user_id: i32,
    pub crate_id: i32,
    pub created_at: Timespec,
}

#[derive(Insertable, Clone, Copy, Debug)]
#[table_name = "crate_owner_invitations"]
pub struct NewCrateOwnerInvitation {
    pub invited_user_id: i32,
    pub invited_by_user_id: i32,
    pub crate_id: i32,
}

impl CrateOwnerInvitation {
    pub fn invited_by_username(&self, conn: &PgConnection) -> String {
        users::table
            .find(self.invited_by_user_id)
            .select(users::gh_login)
            .first(&*conn)
            .unwrap_or_else(|_| String::from("(unknown username)"))
    }

    pub fn crate_name(&self, conn: &PgConnection) -> String {
        crates::table
            .find(self.crate_id)
            .select(crates::name)
            .first(&*conn)
            .unwrap_or_else(|_| String::from("(unknown crate name)"))
    }

    pub fn encodable(self, conn: &PgConnection) -> EncodableCrateOwnerInvitation {
        EncodableCrateOwnerInvitation {
            invited_by_username: self.invited_by_username(conn),
            crate_name: self.crate_name(conn),
            crate_id: self.crate_id,
            created_at: ::encode_time(self.created_at),
        }
    }
}

/// The serialization format for the `CrateOwnerInvitation` model.
#[derive(Deserialize, Serialize, Debug)]
pub struct EncodableCrateOwnerInvitation {
    pub invited_by_username: String,
    pub crate_name: String,
    pub crate_id: i32,
    pub created_at: String,
}

/// Handles the `GET /me/crate_owner_invitations` route.
pub fn list(req: &mut Request) -> CargoResult<Response> {
    let conn = &*req.db_conn()?;
    let user_id = req.user()?.id;

    let crate_owner_invitations = crate_owner_invitations::table
        .filter(crate_owner_invitations::invited_user_id.eq(user_id))
        .load::<CrateOwnerInvitation>(&*conn)?
        .into_iter()
        .map(|i| i.encodable(conn))
        .collect();

    #[derive(Serialize)]
    struct R {
        crate_owner_invitations: Vec<EncodableCrateOwnerInvitation>,
    }
    Ok(req.json(&R { crate_owner_invitations }))
}

#[derive(Deserialize)]
struct OwnerInvitation {
    crate_owner_invite: InvitationResponse,
}

#[derive(Deserialize, Serialize, Debug, Copy, Clone)]
pub struct InvitationResponse {
    pub crate_id: i32,
    pub accepted: bool,
}

/// Handles the `PUT /me/crate_owner_invitations/:crate_id` route.
pub fn handle_invite(req: &mut Request) -> CargoResult<Response> {

    let conn = &*req.db_conn()?;

    let mut body = String::new();
    req.body().read_to_string(&mut body)?;

    let crate_invite: OwnerInvitation = serde_json::from_str(&body).map_err(|_| {
        human("invalid json request")
    })?;

    let crate_invite = crate_invite.crate_owner_invite;

    if crate_invite.accepted {
        accept_invite(req, conn, crate_invite)
    } else {
        #[derive(Serialize)]
        struct R {
            crate_owner_invitation: InvitationResponse,
        }
        Ok(req.json(&R { crate_owner_invitation: crate_invite }))
    }
}

fn accept_invite(
    req: &mut Request,
    conn: &PgConnection,
    crate_invite: InvitationResponse,
) -> CargoResult<Response> {
    let user_id = req.user()?.id;
    use diesel::{insert, delete};
    let pending_crate_owner = crate_owner_invitations::table
        .filter(crate_owner_invitations::crate_id.eq(crate_invite.crate_id))
        .filter(crate_owner_invitations::invited_user_id.eq(user_id))
        .first::<CrateOwnerInvitation>(&*conn)?;

    let owner = CrateOwner {
        crate_id: crate_invite.crate_id,
        owner_id: user_id,
        created_by: pending_crate_owner.invited_by_user_id,
        owner_kind: OwnerKind::User as i32,
    };

    conn.transaction(|| {
        insert(&owner).into(crate_owners::table).execute(conn)?;
        delete(
            crate_owner_invitations::table
                .filter(crate_owner_invitations::crate_id.eq(crate_invite.crate_id))
                .filter(crate_owner_invitations::invited_user_id.eq(user_id)),
        ).execute(conn)?;

        #[derive(Serialize)]
        struct R {
            crate_owner_invitation: InvitationResponse,
        }
        Ok(req.json(&R { crate_owner_invitation: crate_invite }))
    })
}
<|endoftext|>


<commit_before>/*
* Huffman table generation:
* HuffTable::huff_decode()
* HuffTable::initialize()
* and used data structures are originally grabbed from the IJG software,
* and adapted by Hubert Figuiere.
*
* Copyright (C) 1991, 1992, Thomas G. Lane.
* Part of the Independent JPEG Group's software.
* See the file Copyright for more details.
*
* Copyright (c) 1993 Brian C. Smith, The Regents of the University
* of California
* All rights reserved.
*
* Copyright (c) 1994 Kongji Huang and Brian C. Smith.
* Cornell University
* All rights reserved.
*
* Permission to use, copy, modify, and distribute this software and its
* documentation for any purpose, without fee, and without written agreement is
* hereby granted, provided that the above copyright notice and the following
* two paragraphs appear in all copies of this software.
*
* IN NO EVENT SHALL CORNELL UNIVERSITY BE LIABLE TO ANY PARTY FOR
* DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT
* OF THE USE OF THIS SOFTWARE AND ITS DOCUMENTATION, EVEN IF CORNELL
* UNIVERSITY HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
* CORNELL UNIVERSITY SPECIFICALLY DISCLAIMS ANY WARRANTIES,
* INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY
* AND FITNESS FOR A PARTICULAR PURPOSE.  THE SOFTWARE PROVIDED HEREUNDER IS
* ON AN "AS IS" BASIS, AND CORNELL UNIVERSITY HAS NO OBLIGATION TO
* PROVIDE MAINTENANCE, SUPPORT, UPDATES, ENHANCEMENTS, OR MODIFICATIONS.
*/

use std::fmt;
use decoders::basics::*;

const HUFF_BITMASK: [u32;32] = [0xffffffff, 0x7fffffff,
                                0x3fffffff, 0x1fffffff,
                                0x0fffffff, 0x07ffffff,
                                0x03ffffff, 0x01ffffff,
                                0x00ffffff, 0x007fffff,
                                0x003fffff, 0x001fffff,
                                0x000fffff, 0x0007ffff,
                                0x0003ffff, 0x0001ffff,
                                0x0000ffff, 0x00007fff,
                                0x00003fff, 0x00001fff,
                                0x00000fff, 0x000007ff,
                                0x000003ff, 0x000001ff,
                                0x000000ff, 0x0000007f,
                                0x0000003f, 0x0000001f,
                                0x0000000f, 0x00000007,
                                0x00000003, 0x00000001];

pub struct HuffTable {
  // These two fields directly represent the contents of a JPEG DHT marker
  pub bits: [u32;17],
  pub huffval: [u32;256],

  // The remaining fields are computed from the above to allow more
  // efficient coding and decoding and thus private
  mincode: [u16;17],
  maxcode: [i32;18],
  valptr: [i16;17],
  numbits: [usize;256],
  bigtable: Vec<i32>,
  precision: usize,
  pub use_bigtable: bool,
  pub dng_compatible: bool,
  pub initialized: bool,
}

impl HuffTable {
  pub fn empty(precision: usize) -> HuffTable {
    HuffTable {
      bits: [0;17],
      huffval: [0;256],
      mincode: [0;17],
      maxcode: [0;18],
      valptr: [0;17],
      numbits: [0;256],
      bigtable: Vec::new(),
      precision: precision,
      use_bigtable: true,
      dng_compatible: true,
      initialized: false,
    }
  }

  pub fn initialize(&mut self, use_bigtable: bool) -> Result<(), String> {
    // Figure C.1: make table of Huffman code length for each symbol
    // Note that this is in code-length order.
    let mut p = 0;
    let mut huffsize: [u8;257] = [0;257];
    for l in 1..17 {
      for _ in 1..((self.bits[l] as usize)+1) {
        huffsize[p] = l as u8;
        p += 1;
        if p > 256 {
          return Err("ljpeg: Code length too long. Corrupt data.".to_string())
        }
      }
    }
    huffsize[p] = 0;
    let lastp = p;

    // Figure C.2: generate the codes themselves
    // Note that this is in code-length order.
    let mut code: u16 = 0;
    let mut si: u32 = huffsize[0] as u32;
    let mut huffcode: [u16;257] = [0;257];
    p = 0;
    while huffsize[p] > 0 {
      while (huffsize[p] as u32) == si {
        huffcode[p] = code;
        p += 1;
        code += 1;
      }
      code <<= 1;
      si += 1;
      if p > 256 {
        return Err("ljpeg: Code length too long. Corrupt data.".to_string())
      }
    }

    //Figure F.15: generate decoding tables
    self.mincode[0] = 0;
    self.maxcode[0] = 0;
    p = 0;
    for l in 1..17 {
      if self.bits[l] > 0 {
        self.valptr[l] = p as i16;
        self.mincode[l] = huffcode[p];
        p += self.bits[l] as usize;
        self.maxcode[l] = huffcode[p - 1] as i32;
      } else {
        self.valptr[l] = 0xff;   // This check must be present to avoid crash on junk
        self.maxcode[l] = -1;
      }
      if p > 256 {
        return Err("ljpeg: Code length too long. Corrupt data.".to_string())
      }
    }

    // We put in this value to ensure HuffDecode terminates
    self.maxcode[17] = 0xFFFFF;

    // Build the numbits, value lookup tables.
    // These table allow us to gather 8 bits from the bits stream,
    // and immediately lookup the size and value of the huffman codes.
    // If size is zero, it means that more than 8 bits are in the huffman
    // code (this happens about 3-4% of the time).
    for p in 0..lastp {
      let size = huffsize[p];
      if size <= 8 {
        let value: i32 = self.huffval[p] as i32;
        let code = huffcode[p];
        let ll: i32 = (code << (8 - size)) as i32;
        let ul: i32 = if size < 8 {
          ll | (HUFF_BITMASK[(24+size) as usize]) as i32
        } else {
          ll
        };
        if ul > 256 || ll > ul {
          return Err("ljpeg: Code length too long. Corrupt data.".to_string())
        }
        for i in ll..(ul+1) {
          self.numbits[i as usize] = (size as usize) | ((value as usize) << 4);
        }
      }
    }

    if use_bigtable {
      try!(self.initialize_bigtable());
    }
    self.initialized = true;
    self.use_bigtable = use_bigtable;
    Ok(())
  }

  fn initialize_bigtable(&mut self) -> Result<(), String> {
    let bits: usize = 14; // HuffDecode functions must be changed, if this is modified.
    let size: usize = 1 << bits;

    self.bigtable = vec![0 as i32;size];
    let mut rv: i32;
    for i in 0..size {
      let input = (i << 2) as u16;
      let mut code: i32 = (input >> 8) as i32;
      let val: u32 = self.numbits[code as usize] as u32;
      let mut l: u32 = val & 15;
      if l > 0 {
        rv = (val >> 4) as i32;
      } else {
        l = 8;
        while code > self.maxcode[l as usize] && l < 16 {
          let temp: i32 = (input >> (15 - l) & 1) as i32;
          code = (code << 1) | temp;
          l += 1;
        }

        //With garbage input we may reach the sentinel value l = 17.
        if l > self.precision as u32 || self.valptr[l as usize] == 0xff {
          self.bigtable[i] = 0xff;
          continue
        } else {
          rv = self.huffval[
            self.valptr[l as usize] as usize +
            (code - (self.mincode[l as usize] as i32)) as usize
          ] as i32;
        }
      }

      if rv == 16 {
        self.bigtable[i] = if self.dng_compatible {
          (-(32768 << 8)) | (16 + l as i32)
        } else {
          (-(32768 << 8)) | l as i32
        };
        continue
      }

      if rv + l as i32 > bits as i32 {
        self.bigtable[i] = 0xff;
        continue
      }

      if rv != 0 {
        let mut x = (input >> (16 - (l as i32) - rv) & ((1 << rv) - 1)) as i32;
        if (x & (1 << (rv - 1))) == 0 {
          x -= (1 << rv) - 1;
        }
        self.bigtable[i] = (x << 8) | ((l as i32) + rv);
      } else {
        self.bigtable[i] = l as i32;
      }
    }
    self.use_bigtable = true;
    Ok(())
  }

  // Taken from Figure F.16: extract next coded symbol from input stream
  pub fn huff_decode(&self, pump: &mut BitPump) -> Result<i32,String> {
    if !self.initialized {
      return Err("ljpeg: Trying to read from uninitialized HuffTable".to_string());
    }

    //First attempt to do complete decode, by using the first 14 bits
    let code: i32 = pump.peek_bits(14) as i32;
    if self.use_bigtable {
      let val: i32 = self.bigtable[code as usize];
      if val & 0xff != 0xff {
        pump.consume_bits((val & 0xff) as u32);
        return Ok(val >> 8)
      }
    }

    let len = try!(self.huff_len(pump));
    Ok(self.huff_diff(pump, len))
  }

  pub fn huff_len(&self, pump: &mut BitPump) -> Result<u32,String> {
    let mut code = pump.peek_bits(8) as usize;
    let val = self.numbits[code as usize] as u32;
    let mut l = val & 15;
    if l != 0 {
      pump.consume_bits(l);
      return Ok(val >> 4)
    }
    pump.consume_bits(8);
    l = 8;
    while code as i32 > self.maxcode[l as usize] {
      let temp = pump.get_bits(1) as usize;
      code = (code << 1) | temp;
      l += 1;
    }

    // With garbage input we may reach the sentinel value l = 17.
    if l > self.precision as u32 || self.valptr[l as usize] == 0xff {
      return Err(format!("ljpeg: bad Huffman code: {}", l).to_string())
    } else {
      return Ok(self.huffval[
        self.valptr[l as usize] as usize +
        (code - (self.mincode[l as usize] as usize)) as usize
      ] as u32);
    }
  }

  pub fn huff_diff(&self, pump: &mut BitPump, len: u32) -> i32 {
    match len {
      0 => 0,
      16 => {
        if self.dng_compatible {
          pump.get_bits(16); // consume can fail because we haven't peeked yet
        }
        -32768
      },
      len => {
        // Section F.2.2.1: decode the difference and
        // Figure F.12: extend sign bit
        let mut x: i32 = pump.get_bits(len as u32) as i32;
        if (x & (1 << (len - 1))) == 0 {
          x -= (1 << len) - 1;
        }
        x
      },
    }
  }
}

impl fmt::Debug for HuffTable {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    if self.initialized {
      write!(f, "HuffTable {{ bits: {:?} huffval: {:?} }}", self.bits, &self.huffval[..])
    } else {
      write!(f, "HuffTable {{ uninitialized }}")
    }
  }
}

<commit_msg>ljpeg: don't test initialization on every decode<commit_after>/*
* Huffman table generation:
* HuffTable::huff_decode()
* HuffTable::initialize()
* and used data structures are originally grabbed from the IJG software,
* and adapted by Hubert Figuiere.
*
* Copyright (C) 1991, 1992, Thomas G. Lane.
* Part of the Independent JPEG Group's software.
* See the file Copyright for more details.
*
* Copyright (c) 1993 Brian C. Smith, The Regents of the University
* of California
* All rights reserved.
*
* Copyright (c) 1994 Kongji Huang and Brian C. Smith.
* Cornell University
* All rights reserved.
*
* Permission to use, copy, modify, and distribute this software and its
* documentation for any purpose, without fee, and without written agreement is
* hereby granted, provided that the above copyright notice and the following
* two paragraphs appear in all copies of this software.
*
* IN NO EVENT SHALL CORNELL UNIVERSITY BE LIABLE TO ANY PARTY FOR
* DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT
* OF THE USE OF THIS SOFTWARE AND ITS DOCUMENTATION, EVEN IF CORNELL
* UNIVERSITY HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
* CORNELL UNIVERSITY SPECIFICALLY DISCLAIMS ANY WARRANTIES,
* INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY
* AND FITNESS FOR A PARTICULAR PURPOSE.  THE SOFTWARE PROVIDED HEREUNDER IS
* ON AN "AS IS" BASIS, AND CORNELL UNIVERSITY HAS NO OBLIGATION TO
* PROVIDE MAINTENANCE, SUPPORT, UPDATES, ENHANCEMENTS, OR MODIFICATIONS.
*/

use std::fmt;
use decoders::basics::*;

const HUFF_BITMASK: [u32;32] = [0xffffffff, 0x7fffffff,
                                0x3fffffff, 0x1fffffff,
                                0x0fffffff, 0x07ffffff,
                                0x03ffffff, 0x01ffffff,
                                0x00ffffff, 0x007fffff,
                                0x003fffff, 0x001fffff,
                                0x000fffff, 0x0007ffff,
                                0x0003ffff, 0x0001ffff,
                                0x0000ffff, 0x00007fff,
                                0x00003fff, 0x00001fff,
                                0x00000fff, 0x000007ff,
                                0x000003ff, 0x000001ff,
                                0x000000ff, 0x0000007f,
                                0x0000003f, 0x0000001f,
                                0x0000000f, 0x00000007,
                                0x00000003, 0x00000001];

pub struct HuffTable {
  // These two fields directly represent the contents of a JPEG DHT marker
  pub bits: [u32;17],
  pub huffval: [u32;256],

  // The remaining fields are computed from the above to allow more
  // efficient coding and decoding and thus private
  mincode: [u16;17],
  maxcode: [i32;18],
  valptr: [i16;17],
  numbits: [usize;256],
  bigtable: Vec<i32>,
  precision: usize,
  pub use_bigtable: bool,
  pub dng_compatible: bool,
  pub initialized: bool,
}

impl HuffTable {
  pub fn empty(precision: usize) -> HuffTable {
    HuffTable {
      bits: [0;17],
      huffval: [0;256],
      mincode: [0;17],
      maxcode: [0;18],
      valptr: [0;17],
      numbits: [0;256],
      bigtable: Vec::new(),
      precision: precision,
      use_bigtable: true,
      dng_compatible: true,
      initialized: false,
    }
  }

  pub fn initialize(&mut self, use_bigtable: bool) -> Result<(), String> {
    // Figure C.1: make table of Huffman code length for each symbol
    // Note that this is in code-length order.
    let mut p = 0;
    let mut huffsize: [u8;257] = [0;257];
    for l in 1..17 {
      for _ in 1..((self.bits[l] as usize)+1) {
        huffsize[p] = l as u8;
        p += 1;
        if p > 256 {
          return Err("ljpeg: Code length too long. Corrupt data.".to_string())
        }
      }
    }
    huffsize[p] = 0;
    let lastp = p;

    // Figure C.2: generate the codes themselves
    // Note that this is in code-length order.
    let mut code: u16 = 0;
    let mut si: u32 = huffsize[0] as u32;
    let mut huffcode: [u16;257] = [0;257];
    p = 0;
    while huffsize[p] > 0 {
      while (huffsize[p] as u32) == si {
        huffcode[p] = code;
        p += 1;
        code += 1;
      }
      code <<= 1;
      si += 1;
      if p > 256 {
        return Err("ljpeg: Code length too long. Corrupt data.".to_string())
      }
    }

    //Figure F.15: generate decoding tables
    self.mincode[0] = 0;
    self.maxcode[0] = 0;
    p = 0;
    for l in 1..17 {
      if self.bits[l] > 0 {
        self.valptr[l] = p as i16;
        self.mincode[l] = huffcode[p];
        p += self.bits[l] as usize;
        self.maxcode[l] = huffcode[p - 1] as i32;
      } else {
        self.valptr[l] = 0xff;   // This check must be present to avoid crash on junk
        self.maxcode[l] = -1;
      }
      if p > 256 {
        return Err("ljpeg: Code length too long. Corrupt data.".to_string())
      }
    }

    // We put in this value to ensure HuffDecode terminates
    self.maxcode[17] = 0xFFFFF;

    // Build the numbits, value lookup tables.
    // These table allow us to gather 8 bits from the bits stream,
    // and immediately lookup the size and value of the huffman codes.
    // If size is zero, it means that more than 8 bits are in the huffman
    // code (this happens about 3-4% of the time).
    for p in 0..lastp {
      let size = huffsize[p];
      if size <= 8 {
        let value: i32 = self.huffval[p] as i32;
        let code = huffcode[p];
        let ll: i32 = (code << (8 - size)) as i32;
        let ul: i32 = if size < 8 {
          ll | (HUFF_BITMASK[(24+size) as usize]) as i32
        } else {
          ll
        };
        if ul > 256 || ll > ul {
          return Err("ljpeg: Code length too long. Corrupt data.".to_string())
        }
        for i in ll..(ul+1) {
          self.numbits[i as usize] = (size as usize) | ((value as usize) << 4);
        }
      }
    }

    if use_bigtable {
      try!(self.initialize_bigtable());
    }
    self.initialized = true;
    self.use_bigtable = use_bigtable;
    Ok(())
  }

  fn initialize_bigtable(&mut self) -> Result<(), String> {
    let bits: usize = 14; // HuffDecode functions must be changed, if this is modified.
    let size: usize = 1 << bits;

    self.bigtable = vec![0 as i32;size];
    let mut rv: i32;
    for i in 0..size {
      let input = (i << 2) as u16;
      let mut code: i32 = (input >> 8) as i32;
      let val: u32 = self.numbits[code as usize] as u32;
      let mut l: u32 = val & 15;
      if l > 0 {
        rv = (val >> 4) as i32;
      } else {
        l = 8;
        while code > self.maxcode[l as usize] && l < 16 {
          let temp: i32 = (input >> (15 - l) & 1) as i32;
          code = (code << 1) | temp;
          l += 1;
        }

        //With garbage input we may reach the sentinel value l = 17.
        if l > self.precision as u32 || self.valptr[l as usize] == 0xff {
          self.bigtable[i] = 0xff;
          continue
        } else {
          rv = self.huffval[
            self.valptr[l as usize] as usize +
            (code - (self.mincode[l as usize] as i32)) as usize
          ] as i32;
        }
      }

      if rv == 16 {
        self.bigtable[i] = if self.dng_compatible {
          (-(32768 << 8)) | (16 + l as i32)
        } else {
          (-(32768 << 8)) | l as i32
        };
        continue
      }

      if rv + l as i32 > bits as i32 {
        self.bigtable[i] = 0xff;
        continue
      }

      if rv != 0 {
        let mut x = (input >> (16 - (l as i32) - rv) & ((1 << rv) - 1)) as i32;
        if (x & (1 << (rv - 1))) == 0 {
          x -= (1 << rv) - 1;
        }
        self.bigtable[i] = (x << 8) | ((l as i32) + rv);
      } else {
        self.bigtable[i] = l as i32;
      }
    }
    self.use_bigtable = true;
    Ok(())
  }

  // Taken from Figure F.16: extract next coded symbol from input stream
  pub fn huff_decode(&self, pump: &mut BitPump) -> Result<i32,String> {
    //First attempt to do complete decode, by using the first 14 bits
    let code: i32 = pump.peek_bits(14) as i32;
    if self.use_bigtable {
      let val: i32 = self.bigtable[code as usize];
      if val & 0xff != 0xff {
        pump.consume_bits((val & 0xff) as u32);
        return Ok(val >> 8)
      }
    }

    let len = try!(self.huff_len(pump));
    Ok(self.huff_diff(pump, len))
  }

  pub fn huff_len(&self, pump: &mut BitPump) -> Result<u32,String> {
    let mut code = pump.peek_bits(8) as usize;
    let val = self.numbits[code as usize] as u32;
    let mut l = val & 15;
    if l != 0 {
      pump.consume_bits(l);
      return Ok(val >> 4)
    }
    pump.consume_bits(8);
    l = 8;
    while code as i32 > self.maxcode[l as usize] {
      let temp = pump.get_bits(1) as usize;
      code = (code << 1) | temp;
      l += 1;
    }

    // With garbage input we may reach the sentinel value l = 17.
    if l > self.precision as u32 || self.valptr[l as usize] == 0xff {
      return Err(format!("ljpeg: bad Huffman code: {}", l).to_string())
    } else {
      return Ok(self.huffval[
        self.valptr[l as usize] as usize +
        (code - (self.mincode[l as usize] as usize)) as usize
      ] as u32);
    }
  }

  pub fn huff_diff(&self, pump: &mut BitPump, len: u32) -> i32 {
    match len {
      0 => 0,
      16 => {
        if self.dng_compatible {
          pump.get_bits(16); // consume can fail because we haven't peeked yet
        }
        -32768
      },
      len => {
        // Section F.2.2.1: decode the difference and
        // Figure F.12: extend sign bit
        let mut x: i32 = pump.get_bits(len as u32) as i32;
        if (x & (1 << (len - 1))) == 0 {
          x -= (1 << len) - 1;
        }
        x
      },
    }
  }
}

impl fmt::Debug for HuffTable {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    if self.initialized {
      write!(f, "HuffTable {{ bits: {:?} huffval: {:?} }}", self.bits, &self.huffval[..])
    } else {
      write!(f, "HuffTable {{ uninitialized }}")
    }
  }
}

<|endoftext|>


<commit_before>use std::ptr;

struct TssContext { handle: u32 }
struct TssTpm<'context> { context: &'context TssContext, handle: u32 }

type TssHContext = u32;
type TssHTPM = u32;
type TssResult = u32;
type TssUnicode = u16;

const TSS_SUCCESS: TssResult = 0;

#[link(name = "tspi")]
extern {
    fn Tspi_Context_Create(phContext: *mut TssHContext) -> TssResult;
    fn Tspi_Context_Connect(phContext: TssHContext, wszDestination: *const TssUnicode) -> TssResult;
    fn Tspi_Context_GetTpmObject(phContext: TssHContext, phTPM: *mut TssHTPM) -> TssResult;
    fn Tspi_Context_FreeMemory(phContext: TssHContext, rgbMemory: *mut u8) -> TssResult;
    fn Tspi_Context_Close(phContext: TssHContext) -> TssResult;
    fn Tspi_TPM_PcrRead(hTPM: TssHTPM, ulPcrIndex: u32, pulPcrValueLength: *mut u32, prgbPcrValue: *mut *mut u8) -> TssResult;
}

impl TssContext {
    fn new() -> Result<TssContext, TssResult> {
        unsafe {
            let mut handle = 0;
            let result = Tspi_Context_Create(&mut handle);
            if result != TSS_SUCCESS {
                Err(result)
            } else {
                Ok(TssContext { handle: handle })
            }
        }
    }

    // TODO: support destination
    fn connect(&self) -> Result<(), TssResult> {
        unsafe {
            let result = Tspi_Context_Connect(self.handle, ptr::null());
            if result != TSS_SUCCESS {
                Err(result)
            } else {
                Ok(())
            }
        }
    }

    fn get_tpm_object(&self) -> Result<TssTpm, TssResult> {
        unsafe {
            let mut handle = 0;
            let result = Tspi_Context_GetTpmObject(self.handle, &mut handle);
            if result != TSS_SUCCESS {
                Err(result)
            } else {
                Ok(TssTpm { context: self, handle: handle })
            }
        }
    }
}

impl Drop for TssContext {
    fn drop(&mut self) {
        unsafe {
            Tspi_Context_FreeMemory(self.handle, 0 as *mut u8);
            Tspi_Context_Close(self.handle);
        }
    }
}

impl<'context> TssTpm<'context> {
    fn pcr_read(&self, pcr_index: u32) -> Result<Vec<u8>, TssResult> {
        unsafe {
            let mut ulPcrValueLength = -1;
            let mut pRgbPcrValue = 0 as *mut u8;
            let result = Tspi_TPM_PcrRead(self.handle, pcr_index, &mut ulPcrValueLength, &mut pRgbPcrValue);
            if result != TSS_SUCCESS {
                Err(result)
            } else {
                let mut vec = Vec::new();
                for i in 0..ulPcrValueLength {
                    vec.push(*pRgbPcrValue.offset(i as isize));
                }
                Tspi_Context_FreeMemory(self.context.handle, pRgbPcrValue);
                Ok(vec)
            }
        }
    }
}

fn main() {
    // TODO: Any cleaner way to write this?
    let contextresult = TssContext::new();
    if let Ok(context) = contextresult {
        if let Ok(_) = context.connect() {
            if let Ok(tpm) = context.get_tpm_object() {
                println!("I'M IN UR TPM READING UR PCRZ (From Rust!)");
                println!("TPM handle: {:?}", tpm.handle);
                for i in 0..24 {
                    if let Ok(vec) = tpm.pcr_read(i) {
                        print!("PCR {:02}", i);
                        for j in 0..19 {
                            print!(" {:02x}", vec[j]);
                        }
                        print!("\n");
                    }
                }
            } else {
                println!("Failed to get TPM handle :(")
            }
        } else {
            println!("Failed to connect :(");
        }
    } else {
        println!("Failed :(");
    }

    /*
    match blah {
        Ok(context) => println!("Context created! {:?}", context.handle),
        Err(e) => println!("Context failed with err {:?}", e),
    }
    println!("Hello world!");
    */
}
<commit_msg>Move code out of unsafe blocks if possible<commit_after>use std::ptr;

struct TssContext { handle: u32 }
struct TssTpm<'context> { context: &'context TssContext, handle: u32 }

type TssHContext = u32;
type TssHTPM = u32;
type TssResult = u32;
type TssUnicode = u16;

const TSS_SUCCESS: TssResult = 0;

#[link(name = "tspi")]
extern {
    fn Tspi_Context_Create(phContext: *mut TssHContext) -> TssResult;
    fn Tspi_Context_Connect(phContext: TssHContext, wszDestination: *const TssUnicode) -> TssResult;
    fn Tspi_Context_GetTpmObject(phContext: TssHContext, phTPM: *mut TssHTPM) -> TssResult;
    fn Tspi_Context_FreeMemory(phContext: TssHContext, rgbMemory: *mut u8) -> TssResult;
    fn Tspi_Context_Close(phContext: TssHContext) -> TssResult;
    fn Tspi_TPM_PcrRead(hTPM: TssHTPM, ulPcrIndex: u32, pulPcrValueLength: *mut u32, prgbPcrValue: *mut *mut u8) -> TssResult;
}

impl TssContext {
    fn new() -> Result<TssContext, TssResult> {
        let mut handle = 0;
        let result = unsafe {
            Tspi_Context_Create(&mut handle)
        };
        if result != TSS_SUCCESS {
            Err(result)
        } else {
            Ok(TssContext { handle: handle })
        }
    }

    // TODO: support destination
    fn connect(&self) -> Result<(), TssResult> {
        let result = unsafe {
            Tspi_Context_Connect(self.handle, ptr::null())
        };
        if result != TSS_SUCCESS {
            Err(result)
        } else {
            Ok(())
        }
    }

    fn get_tpm_object(&self) -> Result<TssTpm, TssResult> {
        let mut handle = 0;
        let result = unsafe {
            Tspi_Context_GetTpmObject(self.handle, &mut handle)
        };
        if result != TSS_SUCCESS {
            Err(result)
        } else {
            Ok(TssTpm { context: self, handle: handle })
        }
    }
}

impl Drop for TssContext {
    fn drop(&mut self) {
        unsafe {
            Tspi_Context_FreeMemory(self.handle, 0 as *mut u8);
            Tspi_Context_Close(self.handle);
        }
    }
}

impl<'context> TssTpm<'context> {
    fn pcr_read(&self, pcr_index: u32) -> Result<Vec<u8>, TssResult> {
        let mut ulPcrValueLength = -1;
        let mut pRgbPcrValue = 0 as *mut u8;
        let result = unsafe {
            Tspi_TPM_PcrRead(self.handle, pcr_index, &mut ulPcrValueLength, &mut pRgbPcrValue)
        };
        if result != TSS_SUCCESS {
            Err(result)
        } else {
            let mut vec = Vec::new();
            unsafe {
                for i in 0..ulPcrValueLength {
                    vec.push(*pRgbPcrValue.offset(i as isize));
                }
                Tspi_Context_FreeMemory(self.context.handle, pRgbPcrValue);
            }
            Ok(vec)
        }
    }
}

fn main() {
    // TODO: Any cleaner way to write this?
    let contextresult = TssContext::new();
    if let Ok(context) = contextresult {
        if let Ok(_) = context.connect() {
            if let Ok(tpm) = context.get_tpm_object() {
                println!("I'M IN UR TPM READING UR PCRZ (From Rust!)");
                println!("TPM handle: {:?}", tpm.handle);
                for i in 0..24 {
                    if let Ok(vec) = tpm.pcr_read(i) {
                        print!("PCR {:02}", i);
                        for j in 0..19 {
                            print!(" {:02x}", vec[j]);
                        }
                        print!("\n");
                    }
                }
            } else {
                println!("Failed to get TPM handle :(")
            }
        } else {
            println!("Failed to connect :(");
        }
    } else {
        println!("Failed :(");
    }

    /*
    match blah {
        Ok(context) => println!("Context created! {:?}", context.handle),
        Err(e) => println!("Context failed with err {:?}", e),
    }
    println!("Hello world!");
    */
}
<|endoftext|>