From dd78f5291375a10117b4fa3d13d0b3a51fd8b0a5 Mon Sep 17 00:00:00 2001
From: Stefan Perovsek <stefan.perovsek@devolo.de>
Date: Mon, 7 Jan 2019 13:26:10 +0100
Subject: [PATCH] add SSH psk client / server callback feature

---
 .../asio/ssl/detail/psk_client_callback.hpp   | 65 +++++++++++++++++++
 .../asio/ssl/detail/psk_server_callback.hpp   | 64 ++++++++++++++++++
 include/asio/ssl/dtls/context.hpp             | 25 +++++++
 include/asio/ssl/dtls/impl/context.hpp        | 37 +++++++++++
 include/asio/ssl/dtls/impl/context.ipp        | 28 ++++++++
 5 files changed, 219 insertions(+)
 create mode 100644 include/asio/ssl/detail/psk_client_callback.hpp
 create mode 100644 include/asio/ssl/detail/psk_server_callback.hpp

diff --git a/include/asio/ssl/detail/psk_client_callback.hpp b/include/asio/ssl/detail/psk_client_callback.hpp
new file mode 100644
index 0000000..fb1d1de
--- /dev/null
+++ b/include/asio/ssl/detail/psk_client_callback.hpp
@@ -0,0 +1,65 @@
+#ifndef ASIO_SSL_DETAIL_PSK_CLIENT_CALLBACK_HPP
+#define ASIO_SSL_DETAIL_PSK_CLIENT_CALLBACK_HPP
+
+#include "asio/detail/config.hpp"
+
+#include <cstddef>
+#include <string>
+#include "asio/ssl/context_base.hpp"
+
+#include "asio/detail/push_options.hpp"
+
+namespace asio {
+namespace ssl {
+namespace detail {
+
+class psk_client_callback_base
+{
+public:
+  virtual ~psk_client_callback_base()
+  {
+  }
+
+  virtual bool call(SSL *ssl, const char *hint, char *identity, unsigned int max_identity_len, unsigned char *psk, unsigned int max_psk_len) = 0;
+  virtual unsigned int result() const = 0;
+  virtual psk_client_callback_base* clone() = 0;
+};
+
+template <typename PskClientCallback>
+class psk_client_callback : public psk_client_callback_base
+{
+public:
+  explicit psk_client_callback(PskClientCallback callback)
+    : _callback(callback), _retval(0)
+  {
+  }
+
+  virtual bool call(SSL *ssl, const char *hint, char *identity, unsigned int max_identity_len, unsigned char *psk, unsigned int max_psk_len)
+  {
+    _retval = _callback(ssl, hint, identity, max_identity_len, psk, max_psk_len);
+    return true;
+  }
+
+  virtual unsigned int result() const
+  {
+    return _retval;
+  }
+
+  virtual psk_client_callback_base* clone()
+  {
+    return new psk_client_callback<PskClientCallback>(_callback);
+  }
+
+private:
+  PskClientCallback _callback;
+  unsigned int      _retval;
+};
+
+} // namespace detail
+} // namespace ssl
+} // namespace asio
+
+
+#include "asio/detail/pop_options.hpp"
+
+#endif // ASIO_SSL_DETAIL_PSK_CLIENT_CALLBACK_HPP
diff --git a/include/asio/ssl/detail/psk_server_callback.hpp b/include/asio/ssl/detail/psk_server_callback.hpp
new file mode 100644
index 0000000..161e2ae
--- /dev/null
+++ b/include/asio/ssl/detail/psk_server_callback.hpp
@@ -0,0 +1,64 @@
+#ifndef ASIO_SSL_DETAIL_PSK_SERVER_CALLBACK_HPP
+#define ASIO_SSL_DETAIL_PSK_SERVER_CALLBACK_HPP
+
+#include "asio/detail/config.hpp"
+
+#include <cstddef>
+#include <string>
+#include "asio/ssl/context_base.hpp"
+
+#include "asio/detail/push_options.hpp"
+
+namespace asio {
+namespace ssl {
+namespace detail {
+
+class psk_server_callback_base
+{
+public:
+  virtual ~psk_server_callback_base()
+  {
+  }
+
+  virtual bool call(SSL *ssl, const char *identity, unsigned char *psk, unsigned int max_psk_len) = 0;
+  virtual unsigned int result() const = 0;
+  virtual psk_server_callback_base* clone() = 0;
+};
+
+template <typename PskServerCallback>
+class psk_server_callback : public psk_server_callback_base
+{
+public:
+  explicit psk_server_callback(PskServerCallback callback)
+    : _callback(callback), _retval(0)
+  {
+  }
+
+  virtual bool call(SSL *ssl, const char *identity, unsigned char *psk, unsigned int max_psk_len)
+  {
+    return _callback(ssl, identity, psk, max_psk_len);
+  }
+
+  virtual unsigned int result() const
+  {
+    return _retval;
+  }
+
+  virtual psk_server_callback_base* clone()
+  {
+    return new psk_server_callback<PskServerCallback>(_callback);
+  }
+
+private:
+  PskServerCallback _callback;
+  unsigned int      _retval;
+};
+
+} // namespace detail
+} // namespace ssl
+} // namespace asio
+
+
+#include "asio/detail/pop_options.hpp"
+
+#endif // ASIO_SSL_DETAIL_PSK_SERVER_CALLBACK_HPP
diff --git a/include/asio/ssl/dtls/context.hpp b/include/asio/ssl/dtls/context.hpp
index 4938463..ddc4279 100644
--- a/include/asio/ssl/dtls/context.hpp
+++ b/include/asio/ssl/dtls/context.hpp
@@ -24,6 +24,8 @@
 #include "asio/ssl/detail/openssl_types.hpp"
 #include "asio/ssl/detail/openssl_init.hpp"
 #include "asio/ssl/detail/password_callback.hpp"
