runtime.cpp

/*
 * Copyright 2018-present Amazon.com, Inc. or its affiliates. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License").
 * You may not use this file except in compliance with the License.
 * A copy of the License is located at
 *
 *  http://aws.amazon.com/apache2.0
 *
 * or in the "license" file accompanying this file. This file is distributed
 * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
 * express or implied. See the License for the specific language governing
 * permissions and limitations under the License.
 */

#include "aws/lambda-runtime/runtime.h"
#include "aws/lambda-runtime/version.h"
#include "aws/lambda-runtime/outcome.h"
#include "aws/logging/logging.h"
#include "aws/http/response.h"

#include <curl/curl.h>
#include <curl/curlver.h>
#include <climits> // for ULONG_MAX
#include <cassert>
#include <chrono>
#include <array>
#include <cstdlib> // for strtoul

#define AWS_LAMBDA_RUNTIME_API __attribute__((visibility("default")))

namespace aws {
namespace lambda_runtime {

static char const LOG_TAG[] = "LAMBDA_RUNTIME";
static char const REQUEST_ID_HEADER[] = "lambda-runtime-aws-request-id";
static char const TRACE_ID_HEADER[] = "lambda-runtime-trace-id";
static char const CLIENT_CONTEXT_HEADER[] = "lambda-runtime-client-context";
static char const COGNITO_IDENTITY_HEADER[] = "lambda-runtime-cognito-identity";
static char const DEADLINE_MS_HEADER[] = "lambda-runtime-deadline-ms";
static char const FUNCTION_ARN_HEADER[] = "lambda-runtime-invoked-function-arn";

enum Endpoints {
    INIT,
    NEXT,
    RESULT,
};

static bool is_success(aws::http::response_code httpcode)
{
    auto const code = static_cast<int>(httpcode);
    return code >= 200 && code <= 299;
}

static size_t write_data(char* ptr, size_t size, size_t nmemb, void* userdata)
{
    if (!ptr) {
        return 0;
    }

    auto const resp = static_cast<http::response*>(userdata);
    assert(size == 1);
    (void)size; // avoid warning in release builds
    assert(resp);
    resp->append_body(ptr, nmemb);
    return nmemb;
}

static inline bool IsSpace(int ch)
{
    if (ch < -1 || ch > 255) {
        return false;
    }

    return ::isspace(ch) != 0;
}

static inline std::string trim(std::string s)
{
    // trim right
    s.erase(std::find_if(s.rbegin(), s.rend(), [](int ch) { return !IsSpace(ch); }).base(), s.end());
    // trim left
    s.erase(s.begin(), std::find_if(s.begin(), s.end(), [](int ch) { return !IsSpace(ch); }));
    return s;
}

static size_t write_header(char* ptr, size_t size, size_t nmemb, void* userdata)
{
    if (!ptr) {
        return 0;
    }

    logging::log_debug(LOG_TAG, "received header: %s", std::string(ptr, nmemb).c_str());

    auto const resp = static_cast<http::response*>(userdata);
    assert(resp);
    for (size_t i = 0; i < nmemb; i++) {
        if (ptr[i] != ':') {
            continue;
        }
        std::string key{ptr, i};
        std::string value{ptr + i + 1, nmemb - i - 1};
        resp->add_header(trim(key), trim(value));
        break;
    }
    return size * nmemb;
}

static std::string const& get_user_agent_header()
{
    static std::string user_agent = std::string("User-Agent: AWS_Lambda_Cpp/") + get_version();
    return user_agent;
}

static size_t read_data(char* buffer, size_t size, size_t nitems, void* userdata)
{
    auto const limit = size * nitems;
    auto ctx = static_cast<std::pair<std::string const&, size_t>*>(userdata);
    assert(ctx);
    auto const unread = ctx->first.length() - ctx->second;
    if (0 == unread) {
        return 0;
    }

    if (unread <= limit) {
        std::copy_n(ctx->first.begin() + ctx->second, unread, buffer);
        ctx->second += unread;
        return unread;
    }

    std::copy_n(ctx->first.begin() + ctx->second, limit, buffer);
    ctx->second += limit;
    return limit;
}

#ifndef NDEBUG
static int rt_curl_debug_callback(CURL* handle, curl_infotype type, char* data, size_t size, void* userdata)
{
    (void)handle;
    (void)type;
    (void)userdata;
    std::string s(data, size);
    logging::log_debug(LOG_TAG, "CURL DBG: %s", s.c_str());
    return 0;
}
#endif

struct no_result {
};

class runtime {
public:
    using next_outcome = aws::lambda_runtime::outcome<invocation_request, aws::http::response_code>;
    using post_outcome = aws::lambda_runtime::outcome<no_result, aws::http::response_code>;

