<commit_before>use js_types::js_type::{JsT, JsType};
use std::vec::Vec;
use std::collections::hash_set::HashSet;
use std::collections::vec_deque::VecDeque;
use uuid::Uuid;

/// A graph node that stores a set of object UIDs, as well as a list of
/// pointers to nodes it can flow to.
/// A variable v is live on edge e if there is:
///     - a node n in the CFG s.t. use[n] contains v
///     - a directed path from e -> n s.t. for every statement s on the path,
///       def[s] does not contain v
#[derive(Clone)]
struct LivenessNode {
    defs: HashSet<Uuid>,
    uses: HashSet<Uuid>,
    ins: HashSet<Uuid>,
    outs: HashSet<Uuid>,
    succ: Vec<Box<LivenessNode>>,
    pred: Vec<Box<LivenessNode>>,
}

struct LivenessGraph {
    nodes: Vec<LivenessNode>,
}

impl LivenessNode {
    fn new() -> LivenessNode {
        LivenessNode {
            defs: HashSet::new(),
            uses: HashSet::new(),
            ins: HashSet::new(),
            outs: HashSet::new(),
            succ: Vec::new(),
            pred: Vec::new(),
        }
    }

    fn node_flow(&mut self) {
        self.ins = self.uses.union(&self.outs.difference(&self.defs).cloned().collect())
                            .cloned().collect();

        self.outs = self.succ.clone().into_iter().flat_map(|s| s.ins).collect();
    }

    // TODO Given a list of AST instructions, parse out what UUIDs are used

}

impl LivenessGraph {
    fn new() -> LivenessGraph {
        LivenessGraph {
            nodes: Vec::new(),
        }
    }

    fn graph_flow(&mut self) {
        let mut node_queue: VecDeque<&mut LivenessNode> = self.nodes.iter_mut().collect();
        while let Some(n) = node_queue.pop_front() {
            let old_ins = n.ins.clone();
            let old_outs = n.outs.clone();
            n.node_flow();
            if old_outs != n.outs {
                n.succ.iter_mut().map(|s| node_queue.push_back(s));
            }
        }
    }
}

pub struct StackFrame<'m> {
    members: Vec<&'m JsT>,
}

impl<'m> StackFrame<'m> {
    pub fn new() -> StackFrame<'m> {
        StackFrame { members: Vec::new(), }
    }

    pub fn alloc_ref(&mut self, refce: &'m mut JsT)  {
        self.members.push(refce);
    }
}
<commit_msg>Add stubs for compute_uses & compute_defs for LivenessNodes<commit_after>use js_types::js_type::{JsT, JsType};
use std::vec::Vec;
use std::collections::hash_set::HashSet;
use std::collections::vec_deque::VecDeque;
use uuid::Uuid;

/// A graph node that stores a set of object UIDs, as well as a list of
/// pointers to nodes it can flow to.
/// A variable v is live on edge e if there is:
///     - a node n in the CFG s.t. use[n] contains v
///     - a directed path from e -> n s.t. for every statement s on the path,
///       def[s] does not contain v
#[derive(Clone)]
struct LivenessNode {
    defs: HashSet<Uuid>,
    uses: HashSet<Uuid>,
    ins: HashSet<Uuid>,
    outs: HashSet<Uuid>,
    succ: Vec<Box<LivenessNode>>,
    pred: Vec<Box<LivenessNode>>,
}

struct LivenessGraph {
    nodes: Vec<LivenessNode>,
}

type Statement = i32; // TODO, placeholder

impl LivenessNode {
    fn new() -> LivenessNode {
        LivenessNode {
            defs: HashSet::new(),
            uses: HashSet::new(),
            ins: HashSet::new(),
            outs: HashSet::new(),
            succ: Vec::new(),
            pred: Vec::new(),
        }
    }

    fn node_flow(&mut self) {
        self.ins = self.uses.union(&self.outs.difference(&self.defs).cloned().collect())
                            .cloned().collect();

        self.outs = self.succ.clone().into_iter().flat_map(|s| s.ins).collect();
    }

    // TODO Given a list of AST instructions, parse out what UUIDs are used
    fn compute_uses(&mut self, statements: Vec<Statement>) {

    }

    fn compute_defs(&mut self, statements: Vec<Statement>) {

    }

}

impl LivenessGraph {
    fn new() -> LivenessGraph {
        LivenessGraph {
            nodes: Vec::new(),
        }
    }

    fn graph_flow(&mut self) {
        let mut node_queue: VecDeque<&mut LivenessNode> = self.nodes.iter_mut().collect();
        while let Some(n) = node_queue.pop_front() {
            let old_ins = n.ins.clone();
            let old_outs = n.outs.clone();
            n.node_flow();
            if old_outs != n.outs {
                n.succ.iter_mut().map(|s| node_queue.push_back(s));
            }
        }
    }
}

pub struct StackFrame<'m> {
    members: Vec<&'m JsT>,
}

impl<'m> StackFrame<'m> {
    pub fn new() -> StackFrame<'m> {
        StackFrame { members: Vec::new(), }
    }

    pub fn alloc_ref(&mut self, refce: &'m mut JsT)  {
        self.members.push(refce);
    }
}
<|endoftext|>


<commit_before>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// OCSP request is defined in
// https://tools.ietf.org/html/rfc6960#section-4.1
//
// OCSP over http is defined in
// https://tools.ietf.org/html/rfc2560#appendix-A.1

use crate::crypto::{CertificateChain, HashAlgorithm};
use crate::fetcher::Fetcher;
use crate::http::{HttpRequest, Method};
use crate::runtime::Runtime;
use anyhow::{anyhow, Error, Result};
use der_parser::{
    ber::{BerObject, BerObjectContent},
    oid,
    oid::Oid,
};
use serde::{Deserialize, Serialize};
use std::time::{Duration, SystemTime};
use tokio::sync::Mutex;
use x509_parser::{
    certificate::X509Certificate,
    extensions::{GeneralName, ParsedExtension},
};

fn create_ocsp_request(cert: &X509Certificate, issuer: &X509Certificate) -> Vec<u8> {
    let hash_algorithm = HashAlgorithm::Sha1;
    let issuer_name = issuer.tbs_certificate.subject.as_raw();
    let issuer_key = issuer
        .tbs_certificate
        .subject_pki
        .subject_public_key
        .data
        .as_ref();
    let issuer_key_hash = hash_algorithm.digest(issuer_key);
    let issuer_name_hash = hash_algorithm.digest(issuer_name);
    let serial_number = cert.tbs_certificate.raw_serial();
    // https://tools.ietf.org/html/rfc6960#section-4.1.1
    // CertID          ::=     SEQUENCE {
    //     hashAlgorithm       AlgorithmIdentifier,
    //     issuerNameHash      OCTET STRING, -- Hash of issuer's DN
    //     issuerKeyHash       OCTET STRING, -- Hash of issuer's public key
    //     serialNumber        CertificateSerialNumber }
    let cert_id = BerObject::from_seq(vec![
        hash_algorithm.to_ber(),
        BerObject::from_obj(BerObjectContent::OctetString(&issuer_name_hash)),
        BerObject::from_obj(BerObjectContent::OctetString(&issuer_key_hash)),
        BerObject::from_obj(BerObjectContent::Integer(serial_number)),
    ]);
    // https://tools.ietf.org/html/rfc6960#section-4.1.1
    // Request         ::=     SEQUENCE {
    //     reqCert                     CertID,
    //     singleRequestExtensions     [0] EXPLICIT Extensions OPTIONAL }
    let request = BerObject::from_seq(vec![cert_id]);
    // https://tools.ietf.org/html/rfc6960#section-4.1.1
    // TBSRequest      ::=     SEQUENCE {
    //     version             [0]     EXPLICIT Version DEFAULT v1,
    //     requestorName       [1]     EXPLICIT GeneralName OPTIONAL,
    //     requestList                 SEQUENCE OF Request,
    //     requestExtensions   [2]     EXPLICIT Extensions OPTIONAL }
    let tbs_request = BerObject::from_seq(vec![BerObject::from_seq(vec![request])]);
    // https://tools.ietf.org/html/rfc6960#section-4.1.1
    // OCSPRequest     ::=     SEQUENCE {
    //     tbsRequest                  TBSRequest,
    //     optionalSignature   [0]     EXPLICIT Signature OPTIONAL }
    let ocsp_request = BerObject::from_seq(vec![tbs_request]);
    ocsp_request.to_vec().unwrap()
}

// https://datatracker.ietf.org/doc/html/rfc4325#section-2
// https://datatracker.ietf.org/doc/html/rfc3280#section-4.2.2.1
const AIA: Oid<'static> = oid!(1.3.6 .1 .5 .5 .7 .1 .1);
// https://www.iana.org/assignments/smi-numbers/smi-numbers.xhtml#smi-numbers-1.3.6.1.5.5.7.48.1
const AIA_OCSP: Oid<'static> = oid!(1.3.6 .1 .5 .5 .7 .48 .1);

