<commit_before>extern crate mws;
use chrono::offset::TimeZone;
use chrono::{DateTime, NaiveDate, Utc};
use std::path::PathBuf;
use structopt::StructOpt;

use mws::client::{Client, Method};

mod env;

use self::env::Env;

#[derive(Debug, StructOpt)]
#[structopt(name = "mws-cli", about = "MWS CLI.")]
struct Opt {
  #[structopt(long = "env", parse(from_os_str))]
  env: Option<PathBuf>,
  #[structopt(subcommand)]
  cmd: Command,
}

#[derive(Debug, StructOpt)]
enum Command {
  ReportRequest {
    #[structopt(long = "report_type")]
    report_type: String,
    #[structopt(long = "start_date", parse(try_from_str))]
    start_date: Option<NaiveDate>,
    #[structopt(long = "end_date", parse(try_from_str))]
    end_date: Option<NaiveDate>,
    #[structopt(long = "report_options")]
    report_options: Option<String>,
    #[structopt(long = "marketplace")]
    marketplace_id_list: Option<Vec<String>>,
  },
  ReportListRequestByIds {
    #[structopt(long = "id")]
    ids: Vec<String>,
  },
  ReportListRequestByTypes {
    #[structopt(long = "type")]
    types: Vec<String>,
  },
  ReportGet {
    #[structopt(long = "id")]
    id: String,
    #[structopt(long = "out", parse(from_os_str))]
    out: PathBuf,
  },
  EncodingConvJp {
    #[structopt(long = "in", parse(from_os_str))]
    input: PathBuf,
    #[structopt(long = "out", parse(from_os_str))]
    out: PathBuf,
  },
  ProductGetLowestPricedOffersForSKU {
    #[structopt(long = "marketplace")]
    marketplace_id: String,
    #[structopt(long = "seller_sku")]
    seller_sku: String,
    #[structopt(long = "condition")]
    condition: String,
  },
  ProductGetMyPriceForASIN {
    #[structopt(long = "marketplace")]
    marketplace_id: String,
    #[structopt(long = "asin")]
    asins: Vec<String>,
    #[structopt(long = "condition")]
    condition: Option<String>,
  },
  SubmitFeed {
    #[structopt(long = "feed_type")]
    feed_type: String,
    #[structopt(long = "contant_file", parse(from_os_str))]
    content_file: PathBuf,
    #[structopt(long = "marketplace")]
    marketplace_id_list: Option<Vec<String>>,
    #[structopt(long = "content_type")]
    content_type: String,
  },
  ListSubscriptions {
    #[structopt(long = "marketplace")]
    marketplace_id: String,
  },
  ListFinancialEvents {
    #[structopt(long = "posted_after", parse(try_from_str))]
    posted_after: NaiveDate,
    #[structopt(long = "posted_before", parse(try_from_str))]
    posted_before: NaiveDate,
    #[structopt(long = "out", parse(from_os_str))]
    outdir: PathBuf,
  },
  ListOrders {
    #[structopt(long = "marketplace")]
    marketplace_id_list: Vec<String>,
    #[structopt(long = "created_after", parse(try_from_str))]
    created_after: NaiveDate,
    #[structopt(long = "out", parse(from_os_str))]
    outdir: PathBuf,
  }
}

fn main() {
  let opt = Opt::from_args();

  let env_path = opt.env.unwrap_or_else(|| PathBuf::from(".env"));

  println!("env path: {:?}", env_path);
  ::dotenv::from_path(&env_path).ok();

  let env = Env::from_env();

  println!("seller id: {}", env.seller_id);
  println!("region id: {}", env.region_id);

  let client = get_client(&env);

  match opt.cmd {
    Command::ReportRequest {
      report_type,
      start_date,
      end_date,
      report_options,
      marketplace_id_list,
    } => {
      use mws::reports::*;
      let res = RequestReport(
        &client,
        RequestReportParameters {
          ReportType: report_type,
          StartDate: start_date.map(get_utc_datetime),
          EndDate: end_date.map(get_utc_datetime),
          ReportOptions: report_options,
          MarketplaceIdList: marketplace_id_list,
        },
      )
      .unwrap();
      println!("{:#?}", res)
    }
    Command::ReportListRequestByIds { ids } => {
      use mws::reports::*;
      let res = GetReportRequestList(
        &client,
        GetReportRequestListParameters {
          ReportRequestIdList: Some(ids),
          ..Default::default()
        },
      )
      .unwrap();

      println!("{:#?}", res)
    }
    Command::ReportListRequestByTypes { types } => {
      use mws::reports::*;
      let res = GetReportRequestList(
        &client,
        GetReportRequestListParameters {
          ReportTypeList: Some(types),
          ..Default::default()
        },
      )
      .unwrap();

      println!("{:#?}", res)
    }
    Command::ReportGet { id, out } => {
      use mws::reports::*;
      let mut out = std::fs::File::create(out).unwrap();
      GetReport(&client, id, &mut out).unwrap();
    }
    Command::EncodingConvJp { input, out } => {
      use encoding_rs::*;
      let bytes = std::fs::read(input).unwrap();
      let (cow, encoding_used, had_errors) = SHIFT_JIS.decode(&bytes);
      if had_errors {
        panic!("decode error.")
      }
      println!("encoding_used: {:?}", encoding_used);
      std::fs::write(out, cow.as_ref()).unwrap();
    }
    Command::ProductGetLowestPricedOffersForSKU {
      marketplace_id,
      seller_sku,
      condition,
    } => {
      use mws::products::*;
      let res = GetLowestPricedOffersForSKU(
        &client,
        GetLowestPricedOffersForSKUParameters {
          MarketplaceId: marketplace_id,
          SellerSKU: seller_sku,
          ItemCondition: ItemCondition::from(&condition as &str),
        },
      )
      .unwrap();
      println!("{:#?}", res)
    }
    Command::ProductGetMyPriceForASIN {
      marketplace_id,
      asins,
      condition,
    } => {
      use mws::products::*;
      let res = GetMyPriceForASIN(
        &client,
        GetMyPriceForASINParameters {
          MarketplaceId: marketplace_id,
          ASINList: asins,
          ItemCondition: condition.map(|v| ItemCondition::from(&v as &str)),
        },
      )
      .unwrap();
      println!("{:#?}", res)
    }
    Command::SubmitFeed {
      feed_type,
      content_file,
      marketplace_id_list,
      content_type,
    } => {
      use mws::feeds::*;
      use std::io::Cursor;
      let content = std::fs::read(content_file).unwrap();
      let digest = md5::compute(&content);
      let b64 = base64::encode(&*digest);
      let res = SubmitFeed(
        &client,
        SubmitFeedParameters {
          FeedType: feed_type,
          MarketplaceIdList: marketplace_id_list,
          PurgeAndReplace: None,
        },
        Cursor::new(content),
        b64,
        content_type,
      )
      .unwrap();
      println!("{:#?}", res)
    }
    Command::ListSubscriptions { marketplace_id } => {
      use mws::subscriptions::*;
      let res = ListSubscriptions(&client, marketplace_id).unwrap();
      println!("{:#?}", res)
    }
    Command::ListFinancialEvents {
      posted_after,
      posted_before,
      outdir,
    } => {
      use std::thread::sleep;
      use std::time::Duration;
      let mut page = 1;
      let mut next_token: Option<String> = None;

      let duration = posted_before - posted_after;
      let days = duration.num_days();

      if days < 0 {
        panic!("invalid date range")
      }

      let chunks = ((days as f64) / 180.0).ceil() as i64;
      let ranges: Vec<[NaiveDate; 2]> = (0..chunks).into_iter().map(|chunk_idx| {
        [posted_after + chrono::Duration::days(chunk_idx * 180), std::cmp::min(posted_before, posted_after + chrono::Duration::days((chunk_idx + 1) * 180))]
      }).collect();

      println!("ranges = {:?}", ranges);

      for [posted_after, posted_before] in ranges {
        loop {
          println!("loading page {} ...", page);

          let res = if let Some(next_token) = next_token.take() {
            client
              .request_xml_generic(
                Method::POST,
                "/Finances",
                "2015-05-01",
                "ListFinancialEventsByNextToken",
                vec![
                  (
                    "NextToken".to_string(),
                    next_token,
                  ),
                ],
              )
          } else {
            client
              .request_xml_generic(
                Method::POST,
                "/Finances",
                "2015-05-01",
                "ListFinancialEvents",
                vec![
                  (
                    "PostedAfter".to_string(),
                    get_utc_datetime(posted_after).to_rfc3339(),
                  ),
                  (
                    "PostedBefore".to_string(),
                    get_utc_datetime(posted_before).to_rfc3339(),
                  ),
                ],
              )
          };

          let res = match res {
            Ok(res) => res,
            Err(err) => {
              if err.should_try_again() {
                eprintln!("RetriableException: {:?}", err);
                continue;
              } else {
                Err(err).unwrap()
              }
            }
          };

          let filename = format!("financial_events_{}_{}_{}.xml", posted_after, posted_before, page);
          let f = std::fs::File::create(outdir.join(filename)).unwrap();
          res.result_element.write(f).unwrap();

          next_token = res.next_token().clone().map(|v| v.to_string());
          if next_token.is_none() {
            break;
          }

          page += 1;

          sleep(Duration::from_secs(1));
        }
      }
    }
    Command::ListOrders {
      marketplace_id_list,
      created_after,
      outdir
    } => {
      use std::thread::sleep;
      use std::time::Duration;
      use mws::orders::*;

      let mut page = 1;
      let mut next_token = None;

      loop {
        println!("loading page = {}", page);

        let res = if let Some(next_token) = next_token.clone() {
          ListOrdersByNextToken(&client, next_token)
        } else {
          ListOrders(&client, ListOrdersParameters {
            MarketplaceId: marketplace_id_list.clone(),
            CreatedAfter: Some(DateTime::<Utc>::from_utc(created_after.and_hms(0, 0, 0).into(), Utc)),
            MaxResultsPerPage: Some(100),
            ..Default::default()
          })
        };

        match res {
          Ok(res) =>{
            next_token = res.NextToken;
            let orders = res.Orders;
            let mut json_orders = vec![];
            println!("orders = {}", orders.len());

            for order in orders {
              println!("loading items: {} {:?}", order.AmazonOrderId, order.PurchaseDate);

              let mut next_token = None;
              let mut items = vec![];
              loop {
                let res = if let Some(next_token) = next_token.clone() {
                  ListOrderItemsByNextToken(&client, next_token)
                } else {
                  ListOrderItems(&client, order.AmazonOrderId.clone())
                };

                match res {
                  Ok(res) => {
                    next_token = res.NextToken;
                    println!("items = {}", res.OrderItems.len());
                    items.extend(res.OrderItems);
                  }
                  Err(err) => {
                    if err.should_try_again() {
                      sleep(Duration::from_secs(10));
                      continue;
                    } else {
                      Err(err).unwrap()
                    }
                  }
                }

                sleep(Duration::from_secs(1));

                if next_token.is_none() {
                  break;
                }
              }

              json_orders.push(serde_json::json!({
                "order": serde_json::to_value(&order).unwrap(),
                "items": serde_json::to_value(&items).unwrap()
              }));
            }

            let filename = format!("orders_{}.json", page);
            let f = std::fs::File::create(outdir.join(filename)).unwrap();
            serde_json::to_writer(f, &json_orders).unwrap();
          },
          Err(err) => {
            if err.should_try_again() {
              sleep(Duration::from_secs(10));
              continue;
            } else {
              Err(err).unwrap()
            }
          }
        }

        if next_token.is_none() {
          break;
        }

        page += 1;
        sleep(Duration::from_secs(5));
      }
    }
  }
}

