Merge pull request #333 from Ganbin/sign-decoded-message

[WalletConnect] Allow signed message to be hex or plain

Author: Rigidity

crates/sage/src/endpoints/wallet_connect.rs

@@ -1,7 +1,7 @@
 use chia::{
     bls::{master_to_wallet_hardened, master_to_wallet_unhardened, sign},
     clvm_utils::ToTreeHash,
-    protocol::{Bytes, Coin, CoinSpend, SpendBundle},
+    protocol::{Coin, CoinSpend, SpendBundle},
     puzzles::{cat::CatArgs, standard::StandardArgs, DeriveSynthetic, Proof},
 };
 use chia_wallet_sdk::driver::{Layer, SpendContext};
@@ -16,7 +16,7 @@ use tracing::{debug, info, warn};
 
 use crate::{
     parse_asset_id, parse_coin_id, parse_did_id, parse_hash, parse_nft_id, parse_program,
-    parse_public_key, parse_signature, Error, Result, Sage,
+    parse_public_key, parse_signature, parse_signature_message, Error, Result, Sage,
 };
 
 impl Sage {
@@ -347,7 +347,7 @@ impl Sage {
         }
         .derive_synthetic();
 
-        let decoded_message = Bytes::from(hex::decode(&req.message)?);
+        let decoded_message = parse_signature_message(req.message)?;
         let signature = sign(
             &secret_key,
             ("Chia Signed Message", decoded_message).tree_hash(),
@@ -384,7 +384,7 @@ impl Sage {
         }
         .derive_synthetic();
 
-        let decoded_message = Bytes::from(hex::decode(&req.message)?);
+        let decoded_message = parse_signature_message(req.message)?;
         let signature = sign(
             &secret_key,
             ("Chia Signed Message", decoded_message).tree_hash(),

crates/sage/src/utils/parse.rs

@@ -99,6 +99,26 @@ pub fn parse_signature(input: String) -> Result<Signature> {
     Ok(Signature::from_bytes(&signature)?)
 }
 
+/// Parse a signature message.
+///
+/// It takes a string and returns a Bytes object.
+///
+/// This function supports hex strings with or without a 0x prefix.
+/// It also supports non-hex strings.
+pub fn parse_signature_message(input: String) -> Result<Bytes> {
+    let stripped = if let Some(stripped) = input.strip_prefix("0x") {
+        stripped
+    } else {
+        &input
+    };
+
+    if stripped.chars().all(|c| c.is_ascii_hexdigit()) && !stripped.is_empty() {
+        Ok(Bytes::from(hex::decode(stripped)?))
+    } else {
+        Ok(Bytes::from(input.as_bytes()))
+    }
+}
+
 pub fn parse_public_key(input: String) -> Result<PublicKey> {
     let stripped = if let Some(stripped) = input.strip_prefix("0x") {
         stripped
@@ -134,3 +154,41 @@ pub fn parse_memos(input: Option<Vec<String>>) -> Result<Option<Vec<Bytes>>> {
         Ok(None)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_parse_signature_message() {
+        // Test hex string with 0x prefix
+        let input = "0x1234567890abcdef";
+        let result = parse_signature_message(input.to_string()).unwrap();
+        let expected = Bytes::from(hex::decode(&input[2..]).unwrap());
+        assert_eq!(result, expected);
+
+        // Test hex string without prefix
+        let input = "1234567890abcdef";
+        let result = parse_signature_message(input.to_string()).unwrap();
+        let expected = Bytes::from(hex::decode(input).unwrap());
+        assert_eq!(result, expected);
+
+        // Test non-hex string
+        let input = "Hello, world!";
+        let result = parse_signature_message(input.to_string()).unwrap();
+        let expected = Bytes::from(input.as_bytes());
+        assert_eq!(result, expected);
+
+        // Test hex string with 0x prefix
+        let input = "0xcafe";
+        let result = parse_signature_message(input.to_string()).unwrap();
+        let expected = Bytes::from(hex::decode(&input[2..]).unwrap());
+        assert_eq!(result, expected);
+
+        // Test hex string without prefix
+        let input = "cafe";
+        let result = parse_signature_message(input.to_string()).unwrap();
+        let expected = Bytes::from(hex::decode(input).unwrap());
+        assert_eq!(result, expected);
+    }
+}

src/contexts/WalletConnectContext.tsx

@@ -19,7 +19,7 @@ import {
 } from '@/components/ui/dialog';
 import { useWallet } from '@/contexts/WalletContext';
 import { useErrors } from '@/hooks/useErrors';
-import { fromMojos } from '@/lib/utils';
+import { fromMojos, decodeHexMessage, isHex } from '@/lib/utils';
 import { useWalletState } from '@/state';
 import {
   Params,
@@ -40,6 +40,7 @@ import {
   useState,
 } from 'react';
 import { formatNumber } from '../i18n';
+import { Switch } from '@/components/ui/switch';
 
 export interface WalletConnectContextType {
   sessions: SessionTypes.Struct[];
@@ -405,6 +406,42 @@ function SignCoinSpendsDialog({
   );
 }
 
+function MessageToSign(params: { message: string }) {
+  const [showDecoded, setShowDecoded] = useState(false);
+  const isHexMessage = isHex(params.message);
+  const message = isHexMessage
+    ? !showDecoded
+      ? params.message
+      : decodeHexMessage(params.message)
+    : params.message;
+
+  return (
+    <div className='space-y-2'>
+      <div className='flex items-center justify-between gap-2 flex-wrap'>
+        <div className='font-medium'>
+          Message{' '}
+          {isHexMessage && showDecoded && (
+            <span className='text-xs text-muted-foreground ml-1'>
+              (Decoded)
+            </span>
+          )}
+        </div>
+        {isHexMessage && (
+          <div className='flex items-center gap-2'>
+            <span className='text-sm text-muted-foreground whitespace-nowrap'>
+              Show decoded
+            </span>
+            <Switch checked={showDecoded} onCheckedChange={setShowDecoded} />
+          </div>
+        )}
+      </div>
+      <div className='text-sm text-muted-foreground break-all font-mono bg-muted p-2 rounded whitespace-pre-wrap'>
+        {message}
+      </div>
+    </div>
+  );
+}
+
 function SignMessageDialog({
   params,
 }: CommandDialogProps<'chip0002_signMessage'>) {
@@ -416,12 +453,23 @@ function SignMessageDialog({
           {params.publicKey}
         </div>
       </div>
+      <MessageToSign message={params.message} />
+    </div>
+  );
+}
+
+function SignMessageByAddressDialog({
+  params,
+}: CommandDialogProps<'chia_signMessageByAddress'>) {
+  return (
+    <div className='space-y-4 p-4'>
       <div className='space-y-2'>
-        <div className='font-medium'>Message</div>
-        <div className='text-sm text-muted-foreground break-all font-mono bg-muted p-2 rounded whitespace-pre-wrap'>
-          {params.message}
+        <div className='font-medium'>Address</div>
+        <div className='text-sm text-muted-foreground break-all font-mono bg-muted p-2 rounded'>
+          {params.address}
         </div>
       </div>
+      <MessageToSign message={params.message} />
     </div>
   );
 }
@@ -586,27 +634,6 @@ function SendDialog({ params }: CommandDialogProps<'chia_send'>) {
   );
 }
 
-function SignMessageByAddressDialog({
-  params,
-}: CommandDialogProps<'chia_signMessageByAddress'>) {
-  return (
-    <div className='space-y-4 p-4'>
-      <div className='space-y-2'>
-        <div className='font-medium'>Address</div>
-        <div className='text-sm text-muted-foreground break-all font-mono bg-muted p-2 rounded'>
-          {params.address}
-        </div>
-      </div>
-      <div className='space-y-2'>
-        <div className='font-medium'>Message</div>
-        <div className='text-sm text-muted-foreground break-all font-mono bg-muted p-2 rounded whitespace-pre-wrap'>
-          {params.message}
-        </div>
-      </div>
-    </div>
-  );
-}
-
 function DefaultCommandDialog({ params }: { params: unknown }) {
   return (
     <div className='p-4'>

src/lib/utils.ts

@@ -170,3 +170,11 @@ function fromHex(hex: string): Uint8Array {
   }
   return i === bytes.length ? bytes : bytes.slice(0, i);
 }
+
+export function decodeHexMessage(hexMessage: string): string {
+  return new TextDecoder().decode(fromHex(sanitizeHex(hexMessage)));
+}
+
+export function isHex(str: string): boolean {
+  return /^(0x)?[0-9a-fA-F]+$/.test(str);
+}