pub async fn fetch_from_ca(
    cert_der: &[u8],
    issuer_der: &[u8],
    fetcher: &dyn Fetcher,
) -> Result<Vec<u8>> {
    let cert = x509_parser::parse_x509_certificate(cert_der)
        .map_err(|e| Error::from(e).context("Failed to parse cert DER"))?
        .1;
    let issuer = x509_parser::parse_x509_certificate(issuer_der)
        .map_err(|e| Error::from(e).context("Failed to parse issuer DER"))?
        .1;
    let aia = cert.extensions().iter().find(|ext| ext.oid == AIA);
    let aia = if let Some(aia) = aia {
        aia
    } else {
        // If the certificate doesn't include an AIA section, it is probably a
        // self-signed certificate. Return a stub OCSP response.
        return Ok(b"ocsp".to_vec());
    };
    let aia = if let ParsedExtension::AuthorityInfoAccess(aia) = aia.parsed_extension() {
        aia
    } else {
        return Err(anyhow!("Failed to parse AIA extension"));
    };
    let url = aia
        .accessdescs
        .iter()
        .find_map(|access_desc| {
            if access_desc.access_method == AIA_OCSP {
                match access_desc.access_location {
                    GeneralName::URI(url) => Some(url),
                    _ => None,
                }
            } else {
                None
            }
        })
        .ok_or_else(|| anyhow!("AIA OCSP responder with type of URI is not found."))?;
    let req = HttpRequest {
        body: create_ocsp_request(&cert, &issuer),
        headers: vec![(
            String::from("content-type"),
            String::from("application/ocsp-request"),
        )],
        method: Method::Post,
        url: url.into(),
    };
    let rsp = fetcher
        .fetch(req)
        .await
        .map_err(|e| e.context("Failed to fetch OCSP"))?;
    Ok(rsp.body)
}

const OCSP_KEY: &str = "OCSP";

#[derive(Serialize, Deserialize)]
struct OcspData {
    pub expiration_time: SystemTime,
    pub recommended_update_time: SystemTime,
    #[serde(with = "crate::serde_helpers::base64")]
    pub value: Vec<u8>,
}

pub enum OcspUpdateStrategy {
    EarlyAsRecommended,
    LazyIfUnexpired,
}

