From 7a25d24ba3c1c39029f1e7c6df1ca2ac9b0af61e Mon Sep 17 00:00:00 2001
From: Yohann D'ANELLO <ynerant@crans.org>
Date: Mon, 21 Dec 2020 16:04:13 +0100
Subject: [PATCH] Ensure that a DataTLV is not too long

Signed-off-by: Yohann D'ANELLO <ynerant@crans.org>
---
 squinnondation/hazel.py          | 12 ++++++++++--
 squinnondation/messages.py       |  7 ++++++-
 squinnondation/squinnondation.py |  2 +-
 3 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/squinnondation/hazel.py b/squinnondation/hazel.py
index 06d1024..832dc5a 100644
--- a/squinnondation/hazel.py
+++ b/squinnondation/hazel.py
@@ -109,7 +109,10 @@ class Squirrel(Hazelnut):
             self.refresh_input()
             if not self.squinnondation.no_emoji:
                 self.refresh_emoji_pad()
-            key = self.squinnondation.screen.getkey(curses.LINES - 1, 3 + len(self.nickname) + self.input_index)
+            try:
+                key = self.squinnondation.screen.getkey(curses.LINES - 1, 3 + len(self.nickname) + self.input_index)
+            except curses.error:
+                continue
 
             if key == "KEY_MOUSE":
                 try:
@@ -169,7 +172,12 @@ class Squirrel(Hazelnut):
             return
         elif key != "\n":
             # Insert the pressed key in the current message
-            self.input_buffer = self.input_buffer[:self.input_index] + key + self.input_buffer[self.input_index:]
+            new_buffer = self.input_buffer[:self.input_index] + key + self.input_buffer[self.input_index:]
+            if len(DataTLV.construct(f"<{self.nickname}> {new_buffer}")) > 255 - 8 - 4:
+                # The message is too long to be sent once. We don't allow the user to type any other character.
+                curses.beep()
+                return
+            self.input_buffer = new_buffer
             self.input_index += 1
             return
 
diff --git a/squinnondation/messages.py b/squinnondation/messages.py
index 1fe9d39..869ae87 100644
--- a/squinnondation/messages.py
+++ b/squinnondation/messages.py
@@ -183,6 +183,11 @@ class DataTLV(TLV):
     nonce: int
     data: bytes
 
+    def validate_data(self) -> bool:
+        if len(self.data) >= 256 - 4 - 8:
+            raise ValueError("The data is too long, the length is larger that one byte.")
+        return True
+
     def unmarshal(self, raw_data: bytes) -> None:
         self.type = raw_data[0]
         self.length = raw_data[1]
@@ -305,7 +310,7 @@ class Packet:
             raise ValueError("The magic code of the packet must be 95, found: {:d}".format(self.magic))
         if self.version != 0:
             raise ValueError("The version of the packet is not supported: {:d}".format(self.version))
-        if not (0 <= self.body_length <= 120):
+        if not (0 <= self.body_length <= 1200):
             raise ValueError("The body length of the packet is negative or too high. It must be between 0 and 1020,"
                              "found: {:d}".format(self.body_length))
         return all(tlv.validate_data() for tlv in self.body)
diff --git a/squinnondation/squinnondation.py b/squinnondation/squinnondation.py
index b80d8c3..1769fbb 100644
--- a/squinnondation/squinnondation.py
+++ b/squinnondation/squinnondation.py
@@ -52,7 +52,7 @@ class Squinnondation:
             instance.screen = screen
             screen.addstr(0, 0, "Enter your nickname: ")
             curses.echo()
-            nickname = screen.getstr().decode("UTF-8")
+            nickname = screen.getstr(225).decode("UTF-8")  # Limit nickname length to be included in a DataTLV
             curses.noecho()
 
             squirrel = Squirrel(instance, nickname)