+#include "asio/ssl/detail/psk_client_callback.hpp"
+#include "asio/ssl/detail/psk_server_callback.hpp"
 #include "asio/ssl/detail/verify_callback.hpp"
 #include "asio/ssl/verify_mode.hpp"
 
@@ -741,6 +743,21 @@ class context
   ASIO_SYNC_OP_VOID set_password_callback(PasswordCallback callback,
       asio::error_code& ec);
 
+
+  template <typename PskClientCallback>
+  void set_psk_client_callback(PskClientCallback callback);
+
+  template <typename PskClientCallback>
+  ASIO_SYNC_OP_VOID set_psk_client_callback(PskClientCallback callback,
+      asio::error_code& ec);
+
+  template <typename PskServerCallback>
+  void set_psk_server_callback(PskServerCallback callback);
+
+  template <typename PskServerCallback>
+  ASIO_SYNC_OP_VOID set_psk_server_callback(PskServerCallback callback,
+      asio::error_code& ec);
+
 private:
   struct bio_cleanup;
   struct x509_cleanup;
@@ -764,6 +781,14 @@ class context
   ASIO_DECL static int password_callback_function(
       char* buf, int size, int purpose, void* data);
 
+  // Helper function used to set a password callback.
+  ASIO_DECL ASIO_SYNC_OP_VOID do_set_psk_client_callback(
+      ssl::detail::psk_client_callback_base* callback, asio::error_code& ec);
+
+  // Helper function used to set a password callback.
+  ASIO_DECL ASIO_SYNC_OP_VOID do_set_psk_server_callback(
+      ssl::detail::psk_server_callback_base* callback, asio::error_code& ec);
+
   // Helper function to set the temporary Diffie-Hellman parameters from a BIO.
   ASIO_DECL ASIO_SYNC_OP_VOID do_use_tmp_dh(
       BIO* bio, asio::error_code& ec);
