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, SemanticFullMessage,
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: "semantic".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 = 3;
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 = Some(encode_frame(
465 &SemanticFullMessage::new(snapshot),
466 self.limits.max_frame_bytes,
467 )?);
468 let commit_frame =
469 encode_frame(&RevisionCommit::new(revision), self.limits.max_frame_bytes)?;
470
471 if let Some(frame) = &tree_frame {
472 self.write_frame(frame)?;
473 }
474 self.write_frame(&commit_frame)?;
475
476 self.revision = revision;
478 if tree_frame.is_some() {
479 self.snapshots_sent += 1;
480 }
481
482 Ok(marker)
483 }
484
485 pub fn snapshots_sent(&self) -> u64 {
487 self.snapshots_sent
488 }
489
490 pub fn logs_dropped(&self) -> u64 {
493 self.logs_dropped
494 }
495
496 pub fn log(&mut self, mut record: LogRecord) -> bool {
513 if self.session_id.is_none() || self.stream.is_none() || self.log_bucket.is_none() {
514 return false;
515 }
516
517 let origin = record.seq;
518 self.log_seq += 1;
519 record.seq = self.log_seq;
520 if record.ts == 0 {
521 record.ts = epoch_millis();
522 }
523 if record.revision.is_none() && self.revision > 0 {
524 record.revision = Some(self.revision);
525 }
526
527 let now = Instant::now();
528 let allowed = self
529 .log_bucket
530 .as_mut()
531 .is_some_and(|bucket| bucket.take(now));
532 if !allowed {
533 self.logs_dropped += 1;
534 return false;
535 }
536 if origin > 0 && record.attrs.len() < MAX_LOG_ATTRS {
537 record
540 .attrs
541 .insert("origin.seq".to_owned(), AttrValue::Int(origin));
542 if record.validate(&self.limits).is_err() {
543 record.attrs.remove("origin.seq");
544 }
545 }
546 if record.validate(&self.limits).is_err() {
547 self.logs_dropped += 1;
550 return false;
551 }
552 self.send(&LogMessage::new(&record)).is_ok()
553 }
554
555 pub fn log_message(&mut self, level: LogLevel, message: impl Into<String>) -> bool {
557 self.log(LogRecord::new(level, message))
558 }
559
560 pub fn poll(&mut self) -> Result<(), Error> {
569 let mut buffer = [0u8; 8192];
570 loop {
571 let read = match self.stream.as_mut() {
572 None => return Ok(()),
573 Some(stream) => read_transport(stream, &mut buffer),
574 };
575 match read {
576 Ok(Incoming::Closed) => {
577 self.close();
578 return Ok(());
579 }
580 Ok(Incoming::Data(count)) => {
581 let frames = self.decoder.push(&buffer[..count])?;
582 for frame in frames {
583 self.handle(&frame.value)?;
584 }
585 }
586 Ok(Incoming::Idle) => return Ok(()),
587 Err(error) if error.kind() == ErrorKind::Interrupted => continue,
588 Err(error) => {
589 self.close();
590 return Err(Error::Io(error));
591 }
592 }
593 }
594 }
595
596 fn handle(&mut self, value: &Value) -> Result<(), Error> {
597 if let Err(error) = parse_driver_message(value, &self.limits) {
598 self.debug_line(
599 Category::Diag,
600 &format!("rejected a driver message: {error}"),
601 );
602 let _ = self.send(&ProtocolErrorMessage::new("malformed", error.to_string()));
603 self.close();
604 return Err(Error::Parse(error));
605 }
606
607 match value.get("type").and_then(Value::as_str) {
608 Some("hello-ack") => {
609 let ack: HelloAck = serde_json::from_value(value.clone()).expect("validated above");
610 self.session_id = Some(ack.session_id);
611 self.limits = ack.limits;
612 self.marker_enabled = ack.marker.enabled;
613 self.log_budget = ack.logs;
614 self.log_bucket = match ack.logs {
615 Some(budget) if budget.enabled => Some(TokenBucket::new(
616 budget.max_records_per_second,
617 budget.burst,
618 Instant::now(),
619 )),
620 _ => None,
621 };
622 self.subscribe = ack.subscribe;
623 if let Some(log) = self.options.debug.as_ref() {
624 let session = self.session_id.clone().unwrap_or_default();
625 log.set_label(&session);
626 log.line(
627 Category::Sem,
628 &format!(
629 "hello-ack session={session} marker={} subscribe={} logs={}",
630 on_off(self.marker_enabled),
631 self.subscribe,
632 on_off(self.log_bucket.is_some())
633 ),
634 );
635 }
636 }
637 Some("error") => {
638 self.debug_line(
639 Category::Diag,
640 &format!(
641 "driver ended the session: {}",
642 value.get("code").and_then(Value::as_str).unwrap_or("?")
643 ),
644 );
645 self.close();
646 }
647 _ => {}
648 }
649 Ok(())
650 }
651
652 fn send<T: serde::Serialize>(&mut self, message: &T) -> Result<(), Error> {
653 let frame = encode_frame(message, self.limits.max_frame_bytes)?;
654 self.write_frame(&frame)
655 }
656
657 pub(crate) fn write_frame(&mut self, frame: &[u8]) -> Result<(), Error> {
658 let Some(stream) = self.stream.as_mut() else {
659 return Ok(());
660 };
661 match write_transport_frame(stream, frame, self.options.write_timeout) {
662 Ok(()) => Ok(()),
663 Err(error) => {
664 let timed_out = matches!(
665 error.kind(),
666 ErrorKind::WouldBlock | ErrorKind::TimedOut | ErrorKind::Interrupted
667 );
668 self.close();
669 if timed_out {
670 self.debug_line(
674 Category::Diag,
675 "write deadline exceeded; session is unrecoverable",
676 );
677 return Err(Error::WriteTimeout);
678 }
679 Err(Error::Io(error))
680 }
681 }
682 }
683
684 pub(crate) fn accept_queued_publication(&mut self, revision: i64, snapshot_sent: bool) {
685 self.revision = revision;
686 if snapshot_sent {
687 self.snapshots_sent += 1;
688 }
689 }
690
691 pub(crate) fn take_evidence_lease(&mut self) -> Option<EvidenceProviderLease> {
692 self.evidence_lease.take()
693 }
694
695 pub(crate) fn publication_config(&self) -> Option<(String, String, Limits, bool, i64)> {
696 Some((
697 self.token.clone(),
698 self.session_id.clone()?,
699 self.limits,
700 self.marker_enabled,
701 self.revision,
702 ))
703 }
704
705 #[cfg(all(test, unix))]
706 pub(crate) fn test_connected(stream: TransportStream) -> Self {
707 let mut client = Self::new("unused", "test-token", Options::new("queue-test", "1"));
708 client.stream = Some(stream);
709 client.session_id = Some("test-session".into());
710 client.marker_enabled = true;
711 client
712 }
713}
714
715#[cfg(unix)]
716fn endpoint_supported(endpoint: &str) -> bool {
717 !endpoint.starts_with(r"\\.\pipe\") && !endpoint.starts_with(r"\\?\pipe\")
718}
719
720#[cfg(windows)]
721fn endpoint_supported(endpoint: &str) -> bool {
722 endpoint.starts_with(r"\\.\pipe\") || endpoint.starts_with(r"\\?\pipe\")
723}
724
725#[cfg(unix)]
726fn connect_transport(
727 endpoint: &str,
728 _dial_timeout: Duration,
729 write_timeout: Option<Duration>,
730) -> std::io::Result<TransportStream> {
731 let stream = UnixStream::connect(endpoint)?;
732 stream.set_read_timeout(Some(Duration::from_millis(50)))?;
733 stream.set_write_timeout(write_timeout)?;
734 Ok(stream)
735}
736
737#[cfg(windows)]
738fn connect_transport(
739 endpoint: &str,
740 dial_timeout: Duration,
741 _write_timeout: Option<Duration>,
742) -> std::io::Result<TransportStream> {
743 let stream = TransportStream::connect_by_path_with_wait_mode(
744 endpoint,
745 ConnectWaitMode::Timeout(dial_timeout),
746 )?;
747 stream.set_nonblocking(true)?;
751 Ok(stream)
752}
753
754enum Incoming {
756 Data(usize),
757 Idle,
759 Closed,
761}
762
763#[cfg(unix)]
764fn read_transport(stream: &mut TransportStream, buffer: &mut [u8]) -> std::io::Result<Incoming> {
765 match stream.read(buffer) {
766 Ok(0) => Ok(Incoming::Closed),
767 Ok(count) => Ok(Incoming::Data(count)),
768 Err(error) if matches!(error.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut) => {
769 Ok(Incoming::Idle)
770 }
771 Err(error) => Err(error),
772 }
773}
774
775#[cfg(windows)]
777const ERROR_NO_DATA: i32 = 232;
778#[cfg(windows)]
780const ERROR_BROKEN_PIPE: i32 = 109;
781
782#[cfg(windows)]
793fn read_transport(stream: &mut TransportStream, buffer: &mut [u8]) -> std::io::Result<Incoming> {
794 match stream.read(buffer) {
795 Ok(0) => Ok(Incoming::Idle),
796 Ok(count) => Ok(Incoming::Data(count)),
797 Err(error) if matches!(error.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut) => {
798 Ok(Incoming::Idle)
799 }
800 Err(error) if error.raw_os_error() == Some(ERROR_NO_DATA) => Ok(Incoming::Idle),
801 Err(error) if error.raw_os_error() == Some(ERROR_BROKEN_PIPE) => Ok(Incoming::Closed),
802 Err(error) => Err(error),
803 }
804}
805
806#[cfg(unix)]
807fn close_transport(stream: TransportStream) {
808 let _ = stream.shutdown(std::net::Shutdown::Both);
809}
810
811#[cfg(windows)]
812fn close_transport(_stream: TransportStream) {
813 }
816
817#[cfg(unix)]
818fn write_transport_frame(
819 stream: &mut TransportStream,
820 frame: &[u8],
821 _timeout: Option<Duration>,
822) -> std::io::Result<()> {
823 stream.write_all(frame).and_then(|()| stream.flush())
824}
825
826#[cfg(windows)]
827fn write_transport_frame(
828 stream: &mut TransportStream,
829 frame: &[u8],
830 timeout: Option<Duration>,
831) -> std::io::Result<()> {
832 let deadline = timeout.map(|duration| Instant::now() + duration);
833 let mut offset = 0;
834 while offset < frame.len() {
835 match stream.write(&frame[offset..]) {
836 Ok(0) => return Err(std::io::Error::from(ErrorKind::WriteZero)),
837 Ok(written) => offset += written,
838 Err(error) if error.kind() == ErrorKind::Interrupted => continue,
839 Err(error) if error.kind() == ErrorKind::WouldBlock => {
840 if deadline.is_some_and(|deadline| Instant::now() >= deadline) {
841 return Err(std::io::Error::from(ErrorKind::TimedOut));
842 }
843 std::thread::yield_now();
844 }
845 Err(error) => return Err(error),
846 }
847 }
848 loop {
849 match stream.flush() {
850 Ok(()) => return Ok(()),
851 Err(error) if error.kind() == ErrorKind::Interrupted => continue,
852 Err(error) if error.kind() == ErrorKind::WouldBlock => {
853 if deadline.is_some_and(|deadline| Instant::now() >= deadline) {
854 return Err(std::io::Error::from(ErrorKind::TimedOut));
855 }
856 std::thread::yield_now();
857 }
858 Err(error) => return Err(error),
859 }
860 }
861}