[clang][bytecode] Check array sizes against step limit (llvm#137679)

Author: tbaederr

clang/lib/AST/ByteCode/Compiler.cpp

@@ -1862,6 +1862,13 @@ bool Compiler<Emitter>::visitInitList(ArrayRef<const Expr *> Inits,
     if (Inits.size() == 1 && QT == Inits[0]->getType())
       return this->delegate(Inits[0]);
 
+    const ConstantArrayType *CAT =
+        Ctx.getASTContext().getAsConstantArrayType(QT);
+    uint64_t NumElems = CAT->getZExtSize();
+
+    if (!this->emitCheckArraySize(NumElems, E))
+      return false;
+
     unsigned ElementIndex = 0;
     for (const Expr *Init : Inits) {
       if (const auto *EmbedS =
@@ -1890,10 +1897,6 @@ bool Compiler<Emitter>::visitInitList(ArrayRef<const Expr *> Inits,
     // Expand the filler expression.
     // FIXME: This should go away.
     if (ArrayFiller) {
-      const ConstantArrayType *CAT =
-          Ctx.getASTContext().getAsConstantArrayType(QT);
-      uint64_t NumElems = CAT->getZExtSize();
-
       for (; ElementIndex != NumElems; ++ElementIndex) {
         if (!this->visitArrayElemInit(ElementIndex, ArrayFiller))
           return false;

clang/lib/AST/ByteCode/Interp.h

@@ -175,6 +175,8 @@ bool handleFixedPointOverflow(InterpState &S, CodePtr OpPC,
 
 bool isConstexprUnknown(const Pointer &P);
 
+inline bool CheckArraySize(InterpState &S, CodePtr OpPC, uint64_t NumElems);
+
 enum class ShiftDir { Left, Right };
 
 /// Checks if the shift operation is legal.
@@ -3110,6 +3112,9 @@ inline bool AllocN(InterpState &S, CodePtr OpPC, PrimType T, const Expr *Source,
   }
   assert(NumElements.isPositive());
 
+  if (!CheckArraySize(S, OpPC, static_cast<uint64_t>(NumElements)))
+    return false;
+
   DynamicAllocator &Allocator = S.getAllocator();
   Block *B =
       Allocator.allocate(Source, T, static_cast<size_t>(NumElements),
@@ -3140,6 +3145,9 @@ inline bool AllocCN(InterpState &S, CodePtr OpPC, const Descriptor *ElementDesc,
   }
   assert(NumElements.isPositive());
 
+  if (!CheckArraySize(S, OpPC, static_cast<uint64_t>(NumElements)))
+    return false;
+
   DynamicAllocator &Allocator = S.getAllocator();
   Block *B =
       Allocator.allocate(ElementDesc, static_cast<size_t>(NumElements),
@@ -3246,6 +3254,17 @@ inline bool CheckDestruction(InterpState &S, CodePtr OpPC) {
   return CheckDestructor(S, OpPC, Ptr);
 }
 
+inline bool CheckArraySize(InterpState &S, CodePtr OpPC, uint64_t NumElems) {
+  uint64_t Limit = S.getLangOpts().ConstexprStepLimit;
+  if (NumElems > Limit) {
+    S.FFDiag(S.Current->getSource(OpPC),
+             diag::note_constexpr_new_exceeds_limits)
+        << NumElems << Limit;
+    return false;
+  }
+  return true;
+}
+
 //===----------------------------------------------------------------------===//
 // Read opcode arguments
 //===----------------------------------------------------------------------===//

clang/lib/AST/ByteCode/InterpBuiltin.cpp

@@ -1533,6 +1533,9 @@ static bool interp__builtin_operator_new(InterpState &S, CodePtr OpPC,
     return false;
   }
 
+  if (!CheckArraySize(S, OpPC, NumElems.getZExtValue()))
+    return false;
+
   bool IsArray = NumElems.ugt(1);
   std::optional<PrimType> ElemT = S.getContext().classify(ElemType);
   DynamicAllocator &Allocator = S.getAllocator();

clang/lib/AST/ByteCode/Opcodes.td

@@ -414,6 +414,8 @@ def CheckLiteralType : Opcode {
   let Args = [ArgTypePtr];
 }
 
+def CheckArraySize : Opcode { let Args = [ArgUint64]; }
+
 // [] -> [Value]
 def GetGlobal : AccessOpcode;
 def GetGlobalUnchecked : AccessOpcode;

clang/test/AST/ByteCode/dynalloc-limits.cpp

@@ -0,0 +1,73 @@
+// RUN: %clang_cc1 -std=c++20 -verify=ref,both      -fconstexpr-steps=1024 -Wvla %s
+// RUN: %clang_cc1 -std=c++20 -verify=both,both -fconstexpr-steps=1024 -Wvla %s -fexperimental-new-constant-interpreter
+
+
+
+
+namespace std {
+  using size_t = decltype(sizeof(0));
+}
+
+void *operator new(std::size_t, void *p) { return p; }
+
+namespace std {
+  template<typename T> struct allocator {
+    constexpr T *allocate(size_t N) {
+      return (T*)operator new(sizeof(T) * N); // #alloc
+    }
+    constexpr void deallocate(void *p) {
+      operator delete(p);
+    }
+  };
+  template<typename T, typename ...Args>
+  constexpr void construct_at(void *p, Args &&...args) { // #construct
+    new (p) T((Args&&)args...);
+  }
+}
+
+template <typename T>
+struct S {
+    constexpr S(unsigned long long N)
+    : data(nullptr){
+        data = alloc.allocate(N);  // #call
+        for(std::size_t i = 0; i < N; i ++)
+            std::construct_at<T>(data + i, i); // #construct_call
+    }
+    constexpr T operator[](std::size_t i) const {
+      return data[i];
+    }
+
+    constexpr ~S() {
+        alloc.deallocate(data);
+    }
+    std::allocator<T> alloc;
+    T* data;
+};
+
+#if __LP64__
+constexpr std::size_t s = S<std::size_t>(~0UL)[42]; // both-error {{constexpr variable 's' must be initialized by a constant expression}} \
+                                           // both-note-re@#call {{in call to 'this->alloc.allocate({{.*}})'}} \
+                                           // both-note-re@#alloc {{cannot allocate array; evaluated array bound {{.*}} is too large}} \
+                                           // both-note-re {{in call to 'S({{.*}})'}}
+#endif
+
+constexpr std::size_t ssmall = S<std::size_t>(100)[42];
+
+constexpr std::size_t s5 = S<std::size_t>(1025)[42]; // both-error {{constexpr variable 's5' must be initialized by a constant expression}} \
+                                   // both-note@#alloc {{cannot allocate array; evaluated array bound 1025 exceeds the limit (1024); use '-fconstexpr-steps' to increase this limit}} \
+                                   // both-note@#call {{in call to 'this->alloc.allocate(1025)'}} \
+                                   // both-note {{in call}}
+
+
+
+template <auto N>
+constexpr int stack_array() {
+    [[maybe_unused]] char BIG[N] = {1};  // both-note {{cannot allocate array; evaluated array bound 1025 exceeds the limit (1024)}}
+    return BIG[N-1];
+}
+
+int c = stack_array<1024>();
+int d = stack_array<1025>();
+constexpr int e = stack_array<1024>();
+constexpr int f = stack_array<1025>(); // both-error {{constexpr variable 'f' must be initialized by a constant expression}} \
+                                       // both-note {{in call}}