fn get_client(env: &Env) -> Client {
  use mws::client::ClientOptions;
  use mws::constants;
  let region = constants::get_region(&env.region_id).expect("invalid region id");
  let opts = ClientOptions {
    endpoint: region.endpoint.to_string(),
    seller_id: env.seller_id.clone(),
    mws_auth_token: env.auth_token.clone(),
    aws_access_key_id: env.access_key_id.clone(),
    secret_key: env.secret_key.clone(),
  };
  Client::new(opts).unwrap()
}

fn get_utc_datetime(date: NaiveDate) -> DateTime<Utc> {
  Utc.from_utc_date(&date).and_hms(0, 0, 0)
}
<commit_msg>do not remove next_token when RetriableException<commit_after>extern crate mws;
use chrono::offset::TimeZone;
use chrono::{DateTime, NaiveDate, Utc};
use std::path::PathBuf;
use structopt::StructOpt;

use mws::client::{Client, Method};

mod env;

use self::env::Env;

#[derive(Debug, StructOpt)]
#[structopt(name = "mws-cli", about = "MWS CLI.")]
struct Opt {
  #[structopt(long = "env", parse(from_os_str))]
  env: Option<PathBuf>,
  #[structopt(subcommand)]
  cmd: Command,
}

#[derive(Debug, StructOpt)]
enum Command {
  ReportRequest {
    #[structopt(long = "report_type")]
    report_type: String,
    #[structopt(long = "start_date", parse(try_from_str))]
    start_date: Option<NaiveDate>,
    #[structopt(long = "end_date", parse(try_from_str))]
    end_date: Option<NaiveDate>,
    #[structopt(long = "report_options")]
    report_options: Option<String>,
    #[structopt(long = "marketplace")]
    marketplace_id_list: Option<Vec<String>>,
  },
  ReportListRequestByIds {
    #[structopt(long = "id")]
    ids: Vec<String>,
  },
  ReportListRequestByTypes {
    #[structopt(long = "type")]
    types: Vec<String>,
  },
  ReportGet {
    #[structopt(long = "id")]
    id: String,
    #[structopt(long = "out", parse(from_os_str))]
    out: PathBuf,
  },
  EncodingConvJp {
    #[structopt(long = "in", parse(from_os_str))]
    input: PathBuf,
    #[structopt(long = "out", parse(from_os_str))]
    out: PathBuf,
  },
  ProductGetLowestPricedOffersForSKU {
    #[structopt(long = "marketplace")]
    marketplace_id: String,
    #[structopt(long = "seller_sku")]
    seller_sku: String,
    #[structopt(long = "condition")]
    condition: String,
  },
  ProductGetMyPriceForASIN {
    #[structopt(long = "marketplace")]
    marketplace_id: String,
    #[structopt(long = "asin")]
    asins: Vec<String>,
    #[structopt(long = "condition")]
    condition: Option<String>,
  },
  SubmitFeed {
    #[structopt(long = "feed_type")]
    feed_type: String,
    #[structopt(long = "contant_file", parse(from_os_str))]
    content_file: PathBuf,
    #[structopt(long = "marketplace")]
    marketplace_id_list: Option<Vec<String>>,
    #[structopt(long = "content_type")]
    content_type: String,
  },
  ListSubscriptions {
    #[structopt(long = "marketplace")]
    marketplace_id: String,
  },
  ListFinancialEvents {
    #[structopt(long = "posted_after", parse(try_from_str))]
    posted_after: NaiveDate,
    #[structopt(long = "posted_before", parse(try_from_str))]
    posted_before: NaiveDate,
    #[structopt(long = "out", parse(from_os_str))]
    outdir: PathBuf,
  },
  ListOrders {
    #[structopt(long = "marketplace")]
    marketplace_id_list: Vec<String>,
    #[structopt(long = "created_after", parse(try_from_str))]
    created_after: NaiveDate,
    #[structopt(long = "out", parse(from_os_str))]
    outdir: PathBuf,
  }
}

fn main() {
  let opt = Opt::from_args();

  let env_path = opt.env.unwrap_or_else(|| PathBuf::from(".env"));

  println!("env path: {:?}", env_path);
  ::dotenv::from_path(&env_path).ok();

  let env = Env::from_env();

  println!("seller id: {}", env.seller_id);
  println!("region id: {}", env.region_id);

  let client = get_client(&env);

  match opt.cmd {
    Command::ReportRequest {
      report_type,
      start_date,
      end_date,
      report_options,
      marketplace_id_list,
    } => {
      use mws::reports::*;
      let res = RequestReport(
        &client,
        RequestReportParameters {
          ReportType: report_type,
          StartDate: start_date.map(get_utc_datetime),
          EndDate: end_date.map(get_utc_datetime),
          ReportOptions: report_options,
          MarketplaceIdList: marketplace_id_list,
        },
      )
      .unwrap();
      println!("{:#?}", res)
    }
    Command::ReportListRequestByIds { ids } => {
      use mws::reports::*;
      let res = GetReportRequestList(
        &client,
        GetReportRequestListParameters {
          ReportRequestIdList: Some(ids),
          ..Default::default()
        },
      )
      .unwrap();

      println!("{:#?}", res)
    }
    Command::ReportListRequestByTypes { types } => {
      use mws::reports::*;
      let res = GetReportRequestList(
        &client,
        GetReportRequestListParameters {
          ReportTypeList: Some(types),
          ..Default::default()
        },
      )
      .unwrap();

      println!("{:#?}", res)
    }
    Command::ReportGet { id, out } => {
      use mws::reports::*;
      let mut out = std::fs::File::create(out).unwrap();
      GetReport(&client, id, &mut out).unwrap();
    }
    Command::EncodingConvJp { input, out } => {
      use encoding_rs::*;
      let bytes = std::fs::read(input).unwrap();
      let (cow, encoding_used, had_errors) = SHIFT_JIS.decode(&bytes);
      if had_errors {
        panic!("decode error.")
      }
      println!("encoding_used: {:?}", encoding_used);
      std::fs::write(out, cow.as_ref()).unwrap();
    }
    Command::ProductGetLowestPricedOffersForSKU {
      marketplace_id,
      seller_sku,
      condition,
    } => {
      use mws::products::*;
      let res = GetLowestPricedOffersForSKU(
        &client,
        GetLowestPricedOffersForSKUParameters {
          MarketplaceId: marketplace_id,
          SellerSKU: seller_sku,
          ItemCondition: ItemCondition::from(&condition as &str),
        },
      )
      .unwrap();
      println!("{:#?}", res)
    }
    Command::ProductGetMyPriceForASIN {
      marketplace_id,
      asins,
      condition,
    } => {
      use mws::products::*;
      let res = GetMyPriceForASIN(
        &client,
        GetMyPriceForASINParameters {
          MarketplaceId: marketplace_id,
          ASINList: asins,
          ItemCondition: condition.map(|v| ItemCondition::from(&v as &str)),
        },
      )
      .unwrap();
      println!("{:#?}", res)
    }
    Command::SubmitFeed {
      feed_type,
      content_file,
      marketplace_id_list,
      content_type,
    } => {
      use mws::feeds::*;
      use std::io::Cursor;
      let content = std::fs::read(content_file).unwrap();
      let digest = md5::compute(&content);
      let b64 = base64::encode(&*digest);
      let res = SubmitFeed(
        &client,
        SubmitFeedParameters {
          FeedType: feed_type,
          MarketplaceIdList: marketplace_id_list,
          PurgeAndReplace: None,
        },
        Cursor::new(content),
        b64,
        content_type,
      )
      .unwrap();
      println!("{:#?}", res)
    }
    Command::ListSubscriptions { marketplace_id } => {
      use mws::subscriptions::*;
      let res = ListSubscriptions(&client, marketplace_id).unwrap();
      println!("{:#?}", res)
    }
    Command::ListFinancialEvents {
      posted_after,
      posted_before,
      outdir,
    } => {
      use std::thread::sleep;
      use std::time::Duration;
      let mut page = 1;
      let mut next_token: Option<String> = None;

      let duration = posted_before - posted_after;
      let days = duration.num_days();

      if days < 0 {
        panic!("invalid date range")
      }

      let chunks = ((days as f64) / 180.0).ceil() as i64;
      let ranges: Vec<[NaiveDate; 2]> = (0..chunks).into_iter().map(|chunk_idx| {
        [posted_after + chrono::Duration::days(chunk_idx * 180), std::cmp::min(posted_before, posted_after + chrono::Duration::days((chunk_idx + 1) * 180))]
      }).collect();

      println!("ranges = {:?}", ranges);

      for [posted_after, posted_before] in ranges {
        loop {
          println!("loading page {} ...", page);

          let res = if let Some(next_token) = next_token.clone() {
            client
              .request_xml_generic(
                Method::POST,
                "/Finances",
                "2015-05-01",
                "ListFinancialEventsByNextToken",
                vec![
                  (
                    "NextToken".to_string(),
                    next_token,
                  ),
                ],
              )
          } else {
            client
              .request_xml_generic(
                Method::POST,
                "/Finances",
                "2015-05-01",
                "ListFinancialEvents",
                vec![
                  (
                    "PostedAfter".to_string(),
                    get_utc_datetime(posted_after).to_rfc3339(),
                  ),
                  (
                    "PostedBefore".to_string(),
                    get_utc_datetime(posted_before).to_rfc3339(),
                  ),
                ],
              )
          };

          let res = match res {
            Ok(res) => res,
            Err(err) => {
              if err.should_try_again() {
                eprintln!("RetriableException: {:?}", err);
                continue;
              } else {
                Err(err).unwrap()
              }
            }
          };

          let filename = format!("financial_events_{}_{}_{}.xml", posted_after, posted_before, page);
          let f = std::fs::File::create(outdir.join(filename)).unwrap();
          res.result_element.write(f).unwrap();

          next_token = res.next_token().clone().map(|v| v.to_string());
          if next_token.is_none() {
            break;
          }

          page += 1;

          sleep(Duration::from_secs(1));
        }
      }
    }
    Command::ListOrders {
      marketplace_id_list,
      created_after,
      outdir
    } => {
      use std::thread::sleep;
      use std::time::Duration;
      use mws::orders::*;

      let mut page = 1;
      let mut next_token = None;

      loop {
        println!("loading page = {}", page);

        let res = if let Some(next_token) = next_token.clone() {
          ListOrdersByNextToken(&client, next_token)
        } else {
          ListOrders(&client, ListOrdersParameters {
            MarketplaceId: marketplace_id_list.clone(),
            CreatedAfter: Some(DateTime::<Utc>::from_utc(created_after.and_hms(0, 0, 0).into(), Utc)),
            MaxResultsPerPage: Some(100),
            ..Default::default()
          })
        };

        match res {
          Ok(res) =>{
            next_token = res.NextToken;
            let orders = res.Orders;
            let mut json_orders = vec![];
            println!("orders = {}", orders.len());

            for order in orders {
              println!("loading items: {} {:?}", order.AmazonOrderId, order.PurchaseDate);

              let mut next_token = None;
              let mut items = vec![];
              loop {
                let res = if let Some(next_token) = next_token.clone() {
                  ListOrderItemsByNextToken(&client, next_token)
                } else {
                  ListOrderItems(&client, order.AmazonOrderId.clone())
                };

                match res {
                  Ok(res) => {
                    next_token = res.NextToken;
                    println!("items = {}", res.OrderItems.len());
                    items.extend(res.OrderItems);
                  }
                  Err(err) => {
                    if err.should_try_again() {
                      sleep(Duration::from_secs(10));
                      continue;
                    } else {
                      Err(err).unwrap()
                    }
                  }
                }

                sleep(Duration::from_secs(1));

                if next_token.is_none() {
                  break;
                }
              }

              json_orders.push(serde_json::json!({
                "order": serde_json::to_value(&order).unwrap(),
                "items": serde_json::to_value(&items).unwrap()
              }));
            }

            let filename = format!("orders_{}.json", page);
            let f = std::fs::File::create(outdir.join(filename)).unwrap();
            serde_json::to_writer(f, &json_orders).unwrap();
          },
          Err(err) => {
            if err.should_try_again() {
              sleep(Duration::from_secs(10));
              continue;
            } else {
              Err(err).unwrap()
            }
          }
        }

        if next_token.is_none() {
          break;
        }

        page += 1;
        sleep(Duration::from_secs(5));
      }
    }
  }
}