/// Reads OCSP in storage, checks the expiration status, and returns latest.
/// If OCSP in storage needs update, fetches it from the server and writes it
/// into storage. The outging traffic to the server is throttled to be a
/// single task.
pub async fn read_and_update_ocsp_in_storage(
    certificate_chain: &CertificateChain,
    runtime: &Runtime,
    strategy: OcspUpdateStrategy,
) -> Result<Vec<u8>> {
    // Checks whether we can directly return the existing OCSP in storage.
    if let Some(old_ocsp) = runtime.storage.read(OCSP_KEY).await? {
        if let Ok(old_ocsp) = serde_json::from_str::<OcspData>(&old_ocsp) {
            match strategy {
                OcspUpdateStrategy::EarlyAsRecommended => {
                    if old_ocsp.recommended_update_time > runtime.now {
                        return Ok(old_ocsp.value);
                    }
                }
                OcspUpdateStrategy::LazyIfUnexpired => {
                    if old_ocsp.expiration_time > runtime.now {
                        return Ok(old_ocsp.value);
                    }
                }
            }
        } else {
            // The existing OCSP in storage can't be parsed as `OcspData`.
        }
    } else {
        // There is no OCSP in storage.
    }
    if certificate_chain.issuers.is_empty() {
        return Err(Error::msg("Certiciate chain contains no issuer."));
    }
    let cert_der = &certificate_chain.end_entity.der;
    let issuer_der = &certificate_chain.issuers[0].der;
    let new_ocsp_value = {
        static SINGLE_TASK: Mutex<()> = Mutex::const_new(());
        let guard = SINGLE_TASK.lock().await;
        let ocsp = fetch_from_ca(cert_der, issuer_der, runtime.fetcher.as_ref()).await?;
        std::mem::drop(guard);
        ocsp
    };
    const SIX_DAYS: Duration = Duration::from_secs(3600 * 24 * 6);
    const ONE_DAY: Duration = Duration::from_secs(3600 * 24);
    let new_ocsp = OcspData {
        expiration_time: runtime.now + SIX_DAYS,
        recommended_update_time: runtime.now + ONE_DAY,
        value: new_ocsp_value,
    };
    runtime
        .storage
        .write(OCSP_KEY, &serde_json::to_string(&new_ocsp)?)
        .await?;
    Ok(new_ocsp.value)
}
<commit_msg>Ignore storage errors for reading and writing OCSP (#354)<commit_after>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// OCSP request is defined in
// https://tools.ietf.org/html/rfc6960#section-4.1
//
// OCSP over http is defined in
// https://tools.ietf.org/html/rfc2560#appendix-A.1

use crate::crypto::{CertificateChain, HashAlgorithm};
use crate::fetcher::Fetcher;
use crate::http::{HttpRequest, Method};
use crate::runtime::Runtime;
use anyhow::{anyhow, Error, Result};
use der_parser::{
    ber::{BerObject, BerObjectContent},
    oid,
    oid::Oid,
};
use serde::{Deserialize, Serialize};
use std::time::{Duration, SystemTime};
use tokio::sync::Mutex;
use x509_parser::{
    certificate::X509Certificate,
    extensions::{GeneralName, ParsedExtension},
};

fn create_ocsp_request(cert: &X509Certificate, issuer: &X509Certificate) -> Vec<u8> {
    let hash_algorithm = HashAlgorithm::Sha1;
    let issuer_name = issuer.tbs_certificate.subject.as_raw();
    let issuer_key = issuer
        .tbs_certificate
        .subject_pki
        .subject_public_key
        .data
        .as_ref();
    let issuer_key_hash = hash_algorithm.digest(issuer_key);
    let issuer_name_hash = hash_algorithm.digest(issuer_name);
    let serial_number = cert.tbs_certificate.raw_serial();
    // https://tools.ietf.org/html/rfc6960#section-4.1.1
    // CertID          ::=     SEQUENCE {
    //     hashAlgorithm       AlgorithmIdentifier,
    //     issuerNameHash      OCTET STRING, -- Hash of issuer's DN
    //     issuerKeyHash       OCTET STRING, -- Hash of issuer's public key
    //     serialNumber        CertificateSerialNumber }
    let cert_id = BerObject::from_seq(vec![
        hash_algorithm.to_ber(),
        BerObject::from_obj(BerObjectContent::OctetString(&issuer_name_hash)),
        BerObject::from_obj(BerObjectContent::OctetString(&issuer_key_hash)),
        BerObject::from_obj(BerObjectContent::Integer(serial_number)),
    ]);
    // https://tools.ietf.org/html/rfc6960#section-4.1.1
    // Request         ::=     SEQUENCE {
    //     reqCert                     CertID,
    //     singleRequestExtensions     [0] EXPLICIT Extensions OPTIONAL }
    let request = BerObject::from_seq(vec![cert_id]);
    // https://tools.ietf.org/html/rfc6960#section-4.1.1
    // TBSRequest      ::=     SEQUENCE {
    //     version             [0]     EXPLICIT Version DEFAULT v1,
    //     requestorName       [1]     EXPLICIT GeneralName OPTIONAL,
    //     requestList                 SEQUENCE OF Request,
    //     requestExtensions   [2]     EXPLICIT Extensions OPTIONAL }
    let tbs_request = BerObject::from_seq(vec![BerObject::from_seq(vec![request])]);
    // https://tools.ietf.org/html/rfc6960#section-4.1.1
    // OCSPRequest     ::=     SEQUENCE {
    //     tbsRequest                  TBSRequest,
    //     optionalSignature   [0]     EXPLICIT Signature OPTIONAL }
    let ocsp_request = BerObject::from_seq(vec![tbs_request]);
    ocsp_request.to_vec().unwrap()
}

// https://datatracker.ietf.org/doc/html/rfc4325#section-2
// https://datatracker.ietf.org/doc/html/rfc3280#section-4.2.2.1
const AIA: Oid<'static> = oid!(1.3.6 .1 .5 .5 .7 .1 .1);
// https://www.iana.org/assignments/smi-numbers/smi-numbers.xhtml#smi-numbers-1.3.6.1.5.5.7.48.1
const AIA_OCSP: Oid<'static> = oid!(1.3.6 .1 .5 .5 .7 .48 .1);

pub async fn fetch_from_ca(
    cert_der: &[u8],
    issuer_der: &[u8],
    fetcher: &dyn Fetcher,
) -> Result<Vec<u8>> {
    let cert = x509_parser::parse_x509_certificate(cert_der)
        .map_err(|e| Error::from(e).context("Failed to parse cert DER"))?
        .1;
    let issuer = x509_parser::parse_x509_certificate(issuer_der)
        .map_err(|e| Error::from(e).context("Failed to parse issuer DER"))?
        .1;
    let aia = cert.extensions().iter().find(|ext| ext.oid == AIA);
    let aia = if let Some(aia) = aia {
        aia
    } else {
        // If the certificate doesn't include an AIA section, it is probably a
        // self-signed certificate. Return a stub OCSP response.
        return Ok(b"ocsp".to_vec());
    };
    let aia = if let ParsedExtension::AuthorityInfoAccess(aia) = aia.parsed_extension() {
        aia
    } else {
        return Err(anyhow!("Failed to parse AIA extension"));
    };
    let url = aia
        .accessdescs
        .iter()
        .find_map(|access_desc| {
            if access_desc.access_method == AIA_OCSP {
                match access_desc.access_location {
                    GeneralName::URI(url) => Some(url),
                    _ => None,
                }
            } else {
                None
            }
        })
        .ok_or_else(|| anyhow!("AIA OCSP responder with type of URI is not found."))?;
    let req = HttpRequest {
        body: create_ocsp_request(&cert, &issuer),
        headers: vec![(
            String::from("content-type"),
            String::from("application/ocsp-request"),
        )],
        method: Method::Post,
        url: url.into(),
    };
    let rsp = fetcher
        .fetch(req)
        .await
        .map_err(|e| e.context("Failed to fetch OCSP"))?;
    Ok(rsp.body)
}

const OCSP_KEY: &str = "OCSP";

#[derive(Serialize, Deserialize)]
struct OcspData {
    pub expiration_time: SystemTime,
    pub recommended_update_time: SystemTime,
    #[serde(with = "crate::serde_helpers::base64")]
    pub value: Vec<u8>,
}

pub enum OcspUpdateStrategy {
    EarlyAsRecommended,
    LazyIfUnexpired,
}

/// Reads OCSP in storage, checks the expiration status, and returns latest.
/// If OCSP in storage needs update, fetches it from the server and writes it
/// into storage. The outging traffic to the server is throttled to be a
/// single task.
/// If there is any error when reading from and writing to storage, the error
/// will be ignored but an error message will be printed to output/log.
pub async fn read_and_update_ocsp_in_storage(
    certificate_chain: &CertificateChain,
    runtime: &Runtime,
    strategy: OcspUpdateStrategy,
) -> Result<Vec<u8>> {
    // Checks whether we can directly return the existing OCSP in storage.
    match runtime.storage.read(OCSP_KEY).await {
        Ok(Some(old_ocsp)) => {
            if let Ok(old_ocsp) = serde_json::from_str::<OcspData>(&old_ocsp) {
                match strategy {
                    OcspUpdateStrategy::EarlyAsRecommended => {
                        if old_ocsp.recommended_update_time > runtime.now {
                            return Ok(old_ocsp.value);
                        }
                    }
                    OcspUpdateStrategy::LazyIfUnexpired => {
                        if old_ocsp.expiration_time > runtime.now {
                            return Ok(old_ocsp.value);
                        }
                    }
                }
            } else {
                // The existing OCSP in storage can't be parsed as `OcspData`.
            }
        }
        Ok(None) => {
            // There is no OCSP in storage.
        }
        Err(e) => {
            println!("Failed to read OCSP from storage. {}", e);
        }
    }
    if certificate_chain.issuers.is_empty() {
        return Err(Error::msg("Certiciate chain contains no issuer."));
    }
    let cert_der = &certificate_chain.end_entity.der;
    let issuer_der = &certificate_chain.issuers[0].der;
    let new_ocsp_value = {
        static SINGLE_TASK: Mutex<()> = Mutex::const_new(());
        let guard = SINGLE_TASK.lock().await;
        let ocsp = fetch_from_ca(cert_der, issuer_der, runtime.fetcher.as_ref()).await?;
        std::mem::drop(guard);
        ocsp
    };
    const SIX_DAYS: Duration = Duration::from_secs(3600 * 24 * 6);
    const ONE_DAY: Duration = Duration::from_secs(3600 * 24);
    let new_ocsp = OcspData {
        expiration_time: runtime.now + SIX_DAYS,
        recommended_update_time: runtime.now + ONE_DAY,
        value: new_ocsp_value,
    };
    let write_result = runtime
        .storage
        .write(OCSP_KEY, &serde_json::to_string(&new_ocsp)?)
        .await;
    if let Err(e) = write_result {
        println!("Failed to write OCSP to storage. {}", e);
    }
    Ok(new_ocsp.value)
}
<|endoftext|>


<commit_before>#[cfg(feature = "trace_log")]
extern crate env_logger;

extern crate radeco_lib;
extern crate r2pipe;
extern crate r2api;
extern crate base64;
extern crate rustyline;

use radeco_lib::analysis::cse::cse::CSE;
use radeco_lib::analysis::interproc::fixcall::CallFixer;
use radeco_lib::analysis::sccp;
use radeco_lib::backend::lang_c::c_cfg::CCFGVerifier;
use radeco_lib::backend::lang_c::c_cfg_builder;
use radeco_lib::frontend::radeco_containers::*;
use radeco_lib::middle::{dce, dot};
use radeco_lib::middle::ir_reader::parse_il;
use radeco_lib::middle::ir_writer;
use radeco_lib::middle::regfile::SubRegisterFile;
use radeco_lib::middle::ssa::ssastorage::SSAStorage;
use radeco_lib::middle::ssa::verifier;
use rustyline::Editor;

use rustyline::error::ReadlineError;

use std::collections::HashMap;
use std::process;
use std::str;
use std::sync::Arc;

// On unix platforms you can use ANSI escape sequences
#[cfg(unix)]
static PROMPT: &'static str = "\x1b[1;32m>>\x1b[0m ";

// Windows consoles typically don't support ANSI escape sequences out
// of the box
#[cfg(windows)]
static PROMPT: &'static str = ">> ";

fn main() {
    let mut rl = Editor::<()>::new();
    let mut proj = None;
    loop {
        let readline = rl.readline(">> ");
        match readline {
            Ok(line) => {
                let mut terms = line.split_whitespace();
                let o1 = terms.next();
                let o2 = terms.next();
                cmd(o1, o2, &mut proj);
            }
            Err(ReadlineError::Interrupted) |
            Err(ReadlineError::Eof) => break,
            Err(err) => {
                println!("Error: {:?}", err);
                break;
            }
        }
    }
}

mod command {
    pub const LOAD: &'static str = "load";
    pub const FNLIST: &'static str = "fn_list";
    pub const ANALYZE: &'static str = "analyze";
    pub const DOT: &'static str = "dot";
    pub const IR: &'static str = "ir";
    pub const DECOMPILE: &'static str = "decompile";
}

fn cmd(op1: Option<&str>, op2: Option<&str>, proj_opt: &mut Option<RadecoProject>) {
    match (op1, op2) {
        (Some(command::LOAD), Some(s)) => {
            let mut p = ProjectLoader::new().path(s).load();
            let regfile = p.regfile().clone();
            for mut xy in p.iter_mut() {
                analyze_mod(regfile.clone(), xy.module);
            }
            *proj_opt = Some(p);
            return;
        }
        _ => {}
    };
    if let Some(proj) = proj_opt {
        match (op1, op2) {
            (Some(command::ANALYZE), Some("*")) => {
                let rfns = proj.iter_mut().map(|i| i.module).flat_map(|rmod| {
                    rmod.functions.values_mut()
                });
                for rfn in rfns {
                    analyze(rfn);
                }
            }
            (Some(command::FNLIST), _) => {
                let funcs = fn_list(&proj);
                println!("{:?}", funcs);
            }
            (Some(command::ANALYZE), Some(f)) => {
                if let Some(rfn) = get_function_mut(f, proj) {
                    analyze(rfn);
                } else {
                    println!("{} is not found", f);
                }
            }
            (Some(command::DOT), Some(f)) => {
                if let Some(rfn) = get_function(f, &proj) {
                    println!("{}", emit_dot(rfn.ssa()));
                } else {
                    println!("{} is not found", f);
                }
            }
            (Some(command::IR), Some(f)) => {
                if let Some(rfn) = get_function(f, &proj) {
                    println!("{}", emit_ir(rfn));
                } else {
                    println!("{} is not found", f);
                }
            }
            (Some(command::DECOMPILE), Some(f)) => {
                if let Some(rfn) = get_function(f, &proj) {
                    let rmod = proj.iter().map(|i| i.module).next().unwrap();
                    let func_name_map = rmod.functions
                        .iter()
                        .map(|(&addr, f)| (addr, f.name.to_string()))
                        .collect();
                    let strings = rmod.strings()
                        .iter()
                        .filter(|ref s| s.vaddr.is_some() && s.string.is_some())
                        .cloned()
                        .map(|s| {
                            let (addr, _s) = (s.vaddr.unwrap(), s.string.unwrap());
                            let bytes = base64::decode(&_s).unwrap_or(Vec::new());
                            let ret_string = match str::from_utf8(bytes.as_slice()) {
                                Ok(v) => v.to_string(),
                                Err(_e) => _s,
                            };
                            (addr, ret_string)
                        })
                        .collect();
                    println!("{}", decompile(rfn, &func_name_map, &strings));
                } else {
                    println!("{} is not found", f);
                }
            }
            _ => {}
        }
    } else {
        println!("Load a project first");
    }
}

fn fn_list(proj: &RadecoProject) -> Vec<String> {
    proj.iter()
        .map(|i| i.module)
        .flat_map(|rmod| rmod.functions.values())
        .map(|s| s.name.to_string())
        .collect()
}

fn get_function<'a>(name: &str, proj: &'a RadecoProject) -> Option<&'a RadecoFunction> {
    proj.iter()
        .map(|i| i.module)
        .flat_map(|rmod| rmod.functions.values())
        .filter(|rfn| rfn.name == name)
        .next()
}

