brooks_lib/hmds/
mod.rs

1// brooks-lib, Copyright 2026, Will Hawkins
2//
3// This file is part of brooks-lib.
4//
5// This file is free software: you can redistribute it and/or modify
6// it under the terms of the GNU General Public License as published by
7// the Free Software Foundation, either version 3 of the License, or
8// (at your option) any later version.
9//
10// This program is distributed in the hope that it will be useful,
11// but WITHOUT ANY WARRANTY; without even the implied warranty of
12// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
13// GNU General Public License for more details.
14//
15// You should have received a copy of the GNU General Public License
16// along with this program. If not, see <https://www.gnu.org/licenses/>.
17
18use std::{collections::HashMap, fs, io, path::PathBuf, str::FromStr};
19
20use chrono::{DateTime, Utc};
21use http::Uri;
22use serde::{Deserialize, Serialize};
23use serde_json::Value;
24
25#[cfg(all(unix, not(feature = "disable-domain")))]
26use std::path::Path;
27#[cfg(all(unix, not(feature = "disable-domain")))]
28use tokio::{
29    io::{AsyncReadExt, AsyncWriteExt},
30    net::UnixStream,
31};
32
33use crate::{
34    cdni::imd::spec::HostMetadata, environment::scope::Scopes, mel::types::Type,
35    verifying::verify::CdniVerificationKey,
36};
37
38#[derive(Debug, Default, Clone)]
39pub struct HmdsServerConfiguration {
40    pub(crate) path: Option<PathBuf>,
41    pub(crate) url: Option<http::Uri>,
42}
43
44impl HmdsServerConfiguration {
45    pub fn new_by_sense(config_str: &str) -> io::Result<Self> {
46        if config_str.starts_with("http://") || config_str.starts_with("https://") {
47            let url = http::Uri::try_from(config_str).map_err(io::Error::other)?;
48
49            Ok(HmdsServerConfiguration {
50                url: Some(url),
51                ..Default::default()
52            })
53        } else {
54            let path = PathBuf::from(config_str);
55
56            let exists = fs::exists(&path)?;
57            if !exists {
58                return Err(io::Error::other(format!(
59                    "Could not open connection to HMD server over UNIX domain socket at path {config_str}"
60                )));
61            }
62
63            Ok(HmdsServerConfiguration {
64                path: Some(path),
65                ..Default::default()
66            })
67        }
68    }
69
70    pub fn new_domain(path: PathBuf) -> Self {
71        Self {
72            path: Some(path),
73            ..Default::default()
74        }
75    }
76
77    pub fn new_http(http: Uri) -> Self {
78        Self {
79            url: Some(http),
80            ..Default::default()
81        }
82    }
83
84    pub fn is_http(&self) -> bool {
85        self.url.is_some()
86    }
87
88    pub fn is_domain(&self) -> bool {
89        self.path.is_some()
90    }
91}
92
93pub type HostMetadataCacheValue = (HostMetadata<CdniVerificationKey>, Scopes<Type>);
94
95pub struct HmdsConfiguration {
96    pub(crate) hmds_server: HmdsServerConfiguration,
97    pub(crate) hmds_cache: HashMap<String, (chrono::DateTime<Utc>, HostMetadataCacheValue)>,
98}
99
100impl From<HmdsServerConfiguration> for HmdsConfiguration {
101    fn from(value: HmdsServerConfiguration) -> Self {
102        HmdsConfiguration {
103            hmds_server: value,
104            hmds_cache: HashMap::new(),
105        }
106    }
107}
108
109#[derive(Serialize, Deserialize)]
110pub struct ExpirableJsonValue {
111    pub expiry: DateTime<Utc>,
112    pub host: String,
113    pub value: Value,
114}
115
116#[cfg(all(unix, not(feature = "disable-domain")))]
117async fn hmds_write_entire(s: &mut UnixStream, d: &[u8]) -> io::Result<usize> {
118    s.write_u64_le(d.len() as u64).await?;
119
120    let mut already_sent = 0usize;
121    loop {
122        s.writable().await?;
123        match s.try_write(&d[already_sent..]) {
124            Ok(n) => {
125                if n == 0 {
126                    return Ok(already_sent);
127                }
128                already_sent += n;
129                if already_sent == d.len() {
130                    return Ok(already_sent);
131                }
132                continue;
133            }
134            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
135                continue;
136            }
137            Err(e) => return Err(e),
138        }
139    }
140}
141
142#[cfg(all(unix, not(feature = "disable-domain")))]
143async fn hmds_read_entire(s: &mut UnixStream, d: &mut [u8]) -> io::Result<usize> {
144    let mut already_read = 0usize;
145    loop {
146        s.readable().await?;
147        match s.try_read(&mut d[already_read..]) {
148            Ok(n) => {
149                if n == 0 {
150                    return Ok(already_read);
151                }
152                already_read += n;
153                if already_read == d.len() {
154                    return Ok(already_read);
155                }
156                continue;
157            }
158            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
159                continue;
160            }
161            Err(e) => return Err(e),
162        }
163    }
164}
165
166pub async fn query_hmds(
167    query: &str,
168    config: &HmdsServerConfiguration,
169) -> io::Result<Option<(DateTime<Utc>, Value)>> {
170    #[cfg(all(unix, not(feature = "disable-domain")))]
171    if config.is_domain() {
172        return query_hmds_domain(query, config.path.as_ref().unwrap().as_path()).await;
173    }
174
175    if config.is_http() {
176        return query_hmds_http(query, config.url.as_ref().unwrap()).await;
177    } else {
178        Err(io::Error::other(
179            "Querying HMDS over HTTP is the only supported option and configuration is missing",
180        ))
181    }
182}
183
184#[cfg(all(unix, not(feature = "disable-domain")))]
185pub async fn query_hmds_domain(
186    query: &str,
187    server_path: &Path,
188) -> io::Result<Option<(DateTime<Utc>, Value)>> {
189    let mut socket = UnixStream::connect(server_path).await?;
190
191    // Wait for the socket to be readable
192    socket.writable().await?;
193
194    hmds_write_entire(&mut socket, query.as_bytes()).await?;
195
196    socket.readable().await?;
197
198    let incoming_buffer_size = socket.read_u64_le().await?;
199
200    let mut buf: Vec<u8> = vec![0; incoming_buffer_size as usize];
201
202    hmds_read_entire(&mut socket, &mut buf).await?;
203
204    let s = String::from_utf8(buf).map_err(|_| io::ErrorKind::InvalidData)?;
205
206    if s.is_empty() {
207        return Ok(None);
208    }
209
210    let result: ExpirableJsonValue =
211        serde_json::from_str(&s).map_err(|_| io::ErrorKind::InvalidData)?;
212
213    Ok(Some((result.expiry, result.value)))
214}
215
216pub async fn query_hmds_http(
217    query: &str,
218    server: &Uri,
219) -> io::Result<Option<(DateTime<Utc>, Value)>> {
220    let query_uri = reqwest::Url::from_str(&(server.to_string() + &format!("qry/?host={query}")))
221        .map_err(io::Error::other)?;
222
223    let query = reqwest::get(query_uri);
224
225    let response = query.await.map_err(io::Error::other)?;
226
227    match response.error_for_status() {
228        Ok(response) => {
229            let body = response.bytes().await.map_err(std::io::Error::other)?;
230
231            let result: ExpirableJsonValue =
232                serde_json::from_slice(&body).map_err(|_| io::ErrorKind::InvalidData)?;
233
234            Ok(Some((result.expiry, result.value)))
235        }
236        Err(e) => Err(io::Error::other(e)),
237    }
238}
brooks_lib/hmds/mod.rs

brooks_lib/hmds/
mod.rs
