<commit_before>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

use std::process::exit;

use chrono::naive::NaiveDateTime;

use libimagdiary::diary::Diary;
use libimagdiary::diaryid::DiaryId;
use libimagdiary::error::DiaryErrorKind as DEK;
use libimagdiary::error::MapErrInto;
use libimagentryedit::edit::Edit;
use libimagrt::runtime::Runtime;
use libimagerror::trace::MapErrTrace;
use libimagerror::into::IntoError;
use libimagtimeui::datetime::DateTime;
use libimagtimeui::parse::Parse;
use libimagutil::warn_exit::warn_exit;
use libimagstore::storeid::IntoStoreId;
use libimagerror::trace::trace_error_exit;

use util::get_diary_name;

pub fn edit(rt: &Runtime) {
    let diaryname = get_diary_name(rt).unwrap_or_else(|| warn_exit("No diary name", 1));

    let datetime : Option<NaiveDateTime> = rt
        .cli()
        .subcommand_matches("edit")
        .unwrap()
        .value_of("datetime")
        .and_then(DateTime::parse)
        .map(|dt| dt.into());

    let to_edit = match datetime {
        Some(dt) => {
            let id = match DiaryId::from_datetime(diaryname.clone(), dt).into_storeid() {
                Ok(id) => id,
                Err(e) => trace_error_exit(&e, 1),
            };
            rt.store().get(id)
        },
        None     => match rt.store().get_youngest_entry_id(&diaryname) {
            Some(Ok(id)) => match id.into_storeid() {
                Ok(sid) => rt.store().get(sid),
                Err(e) => trace_error_exit(&e, 1),
            },
            Some(Err(e)) => trace_error_exit(&e, 1),
            None => {
                error!("No entries in diary. Aborting");
                exit(1)
            }
        }
    };

    to_edit.map(|opte| match opte {
        Some(mut e) => e.edit_content(rt).map_err_into(DEK::IOError),
        None        => Err(DEK::EntryNotInDiary.into_error()),
    })
    .map_err_trace().ok();
}

<commit_msg>Refactor, minify edit() impl<commit_after>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

use std::process::exit;

use chrono::naive::NaiveDateTime;

use libimagdiary::diary::Diary;
use libimagdiary::diaryid::DiaryId;
use libimagdiary::error::DiaryErrorKind as DEK;
use libimagdiary::error::MapErrInto;
use libimagentryedit::edit::Edit;
use libimagrt::runtime::Runtime;
use libimagerror::trace::MapErrTrace;
use libimagerror::into::IntoError;
use libimagtimeui::datetime::DateTime;
use libimagtimeui::parse::Parse;
use libimagutil::warn_exit::warn_exit;
use libimagerror::trace::trace_error_exit;

use util::get_diary_name;

pub fn edit(rt: &Runtime) {
    let diaryname = get_diary_name(rt).unwrap_or_else(|| warn_exit("No diary name", 1));

    rt.cli()
        .subcommand_matches("edit")
        .unwrap()
        .value_of("datetime")
        .and_then(DateTime::parse)
        .map(|dt| dt.into())
        .map(|dt: NaiveDateTime| DiaryId::from_datetime(diaryname.clone(), dt))
        .or_else(|| {
            rt.store()
                .get_youngest_entry_id(&diaryname)
                .map(|optid| match optid {
                    Ok(id) => id,
                    Err(e) => trace_error_exit(&e, 1),
                })
        })
        .ok_or_else(|| {
            error!("No entries in diary. Aborting");
            exit(1)
        })
        .and_then(|id| rt.store().get(id))
        .map(|opte| match opte {
            Some(mut e) => e.edit_content(rt).map_err_into(DEK::IOError),
            None        => Err(DEK::EntryNotInDiary.into_error()),
        })
        .map_err_trace()
        .ok();
}

<|endoftext|>


<commit_before>/*!

# Supported types

| Rust Type               | JSON Serialization     | Notes                                     |
|-------------------------|------------------------|-------------------------------------------|
| `DateTime<FixedOffset>` | RFC3339 string         |                                           |
| `DateTime<Utc>`         | RFC3339 string         |                                           |
| `NaiveDate`             | RFC3339 string         |                                           |
| `NaiveDateTime`         | float (unix timestamp) | JSON numbers (i.e. IEEE doubles) are not  |
|                         |                        | precise enough for nanoseconds.           |
|                         |                        | Values will be truncated to microsecond   |
|                         |                        | resolution.                               |

*/
use chrono::prelude::*;

use Value;

#[doc(hidden)]
pub static RFC3339_FORMAT: &'static str = "%Y-%m-%dT%H:%M:%S%.f%:z";
static RFC3339_PARSE_FORMAT: &'static str = "%+";

graphql_scalar!(DateTime<FixedOffset> as "DateTimeFixedOffset" {
    description: "DateTime"

    resolve(&self) -> Value {
        Value::string(self.to_rfc3339())
    }

    from_input_value(v: &InputValue) -> Option<DateTime<FixedOffset>> {
        v.as_string_value()
         .and_then(|s| DateTime::parse_from_rfc3339(s).ok())
    }
});

graphql_scalar!(DateTime<Utc> as "DateTimeUtc" {
    description: "DateTime"

    resolve(&self) -> Value {
        Value::string(self.to_rfc3339())
    }

    from_input_value(v: &InputValue) -> Option<DateTime<Utc>> {
        v.as_string_value()
         .and_then(|s| (s.parse::<DateTime<Utc>>().ok()))
    }
});

// Don't use `Date` as the docs say:
// "[Date] should be considered ambiguous at best, due to the "
// inherent lack of precision required for the time zone resolution.
// For serialization and deserialization uses, it is best to use
// `NaiveDate` instead."
graphql_scalar!(NaiveDate {
    description: "NaiveDate"

    resolve(&self) -> Value {
        Value::string(self.format(RFC3339_FORMAT).to_string())
    }

    from_input_value(v: &InputValue) -> Option<NaiveDate> {
        v.as_string_value()
         .and_then(|s| NaiveDate::parse_from_str(s, RFC3339_PARSE_FORMAT).ok())
    }
});

/// JSON numbers (i.e. IEEE doubles) are not precise enough for nanosecond
/// datetimes. Values will be truncated to microsecond resolution.
graphql_scalar!(NaiveDateTime {
    description: "NaiveDateTime"

    resolve(&self) -> Value {
        Value::float(self.timestamp() as f64)
    }

    from_input_value(v: &InputValue) -> Option<NaiveDateTime> {
        v.as_float_value()
         .and_then(|f| NaiveDateTime::from_timestamp_opt(f as i64, 0))
    }
});

#[cfg(test)]
mod test {
    use chrono::prelude::*;
    use super::RFC3339_PARSE_FORMAT;

    fn datetime_fixedoffset_test(raw: &'static str) {
        let input = ::InputValue::String(raw.to_string());

        let parsed: DateTime<FixedOffset> = ::FromInputValue::from_input_value(&input).unwrap();
        let expected = DateTime::parse_from_rfc3339(raw).unwrap();

        assert_eq!(parsed, expected);
    }

    #[test]
    fn datetime_fixedoffset_from_input_value() {
        datetime_fixedoffset_test("2014-11-28T21:00:09+09:00");
    }

    #[test]
    fn datetime_fixedoffset_from_input_value_with_z_timezone() {
        datetime_fixedoffset_test("2014-11-28T21:00:09Z");
    }

    #[test]
    fn datetime_fixedoffset_from_input_value_with_fractional_seconds() {
        datetime_fixedoffset_test("2014-11-28T21:00:09.05+09:00");
    }

    fn datetime_utc_test(raw: &'static str) {
        let input = ::InputValue::String(raw.to_string());

        let parsed: DateTime<Utc> = ::FromInputValue::from_input_value(&input).unwrap();
        let expected = DateTime::parse_from_rfc3339(raw)
            .unwrap()
            .with_timezone(&Utc);

        assert_eq!(parsed, expected);
    }

    #[test]
    fn datetime_utc_from_input_value() {
        datetime_utc_test("2014-11-28T21:00:09+09:00")
    }

    #[test]
    fn datetime_utc_from_input_value_with_z_timezone() {
        datetime_utc_test("2014-11-28T21:00:09Z")
    }

    #[test]
    fn datetime_utc_from_input_value_with_fractional_seconds() {
        datetime_utc_test("2014-11-28T21:00:09.005+09:00");
    }

    fn naivedate_test(raw: &'static str, y: i32, m: u32, d: u32) {
        let input = ::InputValue::String(raw.to_string());

        let parsed: NaiveDate = ::FromInputValue::from_input_value(&input).unwrap();
        let expected = NaiveDate::parse_from_str(raw, &RFC3339_PARSE_FORMAT).unwrap();
        let expected_via_datetime = DateTime::parse_from_rfc3339(raw)
            .unwrap()
            .date()
            .naive_utc();
        let expected_via_ymd = NaiveDate::from_ymd(y, m, d);

        assert_eq!(parsed, expected);
        assert_eq!(parsed, expected_via_datetime);
        assert_eq!(parsed, expected_via_ymd);

        assert_eq!(parsed.year(), y);
        assert_eq!(parsed.month(), m);
        assert_eq!(parsed.day(), d);
    }

    #[test]
    fn naivedate_from_input_value() {
        naivedate_test("1996-12-19T16:39:57-08:00", 1996, 12, 19);
    }

    #[test]
    fn naivedate_from_input_value_with_fractional_seconds() {
        naivedate_test("1996-12-19T16:39:57.005-08:00", 1996, 12, 19);
    }

    #[test]
    fn naivedate_from_input_value_with_z_timezone() {
        naivedate_test("1996-12-19T16:39:57Z", 1996, 12, 19);
    }

