<commit_before>// Copyright (C) 2020 Jason Ish
//
// Permission is hereby granted, free of charge, to any person obtaining
// a copy of this software and associated documentation files (the
// "Software"), to deal in the Software without restriction, including
// without limitation the rights to use, copy, modify, merge, publish,
// distribute, sublicense, and/or sell copies of the Software, and to
// permit persons to whom the Software is furnished to do so, subject to
// the following conditions:
//
// The above copyright notice and this permission notice shall be
// included in all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
// LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
// WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

use crate::prelude::*;
use axum::body::Bytes;
use axum::extract::{ContentLengthLimit, Extension};
use axum::http::StatusCode;
use axum::response::IntoResponse;
use axum::Json;
use serde_json::json;
use std::io::BufRead;
use std::sync::Arc;

use crate::eve::eve::EveJson;
use crate::server::ServerContext;

pub(crate) async fn handler_new(
    Extension(context): Extension<Arc<ServerContext>>,
    ContentLengthLimit(body): ContentLengthLimit<Bytes, { 1024 * 5_000 }>, // ~5mb
) -> impl IntoResponse {
    let mut importer = match context.datastore.get_importer() {
        Some(importer) => importer,
        None => {
            return (StatusCode::NOT_IMPLEMENTED, "").into_response();
        }
    };
    let mut errors = Vec::new();

    let mut buf = &body[..];
    let mut count = 0;
    let mut line = String::new();
    loop {
        match buf.read_line(&mut line) {
            Err(err) => {
                errors.push(format!("Failed to read event from request body: {}", err));
                // Failed to read line, can't continue.
                break;
            }
            Ok(n) => {
                if n == 0 {
                    // EOF.
                    break;
                }
                match serde_json::from_str::<EveJson>(&line) {
                    Err(err) => {
                        errors.push(format!(
                            "Failed to decode event from request body ({}): {}",
                            err, line
                        ));
                    }
                    Ok(event) => {
                        count += 1;
                        if let Err(err) = importer.submit(event).await {
                            error!("Failed to submit event to importer: {}", err);
                        }
                    }
                }
            }
        }
        line.truncate(0);
    }

    // I've seen an issue in the Go agent where it sent 0 events, return early if we have
    // nothing to commit.
    if count == 0 {
        // TODO: Log something or return an error to the client.
        let response = json!({
            "Count": 0,
        });
        return Json(response).into_response();
    }

    match importer.commit().await {
        Ok(n) => {
            debug!("Committed {} events (received {})", n, count);
            let response = json!({
                // Kept capitolized for compatibility with the Go agent.
                "Count": n,
            });
            return Json(response).into_response();
        }
        Err(err) => {
            error!("Failed to commit events (received {}): {}", count, err);
            return (StatusCode::INTERNAL_SERVER_ERROR, "").into_response();
        }
    }
}
<commit_msg>submit api: increase content length to 256mb<commit_after>// Copyright (C) 2020 Jason Ish
//
// Permission is hereby granted, free of charge, to any person obtaining
// a copy of this software and associated documentation files (the
// "Software"), to deal in the Software without restriction, including
// without limitation the rights to use, copy, modify, merge, publish,
// distribute, sublicense, and/or sell copies of the Software, and to
// permit persons to whom the Software is furnished to do so, subject to
// the following conditions:
//
// The above copyright notice and this permission notice shall be
// included in all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
// LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
// WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

use crate::prelude::*;
use axum::body::Bytes;
use axum::extract::{ContentLengthLimit, Extension};
use axum::http::StatusCode;
use axum::response::IntoResponse;
use axum::Json;
use serde_json::json;
use std::io::BufRead;
use std::sync::Arc;

use crate::eve::eve::EveJson;
use crate::server::ServerContext;

