feat: rtmp chunk write header alloc memory

ireader · Oct 20, 2024 · 601a372 · 601a372
1 parent 95fe593
commit 601a372
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 2 deletions.
diff --git a/librtmp/include/rtmp-internal.h b/librtmp/include/rtmp-internal.h
@@ -90,6 +90,13 @@ struct rtmp_parser_t
 	struct rtmp_packet_t* pkt;
 };
 
+struct rtmp_chunk_write_header_t
+{
+	uint8_t bufer[MAX_CHUNK_HEADER + 5 * 20]; // 1-full chunk header + 20-type3 chunk header(with timestamp)
+	uint8_t* ptr;
+	uint32_t capacity;
+};
+
 struct rtmp_t
 {
 	uint32_t in_chunk_size; // read from network
@@ -108,6 +115,8 @@ struct rtmp_t
 	struct rtmp_packet_t out_packets[N_CHUNK_STREAM]; // send to network
 	struct rtmp_parser_t parser;
 
+	struct rtmp_chunk_write_header_t chunk_write_header; // rtmp_chunk_write only
+
 	void* param;
 
 	/// @return 0-ok, other-error

diff --git a/librtmp/source/rtmp-chunk-write.c b/librtmp/source/rtmp-chunk-write.c
@@ -1,10 +1,13 @@
 #include "rtmp-chunk-header.h"
 #include "rtmp-internal.h"
 #include "rtmp-util.h"
+#include <stdlib.h>
 #include <string.h>
 #include <assert.h>
 #include <errno.h>
 
+#define MAX_CHECK_TYPE3_WITH_TIMESTAMP 7
+
 // 5.3.1. Chunk Format (p11)
 /* 3-bytes basic header + 11-bytes message header + 4-bytes extended timestamp */
 //#define MAX_CHUNK_HEADER 18
@@ -89,20 +92,33 @@ static const struct rtmp_chunk_header_t* rtmp_chunk_header_zip(struct rtmp_t* rt
 int rtmp_chunk_write(struct rtmp_t* rtmp, const struct rtmp_chunk_header_t* h, const uint8_t* payload)
 {
 	int r = 0;
-	uint8_t p[MAX_CHUNK_HEADER];
-	uint32_t chunkSize, headerSize, payloadSize;
+	uint8_t *p;
+	uint32_t chunkSize, headerSize, payloadSize, N;
 	const struct rtmp_chunk_header_t* header;
 
 	// compression rtmp chunk header
 	header = rtmp_chunk_header_zip(rtmp, h);
 	if (!header || header->length >= 0xFFFFFF)
 		return -EINVAL; // invalid length
 
+	// alloc chunk header
+	N = header->cid < 64 ? 1 : (header->cid < (64 + 256) ? 2 : 3) /*type-3 header*/ + header->timestamp >= 0xFFFFFF ? 4 : 0;
+	if (MAX_CHUNK_HEADER + header->length / rtmp->out_chunk_size * N > rtmp->chunk_write_header.capacity)
+	{
+		p = realloc(rtmp->chunk_write_header.ptr == rtmp->chunk_write_header.bufer ? 0 : rtmp->chunk_write_header.ptr, MAX_CHUNK_HEADER + header->length / rtmp->out_chunk_size * MAX_CHECK_TYPE3_WITH_TIMESTAMP);
+		if (!p)
+			return -ENOMEM;
+		rtmp->chunk_write_header.ptr = p;
+		rtmp->chunk_write_header.capacity = MAX_CHUNK_HEADER + header->length / rtmp->out_chunk_size * MAX_CHECK_TYPE3_WITH_TIMESTAMP;
+	}
+	p = rtmp->chunk_write_header.ptr;
+
 	payloadSize = header->length;
 	headerSize = rtmp_chunk_basic_header_write(p, header->fmt, header->cid);
 	headerSize += rtmp_chunk_message_header_write(p + headerSize, header);
 	if(header->timestamp >= 0xFFFFFF)
 		headerSize += rtmp_chunk_extended_timestamp_write(p + headerSize, header->timestamp);
+	assert(headerSize < MAX_CHUNK_HEADER);
 
 	while (payloadSize > 0 && 0 == r)
 	{
@@ -114,9 +130,12 @@ int rtmp_chunk_write(struct rtmp_t* rtmp, const struct rtmp_chunk_header_t* h, c
 
 		if (payloadSize > 0)
 		{
+			p += headerSize; // next chunk header
+			assert(p + N <= rtmp->chunk_write_header.ptr + rtmp->chunk_write_header.capacity);
 			headerSize = rtmp_chunk_basic_header_write(p, RTMP_CHUNK_TYPE_3, header->cid);
 			if (header->timestamp >= 0xFFFFFF)
 				headerSize += rtmp_chunk_extended_timestamp_write(p + headerSize, header->timestamp);
+			assert(headerSize == N);
 		}
 	}
 

diff --git a/librtmp/source/rtmp-client.c b/librtmp/source/rtmp-client.c
@@ -398,6 +398,9 @@ struct rtmp_client_t* rtmp_client_create(const char* appname, const char* playpa
 	ctx->rtmp.out_packets[RTMP_CHANNEL_AUDIO].header.cid = RTMP_CHANNEL_AUDIO;
 	ctx->rtmp.out_packets[RTMP_CHANNEL_VIDEO].header.cid = RTMP_CHANNEL_VIDEO;
 	ctx->rtmp.out_packets[RTMP_CHANNEL_DATA].header.cid = RTMP_CHANNEL_DATA;
+
+	ctx->rtmp.chunk_write_header.ptr = ctx->rtmp.chunk_write_header.bufer;
+	ctx->rtmp.chunk_write_header.capacity = sizeof(ctx->rtmp.chunk_write_header.bufer);
 	return ctx;
 }
 
@@ -418,6 +421,12 @@ void rtmp_client_destroy(struct rtmp_client_t* ctx)
 		}
 	}
 
+	if (ctx->rtmp.chunk_write_header.ptr && ctx->rtmp.chunk_write_header.ptr != ctx->rtmp.chunk_write_header.bufer)
+	{
+		free(ctx->rtmp.chunk_write_header.ptr);
+		ctx->rtmp.chunk_write_header.capacity = 0;
+	}
+
 #if defined(DEBUG) || defined(_DEBUG)
 	memset(ctx, 0xCC, sizeof(*ctx));
 #endif

diff --git a/librtmp/source/rtmp-server.c b/librtmp/source/rtmp-server.c
@@ -441,6 +441,9 @@ struct rtmp_server_t* rtmp_server_create(void* param, const struct rtmp_server_h
 	ctx->rtmp.out_packets[RTMP_CHANNEL_AUDIO].header.cid = RTMP_CHANNEL_AUDIO;
 	ctx->rtmp.out_packets[RTMP_CHANNEL_VIDEO].header.cid = RTMP_CHANNEL_VIDEO;
 	ctx->rtmp.out_packets[RTMP_CHANNEL_DATA].header.cid = RTMP_CHANNEL_DATA;
+
+	ctx->rtmp.chunk_write_header.ptr = ctx->rtmp.chunk_write_header.bufer;
+	ctx->rtmp.chunk_write_header.capacity = sizeof(ctx->rtmp.chunk_write_header.bufer);
 	return ctx;
 }
 
@@ -455,6 +458,12 @@ void rtmp_server_destroy(struct rtmp_server_t* ctx)
 			free(ctx->rtmp.in_packets[i].payload);
 	}
 
+	if (ctx->rtmp.chunk_write_header.ptr && ctx->rtmp.chunk_write_header.ptr != ctx->rtmp.chunk_write_header.bufer)
+	{
+		free(ctx->rtmp.chunk_write_header.ptr);
+		ctx->rtmp.chunk_write_header.capacity = 0;
+	}
+
 	free(ctx);
 }