/**
* Licensed to Apereo under one or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information regarding copyright ownership. Apereo
* licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use
* this file except in compliance with the License. You may obtain a copy of the License at the
* following location:
*
* <p>http://www.apache.org/licenses/LICENSE-2.0
*
* <p>Unless required by applicable law or agreed to in writing, software distributed under the
* License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
* express or implied. See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apereo.portal.portlet.container.cache;
import com.google.common.base.Function;
import java.io.IOException;
import java.io.OutputStream;
import org.apache.commons.io.output.NullOutputStream;
import org.apereo.portal.utils.TeeOutputStream;
/**
* Subclass of {@link TeeOutputStream} that stops writing to the branch once the limit is hit by
* calling {@link #setBranch(OutputStream)} with {@link NullOutputStream}.
*
* <p>A callback {@link Function} to be executed when the limit is hit can be provided as well.
*
*/
public class LimitingTeeOutputStream extends TeeOutputStream {
private final long maximumBytes;
private final OutputStream branch;
private final Function<LimitingTeeOutputStream, ?> limitReachedCallback;
private long byteCount = 0;
private boolean limitReached = false;
public LimitingTeeOutputStream(long maximumBytes, OutputStream out, OutputStream branch) {
this(maximumBytes, out, branch, null);
}
public LimitingTeeOutputStream(
long maximumBytes,
OutputStream out,
OutputStream branch,
Function<LimitingTeeOutputStream, ?> limitReachedCallback) {
super(out, branch);
this.maximumBytes = maximumBytes;
this.branch = branch;
this.limitReachedCallback = limitReachedCallback;
}
/**
* Sets the byte count back to 0. If {@link #isLimitReached()} is true it is set back to false
* and the original branch {@link OutputStream} is set as the current branch again.
*/
public void resetByteCount() {
this.byteCount = 0;
if (limitReached) {
limitReached = false;
this.setBranch(branch);
}
}
/** @return Number of bytes seen so far */
public long getByteCount() {
return byteCount;
}
/** @return true if the limit has been reached */
public boolean isLimitReached() {
return limitReached;
}
@Override
protected void beforeWrite(int n) throws IOException {
this.byteCount += n;
if (this.maximumBytes > 0 && !this.limitReached && this.byteCount > this.maximumBytes) {
//Hit limit, replace tee'd OutputStream with a null OutputStream
this.limitReached = true;
this.setBranch(NullOutputStream.NULL_OUTPUT_STREAM);
if (this.limitReachedCallback != null) {
this.limitReachedCallback.apply(this);
}
}
}
}