feat: allowed routes to rewrite path and protect by roles (#491)

epam · Sep 23, 2024 · 3cc5e0a · 3cc5e0a
1 parent 7640323
commit 3cc5e0a
Show file tree

Hide file tree

Showing 5 changed files with 104 additions and 2 deletions.
diff --git 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<String> userRoles;
     private String userHash;
     private TokenUsage tokenUsage;
+    private boolean rewritePath;
     private UpstreamRoute upstreamRoute;
     private HttpClientRequest proxyRequest;
     private Map<String, String> 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
@@ -13,9 +13,11 @@ public class Route {
 
     private String name;
     private Response response;
+    private boolean rewritePath;
     private List<Pattern> paths = List.of();
     private Set<HttpMethod> methods = Set.of();
     private List<Upstream> upstreams = List.of();
+    private Set<String> 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
@@ -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<String> allowedRoles = route.getUserRoles();
+        List<String> actualRoles = context.getUserRoles();
+
+        return allowedRoles.isEmpty() || actualRoles.stream()
+                .anyMatch(allowedRoles::contains);
+    }
 }
diff --git a/src/test/java/com/epam/aidial/core/RouteApiTest.java 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<Arguments> 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
@@ -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": {