From 271ee5e39581bc4d8f5d29f797a6cbc7353dd476 Mon Sep 17 00:00:00 2001
From: Xander Bil <xander@biltopia.org>
Date: Sat, 17 Aug 2024 19:26:16 +0200
Subject: [PATCH] Add size restrictions

---
 zns-cli/src/main.rs         |  1 -
 zns-daemon/src/db/models.rs | 13 +++++++++++++
 zns/src/parser.rs           |  7 +++++++
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/zns-cli/src/main.rs b/zns-cli/src/main.rs
index aa29290..a976aba 100644
--- a/zns-cli/src/main.rs
+++ b/zns-cli/src/main.rs
@@ -239,7 +239,6 @@ impl KeyTransformer for OpenSSHKey {
         }?;
 
         let length = reader.read_u32()?;
-        println!("{}", length);
         reader.read(length as usize)?;
 
         Ok(Self {
diff --git a/zns-daemon/src/db/models.rs b/zns-daemon/src/db/models.rs
index b62be9f..647ee29 100644
--- a/zns-daemon/src/db/models.rs
+++ b/zns-daemon/src/db/models.rs
@@ -1,3 +1,5 @@
+use std::fmt::format;
+
 use diesel::prelude::*;
 use diesel::sql_types::Text;
 use zns::{
@@ -90,7 +92,18 @@ impl Record {
     }
 }
 
+const MAX_RDATA_SIZE: usize = 1000;
+
 pub fn insert_into_database(rr: &RR, connection: &mut PgConnection) -> Result<(), ZNSError> {
+    if rr.rdata.len() > MAX_RDATA_SIZE {
+        return Err(ZNSError::Refused {
+            message: format!(
+                "RDATA size of record is bigger then maximum limit: {}",
+                MAX_RDATA_SIZE
+            ),
+        });
+    }
+
     let record = Record {
         name: rr.name.join("."),
         _type: rr._type.clone().into(),
diff --git a/zns/src/parser.rs b/zns/src/parser.rs
index 376febf..ccf6797 100644
--- a/zns/src/parser.rs
+++ b/zns/src/parser.rs
@@ -131,6 +131,13 @@ impl FromBytes for LabelString {
                 })?,
             );
             code = reader.read_u8()?;
+
+            // Set maximum number of labels.
+            if out.len() > 255 {
+                return Err(ZNSError::Refused {
+                    message: String::from("Exceeded maximum number of labels"),
+                });
+            }
         }
 
         if code & 0b11000000 != 0 {