pub(crate) async fn handler_new(
    Extension(context): Extension<Arc<ServerContext>>,
    ContentLengthLimit(body): ContentLengthLimit<Bytes, { 1024 * 1024 * 256 }>,
) -> impl IntoResponse {
    let mut importer = match context.datastore.get_importer() {
        Some(importer) => importer,
        None => {
            return (StatusCode::NOT_IMPLEMENTED, "").into_response();
        }
    };
    let mut errors = Vec::new();

    let mut buf = &body[..];
    let mut count = 0;
    let mut line = String::new();
    loop {
        match buf.read_line(&mut line) {
            Err(err) => {
                errors.push(format!("Failed to read event from request body: {}", err));
                // Failed to read line, can't continue.
                break;
            }
            Ok(n) => {
                if n == 0 {
                    // EOF.
                    break;
                }
                match serde_json::from_str::<EveJson>(&line) {
                    Err(err) => {
                        errors.push(format!(
                            "Failed to decode event from request body ({}): {}",
                            err, line
                        ));
                    }
                    Ok(event) => {
                        count += 1;
                        if let Err(err) = importer.submit(event).await {
                            error!("Failed to submit event to importer: {}", err);
                        }
                    }
                }
            }
        }
        line.truncate(0);
    }

    // I've seen an issue in the Go agent where it sent 0 events, return early if we have
    // nothing to commit.
    if count == 0 {
        // TODO: Log something or return an error to the client.
        let response = json!({
            "Count": 0,
        });
        return Json(response).into_response();
    }

    match importer.commit().await {
        Ok(n) => {
            debug!("Committed {} events (received {})", n, count);
            let response = json!({
                // Kept capitolized for compatibility with the Go agent.
                "Count": n,
            });
            return Json(response).into_response();
        }
        Err(err) => {
            error!("Failed to commit events (received {}): {}", count, err);
            return (StatusCode::INTERNAL_SERVER_ERROR, "").into_response();
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2020 The Tink-Rust Authors
//
// 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.
//
////////////////////////////////////////////////////////////////////////////////

use tink_core::TinkError;
use tink_proto::{EcdsaSignatureEncoding, EllipticCurveType, HashType};

/// Supported signature encodings.  This is a precise subset of the protobuf enum,
/// allowing exact `match`es.
#[derive(Clone, Debug)]
pub enum SignatureEncoding {
    Der,
    IeeeP1363,
}

/// Validate ECDSA parameters.
/// The hash's strength must not be weaker than the curve's strength.
/// Only DER encoding is supported now.
pub fn validate_ecdsa_params(
    hash_alg: tink_proto::HashType,
    curve: tink_proto::EllipticCurveType,
    encoding: tink_proto::EcdsaSignatureEncoding,
) -> Result<SignatureEncoding, TinkError> {
    let encoding = match encoding {
        EcdsaSignatureEncoding::IeeeP1363 => SignatureEncoding::IeeeP1363,
        EcdsaSignatureEncoding::Der => SignatureEncoding::Der,
        _ => return Err("ecdsa: unsupported encoding".into()),
    };
    match curve {
        EllipticCurveType::NistP256 => {
            if hash_alg != HashType::Sha256 {
                return Err("invalid hash type, expect SHA-256".into());
            }
        }
        EllipticCurveType::NistP384 => {
            if hash_alg != HashType::Sha384 && hash_alg != HashType::Sha512 {
                return Err("invalid hash type, expect SHA-384 or SHA-512".into());
            }
        }
        EllipticCurveType::NistP521 => {
            if hash_alg != HashType::Sha512 {
                return Err("invalid hash type, expect SHA-512".into());
            }
        }
        _ => return Err(format!("unsupported curve: {:?}", curve).into()),
    }
    Ok(encoding)
}
<commit_msg>tink-signature: fix comment<commit_after>// Copyright 2020 The Tink-Rust Authors
//
// 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.
//
////////////////////////////////////////////////////////////////////////////////

use tink_core::TinkError;
use tink_proto::{EcdsaSignatureEncoding, EllipticCurveType, HashType};

/// Supported signature encodings.  This is a precise subset of the protobuf enum,
/// allowing exact `match`es.
#[derive(Clone, Debug)]
pub enum SignatureEncoding {
    Der,
    IeeeP1363,
}

/// Validate ECDSA parameters.
/// The hash's strength must not be weaker than the curve's strength.
/// DER and IEEE_P1363 encodings are supported.
pub fn validate_ecdsa_params(
    hash_alg: tink_proto::HashType,
    curve: tink_proto::EllipticCurveType,
    encoding: tink_proto::EcdsaSignatureEncoding,
) -> Result<SignatureEncoding, TinkError> {
    let encoding = match encoding {
        EcdsaSignatureEncoding::IeeeP1363 => SignatureEncoding::IeeeP1363,
        EcdsaSignatureEncoding::Der => SignatureEncoding::Der,
        _ => return Err("ecdsa: unsupported encoding".into()),
    };
    match curve {
        EllipticCurveType::NistP256 => {
            if hash_alg != HashType::Sha256 {
                return Err("invalid hash type, expect SHA-256".into());
            }
        }
        EllipticCurveType::NistP384 => {
            if hash_alg != HashType::Sha384 && hash_alg != HashType::Sha512 {
                return Err("invalid hash type, expect SHA-384 or SHA-512".into());
            }
        }
        EllipticCurveType::NistP521 => {
            if hash_alg != HashType::Sha512 {
                return Err("invalid hash type, expect SHA-512".into());
            }
        }
        _ => return Err(format!("unsupported curve: {:?}", curve).into()),
    }
    Ok(encoding)
}
<|endoftext|>


<commit_before>use std::mem;
use syntax::ast::{Expr, ExprParen, Lit, Stmt, TokenTree, TtDelimited, TtToken};
use syntax::ext::quote::rt::ToTokens;
use syntax::codemap::Span;
use syntax::ext::base::ExtCtxt;
use syntax::parse;
use syntax::parse::parser::Parser as RustParser;
use syntax::parse::token::{self, DelimToken};
use syntax::ptr::P;

use super::render::{Escape, Renderer};

macro_rules! dollar {
    () => (TtToken(_, token::Dollar))
}
macro_rules! dot {
    () => (TtToken(_, token::Dot))
}
macro_rules! eq {
    () => (TtToken(_, token::Eq))
}
macro_rules! not {
    () => (TtToken(_, token::Not))
}
macro_rules! question {
    () => (TtToken(_, token::Question))
}
macro_rules! semi {
    () => (TtToken(_, token::Semi))
}
macro_rules! minus {
    () => (TtToken(_, token::BinOp(token::Minus)))
}
macro_rules! slash {
    () => (TtToken(_, token::BinOp(token::Slash)))
}
macro_rules! literal {
    () => (TtToken(_, token::Literal(..)))
}
macro_rules! ident {
    ($x:pat) => (ident!(_, $x));
    ($sp:pat, $x:pat) => (TtToken($sp, token::Ident($x, token::IdentStyle::Plain)))
}

pub fn parse(cx: &ExtCtxt, input: &[TokenTree], sp: Span) -> P<Expr> {
    let mut parser = Parser {
        in_attr: false,
        input: input,
        span: sp,
        render: Renderer::new(cx),
    };
    parser.markups();
    parser.into_render().into_expr()
}

struct Parser<'cx, 's: 'cx, 'i> {
    in_attr: bool,
    input: &'i [TokenTree],
    span: Span,
    render: Renderer<'cx, 's>,
}

impl<'cx, 's, 'i> Parser<'cx, 's, 'i> {
    /// Finalize the `Parser`, returning the `Renderer` underneath.
    fn into_render(self) -> Renderer<'cx, 's> {
        let Parser { render, .. } = self;
        render
    }

    /// Consume `n` items from the input.
    fn shift(&mut self, n: usize) {
        self.input = &self.input[n..];
    }

    /// Construct a Rust AST parser from the given token tree.
    fn new_rust_parser(&self, tts: Vec<TokenTree>) -> RustParser<'s> {
        parse::tts_to_parser(self.render.cx.parse_sess, tts, self.render.cx.cfg.clone())
    }

    fn markups(&mut self) {
        loop {
            match self.input {
                [] => return,
                [semi!(), ..] => self.shift(1),
                [_, ..] => if !self.markup() { return },
            }
        }
    }

    fn markup(&mut self) -> bool {
        match self.input {
            // Literal
            [minus!(), ref tt @ literal!(), ..] => {
                self.shift(2);
                self.literal(tt, true)
            },
            [ref tt @ literal!(), ..] => {
                self.shift(1);
                self.literal(tt, false)
            },
            // If
            [dollar!(), ident!(sp, name), ..] if name.as_str() == "if" => {
                self.shift(2);
                self.if_expr(sp);
            },
            // For
            [dollar!(), ident!(sp, name), ..] if name.as_str() == "for" => {
                self.shift(2);
                self.for_expr(sp);
            },
            // Splice
            [ref tt @ dollar!(), dollar!(), ..] => {
                self.shift(2);
                let expr = self.splice(tt.get_span());
                self.render.splice(expr, Escape::PassThru);
            },
            [ref tt @ dollar!(), ..] => {
                self.shift(1);
                let expr = self.splice(tt.get_span());
                self.render.splice(expr, Escape::Escape);
            },
            // Element
            [ident!(sp, name), ..] => {
                self.shift(1);
                self.element(sp, name.as_str())
            },
            // Block
            [TtDelimited(sp, ref d), ..] if d.delim == token::DelimToken::Brace => {
                self.shift(1);
                let stmts = self.block(sp, &d.tts);
                self.render.push_stmts(stmts);
            },
            // ???
            _ => {
                if let [ref tt, ..] = self.input {
                    self.render.cx.span_err(tt.get_span(), "invalid syntax");
                } else {
                    self.render.cx.span_err(self.span, "unexpected end of block");
                }
                return false;
            },
        }
        true
    }

    fn literal(&mut self, tt: &TokenTree, minus: bool) {
        let lit = self.new_rust_parser(vec![tt.clone()]).parse_lit();
        match lit_to_string(self.render.cx, lit, minus) {
            Some(s) => self.render.string(&s, Escape::Escape),
            None => {},
        }
    }

    fn if_expr(&mut self, sp: Span) {
        // Parse the initial if
        let mut if_cond = vec![];
        let if_body;
        loop { match self.input {
            [TtDelimited(sp, ref d), ..] if d.delim == DelimToken::Brace => {
                self.shift(1);
                if_body = self.block(sp, &d.tts);
                break;
            },
            [ref tt, ..] => {
                self.shift(1);
                if_cond.push(tt.clone());
            },
            [] => self.render.cx.span_fatal(sp, "expected body for this $if"),
        }}
        // Parse the (optional) else
        let else_body = match self.input {
            [dollar!(), ident!(else_), ..] if else_.as_str() == "else" => {
                self.shift(2);
                match self.input {
                    [ident!(sp, if_), ..] if if_.as_str() == "if" => {
                        self.shift(1);
                        let else_body = {
                            // Parse an if expression, but capture the result
                            // rather than emitting it right away
                            let mut render = self.render.fork();
                            mem::swap(&mut self.render, &mut render);
                            self.if_expr(sp);
                            mem::swap(&mut self.render, &mut render);
                            render.into_stmts()
                        };
                        Some(else_body)
                    },
                    [TtDelimited(sp, ref d), ..] if d.delim == DelimToken::Brace => {
                        self.shift(1);
                        Some(self.block(sp, &d.tts))
                    },
                    _ => self.render.cx.span_fatal(sp, "expected body for this $else"),
                }
            },
            _ => None,
        };
        self.render.emit_if(if_cond, if_body, else_body);
    }

    fn for_expr(&mut self, sp: Span) {
        let mut pattern = vec![];
        loop { match self.input {
            [ident!(in_), ..] if in_.as_str() == "in" => {
                self.shift(1);
                break;
            },
            [ref tt, ..] => {
                self.shift(1);
                pattern.push(tt.clone());
            },
            _ => self.render.cx.span_fatal(sp, "invalid $for"),
        }}
        let pattern = self.new_rust_parser(pattern).parse_pat();
        let mut iterable = vec![];
        let body;
        loop { match self.input {
            [TtDelimited(sp, ref d), ..] if d.delim == DelimToken::Brace => {
                self.shift(1);
                body = self.block(sp, &d.tts);
                break;
            },
            [ref tt, ..] => {
                self.shift(1);
                iterable.push(tt.clone());
            },
            _ => self.render.cx.span_fatal(sp, "invalid $for"),
        }}
        let iterable = self.new_rust_parser(iterable).parse_expr();
        self.render.emit_for(pattern, iterable, body);
    }

    fn splice(&mut self, sp: Span) -> P<Expr> {
        let mut tts = vec![];
        // First, munch a single token tree
        if let [ref tt, ..] = self.input {
            self.shift(1);
            tts.push(tt.clone());
        }
        loop { match self.input {
            // Munch attribute lookups e.g. `$person.address.street`
            [ref dot @ dot!(), ref ident @ ident!(_), ..] => {
                self.shift(2);
                tts.push(dot.clone());
                tts.push(ident.clone());
            },
            // Munch function calls `()` and indexing operations `[]`
            [TtDelimited(sp, ref d), ..] if d.delim != token::DelimToken::Brace => {
                self.shift(1);
                tts.push(TtDelimited(sp, d.clone()));
            },
            _ => break,
        }}
        if tts.is_empty() {
            self.render.cx.span_fatal(sp, "expected expression for this splice");
        } else {
            self.new_rust_parser(tts).parse_expr()
        }
    }

    fn element(&mut self, sp: Span, name: &str) {
        if self.in_attr {
            self.render.cx.span_err(sp, "unexpected element, you silly bumpkin");
            return;
        }
        self.render.element_open_start(name);
        self.attrs();
        self.render.element_open_end();
        if let [slash!(), ..] = self.input {
            self.shift(1);
        } else {
            self.markup();
            self.render.element_close(name);
        }
    }

    fn attrs(&mut self) {
        loop { match self.input {
            [ident!(name), eq!(), ..] => {
                // Non-empty attribute
                self.shift(2);
                self.render.attribute_start(name.as_str());
                {
                    // Parse a value under an attribute context
                    let mut in_attr = true;
                    mem::swap(&mut self.in_attr, &mut in_attr);
                    self.markup();
                    mem::swap(&mut self.in_attr, &mut in_attr);
                }
                self.render.attribute_end();
            },
            [ident!(name), question!(), ..] => {
                // Empty attribute
                self.shift(2);
                if let [ref tt @ eq!(), ..] = self.input {
                    // Toggle the attribute based on a boolean expression
                    self.shift(1);
                    let cond = self.splice(tt.get_span());
                    // Silence "unnecessary parentheses" warnings
                    let cond = strip_outer_parens(cond).to_tokens(self.render.cx);
                    let body = {
                        let mut r = self.render.fork();
                        r.attribute_empty(name.as_str());
                        r.into_stmts()
                    };
                    self.render.emit_if(cond, body, None);
                } else {
                    // Write the attribute unconditionally
                    self.render.attribute_empty(name.as_str());
                }
            },
            _ => return,
        }}
    }

    fn block(&mut self, sp: Span, tts: &[TokenTree]) -> Vec<P<Stmt>> {
        let mut parse = Parser {
            in_attr: self.in_attr,
            input: tts,
            span: sp,
            render: self.render.fork(),
        };
        parse.markups();
        parse.into_render().into_stmts()
    }
}

/// Convert a literal to a string.
fn lit_to_string(cx: &ExtCtxt, lit: Lit, minus: bool) -> Option<String> {
    use syntax::ast::Lit_::*;
    let mut result = String::new();
    if minus {
        result.push('-');
    }
    match lit.node {
        LitStr(s, _) => result.push_str(&s),
        LitBinary(..) | LitByte(..) => {
            cx.span_err(lit.span, "cannot splice binary data");
            return None;
        },
        LitChar(c) => result.push(c),
        LitInt(x, _) => result.push_str(&x.to_string()),
        LitFloat(s, _) | LitFloatUnsuffixed(s) => result.push_str(&s),
        LitBool(b) => result.push_str(if b { "true" } else { "false" }),
    };
    Some(result)
}

/// If the expression is wrapped in parentheses, strip them off.
fn strip_outer_parens(expr: P<Expr>) -> P<Expr> {
    expr.and_then(|expr| match expr {
        Expr { node: ExprParen(inner), .. } => inner,
        expr => P(expr),
    })
}
<commit_msg>Add some missing semicolons<commit_after>use std::mem;
use syntax::ast::{Expr, ExprParen, Lit, Stmt, TokenTree, TtDelimited, TtToken};
use syntax::ext::quote::rt::ToTokens;
use syntax::codemap::Span;
use syntax::ext::base::ExtCtxt;
use syntax::parse;
use syntax::parse::parser::Parser as RustParser;
use syntax::parse::token::{self, DelimToken};
use syntax::ptr::P;

use super::render::{Escape, Renderer};

macro_rules! dollar {
    () => (TtToken(_, token::Dollar))
}
macro_rules! dot {
    () => (TtToken(_, token::Dot))
}
macro_rules! eq {
    () => (TtToken(_, token::Eq))
}
macro_rules! not {
    () => (TtToken(_, token::Not))
}
macro_rules! question {
    () => (TtToken(_, token::Question))
}
macro_rules! semi {
    () => (TtToken(_, token::Semi))
}
macro_rules! minus {
    () => (TtToken(_, token::BinOp(token::Minus)))
}
macro_rules! slash {
    () => (TtToken(_, token::BinOp(token::Slash)))
}
macro_rules! literal {
    () => (TtToken(_, token::Literal(..)))
}
macro_rules! ident {
    ($x:pat) => (ident!(_, $x));
    ($sp:pat, $x:pat) => (TtToken($sp, token::Ident($x, token::IdentStyle::Plain)))
}

pub fn parse(cx: &ExtCtxt, input: &[TokenTree], sp: Span) -> P<Expr> {
    let mut parser = Parser {
        in_attr: false,
        input: input,
        span: sp,
        render: Renderer::new(cx),
    };
    parser.markups();
    parser.into_render().into_expr()
}

struct Parser<'cx, 's: 'cx, 'i> {
    in_attr: bool,
    input: &'i [TokenTree],
    span: Span,
    render: Renderer<'cx, 's>,
}

impl<'cx, 's, 'i> Parser<'cx, 's, 'i> {
    /// Finalize the `Parser`, returning the `Renderer` underneath.
    fn into_render(self) -> Renderer<'cx, 's> {
        let Parser { render, .. } = self;
        render
    }

    /// Consume `n` items from the input.
    fn shift(&mut self, n: usize) {
        self.input = &self.input[n..];
    }

    /// Construct a Rust AST parser from the given token tree.
    fn new_rust_parser(&self, tts: Vec<TokenTree>) -> RustParser<'s> {
        parse::tts_to_parser(self.render.cx.parse_sess, tts, self.render.cx.cfg.clone())
    }

    fn markups(&mut self) {
        loop {
            match self.input {
                [] => return,
                [semi!(), ..] => self.shift(1),
                [_, ..] => if !self.markup() { return },
            }
        }
    }

    fn markup(&mut self) -> bool {
        match self.input {
            // Literal
            [minus!(), ref tt @ literal!(), ..] => {
                self.shift(2);
                self.literal(tt, true);
            },
            [ref tt @ literal!(), ..] => {
                self.shift(1);
                self.literal(tt, false);
            },
            // If
            [dollar!(), ident!(sp, name), ..] if name.as_str() == "if" => {
                self.shift(2);
                self.if_expr(sp);
            },
            // For
            [dollar!(), ident!(sp, name), ..] if name.as_str() == "for" => {
                self.shift(2);
                self.for_expr(sp);
            },
            // Splice
            [ref tt @ dollar!(), dollar!(), ..] => {
                self.shift(2);
                let expr = self.splice(tt.get_span());
                self.render.splice(expr, Escape::PassThru);
            },
            [ref tt @ dollar!(), ..] => {
                self.shift(1);
                let expr = self.splice(tt.get_span());
                self.render.splice(expr, Escape::Escape);
            },
            // Element
            [ident!(sp, name), ..] => {
                self.shift(1);
                self.element(sp, name.as_str());
            },
            // Block
            [TtDelimited(sp, ref d), ..] if d.delim == token::DelimToken::Brace => {
                self.shift(1);
                let stmts = self.block(sp, &d.tts);
                self.render.push_stmts(stmts);
            },
            // ???
            _ => {
                if let [ref tt, ..] = self.input {
                    self.render.cx.span_err(tt.get_span(), "invalid syntax");
                } else {
                    self.render.cx.span_err(self.span, "unexpected end of block");
                }
                return false;
            },
        }
        true
    }

    fn literal(&mut self, tt: &TokenTree, minus: bool) {
        let lit = self.new_rust_parser(vec![tt.clone()]).parse_lit();
        match lit_to_string(self.render.cx, lit, minus) {
            Some(s) => self.render.string(&s, Escape::Escape),
            None => {},
        }
    }

    fn if_expr(&mut self, sp: Span) {
        // Parse the initial if
        let mut if_cond = vec![];
        let if_body;
        loop { match self.input {
            [TtDelimited(sp, ref d), ..] if d.delim == DelimToken::Brace => {
                self.shift(1);
                if_body = self.block(sp, &d.tts);
                break;
            },
            [ref tt, ..] => {
                self.shift(1);
                if_cond.push(tt.clone());
            },
            [] => self.render.cx.span_fatal(sp, "expected body for this $if"),
        }}
        // Parse the (optional) else
        let else_body = match self.input {
            [dollar!(), ident!(else_), ..] if else_.as_str() == "else" => {
                self.shift(2);
                match self.input {
                    [ident!(sp, if_), ..] if if_.as_str() == "if" => {
                        self.shift(1);
                        let else_body = {
                            // Parse an if expression, but capture the result
                            // rather than emitting it right away
                            let mut render = self.render.fork();
                            mem::swap(&mut self.render, &mut render);
                            self.if_expr(sp);
                            mem::swap(&mut self.render, &mut render);
                            render.into_stmts()
                        };
                        Some(else_body)
                    },
                    [TtDelimited(sp, ref d), ..] if d.delim == DelimToken::Brace => {
                        self.shift(1);
                        Some(self.block(sp, &d.tts))
                    },
                    _ => self.render.cx.span_fatal(sp, "expected body for this $else"),
                }
            },
            _ => None,
        };
        self.render.emit_if(if_cond, if_body, else_body);
    }

    fn for_expr(&mut self, sp: Span) {
        let mut pattern = vec![];
        loop { match self.input {
            [ident!(in_), ..] if in_.as_str() == "in" => {
                self.shift(1);
                break;
            },
            [ref tt, ..] => {
                self.shift(1);
                pattern.push(tt.clone());
            },
            _ => self.render.cx.span_fatal(sp, "invalid $for"),
        }}
        let pattern = self.new_rust_parser(pattern).parse_pat();
        let mut iterable = vec![];
        let body;
        loop { match self.input {
            [TtDelimited(sp, ref d), ..] if d.delim == DelimToken::Brace => {
                self.shift(1);
                body = self.block(sp, &d.tts);
                break;
            },
            [ref tt, ..] => {
                self.shift(1);
                iterable.push(tt.clone());
            },
            _ => self.render.cx.span_fatal(sp, "invalid $for"),
        }}
        let iterable = self.new_rust_parser(iterable).parse_expr();
        self.render.emit_for(pattern, iterable, body);
    }

    fn splice(&mut self, sp: Span) -> P<Expr> {
        let mut tts = vec![];
        // First, munch a single token tree
        if let [ref tt, ..] = self.input {
            self.shift(1);
            tts.push(tt.clone());
        }
        loop { match self.input {
            // Munch attribute lookups e.g. `$person.address.street`
            [ref dot @ dot!(), ref ident @ ident!(_), ..] => {
                self.shift(2);
                tts.push(dot.clone());
                tts.push(ident.clone());
            },
            // Munch function calls `()` and indexing operations `[]`
            [TtDelimited(sp, ref d), ..] if d.delim != token::DelimToken::Brace => {
                self.shift(1);
                tts.push(TtDelimited(sp, d.clone()));
            },
            _ => break,
        }}
        if tts.is_empty() {
            self.render.cx.span_fatal(sp, "expected expression for this splice");
        } else {
            self.new_rust_parser(tts).parse_expr()
        }
    }

    fn element(&mut self, sp: Span, name: &str) {
        if self.in_attr {
            self.render.cx.span_err(sp, "unexpected element, you silly bumpkin");
            return;
        }
        self.render.element_open_start(name);
        self.attrs();
        self.render.element_open_end();
        if let [slash!(), ..] = self.input {
            self.shift(1);
        } else {
            self.markup();
            self.render.element_close(name);
        }
    }

    fn attrs(&mut self) {
        loop { match self.input {
            [ident!(name), eq!(), ..] => {
                // Non-empty attribute
                self.shift(2);
                self.render.attribute_start(name.as_str());
                {
                    // Parse a value under an attribute context
                    let mut in_attr = true;
                    mem::swap(&mut self.in_attr, &mut in_attr);
                    self.markup();
                    mem::swap(&mut self.in_attr, &mut in_attr);
                }
                self.render.attribute_end();
            },
            [ident!(name), question!(), ..] => {
                // Empty attribute
                self.shift(2);
                if let [ref tt @ eq!(), ..] = self.input {
                    // Toggle the attribute based on a boolean expression
                    self.shift(1);
                    let cond = self.splice(tt.get_span());
                    // Silence "unnecessary parentheses" warnings
                    let cond = strip_outer_parens(cond).to_tokens(self.render.cx);
                    let body = {
                        let mut r = self.render.fork();
                        r.attribute_empty(name.as_str());
                        r.into_stmts()
                    };
                    self.render.emit_if(cond, body, None);
                } else {
                    // Write the attribute unconditionally
                    self.render.attribute_empty(name.as_str());
                }
            },
            _ => return,
        }}
    }

    fn block(&mut self, sp: Span, tts: &[TokenTree]) -> Vec<P<Stmt>> {
        let mut parse = Parser {
            in_attr: self.in_attr,
            input: tts,
            span: sp,
            render: self.render.fork(),
        };
        parse.markups();
        parse.into_render().into_stmts()
    }
}

/// Convert a literal to a string.
fn lit_to_string(cx: &ExtCtxt, lit: Lit, minus: bool) -> Option<String> {
    use syntax::ast::Lit_::*;
    let mut result = String::new();
    if minus {
        result.push('-');
    }
    match lit.node {
        LitStr(s, _) => result.push_str(&s),
        LitBinary(..) | LitByte(..) => {
            cx.span_err(lit.span, "cannot splice binary data");
            return None;
        },
        LitChar(c) => result.push(c),
        LitInt(x, _) => result.push_str(&x.to_string()),
        LitFloat(s, _) | LitFloatUnsuffixed(s) => result.push_str(&s),
        LitBool(b) => result.push_str(if b { "true" } else { "false" }),
    };
    Some(result)
}

/// If the expression is wrapped in parentheses, strip them off.
fn strip_outer_parens(expr: P<Expr>) -> P<Expr> {
    expr.and_then(|expr| match expr {
        Expr { node: ExprParen(inner), .. } => inner,
        expr => P(expr),
    })
}
<|endoftext|>


<commit_before>//! Simple mailbox implementation using `HashMap`s
//!

use std::collections::HashMap;
use std::fmt::Debug;
use std::cmp::Ordering;

use rustc_serialize::Encodable;

use mpirs::comm_request::{CommRequest, RequestProc};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
struct MailboxKey {
    key_type: MessageTy,
    actor: RequestProc,
    tag: u64,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum MessageTy {
    MSend,
    MRecv,
}

impl MailboxKey {
    fn new(kt: MessageTy, p: RequestProc, t: u64) -> MailboxKey {
        MailboxKey {
            key_type: kt,
            actor: p,
            tag: t,
        }
    }
}

#[derive(Clone, Debug)]
enum KT {
    H1(MailboxKey),
    H2(MailboxKey),
}

#[derive(Clone, Debug)]
pub struct Mail {
    id: usize,
    port: usize,
    req: String,
}

impl PartialEq for Mail {
    fn eq(&self, other: &Mail) -> bool {
        self.id == other.id
    }
}

impl Eq for Mail { }

impl PartialOrd for Mail {
    fn partial_cmp(&self, other: &Mail) -> Option<Ordering> {
        self.id.partial_cmp(&other.id)
    }
}

impl Ord for Mail {
    fn cmp(&self, other: &Mail) -> Ordering {
        self.id.cmp(&other.id)
    }
}

#[derive(Debug, Clone)]
pub struct Mailbox {
    h1: HashMap<MailboxKey, Vec<Mail>>,
    h2: HashMap<MailboxKey, Vec<Mail>>,
}

impl Mailbox {
    pub fn new() -> Mailbox {
        Mailbox {
            h1: HashMap::new(),
            h2: HashMap::new(),
        }
    }

    pub fn pop_matching_mail<T>(&mut self, req: &CommRequest<T>) -> Option<Mail>
        where T: Debug + Clone + Encodable
    {
        unimplemented!()
    }

    pub fn insert_mail<T>(&mut self, mail: Mail, req: &CommRequest<T>)
        where T: Debug + Clone + Encodable
    {
        unimplemented!()
    }

    fn mirror_keys<T>(req: &CommRequest<T>) -> Vec<KT>
        where T: Debug + Clone + Encodable
    {
        if req.is_send() {
            let mut keys = vec![KT::H1(MailboxKey::new(MessageTy::MRecv,
                                                       req.src().unwrap(),
                                                       req.tag())),
                                KT::H1(MailboxKey::new(MessageTy::MRecv,
                                                       RequestProc::Any,
                                                       req.tag()))];
            if !req.is_dst_any() {
                keys.push(KT::H2(MailboxKey::new(MessageTy::MRecv, req.dst().unwrap(), req.tag())));
            }
            keys
        } else if req.is_recv() {
            let mut keys = vec![KT::H2(MailboxKey::new(MessageTy::MSend,
                                                       req.dst().unwrap(),
                                                       req.tag())),
                                KT::H2(MailboxKey::new(MessageTy::MSend,
                                                       RequestProc::Any,
                                                       req.tag()))];
            if !req.is_src_any() {
                keys.push(KT::H1(MailboxKey::new(MessageTy::MSend, req.src().unwrap(), req.tag())));
            }
            keys
        } else {
            panic!("Request is neither a send nor recv");
        }
    }

    fn fast_first_union_intersect(&mut self, keys: Vec<KT>) -> Option<Mail> {
        unimplemented!()
    }

    fn fast_first_union(&mut self, keys: Vec<KT>) -> Option<Mail> {
        unimplemented!()
    }
}
<commit_msg>Implements fast first union<commit_after>//! Simple mailbox implementation using `HashMap`s
//!

use std::collections::{HashMap, VecDeque};
use std::collections::hash_map::Entry;
use std::fmt::Debug;
use std::cmp::Ordering;

use rustc_serialize::Encodable;
use rustc_serialize::json;

use mpirs::comm_request::{CommRequest, RequestProc};

macro_rules! get_value {
    ($m: ident, $k: expr) => {
        match $k {
            KT::H1(ref k) => $m.h1.get(k).clone(),
            KT::H2(ref k) => $m.h2.get(k).clone(),
        }
    }
}

macro_rules! insert_value {
    ($m: ident, $k: expr, $v: expr) => {
        match $k {
            KT::H1(ref k) => $m.h1.insert(k.clone(), $v.clone()),
            KT::H2(ref k) => $m.h2.insert(k.clone(), $v.clone()),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
struct MailboxKey {
    key_type: MessageTy,
    actor: RequestProc,
    tag: u64,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum MessageTy {
    MSend,
    MRecv,
}

impl MailboxKey {
    fn new(kt: MessageTy, p: RequestProc, t: u64) -> MailboxKey {
        MailboxKey {
            key_type: kt,
            actor: p,
            tag: t,
        }
    }
}

#[derive(Clone, Debug)]
enum KT {
    H1(MailboxKey),
    H2(MailboxKey),
}

#[derive(Clone, Debug)]
pub struct Mail {
    id: usize,
    port: usize,
    req: String,
}

impl Mail {
    pub fn new<T>(id: usize, port: usize, req: &CommRequest<T>) -> Mail
        where T: Debug + Clone + Encodable
    {
        Mail {
            id: id,
            port: port,
            req: json::encode(req).unwrap(),
        }
    }
}

impl PartialEq for Mail {
    fn eq(&self, other: &Mail) -> bool {
        self.id == other.id
    }
}

impl Eq for Mail { }

impl PartialOrd for Mail {
    fn partial_cmp(&self, other: &Mail) -> Option<Ordering> {
        self.id.partial_cmp(&other.id)
    }
}

impl Ord for Mail {
    fn cmp(&self, other: &Mail) -> Ordering {
        self.id.cmp(&other.id)
    }
}

#[derive(Debug, Clone)]
pub struct Mailbox {
    h1: HashMap<MailboxKey, VecDeque<Mail>>,
    h2: HashMap<MailboxKey, VecDeque<Mail>>,
    id: usize,
}

impl Mailbox {
    pub fn new() -> Mailbox {
        Mailbox {
            id: 0,
            h1: HashMap::new(),
            h2: HashMap::new(),
        }
    }

    pub fn pop_matching_mail<T>(&mut self, req: &CommRequest<T>) -> Option<Mail>
        where T: Debug + Clone + Encodable
    {
        let keys = Mailbox::mirror_keys(req);
        match keys.len() {
            2 => self.fast_first_union(keys),
            3 => self.fast_first_union_intersect(keys),
            _ => unreachable!(),
        }
    }

    pub fn insert_mail<T>(&mut self, mail: Mail, req: &CommRequest<T>, port: usize)
        where T: Debug + Clone + Encodable
    {
        let mtype = if req.is_send() {
            MessageTy::MSend
        } else {
            MessageTy::MRecv
        };

        let mail = Mail::new(self.id, port, req);
        self.id += 1;

        let h1_key = MailboxKey::new(mtype, req.src().unwrap(), req.tag());
        let h2_key = MailboxKey::new(mtype, req.dst().unwrap(), req.tag());

        if !self.h1.contains_key(&h1_key) {
            self.h1.insert(h1_key, VecDeque::new());
        }

        if let Some(ref mut v) = self.h1.get_mut(&h1_key) {
            v.push_back(mail.clone());
        }

        if !self.h2.contains_key(&h2_key) {
            self.h2.insert(h2_key, VecDeque::new());
        }

        if let Some(ref mut v) = self.h2.get_mut(&h2_key) {
            v.push_back(mail.clone());
        }
    }

    fn mirror_keys<T>(req: &CommRequest<T>) -> Vec<KT>
        where T: Debug + Clone + Encodable
    {
        if req.is_send() {
            let mut keys = vec![KT::H1(MailboxKey::new(MessageTy::MRecv,
                                                       req.src().unwrap(),
                                                       req.tag())),
                                KT::H1(MailboxKey::new(MessageTy::MRecv,
                                                       RequestProc::Any,
                                                       req.tag()))];
            if !req.is_dst_any() {
                keys.push(KT::H2(MailboxKey::new(MessageTy::MRecv, req.dst().unwrap(), req.tag())));
            }
            keys
        } else if req.is_recv() {
            let mut keys = vec![KT::H2(MailboxKey::new(MessageTy::MSend,
                                                       req.dst().unwrap(),
                                                       req.tag())),
                                KT::H2(MailboxKey::new(MessageTy::MSend,
                                                       RequestProc::Any,
                                                       req.tag()))];
            if !req.is_src_any() {
                keys.push(KT::H1(MailboxKey::new(MessageTy::MSend, req.src().unwrap(), req.tag())));
            }
            keys
        } else {
            panic!("Request is neither a send nor recv");
        }
    }

    fn fast_first_union_intersect(&mut self, keys: Vec<KT>) -> Option<Mail> {
        unimplemented!()
    }

    fn fast_first_union(&mut self, keys: Vec<KT>) -> Option<Mail> {
        let (e, v1_, v2_) = {
            let v1 = get_value!(self, keys[0]);
            let v2 = get_value!(self, keys[1]);

            match (v1, v2) {
                (None, None) => (None, None, None),
                (None, Some(e)) => {
                    let mut e_ = e.clone();
                    let ele = e_.pop_front();
                    (ele, None, Some(e_))
                }
                (Some(e), None) => {
                    let mut e_ = e.clone();
                    let ele = e_.pop_front();
                    (ele, Some(e_), None)
                }
                (Some(e), Some(e_)) => {
                    let mut vec1 = e.clone();
                    let mut vec2 = e_.clone();
                    let ele = match vec1[0].cmp(&vec2[0]) {
                        Ordering::Less => vec1.pop_front(),
                        Ordering::Greater => vec2.pop_front(),
                        Ordering::Equal => unreachable!(),
                    };
                    (ele, Some(vec1), Some(vec2))
                },
            }
        };
        if v1_.is_some() {
            insert_value!(self, keys[0], v1_.unwrap());
        }
        if v2_.is_some() {
            insert_value!(self, keys[1], v2_.unwrap());
        }
        e
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate crossbeam_channel;

use std::collections::HashMap;
use std::io::BufReader;
use std::io::Result;
use std::io::{Error, ErrorKind};
use std::net::{TcpListener, TcpStream};
use std::sync::Arc;
use std::sync::Mutex;
use std::thread;
use std::time::Duration;

use crossbeam_channel::{tick, unbounded, Sender};

use birch::client::{ClientConnection, ClientEvent};
use birch::proto::RawMessage;
use birch::socket::IrcReader;
use birch::socket::{IrcChannel, IrcSocket, IrcSocketConfig, IrcWriter};

struct ClientFanoutManager {
    id_counter: usize,
    clients: HashMap<usize, Sender<RawMessage>>,
}

impl ClientFanoutManager {
    fn new() -> Self {
        Self {
            id_counter: 0,
            clients: HashMap::new(),
        }
    }

    fn next_id(&mut self) -> usize {
        let id = self.id_counter;
        self.id_counter += 1;
        id
    }

    fn add_client(&mut self, sender: Sender<RawMessage>) -> usize {
        let id = self.next_id();
        self.clients.insert(id, sender);
        id
    }

    fn remove_client(&mut self, id: usize) -> Option<()> {
        self.clients.remove(&id).map(|_| ())
    }

    fn write_message(&mut self, msg: &RawMessage) {
        let mut dead_clients = vec![];
        for (id, client) in self.clients.iter() {
            if let Err(err) = client.send(msg.clone()) {
                println!("Failed to write to client: {:?}", err);
                dead_clients.push(*id);
            }
        }
        // Prune anything we failed to write to
        for id in dead_clients.into_iter() {
            self.remove_client(id);
        }
    }
}

struct ClientManager {
    fanout_manager: Arc<Mutex<ClientFanoutManager>>,
    fanout: IrcChannel,
}

impl ClientManager {
    fn new() -> Self {
        let (sender, receiver) = IrcChannel::new();
        let manager = Arc::new(Mutex::new(ClientFanoutManager::new()));

        let fanout = Arc::clone(&manager);
        thread::spawn(move || {
            for msg in receiver {
                fanout.lock().expect("mutex poisoned").write_message(&msg);
            }
        });

        Self {
            fanout_manager: manager,
            fanout: sender,
        }
    }

    fn fanout_sender(&self) -> IrcChannel {
        self.fanout.clone()
    }

    fn accept_client_connection(
        &mut self,
        stream: TcpStream,
        new_user_chan: Sender<IrcChannel>,
    ) {
        let to_client = unbounded();
        let from_client = unbounded();

        let read_stream = stream.try_clone().expect("clone stream");
        let mut write_stream = stream.try_clone().expect("clone stream");

        // Read thread - handle input from client
        let client_sender = from_client.0.clone();
        thread::spawn(move || {
            // TODO: This should not be a magic number
            read_stream
                .set_read_timeout(Some(Duration::from_secs(240)))
                .expect("set_read_timeout call failed");

            let mut reader = BufReader::new(read_stream);
            loop {
                match reader.read_message() {
                    Ok(msg) => {
                        if let Err(err) = client_sender.send(msg) {
                            println!("Failed to send: {:?}", err);
                            break;
                        }
                    }
                    Err(err) => {
                        println!("Failed to read client message: {:?}", err);
                        break;
                    }
                }
            }
        });

        let fanout = Arc::clone(&self.fanout_manager);
        thread::spawn(move || {
            let events = unbounded();
            let ping_ticker = tick(Duration::from_secs(120));
            let mut client = ClientConnection::new(events.0);

            let recv_err = || Error::new(ErrorKind::Other, "receiver disconnected");

            loop {
                let result = select! {
                    recv(events.1) -> msg => {
                        msg.map_err(|_| recv_err())
                            .and_then(|msg| match msg {
                                ClientEvent::WriteNetwork(msg) => {
                                    println!("todo; send to network: {:?}", msg);
                                    Ok(())
                                },
                                ClientEvent::WriteClient(msg) =>
                                    to_client.0.send(msg).map_err(|_| recv_err()),

                                ClientEvent::Authenticated => {
                                    fanout.lock().unwrap().add_client(to_client.0.clone());
                                    new_user_chan.send(IrcChannel(to_client.0.clone())).expect("blah");
                                    Ok(())
                                },
                            })
                    },

                    recv(to_client.1) -> msg => {
                        msg.map_err(|_| recv_err())
                            .and_then(|msg| write_stream.write_message(&msg))
                    },

                    recv(from_client.1) -> msg => {
                        msg.map_err(|_| recv_err())
                            .and_then(|msg| client.handle_message(&msg))
                    },

                    recv(ping_ticker) -> _ => client.ping()
                };

                if let Err(err) = result {
                    println!("client failed: {:?}", err);
                    break;
                }
            }
        });
    }
}

fn main() -> Result<()> {
    let mut client_manager = ClientManager::new();

    let config = IrcSocketConfig {
        nick: "ep",
        addr: "irc.freenode.net:6667",
        max_retries: Some(3),
    };
    let mut sock = IrcSocket::new(config);

    let mut fanout_sender = client_manager.fanout_sender();
    let new_user_chan = sock.new_user_channel();

    thread::spawn(move || {
        // TODO: surface any errors here + shutdown hook
        sock.start_loop(&mut fanout_sender)
    });

    // TODO: hook this up to config
    let listener = TcpListener::bind("127.0.0.1:9123")?;
    for stream in listener.incoming() {
        match stream {
            Ok(stream) => {
                println!("Client connected");
                client_manager.accept_client_connection(stream, new_user_chan.clone());
            }
            Err(err) => println!("Failed to accept: {:?}", err),
        }
    }

    Ok(())
}
<commit_msg>Collapse imports<commit_after>#[macro_use]
extern crate crossbeam_channel;

use std::collections::HashMap;
use std::io::{BufReader, Error, ErrorKind, Result};
use std::net::{TcpListener, TcpStream};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;

use crossbeam_channel::{tick, unbounded, Sender};

use birch::client::{ClientConnection, ClientEvent};
use birch::proto::RawMessage;
use birch::socket::{IrcChannel, IrcReader, IrcSocket, IrcSocketConfig, IrcWriter};

struct ClientFanoutManager {
    id_counter: usize,
    clients: HashMap<usize, Sender<RawMessage>>,
}

impl ClientFanoutManager {
    fn new() -> Self {
        Self {
            id_counter: 0,
            clients: HashMap::new(),
        }
    }

    fn next_id(&mut self) -> usize {
        let id = self.id_counter;
        self.id_counter += 1;
        id
    }

    fn add_client(&mut self, sender: Sender<RawMessage>) -> usize {
        let id = self.next_id();
        self.clients.insert(id, sender);
        id
    }

    fn remove_client(&mut self, id: usize) -> Option<()> {
        self.clients.remove(&id).map(|_| ())
    }

    fn write_message(&mut self, msg: &RawMessage) {
        let mut dead_clients = vec![];
        for (id, client) in self.clients.iter() {
            if let Err(err) = client.send(msg.clone()) {
                println!("Failed to write to client: {:?}", err);
                dead_clients.push(*id);
            }
        }
        // Prune anything we failed to write to
        for id in dead_clients.into_iter() {
            self.remove_client(id);
        }
    }
}

struct ClientManager {
    fanout_manager: Arc<Mutex<ClientFanoutManager>>,
    fanout: IrcChannel,
}

impl ClientManager {
    fn new() -> Self {
        let (sender, receiver) = IrcChannel::new();
        let manager = Arc::new(Mutex::new(ClientFanoutManager::new()));

        let fanout = Arc::clone(&manager);
        thread::spawn(move || {
            for msg in receiver {
                fanout.lock().expect("mutex poisoned").write_message(&msg);
            }
        });

        Self {
            fanout_manager: manager,
            fanout: sender,
        }
    }

    fn fanout_sender(&self) -> IrcChannel {
        self.fanout.clone()
    }

    fn accept_client_connection(
        &mut self,
        stream: TcpStream,
        new_user_chan: Sender<IrcChannel>,
    ) {
        let to_client = unbounded();
        let from_client = unbounded();

        let read_stream = stream.try_clone().expect("clone stream");
        let mut write_stream = stream.try_clone().expect("clone stream");

        // Read thread - handle input from client
        let client_sender = from_client.0.clone();
        thread::spawn(move || {
            // TODO: This should not be a magic number
            read_stream
                .set_read_timeout(Some(Duration::from_secs(240)))
                .expect("set_read_timeout call failed");

            let mut reader = BufReader::new(read_stream);
            loop {
                match reader.read_message() {
                    Ok(msg) => {
                        if let Err(err) = client_sender.send(msg) {
                            println!("Failed to send: {:?}", err);
                            break;
                        }
                    }
                    Err(err) => {
                        println!("Failed to read client message: {:?}", err);
                        break;
                    }
                }
            }
        });

        let fanout = Arc::clone(&self.fanout_manager);
        thread::spawn(move || {
            let events = unbounded();
            let ping_ticker = tick(Duration::from_secs(120));
            let mut client = ClientConnection::new(events.0);

            let recv_err = || Error::new(ErrorKind::Other, "receiver disconnected");

            loop {
                let result = select! {
                    recv(events.1) -> msg => {
                        msg.map_err(|_| recv_err())
                            .and_then(|msg| match msg {
                                ClientEvent::WriteNetwork(msg) => {
                                    println!("todo; send to network: {:?}", msg);
                                    Ok(())
                                },
                                ClientEvent::WriteClient(msg) =>
                                    to_client.0.send(msg).map_err(|_| recv_err()),

                                ClientEvent::Authenticated => {
                                    fanout.lock().unwrap().add_client(to_client.0.clone());
                                    new_user_chan.send(IrcChannel(to_client.0.clone())).expect("blah");
                                    Ok(())
                                },
                            })
                    },

                    recv(to_client.1) -> msg => {
                        msg.map_err(|_| recv_err())
                            .and_then(|msg| write_stream.write_message(&msg))
                    },

                    recv(from_client.1) -> msg => {
                        msg.map_err(|_| recv_err())
                            .and_then(|msg| client.handle_message(&msg))
                    },

                    recv(ping_ticker) -> _ => client.ping()
                };

                if let Err(err) = result {
                    println!("client failed: {:?}", err);
                    break;
                }
            }
        });
    }
}

fn main() -> Result<()> {
    let mut client_manager = ClientManager::new();

    let config = IrcSocketConfig {
        nick: "ep",
        addr: "irc.freenode.net:6667",
        max_retries: Some(3),
    };
    let mut sock = IrcSocket::new(config);

    let mut fanout_sender = client_manager.fanout_sender();
    let new_user_chan = sock.new_user_channel();

    thread::spawn(move || {
        // TODO: surface any errors here + shutdown hook
        sock.start_loop(&mut fanout_sender)
    });

    // TODO: hook this up to config
    let listener = TcpListener::bind("127.0.0.1:9123")?;
    for stream in listener.incoming() {
        match stream {
            Ok(stream) => {
                println!("Client connected");
                client_manager.accept_client_connection(stream, new_user_chan.clone());
            }
            Err(err) => println!("Failed to accept: {:?}", err),
        }
    }

    Ok(())
}
<|endoftext|>


<commit_before>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Martin Kysel <code@martinkysel.com>
//  *
//  * For the full copyright and license information, please view the LICENSE file
//  * that was distributed with this source code.

// spell-checker:ignore (methods) isnt

#[macro_use]
extern crate uucore;

use std::{
    fs::File,
    io::{stdin, stdout, BufReader, Read, Stdout, Write},
    path::Path,
    time::Duration,
};

#[cfg(all(unix, not(target_os = "fuchsia")))]
extern crate nix;

use clap::{crate_version, App, Arg};
use crossterm::{
    event::{self, Event, KeyCode, KeyEvent, KeyModifiers},
    execute, queue,
    style::Attribute,
    terminal,
};

use unicode_segmentation::UnicodeSegmentation;
use unicode_width::UnicodeWidthStr;

const BELL: &str = "\x07";

pub mod options {
    pub const SILENT: &str = "silent";
    pub const LOGICAL: &str = "logical";
    pub const NO_PAUSE: &str = "no-pause";
    pub const PRINT_OVER: &str = "print-over";
    pub const CLEAN_PRINT: &str = "clean-print";
    pub const SQUEEZE: &str = "squeeze";
    pub const PLAIN: &str = "plain";
    pub const LINES: &str = "lines";
    pub const NUMBER: &str = "number";
    pub const PATTERN: &str = "pattern";
    pub const FROM_LINE: &str = "from-line";
    pub const FILES: &str = "files";
}

const MULTI_FILE_TOP_PROMPT: &str = "::::::::::::::\n{}\n::::::::::::::\n";

pub fn uumain(args: impl uucore::Args) -> i32 {
    let matches = App::new(executable!())
        .about("A file perusal filter for CRT viewing.")
        .version(crate_version!())
        .arg(
            Arg::with_name(options::SILENT)
                .short("d")
                .long(options::SILENT)
                .help("Display help instead of ringing bell"),
        )
        // The commented arguments below are unimplemented:
        /*
        .arg(
            Arg::with_name(options::LOGICAL)
                .short("f")
                .long(options::LOGICAL)
                .help("Count logical rather than screen lines"),
        )
        .arg(
            Arg::with_name(options::NO_PAUSE)
                .short("l")
                .long(options::NO_PAUSE)
                .help("Suppress pause after form feed"),
        )
        .arg(
            Arg::with_name(options::PRINT_OVER)
                .short("c")
                .long(options::PRINT_OVER)
                .help("Do not scroll, display text and clean line ends"),
        )
        .arg(
            Arg::with_name(options::CLEAN_PRINT)
                .short("p")
                .long(options::CLEAN_PRINT)
                .help("Do not scroll, clean screen and display text"),
        )
        .arg(
            Arg::with_name(options::SQUEEZE)
                .short("s")
                .long(options::SQUEEZE)
                .help("Squeeze multiple blank lines into one"),
        )
        .arg(
            Arg::with_name(options::PLAIN)
                .short("u")
                .long(options::PLAIN)
                .help("Suppress underlining and bold"),
        )
        .arg(
            Arg::with_name(options::LINES)
                .short("n")
                .long(options::LINES)
                .value_name("number")
                .takes_value(true)
                .help("The number of lines per screen full"),
        )
        .arg(
            Arg::with_name(options::NUMBER)
                .allow_hyphen_values(true)
                .long(options::NUMBER)
                .required(false)
                .takes_value(true)
                .help("Same as --lines"),
        )
        .arg(
            Arg::with_name(options::FROM_LINE)
                .short("F")
                .allow_hyphen_values(true)
                .required(false)
                .takes_value(true)
                .value_name("number")
                .help("Display file beginning from line number"),
        )
        .arg(
            Arg::with_name(options::PATTERN)
                .short("P")
                .allow_hyphen_values(true)
                .required(false)
                .takes_value(true)
                .help("Display file beginning from pattern match"),
        )
        */
        .arg(
            Arg::with_name(options::FILES)
                .required(false)
                .multiple(true)
                .help("Path to the files to be read"),
        )
        .get_matches_from(args);

    let mut buff = String::new();
    let silent = matches.is_present(options::SILENT);
    if let Some(files) = matches.values_of(options::FILES) {
        let mut stdout = setup_term();
        let length = files.len();

        let mut files_iter = files.peekable();
        while let (Some(file), next_file) = (files_iter.next(), files_iter.peek()) {
            let file = Path::new(file);
            if file.is_dir() {
                terminal::disable_raw_mode().unwrap();
                show_usage_error!("'{}' is a directory.", file.display());
                return 1;
            }
            if !file.exists() {
                terminal::disable_raw_mode().unwrap();
                show_error!("cannot open {}: No such file or directory", file.display());
                return 1;
            }
            if length > 1 {
                buff.push_str(&MULTI_FILE_TOP_PROMPT.replace("{}", file.to_str().unwrap()));
            }
            let mut reader = BufReader::new(File::open(file).unwrap());
            reader.read_to_string(&mut buff).unwrap();
            more(&buff, &mut stdout, next_file.copied(), silent);
            buff.clear();
        }
        reset_term(&mut stdout);
    } else if atty::isnt(atty::Stream::Stdin) {
        stdin().read_to_string(&mut buff).unwrap();
        let mut stdout = setup_term();
        more(&buff, &mut stdout, None, silent);
        reset_term(&mut stdout);
    } else {
        show_usage_error!("bad usage");
    }
    0
}

#[cfg(not(target_os = "fuchsia"))]
fn setup_term() -> std::io::Stdout {
    let stdout = stdout();
    terminal::enable_raw_mode().unwrap();
    stdout
}

#[cfg(target_os = "fuchsia")]
#[inline(always)]
fn setup_term() -> usize {
    0
}

#[cfg(not(target_os = "fuchsia"))]
fn reset_term(stdout: &mut std::io::Stdout) {
    terminal::disable_raw_mode().unwrap();
    // Clear the prompt
    queue!(stdout, terminal::Clear(terminal::ClearType::CurrentLine)).unwrap();
    // Move cursor to the beginning without printing new line
    print!("\r");
    stdout.flush().unwrap();
}

#[cfg(target_os = "fuchsia")]
#[inline(always)]
fn reset_term(_: &mut usize) {}

fn more(buff: &str, mut stdout: &mut Stdout, next_file: Option<&str>, silent: bool) {
    let (cols, rows) = terminal::size().unwrap();
    let lines = break_buff(buff, usize::from(cols));

    let mut pager = Pager::new(rows, lines, next_file, silent);
    pager.draw(stdout, None);
    if pager.should_close() {
        return;
    }

    loop {
        let mut wrong_key = None;
        if event::poll(Duration::from_millis(10)).unwrap() {
            match event::read().unwrap() {
                Event::Key(KeyEvent {
                    code: KeyCode::Char('q'),
                    modifiers: KeyModifiers::NONE,
                })
                | Event::Key(KeyEvent {
                    code: KeyCode::Char('c'),
                    modifiers: KeyModifiers::CONTROL,
                }) => {
                    reset_term(&mut stdout);
                    std::process::exit(0);
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Down,
                    modifiers: KeyModifiers::NONE,
                })
                | Event::Key(KeyEvent {
                    code: KeyCode::Char(' '),
                    modifiers: KeyModifiers::NONE,
                }) => {
                    if pager.should_close() {
                        return;
                    } else {
                        pager.page_down();
                    }
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Up,
                    modifiers: KeyModifiers::NONE,
                }) => {
                    pager.page_up();
                }
                Event::Resize(col, row) => {
                    pager.page_resize(col, row);
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Char(k),
                    ..
                }) => wrong_key = Some(k),
                _ => continue,
            }

            pager.draw(stdout, wrong_key);
        }
    }
}

struct Pager<'a> {
    // The current line at the top of the screen
    upper_mark: usize,
    // The number of rows that fit on the screen
    content_rows: u16,
    lines: Vec<String>,
    next_file: Option<&'a str>,
    line_count: usize,
    silent: bool,
}

impl<'a> Pager<'a> {
    fn new(rows: u16, lines: Vec<String>, next_file: Option<&'a str>, silent: bool) -> Self {
        let line_count = lines.len();
        Self {
            upper_mark: 0,
            content_rows: rows.saturating_sub(1),
            lines,
            next_file,
            line_count,
            silent,
        }
    }

    fn should_close(&mut self) -> bool {
        self.upper_mark
            .saturating_add(self.content_rows.into())
            .ge(&self.line_count)
    }

    fn page_down(&mut self) {
        self.upper_mark = self
            .upper_mark
            .saturating_add(self.content_rows.into())
            .min(self.line_count.saturating_sub(self.content_rows.into()));
    }

    fn page_up(&mut self) {
        self.upper_mark = self.upper_mark.saturating_sub(self.content_rows.into());
    }

    // TODO: Deal with column size changes.
    fn page_resize(&mut self, _: u16, row: u16) {
        self.content_rows = row.saturating_sub(1);
    }

    fn draw(&self, stdout: &mut std::io::Stdout, wrong_key: Option<char>) {
        let lower_mark = self
            .line_count
            .min(self.upper_mark.saturating_add(self.content_rows.into()));
        self.draw_lines(stdout);
        self.draw_prompt(stdout, lower_mark, wrong_key);
        stdout.flush().unwrap();
    }

    fn draw_lines(&self, stdout: &mut std::io::Stdout) {
        execute!(stdout, terminal::Clear(terminal::ClearType::CurrentLine)).unwrap();
        let displayed_lines = self
            .lines
            .iter()
            .skip(self.upper_mark)
            .take(self.content_rows.into());

        for line in displayed_lines {
            stdout
                .write_all(format!("\r{}\n", line).as_bytes())
                .unwrap();
        }
    }

    fn draw_prompt(&self, stdout: &mut Stdout, lower_mark: usize, wrong_key: Option<char>) {
        let status_inner = if lower_mark == self.line_count {
            format!("Next file: {}", self.next_file.unwrap_or_default())
        } else {
            format!(
                "{}%",
                (lower_mark as f64 / self.line_count as f64 * 100.0).round() as u16
            )
        };

        let status = format!("--More--({})", status_inner);

        let banner = match (self.silent, wrong_key) {
            (true, Some(key)) => {
                format!(
                    "{} [Unknown key: '{}'. Press 'h' for instructions. (unimplemented)]",
                    status, key
                )
            }
            (true, None) => format!("{}[Press space to continue, 'q' to quit.]", status),
            (false, Some(_)) => format!("{}{}", status, BELL),
            (false, None) => status,
        };

        write!(
            stdout,
            "\r{}{}{}",
            Attribute::Reverse,
            banner,
            Attribute::Reset
        )
        .unwrap();
    }
}

// Break the lines on the cols of the terminal
fn break_buff(buff: &str, cols: usize) -> Vec<String> {
    let mut lines = Vec::with_capacity(buff.lines().count());

    for l in buff.lines() {
        lines.append(&mut break_line(l, cols));
    }
    lines
}

fn break_line(line: &str, cols: usize) -> Vec<String> {
    let width = UnicodeWidthStr::width(line);
    let mut lines = Vec::new();
    if width < cols {
        lines.push(line.to_string());
        return lines;
    }

    let gr_idx = UnicodeSegmentation::grapheme_indices(line, true);
    let mut last_index = 0;
    let mut total_width = 0;
    for (index, grapheme) in gr_idx {
        let width = UnicodeWidthStr::width(grapheme);
        total_width += width;

        if total_width > cols {
            lines.push(line[last_index..index].to_string());
            last_index = index;
            total_width = width;
        }
    }

    if last_index != line.len() {
        lines.push(line[last_index..].to_string());
    }
    lines
}

#[cfg(test)]
mod tests {
    use super::break_line;
    use unicode_width::UnicodeWidthStr;

    #[test]
    fn test_break_lines_long() {
        let mut test_string = String::with_capacity(100);
        for _ in 0..200 {
            test_string.push('#');
        }

        let lines = break_line(&test_string, 80);
        let widths: Vec<usize> = lines
            .iter()
            .map(|s| UnicodeWidthStr::width(&s[..]))
            .collect();

        assert_eq!((80, 80, 40), (widths[0], widths[1], widths[2]));
    }

    #[test]
    fn test_break_lines_short() {
        let mut test_string = String::with_capacity(100);
        for _ in 0..20 {
            test_string.push('#');
        }

        let lines = break_line(&test_string, 80);

        assert_eq!(20, lines[0].len());
    }

    #[test]
    fn test_break_line_zwj() {
        let mut test_string = String::with_capacity(1100);
        for _ in 0..20 {
            test_string.push_str("👩🏻🔬");
        }

        let lines = break_line(&test_string, 80);

        let widths: Vec<usize> = lines
            .iter()
            .map(|s| UnicodeWidthStr::width(&s[..]))
            .collect();

        // Each 👩🏻🔬 is 6 character width it break line to the closest number to 80 => 6 * 13 = 78
        assert_eq!((78, 42), (widths[0], widths[1]));
    }
}
<commit_msg>Simplified page down implementation<commit_after>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Martin Kysel <code@martinkysel.com>
//  *
//  * For the full copyright and license information, please view the LICENSE file
//  * that was distributed with this source code.

// spell-checker:ignore (methods) isnt

#[macro_use]
extern crate uucore;

use std::{
    fs::File,
    io::{stdin, stdout, BufReader, Read, Stdout, Write},
    path::Path,
    time::Duration,
};

#[cfg(all(unix, not(target_os = "fuchsia")))]
extern crate nix;

use clap::{crate_version, App, Arg};
use crossterm::{
    event::{self, Event, KeyCode, KeyEvent, KeyModifiers},
    execute, queue,
    style::Attribute,
    terminal,
};

use unicode_segmentation::UnicodeSegmentation;
use unicode_width::UnicodeWidthStr;

const BELL: &str = "\x07";

pub mod options {
    pub const SILENT: &str = "silent";
    pub const LOGICAL: &str = "logical";
    pub const NO_PAUSE: &str = "no-pause";
    pub const PRINT_OVER: &str = "print-over";
    pub const CLEAN_PRINT: &str = "clean-print";
    pub const SQUEEZE: &str = "squeeze";
    pub const PLAIN: &str = "plain";
    pub const LINES: &str = "lines";
    pub const NUMBER: &str = "number";
    pub const PATTERN: &str = "pattern";
    pub const FROM_LINE: &str = "from-line";
    pub const FILES: &str = "files";
}

const MULTI_FILE_TOP_PROMPT: &str = "::::::::::::::\n{}\n::::::::::::::\n";

pub fn uumain(args: impl uucore::Args) -> i32 {
    let matches = App::new(executable!())
        .about("A file perusal filter for CRT viewing.")
        .version(crate_version!())
        .arg(
            Arg::with_name(options::SILENT)
                .short("d")
                .long(options::SILENT)
                .help("Display help instead of ringing bell"),
        )
        // The commented arguments below are unimplemented:
        /*
        .arg(
            Arg::with_name(options::LOGICAL)
                .short("f")
                .long(options::LOGICAL)
                .help("Count logical rather than screen lines"),
        )
        .arg(
            Arg::with_name(options::NO_PAUSE)
                .short("l")
                .long(options::NO_PAUSE)
                .help("Suppress pause after form feed"),
        )
        .arg(
            Arg::with_name(options::PRINT_OVER)
                .short("c")
                .long(options::PRINT_OVER)
                .help("Do not scroll, display text and clean line ends"),
        )
        .arg(
            Arg::with_name(options::CLEAN_PRINT)
                .short("p")
                .long(options::CLEAN_PRINT)
                .help("Do not scroll, clean screen and display text"),
        )
        .arg(
            Arg::with_name(options::SQUEEZE)
                .short("s")
                .long(options::SQUEEZE)
                .help("Squeeze multiple blank lines into one"),
        )
        .arg(
            Arg::with_name(options::PLAIN)
                .short("u")
                .long(options::PLAIN)
                .help("Suppress underlining and bold"),
        )
        .arg(
            Arg::with_name(options::LINES)
                .short("n")
                .long(options::LINES)
                .value_name("number")
                .takes_value(true)
                .help("The number of lines per screen full"),
        )
        .arg(
            Arg::with_name(options::NUMBER)
                .allow_hyphen_values(true)
                .long(options::NUMBER)
                .required(false)
                .takes_value(true)
                .help("Same as --lines"),
        )
        .arg(
            Arg::with_name(options::FROM_LINE)
                .short("F")
                .allow_hyphen_values(true)
                .required(false)
                .takes_value(true)
                .value_name("number")
                .help("Display file beginning from line number"),
        )
        .arg(
            Arg::with_name(options::PATTERN)
                .short("P")
                .allow_hyphen_values(true)
                .required(false)
                .takes_value(true)
                .help("Display file beginning from pattern match"),
        )
        */
        .arg(
            Arg::with_name(options::FILES)
                .required(false)
                .multiple(true)
                .help("Path to the files to be read"),
        )
        .get_matches_from(args);

    let mut buff = String::new();
    let silent = matches.is_present(options::SILENT);
    if let Some(files) = matches.values_of(options::FILES) {
        let mut stdout = setup_term();
        let length = files.len();

        let mut files_iter = files.peekable();
        while let (Some(file), next_file) = (files_iter.next(), files_iter.peek()) {
            let file = Path::new(file);
            if file.is_dir() {
                terminal::disable_raw_mode().unwrap();
                show_usage_error!("'{}' is a directory.", file.display());
                return 1;
            }
            if !file.exists() {
                terminal::disable_raw_mode().unwrap();
                show_error!("cannot open {}: No such file or directory", file.display());
                return 1;
            }
            if length > 1 {
                buff.push_str(&MULTI_FILE_TOP_PROMPT.replace("{}", file.to_str().unwrap()));
            }
            let mut reader = BufReader::new(File::open(file).unwrap());
            reader.read_to_string(&mut buff).unwrap();
            more(&buff, &mut stdout, next_file.copied(), silent);
            buff.clear();
        }
        reset_term(&mut stdout);
    } else if atty::isnt(atty::Stream::Stdin) {
        stdin().read_to_string(&mut buff).unwrap();
        let mut stdout = setup_term();
        more(&buff, &mut stdout, None, silent);
        reset_term(&mut stdout);
    } else {
        show_usage_error!("bad usage");
    }
    0
}

#[cfg(not(target_os = "fuchsia"))]
fn setup_term() -> std::io::Stdout {
    let stdout = stdout();
    terminal::enable_raw_mode().unwrap();
    stdout
}

#[cfg(target_os = "fuchsia")]
#[inline(always)]
fn setup_term() -> usize {
    0
}

#[cfg(not(target_os = "fuchsia"))]
fn reset_term(stdout: &mut std::io::Stdout) {
    terminal::disable_raw_mode().unwrap();
    // Clear the prompt
    queue!(stdout, terminal::Clear(terminal::ClearType::CurrentLine)).unwrap();
    // Move cursor to the beginning without printing new line
    print!("\r");
    stdout.flush().unwrap();
}

#[cfg(target_os = "fuchsia")]
#[inline(always)]
fn reset_term(_: &mut usize) {}

fn more(buff: &str, mut stdout: &mut Stdout, next_file: Option<&str>, silent: bool) {
    let (cols, rows) = terminal::size().unwrap();
    let lines = break_buff(buff, usize::from(cols));

    let mut pager = Pager::new(rows, lines, next_file, silent);
    pager.draw(stdout, None);
    if pager.should_close() {
        return;
    }

    loop {
        let mut wrong_key = None;
        if event::poll(Duration::from_millis(10)).unwrap() {
            match event::read().unwrap() {
                Event::Key(KeyEvent {
                    code: KeyCode::Char('q'),
                    modifiers: KeyModifiers::NONE,
                })
                | Event::Key(KeyEvent {
                    code: KeyCode::Char('c'),
                    modifiers: KeyModifiers::CONTROL,
                }) => {
                    reset_term(&mut stdout);
                    std::process::exit(0);
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Down,
                    modifiers: KeyModifiers::NONE,
                })
                | Event::Key(KeyEvent {
                    code: KeyCode::Char(' '),
                    modifiers: KeyModifiers::NONE,
                }) => {
                    if pager.should_close() {
                        return;
                    } else {
                        pager.page_down();
                    }
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Up,
                    modifiers: KeyModifiers::NONE,
                }) => {
                    pager.page_up();
                }
                Event::Resize(col, row) => {
                    pager.page_resize(col, row);
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Char(k),
                    ..
                }) => wrong_key = Some(k),
                _ => continue,
            }

            pager.draw(stdout, wrong_key);
        }
    }
}

struct Pager<'a> {
    // The current line at the top of the screen
    upper_mark: usize,
    // The number of rows that fit on the screen
    content_rows: u16,
    lines: Vec<String>,
    next_file: Option<&'a str>,
    line_count: usize,
    silent: bool,
}

