<commit_before>use crate::gssapi::*;
use crate::rparser::*;
use crate::sasl::parse_sasl_buffer;
use crate::tls::TlsParser;
use crate::{gen_get_variants, Variant};
use ldap_parser::asn1_rs::{Error, Header};
use ldap_parser::ldap::*;
use ldap_parser::*;
use std::collections::BTreeMap;

pub struct LDAPBuilder {}
impl RBuilder for LDAPBuilder {
    fn build(&self) -> Box<dyn RParser> {
        Box::new(LDAPParser::new(b"LDAP"))
    }
    fn get_l4_probe(&self) -> Option<ProbeL4> {
        Some(ldap_probe)
    }
}

pub struct LDAPParser<'a> {
    _name: Option<&'a [u8]>,
    has_sasl_layers: bool,
    bind_dn: Option<String>,
    sasl_mech: Option<String>,
    // operation tag of current request, if present
    protocol_op: Option<ProtocolOpTag>,
    // message ID of LDAP request, if present
    ldap_request: Option<MessageID>,
    // numeric result code of LDAP response, if present
    ldap_result_code: Option<u32>,
    // message ID of TLS request, if present
    tls_request: Option<MessageID>,
    has_starttls: bool,
    transactions: BTreeMap<MessageID, LDAPTransaction>,

    defrag_buffer: Vec<u8>,
    tls_parser: Option<TlsParser<'a>>,
}

#[allow(unused)]
#[derive(Debug)]
pub struct LDAPTransaction {
    message_id: MessageID,
    tx_op: ProtocolOpTag,
    result_code: Option<u32>,
    bind_dn: Option<String>,
}

impl<'a> RParser for LDAPParser<'a> {
    fn parse_l4(&mut self, data: &[u8], direction: Direction) -> ParseResult {
        if self.has_starttls {
            return self.parse_ldap_starttls(data, direction);
        }
        // debug: check if we missed TLS
        if self.defrag_buffer.is_empty() && data.len() >= 2 && &data[..2] == b"\x16\x03" {
            warn!("LDAP: missed STARTTLS ?!");
        }
        if self.has_sasl_layers {
            self.parse_ldap_sasl(data, direction)
        } else {
            self.parse_ldap_messages(data, direction)
        }
    }

    gen_get_variants! {LDAPParser, "ldap.",
        ldap_request => |s| {
            s.ldap_request.map(|op| op.0.into())
        },
        ldap_result_code      => map,
        protocol_op => |s| {
            s.protocol_op.map(|op| Variant::OwnedStr(op.to_string()))
        },
        bind_dn               => map_as_ref,
        sasl_mech             => map_as_ref,
        has_starttls          => into,
        tls_client_version => |s| {
            s.tls_parser.as_ref().map(|p|
                p.get("tls.client_version")
            ).flatten()
        },
        tls_ssl_record_version => |s| {
            s.tls_parser.as_ref().map(|p|
                p.get("tls.ssl_record_version")
            ).flatten()
        },
        tls_compression => |s| {
            s.tls_parser.as_ref().map(|p|
                p.get("tls.compression")
            ).flatten()
        },
        tls_cipher => |s| {
            s.tls_parser.as_ref().map(|p|
                p.get("tls.cipher")
            ).flatten()
        },
    }
}

impl<'a> LDAPParser<'a> {
    pub fn new(name: &'a [u8]) -> LDAPParser<'a> {
        LDAPParser {
            _name: Some(name),
            has_sasl_layers: false,
            bind_dn: None,
            sasl_mech: None,
            tls_request: None,
            has_starttls: false,
            defrag_buffer: Vec::new(),
            ldap_request: None,
            ldap_result_code: None,
            transactions: BTreeMap::new(),
            protocol_op: None,
            tls_parser: None,
        }
    }

    pub fn parse_ldap_starttls(&mut self, data: &[u8], direction: Direction) -> ParseResult {
        if let Some(parser) = &mut self.tls_parser {
            parser.parse_tcp_level(data, direction)
        } else {
            warn!("Should not happen: could not instantiate TLS parser");
            ParseResult::Fatal
        }
    }

    pub fn parse_ldap_sasl(&mut self, data: &[u8], direction: Direction) -> ParseResult {
        // are we already waiting for more data (defragmentation)?
        if !self.defrag_buffer.is_empty() {
            return self.parse_ldap_messages(data, direction);
        }
        // are we really using SASL?
        // sometimes, SASL is negotiated but the next packets are sent in clear
        //
        // check if first byte is 0 (sasl layer length is 4 bytes, so this is OK if the packet does not exceed 16 MB)
        // if LDAP, first byte should be 0x30 anyway
        if !data.is_empty() && data[0] != 0 {
            warn!("XXX CLEAR TEXT LDAP after SASL");
            return self.parse_ldap_messages(data, direction);
        }
        let mech = match &self.sasl_mech {
            Some(s) => s,
            None => {
                debug!("SASL mechanism unknown, stopping following LDAP session");
                return ParseResult::Stop;
            }
        };
        match mech.as_ref() {
            "GSSAPI" | "GSS-SPNEGO" => {
                // use nom::HexDisplay;
                // debug!("\n{}", data.to_hex(16));
                let (_rem, sasl) = match parse_sasl_buffer(data) {
                    Ok(x) => x,
                    Err(e) => {
                        debug!("LDAP: parsing SASL buffer failed: {:?}", e);
                        return ParseResult::Error;
                    }
                };
                let (_rem, gss) = match parse_gssapi(sasl.0) {
                    Ok(x) => x,
                    Err(e) => {
                        debug!("LDAP: parsing GSS-API buffer failed: {:?}", e);
                        return ParseResult::Error;
                    }
                };
                match &gss {
                    GssApiBuffer::Wrapped(buffer) => {
                        if buffer.token.seal_alg == 0xffff {
                            // no encryption, only integrity
                            // debug!("Integrity-only GSS-API KRB5 encapsulation detected");
                            // XXX remove padding
                            //
                            // integrity-only messages seem to always end up by hex '01', though I could not find why
                            self.parse_ldap_messages(buffer.payload, direction)
                        } else {
                            trace!(
                                "Encrypted GSS-API KRB5 encapsulation detected, activating bypass"
                            );
                            ParseResult::Stop
                        }
                    }
                    GssApiBuffer::GssCfxWrap { token, payload } => {
                        if token.flags & KRB_CFX_FLAG_SEALED != 0 {
                            trace!(
                                "Encrypted GSS-API CFX KRB5 encapsulation detected, activating bypass"
                            );
                            return ParseResult::Stop;
                        }
                        // no encryption, only integrity
                        // debug!("Integrity-only GSS-API KRB5 encapsulation detected");
                        // XXX remove padding
                        //
                        self.parse_ldap_messages(payload, direction)
                    }
                    _ => {
                        trace!("Unknown GSS-API encapsulation, activating bypass");
                        ParseResult::Stop
                    }
                }
            }
            _ => {
                trace!("Unknown SASL layer {}, activating bypass", mech);
                ParseResult::Stop
            }
        }
    }

    pub fn parse_ldap_messages(&mut self, data: &[u8], direction: Direction) -> ParseResult {
        let mut v: Vec<u8>;
        // Check if TCP data is being defragmented
        let defrag_buffer = match self.defrag_buffer.len() {
            0 => data,
            _ => {
                // sanity check vector length to avoid memory exhaustion
                if self.defrag_buffer.len() + data.len() > 16_777_216 {
                    warn!("LDAP: defrag buffer overflow");
                    return ParseResult::Error;
                };
                v = self.defrag_buffer.split_off(0);
                v.extend_from_slice(data);
                v.as_slice()
            }
        };
        let mut cur_i = defrag_buffer;
        loop {
            if cur_i.is_empty() {
                break;
            }
            // read header and see if we have a complete message
            if check_ldap_message_complete(cur_i).is_err() {
                // XXX if bytes are exactly "01", drop? (often happens with GSS-API)
                if cur_i.len() == 1
                    && cur_i[0] == 1
                    && self.sasl_mech == Some("GSS-SPNEGO".to_string())
                {
                    warn!("Fixup for SASL/GSS-API: removing 1 byte");
                    // cur_i = &cur_i[1..];
                    break;
                }
                trace!("Fragmentation required (TCP level)");
                // if not complete, add to defrag buffer
                self.defrag_buffer.extend_from_slice(cur_i);
                // panic!("XXX Fragmented LDAP");
                break;
            }
            match LdapMessage::from_ber(cur_i) {
                Ok((rem, msg)) => {
                    trace!("parse_ldap_message: {:?}", msg);
                    self.handle_message(&msg, direction);
                    cur_i = rem;
                }
                e => {
                    warn!("parse_ldap_message failed: {:?}", e);
                    break;
                }
            };
        }
        trace!("current transactions: {:?}", self.transactions);
        // if !cur_i.is_empty() {
        //     // XXX if bytes are exactly "01", drop? (often happens with GSS-API)
        //     warn!("Extra bytes after LDAP data");
        //     use nom::HexDisplay;
        //     println!("{}", cur_i.to_hex(16));
        // }
        ParseResult::Ok
    }

    pub fn handle_message(&mut self, msg: &LdapMessage, _direction: Direction) {
        let tx_op = msg.protocol_op.tag();
        let bind_dn = &self.bind_dn;
        let mut tx = self
            .transactions
            .entry(msg.message_id)
            .or_insert_with(|| LDAPTransaction::new(msg.message_id, tx_op, bind_dn));
        self.protocol_op = Some(tx_op);
        match &msg.protocol_op {
            ProtocolOp::BindRequest(req) => {
                // store msg_id
                self.ldap_request = Some(msg.message_id);
                // store bind DN
                self.bind_dn = Some(req.name.0.to_string());
                tx.bind_dn = Some(req.name.0.to_string());
                // if SASL, store mechanism
                if let AuthenticationChoice::Sasl(creds) = &req.authentication {
                    self.sasl_mech = Some(creds.mechanism.0.to_string());
                    if let Some(data) = &creds.credentials {
                        if !data.is_empty() {
                            use nom::HexDisplay;
                            println!("XXX\n{}", data.to_hex(16));
                            let (_, xxx) = parse_gssapi(data).unwrap(); // XXX
                            println!("{:?}", xxx);
                        }
                    }
                } else {
                    // XXX one connection can bind multiple times (and/or do stuff unbound)
                    self.sasl_mech = None;
                }
            }
            ProtocolOp::BindResponse(resp) => {
                self.ldap_result_code = Some(resp.result.result_code.0);
                tx.result_code = Some(resp.result.result_code.0);
                // check msg_id
                // check for success
                if resp.result.result_code == ResultCode::Success {
                    // if SASL authentication succeeded, SASL layers are now active
                    if self.sasl_mech.is_some() {
                        self.has_sasl_layers = true;
                    }
                } else {
                    self.bind_dn = None;
                }
            }
            ProtocolOp::ExtendedRequest(req) => {
                self.ldap_request = Some(msg.message_id);
                // check for STARTTLS
                if req.request_name.0 == "1.3.6.1.4.1.1466.20037" {
                    self.tls_request = Some(msg.message_id);
                }
            }
            ProtocolOp::ExtendedResponse(res) => {
                if let Some(req_id) = self.tls_request {
                    self.ldap_result_code = Some(res.result.result_code.0);
                    tx.result_code = Some(res.result.result_code.0);
                    if res.result.result_code == ResultCode::Success && msg.message_id == req_id {
                        self.has_starttls = true;
                        self.tls_parser = Some(TlsParser::new(b"LDAP/TLS"));
                    }
                }
            }
            ProtocolOp::SearchRequest(_)
            | ProtocolOp::ModifyRequest(_)
            | ProtocolOp::AddRequest(_)
            | ProtocolOp::ModDnRequest(_)
            | ProtocolOp::CompareRequest(_)
            | ProtocolOp::AbandonRequest(_) => {
                self.ldap_request = Some(msg.message_id);
            }
            ProtocolOp::ModifyResponse(r) => {
                self.ldap_result_code = Some(r.result.result_code.0);
                tx.result_code = Some(r.result.result_code.0);
            }
            ProtocolOp::SearchResultDone(r)
            | ProtocolOp::AddResponse(r)
            | ProtocolOp::ModDnResponse(r)
            | ProtocolOp::CompareResponse(r) => {
                self.ldap_result_code = Some(r.result_code.0);
                tx.result_code = Some(r.result_code.0);
            }
            _ => (),
        }
    }
}

impl LDAPTransaction {
    pub fn new(message_id: MessageID, tx_op: ProtocolOpTag, bind_dn: &Option<String>) -> Self {
        LDAPTransaction {
            message_id,
            tx_op,
            result_code: None,
            bind_dn: bind_dn.clone(),
        }
    }
}

const COMMON_LDAP_PORTS: &[u16] = &[389, 3268];

pub fn ldap_probe(i: &[u8], l4info: &L4Info) -> ProbeResult {
    if i.len() <= 2 {
        return ProbeResult::Unsure;
    }
    // debug!("Probing for LDAP");
    if LdapMessage::from_ber(i).is_ok() {
        if COMMON_LDAP_PORTS.contains(&l4info.src_port)
            || COMMON_LDAP_PORTS.contains(&l4info.dst_port)
        {
            ProbeResult::Certain
        } else {
            ProbeResult::Unsure
        }
    } else {
        ProbeResult::NotForUs
    }
}

fn check_ldap_message_complete(i: &[u8]) -> asn1_rs::ParseResult<()> {
    let (i, header) = Header::from_ber(i)?;
    let len = header.length().definite()?;
    if i.len() >= len {
        Ok((i, ()))
    } else {
        Err(Error::InvalidLength.into())
    }
}
<commit_msg>Fix clippy warnings (replace .map().flatten() with .and_then())<commit_after>use crate::gssapi::*;
use crate::rparser::*;
use crate::sasl::parse_sasl_buffer;
use crate::tls::TlsParser;
use crate::{gen_get_variants, Variant};
use ldap_parser::asn1_rs::{Error, Header};
use ldap_parser::ldap::*;
use ldap_parser::*;
use std::collections::BTreeMap;

pub struct LDAPBuilder {}
impl RBuilder for LDAPBuilder {
    fn build(&self) -> Box<dyn RParser> {
        Box::new(LDAPParser::new(b"LDAP"))
    }
    fn get_l4_probe(&self) -> Option<ProbeL4> {
        Some(ldap_probe)
    }
}

pub struct LDAPParser<'a> {
    _name: Option<&'a [u8]>,
    has_sasl_layers: bool,
    bind_dn: Option<String>,
    sasl_mech: Option<String>,
    // operation tag of current request, if present
    protocol_op: Option<ProtocolOpTag>,
    // message ID of LDAP request, if present
    ldap_request: Option<MessageID>,
    // numeric result code of LDAP response, if present
    ldap_result_code: Option<u32>,
    // message ID of TLS request, if present
    tls_request: Option<MessageID>,
    has_starttls: bool,
    transactions: BTreeMap<MessageID, LDAPTransaction>,

