Skip to content
rachel/workerd
File

Blob: src/workerd/api/system-streams.c++

16.0 KB
1// Copyright (c) 2017-2022 Cloudflare, Inc.
2// Licensed under the Apache 2.0 license found in the LICENSE file or at:
3//     https://opensource.org/licenses/Apache-2.0
4 
5#include "system-streams.h"
6 
7#include "util.h"
8 
9#include <kj/compat/brotli.h>
10#include <kj/compat/gzip.h>
11#include <kj/one-of.h>
12 
13namespace workerd::api {
14 
15// =======================================================================================
16// EncodedAsyncInputStream
17 
18namespace {
19 
20// A wrapper around a native `kj::AsyncInputStream` which knows the underlying encoding of the
21// stream and whether or not it requires pending event registration.
22class EncodedAsyncInputStream final: public ReadableStreamSource {
23 public:
24  explicit EncodedAsyncInputStream(
25      kj::Own<kj::AsyncInputStream> inner, StreamEncoding encoding, IoContext& context);
26 
27  // Read bytes in identity encoding. If the stream is not already in identity encoding, it will be
28  // converted to identity encoding via an appropriate stream wrapper.
29  kj::Promise<size_t> tryRead(void* buffer, size_t minBytes, size_t maxBytes) override;
30 
31  StreamEncoding getPreferredEncoding() override {
32    return encoding;
33  }
34 
35  // Return the number of bytes, if known, which this input stream will produce if the sink is known
36  // to be of a particular encoding.
37  //
38  // It is likely an error to call this function without immediately following it with a pumpTo()
39  // to a EncodedAsyncOutputStream of that exact encoding.
40  kj::Maybe<uint64_t> tryGetLength(StreamEncoding outEncoding) override;
41 
42  // Consume this stream and return two streams with the same encoding that read the exact same
43  // data.
44  //
45  // This implementation of `tryTee()` is not technically required for correctness, but prevents
46  // re-encoding (and converting Content-Length responses to chunk-encoded responses) gzip and
47  // brotli streams.
48  kj::Maybe<Tee> tryTee(uint64_t limit) override;
49 
50  void cancel(kj::Exception reason) override;
51 
52 private:
53  friend class EncodedAsyncOutputStream;
54 
55  void ensureIdentityEncoding();
56 
57  kj::Own<kj::AsyncInputStream> inner;
58  StreamEncoding encoding;
59  kj::Canceler canceler;
60 
61  IoContext& ioContext;
62};
63 
64EncodedAsyncInputStream::EncodedAsyncInputStream(
65    kj::Own<kj::AsyncInputStream> inner, StreamEncoding encoding, IoContext& context)
66    : inner(kj::mv(inner)),
67      encoding(encoding),
68      ioContext(context) {}
69 
70kj::Promise<size_t> EncodedAsyncInputStream::tryRead(
71    void* buffer, size_t minBytes, size_t maxBytes) {
72  ensureIdentityEncoding();
73 
74  return kj::evalNow([&]() {
75    return canceler.wrap(inner->tryRead(buffer, minBytes, maxBytes))
76        .attach(ioContext.registerPendingEvent());
77  }).catch_([](kj::Exception&& exception) -> kj::Promise<size_t> {
78    KJ_IF_SOME(e,
79        translateKjException(exception,
80            {
81              {"gzip compressed stream ended prematurely"_kj,
82                "Gzip compressed stream ended prematurely."_kj},
83              {"gzip decompression failed"_kj, "Gzip decompression failed."},
84              {"brotli state allocation failed"_kj, "Brotli state allocation failed."},
85              {"invalid brotli window size"_kj, "Invalid brotli window size."},
86              {"invalid brotli compression level"_kj, "Invalid brotli compression level."},
87              {"brotli window size too big"_kj, "Brotli window size too big."},
88              {"brotli decompression failed"_kj, "Brotli decompression failed."},
89              {"brotli compression failed"_kj, "Brotli compression failed."},
90              {"brotli compressed stream ended prematurely"_kj,
91                "Brotli compressed stream ended prematurely."},
92            })) {
93      return kj::mv(e);
94    }
95 
96    // Let the original exception pass through, since it is likely already a jsg.TypeError.
97    return kj::mv(exception);
98  });
99}
100 
101kj::Maybe<uint64_t> EncodedAsyncInputStream::tryGetLength(StreamEncoding outEncoding) {
102  if (outEncoding == encoding) {
103    return inner->tryGetLength();
104  } else {
105    // We have no idea what the length will be once encoded/decoded.
106    return kj::none;
107  }
108}
109 
110kj::Maybe<ReadableStreamSource::Tee> EncodedAsyncInputStream::tryTee(uint64_t limit) {
111  // We tee the stream in its original encoding, because chances are highest that we'll be pumped
112  // to sinks that are of the same encoding, and only read in identity encoding no more than once.
113  //
114  // Additionally, we should propagate the fact that this stream is a native stream to the branches
115  // of the tee, so that branches which fall behind their siblings (and thus are reading from the
116  // tee buffer) still register pending events correctly.
117  auto tee = kj::newTee(kj::mv(inner), limit);
118 
119  Tee result;
120  result.branches[0] = newSystemStream(newTeeErrorAdapter(kj::mv(tee.branches[0])), encoding);
121  result.branches[1] = newSystemStream(newTeeErrorAdapter(kj::mv(tee.branches[1])), encoding);
122  return kj::mv(result);
123}
124 
125void EncodedAsyncInputStream::cancel(kj::Exception reason) {
126  // Cancel any pending read operations. This will cause the wrapped promises to be rejected
127  // with a cancellation exception, which properly cleans up the BlockedRead state in AsyncPipe
128  // before the pipe itself is destroyed.
129  canceler.cancel(kj::mv(reason));
130}
131 
132void EncodedAsyncInputStream::ensureIdentityEncoding() {
133  // Decompression gets added to the stream here if needed based on the content encoding.
134  if (encoding == StreamEncoding::GZIP) {
135    inner = kj::heap<kj::GzipAsyncInputStream>(*inner).attach(kj::mv(inner));
136    encoding = StreamEncoding::IDENTITY;
137  } else if (encoding == StreamEncoding::BROTLI) {
138    inner = kj::heap<kj::BrotliAsyncInputStream>(*inner).attach(kj::mv(inner));
139    encoding = StreamEncoding::IDENTITY;
140  } else {
141    // We currently support gzip and brotli as non-identity content encodings.
142    KJ_ASSERT(encoding == StreamEncoding::IDENTITY);
143  }
144}
145 
146// =======================================================================================
147// EncodedAsyncOutputStream
148 
149// A wrapper around a native `kj::AsyncOutputStream` which knows the underlying encoding of the
150// stream and optimizes pumps from `EncodedAsyncInputStream`.
151//
152// The inner will be held on to right up until either end() or abort() is called.
153// This is important because some AsyncOutputStream implementations perform cleanup
154// operations equivalent to end() in their destructors (for instance HttpChunkedEntityWriter).
155// If we wait to clear the kj::Own when the EncodedAsyncOutputStream is destroyed, and the
156// EncodedAsyncOutputStream is owned (for instance) by an IoOwn, then the lifetime of the
157// inner may be extended past when it should. Eventually, kj::AsyncOutputStream should
158// probably have a distinct end() method of its own that we can defer to, but until it
159// does, it is important for us to release it as soon as end() or abort() are called.
160class EncodedAsyncOutputStream final: public WritableStreamSink {
161 public:
162  explicit EncodedAsyncOutputStream(
163      kj::Own<kj::AsyncOutputStream> inner, StreamEncoding encoding, IoContext& context);
164 
165  kj::Promise<void> write(kj::ArrayPtr<const byte> buffer) override;
166  kj::Promise<void> write(kj::ArrayPtr<const kj::ArrayPtr<const kj::byte>> pieces) override;
167 
168  kj::Maybe<kj::Promise<DeferredProxy<void>>> tryPumpFrom(
169      ReadableStreamSource& input, bool end) override;
170 
171  kj::Promise<void> end() override;
172 
173  void abort(kj::Exception reason) override;
174 
175  StreamEncoding disownEncodingResponsibility() override;
176 
177 private:
178  void ensureIdentityEncoding();
179 
180  // Unwrap `inner` as a `kj::AsyncOutputStream`.
181  kj::AsyncOutputStream& getInner();
182  // TODO(cleanup): Obviously this is polymorphism. We should be able to do better.
183 
184  // A sentinel indicating that the EncodedOutputStream has ended and is no longer usable.
185  struct Ended {};
186 
187  // I use a OneOf here rather than probing with downcasts because end() must be called for
188  // correctness rather than for optimization. I "know" this code will never be compiled w/o RTTI,
189  // but I'm paranoid.
190  kj::OneOf<kj::Own<kj::AsyncOutputStream>,
191      kj::Own<kj::GzipAsyncOutputStream>,
192      kj::Own<kj::BrotliAsyncOutputStream>,
193      Ended>
194      inner;
195 
196  StreamEncoding encoding;
197 
198  IoContext& ioContext;
199};
200 
201EncodedAsyncOutputStream::EncodedAsyncOutputStream(
202    kj::Own<kj::AsyncOutputStream> inner, StreamEncoding encoding, IoContext& context)
203    : inner(kj::mv(inner)),
204      encoding(encoding),
205      ioContext(context) {}
206 
207kj::Promise<void> EncodedAsyncOutputStream::write(kj::ArrayPtr<const byte> buffer) {
208  // Alternatively, we could throw here but this is erring on the side of leniency.
209  if (inner.is<Ended>()) return kj::READY_NOW;
210 
211  ensureIdentityEncoding();
212 
213  return getInner().write(buffer).attach(ioContext.registerPendingEvent());
214}
215 
216kj::Promise<void> EncodedAsyncOutputStream::write(
217    kj::ArrayPtr<const kj::ArrayPtr<const byte>> pieces) {
218  // Alternatively, we could throw here but this is erring on the side of leniency.
219  if (inner.is<Ended>()) return kj::READY_NOW;
220 
221  ensureIdentityEncoding();
222 
223  return getInner().write(pieces).attach(ioContext.registerPendingEvent());
224}
225 
226kj::Maybe<kj::Promise<DeferredProxy<void>>> EncodedAsyncOutputStream::tryPumpFrom(
227    ReadableStreamSource& input, bool end) {
228 
229  // If this output stream has already been ended, then there's nothing more to
230  // pump into it, just return an immediately resolved promise. Alternatively
231  // we could throw here.
232  if (inner.is<Ended>()) {
233    return kj::Promise<DeferredProxy<void>>(DeferredProxy<void>{kj::READY_NOW});
234  }
235 
236  KJ_IF_SOME(nativeInput, kj::dynamicDowncastIfAvailable<EncodedAsyncInputStream>(input)) {
237    // We can avoid putting our inner streams into identity encoding if the input and output both
238    // have the same encoding. Since ReadableStreamSource/WritableStreamSink always pump everything
239    // (there is no `amount` parameter like in the KJ equivalents), we can assume that we will
240    // always stop at a valid endpoint.
241    //
242    // Note that even if we have to pump in identity encoding, there is no reason to return nullptr.
243    // We can still optimize the pump a little by registering only a single pending event rather
244    // than falling back to the heavier weight algorithm in ReadableStreamSource, which depends on
245    // tryRead() and write() registering their own individual events on every call.
246    if (nativeInput.encoding != encoding) {
247      ensureIdentityEncoding();
248      nativeInput.ensureIdentityEncoding();
249    }
250 
251    auto promise = nativeInput.inner->pumpTo(getInner()).ignoreResult();
252    if (end) {
253      // TODO(cleanup): When KJ streams are refactored to have a general end(), this stupid switch
254      //   can go away.
255      KJ_SWITCH_ONEOF(inner) {
256        KJ_CASE_ONEOF(stream, kj::Own<kj::AsyncOutputStream>) {
257          KJ_IF_SOME(ee, kj::dynamicDowncastIfAvailable<capnp::ExplicitEndOutputStream>(*stream)) {
258            promise = promise.then([&ee = ee]() { return ee.end(); });
259          } else KJ_IF_SOME(aio, kj::dynamicDowncastIfAvailable<kj::AsyncIoStream>(*stream)) {
260            promise = promise.then([&aio = aio]() { aio.shutdownWrite(); });
261          }
262        }
263        KJ_CASE_ONEOF(gz, kj::Own<kj::GzipAsyncOutputStream>) {
264          promise = promise.then([&gz = gz]() { return gz->end(); });
265        }
266        KJ_CASE_ONEOF(br, kj::Own<kj::BrotliAsyncOutputStream>) {
267          promise = promise.then([&br = br]() { return br->end(); });
268        }
269        KJ_CASE_ONEOF(e, Ended) {}
270      }
271    }
272 
273    // Since this is a system stream, the pump task is eligible to be deferred past IoContext
274    // lifetime!
275    return kj::Promise<DeferredProxy<void>>(DeferredProxy<void>{kj::mv(promise)});
276  }
277 
278  return kj::none;
279}
280 
281StreamEncoding EncodedAsyncOutputStream::disownEncodingResponsibility() {
282  StreamEncoding result = encoding;
283  encoding = StreamEncoding::IDENTITY;
284  return result;
285}
286 
287kj::Promise<void> EncodedAsyncOutputStream::end() {
288  if (inner.is<Ended>()) return kj::READY_NOW;
289 
290  kj::Promise<void> promise = kj::READY_NOW;
291 
292  // TODO(cleanup): When KJ streams are refactored to have a general end(), this stupid switch
293  //   can go away.
294  KJ_SWITCH_ONEOF(inner) {
295    KJ_CASE_ONEOF(stream, kj::Own<kj::AsyncOutputStream>) {
296      KJ_IF_SOME(ee, kj::dynamicDowncastIfAvailable<capnp::ExplicitEndOutputStream>(*stream)) {
297        promise = ee.end().attach(kj::mv(stream));
298      } else KJ_IF_SOME(aio, kj::dynamicDowncastIfAvailable<kj::AsyncIoStream>(*stream)) {
299        aio.shutdownWrite();
300        promise = promise.attach(kj::mv(stream));
301      }
302    }
303    KJ_CASE_ONEOF(gz, kj::Own<kj::GzipAsyncOutputStream>) {
304      promise = gz->end().attach(kj::mv(gz));
305    }
306    KJ_CASE_ONEOF(br, kj::Own<kj::BrotliAsyncOutputStream>) {
307      promise = br->end().attach(kj::mv(br));
308    }
309    KJ_CASE_ONEOF(e, Ended) {}
310  }
311 
312  inner.init<Ended>();
313 
314  return promise.attach(ioContext.registerPendingEvent());
315}
316 
317void EncodedAsyncOutputStream::abort(kj::Exception reason) {
318  KJ_SWITCH_ONEOF(inner) {
319    KJ_CASE_ONEOF(stream, kj::Own<kj::AsyncOutputStream>) {
320      stream->abortWrite(kj::mv(reason));
321    }
322    KJ_CASE_ONEOF(gz, kj::Own<kj::GzipAsyncOutputStream>) {
323      gz->abortWrite(kj::mv(reason));
324    }
325    KJ_CASE_ONEOF(br, kj::Own<kj::BrotliAsyncOutputStream>) {
326      br->abortWrite(kj::mv(reason));
327    }
328    KJ_CASE_ONEOF(e, Ended) {}
329  }
330  inner.init<Ended>();
331}
332 
333void EncodedAsyncOutputStream::ensureIdentityEncoding() {
334  // Compression gets added to the stream here if needed based on the content encoding.
335  KJ_DASSERT(!inner.is<Ended>(), "the EncodedAsyncOutputStream has been ended or aborted");
336  if (encoding == StreamEncoding::GZIP) {
337    // This is safe because only a kj::AsyncOutputStream can have non-identity encoding.
338    auto& stream = inner.get<kj::Own<kj::AsyncOutputStream>>();
339 
340    inner = kj::heap<kj::GzipAsyncOutputStream>(*stream).attach(kj::mv(stream));
341    encoding = StreamEncoding::IDENTITY;
342  } else if (encoding == StreamEncoding::BROTLI) {
343    auto& stream = inner.get<kj::Own<kj::AsyncOutputStream>>();
344 
345    inner = kj::heap<kj::BrotliAsyncOutputStream>(*stream).attach(kj::mv(stream));
346    encoding = StreamEncoding::IDENTITY;
347  } else {
348    // We currently support gzip and brotli as non-identity content encodings.
349    KJ_ASSERT(encoding == StreamEncoding::IDENTITY);
350  }
351}
352 
353kj::AsyncOutputStream& EncodedAsyncOutputStream::getInner() {
354  KJ_SWITCH_ONEOF(inner) {
355    KJ_CASE_ONEOF(stream, kj::Own<kj::AsyncOutputStream>) {
356      return *stream;
357    }
358    KJ_CASE_ONEOF(gz, kj::Own<kj::GzipAsyncOutputStream>) {
359      return *gz;
360    }
361    KJ_CASE_ONEOF(br, kj::Own<kj::BrotliAsyncOutputStream>) {
362      return *br;
363    }
364    KJ_CASE_ONEOF(ended, Ended) {
365      KJ_FAIL_ASSERT("the EncodedAsyncOutputStream has been ended or aborted.");
366    }
367  }
368 
369  KJ_UNREACHABLE;
370}
371 
372}  // namespace
373 
374kj::Own<ReadableStreamSource> newSystemStream(
375    kj::Own<kj::AsyncInputStream> inner, StreamEncoding encoding, IoContext& context) {
376  return kj::heap<EncodedAsyncInputStream>(kj::mv(inner), encoding, context);
377}
378kj::Own<WritableStreamSink> newSystemStream(
379    kj::Own<kj::AsyncOutputStream> inner, StreamEncoding encoding, IoContext& context) {
380  return kj::heap<EncodedAsyncOutputStream>(kj::mv(inner), encoding, context);
381}
382 
383SystemMultiStream newSystemMultiStream(
384    kj::RefcountedWrapper<kj::Own<kj::AsyncIoStream>>& stream, IoContext& context) {
385 
386  return {.readable = kj::heap<EncodedAsyncInputStream>(
387              stream.addWrappedRef(), StreamEncoding::IDENTITY, context),
388    .writable = kj::heap<EncodedAsyncOutputStream>(
389        stream.addWrappedRef(), StreamEncoding::IDENTITY, context)};
390}
391 
392ContentEncodingOptions::ContentEncodingOptions(CompatibilityFlags::Reader flags)
393    : brotliEnabled(flags.getBrotliContentEncoding()) {}
394 
395StreamEncoding getContentEncoding(IoContext& context,
396    const kj::HttpHeaders& headers,
397    Response::BodyEncoding bodyEncoding,
398    ContentEncodingOptions options) {
399  if (bodyEncoding == Response::BodyEncoding::MANUAL) {
400    return StreamEncoding::IDENTITY;
401  }
402  KJ_IF_SOME(encodingStr, headers.get(context.getHeaderIds().contentEncoding)) {
403    if (encodingStr == "gzip") {
404      return StreamEncoding::GZIP;
405    } else if (options.brotliEnabled && encodingStr == "br") {
406      return StreamEncoding::BROTLI;
407    }
408  }
409  return StreamEncoding::IDENTITY;
410}
411 
412}  // namespace workerd::api