1 /*
2  * JBoss, Home of Professional Open Source.
3  * Copyright 2014 Red Hat, Inc., and individual contributors
4  * as indicated by the @author tags.
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
19 package io.undertow.servlet.spec;
20
21 import javax.servlet.DispatcherType;
22 import java.io.IOException;
23 import java.io.UnsupportedEncodingException;
24 import java.nio.ByteBuffer;
25 import java.nio.CharBuffer;
26 import java.nio.charset.Charset;
27 import java.nio.charset.CharsetEncoder;
28 import java.nio.charset.CoderResult;
29 import java.nio.charset.CodingErrorAction;
30 import java.util.Locale;
31
32 /**
33  * Real servlet print writer functionality, that is not limited by extending
34  * {@link java.io.PrintWriter}
35  * <p>
36  *
37  * @author Stuart Douglas
38  */
39 public class ServletPrintWriter {
40
41     private static final char[] EMPTY_CHAR = {};
42
43     private final ServletOutputStreamImpl outputStream;
44     private final String charset;
45     private CharsetEncoder charsetEncoder;
46     private boolean error = false;
47     private boolean closed = false;
48     private char[] underflow;
49
50     public ServletPrintWriter(final ServletOutputStreamImpl outputStream, final String charset) throws UnsupportedEncodingException {
51         this.charset = charset;
52         this.outputStream = outputStream;
53
54         //for some known charset we get optimistic and hope that
55         //only ascii will be output
56         //in this case we can avoid creating the encoder altogether
57         if (!charset.equalsIgnoreCase("utf-8") &&
58                 !charset.equalsIgnoreCase("iso-8859-1")) {
59             createEncoder();
60         }
61     }
62
63     private void createEncoder() {
64         this.charsetEncoder = Charset.forName(this.charset).newEncoder();
65         //replace malformed and unmappable with question marks
66         this.charsetEncoder.onUnmappableCharacter(CodingErrorAction.REPLACE);
67         this.charsetEncoder.onMalformedInput(CodingErrorAction.REPLACE);
68     }
69
70     public void flush() {
71         try {
72             outputStream.flush();
73         } catch (IOException e) {
74             error = true;
75         }
76     }
77
78     public void close() {
79
80         if (outputStream.getServletRequestContext().getOriginalRequest().getDispatcherType() == DispatcherType.INCLUDE) {
81             return;
82         }
83         if (closed) {
84             return;
85         }
86         closed = true;
87         try {
88             boolean done = false;
89             CharBuffer buffer;
90             if (underflow == null) {
91                 buffer = CharBuffer.wrap(EMPTY_CHAR);
92             } else {
93                 buffer = CharBuffer.wrap(underflow);
94                 underflow = null;
95             }
96             if (charsetEncoder != null) {
97                 do {
98                     ByteBuffer out = outputStream.underlyingBuffer();
99                     if (out == null) {
100                         //servlet output stream has already been closed
101                         error = true;
102                         return;
103                     }
104                     CoderResult result = charsetEncoder.encode(buffer, out, true);
105                     if (result.isOverflow()) {
106                         outputStream.flushInternal();
107                         if (out.remaining() == 0) {
108                             outputStream.close();
109                             error = true;
110                             return;
111                         }
112                     } else {
113                         done = true;
114                     }
115                 } while (!done);
116             }
117             outputStream.close();
118         } catch (IOException e) {
119             error = true;
120         }
121     }
122
123     public boolean checkError() {
124         flush();
125         return error;
126     }
127
128     public void write(final CharBuffer input) {
129         ByteBuffer buffer = outputStream.underlyingBuffer();
130         if (buffer == null) {
131             //stream has been closed
132             error = true;
133             return;
134         }
135         try {
136             if (!buffer.hasRemaining()) {
137                 outputStream.flushInternal();
138                 if (!buffer.hasRemaining()) {
139                     error = true;
140                     return;
141                 }
142             }
143
144             if (charsetEncoder == null) {
145                 createEncoder();
146             }
147             final CharBuffer cb;
148             if (underflow == null) {
149                 cb = input;
150             } else {
151                 char[] newArray = new char[underflow.length + input.remaining()];
152                 System.arraycopy(underflow, 0, newArray, 0, underflow.length);
153                 input.get(newArray, underflow.length, input.remaining());
154                 cb = CharBuffer.wrap(newArray);
155                 underflow = null;
156             }
157             int last = -1;
158             while (cb.hasRemaining()) {
159                 int remaining = buffer.remaining();
160                 CoderResult result = charsetEncoder.encode(cb, buffer, false);
161                 outputStream.updateWritten(remaining - buffer.remaining());
162                 if (result.isOverflow() || !buffer.hasRemaining()) {
163                     outputStream.flushInternal();
164                     if (!buffer.hasRemaining()) {
165                         error = true;
166                         return;
167                     }
168                 }
169                 if (result.isUnderflow()) {
170                     underflow = new char[cb.remaining()];
171                     cb.get(underflow);
172                     return;
173                 }
174                 if (result.isError()) {
175                     error = true;
176                     return;
177                 }
178                 if (result.isUnmappable()) {
179                     //this should not happen
180                     error = true;
181                     return;
182                 }
183                 if (last == cb.remaining()) {
184                     underflow = new char[cb.remaining()];
185                     cb.get(underflow);
186                     return;
187                 }
188                 last = cb.remaining();
189             }
190         } catch (IOException e) {
191             error = true;
192         }
193     }
194
195     public void write(final int c) {
196         write(Character.toString((char)c));
197     }
198
199     public void write(final char[] buf, final int off, final int len) {
200         if(charsetEncoder == null) {
201             try {
202                 ByteBuffer buffer = outputStream.underlyingBuffer();
203                 if(buffer == null) {
204                     //already closed
205                     error = true;
206                     return;
207                 }
208                 //fast path, basically we are hoping this is ascii only
209                 int remaining = buffer.remaining();
210                 boolean ok = true;
211                 //so we have a pure ascii buffer, just write it out and skip all the encoder cost
212
213                 int end = off + len;
214                 int i = off;
215                 int flushPos = i + remaining;
216                 while (ok && i < end) {
217                     int realEnd = Math.min(end, flushPos);
218                     for (; i < realEnd; ++i) {
219                         char c = buf[i];
220                         if (c > 127) {
221                             ok = false;
222                             break;
223                         } else {
224                             buffer.put((byte) c);
225                         }
226                     }
227                     if (i == flushPos) {
228                         outputStream.flushInternal();
229                         flushPos = i + buffer.remaining();
230                     }
231                 }
232                 outputStream.updateWritten(remaining - buffer.remaining());
233                 if (ok) {
234                     return;
235                 }
236                 final CharBuffer cb = CharBuffer.wrap(buf, i, len - (i - off));
237                 write(cb);
238                 return;
239             } catch (IOException e) {
240                 error = false;
241                 return;
242             }
243
244         }
245         final CharBuffer cb = CharBuffer.wrap(buf, off, len);
246         write(cb);
247     }
248
249     public void write(final char[] buf) {
250         write(buf,0, buf.length);
251     }
252
253     public void write(final String s, final int off, final int len) {
254         if(charsetEncoder == null) {
255             try {
256                 ByteBuffer buffer = outputStream.underlyingBuffer();
257                 if(buffer == null) {
258                     //already closed
259                     error = true;
260                     return;
261                 }
262                 //fast path, basically we are hoping this is ascii only
263                 int remaining = buffer.remaining();
264                 boolean ok = true;
265                 //so we have a pure ascii buffer, just write it out and skip all the encoder cost
266
267                 int end = off + len;
268                 int i = off;
269                 int fpos = i + remaining;
270                 for (; i < end; ++i) {
271                     if (i == fpos) {
272                         outputStream.flushInternal();
273                         fpos = i + buffer.remaining();
274                     }
275                     char c = s.charAt(i);
276                     if (c > 127) {
277                         ok = false;
278                         break;
279                     }
280                     buffer.put((byte) c);
281                 }
282                 outputStream.updateWritten(remaining - buffer.remaining());
283                 if (ok) {
284                     return;
285                 }
286                 //wrap(String, off, len) acts wrong in the presence of multi byte characters
287                 final CharBuffer cb = CharBuffer.wrap(s.toCharArray(), i, len - (i - off));
288                 write(cb);
289                 return;
290             } catch (IOException e) {
291                 error = false;
292                 return;
293             }
294
295         }
296         final CharBuffer cb = CharBuffer.wrap(s, off, off + len);
297         write(cb);
298     }
299
300     public void write(final String s) {
301         write(s, 0, s.length());
302     }
303
304     public void print(final boolean b) {
305         write(Boolean.toString(b));
306     }
307
308     public void print(final char c) {
309         write(Character.toString(c));
310     }
311
312     public void print(final int i) {
313         write(Integer.toString(i));
314     }
315
316     public void print(final long l) {
317         write(Long.toString(l));
318     }
319
320     public void print(final float f) {
321         write(Float.toString(f));
322     }
323
324     public void print(final double d) {
325         write(Double.toString(d));
326     }
327
328     public void print(final char[] s) {
329         write(CharBuffer.wrap(s));
330     }
331
332     public void print(final String s) {
333         write(s == null ? "null" : s);
334     }
335
336     public void print(final Object obj) {
337         write(obj == null ? "null" : obj.toString());
338     }
339
340     public void println() {
341         print('\n');
342     }
343
344     public void println(final boolean b) {
345         print(b);
346         print('\n');
347     }
348
349     public void println(final char c) {
350         print(c);
351         print('\n');
352     }
353
354     public void println(final int i) {
355         print(i);
356         print('\n');
357     }
358
359     public void println(final long l) {
360         print(l);
361         print('\n');
362     }
363
364     public void println(final float f) {
365         print(f);
366         print('\n');
367     }
368
369     public void println(final double d) {
370         print(d);
371         print('\n');
372     }
373
374     public void println(final char[] s) {
375         print(s);
376         print('\n');
377     }
378
379     public void println(final String s) {
380         print(s);
381         print('\n');
382     }
383
384     public void println(final Object obj) {
385         print(obj);
386         print('\n');
387     }
388
389     public void printf(final String format, final Object... args) {
390         print(String.format(format, args));
391     }
392
393     public void printf(final Locale l, final String format, final Object... args) {
394         print(String.format(l, format, args));
395     }
396
397
398     public void format(final String format, final Object... args) {
399         printf(format, args);
400     }
401
402     public void format(final Locale l, final String format, final Object... args) {
403         printf(l, format, args);
404     }
405
406     public void append(final CharSequence csq) {
407         if (csq == null) {
408             write("null");
409         } else {
410             write(csq.toString());
411         }
412     }
413
414     public void append(final CharSequence csq, final int start, final int end) {
415         CharSequence cs = (csq == null ? "null" : csq);
416         write(cs.subSequence(start, end).toString());
417     }
418
419     public void append(final char c) {
420         write(c);
421     }
422
423 }
424