001/*
002 * Copyright 2024-2026 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.jakarta;
018
019import jakarta.servlet.ServletOutputStream;
020import jakarta.servlet.WriteListener;
021import org.jspecify.annotations.NonNull;
022import org.jspecify.annotations.Nullable;
023
024import javax.annotation.concurrent.NotThreadSafe;
025import java.io.IOException;
026import java.io.OutputStream;
027import java.util.function.BiConsumer;
028import java.util.function.Consumer;
029
030import static java.lang.String.format;
031import static java.util.Objects.requireNonNull;
032
033/**
034 * Soklet integration implementation of {@link ServletOutputStream}.
035 *
036 * @author <a href="https://www.revetkn.com">Mark Allen</a>
037 */
038@NotThreadSafe
039public final class SokletServletOutputStream extends ServletOutputStream {
040        @NonNull
041        private final OutputStream outputStream;
042        @NonNull
043        private final BiConsumer<@NonNull SokletServletOutputStream, @NonNull Integer> onWriteOccurred;
044        @NonNull
045        private final Consumer<@NonNull SokletServletOutputStream> onWriteFinalized;
046        @NonNull
047        private Boolean writeFinalized;
048        @NonNull
049        private Boolean closed;
050
051        @NonNull
052        public static SokletServletOutputStream fromOutputStream(@NonNull OutputStream outputStream) {
053                requireNonNull(outputStream);
054                return withOutputStream(outputStream).build();
055        }
056
057        @NonNull
058        public static Builder withOutputStream(@NonNull OutputStream outputStream) {
059                return new Builder(outputStream);
060        }
061
062        private SokletServletOutputStream(@NonNull Builder builder) {
063                requireNonNull(builder);
064                requireNonNull(builder.outputStream);
065
066                this.outputStream = builder.outputStream;
067                this.onWriteOccurred = builder.onWriteOccurred != null ? builder.onWriteOccurred : (ignored1, ignored2) -> {};
068                this.onWriteFinalized = builder.onWriteFinalized != null ? builder.onWriteFinalized : (ignored) -> {};
069                this.writeFinalized = false;
070                this.closed = false;
071        }
072
073        /**
074         * Builder used to construct instances of {@link SokletServletOutputStream}.
075         * <p>
076         * This class is intended for use by a single thread.
077         *
078         * @author <a href="https://www.revetkn.com">Mark Allen</a>
079         */
080        @NotThreadSafe
081        public static class Builder {
082                @NonNull
083                private OutputStream outputStream;
084                @Nullable
085                private BiConsumer<@NonNull SokletServletOutputStream, @NonNull Integer> onWriteOccurred;
086                @Nullable
087                private Consumer<@NonNull SokletServletOutputStream> onWriteFinalized;
088
089                @NonNull
090                private Builder(@NonNull OutputStream outputStream) {
091                        requireNonNull(outputStream);
092                        this.outputStream = outputStream;
093                }
094
095                @NonNull
096                public Builder outputStream(@NonNull OutputStream outputStream) {
097                        requireNonNull(outputStream);
098                        this.outputStream = outputStream;
099                        return this;
100                }
101
102                @NonNull
103                public Builder onWriteOccurred(@Nullable BiConsumer<@NonNull SokletServletOutputStream, @NonNull Integer> onWriteOccurred) {
104                        this.onWriteOccurred = onWriteOccurred;
105                        return this;
106                }
107
108                @NonNull
109                public Builder onWriteFinalized(@Nullable Consumer<@NonNull SokletServletOutputStream> onWriteFinalized) {
110                        this.onWriteFinalized = onWriteFinalized;
111                        return this;
112                }
113
114                @NonNull
115                public SokletServletOutputStream build() {
116                        return new SokletServletOutputStream(this);
117                }
118        }
119
120        @NonNull
121        private OutputStream getOutputStream() {
122                return this.outputStream;
123        }
124
125        @NonNull
126        private BiConsumer<@NonNull SokletServletOutputStream, @NonNull Integer> getOnWriteOccurred() {
127                return this.onWriteOccurred;
128        }
129
130        @NonNull
131        private Consumer<@NonNull SokletServletOutputStream> getOnWriteFinalized() {
132                return this.onWriteFinalized;
133        }
134
135        @NonNull
136        private Boolean getWriteFinalized() {
137                return this.writeFinalized;
138        }
139
140        private void setWriteFinalized(@NonNull Boolean writeFinalized) {
141                requireNonNull(writeFinalized);
142                this.writeFinalized = writeFinalized;
143        }
144
145        @NonNull
146        private Boolean getClosed() {
147                return this.closed;
148        }
149
150        private void setClosed(@NonNull Boolean closed) {
151                requireNonNull(closed);
152                this.closed = closed;
153        }
154
155        private void ensureOpen() throws IOException {
156                if (getClosed())
157                        throw new IOException("Stream is closed");
158        }
159
160// Implementation of ServletOutputStream methods below:
161
162        @Override
163        public void write(int b) throws IOException {
164                ensureOpen();
165                getOutputStream().write(b);
166                getOnWriteOccurred().accept(this, 1);
167        }
168
169        @Override
170        public boolean isReady() {
171                return !getClosed();
172        }
173
174        @Override
175        public void write(@NonNull byte[] b,
176                                                                                int off,
177                                                                                int len) throws IOException {
178                requireNonNull(b);
179                ensureOpen();
180                if (len == 0)
181                        return;
182
183                getOutputStream().write(b, off, len);
184                getOnWriteOccurred().accept(this, len);
185        }
186
187        @Override
188        public void flush() throws IOException {
189                ensureOpen();
190                super.flush();
191                getOutputStream().flush();
192
193                if (!getWriteFinalized()) {
194                        setWriteFinalized(true);
195                        getOnWriteFinalized().accept(this);
196                }
197        }
198
199        @Override
200        public void close() throws IOException {
201                if (getClosed())
202                        return;
203
204                try {
205                        super.close();
206                        getOutputStream().close();
207                } finally {
208                        setClosed(true);
209                        if (!getWriteFinalized()) {
210                                setWriteFinalized(true);
211                                getOnWriteFinalized().accept(this);
212                        }
213                }
214        }
215
216        @Override
217        public void setWriteListener(@NonNull WriteListener writeListener) {
218                requireNonNull(writeListener);
219                throw new IllegalStateException(format("%s functionality is not supported", WriteListener.class.getSimpleName()));
220        }
221}