//

// Copyright 2023 Google LLC

//

// Licensed under the Apache License, Version 2.0 (the "License");

// you may not use this file except in compliance with the License.

// You may obtain a copy of the License at

//

//      https://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing, software

// distributed under the License 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 "eval/compiler/comprehension_vulnerability_check.h"

#include <algorithm>

#include <memory>

#include <vector>

#include "absl/status/status.h"

#include "absl/strings/string_view.h"

#include "absl/types/variant.h"

#include "base/builtins.h"

#include "common/ast.h"

#include "common/constant.h"

#include "common/expr.h"

#include "eval/compiler/flat_expr_builder_extensions.h"

namespace google::api::expr::runtime {

namespace {

using ::cel::CallExpr;

using ::cel::ComprehensionExpr;

using ::cel::Constant;

using ::cel::Expr;

using ::cel::IdentExpr;

using ::cel::ListExpr;

using ::cel::MapExpr;

using ::cel::SelectExpr;

using ::cel::StructExpr;

using ::cel::UnspecifiedExpr;

// ComprehensionAccumulationReferences recursively walks an expression to count

// the locations where the given accumulation var_name is referenced.

//

// The purpose of this function is to detect cases where the accumulation

// variable might be used in hand-rolled ASTs that cause exponential memory

// consumption. The var_name is generally not accessible by CEL expression

// writers, only by macro authors. However, a hand-rolled AST makes it possible

// to misuse the accumulation variable.

//

// Limitations:

// - This check only covers standard operators and functions.

//   Extension functions may cause the same issue if they allocate an amount of

//   memory that is dependent on the size of the inputs.

//

// - This check is not exhaustive. There may be ways to construct an AST to

//   trigger exponential memory growth not captured by this check.

//

// The algorithm for reference counting is as follows:

//

//  * Calls - If the call is a concatenation operator, sum the number of places

//            where the variable appears within the call, as this could result

//            in memory explosion if the accumulation variable type is a list

//            or string. Otherwise, return 0.

//

//            accu: ["hello"]

//            expr: accu + accu // memory grows exponentionally

//

//  * CreateList - If the accumulation var_name appears within multiple elements

//            of a CreateList call, this means that the accumulation is

//            generating an ever-expanding tree of values that will likely

//            exhaust memory.

//

//            accu: ["hello"]

//            expr: [accu, accu] // memory grows exponentially

//

//  * CreateStruct - If the accumulation var_name as an entry within the

//            creation of a map or message value, then it's possible that the

//            comprehension is accumulating an ever-expanding tree of values.

//

//            accu: {"key": "val"}

//            expr: {1: accu, 2: accu}

//

//  * Comprehension - If the accumulation var_name is not shadowed by a nested

//            iter_var or accu_var, then it may be accmulating memory within a

//            nested context. The accumulation may occur on either the

//            comprehension loop_step or result step.

//

// Since this behavior generally only occurs within hand-rolled ASTs, it is

// very reasonable to opt-in to this check only when using human authored ASTs.

int ComprehensionAccumulationReferences(const cel::Expr& expr,

                                        absl::string_view var_name) {

  struct Handler {

    const Expr& expr;

    absl::string_view var_name;

    int operator()(const CallExpr& call) {

      int references = 0;

      absl::string_view function = call.function();

      // Return the maximum reference count of each side of the ternary branch.

      if (function == cel::builtin::kTernary && call.args().size() == 3) {

        return std::max(

            ComprehensionAccumulationReferences(call.args()[1], var_name),

            ComprehensionAccumulationReferences(call.args()[2], var_name));

      }

      // Return the number of times the accumulator var_name appears in the add

      // expression. There's no arg size check on the add as it may become a

      // variadic add at a future date.

      if (function == cel::builtin::kAdd) {

        for (int i = 0; i < call.args().size(); i++) {

          references +=

              ComprehensionAccumulationReferences(call.args()[i], var_name);

        }

        return references;

      }

      // Return whether the accumulator var_name is used as the operand in an

      // index expression or in the identity `dyn` function.

      if ((function == cel::builtin::kIndex && call.args().size() == 2) ||

          (function == cel::builtin::kDyn && call.args().size() == 1)) {

        return ComprehensionAccumulationReferences(call.args()[0], var_name);

      }

      return 0;

    }

    int operator()(const ComprehensionExpr& comprehension) {

      absl::string_view accu_var = comprehension.accu_var();

      absl::string_view iter_var = comprehension.iter_var();

      int result_references = 0;

      int loop_step_references = 0;

      int sum_of_accumulator_references = 0;

      // The accumulation or iteration variable shadows the var_name and so will

      // not manipulate the target var_name in a nested comprehension scope.

      if (accu_var != var_name && iter_var != var_name) {

        loop_step_references = ComprehensionAccumulationReferences(

            comprehension.loop_step(), var_name);

      }

      // Accumulator variable (but not necessarily iter var) can shadow an

      // outer accumulator variable in the result sub-expression.

      if (accu_var != var_name) {

        result_references = ComprehensionAccumulationReferences(

            comprehension.result(), var_name);

      }

      // Count the raw number of times the accumulator variable was referenced.

      // This is to account for cases where the outer accumulator is shadowed by

      // the inner accumulator, while the inner accumulator is being used as the

      // iterable range.

      //

      // An equivalent expression to this problem:

      //

      // outer_accu := outer_accu

      // for y in outer_accu:

      //     outer_accu += input

      // return outer_accu

      // If this is overly restrictive (Ex: when generalized reducers is

      // implemented), we may need to revisit this solution

      sum_of_accumulator_references = ComprehensionAccumulationReferences(

          comprehension.accu_init(), var_name);

      sum_of_accumulator_references += ComprehensionAccumulationReferences(

          comprehension.iter_range(), var_name);

      // Count the number of times the accumulator var_name within the loop_step

      // or the nested comprehension result.

      //

      // This doesn't cover cases where the inner accumulator accumulates the

      // outer accumulator then is returned in the inner comprehension result.

      return std::max({loop_step_references, result_references,

                       sum_of_accumulator_references});

    }

    int operator()(const ListExpr& list) {

      // Count the number of times the accumulator var_name appears within a

      // create list expression's elements.

      int references = 0;

      for (int i = 0; i < list.elements().size(); i++) {

        references += ComprehensionAccumulationReferences(

            list.elements()[i].expr(), var_name);

      }

      return references;

    }

    int operator()(const StructExpr& map) {

      // Count the number of times the accumulation variable occurs within

      // entry values.

      int references = 0;

      for (int i = 0; i < map.fields().size(); i++) {

        const auto& entry = map.fields()[i];

        if (entry.has_value()) {

          references +=

              ComprehensionAccumulationReferences(entry.value(), var_name);

        }

      }

      return references;

    }

    int operator()(const MapExpr& map) {

      // Count the number of times the accumulation variable occurs within

      // entry values.

      int references = 0;

      for (int i = 0; i < map.entries().size(); i++) {

        const auto& entry = map.entries()[i];

        if (entry.has_value()) {

          references +=

              ComprehensionAccumulationReferences(entry.value(), var_name);

        }

      }

      return references;

    }

    int operator()(const SelectExpr& select) {

      // Test only expressions have a boolean return and thus cannot easily

      // allocate large amounts of memory.

      if (select.test_only()) {

        return 0;

      }

      // Return whether the accumulator var_name appears within a non-test

      // select operand.

      return ComprehensionAccumulationReferences(select.operand(), var_name);

    }

    int operator()(const IdentExpr& ident) {

      // Return whether the identifier name equals the accumulator var_name.

      return ident.name() == var_name ? 1 : 0;

    }

    int operator()(const Constant& constant) { return 0; }

    int operator()(const UnspecifiedExpr&) { return 0; }

  } handler{expr, var_name};

  return absl::visit(handler, expr.kind());

}

bool ComprehensionHasMemoryExhaustionVulnerability(

    const ComprehensionExpr& comprehension) {

  absl::string_view accu_var = comprehension.accu_var();

  const auto& loop_step = comprehension.loop_step();

  return ComprehensionAccumulationReferences(loop_step, accu_var) >= 2;

}

class ComprehensionVulnerabilityCheck : public ProgramOptimizer {

 public:

  absl::Status OnPreVisit(PlannerContext& context, const Expr& node) override {

    if (node.has_comprehension_expr() &&

        ComprehensionHasMemoryExhaustionVulnerability(

            node.comprehension_expr())) {

      return absl::InvalidArgumentError(

          "Comprehension contains memory exhaustion vulnerability");

    }

    return absl::OkStatus();

  }

  absl::Status OnPostVisit(PlannerContext& context,

                           const cel::Expr& node) override {

    return absl::OkStatus();

  }

};

}  // namespace

ProgramOptimizerFactory CreateComprehensionVulnerabilityCheck() {

  return [](PlannerContext&, const cel::Ast& ast) {

    return std::make_unique<ComprehensionVulnerabilityCheck>();

  };

}

}  // namespace google::api::expr::runtime