001/*
002 * ====================================================================
003 * Licensed to the Apache Software Foundation (ASF) under one
004 * or more contributor license agreements.  See the NOTICE file
005 * distributed with this work for additional information
006 * regarding copyright ownership.  The ASF licenses this file
007 * to you under the Apache License, Version 2.0 (the
008 * "License"); you may not use this file except in compliance
009 * with the License.  You may obtain a copy of the License at
010 *
011 *   http://www.apache.org/licenses/LICENSE-2.0
012 *
013 * Unless required by applicable law or agreed to in writing,
014 * software distributed under the License is distributed on an
015 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
016 * KIND, either express or implied.  See the License for the
017 * specific language governing permissions and limitations
018 * under the License.
019 * ====================================================================
020 *
021 * This software consists of voluntary contributions made by many
022 * individuals on behalf of the Apache Software Foundation.  For more
023 * information on the Apache Software Foundation, please see
024 * <http://www.apache.org/>.
025 *
026 */
027
028package org.apache.http.impl.io;
029
030import java.io.IOException;
031import java.io.OutputStream;
032
033import org.apache.http.io.SessionOutputBuffer;
034import org.apache.http.util.Args;
035
036/**
037 * Output stream that cuts off after a defined number of bytes. This class
038 * is used to send content of HTTP messages where the end of the content entity
039 * is determined by the value of the {@code Content-Length header}.
040 * Entities transferred using this stream can be maximum {@link Long#MAX_VALUE}
041 * long.
042 * <p>
043 * Note that this class NEVER closes the underlying stream, even when close
044 * gets called.  Instead, the stream will be marked as closed and no further
045 * output will be permitted.
046 *
047 * @since 4.0
048 */
049public class ContentLengthOutputStream extends OutputStream {
050
051    /**
052     * Wrapped session output buffer.
053     */
054    private final SessionOutputBuffer out;
055
056    /**
057     * The maximum number of bytes that can be written the stream. Subsequent
058     * write operations will be ignored.
059     */
060    private final long contentLength;
061
062    /** Total bytes written */
063    private long total = 0;
064
065    /** True if the stream is closed. */
066    private boolean closed = false;
067
068    /**
069     * Wraps a session output buffer and cuts off output after a defined number
070     * of bytes.
071     *
072     * @param out The session output buffer
073     * @param contentLength The maximum number of bytes that can be written to
074     * the stream. Subsequent write operations will be ignored.
075     *
076     * @since 4.0
077     */
078    public ContentLengthOutputStream(final SessionOutputBuffer out, final long contentLength) {
079        super();
080        this.out = Args.notNull(out, "Session output buffer");
081        this.contentLength = Args.notNegative(contentLength, "Content length");
082    }
083
084    /**
085     * <p>Does not close the underlying socket output.</p>
086     *
087     * @throws IOException If an I/O problem occurs.
088     */
089    @Override
090    public void close() throws IOException {
091        if (!this.closed) {
092            this.closed = true;
093            this.out.flush();
094        }
095    }
096
097    @Override
098    public void flush() throws IOException {
099        this.out.flush();
100    }
101
102    @Override
103    public void write(final byte[] b, final int off, final int len) throws IOException {
104        if (this.closed) {
105            throw new IOException("Attempted write to closed stream.");
106        }
107        if (this.total < this.contentLength) {
108            final long max = this.contentLength - this.total;
109            int chunk = len;
110            if (chunk > max) {
111                chunk = (int) max;
112            }
113            this.out.write(b, off, chunk);
114            this.total += chunk;
115        }
116    }
117
118    @Override
119    public void write(final byte[] b) throws IOException {
120        write(b, 0, b.length);
121    }
122
123    @Override
124    public void write(final int b) throws IOException {
125        if (this.closed) {
126            throw new IOException("Attempted write to closed stream.");
127        }
128        if (this.total < this.contentLength) {
129            this.out.write(b);
130            this.total++;
131        }
132    }
133
134}