<commit_before>// base💯 - Copyright 2017 Adam Niederer

// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

#![allow(non_upper_case_globals)]
#![cfg_attr(test, feature(test))]
#![cfg_attr(feature = "simd", feature(target_feature))]
#![cfg_attr(feature = "simd", feature(asm))]
#![cfg_attr(feature = "simd", feature(cfg_target_feature))]

#[macro_use] extern crate clap;
#[cfg(feature = "simd")] extern crate stdsimd;

use std::io::{self, Read, Write, BufRead, BufReader, BufWriter};
use std::fs::File;
use std::iter::Iterator;
use clap::App;

fn main() {
    let cli_spec = load_yaml!("cli.yml");
    let cli_args = App::from_yaml(cli_spec).get_matches();

    let mut reader = {
        if let Some(path) = cli_args.value_of("input") {
            Box::new(BufReader::new(match File::open(path) {
                Ok(path) => path,
                _ => {
                    writeln!(io::stderr(), "base💯: no such file: {}", path).expect("base💯: stderr write error");
                    return;
                }
            })) as Box<BufRead>
        } else {
            Box::new(BufReader::new(io::stdin())) as Box<BufRead>
        }
    };

    let mut writer = BufWriter::new(io::stdout());

    if cli_args.is_present("decode") {
        let mut write_buf = [0u8; 0x10000];
        let mut buffer = [0u8; 0x10000];

        while let Ok(num_read) = reader.read(&mut buffer) {
            if num_read == 0 {
                break;
            }

            match writer.write_all(&emoji_to_char(&buffer[..num_read], &mut write_buf)[..num_read / 4]) {
                Ok(_) => (),
                _ => {
                    writeln!(io::stderr(), "base💯: write error").expect("base💯: stderr write error");
                    return;
                }
            }
        }
    } else {
        let mut write_buf = [0u8; 0x40000];
        let mut buffer = [0u8; 0x10000];

        while let Ok(num_read) = reader.read(&mut buffer) {
            if num_read == 0 {
                break;
            }

            match writer.write_all(&char_to_emoji(&buffer[..num_read], &mut write_buf)[..num_read * 4]) {
                Ok(_) => (),
                _ => {
                    writeln!(io::stderr(), "base💯: write error").expect("base💯: stderr write error");
                    return;
                }
            }
        }
    }
    writer.flush().expect("Write error");
}

#[cfg(any(not(feature = "simd"), not(target_arch = "x86_64")))]
fn emoji_to_char<'a, 'b>(buf: &'a[u8], out: &'b mut [u8]) -> &'b[u8] {
    for (i, chunk) in buf.chunks(4).enumerate() {
        out[i] = ((chunk[2].wrapping_sub(143)).wrapping_mul(64)).wrapping_add(chunk[3].wrapping_sub(128)).wrapping_sub(55)
    }
    out
}

#[cfg(feature = "simd")]
#[cfg(target_feature = "avx2")]
pub fn emoji_to_char<'a, 'b>(buf: &'a[u8], out: &'b mut [u8]) -> &'b[u8] {
    use stdsimd::simd::u8x32;
    let mut i = 0;
    for chunk in buf.chunks(128) {
        if chunk.len() < 128 {
            // Non-SIMD implementation for the final <128 bytes
            for chunk in chunk.chunks(4) {
                out[i] = ((chunk[2].wrapping_sub(143)).wrapping_mul(64)).wrapping_add(chunk[3].wrapping_sub(128)).wrapping_sub(55);
                i += 1;
            }
        } else {
            // AVX implementation of decoding algo
            // a, b, c, d contain one garbage word and one useful word
            let a = u8x32::load(chunk, 0);
            let b = u8x32::load(chunk, 32);
            let c = u8x32::load(chunk, 64);
            let d = u8x32::load(chunk, 96);

            // Constant mask for removing low bytes
            let hi_mask = u8x32::new(0xFF, 0x00, 0xFF, 0x00,0xFF, 0x00, 0xFF, 0x00,
                                     0xFF, 0x00, 0xFF, 0x00,0xFF, 0x00, 0xFF, 0x00,
                                     0xFF, 0x00, 0xFF, 0x00,0xFF, 0x00, 0xFF, 0x00,
                                     0xFF, 0x00, 0xFF, 0x00,0xFF, 0x00, 0xFF, 0x00);

            // Results from the vector magic below
            let hi: u8x32;
            let lo: u8x32;

            unsafe {
                asm! {
                    "vpsrld ymm1, ymm1, 16
                     vpsrld ymm2, ymm2, 16
                     vpsrld ymm3, ymm3, 16
                     vpsrld ymm4, ymm4, 16
                     // ymm1 .. ymm4 now contain 0x00, 0x00, 0xhi, 0xlo ...

                     // Pack ymm* into ymm1 and ymm2 and order them properly
                     vpackusdw ymm1, ymm1, ymm2
                     vpackusdw ymm2, ymm3, ymm4
                     // vpackusdw interleaves ymm[13] and ymm[24; we want them end-to-end
                     vpermpd ymm1, ymm1, 0xD8 // [191:128] <-> [127:64]
                     vpermpd ymm2, ymm2, 0xD8 // [191:128] <-> [127:64]
                     // ymm1 and ymm2 now contain interleaved high and low bytes

                     // Store high bytes (0xhi 0x00) in ymm3 and ymm4
                     vpand ymm3, ymm1, ymm6
                     vpand ymm4, ymm2, ymm6
                     // Pack hi bytes into ymm4
                     vpackuswb ymm4, ymm3, ymm4

                     // Remove the high bytes and move low bytes into position
                     vpsrlw ymm1, ymm1, 8
                     vpsrlw ymm2, ymm2, 8
                     // Store low bytes in ymm3
                     vpackuswb ymm3, ymm1, ymm2

                     // Fight interleaving again
                     vpermpd ymm4, ymm4, 0xD8 // [191:128] <-> [127:64]
                     vpermpd ymm3, ymm3, 0xD8 // [191:128] <-> [127:64]"
                     : "={ymm3}"(lo), "={ymm4}"(hi)
                     : "{ymm1}"(a), "{ymm2}"(b), "{ymm3}"(c),
                       "{ymm4}"(d), "{ymm6}"(hi_mask)
                     : "cc"
                     : "intel"
                }
            }

            ((((hi - u8x32::splat(143)) * u8x32::splat(64))
              + lo - u8x32::splat(128)) - u8x32::splat(55))
                .store(out, i);
            i += 32;
        }
    }
    out
}

#[cfg(feature = "simd")]
#[cfg(all(not(target_feature = "avx2"), target_feature = "sse2"))]
pub fn emoji_to_char<'a, 'b>(buf: &'a[u8], out: &'b mut [u8]) -> &'b[u8] {
    use stdsimd::simd::u8x16;
    let mut i = 0;
    for chunk in buf.chunks(64) {
        // Chunk data format: 0x00 0x00 0xhi 0xlo
        if chunk.len() < 64 {
            // Non-SIMD implementation for the final <128 bytes
            for (i, chunk) in chunk.chunks(4).enumerate() {
                out[i] = ((chunk[2].wrapping_sub(143)).wrapping_mul(64)).wrapping_add(chunk[3].wrapping_sub(128)).wrapping_sub(55)
            }
        } else {

            // SSE2 implementation of decoding algo

            // TODO: Use Get rid of these scalar loads
            let msbs = u8x16::new(chunk[2], chunk[6], chunk[10], chunk[14], chunk[18], chunk[22], chunk[26], chunk[30], chunk[34], chunk[38], chunk[42], chunk[46], chunk[50], chunk[54], chunk[58], chunk[62]);
            let lsbs = u8x16::new(chunk[3], chunk[7], chunk[11], chunk[15], chunk[19], chunk[23], chunk[27], chunk[31], chunk[35], chunk[39], chunk[43], chunk[47], chunk[51], chunk[55], chunk[59], chunk[63]);
            ((((msbs - u8x16::splat(143)) * u8x16::splat(64))
              + lsbs - u8x16::splat(128)) - u8x16::splat(55))
                .store(out, i);
            i += 16;
        }
    }
    out
}

#[cfg(any(not(feature = "simd"), not(target_arch = "x86_64")))]
pub fn char_to_emoji<'a, 'b>(buf: &'a[u8], out: &'b mut [u8]) -> &'b [u8] {
    for (i, ch) in buf.iter().enumerate() {
        out[4 * i + 0] = 0xf0;
        out[4 * i + 1] = 0x9f;
        out[4 * i + 2] = (((*ch as u16).wrapping_add(55)) / 64 + 143) as u8;
        out[4 * i + 3] = (ch.wrapping_add(55) % 64).wrapping_add(128);
    }
    out
}

#[cfg(feature = "simd")]
#[cfg(target_feature = "sse2")]
pub fn char_to_emoji<'a, 'b>(buf: &'a[u8], out: &'b mut [u8]) -> &'b [u8] {
    // TODO: SSE this up
    for (i, ch) in buf.iter().enumerate() {
        out[4 * i + 0] = 0xf0;
        out[4 * i + 1] = 0x9f;
        out[4 * i + 2] = (((*ch as u16).wrapping_add(55)) / 64 + 143) as u8;
        out[4 * i + 3] = (ch.wrapping_add(55) % 64).wrapping_add(128);
    }
    out
}
<commit_msg>Prepare for vectorized encoding<commit_after>// base💯 - Copyright 2017 Adam Niederer

// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

#![allow(non_upper_case_globals)]
#![cfg_attr(test, feature(test))]
#![cfg_attr(feature = "simd", feature(target_feature))]
#![cfg_attr(feature = "simd", feature(asm))]
#![cfg_attr(feature = "simd", feature(cfg_target_feature))]

#[macro_use] extern crate clap;
#[cfg(feature = "simd")] extern crate stdsimd;

use std::io::{self, Read, Write, BufRead, BufReader, BufWriter};
use std::fs::File;
use std::iter::Iterator;
use clap::App;

fn main() {
    let cli_spec = load_yaml!("cli.yml");
    let cli_args = App::from_yaml(cli_spec).get_matches();

    let mut reader = {
        if let Some(path) = cli_args.value_of("input") {
            Box::new(BufReader::new(match File::open(path) {
                Ok(path) => path,
                _ => {
                    writeln!(io::stderr(), "base💯: no such file: {}", path).expect("base💯: stderr write error");
                    return;
                }
            })) as Box<BufRead>
        } else {
            Box::new(BufReader::new(io::stdin())) as Box<BufRead>
        }
    };

    let mut writer = BufWriter::new(io::stdout());

    if cli_args.is_present("decode") {
        let mut write_buf = [0u8; 0x10000];
        let mut buffer = [0u8; 0x10000];

        while let Ok(num_read) = reader.read(&mut buffer) {
            if num_read == 0 {
                break;
            }

            match writer.write_all(&emoji_to_char(&buffer[..num_read], &mut write_buf)[..num_read / 4]) {
                Ok(_) => (),
                _ => {
                    writeln!(io::stderr(), "base💯: write error").expect("base💯: stderr write error");
                    return;
                }
            }
        }
    } else {
        let mut write_buf = [0u8; 0x40000];
        let mut buffer = [0u8; 0x10000];

        while let Ok(num_read) = reader.read(&mut buffer) {
            if num_read == 0 {
                break;
            }

            match writer.write_all(&char_to_emoji(&buffer[..num_read], &mut write_buf)[..num_read * 4]) {
                Ok(_) => (),
                _ => {
                    writeln!(io::stderr(), "base💯: write error").expect("base💯: stderr write error");
                    return;
                }
            }
        }
    }
    writer.flush().expect("Write error");
}

#[cfg(any(not(feature = "simd"), not(target_arch = "x86_64")))]
fn emoji_to_char<'a, 'b>(buf: &'a[u8], out: &'b mut [u8]) -> &'b[u8] {
    for (i, chunk) in buf.chunks(4).enumerate() {
        out[i] = ((chunk[2].wrapping_sub(143)).wrapping_mul(64)).wrapping_add(chunk[3].wrapping_sub(128)).wrapping_sub(55)
    }
    out
}

