diff --git a/extensions/auth/opa/impl/src/main/java/org/apache/polaris/extension/auth/opa/OpaPolarisAuthorizer.java b/extensions/auth/opa/impl/src/main/java/org/apache/polaris/extension/auth/opa/OpaPolarisAuthorizer.java index ab985abfea..016e893337 100644 --- a/extensions/auth/opa/impl/src/main/java/org/apache/polaris/extension/auth/opa/OpaPolarisAuthorizer.java +++ b/extensions/auth/opa/impl/src/main/java/org/apache/polaris/extension/auth/opa/OpaPolarisAuthorizer.java @@ -29,10 +29,13 @@ import java.util.Set; import org.apache.hc.client5.http.classic.methods.HttpPost; import org.apache.hc.client5.http.impl.classic.CloseableHttpClient; -import org.apache.hc.client5.http.impl.classic.CloseableHttpResponse; +import org.apache.hc.core5.http.ClassicHttpRequest; +import org.apache.hc.core5.http.ClassicHttpResponse; import org.apache.hc.core5.http.ContentType; +import org.apache.hc.core5.http.HttpException; import org.apache.hc.core5.http.HttpHeaders; import org.apache.hc.core5.http.ParseException; +import org.apache.hc.core5.http.io.HttpClientResponseHandler; import org.apache.hc.core5.http.io.entity.EntityUtils; import org.apache.hc.core5.http.io.entity.StringEntity; import org.apache.iceberg.exceptions.ForbiddenException; @@ -63,8 +66,8 @@ class OpaPolarisAuthorizer implements PolarisAuthorizer { /** * Public constructor that accepts a complete policy URI. * - * @param policyUri The required URI for the OPA endpoint. For example, - * https://opa.example.com/v1/polaris/allow + * @param policyUri The required URI for the OPA endpoint. For example, {@code + * https://opa.example.com/v1/polaris/allow}. * @param httpClient Apache HttpClient (required, injected by CDI). SSL configuration should be * handled by the CDI producer. * @param objectMapper Jackson ObjectMapper for JSON serialization (required). Shared across @@ -174,27 +177,35 @@ private boolean queryOpa( httpPost.setEntity(new StringEntity(inputJson, ContentType.APPLICATION_JSON)); // Execute request - try (CloseableHttpResponse response = httpClient.execute(httpPost)) { - int statusCode = response.getCode(); - if (statusCode != 200) { - return false; - } - - // Read and parse response - String responseBody; - try { - responseBody = EntityUtils.toString(response.getEntity()); - } catch (ParseException e) { - throw new RuntimeException("Failed to parse OPA response", e); - } - ObjectNode respNode = (ObjectNode) objectMapper.readTree(responseBody); - return respNode.path("result").path("allow").asBoolean(false); - } - } catch (IOException e) { + return httpClientExecute(httpPost, this::queryOpaCheckResponse); + } catch (HttpException | IOException e) { throw new RuntimeException("OPA query failed", e); } } + T httpClientExecute( + ClassicHttpRequest request, HttpClientResponseHandler responseHandler) + throws HttpException, IOException { + return httpClient.execute(request, responseHandler); + } + + private boolean queryOpaCheckResponse(ClassicHttpResponse response) throws IOException { + int statusCode = response.getCode(); + if (statusCode != 200) { + return false; + } + + // Read and parse response + String responseBody; + try { + responseBody = EntityUtils.toString(response.getEntity()); + } catch (ParseException e) { + throw new RuntimeException("Failed to parse OPA response", e); + } + ObjectNode respNode = (ObjectNode) objectMapper.readTree(responseBody); + return respNode.path("result").path("allow").asBoolean(false); + } + /** * Builds the OPA input JSON for the authorization query. * diff --git a/extensions/auth/opa/impl/src/test/java/org/apache/polaris/extension/auth/opa/OpaPolarisAuthorizerTest.java b/extensions/auth/opa/impl/src/test/java/org/apache/polaris/extension/auth/opa/OpaPolarisAuthorizerTest.java index 1c359ade2a..8042f7c415 100644 --- a/extensions/auth/opa/impl/src/test/java/org/apache/polaris/extension/auth/opa/OpaPolarisAuthorizerTest.java +++ b/extensions/auth/opa/impl/src/test/java/org/apache/polaris/extension/auth/opa/OpaPolarisAuthorizerTest.java @@ -20,17 +20,13 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatNoException; -import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; import com.sun.net.httpserver.HttpExchange; import com.sun.net.httpserver.HttpHandler; import com.sun.net.httpserver.HttpServer; -import java.io.ByteArrayInputStream; import java.io.IOException; import java.io.OutputStream; import java.net.InetSocketAddress; @@ -42,9 +38,14 @@ import java.util.Set; import org.apache.hc.client5.http.classic.methods.HttpPost; import org.apache.hc.client5.http.impl.classic.CloseableHttpClient; -import org.apache.hc.client5.http.impl.classic.CloseableHttpResponse; import org.apache.hc.client5.http.impl.classic.HttpClients; +import org.apache.hc.core5.http.ClassicHttpRequest; +import org.apache.hc.core5.http.ClassicHttpResponse; import org.apache.hc.core5.http.HttpEntity; +import org.apache.hc.core5.http.HttpException; +import org.apache.hc.core5.http.io.HttpClientResponseHandler; +import org.apache.hc.core5.http.io.entity.HttpEntities; +import org.apache.hc.core5.http.message.BasicClassicHttpResponse; import org.apache.polaris.core.auth.PolarisAuthorizableOperation; import org.apache.polaris.core.auth.PolarisPrincipal; import org.apache.polaris.core.entity.PolarisBaseEntity; @@ -55,7 +56,6 @@ import org.apache.polaris.extension.auth.opa.token.BearerTokenProvider; import org.apache.polaris.extension.auth.opa.token.StaticBearerTokenProvider; import org.junit.jupiter.api.Test; -import org.mockito.ArgumentCaptor; /** * Unit tests for OpaPolarisAuthorizer including basic functionality and bearer token authentication @@ -490,23 +490,27 @@ public void testCreateWithHttpsAndBearerToken() { } @Test - public void testBearerTokenIsAddedToHttpRequest() throws IOException { + public void testBearerTokenIsAddedToHttpRequest() { URI policyUri = URI.create("http://opa.example.com:8181/v1/data/polaris/allow"); - CloseableHttpClient mockHttpClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = mock(CloseableHttpResponse.class); - HttpEntity mockEntity = mock(HttpEntity.class); - - when(mockHttpClient.execute(any(HttpPost.class))).thenReturn(mockResponse); - when(mockResponse.getCode()).thenReturn(200); - when(mockResponse.getEntity()).thenReturn(mockEntity); - when(mockEntity.getContent()) - .thenReturn( - new ByteArrayInputStream( - "{\"result\":{\"allow\":true}}".getBytes(StandardCharsets.UTF_8))); + HttpEntity mockEntity = HttpEntities.create("{\"result\":{\"allow\":true}}"); + @SuppressWarnings("resource") + ClassicHttpResponse mockResponse = new BasicClassicHttpResponse(200); + mockResponse.setEntity(mockEntity); BearerTokenProvider tokenProvider = new StaticBearerTokenProvider("test-bearer-token"); OpaPolarisAuthorizer authorizer = - new OpaPolarisAuthorizer(policyUri, mockHttpClient, new ObjectMapper(), tokenProvider); + new OpaPolarisAuthorizer( + policyUri, mock(CloseableHttpClient.class), new ObjectMapper(), tokenProvider) { + @Override + T httpClientExecute( + ClassicHttpRequest request, HttpClientResponseHandler responseHandler) + throws HttpException, IOException { + // Verify the Authorization header with static bearer token + verifyAuthorizationHeader(request, "test-bearer-token"); + + return responseHandler.handleResponse(mockResponse); + } + }; PolarisPrincipal mockPrincipal = PolarisPrincipal.of("test-user", Map.of(), Collections.emptySet()); @@ -522,32 +526,33 @@ public void testBearerTokenIsAddedToHttpRequest() throws IOException { (PolarisResolvedPathWrapper) null, (PolarisResolvedPathWrapper) null); }); - - // Verify the Authorization header with static bearer token - verifyAuthorizationHeader(mockHttpClient, "test-bearer-token"); } @Test - public void testBearerTokenFromBearerTokenProvider() throws IOException { + public void testBearerTokenFromBearerTokenProvider() { // Mock HTTP client and response - CloseableHttpClient mockHttpClient = mock(CloseableHttpClient.class); - CloseableHttpResponse mockResponse = mock(CloseableHttpResponse.class); - HttpEntity mockEntity = mock(HttpEntity.class); - - when(mockHttpClient.execute(any(HttpPost.class))).thenReturn(mockResponse); - when(mockResponse.getCode()).thenReturn(200); - when(mockResponse.getEntity()).thenReturn(mockEntity); - when(mockEntity.getContent()) - .thenReturn( - new ByteArrayInputStream( - "{\"result\":{\"allow\":true}}".getBytes(StandardCharsets.UTF_8))); + HttpEntity mockEntity = HttpEntities.create("{\"result\":{\"allow\":true}}"); + @SuppressWarnings("resource") + ClassicHttpResponse mockResponse = new BasicClassicHttpResponse(200); + mockResponse.setEntity(mockEntity); // Create token provider that returns a dynamic token BearerTokenProvider tokenProvider = () -> "dynamic-token-12345"; URI policyUri = URI.create("http://opa.example.com:8181/v1/data/polaris/allow"); // Create authorizer with the token provider instead of static token OpaPolarisAuthorizer authorizer = - new OpaPolarisAuthorizer(policyUri, mockHttpClient, new ObjectMapper(), tokenProvider); + new OpaPolarisAuthorizer( + policyUri, mock(CloseableHttpClient.class), new ObjectMapper(), tokenProvider) { + @Override + T httpClientExecute( + ClassicHttpRequest request, HttpClientResponseHandler responseHandler) + throws HttpException, IOException { + // Verify the Authorization header with bearer token from provider + verifyAuthorizationHeader(request, "dynamic-token-12345"); + + return responseHandler.handleResponse(mockResponse); + } + }; // Create mock principal and entities PolarisPrincipal mockPrincipal = @@ -566,9 +571,6 @@ public void testBearerTokenFromBearerTokenProvider() throws IOException { (PolarisResolvedPathWrapper) null, (PolarisResolvedPathWrapper) null); }); - - // Verify the Authorization header with bearer token from provider - verifyAuthorizationHeader(mockHttpClient, "dynamic-token-12345"); } private ResolvedPolarisEntity createResolvedEntity(PolarisEntity entity) { @@ -632,17 +634,13 @@ public void handle(HttpExchange exchange) throws IOException { /** * Helper method to capture and verify HTTP request Authorization header. * - * @param mockHttpClient The mocked HTTP client to verify against + * @param capturedRequest The request issued to the HTTP client to verify against * @param expectedToken The expected bearer token value, or null if no Authorization header * expected */ - private void verifyAuthorizationHeader(CloseableHttpClient mockHttpClient, String expectedToken) - throws IOException { - // Capture the HTTP request to verify bearer token header - ArgumentCaptor httpPostCaptor = ArgumentCaptor.forClass(HttpPost.class); - verify(mockHttpClient).execute(httpPostCaptor.capture()); - - HttpPost capturedRequest = httpPostCaptor.getValue(); + private void verifyAuthorizationHeader(ClassicHttpRequest capturedRequest, String expectedToken) { + // Capture the HTTP request to verify the bearer token header + assertThat(capturedRequest).isInstanceOf(HttpPost.class); if (expectedToken != null) { // Verify the Authorization header is present and contains the expected token