C#: Update the extractor to include the kind of the extracted attributes. Furthermore, include method return attributes.

Author: michaelnebel

csharp/extractor/Semmle.Extraction.CSharp/Entities/Attribute.cs

@@ -1,3 +1,5 @@
+using System;
+using System.Collections.Generic;
 using System.IO;
 using System.Linq;
 using Microsoft.CodeAnalysis;
@@ -6,18 +8,28 @@
 
 namespace Semmle.Extraction.CSharp.Entities
 {
+    internal enum AttributeKind
+    {
+        Default = 0,
+        Return = 1,
+        Assembly = 2,
+        Module = 3,
+    }
+
     internal class Attribute : CachedEntity<AttributeData>, IExpressionParentEntity
     {
         bool IExpressionParentEntity.IsTopLevelParent => true;
 
         private readonly AttributeSyntax? attributeSyntax;
         private readonly IEntity entity;
+        private readonly AttributeKind kind;
 
-        private Attribute(Context cx, AttributeData attributeData, IEntity entity)
+        private Attribute(Context cx, AttributeData attributeData, IEntity entity, AttributeKind kind)
             : base(cx, attributeData)
         {
             this.attributeSyntax = attributeData.ApplicationSyntaxReference?.GetSyntax() as AttributeSyntax;
             this.entity = entity;
+            this.kind = kind;
         }
 
         public override void WriteId(EscapingTextWriter trapFile)
@@ -48,7 +60,7 @@ public sealed override void WriteQuotedId(EscapingTextWriter trapFile)
         public override void Populate(TextWriter trapFile)
         {
             var type = Type.Create(Context, Symbol.AttributeClass);
-            trapFile.attributes(this, type.TypeRef, entity);
+            trapFile.attributes(this, kind, type.TypeRef, entity);
             trapFile.attribute_location(this, Location);
 
             if (attributeSyntax is not null)
@@ -125,26 +137,36 @@ private void ExtractArguments(TextWriter trapFile)
 
         public override bool NeedsPopulation => true;
 
+        private static void ExtractAttributes<T>(Context cx, T symbol, Func<T, IEnumerable<AttributeData>> getAttributes, IEntity entity, AttributeKind kind) where T : ISymbol
+        {
+            foreach (var attribute in getAttributes(symbol))
+            {
+                Create(cx, attribute, entity, kind);
+            }
+        }
+
         public static void ExtractAttributes(Context cx, ISymbol symbol, IEntity entity)
         {
-            foreach (var attribute in symbol.GetAttributes())
+            ExtractAttributes(cx, symbol, s => s.GetAttributes(), entity, AttributeKind.Default);
+            if (symbol is IMethodSymbol method)
             {
-                Create(cx, attribute, entity);
+                ExtractAttributes(cx, method, s => s.GetReturnTypeAttributes(), entity, AttributeKind.Return);
             }
         }
 
-        public static Attribute Create(Context cx, AttributeData attributeData, IEntity entity)
+
+        public static Attribute Create(Context cx, AttributeData attributeData, IEntity entity, AttributeKind kind)
         {
-            var init = (attributeData, entity);
+            var init = (attributeData, entity, kind);
             return AttributeFactory.Instance.CreateEntity(cx, attributeData, init);
         }
 
-        private class AttributeFactory : CachedEntityFactory<(AttributeData attributeData, IEntity receiver), Attribute>
+        private class AttributeFactory : CachedEntityFactory<(AttributeData attributeData, IEntity receiver, AttributeKind kind), Attribute>
         {
             public static readonly AttributeFactory Instance = new AttributeFactory();
 
-            public override Attribute Create(Context cx, (AttributeData attributeData, IEntity receiver) init) =>
-                new Attribute(cx, init.attributeData, init.receiver);
+            public override Attribute Create(Context cx, (AttributeData attributeData, IEntity receiver, AttributeKind kind) init) =>
+                new Attribute(cx, init.attributeData, init.receiver, init.kind);
         }
     }
 }

csharp/extractor/Semmle.Extraction.CSharp/Populators/TypeContainerVisitor.cs

@@ -80,17 +80,26 @@ private void ExtractTypeDeclaration(BaseTypeDeclarationSyntax node)
             Entities.Type.Create(Cx, Cx.GetModel(node).GetDeclaredSymbol(node)).ExtractRecursive(TrapFile, Parent);
         }
 
+        private static Entities.AttributeKind ExtractGlobalTarget(AttributeListSyntax node) =>
+            node.Target?.Identifier.Kind() switch
+            {
+                SyntaxKind.AssemblyKeyword => Entities.AttributeKind.Assembly,
+                SyntaxKind.ModuleKeyword => Entities.AttributeKind.Module,
+                _ => throw new InternalError(node, "Unhandled global target")
+            };
+
         public override void VisitAttributeList(AttributeListSyntax node)
         {
             if (Cx.Extractor.Mode.HasFlag(ExtractorMode.Standalone))
                 return;
 
             var outputAssembly = Assembly.CreateOutputAssembly(Cx);
+            var kind = ExtractGlobalTarget(node);
             foreach (var attribute in node.Attributes)
             {
                 if (attributeLookup.Value(attribute) is AttributeData attributeData)
                 {
-                    var ae = Entities.Attribute.Create(Cx, attributeData, outputAssembly);
+                    var ae = Entities.Attribute.Create(Cx, attributeData, outputAssembly, kind);
                     Cx.BindComments(ae, attribute.GetLocation());
                 }
             }

csharp/extractor/Semmle.Extraction.CSharp/Tuples.cs

@@ -32,8 +32,8 @@ internal static void init_only_accessors(this TextWriter trapFile, Accessor acce
         internal static void array_element_type(this TextWriter trapFile, ArrayType array, int dimension, int rank, Type elementType) =>
             trapFile.WriteTuple("array_element_type", array, dimension, rank, elementType);
 
-        internal static void attributes(this TextWriter trapFile, Attribute attribute, Type attributeType, IEntity entity) =>
-            trapFile.WriteTuple("attributes", attribute, attributeType, entity);
+        internal static void attributes(this TextWriter trapFile, Attribute attribute, AttributeKind kind, Type attributeType, IEntity entity) =>
+            trapFile.WriteTuple("attributes", attribute, (int)kind, attributeType, entity);
 
         internal static void attribute_location(this TextWriter trapFile, Attribute attribute, Location location) =>
             trapFile.WriteTuple("attribute_location", attribute, location);