#[cfg(feature = "simd")]
#[cfg(target_feature = "avx2")]
pub fn emoji_to_char<'a, 'b>(buf: &'a[u8], out: &'b mut [u8]) -> &'b[u8] {
    use stdsimd::simd::u8x32;
    let mut i = 0;
    for chunk in buf.chunks(128) {
        if chunk.len() < 128 {
            // Non-SIMD implementation for the final <128 bytes
            for chunk in chunk.chunks(4) {
                out[i] = ((chunk[2].wrapping_sub(143)).wrapping_mul(64)).wrapping_add(chunk[3].wrapping_sub(128)).wrapping_sub(55);
                i += 1;
            }
        } else {
            // AVX implementation of decoding algo
            // a, b, c, d contain one garbage word and one useful word
            let a = u8x32::load(chunk, 0);
            let b = u8x32::load(chunk, 32);
            let c = u8x32::load(chunk, 64);
            let d = u8x32::load(chunk, 96);

            // Constant mask for removing low bytes
            let hi_mask = u8x32::new(0xFF, 0x00, 0xFF, 0x00,0xFF, 0x00, 0xFF, 0x00,
                                     0xFF, 0x00, 0xFF, 0x00,0xFF, 0x00, 0xFF, 0x00,
                                     0xFF, 0x00, 0xFF, 0x00,0xFF, 0x00, 0xFF, 0x00,
                                     0xFF, 0x00, 0xFF, 0x00,0xFF, 0x00, 0xFF, 0x00);

            // Results from the vector magic below
            let hi: u8x32;
            let lo: u8x32;

            unsafe {
                asm! {
                    "vpsrld ymm1, ymm1, 16
                     vpsrld ymm2, ymm2, 16
                     vpsrld ymm3, ymm3, 16
                     vpsrld ymm4, ymm4, 16
                     // ymm1 .. ymm4 now contain 0x00, 0x00, 0xhi, 0xlo ...

                     // Pack ymm* into ymm1 and ymm2 and order them properly
                     vpackusdw ymm1, ymm1, ymm2
                     vpackusdw ymm2, ymm3, ymm4
                     // vpackusdw interleaves ymm[13] and ymm[24; we want them end-to-end
                     vpermpd ymm1, ymm1, 0xD8 // [191:128] <-> [127:64]
                     vpermpd ymm2, ymm2, 0xD8 // [191:128] <-> [127:64]
                     // ymm1 and ymm2 now contain interleaved high and low bytes

                     // Store high bytes (0xhi 0x00) in ymm3 and ymm4
                     vpand ymm3, ymm1, ymm6
                     vpand ymm4, ymm2, ymm6
                     // Pack hi bytes into ymm4
                     vpackuswb ymm4, ymm3, ymm4

                     // Remove the high bytes and move low bytes into position
                     vpsrlw ymm1, ymm1, 8
                     vpsrlw ymm2, ymm2, 8
                     // Store low bytes in ymm3
                     vpackuswb ymm3, ymm1, ymm2

                     // Fight interleaving again
                     vpermpd ymm4, ymm4, 0xD8 // [191:128] <-> [127:64]
                     vpermpd ymm3, ymm3, 0xD8 // [191:128] <-> [127:64]"
                     : "={ymm3}"(lo), "={ymm4}"(hi)
                     : "{ymm1}"(a), "{ymm2}"(b), "{ymm3}"(c),
                       "{ymm4}"(d), "{ymm6}"(hi_mask)
                     : "cc"
                     : "intel"
                }
            }

            ((((hi - u8x32::splat(143)) * u8x32::splat(64))
              + lo - u8x32::splat(128)) - u8x32::splat(55))
                .store(out, i);
            i += 32;
        }
    }
    out
}

#[cfg(feature = "simd")]
#[cfg(all(not(target_feature = "avx2"), target_feature = "sse2"))]
pub fn emoji_to_char<'a, 'b>(buf: &'a[u8], out: &'b mut [u8]) -> &'b[u8] {
    use stdsimd::simd::u8x16;
    let mut i = 0;
    for chunk in buf.chunks(64) {
        // Chunk data format: 0x00 0x00 0xhi 0xlo
        if chunk.len() < 64 {
            // Non-SIMD implementation for the final <128 bytes
            for (i, chunk) in chunk.chunks(4).enumerate() {
                out[i] = ((chunk[2].wrapping_sub(143)).wrapping_mul(64)).wrapping_add(chunk[3].wrapping_sub(128)).wrapping_sub(55)
            }
        } else {

            // SSE2 implementation of decoding algo

            // TODO: Use Get rid of these scalar loads
            let msbs = u8x16::new(chunk[2], chunk[6], chunk[10], chunk[14], chunk[18], chunk[22], chunk[26], chunk[30], chunk[34], chunk[38], chunk[42], chunk[46], chunk[50], chunk[54], chunk[58], chunk[62]);
            let lsbs = u8x16::new(chunk[3], chunk[7], chunk[11], chunk[15], chunk[19], chunk[23], chunk[27], chunk[31], chunk[35], chunk[39], chunk[43], chunk[47], chunk[51], chunk[55], chunk[59], chunk[63]);
            ((((msbs - u8x16::splat(143)) * u8x16::splat(64))
              + lsbs - u8x16::splat(128)) - u8x16::splat(55))
                .store(out, i);
            i += 16;
        }
    }
    out
}

#[cfg(any(not(feature = "simd"), not(target_arch = "x86_64")))]
pub fn char_to_emoji<'a, 'b>(buf: &'a[u8], out: &'b mut [u8]) -> &'b [u8] {
    for (i, ch) in buf.iter().enumerate() {
        out[4 * i + 0] = 0xf0;
        out[4 * i + 1] = 0x9f;
        out[4 * i + 2] = ((((*ch as u16).wrapping_add(55)) >> 6) + 143) as u8;
        out[4 * i + 3] = ((ch.wrapping_add(55) & 0x3f) - 64).wrapping_add(128);
    }
    out
}

#[cfg(feature = "simd")]
#[cfg(target_feature = "sse2")]
pub fn char_to_emoji<'a, 'b>(buf: &'a[u8], out: &'b mut [u8]) -> &'b [u8] {
    // TODO: SSE this up
    for (i, ch) in buf.iter().enumerate() {
        out[4 * i + 0] = 0xf0;
        out[4 * i + 1] = 0x9f;
        // (ch + 55) >> 6 approximates (ch + 55) / 64
        out[4 * i + 2] = ((((*ch as u16).wrapping_add(55)) >> 6) + 143) as u8;
        // (ch + 55) & 0x3f approximates (ch + 55) % 64
        out[4 * i + 3] = (ch.wrapping_add(55) & 0x3f).wrapping_add(128);
    }
    out
}
<|endoftext|>


<commit_before>#![forbid(unsafe_code)]

pub mod api;
pub mod db;
pub mod importer;
pub mod indexer;
pub mod model;
pub mod opening;
pub mod util;

use std::{mem, net::SocketAddr, path::PathBuf, sync::Arc, time::Duration};

use axum::{
    extract::{Extension, Path, Query},
    handler::{get, put},
    http::StatusCode,
    AddExtensionLayer, Json, Router,
};
use clap::Parser;
use futures_util::stream::Stream;
use serde::Deserialize;
use serde_with::{serde_as, DisplayFromStr};
use shakmaty::{
    fen::Fen,
    san::{San, SanPlus},
    uci::Uci,
    variant::{Variant, VariantPosition},
    zobrist::Zobrist,
    CastlingMode,
};
use tokio::sync::watch;

use crate::{
    api::{
        Error, ExplorerGame, ExplorerGameWithUci, ExplorerMove, ExplorerResponse, LichessQuery,
        Limits, MastersQuery, NdJson, PersonalQuery, PersonalQueryFilter,
    },
    db::{Database, LichessDatabase},
    importer::{LichessGame, LichessImporter, MastersImporter},
    indexer::{IndexerOpt, IndexerStub},
    model::{GameId, KeyBuilder, KeyPrefix, MastersGame, MastersGameWithId, PreparedMove, UserId},
    opening::{Opening, Openings},
    util::DedupStreamExt as _,
};

#[derive(Parser)]
struct Opt {
    #[clap(long, default_value = "127.0.0.1:9000")]
    bind: SocketAddr,
    #[clap(long, default_value = "_db")]
    db: PathBuf,
    #[clap(long)]
    cors: bool,
    #[clap(flatten)]
    indexer: IndexerOpt,
}

#[tokio::main]
async fn main() {
    env_logger::Builder::from_env(
        env_logger::Env::new()
            .filter("EXPLORER_LOG")
            .write_style("EXPLORER_LOG_STYLE"),
    )
    .format_timestamp(None)
    .format_module_path(false)
    .format_target(false)
    .init();

    let opt = Opt::parse();

    let openings: &'static Openings = Box::leak(Box::new(Openings::build_table()));
    let db = Arc::new(Database::open(opt.db).expect("db"));
    let (indexer, join_handles) = IndexerStub::spawn(Arc::clone(&db), opt.indexer);
    let masters_importer = MastersImporter::new(Arc::clone(&db));
    let lichess_importer = LichessImporter::new(Arc::clone(&db));

    let app = Router::new()
        .route("/admin/explorer.indexing", get(num_indexing))
        .route("/import/masters", put(masters_import))
        .route("/import/lichess", put(lichess_import))
        .route("/masters/pgn/:id", get(masters_pgn))
        .route("/masters", get(masters))
        .route("/master", get(masters)) // bc
        .route("/personal", get(personal)) // bc
        .route("/player", get(personal))
        .route("/lichess", get(lichess))
        .layer(AddExtensionLayer::new(openings))
        .layer(AddExtensionLayer::new(db))
        .layer(AddExtensionLayer::new(masters_importer))
        .layer(AddExtensionLayer::new(lichess_importer))
        .layer(AddExtensionLayer::new(indexer));

    let app = if opt.cors {
        app.layer(tower_http::set_header::SetResponseHeaderLayer::<
            _,
            axum::body::Body,
        >::if_not_present(
            axum::http::header::ACCESS_CONTROL_ALLOW_ORIGIN,
            axum::http::HeaderValue::from_static("*"),
        ))
    } else {
        app
    };

    axum::Server::bind(&opt.bind)
        .serve(app.into_make_service())
        .await
        .expect("bind");

    for join_handle in join_handles {
        join_handle.await.expect("indexer");
    }
}

async fn num_indexing(Extension(indexer): Extension<IndexerStub>) -> String {
    indexer.num_indexing().await.to_string()
}

struct PersonalStreamState {
    indexing: Option<watch::Receiver<()>>,
    key: KeyPrefix,
    db: Arc<Database>,
    filter: PersonalQueryFilter,
    limits: Limits,
    pos: VariantPosition,
    opening: Option<&'static Opening>,
    first: bool,
    done: bool,
}

fn finalize_lichess_moves(
    moves: Vec<PreparedMove>,
    pos: &VariantPosition,
    lichess_db: &LichessDatabase,
) -> Vec<ExplorerMove> {
    moves
        .into_iter()
        .map(|p| ExplorerMove {
            stats: p.stats,
            san: p.uci.to_move(pos).map_or(
                SanPlus {
                    san: San::Null,
                    suffix: None,
                },
                |m| SanPlus::from_move(pos.clone(), &m),
            ),
            uci: p.uci,
            average_rating: p.average_rating,
            average_opponent_rating: p.average_opponent_rating,
            game: p.game.and_then(|id| {
                lichess_db
                    .game(id)
                    .expect("get game")
                    .map(|info| ExplorerGame::from_lichess(id, info))
            }),
        })
        .collect()
}

