<commit_before>//! Test command for testing the binary machine code emission.
//!
//! The `binemit` test command generates binary machine code for every instruction in the input
//! functions and compares the results to the expected output.

use cretonne_codegen::binemit;
use cretonne_codegen::binemit::RegDiversions;
use cretonne_codegen::dbg::DisplayList;
use cretonne_codegen::ir;
use cretonne_codegen::ir::entities::AnyEntity;
use cretonne_codegen::print_errors::pretty_error;
use cretonne_codegen::settings::OptLevel;
use cretonne_reader::TestCommand;
use match_directive::match_directive;
use std::borrow::Cow;
use std::collections::HashMap;
use std::fmt::Write;
use subtest::{Context, Result, SubTest};

struct TestBinEmit;

pub fn subtest(parsed: &TestCommand) -> Result<Box<SubTest>> {
    assert_eq!(parsed.command, "binemit");
    if !parsed.options.is_empty() {
        Err(format!("No options allowed on {}", parsed))
    } else {
        Ok(Box::new(TestBinEmit))
    }
}

/// Code sink that generates text.
struct TextSink {
    offset: binemit::CodeOffset,
    text: String,
}

impl TextSink {
    /// Create a new empty TextSink.
    pub fn new() -> Self {
        Self {
            offset: 0,
            text: String::new(),
        }
    }
}

impl binemit::CodeSink for TextSink {
    fn offset(&self) -> binemit::CodeOffset {
        self.offset
    }

    fn put1(&mut self, x: u8) {
        write!(self.text, "{:02x} ", x).unwrap();
        self.offset += 1;
    }

    fn put2(&mut self, x: u16) {
        write!(self.text, "{:04x} ", x).unwrap();
        self.offset += 2;
    }

    fn put4(&mut self, x: u32) {
        write!(self.text, "{:08x} ", x).unwrap();
        self.offset += 4;
    }

    fn put8(&mut self, x: u64) {
        write!(self.text, "{:016x} ", x).unwrap();
        self.offset += 8;
    }

    fn reloc_ebb(&mut self, reloc: binemit::Reloc, ebb_offset: binemit::CodeOffset) {
        write!(self.text, "{}({}) ", reloc, ebb_offset).unwrap();
    }

    fn reloc_external(
        &mut self,
        reloc: binemit::Reloc,
        name: &ir::ExternalName,
        addend: binemit::Addend,
    ) {
        write!(self.text, "{}({}", reloc, name,).unwrap();
        if addend != 0 {
            write!(self.text, "{:+}", addend,).unwrap();
        }
        write!(self.text, ") ",).unwrap();
    }

    fn reloc_jt(&mut self, reloc: binemit::Reloc, jt: ir::JumpTable) {
        write!(self.text, "{}({}) ", reloc, jt).unwrap();
    }

    fn trap(&mut self, code: ir::TrapCode, _srcloc: ir::SourceLoc) {
        write!(self.text, "{} ", code).unwrap();
    }
}

impl SubTest for TestBinEmit {
    fn name(&self) -> Cow<str> {
        Cow::from("binemit")
    }

    fn is_mutating(&self) -> bool {
        true
    }

    fn needs_isa(&self) -> bool {
        true
    }

    fn run(&self, func: Cow<ir::Function>, context: &Context) -> Result<()> {
        let isa = context.isa.expect("binemit needs an ISA");
        let encinfo = isa.encoding_info();
        // TODO: Run a verifier pass over the code first to detect any bad encodings or missing/bad
        // value locations. The current error reporting is just crashing...
        let mut func = func.into_owned();

        // Fix the stack frame layout so we can test spill/fill encodings.
        let min_offset = func.stack_slots
            .values()
            .map(|slot| slot.offset.unwrap())
            .min();
        func.stack_slots.frame_size = min_offset.map(|off| (-off) as u32);

        let opt_level = isa.flags().opt_level();

        // Give an encoding to any instruction that doesn't already have one.
        let mut divert = RegDiversions::new();
        for ebb in func.layout.ebbs() {
            divert.clear();
            for inst in func.layout.ebb_insts(ebb) {
                if !func.encodings[inst].is_legal() {
                    // Find an encoding that satisfies both immediate field and register
                    // constraints.
                    if let Some(enc) = {
                        let mut legal_encodings = isa.legal_encodings(
                            &func,
                            &func.dfg[inst],
                            func.dfg.ctrl_typevar(inst),
                        ).filter(|e| {
                                let recipe_constraints = &encinfo.constraints[e.recipe()];
                                recipe_constraints.satisfied(inst, &divert, &func)
                            });

                        if opt_level == OptLevel::Best {
                            // Get the smallest legal encoding
                            legal_encodings.min_by_key(|&e| encinfo.bytes(e))
                        } else {
                            // If not optimizing, just use the first encoding.
                            legal_encodings.next()
                        }
                    }
                    {
                        func.encodings[inst] = enc;
                    }
                }
                divert.apply(&func.dfg[inst]);
            }
        }

        // Relax branches and compute EBB offsets based on the encodings.
        let code_size = binemit::relax_branches(&mut func, isa).map_err(|e| {
            pretty_error(&func, context.isa, e)
        })?;

        // Collect all of the 'bin:' directives on instructions.
        let mut bins = HashMap::new();
        for comment in &context.details.comments {
            if let Some(want) = match_directive(comment.text, "bin:") {
                match comment.entity {
                    AnyEntity::Inst(inst) => {
                        if let Some(prev) = bins.insert(inst, want) {
                            return Err(format!(
                                "multiple 'bin:' directives on {}: '{}' and '{}'",
                                func.dfg.display_inst(inst, isa),
                                prev,
                                want
                            ));
                        }
                    }
                    _ => {
                        return Err(format!(
                            "'bin:' directive on non-inst {}: {}",
                            comment.entity,
                            comment.text
                        ))
                    }
                }
            }
        }
        if bins.is_empty() {
            return Err("No 'bin:' directives found".to_string());
        }

        // Now emit all instructions.
        let mut sink = TextSink::new();
        for ebb in func.layout.ebbs() {
            divert.clear();
            // Correct header offsets should have been computed by `relax_branches()`.
            assert_eq!(
                sink.offset,
                func.offsets[ebb],
                "Inconsistent {} header offset",
                ebb
            );
            for (offset, inst, enc_bytes) in func.inst_offsets(ebb, &encinfo) {
                assert_eq!(sink.offset, offset);
                sink.text.clear();
                let enc = func.encodings[inst];

                // Send legal encodings into the emitter.
                if enc.is_legal() {
                    // Generate a better error message if output locations are not specified.
                    if let Some(&v) = func.dfg.inst_results(inst).iter().find(|&&v| {
                        !func.locations[v].is_assigned()
                    })
                    {
                        return Err(format!(
                            "Missing register/stack slot for {} in {}",
                            v,
                            func.dfg.display_inst(inst, isa)
                        ));
                    }
                    let before = sink.offset;
                    isa.emit_inst(&func, inst, &mut divert, &mut sink);
                    let emitted = sink.offset - before;
                    // Verify the encoding recipe sizes against the ISAs emit_inst implementation.
                    assert_eq!(
                        emitted,
                        enc_bytes,
                        "Inconsistent size for [{}] {}",
                        encinfo.display(enc),
                        func.dfg.display_inst(inst, isa)
                    );
                }

                // Check against bin: directives.
                if let Some(want) = bins.remove(&inst) {
                    if !enc.is_legal() {
                        // A possible cause of an unencoded instruction is a missing location for
                        // one of the input operands.
                        if let Some(&v) = func.dfg.inst_args(inst).iter().find(|&&v| {
                            !func.locations[v].is_assigned()
                        })
                        {
                            return Err(format!(
                                "Missing register/stack slot for {} in {}",
                                v,
                                func.dfg.display_inst(inst, isa)
                            ));
                        }

                        // Do any encodings exist?
                        let encodings = isa.legal_encodings(
                            &func,
                            &func.dfg[inst],
                            func.dfg.ctrl_typevar(inst),
                        ).map(|e| encinfo.display(e))
                            .collect::<Vec<_>>();

                        if encodings.is_empty() {
                            return Err(format!(
                                "No encodings found for: {}",
                                func.dfg.display_inst(inst, isa)
                            ));
                        }
                        return Err(format!(
                            "No matching encodings for {} in {}",
                            func.dfg.display_inst(inst, isa),
                            DisplayList(&encodings),
                        ));
                    }
                    let have = sink.text.trim();
                    if have != want {
                        return Err(format!(
                            "Bad machine code for {}: {}\nWant: {}\nGot:  {}",
                            inst,
                            func.dfg.display_inst(inst, isa),
                            want,
                            have
                        ));
                    }
                }
            }
        }

        if sink.offset != code_size {
            return Err(format!(
                "Expected code size {}, got {}",
                code_size,
                sink.offset
            ));
        }

        Ok(())
    }
}
<commit_msg>Tidy up unneeded commmas.<commit_after>//! Test command for testing the binary machine code emission.
//!
//! The `binemit` test command generates binary machine code for every instruction in the input
//! functions and compares the results to the expected output.

use cretonne_codegen::binemit;
use cretonne_codegen::binemit::RegDiversions;
use cretonne_codegen::dbg::DisplayList;
use cretonne_codegen::ir;
use cretonne_codegen::ir::entities::AnyEntity;
use cretonne_codegen::print_errors::pretty_error;
use cretonne_codegen::settings::OptLevel;
use cretonne_reader::TestCommand;
use match_directive::match_directive;
use std::borrow::Cow;
use std::collections::HashMap;
use std::fmt::Write;
use subtest::{Context, Result, SubTest};

struct TestBinEmit;

pub fn subtest(parsed: &TestCommand) -> Result<Box<SubTest>> {
    assert_eq!(parsed.command, "binemit");
    if !parsed.options.is_empty() {
        Err(format!("No options allowed on {}", parsed))
    } else {
        Ok(Box::new(TestBinEmit))
    }
}

/// Code sink that generates text.
struct TextSink {
    offset: binemit::CodeOffset,
    text: String,
}

impl TextSink {
    /// Create a new empty TextSink.
    pub fn new() -> Self {
        Self {
            offset: 0,
            text: String::new(),
        }
    }
}

impl binemit::CodeSink for TextSink {
    fn offset(&self) -> binemit::CodeOffset {
        self.offset
    }

    fn put1(&mut self, x: u8) {
        write!(self.text, "{:02x} ", x).unwrap();
        self.offset += 1;
    }

    fn put2(&mut self, x: u16) {
        write!(self.text, "{:04x} ", x).unwrap();
        self.offset += 2;
    }

    fn put4(&mut self, x: u32) {
        write!(self.text, "{:08x} ", x).unwrap();
        self.offset += 4;
    }

    fn put8(&mut self, x: u64) {
        write!(self.text, "{:016x} ", x).unwrap();
        self.offset += 8;
    }

    fn reloc_ebb(&mut self, reloc: binemit::Reloc, ebb_offset: binemit::CodeOffset) {
        write!(self.text, "{}({}) ", reloc, ebb_offset).unwrap();
    }

    fn reloc_external(
        &mut self,
        reloc: binemit::Reloc,
        name: &ir::ExternalName,
        addend: binemit::Addend,
    ) {
        write!(self.text, "{}({}", reloc, name).unwrap();
        if addend != 0 {
            write!(self.text, "{:+}", addend).unwrap();
        }
        write!(self.text, ") ").unwrap();
    }