impl<'a> Pager<'a> {
    fn new(rows: u16, lines: Vec<String>, next_file: Option<&'a str>, silent: bool) -> Self {
        let line_count = lines.len();
        Self {
            upper_mark: 0,
            content_rows: rows.saturating_sub(1),
            lines,
            next_file,
            line_count,
            silent,
        }
    }

    fn should_close(&mut self) -> bool {
        self.upper_mark
            .saturating_add(self.content_rows.into())
            .ge(&self.line_count)
    }

    fn page_down(&mut self) {
        self.upper_mark = self.upper_mark.saturating_add(self.content_rows.into());
    }

    fn page_up(&mut self) {
        self.upper_mark = self.upper_mark.saturating_sub(self.content_rows.into());
    }

    // TODO: Deal with column size changes.
    fn page_resize(&mut self, _: u16, row: u16) {
        self.content_rows = row.saturating_sub(1);
    }

    fn draw(&self, stdout: &mut std::io::Stdout, wrong_key: Option<char>) {
        let lower_mark = self
            .line_count
            .min(self.upper_mark.saturating_add(self.content_rows.into()));
        self.draw_lines(stdout);
        self.draw_prompt(stdout, lower_mark, wrong_key);
        stdout.flush().unwrap();
    }

    fn draw_lines(&self, stdout: &mut std::io::Stdout) {
        execute!(stdout, terminal::Clear(terminal::ClearType::CurrentLine)).unwrap();
        let displayed_lines = self
            .lines
            .iter()
            .skip(self.upper_mark)
            .take(self.content_rows.into());

        for line in displayed_lines {
            stdout
                .write_all(format!("\r{}\n", line).as_bytes())
                .unwrap();
        }
    }

