rsa: add support for 4096 and 8192 bit keys

Allow to use larger RSA keys by setting CONFIG_RSA_KEY_SIZE to 4096 or
8192 rather than using the default 2048-bit size.

It's mainly for benchmarking purpose right now as we don't have the RAM
to store the 3x key size buffer and the flash space for the public key
structure.

Signed-off-by: Vincent Palatin <vpalatin@chromium.org>

BRANCH=samus
BUG=none
TEST=build Zinger with CONFIG_RSA_KEY_SIZE equals to 4096 and run it.

Change-Id: I9839121bf158d0a30dde1e48d875f345191bfec2
Reviewed-on: https://chromium-review.googlesource.com/228925
Reviewed-by: Randall Spangler <rspangler@chromium.org>
Commit-Queue: Vincent Palatin <vpalatin@chromium.org>
Tested-by: Vincent Palatin <vpalatin@chromium.org>

9 files changed, 102 insertions, 21 deletions

diff --git a/Makefile b/Makefile
index 13c44d8b76..a2d321b17b 100644
--- a/Makefile
+++ b/Makefile
@@ -95,6 +95,11 @@ _rw_size_str:=$(shell echo "CONFIG_FW_RW_SIZE" | $(CPP) $(CPPFLAGS) -P \
 		-Ichip/$(CHIP) -Iboard/$(BOARD) -imacros include/config.h)
 _rw_size:=$(shell echo "$$(($(_rw_size_str)))")
 
+# Get RSA key size from board defines
+_rsa_size:=$(shell echo "CONFIG_RSA_KEY_SIZE" | $(CPP) $(CPPFLAGS) -P \
+            -Ichip/$(CHIP) -Iboard/$(BOARD) -imacros include/rsa.h)
+$(eval RSA_KEY_SIZE=$(_rsa_size))
+
 $(eval BOARD_$(UC_BOARD)=y)
 $(eval CHIP_$(UC_CHIP)=y)
 $(eval CHIP_VARIANT_$(UC_CHIP_VARIANT)=y)
diff --git a/Makefile.rules b/Makefile.rules
index 2c6cd7b4cd..325e6066a7 100644
--- a/Makefile.rules
+++ b/Makefile.rules
@@ -61,7 +61,7 @@ cmd_copyrw-y = cd $(out) && cp $(PROJECT).RW.flat $(PROJECT).RW.bin
 
 # commands for RSA signature
 cmd_pubkey = ./util/pem_extract_pubkey.py $(PEM) > $@
-cmd_rsasign = ./util/ec_sign_rsa.py $(PEM) $(out)/$*.bin.tmp
+cmd_rsasign = ./util/ec_sign_rsa.py --$(RSA_KEY_SIZE) $(PEM) $(out)/$*.bin.tmp
 
 # commands to build optional xref files
 cmd_deps_to_list = cat $(deps) | tr -d ':\\' | tr ' ' '\012' \
diff --git a/chip/stm32/config-stm32f03x.h b/chip/stm32/config-stm32f03x.h
index 32238c5361..a1211d38a7 100644
--- a/chip/stm32/config-stm32f03x.h
+++ b/chip/stm32/config-stm32f03x.h
@@ -32,7 +32,8 @@
  * contiguous.
  */
 #if defined(BOARD_ZINGER) || defined(BOARD_MINIMUFFIN)
-#define CONFIG_FW_PSTATE_SIZE   0
+/* Not using pstate but keep some space for the public key */
+#define CONFIG_FW_PSTATE_SIZE   544
 #else
 #define CONFIG_FW_PSTATE_SIZE   CONFIG_FLASH_BANK_SIZE
 #endif
