From 273d6d14cc9ea39398b763f08b0dda9c0789232a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20Vejpustek?= <ondrej.vejpustek@satoshilabs.com>
Date: Wed, 20 Oct 2021 15:34:18 +0200
Subject: [PATCH] refactor(crypto): propagate error from
 zkp_context_acquire_writable()

---
 crypto/zkp_bip340.c  | 36 ++++++++++++++++++++++++------------
 crypto/zkp_context.c |  6 +++++-
 crypto/zkp_ecdsa.c   | 40 ++++++++++++++++++++++++++--------------
 3 files changed, 55 insertions(+), 27 deletions(-)

diff --git a/crypto/zkp_bip340.c b/crypto/zkp_bip340.c
index 8ee126e08..d0fdab885 100644
--- a/crypto/zkp_bip340.c
+++ b/crypto/zkp_bip340.c
@@ -44,12 +44,16 @@ int zkp_bip340_get_public_key(const uint8_t *private_key_bytes,
 
   if (result == 0) {
     secp256k1_context *context_writable = zkp_context_acquire_writable();
-    secp256k1_context_writable_randomize(context_writable);
-    if (secp256k1_ec_pubkey_create(context_writable, &pubkey,
-                                   private_key_bytes) != 1) {
+    if (context_writable) {
+      secp256k1_context_writable_randomize(context_writable);
+      if (secp256k1_ec_pubkey_create(context_writable, &pubkey,
+                                     private_key_bytes) != 1) {
+        result = -1;
+      }
+      zkp_context_release_writable();
+    } else {
       result = -1;
     }
-    zkp_context_release_writable();
   }
 
   secp256k1_xonly_pubkey xonly_pubkey = {0};
@@ -91,22 +95,30 @@ int zkp_bip340_sign_digest(const uint8_t *private_key_bytes,
 
   if (result == 0) {
     secp256k1_context *context_writable = zkp_context_acquire_writable();
-    secp256k1_context_writable_randomize(context_writable);
-    if (secp256k1_keypair_create(context_writable, &keypair,
-                                 private_key_bytes) != 1) {
+    if (context_writable) {
+      secp256k1_context_writable_randomize(context_writable);
+      if (secp256k1_keypair_create(context_writable, &keypair,
+                                   private_key_bytes) != 1) {
+        result = -1;
+      }
+      zkp_context_release_writable();
+    } else {
       result = -1;
     }
-    zkp_context_release_writable();
   }
 
   if (result == 0) {
     secp256k1_context *context_writable = zkp_context_acquire_writable();
-    secp256k1_context_writable_randomize(context_writable);
-    if (secp256k1_schnorrsig_sign(context_writable, signature_bytes, digest,
-                                  &keypair, NULL, auxiliary_data) != 1) {
+    if (context_writable) {
+      secp256k1_context_writable_randomize(context_writable);
+      if (secp256k1_schnorrsig_sign(context_writable, signature_bytes, digest,
+                                    &keypair, NULL, auxiliary_data) != 1) {
+        result = -1;
+      }
+      zkp_context_release_writable();
+    } else {
       result = -1;
     }
-    zkp_context_release_writable();
   }
 
   memzero(&keypair, sizeof(keypair));
diff --git a/crypto/zkp_context.c b/crypto/zkp_context.c
index 0208817ae..022d7e7ac 100644
--- a/crypto/zkp_context.c
+++ b/crypto/zkp_context.c
@@ -81,11 +81,15 @@ const secp256k1_context *zkp_context_get_read_only() {
   return context;
 }
 
+// returns NULL if context cannot be acquired
 secp256k1_context *zkp_context_acquire_writable() {
   assert(initialized == true);
 
   // We don't expect the context to be used by multiple threads
-  assert(!atomic_flag_test_and_set(&locked));
+  if (atomic_flag_test_and_set(&locked)) {
+    return NULL;
+  }
+
   return context;
 }
 
diff --git a/crypto/zkp_ecdsa.c b/crypto/zkp_ecdsa.c
index 7c5676b0e..7c20a66eb 100644
--- a/crypto/zkp_ecdsa.c
+++ b/crypto/zkp_ecdsa.c
@@ -55,12 +55,16 @@ int zkp_ecdsa_get_public_key33(const ecdsa_curve *curve,
 
   if (result == 0) {
     secp256k1_context *context_writable = zkp_context_acquire_writable();
-    secp256k1_context_writable_randomize(context_writable);
-    if (secp256k1_ec_pubkey_create(context_writable, &public_key,
-                                   private_key_bytes) != 1) {
+    if (context_writable) {
+      secp256k1_context_writable_randomize(context_writable);
+      if (secp256k1_ec_pubkey_create(context_writable, &public_key,
+                                     private_key_bytes) != 1) {
+        result = 1;
+      }
+      zkp_context_release_writable();
+    } else {
       result = 1;
     }
-    zkp_context_release_writable();
   }
 
   if (result == 0) {
@@ -94,12 +98,16 @@ int zkp_ecdsa_get_public_key65(const ecdsa_curve *curve,
 
   if (result == 0) {
     secp256k1_context *context_writable = zkp_context_acquire_writable();
-    secp256k1_context_writable_randomize(context_writable);
-    if (secp256k1_ec_pubkey_create(context_writable, &public_key,
-                                   private_key_bytes) != 1) {
+    if (context_writable) {
+      secp256k1_context_writable_randomize(context_writable);
+      if (secp256k1_ec_pubkey_create(context_writable, &public_key,
+                                     private_key_bytes) != 1) {
+        result = 1;
+      }
+      zkp_context_release_writable();
+    } else {
       result = 1;
     }
-    zkp_context_release_writable();
   }
 
   if (result == 0) {
@@ -147,14 +155,18 @@ int zkp_ecdsa_sign_digest(
   secp256k1_ecdsa_recoverable_signature recoverable_signature = {0};
 
   if (result == 0) {
-    secp256k1_context *ctx_writable = zkp_context_acquire_writable();
-    secp256k1_context_writable_randomize(ctx_writable);
-    if (secp256k1_ecdsa_sign_recoverable(ctx_writable, &recoverable_signature,
-                                         digest, private_key_bytes, NULL,
-                                         NULL) != 1) {
+    secp256k1_context *context_writable = zkp_context_acquire_writable();
+    if (context_writable) {
+      secp256k1_context_writable_randomize(context_writable);
+      if (secp256k1_ecdsa_sign_recoverable(
+              context_writable, &recoverable_signature, digest,
+              private_key_bytes, NULL, NULL) != 1) {
+        result = 1;
+      }
+      zkp_context_release_writable();
+    } else {
       result = 1;
     }
-    zkp_context_release_writable();
   }
 
   if (result == 0) {