1 /*
2  * Copyright 2008-2019 by Emeric Vernat
3  *
4  *     This file is part of Java Melody.
5  *
6  * Licensed under the Apache License, Version 2.0 (the "License");
7  * you may not use this file except in compliance with the License.
8  * You may obtain a copy of the License at
9  *
10  *     http://www.apache.org/licenses/LICENSE-2.0
11  *
12  * Unless required by applicable law or agreed to in writing, software
13  * distributed under the License is distributed on an "AS IS" BASIS,
14  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15  * See the License for the specific language governing permissions and
16  * limitations under the License.
17  */
18 package net.bull.javamelody.internal.web;
19
20 import java.io.ByteArrayOutputStream;
21 import java.io.IOException;
22 import java.io.OutputStream;
23 import java.util.zip.GZIPOutputStream;
24
25 import javax.servlet.ServletOutputStream;
26 import javax.servlet.WriteListener;
27 import javax.servlet.http.HttpServletResponse;
28
29 /**
30  * Implémentation de ServletOutputStream qui fonctionne avec le {@link CompressionServletResponseWrapper}.
31  * @author Emeric Vernat
32  */
33 class CompressionResponseStream extends ServletOutputStream {
34     private final int compressionThreshold;
35     private final HttpServletResponse response;
36     private OutputStream stream;
37
38     /**
39      * Construit un servlet output stream associé avec la réponse spécifiée.
40      * @param response HttpServletResponse
41      * @param compressionThreshold int
42      */
43     CompressionResponseStream(HttpServletResponse response, int compressionThreshold) {
44         super();
45         assert response != null;
46         assert compressionThreshold >= 0;
47         this.response = response;
48         this.compressionThreshold = compressionThreshold;
49         this.stream = new ByteArrayOutputStream(compressionThreshold);
50     }
51
52     /** {@inheritDoc} */
53     @Override
54     public void close() throws IOException {
55         if (stream instanceof ByteArrayOutputStream) {
56             final byte[] bytes = ((ByteArrayOutputStream) stream).toByteArray();
57             response.getOutputStream().write(bytes);
58             stream = response.getOutputStream();
59         }
60         stream.close();
61     }
62
63     /**
64      * Flushe les données bufferisées de cet output stream.
65      * @throws IOException e
66      */
67     @Override
68     public void flush() throws IOException {
69         stream.flush();
70     }
71
72     private void checkBufferSize(int length) throws IOException {
73         // check if we are buffering too large of a file
74         if (stream instanceof ByteArrayOutputStream) {
75             final ByteArrayOutputStream baos = (ByteArrayOutputStream) stream;
76             if (baos.size() + length > compressionThreshold) {
77                 // files too large to keep in memory are sent to the client
78                 flushToGZIP();
79             }
80         }
81     }
82
83     private void flushToGZIP() throws IOException {
84         if (stream instanceof ByteArrayOutputStream) {
85             // indication de compression,
86             // on utilise setHeader et non addHeader pour être compatible avec PJL compression filter
87             // en particulier dans le plugin grails vis à vis de l'autre plugin grails UiPerformance
88             response.setHeader("Content-Encoding", "gzip");
89             response.setHeader("Vary", "Accept-Encoding");
90
91             // make new gzip stream using the response output stream (content-encoding is in constructor)
92             final GZIPOutputStream gzipstream = new GZIPOutputStream(response.getOutputStream(),
93                     compressionThreshold);
94             // get existing bytes
95             final byte[] bytes = ((ByteArrayOutputStream) stream).toByteArray();
96             gzipstream.write(bytes);
97             // we are no longer buffering, send content via gzipstream
98             stream = gzipstream;
99         }
100     }
101
102     /** {@inheritDoc} */
103     @Override
104     public void write(int i) throws IOException {
105         // make sure we aren't over the buffer's limit
106         checkBufferSize(1);
107         // write the byte to the temporary output
108         stream.write(i);
109     }
110
111     /** {@inheritDoc} */
112     @Override
113     public void write(byte[] bytes) throws IOException {
114         write(bytes, 0, bytes.length);
115     }
116
117     /** {@inheritDoc} */
118     @Override
119     public void write(byte[] bytes, int off, int len) throws IOException {
120         if (len == 0) {
121             return;
122         }
123
124         // make sure we aren't over the buffer's limit
125         checkBufferSize(len);
126         // write the content to the buffer
127         stream.write(bytes, off, len);
128     }
129
130     @Override
131     public boolean isReady() {
132         return true;
133     }
134
135     @Override
136     public void setWriteListener(WriteListener writeListener) {
137         // nothing to do
138     }
139 }
140