    fn reloc_jt(&mut self, reloc: binemit::Reloc, jt: ir::JumpTable) {
        write!(self.text, "{}({}) ", reloc, jt).unwrap();
    }

    fn trap(&mut self, code: ir::TrapCode, _srcloc: ir::SourceLoc) {
        write!(self.text, "{} ", code).unwrap();
    }
}

impl SubTest for TestBinEmit {
    fn name(&self) -> Cow<str> {
        Cow::from("binemit")
    }

    fn is_mutating(&self) -> bool {
        true
    }

    fn needs_isa(&self) -> bool {
        true
    }

    fn run(&self, func: Cow<ir::Function>, context: &Context) -> Result<()> {
        let isa = context.isa.expect("binemit needs an ISA");
        let encinfo = isa.encoding_info();
        // TODO: Run a verifier pass over the code first to detect any bad encodings or missing/bad
        // value locations. The current error reporting is just crashing...
        let mut func = func.into_owned();

        // Fix the stack frame layout so we can test spill/fill encodings.
        let min_offset = func.stack_slots
            .values()
            .map(|slot| slot.offset.unwrap())
            .min();
        func.stack_slots.frame_size = min_offset.map(|off| (-off) as u32);

        let opt_level = isa.flags().opt_level();

        // Give an encoding to any instruction that doesn't already have one.
        let mut divert = RegDiversions::new();
        for ebb in func.layout.ebbs() {
            divert.clear();
            for inst in func.layout.ebb_insts(ebb) {
                if !func.encodings[inst].is_legal() {
                    // Find an encoding that satisfies both immediate field and register
                    // constraints.
                    if let Some(enc) = {
                        let mut legal_encodings = isa.legal_encodings(
                            &func,
                            &func.dfg[inst],
                            func.dfg.ctrl_typevar(inst),
                        ).filter(|e| {
                                let recipe_constraints = &encinfo.constraints[e.recipe()];
                                recipe_constraints.satisfied(inst, &divert, &func)
                            });

                        if opt_level == OptLevel::Best {
                            // Get the smallest legal encoding
                            legal_encodings.min_by_key(|&e| encinfo.bytes(e))
                        } else {
                            // If not optimizing, just use the first encoding.
                            legal_encodings.next()
                        }
                    }
                    {
                        func.encodings[inst] = enc;
                    }
                }
                divert.apply(&func.dfg[inst]);
            }
        }

        // Relax branches and compute EBB offsets based on the encodings.
        let code_size = binemit::relax_branches(&mut func, isa).map_err(|e| {
            pretty_error(&func, context.isa, e)
        })?;

        // Collect all of the 'bin:' directives on instructions.
        let mut bins = HashMap::new();
        for comment in &context.details.comments {
            if let Some(want) = match_directive(comment.text, "bin:") {
                match comment.entity {
                    AnyEntity::Inst(inst) => {
                        if let Some(prev) = bins.insert(inst, want) {
                            return Err(format!(
                                "multiple 'bin:' directives on {}: '{}' and '{}'",
                                func.dfg.display_inst(inst, isa),
                                prev,
                                want
                            ));
                        }
                    }
                    _ => {
                        return Err(format!(
                            "'bin:' directive on non-inst {}: {}",
                            comment.entity,
                            comment.text
                        ))
                    }
                }
            }
        }
        if bins.is_empty() {
            return Err("No 'bin:' directives found".to_string());
        }

        // Now emit all instructions.
        let mut sink = TextSink::new();
        for ebb in func.layout.ebbs() {
            divert.clear();
            // Correct header offsets should have been computed by `relax_branches()`.
            assert_eq!(
                sink.offset,
                func.offsets[ebb],
                "Inconsistent {} header offset",
                ebb
            );
            for (offset, inst, enc_bytes) in func.inst_offsets(ebb, &encinfo) {
                assert_eq!(sink.offset, offset);
                sink.text.clear();
                let enc = func.encodings[inst];

                // Send legal encodings into the emitter.
                if enc.is_legal() {
                    // Generate a better error message if output locations are not specified.
                    if let Some(&v) = func.dfg.inst_results(inst).iter().find(|&&v| {
                        !func.locations[v].is_assigned()
                    })
                    {
                        return Err(format!(
                            "Missing register/stack slot for {} in {}",
                            v,
                            func.dfg.display_inst(inst, isa)
                        ));
                    }
                    let before = sink.offset;
                    isa.emit_inst(&func, inst, &mut divert, &mut sink);
                    let emitted = sink.offset - before;
                    // Verify the encoding recipe sizes against the ISAs emit_inst implementation.
                    assert_eq!(
                        emitted,
                        enc_bytes,
                        "Inconsistent size for [{}] {}",
                        encinfo.display(enc),
                        func.dfg.display_inst(inst, isa)
                    );
                }

                // Check against bin: directives.
                if let Some(want) = bins.remove(&inst) {
                    if !enc.is_legal() {
                        // A possible cause of an unencoded instruction is a missing location for
                        // one of the input operands.
                        if let Some(&v) = func.dfg.inst_args(inst).iter().find(|&&v| {
                            !func.locations[v].is_assigned()
                        })
                        {
                            return Err(format!(
                                "Missing register/stack slot for {} in {}",
                                v,
                                func.dfg.display_inst(inst, isa)
                            ));
                        }

                        // Do any encodings exist?
                        let encodings = isa.legal_encodings(
                            &func,
                            &func.dfg[inst],
                            func.dfg.ctrl_typevar(inst),
                        ).map(|e| encinfo.display(e))
                            .collect::<Vec<_>>();

                        if encodings.is_empty() {
                            return Err(format!(
                                "No encodings found for: {}",
                                func.dfg.display_inst(inst, isa)
                            ));
                        }
                        return Err(format!(
                            "No matching encodings for {} in {}",
                            func.dfg.display_inst(inst, isa),
                            DisplayList(&encodings),
                        ));
                    }
                    let have = sink.text.trim();
                    if have != want {
                        return Err(format!(
                            "Bad machine code for {}: {}\nWant: {}\nGot:  {}",
                            inst,
                            func.dfg.display_inst(inst, isa),
                            want,
                            have
                        ));
                    }
                }
            }
        }

        if sink.offset != code_size {
            return Err(format!(
                "Expected code size {}, got {}",
                code_size,
                sink.offset
            ));
        }

        Ok(())
    }
}
<|endoftext|>


<commit_before>#[macro_export]
macro_rules! decl_enum {
    ($name:ident) => {
        #[derive(Clone, Eq, PartialEq, Debug)]
        pub struct $name(i32);
    }
}

#[macro_export]
macro_rules! impl_enum {
    ($name:ident; $($variant:ident = $value:expr,)*) => {
        #[allow(non_upper_case_globals)]
        impl $name {
            $(pub const $variant: i32 = $value;)*
        }
    }
}

#[macro_export]
macro_rules! make_enum {
    ($name:ident; $($variant:ident = $value:expr,)*) => {
        decl_enum!{$name}
        impl_enum!{$name; $($variant = $value,)*}
    }
}
<commit_msg>fix type of associated constants<commit_after>#[macro_export]
macro_rules! decl_enum {
    ($name:ident) => {
        #[derive(Clone, Eq, PartialEq, Debug)]
        pub struct $name(pub i32);
    }
}

#[macro_export]
macro_rules! impl_enum {
    ($name:ident; $($variant:ident = $value:expr,)*) => {
        #[allow(non_upper_case_globals)]
        impl $name {
            $(pub const $variant: $name = $name($value);)*
        }
    }
}

#[macro_export]
macro_rules! make_enum {
    ($name:ident; $($variant:ident = $value:expr,)*) => {
        decl_enum!{$name}
        impl_enum!{$name; $($variant = $value,)*}
    }
}
<|endoftext|>


<commit_before>use std::collections::{HashMap, HashSet};
use std::fmt;
use std::time::Duration;

use prost::Message;
use tokio::sync::oneshot;
use tracing::{debug, info, warn};

use crate::xds::istio::workload::Workload;
use crate::xds::service::discovery::v3::aggregated_discovery_service_client::AggregatedDiscoveryServiceClient;
use crate::xds::service::discovery::v3::Resource as ProtoResource;
use crate::xds::service::discovery::v3::*;
use crate::{tls, xds};

use super::Error;

#[derive(Eq, Hash, PartialEq, Debug, Clone)]
pub struct ResourceKey {
    pub name: String,
    pub type_url: String,
}

#[derive(Default)]
pub struct HandlerContext {
    rejects: Vec<RejectedConfig>,
}

impl HandlerContext {
    pub fn new() -> HandlerContext {
        HandlerContext {
            ..Default::default()
        }
    }
}

struct RejectedConfig {
    name: String,
    reason: anyhow::Error,
}

impl fmt::Display for RejectedConfig {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}: {}", self.name, self.reason,)
    }
}

impl HandlerContext {
    pub fn reject(&mut self, name: String, reason: anyhow::Error) {
        self.rejects.push(RejectedConfig { name, reason })
    }
}

pub trait Handler<T: prost::Message>: Send + Sync + 'static {
    fn handle(&self, ctx: &mut HandlerContext, res: Vec<XdsUpdate<T>>);
}

struct NopHandler {}

impl<T: prost::Message> Handler<T> for NopHandler {
    fn handle(&self, _ctx: &mut HandlerContext, _res: Vec<XdsUpdate<T>>) {}
}

pub struct Config {
    workload_handler: Box<dyn Handler<Workload>>,
    initial_watches: Vec<String>,
    on_demand: bool,
}

impl Config {
    pub fn new(config: crate::config::Config) -> Config {
        Config {
            workload_handler: Box::new(NopHandler {}),
            initial_watches: Vec::new(),
            on_demand: config.xds_on_demand,
        }
    }

    pub fn with_workload_handler(mut self, f: impl Handler<Workload>) -> Config {
        self.workload_handler = Box::new(f);
        self
    }

    pub fn watch(mut self, type_url: String) -> Config {
        self.initial_watches.push(type_url);
        self
    }

    pub fn build(self) -> AdsClient {
        let (tx, rx) = tokio::sync::mpsc::channel(100);
        AdsClient {
            config: self,
            workloads: HashSet::new(),
            pending: Default::default(),
            demand: rx,
            demand_tx: tx,
        }
    }
}

pub struct AdsClient {
    config: Config,
    /// Stores all known workload resources (by name)
    workloads: HashSet<String>,

    /// pending stores a list of all resources that are pending and XDS push
    pending: HashMap<ResourceKey, oneshot::Sender<()>>,

    demand: tokio::sync::mpsc::Receiver<(oneshot::Sender<()>, ResourceKey)>,
    demand_tx: tokio::sync::mpsc::Sender<(oneshot::Sender<()>, ResourceKey)>,
}

/// Demanded allows awaiting for an on-demand XDS resource
pub struct Demanded {
    b: oneshot::Receiver<()>,
}

impl Demanded {
    /// recv awaits for the requested resource
    /// Note: the actual resource is not directly returned. Instead, callers are notified that the event
    /// has been handled through the configured resource handler.
    pub async fn recv(self) {
        let _ = self.b.await;
    }
}

