1use std::io::{self, Read, Write};
4use std::thread;
5use std::time::{Duration, Instant};
6
7use prost::Message;
8
9use crate::broker::backend_lifecycle::identity::{DaemonProcess, IdentityError};
10use crate::broker::backend_lifecycle::verify_pid::{self, ProcessHandle, VerifyPidError};
11use crate::broker::protocol::{
12 self, read_frame, write_frame, Endpoint, Frame, FrameKind, FramingError, PayloadEncoding,
13 ENVELOPE_VERSION, MAX_FRAME_BYTES, PROTOCOL_VERSION,
14};
15
16pub const PROBE_NONCE_BYTES: usize = 32;
18const NONBLOCKING_POLL_INTERVAL: Duration = Duration::from_millis(5);
19
20pub use crate::broker::protocol::registry::BACKEND_HANDLE_PROBE_PAYLOAD_PROTOCOL;
26
27pub const DEFAULT_ENDPOINT_PROBE_TIMEOUT: Duration = Duration::from_millis(500);
29
30pub fn probe_endpoint(
32 endpoint: &Endpoint,
33 expected: &DaemonProcess,
34) -> Result<ProcessHandle, ProbeError> {
35 probe_endpoint_with_timeout(endpoint, expected, DEFAULT_ENDPOINT_PROBE_TIMEOUT)
36}
37
38pub fn probe_endpoint_with_timeout(
49 endpoint: &Endpoint,
50 expected: &DaemonProcess,
51 timeout: Duration,
52) -> Result<ProcessHandle, ProbeError> {
53 if !same_endpoint(endpoint, &expected.ipc_endpoint) {
54 return Err(ProbeError::EndpointMismatch);
55 }
56 let process_handle =
57 verify_pid::verify_daemon_process(expected).map_err(ProbeError::VerifyPid)?;
58 probe_endpoint_response_with_timeout(endpoint, expected, timeout)?;
59 Ok(process_handle)
60}
61
62pub fn same_endpoint(left: &Endpoint, right: &Endpoint) -> bool {
64 left.namespace_id == right.namespace_id && left.path == right.path
65}
66
67pub fn probe_endpoint_response(
74 endpoint: &Endpoint,
75 expected: &DaemonProcess,
76) -> Result<(), EndpointProbeError> {
77 probe_endpoint_response_with_timeout(endpoint, expected, DEFAULT_ENDPOINT_PROBE_TIMEOUT)
78}
79
80pub fn probe_endpoint_response_with_timeout(
82 endpoint: &Endpoint,
83 expected: &DaemonProcess,
84 timeout: Duration,
85) -> Result<(), EndpointProbeError> {
86 let mut nonce = [0_u8; PROBE_NONCE_BYTES];
87 getrandom::fill(&mut nonce).map_err(EndpointProbeError::Random)?;
88 let request_id = u64::from_le_bytes(nonce[..8].try_into().expect("nonce has 8 bytes"));
89 let request_frame = endpoint_probe_request_frame(request_id, &nonce);
90 let mut request_bytes = Vec::new();
91 request_frame
92 .encode(&mut request_bytes)
93 .map_err(EndpointProbeError::EncodeFrame)?;
94
95 let deadline = Instant::now() + timeout;
96 let mut stream = connect_endpoint_with_deadline(endpoint, deadline)?;
97 stream
98 .set_nonblocking(true)
99 .map_err(EndpointProbeError::ConfigureNonblocking)?;
100 write_probe_frame_with_deadline(&mut stream, &request_bytes, deadline)?;
101
102 let response_bytes = read_probe_frame_with_deadline(&mut stream, deadline)?;
103 let response_frame =
104 Frame::decode(response_bytes.as_slice()).map_err(EndpointProbeError::DecodeFrame)?;
105 validate_endpoint_probe_response_frame(&response_frame, request_id)?;
106 let actual = decode_response_identity(&response_frame.payload, &nonce)?;
107 if !same_daemon_identity(&actual, expected) {
108 return Err(identity_mismatch(expected, &actual));
109 }
110 Ok(())
111}
112
113#[derive(Debug, Clone, PartialEq, Eq)]
115pub struct EndpointProbeRequest {
116 pub request_id: u64,
118 pub nonce: [u8; PROBE_NONCE_BYTES],
120 pub traceparent: String,
122 pub tracestate: String,
124}
125
126pub fn read_endpoint_probe_request<S: Read>(
128 stream: &mut S,
129) -> Result<EndpointProbeRequest, EndpointProbeServerError> {
130 let request_bytes = read_frame(stream)?;
131 let frame =
132 Frame::decode(request_bytes.as_slice()).map_err(EndpointProbeServerError::DecodeFrame)?;
133 endpoint_probe_request_from_frame(&frame)
134}
135
136pub fn endpoint_probe_request_from_frame(
142 frame: &Frame,
143) -> Result<EndpointProbeRequest, EndpointProbeServerError> {
144 validate_endpoint_probe_request_frame(frame)?;
145 let nonce = frame
146 .payload
147 .as_slice()
148 .try_into()
149 .map_err(|_| EndpointProbeServerError::MalformedPayload("nonce must be 32 bytes"))?;
150 Ok(EndpointProbeRequest {
151 request_id: frame.request_id,
152 nonce,
153 traceparent: frame.traceparent.clone(),
154 tracestate: frame.tracestate.clone(),
155 })
156}
157
158pub fn write_endpoint_probe_response<S: Write>(
160 stream: &mut S,
161 request: &EndpointProbeRequest,
162 daemon: &DaemonProcess,
163) -> Result<(), EndpointProbeServerError> {
164 let response_frame = endpoint_probe_response_frame(request, daemon);
165 let mut response_bytes = Vec::new();
166 response_frame
167 .encode(&mut response_bytes)
168 .map_err(EndpointProbeServerError::EncodeFrame)?;
169 write_frame(stream, &response_bytes)?;
170 Ok(())
171}
172
173pub fn handle_endpoint_probe<S: Read + Write>(
175 stream: &mut S,
176 daemon: &DaemonProcess,
177) -> Result<(), EndpointProbeServerError> {
178 let request = read_endpoint_probe_request(stream)?;
179 write_endpoint_probe_response(stream, &request, daemon)
180}
181
182#[derive(Debug, thiserror::Error)]
184pub enum ProbeError {
185 #[error("endpoint does not match expected daemon identity")]
187 EndpointMismatch,
188 #[error(transparent)]
190 EndpointResponse(#[from] EndpointProbeError),
191 #[error(transparent)]
193 VerifyPid(#[from] VerifyPidError),
194}
195
196#[derive(Debug, thiserror::Error)]
198pub enum EndpointProbeError {
199 #[error("backend endpoint probe random generation failed: {0}")]
201 Random(getrandom::Error),
202 #[error("backend endpoint probe local-socket name failed: {0}")]
204 LocalSocketName(io::Error),
205 #[error("backend endpoint probe connect failed: {0}")]
207 Connect(io::Error),
208 #[error("backend endpoint probe nonblocking setup failed: {0}")]
210 ConfigureNonblocking(io::Error),
211 #[error("backend endpoint probe timed out")]
213 Timeout,
214 #[error("backend endpoint probe I/O failed: {0}")]
216 Io(io::Error),
217 #[error("backend endpoint probe unsupported framing version: got {got}, expected {expected}")]
219 UnsupportedFramingVersion {
220 got: u8,
222 expected: u8,
224 },
225 #[error("backend endpoint probe frame body too large: {body_length} bytes exceeds cap {cap}")]
227 FrameTooLarge {
228 body_length: usize,
230 cap: usize,
232 },
233 #[error("failed to encode endpoint probe frame: {0}")]
235 EncodeFrame(prost::EncodeError),
236 #[error("failed to decode endpoint probe response Frame: {0}")]
238 DecodeFrame(prost::DecodeError),
239 #[error("unexpected endpoint probe response: {0}")]
241 UnexpectedFrame(&'static str),
242 #[error("endpoint probe response payload is malformed: {0}")]
244 MalformedPayload(&'static str),
245 #[error("failed to decode endpoint probe daemon identity: {0}")]
247 DecodeDaemonProcess(prost::DecodeError),
248 #[error(transparent)]
250 Identity(#[from] IdentityError),
251 #[error("endpoint probe response identity did not match expected daemon identity: {field}")]
253 IdentityMismatch {
254 field: &'static str,
256 },
257}
258
259#[derive(Debug, thiserror::Error)]
261pub enum EndpointProbeServerError {
262 #[error(transparent)]
264 Framing(#[from] FramingError),
265 #[error("failed to decode endpoint probe request Frame: {0}")]
267 DecodeFrame(prost::DecodeError),
268 #[error("failed to encode endpoint probe response Frame: {0}")]
270 EncodeFrame(prost::EncodeError),
271 #[error("unexpected endpoint probe request: {0}")]
273 UnexpectedFrame(&'static str),
274 #[error("endpoint probe request payload is malformed: {0}")]
276 MalformedPayload(&'static str),
277}
278
279fn endpoint_probe_request_frame(request_id: u64, nonce: &[u8; PROBE_NONCE_BYTES]) -> Frame {
280 Frame {
281 envelope_version: PROTOCOL_VERSION,
282 kind: FrameKind::Request as i32,
283 payload_protocol: BACKEND_HANDLE_PROBE_PAYLOAD_PROTOCOL,
284 payload: nonce.to_vec(),
285 request_id,
286 payload_encoding: PayloadEncoding::None as i32,
287 deadline_unix_ms: 0,
288 traceparent: String::new(),
289 tracestate: String::new(),
290 }
291}
292
293pub fn endpoint_probe_response_frame(
299 request: &EndpointProbeRequest,
300 daemon: &DaemonProcess,
301) -> Frame {
302 let mut payload = Vec::with_capacity(PROBE_NONCE_BYTES + 128);
303 payload.extend_from_slice(&request.nonce);
304 daemon.encode_probe_identity(&mut payload).expect(
305 "prost encoding DaemonProcess into Vec cannot fail because Vec writes are infallible",
306 );
307
308 Frame {
309 envelope_version: PROTOCOL_VERSION,
310 kind: FrameKind::Response as i32,
311 payload_protocol: BACKEND_HANDLE_PROBE_PAYLOAD_PROTOCOL,
312 payload,
313 request_id: request.request_id,
314 payload_encoding: PayloadEncoding::None as i32,
315 deadline_unix_ms: 0,
316 traceparent: request.traceparent.clone(),
317 tracestate: request.tracestate.clone(),
318 }
319}
320
321pub fn validate_endpoint_probe_request_frame(
325 frame: &Frame,
326) -> Result<(), EndpointProbeServerError> {
327 if frame.envelope_version != PROTOCOL_VERSION {
328 return Err(EndpointProbeServerError::UnexpectedFrame(
329 "envelope_version is not v1",
330 ));
331 }
332 if FrameKind::try_from(frame.kind) != Ok(FrameKind::Request) {
333 return Err(EndpointProbeServerError::UnexpectedFrame(
334 "kind is not REQUEST",
335 ));
336 }
337 if frame.payload_protocol != BACKEND_HANDLE_PROBE_PAYLOAD_PROTOCOL {
338 return Err(EndpointProbeServerError::UnexpectedFrame(
339 "payload_protocol is not endpoint probe",
340 ));
341 }
342 if PayloadEncoding::try_from(frame.payload_encoding) != Ok(PayloadEncoding::None) {
343 return Err(EndpointProbeServerError::UnexpectedFrame(
344 "payload is compressed",
345 ));
346 }
347 if frame.payload.len() != PROBE_NONCE_BYTES {
348 return Err(EndpointProbeServerError::MalformedPayload(
349 "nonce must be 32 bytes",
350 ));
351 }
352 Ok(())
353}
354
355fn validate_endpoint_probe_response_frame(
356 frame: &Frame,
357 request_id: u64,
358) -> Result<(), EndpointProbeError> {
359 if frame.envelope_version != PROTOCOL_VERSION {
360 return Err(EndpointProbeError::UnexpectedFrame(
361 "envelope_version is not v1",
362 ));
363 }
364 if FrameKind::try_from(frame.kind) != Ok(FrameKind::Response) {
365 return Err(EndpointProbeError::UnexpectedFrame("kind is not RESPONSE"));
366 }
367 if frame.payload_protocol != BACKEND_HANDLE_PROBE_PAYLOAD_PROTOCOL {
368 return Err(EndpointProbeError::UnexpectedFrame(
369 "payload_protocol is not endpoint probe",
370 ));
371 }
372 if frame.request_id != request_id {
373 return Err(EndpointProbeError::UnexpectedFrame(
374 "request_id does not match endpoint probe request",
375 ));
376 }
377 if PayloadEncoding::try_from(frame.payload_encoding) != Ok(PayloadEncoding::None) {
378 return Err(EndpointProbeError::UnexpectedFrame("payload is compressed"));
379 }
380 Ok(())
381}
382
383pub fn decode_response_identity(
391 payload: &[u8],
392 expected_nonce: &[u8; PROBE_NONCE_BYTES],
393) -> Result<DaemonProcess, EndpointProbeError> {
394 if payload.len() < PROBE_NONCE_BYTES {
395 return Err(EndpointProbeError::MalformedPayload(
396 "payload shorter than nonce",
397 ));
398 }
399 let (nonce, identity_bytes) = payload.split_at(PROBE_NONCE_BYTES);
400 if nonce != expected_nonce {
401 return Err(EndpointProbeError::UnexpectedFrame(
402 "nonce does not match endpoint probe request",
403 ));
404 }
405 let proto_identity = protocol::DaemonProcess::decode(identity_bytes)
406 .map_err(EndpointProbeError::DecodeDaemonProcess)?;
407 DaemonProcess::try_from(proto_identity).map_err(EndpointProbeError::Identity)
408}
409
410fn identity_mismatch(expected: &DaemonProcess, actual: &DaemonProcess) -> EndpointProbeError {
411 let field = if actual.pid != expected.pid {
412 "pid"
413 } else if actual.exe_path != expected.exe_path {
414 "exe_path"
415 } else if actual.exe_hash != expected.exe_hash {
416 "exe_hash"
417 } else if actual.boot_id != expected.boot_id {
418 "boot_id"
419 } else if !same_endpoint(&actual.ipc_endpoint, &expected.ipc_endpoint) {
420 "ipc_endpoint"
421 } else {
422 "unknown"
423 };
424 EndpointProbeError::IdentityMismatch { field }
425}
426
427fn same_daemon_identity(left: &DaemonProcess, right: &DaemonProcess) -> bool {
428 left.pid == right.pid
429 && left.exe_path == right.exe_path
430 && left.exe_hash == right.exe_hash
431 && left.boot_id == right.boot_id
432 && same_endpoint(&left.ipc_endpoint, &right.ipc_endpoint)
433}
434
435fn connect_endpoint_with_deadline(
446 endpoint: &Endpoint,
447 deadline: Instant,
448) -> Result<crate::platform::ipc::Stream, EndpointProbeError> {
449 if endpoint.path.is_empty() {
450 return Err(EndpointProbeError::Connect(io::Error::new(
451 io::ErrorKind::InvalidInput,
452 "backend endpoint path is empty",
453 )));
454 }
455 let endpoint = crate::platform::ipc::Endpoint::new(endpoint.path.clone())
459 .map_err(EndpointProbeError::LocalSocketName)?;
460
461 let dial_endpoint = endpoint.clone();
462 let (tx, rx) = std::sync::mpsc::channel();
463 thread::Builder::new()
464 .name("rp-endpoint-probe-connect".to_string())
465 .spawn(move || {
466 let _ = tx.send(crate::platform::ipc::Stream::connect(&dial_endpoint));
468 })
469 .map_err(EndpointProbeError::Connect)?;
470
471 let remaining = deadline.saturating_duration_since(Instant::now());
472 match rx.recv_timeout(remaining) {
473 Ok(Ok(stream)) => Ok(stream),
474 Ok(Err(err)) => Err(EndpointProbeError::Connect(err)),
475 Err(_) => Err(EndpointProbeError::Connect(io::Error::new(
476 io::ErrorKind::TimedOut,
477 format!(
478 "backend endpoint probe connect timed out after the probe deadline \
479 (endpoint {}): the listener exists but never completed the connection",
480 endpoint.display()
481 ),
482 ))),
483 }
484}
485
486fn write_probe_frame_with_deadline(
487 stream: &mut crate::platform::ipc::Stream,
488 body: &[u8],
489 deadline: Instant,
490) -> Result<(), EndpointProbeError> {
491 if body.len() > MAX_FRAME_BYTES {
492 return Err(EndpointProbeError::FrameTooLarge {
493 body_length: body.len(),
494 cap: MAX_FRAME_BYTES,
495 });
496 }
497 let mut wire = Vec::with_capacity(1 + 4 + body.len());
498 wire.push(ENVELOPE_VERSION);
499 wire.extend_from_slice(&(body.len() as u32).to_le_bytes());
500 wire.extend_from_slice(body);
501 write_all_with_deadline(stream, &wire, deadline)?;
502 flush_with_deadline(stream, deadline)
503}
504
505fn read_probe_frame_with_deadline(
506 stream: &mut crate::platform::ipc::Stream,
507 deadline: Instant,
508) -> Result<Vec<u8>, EndpointProbeError> {
509 parse_probe_frame(|buf| read_exact_with_deadline(stream, buf, deadline))
510}
511
512pub fn read_probe_frame<R: Read>(reader: &mut R) -> Result<Vec<u8>, EndpointProbeError> {
520 parse_probe_frame(|buf| reader.read_exact(buf).map_err(EndpointProbeError::Io))
521}
522
523fn parse_probe_frame(
528 mut read_exact: impl FnMut(&mut [u8]) -> Result<(), EndpointProbeError>,
529) -> Result<Vec<u8>, EndpointProbeError> {
530 let mut version = [0_u8; 1];
531 read_exact(&mut version)?;
532 if version[0] != ENVELOPE_VERSION {
533 return Err(EndpointProbeError::UnsupportedFramingVersion {
534 got: version[0],
535 expected: ENVELOPE_VERSION,
536 });
537 }
538
539 let mut len = [0_u8; 4];
540 read_exact(&mut len)?;
541 let body_length = u32::from_le_bytes(len) as usize;
542 if body_length > MAX_FRAME_BYTES {
543 return Err(EndpointProbeError::FrameTooLarge {
544 body_length,
545 cap: MAX_FRAME_BYTES,
546 });
547 }
548
549 let mut body = vec![0_u8; body_length];
550 if body_length > 0 {
551 read_exact(&mut body)?;
552 }
553 Ok(body)
554}
555
556fn write_all_with_deadline<W: Write>(
557 writer: &mut W,
558 mut buf: &[u8],
559 deadline: Instant,
560) -> Result<(), EndpointProbeError> {
561 while !buf.is_empty() {
562 match writer.write(buf) {
563 Ok(0) => {
564 return Err(EndpointProbeError::Io(io::Error::new(
565 io::ErrorKind::WriteZero,
566 "endpoint probe write returned zero bytes",
567 )));
568 }
569 Ok(written) => buf = &buf[written..],
570 Err(err) if err.kind() == io::ErrorKind::WouldBlock => wait_for_io(deadline)?,
571 Err(err) => return Err(EndpointProbeError::Io(err)),
572 }
573 }
574 Ok(())
575}
576
577fn read_exact_with_deadline<R: Read>(
578 reader: &mut R,
579 mut buf: &mut [u8],
580 deadline: Instant,
581) -> Result<(), EndpointProbeError> {
582 while !buf.is_empty() {
583 match reader.read(buf) {
584 Ok(0) => wait_for_io(deadline)?,
585 Ok(read) => {
586 let tmp = buf;
587 buf = &mut tmp[read..];
588 }
589 Err(err) if err.kind() == io::ErrorKind::WouldBlock => wait_for_io(deadline)?,
590 Err(err) => return Err(EndpointProbeError::Io(err)),
591 }
592 }
593 Ok(())
594}
595
596fn flush_with_deadline<W: Write>(
597 writer: &mut W,
598 deadline: Instant,
599) -> Result<(), EndpointProbeError> {
600 loop {
601 match writer.flush() {
602 Ok(()) => return Ok(()),
603 Err(err) if err.kind() == io::ErrorKind::WouldBlock => wait_for_io(deadline)?,
604 Err(err) => return Err(EndpointProbeError::Io(err)),
605 }
606 }
607}
608
609fn wait_for_io(deadline: Instant) -> Result<(), EndpointProbeError> {
610 if Instant::now() >= deadline {
611 return Err(EndpointProbeError::Timeout);
612 }
613 let remaining = deadline.saturating_duration_since(Instant::now());
614 thread::sleep(remaining.min(NONBLOCKING_POLL_INTERVAL));
615 Ok(())
616}