Copilot commented on code in PR #3535: URL: https://github.com/apache/brpc/pull/3535#discussion_r3989322747
########## src/brpc/adapter_transport.h: ########## @@ -0,0 +1,101 @@ +// 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. + +#ifndef BRPC_ADAPTER_TRANSPORT_H +#define BRPC_ADAPTER_TRANSPORT_H + +#include <memory> + +#include "brpc/socket_mode.h" +#include "brpc/transport.h" +#include "brpc/transport_handshake.h" +#include "brpc/parse_result.h" + +namespace brpc { + +class TcpTransport; +class RdmaTransport; +class UBShmTransport; + +// The top-level Transport installed in Socket. It starts on TcpTransport and +// may switch to an independent RDMA/URMA/UBSHM Transport after a successful +// handshake. TCP remains usable before negotiation and after fallback. +class AdapterTransport : public Transport { + friend class TransportFactory; + friend class RdmaTransport; + friend class UBShmTransport; +public: + void Init(Socket* socket, const SocketOptions& options) override; + void Release() override; + int Reset(int32_t expected_nref) override; + std::shared_ptr<AppConnect> Connect() override; + int CutFromIOBuf(butil::IOBuf* buf) override; + ssize_t CutFromIOBufList(butil::IOBuf** buf, size_t ndata) override; + int WaitEpollOut(butil::atomic<int>* epollout_butex, + bool pollin, timespec duetime) override; + void ProcessEvent(bthread_attr_t attr) override; + void QueueMessage(InputMessageClosure& input_msg, + int* num_bthread_created, bool last_msg) override; + void Debug(std::ostream& os) override; + + int handshake_phase() const { return _handshake.phase(); } + int handshake_version() const { return _handshake.protocol_version(); } + handshake::HandshakeSession* handshake_session() { return &_handshake; } + Transport* high_speed_transport() const { + return _high_speed_transport.get(); + } + bool upgrade_capable() const { return _high_speed_transport != NULL; } Review Comment: This only checks whether some high-speed transport exists, but the RDMA and UBSHM server adapters use it to decide whether their concrete downcast is safe. For example, an UBSHM socket receiving a non-UB handshake dispatches to the RDMA adapter, which then treats `UBShmTransport` as `RdmaTransport` (and vice versa); the resulting endpoint access is undefined and can crash or corrupt state. Make the capability check transport-mode-specific before either adapter runs. ########## src/brpc/transport_handshake.h: ########## @@ -0,0 +1,197 @@ +// 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. + +#ifndef BRPC_TRANSPORT_HANDSHAKE_H +#define BRPC_TRANSPORT_HANDSHAKE_H + +#include <cstddef> +#include <functional> +#include <string> +#include <vector> + +#include "butil/atomicops.h" +#include "butil/macros.h" +#include "brpc/destroyable.h" +#include "brpc/handshake/handshake_frame.h" + +namespace brpc { + +class Socket; + +namespace handshake { + +class HandshakeAdapter; + +// Context retained by InputMessenger between the hello and ACK parse calls. +// Remembering the selected stateless adapter is necessary because ACK frames +// have no magic and cannot be dispatched from their bytes alone. +struct ServerHandshakeContext : public Destroyable { + ServerHandshakeContext() : _adapter(NULL) {} + static ServerHandshakeContext* Create(HandshakeAdapter* adapter); + HandshakeAdapter* adapter() const { return _adapter; } + void Destroy() override; + +private: + HandshakeAdapter* _adapter; +}; + +// Protocol adapters may use transport-specific intermediate values, but the +// terminal values are shared so that AdapterTransport can make the same +// acquire-side decision for RDMA, URMA and UBSHM. +enum Phase { + UNINITIALIZED = 0, + PREPARING = 1, + HELLO_SEND = 2, + HELLO_WAIT = 3, + NEGOTIATING = 4, + ACK_SEND = 5, + ACK_WAIT = 6, + ESTABLISHED = 0x100, + FALLBACK_TCP = 0x200, + FAILED = 0x300, +}; + +enum StepResult { + STEP_OK = 0, + STEP_FALLBACK, + STEP_NEED_MORE, + STEP_NOT_MINE, + STEP_ERROR, +}; + +// A protocol describes only its fields and resource-independent wire values. +// HandshakeSession owns framing and I/O through FrameCodec. The callbacks may +// retain strongly typed parsed state in their protocol adapter. +struct HandshakeCodec { + int protocol_version; + FrameSpec hello_frame; + FrameSpec ack_frame; + std::function<StepResult(bool, std::string*)> build_hello; + std::function<StepResult(const std::string&)> parse_hello; + std::function<StepResult(bool, std::string*)> build_ack; + std::function<StepResult(const std::string&, bool*)> parse_ack; +}; + +// Resource-specific operations supplied by a Transport and invoked by the +// common coordinator. Wire I/O and field codec invocation remain owned by +// HandshakeSession. +struct TransportUpgradeOps { + std::function<StepResult()> prepare_resources; + std::function<StepResult()> negotiate_resources; + std::function<void()> set_high_speed_active; + std::function<void()> set_tcp_active; + std::function<void()> on_failed; +}; + +struct ClientHandshakeCallbacks { + HandshakeCodec codec; + TransportUpgradeOps transport; +}; + +// The server driver is independent of the input mode. A parser callback can +// return STEP_NEED_MORE, while a blocking callback waits before returning. +struct ServerHandshakeCallbacks { + bool fallback_on_not_mine; + // Buffered parsers may offer multiple codecs (RDMA v2/v3). Blocking + // server handshakes currently provide exactly one codec. + std::vector<HandshakeCodec> codecs; + HandshakeInput* input; + TransportUpgradeOps transport; + std::function<StepResult()> validate_established; +}; + +// Owns one connection-upgrade attempt, invokes the protocol field codec and +// resource callbacks, and provides common framing, TCP control-plane I/O, +// lifecycle and publication ordering. +class HandshakeSession { +public: + explicit HandshakeSession(Socket* socket = NULL) + : _socket_io(socket), _io(&_socket_io), _phase(UNINITIALIZED), + _protocol_version(0), _local_enabled(false) {} + + void Reset(Socket* socket) { + _socket_io.Reset(socket); + _io = &_socket_io; + _protocol_version = 0; + _local_enabled = false; + _phase.store(UNINITIALIZED, butil::memory_order_relaxed); + } + + int phase(butil::memory_order order = butil::memory_order_acquire) const { + return _phase.load(order); + } + + void SetPhase(int phase) { + _phase.store(phase, butil::memory_order_relaxed); Review Comment: The client handshake publishes `HELLO_WAIT` with this relaxed store, while `AdapterTransport::ProcessTcpEvent` uses an acquire load and only wakes `SocketHandshakeIO` when it observes a non-`UNINITIALIZED` pre-establishment phase. Without a release publication, the event thread may observe the old `UNINITIALIZED` value and skip `NotifyReadable`, leaving `ReadExact` blocked even though the peer's response is readable. Publish intermediate phases with release semantics (or otherwise synchronize the phase transition). ########## src/brpc/handshake/ubshm_handshake.cpp: ########## @@ -0,0 +1,407 @@ +// 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/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 = 2; +static const uint16_t IMPL_VERSION = 1; +#endif // BRPC_WITH_UBRING +static const uint32_t ACK_OK = 0x1; + +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( + NULL, 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); + +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 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); +} + +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)); +} + +handshake::HandshakeCodec UBShmHandshakeAdapter::MakeCodec() const { + handshake::HandshakeCodec codec{}; + codec.protocol_version = 2; + codec.hello_frame = handshake::ubshm_wire::HelloFrameSpec(); + codec.ack_frame = handshake::ubshm_wire::AckFrameSpec(); + codec.build_ack = [](bool enabled, std::string* payload) { + const uint32_t flags_be = butil::HostToNet32( + enabled ? handshake::ubshm_wire::ACK_OK : 0); + payload->assign(reinterpret_cast<const char*>(&flags_be), + sizeof(flags_be)); + return handshake::STEP_OK; + }; + codec.parse_ack = [](const std::string& payload, bool* enabled) { + if (payload.size() != handshake::ubshm_wire::ACK_LEN) { + errno = EPROTO; + return handshake::STEP_ERROR; + } + uint32_t flags_be = 0; + memcpy(&flags_be, payload.data(), sizeof(flags_be)); + *enabled = (butil::NetToHost32(flags_be) & + handshake::ubshm_wire::ACK_OK) != 0; + return handshake::STEP_OK; + }; + return codec; +} + +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; + memcpy(message.shm_name, shm_name, SHM_MAX_NAME_BUFF_LEN); Review Comment: The client passes `shm_name_str.c_str()` from a short `std::string` here, but this unconditionally reads 48 bytes from that pointer. When the name is shorter than `SHM_MAX_NAME_BUFF_LEN`, the copy reads past the string allocation during every UBSHM client handshake; use a length-aware hello API or copy from a fixed-size, zero-padded buffer at the call sites. This issue also appears on line 368 of the same file. ########## src/brpc/handshake/ubshm_handshake.cpp: ########## @@ -0,0 +1,407 @@ +// 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/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 = 2; +static const uint16_t IMPL_VERSION = 1; +#endif // BRPC_WITH_UBRING +static const uint32_t ACK_OK = 0x1; + +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( + NULL, 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); + +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 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); +} + +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)); +} + +handshake::HandshakeCodec UBShmHandshakeAdapter::MakeCodec() const { + handshake::HandshakeCodec codec{}; + codec.protocol_version = 2; + codec.hello_frame = handshake::ubshm_wire::HelloFrameSpec(); + codec.ack_frame = handshake::ubshm_wire::AckFrameSpec(); + codec.build_ack = [](bool enabled, std::string* payload) { + const uint32_t flags_be = butil::HostToNet32( + enabled ? handshake::ubshm_wire::ACK_OK : 0); + payload->assign(reinterpret_cast<const char*>(&flags_be), + sizeof(flags_be)); + return handshake::STEP_OK; + }; + codec.parse_ack = [](const std::string& payload, bool* enabled) { + if (payload.size() != handshake::ubshm_wire::ACK_LEN) { + errno = EPROTO; + return handshake::STEP_ERROR; + } + uint32_t flags_be = 0; + memcpy(&flags_be, payload.data(), sizeof(flags_be)); + *enabled = (butil::NetToHost32(flags_be) & + handshake::ubshm_wire::ACK_OK) != 0; + return handshake::STEP_OK; + }; + return codec; +} + +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; + memcpy(message.shm_name, shm_name, SHM_MAX_NAME_BUFF_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; + } + return NegotiationValid(*message) ? + handshake::STEP_OK : handshake::STEP_FALLBACK; +} + +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); +}; + + +static HandshakeCodec MakeUBShmFallbackCodec() { + HandshakeCodec codec{}; + codec.protocol_version = 2; + codec.hello_frame = ubshm_wire::HelloFrameSpec(); + codec.ack_frame = ubshm_wire::AckFrameSpec(); + codec.parse_hello = [](const std::string&) { + return STEP_FALLBACK; + }; + codec.build_hello = [](bool enabled, std::string* payload) { + 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; + }; + codec.build_ack = [](bool enabled, std::string* payload) { + const uint32_t flags_be = butil::HostToNet32( + enabled ? ubshm_wire::ACK_OK : 0); + payload->assign(reinterpret_cast<const char*>(&flags_be), + sizeof(flags_be)); + return STEP_OK; + }; + codec.parse_ack = [](const std::string& payload, bool* enabled) { + if (payload.size() != ubshm_wire::ACK_LEN) { + errno = EPROTO; + return STEP_ERROR; + } + *enabled = false; + return STEP_OK; + }; + return codec; +} + +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); + ServerHandshakeCallbacks callbacks{}; + callbacks.fallback_on_not_mine = false; + callbacks.codecs.push_back(MakeUBShmFallbackCodec()); + callbacks.input = &input; + callbacks.transport.prepare_resources = []() { return STEP_OK; }; + callbacks.transport.negotiate_resources = []() { return STEP_OK; }; + callbacks.transport.set_high_speed_active = []() {}; + callbacks.transport.set_tcp_active = []() {}; + callbacks.transport.on_failed = []() {}; + return GetSession(socket)->RunServer(callbacks); +} + +#if BRPC_WITH_UBRING +StepResult UBShmServerHandshakeAdapter::RunUBShmServerHandshake( + butil::IOBuf* source, Socket* socket) { + UBShmTransport* transport = UBShmTransport::Get(socket); + CHECK(transport->GetUBShmEp() != NULL); + + ubring::HelloMessage remote{}; + ubring::UBShmHandshakeAdapter wire; + IOBufHandshakeInput input(source); + ServerHandshakeCallbacks callbacks{}; + callbacks.fallback_on_not_mine = false; + callbacks.input = &input; + HandshakeCodec codec = wire.MakeCodec(); + codec.parse_hello = [&](const std::string& payload) { + const StepResult result = wire.ParseHello(payload, &remote); + if (result == STEP_OK || result == STEP_FALLBACK) { + LOG_IF(INFO, ubring::FLAGS_ub_trace_verbose) + << "server receive handshake message : " + << remote.toString(); + } + if (result == STEP_FALLBACK) { + transport->DeactivateUpgrade(); + } + return result; + }; + codec.build_hello = [&](bool enabled, std::string* payload) { + const uint64_t len = enabled + ? static_cast<uint64_t>(ubring::FLAGS_data_queue_size) * + MB_TO_BYTE + : 0; + return wire.BuildHello( + enabled, len, enabled ? remote.shm_name : NULL, payload); + }; + callbacks.codecs.push_back(codec); + callbacks.transport.prepare_resources = [&]() { + if (!ubring::IsUBAvailable()) { + transport->DeactivateUpgrade(); + return STEP_FALLBACK; + } + ubring::SHM remote_trx_shm = { + NULL, remote.len, 0, {0}, + static_cast<uint32_t>(socket->fd())}; + strncpy(remote_trx_shm.name, remote.shm_name, + SHM_MAX_NAME_BUFF_LEN); + + const size_t local_shm_len = + static_cast<size_t>(ubring::FLAGS_data_queue_size) * MB_TO_BYTE; + ubring::SHM local_trx_shm = { + NULL, 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 != NULL) { + *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) < 0) { + LOG(WARNING) + << "Fail to allocate ub resources, fallback to tcp:" + << socket->description(); + transport->DeactivateUpgrade(); + return STEP_FALLBACK; + } + return STEP_OK; + }; + callbacks.transport.negotiate_resources = []() { return STEP_OK; }; + callbacks.validate_established = [&]() { + if (!source->empty() || + !transport->UpgradeActive()) { + return STEP_ERROR; + } + return STEP_OK; + }; + callbacks.transport.set_high_speed_active = [transport]() { + transport->ActivateUpgrade(); + }; + callbacks.transport.set_tcp_active = [transport]() { + transport->DeactivateUpgrade(); + }; + callbacks.transport.on_failed = []() {}; + const StepResult result = GetSession(socket)->RunServer(callbacks); + if (result == STEP_OK) { + transport->FinishUpgrade(); + LOG_IF(INFO, ubring::FLAGS_ub_trace_verbose) + << "Server handshake ends (use ubring) on " + << socket->description(); + } else if (result == STEP_FALLBACK) { + LOG_IF(INFO, ubring::FLAGS_ub_trace_verbose) + << "Server handshake ends (use tcp) on " + << socket->description(); + } + return result; +} +#endif + +StepResult UBShmServerHandshakeAdapter::RunServerStep( + butil::IOBuf* source, Socket* socket) { +#if BRPC_WITH_UBRING + if (AdapterTransport::Get(socket)->upgrade_capable()) { + return RunUBShmServerHandshake(source, socket); Review Comment: `upgrade_capable()` only checks that some high-speed transport exists, so this branch also runs on an RDMA-mode socket when the first byte is `U`. `RunUBShmServerHandshake` then calls `UBShmTransport::Get`, which static-casts the RDMA transport to `UBShmTransport`; an RDMA hello on an UBSHM-mode socket has the inverse problem. A client using the other upgrade protocol can therefore crash the server. Gate this branch on the configured socket mode as well. ########## src/brpc/ubshm_transport.cpp: ########## @@ -64,27 +74,47 @@ int UBShmTransport::Reset(int32_t expected_nref) { } std::shared_ptr<AppConnect> UBShmTransport::Connect() { - if (_default_connect == nullptr) { - return std::make_shared<ubring::UBConnect>(); - } return _default_connect; } -int UBShmTransport::CutFromIOBuf(butil::IOBuf *buf) { - if (_ub_ep && _ub_state != UB_OFF) { - butil::IOBuf *data_arr[1] = {buf}; - return _ub_ep->CutFromIOBufList(data_arr, 1); - } else { - return _tcp_transport->CutFromIOBuf(buf); +void UBShmTransport::SetHighSpeedAvailable(bool available) { + _ub_state = available ? UB_ON : UB_OFF; } + +int UBShmTransport::PrepareUpgradeResources(ubring::SHM *local_trx_shm, + const char *shm_name) { + return _ub_ep->AllocateClientResources(local_trx_shm, shm_name); +} + +int UBShmTransport::NegotiateUpgradeResources(ubring::SHM *local_trx_shm, + const char *shm_name) { + return _ub_ep->_ub_ring->UbrMapRemoteShm(local_trx_shm, shm_name); +} + +int UBShmTransport::PrepareServerUpgradeResources(ubring::SHM *remote_trx_shm, + ubring::SHM *local_trx_shm) { + return _ub_ep->AllocateServerResources(remote_trx_shm, local_trx_shm); +} + +void UBShmTransport::ActivateUpgrade() { SetHighSpeedAvailable(true); } + +void UBShmTransport::DeactivateUpgrade() { SetHighSpeedAvailable(false); } Review Comment: UBSHM fallback likewise only flips `_ub_state`, leaving the ring, shared-memory mappings, and registered poller alive after the common handshake selects TCP. The next TCP connection can therefore retain or process stale UBSHM resources; deactivation needs to reset/deallocate the endpoint before fallback. ########## src/brpc/handshake/handshake_io.cpp: ########## @@ -0,0 +1,150 @@ +// 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/handshake_io.h" + +#include <cstdint> +#include <errno.h> +#include <unistd.h> + +#include "bthread/butex.h" +#include "butil/time.h" +#include "brpc/errno.pb.h" +#include "brpc/socket.h" + +namespace brpc { +namespace handshake { + +size_t IOBufHandshakeInput::Size() const { + return _source != NULL ? _source->size() : 0; +} + +bool IOBufHandshakeInput::CopyTo(void* data, size_t len) const { + return _source != NULL && _source->copy_to(data, len) == len; +} + +bool IOBufHandshakeInput::Consume(size_t len) { + return _source != NULL && _source->pop_front(len) == len; +} + +static const int WAIT_TIMEOUT_MS = 50; + +SocketHandshakeIO::SocketHandshakeIO(Socket* socket) + : _socket(socket) + , _read_butex(bthread::butex_create_checked<butil::atomic<int> >()) { +} + +SocketHandshakeIO::~SocketHandshakeIO() { + bthread::butex_destroy(_read_butex); +} + +void SocketHandshakeIO::Reset(Socket* socket) { + _socket = socket; +} + +void SocketHandshakeIO::NotifyReadable() { + _read_butex->fetch_add(1, butil::memory_order_release); + bthread::butex_wake(_read_butex); +} + +template <typename ReadOnce> +static int ReadExactLoop(butil::atomic<int>* read_butex, + size_t len, ReadOnce read_once) { + size_t received = 0; + while (received < len) { + const int expected = read_butex->load(butil::memory_order_acquire); + const timespec duetime = butil::milliseconds_from_now(WAIT_TIMEOUT_MS); + const ssize_t nr = read_once(received, len - received); + if (nr < 0) { + if (errno != EAGAIN) { + return -1; + } Review Comment: A signal can interrupt `read(2)` and set `errno` to `EINTR`; this loop treats that normal retry condition as a fatal handshake I/O error. Retry the read before handling `EAGAIN` so transient signals do not fail otherwise healthy client handshakes. This issue also appears on line 115 of the same file. ########## src/brpc/handshake/rdma_handshake.cpp: ########## @@ -0,0 +1,576 @@ +// 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/rdma_handshake.h" + +#include <errno.h> +#include <limits> +#include <string> + +#include "butil/logging.h" +#include "butil/raw_pack.h" +#include "butil/sys_byteorder.h" +#include "brpc/adapter_transport.h" +#include "brpc/handshake/rdma_handshake_constants.h" +#include "brpc/rdma_handshake.pb.h" +#include "brpc/socket.h" + +#if BRPC_WITH_RDMA + +#include <cstring> + +#include <gflags/gflags.h> + +#include "brpc/rdma_transport.h" + +namespace brpc { +namespace rdma { + +DEFINE_int32(rdma_client_handshake_version, 2, + "RDMA handshake protocol version used by client. " + "2 = legacy 'RDMA' magic (default, compatible with all servers); " + "3 = new 'RDM3' protobuf-based handshake " + "(MUST only be enabled after target servers support v3)."); +DECLARE_bool(rdma_trace_verbose); + +extern const uint16_t MIN_QP_SIZE; +extern const uint16_t MIN_BLOCK_SIZE; +extern bool g_skip_rdma_init; + +DEFINE_bool(rdma_ece, false, + "Enable end-to-end ECE negotiation in the RDMA v3 handshake"); + +void RdmaHandshakeAdapter::FillLocalHello(ParsedHello* local) const { + _ep->GetLocalConnectionInfo(local); +} + +void RdmaHandshakeAdapter::PrepareClientEce() { + if (!FLAGS_rdma_ece) { + return; + } + ibv_ece ece; + const int rc = _ep->QueryLocalEce(&ece); + if (rc == 0) { + _ep->SetOutgoingEce(ece); + } else if (rc < 0) { + LOG_IF(WARNING, FLAGS_rdma_trace_verbose) + << "Fail to IbvQueryEce on client, ECE not advertised"; + } +} + +handshake::HandshakeCodec RdmaHandshakeAdapter::MakeCodec( + ParsedHello* remote) { + handshake::HandshakeCodec codec{}; + codec.protocol_version = ProtocolVersion(); + codec.hello_frame = HelloFrameSpec(); + codec.ack_frame = RdmaAckFrameSpec(); + codec.build_hello = [this](bool enabled, std::string* payload) { + return BuildLocalHello(enabled, payload); + }; + codec.parse_hello = [this, remote](const std::string& payload) { + return ParseRemoteHello(payload, remote); + }; + codec.build_ack = [](bool enabled, std::string* payload) { + const uint32_t flags_be = butil::HostToNet32( + enabled ? HELLO_ACK_RDMA_OK : 0); + payload->assign(reinterpret_cast<const char*>(&flags_be), + sizeof(flags_be)); + return handshake::STEP_OK; + }; + codec.parse_ack = [](const std::string& payload, bool* enabled) { + if (payload.size() != HELLO_ACK_LEN) { + errno = EPROTO; + return handshake::STEP_ERROR; + } + uint32_t flags_be = 0; + memcpy(&flags_be, payload.data(), sizeof(flags_be)); + *enabled = (butil::NetToHost32(flags_be) & HELLO_ACK_RDMA_OK) != 0; + return handshake::STEP_OK; + }; + return codec; +} + +namespace v2_wire { + +void HelloMessage::Serialize(void* data) const { + butil::RawPacker(data) + .pack16(msg_len) + .pack16(hello_ver) + .pack16(impl_ver) + .pack32(block_size) + .pack16(sq_size) + .pack16(rq_size) + .pack16(lid) + .pack_bytes(gid.raw, sizeof(gid.raw)) + .pack32(qp_num); +} + +void HelloMessage::Deserialize(const void* data) { + butil::RawUnpacker(data) + .unpack16(msg_len) + .unpack16(hello_ver) + .unpack16(impl_ver) + .unpack32(block_size) + .unpack16(sq_size) + .unpack16(rq_size) + .unpack16(lid) + .unpack_bytes(gid.raw, sizeof(gid.raw)) + .unpack32(qp_num); +} + +static bool ValidHelloMessage(const HelloMessage& msg) { + return msg.hello_ver == HELLO_V2_VERSION && + msg.impl_ver == IMPL_V2_VERSION && + msg.block_size >= MIN_BLOCK_SIZE && + msg.sq_size >= MIN_QP_SIZE && + msg.rq_size >= MIN_QP_SIZE; +} + +static void TranslateHello(const HelloMessage& msg, ParsedHello* out) { + out->block_size = msg.block_size; + out->sq_size = msg.sq_size; + out->rq_size = msg.rq_size; + out->lid = msg.lid; + out->gid = msg.gid; + out->qp_num = msg.qp_num; +} + +static void FillMessage(const ParsedHello& local, HelloMessage* msg) { + msg->msg_len = HELLO_V2_MSG_LEN_MIN; + msg->hello_ver = HELLO_V2_VERSION; + msg->impl_ver = IMPL_V2_VERSION; + msg->block_size = local.block_size; + msg->sq_size = local.sq_size; + msg->rq_size = local.rq_size; + msg->lid = local.lid; + msg->gid = local.gid; + msg->qp_num = local.qp_num; +} + +static handshake::StepResult SerializePayload( + const HelloMessage& msg, std::string* payload) { + uint8_t body[HELLO_V2_MSG_LEN_MIN - HELLO_MAGIC_LEN]; + msg.Serialize(body); + // FrameCodec owns msg_len, so the protocol payload starts after it. + payload->assign(reinterpret_cast<const char*>(body + sizeof(uint16_t)), + sizeof(body) - sizeof(uint16_t)); + return handshake::STEP_OK; +} + +static handshake::StepResult ParsePayload( + const std::string& payload, ParsedHello* remote) { + const size_t base_payload_len = + HELLO_V2_MSG_LEN_MIN - HELLO_MAGIC_LEN - sizeof(uint16_t); + if (payload.size() < base_payload_len) { + errno = EPROTO; + return handshake::STEP_ERROR; + } + uint8_t body[HELLO_V2_MSG_LEN_MIN - HELLO_MAGIC_LEN]; + const uint16_t total_be = butil::HostToNet16( + static_cast<uint16_t>(HELLO_MAGIC_LEN + sizeof(uint16_t) + + payload.size())); + memcpy(body, &total_be, sizeof(total_be)); + memcpy(body + sizeof(total_be), payload.data(), base_payload_len); + + HelloMessage msg{}; + msg.Deserialize(body); + if (!ValidHelloMessage(msg)) { + return handshake::STEP_FALLBACK; + } + TranslateHello(msg, remote); + return handshake::STEP_OK; +} + +} // namespace v2_wire + +const handshake::FrameSpec& +RdmaClientHandshakeAdapterV2::HelloFrameSpec() const { + return RdmaHelloFrameSpec(2); +} + +handshake::StepResult RdmaClientHandshakeAdapterV2::BuildLocalHello( + bool enabled, std::string* payload) { + CHECK(enabled); + ParsedHello local{}; + FillLocalHello(&local); + v2_wire::HelloMessage msg{}; + v2_wire::FillMessage(local, &msg); + return v2_wire::SerializePayload(msg, payload); +} + +handshake::StepResult RdmaClientHandshakeAdapterV2::ParseRemoteHello( + const std::string& payload, ParsedHello* remote) { + return v2_wire::ParsePayload(payload, remote); +} + +const handshake::FrameSpec& +RdmaServerHandshakeAdapterV2::HelloFrameSpec() const { + return RdmaHelloFrameSpec(2); +} + +handshake::StepResult RdmaServerHandshakeAdapterV2::BuildLocalHello( + bool enabled, std::string* payload) { + v2_wire::HelloMessage msg{}; + msg.msg_len = HELLO_V2_MSG_LEN_MIN; + if (enabled) { + ParsedHello local{}; + FillLocalHello(&local); + v2_wire::FillMessage(local, &msg); + } + return v2_wire::SerializePayload(msg, payload); +} + +handshake::StepResult RdmaServerHandshakeAdapterV2::ParseRemoteHello( + const std::string& payload, ParsedHello* remote) { + return v2_wire::ParsePayload(payload, remote); +} + +namespace v3_wire { + +static bool ValidRdmaHello(const RdmaHello& msg) { + if (msg.gid().size() != sizeof(ibv_gid)) { + return false; + } + const uint16_t max_uint16 = std::numeric_limits<uint16_t>::max(); + if (msg.sq_size() > max_uint16 || msg.rq_size() > max_uint16 || + msg.lid() > max_uint16) { + return false; + } + if (msg.block_size() < MIN_BLOCK_SIZE || msg.sq_size() < MIN_QP_SIZE || + msg.rq_size() < MIN_QP_SIZE) { + return false; + } + return msg.qp_num() != 0 || g_skip_rdma_init; +} + +static void FillLocalRdmaHello(const ParsedHello& local, RdmaHello* msg) { + msg->set_block_size(local.block_size); + msg->set_sq_size(local.sq_size); + msg->set_rq_size(local.rq_size); + msg->set_lid(local.lid); + msg->set_gid(reinterpret_cast<const char*>(local.gid.raw), + sizeof(local.gid.raw)); + msg->set_qp_num(local.qp_num); + if (FLAGS_rdma_ece && local.ece.has_value()) { + RdmaEce* ece = msg->mutable_ece(); + ece->set_vendor_id(local.ece->vendor_id); + ece->set_options(local.ece->options); + ece->set_comp_mask(local.ece->comp_mask); + } +} + +static void TranslateHello(const RdmaHello& msg, ParsedHello* out) { + out->block_size = msg.block_size(); + out->sq_size = static_cast<uint16_t>(msg.sq_size()); + out->rq_size = static_cast<uint16_t>(msg.rq_size()); + out->lid = static_cast<uint16_t>(msg.lid()); + fast_memcpy(out->gid.raw, msg.gid().data(), sizeof(out->gid.raw)); + out->qp_num = msg.qp_num(); + if (FLAGS_rdma_ece && msg.has_ece()) { + ibv_ece ece; + ece.vendor_id = msg.ece().vendor_id(); + ece.options = msg.ece().options(); + ece.comp_mask = msg.ece().comp_mask(); + out->ece = ece; + } +} + +static handshake::StepResult SerializePayload( + const RdmaHello& msg, std::string* payload) { + if (!msg.SerializeToString(payload) || + payload->size() > HELLO_V3_MAX_PB_SIZE) { + errno = EPROTO; + return handshake::STEP_ERROR; + } + return handshake::STEP_OK; +} + +static handshake::StepResult ParsePayload( + const std::string& payload, ParsedHello* remote) { + RdmaHello msg; + if (!msg.ParseFromArray(payload.data(), static_cast<int>(payload.size()))) { + errno = EPROTO; + return handshake::STEP_ERROR; + } + if (!ValidRdmaHello(msg)) { + return handshake::STEP_FALLBACK; + } + TranslateHello(msg, remote); + return handshake::STEP_OK; +} + +static void FillDisabledHello(RdmaHello* msg) { + msg->set_block_size(0); + msg->set_sq_size(0); + msg->set_rq_size(0); + msg->set_lid(0); + msg->set_gid(std::string(sizeof(ibv_gid), '\0')); + msg->set_qp_num(0); +} + +} // namespace v3_wire + +const handshake::FrameSpec& +RdmaClientHandshakeAdapterV3::HelloFrameSpec() const { + return RdmaHelloFrameSpec(3); +} + +handshake::StepResult RdmaClientHandshakeAdapterV3::BuildLocalHello( + bool enabled, std::string* payload) { + CHECK(enabled); + PrepareClientEce(); + ParsedHello local{}; + FillLocalHello(&local); + RdmaHello msg; + v3_wire::FillLocalRdmaHello(local, &msg); + return v3_wire::SerializePayload(msg, payload); +} + +handshake::StepResult RdmaClientHandshakeAdapterV3::ParseRemoteHello( + const std::string& payload, ParsedHello* remote) { + return v3_wire::ParsePayload(payload, remote); +} + +const handshake::FrameSpec& +RdmaServerHandshakeAdapterV3::HelloFrameSpec() const { + return RdmaHelloFrameSpec(3); +} + +handshake::StepResult RdmaServerHandshakeAdapterV3::BuildLocalHello( + bool enabled, std::string* payload) { + RdmaHello msg; + if (enabled) { + ParsedHello local{}; + FillLocalHello(&local); + v3_wire::FillLocalRdmaHello(local, &msg); + } else { + v3_wire::FillDisabledHello(&msg); + } + return v3_wire::SerializePayload(msg, payload); +} + +handshake::StepResult RdmaServerHandshakeAdapterV3::ParseRemoteHello( + const std::string& payload, ParsedHello* remote) { + return v3_wire::ParsePayload(payload, remote); +} + +std::unique_ptr<RdmaHandshakeAdapter> CreateClientHandshakeAdapter( + RdmaEndpoint* ep) { + if (FLAGS_rdma_client_handshake_version == 3) { + return std::unique_ptr<RdmaHandshakeAdapter>( + new RdmaClientHandshakeAdapterV3(ep)); + } + return std::unique_ptr<RdmaHandshakeAdapter>( + new RdmaClientHandshakeAdapterV2(ep)); +} + +std::vector<std::unique_ptr<RdmaHandshakeAdapter> > +CreateServerHandshakeAdapters(RdmaEndpoint* ep) { + std::vector<std::unique_ptr<RdmaHandshakeAdapter> > adapters; + adapters.emplace_back(new RdmaServerHandshakeAdapterV2(ep)); + adapters.emplace_back(new RdmaServerHandshakeAdapterV3(ep)); + return adapters; +} + +} // namespace rdma +} // namespace brpc + +#endif // BRPC_WITH_RDMA + +namespace brpc { +namespace handshake { + +class RdmaServerHandshakeAdapter : public StandardHandshakeAdapter { +public: + RdmaServerHandshakeAdapter() = 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_RDMA + StepResult RunRdmaServerHandshake( + butil::IOBuf* source, Socket* socket); +#endif + + DISALLOW_COPY_AND_ASSIGN(RdmaServerHandshakeAdapter); +}; + +static constexpr uint16_t V2_HELLO_VERSION_INVALID = + std::numeric_limits<uint16_t>::max(); +static constexpr size_t V3_GID_LEN = 16; + +static HandshakeCodec MakeRdmaFallbackCodec(int version) { + HandshakeCodec codec{}; + codec.protocol_version = version; + codec.hello_frame = rdma::RdmaHelloFrameSpec(version); + codec.ack_frame = rdma::RdmaAckFrameSpec(); + codec.parse_hello = [](const std::string&) { + return STEP_FALLBACK; + }; + codec.build_hello = [version](bool enabled, std::string* payload) { + if (enabled) { + errno = EPROTO; + return STEP_ERROR; + } + if (version == 2) { + payload->assign( + rdma::HELLO_V2_MSG_LEN_MIN - rdma::HELLO_MAGIC_LEN - + sizeof(uint16_t), + '\0'); + butil::RawPacker(&(*payload)[0]) + .pack16(V2_HELLO_VERSION_INVALID); + return STEP_OK; + } + + rdma::RdmaHello reply; + reply.set_block_size(0); + reply.set_sq_size(0); + reply.set_rq_size(0); + reply.set_lid(0); + reply.set_gid(std::string(V3_GID_LEN, '\0')); + reply.set_qp_num(0); + if (!reply.SerializeToString(payload)) { + errno = EPROTO; + return STEP_ERROR; + } + return STEP_OK; + }; + codec.build_ack = [](bool enabled, std::string* payload) { + const uint32_t flags_be = butil::HostToNet32( + enabled ? rdma::HELLO_ACK_RDMA_OK : 0); + payload->assign(reinterpret_cast<const char*>(&flags_be), + sizeof(flags_be)); + return STEP_OK; + }; + codec.parse_ack = [](const std::string& payload, bool* enabled) { + if (payload.size() != rdma::HELLO_ACK_LEN) { + errno = EPROTO; + return STEP_ERROR; + } + *enabled = false; + return STEP_OK; + }; + return codec; +} + +HandshakeAdapter* GetRdmaServerHandshakeAdapter() { + static RdmaServerHandshakeAdapter adapter; + return &adapter; +} + +HandshakeSession* RdmaServerHandshakeAdapter::GetSession( + Socket* socket) const { + return AdapterTransport::Get(socket)->handshake_session(); +} + +StepResult RdmaServerHandshakeAdapter::RunFallbackServerHandshake( + butil::IOBuf* source, Socket* socket) { + IOBufHandshakeInput input(source); + ServerHandshakeCallbacks callbacks{}; + callbacks.fallback_on_not_mine = false; + callbacks.codecs.push_back(MakeRdmaFallbackCodec(2)); + callbacks.codecs.push_back(MakeRdmaFallbackCodec(3)); + callbacks.input = &input; + callbacks.transport.prepare_resources = []() { return STEP_OK; }; + callbacks.transport.negotiate_resources = []() { return STEP_OK; }; + callbacks.transport.set_high_speed_active = []() {}; + callbacks.transport.set_tcp_active = []() {}; + callbacks.transport.on_failed = []() {}; + return GetSession(socket)->RunServer(callbacks); +} + +#if BRPC_WITH_RDMA +StepResult RdmaServerHandshakeAdapter::RunRdmaServerHandshake( + butil::IOBuf* source, Socket* socket) { + RdmaTransport* transport = RdmaTransport::Get(socket); + CHECK(transport->GetRdmaEp() != NULL); + + rdma::ParsedHello remote{}; + std::vector<std::unique_ptr<rdma::RdmaHandshakeAdapter> > protocols = + transport->CreateServerHandshakeAdapters(); + IOBufHandshakeInput input(source); + ServerHandshakeCallbacks callbacks{}; + callbacks.fallback_on_not_mine = false; + callbacks.input = &input; + for (size_t i = 0; i < protocols.size(); ++i) { + HandshakeCodec codec = protocols[i]->MakeCodec(&remote); + const std::function<StepResult(const std::string&)> parse_hello = + codec.parse_hello; + codec.parse_hello = [transport, parse_hello]( + const std::string& payload) { + const StepResult result = parse_hello(payload); + if (result == STEP_FALLBACK) { + transport->DeactivateUpgrade(); + } + return result; + }; + callbacks.codecs.push_back(codec); + } + callbacks.transport.prepare_resources = [&]() { + if (transport->PrepareUpgradeResources() < 0) { + PLOG(WARNING) + << "Fail to allocate rdma resources, fallback to tcp:" + << socket->description(); + transport->DeactivateUpgrade(); + return STEP_FALLBACK; + } + return STEP_OK; + }; + callbacks.transport.negotiate_resources = [&]() { + if (transport->NegotiateUpgradeResources(remote, true) < 0) { + PLOG(WARNING) + << "Fail to negotiate rdma resources, fallback to tcp:" + << socket->description(); + transport->DeactivateUpgrade(); + return STEP_FALLBACK; + } + return STEP_OK; + }; + callbacks.validate_established = [&]() { + if (!source->empty()) { + return STEP_ERROR; + } + return STEP_OK; + }; + callbacks.transport.set_high_speed_active = [transport]() { + transport->ActivateUpgrade(); + }; + callbacks.transport.set_tcp_active = [transport]() { + transport->DeactivateUpgrade(); + }; + callbacks.transport.on_failed = []() {}; + return GetSession(socket)->RunServer(callbacks); +} +#endif + +StepResult RdmaServerHandshakeAdapter::RunServerStep( + butil::IOBuf* source, Socket* socket) { +#if BRPC_WITH_RDMA + if (AdapterTransport::Get(socket)->upgrade_capable()) { + return RunRdmaServerHandshake(source, socket); Review Comment: `upgrade_capable()` only checks that some high-speed transport exists, so this branch also runs on an UBSHM-mode socket when the first byte is not `U`. `RunRdmaServerHandshake` then calls `RdmaTransport::Get`, which static-casts the UBSHM transport to `RdmaTransport`; the inverse is possible in the UBSHM adapter. A client using the other upgrade protocol can therefore crash the server. Gate this branch on the configured socket mode as well. ########## src/brpc/rdma_transport.cpp: ########## @@ -70,31 +68,49 @@ int RdmaTransport::Reset(int32_t expected_nref) { } std::shared_ptr<AppConnect> RdmaTransport::Connect() { - if (_default_connect == nullptr) { - return std::make_shared<rdma::RdmaConnect>(); - } - return _default_connect; + return _default_connect; +} + +void RdmaTransport::SetHighSpeedAvailable(bool available) { + _rdma_state = available ? RDMA_ON : RDMA_OFF; +} + +int RdmaTransport::PrepareUpgradeResources() { + return _rdma_ep->AllocateResources(); +} + +int RdmaTransport::NegotiateUpgradeResources( + const rdma::RdmaConnectionInfo &remote, bool server) { + _rdma_ep->ApplyRemoteInfo(remote); + return _rdma_ep->BringUpQp(remote, server); } +int RdmaTransport::StartUpgradeEvents() { + return _rdma_ep->StartCqEvents(); +} + +std::unique_ptr<rdma::RdmaHandshakeAdapter> +RdmaTransport::CreateClientHandshakeAdapter() { + return rdma::CreateClientHandshakeAdapter(_rdma_ep); +} + +std::vector<std::unique_ptr<rdma::RdmaHandshakeAdapter>> +RdmaTransport::CreateServerHandshakeAdapters() { + return rdma::CreateServerHandshakeAdapters(_rdma_ep); +} + +void RdmaTransport::ActivateUpgrade() { SetHighSpeedAvailable(true); } + +void RdmaTransport::DeactivateUpgrade() { SetHighSpeedAvailable(false); } Review Comment: Fallback only changes `_rdma_state`; it does not release resources allocated by `PrepareUpgradeResources`/`NegotiateUpgradeResources`. When the peer rejects the upgrade or negotiation fails, the socket continues as TCP while the RDMA QP/CQ and buffers remain attached until socket destruction, and repeated fallback connections can exhaust RDMA resources. Deactivation must reset/deallocate the endpoint before selecting TCP. -- 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]
