feat(legacy): add support for aptos message sign (#219)

somebodyLi · web-flow · commit d5e105aae6dc · 2022-11-03T16:15:11.000+08:00
diff --git a/common/protob/messages-aptos.proto b/common/protob/messages-aptos.proto
@@ -41,3 +41,30 @@ message AptosSignedTx {
     required bytes public_key = 1;  // public key for the private key used to sign tx
     required bytes signature = 2;           // the signature of the raw transaction
 }
+
+/**
+ * Request: ask device to sign Aptos message
+ * @start
+ * @next AptosMessageSignature
+ */
+ message AptosSignMessage {
+    repeated uint32 address_n = 1;
+    required AptosMessagePayload payload = 2;
+
+    message AptosMessagePayload {
+        optional string address = 2;
+        optional string chain_id = 3;
+        optional string application = 4;
+        required string nonce = 5;
+        required string message = 6;
+    }
+}
+
+/**
+* Response: signature for message
+* @end
+**/
+message AptosMessageSignature {
+    required bytes signature = 1;
+    required string address = 2;
+}
diff --git a/common/protob/messages.proto b/common/protob/messages.proto
@@ -367,6 +367,8 @@ enum MessageType {
     MessageType_AptosAddress = 10601 [(wire_out) = true];
     MessageType_AptosSignTx = 10602 [(wire_in) = true];
     MessageType_AptosSignedTx = 10603 [(wire_out) = true];
+    MessageType_AptosSignMessage = 10604 [(wire_in) = true];
+    MessageType_AptosMessageSignature = 10605 [(wire_out) = true];
 
     // WebAuthn
     MessageType_WebAuthnListResidentCredentials = 800 [(wire_in) = true];
diff --git a/core/src/trezor/enums/MessageType.py b/core/src/trezor/enums/MessageType.py
@@ -240,6 +240,8 @@
     AptosAddress = 10601
     AptosSignTx = 10602
     AptosSignedTx = 10603
+    AptosSignMessage = 10604
+    AptosMessageSignature = 10605
     WebAuthnListResidentCredentials = 800
     WebAuthnCredentials = 801
     WebAuthnAddResidentCredential = 802
diff --git a/core/src/trezor/enums/__init__.py b/core/src/trezor/enums/__init__.py
@@ -245,6 +245,8 @@ class MessageType(IntEnum):
         AptosAddress = 10601
         AptosSignTx = 10602
         AptosSignedTx = 10603
+        AptosSignMessage = 10604
+        AptosMessageSignature = 10605
         WebAuthnListResidentCredentials = 800
         WebAuthnCredentials = 801
         WebAuthnAddResidentCredential = 802
diff --git a/core/src/trezor/messages.py b/core/src/trezor/messages.py
@@ -6254,6 +6254,60 @@ def __init__(
         def is_type_of(cls, msg: protobuf.MessageType) -> TypeGuard["AptosSignedTx"]:
             return isinstance(msg, cls)
 
+    class AptosSignMessage(protobuf.MessageType):
+        address_n: "list[int]"
+        payload: "AptosMessagePayload"
+
+        def __init__(
+            self,
+            *,
+            payload: "AptosMessagePayload",
+            address_n: "list[int] | None" = None,
+        ) -> None:
+            pass
+
+        @classmethod
+        def is_type_of(cls, msg: protobuf.MessageType) -> TypeGuard["AptosSignMessage"]:
+            return isinstance(msg, cls)
+
+    class AptosMessageSignature(protobuf.MessageType):
+        signature: "bytes"
+        address: "str"
+
+        def __init__(
+            self,
+            *,
+            signature: "bytes",
+            address: "str",
+        ) -> None:
+            pass
+
+        @classmethod
+        def is_type_of(cls, msg: protobuf.MessageType) -> TypeGuard["AptosMessageSignature"]:
+            return isinstance(msg, cls)
+
+    class AptosMessagePayload(protobuf.MessageType):
+        address: "str | None"
+        chain_id: "str | None"
+        application: "str | None"
+        nonce: "str"
+        message: "str"
+
+        def __init__(
+            self,
+            *,
+            nonce: "str",
+            message: "str",
+            address: "str | None" = None,
+            chain_id: "str | None" = None,
+            application: "str | None" = None,
+        ) -> None:
+            pass
+
+        @classmethod
+        def is_type_of(cls, msg: protobuf.MessageType) -> TypeGuard["AptosMessagePayload"]:
+            return isinstance(msg, cls)
+
     class NearGetAddress(protobuf.MessageType):
         address_n: "list[int]"
         show_display: "bool | None"
diff --git a/legacy/firmware/aptos.c b/legacy/firmware/aptos.c
@@ -15,6 +15,8 @@ static const uint8_t APTOS_RAW_TX_PREFIX[32] = {
     54,  67,  169, 188, 111, 102, 147, 189, 220, 26,  159,
     236, 158, 103, 74,  70,  30,  170, 0,   177, 147};
 
+static const char *MESSAGE_PREFIX = "APTOS\n";
+
 void aptos_get_address_from_public_key(const uint8_t *public_key,
                                        char *address) {
   uint8_t buf[SIZE_PUBKEY] = {0};
@@ -52,3 +54,56 @@ void aptos_sign_tx(const AptosSignTx *msg, const HDNode *node,
   resp->public_key.size = 32;
   msg_write(MessageType_MessageType_AptosSignedTx, resp);
 }
+
+void aptos_sign_message(const AptosSignMessage *msg, const HDNode *node,
+                        AptosMessageSignature *resp) {
+  AptosMessagePayload payload = msg->payload;
+  char full_message[sizeof(AptosMessagePayload) + 58] = {0};
+
+  strcat(full_message, MESSAGE_PREFIX);
+  if (payload.has_address) {
+    char *address = payload.address;
+    strcat(full_message, "address: ");
+    strcat(full_message, address);
+    strcat(full_message, "\n");
+  }
+  if (payload.has_application) {
+    char *application = payload.application;
+    strcat(full_message, "application: ");
+    strcat(full_message, application);
+    strcat(full_message, "\n");
+  }
+  if (payload.has_chain_id) {
+    char *chain_id = payload.chain_id;
+    strcat(full_message, "chainId: ");
+    strcat(full_message, chain_id);
+    strcat(full_message, "\n");
+  }
+  char *message = payload.message;
+  strcat(full_message, "message: ");
+  strcat(full_message, message);
+  strcat(full_message, "\n");
+  char *nonce = payload.nonce;
+  strcat(full_message, "nonce: ");
+  strcat(full_message, nonce);
+
+  aptos_get_address_from_public_key(node->public_key + 1, resp->address);
+  // display here
+  layoutVerifyAddress(NULL, resp->address);
+  if (!protectButton(ButtonRequestType_ButtonRequest_Other, false)) {
+    fsm_sendFailure(FailureType_Failure_ActionCancelled, NULL);
+    layoutHome();
+    return;
+  }
+  if (!fsm_layoutSignMessage((const uint8_t *)full_message,
+                             strlen(full_message))) {
+    fsm_sendFailure(FailureType_Failure_ActionCancelled, NULL);
+    layoutHome();
+    return;
+  }
+  ed25519_sign((const uint8_t *)full_message, strlen(full_message),
+               node->private_key, node->public_key + 1, resp->signature.bytes);
+
+  resp->signature.size = 64;
+  msg_write(MessageType_MessageType_AptosMessageSignature, resp);
+}
diff --git a/legacy/firmware/aptos.h b/legacy/firmware/aptos.h
@@ -32,4 +32,6 @@ void aptos_get_address_from_public_key(const uint8_t *public_key,
                                        char *address);
 void aptos_sign_tx(const AptosSignTx *msg, const HDNode *node,
                    AptosSignedTx *resp);
+void aptos_sign_message(const AptosSignMessage *msg, const HDNode *node,
+                        AptosMessageSignature *resp);
 #endif  // __APTOS_H__
diff --git a/legacy/firmware/fsm.h b/legacy/firmware/fsm.h
@@ -176,6 +176,7 @@ void fsm_msgTronSignTx(TronSignTx *msg);
 // aptos
 void fsm_msgAptosGetAddress(const AptosGetAddress *msg);
 void fsm_msgAptosSignTx(const AptosSignTx *msg);
+void fsm_msgAptosSignMessage(const AptosSignMessage *msg);
 
 // near
 void fsm_msgNearGetAddress(NearGetAddress *msg);
diff --git a/legacy/firmware/fsm_msg_aptos.h b/legacy/firmware/fsm_msg_aptos.h
@@ -62,3 +62,20 @@ void fsm_msgAptosSignTx(const AptosSignTx *msg) {
 
   layoutHome();
 }
+
+void fsm_msgAptosSignMessage(const AptosSignMessage *msg) {
+  CHECK_INITIALIZED
+
+  CHECK_PIN
+  RESP_INIT(AptosMessageSignature)
+
+  HDNode *node = fsm_getDerivedNode(ED25519_NAME, msg->address_n,
+                                    msg->address_n_count, NULL);
+  if (!node) return;
+
+  hdnode_fill_public_key(node);
+
+  aptos_sign_message(msg, node, resp);
+
+  layoutHome();
+}
diff --git a/legacy/firmware/protob/messages-aptos.options b/legacy/firmware/protob/messages-aptos.options
@@ -6,3 +6,13 @@ AptosSignTx.raw_tx                                          max_size:1024
 
 AptosSignedTx.public_key                                    max_size:32
 AptosSignedTx.signature                                     max_size:64
+
+AptosSignMessage.address_n                                  max_count:8
+AptosMessagePayload.address                                      max_size:67
+AptosMessagePayload.chain_id                                     max_size: 8
+AptosMessagePayload.application                                  max_size:128
+AptosMessagePayload.nonce                                        max_size:64
+AptosMessagePayload.message                                      max_size:1024
+
+AptosMessageSignature.signature                             max_size:64
+AptosMessageSignature.address                               max_size:67
diff --git a/python/src/trezorlib/aptos.py b/python/src/trezorlib/aptos.py
@@ -16,7 +16,7 @@
 # along with this library.  If not, see <http://www.gnu.org/licenses/>.
 
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Dict
 
 from . import messages
 from .tools import expect
@@ -39,3 +39,19 @@ def get_address(
 @expect(messages.AptosSignedTx)
 def sign_tx(client: "TrezorClient", address_n: "Address", rawtx: bytes):
     return client.call(messages.AptosSignTx(address_n=address_n, raw_tx=rawtx))
+
+
+@expect(messages.AptosMessageSignature)
+def sign_message(client: "TrezorClient", address_n: "Address", payload: Dict):
+    return client.call(
+        messages.AptosSignMessage(
+            address_n=address_n,
+            payload=messages.AptosMessagePayload(
+                address=payload["address"],
+                chain_id=payload["chain_id"],
+                application=payload["application"],
+                nonce=payload["nonce"],
+                message=payload["message"],
+            ),
+        )
+    )
diff --git a/python/src/trezorlib/cli/aptos.py b/python/src/trezorlib/cli/aptos.py
@@ -56,3 +56,40 @@ def sign_raw_tx(client: "TrezorClient", address: str, message: str):
         "signature": f"0x{resp.signature.hex()}",
     }
     return result
+
+
+@cli.command()
+@click.option("-n", "--address", required=True, help=PATH_HELP)
+@click.option("-a", "--addr")
+@click.option("-c", "--chain_id")
+@click.option("-app", "--application")
+@click.argument("nonce")
+@click.argument("message")
+@with_client
+def sign_message(
+    client: "TrezorClient",
+    address: str,
+    addr: str,
+    chain_id: str,
+    application: str,
+    nonce: str,
+    message: str,
+):
+    """Sign formatted message."""
+    address_n = tools.parse_path(address)
+    resp = aptos.sign_message(
+        client,
+        address_n,
+        {
+            "address": addr,
+            "chain_id": chain_id,
+            "application": application,
+            "nonce": nonce,
+            "message": message,
+        },
+    )
+    result = {
+        "address": resp.address,
+        "signature": f"0x{resp.signature.hex()}",
+    }
+    return result
diff --git a/python/src/trezorlib/messages.py b/python/src/trezorlib/messages.py
@@ -266,6 +266,8 @@ class MessageType(IntEnum):
     AptosAddress = 10601
     AptosSignTx = 10602
     AptosSignedTx = 10603
+    AptosSignMessage = 10604
+    AptosMessageSignature = 10605
     WebAuthnListResidentCredentials = 800
     WebAuthnCredentials = 801
     WebAuthnAddResidentCredential = 802
@@ -666,6 +668,66 @@ def __init__(
         self.signature = signature
 
 
+class AptosSignMessage(protobuf.MessageType):
+    MESSAGE_WIRE_TYPE = 10604
+    FIELDS = {
+        1: protobuf.Field("address_n", "uint32", repeated=True, required=False),
+        2: protobuf.Field("payload", "AptosMessagePayload", repeated=False, required=True),
+    }
+
+    def __init__(
+        self,
+        *,
+        payload: "AptosMessagePayload",
+        address_n: Optional[Sequence["int"]] = None,
+    ) -> None:
+        self.address_n: Sequence["int"] = address_n if address_n is not None else []
+        self.payload = payload
+
+
+class AptosMessageSignature(protobuf.MessageType):
+    MESSAGE_WIRE_TYPE = 10605
+    FIELDS = {
+        1: protobuf.Field("signature", "bytes", repeated=False, required=True),
+        2: protobuf.Field("address", "string", repeated=False, required=True),
+    }
+
+    def __init__(
+        self,
+        *,
+        signature: "bytes",
+        address: "str",
+    ) -> None:
+        self.signature = signature
+        self.address = address
+
+
+class AptosMessagePayload(protobuf.MessageType):
+    MESSAGE_WIRE_TYPE = None
+    FIELDS = {
+        2: protobuf.Field("address", "string", repeated=False, required=False),
+        3: protobuf.Field("chain_id", "string", repeated=False, required=False),
+        4: protobuf.Field("application", "string", repeated=False, required=False),
+        5: protobuf.Field("nonce", "string", repeated=False, required=True),
+        6: protobuf.Field("message", "string", repeated=False, required=True),
+    }
+
+    def __init__(
+        self,
+        *,
+        nonce: "str",
+        message: "str",
+        address: Optional["str"] = None,
+        chain_id: Optional["str"] = None,
+        application: Optional["str"] = None,
+    ) -> None:
+        self.nonce = nonce
+        self.message = message
+        self.address = address
+        self.chain_id = chain_id
+        self.application = application
+
+
 class BinanceGetAddress(protobuf.MessageType):
     MESSAGE_WIRE_TYPE = 700
     FIELDS = {
