From 71a39bc0d7d6270a936445fecbb50eb6ff7e7258 Mon Sep 17 00:00:00 2001
From: Tomas Susanka <tsusanka@gmail.com>
Date: Fri, 13 Mar 2020 08:50:55 +0000
Subject: [PATCH] legacy: check inputs' and outputs' script types

---
 legacy/firmware/signing.c            | 74 +++++++++++++++++++++++++---
 tests/device_tests/test_op_return.py | 21 ++------
 2 files changed, 73 insertions(+), 22 deletions(-)

diff --git a/legacy/firmware/signing.c b/legacy/firmware/signing.c
index 01d93f885..4932621d6 100644
--- a/legacy/firmware/signing.c
+++ b/legacy/firmware/signing.c
@@ -588,6 +588,26 @@ void signing_init(const SignTx *msg, const CoinInfo *_coin,
 #define MIN(a, b) (((a) < (b)) ? (a) : (b))
 
 static bool signing_check_input(const TxInputType *txinput) {
+  if (txinput->has_multisig &&
+      (txinput->script_type != InputScriptType_SPENDMULTISIG &&
+       txinput->script_type != InputScriptType_SPENDP2SHWITNESS &&
+       txinput->script_type != InputScriptType_SPENDWITNESS)) {
+    fsm_sendFailure(FailureType_Failure_ProcessError,
+                    _("Multisig field provided but not expected."));
+    signing_abort();
+    return false;
+  }
+  if (txinput->address_n_count > 0 &&
+      (txinput->script_type != InputScriptType_SPENDADDRESS &&
+       txinput->script_type != InputScriptType_SPENDMULTISIG &&
+       txinput->script_type != InputScriptType_SPENDWITNESS &&
+       txinput->script_type != InputScriptType_SPENDP2SHWITNESS)) {
+    fsm_sendFailure(FailureType_Failure_DataError,
+                    "Input's address_n provided but not expected.");
+    signing_abort();
+    return false;
+  }
+
   /* compute multisig fingerprint */
   /* (if all input share the same fingerprint, outputs having the same
    * fingerprint will be considered as change outputs) */
@@ -662,6 +682,36 @@ static bool signing_check_output(TxOutputType *txoutput) {
   //   add it to hash_outputs
   //   ask user for permission
 
+  if (txoutput->has_multisig &&
+      (txoutput->script_type != OutputScriptType_PAYTOMULTISIG &&
+       txoutput->script_type != OutputScriptType_PAYTOP2SHWITNESS &&
+       txoutput->script_type != OutputScriptType_PAYTOWITNESS)) {
+    fsm_sendFailure(FailureType_Failure_DataError,
+                    _("Multisig field provided but not expected."));
+    signing_abort();
+    return false;
+  }
+
+  if (txoutput->address_n_count > 0 &&
+      (txoutput->script_type != OutputScriptType_PAYTOADDRESS &&
+       txoutput->script_type != OutputScriptType_PAYTOMULTISIG &&
+       txoutput->script_type != OutputScriptType_PAYTOWITNESS &&
+       txoutput->script_type != OutputScriptType_PAYTOP2SHWITNESS)) {
+    fsm_sendFailure(FailureType_Failure_DataError,
+                    _("Output's address_n provided but not expected."));
+    signing_abort();
+    return false;
+  }
+
+  if (txoutput->has_op_return_data &&
+      (txoutput->script_type != OutputScriptType_PAYTOOPRETURN)) {
+    fsm_sendFailure(
+        FailureType_Failure_DataError,
+        _("OP RETURN data provided but not OP RETURN script type."));
+    signing_abort();
+    return false;
+  }
+
   if (txoutput->script_type == OutputScriptType_PAYTOOPRETURN) {
     if (txoutput->has_address || (txoutput->address_n_count > 0) ||
         txoutput->has_multisig) {
@@ -670,16 +720,28 @@ static bool signing_check_output(TxOutputType *txoutput) {
       signing_abort();
       return false;
     }
-  }
-  // check for change address
-  bool is_change = false;
-  if (txoutput->address_n_count > 0) {
-    if (txoutput->has_address) {
+    if (txoutput->amount != 0) {
       fsm_sendFailure(FailureType_Failure_DataError,
-                      _("Address in change output"));
+                      _("OP_RETURN output with non-zero amount"));
       signing_abort();
       return false;
     }
+  } else {
+    if (txoutput->has_address && txoutput->address_n_count > 0) {
+      fsm_sendFailure(FailureType_Failure_DataError,
+                      _("Both address and address_n provided."));
+      signing_abort();
+      return false;
+    } else if (!txoutput->has_address && txoutput->address_n_count == 0) {
+      fsm_sendFailure(FailureType_Failure_DataError, _("Missing address"));
+      signing_abort();
+      return false;
+    }
+  }
+
+  // check for change address
+  bool is_change = false;
+  if (txoutput->address_n_count > 0) {
     /*
      * For multisig check that all inputs are multisig
      */
diff --git a/tests/device_tests/test_op_return.py b/tests/device_tests/test_op_return.py
index 25913f9a9..3ec7db8a0 100644
--- a/tests/device_tests/test_op_return.py
+++ b/tests/device_tests/test_op_return.py
@@ -172,14 +172,8 @@ class TestOpReturn:
             with pytest.raises(CallException) as exc:
                 btc.sign_tx(client, "Bitcoin", [inp1], [out1], prev_txes=TX_API)
 
-            if client.features.model == "1":
-                assert exc.value.args[0] == proto.FailureType.ProcessError
-                assert exc.value.args[1].endswith("Failed to compile output")
-            else:
-                assert exc.value.args[0] == proto.FailureType.DataError
-                assert exc.value.args[1].endswith(
-                    "OP_RETURN output with non-zero amount"
-                )
+            assert exc.value.args[0] == proto.FailureType.DataError
+            assert exc.value.args[1].endswith("OP_RETURN output with non-zero amount")
 
     @pytest.mark.skip_ui
     def test_opreturn_address(self, client):
@@ -236,11 +230,6 @@ class TestOpReturn:
                 )
 
             assert exc.value.args[0] == proto.FailureType.DataError
-            if client.features.model == "1":
-                assert exc.value.args[1].endswith(
-                    "OP_RETURN output with address or multisig"
-                )
-            else:
-                assert (
-                    exc.value.args[1] == "Output's address_n provided but not expected."
-                )
+            assert exc.value.args[1].endswith(
+                "Output's address_n provided but not expected."
+            )