Block invalid calls in Await (#2225)

Author: Quinn-With-Two-Ns

temporal-sdk/src/main/java/io/temporal/internal/sync/WorkflowInternal.java

@@ -498,13 +498,34 @@ public static <R> R executeActivity(
   public static void await(String reason, Supplier<Boolean> unblockCondition)
       throws DestroyWorkflowThreadError {
     assertNotReadOnly(reason);
-    getWorkflowOutboundInterceptor().await(reason, unblockCondition);
+    getWorkflowOutboundInterceptor()
+        .await(
+            reason,
+            () -> {
+              getRootWorkflowContext().setReadOnly(true);
+              try {
+                return unblockCondition.get();
+              } finally {
+                getRootWorkflowContext().setReadOnly(false);
+              }
+            });
   }
 
   public static boolean await(Duration timeout, String reason, Supplier<Boolean> unblockCondition)
       throws DestroyWorkflowThreadError {
     assertNotReadOnly(reason);
-    return getWorkflowOutboundInterceptor().await(timeout, reason, unblockCondition);
+    return getWorkflowOutboundInterceptor()
+        .await(
+            timeout,
+            reason,
+            () -> {
+              getRootWorkflowContext().setReadOnly(true);
+              try {
+                return unblockCondition.get();
+              } finally {
+                getRootWorkflowContext().setReadOnly(false);
+              }
+            });
   }
 
   public static <R> R sideEffect(Class<R> resultClass, Type resultType, Func<R> func) {

temporal-sdk/src/test/java/io/temporal/workflow/BadAwaitTest.java

@@ -0,0 +1,65 @@
+/*
+ * Copyright (C) 2022 Temporal Technologies, Inc. All Rights Reserved.
+ *
+ * Copyright (C) 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Modifications copyright (C) 2017 Uber Technologies, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this material except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.temporal.workflow;
+
+import static org.junit.Assert.assertThrows;
+
+import io.temporal.client.WorkflowFailedException;
+import io.temporal.testing.internal.SDKTestWorkflowRule;
+import io.temporal.worker.WorkflowImplementationOptions;
+import io.temporal.workflow.shared.TestWorkflows;
+import io.temporal.workflow.shared.TestWorkflows.TestWorkflow1;
+import org.junit.Rule;
+import org.junit.Test;
+
+public class BadAwaitTest {
+
+  @Rule
+  public SDKTestWorkflowRule testWorkflowRule =
+      SDKTestWorkflowRule.newBuilder()
+          .setWorkflowTypes(
+              WorkflowImplementationOptions.newBuilder()
+                  .setFailWorkflowExceptionTypes(RuntimeException.class)
+                  .build(),
+              TestAwait.class)
+          .build();
+
+  @Test
+  public void testBadAwait() {
+    for (String testCase : TestWorkflows.illegalCallCases) {
+      TestWorkflows.TestWorkflow1 workflowStub =
+          testWorkflowRule.newWorkflowStubTimeoutOptions(TestWorkflows.TestWorkflow1.class);
+      assertThrows(WorkflowFailedException.class, () -> workflowStub.execute(testCase));
+    }
+  }
+
+  public static class TestAwait implements TestWorkflow1 {
+    @Override
+    public String execute(String testCase) {
+      Workflow.await(
+          () -> {
+            TestWorkflows.illegalCalls(testCase);
+            return true;
+          });
+      return "fail";
+    }
+  }
+}