add AlignedWriter for ZonefileFmt

NLnetLabs · Nov 22, 2024 · b34e005 · b34e005
1 parent 8f6e443
commit b34e005
Show file tree

Hide file tree

Showing 2 changed files with 112 additions and 23 deletions.
diff --git a/src/base/dig_printer.rs b/src/base/dig_printer.rs
@@ -2,7 +2,7 @@ use core::fmt;
 
 use crate::rdata::AllRecordData;
 
-use super::zonefile_fmt::ZonefileFmt;
+use super::zonefile_fmt::{DisplayKind, ZonefileFmt};
 use super::ParsedRecord;
 use super::{opt::AllOptData, Message, Rtype};
 
@@ -24,7 +24,7 @@ impl<'a, Octs: AsRef<[u8]>> fmt::Display for DigPrinter<'a, Octs> {
         writeln!(
             f,
             ";; ->>HEADER<<- opcode: {}, rcode: {}, id: {}",
-            header.opcode().display_zonefile(false),
+            header.opcode().display_zonefile(DisplayKind::Simple),
             header.rcode(),
             header.id()
         )?;
@@ -161,7 +161,9 @@ fn write_record_item(
     let parsed = item.to_any_record::<AllRecordData<_, _>>();
 
     match parsed {
-        Ok(item) => writeln!(f, "{}", item.display_zonefile(false)),
+        Ok(item) => {
+            writeln!(f, "{}", item.display_zonefile(DisplayKind::Simple))
+        }
         Err(_) => writeln!(
             f,
             "; {} {} {} {} <invalid data>",

diff --git a/src/base/zonefile_fmt.rs b/src/base/zonefile_fmt.rs
@@ -11,21 +11,32 @@ impl From<fmt::Error> for Error {
 
 pub type Result = core::result::Result<(), Error>;
 
+pub enum DisplayKind {
+    Simple,
+    Aligned,
+    Multiline,
+}
+
 pub struct ZoneFileDisplay<'a, T: ?Sized> {
     inner: &'a T,
-    pretty: bool,
+    kind: DisplayKind,
 }
 
 impl<T: ZonefileFmt + ?Sized> fmt::Display for ZoneFileDisplay<'_, T> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        if self.pretty {
-            self.inner
-                .fmt(&mut MultiLineWriter::new(f))
-                .map_err(|_| fmt::Error)
-        } else {
-            self.inner
+        match self.kind {
+            DisplayKind::Simple => self
+                .inner
                 .fmt(&mut SimpleWriter::new(f))
-                .map_err(|_| fmt::Error)
+                .map_err(|_| fmt::Error),
+            DisplayKind::Aligned => self
+                .inner
+                .fmt(&mut AlignedWriter::new(f))
+                .map_err(|_| fmt::Error),
+            DisplayKind::Multiline => self
+                .inner
+                .fmt(&mut MultiLineWriter::new(f))
+                .map_err(|_| fmt::Error),
         }
     }
 }
@@ -41,10 +52,13 @@ pub trait ZonefileFmt {
     ///
     /// The returned object will be displayed as zonefile when printed or
     /// written using `fmt::Display`.
-    fn display_zonefile(&self, pretty: bool) -> ZoneFileDisplay<'_, Self> {
+    fn display_zonefile(
+        &self,
+        display_kind: DisplayKind,
+    ) -> ZoneFileDisplay<'_, Self> {
         ZoneFileDisplay {
             inner: self,
-            pretty,
+            kind: display_kind,
         }
     }
 }
@@ -128,6 +142,59 @@ impl<W: fmt::Write> FormatWriter for SimpleWriter<W> {
     }
 }
 
