-
Notifications
You must be signed in to change notification settings - Fork 172
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add authorization code redirect request handler & tests
- Loading branch information
1 parent
db9949c
commit 2b3328f
Showing
4 changed files
with
166 additions
and
50 deletions.
There are no files selected for viewing
87 changes: 87 additions & 0 deletions
87
...in/java/net/snowflake/client/core/auth/oauth/AuthorizationCodeRedirectRequestHandler.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,87 @@ | ||
| /* | ||
| * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. | ||
| */ | ||
|
|
||
| package net.snowflake.client.core.auth.oauth; | ||
|
|
||
| import com.amazonaws.util.StringUtils; | ||
| import com.nimbusds.oauth2.sdk.id.State; | ||
| import com.sun.net.httpserver.HttpExchange; | ||
| import com.sun.net.httpserver.HttpHandler; | ||
| import java.io.IOException; | ||
| import java.nio.charset.StandardCharsets; | ||
| import java.util.Map; | ||
| import java.util.concurrent.CompletableFuture; | ||
| import java.util.stream.Collectors; | ||
| import net.snowflake.client.core.SFException; | ||
| import net.snowflake.client.jdbc.ErrorCode; | ||
| import net.snowflake.client.log.SFLogger; | ||
| import net.snowflake.client.log.SFLoggerFactory; | ||
| import org.apache.http.NameValuePair; | ||
| import org.apache.http.client.utils.URLEncodedUtils; | ||
|
|
||
| class AuthorizationCodeRedirectRequestHandler implements HttpHandler { | ||
|
|
||
| private static final SFLogger logger = | ||
| SFLoggerFactory.getLogger(AuthorizationCodeRedirectRequestHandler.class); | ||
|
|
||
| private final CompletableFuture<String> authorizationCodeFuture; | ||
| private final State expectedState; | ||
|
|
||
| AuthorizationCodeRedirectRequestHandler( | ||
| CompletableFuture<String> authorizationCodeFuture, State expectedState) { | ||
| this.authorizationCodeFuture = authorizationCodeFuture; | ||
| this.expectedState = expectedState; | ||
| } | ||
|
|
||
| @Override | ||
| public void handle(HttpExchange exchange) throws IOException { | ||
| Map<String, String> urlParams = | ||
| URLEncodedUtils.parse(exchange.getRequestURI(), StandardCharsets.UTF_8).stream() | ||
| .collect(Collectors.toMap(NameValuePair::getName, NameValuePair::getValue)); | ||
| String response = handleRedirectRequest(urlParams, authorizationCodeFuture, expectedState); | ||
| exchange.sendResponseHeaders(200, response.length()); | ||
| exchange.getResponseBody().write(response.getBytes(StandardCharsets.UTF_8)); | ||
| exchange.getResponseBody().close(); | ||
| } | ||
|
|
||
| static String handleRedirectRequest( | ||
| Map<String, String> urlParams, | ||
| CompletableFuture<String> authorizationCodeFuture, | ||
| State expectedState) { | ||
| String response; | ||
| if (urlParams.containsKey("error")) { | ||
| response = "Authorization error: " + urlParams.get("error"); | ||
| authorizationCodeFuture.completeExceptionally( | ||
| new SFException( | ||
| ErrorCode.OAUTH_AUTHORIZATION_CODE_FLOW_ERROR, | ||
| String.format( | ||
| "Error during authorization: %s, %s", | ||
| urlParams.get("error"), urlParams.get("error_description")))); | ||
| } else if (!expectedState.getValue().equals(urlParams.get("state"))) { | ||
| authorizationCodeFuture.completeExceptionally( | ||
| new SFException( | ||
| ErrorCode.OAUTH_AUTHORIZATION_CODE_FLOW_ERROR, | ||
| String.format( | ||
| "Invalid authorization request redirection state: %s, expected: %s", | ||
| urlParams.get("state"), expectedState.getValue()))); | ||
| response = "Authorization error: invalid authorization request redirection state"; | ||
| } else { | ||
| String authorizationCode = urlParams.get("code"); | ||
| if (!StringUtils.isNullOrEmpty(authorizationCode)) { | ||
| logger.debug("Received authorization code on redirect URI"); | ||
| response = "Authorization completed successfully."; | ||
| authorizationCodeFuture.complete(authorizationCode); | ||
| } else { | ||
| authorizationCodeFuture.completeExceptionally( | ||
| new SFException( | ||
| ErrorCode.OAUTH_AUTHORIZATION_CODE_FLOW_ERROR, | ||
| String.format( | ||
| "Authorization code redirect URI server received request without authorization code; queryParams: %s", | ||
| urlParams))); | ||
| response = "Authorization error: authorization code has not been returned to the driver."; | ||
| } | ||
| } | ||
| return response; | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
75 changes: 75 additions & 0 deletions
75
...ava/net/snowflake/client/core/auth/oauth/AuthorizationCodeRedirectRequestHandlerTest.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,75 @@ | ||
| /* | ||
| * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. | ||
| */ | ||
|
|
||
| package net.snowflake.client.core.auth.oauth; | ||
|
|
||
| import com.nimbusds.oauth2.sdk.id.State; | ||
| import java.util.HashMap; | ||
| import java.util.Map; | ||
| import java.util.concurrent.CompletableFuture; | ||
| import net.snowflake.client.core.SFException; | ||
| import org.junit.Test; | ||
| import org.junit.jupiter.api.Assertions; | ||
| import org.mockito.Mockito; | ||
|
|
||
| public class AuthorizationCodeRedirectRequestHandlerTest { | ||
|
|
||
| CompletableFuture<String> authorizationCodeFutureMock = Mockito.mock(CompletableFuture.class); | ||
|
|
||
| @Test | ||
| public void shouldReturnSuccessResponse() { | ||
| Map<String, String> params = new HashMap<>(); | ||
| params.put("code", "some authorization code"); | ||
| params.put("state", "abc"); | ||
|
|
||
| String response = | ||
| AuthorizationCodeRedirectRequestHandler.handleRedirectRequest( | ||
| params, authorizationCodeFutureMock, new State("abc")); | ||
| Mockito.verify(authorizationCodeFutureMock).complete("some authorization code"); | ||
| Assertions.assertEquals("Authorization completed successfully.", response); | ||
| } | ||
|
|
||
| @Test | ||
| public void shouldReturnRandomErrorResponse() { | ||
| Map<String, String> params = new HashMap<>(); | ||
| params.put("error", "some random error"); | ||
|
|
||
| String response = | ||
| AuthorizationCodeRedirectRequestHandler.handleRedirectRequest( | ||
| params, authorizationCodeFutureMock, new State("abc")); | ||
| Mockito.verify(authorizationCodeFutureMock) | ||
| .completeExceptionally(Mockito.any(SFException.class)); | ||
| Assertions.assertEquals("Authorization error: some random error", response); | ||
| } | ||
|
|
||
| @Test | ||
| public void shouldReturnInvalidStateErrorResponse() { | ||
| Map<String, String> params = new HashMap<>(); | ||
| params.put("authorization_code", "some authorization code"); | ||
| params.put("state", "invalid state"); | ||
|
|
||
| String response = | ||
| AuthorizationCodeRedirectRequestHandler.handleRedirectRequest( | ||
| params, authorizationCodeFutureMock, new State("abc")); | ||
| Mockito.verify(authorizationCodeFutureMock) | ||
| .completeExceptionally(Mockito.any(SFException.class)); | ||
| Assertions.assertEquals( | ||
| "Authorization error: invalid authorization request redirection state", response); | ||
| } | ||
|
|
||
| @Test | ||
| public void shouldReturnAuthorizationCodeAbsentErrorResponse() { | ||
| Map<String, String> params = new HashMap<>(); | ||
| params.put("state", "abc"); | ||
| params.put("some-random-param", "some-value"); | ||
|
|
||
| String response = | ||
| AuthorizationCodeRedirectRequestHandler.handleRedirectRequest( | ||
| params, authorizationCodeFutureMock, new State("abc")); | ||
| Mockito.verify(authorizationCodeFutureMock) | ||
| .completeExceptionally(Mockito.any(SFException.class)); | ||
| Assertions.assertEquals( | ||
| "Authorization error: authorization code has not been returned to the driver.", response); | ||
| } | ||
| } |