From 3cc5e0accff5da52161e2b9a27345577e1011310 Mon Sep 17 00:00:00 2001 From: Aleksey Vdovenko Date: Mon, 23 Sep 2024 15:58:54 +0200 Subject: [PATCH] feat: allowed routes to rewrite path and protect by roles (#491) --- .../com/epam/aidial/core/ProxyContext.java | 1 + .../com/epam/aidial/core/config/Route.java | 2 + .../core/controller/RouteController.java | 29 +++++++++++- .../com/epam/aidial/core/RouteApiTest.java | 47 +++++++++++++++++++ src/test/resources/aidial.config.json | 27 +++++++++++ 5 files changed, 104 insertions(+), 2 deletions(-) create mode 100644 src/test/java/com/epam/aidial/core/RouteApiTest.java diff --git a/src/main/java/com/epam/aidial/core/ProxyContext.java b/src/main/java/com/epam/aidial/core/ProxyContext.java index c6229ab7c..de060043e 100644 --- a/src/main/java/com/epam/aidial/core/ProxyContext.java +++ b/src/main/java/com/epam/aidial/core/ProxyContext.java @@ -67,6 +67,7 @@ public class ProxyContext { private List userRoles; private String userHash; private TokenUsage tokenUsage; + private boolean rewritePath; private UpstreamRoute upstreamRoute; private HttpClientRequest proxyRequest; private Map requestHeaders = Map.of(); diff --git a/src/main/java/com/epam/aidial/core/config/Route.java b/src/main/java/com/epam/aidial/core/config/Route.java index 85492a141..4facc260e 100644 --- a/src/main/java/com/epam/aidial/core/config/Route.java +++ b/src/main/java/com/epam/aidial/core/config/Route.java @@ -13,9 +13,11 @@ public class Route { private String name; private Response response; + private boolean rewritePath; private List paths = List.of(); private Set methods = Set.of(); private List upstreams = List.of(); + private Set userRoles = Set.of(); @Data public static class Response { diff --git a/src/main/java/com/epam/aidial/core/controller/RouteController.java b/src/main/java/com/epam/aidial/core/controller/RouteController.java index 6495bbb6b..f272b243f 100644 --- a/src/main/java/com/epam/aidial/core/controller/RouteController.java +++ b/src/main/java/com/epam/aidial/core/controller/RouteController.java @@ -23,8 +23,8 @@ import lombok.RequiredArgsConstructor; import lombok.SneakyThrows; import lombok.extern.slf4j.Slf4j; +import org.apache.http.client.utils.URIBuilder; -import java.net.URL; import java.util.List; import java.util.Objects; import java.util.Set; @@ -46,6 +46,13 @@ public Future handle() { return Future.succeededFuture(); } + if (!hasAccess(route)) { + log.error("Forbidden route {}. Trace: {}. Span: {}. Key: {}. User sub: {}.", + route.getName(), context.getTraceId(), context.getSpanId(), context.getProject(), context.getUserSub()); + context.respond(HttpStatus.FORBIDDEN, "Forbidden route"); + return Future.succeededFuture(); + } + Route.Response response = route.getResponse(); if (response == null) { UpstreamProvider upstreamProvider = new RouteEndpointProvider(route); @@ -57,6 +64,7 @@ public Future handle() { return Future.succeededFuture(); } + context.setRewritePath(route.isRewritePath()); context.setUpstreamRoute(upstreamRoute); } else { context.getResponse().setStatusCode(response.getStatus()); @@ -87,7 +95,7 @@ private Future sendRequest() { Upstream upstream = route.get(); Objects.requireNonNull(upstream); RequestOptions options = new RequestOptions() - .setAbsoluteURI(new URL(upstream.getEndpoint())) + .setAbsoluteURI(getEndpointUri(upstream)) .setMethod(request.method()); return proxy.getClient().request(options) @@ -241,4 +249,21 @@ private Route selectRoute() { return null; } + + @SneakyThrows + private String getEndpointUri(Upstream upstream) { + URIBuilder uriBuilder = new URIBuilder(upstream.getEndpoint()); + if (context.isRewritePath()) { + uriBuilder.setPath(context.getRequest().path()); + } + return uriBuilder.toString(); + } + + private boolean hasAccess(Route route) { + Set allowedRoles = route.getUserRoles(); + List actualRoles = context.getUserRoles(); + + return allowedRoles.isEmpty() || actualRoles.stream() + .anyMatch(allowedRoles::contains); + } } \ No newline at end of file diff --git a/src/test/java/com/epam/aidial/core/RouteApiTest.java b/src/test/java/com/epam/aidial/core/RouteApiTest.java new file mode 100644 index 000000000..bc24b5548 --- /dev/null +++ b/src/test/java/com/epam/aidial/core/RouteApiTest.java @@ -0,0 +1,47 @@ +package com.epam.aidial.core; + +import io.vertx.core.Vertx; +import io.vertx.core.http.HttpMethod; +import io.vertx.junit5.VertxExtension; +import io.vertx.junit5.VertxTestContext; +import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +import java.util.List; +import java.util.UUID; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +@ExtendWith(VertxExtension.class) +class RouteApiTest extends ResourceBaseTest { + + @ParameterizedTest + @MethodSource("datasource") + void route(HttpMethod method, String path, String apiKey, int expectedStatus, String expectedResponse, + Vertx vertx, VertxTestContext context) { + + var targetServer = vertx.createHttpServer() + .requestHandler(req -> req.response().end(req.path())); + targetServer.listen(9876) + .onComplete(context.succeedingThenComplete()); + + var reqBody = method == HttpMethod.POST ? UUID.randomUUID().toString() : null; + var resp = send(method, path, null, reqBody, "api-key", apiKey); + + assertEquals(expectedStatus, resp.status()); + assertEquals(expectedResponse, resp.body()); + } + + private static List datasource() { + return List.of( + Arguments.of(HttpMethod.GET, "/v1/plain", "vstore_user_key", 200, "/"), + Arguments.of(HttpMethod.GET, "/v1/plain", "vstore_admin_key", 200, "/"), + Arguments.of(HttpMethod.GET, "/v1/vector_store/1", "vstore_user_key", 200, "/v1/vector_store/1"), + Arguments.of(HttpMethod.GET, "/v1/vector_store/1", "vstore_admin_key", 200, "/v1/vector_store/1"), + Arguments.of(HttpMethod.POST, "/v1/vector_store/1", "vstore_user_key", 403, "Forbidden route"), + Arguments.of(HttpMethod.POST, "/v1/vector_store/1", "vstore_admin_key", 200, "/v1/vector_store/1") + ); + } +} diff --git a/src/test/resources/aidial.config.json b/src/test/resources/aidial.config.json index 1890a836c..0daac1574 100644 --- a/src/test/resources/aidial.config.json +++ b/src/test/resources/aidial.config.json @@ -6,6 +6,25 @@ "response" : { "status": 200 } + }, + "plain": { + "paths": ["/v1/plain"], + "methods": ["GET"], + "upstreams": [{"endpoint": "http://localhost:9876"}] + }, + "vector_store_query": { + "paths": ["/v1/vector_store(/[^/]+)*$"], + "rewritePath": true, + "methods": ["GET"], + "userRoles": ["vstore_user", "vstore_admin"], + "upstreams": [{"endpoint": "http://localhost:9876"}] + }, + "vector_store_mutation": { + "paths": ["/v1/vector_store(/[^/]+)*$"], + "rewritePath": true, + "methods": ["POST", "PUT", "DELETE"], + "userRoles": ["vstore_admin"], + "upstreams": [{"endpoint": "http://localhost:9876"}] } }, "addons": { @@ -101,6 +120,14 @@ "proxyKey2": { "project": "EPM-RTC-RAIL", "role": "default" + }, + "vstore_user_key": { + "project": "test", + "role": "vstore_user" + }, + "vstore_admin_key": { + "project": "test", + "role": "vstore_admin" } }, "roles": {