#include "module.h"

#include <yweb/webdaemons/icookiedaemon/icookie_lib/icookie.h>
#include <yweb/webdaemons/icookiedaemon/icookie_lib/process.h>
#include <yweb/webdaemons/icookiedaemon/icookie_lib/keys.h>
#include <yweb/webdaemons/icookiedaemon/icookie_lib/utils/sae_cookie/sae.h>
#include <yweb/webdaemons/icookiedaemon/icookie_lib/blacklist/default/blacklist_default.h>

#include <quality/ab_testing/usersplit_lib/hash_calc/hash_calc.h>

#include <balancer/kernel/client_hints/impl/client_hints.h>
#include <balancer/kernel/cookie/cookie.h>
#include <balancer/kernel/cookie/gdpr/cookie_type.h>
#include <balancer/kernel/cookie/gdpr/lifetime.h>
#include <balancer/kernel/custom_io/stream.h>
#include <balancer/kernel/fs/shared_file_exists_checker.h>
#include <balancer/kernel/helpers/default_instance.h>
#include <balancer/kernel/http/parser/common_headers.h>
#include <balancer/kernel/http/parser/header_validation.h>
#include <balancer/kernel/http/parser/http.h>
#include <balancer/kernel/memory/chunks.h>
#include <balancer/kernel/module/module.h>
#include <balancer/kernel/regexp/regexp_pire.h>
#include <balancer/kernel/ssl/sslio.h>

#include <library/cpp/string_utils/quote/quote.h>
#include <library/cpp/cgiparam/cgiparam.h>

#include <util/generic/ptr.h>
#include <util/generic/strbuf.h>
#include <util/generic/string.h>
#include <util/string/cast.h>
#include <util/string/split.h>


using namespace NSrvKernel;
using namespace NRegExp;

namespace NModIcookie {
    struct TProcessBufs {
        TString HostBuf;
        TString RandomUidBuf;
        TString Uuid;
        TString Transport;
        TCgiParameters Cgi;
        TString ExtUidDecryptedFromParent;
    };

    class TUserAgentFsm : public TFsm, public TWithDefaultInstance<TUserAgentFsm> {
    public:
        TUserAgentFsm() noexcept
            : TFsm("User-Agent", TFsm::TOptions().SetCaseInsensitive(true))
        {}
    };

    class TSearchappUserAgentFsm : public TFsm, public TWithDefaultInstance<TSearchappUserAgentFsm> {
    public:
        TSearchappUserAgentFsm() noexcept
            : TFsm(NIcookie::SEARCHAPP_USERAGENT_REGEX, TFsm::TOptions().SetSurround(true))
        {}
    };

    static const TStringBuf SEARCHAPP_UUID_REGEX = "[a-f0-9]{32}";

    class TSearchappUuidFsm : public TFsm, public TWithDefaultInstance<TSearchappUuidFsm> {
    public:
        TSearchappUuidFsm() noexcept
            : TFsm(SEARCHAPP_UUID_REGEX)
        {}
    };

    class TICookieFsm : public TFsm, public TWithDefaultInstance<TICookieFsm> {
    public:
        TICookieFsm() noexcept
            : TFsm(TFsm("i") | TFsm("yandexuid") | TFsm("yandex_login") | TFsm("sae"))
        {}

        static const size_t I = 0;
        static const size_t YandexUid = 1;
        static const size_t YandexLogin = 2;
        static const size_t Sae = 3;
    };

    namespace NExpPrivate {
        static constexpr const size_t SLOTS_COUNT = 100;
        static constexpr const size_t BUCKETS_COUNT = 100;

        bool IsValidTestid(TStringBuf testidStr) {
            ui32 testid;
            return TryFromString(testidStr, testid);
        }

