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.server.util;
018
019import static org.apache.juneau.commons.utils.Shorts.*;
020
021import java.io.*;
022
023import jakarta.servlet.*;
024
025/**
026 * ServletInputStream wrapper around a normal input stream with support for limiting input.
027 *
028 */
029@SuppressWarnings({
030   "resource" // is and sis are servlet-container-managed streams passed in by the caller; closed by the container
031})
032public class BoundedServletInputStream extends ServletInputStream {
033
034   private final InputStream is;
035   private final ServletInputStream sis;
036   private long remain;
037
038   /**
039    * Wraps the specified byte array.
040    *
041    * @param b The byte contents of the stream. Must not be <jk>null</jk>.
042    */
043   public BoundedServletInputStream(byte[] b) {
044      this(new ByteArrayInputStream(b), Long.MAX_VALUE);
045   }
046
047   /**
048    * Wraps the specified input stream.
049    *
050    * @param is The input stream to wrap. Must not be <jk>null</jk>.
051    * @param max The maximum number of bytes to read from the stream.
052    */
053   public BoundedServletInputStream(InputStream is, long max) {
054      this.is = is;
055      this.sis = null;
056      this.remain = max;
057   }
058
059   /**
060    * Wraps the specified input stream.
061    *
062    * @param sis The input stream to wrap. Must not be <jk>null</jk>.
063    * @param max The maximum number of bytes to read from the stream.
064    */
065   public BoundedServletInputStream(ServletInputStream sis, long max) {
066      this.sis = sis;
067      this.is = sis;
068      this.remain = max;
069   }
070
071   @Override /* Overridden from InputStream */
072   public int available() throws IOException {
073      if (remain <= 0)
074         return 0;
075      return is.available();
076   }
077
078   @Override /* Overridden from InputStream */
079   public void close() throws IOException {
080      is.close();
081   }
082
083   @Override /* Overridden from ServletInputStream */
084   public boolean isFinished() { return sis != null && sis.isFinished(); }
085
086   @Override /* Overridden from ServletInputStream */
087   public boolean isReady() { return sis == null || sis.isReady(); }
088
089   @Override /* Overridden from InputStream */
090   public synchronized void mark(int limit) {
091      is.mark(limit);
092   }
093
094   @Override /* Overridden from InputStream */
095   public boolean markSupported() {
096      return is.markSupported();
097   }
098
099   @Override /* Overridden from InputStream */
100   public int read() throws IOException {
101      if (remain <= 0)
102         return probeAtLimit();
103      var r = is.read();
104      if (r != -1)
105         remain--;
106      return r;
107   }
108
109   @Override /* Overridden from InputStream */
110   public int read(byte[] b) throws IOException {
111      return read(b, 0, b.length);
112   }
113
114   @Override /* Overridden from InputStream */
115   public int read(byte[] b, int off, int len) throws IOException {
116      if (len == 0)
117         return 0;
118      if (remain <= 0)
119         return probeAtLimit();
120      int r = is.read(b, off, (int)Math.min(len, remain));
121      if (r > 0)
122         remain -= r;
123      return r;
124   }
125
126   @Override /* Overridden from InputStream */
127   public synchronized void reset() throws IOException {
128      is.reset();
129   }
130
131   @Override /* Overridden from ServletInputStream */
132   public void setReadListener(ReadListener arg0) {
133      if (nn(sis))
134         sis.setReadListener(arg0);
135   }
136
137   @Override /* Overridden from InputStream */
138   public long skip(long n) throws IOException {
139      long toSkip = Math.min(n, remain);
140      long r = is.skip(toSkip);
141      decrement(r);
142      return r;
143   }
144
145   /**
146    * Called when the limit has been reached: returns -1 if the underlying stream is also at end-of-stream
147    * (i.e. the body was exactly the maximum size), otherwise the limit has been exceeded.
148    */
149   private int probeAtLimit() throws IOException {
150      if (is.read() == -1)
151         return -1;
152      throw ioex("Input limit exceeded.  See @Rest(maxInput).");
153   }
154
155   private void decrement(long count) throws IOException {
156      remain -= count;
157      if (remain < 0)
158         throw ioex("Input limit exceeded.  See @Rest(maxInput).");
159   }
160}