001/*
002 * Licensed to the Apache Software Foundation (ASF) under one or more
003 * contributor license agreements.  See the NOTICE file distributed with
004 * this work for additional information regarding copyright ownership.
005 * The ASF licenses this file to You under the Apache License, Version 2.0
006 * (the "License"); you may not use this file except in compliance with
007 * the License.  You may obtain a copy of the License at
008 *
009 *      http://www.apache.org/licenses/LICENSE-2.0
010 *
011 * Unless required by applicable law or agreed to in writing, software
012 * distributed under the License is distributed on an "AS IS" BASIS,
013 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
014 * See the License for the specific language governing permissions and
015 * limitations under the License.
016 */
017package org.apache.juneau.rest.util;
018
019import java.io.*;
020
021import jakarta.servlet.*;
022import jakarta.servlet.http.*;
023
024/**
025 * Wraps an {@link HttpServletResponse} and caches the output stream in a separate buffer for debugging purposes.
026 *
027 */
028public class CachingHttpServletResponse extends HttpServletResponseWrapper {
029
030   /**
031    * Wraps the specified response inside a {@link CachingHttpServletResponse} if it isn't already.
032    *
033    * @param res The response to wrap.
034    * @return The wrapped request.
035    * @throws IOException Thrown by underlying content stream.
036    */
037   public static CachingHttpServletResponse wrap(HttpServletResponse res) throws IOException {
038      if (res instanceof CachingHttpServletResponse res2)
039         return res2;
040      return new CachingHttpServletResponse(res);
041   }
042
043   final ByteArrayOutputStream baos = new ByteArrayOutputStream();
044
045   final ServletOutputStream os;
046
047   /**
048    * Constructor.
049    *
050    * @param res The wrapped servlet response.
051    * @throws IOException Thrown by underlying stream.
052    */
053   protected CachingHttpServletResponse(HttpServletResponse res) throws IOException {
054      super(res);
055      os = res.getOutputStream();
056   }
057
058   /**
059    * Returns the content of the servlet response without consuming the stream.
060    *
061    * @return The content of the response.
062    */
063   public byte[] getContent() { return baos.toByteArray(); }
064
065   @Override
066   public ServletOutputStream getOutputStream() throws IOException {
067      return new ServletOutputStream() {
068
069         @Override
070         public void close() throws IOException {
071            os.close();
072         }
073
074         @Override
075         public void flush() throws IOException {
076            os.flush();
077         }
078
079         @Override
080         public boolean isReady() { return os.isReady(); }
081
082         @Override
083         public void setWriteListener(WriteListener writeListener) {
084            os.setWriteListener(writeListener);
085         }
086
087         @Override
088         public void write(int b) throws IOException {
089            baos.write(b);
090            os.write(b);
091         }
092      };
093   }
094}