/// Demander allows requesting XDS resources on-demand
#[derive(Debug, Clone)]
pub struct Demander {
    demand: tokio::sync::mpsc::Sender<(oneshot::Sender<()>, ResourceKey)>,
}

impl Demander {
    /// Demand requests a given workload by name
    pub async fn demand(&self, name: String) -> Demanded {
        let (tx, rx) = oneshot::channel::<()>();
        self.demand
            .send((
                tx,
                ResourceKey {
                    name,
                    type_url: xds::WORKLOAD_TYPE.to_string(),
                },
            ))
            .await
            .unwrap();
        Demanded { b: rx }
    }
}

impl AdsClient {
    /// demander returns a Demander instance which can be used to request resources on-demand
    pub fn demander(&self) -> Demander {
        Demander {
            demand: self.demand_tx.clone(),
        }
    }

    pub async fn run(mut self) -> Result<(), Error> {
        let mut backoff = Duration::from_millis(10);
        let max_backoff = Duration::from_secs(15);
        loop {
            let res = self.run_internal().await;
            match res {
                Err(e @ Error::Connection(_)) => {
                    // For connection errors, we add backoff
                    backoff = std::cmp::min(max_backoff, backoff * 2);
                    warn!("XDS client error: {}, retrying in {:?}", e, backoff);
                    tokio::time::sleep(backoff).await;
                }
                Err(e) => {
                    // For other errors, we connect immediately
                    // TODO: we may need more nuance here; if we fail due to invalid initial request we may overload
                    // But we want to reconnect from MaxConnectionAge immediately.
                    warn!("XDS client error: {}, retrying", e);
                }
                Ok(_) => {
                    warn!("XDS client complete");
                    backoff = Duration::from_millis(10);
                }
            }
        }
    }

    async fn run_internal(&mut self) -> Result<(), Error> {
        let initial_watches = &self.config.initial_watches;
        let workload_handler = &self.config.workload_handler;
        let workloads = &mut self.workloads;
        match workloads.len() {
            0 => info!("Starting initial ADS client"),
            n => info!("Starting ADS client with {n} workloads"),
        };

        // TODO: copy JWT auth logic from CA client and use TLS here
        let _address = if std::env::var("KUBERNETES_SERVICE_HOST").is_ok() {
            "https://istiod.istio-system:15012"
        } else {
            "https://localhost:15012"
        };
        let address = if std::env::var("KUBERNETES_SERVICE_HOST").is_ok() {
            "http://istiod.istio-system:15010"
        } else {
            "http://localhost:15010"
        };
        let svc = tls::grpc_connector(address).unwrap();
        let mut client = AggregatedDiscoveryServiceClient::new(svc);
        let (discovery_req_tx, mut discovery_req_rx) =
            tokio::sync::mpsc::channel::<DeltaDiscoveryRequest>(100);
        let watches = initial_watches.clone();
        let irv: HashMap<String, String> = workloads
            .iter()
            .map(|n| (n.clone(), "".to_string()))
            .collect();

        let (sub, unsub) = if self.config.on_demand {
            // XDS doesn't have a way to subscribe to zero resources. We workaround this by subscribing and unsubscribing
            // in one event, effectively giving us "subscribe to nothing".
            (vec!["*".to_string()], vec!["*".to_string()])
        } else {
            (vec![], vec![])
        };
        let outbound = async_stream::stream! {
            for request_type in watches {
                let irv = irv.clone();
                let initial = DeltaDiscoveryRequest {
                    type_url: request_type.clone(),
                    node: Some(Node{
                        id: format!("{}~{}~{}.{}~{}.svc.cluster.local", "sidecar", "1.1.1.1", "test", "test", "test"),
                        ..Default::default()
                    }),
                    initial_resource_versions: irv,
                    resource_names_subscribe: sub.clone(),
                    resource_names_unsubscribe: unsub.clone(),
                    ..Default::default()
                };
                info!(type_url=request_type, "sending initial request");
                yield initial;
            }
            while let Some(message) = discovery_req_rx.recv().await {
                info!(type_url=message.type_url, "sending request");
                yield message
            }
            warn!("outbound stream complete");
        };

        info!("Starting stream");
        let mut respones_stream = client
            .delta_aggregated_resources(tonic::Request::new(outbound))
            .await
            .map_err(Error::Connection)?
            .into_inner();

        info!("Stream established");

        let demand = &mut self.demand;
        loop {
            tokio::select! {
                demand_event = demand.recv() => {
                    info!("received on demand request {demand_event:?}");
                    let (tx, demand_event) = match demand_event {
                        None => continue,
                        Some(d) => d,
                    };
                    let ResourceKey { type_url, name } = demand_event.clone();
                    self.pending.insert(demand_event, tx);
                    discovery_req_tx
                        .send(DeltaDiscoveryRequest {
                            type_url,
                            node: Some(Node {
                                id: format!(
                                    "{}~{}~{}.{}~{}.svc.cluster.local",
                                    "sidecar", "1.1.1.1", "test", "test", "test"
                                ),
                                ..Default::default()
                            }),
                            resource_names_subscribe: vec![name],
                            ..Default::default()
                        })
                        .await
                        .map_err(|e| Error::RequestFailure(Box::new(e)))?;
                }
                msg = respones_stream.message() =>{
                    let response = match msg? {
                        None => break,
                        Some(r) => r
                    };
                    let type_url = response.type_url;
                    info!(
                        type_url = type_url.clone(),
                        size = response.resources.len(),
                        "received response"
                    );
                    let mut updates: Vec<XdsUpdate<Workload>> = Vec::new();
                    for res in response.resources.iter().cloned() {
                        let resource = Resource::try_from(res).unwrap();
                        let _cpy = resource.clone();
                        match resource {
                            Resource::Workload(inner) => {
                                let name = inner.name.clone();
                                let up = XdsUpdate::Update(inner);
                                updates.push(up);
                                workloads.insert(name.clone());
                                let pending = {
                                    self.pending.remove(&ResourceKey {
                                        name: name.clone(),
                                        type_url: type_url.clone(),
                                    })
                                };
                                if let Some(send) = pending {
                                debug!("on demand notify {}", name);
                                    send.send(()).map_err(|_| Error::OnDemandSend())?;
                                }
                            }
                        }
                    }
                    for res in response.removed_resources {
                        workloads.remove(&res.clone());
                        debug!("received delete resource {:#?}", res);
                            let pending = {
                                self.pending.remove(&ResourceKey {
                                    name: res.clone(),
                                    type_url: type_url.clone(),
                                })
                            };
                            if let Some(send) = pending {
                                debug!("on demand notify {}", res);
                                send.send(()).map_err(|_| Error::OnDemandSend())?;
                            }
                        match type_url.clone().as_str() {
                            xds::WORKLOAD_TYPE => {
                                updates.push(XdsUpdate::Remove(res));
                            }
                            _ => warn!("ignoring unwatched type {}", type_url),
                        }
                    }

                    let mut ctx = HandlerContext::new();
                    workload_handler.handle(&mut ctx, updates);

                    let error_detail = match ctx.rejects.len() {
                        0 => None,
                        _ => Some(
                            ctx.rejects
                                .into_iter()
                                .map(|reject| reject.to_string())
                                .collect::<Vec<String>>()
                                .join("; "),
                        ),
                    };
                    info!(
                        type_url = type_url.clone(),
                        none = response.nonce,
                        "sending {}",
                        if error_detail.is_some() {
                            "NACK"
                        } else {
                            "ACK"
                        }
                    );
                    discovery_req_tx
                        .send(DeltaDiscoveryRequest {
                            type_url: type_url.clone(),
                            node: Some(Node {
                                id: format!(
                                    "{}~{}~{}.{}~{}.svc.cluster.local",
                                    "sidecar", "1.1.1.1", "test", "test", "test"
                                ),
                                ..Default::default()
                            }),
                            response_nonce: response.nonce.clone(),
                            error_detail: error_detail.map(|msg| Status {
                                message: msg,
                                ..Default::default()
                            }),
                            ..Default::default()
                        })
                        .await
                        .map_err(|e| Error::RequestFailure(Box::new(e)))?;

                }

            }
        }
        info!("Stream terminate");
        Ok(())
    }
}

#[derive(Clone, Debug)]
pub struct XdsResource<T: prost::Message> {
    pub name: String,
    pub resource: T,
}

#[derive(Debug)]
pub enum XdsUpdate<T: prost::Message> {
    Update(XdsResource<T>),
    Remove(String),
}

impl<T: prost::Message> XdsUpdate<T> {
    pub fn name(&self) -> String {
        match self {
            XdsUpdate::Update(ref r) => r.name.clone(),
            XdsUpdate::Remove(n) => n.to_string(),
        }
    }
}

// TODO: consider using https://crates.io/crates/prost-reflect to dynamically create types allowing any
// compiled protos. For now, we just hardcode
#[derive(Clone, Debug)]
pub enum Resource {
    Workload(XdsResource<Workload>),
}

impl TryFrom<ProtoResource> for Resource {
    type Error = AdsError;

    fn try_from(resource: ProtoResource) -> Result<Self, Self::Error> {
        let res = resource.resource.unwrap();
        Ok(match &*res.type_url {
            xds::WORKLOAD_TYPE => {
                let inner = <Workload>::decode(&*res.value).unwrap();
                Resource::Workload(XdsResource {
                    name: resource.name,
                    resource: inner,
                })
            }
            url => return Err(AdsError::UnknownResourceType(url.into())),
        })
    }
}

#[derive(Clone, Debug)]
pub enum AdsError {
    UnknownResourceType(String),
}
<commit_msg>xds: break out huge tokio::select statement (#57)<commit_after>use std::collections::{HashMap, HashSet};
use std::fmt;
use std::time::Duration;

use prost::Message;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tracing::{debug, info, warn};

use crate::xds::istio::workload::Workload;
use crate::xds::service::discovery::v3::aggregated_discovery_service_client::AggregatedDiscoveryServiceClient;
use crate::xds::service::discovery::v3::Resource as ProtoResource;
use crate::xds::service::discovery::v3::*;
use crate::{tls, xds};

use super::Error;

#[derive(Eq, Hash, PartialEq, Debug, Clone)]
pub struct ResourceKey {
    pub name: String,
    pub type_url: String,
}

#[derive(Default)]
pub struct HandlerContext {
    rejects: Vec<RejectedConfig>,
}

impl HandlerContext {
    pub fn new() -> HandlerContext {
        HandlerContext {
            ..Default::default()
        }
    }
}

struct RejectedConfig {
    name: String,
    reason: anyhow::Error,
}

impl fmt::Display for RejectedConfig {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}: {}", self.name, self.reason,)
    }
}

impl HandlerContext {
    pub fn reject(&mut self, name: String, reason: anyhow::Error) {
        self.rejects.push(RejectedConfig { name, reason })
    }
}

pub trait Handler<T: prost::Message>: Send + Sync + 'static {
    fn handle(&self, ctx: &mut HandlerContext, res: Vec<XdsUpdate<T>>);
}

struct NopHandler {}

impl<T: prost::Message> Handler<T> for NopHandler {
    fn handle(&self, _ctx: &mut HandlerContext, _res: Vec<XdsUpdate<T>>) {}
}

pub struct Config {
    workload_handler: Box<dyn Handler<Workload>>,
    initial_watches: Vec<String>,
    on_demand: bool,
}

