package org.apache.catalina.tomcat.util.http;

import junit.framework.Test;

import junit.framework.TestCase;

import junit.framework.TestSuite;

import junit.textui.TestRunner;

import org.apache.coyote.Response;

import org.apache.coyote.http11.InternalOutputBuffer;

import org.apache.coyote.http11.filters.GzipOutputFilter;

import org.apache.tomcat.util.buf.ByteChunk;

import java.io.ByteArrayOutputStream;

import java.util.zip.GZIPOutputStream;

/**

 * User: Jiong Wang (jiwang@linkedin.com)

 * Date: Feb 13, 2010

 * Time: 12:04:38 AM

 *

 * Test case to demonstrate the interaction between gzip and flushing in the output filter.

 */

public class TestGzipOutputFilter extends TestCase

{

  public static void main( String args[] ) {

     TestRunner.run(suite());

  }

  public static Test suite() {

     TestSuite suite = new TestSuite();

     suite.addTest(new TestSuite(TestGzipOutputFilter.class));

     return suite;

  }

  public TestGzipOutputFilter()

  {

  }

  public TestGzipOutputFilter(String s)

  {

    super(s);

  }

  /**

   * Test the interaction betwen gzip and flushing.

   *

   * The idea is to:

   * 1. create a internal output buffer, response, and attach an active gzipoutputfilter to the output buffer

   * 2. set the output stream of the internal buffer to be a ByteArrayOutputStream so we can inspect the output bytes

   * 3. write a chunk out using the gzipoutputfilter and invoke a flush on the InternalOutputBuffer

   * 4. read from the ByteArrayOutputStream to find out what's being written out (flushed)

   * 5. find out what's expected by wrting to GZIPOutputStream and close it (to force flushing)

   * 6. Compare the size of the two arrays, they should be close (instead of one being much shorter than the other one)

   * @throws Exception

   */

  public void testFlushingWithGzip() throws Exception

  {

    // set up response, InternalOutputBuffer, and ByteArrayOutputStream

    Response res = new Response();

    InternalOutputBuffer iob = new InternalOutputBuffer(res);

    ByteArrayOutputStream bos = new ByteArrayOutputStream();

    iob.setOutputStream(bos);

    res.setOutputBuffer(iob);

    // set up GzipOutputFilter to attach to the InternalOutputBuffer

    GzipOutputFilter gf = new GzipOutputFilter();

    iob.addFilter(gf);

    iob.addActiveFilter(gf);

    // write a chunk out

    ByteChunk chunk = new ByteChunk(1024);

    byte[] d = "Hello there tomcat developers, there is a bug in JDK".getBytes();

    chunk.append(d, 0, d.length);

    iob.doWrite(chunk, res);

    // flush the InternalOutputBuffer

    iob.flush();

    // read from the ByteArrayOutputStream to find out what's being written out (flushed)

    byte[] dataFound = bos.toByteArray();

    // find out what's expected by wrting to GZIPOutputStream and close it (to force flushing)

    ByteArrayOutputStream gbos = new ByteArrayOutputStream(1024);

    GZIPOutputStream gos = new GZIPOutputStream(gbos);

    gos.write(d);

    gos.close();

    // read the expected data

    byte[] dataExpected = gbos.toByteArray();

    // most of the data should have been flushed out

    assertTrue(dataFound.length >= (dataExpected.length - 20));

//    System.out.println("dataFound = " + Arrays.toString(dataFound));

//    System.out.println("dataExpected = " + Arrays.toString(dataExpected));

  }

}

import org.apache.coyote.http11.filters.GzipOutputFilter;

  /**

   * Logger.

   */

  protected static org.apache.juli.logging.Log log

      = org.apache.juli.logging.LogFactory.getLog(InternalOutputBuffer.class);

      // go through the filters and if there is gzip filter

      // invoke it to flush

      for (int i = 0; i <= lastActiveFilter; i++) {

        if (activeFilters[i] instanceof GzipOutputFilter)

        {

          if (log.isDebugEnabled())

          {

            log.debug("Flushing the gzip filter at position " + i + " of the filter chain...");

          }

          ((GzipOutputFilter)activeFilters[i]).flush();

          break; 

        }

      }

import org.apache.coyote.OutputBuffer;

import org.apache.coyote.Response;

import org.apache.coyote.http11.OutputFilter;

import org.apache.tomcat.util.buf.ByteChunk;

  /**

   * Logger.

   */

  protected static org.apache.juli.logging.Log log

      = org.apache.juli.logging.LogFactory.getLog(GzipOutputFilter.class);

  /**

   * Added to allow flushing to happen for the gzip'ed outputstream

   */

  public void flush()

  {

    if (compressionStream != null)

    {

      try

      {

        if (log.isDebugEnabled())

        {

          log.debug("Flushing the compression stream!");

        }

        compressionStream.flush();

      }

      catch (IOException e)

      {

        if (log.isDebugEnabled())

        {

          log.debug("Ignored exception while flushing gzip filter", e);

        }

      }

    }

  }

package org.apache.coyote.http11.filters;

import java.util.zip.GZIPOutputStream;

import java.util.zip.Deflater;

import java.io.OutputStream;

import java.io.IOException;

/**

 * User: Jiong Wang (jiwang@linkedin.com)

 * Date: Jan 6, 2010

 * Time: 10:30:48 AM

 *

 */

public class FlushableGZIPOutputStream extends GZIPOutputStream

 {

   public FlushableGZIPOutputStream(OutputStream os) throws IOException

   {

     super(os);

   }

   private static final byte[] EMPTYBYTEARRAY = new byte[0];

   private boolean hasData = false;

   /**

    * Here we make sure we have received data, so that the header has

    * been for sure written to the output stream already.

    */

   @Override

   public synchronized void write(byte[] bytes, int i, int i1) throws IOException

   {

     super.write(bytes, i, i1);

     hasData = true;

   }

   @Override

   public synchronized void write(int i) throws IOException

   {

     super.write(i);

     hasData = true;

   }

   @Override

   public synchronized void write(byte[] bytes) throws IOException

   {

     super.write(bytes);

     hasData = true;

   }

   @Override

   public synchronized void flush() throws IOException

   {

     if (!hasData)

     {

       return; // do not allow the gzip header to be flushed on its own

     }

     // trick the deflater to flush

     /**

      * Now this is tricky: We force the Deflater to flush its data by

      * switching compression level. As yet, a perplexingly simple workaround

      * for http://developer.java.sun.com/developer/bugParade/bugs/4255743.html

      */

     if (!def.finished()) {

       def.setInput(EMPTYBYTEARRAY, 0, 0);

       def.setLevel(Deflater.NO_COMPRESSION);

       deflate();

       def.setLevel(Deflater.DEFAULT_COMPRESSION);

       deflate();

       out.flush();

     }

     hasData = false; // no more data to flush 

   }

   /*

    * Keep on calling deflate until it runs dry. The default implementation only does it once and can therefore

    * hold onto data when they need to be flushed out.

    */

  protected void deflate() throws IOException

  {

    int len;

    do

    {

      len = def.deflate(buf, 0, buf.length);

      if (len > 0)

      {

        out.write(buf, 0, len);

      }

    }

    while (len != 0);

  }

 }