Expanding traffic_annotation_extractor clang tool to extract network calls.

tools/clang/traffic_annotation_extractor/traffic_annotation_extractor.cpp

 // Copyright 2017 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 // This clang tool finds all instances of the following functions:
 //   - net::DefineNetworkTrafficAnnotation
 //   - net::DefinePartialNetworkTrafficAnnotation
 //   - net::CompleteNetworkTrafficAnnotation
 //   - net::BranchedCompleteNetworkTrafficAnnotation
 // It extracts the location info and content of annotation tags, and outputs
-// them to llvm::outs. Please refer to README.md for build and usage
-// instructions.
+// them to llvm::outs. It also extracts all calls the following network request
+// creation functions and returns their source location and availability of a
+// net::[Partial]NetworkTrafficAnnotation parameter in them:
+//   - SSLClientSocket::SSLClientSocket
+//   - TCPClientSocket::TCPClientSocket
+//   - UDPClientSocket::UDPClientSocket
+//   - URLFetcher::Create
+//   - ClientSocketFactory::CreateDatagramClientSocket
+//   - ClientSocketFactory::CreateSSLClientSocket
+//   - ClientSocketFactory::CreateTransportClientSocket
+//   - URLRequestContext::CreateRequest
+// Please refer to README.md for build and usage instructions.
 
 #include <memory>
 #include <vector>
 
 #include "clang/ASTMatchers/ASTMatchFinder.h"
 #include "clang/ASTMatchers/ASTMatchers.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Frontend/FrontendActions.h"
 #include "clang/Lex/Lexer.h"
 #include "clang/Tooling/CommonOptionsParser.h"
 #include "clang/Tooling/Refactoring.h"
 #include "clang/Tooling/Tooling.h"
 #include "llvm/Support/CommandLine.h"
 
 using namespace clang::ast_matchers;
 
 namespace {
 
+// Information about location of a line of code.
+struct Location {
+  std::string file_path;
+  int line_number = -1;
+
+  // Name of the function including this line. E.g., in the following code,
+  // |function_name| will be 'foo' for all |line_number| values 101-103.
+  //
+  // 100 void foo() {
+  // 101   NetworkTrafficAnnotationTag baz =
+  // 102       net::DefineNetworkTrafficAnnotation(...); }
+  // 103   bar(baz);
+  // 104 }
+  // If no function is found, 'Global Namespace' will be returned.
+  std::string function_name;
+};
+
 // An instance of a call to either of the 4 network traffic annotation
 // definition functions.
 struct NetworkAnnotationInstance {
-  // Information about where the call has happened.
-  struct Location {
-    std::string file_path;
-    int line_number = -1;
-
-    // Name of the function including this line. E.g., in the following code,
-    // |function_name| will be 'foo' for all |line_number| values 101-103.
-    //
-    // 100 void foo() {
-    // 101   NetworkTrafficAnnotationTag baz =
-    // 102       net::DefineNetworkTrafficAnnotation(...); }
-    // 103   bar(baz);
-    // 104 }
-    // If no function is found, 'Global Namespace' will be returned.
-    std::string function_name;
-  };
-
   // Annotation content. These are the arguments of the call to either of the 4
   // network traffic annotation definition functions.
   struct Annotation {
     std::string unique_id;
     std::string text;
 
     // |extra_id| will have |completing_id| for
     // net::DefinePartialNetworkTrafficAnnotation and |group_id| for
     // net::BranchedCompleteNetworkTrafficAnnotation. It will be empty in other
     // cases.
@@ skipping 22 matching lines @@
       case kCompleting:
         return "Completing";
       case kBranchedCompleting:
         return "BranchedCompleting";
     }
     assert(false);
     return "";
   }
 };
 