fn get_client(env: &Env) -> Client {
  use mws::client::ClientOptions;
  use mws::constants;
  let region = constants::get_region(&env.region_id).expect("invalid region id");
  let opts = ClientOptions {
    endpoint: region.endpoint.to_string(),
    seller_id: env.seller_id.clone(),
    mws_auth_token: env.auth_token.clone(),
    aws_access_key_id: env.access_key_id.clone(),
    secret_key: env.secret_key.clone(),
  };
  Client::new(opts).unwrap()
}

fn get_utc_datetime(date: NaiveDate) -> DateTime<Utc> {
  Utc.from_utc_date(&date).and_hms(0, 0, 0)
}
<|endoftext|>


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

//! Transaction pool implementation leveraging txhashset for chain state
//! validation. It is a valid operation to add a tx to the tx pool if the
//! resulting tx pool can be added to the current chain state to produce a
//! valid chain state.

use std::collections::VecDeque;
use std::sync::Arc;
use util::RwLock;

use chrono::prelude::Utc;

use core::core::hash::{Hash, Hashed};
use core::core::id::ShortId;
use core::core::verifier_cache::VerifierCache;
use core::core::{transaction, Block, BlockHeader, Transaction};
use pool::Pool;
use types::{BlockChain, PoolAdapter, PoolConfig, PoolEntry, PoolEntryState, PoolError, TxSource};

// Cache this many txs to handle a potential fork and re-org.
const REORG_CACHE_SIZE: usize = 100;

/// Transaction pool implementation.
pub struct TransactionPool {
	/// Pool Config
	pub config: PoolConfig,
	/// Our transaction pool.
	pub txpool: Pool,
	/// Our Dandelion "stempool".
	pub stempool: Pool,
	/// Cache of previous txs in case of a re-org.
	pub reorg_cache: Arc<RwLock<VecDeque<PoolEntry>>>,
	/// The blockchain
	pub blockchain: Arc<BlockChain>,
	pub verifier_cache: Arc<RwLock<VerifierCache>>,
	/// The pool adapter
	pub adapter: Arc<PoolAdapter>,
}

impl TransactionPool {
	/// Create a new transaction pool
	pub fn new(
		config: PoolConfig,
		chain: Arc<BlockChain>,
		verifier_cache: Arc<RwLock<VerifierCache>>,
		adapter: Arc<PoolAdapter>,
	) -> TransactionPool {
		TransactionPool {
			config,
			txpool: Pool::new(chain.clone(), verifier_cache.clone(), "txpool".to_string()),
			stempool: Pool::new(
				chain.clone(),
				verifier_cache.clone(),
				"stempool".to_string(),
			),
			reorg_cache: Arc::new(RwLock::new(VecDeque::new())),
			blockchain: chain,
			verifier_cache,
			adapter,
		}
	}

	pub fn chain_head(&self) -> Result<BlockHeader, PoolError> {
		self.blockchain.chain_head()
	}

	fn add_to_stempool(&mut self, entry: PoolEntry, header: &BlockHeader) -> Result<(), PoolError> {
		// Add tx to stempool (passing in all txs from txpool to validate against).
		self.stempool
			.add_to_pool(entry, self.txpool.all_transactions(), header)?;

		// Note: we do not notify the adapter here,
		// we let the dandelion monitor handle this.
		Ok(())
	}

	fn add_to_reorg_cache(&mut self, entry: PoolEntry) -> Result<(), PoolError> {
		let mut cache = self.reorg_cache.write();
		cache.push_back(entry);
		if cache.len() > REORG_CACHE_SIZE {
			cache.pop_front();
		}
		debug!("added tx to reorg_cache: size now {}", cache.len());
		Ok(())
	}

	fn add_to_txpool(
		&mut self,
		mut entry: PoolEntry,
		header: &BlockHeader,
	) -> Result<(), PoolError> {
		// First deaggregate the tx based on current txpool txs.
		if entry.tx.kernels().len() > 1 {
			let txs = self.txpool.find_matching_transactions(entry.tx.kernels());
			if !txs.is_empty() {
				let tx = transaction::deaggregate(entry.tx, txs)?;
				tx.validate(self.verifier_cache.clone())?;
				entry.tx = tx;
				entry.src.debug_name = "deagg".to_string();
			}
		}
		self.txpool.add_to_pool(entry.clone(), vec![], header)?;

		// We now need to reconcile the stempool based on the new state of the txpool.
		// Some stempool txs may no longer be valid and we need to evict them.
		{
			let txpool_tx = self.txpool.aggregate_transaction()?;
			self.stempool.reconcile(txpool_tx, header)?;
		}

		self.adapter.tx_accepted(&entry.tx);
		Ok(())
	}

	/// Add the given tx to the pool, directing it to either the stempool or
	/// txpool based on stem flag provided.
	pub fn add_to_pool(
		&mut self,
		src: TxSource,
		tx: Transaction,
		stem: bool,
		header: &BlockHeader,
	) -> Result<(), PoolError> {
		// Quick check to deal with common case of seeing the *same* tx
		// broadcast from multiple peers simultaneously.
		if !stem && self.txpool.contains_tx(tx.hash()) {
			return Err(PoolError::DuplicateTx);
		}

		// Do we have the capacity to accept this transaction?
		self.is_acceptable(&tx, stem)?;

		// Make sure the transaction is valid before anything else.
		tx.validate(self.verifier_cache.clone())
			.map_err(PoolError::InvalidTx)?;

		// Check the tx lock_time is valid based on current chain state.
		self.blockchain.verify_tx_lock_height(&tx)?;

		// Check coinbase maturity before we go any further.
		self.blockchain.verify_coinbase_maturity(&tx)?;

		let entry = PoolEntry {
			state: PoolEntryState::Fresh,
			src,
			tx_at: Utc::now(),
			tx,
		};

		if stem {
			// TODO - what happens to txs in the stempool in a re-org scenario?
			self.add_to_stempool(entry, header)?;
		} else {
			self.add_to_txpool(entry.clone(), header)?;
			self.add_to_reorg_cache(entry)?;
		}
		Ok(())
	}

	fn reconcile_reorg_cache(&mut self, header: &BlockHeader) -> Result<(), PoolError> {
		let entries = self.reorg_cache.read().iter().cloned().collect::<Vec<_>>();
		debug!("reconcile_reorg_cache: size: {} ...", entries.len());
		for entry in entries {
			let _ = &self.add_to_txpool(entry.clone(), header);
		}
		debug!("reconcile_reorg_cache: ... done.");
		Ok(())
	}

	/// Reconcile the transaction pool (both txpool and stempool) against the
	/// provided block.
	pub fn reconcile_block(&mut self, block: &Block) -> Result<(), PoolError> {
		// First reconcile the txpool.
		self.txpool.reconcile_block(block);
		self.txpool.reconcile(None, &block.header)?;

		// Take our "reorg_cache" and see if this block means
		// we need to (re)add old txs due to a fork and re-org.
		self.reconcile_reorg_cache(&block.header)?;

		// Now reconcile our stempool, accounting for the updated txpool txs.
		self.stempool.reconcile_block(block);
		{
			let txpool_tx = self.txpool.aggregate_transaction()?;
			self.stempool.reconcile(txpool_tx, &block.header)?;
		}

		Ok(())
	}

	/// Retrieve individual transaction for the given kernel hash.
	pub fn retrieve_tx_by_kernel_hash(&self, hash: Hash) -> Option<Transaction> {
		self.txpool.retrieve_tx_by_kernel_hash(hash)
	}

	/// Retrieve all transactions matching the provided "compact block"
	/// based on the kernel set.
	/// Note: we only look in the txpool for this (stempool is under embargo).
	pub fn retrieve_transactions(
		&self,
		hash: Hash,
		nonce: u64,
		kern_ids: &[ShortId],
	) -> (Vec<Transaction>, Vec<ShortId>) {
		self.txpool.retrieve_transactions(hash, nonce, kern_ids)
	}

	/// Whether the transaction is acceptable to the pool, given both how
	/// full the pool is and the transaction weight.
	fn is_acceptable(&self, tx: &Transaction, stem: bool) -> Result<(), PoolError> {
		if self.total_size() > self.config.max_pool_size {
			// TODO evict old/large transactions instead
			return Err(PoolError::OverCapacity);
		}

		// Check that the stempool can accept this transaction
		if stem {
			if self.stempool.size() > self.config.max_stempool_size {
				// TODO evict old/large transactions instead
				return Err(PoolError::OverCapacity);
			}
		}

		// for a basic transaction (1 input, 2 outputs) -
		// (-1 * 1) + (4 * 2) + 1 = 8
		// 8 * 10 = 80
		if self.config.accept_fee_base > 0 {
			let threshold = (tx.tx_weight() as u64) * self.config.accept_fee_base;
			if tx.fee() < threshold {
				return Err(PoolError::LowFeeTransaction(threshold));
			}
		}
		Ok(())
	}

