From fb9ec95cf0c2ed6e3b5273b2274eac70dd8bd3d4 Mon Sep 17 00:00:00 2001
From: Peter Klausler <pklausler@nvidia.com>
Date: Thu, 28 Apr 2022 14:23:31 -0700
Subject: [PATCH] [flang][runtime] Enforce restrictions on unlimited format
 repetition

A repeated format item group with an unlimited ('*') repetition count
can appear only as the last item at the top level of a format; it can't
be nested in more parentheses and it can't be followed by anything
else.

Differential Revision: https://reviews.llvm.org/D125054
---
 flang/runtime/format-implementation.h | 17 ++++++++++++++++-
 flang/unittests/Runtime/Format.cpp    |  3 +--
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/flang/runtime/format-implementation.h b/flang/runtime/format-implementation.h
index 11f3ad7d18e4..ab2fad3188bd 100644
--- a/flang/runtime/format-implementation.h
+++ b/flang/runtime/format-implementation.h
@@ -207,6 +207,13 @@ int FormatControl<CONTEXT>::CueUpNextDataEdit(Context &context, bool stop) {
             maybeReversionPoint);
         return 0;
       }
+      if (height_ != 1) {
+        ReportBadFormat(context,
+            "Invalid FORMAT: '*' must be nested in exactly one set of "
+            "parentheses",
+            maybeReversionPoint);
+        return 0;
+      }
     }
     ch = Capitalize(ch);
     if (ch == '(') {
@@ -251,12 +258,20 @@ int FormatControl<CONTEXT>::CueUpNextDataEdit(Context &context, bool stop) {
         ++restart;
       }
       if (stack_[height_ - 1].remaining == Iteration::unlimited) {
-        offset_ = restart;
+        if (height_ > 1 && GetNextChar(context) != ')') {
+          ReportBadFormat(context,
+              "Unlimited repetition in FORMAT may not be followed by more "
+              "items",
+              restart);
+          return 0;
+        }
         if (offset_ == unlimitedLoopCheck) {
           ReportBadFormat(context,
               "Unlimited repetition in FORMAT lacks data edit descriptors",
               restart);
+          return 0;
         }
+        offset_ = restart;
       } else if (stack_[height_ - 1].remaining-- > 0) {
         offset_ = restart;
       } else {
diff --git a/flang/unittests/Runtime/Format.cpp b/flang/unittests/Runtime/Format.cpp
index 5391c48fdf87..963820d40efc 100644
--- a/flang/unittests/Runtime/Format.cpp
+++ b/flang/unittests/Runtime/Format.cpp
@@ -107,8 +107,7 @@ TEST(FormatTests, FormatStringTraversal) {
           ResultsTy{"'PI='", "F9.7", "'PI='", "F9.7", "'done'"}, 1},
       {2, "(3('PI=',F9.7,:),'tooFar')",
           ResultsTy{"'PI='", "F9.7", "'PI='", "F9.7"}, 1},
-      {2, "(*('PI=',F9.7,:),'tooFar')",
-          ResultsTy{"'PI='", "F9.7", "'PI='", "F9.7"}, 1},
+      {2, "(*('PI=',F9.7,:))", ResultsTy{"'PI='", "F9.7", "'PI='", "F9.7"}, 1},
       {1, "(3F9.7)", ResultsTy{"2*F9.7"}, 2},
   };