impl Config {
    pub fn new(config: crate::config::Config) -> Config {
        Config {
            workload_handler: Box::new(NopHandler {}),
            initial_watches: Vec::new(),
            on_demand: config.xds_on_demand,
        }
    }

    pub fn with_workload_handler(mut self, f: impl Handler<Workload>) -> Config {
        self.workload_handler = Box::new(f);
        self
    }

    pub fn watch(mut self, type_url: String) -> Config {
        self.initial_watches.push(type_url);
        self
    }

    pub fn build(self) -> AdsClient {
        let (tx, rx) = mpsc::channel(100);
        AdsClient {
            config: self,
            workloads: HashSet::new(),
            pending: Default::default(),
            demand: rx,
            demand_tx: tx,
        }
    }
}

pub struct AdsClient {
    config: Config,
    /// Stores all known workload resources (by name)
    workloads: HashSet<String>,

    /// pending stores a list of all resources that are pending and XDS push
    pending: HashMap<ResourceKey, oneshot::Sender<()>>,

    demand: mpsc::Receiver<(oneshot::Sender<()>, ResourceKey)>,
    demand_tx: mpsc::Sender<(oneshot::Sender<()>, ResourceKey)>,
}

/// Demanded allows awaiting for an on-demand XDS resource
pub struct Demanded {
    b: oneshot::Receiver<()>,
}

impl Demanded {
    /// recv awaits for the requested resource
    /// Note: the actual resource is not directly returned. Instead, callers are notified that the event
    /// has been handled through the configured resource handler.
    pub async fn recv(self) {
        let _ = self.b.await;
    }
}

/// Demander allows requesting XDS resources on-demand
#[derive(Debug, Clone)]
pub struct Demander {
    demand: mpsc::Sender<(oneshot::Sender<()>, ResourceKey)>,
}

impl Demander {
    /// Demand requests a given workload by name
    pub async fn demand(&self, name: String) -> Demanded {
        let (tx, rx) = oneshot::channel::<()>();
        self.demand
            .send((
                tx,
                ResourceKey {
                    name,
                    type_url: xds::WORKLOAD_TYPE.to_string(),
                },
            ))
            .await
            .unwrap();
        Demanded { b: rx }
    }
}

impl AdsClient {
    /// demander returns a Demander instance which can be used to request resources on-demand
    pub fn demander(&self) -> Demander {
        Demander {
            demand: self.demand_tx.clone(),
        }
    }

    pub async fn run(mut self) -> Result<(), Error> {
        let mut backoff = Duration::from_millis(10);
        let max_backoff = Duration::from_secs(15);
        loop {
            let res = self.run_internal().await;
            match res {
                Err(e @ Error::Connection(_)) => {
                    // For connection errors, we add backoff
                    backoff = std::cmp::min(max_backoff, backoff * 2);
                    warn!("XDS client error: {}, retrying in {:?}", e, backoff);
                    tokio::time::sleep(backoff).await;
                }
                Err(e) => {
                    // For other errors, we connect immediately
                    // TODO: we may need more nuance here; if we fail due to invalid initial request we may overload
                    // But we want to reconnect from MaxConnectionAge immediately.
                    warn!("XDS client error: {}, retrying", e);
                    backoff = Duration::from_millis(10);
                }
                Ok(_) => {
                    warn!("XDS client complete");
                    backoff = Duration::from_millis(10);
                }
            }
        }
    }

    async fn run_internal(&mut self) -> Result<(), Error> {
        let initial_watches = &self.config.initial_watches;
        let workloads = &mut self.workloads;
        match workloads.len() {
            0 => info!("Starting initial ADS client"),
            n => info!("Starting ADS client with {n} workloads"),
        };

        // TODO: copy JWT auth logic from CA client and use TLS here
        let _address = if std::env::var("KUBERNETES_SERVICE_HOST").is_ok() {
            "https://istiod.istio-system:15012"
        } else {
            "https://localhost:15012"
        };
        let address = if std::env::var("KUBERNETES_SERVICE_HOST").is_ok() {
            "http://istiod.istio-system:15010"
        } else {
            "http://localhost:15010"
        };
        let svc = tls::grpc_connector(address).unwrap();
        let mut client = AggregatedDiscoveryServiceClient::new(svc);
        let (discovery_req_tx, mut discovery_req_rx) = mpsc::channel::<DeltaDiscoveryRequest>(100);
        let watches = initial_watches.clone();
        let irv: HashMap<String, String> = workloads
            .iter()
            .map(|n| (n.clone(), "".to_string()))
            .collect();

        let (sub, unsub) = if self.config.on_demand {
            // XDS doesn't have a way to subscribe to zero resources. We workaround this by subscribing and unsubscribing
            // in one event, effectively giving us "subscribe to nothing".
            (vec!["*".to_string()], vec!["*".to_string()])
        } else {
            (vec![], vec![])
        };
        let outbound = async_stream::stream! {
            for request_type in watches {
                let irv = irv.clone();
                let initial = DeltaDiscoveryRequest {
                    type_url: request_type.clone(),
                    node: Some(Node{
                        id: format!("{}~{}~{}.{}~{}.svc.cluster.local", "sidecar", "1.1.1.1", "test", "test", "test"),
                        ..Default::default()
                    }),
                    initial_resource_versions: irv,
                    resource_names_subscribe: sub.clone(),
                    resource_names_unsubscribe: unsub.clone(),
                    ..Default::default()
                };
                info!(type_url=request_type, "sending initial request");
                yield initial;
            }
            while let Some(message) = discovery_req_rx.recv().await {
                info!(type_url=message.type_url, "sending request");
                yield message
            }
            warn!("outbound stream complete");
        };

        info!("Starting stream");
        let mut response_stream = client
            .delta_aggregated_resources(tonic::Request::new(outbound))
            .await
            .map_err(Error::Connection)?
            .into_inner();

        info!("Stream established");

        loop {
            tokio::select! {
                _demand_event = self.demand.recv() => {
                    self.handle_demand_event(None, &discovery_req_tx).await?;
                }
                msg = response_stream.message() =>{
                    self.handle_stream_event(msg?, &discovery_req_tx).await?;
                }

            }
        }
    }

    async fn handle_stream_event(
        &mut self,
        stream_event: Option<DeltaDiscoveryResponse>,
        send: &mpsc::Sender<DeltaDiscoveryRequest>,
    ) -> Result<(), Error> {
        let Some(response) = stream_event else {
            return Ok(());
        };
        let type_url = response.type_url;
        info!(
            type_url = type_url.clone(),
            size = response.resources.len(),
            "received response"
        );
        let mut updates: Vec<XdsUpdate<Workload>> = Vec::new();
        for res in response.resources.iter().cloned() {
            let resource = Resource::try_from(res).unwrap();
            let _cpy = resource.clone();
            match resource {
                Resource::Workload(inner) => {
                    let name = inner.name.clone();
                    let up = XdsUpdate::Update(inner);
                    updates.push(up);
                    self.workloads.insert(name.clone());
                    let pending = {
                        self.pending.remove(&ResourceKey {
                            name: name.clone(),
                            type_url: type_url.clone(),
                        })
                    };
                    if let Some(send) = pending {
                        debug!("on demand notify {}", name);
                        send.send(()).map_err(|_| Error::OnDemandSend())?;
                    }
                }
            }
        }
        for res in response.removed_resources {
            self.workloads.remove(&res.clone());
            debug!("received delete resource {:#?}", res);
            let pending = {
                self.pending.remove(&ResourceKey {
                    name: res.clone(),
                    type_url: type_url.clone(),
                })
            };
            if let Some(send) = pending {
                debug!("on demand notify {}", res);
                send.send(()).map_err(|_| Error::OnDemandSend())?;
            }
            match type_url.clone().as_str() {
                xds::WORKLOAD_TYPE => {
                    updates.push(XdsUpdate::Remove(res));
                }
                _ => warn!("ignoring unwatched type {}", type_url),
            }
        }

        let mut ctx = HandlerContext::new();
        self.config.workload_handler.handle(&mut ctx, updates);

        let error_detail = match ctx.rejects.len() {
            0 => None,
            _ => Some(
                ctx.rejects
                    .into_iter()
                    .map(|reject| reject.to_string())
                    .collect::<Vec<String>>()
                    .join("; "),
            ),
        };
        info!(
            type_url = type_url.clone(),
            none = response.nonce,
            "sending {}",
            if error_detail.is_some() {
                "NACK"
            } else {
                "ACK"
            }
        );
        send.send(DeltaDiscoveryRequest {
            type_url: type_url.clone(),
            node: Some(Node {
                id: format!(
                    "{}~{}~{}.{}~{}.svc.cluster.local",
                    "sidecar", "1.1.1.1", "test", "test", "test"
                ),
                ..Default::default()
            }),
            response_nonce: response.nonce.clone(),
            error_detail: error_detail.map(|msg| Status {
                message: msg,
                ..Default::default()
            }),
            ..Default::default()
        })
        .await
        .map_err(|e| Error::RequestFailure(Box::new(e)))
    }

    async fn handle_demand_event(
        &mut self,
        demand_event: Option<(oneshot::Sender<()>, ResourceKey)>,
        send: &mpsc::Sender<DeltaDiscoveryRequest>,
    ) -> Result<(), Error> {
        info!("received on demand request {demand_event:?}");
        let Some((tx, demand_event)) = demand_event else {
            return Ok(());
        };
        let ResourceKey { type_url, name } = demand_event.clone();
        self.pending.insert(demand_event, tx);
        send.send(DeltaDiscoveryRequest {
            type_url,
            node: Some(Node {
                id: format!(
                    "{}~{}~{}.{}~{}.svc.cluster.local",
                    "sidecar", "1.1.1.1", "test", "test", "test"
                ),
                ..Default::default()
            }),
            resource_names_subscribe: vec![name],
            ..Default::default()
        })
        .await
        .map_err(|e| Error::RequestFailure(Box::new(e)))?;
        Ok(())
    }
}

#[derive(Clone, Debug)]
pub struct XdsResource<T: prost::Message> {
    pub name: String,
    pub resource: T,
}

#[derive(Debug)]
pub enum XdsUpdate<T: prost::Message> {
    Update(XdsResource<T>),
    Remove(String),
}

impl<T: prost::Message> XdsUpdate<T> {
    pub fn name(&self) -> String {
        match self {
            XdsUpdate::Update(ref r) => r.name.clone(),
            XdsUpdate::Remove(n) => n.to_string(),
        }
    }
}

// TODO: consider using https://crates.io/crates/prost-reflect to dynamically create types allowing any
// compiled protos. For now, we just hardcode
#[derive(Clone, Debug)]
pub enum Resource {
    Workload(XdsResource<Workload>),
}

impl TryFrom<ProtoResource> for Resource {
    type Error = AdsError;

    fn try_from(resource: ProtoResource) -> Result<Self, Self::Error> {
        let res = resource.resource.unwrap();
        Ok(match &*res.type_url {
            xds::WORKLOAD_TYPE => {
                let inner = <Workload>::decode(&*res.value).unwrap();
                Resource::Workload(XdsResource {
                    name: resource.name,
                    resource: inner,
                })
            }
            url => return Err(AdsError::UnknownResourceType(url.into())),
        })
    }
}

