memflow_registry/storage/
database.rs

1use std::{cmp::Reverse, collections::HashMap};
2
3use chrono::NaiveDateTime;
4use log::info;
5use memflow::plugins::plugin_analyzer::{PluginArchitecture, PluginDescriptorInfo, PluginFileType};
6use serde::{Deserialize, Serialize};
7
8use crate::{error::Result, rest::models::PluginInfo};
9
10use super::PluginMetadata;
11
12const DEFAULT_PLUGIN_VARIANTS: usize = 5;
13const MAX_PLUGIN_VARIANTS: usize = 50;
14
15#[derive(Default)]
16pub struct PluginDatabase {
17    plugins: HashMap<String, Vec<PluginVariant>>,
18}
19
20#[derive(Clone, Serialize, Deserialize)]
21pub struct PluginVariant {
22    pub digest: String,
23    pub signature: String,
24    pub created_at: NaiveDateTime,
25    pub descriptor: PluginDescriptorInfo,
26}
27
28#[derive(Debug, Default, Clone, Deserialize)]
29pub struct PluginDatabaseFindParams {
30    pub version: Option<String>,
31    pub memflow_plugin_version: Option<i32>,
32    pub file_type: Option<PluginFileType>,
33    pub architecture: Option<PluginArchitecture>,
34
35    // pagination parameters
36    pub skip: Option<usize>,
37    pub limit: Option<usize>,
38}
39
40impl PluginDatabase {
41    pub fn new() -> Self {
42        Self::default()
43    }
44
45    /// Inserts all plugin variants of this file into the database
46    pub fn insert_all(&mut self, metadata: &PluginMetadata) -> Result<()> {
47        for descriptor in metadata.descriptors.iter() {
48            info!(
49                "adding plugin variant to db: digest={}; created_at={}; descriptor={:?}",
50                metadata.digest, metadata.created_at, descriptor
51            );
52
53            let entry = self.plugins.entry(descriptor.name.clone()).or_default();
54
55            // sort by plugin_version and created_at
56            // metadata is guaranteed to contain at least one descriptor and the plugin_version is identical for all connectors of a file.
57            let search_key = (
58                metadata.descriptors.first().unwrap().plugin_version,
59                metadata.created_at,
60            );
61            match entry.binary_search_by_key(&Reverse(search_key), |entry| {
62                Reverse((entry.descriptor.plugin_version, entry.created_at))
63            }) {
64                Ok(_) => unreachable!(), // element already in vector @ `pos` // TODO: check for duplicate entries
65                Err(pos) => entry.insert(
66                    pos,
67                    PluginVariant {
68                        digest: metadata.digest.clone(),
69                        signature: metadata.signature.clone(),
70                        created_at: metadata.created_at,
71                        descriptor: descriptor.clone(),
72                    },
73                ),
74            }
75        }
76
77        Ok(())
78    }
79
80    /// Returns a list of all plugin names and their descriptions.
81    pub fn plugins(&self) -> Vec<PluginInfo> {
82        let mut plugins = self
83            .plugins
84            .iter()
85            .flat_map(|(key, variants)| {
86                variants.iter().map(|variant| PluginInfo {
87                    name: key.to_owned(),
88                    description: variant.descriptor.description.clone(),
89                })
90            })
91            .collect::<Vec<_>>();
92        plugins.sort_by(|a, b| a.name.cmp(&b.name));
93        plugins.dedup_by(|a, b| a.name == b.name);
94        plugins
95    }
96
97    /// Retrieves a specific digest
98    #[allow(unused)]
99    pub fn find_by_digest(&self, digest: &str) -> Option<PluginVariant> {
100        self.plugins
101            .iter()
102            .find_map(|(_, variants)| variants.iter().find(|variant| variant.digest == digest))
103            .cloned()
104    }
105
106    /// Removes all entries with the specified digest from the database
107    pub fn delete_by_digest(&mut self, digest: &str) {
108        for plugin in self.plugins.iter_mut() {
109            plugin.1.retain(|variant| variant.digest != digest);
110        }
111    }
112
113    /// Retrieves a list of variants for a specific plugin.
114    /// Additional search parameters can be specified.
115    pub fn plugin_variants(
116        &self,
117        plugin_name: &str,
118        params: PluginDatabaseFindParams,
119    ) -> Vec<PluginVariant> {
120        self.plugins
121            .get(plugin_name)
122            .map(|variants| {
123                variants
124                    .iter()
125                    .skip(params.skip.unwrap_or(0))
126                    .filter(|p| p.descriptor.name == plugin_name)
127                    .filter(|p| {
128                        if let Some(version) = &params.version {
129                            // version can match the version directly or the corresponding digest
130                            if *version != p.descriptor.version
131                                && *version != p.digest[..version.len()]
132                            {
133                                return false;
134                            }
135                        }
136
137                        if let Some(memflow_plugin_version) = params.memflow_plugin_version {
138                            if memflow_plugin_version != p.descriptor.plugin_version {
139                                return false;
140                            }
141                        }
142
143                        if let Some(file_type) = params.file_type {
144                            if file_type != p.descriptor.file_type {
145                                return false;
146                            }
147                        }
148
149                        if let Some(architecture) = params.architecture {
150                            if architecture != p.descriptor.architecture {
151                                return false;
152                            }
153                        }
154
155                        true
156                    })
157                    .take(
158                        params
159                            .limit
160                            .unwrap_or(DEFAULT_PLUGIN_VARIANTS)
161                            .min(MAX_PLUGIN_VARIANTS),
162                    )
163                    .cloned()
164                    .collect::<Vec<_>>()
165            })
166            .unwrap_or_default()
167    }
168}
memflow_registry/storage/database.rs

memflow_registry/storage/
database.rs
