Annotation features

Generate features from parameter values in annotations.

documentation/website/documentation/models.md

@@ -415,6 +415,117 @@ Note that this only works for numeric and string literals. In cases where the ar
 }
 ```
 
+#### Annotation Features
+
+In model generators we can also use annotation features, which translate to regular user features based on annotation parameter values.
+
+Config example:
+```json
+{
+{
+  "model_generators": [
+    {
+      "find": "methods",
+      "where": [
+        {
+          "constraint": "signature_match",
+          "parent": "Lcom/facebook/marianatrench/integrationtests/AnnotationFeature;",
+          "name": "getSourceWithMethodAnnotation"
+        }
+      ],
+      "model": {
+        "generations": [
+          {
+            "kind": "Source",
+            "port": "Return",
+            "via_annotation": [
+              {
+                "type": "Lcom/facebook/marianatrench/integrationtests/Path;",
+                "port": "Return"
+              }
+            ]
+          }
+        ]
+      }
+    },
+    {
+      "find": "methods",
+      "where": [
+        {
+          "constraint": "signature_match",
+          "parent": "Lcom/facebook/marianatrench/integrationtests/AnnotationFeature;",
+          "name": "getSourceWithParameterAnnotation"
+        }
+      ],
+      "model": {
+        "generations": [
+          {
+            "kind": "Source",
+            "port": "Return",
+            "via_annotation": [
+              {
+                "type": "Lcom/facebook/marianatrench/integrationtests/QueryParam;",
+                "port": "Argument(1)"
+              }
+            ]
+          }
+        ]
+      }
+    }
+  ]
+}
+```
+
+Java class with annotations:
+```java
+public class AnnotationFeature {
+
+  @Path("/issue_1")
+  Object getSourceWithMethodAnnotation() {
+    return new Object();
+  }
+
+  Object getSourceWithParameterAnnotation(@QueryParam("query_param_name") String value) {
+    return "unrelated";
+  }
+
+  void testSourceWithMethodAnnotation() {
+    Object source = getSourceWithMethodAnnotation();
+    Origin.sink(source);
+  }
+
+  void testSourceWithParameterAnnotation() {
+    Object source = getSourceWithParameterAnnotation("argument_value");
+    Origin.sink(source);
+  }
+}
+```
+
+Resulting issues:
+
+```json
+{
+  "issues" : 
+  [
+    {
+      "always_features" : 
+      [
+        "/issue_1"
+      ],
+      "callee" : "Lcom/facebook/marianatrench/integrationtests/Origin;.sink:(Ljava/lang/Object;)V",
+...
+{
+  "issues" : 
+  [
+    {
+      "always_features" : 
+      [
+        "query_param_name"
+      ],
+      "callee" : "Lcom/facebook/marianatrench/integrationtests/Origin;.sink:(Ljava/lang/Object;)V",
+...
+```
+
 ### Taint Broadening
 
 **Taint broadening** (also called **collapsing**) happens when Mariana Trench needs to make an approximation about a taint flow. It is the operation of reducing a **taint tree** into a single element. A **taint tree** is a tree where edges are field names and nodes are taint element. This is how Mariana Trench represents internally which fields (or sequence of fields) are tainted.

source/AnnotationFeature.cpp

@@ -0,0 +1,58 @@
+#include <mariana-trench/AnnotationFeature.h>
+#include <mariana-trench/Context.h>
+#include <mariana-trench/FeatureFactory.h>
+#include <mariana-trench/JsonValidation.h>
+
+namespace marianatrench {
+
+AnnotationFeature::AnnotationFeature(Root port, const DexType* dex_type, std::optional<std::string> label) : port_(std::move(port)), dex_type_(dex_type), label_(label) {
+}
+
+const AnnotationFeature* AnnotationFeature::from_json(const Json::Value& value, Context& context) {
+  JsonValidation::validate_object(value);
+  Root port = Root::from_json(value["port"]);
+  const DexType* dex_type = JsonValidation::dex_type(value, "type");
+  Json::Value configured_label = value["label"];
+  std::optional<std::string> label;
+  if (!configured_label.isNull()) {
+    label = JsonValidation::string(configured_label);
+  }
+  return context.feature_factory->get_unique_annotation_feature(AnnotationFeature{std::move(port), dex_type, std::move(label)});
+}
+
+std::size_t AnnotationFeature::hash() const {
+  std::size_t seed = 0;
+  boost::hash_combine(seed, std::hash<Root>{}(port_));
+  boost::hash_combine(seed, dex_type_);
+  boost::hash_combine(seed, label_);
+  return seed;
+}
+
+bool AnnotationFeature::operator==(const AnnotationFeature& other) const {
+  return port_ == other.port_ && dex_type_ == other.dex_type_ && label_ == other.label_;
+}
+
+const Root& AnnotationFeature::port() const {
+  return port_;
+}
+
+const DexType* AnnotationFeature::dex_type() const {
+  return dex_type_;
+}
+
+const std::optional<std::string>& AnnotationFeature::label() const {
+  return label_;
+}
+
+std::ostream& operator<<(std::ostream& out, const AnnotationFeature& annotation_feature) {
+  out << "AnnotationFeature(port=`"
+      << annotation_feature.port()
+      << "`, dex_type=`"
+      << annotation_feature.dex_type()->str();
+  if (annotation_feature.label()) {
+    out << "`, label=`" << *annotation_feature.label();
+  }
+  return out << "`)";
+}
+
+}

source/AnnotationFeature.h

@@ -0,0 +1,67 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <functional>
+
+#include <mariana-trench/Access.h>
+
+namespace marianatrench {
+
+/**
+ * Annotation features are mapped to regular user features at model template
+ * instantiation time. They consist of an annotation type to find and a root
+ * port to indicate whether it is a method annotation (Root::Return) or one of
+ * its parameters. The usere label content is set to the content of the
+ * annotation's `value()` parameter, if present.
+ */
+class AnnotationFeature {
+ public:
+  /** Public for testing purposes only, prod only uses from_json. */
+  AnnotationFeature(Root port, const class DexType* dex_type, std::optional<std::string> label);
+
+  /**
+   * Parses the given JSON object, returning a feature created using a
+   * UniquePointerFactory.
+   */
+  static const AnnotationFeature* from_json(const Json::Value& value, class Context& context);
+
+  /** Used to store in a UniquePointerFactory. */
+  std::size_t hash() const;
+  bool operator==(const AnnotationFeature& other) const;
+
+  /**
+   * Annotation location. Root::Return for method, Root::Argument for
+   * parameters.
+   */
+  const Root& port() const;
+
+  /** Type of the annotation. */
+  const DexType* dex_type() const;
+
+  /** Label to use for annotation feature value. */
+  const std::optional<std::string>& label() const;
+
+  friend std::ostream& operator<<(
+      std::ostream& out,
+      const AnnotationFeature& annotation_feature);
+
+ private:
+  Root port_;
+  const DexType* dex_type_;
+  std::optional<std::string> label_;
+};
+
+}
+
+template <>
+struct std::hash<marianatrench::AnnotationFeature> {
+  std::size_t operator()(const marianatrench::AnnotationFeature& annotation_feature) const {
+    return annotation_feature.hash();
+  }
+};

source/AnnotationFeatureSet.cpp

@@ -0,0 +1,30 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <Show.h>
+
+#include <mariana-trench/AnnotationFeatureSet.h>
+#include <mariana-trench/JsonValidation.h>
+#include <mariana-trench/Show.h>
+
+namespace marianatrench {
+
+AnnotationFeatureSet::AnnotationFeatureSet(Set set) : set_(std::move(set)) {}
+
+AnnotationFeatureSet AnnotationFeatureSet::from_json(const Json::Value& value, Context& context) {
+  AnnotationFeatureSet features;
+  for (const auto& feature_value : JsonValidation::null_or_array(value)) {
+    features.add(AnnotationFeature::from_json(feature_value, context));
+  }
+  return features;
+}
+
+std::ostream& operator<<(std::ostream& out, const AnnotationFeatureSet& annotation_features) {
+    return show_set(out, annotation_features);
+}
+
+} // namespace marianatrench

source/AnnotationFeatureSet.h

@@ -0,0 +1,55 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <sparta/AbstractDomain.h>
+
+#include <mariana-trench/AnnotationFeature.h>
+#include <mariana-trench/IncludeMacros.h>
+#include <mariana-trench/PatriciaTreeSetAbstractDomain.h>
+
+namespace marianatrench {
+
+/**
+ * Used to store annotation features in a TaintConfig. Annotation features are
+ * incomplete user feature templates, and are thus not copied to a Frame.
+ * Instead they are instantiated as user features during template model
+ * instantiation.
+ */
+class AnnotationFeatureSet final : public sparta::AbstractDomain<AnnotationFeatureSet> {
+ private:
+  using Set = PatriciaTreeSetAbstractDomain<
+      const AnnotationFeature*,
+      /* bottom_is_empty */ true,
+      /* with_top */ false>;
+
+ public:
+  INCLUDE_SET_MEMBER_TYPES(Set, const AnnotationFeature*)
+
+ private:
+  explicit AnnotationFeatureSet(Set set);
+
+ public:
+  /* Create the bottom (i.e, empty) feature set. */
+  AnnotationFeatureSet() = default;
+
+  INCLUDE_DEFAULT_COPY_CONSTRUCTORS_AND_ASSIGNMENTS(AnnotationFeatureSet)
+  INCLUDE_ABSTRACT_DOMAIN_METHODS(AnnotationFeatureSet, Set, set_)
+  INCLUDE_SET_METHODS(AnnotationFeatureSet, Set, set_, const AnnotationFeature*, iterator)
+
+  static AnnotationFeatureSet from_json(const Json::Value& value, Context& context);
+
+  friend std::ostream& operator<<(
+      std::ostream& out,
+      const AnnotationFeatureSet& annotation_features);
+
+ public:
+  Set set_;
+};
+
+} // namespace marianatrench

source/ArtificialMethods.cpp

@@ -56,6 +56,7 @@ std::vector<Model> ArtificialMethods::models(Context& context) const {
           /* field origins */ {},
           /* inferred_features */ FeatureMayAlwaysSet::bottom(),
           /* user_features */ FeatureSet::bottom(),
+          /* annotation_features */ AnnotationFeatureSet::bottom(),
           /* via_type_of_ports */ {},
           /* via_value_of_ports */ {},
           /* canonical_names */ {},

source/BackwardTaintTransfer.cpp

@@ -619,6 +619,7 @@ void check_flows_to_array_allocation(
       /* inferred features */ {},
       /* locally_inferred_features */ {},
       /* user features */ {},
+      /* annotation_features */ {},
       /* via_type_of_ports */ {},
       /* via_value_of_ports */ {},
       /* canonical_names */ {},

source/CallPositionFrames.cpp

@@ -148,6 +148,7 @@ CallPositionFrames CallPositionFrames::attach_position(
                 frame.field_origins(),
                 inferred_features,
                 /* user_features */ FeatureSet::bottom(),
+                /* annotation_features */ AnnotationFeatureSet::bottom(),
                 /* via_type_of_ports */ {},
                 /* via_value_of_ports */ {},
                 frame.canonical_names(),
@@ -219,6 +220,7 @@ CallPositionFrames::map_positions(
           frame.field_origins(),
           frame.inferred_features(),
           frame.user_features(),
+          /* annotation_features */ AnnotationFeatureSet::bottom(),
           frame.via_type_of_ports(),
           frame.via_value_of_ports(),
           frame.canonical_names(),

source/FeatureFactory.cpp

@@ -81,6 +81,10 @@ const Feature* FeatureFactory::get_invalid_path_broadening() const {
   return factory_.create("via-invalid-path-broadening");
 }
 
+const AnnotationFeature* FeatureFactory::get_unique_annotation_feature(const AnnotationFeature& annotation_feature) const {
+  return annotation_factory_.create(annotation_feature);
+}
+
 const FeatureFactory& FeatureFactory::singleton() {
   // Thread-safe global variable, initialized on first call.
   static FeatureFactory instance;