Merge branch 'master' into scx1332/fix_pretty_table

agent/provider/src/cli/preset.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::BTreeMap;
 
 use anyhow::{anyhow, bail, Result};
 use dialoguer::{Input, Select};
@@ -131,10 +131,10 @@ impl PresetUpdater {
 
     pub fn update_metrics(&mut self, config: &ProviderConfig) -> Result<()> {
         let registry = config.registry()?;
-        let mut usage_coeffs: HashMap<String, f64> = Default::default();
+        let mut usage_coeffs: BTreeMap<String, f64> = Default::default();
         let exe_unit_desc = registry.find_exeunit(&self.preset.exeunit_name)?;
 
-        fn get_usage(m: &HashMap<String, f64>, k1: &str, k2: &str) -> f64 {
+        fn get_usage(m: &BTreeMap<String, f64>, k1: &str, k2: &str) -> f64 {
             m.get(k1)
                 .cloned()
                 .unwrap_or_else(|| m.get(k2).cloned().unwrap_or(0.))

agent/provider/src/cli/rule.rs

@@ -1,7 +1,12 @@
 use std::collections::HashMap;
+use std::path::PathBuf;
 
-use anyhow::Result;
+use anyhow::{anyhow, Result};
 use structopt::StructOpt;
+use strum::VariantNames;
+use ya_manifest_utils::policy::CertPermissions;
+use ya_manifest_utils::util::cert_to_id;
+use ya_manifest_utils::{KeystoreLoadResult, KeystoreManager};
 use ya_utils_cli::{CommandOutput, ResponseTable};
 
 use crate::rules::CertRule;
@@ -28,19 +33,28 @@ pub enum SetOutboundRule {
     Disable,
     Enable,
     Everyone {
-        #[structopt(subcommand)]
+        #[structopt(short, long, possible_values = Mode::VARIANTS)]
         mode: Mode,
     },
     AuditedPayload {
         #[structopt(long)]
         cert_id: Option<String>,
-        #[structopt(subcommand)]
+        #[structopt(short, long, possible_values = Mode::VARIANTS)]
         mode: Mode,
     },
-    Partner {
-        #[structopt(long)]
+    Partner(RuleWithCert),
+}
+
+#[derive(StructOpt, Clone, Debug)]
+pub enum RuleWithCert {
+    CertId {
         cert_id: String,
-        #[structopt(subcommand)]
+        #[structopt(short, long, possible_values = Mode::VARIANTS)]
+        mode: Mode,
+    },
+    ImportCert {
+        import_cert: PathBuf,
+        #[structopt(short, long, possible_values = Mode::VARIANTS)]
         mode: Mode,
     },
 }
@@ -70,7 +84,36 @@ fn set(set_rule: SetRule, config: ProviderConfig) -> Result<()> {
                 Some(_) => todo!("Setting rule for specific certificate isn't implemented yet"),
                 None => rules.set_default_audited_payload_mode(mode),
             },
-            SetOutboundRule::Partner { cert_id, mode } => rules.set_partner_mode(cert_id, mode),
+            SetOutboundRule::Partner(RuleWithCert::CertId { cert_id, mode }) => {
+                rules.set_partner_mode(cert_id, mode)
+            }
+            SetOutboundRule::Partner(RuleWithCert::ImportCert { import_cert, mode }) => {
+                let keystore_manager = KeystoreManager::try_new(&rules.cert_dir)?;
+
+                let KeystoreLoadResult { loaded, skipped } =
+                    keystore_manager.load_certs(&vec![import_cert])?;
+
+                //TODO it will be removed after backward compatibility is done
+                rules.keystore.permissions_manager().set_many(
+                    &loaded.iter().chain(skipped.iter()).cloned().collect(),
+                    vec![CertPermissions::All],
+                    true,
+                );
+                rules
+                    .keystore
+                    .permissions_manager()
+                    .save(&rules.cert_dir)
+                    .map_err(|e| anyhow!("Failed to save permissions file: {e}"))?;
+
+                rules.keystore.reload(&rules.cert_dir)?;
+
+                for cert in loaded.into_iter().chain(skipped) {
+                    let cert_id = cert_to_id(&cert)?;
+                    rules.set_partner_mode(cert_id, mode.clone())?;
+                }
+
+                Ok(())
+            }
         },
     }
 }

agent/provider/src/config/presets.rs

@@ -1,4 +1,4 @@
-use std::collections::{HashMap, HashSet};
+use std::collections::{BTreeMap, HashSet};
 use std::path::Path;
 
 use anyhow::anyhow;
@@ -14,14 +14,15 @@ pub struct PresetV0 {
     pub name: String,
     pub exeunit_name: String,
     pub pricing_model: String,
-    pub usage_coeffs: HashMap<String, f64>,
+    pub usage_coeffs: BTreeMap<String, f64>,
 }
 
 #[derive(Serialize, Deserialize, Clone, Debug, Default)]
 #[serde(rename_all = "kebab-case")]
 pub struct Presets {
     pub active: Vec<String>,
-    pub presets: HashMap<String, Preset>,
+    // It's important that all values are sorted, so that other tools can easily detect changes.
+    pub presets: BTreeMap<String, Preset>,
 }
 
 #[derive(Serialize, Deserialize)]

agent/provider/src/market/presets.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::BTreeMap;
 use std::fmt;
 use std::fmt::Formatter;
 use std::path::Path;
@@ -21,7 +21,8 @@ pub struct Preset {
     pub exeunit_name: String,
     pub pricing_model: String,
     pub initial_price: f64,
-    pub usage_coeffs: HashMap<String, f64>,
+    // It's important that all values are sorted, so that other tools can easily detect changes.
+    pub usage_coeffs: BTreeMap<String, f64>,
 }
 
 impl Preset {

agent/provider/src/provider_agent.rs

@@ -643,7 +643,7 @@ mod tests {
 
         let preset = Preset {
             pricing_model: "linear".to_string(),
-            usage_coeffs: std::collections::HashMap::from([("test_coefficient".to_string(), 1.0)]),
+            usage_coeffs: std::collections::BTreeMap::from([("test_coefficient".to_string(), 1.0)]),
             ..Default::default()
         };
 

agent/provider/src/rules.rs

@@ -11,7 +11,7 @@ use std::{
 use anyhow::{anyhow, Result};
 use serde::{Deserialize, Serialize};
 use structopt::StructOpt;
-use strum::Display;
+use strum::{Display, EnumString, EnumVariantNames};
 use url::Url;
 use ya_manifest_utils::{
     matching::{
@@ -27,10 +27,10 @@ use crate::startup_config::FileMonitor;
 #[derive(Clone, Debug)]
 pub struct RulesManager {
     pub rulestore: Rulestore,
+    pub keystore: Keystore,
+    pub cert_dir: PathBuf,
     whitelist: DomainWhitelistState,
-    keystore: Keystore,
     whitelist_file: PathBuf,
-    cert_dir: PathBuf,
 }
 
 impl RulesManager {
@@ -431,8 +431,20 @@ pub struct CertRule {
     pub description: String,
 }
 
-#[derive(StructOpt, Clone, Debug, Serialize, Deserialize, Eq, PartialEq, Display)]
+#[derive(
+    StructOpt,
+    Clone,
+    Debug,
+    Serialize,
+    Deserialize,
+    Eq,
+    PartialEq,
+    Display,
+    EnumString,
+    EnumVariantNames,
+)]
 #[serde(rename_all = "kebab-case")]
+#[strum(serialize_all = "kebab-case")]
 pub enum Mode {
     All,
     None,

agent/provider/tests/rule_cli.rs

@@ -88,6 +88,7 @@ fn rule_set_should_edit_everyone_mode(mode: &str) {
         .arg("set")
         .arg("outbound")
         .arg(rule)
+        .arg("--mode")
         .arg(mode)
         .assert()
         .success();
@@ -110,6 +111,7 @@ fn rule_set_should_edit_default_modes_for_certificate_rules(rule: &str, mode: &s
         .arg("set")
         .arg("outbound")
         .arg(rule)
+        .arg("--mode")
         .arg(mode)
         .assert()
         .success();
@@ -136,12 +138,13 @@ fn adding_rule_for_non_existing_certificate_should_fail(rule: &str) {
         .arg("set")
         .arg("outbound")
         .arg(rule)
-        .arg("--cert-id")
+        .arg("cert-id")
         .arg(cert_id)
+        .arg("--mode")
         .arg("all")
         .assert()
         .stderr(format!(
-            "Error: Setting Partner mode All failed: No cert id: {cert_id} found in keystore\n"
+            "Error: Setting Partner mode all failed: No cert id: {cert_id} found in keystore\n"
         ));
 }
 
@@ -161,8 +164,9 @@ fn rule_set_should_edit_certificate_rules(rule: &str, mode: &str) {
         .arg("set")
         .arg("outbound")
         .arg(rule)
-        .arg("--cert-id")
+        .arg("cert-id")
         .arg(&cert_id)
+        .arg("--mode")
         .arg(mode)
         .assert()
         .success();
@@ -172,6 +176,35 @@ fn rule_set_should_edit_certificate_rules(rule: &str, mode: &str) {
     assert_eq!(&result["outbound"][rule][&cert_id]["mode"], mode);
 }
 
+#[test_case("partner", "all")]
+#[test_case("partner", "none")]
+#[test_case("partner", "whitelist")]
+#[serial_test::serial]
+fn rule_set_with_import_cert_should_add_to_keystore_and_rulestore(rule: &str, mode: &str) {
+    let (data_dir, resource_cert_dir) = prepare_test_dir_with_cert_resources();
+
+    Command::cargo_bin("ya-provider")
+        .unwrap()
+        .env("DATA_DIR", data_dir.path().to_str().unwrap())
+        .arg("rule")
+        .arg("set")
+        .arg("outbound")
+        .arg(rule)
+        .arg("import-cert")
+        .arg(resource_cert_dir.join("foo_ca-chain.cert.pem"))
+        .arg("--mode")
+        .arg(mode)
+        .assert()
+        .success();
+
+    let result = list_rules_command(data_dir.path());
+    let added_certs = list_certs(data_dir.path());
+
+    for cert in added_certs {
+        assert_eq!(result["outbound"][rule][cert]["mode"], mode);
+    }
+}
+
 #[test]
 #[serial_test::serial]
 fn removing_cert_should_also_remove_its_rule() {
@@ -188,8 +221,9 @@ fn removing_cert_should_also_remove_its_rule() {
         .arg("set")
         .arg("outbound")
         .arg(rule)
-        .arg("--cert-id")
+        .arg("cert-id")
         .arg(&cert_id)
+        .arg("--mode")
         .arg("all")
         .assert()
         .success();
@@ -235,6 +269,10 @@ fn add_certificate_to_keystore(data_dir: &Path, resource_cert_dir: &Path) -> Str
         .assert()
         .success();
 
+    list_certs(data_dir)[0].clone()
+}
+
+fn list_certs(data_dir: &Path) -> Vec<String> {
     let output = Command::cargo_bin("ya-provider")
         .unwrap()
         .env("DATA_DIR", data_dir.to_str().unwrap())
@@ -244,8 +282,12 @@ fn add_certificate_to_keystore(data_dir: &Path, resource_cert_dir: &Path) -> Str
         .output()
         .unwrap();
     let result: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap();
-
-    result[0]["ID"].as_str().unwrap().to_string()
+    result
+        .as_array()
+        .unwrap()
+        .iter()
+        .map(|v| v["ID"].as_str().unwrap().to_string())
+        .collect()
 }
 
 fn prepare_test_dir() -> TempDir {