diff --git a/lib/main.js b/lib/main.js
index 90d9206a4..3862cd81a 100644
--- a/lib/main.js
+++ b/lib/main.js
@@ -96,6 +96,20 @@ function main() {
             var tenantId = core.getInput('tenant-id', { required: false });
             var subscriptionId = core.getInput('subscription-id', { required: false });
             var resourceManagerEndpointUrl = "https://management.azure.com/";
+            switch(environment){
+                case 'azurecloud':
+                    resourceManagerEndpointUrl = "https://management.azure.com/";
+                    break;
+                case 'azureusgovernment':
+                    resourceManagerEndpointUrl = "https://management.usgovcloudapi.net/";
+                    break;
+                case 'azurechinacloud':
+                    resourceManagerEndpointUrl = "https://management.chinacloudapi.cn/";
+                    break;
+                default:
+                    resourceManagerEndpointUrl = "https://management.azure.com/";
+                    break;
+            }
             var enableOIDC = true;
             var federatedToken = null;
             // If any of the individual credentials (clent_id, tenat_id, subscription_id) is present.
@@ -137,7 +151,7 @@ function main() {
                     let audience = core.getInput('audience', { required: false });
                     federatedToken = yield core.getIDToken(audience);
                     if (!!federatedToken) {
-                        if (environment != "azurecloud")
+                        if (environment != "azurecloud" || "azureusgovernment" || "azurechinacloud")
                             throw new Error(`Your current environment - "${environment}" is not supported for OIDC login.`);
                         let [issuer, subjectClaim] = yield jwtParser(federatedToken);
                         console.log("Federated token details: \n issuer - " + issuer + " \n subject claim - " + subjectClaim);
diff --git a/src/main.ts b/src/main.ts
index 13ebea0cb..4612d3763 100644
--- a/src/main.ts
+++ b/src/main.ts
@@ -71,6 +71,20 @@ async function main() {
         var tenantId = core.getInput('tenant-id', { required: false });
         var subscriptionId = core.getInput('subscription-id', { required: false });
         var resourceManagerEndpointUrl = "https://management.azure.com/";
+        switch(environment){
+            case 'azurecloud':
+                resourceManagerEndpointUrl = "https://management.azure.com/";
+                break;
+            case 'azureusgovernment':
+                resourceManagerEndpointUrl = "https://management.usgovcloudapi.net/";
+                break;
+            case 'azurechinacloud':
+                resourceManagerEndpointUrl = "https://management.chinacloudapi.cn/";
+                break;
+            default:
+                resourceManagerEndpointUrl = "https://management.azure.com/";
+                break;
+        }
         var enableOIDC = true;
         var federatedToken = null;
 
@@ -115,7 +129,7 @@ async function main() {
                 let audience = core.getInput('audience', { required: false });
                 federatedToken = await core.getIDToken(audience);
                 if (!!federatedToken) {
-                    if (environment != "azurecloud")
+                    if (environment != "azurecloud" || "azureusgovernment" || "azurechinacloud")
                         throw new Error(`Your current environment - "${environment}" is not supported for OIDC login.`);
                     let [issuer, subjectClaim] = await jwtParser(federatedToken);
                     console.log("Federated token details: \n issuer - " + issuer + " \n subject claim - " + subjectClaim);