	/// Get the total size of the pool.
	/// Note: we only consider the txpool here as stempool is under embargo.
	pub fn total_size(&self) -> usize {
		self.txpool.size()
	}

	/// Returns a vector of transactions from the txpool so we can build a
	/// block from them.
	pub fn prepare_mineable_transactions(&self) -> Result<Vec<Transaction>, PoolError> {
		self.txpool.prepare_mineable_transactions(self.config.mineable_max_weight)
	}
}
<commit_msg>Reorg cache time based (cache 30mins of txs) (#2002)<commit_after>// Copyright 2018 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Transaction pool implementation leveraging txhashset for chain state
//! validation. It is a valid operation to add a tx to the tx pool if the
//! resulting tx pool can be added to the current chain state to produce a
//! valid chain state.

use std::collections::VecDeque;
use std::sync::Arc;
use util::RwLock;

use chrono::prelude::*;
use chrono::Duration;

use core::core::hash::{Hash, Hashed};
use core::core::id::ShortId;
use core::core::verifier_cache::VerifierCache;
use core::core::{transaction, Block, BlockHeader, Transaction};
use pool::Pool;
use types::{BlockChain, PoolAdapter, PoolConfig, PoolEntry, PoolEntryState, PoolError, TxSource};

/// Transaction pool implementation.
pub struct TransactionPool {
	/// Pool Config
	pub config: PoolConfig,
	/// Our transaction pool.
	pub txpool: Pool,
	/// Our Dandelion "stempool".
	pub stempool: Pool,
	/// Cache of previous txs in case of a re-org.
	pub reorg_cache: Arc<RwLock<VecDeque<PoolEntry>>>,
	/// The blockchain
	pub blockchain: Arc<BlockChain>,
	pub verifier_cache: Arc<RwLock<VerifierCache>>,
	/// The pool adapter
	pub adapter: Arc<PoolAdapter>,
}

impl TransactionPool {
	/// Create a new transaction pool
	pub fn new(
		config: PoolConfig,
		chain: Arc<BlockChain>,
		verifier_cache: Arc<RwLock<VerifierCache>>,
		adapter: Arc<PoolAdapter>,
	) -> TransactionPool {
		TransactionPool {
			config,
			txpool: Pool::new(chain.clone(), verifier_cache.clone(), "txpool".to_string()),
			stempool: Pool::new(
				chain.clone(),
				verifier_cache.clone(),
				"stempool".to_string(),
			),
			reorg_cache: Arc::new(RwLock::new(VecDeque::new())),
			blockchain: chain,
			verifier_cache,
			adapter,
		}
	}

	pub fn chain_head(&self) -> Result<BlockHeader, PoolError> {
		self.blockchain.chain_head()
	}

	fn add_to_stempool(&mut self, entry: PoolEntry, header: &BlockHeader) -> Result<(), PoolError> {
		// Add tx to stempool (passing in all txs from txpool to validate against).
		self.stempool
			.add_to_pool(entry, self.txpool.all_transactions(), header)?;

		// Note: we do not notify the adapter here,
		// we let the dandelion monitor handle this.
		Ok(())
	}

	fn add_to_reorg_cache(&mut self, entry: PoolEntry) {
		let mut cache = self.reorg_cache.write();
		cache.push_back(entry);

		// We cache 30 mins of txs but we have a hard limit to avoid catastrophic failure.
		// For simplicity use the same value as the actual tx pool limit.
		if cache.len() > self.config.max_pool_size {
			let _ = cache.pop_front();
		}
		debug!("added tx to reorg_cache: size now {}", cache.len());
	}

	// Old txs will "age out" after 30 mins.
	fn truncate_reorg_cache(&mut self, cutoff: DateTime<Utc>) {
		let mut cache = self.reorg_cache.write();

		while cache.front().map(|x| x.tx_at < cutoff).unwrap_or(false) {
			let _ = cache.pop_front();
		}
	}

	fn add_to_txpool(
		&mut self,
		mut entry: PoolEntry,
		header: &BlockHeader,
	) -> Result<(), PoolError> {
		// First deaggregate the tx based on current txpool txs.
		if entry.tx.kernels().len() > 1 {
			let txs = self.txpool.find_matching_transactions(entry.tx.kernels());
			if !txs.is_empty() {
				let tx = transaction::deaggregate(entry.tx, txs)?;
				tx.validate(self.verifier_cache.clone())?;
				entry.tx = tx;
				entry.src.debug_name = "deagg".to_string();
			}
		}
		self.txpool.add_to_pool(entry.clone(), vec![], header)?;

		// We now need to reconcile the stempool based on the new state of the txpool.
		// Some stempool txs may no longer be valid and we need to evict them.
		{
			let txpool_tx = self.txpool.aggregate_transaction()?;
			self.stempool.reconcile(txpool_tx, header)?;
		}

		self.adapter.tx_accepted(&entry.tx);
		Ok(())
	}

	/// Add the given tx to the pool, directing it to either the stempool or
	/// txpool based on stem flag provided.
	pub fn add_to_pool(
		&mut self,
		src: TxSource,
		tx: Transaction,
		stem: bool,
		header: &BlockHeader,
	) -> Result<(), PoolError> {
		// Quick check to deal with common case of seeing the *same* tx
		// broadcast from multiple peers simultaneously.
		if !stem && self.txpool.contains_tx(tx.hash()) {
			return Err(PoolError::DuplicateTx);
		}

		// Do we have the capacity to accept this transaction?
		self.is_acceptable(&tx, stem)?;

		// Make sure the transaction is valid before anything else.
		tx.validate(self.verifier_cache.clone())
			.map_err(PoolError::InvalidTx)?;

		// Check the tx lock_time is valid based on current chain state.
		self.blockchain.verify_tx_lock_height(&tx)?;

		// Check coinbase maturity before we go any further.
		self.blockchain.verify_coinbase_maturity(&tx)?;

		let entry = PoolEntry {
			state: PoolEntryState::Fresh,
			src,
			tx_at: Utc::now(),
			tx,
		};

		if stem {
			// TODO - what happens to txs in the stempool in a re-org scenario?
			self.add_to_stempool(entry, header)?;
		} else {
			self.add_to_txpool(entry.clone(), header)?;
			self.add_to_reorg_cache(entry);
		}
		Ok(())
	}

	fn reconcile_reorg_cache(&mut self, header: &BlockHeader) -> Result<(), PoolError> {
		// First "age out" any old txs in the reorg cache.
		let cutoff = Utc::now() - Duration::minutes(30);
		self.truncate_reorg_cache(cutoff);

		let entries = self.reorg_cache.read().iter().cloned().collect::<Vec<_>>();
		debug!("reconcile_reorg_cache: size: {} ...", entries.len());
		for entry in entries {
			let _ = &self.add_to_txpool(entry.clone(), header);
		}
		debug!("reconcile_reorg_cache: ... done.");
		Ok(())
	}

	/// Reconcile the transaction pool (both txpool and stempool) against the
	/// provided block.
	pub fn reconcile_block(&mut self, block: &Block) -> Result<(), PoolError> {
		// First reconcile the txpool.
		self.txpool.reconcile_block(block);
		self.txpool.reconcile(None, &block.header)?;

		// Take our "reorg_cache" and see if this block means
		// we need to (re)add old txs due to a fork and re-org.
		self.reconcile_reorg_cache(&block.header)?;

		// Now reconcile our stempool, accounting for the updated txpool txs.
		self.stempool.reconcile_block(block);
		{
			let txpool_tx = self.txpool.aggregate_transaction()?;
			self.stempool.reconcile(txpool_tx, &block.header)?;
		}

		Ok(())
	}

	/// Retrieve individual transaction for the given kernel hash.
	pub fn retrieve_tx_by_kernel_hash(&self, hash: Hash) -> Option<Transaction> {
		self.txpool.retrieve_tx_by_kernel_hash(hash)
	}

	/// Retrieve all transactions matching the provided "compact block"
	/// based on the kernel set.
	/// Note: we only look in the txpool for this (stempool is under embargo).
	pub fn retrieve_transactions(
		&self,
		hash: Hash,
		nonce: u64,
		kern_ids: &[ShortId],
	) -> (Vec<Transaction>, Vec<ShortId>) {
		self.txpool.retrieve_transactions(hash, nonce, kern_ids)
	}

	/// Whether the transaction is acceptable to the pool, given both how
	/// full the pool is and the transaction weight.
	fn is_acceptable(&self, tx: &Transaction, stem: bool) -> Result<(), PoolError> {
		if self.total_size() > self.config.max_pool_size {
			// TODO evict old/large transactions instead
			return Err(PoolError::OverCapacity);
		}

		// Check that the stempool can accept this transaction
		if stem {
			if self.stempool.size() > self.config.max_stempool_size {
				// TODO evict old/large transactions instead
				return Err(PoolError::OverCapacity);
			}
		}

		// for a basic transaction (1 input, 2 outputs) -
		// (-1 * 1) + (4 * 2) + 1 = 8
		// 8 * 10 = 80
		if self.config.accept_fee_base > 0 {
			let threshold = (tx.tx_weight() as u64) * self.config.accept_fee_base;
			if tx.fee() < threshold {
				return Err(PoolError::LowFeeTransaction(threshold));
			}
		}
		Ok(())
	}

	/// Get the total size of the pool.
	/// Note: we only consider the txpool here as stempool is under embargo.
	pub fn total_size(&self) -> usize {
		self.txpool.size()
	}

	/// Returns a vector of transactions from the txpool so we can build a
	/// block from them.
	pub fn prepare_mineable_transactions(&self) -> Result<Vec<Transaction>, PoolError> {
		self.txpool
			.prepare_mineable_transactions(self.config.mineable_max_weight)
	}
}
<|endoftext|>


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

use anyhow::{anyhow, Result};
use erased_serde::Serialize;
use nix::unistd::isatty;
use raw_tty::TtyModeGuard;
use regex::Regex;
use std::any::Any;
use std::fs::File;
use std::io::{Read, Write};
use std::os::unix::io::AsRawFd;
use std::time::{Duration, Instant};
use structopt::StructOpt;

use opentitanlib::app::command::CommandDispatch;
use opentitanlib::app::TransportWrapper;
use opentitanlib::io::uart::{Uart, UartParams};
use opentitanlib::transport::Capability;
use opentitanlib::util::file;

#[derive(Debug, StructOpt)]
pub struct Console {
    #[structopt(flatten)]
    params: UartParams,

    #[structopt(short, long, help = "Do not print console start end exit messages.")]
    quiet: bool,

    #[structopt(short, long, help = "Log console output to a file")]
    logfile: Option<String>,

    #[structopt(short, long, help = "Exit after a timeout in seconds.")]
    timeout: Option<u64>,

    #[structopt(long, help = "Exit with success if the specified regex is matched.")]
    exit_success: Option<String>,

    #[structopt(long, help = "Exit with failure if the specified regex is matched.")]
    exit_failure: Option<String>,
}

impl CommandDispatch for Console {
    fn run(
        &self,
        _context: &dyn Any,
        transport: &TransportWrapper,
    ) -> Result<Option<Box<dyn Serialize>>> {
        // We need the UART for the console command to operate.
        transport.capabilities()?.request(Capability::UART).ok()?;
        let mut stdout = std::io::stdout();
        let mut stdin = std::io::stdin();

        // Set up resources specified by the command line parameters.
        let mut console = InnerConsole {
            logfile: self.logfile.as_ref().map(File::create).transpose()?,
            deadline: self.timeout.map(|t| Instant::now() + Duration::new(t, 0)),
            exit_success: self
                .exit_success
                .as_ref()
                .map(|s| Regex::new(s.as_str()))
                .transpose()?,
            exit_failure: self
                .exit_failure
                .as_ref()
                .map(|s| Regex::new(s.as_str()))
                .transpose()?,
            ..Default::default()
        };

        if !self.quiet {
            println!("Starting interactive console");
            println!("[CTRL+C] to exit.\n");
        }
        {
            // Put the terminal into raw mode.  The tty guard will restore the
            // console settings when it goes out of scope.
            let _stdin_guard = if isatty(stdin.as_raw_fd())? {
                let mut guard = TtyModeGuard::new(stdin.as_raw_fd())?;
                guard.set_raw_mode()?;
                Some(guard)
            } else {
                None
            };
            let _stdout_guard = if isatty(stdout.as_raw_fd())? {
                let mut guard = TtyModeGuard::new(stdout.as_raw_fd())?;
                guard.set_raw_mode()?;
                Some(guard)
            } else {
                None
            };
            let uart = self.params.create(transport)?;
            console.interact(&*uart, &mut stdin, &mut stdout)?;
        }
        if !self.quiet {
            println!("\n\nExiting interactive console.");
        }
        Ok(None)
    }
}

#[derive(Default)]
struct InnerConsole {
    logfile: Option<File>,
    deadline: Option<Instant>,
    exit_success: Option<Regex>,
    exit_failure: Option<Regex>,
    buffer: String,
}

enum ExitStatus {
    None,
    ExitSuccess,
    ExitFailure,
}

impl InnerConsole {
    const CTRL_C: u8 = 3;
    const BUFFER_LEN: usize = 1024;

    // Runs an interactive console until CTRL_C is received.
    fn interact(
        &mut self,
        uart: &dyn Uart,
        stdin: &mut (impl Read + AsRawFd),
        stdout: &mut impl Write,
    ) -> Result<()> {
        let mut buf = [0u8; 256];

        loop {
            if let Some(deadline) = self.deadline {
                if Instant::now() > deadline {
                    // If we have an exit success condition, then a timeout
                    // should be an error.
                    if self.exit_success.is_some() {
                        return Err(anyhow!("Console timeout exceeded"));
                    } else {
                        break;
                    }
                }
            }

            // This loop _should_ really use unix `poll` in the conventional way
            // to learn when the console or uart file descriptors become ready,
            // but some UART backends will bury their implementation in libusb
            // and make discovering the file descriptor difficult or impossible.
            //
            // As a pragmatic implementation detail, we wait for the UART
            // for a short period of time and then service the console.
            //
            // TODO: as we write more backends, re-evaluate whether there is a
            // better way to approach waiting on the UART and keyboard.

            // Check for input on the uart.
            match self.uart_read(uart, Duration::from_millis(10), stdout)? {
                ExitStatus::None => {}
                ExitStatus::ExitSuccess => {
                    break;
                }
                ExitStatus::ExitFailure => {
                    return Err(anyhow!("Matched exit_failure expression"));
                }
            };

            // Wait for input from the user.
            if file::wait_read_timeout(&*stdin, Duration::from_millis(0)).is_ok() {
                let len = stdin.read(&mut buf)?;
                if len == 1 && buf[0] == InnerConsole::CTRL_C {
                    break;
                }
                uart.write(&buf[..len])?;
            }
        }
        Ok(())
    }

    // Maintain a buffer for the exit regexes to match against.
    fn append_buffer(&mut self, data: &[u8]) {
        self.buffer.push_str(&String::from_utf8_lossy(data));
        while self.buffer.len() > InnerConsole::BUFFER_LEN {
            self.buffer.remove(0);
        }
    }

    // Read from the uart and process the data read.
    fn uart_read(
        &mut self,
        uart: &dyn Uart,
        timeout: Duration,
        stdout: &mut impl Write,
    ) -> Result<ExitStatus> {
        let mut buf = [0u8; 256];
        let len = uart.read_timeout(&mut buf, timeout)?;
        if len > 0 {
            stdout.write_all(&buf[..len])?;
            stdout.flush()?;

            // If we're logging, save it to the logfile.
            if let Some(logfile) = &mut self.logfile {
                logfile.write_all(&buf[..len])?;
            }

            // If we have exit condition regexes check them.
            self.append_buffer(&buf[..len]);
            if self
                .exit_success
                .as_ref()
                .map(|rx| rx.is_match(&self.buffer))
                == Some(true)
            {
                return Ok(ExitStatus::ExitSuccess);
            }
            if self
                .exit_failure
                .as_ref()
                .map(|rx| rx.is_match(&self.buffer))
                == Some(true)
            {
                return Ok(ExitStatus::ExitFailure);
            }
        }
        Ok(ExitStatus::None)
    }
}
<commit_msg>[opentitantool] Add timestamps to console output<commit_after>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

use anyhow::{anyhow, Result};
use erased_serde::Serialize;
use nix::unistd::isatty;
use raw_tty::TtyModeGuard;
use regex::Regex;
use std::any::Any;
use std::fs::File;
use std::io::{Read, Write};
use std::os::unix::io::AsRawFd;
use std::time::{Duration, Instant, SystemTime};
use structopt::StructOpt;

use opentitanlib::app::command::CommandDispatch;
use opentitanlib::app::TransportWrapper;
use opentitanlib::io::uart::{Uart, UartParams};
use opentitanlib::transport::Capability;
use opentitanlib::util::file;

#[derive(Debug, StructOpt)]
pub struct Console {
    #[structopt(flatten)]
    params: UartParams,

    #[structopt(short, long, help = "Do not print console start end exit messages.")]
    quiet: bool,

    #[structopt(short, long, help = "Log console output to a file")]
    logfile: Option<String>,

    #[structopt(short, long, help = "Exit after a timeout in seconds.")]
    timeout: Option<u64>,

    #[structopt(long, help = "Print a timestamp on each line of console output.")]
    timestamp: bool,

    #[structopt(long, help = "Exit with success if the specified regex is matched.")]
    exit_success: Option<String>,

    #[structopt(long, help = "Exit with failure if the specified regex is matched.")]
    exit_failure: Option<String>,
}

impl CommandDispatch for Console {
    fn run(
        &self,
        _context: &dyn Any,
        transport: &TransportWrapper,
    ) -> Result<Option<Box<dyn Serialize>>> {
        // We need the UART for the console command to operate.
        transport.capabilities()?.request(Capability::UART).ok()?;
        let mut stdout = std::io::stdout();
        let mut stdin = std::io::stdin();

        // Set up resources specified by the command line parameters.
        let mut console = InnerConsole {
            logfile: self.logfile.as_ref().map(File::create).transpose()?,
            deadline: self.timeout.map(|t| Instant::now() + Duration::new(t, 0)),
            exit_success: self
                .exit_success
                .as_ref()
                .map(|s| Regex::new(s.as_str()))
                .transpose()?,
            exit_failure: self
                .exit_failure
                .as_ref()
                .map(|s| Regex::new(s.as_str()))
                .transpose()?,
            timestamp: self.timestamp,
            newline: true,
            ..Default::default()
        };

        if !self.quiet {
            println!("Starting interactive console");
            println!("[CTRL+C] to exit.\n");
        }
        {
            // Put the terminal into raw mode.  The tty guard will restore the
            // console settings when it goes out of scope.
            let _stdin_guard = if isatty(stdin.as_raw_fd())? {
                let mut guard = TtyModeGuard::new(stdin.as_raw_fd())?;
                guard.set_raw_mode()?;
                Some(guard)
            } else {
                None
            };
            let _stdout_guard = if isatty(stdout.as_raw_fd())? {
                let mut guard = TtyModeGuard::new(stdout.as_raw_fd())?;
                guard.set_raw_mode()?;
                Some(guard)
            } else {
                None
            };
            let uart = self.params.create(transport)?;
            console.interact(&*uart, &mut stdin, &mut stdout)?;
        }
        if !self.quiet {
            println!("\n\nExiting interactive console.");
        }
        Ok(None)
    }
}

#[derive(Default)]
struct InnerConsole {
    logfile: Option<File>,
    deadline: Option<Instant>,
    exit_success: Option<Regex>,
    exit_failure: Option<Regex>,
    buffer: String,
    timestamp: bool,
    newline: bool,
}

enum ExitStatus {
    None,
    ExitSuccess,
    ExitFailure,
}

impl InnerConsole {
    const CTRL_C: u8 = 3;
    const BUFFER_LEN: usize = 1024;

    // Runs an interactive console until CTRL_C is received.
    fn interact(
        &mut self,
        uart: &dyn Uart,
        stdin: &mut (impl Read + AsRawFd),
        stdout: &mut impl Write,
    ) -> Result<()> {
        let mut buf = [0u8; 256];

        loop {
            if let Some(deadline) = self.deadline {
                if Instant::now() > deadline {
                    // If we have an exit success condition, then a timeout
                    // should be an error.
                    if self.exit_success.is_some() {
                        return Err(anyhow!("Console timeout exceeded"));
                    } else {
                        break;
                    }
                }
            }

            // This loop _should_ really use unix `poll` in the conventional way
            // to learn when the console or uart file descriptors become ready,
            // but some UART backends will bury their implementation in libusb
            // and make discovering the file descriptor difficult or impossible.
            //
            // As a pragmatic implementation detail, we wait for the UART
            // for a short period of time and then service the console.
            //
            // TODO: as we write more backends, re-evaluate whether there is a
            // better way to approach waiting on the UART and keyboard.

            // Check for input on the uart.
            match self.uart_read(uart, Duration::from_millis(10), stdout)? {
                ExitStatus::None => {}
                ExitStatus::ExitSuccess => {
                    break;
                }
                ExitStatus::ExitFailure => {
                    return Err(anyhow!("Matched exit_failure expression"));
                }
            };

            // Wait for input from the user.
            if file::wait_read_timeout(&*stdin, Duration::from_millis(0)).is_ok() {
                let len = stdin.read(&mut buf)?;
                if len == 1 && buf[0] == InnerConsole::CTRL_C {
                    break;
                }
                uart.write(&buf[..len])?;
            }
        }
        Ok(())
    }

    // Maintain a buffer for the exit regexes to match against.
    fn append_buffer(&mut self, data: &[u8]) {
        self.buffer.push_str(&String::from_utf8_lossy(data));
        while self.buffer.len() > InnerConsole::BUFFER_LEN {
            self.buffer.remove(0);
        }
    }

    // Read from the uart and process the data read.
    fn uart_read(
        &mut self,
        uart: &dyn Uart,
        timeout: Duration,
        stdout: &mut impl Write,
    ) -> Result<ExitStatus> {
        let mut buf = [0u8; 256];
        let len = uart.read_timeout(&mut buf, timeout)?;
        if len > 0 {
            for i in 0..len {
                if self.timestamp && self.newline {
                    let t = humantime::format_rfc3339_millis(SystemTime::now());
                    stdout.write_fmt(format_args!("[{}]", t))?;
                    self.newline = false;
                }
                stdout.write_all(&buf[i..i + 1])?;
                self.newline = buf[i] == b'\n';
            }
            stdout.flush()?;

            // If we're logging, save it to the logfile.
            if let Some(logfile) = &mut self.logfile {
                logfile.write_all(&buf[..len])?;
            }

            // If we have exit condition regexes check them.
            self.append_buffer(&buf[..len]);
            if self
                .exit_success
                .as_ref()
                .map(|rx| rx.is_match(&self.buffer))
                == Some(true)
            {
                return Ok(ExitStatus::ExitSuccess);
            }
            if self
                .exit_failure
                .as_ref()
                .map(|rx| rx.is_match(&self.buffer))
                == Some(true)
            {
                return Ok(ExitStatus::ExitFailure);
            }
        }
        Ok(ExitStatus::None)
    }
}
<|endoftext|>


<commit_before>use env_logger;
use failure;

use gluon;
use hubcaps;

#[macro_use]
extern crate log;
#[macro_use]
extern crate serde_derive;

#[allow(unused_imports)]
#[macro_use]
extern crate structopt;
use tokio;

use gluon_crates_io;
use gluon_master;

#[macro_use]
extern crate gluon_vm;
#[macro_use]
extern crate gluon_codegen;

use std::fs;
use std::ops::Deref;

use futures::{future, prelude::*};

use hyper::client::HttpConnector;
use hyper_tls::HttpsConnector;

use tokio_signal::unix::{Signal, SIGINT, SIGTERM};

use gluon::{
    vm::{
        self,
        api::{OwnedFunction, RuntimeResult, IO},
        ExternModule,
    },
    Thread,
};

use structopt::StructOpt;

pub fn load_master(thread: &Thread) -> vm::Result<ExternModule> {
    #[derive(Debug, Userdata)]
    pub struct TryThread(gluon_master::RootedThread);

    impl Deref for TryThread {
        type Target = gluon_master::Thread;

        fn deref(&self) -> &Self::Target {
            &self.0
        }
    }

    thread.register_type::<TryThread>("MasterTryThread", &[])?;

    ExternModule::new(
        thread,
        record! {
            make_eval_vm => primitive!(1, "make_eval_vm", |()| {
                RuntimeResult::from(gluon_master::make_eval_vm().map(TryThread))
            }),
            eval => primitive!(2, "eval", |t: &TryThread, s: &str| gluon_master::eval(t, s)),
            format_expr => primitive!(2, |t: &TryThread, s: &str| gluon_master::format_expr(t, s))
        },
    )
}

pub fn load(thread: &Thread) -> vm::Result<ExternModule> {
    #[derive(Debug, Userdata)]
    pub struct TryThread(gluon_crates_io::RootedThread);

    impl Deref for TryThread {
        type Target = gluon_crates_io::Thread;

        fn deref(&self) -> &Self::Target {
            &self.0
        }
    }

    thread.register_type::<TryThread>("TryThread", &[])?;

    ExternModule::new(
        thread,
        record! {
            make_eval_vm => primitive!(1, "make_eval_vm", |()| {
                RuntimeResult::from(gluon_crates_io::make_eval_vm().map(TryThread))
            }),
            eval => primitive!(2, "eval", |t: &TryThread, s: &str| gluon_crates_io::eval(t, s)),
            format_expr => primitive!(2, |t: &TryThread, s: &str| gluon_crates_io::format_expr(t, s))
        },
    )
}

#[derive(Debug, Default, Getable, VmType)]
pub struct Gist<'a> {
    pub code: &'a str,
}

