http server, add subject.doAs handler wrapper for exchange attribute …

prometheus-metrics-exporter-httpserver/src/main/java/io/prometheus/metrics/exporter/httpserver/HTTPServer.java

@@ -2,6 +2,7 @@
 
 import com.sun.net.httpserver.Authenticator;
 import com.sun.net.httpserver.HttpContext;
+import com.sun.net.httpserver.HttpExchange;
 import com.sun.net.httpserver.HttpHandler;
 import com.sun.net.httpserver.HttpServer;
 import com.sun.net.httpserver.HttpsConfigurator;
@@ -10,14 +11,18 @@
 import io.prometheus.metrics.model.registry.PrometheusRegistry;
 import java.io.Closeable;
 import java.io.IOException;
+import java.io.InputStream;
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
+import java.security.PrivilegedActionException;
+import java.security.PrivilegedExceptionAction;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.RejectedExecutionHandler;
 import java.util.concurrent.SynchronousQueue;
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
+import javax.security.auth.Subject;
 
 /**
  * Expose Prometheus metrics using a plain Java HttpServer.
@@ -51,16 +56,25 @@ private HTTPServer(
       HttpServer httpServer,
       PrometheusRegistry registry,
       Authenticator authenticator,
+      String authenticatedSubjectAttributeName,
       HttpHandler defaultHandler) {
     if (httpServer.getAddress() == null) {
       throw new IllegalArgumentException("HttpServer hasn't been bound to an address");
     }
     this.server = httpServer;
     this.executorService = executorService;
     registerHandler(
-        "/", defaultHandler == null ? new DefaultHandler() : defaultHandler, authenticator);
-    registerHandler("/metrics", new MetricsHandler(config, registry), authenticator);
-    registerHandler("/-/healthy", new HealthyHandler(), authenticator);
+        "/",
+        defaultHandler == null ? new DefaultHandler() : defaultHandler,
+        authenticator,
+        authenticatedSubjectAttributeName);
+    registerHandler(
+        "/metrics",
+        new MetricsHandler(config, registry),
+        authenticator,
+        authenticatedSubjectAttributeName);
+    registerHandler(
+        "/-/healthy", new HealthyHandler(), authenticator, authenticatedSubjectAttributeName);
     try {
       // HttpServer.start() starts the HttpServer in a new background thread.
       // If we call HttpServer.start() from a thread of the executorService,
@@ -74,13 +88,54 @@ private HTTPServer(
     }
   }
 
-  private void registerHandler(String path, HttpHandler handler, Authenticator authenticator) {
-    HttpContext context = server.createContext(path, handler);
+  private void registerHandler(
+      String path, HttpHandler handler, Authenticator authenticator, String subjectAttributeName) {
+    HttpContext context = server.createContext(path, wrapWithDoAs(handler, subjectAttributeName));
     if (authenticator != null) {
       context.setAuthenticator(authenticator);
     }
   }
 
+  private HttpHandler wrapWithDoAs(HttpHandler handler, String subjectAttributeName) {
+    if (subjectAttributeName == null) {
+      return handler;
+    }
+
+    // invoke handler using the subject.doAs from the named attribute
+    return new HttpHandler() {
+      @Override
+      public void handle(HttpExchange exchange) throws IOException {
+        Object authSubject = exchange.getAttribute(subjectAttributeName);
+        if (authSubject instanceof Subject) {
+          try {
+            Subject.doAs(
+                (Subject) authSubject,
+                (PrivilegedExceptionAction<IOException>)
+                    () -> {
+                      handler.handle(exchange);
+                      return null;
+                    });
+          } catch (PrivilegedActionException e) {
+            if (e.getException() != null) {
+              throw new IOException(e.getException());
+            } else throw new IOException(e);
+          }
+        } else {
+          drainInputAndClose(exchange);
+          exchange.sendResponseHeaders(403, -1);
+        }
+      }
+    };
+  }
+
+  private void drainInputAndClose(HttpExchange httpExchange) throws IOException {
+    InputStream inputStream = httpExchange.getRequestBody();
+    byte[] b = new byte[4096];
+    while (inputStream.read(b) != -1)
+      ;
+    inputStream.close();
+  }
+
   /** Stop the HTTP server. Same as {@link #close()}. */
   public void stop() {
     close();
@@ -120,6 +175,7 @@ public static class Builder {
     private Authenticator authenticator = null;
     private HttpsConfigurator httpsConfigurator = null;
     private HttpHandler defaultHandler = null;
+    private String authenticatedSubjectAttributeName = null;
 
     private Builder(PrometheusProperties config) {
       this.config = config;
@@ -171,6 +227,12 @@ public Builder authenticator(Authenticator authenticator) {
       return this;
     }
 
+    /** Optional: the attribute name of a Subject from a custom authenticator. */
+    public Builder authenticatedSubjectAttributeName(String authenticatedSubjectAttributeName) {
+      this.authenticatedSubjectAttributeName = authenticatedSubjectAttributeName;
+      return this;
+    }
+
     /** Optional: {@link HttpsConfigurator} for TLS/SSL */
     public Builder httpsConfigurator(HttpsConfigurator configurator) {
       this.httpsConfigurator = configurator;
@@ -201,7 +263,13 @@ public HTTPServer buildAndStart() throws IOException {
       ExecutorService executorService = makeExecutorService();
       httpServer.setExecutor(executorService);
       return new HTTPServer(
-          config, executorService, httpServer, registry, authenticator, defaultHandler);
+          config,
+          executorService,
+          httpServer,
+          registry,
+          authenticator,
+          authenticatedSubjectAttributeName,
+          defaultHandler);
     }
 
     private InetSocketAddress makeInetSocketAddress() {

prometheus-metrics-exporter-httpserver/src/test/java/io/prometheus/metrics/exporter/httpserver/HTTPServerTest.java

@@ -0,0 +1,79 @@
+package io.prometheus.metrics.exporter.httpserver;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import com.sun.net.httpserver.Authenticator;
+import com.sun.net.httpserver.HttpExchange;
+import com.sun.net.httpserver.HttpHandler;
+import com.sun.net.httpserver.HttpPrincipal;
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.net.Socket;
+import java.nio.charset.StandardCharsets;
+import java.security.AccessController;
+import java.security.Principal;
+import javax.security.auth.Subject;
+import org.junit.jupiter.api.Test;
+
+public class HTTPServerTest {
+
+  @Test
+  public void testSubjectDoAs() throws Exception {
+
+    final String user = "joe";
+    final Subject subject = new Subject();
+    subject.getPrincipals().add(() -> (user));
+
+    Authenticator authenticator =
+        new Authenticator() {
+          @Override
+          public Result authenticate(HttpExchange exchange) {
+            exchange.setAttribute("aa", subject);
+            return new Success(new HttpPrincipal(user, "/"));
+          }
+        };
+
+    HttpHandler handler =
+        exchange -> {
+          boolean found = false;
+          Subject current = Subject.getSubject(AccessController.getContext());
+          for (Principal p : current.getPrincipals()) {
+            if (user.equals(p.getName())) {
+              found = true;
+              break;
+            }
+          }
+          if (!found) {
+            throw new IOException("Expected validated user joe!");
+          }
+          exchange.sendResponseHeaders(204, -1);
+        };
+    HTTPServer server =
+        HTTPServer.builder()
+            .port(0)
+            .authenticator(authenticator)
+            .defaultHandler(handler)
+            .authenticatedSubjectAttributeName("aa")
+            .buildAndStart();
+
+    Socket socket = new Socket();
+    try {
+      socket.connect(new InetSocketAddress("localhost", server.getPort()));
+
+      socket.getOutputStream().write("GET / HTTP/1.1 \r\n".getBytes(StandardCharsets.UTF_8));
+      socket.getOutputStream().write("HOST: localhost \r\n\r\n".getBytes(StandardCharsets.UTF_8));
+      socket.getOutputStream().flush();
+
+      String actualResponse = "";
+      byte[] resp = new byte[500];
+      int read = socket.getInputStream().read(resp, 0, resp.length);
+      if (read > 0) {
+        actualResponse = new String(resp, 0, read);
+      }
+      assertThat(actualResponse).contains("204");
+
+    } finally {
+      socket.close();
+    }
+  }
+}