From 9b486e22588ad91dbcd61f12dfe1e451b03cd44e Mon Sep 17 00:00:00 2001
From: Brian White <mscdex@mscdex.net>
Date: Sat, 29 Jun 2024 06:16:58 -0400
Subject: [PATCH] tls: add setKeyCert() to tls.Socket

---
 doc/api/tls.md                              | 14 ++++++
 lib/_tls_wrap.js                            | 11 ++++
 src/crypto/crypto_tls.cc                    | 28 +++++++++++
 src/crypto/crypto_tls.h                     |  1 +
 test/parallel/test-tls-server-setkeycert.js | 56 +++++++++++++++++++++
 5 files changed, 110 insertions(+)
 create mode 100644 test/parallel/test-tls-server-setkeycert.js

diff --git a/doc/api/tls.md b/doc/api/tls.md
index ed9c3a489be847..196c0525fd509b 100644
--- a/doc/api/tls.md
+++ b/doc/api/tls.md
@@ -1533,6 +1533,20 @@ When running as the server, the socket will be destroyed with an error after
 For TLSv1.3, renegotiation cannot be initiated, it is not supported by the
 protocol.
 
+### `tlsSocket.setKeyCert(context)`
+
+<!-- YAML
+added: REPLACEME
+-->
+
+* `context` {Object|tls.SecureContext} An object containing at least `key` and
+  `cert` properties from the [`tls.createSecureContext()`][] `options`, or a
+  TLS context object created with [`tls.createSecureContext()`][] itself.
+
+The `tlsSocket.setKeyCert()` method sets the private key and certificate to use
+for the socket. This is mainly useful if you wish to select a server certificate
+from a TLS server's `ALPNCallback`.
+
 ### `tlsSocket.setMaxSendFragment(size)`
 
 <!-- YAML
diff --git a/lib/_tls_wrap.js b/lib/_tls_wrap.js
index cc79b30aacac7d..dfa950ff3a601c 100644
--- a/lib/_tls_wrap.js
+++ b/lib/_tls_wrap.js
@@ -1153,6 +1153,17 @@ TLSSocket.prototype.getX509Certificate = function() {
   return cert ? new InternalX509Certificate(cert) : undefined;
 };
 
