improve the overall logic and fix several bugs based on comments

Signed-off-by: Chenyang Ji <[email protected]>
opensearch-project · Jun 4, 2024 · ff0eb39 · ff0eb39
1 parent deab414
commit ff0eb39
Show file tree

Hide file tree

Showing 8 changed files with 26 additions and 33 deletions.
diff --git a/...hts/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java b/...hts/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java
@@ -139,7 +139,7 @@ public void onRequestEnd(final SearchPhaseContext context, final SearchRequestCo
             attributes.put(Attribute.TOTAL_SHARDS, context.getNumShards());
             attributes.put(Attribute.INDICES, request.indices());
             attributes.put(Attribute.PHASE_LATENCY_MAP, searchRequestContext.phaseTookMap());
-            attributes.put(Attribute.USER_NAME, request.source().labels().get(DefaultUserInfoLabelingRule.USER_NAME));
+            attributes.put(Attribute.LABELS, request.source().labels());
             SearchQueryRecord record = new SearchQueryRecord(request.getOrCreateAbsoluteStartMillis(), measurements, attributes);
             queryInsightsService.addRecord(record);
         } catch (Exception e) {

diff --git a/...ns/query-insights/src/main/java/org/opensearch/plugin/insights/rules/model/Attribute.java b/...ns/query-insights/src/main/java/org/opensearch/plugin/insights/rules/model/Attribute.java
@@ -45,13 +45,9 @@ public enum Attribute {
      */
     NODE_ID,
     /**
-     * User associated with this request
+     * Custom labels
      */
-    USER_NAME,
-    /**
-     * Custom tenant tags
-     */
-    CUSTOMIZED_TAG;
+    LABELS;
 
     /**
      * Read an Attribute from a StreamInput

diff --git a/server/src/main/java/org/opensearch/action/search/SearchRequestOperationsListener.java b/server/src/main/java/org/opensearch/action/search/SearchRequestOperationsListener.java
@@ -41,11 +41,11 @@ protected SearchRequestOperationsListener(final boolean enabled) {
         this.enabled = enabled;
     }
 
-    protected abstract void onPhaseStart(SearchPhaseContext context);
+    protected void onPhaseStart(SearchPhaseContext context) {};
 
-    protected abstract void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext);
+    protected void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext) {};
 
-    protected abstract void onPhaseFailure(SearchPhaseContext context, Throwable cause);
+    protected void onPhaseFailure(SearchPhaseContext context, Throwable cause) {};
 
     protected void onRequestStart(SearchRequestContext searchRequestContext) {}
 

diff --git a/server/src/main/java/org/opensearch/search/builder/SearchSourceBuilder.java b/server/src/main/java/org/opensearch/search/builder/SearchSourceBuilder.java
@@ -1141,14 +1141,6 @@ public Map<String, Object> labels() {
         return labels;
     }
 
-    /**
-     * Define labels within this search request.
-     */
-    public SearchSourceBuilder labels(Map<String, Object> labels) {
-        this.labels = labels;
-        return this;
-    }
-
     /**
      * Add labels within this search request.
      */

diff --git a/server/src/main/java/org/opensearch/search/labels/RuleBasedLabelingService.java b/server/src/main/java/org/opensearch/search/labels/RuleBasedLabelingService.java
@@ -47,6 +47,12 @@ public void applyAllRules(final ThreadContext threadContext, final SearchRequest
             .map(rule -> rule.evaluate(threadContext, searchRequest))
             .flatMap(m -> m.entrySet().stream())
             .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+        // Handling potential spoofing by checking if any conflicts exist between user-supplied labels and the computed labels
+        for (String key : searchRequest.source().labels().keySet()) {
+            if (labels.containsKey(key)) {
+                throw new IllegalArgumentException("Unexpected label found: " + key);
+            }
+        }
         searchRequest.source().addLabels(labels);
     }
 }
diff --git a/server/src/main/java/org/opensearch/search/labels/SearchRequestLabelingListener.java b/server/src/main/java/org/opensearch/search/labels/SearchRequestLabelingListener.java
@@ -27,15 +27,6 @@ public SearchRequestLabelingListener(final ThreadPool threadPool, final RuleBase
         this.ruleBasedLabelingService = ruleBasedLabelingService;
     }
 
-    @Override
-    protected void onPhaseStart(SearchPhaseContext context) {}
-
-    @Override
-    protected void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext) {}
-
-    @Override
-    protected void onPhaseFailure(SearchPhaseContext context, Throwable cause) {}
-
     @Override
     public void onRequestStart(SearchRequestContext searchRequestContext) {
         // add tags to search request

diff --git a/server/src/main/java/org/opensearch/search/labels/rules/DefaultUserInfoLabelingRule.java b/server/src/main/java/org/opensearch/search/labels/rules/DefaultUserInfoLabelingRule.java
@@ -17,7 +17,7 @@
 import java.util.Map;
 
 /**
- * Rules to get user info labels, specifically, the info is injected by the security plugin.
+ * Rules to get user info labels, specifically, the info that is injected by the security plugin.
  */
 public class DefaultUserInfoLabelingRule implements Rule {
     /**
@@ -56,16 +56,16 @@ private Map<String, Object> getUserInfoFromThreadContext(ThreadContext threadCon
         if (threadContext == null) {
             return userInfoMap;
         }
-        Object userInfoObj = threadContext.getTransient(REQUEST_HEADER_USER_INFO);
-        if (userInfoObj == null) {
-            return userInfoMap;
-        }
-        String userInfoStr = userInfoObj.toString();
         Object remoteAddressObj = threadContext.getTransient(REQUEST_HEADER_REMOTE_ADDRESS);
         if (remoteAddressObj != null) {
             userInfoMap.put(REMOTE_ADDRESS, remoteAddressObj.toString());
         }
 
+        Object userInfoObj = threadContext.getTransient(REQUEST_HEADER_USER_INFO);
+        if (userInfoObj == null) {
+            return userInfoMap;
+        }
+        String userInfoStr = userInfoObj.toString();
         String[] userInfo = userInfoStr.split("\\|");
         if ((userInfo.length == 0) || (Strings.isNullOrEmpty(userInfo[0]))) {
             return userInfoMap;

diff --git a/server/src/test/java/org/opensearch/search/labels/DefaultUserInfoLabelingRuleTests.java b/server/src/test/java/org/opensearch/search/labels/DefaultUserInfoLabelingRuleTests.java
@@ -44,6 +44,14 @@ public void testGetUserInfoFromThreadContext() {
         assertEquals(expectedUserInfoMap, actualUserInfoMap);
     }
 
+    public void testGetPartialInfoFromThreadContext() {
+        threadContext.putTransient("_opendistro_security_remote_address", "127.0.0.1");
+        Map<String, Object> expectedUserInfoMap = new HashMap<>();
+        expectedUserInfoMap.put(DefaultUserInfoLabelingRule.REMOTE_ADDRESS, "127.0.0.1");
+        Map<String, Object> actualUserInfoMap = defaultUserInfoLabelingRule.evaluate(threadContext, searchRequest);
+        assertEquals(expectedUserInfoMap, actualUserInfoMap);
+    }
+
     public void testGetUserInfoFromThreadContext_EmptyUserInfo() {
         Map<String, Object> actualUserInfoMap = defaultUserInfoLabelingRule.evaluate(threadContext, searchRequest);
         assertTrue(actualUserInfoMap.isEmpty());