Repository: pinpoint-apm/pinpoint
Issue: Propagate security context to web socket background tasks 
Labels: enhancement
module:web

Timer tasks started and scheduled by web socket messages do not have access to security context.
Preserve the security context of whoever starts a *PinpointWebSocketResponseAggregator* and bind it to it's life cycle so that timer tasks that are created within them may have access to it.

Issue comments:

Pull request review comments:

Patch that resolved the issue:
diff --git a/web/src/main/java/com/navercorp/pinpoint/web/task/TimerTaskDecorator.java b/web/src/main/java/com/navercorp/pinpoint/web/task/TimerTaskDecorator.java
new file mode 100644
index 000000000000..dc4180767d44
--- /dev/null
+++ b/web/src/main/java/com/navercorp/pinpoint/web/task/TimerTaskDecorator.java
@@ -0,0 +1,27 @@
+/*
+ * Copyright 2018 NAVER Corp.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.navercorp.pinpoint.web.task;
+
+import java.util.TimerTask;
+
+/**
+ * @author HyunGil Jeong
+ */
+public interface TimerTaskDecorator {
+
+    TimerTask decorate(TimerTask timerTask);
+}
diff --git a/web/src/main/java/com/navercorp/pinpoint/web/task/TimerTaskDecoratorFactory.java b/web/src/main/java/com/navercorp/pinpoint/web/task/TimerTaskDecoratorFactory.java
new file mode 100644
index 000000000000..9394dcacd138
--- /dev/null
+++ b/web/src/main/java/com/navercorp/pinpoint/web/task/TimerTaskDecoratorFactory.java
@@ -0,0 +1,25 @@
+/*
+ * Copyright 2018 NAVER Corp.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.navercorp.pinpoint.web.task;
+
+/**
+ * @author HyunGil Jeong
+ */
+public interface TimerTaskDecoratorFactory {
+
+    TimerTaskDecorator createTimerTaskDecorator();
+}
diff --git a/web/src/main/java/com/navercorp/pinpoint/web/websocket/ActiveThreadCountHandler.java b/web/src/main/java/com/navercorp/pinpoint/web/websocket/ActiveThreadCountHandler.java
index 0b75c87a3a47..2433f4144103 100644
--- a/web/src/main/java/com/navercorp/pinpoint/web/websocket/ActiveThreadCountHandler.java
+++ b/web/src/main/java/com/navercorp/pinpoint/web/websocket/ActiveThreadCountHandler.java
@@ -22,6 +22,8 @@
 import com.navercorp.pinpoint.rpc.util.MapUtils;
 import com.navercorp.pinpoint.web.security.ServerMapDataFilter;
 import com.navercorp.pinpoint.web.service.AgentService;
+import com.navercorp.pinpoint.web.task.TimerTaskDecorator;
+import com.navercorp.pinpoint.web.task.TimerTaskDecoratorFactory;
 import com.navercorp.pinpoint.web.util.SimpleOrderedThreadPool;
 import com.navercorp.pinpoint.web.websocket.message.PinpointWebSocketMessage;
 import com.navercorp.pinpoint.web.websocket.message.PinpointWebSocketMessageConverter;
@@ -68,6 +70,7 @@ public class ActiveThreadCountHandler extends TextWebSocketHandler implements Pi
 
     private final AtomicBoolean onTimerTask = new AtomicBoolean(false);
 
+
     private SimpleOrderedThreadPool webSocketFlushExecutor;
 
     private Timer flushTimer;
@@ -79,10 +82,13 @@ public class ActiveThreadCountHandler extends TextWebSocketHandler implements Pi
     private final long healthCheckDelay;
 
     private Timer reactiveTimer;
-    
+
     @Autowired(required=false)
     ServerMapDataFilter serverMapDataFilter;
 
+    @Autowired(required = false)
+    private TimerTaskDecoratorFactory timerTaskDecoratorFactory = new PinpointWebSocketTimerTaskDecoratorFactory();
+
     public ActiveThreadCountHandler(AgentService agentService) {
         this(DEFAULT_REQUEST_MAPPING, agentService);
     }
