From c09c5611b8c284d3af9e8e83814b24f8bcff0ac0 Mon Sep 17 00:00:00 2001
From: Jason Tedor <jason@tedor.me>
Date: Sun, 30 Aug 2015 17:49:56 -0400
Subject: [PATCH] Add override of handler registration method to wrap handler

Original commit: elastic/x-pack-elasticsearch@4dcfb637c2cf2888bb26a8a126f50449634f0045
---
 .../shield/transport/ShieldServerTransportService.java     | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/shield/src/main/java/org/elasticsearch/shield/transport/ShieldServerTransportService.java b/shield/src/main/java/org/elasticsearch/shield/transport/ShieldServerTransportService.java
index 017ea32dfa1..b1976f3de33 100644
--- a/shield/src/main/java/org/elasticsearch/shield/transport/ShieldServerTransportService.java
+++ b/shield/src/main/java/org/elasticsearch/shield/transport/ShieldServerTransportService.java
@@ -20,6 +20,7 @@ import org.elasticsearch.transport.netty.NettyTransport;
 
 import java.util.Collections;
 import java.util.Map;
+import java.util.concurrent.Callable;
 
 import static org.elasticsearch.shield.transport.netty.ShieldNettyTransport.*;
 
@@ -67,6 +68,12 @@ public class ShieldServerTransportService extends TransportService {
         super.registerRequestHandler(action, request, executor, forceExecution, wrappedHandler);
     }
 
+    @Override
+    public <Request extends TransportRequest> void registerRequestHandler(String action, Callable<Request> requestFactory, String executor, TransportRequestHandler<Request> handler) {
+        TransportRequestHandler<Request> wrappedHandler = new ProfileSecuredRequestHandler<>(action, handler, profileFilters);
+        super.registerRequestHandler(action, requestFactory, executor, wrappedHandler);
+    }
+
     protected Map<String, ServerTransportFilter> initializeProfileFilters() {
         if (!(transport instanceof ShieldNettyTransport)) {
             return Collections.<String, ServerTransportFilter>singletonMap(NettyTransport.DEFAULT_PROFILE, new ServerTransportFilter.NodeProfile(authcService, authzService, actionMapper, false));