Image Size Filtering for PhotoRec

src/Makefile.am

@@ -50,6 +50,7 @@ testdisk_ncurses_H	= addpart.h addpartn.h adv.h askloc.h chgarch.h chgarchn.h ch
 testdisk_SOURCES	= $(base_C) $(base_H) $(fs_C) $(fs_H) $(testdisk_ncurses_C) $(testdisk_ncurses_H) dir.c dir.h dir_common.h exfat_dir.c exfat_dir.h ext2_dir.c ext2_dir.h ext2_inc.h fat_dir.c fat_dir.h ntfs_dir.c ntfs_dir.h ntfs_inc.h partgptw.c rfs_dir.c rfs_dir.h $(ICON_TESTDISK) next.c next.h
 
 file_C			= filegen.c \
+			  image_filter.c \
 			  file_list.c \
 			  file_1cd.c \
 			  file_3dm.c \
@@ -402,7 +403,7 @@ file_C			= filegen.c \
 			  file_zpr.c \
 			  utfsize.c
 
-file_H			= ext2.h hfsp_struct.h filegen.h file_doc.h file_jpg.h file_gz.h file_riff.h file_sp3.h file_tar.h file_tiff.h luks_struct.h ntfs_struct.h ole.h pe.h suspend.h utfsize.h xfs_struct.h
+file_H			= ext2.h hfsp_struct.h filegen.h image_filter.h file_doc.h file_jpg.h file_gz.h file_riff.h file_sp3.h file_tar.h file_tiff.h luks_struct.h ntfs_struct.h ole.h pe.h suspend.h utfsize.h xfs_struct.h
 
 photorec_C		= photorec.c phcfg.c addpart.c chgarch.c chgtype.c dir.c exfatp.c ext2grp.c ext2_dir.c ext2p.c fat_dir.c fatp.c file_found.c geometry.c ntfs_dir.c ntfsp.c pdisksel.c poptions.c sessionp.c dfxml.c partgptro.c json_log.c
 

src/file_jpg.c

@@ -60,6 +60,7 @@
 #endif
 #include "file_tiff.h"
 #include "setdate.h"
+#include "image_filter.h"
 #if defined(__FRAMAC__)
 #include "__fc_builtin.h"
 #endif
@@ -90,6 +91,7 @@ const file_hint_t file_hint_jpg= {
   .max_filesize=50*1024*1024,
   .recover=1,
   .enable_by_default=1,
+  .is_image=1,
   .register_header_check=&register_header_check_jpg
 };
 
@@ -863,6 +865,7 @@ static time_t jpg_get_date(const unsigned char *buffer, const unsigned int buffe
   return 0;
 }
 
+static int jpg_maches_image_filtering(const unsigned char *buffer, const unsigned int buffer_size, file_recovery_t *file_recovery);
 
 /*@
   @ requires PHOTOREC_MAX_BLOCKSIZE >= buffer_size >= 10;
@@ -1052,6 +1055,8 @@ static int header_check_jpg(const unsigned char *buffer, const unsigned int buff
   file_recovery_new->time=jpg_time;
   file_recovery_new->extension=file_hint_jpg.extension;
   file_recovery_new->file_check=&file_check_jpg;
+  file_recovery_new->file_check_presave=&jpg_maches_image_filtering;
+  file_recovery_new->image_filter=file_recovery->image_filter;
   if(buffer_size >= 4)
     file_recovery_new->data_check=&data_check_jpg;
   /*@ assert valid_read_string(file_recovery_new->extension); */
@@ -1926,6 +1931,58 @@ struct sof_header
 #endif
 } __attribute__ ((gcc_struct, __packed__));
 