@@ -272,7 +278,8 @@ private void bindingResponseAggregator(WebSocketSession webSocketSession, WebSoc
 
         PinpointWebSocketResponseAggregator responseAggregator = aggregatorRepository.get(applicationName);
         if (responseAggregator == null) {
-            responseAggregator = new ActiveThreadCountResponseAggregator(applicationName, agentService, reactiveTimer);
+            TimerTaskDecorator timerTaskDecorator = timerTaskDecoratorFactory.createTimerTaskDecorator();
+            responseAggregator = new ActiveThreadCountResponseAggregator(applicationName, agentService, reactiveTimer, timerTaskDecorator);
             responseAggregator.start();
             aggregatorRepository.put(applicationName, responseAggregator);
         }
diff --git a/web/src/main/java/com/navercorp/pinpoint/web/websocket/ActiveThreadCountResponseAggregator.java b/web/src/main/java/com/navercorp/pinpoint/web/websocket/ActiveThreadCountResponseAggregator.java
index 766c233504f9..de1e1901bc03 100644
--- a/web/src/main/java/com/navercorp/pinpoint/web/websocket/ActiveThreadCountResponseAggregator.java
+++ b/web/src/main/java/com/navercorp/pinpoint/web/websocket/ActiveThreadCountResponseAggregator.java
@@ -19,6 +19,7 @@
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.navercorp.pinpoint.common.server.util.AgentLifeCycleState;
 import com.navercorp.pinpoint.web.service.AgentService;
+import com.navercorp.pinpoint.web.task.TimerTaskDecorator;
 import com.navercorp.pinpoint.web.vo.AgentActiveThreadCount;
 import com.navercorp.pinpoint.web.vo.AgentActiveThreadCountList;
 import com.navercorp.pinpoint.web.vo.AgentInfo;
@@ -32,6 +33,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Timer;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
@@ -55,6 +57,7 @@ public class ActiveThreadCountResponseAggregator implements PinpointWebSocketRes
     private final String applicationName;
     private final AgentService agentService;
     private final Timer timer;
+    private final TimerTaskDecorator timerTaskDecorator;
 
     private final Object workerManagingLock = new Object();
     private final List<WebSocketSession> webSocketSessions = new CopyOnWriteArrayList<>();
@@ -70,11 +73,12 @@ public class ActiveThreadCountResponseAggregator implements PinpointWebSocketRes
 
     private Map<String, AgentActiveThreadCount> activeThreadCountMap = new HashMap<>();
 
-    public ActiveThreadCountResponseAggregator(String applicationName, AgentService agentService, Timer timer) {
-        this.applicationName = applicationName;
-        this.agentService = agentService;
+    public ActiveThreadCountResponseAggregator(String applicationName, AgentService agentService, Timer timer, TimerTaskDecorator timerTaskDecorator) {
+        this.applicationName = Objects.requireNonNull(applicationName, "applicationName must not be null");
+        this.agentService = Objects.requireNonNull(agentService, "agentService must not be null");
 
-        this.timer = timer;
+        this.timer = Objects.requireNonNull(timer, "timer must not be null");
+        this.timerTaskDecorator = Objects.requireNonNull(timerTaskDecorator, "timerTaskDecorator must not be null");
 
         this.messageConverter = new PinpointWebSocketMessageConverter();
     }
@@ -82,7 +86,7 @@ public ActiveThreadCountResponseAggregator(String applicationName, AgentService
     @Override
     public void start() {
         synchronized (workerManagingLock) {
-            workerActiveManager = new WorkerActiveManager(this, agentService, timer);
+            workerActiveManager = new WorkerActiveManager(this, agentService, timer, timerTaskDecorator);
         }
     }
 
diff --git a/web/src/main/java/com/navercorp/pinpoint/web/websocket/ActiveThreadCountWorker.java b/web/src/main/java/com/navercorp/pinpoint/web/websocket/ActiveThreadCountWorker.java
index 99bcbdbc8d7f..03008185a0a1 100644
--- a/web/src/main/java/com/navercorp/pinpoint/web/websocket/ActiveThreadCountWorker.java
+++ b/web/src/main/java/com/navercorp/pinpoint/web/websocket/ActiveThreadCountWorker.java
@@ -39,6 +39,8 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.Objects;
+
 /**
  * @author Taejin Koo
  */
@@ -76,13 +78,12 @@ public ActiveThreadCountWorker(AgentService agentService, AgentInfo agentInfo, P
     }
 
     public ActiveThreadCountWorker(AgentService agentService, String applicationName, String agentId, PinpointWebSocketResponseAggregator webSocketResponseAggregator, WorkerActiveManager workerActiveManager) {
-        this.agentService = agentService;
-
-        this.applicationName = applicationName;
-        this.agentId = agentId;
+        this.agentService = Objects.requireNonNull(agentService, "agentService must not be null");
+        this.applicationName = Objects.requireNonNull(applicationName, "applicationName must not be null");
+        this.agentId = Objects.requireNonNull(agentId, "agentId must not be null");
 
-        this.responseAggregator = webSocketResponseAggregator;
-        this.workerActiveManager = workerActiveManager;
+        this.responseAggregator = Objects.requireNonNull(webSocketResponseAggregator, "responseAggregator must not be null");
+        this.workerActiveManager = Objects.requireNonNull(workerActiveManager, "workerActiveManager must not be null");
 
         AgentActiveThreadCountFactory failResponseFactory = new AgentActiveThreadCountFactory();
         failResponseFactory.setAgentId(agentId);
diff --git a/web/src/main/java/com/navercorp/pinpoint/web/websocket/PinpointWebSocketTimerTaskDecoratorFactory.java b/web/src/main/java/com/navercorp/pinpoint/web/websocket/PinpointWebSocketTimerTaskDecoratorFactory.java
new file mode 100644
index 000000000000..97a13170d317
--- /dev/null
+++ b/web/src/main/java/com/navercorp/pinpoint/web/websocket/PinpointWebSocketTimerTaskDecoratorFactory.java
@@ -0,0 +1,67 @@
+/*
+ * Copyright 2018 NAVER Corp.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.navercorp.pinpoint.web.websocket;
+
+import com.navercorp.pinpoint.web.task.TimerTaskDecorator;
+import com.navercorp.pinpoint.web.task.TimerTaskDecoratorFactory;
+import org.springframework.security.core.Authentication;
+import org.springframework.security.core.context.SecurityContext;
+import org.springframework.security.core.context.SecurityContextHolder;
+
+import java.util.TimerTask;
+
+/**
+ * @author HyunGil Jeong
+ */
+public class PinpointWebSocketTimerTaskDecoratorFactory implements TimerTaskDecoratorFactory {
+
+    @Override
+    public TimerTaskDecorator createTimerTaskDecorator() {
+        return new SecurityContextPreservingTimerTaskDecorator();
+    }
+
+    private static class SecurityContextPreservingTimerTaskDecorator implements TimerTaskDecorator {
+
+        private final SecurityContext securityContext;
+
+        private SecurityContextPreservingTimerTaskDecorator() {
+            SecurityContext securityContext = SecurityContextHolder.createEmptyContext();
+            SecurityContext currentSecurityContext = SecurityContextHolder.getContext();
+            Authentication authentication = currentSecurityContext.getAuthentication();
+            if (authentication != null) {
+                securityContext.setAuthentication(authentication);
+            }
+            this.securityContext = securityContext;
+        }
+
+        @Override
+        public TimerTask decorate(TimerTask timerTask) {
+            return new TimerTask() {
+                @Override
+                public void run() {
+                    SecurityContext previousSecurityContext = SecurityContextHolder.getContext();
+                    try {
+                        SecurityContextHolder.setContext(securityContext);
+                        timerTask.run();
+                    } finally {
+                        SecurityContextHolder.setContext(previousSecurityContext);
+                    }
+                }
+            };
+        }
+    }
+}
diff --git a/web/src/main/java/com/navercorp/pinpoint/web/websocket/WorkerActiveManager.java b/web/src/main/java/com/navercorp/pinpoint/web/websocket/WorkerActiveManager.java
index 212295d10e49..f84bf24fd321 100644
--- a/web/src/main/java/com/navercorp/pinpoint/web/websocket/WorkerActiveManager.java
+++ b/web/src/main/java/com/navercorp/pinpoint/web/websocket/WorkerActiveManager.java
@@ -18,6 +18,7 @@
 
 import com.navercorp.pinpoint.common.server.util.AgentLifeCycleState;
 import com.navercorp.pinpoint.web.service.AgentService;
+import com.navercorp.pinpoint.web.task.TimerTaskDecorator;
 import com.navercorp.pinpoint.web.vo.AgentInfo;
 import com.navercorp.pinpoint.web.vo.AgentStatus;
 import org.slf4j.Logger;
@@ -26,6 +27,7 @@
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Objects;
 import java.util.Set;
 import java.util.Timer;
 import java.util.TimerTask;
@@ -39,7 +41,7 @@
 public class WorkerActiveManager {
 
     private static final long DEFAULT_RECONNECT_DELAY = 5000;
-    private static final long DEFAULT_AGENT_CHECk_DELAY = 10000;
+    private static final long DEFAULT_AGENT_CHECK_DELAY = 10000;
 
     private final Logger logger = LoggerFactory.getLogger(this.getClass());
 
@@ -49,6 +51,7 @@ public class WorkerActiveManager {
     private final AgentService agentService;
 
     private final Timer timer;
+    private final TimerTaskDecorator timerTaskDecorator;
 
     private final AtomicBoolean isStopped = new AtomicBoolean();
 
@@ -60,13 +63,14 @@ public class WorkerActiveManager {
     private final AtomicBoolean onAgentCheckTimerTask = new AtomicBoolean(false);
     private final List<String> defaultAgentIdList = new CopyOnWriteArrayList<>();
 
-    public WorkerActiveManager(PinpointWebSocketResponseAggregator responseAggregator, AgentService agentService, Timer timer) {
-        this.responseAggregator = responseAggregator;
+    public WorkerActiveManager(PinpointWebSocketResponseAggregator responseAggregator, AgentService agentService, Timer timer, TimerTaskDecorator timerTaskDecorator) {
+        this.responseAggregator = Objects.requireNonNull(responseAggregator, "responseAggregator must not be null");
+        this.agentService = Objects.requireNonNull(agentService, "agentService must not be null");
 
-        this.applicationName = responseAggregator.getApplicationName();
-        this.agentService = agentService;
+        this.timer = Objects.requireNonNull(timer, "timer must not be null");
+        this.timerTaskDecorator = Objects.requireNonNull(timerTaskDecorator, "timerTaskDecorator must not be null");
 
-        this.timer = timer;
+        this.applicationName = this.responseAggregator.getApplicationName();
     }
 
     public void close() {
@@ -100,7 +104,8 @@ public void addReactiveWorker(String agentId) {
             boolean turnOn = onReconnectTimerTask.compareAndSet(false, true);
             logger.info("addReactiveWorker turnOn:{}", turnOn);
             if (turnOn) {
-                timer.schedule(new ReactiveTimerTask(), DEFAULT_RECONNECT_DELAY);
+                TimerTask reactiveTimerTask = timerTaskDecorator.decorate(new ReactiveTimerTask());
+                timer.schedule(reactiveTimerTask, DEFAULT_RECONNECT_DELAY);
             }
         }
     }
@@ -110,7 +115,8 @@ public void startAgentCheckJob() {
 
         boolean turnOn = onAgentCheckTimerTask.compareAndSet(false, true);
         if (turnOn) {
-            timer.schedule(new AgentCheckTimerTask(), DEFAULT_AGENT_CHECk_DELAY);
+            TimerTask agentCheckTimerTask = timerTaskDecorator.decorate(new AgentCheckTimerTask());
+            timer.schedule(agentCheckTimerTask, DEFAULT_AGENT_CHECK_DELAY);
         }
     }
 
@@ -170,7 +176,8 @@ public void run() {
                 }
             } finally {
                 if (timer != null && onAgentCheckTimerTask.get() && !isStopped.get()) {
-                    timer.schedule(new AgentCheckTimerTask(), DEFAULT_AGENT_CHECk_DELAY);
+                    TimerTask agentCheckTimerTask = timerTaskDecorator.decorate(new AgentCheckTimerTask());
+                    timer.schedule(agentCheckTimerTask, DEFAULT_AGENT_CHECK_DELAY);
                 }
             }
         }


Test patch:
diff --git a/web/src/test/java/com/navercorp/pinpoint/web/websocket/PinpointWebSocketTimerTaskDecoratorTest.java b/web/src/test/java/com/navercorp/pinpoint/web/websocket/PinpointWebSocketTimerTaskDecoratorTest.java
new file mode 100644
index 000000000000..4823408fc858
--- /dev/null
+++ b/web/src/test/java/com/navercorp/pinpoint/web/websocket/PinpointWebSocketTimerTaskDecoratorTest.java
@@ -0,0 +1,92 @@
+/*
+ * Copyright 2018 NAVER Corp.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.navercorp.pinpoint.web.websocket;
+
+import com.navercorp.pinpoint.web.task.TimerTaskDecoratorFactory;
+import org.junit.Assert;
+import org.junit.Test;
+import org.springframework.security.authentication.TestingAuthenticationToken;
+import org.springframework.security.core.Authentication;
+import org.springframework.security.core.context.SecurityContext;
+import org.springframework.security.core.context.SecurityContextHolder;
+import org.springframework.security.core.context.SecurityContextImpl;
+
+import java.util.Objects;
+import java.util.Timer;
+import java.util.TimerTask;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * @author HyunGil Jeong
+ */
+public class PinpointWebSocketTimerTaskDecoratorTest {
+
+    private static final long DELAY_MS = 1000L;
+
+    private final TimerTaskDecoratorFactory timerTaskDecoratorFactory = new PinpointWebSocketTimerTaskDecoratorFactory();
+
+    @Test
+    public void testAuthenticationPropagation() throws InterruptedException {
+        final int numThreads = 3;
+        final Authentication[] authentications = new Authentication[numThreads];
+        for (int i = 0; i < authentications.length; i++) {
+            final String principal = "principal" + i;
+            final String credential = "credential" + i;
+            authentications[i] = new TestingAuthenticationToken(principal, credential);
+        }
+        final CountDownLatch schedulerLatch = new CountDownLatch(numThreads);
+        final Timer timer = new Timer();
+
+        for (Authentication authentication : authentications) {
+            new Thread(new Runnable() {
+                @Override
+                public void run() {
+                    SecurityContext securityContext = new SecurityContextImpl();
+                    securityContext.setAuthentication(authentication);
+                    SecurityContextHolder.setContext(securityContext);
+                    TimerTask timerTask = timerTaskDecoratorFactory.createTimerTaskDecorator().decorate(new TestTimerTask(schedulerLatch, authentication));
+                    timer.schedule(timerTask, DELAY_MS);
+                }
+            }).start();
+        }
+        Assert.assertTrue("Timed out waiting for timer task completion", schedulerLatch.await(2 * DELAY_MS, TimeUnit.MILLISECONDS));
+    }
+
+    private static class TestTimerTask extends TimerTask {
+
+        private final CountDownLatch executeLatch;
+        private final Authentication expectedAuthentication;
+
+        private TestTimerTask(CountDownLatch executeLatch, Authentication expectedAuthentication) {
+            this.executeLatch = Objects.requireNonNull(executeLatch, "executeLatch must not be null");
+            this.expectedAuthentication = Objects.requireNonNull(expectedAuthentication, "expectedAuthentication must not be null");
+        }
+
+        @Override
+        public void run() {
+            try {
+                SecurityContext securityContext = SecurityContextHolder.getContext();
+                Assert.assertNotNull(securityContext);
+                Authentication actualAuthentication = securityContext.getAuthentication();
+                Assert.assertSame(expectedAuthentication, actualAuthentication);
+            } finally {
+                executeLatch.countDown();
+            }
+        }
+    }
+}