From f1381776bf04da9b9c115c72ee083edd47e022e0 Mon Sep 17 00:00:00 2001
From: James Zuccon <zuccon@gmail.com>
Date: Thu, 7 Nov 2024 08:57:59 +1100
Subject: [PATCH] Validate decodePrivateKeyWif version and length

Signed-off-by: James Zuccon <zuccon@gmail.com>
---
 src/lib/key/wallet-import-format.spec.ts | 25 ++++++++++++++++++++++++
 src/lib/key/wallet-import-format.ts      | 12 ++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/src/lib/key/wallet-import-format.spec.ts b/src/lib/key/wallet-import-format.spec.ts
index db0bbb58..3d6b353e 100644
--- a/src/lib/key/wallet-import-format.spec.ts
+++ b/src/lib/key/wallet-import-format.spec.ts
@@ -8,6 +8,7 @@ import {
   encodePrivateKeyWif,
   hexToBin,
   sha256 as internalSha256,
+  WalletImportFormatError,
 } from '../lib.js';
 
 test('decodePrivateKeyWif: pass through errors', (t) => {
@@ -17,6 +18,30 @@ test('decodePrivateKeyWif: pass through errors', (t) => {
   );
 });
 
+test('decodePrivateKeyWif: invalid version (base58/legacy address)', (t) => {
+  t.deepEqual(
+    // cspell: disable-next-line
+    decodePrivateKeyWif('1CQbfkN8cADaJWk29ARtaa55UNdBa1kLaA'),
+    WalletImportFormatError.invalidVersion,
+  );
+});
+
+test('decodePrivateKeyWif: incorrect length (mainnet wif version, 20 bytes)', (t) => {
+  t.deepEqual(
+    // cspell: disable-next-line
+    decodePrivateKeyWif('tWGD2u9st6K9gUr68hdo53qhZZyk3JoQAF'),
+    WalletImportFormatError.incorrectLength,
+  );
+});
+
+test('decodePrivateKeyWif: incorrect length (testnet wif version, 20 bytes)', (t) => {
+  t.deepEqual(
+    // cspell: disable-next-line
+    decodePrivateKeyWif('2fAnAKxErh7kQTbKhrGcAty42RZaSyLdLhp'),
+    WalletImportFormatError.incorrectLength,
+  );
+});
+
 const wifVectors = test.macro<[WalletImportFormatType, string, string]>({
   // eslint-disable-next-line @typescript-eslint/max-params
   exec: (t, type, wif, key) => {
diff --git a/src/lib/key/wallet-import-format.ts b/src/lib/key/wallet-import-format.ts
index 824d681e..3206a69a 100644
--- a/src/lib/key/wallet-import-format.ts
+++ b/src/lib/key/wallet-import-format.ts
@@ -8,6 +8,7 @@ import type { Sha256 } from '../lib.js';
 
 export enum WalletImportFormatError {
   incorrectLength = 'The WIF private key payload is not the correct length.',
+  invalidVersion = 'The WIF private key does not have a valid version',
 }
 
 /**
@@ -89,9 +90,20 @@ export const decodePrivateKeyWif = (
   wifKey: string,
   sha256: { hash: Sha256['hash'] } = internalSha256,
 ) => {
+  const uncompressedPayloadLength = 32;
   const compressedPayloadLength = 33;
   const decoded = decodeBase58AddressFormat(wifKey, sha256);
   if (typeof decoded === 'string') return decoded;
+  if (
+    decoded.version !== Base58AddressFormatVersion.wif &&
+    decoded.version !== Base58AddressFormatVersion.wifTestnet
+  )
+    return WalletImportFormatError.invalidVersion;
+  if (
+    decoded.payload.length !== uncompressedPayloadLength &&
+    decoded.payload.length !== compressedPayloadLength
+  )
+    return WalletImportFormatError.incorrectLength;
   const mainnet = decoded.version === Base58AddressFormatVersion.wif;
   const compressed = decoded.payload.length === compressedPayloadLength;
   const privateKey = compressed