Implement asynchronous private key operations for client auth.

This adds a new API, SSL_set_private_key_method, which allows the consumer to
customize private key operations. For simplicity, it is incompatible with the
multiple slots feature (which will hopefully go away) but does not, for now,
break it.

The new method is only routed up for the client for now. The server will
require a decrypt hook as well for the plain RSA key exchange.

BUG=347404

Change-Id: I35d69095c29134c34c2af88c613ad557d6957614
Reviewed-on: https://boringssl-review.googlesource.com/5049
Reviewed-by: Adam Langley <agl@google.com>
diff --git a/ssl/test/bssl_shim.cc b/ssl/test/bssl_shim.cc
index 3b95d7e..f4ae982 100644
--- a/ssl/test/bssl_shim.cc
+++ b/ssl/test/bssl_shim.cc
@@ -90,6 +90,12 @@
   ScopedSSL_SESSION pending_session;
   bool early_callback_called = false;
   bool handshake_done = false;
+  // private_key is the underlying private key used when testing custom keys.
+  ScopedEVP_PKEY private_key;
+  std::vector<uint8_t> signature;
+  // signature_retries is the number of times an asynchronous sign operation has
+  // been retried.
+  unsigned signature_retries = 0;
 };
 
 static void TestStateExFree(void *parent, void *ptr, CRYPTO_EX_DATA *ad,
@@ -129,12 +135,100 @@
   return pkey;
 }
 
+static int AsyncPrivateKeyType(SSL *ssl) {
+  return EVP_PKEY_id(GetTestState(ssl)->private_key.get());
+}
+
+static int AsyncPrivateKeySupportsDigest(SSL *ssl, const EVP_MD *md) {
+  return EVP_PKEY_supports_digest(GetTestState(ssl)->private_key.get(), md);
+}
+
+static size_t AsyncPrivateKeyMaxSignatureLen(SSL *ssl) {
+  return EVP_PKEY_size(GetTestState(ssl)->private_key.get());
+}
+
+static ssl_private_key_result_t AsyncPrivateKeySign(
+    SSL *ssl, uint8_t *out, size_t *out_len, size_t max_out,
+    const EVP_MD *md, const uint8_t *in, size_t in_len) {
+  TestState *test_state = GetTestState(ssl);
+  if (!test_state->signature.empty()) {
+    fprintf(stderr, "AsyncPrivateKeySign called with operation pending.\n");
+    abort();
+  }
+
+  ScopedEVP_PKEY_CTX ctx(EVP_PKEY_CTX_new(test_state->private_key.get(),
+                                          nullptr));
+  if (!ctx) {
+    return ssl_private_key_failure;
+  }
+
+  // Write the signature into |test_state|.
+  size_t len = 0;
+  if (!EVP_PKEY_sign_init(ctx.get()) ||
+      !EVP_PKEY_CTX_set_signature_md(ctx.get(), md) ||
+      !EVP_PKEY_sign(ctx.get(), nullptr, &len, in, in_len)) {
+    return ssl_private_key_failure;
+  }
+  test_state->signature.resize(len);
+  if (!EVP_PKEY_sign(ctx.get(), bssl::vector_data(&test_state->signature), &len,
+                     in, in_len)) {
+    return ssl_private_key_failure;
+  }
+  test_state->signature.resize(len);
+
+  // The signature will be released asynchronously in |AsyncPrivateKeySignComplete|.
+  return ssl_private_key_retry;
+}
+
+static ssl_private_key_result_t AsyncPrivateKeySignComplete(
+    SSL *ssl, uint8_t *out, size_t *out_len, size_t max_out) {
+  TestState *test_state = GetTestState(ssl);
+  if (test_state->signature.empty()) {
+    fprintf(stderr,
+            "AsyncPrivateKeySignComplete called without operation pending.\n");
+    abort();
+  }
+
+  if (test_state->signature_retries < 2) {
+    // Only return the signature on the second attempt, to test both incomplete
+    // |sign| and |sign_complete|.
+    return ssl_private_key_retry;
+  }
+
+  if (max_out < test_state->signature.size()) {
+    fprintf(stderr, "Output buffer too small.\n");
+    return ssl_private_key_failure;
+  }
+  memcpy(out, bssl::vector_data(&test_state->signature),
+         test_state->signature.size());
+
+  test_state->signature.clear();
+  test_state->signature_retries = 0;
+  return ssl_private_key_success;
+}
+
+static const SSL_PRIVATE_KEY_METHOD g_async_private_key_method = {
+    AsyncPrivateKeyType,
+    AsyncPrivateKeySupportsDigest,
+    AsyncPrivateKeyMaxSignatureLen,
+    AsyncPrivateKeySign,
+    AsyncPrivateKeySignComplete,
+};
+
 static bool InstallCertificate(SSL *ssl) {
   const TestConfig *config = GetConfigPtr(ssl);
-  if (!config->key_file.empty() &&
-      !SSL_use_PrivateKey_file(ssl, config->key_file.c_str(),
-                               SSL_FILETYPE_PEM)) {
-    return false;
+  TestState *test_state = GetTestState(ssl);
+  if (!config->key_file.empty()) {
+    if (config->use_async_private_key) {
+      test_state->private_key = LoadPrivateKey(config->key_file.c_str());
+      if (!test_state->private_key) {
+        return false;
+      }
+      SSL_set_private_key_method(ssl, &g_async_private_key_method);
+    } else if (!SSL_use_PrivateKey_file(ssl, config->key_file.c_str(),
+                                        SSL_FILETYPE_PEM)) {
+      return false;
+    }
   }
   if (!config->cert_file.empty() &&
       !SSL_use_certificate_file(ssl, config->cert_file.c_str(),
@@ -500,6 +594,9 @@
     case SSL_ERROR_PENDING_CERTIFICATE:
       // The handshake will resume without a second call to the early callback.
       return InstallCertificate(ssl);
+    case SSL_ERROR_WANT_PRIVATE_KEY_OPERATION:
+      test_state->signature_retries++;
+      return true;
     default:
       return false;
   }
diff --git a/ssl/test/runner/runner.go b/ssl/test/runner/runner.go
index b0eef42..1186313 100644
--- a/ssl/test/runner/runner.go
+++ b/ssl/test/runner/runner.go
@@ -2202,6 +2202,20 @@
 			"-key-file", path.Join(*resourceDir, rsaKeyFile),
 		},
 	})
