mindspore2022/mindspore/ccsrc/debug/debugger/grpc_client.cc

253 lines
8.4 KiB
C++

/**
* Copyright 2020 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "debug/debugger/grpc_client.h"
#include <stdio.h>
#include <stdlib.h>
#include <openssl/pem.h>
#include <openssl/err.h>
#include <openssl/pkcs12.h>
#include <openssl/x509.h>
#include <openssl/evp.h>
#include <thread>
#include <vector>
#include "utils/log_adapter.h"
using debugger::Chunk;
using debugger::EventListener;
using debugger::EventReply;
using debugger::EventReply_Status_FAILED;
using debugger::GraphProto;
using debugger::Metadata;
using debugger::TensorProto;
using debugger::WatchpointHit;
#define CHUNK_SIZE 1024 * 1024 * 3
namespace mindspore {
GrpcClient::GrpcClient(const std::string &host, const std::string &port, const bool &ssl_certificate,
const std::string &certificate_dir, const std::string &certificate_passphrase)
: stub_(nullptr) {
Init(host, port, ssl_certificate, certificate_dir, certificate_passphrase);
}
void GrpcClient::Init(const std::string &host, const std::string &port, const bool &ssl_certificate,
const std::string &certificate_dir, const std::string &certificate_passphrase) {
std::string target_str = host + ":" + port;
MS_LOG(INFO) << "GrpcClient connecting to: " << target_str;
std::shared_ptr<grpc::Channel> channel;
if (ssl_certificate) {
FILE *fp;
EVP_PKEY *pkey = NULL;
X509 *cert = NULL;
STACK_OF(X509) *ca = NULL;
PKCS12 *p12 = NULL;
if ((fp = fopen(certificate_dir.c_str(), "rb")) == NULL) {
MS_LOG(ERROR) << "Error opening file: " << certificate_dir;
exit(EXIT_FAILURE);
}
p12 = d2i_PKCS12_fp(fp, NULL);
fclose(fp);
if (p12 == NULL) {
MS_LOG(ERROR) << "Error reading PKCS#12 file";
X509_free(cert);
EVP_PKEY_free(pkey);
sk_X509_pop_free(ca, X509_free);
exit(EXIT_FAILURE);
}
if (!PKCS12_parse(p12, certificate_passphrase.c_str(), &pkey, &cert, &ca)) {
MS_LOG(ERROR) << "Error parsing PKCS#12 file";
X509_free(cert);
EVP_PKEY_free(pkey);
sk_X509_pop_free(ca, X509_free);
exit(EXIT_FAILURE);
}
std::string strca;
std::string strcert;
std::string strkey;
if (pkey == NULL || cert == NULL || ca == NULL) {
MS_LOG(ERROR) << "Error private key or cert or CA certificate.";
X509_free(cert);
EVP_PKEY_free(pkey);
sk_X509_pop_free(ca, X509_free);
exit(EXIT_FAILURE);
} else {
ASN1_TIME *validtime = X509_getm_notAfter(cert);
if (X509_cmp_current_time(validtime) < 0) {
MS_LOG(ERROR) << "This certificate is over its valid time, please use a new certificate.";
X509_free(cert);
EVP_PKEY_free(pkey);
sk_X509_pop_free(ca, X509_free);
exit(EXIT_FAILURE);
}
int nid = X509_get_signature_nid(cert);
int keybit = EVP_PKEY_bits(pkey);
if (nid == NID_sha1) {
MS_LOG(WARNING) << "Signature algrithm is sha1, which maybe not secure enough.";
} else if (keybit < 2048) {
MS_LOG(WARNING) << "The private key bits is: " << keybit << ", which maybe not secure enough.";
}
int dwPriKeyLen = i2d_PrivateKey(pkey, NULL); // get the length of private key
unsigned char *pribuf = (unsigned char *)malloc(sizeof(unsigned char) * dwPriKeyLen);
i2d_PrivateKey(pkey, &pribuf); // PrivateKey DER code
strkey = std::string(reinterpret_cast<char const *>(pribuf), dwPriKeyLen);
int dwcertLen = i2d_X509(cert, NULL); // get the length of private key
unsigned char *certbuf = (unsigned char *)malloc(sizeof(unsigned char) * dwcertLen);
i2d_X509(cert, &certbuf); // PrivateKey DER code
strcert = std::string(reinterpret_cast<char const *>(certbuf), dwcertLen);
int dwcaLen = i2d_X509(sk_X509_value(ca, 0), NULL); // get the length of private key
unsigned char *cabuf = (unsigned char *)malloc(sizeof(unsigned char) * dwcaLen);
i2d_X509(sk_X509_value(ca, 0), &cabuf); // PrivateKey DER code
strca = std::string(reinterpret_cast<char const *>(cabuf), dwcaLen);
free(pribuf);
free(certbuf);
free(cabuf);
}
grpc::SslCredentialsOptions opts = {strca, strkey, strcert};
channel = grpc::CreateChannel(target_str, grpc::SslCredentials(opts));
} else {
channel = grpc::CreateChannel(target_str, grpc::InsecureChannelCredentials());
}
stub_ = EventListener::NewStub(channel);
}
void GrpcClient::Reset() { stub_ = nullptr; }
EventReply GrpcClient::WaitForCommand(const Metadata &metadata) {
EventReply reply;
grpc::ClientContext context;
grpc::Status status = stub_->WaitCMD(&context, metadata, &reply);
if (!status.ok()) {
MS_LOG(ERROR) << "RPC failed: WaitForCommand";
MS_LOG(ERROR) << status.error_code() << ": " << status.error_message();
reply.set_status(EventReply_Status_FAILED);
}
return reply;
}
EventReply GrpcClient::SendMetadata(const Metadata &metadata) {
EventReply reply;
grpc::ClientContext context;
grpc::Status status = stub_->SendMetadata(&context, metadata, &reply);
if (!status.ok()) {
MS_LOG(ERROR) << "RPC failed: SendMetadata";
MS_LOG(ERROR) << status.error_code() << ": " << status.error_message();
reply.set_status(EventReply_Status_FAILED);
}
return reply;
}
std::vector<std::string> ChunkString(std::string str, int graph_size) {
std::vector<std::string> buf;
int size_iter = 0;
while (size_iter < graph_size) {
int chunk_size = CHUNK_SIZE;
if (graph_size - size_iter < CHUNK_SIZE) {
chunk_size = graph_size - size_iter;
}
std::string buffer;
buffer.resize(chunk_size);
memcpy(reinterpret_cast<char *>(buffer.data()), str.data() + size_iter, chunk_size);
buf.push_back(buffer);
size_iter += CHUNK_SIZE;
}
return buf;
}
EventReply GrpcClient::SendGraph(const GraphProto &graph) {
EventReply reply;
grpc::ClientContext context;
Chunk chunk;
std::unique_ptr<grpc::ClientWriter<Chunk> > writer(stub_->SendGraph(&context, &reply));
std::string str = graph.SerializeAsString();
int graph_size = graph.ByteSize();
auto buf = ChunkString(str, graph_size);
for (unsigned int i = 0; i < buf.size(); i++) {
MS_LOG(INFO) << "RPC:sending the " << i << "chunk in graph";
chunk.set_buffer(buf[i]);
if (!writer->Write(chunk)) {
break;
}
std::this_thread::sleep_for(std::chrono::milliseconds(1));
}
writer->WritesDone();
grpc::Status status = writer->Finish();
if (!status.ok()) {
MS_LOG(ERROR) << "RPC failed: SendGraph";
MS_LOG(ERROR) << status.error_code() << ": " << status.error_message();
reply.set_status(EventReply_Status_FAILED);
}
return reply;
}
EventReply GrpcClient::SendTensors(const std::list<TensorProto> &tensors) {
EventReply reply;
grpc::ClientContext context;
std::unique_ptr<grpc::ClientWriter<TensorProto> > writer(stub_->SendTensors(&context, &reply));
for (const auto &tensor : tensors) {
if (!writer->Write(tensor)) {
break;
}
std::this_thread::sleep_for(std::chrono::milliseconds(1));
}
writer->WritesDone();
grpc::Status status = writer->Finish();
if (!status.ok()) {
MS_LOG(ERROR) << "RPC failed: SendTensors";
MS_LOG(ERROR) << status.error_code() << ": " << status.error_message();
reply.set_status(EventReply_Status_FAILED);
}
return reply;
}
EventReply GrpcClient::SendWatchpointHits(const std::list<WatchpointHit> &watchpoints) {
EventReply reply;
grpc::ClientContext context;
std::unique_ptr<grpc::ClientWriter<WatchpointHit> > writer(stub_->SendWatchpointHits(&context, &reply));
for (const auto &watchpoint : watchpoints) {
if (!writer->Write(watchpoint)) {
break;
}
std::this_thread::sleep_for(std::chrono::milliseconds(1));
}
writer->WritesDone();
grpc::Status status = writer->Finish();
if (!status.ok()) {
MS_LOG(ERROR) << "RPC failed: SendWatchpointHits";
MS_LOG(ERROR) << status.error_code() << ": " << status.error_message();
reply.set_status(EventReply_Status_FAILED);
}
return reply;
}
} // namespace mindspore