fn get_function_mut<'a>(name: &str, proj: &'a mut RadecoProject) -> Option<&'a mut RadecoFunction> {
    proj.iter_mut()
        .map(|i| i.module)
        .flat_map(|rmod| rmod.functions.values_mut())
        .filter(|rfn| rfn.name == name)
        .next()
}

fn analyze_mod(regfile: Arc<SubRegisterFile>, rmod: &mut RadecoModule) {
    // Analyze preserved for all functions.
    {
        println!("[*] Fixing Callee Information");
        let bp_name = regfile.get_name_by_alias(&"BP".to_string());
        let bp_name = bp_name.map(|s| s.to_owned());
        let sp_name = regfile.get_name_by_alias(&"SP".to_string());
        let sp_name = sp_name.map(|s| s.to_owned());
        let mut callfixer = CallFixer::new(rmod, bp_name, sp_name);
        callfixer.rounded_analysis();
    }

    // Fix call sites
    radeco_lib::analysis::functions::fix_ssa_opcalls::go(rmod);

    // Infer calling conventions
    radeco_lib::analysis::functions::infer_regusage::run(rmod, &*regfile);
}

fn analyze(rfn: &mut RadecoFunction) {
    println!("[+] Analyzing: {} @ {:#x}", rfn.name, rfn.offset);
    {
        println!("  [*] Eliminating Dead Code");
        dce::collect(rfn.ssa_mut());
    }
    let mut ssa = {
        // Constant Propagation (sccp)
        println!("  [*] Propagating Constants");
        let mut analyzer = sccp::Analyzer::new(rfn.ssa_mut());
        analyzer.analyze();
        analyzer.emit_ssa()
    };
    {
        println!("  [*] Eliminating More DeadCode");
        dce::collect(&mut ssa);
    }
    *rfn.ssa_mut() = ssa;
    {
        // Common SubExpression Elimination (cse)
        println!("  [*] Eliminating Common SubExpressions");
        let mut cse = CSE::new(rfn.ssa_mut());
        cse.run();
    }
    {
        // Verify SSA
        println!("  [*] Verifying SSA's Validity");
        match verifier::verify(rfn.ssa()) {
            Err(e) => {
                println!("  [*] Found Error: {}", e);
                process::exit(255);
            }
            Ok(_) => {}
        }
    }
}

fn emit_ir(rfn: &RadecoFunction) -> String {
    println!("  [*] Writing out IR");
    let mut res = String::new();
    ir_writer::emit_il(&mut res, Some(rfn.name.to_string()), rfn.ssa()).unwrap();
    res
}

fn parse_ir(name: &str, ir_str: &str, regfile: Arc<SubRegisterFile>) {
    println!("  [*] Testing IR parser");
    let parsed = parse_il(ir_str, regfile.clone());
    let mut res = String::new();
    ir_writer::emit_il(&mut res, Some(name.to_string()), &parsed).unwrap();
    for (orig, roundtrip) in ir_str.lines().zip(res.lines()) {
        if orig != roundtrip {
            println!("  FAILED TO ROUND-TRIP: \"{}\" => \"{}\"", orig, roundtrip);
        }
    }
}

fn emit_dot(ssa: &SSAStorage) -> String {
    dot::emit_dot(ssa)
}

fn decompile(
    rfn: &RadecoFunction,
    func_name_map: &HashMap<u64, String>,
    strings: &HashMap<u64, String>,
) -> String {
    let c_cfg = c_cfg_builder::recover_c_cfg(rfn, func_name_map, strings);
    if let Err(err) = CCFGVerifier::verify(&c_cfg) {
        println!("CCFG verification failed");
        println!("{}", err);
    }
    c_cfg.to_c_ast().print()
}
<commit_msg>Minor fixes<commit_after>#[cfg(feature = "trace_log")]
extern crate env_logger;

extern crate radeco_lib;
extern crate r2pipe;
extern crate r2api;
extern crate base64;
extern crate rustyline;

use radeco_lib::analysis::cse::cse::CSE;
use radeco_lib::analysis::interproc::fixcall::CallFixer;
use radeco_lib::analysis::sccp;
use radeco_lib::backend::lang_c::c_cfg::CCFGVerifier;
use radeco_lib::backend::lang_c::c_cfg_builder;
use radeco_lib::frontend::radeco_containers::*;
use radeco_lib::middle::{dce, dot};
use radeco_lib::middle::ir_reader::parse_il;
use radeco_lib::middle::ir_writer;
use radeco_lib::middle::regfile::SubRegisterFile;
use radeco_lib::middle::ssa::ssastorage::SSAStorage;
use radeco_lib::middle::ssa::verifier;
use rustyline::Editor;

use rustyline::error::ReadlineError;

use std::collections::HashMap;
use std::process;
use std::str;
use std::sync::Arc;

// On unix platforms you can use ANSI escape sequences
#[cfg(unix)]
static PROMPT: &'static str = "\x1b[1;32m>>\x1b[0m ";

// Windows consoles typically don't support ANSI escape sequences out
// of the box
#[cfg(windows)]
static PROMPT: &'static str = ">> ";

fn main() {
    let mut rl = Editor::<()>::new();
    let mut proj = None;
    loop {
        let readline = rl.readline(">> ");
        match readline {
            Ok(line) => {
                let mut terms = line.split_whitespace();
                let o1 = terms.next();
                let o2 = terms.next();
                cmd(o1, o2, &mut proj);
            }
            Err(ReadlineError::Interrupted) |
            Err(ReadlineError::Eof) => break,
            Err(err) => {
                println!("Error: {:?}", err);
                break;
            }
        }
    }
}

mod command {
    pub const LOAD: &'static str = "load";
    pub const FNLIST: &'static str = "fn_list";
    pub const ANALYZE: &'static str = "analyze";
    pub const DOT: &'static str = "dot";
    pub const IR: &'static str = "ir";
    pub const DECOMPILE: &'static str = "decompile";
}

fn cmd(op1: Option<&str>, op2: Option<&str>, proj_opt: &mut Option<RadecoProject>) {
    match (op1, op2) {
        (Some(command::LOAD), Some(s)) => {
            let mut p = ProjectLoader::new().path(s).load();
            let regfile = p.regfile().clone();
            for mut xy in p.iter_mut() {
                analyze_mod(regfile.clone(), xy.module);
            }
            *proj_opt = Some(p);
            return;
        }
        _ => {}
    };
    if proj_opt.is_none() {
        println!("Load a project first");
        return;
    }
    let mut proj = proj_opt.as_mut().unwrap();
    match (op1, op2) {
        (Some(command::ANALYZE), Some("*")) => {
            let rfns = proj.iter_mut().map(|i| i.module).flat_map(|rmod| {
                rmod.functions.values_mut()
            });
            for rfn in rfns {
                analyze(rfn);
            }
        }
        (Some(command::FNLIST), _) => {
            let funcs = fn_list(&proj);
            println!("{}", funcs.join("\n"));
        }
        (Some(command::ANALYZE), Some(f)) => {
            if let Some(rfn) = get_function_mut(f, &mut proj) {
                analyze(rfn);
            } else {
                println!("{} is not found", f);
            }
        }
        (Some(command::DOT), Some(f)) => {
            if let Some(rfn) = get_function(f, &proj) {
                println!("{}", emit_dot(rfn.ssa()));
            } else {
                println!("{} is not found", f);
            }
        }
        (Some(command::IR), Some(f)) => {
            if let Some(rfn) = get_function(f, &proj) {
                println!("{}", emit_ir(rfn));
            } else {
                println!("{} is not found", f);
            }
        }
        (Some(command::DECOMPILE), Some(f)) => {
            if let Some(rfn) = get_function(f, &proj) {
                let rmod = proj.iter().map(|i| i.module).next().unwrap();
                let func_name_map = rmod.functions
                    .iter()
                    .map(|(&addr, f)| (addr, f.name.to_string()))
                    .collect();
                let strings = rmod.strings()
                    .iter()
                    .filter(|ref s| s.vaddr.is_some() && s.string.is_some())
                    .cloned()
                    .map(|s| {
                        let (addr, _s) = (s.vaddr.unwrap(), s.string.unwrap());
                        let bytes = base64::decode(&_s).unwrap_or(Vec::new());
                        let ret_string = match str::from_utf8(bytes.as_slice()) {
                            Ok(v) => v.to_string(),
                            Err(_e) => _s,
                        };
                        (addr, ret_string)
                    })
                    .collect();
                println!("{}", decompile(rfn, &func_name_map, &strings));
            } else {
                println!("{} is not found", f);
            }
        }
        _ => {}
    }
}