#[derive(Debug, Default, Serialize, Pushable, VmType)]
pub struct PostGist {
    pub id: String,
    pub html_url: String,
}

#[derive(Debug, Userdata)]
struct Github(hubcaps::Github<HttpsConnector<HttpConnector>>);

fn new_github(gist_access_token: &str) -> Github {
    Github(hubcaps::Github::new(
        "try_gluon".to_string(),
        hubcaps::Credentials::Token(gist_access_token.into()),
    ))
}

fn share(
    github: &Github,
    gist: Gist<'_>,
) -> impl Future<Item = Result<PostGist, String>, Error = vm::Error> {
    info!("Share: `{}`", gist.code);

    github
        .0
        .gists()
        .create(&hubcaps::gists::GistOptions {
            description: Some("Gluon code shared from try_gluon".into()),
            public: Some(true),
            files: Some((
                "try_gluon.glu".into(),
                hubcaps::gists::Content {
                    filename: None,
                    content: gist.code.into(),
                },
            ))
            .into_iter()
            .collect(),
        })
        .map_err(|err| err.to_string())
        .map(|response| PostGist {
            id: response.id,
            html_url: response.html_url,
        })
        .then(Ok)
}

#[derive(StructOpt, Pushable, VmType)]
struct Opts {
    #[structopt(
        long = "gist-access-token",
        env = "GIST_ACCESS_TOKEN",
        help = "The access tokens used to create gists"
    )]
    gist_access_token: Option<String>,
    #[structopt(short = "p", long = "port", help = "The port to start the server on")]
    port: Option<u16>,
    #[structopt(long = "https", help = "Whether to run the server with https")]
    https: bool,
    #[structopt(
        long = "host",
        default_value = "gluon-lang.org",
        help = "The hostname for the server"
    )]
    host: String,
    #[structopt(
        long = "staging",
        help = "Whether to use letsencrypt's staging environment"
    )]
    staging: bool,
}