fn finalize_lichess_games(
    games: Vec<(Uci, GameId)>,
    lichess_db: &LichessDatabase,
) -> Vec<ExplorerGameWithUci> {
    games
        .into_iter()
        .flat_map(|(uci, id)| {
            lichess_db
                .game(id)
                .expect("get game")
                .map(|info| ExplorerGameWithUci {
                    uci,
                    row: ExplorerGame::from_lichess(id, info),
                })
        })
        .collect()
}

async fn personal(
    Extension(openings): Extension<&'static Openings>,
    Extension(db): Extension<Arc<Database>>,
    Extension(indexer): Extension<IndexerStub>,
    Query(query): Query<PersonalQuery>,
) -> Result<NdJson<impl Stream<Item = ExplorerResponse>>, Error> {
    let player = UserId::from(query.player);
    let indexing = indexer.index_player(&player).await;

    let variant = query.variant.into();

    let mut pos = Zobrist::new(match query.fen {
        Some(fen) => VariantPosition::from_setup(variant, &Fen::from(fen), CastlingMode::Chess960)?,
        None => VariantPosition::new(variant),
    });

    let opening = openings.classify_and_play(&mut pos, query.play)?;

    let key = KeyBuilder::personal(&player, query.color).with_zobrist(variant, pos.zobrist_hash());

    let state = PersonalStreamState {
        filter: query.filter,
        limits: query.limits,
        db,
        indexing,
        opening,
        key,
        pos: pos.into_inner(),
        first: true,
        done: false,
    };

    Ok(NdJson(futures_util::stream::unfold(
        state,
        |mut state| async move {
            if state.done {
                return None;
            }

            let first = mem::replace(&mut state.first, false);
            state.done = match state.indexing {
                Some(ref mut indexing) => {
                    tokio::select! {
                        _ = indexing.changed() => true,
                        _ = tokio::time::sleep(Duration::from_millis(if first { 0 } else { 1000 })) => false,
                    }
                }
                None => true,
            };

            let lichess_db = state.db.lichess();
            let mut filtered = lichess_db
                .read_player(&state.key, state.filter.since, state.filter.until)
                .expect("get personal")
                .prepare(&state.filter);

            filtered.moves.truncate(state.limits.moves.unwrap_or(usize::MAX));
            filtered.recent_games.truncate(state.limits.recent_games);

            Some((
                ExplorerResponse {
                    total: filtered.total,
                    moves: finalize_lichess_moves(filtered.moves, &state.pos, &lichess_db),
                    recent_games: Some(finalize_lichess_games(filtered.recent_games, &lichess_db)),
                    top_games: None,
                    opening: state.opening,
                },
                state,
            ))
        },
    ).dedup_by_key(|res| res.total.total())))
}

async fn masters_import(
    Json(body): Json<MastersGameWithId>,
    Extension(importer): Extension<MastersImporter>,
) -> Result<(), Error> {
    importer.import(body).await
}

