[clang] [Clang] Fix 'counted_by' for nested struct pointers (PR #110437)

[Jan Hendrik Farr via cfe-commits]

https://github.com/Cydox created https://github.com/llvm/llvm-project/pull/110437

Fixes #110385

Fix counted_by attribute for cases where the flexible array member is accessed through struct pointer inside another struct:

struct variable {
        int a;
        int b;
        int length;
        short array[] __attribute__((counted_by(length)));
};

struct bucket {
        int a;
        struct variable *growable;
        int b;
};

__builtin_dynamic_object_size(p->growable->array, 0);

This commit makes sure that if the StructBase is both a MemberExpr and a pointer, it is treated as a pointer. Otherwise clang will generate to code to access the address of p->growable intead of loading the value of p->growable->length.

>From 26aaa04f5570d05f15e0ac9882870467b89c2690 Mon Sep 17 00:00:00 2001
From: Jan Hendrik Farr <kernel at jfarr.cc>
Date: Sun, 29 Sep 2024 21:38:13 +0200
Subject: [PATCH] [Clang] Fix 'counted_by' for nested struct pointers

Fixes #110385

Fix counted_by attribute for cases where the flexible array member is accessed
through struct pointer inside another struct:

struct variable {
        int a;
        int b;
        int length;
        short array[] __attribute__((counted_by(length)));
};

struct bucket {
        int a;
        struct variable *growable;
        int b;
};

__builtin_dynamic_object_size(p->growable->array, 0);

This commit makes sure that if the StructBase is both a MemberExpr and a
pointer, it is treated as a pointer. Otherwise clang will generate to
code to access the address of p->growable intead of loading the value of
p->growable->length.
---
 clang/lib/CodeGen/CGExpr.cpp                  |  8 ++--
 clang/test/CodeGen/attr-counted-by-pr110385.c | 40 +++++++++++++++++++
 2 files changed, 44 insertions(+), 4 deletions(-)
 create mode 100644 clang/test/CodeGen/attr-counted-by-pr110385.c

diff --git a/clang/lib/CodeGen/CGExpr.cpp b/clang/lib/CodeGen/CGExpr.cpp
index df4994ba9af6e1..2875cf18d4f6c9 100644
--- a/clang/lib/CodeGen/CGExpr.cpp
+++ b/clang/lib/CodeGen/CGExpr.cpp
@@ -1165,15 +1165,15 @@ llvm::Value *CodeGenFunction::EmitLoadOfCountedByField(
     Res = EmitDeclRefLValue(DRE).getPointer(*this);
     Res = Builder.CreateAlignedLoad(ConvertType(DRE->getType()), Res,
                                     getPointerAlign(), "dre.load");
-  } else if (const MemberExpr *ME = dyn_cast<MemberExpr>(StructBase)) {
-    LValue LV = EmitMemberExpr(ME);
-    Address Addr = LV.getAddress();
-    Res = Addr.emitRawPointer(*this);
   } else if (StructBase->getType()->isPointerType()) {
     LValueBaseInfo BaseInfo;
     TBAAAccessInfo TBAAInfo;
     Address Addr = EmitPointerWithAlignment(StructBase, &BaseInfo, &TBAAInfo);
     Res = Addr.emitRawPointer(*this);
+  } else if (const MemberExpr *ME = dyn_cast<MemberExpr>(StructBase)) {
+    LValue LV = EmitMemberExpr(ME);
+    Address Addr = LV.getAddress();
+    Res = Addr.emitRawPointer(*this);
   } else {
     return nullptr;
   }
diff --git a/clang/test/CodeGen/attr-counted-by-pr110385.c b/clang/test/CodeGen/attr-counted-by-pr110385.c
new file mode 100644
index 00000000000000..49a08c5965ef94
--- /dev/null
+++ b/clang/test/CodeGen/attr-counted-by-pr110385.c
@@ -0,0 +1,40 @@
+// NOTE: Assertions have been autogenerated by utils/update_cc_test_checks.py UTC_ARGS: --version 4
+// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -O2 -Wno-missing-declarations -emit-llvm -o - %s | FileCheck %s
+
+// See #110385
+// Based on reproducer from Kees Cook:
+// https://lore.kernel.org/all/202409170436.C3C6E7F7A@keescook/
+
+struct variable {
+        int a;
+        int b;
+        int length;
+        short array[] __attribute__((counted_by(length)));
+};
+
+struct bucket {
+        int a;
+        struct variable *growable;
+        int b;
+};
+
+void init(void * __attribute__((pass_dynamic_object_size(0))));
+
+// CHECK-LABEL: define dso_local void @test1(
+// CHECK-SAME: ptr nocapture noundef readonly [[FOO:%.*]]) local_unnamed_addr #[[ATTR0:[0-9]+]] {
+// CHECK-NEXT:  entry:
+// CHECK-NEXT:    [[GROWABLE:%.*]] = getelementptr inbounds nuw i8, ptr [[FOO]], i64 8
+// CHECK-NEXT:    [[TMP0:%.*]] = load ptr, ptr [[GROWABLE]], align 8, !tbaa [[TBAA2:![0-9]+]]
+// CHECK-NEXT:    [[ARRAY:%.*]] = getelementptr inbounds nuw i8, ptr [[TMP0]], i64 12
+// CHECK-NEXT:    [[DOT_COUNTED_BY_GEP:%.*]] = getelementptr inbounds i8, ptr [[TMP0]], i64 8
+// CHECK-NEXT:    [[DOT_COUNTED_BY_LOAD:%.*]] = load i32, ptr [[DOT_COUNTED_BY_GEP]], align 4
+// CHECK-NEXT:    [[TMP1:%.*]] = sext i32 [[DOT_COUNTED_BY_LOAD]] to i64
+// CHECK-NEXT:    [[TMP2:%.*]] = shl nsw i64 [[TMP1]], 1
+// CHECK-NEXT:    [[TMP3:%.*]] = icmp sgt i32 [[DOT_COUNTED_BY_LOAD]], -1
+// CHECK-NEXT:    [[TMP4:%.*]] = select i1 [[TMP3]], i64 [[TMP2]], i64 0
+// CHECK-NEXT:    tail call void @init(ptr noundef nonnull [[ARRAY]], i64 noundef [[TMP4]]) #[[ATTR2:[0-9]+]]
+// CHECK-NEXT:    ret void
+//
+void test1(struct bucket *foo) {
+        init(foo->growable->array);
+}