fn main() {
    if let Err(err) = main_() {
        eprintln!("{}\n{}", err, err.backtrace());
    }
}

fn main_() -> Result<(), failure::Error> {
    env_logger::init();

    let opts = Opts::from_args();

    let mut runtime = tokio::runtime::Runtime::new()?;

    let vm = gluon::new_vm();
    gluon::import::add_extern_module(&vm, "gluon.try", load);
    gluon::import::add_extern_module(&vm, "gluon.try.master", load_master);
    gluon::import::add_extern_module(&vm, "gluon.http_server", |vm| {
        ExternModule::new(
            vm,
            record! {
                type Opts => Opts
            },
        )
    });
    gluon::import::add_extern_module(&vm, "github", |vm| {
        vm.register_type::<Github>("Github", &[])?;
        ExternModule::new(
            vm,
            record! {
                new_github => primitive!(1, new_github),
                share => primitive!(2, async fn share)
            },
        )
    });

    let server_source = fs::read_to_string("src/app/server.glu")?;

    let signal = Signal::new(SIGINT)
        .flatten_stream()
        .select(Signal::new(SIGTERM).flatten_stream())
        .into_future()
        .map(|_| {
            eprintln!("Signal received. Shutting down");
        })
        .map_err(|(err, _)| failure::format_err!("{}", err));

    runtime.block_on(
        future::lazy(move || {
            gluon::Compiler::new()
                .run_expr_async::<OwnedFunction<fn(Opts) -> IO<()>>>(
                    &vm,
                    "src.app.server",
                    &server_source,
                )
                .and_then(|(mut f, _)| f.call_async(opts).from_err())
                .map_err(|err| failure::Error::from(err))
                .map(|_| ())
        })
        .select(signal)
        .map_err(|(err, _)| err),
    )?;

    Ok(())
}
<commit_msg>fix: Get the code compiling on windows again<commit_after>use env_logger;
use failure;

use gluon;
use hubcaps;

#[macro_use]
extern crate log;
#[macro_use]
extern crate serde_derive;

#[allow(unused_imports)]
#[macro_use]
extern crate structopt;
use tokio;

use gluon_crates_io;
use gluon_master;

#[macro_use]
extern crate gluon_vm;
#[macro_use]
extern crate gluon_codegen;

use std::fs;
use std::ops::Deref;

use futures::{future, prelude::*};

use hyper::client::HttpConnector;
use hyper_tls::HttpsConnector;

use gluon::{
    vm::{
        self,
        api::{OwnedFunction, RuntimeResult, IO},
        ExternModule,
    },
    Thread,
};

use structopt::StructOpt;

pub fn load_master(thread: &Thread) -> vm::Result<ExternModule> {
    #[derive(Debug, Userdata)]
    pub struct TryThread(gluon_master::RootedThread);

    impl Deref for TryThread {
        type Target = gluon_master::Thread;

        fn deref(&self) -> &Self::Target {
            &self.0
        }
    }

    thread.register_type::<TryThread>("MasterTryThread", &[])?;

    ExternModule::new(
        thread,
        record! {
            make_eval_vm => primitive!(1, "make_eval_vm", |()| {
                RuntimeResult::from(gluon_master::make_eval_vm().map(TryThread))
            }),
            eval => primitive!(2, "eval", |t: &TryThread, s: &str| gluon_master::eval(t, s)),
            format_expr => primitive!(2, |t: &TryThread, s: &str| gluon_master::format_expr(t, s))
        },
    )
}

pub fn load(thread: &Thread) -> vm::Result<ExternModule> {
    #[derive(Debug, Userdata)]
    pub struct TryThread(gluon_crates_io::RootedThread);

    impl Deref for TryThread {
        type Target = gluon_crates_io::Thread;

        fn deref(&self) -> &Self::Target {
            &self.0
        }
    }

    thread.register_type::<TryThread>("TryThread", &[])?;

    ExternModule::new(
        thread,
        record! {
            make_eval_vm => primitive!(1, "make_eval_vm", |()| {
                RuntimeResult::from(gluon_crates_io::make_eval_vm().map(TryThread))
            }),
            eval => primitive!(2, "eval", |t: &TryThread, s: &str| gluon_crates_io::eval(t, s)),
            format_expr => primitive!(2, |t: &TryThread, s: &str| gluon_crates_io::format_expr(t, s))
        },
    )
}

#[derive(Debug, Default, Getable, VmType)]
pub struct Gist<'a> {
    pub code: &'a str,
}

