Merge pull request #29373 from vespa-engine/bjorncs/restapi

Bjorncs/restapi

11 files changed, 70 insertions, 3 deletions

diff --git a/application/abi-spec.json b/application/abi-spec.json
index 95a9d2a524a..27b01d02b5f 100644
--- a/application/abi-spec.json
+++ b/application/abi-spec.json
@@ -324,6 +324,7 @@
       "public com.yahoo.application.container.Processing processing()",
       "public com.yahoo.application.container.DocumentProcessing documentProcessing()",
       "public com.yahoo.component.provider.ComponentRegistry components()",
+      "public com.yahoo.component.provider.ComponentRegistry handlers()",
       "public com.yahoo.application.container.handler.Response handleRequest(com.yahoo.application.container.handler.Request)",
       "public void close()"
     ],
@@ -450,11 +451,13 @@
       "public void <init>(java.lang.String, java.lang.String)",
       "public void <init>(java.lang.String, byte[], com.yahoo.application.container.handler.Request$Method)",
       "public void <init>(java.lang.String, byte[], com.yahoo.application.container.handler.Request$Method, java.security.Principal)",
+      "public void <init>(java.lang.String, byte[], com.yahoo.application.container.handler.Request$Method, java.security.Principal, java.net.SocketAddress)",
       "public void <init>(java.lang.String, java.lang.String, com.yahoo.application.container.handler.Request$Method)",
       "public com.yahoo.application.container.handler.Headers getHeaders()",
       "public byte[] getBody()",
       "public java.lang.String getUri()",
       "public java.util.Map getAttributes()",
+      "public java.util.Optional remoteAddress()",
       "public java.lang.String toString()",
       "public com.yahoo.application.container.handler.Request$Method getMethod()",
       "public java.util.Optional getUserPrincipal()"
diff --git a/application/src/main/java/com/yahoo/application/container/JDisc.java b/application/src/main/java/com/yahoo/application/container/JDisc.java
index 162a5f343a1..3223c8d31c9 100644
--- a/application/src/main/java/com/yahoo/application/container/JDisc.java
+++ b/application/src/main/java/com/yahoo/application/container/JDisc.java
@@ -169,6 +169,9 @@ public final class JDisc implements AutoCloseable {
         return container.getComponentRegistry();
     }
 
