001/*
002 * Copyright 2024 Revetware LLC.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 * http://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016
017package com.soklet.servlet.javax;
018
019import javax.annotation.Nonnull;
020import javax.annotation.Nullable;
021import javax.annotation.concurrent.NotThreadSafe;
022import javax.servlet.ServletOutputStream;
023import javax.servlet.WriteListener;
024import java.io.IOException;
025import java.io.OutputStream;
026import java.util.function.Consumer;
027
028import static java.lang.String.format;
029import static java.util.Objects.requireNonNull;
030
031/**
032 * @author <a href="https://www.revetkn.com">Mark Allen</a>
033 */
034@NotThreadSafe
035public class SokletServletOutputStream extends ServletOutputStream {
036        @Nonnull
037        private final OutputStream outputStream;
038        @Nonnull
039        private final Consumer<SokletServletOutputStream> writeOccurredCallback;
040        @Nonnull
041        private final Consumer<SokletServletOutputStream> writeFinalizedCallback;
042        @Nonnull
043        private Boolean writeFinalized;
044
045        public SokletServletOutputStream(@Nonnull OutputStream outputStream) {
046                this(requireNonNull(outputStream), null, null);
047        }
048
049        public SokletServletOutputStream(@Nonnull OutputStream outputStream,
050                                                                                                                                         @Nullable Consumer<SokletServletOutputStream> writeOccurredCallback,
051                                                                                                                                         @Nullable Consumer<SokletServletOutputStream> writeFinalizedCallback) {
052                super();
053                requireNonNull(outputStream);
054
055                if (writeOccurredCallback == null)
056                        writeOccurredCallback = (ignored) -> {};
057
058                if (writeFinalizedCallback == null)
059                        writeFinalizedCallback = (ignored) -> {};
060
061                this.outputStream = outputStream;
062                this.writeOccurredCallback = writeOccurredCallback;
063                this.writeFinalizedCallback = writeFinalizedCallback;
064                this.writeFinalized = false;
065        }
066
067        @Nonnull
068        protected OutputStream getOutputStream() {
069                return this.outputStream;
070        }
071
072        @Nonnull
073        protected Consumer<SokletServletOutputStream> getWriteOccurredCallback() {
074                return this.writeOccurredCallback;
075        }
076
077        @Nonnull
078        protected Consumer<SokletServletOutputStream> getWriteFinalizedCallback() {
079                return this.writeFinalizedCallback;
080        }
081
082        @Nonnull
083        protected Boolean getWriteFinalized() {
084                return this.writeFinalized;
085        }
086
087        protected void setWriteFinalized(@Nonnull Boolean writeFinalized) {
088                requireNonNull(writeFinalized);
089                this.writeFinalized = writeFinalized;
090        }
091
092// Implementation of ServletOutputStream methods below:
093
094        @Override
095        public void write(int b) throws IOException {
096                getOutputStream().write(b);
097                getWriteOccurredCallback().accept(this);
098        }
099
100        @Override
101        public boolean isReady() {
102                return !getWriteFinalized();
103        }
104
105        @Override
106        public void flush() throws IOException {
107                super.flush();
108                getOutputStream().flush();
109
110                if (!getWriteFinalized()) {
111                        setWriteFinalized(true);
112                        getWriteFinalizedCallback().accept(this);
113                }
114        }
115
116        @Override
117        public void close() throws IOException {
118                super.close();
119                getOutputStream().close();
120
121                if (!getWriteFinalized()) {
122                        setWriteFinalized(true);
123                        getWriteFinalizedCallback().accept(this);
124                }
125        }
126
127        @Override
128        public void setWriteListener(@Nonnull WriteListener writeListener) {
129                requireNonNull(writeListener);
130                throw new IllegalStateException(format("%s functionality is not supported", WriteListener.class.getSimpleName()));
131        }
132}