001/*
002 * Copyright 2022-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;
018
019import com.soklet.internal.spring.LinkedCaseInsensitiveMap;
020import org.jspecify.annotations.NonNull;
021import org.jspecify.annotations.Nullable;
022
023import javax.annotation.concurrent.NotThreadSafe;
024import javax.annotation.concurrent.ThreadSafe;
025import java.io.EOFException;
026import java.io.IOException;
027import java.io.UncheckedIOException;
028import java.nio.ByteBuffer;
029import java.nio.channels.FileChannel;
030import java.nio.file.Files;
031import java.nio.file.Path;
032import java.util.LinkedHashSet;
033import java.util.Map;
034import java.util.Map.Entry;
035import java.util.Optional;
036import java.util.Set;
037import java.util.function.Consumer;
038
039import static java.lang.String.format;
040import static java.nio.file.StandardOpenOption.READ;
041import static java.util.Objects.requireNonNull;
042
043/**
044 * A finalized representation of a {@link Response}, suitable for sending to clients over the wire.
045 * <p>
046 * Your application's {@link ResponseMarshaler} is responsible for taking the {@link Response} returned by a <em>Resource Method</em> as input
047 * and converting its {@link Response#getBody()} to a {@link MarshaledResponseBody}.
048 * <p>
049 * For example, if a {@link Response} were to specify a body of {@code List.of("one", "two")}, a {@link ResponseMarshaler} might
050 * convert it to the JSON string {@code ["one", "two"]} and provide as output a corresponding {@link MarshaledResponse} with a byte-array-backed body containing UTF-8 bytes that represent {@code ["one", "two"]}.
051 * <p>
052 * Alternatively, your <em>Resource Method</em> might want to directly serve bytes to clients (e.g. an image or PDF) and skip the {@link ResponseMarshaler} entirely.
053 * To accomplish this, just have your <em>Resource Method</em> return a {@link MarshaledResponse} instance: this tells Soklet "I already know exactly what bytes I want to send; don't go through the normal marshaling process".
054 * <p>
055 * Instances can be acquired via the {@link #withResponse(Response)} or {@link #withStatusCode(Integer)} builder factory methods.
056 * Convenience instance factories are also available via {@link #fromResponse(Response)} and {@link #fromStatusCode(Integer)}.
057 * <p>
058 * Full documentation is available at <a href="https://www.soklet.com/docs/response-writing">https://www.soklet.com/docs/response-writing</a>.
059 *
060 * @author <a href="https://www.revetkn.com">Mark Allen</a>
061 */
062@ThreadSafe
063public final class MarshaledResponse {
064        @NonNull
065        private final Integer statusCode;
066        @NonNull
067        private final Map<@NonNull String, @NonNull Set<@NonNull String>> headers;
068        @NonNull
069        private final Set<@NonNull ResponseCookie> cookies;
070        @Nullable
071        private final MarshaledResponseBody body;
072
073        /**
074         * Acquires a builder for {@link MarshaledResponse} instances.
075         *
076         * @param response the logical response whose values are used to prime this builder
077         * @return the builder
078         */
079        @NonNull
080        public static Builder withResponse(@NonNull Response response) {
081                requireNonNull(response);
082
083                Object rawBody = response.getBody().orElse(null);
084                MarshaledResponseBody body = null;
085
086                if (rawBody != null && rawBody instanceof byte[] byteArrayBody)
087                        body = new MarshaledResponseBody.Bytes(byteArrayBody);
088
089                Builder builder = new Builder(response.getStatusCode())
090                                .headers(response.getHeaders())
091                                .cookies(response.getCookies());
092
093                if (body != null)
094                        builder.body(body);
095
096                return builder;
097        }
098
099        /**
100         * Creates a {@link MarshaledResponse} from a logical {@link Response} without additional customization.
101         *
102         * @param response the logical response whose values are used to construct this instance
103         * @return a {@link MarshaledResponse} instance
104         */
105        @NonNull
106        public static MarshaledResponse fromResponse(@NonNull Response response) {
107                return withResponse(response).build();
108        }
109
110        /**
111         * Acquires a builder for {@link MarshaledResponse} instances.
112         *
113         * @param statusCode the HTTP status code for this response
114         * @return the builder
115         */
116        @NonNull
117        public static Builder withStatusCode(@NonNull Integer statusCode) {
118                requireNonNull(statusCode);
119                return new Builder(statusCode);
120        }
121
122        /**
123         * Creates a {@link MarshaledResponse} with the given status code and no additional customization.
124         *
125         * @param statusCode the HTTP status code for this response
126         * @return a {@link MarshaledResponse} instance
127         */
128        @NonNull
129        public static MarshaledResponse fromStatusCode(@NonNull Integer statusCode) {
130                return withStatusCode(statusCode).build();
131        }
132
133        /**
134         * Vends a mutable copier seeded with this instance's data, suitable for building new instances.
135         *
136         * @return a copier for this instance
137         */
138        @NonNull
139        public Copier copy() {
140                return new Copier(this);
141        }
142
143        protected MarshaledResponse(@NonNull Builder builder) {
144                requireNonNull(builder);
145
146                this.statusCode = builder.statusCode;
147                this.headers = builder.headers == null ? Map.of() : new LinkedCaseInsensitiveMap<>(builder.headers);
148                this.cookies = builder.cookies == null ? Set.of() : new LinkedHashSet<>(builder.cookies);
149                this.body = builder.body;
150
151                // Verify headers are legal
152                for (Entry<String, Set<String>> entry : this.headers.entrySet()) {
153                        String headerName = entry.getKey();
154                        Set<String> headerValues = entry.getValue();
155
156                        for (String headerValue : headerValues)
157                                Utilities.validateHeaderNameAndValue(headerName, headerValue);
158                }
159        }
160
161        @Override
162        public String toString() {
163                return format("%s{statusCode=%s, headers=%s, cookies=%s, body=%s}", getClass().getSimpleName(),
164                                getStatusCode(), getHeaders(), getCookies(),
165                                format("%d bytes", getBodyLength()));
166        }
167
168        /**
169         * The HTTP status code for this response.
170         *
171         * @return the status code
172         */
173        @NonNull
174        public Integer getStatusCode() {
175                return this.statusCode;
176        }
177
178        /**
179         * The HTTP headers to write for this response.
180         * <p>
181         * Soklet writes one header line per value. If order matters, provide either a {@link java.util.SortedSet} or
182         * {@link java.util.LinkedHashSet} to preserve the desired ordering; otherwise values are naturally sorted for consistency.
183         *
184         * @return the headers to write
185         */
186        @NonNull
187        public Map<@NonNull String, @NonNull Set<@NonNull String>> getHeaders() {
188                return this.headers;
189        }
190
191        /**
192         * The HTTP cookies to write for this response.
193         *
194         * @return the cookies to write
195         */
196        @NonNull
197        public Set<@NonNull ResponseCookie> getCookies() {
198                return this.cookies;
199        }
200
201        /**
202         * The finalized HTTP response body to write, if available.
203         *
204         * @return the response body to write, or {@link Optional#empty()}) if no body should be written
205         */
206        @NonNull
207        public Optional<MarshaledResponseBody> getBody() {
208                return Optional.ofNullable(this.body);
209        }
210
211        /**
212         * The number of bytes this response body will write.
213         *
214         * @return the body length, or {@code 0} if no body is present
215         */
216        @NonNull
217        public Long getBodyLength() {
218                MarshaledResponseBody body = getBody().orElse(null);
219                return body == null ? 0L : body.getLength();
220        }
221
222        @Nullable
223        byte[] bodyBytesOrNull() {
224                MarshaledResponseBody body = getBody().orElse(null);
225
226                if (body == null)
227                        return null;
228
229                if (body instanceof MarshaledResponseBody.Bytes bytes)
230                        return bytes.getBytes();
231
232                if (body instanceof MarshaledResponseBody.File file)
233                        return materializeFile(file.getPath(), file.getOffset(), file.getCount());
234
235                if (body instanceof MarshaledResponseBody.FileChannel fileChannel)
236                        return materializeFileChannel(fileChannel.getChannel(), fileChannel.getOffset(), fileChannel.getCount(), fileChannel.getCloseOnComplete());
237
238                if (body instanceof MarshaledResponseBody.ByteBuffer byteBuffer)
239                        return materializeByteBuffer(byteBuffer.getBuffer());
240
241                throw new IllegalStateException(format("Unsupported marshaled response body type: %s", body.getClass().getName()));
242        }
243
244        @NonNull
245        byte[] bodyBytesOrEmpty() {
246                byte[] bytes = bodyBytesOrNull();
247                return bytes == null ? Utilities.emptyByteArray() : bytes;
248        }
249
250        /**
251         * Builder used to construct instances of {@link MarshaledResponse} via {@link MarshaledResponse#withResponse(Response)} or {@link MarshaledResponse#withStatusCode(Integer)}.
252         * <p>
253         * This class is intended for use by a single thread.
254         *
255         * @author <a href="https://www.revetkn.com">Mark Allen</a>
256         */
257        @NotThreadSafe
258        public static final class Builder {
259                @NonNull
260                private Integer statusCode;
261                @Nullable
262                private Set<@NonNull ResponseCookie> cookies;
263                @Nullable
264                private Map<@NonNull String, @NonNull Set<@NonNull String>> headers;
265                @Nullable
266                private MarshaledResponseBody body;
267
268                protected Builder(@NonNull Integer statusCode) {
269                        requireNonNull(statusCode);
270                        this.statusCode = statusCode;
271                }
272
273                @NonNull
274                public Builder statusCode(@NonNull Integer statusCode) {
275                        requireNonNull(statusCode);
276                        this.statusCode = statusCode;
277                        return this;
278                }
279
280                @NonNull
281                public Builder cookies(@Nullable Set<@NonNull ResponseCookie> cookies) {
282                        this.cookies = cookies;
283                        return this;
284                }
285
286                @NonNull
287                public Builder headers(@Nullable Map<@NonNull String, @NonNull Set<@NonNull String>> headers) {
288                        this.headers = headers;
289                        return this;
290                }
291
292                /**
293                 * Sets a byte-array-backed response body, or removes any current body if {@code bytes} is {@code null}.
294                 *
295                 * @param bytes the response bytes to write, or {@code null} for no body
296                 * @return this builder
297                 */
298                @NonNull
299                public Builder body(@Nullable byte[] bytes) {
300                        return bytes == null
301                                        ? withoutBody()
302                                        : body(new MarshaledResponseBody.Bytes(bytes));
303                }
304
305                /**
306                 * Sets a response body descriptor, or removes any current body if {@code body} is {@code null}.
307                 *
308                 * @param body the response body descriptor to write, or {@code null} for no body
309                 * @return this builder
310                 */
311                @NonNull
312                public Builder body(@Nullable MarshaledResponseBody body) {
313                        if (body == null)
314                                return withoutBody();
315
316                        this.body = body;
317                        return this;
318                }
319
320                /**
321                 * Sets a path-backed response body, or removes any current body if {@code path} is {@code null}.
322                 *
323                 * @param path the file path to write, or {@code null} for no body
324                 * @return this builder
325                 */
326                @NonNull
327                public Builder body(@Nullable Path path) {
328                        if (path == null)
329                                return withoutBody();
330
331                        this.body = fileBody(path);
332                        return this;
333                }
334
335                /**
336                 * Sets a ranged path-backed response body.
337                 *
338                 * @param path   the file path to write
339                 * @param offset the zero-based file offset from which response bytes should be written
340                 * @param count  the number of file bytes to write
341                 * @return this builder
342                 */
343                @NonNull
344                public Builder body(@NonNull Path path, @NonNull Long offset, @NonNull Long count) {
345                        requireNonNull(path);
346                        this.body = fileBody(path, offset, count);
347                        return this;
348                }
349
350                /**
351                 * Sets a file-channel-backed response body.
352                 *
353                 * @param fileChannel     the file channel to write
354                 * @param offset          the zero-based channel offset from which response bytes should be written
355                 * @param count           the number of channel bytes to write
356                 * @param closeOnComplete whether Soklet should close the channel after response completion
357                 * @return this builder
358                 */
359                @NonNull
360                public Builder body(@NonNull FileChannel fileChannel,
361                                                                                                @NonNull Long offset,
362                                                                                                @NonNull Long count,
363                                                                                                @NonNull Boolean closeOnComplete) {
364                        requireNonNull(fileChannel);
365                        this.body = fileChannelBody(fileChannel, offset, count, closeOnComplete);
366                        return this;
367                }
368
369                /**
370                 * Sets a byte-buffer-backed response body, or removes any current body if {@code byteBuffer} is {@code null}.
371                 *
372                 * @param byteBuffer the byte buffer to write, or {@code null} for no body
373                 * @return this builder
374                 */
375                @NonNull
376                public Builder body(@Nullable ByteBuffer byteBuffer) {
377                        if (byteBuffer == null)
378                                return withoutBody();
379
380                        this.body = new MarshaledResponseBody.ByteBuffer(byteBuffer);
381                        return this;
382                }
383
384                /**
385                 * Removes the response body from this builder.
386                 * <p>
387                 * If the current body owns a caller-supplied {@link FileChannel}, the channel is closed before it is
388                 * removed. Path-backed file bodies are lazy and do not hold open resources at this point.
389                 *
390                 * @return this builder
391                 */
392                @NonNull
393                public Builder withoutBody() {
394                        releaseBodyResources(this.body);
395                        this.body = null;
396                        return this;
397                }
398
399                @NonNull
400                public MarshaledResponse build() {
401                        return new MarshaledResponse(this);
402                }
403        }
404
405        /**
406         * Builder used to copy instances of {@link MarshaledResponse} via {@link MarshaledResponse#copy()}.
407         * <p>
408         * This class is intended for use by a single thread.
409         *
410         * @author <a href="https://www.revetkn.com">Mark Allen</a>
411         */
412        @NotThreadSafe
413        public static final class Copier {
414                @NonNull
415                private final Builder builder;
416
417                Copier(@NonNull MarshaledResponse marshaledResponse) {
418                        requireNonNull(marshaledResponse);
419
420                        this.builder = new Builder(marshaledResponse.getStatusCode())
421                                        .headers(new LinkedCaseInsensitiveMap<>(marshaledResponse.getHeaders()))
422                                        .cookies(new LinkedHashSet<>(marshaledResponse.getCookies()));
423
424                        marshaledResponse.getBody().ifPresent(this.builder::body);
425                }
426
427                @NonNull
428                public Copier statusCode(@NonNull Integer statusCode) {
429                        requireNonNull(statusCode);
430                        this.builder.statusCode(statusCode);
431                        return this;
432                }
433
434                @NonNull
435                public Copier headers(@NonNull Map<@NonNull String, @NonNull Set<@NonNull String>> headers) {
436                        this.builder.headers(headers);
437                        return this;
438                }
439
440                // Convenience method for mutation
441                @NonNull
442                public Copier headers(@NonNull Consumer<Map<@NonNull String, @NonNull Set<@NonNull String>>> headersConsumer) {
443                        requireNonNull(headersConsumer);
444
445                        if (this.builder.headers == null)
446                                this.builder.headers(new LinkedCaseInsensitiveMap<>());
447
448                        headersConsumer.accept(this.builder.headers);
449                        return this;
450                }
451
452                @NonNull
453                public Copier cookies(@Nullable Set<@NonNull ResponseCookie> cookies) {
454                        this.builder.cookies(cookies);
455                        return this;
456                }
457
458                // Convenience method for mutation
459                @NonNull
460                public Copier cookies(@NonNull Consumer<Set<@NonNull ResponseCookie>> cookiesConsumer) {
461                        requireNonNull(cookiesConsumer);
462
463                        if (this.builder.cookies == null)
464                                this.builder.cookies(new LinkedHashSet<>());
465
466                        cookiesConsumer.accept(this.builder.cookies);
467                        return this;
468                }
469
470                @NonNull
471                public Copier body(@Nullable byte[] bytes) {
472                        this.builder.body(bytes);
473                        return this;
474                }
475
476                @NonNull
477                public Copier body(@Nullable MarshaledResponseBody body) {
478                        this.builder.body(body);
479                        return this;
480                }
481
482                @NonNull
483                public Copier body(@Nullable Path path) {
484                        this.builder.body(path);
485                        return this;
486                }
487
488                @NonNull
489                public Copier body(@NonNull Path path, @NonNull Long offset, @NonNull Long count) {
490                        this.builder.body(path, offset, count);
491                        return this;
492                }
493
494                @NonNull
495                public Copier body(@NonNull FileChannel fileChannel,
496                                                                                         @NonNull Long offset,
497                                                                                         @NonNull Long count,
498                                                                                         @NonNull Boolean closeOnComplete) {
499                        this.builder.body(fileChannel, offset, count, closeOnComplete);
500                        return this;
501                }
502
503                @NonNull
504                public Copier body(@Nullable ByteBuffer byteBuffer) {
505                        this.builder.body(byteBuffer);
506                        return this;
507                }
508
509                /**
510                 * Removes the response body from this copier.
511                 * <p>
512                 * If the current body owns a caller-supplied {@link FileChannel}, the channel is closed before it is
513                 * removed. Path-backed file bodies are lazy and do not hold open resources at this point.
514                 *
515                 * @return this copier
516                 */
517                @NonNull
518                public Copier withoutBody() {
519                        this.builder.withoutBody();
520                        return this;
521                }
522
523                @NonNull
524                public MarshaledResponse finish() {
525                        return this.builder.build();
526                }
527        }
528
529        private static MarshaledResponseBody.File fileBody(@NonNull Path path) {
530                Long size = fileSize(path);
531                return new MarshaledResponseBody.File(path, 0L, size);
532        }
533
534        private static MarshaledResponseBody.File fileBody(@NonNull Path path, @NonNull Long offset, @NonNull Long count) {
535                Long size = fileSize(path);
536                validateRangeWithinLength(offset, count, size);
537                return new MarshaledResponseBody.File(path, offset, count);
538        }
539
540        private static MarshaledResponseBody.FileChannel fileChannelBody(@NonNull FileChannel fileChannel,
541                                                                                                                                                                                                                                                                         @NonNull Long offset,
542                                                                                                                                                                                                                                                                         @NonNull Long count,
543                                                                                                                                                                                                                                                                         @NonNull Boolean closeOnComplete) {
544                requireNonNull(fileChannel);
545                requireNonNull(closeOnComplete);
546                validateRangeWithinLength(offset, count, fileChannelSize(fileChannel));
547                return new MarshaledResponseBody.FileChannel(fileChannel, offset, count, closeOnComplete);
548        }
549
550        private static void releaseBodyResources(@Nullable MarshaledResponseBody body) {
551                if (!(body instanceof MarshaledResponseBody.FileChannel fileChannelBody) || !fileChannelBody.getCloseOnComplete())
552                        return;
553
554                try {
555                        fileChannelBody.getChannel().close();
556                } catch (IOException e) {
557                        throw new UncheckedIOException("Unable to close file-channel response body.", e);
558                }
559        }
560
561        @NonNull
562        private static Long fileSize(@NonNull Path path) {
563                requireNonNull(path);
564
565                if (!Files.isRegularFile(path))
566                        throw new IllegalArgumentException(format("File body path must reference a regular file: %s", path));
567
568                if (!Files.isReadable(path))
569                        throw new IllegalArgumentException(format("File body path must be readable: %s", path));
570
571                try {
572                        return Files.size(path);
573                } catch (IOException e) {
574                        throw new UncheckedIOException(format("Unable to determine file body length for path: %s", path), e);
575                }
576        }
577
578        @NonNull
579        private static Long fileChannelSize(@NonNull FileChannel fileChannel) {
580                requireNonNull(fileChannel);
581
582                try {
583                        return fileChannel.size();
584                } catch (IOException e) {
585                        throw new UncheckedIOException("Unable to determine file-channel body length.", e);
586                }
587        }
588
589        private static void validateRangeWithinLength(@NonNull Long offset, @NonNull Long count, @NonNull Long length) {
590                requireNonNull(offset);
591                requireNonNull(count);
592                requireNonNull(length);
593
594                if (offset < 0)
595                        throw new IllegalArgumentException("Offset must be >= 0.");
596
597                if (count < 0)
598                        throw new IllegalArgumentException("Count must be >= 0.");
599
600                if (Long.MAX_VALUE - offset < count)
601                        throw new IllegalArgumentException("Offset plus count exceeds maximum supported file position.");
602
603                if (offset + count > length)
604                        throw new IllegalArgumentException(format("Offset plus count must be <= body length %d.", length));
605        }
606
607        @NonNull
608        private static byte[] materializeFile(@NonNull Path path, @NonNull Long offset, @NonNull Long count) {
609                try (FileChannel fileChannel = FileChannel.open(path, READ)) {
610                        return materializeFileChannel(fileChannel, offset, count, false);
611                } catch (IOException e) {
612                        throw new UncheckedIOException(format("Unable to read file response body: %s", path), e);
613                }
614        }
615
616        @NonNull
617        private static byte[] materializeFileChannel(@NonNull FileChannel fileChannel,
618                                                                                                                                                                                         @NonNull Long offset,
619                                                                                                                                                                                         @NonNull Long count,
620                                                                                                                                                                                         @NonNull Boolean closeOnComplete) {
621                requireNonNull(fileChannel);
622                requireNonNull(offset);
623                requireNonNull(count);
624                requireNonNull(closeOnComplete);
625
626                if (count > Integer.MAX_VALUE)
627                        throw new IllegalStateException("Response body is too large to materialize as a byte array.");
628
629                byte[] bytes = new byte[Math.toIntExact(count)];
630                ByteBuffer buffer = ByteBuffer.wrap(bytes);
631                long position = offset;
632
633                try {
634                        while (buffer.hasRemaining()) {
635                                int read = fileChannel.read(buffer, position);
636                                if (read < 0)
637                                        throw new EOFException("File ended before the expected response body length was read.");
638                                if (read == 0)
639                                        throw new EOFException("File did not provide the expected response body length.");
640                                position += read;
641                        }
642                        return bytes;
643                } catch (IOException e) {
644                        throw new UncheckedIOException("Unable to materialize file-channel response body.", e);
645                } finally {
646                        if (closeOnComplete) {
647                                try {
648                                        fileChannel.close();
649                                } catch (IOException e) {
650                                        throw new UncheckedIOException("Unable to close file-channel response body.", e);
651                                }
652                        }
653                }
654        }
655
656        @NonNull
657        private static byte[] materializeByteBuffer(@NonNull ByteBuffer byteBuffer) {
658                requireNonNull(byteBuffer);
659
660                if (byteBuffer.remaining() > Integer.MAX_VALUE)
661                        throw new IllegalStateException("Response body is too large to materialize as a byte array.");
662
663                ByteBuffer duplicate = byteBuffer.asReadOnlyBuffer();
664                byte[] bytes = new byte[duplicate.remaining()];
665                duplicate.get(bytes);
666                return bytes;
667        }
668}