fn fn_list(proj: &RadecoProject) -> Vec<String> {
    proj.iter()
        .map(|i| i.module)
        .flat_map(|rmod| rmod.functions.values())
        .map(|s| s.name.to_string())
        .collect()
}

fn get_function<'a>(name: &str, proj: &'a RadecoProject) -> Option<&'a RadecoFunction> {
    proj.iter()
        .map(|i| i.module)
        .flat_map(|rmod| rmod.functions.values())
        .filter(|rfn| rfn.name == name)
        .next()
}

fn get_function_mut<'a>(name: &str, proj: &'a mut RadecoProject) -> Option<&'a mut RadecoFunction> {
    proj.iter_mut()
        .map(|i| i.module)
        .flat_map(|rmod| rmod.functions.values_mut())
        .filter(|rfn| rfn.name == name)
        .next()
}

fn analyze_mod(regfile: Arc<SubRegisterFile>, rmod: &mut RadecoModule) {
    // Analyze preserved for all functions.
    {
        println!("[*] Fixing Callee Information");
        let bp_name = regfile.get_name_by_alias(&"BP".to_string());
        let bp_name = bp_name.map(|s| s.to_owned());
        let sp_name = regfile.get_name_by_alias(&"SP".to_string());
        let sp_name = sp_name.map(|s| s.to_owned());
        let mut callfixer = CallFixer::new(rmod, bp_name, sp_name);
        callfixer.rounded_analysis();
    }

    // Fix call sites
    radeco_lib::analysis::functions::fix_ssa_opcalls::go(rmod);

    // Infer calling conventions
    radeco_lib::analysis::functions::infer_regusage::run(rmod, &*regfile);
}

fn analyze(rfn: &mut RadecoFunction) {
    println!("[+] Analyzing: {} @ {:#x}", rfn.name, rfn.offset);
    {
        println!("  [*] Eliminating Dead Code");
        dce::collect(rfn.ssa_mut());
    }
    let mut ssa = {
        // Constant Propagation (sccp)
        println!("  [*] Propagating Constants");
        let mut analyzer = sccp::Analyzer::new(rfn.ssa_mut());
        analyzer.analyze();
        analyzer.emit_ssa()
    };
    {
        println!("  [*] Eliminating More DeadCode");
        dce::collect(&mut ssa);
    }
    *rfn.ssa_mut() = ssa;
    {
        // Common SubExpression Elimination (cse)
        println!("  [*] Eliminating Common SubExpressions");
        let mut cse = CSE::new(rfn.ssa_mut());
        cse.run();
    }
    {
        // Verify SSA
        println!("  [*] Verifying SSA's Validity");
        match verifier::verify(rfn.ssa()) {
            Err(e) => {
                println!("  [*] Found Error: {}", e);
                process::exit(255);
            }
            Ok(_) => {}
        }
    }
}

fn emit_ir(rfn: &RadecoFunction) -> String {
    println!("  [*] Writing out IR");
    let mut res = String::new();
    ir_writer::emit_il(&mut res, Some(rfn.name.to_string()), rfn.ssa()).unwrap();
    res
}

fn parse_ir(name: &str, ir_str: &str, regfile: Arc<SubRegisterFile>) {
    println!("  [*] Testing IR parser");
    let parsed = parse_il(ir_str, regfile.clone());
    let mut res = String::new();
    ir_writer::emit_il(&mut res, Some(name.to_string()), &parsed).unwrap();
    for (orig, roundtrip) in ir_str.lines().zip(res.lines()) {
        if orig != roundtrip {
            println!("  FAILED TO ROUND-TRIP: \"{}\" => \"{}\"", orig, roundtrip);
        }
    }
}

fn emit_dot(ssa: &SSAStorage) -> String {
    dot::emit_dot(ssa)
}

fn decompile(
    rfn: &RadecoFunction,
    func_name_map: &HashMap<u64, String>,
    strings: &HashMap<u64, String>,
) -> String {
    let c_cfg = c_cfg_builder::recover_c_cfg(rfn, func_name_map, strings);
    if let Err(err) = CCFGVerifier::verify(&c_cfg) {
        println!("CCFG verification failed");
        println!("{}", err);
    }
    c_cfg.to_c_ast().print()
}
<|endoftext|>


<commit_before>
use autograd::{Variable, Function, FuncId};
use tensor::Tensor;
use std::collections::{HashSet, HashMap};

#[derive(Default, Clone)]
pub struct ExecutionEngine {}

type FnRefs = HashMap<Function, u32>;
type FnDependencies = HashMap<Function, FnRefs>;

impl ExecutionEngine {
    fn _compute_dependencies(function: FuncId) -> FnDependencies {
        unimplemented!()
    }
    fn _free_backward_dependency(dependencies: FnDependencies,
                                 prev_fn: FuncId,
                                 fn_: FuncId,
                                 arg_id: u32)
                                 -> usize {
        unimplemented!();
    }
    fn _is_ready_for_backward(dependencies: FnDependencies, function: FuncId) -> bool {
        unimplemented!()
    }
    fn _add_grad<T>(need_copy: HashSet<Function>,
                    prev_grad: HashMap<u32, &Tensor<T>>,
                    output_nr: u32,
                    d_prev_grad: &Tensor<T>) {
        unimplemented!()
    }
    pub fn run_backward<T>(arg: &mut Variable<T>,
                           gradient: &mut Tensor<T>,
                           retain_variables: bool) {
        unimplemented!()
    }
}
<commit_msg>partially implement run_backward<commit_after>
// silence warnings while still a WIP
#![allow(unused_variables)]
#![allow(dead_code)]
#![allow(unused_assignments)]

use autograd::{Variable, Function, FuncId, FuncIntf};
use tensor::{Tensor, TensorKind};
use std::collections::{HashSet, HashMap, VecDeque};

pub struct ExecutionEngine {}

type FnRefs = HashMap<Function, u32>;
type FnDependencies = HashMap<Function, FnRefs>;

impl ExecutionEngine {
    fn _compute_dependencies(function: &Function) -> FnDependencies {
        unimplemented!()
    }
    fn _free_backward_dependency(dependencies: FnDependencies,
                                 prev_fn: FuncId,
                                 fn_: FuncId,
                                 arg_id: u32)
                                 -> usize {
        unimplemented!();
    }
    fn _is_ready_for_backward(dependencies: FnDependencies, function: FuncId) -> bool {
        unimplemented!()
    }
    fn _add_grad<T>(need_copy: HashSet<Function>,
                    prev_grad: HashMap<u32, &Tensor<T>>,
                    output_nr: u32,
                    d_prev_grad: &Tensor<T>) {
        unimplemented!()
    }
    pub fn run_backward<T: Copy>(var: &mut Variable<T>, grad: TensorKind, retain_variables: bool) {
        let grad_fn;
        match var.grad_fn() {
            Some(v) => grad_fn = v,
            None => {
                let grad_tensor = Tensor::<T>::from(grad);
                var._do_backward(grad_tensor[0]);
                return;
            }
        }
        let mut ready = VecDeque::new();
        ready.push_back((grad_fn.clone(), vec![grad]));
        //let mut need_copy = HashSet::new();

        let dependencies = Self::_compute_dependencies(&grad_fn);
        while !ready.is_empty() {
            let (func, grad) = ready.pop_front().unwrap();
            let grad_input = func.clone()._do_backward(&grad, retain_variables);
        }
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0.
 */

use aws_config::meta::region::RegionProviderChain;
use aws_sdk_ec2::{Client, Error, Region, PKG_VERSION};
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
struct Opt {
    /// The AWS Region.
    #[structopt(short, long)]
    region: Option<String>,

    /// The ID of the instance to monitor.
    #[structopt(short, long)]
    instance_id: String,

    /// Whether to display additional information.
    #[structopt(short, long)]
    verbose: bool,
}

// Enables monitoring for an instance.
// snippet-start:[ec2.rust.monitor-instance]
async fn enable_monitoring(
    client: &aws_sdk_ec2::Client,
    id: &str,
) -> Result<(), aws_sdk_ec2::Error> {
    client.monitor_instances().instance_ids(id).send().await?;

    println!("Enabled monitoring");

    Ok(())
}
// snippet-end:[ec2.rust.monitor-instance]

/// Enables monitoring for an Amazon EC2 instance.
/// # Arguments
///
/// * `-i INSTANCE-ID` - The ID of the instances to monitor.
/// * `[-r REGION]` - The Region in which the client is created.
///   If not supplied, uses the value of the **AWS_REGION** environment variable.
///   If the environment variable is not set, defaults to **us-west-2**.
/// * `[-v]` - Whether to display additional information.
#[tokio::main]
async fn main() -> Result<(), Error> {
    tracing_subscriber::fmt::init();
    let Opt {
        region,
        instance_id,
        verbose,
    } = Opt::from_args();

    let region_provider = RegionProviderChain::first_try(region.map(Region::new))
        .or_default_provider()
        .or_else(Region::new("us-west-2"));
    println!();

    if verbose {
        println!("EC2 client version: {}", PKG_VERSION);
        println!(
            "Region:             {}",
            region_provider.region().await.unwrap().as_ref()
        );
        println!("Instance ID:        {}", instance_id);
        println!();
    }

    let shared_config = aws_config::from_env().region(region_provider).load().await;
    let client = Client::new(&shared_config);

    enable_monitoring(&client, &instance_id).await.unwrap();

    Ok(())
}

#[actix_rt::test]
async fn test_enable_monitoring() {
    let shared_config = aws_config::load_from_env().await;
    let client = Client::new(&shared_config);

    client.monitor_instances().instance_ids("id");
}
<commit_msg>Update monitor-instance.rs<commit_after>/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0.
 */

use aws_config::meta::region::RegionProviderChain;
use aws_sdk_ec2::{Client, Error, Region, PKG_VERSION};
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
struct Opt {
    /// The AWS Region.
    #[structopt(short, long)]
    region: Option<String>,

    /// The ID of the instance to monitor.
    #[structopt(short, long)]
    instance_id: String,

    /// Whether to display additional information.
    #[structopt(short, long)]
    verbose: bool,
}

// Enables monitoring for an instance.
// snippet-start:[ec2.rust.monitor-instance]
async fn enable_monitoring(
    client: &aws_sdk_ec2::Client,
    id: &str,
) -> Result<(), aws_sdk_ec2::Error> {
    client.monitor_instances().instance_ids(id).send().await?;

    println!("Enabled monitoring");

    Ok(())
}
// snippet-end:[ec2.rust.monitor-instance]

/// Enables monitoring for an Amazon EC2 instance.
/// # Arguments
///
/// * `-i INSTANCE-ID` - The ID of the instances to monitor.
/// * `[-r REGION]` - The Region in which the client is created.
///   If not supplied, uses the value of the **AWS_REGION** environment variable.
///   If the environment variable is not set, defaults to **us-west-2**.
/// * `[-v]` - Whether to display additional information.
#[tokio::main]
async fn main() -> Result<(), Error> {
    tracing_subscriber::fmt::init();
    let Opt {
        region,
        instance_id,
        verbose,
    } = Opt::from_args();

    let region_provider = RegionProviderChain::first_try(region.map(Region::new))
        .or_default_provider()
        .or_else(Region::new("us-west-2"));
    println!();

    if verbose {
        println!("EC2 client version: {}", PKG_VERSION);
        println!(
            "Region:             {}",
            region_provider.region().await.unwrap().as_ref()
        );
        println!("Instance ID:        {}", instance_id);
        println!();
    }

    let shared_config = aws_config::from_env().region(region_provider).load().await;
    let client = Client::new(&shared_config);

    enable_monitoring(&client, &instance_id).await
}
<|endoftext|>


<commit_before>use super::{Document,Root,RootChild,Element};

struct Parser;

impl Parser {
    fn new() -> Parser {
        Parser
    }

    fn parse_preamble<'a>(&self, xml: &'a str) -> &'a str {
        // Parse the preamble
        let idx = xml.find_str("?>").expect("No preamble end");
        let end_of_preamble = idx + "?>".len();
        xml.slice_from(end_of_preamble)
    }

