From ba610e5710b962cb2612a1a5d9054ec08aebd1b9 Mon Sep 17 00:00:00 2001
From: lochana-chathura <39232462+lochana-chathura@users.noreply.github.com>
Date: Fri, 5 Apr 2024 00:55:39 +0530
Subject: [PATCH] Disallow worker receive inside on-fail-clause

---
 .../compiler/semantics/analyzer/CodeAnalyzer.java          | 6 ++++--
 .../java/org/ballerinalang/test/worker/WorkerFailTest.java | 3 +++
 .../test-src/workers/send-receive-allowed-positions.bal    | 7 ++++++-
 3 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/semantics/analyzer/CodeAnalyzer.java b/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/semantics/analyzer/CodeAnalyzer.java
index ffb58a3b8aad..cbcaace973ce 100644
--- a/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/semantics/analyzer/CodeAnalyzer.java
+++ b/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/semantics/analyzer/CodeAnalyzer.java
@@ -1966,12 +1966,14 @@ private boolean isReceiveAllowedLocation(BLangFunctionBody enclInvokableBody, BL
     }
 
     private boolean isSendAllowedContext(BLangNode bLangNode) {
-        return isReceiveAllowedContext(bLangNode) || bLangNode.getKind() == NodeKind.IF;
+        return isReceiveAllowedContext(bLangNode) ||
+                bLangNode.getKind() == NodeKind.IF ||
+                bLangNode.getKind() == NodeKind.ON_FAIL;
     }
 
     private boolean isReceiveAllowedContext(BLangNode bLangNode) {
         return switch (bLangNode.getKind()) {
-            case BLOCK_FUNCTION_BODY, BLOCK, ON_FAIL, DO_STMT -> true;
+            case BLOCK_FUNCTION_BODY, BLOCK, DO_STMT -> true;
             default -> false;
         };
     }
diff --git a/tests/jballerina-unit-test/src/test/java/org/ballerinalang/test/worker/WorkerFailTest.java b/tests/jballerina-unit-test/src/test/java/org/ballerinalang/test/worker/WorkerFailTest.java
index 6d2cc3de5251..1442a0d1d61c 100644
--- a/tests/jballerina-unit-test/src/test/java/org/ballerinalang/test/worker/WorkerFailTest.java
+++ b/tests/jballerina-unit-test/src/test/java/org/ballerinalang/test/worker/WorkerFailTest.java
@@ -49,6 +49,9 @@ public void testSendReceiveAllowedSyntacticPositions() {
                 "possible deadlocks", 74, 17);
         validateError(result, index++, receiveNotAllowedError, 79, 15);
         validateError(result, index++, receiveNotAllowedError, 81, 21);
+        validateError(result, index++, receiveNotAllowedError, 93, 13);
+        validateError(result, index++, receiveNotAllowedError, 94, 16);
+        validateError(result, index++, receiveNotAllowedError, 94, 19);
         Assert.assertEquals(result.getErrorCount(), index);
     }
 
diff --git a/tests/jballerina-unit-test/src/test/resources/test-src/workers/send-receive-allowed-positions.bal b/tests/jballerina-unit-test/src/test/resources/test-src/workers/send-receive-allowed-positions.bal
index c7806fe03161..a6ef9faa5582 100644
--- a/tests/jballerina-unit-test/src/test/resources/test-src/workers/send-receive-allowed-positions.bal
+++ b/tests/jballerina-unit-test/src/test/resources/test-src/workers/send-receive-allowed-positions.bal
@@ -90,7 +90,12 @@ public function testReceiveAllowedLocations(boolean b) returns error? {
             _ = <- function; // OK
         }
     } on fail {
-        _ = <- function; // OK
+        _ = <- function; // error: position not allowed
+        _ = <- w2|function; // error: position not allowed
+    }
+
+    worker w2 {
+        1 -> w1;
     }
 
     1 -> w1;