    #[test]
    fn naivedatetime_from_input_value() {
        let raw = 1_000_000_000_f64;
        let input = ::InputValue::Float(raw);

        let parsed: NaiveDateTime = ::FromInputValue::from_input_value(&input).unwrap();
        let expected = NaiveDateTime::from_timestamp_opt(raw as i64, 0).unwrap();

        assert_eq!(parsed, expected);
        assert_eq!(raw, expected.timestamp() as f64);
    }
}
<commit_msg>Change serialization of Chronos NaiveDate (#151)<commit_after>/*!

# Supported types

| Rust Type               | JSON Serialization     | Notes                                     |
|-------------------------|------------------------|-------------------------------------------|
| `DateTime<FixedOffset>` | RFC3339 string         |                                           |
| `DateTime<Utc>`         | RFC3339 string         |                                           |
| `NaiveDate`             | YYYY-MM-DD             |                                           |
| `NaiveDateTime`         | float (unix timestamp) | JSON numbers (i.e. IEEE doubles) are not  |
|                         |                        | precise enough for nanoseconds.           |
|                         |                        | Values will be truncated to microsecond   |
|                         |                        | resolution.                               |

*/
use chrono::prelude::*;

use Value;

#[doc(hidden)]
pub static RFC3339_FORMAT: &'static str = "%Y-%m-%dT%H:%M:%S%.f%:z";

graphql_scalar!(DateTime<FixedOffset> as "DateTimeFixedOffset" {
    description: "DateTime"

    resolve(&self) -> Value {
        Value::string(self.to_rfc3339())
    }

    from_input_value(v: &InputValue) -> Option<DateTime<FixedOffset>> {
        v.as_string_value()
         .and_then(|s| DateTime::parse_from_rfc3339(s).ok())
    }
});

graphql_scalar!(DateTime<Utc> as "DateTimeUtc" {
    description: "DateTime"

    resolve(&self) -> Value {
        Value::string(self.to_rfc3339())
    }

    from_input_value(v: &InputValue) -> Option<DateTime<Utc>> {
        v.as_string_value()
         .and_then(|s| (s.parse::<DateTime<Utc>>().ok()))
    }
});

// Don't use `Date` as the docs say:
// "[Date] should be considered ambiguous at best, due to the "
// inherent lack of precision required for the time zone resolution.
// For serialization and deserialization uses, it is best to use
// `NaiveDate` instead."
graphql_scalar!(NaiveDate {
    description: "NaiveDate"

    resolve(&self) -> Value {
        Value::string(self.format("%Y-%m-%d").to_string())
    }

    from_input_value(v: &InputValue) -> Option<NaiveDate> {
        v.as_string_value()
         .and_then(|s| NaiveDate::parse_from_str(s, "%Y-%m-%d").ok())
    }
});

/// JSON numbers (i.e. IEEE doubles) are not precise enough for nanosecond
/// datetimes. Values will be truncated to microsecond resolution.
graphql_scalar!(NaiveDateTime {
    description: "NaiveDateTime"

    resolve(&self) -> Value {
        Value::float(self.timestamp() as f64)
    }

    from_input_value(v: &InputValue) -> Option<NaiveDateTime> {
        v.as_float_value()
         .and_then(|f| NaiveDateTime::from_timestamp_opt(f as i64, 0))
    }
});

#[cfg(test)]
mod test {
    use chrono::prelude::*;

    fn datetime_fixedoffset_test(raw: &'static str) {
        let input = ::InputValue::String(raw.to_string());

        let parsed: DateTime<FixedOffset> = ::FromInputValue::from_input_value(&input).unwrap();
        let expected = DateTime::parse_from_rfc3339(raw).unwrap();

        assert_eq!(parsed, expected);
    }

    #[test]
    fn datetime_fixedoffset_from_input_value() {
        datetime_fixedoffset_test("2014-11-28T21:00:09+09:00");
    }

    #[test]
    fn datetime_fixedoffset_from_input_value_with_z_timezone() {
        datetime_fixedoffset_test("2014-11-28T21:00:09Z");
    }

    #[test]
    fn datetime_fixedoffset_from_input_value_with_fractional_seconds() {
        datetime_fixedoffset_test("2014-11-28T21:00:09.05+09:00");
    }

    fn datetime_utc_test(raw: &'static str) {
        let input = ::InputValue::String(raw.to_string());

        let parsed: DateTime<Utc> = ::FromInputValue::from_input_value(&input).unwrap();
        let expected = DateTime::parse_from_rfc3339(raw)
            .unwrap()
            .with_timezone(&Utc);

        assert_eq!(parsed, expected);
    }

    #[test]
    fn datetime_utc_from_input_value() {
        datetime_utc_test("2014-11-28T21:00:09+09:00")
    }

    #[test]
    fn datetime_utc_from_input_value_with_z_timezone() {
        datetime_utc_test("2014-11-28T21:00:09Z")
    }

    #[test]
    fn datetime_utc_from_input_value_with_fractional_seconds() {
        datetime_utc_test("2014-11-28T21:00:09.005+09:00");
    }

    #[test]
    fn naivedate_from_input_value() {
        let input = ::InputValue::String("1996-12-19".to_string());
        let y = 1996;
        let m = 12;
        let d = 19;

        let parsed: NaiveDate = ::FromInputValue::from_input_value(&input).unwrap();
        let expected = NaiveDate::from_ymd(y, m, d);

        assert_eq!(parsed, expected);

        assert_eq!(parsed.year(), y);
        assert_eq!(parsed.month(), m);
        assert_eq!(parsed.day(), d);
    }

    #[test]
    fn naivedatetime_from_input_value() {
        let raw = 1_000_000_000_f64;
        let input = ::InputValue::Float(raw);

        let parsed: NaiveDateTime = ::FromInputValue::from_input_value(&input).unwrap();
        let expected = NaiveDateTime::from_timestamp_opt(raw as i64, 0).unwrap();

        assert_eq!(parsed, expected);
        assert_eq!(raw, expected.timestamp() as f64);
    }
}

#[cfg(test)]
mod integration_test {
    use chrono::prelude::*;
    use chrono::Utc;

    use executor::Variables;
    use value::Value;
    use schema::model::RootNode;
    use types::scalars::EmptyMutation;

    #[test]
    fn test_serialization() {
        struct Root {}
        graphql_object!(Root: () |&self| {
            field exampleNaiveDate() -> NaiveDate {
                NaiveDate::from_ymd(2015, 3, 14)
            }
            field exampleNaiveDateTime() -> NaiveDateTime {
                NaiveDate::from_ymd(2016, 7, 8).and_hms(9, 10, 11)
            }
            field exampleDateTimeFixedOffset() -> DateTime<FixedOffset> {
              DateTime::parse_from_rfc3339("1996-12-19T16:39:57-08:00").unwrap()
            }
            field exampleDateTimeUtc() -> DateTime<Utc> {
              Utc.timestamp(61, 0)
            }
        });

        let doc = r#"
        {
            exampleNaiveDate,
            exampleNaiveDateTime,
            exampleDateTimeFixedOffset,
            exampleDateTimeUtc,
        }
        "#;

        let schema = RootNode::new(Root {}, EmptyMutation::<()>::new());

        let (result, errs) =
            ::execute(doc, None, &schema, &Variables::new(), &()).expect("Execution failed");

        assert_eq!(errs, []);

        assert_eq!(
            result,
            Value::object(
                vec![
                    ("exampleNaiveDate", Value::string("2015-03-14")),
                    ("exampleNaiveDateTime", Value::float(1467969011.0)),
                    (
                        "exampleDateTimeFixedOffset",
                        Value::string("1996-12-19T16:39:57-08:00"),
                    ),
                    (
                        "exampleDateTimeUtc",
                        Value::string("1970-01-01T00:01:01+00:00"),
                    ),
                ].into_iter()
                    .collect()
            )
        );
    }
}
<|endoftext|>


<commit_before>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

#![deny(warnings)]
#![deny(unused)]
#![deny(unsafe_code)]

use std::convert::TryFrom;
use std::env;
use std::fs;
use std::mem::size_of;
use std::path::PathBuf;

use mundane::hash::Sha256;
use mundane::public::rsa::RsaPkcs1v15;
use mundane::public::rsa::B3072;
use mundane::public::DerPrivateKey;
use mundane::public::Signature;

use anyhow::bail;
use anyhow::ensure;
use anyhow::Context;
use anyhow::Result;

use zerocopy::AsBytes;

use rom_ext_image::image::Image;
use rom_ext_image::image::ManifestBuffer;
use rom_ext_image::manifest;
use rom_ext_image::manifest::Manifest;

use object::read::elf::ElfFile32;
use object::read::ObjectSection;
use object::Object;

// Type aliases for convenience.
type ImageSignature =
    mundane::public::rsa::RsaSignature<B3072, RsaPkcs1v15, Sha256>;
type PrivateKey = mundane::public::rsa::RsaPrivKey<B3072>;

// TODO: Remove this struct when this functionality is integrated into `opentitantool`.
struct Args {
    input_image: PathBuf,
    priv_key: PathBuf,
    elf_file: PathBuf,
    output_image: PathBuf,
}

impl Args {
    pub fn new(args: env::ArgsOs) -> Result<Args> {
        let args = args.skip(1).collect::<Vec<_>>();
        match args.as_slice() {
            [input_image, priv_key, elf_file, output_image] => Ok(Args {
                    input_image: PathBuf::from(input_image),
                    priv_key: PathBuf::from(priv_key),
                    elf_file: PathBuf::from(elf_file),
                    output_image: PathBuf::from(output_image),
                }),
            args => bail!("Expected exactly 4 positional arguments: input image, private key, elf file, and output image, got: {}.", args.len()),
        }
    }
}

// FIXME: Keeping for now, can be removed if not used in opentitantool.
/// Parses an unsigned big-endian hex string into a little-endian byte vector.
fn _parse_hex_str(hex_str: &str) -> Result<Vec<u8>> {
    ensure!(
        hex_str.starts_with("0x")
            && hex_str.len() > 2
            && hex_str.len() % 2 == 0,
        "Invalid hex string: {}",
        hex_str
    );
    let bytes = (2..hex_str.len())
        .step_by(2)
        .map(|i| u8::from_str_radix(&hex_str[i..i + 2], 16))
        .rev()
        .collect::<Result<_, _>>()?;
    Ok(bytes)
}

/// Updates the manifest of an image.
// TODO: This function must use a public key.
fn update_image_manifest(
    image: &mut Image,
    key: &PrivateKey,
    elf: &ElfFile32,
) -> Result<()> {
    let manifest_addr = u32::try_from(
        elf.section_by_name(".manifest")
            .context("Could not find the `.manifest` section.")?
            .address(),
    )?;

    *image.manifest = Manifest {
        identifier: 0x4552544f,
        image_length: u32::try_from(image.bytes().len())?,
        entry_point: {
            let addr = u32::try_from(elf.entry())?;
            addr.checked_sub(manifest_addr).context("Overflow")?
        },
        code_start: {
            let addr = u32::try_from(
                elf.section_by_name(".vectors")
                    .context("Could not find the `.vectors` section.")?
                    .address(),
            )?;
            addr.checked_sub(manifest_addr).context("Overflow")?
        },
        code_end: {
            let text = elf
                .section_by_name(".text")
                .context("Could not find the `.text` section.")?;
            let addr = u32::try_from(
                text.address()
                    .checked_add(text.size())
                    .context("Overflow")?,
            )?;
            addr.checked_sub(manifest_addr).context("Overflow")?
        },
        ..Default::default()
    };

    let exponent_be = key.public_exponent_be();
    let dest = image.manifest.exponent.as_bytes_mut().iter_mut();
    let src = exponent_be.iter().rev().copied();
    ensure!(dest.len() >= src.len(), "Unexpected exponent length.");
    for (d, s) in Iterator::zip(dest, src) {
        *d = s;
    }

    let modulus_be = key.public_modulus_be();
    let dest = image.manifest.modulus.as_bytes_mut().iter_mut();
    let src = modulus_be.iter().rev().copied();
    ensure!(dest.len() == src.len(), "Unexpected modulus length.");
    for (d, s) in Iterator::zip(dest, src) {
        *d = s;
    }

    Ok(())
}

/// Calculates the signature for the signed portion of an image.
fn calculate_image_signature(
    image: &Image,
    private_key: &PrivateKey,
) -> Result<ImageSignature> {
    ImageSignature::sign(&private_key, &image.signed_bytes())
        .context("Failed to calculate image signature.")
}

/// Updates the signature of an image.
fn update_image_signature(
    image: &mut Image,
    sig: ImageSignature,
) -> Result<()> {
    let dest = image.manifest.signature.as_bytes_mut().iter_mut();
    let src = sig.bytes().iter().rev().copied();
    ensure!(dest.len() == src.len(), "Unexpected signature length.");
    for (d, s) in Iterator::zip(dest, src) {
        *d = s;
    }
    Ok(())
}

fn main() -> Result<()> {
    // TODO(#6915): Convert this to a unit test after we start running rust tests during our
    // builds.
    manifest::check_manifest_layout();
    let args = Args::new(env::args_os())?;

    // We use a separate buffer for manifest because it must have the same alignment as `Manifest`
    // to be able to use `LayoutVerified::new()` and the approach we use to ensure this requires
    // its size to be known at compile time.
    let payload = &fs::read(&args.input_image).with_context(|| {
        format!("Failed to read {}", args.input_image.display())
    })?[size_of::<Manifest>()..];
    let mut manifest_buffer = ManifestBuffer::new();
    let mut image = Image::new(&mut manifest_buffer, payload)?;

    let key = fs::read(&args.priv_key).with_context(|| {
        format!("Failed to read the key from `{}`.", args.priv_key.display())
    })?;
    let key =
        PrivateKey::parse_from_der(&key).context("Failed to parse the key.")?;

    let elf = fs::read(&args.elf_file)?;
    let elf = ElfFile32::parse(elf.as_slice())?;

    update_image_manifest(&mut image, &key, &elf)?;
    let sig = calculate_image_signature(&image, &key)?;
    update_image_signature(&mut image, sig)?;

    fs::write(&args.output_image, image.bytes()).with_context(|| {
        format!(
            "Failed to write the image to {}",
            args.output_image.display()
        )
    })?;
    Ok(())
}
<commit_msg>[rom_ext_image_tools] Check that code region is aligned correctly<commit_after>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

#![deny(warnings)]
#![deny(unused)]
#![deny(unsafe_code)]

use std::convert::TryFrom;
use std::env;
use std::fs;
use std::mem::size_of;
use std::path::PathBuf;

use mundane::hash::Sha256;
use mundane::public::rsa::RsaPkcs1v15;
use mundane::public::rsa::B3072;
use mundane::public::DerPrivateKey;
use mundane::public::Signature;

use anyhow::bail;
use anyhow::ensure;
use anyhow::Context;
use anyhow::Result;

use zerocopy::AsBytes;

use rom_ext_image::image::Image;
use rom_ext_image::image::ManifestBuffer;
use rom_ext_image::manifest;
use rom_ext_image::manifest::Manifest;

use object::read::elf::ElfFile32;
use object::read::ObjectSection;
use object::Object;

// Type aliases for convenience.
type ImageSignature =
    mundane::public::rsa::RsaSignature<B3072, RsaPkcs1v15, Sha256>;
type PrivateKey = mundane::public::rsa::RsaPrivKey<B3072>;

// TODO: Remove this struct when this functionality is integrated into `opentitantool`.
struct Args {
    input_image: PathBuf,
    priv_key: PathBuf,
    elf_file: PathBuf,
    output_image: PathBuf,
}

impl Args {
    pub fn new(args: env::ArgsOs) -> Result<Args> {
        let args = args.skip(1).collect::<Vec<_>>();
        match args.as_slice() {
            [input_image, priv_key, elf_file, output_image] => Ok(Args {
                    input_image: PathBuf::from(input_image),
                    priv_key: PathBuf::from(priv_key),
                    elf_file: PathBuf::from(elf_file),
                    output_image: PathBuf::from(output_image),
                }),
            args => bail!("Expected exactly 4 positional arguments: input image, private key, elf file, and output image, got: {}.", args.len()),
        }
    }
}

// FIXME: Keeping for now, can be removed if not used in opentitantool.
/// Parses an unsigned big-endian hex string into a little-endian byte vector.
fn _parse_hex_str(hex_str: &str) -> Result<Vec<u8>> {
    ensure!(
        hex_str.starts_with("0x")
            && hex_str.len() > 2
            && hex_str.len() % 2 == 0,
        "Invalid hex string: {}",
        hex_str
    );
    let bytes = (2..hex_str.len())
        .step_by(2)
        .map(|i| u8::from_str_radix(&hex_str[i..i + 2], 16))
        .rev()
        .collect::<Result<_, _>>()?;
    Ok(bytes)
}

/// Updates the manifest of an image.
// TODO: This function must use a public key.
fn update_image_manifest(
    image: &mut Image,
    key: &PrivateKey,
    elf: &ElfFile32,
) -> Result<()> {
    let manifest_addr = u32::try_from(
        elf.section_by_name(".manifest")
            .context("Could not find the `.manifest` section.")?
            .address(),
    )?;

    *image.manifest = Manifest {
        identifier: 0x4552544f,
        image_length: u32::try_from(image.bytes().len())?,
        entry_point: {
            let addr = u32::try_from(elf.entry())?;
            addr.checked_sub(manifest_addr).context("Overflow")?
        },
        code_start: {
            let addr = u32::try_from(
                elf.section_by_name(".vectors")
                    .context("Could not find the `.vectors` section.")?
                    .address(),
            )?;
            addr.checked_sub(manifest_addr).context("Overflow")?
        },
        code_end: {
            let text = elf
                .section_by_name(".text")
                .context("Could not find the `.text` section.")?;
            let addr = u32::try_from(
                text.address()
                    .checked_add(text.size())
                    .context("Overflow")?,
            )?;
            addr.checked_sub(manifest_addr).context("Overflow")?
        },
        ..Default::default()
    };

    // Code region must be 32-bit aligned due to PMP constraints.
    fn is_word_aligned(addr: u32) -> bool {
        return addr % 4 == 0;
    }
    ensure!(
        is_word_aligned(image.manifest.code_start),
        "Unaligned code start."
    );
    ensure!(
        is_word_aligned(image.manifest.code_end),
        "Unaligned code end."
    );

    let exponent_be = key.public_exponent_be();
    let dest = image.manifest.exponent.as_bytes_mut().iter_mut();
    let src = exponent_be.iter().rev().copied();
    ensure!(dest.len() >= src.len(), "Unexpected exponent length.");
    for (d, s) in Iterator::zip(dest, src) {
        *d = s;
    }

    let modulus_be = key.public_modulus_be();
    let dest = image.manifest.modulus.as_bytes_mut().iter_mut();
    let src = modulus_be.iter().rev().copied();
    ensure!(dest.len() == src.len(), "Unexpected modulus length.");
    for (d, s) in Iterator::zip(dest, src) {
        *d = s;
    }

    Ok(())
}

/// Calculates the signature for the signed portion of an image.
fn calculate_image_signature(
    image: &Image,
    private_key: &PrivateKey,
) -> Result<ImageSignature> {
    ImageSignature::sign(&private_key, &image.signed_bytes())
        .context("Failed to calculate image signature.")
}

/// Updates the signature of an image.
fn update_image_signature(
    image: &mut Image,
    sig: ImageSignature,
) -> Result<()> {
    let dest = image.manifest.signature.as_bytes_mut().iter_mut();
    let src = sig.bytes().iter().rev().copied();
    ensure!(dest.len() == src.len(), "Unexpected signature length.");
    for (d, s) in Iterator::zip(dest, src) {
        *d = s;
    }
    Ok(())
}

fn main() -> Result<()> {
    // TODO(#6915): Convert this to a unit test after we start running rust tests during our
    // builds.
    manifest::check_manifest_layout();
    let args = Args::new(env::args_os())?;

    // We use a separate buffer for manifest because it must have the same alignment as `Manifest`
    // to be able to use `LayoutVerified::new()` and the approach we use to ensure this requires
    // its size to be known at compile time.
    let payload = &fs::read(&args.input_image).with_context(|| {
        format!("Failed to read {}", args.input_image.display())
    })?[size_of::<Manifest>()..];
    let mut manifest_buffer = ManifestBuffer::new();
    let mut image = Image::new(&mut manifest_buffer, payload)?;

    let key = fs::read(&args.priv_key).with_context(|| {
        format!("Failed to read the key from `{}`.", args.priv_key.display())
    })?;
    let key =
        PrivateKey::parse_from_der(&key).context("Failed to parse the key.")?;

    let elf = fs::read(&args.elf_file)?;
    let elf = ElfFile32::parse(elf.as_slice())?;

    update_image_manifest(&mut image, &key, &elf)?;
    let sig = calculate_image_signature(&image, &key)?;
    update_image_signature(&mut image, sig)?;

    fs::write(&args.output_image, image.bytes()).with_context(|| {
        format!(
            "Failed to write the image to {}",
            args.output_image.display()
        )
    })?;
    Ok(())
}
<|endoftext|>


<commit_before>use std::io::timer;
use std::io::stdio;

struct Shell {
    cmd_prompt: ~str,
}

impl Shell {
    fn new(prompt_str: &str) -> Shell {
        Shell {
            cmd_prompt: prompt_str.to_owned(),
        }
    }

    fn start(&mut self) {
        let mut stdin = stdio::stdin();
        loop {
            stdio::print(self.cmd_prompt);
            stdio::flush();

            let line = stdin.read_line().unwrap();
            let cmd_line = line.trim().to_owned();
            let program = cmd_line.splitn(' ', 1).nth(0).expect("no program");

            match program {
                ""      =>  { continue; }
                "\u274c"  =>  { return; }
                // _       =>  { self.line(cmd_line); }
                _ => { println!("You just ran {}", cmd_line) }
            }
        }
    }
}

fn main() {
    Shell::new("\U0001f41a \U0001f4b2 ").start();
}
<commit_msg>Add running commands and check to see if command exists<commit_after>use std::io::stdio;
use std::io::Process;
use std::io::process;

struct Shell {
    cmd_prompt: ~str,
}

impl Shell {
    fn new(prompt_str: &str) -> Shell {
        Shell {
            cmd_prompt: prompt_str.to_owned(),
        }
    }

    fn start(&mut self) {
        let mut stdin = stdio::stdin();
        loop {
            stdio::print(self.cmd_prompt);
            stdio::flush();

            let line = stdin.read_line().unwrap();
            let cmd_line = line.trim().to_owned();
            let program = cmd_line.splitn(' ', 1).nth(0).expect("no program");

            match program {
                ""      =>  { continue; }
                "\u274c"  =>  { return; }
                _       =>  { self.cmd(cmd_line.clone()); }
                // _ => { println!("You just ran {}", cmd_line) }
            }
        }
    }

    fn cmd(&mut self, cmd_line: ~str) {
        let mut arg_vec = ~[];
        let mut cmd = ~"";
        for (index, s) in cmd_line.split_str(" ").enumerate() {
            if index == 0 { cmd = s.to_owned(); } else { arg_vec.push(s.to_owned()); }
        }
        self.run_cmd(cmd, arg_vec);
    }

    fn run_cmd(&mut self, cmd: ~str, argv: ~[~str]) {
        if self.cmd_exists(cmd) {
            println!("Executing {}", cmd);
            let p = Process::new(cmd, argv);
            let status = p.unwrap().wait();
            match status {
                process::ExitSignal(st) => println!("Exited with signal {}", st),
                process::ExitStatus(st) => println!("Exited with status {}", st)
            }
        }
        else {
            println!("Command {} not found.", cmd)
        }
    }

    fn cmd_exists(&mut self, cmd: &str) -> bool {
        let p = Process::new("which", [cmd.to_owned()]);
        let status = p.unwrap().wait();
        match status {
            process::ExitSignal(st) => st == 0,
            process::ExitStatus(st) => st == 0
        }
    }
}

fn main() {
    Shell::new("\U0001f41a  ").start();
}
<|endoftext|>


<commit_before>use mio::*;

use super::selector::{EventHandler,Selector};
use super::tunnel_server::TunnelServer;

const TAG: &'static str = "Relay";

pub struct Relay {
    port: u16,
}

impl Relay {
    pub fn new(port: u16) -> Relay {
        Relay {
            port: port,
        }
    }

    pub fn start(&self) {
        info!(target: TAG, "Starting server...");
        let mut selector = Selector::new().unwrap();
        let _tunnel_server = TunnelServer::new(self.port, &mut selector);
        self.poll_loop(&mut selector);
    }

    fn poll_loop(&self, selector: &mut Selector) {
        loop {
            let mut events = Events::with_capacity(1024);
            selector.poll.poll(&mut events, None).expect("Cannot poll");

            for event in &events {
                println!("event={:?}", event);
                selector.run_handler(event);
            }
        }
    }
}

<commit_msg>Move events out of the loop<commit_after>use mio::*;

use super::selector::{EventHandler,Selector};
use super::tunnel_server::TunnelServer;

const TAG: &'static str = "Relay";

pub struct Relay {
    port: u16,
}

impl Relay {
    pub fn new(port: u16) -> Relay {
        Relay {
            port: port,
        }
    }

    pub fn start(&self) {
        info!(target: TAG, "Starting server...");
        let mut selector = Selector::new().unwrap();
        let _tunnel_server = TunnelServer::new(self.port, &mut selector);
        self.poll_loop(&mut selector);
    }

    fn poll_loop(&self, selector: &mut Selector) {
        let mut events = Events::with_capacity(1024);
        loop {
            selector.poll.poll(&mut events, None).expect("Cannot poll");

            for event in &events {
                println!("event={:?}", event);
                selector.run_handler(event);
            }
        }
    }
}

<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use serde_json::{self, Value};

macro_rules! btreemap {
    () => {
        BTreeMap::new()
    };
    ($($key:expr => $value:expr),+) => {
        {
            let mut map = BTreeMap::new();
            $(map.insert($key, $value);)+
            map
        }
    }
}

/*
trait Trait {
    type Type;
}
*/

#[derive(Debug, PartialEq, Serialize, Deserialize)]
struct NamedUnit;

#[derive(Debug, PartialEq, Serialize)]
struct SerNamedTuple<'a, 'b, A: 'a, B: 'b, C>(&'a A, &'b mut B, C);

#[derive(Debug, PartialEq, Deserialize)]
struct DeNamedTuple<A, B, C>(A, B, C);

#[derive(Debug, PartialEq, Serialize)]
struct SerNamedMap<'a, 'b, A: 'a, B: 'b, C> {
    a: &'a A,
    b: &'b mut B,
    c: C,
}

#[derive(Debug, PartialEq, Deserialize)]
struct DeNamedMap<A, B, C> {
    a: A,
    b: B,
    c: C,
}

#[derive(Debug, PartialEq, Serialize)]
enum SerEnum<'a, B: 'a, C: /* Trait + */ 'a, D> where D: /* Trait + */ 'a {
    Unit,
    Seq(
        i8,
        B,
        &'a C,
        //C::Type,
        &'a mut D,
        //<D as Trait>::Type,
    ),
    Map {
        a: i8,
        b: B,
        c: &'a C,
        //d: C::Type,
        e: &'a mut D,
        //f: <D as Trait>::Type,
    },

    // Make sure we can support more than one variant.
    _Unit2,
    _Seq2(
        i8,
        B,
        &'a C,
        //C::Type,
        &'a mut D,
        //<D as Trait>::Type,
    ),
    _Map2 {
        a: i8,
        b: B,
        c: &'a C,
        //d: C::Type,
        e: &'a mut D,
        //f: <D as Trait>::Type,
    },
}

#[derive(Debug, PartialEq, Deserialize)]
enum DeEnum<B, C: /* Trait */, D> /* where D: Trait */ {
    Unit,
    Seq(
        i8,
        B,
        C,
        //C::Type,
        D,
        //<D as Trait>::Type,
    ),
    Map {
        a: i8,
        b: B,
        c: C,
        //d: C::Type,
        e: D,
        //f: <D as Trait>::Type,
    },

    // Make sure we can support more than one variant.
    _Unit2,
    _Seq2(
        i8,
        B,
        C,
        //C::Type,
        D,
        //<D as Trait>::Type,
    ),
    _Map2 {
        a: i8,
        b: B,
        c: C,
        //d: C::Type,
        e: D,
        //f: <D as Trait>::Type,
    },
}

#[derive(Serialize)]
enum Lifetimes<'a> {
    LifetimeSeq(&'a i32),
    NoLifetimeSeq(i32),
    LifetimeMap { a: &'a i32 },
    NoLifetimeMap { a: i32 },
}

#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct GenericStruct<T> {
    x: T,
}

#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct GenericTupleStruct<T>(T);

#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub enum GenericEnumSeq<T, E> {
    Ok(T),
    Err(E),
}

