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
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
29 import org.apache.coyote.Response;
30 import org.apache.coyote.http11.HttpOutputBuffer;
31 import org.apache.coyote.http11.OutputFilter;
32 import org.apache.tomcat.util.buf.HexUtils;
33 import org.apache.tomcat.util.http.fileupload.ByteArrayOutputStream;
34
35 /**
36  * Chunked output filter.
37  *
38  * @author Remy Maucherat
39  */
40 public class ChunkedOutputFilter implements OutputFilter {
41
42     private static final byte[] LAST_CHUNK_BYTES = {(byte) '0', (byte) '\r', (byte) '\n'};
43     private static final byte[] CRLF_BYTES = {(byte) '\r', (byte) '\n'};
44     private static final byte[] END_CHUNK_BYTES =
45         {(byte) '0', (byte) '\r', (byte) '\n', (byte) '\r', (byte) '\n'};
46
47     private static final Set<String> disallowedTrailerFieldNames = new HashSet<>();
48
49     static {
50         // Always add these in lower case
51         disallowedTrailerFieldNames.add("age");
52         disallowedTrailerFieldNames.add("cache-control");
53         disallowedTrailerFieldNames.add("content-length");
54         disallowedTrailerFieldNames.add("content-encoding");
55         disallowedTrailerFieldNames.add("content-range");
56         disallowedTrailerFieldNames.add("content-type");
57         disallowedTrailerFieldNames.add("date");
58         disallowedTrailerFieldNames.add("expires");
59         disallowedTrailerFieldNames.add("location");
60         disallowedTrailerFieldNames.add("retry-after");
61         disallowedTrailerFieldNames.add("trailer");
62         disallowedTrailerFieldNames.add("transfer-encoding");
63         disallowedTrailerFieldNames.add("vary");
64         disallowedTrailerFieldNames.add("warning");
65     }
66
67     /**
68      * Next buffer in the pipeline.
69      */
70     protected HttpOutputBuffer buffer;
71
72
73     /**
74      * Chunk header.
75      */
76     protected final ByteBuffer chunkHeader = ByteBuffer.allocate(10);
77
78
79     protected final ByteBuffer lastChunk = ByteBuffer.wrap(LAST_CHUNK_BYTES);
80     protected final ByteBuffer crlfChunk = ByteBuffer.wrap(CRLF_BYTES);
81     /**
82      * End chunk.
83      */
84     protected final ByteBuffer endChunk = ByteBuffer.wrap(END_CHUNK_BYTES);
85
86
87     private Response response;
88
89
90     public ChunkedOutputFilter() {
91         chunkHeader.put(8, (byte) '\r');
92         chunkHeader.put(9, (byte) '\n');
93     }
94
95
96     // --------------------------------------------------- OutputBuffer Methods
97
98     @Override
99     public int doWrite(ByteBuffer chunk) throws IOException {
100
101         int result = chunk.remaining();
102
103         if (result <= 0) {
104             return 0;
105         }
106
107         int pos = calculateChunkHeader(result);
108
109         chunkHeader.position(pos).limit(10);
110         buffer.doWrite(chunkHeader);
111
112         buffer.doWrite(chunk);
113
114         chunkHeader.position(8).limit(10);
115         buffer.doWrite(chunkHeader);
116
117         return result;
118     }
119
120
121     private int calculateChunkHeader(int len) {
122         // Calculate chunk header
123         int pos = 8;
124         int current = len;
125         while (current > 0) {
126             int digit = current % 16;
127             current = current / 16;
128             chunkHeader.put(--pos, HexUtils.getHex(digit));
129         }
130         return pos;
131     }
132
133
134     @Override
135     public long getBytesWritten() {
136         return buffer.getBytesWritten();
137     }
138
139
140     // --------------------------------------------------- OutputFilter Methods
141
142     @Override
143     public void setResponse(Response response) {
144         this.response = response;
145     }
146
147
148     @Override
149     public void setBuffer(HttpOutputBuffer buffer) {
150         this.buffer = buffer;
151     }
152
153
154     @Override
155     public void flush() throws IOException {
156         // No data buffered in this filter. Flush next buffer.
157         buffer.flush();
158     }
159
160
161     @Override
162     public void end() throws IOException {
163
164         Supplier<Map<String,String>> trailerFieldsSupplier = response.getTrailerFields();
165         Map<String,String> trailerFields = null;
166
167         if (trailerFieldsSupplier != null) {
168             trailerFields = trailerFieldsSupplier.get();
169         }
170
171         if (trailerFields == null) {
172             // Write end chunk
173             buffer.doWrite(endChunk);
174             endChunk.position(0).limit(endChunk.capacity());
175         } else {
176             buffer.doWrite(lastChunk);
177             lastChunk.position(0).limit(lastChunk.capacity());
178
179            ByteArrayOutputStream baos = new ByteArrayOutputStream(1024);
180            OutputStreamWriter osw = new OutputStreamWriter(baos, StandardCharsets.ISO_8859_1);
181             for (Map.Entry<String,String> trailerField : trailerFields.entrySet()) {
182                 // Ignore disallowed headers
183                 if (disallowedTrailerFieldNames.contains(
184                         trailerField.getKey().toLowerCase(Locale.ENGLISH))) {
185                     continue;
186                 }
187                 osw.write(trailerField.getKey());
188                 osw.write(':');
189                 osw.write(' ');
190                 osw.write(trailerField.getValue());
191                 osw.write("\r\n");
192             }
193             osw.close();
194             buffer.doWrite(ByteBuffer.wrap(baos.toByteArray()));
195
196             buffer.doWrite(crlfChunk);
197             crlfChunk.position(0).limit(crlfChunk.capacity());
198         }
199         buffer.end();
200     }
201
202
203     @Override
204     public void recycle() {
205         response = null;
206     }
207 }
208