Fixed Security Exceptions For Requests Wihtout PPL Access

Signed-off-by: Vamsi Manohar <reddyvam@amazon.com>
opensearch-project · Aug 29, 2023 · aa650f5 · aa650f5
1 parent 627189b
commit aa650f5
Show file tree

Hide file tree

Showing 9 changed files with 71 additions and 16 deletions.
diff --git a/integ-test/src/test/java/org/opensearch/sql/ppl/ResourceMonitorIT.java b/integ-test/src/test/java/org/opensearch/sql/ppl/ResourceMonitorIT.java
@@ -31,7 +31,7 @@ public void queryExceedResourceLimitShouldFail() throws IOException {
     String query = String.format("search source=%s age=20", TEST_INDEX_DOG);
 
     ResponseException exception = expectThrows(ResponseException.class, () -> executeQuery(query));
-    assertEquals(503, exception.getResponse().getStatusLine().getStatusCode());
+    assertEquals(500, exception.getResponse().getStatusLine().getStatusCode());
     assertThat(
         exception.getMessage(),
         Matchers.containsString("resource is not enough to run the" + " query, quit."));

diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestPPLQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestPPLQueryAction.java
@@ -8,7 +8,6 @@
 import static org.opensearch.core.rest.RestStatus.BAD_REQUEST;
 import static org.opensearch.core.rest.RestStatus.INTERNAL_SERVER_ERROR;
 import static org.opensearch.core.rest.RestStatus.OK;
-import static org.opensearch.core.rest.RestStatus.SERVICE_UNAVAILABLE;
 
 import com.google.common.collect.ImmutableList;
 import java.util.Arrays;
@@ -18,6 +17,7 @@
 import java.util.function.Supplier;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.opensearch.OpenSearchSecurityException;
 import org.opensearch.client.node.NodeClient;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.rest.RestStatus;
@@ -38,6 +38,7 @@
 import org.opensearch.sql.plugin.transport.PPLQueryAction;
 import org.opensearch.sql.plugin.transport.TransportPPLQueryRequest;
 import org.opensearch.sql.plugin.transport.TransportPPLQueryResponse;
+import org.opensearch.sql.prometheus.exceptions.PrometheusClientException;
 
 public class RestPPLQueryAction extends BaseRestHandler {
   public static final String QUERY_API_ENDPOINT = "/_plugins/_ppl";
@@ -134,6 +135,11 @@ public void onFailure(Exception e) {
                       "Failed to explain the query due to error: " + e.getMessage());
                 } else if (e instanceof IllegalAccessException) {
                   reportError(channel, e, BAD_REQUEST);
+                } else if (e instanceof PrometheusClientException) {
+                  reportError(channel, e, BAD_REQUEST);
+                } else if (e instanceof OpenSearchSecurityException) {
+                  OpenSearchSecurityException exception = (OpenSearchSecurityException) e;
+                  reportError(channel, exception, exception.status());
                 } else {
                   LOG.error("Error happened during query handling", e);
                   if (isClientError(e)) {
@@ -145,7 +151,7 @@ public void onFailure(Exception e) {
                     Metrics.getInstance()
                         .getNumericalMetric(MetricName.PPL_FAILED_REQ_COUNT_SYS)
                         .increment();
-                    reportError(channel, e, SERVICE_UNAVAILABLE);
+                    reportError(channel, e, INTERNAL_SERVER_ERROR);
                   }
                 }
               }

diff --git a/prometheus/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClientImpl.java b/prometheus/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClientImpl.java
@@ -22,7 +22,9 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.json.JSONArray;
+import org.json.JSONException;
 import org.json.JSONObject;
+import org.opensearch.sql.prometheus.exceptions.PrometheusClientException;
 import org.opensearch.sql.prometheus.request.system.model.MetricMetadata;
 
 public class PrometheusClientImpl implements PrometheusClient {
@@ -111,14 +113,21 @@ private List<String> toListOfLabels(JSONArray array) {
 
   private JSONObject readResponse(Response response) throws IOException {
     if (response.isSuccessful()) {
-      JSONObject jsonObject = new JSONObject(Objects.requireNonNull(response.body()).string());
+      JSONObject jsonObject;
+      try {
+        jsonObject = new JSONObject(Objects.requireNonNull(response.body()).string());
+      } catch (JSONException jsonException) {
+        throw new PrometheusClientException(
+            "Prometheus returned unexpected body, "
+                + "please verify your prometheus server setup.");
+      }
       if ("success".equals(jsonObject.getString("status"))) {
         return jsonObject;
       } else {
-        throw new RuntimeException(jsonObject.getString("error"));
+        throw new PrometheusClientException(jsonObject.getString("error"));
       }
     } else {
-      throw new RuntimeException(
+      throw new PrometheusClientException(
           String.format(
               "Request to Prometheus is Unsuccessful with : %s",
               Objects.requireNonNull(response.body(), "Response body can't be null").string()));

diff --git a/...eus/src/main/java/org/opensearch/sql/prometheus/exceptions/PrometheusClientException.java b/...eus/src/main/java/org/opensearch/sql/prometheus/exceptions/PrometheusClientException.java
@@ -0,0 +1,15 @@
+/*
+ *
+ *  * Copyright OpenSearch Contributors
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package org.opensearch.sql.prometheus.exceptions;
+
+/** PrometheusClientException. */
+public class PrometheusClientException extends RuntimeException {
+  public PrometheusClientException(String message) {
+    super(message);
+  }
+}
diff --git a/...in/java/org/opensearch/sql/prometheus/request/system/PrometheusDescribeMetricRequest.java b/...in/java/org/opensearch/sql/prometheus/request/system/PrometheusDescribeMetricRequest.java
@@ -27,6 +27,7 @@
 import org.opensearch.sql.data.type.ExprCoreType;
 import org.opensearch.sql.data.type.ExprType;
 import org.opensearch.sql.prometheus.client.PrometheusClient;
+import org.opensearch.sql.prometheus.exceptions.PrometheusClientException;
 import org.opensearch.sql.prometheus.storage.PrometheusMetricDefaultSchema;
 
 /**
@@ -80,7 +81,7 @@ public Map<String, ExprType> getFieldTypes() {
                     "Error while fetching labels for {} from prometheus: {}",
                     metricName,
                     e.getMessage());
-                throw new RuntimeException(
+                throw new PrometheusClientException(
                     String.format(
                         "Error while fetching labels " + "for %s from prometheus: %s",
                         metricName, e.getMessage()));

diff --git a/prometheus/src/main/java/org/opensearch/sql/prometheus/storage/PrometheusStorageFactory.java b/prometheus/src/main/java/org/opensearch/sql/prometheus/storage/PrometheusStorageFactory.java
@@ -99,6 +99,7 @@ private OkHttpClient getHttpClient(Map<String, String> config) {
     OkHttpClient.Builder okHttpClient = new OkHttpClient.Builder();
     okHttpClient.callTimeout(1, TimeUnit.MINUTES);
     okHttpClient.connectTimeout(30, TimeUnit.SECONDS);
+    okHttpClient.followRedirects(false);
     if (config.get(AUTH_TYPE) != null) {
       AuthenticationType authenticationType = AuthenticationType.get(config.get(AUTH_TYPE));
       if (AuthenticationType.BASICAUTH.equals(authenticationType)) {
@@ -162,8 +163,8 @@ private void validateURI(Map<String, String> config) throws URISyntaxException {
       if (!matcher.matches()) {
         throw new IllegalArgumentException(
             String.format(
-                "Disallowed hostname in the uri: %s. Validate with %s config",
-                config.get(URI), Settings.Key.DATASOURCES_URI_ALLOWHOSTS.getKeyValue()));
+                "Disallowed hostname in the uri. Validate with %s config",
+                Settings.Key.DATASOURCES_URI_ALLOWHOSTS.getKeyValue()));
       }
     }
   }

diff --git a/prometheus/src/test/java/org/opensearch/sql/prometheus/client/PrometheusClientImplTest.java b/prometheus/src/test/java/org/opensearch/sql/prometheus/client/PrometheusClientImplTest.java
@@ -35,6 +35,7 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.junit.jupiter.MockitoExtension;
+import org.opensearch.sql.prometheus.exceptions.PrometheusClientException;
 import org.opensearch.sql.prometheus.request.system.model.MetricMetadata;
 
 @ExtendWith(MockitoExtension.class)
@@ -73,11 +74,30 @@ void testQueryRangeWith2xxStatusAndError() {
             .addHeader("Content-Type", "application/json; charset=utf-8")
             .setBody(getJson("error_response.json"));
     mockWebServer.enqueue(mockResponse);
-    RuntimeException runtimeException =
+    PrometheusClientException prometheusClientException =
         assertThrows(
-            RuntimeException.class,
+            PrometheusClientException.class,
             () -> prometheusClient.queryRange(QUERY, STARTTIME, ENDTIME, STEP));
-    assertEquals("Error", runtimeException.getMessage());
+    assertEquals("Error", prometheusClientException.getMessage());
+    RecordedRequest recordedRequest = mockWebServer.takeRequest();
+    verifyQueryRangeCall(recordedRequest);
+  }
+
+  @Test
+  @SneakyThrows
+  void testQueryRangeWithNonJsonResponse() {
+    MockResponse mockResponse =
+        new MockResponse()
+            .addHeader("Content-Type", "application/json; charset=utf-8")
+            .setBody(getJson("non_json_response.json"));
+    mockWebServer.enqueue(mockResponse);
+    PrometheusClientException prometheusClientException =
+        assertThrows(
+            PrometheusClientException.class,
+            () -> prometheusClient.queryRange(QUERY, STARTTIME, ENDTIME, STEP));
+    assertEquals(
+        "Prometheus returned unexpected body, " + "please verify your prometheus server setup.",
+        prometheusClientException.getMessage());
     RecordedRequest recordedRequest = mockWebServer.takeRequest();
     verifyQueryRangeCall(recordedRequest);
   }
@@ -90,12 +110,14 @@ void testQueryRangeWithNon2xxError() {
             .addHeader("Content-Type", "application/json; charset=utf-8")
             .setResponseCode(400);
     mockWebServer.enqueue(mockResponse);
-    RuntimeException runtimeException =
+    PrometheusClientException prometheusClientException =
         assertThrows(
-            RuntimeException.class,
+            PrometheusClientException.class,
             () -> prometheusClient.queryRange(QUERY, STARTTIME, ENDTIME, STEP));
     assertTrue(
-        runtimeException.getMessage().contains("Request to Prometheus is Unsuccessful with :"));
+        prometheusClientException
+            .getMessage()
+            .contains("Request to Prometheus is Unsuccessful with :"));
     RecordedRequest recordedRequest = mockWebServer.takeRequest();
     verifyQueryRangeCall(recordedRequest);
   }

diff --git a/...eus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusStorageFactoryTest.java b/...eus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusStorageFactoryTest.java
@@ -269,7 +269,7 @@ void createDataSourceWithHostnameNotMatchingWithAllowHostsConfig() {
         exception
             .getMessage()
             .contains(
-                "Disallowed hostname in the uri: http://localhost.com:9090. "
+                "Disallowed hostname in the uri. "
                     + "Validate with plugins.query.datasources.uri.allowhosts config"));
   }
 

diff --git a/prometheus/src/test/resources/non_json_response.json b/prometheus/src/test/resources/non_json_response.json
@@ -0,0 +1 @@
+fadsfadsfasdfasdfadsfasdfadsfadsf