PhotoRec: detect MPEG filesize

src/file_mpg.c

@@ -29,6 +29,7 @@
 #include <stdio.h>
 #include "types.h"
 #include "filegen.h"
+#include "log.h"
 
 static void register_header_check_mpg(file_stat_t *file_stat);
 
@@ -42,36 +43,125 @@ const file_hint_t file_hint_mpg= {
   .register_header_check=&register_header_check_mpg
 };
 
+static int calculate_packet_size(const unsigned char *buffer)
+{
+  /* http://dvd.sourceforge.net/dvdinfo/mpeghdrs.html */
+  if(buffer[0]!=0 || buffer[1]!=0 || buffer[2]!=1)
+    return 0;
+  switch(buffer[3])
+  {
+    /* Pack header: */
+    case 0xBA:
+      if((buffer[4]&0xc4)==0x44 && (buffer[6]&4)==4 && (buffer[8]&4)==4 &&
+	  (buffer[9]&1)==1 && (buffer[12]&3)==3)
+	return (buffer[13] & 0x7) + 14;
+      if((buffer[4]&0xF1)==0x21 && (buffer[6]&1)==1 && (buffer[8]&1)==1 &&
+	  (buffer[9]&0x80)==0x80 && (buffer[11]&1)==1)
+	return 12;
+      return 0;
+      /* Sequence Header */
+    case 0xB3:
+      if((buffer[10]&0x20)==0x20)
+      {
+	if((buffer[11]&3)!=0)
+	  return 12+64;	/* quantiser matrix */
+	return 12;
+      }
+      return 0;
+      /* Extension */
+    case 0xB5:
+      /* Sequence_Extension */
+      if((buffer[4]&0xF0)==0x10 && (buffer[7]&1)==1)	
+	return 10;
+      /* Sequence_Display_Extension without color description */
+      if((buffer[4]&0xF0)==0x20 && (buffer[4]&1)==0 && (buffer[6]&2)==2)	
+	return 9;
+      /* Sequence_Display_Extension with color description */
+      if((buffer[4]&0xF0)==0x20 && (buffer[4]&1)==1 && (buffer[9]&2)==2)	
+	return 12;
+      /* Picture_Coding_Extension */
+      if((buffer[4]&0xF0)==0x40)
+      {
+	if((buffer[8]&0x40)==0)
+	  return 9;
+	else
+	  return 11;
+      }
+      return 0;
+    case 0xB8: /* Group of Pictures */
+      if((buffer[5]&0x40)==0x40)
+	return 8;
+      return 0;
+    case 0xB9: /* EOC */
+      return 4;
+    case 0xBD:		/* Private stream 1 (non MPEG audio, subpictures) */
+    case 0xC0 ... 0xDF: /* Mpeg Audio stream */
+    case 0xE0 ... 0xEF: /* Mpeg Video stream */
+#if 0
+      {
+	uint32_t pts = 0;
+	// This is mpeg 2:
+	if((buffer[6] & 0xC0) == 0x80 &&
+	    // PTS DTS flags
+	    (buffer[7] >> 7)==1)
+	{
+	  pts = ((buffer[13] | (buffer[12] << 8) ) >> 1) |
+	    ((buffer[11] | (buffer[10] << 8) ) >> 1) << 15;
+
+	  //      log_debug("MPG2 (%u - 0x%02X)PTS is 0x%08X\n", current->id,buffer[3], pts);
+	  // This is mpeg 1. The PTS goes right after the header and must
+	  // have the bits 0x21 set:
+	}
+	else if((buffer[6] & 0x21)==0x21)
+	{
+	  pts = ((buffer[10] | (buffer[9] << 8) ) >> 1) | ((buffer[8] | (buffer[7] << 8) ) >> 1) << 15;
+	  //log_debug("MPG1 (%u - 0x%02X)PTS is 0x%08X\n", current->id,buffer[3], pts);
+	};
+      }
+#endif
+      return (buffer[4] << 8) + buffer[5] + 6;
+    case 0xBB:	/* System header */
+    case 0xBE:	/* Padding stream */
+    case 0xBF:	/* Private Stream 2 */
+      return (buffer[4] << 8) + buffer[5] + 6;
+    case 0:
+      return 0;
+    default:
+#ifdef DEBUG_MPG
+      log_info("I dont know how to handle 0x%02X\n", buffer[3]);
+#endif
+      return 0;
+  }
+}
+
+
 static data_check_t data_check_mpg(const unsigned char *buffer, const unsigned int buffer_size, file_recovery_t *file_recovery)
 {
-  const unsigned char padding_iso_end[8]=     {0xFF, 0xFF, 0xFF, 0xFF, 0x00, 0x00, 0x01, 0xB9};
+  while(file_recovery->calculated_file_size + buffer_size/2  >= file_recovery->file_size &&
-  /* search padding + end code */
+      file_recovery->calculated_file_size + 14 < file_recovery->file_size + buffer_size/2)
-  if(buffer_size>=8 && memcmp(&buffer[buffer_size/2-4], padding_iso_end, sizeof(padding_iso_end))==0)
   {
-    file_recovery->calculated_file_size=file_recovery->file_size+4;
+    const unsigned int i=file_recovery->calculated_file_size - file_recovery->file_size + buffer_size/2;
-    return DC_STOP;
+    const unsigned int ret=calculate_packet_size(&buffer[i]);
+#ifdef DEBUG_MPG
+    log_info("data_check_mpg %llu 0x%02x %u\n", (long long unsigned)file_recovery->calculated_file_size, buffer[i+3], ret);
+#endif
+    if(ret==0)
+      return DC_STOP;
+    file_recovery->calculated_file_size+=ret;
   }
-  /* search video sequence end followed by iso end code*/
-  if(buffer_size>=14)
-  {
-    unsigned int i;
-    for(i=buffer_size/2-7; i<buffer_size-7; i++)
-    {
-      const unsigned char sequence_end_iso_end[8]={0x00, 0x00, 0x01, 0xB7, 0x00, 0x00, 0x01, 0xB9};
-      if(buffer[i]==0x00 && memcmp(&buffer[i], sequence_end_iso_end, sizeof(sequence_end_iso_end))==0)
-      {
-	file_recovery->calculated_file_size=file_recovery->file_size+i+sizeof(sequence_end_iso_end)-buffer_size/2;
-	return DC_STOP;
-      }
-    }
-  }
-  /* some files don't end by iso end code, so continue... */
-  file_recovery->calculated_file_size=file_recovery->file_size+(buffer_size/2);
   return DC_CONTINUE;
 }
 
 static int header_check_mpg_Pack(const unsigned char *buffer, const unsigned int buffer_size, const unsigned int safe_header_only, const file_recovery_t *file_recovery, file_recovery_t *file_recovery_new)
 {
+  unsigned int i=0;
+  while(i<buffer_size && i+14<512)
+  {
+    const unsigned int ret=calculate_packet_size(&buffer[i]);
+    if(ret==0)
+      return 0;
+    i+=ret;
+  }
   /* MPEG-1 http://andrewduncan.ws/MPEG/MPEG-1.ps */
   /* pack start code 0x1BA + MPEG-1 + SCR=0 */
   if((buffer[4]&0xF1)==0x21 && (buffer[6]&1)==1 && (buffer[8]&1)==1 &&
@@ -95,7 +185,7 @@ static int header_check_mpg_Pack(const unsigned char *buffer, const unsigned int
   }
   /* MPEG-2 Program stream http://neuron2.net/library/mpeg2/iso13818-1.pdf */
   /* MPEG2 system header start code, several per file */
-  if((buffer[4]&0xc4)==0x44 && (buffer[6]&4)==4 && (buffer[8]&4)==4)
+  if((buffer[4]&0xc4)==0x44 && (buffer[6]&4)==4 && (buffer[8]&4)==4 && (buffer[9]&1)==1 && (buffer[12]&3)==3)
   {
     /*
      *       '01'                                            2	01
@@ -117,7 +207,7 @@ static int header_check_mpg_Pack(const unsigned char *buffer, const unsigned int
      */
 
     if(buffer[4]==0x44 && buffer[5]==0 && buffer[6]==4 && buffer[7]==0 && (buffer[8]&0xfc)==4)
-    {
+    { /* SCR=0 */
       reset_file_recovery(file_recovery_new);
       file_recovery_new->extension=file_hint_mpg.extension;
       file_recovery_new->data_check=&data_check_mpg;
@@ -160,6 +250,14 @@ static int header_check_mpg_System(const unsigned char *buffer, const unsigned i
   /* MPEG-1 system header start code */
   if((buffer[6]&0x80)==0x80 && (buffer[8]&0x01)==0x01 && buffer[11]==0xff)
   {
+    unsigned int i=0;
+    while(i<buffer_size && i+14<512)
+    {
+      const unsigned int ret=calculate_packet_size(&buffer[i]);
+      if(ret==0)
+	return 0;
+      i+=ret;
+    }
     reset_file_recovery(file_recovery_new);
     file_recovery_new->extension=file_hint_mpg.extension;
     file_recovery_new->data_check=&data_check_mpg;
@@ -173,7 +271,7 @@ static int header_check_mpg_Sequence(const unsigned char *buffer, const unsigned
 {
   if(file_recovery!=NULL && file_recovery->file_stat!=NULL && file_recovery->file_stat->file_hint==&file_hint_mpg)
     return 0;
-  /* MPEG-1 sequence header code */
+  /* MPEG-1 sequence header code 0x1B3 */
   /* horizontal size>0 */
   if((buffer[4]<<4)+(buffer[5]>>4)>0 &&
       /* vertical size>0 */
@@ -186,8 +284,15 @@ static int header_check_mpg_Sequence(const unsigned char *buffer, const unsigned
       (buffer[8]!=0 || buffer[9]!=0 || (buffer[10]&0xc0)!=0) &&
       /* marker */
       (buffer[10]&0x20)==0x20)
-
   {
+    unsigned int i=0;
+    while(i<buffer_size && i+14<512)
+    {
+      const unsigned int ret=calculate_packet_size(&buffer[i]);
+      if(ret==0)
+	return 0;
+      i+=ret;
+    }
     reset_file_recovery(file_recovery_new);
     file_recovery_new->extension=file_hint_mpg.extension;
     file_recovery_new->data_check=&data_check_mpg;
@@ -212,6 +317,14 @@ static int header_check_mpg4_ElemVideo(const unsigned char *buffer, const unsign
       (buffer[5]>>4)!=0 && (buffer[5]>>4)!=0x0f
     )
   {
+    unsigned int i=0;
+    while(i<buffer_size && i+14<512)
+    {
+      const unsigned int ret=calculate_packet_size(&buffer[i]);
+      if(ret==0)
+	return 0;
+      i+=ret;
+    }
     reset_file_recovery(file_recovery_new);
     file_recovery_new->extension=file_hint_mpg.extension;
     file_recovery_new->data_check=&data_check_mpg;