#[derive(Clone, Debug)]
pub enum AdsError {
    UnknownResourceType(String),
}
<|endoftext|>


<commit_before>#![allow(dead_code, unused_imports, unused_variables)]
#[macro_use]
extern crate derive_more;

use std::path::PathBuf;

// Here just to make sure that this doesn't conflict with
// the derives in some way
use std::fmt::Binary;

#[derive(Display)]
struct MyInt(i32);

#[derive(Display)]
#[display(fmt = "({}, {})", x, y)]
struct Point2D {
    x: i32,
    y: i32,
}

#[derive(Display)]
enum E {
    Uint(u32),
    #[display(fmt = "I am B {:b}", i)]
    Binary {
        i: i8,
    },
    #[cfg(feature = "nightly")]
    #[display(fmt = "I am C {}", "_0.display()")]
    Path(PathBuf),
}

#[derive(Display)]
#[display(fmt = "Java EE")]
enum EE {
    A,
    B,
}

#[derive(Display)]
#[display(fmt = "Hello there!")]
union U {
    i: u32,
}

#[derive(Octal)]
#[octal(fmt = "7")]
struct S;

#[derive(UpperHex)]
#[upper_hex(fmt = "UpperHex")]
struct UH;

#[derive(Display)]
struct Unit;

#[derive(Display)]
struct UnitStruct {}

#[derive(Display)]
enum EmptyEnum {}

#[derive(Display)]
#[display(fmt = "Generic")]
struct Generic<T>(T);

#[test]
fn check_display() {
    assert_eq!(MyInt(-2).to_string(), "-2");
    assert_eq!(Point2D { x: 3, y: 4 }.to_string(), "(3, 4)");
    assert_eq!(E::Uint(2).to_string(), "2");
    assert_eq!(E::Binary { i: -2 }.to_string(), "I am B 11111110");
    #[cfg(feature = "nightly")]
    assert_eq!(E::Path("abc".into()).to_string(), "I am C abc");
    assert_eq!(EE::A.to_string(), "Java EE");
    assert_eq!(EE::B.to_string(), "Java EE");
    assert_eq!(U { i: 2 }.to_string(), "Hello there!");
    assert_eq!(format!("{:o}", S), "7");
    assert_eq!(format!("{:X}", UH), "UpperHex");
    assert_eq!(Unit.to_string(), "Unit");
    assert_eq!(UnitStruct {}.to_string(), "UnitStruct");
    assert_eq!(Generic(()).to_string(), "Generic");
}
<commit_msg>Add some extra tests for display<commit_after>#![allow(dead_code, unused_imports, unused_variables)]
#[macro_use]
extern crate derive_more;

use std::path::PathBuf;

// Here just to make sure that this doesn't conflict with
// the derives in some way
use std::fmt::Binary;

#[derive(Display, Octal, Binary)]
struct MyInt(i32);

#[derive(Display)]
#[display(fmt = "({}, {})", x, y)]
struct Point2D {
    x: i32,
    y: i32,
}

#[derive(Display)]
enum E {
    Uint(u32),
    #[display(fmt = "I am B {:b}", i)]
    Binary {
        i: i8,
    },
    #[cfg(feature = "nightly")]
    #[display(fmt = "I am C {}", "_0.display()")]
    Path(PathBuf),
}

#[derive(Display)]
#[display(fmt = "Java EE")]
enum EE {
    A,
    B,
}

#[derive(Display)]
#[display(fmt = "Hello there!")]
union U {
    i: u32,
}

#[derive(Octal)]
#[octal(fmt = "7")]
struct S;

#[derive(UpperHex)]
#[upper_hex(fmt = "UpperHex")]
struct UH;

#[derive(Display)]
struct Unit;

#[derive(Display)]
struct UnitStruct {}

#[derive(Display)]
enum EmptyEnum {}

#[derive(Display)]
#[display(fmt = "Generic")]
struct Generic<T>(T);

#[test]
fn check_display() {
    assert_eq!(MyInt(-2).to_string(), "-2");
    assert_eq!(format!("{:b}", MyInt(9)), "1001");
    assert_eq!(format!("{:o}", MyInt(9)), "11");
    assert_eq!(Point2D { x: 3, y: 4 }.to_string(), "(3, 4)");
    assert_eq!(E::Uint(2).to_string(), "2");
    assert_eq!(E::Binary { i: -2 }.to_string(), "I am B 11111110");
    #[cfg(feature = "nightly")]
    assert_eq!(E::Path("abc".into()).to_string(), "I am C abc");
    assert_eq!(EE::A.to_string(), "Java EE");
    assert_eq!(EE::B.to_string(), "Java EE");
    assert_eq!(U { i: 2 }.to_string(), "Hello there!");
    assert_eq!(format!("{:o}", S), "7");
    assert_eq!(format!("{:X}", UH), "UpperHex");
    assert_eq!(Unit.to_string(), "Unit");
    assert_eq!(UnitStruct {}.to_string(), "UnitStruct");
    assert_eq!(Generic(()).to_string(), "Generic");
}
<|endoftext|>


<commit_before>extern crate perfcnt;
extern crate x86;

use perfcnt::{PerfCounter, PerfCounterBuilderLinux, AbstractPerfCounter};

pub fn main() {
    let counter_description = x86::perfcnt::core_counters().unwrap().get("BR_INST_RETIRED.ALL_BRANCHES").unwrap();
    let pc: PerfCounter = PerfCounterBuilderLinux::from_raw_intel_counter_description(counter_description)
        .exclude_idle()
        .exclude_kernel()
        .finish()
        .expect("Could not create counter");

    pc.start().expect("Can not start the counter");
    println!("");
    pc.stop().expect("Can not stop the counter");;

    println!("{}: {:?}", counter_description.brief_description, pc.read().expect("Can not read counter"));
    pc.reset().expect("Can not reset the counter");
}
<commit_msg>Fix example.<commit_after>extern crate perfcnt;
extern crate x86;

use perfcnt::{PerfCounter, PerfCounterBuilderLinux, AbstractPerfCounter};

pub fn main() {
    let counter_description = x86::perfcnt::core_counters().unwrap().get("BR_INST_RETIRED.ALL_BRANCHES").unwrap();
    let pc: PerfCounter = PerfCounterBuilderLinux::from_intel_event_description(counter_description)
        .exclude_idle()
        .exclude_kernel()
        .finish()
        .expect("Could not create counter");

    pc.start().expect("Can not start the counter");
    println!("");
    pc.stop().expect("Can not stop the counter");;

    println!("{}: {:?}", counter_description.brief_description, pc.read().expect("Can not read counter"));
    pc.reset().expect("Can not reset the counter");
}
<|endoftext|>


<commit_before>extern crate clap;
extern crate regex;

include!("../clap-test.rs");

use clap::{App, ArgMatches, Arg, ErrorKind};

#[test]
fn indices_mult_opts() {
	let m = App::new("ind")
		.arg(Arg::with_name("exclude")
			.short("e")
			.takes_value(true)
			.multiple(true))
		.arg(Arg::with_name("include")
			.short("i")
			.takes_value(true)
			.multiple(true))
		.get_matches_from(vec!["ind", "-e", "A", "B", "-i", "B", "C", "-e", "C"]);

		assert_eq!(m.indices_of("exclude").unwrap().collect::<Vec<_>>(), &[2, 3, 8]);
		assert_eq!(m.indices_of("include").unwrap().collect::<Vec<_>>(), &[5, 6]);
}

#[test]
fn index_mult_opts() {
	let m = App::new("ind")
		.arg(Arg::with_name("exclude")
			.short("e")
			.takes_value(true)
			.multiple(true))
		.arg(Arg::with_name("include")
			.short("i")
			.takes_value(true)
			.multiple(true))
		.get_matches_from(vec!["ind", "-e", "A", "B", "-i", "B", "C", "-e", "C"]);

		assert_eq!(m.index_of("exclude"), Some(2));
		assert_eq!(m.index_of("include"), Some(5));
}

#[test]
fn index_flag() {
	let m = App::new("ind")
		.arg(Arg::with_name("exclude")
			.short("e"))
		.arg(Arg::with_name("include")
			.short("i"))
		.get_matches_from(vec!["ind", "-e", "-i"]);

		assert_eq!(m.index_of("exclude"), Some(1));
		assert_eq!(m.index_of("include"), Some(2));
}

#[test]
fn index_flags() {
	let m = App::new("ind")
		.arg(Arg::with_name("exclude")
			.short("e")
			.multiple(true))
		.arg(Arg::with_name("include")
			.short("i")
			.multiple(true))
		.get_matches_from(vec!["ind", "-e", "-i", "-e", "-e", "-i"]);

		assert_eq!(m.index_of("exclude"), Some(1));
		assert_eq!(m.index_of("include"), Some(2));
}

#[test]
fn indices_mult_flags() {
	let m = App::new("ind")
		.arg(Arg::with_name("exclude")
			.short("e")
			.multiple(true))
		.arg(Arg::with_name("include")
			.short("i")
			.multiple(true))
		.get_matches_from(vec!["ind", "-e", "-i", "-e", "-e", "-i"]);

		assert_eq!(m.indices_of("exclude").unwrap().collect::<Vec<_>>(), &[1, 3, 4]);
		assert_eq!(m.indices_of("include").unwrap().collect::<Vec<_>>(), &[2, 5]);
}

#[test]
fn indices_mult_flags_combined() {
	let m = App::new("ind")
		.arg(Arg::with_name("exclude")
			.short("e")
			.multiple(true))
		.arg(Arg::with_name("include")
			.short("i")
			.multiple(true))
		.get_matches_from(vec!["ind", "-eieei"]);

		assert_eq!(m.indices_of("exclude").unwrap().collect::<Vec<_>>(), &[1, 3, 4]);
		assert_eq!(m.indices_of("include").unwrap().collect::<Vec<_>>(), &[2, 5]);
}

#[test]
fn indices_mult_flags_opt_combined() {
	let m = App::new("ind")
		.arg(Arg::with_name("exclude")
			.short("e")
			.multiple(true))
		.arg(Arg::with_name("include")
			.short("i")
			.multiple(true))
		.arg(Arg::with_name("option")
			.short("0")
			.takes_value(true))
		.get_matches_from(vec!["ind", "-eieeio", "val"]);

		assert_eq!(m.indices_of("exclude").unwrap().collect::<Vec<_>>(), &[1, 3, 4]);
		assert_eq!(m.indices_of("include").unwrap().collect::<Vec<_>>(), &[2, 5]);
		assert_eq!(m.indices_of("option").unwrap().collect::<Vec<_>>(), &[7]);
}

#[test]
fn indices_mult_flags_opt_combined_eq() {
	let m = App::new("ind")
		.arg(Arg::with_name("exclude")
			.short("e")
			.multiple(true))
		.arg(Arg::with_name("include")
			.short("i")
			.multiple(true))
		.arg(Arg::with_name("option")
			.short("0")
			.takes_value(true))
		.get_matches_from(vec!["ind", "-eieeio=val"]);

		assert_eq!(m.indices_of("exclude").unwrap().collect::<Vec<_>>(), &[1, 3, 4]);
		assert_eq!(m.indices_of("include").unwrap().collect::<Vec<_>>(), &[2, 5]);
		assert_eq!(m.indices_of("option").unwrap().collect::<Vec<_>>(), &[7]);
}

