ResponseBuilder.h

Go to the documentation of this file.

/*
 *  Copyright (c) 2015-present, Facebook, Inc.
 *  All rights reserved.
 *
 *  This source code is licensed under the BSD-style license found in the
 *  LICENSE file in the root directory of this source tree. An additional grant
 *  of patent rights can be found in the PATENTS file in the same directory.
 *
 */
#pragma once

#include <folly/ScopeGuard.h>
#include <proxygen/httpserver/ResponseHandler.h>

namespace proxygen {

class ResponseBuilder {
 public:
  explicit ResponseBuilder(ResponseHandler* txn): txn_(txn) {
  }

  ResponseBuilder& promise(const std::string& url, const std::string& host) {
    headers_ = std::make_unique<HTTPMessage>();
    headers_->setHTTPVersion(1, 1);
    headers_->setURL(url);
    headers_->getHeaders().set(HTTP_HEADER_HOST, host);
    return *this;
  }

  ResponseBuilder& status(uint16_t code, const std::string& message) {
    headers_ = std::make_unique<HTTPMessage>();
    headers_->setHTTPVersion(1, 1);
    headers_->setStatusCode(code);
    headers_->setStatusMessage(message);
    return *this;
  }

  template <typename T>
  ResponseBuilder& header(const std::string& headerIn, const T& value) {
    CHECK(headers_) << "You need to call `status` before adding headers";
    headers_->getHeaders().add(headerIn, value);
    return *this;
  }

  template <typename T>
  ResponseBuilder& header(HTTPHeaderCode code, const T& value) {
    CHECK(headers_) << "You need to call `status` before adding headers";
    headers_->getHeaders().add(code, value);
    return *this;
  }

  ResponseBuilder& body(std::unique_ptr<folly::IOBuf> bodyIn) {
    if (bodyIn) {
      if (body_) {
        body_->prependChain(std::move(bodyIn));
      } else {
        body_ = std::move(bodyIn);
      }
    }

    return *this;
  }

  template <typename T>
  ResponseBuilder& body(T&& t) {
    return body(folly::IOBuf::maybeCopyBuffer(
        folly::to<std::string>(std::forward<T>(t))));
  }

  ResponseBuilder& closeConnection() {
    return header(HTTP_HEADER_CONNECTION, "close");
  }

  ResponseBuilder& trailers(const HTTPHeaders& trailers) {
    trailers_.reset(new HTTPHeaders(trailers));
    return *this;
  }

  void sendWithEOM() {
    sendEOM_ = true;
    send();
  }

  void send() {
    // Once we send them, we don't want to send them again
    SCOPE_EXIT { headers_.reset(); };

    // By default, chunked
    bool chunked = true;

    // If we have complete response, we can use Content-Length and get done
    if (headers_ && sendEOM_) {
      chunked = false;
    }

    if (headers_) {
      // We don't need to add Content-Length or Encoding for 1xx responses
      if (headers_->isResponse() && headers_->getStatusCode() >= 200) {
        if (chunked) {
          headers_->setIsChunked(true);
        } else {
          const auto len = body_ ? body_->computeChainDataLength() : 0;
          headers_->getHeaders().add(
              HTTP_HEADER_CONTENT_LENGTH,
              folly::to<std::string>(len));
        }
      }

      txn_->sendHeaders(*headers_);
    }

    if (body_) {
      if (chunked) {
        txn_->sendChunkHeader(body_->computeChainDataLength());
        txn_->sendBody(std::move(body_));
        txn_->sendChunkTerminator();
      } else {
        txn_->sendBody(std::move(body_));
      }
    }

    if (sendEOM_) {
      if (trailers_) {
        auto txn = txn_->getTransaction();
        if (txn) {
          txn->sendTrailers(*trailers_);
        }
        trailers_.reset();
      }
      txn_->sendEOM();
    }
  }

  enum class UpgradeType {
    CONNECT_REQUEST = 0,
    HTTP_UPGRADE,
  };

  void acceptUpgradeRequest(UpgradeType upgradeType,
                            const std::string upgradeProtocol = "") {
    headers_ = std::make_unique<HTTPMessage>();
    if (upgradeType == UpgradeType::CONNECT_REQUEST) {
      headers_->constructDirectResponse({1, 1}, 200, "OK");
    } else {
      CHECK(!upgradeProtocol.empty());
      headers_->constructDirectResponse({1, 1}, 101, "Switching Protocols");
      headers_->getHeaders().add(HTTP_HEADER_UPGRADE, upgradeProtocol);
      headers_->getHeaders().add(HTTP_HEADER_CONNECTION, "Upgrade");
    }
    txn_->sendHeaders(*headers_);
  }

  void rejectUpgradeRequest() {
    headers_ = std::make_unique<HTTPMessage>();
    headers_->constructDirectResponse({1, 1}, 400, "Bad Request");
    txn_->sendHeaders(*headers_);
    txn_->sendEOM();
  }

  ResponseBuilder& setEgressWebsocketHeaders() {
    headers_->setEgressWebsocketUpgrade();
    return *this;
  }

 private:
  ResponseHandler* const txn_{nullptr};

  std::unique_ptr<HTTPMessage> headers_;
  std::unique_ptr<folly::IOBuf> body_;
  std::unique_ptr<HTTPHeaders> trailers_;

  // If true, sends EOM.
  bool sendEOM_{false};
};

}