From ebd56306e951e751883e9f173e1b6235846e78bf Mon Sep 17 00:00:00 2001
From: Mikko Rasa <tdb@tdb.fi>
Date: Fri, 8 Nov 2013 00:23:32 +0200
Subject: [PATCH] Add some more range checks in ArduControl and improve
 existing ones

---
 source/libr2c2/arducontrol.cpp | 27 ++++++++++++++++-----------
 source/libr2c2/arducontrol.h   | 10 +++++++++-
 2 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/source/libr2c2/arducontrol.cpp b/source/libr2c2/arducontrol.cpp
index 9cf8383..e940a2a 100644
--- a/source/libr2c2/arducontrol.cpp
+++ b/source/libr2c2/arducontrol.cpp
@@ -9,6 +9,11 @@ using namespace Msp;
 
 namespace R2C2 {
 
+ArduControl::ProtocolInfo ArduControl::protocol_info[2] =
+{
+	{ 79, 14, 4 }  // MM
+};
+
 ArduControl::ArduControl(const string &dev):
 	serial(dev),
 	debug(true),
@@ -57,17 +62,15 @@ ArduControl::Protocol ArduControl::map_protocol(const string &proto_name)
 {
 	if(proto_name=="MM")
 		return MM;
+	else if(proto_name=="MFX")
+		return MFX;
 	else
 		throw invalid_argument("ArduControl::map_protocol");
 }
 
 unsigned ArduControl::get_protocol_speed_steps(const string &proto_name) const
 {
-	Protocol proto = map_protocol(proto_name);
-	if(proto==MM)
-		return 14;
-	else
-		return 0;
+	return protocol_info[map_protocol(proto_name)].max_speed;
 }
 
 unsigned ArduControl::add_loco(unsigned addr, const string &proto_name, const VehicleType &)
@@ -76,12 +79,8 @@ unsigned ArduControl::add_loco(unsigned addr, const string &proto_name, const Ve
 		throw invalid_argument("ArduControl::add_loco");
 
 	Protocol proto = map_protocol(proto_name);
-
-	if(proto==MM)
-	{
-		if(addr>=80)
-			throw invalid_argument("ArduControl::add_loco");
-	}
+	if(addr>protocol_info[proto].max_address)
+		throw invalid_argument("ArduControl::add_loco");
 
 	Locomotive loco(proto, addr);
 	insert_unique(locomotives, loco.id, loco);
@@ -99,6 +98,9 @@ void ArduControl::remove_loco(unsigned id)
 void ArduControl::set_loco_speed(unsigned id, unsigned speed)
 {
 	Locomotive &loco = get_item(locomotives, id);
+	if(speed>protocol_info[loco.proto].max_speed)
+		throw invalid_argument("ArduControl::set_loco_speed");
+
 	if(loco.speed.set(speed))
 	{
 		QueuedCommand cmd(loco, Locomotive::SPEED);
@@ -123,6 +125,9 @@ void ArduControl::set_loco_reverse(unsigned id, bool rev)
 void ArduControl::set_loco_function(unsigned id, unsigned func, bool state)
 {
 	Locomotive &loco = get_item(locomotives, id);
+	if(func>protocol_info[loco.proto].max_func)
+		throw invalid_argument("ArduControl::set_loco_function");
+
 	unsigned mask = 1<<func;
 	if(loco.funcs.set((loco.funcs&~mask)|(mask*state)))
 	{
diff --git a/source/libr2c2/arducontrol.h b/source/libr2c2/arducontrol.h
index f0c14f9..fb0cb68 100644
--- a/source/libr2c2/arducontrol.h
+++ b/source/libr2c2/arducontrol.h
@@ -66,10 +66,16 @@ private:
 
 	enum Protocol
 	{
-		NONE,
 		MM
 	};
 
+	struct ProtocolInfo
+	{
+		unsigned max_address;
+		unsigned max_speed;
+		unsigned max_func;
+	};
+
 	template<typename T>
 	struct ControlledVariable
 	{
@@ -202,6 +208,8 @@ private:
 	Msp::Mutex mutex;
 	ControlThread thread;
 
+	static ProtocolInfo protocol_info[2];
+
 public:
 	ArduControl(const std::string &);
 	~ArduControl();
-- 
2.45.2