#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub enum GenericEnumMap<T, E> {
    Ok { x: T },
    Err { x: E },
}

#[test]
fn test_named_unit() {
    let named_unit = NamedUnit;

    assert_eq!(
        serde_json::to_string(&named_unit).unwrap(),
        "null".to_string()
    );

    assert_eq!(
        serde_json::to_value(&named_unit),
        Value::Null
    );

    let v: NamedUnit = serde_json::from_str("null").unwrap();
    assert_eq!(v, named_unit);

    let v: NamedUnit = serde_json::from_value(Value::Null).unwrap();
    assert_eq!(v, named_unit);
}

#[test]
fn test_ser_named_tuple() {
    let a = 5;
    let mut b = 6;
    let c = 7;
    let named_tuple = SerNamedTuple(&a, &mut b, c);

    assert_eq!(
        serde_json::to_string(&named_tuple).unwrap(),
        "[5,6,7]"
    );

    assert_eq!(
        serde_json::to_value(&named_tuple),
        Value::Array(vec![Value::U64(5), Value::U64(6), Value::U64(7)])
    );
}

#[test]
fn test_de_named_tuple() {
    let v: DeNamedTuple<i32, i32, i32> = serde_json::from_str("[1,2,3]").unwrap();
    assert_eq!(
        v,
        DeNamedTuple(1, 2, 3)
    );

    let v: Value = serde_json::from_str("[1,2,3]").unwrap();
    assert_eq!(
        v,
        Value::Array(vec![
            Value::U64(1),
            Value::U64(2),
            Value::U64(3),
        ])
    );
}

