1use std::collections::VecDeque;
15use std::io::{ErrorKind, Read, Write};
16use std::os::unix::net::UnixStream;
17use std::sync::Arc;
18use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
19
20use serde_json::value::RawValue;
21use serde_json::Value;
22
23use crate::debug::{describe_endpoint, error_label, join_capabilities, on_off, Category, DebugLog};
24use crate::diffing::build_delta;
25use crate::error::Error;
26use crate::framing::{encode_frame, FrameDecoder};
27use crate::limits::{Limits, DEFAULT_LIMITS};
28use crate::logs::{AttrValue, LogLevel, LogRecord, MAX_LOG_ATTRS};
29use crate::marker::encode_marker;
30use crate::messages::{
31 default_capabilities, parse_driver_message, GetTree, GetTreeResult, Hello, HelloAck,
32 LogMessage, ProbeInfo, ProtocolErrorMessage, RevisionCommit, SnapshotMessage,
33};
34use crate::roles::Capability;
35use crate::tree::Snapshot;
36use crate::validate::validate_snapshot;
37
38pub const ENV_ENDPOINT: &str = "TERMWRIGHT_ENDPOINT";
40pub const ENV_TOKEN: &str = "TERMWRIGHT_TOKEN";
42pub const ENV_PROTOCOL: &str = "TERMWRIGHT_PROTOCOL";
44
45pub const DIAL_TIMEOUT: Duration = Duration::from_secs(5);
47
48pub const WRITE_TIMEOUT: Duration = Duration::from_millis(250);
56
57const SNAPSHOT_HISTORY: usize = 8;
59
60#[derive(Debug, Clone)]
62pub struct Options {
63 pub adapter_name: String,
65 pub adapter_version: String,
67 pub capabilities: Vec<Capability>,
69 pub limits: Limits,
71 pub write_timeout: Option<Duration>,
74 pub probe: Option<ProbeInfo>,
77 pub debug: Option<Arc<DebugLog>>,
82}
83
84impl Options {
85 pub fn with_logs(adapter_name: impl Into<String>, adapter_version: impl Into<String>) -> Self {
90 let mut options = Self::new(adapter_name, adapter_version);
91 options.capabilities.push(Capability::Logs);
92 options
93 }
94
95 pub fn new(adapter_name: impl Into<String>, adapter_version: impl Into<String>) -> Self {
97 Self {
98 adapter_name: adapter_name.into(),
99 adapter_version: adapter_version.into(),
100 capabilities: default_capabilities(),
101 limits: DEFAULT_LIMITS,
102 write_timeout: Some(WRITE_TIMEOUT),
103 probe: None,
104 debug: None,
108 }
109 }
110}
111
112fn epoch_millis() -> i64 {
115 SystemTime::now()
116 .duration_since(UNIX_EPOCH)
117 .map(|since| since.as_millis() as i64)
118 .unwrap_or(0)
119}
120
121#[derive(Debug)]
127struct TokenBucket {
128 per_second: f64,
129 capacity: f64,
130 tokens: f64,
131 updated: Instant,
132}
133
134impl TokenBucket {
135 fn new(per_second: i64, burst: i64, now: Instant) -> Self {
136 let rate = per_second.max(0) as f64;
137 let capacity = rate + burst.max(0) as f64;
138 Self {
139 per_second: rate,
140 capacity,
141 tokens: capacity,
142 updated: now,
143 }
144 }
145
146 fn take(&mut self, now: Instant) -> bool {
148 if self.per_second <= 0.0 {
149 return false;
150 }
151 let elapsed = now.saturating_duration_since(self.updated).as_secs_f64();
152 self.updated = now;
153 self.tokens = (self.tokens + elapsed * self.per_second).min(self.capacity);
154 if self.tokens < 1.0 {
155 return false;
156 }
157 self.tokens -= 1.0;
158 true
159 }
160}
161
162#[derive(Debug)]
166pub struct Client {
167 endpoint: String,
168 token: String,
169 options: Options,
170 protocol: String,
171
172 stream: Option<UnixStream>,
173 decoder: FrameDecoder,
174 limits: Limits,
175 session_id: Option<String>,
176 revision: i64,
177 marker_enabled: bool,
178 log_budget: Option<crate::messages::LogBudget>,
179 published: Option<Value>,
180 deltas_sent: u64,
181 snapshots_sent: u64,
182 log_seq: i64,
183 log_bucket: Option<TokenBucket>,
184 logs_dropped: u64,
185 subscribe: String,
186 history: VecDeque<(i64, Box<RawValue>)>,
187 force_full: bool,
188}
189
190impl Client {
191 pub fn new(endpoint: impl Into<String>, token: impl Into<String>, options: Options) -> Self {
193 let limits = options.limits;
194 Self {
195 endpoint: endpoint.into(),
196 token: token.into(),
197 options,
198 protocol: crate::messages::PROTOCOL_ID.into(),
199 stream: None,
200 decoder: FrameDecoder::new(limits.max_frame_bytes, limits.max_depth),
201 limits,
202 session_id: None,
203 revision: 0,
204 marker_enabled: false,
205 log_budget: None,
206 published: None,
207 deltas_sent: 0,
208 snapshots_sent: 0,
209 log_seq: 0,
210 log_bucket: None,
211 logs_dropped: 0,
212 subscribe: "snapshots".to_owned(),
213 history: VecDeque::new(),
214 force_full: false,
215 }
216 }
217
218 #[must_use]
223 pub fn qualified_observations(&self) -> bool {
224 self.protocol == crate::messages::PROTOCOL_V2_ID
225 }
226
227 pub fn from_env(mut options: Options) -> Option<Self> {
232 if options.debug.is_none() {
233 options.debug = DebugLog::from_env(&options.adapter_name).map(Arc::new);
234 }
235 Self::from_values(
236 std::env::var(ENV_ENDPOINT).ok().as_deref(),
237 std::env::var(ENV_TOKEN).ok().as_deref(),
238 std::env::var(ENV_PROTOCOL).ok().as_deref(),
239 options,
240 )
241 }
242
243 pub fn from_values(
250 endpoint: Option<&str>,
251 token: Option<&str>,
252 protocol: Option<&str>,
253 options: Options,
254 ) -> Option<Self> {
255 let endpoint = endpoint.filter(|value| !value.is_empty());
256 let token = token.filter(|value| !value.is_empty());
257 let (Some(endpoint), Some(token)) = (endpoint, token) else {
258 if let Some(log) = options.debug.as_ref() {
259 let mut missing = Vec::new();
260 if endpoint.is_none() {
261 missing.push(ENV_ENDPOINT);
262 }
263 if token.is_none() {
264 missing.push(ENV_TOKEN);
265 }
266 log.line(
267 Category::Diag,
268 &format!("dormant: {} not set", missing.join(" and ")),
269 );
270 }
271 return None;
272 };
273 if let Some(protocol) = protocol.filter(|value| !value.is_empty()) {
274 if protocol != crate::messages::PROTOCOL_ID
275 && protocol != crate::messages::PROTOCOL_V2_ID
276 && protocol != "1"
277 && protocol != "2"
278 {
279 if let Some(log) = options.debug.as_ref() {
280 log.line(
281 Category::Diag,
282 &format!(
283 "dormant: {ENV_PROTOCOL}={protocol:?} is not {:?}",
284 crate::messages::PROTOCOL_ID
285 ),
286 );
287 }
288 return None;
289 }
290 }
291 if endpoint.starts_with(r"\\.\pipe\") || endpoint.starts_with(r"\\?\pipe\") {
292 if let Some(log) = options.debug.as_ref() {
295 log.line(
296 Category::Diag,
297 &format!(
298 "dormant: {} needs a Windows transport this client does not have",
299 describe_endpoint(endpoint)
300 ),
301 );
302 }
303 return None;
304 }
305 let v2 = matches!(protocol, Some(crate::messages::PROTOCOL_V2_ID) | Some("2"));
306 let mut client = Self::new(endpoint, token, options);
307 if v2 {
308 client.protocol = crate::messages::PROTOCOL_V2_ID.into();
309 if !client
310 .options
311 .capabilities
312 .contains(&Capability::QualifiedObservations)
313 {
314 client
315 .options
316 .capabilities
317 .push(Capability::QualifiedObservations);
318 }
319 }
320 Some(client)
321 }
322
323 pub fn connect(&mut self, timeout: Duration) -> Result<(), Error> {
331 self.debug_line(
332 Category::Sem,
333 &format!(
334 "dial {} timeout={}ms",
335 describe_endpoint(&self.endpoint),
336 timeout.as_millis()
337 ),
338 );
339 let stream = match UnixStream::connect(&self.endpoint) {
340 Ok(stream) => stream,
341 Err(error) => {
342 self.debug_line(
343 Category::Diag,
344 &format!("dial failed, staying dormant: {}", error_label(&error)),
345 );
346 return Err(error.into());
347 }
348 };
349 stream.set_read_timeout(Some(Duration::from_millis(50)))?;
350 stream.set_write_timeout(self.options.write_timeout)?;
351 self.stream = Some(stream);
352
353 let mut hello = Hello::new(
354 &self.token,
355 &self.options.adapter_name,
356 &self.options.adapter_version,
357 self.options.capabilities.clone(),
358 );
359 hello.protocol = self.protocol.clone();
360 if let Some(probe) = self.options.probe.clone() {
361 hello = hello.with_probe(probe);
362 }
363 self.send(&hello)?;
364 self.debug_line(
365 Category::Sem,
366 &format!(
367 "hello sent adapter={}/{} caps={}",
368 self.options.adapter_name,
369 self.options.adapter_version,
370 join_capabilities(&self.options.capabilities)
371 ),
372 );
373
374 let deadline = Instant::now() + timeout;
375 while self.session_id.is_none() {
376 if Instant::now() >= deadline {
377 self.debug_line(
378 Category::Diag,
379 &format!(
380 "no hello-ack within {}ms, staying dormant",
381 timeout.as_millis()
382 ),
383 );
384 self.close();
385 return Err(Error::HandshakeTimeout);
386 }
387 self.poll()?;
388 }
389 Ok(())
390 }
391
392 fn debug_line(&self, category: Category, message: &str) {
399 if let Some(log) = self.options.debug.as_ref() {
400 log.line(category, message);
401 }
402 }
403
404 pub fn connected(&self) -> bool {
406 self.session_id.is_some() && self.stream.is_some()
407 }
408
409 pub fn session_id(&self) -> Option<&str> {
411 self.session_id.as_deref()
412 }
413
414 pub fn revision(&self) -> i64 {
416 self.revision
417 }
418
419 pub fn log_budget(&self) -> Option<crate::messages::LogBudget> {
422 self.log_budget
423 }
424
425 pub fn limits(&self) -> &Limits {
427 &self.limits
428 }
429
430 pub fn require_full_snapshot(&mut self) {
438 self.force_full = true;
439 }
440
441 #[must_use]
443 pub fn full_snapshot_required(&self) -> bool {
444 self.force_full
445 }
446
447 pub fn close(&mut self) {
449 if let Some(stream) = self.stream.take() {
450 self.debug_line(
451 Category::Sem,
452 &format!(
453 "close r{} snapshots={} deltas={} logs_dropped={}",
454 self.revision, self.snapshots_sent, self.deltas_sent, self.logs_dropped
455 ),
456 );
457 let _ = stream.shutdown(std::net::Shutdown::Both);
458 }
459 self.session_id = None;
460 }
461
462 pub fn publish(&mut self, snapshot: &mut Snapshot) -> Result<Option<String>, Error> {
476 let result = self.publish_inner(snapshot);
477 if matches!(result, Err(Error::Protocol(_) | Error::Validation(_))) {
478 self.require_full_snapshot();
482 }
483 result
484 }
485
486 fn publish_inner(&mut self, snapshot: &mut Snapshot) -> Result<Option<String>, Error> {
487 let Some(session_id) = self.session_id.clone() else {
488 return Ok(None);
489 };
490 if self.stream.is_none() {
491 return Ok(None);
492 }
493
494 let revision = self.revision + 1;
495 snapshot.v = if self.protocol == crate::messages::PROTOCOL_V2_ID {
496 2
497 } else {
498 1
499 };
500 snapshot.session_id = session_id.clone();
501 snapshot.revision = revision;
502
503 let body = serde_json::to_string(&snapshot).map_err(|_| {
504 Error::Protocol(crate::error::Violation::new(
505 "frame-malformed",
506 "snapshot is not JSON-serialisable",
507 ))
508 })?;
509 let parsed: Value = serde_json::from_str(&body).expect("just serialised");
510 validate_snapshot(&parsed, &self.limits)?;
511
512 let marker = if self.marker_enabled {
513 Some(encode_marker(&self.token, &session_id, revision)?)
514 } else {
515 None
516 };
517
518 let mut sent_delta = false;
522 let tree_frame = if self.subscribe != "revisions" {
523 let forced = self.force_full;
531 if forced {
532 self.debug_line(
533 Category::Io,
534 &format!("r{revision} full snapshot: the producer reported a gap"),
535 );
536 }
537 let delta = if self.subscribe == "diffs" && !forced {
538 self.published
539 .as_ref()
540 .and_then(|base| build_delta(base, &parsed))
541 } else {
542 None
543 };
544 match delta {
545 Some(delta) => {
546 sent_delta = true;
547 Some(encode_frame(&delta, self.limits.max_frame_bytes)?)
548 }
549 None => Some(encode_frame(
550 &SnapshotMessage::new(snapshot),
551 self.limits.max_frame_bytes,
552 )?),
553 }
554 } else {
555 None
556 };
557 let commit_frame =
558 encode_frame(&RevisionCommit::new(revision), self.limits.max_frame_bytes)?;
559
560 if let Some(frame) = &tree_frame {
561 self.write_frame(frame)?;
562 }
563 self.write_frame(&commit_frame)?;
564
565 self.revision = revision;
569 self.remember(revision, RawValue::from_string(body).expect("valid JSON"));
570 if tree_frame.is_some() {
571 self.published = Some(parsed);
572 self.force_full = false;
573 if sent_delta {
574 self.deltas_sent += 1;
575 } else {
576 self.snapshots_sent += 1;
577 }
578 }
579
580 Ok(marker)
581 }
582
583 pub fn deltas_sent(&self) -> u64 {
585 self.deltas_sent
586 }
587
588 pub fn snapshots_sent(&self) -> u64 {
590 self.snapshots_sent
591 }
592
593 pub fn logs_dropped(&self) -> u64 {
596 self.logs_dropped
597 }
598
599 pub fn log(&mut self, mut record: LogRecord) -> bool {
616 if self.session_id.is_none() || self.stream.is_none() || self.log_bucket.is_none() {
617 return false;
618 }
619
620 let origin = record.seq;
621 self.log_seq += 1;
622 record.seq = self.log_seq;
623 if record.ts == 0 {
624 record.ts = epoch_millis();
625 }
626 if record.revision.is_none() && self.revision > 0 {
627 record.revision = Some(self.revision);
628 }
629
630 let now = Instant::now();
631 let allowed = self
632 .log_bucket
633 .as_mut()
634 .is_some_and(|bucket| bucket.take(now));
635 if !allowed {
636 self.logs_dropped += 1;
637 return false;
638 }
639 if origin > 0 && record.attrs.len() < MAX_LOG_ATTRS {
640 record
643 .attrs
644 .insert("origin.seq".to_owned(), AttrValue::Int(origin));
645 if record.validate(&self.limits).is_err() {
646 record.attrs.remove("origin.seq");
647 }
648 }
649 if record.validate(&self.limits).is_err() {
650 self.logs_dropped += 1;
653 return false;
654 }
655 self.send(&LogMessage::new(&record)).is_ok()
656 }
657
658 pub fn log_message(&mut self, level: LogLevel, message: impl Into<String>) -> bool {
660 self.log(LogRecord::new(level, message))
661 }
662
663 pub fn poll(&mut self) -> Result<(), Error> {
672 let mut buffer = [0u8; 8192];
673 loop {
674 let read = match self.stream.as_mut() {
675 None => return Ok(()),
676 Some(stream) => stream.read(&mut buffer),
677 };
678 match read {
679 Ok(0) => {
680 self.close();
681 return Ok(());
682 }
683 Ok(count) => {
684 let frames = self.decoder.push(&buffer[..count])?;
685 for frame in frames {
686 self.handle(&frame.value)?;
687 }
688 }
689 Err(error)
690 if matches!(error.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut) =>
691 {
692 return Ok(())
693 }
694 Err(error) if error.kind() == ErrorKind::Interrupted => continue,
695 Err(error) => {
696 self.close();
697 return Err(Error::Io(error));
698 }
699 }
700 }
701 }
702
703 fn handle(&mut self, value: &Value) -> Result<(), Error> {
704 if let Err(error) = parse_driver_message(value, &self.limits) {
705 self.debug_line(
706 Category::Diag,
707 &format!("rejected a driver message: {error}"),
708 );
709 let _ = self.send(&ProtocolErrorMessage::new("malformed", error.to_string()));
710 self.close();
711 return Err(Error::Parse(error));
712 }
713
714 match value.get("type").and_then(Value::as_str) {
715 Some("hello-ack") => {
716 let ack: HelloAck = serde_json::from_value(value.clone()).expect("validated above");
717 self.session_id = Some(ack.session_id);
718 self.limits = ack.limits;
719 self.marker_enabled = ack.marker.enabled;
720 self.log_budget = ack.logs;
721 self.log_bucket = match ack.logs {
722 Some(budget) if budget.enabled => Some(TokenBucket::new(
723 budget.max_records_per_second,
724 budget.burst,
725 Instant::now(),
726 )),
727 _ => None,
728 };
729 self.subscribe = ack.subscribe;
730 if let Some(log) = self.options.debug.as_ref() {
731 let session = self.session_id.clone().unwrap_or_default();
732 log.set_label(&session);
733 log.line(
734 Category::Sem,
735 &format!(
736 "hello-ack session={session} marker={} subscribe={} logs={}",
737 on_off(self.marker_enabled),
738 self.subscribe,
739 on_off(self.log_bucket.is_some())
740 ),
741 );
742 }
743 }
744 Some("get-tree") => {
745 let request: GetTree =
746 serde_json::from_value(value.clone()).expect("validated above");
747 let wanted = request.revision.unwrap_or(self.revision);
748 let held = self
749 .history
750 .iter()
751 .find(|(revision, _)| *revision == wanted)
752 .map(|(_, body)| body.clone());
753 let answer = match held {
754 Some(body) => GetTreeResult::found(request.request_id, body),
755 None => GetTreeResult::missing(
756 request.request_id,
757 format!("revision {wanted} is not retained"),
758 ),
759 };
760 self.send(&answer)?;
761 }
762 Some("error") => {
763 self.debug_line(
764 Category::Diag,
765 &format!(
766 "driver ended the session: {}",
767 value.get("code").and_then(Value::as_str).unwrap_or("?")
768 ),
769 );
770 self.close();
771 }
772 _ => {}
773 }
774 Ok(())
775 }
776
777 fn remember(&mut self, revision: i64, body: Box<RawValue>) {
778 self.history.push_back((revision, body));
779 while self.history.len() > SNAPSHOT_HISTORY {
780 self.history.pop_front();
781 }
782 }
783
784 fn send<T: serde::Serialize>(&mut self, message: &T) -> Result<(), Error> {
785 let frame = encode_frame(message, self.limits.max_frame_bytes)?;
786 self.write_frame(&frame)
787 }
788
789 fn write_frame(&mut self, frame: &[u8]) -> Result<(), Error> {
790 let Some(stream) = self.stream.as_mut() else {
791 return Ok(());
792 };
793 match stream.write_all(frame).and_then(|()| stream.flush()) {
794 Ok(()) => Ok(()),
795 Err(error) => {
796 let timed_out = matches!(
797 error.kind(),
798 ErrorKind::WouldBlock | ErrorKind::TimedOut | ErrorKind::Interrupted
799 );
800 self.close();
801 if timed_out {
802 self.debug_line(
806 Category::Diag,
807 "write deadline exceeded; session is unrecoverable",
808 );
809 return Err(Error::WriteTimeout);
810 }
811 Err(Error::Io(error))
812 }
813 }
814 }
815}