diff --git a/include/asio/ssl/dtls/impl/context.hpp b/include/asio/ssl/dtls/impl/context.hpp
index 3f40a95..21cd5ea 100644
--- a/include/asio/ssl/dtls/impl/context.hpp
+++ b/include/asio/ssl/dtls/impl/context.hpp
@@ -60,6 +60,43 @@ ASIO_SYNC_OP_VOID context::set_password_callback(
   ASIO_SYNC_OP_VOID_RETURN(ec);
 }
 
+template <typename PskClientCallback>
+void context::set_psk_client_callback(PskClientCallback callback)
+{
+  asio::error_code ec;
+  this->set_psk_client_callback(callback, ec);
+  asio::detail::throw_error(ec, "set_psk_client_callback");
+}
+
+template <typename PskClientCallback>
+ASIO_SYNC_OP_VOID context::set_psk_client_callback(
+    PskClientCallback callback, asio::error_code& ec)
+{
+  do_set_psk_client_callback(
+      new detail::psk_client_callback<PskClientCallback>(callback), ec);
+  ASIO_SYNC_OP_VOID_RETURN(ec);
+}
+
+/*
+template <typename PskServerCallback>
+ASIO_SYNC_OP_VOID void context::set_psk_server_callback(PskServerCallback callback)
+{
+  asio::error_code ec;
+  this->set_psk_server_callback(callback, ec);
+  asio::detail::throw_error(ec, "set_psk_server_callback");
+}
+
+template <typename PskServerCallback>
+ASIO_SYNC_OP_VOID context::set_psk_server_callback(
+    PskServerCallback callback, asio::error_code& ec)
+{
+  do_set_psk_server_callback(
+      new detail::psk_server_callback<PskServerCallback>(callback), ec);
+  ASIO_SYNC_OP_VOID_RETURN(ec);
+}
+*/
+
+
 } // namespace dtls
 } // namespace ssl
 } // namespace asio
diff --git a/include/asio/ssl/dtls/impl/context.ipp b/include/asio/ssl/dtls/impl/context.ipp
index 2586fff..a5d0bcc 100644
--- a/include/asio/ssl/dtls/impl/context.ipp
+++ b/include/asio/ssl/dtls/impl/context.ipp
@@ -953,6 +953,34 @@ int context::password_callback_function(
   return 0;
 }
 
+
+ASIO_SYNC_OP_VOID context::do_set_psk_client_callback(
+    ssl::detail::psk_client_callback_base* callback, asio::error_code& ec)
+{
+  SSL_CTX_set_cipher_list(handle_, "PSK");
+  static ssl::detail::psk_client_callback_base* cb = callback;
+
+  SSL_CTX_set_psk_client_callback(handle_, [](SSL *ssl, const char *hint, char *identity, unsigned int max_identity_len, unsigned char *psk, unsigned int max_psk_len) -> unsigned int {
+        if ( cb ) { cb->call(ssl, hint, identity, max_identity_len, psk, max_psk_len); return cb->result(); } return 0; });
+
+  ec = asio::error_code();
+  ASIO_SYNC_OP_VOID_RETURN(ec);
+}
+
+ASIO_SYNC_OP_VOID context::do_set_psk_server_callback(
+    ssl::detail::psk_server_callback_base* callback, asio::error_code& ec)
+{
+  SSL_CTX_set_cipher_list(handle_, "PSK");
+  static ssl::detail::psk_server_callback_base* cb = callback;
+
+  SSL_CTX_set_psk_server_callback(handle_, [](SSL *ssl, const char *identity, unsigned char *psk, unsigned int max_psk_len) -> unsigned int {
+        if ( cb) { cb->call(ssl, identity, psk, max_psk_len); return cb->result(); } return 0; });
+
+  ec = asio::error_code();
+  ASIO_SYNC_OP_VOID_RETURN(ec);
+}
+
+
 BIO* context::make_buffer_bio(const const_buffer& b)
 {
   return ::BIO_new_mem_buf(