+	if async {
+		tests = append(tests, testCase{
+			testType: clientTest,
+			name:     "ClientAuth-Client-AsyncKey",
+			config: Config{
+				ClientAuth: RequireAnyClientCert,
+			},
+			flags: []string{
+				"-cert-file", rsaCertificateFile,
+				"-key-file", rsaKeyFile,
+				"-use-async-private-key",
+			},
+		})
+	}
 	tests = append(tests, testCase{
 		testType: serverTest,
 		name:     "ClientAuth-Server",
diff --git a/ssl/test/test_config.cc b/ssl/test/test_config.cc
index 363b6f3..031ad93 100644
--- a/ssl/test/test_config.cc
+++ b/ssl/test/test_config.cc
@@ -82,6 +82,7 @@
   { "-reject-peer-renegotiations", &TestConfig::reject_peer_renegotiations },
   { "-no-legacy-server-connect", &TestConfig::no_legacy_server_connect },
   { "-tls-unique", &TestConfig::tls_unique },
+  { "-use-async-private-key", &TestConfig::use_async_private_key },
 };
 
 const Flag<std::string> kStringFlags[] = {
diff --git a/ssl/test/test_config.h b/ssl/test/test_config.h
index 5d753c8..e9af0de 100644
--- a/ssl/test/test_config.h
+++ b/ssl/test/test_config.h
@@ -79,6 +79,7 @@
   bool reject_peer_renegotiations = false;
   bool no_legacy_server_connect = false;
   bool tls_unique = false;
+  bool use_async_private_key = false;
 };
 
 bool ParseConfig(int argc, char **argv, TestConfig *out_config);