    defrag_buffer: Vec<u8>,
    tls_parser: Option<TlsParser<'a>>,
}

#[allow(unused)]
#[derive(Debug)]
pub struct LDAPTransaction {
    message_id: MessageID,
    tx_op: ProtocolOpTag,
    result_code: Option<u32>,
    bind_dn: Option<String>,
}

impl<'a> RParser for LDAPParser<'a> {
    fn parse_l4(&mut self, data: &[u8], direction: Direction) -> ParseResult {
        if self.has_starttls {
            return self.parse_ldap_starttls(data, direction);
        }
        // debug: check if we missed TLS
        if self.defrag_buffer.is_empty() && data.len() >= 2 && &data[..2] == b"\x16\x03" {
            warn!("LDAP: missed STARTTLS ?!");
        }
        if self.has_sasl_layers {
            self.parse_ldap_sasl(data, direction)
        } else {
            self.parse_ldap_messages(data, direction)
        }
    }

    gen_get_variants! {LDAPParser, "ldap.",
        ldap_request => |s| {
            s.ldap_request.map(|op| op.0.into())
        },
        ldap_result_code      => map,
        protocol_op => |s| {
            s.protocol_op.map(|op| Variant::OwnedStr(op.to_string()))
        },
        bind_dn               => map_as_ref,
        sasl_mech             => map_as_ref,
        has_starttls          => into,
        tls_client_version => |s| {
            s.tls_parser.as_ref().and_then(|p|
                p.get("tls.client_version"))
        },
        tls_ssl_record_version => |s| {
            s.tls_parser.as_ref().and_then(|p|
                p.get("tls.ssl_record_version"))
        },
        tls_compression => |s| {
            s.tls_parser.as_ref().and_then(|p|
                p.get("tls.compression"))
        },
        tls_cipher => |s| {
            s.tls_parser.as_ref().and_then(|p|
                p.get("tls.cipher"))
        },
    }
}

impl<'a> LDAPParser<'a> {
    pub fn new(name: &'a [u8]) -> LDAPParser<'a> {
        LDAPParser {
            _name: Some(name),
            has_sasl_layers: false,
            bind_dn: None,
            sasl_mech: None,
            tls_request: None,
            has_starttls: false,
            defrag_buffer: Vec::new(),
            ldap_request: None,
            ldap_result_code: None,
            transactions: BTreeMap::new(),
            protocol_op: None,
            tls_parser: None,
        }
    }

    pub fn parse_ldap_starttls(&mut self, data: &[u8], direction: Direction) -> ParseResult {
        if let Some(parser) = &mut self.tls_parser {
            parser.parse_tcp_level(data, direction)
        } else {
            warn!("Should not happen: could not instantiate TLS parser");
            ParseResult::Fatal
        }
    }

    pub fn parse_ldap_sasl(&mut self, data: &[u8], direction: Direction) -> ParseResult {
        // are we already waiting for more data (defragmentation)?
        if !self.defrag_buffer.is_empty() {
            return self.parse_ldap_messages(data, direction);
        }
        // are we really using SASL?
        // sometimes, SASL is negotiated but the next packets are sent in clear
        //
        // check if first byte is 0 (sasl layer length is 4 bytes, so this is OK if the packet does not exceed 16 MB)
        // if LDAP, first byte should be 0x30 anyway
        if !data.is_empty() && data[0] != 0 {
            warn!("XXX CLEAR TEXT LDAP after SASL");
            return self.parse_ldap_messages(data, direction);
        }
        let mech = match &self.sasl_mech {
            Some(s) => s,
            None => {
                debug!("SASL mechanism unknown, stopping following LDAP session");
                return ParseResult::Stop;
            }
        };
        match mech.as_ref() {
            "GSSAPI" | "GSS-SPNEGO" => {
                // use nom::HexDisplay;
                // debug!("\n{}", data.to_hex(16));
                let (_rem, sasl) = match parse_sasl_buffer(data) {
                    Ok(x) => x,
                    Err(e) => {
                        debug!("LDAP: parsing SASL buffer failed: {:?}", e);
                        return ParseResult::Error;
                    }
                };
                let (_rem, gss) = match parse_gssapi(sasl.0) {
                    Ok(x) => x,
                    Err(e) => {
                        debug!("LDAP: parsing GSS-API buffer failed: {:?}", e);
                        return ParseResult::Error;
                    }
                };
                match &gss {
                    GssApiBuffer::Wrapped(buffer) => {
                        if buffer.token.seal_alg == 0xffff {
                            // no encryption, only integrity
                            // debug!("Integrity-only GSS-API KRB5 encapsulation detected");
                            // XXX remove padding
                            //
                            // integrity-only messages seem to always end up by hex '01', though I could not find why
                            self.parse_ldap_messages(buffer.payload, direction)
                        } else {
                            trace!(
                                "Encrypted GSS-API KRB5 encapsulation detected, activating bypass"
                            );
                            ParseResult::Stop
                        }
                    }
                    GssApiBuffer::GssCfxWrap { token, payload } => {
                        if token.flags & KRB_CFX_FLAG_SEALED != 0 {
                            trace!(
                                "Encrypted GSS-API CFX KRB5 encapsulation detected, activating bypass"
                            );
                            return ParseResult::Stop;
                        }
                        // no encryption, only integrity
                        // debug!("Integrity-only GSS-API KRB5 encapsulation detected");
                        // XXX remove padding
                        //
                        self.parse_ldap_messages(payload, direction)
                    }
                    _ => {
                        trace!("Unknown GSS-API encapsulation, activating bypass");
                        ParseResult::Stop
                    }
                }
            }
            _ => {
                trace!("Unknown SASL layer {}, activating bypass", mech);
                ParseResult::Stop
            }
        }
    }

    pub fn parse_ldap_messages(&mut self, data: &[u8], direction: Direction) -> ParseResult {
        let mut v: Vec<u8>;
        // Check if TCP data is being defragmented
        let defrag_buffer = match self.defrag_buffer.len() {
            0 => data,
            _ => {
                // sanity check vector length to avoid memory exhaustion
                if self.defrag_buffer.len() + data.len() > 16_777_216 {
                    warn!("LDAP: defrag buffer overflow");
                    return ParseResult::Error;
                };
                v = self.defrag_buffer.split_off(0);
                v.extend_from_slice(data);
                v.as_slice()
            }
        };
        let mut cur_i = defrag_buffer;
        loop {
            if cur_i.is_empty() {
                break;
            }
            // read header and see if we have a complete message
            if check_ldap_message_complete(cur_i).is_err() {
                // XXX if bytes are exactly "01", drop? (often happens with GSS-API)
                if cur_i.len() == 1
                    && cur_i[0] == 1
                    && self.sasl_mech == Some("GSS-SPNEGO".to_string())
                {
                    warn!("Fixup for SASL/GSS-API: removing 1 byte");
                    // cur_i = &cur_i[1..];
                    break;
                }
                trace!("Fragmentation required (TCP level)");
                // if not complete, add to defrag buffer
                self.defrag_buffer.extend_from_slice(cur_i);
                // panic!("XXX Fragmented LDAP");
                break;
            }
            match LdapMessage::from_ber(cur_i) {
                Ok((rem, msg)) => {
                    trace!("parse_ldap_message: {:?}", msg);
                    self.handle_message(&msg, direction);
                    cur_i = rem;
                }
                e => {
                    warn!("parse_ldap_message failed: {:?}", e);
                    break;
                }
            };
        }
        trace!("current transactions: {:?}", self.transactions);
        // if !cur_i.is_empty() {
        //     // XXX if bytes are exactly "01", drop? (often happens with GSS-API)
        //     warn!("Extra bytes after LDAP data");
        //     use nom::HexDisplay;
        //     println!("{}", cur_i.to_hex(16));
        // }
        ParseResult::Ok
    }

    pub fn handle_message(&mut self, msg: &LdapMessage, _direction: Direction) {
        let tx_op = msg.protocol_op.tag();
        let bind_dn = &self.bind_dn;
        let mut tx = self
            .transactions
            .entry(msg.message_id)
            .or_insert_with(|| LDAPTransaction::new(msg.message_id, tx_op, bind_dn));
        self.protocol_op = Some(tx_op);
        match &msg.protocol_op {
            ProtocolOp::BindRequest(req) => {
                // store msg_id
                self.ldap_request = Some(msg.message_id);
                // store bind DN
                self.bind_dn = Some(req.name.0.to_string());
                tx.bind_dn = Some(req.name.0.to_string());
                // if SASL, store mechanism
                if let AuthenticationChoice::Sasl(creds) = &req.authentication {
                    self.sasl_mech = Some(creds.mechanism.0.to_string());
                    if let Some(data) = &creds.credentials {
                        if !data.is_empty() {
                            use nom::HexDisplay;
                            println!("XXX\n{}", data.to_hex(16));
                            let (_, xxx) = parse_gssapi(data).unwrap(); // XXX
                            println!("{:?}", xxx);
                        }
                    }
                } else {
                    // XXX one connection can bind multiple times (and/or do stuff unbound)
                    self.sasl_mech = None;
                }
            }
            ProtocolOp::BindResponse(resp) => {
                self.ldap_result_code = Some(resp.result.result_code.0);
                tx.result_code = Some(resp.result.result_code.0);
                // check msg_id
                // check for success
                if resp.result.result_code == ResultCode::Success {
                    // if SASL authentication succeeded, SASL layers are now active
                    if self.sasl_mech.is_some() {
                        self.has_sasl_layers = true;
                    }
                } else {
                    self.bind_dn = None;
                }
            }
            ProtocolOp::ExtendedRequest(req) => {
                self.ldap_request = Some(msg.message_id);
                // check for STARTTLS
                if req.request_name.0 == "1.3.6.1.4.1.1466.20037" {
                    self.tls_request = Some(msg.message_id);
                }
            }
            ProtocolOp::ExtendedResponse(res) => {
                if let Some(req_id) = self.tls_request {
                    self.ldap_result_code = Some(res.result.result_code.0);
                    tx.result_code = Some(res.result.result_code.0);
                    if res.result.result_code == ResultCode::Success && msg.message_id == req_id {
                        self.has_starttls = true;
                        self.tls_parser = Some(TlsParser::new(b"LDAP/TLS"));
                    }
                }
            }
            ProtocolOp::SearchRequest(_)
            | ProtocolOp::ModifyRequest(_)
            | ProtocolOp::AddRequest(_)
            | ProtocolOp::ModDnRequest(_)
            | ProtocolOp::CompareRequest(_)
            | ProtocolOp::AbandonRequest(_) => {
                self.ldap_request = Some(msg.message_id);
            }
            ProtocolOp::ModifyResponse(r) => {
                self.ldap_result_code = Some(r.result.result_code.0);
                tx.result_code = Some(r.result.result_code.0);
            }
            ProtocolOp::SearchResultDone(r)
            | ProtocolOp::AddResponse(r)
            | ProtocolOp::ModDnResponse(r)
            | ProtocolOp::CompareResponse(r) => {
                self.ldap_result_code = Some(r.result_code.0);
                tx.result_code = Some(r.result_code.0);
            }
            _ => (),
        }
    }
}

impl LDAPTransaction {
    pub fn new(message_id: MessageID, tx_op: ProtocolOpTag, bind_dn: &Option<String>) -> Self {
        LDAPTransaction {
            message_id,
            tx_op,
            result_code: None,
            bind_dn: bind_dn.clone(),
        }
    }
}

const COMMON_LDAP_PORTS: &[u16] = &[389, 3268];

pub fn ldap_probe(i: &[u8], l4info: &L4Info) -> ProbeResult {
    if i.len() <= 2 {
        return ProbeResult::Unsure;
    }
    // debug!("Probing for LDAP");
    if LdapMessage::from_ber(i).is_ok() {
        if COMMON_LDAP_PORTS.contains(&l4info.src_port)
            || COMMON_LDAP_PORTS.contains(&l4info.dst_port)
        {
            ProbeResult::Certain
        } else {
            ProbeResult::Unsure
        }
    } else {
        ProbeResult::NotForUs
    }
}

fn check_ldap_message_complete(i: &[u8]) -> asn1_rs::ParseResult<()> {
    let (i, header) = Header::from_ber(i)?;
    let len = header.length().definite()?;
    if i.len() >= len {
        Ok((i, ()))
    } else {
        Err(Error::InvalidLength.into())
    }
}
<|endoftext|>


<commit_before>use id::{Id, IdIndex, MAXIMUM_CAPACITY};
use flat::{FlatAccess, FlatAccessMut, Flat, FlatGet, FlatGetMut};

pub type IdMapVec<E, T> = IdMap<E, Vec<T>>;

#[derive(Debug)]
pub struct IdMap<E, F: Flat> {
    lookup: Vec<Id<E>>,
    reverse_lookup: Vec<IdIndex>,
    flat: F,
}

impl<E, F: Flat> IdMap<E, F> {
    pub fn new() -> Self {
        IdMap {
            lookup: Vec::new(),
            reverse_lookup: Vec::new(),
            flat: F::new(),
        }
    }

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

