Skip to main content

running_process/broker/backend_lifecycle/
probe.rs

1//! Endpoint and process identity checks for backend handles.
2
3use 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
16/// Byte length of the random challenge carried by endpoint probe requests.
17pub const PROBE_NONCE_BYTES: usize = 32;
18const NONBLOCKING_POLL_INTERVAL: Duration = Duration::from_millis(5);
19
20/// Payload protocol reserved for `BackendHandle` endpoint identity probes.
21///
22/// Re-exported from the authoritative
23/// [`registry`](crate::broker::protocol::registry), which owns every v1
24/// payload-protocol ID (#375).
25pub use crate::broker::protocol::registry::BACKEND_HANDLE_PROBE_PAYLOAD_PROTOCOL;
26
27/// Default deadline for the active endpoint-response proof.
28pub const DEFAULT_ENDPOINT_PROBE_TIMEOUT: Duration = Duration::from_millis(500);
29
30/// Verify that an endpoint refers to the expected daemon process.
31pub 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
38/// [`probe_endpoint`] with a caller-chosen deadline for the response proof.
39///
40/// The default is a budget for a backend answering an identity probe, and it
41/// is the right budget for one running normally. A caller that knows its
42/// backend is running slower than normal for a reason unrelated to health --
43/// coverage instrumentation being the case this exists for (#1114) -- can say
44/// so here instead of the default being raised for everyone.
45///
46/// Only the response proof is bounded by `timeout`. The endpoint comparison
47/// and the process identity checks do no waiting.
48pub 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
62/// Compare two endpoint identities exactly.
63pub fn same_endpoint(left: &Endpoint, right: &Endpoint) -> bool {
64    left.namespace_id == right.namespace_id && left.path == right.path
65}
66
67/// Actively probe a backend endpoint and verify that it returns the expected
68/// daemon identity.
69///
70/// The probe uses the broker v1 frame layout with a dedicated payload protocol.
71/// Requests carry a 32-byte nonce. Responses must echo that nonce and include a
72/// prost-encoded `DaemonProcess` payload that exactly matches `expected`.
73pub 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
80/// Timed variant of [`probe_endpoint_response`] used by tests and diagnostics.
81pub 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/// Decoded endpoint probe request for backend-side responders.
114#[derive(Debug, Clone, PartialEq, Eq)]
115pub struct EndpointProbeRequest {
116    /// Request frame ID that the response must echo.
117    pub request_id: u64,
118    /// Random challenge that the response must echo.
119    pub nonce: [u8; PROBE_NONCE_BYTES],
120    /// Trace context copied from the request frame, if any.
121    pub traceparent: String,
122    /// Trace state copied from the request frame, if any.
123    pub tracestate: String,
124}
125
126/// Read and validate one endpoint probe request from an accepted IPC stream.
127pub 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
136/// Validate an already-decoded frame as an endpoint probe request.
137///
138/// Exposed for the [`crate::broker::backend_sdk`] endpoint mux (#412),
139/// which decodes frames from a consumer-owned read buffer instead of an
140/// exclusive stream.
141pub 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
158/// Write one endpoint probe response for a validated request.
159pub 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
173/// Serve exactly one endpoint probe request on an already-accepted IPC stream.
174pub 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/// Errors returned while probing a backend endpoint.
183#[derive(Debug, thiserror::Error)]
184pub enum ProbeError {
185    /// The caller-provided endpoint did not match the expected daemon endpoint.
186    #[error("endpoint does not match expected daemon identity")]
187    EndpointMismatch,
188    /// The endpoint did not answer the active identity probe as expected.
189    #[error(transparent)]
190    EndpointResponse(#[from] EndpointProbeError),
191    /// The daemon process identity could not be verified.
192    #[error(transparent)]
193    VerifyPid(#[from] VerifyPidError),
194}
195
196/// Errors returned by the active endpoint-response probe.
197#[derive(Debug, thiserror::Error)]
198pub enum EndpointProbeError {
199    /// The probe nonce could not be generated.
200    #[error("backend endpoint probe random generation failed: {0}")]
201    Random(getrandom::Error),
202    /// The endpoint path/name could not be converted to a local socket name.
203    #[error("backend endpoint probe local-socket name failed: {0}")]
204    LocalSocketName(io::Error),
205    /// Connecting to the endpoint failed.
206    #[error("backend endpoint probe connect failed: {0}")]
207    Connect(io::Error),
208    /// The stream could not be switched to nonblocking mode for deadline I/O.
209    #[error("backend endpoint probe nonblocking setup failed: {0}")]
210    ConfigureNonblocking(io::Error),
211    /// Probe I/O exceeded the configured deadline.
212    #[error("backend endpoint probe timed out")]
213    Timeout,
214    /// Raw probe I/O failed.
215    #[error("backend endpoint probe I/O failed: {0}")]
216    Io(io::Error),
217    /// The peer used the wrong broker framing byte.
218    #[error("backend endpoint probe unsupported framing version: got {got}, expected {expected}")]
219    UnsupportedFramingVersion {
220        /// Framing byte received from the peer.
221        got: u8,
222        /// Framing byte expected by v1.
223        expected: u8,
224    },
225    /// The peer advertised a frame that exceeds the v1 frame cap.
226    #[error("backend endpoint probe frame body too large: {body_length} bytes exceeds cap {cap}")]
227    FrameTooLarge {
228        /// Advertised frame body length.
229        body_length: usize,
230        /// Maximum accepted frame body length.
231        cap: usize,
232    },
233    /// The outbound probe request frame could not be encoded.
234    #[error("failed to encode endpoint probe frame: {0}")]
235    EncodeFrame(prost::EncodeError),
236    /// The response frame could not be decoded.
237    #[error("failed to decode endpoint probe response Frame: {0}")]
238    DecodeFrame(prost::DecodeError),
239    /// The response frame did not match the endpoint-probe contract.
240    #[error("unexpected endpoint probe response: {0}")]
241    UnexpectedFrame(&'static str),
242    /// The response payload did not match the endpoint-probe contract.
243    #[error("endpoint probe response payload is malformed: {0}")]
244    MalformedPayload(&'static str),
245    /// The response daemon identity could not be decoded.
246    #[error("failed to decode endpoint probe daemon identity: {0}")]
247    DecodeDaemonProcess(prost::DecodeError),
248    /// The response daemon identity was malformed.
249    #[error(transparent)]
250    Identity(#[from] IdentityError),
251    /// The response daemon identity did not match the expected identity.
252    #[error("endpoint probe response identity did not match expected daemon identity: {field}")]
253    IdentityMismatch {
254        /// First mismatched identity field.
255        field: &'static str,
256    },
257}
258
259/// Errors returned by backend-side endpoint probe responders.
260#[derive(Debug, thiserror::Error)]
261pub enum EndpointProbeServerError {
262    /// v1 framing failed.
263    #[error(transparent)]
264    Framing(#[from] FramingError),
265    /// The request frame could not be decoded.
266    #[error("failed to decode endpoint probe request Frame: {0}")]
267    DecodeFrame(prost::DecodeError),
268    /// The response frame could not be encoded.
269    #[error("failed to encode endpoint probe response Frame: {0}")]
270    EncodeFrame(prost::EncodeError),
271    /// The request frame did not match the endpoint-probe contract.
272    #[error("unexpected endpoint probe request: {0}")]
273    UnexpectedFrame(&'static str),
274    /// The request payload did not match the endpoint-probe contract.
275    #[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
293/// Build the endpoint-probe response frame for a validated request.
294///
295/// Exposed for the [`crate::broker::backend_sdk`] endpoint mux (#412),
296/// which answers probes from a consumer-owned read buffer instead of an
297/// exclusive stream.
298pub 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
321/// Validate one frame against the endpoint-probe request contract.
322///
323/// Exposed for the [`crate::broker::backend_sdk`] endpoint mux (#412).
324pub 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
383/// Decode the identity payload of one endpoint probe response.
384///
385/// The payload is untrusted (it comes from whatever process answered the
386/// probed endpoint — this is the squat-detection path): a 32-byte nonce echo
387/// followed by a prost-encoded [`protocol::DaemonProcess`]. The nonce must
388/// match `expected_nonce` before the identity bytes are decoded and
389/// normalized through [`DaemonProcess::try_from`]. Exposed for fuzzing.
390pub 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
435/// Connect to the probe endpoint with a hard deadline.
436///
437/// The facade's `Stream::connect` is a blocking syscall with no portable
438/// timeout: on macOS a bound-but-never-accepted Unix socket can
439/// park the caller in `connect(2)` indefinitely once the (tiny) listen
440/// backlog is full, which would silently wedge the broker serve thread
441/// before it ever binds its own control socket (#399). Run the blocking
442/// connect on a helper thread and bound the wait with the probe deadline;
443/// on timeout the helper thread owns (and eventually drops) the abandoned
444/// stream — the same leak-on-timeout pattern as the client handoff wait.
445fn 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    // Resolve the name synchronously so naming errors keep their own variant.
456    // The facade endpoint owns its path, so the helper thread can hold one
457    // outright rather than re-deriving a borrowed name from a cloned string.
458    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            // Receiver gone means the probe timed out; drop the stream here.
467            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
512/// Read one length-prefixed probe frame from an in-memory or blocking reader.
513///
514/// This drives the same byte-level parser as the nonblocking
515/// deadline-enforcing read used by [`probe_endpoint_response`]; it is exposed
516/// so fuzzing and tests can feed the framing logic from a
517/// [`std::io::Cursor`]. EOF surfaces as [`EndpointProbeError::Io`] instead of
518/// being retried against a deadline.
519pub 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
523/// Pure byte-level probe frame parse shared by the deadline-enforcing read
524/// and the fuzzing seam: a 1-byte envelope version ([`ENVELOPE_VERSION`]), a
525/// little-endian `u32` body length capped at [`MAX_FRAME_BYTES`], then the
526/// body bytes.
527fn 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}