Skip to content

UCT/TCP: Validate AM receive length #12081

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
roiedanino merged 3 commits into from
Oct 18, 2026
Merged

UCT/TCP: Validate AM receive length #12081

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
2a7a1c8
UCT/TCP: Validate AM receive length
roiedanino Oct 13, 2026
c619274
UCT/TCP: Improve oversized AM diagnostics
roiedanino Oct 14, 2026
aa7024e
UCT/TCP: Reuse maximum AM length
roiedanino Oct 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 17 additions & 4 deletions src/uct/tcp/tcp_ep.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1280,6 +1280,16 @@ static inline void uct_tcp_ep_handle_recv_err(uct_tcp_ep_t *ep,
}
}

static UCS_F_NOINLINE void
uct_tcp_ep_handle_invalid_am_length(uct_tcp_ep_t *ep, uint32_t length,
size_t max_length)
{
ucs_diag("tcp_ep %p (state %s): received AM length %u exceeds maximum %zu",

@svc-nvidia-pr-review svc-nvidia-pr-review Oct 15, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe use ucs_error() here? diag is below the default warn log level so it is not printed by default, and this path also triggers for a legit peer configured with a bigger UCX_TCP_RX_SEG_SIZE (the sender bounds AM zcopy by its own rx_seg_size), so the user would see only a connection reset without the reason.

ep, uct_tcp_ep_cm_state[ep->conn_state].name, length, max_length);
uct_tcp_ep_ctx_reset(&ep->rx);
uct_tcp_ep_handle_disconnected(ep, UCS_ERR_CONNECTION_RESET);
}

static inline unsigned uct_tcp_ep_recv(uct_tcp_ep_t *ep, size_t recv_length)
{
uct_tcp_iface_t UCS_V_UNUSED *iface = ucs_derived_of(ep->super.super.iface,
Expand Down Expand Up @@ -1435,6 +1445,8 @@ static unsigned uct_tcp_ep_progress_am_rx(uct_tcp_ep_t *ep)
{
uct_tcp_iface_t *iface = ucs_derived_of(ep->super.super.iface,
uct_tcp_iface_t);
size_t max_am_length = iface->config.rx_seg_size -
sizeof(uct_tcp_am_hdr_t);
unsigned handled = 0;
uct_tcp_am_hdr_t *hdr;
size_t recv_length;
Expand Down Expand Up @@ -1485,10 +1497,11 @@ static unsigned uct_tcp_ep_progress_am_rx(uct_tcp_ep_t *ep)
}

hdr = UCS_PTR_BYTE_OFFSET(ep->rx.buf, ep->rx.offset);
ucs_assertv(hdr->length <= (iface->config.rx_seg_size - sizeof(*hdr)),
"tcp_ep %p (conn state - %s): %u vs %zu",
ep, uct_tcp_ep_cm_state[ep->conn_state].name, hdr->length,
(iface->config.rx_seg_size - sizeof(*hdr)));
if (ucs_unlikely(hdr->length > max_am_length)) {

@svc-nvidia-pr-review svc-nvidia-pr-review Oct 15, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we validate the length of the service AM ids here as well? hdr->length for PUT_REQ/PUT_ACK/CM is still only checked by ucs_assert, e.g. a PUT_REQ with length != sizeof(uct_tcp_ep_put_req_hdr_t) makes uct_tcp_ep_handle_put_req() read put_req->addr from stale buffer bytes and memcpy() to a garbage address in release builds.

uct_tcp_ep_handle_invalid_am_length(ep, hdr->length,
max_am_length);
goto out;
}

if (remaining < (sizeof(*hdr) + hdr->length)) {
handled++;
Expand Down
24 changes: 24 additions & 0 deletions test/gtest/uct/tcp/test_tcp.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -259,6 +259,30 @@ UCS_TEST_P(test_uct_tcp, listener_flood_connect_and_close) {
test_listener_flood(*m_ent, max_conn, 0);

@svc-nvidia-pr-review svc-nvidia-pr-review Oct 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Residual coverage gap: the new test only covers the single-header case in ACCEPTING state. The partial-message path (remaining < sizeof(*hdr) + hdr->length followed by a second uct_tcp_ep_recv()), which is the one that actually overflowed before this fix, is not directly exercised; it is only covered indirectly because the check now runs before that exit.

@roiedanino roiedanino Oct 14, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The test intentionally sends only the oversized header. Without the fix it enters the partial-message path and subsequent progress calls attempt the second receive, but no payload is sent, so it fails safely by leaving the endpoint active instead of corrupting memory. The exact second-receive overwrite was reproduced separately under ASan. With the fix, every complete header is rejected before it can be retained, independent of TCP segmentation.

@svc-nvidia-pr-review svc-nvidia-pr-review Oct 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Residual coverage gap: no coverage for a TX-capable EP hitting this path (error-handler flow via uct_iface_handle_ep_err), only the RX-only destroy-from-progress flow.

@roiedanino roiedanino Oct 14, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new check only selects the existing uct_tcp_ep_handle_disconnected() path. The TX-capable branch and uct_iface_handle_ep_err() are existing peer-failure plumbing covered by the generic peer-failure tests, while this focused regression covers the new oversized-header decision and the RX-only destruction path. I do not think a duplicate TX-capable fixture is needed here.

}

UCS_TEST_P(test_uct_tcp, reject_oversized_am)
{
constexpr uint64_t magic_number = UCT_TCP_MAGIC_NUMBER;
const uct_tcp_am_hdr_t hdr = {
0, std::numeric_limits<uint32_t>::max()
};
std::vector<char> buf(sizeof(magic_number) + sizeof(hdr));
std::vector<int> fds;

memcpy(&buf[0], &magic_number, sizeof(magic_number));
memcpy(&buf[sizeof(magic_number)], &hdr, sizeof(hdr));

setup_conns_to_entity(*m_ent, 1, fds);
post_send(fds.front(), buf);

wait_for_cond([this]() {
return ucs_list_is_empty(&m_tcp_iface->ep_list);
},
[this]() { progress(); }, 1.0);
EXPECT_TRUE(ucs_list_is_empty(&m_tcp_iface->ep_list));

close(fds.front());
}

UCS_TEST_P(test_uct_tcp, check_addr_len)
{
uct_iface_attr_t iface_attr;
Expand Down
Loading