ChunkedOutputFilter.java

  1. /*
  2.  *  Licensed to the Apache Software Foundation (ASF) under one or more
  3.  *  contributor license agreements.  See the NOTICE file distributed with
  4.  *  this work for additional information regarding copyright ownership.
  5.  *  The ASF licenses this file to You under the Apache License, Version 2.0
  6.  *  (the "License"); you may not use this file except in compliance with
  7.  *  the License.  You may obtain a copy of the License at
  8.  *
  9.  *      http://www.apache.org/licenses/LICENSE-2.0
  10.  *
  11.  *  Unless required by applicable law or agreed to in writing, software
  12.  *  distributed under the License is distributed on an "AS IS" BASIS,
  13.  *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14.  *  See the License for the specific language governing permissions and
  15.  *  limitations under the License.
  16.  */
  17. package org.apache.coyote.http11.filters;

  18. import java.io.ByteArrayOutputStream;
  19. import java.io.IOException;
  20. import java.io.OutputStreamWriter;
  21. import java.nio.ByteBuffer;
  22. import java.nio.charset.StandardCharsets;
  23. import java.util.HashSet;
  24. import java.util.Locale;
  25. import java.util.Map;
  26. import java.util.Set;
  27. import java.util.function.Supplier;

  28. import org.apache.coyote.Response;
  29. import org.apache.coyote.http11.HttpOutputBuffer;
  30. import org.apache.coyote.http11.OutputFilter;
  31. import org.apache.tomcat.util.buf.HexUtils;

  32. /**
  33.  * Chunked output filter.
  34.  *
  35.  * @author Remy Maucherat
  36.  */
  37. public class ChunkedOutputFilter implements OutputFilter {

  38.     private static final byte[] LAST_CHUNK_BYTES = { (byte) '0', (byte) '\r', (byte) '\n' };
  39.     private static final byte[] CRLF_BYTES = { (byte) '\r', (byte) '\n' };
  40.     private static final byte[] END_CHUNK_BYTES = { (byte) '0', (byte) '\r', (byte) '\n', (byte) '\r', (byte) '\n' };

  41.     private static final Set<String> disallowedTrailerFieldNames = new HashSet<>();

  42.     static {
  43.         // Always add these in lower case
  44.         disallowedTrailerFieldNames.add("age");
  45.         disallowedTrailerFieldNames.add("cache-control");
  46.         disallowedTrailerFieldNames.add("content-length");
  47.         disallowedTrailerFieldNames.add("content-encoding");
  48.         disallowedTrailerFieldNames.add("content-range");
  49.         disallowedTrailerFieldNames.add("content-type");
  50.         disallowedTrailerFieldNames.add("date");
  51.         disallowedTrailerFieldNames.add("expires");
  52.         disallowedTrailerFieldNames.add("location");
  53.         disallowedTrailerFieldNames.add("retry-after");
  54.         disallowedTrailerFieldNames.add("trailer");
  55.         disallowedTrailerFieldNames.add("transfer-encoding");
  56.         disallowedTrailerFieldNames.add("vary");
  57.         disallowedTrailerFieldNames.add("warning");
  58.     }

  59.     /**
  60.      * Next buffer in the pipeline.
  61.      */
  62.     protected HttpOutputBuffer buffer;


  63.     /**
  64.      * Chunk header.
  65.      */
  66.     protected final ByteBuffer chunkHeader = ByteBuffer.allocate(10);


  67.     protected final ByteBuffer lastChunk = ByteBuffer.wrap(LAST_CHUNK_BYTES);
  68.     protected final ByteBuffer crlfChunk = ByteBuffer.wrap(CRLF_BYTES);
  69.     /**
  70.      * End chunk.
  71.      */
  72.     protected final ByteBuffer endChunk = ByteBuffer.wrap(END_CHUNK_BYTES);


  73.     private Response response;


  74.     public ChunkedOutputFilter() {
  75.         chunkHeader.put(8, (byte) '\r');
  76.         chunkHeader.put(9, (byte) '\n');
  77.     }


  78.     // --------------------------------------------------- OutputBuffer Methods

  79.     @Override
  80.     public int doWrite(ByteBuffer chunk) throws IOException {

  81.         int result = chunk.remaining();

  82.         if (result <= 0) {
  83.             return 0;
  84.         }

  85.         int pos = calculateChunkHeader(result);

  86.         chunkHeader.position(pos).limit(10);
  87.         buffer.doWrite(chunkHeader);

  88.         buffer.doWrite(chunk);

  89.         chunkHeader.position(8).limit(10);
  90.         buffer.doWrite(chunkHeader);

  91.         return result;
  92.     }


  93.     private int calculateChunkHeader(int len) {
  94.         // Calculate chunk header
  95.         int pos = 8;
  96.         int current = len;
  97.         while (current > 0) {
  98.             int digit = current % 16;
  99.             current = current / 16;
  100.             chunkHeader.put(--pos, HexUtils.getHex(digit));
  101.         }
  102.         return pos;
  103.     }


  104.     @Override
  105.     public long getBytesWritten() {
  106.         return buffer.getBytesWritten();
  107.     }


  108.     // --------------------------------------------------- OutputFilter Methods

  109.     @Override
  110.     public void setResponse(Response response) {
  111.         this.response = response;
  112.     }


  113.     @Override
  114.     public void setBuffer(HttpOutputBuffer buffer) {
  115.         this.buffer = buffer;
  116.     }


  117.     @Override
  118.     public void flush() throws IOException {
  119.         // No data buffered in this filter. Flush next buffer.
  120.         buffer.flush();
  121.     }


  122.     @Override
  123.     public void end() throws IOException {

  124.         Supplier<Map<String,String>> trailerFieldsSupplier = response.getTrailerFields();
  125.         Map<String,String> trailerFields = null;

  126.         if (trailerFieldsSupplier != null) {
  127.             trailerFields = trailerFieldsSupplier.get();
  128.         }

  129.         if (trailerFields == null) {
  130.             // Write end chunk
  131.             buffer.doWrite(endChunk);
  132.             endChunk.position(0).limit(endChunk.capacity());
  133.         } else {
  134.             buffer.doWrite(lastChunk);
  135.             lastChunk.position(0).limit(lastChunk.capacity());

  136.             ByteArrayOutputStream baos = new ByteArrayOutputStream(1024);

  137.             try (OutputStreamWriter osw = new OutputStreamWriter(baos, StandardCharsets.ISO_8859_1)) {
  138.                 for (Map.Entry<String,String> trailerField : trailerFields.entrySet()) {
  139.                     // Ignore disallowed headers
  140.                     if (disallowedTrailerFieldNames.contains(trailerField.getKey().toLowerCase(Locale.ENGLISH))) {
  141.                         continue;
  142.                     }
  143.                     osw.write(trailerField.getKey());
  144.                     osw.write(':');
  145.                     osw.write(' ');
  146.                     osw.write(trailerField.getValue());
  147.                     osw.write("\r\n");
  148.                 }
  149.             }

  150.             buffer.doWrite(ByteBuffer.wrap(baos.toByteArray()));

  151.             buffer.doWrite(crlfChunk);
  152.             crlfChunk.position(0).limit(crlfChunk.capacity());
  153.         }
  154.         buffer.end();
  155.     }


  156.     @Override
  157.     public void recycle() {
  158.         response = null;
  159.     }
  160. }