#[test]
fn test_ser_named_map() {
    let a = 5;
    let mut b = 6;
    let c = 7;
    let named_map = SerNamedMap {
        a: &a,
        b: &mut b,
        c: c,
    };

    assert_eq!(
        serde_json::to_string(&named_map).unwrap(),
        "{\"a\":5,\"b\":6,\"c\":7}"
    );

    assert_eq!(
        serde_json::to_value(&named_map),
        Value::Object(btreemap![
            "a".to_string() => Value::U64(5),
            "b".to_string() => Value::U64(6),
            "c".to_string() => Value::U64(7)
        ])
    );
}

#[test]
fn test_de_named_map() {
    let v = DeNamedMap {
        a: 5,
        b: 6,
        c: 7,
    };

    let v2: DeNamedMap<i32, i32, i32> = serde_json::from_str(
        "{\"a\":5,\"b\":6,\"c\":7}"
    ).unwrap();
    assert_eq!(v, v2);

    let v2 = serde_json::from_value(Value::Object(btreemap![
        "a".to_string() => Value::U64(5),
        "b".to_string() => Value::U64(6),
        "c".to_string() => Value::U64(7)
    ])).unwrap();
    assert_eq!(v, v2);
}

#[test]
fn test_ser_enum_unit() {
    assert_eq!(
        serde_json::to_string(&SerEnum::Unit::<u32, u32, u32>).unwrap(),
        "{\"Unit\":[]}"
    );

    assert_eq!(
        serde_json::to_value(&SerEnum::Unit::<u32, u32, u32>),
        Value::Object(btreemap!(
            "Unit".to_string() => Value::Array(vec![]))
        )
    );
}

#[test]
fn test_ser_enum_seq() {
    let a = 1;
    let b = 2;
    let c = 3;
    //let d = 4;
    let mut e = 5;
    //let f = 6;

    assert_eq!(
        serde_json::to_string(&SerEnum::Seq(
            a,
            b,
            &c,
            //d,
            &mut e,
            //f,
        )).unwrap(),
        "{\"Seq\":[1,2,3,5]}".to_string()
    );

    assert_eq!(
        serde_json::to_value(&SerEnum::Seq(
            a,
            b,
            &c,
            //d,
            &mut e,
            //e,
        )),
        Value::Object(btreemap!(
            "Seq".to_string() => Value::Array(vec![
                Value::U64(1),
                Value::U64(2),
                Value::U64(3),
                //Value::U64(4),
                Value::U64(5),
                //Value::U64(6),
            ])
        ))
    );
}

#[test]
fn test_ser_enum_map() {
    let a = 1;
    let b = 2;
    let c = 3;
    //let d = 4;
    let mut e = 5;
    //let f = 6;

    assert_eq!(
        serde_json::to_string(&SerEnum::Map {
            a: a,
            b: b,
            c: &c,
            //d: d,
            e: &mut e,
            //f: f,
        }).unwrap(),
        "{\"Map\":{\"a\":1,\"b\":2,\"c\":3,\"e\":5}}".to_string()
    );

    assert_eq!(
        serde_json::to_value(&SerEnum::Map {
            a: a,
            b: b,
            c: &c,
            //d: d,
            e: &mut e,
            //f: f,
        }),
        Value::Object(btreemap!(
            "Map".to_string() => Value::Object(btreemap![
                "a".to_string() => Value::U64(1),
                "b".to_string() => Value::U64(2),
                "c".to_string() => Value::U64(3),
                //"d".to_string() => Value::U64(4)
                "e".to_string() => Value::U64(5)
                //"f".to_string() => Value::U64(6)
            ])
        ))
    );
}

#[test]
fn test_de_enum_unit() {
    let v: DeEnum<_, _, _> = serde_json::from_str("{\"Unit\":[]}").unwrap();
    assert_eq!(
        v,
        DeEnum::Unit::<u32, u32, u32>
    );

    let v: DeEnum<_, _, _> = serde_json::from_value(Value::Object(btreemap!(
        "Unit".to_string() => Value::Array(vec![]))
    )).unwrap();
    assert_eq!(
        v,
        DeEnum::Unit::<u32, u32, u32>
    );
}

#[test]
fn test_de_enum_seq() {
    let a = 1;
    let b = 2;
    let c = 3;
    //let d = 4;
    let e = 5;
    //let f = 6;

    let v: DeEnum<_, _, _> = serde_json::from_str("{\"Seq\":[1,2,3,5]}").unwrap();
    assert_eq!(
        v,
        DeEnum::Seq(
            a,
            b,
            c,
            //d,
            e,
            //f,
        )
    );

    let v: DeEnum<_, _, _> = serde_json::from_value(Value::Object(btreemap!(
        "Seq".to_string() => Value::Array(vec![
            Value::U64(1),
            Value::U64(2),
            Value::U64(3),
            //Value::U64(4),
            Value::U64(5),
            //Value::U64(6),
        ])
    ))).unwrap();
    assert_eq!(
        v,
        DeEnum::Seq(
            a,
            b,
            c,
            //d,
            e,
            //e,
        )
    );
}

#[test]
fn test_de_enum_map() {
    let a = 1;
    let b = 2;
    let c = 3;
    //let d = 4;
    let e = 5;
    //let f = 6;

    let v: DeEnum<_, _, _> = serde_json::from_str(
        "{\"Map\":{\"a\":1,\"b\":2,\"c\":3,\"e\":5}}"
    ).unwrap();
    assert_eq!(
        v,
        DeEnum::Map {
            a: a,
            b: b,
            c: c,
            //d: d,
            e: e,
            //f: f,
        }
    );

    let v: DeEnum<_, _, _> = serde_json::from_value(Value::Object(btreemap!(
        "Map".to_string() => Value::Object(btreemap![
            "a".to_string() => Value::U64(1),
            "b".to_string() => Value::U64(2),
            "c".to_string() => Value::U64(3),
            //"d".to_string() => Value::U64(4)
            "e".to_string() => Value::U64(5)
            //"f".to_string() => Value::U64(6)
        ])
    ))).unwrap();

    assert_eq!(
        v,
        DeEnum::Map {
            a: a,
            b: b,
            c: c,
            //d: d,
            e: e,
            //f: f,
        }
    );
}

#[test]
fn test_lifetimes() {
    let value = 5;
    let lifetime = Lifetimes::LifetimeSeq(&value);
    assert_eq!(
        serde_json::to_string(&lifetime).unwrap(),
        "{\"LifetimeSeq\":5}"
    );

    let lifetime = Lifetimes::NoLifetimeSeq(5);
    assert_eq!(
        serde_json::to_string(&lifetime).unwrap(),
        "{\"NoLifetimeSeq\":5}"
    );

    let value = 5;
    let lifetime = Lifetimes::LifetimeMap { a: &value };
    assert_eq!(
        serde_json::to_string(&lifetime).unwrap(),
        "{\"LifetimeMap\":{\"a\":5}}"
    );

    let lifetime = Lifetimes::NoLifetimeMap { a: 5 };
    assert_eq!(
        serde_json::to_string(&lifetime).unwrap(),
        "{\"NoLifetimeMap\":{\"a\":5}}"
    );
}