    runtime(std::string const& endpoint);
    ~runtime();

    /**
     * Ask lambda for an invocation.
     */
    next_outcome get_next();

    /**
     * Tells lambda that the function has succeeded.
     */
    post_outcome post_success(std::string const& request_id, invocation_response const& handler_response);

    /**
     * Tells lambda that the function has failed.
     */
    post_outcome post_failure(std::string const& request_id, invocation_response const& handler_response);

private:
    void set_curl_next_options();
    void set_curl_post_result_options();
    post_outcome do_post(
        std::string const& url,
        std::string const& request_id,
        invocation_response const& handler_response);

private:
    std::array<std::string const, 3> const m_endpoints;
    CURL* const m_curl_handle;
};

runtime::runtime(std::string const& endpoint)
    : m_endpoints{{endpoint + "/2018-06-01/runtime/init/error",
                   endpoint + "/2018-06-01/runtime/invocation/next",
                   endpoint + "/2018-06-01/runtime/invocation/"}},
      m_curl_handle(curl_easy_init())
{
    if (!m_curl_handle) {
        logging::log_error(LOG_TAG, "Failed to acquire curl easy handle for next.");
    }
}

runtime::~runtime()
{
    curl_easy_cleanup(m_curl_handle);
}

void runtime::set_curl_next_options()
{
    // lambda freezes the container when no further tasks are available. The freezing period could be longer than the
    // request timeout, which causes the following get_next request to fail with a timeout error.
    curl_easy_reset(m_curl_handle);
    curl_easy_setopt(m_curl_handle, CURLOPT_TIMEOUT, 0L);
    curl_easy_setopt(m_curl_handle, CURLOPT_CONNECTTIMEOUT, 1L);
    curl_easy_setopt(m_curl_handle, CURLOPT_NOSIGNAL, 1L);
    curl_easy_setopt(m_curl_handle, CURLOPT_TCP_NODELAY, 1L);
    curl_easy_setopt(m_curl_handle, CURLOPT_HTTP_VERSION, CURL_HTTP_VERSION_1_1);

    curl_easy_setopt(m_curl_handle, CURLOPT_HTTPGET, 1L);
    curl_easy_setopt(m_curl_handle, CURLOPT_URL, m_endpoints[Endpoints::NEXT].c_str());

    curl_easy_setopt(m_curl_handle, CURLOPT_WRITEFUNCTION, write_data);
    curl_easy_setopt(m_curl_handle, CURLOPT_HEADERFUNCTION, write_header);

#ifndef NDEBUG
    curl_easy_setopt(m_curl_handle, CURLOPT_VERBOSE, 1);
    curl_easy_setopt(m_curl_handle, CURLOPT_DEBUGFUNCTION, rt_curl_debug_callback);
#endif
}

void runtime::set_curl_post_result_options()
{
    curl_easy_reset(m_curl_handle);
    curl_easy_setopt(m_curl_handle, CURLOPT_TIMEOUT, 0L);
    curl_easy_setopt(m_curl_handle, CURLOPT_CONNECTTIMEOUT, 1L);
    curl_easy_setopt(m_curl_handle, CURLOPT_NOSIGNAL, 1L);
    curl_easy_setopt(m_curl_handle, CURLOPT_TCP_NODELAY, 1L);
    curl_easy_setopt(m_curl_handle, CURLOPT_HTTP_VERSION, CURL_HTTP_VERSION_1_1);

    curl_easy_setopt(m_curl_handle, CURLOPT_POST, 1L);
    curl_easy_setopt(m_curl_handle, CURLOPT_READFUNCTION, read_data);
    curl_easy_setopt(m_curl_handle, CURLOPT_WRITEFUNCTION, write_data);
    curl_easy_setopt(m_curl_handle, CURLOPT_HEADERFUNCTION, write_header);

#ifndef NDEBUG
    curl_easy_setopt(m_curl_handle, CURLOPT_VERBOSE, 1);
    curl_easy_setopt(m_curl_handle, CURLOPT_DEBUGFUNCTION, rt_curl_debug_callback);
#endif
}

runtime::next_outcome runtime::get_next()
{
    http::response resp;
    set_curl_next_options();
    curl_easy_setopt(m_curl_handle, CURLOPT_WRITEDATA, &resp);
    curl_easy_setopt(m_curl_handle, CURLOPT_HEADERDATA, &resp);

    curl_slist* headers = nullptr;
    headers = curl_slist_append(headers, get_user_agent_header().c_str());

    logging::log_debug(LOG_TAG, "Making request to %s", m_endpoints[Endpoints::NEXT].c_str());
    CURLcode curl_code = curl_easy_perform(m_curl_handle);
    logging::log_debug(LOG_TAG, "Completed request to %s", m_endpoints[Endpoints::NEXT].c_str());
    curl_slist_free_all(headers);

    if (curl_code != CURLE_OK) {
        logging::log_debug(LOG_TAG, "CURL returned error code %d - %s", curl_code, curl_easy_strerror(curl_code));
        logging::log_error(LOG_TAG, "Failed to get next invocation. No Response from endpoint");
        return aws::http::response_code::REQUEST_NOT_MADE;
    }

    {
        long resp_code;
        curl_easy_getinfo(m_curl_handle, CURLINFO_RESPONSE_CODE, &resp_code);
        resp.set_response_code(static_cast<aws::http::response_code>(resp_code));
    }

    {
        char* content_type = nullptr;
        curl_easy_getinfo(m_curl_handle, CURLINFO_CONTENT_TYPE, &content_type);
        resp.set_content_type(content_type);
    }

    if (!is_success(resp.get_response_code())) {
        logging::log_error(
            LOG_TAG,
            "Failed to get next invocation. Http Response code: %d",
            static_cast<int>(resp.get_response_code()));
        return resp.get_response_code();
    }

    if (!resp.has_header(REQUEST_ID_HEADER)) {
        logging::log_error(LOG_TAG, "Failed to find header %s in response", REQUEST_ID_HEADER);
        return aws::http::response_code::REQUEST_NOT_MADE;
    }
    invocation_request req;
    req.payload = resp.get_body();
    req.request_id = resp.get_header(REQUEST_ID_HEADER);

    if (resp.has_header(TRACE_ID_HEADER)) {
        req.xray_trace_id = resp.get_header(TRACE_ID_HEADER);
    }

    if (resp.has_header(CLIENT_CONTEXT_HEADER)) {
        req.client_context = resp.get_header(CLIENT_CONTEXT_HEADER);
    }

    if (resp.has_header(COGNITO_IDENTITY_HEADER)) {
        req.cognito_identity = resp.get_header(COGNITO_IDENTITY_HEADER);
    }

    if (resp.has_header(FUNCTION_ARN_HEADER)) {
        req.function_arn = resp.get_header(FUNCTION_ARN_HEADER);
    }

    if (resp.has_header(DEADLINE_MS_HEADER)) {
        auto const& deadline_string = resp.get_header(DEADLINE_MS_HEADER);
        unsigned long ms = strtoul(deadline_string.c_str(), nullptr, 10);
        assert(ms > 0);
        assert(ms < ULONG_MAX);
        req.deadline += std::chrono::milliseconds(ms);
        logging::log_info(
            LOG_TAG,
            "Received payload: %s\nTime remaining: %ld",
            req.payload.c_str(),
            req.get_time_remaining().count());
    }
    return next_outcome(req);
}

runtime::post_outcome runtime::post_success(std::string const& request_id, invocation_response const& handler_response)
{
    std::string const url = m_endpoints[Endpoints::RESULT] + request_id + "/response";
    return do_post(url, request_id, handler_response);
}

runtime::post_outcome runtime::post_failure(std::string const& request_id, invocation_response const& handler_response)
{
    std::string const url = m_endpoints[Endpoints::RESULT] + request_id + "/error";
    return do_post(url, request_id, handler_response);
}

runtime::post_outcome runtime::do_post(
    std::string const& url,
    std::string const& request_id,
    invocation_response const& handler_response)
{
    set_curl_post_result_options();
    curl_easy_setopt(m_curl_handle, CURLOPT_URL, url.c_str());
    logging::log_info(LOG_TAG, "Making request to %s", url.c_str());

    curl_slist* headers = nullptr;
    if (handler_response.get_content_type().empty()) {
        headers = curl_slist_append(headers, "content-type: text/html");
    }
    else {
        headers = curl_slist_append(headers, ("content-type: " + handler_response.get_content_type()).c_str());
    }

    headers = curl_slist_append(headers, "Expect:");
    headers = curl_slist_append(headers, "transfer-encoding:");
    headers = curl_slist_append(headers, get_user_agent_header().c_str());
    auto const& payload = handler_response.get_payload();
    logging::log_debug(
        LOG_TAG, "calculating content length... %s", ("content-length: " + std::to_string(payload.length())).c_str());
    headers = curl_slist_append(headers, ("content-length: " + std::to_string(payload.length())).c_str());

    std::pair<std::string const&, size_t> ctx{payload, 0};
    aws::http::response resp;
    curl_easy_setopt(m_curl_handle, CURLOPT_WRITEDATA, &resp);
    curl_easy_setopt(m_curl_handle, CURLOPT_HEADERDATA, &resp);
    curl_easy_setopt(m_curl_handle, CURLOPT_READDATA, &ctx);
    curl_easy_setopt(m_curl_handle, CURLOPT_HTTPHEADER, headers);
    CURLcode curl_code = curl_easy_perform(m_curl_handle);
    curl_slist_free_all(headers);

    if (curl_code != CURLE_OK) {
        logging::log_debug(
            LOG_TAG,
            "CURL returned error code %d - %s, for invocation %s",
            curl_code,
            curl_easy_strerror(curl_code),
            request_id.c_str());
        return aws::http::response_code::REQUEST_NOT_MADE;
    }

    long http_response_code;
    curl_easy_getinfo(m_curl_handle, CURLINFO_RESPONSE_CODE, &http_response_code);

    if (!is_success(aws::http::response_code(http_response_code))) {
        logging::log_error(
            LOG_TAG, "Failed to post handler success response. Http response code: %ld.", http_response_code);
        return aws::http::response_code(http_response_code);
    }

    return post_outcome(no_result{});
}

static bool handle_post_outcome(runtime::post_outcome const& o, std::string const& request_id)
{
    if (o.is_success()) {
        return true;
    }

    if (o.get_failure() == aws::http::response_code::REQUEST_NOT_MADE) {
        logging::log_error(LOG_TAG, "Failed to send HTTP request for invocation %s.", request_id.c_str());
        return false;
    }

    logging::log_info(
        LOG_TAG,
        "HTTP Request for invocation %s was not successful. HTTP response code: %d.",
        request_id.c_str(),
        static_cast<int>(o.get_failure()));
    return false;
}

AWS_LAMBDA_RUNTIME_API
void run_handler(std::function<invocation_response(invocation_request const&)> const& handler)
{
    logging::log_info(LOG_TAG, "Initializing the C++ Lambda Runtime.");
    std::string endpoint("http://");
    if (auto ep = std::getenv("AWS_LAMBDA_RUNTIME_API")) {
        assert(ep);
        logging::log_debug(LOG_TAG, "LAMBDA_SERVER_ADDRESS defined in environment as: %s", ep);
        endpoint += ep;
    }

    runtime rt(endpoint);

    size_t retries = 0;
    size_t const max_retries = 3;

    while (retries < max_retries) {
        const auto next_outcome = rt.get_next();
        if (!next_outcome.is_success()) {
            if (next_outcome.get_failure() == aws::http::response_code::REQUEST_NOT_MADE) {
                ++retries;
                continue;
            }

            logging::log_info(
                LOG_TAG,
                "HTTP request was not successful. HTTP response code: %d. Retrying..",
                static_cast<int>(next_outcome.get_failure()));
            ++retries;
            continue;
        }

        retries = 0;

        auto const& req = next_outcome.get_result();
        logging::log_info(LOG_TAG, "Invoking user handler");
        invocation_response res = handler(req);
        logging::log_info(LOG_TAG, "Invoking user handler completed.");

        if (res.is_success()) {
            const auto post_outcome = rt.post_success(req.request_id, res);
            if (!handle_post_outcome(post_outcome, req.request_id)) {
                return; // TODO: implement a better retry strategy
            }
        }
        else {
            const auto post_outcome = rt.post_failure(req.request_id, res);
            if (!handle_post_outcome(post_outcome, req.request_id)) {
                return; // TODO: implement a better retry strategy
            }
        }
    }

    if (retries == max_retries) {
        logging::log_error(
            LOG_TAG, "Exhausted all retries. This is probably a bug in libcurl v" LIBCURL_VERSION " Exiting!");
    }
}

static std::string json_escape(std::string const& in)
{
    std::string out;
    out.reserve(in.length()); // most strings will end up identical
    for (char ch : in) {
        if (ch > 31 && ch != '\"' && ch != '\\') {
            out.append(1, ch);
        }
        else {
            out.append(1, '\\');
            switch (ch) {
                case '\\':
                    out.append(1, '\\');
                    break;
                case '"':
                    out.append(1, '"');
                    break;
                case '\b':
                    out.append(1, 'b');
                    break;
                case '\f':
                    out.append(1, 'f');
                    break;
                case '\n':
                    out.append(1, 'n');
                    break;
                case '\r':
                    out.append(1, 'r');
                    break;
                case '\t':
                    out.append(1, 't');
                    break;
                default:
                    // escape and print as unicode codepoint
                    char buf[6]; // 4 hex + letter 'u' + \0
                    sprintf(buf, "u%04x", ch);
                    out.append(buf, 5); // add only five, discarding the null terminator.
                    break;
            }
        }
    }
    return out;
}

AWS_LAMBDA_RUNTIME_API
invocation_response invocation_response::success(std::string const& payload, std::string const& content_type)
{
    invocation_response r;
    r.m_success = true;
    r.m_content_type = content_type;
    r.m_payload = payload;
    return r;
}

AWS_LAMBDA_RUNTIME_API
invocation_response invocation_response::failure(std::string const& error_message, std::string const& error_type)
{
    invocation_response r;
    r.m_success = false;
    r.m_content_type = "application/json";
    r.m_payload = R"({"errorMessage":")" + json_escape(error_message) + R"(","errorType":")" + json_escape(error_type) +
                  R"(", "stackTrace":[]})";
    return r;
}

} // namespace lambda_runtime
} // namespace aws