From 75dd8d7979691d193d3e0d3ccaec91255597ab28 Mon Sep 17 00:00:00 2001
From: Mark Elliot <mark.elliot@gmail.com>
Date: Wed, 1 Nov 2023 21:24:34 -0400
Subject: [PATCH] Minor cleanup

---
 .../main/java/com/markelliot/barista/Server.java    | 13 ++++---------
 .../markelliot/barista/handlers/HandlerChain.java   |  7 +++++++
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/barista/src/main/java/com/markelliot/barista/Server.java b/barista/src/main/java/com/markelliot/barista/Server.java
index 3db4153..7a0c442 100644
--- a/barista/src/main/java/com/markelliot/barista/Server.java
+++ b/barista/src/main/java/com/markelliot/barista/Server.java
@@ -186,16 +186,11 @@ public Server start() {
                 Spans.register("barista", span -> tracing.info("TRACING {}", span));
             }
 
-            HandlerChain handlerChain = HandlerChain.of(DispatchFromIoThreadHandler::new)
+            HttpHandler handler = HandlerChain.of(DispatchFromIoThreadHandler::new)
                     .then(h -> new CorsHandler(allowAllOrigins, allowedOrigins, h))
-                    .then(h -> new TracingHandler(tracingRate, h));
-
-            if (strictTransportSecurity) {
-                handlerChain.then(StrictTransportSecurityHandler::new);
-            }
-
-            HttpHandler handler = handlerChain.last(
-                    new EndpointHandlerBuilder(serde, authz, fallbackHandler).build(endpointHandlers));
+                    .then(h -> new TracingHandler(tracingRate, h))
+                    .then(StrictTransportSecurityHandler::new, strictTransportSecurity)
+                    .last(new EndpointHandlerBuilder(serde, authz, fallbackHandler).build(endpointHandlers));
             GracefulShutdownHandler shutdownHandler = new GracefulShutdownHandler(handler);
             Undertow undertow = Undertow.builder()
                     .setHandler(new DispatchFromIoThreadHandler(shutdownHandler))
diff --git a/barista/src/main/java/com/markelliot/barista/handlers/HandlerChain.java b/barista/src/main/java/com/markelliot/barista/handlers/HandlerChain.java
index 6e9f754..6a3b882 100644
--- a/barista/src/main/java/com/markelliot/barista/handlers/HandlerChain.java
+++ b/barista/src/main/java/com/markelliot/barista/handlers/HandlerChain.java
@@ -53,6 +53,13 @@ public HandlerChain then(DelegatingHandler next) {
         return this;
     }
 
+    public HandlerChain then(DelegatingHandler next, boolean onlyIf) {
+        if (onlyIf) {
+            handlers.add(next);
+        }
+        return this;
+    }
+
     public HandlerChain then(Set<DelegatingHandler> next) {
         for (DelegatingHandler fn : next) {
             then(fn);