#[test]
fn test_generic() {
    macro_rules! declare_tests {
        ($($ty:ty : $value:expr => $str:expr,)+) => {
            $({
                let value: $ty = $value;

                let string = serde_json::to_string(&value).unwrap();
                assert_eq!(string, $str);

                let expected: $ty = serde_json::from_str(&string).unwrap();
                assert_eq!(value, expected);
            })+
        }
    }

    declare_tests!(
        GenericStruct<u32> : GenericStruct { x: 5 } => "{\"x\":5}",
        GenericTupleStruct<u32> : GenericTupleStruct(5) => "[5]",
        GenericEnumSeq<u32, u32> : GenericEnumSeq::Ok(5) => "{\"Ok\":[5]}",
        GenericEnumSeq<u32, u32> : GenericEnumSeq::Err(5) => "{\"Err\":[5]}",
        GenericEnumMap<u32, u32> : GenericEnumMap::Ok { x: 5 } => "{\"Ok\":{\"x\":5}}",
        GenericEnumMap<u32, u32> : GenericEnumMap::Err { x: 5 } => "{\"Err\":{\"x\":5}}",
    );
}
<commit_msg>rebased again<commit_after>use std::collections::BTreeMap;
use serde_json::{self, Value};

macro_rules! btreemap {
    () => {
        BTreeMap::new()
    };
    ($($key:expr => $value:expr),+) => {
        {
            let mut map = BTreeMap::new();
            $(map.insert($key, $value);)+
            map
        }
    }
}

/*
trait Trait {
    type Type;
}
*/

#[derive(Debug, PartialEq, Serialize, Deserialize)]
struct NamedUnit;

#[derive(Debug, PartialEq, Serialize)]
struct SerNamedTuple<'a, 'b, A: 'a, B: 'b, C>(&'a A, &'b mut B, C);

#[derive(Debug, PartialEq, Deserialize)]
struct DeNamedTuple<A, B, C>(A, B, C);

#[derive(Debug, PartialEq, Serialize)]
struct SerNamedMap<'a, 'b, A: 'a, B: 'b, C> {
    a: &'a A,
    b: &'b mut B,
    c: C,
}

#[derive(Debug, PartialEq, Deserialize)]
struct DeNamedMap<A, B, C> {
    a: A,
    b: B,
    c: C,
}

#[derive(Debug, PartialEq, Serialize)]
enum SerEnum<'a, B: 'a, C: /* Trait + */ 'a, D> where D: /* Trait + */ 'a {
    Unit,
    Seq(
        i8,
        B,
        &'a C,
        //C::Type,
        &'a mut D,
        //<D as Trait>::Type,
    ),
    Map {
        a: i8,
        b: B,
        c: &'a C,
        //d: C::Type,
        e: &'a mut D,
        //f: <D as Trait>::Type,
    },

    // Make sure we can support more than one variant.
    _Unit2,
    _Seq2(
        i8,
        B,
        &'a C,
        //C::Type,
        &'a mut D,
        //<D as Trait>::Type,
    ),
    _Map2 {
        a: i8,
        b: B,
        c: &'a C,
        //d: C::Type,
        e: &'a mut D,
        //f: <D as Trait>::Type,
    },
}

#[derive(Debug, PartialEq, Deserialize)]
enum DeEnum<B, C: /* Trait */, D> /* where D: Trait */ {
    Unit,
    Seq(
        i8,
        B,
        C,
        //C::Type,
        D,
        //<D as Trait>::Type,
    ),
    Map {
        a: i8,
        b: B,
        c: C,
        //d: C::Type,
        e: D,
        //f: <D as Trait>::Type,
    },

    // Make sure we can support more than one variant.
    _Unit2,
    _Seq2(
        i8,
        B,
        C,
        //C::Type,
        D,
        //<D as Trait>::Type,
    ),
    _Map2 {
        a: i8,
        b: B,
        c: C,
        //d: C::Type,
        e: D,
        //f: <D as Trait>::Type,
    },
}

#[derive(Serialize)]
enum Lifetimes<'a> {
    LifetimeSeq(&'a i32),
    NoLifetimeSeq(i32),
    LifetimeMap { a: &'a i32 },
    NoLifetimeMap { a: i32 },
}

#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct GenericStruct<T> {
    x: T,
}

#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct GenericTupleStruct<T>(T);

#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub enum GenericEnumSeq<T, E> {
    Ok(T),
    Err(E),
}

#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub enum GenericEnumMap<T, E> {
    Ok { x: T },
    Err { x: E },
}

#[test]
fn test_named_unit() {
    let named_unit = NamedUnit;

    assert_eq!(
        serde_json::to_string(&named_unit).unwrap(),
        "null".to_string()
    );

    assert_eq!(
        serde_json::to_value(&named_unit),
        Value::Null
    );

    let v: NamedUnit = serde_json::from_str("null").unwrap();
    assert_eq!(v, named_unit);

    let v: NamedUnit = serde_json::from_value(Value::Null).unwrap();
    assert_eq!(v, named_unit);
}

#[test]
fn test_ser_named_tuple() {
    let a = 5;
    let mut b = 6;
    let c = 7;
    let named_tuple = SerNamedTuple(&a, &mut b, c);

    assert_eq!(
        serde_json::to_string(&named_tuple).unwrap(),
        "[5,6,7]"
    );

    assert_eq!(
        serde_json::to_value(&named_tuple),
        Value::Array(vec![Value::U64(5), Value::U64(6), Value::U64(7)])
    );
}

#[test]
fn test_de_named_tuple() {
    let v: DeNamedTuple<i32, i32, i32> = serde_json::from_str("[1,2,3]").unwrap();
    assert_eq!(
        v,
        DeNamedTuple(1, 2, 3)
    );

    let v: Value = serde_json::from_str("[1,2,3]").unwrap();
    assert_eq!(
        v,
        Value::Array(vec![
            Value::U64(1),
            Value::U64(2),
            Value::U64(3),
        ])
    );
}

#[test]
fn test_ser_named_map() {
    let a = 5;
    let mut b = 6;
    let c = 7;
    let named_map = SerNamedMap {
        a: &a,
        b: &mut b,
        c: c,
    };

    assert_eq!(
        serde_json::to_string(&named_map).unwrap(),
        "{\"a\":5,\"b\":6,\"c\":7}"
    );

    assert_eq!(
        serde_json::to_value(&named_map),
        Value::Object(btreemap![
            "a".to_string() => Value::U64(5),
            "b".to_string() => Value::U64(6),
            "c".to_string() => Value::U64(7)
        ])
    );
}

#[test]
fn test_de_named_map() {
    let v = DeNamedMap {
        a: 5,
        b: 6,
        c: 7,
    };

    let v2: DeNamedMap<i32, i32, i32> = serde_json::from_str(
        "{\"a\":5,\"b\":6,\"c\":7}"
    ).unwrap();
    assert_eq!(v, v2);

    let v2 = serde_json::from_value(Value::Object(btreemap![
        "a".to_string() => Value::U64(5),
        "b".to_string() => Value::U64(6),
        "c".to_string() => Value::U64(7)
    ])).unwrap();
    assert_eq!(v, v2);
}

#[test]
fn test_ser_enum_unit() {
    assert_eq!(
        serde_json::to_string(&SerEnum::Unit::<u32, u32, u32>).unwrap(),
        "{\"Unit\":[]}"
    );

    assert_eq!(
        serde_json::to_value(&SerEnum::Unit::<u32, u32, u32>),
        Value::Object(btreemap!(
            "Unit".to_string() => Value::Array(vec![]))
        )
    );
}

#[test]
fn test_ser_enum_seq() {
    let a = 1;
    let b = 2;
    let c = 3;
    //let d = 4;
    let mut e = 5;
    //let f = 6;

    assert_eq!(
        serde_json::to_string(&SerEnum::Seq(
            a,
            b,
            &c,
            //d,
            &mut e,
            //f,
        )).unwrap(),
        "{\"Seq\":[1,2,3,5]}".to_string()
    );

    assert_eq!(
        serde_json::to_value(&SerEnum::Seq(
            a,
            b,
            &c,
            //d,
            &mut e,
            //e,
        )),
        Value::Object(btreemap!(
            "Seq".to_string() => Value::Array(vec![
                Value::U64(1),
                Value::U64(2),
                Value::U64(3),
                //Value::U64(4),
                Value::U64(5),
                //Value::U64(6),
            ])
        ))
    );
}

#[test]
fn test_ser_enum_map() {
    let a = 1;
    let b = 2;
    let c = 3;
    //let d = 4;
    let mut e = 5;
    //let f = 6;

    assert_eq!(
        serde_json::to_string(&SerEnum::Map {
            a: a,
            b: b,
            c: &c,
            //d: d,
            e: &mut e,
            //f: f,
        }).unwrap(),
        "{\"Map\":{\"a\":1,\"b\":2,\"c\":3,\"e\":5}}".to_string()
    );

    assert_eq!(
        serde_json::to_value(&SerEnum::Map {
            a: a,
            b: b,
            c: &c,
            //d: d,
            e: &mut e,
            //f: f,
        }),
        Value::Object(btreemap!(
            "Map".to_string() => Value::Object(btreemap![
                "a".to_string() => Value::U64(1),
                "b".to_string() => Value::U64(2),
                "c".to_string() => Value::U64(3),
                //"d".to_string() => Value::U64(4)
                "e".to_string() => Value::U64(5)
                //"f".to_string() => Value::U64(6)
            ])
        ))
    );
}

#[test]
fn test_de_enum_unit() {
    let v: DeEnum<_, _, _> = serde_json::from_str("{\"Unit\":[]}").unwrap();
    assert_eq!(
        v,
        DeEnum::Unit::<u32, u32, u32>
    );

    let v: DeEnum<_, _, _> = serde_json::from_value(Value::Object(btreemap!(
        "Unit".to_string() => Value::Array(vec![]))
    )).unwrap();
    assert_eq!(
        v,
        DeEnum::Unit::<u32, u32, u32>
    );
}

#[test]
fn test_de_enum_seq() {
    let a = 1;
    let b = 2;
    let c = 3;
    //let d = 4;
    let e = 5;
    //let f = 6;

    let v: DeEnum<_, _, _> = serde_json::from_str("{\"Seq\":[1,2,3,5]}").unwrap();
    assert_eq!(
        v,
        DeEnum::Seq(
            a,
            b,
            c,
            //d,
            e,
            //f,
        )
    );

    let v: DeEnum<_, _, _> = serde_json::from_value(Value::Object(btreemap!(
        "Seq".to_string() => Value::Array(vec![
            Value::U64(1),
            Value::U64(2),
            Value::U64(3),
            //Value::U64(4),
            Value::U64(5),
            //Value::U64(6),
        ])
    ))).unwrap();
    assert_eq!(
        v,
        DeEnum::Seq(
            a,
            b,
            c,
            //d,
            e,
            //e,
        )
    );
}

#[test]
fn test_de_enum_map() {
    let a = 1;
    let b = 2;
    let c = 3;
    //let d = 4;
    let e = 5;
    //let f = 6;

    let v: DeEnum<_, _, _> = serde_json::from_str(
        "{\"Map\":{\"a\":1,\"b\":2,\"c\":3,\"e\":5}}"
    ).unwrap();
    assert_eq!(
        v,
        DeEnum::Map {
            a: a,
            b: b,
            c: c,
            //d: d,
            e: e,
            //f: f,
        }
    );

    let v: DeEnum<_, _, _> = serde_json::from_value(Value::Object(btreemap!(
        "Map".to_string() => Value::Object(btreemap![
            "a".to_string() => Value::U64(1),
            "b".to_string() => Value::U64(2),
            "c".to_string() => Value::U64(3),
            //"d".to_string() => Value::U64(4)
            "e".to_string() => Value::U64(5)
            //"f".to_string() => Value::U64(6)
        ])
    ))).unwrap();

    assert_eq!(
        v,
        DeEnum::Map {
            a: a,
            b: b,
            c: c,
            //d: d,
            e: e,
            //f: f,
        }
    );
}

