1use std::collections::HashMap;
9use std::sync::atomic::{AtomicU64, Ordering};
10use std::sync::Mutex;
11use std::time::{Duration, Instant};
12
13use prost::Message;
14
15use crate::broker::capabilities::{handoff_transport_available, CAP_HANDLE_PASSING};
16use crate::broker::lifecycle::names::{validate_service_name, validate_version, PipePathError};
17use crate::broker::protocol::{
18 hello_reply::Result as HelloReplyResult, validate_frame_envelope, ErrorCode, Frame, FrameKind,
19 FrameValidationError, Hello, HelloReply, Negotiated, Refused, ServiceDefinition,
20 CONTROL_PAYLOAD_PROTOCOL, PROTOCOL_VERSION,
21};
22use crate::broker::server::handoff::{
23 AcknowledgedHandoff, ExpiredHandoff, HandoffAckError, HandoffAckRegistry, HandoffToken,
24 HandoffTokenStore, PendingHandoffBackend,
25};
26use crate::broker::server::session_token::{SessionTokenAuthority, SessionTokenRejection};
27use crate::broker::server::version_allow_list::{check_version_allowed, VersionPolicyBlock};
28use crate::broker::server::TraceContext;
29
30const DEFAULT_KEEPALIVE_SECS: u64 = 30 * 60;
31const DEFAULT_RATE_LIMIT_MAX_PER_WINDOW: u32 = 256;
32const DEFAULT_RATE_LIMIT_WINDOW: Duration = Duration::from_secs(1);
33
34#[derive(Clone, Debug, PartialEq, Eq)]
36pub struct PeerIdentity {
37 pub pid: u32,
39 pub uid_or_sid: String,
41}
42
43#[derive(Clone, Debug)]
45pub struct HelloRequest {
46 pub frame: Frame,
48 pub hello: Hello,
50 pub peer: PeerIdentity,
52}
53
54impl HelloRequest {
55 pub fn decode(frame: Frame, peer: PeerIdentity) -> Result<Self, Refused> {
57 validate_frame_envelope(&frame, FrameKind::Request, CONTROL_PAYLOAD_PROTOCOL).map_err(
58 |error| match error {
59 FrameValidationError::EnvelopeVersion { .. } => refused(
60 ErrorCode::ErrorVersionUnsupported,
61 "frame envelope_version is not v1",
62 0,
63 ),
64 FrameValidationError::Kind { .. } => refused(
65 ErrorCode::ErrorPeerRejected,
66 "Hello frame kind must be REQUEST",
67 0,
68 ),
69 FrameValidationError::PayloadProtocol { .. } => refused(
70 ErrorCode::ErrorPeerRejected,
71 "Hello frame payload_protocol must be control-plane",
72 0,
73 ),
74 FrameValidationError::PayloadEncoding { .. } => refused(
75 ErrorCode::ErrorPeerRejected,
76 "Hello payload must not be compressed",
77 0,
78 ),
79 },
80 )?;
81 let hello = Hello::decode(frame.payload.as_slice())
82 .map_err(|_| refused(ErrorCode::ErrorPeerRejected, "malformed Hello payload", 0))?;
83 Ok(Self { frame, hello, peer })
84 }
85
86 pub fn trace_context(&self) -> TraceContext {
88 TraceContext::from_frame(&self.frame)
89 }
90}
91
92#[derive(Clone, Debug)]
94pub struct RegisteredBackend {
95 pub service_definition: ServiceDefinition,
97 pub daemon_version: String,
99 pub backend_pipe: String,
101 pub server_capabilities: u64,
103}
104
105#[derive(Debug)]
107pub struct HelloHandler {
108 backends: HashMap<String, RegisteredBackend>,
109 next_connection_id: AtomicU64,
110 rate_limiter: PeerRateLimiter,
111 handoff_tokens: Mutex<HandoffTokenStore>,
112 handoff_acks: Mutex<HandoffAckRegistry>,
113 session_tokens: Option<Mutex<SessionTokenAuthority>>,
122}
123
124impl HelloHandler {
125 pub fn new() -> Self {
127 Self {
128 backends: HashMap::new(),
129 next_connection_id: AtomicU64::new(1),
130 rate_limiter: PeerRateLimiter::default(),
131 handoff_tokens: Mutex::new(HandoffTokenStore::new()),
132 handoff_acks: Mutex::new(HandoffAckRegistry::new()),
133 session_tokens: None,
134 }
135 }
136
137 pub fn with_session_token_authority(mut self, authority: SessionTokenAuthority) -> Self {
142 self.session_tokens = Some(Mutex::new(authority));
143 self
144 }
145
146 pub fn session_token_authority(
150 &self,
151 ) -> Option<std::sync::MutexGuard<'_, SessionTokenAuthority>> {
152 self.session_tokens
153 .as_ref()
154 .map(|m| m.lock().unwrap_or_else(|poisoned| poisoned.into_inner()))
155 }
156
157 pub fn with_handoff_ack_deadline(self, ack_deadline: Duration) -> Self {
159 *self.handoff_ack_registry() = HandoffAckRegistry::with_ack_deadline(ack_deadline);
160 self
161 }
162
163 pub fn handoff_token_store(&self) -> std::sync::MutexGuard<'_, HandoffTokenStore> {
169 self.handoff_tokens
170 .lock()
171 .unwrap_or_else(|poisoned| poisoned.into_inner())
172 }
173
174 pub fn handoff_ack_registry(&self) -> std::sync::MutexGuard<'_, HandoffAckRegistry> {
181 self.handoff_acks
182 .lock()
183 .unwrap_or_else(|poisoned| poisoned.into_inner())
184 }
185
186 pub fn acknowledge_handoff(
191 &self,
192 token: &HandoffToken,
193 now: Instant,
194 ) -> Result<AcknowledgedHandoff, HandoffAckError> {
195 let mut acks = self.handoff_ack_registry();
196 let mut tokens = self.handoff_token_store();
197 acks.acknowledge(&mut tokens, token, now)
198 }
199
200 pub fn expire_overdue_handoffs(&self, now: Instant) -> Vec<ExpiredHandoff> {
205 let mut acks = self.handoff_ack_registry();
206 let mut tokens = self.handoff_token_store();
207 acks.expire_overdue(&mut tokens, now)
208 }
209
210 pub fn with_rate_limit(mut self, max_per_window: u32, window: Duration) -> Self {
212 self.rate_limiter = PeerRateLimiter::new(max_per_window, window);
213 self
214 }
215
216 pub fn with_backend(mut self, backend: RegisteredBackend) -> Result<Self, HelloHandlerError> {
218 validate_service_name_for_result(&backend.service_definition.service_name)?;
219 if !backend.service_definition.min_version.is_empty() {
220 validate_version_for_result(&backend.service_definition.min_version)?;
221 }
222 for version in &backend.service_definition.version_allow_list {
223 validate_version_for_result(version)?;
224 }
225 self.backends
226 .insert(backend.service_definition.service_name.clone(), backend);
227 Ok(self)
228 }
229
230 pub fn handle_frame(&self, frame: Frame, peer: PeerIdentity) -> HelloReply {
232 match HelloRequest::decode(frame, peer) {
233 Ok(request) => self.handle_request(&request),
234 Err(refused) => refused_reply(refused),
235 }
236 }
237
238 pub fn handle_request(&self, request: &HelloRequest) -> HelloReply {
240 let hello = &request.hello;
241 if let Some(refused) = validate_hello_shape(hello, &request.peer) {
242 return refused_reply(refused);
243 }
244 if let Some(retry_after) = self.rate_limiter.check(request.peer.pid) {
245 return refused_reply(refused(
246 ErrorCode::ErrorRateLimited,
247 "Hello rate limit exceeded",
248 duration_to_retry_ms(retry_after),
249 ));
250 }
251
252 let Some(backend) = self.backends.get(&hello.service_name) else {
253 return refused_reply(refused(
254 ErrorCode::ErrorServiceUnknown,
255 "service is not registered",
256 0,
257 ));
258 };
259
260 if let Some(refused) = validate_version_policy(hello, &backend.service_definition) {
261 return refused_reply(refused);
262 }
263
264 if let Some(refused) = self.validate_session_token(hello) {
265 return refused_reply(refused);
266 }
267
268 let connection_id = self.next_connection_id.fetch_add(1, Ordering::Relaxed);
269 let handle_passed_token =
270 self.issue_handoff_token(hello.client_capabilities, &hello.service_name);
271 let mut server_capabilities = backend.server_capabilities;
272 if !handle_passed_token.is_empty() {
273 server_capabilities |= CAP_HANDLE_PASSING;
274 }
275 refused_or_negotiated(HelloReplyResult::Negotiated(Negotiated {
276 negotiated_protocol: PROTOCOL_VERSION,
277 daemon_version: backend.daemon_version.clone(),
278 backend_pipe: backend.backend_pipe.clone(),
279 warnings: Vec::new(),
280 server_capabilities,
281 keepalive_interval_secs: if hello.client_keepalive_secs == 0 {
282 DEFAULT_KEEPALIVE_SECS
283 } else {
284 hello.client_keepalive_secs
285 },
286 handle_passed_token,
287 connection_id,
288 }))
289 }
290
291 fn validate_session_token(&self, hello: &Hello) -> Option<Refused> {
300 let authority = self.session_tokens.as_ref()?;
301 let authority = authority
302 .lock()
303 .unwrap_or_else(|poisoned| poisoned.into_inner());
304 match authority.validate(&hello.auth_token, &hello.service_name) {
305 Ok(()) => None,
306 Err(rejection) => Some(refused(
307 ErrorCode::ErrorPeerRejected,
308 session_token_rejection_reason(rejection),
309 0,
310 )),
311 }
312 }
313
314 fn issue_handoff_token(&self, client_capabilities: u64, service_name: &str) -> Vec<u8> {
327 if client_capabilities & CAP_HANDLE_PASSING == 0 || !handoff_transport_available() {
328 return Vec::new();
329 }
330 let now = Instant::now();
331 let mut acks = self.handoff_ack_registry();
333 let mut tokens = self.handoff_token_store();
334 match tokens.issue(now) {
335 Ok(token) => {
336 acks.register(token, PendingHandoffBackend::for_service(service_name), now);
337 token.into_bytes().to_vec()
338 }
339 Err(_) => Vec::new(),
340 }
341 }
342}
343
344impl Default for HelloHandler {
345 fn default() -> Self {
346 Self::new()
347 }
348}
349
350#[derive(Debug, thiserror::Error)]
352pub enum HelloHandlerError {
353 #[error(transparent)]
355 PipePath(#[from] PipePathError),
356}
357
358#[derive(Debug)]
360struct PeerRateLimiter {
361 max_per_window: u32,
362 window: Duration,
363 entries: Mutex<HashMap<u32, PeerRateWindow>>,
364}
365
366impl PeerRateLimiter {
367 fn new(max_per_window: u32, window: Duration) -> Self {
368 Self {
369 max_per_window: max_per_window.max(1),
370 window: if window.is_zero() {
371 Duration::from_millis(1)
372 } else {
373 window
374 },
375 entries: Mutex::new(HashMap::new()),
376 }
377 }
378
379 fn check(&self, pid: u32) -> Option<Duration> {
380 if pid == 0 {
381 return None;
382 }
383
384 let now = Instant::now();
385 let mut entries = self
386 .entries
387 .lock()
388 .unwrap_or_else(|poisoned| poisoned.into_inner());
389 let entry = entries.entry(pid).or_insert(PeerRateWindow {
390 started_at: now,
391 count: 0,
392 });
393 let elapsed = now.duration_since(entry.started_at);
394 if elapsed >= self.window {
395 entry.started_at = now;
396 entry.count = 0;
397 }
398
399 if entry.count < self.max_per_window {
400 entry.count += 1;
401 None
402 } else {
403 Some(self.window.saturating_sub(elapsed))
404 }
405 }
406}
407
408impl Default for PeerRateLimiter {
409 fn default() -> Self {
410 Self::new(DEFAULT_RATE_LIMIT_MAX_PER_WINDOW, DEFAULT_RATE_LIMIT_WINDOW)
411 }
412}
413
414#[derive(Debug)]
415struct PeerRateWindow {
416 started_at: Instant,
417 count: u32,
418}
419
420pub(crate) fn validate_hello_shape(hello: &Hello, peer: &PeerIdentity) -> Option<Refused> {
428 if hello.client_min_protocol > PROTOCOL_VERSION || hello.client_max_protocol < PROTOCOL_VERSION
429 {
430 return Some(refused(
431 ErrorCode::ErrorVersionUnsupported,
432 "client protocol range does not include v1",
433 0,
434 ));
435 }
436 if validate_service_name(&hello.service_name).is_err() {
437 return Some(refused(
438 ErrorCode::ErrorPeerRejected,
439 "invalid service_name",
440 0,
441 ));
442 }
443 if hello.wanted_version.len() > 64 || validate_version(&hello.wanted_version).is_err() {
444 return Some(refused(
445 ErrorCode::ErrorPeerRejected,
446 "invalid wanted_version",
447 0,
448 ));
449 }
450 if hello.client_version.len() > 128 {
451 return Some(refused(
452 ErrorCode::ErrorPeerRejected,
453 "client_version exceeds 128 bytes",
454 0,
455 ));
456 }
457 if hello.client_lib_name.len() > 64 || hello.client_lib_version.len() > 64 {
458 return Some(refused(
459 ErrorCode::ErrorPeerRejected,
460 "client_lib fields exceed 64 bytes",
461 0,
462 ));
463 }
464 if hello.peer_pid != 0 && peer.pid != 0 && hello.peer_pid != peer.pid {
467 return Some(refused(
468 ErrorCode::ErrorPeerRejected,
469 "peer_pid does not match verified peer",
470 0,
471 ));
472 }
473 None
474}
475
476fn validate_version_policy(hello: &Hello, service: &ServiceDefinition) -> Option<Refused> {
477 match check_version_allowed(&hello.wanted_version, service) {
478 Ok(()) => None,
479 Err(VersionPolicyBlock::BelowMinVersion) => Some(refused(
480 ErrorCode::ErrorVersionBlocked,
481 "wanted_version is below min_version",
482 30_000,
483 )),
484 Err(VersionPolicyBlock::OutsideAllowList) => Some(refused(
485 ErrorCode::ErrorVersionBlocked,
486 "wanted_version is not in version_allow_list",
487 30_000,
488 )),
489 }
490}
491
492fn validate_service_name_for_result(name: &str) -> Result<(), HelloHandlerError> {
493 validate_service_name(name).map_err(HelloHandlerError::PipePath)
494}
495
496fn validate_version_for_result(version: &str) -> Result<(), HelloHandlerError> {
497 validate_version(version).map_err(HelloHandlerError::PipePath)
498}
499
500fn duration_to_retry_ms(duration: Duration) -> u64 {
501 let millis = duration.as_millis().max(1);
502 u64::try_from(millis).unwrap_or(u64::MAX)
503}
504
505fn refused(code: ErrorCode, reason: impl Into<String>, retry_after_ms: u64) -> Refused {
506 Refused {
507 reason: reason.into(),
508 daemon_min_protocol: PROTOCOL_VERSION,
509 daemon_max_protocol: PROTOCOL_VERSION,
510 code: code as i32,
511 details: HashMap::new(),
512 retry_after_ms,
513 }
514}
515
516fn refused_reply(refused: Refused) -> HelloReply {
517 refused_or_negotiated(HelloReplyResult::Refused(refused))
518}
519
520fn session_token_rejection_reason(rejection: SessionTokenRejection) -> &'static str {
524 match rejection {
525 SessionTokenRejection::MalformedLength => "session token malformed: wrong byte length",
526 SessionTokenRejection::BrokerHalfMismatch => {
527 "session invalidated: broker was restarted or rotated since this token was issued"
528 }
529 SessionTokenRejection::DaemonUnknown => {
530 "session invalidated: daemon is no longer registered"
531 }
532 SessionTokenRejection::DaemonHalfMismatch => {
533 "session invalidated: daemon was restarted since this token was issued"
534 }
535 }
536}
537
538fn refused_or_negotiated(result: HelloReplyResult) -> HelloReply {
539 HelloReply {
540 result: Some(result),
541 }
542}