    fn draw_prompt(&self, stdout: &mut Stdout, lower_mark: usize, wrong_key: Option<char>) {
        let status_inner = if lower_mark == self.line_count {
            format!("Next file: {}", self.next_file.unwrap_or_default())
        } else {
            format!(
                "{}%",
                (lower_mark as f64 / self.line_count as f64 * 100.0).round() as u16
            )
        };

        let status = format!("--More--({})", status_inner);

        let banner = match (self.silent, wrong_key) {
            (true, Some(key)) => {
                format!(
                    "{} [Unknown key: '{}'. Press 'h' for instructions. (unimplemented)]",
                    status, key
                )
            }
            (true, None) => format!("{}[Press space to continue, 'q' to quit.]", status),
            (false, Some(_)) => format!("{}{}", status, BELL),
            (false, None) => status,
        };

        write!(
            stdout,
            "\r{}{}{}",
            Attribute::Reverse,
            banner,
            Attribute::Reset
        )
        .unwrap();
    }
}

// Break the lines on the cols of the terminal
fn break_buff(buff: &str, cols: usize) -> Vec<String> {
    let mut lines = Vec::with_capacity(buff.lines().count());

    for l in buff.lines() {
        lines.append(&mut break_line(l, cols));
    }
    lines
}

fn break_line(line: &str, cols: usize) -> Vec<String> {
    let width = UnicodeWidthStr::width(line);
    let mut lines = Vec::new();
    if width < cols {
        lines.push(line.to_string());
        return lines;
    }

    let gr_idx = UnicodeSegmentation::grapheme_indices(line, true);
    let mut last_index = 0;
    let mut total_width = 0;
    for (index, grapheme) in gr_idx {
        let width = UnicodeWidthStr::width(grapheme);
        total_width += width;

        if total_width > cols {
            lines.push(line[last_index..index].to_string());
            last_index = index;
            total_width = width;
        }
    }

    if last_index != line.len() {
        lines.push(line[last_index..].to_string());
    }
    lines
}

