Skip to main content

noq_proto/config/
qlog.rs

1#[cfg(feature = "qlog")]
2use std::io;
3use std::{io::BufWriter, net::SocketAddr, path::PathBuf, time::SystemTime};
4
5use tracing::{trace, warn};
6
7use crate::{ConnectionId, Instant, Side};
8
9/// Constructs a [`QlogConfig`] for individual connections.
10///
11/// This is set via [`TransportConfig::qlog_factory`].
12///
13/// [`TransportConfig::qlog_factory`]: crate::TransportConfig::qlog_factory
14pub trait QlogFactory: Send + Sync + 'static {
15    /// Returns a [`QlogConfig`] for a connection, if logging should be enabled.
16    ///
17    /// If `None` is returned, qlog capture is disabled for the connection.
18    fn for_connection(
19        &self,
20        side: Side,
21        remote: SocketAddr,
22        initial_dst_cid: ConnectionId,
23        now: Instant,
24    ) -> Option<QlogConfig>;
25}
26
27/// Configuration for qlog trace logging.
28///
29/// This struct is returned from [`QlogFactory::for_connection`] if qlog logging should
30/// be enabled for a connection. It allows to set metadata for the qlog trace.
31///
32/// The trace will be written to the provided writer in the [`JSON-SEQ format`] defined in the qlog
33/// spec.
34///
35/// [`JSON-SEQ format`](https://www.ietf.org/archive/id/draft-ietf-quic-qlog-main-schema-13.html#section-5)
36#[cfg(feature = "qlog")]
37pub struct QlogConfig {
38    pub(crate) writer: Box<dyn io::Write + Send + Sync>,
39    pub(crate) title: Option<String>,
40    pub(crate) description: Option<String>,
41    pub(crate) start_time: Option<Instant>,
42}
43
44#[cfg(feature = "qlog")]
45impl QlogConfig {
46    /// Creates a new [`QlogConfig`] that writes a qlog trace to the specified `writer`.
47    pub fn new(writer: Box<dyn io::Write + Send + Sync>) -> Self {
48        Self {
49            writer,
50            title: None,
51            description: None,
52            start_time: None,
53        }
54    }
55
56    /// Title to record in the qlog capture
57    pub fn title(&mut self, title: Option<String>) -> &mut Self {
58        self.title = title;
59        self
60    }
61
62    /// Description to record in the qlog capture
63    pub fn description(&mut self, description: Option<String>) -> &mut Self {
64        self.description = description;
65        self
66    }
67
68    /// Epoch qlog event times are recorded relative to
69    ///
70    /// If unset, the start of the connection is used.
71    pub fn start_time(&mut self, start_time: Instant) -> &mut Self {
72        self.start_time = Some(start_time);
73        self
74    }
75}
76
77/// Enables writing qlog traces to a directory.
78#[derive(Debug)]
79pub struct QlogFileFactory {
80    dir: Option<PathBuf>,
81    prefix: Option<String>,
82    start_instant: Option<Instant>,
83}
84
85impl QlogFileFactory {
86    /// Creates a new qlog factory that writes files into the specified directory.
87    pub fn new(dir: PathBuf) -> Self {
88        Self {
89            dir: Some(dir),
90            prefix: None,
91            start_instant: None,
92        }
93    }
94
95    /// Creates a new qlog factory that writes files into `QLOGDIR`, if set.
96    ///
97    /// If the environment variable `QLOGDIR` is set, qlog traces for all connections handled
98    /// by this endpoint will be written into that directory.
99    /// If the directory doesn't exist it will be created.
100    pub fn from_env() -> Self {
101        let dir = match std::env::var("QLOGDIR") {
102            Ok(dir) => {
103                if let Err(err) = std::fs::create_dir_all(&dir) {
104                    warn!("qlog not enabled: failed to create qlog directory at {dir}: {err}",);
105                    None
106                } else {
107                    Some(PathBuf::from(dir))
108                }
109            }
110            Err(_) => None,
111        };
112        Self {
113            dir,
114            prefix: None,
115            start_instant: None,
116        }
117    }
118
119    /// Sets a prefix to the filename of the generated files.
120    pub fn with_prefix(mut self, prefix: impl ToString) -> Self {
121        self.prefix = Some(prefix.to_string());
122        self
123    }
124
125    /// Override the instant relative to which all events are recorded.
126    ///
127    /// If not set, events will be recorded relative to the start of the connection.
128    pub fn with_start_instant(mut self, start: Instant) -> Self {
129        self.start_instant = Some(start);
130        self
131    }
132}
133
134impl QlogFactory for QlogFileFactory {
135    fn for_connection(
136        &self,
137        side: Side,
138        _remote: SocketAddr,
139        initial_dst_cid: ConnectionId,
140        now: Instant,
141    ) -> Option<QlogConfig> {
142        let dir = self.dir.as_ref()?;
143
144        let name = {
145            let timestamp = SystemTime::now()
146                .checked_sub(Instant::now().duration_since(now))?
147                .duration_since(SystemTime::UNIX_EPOCH)
148                .ok()?
149                .as_millis();
150            let prefix = self
151                .prefix
152                .as_ref()
153                .filter(|prefix| !prefix.is_empty())
154                .map(|prefix| format!("{prefix}-"))
155                .unwrap_or_default();
156            let side = format!("{side:?}").to_lowercase();
157            format!("{prefix}{timestamp}-{initial_dst_cid}-{side}.qlog")
158        };
159        let path = dir.join(name);
160        let file = std::fs::File::create(&path)
161            .inspect_err(|err| warn!("Failed to create qlog file at {}: {err}", path.display()))
162            .ok()?;
163        trace!(
164            "Initialized qlog file for connection {initial_dst_cid} at {}",
165            path.display()
166        );
167        let writer = BufWriter::new(file);
168        let mut config = QlogConfig::new(Box::new(writer));
169        if let Some(instant) = self.start_instant {
170            config.start_time(instant);
171        }
172        Some(config)
173    }
174}