    pub fn with_capacity(capacity: usize) -> Self {
        assert!(capacity <= MAXIMUM_CAPACITY);
        IdMap {
            lookup: Vec::with_capacity(capacity),
            reverse_lookup: Vec::with_capacity(capacity),
            flat: F::with_capacity(capacity),
        }
    }

    #[inline]
    pub fn insert(&mut self, id: Id<E>, mut element: F::Element) -> Option<F::Element> {
        let IdMap {
            ref mut lookup,
            ref mut reverse_lookup,
            ref mut flat,
        } = *self;
        let new_index = {
            let len = reverse_lookup.len();
            assert!(len <= MAXIMUM_CAPACITY);
            len as IdIndex
        };

        let usize_index = id.index as usize;
        let lookup_id = if usize_index < lookup.len() {
            let lookup_id = &mut lookup[usize_index];
            element = match flat.replace(lookup_id.index as usize, element) {
                Ok(value) => {
                    return if id.tag == lookup_id.tag {
                        Some(value)
                    } else {
                        lookup_id.tag = id.tag;
                        None
                    }
                }
                Err(element) => element,
            };
            lookup_id
        } else {
            lookup.resize(usize_index + 1, Id::invalid());
            &mut lookup[usize_index]
        };
        *lookup_id = Id {
            index: new_index,
            ..id
        };
        reverse_lookup.push(id.index);
        flat.push(element);
        None
    }

    #[inline]
    pub fn remove(&mut self, id: Id<E>) -> Option<F::Element> {
        let IdMap {
            ref mut lookup,
            ref mut reverse_lookup,
            ref mut flat,
        } = *self;

        let usize_index = id.index as usize;
        let lookup_index = match lookup.get_mut(usize_index) {
            Some(lookup_id) => {
                if lookup_id.tag != id.tag {
                    return None;
                }
                let index = lookup_id.index;
                *lookup_id = Id::invalid();
                index
            }
            None => return None,
        };
        let usize_lookup_index = lookup_index as usize;
        let old_value = match flat.swap_remove(usize_lookup_index) {
            Some(old_value) => old_value,
            None => return None,
        };
        reverse_lookup.swap_remove(usize_lookup_index);
        reverse_lookup.get(usize_lookup_index).map(
            |&reverse_index| {
                lookup[reverse_index as usize].index = lookup_index
            },
        );
        Some(old_value)
    }

    #[inline]
    pub fn remove_by_index(&mut self, index: usize) -> Option<F::Element> {
        let IdMap {
            ref mut lookup,
            ref mut reverse_lookup,
            ref mut flat,
        } = *self;
        let old_value = match flat.swap_remove(index) {
            Some(old_value) => old_value,
            None => return None,
        };
        reverse_lookup.swap_remove(index);
        reverse_lookup.get(index).map(
            |&reverse_index| {
                lookup[reverse_index as usize].index = index as u32;
            },
        );
        Some(old_value)
    }

    #[inline]
    pub fn get<'a>(&'a self, id: Id<E>) -> Option<<&'a F as FlatGet>::ElementRef>
    where
        &'a F: FlatGet,
    {
        match self.lookup.get(id.index as usize) {
            Some(lookup_id) if lookup_id.tag == id.tag => {
                self.flat.flat_get(lookup_id.index as usize)
            }
            _ => None,
        }
    }

    #[inline]
    pub fn get_mut<'a>(&'a mut self, id: Id<E>) -> Option<<&'a mut F as FlatGetMut>::ElementRefMut>
    where
        &'a mut F: FlatGetMut,
    {
        match self.lookup.get(id.index as usize) {
            Some(lookup_id) if lookup_id.tag == id.tag => {
                self.flat.flat_get_mut(lookup_id.index as usize)
            }
            _ => None,
        }
    }

    #[inline]
    pub fn get_by_index<'a>(&'a self, index: usize) -> Option<<&'a F as FlatGet>::ElementRef>
    where
        &'a F: FlatGet,
    {
        self.flat.flat_get(index)
    }

    #[inline]
    pub fn get_mut_by_index<'a>(
        &'a mut self,
        index: usize,
    ) -> Option<<&'a mut F as FlatGetMut>::ElementRefMut>
    where
        &'a mut F: FlatGetMut,
    {
        self.flat.flat_get_mut(index)
    }

    #[inline]
    pub fn swap_indices(&mut self, index_a: usize, index_b: usize) -> bool {
        if self.flat.swap(index_a, index_b) {
            self.reverse_lookup.swap(index_a, index_b);
            self.lookup[self.reverse_lookup[index_a] as usize].index = index_a as IdIndex;
            self.lookup[self.reverse_lookup[index_b] as usize].index = index_b as IdIndex;
            true
        } else {
            false
        }
    }

    #[inline]
    pub fn id_to_index(&self, id: Id<E>) -> Option<usize> {
        match self.lookup.get(id.index as usize) {
            Some(lookup_id) if lookup_id.tag == id.tag => Some(lookup_id.index as usize),
            _ => None,
        }
    }

    #[inline]
    pub fn index_to_id(&self, index: usize) -> Option<Id<E>> {
        match self.reverse_lookup.get(index) {
            Some(&lookup_index) => {
                Some(Id {
                    index: lookup_index,
                    ..self.lookup[lookup_index as usize]
                })
            }
            _ => None,
        }
    }

    #[inline]
    pub fn access<'a>(&'a self) -> <&'a F as FlatAccess>::Access
    where
        &'a F: FlatAccess,
    {
        self.flat.flat_access()
    }

    #[inline]
    pub fn access_mut<'a>(&'a mut self) -> <&'a mut F as FlatAccessMut>::AccessMut
    where
        &'a mut F: FlatAccessMut,
    {
        self.flat.flat_access_mut()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use slab::IdSlab;

    fn setup() -> (IdSlab<()>, IdMapVec<(), u32>) {
        (IdSlab::new(), IdMapVec::new())
    }

    #[test]
    fn insert_twice_returns_first_and_respects_tag() {
        let (mut slab, mut map) = setup();
        let id1 = slab.insert(());
        assert_eq!(map.insert(id1, 1), None);
        assert_eq!(map.insert(id1, 2), Some(1));
        slab.remove(id1);
        let id2 = slab.insert(());
        assert_eq!(map.insert(id2, 3), None);
        assert_eq!(map.insert(id2, 4), Some(3));
        assert_eq!(map.insert(id1, 4), None);
    }

    #[test]
    fn insert_get_returns_correct() {
        let (mut slab, mut map) = setup();
        let id1 = slab.insert(());
        let id2 = slab.insert(());
        assert_eq!(map.get(id1), None);
        assert_eq!(map.get(id2), None);

        assert_eq!(map.insert(id1, 1), None);
        assert_eq!(map.insert(id2, 2), None);

        assert_eq!(map.get(id1), Some(&1));
        assert_eq!(map.get(id2), Some(&2));

        assert_eq!(map.get_mut(id1).map(|x| *x), Some(1));
        assert_eq!(map.get_mut(id2).map(|x| *x), Some(2));

        slab.remove(id1);
        slab.remove(id2);
        let id3 = slab.insert(());
        let id4 = slab.insert(());
        assert_eq!(map.get(id3), None);
        assert_eq!(map.get(id4), None);

        assert_eq!(map.insert(id3, 3), None);
        assert_eq!(map.insert(id4, 4), None);

        assert_eq!(map.get(id1), None);
        assert_eq!(map.get(id2), None);

        assert_eq!(map.get(id3), Some(&3));
        assert_eq!(map.get(id4), Some(&4));
        assert_eq!(map.get_mut(id3).map(|x| *x), Some(3));
        assert_eq!(map.get_mut(id4).map(|x| *x), Some(4));
    }

    #[test]
    fn modify_through_get_mut() {
        let (mut slab, mut map) = setup();
        let id1 = slab.insert(());
        assert_eq!(map.get(id1), None);
        assert_eq!(map.insert(id1, 1), None);
        *map.get_mut(id1).unwrap() = 10;
        assert_eq!(map.get(id1), Some(&10));
        assert_eq!(map.insert(id1, 20), Some(10));
    }

    #[test]
    fn insert_remove_get() {
        let (mut slab, mut map) = setup();
        let id1 = slab.insert(());
        let id2 = slab.insert(());

        assert_eq!(map.insert(id1, 1), None);
        assert_eq!(map.remove(id1), Some(1));
        assert_eq!(map.remove(id1), None);
        assert_eq!(map.get(id1), None);

        assert_eq!(map.insert(id2, 20), None);
        assert_eq!(map.insert(id1, 10), None);
        assert_eq!(map.remove(id1), Some(10));
        assert_eq!(map.get(id2), Some(&20));
        assert_eq!(map.remove(id2), Some(20));
        assert_eq!(map.get(id2), None);
    }

    #[test]
    fn swap() {
        let (mut slab, mut map) = setup();
        let id1 = slab.insert(());
        let id2 = slab.insert(());

        assert_eq!(map.insert(id1, 1), None);
        assert_eq!(map.insert(id2, 2), None);

        let i1 = map.id_to_index(id1).unwrap();
        let i2 = map.id_to_index(id2).unwrap();
        map.swap_indices(i1, i2);

        assert_eq!(map.id_to_index(id1), Some(i2));
        assert_eq!(map.id_to_index(id2), Some(i1));

        assert_eq!(map.get_by_index(i1), Some(&2));
        assert_eq!(map.get_by_index(i2), Some(&1));

        assert_eq!(map.get(id1), Some(&1));
        assert_eq!(map.get(id2), Some(&2));

        let id3 = slab.insert(());
        assert_eq!(map.insert(id3, 3), None);

        assert_eq!(map.get(id1), Some(&1));
        assert_eq!(map.get(id2), Some(&2));
        assert_eq!(map.get(id3), Some(&3));

        assert_eq!(map.remove(id1), Some(1));
        assert_eq!(map.get(id1), None);
        assert_eq!(map.get(id2), Some(&2));
        assert_eq!(map.get(id3), Some(&3));

        assert_eq!(map.remove(id2), Some(2));
        assert_eq!(map.get(id1), None);
        assert_eq!(map.get(id2), None);
        assert_eq!(map.get(id3), Some(&3));

        assert_eq!(map.remove(id3), Some(3));
        assert_eq!(map.get(id1), None);
        assert_eq!(map.get(id2), None);
        assert_eq!(map.get(id3), None);
    }
}
<commit_msg>fix map remove_by_index<commit_after>use id::{Id, IdIndex, MAXIMUM_CAPACITY};
use flat::{FlatAccess, FlatAccessMut, Flat, FlatGet, FlatGetMut};

