From 599afed56f0d008754c583b68280a897aa5fd290 Mon Sep 17 00:00:00 2001
From: Joshua Scott <joshsoragaming@gmail.com>
Date: Wed, 11 Apr 2018 00:50:10 +0100
Subject: [PATCH] control: Catch DML errors when parsing control messages

---
 src/protocol/control/ClientHello.cpp | 14 ++++++++++++--
 src/protocol/control/Ping.cpp        | 14 ++++++++++++--
 src/protocol/control/ServerHello.cpp | 14 ++++++++++++--
 3 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/src/protocol/control/ClientHello.cpp b/src/protocol/control/ClientHello.cpp
index 6d6e819..007ec21 100644
--- a/src/protocol/control/ClientHello.cpp
+++ b/src/protocol/control/ClientHello.cpp
@@ -1,5 +1,6 @@
 #include "ki/protocol/control/ClientHello.h"
 #include "ki/dml/Record.h"
+#include "ki/protocol/exception.h"
 
 namespace ki
 {
@@ -64,7 +65,16 @@ namespace control
 		auto *timestamp = record.add_field<dml::INT>("m_timestamp");
 		auto *milliseconds = record.add_field<dml::UINT>("m_milliseconds");
 		auto *session_id = record.add_field<dml::USHRT>("m_session_id");
-		record.read_from(istream);
+		try
+		{
+			record.read_from(istream);
+		}
+		catch (dml::parse_error &e)
+		{
+			std::ostringstream oss;
+			oss << "Error reading ClientHello payload: " << e.what();
+			throw parse_error(oss.str());
+		}
 
 		m_timestamp = timestamp->get_value();
 		m_milliseconds = milliseconds->get_value();
@@ -79,4 +89,4 @@ namespace control
 	}
 }
 }
-}
\ No newline at end of file
+}
diff --git a/src/protocol/control/Ping.cpp b/src/protocol/control/Ping.cpp
index 6260b6b..5e7987b 100644
--- a/src/protocol/control/Ping.cpp
+++ b/src/protocol/control/Ping.cpp
@@ -1,5 +1,6 @@
 #include "ki/protocol/control/Ping.h"
 #include "ki/dml/Record.h"
+#include "ki/protocol/exception.h"
 
 namespace ki
 {
@@ -59,7 +60,16 @@ namespace control
 		auto *session_id = record.add_field<dml::USHRT>("m_session_id");
 		auto *milliseconds = record.add_field<dml::USHRT>("m_milliseconds");
 		auto *minutes = record.add_field<dml::UBYT>("m_minutes");
-		record.read_from(istream);
+		try
+		{
+			record.read_from(istream);
+		}
+		catch (dml::parse_error &e)
+		{
+			std::ostringstream oss;
+			oss << "Error reading Ping payload: " << e.what();
+			throw parse_error(oss.str());
+		}
 
 		m_session_id = session_id->get_value();
 		m_milliseconds = milliseconds->get_value();
@@ -73,4 +83,4 @@ namespace control
 	}
 }
 }
-}
\ No newline at end of file
+}
diff --git a/src/protocol/control/ServerHello.cpp b/src/protocol/control/ServerHello.cpp
index 5f25d89..450a5b5 100644
--- a/src/protocol/control/ServerHello.cpp
+++ b/src/protocol/control/ServerHello.cpp
@@ -1,5 +1,6 @@
 #include "ki/protocol/control/ServerHello.h"
 #include "ki/dml/Record.h"
+#include "ki/protocol/exception.h"
 
 namespace ki
 {
@@ -62,7 +63,16 @@ namespace control
 		record.add_field<dml::UINT>("unknown");
 		auto *timestamp = record.add_field<dml::INT>("m_timestamp");
 		auto *milliseconds = record.add_field<dml::UINT>("m_milliseconds");
-		record.read_from(istream);
+		try
+		{
+			record.read_from(istream);
+		}
+		catch (dml::parse_error &e)
+		{
+			std::ostringstream oss;
+			oss << "Error reading ServerHello payload: " << e.what();
+			throw parse_error(oss.str());
+		}
 
 		m_session_id = session_id->get_value();
 		m_timestamp = timestamp->get_value();
@@ -76,4 +86,4 @@ namespace control
 	}
 }
 }
-}
\ No newline at end of file
+}