first cut at brotli support

dropbox · Sep 2, 2017 · d62a8c0 · d62a8c0
1 parent a501bfc
commit d62a8c0
Show file tree

Hide file tree

Showing 5 changed files with 53 additions and 19 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -220,6 +220,8 @@ set(LEPTON_SOURCES
    src/io/BufferedIO.hh
    src/io/ZlibCompression.cc
    src/io/ZlibCompression.hh
+   src/io/BrotliCompression.cc
+   src/io/BrotliCompression.hh
    src/io/Seccomp.hh
    src/io/Seccomp.cc
    src/io/seccomp-bpf.hh

diff --git a/Makefile.am b/Makefile.am
@@ -56,6 +56,8 @@ lepton_SOURCES = \
    src/io/BufferedIO.hh \
    src/io/ZlibCompression.cc \
    src/io/ZlibCompression.hh \
+   src/io/BrotliCompression.cc \
+   src/io/BrotliCompression.hh \
    src/io/Seccomp.hh \
    src/io/Seccomp.cc \
    src/io/seccomp-bpf.hh \

diff --git a/src/io/ZlibCompression.cc b/src/io/ZlibCompression.cc
@@ -74,7 +74,8 @@ std::vector<uint8_t,
     return retval;
 }
 std::pair<std::vector<uint8_t, JpegAllocator<uint8_t> >,
