github-advanced-security[bot] commented on code in PR #3037: URL: https://github.com/apache/drill/pull/3037#discussion_r3560287205
########## exec/java-exec/src/main/java/org/apache/drill/exec/server/rest/ai/AnthropicProvider.java: ########## @@ -0,0 +1,438 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.drill.exec.server.rest.ai; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +import okhttp3.MediaType; +import okhttp3.OkHttpClient; +import okhttp3.Request; +import okhttp3.RequestBody; +import okhttp3.Response; +import okhttp3.ResponseBody; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.BufferedReader; +import java.io.InputStreamReader; +import java.io.OutputStream; +import java.nio.charset.StandardCharsets; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * LLM provider for the Anthropic Claude API. + * Calls /v1/messages with stream:true, translates Anthropic SSE events + * (message_start, content_block_delta, etc.) to the normalized format. + */ +public class AnthropicProvider implements LlmProvider { + + private static final Logger logger = LoggerFactory.getLogger(AnthropicProvider.class); + private static final ObjectMapper MAPPER = new ObjectMapper(); + private static final MediaType JSON_TYPE = MediaType.parse("application/json"); + private static final String DEFAULT_ENDPOINT = "https://api.anthropic.com";; + private static final String ANTHROPIC_VERSION = "2023-06-01"; + + @Override + public String getId() { + return "anthropic"; + } + + @Override + public String getDisplayName() { + return "Anthropic Claude"; + } + + @Override + public LlmCallResult streamChatCompletion(LlmConfig config, List<ChatMessage> messages, + List<ToolDefinition> tools, OutputStream out, UsageObserver usageObserver) + throws Exception { + UsageObserver observer = usageObserver != null ? usageObserver : UsageObserver.NOOP; + + // Create HTTP client with enterprise configuration + OkHttpClient httpClient = HttpClientFactory.createClient(config); + + String endpoint = LlmProvider.normalizeEndpoint(config.getApiEndpoint(), DEFAULT_ENDPOINT); + String url = endpoint + "/v1/messages"; + + ObjectNode requestBody = buildRequestBody(config, messages, tools); + + Request request = new Request.Builder() + .url(url) + .post(RequestBody.create(requestBody.toString(), JSON_TYPE)) + .addHeader("Content-Type", "application/json") + .addHeader("x-api-key", config.getApiKey() != null ? config.getApiKey() : "") + .addHeader("anthropic-version", ANTHROPIC_VERSION) + .build(); + + LlmCallResult result = new LlmCallResult(); + try (Response response = httpClient.newCall(request).execute()) { + if (!response.isSuccessful()) { + String errorBody = ""; + ResponseBody body = response.body(); + if (body != null) { + errorBody = body.string(); + } + String errorMsg = "Anthropic API error " + response.code() + ": " + errorBody; + logger.error(errorMsg); + writeSseEvent(out, "error", "{\"message\":" + MAPPER.writeValueAsString(errorMsg) + "}"); + return result; + } + + ResponseBody body = response.body(); + if (body == null) { + writeSseEvent(out, "error", "{\"message\":\"Empty response from Anthropic API\"}"); + return result; + } + + try (BufferedReader reader = new BufferedReader( + new InputStreamReader(body.byteStream(), StandardCharsets.UTF_8))) { + processAnthropicStream(reader, out, result, observer); + } + } + return result; + } + + @Override + public ValidationResult validateConfig(LlmConfig config) { + if (config.getApiKey() == null || config.getApiKey().isEmpty()) { + return ValidationResult.error("API key is required for Anthropic"); + } + if (config.getModel() == null || config.getModel().isEmpty()) { + return ValidationResult.error("Model name is required"); + } + + // Send a minimal, non-streaming request so we surface real failures (bad key, + // unknown model, unreachable endpoint) instead of reporting success blindly. + String endpoint = LlmProvider.normalizeEndpoint(config.getApiEndpoint(), DEFAULT_ENDPOINT); + String url = endpoint + "/v1/messages"; + + ObjectNode probeBody = MAPPER.createObjectNode(); + probeBody.put("model", config.getModel()); + probeBody.put("max_tokens", 1); + ArrayNode probeMessages = probeBody.putArray("messages"); + ObjectNode probeMsg = probeMessages.addObject(); + probeMsg.put("role", "user"); + probeMsg.put("content", "ping"); + + Request request = new Request.Builder() + .url(url) Review Comment: ## CodeQL / Server-side request forgery Potential server-side request forgery due to a [user-provided value](1). [Show more details](https://github.com/apache/drill/security/code-scanning/81) ########## exec/java-exec/src/main/java/org/apache/drill/exec/server/rest/ai/HttpClientFactory.java: ########## @@ -0,0 +1,291 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.drill.exec.server.rest.ai; + +import okhttp3.Credentials; +import okhttp3.HttpUrl; +import okhttp3.OkHttpClient; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.net.InetSocketAddress; +import java.net.Proxy; +import javax.net.ssl.KeyManager; +import javax.net.ssl.KeyManagerFactory; +import javax.net.ssl.SSLContext; +import javax.net.ssl.TrustManager; +import javax.net.ssl.TrustManagerFactory; +import javax.net.ssl.X509TrustManager; +import java.io.ByteArrayInputStream; +import java.io.FileInputStream; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Paths; +import java.security.KeyFactory; +import java.security.KeyStore; +import java.security.PrivateKey; +import java.security.SecureRandom; +import java.security.cert.Certificate; +import java.security.cert.CertificateFactory; +import java.security.spec.PKCS8EncodedKeySpec; +import java.util.ArrayList; +import java.util.Base64; +import java.util.List; +import java.util.concurrent.TimeUnit; +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +/** + * Factory for creating OkHttpClient instances with enterprise configuration. + * Centralizes HTTP client setup including proxies, SSL/TLS, custom headers, and timeouts. + */ +public class HttpClientFactory { + private static final Logger logger = LoggerFactory.getLogger(HttpClientFactory.class); + + private static final Pattern CERT_PATTERN = Pattern.compile( + "-----BEGIN CERTIFICATE-----(.+?)-----END CERTIFICATE-----", Pattern.DOTALL); + private static final Pattern KEY_PATTERN = Pattern.compile( + "-----BEGIN PRIVATE KEY-----(.+?)-----END PRIVATE KEY-----", Pattern.DOTALL); + + // Default timeouts (in seconds) - matching existing hardcoded values + private static final int DEFAULT_CONNECT_TIMEOUT = 30; + private static final int DEFAULT_READ_TIMEOUT = 120; + private static final int DEFAULT_WRITE_TIMEOUT = 30; + + /** + * Create an OkHttpClient configured with all enterprise settings from LlmConfig. + * + * @param config The LLM configuration + * @return Configured OkHttpClient + * @throws Exception if SSL/TLS configuration or file loading fails + */ + public static OkHttpClient createClient(LlmConfig config) throws Exception { + OkHttpClient.Builder builder = new OkHttpClient.Builder(); + + // 1. Configure timeouts + int connectTimeout = config.getConnectTimeoutSeconds() != null + ? config.getConnectTimeoutSeconds() : DEFAULT_CONNECT_TIMEOUT; + int readTimeout = config.getReadTimeoutSeconds() != null + ? config.getReadTimeoutSeconds() : DEFAULT_READ_TIMEOUT; + int writeTimeout = config.getWriteTimeoutSeconds() != null + ? config.getWriteTimeoutSeconds() : DEFAULT_WRITE_TIMEOUT; + + builder.connectTimeout(connectTimeout, TimeUnit.SECONDS) + .readTimeout(readTimeout, TimeUnit.SECONDS) + .writeTimeout(writeTimeout, TimeUnit.SECONDS); + + // 2. Configure proxy + if (config.getProxyUrl() != null && !config.getProxyUrl().isEmpty()) { + configureProxy(builder, config); + } + + // 3. Configure SSL/TLS + configureSSL(builder, config); + + // 4. Add custom headers interceptor + if (config.getCustomHeaders() != null && !config.getCustomHeaders().isEmpty()) { + builder.addInterceptor(new CustomHeadersInterceptor(config.getCustomHeaders())); + } + + // 5. Add logging/debugging interceptor (redacts sensitive headers) + if (logger.isDebugEnabled()) { + builder.addInterceptor(new SensitiveHeaderRedactor()); + } + + return builder.build(); + } + + /** + * Configure HTTP proxy with optional authentication. + */ + private static void configureProxy(OkHttpClient.Builder builder, LlmConfig config) { + HttpUrl proxyUrl = HttpUrl.parse(config.getProxyUrl()); + if (proxyUrl == null) { + throw new IllegalArgumentException("Invalid proxy URL: " + config.getProxyUrl()); + } + + Proxy proxy = new Proxy(Proxy.Type.HTTP, + new InetSocketAddress(proxyUrl.host(), proxyUrl.port())); + builder.proxy(proxy); + + // Proxy authentication + if (config.getProxyUsername() != null && !config.getProxyUsername().isEmpty()) { + builder.proxyAuthenticator((route, response) -> { + String credential = Credentials.basic( + config.getProxyUsername(), + config.getProxyPassword() != null ? config.getProxyPassword() : "" + ); + return response.request().newBuilder() + .header("Proxy-Authorization", credential) + .build(); + }); + } + } + + /** + * Configure SSL/TLS with optional custom truststore (CA certificates) and keystore (client certificates). + */ + private static void configureSSL(OkHttpClient.Builder builder, LlmConfig config) + throws Exception { + // Check if custom SSL configuration is needed + boolean hasKeystore = config.getKeystorePath() != null && !config.getKeystorePath().isEmpty(); + boolean hasTruststore = config.getTruststorePath() != null && !config.getTruststorePath().isEmpty(); + boolean hasPemCert = config.getClientCertPath() != null && !config.getClientCertPath().isEmpty(); + + if (hasKeystore || hasTruststore || hasPemCert) { + SSLContext sslContext = SSLContext.getInstance("TLS"); + + KeyManager[] keyManagers = null; + if (hasKeystore) { + keyManagers = loadKeyManagers(config); + } else if (hasPemCert) { + keyManagers = loadKeyManagersFromPem(config); + } + + TrustManager[] trustManagers = null; + if (hasTruststore) { + trustManagers = loadTrustManagers(config); + } + + sslContext.init(keyManagers, trustManagers, new SecureRandom()); + + X509TrustManager trustManager = trustManagers != null + ? (X509TrustManager) trustManagers[0] + : getDefaultTrustManager(); + + builder.sslSocketFactory(sslContext.getSocketFactory(), trustManager); + } + + // Disable SSL verification if explicitly set to false (INSECURE - log warning) + if (Boolean.FALSE.equals(config.getVerifySSL())) { + logger.warn("SSL certificate verification is disabled! Use only for testing/development."); + builder.hostnameVerifier((hostname, session) -> true); + } + } + + /** + * Load key managers from keystore for mTLS (mutual TLS / client certificate) support. + */ + private static KeyManager[] loadKeyManagers(LlmConfig config) throws Exception { + String keystorePath = config.getKeystorePath(); + String keystorePassword = config.getKeystorePassword(); + String keystoreType = config.getKeystoreType() != null ? config.getKeystoreType() : "JKS"; + + KeyStore keyStore = KeyStore.getInstance(keystoreType); + try (FileInputStream fis = new FileInputStream(keystorePath)) { + keyStore.load(fis, keystorePassword != null ? keystorePassword.toCharArray() : null); + } + + KeyManagerFactory kmf = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm()); + kmf.init(keyStore, keystorePassword != null ? keystorePassword.toCharArray() : null); + + return kmf.getKeyManagers(); + } + + /** + * Load trust managers from truststore for custom CA certificate support. + */ + private static TrustManager[] loadTrustManagers(LlmConfig config) throws Exception { + String truststorePath = config.getTruststorePath(); + String truststorePassword = config.getTruststorePassword(); + String truststoreType = config.getTruststoreType() != null ? config.getTruststoreType() : "JKS"; + + KeyStore trustStore = KeyStore.getInstance(truststoreType); + try (FileInputStream fis = new FileInputStream(truststorePath)) { + trustStore.load(fis, truststorePassword != null ? truststorePassword.toCharArray() : null); + } + + TrustManagerFactory tmf = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm()); + tmf.init(trustStore); + + return tmf.getTrustManagers(); + } + + /** + * Load key managers from a PEM file holding the client certificate chain plus an + * unencrypted PKCS#8 private key (the single-file form that Python httpx/requests + * accepts as {@code cert=}). Used for mTLS against enterprise gateways. + * + * <p>PKCS#1 ({@code BEGIN RSA PRIVATE KEY}) and SEC1 ({@code BEGIN EC PRIVATE KEY}) keys + * are not supported by the JDK directly; the error message tells the user how to convert. + */ + private static KeyManager[] loadKeyManagersFromPem(LlmConfig config) throws Exception { + String pemPath = config.getClientCertPath(); + String pem = new String(Files.readAllBytes(Paths.get(pemPath)), StandardCharsets.UTF_8); Review Comment: ## CodeQL / Uncontrolled data used in path expression This path depends on a [user-provided value](1). [Show more details](https://github.com/apache/drill/security/code-scanning/85) ########## exec/java-exec/src/main/java/org/apache/drill/exec/server/rest/ai/OAuth2GatewayProvider.java: ########## @@ -0,0 +1,225 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.drill.exec.server.rest.ai; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import okhttp3.MediaType; +import okhttp3.Request; +import okhttp3.RequestBody; +import okhttp3.Response; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.nio.charset.StandardCharsets; +import java.time.LocalDateTime; +import java.time.format.DateTimeFormatter; +import java.util.Base64; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.UUID; + +/** + * Provider for OpenAI-compatible chat APIs fronted by an Apigee-style enterprise gateway + * that authenticates with OAuth2 client-credentials and mTLS. The chat API itself is + * OpenAI-compatible, so this reuses the entire request-building and SSE-parsing path of + * {@link OpenAiCompatibleProvider} and only overrides authentication: + * + * <ol> + * <li>Fetch a bearer token from the auth URL via OAuth2 client-credentials — Basic auth + * with {@code base64(consumerKey:consumerSecret)}, body {@code grant_type=client_credentials}. + * <li>Cache the token until shortly before its {@code expires_in}. + * <li>Attach the headers defined by {@code gatewayHeaders} (see {@link #DEFAULT_GATEWAY_HEADERS}). + * </ol> + * + * <p>Header names are fully configurable via {@code gatewayHeaders}, a map of header name to a + * value template. Templates may contain the placeholders {@code {token}} (the fetched bearer + * token), {@code {uuid}} (a fresh UUID per request), {@code {timestamp}} (a local microsecond + * ISO-8601 timestamp), {@code {apiKey}}, {@code {clientId}}, {@code {usecaseId}} and + * {@code {model}}; anything else is sent literally. See {@link #DEFAULT_GATEWAY_HEADERS} for + * the mapping used when none is configured. + * + * <p>mTLS (the client {@code .pem}) is handled by {@link HttpClientFactory} via + * {@code clientCertPath}, so it applies to both the token call and the endpoint call. + */ +public class OAuth2GatewayProvider extends OpenAiCompatibleProvider { + + private static final Logger logger = LoggerFactory.getLogger(OAuth2GatewayProvider.class); + private static final ObjectMapper MAPPER = new ObjectMapper(); + private static final MediaType FORM = MediaType.parse("application/x-www-form-urlencoded"); + + // Matches Python's datetime.now().isoformat(): local time, microsecond precision, no zone. + // ponytail: local time to mirror the reference client; switch to UTC here if the gateway + // rejects the date as stale. + private static final DateTimeFormatter REQUEST_DATE = + DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss.SSSSSS"); + + // Refresh a bit before the stated expiry to absorb clock skew and request latency. + private static final long EXPIRY_MARGIN_MILLIS = 60_000L; + + /** + * Default header mapping used when {@code gatewayHeaders} is not configured. Keys are + * header names; values are templates resolved by {@link #resolveTemplate}. + */ + static final Map<String, String> DEFAULT_GATEWAY_HEADERS; + + static { + Map<String, String> defaults = new LinkedHashMap<>(); + defaults.put("Authorization", "Bearer {token}"); + defaults.put("client-id", "{clientId}"); + defaults.put("usecase-id", "{usecaseId}"); + defaults.put("api-key", "{apiKey}"); + defaults.put("x-request-id", "{uuid}"); + defaults.put("x-correlation-id", "{uuid}"); + defaults.put("request-date", "{timestamp}"); + DEFAULT_GATEWAY_HEADERS = Collections.unmodifiableMap(defaults); + } + + private String cachedToken; + private String cachedFor; + private long cachedExpiryMillis; + + @Override + public String getId() { + return "oauth2"; + } + + @Override + public String getDisplayName() { + return "OAuth2 Gateway"; + } + + @Override + public ValidationResult validateConfig(LlmConfig config) { + String missing = firstMissingField(config); + if (missing != null) { + return ValidationResult.error(missing + " is required"); + } + // probe() drives the real flow: token fetch -> endpoint call, so a failure tells the + // user which leg broke (auth vs. mTLS vs. endpoint) via the ValidationResult details. + return probe(config); + } + + @Override + protected void decorateRequest(Request.Builder builder, LlmConfig config) throws Exception { + Map<String, String> headers = config.getGatewayHeaders(); + if (headers == null || headers.isEmpty()) { + headers = DEFAULT_GATEWAY_HEADERS; + } + // Fetch the token only if some header actually references it. + String token = null; + for (Map.Entry<String, String> header : headers.entrySet()) { + String template = header.getValue(); + if (template == null) { + continue; + } + if (token == null && template.contains("{token}")) { + token = getAccessToken(config); + } + builder.addHeader(header.getKey(), resolveTemplate(template, config, token)); + } + } + + /** Substitute the supported placeholders in a header value template. */ + private String resolveTemplate(String template, LlmConfig config, String token) { + return template + .replace("{token}", orEmpty(token)) + .replace("{uuid}", UUID.randomUUID().toString()) + .replace("{timestamp}", LocalDateTime.now().format(REQUEST_DATE)) + .replace("{apiKey}", orEmpty(config.getApiKey())) + .replace("{clientId}", orEmpty(config.getClientId())) + .replace("{usecaseId}", orEmpty(config.getUsecaseId())) + .replace("{model}", orEmpty(config.getModel())); + } + + private static String orEmpty(String value) { + return value == null ? "" : value; + } + + /** + * Fetch and cache an OAuth2 client-credentials bearer token. Synchronized because the + * singleton provider serves concurrent chat requests; contention only occurs on refresh. + */ + private synchronized String getAccessToken(LlmConfig config) throws Exception { + long now = System.currentTimeMillis(); + if (cachedToken != null + && config.getConsumerKey().equals(cachedFor) + && now < cachedExpiryMillis - EXPIRY_MARGIN_MILLIS) { + return cachedToken; + } + + String basic = Base64.getEncoder().encodeToString( + (config.getConsumerKey() + ":" + config.getConsumerSecret()) + .getBytes(StandardCharsets.UTF_8)); + + Request request = new Request.Builder() + .url(config.getAuthUrl()) Review Comment: ## CodeQL / Server-side request forgery Potential server-side request forgery due to a [user-provided value](1). [Show more details](https://github.com/apache/drill/security/code-scanning/84) ########## exec/java-exec/src/main/java/org/apache/drill/exec/server/rest/ai/EnterpriseProvider.java: ########## @@ -0,0 +1,474 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.drill.exec.server.rest.ai; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import okhttp3.MediaType; +import okhttp3.OkHttpClient; +import okhttp3.Request; +import okhttp3.RequestBody; +import okhttp3.Response; +import okhttp3.ResponseBody; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.BufferedReader; +import java.io.InputStreamReader; +import java.io.OutputStream; +import java.nio.charset.StandardCharsets; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * LLM provider for custom/enterprise APIs with non-standard request/response formats. + * + * Uses configurable request templates and response mappings to support arbitrary API formats. + * Request templates support simple variable substitution: {model}, {temperature}, {maxTokens}, {messages}, etc. + * Response mappings use JSONPath-like syntax to extract deltas, tool calls, and completion status. + */ +public class EnterpriseProvider implements LlmProvider { + + private static final Logger logger = LoggerFactory.getLogger(EnterpriseProvider.class); + private static final ObjectMapper MAPPER = new ObjectMapper(); + private static final MediaType JSON_TYPE = MediaType.parse("application/json"); + + @Override + public String getId() { + return "enterprise"; + } + + @Override + public String getDisplayName() { + return "Enterprise (Custom API)"; + } + + @Override + public ValidationResult validateConfig(LlmConfig config) { + if (config.getModel() == null || config.getModel().isEmpty()) { + return ValidationResult.error("Model name is required"); + } + + if (config.getApiEndpoint() == null || config.getApiEndpoint().isEmpty()) { + return ValidationResult.error("API endpoint is required"); + } + + if (config.getRequestTemplate() == null || config.getRequestTemplate().isEmpty()) { + return ValidationResult.error("Request template is required for enterprise provider"); + } + + // Validate that request template is valid JSON + try { + MAPPER.readTree(config.getRequestTemplate()); + } catch (Exception e) { + return ValidationResult.error("Invalid request template JSON: " + e.getMessage()); + } + + // Validate that response mapping is valid JSON (if provided) + if (config.getResponseMapping() != null && !config.getResponseMapping().isEmpty()) { + try { + MAPPER.readTree(config.getResponseMapping()); + } catch (Exception e) { + return ValidationResult.error("Invalid response mapping JSON: " + e.getMessage()); + } + } + + // Send a real probe to the endpoint (using the configured template) so "Test + // Connection" actually exercises DNS, TLS, proxy, and auth instead of only + // checking that the templates parse. This is the same request path the live + // chat uses, so a failure here surfaces the real enterprise-gateway error. + String url = config.getApiEndpoint(); + String requestBody; + try { + requestBody = buildCustomRequest(config, List.of(ChatMessage.user("ping")), null); + } catch (Exception e) { + logger.warn("Enterprise config test could not build request from template", e); + return ValidationResult.error("Failed to build request from template: " + e.getMessage(), + LlmProvider.describeException(url, e)); + } + + Request.Builder reqBuilder = new Request.Builder() + .url(url) Review Comment: ## CodeQL / Server-side request forgery Potential server-side request forgery due to a [user-provided value](1). [Show more details](https://github.com/apache/drill/security/code-scanning/82) ########## exec/java-exec/src/main/java/org/apache/drill/exec/server/rest/ai/OpenAiCompatibleProvider.java: ########## @@ -0,0 +1,413 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.drill.exec.server.rest.ai; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +import okhttp3.MediaType; +import okhttp3.OkHttpClient; +import okhttp3.Request; +import okhttp3.RequestBody; +import okhttp3.Response; +import okhttp3.ResponseBody; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.BufferedReader; +import java.io.InputStreamReader; +import java.io.OutputStream; +import java.nio.charset.StandardCharsets; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * LLM provider for OpenAI-compatible APIs (OpenAI, Azure OpenAI, Ollama, etc.). + * Uses OkHttp to call /chat/completions with stream:true, + * reads SSE line-by-line, and normalizes deltas to the common wire format. + */ +public class OpenAiCompatibleProvider implements LlmProvider { + + private static final Logger logger = LoggerFactory.getLogger(OpenAiCompatibleProvider.class); + private static final ObjectMapper MAPPER = new ObjectMapper(); + private static final MediaType JSON_TYPE = MediaType.parse("application/json"); + private static final String DEFAULT_ENDPOINT = "https://api.openai.com/v1";; + + @Override + public String getId() { + return "openai"; + } + + @Override + public String getDisplayName() { + return "OpenAI Compatible"; + } + + @Override + public LlmCallResult streamChatCompletion(LlmConfig config, List<ChatMessage> messages, + List<ToolDefinition> tools, OutputStream out, UsageObserver usageObserver) + throws Exception { + UsageObserver observer = usageObserver != null ? usageObserver : UsageObserver.NOOP; + + // Create HTTP client with enterprise configuration + OkHttpClient httpClient = HttpClientFactory.createClient(config); + + String url = chatUrl(config); + + ObjectNode requestBody = buildRequestBody(config, messages, tools); + + Request.Builder reqBuilder = new Request.Builder() + .url(url) + .post(RequestBody.create(requestBody.toString(), JSON_TYPE)) + .addHeader("Content-Type", "application/json"); + decorateRequest(reqBuilder, config); + + Request request = reqBuilder.build(); + + LlmCallResult result = new LlmCallResult(); + try (Response response = httpClient.newCall(request).execute()) { + if (!response.isSuccessful()) { + String errorBody = ""; + ResponseBody body = response.body(); + if (body != null) { + errorBody = body.string(); + } + String errorMsg = "LLM API error " + response.code() + ": " + errorBody; + logger.error(errorMsg); + writeSseEvent(out, "error", "{\"message\":" + MAPPER.writeValueAsString(errorMsg) + "}"); + return result; + } + + ResponseBody body = response.body(); + if (body == null) { + writeSseEvent(out, "error", "{\"message\":\"Empty response from LLM API\"}"); + return result; + } + + try (BufferedReader reader = new BufferedReader( + new InputStreamReader(body.byteStream(), StandardCharsets.UTF_8))) { + processOpenAiStream(reader, out, result, observer); + } + } + return result; + } + + @Override + public ValidationResult validateConfig(LlmConfig config) { + if (config.getApiKey() == null || config.getApiKey().isEmpty()) { + // Ollama doesn't require an API key, so only warn + String endpoint = config.getApiEndpoint(); + if (endpoint != null && !endpoint.contains("localhost") && !endpoint.contains("127.0.0.1")) { + return ValidationResult.error("API key is required for non-local endpoints"); + } + } + if (config.getModel() == null || config.getModel().isEmpty()) { + return ValidationResult.error("Model name is required"); + } + return probe(config); + } + + /** + * Send a minimal, non-streaming request so we surface real failures (bad key/token, + * unknown model, unreachable endpoint, TLS) instead of reporting success blindly. + * Uses {@link #chatUrl} and {@link #decorateRequest}, so subclasses that override auth + * (e.g. an OAuth gateway) exercise their real auth path here too. + */ + protected ValidationResult probe(LlmConfig config) { + String url = chatUrl(config); + + ObjectNode probeBody = MAPPER.createObjectNode(); + probeBody.put("model", config.getModel()); + probeBody.put("max_tokens", 1); + ArrayNode probeMessages = probeBody.putArray("messages"); + ObjectNode probeMsg = probeMessages.addObject(); + probeMsg.put("role", "user"); + probeMsg.put("content", "ping"); + + try { + Request.Builder reqBuilder = new Request.Builder() + .url(url) Review Comment: ## CodeQL / Server-side request forgery Potential server-side request forgery due to a [user-provided value](1). [Show more details](https://github.com/apache/drill/security/code-scanning/83) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
