From 2aae0b01ecb266f7da7068d373689fb6799cdf9d Mon Sep 17 00:00:00 2001 From: chendongsheng Date: Thu, 22 Apr 2021 21:21:51 +0800 Subject: [PATCH] tcp support ssl --- mindspore/ccsrc/pipeline/jit/init.cc | 3 +- mindspore/ccsrc/ps/constants.h | 5 ++- .../ps/core/communicator/http_communicator.cc | 1 - .../ps/core/communicator/http_communicator.h | 5 ++- .../core/communicator/http_request_handler.cc | 40 ++++++++++--------- .../core/communicator/http_request_handler.h | 1 + .../ccsrc/ps/core/communicator/ssl_wrapper.cc | 14 ++++++- .../ccsrc/ps/core/communicator/ssl_wrapper.h | 3 +- .../ccsrc/ps/core/communicator/tcp_client.cc | 27 ++++++++++++- .../ccsrc/ps/core/communicator/tcp_client.h | 4 ++ .../ccsrc/ps/core/communicator/tcp_server.cc | 25 +++++++++++- .../ccsrc/ps/core/communicator/tcp_server.h | 4 ++ mindspore/ccsrc/ps/ps_context.cc | 4 ++ mindspore/ccsrc/ps/ps_context.h | 4 ++ mindspore/parallel/_ps_context.py | 3 +- 15 files changed, 114 insertions(+), 29 deletions(-) diff --git a/mindspore/ccsrc/pipeline/jit/init.cc b/mindspore/ccsrc/pipeline/jit/init.cc index c2532811ee5..d18e05caa93 100644 --- a/mindspore/ccsrc/pipeline/jit/init.cc +++ b/mindspore/ccsrc/pipeline/jit/init.cc @@ -351,7 +351,8 @@ PYBIND11_MODULE(_c_expression, m) { .def("set_client_epoch_num", &PSContext::set_client_epoch_num, "Set federated learning client epoch number.") .def("set_client_batch_size", &PSContext::set_client_batch_size, "Set federated learning client batch size.") .def("set_secure_aggregation", &PSContext::set_secure_aggregation, - "Set federated learning client using secure aggregation."); + "Set federated learning client using secure aggregation.") + .def("set_enable_ssl", &PSContext::enable_ssl, "Set PS SSL mode enabled or disabled."); (void)py::class_>(m, "OpInfoLoaderPy") .def(py::init()) diff --git a/mindspore/ccsrc/ps/constants.h b/mindspore/ccsrc/ps/constants.h index 1189ebc2275..92cb60280ac 100644 --- a/mindspore/ccsrc/ps/constants.h +++ b/mindspore/ccsrc/ps/constants.h @@ -54,7 +54,8 @@ constexpr char kSparseLazyAdamOp[] = "LazyAdam"; constexpr char kSparseFtrlOp[] = "FTRL"; constexpr char kCertificateChain[] = "server.crt"; -constexpr char kPrivateKey[] = "server.key.unsecure"; +constexpr char kPrivateKey[] = "server.key"; +constexpr char kCAcrt[] = "ca.crt"; constexpr int64_t kInitWeightsCmd = 10; constexpr int64_t kInitWeightToOptimIdCmd = 11; @@ -82,6 +83,8 @@ constexpr int64_t kSubmitTimeOutInMs = 30000; constexpr int64_t kRetryCount = 60; constexpr int64_t kRetryIntervalInMs = 10; +constexpr int64_t kThreadNum = 32; + using DataPtr = std::shared_ptr; using VectorPtr = std::shared_ptr>; using Key = uint64_t; diff --git a/mindspore/ccsrc/ps/core/communicator/http_communicator.cc b/mindspore/ccsrc/ps/core/communicator/http_communicator.cc index f4ed071977b..cdc8556b00f 100644 --- a/mindspore/ccsrc/ps/core/communicator/http_communicator.cc +++ b/mindspore/ccsrc/ps/core/communicator/http_communicator.cc @@ -23,7 +23,6 @@ namespace ps { namespace core { bool HttpCommunicator::Start() { MS_LOG(INFO) << "Initialize http server IP:" << ip_ << ", PORT:" << port_; - http_server_ = std::make_shared(ip_, port_, 32); http_server_->InitServer(); MS_EXCEPTION_IF_NULL(http_server_); if (!http_server_->Start()) { diff --git a/mindspore/ccsrc/ps/core/communicator/http_communicator.h b/mindspore/ccsrc/ps/core/communicator/http_communicator.h index 64ff99cb9c5..e03f32e4e19 100644 --- a/mindspore/ccsrc/ps/core/communicator/http_communicator.h +++ b/mindspore/ccsrc/ps/core/communicator/http_communicator.h @@ -34,7 +34,10 @@ class HttpCommunicator : public CommunicatorBase { public: explicit HttpCommunicator(const std::string &ip, std::int16_t port, const std::shared_ptr &task_executor) - : task_executor_(task_executor), http_server_(nullptr), ip_(ip), port_(port) {} + : task_executor_(task_executor), http_server_(nullptr), ip_(ip), port_(port) { + http_server_ = std::make_shared(ip_, port_, kThreadNum); + } + ~HttpCommunicator() = default; bool Start() override; diff --git a/mindspore/ccsrc/ps/core/communicator/http_request_handler.cc b/mindspore/ccsrc/ps/core/communicator/http_request_handler.cc index 9e003346404..ccf2bf5519d 100644 --- a/mindspore/ccsrc/ps/core/communicator/http_request_handler.cc +++ b/mindspore/ccsrc/ps/core/communicator/http_request_handler.cc @@ -25,28 +25,30 @@ bool HttpRequestHandler::Initialize(int fd, const std::unordered_mapenable_ssl()) { + SSL_CTX_set_options(SSLWrapper::GetInstance().GetSSLCtx(), + SSL_OP_SINGLE_DH_USE | SSL_OP_SINGLE_ECDH_USE | SSL_OP_NO_SSLv2); + EC_KEY *ecdh = EC_KEY_new_by_curve_name(NID_X9_62_prime256v1); + MS_EXCEPTION_IF_NULL(ecdh); - if (!SSL_CTX_use_certificate_chain_file(SSLWrapper::GetInstance().GetSSLCtx(), kCertificateChain)) { - MS_LOG(ERROR) << "SSL use certificate chain file failed!"; - return false; + if (!SSL_CTX_use_certificate_chain_file(SSLWrapper::GetInstance().GetSSLCtx(), kCertificateChain)) { + MS_LOG(ERROR) << "SSL use certificate chain file failed!"; + return false; + } + + if (!SSL_CTX_use_PrivateKey_file(SSLWrapper::GetInstance().GetSSLCtx(), kPrivateKey, SSL_FILETYPE_PEM)) { + MS_LOG(ERROR) << "SSL use private key file failed!"; + return false; + } + + if (!SSL_CTX_check_private_key(SSLWrapper::GetInstance().GetSSLCtx())) { + MS_LOG(ERROR) << "SSL check private key file failed!"; + return false; + } + + evhttp_set_bevcb(http, BuffereventCallback, SSLWrapper::GetInstance().GetSSLCtx()); } - if (!SSL_CTX_use_PrivateKey_file(SSLWrapper::GetInstance().GetSSLCtx(), kPrivateKey, SSL_FILETYPE_PEM)) { - MS_LOG(ERROR) << "SSL use private key file failed!"; - return false; - } - - if (!SSL_CTX_check_private_key(SSLWrapper::GetInstance().GetSSLCtx())) { - MS_LOG(ERROR) << "SSL check private key file failed!"; - return false; - } - - evhttp_set_bevcb(http, BuffereventCallback, SSLWrapper::GetInstance().GetSSLCtx()); - int result = evhttp_accept_socket(http, fd); if (result < 0) { MS_LOG(ERROR) << "Evhttp accept socket failed!"; diff --git a/mindspore/ccsrc/ps/core/communicator/http_request_handler.h b/mindspore/ccsrc/ps/core/communicator/http_request_handler.h index 0bf47562e86..8dda31a5c3f 100644 --- a/mindspore/ccsrc/ps/core/communicator/http_request_handler.h +++ b/mindspore/ccsrc/ps/core/communicator/http_request_handler.h @@ -31,6 +31,7 @@ #include "ps/core/communicator/http_message_handler.h" #include "ps/core/communicator/ssl_wrapper.h" #include "ps/constants.h" +#include "ps/ps_context.h" namespace mindspore { namespace ps { diff --git a/mindspore/ccsrc/ps/core/communicator/ssl_wrapper.cc b/mindspore/ccsrc/ps/core/communicator/ssl_wrapper.cc index 28df9165b66..4c33354affc 100644 --- a/mindspore/ccsrc/ps/core/communicator/ssl_wrapper.cc +++ b/mindspore/ccsrc/ps/core/communicator/ssl_wrapper.cc @@ -20,7 +20,7 @@ namespace mindspore { namespace ps { namespace core { -SSLWrapper::SSLWrapper() : ssl_ctx_(nullptr) { InitSSL(); } +SSLWrapper::SSLWrapper() : ssl_ctx_(nullptr), client_ssl_ctx_(nullptr) { InitSSL(); } SSLWrapper::~SSLWrapper() { CleanSSL(); } @@ -41,6 +41,10 @@ void SSLWrapper::InitSSL() { if (X509_STORE_set_default_paths(store) != 1) { MS_LOG(ERROR) << "X509_STORE_set_default_paths failed"; } + client_ssl_ctx_ = SSL_CTX_new(SSLv23_client_method()); + if (!ssl_ctx_) { + MS_LOG(ERROR) << "SSL_CTX_new failed"; + } } void SSLWrapper::CleanSSL() { @@ -53,7 +57,13 @@ void SSLWrapper::CleanSSL() { CRYPTO_cleanup_all_ex_data(); } -SSL_CTX *SSLWrapper::GetSSLCtx() { return ssl_ctx_; } +SSL_CTX *SSLWrapper::GetSSLCtx(bool is_server) { + if (is_server) { + return ssl_ctx_; + } else { + return client_ssl_ctx_; + } +} } // namespace core } // namespace ps } // namespace mindspore diff --git a/mindspore/ccsrc/ps/core/communicator/ssl_wrapper.h b/mindspore/ccsrc/ps/core/communicator/ssl_wrapper.h index 147b4aff086..843af7b8087 100644 --- a/mindspore/ccsrc/ps/core/communicator/ssl_wrapper.h +++ b/mindspore/ccsrc/ps/core/communicator/ssl_wrapper.h @@ -32,7 +32,7 @@ class SSLWrapper { static SSLWrapper instance; return instance; } - SSL_CTX *GetSSLCtx(); + SSL_CTX *GetSSLCtx(bool is_server = true); private: SSLWrapper(); @@ -44,6 +44,7 @@ class SSLWrapper { void CleanSSL(); SSL_CTX *ssl_ctx_; + SSL_CTX *client_ssl_ctx_; }; } // namespace core } // namespace ps diff --git a/mindspore/ccsrc/ps/core/communicator/tcp_client.cc b/mindspore/ccsrc/ps/core/communicator/tcp_client.cc index 968ade6b5b2..698b352e3fa 100644 --- a/mindspore/ccsrc/ps/core/communicator/tcp_client.cc +++ b/mindspore/ccsrc/ps/core/communicator/tcp_client.cc @@ -106,7 +106,32 @@ void TcpClient::Init() { sin.sin_addr.s_addr = inet_addr(server_address_.c_str()); sin.sin_port = htons(server_port_); - buffer_event_ = bufferevent_socket_new(event_base_, -1, BEV_OPT_CLOSE_ON_FREE | BEV_OPT_THREADSAFE); + if (!PSContext::instance()->enable_ssl()) { + buffer_event_ = bufferevent_socket_new(event_base_, -1, BEV_OPT_CLOSE_ON_FREE | BEV_OPT_THREADSAFE); + } else { + SSL *ssl = SSL_new(SSLWrapper::GetInstance().GetSSLCtx(false)); + if (!SSL_CTX_use_certificate_chain_file(SSLWrapper::GetInstance().GetSSLCtx(false), kCertificateChain)) { + MS_LOG(EXCEPTION) << "SSL use certificate chain file failed!"; + } + + if (!SSL_CTX_use_PrivateKey_file(SSLWrapper::GetInstance().GetSSLCtx(false), kPrivateKey, SSL_FILETYPE_PEM)) { + MS_LOG(EXCEPTION) << "SSL use private key file failed!"; + } + + if (!SSL_CTX_check_private_key(SSLWrapper::GetInstance().GetSSLCtx(false))) { + MS_LOG(EXCEPTION) << "SSL check private key file failed!"; + } + + if (!SSL_CTX_load_verify_locations(SSLWrapper::GetInstance().GetSSLCtx(false), kCAcrt, nullptr)) { + MS_LOG(EXCEPTION) << "SSL load ca location failed!"; + } + + SSL_CTX_set_options(SSLWrapper::GetInstance().GetSSLCtx(false), SSL_OP_NO_SSLv2); + + buffer_event_ = bufferevent_openssl_socket_new(event_base_, -1, ssl, BUFFEREVENT_SSL_CONNECTING, + BEV_OPT_CLOSE_ON_FREE | BEV_OPT_THREADSAFE); + } + MS_EXCEPTION_IF_NULL(buffer_event_); bufferevent_setcb(buffer_event_, ReadCallback, nullptr, EventCallback, this); diff --git a/mindspore/ccsrc/ps/core/communicator/tcp_client.h b/mindspore/ccsrc/ps/core/communicator/tcp_client.h index 71748bab4e4..84704b840cf 100644 --- a/mindspore/ccsrc/ps/core/communicator/tcp_client.h +++ b/mindspore/ccsrc/ps/core/communicator/tcp_client.h @@ -22,6 +22,7 @@ #include #include #include +#include #include #include @@ -35,6 +36,9 @@ #include "ps/core/cluster_metadata.h" #include "utils/convert_utils_base.h" #include "ps/core/comm_util.h" +#include "ps/core/communicator/ssl_wrapper.h" +#include "ps/constants.h" +#include "ps/ps_context.h" namespace mindspore { namespace ps { diff --git a/mindspore/ccsrc/ps/core/communicator/tcp_server.cc b/mindspore/ccsrc/ps/core/communicator/tcp_server.cc index c23c2a607ff..9982d35176f 100644 --- a/mindspore/ccsrc/ps/core/communicator/tcp_server.cc +++ b/mindspore/ccsrc/ps/core/communicator/tcp_server.cc @@ -282,7 +282,30 @@ void TcpServer::ListenerCallback(struct evconnlistener *, evutil_socket_t fd, st MS_EXCEPTION_IF_NULL(base); MS_EXCEPTION_IF_NULL(sockaddr); - struct bufferevent *bev = bufferevent_socket_new(base, fd, BEV_OPT_CLOSE_ON_FREE | BEV_OPT_THREADSAFE); + struct bufferevent *bev = nullptr; + + if (!PSContext::instance()->enable_ssl()) { + bev = bufferevent_socket_new(base, fd, BEV_OPT_CLOSE_ON_FREE | BEV_OPT_THREADSAFE); + } else { + SSL *ssl = SSL_new(SSLWrapper::GetInstance().GetSSLCtx()); + + if (!SSL_CTX_use_certificate_chain_file(SSLWrapper::GetInstance().GetSSLCtx(), kCertificateChain)) { + MS_LOG(EXCEPTION) << "SSL use certificate chain file failed!"; + } + + if (!SSL_CTX_use_PrivateKey_file(SSLWrapper::GetInstance().GetSSLCtx(), kPrivateKey, SSL_FILETYPE_PEM)) { + MS_LOG(EXCEPTION) << "SSL use private key file failed!"; + } + + if (!SSL_CTX_check_private_key(SSLWrapper::GetInstance().GetSSLCtx())) { + MS_LOG(EXCEPTION) << "SSL check private key file failed!"; + } + SSL_CTX_set_options(SSLWrapper::GetInstance().GetSSLCtx(), SSL_OP_NO_SSLv2); + + bev = bufferevent_openssl_socket_new(base, fd, ssl, BUFFEREVENT_SSL_ACCEPTING, + BEV_OPT_CLOSE_ON_FREE | BEV_OPT_THREADSAFE); + } + if (bev == nullptr) { MS_LOG(ERROR) << "Error constructing buffer event!"; int ret = event_base_loopbreak(base); diff --git a/mindspore/ccsrc/ps/core/communicator/tcp_server.h b/mindspore/ccsrc/ps/core/communicator/tcp_server.h index 64c36cb6e1a..0a0edf25da0 100644 --- a/mindspore/ccsrc/ps/core/communicator/tcp_server.h +++ b/mindspore/ccsrc/ps/core/communicator/tcp_server.h @@ -22,6 +22,7 @@ #include #include #include +#include #include #include @@ -35,9 +36,12 @@ #include #include "ps/core/communicator/tcp_message_handler.h" +#include "ps/core/communicator/ssl_wrapper.h" #include "ps/core/cluster_metadata.h" #include "utils/convert_utils_base.h" #include "ps/core/comm_util.h" +#include "ps/constants.h" +#include "ps/ps_context.h" namespace mindspore { namespace ps { diff --git a/mindspore/ccsrc/ps/ps_context.cc b/mindspore/ccsrc/ps/ps_context.cc index 937bf64baf2..6e1d937d8ab 100644 --- a/mindspore/ccsrc/ps/ps_context.cc +++ b/mindspore/ccsrc/ps/ps_context.cc @@ -286,5 +286,9 @@ uint64_t PSContext::worker_upload_weights() const { return worker_upload_weights void PSContext::set_secure_aggregation(bool secure_aggregation) { secure_aggregation_ = secure_aggregation; } bool PSContext::secure_aggregation() const { return secure_aggregation_; } + +bool PSContext::enable_ssl() const { return enable_ssl_; } + +void PSContext::set_enable_ssl(bool enabled) { enable_ssl_ = enabled; } } // namespace ps } // namespace mindspore diff --git a/mindspore/ccsrc/ps/ps_context.h b/mindspore/ccsrc/ps/ps_context.h index f1199c1d307..70f9f45939a 100644 --- a/mindspore/ccsrc/ps/ps_context.h +++ b/mindspore/ccsrc/ps/ps_context.h @@ -80,6 +80,8 @@ class PSContext { void CloneHashTable(const std::string &dest_param_name, const std::string &src_param_name) const; void set_cache_enable(bool cache_enable) const; void set_rank_id(int rank_id) const; + bool enable_ssl() const; + void set_enable_ssl(bool enabled); // In new server framework, process role, worker number, server number, scheduler ip and scheduler port should be set // by ps_context. @@ -145,6 +147,7 @@ class PSContext { is_worker_(false), is_pserver_(false), is_sched_(false), + enable_ssl_(false), rank_id_(-1), worker_num_(0), server_num_(0), @@ -166,6 +169,7 @@ class PSContext { bool is_worker_; bool is_pserver_; bool is_sched_; + bool enable_ssl_; int rank_id_; uint32_t worker_num_; uint32_t server_num_; diff --git a/mindspore/parallel/_ps_context.py b/mindspore/parallel/_ps_context.py index df4f11beb79..474fe8a6354 100644 --- a/mindspore/parallel/_ps_context.py +++ b/mindspore/parallel/_ps_context.py @@ -47,7 +47,8 @@ _set_ps_context_func_map = { "fl_iteration_num": ps_context().set_fl_iteration_num, "client_epoch_num": ps_context().set_client_epoch_num, "client_batch_size": ps_context().set_client_batch_size, - "secure_aggregation": ps_context().set_secure_aggregation + "secure_aggregation": ps_context().set_secure_aggregation, + "enable_ps_ssl": ps_context().set_enable_ssl } _get_ps_context_func_map = {