#[test]
fn indices_mult_opt_value_delim_eq() {
    let m = App::new("myapp")
	    .arg(Arg::with_name("option")
		    .short("o")
		    .takes_value(true)
		    .multiple(true))
	    .get_matches_from(vec!["myapp", "-o=val1,val2,val3"]);
    assert_eq!(m.indices_of("option").unwrap().collect::<Vec<_>>(), &[2, 3, 4]);
}

#[test]
fn indices_mult_opt_mult_flag() {
    let m = App::new("myapp")
    .arg(Arg::with_name("option")
	    .short("o")
	    .takes_value(true)
	    .multiple(true))
    .arg(Arg::with_name("flag")
	    .short("f")
	    .multiple(true))
    .get_matches_from(vec!["myapp", "-o", "val1", "-f", "-o", "val2", "-f"]);

    assert_eq!(m.indices_of("option").unwrap().collect::<Vec<_>>(), &[2, 5]);
    assert_eq!(m.indices_of("flag").unwrap().collect::<Vec<_>>(), &[3, 6]);
}<commit_msg>fixup! tests: adds tests for querying indices<commit_after>extern crate clap;
extern crate regex;

include!("../clap-test.rs");

use clap::{App, Arg};

#[test]
fn indices_mult_opts() {
	let m = App::new("ind")
		.arg(Arg::with_name("exclude")
			.short("e")
			.takes_value(true)
			.multiple(true))
		.arg(Arg::with_name("include")
			.short("i")
			.takes_value(true)
			.multiple(true))
		.get_matches_from(vec!["ind", "-e", "A", "B", "-i", "B", "C", "-e", "C"]);

		assert_eq!(m.indices_of("exclude").unwrap().collect::<Vec<_>>(), &[2, 3, 8]);
		assert_eq!(m.indices_of("include").unwrap().collect::<Vec<_>>(), &[5, 6]);
}

#[test]
fn index_mult_opts() {
	let m = App::new("ind")
		.arg(Arg::with_name("exclude")
			.short("e")
			.takes_value(true)
			.multiple(true))
		.arg(Arg::with_name("include")
			.short("i")
			.takes_value(true)
			.multiple(true))
		.get_matches_from(vec!["ind", "-e", "A", "B", "-i", "B", "C", "-e", "C"]);

		assert_eq!(m.index_of("exclude"), Some(2));
		assert_eq!(m.index_of("include"), Some(5));
}

#[test]
fn index_flag() {
	let m = App::new("ind")
		.arg(Arg::with_name("exclude")
			.short("e"))
		.arg(Arg::with_name("include")
			.short("i"))
		.get_matches_from(vec!["ind", "-e", "-i"]);

		assert_eq!(m.index_of("exclude"), Some(1));
		assert_eq!(m.index_of("include"), Some(2));
}

#[test]
fn index_flags() {
	let m = App::new("ind")
		.arg(Arg::with_name("exclude")
			.short("e")
			.multiple(true))
		.arg(Arg::with_name("include")
			.short("i")
			.multiple(true))
		.get_matches_from(vec!["ind", "-e", "-i", "-e", "-e", "-i"]);

		assert_eq!(m.index_of("exclude"), Some(1));
		assert_eq!(m.index_of("include"), Some(2));
}

#[test]
fn indices_mult_flags() {
	let m = App::new("ind")
		.arg(Arg::with_name("exclude")
			.short("e")
			.multiple(true))
		.arg(Arg::with_name("include")
			.short("i")
			.multiple(true))
		.get_matches_from(vec!["ind", "-e", "-i", "-e", "-e", "-i"]);

		assert_eq!(m.indices_of("exclude").unwrap().collect::<Vec<_>>(), &[1, 3, 4]);
		assert_eq!(m.indices_of("include").unwrap().collect::<Vec<_>>(), &[2, 5]);
}

#[test]
fn indices_mult_flags_combined() {
	let m = App::new("ind")
		.arg(Arg::with_name("exclude")
			.short("e")
			.multiple(true))
		.arg(Arg::with_name("include")
			.short("i")
			.multiple(true))
		.get_matches_from(vec!["ind", "-eieei"]);

		assert_eq!(m.indices_of("exclude").unwrap().collect::<Vec<_>>(), &[1, 3, 4]);
		assert_eq!(m.indices_of("include").unwrap().collect::<Vec<_>>(), &[2, 5]);
}

#[test]
fn indices_mult_flags_opt_combined() {
	let m = App::new("ind")
		.arg(Arg::with_name("exclude")
			.short("e")
			.multiple(true))
		.arg(Arg::with_name("include")
			.short("i")
			.multiple(true))
		.arg(Arg::with_name("option")
			.short("o")
			.takes_value(true))
		.get_matches_from(vec!["ind", "-eieeio", "val"]);

		assert_eq!(m.indices_of("exclude").unwrap().collect::<Vec<_>>(), &[1, 3, 4]);
		assert_eq!(m.indices_of("include").unwrap().collect::<Vec<_>>(), &[2, 5]);
		assert_eq!(m.indices_of("option").unwrap().collect::<Vec<_>>(), &[7]);
}

#[test]
fn indices_mult_flags_opt_combined_eq() {
	let m = App::new("ind")
		.arg(Arg::with_name("exclude")
			.short("e")
			.multiple(true))
		.arg(Arg::with_name("include")
			.short("i")
			.multiple(true))
		.arg(Arg::with_name("option")
			.short("o")
			.takes_value(true))
		.get_matches_from(vec!["ind", "-eieeio=val"]);

		assert_eq!(m.indices_of("exclude").unwrap().collect::<Vec<_>>(), &[1, 3, 4]);
		assert_eq!(m.indices_of("include").unwrap().collect::<Vec<_>>(), &[2, 5]);
		assert_eq!(m.indices_of("option").unwrap().collect::<Vec<_>>(), &[7]);
}

#[test]
fn indices_mult_opt_value_delim_eq() {
    let m = App::new("myapp")
	    .arg(Arg::with_name("option")
		    .short("o")
		    .takes_value(true)
		    .use_delimiter(true)
		    .multiple(true))
	    .get_matches_from(vec!["myapp", "-o=val1,val2,val3"]);
    assert_eq!(m.indices_of("option").unwrap().collect::<Vec<_>>(), &[2, 3, 4]);
}

#[test]
fn indices_mult_opt_value_no_delim_eq() {
    let m = App::new("myapp")
	    .arg(Arg::with_name("option")
		    .short("o")
		    .takes_value(true)
		    .multiple(true))
	    .get_matches_from(vec!["myapp", "-o=val1,val2,val3"]);
    assert_eq!(m.indices_of("option").unwrap().collect::<Vec<_>>(), &[2]);
}

#[test]
fn indices_mult_opt_mult_flag() {
    let m = App::new("myapp")
    .arg(Arg::with_name("option")
	    .short("o")
	    .takes_value(true)
	    .multiple(true))
    .arg(Arg::with_name("flag")
	    .short("f")
	    .multiple(true))
    .get_matches_from(vec!["myapp", "-o", "val1", "-f", "-o", "val2", "-f"]);

    assert_eq!(m.indices_of("option").unwrap().collect::<Vec<_>>(), &[2, 5]);
    assert_eq!(m.indices_of("flag").unwrap().collect::<Vec<_>>(), &[3, 6]);
}<|endoftext|>


<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{json_rpc::JsonRpcClientWrapper, TransactionContext};
use libra_crypto::{ed25519::Ed25519PublicKey, x25519};
use libra_global_constants::{
    CONSENSUS_KEY, FULLNODE_NETWORK_KEY, OPERATOR_ACCOUNT, OWNER_ACCOUNT, VALIDATOR_NETWORK_KEY,
};
use libra_management::{error::Error, storage::to_x25519};
use libra_network_address::{
    encrypted::{EncNetworkAddress, Key, RawEncNetworkAddress, TEST_SHARED_VAL_NETADDR_KEY},
    NetworkAddress, Protocol, RawNetworkAddress,
};
use libra_types::account_address::AccountAddress;
use serde::Serialize;
use std::convert::TryFrom;
use structopt::StructOpt;

// TODO: Load all chain IDs from the host
#[derive(Debug, StructOpt)]
pub struct SetValidatorConfig {
    /// JSON-RPC Endpoint (e.g. http://localhost:8080)
    #[structopt(long, required_unless = "config")]
    json_server: Option<String>,
    #[structopt(flatten)]
    validator_config: libra_management::validator_config::ValidatorConfig,
    #[structopt(long, help = "Validator Network Address")]
    validator_address: Option<NetworkAddress>,
    #[structopt(long, help = "Full Node Network Address")]
    fullnode_address: Option<NetworkAddress>,
}

impl SetValidatorConfig {
    pub fn execute(self) -> Result<TransactionContext, Error> {
        let config = self
            .validator_config
            .config
            .load()?
            .override_chain_id(self.validator_config.chain_id)
            .override_json_server(&self.json_server)
            .override_validator_backend(
                &self.validator_config.validator_backend.validator_backend,
            )?;
        let storage = config.validator_backend();

        let operator_account = storage.account_address(OPERATOR_ACCOUNT)?;
        let owner_account = storage.account_address(OWNER_ACCOUNT)?;
        let client = JsonRpcClientWrapper::new(config.json_server);
        let sequence_number = client.sequence_number(operator_account)?;

        // Retrieve the current validator / fullnode addresses and update accordingly
        let validator_config = client.validator_config(owner_account)?;
        let validator_address = if let Some(validator_address) = self.validator_address {
            validator_address
        } else {
            decode_validator_address(
                &validator_config.validator_network_address,
                &owner_account,
                &TEST_SHARED_VAL_NETADDR_KEY,
                0, // addr_idx
            )?
        };

        let fullnode_address = if let Some(fullnode_address) = self.fullnode_address {
            fullnode_address
        } else {
            decode_address(&validator_config.full_node_network_address)?
        };

        let txn = self.validator_config.build_transaction(
            sequence_number,
            fullnode_address,
            validator_address,
            true,
        )?;
        client.submit_transaction(txn.as_signed_user_txn().unwrap().clone())
    }
}

#[derive(Debug, StructOpt)]
pub struct RotateKey {
    /// JSON-RPC Endpoint (e.g. http://localhost:8080)
    #[structopt(long, required_unless = "config")]
    json_server: Option<String>,
    #[structopt(flatten)]
    validator_config: libra_management::validator_config::ValidatorConfig,
}

