/* This software is published under the terms of the OpenSymphony Software
    * License version 1.1, of which a copy has been included with this
    * distribution in the LICENSE.txt file. */
   package com.atlassian.gzipfilter;
   
   import com.atlassian.gzipfilter.selector.GzipCompatibilitySelector;
   import com.atlassian.gzipfilter.util.HttpContentType;
   import org.slf4j.Logger;
   import org.slf4j.LoggerFactory;
  
  import java.io.IOException;
  import java.io.PrintWriter;
  import javax.servlet.ServletOutputStream;
  import javax.servlet.ServletResponse;
  import javax.servlet.http.HttpServletResponse;
  import javax.servlet.http.HttpServletResponseWrapper;
  
  /**
   * Implementation of HttpServletResponseWrapper that captures page data instead of
   * sending to the writer.
   * <p/>
   * <p>Should be used in filter-chains or when forwarding/including pages
   * using a RequestDispatcher.</p>
   *
   * @author <a href="mailto:joe@truemesh.com">Joe Walnes</a>
   * @author <a href="mailto:scott@atlassian.com">Scott Farquhar</a>
   */
  public class SelectingResponseWrapper extends HttpServletResponseWrapper
  {
      private static final Logger log = LoggerFactory.getLogger(SelectingResponseWrapper.class);
  
      private final RoutablePrintWriter routablePrintWriter;
      private final RoutableServletOutputStream routableServletOutputStream;
      private final GzipCompatibilitySelector compatibilitySelector;
  
      private final GzipResponseWrapper wrappedResponse;
  
      private boolean gzippablePage = false;
      private boolean headersCommitted = false;
  
      public SelectingResponseWrapper(final HttpServletResponse unWrappedResponse, GzipCompatibilitySelector compatibilitySelector, String defaultEncoding)
      {
          super(unWrappedResponse);
          this.wrappedResponse = new GzipResponseWrapper(unWrappedResponse, defaultEncoding);
          this.compatibilitySelector = compatibilitySelector;
  
          Runnable gzipHeadersCommitter = new Runnable()
          {
              /**
               * This callback is our last chance to commit headers before response
               * is committed by using output stream
               */
              public void run()
              {
                  commitGzipHeaders();
              }
          };
          routablePrintWriter = new RoutablePrintWriter(
                  new RoutablePrintWriterDestinationFactory(unWrappedResponse),
                  gzipHeadersCommitter);
          routableServletOutputStream = new RoutableServletOutputStream(
                  new RoutableServletOutputStreamDestinationFactory(unWrappedResponse),
                  gzipHeadersCommitter);
      }
  
      public void setContentType(String type)
      {
          super.setContentType(type);
  
          if (type != null)
          {
              final HttpContentType httpContentType = new HttpContentType(type);
              if (compatibilitySelector.shouldGzip(httpContentType.getType()))
              {
                  activateGzip(httpContentType.getEncoding());
              }
              else
              {
                  deactivateGzip();
              }
          }
      }
  
      /**
       * If a redirect is set, cancel the gzip header.  The JWebunit test http client fails when it tries to
       * unzip an empty body from a redirect and other clients might also fail.  Since redirects are unlikely to have
       * much of a body anyway, it makes sense not to gzip them.
       */
      public void sendRedirect(String location) throws IOException
      {
          if (!wrappedResponse.isCommitted() && gzippablePage)
              deactivateGzip();
          super.sendRedirect(location);
      }
  
      /**
       * <p>In addition to its normal behaviour this will also deactivate gzip
       * encoding for those responses that MUST NOT have a body.</p>
       * <p>This is done according to <a href="http://www.w3.org/Protocols/rfc2616/rfc2616-sec10.html">RFC 2616, section 10</a></p>
      *
      * @see HttpServletResponseWrapper#setStatus(int, String)
      * @see #setStatus(int)
      */
     public void setStatus(int statusCode, String sm)
     {
         super.setStatus(statusCode, sm);
         if (!shouldGzip(statusCode))
         {
             deactivateGzip();
         }
     }
 
     /**
      * <p>In addition to its normal behaviour this will also deactivate gzip
      * encoding for those responses that MUST NOT have a body.</p>
      * <p>This is done according to <a href="http://www.w3.org/Protocols/rfc2616/rfc2616-sec10.html">RFC 2616, section 10</a></p>
      *
      * @see #setStatus(int, String)
      * @see HttpServletResponseWrapper#setStatus(int)
      */
     public void setStatus(int statusCode)
     {
         super.setStatus(statusCode);
         if (!shouldGzip(statusCode))
         {
             deactivateGzip();
         }
     }
 
     public void sendError(int sc, String msg) throws IOException
     {
         if (gzippablePage)
         {
             deactivateGzip();
         }
 
         super.sendError(sc, msg);
     }
 
     public void sendError(int sc) throws IOException
     {
         if (gzippablePage)
         {
             deactivateGzip();
         }
 
         super.sendError(sc);
     }
 
     /**
      * Defines whether or not response with the given code should be gzipped.
      * @param statusCode the response status code.
      * @return <code>true</code> if the status code is good to be gzipped, <code>false</code> otherwise.
      */
     private boolean shouldGzip(int statusCode)
     {
         return statusCode != SC_NO_CONTENT && statusCode != SC_NOT_MODIFIED;
     }
 
     /**
      * Sets header for gzipped responses if it is possible
      */
     private void commitGzipHeaders()
     {
         if (headersCommitted) {
             return;
         }
         if (!gzippablePage) {
             log.trace("Not a gzippable page");
             return;
         }
         if (wrappedResponse.isCommitted())
         {
             log.debug("Response is committed, can't set gzip headers");
             return;
         }
 
         log.debug("Setting gzip headers");
         wrappedResponse.setHeader("Content-Encoding", "gzip");
         // Technically, we should add a 'vary' header here, as we are sending a different response
         // based on both the user-agent and the accept encoding.  This is so that proxies can
         // cache multiple versions of a page, and send the right one back to clients.
         // We should add 'Vary: user-agent, accept-encoding', but we can't.
         //
         // References:
         //    http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html#sec13.6
         //    http://www.w3.org/Protocols/rfc2616/rfc2616-sec14.html#sec14.44
         //    http://www.port80software.com/200ok/archive/2005/01/21/272.aspx
         //
         //
         // However, there is a *huge* bug in all versions of Internet Explorer, such that if a vary header
         // is set, it will refuse to cache the page (even on disk!).  The only Vary header that
         // internet explorer accepts is 'Vary: user-agent'.
         //
         // Reference:
         //    http://lists.over.net/pipermail/mod_gzip/2002-December/006826.html
         //    http://mail-archives.apache.org/mod_mbox/httpd-dev/200511.mbox/%3C198.4b088455.30a1c19b@aol.com%3E
         //    http://httpd.apache.org/docs/2.0/misc/known_client_problems.html
         //    http://support.microsoft.com/kb/824847/en-us?spid=8722&sid=global
 
 
         wrappedResponse.setHeader("Vary", "User-Agent");
         //wrappedResponse.addHeader("Vary", "Accept-Encoding"); - Can't do due to internet explorer bugs
 
         // Another way to avoid this broken-ness is to add a cache-control header.
         // set elsewhere, so we dno'
         // So - instead, to avoid this brokenness, let's just ensure that this content isn't
         // cached by any shared caches, and add a 'Cache-control: private' header
         // Reference:
         //    http://www.w3.org/Protocols/rfc2616/rfc2616-sec14.html#sec14.9.1
         //
         // However - this may conflict with other cache-control headers we may have, so let's not do it
         // See Scott and Anton for long discussion.
 
         //wrappedResponse.addHeader("Cache-control", "private");
         headersCommitted = true;
     }
 
     private void activateGzip(String encoding)
     {
         if (gzippablePage)
         {
             return; // already activated
         }
 
         if (wrappedResponse.isCommitted()) {
             log.debug("Response is committed, gzip can not be activated");
             return;
         }
 
         if (headersCommitted) {
             log.debug("Headers are committed, gzip can not be activated");
             return;
         }
 
         //JRA-35223 - Cannot activate gzip for a response when someone already have set Content-Length header
         //because this will produce a response with incorrect Content-Length (Content-Length set to size of
         //uncompressed data, but compressed content which is smaller)
         if(wrappedResponse.containsHeader("Content-Length"))
         {
             log.debug("Gzip compression can not be activated when the Content-Length header has already been set "
                     + "on the response, and therefore uncompressed content will be sent instead");
             return;
         }
 
         if (encoding != null) {
             wrappedResponse.setEncoding(encoding);
         }
 
         routablePrintWriter.updateDestination(new RoutablePrintWriterDestinationFactory(wrappedResponse));
         routableServletOutputStream.updateDestination(new RoutableServletOutputStreamDestinationFactory(wrappedResponse));
         gzippablePage = true;
         log.debug("gzip activated");
     }
 
     private void deactivateGzip()
     {
         gzippablePage = false;
 
         routablePrintWriter.updateDestination(new RoutablePrintWriterDestinationFactory(getResponse()));
         routableServletOutputStream.updateDestination(new RoutableServletOutputStreamDestinationFactory(getResponse()));
         log.debug("gzip deactivated");
     }
 
     /**
      * Prevent content-length being set if page is parseable.
      */
     @Override
     public void setContentLength(int contentLength)
     {
         if (!gzippablePage) { super.setContentLength(contentLength); }
     }
 
     /**
      * Prevent buffer from being flushed if this is a page being parsed.
      */
     @Override
     public void flushBuffer() throws IOException
     {
         if (!gzippablePage)
         {
             log.debug("Flushing buffer");
             super.flushBuffer();
         }
     }
 
     /**
      * Prevent content-length being set if page is parseable.
      */
     @Override
     public void setHeader(String name, String value)
     {
         if (name.toLowerCase().equals("content-type"))
         { // ensure ContentType is always set through setContentType()
             setContentType(value);
         }
         else if (!gzippablePage || !name.toLowerCase().equals("content-length"))
         {
             super.setHeader(name, value);
         }
     }
 
     /**
      * Prevent content-length being set if page is parseable.
      */
     @Override
     public void addHeader(String name, String value)
     {
         if (name.toLowerCase().equals("content-type"))
         { // ensure ContentType is always set through setContentType()
             setContentType(value);
         }
         else if (!gzippablePage || !name.toLowerCase().equals("content-length"))
         {
             super.addHeader(name, value);
         }
     }
 
     @Override
     public ServletOutputStream getOutputStream()
     {
         return routableServletOutputStream;
     }
 
     @Override
     public PrintWriter getWriter()
     {
         return routablePrintWriter;
     }
 
     public void finishResponse()
     {
         if (gzippablePage)
         {
             commitGzipHeaders();
             wrappedResponse.finishResponse();
         }
     }
 
     private static class RoutablePrintWriterDestinationFactory implements RoutablePrintWriter.DestinationFactory
     {
         private final ServletResponse servletResponse;
 
         public RoutablePrintWriterDestinationFactory(final ServletResponse servletResponse)
         {
             this.servletResponse = servletResponse;
         }
 
         public PrintWriter activateDestination() throws IOException
         {
             return servletResponse.getWriter();
         }
     }
 
     private static class RoutableServletOutputStreamDestinationFactory implements RoutableServletOutputStream.DestinationFactory
     {
         private final ServletResponse servletResponse;
 
         public RoutableServletOutputStreamDestinationFactory (final ServletResponse servletResponse)
         {
             this.servletResponse = servletResponse;
         }
 
         public ServletOutputStream create() throws IOException
         {
             return servletResponse.getOutputStream();
         }
     }
 }