Skip to content

Commit ce4aee3

Browse files
DmytroZaichenkoDevDmytro Zaichenko
andauthored
fix: Static OAuth endpoints are overwritten by metadata discovery results (#1414)
Co-authored-by: Dmytro Zaichenko <dmytro_zaichenko@epam.com>
1 parent 926c857 commit ce4aee3

File tree

3 files changed

+140
-7
lines changed

3 files changed

+140
-7
lines changed

credentials/src/main/java/com/epam/aidial/core/credentials/service/registration/StaticResourceRegistrationStrategy.java

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ public class StaticResourceRegistrationStrategy implements ResourceRegistrationS
5050
* @param resourceAuthSettings The static authentication settings provided for the resource.
5151
* @return A {@link ClientRegistration} containing the complete registration details.
5252
* @throws IllegalArgumentException If any of the required endpoints (authorizationEndpoint,
53-
* tokenEndpoint, codeChallengeMethod) cannot be resolved.
53+
* tokenEndpoint) cannot be resolved.
5454
*/
5555
@Override
5656
public ClientRegistration register(String resourceId, String resourceEndpoint, ResourceAuthSettings resourceAuthSettings) {
@@ -78,12 +78,12 @@ public ClientRegistration register(String resourceId, String resourceEndpoint, R
7878
}
7979

8080
if (authServerMetadata != null) {
81-
authorizationEndpoint = Optional.ofNullable(authServerMetadata.getAuthorizationEndpoint())
82-
.orElse(authorizationEndpoint);
83-
tokenEndpoint = Optional.ofNullable(authServerMetadata.getTokenEndpoint())
84-
.orElse(tokenEndpoint);
85-
codeChallengeMethod = getCodeChallengeMethod(authServerMetadata)
86-
.orElse(codeChallengeMethod);
81+
authorizationEndpoint = Optional.ofNullable(authorizationEndpoint)
82+
.orElse(authServerMetadata.getAuthorizationEndpoint());
83+
tokenEndpoint = Optional.ofNullable(tokenEndpoint)
84+
.orElse(authServerMetadata.getTokenEndpoint());
85+
codeChallengeMethod = Optional.ofNullable(codeChallengeMethod)
86+
.orElse(getCodeChallengeMethod(authServerMetadata).orElse(null));
8787
}
8888
}
8989

credentials/src/test/java/com/epam/aidial/core/credentials/service/registration/StaticResourceRegistrationStrategyTest.java

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
import static org.junit.jupiter.api.Assertions.assertEquals;
1818
import static org.junit.jupiter.api.Assertions.assertNotNull;
19+
import static org.junit.jupiter.api.Assertions.assertNull;
1920
import static org.mockito.Mockito.mock;
2021
import static org.mockito.Mockito.verify;
2122
import static org.mockito.Mockito.verifyNoInteractions;
@@ -122,4 +123,68 @@ void testCreateStaticResourceRegistrationWithDiscoveredData() {
122123
verify(authorizationServerMetadataService).getAuthorizationServerMetadata(
123124
resourceId, resourceEndpoint, protectedResourceMetadata, false);
124125
}
126+
127+
@Test
128+
void testStaticEndpointsNotOverwrittenByDiscovery() {
129+
// Given
130+
String resourceId = "staticResource";
131+
String resourceEndpoint = "https://test.endpoint.com";
132+
String clientId = "staticClientId";
133+
String clientSecret = "staticClientSecret";
134+
String redirectUri = "https://static.redirect.uri";
135+
136+
// Static endpoints provided by the user
137+
String staticAuthorizationEndpoint = "https://static.auth.endpoint";
138+
String staticTokenEndpoint = "https://static.token.endpoint";
139+
List<String> scopesSupported = List.of("scope1", "scope2");
140+
141+
// Discovery returns different endpoints and a codeChallengeMethod
142+
String discoveredAuthorizationEndpoint = "https://discovered.auth.endpoint";
143+
String discoveredTokenEndpoint = "https://discovered.token.endpoint";
144+
String discoveredCodeChallengeMethod = "S256";
145+
146+
// codeChallengeMethod is null (simulates UI-created toolset), triggering discovery
147+
ResourceAuthSettings resourceAuthSettings = ResourceAuthSettings.builder()
148+
.clientId(clientId)
149+
.clientSecret(clientSecret)
150+
.redirectUri(redirectUri)
151+
.authorizationEndpoint(staticAuthorizationEndpoint)
152+
.tokenEndpoint(staticTokenEndpoint)
153+
.scopesSupported(scopesSupported)
154+
.build();
155+
156+
assertNull(resourceAuthSettings.getCodeChallengeMethod());
157+
158+
AuthorizationServerProtectedResourceMetadata protectedResourceMetadata = mock(AuthorizationServerProtectedResourceMetadata.class);
159+
when(protectedResourceMetadataService.getProtectedResourceMetadata(resourceId, resourceEndpoint))
160+
.thenReturn(protectedResourceMetadata);
161+
162+
AuthorizationServerMetadata authorizationServerMetadata = mock(AuthorizationServerMetadata.class);
163+
when(authorizationServerMetadata.getAuthorizationEndpoint()).thenReturn(discoveredAuthorizationEndpoint);
164+
when(authorizationServerMetadata.getTokenEndpoint()).thenReturn(discoveredTokenEndpoint);
165+
when(authorizationServerMetadata.getCodeChallengeMethodsSupported()).thenReturn(List.of(discoveredCodeChallengeMethod));
166+
167+
when(authorizationServerMetadataService.getAuthorizationServerMetadata(
168+
resourceId, resourceEndpoint, protectedResourceMetadata, false))
169+
.thenReturn(authorizationServerMetadata);
170+
171+
// When
172+
ClientRegistration result = resourceRegistrationStrategy.register(resourceId, resourceEndpoint, resourceAuthSettings);
173+
174+
// Then
175+
assertNotNull(result);
176+
assertEquals(clientId, result.getClientId());
177+
assertEquals(clientSecret, result.getClientSecret());
178+
assertEquals(redirectUri, result.getRedirectUri());
179+
// Static endpoints must be preserved (not overwritten by discovery)
180+
assertEquals(staticAuthorizationEndpoint, result.getAuthorizationEndpoint());
181+
assertEquals(staticTokenEndpoint, result.getTokenEndpoint());
182+
// codeChallengeMethod should be filled from discovery
183+
assertEquals(discoveredCodeChallengeMethod, result.getCodeChallengeMethod());
184+
assertEquals(scopesSupported, result.getScopesSupported());
185+
// Verify discovery was called (because codeChallengeMethod was null)
186+
verify(protectedResourceMetadataService).getProtectedResourceMetadata(resourceId, resourceEndpoint);
187+
verify(authorizationServerMetadataService).getAuthorizationServerMetadata(
188+
resourceId, resourceEndpoint, protectedResourceMetadata, false);
189+
}
125190
}