#[test]
fn test_lifetimes() {
    let value = 5;
    let lifetime = Lifetimes::LifetimeSeq(&value);
    assert_eq!(
        serde_json::to_string(&lifetime).unwrap(),
        "{\"LifetimeSeq\":5}"
    );

    let lifetime = Lifetimes::NoLifetimeSeq(5);
    assert_eq!(
        serde_json::to_string(&lifetime).unwrap(),
        "{\"NoLifetimeSeq\":5}"
    );

    let value = 5;
    let lifetime = Lifetimes::LifetimeMap { a: &value };
    assert_eq!(
        serde_json::to_string(&lifetime).unwrap(),
        "{\"LifetimeMap\":{\"a\":5}}"
    );

    let lifetime = Lifetimes::NoLifetimeMap { a: 5 };
    assert_eq!(
        serde_json::to_string(&lifetime).unwrap(),
        "{\"NoLifetimeMap\":{\"a\":5}}"
    );
}

#[test]
fn test_generic() {
    macro_rules! declare_tests {
        ($($ty:ty : $value:expr => $str:expr,)+) => {
            $({
                let value: $ty = $value;

                let string = serde_json::to_string(&value).unwrap();
                assert_eq!(string, $str);

                let expected: $ty = serde_json::from_str(&string).unwrap();
                assert_eq!(value, expected);
            })+
        }
    }

    declare_tests!(
        GenericStruct<u32> : GenericStruct { x: 5 } => "{\"x\":5}",
        GenericTupleStruct<u32> : GenericTupleStruct(5) => "[5]",
        GenericEnumSeq<u32, u32> : GenericEnumSeq::Ok(5) => "{\"Ok\":5}",
        GenericEnumSeq<u32, u32> : GenericEnumSeq::Err(5) => "{\"Err\":5}",
        GenericEnumMap<u32, u32> : GenericEnumMap::Ok { x: 5 } => "{\"Ok\":{\"x\":5}}",
        GenericEnumMap<u32, u32> : GenericEnumMap::Err { x: 5 } => "{\"Err\":{\"x\":5}}",
    );
}
<|endoftext|>


<commit_before>use std::env;
use std::fs::File;
use std::io::prelude::*;

fn main() {
    if let Some(path) = env::args().nth(1) {
        let mut contents = String::new();
        match File::open(path) {
            Ok(mut file) => {
                match file.read_to_string(&mut contents) {
                    Ok(_) => println!("Would read and eval here."),
                    Err(_) => println!("Found the file but failed to read it."),
                }
            }
            Err(_) => println!("Failed to open the file for reading."),
        };
    } else {
        println!("Please provide a path to some brainfuck source as the first argument.")
    }
}

mod brainfuck {
    #[derive(PartialEq)]
    #[derive(Debug)]
    pub enum Command {
        IncrementPointer,
        DecrementPointer,
        IncrementValue,
        DecrementValue,
        OutputValue,
        InputValue,
        ForwardsTo(usize),
        BackwardsTo(usize),
    }

    pub fn read(source: &str) -> Result<Vec<Command>, &'static str> {
        let mut forward_pointers: Vec<usize> = vec![];
        let mut backward_pointers: Vec<usize> = vec![];
        let mut position: usize = 0;
        let mut program: Vec<Command> = source
            .chars()
            .filter_map(|token| {
                let command = match token {
                    '>' => Some(Command::IncrementPointer),
                    '<' => Some(Command::DecrementPointer),
                    '+' => Some(Command::IncrementValue),
                    '-' => Some(Command::DecrementValue),
                    '.' => Some(Command::OutputValue),
                    ',' => Some(Command::InputValue),
                    '[' => {
                        forward_pointers.push(position);
                        Some(Command::BackwardsTo(position))
                    }
                    ']' => {
                        backward_pointers.push(position);
                        Some(Command::ForwardsTo(position))
                    }
                    _ => None,
                };

                if command.is_some() {
                    position += 1
                }

                command
            })
            .collect();

        if forward_pointers.len() == backward_pointers.len() {
            for (a, b) in forward_pointers.iter().zip(backward_pointers) {
                program.swap(*a, b)
            }

            Ok(program)
        } else {
            Err("Found an unmatched square brace.")
        }
    }

    // pub fn eval(program: Vec<Command>) {
    // }
}

#[cfg(test)]
mod tests {
    use brainfuck::*;
    use brainfuck::Command::*;

    #[test]
    fn read_empty() {
        let source = "";
        let expected: Vec<Command> = vec![];

        match read(source) {
            Ok(actual) => assert_eq!(expected, actual),
            Err(_) => assert!(false),
        };
    }

    #[test]
    fn read_simple() {
        let source = "+>-<.,<.-+>,";
        let expected = vec![IncrementValue,
                            IncrementPointer,
                            DecrementValue,
                            DecrementPointer,
                            OutputValue,
                            InputValue,
                            DecrementPointer,
                            OutputValue,
                            DecrementValue,
                            IncrementValue,
                            IncrementPointer,
                            InputValue];

        match read(source) {
            Ok(actual) => assert_eq!(expected, actual),
            Err(_) => assert!(false),
        };
    }

    #[test]
    fn read_garbage() {
        let source = ",.+lol,hey.>there<";
        let expected = vec![InputValue,
                            OutputValue,
                            IncrementValue,
                            InputValue,
                            OutputValue,
                            IncrementPointer,
                            DecrementPointer];

        match read(source) {
            Ok(actual) => assert_eq!(expected, actual),
            Err(_) => assert!(false),
        };
    }

    #[test]
    fn read_loop() {
        let source = "[->+<]";
        let expected = vec![ForwardsTo(5),
                            DecrementValue,
                            IncrementPointer,
                            IncrementValue,
                            DecrementPointer,
                            BackwardsTo(0)];

        match read(source) {
            Ok(actual) => assert_eq!(expected, actual),
            Err(_) => assert!(false),
        };
    }

    #[test]
    fn read_bad_loop() {
        let source = "[,.[,.>+-][./.[]]";
        let expected = "Found an unmatched square brace.";

        match read(source) {
            Ok(_) => assert!(false),
            Err(actual) => assert_eq!(expected, actual),
        };
    }

    // #[test]
    // fn eval_empty() {
    //     let program = vec![];
    //     eval(program);
    //     assert_eq!(expected, output)
    // }
}
<commit_msg>Almost have whole eval and main implemented<commit_after>use std::env;
use std::io;
use std::fs::File;
use std::io::prelude::*;

fn main() {
    if let Some(path) = env::args().nth(1) {
        let mut contents = String::new();
        match File::open(path) {
            Ok(mut file) => {
                match file.read_to_string(&mut contents) {
                    Ok(_) => {
                        match brainfuck::read(&contents) {
                            Ok(program) => {
                                let stdio = io::stdin();
                                let input = stdio.lock();
                                let output = io::stdout();
                                brainfuck::eval(program, input, output);
                            }
                            Err(msg) => {
                                println!("Read the file, but could not parse it: {}", msg)
                            }
                        }
                    }
                    Err(_) => println!("Found the file but failed to read it."),
                }
            }
            Err(_) => println!("Failed to open the file for reading."),
        };
    } else {
        println!("Please provide a path to some brainfuck source as the first argument.")
    }
}

mod brainfuck {
    use std::io::{BufRead, Write};

    #[derive(PartialEq)]
    #[derive(Debug)]
    pub enum Command {
        IncrementPointer,
        DecrementPointer,
        IncrementValue,
        DecrementValue,
        OutputValue,
        InputValue,
        ForwardsTo(usize),
        BackwardsTo(usize),
    }

    pub fn read(source: &str) -> Result<Vec<Command>, &'static str> {
        let mut forward_pointers: Vec<usize> = vec![];
        let mut backward_pointers: Vec<usize> = vec![];
        let mut position: usize = 0;
        let mut program: Vec<Command> = source
            .chars()
            .filter_map(|token| {
                let command = match token {
                    '>' => Some(Command::IncrementPointer),
                    '<' => Some(Command::DecrementPointer),
                    '+' => Some(Command::IncrementValue),
                    '-' => Some(Command::DecrementValue),
                    '.' => Some(Command::OutputValue),
                    ',' => Some(Command::InputValue),
                    '[' => {
                        forward_pointers.push(position);
                        Some(Command::BackwardsTo(position))
                    }
                    ']' => {
                        backward_pointers.push(position);
                        Some(Command::ForwardsTo(position))
                    }
                    _ => None,
                };

                if command.is_some() {
                    position += 1
                }

                command
            })
            .collect();

        if forward_pointers.len() == backward_pointers.len() {
            for (a, b) in forward_pointers.iter().zip(backward_pointers) {
                program.swap(*a, b)
            }

            Ok(program)
        } else {
            Err("Found an unmatched square brace.")
        }
    }

    pub fn eval<R, W>(program: Vec<Command>, mut reader: R, mut writer: W)
        where R: BufRead,
              W: Write
    {
        let mut memory: Vec<i8> = vec![0; 30000];
        let mut mem_pointer: usize = 0;
        let mut prog_pointer: usize = 0;
        let prog_len = program.len();

        while prog_pointer < prog_len {
            match program[prog_pointer] {
                Command::IncrementPointer => {
                    mem_pointer += 1;
                }
                Command::DecrementPointer => {
                    mem_pointer -= 1;
                }
                Command::IncrementValue => {
                    memory[mem_pointer] += 1;
                }
                Command::DecrementValue => {
                    memory[mem_pointer] -= 1;
                }
                Command::OutputValue => {
                    write!(&mut writer, "{}", memory[mem_pointer]).expect("Unable to write");
                }
                Command::InputValue => {
                    let mut input = String::new();
                    reader.read_line(&mut input).expect("Unable to read");
                    memory[mem_pointer] = input.as_bytes()[0] as i8;
                }
                Command::ForwardsTo(next) => {
                    if memory[mem_pointer] == 0 {
                        prog_pointer = next;
                    }
                }
                Command::BackwardsTo(next) => {
                    if memory[mem_pointer] != 0 {
                        prog_pointer = next;
                    }
                }
            };

            prog_pointer += 1;
        }
    }
}

#[cfg(test)]
mod tests {
    use brainfuck::*;
    use brainfuck::Command::*;
    use std::io::Cursor;

    #[test]
    fn read_empty() {
        let source = "";
        let expected: Vec<Command> = vec![];

        match read(source) {
            Ok(actual) => assert_eq!(expected, actual),
            Err(_) => assert!(false),
        };
    }

    #[test]
    fn read_simple() {
        let source = "+>-<.,<.-+>,";
        let expected = vec![IncrementValue,
                            IncrementPointer,
                            DecrementValue,
                            DecrementPointer,
                            OutputValue,
                            InputValue,
                            DecrementPointer,
                            OutputValue,
                            DecrementValue,
                            IncrementValue,
                            IncrementPointer,
                            InputValue];

        match read(source) {
            Ok(actual) => assert_eq!(expected, actual),
            Err(_) => assert!(false),
        };
    }

    #[test]
    fn read_garbage() {
        let source = ",.+lol,hey.>there<";
        let expected = vec![InputValue,
                            OutputValue,
                            IncrementValue,
                            InputValue,
                            OutputValue,
                            IncrementPointer,
                            DecrementPointer];

        match read(source) {
            Ok(actual) => assert_eq!(expected, actual),
            Err(_) => assert!(false),
        };
    }

    #[test]
    fn read_loop() {
        let source = "[->+<]";
        let expected = vec![ForwardsTo(5),
                            DecrementValue,
                            IncrementPointer,
                            IncrementValue,
                            DecrementPointer,
                            BackwardsTo(0)];

        match read(source) {
            Ok(actual) => assert_eq!(expected, actual),
            Err(_) => assert!(false),
        };
    }

    #[test]
    fn read_bad_loop() {
        let source = "[,.[,.>+-][./.[]]";
        let expected = "Found an unmatched square brace.";

        match read(source) {
            Ok(_) => assert!(false),
            Err(actual) => assert_eq!(expected, actual),
        };
    }

    #[test]
    fn eval_empty() {
        let input = Cursor::new(&b""[..]);
        let mut output = Cursor::new(vec![]);

        match read("") {
            Ok(program) => eval(program, input, &mut output),
            Err(_) => assert!(false),
        }

        let expected = "";
        let actual = String::from_utf8(output.into_inner()).expect("Not UTF-8");

        assert_eq!(expected, actual);
    }

