apollo_environment_detector/
lib.rs

1//! # Compute Environment Detector
2//!
3//! This library provides two functions for easily detecting a [`ComputeEnvironment`] based on a
4//! given weighted threshold.
5//!
6//! # Examples
7//!
8//! ```
9//! use apollo_environment_detector::{detect, detect_one, MAX_INDIVIDUAL_WEIGHTING};
10//!
11//! // Attempt to detect multiple environments based on a weighting.
12//! let compute_envs = detect(MAX_INDIVIDUAL_WEIGHTING);
13//! println!("{:?}", compute_envs);
14//!
15//! // Attempt to detect a single environment based on a weighting.
16//! let compute_env = detect_one(MAX_INDIVIDUAL_WEIGHTING);
17//! println!("{:?}", compute_env);
18//! ```
19
20#![warn(missing_docs)]
21
22use std::{cmp::Ordering, collections::HashSet, ops::Deref};
23
24mod detector;
25use detector::Detector;
26mod env_vars;
27mod environment;
28pub use environment::{CloudProvider, ComputeEnvironment};
29mod smbios;
30use smbios::Smbios;
31mod specificity;
32use specificity::Specificity as _;
33
34/// Represents the maximum weighting of all supported detectors (`2^15`).
35///
36/// This maximum weighting was chosen in order to have enough buffer compared
37/// to avoid thresholding and overflows when using multiple detectors.
38pub const MAX_TOTAL_WEIGHTING: u16 = 2 << 14;
39
40/// Represents the maximum individual detector weighting.
41///
42/// There are currently 2 supported detectors:
43/// - SMBIOS
44/// - Environment Variables
45pub const MAX_INDIVIDUAL_WEIGHTING: u16 = MAX_TOTAL_WEIGHTING / 2;
46
47/// Detect a single, most likely [`ComputeEnvironment`] above a certain weighted threshold.
48pub fn detect_one(threshold: u16) -> Option<ComputeEnvironment> {
49    detect(threshold).first().copied()
50}
51
52/// Detect potential [`ComputeEnvironment`]s above a certain weighted threshold.
53///
54/// Returns an ordered [`Vec`] with the highest weighted candidates first.
55pub fn detect(threshold: u16) -> Vec<ComputeEnvironment> {
56    let detectors: Vec<_> = ComputeEnvironment::iter().map(|ce| ce.detector()).collect();
57
58    // Read current environment variables
59    let env_vars: HashSet<_> = detectors
60        .iter()
61        .flat_map(|detector| detector.env_vars)
62        .filter(|var| env_vars::hasenv(var))
63        .map(Deref::deref)
64        .collect();
65
66    // Read SMBIOS data
67    let smbios = Smbios::detect();
68
69    // Run detectors against env vars and SMBIOS data
70    detect_inner(detectors, smbios, env_vars, threshold)
71}
72
73fn detect_inner(
74    detectors: Vec<Detector>,
75    smbios: Smbios,
76    env_vars: HashSet<&'static str>,
77    threshold: u16,
78) -> Vec<ComputeEnvironment> {
79    let mut detectors: Vec<_> = detectors
80        .into_iter()
81        .filter_map(|detector| {
82            let score = detector.detect(&smbios, &env_vars);
83            if score >= threshold {
84                Some((detector, score))
85            } else {
86                None
87            }
88        })
89        .collect();
90
91    detectors.sort_by(|(left, left_score), (right, right_score)| {
92        match Ord::cmp(left_score, right_score) {
93            Ordering::Equal => left
94                .specificity_cmp(right)
95                .unwrap_or(Ordering::Equal)
96                .reverse(),
97            o => o.reverse(),
98        }
99    });
100
101    detectors
102        .into_iter()
103        .map(|(detector, _)| detector.environment)
104        .collect()
105}
106
107#[cfg(test)]
108mod tests {
109    use rstest::{fixture, rstest};
110
111    use super::*;
112
113    #[fixture]
114    fn detectors() -> Vec<Detector> {
115        ComputeEnvironment::iter().map(|ce| ce.detector()).collect()
116    }
117
118    #[rstest]
119    fn test_complete(
120        #[values(
121            ComputeEnvironment::AwsEc2,
122            ComputeEnvironment::AwsEcs,
123            ComputeEnvironment::AwsLambda,
124            ComputeEnvironment::AwsKubernetes,
125            ComputeEnvironment::AwsNomad,
126            ComputeEnvironment::AzureContainerApps,
127            ComputeEnvironment::AzureContainerAppsJob,
128            ComputeEnvironment::AzureContainerInstance,
129            ComputeEnvironment::AzureKubernetes,
130            ComputeEnvironment::AzureVM,
131            ComputeEnvironment::AzureNomad,
132            ComputeEnvironment::GcpCloudRunGen1,
133            ComputeEnvironment::GcpCloudRunGen2,
134            ComputeEnvironment::GcpCloudRunJob,
135            ComputeEnvironment::GcpComputeEngine,
136            ComputeEnvironment::GcpKubernetes,
137            ComputeEnvironment::GcpNomad,
138            ComputeEnvironment::Kubernetes,
139            ComputeEnvironment::Nomad,
140            ComputeEnvironment::Qemu
141        )]
142        environment: ComputeEnvironment,
143        detectors: Vec<Detector>,
144    ) {
145        let smbios: Smbios = environment.detector().smbios.clone().into();
146        let env_vars: HashSet<_> = environment
147            .detector()
148            .env_vars
149            .iter()
150            .map(Deref::deref)
151            .collect();
152
153        let result = detect_inner(detectors, smbios, env_vars, u16::MIN);
154
155        assert_eq!(result.first(), Some(&environment));
156    }
157
158    #[rstest]
159    fn test_missing_1_env_var(
160        #[values(
161            ComputeEnvironment::AwsEc2,
162            ComputeEnvironment::AwsEcs,
163            ComputeEnvironment::AwsLambda,
164            ComputeEnvironment::AwsKubernetes,
165            ComputeEnvironment::AwsNomad,
166            // Accepted risk: these tests will fail if we remove one of the env var specific to
167            // Azure Container Apps
168            // ComputeEnvironment::AzureContainerApps,
169            // ComputeEnvironment::AzureContainerAppsJob,
170            ComputeEnvironment::AzureContainerInstance,
171            ComputeEnvironment::AzureKubernetes,
172            ComputeEnvironment::AzureVM,
173            ComputeEnvironment::AzureNomad,
174            ComputeEnvironment::GcpCloudRunGen1,
175            ComputeEnvironment::GcpCloudRunGen2,
176            ComputeEnvironment::GcpCloudRunJob,
177            ComputeEnvironment::GcpComputeEngine,
178            ComputeEnvironment::GcpKubernetes,
179            ComputeEnvironment::GcpNomad,
180            ComputeEnvironment::Kubernetes,
181            ComputeEnvironment::Nomad,
182            ComputeEnvironment::Qemu
183        )]
184        environment: ComputeEnvironment,
185        detectors: Vec<Detector>,
186    ) {
187        let smbios: Smbios = environment.detector().smbios.clone().into();
188        let env_vars = environment.detector().env_vars.to_vec();
189
190        for i in 0..(env_vars.len()) {
191            let mut env_vars = env_vars.clone();
192            let removed = env_vars.remove(i);
193            let env_vars = env_vars.into_iter().collect();
194
195            let result = detect_inner(detectors.clone(), smbios.clone(), env_vars, u16::MIN);
196
197            assert_eq!(
198                result.first(),
199                Some(&environment),
200                "mismatch with {removed} removed"
201            );
202        }
203    }
204
205    #[rstest]
206    fn test_missing_2_env_var(
207        #[values(
208            ComputeEnvironment::AwsEc2,
209            ComputeEnvironment::AwsEcs,
210            ComputeEnvironment::AwsLambda,
211            ComputeEnvironment::AwsKubernetes,
212            ComputeEnvironment::AwsNomad,
213            // Accepted risk: these tests will fail if we remove two of the env vars specific to
214            // Azure Container Apps
215            // ComputeEnvironment::AzureContainerApps,
216            // ComputeEnvironment::AzureContainerAppsJob,
217            ComputeEnvironment::AzureContainerInstance,
218            ComputeEnvironment::AzureKubernetes,
219            ComputeEnvironment::AzureVM,
220            ComputeEnvironment::AzureNomad,
221            ComputeEnvironment::GcpCloudRunGen1,
222            ComputeEnvironment::GcpCloudRunGen2,
223            ComputeEnvironment::GcpCloudRunJob,
224            ComputeEnvironment::GcpComputeEngine,
225            ComputeEnvironment::GcpKubernetes,
226            ComputeEnvironment::GcpNomad,
227            ComputeEnvironment::Kubernetes,
228            ComputeEnvironment::Nomad,
229            ComputeEnvironment::Qemu
230        )]
231        environment: ComputeEnvironment,
232        detectors: Vec<Detector>,
233    ) {
234        let smbios: Smbios = environment.detector().smbios.clone().into();
235        let env_vars = environment.detector().env_vars.to_vec();
236
237        for i in 0..(env_vars.len()) {
238            for j in 0..(env_vars.len() - 1) {
239                let mut env_vars = env_vars.clone();
240                let removed_1 = env_vars.remove(i);
241                let removed_2 = env_vars.remove(j);
242                let env_vars = env_vars.into_iter().collect();
243
244                let result = detect_inner(detectors.clone(), smbios.clone(), env_vars, u16::MIN);
245
246                assert_eq!(
247                    result.first(),
248                    Some(&environment),
249                    "mismatch with {removed_1} and {removed_2} removed"
250                );
251            }
252        }
253    }
254}
apollo_environment_detector/lib.rs

apollo_environment_detector/
lib.rs