pub type IdMapVec<E, T> = IdMap<E, Vec<T>>;

#[derive(Debug)]
pub struct IdMap<E, F: Flat> {
    lookup: Vec<Id<E>>,
    reverse_lookup: Vec<IdIndex>,
    flat: F,
}

impl<E, F: Flat> IdMap<E, F> {
    pub fn new() -> Self {
        IdMap {
            lookup: Vec::new(),
            reverse_lookup: Vec::new(),
            flat: F::new(),
        }
    }

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

    pub fn with_capacity(capacity: usize) -> Self {
        assert!(capacity <= MAXIMUM_CAPACITY);
        IdMap {
            lookup: Vec::with_capacity(capacity),
            reverse_lookup: Vec::with_capacity(capacity),
            flat: F::with_capacity(capacity),
        }
    }

    #[inline]
    pub fn insert(&mut self, id: Id<E>, mut element: F::Element) -> Option<F::Element> {
        let IdMap {
            ref mut lookup,
            ref mut reverse_lookup,
            ref mut flat,
        } = *self;
        let new_index = {
            let len = reverse_lookup.len();
            assert!(len <= MAXIMUM_CAPACITY);
            len as IdIndex
        };

        let usize_index = id.index as usize;
        let lookup_id = if usize_index < lookup.len() {
            let lookup_id = &mut lookup[usize_index];
            element = match flat.replace(lookup_id.index as usize, element) {
                Ok(value) => {
                    return if id.tag == lookup_id.tag {
                        Some(value)
                    } else {
                        lookup_id.tag = id.tag;
                        None
                    }
                }
                Err(element) => element,
            };
            lookup_id
        } else {
            lookup.resize(usize_index + 1, Id::invalid());
            &mut lookup[usize_index]
        };
        *lookup_id = Id {
            index: new_index,
            ..id
        };
        reverse_lookup.push(id.index);
        flat.push(element);
        None
    }

    #[inline]
    pub fn remove(&mut self, id: Id<E>) -> Option<F::Element> {
        let IdMap {
            ref mut lookup,
            ref mut reverse_lookup,
            ref mut flat,
        } = *self;

        let usize_index = id.index as usize;
        let lookup_index = match lookup.get_mut(usize_index) {
            Some(lookup_id) => {
                if lookup_id.tag != id.tag {
                    return None;
                }
                let index = lookup_id.index;
                *lookup_id = Id::invalid();
                index
            }
            None => return None,
        };
        let usize_lookup_index = lookup_index as usize;
        let old_value = match flat.swap_remove(usize_lookup_index) {
            Some(old_value) => old_value,
            None => return None,
        };
        reverse_lookup.swap_remove(usize_lookup_index);
        reverse_lookup.get(usize_lookup_index).map(
            |&reverse_index| {
                lookup[reverse_index as usize].index = lookup_index
            },
        );
        Some(old_value)
    }

    #[inline]
    pub fn remove_by_index(&mut self, index: usize) -> Option<F::Element> {
        let IdMap {
            ref mut lookup,
            ref mut reverse_lookup,
            ref mut flat,
        } = *self;
        let old_value = match flat.swap_remove(index) {
            Some(old_value) => old_value,
            None => return None,
        };
        lookup[reverse_lookup[index] as usize] = Id::invalid();
        reverse_lookup.swap_remove(index);
        reverse_lookup.get(index).map(|&reverse_index| {
            lookup[reverse_index as usize].index = index as u32;
        });
        Some(old_value)
    }

    #[inline]
    pub fn get<'a>(&'a self, id: Id<E>) -> Option<<&'a F as FlatGet>::ElementRef>
    where
        &'a F: FlatGet,
    {
        match self.lookup.get(id.index as usize) {
            Some(lookup_id) if lookup_id.tag == id.tag => {
                self.flat.flat_get(lookup_id.index as usize)
            }
            _ => None,
        }
    }

    #[inline]
    pub fn get_mut<'a>(&'a mut self, id: Id<E>) -> Option<<&'a mut F as FlatGetMut>::ElementRefMut>
    where
        &'a mut F: FlatGetMut,
    {
        match self.lookup.get(id.index as usize) {
            Some(lookup_id) if lookup_id.tag == id.tag => {
                self.flat.flat_get_mut(lookup_id.index as usize)
            }
            _ => None,
        }
    }

    #[inline]
    pub fn get_by_index<'a>(&'a self, index: usize) -> Option<<&'a F as FlatGet>::ElementRef>
    where
        &'a F: FlatGet,
    {
        self.flat.flat_get(index)
    }

    #[inline]
    pub fn get_mut_by_index<'a>(
        &'a mut self,
        index: usize,
    ) -> Option<<&'a mut F as FlatGetMut>::ElementRefMut>
    where
        &'a mut F: FlatGetMut,
    {
        self.flat.flat_get_mut(index)
    }

    #[inline]
    pub fn swap_indices(&mut self, index_a: usize, index_b: usize) -> bool {
        if self.flat.swap(index_a, index_b) {
            self.reverse_lookup.swap(index_a, index_b);
            self.lookup[self.reverse_lookup[index_a] as usize].index = index_a as IdIndex;
            self.lookup[self.reverse_lookup[index_b] as usize].index = index_b as IdIndex;
            true
        } else {
            false
        }
    }

    #[inline]
    pub fn id_to_index(&self, id: Id<E>) -> Option<usize> {
        match self.lookup.get(id.index as usize) {
            Some(lookup_id) if lookup_id.tag == id.tag => Some(lookup_id.index as usize),
            _ => None,
        }
    }

    #[inline]
    pub fn index_to_id(&self, index: usize) -> Option<Id<E>> {
        match self.reverse_lookup.get(index) {
            Some(&lookup_index) => {
                Some(Id {
                    index: lookup_index,
                    ..self.lookup[lookup_index as usize]
                })
            }
            _ => None,
        }
    }

    #[inline]
    pub fn access<'a>(&'a self) -> <&'a F as FlatAccess>::Access
    where
        &'a F: FlatAccess,
    {
        self.flat.flat_access()
    }

    #[inline]
    pub fn access_mut<'a>(&'a mut self) -> <&'a mut F as FlatAccessMut>::AccessMut
    where
        &'a mut F: FlatAccessMut,
    {
        self.flat.flat_access_mut()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use slab::IdSlab;

    fn setup() -> (IdSlab<()>, IdMapVec<(), u32>) {
        (IdSlab::new(), IdMapVec::new())
    }

    #[test]
    fn insert_twice_returns_first_and_respects_tag() {
        let (mut slab, mut map) = setup();
        let id1 = slab.insert(());
        assert_eq!(map.insert(id1, 1), None);
        assert_eq!(map.insert(id1, 2), Some(1));
        slab.remove(id1);
        let id2 = slab.insert(());
        assert_eq!(map.insert(id2, 3), None);
        assert_eq!(map.insert(id2, 4), Some(3));
        assert_eq!(map.insert(id1, 4), None);
    }

    #[test]
    fn insert_get_returns_correct() {
        let (mut slab, mut map) = setup();
        let id1 = slab.insert(());
        let id2 = slab.insert(());
        assert_eq!(map.get(id1), None);
        assert_eq!(map.get(id2), None);

        assert_eq!(map.insert(id1, 1), None);
        assert_eq!(map.insert(id2, 2), None);

        assert_eq!(map.get(id1), Some(&1));
        assert_eq!(map.get(id2), Some(&2));

        assert_eq!(map.get_mut(id1).map(|x| *x), Some(1));
        assert_eq!(map.get_mut(id2).map(|x| *x), Some(2));

        slab.remove(id1);
        slab.remove(id2);
        let id3 = slab.insert(());
        let id4 = slab.insert(());
        assert_eq!(map.get(id3), None);
        assert_eq!(map.get(id4), None);

        assert_eq!(map.insert(id3, 3), None);
        assert_eq!(map.insert(id4, 4), None);

        assert_eq!(map.get(id1), None);
        assert_eq!(map.get(id2), None);

        assert_eq!(map.get(id3), Some(&3));
        assert_eq!(map.get(id4), Some(&4));
        assert_eq!(map.get_mut(id3).map(|x| *x), Some(3));
        assert_eq!(map.get_mut(id4).map(|x| *x), Some(4));
    }

    #[test]
    fn modify_through_get_mut() {
        let (mut slab, mut map) = setup();
        let id1 = slab.insert(());
        assert_eq!(map.get(id1), None);
        assert_eq!(map.insert(id1, 1), None);
        *map.get_mut(id1).unwrap() = 10;
        assert_eq!(map.get(id1), Some(&10));
        assert_eq!(map.insert(id1, 20), Some(10));
    }

    #[test]
    fn insert_remove_get() {
        let (mut slab, mut map) = setup();
        let id1 = slab.insert(());
        let id2 = slab.insert(());

        assert_eq!(map.insert(id1, 1), None);
        assert_eq!(map.remove(id1), Some(1));
        assert_eq!(map.remove(id1), None);
        assert_eq!(map.get(id1), None);

        assert_eq!(map.insert(id2, 20), None);
        assert_eq!(map.insert(id1, 10), None);
        assert_eq!(map.remove(id1), Some(10));
        assert_eq!(map.get(id2), Some(&20));
        assert_eq!(map.remove(id2), Some(20));
        assert_eq!(map.get(id2), None);
    }

    #[test]
    fn swap() {
        let (mut slab, mut map) = setup();
        let id1 = slab.insert(());
        let id2 = slab.insert(());

        assert_eq!(map.insert(id1, 1), None);
        assert_eq!(map.insert(id2, 2), None);

        let i1 = map.id_to_index(id1).unwrap();
        let i2 = map.id_to_index(id2).unwrap();
        map.swap_indices(i1, i2);

        assert_eq!(map.id_to_index(id1), Some(i2));
        assert_eq!(map.id_to_index(id2), Some(i1));

        assert_eq!(map.get_by_index(i1), Some(&2));
        assert_eq!(map.get_by_index(i2), Some(&1));

        assert_eq!(map.get(id1), Some(&1));
        assert_eq!(map.get(id2), Some(&2));

        let id3 = slab.insert(());
        assert_eq!(map.insert(id3, 3), None);

        assert_eq!(map.get(id1), Some(&1));
        assert_eq!(map.get(id2), Some(&2));
        assert_eq!(map.get(id3), Some(&3));

        assert_eq!(map.remove(id1), Some(1));
        assert_eq!(map.get(id1), None);
        assert_eq!(map.get(id2), Some(&2));
        assert_eq!(map.get(id3), Some(&3));

        assert_eq!(map.remove(id2), Some(2));
        assert_eq!(map.get(id1), None);
        assert_eq!(map.get(id2), None);
        assert_eq!(map.get(id3), Some(&3));

        assert_eq!(map.remove(id3), Some(3));
        assert_eq!(map.get(id1), None);
        assert_eq!(map.get(id2), None);
        assert_eq!(map.get(id3), None);
    }
}
<|endoftext|>