+TLSSocket.prototype.setKeyCert = function(context) {
+  if (this._handle) {
+    let secureContext;
+    if (context instanceof common.SecureContext)
+      secureContext = context;
+    else
+      secureContext = tls.createSecureContext(context);
+    this._handle.setKeyCert(secureContext.context);
+  }
+};
+
 // Proxy TLSSocket handle methods
 function makeSocketMethodProxy(name) {
   return function socketMethodProxy(...args) {
diff --git a/src/crypto/crypto_tls.cc b/src/crypto/crypto_tls.cc
index 34382af9026781..9fb567f89c11d4 100644
--- a/src/crypto/crypto_tls.cc
+++ b/src/crypto/crypto_tls.cc
@@ -1595,6 +1595,33 @@ void TLSWrap::SetALPNProtocols(const FunctionCallbackInfo<Value>& args) {
   }
 }
 
+void TLSWrap::SetKeyCert(const FunctionCallbackInfo<Value>& args) {
+  TLSWrap* w;
+  ASSIGN_OR_RETURN_UNWRAP(&w, args.This());
+  Environment* env = w->env();
+
+  if (w->is_client()) return;
+
+  if (args.Length() < 1 || !args[0]->IsObject())
+    return env->ThrowTypeError("Must give a SecureContext as first argument");
+
+  Local<Value> ctx = args[0];
+  if (UNLIKELY(ctx.IsEmpty())) return;
+
+  Local<FunctionTemplate> cons = env->secure_context_constructor_template();
+  if (cons->HasInstance(ctx)) {
+    SecureContext* sc = Unwrap<SecureContext>(ctx.As<Object>());
+    CHECK_NOT_NULL(sc);
+    if (!UseSNIContext(w->ssl_, BaseObjectPtr<SecureContext>(sc)) ||
+        !w->SetCACerts(sc)) {
+      unsigned long err = ERR_get_error();  // NOLINT(runtime/int)
+      return ThrowCryptoError(env, err, "SetKeyCert");
+    }
+  } else {
+    return env->ThrowTypeError("Must give a SecureContext as first argument");
+  }
+}
+
 void TLSWrap::GetPeerCertificate(const FunctionCallbackInfo<Value>& args) {
   TLSWrap* w;
   ASSIGN_OR_RETURN_UNWRAP(&w, args.This());
@@ -2130,6 +2157,7 @@ void TLSWrap::Initialize(
   SetProtoMethod(isolate, t, "renegotiate", Renegotiate);
   SetProtoMethod(isolate, t, "requestOCSP", RequestOCSP);
   SetProtoMethod(isolate, t, "setALPNProtocols", SetALPNProtocols);
+  SetProtoMethod(isolate, t, "setKeyCert", SetKeyCert);
   SetProtoMethod(isolate, t, "setOCSPResponse", SetOCSPResponse);
   SetProtoMethod(isolate, t, "setServername", SetServername);
   SetProtoMethod(isolate, t, "setSession", SetSession);
diff --git a/src/crypto/crypto_tls.h b/src/crypto/crypto_tls.h
index b65cea22255248..1faad67a4a2886 100644
--- a/src/crypto/crypto_tls.h
+++ b/src/crypto/crypto_tls.h
@@ -213,6 +213,7 @@ class TLSWrap : public AsyncWrap,
   static void Renegotiate(const v8::FunctionCallbackInfo<v8::Value>& args);
   static void RequestOCSP(const v8::FunctionCallbackInfo<v8::Value>& args);
   static void SetALPNProtocols(const v8::FunctionCallbackInfo<v8::Value>& args);
+  static void SetKeyCert(const v8::FunctionCallbackInfo<v8::Value>& args);
   static void SetOCSPResponse(const v8::FunctionCallbackInfo<v8::Value>& args);
   static void SetServername(const v8::FunctionCallbackInfo<v8::Value>& args);
   static void SetSession(const v8::FunctionCallbackInfo<v8::Value>& args);
diff --git a/test/parallel/test-tls-server-setkeycert.js b/test/parallel/test-tls-server-setkeycert.js
new file mode 100644
index 00000000000000..2d47c5275e02ca
--- /dev/null
+++ b/test/parallel/test-tls-server-setkeycert.js
@@ -0,0 +1,56 @@
+'use strict';
+const common = require('../common');
+
+if (!common.hasCrypto)
+  common.skip('missing crypto');
+
+const assert = require('assert');
+const { X509Certificate } = require('crypto');
+const tls = require('tls');
+const fixtures = require('../common/fixtures');
+
+const altKeyCert = {
+  key: fixtures.readKey('agent2-key.pem'),
+  cert: fixtures.readKey('agent2-cert.pem'),
+  minVersion: 'TLSv1.2',
+};
+
+const altKeyCertVals = [
+  altKeyCert,
+  tls.createSecureContext(altKeyCert),
+];
+
+(function next() {
+  if (!altKeyCertVals.length)
+    return;
+  const altKeyCertVal = altKeyCertVals.shift();
+  console.log('trying', altKeyCertVal);
+  const options = {
+    key: fixtures.readKey('agent1-key.pem'),
+    cert: fixtures.readKey('agent1-cert.pem'),
+    minVersion: 'TLSv1.3',
+    ALPNCallback: common.mustCall(function({ servername, protocols }) {
+      this.setKeyCert(altKeyCertVal);
+      assert.deepStrictEqual(protocols, ['acme-tls/1']);
+      return protocols[0];
+    }),
+  };
+
+  tls.createServer(options, (s) => s.end()).listen(0, function() {
+    this.on('connection', common.mustCall((socket) => this.close()));
+
+    tls.connect({
+      port: this.address().port,
+      rejectUnauthorized: false,
+      ALPNProtocols: ['acme-tls/1'],
+    }, common.mustCall(function() {
+      const altCert = new X509Certificate(altKeyCert.cert);
+      assert.strictEqual(
+        this.getPeerX509Certificate().raw.equals(altCert.raw),
+        true
+      );
+      this.end();
+      next();
+    }));
+  });
+})();