#[cfg(test)]
mod tests {
    use super::break_line;
    use unicode_width::UnicodeWidthStr;

    #[test]
    fn test_break_lines_long() {
        let mut test_string = String::with_capacity(100);
        for _ in 0..200 {
            test_string.push('#');
        }

        let lines = break_line(&test_string, 80);
        let widths: Vec<usize> = lines
            .iter()
            .map(|s| UnicodeWidthStr::width(&s[..]))
            .collect();

        assert_eq!((80, 80, 40), (widths[0], widths[1], widths[2]));
    }

    #[test]
    fn test_break_lines_short() {
        let mut test_string = String::with_capacity(100);
        for _ in 0..20 {
            test_string.push('#');
        }

        let lines = break_line(&test_string, 80);

        assert_eq!(20, lines[0].len());
    }

    #[test]
    fn test_break_line_zwj() {
        let mut test_string = String::with_capacity(1100);
        for _ in 0..20 {
            test_string.push_str("👩🏻🔬");
        }

        let lines = break_line(&test_string, 80);

        let widths: Vec<usize> = lines
            .iter()
            .map(|s| UnicodeWidthStr::width(&s[..]))
            .collect();

        // Each 👩🏻🔬 is 6 character width it break line to the closest number to 80 => 6 * 13 = 78
        assert_eq!((78, 42), (widths[0], widths[1]));
    }
}
<|endoftext|>