diff --git a/common/rsa.c b/common/rsa.c
index f3ca8601fb..5a411cccef 100644
--- a/common/rsa.c
+++ b/common/rsa.c
@@ -182,11 +182,11 @@ static int check_padding(const uint8_t *sig)
 }
 
 /*
- * Verify a 2048 bit SHA256WithRSA PKCS#1 v1.5 signature against an expected
+ * Verify a SHA256WithRSA PKCS#1 v1.5 signature against an expected
  * SHA256 hash.
  *
  * @param key           RSA public key
- * @param signature     2048-bit RSA signature
+ * @param signature     RSA signature
  * @param sha           SHA-256 digest of the content to verify
  * @param workbuf32     Work buffer; caller must verify this is
  *                      3 x RSANUMWORDS elements long.
diff --git a/core/cortex-m/ec.lds.S b/core/cortex-m/ec.lds.S
index e7df0c2516..bd813a8b22 100644
--- a/core/cortex-m/ec.lds.S
+++ b/core/cortex-m/ec.lds.S
@@ -3,6 +3,7 @@
  * found in the LICENSE file.
  */
 #include "config.h"
+#include "rsa.h"
 
 #define FW_OFF_(section) CONFIG_FW_##section##_OFF
 #define FW_OFF(section) (CONFIG_FLASH_BASE + FW_OFF_(section))
@@ -18,6 +19,10 @@ MEMORY
 {
     FLASH (rx) : ORIGIN = FW_OFF(SECTION), LENGTH = FW_SIZE(SECTION)
     IRAM (rw)  : ORIGIN = CONFIG_RAM_BASE, LENGTH = CONFIG_RAM_SIZE
+#ifdef RSA_PUBLIC_KEY_SIZE
+    PSTATE(r)  : ORIGIN = FW_OFF(SECTION) + FW_SIZE(SECTION), \
+                 LENGTH = CONFIG_FW_PSTATE_SIZE
+#endif
 #ifdef CONFIG_USB
     USB_RAM (rw) : \
         ORIGIN = CONFIG_USB_RAM_BASE, \
@@ -216,12 +221,23 @@ SECTIONS
 	/* NOTHING MAY GO AFTER THIS! */
     } > IRAM
 