-using Collector = std::vector<NetworkAnnotationInstance>;
+// An instance of a call to one of the monitored function.
+struct CallInstance {
+  // Location of the call.
+  Location location;
+
+  // Whether the function is annotated.
+  bool has_annotation = false;
+
+  // Name of the called function.
+  std::string called_function_name;
+};
+
+// A structure to keep detected annotation and call instances.
+struct Collector {
+  std::vector<NetworkAnnotationInstance> annotations;
+  std::vector<CallInstance> calls;
+};
 
 // This class implements the call back functions for AST Matchers. The matchers
 // are defined in RunMatchers function. When a pattern is found there,
 // the run function in this class is called back with information on the matched
 // location and description of the matched pattern.
 class NetworkAnnotationTagCallback : public MatchFinder::MatchCallback {
  public:
   explicit NetworkAnnotationTagCallback(Collector* collector)
       : collector_(collector) {}
   ~NetworkAnnotationTagCallback() override = default;
 
   // Is called on any pattern found by ASTMathers that are defined in RunMathers
   // function.
   virtual void run(const MatchFinder::MatchResult& result) override {
+    if (const clang::CallExpr* call_expr =
+            result.Nodes.getNodeAs<clang::CallExpr>("monitored_function")) {
+      AddFunction(call_expr, result);
+    } else {
+      AddAnnotation(result);
+    }
+  }
+
+  void GetInstanceLocation(const MatchFinder::MatchResult& result,
+                           const clang::CallExpr* call_expr,
+                           const clang::FunctionDecl* ancestor,
+                           Location* instance_location) {
+    clang::SourceLocation source_location = call_expr->getLocStart();
+    if (source_location.isMacroID()) {
+      source_location =
+          result.SourceManager->getImmediateMacroCallerLoc(source_location);
+    }
+    instance_location->file_path =
+        result.SourceManager->getFilename(source_location);
+    instance_location->line_number =
+        result.SourceManager->getSpellingLineNumber(source_location);
+    if (ancestor)
+      instance_location->function_name = ancestor->getQualifiedNameAsString();
+    else
+      instance_location->function_name = "Global Namespace";
+
+    std::replace(instance_location->file_path.begin(),
+                 instance_location->file_path.end(), '\\', '/');
+
+    // Trim leading "../"s from file path.
+    while (instance_location->file_path.length() > 3 &&
+           instance_location->file_path.substr(0, 3) == "../") {
+      instance_location->file_path = instance_location->file_path.substr(
+          3, instance_location->file_path.length() - 3);
+    }
+  }
+
+  // Stores a function call that should be monitored.
+  void AddFunction(const clang::CallExpr* call_expr,
+                   const MatchFinder::MatchResult& result) {
+    CallInstance instance;
+
+    const clang::FunctionDecl* ancestor =
+        result.Nodes.getNodeAs<clang::FunctionDecl>("function_context");
+    GetInstanceLocation(result, call_expr, ancestor, &instance.location);
+    instance.called_function_name =
+        call_expr->getDirectCallee()->getQualifiedNameAsString();
+    instance.has_annotation =
+        (result.Nodes.getNodeAs<clang::RecordDecl>("annotation") != NULL);

[dcheng, 2017/05/29 08:28:00]

Nit: nullptr instead of NULL

[Ramin Halavati, 2017/05/29 09:14:02]

Done.

+    collector_->calls.push_back(instance);
+  }
+
+  // Stores an annotation.
+  void AddAnnotation(const MatchFinder::MatchResult& result) {
     NetworkAnnotationInstance instance;
 
     const clang::StringLiteral* unique_id =
         result.Nodes.getNodeAs<clang::StringLiteral>("unique_id");
     const clang::StringLiteral* annotation_text =
         result.Nodes.getNodeAs<clang::StringLiteral>("annotation_text");
     const clang::FunctionDecl* ancestor =
         result.Nodes.getNodeAs<clang::FunctionDecl>("function_context");
     const clang::StringLiteral* group_id =
         result.Nodes.getNodeAs<clang::StringLiteral>("group_id");
@@ skipping 18 matching lines @@
       assert(group_id);
       instance.annotation.extra_id = group_id->getString();
     } else {
       assert(false);
     }
 
     assert(unique_id && annotation_text);
     instance.annotation.unique_id = unique_id->getString();
     instance.annotation.text = annotation_text->getString();
 
-    // Get annotation location.
-    clang::SourceLocation source_location = call_expr->getLocStart();
-    if (source_location.isMacroID()) {
-      source_location =
-          result.SourceManager->getImmediateMacroCallerLoc(source_location);
-    }
-    instance.location.file_path =
-        result.SourceManager->getFilename(source_location);
-    instance.location.line_number =
-        result.SourceManager->getSpellingLineNumber(source_location);
-    if (ancestor)
-      instance.location.function_name = ancestor->getQualifiedNameAsString();
-    else
-      instance.location.function_name = "Global Namespace";
+    GetInstanceLocation(result, call_expr, ancestor, &instance.location);
 
-    // Trim leading "../"s from file path.
-    std::replace(instance.location.file_path.begin(),
-                 instance.location.file_path.end(), '\\', '/');
-    while (instance.location.file_path.length() > 3 &&
-           instance.location.file_path.substr(0, 3) == "../") {
-      instance.location.file_path = instance.location.file_path.substr(
-          3, instance.location.file_path.length() - 3);
-    }
-
-    collector_->push_back(instance);
+    collector_->annotations.push_back(instance);
   }
 
  private:
   Collector* collector_;
 };
 
 // Sets up an ASTMatcher and runs clang tool to populate collector. Returns the
 // result of running the clang tool.
 int RunMatchers(clang::tooling::ClangTool* clang_tool, Collector* collector) {
   NetworkAnnotationTagCallback callback(collector);
@@ skipping 36 matching lines @@
       callExpr(hasDeclaration(functionDecl(anyOf(
                    hasName("BranchedCompleteNetworkTrafficAnnotation"),
                    hasName("net::BranchedCompleteNetworkTrafficAnnotation")))),
                hasArgument(0, stringLiteral().bind("unique_id")),
                hasArgument(1, stringLiteral().bind("group_id")),
                hasArgument(3, stringLiteral().bind("annotation_text")),
                anyOf(hasAncestor(functionDecl().bind("function_context")),
                      unless(hasAncestor(functionDecl()))))
           .bind("branched_completing_function"),
       &callback);