<commit_before>#![allow(non_upper_case_globals)]
extern crate rustc_serialize;
use rustc_serialize::json::{self, ToJson};

extern crate rs_es;
use rs_es::Client;

extern crate iron;
use iron::prelude::*;
use iron::status;
use iron::mime::Mime;

extern crate logger;
use logger::Logger;

extern crate router;
use router::Router;

extern crate params;
use params::*;

extern crate searchspot;
use searchspot::config::*;
use searchspot::search::SearchResult;

#[macro_use] pub mod macros;
#[macro_use] extern crate lazy_static;
#[macro_use] extern crate maplit;

extern crate chrono;
mod resources;
use resources::user::Talent;

use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::io::Read;
use std::env;

lazy_static! {
  static ref config: Config = match env::args().nth(1) {
    Some(file) => Config::from_file(file),
    None       => Config::from_env()
  };
}

fn main() {
  let host = format!("{}:{}", config.http.host, config.http.port);

  println!("Searchspot v{}\n{}\n{}\n", env!("CARGO_PKG_VERSION"), config.es, config.http);

  let mut router = Router::new();
  handle_talents_search(&mut router);
  handle_talents_indexing(&mut router);
  handle_talents_reset(&mut router);

  let mut chain = Chain::new(router);

  // for some reasons, this link makes heroku crash
  if env::var("DYNO").is_err() {
    chain.link(Logger::new(None));
  }

  Iron::new(chain).http(&*host).unwrap();
}

fn handle_talents_search(mut router: &mut Router) {
  let es = Arc::new(Mutex::new(Client::new(&*config.es.host, config.es.port)));

  router.get("/talents", move |r: &mut Request| {
    search_talents(r, &mut es.lock().unwrap(), &*config.es.index)
  });
}

fn handle_talents_indexing(mut router: &mut Router) {
  let es = Arc::new(Mutex::new(Client::new(&*config.es.host, config.es.port)));

  router.post("/talents", move |r: &mut Request| {
    index_talents(r, &mut es.lock().unwrap(), &*config.es.index)
  });
}

fn handle_talents_reset(mut router: &mut Router) {
  let es = Arc::new(Mutex::new(Client::new(&*config.es.host, config.es.port)));

  router.delete("/talents", move |r: &mut Request| {
    reset_talents(r, &mut es.lock().unwrap(), &*config.es.index)
  });
}

macro_rules! try_or_422 {
  ($expr:expr) => (match $expr {
    Ok(val)  => val,
    Err(err) => {
      let content_type = "application/json".parse::<Mime>().unwrap();
      let mut error = HashMap::new();
      error.insert("error", format!("{}", err));

      return Ok(Response::with(
        (content_type, status::UnprocessableEntity, json::encode(&error).unwrap())
      ))
    }
  })
}

fn search_talents(req: &mut Request, mut es: &mut Client, index: &str) -> IronResult<Response> {
  let params   = try_or_422!(req.get_ref::<Params>());
  let response = SearchResult {
    results: Talent::search(&mut es, index, params),
    params:  params.clone()
  };

  let content_type = "application/json".parse::<Mime>().unwrap();
  Ok(Response::with(
    (content_type, status::Ok, try_or_422!(json::encode(&response.to_json())))
  ))
}

fn index_talents(req: &mut Request, mut es: &mut Client, index: &str) -> IronResult<Response> {
  let mut payload = String::new();
  req.body.read_to_string(&mut payload).unwrap();

  let talents: Vec<Talent> = try_or_422!(json::decode(&payload));
  for talent in talents {
    try_or_422!(talent.index(es, index));
  }

  Ok(Response::with(status::Ok))
}

fn reset_talents(_: &mut Request, mut es: &mut Client, index: &str) -> IronResult<Response> {
  match Talent::reset_index(&mut es, index) {
    Ok(_)  => Ok(Response::with(status::Ok)),
    Err(_) => Ok(Response::with(status::UnprocessableEntity))
  }
}
<commit_msg>main: remove shared ES client<commit_after>#![allow(non_upper_case_globals)]
extern crate rustc_serialize;
use rustc_serialize::json::{self, ToJson};

extern crate rs_es;
use rs_es::Client;

extern crate iron;
use iron::prelude::*;
use iron::status;
use iron::mime::Mime;

extern crate logger;
use logger::Logger;

extern crate router;
use router::Router;

extern crate params;
use params::*;

extern crate searchspot;
use searchspot::config::*;
use searchspot::search::SearchResult;

#[macro_use] pub mod macros;
#[macro_use] extern crate lazy_static;
#[macro_use] extern crate maplit;

extern crate chrono;
mod resources;
use resources::user::Talent;

use std::collections::HashMap;
use std::io::Read;
use std::env;

lazy_static! {
  static ref config: Config = match env::args().nth(1) {
    Some(file) => Config::from_file(file),
    None       => Config::from_env()
  };
}

fn main() {
  let host = format!("{}:{}", config.http.host, config.http.port);

  println!("Searchspot v{}\n{}\n{}\n", env!("CARGO_PKG_VERSION"), config.es, config.http);

  let mut router = Router::new();
  handle_talents_search(&mut router);
  handle_talents_indexing(&mut router);
  handle_talents_reset(&mut router);

  let mut chain = Chain::new(router);

  // for some reasons, this link makes heroku crash
  if env::var("DYNO").is_err() {
    chain.link(Logger::new(None));
  }

  Iron::new(chain).http(&*host).unwrap();
}

fn handle_talents_search(mut router: &mut Router) {
  router.get("/talents", move |r: &mut Request| {
    let mut es = Client::new(&*config.es.host, config.es.port);
    search_talents(r, &mut es, &*config.es.index)
  });
}

fn handle_talents_indexing(mut router: &mut Router) {
  router.post("/talents", move |r: &mut Request| {
    let mut es = Client::new(&*config.es.host, config.es.port);
    index_talents(r, &mut es, &*config.es.index)
  });
}

fn handle_talents_reset(mut router: &mut Router) {
  router.delete("/talents", move |r: &mut Request| {
    let mut es = Client::new(&*config.es.host, config.es.port);
    reset_talents(r, &mut es, &*config.es.index)
  });
}

macro_rules! try_or_422 {
  ($expr:expr) => (match $expr {
    Ok(val)  => val,
    Err(err) => {
      let content_type = "application/json".parse::<Mime>().unwrap();
      let mut error = HashMap::new();
      error.insert("error", format!("{}", err));

      return Ok(Response::with(
        (content_type, status::UnprocessableEntity, json::encode(&error).unwrap())
      ))
    }
  })
}

fn search_talents(req: &mut Request, mut es: &mut Client, index: &str) -> IronResult<Response> {
  let params   = try_or_422!(req.get_ref::<Params>());
  let response = SearchResult {
    results: Talent::search(&mut es, index, params),
    params:  params.clone()
  };

  let content_type = "application/json".parse::<Mime>().unwrap();
  Ok(Response::with(
    (content_type, status::Ok, try_or_422!(json::encode(&response.to_json())))
  ))
}

fn index_talents(req: &mut Request, mut es: &mut Client, index: &str) -> IronResult<Response> {
  let mut payload = String::new();
  req.body.read_to_string(&mut payload).unwrap();

  let talents: Vec<Talent> = try_or_422!(json::decode(&payload));
  for talent in talents {
    try_or_422!(talent.index(es, index));
  }

  Ok(Response::with(status::Ok))
}