<commit_before>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Alex Lyon <arcterus@mail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.

// spell-checker:ignore (ToDO) getpriority execvp setpriority nstr PRIO cstrs ENOENT

#[macro_use]
extern crate uucore;

use libc::{c_char, c_int, execvp, PRIO_PROCESS};
use std::ffi::CString;
use std::io::Error;
use std::ptr;

use clap::{crate_version, Arg, Command};
use uucore::{
    error::{set_exit_code, UClapError, UResult, USimpleError, UUsageError},
    format_usage,
};

pub mod options {
    pub static ADJUSTMENT: &str = "adjustment";
    pub static COMMAND: &str = "COMMAND";
}

const ABOUT: &str = "\
    Run COMMAND with an adjusted niceness, which affects process scheduling. \
    With no COMMAND, print the current niceness.  Niceness values range from at \
    least -20 (most favorable to the process) to 19 (least favorable to the \
    process).";
const USAGE: &str = "{} [OPTIONS] [COMMAND [ARGS]]";

#[uucore::main]
pub fn uumain(args: impl uucore::Args) -> UResult<()> {
    let matches = uu_app().try_get_matches_from(args).with_exit_code(125)?;

    let mut niceness = unsafe {
        nix::errno::Errno::clear();
        libc::getpriority(PRIO_PROCESS, 0)
    };
    if Error::last_os_error().raw_os_error().unwrap() != 0 {
        return Err(USimpleError::new(
            125,
            format!("getpriority: {}", Error::last_os_error()),
        ));
    }

    let adjustment = match matches.value_of(options::ADJUSTMENT) {
        Some(nstr) => {
            if !matches.is_present(options::COMMAND) {
                return Err(UUsageError::new(
                    125,
                    "A command must be given with an adjustment.",
                ));
            }
            match nstr.parse() {
                Ok(num) => num,
                Err(e) => {
                    return Err(USimpleError::new(
                        125,
                        format!("\"{}\" is not a valid number: {}", nstr, e),
                    ))
                }
            }
        }
        None => {
            if !matches.is_present(options::COMMAND) {
                println!("{}", niceness);
                return Ok(());
            }
            10_i32
        }
    };

    niceness += adjustment;
    if unsafe { libc::setpriority(PRIO_PROCESS, 0, niceness) } == -1 {
        show_warning!("setpriority: {}", Error::last_os_error());
    }

    let cstrs: Vec<CString> = matches
        .values_of(options::COMMAND)
        .unwrap()
        .map(|x| CString::new(x.as_bytes()).unwrap())
        .collect();

    let mut args: Vec<*const c_char> = cstrs.iter().map(|s| s.as_ptr()).collect();
    args.push(ptr::null::<c_char>());
    unsafe {
        execvp(args[0], args.as_mut_ptr());
    }

    show_error!("execvp: {}", Error::last_os_error());
    let exit_code = if Error::last_os_error().raw_os_error().unwrap() as c_int == libc::ENOENT {
        127
    } else {
        126
    };
    set_exit_code(exit_code);
    Ok(())
}

