1use std::io::{ErrorKind, Read, Write};
15#[cfg(unix)]
16use std::os::unix::net::UnixStream;
17use std::sync::Arc;
18use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
19
20use serde_json::Value;
21
22use crate::debug::{describe_endpoint, error_label, join_capabilities, on_off, Category, DebugLog};
23use crate::error::Error;
24use crate::evidence::{Lease as EvidenceProviderLease, Registry as EvidenceProviderRegistry};
25use crate::framing::{encode_frame, FrameDecoder};
26use crate::limits::{Limits, DEFAULT_LIMITS};
27use crate::logs::{AttrValue, LogLevel, LogRecord, MAX_LOG_ATTRS};
28use crate::marker::encode_marker;
29use crate::messages::{
30 default_capabilities, parse_driver_message, Hello, HelloAck, LogMessage, ProbeInfo,
31 ProtocolErrorMessage, RevisionCommit, SnapshotMessage,
32};
33use crate::roles::Capability;
34use crate::tree::Snapshot;
35use crate::validate::validate_snapshot;
36
37#[cfg(unix)]
38type TransportStream = UnixStream;
39
40#[cfg(windows)]
41use interprocess::{
42 os::windows::named_pipe::{pipe_mode, DuplexPipeStream},
43 ConnectWaitMode,
44};
45#[cfg(windows)]
46type TransportStream = DuplexPipeStream<pipe_mode::Bytes>;
47
48pub const ENV_ENDPOINT: &str = "TERMWRIGHT_ENDPOINT";
50pub const ENV_TOKEN: &str = "TERMWRIGHT_TOKEN";
52pub const DIAL_TIMEOUT: Duration = Duration::from_secs(5);
54
55pub const WRITE_TIMEOUT: Duration = Duration::from_millis(250);
63
64#[derive(Debug, Clone)]
66pub struct Options {
67 pub adapter_name: String,
69 pub adapter_version: String,
71 pub capabilities: Vec<Capability>,
73 pub limits: Limits,
75 pub write_timeout: Option<Duration>,
78 pub probe: Option<ProbeInfo>,
81 pub debug: Option<Arc<DebugLog>>,
86 pub evidence_registry: Option<EvidenceProviderRegistry>,
88}
89
90impl Options {
91 pub fn with_logs(adapter_name: impl Into<String>, adapter_version: impl Into<String>) -> Self {
96 let mut options = Self::new(adapter_name, adapter_version);
97 options.capabilities.push(Capability::Logs);
98 options
99 }
100
101 pub fn new(adapter_name: impl Into<String>, adapter_version: impl Into<String>) -> Self {
103 Self {
104 adapter_name: adapter_name.into(),
105 adapter_version: adapter_version.into(),
106 capabilities: default_capabilities(),
107 limits: DEFAULT_LIMITS,
108 write_timeout: Some(WRITE_TIMEOUT),
109 probe: None,
110 debug: None,
114 evidence_registry: None,
115 }
116 }
117}
118
119fn epoch_millis() -> i64 {
122 SystemTime::now()
123 .duration_since(UNIX_EPOCH)
124 .map(|since| since.as_millis() as i64)
125 .unwrap_or(0)
126}
127
128#[derive(Debug)]
134struct TokenBucket {
135 per_second: f64,
136 capacity: f64,
137 tokens: f64,
138 updated: Instant,
139}
140
141impl TokenBucket {
142 fn new(per_second: i64, burst: i64, now: Instant) -> Self {
143 let rate = per_second.max(0) as f64;
144 let capacity = rate + burst.max(0) as f64;
145 Self {
146 per_second: rate,
147 capacity,
148 tokens: capacity,
149 updated: now,
150 }
151 }
152
153 fn take(&mut self, now: Instant) -> bool {
155 if self.per_second <= 0.0 {
156 return false;
157 }
158 let elapsed = now.saturating_duration_since(self.updated).as_secs_f64();
159 self.updated = now;
160 self.tokens = (self.tokens + elapsed * self.per_second).min(self.capacity);
161 if self.tokens < 1.0 {
162 return false;
163 }
164 self.tokens -= 1.0;
165 true
166 }
167}
168
169#[derive(Debug)]
173pub struct Client {
174 endpoint: String,
175 token: String,
176 options: Options,
177 stream: Option<TransportStream>,
178 decoder: FrameDecoder,
179 limits: Limits,
180 session_id: Option<String>,
181 revision: i64,
182 marker_enabled: bool,
183 log_budget: Option<crate::messages::LogBudget>,
184 snapshots_sent: u64,
185 log_seq: i64,
186 log_bucket: Option<TokenBucket>,
187 logs_dropped: u64,
188 subscribe: String,
189 evidence_lease: Option<EvidenceProviderLease>,
190}
191
192impl Client {
193 pub fn new(endpoint: impl Into<String>, token: impl Into<String>, options: Options) -> Self {
195 let limits = options.limits;
196 Self {
197 endpoint: endpoint.into(),
198 token: token.into(),
199 options,
200 stream: None,
201 decoder: FrameDecoder::new(limits.max_frame_bytes, limits.max_depth),
202 limits,
203 session_id: None,
204 revision: 0,
205 marker_enabled: false,
206 log_budget: None,
207 snapshots_sent: 0,
208 log_seq: 0,
209 log_bucket: None,
210 logs_dropped: 0,
211 subscribe: "snapshots".to_owned(),
212 evidence_lease: None,
213 }
214 }
215
216 pub fn from_env(mut options: Options) -> Option<Self> {
221 if options.debug.is_none() {
222 options.debug = DebugLog::from_env(&options.adapter_name).map(Arc::new);
223 }
224 Self::from_values(
225 std::env::var(ENV_ENDPOINT).ok().as_deref(),
226 std::env::var(ENV_TOKEN).ok().as_deref(),
227 options,
228 )
229 }
230
231 pub fn from_values(
237 endpoint: Option<&str>,
238 token: Option<&str>,
239 options: Options,
240 ) -> Option<Self> {
241 let endpoint = endpoint.filter(|value| !value.is_empty());
242 let token = token.filter(|value| !value.is_empty());
243 let (Some(endpoint), Some(token)) = (endpoint, token) else {
244 if let Some(log) = options.debug.as_ref() {
245 let mut missing = Vec::new();
246 if endpoint.is_none() {
247 missing.push(ENV_ENDPOINT);
248 }
249 if token.is_none() {
250 missing.push(ENV_TOKEN);
251 }
252 log.line(
253 Category::Diag,
254 &format!("dormant: {} not set", missing.join(" and ")),
255 );
256 }
257 return None;
258 };
259 if !endpoint_supported(endpoint) {
260 if let Some(log) = options.debug.as_ref() {
261 log.line(
262 Category::Diag,
263 &format!(
264 "dormant: {} is not a local endpoint for this platform",
265 describe_endpoint(endpoint)
266 ),
267 );
268 }
269 return None;
270 }
271 Some(Self::new(endpoint, token, options))
272 }
273
274 pub fn connect(&mut self, timeout: Duration) -> Result<(), Error> {
282 if let Some(probe) = self.options.probe.as_ref() {
283 probe.validate()?;
284 }
285 self.debug_line(
286 Category::Sem,
287 &format!(
288 "dial {} timeout={}ms",
289 describe_endpoint(&self.endpoint),
290 timeout.as_millis()
291 ),
292 );
293 let stream = match connect_transport(&self.endpoint, timeout, self.options.write_timeout) {
294 Ok(stream) => stream,
295 Err(error) => {
296 self.debug_line(
297 Category::Diag,
298 &format!("dial failed, staying dormant: {}", error_label(&error)),
299 );
300 return Err(error.into());
301 }
302 };
303 self.stream = Some(stream);
304
305 let mut hello = Hello::new(
306 &self.token,
307 &self.options.adapter_name,
308 &self.options.adapter_version,
309 self.options.capabilities.clone(),
310 );
311 if let Some(probe) = self.options.probe.clone() {
312 hello = hello.with_probe(probe);
313 }
314 if let Some(registry) = self.options.evidence_registry.as_ref() {
315 let lease = registry.freeze();
316 hello = hello.with_providers(lease.registrations());
317 self.evidence_lease = Some(lease);
318 }
319 self.send(&hello)?;
320 self.debug_line(
321 Category::Sem,
322 &format!(
323 "hello sent adapter={}/{} caps={}",
324 self.options.adapter_name,
325 self.options.adapter_version,
326 join_capabilities(&self.options.capabilities)
327 ),
328 );
329
330 let deadline = Instant::now() + timeout;
331 while self.session_id.is_none() {
332 if Instant::now() >= deadline {
333 self.debug_line(
334 Category::Diag,
335 &format!(
336 "no hello-ack within {}ms, staying dormant",
337 timeout.as_millis()
338 ),
339 );
340 self.close();
341 return Err(Error::HandshakeTimeout);
342 }
343 self.poll()?;
344 std::thread::yield_now();
345 }
346 Ok(())
347 }
348
349 fn debug_line(&self, category: Category, message: &str) {
356 if let Some(log) = self.options.debug.as_ref() {
357 log.line(category, message);
358 }
359 }
360
361 pub fn connected(&self) -> bool {
363 self.session_id.is_some() && self.stream.is_some()
364 }
365
366 pub fn session_id(&self) -> Option<&str> {
368 self.session_id.as_deref()
369 }
370
371 pub fn revision(&self) -> i64 {
373 self.revision
374 }
375
376 pub fn log_budget(&self) -> Option<crate::messages::LogBudget> {
379 self.log_budget
380 }
381
382 pub fn limits(&self) -> &Limits {
384 &self.limits
385 }
386
387 pub fn close(&mut self) {
389 if let Some(stream) = self.stream.take() {
390 self.debug_line(
391 Category::Sem,
392 &format!(
393 "close r{} snapshots={} logs_dropped={}",
394 self.revision, self.snapshots_sent, self.logs_dropped
395 ),
396 );
397 close_transport(stream);
398 }
399 self.session_id = None;
400 if let Some(mut lease) = self.evidence_lease.take() {
401 lease.close();
402 }
403 }
404
405 pub fn fail(&mut self, code: &str, message: impl Into<String>) -> Result<(), Error> {
407 let result = self.send(&ProtocolErrorMessage::new(code, message));
408 self.close();
409 result
410 }
411
412 pub fn publish(&mut self, snapshot: &mut Snapshot) -> Result<Option<String>, Error> {
426 self.publish_inner(snapshot)
427 }
428
429 fn publish_inner(&mut self, snapshot: &mut Snapshot) -> Result<Option<String>, Error> {
430 let Some(session_id) = self.session_id.clone() else {
431 return Ok(None);
432 };
433 if self.stream.is_none() {
434 return Ok(None);
435 }
436
437 let revision = self.revision + 1;
438 snapshot.v = 2;
439 snapshot.session_id = session_id.clone();
440 snapshot.revision = revision;
441 if let Some(lease) = self.evidence_lease.as_ref() {
442 snapshot.provider_evidence =
443 lease.collect(&session_id, revision, snapshot.columns, snapshot.rows);
444 }
445
446 let body = serde_json::to_string(&snapshot).map_err(|_| {
447 Error::Protocol(crate::error::Violation::new(
448 "frame-malformed",
449 "snapshot is not JSON-serialisable",
450 ))
451 })?;
452 let parsed: Value = serde_json::from_str(&body).expect("just serialised");
453 validate_snapshot(&parsed, &self.limits)?;
454
455 let marker = if self.marker_enabled {
456 Some(encode_marker(&self.token, &session_id, revision)?)
457 } else {
458 None
459 };
460
461 let tree_frame = if self.subscribe != "revisions" {
465 Some(encode_frame(
466 &SnapshotMessage::new(snapshot),
467 self.limits.max_frame_bytes,
468 )?)
469 } else {
470 None
471 };
472 let commit_frame =
473 encode_frame(&RevisionCommit::new(revision), self.limits.max_frame_bytes)?;
474
475 if let Some(frame) = &tree_frame {
476 self.write_frame(frame)?;
477 }
478 self.write_frame(&commit_frame)?;
479
480 self.revision = revision;
482 if tree_frame.is_some() {
483 self.snapshots_sent += 1;
484 }
485
486 Ok(marker)
487 }
488
489 pub fn snapshots_sent(&self) -> u64 {
491 self.snapshots_sent
492 }
493
494 pub fn logs_dropped(&self) -> u64 {
497 self.logs_dropped
498 }
499
500 pub fn log(&mut self, mut record: LogRecord) -> bool {
517 if self.session_id.is_none() || self.stream.is_none() || self.log_bucket.is_none() {
518 return false;
519 }
520
521 let origin = record.seq;
522 self.log_seq += 1;
523 record.seq = self.log_seq;
524 if record.ts == 0 {
525 record.ts = epoch_millis();
526 }
527 if record.revision.is_none() && self.revision > 0 {
528 record.revision = Some(self.revision);
529 }
530
531 let now = Instant::now();
532 let allowed = self
533 .log_bucket
534 .as_mut()
535 .is_some_and(|bucket| bucket.take(now));
536 if !allowed {
537 self.logs_dropped += 1;
538 return false;
539 }
540 if origin > 0 && record.attrs.len() < MAX_LOG_ATTRS {
541 record
544 .attrs
545 .insert("origin.seq".to_owned(), AttrValue::Int(origin));
546 if record.validate(&self.limits).is_err() {
547 record.attrs.remove("origin.seq");
548 }
549 }
550 if record.validate(&self.limits).is_err() {
551 self.logs_dropped += 1;
554 return false;
555 }
556 self.send(&LogMessage::new(&record)).is_ok()
557 }
558
559 pub fn log_message(&mut self, level: LogLevel, message: impl Into<String>) -> bool {
561 self.log(LogRecord::new(level, message))
562 }
563
564 pub fn poll(&mut self) -> Result<(), Error> {
573 let mut buffer = [0u8; 8192];
574 loop {
575 let read = match self.stream.as_mut() {
576 None => return Ok(()),
577 Some(stream) => read_transport(stream, &mut buffer),
578 };
579 match read {
580 Ok(Incoming::Closed) => {
581 self.close();
582 return Ok(());
583 }
584 Ok(Incoming::Data(count)) => {
585 let frames = self.decoder.push(&buffer[..count])?;
586 for frame in frames {
587 self.handle(&frame.value)?;
588 }
589 }
590 Ok(Incoming::Idle) => return Ok(()),
591 Err(error) if error.kind() == ErrorKind::Interrupted => continue,
592 Err(error) => {
593 self.close();
594 return Err(Error::Io(error));
595 }
596 }
597 }
598 }
599
600 fn handle(&mut self, value: &Value) -> Result<(), Error> {
601 if let Err(error) = parse_driver_message(value, &self.limits) {
602 self.debug_line(
603 Category::Diag,
604 &format!("rejected a driver message: {error}"),
605 );
606 let _ = self.send(&ProtocolErrorMessage::new("malformed", error.to_string()));
607 self.close();
608 return Err(Error::Parse(error));
609 }
610
611 match value.get("type").and_then(Value::as_str) {
612 Some("hello-ack") => {
613 let ack: HelloAck = serde_json::from_value(value.clone()).expect("validated above");
614 self.session_id = Some(ack.session_id);
615 self.limits = ack.limits;
616 self.marker_enabled = ack.marker.enabled;
617 self.log_budget = ack.logs;
618 self.log_bucket = match ack.logs {
619 Some(budget) if budget.enabled => Some(TokenBucket::new(
620 budget.max_records_per_second,
621 budget.burst,
622 Instant::now(),
623 )),
624 _ => None,
625 };
626 self.subscribe = ack.subscribe;
627 if let Some(log) = self.options.debug.as_ref() {
628 let session = self.session_id.clone().unwrap_or_default();
629 log.set_label(&session);
630 log.line(
631 Category::Sem,
632 &format!(
633 "hello-ack session={session} marker={} subscribe={} logs={}",
634 on_off(self.marker_enabled),
635 self.subscribe,
636 on_off(self.log_bucket.is_some())
637 ),
638 );
639 }
640 }
641 Some("error") => {
642 self.debug_line(
643 Category::Diag,
644 &format!(
645 "driver ended the session: {}",
646 value.get("code").and_then(Value::as_str).unwrap_or("?")
647 ),
648 );
649 self.close();
650 }
651 _ => {}
652 }
653 Ok(())
654 }
655
656 fn send<T: serde::Serialize>(&mut self, message: &T) -> Result<(), Error> {
657 let frame = encode_frame(message, self.limits.max_frame_bytes)?;
658 self.write_frame(&frame)
659 }
660
661 pub(crate) fn write_frame(&mut self, frame: &[u8]) -> Result<(), Error> {
662 let Some(stream) = self.stream.as_mut() else {
663 return Ok(());
664 };
665 match write_transport_frame(stream, frame, self.options.write_timeout) {
666 Ok(()) => Ok(()),
667 Err(error) => {
668 let timed_out = matches!(
669 error.kind(),
670 ErrorKind::WouldBlock | ErrorKind::TimedOut | ErrorKind::Interrupted
671 );
672 self.close();
673 if timed_out {
674 self.debug_line(
678 Category::Diag,
679 "write deadline exceeded; session is unrecoverable",
680 );
681 return Err(Error::WriteTimeout);
682 }
683 Err(Error::Io(error))
684 }
685 }
686 }
687
688 pub(crate) fn accept_queued_publication(&mut self, revision: i64, snapshot_sent: bool) {
689 self.revision = revision;
690 if snapshot_sent {
691 self.snapshots_sent += 1;
692 }
693 }
694
695 pub(crate) fn take_evidence_lease(&mut self) -> Option<EvidenceProviderLease> {
696 self.evidence_lease.take()
697 }
698
699 pub(crate) fn publication_config(&self) -> Option<(String, String, Limits, String, bool, i64)> {
700 Some((
701 self.token.clone(),
702 self.session_id.clone()?,
703 self.limits,
704 self.subscribe.clone(),
705 self.marker_enabled,
706 self.revision,
707 ))
708 }
709
710 #[cfg(all(test, unix))]
711 pub(crate) fn test_connected(stream: TransportStream) -> Self {
712 let mut client = Self::new("unused", "test-token", Options::new("queue-test", "1"));
713 client.stream = Some(stream);
714 client.session_id = Some("test-session".into());
715 client.marker_enabled = true;
716 client
717 }
718}
719
720#[cfg(unix)]
721fn endpoint_supported(endpoint: &str) -> bool {
722 !endpoint.starts_with(r"\\.\pipe\") && !endpoint.starts_with(r"\\?\pipe\")
723}
724
725#[cfg(windows)]
726fn endpoint_supported(endpoint: &str) -> bool {
727 endpoint.starts_with(r"\\.\pipe\") || endpoint.starts_with(r"\\?\pipe\")
728}
729
730#[cfg(unix)]
731fn connect_transport(
732 endpoint: &str,
733 _dial_timeout: Duration,
734 write_timeout: Option<Duration>,
735) -> std::io::Result<TransportStream> {
736 let stream = UnixStream::connect(endpoint)?;
737 stream.set_read_timeout(Some(Duration::from_millis(50)))?;
738 stream.set_write_timeout(write_timeout)?;
739 Ok(stream)
740}
741
742#[cfg(windows)]
743fn connect_transport(
744 endpoint: &str,
745 dial_timeout: Duration,
746 _write_timeout: Option<Duration>,
747) -> std::io::Result<TransportStream> {
748 let stream = TransportStream::connect_by_path_with_wait_mode(
749 endpoint,
750 ConnectWaitMode::Timeout(dial_timeout),
751 )?;
752 stream.set_nonblocking(true)?;
756 Ok(stream)
757}
758
759enum Incoming {
761 Data(usize),
762 Idle,
764 Closed,
766}
767
768#[cfg(unix)]
769fn read_transport(stream: &mut TransportStream, buffer: &mut [u8]) -> std::io::Result<Incoming> {
770 match stream.read(buffer) {
771 Ok(0) => Ok(Incoming::Closed),
772 Ok(count) => Ok(Incoming::Data(count)),
773 Err(error) if matches!(error.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut) => {
774 Ok(Incoming::Idle)
775 }
776 Err(error) => Err(error),
777 }
778}
779
780#[cfg(windows)]
782const ERROR_NO_DATA: i32 = 232;
783#[cfg(windows)]
785const ERROR_BROKEN_PIPE: i32 = 109;
786
787#[cfg(windows)]
798fn read_transport(stream: &mut TransportStream, buffer: &mut [u8]) -> std::io::Result<Incoming> {
799 match stream.read(buffer) {
800 Ok(0) => Ok(Incoming::Idle),
801 Ok(count) => Ok(Incoming::Data(count)),
802 Err(error) if matches!(error.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut) => {
803 Ok(Incoming::Idle)
804 }
805 Err(error) if error.raw_os_error() == Some(ERROR_NO_DATA) => Ok(Incoming::Idle),
806 Err(error) if error.raw_os_error() == Some(ERROR_BROKEN_PIPE) => Ok(Incoming::Closed),
807 Err(error) => Err(error),
808 }
809}
810
811#[cfg(unix)]
812fn close_transport(stream: TransportStream) {
813 let _ = stream.shutdown(std::net::Shutdown::Both);
814}
815
816#[cfg(windows)]
817fn close_transport(_stream: TransportStream) {
818 }
821
822#[cfg(unix)]
823fn write_transport_frame(
824 stream: &mut TransportStream,
825 frame: &[u8],
826 _timeout: Option<Duration>,
827) -> std::io::Result<()> {
828 stream.write_all(frame).and_then(|()| stream.flush())
829}
830
831#[cfg(windows)]
832fn write_transport_frame(
833 stream: &mut TransportStream,
834 frame: &[u8],
835 timeout: Option<Duration>,
836) -> std::io::Result<()> {
837 let deadline = timeout.map(|duration| Instant::now() + duration);
838 let mut offset = 0;
839 while offset < frame.len() {
840 match stream.write(&frame[offset..]) {
841 Ok(0) => return Err(std::io::Error::from(ErrorKind::WriteZero)),
842 Ok(written) => offset += written,
843 Err(error) if error.kind() == ErrorKind::Interrupted => continue,
844 Err(error) if error.kind() == ErrorKind::WouldBlock => {
845 if deadline.is_some_and(|deadline| Instant::now() >= deadline) {
846 return Err(std::io::Error::from(ErrorKind::TimedOut));
847 }
848 std::thread::yield_now();
849 }
850 Err(error) => return Err(error),
851 }
852 }
853 loop {
854 match stream.flush() {
855 Ok(()) => return Ok(()),
856 Err(error) if error.kind() == ErrorKind::Interrupted => continue,
857 Err(error) if error.kind() == ErrorKind::WouldBlock => {
858 if deadline.is_some_and(|deadline| Instant::now() >= deadline) {
859 return Err(std::io::Error::from(ErrorKind::TimedOut));
860 }
861 std::thread::yield_now();
862 }
863 Err(error) => return Err(error),
864 }
865 }
866}