+static int jpg_maches_image_filtering(const unsigned char *buffer, const unsigned int buffer_size, file_recovery_t *file_recovery)
+{
+  if(!file_recovery->image_filter)
+    return 1;
+
+  if(buffer_size < 20)
+    return 1;
+
+  if(!(buffer[0] == 0xff && buffer[1] == 0xd8 && buffer[2] == 0xff))
+    return 1;
+
+  const unsigned char *check_buffer = buffer;
+  unsigned int check_size = buffer_size;
+
+  uint64_t estimated_file_size = 0;
+  if(check_size > 2) {
+    for(unsigned int i = check_size - 2; i > 2; i--)
+    {
+      if(check_buffer[i] == 0xff && check_buffer[i+1] == 0xd9)
+      {
+        estimated_file_size = i + 2;
+        break;
+      }
+    }
+  }
+
+  // Apply file size filter if we found end marker
+  if(estimated_file_size > 0 && file_recovery->image_filter && should_skip_image_by_filesize(file_recovery->image_filter, estimated_file_size))
+    return 0;
+
+  // Check dimensions in the buffer we're examining
+  if(check_size > 10)
+  {
+    for(unsigned int i = 0; i < check_size - 10; i++)
+    {
+      if(check_buffer[i] == 0xff && check_buffer[i+1] == 0xc0)
+      {
+        if(i + 10 < check_size)
+        {
+          const struct sof_header *sof = (const struct sof_header *)&check_buffer[i];
+          const unsigned int width = be16(sof->width);
+          const unsigned int height = be16(sof->height);
+          if(file_recovery->image_filter && should_skip_image_by_dimensions(file_recovery->image_filter, width, height))
+            return 0;
+        }
+        break;
+      }
+    }
+  }
+  return 1;
+}
+
 /*@
   @ requires \valid_read(buffer + (0..buffer_size-1));
   @ terminates \true;
@@ -2299,6 +2356,10 @@ static int jpg_check_app1(file_recovery_t *file_recovery, const unsigned int ext
     return 1;
   if(thumb_offset+thumb_size > nbytes)
     return 1;
+
+  if (file_recovery->image_filter && !jpg_maches_image_filtering((const unsigned char *)(buffer + thumb_offset), thumb_size, file_recovery))
+    return 1;
+
   /*@ assert thumb_offset + thumb_size <= nbytes; */
   /*@ assert 0 < thumb_size; */
   /*@ assert thumb_offset < nbytes; */
@@ -2457,6 +2518,14 @@ static void file_check_jpg(file_recovery_t *file_recovery)
 #endif
   if(file_recovery->offset_error!=0)
     return ;
+// we could simply disable jpg thumb extraction when image filtering is active
+// but for now they're passed through jpg_maches_image_filtering like normal photos
+//   const unsigned int extract_thumb = file_recovery->image_filter ? 0 : 1;
+// #ifdef DEBUG_JPEG
+//   if(file_recovery->image_filter)
+//     log_info("skipping thumbnail extraction because image filtering is enabled\n");
+// #endif
+//   thumb_offset=jpg_check_structure(file_recovery, extract_thumb);
   thumb_offset=jpg_check_structure(file_recovery, 1);
 #ifdef DEBUG_JPEG
   log_info("jpg_check_structure error at %llu\n", (long long unsigned)file_recovery->offset_error);
@@ -2489,11 +2558,22 @@ static void file_check_jpg(file_recovery_t *file_recovery)
 #endif
   if(file_recovery->offset_error!=0)
     return ;
+
 #if defined(HAVE_LIBJPEG) && defined(HAVE_JPEGLIB_H)
   jpg_check_picture(file_recovery);
 #else
   file_recovery->file_size=file_recovery->calculated_file_size;
 #endif
+  if(file_recovery->image_filter && file_recovery->handle) {
+    fseek(file_recovery->handle, 0, SEEK_SET);
+    char buffer[512];
+    if(fread(buffer, 1, sizeof(buffer), file_recovery->handle) > 0) {
+      unsigned int width = 0, height = 0;
+      jpg_get_size((unsigned char*)buffer, sizeof(buffer), &height, &width);
+      file_recovery->image_data.width = width;
+      file_recovery->image_data.height = height;
+    }
+  }
 #if 0
     /* FIXME REMOVE ME */
   if(file_recovery->offset_error!=0)