    fn parse_element<'a>(&self, doc: Document, xml: &'a str) -> (Element, &'a str) {
        let (_, after_start_brace) = xml.slice_literal("<").expect("no start brace");

        let (name, after_name) = after_start_brace.slice_name().expect("failed to parse a name!");

        let (_, after_space) = after_name.slice_space().expect("no space after name");

        let (_, after_end_brace) = after_space.slice_literal("/>").expect("no end brace");

        let e = doc.new_element(name.to_string());

        (e, after_end_brace)
    }

    fn parse(&self, xml: &str) -> Document {
        let doc = Document::new();

        let after_preamble = self.parse_preamble(xml);

        let (element, _tail) = self.parse_element(doc.clone(), after_preamble);
        doc.root().append_child(element);

        doc
    }
}

trait XmlStr<'a> {
    fn slice_at(&self, position: uint) -> (&'a str, &'a str);
    fn slice_literal(&self, expected: &str) -> Option<(&'a str, &'a str)>;
    fn slice_start_rest(&self, is_first: |char| -> bool, is_rest: |char| -> bool) -> Option<(&'a str, &'a str)>;
    fn slice_name(&self) -> Option<(&'a str, &'a str)>;
    fn slice_space(&self) -> Option<(&'a str, &'a str)>;
}

impl<'a> XmlStr<'a> for &'a str {
    fn slice_at(&self, position: uint) -> (&'a str, &'a str) {
        (self.slice_to(position), self.slice_from(position))
    }

    fn slice_literal(&self, expected: &str) -> Option<(&'a str, &'a str)> {
        if self.starts_with(expected) {
            Some(self.slice_at(expected.len()))
        } else {
            None
        }
    }

    fn slice_start_rest(&self,
                        is_first: |char| -> bool,
                        is_rest: |char| -> bool)
                        -> Option<(&'a str, &'a str)>
    {
        let mut positions = self.char_indices();

        match positions.next() {
            Some((_, c)) if is_first(c) => (),
            Some((_, _)) => return None,
            None => return None,
        };

        let mut positions = positions.skip_while(|&(_, c)| is_rest(c));
        match positions.next() {
            Some((offset, _)) => Some(self.slice_at(offset)),
            None => Some((self.clone(), "")),
        }
    }

    fn slice_name(&self) -> Option<(&'a str, &'a str)> {
        self.slice_start_rest(|c| c.is_name_start_char(), |c| c.is_name_char())
    }

    fn slice_space(&self) -> Option<(&'a str, &'a str)> {
        self.slice_start_rest(|c| c.is_space_char(), |c| c.is_space_char())
    }
}

trait XmlChar {
    fn is_name_start_char(&self) -> bool;
    fn is_name_char(&self) -> bool;
    fn is_space_char(&self) -> bool;
}

impl XmlChar for char {
    fn is_name_start_char(&self) -> bool {
        match *self {
            ':'                        |
            'A'..'Z'                   |
            '_'                        |
            'a'..'z'                   |
            '\U000000C0'..'\U000000D6' |
            '\U000000D8'..'\U000000F6' |
            '\U000000F8'..'\U000002FF' |
            '\U00000370'..'\U0000037D' |
            '\U0000037F'..'\U00001FFF' |
            '\U0000200C'..'\U0000200D' |
            '\U00002070'..'\U0000218F' |
            '\U00002C00'..'\U00002FEF' |
            '\U00003001'..'\U0000D7FF' |
            '\U0000F900'..'\U0000FDCF' |
            '\U0000FDF0'..'\U0000FFFD' |
            '\U00010000'..'\U000EFFFF' => true,
            _ => false,
        }
    }

    fn is_name_char(&self) -> bool {
        if self.is_name_start_char() { return true; }
        match *self {
            '-'                |
            '.'                |
            '0'..'9'           |
            '\u00B7'           |
            '\u0300'..'\u036F' |
            '\u203F'..'\u2040' => true,
            _ => false
        }
    }

    fn is_space_char(&self) -> bool {
        match *self {
            '\x20' |
            '\x09' |
            '\x0D' |
            '\x0A' => true,
            _ => false,
        }
    }
}

trait Hax {
    fn first_child(&self) -> Option<RootChild>;
}

impl Hax for Root {
    fn first_child(&self) -> Option<RootChild> {
        self.children().remove(0)
    }
}

#[test]
fn parses_a_document_with_a_single_element() {
    let parser = Parser::new();
    let doc = parser.parse("<?xml version='1.0' ?><hello />");
    let top = doc.root().first_child().unwrap().element().unwrap();

    assert_eq!(top.name().as_slice(), "hello");
}
<commit_msg>Parse an element with a single attribute<commit_after>use super::{Document,Root,RootChild,Element};

struct Parser;

struct ParsedAttribute<'a> {
    name: &'a str,
    value: &'a str,
}

impl Parser {
    fn new() -> Parser {
        Parser
    }

