agent_first_data/document/format/
dotenv.rs1use crate::document::{DocumentError, DocumentResult, Value};
4use std::collections::BTreeMap;
5
6const MAX_DOTENV_VALUE_BYTES: usize = 1024 * 1024;
7
8#[derive(Debug)]
9pub struct DotenvDocument<'a> {
10 source: &'a str,
11 keys: Vec<String>,
12}
13
14impl<'a> DotenvDocument<'a> {
15 pub fn parse(source: &'a str) -> DocumentResult<Self> {
16 let value = parse_semantic(source)?;
17 let keys = value
18 .as_object()
19 .map(|object| object.keys().cloned().collect())
20 .unwrap_or_default();
21 Ok(Self { source, keys })
22 }
23
24 fn has_key(&self, key: &str) -> bool {
25 self.keys.iter().any(|candidate| candidate == key)
26 }
27
28 #[allow(dead_code)]
29 fn source(&self) -> &'a str {
30 self.source
31 }
32}
33
34pub fn load(content: &str) -> DocumentResult<Value> {
36 parse_semantic(content)
37}
38
39fn parse_semantic(content: &str) -> DocumentResult<Value> {
40 let mut values = BTreeMap::new();
41
42 for (line_number, raw_line) in logical_lines(content)? {
43 let line = raw_line
44 .strip_suffix('\r')
45 .unwrap_or(&raw_line)
46 .trim_start();
47 if line.is_empty() || line.starts_with('#') {
48 continue;
49 }
50
51 let assignment = line
52 .strip_prefix("export")
53 .and_then(|rest| {
54 rest.starts_with(char::is_whitespace)
55 .then(|| rest.trim_start())
56 })
57 .unwrap_or(line);
58 let Some((raw_key, raw_value)) = assignment.split_once('=') else {
59 return Err(parse_error(line_number, "expected KEY=value assignment"));
60 };
61 let key = raw_key.trim();
62 if !valid_key(key) {
63 return Err(parse_error(line_number, "invalid variable name"));
64 }
65
66 let value = parse_value(raw_value.trim_start(), line_number)?;
67 if values.contains_key(key) {
68 return Err(parse_error(line_number, "duplicate variable name"));
69 }
70 values.insert(key.to_string(), Value::String(value));
71 }
72
73 Ok(Value::Object(values))
74}
75
76fn logical_lines(content: &str) -> DocumentResult<Vec<(usize, String)>> {
86 let mut lines = Vec::new();
87 let mut buffer = String::new();
88 let mut first_line = 1;
89 let mut line_number = 1;
90 let mut quote = None;
91 let mut quote_opened_at = 0;
92 let mut escaped = false;
93 let mut value_may_open = false;
94 let mut in_comment = false;
95 let mut at_line_start = true;
96
97 for character in content.chars() {
98 if quote.is_none() {
99 if at_line_start && !character.is_whitespace() {
100 at_line_start = false;
101 in_comment = character == '#';
102 }
103 if !in_comment {
104 if value_may_open
105 && character.is_whitespace()
106 && character != '\n'
107 && character != '\r'
108 {
109 buffer.push(character);
110 continue;
111 }
112 if value_may_open && (character == '\'' || character == '"') {
113 quote = Some(character);
114 quote_opened_at = line_number;
115 value_may_open = false;
116 } else if character == '=' {
117 value_may_open = true;
118 } else if !character.is_whitespace() {
119 value_may_open = false;
122 }
123 }
124 } else if escaped {
125 escaped = false;
126 } else if character == '\\' {
127 escaped = true;
128 } else if Some(character) == quote {
129 quote = None;
130 }
131 if character == '\n' && quote.is_none() {
132 lines.push((
133 first_line,
134 buffer.strip_suffix('\n').unwrap_or(&buffer).to_string(),
135 ));
136 buffer.clear();
137 first_line = line_number + 1;
138 } else {
139 buffer.push(character);
140 }
141 if character == '\n' {
142 line_number += 1;
143 if quote.is_none() {
144 value_may_open = false;
145 in_comment = false;
146 at_line_start = true;
147 }
148 }
149 }
150 if quote.is_some() {
151 return Err(parse_error(quote_opened_at, "unterminated quoted value"));
155 }
156 if !buffer.is_empty() || content.is_empty() {
157 lines.push((first_line, buffer));
158 }
159 Ok(lines)
160}
161
162pub fn save(_value: &Value) -> DocumentResult<String> {
164 Err(DocumentError::UnsupportedOperation {
165 format: "dotenv".to_string(),
166 operation: "save".to_string(),
167 detail: "dotenv files are read-only because rewriting would lose comments, ordering, and quoting"
168 .to_string(),
169 })
170}
171
172pub fn set_preserving(content: &str, key: &str, value: &Value) -> DocumentResult<String> {
174 let Value::String(value) = value else {
175 return Err(unsupported("set", "dotenv values are strings"));
176 };
177 let replacement = encode_value(value);
178 edit_line(content, key, Some(&replacement))
179}
180
181pub fn unset_preserving(content: &str, key: &str) -> DocumentResult<String> {
183 edit_line(content, key, None)
184}
185
186fn edit_line(content: &str, key: &str, replacement: Option<&str>) -> DocumentResult<String> {
187 if !valid_key(key) {
188 return Err(parse_error(0, "invalid variable name"));
189 }
190 let document = DotenvDocument::parse(content).map_err(|error| with_path(error, key))?;
191 let mut output = String::with_capacity(content.len());
192 let found = document.has_key(key);
193 for line in content.split_inclusive('\n') {
194 let body = line.strip_suffix('\n').unwrap_or(line);
195 let bare = body.strip_suffix('\r').unwrap_or(body);
196 let trimmed = bare.trim_start();
197 let assignment = trimmed
198 .strip_prefix("export")
199 .and_then(|rest| {
200 rest.starts_with(char::is_whitespace)
201 .then(|| rest.trim_start())
202 })
203 .unwrap_or(trimmed);
204 let matches = assignment
205 .split_once('=')
206 .is_some_and(|(name, _)| name.trim() == key);
207 if matches {
208 if let Some(value) = replacement {
209 let prefix_len = bare.len() - assignment.len();
210 let prefix = &body[..prefix_len];
211 let (name, old_value) = assignment.split_once('=').unwrap_or((key, ""));
212 let suffix = old_value
213 .find(" #")
214 .map(|index| &old_value[index..])
215 .unwrap_or("");
216 output.push_str(prefix);
217 output.push_str(name);
218 output.push('=');
219 output.push_str(value);
220 output.push_str(suffix);
221 if body.ends_with('\r') {
222 output.push('\r');
223 }
224 if line.ends_with('\n') {
225 output.push('\n');
226 }
227 }
228 } else {
229 output.push_str(line);
230 }
231 }
232 if !found {
233 if let Some(value) = replacement {
234 let newline = if content.contains("\r\n") {
235 "\r\n"
236 } else {
237 "\n"
238 };
239 if !output.is_empty() && !output.ends_with('\n') {
240 output.push_str(newline);
241 }
242 output.push_str(key);
243 output.push('=');
244 output.push_str(value);
245 if content.ends_with('\n') || !output.ends_with('\n') {
246 output.push_str(newline);
247 }
248 return Ok(output);
249 }
250 return Err(DocumentError::PathNotFound {
251 path: key.to_string(),
252 });
253 }
254 Ok(output)
255}
256
257fn encode_value(value: &str) -> String {
258 if !value.is_empty()
259 && value
260 .chars()
261 .all(|c| !c.is_whitespace() && c != '#' && c != '\\')
262 {
263 return value.to_string();
264 }
265 format!(
266 "\"{}\"",
267 value
268 .replace('\\', "\\\\")
269 .replace('"', "\\\"")
270 .replace('\n', "\\n")
271 )
272}
273
274fn unsupported(operation: &str, detail: &str) -> DocumentError {
275 DocumentError::UnsupportedOperation {
276 format: "dotenv".to_string(),
277 operation: operation.to_string(),
278 detail: detail.to_string(),
279 }
280}
281
282fn valid_key(key: &str) -> bool {
283 let mut chars = key.chars();
284 chars
285 .next()
286 .is_some_and(|first| first == '_' || first.is_ascii_alphabetic())
287 && chars.all(|character| character == '_' || character.is_ascii_alphanumeric())
288}
289
290fn parse_value(input: &str, line_number: usize) -> DocumentResult<String> {
291 if input.len() > MAX_DOTENV_VALUE_BYTES {
292 return Err(parse_error(line_number, "value exceeds 1 MiB"));
293 }
294 match input.chars().next() {
295 Some('\'') => parse_quoted(input, '\'', line_number),
296 Some('"') => parse_quoted(input, '"', line_number),
297 _ => parse_unquoted(input),
298 }
299}
300
301fn parse_quoted(input: &str, quote: char, line_number: usize) -> DocumentResult<String> {
302 let mut output = String::new();
303 let mut escaped = false;
304
305 for (offset, character) in input[quote.len_utf8()..].char_indices() {
306 if escaped {
307 match (quote, character) {
308 ('"', 'n') => output.push('\n'),
309 ('"', 'r') => output.push('\r'),
310 ('"', 't') => output.push('\t'),
311 ('"', '\\') => output.push('\\'),
312 ('"', '"') => output.push('"'),
313 ('\'', '\'') => output.push('\''),
314 (_, other) => {
315 output.push('\\');
316 output.push(other);
317 }
318 }
319 escaped = false;
320 continue;
321 }
322 if character == '\\' {
323 escaped = true;
324 continue;
325 }
326 if character == quote {
327 let remainder_index = quote.len_utf8() + offset + character.len_utf8();
328 let remainder = input[remainder_index..].trim_start();
329 if !remainder.is_empty() && !remainder.starts_with('#') {
330 return Err(parse_error(
331 line_number,
332 "unexpected content after quoted value",
333 ));
334 }
335 return Ok(output);
336 }
337 output.push(character);
338 }
339
340 Err(parse_error(line_number, "unterminated quoted value"))
341}
342
343fn parse_unquoted(input: &str) -> DocumentResult<String> {
344 let mut output = String::new();
345 let mut escaped = false;
346
347 for character in input.chars() {
348 if escaped {
349 output.push(character);
350 escaped = false;
351 } else if character == '\\' {
352 escaped = true;
353 } else if character == '#' && output.chars().last().is_some_and(char::is_whitespace) {
354 break;
355 } else {
356 output.push(character);
357 }
358 }
359 if escaped {
360 output.push('\\');
361 }
362 Ok(output.trim_end().to_string())
363}
364
365fn parse_error(line_number: usize, detail: &str) -> DocumentError {
366 DocumentError::ParseError {
367 format: "dotenv".to_string(),
368 detail: format!("line {line_number}: {detail}"),
369 }
370}
371
372fn with_path(error: DocumentError, path: &str) -> DocumentError {
373 match error {
374 DocumentError::ParseError { format, detail } => DocumentError::ParseError {
375 format,
376 detail: format!("path `{path}`: {detail}"),
377 },
378 other => other,
379 }
380}
381
382#[cfg(test)]
383mod tests {
384 use super::*;
385
386 #[test]
387 fn a_comment_does_not_swallow_the_file() {
388 let source =
393 "# set KEY=value if it isn't already there\nAPI_HOST=example.com\nAPI_PORT=8080\n";
394 let Value::Object(values) = load(source).unwrap() else {
395 panic!("expected an object");
396 };
397 assert_eq!(values.len(), 2);
398 assert_eq!(
399 values.get("API_HOST"),
400 Some(&Value::String("example.com".to_string()))
401 );
402 }
403
404 #[test]
405 fn a_quote_after_the_value_started_is_literal() {
406 let Value::Object(values) = load("MESSAGE=it isn't quoted\n").unwrap() else {
407 panic!("expected an object");
408 };
409 assert_eq!(
410 values.get("MESSAGE"),
411 Some(&Value::String("it isn't quoted".to_string()))
412 );
413 }
414
415 #[test]
416 fn quoted_values_still_span_lines() {
417 let Value::Object(values) = load("A=\"line1\nline2\"\nB=after\n").unwrap() else {
418 panic!("expected an object");
419 };
420 assert_eq!(
421 values.get("A"),
422 Some(&Value::String("line1\nline2".to_string()))
423 );
424 assert_eq!(values.get("B"), Some(&Value::String("after".to_string())));
425 }
426
427 #[test]
428 fn an_unterminated_quote_is_a_parse_error() {
429 let error = load("A=\"unterminated\nB=after\n").unwrap_err();
432 assert_eq!(error.code(), "document_parse_failed");
433 }
434}