From 89b4b7d949bc70784b4570328f1b8b5940c026b6 Mon Sep 17 00:00:00 2001
From: Michael Kuc <michaelkuc6@gmail.com>
Date: Sun, 8 Sep 2019 20:27:47 +0100
Subject: [PATCH] Additional error handling for rarer states.

Ensure std::map::at is not blindly applied without checking that the key
actually exists.

Fixes crashes where u2f client on PC/browser expects channel open, but
service has been restarted.
---
 Controller.cpp            | 16 +++++++++-------
 Streams.cpp               |  5 ++++-
 U2F_Authenticate_APDU.cpp |  2 +-
 U2F_Msg_CMD.cpp           |  7 ++++++-
 4 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/Controller.cpp b/Controller.cpp
index 1708958..91560ff 100644
--- a/Controller.cpp
+++ b/Controller.cpp
@@ -48,17 +48,19 @@ void Controller::handleTransaction(const U2FMessage& msg) {
 
 	lastMessage = chrono::system_clock::now();
 
-	auto opChannel = msg.cid;
+	auto opChannelID = msg.cid;
 
 	if (msg.cmd == U2FHID_INIT) {
-		opChannel = nextChannel();
-		auto channel = Channel{ opChannel };
-
+		opChannelID = nextChannel();
+		auto channel = Channel{ opChannelID };
 		try {
-			channels.emplace(opChannel, channel); // In case of wrap-around replace existing one
+			channels.emplace(opChannelID, channel); // In case of wrap-around replace existing one
 		} catch (...) {
-			channels.insert(make_pair(opChannel, channel));
+			channels.insert(make_pair(opChannelID, channel));
 		}
+	} else if (channels.find(opChannelID) == channels.end()) {
+		U2FMessage::error(opChannelID, ERR_CHANNEL_BUSY);
+		return;
 	}
 
 #ifdef DEBUG_MSGS
@@ -66,7 +68,7 @@ void Controller::handleTransaction(const U2FMessage& msg) {
 	clog << "cid: " << msg.cid << ", cmd: " << static_cast<unsigned int>(msg.cmd) << endl;
 #endif
 
-	channels.at(opChannel).handle(msg);
+	channels.at(opChannelID).handle(msg);
 }
 
 uint32_t Controller::nextChannel() {
diff --git a/Streams.cpp b/Streams.cpp
index 1e32370..077ba91 100644
--- a/Streams.cpp
+++ b/Streams.cpp
@@ -186,7 +186,10 @@ FILE* initHTML(FILE* fPtr, const string& title) {
 void closeHTML(FILE* fPtr) {
 	fprintf(fPtr, "\t</body>\n"
 	              "</html>");
-	fclose(fPtr);
+	int successCode = fclose(fPtr);
+
+	if (successCode != 0)
+		cerr << "File closing error: " << errno << endl;
 }
 
 #endif
diff --git a/U2F_Authenticate_APDU.cpp b/U2F_Authenticate_APDU.cpp
index 4a6d3fa..60c18c3 100644
--- a/U2F_Authenticate_APDU.cpp
+++ b/U2F_Authenticate_APDU.cpp
@@ -60,7 +60,7 @@ void U2F_Authenticate_APDU::respond(const uint32_t channelID) const {
 	if (Storage::appParams.find(keyHB) == Storage::appParams.end()) {
 		// Respond with error code - key handle doesn't exist in storage
 		cerr << "Invalid key handle" << endl;
-		this->error(channelID, SW_WRONG_DATA);
+		this->error(channelID, APDU_STATUS::SW_WRONG_DATA);
 		return;
 	}
 
diff --git a/U2F_Msg_CMD.cpp b/U2F_Msg_CMD.cpp
index 20dac7f..66533d1 100644
--- a/U2F_Msg_CMD.cpp
+++ b/U2F_Msg_CMD.cpp
@@ -81,7 +81,12 @@ shared_ptr<U2F_Msg_CMD> U2F_Msg_CMD::generate(const U2FMessage& uMsg) {
 	const uint32_t cBCount = data.size();
 	auto startPtr = data.begin(), endPtr = data.end();
 
-	if (usesData.at(cmd.ins) || data.size() > 3) {
+	const auto cmdUsesData = usesData.find(cmd.ins);
+
+	if (cmdUsesData == usesData.end()) {
+		U2F_Msg_CMD::error(uMsg.cid, APDU_STATUS::SW_INS_NOT_SUPPORTED);
+		throw runtime_error{ "Unknown instruction: unsure if uses data" };
+	} else if (cmdUsesData->second || data.size() > 3) {
 		if (cBCount == 0) {
 			U2F_Msg_CMD::error(uMsg.cid, APDU_STATUS::SW_WRONG_LENGTH);
 			throw runtime_error{ "Invalid command - should have attached data" };