From efeedbe69b7fa9845462775c7a528c4f1838b7d8 Mon Sep 17 00:00:00 2001
From: Pavol Rusnak <stick@gk2.sk>
Date: Thu, 6 Sep 2018 19:09:00 +0200
Subject: [PATCH] fsm: implement script_type handling in GetPublicKey

---
 firmware/coin_info.c.mako |  3 ++-
 firmware/coins.h          |  3 ++-
 firmware/fsm_msg_coin.h   | 18 +++++++++++++++++-
 3 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/firmware/coin_info.c.mako b/firmware/coin_info.c.mako
index e504bde2c..80ba716ff 100644
--- a/firmware/coin_info.c.mako
+++ b/firmware/coin_info.c.mako
@@ -38,7 +38,8 @@ const CoinInfo coins[COINS_COUNT] = {
 	.address_type = ${c.address_type},
 	.address_type_p2sh = ${c.address_type_p2sh},
 	.xpub_magic = ${hex(c.xpub_magic)},
-	.xprv_magic = ${hex(c.xprv_magic)},
+	.xpub_magic_segwit_p2sh = ${hex(c.xpub_magic_segwit_p2sh)},
+	.xpub_magic_segwit_native = ${hex(c.xpub_magic_segwit_native)},
 	.fork_id = ${c_int(c.fork_id)},
 	.version_group_id = ${hex(c.version_group_id)},
 	.bech32_prefix = ${c_str(c.bech32_prefix)},
diff --git a/firmware/coins.h b/firmware/coins.h
index 3a3b8116b..b313c6c8d 100644
--- a/firmware/coins.h
+++ b/firmware/coins.h
@@ -41,7 +41,8 @@ typedef struct _CoinInfo {
 	uint32_t address_type;
 	uint32_t address_type_p2sh;
 	uint32_t xpub_magic;
-	uint32_t xprv_magic;
+	uint32_t xpub_magic_segwit_p2sh;
+	uint32_t xpub_magic_segwit_native;
 	uint32_t fork_id;
 	uint32_t version_group_id;
 	const char *bech32_prefix;
diff --git a/firmware/fsm_msg_coin.h b/firmware/fsm_msg_coin.h
index 6152edb4a..8c810205b 100644
--- a/firmware/fsm_msg_coin.h
+++ b/firmware/fsm_msg_coin.h
@@ -25,6 +25,8 @@ void fsm_msgGetPublicKey(GetPublicKey *msg)
 
 	CHECK_PIN
 
+	InputScriptType script_type = msg->has_script_type ? msg->script_type : InputScriptType_SPENDADDRESS;
+
 	const CoinInfo *coin = fsm_getCoin(msg->has_coin_name, msg->coin_name);
 	if (!coin) return;
 
@@ -60,7 +62,21 @@ void fsm_msgGetPublicKey(GetPublicKey *msg)
 		resp->node.public_key.bytes[0] = 0;
 	}
 	resp->has_xpub = true;
-	hdnode_serialize_public(node, fingerprint, coin->xpub_magic, resp->xpub, sizeof(resp->xpub));
+
+	if (coin->xpub_magic && script_type == InputScriptType_SPENDADDRESS) {
+		hdnode_serialize_public(node, fingerprint, coin->xpub_magic, resp->xpub, sizeof(resp->xpub));
+	} else
+	if (coin->has_segwit && coin->xpub_magic_segwit_p2sh && script_type == InputScriptType_SPENDP2SHWITNESS) {
+		hdnode_serialize_public(node, fingerprint, coin->xpub_magic_segwit_p2sh, resp->xpub, sizeof(resp->xpub));
+	} else
+	if (coin->has_segwit && coin->xpub_magic_segwit_native && script_type == InputScriptType_SPENDWITNESS) {
+		hdnode_serialize_public(node, fingerprint, coin->xpub_magic_segwit_native, resp->xpub, sizeof(resp->xpub));
+	} else {
+		fsm_sendFailure(FailureType_Failure_DataError, _("Invalid combination of coin and script_type"));
+		layoutHome();
+		return;
+	}
+
 	msg_write(MessageType_MessageType_PublicKey, resp);
 	layoutHome();
 }