        THashSet<ui32> ParseSlots(TStringBuf slotsStr) {
            if (slotsStr.empty()) {
                ythrow TConfigParseError() << " slots are not specified";
            }

            THashSet<ui32> slots;
            auto slotsArrayOfStrings = StringSplitter(slotsStr).Split(',');
            for (const auto& slotStr: slotsArrayOfStrings) {
                ui32 slot;
                if (!TryFromString(slotStr, slot)) {
                    ythrow TConfigParseError() << " cannot parse slots from '" << slotsStr << "'";
                }
                if (slot >= SLOTS_COUNT) {
                    ythrow TConfigParseError() << " slot should be in range [0;" << SLOTS_COUNT << ")";
                } else if (slots.contains(slot)) {
                    ythrow TConfigParseError() << " found duplicate for slot [" << slot << "]";
                }
                slots.insert(slot);
            }

            return slots;
        }
    }

Y_TLS(icookie) {
    bool KillSwitchFileExists() const noexcept {
        return KillSwitchChecker.Exists();
    }

    THolder<NIcookie::TIcookieEncrypter> Encrypter;
    TSharedFileExistsChecker KillSwitchChecker;
};

MODULE_WITH_TLS_BASE(icookie, TModuleWithSubModule) {
public:
    TModule(const TModuleParams& params)
        : TModuleBase(params)
    {
        Y_ENSURE(GdprCookieType_ == NGdprCookie::ECookieType::Tech,
            "i cookie is no more gdpr-safe, must implement proper gdpr check");

        GetCachedUATraits();

        InitSsl();

        Config->ForEach(this);

        if (KeysFilename_ && UseDefaultKeys_) {
            ythrow TConfigParseError() << " can't use 'keys_file' and 'use_default_keys' simultaneously";
        }
        if (!KeysFilename_ && !UseDefaultKeys_) {
            ythrow TConfigParseError() << " must specify either 'keys_file' or 'use_default_keys'";
        }
        if (!KeysSet_) {
            ythrow TConfigParseError() << " keys not initialized";
        }
        if (KeysSet_->GetKeys().empty()) {
            ythrow TConfigParseError() << " keys are empty";
        }

        if (Domains_.empty()) {
            ythrow TConfigParseError() << " 'domains' option is empty or not specified";
        }

        if (!CheckHeaderName(DecryptedExtUidHeader_)) {
            ythrow TConfigParseError() << " bad value for 'decrypted_ext_uid_header'";
        } else if (!CheckRestrictedHeaderName(DecryptedExtUidHeader_)) {
            ythrow TConfigParseError{} << "\"decrypted_ext_uid_header\" value " <<  DecryptedExtUidHeader_.Quote()
                << " contains one of the restricted headers " << RestrictedHeadersListString() << "\n";
        }

        if (!CheckHeaderName(DecryptedUidHeader_)) {
            ythrow TConfigParseError() << " bad value for 'decrypted_uid_header'";
        } else if (!CheckRestrictedHeaderName(DecryptedUidHeader_)) {
            ythrow TConfigParseError{} << "\"decrypted_uid_header\" value " <<  DecryptedUidHeader_.Quote()
                << " contains one of the restricted headers " << RestrictedHeadersListString() << "\n";
        }
        if (!CheckHeaderName(DecryptedLoginHashHeader_)) {
            ythrow TConfigParseError() << " bad value for 'decrypted_login_hash_header'";
        } else if (!CheckRestrictedHeaderName(DecryptedLoginHashHeader_)) {
            ythrow TConfigParseError{} << "\"decrypted_login_hash_header\" value " <<  DecryptedLoginHashHeader_.Quote()
                << " contains one of the restricted headers " << RestrictedHeadersListString() << "\n";
        }
        if (!CheckHeaderName(SrcICookieHeader_)) {
            ythrow TConfigParseError() << " bad value for 'src_icookie_header'";
        }
        if (!CheckHeaderName(ErrorHeader_)) {
            ythrow TConfigParseError() << " bad value for 'error_header'";
        } else if (!CheckRestrictedHeaderName(ErrorHeader_)) {
            ythrow TConfigParseError{} << "\"error_header\" value " <<  ErrorHeader_.Quote()
                << " contains one of the restricted headers " << RestrictedHeadersListString() << "\n";
        }
        if (!CheckHeaderName(InfoHeader_)) {
            ythrow TConfigParseError() << " bad value for 'info_header'";
        } else if (!CheckRestrictedHeaderName(InfoHeader_)) {
            ythrow TConfigParseError{} << "\"info_header\" value " <<  InfoHeader_.Quote()
                << " contains one of the restricted headers " << RestrictedHeadersListString() << "\n";
        }
        if (!EncryptedHeader_.empty() && !CheckHeaderName(EncryptedHeader_)) {
            ythrow TConfigParseError() << " bad value for 'encrypted_header'";
        } else if (!EncryptedHeader_.empty() && !CheckRestrictedHeaderName(EncryptedHeader_)) {
            ythrow TConfigParseError{} << "\"encrypted_header\" value " <<  EncryptedHeader_.Quote()
                << " contains one of the restricted headers " << RestrictedHeadersListString() << "\n";
        }
        if (!TakeRandomUidFrom_.empty() && !CheckHeaderName(TakeRandomUidFrom_)) {
            ythrow TConfigParseError() << " bad value for 'take_randomuid_from'";
        } else if (!TakeRandomUidFrom_.empty() && !CheckRestrictedHeaderName(TakeRandomUidFrom_)) {
            ythrow TConfigParseError{} << "\"take_randomuid_from\" value " <<  TakeRandomUidFrom_.Quote()
                << " contains one of the restricted headers " << RestrictedHeadersListString() << "\n";
        }
        if (ForceGenerateFromTransport_ && MaxTransportAge_ == 0) {
            ythrow TConfigParseError() << " non-zero 'max_transport_age' required for 'force_generate_from_transport'";
        }

        if (!ExpSalt_.empty()) {
            if (!NExpPrivate::IsValidTestid(ExpATestid_)) {
                ythrow TConfigParseError() << " bad 'A' testid specified '" << ExpATestid_ << "'";
            }

            if (!NExpPrivate::IsValidTestid(ExpBTestid_)) {
                ythrow TConfigParseError() << " bad 'B' testid specified '" << ExpBTestid_ << "'";
            }

            ExpASlots_ = NExpPrivate::ParseSlots(ExpASlotsStr_);
            ExpBSlots_ = NExpPrivate::ParseSlots(ExpBSlotsStr_);

            THashSet<ui32> slotsIntersection;
            std::set_intersection(ExpASlots_.begin(), ExpASlots_.end(),
                                  ExpBSlots_.begin(), ExpBSlots_.end(),
                                  std::inserter(slotsIntersection, slotsIntersection.begin()));

            if (!slotsIntersection.empty()) {
                ythrow TConfigParseError() << "slots of A and B testids should not be intersects";
            }
        }

        IcookieHeadersFsm_.Reset(new TFsm(
              TFsm(DecryptedUidHeader_, TFsm::TOptions().SetCaseInsensitive(true))
            | TFsm(DecryptedLoginHashHeader_, TFsm::TOptions().SetCaseInsensitive(true))
            | TFsm(ErrorHeader_, TFsm::TOptions().SetCaseInsensitive(true))
            | TFsm(InfoHeader_, TFsm::TOptions().SetCaseInsensitive(true))
            | TFsm(SrcICookieHeader_, TFsm::TOptions().SetCaseInsensitive(true))
        ));

        AllHeadersFsm_.Reset(new TFsm(
              TCookieFsm::Instance()                                                // 0
            | THostFsm::Instance()                                                  // 1
            | TUserAgentFsm::Instance()                                             // 2
            | TFsm(DecryptedUidHeader_, TFsm::TOptions().SetCaseInsensitive(true))  // 3
            | TFsm(ErrorHeader_, TFsm::TOptions().SetCaseInsensitive(true))         // 4
            | TFsm(InfoHeader_, TFsm::TOptions().SetCaseInsensitive(true))          // 5
            | TFsm(DecryptedLoginHashHeader_, TFsm::TOptions().SetCaseInsensitive(true)) // 6
            | TFsm(SrcICookieHeader_, TFsm::TOptions().SetCaseInsensitive(true))    // 7
            | TFsm(EncryptedHeader_, TFsm::TOptions().SetCaseInsensitive(true))     // 8
            | TFsm(DecryptedExtUidHeader_, TFsm::TOptions().SetCaseInsensitive(true))     // 9
        ));

        if (!TakeRandomUidFrom_.empty()) {
            *AllHeadersFsm_ = (
                  *AllHeadersFsm_
                | TFsm(TakeRandomUidFrom_, TFsm::TOptions().SetCaseInsensitive(true))   // 10
            );
        }

        Policy_.SetEnableSetcookie(EnableSetCookie_);
        Policy_.SetDomains(Domains_);
        Policy_.SetFlagSecure(FlagSecure_);
        Policy_.SetSchemeBitmask(IcookieSchemeBitmask_);
        Policy_.SetForceEqualToYandexuid(ForceEqualToYandexuid_);
        Policy_.SetForceGenerateFromSearchappUuid(ForceGenerateFromSearchappUuid_);
        Policy_.SetForceGenerateFromYandexBrowserUuid(ForceGenerateFromYandexBrowserUuid_);
        Policy_.SetForceGenerateFromTransport(ForceGenerateFromTransport_);
        Policy_.SetMaxTransportAge(MaxTransportAge_);
        Policy_.SetEnableGuessSearchapp(EnableGuessSearchapp_);
        Policy_.SetFlagSameSiteNone(true);
        Policy_.SetReencryptOnOldFlags(false);
        Policy_.SetCookieLifetime(NGdprCookie::MaxAllowedAge.Seconds());
        Policy_.SetEnableIcookieEncryptedHeader(!EncryptedHeader_.empty());

        Blacklist_ = &NIcookie::TDefaultYandexBlacklist::Instance();

        // TODO(velavokr): a temporary fix for BALANCER-3031. A more appropriate approach is described in BALANCER-3034
        NeedIdempotence_ = CheckParents([&](TStringBuf name) {
            if (name.StartsWith("balancer") || name == "shared") {
                return true;
            }
            return false;
        });
    }

private:
    START_PARSE {
        ON_KEY("file_switch", KillSwitchFile_) {
            return;
        }

        ON_KEY("trust_parent", TrustParent_) {
            return;
        }

        ON_KEY("trust_children", TrustChildren_) {
            return;
        }

        ON_KEY("keys_file", KeysFilename_) {
            LoadKeysFromFile();
            return;
        }

        ON_KEY("use_default_keys", UseDefaultKeys_) {
            LoadDefaultKeys();
            return;
        }

        ON_KEY("enable_set_cookie", EnableSetCookie_) {
            return;
        }

        ON_KEY("enable_decrypting", EnableDecrypting_) {
            return;
        }

        ON_KEY("decrypted_uid_header", DecryptedUidHeader_) {
            return;
        }

        ON_KEY("decrypted_ext_uid_header", DecryptedExtUidHeader_) {
            return;
        }

        ON_KEY("decrypted_login_hash_header", DecryptedLoginHashHeader_) {
            return;
        }

        ON_KEY("src_icookie_header", SrcICookieHeader_) {
            return;
        }

        ON_KEY("error_header", ErrorHeader_) {
            return;
        }

        ON_KEY("info_header", InfoHeader_) {
            return;
        }

        ON_KEY("encrypted_header", EncryptedHeader_) {
            return;
        }

        ON_KEY("take_randomuid_from", TakeRandomUidFrom_) {
            return;
        }

        ON_KEY("flag_secure", FlagSecure_) {
            return;
        }

        ON_KEY("scheme_bitmask", IcookieSchemeBitmask_) {
            return;
        }

        ON_KEY("force_equal_to_yandexuid", ForceEqualToYandexuid_) {
            return;
        }

        ON_KEY("force_generate_from_searchapp_uuid", ForceGenerateFromSearchappUuid_) {
            return;
        }

        ON_KEY("enable_parse_searchapp_uuid", EnableParseSearchappUuid_) {
            return;
        }

        ON_KEY("force_generate_from_yandex_browser_uuid", ForceGenerateFromYandexBrowserUuid_) {
            return;
        }

        ON_KEY("force_generate_from_transport", ForceGenerateFromTransport_) {
            return;
        }

        ON_KEY("max_transport_age", MaxTransportAge_) {
            return;
        }

        ON_KEY("exp_type", ExpType_) {
            return;
        }

        ON_KEY("exp_salt", ExpSalt_) {
            return;
        }

        ON_KEY("exp_A_testid", ExpATestid_) {
            return;
        }

        ON_KEY("exp_A_slots", ExpASlotsStr_) {
            return;
        }

        ON_KEY("exp_B_testid", ExpBTestid_) {
            return;
        }

        ON_KEY("exp_B_slots", ExpBSlotsStr_) {
            return;
        }

        ON_KEY("enable_guess_searchapp", EnableGuessSearchapp_) {
            return;
        }

        if (key == "domains") try {
            const TString domainsValue = value->AsString();
            for (const auto& i: StringSplitter(domainsValue).Split(',')) {
                const TString domain = ToString(i.Token());

                if (!NIcookie::CheckDomainName(domain)) {
                    ythrow TConfigParseError() << " bad domain name: '" << domain << '\'';
                }

                Domains_.push_back(domain);
            }
            return;
        } catch (...) {
            ythrow TConfigParseError() << "error parsing " << key.Quote() << ": " << CurrentExceptionMessage();
        }

        Submodule_.Reset(Loader->MustLoad(key, Copy(value->AsSubConfig())).Release());
        return;
    } END_PARSE

    void LoadKeysFromFile() {
        KeysSet_ = MakeHolder<NIcookie::TKeysSet>(NIcookie::ReadKeysFromFile(KeysFilename_));
    }

    void LoadDefaultKeys() {
        KeysSet_ = MakeHolder<NIcookie::TKeysSet>(NIcookie::GetDefaultYandexKeys());
    }

    class TIcookieSetter {
    public:
        TIcookieSetter(const TModule& parent) noexcept
            : Parent_(parent)
        {}

        void SetSetCookieHeader(TString&& setCookieHeader) noexcept {
            SetCookieHeader_ = std::move(setCookieHeader);
        }

        bool NeedToHandle() noexcept {
            return !Parent_.TrustChildren_ || !SetCookieHeader_.empty();
        }

        void ProcessResponse(TResponse& response) const noexcept {
            THeaders* headers = &response.Headers();
            if (!HasSetCookieAlready(headers) &&
                Parent_.EnableSetCookie_ &&
                !SetCookieHeader_.empty())
            {
                headers->Add("Set-Cookie", SetCookieHeader_);
            }
        }

    private:
        bool HasSetCookieAlready(THeaders* headers) const noexcept {
            bool hasSetCookie = false;
            static const TString prefix = NIcookie::ICOOKIE_FIELD + TString("=");
            EraseNodesIf(*headers, [&](auto& header) {
                if (Match(TSetCookieFsm::Instance(), header.first.AsStringBuf())) {
                    auto& headerValues = header.second;

                    // remove headerValues that start with prefix.
                    auto newEnd = std::remove_if(headerValues.begin(), headerValues.end(), [&](auto& headerValue) {
                        if (headerValue.AsStringBuf().StartsWith(prefix)) {
                            if (Parent_.TrustChildren_) {
                                hasSetCookie = true;
                                return false;
                            } else {
                                return true;
                            }
                        }
                        return false;
                    });

                    headerValues.erase(newEnd, headerValues.end());
                }

                if (header.second.empty()) {
                    return true; // if there's no headervalues left, then delete header from THeaders.
                } else {
                    return false;
                }
            });

            return hasSetCookie;
        }

    private:
        const TModule& Parent_;
        TString SetCookieHeader_;
    };

    THolder<TTls> DoInitTls(IWorkerCtl* process) override {
        auto tls = MakeHolder<TTls>();
        tls->Encrypter.Reset(MakeHolder<NIcookie::TIcookieEncrypter>(*KeysSet_));
        if (KillSwitchFile_) {
            tls->KillSwitchChecker = process->SharedFiles()->FileChecker(KillSwitchFile_, TDuration::Seconds(1));
        }
        return tls;
    }

    struct TSSNChecker : NIcookie::TProcessContext {
        TSSNChecker(const TConnDescr& descr)
            : Descr_(descr)
        {}
        bool SupportsSameSiteNone() const noexcept override {
            return GetOrSetUATraits(Descr_).SameSiteNoneSupport;
        }
        const TConnDescr& Descr_;
    };

    TMaybe<std::pair<bool, NUserSplit::TBucket>> GetExpInfo(const TStringBuf data) const {
        if (data.empty()) {
            return Nothing();
        }

        NUserSplit::THashCalculator hashCalculator(NExpPrivate::BUCKETS_COUNT, NExpPrivate::SLOTS_COUNT, ExpSalt_);

        NUserSplit::TBucket bucket;
        NUserSplit::TSlot slot;

        hashCalculator.GetBucketAndSlot(data, bucket, slot);
        if (ExpASlots_.contains(slot)) {
            return std::make_pair(false, bucket);
        } else if (ExpBSlots_.contains(slot)) {
            return std::make_pair(true, bucket);
        } else {
            return Nothing();
        }
    }

    TError DoRun(const TConnDescr& descr, TTls& tls) const noexcept override {
        if (Y_UNLIKELY(descr.Request == nullptr)) {
            return Submodule_->Run(descr);
        }

        // these are required here since some structs capture TStringBuf's of TChunkPtr
        // who should live long enough to set cookie
        TProcessBufs processBufs;
        TIcookieSetter setter(*this);

        const TConnDescr* newDescr = &descr;
        TMaybe<TConnDescr> maybeNewConnDescr;
        TMaybe<TRequest> maybeNewRequest;

        if (NeedIdempotence_) {
            maybeNewRequest = *descr.Request;
            maybeNewConnDescr = descr.Copy(maybeNewRequest.Get());
            newDescr = maybeNewConnDescr.Get();
        }

        if (!tls.KillSwitchFileExists()) {
            TModuleProcessContext context(descr);

            context.IcookieContext.IsConnectionSsl = newDescr->Properties->UserConnIsSsl;

            ProcessHeaders(newDescr->Request->Headers(), context, processBufs);
            ProcessYandexuid(newDescr->Request->Headers(), context);
            TryParseSearchappUuidFromSae(context);
            ProcessRequestLine(*newDescr, context, processBufs);

            // USEREXP-10460 TODO: remove after experiment
            CheckExperimentAndWriteExperimentHeader(context, processBufs, newDescr->Request->Headers());

            DoProcessSrcICookie(context, newDescr->Request->Headers(), processBufs, tls);
            DoProcess(context, newDescr->Request->Headers(), setter, tls);
        } else if (!TrustParent_) {
            newDescr->Request->Headers().Delete(*IcookieHeadersFsm_);
        }

        if (setter.NeedToHandle()) {
            auto output = MakeHttpOutput([&](TResponse&& response, const bool forceClose, TInstant deadline) {
                setter.ProcessResponse(response);
                return descr.Output->SendHead(std::move(response), forceClose, deadline);
            }, [&](TChunkList lst, TInstant deadline) {
                return descr.Output->Send(std::move(lst), deadline);
            }, [&](THeaders&& trailers, TInstant deadline) {
                return descr.Output->SendTrailers(std::move(trailers), deadline);
            });
            if (!maybeNewConnDescr) {
                maybeNewConnDescr = descr.Copy();
                newDescr = maybeNewConnDescr.Get();
            }
            maybeNewConnDescr->Output = &output;

            return Submodule_->Run(*newDescr);
        }

        return Submodule_->Run(*newDescr);
    }

    struct TModuleProcessContext : TNonCopyable {
        TModuleProcessContext(const TConnDescr& descr)
            : IcookieContext(descr)
        {}

        TSSNChecker IcookieContext;
        TStringBuf UidEncrypted;
        TStringBuf ExtUidDecryptedFromParent;
        TStringBuf SaeCookie;
        bool ReceivedDecryptedUidFromParent = false;
        bool ReceivedDecryptedLoginHashFromParent = false;
        bool ReceivedIcookieInfoFromParent = false;
        bool ReceivedSrcICookieHeader = false;
        bool ReceivedEncryptedUidFromParent = false;
        bool ReceivedDecryptedExtUidFromParent = false;
    };

    // USEREXP-10460
    void CheckExperimentAndWriteExperimentHeader(TModuleProcessContext& context, TProcessBufs& bufs, THeaders& headers) const noexcept {
        if ((ExpSalt_.empty() || ExpType_ != NIcookie::ET_BROWSER_UUID) && !ForceGenerateFromYandexBrowserUuid_) {
            return;
        }

        if (context.IcookieContext.ClientType != NIcookie::CT_YANDEX_BROWSER) {
            return;
        }

        auto maybeParsed = NSaeCookie::TSaeCookie::FromRaw(context.SaeCookie);
        if (maybeParsed.Defined()) {
            TStringBuf platform = maybeParsed.GetRef().Value(NSaeCookie::ESaeCookiePart::PLATFORM, Default<TStringBuf>());
            TStringBuf deviceType = maybeParsed.GetRef().Value(NSaeCookie::ESaeCookiePart::DEVICE_TYPE, Default<TStringBuf>());

            // USEREXP-12931 - only for android browser
            if (platform == "android" && deviceType == "phone") {
                bufs.Uuid = maybeParsed.GetRef().Value(NSaeCookie::ESaeCookiePart::UUID, Default<TStringBuf>());
            }
        }

        if (ForceGenerateFromYandexBrowserUuid_) {
            context.IcookieContext.Uuid = bufs.Uuid;
        } else if (!ExpSalt_.empty()) {
            auto expInfo = GetExpInfo(bufs.Uuid);
            if (expInfo.Defined()) {
                if (expInfo.GetRef().first) {
                    context.IcookieContext.Uuid = bufs.Uuid;
                    context.IcookieContext.ExperimentType = NIcookie::ET_BROWSER_UUID;
                }

                headers.Add("Y-Balancer-Experiments", TString::Join((expInfo.GetRef().first ? ExpBTestid_ : ExpATestid_), ",0,", ToString(expInfo.GetRef().second)));
            }
        }
    }

    // BROWSER-138806: Try to get uuid from sae if possible.
    // Requires headers to be processed. Should be called before request line
    // is processed, because uuid from sae is higher importance.
    void TryParseSearchappUuidFromSae(TModuleProcessContext& context) const {
        if (EnableParseSearchappUuid_ && context.IcookieContext.ClientType == NIcookie::EClientType::CT_SEARCHAPP) {
            auto maybeParsed = NSaeCookie::TSaeCookie::FromRaw(context.SaeCookie);
            if (maybeParsed.Defined()) {
                context.IcookieContext.Uuid = maybeParsed.GetRef().Value(NSaeCookie::ESaeCookiePart::UUID, Default<TStringBuf>());
            }
        }
    }

    void ProcessHeaders(THeaders& headers, TModuleProcessContext& context, TProcessBufs& bufs) const noexcept {
        for (auto it = headers.begin(); it != headers.end();) {
            auto& header = *it;
            TMatcher matcher{ *AllHeadersFsm_ };

            if (Match(matcher, header.first.AsStringBuf()).Final()) {
                switch (*matcher.MatchedRegexps().first) {
                    case 0:     // Cookie
                        for (const auto& value : header.second) {
                            ProcessCookiesHeader(value.AsStringBuf(), context);
                        }
                        break;
                    case 1:     // Host
                        bufs.HostBuf = header.second.back().AsStringBuf();
                        context.IcookieContext.Host = bufs.HostBuf;
                        break;
                    case 2:     // User-Agent
                        if (ForceGenerateFromSearchappUuid_ || ForceGenerateFromYandexBrowserUuid_ ||
                                ForceGenerateFromTransport_ || !ExpSalt_.empty()) {
                            context.IcookieContext.ClientType = NIcookie::GetClientType(header.second.back().AsStringBuf());
                        }
                        break;
                    case 3:     // X-Yandex-ICookie
                        context.ReceivedDecryptedUidFromParent = true;
                        if (!TrustParent_) {
                            headers.erase(it++);
                            continue;
                        }
                        break;
                    case 4:     // X-Yandex-ICookie-Error
                        if (!TrustParent_) {
                            headers.erase(it++);
                            continue;
                        }
                        break;
                    case 5:     // X-Yandex-ICookie-Info
                        context.ReceivedIcookieInfoFromParent = true;
                        if (!TrustParent_) {
                            headers.erase(it++);
                            continue;
                        }
                        break;
                    case 6:     // X-Yandex-LoginHash
                        context.ReceivedDecryptedLoginHashFromParent = true;
                        if (!TrustParent_) {
                            headers.erase(it++);
                            continue;
                        }
                        break;
                    case 7:     // X-Yandex-Src-ICookie
                        context.ReceivedSrcICookieHeader = true;
                        if (!TrustParent_) {
                            headers.erase(it++);
                            continue;
                        }
                        break;
                    case 8:     // X-Yandex-ICookie-Encrypted
                        context.ReceivedEncryptedUidFromParent = true;
                        if (!TrustParent_) {
                            headers.erase(it++);
                            continue;
                        }
                        break;
                    case 9:  // X-Yandex-ICookie-Ext
                        context.ReceivedDecryptedExtUidFromParent = true;
                        if (!TrustParent_) {
                            headers.erase(it++);
                            continue;
                        } else {
                            bufs.ExtUidDecryptedFromParent = header.second.back().AsStringBuf();
                            context.ExtUidDecryptedFromParent = bufs.ExtUidDecryptedFromParent;
                        }
                        break;
                    case 10:     // X-Yandex-RandomUID
                        bufs.RandomUidBuf = header.second.back().AsStringBuf();
                        context.IcookieContext.YandexRandomUID = bufs.RandomUidBuf;
                        break;
                }
            }
            ++it;
        }
    }

    void ProcessCookiesHeader(TStringBuf cookiesHeader, TModuleProcessContext& context) const noexcept {
        bool haveI = false;
        bool haveYandexuid = false;
        bool haveLogin = false;
        bool haveSae = false;
        FindCookieKFsm(cookiesHeader, TICookieFsm::Instance(), [&](const TNameValue& nv, size_t idx) {
            switch (idx) {
            case TICookieFsm::I:
                haveI = true;
                context.UidEncrypted = nv.Value;
                break;
            case TICookieFsm::YandexUid:
                haveYandexuid = true;
                context.IcookieContext.Yandexuid = nv.Value;
                break;
            case TICookieFsm::YandexLogin:
                haveLogin = true;
                context.IcookieContext.Login = nv.Value;
                break;
            case TICookieFsm::Sae:
                haveSae = true;
                context.SaeCookie = nv.Value;
                break;
            }
            return (haveI && haveYandexuid && haveLogin && haveSae) ? ECookieIter::Stop : ECookieIter::Continue;
        });
    }

    void ProcessYandexuid(THeaders& headers, TModuleProcessContext& context) const noexcept {
        if (Y_LIKELY(!Blacklist_ || !Blacklist_->Has(context.IcookieContext.Yandexuid))) {
            return;
        }

        context.IcookieContext.Yandexuid = "";

        auto oldCookie = headers.GetValuesRef("cookie");

        size_t size = 0;

        for (const auto& item : oldCookie) {
            size += item.size();
        }

        TString newCookie(Reserve(size));

        for (const auto& item : oldCookie) {
            EveryCookieKV([&](TNameValue nv) {
                if (nv.Name.Defined() && *nv.Name == "yandexuid") {
                    return;
                }
                AppendCookieKV(newCookie, nv);
            }, item.AsStringBuf());
        }

        headers.Replace("cookie", std::move(newCookie));
    }

    void ProcessRequestLine(const TConnDescr& newDescr, TModuleProcessContext& context, TProcessBufs& bufs) const noexcept {
        TStringBuf cgi = newDescr.Request->RequestLine().CGI.AsStringBuf();

        if (!cgi.empty() && cgi.front() == '?') {
            cgi = cgi.substr(1);
        }

        bufs.Cgi.ScanAddUnescaped(cgi);

        if (EnableParseSearchappUuid_ && context.IcookieContext.ClientType == NIcookie::EClientType::CT_SEARCHAPP && context.IcookieContext.Uuid.empty()) {
            TStringBuf uuid = bufs.Cgi.Get("uuid");

            if (Match(TSearchappUuidFsm::Instance(), uuid)) {
                context.IcookieContext.Uuid = uuid;
            }
        }
    }

    void DoProcessSrcICookie(TModuleProcessContext& context, THeaders& headers, TProcessBufs& bufs, TTls& tls) const noexcept {
        if ((TrustParent_ && context.ReceivedSrcICookieHeader) || !EnableDecrypting_) {
            return;
        }

        TStringBuf srcICookie = bufs.Cgi.Get("icookie");

        // making best effort of recognizing src-icookie - both escaped and unescaped
        if (srcICookie.size() != NIcookie::FINAL_BASE64_SIZE) {
            bufs.Transport = CGIUnescapeRet(srcICookie);
            srcICookie = bufs.Transport;
        }

        const auto decryptResult = tls.Encrypter->Decrypt(srcICookie);

        if (!decryptResult.IsOk()) {
            return;
        }

        if (MaxTransportAge_ > 0) {
            auto encryptionAge = NIcookie::CalculateAge(decryptResult.GetEncryptionTimestamp());

            if (encryptionAge > MaxTransportAge_) {
                return;
            }
        }

        const auto& uid = decryptResult.GetUid();

        if (!uid.Defined()) {
            return;
        }

        TString result = uid->BuildShortRepresentation();

        headers.Add(SrcICookieHeader_, result);

        context.IcookieContext.Transport = srcICookie;

        if (!ForceGenerateFromTransport_ && !ExpSalt_.empty() && ExpType_ == NIcookie::ET_TRANSPORT) {
            auto expInfo = GetExpInfo(result);

            if (expInfo.Defined()) {
                if (expInfo.GetRef().first) {
                    context.IcookieContext.ExperimentType = NIcookie::ET_TRANSPORT;
                }

                headers.Add("Y-Balancer-Experiments", TString::Join((expInfo.GetRef().first ? ExpBTestid_ : ExpATestid_), ",0,", ToString(expInfo.GetRef().second)));
            }
        }
    }

    void DoProcess(const TModuleProcessContext& context, THeaders& headers, TIcookieSetter& setter, TTls& tls) const noexcept {
        if (TrustParent_ && context.ReceivedDecryptedUidFromParent &&
            context.ReceivedDecryptedLoginHashFromParent && context.ReceivedDecryptedExtUidFromParent &&
            !EncryptedHeader_.empty() && Policy_.GetEnableIcookieEncryptedHeader() &&
            !context.ReceivedEncryptedUidFromParent

        ) {
            try {
                auto value = NIcookie::EncryptIcookie(NIcookie::TUid(context.ExtUidDecryptedFromParent), NIcookie::FLAG_TRANSPORT, *tls.Encrypter);
                if (value.Defined()) {
                    headers.Add(EncryptedHeader_, *value);
                }
            } catch (const NIcookie::TIcookieException&) { // we assume that this is the only kind of exception that can be thrown
                if (EnableDecrypting_) {
                    headers.Add(ErrorHeader_, "Exception");
                }
            }
        }

        if ((TrustParent_ && context.ReceivedDecryptedUidFromParent && context.ReceivedDecryptedLoginHashFromParent) ||
            (!EnableSetCookie_ && !EnableDecrypting_))
        {
            return;
        }

        try {
            auto result = NIcookie::ProcessCookie(context.UidEncrypted,
                                                  context.IcookieContext,
                                                  Policy_,
                                                  *tls.Encrypter,
                                                  nullptr,
                                                  Blacklist_);

            if (Y_LIKELY(result.IsSuccess())) {
                if (EnableSetCookie_ && result.IsSetCookieNeeded()) {
                    setter.SetSetCookieHeader(std::move(result.GetSetCookieHeader()));
                }
                if (EnableDecrypting_) {
                    if (!TrustParent_ || !context.ReceivedDecryptedUidFromParent) {
                        headers.Add(DecryptedUidHeader_, result.GetUid().BuildShortRepresentation());
                    }
                    if (!TrustParent_ || !context.ReceivedDecryptedExtUidFromParent) {
                        headers.Add(DecryptedExtUidHeader_, ToString(result.GetUid()));
                    }

                    if (!TrustParent_ || !context.ReceivedDecryptedLoginHashFromParent) {
                        NIcookie::TLoginHash loginHash;

                        if (NIcookie::GenerateLoginHash(context.IcookieContext.Login, loginHash)) {
                            headers.Add(DecryptedLoginHashHeader_, ToString(loginHash));
                        }
                    }
                    if (!TrustParent_ ||
                        (!context.ReceivedDecryptedUidFromParent && !context.ReceivedIcookieInfoFromParent))
                    {
                        headers.Add(InfoHeader_, result.GetIcookieInfoHeader());
                    }
                }

                if (!EncryptedHeader_.empty()) {
                    auto value = result.GetIcookieEncryptedHeader();

                    if (value.Defined()) {
                        headers.Add(EncryptedHeader_, *value);
                    }
                }
            } else {
                if (EnableDecrypting_) {
                    headers.Add(ErrorHeader_, "Cypher");
                }
            }
        } catch (const NIcookie::TIcookieException&) { // we assume that this is the only kind of exception that can be thrown
            if (EnableDecrypting_) {
                headers.Add(ErrorHeader_, "Exception");
            }
        }
    }

private:
    NIcookie::TIcookieProcessPolicy Policy_;
    const NIcookie::TBlacklist* Blacklist_;

    TVector<TString> Domains_;

    THolder<NIcookie::TKeysSet> KeysSet_;

    TString KillSwitchFile_;

    THolder<TFsm> IcookieHeadersFsm_;
    THolder<TFsm> AllHeadersFsm_;

    TString DecryptedUidHeader_{ NIcookie::ICOOKIE_DECRYPTED_HEADER };
    TString DecryptedExtUidHeader_{ NIcookie::EXT_ICOOKIE_DECRYPTED_HEADER };
    TString ErrorHeader_{ NIcookie::ICOOKIE_ERROR_HEADER };
    TString InfoHeader_{ NIcookie::ICOOKIE_INFO_HEADER };
    TString EncryptedHeader_;
    TString TakeRandomUidFrom_;

    TString DecryptedLoginHashHeader_{ NIcookie::LOGIN_HASH_DECRYPTED_HEADER };
    TString SrcICookieHeader_{ NIcookie::SRC_ICOOKIE_HEADER };

    TString KeysFilename_;

    NIcookie::TIcookieProcessPolicy::TScheme IcookieSchemeBitmask_{NIcookie::TIcookieProcessPolicy::SchemeSsl};

    bool UseDefaultKeys_{ false };  // should be set explicitly

    bool TrustParent_{ false };     // trust parent 'X-Yandex-ICookie'
    bool TrustChildren_{ false };   // trust children 'Set-Cookie: i=...'

    bool EnableSetCookie_{ true };
    bool EnableDecrypting_{ true };

    bool FlagSecure_{ true };       // USEREXP-4244

    bool ForceEqualToYandexuid_{ false };   // USEREXP-5271
    bool ForceGenerateFromSearchappUuid_{ false };  // USEREXP-5030
    bool ForceGenerateFromYandexBrowserUuid_{ false };
    bool ForceGenerateFromTransport_{ false };  // USEREXP-11685

    ui32 MaxTransportAge_{ 0 };     // USEREXP-11685

    TString ExpSalt_;
    TString ExpATestid_;
    TString ExpBTestid_;
    TString ExpASlotsStr_;
    TString ExpBSlotsStr_;
    ui32 ExpType_{ NIcookie::ET_BROWSER_UUID };    // backwards compatibility

    THashSet<ui32> ExpASlots_;
    THashSet<ui32> ExpBSlots_;

    bool EnableParseSearchappUuid_{ true }; // USEREXP-5030
    bool EnableGuessSearchapp_{ true };     // USEREXP-6946
    bool NeedIdempotence_{ true };

    NGdprCookie::ECookieType GdprCookieType_ = NGdprCookie::CookieType("i");
};

IModuleHandle* Handle() {
    return TModule::Handle();
}
}