    #[test]
    fn eval_cat() {
        let input = Cursor::new(&b"Hello, World!"[..]);
        let mut output = Cursor::new(vec![]);

        match read(",[.,]") {
            Ok(program) => eval(program, input, &mut output),
            Err(_) => assert!(false),
        }

        let expected = "Hello, World!";
        let actual = String::from_utf8(output.into_inner()).expect("Not UTF-8");

        assert_eq!(expected, actual);
    }
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::io::Read;
use std::path::Path;

use lsp_types::{Location, Range, Url};

use mojom_syntax::{preorder, Traversal};

use crate::definition::create_lsp_range;
use crate::mojomast::MojomAst;

#[derive(Debug)]
struct ImportDefinition {
    pub ident: String,
    pub range: Range,
}

#[derive(Debug)]
struct Import {
    uri: Url,
    module_name: Option<String>,
    definitions: Vec<ImportDefinition>,
}

#[derive(Debug)]
enum ImportError {
    IoError(std::io::Error),
    SyntaxError(String),
}

impl From<std::io::Error> for ImportError {
    fn from(err: std::io::Error) -> Self {
        ImportError::IoError(err)
    }
}

type ImportResult = std::result::Result<Import, ImportError>;

#[derive(Debug)]
pub(crate) struct ImportedFiles {
    parsed_imports: Vec<ImportResult>,
}

impl ImportedFiles {
    pub(crate) fn find_definition(&self, ident: &str) -> Option<Location> {
        let valid_imports = self.parsed_imports.iter().filter_map(|i| i.as_ref().ok());
        for imported in valid_imports {
            for definition in &imported.definitions {
                if definition.ident == ident {
                    let loc = Location::new(imported.uri.clone(), definition.range.clone());
                    return Some(loc);
                }

                if let Some(module_name) = &imported.module_name {
                    let canonocal_name = format!("{}.{}", module_name, definition.ident);
                    if canonocal_name == ident {
                        let loc = Location::new(imported.uri.clone(), definition.range.clone());
                        return Some(loc);
                    }
                }
            }
        }
        None
    }
}

pub(crate) fn check_imports<P: AsRef<Path>>(root_path: P, ast: &MojomAst) -> ImportedFiles {
    let root_path = root_path.as_ref();
    let mut parsed_imports = Vec::new();
    for stmt in &ast.mojom.stmts {
        match stmt {
            mojom_syntax::Statement::Import(stmt) => {
                let path = ast.text(&stmt.path);
                let path = root_path.join(&path[1..path.len() - 1]);
                let imported = parse_imported(&path);
                parsed_imports.push(imported);
            }
            _ => (),
        }
    }

    ImportedFiles {
        parsed_imports: parsed_imports,
    }
}

fn add_definition<'a, 'b, 'c>(
    field: &'a mojom_syntax::Range,
    ast: &'b MojomAst,
    path: &'c mut Vec<&'b str>,
    definitions: &'c mut Vec<ImportDefinition>,
) {
    let name = ast.text(field);
    path.push(name);
    let ident = path.join(".");
    path.pop();
    let range = create_lsp_range(&ast, field);
    definitions.push(ImportDefinition {
        ident: ident,
        range: range,
    });
}

fn parse_imported<P: AsRef<Path>>(path: P) -> ImportResult {
    let mut text = String::new();
    File::open(path.as_ref())?.read_to_string(&mut text)?;

    let mojom =
        mojom_syntax::parse(&text).map_err(|err| ImportError::SyntaxError(err.to_string()))?;

    // Unwrap shoud be safe because we opened file already.
    let path = path.as_ref().canonicalize().unwrap();
    let uri = Url::from_file_path(&path).unwrap();

    let ast = MojomAst::from_mojom(uri, text, mojom);

    let mut path = Vec::new();
    let mut definitions: Vec<ImportDefinition> = Vec::new();
    for traversal in preorder(&ast.mojom) {
        match traversal {
            Traversal::EnterInterface(node) => {
                add_definition(&node.name, &ast, &mut path, &mut definitions);
                let name = ast.text(&node.name);
                path.push(name);
            }
            Traversal::LeaveInterface(_) => {
                path.pop();
            }
            Traversal::EnterStruct(node) => {
                add_definition(&node.name, &ast, &mut path, &mut definitions);
                let name = ast.text(&node.name);
                path.push(name);
            }
            Traversal::LeaveStruct(_) => {
                path.pop();
            }
            Traversal::Union(node) => add_definition(&node.name, &ast, &mut path, &mut definitions),
            Traversal::Enum(node) => add_definition(&node.name, &ast, &mut path, &mut definitions),
            Traversal::Const(node) => add_definition(&node.name, &ast, &mut path, &mut definitions),
            _ => (),
        }
    }

    let module_name = ast.module_name().map(|name| name.to_owned());

    Ok(Import {
        uri: ast.uri.clone(),
        module_name: module_name,
        definitions: definitions,
    })
}

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

    fn create_uri<P: AsRef<Path>>(path: P) -> Url {
        let path = path.as_ref().canonicalize().unwrap();
        Url::from_file_path(path).unwrap()
    }

    #[test]
    fn test_parse_imported() {
        let res = parse_imported("../testdata/my_interface.mojom");
        assert!(res.is_ok());
    }

    #[test]
    fn test_check_imports() {
        let root_path = "../testdata";
        let file_path = "../testdata/my_service.mojom";
        let mut text = String::new();
        File::open(&file_path)
            .unwrap()
            .read_to_string(&mut text)
            .unwrap();
        let uri = create_uri(&file_path);
        let ast = MojomAst::new(uri, text).unwrap();

        let imports = check_imports(&root_path, &ast);

        let res = imports.find_definition("FooStruct.FooEnum");
        assert!(res.is_some());
    }
}
<commit_msg>Add NotFound error for imported files<commit_after>use std::fs::File;
use std::io::Read;
use std::path::Path;

use lsp_types::{Location, Range, Url};

use mojom_syntax::{preorder, Traversal};

use crate::definition::create_lsp_range;
use crate::mojomast::MojomAst;

#[derive(Debug)]
struct ImportDefinition {
    pub ident: String,
    pub range: Range,
}

#[derive(Debug)]
struct Import {
    uri: Url,
    module_name: Option<String>,
    definitions: Vec<ImportDefinition>,
}

#[derive(Debug)]
enum ImportError {
    IoError(std::io::Error),
    NotFound(String /* path */),
    SyntaxError(String),
}

impl From<std::io::Error> for ImportError {
    fn from(err: std::io::Error) -> Self {
        if err.kind() == std::io::ErrorKind::NotFound {
            return ImportError::NotFound(err.to_string());
        }
        ImportError::IoError(err)
    }
}

type ImportResult = std::result::Result<Import, ImportError>;

#[derive(Debug)]
pub(crate) struct ImportedFiles {
    parsed_imports: Vec<ImportResult>,
}

impl ImportedFiles {
    pub(crate) fn find_definition(&self, ident: &str) -> Option<Location> {
        let valid_imports = self.parsed_imports.iter().filter_map(|i| i.as_ref().ok());
        for imported in valid_imports {
            for definition in &imported.definitions {
                if definition.ident == ident {
                    let loc = Location::new(imported.uri.clone(), definition.range.clone());
                    return Some(loc);
                }

                if let Some(module_name) = &imported.module_name {
                    let canonocal_name = format!("{}.{}", module_name, definition.ident);
                    if canonocal_name == ident {
                        let loc = Location::new(imported.uri.clone(), definition.range.clone());
                        return Some(loc);
                    }
                }
            }
        }
        None
    }
}

pub(crate) fn check_imports<P: AsRef<Path>>(root_path: P, ast: &MojomAst) -> ImportedFiles {
    let root_path = root_path.as_ref();
    let mut parsed_imports = Vec::new();
    for stmt in &ast.mojom.stmts {
        match stmt {
            mojom_syntax::Statement::Import(stmt) => {
                let path = ast.text(&stmt.path);
                let path = root_path.join(&path[1..path.len() - 1]);
                let imported = parse_imported(&path);
                parsed_imports.push(imported);
            }
            _ => (),
        }
    }

    ImportedFiles {
        parsed_imports: parsed_imports,
    }
}

fn add_definition<'a, 'b, 'c>(
    field: &'a mojom_syntax::Range,
    ast: &'b MojomAst,
    path: &'c mut Vec<&'b str>,
    definitions: &'c mut Vec<ImportDefinition>,
) {
    let name = ast.text(field);
    path.push(name);
    let ident = path.join(".");
    path.pop();
    let range = create_lsp_range(&ast, field);
    definitions.push(ImportDefinition {
        ident: ident,
        range: range,
    });
}

fn parse_imported<P: AsRef<Path>>(path: P) -> ImportResult {
    let mut text = String::new();
    File::open(path.as_ref())?.read_to_string(&mut text)?;

    let mojom =
        mojom_syntax::parse(&text).map_err(|err| ImportError::SyntaxError(err.to_string()))?;

    // Unwrap shoud be safe because we opened file already.
    let path = path.as_ref().canonicalize().unwrap();
    let uri = Url::from_file_path(&path).unwrap();

    let ast = MojomAst::from_mojom(uri, text, mojom);

    let mut path = Vec::new();
    let mut definitions: Vec<ImportDefinition> = Vec::new();
    for traversal in preorder(&ast.mojom) {
        match traversal {
            Traversal::EnterInterface(node) => {
                add_definition(&node.name, &ast, &mut path, &mut definitions);
                let name = ast.text(&node.name);
                path.push(name);
            }
            Traversal::LeaveInterface(_) => {
                path.pop();
            }
            Traversal::EnterStruct(node) => {
                add_definition(&node.name, &ast, &mut path, &mut definitions);
                let name = ast.text(&node.name);
                path.push(name);
            }
            Traversal::LeaveStruct(_) => {
                path.pop();
            }
            Traversal::Union(node) => add_definition(&node.name, &ast, &mut path, &mut definitions),
            Traversal::Enum(node) => add_definition(&node.name, &ast, &mut path, &mut definitions),
            Traversal::Const(node) => add_definition(&node.name, &ast, &mut path, &mut definitions),
            _ => (),
        }
    }

    let module_name = ast.module_name().map(|name| name.to_owned());

    Ok(Import {
        uri: ast.uri.clone(),
        module_name: module_name,
        definitions: definitions,
    })
}

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

    fn create_uri<P: AsRef<Path>>(path: P) -> Url {
        let path = path.as_ref().canonicalize().unwrap();
        Url::from_file_path(path).unwrap()
    }

    #[test]
    fn test_parse_imported() {
        let res = parse_imported("../testdata/my_interface.mojom");
        assert!(res.is_ok());
    }

    #[test]
    fn test_check_imports() {
        let root_path = "../testdata";
        let file_path = "../testdata/my_service.mojom";
        let mut text = String::new();
        File::open(&file_path)
            .unwrap()
            .read_to_string(&mut text)
            .unwrap();
        let uri = create_uri(&file_path);
        let ast = MojomAst::new(uri, text).unwrap();

        let imports = check_imports(&root_path, &ast);

        let res = imports.find_definition("FooStruct.FooEnum");
        assert!(res.is_some());
    }
}
<|endoftext|>


