Spicy-cream commented on code in PR #3535: URL: https://github.com/apache/brpc/pull/3535#discussion_r4450273788
########## src/brpc/handshake/ubshm_handshake.cpp: ########## @@ -0,0 +1,486 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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 "brpc/handshake/ubshm_handshake.h" + +#include <errno.h> +#include <cstdio> + +#include "butil/raw_pack.h" +#include "butil/sys_byteorder.h" +#include "brpc/adapter_transport.h" +#include "brpc/socket.h" + +#if BRPC_WITH_UBRING + +#include <array> +#include <cstring> + +#include "butil/logging.h" +#include "brpc/reloadable_flags.h" +#include "brpc/ubshm/common/common.h" +#include "brpc/ubshm/shm/shm_mgr.h" +#include "brpc/ubshm/ub_endpoint.h" +#include "brpc/ubshm/ub_helper.h" +#include "brpc/ubshm/ubr_trx.h" +#include "brpc/ubshm_transport.h" + +#endif + +namespace brpc { +namespace handshake { +namespace ubshm_wire { + +static const char* const MAGIC = "UB"; +static const size_t MAGIC_LEN = 2; +static const size_t HELLO_LEN = 64; +static const size_t ACK_LEN = 4; +#if BRPC_WITH_UBRING +static const uint16_t HELLO_VERSION = 3; +static const uint16_t IMPL_VERSION = 1; +#endif // BRPC_WITH_UBRING +static const FrameSpec& HelloFrameSpec() { + static const FrameSpec spec( + MAGIC, MAGIC_LEN, HELLO_LEN, HELLO_LEN, FrameSpec::FIXED); + return spec; +} + +static const FrameSpec& AckFrameSpec() { + static const FrameSpec spec( + nullptr, 0, ACK_LEN, ACK_LEN, FrameSpec::FIXED); + return spec; +} + +} // namespace ubshm_wire +} // namespace handshake +} // namespace brpc + +#if BRPC_WITH_UBRING + +namespace brpc { +namespace ubring { + +DEFINE_int32(data_queue_size, 4, "data queue size for UB"); +DEFINE_bool(ub_trace_verbose, false, "Print log message verbosely"); +BRPC_VALIDATE_GFLAG(ub_trace_verbose, brpc::PassValidate); + +UbrDataFormat PreferredDataFormatForShmType(SHM_TYPE shm_type) { + switch (shm_type) { + case SHM_TYPE_IPC: + return UBR_DATA_FORMAT_IPC_V2; + case SHM_TYPE_UBS: + return UBR_DATA_FORMAT_LEGACY_64; + default: + return UBR_DATA_FORMAT_NONE; + } +} + +UbrDataFormat SelectDataFormat(UbrDataFormat local_format, + uint16_t remote_format_id) { + if ((local_format != UBR_DATA_FORMAT_LEGACY_64 && + local_format != UBR_DATA_FORMAT_IPC_V2) || + remote_format_id != static_cast<uint16_t>(local_format)) { + return UBR_DATA_FORMAT_NONE; + } + return local_format; +} + +void HelloMessage::Serialize(void* data) const { + char* current_pos = static_cast<char*>(data); + const uint16_t net_msg_len = butil::HostToNet16(msg_len); + memcpy(current_pos, &net_msg_len, sizeof(net_msg_len)); + current_pos += sizeof(net_msg_len); + const uint16_t net_hello_ver = butil::HostToNet16(hello_ver); + memcpy(current_pos, &net_hello_ver, sizeof(net_hello_ver)); + current_pos += sizeof(net_hello_ver); + const uint16_t net_impl_ver = butil::HostToNet16(impl_ver); + memcpy(current_pos, &net_impl_ver, sizeof(net_impl_ver)); + current_pos += sizeof(net_impl_ver); + const uint64_t net_len = butil::HostToNet64(len); + memcpy(current_pos, &net_len, sizeof(net_len)); + current_pos += sizeof(net_len); + memcpy(current_pos, shm_name, SHM_MAX_NAME_BUFF_LEN); +} + +void HelloFormatExtension::Serialize(void* data) const { + char* current = static_cast<char*>(data); + const uint16_t length = butil::HostToNet16(extension_len); + const uint16_t format = butil::HostToNet16(format_id); + memcpy(current, &length, sizeof(length)); + memcpy(current + sizeof(length), &format, sizeof(format)); +} + +void HelloFormatExtension::Deserialize(const void* data) { + const char* current = static_cast<const char*>(data); + uint16_t length; + uint16_t format; + memcpy(&length, current, sizeof(length)); + memcpy(&format, current + sizeof(length), sizeof(format)); + extension_len = butil::NetToHost16(length); + format_id = butil::NetToHost16(format); +} + +void HelloMessage::Deserialize(const void* data) { + const char* current_pos = static_cast<const char*>(data); + uint16_t net_msg_len; + memcpy(&net_msg_len, current_pos, sizeof(net_msg_len)); + msg_len = butil::NetToHost16(net_msg_len); + current_pos += sizeof(net_msg_len); + uint16_t net_hello_ver; + memcpy(&net_hello_ver, current_pos, sizeof(net_hello_ver)); + hello_ver = butil::NetToHost16(net_hello_ver); + current_pos += sizeof(net_hello_ver); + uint16_t net_impl_ver; + memcpy(&net_impl_ver, current_pos, sizeof(net_impl_ver)); + impl_ver = butil::NetToHost16(net_impl_ver); + current_pos += sizeof(net_impl_ver); + uint64_t net_len; + memcpy(&net_len, current_pos, sizeof(net_len)); + len = butil::NetToHost64(net_len); + current_pos += sizeof(net_len); + memcpy(shm_name, current_pos, SHM_MAX_NAME_BUFF_LEN); + // The decoded name is also used directly as a C-string by callers. + shm_name[SHM_MAX_NAME_BUFF_LEN - 1] = '\0'; +} + +std::string HelloMessage::toString() const { + constexpr size_t MAX_LEN = + 16 + 6 + 16 + 6 + 16 + 6 + 20 + 6 + SHM_MAX_NAME_BUFF_LEN + 32; + std::array<char, MAX_LEN> buf; + const int n = snprintf( + buf.data(), buf.size(), + "msg_len=%u, hello_ver=%u, impl_ver=%u, len=%lu, shm_name=%.*s", + msg_len, hello_ver, impl_ver, + static_cast<unsigned long>(len), + static_cast<int>(SHM_MAX_NAME_BUFF_LEN), shm_name); + return std::string(buf.data(), static_cast<size_t>(n)); +} + +const handshake::FrameSpec& UBShmHandshakeAdapter::HelloFrameSpec() const { + return handshake::ubshm_wire::HelloFrameSpec(); +} + +const handshake::FrameSpec& UBShmHandshakeAdapter::AckFrameSpec() const { + return handshake::ubshm_wire::AckFrameSpec(); +} + +const handshake::FrameSpec& +UBShmHandshakeAdapter::ExtensionFrameSpec() const { + static const handshake::FrameSpec spec( + nullptr, 0, HelloFormatExtension::WIRE_SIZE, + HelloFormatExtension::WIRE_SIZE, handshake::FrameSpec::FIXED); + return spec; +} + +void UBShmHandshakeAdapter::ConfigureClientHello( + uint64_t len, const char* shm_name) { + _local_len = len; + _local_name = shm_name != nullptr ? shm_name : ""; + _server_reply = false; +} + +handshake::StepResult UBShmHandshakeAdapter::BuildHello( + bool enabled, std::string* payload) { + const char* name = nullptr; + if (enabled) { + name = _server_reply ? _remote.shm_name : _local_name.c_str(); + } + return BuildHello(enabled, enabled ? _local_len : 0, name, payload); +} + +handshake::StepResult UBShmHandshakeAdapter::ParseHello( + const std::string& payload) { + return ParseHello(payload, &_remote); +} + +handshake::StepResult UBShmHandshakeAdapter::BuildExtension( + bool enabled, std::string* payload) { + const HelloFormatExtension extension = { + HelloFormatExtension::WIRE_SIZE, + static_cast<uint16_t>(enabled ? _local_format + : UBR_DATA_FORMAT_NONE)}; + payload->resize(HelloFormatExtension::WIRE_SIZE); + extension.Serialize(&(*payload)[0]); + return handshake::STEP_OK; +} + +handshake::StepResult UBShmHandshakeAdapter::ParseExtension( + const std::string& payload) { + if (payload.size() != HelloFormatExtension::WIRE_SIZE) { + errno = EPROTO; + return handshake::STEP_ERROR; + } + HelloFormatExtension extension{}; + extension.Deserialize(payload.data()); + return extension.extension_len == HelloFormatExtension::WIRE_SIZE && + SelectDataFormat(_local_format, extension.format_id) != + UBR_DATA_FORMAT_NONE + ? handshake::STEP_OK : handshake::STEP_FALLBACK; +} + +handshake::StepResult UBShmHandshakeAdapter::BuildHello( + bool enabled, uint64_t len, const char* shm_name, + std::string* payload) const { + HelloMessage message{}; + message.msg_len = static_cast<uint16_t>( + handshake::ubshm_wire::HELLO_LEN); + if (enabled) { + message.hello_ver = handshake::ubshm_wire::HELLO_VERSION; + message.impl_ver = handshake::ubshm_wire::IMPL_VERSION; + message.len = len; + if (shm_name == nullptr) { + errno = EINVAL; + return handshake::STEP_ERROR; + } + const size_t shm_name_len = + strnlen(shm_name, SHM_MAX_NAME_LEN); + memcpy(message.shm_name, shm_name, shm_name_len); + } + payload->assign( + handshake::ubshm_wire::HELLO_LEN - + handshake::ubshm_wire::MAGIC_LEN, + '\0'); + message.Serialize(&(*payload)[0]); + return handshake::STEP_OK; +} + +handshake::StepResult UBShmHandshakeAdapter::ParseHello( + const std::string& payload, HelloMessage* message) const { + if (payload.size() != handshake::ubshm_wire::HELLO_LEN - + handshake::ubshm_wire::MAGIC_LEN) { + errno = EPROTO; + return handshake::STEP_ERROR; + } + message->Deserialize(payload.data()); + if (message->msg_len != handshake::ubshm_wire::HELLO_LEN) { + errno = EPROTO; + return handshake::STEP_ERROR; + } + if (!NegotiationValid(*message)) { + return handshake::STEP_FALLBACK; + } + // Deserialize terminates the decoded name for safe C-string use. Inspect + // the original wire bytes so that truncation cannot hide a malformed name. + const char* wire_name = + payload.data() + payload.size() - SHM_MAX_NAME_BUFF_LEN; + if (memchr(wire_name, '\0', SHM_MAX_NAME_BUFF_LEN) == nullptr) { + errno = EPROTO; + return handshake::STEP_ERROR; + } + return handshake::STEP_OK; +} + +bool UBShmHandshakeAdapter::NegotiationValid( + const HelloMessage& message) const { + return message.hello_ver == handshake::ubshm_wire::HELLO_VERSION && + message.impl_ver == handshake::ubshm_wire::IMPL_VERSION; +} + +} // namespace ubring +} // namespace brpc + +#endif // BRPC_WITH_UBRING + +namespace brpc { +namespace handshake { + +class UBShmServerHandshakeAdapter : public StandardHandshakeAdapter { +public: + UBShmServerHandshakeAdapter() = default; + +protected: + StepResult RunServerStep( + butil::IOBuf* source, Socket* socket) override; + HandshakeSession* GetSession(Socket* socket) const override; + +private: + StepResult RunFallbackServerHandshake( + butil::IOBuf* source, Socket* socket); +#if BRPC_WITH_UBRING + StepResult RunUBShmServerHandshake( + butil::IOBuf* source, Socket* socket); +#endif + + DISALLOW_COPY_AND_ASSIGN(UBShmServerHandshakeAdapter); +}; + + +class UBShmFallbackProtocol : public FallbackHandshakeProtocol { +public: + int ProtocolVersion() const override { return 2; } + const FrameSpec& HelloFrameSpec() const override { + return ubshm_wire::HelloFrameSpec(); + } + const FrameSpec& AckFrameSpec() const override { + return ubshm_wire::AckFrameSpec(); + } + StepResult BuildHello(bool enabled, std::string* payload) override { + if (enabled) { + errno = EPROTO; + return STEP_ERROR; + } + payload->assign( + ubshm_wire::HELLO_LEN - ubshm_wire::MAGIC_LEN, '\0'); + butil::RawPacker(&(*payload)[0]) + .pack16(static_cast<uint16_t>(ubshm_wire::HELLO_LEN)); + return STEP_OK; + } +}; + +HandshakeAdapter* GetUBShmServerHandshakeAdapter() { + static UBShmServerHandshakeAdapter adapter; + return &adapter; +} + +HandshakeSession* UBShmServerHandshakeAdapter::GetSession( + Socket* socket) const { + return AdapterTransport::Get(socket)->handshake_session(); +} + +StepResult UBShmServerHandshakeAdapter::RunFallbackServerHandshake( + butil::IOBuf* source, Socket* socket) { + IOBufHandshakeInput input(source); + UBShmFallbackProtocol protocol; + std::vector<HandshakeProtocol*> protocols(1, &protocol); + FallbackHandshakeTransport transport; + return GetSession(socket)->RunServer( + protocols, &input, &transport, false); +} + +#if BRPC_WITH_UBRING +class UBShmServerHandshakeTransport : public HandshakeTransport { +public: + UBShmServerHandshakeTransport( + UBShmTransport* transport, ubring::UBShmHandshakeAdapter* protocol, + Socket* socket, butil::IOBuf* source) + : _transport(transport), _protocol(protocol), _socket(socket), + _source(source) {} + + void OnProtocolSelected(HandshakeProtocol*) override { + LOG_IF(INFO, ubring::FLAGS_ub_trace_verbose) + << "server receive handshake message : " + << _protocol->remote().toString(); + } + + StepResult PrepareResources() override { + if (!ubring::IsUBAvailable()) { + _transport->DeactivateUpgrade(); + return STEP_FALLBACK; + } + const ubring::HelloMessage& remote = _protocol->remote(); + const size_t remote_name_len = + strnlen(remote.shm_name, SHM_MAX_NAME_BUFF_LEN); + + ubring::SHM remote_trx_shm = { + nullptr, remote.len, 0, {0}, + static_cast<uint32_t>(_socket->fd())}; + memcpy(remote_trx_shm.name, remote.shm_name, remote_name_len + 1); + + const size_t local_shm_len = + static_cast<size_t>(ubring::FLAGS_data_queue_size) * MB_TO_BYTE; + ubring::SHM local_trx_shm = { + nullptr, local_shm_len, 0, {0}, + static_cast<uint32_t>(_socket->fd())}; + char client_name[SHM_MAX_NAME_BUFF_LEN + 1]; + memcpy(client_name, remote.shm_name, SHM_MAX_NAME_BUFF_LEN); + client_name[SHM_MAX_NAME_BUFF_LEN] = '\0'; + char* client_ip_port = strrchr(client_name, '_'); + if (client_ip_port != nullptr) { + *client_ip_port = '\0'; + } + const int result = snprintf( + local_trx_shm.name, SHM_MAX_NAME_BUFF_LEN, "%s_%s", + client_name, SERVER_SHM_NAME_SUFFIX); + if (UNLIKELY(result < 0)) { + _transport->DeactivateUpgrade(); + return STEP_FALLBACK; + } + if (_transport->PrepareServerUpgradeResources( + &remote_trx_shm, &local_trx_shm, + _protocol->local_format()) < 0) { + LOG(WARNING) + << "Fail to allocate ub resources, fallback to tcp:" + << _socket->description(); + _transport->DeactivateUpgrade(); + return STEP_FALLBACK; + } + return STEP_OK; + } + + StepResult NegotiateResources() override { return STEP_OK; } + void OnEstablished() override { + _transport->GetUBShmEp()->SetNegotiatedDataFormat( + _protocol->local_format()); + _transport->ActivateUpgrade(); + } + void OnFallback() override { _transport->DeactivateUpgrade(); } Review Comment: fixed -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
