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
#pragma once
6
 
7
#include <kj/compat/http.h>
8
#include <kj/debug.h>
9
 
10
namespace workerd {
11
 
12
// Attaches the given object to a `Request` so that it lives as long as the request's properties.
13
// The given object must support `kj::addRef()` (e.g. `kj::Refcount`).
14
template <typename T>
15
kj::HttpClient::Request attachToRequest(kj::HttpClient::Request req, T&& rcAttachment) {
16
  req.body = req.body.attach(kj::addRef(*rcAttachment));
17
  req.response = req.response.then(
18
      [rcAttachment = kj::mv(rcAttachment)](kj::HttpClient::Response&& response) mutable {
19
    response.body = response.body.attach(kj::mv(rcAttachment));
20
    return kj::mv(response);
21
  });
22
  return req;
23
}
24
 
25
// Attaches the given object to a `WebSocketResponse` promise so that it lives as long as the
26
// returned response's properties.
27
template <typename T>
28
kj::Promise<kj::HttpClient::WebSocketResponse> attachToWebSocketResponse(
29
    kj::Promise<kj::HttpClient::WebSocketResponse> promise, T&& attachment) {
30
  return promise.then(
31
      [attachment = kj::mv(attachment)](kj::HttpClient::WebSocketResponse&& response) mutable {
32
    KJ_SWITCH_ONEOF(response.webSocketOrBody) {
33
      KJ_CASE_ONEOF(stream, kj::Own<kj::AsyncInputStream>) {
34
        response.webSocketOrBody = stream.attach(kj::mv(attachment));
35
      }
36
      KJ_CASE_ONEOF(ws, kj::Own<kj::WebSocket>) {
37
        response.webSocketOrBody = ws.attach(kj::mv(attachment));
38
      }
39
    }
40
    return kj::mv(response);
41
  });
42
}
43
 
44
// A Response kj::HttpService::Response implementation that records the status
45
// code on the response.
46
class SimpleResponseObserver final: public kj::HttpService::Response {
47
 public:
48
  SimpleResponseObserver(kj::uint* statusCode, kj::HttpService::Response& response)
49
      : inner(response),
50
        statusCode(statusCode) {}
51
  KJ_DISALLOW_COPY_AND_MOVE(SimpleResponseObserver);
52
 
53
  kj::Own<kj::AsyncOutputStream> send(kj::uint status,
54
      kj::StringPtr statusText,
55
      const kj::HttpHeaders& headers,
56
      kj::Maybe<uint64_t> expectedBodySize) override {
57
    *statusCode = status;
58
    return inner.send(status, statusText, headers, expectedBodySize);
59
  }
60
 
61
  kj::Own<kj::WebSocket> acceptWebSocket(const kj::HttpHeaders& headers) override {
62
    return inner.acceptWebSocket(headers);
63
  }
64
 
65
 private:
66
  kj::HttpService::Response& inner;
67
  kj::uint* statusCode;
68
};
69
 
70
}  // namespace workerd