From a20c190fea18fd91e52b53e9d2c5aa968f308572 Mon Sep 17 00:00:00 2001
From: Michael Kuc <michaelkuc6@gmail.com>
Date: Fri, 10 Aug 2018 14:24:26 +0000
Subject: [PATCH] Improved U2F HID error handling.

Fixed U2FMessage's error reporting.
Implemented better error reporting.
---
 U2FMessage.cpp   |  5 ++---
 U2F_CMD.cpp      | 26 +++++++++++++-------------
 U2F_Init_CMD.cpp |  8 ++++++++
 3 files changed, 23 insertions(+), 16 deletions(-)

diff --git a/U2FMessage.cpp b/U2FMessage.cpp
index 02c063e..870f1fa 100644
--- a/U2FMessage.cpp
+++ b/U2FMessage.cpp
@@ -160,11 +160,10 @@ U2FMessage::U2FMessage(const uint32_t nCID, const uint8_t nCMD)
 	: cid{ nCID }, cmd{ nCMD }
 {}
 
-void U2FMessage::error(const uint32_t tCID, const uint16_t tErr)
+void U2FMessage::error(const uint32_t tCID, const uint8_t tErr)
 {
 	U2FMessage msg{};
 	msg.cid = tCID;
 	msg.cmd = U2FHID_ERROR;
-	msg.data.push_back((tErr >> 8) & 0xFF);
-	msg.data.push_back(tErr & 0xFF);
+	msg.data.push_back(tErr);
 }
diff --git a/U2F_CMD.cpp b/U2F_CMD.cpp
index 75c2a59..d27898a 100644
--- a/U2F_CMD.cpp
+++ b/U2F_CMD.cpp
@@ -7,20 +7,20 @@ using namespace std;
 
 shared_ptr<U2F_CMD> U2F_CMD::get(const shared_ptr<U2FMessage> uMsg)
 {
-	switch (uMsg->cmd)
+	try
 	{
-		case U2FHID_MSG:
-			try
-			{
-				return U2F_Msg_CMD::generate(uMsg);
-			}
-			catch (runtime_error)
-			{
+		switch (uMsg->cmd)
+		{
+			case U2FHID_MSG:
+					return U2F_Msg_CMD::generate(uMsg);
+			case U2FHID_INIT:
+				return make_shared<U2F_Init_CMD>(uMsg);
+			default:
 				return {};
-			}
-		case U2FHID_INIT:
-			return make_shared<U2F_Init_CMD>(uMsg);
-		default:
-			return {};
+		}
+	}
+	catch (runtime_error)
+	{
+		return {};
 	}
 }
diff --git a/U2F_Init_CMD.cpp b/U2F_Init_CMD.cpp
index 5cfea69..a7c0bf1 100644
--- a/U2F_Init_CMD.cpp
+++ b/U2F_Init_CMD.cpp
@@ -9,8 +9,16 @@ U2F_Init_CMD::U2F_Init_CMD(const shared_ptr<U2FMessage> uMsg)
 {
 	if (uMsg->cmd != U2FHID_INIT)
 		throw runtime_error{ "Failed to get U2F Init message" };
+	else if (uMsg->cid != CID_BROADCAST)
+	{
+		U2FMessage::error(uMsg->cid, ERR_OTHER);
+		throw runtime_error{ "Invalid CID for init command" };
+	}
 	else if (uMsg->data.size() != INIT_NONCE_SIZE)
+	{
+		U2FMessage::error(uMsg->cid, ERR_INVALID_LEN);
 		throw runtime_error{ "Init nonce is incorrect size" };
+	}
 
 	this->nonce = *reinterpret_cast<const uint64_t*>(uMsg->data.data());
 }