    fn parse_preamble<'a>(&self, xml: &'a str) -> &'a str {
        // Parse the preamble
        let idx = xml.find_str("?>").expect("No preamble end");
        let end_of_preamble = idx + "?>".len();
        xml.slice_from(end_of_preamble)
    }

    fn optional_space<'a>(&self, xml: &'a str) -> &'a str {
        match xml.slice_space() {
            Some((_, next_xml)) => next_xml,
            None => xml,
        }
    }

    fn parse_attribute<'a>(&self, xml: &'a str) -> Option<(ParsedAttribute<'a>, &'a str)> {
        let (name, xml) = match xml.slice_name() {
            Some(x) => x,
            None => return None,
        };

        let xml = self.optional_space(xml);
        let (_, xml) = xml.slice_literal("=").expect("No equal sign");
        let xml = self.optional_space(xml);

        let (_, xml) = xml.slice_literal("'").expect("No quote"); // single and double quotes
        let (value, xml) = xml.slice_until("'").expect("No value"); // single and double quotes
        let (_, xml) = xml.slice_literal("'").expect("No quote"); // single and double quotes

        Some((ParsedAttribute{name: name, value: value}, xml))
    }

    fn parse_element<'a>(&self, doc: Document, xml: &'a str) -> (Element, &'a str) {
        let (_, after_start_brace) = xml.slice_literal("<").expect("no start brace");

        let (name, after_name) = after_start_brace.slice_name().expect("failed to parse a name!");

        let (_, after_space) = after_name.slice_space().expect("no space after name");

        let mut attrs = Vec::new();
        let after_attr = match self.parse_attribute(after_space) {
            None => after_space,
            Some((attr, after_attr)) => {
                attrs.push(attr);
                after_attr
            },
        };

        let after_attr = self.optional_space(after_attr);

        let (_, after_end_brace) = after_attr.slice_literal("/>").expect("no end brace");

        let e = doc.new_element(name.to_string());
        for attr in attrs.iter() {
            e.set_attribute(attr.name.to_string(), attr.value.to_string());
        }

        (e, after_end_brace)
    }

    fn parse(&self, xml: &str) -> Document {
        let doc = Document::new();

        let after_preamble = self.parse_preamble(xml);

        let (element, _tail) = self.parse_element(doc.clone(), after_preamble);
        doc.root().append_child(element);

        doc
    }
}

trait XmlStr<'a> {
    fn slice_at(&self, position: uint) -> (&'a str, &'a str);
    fn slice_until(&self, s: &str) -> Option<(&'a str, &'a str)>;
    fn slice_literal(&self, expected: &str) -> Option<(&'a str, &'a str)>;
    fn slice_start_rest(&self, is_first: |char| -> bool, is_rest: |char| -> bool) -> Option<(&'a str, &'a str)>;
    fn slice_name(&self) -> Option<(&'a str, &'a str)>;
    fn slice_space(&self) -> Option<(&'a str, &'a str)>;
}

impl<'a> XmlStr<'a> for &'a str {
    fn slice_at(&self, position: uint) -> (&'a str, &'a str) {
        (self.slice_to(position), self.slice_from(position))
    }

    fn slice_until(&self, s: &str) -> Option<(&'a str, &'a str)> {
        match self.find_str(s) {
            Some(position) => Some(self.slice_at(position)),
            None => None
        }
    }

    fn slice_literal(&self, expected: &str) -> Option<(&'a str, &'a str)> {
        if self.starts_with(expected) {
            Some(self.slice_at(expected.len()))
        } else {
            None
        }
    }

    fn slice_start_rest(&self,
                        is_first: |char| -> bool,
                        is_rest: |char| -> bool)
                        -> Option<(&'a str, &'a str)>
    {
        let mut positions = self.char_indices();

        match positions.next() {
            Some((_, c)) if is_first(c) => (),
            Some((_, _)) => return None,
            None => return None,
        };

        let mut positions = positions.skip_while(|&(_, c)| is_rest(c));
        match positions.next() {
            Some((offset, _)) => Some(self.slice_at(offset)),
            None => Some((self.clone(), "")),
        }
    }

    fn slice_name(&self) -> Option<(&'a str, &'a str)> {
        self.slice_start_rest(|c| c.is_name_start_char(), |c| c.is_name_char())
    }

    fn slice_space(&self) -> Option<(&'a str, &'a str)> {
        self.slice_start_rest(|c| c.is_space_char(), |c| c.is_space_char())
    }
}

trait XmlChar {
    fn is_name_start_char(&self) -> bool;
    fn is_name_char(&self) -> bool;
    fn is_space_char(&self) -> bool;
}

impl XmlChar for char {
    fn is_name_start_char(&self) -> bool {
        match *self {
            ':'                        |
            'A'..'Z'                   |
            '_'                        |
            'a'..'z'                   |
            '\U000000C0'..'\U000000D6' |
            '\U000000D8'..'\U000000F6' |
            '\U000000F8'..'\U000002FF' |
            '\U00000370'..'\U0000037D' |
            '\U0000037F'..'\U00001FFF' |
            '\U0000200C'..'\U0000200D' |
            '\U00002070'..'\U0000218F' |
            '\U00002C00'..'\U00002FEF' |
            '\U00003001'..'\U0000D7FF' |
            '\U0000F900'..'\U0000FDCF' |
            '\U0000FDF0'..'\U0000FFFD' |
            '\U00010000'..'\U000EFFFF' => true,
            _ => false,
        }
    }

    fn is_name_char(&self) -> bool {
        if self.is_name_start_char() { return true; }
        match *self {
            '-'                |
            '.'                |
            '0'..'9'           |
            '\u00B7'           |
            '\u0300'..'\u036F' |
            '\u203F'..'\u2040' => true,
            _ => false
        }
    }

    fn is_space_char(&self) -> bool {
        match *self {
            '\x20' |
            '\x09' |
            '\x0D' |
            '\x0A' => true,
            _ => false,
        }
    }
}

trait Hax {
    fn first_child(&self) -> Option<RootChild>;
}

impl Hax for Root {
    fn first_child(&self) -> Option<RootChild> {
        self.children().remove(0)
    }
}

#[test]
fn parses_a_document_with_a_single_element() {
    let parser = Parser::new();
    let doc = parser.parse("<?xml version='1.0' ?><hello />");
    let top = doc.root().first_child().unwrap().element().unwrap();

    assert_eq!(top.name().as_slice(), "hello");
}

#[test]
fn parses_an_element_with_an_attribute() {
    let parser = Parser::new();
    let doc = parser.parse("<?xml version='1.0' ?><hello scope='world'/>");
    let top = doc.root().first_child().unwrap().element().unwrap();

    assert_eq!(top.get_attribute("scope").unwrap().as_slice(), "world");
}
<|endoftext|>


<commit_before>#![deny(warnings)]

extern crate raw_cpuid;
extern crate syscall;

use std::env;
use std::fs::File;
use std::io::{self, Read, Write};

// std::fmt::Write conflicts with std::io::Write, hence the alias
use std::fmt::Write as FmtWrite;

const SECONDS_PER_MINUTE: i64 = 60;
const SECONDS_PER_HOUR: i64 = 3600;
const SECONDS_PER_DAY: i64 = 86400;

fn main() {
    let user = env::var("USER").unwrap_or(String::new());
    let mut hostname = String::new();
    if let Ok(mut file) = File::open("/etc/hostname") {
        let _ = file.read_to_string(&mut hostname);
    }

    let uptime;
    let uptime_secs;
    let uptime_mins;
    let uptime_hours;
    let uptime_days;

    let mut uptime_str = String::new();

    let mut ts = syscall::TimeSpec::default();
    if syscall::clock_gettime(syscall::CLOCK_MONOTONIC, &mut ts).is_ok() {
        uptime = ts.tv_sec;
        uptime_secs = uptime % 60;
        uptime_mins = (uptime / SECONDS_PER_MINUTE) % 60;
        uptime_hours = (uptime / SECONDS_PER_HOUR) % 24;
        uptime_days = (uptime / SECONDS_PER_DAY) % 365;

        let fmt_result;
        if uptime_days > 0 {
            fmt_result = write!(&mut uptime_str, "{}d {}h {}m {}s", uptime_days,
                                uptime_hours, uptime_mins, uptime_secs);
        } else if uptime_hours > 0 {
            fmt_result = write!(&mut uptime_str, "{}h {}m {}s", uptime_hours,
                                uptime_mins, uptime_secs);
        } else if uptime_mins > 0 {
            fmt_result = write!(&mut uptime_str, "{}m {}s", uptime_mins,
                                uptime_secs);
        } else {
            fmt_result = write!(&mut uptime_str, "{}s", uptime_secs);
        }

        if let Err(_) = fmt_result {
            // We probably don't want to panic! on this error
            println!("error: couldn't parse uptime");
        }
    }

    let mut width = 0;
    let mut height = 0;
    if let Ok(display_name) = env::var("DISPLAY") {
        if let Ok(display) = syscall::open(&display_name, syscall::O_STAT) {
            let mut buf: [u8; 4096] = [0; 4096];
            if let Ok(count) = syscall::fpath(display, &mut buf) {
                let path = unsafe { String::from_utf8_unchecked(Vec::from(&buf[..count])) };
                let res = path.split(":").nth(1).unwrap_or("");
                width = res.split("/").nth(1).unwrap_or("").parse::<i32>().unwrap_or(0);
                height = res.split("/").nth(2).unwrap_or("").parse::<i32>().unwrap_or(0);
            }
            let _ = syscall::close(display);
        }
    }

    let mut cpu = String::new();
    {
        let cpuid = raw_cpuid::CpuId::new();
        if let Some(info) = cpuid.get_extended_function_info() {
            if let Some(brand) = info.processor_brand_string() {
                cpu = brand.trim().to_string();
            }
        }
    }

    let mut ram = String::new();
    {
        let mut stat = syscall::StatVfs::default();
        if let Ok(fd) = syscall::open("memory:", syscall::O_STAT) {
            if syscall::fstatvfs(fd, &mut stat).is_ok() {
                let size = stat.f_blocks * stat.f_bsize as u64;
                let used = (stat.f_blocks - stat.f_bfree) * stat.f_bsize as u64;

                ram = format!("{}MB / {}MB", (used + 1048575)/1048576, (size + 1048575)/1048576);
            }
            let _ = syscall::close(fd);
        }
    }

    let left = [
        "              :+yMMMMy+:              ",
        "         .+dddNMMMMMMMMNddd+.         ",
        "       sydMMMMo/sMMMMs/oMMMMdys       ",
        "    .oMMMdso              osdMMMo.    ",
        "   .+MMd/`   -:::::::`      `/dMM+.   ",
        "  +dMMN.     NMMNNNMMdyo`     .NMMd+  ",
        "  yMMN       NM+   oomMMN       NMMy  ",
        " hNMd.       NM+     `oMN       .dMNh ",
        " dMMh        NM+     `oMN        hMMd ",
        " dMMh        NM+-oooodMMN        hMMd ",
        " dMMh        NM+/MMMMdhs`        hMMd ",
        " hNMd.       NM+`/mMMm+         .dMNh ",
        "  yMMN       NM+   oNMd+        NMMy  ",
        "  +dMMN.     NM+    omMMm     .NMMd+  ",
        "   .+MMd/`   --.     .---   `/dMM+.   ",
        "    .oMMMdso              osdMMMo.    ",
        "       sydMMMMo////////oMMMMdys       ",
        "         .+dddNMMMMMMMMNddd+.         ",
        "              :++++++++:              ",
    ];

    let right = [
        format!("\x1B[1;38;5;75m{}\x1B[0m@\x1B[1;38;5;75m{}\x1B[0m", user, hostname.trim()),
        format!("\x1B[1;38;5;75mOS:\x1B[0m redox-os"),
        format!("\x1B[1;38;5;75mKernel:\x1B[0m redox"),
        format!("\x1B[1;38;5;75mUptime:\x1B[0m {}", uptime_str),
        format!("\x1B[1;38;5;75mShell:\x1B[0m ion"),
        format!("\x1B[1;38;5;75mResolution:\x1B[0m {}x{}", width, height),
        format!("\x1B[1;38;5;75mDE:\x1B[0m orbital"),
        format!("\x1B[1;38;5;75mWM:\x1B[0m orbital"),
        format!("\x1B[1;38;5;75mFont:\x1B[0m unifont"),
        format!("\x1B[1;38;5;75mCPU:\x1B[0m {}", cpu),
        format!("\x1B[1;38;5;75mRAM:\x1B[0m {}", ram)
    ];

    let mut string = String::new();
    for i in 0..left.len() {
        string.push_str("\x1B[1;38;5;75m");
        string.push_str(left[i]);
        string.push_str("  \x1B[0m");
        if let Some(r) = right.get(i) {
            string.push_str(r);
        }
        string.push_str("\n");
    }

    io::stdout().write(string.as_bytes()).unwrap();
}
<commit_msg>Remove unnecessary variable declarations<commit_after>#![deny(warnings)]