-          JpegError > ZlibDecoderDecompressionReader::Decompress(const uint8_t *buffer, size_t size, const JpegAllocator<uint8_t> &alloc) {
+          JpegError > ZlibDecoderDecompressionReader::Decompress(const uint8_t *buffer, size_t size, const JpegAllocator<uint8_t> &alloc,
+                                                                 size_t max_file_size) {
     z_stream strm;
     memset(&strm, 0, sizeof(z_stream));
     JpegAllocator<uint8_t> local_alloc;
@@ -113,7 +114,11 @@ std::pair<std::vector<uint8_t, JpegAllocator<uint8_t> >,
             }
             if (strm.avail_out == 0) {
                 retval_size += avail_bytes - strm.avail_out;
-                retval.first.resize(retval.first.size() * 2);
+                if (retval.first.size() == max_file_size) {
+                    retval.second = JpegError::errShortHuffmanData();
+                    break;
+                }
+                retval.first.resize(std::min(retval.first.size() * 2, max_file_size));
                 avail_bytes = retval.first.size() - retval_size;
 
                 strm.next_out = retval.first.data() + retval_size;

diff --git a/src/io/ZlibCompression.hh b/src/io/ZlibCompression.hh
@@ -47,7 +47,8 @@ public:
                                  JpegAllocator<uint8_t> >,
                      JpegError> Decompress(const uint8_t *buffer,
                                              size_t size,
-                                             const JpegAllocator<uint8_t> &alloc);
+                                           const JpegAllocator<uint8_t> &alloc,
+                         size_t max_size);
     ZlibDecoderDecompressionReader(DecoderReader *r, bool concatenated, const JpegAllocator<uint8_t> &alloc);
     virtual std::pair<uint32, JpegError> Read(uint8*data, unsigned int size);
     virtual ~ZlibDecoderDecompressionReader();

diff --git a/src/lepton/jpgcoder.cc b/src/lepton/jpgcoder.cc
@@ -78,6 +78,7 @@ volatile int volatile1024 = 1024;
 #include "socket_serve.hh"
 #include "validation.hh"
 #include "../io/ZlibCompression.hh"
+#include "../io/BrotliCompression.hh"
 #include "../io/MemReadWriter.hh"
 #include "../io/BufferedIO.hh"
 #include "../io/Zlib0.hh"
@@ -1080,8 +1081,9 @@ int initialize_options( int argc, const char*const * argv )
             g_skip_validation = false;
         } else if ( strcmp((*argv), "-roundtrip") == 0 ) {
             g_skip_validation = false;
-        }
-        else if ( strncmp((*argv), "-maxchildren=", strlen("-maxchildren=") ) == 0 ) {
+        } else if ( strcmp((*argv), "-brotliheader") == 0 ) {
+            ujgversion = 2; // use brotli to compress the header and trailer rather than zlib
+        } else if ( strncmp((*argv), "-maxchildren=", strlen("-maxchildren=") ) == 0 ) {
             g_socketserve_info.max_children = strtol((*argv) + strlen("-maxchildren="), NULL, 10);
         }
         else if ( strncmp((*argv), "-listenbacklog=", strlen("-listenbacklog=") ) == 0 ) {
@@ -1952,7 +1954,7 @@ unsigned char read_fixed_ujpg_header() {
         custom_exit(ExitCode::SHORT_READ);
     }
     // check version number
-    if (header[0] != 1 && header[0] != 2 && header[0] != ujgversion) {
+    if (header[0] != 1 && header[0] != 2 && header[0] != 3 && header[0] != ujgversion) {
         // let us roll out a new version gently
         fprintf( stderr, "incompatible file, use %s v%i.%i",
             appname, header[ 0 ] / 10, header[ 0 ] % 10 );
@@ -3791,11 +3793,16 @@ bool write_ujpg(std::vector<ThreadHandoff> row_thread_handoffs,
         //custom_exit(ExitCode::HEADER_TOO_LARGE);
     }
     std::vector<uint8_t, Sirikata::JpegAllocator<uint8_t> > compressed_header;
-    compressed_header =
+    if (ujgversion == 1) {
+        compressed_header =
             Sirikata::ZlibDecoderCompressionWriter::Compress(mrw.buffer().data(),
                                                              mrw.buffer().size(),
                                                              Sirikata::JpegAllocator<uint8_t>());
-
+    } else {
+        compressed_header = Sirikata::BrotliCodec::Compress(mrw.buffer().data(),
+                                                            mrw.buffer().size(),
+                                                            Sirikata::JpegAllocator<uint8_t>());
+    }
     write_byte_bill(Billing::HEADER, false, 2 + hdrs + prefix_grbs + grbs);
     static_assert(MAX_NUM_THREADS <= 255, "We only have a single byte for num threads");
     always_assert(NUM_THREADS <= 255);
@@ -3903,18 +3910,35 @@ bool read_ujpg( void )
                                                  &mem_realloc_nop,
                                                  &MemMgrAllocatorMsize);
 #endif
-        std::pair<std::vector<uint8_t,
-                              Sirikata::JpegAllocator<uint8_t> >,
-                  JpegError> uncompressed_header_buffer(
-                ZlibDecoderDecompressionReader::Decompress(compressed_header_buffer.data(),
-                                                         compressed_header_buffer.size(),
-                                                           no_free_allocator));
-        if (uncompressed_header_buffer.second) {
-            always_assert(false && "Data not properly zlib coded");
-            return false;
+        if (ujgversion == 1) {
+            std::pair<std::vector<uint8_t,
+                                  Sirikata::JpegAllocator<uint8_t> >,
+                      JpegError> uncompressed_header_buffer(
+                          ZlibDecoderDecompressionReader::Decompress(compressed_header_buffer.data(),
+                                                                     compressed_header_buffer.size(),
+                                                                     no_free_allocator,
+                                                                     max_file_size));
+            if (uncompressed_header_buffer.second) {
+                always_assert(false && "Data not properly zlib coded");
+                return false;
+            }
+            zlib_hdrs = compressed_header_buffer.size();
+            header_reader.SwapIn(uncompressed_header_buffer.first, 0);
+        } else {
+            std::pair<std::vector<uint8_t,
+                                  Sirikata::JpegAllocator<uint8_t> >,
+                      JpegError> uncompressed_header_buffer(
+                          Sirikata::BrotliCodec::Decompress(compressed_header_buffer.data(),
+                                                            compressed_header_buffer.size(),
+                                                            no_free_allocator,
+                              max_file_size));
+            if (uncompressed_header_buffer.second) {
+                always_assert(false && "Data not properly zlib coded");
+                return false;
+            }
+            zlib_hdrs = compressed_header_buffer.size();
+            header_reader.SwapIn(uncompressed_header_buffer.first, 0);            
         }
-        zlib_hdrs = compressed_header_buffer.size();
-        header_reader.SwapIn(uncompressed_header_buffer.first, 0);
     }
     grbs = sizeof(EOI);
     grbgdata = EOI; // if we don't have any garbage, assume FFD9 EOI