Merge pull request #1177 from paullouisageneau/websocket-certificate-…

…from-string Add support for loading WebSocket certificate from PEM string
paullouisageneau · May 11, 2024 · 3f65c13 · 3f65c13
2 parents 3a11fec + d3c94b7
commit 3f65c13
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 17 deletions.
diff --git a/src/impl/websocket.cpp b/src/impl/websocket.cpp
@@ -34,11 +34,28 @@ using namespace std::placeholders;
 using namespace std::chrono_literals;
 using std::chrono::milliseconds;
 
+const string PemBeginCertificateTag = "-----BEGIN CERTIFICATE-----";
+
 WebSocket::WebSocket(optional<Configuration> optConfig, certificate_ptr certificate)
     : config(optConfig ? std::move(*optConfig) : Configuration()),
-      mCertificate(certificate ? std::move(certificate) : std::move(loadCertificate(config))),
-      mIsSecure(mCertificate != nullptr), mRecvQueue(RECV_QUEUE_LIMIT, message_size_func) {
+      mRecvQueue(RECV_QUEUE_LIMIT, message_size_func) {
 	PLOG_VERBOSE << "Creating WebSocket";
+
+	if (certificate) {
+		mCertificate = std::move(certificate);
+	} else if (config.certificatePemFile && config.keyPemFile) {
+		mCertificate = std::make_shared<Certificate>(
+		    config.certificatePemFile->find(PemBeginCertificateTag) != string::npos
+		        ? Certificate::FromString(*config.certificatePemFile, *config.keyPemFile)
+		        : Certificate::FromFile(*config.certificatePemFile, *config.keyPemFile,
+		                                config.keyPemPass.value_or("")));
+	} else if (config.certificatePemFile || config.keyPemFile) {
+		throw std::invalid_argument(
+		    "Either none or both certificate and key PEM files must be specified");
+	}
+
+	mIsSecure = mCertificate != nullptr;
+
 	if (config.proxyServer) {
 		if (config.proxyServer->type == ProxyServer::Type::Socks5)
 			throw std::invalid_argument(
@@ -49,19 +66,6 @@ WebSocket::WebSocket(optional<Configuration> optConfig, certificate_ptr certific
 	}
 }
 
-certificate_ptr WebSocket::loadCertificate(const Configuration& config) {
-	if (!config.certificatePemFile)
-		return nullptr;
-
-	if (config.keyPemFile)
-		return std::make_shared<Certificate>(
-			Certificate::FromFile(*config.certificatePemFile, *config.keyPemFile,
-										config.keyPemPass.value_or("")));
-
-	throw std::invalid_argument(
-		"Either none or both certificate and key PEM files must be specified");
-}
-
 WebSocket::~WebSocket() { PLOG_VERBOSE << "Destroying WebSocket"; }
 
 void WebSocket::open(const string &url) {
@@ -156,7 +160,9 @@ bool WebSocket::isOpen() const { return state == State::Open; }
 
 bool WebSocket::isClosed() const { return state == State::Closed; }
 
-size_t WebSocket::maxMessageSize() const { return config.maxMessageSize.value_or(DEFAULT_WS_MAX_MESSAGE_SIZE); }
+size_t WebSocket::maxMessageSize() const {
+	return config.maxMessageSize.value_or(DEFAULT_WS_MAX_MESSAGE_SIZE);
+}
 
 optional<message_variant> WebSocket::receive() {
 	auto next = mRecvQueue.pop();

diff --git a/src/impl/websocket.hpp b/src/impl/websocket.hpp
@@ -73,7 +73,7 @@ struct WebSocket final : public Channel, public std::enable_shared_from_this<Web
 
 	const init_token mInitToken = Init::Instance().token();
 
-	const certificate_ptr mCertificate;
+	certificate_ptr mCertificate;
 	bool mIsSecure;
 
 	optional<string> mHostname; // for TLS SNI and Proxy