extern crate raw_cpuid;
extern crate syscall;

use std::env;
use std::fs::File;
use std::io::{self, Read, Write};

// std::fmt::Write conflicts with std::io::Write, hence the alias
use std::fmt::Write as FmtWrite;

const SECONDS_PER_MINUTE: i64 = 60;
const SECONDS_PER_HOUR: i64 = 3600;
const SECONDS_PER_DAY: i64 = 86400;

fn main() {
    let user = env::var("USER").unwrap_or(String::new());
    let mut hostname = String::new();
    if let Ok(mut file) = File::open("/etc/hostname") {
        let _ = file.read_to_string(&mut hostname);
    }

    let mut uptime_str = String::new();

    let mut ts = syscall::TimeSpec::default();
    if syscall::clock_gettime(syscall::CLOCK_MONOTONIC, &mut ts).is_ok() {
        let uptime = ts.tv_sec;
        let uptime_secs = uptime % 60;
        let uptime_mins = (uptime / SECONDS_PER_MINUTE) % 60;
        let uptime_hours = (uptime / SECONDS_PER_HOUR) % 24;
        let uptime_days = (uptime / SECONDS_PER_DAY) % 365;

        let fmt_result;
        if uptime_days > 0 {
            fmt_result = write!(&mut uptime_str, "{}d {}h {}m {}s", uptime_days,
                                uptime_hours, uptime_mins, uptime_secs);
        } else if uptime_hours > 0 {
            fmt_result = write!(&mut uptime_str, "{}h {}m {}s", uptime_hours,
                                uptime_mins, uptime_secs);
        } else if uptime_mins > 0 {
            fmt_result = write!(&mut uptime_str, "{}m {}s", uptime_mins,
                                uptime_secs);
        } else {
            fmt_result = write!(&mut uptime_str, "{}s", uptime_secs);
        }

        if let Err(_) = fmt_result {
            // We probably don't want to panic! on this error
            println!("error: couldn't parse uptime");
        }
    }

    let mut width = 0;
    let mut height = 0;
    if let Ok(display_name) = env::var("DISPLAY") {
        if let Ok(display) = syscall::open(&display_name, syscall::O_STAT) {
            let mut buf: [u8; 4096] = [0; 4096];
            if let Ok(count) = syscall::fpath(display, &mut buf) {
                let path = unsafe { String::from_utf8_unchecked(Vec::from(&buf[..count])) };
                let res = path.split(":").nth(1).unwrap_or("");
                width = res.split("/").nth(1).unwrap_or("").parse::<i32>().unwrap_or(0);
                height = res.split("/").nth(2).unwrap_or("").parse::<i32>().unwrap_or(0);
            }
            let _ = syscall::close(display);
        }
    }

    let mut cpu = String::new();
    {
        let cpuid = raw_cpuid::CpuId::new();
        if let Some(info) = cpuid.get_extended_function_info() {
            if let Some(brand) = info.processor_brand_string() {
                cpu = brand.trim().to_string();
            }
        }
    }

    let mut ram = String::new();
    {
        let mut stat = syscall::StatVfs::default();
        if let Ok(fd) = syscall::open("memory:", syscall::O_STAT) {
            if syscall::fstatvfs(fd, &mut stat).is_ok() {
                let size = stat.f_blocks * stat.f_bsize as u64;
                let used = (stat.f_blocks - stat.f_bfree) * stat.f_bsize as u64;

                ram = format!("{}MB / {}MB", (used + 1048575)/1048576, (size + 1048575)/1048576);
            }
            let _ = syscall::close(fd);
        }
    }

    let left = [
        "              :+yMMMMy+:              ",
        "         .+dddNMMMMMMMMNddd+.         ",
        "       sydMMMMo/sMMMMs/oMMMMdys       ",
        "    .oMMMdso              osdMMMo.    ",
        "   .+MMd/`   -:::::::`      `/dMM+.   ",
        "  +dMMN.     NMMNNNMMdyo`     .NMMd+  ",
        "  yMMN       NM+   oomMMN       NMMy  ",
        " hNMd.       NM+     `oMN       .dMNh ",
        " dMMh        NM+     `oMN        hMMd ",
        " dMMh        NM+-oooodMMN        hMMd ",
        " dMMh        NM+/MMMMdhs`        hMMd ",
        " hNMd.       NM+`/mMMm+         .dMNh ",
        "  yMMN       NM+   oNMd+        NMMy  ",
        "  +dMMN.     NM+    omMMm     .NMMd+  ",
        "   .+MMd/`   --.     .---   `/dMM+.   ",
        "    .oMMMdso              osdMMMo.    ",
        "       sydMMMMo////////oMMMMdys       ",
        "         .+dddNMMMMMMMMNddd+.         ",
        "              :++++++++:              ",
    ];

    let right = [
        format!("\x1B[1;38;5;75m{}\x1B[0m@\x1B[1;38;5;75m{}\x1B[0m", user, hostname.trim()),
        format!("\x1B[1;38;5;75mOS:\x1B[0m redox-os"),
        format!("\x1B[1;38;5;75mKernel:\x1B[0m redox"),
        format!("\x1B[1;38;5;75mUptime:\x1B[0m {}", uptime_str),
        format!("\x1B[1;38;5;75mShell:\x1B[0m ion"),
        format!("\x1B[1;38;5;75mResolution:\x1B[0m {}x{}", width, height),
        format!("\x1B[1;38;5;75mDE:\x1B[0m orbital"),
        format!("\x1B[1;38;5;75mWM:\x1B[0m orbital"),
        format!("\x1B[1;38;5;75mFont:\x1B[0m unifont"),
        format!("\x1B[1;38;5;75mCPU:\x1B[0m {}", cpu),
        format!("\x1B[1;38;5;75mRAM:\x1B[0m {}", ram)
    ];

    let mut string = String::new();
    for i in 0..left.len() {
        string.push_str("\x1B[1;38;5;75m");
        string.push_str(left[i]);
        string.push_str("  \x1B[0m");
        if let Some(r) = right.get(i) {
            string.push_str(r);
        }
        string.push_str("\n");
    }

    io::stdout().write(string.as_bytes()).unwrap();
}
<|endoftext|>