From f0880ef4853089ce816c49219332792c139c7487 Mon Sep 17 00:00:00 2001
From: tao3 <tao3@outlook.com>
Date: Sat, 7 Nov 2020 09:24:50 +0800
Subject: [PATCH] feat: TS dynamic alloc PMT

---
 libmpeg/include/mpeg-ts-proto.h |  5 ++-
 libmpeg/source/mpeg-pat.c       | 57 ++++++++++++++++++++++++++-------
 libmpeg/source/mpeg-ts-dec.c    |  3 ++
 libmpeg/source/mpeg-ts-enc.c    | 15 +++++----
 4 files changed, 61 insertions(+), 19 deletions(-)

diff --git a/libmpeg/include/mpeg-ts-proto.h b/libmpeg/include/mpeg-ts-proto.h
index aa8d9d0d..bad81065 100644
--- a/libmpeg/include/mpeg-ts-proto.h
+++ b/libmpeg/include/mpeg-ts-proto.h
@@ -81,7 +81,9 @@ struct pat_t
 	unsigned int cc;	//continuity_counter : 4;
 
 	unsigned int pmt_count;
-	struct pmt_t pmts[4];
+	unsigned int pmt_capacity;
+	struct pmt_t pmt_default[1];
+	struct pmt_t* pmts;
 };
 
 // Table 2-3 - PID table(p36)
@@ -187,6 +189,7 @@ enum
     MPEG_FLAG_H264_H265_WITH_AUD = 0x8000,
 };
 
+struct pmt_t* pat_alloc_pmt(struct pat_t* pat);
 struct pmt_t* pat_find(struct pat_t* pat, uint16_t pn);
 size_t pat_read(struct pat_t *pat, const uint8_t* data, size_t bytes);
 size_t pat_write(const struct pat_t *pat, uint8_t *data);
diff --git a/libmpeg/source/mpeg-pat.c b/libmpeg/source/mpeg-pat.c
index 43150ede..b68da454 100644
--- a/libmpeg/source/mpeg-pat.c
+++ b/libmpeg/source/mpeg-pat.c
@@ -4,25 +4,61 @@
 
 #include "mpeg-ts-proto.h"
 #include "mpeg-util.h"
+#include <stdlib.h>
+#include <string.h>
 #include <assert.h>
 
+struct pmt_t* pat_alloc_pmt(struct pat_t* pat)
+{
+	void* ptr;
+	unsigned int n;
+
+	if (NULL == pat->pmts)
+	{
+		assert(0 == pat->pmt_count);
+		assert(0 == pat->pmt_capacity);
+		pat->pmts = pat->pmt_default;
+		pat->pmt_capacity = sizeof(pat->pmt_default) / sizeof(pat->pmt_default[0]);
+	}
+
+	if (pat->pmt_count >= pat->pmt_capacity)
+	{
+		if (pat->pmt_count + 1 > 65535)
+		{
+			assert(0);
+			return NULL;
+		}
+
+		n = pat->pmt_capacity + pat->pmt_capacity / 4 + 4;
+		ptr = realloc(pat->pmts == pat->pmt_default ? NULL : pat->pmts, sizeof(pat->pmts[0]) * n);
+		if (!ptr)
+			return NULL;
+
+		if (pat->pmts == pat->pmt_default)
+			memmove(ptr, pat->pmts, sizeof(pat->pmt_default));
+		pat->pmts = (struct pmt_t*)ptr;
+		pat->pmt_capacity = n;
+	}
+
+	// new pmt
+	memset(&pat->pmts[pat->pmt_count], 0, sizeof(pat->pmts[0]));
+	return &pat->pmts[pat->pmt_count];
+}
+
 static struct pmt_t* pat_fetch(struct pat_t* pat, uint16_t pid)
 {
-    unsigned int i;
-    for(i = 0; i < pat->pmt_count; i++)
+	unsigned int i;
+	struct pmt_t* pmt;
+	for(i = 0; i < pat->pmt_count; i++)
     {
         if(pat->pmts[i].pid == pid)
             return &pat->pmts[i];
     }
-    
-    if(pat->pmt_count >= sizeof(pat->pmts) / sizeof(pat->pmts[0]))
-    {
-        assert(0);
-        return NULL;
-    }
-    
+	
     // new pmt
-    return &pat->pmts[pat->pmt_count++];
+	pmt = pat_alloc_pmt(pat);
+	pat->pmt_count++;
+	return pmt;
 }
 
 size_t pat_read(struct pat_t *pat, const uint8_t* data, size_t bytes)
@@ -65,7 +101,6 @@ size_t pat_read(struct pat_t *pat, const uint8_t* data, size_t bytes)
         
         if(0 == pn)
             continue; // ignore NIT info
-        assert(pat->pmt_count <= sizeof(pat->pmts)/sizeof(pat->pmts[0]));
         pmt = pat_fetch(pat, pid);
         if(NULL == pmt)
             continue;
diff --git a/libmpeg/source/mpeg-ts-dec.c b/libmpeg/source/mpeg-ts-dec.c
index 9430c243..94accdff 100644
--- a/libmpeg/source/mpeg-ts-dec.c
+++ b/libmpeg/source/mpeg-ts-dec.c
@@ -303,6 +303,9 @@ int ts_demuxer_destroy(struct ts_demuxer_t* ts)
         }
     }
 
+	if (ts->pat.pmts && ts->pat.pmts != ts->pat.pmt_default)
+		free(ts->pat.pmts);
+
     free(ts);
     return 0;
 }
diff --git a/libmpeg/source/mpeg-ts-enc.c b/libmpeg/source/mpeg-ts-enc.c
index 8afb18e7..9372befa 100644
--- a/libmpeg/source/mpeg-ts-enc.c
+++ b/libmpeg/source/mpeg-ts-enc.c
@@ -365,6 +365,8 @@ int mpeg_ts_destroy(void* ts)
 		mpeg_ts_pmt_destroy(pmt);
 	}
 
+	if (tsctx->pat.pmts && tsctx->pat.pmts != tsctx->pat.pmt_default)
+		free(tsctx->pat.pmts);
 	free(tsctx);
 	return 0;
 }
@@ -385,6 +387,9 @@ int mpeg_ts_add_program(void* ts, uint16_t pn, const void* info, int bytes)
 	struct pmt_t* pmt;
 	mpeg_ts_enc_context_t* tsctx;
 
+	if (pn < 1 || bytes < 0 || bytes >= (1 << 12))
+		return -1; // EINVAL: pminfo-len 12-bits
+
 	tsctx = (mpeg_ts_enc_context_t*)ts;
 	for (i = 0; i < tsctx->pat.pmt_count; i++)
 	{
@@ -393,15 +398,11 @@ int mpeg_ts_add_program(void* ts, uint16_t pn, const void* info, int bytes)
 			return -1; // EEXIST
 	}
 
-	if (i >= sizeof(tsctx->pat.pmts) / sizeof(tsctx->pat.pmts[0]))
+	assert(tsctx->pat.pmt_count == i);
+	pmt = pat_alloc_pmt(&tsctx->pat);
+	if (!pmt)
 		return -1; // E2BIG
 
-	if (pn < 1 || bytes < 0 || bytes >= (1 << 12))
-		return -1; // EINVAL: pminfo-len 12-bits
-
-	assert(tsctx->pat.pmt_count == i);
-	pmt = &tsctx->pat.pmts[tsctx->pat.pmt_count];
-	memset(pmt, 0, sizeof(*pmt));
 	pmt->pid = tsctx->pid++;
 	pmt->pn = pn;
 	pmt->ver = 0x00;