impl RotateKey {
    pub fn execute(
        self,
        key_name: &'static str,
    ) -> Result<(TransactionContext, Ed25519PublicKey), Error> {
        // Load the config, storage backend and create a json rpc client.
        let config = self
            .validator_config
            .config()?
            .override_json_server(&self.json_server);
        let mut storage = config.validator_backend();
        let client = JsonRpcClientWrapper::new(config.json_server);

        // Fetch the current on-chain validator config for the node
        let owner_account = storage.account_address(OWNER_ACCOUNT)?;
        let validator_config = client.validator_config(owner_account)?;

        // Check that the key held in storage matches the key registered on-chain in the validator
        // config. If so, rotate the key in storage. If not, don't rotate the key in storage and
        // rather allow the next step to resubmit the set_validator_config transaction with the
        // current key (to resynchronize the validator config on the blockchain).
        let mut storage_key = storage.ed25519_public_from_private(key_name)?;
        let keys_match = match key_name {
            CONSENSUS_KEY => storage_key == validator_config.consensus_public_key,
            VALIDATOR_NETWORK_KEY => {
                to_x25519(storage_key.clone())?
                    == validator_config.validator_network_identity_public_key
            }
            FULLNODE_NETWORK_KEY => {
                to_x25519(storage_key.clone())?
                    == validator_config.full_node_network_identity_public_key
            }
            _ => {
                return Err(Error::UnexpectedError(
                    "Rotate key was called with an unknown key name!".into(),
                ));
            }
        };
        if keys_match {
            storage_key = storage.rotate_key(key_name)?;
        }

        // Create and set the validator config state on the blockchain.
        let set_validator_config = SetValidatorConfig {
            json_server: self.json_server,
            validator_config: self.validator_config,
            validator_address: None,
            fullnode_address: None,
        };
        set_validator_config
            .execute()
            .map(|txn_ctx| (txn_ctx, storage_key))
    }
}

#[derive(Debug, StructOpt)]
pub struct RotateConsensusKey {
    #[structopt(flatten)]
    rotate_key: RotateKey,
}

impl RotateConsensusKey {
    pub fn execute(self) -> Result<(TransactionContext, Ed25519PublicKey), Error> {
        self.rotate_key.execute(CONSENSUS_KEY)
    }
}

#[derive(Debug, StructOpt)]
pub struct RotateValidatorNetworkKey {
    #[structopt(flatten)]
    rotate_key: RotateKey,
}

impl RotateValidatorNetworkKey {
    pub fn execute(self) -> Result<(TransactionContext, x25519::PublicKey), Error> {
        let (txn_ctx, key) = self.rotate_key.execute(VALIDATOR_NETWORK_KEY)?;
        Ok((txn_ctx, to_x25519(key)?))
    }
}

#[derive(Debug, StructOpt)]
pub struct RotateFullNodeNetworkKey {
    #[structopt(flatten)]
    rotate_key: RotateKey,
}

impl RotateFullNodeNetworkKey {
    pub fn execute(self) -> Result<(TransactionContext, x25519::PublicKey), Error> {
        let (txn_ctx, key) = self.rotate_key.execute(FULLNODE_NETWORK_KEY)?;
        Ok((txn_ctx, to_x25519(key)?))
    }
}

pub fn decode_validator_address(
    address: &RawEncNetworkAddress,
    account: &AccountAddress,
    key: &Key,
    addr_idx: u32,
) -> Result<NetworkAddress, Error> {
    let enc_addr = EncNetworkAddress::try_from(address).map_err(|e| {
        Error::UnexpectedError(format!(
            "Failed to decode network address {}",
            e.to_string()
        ))
    })?;
    let raw_addr = enc_addr.decrypt(key, account, addr_idx).map_err(|e| {
        Error::UnexpectedError(format!(
            "Failed to decrypt network address {}",
            e.to_string()
        ))
    })?;
    decode_address(&raw_addr)
}

pub fn decode_address(raw_address: &RawNetworkAddress) -> Result<NetworkAddress, Error> {
    let network_address = NetworkAddress::try_from(raw_address).map_err(|e| {
        Error::UnexpectedError(format!(
            "Failed to decode network address {}",
            e.to_string()
        ))
    })?;
    let protocols = network_address
        .as_slice()
        .iter()
        .filter(|protocol| match protocol {
            Protocol::Dns(_)
            | Protocol::Dns4(_)
            | Protocol::Dns6(_)
            | Protocol::Ip4(_)
            | Protocol::Ip6(_)
            | Protocol::Memory(_)
            | Protocol::Tcp(_) => true,
            _ => false,
        })
        .cloned()
        .collect::<Vec<_>>();
    Ok(NetworkAddress::try_from(protocols).unwrap())
}

#[derive(Debug, StructOpt)]
pub struct ValidatorConfig {
    #[structopt(long, help = "Validator account address to display the config")]
    account_address: AccountAddress,
    #[structopt(flatten)]
    config: libra_management::config::ConfigPath,
    /// JSON-RPC Endpoint (e.g. http://localhost:8080)
    #[structopt(long, required_unless = "config")]
    json_server: Option<String>,
}

impl ValidatorConfig {
    pub fn execute(self) -> Result<DecryptedValidatorConfig, Error> {
        let config = self.config.load()?.override_json_server(&self.json_server);
        let client = JsonRpcClientWrapper::new(config.json_server);
        client.validator_config(self.account_address).map(|config| {
            DecryptedValidatorConfig::from_validator_config(
                self.account_address,
                &TEST_SHARED_VAL_NETADDR_KEY,
                0,
                &config,
            )
            .unwrap()
        })
    }
}

#[derive(Serialize)]
pub struct DecryptedValidatorConfig {
    pub consensus_public_key: Ed25519PublicKey,
    pub validator_network_key: x25519::PublicKey,
    pub validator_network_address: NetworkAddress,
    pub full_node_network_key: x25519::PublicKey,
    pub full_node_network_address: NetworkAddress,
}

impl DecryptedValidatorConfig {
    pub fn from_validator_config(
        account: AccountAddress,
        address_encryption_key: &libra_network_address::encrypted::Key,
        addr_idx: u32,
        validator_config: &libra_types::validator_config::ValidatorConfig,
    ) -> Result<DecryptedValidatorConfig, Error> {
        let full_node_network_address =
            crate::validator_config::decode_address(&validator_config.full_node_network_address)?;

        let validator_network_address = crate::validator_config::decode_validator_address(
            &validator_config.validator_network_address,
            &account,
            address_encryption_key,
            addr_idx,
        )?;

        Ok(DecryptedValidatorConfig {
            consensus_public_key: validator_config.consensus_public_key.clone(),
            validator_network_key: validator_config.validator_network_identity_public_key,
            validator_network_address,
            full_node_network_key: validator_config.full_node_network_identity_public_key,
            full_node_network_address,
        })
    }
}
<commit_msg>[management] Prevent from sending no change validator configs<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{json_rpc::JsonRpcClientWrapper, TransactionContext};
use libra_crypto::{ed25519::Ed25519PublicKey, x25519};
use libra_global_constants::{
    CONSENSUS_KEY, FULLNODE_NETWORK_KEY, OPERATOR_ACCOUNT, OWNER_ACCOUNT, VALIDATOR_NETWORK_KEY,
};
use libra_management::{error::Error, storage::to_x25519};
use libra_network_address::{
    encrypted::{EncNetworkAddress, Key, RawEncNetworkAddress, TEST_SHARED_VAL_NETADDR_KEY},
    NetworkAddress, Protocol, RawNetworkAddress,
};
use libra_types::account_address::AccountAddress;
use serde::Serialize;
use std::convert::TryFrom;
use structopt::StructOpt;

// TODO: Load all chain IDs from the host
#[derive(Debug, StructOpt)]
pub struct SetValidatorConfig {
    /// JSON-RPC Endpoint (e.g. http://localhost:8080)
    #[structopt(long, required_unless = "config")]
    json_server: Option<String>,
    #[structopt(flatten)]
    validator_config: libra_management::validator_config::ValidatorConfig,
    #[structopt(
        long,
        required_unless = "fullnode-address",
        help = "Validator Network Address"
    )]
    validator_address: Option<NetworkAddress>,
    #[structopt(
        long,
        required_unless = "validator-address",
        help = "Full Node Network Address"
    )]
    fullnode_address: Option<NetworkAddress>,
}

impl SetValidatorConfig {
    pub fn execute(self) -> Result<TransactionContext, Error> {
        let config = self
            .validator_config
            .config
            .load()?
            .override_chain_id(self.validator_config.chain_id)
            .override_json_server(&self.json_server)
            .override_validator_backend(
                &self.validator_config.validator_backend.validator_backend,
            )?;
        let storage = config.validator_backend();

        let operator_account = storage.account_address(OPERATOR_ACCOUNT)?;
        let owner_account = storage.account_address(OWNER_ACCOUNT)?;
        let client = JsonRpcClientWrapper::new(config.json_server);
        let sequence_number = client.sequence_number(operator_account)?;

        // Retrieve the current validator / fullnode addresses and update accordingly
        let validator_config = client.validator_config(owner_account)?;
        let validator_address = if let Some(validator_address) = self.validator_address {
            validator_address
        } else {
            decode_validator_address(
                &validator_config.validator_network_address,
                &owner_account,
                &TEST_SHARED_VAL_NETADDR_KEY,
                0, // addr_idx
            )?
        };

        let fullnode_address = if let Some(fullnode_address) = self.fullnode_address {
            fullnode_address
        } else {
            decode_address(&validator_config.full_node_network_address)?
        };

        let txn = self.validator_config.build_transaction(
            sequence_number,
            fullnode_address,
            validator_address,
            true,
        )?;
        client.submit_transaction(txn.as_signed_user_txn().unwrap().clone())
    }
}

#[derive(Debug, StructOpt)]
pub struct RotateKey {
    /// JSON-RPC Endpoint (e.g. http://localhost:8080)
    #[structopt(long, required_unless = "config")]
    json_server: Option<String>,
    #[structopt(flatten)]
    validator_config: libra_management::validator_config::ValidatorConfig,
}

impl RotateKey {
    pub fn execute(
        self,
        key_name: &'static str,
    ) -> Result<(TransactionContext, Ed25519PublicKey), Error> {
        // Load the config, storage backend and create a json rpc client.
        let config = self
            .validator_config
            .config()?
            .override_json_server(&self.json_server);
        let mut storage = config.validator_backend();
        let client = JsonRpcClientWrapper::new(config.json_server);

        // Fetch the current on-chain validator config for the node
        let owner_account = storage.account_address(OWNER_ACCOUNT)?;
        let validator_config = client.validator_config(owner_account)?;

        // Check that the key held in storage matches the key registered on-chain in the validator
        // config. If so, rotate the key in storage. If not, don't rotate the key in storage and
        // rather allow the next step to resubmit the set_validator_config transaction with the
        // current key (to resynchronize the validator config on the blockchain).
        let mut storage_key = storage.ed25519_public_from_private(key_name)?;
        let keys_match = match key_name {
            CONSENSUS_KEY => storage_key == validator_config.consensus_public_key,
            VALIDATOR_NETWORK_KEY => {
                to_x25519(storage_key.clone())?
                    == validator_config.validator_network_identity_public_key
            }
            FULLNODE_NETWORK_KEY => {
                to_x25519(storage_key.clone())?
                    == validator_config.full_node_network_identity_public_key
            }
            _ => {
                return Err(Error::UnexpectedError(
                    "Rotate key was called with an unknown key name!".into(),
                ));
            }
        };
        if keys_match {
            storage_key = storage.rotate_key(key_name)?;
        }

        // Create and set the validator config state on the blockchain.
        let set_validator_config = SetValidatorConfig {
            json_server: self.json_server,
            validator_config: self.validator_config,
            validator_address: None,
            fullnode_address: None,
        };
        set_validator_config
            .execute()
            .map(|txn_ctx| (txn_ctx, storage_key))
    }
}

#[derive(Debug, StructOpt)]
pub struct RotateConsensusKey {
    #[structopt(flatten)]
    rotate_key: RotateKey,
}

