From aeb9c89be81f3e5d5769e15a709e4e9f5926c16c Mon Sep 17 00:00:00 2001
From: Michael Kuc <mdk3018@ic.ac.uk>
Date: Thu, 29 Aug 2019 16:20:42 +0100
Subject: [PATCH] Save immediately on persistent state modification.

Modifications for #6.
---
 Channel.cpp      | 9 +++++++--
 Storage.cpp      | 2 +-
 U2F_CMD.hpp      | 1 +
 U2F_Init_CMD.cpp | 4 ++++
 U2F_Init_CMD.hpp | 1 +
 U2F_Msg_CMD.cpp  | 7 +++++++
 U2F_Msg_CMD.hpp  | 3 ++-
 U2F_Ping_CMD.cpp | 4 ++++
 U2F_Ping_CMD.hpp | 1 +
 9 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/Channel.cpp b/Channel.cpp
index 813aae4..03e2393 100644
--- a/Channel.cpp
+++ b/Channel.cpp
@@ -17,6 +17,7 @@ along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */
 
 #include "Channel.hpp"
+#include "Storage.hpp"
 #include "U2F_CMD.hpp"
 #include "u2f.hpp"
 #include <iostream>
@@ -46,8 +47,12 @@ void Channel::handle(const shared_ptr<U2FMessage> uMsg) {
 
 	auto cmd = U2F_CMD::get(uMsg);
 
-	if (cmd)
-		return cmd->respond(this->cid);
+	if (cmd) {
+		cmd->respond(this->cid);
+
+		if (cmd->modifiesPersistentState())
+			Storage::save();
+	}
 }
 
 void Channel::init(const ChannelInitState newInitState) {
diff --git a/Storage.cpp b/Storage.cpp
index 2260196..e03b941 100644
--- a/Storage.cpp
+++ b/Storage.cpp
@@ -80,7 +80,7 @@ void Storage::init(const string& dirPrefix) {
 void Storage::save() {
 	ofstream file{ Storage::filename };
 
-	for (auto& keypair : Storage::appParams) {
+	for (const auto& keypair : Storage::appParams) {
 		const auto& keyID = keypair.first;
 		const auto& appParam = keypair.second;
 		const auto& privKey = Storage::privKeys[keyID];
diff --git a/U2F_CMD.hpp b/U2F_CMD.hpp
index 75d3a1e..d255a52 100644
--- a/U2F_CMD.hpp
+++ b/U2F_CMD.hpp
@@ -28,4 +28,5 @@ public:
 	virtual ~U2F_CMD() = default;
 	static std::shared_ptr<U2F_CMD> get(const std::shared_ptr<U2FMessage> uMsg);
 	virtual void respond(const uint32_t channelID) const = 0;
+	virtual bool modifiesPersistentState() const = 0;
 }; // For polymorphic type casting
diff --git a/U2F_Init_CMD.cpp b/U2F_Init_CMD.cpp
index c43bb98..7da0d8d 100644
--- a/U2F_Init_CMD.cpp
+++ b/U2F_Init_CMD.cpp
@@ -52,3 +52,7 @@ void U2F_Init_CMD::respond(const uint32_t channelID) const {
 
 	msg.write();
 }
+
+bool U2F_Init_CMD::modifiesPersistentState() const {
+	return false;
+}
diff --git a/U2F_Init_CMD.hpp b/U2F_Init_CMD.hpp
index 9780826..5d600d9 100644
--- a/U2F_Init_CMD.hpp
+++ b/U2F_Init_CMD.hpp
@@ -28,4 +28,5 @@ struct U2F_Init_CMD : U2F_CMD {
 public:
 	U2F_Init_CMD(const std::shared_ptr<U2FMessage> uMsg);
 	virtual void respond(const uint32_t channelID) const override;
+	virtual bool modifiesPersistentState() const override;
 };
diff --git a/U2F_Msg_CMD.cpp b/U2F_Msg_CMD.cpp
index 71ddfeb..c164f58 100644
--- a/U2F_Msg_CMD.cpp
+++ b/U2F_Msg_CMD.cpp
@@ -192,3 +192,10 @@ const map<uint8_t, bool> U2F_Msg_CMD::usesData = { { U2F_REG, true },
 void U2F_Msg_CMD::respond(const uint32_t channelID) const {
 	U2F_Msg_CMD::error(channelID, static_cast<uint16_t>(APDU_STATUS::SW_INS_NOT_SUPPORTED));
 }
+
+bool U2F_Msg_CMD::modifiesPersistentState() const {
+	const auto usesPersistentState = usesData.find(ins);
+
+	// Be conservative for future instructions. Assume that they do modify persist state.
+	return (usesPersistentState == usesData.end()) || usesPersistentState->second;
+}
diff --git a/U2F_Msg_CMD.hpp b/U2F_Msg_CMD.hpp
index 010982b..f0daae0 100644
--- a/U2F_Msg_CMD.hpp
+++ b/U2F_Msg_CMD.hpp
@@ -40,5 +40,6 @@ protected:
 public:
 	static std::shared_ptr<U2F_Msg_CMD> generate(const std::shared_ptr<U2FMessage> uMsg);
 	static void error(const uint32_t channelID, const uint16_t errCode);
-	void respond(const uint32_t channelID) const;
+	void respond(const uint32_t channelID) const override;
+	virtual bool modifiesPersistentState() const override;
 };
diff --git a/U2F_Ping_CMD.cpp b/U2F_Ping_CMD.cpp
index e3777c5..33b265d 100644
--- a/U2F_Ping_CMD.cpp
+++ b/U2F_Ping_CMD.cpp
@@ -33,3 +33,7 @@ void U2F_Ping_CMD::respond(const uint32_t channelID) const {
 	msg.data = nonce;
 	msg.write();
 }
+
+bool U2F_Ping_CMD::modifiesPersistentState() const {
+	return false;
+}
diff --git a/U2F_Ping_CMD.hpp b/U2F_Ping_CMD.hpp
index c32c8a8..204d3a3 100644
--- a/U2F_Ping_CMD.hpp
+++ b/U2F_Ping_CMD.hpp
@@ -28,4 +28,5 @@ struct U2F_Ping_CMD : U2F_CMD {
 public:
 	U2F_Ping_CMD(const std::shared_ptr<U2FMessage> uMsg);
 	virtual void respond(const uint32_t channelID) const override;
+	virtual bool modifiesPersistentState() const override;
 };