From 4fc96f2461d94c0fac76a1b2aa321787f8d17ed9 Mon Sep 17 00:00:00 2001
From: Iwan Timmer <irtimmer@gmail.com>
Date: Wed, 12 Aug 2015 17:13:26 +0200
Subject: [PATCH] Fix heap corruption in SPS

---
 libgamestream/sps.c | 29 +++++++++++++++++------------
 libgamestream/sps.h |  2 +-
 src/video/omx.c     |  3 ++-
 3 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/libgamestream/sps.c b/libgamestream/sps.c
index 91742b5..5b9e5da 100644
--- a/libgamestream/sps.c
+++ b/libgamestream/sps.c
@@ -32,26 +32,24 @@ void gs_sps_init(int width, int height) {
   initial_height = height;
 }
 
-PLENTRY gs_sps_fix(PLENTRY entry, int flags) {
+PLENTRY gs_sps_fix(PLENTRY *head, int flags) {
+  PLENTRY entry = *head;
   if (replay_sps == 1) {
-    PLENTRY replay_entry = malloc(sizeof(LENTRY));
-    char *buffer = malloc(128);
-    if (replay_entry == NULL || buffer == NULL)
+    PLENTRY replay_entry = (PLENTRY) malloc(sizeof(*replay_entry) + 128);
+    if (replay_entry == NULL)
       return NULL;
 
+    replay_entry->data = (char *) (entry + 1);
     char spsData[] = {0x00, 0x00, 0x00, 0x01, 0x67};
-    memcpy(buffer, spsData, sizeof(spsData));
+    memcpy(replay_entry->data, spsData, sizeof(spsData));
     h264_stream->sps->profile_idc = H264_PROFILE_HIGH;
-    replay_entry->length = write_nal_unit(h264_stream, buffer+4, 124) + 4;
+    replay_entry->length = write_nal_unit(h264_stream, replay_entry->data+4, 124) + 4;
 
-    replay_entry->data = buffer;
     replay_entry->next = entry;
     entry = replay_entry;
     replay_sps = 2;
-    return replay_entry;
   } else if ((entry->data[4] & 0x1F) == NAL_UNIT_TYPE_SPS) {
     read_nal_unit(h264_stream, entry->data+4, entry->length-4);
-    free(entry->data);
 
     // Some decoders rely on H264 level to decide how many buffers are needed
     // Since we only need one frame buffered, we'll set level as low as we can
@@ -96,15 +94,22 @@ PLENTRY gs_sps_fix(PLENTRY entry, int flags) {
     if ((flags & GS_SPS_BASELINE_HACK) == GS_SPS_BASELINE_HACK && !replay_sps)
       h264_stream->sps->profile_idc = H264_PROFILE_BASELINE;
 
-    entry->data = malloc(128);
-    if (entry->data == NULL)
+    PLENTRY sps_entry = (PLENTRY) malloc(sizeof(*sps_entry) + 128);
+    if (sps_entry == NULL)
       return NULL;
 
-    entry->length = write_nal_unit(h264_stream, entry->data+4, 124) + 4;
+    PLENTRY next = entry->next;
+    free(entry);
+    sps_entry->data = (char*) (entry + 1);
+    sps_entry->length = write_nal_unit(h264_stream, sps_entry->data+4, 124) + 4;
+    printf("Writen %d\n", sps_entry->length);
+    sps_entry->next = next;
+    entry = sps_entry;
   } else if ((entry->data[4] & 0x1F) == NAL_UNIT_TYPE_PPS) {
     if ((flags & GS_SPS_BASELINE_HACK) == GS_SPS_BASELINE_HACK && !replay_sps)
       replay_sps = 1;
 
   }
+  *head = entry;
   return entry;
 }
diff --git a/libgamestream/sps.h b/libgamestream/sps.h
index 0076205..bd2cf46 100644
--- a/libgamestream/sps.h
+++ b/libgamestream/sps.h
@@ -23,4 +23,4 @@
 #define GS_SPS_BASELINE_HACK 0x02
 
 void gs_sps_init();
-PLENTRY gs_sps_fix(PLENTRY entry, int flags);
+PLENTRY gs_sps_fix(PLENTRY *entry, int flags);
diff --git a/src/video/omx.c b/src/video/omx.c
index 4faabeb..35b062a 100644
--- a/src/video/omx.c
+++ b/src/video/omx.c
@@ -176,7 +176,8 @@ static int decoder_renderer_submit_decode_unit(PDECODE_UNIT decodeUnit) {
     first_packet = 0;
   }
 
-  PLENTRY entry = gs_sps_fix(decodeUnit->bufferList, GS_SPS_BITSTREAM_FIXUP);
+  PLENTRY entry = gs_sps_fix(&decodeUnit->bufferList, GS_SPS_BITSTREAM_FIXUP);
+  decodeUnit->bufferList = entry;
   while (entry != NULL) {
     memcpy(dest, entry->data, entry->length);
     buf->nFilledLen += entry->length;