@@ -2630,6 +2710,10 @@ static data_check_t data_check_jpg(const unsigned char *buffer, const unsigned i
   /* Skip the SOI */
   if(file_recovery->calculated_file_size<2)
     file_recovery->calculated_file_size=2;
+
+  if (file_recovery->image_filter && !jpg_maches_image_filtering(buffer, buffer_size, file_recovery))
+    return DC_STOP;
+
   /*@ assert file_recovery->calculated_file_size >= 2; */
   /*@ assert file_recovery->data_check == &data_check_jpg; */
   /* Search SOS */

src/file_png.c

@@ -38,6 +38,7 @@
 #include "types.h"
 #include "common.h"
 #include "filegen.h"
+#include "image_filter.h"
 
 extern const file_hint_t file_hint_doc;
 
@@ -50,6 +51,7 @@ const file_hint_t file_hint_png= {
   .max_filesize=PHOTOREC_MAX_FILE_SIZE,
   .recover=1,
   .enable_by_default=1,
+  .is_image=1,
   .register_header_check=&register_header_check_png
 };
 
@@ -167,6 +169,10 @@ static void file_check_png(file_recovery_t *fr)
 	fr->file_size=0;
 	return ;
       }
+      if(fr->image_filter) {
+        fr->image_data.width = be32(ihdr->width);
+        fr->image_data.height = be32(ihdr->height);
+      }
     }
   }
 }
@@ -313,6 +319,49 @@ static int header_check_mng(const unsigned char *buffer, const unsigned int buff
   return 1;
 }
 
+static int png_maches_image_filtering(const unsigned char *buffer, const unsigned int buffer_size, file_recovery_t *file_recovery)
+{
+  if(!file_recovery->image_filter)
+    return 1;
+
+  if(buffer_size < 24)
+    return 1;
+
+  if(!(buffer[0] == 0x89 && buffer[1] == 'P' && buffer[2] == 'N' && buffer[3] == 'G'))
+    return 1;
+
+  const unsigned char *check_buffer = buffer;
+  unsigned int check_size = buffer_size;
+
+  // Estimate file size by finding PNG IEND chunk
+  uint64_t estimated_file_size = 0;
+  if(check_size > 12) {
+    for(unsigned int i = check_size - 12; i > 8; i--)
+    {
+      if(check_buffer[i] == 'I' && check_buffer[i+1] == 'E' && check_buffer[i+2] == 'N' && check_buffer[i+3] == 'D')
+      {
+        estimated_file_size = i + 8; // IEND chunk + 4-byte CRC
+        break;
+      }
+    }
+  }
+
+  // Apply file size filter if we found IEND
+  if(estimated_file_size > 0 && file_recovery->image_filter && should_skip_image_by_filesize(file_recovery->image_filter, estimated_file_size))
+    return 0;
+
+  if(check_size >= 16 + sizeof(struct png_ihdr))
+  {
+    const struct png_ihdr *ihdr = (const struct png_ihdr *)&check_buffer[16];
+    const unsigned int width = be32(ihdr->width);
+    const unsigned int height = be32(ihdr->height);
+    if(should_skip_image_by_dimensions(file_recovery->image_filter, width, height))
+      return 0;
+  }
+
+  return 1;
+}
+
 /*@
   @ requires buffer_size >= 16 + sizeof(struct png_ihdr);
   @ requires separation: \separated(&file_hint_png, buffer+(..), file_recovery, file_recovery_new);
@@ -350,6 +399,8 @@ static int header_check_png(const unsigned char *buffer, const unsigned int buff
   file_recovery_new->calculated_file_size=8;
   file_recovery_new->data_check=&data_check_png;
   file_recovery_new->file_check=&file_check_png;
+  file_recovery_new->file_check_presave=&png_maches_image_filtering;
+  file_recovery_new->image_filter=file_recovery->image_filter;
   /*@ assert valid_file_recovery(file_recovery_new); */
   return 1;
 }