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
+41 −4
Merged

UCT/TCP: Validate AM receive length #12081

Show file tree
Hide file tree
Changes from 1 commit
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_debug("tcp_ep %p (state %s): received AM length %u exceeds maximum %zu",

@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.

ucs_debug hides this entirely in release runs. Peers do not exchange segment sizes, so a peer configured with a larger RX_SEG_SIZE will legitimately hit this and the user only sees a connection reset with no hint. Can we use ucs_diag here? The connection is reset right after, so there is at most one message per connection.

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 @@ -1485,10 +1495,13 @@ 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 >

@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.

hdr->length is still only ucs_assert-ed for the internal AM IDs further down (UCT_TCP_EP_PUT_REQ_AM_ID, UCT_TCP_EP_PUT_ACK_AM_ID, uct_tcp_cm_handle_conn_pkt). A crafted header passes this new check and in release builds uct_tcp_ep_handle_put_req() then memcpy's to put_req->addr. Should those lengths be validated here as well, or is that intentionally out of scope?

@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.

This is intentionally out of scope for this change. This PR prevents an oversized AM length from being retained and used to size a later receive beyond the RX buffer. Exact semantic validation of internal PUT and CM messages is a separate parser-hardening issue that needs event-specific checks and focused tests.

(iface->config.rx_seg_size - sizeof(*hdr)))) {
uct_tcp_ep_handle_invalid_am_length(

@svc-nvidia-pr-review svc-nvidia-pr-review Oct 14, 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.

minor: iface->config.rx_seg_size - sizeof(*hdr) is computed twice here; maybe derive it inside the helper from the ep's iface and drop the max_length parameter.

ep, hdr->length,
iface->config.rx_seg_size - sizeof(*hdr));
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)
{
const uint64_t magic_number = UCT_TCP_MAGIC_NUMBER;

@gleon99 gleon99 Oct 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

constexpr?

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