diff --git a/dotenv/src/iter.rs b/dotenv/src/iter.rs index 40506e2..8cfd573 100644 --- a/dotenv/src/iter.rs +++ b/dotenv/src/iter.rs @@ -1,20 +1,22 @@ use std::collections::HashMap; use std::env; use std::io::prelude::*; -use std::io::{BufReader, Lines}; +use std::io::BufReader; use crate::errors::*; use crate::parse; pub struct Iter { - lines: Lines>, + lines: QuotedLines>, substitution_data: HashMap>, } impl Iter { pub fn new(reader: R) -> Iter { Iter { - lines: BufReader::new(reader).lines(), + lines: QuotedLines { + buf: BufReader::new(reader), + }, substitution_data: HashMap::new(), } } @@ -31,6 +33,62 @@ impl Iter { } } +struct QuotedLines { + buf: B, +} + +fn is_complete(buf: &String) -> bool { + let mut escape = false; + let mut strong_quote = false; + let mut weak_quote = false; + let mut count = 0_u32; + + for c in buf.chars() { + if escape { + escape = false + } else { + match c { + '\\' => escape = true, + '"' if !strong_quote => { + count += 1; + weak_quote = true + } + '\'' if !weak_quote => { + count += 1; + strong_quote = true + } + _ => (), + } + } + } + count % 2 == 0 +} + +impl Iterator for QuotedLines { + type Item = Result; + + fn next(&mut self) -> Option> { + let mut buf = String::new(); + loop { + match self.buf.read_line(&mut buf) { + Ok(0) => return None, + Ok(_n) => { + if is_complete(&buf) { + if buf.ends_with('\n') { + buf.pop(); + if buf.ends_with('\r') { + buf.pop(); + } + } + return Some(Ok(buf)); + } + } + Err(e) => return Some(Err(Error::Io(e))), + } + } + } +} + impl Iterator for Iter { type Item = Result<(String, String)>; @@ -38,7 +96,7 @@ impl Iterator for Iter { loop { let line = match self.lines.next() { Some(Ok(line)) => line, - Some(Err(err)) => return Some(Err(Error::Io(err))), + Some(Err(err)) => return Some(Err(err)), None => return None, }; diff --git a/dotenv/tests/test-multiline.rs b/dotenv/tests/test-multiline.rs new file mode 100644 index 0000000..2a6eec8 --- /dev/null +++ b/dotenv/tests/test-multiline.rs @@ -0,0 +1,20 @@ +mod common; + +use dotenv::*; +use std::env; + +use crate::common::*; + +#[test] +fn test_multiline() { + let value = "-----BEGIN PRIVATE KEY-----\n-----END PRIVATE KEY-----\\n\\\"QUOTED\\\""; + let weak = "-----BEGIN PRIVATE KEY-----\n-----END PRIVATE KEY-----\n\"QUOTED\""; + let dir = tempdir_with_dotenv(&format!("WEAK=\"{}\"\nSTRONG='{}'", value, value)).unwrap(); + + dotenv().ok(); + assert_eq!(var("WEAK").unwrap(), weak); + assert_eq!(var("STRONG").unwrap(), value); + + env::set_current_dir(dir.path().parent().unwrap()).unwrap(); + dir.close().unwrap(); +}