fn reset_talents(_: &mut Request, mut es: &mut Client, index: &str) -> IronResult<Response> {
  match Talent::reset_index(&mut es, index) {
    Ok(_)  => Ok(Response::with(status::Ok)),
    Err(_) => Ok(Response::with(status::UnprocessableEntity))
  }
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::io::{BufReader, BufRead, Write};
use std::collections::{HashMap};

use std::env;
use std::io;

const WHITESPACE_FACTOR: isize = 5;
const WHITESPACE_REDUCE: isize = 2;
const CLASS_FACTOR: isize = 5;
const CLASS_REDUCE: isize = 2;

#[derive(Debug)]
struct LineInfo {
    line: String,
    char_map: HashMap<char, Vec<usize>>,
    heatmap: Vec<isize>
}

#[derive(PartialEq)]
enum CharClass {
    Whitespace,
    Numeric,
    Alphabetic,
    Other
}

impl<T: Into<String>> From<T> for LineInfo {
    fn from(item: T) -> LineInfo {
        let mut map: HashMap<char, Vec<usize>> = HashMap::new();
        let mut heat = vec![];
        let line = item.into();

        let mut ws_score = 0;
        let mut cs_score = 0;
        let mut cur_class = CharClass::Whitespace;

        for (idx, c) in line.chars().enumerate() {
            // don't map whitespace
            if !c.is_whitespace() {
                // update the character class change score if needed
                if c.is_numeric() {
                    if cur_class != CharClass::Numeric {
                        cur_class = CharClass::Numeric;
                        cs_score += CLASS_FACTOR;
                    }
                } else if c.is_alphabetic() {
                    if cur_class != CharClass::Alphabetic {
                        cur_class = CharClass::Alphabetic;
                        cs_score += CLASS_FACTOR;
                    }
                } else {
                    if cur_class != CharClass::Other {
                        cur_class = CharClass::Other;
                        cs_score += CLASS_FACTOR;
                    }
                }

                // add an entry in the character map
                map.entry(c).or_insert(Vec::new()).push(idx);
                if c.is_uppercase() {
                    for lc in c.to_lowercase() {
                        // also insert all lowercase equivalents of this character
                        // but not the other way around, so that typing something
                        // uppercase specifies to match uppercase
                        map.entry(lc).or_insert(Vec::new()).push(idx);
                    }
                }
            } else {
                // whitespace is treated differently
                cur_class = CharClass::Whitespace;
                ws_score = WHITESPACE_FACTOR;
            }

            // push to the heatmap
            heat.push(ws_score + cs_score);

            // reduce things
            ws_score /= WHITESPACE_REDUCE;
            cs_score /= CLASS_REDUCE;
        }

        LineInfo {
            line: line,
            char_map: map,
            heatmap: heat
        }
    }
}

fn main() {
    let history_path = match env::var("HISTFILE") {
        Ok(p) => p,
        Err(e) => panic!("Failed to get bash history file: {}", e)
    };

    let input_file = match File::open(&history_path) {
        Ok(f) => BufReader::new(f),
        Err(e) => panic!("Could not open history file: {}", e)
    };

    let mut query = String::new();

    print!("Match input: ");
    match io::stdout().flush() {
        Err(e) => panic!("Failed to flush stdout: {}", e),
        Ok(_) => {}
    }

    match io::stdin().read_line(&mut query) {
        Ok(_) => {},
        Err(e) => panic!("Failed to read input line: {}", e)
    }

    match query.pop() {
        Some('\n') => {/* pop off trailing newline */},
        Some(c) => query.push(c),
        None => {/* Do nothing with an empty query */}
    }

    let mut positions = vec![];

    for m_line in input_file.lines() {
        let line = match m_line {
            Ok(l) => l,
            Err(e) => panic!("Failed to read line: {}", e)
        };

        let info = LineInfo::from(line.clone());
        positions.clear();

        for c in query.chars() {
            match info.char_map.get(&c) {
                None => break,
                Some(list) => {
                    positions.push(list.clone());
                }
            }
        }

        if positions.len() == query.len() {
            // matching line
            println!("Line {:?}: positions {:?}", line, positions);
        }
    }
}
<commit_msg>Create permutation function<commit_after>use std::fs::File;
use std::io::{BufReader, BufRead, Write};
use std::collections::{HashMap};

use std::env;
use std::io;

const WHITESPACE_FACTOR: isize = 5;
const WHITESPACE_REDUCE: isize = 2;
const CLASS_FACTOR: isize = 5;
const CLASS_REDUCE: isize = 2;

#[derive(Debug)]
struct LineInfo {
    line: String,
    char_map: HashMap<char, Vec<usize>>,
    heatmap: Vec<isize>
}

#[derive(PartialEq)]
enum CharClass {
    Whitespace,
    Numeric,
    Alphabetic,
    Other
}

impl<T: Into<String>> From<T> for LineInfo {
    fn from(item: T) -> LineInfo {
        let mut map: HashMap<char, Vec<usize>> = HashMap::new();
        let mut heat = vec![];
        let line = item.into();

        let mut ws_score = 0;
        let mut cs_score = 0;
        let mut cur_class = CharClass::Whitespace;

        for (idx, c) in line.chars().enumerate() {
            // don't map whitespace
            if !c.is_whitespace() {
                // update the character class change score if needed
                if c.is_numeric() {
                    if cur_class != CharClass::Numeric {
                        cur_class = CharClass::Numeric;
                        cs_score += CLASS_FACTOR;
                    }
                } else if c.is_alphabetic() {
                    if cur_class != CharClass::Alphabetic {
                        cur_class = CharClass::Alphabetic;
                        cs_score += CLASS_FACTOR;
                    }
                } else {
                    if cur_class != CharClass::Other {
                        cur_class = CharClass::Other;
                        cs_score += CLASS_FACTOR;
                    }
                }

                // add an entry in the character map
                map.entry(c).or_insert(Vec::new()).push(idx);
                if c.is_uppercase() {
                    for lc in c.to_lowercase() {
                        // also insert all lowercase equivalents of this character
                        // but not the other way around, so that typing something
                        // uppercase specifies to match uppercase
                        map.entry(lc).or_insert(Vec::new()).push(idx);
                    }
                }
            } else {
                // whitespace is treated differently
                cur_class = CharClass::Whitespace;
                ws_score = WHITESPACE_FACTOR;
            }

            // push to the heatmap
            heat.push(ws_score + cs_score);

            // reduce things
            ws_score /= WHITESPACE_REDUCE;
            cs_score /= CLASS_REDUCE;
        }

        LineInfo {
            line: line,
            char_map: map,
            heatmap: heat
        }
    }
}

impl LineInfo {
    fn query_sequence<T: AsRef<str>>(&self, query_item: T) -> Option<Vec<Vec<usize>>> {
        let query = query_item.as_ref();
        let mut positions: Vec<Vec<usize>> = vec![];

        for c in query.chars() {
            match self.char_map.get(&c) {
                None => break,
                Some(list) => {
                    let to_push;
                    match positions.last() {
                        None => {
                            to_push = list.clone();
                        },
                        Some(item) => {
                            match list.binary_search(&item[0]) {
                                Ok(idx) => {
                                    if idx >= list.len() - 1 {
                                        // line is non-matching
                                        break;
                                    } else {
                                        to_push = list.split_at(idx + 1).1.into();
                                    }
                                },
                                Err(idx) => {
                                    if idx >= list.len() {
                                        // line is non-matching
                                        break;
                                    } else {
                                        to_push = list.split_at(idx).1.into();
                                    }
                                }
                            }
                        }
                    }
                    positions.push(to_push);
                }
            }
        }

        if positions.len() == query.len() {
            Some(positions)
        } else {
            None
        }
    }

    fn query_positions<T: AsRef<str>>(&self, query: T) -> Option<Vec<Vec<usize>>> {
        match self.query_sequence(query) {
            None => None,
            Some(positions) => {
                // matching line
                // create our idx vector
                let mut idx = vec![0; positions.len()];
                let mut result = vec![];
                loop {
                    // check that current configuration is strictly increasing
                    let mut ignore = false;
                    {
                        let mut last_pos = None;
                        for (i, pos) in idx.iter().enumerate() {
                            match last_pos {
                                None => last_pos = Some(positions[i][*pos]),
                                Some(other) => {
                                    if other >= positions[i][*pos] {
                                        ignore = true;
                                        break;
                                    } else {
                                        last_pos = Some(positions[i][*pos]);
                                    }
                                }
                            }
                        }
                    }

                    if !ignore {
                        // add the configuration to the list
                        result.push(idx.iter().enumerate().map(|(i, pos)| {positions[i][*pos]}).collect());
                    }

                    // update our position vector
                    let mut update_idx = idx.len() - 1;
                    let mut finished = false;
                    loop {
                        idx[update_idx] += 1;
                        if idx[update_idx] >= positions[update_idx].len() {
                            if update_idx == 0 {
                                // we're finished with all permutations
                                finished = true;
                                break;
                            } else {
                                idx[update_idx] = 0;
                                update_idx -= 1;
                            }
                        } else {
                            // finished updating for this permutation
                            break;
                        }
                    }
                    if finished {
                        // finished with everything
                        break;
                    }
                }

                // return result
                Some(result)
            }
        }
    }
}

fn main() {
    let history_path = match env::var("HISTFILE") {
        Ok(p) => p,
        Err(e) => panic!("Failed to get bash history file: {}", e)
    };

    let input_file = match File::open(&history_path) {
        Ok(f) => BufReader::new(f),
        Err(e) => panic!("Could not open history file: {}", e)
    };

    let mut query = String::new();

    print!("Match input: ");
    match io::stdout().flush() {
        Err(e) => panic!("Failed to flush stdout: {}", e),
        Ok(_) => {}
    }

    match io::stdin().read_line(&mut query) {
        Ok(_) => {},
        Err(e) => panic!("Failed to read input line: {}", e)
    }

    match query.pop() {
        Some('\n') => {/* pop off trailing newline */},
        Some(c) => query.push(c),
        None => {/* Do nothing with an empty query */}
    }

    for m_line in input_file.lines() {
        let line = match m_line {
            Ok(l) => l,
            Err(e) => panic!("Failed to read line: {}", e)
        };

        let info = LineInfo::from(line.clone());

        match info.query_positions(&query) {
            None => {
                // non-matching line
            },
            Some(positions) => {
                println!("Matching line {:?} with positions {:?}", line, positions);
            }
        }
    }
}
<|endoftext|>


<commit_before>// ---------------------------------------------------------------------------------------
// Ruthenium, an ack-like searcher, (c) 2015 Georg Brandl.
// Licensed under the MIT license.
// ---------------------------------------------------------------------------------------

#[macro_use]
extern crate clap;
extern crate libc;
extern crate regex_dfa;
extern crate walkdir;
extern crate memmap;
extern crate scoped_threadpool;
extern crate num_cpus;
extern crate glob;

mod search;
mod ignore;
mod display;
mod options;

use std::cmp::max;
use std::sync::mpsc::{channel, Sender};
use std::thread;
use memmap::{Mmap, Protection};
use scoped_threadpool::Pool;
use walkdir::WalkDirIterator;

use display::DisplayMode;
use search::FileResult;
use options::Opts;

fn walk(chan: Sender<FileResult>, opts: &Opts) {
    let mut pool = Pool::new(max(opts.workers - 1, 1));
    let regex = search::create_rx(&opts);

    let walker = walkdir::WalkDir::new(&opts.path)
        .follow_links(opts.follow_links)
        .max_depth(opts.depth);
    pool.scoped(|scope| {
        let rx = &regex;
        let mut parent_stack: Vec<::std::path::PathBuf> = Vec::new();
        // stack of Ignore structs per directory, they accumulate
        // XXX: add global ignores from cmdline and a config file here
        let mut ignore_stack = Vec::new();
        let walker = walker.into_iter().filter_entry(|entry| {
            {
                let new_parent = entry.path().parent().unwrap();
                //println!("new_parent: {:?}", new_parent);
                // remove parents that are not applicable anymore
                while !parent_stack.is_empty() &&
                    parent_stack.last().unwrap().as_path() != new_parent
                {
                    //println!("popping");
                    ignore_stack.pop();
                    parent_stack.pop();
                }
            }
            // weed out hidden files
            let path = entry.path();
            if let Some(fname) = path.file_name() {
                if !opts.do_hidden && fname.to_string_lossy().starts_with(".") {
                    return false;
                }
            }
            if opts.check_ignores && ignore::match_patterns(path, &ignore_stack) {
                return false;
            }
            // we got a new dir?
            if entry.file_type().is_dir() {
                let new_path = entry.path().to_path_buf();
                // read ignore patterns specific to this directory
                ignore_stack.push(ignore::read_patterns(&new_path));
                parent_stack.push(new_path);
            }
            //println!("{:?} {:?}", entry.path(), parent_stack);
            true
        });
        for entry in walker {
            if let Ok(entry) = entry {
                // weed out further special files
                if !entry.file_type().is_file() {
                    continue;
                }
                // open and search file in one of the worker threads
                let ch = chan.clone();
                scope.execute(move || {
                    let path = entry.path();
                    if let Ok(map) = Mmap::open_path(path, Protection::Read) {
                        let buf = unsafe { map.as_slice() };
                        let res = search::search(rx, &opts, path, buf);
                        ch.send(res).unwrap();
                    }
                });
            }
        }
    });
}

fn run<D: DisplayMode>(display: &mut D, opts: Opts) {
    let (w_chan, r_chan) = channel();
    thread::spawn(move || {
        walk(w_chan, &opts);
    });
    while let Ok(r) = r_chan.recv() {
        display.print_result(r);
    }
}

fn main() {
    let mut opts = Opts::from_cmdline();
    let colors = opts.colors.take().unwrap();
    if opts.only_count {
        run(&mut display::CountMode::new(colors), opts);
    } else if opts.only_files == Some(true) {
        run(&mut display::FilesOnlyMode::new(colors, true), opts);
    } else if opts.only_files == Some(false) {
        run(&mut display::FilesOnlyMode::new(colors, false), opts);
    } else if opts.ackmate_format {
        run(&mut display::AckMateMode::new(), opts);
    } else if opts.vimgrep_format {
        run(&mut display::VimGrepMode, opts);
    } else if !opts.show_heading {
        run(&mut display::OneLineMode::new(colors, opts.show_break), opts);
    } else {
        run(&mut display::DefaultMode::new(colors, opts.show_break), opts);
    }
}
<commit_msg>Remove unnecessary scope.<commit_after>// ---------------------------------------------------------------------------------------
// Ruthenium, an ack-like searcher, (c) 2015 Georg Brandl.
// Licensed under the MIT license.
// ---------------------------------------------------------------------------------------

#[macro_use]
extern crate clap;
extern crate libc;
extern crate regex_dfa;
extern crate walkdir;
extern crate memmap;
extern crate scoped_threadpool;
extern crate num_cpus;
extern crate glob;

mod search;
mod ignore;
mod display;
mod options;

use std::cmp::max;
use std::sync::mpsc::{channel, Sender};
use std::thread;
use memmap::{Mmap, Protection};
use scoped_threadpool::Pool;
use walkdir::WalkDirIterator;

use display::DisplayMode;
use search::FileResult;
use options::Opts;

fn walk(chan: Sender<FileResult>, opts: &Opts) {
    let mut pool = Pool::new(max(opts.workers - 1, 1));
    let regex = search::create_rx(&opts);

    let walker = walkdir::WalkDir::new(&opts.path)
        .follow_links(opts.follow_links)
        .max_depth(opts.depth);
    pool.scoped(|scope| {
        let rx = &regex;
        let mut parent_stack: Vec<::std::path::PathBuf> = Vec::new();
        // stack of Ignore structs per directory, they accumulate
        // XXX: add global ignores from cmdline and a config file here
        let mut ignore_stack = Vec::new();
        let walker = walker.into_iter().filter_entry(|entry| {
            let new_parent = entry.path().parent().unwrap();
            // remove parents that are not applicable anymore
            while !parent_stack.is_empty() &&
                parent_stack.last().unwrap().as_path() != new_parent
            {
                //println!("popping");
                ignore_stack.pop();
                parent_stack.pop();
            }
            // weed out hidden files
            let path = entry.path();
            if let Some(fname) = path.file_name() {
                if !opts.do_hidden && fname.to_string_lossy().starts_with(".") {
                    return false;
                }
            }
            if opts.check_ignores && ignore::match_patterns(path, &ignore_stack) {
                return false;
            }
            // we got a new dir?
            if entry.file_type().is_dir() {
                let new_path = entry.path().to_path_buf();
                // read ignore patterns specific to this directory
                ignore_stack.push(ignore::read_patterns(&new_path));
                parent_stack.push(new_path);
            }
            //println!("{:?} {:?}", entry.path(), parent_stack);
            true
        });
        for entry in walker {
            if let Ok(entry) = entry {
                // weed out further special files
                if !entry.file_type().is_file() {
                    continue;
                }
                // open and search file in one of the worker threads
                let ch = chan.clone();
                scope.execute(move || {
                    let path = entry.path();
                    if let Ok(map) = Mmap::open_path(path, Protection::Read) {
                        let buf = unsafe { map.as_slice() };
                        let res = search::search(rx, &opts, path, buf);
                        ch.send(res).unwrap();
                    }
                });
            }
        }
    });
}

fn run<D: DisplayMode>(display: &mut D, opts: Opts) {
    let (w_chan, r_chan) = channel();
    thread::spawn(move || {
        walk(w_chan, &opts);
    });
    while let Ok(r) = r_chan.recv() {
        display.print_result(r);
    }
}

fn main() {
    let mut opts = Opts::from_cmdline();
    let colors = opts.colors.take().unwrap();
    if opts.only_count {
        run(&mut display::CountMode::new(colors), opts);
    } else if opts.only_files == Some(true) {
        run(&mut display::FilesOnlyMode::new(colors, true), opts);
    } else if opts.only_files == Some(false) {
        run(&mut display::FilesOnlyMode::new(colors, false), opts);
    } else if opts.ackmate_format {
        run(&mut display::AckMateMode::new(), opts);
    } else if opts.vimgrep_format {
        run(&mut display::VimGrepMode, opts);
    } else if !opts.show_heading {
        run(&mut display::OneLineMode::new(colors, opts.show_break), opts);
    } else {
        run(&mut display::DefaultMode::new(colors, opts.show_break), opts);
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::{error, fmt};
use std::io::{self, Write};
use std::process::{Command, ExitStatus};

extern crate itertools;
use itertools::Itertools;

// TODO(tsion): Use the readline library.
fn prompt(line: &mut String) -> io::Result<usize> {
    let current_dir = try!(std::env::current_dir());
    print!("{}> ", current_dir.display());
    try!(io::stdout().flush());
    io::stdin().read_line(line)
}

#[derive(Clone, Debug, Eq, PartialEq)]
enum Expr {
    Text(String),
}

#[derive(Clone, Debug, Eq, PartialEq)]
enum Ast {
    Empty,
    Call { command: String, args: Vec<Vec<Expr>> }
}

#[derive(Clone, Debug, Eq, PartialEq)]
enum Token {
    Newline,
    Whitespace,
    Text(String),
}

#[derive(Clone, Debug, Eq, PartialEq)]
struct Lexer<'src> {
    source: &'src str,
    position: usize,
}

#[derive(Clone, Debug, Eq, PartialEq)]
enum ParseError {
    UnclosedDelimiter,
    UnexpectedChar,
    UnexpectedEnd,
}

type ParseResult<T> = Result<T, ParseError>;

impl error::Error for ParseError {
    fn description(&self) -> &str {
        match *self {
            ParseError::UnclosedDelimiter => "unclosed delimiter",
            ParseError::UnexpectedChar    => "unexpected character",
            ParseError::UnexpectedEnd     => "unexpected end of input",
        }
    }
}

impl fmt::Display for ParseError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", error::Error::description(self))
    }
}

