fix: make HcEngine plugin lookup use {publisher}/{plugin} as key

mitre · Sep 9, 2024 · 70f2bfd · 70f2bfd
1 parent 7574fcf
commit 70f2bfd
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 17 deletions.
diff --git a/hipcheck/src/engine.rs b/hipcheck/src/engine.rs
@@ -4,7 +4,10 @@ pub use crate::plugin::{HcPluginCore, PluginExecutor, PluginWithConfig};
 use crate::{
 	cache::plugin_cache::HcPluginCache,
 	hc_error,
-	plugin::{retrieve_plugins, Plugin, PluginManifest, PluginResponse, QueryResult, CURRENT_ARCH},
+	plugin::{
+		get_plugin_key, retrieve_plugins, Plugin, PluginManifest, PluginResponse, QueryResult,
+		CURRENT_ARCH,
+	},
 	policy::PolicyFile,
 	util::fs::{find_file_by_name, read_string},
 	Result,
@@ -51,9 +54,9 @@ fn default_policy_expr(
 	plugin: String,
 ) -> Result<Option<String>> {
 	let core = db.core();
-	// @Todo - plugins map should be keyed on publisher too
-	let Some(p_handle) = core.plugins.get(&plugin) else {
-		return Err(hc_error!("No such plugin {}::{}", publisher, plugin));
+	let key = get_plugin_key(publisher.as_str(), plugin.as_str());
+	let Some(p_handle) = core.plugins.get(&key) else {
+		return Err(hc_error!("No such plugin {}", key));
 	};
 	Ok(p_handle.get_default_policy_expr().cloned())
 }
@@ -64,9 +67,9 @@ fn default_query_explanation(
 	plugin: String,
 ) -> Result<Option<String>> {
 	let core = db.core();
-	// @Todo - plugins map should be keyed on publisher too
+	let key = get_plugin_key(publisher.as_str(), plugin.as_str());
 	let Some(p_handle) = core.plugins.get(&plugin) else {
-		return Err(hc_error!("No such plugin {}::{}", publisher, plugin));
+		return Err(hc_error!("No such plugin {}", key));
 	};
 	Ok(p_handle.get_default_query_explanation().cloned())
 }
@@ -80,9 +83,10 @@ fn query(
 ) -> Result<QueryResult> {
 	let runtime = RUNTIME.handle();
 	let core = db.core();
+	let hash_key = get_plugin_key(publisher.as_str(), plugin.as_str());
 	// Find the plugin
-	let Some(p_handle) = core.plugins.get(&plugin) else {
-		return Err(hc_error!("No such plugin {}::{}", publisher, plugin));
+	let Some(p_handle) = core.plugins.get(&hash_key) else {
+		return Err(hc_error!("No such plugin {}", hash_key));
 	};
 	// Initiate the query. If remote closed or we got our response immediately,
 	// return
@@ -128,8 +132,9 @@ pub fn async_query(
 ) -> BoxFuture<'static, Result<QueryResult>> {
 	async move {
 		// Find the plugin
-		let Some(p_handle) = core.plugins.get(&plugin) else {
-			return Err(hc_error!("No such plugin {}::{}", publisher, plugin));
+		let hash_key = get_plugin_key(publisher.as_str(), plugin.as_str());
+		let Some(p_handle) = core.plugins.get(&hash_key) else {
+			return Err(hc_error!("No such plugin {}", hash_key));
 		};
 		// Initiate the query. If remote closed or we got our response immediately,
 		// return
@@ -248,7 +253,7 @@ pub fn start_plugins(
 			})?;
 
 		let plugin = Plugin {
-			name: plugin_id.name.0.clone(),
+			name: plugin_id.to_policy_file_plugin_identifier(),
 			entrypoint,
 		};
 

diff --git a/hipcheck/src/main.rs b/hipcheck/src/main.rs
@@ -588,11 +588,11 @@ fn cmd_plugin(args: PluginArgs) {
 	let entrypoint1 = pathbuf![tgt_dir, "dummy_rand_data"];
 	let entrypoint2 = pathbuf![tgt_dir, "dummy_sha256"];
 	let plugin1 = Plugin {
-		name: "rand_data".to_owned(),
+		name: "dummy/rand_data".to_owned(),
 		entrypoint: entrypoint1.display().to_string(),
 	};
 	let plugin2 = Plugin {
-		name: "sha256".to_owned(),
+		name: "dummy/sha256".to_owned(),
 		entrypoint: entrypoint2.display().to_string(),
 	};
 	let plugin_executor = PluginExecutor::new(
@@ -627,7 +627,7 @@ fn cmd_plugin(args: PluginArgs) {
 				println!("Spawning");
 				futs.spawn(async_query(
 					arc_core,
-					"MITRE".to_owned(),
+					"dummy".to_owned(),
 					"rand_data".to_owned(),
 					"rand_data".to_owned(),
 					serde_json::json!(i),
@@ -639,7 +639,7 @@ fn cmd_plugin(args: PluginArgs) {
 		});
 	} else {
 		let res = engine.query(
-			"MITRE".to_owned(),
+			"dummy".to_owned(),
 			"rand_data".to_owned(),
 			"rand_data".to_owned(),
 			serde_json::json!(1),

diff --git a/hipcheck/src/plugin/mod.rs b/hipcheck/src/plugin/mod.rs
@@ -9,7 +9,7 @@ mod supported_arch;
 mod types;
 
 use crate::error::Result;
-pub use crate::plugin::{manager::*, plugin_id::PluginId, types::*};
+pub use crate::plugin::{get_plugin_key, manager::*, plugin_id::PluginId, types::*};
 pub use download_manifest::{ArchiveFormat, DownloadManifest, HashAlgorithm, HashWithDigest};
 pub use plugin_manifest::{PluginManifest, PluginName, PluginPublisher, PluginVersion};
 pub use retrieval::retrieve_plugins;

diff --git a/hipcheck/src/plugin/types.rs b/hipcheck/src/plugin/types.rs
@@ -620,3 +620,7 @@ impl From<Query> for PluginResponse {
 		}
 	}
 }
+
+pub fn get_plugin_key(publisher: &str, plugin: &str) -> String {
+	format!("{publisher}/{plugin}")
+}
diff --git a/plugins/dummy_rand_data/src/main.rs b/plugins/dummy_rand_data/src/main.rs
@@ -37,7 +37,7 @@ pub async fn handle_rand_data(mut session: QuerySession, key: u64) -> Result<()>
 
 	let sha_req = Query {
 		direction: QueryDirection::Request,
-		publisher: "MITRE".to_owned(),
+		publisher: "dummy".to_owned(),
 		plugin: "sha256".to_owned(),
 		query: "sha256".to_owned(),
 		key: json!(vec![sha_input]),
-Original file line number
+Diff line change
@@ Expand Up / @@ -620,3 +620,7 @@ impl From<Query> for PluginResponse { @@
     		}
     	}
     }
+    pub fn get_plugin_key(publisher: &str, plugin: &str) -> String {
+    	format!("{publisher}/{plugin}")
+    }