Adding dns lookup

jonathan-buttner · jonathan-buttner · commit 04e1a648a945 · 2025-02-28T16:36:38.000-05:00
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java
@@ -277,10 +277,7 @@ public Collection<?> createComponents(PluginServices services) {
         var inferenceServiceSettings = new ElasticInferenceServiceSettings(settings);
         inferenceServiceSettings.init(services.clusterService());
 
-        var authorizationHandler = new ElasticInferenceServiceAuthorizationRequestHandler(
-            inferenceServiceSettings.getElasticInferenceServiceUrl(),
-            services.threadPool()
-        );
+        var authorizationHandler = new ElasticInferenceServiceAuthorizationRequestHandler(inferenceServiceSettings, services.threadPool());
 
         inferenceServices.add(
             () -> List.of(
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceSettings.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceSettings.java
@@ -33,6 +33,16 @@ public class ElasticInferenceServiceSettings {
         Setting.Property.NodeScope
     );
 
+    /**
+     * This controls whether the authorization logic will look up the ip address of the gateway.
+     * This should only be enabled for testing and debugging purposes.
+     */
+    static final Setting<Boolean> AUTH_DEBUGGING_DNS_LOOKUP_ENABLED = Setting.boolSetting(
+        "xpack.inference.elastic.authorization_debugging_dns_lookup_enabled",
+        false,
+        Setting.Property.NodeScope
+    );
+
     /**
      * This setting is for testing only. It controls whether authorization is only performed once at bootup. If set to true, an
      * authorization request will be made repeatedly on an interval.
@@ -75,6 +85,7 @@ public class ElasticInferenceServiceSettings {
 
     private final String elasticInferenceServiceUrl;
     private final boolean periodicAuthorizationEnabled;
+    private final boolean authDebuggingDnsLookupEnabled;
     private volatile TimeValue authRequestInterval;
     private volatile TimeValue maxAuthorizationRequestJitter;
 
@@ -84,6 +95,7 @@ public ElasticInferenceServiceSettings(Settings settings) {
         periodicAuthorizationEnabled = PERIODIC_AUTHORIZATION_ENABLED.get(settings);
         authRequestInterval = AUTHORIZATION_REQUEST_INTERVAL.get(settings);
         maxAuthorizationRequestJitter = MAX_AUTHORIZATION_REQUEST_JITTER.get(settings);
+        authDebuggingDnsLookupEnabled = AUTH_DEBUGGING_DNS_LOOKUP_ENABLED.get(settings);
     }
 
     /**
@@ -115,6 +127,14 @@ public TimeValue getMaxAuthorizationRequestJitter() {
         return maxAuthorizationRequestJitter;
     }
 
+    public boolean isAuthDebuggingDnsLookupEnabled() {
+        return authDebuggingDnsLookupEnabled;
+    }
+
+    public boolean isPeriodicAuthorizationEnabled() {
+        return periodicAuthorizationEnabled;
+    }
+
     public static List<Setting<?>> getSettingsDefinitions() {
         ArrayList<Setting<?>> settings = new ArrayList<>();
         settings.add(EIS_GATEWAY_URL);
@@ -124,14 +144,11 @@ public static List<Setting<?>> getSettingsDefinitions() {
         settings.add(PERIODIC_AUTHORIZATION_ENABLED);
         settings.add(AUTHORIZATION_REQUEST_INTERVAL);
         settings.add(MAX_AUTHORIZATION_REQUEST_JITTER);
+        settings.add(AUTH_DEBUGGING_DNS_LOOKUP_ENABLED);
         return settings;
     }
 
     public String getElasticInferenceServiceUrl() {
         return Strings.isEmpty(elasticInferenceServiceUrl) ? eisGatewayUrl : elasticInferenceServiceUrl;
     }
-
-    public boolean isPeriodicAuthorizationEnabled() {
-        return periodicAuthorizationEnabled;
-    }
 }
diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/authorization/ElasticInferenceServiceAuthorizationRequestHandler.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/authorization/ElasticInferenceServiceAuthorizationRequestHandler.java
@@ -12,7 +12,6 @@
 import org.elasticsearch.ElasticsearchWrapperException;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.common.Strings;
-import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.inference.InferenceServiceResults;
 import org.elasticsearch.tasks.Task;
@@ -22,8 +21,11 @@
 import org.elasticsearch.xpack.inference.external.http.sender.Sender;
 import org.elasticsearch.xpack.inference.external.request.elastic.ElasticInferenceServiceAuthorizationRequest;
 import org.elasticsearch.xpack.inference.external.response.elastic.ElasticInferenceServiceAuthorizationResponseEntity;
+import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceSettings;
 import org.elasticsearch.xpack.inference.telemetry.TraceContext;
 
+import java.net.InetAddress;
+import java.net.URI;
 import java.util.Locale;
 import java.util.Objects;
 import java.util.concurrent.CountDownLatch;
@@ -48,20 +50,25 @@ private static ResponseHandler createAuthResponseHandler() {
         );
     }
 
-    private final String baseUrl;
+    private final ElasticInferenceServiceSettings elasticInferenceServiceSettings;
     private final ThreadPool threadPool;
     private final Logger logger;
     private final CountDownLatch requestCompleteLatch = new CountDownLatch(1);
 
-    public ElasticInferenceServiceAuthorizationRequestHandler(@Nullable String baseUrl, ThreadPool threadPool) {
-        this.baseUrl = baseUrl;
-        this.threadPool = Objects.requireNonNull(threadPool);
-        logger = LogManager.getLogger(ElasticInferenceServiceAuthorizationRequestHandler.class);
+    public ElasticInferenceServiceAuthorizationRequestHandler(
+        ElasticInferenceServiceSettings elasticInferenceServiceSettings,
+        ThreadPool threadPool
+    ) {
+        this(elasticInferenceServiceSettings, threadPool, LogManager.getLogger(ElasticInferenceServiceAuthorizationRequestHandler.class));
     }
 
     // only use for testing
-    ElasticInferenceServiceAuthorizationRequestHandler(@Nullable String baseUrl, ThreadPool threadPool, Logger logger) {
-        this.baseUrl = baseUrl;
+    ElasticInferenceServiceAuthorizationRequestHandler(
+        ElasticInferenceServiceSettings elasticInferenceServiceSettings,
+        ThreadPool threadPool,
+        Logger logger
+    ) {
+        this.elasticInferenceServiceSettings = Objects.requireNonNull(elasticInferenceServiceSettings);
         this.threadPool = Objects.requireNonNull(threadPool);
         this.logger = Objects.requireNonNull(logger);
     }
@@ -75,7 +82,7 @@ public void getAuthorization(ActionListener<ElasticInferenceServiceAuthorization
         try {
             logger.debug("Retrieving authorization information from the Elastic Inference Service.");
 
-            if (Strings.isNullOrEmpty(baseUrl)) {
+            if (Strings.isNullOrEmpty(elasticInferenceServiceSettings.getElasticInferenceServiceUrl())) {
                 logger.debug("The base URL for the authorization service is not valid, rejecting authorization.");
                 listener.onResponse(ElasticInferenceServiceAuthorizationModel.newDisabledService());
                 return;
@@ -108,7 +115,28 @@ public void getAuthorization(ActionListener<ElasticInferenceServiceAuthorization
                 requestCompleteLatch.countDown();
             });
 
-            var request = new ElasticInferenceServiceAuthorizationRequest(baseUrl, getCurrentTraceInfo());
+            if (logger.isDebugEnabled() && elasticInferenceServiceSettings.isAuthDebuggingDnsLookupEnabled()) {
+                logger.debug("Attempting to look up authorization gateway host info");
+                try {
+                    var baseUri = new URI(elasticInferenceServiceSettings.getElasticInferenceServiceUrl());
+                    logger.debug(
+                        Strings.format(
+                            "Looking up authorization gateway ip for url: %s, host: %s",
+                            elasticInferenceServiceSettings.getElasticInferenceServiceUrl(),
+                            baseUri.getHost()
+                        )
+                    );
+                    var gatewayAddress = InetAddress.getByName(baseUri.getHost());
+                    logger.debug(Strings.format("Gateway address: %s", gatewayAddress.getHostAddress()));
+                } catch (Exception e) {
+                    logger.debug("Failed to resolve gateway address", e);
+                }
+            }
+
+            var request = new ElasticInferenceServiceAuthorizationRequest(
+                elasticInferenceServiceSettings.getElasticInferenceServiceUrl(),
+                getCurrentTraceInfo()
+            );
 
             sender.sendWithoutQueuing(logger, request, AUTH_RESPONSE_HANDLER, DEFAULT_AUTH_TIMEOUT, newListener);
         } catch (Exception e) {
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceSettingsTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceSettingsTests.java
@@ -42,6 +42,15 @@ public static ElasticInferenceServiceSettings create(
         return new ElasticInferenceServiceSettings(settings);
     }
 
+    public static ElasticInferenceServiceSettings createWithUrlAndDnsLookup(String elasticInferenceServiceUrl) {
+        var settings = Settings.builder()
+            .put(ElasticInferenceServiceSettings.ELASTIC_INFERENCE_SERVICE_URL.getKey(), elasticInferenceServiceUrl)
+            .put(ElasticInferenceServiceSettings.AUTH_DEBUGGING_DNS_LOOKUP_ENABLED.getKey(), true)
+            .build();
+
+        return new ElasticInferenceServiceSettings(settings);
+    }
+
     public void testGetElasticInferenceServiceUrl_WithUrlSetting() {
         var settings = Settings.builder()
             .put(ElasticInferenceServiceSettings.ELASTIC_INFERENCE_SERVICE_URL.getKey(), ELASTIC_INFERENCE_SERVICE_URL)
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceTests.java
@@ -1163,7 +1163,7 @@ private ElasticInferenceService createServiceWithAuthHandler(HttpRequestSender.F
             createWithEmptySettings(threadPool),
             ElasticInferenceServiceSettingsTests.create(eisGatewayUrl),
             mockModelRegistry(),
-            new ElasticInferenceServiceAuthorizationRequestHandler(eisGatewayUrl, threadPool)
+            new ElasticInferenceServiceAuthorizationRequestHandler(ElasticInferenceServiceSettingsTests.create(eisGatewayUrl), threadPool)
         );
     }
 
@@ -1177,7 +1177,7 @@ public static ElasticInferenceService createServiceWithAuthHandler(
             createWithEmptySettings(threadPool),
             ElasticInferenceServiceSettingsTests.create(eisGatewayUrl),
             mockModelRegistry(threadPool),
-            new ElasticInferenceServiceAuthorizationRequestHandler(eisGatewayUrl, threadPool)
+            new ElasticInferenceServiceAuthorizationRequestHandler(ElasticInferenceServiceSettingsTests.create(eisGatewayUrl), threadPool)
         );
     }
 }
diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/authorization/ElasticInferenceServiceAuthorizationRequestHandlerTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/authorization/ElasticInferenceServiceAuthorizationRequestHandlerTests.java
@@ -24,6 +24,7 @@
 import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderTests;
 import org.elasticsearch.xpack.inference.external.http.sender.Sender;
 import org.elasticsearch.xpack.inference.logging.ThrottlerManager;
+import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceSettingsTests;
 import org.junit.After;
 import org.junit.Before;
 import org.mockito.ArgumentCaptor;
@@ -44,7 +45,6 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
 
 public class ElasticInferenceServiceAuthorizationRequestHandlerTests extends ESTestCase {
@@ -71,7 +71,11 @@ public void shutdown() throws IOException {
     public void testDoesNotAttempt_ToRetrieveAuthorization_IfBaseUrlIsNull() throws Exception {
         var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager);
         var logger = mock(Logger.class);
-        var authHandler = new ElasticInferenceServiceAuthorizationRequestHandler(null, threadPool, logger);
+        var authHandler = new ElasticInferenceServiceAuthorizationRequestHandler(
+            ElasticInferenceServiceSettingsTests.create(null),
+            threadPool,
+            logger
+        );
 
         try (var sender = senderFactory.createSender()) {
             PlainActionFuture<ElasticInferenceServiceAuthorizationModel> listener = new PlainActionFuture<>();
@@ -93,7 +97,11 @@ public void testDoesNotAttempt_ToRetrieveAuthorization_IfBaseUrlIsNull() throws
     public void testDoesNotAttempt_ToRetrieveAuthorization_IfBaseUrlIsEmpty() throws Exception {
         var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager);
         var logger = mock(Logger.class);
-        var authHandler = new ElasticInferenceServiceAuthorizationRequestHandler("", threadPool, logger);
+        var authHandler = new ElasticInferenceServiceAuthorizationRequestHandler(
+            ElasticInferenceServiceSettingsTests.create(""),
+            threadPool,
+            logger
+        );
 
         try (var sender = senderFactory.createSender()) {
             PlainActionFuture<ElasticInferenceServiceAuthorizationModel> listener = new PlainActionFuture<>();
@@ -116,7 +124,11 @@ public void testGetAuthorization_FailsWhenAnInvalidFieldIsFound() throws IOExcep
         var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager);
         var eisGatewayUrl = getUrl(webServer);
         var logger = mock(Logger.class);
-        var authHandler = new ElasticInferenceServiceAuthorizationRequestHandler(eisGatewayUrl, threadPool, logger);
+        var authHandler = new ElasticInferenceServiceAuthorizationRequestHandler(
+            ElasticInferenceServiceSettingsTests.create(eisGatewayUrl),
+            threadPool,
+            logger
+        );
 
         try (var sender = senderFactory.createSender()) {
             String responseJson = """
@@ -167,7 +179,11 @@ public void testGetAuthorization_ReturnsAValidResponse() throws IOException {
         var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager);
         var eisGatewayUrl = getUrl(webServer);
         var logger = mock(Logger.class);
-        var authHandler = new ElasticInferenceServiceAuthorizationRequestHandler(eisGatewayUrl, threadPool, logger);
+        var authHandler = new ElasticInferenceServiceAuthorizationRequestHandler(
+            ElasticInferenceServiceSettingsTests.create(eisGatewayUrl),
+            threadPool,
+            logger
+        );
 
         try (var sender = senderFactory.createSender()) {
             String responseJson = """
@@ -196,7 +212,48 @@ public void testGetAuthorization_ReturnsAValidResponse() throws IOException {
 
             var message = loggerArgsCaptor.getValue();
             assertThat(message, is("Retrieving authorization information from the Elastic Inference Service."));
-            verifyNoMoreInteractions(logger);
+        }
+    }
+
+    public void testGetAuthorization_PerformsADnsLookup() throws IOException {
+        var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager);
+        var eisGatewayUrl = getUrl(webServer);
+        var logger = mock(Logger.class);
+        when(logger.isDebugEnabled()).thenReturn(true);
+        var authHandler = new ElasticInferenceServiceAuthorizationRequestHandler(
+            ElasticInferenceServiceSettingsTests.createWithUrlAndDnsLookup(eisGatewayUrl),
+            threadPool,
+            logger
+        );
+
+        try (var sender = senderFactory.createSender()) {
+            String responseJson = """
+                {
+                    "models": [
+                        {
+                          "model_name": "model-a",
+                          "task_types": ["embed/text/sparse", "chat"]
+                        }
+                    ]
+                }
+                """;
+
+            webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson));
+
+            PlainActionFuture<ElasticInferenceServiceAuthorizationModel> listener = new PlainActionFuture<>();
+            authHandler.getAuthorization(listener, sender);
+
+            var authResponse = listener.actionGet(TIMEOUT);
+            assertThat(authResponse.getAuthorizedTaskTypes(), is(EnumSet.of(TaskType.SPARSE_EMBEDDING, TaskType.CHAT_COMPLETION)));
+            assertThat(authResponse.getAuthorizedModelIds(), is(Set.of("model-a")));
+            assertTrue(authResponse.isAuthorized());
+
+            var loggerArgsCaptor = ArgumentCaptor.forClass(String.class);
+            verify(logger, times(4)).debug(loggerArgsCaptor.capture());
+
+            var messages = loggerArgsCaptor.getAllValues();
+            assertThat(messages.size(), is(4));
+            assertThat(messages.get(3), is("Gateway address: 127.0.0.1"));
         }
     }
 
@@ -205,7 +262,11 @@ public void testGetAuthorization_OnResponseCalledOnce() throws IOException {
         var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager);
         var eisGatewayUrl = getUrl(webServer);
         var logger = mock(Logger.class);
-        var authHandler = new ElasticInferenceServiceAuthorizationRequestHandler(eisGatewayUrl, threadPool, logger);
+        var authHandler = new ElasticInferenceServiceAuthorizationRequestHandler(
+            ElasticInferenceServiceSettingsTests.create(eisGatewayUrl),
+            threadPool,
+            logger
+        );
 
         ActionListener<ElasticInferenceServiceAuthorizationModel> listener = mock(ActionListener.class);
         String responseJson = """
@@ -230,7 +291,6 @@ public void testGetAuthorization_OnResponseCalledOnce() throws IOException {
 
             var message = loggerArgsCaptor.getValue();
             assertThat(message, is("Retrieving authorization information from the Elastic Inference Service."));
-            verifyNoMoreInteractions(logger);
         }
     }
 
@@ -246,7 +306,11 @@ public void testGetAuthorization_InvalidResponse() throws IOException {
         }).when(senderMock).sendWithoutQueuing(any(), any(), any(), any(), any());
 
         var logger = mock(Logger.class);
-        var authHandler = new ElasticInferenceServiceAuthorizationRequestHandler("abc", threadPool, logger);
+        var authHandler = new ElasticInferenceServiceAuthorizationRequestHandler(
+            ElasticInferenceServiceSettingsTests.create("abc"),
+            threadPool,
+            logger
+        );
 
         try (var sender = senderFactory.createSender()) {
             PlainActionFuture<ElasticInferenceServiceAuthorizationModel> listener = new PlainActionFuture<>();

Original file line number	Diff line number	Diff line change
`@@ -1163,7 +1163,7 @@ private ElasticInferenceService createServiceWithAuthHandler(HttpRequestSender.F`
`1163`	`1163`	`createWithEmptySettings(threadPool),`
`1164`	`1164`	`ElasticInferenceServiceSettingsTests.create(eisGatewayUrl),`
`1165`	`1165`	`mockModelRegistry(),`
`1166`		`- new ElasticInferenceServiceAuthorizationRequestHandler(eisGatewayUrl, threadPool)`
	`1166`	`+ new ElasticInferenceServiceAuthorizationRequestHandler(ElasticInferenceServiceSettingsTests.create(eisGatewayUrl), threadPool)`
`1167`	`1167`	`);`
`1168`	`1168`	`}`
`1169`	`1169`
`@@ -1177,7 +1177,7 @@ public static ElasticInferenceService createServiceWithAuthHandler(`
`1177`	`1177`	`createWithEmptySettings(threadPool),`
`1178`	`1178`	`ElasticInferenceServiceSettingsTests.create(eisGatewayUrl),`
`1179`	`1179`	`mockModelRegistry(threadPool),`
`1180`		`- new ElasticInferenceServiceAuthorizationRequestHandler(eisGatewayUrl, threadPool)`
	`1180`	`+ new ElasticInferenceServiceAuthorizationRequestHandler(ElasticInferenceServiceSettingsTests.create(eisGatewayUrl), threadPool)`
`1181`	`1181`	`);`
`1182`	`1182`	`}`
`1183`	`1183`	`}`