+/// A simple writer that aligns some bits
+struct AlignedWriter<W> {
+    first: bool,
+    blocks: usize,
+    writer: W,
+}
+
+impl<W> AlignedWriter<W> {
+    fn new(writer: W) -> Self {
+        Self {
+            first: true,
+            blocks: 0,
+            writer,
+        }
+    }
+}
+
+impl<W: fmt::Write> FormatWriter for AlignedWriter<W> {
+    fn fmt_token(&mut self, args: fmt::Arguments<'_>) -> Result {
+        if !self.first {
+            let c = if self.blocks == 0 { '\t' } else { ' ' };
+            self.writer.write_char(c)?;
+        }
+        self.first = false;
+        self.writer.write_fmt(args)?;
+        Ok(())
+    }
+
+    fn begin_block(&mut self) -> Result {
+        self.blocks += 1;
+        Ok(())
+    }
+
+    fn end_block(&mut self) -> Result {
+        self.blocks -= 1;
+        Ok(())
+    }
+
+    fn fmt_comment(&mut self, _args: fmt::Arguments<'_>) -> Result {
+        Ok(())
+    }
+
+    fn newline(&mut self) -> Result {
+        self.writer.write_char('\n')?;
+        self.first = true;
+
+        // Little sanity thing, it _should_ already be 0
+        self.blocks = 0;
+
+        Ok(())
+    }
+}
+
 struct MultiLineWriter<W> {
     current_column: usize,
     block_indent: Option<usize>,
@@ -237,7 +304,7 @@ mod test {
     use std::vec::Vec;
 
     use crate::base::iana::{Class, DigestAlg, SecAlg};
-    use crate::base::zonefile_fmt::ZonefileFmt;
+    use crate::base::zonefile_fmt::{DisplayKind, ZonefileFmt};
     use crate::base::{Name, Record, Ttl};
     use crate::rdata::{Cds, Cname, Ds, Mx, Txt, A};
 
@@ -251,7 +318,7 @@ mod test {
         let record = create_record(A::new("128.140.76.106".parse().unwrap()));
         assert_eq!(
             "example.com. 3600 IN A 128.140.76.106",
-            record.display_zonefile(false).to_string()
+            record.display_zonefile(DisplayKind::Simple).to_string()
         );
     }
 
@@ -262,7 +329,7 @@ mod test {
         ));
         assert_eq!(
             "example.com. 3600 IN CNAME example.com.",
-            record.display_zonefile(false).to_string()
+            record.display_zonefile(DisplayKind::Simple).to_string()
         );
     }
 
@@ -279,7 +346,7 @@ mod test {
         );
         assert_eq!(
             "example.com. 3600 IN DS 5414 15 2 DEADBEEF",
-            record.display_zonefile(false).to_string()
+            record.display_zonefile(DisplayKind::Simple).to_string()
         );
         assert_eq!(
             [
@@ -289,7 +356,7 @@ mod test {
                 "                          DEADBEEF )",
             ]
             .join("\n"),
-            record.display_zonefile(true).to_string()
+            record.display_zonefile(DisplayKind::Multiline).to_string()
         );
     }
 
@@ -306,7 +373,7 @@ mod test {
         );
         assert_eq!(
             "example.com. 3600 IN CDS 5414 15 2 DEADBEEF",
-            record.display_zonefile(false).to_string()
+            record.display_zonefile(DisplayKind::Simple).to_string()
         );
     }
 
@@ -318,7 +385,7 @@ mod test {
         ));
         assert_eq!(
             "example.com. 3600 IN MX 20 example.com.",
-            record.display_zonefile(false).to_string()
+            record.display_zonefile(DisplayKind::Simple).to_string()
         );
     }
 
@@ -338,7 +405,7 @@ mod test {
             more like a silly monkey with a typewriter accidentally writing \
             some shakespeare along the way but it feels like I have to type \
             e\" \"ven longer to hit that limit!\"",
-            record.display_zonefile(false).to_string()
+            record.display_zonefile(DisplayKind::Simple).to_string()
         );
     }
 
@@ -351,7 +418,7 @@ mod test {
         ));
         assert_eq!(
             "example.com. 3600 IN HINFO \"Windows\" \"Windows Server\"",
-            record.display_zonefile(false).to_string()
+            record.display_zonefile(DisplayKind::Simple).to_string()
         );
     }
 
@@ -368,7 +435,27 @@ mod test {
         ));
         assert_eq!(
             r#"example.com. 3600 IN NAPTR 100 50 "a" "z3950+N2L+N2C" "!^urn:cid:.+@([^\\.]+\\.)(.*)$!\\2!i" cidserver.example.com."#,
-            record.display_zonefile(false).to_string()
+            record.display_zonefile(DisplayKind::Simple).to_string()
+        );
+    }
+
+    #[test]
+    fn aligned() {
+        let record = create_record(
+            Cds::new(
+                5414,
+                SecAlg::ED25519,
+                DigestAlg::SHA256,
+                &[0xDE, 0xAD, 0xBE, 0xEF],
+            )
+            .unwrap(),
+        );
+
+        // The name, ttl, class and rtype should be separated by \t, but the
+        // rdata shouldn't.
+        assert_eq!(
+            "example.com.\t3600\tIN\tCDS 5414 15 2 DEADBEEF",
+            record.display_zonefile(DisplayKind::Aligned).to_string()
         );
     }
 }