<commit_before>/**
 *   Mles server peer
 *
 *   Copyright (C) 2017  Juhamatti Kuusisaari / Mles developers
 *
 *   This program is free software: you can redistribute it and/or modify
 *   it under the terms of the GNU General Public License as published by
 *   the Free Software Foundation, either version 3 of the License, or
 *   (at your option) any later version.
 *
 *   This program is distributed in the hope that it will be useful,
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *   GNU General Public License for more details.
 *
 *   You should have received a copy of the GNU General Public License
 *   along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
extern crate tokio_core;
extern crate tokio_io;
extern crate futures;

use std::rc::Rc;
use std::cell::RefCell;
use std::iter;
use std::io::Error;
use std::net::SocketAddr;
use std::thread;
use std::time::Duration;

use tokio_core::net::TcpStream;
use tokio_core::reactor::Core;
use tokio_io::io;
use tokio_io::AsyncRead;

use futures::Future;
use futures::stream::{self, Stream};
use futures::sync::mpsc::{unbounded, UnboundedSender};

use local_db::*;
use frame::*;

const HDRL: usize = 4; //hdr len
const KEYL: usize = 8; //key len
const HDRKEYL: usize = HDRL + KEYL;

const MAXWAIT: u64 = 10*60;
const WAITTIME: u64 = 5;

pub fn peer_conn(peer: SocketAddr, peer_cnt: u64, channel: String, msg: Vec<u8>,
                 tx_peer_for_msgs: UnboundedSender<(u64, String, UnboundedSender<Vec<u8>>, UnboundedSender<UnboundedSender<Vec<u8>>>)>)
{
    let mut core = Core::new().unwrap();
    let waittime = WAITTIME;
    let mut loopcnt = 1;

    println!("Peer channel thread for channel {}", channel);
    loop {
        let mles_peer_db = Rc::new(RefCell::new(MlesPeerDb::new()));
        let handle = core.handle();
        let channel = channel.clone();
        let channel2 = channel.clone();
        let tx_peer_for_msgs = tx_peer_for_msgs.clone();
        let msg = msg.clone();

        let tcp = TcpStream::connect(&peer, &handle);

        let (tx_orig_chan, rx_orig_chan) = unbounded();

        let client = tcp.and_then(move |pstream| {
            let _val = pstream.set_nodelay(true).map_err(|_| panic!("Cannot set peer to no delay"));
            println!("Successfully connected to peer");

            //save writes to db
            let (reader, writer) = pstream.split();
            let (tx, rx) = unbounded();
            let _res = tx_peer_for_msgs.send((peer_cnt, channel, tx.clone(), tx_orig_chan.clone())).map_err(|err| { println!("Cannot send from peer: {}", err); () });
            let _res = tx.send(msg).map_err(|err| { println!("Cannot write to tx: {}", err); });

            let mles_peer_db_inner = mles_peer_db.clone();
            let psocket_writer = rx.fold(writer, move |writer, msg| {
                //push message to history
                let mut mles_peer_db = mles_peer_db_inner.borrow_mut();
                mles_peer_db.add_message(msg.clone());

                //send message forward
                let amt = io::write_all(writer, msg);
                let amt = amt.map(|(writer, _)| writer);
                amt.map_err(|_| ())
            });
            handle.spawn(psocket_writer.then(|_| {
                Ok(())
            }));

            let mles_peer_db_inner = mles_peer_db.clone();
            let tx_origs_reader = rx_orig_chan.for_each(move |tx_orig| {
                //save receiver side tx to db
                let mut mles_peer_db_once = mles_peer_db_inner.borrow_mut();
                mles_peer_db_once.add_channel(tx_orig.clone());

                //push history to client if not the first one (as peer will send the history then)
                if mles_peer_db_once.get_messages_len() > 1 {
                    for msg in mles_peer_db_once.get_messages().iter() {
                        let _res = tx_orig.send(msg.clone()).map_err(|err| { println!("Failed to send history from peer: {}", err); () });
                    }
                }
                Ok(())
            });
            handle.spawn(tx_origs_reader.then(|_| {
                Ok(())
            }));

            let mles_peer_db_inner = mles_peer_db.clone();
            let iter = stream::iter(iter::repeat(()).map(Ok::<(), Error>));
            iter.fold(reader, move |reader, _| {
                let frame = io::read_exact(reader, vec![0;HDRKEYL]);
                let frame = frame.and_then(move |(reader, hdr_key)| process_hdr_dummy_key(reader, hdr_key));

                let frame = frame.and_then(move |(reader, hdr_key, hdr_len)| {
                    let tframe = io::read_exact(reader, vec![0;hdr_len]);
                    tframe.and_then(move |(reader, message)| process_msg(reader, hdr_key, message))
                });

                let mles_peer_db_frame = mles_peer_db_inner.clone();
                frame.map(move |(reader, mut hdr_key, message)| {
                    hdr_key.extend(message);

                    //send message forward
                    let mut mles_peer_db = mles_peer_db_frame.borrow_mut();
                    for tx_orig in mles_peer_db.get_channels().iter() {
                        let _res = tx_orig.send(hdr_key.clone()).map_err(|err| { println!("Failed to send from peer: {}", err); () });
                    }
                    //push message to history
                    mles_peer_db.add_message(hdr_key);

                    reader
                })
            })
        });

        // execute server
        let _res = core.run(client).map_err(|err| { println!("Peer: {}", err); () });

        let mut wait = waittime * loopcnt;
        if wait > MAXWAIT {
            wait = MAXWAIT;
        }
        loopcnt *= 2;

        println!("Connection for channel {} failed. Please check for proper key. Retrying in {} s.", channel2, wait);
        thread::sleep(Duration::from_secs(wait));
    }
}

pub fn inc_peer_cnt(cnt: u64) -> u64 {
    let mut val = cnt;
    val = val >> 32;
    val += 1;
    val << 32
}

pub fn has_peer(peer: &SocketAddr) -> bool {
    0 != peer.port()
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::{IpAddr, Ipv4Addr};

    #[test]
    fn test_peer_inc_cnt() {
        let val: u64 = 1 << 32;
        assert_eq!(2 << 32, inc_peer_cnt(val));
    }

    #[test]
    fn test_has_peer() {
        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0);
        assert_eq!(false, has_peer(&addr));
    }
}
<commit_msg>Avoid losing first messages in case TCP connection peer is not established immediately. (#62)<commit_after>/**
 *   Mles server peer
 *
 *   Copyright (C) 2017  Juhamatti Kuusisaari / Mles developers
 *
 *   This program is free software: you can redistribute it and/or modify
 *   it under the terms of the GNU General Public License as published by
 *   the Free Software Foundation, either version 3 of the License, or
 *   (at your option) any later version.
 *
 *   This program is distributed in the hope that it will be useful,
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *   GNU General Public License for more details.
 *
 *   You should have received a copy of the GNU General Public License
 *   along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
extern crate tokio_core;
extern crate tokio_io;
extern crate futures;

use std::rc::Rc;
use std::cell::RefCell;
use std::iter;
use std::io::Error;
use std::net::SocketAddr;
use std::thread;
use std::time::Duration;

use tokio_core::net::TcpStream;
use tokio_core::reactor::Core;
use tokio_io::io;
use tokio_io::AsyncRead;

use futures::Future;
use futures::stream::{self, Stream};
use futures::sync::mpsc::{unbounded, UnboundedSender};

use local_db::*;
use frame::*;

const HDRL: usize = 4; //hdr len
const KEYL: usize = 8; //key len
const HDRKEYL: usize = HDRL + KEYL;

const MAXWAIT: u64 = 10*60;
const WAITTIME: u64 = 5;

pub fn peer_conn(peer: SocketAddr, peer_cnt: u64, channel: String, msg: Vec<u8>,
                 tx_peer_for_msgs: UnboundedSender<(u64, String, UnboundedSender<Vec<u8>>, UnboundedSender<UnboundedSender<Vec<u8>>>)>)
{
    let mut core = Core::new().unwrap();
    let waittime = WAITTIME;
    let mut loopcnt = 1;

    println!("Peer channel thread for channel {}", channel);
    loop {
        let mles_peer_db = Rc::new(RefCell::new(MlesPeerDb::new()));
        let handle = core.handle();
        let channel = channel.clone();
        let channel2 = channel.clone();
        let tx_peer_for_msgs = tx_peer_for_msgs.clone();
        let msg = msg.clone();

        let tcp = TcpStream::connect(&peer, &handle);

        let (tx_orig_chan, rx_orig_chan) = unbounded();
        let (tx, rx) = unbounded();

        //distribute channels
        let _res = tx_peer_for_msgs.send((peer_cnt, channel, tx.clone(), tx_orig_chan.clone())).map_err(|err| { println!("Cannot send from peer: {}", err); () });
        let _res = tx.send(msg).map_err(|err| { println!("Cannot write to tx: {}", err); });

        let client = tcp.and_then(move |pstream| {
            let _val = pstream.set_nodelay(true).map_err(|_| panic!("Cannot set peer to no delay"));
            println!("Successfully connected to peer");

            let (reader, writer) = pstream.split();

            let mles_peer_db_inner = mles_peer_db.clone();
            let psocket_writer = rx.fold(writer, move |writer, msg| {
                //push message to history
                let mut mles_peer_db = mles_peer_db_inner.borrow_mut();
                mles_peer_db.add_message(msg.clone());

                //send message forward
                let amt = io::write_all(writer, msg);
                let amt = amt.map(|(writer, _)| writer);
                amt.map_err(|_| ())
            });
            handle.spawn(psocket_writer.then(|_| {
                Ok(())
            }));

            let mles_peer_db_inner = mles_peer_db.clone();
            let tx_origs_reader = rx_orig_chan.for_each(move |tx_orig| {
                //save receiver side tx to db
                let mut mles_peer_db_once = mles_peer_db_inner.borrow_mut();
                mles_peer_db_once.add_channel(tx_orig.clone());

                //push history to client if not the first one (as peer will send the history then)
                if mles_peer_db_once.get_messages_len() > 1 {
                    for msg in mles_peer_db_once.get_messages().iter() {
                        let _res = tx_orig.send(msg.clone()).map_err(|err| { println!("Failed to send history from peer: {}", err); () });
                    }
                }
                Ok(())
            });
            handle.spawn(tx_origs_reader.then(|_| {
                Ok(())
            }));

            let mles_peer_db_inner = mles_peer_db.clone();
            let iter = stream::iter(iter::repeat(()).map(Ok::<(), Error>));
            iter.fold(reader, move |reader, _| {
                let frame = io::read_exact(reader, vec![0;HDRKEYL]);
                let frame = frame.and_then(move |(reader, hdr_key)| process_hdr_dummy_key(reader, hdr_key));

                let frame = frame.and_then(move |(reader, hdr_key, hdr_len)| {
                    let tframe = io::read_exact(reader, vec![0;hdr_len]);
                    tframe.and_then(move |(reader, message)| process_msg(reader, hdr_key, message))
                });

                let mles_peer_db_frame = mles_peer_db_inner.clone();
                frame.map(move |(reader, mut hdr_key, message)| {
                    hdr_key.extend(message);

                    //send message forward
                    let mut mles_peer_db = mles_peer_db_frame.borrow_mut();
                    for tx_orig in mles_peer_db.get_channels().iter() {
                        let _res = tx_orig.send(hdr_key.clone()).map_err(|err| { println!("Failed to send from peer: {}", err); () });
                    }
                    //push message to history
                    mles_peer_db.add_message(hdr_key);

                    reader
                })
            })
        });

        // execute server
        let _res = core.run(client).map_err(|err| { println!("Peer: {}", err); () });

        let mut wait = waittime * loopcnt;
        if wait > MAXWAIT {
            wait = MAXWAIT;
        }
        loopcnt *= 2;

        println!("Connection for channel {} failed. Please check for proper key. Retrying in {} s.", channel2, wait);
        thread::sleep(Duration::from_secs(wait));
    }
}

pub fn inc_peer_cnt(cnt: u64) -> u64 {
    let mut val = cnt;
    val = val >> 32;
    val += 1;
    val << 32
}

pub fn has_peer(peer: &SocketAddr) -> bool {
    0 != peer.port()
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::{IpAddr, Ipv4Addr};

    #[test]
    fn test_peer_inc_cnt() {
        let val: u64 = 1 << 32;
        assert_eq!(2 << 32, inc_peer_cnt(val));
    }

    #[test]
    fn test_has_peer() {
        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0);
        assert_eq!(false, has_peer(&addr));
    }
}
<|endoftext|>