server/src/test/java/com/epam/aidial/core/server/ToolSetApiTest.java

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1090,6 +1090,74 @@ void testSignOutFromApiCreatedToolsetWithIncorrectAuthType() {
10901090
verify(response, 400, "Wrong authentication_type. Expected type: API_KEY, provided: OAUTH");
10911091
}
10921092

1093+
@Test
1094+
void testCreateToolSetWithOauthStaticEndpointsPreservedAfterDiscovery() throws JsonProcessingException {
1095+
String requestBody = """
1096+
{
1097+
"endpoint": "http://localhost:9876/mcp",
1098+
"transport": "HTTP",
1099+
"allowedTools": [],
1100+
"auth_settings": {
1101+
"authentication_type": "OAUTH",
1102+
"client_id": "my-client-id",
1103+
"client_secret": "my-client-secret",
1104+
"redirect_uri": "http://localhost:3000/auth/signin",
1105+
"authorization_endpoint": "https://static.auth.example.com/authorize",
1106+
"token_endpoint": "https://static.auth.example.com/token"
1107+
}
1108+
}
1109+
""";
1110+
1111+
String protectedResourceMetadata = """
1112+
{
1113+
"resource": "http://localhost:9876/mcp",
1114+
"authorization_servers": ["http://localhost:9876"],
1115+
"scopes_supported": ["read", "write"]
1116+
}
1117+
""";
1118+
1119+
String authServerMetadata = """
1120+
{
1121+
"issuer": "http://localhost:9876",
1122+
"authorization_endpoint": "https://discovered.auth.example.com/authorize",
1123+
"token_endpoint": "https://discovered.auth.example.com/token",
1124+
"code_challenge_methods_supported": ["S256"]
1125+
}
1126+
""";
1127+
1128+
try (TestWebServer server = new TestWebServer(9876)) {
1129+
// MCP endpoint returns 401 without WWW-Authenticate to trigger well-known discovery
1130+
server.map(HttpMethod.POST, "/mcp", 401, "");
1131+
1132+
// Protected resource metadata discovery
1133+
server.map(HttpMethod.GET, "/.well-known/oauth-protected-resource/mcp",
1134+
200, protectedResourceMetadata, "Content-Type", "application/json");
1135+
1136+
// Authorization server metadata discovery (on the auth server from authorization_servers)
1137+
server.map(HttpMethod.GET, "/.well-known/oauth-authorization-server",
1138+
200, authServerMetadata, "Content-Type", "application/json");
1139+
1140+
Response response = send(HttpMethod.PUT, "/v1/toolsets/4X25dj1mja51jykqxsXnCH/toolset-oauth-static@",
1141+
null, requestBody, "authorization", "admin");
1142+
1143+
assertEquals(200, response.status());
1144+
1145+
// GET the created toolset to verify auth_settings
1146+
response = send(HttpMethod.GET, "/v1/toolsets/4X25dj1mja51jykqxsXnCH/toolset-oauth-static@",
1147+
null, null, "authorization", "admin");
1148+
1149+
assertEquals(200, response.status());
1150+
JsonNode toolset = ProxyUtil.MAPPER.readTree(response.body());
1151+
JsonNode authSettings = toolset.get("auth_settings");
1152+
1153+
// Static endpoints must be preserved (not overwritten by discovered ones)
1154+
assertEquals("https://static.auth.example.com/authorize", authSettings.get("authorization_endpoint").asText());
1155+
assertEquals("https://static.auth.example.com/token", authSettings.get("token_endpoint").asText());
1156+
// codeChallengeMethod should be filled from discovery
1157+
assertEquals("S256", authSettings.get("code_challenge_method").asText());
1158+
}
1159+
}
1160+
10931161
private String replaceValueInJsonArray(
10941162
String originalJson,
10951163
String arrayFieldName,

0 commit comments

Comments
 (0)