+    /** @return registry of all request handlers configured */
+    public ComponentRegistry<RequestHandler> handlers() { return container.getRequestHandlerRegistry(); }
+
     /**
      * Handles the given {@link com.yahoo.application.container.handler.Request} by passing it to the {@link RequestHandler}
      * that is bound to the request's URI.
diff --git a/application/src/main/java/com/yahoo/application/container/SynchronousRequestResponseHandler.java b/application/src/main/java/com/yahoo/application/container/SynchronousRequestResponseHandler.java
index c54b3f60cf9..1b4862c75c0 100644
--- a/application/src/main/java/com/yahoo/application/container/SynchronousRequestResponseHandler.java
+++ b/application/src/main/java/com/yahoo/application/container/SynchronousRequestResponseHandler.java
@@ -90,6 +90,7 @@ final class SynchronousRequestResponseHandler {
                                                                             URI.create(request.getUri()),
                                                                             com.yahoo.jdisc.http.HttpRequest.Method.valueOf(request.getMethod().name()));
             request.getUserPrincipal().ifPresent(httpRequest::setUserPrincipal);
+            request.remoteAddress().ifPresent(httpRequest::setRemoteAddress);
             discRequest = httpRequest;
         } else {
             discRequest = new com.yahoo.jdisc.Request(currentContainer, URI.create(request.getUri()));
diff --git a/application/src/main/java/com/yahoo/application/container/handler/Request.java b/application/src/main/java/com/yahoo/application/container/handler/Request.java
index b6dff44269b..d877258cb15 100644
--- a/application/src/main/java/com/yahoo/application/container/handler/Request.java
+++ b/application/src/main/java/com/yahoo/application/container/handler/Request.java
@@ -3,6 +3,7 @@ package com.yahoo.application.container.handler;
 
 import com.yahoo.api.annotations.Beta;
 
+import java.net.SocketAddress;
 import java.nio.charset.StandardCharsets;
 import java.security.Principal;
 import java.util.Map;
@@ -24,6 +25,7 @@ public class Request {
     private final Method method;
     private final Map<String, Object> attributes = new ConcurrentHashMap<>();
     private final Principal userPrincipal;
+    private final SocketAddress remoteAddress;
 
     /**
      * Creates a Request with an empty body.
@@ -73,10 +75,24 @@ public class Request {
      * @param principal the user principal of the request
      */
     public Request(String uri, byte[] body, Method method, Principal principal) {
+        this(uri, body, method, principal, null);
+    }
+
+    /**
+     * Creates a Request with a message body, method and user principal.
+     *
+     * @param uri the URI of the request
+     * @param body the body of the request
+     * @param method the method of the request
+     * @param principal the user principal of the request
+     * @param remoteAddress the remote address of the request
+     */
+    public Request(String uri, byte[] body, Method method, Principal principal, SocketAddress remoteAddress) {
         this.uri = uri;
         this.body = body;
         this.method = method;
         this.userPrincipal = principal;
+        this.remoteAddress = remoteAddress;
     }
 
     /**
@@ -122,6 +138,8 @@ public class Request {
         return attributes;
     }
 
+    public Optional<SocketAddress> remoteAddress() { return Optional.ofNullable(remoteAddress); }
+
     @Override
     public String toString() {
         String bodyStr = (body == null || body.length == 0) ? "[empty]" : "[omitted]";
diff --git a/container-core/abi-spec.json b/container-core/abi-spec.json
index 3c11abb9d08..d07d21ae71f 100644
--- a/container-core/abi-spec.json
+++ b/container-core/abi-spec.json
@@ -660,6 +660,10 @@
       "public com.yahoo.container.jdisc.HttpRequestBuilder withRequestContent(java.io.InputStream)",
       "public com.yahoo.container.jdisc.HttpRequestBuilder withScheme(java.lang.String)",
       "public com.yahoo.container.jdisc.HttpRequestBuilder withHostname(java.lang.String)",
+      "public com.yahoo.container.jdisc.HttpRequestBuilder withPrincipal(java.security.Principal)",
+      "public com.yahoo.container.jdisc.HttpRequestBuilder withRemoteAddress(java.net.SocketAddress)",
+      "public com.yahoo.container.jdisc.HttpRequestBuilder withAttribute(java.lang.String, java.lang.Object)",
+      "public com.yahoo.container.jdisc.HttpRequestBuilder withPort(int)",
       "public com.yahoo.container.jdisc.HttpRequest build()"
     ],
     "fields" : [ ]
diff --git a/container-core/src/main/java/com/yahoo/container/jdisc/HttpRequestBuilder.java b/container-core/src/main/java/com/yahoo/container/jdisc/HttpRequestBuilder.java
index a2d792e6ae0..147f388e08c 100644
--- a/container-core/src/main/java/com/yahoo/container/jdisc/HttpRequestBuilder.java
+++ b/container-core/src/main/java/com/yahoo/container/jdisc/HttpRequestBuilder.java
@@ -4,6 +4,8 @@ package com.yahoo.container.jdisc;
 import com.yahoo.jdisc.http.HttpRequest.Method;
 
 import java.io.InputStream;
+import java.net.SocketAddress;
+import java.security.Principal;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -19,9 +21,13 @@ public class HttpRequestBuilder {
     private final String path;
     private final Map<String, List<String>> queryParameters = new TreeMap<>();
     private final Map<String, String> headers = new TreeMap<>();
+    private final Map<String, Object> attributes = new TreeMap<>();
     private String scheme;
     private String hostname;
     private InputStream content;
+    private Principal principal;
+    private SocketAddress socketAddress;
+    private int port = -1;
 
     private HttpRequestBuilder(Method method, String path) {
         this.method = method;
@@ -43,10 +49,20 @@ public class HttpRequestBuilder {
 
     public HttpRequestBuilder withHostname(String hostname) { this.hostname = hostname; return this; }
 
+    public HttpRequestBuilder withPrincipal(Principal p) { principal = p; return this; }
+
+    public HttpRequestBuilder withRemoteAddress(SocketAddress sa) { socketAddress = sa; return this; }
+
+    public HttpRequestBuilder withAttribute(String name, Object value) { attributes.put(name, value); return this; }
+
+    public HttpRequestBuilder withPort(int port) { this.port = port; return this; }
+
     public HttpRequest build() {
         String scheme = this.scheme != null ? this.scheme : "http";
         String hostname = this.hostname != null ? this.hostname : "localhost";
-        StringBuilder uriBuilder = new StringBuilder(scheme).append("://").append(hostname).append(path);
+        StringBuilder uriBuilder = new StringBuilder(scheme).append("://").append(hostname);
+        if (port > 0) uriBuilder.append(':').append(port);
+        uriBuilder.append(path);
         if (queryParameters.size() > 0) {
             uriBuilder.append('?');
             queryParameters.forEach((name, values) -> {
@@ -66,6 +82,9 @@ public class HttpRequestBuilder {
             request = HttpRequest.createTestRequest(uriBuilder.toString(), method);
         }
         headers.forEach((name, value) -> request.getJDiscRequest().headers().put(name, value));
+        if (principal != null) request.getJDiscRequest().setUserPrincipal(principal);
+        if (socketAddress != null) request.getJDiscRequest().setRemoteAddress(socketAddress);
+        request.getJDiscRequest().context().putAll(attributes);
         return request;
     }
 }
diff --git a/container-core/src/main/java/com/yahoo/restapi/RestApi.java b/container-core/src/main/java/com/yahoo/restapi/RestApi.java
index 18d8d8c49b4..ee5628988c9 100644
--- a/container-core/src/main/java/com/yahoo/restapi/RestApi.java
+++ b/container-core/src/main/java/com/yahoo/restapi/RestApi.java
@@ -15,6 +15,7 @@ import com.yahoo.security.tls.ConnectionAuthContext;
 
 import javax.net.ssl.SSLSession;
 import java.io.InputStream;
+import java.net.InetSocketAddress;
 import java.security.Principal;
 import java.util.List;
 import java.util.Optional;
@@ -153,6 +154,7 @@ public interface RestApi {
         Principal userPrincipalOrThrow();
         Optional<SSLSession> sslSession();
         Optional<ConnectionAuthContext> connectionAuthContext();
+        InetSocketAddress remoteAddress();
 
         interface Parameters {
             Optional<String> getString(String name);
@@ -193,6 +195,7 @@ public interface RestApi {
     interface FilterContext {
         RequestContext requestContext();
         String route();
+        void setPrincipal(Principal principal);
         HttpResponse executeNext();
     }
 }
diff --git a/container-core/src/main/java/com/yahoo/restapi/RestApiException.java b/container-core/src/main/java/com/yahoo/restapi/RestApiException.java
index e3acf4258f1..3a44bc4da5f 100644
--- a/container-core/src/main/java/com/yahoo/restapi/RestApiException.java
+++ b/container-core/src/main/java/com/yahoo/restapi/RestApiException.java
@@ -57,6 +57,7 @@ public class RestApiException extends RuntimeException {
     }
 
     public static class BadRequest extends RestApiException {
+        public BadRequest() { this("Bad request"); }
         public BadRequest(String message) { this(message, null); }
         public BadRequest(Throwable cause) { this(cause.getMessage(), cause); }
         public BadRequest(String message, Throwable cause) { super(ErrorResponse::badRequest, message, cause); }
@@ -69,6 +70,7 @@ public class RestApiException extends RuntimeException {
     }
 
     public static class Forbidden extends RestApiException {
+        public Forbidden() { this("Forbidden"); }
         public Forbidden(String message) { super(ErrorResponse::forbidden, message, null); }
         public Forbidden(String message, Throwable cause) { super(ErrorResponse::forbidden, message, cause); }
     }
diff --git a/container-core/src/main/java/com/yahoo/restapi/RestApiImpl.java b/container-core/src/main/java/com/yahoo/restapi/RestApiImpl.java
index 196a57d23bf..090e06c221f 100644
--- a/container-core/src/main/java/com/yahoo/restapi/RestApiImpl.java
+++ b/container-core/src/main/java/com/yahoo/restapi/RestApiImpl.java
@@ -21,6 +21,7 @@ import com.yahoo.security.tls.TransportSecurityUtils;
 
 import javax.net.ssl.SSLSession;
 import java.io.InputStream;
+import java.net.InetSocketAddress;
 import java.net.URI;
 import java.security.Principal;
 import java.util.ArrayList;
@@ -491,7 +492,7 @@ class RestApiImpl implements RestApi {
         @Override public Optional<ConnectionAuthContext> connectionAuthContext() {
             return sslSession().flatMap(TransportSecurityUtils::getConnectionAuthContext);
         }
-
+        @Override public InetSocketAddress remoteAddress() { return (InetSocketAddress) request.getJDiscRequest().getRemoteAddress(); }
 
         private class PathParametersImpl implements RestApi.RequestContext.PathParameters {
             @Override
@@ -558,6 +559,7 @@ class RestApiImpl implements RestApi {
 
         @Override public RestApi.RequestContext requestContext() { return requestContext; }
         @Override public String route() { return route.name != null ? route.name : route.pathPattern; }
+        @Override public void setPrincipal(Principal p) { requestContext.request.getJDiscRequest().setUserPrincipal(p); }
 
         HttpResponse executeFirst() { return filter.filterRequest(this); }
 
diff --git a/container-core/src/test/java/com/yahoo/restapi/RestApiImplTest.java b/container-core/src/test/java/com/yahoo/restapi/RestApiImplTest.java
index 95ccff1a075..d27e04bbd7a 100644
--- a/container-core/src/test/java/com/yahoo/restapi/RestApiImplTest.java
+++ b/container-core/src/test/java/com/yahoo/restapi/RestApiImplTest.java
@@ -182,6 +182,18 @@ class RestApiImplTest {
                             "{\"error-code\":\"BAD_REQUEST\",\"message\":\"Query parameter 'missing' is missing\"}");
     }
 
+    @Test
+    void principal_from_filter_is_visible_to_handler() {
+        var restApi = RestApi.builder()
+                .addRoute(route("/api1").get(ctx -> ctx.userPrincipalOrThrow().getName()))
+                .addFilter(ctx -> {
+                    ctx.setPrincipal(() -> "my-principal-name");
+                    return ctx.executeNext();
+                })
+                .build();
+        verifyJsonResponse(restApi, Method.GET, "/api1", null, 200, "{\"message\":\"my-principal-name\"}");
+    }
+
     private static void verifyJsonResponse(
             RestApi restApi, Method method, String path, String requestContent, int expectedStatusCode,
             String expectedJson) {
diff --git a/searchlib/abi-spec.json b/searchlib/abi-spec.json
index 0b1cb7a103c..2d67abb0e04 100644
--- a/searchlib/abi-spec.json
+++ b/searchlib/abi-spec.json
@@ -1728,4 +1728,4 @@
     ],
     "fields" : [ ]
   }
-}
+}
\ No newline at end of file