#pragma once

#include <atomic>

#include <map>

#include <memory>

#include "envoy/access_log/access_log.h"

#include "envoy/buffer/buffer.h"

#include "envoy/extensions/wasm/v3/wasm.pb.validate.h"

#include "envoy/http/filter.h"

#include "envoy/stats/sink.h"

#include "envoy/upstream/cluster_manager.h"

#include "common/common/assert.h"

#include "common/common/logger.h"

#include "extensions/common/wasm/wasm_state.h"

#include "extensions/filters/common/expr/evaluator.h"

#include "eval/public/activation.h"

#include "include/proxy-wasm/wasm.h"

namespace Envoy {

namespace Extensions {

namespace Common {

namespace Wasm {

using proxy_wasm::BufferInterface;

using proxy_wasm::CloseType;

using proxy_wasm::ContextBase;

using proxy_wasm::Pairs;

using proxy_wasm::PairsWithStringValues;

using proxy_wasm::PluginBase;

using proxy_wasm::SharedQueueDequeueToken;

using proxy_wasm::SharedQueueEnqueueToken;

using proxy_wasm::WasmBase;

using proxy_wasm::WasmBufferType;

using proxy_wasm::WasmHandleBase;

using proxy_wasm::WasmHeaderMapType;

using proxy_wasm::WasmResult;

using proxy_wasm::WasmStreamType;

using VmConfig = envoy::extensions::wasm::v3::VmConfig;

using GrpcService = envoy::config::core::v3::GrpcService;

class Wasm;

using WasmHandleBaseSharedPtr = std::shared_ptr<WasmHandleBase>;

// Opaque context object.

class StorageObject {

public:

virtual ~StorageObject() = default;

};

class Buffer : public proxy_wasm::BufferBase {

public:

Buffer() = default;

// proxy_wasm::BufferInterface

size_t size() const override;

WasmResult copyTo(WasmBase* wasm, size_t start, size_t length, uint64_t ptr_ptr,

uint64_t size_ptr) const override;

WasmResult copyFrom(size_t start, size_t length, absl::string_view data) override;

// proxy_wasm::BufferBase

void clear() override {

proxy_wasm::BufferBase::clear();

const_buffer_instance_ = nullptr;

buffer_instance_ = nullptr;

}

Buffer* set(absl::string_view data) {

return static_cast<Buffer*>(proxy_wasm::BufferBase::set(data));

}

Buffer* set(std::unique_ptr<char[]> owned_data, uint32_t owned_data_size) {

return static_cast<Buffer*>(

proxy_wasm::BufferBase::set(std::move(owned_data), owned_data_size));

}

Buffer* set(::Envoy::Buffer::Instance* buffer_instance) {

clear();

buffer_instance_ = buffer_instance;

const_buffer_instance_ = buffer_instance;

return this;

}

Buffer* set(const ::Envoy::Buffer::Instance* buffer_instance) {

clear();

const_buffer_instance_ = buffer_instance;

return this;

}

private:

const ::Envoy::Buffer::Instance* const_buffer_instance_{};

::Envoy::Buffer::Instance* buffer_instance_{};

};

// Plugin contains the information for a filter/service.

struct Plugin : public PluginBase {

Plugin(absl::string_view name, absl::string_view root_id, absl::string_view vm_id,

absl::string_view runtime, absl::string_view plugin_configuration, bool fail_open,

envoy::config::core::v3::TrafficDirection direction,

const LocalInfo::LocalInfo& local_info,

const envoy::config::core::v3::Metadata* listener_metadata)

: PluginBase(name, root_id, vm_id, runtime, plugin_configuration, fail_open),

direction_(direction), local_info_(local_info), listener_metadata_(listener_metadata) {}

envoy::config::core::v3::TrafficDirection direction_;

const LocalInfo::LocalInfo& local_info_;

const envoy::config::core::v3::Metadata* listener_metadata_;

};

using PluginSharedPtr = std::shared_ptr<Plugin>;

// A context which will be the target of callbacks for a particular session

// e.g. a handler of a stream.

class Context : public proxy_wasm::ContextBase,

public Logger::Loggable<Logger::Id::wasm>,

public AccessLog::Instance,

public Http::StreamFilter,

public Network::ConnectionCallbacks,

public Network::Filter,

public google::api::expr::runtime::BaseActivation,

public std::enable_shared_from_this<Context> {

public:

Context(); // Testing.

Context(Wasm* wasm); // Vm Context.

Context(Wasm* wasm, const PluginSharedPtr& plugin); // Root Context.

Context(Wasm* wasm, uint32_t root_context_id, const PluginSharedPtr& plugin); // Stream context.

~Context() override;

Wasm* wasm() const;

Plugin* plugin() const;

Context* rootContext() const;

Upstream::ClusterManager& clusterManager() const;

// proxy_wasm::ContextBase

void error(absl::string_view message) override;

// Retrieves the stream info associated with the request (a.k.a active stream).

// It selects a value based on the following order: encoder callback, decoder

// callback, log callback, network read filter callback, network write filter

// callback. As long as any one of the callbacks is invoked, the value should be

// available.

const StreamInfo::StreamInfo* getConstRequestStreamInfo() const;

StreamInfo::StreamInfo* getRequestStreamInfo() const;

// Retrieves the connection object associated with the request (a.k.a active stream).

// It selects a value based on the following order: encoder callback, decoder

// callback. As long as any one of the callbacks is invoked, the value should be

// available.

const Network::Connection* getConnection() const;

//

// VM level down-calls into the Wasm code on Context(id == 0).

//

virtual bool validateConfiguration(absl::string_view configuration,

const std::shared_ptr<PluginBase>& plugin); // deprecated

// AccessLog::Instance

void log(const Http::RequestHeaderMap* request_headers,

const Http::ResponseHeaderMap* response_headers,

const Http::ResponseTrailerMap* response_trailers,

const StreamInfo::StreamInfo& stream_info) override;

uint32_t getLogLevel() override;

// Network::ConnectionCallbacks

void onEvent(Network::ConnectionEvent event) override;

void onAboveWriteBufferHighWatermark() override {}

void onBelowWriteBufferLowWatermark() override {}

// Network::ReadFilter

Network::FilterStatus onNewConnection() override;

Network::FilterStatus onData(::Envoy::Buffer::Instance& data, bool end_stream) override;

void initializeReadFilterCallbacks(Network::ReadFilterCallbacks& callbacks) override;

// Network::WriteFilter

Network::FilterStatus onWrite(::Envoy::Buffer::Instance& data, bool end_stream) override;

void initializeWriteFilterCallbacks(Network::WriteFilterCallbacks& callbacks) override;

// proxy_wasm::ContextBase

void onDownstreamConnectionClose(CloseType) override;

void onUpstreamConnectionClose(CloseType) override;

// Http::StreamFilterBase. Note: This calls onDone() in Wasm.

void onDestroy() override;

// Http::StreamDecoderFilter

Http::FilterHeadersStatus decodeHeaders(Http::RequestHeaderMap& headers,

bool end_stream) override;

Http::FilterDataStatus decodeData(::Envoy::Buffer::Instance& data, bool end_stream) override;

Http::FilterTrailersStatus decodeTrailers(Http::RequestTrailerMap& trailers) override;

Http::FilterMetadataStatus decodeMetadata(Http::MetadataMap& metadata_map) override;

void setDecoderFilterCallbacks(Envoy::Http::StreamDecoderFilterCallbacks& callbacks) override;

// Http::StreamEncoderFilter

Http::FilterHeadersStatus encode100ContinueHeaders(Http::ResponseHeaderMap&) override;

Http::FilterHeadersStatus encodeHeaders(Http::ResponseHeaderMap& headers,

bool end_stream) override;

Http::FilterDataStatus encodeData(::Envoy::Buffer::Instance& data, bool end_stream) override;

Http::FilterTrailersStatus encodeTrailers(Http::ResponseTrailerMap& trailers) override;

Http::FilterMetadataStatus encodeMetadata(Http::MetadataMap& metadata_map) override;

void setEncoderFilterCallbacks(Envoy::Http::StreamEncoderFilterCallbacks& callbacks) override;

// VM calls out to host.

// proxy_wasm::ContextBase

// General

WasmResult log(uint32_t level, absl::string_view message) override;

uint64_t getCurrentTimeNanoseconds() override;

absl::string_view getConfiguration() override;

std::pair<uint32_t, absl::string_view> getStatus() override;

// State accessors

WasmResult getProperty(absl::string_view path, std::string* result) override;

WasmResult setProperty(absl::string_view path, absl::string_view value) override;

WasmResult declareProperty(absl::string_view path,

std::unique_ptr<const WasmStatePrototype> state_prototype);

// Continue

WasmResult continueStream(WasmStreamType stream_type) override;

WasmResult closeStream(WasmStreamType stream_type) override;

WasmResult sendLocalResponse(uint32_t response_code, absl::string_view body_text,

Pairs additional_headers, uint32_t grpc_status,

absl::string_view details) override;

void clearRouteCache() override {

if (decoder_callbacks_) {

decoder_callbacks_->clearRouteCache();

}

}

// Header/Trailer/Metadata Maps

WasmResult addHeaderMapValue(WasmHeaderMapType type, absl::string_view key,

absl::string_view value) override;

WasmResult getHeaderMapValue(WasmHeaderMapType type, absl::string_view key,

absl::string_view* value) override;

WasmResult getHeaderMapPairs(WasmHeaderMapType type, Pairs* result) override;

WasmResult setHeaderMapPairs(WasmHeaderMapType type, const Pairs& pairs) override;

WasmResult removeHeaderMapValue(WasmHeaderMapType type, absl::string_view key) override;

WasmResult replaceHeaderMapValue(WasmHeaderMapType type, absl::string_view key,

absl::string_view value) override;

WasmResult getHeaderMapSize(WasmHeaderMapType type, uint32_t* size) override;

// Buffer

BufferInterface* getBuffer(WasmBufferType type) override;

// TODO: use stream_type.

bool endOfStream(WasmStreamType /* stream_type */) override { return end_of_stream_; }

// HTTP

WasmResult httpCall(absl::string_view cluster, const Pairs& request_headers,

absl::string_view request_body, const Pairs& request_trailers,

int timeout_millisconds, uint32_t* token_ptr) override;

// Stats/Metrics

WasmResult defineMetric(uint32_t type, absl::string_view name, uint32_t* metric_id_ptr) override;

WasmResult incrementMetric(uint32_t metric_id, int64_t offset) override;

WasmResult recordMetric(uint32_t metric_id, uint64_t value) override;

WasmResult getMetric(uint32_t metric_id, uint64_t* value_ptr) override;

// gRPC

WasmResult grpcCall(absl::string_view grpc_service, absl::string_view service_name,

absl::string_view method_name, const Pairs& initial_metadata,

absl::string_view request, std::chrono::milliseconds timeout,

uint32_t* token_ptr) override;

WasmResult grpcStream(absl::string_view grpc_service, absl::string_view service_name,

absl::string_view method_name, const Pairs& initial_metadat,

uint32_t* token_ptr) override;

WasmResult grpcClose(uint32_t token) override;

WasmResult grpcCancel(uint32_t token) override;

WasmResult grpcSend(uint32_t token, absl::string_view message, bool end_stream) override;

// Envoy specific ABI

void onResolveDns(uint32_t token, Envoy::Network::DnsResolver::ResolutionStatus status,

std::list<Envoy::Network::DnsResponse>&& response);

void onStatsUpdate(Envoy::Stats::MetricSnapshot& snapshot);

// CEL evaluation

std::vector<const google::api::expr::runtime::CelFunction*>

FindFunctionOverloads(absl::string_view) const override {

return {};

}

absl::optional<google::api::expr::runtime::CelValue>

findValue(absl::string_view name, Protobuf::Arena* arena, bool last) const;

absl::optional<google::api::expr::runtime::CelValue>

FindValue(absl::string_view name, Protobuf::Arena* arena) const override {

return findValue(name, arena, false);

}

bool IsPathUnknown(absl::string_view) const override { return false; }

const std::vector<google::api::expr::runtime::CelAttributePattern>&

unknown_attribute_patterns() const override {

static const std::vector<google::api::expr::runtime::CelAttributePattern> empty;

return empty;

}

const Protobuf::FieldMask& unknown_paths() const override {

return Protobuf::FieldMask::default_instance();

}

// Foreign function state

virtual void setForeignData(absl::string_view data_name, std::unique_ptr<StorageObject> data) {

data_storage_[data_name] = std::move(data);

}

template <typename T> T* getForeignData(absl::string_view data_name) {

const auto& it = data_storage_.find(data_name);

if (it == data_storage_.end()) {

return nullptr;

}

return dynamic_cast<T*>(it->second.get());

}

uint32_t nextGrpcCallToken();

uint32_t nextGrpcStreamToken();

uint32_t nextHttpCallToken();

void setNextGrpcTokenForTesting(uint32_t token) { next_grpc_token_ = token; }

void setNextHttpCallTokenForTesting(uint32_t token) { next_http_call_token_ = token; }

protected:

friend class Wasm;

void addAfterVmCallAction(std::function<void()> f);

void onCloseTCP();

struct AsyncClientHandler : public Http::AsyncClient::Callbacks {

// Http::AsyncClient::Callbacks

void onSuccess(const Http::AsyncClient::Request&,

Envoy::Http::ResponseMessagePtr&& response) override {

context_->onHttpCallSuccess(token_, std::move(response));

}

void onFailure(const Http::AsyncClient::Request&,

Http::AsyncClient::FailureReason reason) override {

context_->onHttpCallFailure(token_, reason);

}

void

onBeforeFinalizeUpstreamSpan(Envoy::Tracing::Span& /* span */,

const Http::ResponseHeaderMap* /* response_headers */) override {}

Context* context_;

uint32_t token_;

Http::AsyncClient::Request* request_;

};

struct GrpcCallClientHandler : public Grpc::RawAsyncRequestCallbacks {

// Grpc::AsyncRequestCallbacks

void onCreateInitialMetadata(Http::RequestHeaderMap& initial_metadata) override {

context_->onGrpcCreateInitialMetadata(token_, initial_metadata);

}

void onSuccessRaw(::Envoy::Buffer::InstancePtr&& response, Tracing::Span& /* span */) override {

context_->onGrpcReceiveWrapper(token_, std::move(response));

}

void onFailure(Grpc::Status::GrpcStatus status, const std::string& message,

Tracing::Span& /* span */) override {

context_->onGrpcCloseWrapper(token_, status, message);

}

Context* context_;

uint32_t token_;

Grpc::RawAsyncClientPtr client_;

Grpc::AsyncRequest* request_;

};

struct GrpcStreamClientHandler : public Grpc::RawAsyncStreamCallbacks {

// Grpc::AsyncStreamCallbacks

void onCreateInitialMetadata(Http::RequestHeaderMap&) override {}

void onReceiveInitialMetadata(Http::ResponseHeaderMapPtr&& metadata) override {

context_->onGrpcReceiveInitialMetadataWrapper(token_, std::move(metadata));

}

bool onReceiveMessageRaw(::Envoy::Buffer::InstancePtr&& response) override {

context_->onGrpcReceiveWrapper(token_, std::move(response));

return true;

}

void onReceiveTrailingMetadata(Http::ResponseTrailerMapPtr&& metadata) override {

context_->onGrpcReceiveTrailingMetadataWrapper(token_, std::move(metadata));

}

void onRemoteClose(Grpc::Status::GrpcStatus status, const std::string& message) override {

remote_closed_ = true;

context_->onGrpcCloseWrapper(token_, status, message);

}

Context* context_;

uint32_t token_;

Grpc::RawAsyncClientPtr client_;

Grpc::RawAsyncStream* stream_;

bool local_closed_ = false;

bool remote_closed_ = false;

};

void onHttpCallSuccess(uint32_t token, Envoy::Http::ResponseMessagePtr&& response);

void onHttpCallFailure(uint32_t token, Http::AsyncClient::FailureReason reason);

void onGrpcCreateInitialMetadata(uint32_t token, Http::RequestHeaderMap& metadata);

void onGrpcReceiveInitialMetadataWrapper(uint32_t token, Http::HeaderMapPtr&& metadata);

void onGrpcReceiveWrapper(uint32_t token, ::Envoy::Buffer::InstancePtr response);

void onGrpcReceiveTrailingMetadataWrapper(uint32_t token, Http::HeaderMapPtr&& metadata);

void onGrpcCloseWrapper(uint32_t token, const Grpc::Status::GrpcStatus& status,

const absl::string_view message);

bool isGrpcStreamToken(uint32_t token) { return (token & 1) == 0; }

bool isGrpcCallToken(uint32_t token) { return (token & 1) == 1; }

Http::HeaderMap* getMap(WasmHeaderMapType type);

const Http::HeaderMap* getConstMap(WasmHeaderMapType type);

const LocalInfo::LocalInfo* root_local_info_{nullptr}; // set only for root_context.

uint32_t next_http_call_token_ = 1;

uint32_t next_grpc_token_ = 1; // Odd tokens are for Calls even for Streams.

// Network callbacks.

Network::ReadFilterCallbacks* network_read_filter_callbacks_{};

Network::WriteFilterCallbacks* network_write_filter_callbacks_{};

// HTTP callbacks.

Envoy::Http::StreamDecoderFilterCallbacks* decoder_callbacks_{};

Envoy::Http::StreamEncoderFilterCallbacks* encoder_callbacks_{};

// Status.

uint32_t status_code_{0};

absl::string_view status_message_;

// Network filter state.

::Envoy::Buffer::Instance* network_downstream_data_buffer_{};

::Envoy::Buffer::Instance* network_upstream_data_buffer_{};

// HTTP filter state.

bool http_request_started_ = false; // When decodeHeaders() is called the request is "started".

Http::RequestHeaderMap* request_headers_{};

Http::ResponseHeaderMap* response_headers_{};

::Envoy::Buffer::Instance* request_body_buffer_{};

::Envoy::Buffer::Instance* response_body_buffer_{};

Http::RequestTrailerMap* request_trailers_{};

Http::ResponseTrailerMap* response_trailers_{};

Http::MetadataMap* request_metadata_{};

Http::MetadataMap* response_metadata_{};

// Only available during onHttpCallResponse.

Envoy::Http::ResponseMessagePtr* http_call_response_{};

Http::HeaderMapPtr grpc_receive_initial_metadata_{};

Http::HeaderMapPtr grpc_receive_trailing_metadata_{};

// Only available (non-nullptr) during onGrpcReceive.

::Envoy::Buffer::InstancePtr grpc_receive_buffer_;

// Only available (non-nullptr) during grpcCall and grpcStream.

Http::RequestHeaderMapPtr grpc_initial_metadata_;

// Access log state.

const StreamInfo::StreamInfo* access_log_stream_info_{};

const Http::RequestHeaderMap* access_log_request_headers_{};

const Http::ResponseHeaderMap* access_log_response_headers_{};

const Http::ResponseTrailerMap* access_log_response_trailers_{};

// Temporary state.

ProtobufWkt::Struct temporary_metadata_;

bool end_of_stream_;

bool buffering_request_body_ = false;

bool buffering_response_body_ = false;

Buffer buffer_;

// MB: must be a node-type map as we take persistent references to the entries.

std::map<uint32_t, AsyncClientHandler> http_request_;

std::map<uint32_t, GrpcCallClientHandler> grpc_call_request_;

std::map<uint32_t, GrpcStreamClientHandler> grpc_stream_;

// Opaque state.

absl::flat_hash_map<std::string, std::unique_ptr<StorageObject>> data_storage_;

// TCP State.

bool upstream_closed_ = false;

bool downstream_closed_ = false;

bool tcp_connection_closed_ = false;

// Filter state prototype declaration.

absl::flat_hash_map<std::string, std::unique_ptr<const WasmStatePrototype>> state_prototypes_;

};

using ContextSharedPtr = std::shared_ptr<Context>;

WasmResult serializeValue(Filters::Common::Expr::CelValue value, std::string* result);

} // namespace Wasm

} // namespace Common

} // namespace Extensions

} // namespace Envoy