1 files changed, 55 insertions, 36 deletions

diff --git a/src/gcontainer/gcontainer.cc b/src/gcontainer/gcontainer.cc
index 1179778..b97fdd7 100644
--- a/src/gcontainer/gcontainer.cc
+++ b/src/gcontainer/gcontainer.cc
@@ -4,6 +4,7 @@
 
 #include "image_io/base/data_segment.h"
 #include "image_io/base/data_segment_data_source.h"
+#include "image_io/base/istream_data_source.h"
 #include "image_io/base/message_handler.h"
 #include "image_io/base/ostream_data_destination.h"
 #include "image_io/jpeg/jpeg_info.h"
@@ -19,44 +20,44 @@ namespace {
 using photos_editing_formats::image_io::DataRange;
 using photos_editing_formats::image_io::DataSegment;
 using photos_editing_formats::image_io::DataSegmentDataSource;
+using photos_editing_formats::image_io::IStreamRefDataSource;
 using photos_editing_formats::image_io::JpegInfoBuilder;
 using photos_editing_formats::image_io::JpegScanner;
 using photos_editing_formats::image_io::Message;
 using photos_editing_formats::image_io::MessageHandler;
 using photos_editing_formats::image_io::OStreamDataDestination;
-using photos_editing_formats::image_io::ReportErrorPolicy;
 using std::string;
 
 // Populates first_image_range with the first image (from the header metadata
 // to the EOI marker) present in the JPEG file input_file_name. Returns true if
 // such a first image is found, false otherwise.
 //
-// input_file_name must be a JPEG file.
+// input_jpeg_stream must be a JPEG stream.
 // image_data_segment is populated with the DataSegment for
 // input_file_name, and is populated only in the successful case.
 // first_image_range is populated with the first image found in the input file,
 // only if such an image is found.
-bool ExtractFirstImageInJpeg(const string& input_file_name,
-                             std::shared_ptr<DataSegment>* image_data_segment,
+
+bool ExtractFirstImageInJpeg(std::istream& input_jpeg_stream,
+                             MessageHandler* message_handler,
                              DataRange* first_image_range) {
   if (first_image_range == nullptr) {
     return false;
   }
 
   // Get the input and output setup.
-  MessageHandler::Get()->ClearMessages();
-  auto data_segment =
-      ReadEntireFile(input_file_name, ReportErrorPolicy::kReportError);
-  if (!data_segment) {
-    return false;
+  if (message_handler) {
+    message_handler->ClearMessages();
   }
 
   // Get the jpeg info and first image range from the input.
-  DataSegmentDataSource data_source(data_segment);
+  IStreamRefDataSource data_source(input_jpeg_stream);
   JpegInfoBuilder jpeg_info_builder;
   jpeg_info_builder.SetImageLimit(1);
-  JpegScanner jpeg_scanner;
+  JpegScanner jpeg_scanner(message_handler);
   jpeg_scanner.Run(&data_source, &jpeg_info_builder);
+  data_source.Reset();
+
   if (jpeg_scanner.HasError()) {
     return false;
   }
@@ -64,12 +65,13 @@ bool ExtractFirstImageInJpeg(const string& input_file_name,
   const auto& jpeg_info = jpeg_info_builder.GetInfo();
   const auto& image_ranges = jpeg_info.GetImageRanges();
   if (image_ranges.empty()) {
-    MessageHandler::Get()->ReportMessage(Message::kPrematureEndOfDataError,
-                                         "No Images Found");
+    if (message_handler) {
+      message_handler->ReportMessage(Message::kPrematureEndOfDataError,
+                                     "No Images Found");
+    }
     return false;
   }
 
-  *image_data_segment = data_segment;
   *first_image_range = image_ranges[0];
   return true;
 }
@@ -79,23 +81,27 @@ bool ExtractFirstImageInJpeg(const string& input_file_name,
 bool WriteImageAndFiles(const string& input_file_name,
                         const std::vector<string>& other_files,
                         const string& output_file_name) {
-  auto output_stream =
-      OpenOutputFile(output_file_name, ReportErrorPolicy::kReportError);
+  MessageHandler message_handler;
+  auto output_stream = OpenOutputFile(output_file_name, &message_handler);
   if (!output_stream) {
     return false;
   }
 
-  OStreamDataDestination output_destination(std::move(output_stream));
+  OStreamDataDestination output_destination(std::move(output_stream),
+                                            &message_handler);
   output_destination.SetName(output_file_name);
 
   DataRange image_range;
-  std::shared_ptr<DataSegment> data_segment;
-  if (!ExtractFirstImageInJpeg(input_file_name, &data_segment, &image_range)) {
+  std::unique_ptr<std::istream> input_stream =
+      OpenInputFile(input_file_name, &message_handler);
+
+  if (!ExtractFirstImageInJpeg(*input_stream, &message_handler, &image_range)) {
     return false;
   }
 
   output_destination.StartTransfer();
-  DataSegmentDataSource data_source(data_segment);
+  IStreamDataSource data_source(
+      OpenInputFile(input_file_name, &message_handler));
   data_source.TransferData(image_range, image_range.GetLength(),
                            &output_destination);
 
@@ -104,8 +110,7 @@ bool WriteImageAndFiles(const string& input_file_name,
     if (tack_on_file.empty()) {
       continue;
     }
-    auto tack_on_data_segment =
-        ReadEntireFile(tack_on_file, ReportErrorPolicy::kReportError);
+    auto tack_on_data_segment = ReadEntireFile(tack_on_file, &message_handler);
     if (!tack_on_data_segment) {
       continue;
     }
@@ -122,25 +127,40 @@ bool WriteImageAndFiles(const string& input_file_name,
          !output_destination.HasError();
 }
 
-bool ParseFileAfterImage(const string& input_file_name,
+bool ParseFileAfterImage(const std::string& input_file_name,
                          size_t file_start_offset, size_t file_length,
-                         string* out_file_contents) {
-  if (out_file_contents == nullptr || file_start_offset < 0 ||
-      file_length == 0) {
+                         std::string* out_file_contents) {
+  std::ifstream input_stream(input_file_name);
+  if (!input_stream.is_open()) {
     return false;
   }
+  return ParseFileAfterImageFromStream(file_start_offset, file_length,
+                                       input_stream, out_file_contents);
+}
+
+bool ParseFileAfterImageFromStream(size_t start_offset, size_t length,
+                                   std::istream& input_jpeg_stream,
+                                   std::string* out_contents) {
+  if (out_contents == nullptr || start_offset < 0 || length == 0) {
+    return false;
+  }
+
+  size_t curr_posn = input_jpeg_stream.tellg();
+  input_jpeg_stream.seekg(0, input_jpeg_stream.end);
+  size_t stream_size = input_jpeg_stream.tellg();
+  input_jpeg_stream.seekg(curr_posn, input_jpeg_stream.beg);
 
   DataRange image_range;
-  std::shared_ptr<DataSegment> data_segment;
-  if (!ExtractFirstImageInJpeg(input_file_name, &data_segment, &image_range)) {
+  MessageHandler message_handler;
+  if (!ExtractFirstImageInJpeg(input_jpeg_stream, &message_handler,
+                               &image_range)) {
     return false;
   }
 
   size_t image_bytes_end_offset = image_range.GetEnd();
-  size_t image_file_end = data_segment->GetEnd();
-  size_t file_start_in_image = image_bytes_end_offset + file_start_offset;
-  size_t file_end_in_image = file_start_in_image + file_length;
-  if (image_file_end < file_end_in_image) {
+  size_t file_start_in_image = image_bytes_end_offset + start_offset;
+  size_t file_end_in_image = file_start_in_image + length;
+  if (stream_size < file_end_in_image) {
     // Requested file is past the end of the image file.
     return false;
   }
@@ -151,10 +171,9 @@ bool ParseFileAfterImage(const string& input_file_name,
   // TODO(miraleung): Consider subclassing image_io/data_destination.h and
   // transferring bytes directly into the string. TBD pending additional mime
   // type getters.
-  std::ifstream input_file_stream(input_file_name);
-  input_file_stream.seekg(file_range.GetBegin());
-  out_file_contents->resize(file_range_size);
-  input_file_stream.read(&(*out_file_contents)[0], file_range_size);
+  input_jpeg_stream.seekg(file_range.GetBegin(), input_jpeg_stream.beg);
+  out_contents->resize(file_range_size);
+  input_jpeg_stream.read(&(*out_contents)[0], file_range_size);
   return true;
 }