+
+  // Setup patterns to find functions that should be monitored.
+  match_finder.addMatcher(
+      callExpr(
+          hasDeclaration(functionDecl(
+              anyOf(hasName("SSLClientSocket::SSLClientSocket"),
+                    hasName("TCPClientSocket::TCPClientSocket"),
+                    hasName("UDPClientSocket::UDPClientSocket"),
+                    hasName("URLFetcher::Create"),
+                    hasName("ClientSocketFactory::CreateDatagramClientSocket"),
+                    hasName("ClientSocketFactory::CreateSSLClientSocket"),
+                    hasName("ClientSocketFactory::CreateTransportClientSocket"),
+                    hasName("URLRequestContext::CreateRequest")),
+              anyOf(
+                  hasAnyParameter(hasType(
+                      recordDecl(
+                          anyOf(hasName("net::NetworkTrafficAnnotationTag"),
+                                hasName(
+                                    "net::PartialNetworkTrafficAnnotationTag")))
+                          .bind("annotation"))),
+                  unless(hasAnyParameter(hasType(recordDecl(anyOf(
+                      hasName("net::NetworkTrafficAnnotationTag"),
+                      hasName(
+                          "net::PartialNetworkTrafficAnnotationTag"))))))))),
+          anyOf(hasAncestor(functionDecl().bind("function_context")),
+                unless(hasAncestor(functionDecl()))))

[dcheng, 2017/05/29 08:28:00]

One thing to consider is to move some of these expressions to a submatcher
that's reused.

e.g.

auto is_network_request_function = callExpr(
    hasDeclaration(functionDecl(...));

auto bind_function_context_if_present =
    expr(anyOf(
        hasAncestor(
            functionDecl().bind("function_context")), 
        unless(...));

auto has_annotation_parameter = callExpr(
  anyOf(hasAnyParameter(...)), unless(...));

then combine it with:
match_finder.addMatcher(
    callExpr(
        is_network_request_function,
        has_annotation_parameter,
        bind_function_context_if_present), <callback>);

[Ramin Halavati, 2017/05/29 09:14:02]

Thank you very much, Done.

+          .bind("monitored_function"),
+      &callback);
+
   std::unique_ptr<clang::tooling::FrontendActionFactory> frontend_factory =
       clang::tooling::newFrontendActionFactory(&match_finder);
   return clang_tool->run(frontend_factory.get());
 }
 
 }  // namespace
 
 static llvm::cl::OptionCategory ToolCategory(
     "traffic_annotation_extractor: Extract traffic annotation texts");
 static llvm::cl::extrahelp CommonHelp(
     clang::tooling::CommonOptionsParser::HelpMessage);
 
 int main(int argc, const char* argv[]) {
   clang::tooling::CommonOptionsParser options(argc, argv, ToolCategory);
   clang::tooling::ClangTool tool(options.getCompilations(),
                                  options.getSourcePathList());
   Collector collector;
 
   int result = RunMatchers(&tool, &collector);
 
   if (result != 0)
     return result;
 
   // For each call to any of the functions that define a network traffic
   // annotation, write annotation text and relevant meta data into llvm::outs().
-  for (const NetworkAnnotationInstance& call : collector) {
+  for (const NetworkAnnotationInstance& instance : collector.annotations) {
     llvm::outs() << "==== NEW ANNOTATION ====\n";
-    llvm::outs() << call.location.file_path << "\n";
-    llvm::outs() << call.location.function_name << "\n";
-    llvm::outs() << call.location.line_number << "\n";
-    llvm::outs() << call.GetTypeName() << "\n";
-    llvm::outs() << call.annotation.unique_id << "\n";
-    llvm::outs() << call.annotation.extra_id << "\n";
-    llvm::outs() << call.annotation.text << "\n";
+    llvm::outs() << instance.location.file_path << "\n";
+    llvm::outs() << instance.location.function_name << "\n";
+    llvm::outs() << instance.location.line_number << "\n";
+    llvm::outs() << instance.GetTypeName() << "\n";
+    llvm::outs() << instance.annotation.unique_id << "\n";
+    llvm::outs() << instance.annotation.extra_id << "\n";
+    llvm::outs() << instance.annotation.text << "\n";
     llvm::outs() << "==== ANNOTATION ENDS ====\n";
   }
 
+  // For each call, write annotation text and relevant meta data.
+  for (const CallInstance& instance : collector.calls) {
+    llvm::outs() << "==== NEW CALL ====\n";
+    llvm::outs() << instance.location.file_path << "\n";
+    llvm::outs() << instance.location.function_name << "\n";
+    llvm::outs() << instance.location.line_number << "\n";
+    llvm::outs() << instance.called_function_name << "\n";
+    llvm::outs() << instance.has_annotation << "\n";
+    llvm::outs() << "==== CALL ENDS ====\n";
+  }
+
   return 0;
 }