impl RotateConsensusKey {
    pub fn execute(self) -> Result<(TransactionContext, Ed25519PublicKey), Error> {
        self.rotate_key.execute(CONSENSUS_KEY)
    }
}

#[derive(Debug, StructOpt)]
pub struct RotateValidatorNetworkKey {
    #[structopt(flatten)]
    rotate_key: RotateKey,
}

impl RotateValidatorNetworkKey {
    pub fn execute(self) -> Result<(TransactionContext, x25519::PublicKey), Error> {
        let (txn_ctx, key) = self.rotate_key.execute(VALIDATOR_NETWORK_KEY)?;
        Ok((txn_ctx, to_x25519(key)?))
    }
}

#[derive(Debug, StructOpt)]
pub struct RotateFullNodeNetworkKey {
    #[structopt(flatten)]
    rotate_key: RotateKey,
}

impl RotateFullNodeNetworkKey {
    pub fn execute(self) -> Result<(TransactionContext, x25519::PublicKey), Error> {
        let (txn_ctx, key) = self.rotate_key.execute(FULLNODE_NETWORK_KEY)?;
        Ok((txn_ctx, to_x25519(key)?))
    }
}

pub fn decode_validator_address(
    address: &RawEncNetworkAddress,
    account: &AccountAddress,
    key: &Key,
    addr_idx: u32,
) -> Result<NetworkAddress, Error> {
    let enc_addr = EncNetworkAddress::try_from(address).map_err(|e| {
        Error::UnexpectedError(format!(
            "Failed to decode network address {}",
            e.to_string()
        ))
    })?;
    let raw_addr = enc_addr.decrypt(key, account, addr_idx).map_err(|e| {
        Error::UnexpectedError(format!(
            "Failed to decrypt network address {}",
            e.to_string()
        ))
    })?;
    decode_address(&raw_addr)
}

pub fn decode_address(raw_address: &RawNetworkAddress) -> Result<NetworkAddress, Error> {
    let network_address = NetworkAddress::try_from(raw_address).map_err(|e| {
        Error::UnexpectedError(format!(
            "Failed to decode network address {}",
            e.to_string()
        ))
    })?;
    let protocols = network_address
        .as_slice()
        .iter()
        .filter(|protocol| match protocol {
            Protocol::Dns(_)
            | Protocol::Dns4(_)
            | Protocol::Dns6(_)
            | Protocol::Ip4(_)
            | Protocol::Ip6(_)
            | Protocol::Memory(_)
            | Protocol::Tcp(_) => true,
            _ => false,
        })
        .cloned()
        .collect::<Vec<_>>();
    Ok(NetworkAddress::try_from(protocols).unwrap())
}

#[derive(Debug, StructOpt)]
pub struct ValidatorConfig {
    #[structopt(long, help = "Validator account address to display the config")]
    account_address: AccountAddress,
    #[structopt(flatten)]
    config: libra_management::config::ConfigPath,
    /// JSON-RPC Endpoint (e.g. http://localhost:8080)
    #[structopt(long, required_unless = "config")]
    json_server: Option<String>,
}

impl ValidatorConfig {
    pub fn execute(self) -> Result<DecryptedValidatorConfig, Error> {
        let config = self.config.load()?.override_json_server(&self.json_server);
        let client = JsonRpcClientWrapper::new(config.json_server);
        client.validator_config(self.account_address).map(|config| {
            DecryptedValidatorConfig::from_validator_config(
                self.account_address,
                &TEST_SHARED_VAL_NETADDR_KEY,
                0,
                &config,
            )
            .unwrap()
        })
    }
}

#[derive(Serialize)]
pub struct DecryptedValidatorConfig {
    pub consensus_public_key: Ed25519PublicKey,
    pub validator_network_key: x25519::PublicKey,
    pub validator_network_address: NetworkAddress,
    pub full_node_network_key: x25519::PublicKey,
    pub full_node_network_address: NetworkAddress,
}

impl DecryptedValidatorConfig {
    pub fn from_validator_config(
        account: AccountAddress,
        address_encryption_key: &libra_network_address::encrypted::Key,
        addr_idx: u32,
        validator_config: &libra_types::validator_config::ValidatorConfig,
    ) -> Result<DecryptedValidatorConfig, Error> {
        let full_node_network_address =
            crate::validator_config::decode_address(&validator_config.full_node_network_address)?;

        let validator_network_address = crate::validator_config::decode_validator_address(
            &validator_config.validator_network_address,
            &account,
            address_encryption_key,
            addr_idx,
        )?;

        Ok(DecryptedValidatorConfig {
            consensus_public_key: validator_config.consensus_public_key.clone(),
            validator_network_key: validator_config.validator_network_identity_public_key,
            validator_network_address,
            full_node_network_key: validator_config.full_node_network_identity_public_key,
            full_node_network_address,
        })
    }
}
<|endoftext|>


<commit_before>//! [PUT /_matrix/federation/v2/send_leave/{roomId}/{eventId}](https://matrix.org/docs/spec/server_server/r0.1.4#put-matrix-federation-v2-send-leave-roomid-eventid)

use std::time::SystemTime;

use js_int::UInt;
use ruma_api::ruma_api;
use ruma_events::{room::member::MemberEventContent, EventType};
use ruma_identifiers::{EventId, RoomId, ServerName, UserId};
use ruma_serde::Raw;

ruma_api! {
    metadata: {
        description: "Submits a signed leave event to the receiving server for it to accept it into the room's graph.",
        name: "create_leave_event",
        method: PUT,
        path: "/_matrix/federation/v1/send_leave/:room_id/:event_id",
        rate_limited: false,
        authentication: ServerSignatures,
    }

    request: {
        /// The room ID that is about to be left.
        #[ruma_api(path)]
        pub room_id: &'a RoomId,

        /// The event ID for the leave event.
        #[ruma_api(path)]
        pub event_id: &'a EventId,

        /// The user ID of the leaving member.
        #[ruma_api(query)]
        pub sender: &'a UserId,

        /// The name of the leaving homeserver.
        #[ruma_api(query)]
        pub origin: &'a ServerName,

        /// A timestamp added by the leaving homeserver.
        #[ruma_api(query)]
        #[serde(with = "ruma_serde::time::ms_since_unix_epoch")]
        pub origin_server_ts: SystemTime,

        /// The value `m.room.member`.
        #[ruma_api(query)]
        #[serde(rename = "type")]
        pub event_type: EventType,

        /// The user ID of the leaving member.
        #[ruma_api(query)]
        pub state_key: &'a str,

        /// The content of the event.
        #[ruma_api(query)]
        pub content: Raw<MemberEventContent>,

        /// This field must be present but is ignored; it may be 0.
        #[ruma_api(query)]
        pub depth: UInt,
    }

    #[derive(Default)]
    response: {}
}

impl<'a> Request<'a> {
    /// Creates a new `Request` with:
    /// * the room ID that is about to be left.
    /// * the event ID for the leave event.
    /// * the user ID of the leaving member.
    /// * the name of the leaving homeserver.
    /// * a timestamp added by the leaving homeserver.
    /// * the value `m.room.member`.
    /// * the user ID of the leaving member.
    /// * the content of the event.
    /// * this field must be present but is ignored; it may be 0.
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        room_id: &'a RoomId,
        event_id: &'a EventId,
        sender: &'a UserId,
        origin: &'a ServerName,
        origin_server_ts: SystemTime,
        event_type: EventType,
        state_key: &'a str,
        content: Raw<MemberEventContent>,
        depth: UInt,
    ) -> Self {
        Self {
            room_id,
            event_id,
            sender,
            origin,
            origin_server_ts,
            event_type,
            state_key,
            content,
            depth,
        }
    }
}

impl Response {
    /// Creates an empty `Response`.
    pub fn new() -> Self {
        Self
    }
}

#[cfg(all(test, feature = "server"))]
mod tests {
    use std::convert::TryInto;

    use super::Response;

    #[test]
    fn response_body() {
        let res: http::Response<Vec<u8>> = Response::new().try_into().unwrap();

        assert_eq!(res.body(), b"{}");
    }
}
<commit_msg>Fix test using try_into -> try_into_http_response<commit_after>//! [PUT /_matrix/federation/v2/send_leave/{roomId}/{eventId}](https://matrix.org/docs/spec/server_server/r0.1.4#put-matrix-federation-v2-send-leave-roomid-eventid)

use std::time::SystemTime;

use js_int::UInt;
use ruma_api::ruma_api;
use ruma_events::{room::member::MemberEventContent, EventType};
use ruma_identifiers::{EventId, RoomId, ServerName, UserId};
use ruma_serde::Raw;

ruma_api! {
    metadata: {
        description: "Submits a signed leave event to the receiving server for it to accept it into the room's graph.",
        name: "create_leave_event",
        method: PUT,
        path: "/_matrix/federation/v1/send_leave/:room_id/:event_id",
        rate_limited: false,
        authentication: ServerSignatures,
    }

    request: {
        /// The room ID that is about to be left.
        #[ruma_api(path)]
        pub room_id: &'a RoomId,

        /// The event ID for the leave event.
        #[ruma_api(path)]
        pub event_id: &'a EventId,

        /// The user ID of the leaving member.
        #[ruma_api(query)]
        pub sender: &'a UserId,

        /// The name of the leaving homeserver.
        #[ruma_api(query)]
        pub origin: &'a ServerName,

        /// A timestamp added by the leaving homeserver.
        #[ruma_api(query)]
        #[serde(with = "ruma_serde::time::ms_since_unix_epoch")]
        pub origin_server_ts: SystemTime,

        /// The value `m.room.member`.
        #[ruma_api(query)]
        #[serde(rename = "type")]
        pub event_type: EventType,

        /// The user ID of the leaving member.
        #[ruma_api(query)]
        pub state_key: &'a str,

        /// The content of the event.
        #[ruma_api(query)]
        pub content: Raw<MemberEventContent>,

        /// This field must be present but is ignored; it may be 0.
        #[ruma_api(query)]
        pub depth: UInt,
    }

    #[derive(Default)]
    response: {}
}

impl<'a> Request<'a> {
    /// Creates a new `Request` with:
    /// * the room ID that is about to be left.
    /// * the event ID for the leave event.
    /// * the user ID of the leaving member.
    /// * the name of the leaving homeserver.
    /// * a timestamp added by the leaving homeserver.
    /// * the value `m.room.member`.
    /// * the user ID of the leaving member.
    /// * the content of the event.
    /// * this field must be present but is ignored; it may be 0.
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        room_id: &'a RoomId,
        event_id: &'a EventId,
        sender: &'a UserId,
        origin: &'a ServerName,
        origin_server_ts: SystemTime,
        event_type: EventType,
        state_key: &'a str,
        content: Raw<MemberEventContent>,
        depth: UInt,
    ) -> Self {
        Self {
            room_id,
            event_id,
            sender,
            origin,
            origin_server_ts,
            event_type,
            state_key,
            content,
            depth,
        }
    }
}

impl Response {
    /// Creates an empty `Response`.
    pub fn new() -> Self {
        Self
    }
}

#[cfg(all(test, feature = "server"))]
mod tests {
    use ruma_api::OutgoingResponse;

    use super::Response;

    #[test]
    fn response_body() {
        let res = Response::new().try_into_http_response().unwrap();

        assert_eq!(res.body(), b"{}");
    }
}
<|endoftext|>