pub fn uu_app<'a>() -> Command<'a> {
    Command::new(uucore::util_name())
        .about(ABOUT)
        .override_usage(format_usage(USAGE))
        .trailing_var_arg(true)
        .infer_long_args(true)
        .version(crate_version!())
        .arg(
            Arg::new(options::ADJUSTMENT)
                .short('n')
                .long(options::ADJUSTMENT)
                .help("add N to the niceness (default is 10)")
                .takes_value(true)
                .allow_hyphen_values(true),
        )
        .arg(
            Arg::new(options::COMMAND)
                .multiple_occurrences(true)
                .value_hint(clap::ValueHint::CommandName),
        )
}
<commit_msg>nice: Move call to Errno::clear() outside of unsafe block<commit_after>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Alex Lyon <arcterus@mail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.

// spell-checker:ignore (ToDO) getpriority execvp setpriority nstr PRIO cstrs ENOENT

#[macro_use]
extern crate uucore;

use libc::{c_char, c_int, execvp, PRIO_PROCESS};
use std::ffi::CString;
use std::io::Error;
use std::ptr;

use clap::{crate_version, Arg, Command};
use uucore::{
    error::{set_exit_code, UClapError, UResult, USimpleError, UUsageError},
    format_usage,
};

pub mod options {
    pub static ADJUSTMENT: &str = "adjustment";
    pub static COMMAND: &str = "COMMAND";
}

