tools/telemetry/third_party/webpagereplay/httpzlib.py

   1 #!/usr/bin/env python
   2 # Copyright 2011 Google Inc. All Rights Reserved.
   3 #
   4 # Licensed under the Apache License, Version 2.0 (the "License");
   5 # you may not use this file except in compliance with the License.
   6 # You may obtain a copy of the License at
   7 #
   8 #      http://www.apache.org/licenses/LICENSE-2.0
   9 #
  10 # Unless required by applicable law or agreed to in writing, software
  11 # distributed under the License is distributed on an "AS IS" BASIS,
  12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 # See the License for the specific language governing permissions and
  14 # limitations under the License.
  15
  16 """Apply gzip/deflate to separate chunks of data."""
  17
  18 import struct
  19 import zlib
  20
  21 GZIP_HEADER = (
  22     '\037\213'             # magic header
  23     '\010'                 # compression method
  24     '\000'                 # flags (none)
  25     '\000\000\000\000'     # packed time (use zero)
  26     '\002'
  27     '\377')
  28
  29
  30 def compress_chunks(uncompressed_chunks, use_gzip):
  31   """Compress a list of data with gzip or deflate.
  32
  33   The returned chunks may be used with HTTP chunked encoding.
  34
  35   Args:
  36     uncompressed_chunks: a list of strings
  37        (e.g. ["this is the first chunk", "and the second"])
  38     use_gzip: if True, compress with gzip. Otherwise, use deflate.
  39
  40   Returns:
  41     [compressed_chunk_1, compressed_chunk_2, ...]
  42   """
  43   if use_gzip:
  44     size = 0
  45     crc = zlib.crc32("") & 0xffffffffL
  46     compressor = zlib.compressobj(
  47         6, zlib.DEFLATED, -zlib.MAX_WBITS, zlib.DEF_MEM_LEVEL, 0)
  48   else:
  49     compressor = zlib.compressobj()
  50   compressed_chunks = []
  51   last_index = len(uncompressed_chunks) - 1
  52   for index, data in enumerate(uncompressed_chunks):
  53     chunk = ''
  54     if use_gzip:
  55       size += len(data)
  56       crc = zlib.crc32(data, crc) & 0xffffffffL
  57       if index == 0:
  58         chunk += GZIP_HEADER
  59     chunk += compressor.compress(data)
  60     if index < last_index:
  61       chunk += compressor.flush(zlib.Z_SYNC_FLUSH)
  62     else:
  63       chunk += (compressor.flush(zlib.Z_FULL_FLUSH) +
  64                 compressor.flush())
  65       if use_gzip:
  66         chunk += (struct.pack("<L", long(crc)) +
  67                   struct.pack("<L", long(size)))
  68     compressed_chunks.append(chunk)
  69   return compressed_chunks
  70
  71
  72 def uncompress_chunks(compressed_chunks, use_gzip):
  73   """Uncompress a list of data compressed with gzip or deflate.
  74
  75   Args:
  76     compressed_chunks: a list of compressed data
  77     use_gzip: if True, uncompress with gzip. Otherwise, use deflate.
  78
  79   Returns:
  80     [uncompressed_chunk_1, uncompressed_chunk_2, ...]
  81   """
  82   if use_gzip:
  83     decompress = zlib.decompressobj(16 + zlib.MAX_WBITS).decompress
  84   else:
  85     decompress = zlib.decompressobj(-zlib.MAX_WBITS).decompress
  86   return [decompress(c) for c in compressed_chunks]