#[serde_as]
#[derive(Deserialize)]
struct MastersGameId(#[serde_as(as = "DisplayFromStr")] GameId);

async fn masters_pgn(
    Path(MastersGameId(id)): Path<MastersGameId>,
    Extension(db): Extension<Arc<Database>>,
) -> Result<MastersGame, StatusCode> {
    match db.masters().game(id).expect("get masters game") {
        Some(game) => Ok(game),
        None => Err(StatusCode::NOT_FOUND),
    }
}

async fn masters(
    Extension(openings): Extension<&'static Openings>,
    Extension(db): Extension<Arc<Database>>,
    Query(query): Query<MastersQuery>,
) -> Result<Json<ExplorerResponse>, Error> {
    let mut pos = Zobrist::new(match query.fen {
        Some(fen) => {
            VariantPosition::from_setup(Variant::Chess, &Fen::from(fen), CastlingMode::Chess960)?
        }
        None => VariantPosition::new(Variant::Chess),
    });

    let opening = openings.classify_and_play(&mut pos, query.play)?;
    let key = KeyBuilder::masters().with_zobrist(Variant::Chess, pos.zobrist_hash());
    let masters_db = db.masters();
    let mut entry = masters_db
        .read(key, query.since, query.until)
        .expect("get masters")
        .prepare();

    entry.moves.truncate(query.limits.moves.unwrap_or(12));
    entry.top_games.truncate(query.limits.top_games);

    Ok(Json(ExplorerResponse {
        total: entry.total,
        moves: entry
            .moves
            .into_iter()
            .map(|p| ExplorerMove {
                san: p.uci.to_move(&pos).map_or(
                    SanPlus {
                        san: San::Null,
                        suffix: None,
                    },
                    |m| SanPlus::from_move(pos.clone(), &m),
                ),
                uci: p.uci,
                average_rating: p.average_rating,
                average_opponent_rating: p.average_opponent_rating,
                stats: p.stats,
                game: p.game.and_then(|id| {
                    masters_db
                        .game(id)
                        .expect("get masters game")
                        .map(|info| ExplorerGame::from_masters(id, info))
                }),
            })
            .collect(),
        top_games: Some(
            entry
                .top_games
                .into_iter()
                .flat_map(|(uci, id)| {
                    masters_db
                        .game(id)
                        .expect("get masters game")
                        .map(|info| ExplorerGameWithUci {
                            uci,
                            row: ExplorerGame::from_masters(id, info),
                        })
                })
                .collect(),
        ),
        opening,
        recent_games: None,
    }))
}

async fn lichess_import(
    Json(body): Json<Vec<LichessGame>>,
    Extension(importer): Extension<LichessImporter>,
) -> Result<(), Error> {
    for game in body {
        importer.import(game).await?;
    }
    Ok(())
}

async fn lichess(
    Extension(openings): Extension<&'static Openings>,
    Extension(db): Extension<Arc<Database>>,
    Query(query): Query<LichessQuery>,
) -> Result<Json<ExplorerResponse>, Error> {
    let variant = Variant::from(query.variant);
    let mut pos = Zobrist::new(match query.fen {
        Some(fen) => VariantPosition::from_setup(variant, &Fen::from(fen), CastlingMode::Chess960)?,
        None => VariantPosition::new(variant),
    });

    let opening = openings.classify_and_play(&mut pos, query.play)?;

    let key = KeyBuilder::lichess().with_zobrist(variant, pos.zobrist_hash());
    let lichess_db = db.lichess();
    let mut filtered = lichess_db
        .read_lichess(&key, query.filter.since, query.filter.until)
        .expect("get lichess")
        .prepare(&dbg!(query.filter));

    filtered.moves.truncate(query.limits.moves.unwrap_or(12));
    filtered.recent_games.truncate(query.limits.recent_games);
    filtered.top_games.truncate(query.limits.top_games);

    Ok(Json(ExplorerResponse {
        total: filtered.total,
        moves: finalize_lichess_moves(filtered.moves, pos.as_inner(), &lichess_db),
        recent_games: Some(finalize_lichess_games(filtered.recent_games, &lichess_db)),
        top_games: Some(finalize_lichess_games(filtered.top_games, &lichess_db)),
        opening,
    }))
}
<commit_msg>reintroduce a way to read column family props<commit_after>#![forbid(unsafe_code)]

pub mod api;
pub mod db;
pub mod importer;
pub mod indexer;
pub mod model;
pub mod opening;
pub mod util;

use std::{mem, net::SocketAddr, path::PathBuf, sync::Arc, time::Duration};

use axum::{
    extract::{Extension, Path, Query},
    handler::{get, put},
    http::StatusCode,
    AddExtensionLayer, Json, Router,
};
use clap::Parser;
use futures_util::stream::Stream;
use serde::Deserialize;
use serde_with::{serde_as, DisplayFromStr};
use shakmaty::{
    fen::Fen,
    san::{San, SanPlus},
    uci::Uci,
    variant::{Variant, VariantPosition},
    zobrist::Zobrist,
    CastlingMode,
};
use tokio::sync::watch;

use crate::{
    api::{
        Error, ExplorerGame, ExplorerGameWithUci, ExplorerMove, ExplorerResponse, LichessQuery,
        Limits, MastersQuery, NdJson, PersonalQuery, PersonalQueryFilter,
    },
    db::{Database, LichessDatabase},
    importer::{LichessGame, LichessImporter, MastersImporter},
    indexer::{IndexerOpt, IndexerStub},
    model::{GameId, KeyBuilder, KeyPrefix, MastersGame, MastersGameWithId, PreparedMove, UserId},
    opening::{Opening, Openings},
    util::DedupStreamExt as _,
};

#[derive(Parser)]
struct Opt {
    #[clap(long, default_value = "127.0.0.1:9000")]
    bind: SocketAddr,
    #[clap(long, default_value = "_db")]
    db: PathBuf,
    #[clap(long)]
    cors: bool,
    #[clap(flatten)]
    indexer: IndexerOpt,
}

#[tokio::main]
async fn main() {
    env_logger::Builder::from_env(
        env_logger::Env::new()
            .filter("EXPLORER_LOG")
            .write_style("EXPLORER_LOG_STYLE"),
    )
    .format_timestamp(None)
    .format_module_path(false)
    .format_target(false)
    .init();

    let opt = Opt::parse();

    let openings: &'static Openings = Box::leak(Box::new(Openings::build_table()));
    let db = Arc::new(Database::open(opt.db).expect("db"));
    let (indexer, join_handles) = IndexerStub::spawn(Arc::clone(&db), opt.indexer);
    let masters_importer = MastersImporter::new(Arc::clone(&db));
    let lichess_importer = LichessImporter::new(Arc::clone(&db));

    let app = Router::new()
        .route("/cf/:cf/:prop", get(cf_prop))
        .route("/admin/explorer.indexing", get(num_indexing))
        .route("/import/masters", put(masters_import))
        .route("/import/lichess", put(lichess_import))
        .route("/masters/pgn/:id", get(masters_pgn))
        .route("/masters", get(masters))
        .route("/master", get(masters)) // bc
        .route("/personal", get(personal)) // bc
        .route("/player", get(personal))
        .route("/lichess", get(lichess))
        .layer(AddExtensionLayer::new(openings))
        .layer(AddExtensionLayer::new(db))
        .layer(AddExtensionLayer::new(masters_importer))
        .layer(AddExtensionLayer::new(lichess_importer))
        .layer(AddExtensionLayer::new(indexer));

    let app = if opt.cors {
        app.layer(tower_http::set_header::SetResponseHeaderLayer::<
            _,
            axum::body::Body,
        >::if_not_present(
            axum::http::header::ACCESS_CONTROL_ALLOW_ORIGIN,
            axum::http::HeaderValue::from_static("*"),
        ))
    } else {
        app
    };

    axum::Server::bind(&opt.bind)
        .serve(app.into_make_service())
        .await
        .expect("bind");

    for join_handle in join_handles {
        join_handle.await.expect("indexer");
    }
}

#[derive(Deserialize)]
struct ColumnFamilyProp {
    cf: String,
    prop: String,
}

async fn cf_prop(
    Path(path): Path<ColumnFamilyProp>,
    Extension(db): Extension<Arc<Database>>,
) -> Result<String, StatusCode> {
    db.inner
        .cf_handle(&path.cf)
        .and_then(|cf| {
            db.inner
                .property_value_cf(cf, &path.prop)
                .expect("property value")
        })
        .ok_or(StatusCode::NOT_FOUND)
}

async fn num_indexing(Extension(indexer): Extension<IndexerStub>) -> String {
    indexer.num_indexing().await.to_string()
}

struct PersonalStreamState {
    indexing: Option<watch::Receiver<()>>,
    key: KeyPrefix,
    db: Arc<Database>,
    filter: PersonalQueryFilter,
    limits: Limits,
    pos: VariantPosition,
    opening: Option<&'static Opening>,
    first: bool,
    done: bool,
}

fn finalize_lichess_moves(
    moves: Vec<PreparedMove>,
    pos: &VariantPosition,
    lichess_db: &LichessDatabase,
) -> Vec<ExplorerMove> {
    moves
        .into_iter()
        .map(|p| ExplorerMove {
            stats: p.stats,
            san: p.uci.to_move(pos).map_or(
                SanPlus {
                    san: San::Null,
                    suffix: None,
                },
                |m| SanPlus::from_move(pos.clone(), &m),
            ),
            uci: p.uci,
            average_rating: p.average_rating,
            average_opponent_rating: p.average_opponent_rating,
            game: p.game.and_then(|id| {
                lichess_db
                    .game(id)
                    .expect("get game")
                    .map(|info| ExplorerGame::from_lichess(id, info))
            }),
        })
        .collect()
}

fn finalize_lichess_games(
    games: Vec<(Uci, GameId)>,
    lichess_db: &LichessDatabase,
) -> Vec<ExplorerGameWithUci> {
    games
        .into_iter()
        .flat_map(|(uci, id)| {
            lichess_db
                .game(id)
                .expect("get game")
                .map(|info| ExplorerGameWithUci {
                    uci,
                    row: ExplorerGame::from_lichess(id, info),
                })
        })
        .collect()
}

async fn personal(
    Extension(openings): Extension<&'static Openings>,
    Extension(db): Extension<Arc<Database>>,
    Extension(indexer): Extension<IndexerStub>,
    Query(query): Query<PersonalQuery>,
) -> Result<NdJson<impl Stream<Item = ExplorerResponse>>, Error> {
    let player = UserId::from(query.player);
    let indexing = indexer.index_player(&player).await;

    let variant = query.variant.into();

    let mut pos = Zobrist::new(match query.fen {
        Some(fen) => VariantPosition::from_setup(variant, &Fen::from(fen), CastlingMode::Chess960)?,
        None => VariantPosition::new(variant),
    });

    let opening = openings.classify_and_play(&mut pos, query.play)?;

    let key = KeyBuilder::personal(&player, query.color).with_zobrist(variant, pos.zobrist_hash());

    let state = PersonalStreamState {
        filter: query.filter,
        limits: query.limits,
        db,
        indexing,
        opening,
        key,
        pos: pos.into_inner(),
        first: true,
        done: false,
    };

    Ok(NdJson(futures_util::stream::unfold(
        state,
        |mut state| async move {
            if state.done {
                return None;
            }

            let first = mem::replace(&mut state.first, false);
            state.done = match state.indexing {
                Some(ref mut indexing) => {
                    tokio::select! {
                        _ = indexing.changed() => true,
                        _ = tokio::time::sleep(Duration::from_millis(if first { 0 } else { 1000 })) => false,
                    }
                }
                None => true,
            };

            let lichess_db = state.db.lichess();
            let mut filtered = lichess_db
                .read_player(&state.key, state.filter.since, state.filter.until)
                .expect("get personal")
                .prepare(&state.filter);

            filtered.moves.truncate(state.limits.moves.unwrap_or(usize::MAX));
            filtered.recent_games.truncate(state.limits.recent_games);

            Some((
                ExplorerResponse {
                    total: filtered.total,
                    moves: finalize_lichess_moves(filtered.moves, &state.pos, &lichess_db),
                    recent_games: Some(finalize_lichess_games(filtered.recent_games, &lichess_db)),
                    top_games: None,
                    opening: state.opening,
                },
                state,
            ))
        },
    ).dedup_by_key(|res| res.total.total())))
}

async fn masters_import(
    Json(body): Json<MastersGameWithId>,
    Extension(importer): Extension<MastersImporter>,
) -> Result<(), Error> {
    importer.import(body).await
}

#[serde_as]
#[derive(Deserialize)]
struct MastersGameId(#[serde_as(as = "DisplayFromStr")] GameId);

async fn masters_pgn(
    Path(MastersGameId(id)): Path<MastersGameId>,
    Extension(db): Extension<Arc<Database>>,
) -> Result<MastersGame, StatusCode> {
    match db.masters().game(id).expect("get masters game") {
        Some(game) => Ok(game),
        None => Err(StatusCode::NOT_FOUND),
    }
}

async fn masters(
    Extension(openings): Extension<&'static Openings>,
    Extension(db): Extension<Arc<Database>>,
    Query(query): Query<MastersQuery>,
) -> Result<Json<ExplorerResponse>, Error> {
    let mut pos = Zobrist::new(match query.fen {
        Some(fen) => {
            VariantPosition::from_setup(Variant::Chess, &Fen::from(fen), CastlingMode::Chess960)?
        }
        None => VariantPosition::new(Variant::Chess),
    });

    let opening = openings.classify_and_play(&mut pos, query.play)?;
    let key = KeyBuilder::masters().with_zobrist(Variant::Chess, pos.zobrist_hash());
    let masters_db = db.masters();
    let mut entry = masters_db
        .read(key, query.since, query.until)
        .expect("get masters")
        .prepare();

    entry.moves.truncate(query.limits.moves.unwrap_or(12));
    entry.top_games.truncate(query.limits.top_games);

    Ok(Json(ExplorerResponse {
        total: entry.total,
        moves: entry
            .moves
            .into_iter()
            .map(|p| ExplorerMove {
                san: p.uci.to_move(&pos).map_or(
                    SanPlus {
                        san: San::Null,
                        suffix: None,
                    },
                    |m| SanPlus::from_move(pos.clone(), &m),
                ),
                uci: p.uci,
                average_rating: p.average_rating,
                average_opponent_rating: p.average_opponent_rating,
                stats: p.stats,
                game: p.game.and_then(|id| {
                    masters_db
                        .game(id)
                        .expect("get masters game")
                        .map(|info| ExplorerGame::from_masters(id, info))
                }),
            })
            .collect(),
        top_games: Some(
            entry
                .top_games
                .into_iter()
                .flat_map(|(uci, id)| {
                    masters_db
                        .game(id)
                        .expect("get masters game")
                        .map(|info| ExplorerGameWithUci {
                            uci,
                            row: ExplorerGame::from_masters(id, info),
                        })
                })
                .collect(),
        ),
        opening,
        recent_games: None,
    }))
}

async fn lichess_import(
    Json(body): Json<Vec<LichessGame>>,
    Extension(importer): Extension<LichessImporter>,
) -> Result<(), Error> {
    for game in body {
        importer.import(game).await?;
    }
    Ok(())
}

async fn lichess(
    Extension(openings): Extension<&'static Openings>,
    Extension(db): Extension<Arc<Database>>,
    Query(query): Query<LichessQuery>,
) -> Result<Json<ExplorerResponse>, Error> {
    let variant = Variant::from(query.variant);
    let mut pos = Zobrist::new(match query.fen {
        Some(fen) => VariantPosition::from_setup(variant, &Fen::from(fen), CastlingMode::Chess960)?,
        None => VariantPosition::new(variant),
    });

    let opening = openings.classify_and_play(&mut pos, query.play)?;

    let key = KeyBuilder::lichess().with_zobrist(variant, pos.zobrist_hash());
    let lichess_db = db.lichess();
    let mut filtered = lichess_db
        .read_lichess(&key, query.filter.since, query.filter.until)
        .expect("get lichess")
        .prepare(&dbg!(query.filter));

    filtered.moves.truncate(query.limits.moves.unwrap_or(12));
    filtered.recent_games.truncate(query.limits.recent_games);
    filtered.top_games.truncate(query.limits.top_games);

    Ok(Json(ExplorerResponse {
        total: filtered.total,
        moves: finalize_lichess_moves(filtered.moves, pos.as_inner(), &lichess_db),
        recent_games: Some(finalize_lichess_games(filtered.recent_games, &lichess_db)),
        top_games: Some(finalize_lichess_games(filtered.top_games, &lichess_db)),
        opening,
    }))
}
<|endoftext|>


<commit_before>//! SEGIMAP is an IMAP server implementation.
#![feature(macro_rules)]
#![deny(non_camel_case_types)]
#![feature(phase)]

extern crate "rust-crypto" as crypto;
extern crate regex;
#[phase(plugin)] extern crate regex_macros;
#[phase(plugin, link)] extern crate log;
#[phase(plugin)] extern crate peg_syntax_ext;
extern crate serialize;
extern crate toml;

pub use config::Config;
pub use email::Email;
pub use login::LoginData;
pub use message::Message;
pub use server::Server;
pub use session::Session;
pub use user::User;
pub use sequence_set::SequenceSet;

use std::io::{Listener, Acceptor, BufferedStream};
use std::sync::Arc;

mod auth;
mod config;
mod email;
mod error;
mod folder;
mod login;
mod message;
mod sequence_set;
mod server;
mod session;
mod user;

/// The file in which user data is stored.
// TODO: add the ability for the user to specify the user data file as an
// argument.
static USER_DATA_FILE: &'static str = "./users.json";

fn main() {
    // Load configuration.
    let config = Config::new();

    // let mut users = Vec::new();
    // users.push(User::new(Email::new("will".to_string(), "xqz.ca".to_string()), "54321".to_string(), "./maildir".to_string()));
    // users.push(User::new(Email::new("nikitapekin".to_string(), "gmail.com".to_string()), "12345".to_string(), "./maildir".to_string()));

    // user::save_users(USER_DATA_FILE.to_string(), users);

    // Load the user data from the user data file.
    // TODO: figure out what to do for error handling.
    let users = user::load_users(USER_DATA_FILE.to_string()).unwrap();

    let multipart_message = Message::parse(&Path::new("maildir/cur/12345:2,FRS")).unwrap();
    let html_message = Message::parse(&Path::new("maildir/cur/54321:2,FS")).unwrap();

    // Avoid unused variable notices temporarily.
    //println!("Config: {}", config);
    //println!("Users: {}", users);
    //println!("Message: {}", multipart_message);
    //println!("Message: {}", html_message);

    let serv = Arc::new(Server::new(config, users));
    match serv.imap_listener() {
        Err(e) => {
            error!("Error listening on IMAP port: {}", e);
        }
        Ok(v) => {
            let mut acceptor = v.listen();
            for stream in acceptor.incoming() {
                match stream {
                    Err(e) => {
                        error!("Error accepting incoming connection: {}", e)
                    }
                    Ok(stream) => {
                        let session_serv = serv.clone();
                        spawn(proc() {
                            let mut session = Session::new(BufferedStream::new(stream), session_serv);
                            session.handle();
                        });
                    }
                }
            }
            drop(acceptor);
        }
    }
}

<commit_msg>Fix merge errorsg<commit_after>//! SEGIMAP is an IMAP server implementation.
#![feature(macro_rules)]
#![deny(non_camel_case_types)]
#![feature(phase)]

extern crate "rust-crypto" as crypto;
extern crate regex;
#[phase(plugin)] extern crate regex_macros;
#[phase(plugin, link)] extern crate log;
#[phase(plugin)] extern crate peg_syntax_ext;
extern crate serialize;
extern crate toml;

pub use config::Config;
pub use email::Email;
pub use login::LoginData;
pub use message::Message;
pub use server::Server;
pub use session::Session;
pub use user::User;
pub use sequence_set::SequenceSet;

use std::io::{Listener, Acceptor, BufferedStream};
use std::sync::Arc;

mod auth;
mod config;
mod email;
mod error;
mod folder;
mod login;
mod message;
mod sequence_set;
mod server;
mod session;
mod user;

/// The file in which user data is stored.
// TODO: add the ability for the user to specify the user data file as an
// argument.
static USER_DATA_FILE: &'static str = "./users.json";

fn main() {
    // Load configuration.
    let config = Config::new();

    // let mut users = Vec::new();
    // users.push(User::new(Email::new("will".to_string(), "xqz.ca".to_string()), "54321".to_string(), "./maildir".to_string()));
    // users.push(User::new(Email::new("nikitapekin".to_string(), "gmail.com".to_string()), "12345".to_string(), "./maildir".to_string()));

    // user::save_users(USER_DATA_FILE.to_string(), users);

    // Load the user data from the user data file.
    // TODO: figure out what to do for error handling.
    let users = user::load_users(USER_DATA_FILE.to_string()).unwrap();

    let multipart_message = Message::parse(&Path::new("maildir/new/12345")).unwrap();
    let html_message = Message::parse(&Path::new("maildir/cur/54321:2,FS")).unwrap();

    // Avoid unused variable notices temporarily.
    //println!("Config: {}", config);
    //println!("Users: {}", users);
    //println!("Message: {}", multipart_message);
    //println!("Message: {}", html_message);

    let serv = Arc::new(Server::new(config, users));
    match serv.imap_listener() {
        Err(e) => {
            error!("Error listening on IMAP port: {}", e);
        }
        Ok(v) => {
            let mut acceptor = v.listen();
            for stream in acceptor.incoming() {
                match stream {
                    Err(e) => {
                        error!("Error accepting incoming connection: {}", e)
                    }
                    Ok(stream) => {
                        let session_serv = serv.clone();
                        spawn(proc() {
                            let mut session = Session::new(BufferedStream::new(stream), session_serv);
                            session.handle();
                        });
                    }
                }
            }
            drop(acceptor);
        }
    }
}

<|endoftext|>


<commit_before>extern crate mio;

use std::io::Write;
use std::io::Read;
use std::rc::{Rc};
use std::cell::{RefCell};
use mio::tcp::*;
mod server;
mod broker;
mod message;

const MQTT_SERVER_TOKEN: mio::Token = mio::Token(0);

fn main() {
    let address = "0.0.0.0:1883".parse().unwrap();
    let listener = TcpListener::bind(&address).expect(&format!("Could not bind to {}", address));
    let mut event_loop = mio::EventLoop::new().expect("Could not create MIO event loop");
    event_loop.register(&listener, MQTT_SERVER_TOKEN).expect("Could not register listener");
    event_loop.run(&mut MioHandler::new(listener)).expect("Could not run event loop");
}

struct MioHandler {
    listener: TcpListener,
    connections: mio::util::Slab<Rc<RefCell<Connection>>>,
    mqtt_streams: mio::util::Slab<server::Stream>,
    server: server::Server<Connection>,
}

struct Connection {
    socket: mio::tcp::TcpStream,
}

impl MioHandler {
    fn new(listener: TcpListener) -> Self {
        let max_conns = 1024 * 32;
        let connections_slab = mio::util::Slab::new_starting_at(mio::Token(1), max_conns);
        let mqtt_stream_slab = mio::util::Slab::new_starting_at(mio::Token(1), max_conns);

        MioHandler {
            listener: listener,
            connections: connections_slab,
            mqtt_streams: mqtt_stream_slab,
            server: server::Server::new(),
        }
    }
}

impl mio::Handler for MioHandler {
    type Timeout = ();
    type Message = ();

    fn ready(&mut self,
             event_loop: &mut mio::EventLoop<MioHandler>,
             token: mio::Token,
             events: mio::EventSet) {
        match token {
            MQTT_SERVER_TOKEN => {
                assert!(events.is_readable());

                match self.listener.accept() {
                    Ok(Some(socket)) => {
                        let token = self.connections
                            .insert_with(|_| Rc::new(RefCell::new(Connection::new(socket))))
                            .expect("Could not insert new connection in slab");
                        self.mqtt_streams.
                            insert_with(|_| server::Stream::new())
                            .expect("Could not insert new stream into slab");
                        let connection = &self.connections[token].clone();
                        event_loop.register_opt(
                            &connection.borrow().socket,
                            token,
                            mio::EventSet::readable(),
                            mio::PollOpt::edge()).expect("Could not register connection with event loop");
                    }
                    Ok(None) => {
                        println!("The server socket wasn't actually ready");
                    }
                    Err(e) => {
                        println!("listener.accept errored: {}", e);
                        event_loop.shutdown();
                    }
                }
            }
            _ => {
                assert!(events.is_readable());
                let still_connected = connection_ready(&mut self.server,
                                                       &mut self.mqtt_streams[token],
                                                       self.connections[token].clone());
                if !still_connected {
                    event_loop.deregister(&self.connections[token].borrow().socket)
                        .expect("Could not deregister connection with event loop");
                    self.server.unsubscribe_all(self.connections[token].clone());
                    self.connections.remove(token).expect("Could not remove connection from slab");
                }
            }
        }
    }
}

fn connection_ready(server: &mut server::Server<Connection>,
                    stream: &mut server::Stream,
                    connection: Rc<RefCell<Connection>>) -> bool {
    let connection = connection.clone();
    let read_result = connection.borrow_mut().read(stream.buffer());

    match read_result {
        Ok(length) => {
            stream.handle_messages(length, server, connection.clone())
        },
        _ => panic!("Error reading bytes from stream"),
    }
}

impl Connection {
    fn new(socket: mio::tcp::TcpStream) -> Self {
        Connection { socket: socket }
    }

    fn read(&mut self, buffer: &mut [u8]) -> std::io::Result<usize> {
        self.socket.read(buffer)
    }
}

impl broker::Subscriber for Connection {
    fn new_message(&mut self, bytes: &[u8]) {
        self.socket.write(bytes).expect("Error writing to socket");
    }
}
<commit_msg>Panics if too many bytes to handle<commit_after>extern crate mio;

use std::io::Write;
use std::io::Read;
use std::rc::{Rc};
use std::cell::{RefCell};
use mio::tcp::*;
mod server;
mod broker;
mod message;

const MQTT_SERVER_TOKEN: mio::Token = mio::Token(0);

fn main() {
    let address = "0.0.0.0:1883".parse().unwrap();
    let listener = TcpListener::bind(&address).expect(&format!("Could not bind to {}", address));
    let mut event_loop = mio::EventLoop::new().expect("Could not create MIO event loop");
    event_loop.register(&listener, MQTT_SERVER_TOKEN).expect("Could not register listener");
    event_loop.run(&mut MioHandler::new(listener)).expect("Could not run event loop");
}

struct MioHandler {
    listener: TcpListener,
    connections: mio::util::Slab<Rc<RefCell<Connection>>>,
    mqtt_streams: mio::util::Slab<server::Stream>,
    server: server::Server<Connection>,
}

struct Connection {
    socket: mio::tcp::TcpStream,
}

impl MioHandler {
    fn new(listener: TcpListener) -> Self {
        let max_conns = 1024 * 32;
        let connections_slab = mio::util::Slab::new_starting_at(mio::Token(1), max_conns);
        let mqtt_stream_slab = mio::util::Slab::new_starting_at(mio::Token(1), max_conns);

        MioHandler {
            listener: listener,
            connections: connections_slab,
            mqtt_streams: mqtt_stream_slab,
            server: server::Server::new(),
        }
    }
}

impl mio::Handler for MioHandler {
    type Timeout = ();
    type Message = ();

    fn ready(&mut self,
             event_loop: &mut mio::EventLoop<MioHandler>,
             token: mio::Token,
             events: mio::EventSet) {
        match token {
            MQTT_SERVER_TOKEN => {
                assert!(events.is_readable());

                match self.listener.accept() {
                    Ok(Some(socket)) => {
                        let token = self.connections
                            .insert_with(|_| Rc::new(RefCell::new(Connection::new(socket))))
                            .expect("Could not insert new connection in slab");
                        self.mqtt_streams.
                            insert_with(|_| server::Stream::new())
                            .expect("Could not insert new stream into slab");
                        let connection = &self.connections[token].clone();
                        event_loop.register_opt(
                            &connection.borrow().socket,
                            token,
                            mio::EventSet::readable(),
                            mio::PollOpt::edge()).expect("Could not register connection with event loop");
                    }
                    Ok(None) => {
                        println!("The server socket wasn't actually ready");
                    }
                    Err(e) => {
                        println!("listener.accept errored: {}", e);
                        event_loop.shutdown();
                    }
                }
            }
            _ => {
                assert!(events.is_readable());
                let still_connected = connection_ready(&mut self.server,
                                                       &mut self.mqtt_streams[token],
                                                       self.connections[token].clone());
                if !still_connected {
                    event_loop.deregister(&self.connections[token].borrow().socket)
                        .expect("Could not deregister connection with event loop");
                    self.server.unsubscribe_all(self.connections[token].clone());
                    self.connections.remove(token).expect("Could not remove connection from slab");
                }
            }
        }
    }
}

fn connection_ready(server: &mut server::Server<Connection>,
                    stream: &mut server::Stream,
                    connection: Rc<RefCell<Connection>>) -> bool {
    let connection = connection.clone();
    let read_result = connection.borrow_mut().read(stream.buffer());

    match read_result {
        Ok(length) => {
            if length >= stream.total_buffer_len() {
                panic!(format!("Too many bytes ({}) for puny stream buffer ({})",
                               length, stream.total_buffer_len()));
            }
            stream.handle_messages(length, server, connection.clone())
        },
        _ => panic!("Error reading bytes from stream"),
    }
}

impl Connection {
    fn new(socket: mio::tcp::TcpStream) -> Self {
        Connection { socket: socket }
    }

    fn read(&mut self, buffer: &mut [u8]) -> std::io::Result<usize> {
        self.socket.read(buffer)
    }
}

impl broker::Subscriber for Connection {
    fn new_message(&mut self, bytes: &[u8]) {
        self.socket.write(bytes).expect("Error writing to socket");
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate clap;
use clap::{App, Arg};

arg_enum! {
    #[derive(Debug)]
    enum Method
    {
        install,
        remove
    }
}

fn main() {
    let arg_matches = App::new("wiz")
                        .version("0.1.0")

                        .arg(Arg::with_name("method")
                            .help("Specify what to do with the specified package")
                            .index(1)
                            .possible_values(&Method::variants())
                            .required(true))

                        .arg(Arg::with_name("package")
                            .help("Specify package")
                            .index(2)
                            .required(true))

                        .get_matches();

    let method = value_t!(arg_matches.value_of("method"), Method).unwrap();
    let package = arg_matches.value_of("package").unwrap_or("invalid");

    match method {
        Method::install => println!("install: {}", package),
        Method::remove => println!("remove: {}", package)
    }
}
<commit_msg>camal case ?<commit_after>#[macro_use]
extern crate clap;
use clap::{App, Arg};

arg_enum! {
    #[derive(Debug)]
    enum Method {
        Install,
        Remove
    }
}

fn main() {
    let arg_matches = App::new("wiz")
        .version("0.1.0")

        .arg(Arg::with_name("method")
             .help("Specify what to do with the specified package")
             .index(1)
             .possible_values(&Method::variants())
             .required(true))

        .arg(Arg::with_name("package")
             .help("Specify package")
             .index(2)
             .required(true))

        .get_matches();

    let method = value_t!(arg_matches.value_of("method"), Method).unwrap();
    let package = arg_matches.value_of("package").unwrap_or("invalid");

    match method {
        Method::Install => println!("install: {}", package),
        Method::Remove => println!("remove: {}", package)
    }
}
<|endoftext|>


<commit_before>extern crate lp_modeler;

use lp_modeler::operations::LpOperations;
use lp_modeler::problem::{LpProblem, LpObjective, LpFileFormat};
use lp_modeler::solvers::*;
use lp_modeler::variables::{LpInteger};

fn main() {

    let ref a = LpInteger::new("a");
    let ref b = LpInteger::new("b");
    let ref c = LpInteger::new("c");

    let mut problem = LpProblem::new("Problem", LpObjective::Maximize);

    /*
    problem += 10.0 * a + 20.0 * b;

    problem += (300*(a-b)).ge(100);
    problem += (300*(-a+b)).le(100);
    problem += (a+b).le(10);
    */
    let cacou = ((( 6 + 3*a) + 3*b ) * a );
    let coucou = (3 * (2+a+b) * (2+a));
    println!("{}", (coucou).to_lp_file_format());
    println!("{}", cacou.to_lp_file_format());

   // problem += ((a+b)*300).le(100);

    // Suboptimal instead of unfeasible
    // problem += (a+b).equal(10);

    // let solver = GurobiSolver
    // solver <<= BaseDirectory("/opt/gurobi1.2/...")
    // solver <<= Config().arg("-thread 2").arg("...")

    let _ = problem.write_lp("toto.lp");

    let solver = GurobiSolver::new();

    match solver.run(&problem) {
        Ok((status, res)) => {
            println!("Status {:?}", status);
            for (name, value) in res.iter() {
                println!("value of {} = {}", name, value);
            }
        },
        Err(msg) => println!("{}", msg),
    }

}
<commit_msg>Add some case to test commut. property<commit_after>extern crate lp_modeler;

use lp_modeler::operations::LpOperations;
use lp_modeler::problem::{LpProblem, LpObjective, LpFileFormat};
use lp_modeler::solvers::*;
use lp_modeler::variables::{LpInteger};

fn main() {

    let ref a = LpInteger::new("a");
    let ref b = LpInteger::new("b");
    let ref c = LpInteger::new("c");

    let mut problem = LpProblem::new("Problem", LpObjective::Maximize);

    /*
    problem += 10.0 * a + 20.0 * b;

    problem += (300*(a-b)).ge(100);
    problem += (300*(-a+b)).le(100);
    problem += (a+b).le(10);
    */
    let cacou = ((( 6 + 3*a) + 3*b ) * a );
    let coucou = (3 * (2+a+b) * (2+a));
    println!("{}", coucou.to_lp_file_format());
    println!("{}", (coucou*15).to_lp_file_format());
    println!("{}", (15*((3*a)*a)).to_lp_file_format());
    println!("{}", cacou.to_lp_file_format());
    let test = (a * (2 + b) * 3);
    println!("{}", (10 * test).to_lp_file_format());

    problem += ((a+b)*300).le(100);

    // Suboptimal instead of unfeasible
    // problem += (a+b).equal(10);

    // let solver = GurobiSolver
    // solver <<= BaseDirectory("/opt/gurobi1.2/...")
    // solver <<= Config().arg("-thread 2").arg("...")

    let _ = problem.write_lp("toto.lp");

    let solver = GurobiSolver::new();

    match solver.run(&problem) {
        Ok((status, res)) => {
            println!("Status {:?}", status);
            for (name, value) in res.iter() {
                println!("value of {} = {}", name, value);
            }
        },
        Err(msg) => println!("{}", msg),
    }

}
<|endoftext|>


<commit_before>use std::io;
use std::io::prelude::*;

extern crate forth;

use forth::interpreter::Interpreter;

fn main() {
    println!("forth-rs version 0.1.0");
    println!("Type 'bye' to exit");
    let mut interp = Interpreter::new();
    let mut stdin = io::stdin();
    loop {
        io::stdout().flush().ok().expect("Could not flush stdout");
        let mut line = String::new();
        stdin.read_line(&mut line).ok().expect("Unable to read from stdin");
        match interp.exec(line.trim()) {
            Err(e) => println!("{:?}", e),
            Ok(_) => {},
        }
    }
}
<commit_msg>More standard forth prompt<commit_after>use std::io;
use std::io::prelude::*;

extern crate forth;

use forth::interpreter::Interpreter;

fn main() {
    println!("forth-rs version 0.1.0");
    println!("Type 'bye' to exit");
    let mut interp = Interpreter::new();
    let mut stdin = io::stdin();
    loop {
        io::stdout().flush().ok().expect("Could not flush stdout");
        let mut line = String::new();
        stdin.read_line(&mut line).ok().expect("Unable to read from stdin");
        match interp.exec(line.trim()) {
            Err(e) => print!("{:?}", e),
            Ok(_) => print!("ok"),
        }
    }
}
<|endoftext|>


<commit_before>// disable warnings for dead code when running debug
#![cfg_attr(debug_assertions, allow(dead_code))]

extern crate net2;
extern crate bincode;
extern crate rustc_serialize;
#[macro_use]
extern crate log;
extern crate env_logger;
extern crate byteorder;

use std::net::{Ipv4Addr, Ipv6Addr};
#[cfg(not(windows))]
use net2::unix::UnixUdpBuilderExt;
use net2::UdpBuilder;
use bincode::rustc_serialize::decode;
use byteorder::{ByteOrder, NetworkEndian};

#[repr(C)]
#[repr(packed)]
#[derive(Default)]
#[derive(Debug)]
#[derive(RustcEncodable, RustcDecodable)]
struct MDNSPacketHeader {
    id: u16, // transaction ID
    flags: u16,
    num_qn: u16,
    num_ans_rr: u16,
    num_auth_rr: u16,
    num_add_rr: u16,
}

#[derive(Debug)]
enum MDNSType {
    A(Ipv4Addr),
    PTR,
    TXT(Vec<String>),
    AAAA(Ipv6Addr),
    SRV,
    NSEC,
    ANY,
    UNKNOW(u16),
}

impl MDNSType {
    #[doc(hidden)]
    pub fn to_u16(&self) -> u16 {
        match *self {
            MDNSType::A(_) => 0x01,
            MDNSType::PTR => 0x0C,
            MDNSType::TXT(_) => 0x10,
            MDNSType::AAAA(_) => 0x1C,
            MDNSType::SRV => 0x21,
            MDNSType::NSEC => 0x2F,
            MDNSType::ANY => 0xFF,
            MDNSType::UNKNOW(n) => n,
        }
    }

    pub fn from_data(mdns_type_id: u16, packet_data: &[u8]) -> Result<MDNSType, String> {
        match mdns_type_id {
            0x10 => {
                let mut txt_map = vec![];
                let mut data_offset: usize = 0;
                while packet_data.len() > data_offset {
                    let (label_string, label_size) = label_to_string(packet_data, data_offset);
                    trace!("TXT: {} {}", label_string, label_size);
                    txt_map.push(label_string);
                    data_offset += label_size;
                }
                trace!("{:?}", txt_map);
                Ok(MDNSType::TXT(txt_map))
            }
            0x01 => {
                if packet_data.len() < 4 {
                    return Err("Data section to small for IPV4 address".to_owned());
                }

                let ip = NetworkEndian::read_u32(&packet_data[0..4]);
                let address = Ipv4Addr::new(((ip >> 24) & 0xFF) as u8,
                                            ((ip >> 16) & 0xFF) as u8,
                                            ((ip >> 8) & 0xFF) as u8,
                                            (ip & 0xFF) as u8);

                Ok(MDNSType::A(address))
            }
            0x1C => {
                if packet_data.len() < (8 * 2) {
                    return Err("Data section to small for IPV6 address".to_owned());
                }

                let mut ipv6_octets = [0u16; 8];
                for ip_index in 0..8 {
                    ipv6_octets[ip_index] =
                        NetworkEndian::read_u16(&packet_data[(ip_index * 2)..((ip_index + 1) * 2)]);

                }

                Ok(MDNSType::AAAA(Ipv6Addr::new(ipv6_octets[0],
                                                ipv6_octets[1],
                                                ipv6_octets[2],
                                                ipv6_octets[3],
                                                ipv6_octets[4],
                                                ipv6_octets[5],
                                                ipv6_octets[6],
                                                ipv6_octets[7])))
            }
            _ => Ok(MDNSType::UNKNOW(mdns_type_id)),
        }
    }
}

#[derive(Debug)]
struct MDNSData {
    id: u16,
    flags: u16,
    questions: Vec<MDNSQuestion>,
    answers: Vec<MDNSAnswer>,
}

#[derive(Debug)]
struct MDNSQuestion {
    name: String,
    rr_type_id: u16,
    rr_class: u16,
}

#[derive(Debug)]
struct MDNSAnswer {
    name: String,
    rr_type: MDNSType,
    rr_class: u16,
    ttl: u32,
}

#[inline]
fn label_to_string(packet: &[u8], offset: usize) -> (String, usize) {
    let string_size = packet[offset] as usize;

    let full_string = String::from_utf8_lossy(&packet[(offset + 1)..(offset + 1 + string_size)])
        .into_owned();

    (full_string, string_size + 1)
}

fn decompress_label(packet: &[u8], mut offset: usize) -> (String, usize) {
    // create the string
    let mut full_string = String::new();
    // second offset to let the caller know how much we moved
    // from the current offset
    let mut label_offset = 0;

    // helper that gets true as soon as we have hit a compress token
    // the label_offset calculation is different when we have have
    // a commpressed label
    let mut compressed = false;

    // Now go until we found a string with ZERO length
    while 0 != packet[offset] {
        // if this is not the first string we have found,
        // add a . to join them
        if full_string.len() > 0 {
            full_string.push('.');
        }

        trace!("size buffer: {:?}", &packet[offset..(offset + 2)]);

        // is this a compressed token?
        if 0xC0 == (packet[offset] & 0xC0) {
            // yes
            if false == compressed {
                compressed = true;
                // we have moved 2 bytes from the current position in memory
                label_offset += 2;
            }

            // now the second byte after the token
            // is the lower size of the uint16 offset
            // from the packet start!!!
            let small_size = packet[offset + 1] as usize;
            trace!("In BIG SIZE: {}", small_size);

            // get the higher size
            offset = (packet[offset] & 0x3F) as usize;
            trace!("In BIG SIZE_1: {}", offset);

            // shift 8 bits to the left
            offset *= 256;
            trace!("In BIG SIZE_2: {}", offset);

            // add the lower part
            offset += small_size;
        }

        // get size of the string
        let string_size = packet[offset] as usize;
        // jump over the size byte
        offset += 1;

        trace!("Before UTF8_2: {} {}", offset, string_size);
        trace!("Before UTF8: {:?}", &packet[offset..(offset + string_size)]);

        // now convert the bytes to a string
        // TODO: UTF8???? Or ASCII?
        full_string += std::str::from_utf8(&packet[offset..(offset + string_size)]).unwrap();

        // move to the next token
        offset += string_size;

        if false == compressed {
            // if we haven't had a compress token yet
            // just a the size byte and the string length
            label_offset += string_size;
            label_offset += 1;
        }
    }

    // skip the last ZERO byte indicating the end of the string
    if false == compressed {
        // but only if we have NO compressed string
        label_offset += 1;
    }

    trace!("{:?} {}", full_string, label_offset);

    // done, return the string and the offset we moved
    // in memory to get the String
    (full_string, label_offset)
}

fn parse_packet(packet: &[u8]) -> Result<MDNSData, String> {
    let decoded: MDNSPacketHeader = try!(decode(&packet[0..12]).map_err(|e| e.to_string()));

    let mut result = MDNSData {
        id: decoded.id,
        flags: decoded.flags,
        questions: vec![],
        answers: vec![],
    };

    let mut decode_position = std::mem::size_of::<MDNSPacketHeader>();

    // parse the questions
    for _ in 0..decoded.num_qn {
        // first read the label
        let (label_string, label_offset) = decompress_label(packet, decode_position);
        // move the offset
        decode_position += label_offset;

        // now read the uint16 type
        let rr_type = NetworkEndian::read_u16(&packet[decode_position..(decode_position + 2)]);
        decode_position += 2;
        // and the class (uint16)
        let class_type = NetworkEndian::read_u16(&packet[decode_position..(decode_position + 2)]);
        decode_position += 2;

        // Add the create and push the question struct
        result.questions.push(MDNSQuestion {
            name: label_string,
            rr_type_id: rr_type,
            rr_class: class_type,
        });
    }

    for _ in 0..decoded.num_ans_rr {
        // first read the label
        let (label_string, label_offset) = decompress_label(packet, decode_position);
        // move the offset
        decode_position += label_offset;

        // now read the uint16 type
        let rr_type = NetworkEndian::read_u16(&packet[decode_position..(decode_position + 2)]);
        decode_position += 2;
        // and the class (uint16)
        let class_type = NetworkEndian::read_u16(&packet[decode_position..(decode_position + 2)]);
        decode_position += 2;

        let ttl = NetworkEndian::read_u32(&packet[decode_position..(decode_position + 4)]);
        decode_position += 4;

        // RR data
        let data_len =
            NetworkEndian::read_u16(&packet[decode_position..(decode_position + 2)]) as usize;
        decode_position += 2;

        trace!("Data len: {}, type: {:x}", data_len, rr_type);

        match MDNSType::from_data(rr_type,
                                  &packet[decode_position..(decode_position + data_len)]) {
            Ok(mdns_type) => {

                // Add the create and push the question struct
                result.answers.push(MDNSAnswer {
                    name: label_string,
                    rr_type: mdns_type,
                    rr_class: class_type,
                    ttl: ttl,
                });
            }
            Err(e) => {
                warn!("{}", e);
            }
        }

        decode_position += data_len;

    }

    debug!("{:?}", decoded);

    Ok(result)
}

fn main() {
    env_logger::init().unwrap();

    let udp_builder = UdpBuilder::new_v4().unwrap();

    match udp_builder.reuse_address(true) {
        Ok(s) => s,
        Err(e) => panic!("couldn't reuse_address: {}", e),
    };

    match udp_builder.reuse_port(true) {
        Ok(s) => s,
        Err(e) => panic!("couldn't reuse_address: {}", e),
    };

    let udp_socket = match udp_builder.bind("224.0.0.251:5353") {
        Ok(s) => s,
        Err(e) => panic!("couldn't bind: {}", e),
    };

    let ip_mdns = Ipv4Addr::new(224, 0, 0, 251);
    let ip_any = Ipv4Addr::new(0, 0, 0, 0);

    match udp_socket.join_multicast_v4(&ip_mdns, &ip_any) {
        Ok(s) => s,
        Err(e) => panic!("couldn't join_multicast_v4: {}", e),
    };

    match udp_socket.set_multicast_loop_v4(true) {
        Ok(s) => s,
        Err(e) => panic!("couldn't set_multicast_loop_v4: {}", e),
    };

    info!("start reading");

    loop {
        // 	let mut read_packet: mdns_packet = Default::default();
        let mut byte_array = [0u8; 65536];

        let (amt, _) = udp_socket.recv_from(&mut byte_array).unwrap();
        // println!("{:?} {:?}", amt, src);

        if let Ok(mdns_data) = parse_packet(&byte_array[0..amt]) {
            debug!("{:?}", mdns_data);
        }
    }
}
<commit_msg>Added PTR parsing<commit_after>// disable warnings for dead code when running debug
#![cfg_attr(debug_assertions, allow(dead_code))]

extern crate net2;
extern crate bincode;
extern crate rustc_serialize;
#[macro_use]
extern crate log;
extern crate env_logger;
extern crate byteorder;

use std::net::{Ipv4Addr, Ipv6Addr};
#[cfg(not(windows))]
use net2::unix::UnixUdpBuilderExt;
use net2::UdpBuilder;
use bincode::rustc_serialize::decode;
use byteorder::{ByteOrder, NetworkEndian};

#[repr(C)]
#[repr(packed)]
#[derive(Default)]
#[derive(Debug)]
#[derive(RustcEncodable, RustcDecodable)]
struct MDNSPacketHeader {
    id: u16, // transaction ID
    flags: u16,
    num_qn: u16,
    num_ans_rr: u16,
    num_auth_rr: u16,
    num_add_rr: u16,
}

#[derive(Debug)]
enum MDNSType {
    A(Ipv4Addr),
    PTR(String),
    TXT(Vec<String>),
    AAAA(Ipv6Addr),
    // SRV,
    // NSEC,
    // ANY,
    UNKNOW(u16),
}

impl MDNSType {
    #[doc(hidden)]
    pub fn to_u16(&self) -> u16 {
        match *self {
            MDNSType::A(_) => 0x01,
            MDNSType::PTR(_) => 0x0C,
            MDNSType::TXT(_) => 0x10,
            MDNSType::AAAA(_) => 0x1C,
            // MDNSType::SRV => 0x21,
            // MDNSType::NSEC => 0x2F,
            // MDNSType::ANY => 0xFF,
            MDNSType::UNKNOW(n) => n,
        }
    }

    pub fn from_data(mdns_type_id: u16,
                     full_packet: &[u8],
                     data_len: usize,
                     offset: usize)
                     -> Result<MDNSType, String> {
        match mdns_type_id {
            0x10 => {
                let mut txt_map = vec![];
                let mut data_offset: usize = 0;
                while data_len > data_offset {
                    let (label_string, label_size) = label_to_string(full_packet,
                                                                     offset + data_offset);
                    trace!("TXT: {} {}", label_string, label_size);
                    txt_map.push(label_string);
                    data_offset += label_size;
                }
                trace!("{:?}", txt_map);
                Ok(MDNSType::TXT(txt_map))
            }
            0x01 => {
                if data_len < 4 {
                    return Err("Data section to small for IPV4 address".to_owned());
                }

                let ip = NetworkEndian::read_u32(&full_packet[offset..(offset + 4)]);
                let address = Ipv4Addr::new(((ip >> 24) & 0xFF) as u8,
                                            ((ip >> 16) & 0xFF) as u8,
                                            ((ip >> 8) & 0xFF) as u8,
                                            (ip & 0xFF) as u8);

                Ok(MDNSType::A(address))
            }
            0x1C => {
                if data_len < (8 * 2) {
                    return Err("Data section to small for IPV6 address".to_owned());
                }

                let mut ipv6_octets = [0u16; 8];
                for ip_index in 0..8 {
                    ipv6_octets[ip_index] =
                        NetworkEndian::read_u16(&full_packet[(offset+(ip_index * 2))..(offset +((ip_index + 1) * 2))]);

                }

                Ok(MDNSType::AAAA(Ipv6Addr::new(ipv6_octets[0],
                                                ipv6_octets[1],
                                                ipv6_octets[2],
                                                ipv6_octets[3],
                                                ipv6_octets[4],
                                                ipv6_octets[5],
                                                ipv6_octets[6],
                                                ipv6_octets[7])))
            }
            0x0C => Ok(MDNSType::PTR(decompress_label(full_packet, offset).0)),
            _ => Ok(MDNSType::UNKNOW(mdns_type_id)),
        }
    }
}

#[derive(Debug)]
struct MDNSData {
    id: u16,
    flags: u16,
    questions: Vec<MDNSQuestion>,
    answers: Vec<MDNSAnswer>,
}

#[derive(Debug)]
struct MDNSQuestion {
    name: String,
    rr_type_id: u16,
    rr_class: u16,
}

#[derive(Debug)]
struct MDNSAnswer {
    name: String,
    rr_type: MDNSType,
    rr_class: u16,
    ttl: u32,
}

#[inline]
fn label_to_string(packet: &[u8], offset: usize) -> (String, usize) {
    let string_size = packet[offset] as usize;

    let full_string = String::from_utf8_lossy(&packet[(offset + 1)..(offset + 1 + string_size)])
        .into_owned();

    (full_string, string_size + 1)
}

fn decompress_label(packet: &[u8], mut offset: usize) -> (String, usize) {
    // create the string
    let mut full_string = String::new();
    // second offset to let the caller know how much we moved
    // from the current offset
    let mut label_offset = 0;

    // helper that gets true as soon as we have hit a compress token
    // the label_offset calculation is different when we have have
    // a commpressed label
    let mut compressed = false;

    // Now go until we found a string with ZERO length
    while 0 != packet[offset] {
        // if this is not the first string we have found,
        // add a . to join them
        if full_string.len() > 0 {
            full_string.push('.');
        }

        trace!("size buffer: {:?}", &packet[offset..(offset + 2)]);

        // is this a compressed token?
        if 0xC0 == (packet[offset] & 0xC0) {
            // yes
            if false == compressed {
                compressed = true;
                // we have moved 2 bytes from the current position in memory
                label_offset += 2;
            }

            // now the second byte after the token
            // is the lower size of the uint16 offset
            // from the packet start!!!
            let small_size = packet[offset + 1] as usize;
            trace!("In BIG SIZE: {}", small_size);

            // get the higher size
            offset = (packet[offset] & 0x3F) as usize;
            trace!("In BIG SIZE_1: {}", offset);

            // shift 8 bits to the left
            offset *= 256;
            trace!("In BIG SIZE_2: {}", offset);

            // add the lower part
            offset += small_size;
        }

        // get size of the string
        let string_size = packet[offset] as usize;
        // jump over the size byte
        offset += 1;

        trace!("Before UTF8_2: {} {}", offset, string_size);
        trace!("Before UTF8: {:?}", &packet[offset..(offset + string_size)]);

        // now convert the bytes to a string
        // TODO: UTF8???? Or ASCII?
        full_string += std::str::from_utf8(&packet[offset..(offset + string_size)]).unwrap();

        // move to the next token
        offset += string_size;

        if false == compressed {
            // if we haven't had a compress token yet
            // just a the size byte and the string length
            label_offset += string_size;
            label_offset += 1;
        }
    }

    // skip the last ZERO byte indicating the end of the string
    if false == compressed {
        // but only if we have NO compressed string
        label_offset += 1;
    }

    trace!("{:?} {}", full_string, label_offset);

    // done, return the string and the offset we moved
    // in memory to get the String
    (full_string, label_offset)
}

fn parse_packet(packet: &[u8]) -> Result<MDNSData, String> {
    let decoded: MDNSPacketHeader = try!(decode(&packet[0..12]).map_err(|e| e.to_string()));

    let mut result = MDNSData {
        id: decoded.id,
        flags: decoded.flags,
        questions: vec![],
        answers: vec![],
    };

    let mut decode_position = std::mem::size_of::<MDNSPacketHeader>();

    // parse the questions
    for _ in 0..decoded.num_qn {
        // first read the label
        let (label_string, label_offset) = decompress_label(packet, decode_position);
        // move the offset
        decode_position += label_offset;

        // now read the uint16 type
        let rr_type = NetworkEndian::read_u16(&packet[decode_position..(decode_position + 2)]);
        decode_position += 2;
        // and the class (uint16)
        let class_type = NetworkEndian::read_u16(&packet[decode_position..(decode_position + 2)]);
        decode_position += 2;

        // Add the create and push the question struct
        result.questions.push(MDNSQuestion {
            name: label_string,
            rr_type_id: rr_type,
            rr_class: class_type,
        });
    }

    for _ in 0..decoded.num_ans_rr {
        // first read the label
        let (label_string, label_offset) = decompress_label(packet, decode_position);
        // move the offset
        decode_position += label_offset;

        // now read the uint16 type
        let rr_type = NetworkEndian::read_u16(&packet[decode_position..(decode_position + 2)]);
        decode_position += 2;
        // and the class (uint16)
        let class_type = NetworkEndian::read_u16(&packet[decode_position..(decode_position + 2)]);
        decode_position += 2;

        let ttl = NetworkEndian::read_u32(&packet[decode_position..(decode_position + 4)]);
        decode_position += 4;

        // RR data
        let data_len =
            NetworkEndian::read_u16(&packet[decode_position..(decode_position + 2)]) as usize;
        decode_position += 2;

        trace!("Data len: {}, type: {:x}", data_len, rr_type);

        match MDNSType::from_data(rr_type, packet, data_len, decode_position) {
            Ok(mdns_type) => {

                // Add the create and push the question struct
                result.answers.push(MDNSAnswer {
                    name: label_string,
                    rr_type: mdns_type,
                    rr_class: class_type,
                    ttl: ttl,
                });
            }
            Err(e) => {
                warn!("{}", e);
            }
        }

        decode_position += data_len;

    }

    debug!("{:?}", decoded);

    Ok(result)
}

fn main() {
    env_logger::init().unwrap();

    let udp_builder = UdpBuilder::new_v4().unwrap();

    match udp_builder.reuse_address(true) {
        Ok(s) => s,
        Err(e) => panic!("couldn't reuse_address: {}", e),
    };

    match udp_builder.reuse_port(true) {
        Ok(s) => s,
        Err(e) => panic!("couldn't reuse_address: {}", e),
    };

    let udp_socket = match udp_builder.bind("224.0.0.251:5353") {
        Ok(s) => s,
        Err(e) => panic!("couldn't bind: {}", e),
    };

    let ip_mdns = Ipv4Addr::new(224, 0, 0, 251);
    let ip_any = Ipv4Addr::new(0, 0, 0, 0);

    match udp_socket.join_multicast_v4(&ip_mdns, &ip_any) {
        Ok(s) => s,
        Err(e) => panic!("couldn't join_multicast_v4: {}", e),
    };

    match udp_socket.set_multicast_loop_v4(true) {
        Ok(s) => s,
        Err(e) => panic!("couldn't set_multicast_loop_v4: {}", e),
    };

    info!("start reading");

    loop {
        // 	let mut read_packet: mdns_packet = Default::default();
        let mut byte_array = [0u8; 65536];

        let (amt, _) = udp_socket.recv_from(&mut byte_array).unwrap();
        // println!("{:?} {:?}", amt, src);

        if let Ok(mdns_data) = parse_packet(&byte_array[0..amt]) {
            debug!("{:?}", mdns_data);
        }
    }
}
<|endoftext|>


<commit_before>
mod curses;

use std::ffi::CString;

use curses::*;

enum DrawMe {
    //Circle { radius : i32, x : i32, y : i32 },
    Rec { width : i32, height : i32, x : i32, y : i32 },
    /*VHall { length : i32, x : i32, y : i32 },
    HHall { length : i32, x : i32, y : i32 },*/
    // TODO add color
    // TODO add background color
    // TODO add point case
    // TODO add overlap
    // TODO can probably have a struct represent common items (x,y, color, etc)
}

unsafe fn render_shape( s : &DrawMe ) {
    match s {
        &DrawMe::Rec { width, height, x, y } => render_rec( width, height, x, y ),
        _ => (),
    }
}

unsafe fn render_rec( w : i32, h : i32, x : i32, y : i32 ) {
    let dot = CString::new( "." ).unwrap();
    let h_line = CString::new( "-" ).unwrap();
    let v_line = CString::new( "|" ).unwrap();
    let plus = CString::new( "+" ).unwrap();

    let mut h2 = h + 2;
    while h2 > 0 {
        let d : &CString = if h2 == h + 2 {
            &h_line
        } else if h2 == 1 {
            &h_line
        } else {
            &dot
        };
        let mut w2 = w;

        let other_d : &CString = if h2 == h + 2 {
            &plus
        } else if h2 == 1 {
            &plus
        } else {
            &v_line
        };
        mvprintw( y + h2, x + w2 + 1, other_d.as_ptr() );
        while w2 > 0 {
            mvprintw( y + h2, x + w2, d.as_ptr() );
            w2 -= 1;
        }
        mvprintw( y + h2, x, other_d.as_ptr() );
        h2 -= 1;
    }
}

fn main() {
    let mut x = 0;
    let mut y = 0;
    unsafe {
        initscr();
        cbreak();
        noecho();

        let mut c = getch();
        let mut shapes : Vec<DrawMe> = vec! [];

        while c as u8 as char != 'q' {
            let c2 = c as u8 as char;

            match c2 {
                /*'c' => {
                    shapes.push( DrawMe::Circle { radius: 5, x: x, y: y } );
                },*/
                'r' => {
                    shapes.push( DrawMe::Rec { width: 5, height: 8, x: x, y: y } );
                },
                'j' => {
                    y+=1;
                },
                'k' => {
                    y-=1;
                },
                'l' => {
                    x+=1;
                },
                'h' => {
                    x-=1;
                },
                _ => (),
            }

            for s in &shapes {
                render_shape( s );
            }

            cursor_move( y, x );
      //      clear();
            refresh();

            c = getch();
        }

        endwin();
    }
}

<commit_msg>TODOs<commit_after>
mod curses;

use std::ffi::CString;

use curses::*;

enum DrawMe {
    Rec { width : i32, height : i32, x : i32, y : i32 },
    /*VHall { length : i32, x : i32, y : i32 },
    HHall { length : i32, x : i32, y : i32 },
    Character
    VWall
    HWall
    River
    Lava*/
    // TODO add color
    // TODO add background color
    // TODO can probably have a struct represent common items (x,y, color, etc)
}

unsafe fn render_shape( s : &DrawMe ) {
    match s {
        &DrawMe::Rec { width, height, x, y } => render_rec( width, height, x, y ),
        _ => (),
    }
}

unsafe fn render_rec( w : i32, h : i32, x : i32, y : i32 ) {
    let dot = CString::new( "." ).unwrap();
    let h_line = CString::new( "-" ).unwrap();
    let v_line = CString::new( "|" ).unwrap();
    let plus = CString::new( "+" ).unwrap();

    let mut h2 = h + 2;
    while h2 > 0 {
        let d : &CString = if h2 == h + 2 {
            &h_line
        } else if h2 == 1 {
            &h_line
        } else {
            &dot
        };
        let mut w2 = w;

        let other_d : &CString = if h2 == h + 2 {
            &plus
        } else if h2 == 1 {
            &plus
        } else {
            &v_line
        };
        mvprintw( y + h2, x + w2 + 1, other_d.as_ptr() );
        while w2 > 0 {
            mvprintw( y + h2, x + w2, d.as_ptr() );
            w2 -= 1;
        }
        mvprintw( y + h2, x, other_d.as_ptr() );
        h2 -= 1;
    }
}

fn main() {
    let mut x = 0;
    let mut y = 0;
    unsafe {
        initscr();
        cbreak();
        noecho();

        let mut c = getch();
        let mut shapes : Vec<DrawMe> = vec! [];

        // TODO need mode for entering rech, etc
        // TODO need mode for changing color

        while c as u8 as char != 'q' {
            let c2 = c as u8 as char;

            match c2 {
                /*'c' => {
                    shapes.push( DrawMe::Circle { radius: 5, x: x, y: y } );
                },*/
                'r' => {
                    shapes.push( DrawMe::Rec { width: 5, height: 8, x: x, y: y } );
                },
                'j' => {
                    y+=1;
                },
                'k' => {
                    y-=1;
                },
                'l' => {
                    x+=1;
                },
                'h' => {
                    x-=1;
                },
                _ => (),
            }

            for s in &shapes {
                render_shape( s );
            }

            cursor_move( y, x );
            refresh();

            c = getch();
        }

        endwin();
    }
}

<|endoftext|>


<commit_before>// Copyright 2017 Jeremy Wall <jeremy@marzhillstudios.com>
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
#[macro_use]
extern crate clap;
extern crate ucglib;

use std::fs::File;
use std::rc::Rc;
use std::io;
use std::process;

use ucglib::build::Val;
use ucglib::build;
use ucglib::convert::ConverterRunner;

// TODO(jwall): List the target output types automatically.
fn do_flags<'a>() -> clap::ArgMatches<'a> {
    clap_app!(
        ucg =>
            (version: crate_version!())
            (author: crate_authors!())
            (about: "Universal Configuration Grammar compiler.")
            (@subcommand build =>
             (about: "Compile a specific ucg file.")
             (@arg sym: --sym +takes_value "Specify a specific let binding in the ucg file to output.")
             (@arg target: --target -t +required +takes_value "Target output type. (flags, json, env)")
             (@arg out: --out -o +takes_value "Output file to write to.")
             (@arg INPUT: +required "Input ucg file to build.")
            )
            (@subcommand validate =>
             (about: "Check a specific ucg file for errors.")
             (@arg INPUT: +required "Input ucg file to validate.")
            )
    ).get_matches()
}

fn run_converter(c: ConverterRunner, v: Rc<Val>, f: &str) -> io::Result<()> {
    let file = File::create(f);
    c.convert(v, Box::new(file.unwrap()))
}

fn main() {
    // TODO(jwall): Read and build an actual file.
    let app = do_flags();
    if let Some(matches) = app.subcommand_matches("build") {
        let file = matches.value_of("INPUT").unwrap();
        let out = matches.value_of("out").unwrap();
        let sym = matches.value_of("sym");
        let target = matches.value_of("target").unwrap();
        let mut builder = build::Builder::new();
        match ConverterRunner::new(target) {
            Ok(converter) => {
                let result = builder.build_file(file);
                if !result.is_ok() {
                    eprintln!("{:?}", result.err());
                    process::exit(1);
                }
                let val = match sym {
                    Some(sym_name) => builder.get_out_by_name(sym_name),
                    None => builder.last,
                };
                match val {
                    Some(value) => {
                        run_converter(converter, value, out).unwrap();
                        println!("Build successful");
                        process::exit(0);
                    }
                    None => {
                        eprintln!("Build results in no value.");
                        process::exit(1);
                    }
                }
            }
            Err(msg) => {
                eprintln!("{}", msg);
                process::exit(1);
            }
        }
    } else if let Some(matches) = app.subcommand_matches("validate") {
        let file = matches.value_of("INPUT").unwrap();
        let mut builder = build::Builder::new();
        builder.build_file(file).unwrap();
        println!("File Validates");
        process::exit(0);
    }
}
<commit_msg>Make the out flag optional<commit_after>// Copyright 2017 Jeremy Wall <jeremy@marzhillstudios.com>
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
#[macro_use]
extern crate clap;
extern crate ucglib;

use std::fs::File;
use std::rc::Rc;
use std::io;
use std::process;

use ucglib::build::Val;
use ucglib::build;
use ucglib::convert::ConverterRunner;

// TODO(jwall): List the target output types automatically.
fn do_flags<'a>() -> clap::ArgMatches<'a> {
    clap_app!(
        ucg =>
            (version: crate_version!())
            (author: crate_authors!())
            (about: "Universal Configuration Grammar compiler.")
            (@subcommand build =>
             (about: "Compile a specific ucg file.")
             (@arg sym: --sym +takes_value +required "Specify a specific let binding in the ucg file to output.")
             (@arg target: --target -t +required +takes_value "Target output type. (flags, json, env)")
             (@arg out: --out -o +takes_value "Output file to write to.")
             (@arg INPUT: +required "Input ucg file to build.")
            )
            (@subcommand validate =>
             (about: "Check a specific ucg file for errors.")
             (@arg INPUT: +required "Input ucg file to validate.")
            )
    ).get_matches()
}

fn run_converter(c: ConverterRunner, v: Rc<Val>, f: Option<&str>) -> io::Result<()> {
    let file: Box<std::io::Write> = match f {
        Some(f) => Box::new(try!(File::create(f))),
        None => Box::new(io::stdout()),
    };
    c.convert(v, file)
}

fn main() {
    // TODO(jwall): Read and build an actual file.
    let app = do_flags();
    if let Some(matches) = app.subcommand_matches("build") {
        let file = matches.value_of("INPUT").unwrap();
        let out = matches.value_of("out");
        let sym = matches.value_of("sym");
        let target = matches.value_of("target").unwrap();
        let mut builder = build::Builder::new();
        match ConverterRunner::new(target) {
            Ok(converter) => {
                let result = builder.build_file(file);
                if !result.is_ok() {
                    eprintln!("{:?}", result.err());
                    process::exit(1);
                }
                let val = match sym {
                    Some(sym_name) => builder.get_out_by_name(sym_name),
                    None => builder.last,
                };
                match val {
                    Some(value) => {
                        run_converter(converter, value, out).unwrap();
                        eprintln!("Build successful");
                        process::exit(0);
                    }
                    None => {
                        eprintln!("Build results in no value.");
                        process::exit(1);
                    }
                }
            }
            Err(msg) => {
                eprintln!("{}", msg);
                process::exit(1);
            }
        }
    } else if let Some(matches) = app.subcommand_matches("validate") {
        let file = matches.value_of("INPUT").unwrap();
        let mut builder = build::Builder::new();
        builder.build_file(file).unwrap();
        println!("File Validates");
        process::exit(0);
    }
}
<|endoftext|>