-    .flash_suffix  : AT(LOADADDR(.data) + SIZEOF(.data)) {
+#ifdef RSA_PUBLIC_KEY_SIZE
+#ifdef SECTION_IS_RO
+    .flash_suffix : {
         FILL(0xff);
-        /* Put the public key coefficients at the end of the partition */
-        . = ORIGIN(FLASH) + LENGTH(FLASH) - 528;
+        /*
+	 * Put the public key coefficients at the end of the partition
+	 * after the pstate bits.
+	 */
+        . = ORIGIN(PSTATE) + LENGTH(PSTATE) - RSA_PUBLIC_KEY_SIZE;
+        *(.rsa_pubkey)
+    } > PSTATE
+#else /* RW section: we don't need the RSA public key, put it anywhere */
+    .flash_suffix  : AT(LOADADDR(.data) + SIZEOF(.data)) {
         *(.rsa_pubkey)
     } > FLASH
+#endif
+#endif
 
     /* The linker won't notice if the .data section is too big to fit,
      * apparently because we're sending it into IRAM, not FLASH. The following
diff --git a/core/cortex-m0/ec.lds.S b/core/cortex-m0/ec.lds.S
index a72193a7ae..4a4ddd23f4 100644
--- a/core/cortex-m0/ec.lds.S
+++ b/core/cortex-m0/ec.lds.S
@@ -3,6 +3,7 @@
  * found in the LICENSE file.
  */
 #include "config.h"
+#include "rsa.h"
 
 #define FW_OFF_(section) CONFIG_FW_##section##_OFF
 #define FW_OFF(section) (CONFIG_FLASH_BASE + FW_OFF_(section))
@@ -18,6 +19,10 @@ MEMORY
 {
     FLASH (rx) : ORIGIN = FW_OFF(SECTION), LENGTH = FW_SIZE(SECTION)
     IRAM (rw)  : ORIGIN = CONFIG_RAM_BASE, LENGTH = CONFIG_RAM_SIZE
+#ifdef RSA_PUBLIC_KEY_SIZE
+    PSTATE(r)  : ORIGIN = FW_OFF(SECTION) + FW_SIZE(SECTION), \
+                 LENGTH = CONFIG_FW_PSTATE_SIZE
+#endif
 #ifdef CONFIG_USB
     USB_RAM (rw) : \
         ORIGIN = CONFIG_USB_RAM_BASE, \
@@ -204,12 +209,23 @@ SECTIONS
 	/* NOTHING MAY GO AFTER THIS! */
     } > IRAM
 
-    .flash_suffix  : AT(LOADADDR(.data) + SIZEOF(.data)) {
+#ifdef RSA_PUBLIC_KEY_SIZE
+#ifdef SECTION_IS_RO
+    .flash_suffix : {
         FILL(0xff);
-        /* Put the public key coefficients at the end of the partition */
-        . = ORIGIN(FLASH) + LENGTH(FLASH) - 528;
+        /*
+	 * Put the public key coefficients at the end of the partition
+	 * after the pstate bits.
+	 */
+        . = ORIGIN(PSTATE) + LENGTH(PSTATE) - RSA_PUBLIC_KEY_SIZE;
+        *(.rsa_pubkey)
+    } > PSTATE
+#else /* RW section: we don't need the RSA public key, put it anywhere */
+    .flash_suffix  : AT(LOADADDR(.data) + SIZEOF(.data)) {
         *(.rsa_pubkey)
     } > FLASH
+#endif
+#endif
 
     /* The linker won't notice if the .data section is too big to fit,
      * apparently because we're sending it into IRAM, not FLASH. The following
diff --git a/include/rsa.h b/include/rsa.h
index 95e17450c6..dbf00fa13f 100644
--- a/include/rsa.h
+++ b/include/rsa.h
@@ -6,12 +6,38 @@
 #ifndef _INCLUDE_RSA_H
 #define _INCLUDE_RSA_H
 
-#include "common.h"
+#include "config.h"
+
+#ifndef CONFIG_RSA_KEY_SIZE
+#define CONFIG_RSA_KEY_SIZE 2048 /* default to 2048-bit key length */
+#endif
 
-#define RSANUMBYTES 256  /* 2048 bit key length */
+#define RSANUMBYTES ((CONFIG_RSA_KEY_SIZE)/8)
 #define RSANUMWORDS (RSANUMBYTES / sizeof(uint32_t))
 
-/* 2048-bit RSA public key definition */
+#ifdef CONFIG_RSA /* reserve space for public key only if used */
+/*
+ * The size of the public key structure is
+ * 2 x RSANUMBYTES for n and rr fields
+ * plus 4 for n0inv, aligned on a multiple of 16
+ * Put numerical constants here to please the linker script.
+ */
+#if CONFIG_RSA_KEY_SIZE == 2048
+#define RSA_PUBLIC_KEY_SIZE 528
+#elif CONFIG_RSA_KEY_SIZE == 4096
+#define RSA_PUBLIC_KEY_SIZE 1040
+#elif CONFIG_RSA_KEY_SIZE == 8192
+#define RSA_PUBLIC_KEY_SIZE 2064
+#else
+#error Unsupported RSA key size
+#endif
+#endif /* CONFIG_RSA */
+
+#ifndef __ASSEMBLER__
+
+#include "common.h"
+
+/* RSA public key definition */
 struct rsa_public_key {
 	uint32_t n[RSANUMWORDS];  /* modulus as little endian array */
 	uint32_t rr[RSANUMWORDS]; /* R^2 as little endian array */
@@ -23,4 +49,6 @@ int rsa_verify(const struct rsa_public_key *key,
 	       const uint8_t *sha,
 	       uint32_t *workbuf32);
 
+#endif /* !__ASSEMBLER__ */
+
 #endif /* _INCLUDE_RSA_H */
diff --git a/util/ec_sign_rsa.py b/util/ec_sign_rsa.py
index 661ed37653..9c23f9edbe 100755
--- a/util/ec_sign_rsa.py
+++ b/util/ec_sign_rsa.py
@@ -18,26 +18,36 @@ import sys
 from subprocess import Popen, PIPE
 from pem_extract_pubkey import extract_pubkey
 
-# Size of a 2048-bit RSA signature
-RSANUMBYTES = 256
 # OpenSSL command to sign with SHA256andRSA
 RSA_CMD = ["openssl", "dgst", "-sha256", "-sign"]
 
-# Length reserved at the end of the RO partition for the public key
-PUBKEY_RESERVED_SPACE = 528
+# supported RSA key sizes
+RSA_KEY_SIZES=[2048, 4096, 8192]
+
+def align16(v):
+  return (v + 15) / 16 * 16
 
 def main():
   # Parse command line arguments
   if len(sys.argv) < 3:
-    sys.stderr.write("Usage: %s [--rw] <pem> <ecfile>\n" % sys.argv[0])
+    sys.stderr.write("Usage: %s [--rw] [--4096|--8192] <pem> <ecfile>\n" % sys.argv[0])
     sys.exit(-1)
   if "--rw" in sys.argv:
     sys.argv.remove("--rw")
     has_ro = False
   else:
     has_ro = True
+  # Default to a 2048-bit RSA signature
+  RSANUMBYTES = 2048 / 8
+  for sz in RSA_KEY_SIZES:
+    param = "--%d" % (sz)
+    if param in sys.argv:
+      sys.argv.remove(param)
+      RSANUMBYTES = sz / 8
   pemfile = sys.argv[1]
   ecfile = sys.argv[2]
+  # Length reserved at the end of the RO partition for the public key
+  PUBKEY_RESERVED_SPACE = align16(2 * RSANUMBYTES + 4)
 
   # Get EC firmware content
   try:
diff --git a/util/pem_extract_pubkey.py b/util/pem_extract_pubkey.py
index 739e1bea57..d368da8ee7 100755
--- a/util/pem_extract_pubkey.py
+++ b/util/pem_extract_pubkey.py
@@ -48,6 +48,9 @@ RSAPrivateKey ::= SEQUENCE {
 PEM_HEADER='-----BEGIN RSA PRIVATE KEY-----'
 PEM_FOOTER='-----END RSA PRIVATE KEY-----'
 
+# supported RSA key sizes
+RSA_KEY_SIZES=[2048, 4096, 8192]
+
 class PEMError(Exception):
   """Exception class for pem_extract_pubkey utility."""
 
@@ -130,8 +133,9 @@ def pem_get_mod(filename):
   if mod["tag"] != DER_INTEGER:
     raise PEMError('modulus field should be an integer')
   # 2048 bits + mandatory ASN.1 sign (0) => 257 Bytes
-  if mod["length"] != 257 or mod["data"][0] != '\x00':
-    raise PEMError('Invalid key length (expecting 2048 bits)')
+  modSize = (mod["length"] - 1) * 8
+  if modSize not in RSA_KEY_SIZES or mod["data"][0] != '\x00':
+    raise PEMError('Invalid key length : %d bits' % (modSize))
 
   # 3rd field is Public Exponent
   exp = seq.get_tag()
@@ -171,6 +175,7 @@ def compute_mod_parameters(modulus):
   '''
   # create an array of uint32_t to store the modulus but skip the sign byte
   w = array.array("I",modulus[1:])
+  wordCount = (len(modulus) - 1) / 4
   # all integers in DER encoded .pem file are big endian.
   w.reverse()
   w.byteswap()
@@ -183,7 +188,7 @@ def compute_mod_parameters(modulus):
   n0inv = B - modinv(w[0], B)
   # R = 2^(modulo size); RR = (R * R) % N
   RR = pow(2, 4096, N)
-  rr_words = to_words(RR, 64)
+  rr_words = to_words(RR, wordCount)
 
   return {'mod':w, 'rr':rr_words, 'n0inv':n0inv}