impl<'src> Lexer<'src> {
    fn new(source: &str) -> Lexer {
        Lexer {
            source: source,
            position: 0,
        }
    }

    fn read_char(&mut self) -> Option<char> {
        let opt_c = self.source[self.position..].chars().next();

        if let Some(c) = opt_c {
            self.position += c.len_utf8();
        }

        opt_c
    }

    /// Step backwards one `char` in the input. Must not be called more times than `read_char` has
    /// been called.
    fn unread_char(&mut self) {
        assert!(self.position != 0);
        let (prev_pos, _) = self.source[..self.position].char_indices().next_back().unwrap();
        self.position = prev_pos;
    }

    fn is_whitespace(c: char) -> bool {
        c == ' ' || c == '\t'
    }

    fn is_unquoted_text(c: char) -> bool {
        match c {
            'a'...'z' | 'A'...'Z' | '0'...'9' | '-' | '+' | '/' | '_' | '.' | ',' => true,
            _ => false,
        }
    }

    fn skip_while<F>(&mut self, mut predicate: F) where F: FnMut(char) -> bool {
        while let Some(c) = self.read_char() {
            if !predicate(c) {
                self.unread_char();
                break;
            }
        }
    }

    fn lex_whitespace(&mut self) -> ParseResult<Token> {
        self.skip_while(Lexer::is_whitespace);
        Ok(Token::Whitespace)
    }

    fn lex_unquoted_text(&mut self) -> ParseResult<Token> {
        let start = self.position;
        self.skip_while(Lexer::is_unquoted_text);
        let end = self.position;

        let text = String::from(&self.source[start..end]);
        Ok(Token::Text(text))
    }

    fn lex_double_quoted_text(&mut self) -> ParseResult<Token> {
        let mut text = String::new();

        while let Some(c) = self.read_char() {
            match c {
                '"'  => return Ok(Token::Text(text)),
                '\\' => try!(self.lex_double_quote_escape(&mut text)),
                c => text.push(c),
            };
        }

        Err(ParseError::UnclosedDelimiter)
    }

    fn lex_double_quote_escape(&mut self, text: &mut String) -> ParseResult<()> {
        let escaped = try!(self.read_char().ok_or(ParseError::UnexpectedEnd));

        match escaped {
            '\\' | '"' => text.push(escaped),
            c => {
                text.push('\\');
                text.push(c);
            }
        }

        Ok(())
    }
}

impl<'src> Iterator for Lexer<'src> {
    type Item = ParseResult<Token>;

    fn next(&mut self) -> Option<ParseResult<Token>> {
        self.read_char().map(|c| {
            match c {
                c if Lexer::is_whitespace(c)    => self.lex_whitespace(),
                c if Lexer::is_unquoted_text(c) => {
                    self.unread_char();
                    self.lex_unquoted_text()
                },
                '\r' | '\n'                     => Ok(Token::Newline),
                '"'                             => self.lex_double_quoted_text(),
                _                               => Err(ParseError::UnexpectedChar),
            }
        })
    }
}

#[derive(Clone)]
struct Parser<'src> {
    lexer: Lexer<'src>,
}

impl<'src> Parser<'src> {
    fn new(input: &str) -> Parser {
        Parser { lexer: Lexer::new(input) }
    }

    fn parse(&mut self) -> ParseResult<Ast> {
        if let Some(token_result) = self.lexer.next() {
            let token = try!(token_result);
            match token {
                Token::Whitespace | Token::Newline => self.parse(),
                Token::Text(command) => self.parse_call(command),
            }
        } else {
            Ok(Ast::Empty)
        }
    }

    fn parse_call(&mut self, command: String) -> ParseResult<Ast> {
        let mut args = vec![];
        let mut current_arg = vec![];

        for token_result in &mut self.lexer {
            let token = try!(token_result);
            match token {
                Token::Newline => {
                    if !current_arg.is_empty() {
                        args.push(current_arg);
                    }
                    break;
                },

                Token::Whitespace => {
                    if !current_arg.is_empty() {
                        args.push(current_arg);
                        current_arg = vec![];
                    }
                },

                Token::Text(text) => {
                    current_arg.push(Expr::Text(text));
                },
            }
        }

        Ok(Ast::Call { command: command, args: args })
    }
}

struct Builtin {
    name: &'static str,
    min_args: usize,
    max_args: usize,
    func: fn(&[String]) -> i32,
}

fn result_to_exit_code(cmd: &'static str, result: io::Result<()>) -> i32 {
    match result {
        Ok(()) => 0,
        Err(e) => {
            writeln!(&mut io::stderr(), "shroom: {}: {}", cmd, e).unwrap();
            1
        },
    }
}

fn builtin_cd(args: &[String]) -> i32 {
    if let Some(path) = args.get(0) {
        result_to_exit_code("cd", std::env::set_current_dir(path))
    } else if let Some(home) = std::env::home_dir() {
        result_to_exit_code("cd", std::env::set_current_dir(home))
    } else {
        writeln!(&mut io::stderr(), "shroom: cd: couldn't find home dir").unwrap();
        1
    }
}

fn builtin_exit(args: &[String]) -> i32 {
    if let Some(exit_code_str) = args.get(0) {
        match exit_code_str.parse() {
            Ok(exit_code) => std::process::exit(exit_code),
            Err(e) => {
                writeln!(&mut io::stderr(), "shroom: exit: can't parse exit code: {}", e).unwrap();
                1
            },
        }
    } else {
        std::process::exit(0);
    }
}