const ABOUT: &str = "\
    Run COMMAND with an adjusted niceness, which affects process scheduling. \
    With no COMMAND, print the current niceness.  Niceness values range from at \
    least -20 (most favorable to the process) to 19 (least favorable to the \
    process).";
const USAGE: &str = "{} [OPTIONS] [COMMAND [ARGS]]";

#[uucore::main]
pub fn uumain(args: impl uucore::Args) -> UResult<()> {
    let matches = uu_app().try_get_matches_from(args).with_exit_code(125)?;

    nix::errno::Errno::clear();
    let mut niceness = unsafe { libc::getpriority(PRIO_PROCESS, 0) };
    if Error::last_os_error().raw_os_error().unwrap() != 0 {
        return Err(USimpleError::new(
            125,
            format!("getpriority: {}", Error::last_os_error()),
        ));
    }

    let adjustment = match matches.value_of(options::ADJUSTMENT) {
        Some(nstr) => {
            if !matches.is_present(options::COMMAND) {
                return Err(UUsageError::new(
                    125,
                    "A command must be given with an adjustment.",
                ));
            }
            match nstr.parse() {
                Ok(num) => num,
                Err(e) => {
                    return Err(USimpleError::new(
                        125,
                        format!("\"{}\" is not a valid number: {}", nstr, e),
                    ))
                }
            }
        }
        None => {
            if !matches.is_present(options::COMMAND) {
                println!("{}", niceness);
                return Ok(());
            }
            10_i32
        }
    };

    niceness += adjustment;
    if unsafe { libc::setpriority(PRIO_PROCESS, 0, niceness) } == -1 {
        show_warning!("setpriority: {}", Error::last_os_error());
    }

    let cstrs: Vec<CString> = matches
        .values_of(options::COMMAND)
        .unwrap()
        .map(|x| CString::new(x.as_bytes()).unwrap())
        .collect();

    let mut args: Vec<*const c_char> = cstrs.iter().map(|s| s.as_ptr()).collect();
    args.push(ptr::null::<c_char>());
    unsafe {
        execvp(args[0], args.as_mut_ptr());
    }

    show_error!("execvp: {}", Error::last_os_error());
    let exit_code = if Error::last_os_error().raw_os_error().unwrap() as c_int == libc::ENOENT {
        127
    } else {
        126
    };
    set_exit_code(exit_code);
    Ok(())
}

pub fn uu_app<'a>() -> Command<'a> {
    Command::new(uucore::util_name())
        .about(ABOUT)
        .override_usage(format_usage(USAGE))
        .trailing_var_arg(true)
        .infer_long_args(true)
        .version(crate_version!())
        .arg(
            Arg::new(options::ADJUSTMENT)
                .short('n')
                .long(options::ADJUSTMENT)
                .help("add N to the niceness (default is 10)")
                .takes_value(true)
                .allow_hyphen_values(true),
        )
        .arg(
            Arg::new(options::COMMAND)
                .multiple_occurrences(true)
                .value_hint(clap::ValueHint::CommandName),
        )
}
<|endoftext|>


<commit_before>use crate::logger;
use crate::theme;
use crate::vec::Vec2;
use crate::view::View;
use crate::Printer;

use unicode_width::UnicodeWidthStr as _;

/// View used for debugging, showing logs.
pub struct DebugView {
    // TODO: wrap log lines if needed, and save the line splits here.
}

impl DebugView {
    /// Creates a new DebugView.
    pub fn new() -> Self {
        DebugView {}
    }
}

impl View for DebugView {
    fn draw(&self, printer: &Printer<'_, '_>) {
        let logs = logger::LOGS.lock().unwrap();
        // Only print the last logs, so skip what doesn't fit
        let skipped = logs.len().saturating_sub(printer.size.y);

        for (i, record) in logs.iter().skip(skipped).enumerate() {
            // TODO: Apply style to message? (Ex: errors in bold?)
            // TODO: customizable time format? (24h/AM-PM)
            printer.print(
                (0, i),
                &format!(
                    "{} | [     ] {}",
                    record.time.with_timezone(&chrono::Local).format("%T%.3f"),
                    record.message
                ),
            );
            let color = match record.level {
                log::Level::Error => theme::BaseColor::Red.dark(),
                log::Level::Warn => theme::BaseColor::Yellow.dark(),
                log::Level::Info => theme::BaseColor::Black.light(),
                log::Level::Debug => theme::BaseColor::Green.dark(),
                log::Level::Trace => theme::BaseColor::Blue.dark(),
            };
            printer.with_color(color.into(), |printer| {
                printer.print((16, i), &format!("{:5}", record.level))
            });
        }
    }

    fn required_size(&mut self, _constraint: Vec2) -> Vec2 {
        // TODO: read the logs, and compute the required size to print it.
        let logs = logger::LOGS.lock().unwrap();

        let level_width = 8; // Width of "[ERROR] "
        let time_width = 16; // Width of "23:59:59.123 | "

        // The longest line sets the width
        let w = logs
            .iter()
            .map(|record| record.message.width() + level_width + time_width)
            .max()
            .unwrap_or(1);
        let h = logs.len();

        Vec2::new(w, h)
    }

    fn layout(&mut self, _size: Vec2) {
        // Uh?
    }
}
<commit_msg>Revert "use ... as _" to stay compatible with rust 1.31<commit_after>use crate::logger;
use crate::theme;
use crate::vec::Vec2;
use crate::view::View;
use crate::Printer;

use unicode_width::UnicodeWidthStr;

/// View used for debugging, showing logs.
pub struct DebugView {
    // TODO: wrap log lines if needed, and save the line splits here.
}

impl DebugView {
    /// Creates a new DebugView.
    pub fn new() -> Self {
        DebugView {}
    }
}

impl View for DebugView {
    fn draw(&self, printer: &Printer<'_, '_>) {
        let logs = logger::LOGS.lock().unwrap();
        // Only print the last logs, so skip what doesn't fit
        let skipped = logs.len().saturating_sub(printer.size.y);

        for (i, record) in logs.iter().skip(skipped).enumerate() {
            // TODO: Apply style to message? (Ex: errors in bold?)
            // TODO: customizable time format? (24h/AM-PM)
            printer.print(
                (0, i),
                &format!(
                    "{} | [     ] {}",
                    record.time.with_timezone(&chrono::Local).format("%T%.3f"),
                    record.message
                ),
            );
            let color = match record.level {
                log::Level::Error => theme::BaseColor::Red.dark(),
                log::Level::Warn => theme::BaseColor::Yellow.dark(),
                log::Level::Info => theme::BaseColor::Black.light(),
                log::Level::Debug => theme::BaseColor::Green.dark(),
                log::Level::Trace => theme::BaseColor::Blue.dark(),
            };
            printer.with_color(color.into(), |printer| {
                printer.print((16, i), &format!("{:5}", record.level))
            });
        }
    }

    fn required_size(&mut self, _constraint: Vec2) -> Vec2 {
        // TODO: read the logs, and compute the required size to print it.
        let logs = logger::LOGS.lock().unwrap();

        let level_width = 8; // Width of "[ERROR] "
        let time_width = 16; // Width of "23:59:59.123 | "

        // The longest line sets the width
        let w = logs
            .iter()
            .map(|record| record.message.width() + level_width + time_width)
            .max()
            .unwrap_or(1);
        let h = logs.len();

        Vec2::new(w, h)
    }

    fn layout(&mut self, _size: Vec2) {
        // Uh?
    }
}
<|endoftext|>