diff --git a/include/rootstream.h b/include/rootstream.h index 1b123af..820748c 100644 --- a/include/rootstream.h +++ b/include/rootstream.h @@ -84,6 +84,7 @@ typedef struct { typedef struct { uint8_t *data; /* Frame pixel data (RGBA) */ uint32_t size; /* Total size in bytes */ + uint32_t capacity; /* Allocated buffer size in bytes */ uint32_t width; /* Frame width */ uint32_t height; /* Frame height */ uint32_t pitch; /* Bytes per row (stride) */ @@ -218,6 +219,16 @@ typedef PACKED_STRUCT { } control_packet_t; PACKED_STRUCT_END +/* Fragmented video payload header (inside encrypted payload) */ +typedef PACKED_STRUCT { + uint32_t frame_id; /* Frame sequence number */ + uint32_t total_size; /* Total encoded frame size */ + uint32_t offset; /* Offset of this chunk */ + uint16_t chunk_size; /* Size of this chunk */ + uint16_t flags; /* Reserved for future use */ +} video_chunk_header_t; +PACKED_STRUCT_END + /* Encrypted input event payload */ typedef PACKED_STRUCT { uint8_t type; /* EV_KEY, EV_REL, etc */ @@ -250,6 +261,12 @@ typedef struct { uint64_t handshake_sent_time; /* Handshake timestamp for timeout */ char hostname[64]; /* Peer hostname */ bool is_streaming; /* Currently streaming? */ + uint32_t video_tx_frame_id; /* Outgoing video frame counter */ + uint32_t video_rx_frame_id; /* Current incoming frame id */ + uint8_t *video_rx_buffer; /* Reassembly buffer */ + size_t video_rx_capacity; /* Reassembly buffer size */ + size_t video_rx_expected; /* Expected frame size */ + size_t video_rx_received; /* Bytes received so far */ } peer_t; /* ============================================================================ @@ -472,6 +489,8 @@ void recording_cleanup(rootstream_ctx_t *ctx); int rootstream_net_init(rootstream_ctx_t *ctx, uint16_t port); int rootstream_net_send_encrypted(rootstream_ctx_t *ctx, peer_t *peer, uint8_t type, const void *data, size_t size); +int rootstream_net_send_video(rootstream_ctx_t *ctx, peer_t *peer, + const uint8_t *data, size_t size); int rootstream_net_recv(rootstream_ctx_t *ctx, int timeout_ms); int rootstream_net_handshake(rootstream_ctx_t *ctx, peer_t *peer); diff --git a/src/decoder_mf.c b/src/decoder_mf.c index 4f68227..333639c 100644 --- a/src/decoder_mf.c +++ b/src/decoder_mf.c @@ -461,6 +461,7 @@ int rootstream_decode_frame(rootstream_ctx_t *ctx, /* Fill output frame info */ out->data = mf->frame_buffer; out->size = (uint32_t)copy_size; + out->capacity = (uint32_t)mf->frame_buffer_size; out->width = mf->width; out->height = mf->height; out->pitch = mf->width; /* NV12 Y-plane pitch */ diff --git a/src/drm_capture.c b/src/drm_capture.c index dc237a6..c7d49e9 100644 --- a/src/drm_capture.c +++ b/src/drm_capture.c @@ -240,6 +240,7 @@ int rootstream_capture_init(rootstream_ctx_t *ctx) { ctx->current_frame.width = ctx->display.width; ctx->current_frame.height = ctx->display.height; ctx->current_frame.size = frame_size; + ctx->current_frame.capacity = frame_size; ctx->current_frame.format = 0x34325258; /* DRM_FORMAT_XRGB8888 */ printf("✓ DRM capture initialized: %dx%d @ %d Hz\n", diff --git a/src/network.c b/src/network.c index 3ce36d4..ea97b80 100644 --- a/src/network.c +++ b/src/network.c @@ -57,6 +57,69 @@ #define PACKET_MAGIC 0x524F4F54 /* "ROOT" */ #define DEFAULT_PORT 9876 +#define MAX_VIDEO_FRAME_SIZE (16 * 1024 * 1024) + +static size_t max_plain_payload_size(void) { + size_t max_packet = MAX_PACKET_SIZE; + if (max_packet <= sizeof(packet_header_t) + crypto_aead_chacha20poly1305_IETF_ABYTES) { + return 0; + } + return max_packet - sizeof(packet_header_t) - crypto_aead_chacha20poly1305_IETF_ABYTES; +} + +int rootstream_net_send_video(rootstream_ctx_t *ctx, peer_t *peer, + const uint8_t *data, size_t size) { + if (!ctx || !peer || !data || size == 0) { + fprintf(stderr, "ERROR: Invalid arguments to send_video\n"); + return -1; + } + + size_t max_plain = max_plain_payload_size(); + if (max_plain <= sizeof(video_chunk_header_t)) { + fprintf(stderr, "ERROR: Payload size too small for video chunks\n"); + return -1; + } + + size_t max_chunk = max_plain - sizeof(video_chunk_header_t); + uint8_t *payload = malloc(sizeof(video_chunk_header_t) + max_chunk); + if (!payload) { + fprintf(stderr, "ERROR: Cannot allocate video payload buffer\n"); + return -1; + } + + uint32_t frame_id = peer->video_tx_frame_id++; + size_t offset = 0; + int result = 0; + + while (offset < size) { + size_t chunk_size = size - offset; + if (chunk_size > max_chunk) { + chunk_size = max_chunk; + } + + video_chunk_header_t header = { + .frame_id = frame_id, + .total_size = (uint32_t)size, + .offset = (uint32_t)offset, + .chunk_size = (uint16_t)chunk_size, + .flags = 0 + }; + + memcpy(payload, &header, sizeof(header)); + memcpy(payload + sizeof(header), data + offset, chunk_size); + + if (rootstream_net_send_encrypted(ctx, peer, PKT_VIDEO, + payload, sizeof(header) + chunk_size) < 0) { + result = -1; + break; + } + + offset += chunk_size; + } + + free(payload); + return result; +} /* * Initialize UDP socket for listening and sending @@ -155,7 +218,7 @@ int rootstream_net_init(rootstream_ctx_t *ctx, uint16_t port) { */ int rootstream_net_send_encrypted(rootstream_ctx_t *ctx, peer_t *peer, uint8_t type, const void *data, size_t size) { - if (!ctx || !peer || !data) { + if (!ctx || !peer || (!data && size > 0)) { fprintf(stderr, "ERROR: Invalid arguments to send_encrypted\n"); return -1; } @@ -167,6 +230,13 @@ int rootstream_net_send_encrypted(rootstream_ctx_t *ctx, peer_t *peer, return -1; } + size_t max_plain = max_plain_payload_size(); + if (size > max_plain) { + fprintf(stderr, "ERROR: Payload too large for single packet (%zu > %zu)\n", + size, max_plain); + return -1; + } + /* Allocate packet buffer */ size_t max_cipher_len = size + crypto_aead_chacha20poly1305_IETF_ABYTES; size_t packet_len = sizeof(packet_header_t) + max_cipher_len; @@ -387,22 +457,58 @@ int rootstream_net_recv(rootstream_ctx_t *ctx, int timeout_ms) { rootstream_input_process(ctx, input); } else if (hdr->type == PKT_VIDEO) { - /* Store video frame for client loop to consume */ - if (decrypted_len <= sizeof(ctx->current_frame.data)) { - if (!ctx->current_frame.data) { - /* Allocate frame buffer on first use */ - ctx->current_frame.data = malloc(MAX_PACKET_SIZE * 10); /* Larger buffer for frames */ - if (!ctx->current_frame.data) { - fprintf(stderr, "ERROR: Failed to allocate frame buffer\n"); - break; - } + if (decrypted_len < sizeof(video_chunk_header_t)) { + fprintf(stderr, "WARNING: Video chunk too small: %zu bytes\n", decrypted_len); + break; + } + + video_chunk_header_t header; + memcpy(&header, decrypted, sizeof(header)); + + if (header.total_size == 0 || header.total_size > MAX_VIDEO_FRAME_SIZE) { + fprintf(stderr, "WARNING: Invalid video frame size: %u bytes\n", + header.total_size); + break; + } + + if ((size_t)header.offset + header.chunk_size > header.total_size) { + fprintf(stderr, "WARNING: Video chunk out of range (offset=%u size=%u total=%u)\n", + header.offset, header.chunk_size, header.total_size); + break; + } + + if (decrypted_len != sizeof(video_chunk_header_t) + header.chunk_size) { + fprintf(stderr, "WARNING: Video chunk size mismatch\n"); + break; + } + + if (peer->video_rx_frame_id != header.frame_id) { + peer->video_rx_frame_id = header.frame_id; + peer->video_rx_received = 0; + peer->video_rx_expected = header.total_size; + } + + if (peer->video_rx_capacity < peer->video_rx_expected) { + uint8_t *new_buf = realloc(peer->video_rx_buffer, peer->video_rx_expected); + if (!new_buf) { + fprintf(stderr, "ERROR: Failed to allocate video reassembly buffer\n"); + break; } - memcpy(ctx->current_frame.data, decrypted, decrypted_len); - ctx->current_frame.size = decrypted_len; + peer->video_rx_buffer = new_buf; + peer->video_rx_capacity = peer->video_rx_expected; + } + + memcpy(peer->video_rx_buffer + header.offset, + decrypted + sizeof(video_chunk_header_t), + header.chunk_size); + peer->video_rx_received += header.chunk_size; + + if (peer->video_rx_received >= peer->video_rx_expected) { + ctx->current_frame.data = peer->video_rx_buffer; + ctx->current_frame.size = peer->video_rx_expected; + ctx->current_frame.capacity = peer->video_rx_capacity; ctx->current_frame.timestamp = get_timestamp_us(); ctx->frames_received++; - } else { - fprintf(stderr, "WARNING: Video frame too large: %zu bytes\n", decrypted_len); } } else if (hdr->type == PKT_AUDIO) { @@ -540,11 +646,11 @@ int rootstream_net_handshake(rootstream_ctx_t *ctx, peer_t *peer) { hdr.type = PKT_HANDSHAKE; hdr.payload_size = payload_len; - uint8_t packet[sizeof(packet_header_t) + payload_len]; + uint8_t packet[sizeof(packet_header_t) + 256]; /* Fixed size for MSVC compatibility */ memcpy(packet, &hdr, sizeof(hdr)); memcpy(packet + sizeof(hdr), payload, payload_len); - int sent = rs_socket_sendto(ctx->sock_fd, packet, sizeof(packet), 0, + int sent = rs_socket_sendto(ctx->sock_fd, packet, sizeof(hdr) + payload_len, 0, (struct sockaddr*)&peer->addr, peer->addr_len); if (sent < 0) { @@ -665,6 +771,12 @@ peer_t* rootstream_add_peer(rootstream_ctx_t *ctx, const char *code) { } peer->state = PEER_DISCOVERED; + peer->video_tx_frame_id = 1; + peer->video_rx_frame_id = 0; + peer->video_rx_buffer = NULL; + peer->video_rx_capacity = 0; + peer->video_rx_expected = 0; + peer->video_rx_received = 0; ctx->num_peers++; char fingerprint[32]; diff --git a/src/service.c b/src/service.c index 9147a8b..02812c3 100644 --- a/src/service.c +++ b/src/service.c @@ -269,8 +269,8 @@ int service_run_host(rootstream_ctx_t *ctx) { peer_t *peer = &ctx->peers[i]; if (peer->state == PEER_CONNECTED && peer->is_streaming) { /* Send video */ - if (rootstream_net_send_encrypted(ctx, peer, PKT_VIDEO, - enc_buf, enc_size) < 0) { + if (enc_size > 0 && + rootstream_net_send_video(ctx, peer, enc_buf, enc_size) < 0) { fprintf(stderr, "ERROR: Video send failed (peer=%s)\n", peer->hostname); } @@ -370,8 +370,6 @@ int service_run_client(rootstream_ctx_t *ctx) { /* Check if we received a video frame */ if (ctx->current_frame.data && ctx->current_frame.size > 0) { - ctx->frames_received++; - /* Decode frame */ uint64_t decode_start_us = get_timestamp_us(); if (rootstream_decode_frame(ctx, ctx->current_frame.data, diff --git a/src/vaapi_decoder.c b/src/vaapi_decoder.c index 2626d17..26ba5a4 100644 --- a/src/vaapi_decoder.c +++ b/src/vaapi_decoder.c @@ -320,15 +320,17 @@ int rootstream_decode_frame(rootstream_ctx_t *ctx, out->timestamp = get_timestamp_us(); /* Allocate output buffer if needed */ - if (!out->data) { - out->data = malloc(out->size); - if (!out->data) { + if (!out->data || out->capacity < out->size) { + uint8_t *new_buf = realloc(out->data, out->size); + if (!new_buf) { fprintf(stderr, "ERROR: Cannot allocate output buffer\n"); vaUnmapBuffer(dec->display, image.buf); vaDestroyImage(dec->display, image.image_id); vaDestroyBuffer(dec->display, slice_data_buf); return -1; } + out->data = new_buf; + out->capacity = out->size; } /* Copy pixel data */