fn execute(ast: &Ast) -> i32 {
    let mut builtins = HashMap::new();

    builtins.insert("cd", Builtin {
        name: "cd",
        min_args: 0,
        max_args: 1,
        func: builtin_cd,
    });

    builtins.insert("exit", Builtin {
        name: "exit",
        min_args: 0,
        max_args: 1,
        func: builtin_exit,
    });

    match *ast {
        Ast::Empty => 0,

        Ast::Call { ref command, ref args } => {
            // Evaluate argument expressions.
            let evaluated_args: Vec<String> = args.iter().map(|arg| {
                arg.iter().map(|expr| {
                    match *expr {
                        Expr::Text(ref text) => text,
                    }
                }).join("")
            }).collect();

            if let Some(builtin) = builtins.get(&command[..]) {
                if args.len() < builtin.min_args {
                    writeln!(&mut io::stderr(), "shroom: {}: not enough arguments",
                             builtin.name).unwrap();
                    1
                } else if args.len() > builtin.max_args {
                    writeln!(&mut io::stderr(), "shroom: {}: too many arguments",
                             builtin.name).unwrap();
                    1
                } else {
                    (builtin.func)(&evaluated_args)
                }
            } else {
                match Command::new(command).args(&evaluated_args).status() {
                    Ok(exit_status) => {
                        #[cfg(unix)]
                        fn exit_signal(exit_status: &ExitStatus) -> Option<i32> {
                            use std::os::unix::process::ExitStatusExt;
                            exit_status.signal()
                        }

                        #[cfg(not(unix))]
                        fn exit_signal(_exit_status: &ExitStatus) -> Option<i32> {
                            None
                        }

                        if let Some(code) = exit_status.code() {
                            code
                        } else if let Some(signal) = exit_signal(&exit_status) {
                            128 + signal
                        } else {
                            127
                        }
                    },

                    Err(e) => {
                        writeln!(&mut io::stderr(), "shroom: {}: {}", command, e).unwrap();
                        127
                    },
                }
            }
        },
    }
}

fn main() {
    let mut line = String::new();
    loop {
        prompt(&mut line).unwrap();

        match Parser::new(&line).parse() {
            Ok(ast) => {
                let exit_code = execute(&ast);
                if exit_code != 0 {
                    println!("shroom: exit code: {}", exit_code);
                }
            },

            Err(parse_error) => {
                println!("shroom: parse error: {}", parse_error);
            },
        }

        line.clear();
    }
}
<commit_msg>Conservatively reserve comma for potential syntax use.<commit_after>use std::collections::HashMap;
use std::{error, fmt};
use std::io::{self, Write};
use std::process::{Command, ExitStatus};

extern crate itertools;
use itertools::Itertools;

// TODO(tsion): Use the readline library.
fn prompt(line: &mut String) -> io::Result<usize> {
    let current_dir = try!(std::env::current_dir());
    print!("{}> ", current_dir.display());
    try!(io::stdout().flush());
    io::stdin().read_line(line)
}

#[derive(Clone, Debug, Eq, PartialEq)]
enum Expr {
    Text(String),
}

#[derive(Clone, Debug, Eq, PartialEq)]
enum Ast {
    Empty,
    Call { command: String, args: Vec<Vec<Expr>> }
}

#[derive(Clone, Debug, Eq, PartialEq)]
enum Token {
    Newline,
    Whitespace,
    Text(String),
}

#[derive(Clone, Debug, Eq, PartialEq)]
struct Lexer<'src> {
    source: &'src str,
    position: usize,
}

#[derive(Clone, Debug, Eq, PartialEq)]
enum ParseError {
    UnclosedDelimiter,
    UnexpectedChar,
    UnexpectedEnd,
}

type ParseResult<T> = Result<T, ParseError>;

impl error::Error for ParseError {
    fn description(&self) -> &str {
        match *self {
            ParseError::UnclosedDelimiter => "unclosed delimiter",
            ParseError::UnexpectedChar    => "unexpected character",
            ParseError::UnexpectedEnd     => "unexpected end of input",
        }
    }
}

impl fmt::Display for ParseError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", error::Error::description(self))
    }
}

impl<'src> Lexer<'src> {
    fn new(source: &str) -> Lexer {
        Lexer {
            source: source,
            position: 0,
        }
    }

    fn read_char(&mut self) -> Option<char> {
        let opt_c = self.source[self.position..].chars().next();

        if let Some(c) = opt_c {
            self.position += c.len_utf8();
        }

        opt_c
    }

    /// Step backwards one `char` in the input. Must not be called more times than `read_char` has
    /// been called.
    fn unread_char(&mut self) {
        assert!(self.position != 0);
        let (prev_pos, _) = self.source[..self.position].char_indices().next_back().unwrap();
        self.position = prev_pos;
    }

    fn is_whitespace(c: char) -> bool {
        c == ' ' || c == '\t'
    }

    fn is_unquoted_text(c: char) -> bool {
        match c {
            'a'...'z' | 'A'...'Z' | '0'...'9' | '-' | '+' | '/' | '_' | '.' => true,
            _ => false,
        }
    }

    fn skip_while<F>(&mut self, mut predicate: F) where F: FnMut(char) -> bool {
        while let Some(c) = self.read_char() {
            if !predicate(c) {
                self.unread_char();
                break;
            }
        }
    }

    fn lex_whitespace(&mut self) -> ParseResult<Token> {
        self.skip_while(Lexer::is_whitespace);
        Ok(Token::Whitespace)
    }

    fn lex_unquoted_text(&mut self) -> ParseResult<Token> {
        let start = self.position;
        self.skip_while(Lexer::is_unquoted_text);
        let end = self.position;

        let text = String::from(&self.source[start..end]);
        Ok(Token::Text(text))
    }

    fn lex_double_quoted_text(&mut self) -> ParseResult<Token> {
        let mut text = String::new();

        while let Some(c) = self.read_char() {
            match c {
                '"'  => return Ok(Token::Text(text)),
                '\\' => try!(self.lex_double_quote_escape(&mut text)),
                c => text.push(c),
            };
        }

        Err(ParseError::UnclosedDelimiter)
    }

    fn lex_double_quote_escape(&mut self, text: &mut String) -> ParseResult<()> {
        let escaped = try!(self.read_char().ok_or(ParseError::UnexpectedEnd));

        match escaped {
            '\\' | '"' => text.push(escaped),
            c => {
                text.push('\\');
                text.push(c);
            }
        }

        Ok(())
    }
}

impl<'src> Iterator for Lexer<'src> {
    type Item = ParseResult<Token>;

    fn next(&mut self) -> Option<ParseResult<Token>> {
        self.read_char().map(|c| {
            match c {
                c if Lexer::is_whitespace(c)    => self.lex_whitespace(),
                c if Lexer::is_unquoted_text(c) => {
                    self.unread_char();
                    self.lex_unquoted_text()
                },
                '\r' | '\n'                     => Ok(Token::Newline),
                '"'                             => self.lex_double_quoted_text(),
                _                               => Err(ParseError::UnexpectedChar),
            }
        })
    }
}

#[derive(Clone)]
struct Parser<'src> {
    lexer: Lexer<'src>,
}

impl<'src> Parser<'src> {
    fn new(input: &str) -> Parser {
        Parser { lexer: Lexer::new(input) }
    }

    fn parse(&mut self) -> ParseResult<Ast> {
        if let Some(token_result) = self.lexer.next() {
            let token = try!(token_result);
            match token {
                Token::Whitespace | Token::Newline => self.parse(),
                Token::Text(command) => self.parse_call(command),
            }
        } else {
            Ok(Ast::Empty)
        }
    }

    fn parse_call(&mut self, command: String) -> ParseResult<Ast> {
        let mut args = vec![];
        let mut current_arg = vec![];

        for token_result in &mut self.lexer {
            let token = try!(token_result);
            match token {
                Token::Newline => {
                    if !current_arg.is_empty() {
                        args.push(current_arg);
                    }
                    break;
                },

                Token::Whitespace => {
                    if !current_arg.is_empty() {
                        args.push(current_arg);
                        current_arg = vec![];
                    }
                },

                Token::Text(text) => {
                    current_arg.push(Expr::Text(text));
                },
            }
        }

        Ok(Ast::Call { command: command, args: args })
    }
}

struct Builtin {
    name: &'static str,
    min_args: usize,
    max_args: usize,
    func: fn(&[String]) -> i32,
}

fn result_to_exit_code(cmd: &'static str, result: io::Result<()>) -> i32 {
    match result {
        Ok(()) => 0,
        Err(e) => {
            writeln!(&mut io::stderr(), "shroom: {}: {}", cmd, e).unwrap();
            1
        },
    }
}

fn builtin_cd(args: &[String]) -> i32 {
    if let Some(path) = args.get(0) {
        result_to_exit_code("cd", std::env::set_current_dir(path))
    } else if let Some(home) = std::env::home_dir() {
        result_to_exit_code("cd", std::env::set_current_dir(home))
    } else {
        writeln!(&mut io::stderr(), "shroom: cd: couldn't find home dir").unwrap();
        1
    }
}

fn builtin_exit(args: &[String]) -> i32 {
    if let Some(exit_code_str) = args.get(0) {
        match exit_code_str.parse() {
            Ok(exit_code) => std::process::exit(exit_code),
            Err(e) => {
                writeln!(&mut io::stderr(), "shroom: exit: can't parse exit code: {}", e).unwrap();
                1
            },
        }
    } else {
        std::process::exit(0);
    }
}

fn execute(ast: &Ast) -> i32 {
    let mut builtins = HashMap::new();

    builtins.insert("cd", Builtin {
        name: "cd",
        min_args: 0,
        max_args: 1,
        func: builtin_cd,
    });

    builtins.insert("exit", Builtin {
        name: "exit",
        min_args: 0,
        max_args: 1,
        func: builtin_exit,
    });

    match *ast {
        Ast::Empty => 0,

        Ast::Call { ref command, ref args } => {
            // Evaluate argument expressions.
            let evaluated_args: Vec<String> = args.iter().map(|arg| {
                arg.iter().map(|expr| {
                    match *expr {
                        Expr::Text(ref text) => text,
                    }
                }).join("")
            }).collect();

            if let Some(builtin) = builtins.get(&command[..]) {
                if args.len() < builtin.min_args {
                    writeln!(&mut io::stderr(), "shroom: {}: not enough arguments",
                             builtin.name).unwrap();
                    1
                } else if args.len() > builtin.max_args {
                    writeln!(&mut io::stderr(), "shroom: {}: too many arguments",
                             builtin.name).unwrap();
                    1
                } else {
                    (builtin.func)(&evaluated_args)
                }
            } else {
                match Command::new(command).args(&evaluated_args).status() {
                    Ok(exit_status) => {
                        #[cfg(unix)]
                        fn exit_signal(exit_status: &ExitStatus) -> Option<i32> {
                            use std::os::unix::process::ExitStatusExt;
                            exit_status.signal()
                        }

                        #[cfg(not(unix))]
                        fn exit_signal(_exit_status: &ExitStatus) -> Option<i32> {
                            None
                        }

                        if let Some(code) = exit_status.code() {
                            code
                        } else if let Some(signal) = exit_signal(&exit_status) {
                            128 + signal
                        } else {
                            127
                        }
                    },

                    Err(e) => {
                        writeln!(&mut io::stderr(), "shroom: {}: {}", command, e).unwrap();
                        127
                    },
                }
            }
        },
    }
}

fn main() {
    let mut line = String::new();
    loop {
        prompt(&mut line).unwrap();

        match Parser::new(&line).parse() {
            Ok(ast) => {
                let exit_code = execute(&ast);
                if exit_code != 0 {
                    println!("shroom: exit code: {}", exit_code);
                }
            },

            Err(parse_error) => {
                println!("shroom: parse error: {}", parse_error);
            },
        }

        line.clear();
    }
}
<|endoftext|>