#[derive(Debug, Default, Serialize, Pushable, VmType)]
pub struct PostGist {
    pub id: String,
    pub html_url: String,
}

#[derive(Debug, Userdata)]
struct Github(hubcaps::Github<HttpsConnector<HttpConnector>>);

fn new_github(gist_access_token: &str) -> Github {
    Github(hubcaps::Github::new(
        "try_gluon".to_string(),
        hubcaps::Credentials::Token(gist_access_token.into()),
    ))
}

fn share(
    github: &Github,
    gist: Gist<'_>,
) -> impl Future<Item = Result<PostGist, String>, Error = vm::Error> {
    info!("Share: `{}`", gist.code);

    github
        .0
        .gists()
        .create(&hubcaps::gists::GistOptions {
            description: Some("Gluon code shared from try_gluon".into()),
            public: Some(true),
            files: Some((
                "try_gluon.glu".into(),
                hubcaps::gists::Content {
                    filename: None,
                    content: gist.code.into(),
                },
            ))
            .into_iter()
            .collect(),
        })
        .map_err(|err| err.to_string())
        .map(|response| PostGist {
            id: response.id,
            html_url: response.html_url,
        })
        .then(Ok)
}

#[derive(StructOpt, Pushable, VmType)]
struct Opts {
    #[structopt(
        long = "gist-access-token",
        env = "GIST_ACCESS_TOKEN",
        help = "The access tokens used to create gists"
    )]
    gist_access_token: Option<String>,
    #[structopt(short = "p", long = "port", help = "The port to start the server on")]
    port: Option<u16>,
    #[structopt(long = "https", help = "Whether to run the server with https")]
    https: bool,
    #[structopt(
        long = "host",
        default_value = "gluon-lang.org",
        help = "The hostname for the server"
    )]
    host: String,
    #[structopt(
        long = "staging",
        help = "Whether to use letsencrypt's staging environment"
    )]
    staging: bool,
}

fn main() {
    if let Err(err) = main_() {
        eprintln!("{}\n{}", err, err.backtrace());
    }
}

#[cfg(unix)]
fn exit_server() -> impl Future<Item = (), Error = failure::Error> {
    use tokio_signal::unix::{Signal, SIGINT, SIGTERM};
    Signal::new(SIGINT)
        .flatten_stream()
        .select(Signal::new(SIGTERM).flatten_stream())
        .into_future()
        .map(|_| {
            eprintln!("Signal received. Shutting down");
        })
        .map_err(|(err, _)| failure::format_err!("{}", err))
}

#[cfg(not(unix))]
fn exit_server() -> impl Future<Item = (), Error = failure::Error> {
    tokio_signal::ctrl_c()
        .flatten_stream()
        .into_future()
        .map(|_| ())
        .map_err(|(err, _)| failure::format_err!("{}", err))
}

