1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
// SPDX-License-Identifier: Apache-2.0
//! The seam through which the local proxy serves the subscription models
//! (mz6, `docs/model-access-subscription.md` §4).
//!
//! The built-in registry provider [`PROVIDER_ID`] forwards to the Hive models
//! gateway with the instance's short-lived `hive_i_` access token. The proxy
//! and the Hive connector are sibling crates, so the connector registers the
//! one process-wide [`HiveTokenSource`] (as it does the runtime
//! `LeaseSource`) and the proxy asks it, never the Hive crate directly.
//!
//! The token opens every engine route of Hive, so it may only go to the
//! models origin [`models_origin`] derives from the enrolled Hive URL.
use std::sync::{Arc, RwLock};
/// Registry id of the built-in subscription-model provider.
pub const PROVIDER_ID: &str = "leanctx";
/// Error code of the local 403 while the lease does not grant the right
/// `models.subscription` (or the models origin is refused).
pub const UNAVAILABLE_CODE: &str = "subscription_models_unavailable";
/// Error code of the local 503 when no instance token can be had.
pub const TOKEN_UNAVAILABLE_CODE: &str = "hive_token_unavailable";
/// What a request to the gateway may carry now.
pub struct ModelsAdmission {
/// The models origin (`https://models.<hive host>`), no trailing slash.
pub origin: String,
/// The instance access token (`hive_i_…`).
pub token: zeroize::Zeroizing<String>,
/// A held route lets this one request through as a probe: its success
/// must be reported ([`HiveTokenSource::note_passed`]) to end the hold.
pub probe: bool,
}
// The token is a credential: debug output never shows it.
impl std::fmt::Debug for ModelsAdmission {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ModelsAdmission")
.field("origin", &self.origin)
.field("token", &"[redacted]")
.field("probe", &self.probe)
.finish()
}
}
/// Why a request is answered locally instead of going to the gateway.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ModelsRefusal {
/// Not connected, Hive paused, revoked, lease lapsed or without the
/// right `models.subscription`: 403 [`UNAVAILABLE_CODE`].
Unavailable,
/// The configured `[hive] models_url` is not on the Hive host: 403
/// [`UNAVAILABLE_CODE`]; checked before any token is read or minted.
OriginRefused(String),
/// An earlier gateway refusal (402/403/429) holds the `leanctx` routes
/// until `retry_after_s` from now (throttle style).
Held {
status: u16,
code: String,
retry_after_s: u64,
line: String,
},
/// No token could be minted: 503 [`TOKEN_UNAVAILABLE_CODE`].
TokenUnavailable,
}
/// Where the subscription-model access comes from (the Hive connector).
/// Every method may block (file reads, a token mint); the proxy calls them
/// off the async runtime and bounds the wait.
pub trait HiveTokenSource: Send + Sync {
/// Lease right, origin pin, holds, then the token; `fresh_token` mints
/// a new one (the gateway answered 401 to the cached one).
fn admit_route(&self, fresh_token: bool) -> Result<ModelsAdmission, ModelsRefusal>;
/// The gateway answered `status` (non-2xx) with `body`: record the hold
/// it sets, if any.
fn note_answer(&self, status: u16, retry_after: Option<&str>, body: &[u8]);
/// A probe passed: the holds of the `leanctx` routes end.
fn note_passed(&self);
}
static SOURCE: RwLock<Option<Arc<dyn HiveTokenSource>>> = RwLock::new(None);
/// Registers the process-wide source. A later call replaces it.
pub fn set_hive_token_source(source: Arc<dyn HiveTokenSource>) {
*SOURCE
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner) = Some(source);
}
/// The registered source; `None` without a Hive connector.
pub fn hive_token_source() -> Option<Arc<dyn HiveTokenSource>> {
SOURCE
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone()
}
fn is_loopback(host: &str) -> bool {
let host = host.trim_start_matches('[').trim_end_matches(']');
host == "localhost"
|| host
.parse::<std::net::IpAddr>()
.is_ok_and(|ip| ip.is_loopback())
}
fn origin_of(url: &reqwest::Url) -> String {
url.origin().ascii_serialization()
}
/// The models origin for an instance enrolled at `hive_url`.
///
/// Without `models_url` it is `https://models.<hive host>` (a loopback Hive,
/// as in tests, serves the models itself). A configured `models_url` is
/// honoured only when it is a bare origin whose host equals the Hive host or
/// is a direct subdomain of it, over HTTPS; loopback is allowed for tests.
/// Anything else is refused: the token must never reach a foreign host.
pub fn models_origin(hive_url: &str, models_url: Option<&str>) -> Result<String, String> {
let hive = reqwest::Url::parse(hive_url.trim()).map_err(|e| format!("Hive URL: {e}"))?;
let hive_host = hive
.host_str()
.ok_or("Hive URL has no host")?
.to_ascii_lowercase();
let Some(raw) = models_url.map(str::trim).filter(|s| !s.is_empty()) else {
if is_loopback(&hive_host) {
return Ok(origin_of(&hive));
}
if hive.scheme() != "https" {
return Err("Hive URL is not https".into());
}
let port = hive.port().map(|p| format!(":{p}")).unwrap_or_default();
return Ok(format!("https://models.{hive_host}{port}"));
};
let url = reqwest::Url::parse(raw).map_err(|e| format!("models_url: {e}"))?;
if !matches!(url.path(), "" | "/")
|| url.query().is_some()
|| url.fragment().is_some()
|| !url.username().is_empty()
|| url.password().is_some()
{
return Err("models_url must be a bare origin (scheme, host, port)".into());
}
let host = url
.host_str()
.ok_or("models_url has no host")?
.to_ascii_lowercase();
if is_loopback(&host) && matches!(url.scheme(), "http" | "https") {
return Ok(origin_of(&url));
}
if url.scheme() != "https" {
return Err("models_url must use https".into());
}
let direct_subdomain = host
.strip_suffix(&hive_host)
.and_then(|rest| rest.strip_suffix('.'))
.is_some_and(|label| !label.is_empty() && !label.contains('.'));
if host == hive_host || direct_subdomain {
Ok(origin_of(&url))
} else {
Err(format!(
"models_url host {host} is not the Hive host {hive_host} or a direct subdomain of it"
))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn the_models_origin_stays_on_the_hive_host() {
let hive = "https://hive.example.com";
assert_eq!(
models_origin(hive, None).unwrap(),
"https://models.hive.example.com"
);
for ok in [
"https://models.hive.example.com",
"https://gw.hive.example.com/",
"https://hive.example.com:8443",
"http://127.0.0.1:9000",
] {
assert!(
models_origin(hive, Some(ok)).is_ok(),
"{ok} must be honoured"
);
}
for refused in [
"https://models.evil.example",
"https://hive.example.com.evil.example",
"https://a.b.hive.example.com",
"https://evilhive.example.com",
"http://models.hive.example.com",
"https://models.hive.example.com/v1",
"https://user@models.hive.example.com",
] {
assert!(
models_origin(hive, Some(refused)).is_err(),
"{refused} must be refused"
);
}
assert_eq!(
models_origin("http://127.0.0.1:7000", None).unwrap(),
"http://127.0.0.1:7000"
);
assert!(models_origin("http://hive.example.com", None).is_err());
}
}