package org.restheart.handlers;

import com.google.common.annotations.VisibleForTesting;

import org.restheart.Bootstrapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;

import java.util.List;
import java.util.Optional;
import java.util.Random;

import io.undertow.server.HttpServerExchange;
import io.undertow.util.HeaderValues;
import io.undertow.util.HexConverter;
import io.undertow.util.HttpString;

/**
 * Handler to write tracing headers to the logging MDC. Pick it up via the default way with "%X{name}", e.g. "%X{x-b3-traceid}".
 */
public class TracingInstrumentationHandler extends PipedHttpHandler {
    public TracingInstrumentationHandler(PipedHttpHandler next) {
        super(next);
    }

    @Override
    public void handleRequest(HttpServerExchange exchange, RequestContext context) throws Exception {
        for (String traceIdHeader: Bootstrapper.getConfiguration().getTraceHeaders()) {
            Optional.ofNullable(exchange.getRequestHeaders().get(traceIdHeader))
                .flatMap(x -> Optional.ofNullable(x.peekFirst()))
                .ifPresent(value -> {
                    MDC.put(traceIdHeader, value);
                    exchange.getResponseHeaders().put(HttpString.tryFromString(traceIdHeader), value);
                });
        }

        if (!exchange.isResponseComplete() && getNext() != null) {
            next(exchange, context);
        }
    }
}

@@ -37,195 +37,58 @@
 ort 
-com.google.common.annotations.VisibleForTesting;%0A%0Aimport org.restheart.Bootstrapper;%0Aimport org.slf4j.Logger;%0Aimport org.slf4j.LoggerFactory;%0Aimport org.slf4j.MDC;%0A%0Aimport java.util.List;
+org.restheart.Bootstrapper;%0Aimport org.slf4j.MDC;%0A
 %0Aimp
@@ -114,33 +114,8 @@
 nal;
-%0Aimport java.util.Random;
 %0A%0Aim
@@ -162,84 +162,8 @@
 ge;%0A
-import io.undertow.util.HeaderValues;%0Aimport io.undertow.util.HexConverter;%0A
 impo
@@ -273,16 +273,19 @@
  via the
+%0A *
  default
@@ -648,16 +648,17 @@
 IdHeader
+ 
 : Bootst
@@ -801,16 +801,20 @@
         
+    
 .flatMap
@@ -871,16 +871,20 @@
         
+    
 .ifPrese
@@ -917,16 +917,20 @@
         
+  
+  
 MDC.put(
@@ -944,32 +944,36 @@
 Header, value);%0A
+    
                 
@@ -1051,32 +1051,36 @@
 eader), value);%0A
+    
                 

Remove some unused imports