fn main_() -> Result<(), failure::Error> {
    env_logger::init();

    let opts = Opts::from_args();

    let mut runtime = tokio::runtime::Runtime::new()?;

    let vm = gluon::new_vm();
    gluon::import::add_extern_module(&vm, "gluon.try", load);
    gluon::import::add_extern_module(&vm, "gluon.try.master", load_master);
    gluon::import::add_extern_module(&vm, "gluon.http_server", |vm| {
        ExternModule::new(
            vm,
            record! {
                type Opts => Opts
            },
        )
    });
    gluon::import::add_extern_module(&vm, "github", |vm| {
        vm.register_type::<Github>("Github", &[])?;
        ExternModule::new(
            vm,
            record! {
                new_github => primitive!(1, new_github),
                share => primitive!(2, async fn share)
            },
        )
    });

    let server_source = fs::read_to_string("src/app/server.glu")?;

    runtime.block_on(
        future::lazy(move || {
            gluon::Compiler::new()
                .run_expr_async::<OwnedFunction<fn(Opts) -> IO<()>>>(
                    &vm,
                    "src.app.server",
                    &server_source,
                )
                .and_then(|(mut f, _)| f.call_async(opts).from_err())
                .map_err(|err| failure::Error::from(err))
                .map(|_| ())
        })
        .select(exit_server())
        .map_err(|(err, _)| err),
    )?;

    Ok(())
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate crypto;
use crypto::digest::Digest;
use crypto::sha2::Sha256;
extern crate base64;
extern crate u2fhid;
use std::{io, thread, time};
use std::sync::mpsc::{channel, Sender, Receiver, RecvTimeoutError};
use std::sync::{Arc, Mutex};

const PARAMETER_SIZE : usize = 32;

pub struct WorkUnit
{
    timeout_sec: u8,
    challenge: Vec<u8>,
    application: Vec<u8>,
    key_handle: Option<Vec<u8>>,
    result_tx: Sender<io::Result<Vec<u8>>>
}

pub struct U2FManager {
    work_tx: Sender<WorkUnit>,
}

pub fn open_u2f_manager() -> io::Result<U2FManager> {
    let (mut tx, rx) = channel::<WorkUnit>();
    let manager = U2FManager{
        work_tx: tx,
    };

    if let Err(e) = thread::Builder::new().name("HID Runloop".to_string()).spawn(move || {
        U2FManager::worker_loop(rx);
    }) {
        return Err(e);
    }
    Ok(manager)
}

impl U2FManager {
    fn worker_loop(work_rx: Receiver<WorkUnit>) {
        let mut security_keys = Vec::new();

        let platform = match u2fhid::platform::open_platform_manager() {
            Ok(v) => v,
            Err(e) => panic!("Failure to open platform HID support: {}", e),
        };

        for mut device_obj in platform.find_keys().unwrap() {
            println!("Device seen {}", device_obj);
            if let Err(_) = u2fhid::init_device(&mut device_obj) {
                continue
            }
            println!("Device inited {}", device_obj);
            if let Err(_) = u2fhid::ping_device(&mut device_obj) {
                continue
            }
            println!("Device pinged {}", device_obj);
            if let Err(_) = u2fhid::u2f_version_is_v2(&mut device_obj){
                continue
            }
            println!("Device initialized {}", device_obj);
            security_keys.push(device_obj);
        }

        loop {
            println!("Doing work!");
            let work = match work_rx.recv() {
                Ok(v) => v,
                Err(e) => panic!("Receiver failed {}", e),
            };

            work.result_tx.send(Ok(Vec::new()));
        }

        // let mut iteration_count = 0;
        // while iteration_count < timeout_sec {
        //     if let Ok(ref mut list) = devices.lock() {
        //         // This indexing is because I don't know how to use .iter() and
        //         // keep the U2FDevice traits
        //         for idx in 0..list.len() {
        //             println!("Register {}", list[idx]);
        //             match u2fhid::u2f_register(&mut list[idx], challenge, application) {
        //                 Ok(v) => {
        //                     println!("Register OK {:?}", v);
        //                     // First to complete, we return
        //                     // TODO: Cancel the others?
        //                     for cancel_idx in 0..list.len() {
        //                         if cancel_idx == idx {
        //                             continue;
        //                         }
        //                         let _ = u2fhid::u2f_cancel(&mut list[cancel_idx]);
        //                     }
        //                     return Ok(v)
        //                 },
        //                 Err(_) => continue,
        //             }
        //         }
        //     }

        //     // We've tried all attached security keys
        //     iteration_count += 1;
        //     thread::sleep(time::Duration::from_secs(1));
        // }

        // if let Ok(ref mut list) = devices.lock() {
        //     for cancel_idx in 0..list.len() {
        //         let _ = u2fhid::u2f_cancel(&mut list[cancel_idx]);
        //     }
        // }

        // Err(io::Error::new(io::ErrorKind::TimedOut, "Timed out"))
    }

    pub fn register<F>(&self, timeout_sec: u8, challenge: Vec<u8>, application: Vec<u8>, callback: F)
        where F: FnOnce(io::Result<Vec<u8>>), F: Send + 'static
    {
        if challenge.len() != PARAMETER_SIZE || application.len() != PARAMETER_SIZE {
            callback(Err(io::Error::new(io::ErrorKind::InvalidInput, "Invalid parameter sizes")));
            return;
        }

        let (result_tx, result_rx) = channel::<io::Result<Vec<u8>>>();

        if let Err(e) = self.work_tx.send(WorkUnit{
            timeout_sec: timeout_sec,
            challenge: challenge,
            application: application,
            key_handle: None,
            result_tx: result_tx,
        }) {
            callback(Err(io::Error::new(io::ErrorKind::Other, format!("Send error {}", e))));
            return;
        }

        match result_rx.recv() {
            Ok(v) => callback(v),
            Err(e) => callback(Err(io::Error::new(io::ErrorKind::Other, e))),
        }
    }

    fn sign_thread(platform: &u2fhid::platform::PlatformManager, timeout_sec: u8,
                       challenge: &Vec<u8>, application: &Vec<u8>, key_handle: &Vec<u8>) -> io::Result<Vec<u8>>
    {
        let device_mutex: Arc<Mutex<Vec<u2fhid::platform::Device>>> = Arc::new(Mutex::new(Vec::new()));
        let devices = device_mutex.clone();

        if let Ok(ref mut list) = devices.lock() {
            for mut device_obj in try!(platform.find_keys()) {
                if let Err(_) = u2fhid::init_device(&mut device_obj) {
                    continue
                }
                if let Err(_) = u2fhid::ping_device(&mut device_obj) {
                    continue
                }
                if let Err(_) = u2fhid::u2f_version_is_v2(&mut device_obj){
                    continue
                }
                list.push(device_obj);
            }
        }

        let mut iteration_count = 0;
        while iteration_count < timeout_sec {
            if let Ok(ref mut list) = devices.lock() {
                // This indexing is because I don't know how to use .iter() and
                // keep the U2FDevice traits
                for idx in 0..list.len() {
                    match u2fhid::u2f_sign(&mut list[idx], challenge, application, key_handle) {
                        Ok(v) => {
                            // First to complete, we return
                            // TODO: Cancel the others?
                            for cancel_idx in 0..list.len() {
                                if cancel_idx == idx {
                                    continue;
                                }
                                let _ = u2fhid::u2f_cancel(&mut list[cancel_idx]);
                            }
                            return Ok(v)
                        },
                        Err(_) => continue,
                    }
                }
            }

            // We've tried all attached security keys
            iteration_count += 1;
            thread::sleep(time::Duration::from_secs(1));
        }

        if let Ok(ref mut list) = devices.lock() {
            for cancel_idx in 0..list.len() {
                let _ = u2fhid::u2f_cancel(&mut list[cancel_idx]);
            }
        }

        Err(io::Error::new(io::ErrorKind::TimedOut, "Timed out"))
    }

    pub fn sign<F>(&self, timeout_sec: u8, challenge: Vec<u8>, application: Vec<u8>, key_handle: Vec<u8>, callback: F)
        where F: FnOnce(io::Result<Vec<u8>>), F: Send + 'static
    {
        if challenge.len() != PARAMETER_SIZE || application.len() != PARAMETER_SIZE {
            callback(Err(io::Error::new(io::ErrorKind::InvalidInput, "Invalid parameter sizes")));
            return;
        }

        if let Err(e) = thread::Builder::new().spawn(move || {
            let platform = match u2fhid::platform::open_platform_manager() {
                Ok(v) => v,
                Err(e) => panic!("Failure to open platform HID support: {}", e),
            };

            let result = U2FManager::sign_thread(&platform, timeout_sec, &challenge, &application, &key_handle);
            platform.close();
            callback(result);
        }) {
            panic!("Failed to spawn thread: {}", e);
        }
    }
}

fn u2f_get_key_handle_from_register_response(register_response: &Vec<u8>) -> io::Result<Vec<u8>>
{
    if register_response[0] != 0x05 {
        return Err(io::Error::new(io::ErrorKind::InvalidData, "Reserved byte not set correctly"));
    }

    let key_handle_len = register_response[66] as usize;
    let mut public_key = register_response.clone();
    let mut key_handle = public_key.split_off(67);
    // let attestation = key_handle.split_off(key_handle_len);

    Ok(key_handle)
}

fn main() {
    println!("Searching for keys...");

    let mut challenge = Sha256::new();
    challenge.input_str(r#"{"challenge": "1vQ9mxionq0ngCnjD-wTsv1zUSrGRtFqG2xP09SbZ70", "version": "U2F_V2", "appId": "http://demo.yubico.com"}"#);
    let mut chall_bytes: Vec<u8> = vec![0; challenge.output_bytes()];
    challenge.result(&mut chall_bytes);

    let mut application = Sha256::new();
    application.input_str("http://demo.yubico.com");
    let mut app_bytes: Vec<u8> = vec![0; application.output_bytes()];
    application.result(&mut app_bytes);

    let (tx, rx) = channel();
    let manager = open_u2f_manager().unwrap();

    manager.register(15, chall_bytes, app_bytes, move |result| {
        // Ship back to the main thread
        if let Err(e) = tx.send(result) {
            panic!("Could not send: {}", e);
        }
    });

    let thread_result = match rx.recv() {
        Ok(v) => v,
        Err(e) => panic!("Couldn't read data: {}", e),
    };

    let register_data = match thread_result {
        Ok(v) => v,
        Err(e) => panic!("Register failure: {}", e),
    };

    println!("Register result: {}", base64::encode(&register_data));

    // let key_handle = u2f_get_key_handle_from_register_response(&register_data).unwrap();

    // let mut chall_bytes: Vec<u8> = vec![0; challenge.output_bytes()];
    // challenge.result(&mut chall_bytes);
    // let mut app_bytes: Vec<u8> = vec![0; application.output_bytes()];
    // application.result(&mut app_bytes);

    // manager.sign(15, chall_bytes, app_bytes, key_handle, move|result| {
    //     println!("Sign result: {}", base64::encode(&result.unwrap()));
    // });

    println!("Done.");
}
<commit_msg>Move devices to be initialized by main, not find_keys<commit_after>#[macro_use]
extern crate crypto;
use crypto::digest::Digest;
use crypto::sha2::Sha256;
extern crate base64;
extern crate u2fhid;
use std::{io, thread, time};
use std::io::{Read, Write};
use std::sync::mpsc::{channel, Sender, Receiver, RecvTimeoutError};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use u2fhid::U2FDevice;

const PARAMETER_SIZE : usize = 32;

enum Command {
    Register,
    Sign,
}

pub struct WorkUnit
{
    timeout: Duration,
    start_time: Instant,
    command: Command,
    challenge: Vec<u8>,
    application: Vec<u8>,
    key_handle: Option<Vec<u8>>,
    result_tx: Sender<io::Result<Vec<u8>>>
}

pub struct U2FManager {
    work_tx: Sender<WorkUnit>,
}

pub fn open_u2f_manager() -> io::Result<U2FManager> {
    let (mut tx, rx) = channel::<WorkUnit>();
    let manager = U2FManager{
        work_tx: tx,
    };

    if let Err(e) = thread::Builder::new().name("HID Runloop".to_string()).spawn(move || {
        U2FManager::worker_loop(rx);
    }) {
        return Err(e);
    }
    Ok(manager)
}

impl U2FManager {
    fn worker_loop(work_rx: Receiver<WorkUnit>) {
        let mut security_keys = Vec::new();

        let platform = match u2fhid::platform::open_platform_manager() {
            Ok(v) => v,
            Err(e) => panic!("Failure to open platform HID support: {}", e),
        };

        for mut device_obj in platform.find_keys().unwrap() {
            println!("Device seen {}", device_obj);
            if let Err(_) = u2fhid::init_device(&mut device_obj) {
                continue
            }
            println!("Device inited {}", device_obj);
            if let Err(_) = u2fhid::ping_device(&mut device_obj) {
                continue
            }
            println!("Device pinged {}", device_obj);
            if let Err(_) = u2fhid::u2f_version_is_v2(&mut device_obj) {
                continue
            }
            println!("Device initialized {}", device_obj);
            security_keys.push(device_obj);
        }

        let mut current_job : Option<WorkUnit> = None;

        loop {
            println!("Doing work");

            // Add new devices

            // Get new work
            match work_rx.recv_timeout(Duration::from_millis(100)) {
                Ok(v) => {
                    // TODO: Cancel existing job?
                    current_job = Some(v)
                },
                Err(e) => {
                    if e != RecvTimeoutError::Timeout {
                        panic!("whoa now {}", e);
                    }
                },
            };

            current_job = match current_job {
                Some(job) => {
                    let mut done = false;
                    for idx in 0..security_keys.len() {
                        println!("iterating now");

                        if let Ok(_) = U2FManager::perform_job_for_key(&mut security_keys[idx], &job) {
                            done = true;
                            break;
                        }
                    }

                    if done || job.start_time.elapsed() > job.timeout {
                        None
                    } else {
                        Some(job)
                    }
                },
                None => None, // Nothing to do
            }

        }

        // Err(io::Error::new(io::ErrorKind::TimedOut, "Timed out"))
    }

    pub fn perform_job_for_key<T>(dev: &mut T, job: &WorkUnit) -> io::Result<()>
        where T: U2FDevice + Read + Write
    {
        match job.command {
            Command::Register => {
                if let Ok(bytes) = u2fhid::u2f_register(dev, &job.challenge, &job.application) {
                    println!("Register OK {:?}", bytes);
                    // First to complete, we return
                    job.result_tx.send(Ok(bytes));
                    return Ok(())
                }
            },
            Command::Sign => {
                // It'd be an error if key_handle was None here
                let keybytes = job.key_handle.as_ref().unwrap();

                if let Ok(bytes) = u2fhid::u2f_sign(dev, &job.challenge, &job.application, keybytes) {
                    println!("Sign OK {:?}", bytes);
                    // First to complete, we return
                    job.result_tx.send(Ok(bytes));
                    return Ok(())
                }
            },
        };

        Ok(())
    }

    pub fn register<F>(&self, timeout_sec: u8, challenge: Vec<u8>, application: Vec<u8>, callback: F)
        where F: FnOnce(io::Result<Vec<u8>>), F: Send + 'static
    {
        if challenge.len() != PARAMETER_SIZE || application.len() != PARAMETER_SIZE {
            callback(Err(io::Error::new(io::ErrorKind::InvalidInput, "Invalid parameter sizes")));
            return;
        }

        let (result_tx, result_rx) = channel::<io::Result<Vec<u8>>>();

        if let Err(e) = self.work_tx.send(WorkUnit{
            timeout: Duration::from_secs(timeout_sec as u64),
            start_time: Instant::now(),
            command: Command::Register,
            challenge: challenge,
            application: application,
            key_handle: None,
            result_tx: result_tx,
        }) {
            callback(Err(io::Error::new(io::ErrorKind::Other, format!("Send error {}", e))));
            return;
        }

        match result_rx.recv() {
            Ok(v) => callback(v),
            Err(e) => callback(Err(io::Error::new(io::ErrorKind::Other, e))),
        }
    }

    pub fn sign<F>(&self, timeout_sec: u8, challenge: Vec<u8>, application: Vec<u8>, key_handle: Vec<u8>, callback: F)
        where F: FnOnce(io::Result<Vec<u8>>), F: Send + 'static
    {
        if challenge.len() != PARAMETER_SIZE || application.len() != PARAMETER_SIZE {
            callback(Err(io::Error::new(io::ErrorKind::InvalidInput, "Invalid parameter sizes")));
            return;
        }

    }
}

fn u2f_get_key_handle_from_register_response(register_response: &Vec<u8>) -> io::Result<Vec<u8>>
{
    if register_response[0] != 0x05 {
        return Err(io::Error::new(io::ErrorKind::InvalidData, "Reserved byte not set correctly"));
    }

    let key_handle_len = register_response[66] as usize;
    let mut public_key = register_response.clone();
    let mut key_handle = public_key.split_off(67);
    // let attestation = key_handle.split_off(key_handle_len);

    Ok(key_handle)
}

fn main() {
    println!("Searching for keys...");

    let mut challenge = Sha256::new();
    challenge.input_str(r#"{"challenge": "1vQ9mxionq0ngCnjD-wTsv1zUSrGRtFqG2xP09SbZ70", "version": "U2F_V2", "appId": "http://demo.yubico.com"}"#);
    let mut chall_bytes: Vec<u8> = vec![0; challenge.output_bytes()];
    challenge.result(&mut chall_bytes);

    let mut application = Sha256::new();
    application.input_str("http://demo.yubico.com");
    let mut app_bytes: Vec<u8> = vec![0; application.output_bytes()];
    application.result(&mut app_bytes);

    let (tx, rx) = channel();
    let manager = open_u2f_manager().unwrap();

    manager.register(15, chall_bytes, app_bytes, move |result| {
        // Ship back to the main thread
        if let Err(e) = tx.send(result) {
            panic!("Could not send: {}", e);
        }
    });

    let thread_result = match rx.recv() {
        Ok(v) => v,
        Err(e) => panic!("Couldn't read data: {}", e),
    };

    let register_data = match thread_result {
        Ok(v) => v,
        Err(e) => panic!("Register failure: {}", e),
    };

    println!("Register result: {}", base64::encode(&register_data));

    // let key_handle = u2f_get_key_handle_from_register_response(&register_data).unwrap();

    // let mut chall_bytes: Vec<u8> = vec![0; challenge.output_bytes()];
    // challenge.result(&mut chall_bytes);
    // let mut app_bytes: Vec<u8> = vec![0; application.output_bytes()];
    // application.result(&mut app_bytes);

    // manager.sign(15, chall_bytes, app_bytes, key_handle, move|result| {
    //     println!("Sign result: {}", base64::encode(&result.unwrap()));
    // });

    println!("Done.");
}
<|endoftext|>