Add SecurityContext to delegating TaskScheduler

Wrap DelegatingSecurityContextTaskScheduler's Runnable tasks in
DelegatingSecurityContextRunnables, allowing to specify a
SecurityContext to use for tasks execution.

- Renamed private variable taskScheduler to delegate
- Removed unused local variable in unit test
- Add SecurityContext tests for delegating TaskScheduler

Closes gh-9514

core/src/test/java/org/springframework/security/scheduling/DelegatingSecurityContextTaskSchedulerTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2018 the original author or authors.
+ * Copyright 2002-2021 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,9 +16,9 @@
 
 package org.springframework.security.scheduling;
 
-import java.time.Duration;
 import java.time.Instant;
 import java.util.Date;
+import java.util.concurrent.ScheduledFuture;
 
 import org.junit.After;
 import org.junit.Before;
@@ -28,11 +28,16 @@ import org.mockito.MockitoAnnotations;
 
 import org.springframework.scheduling.TaskScheduler;
 import org.springframework.scheduling.Trigger;
+import org.springframework.scheduling.concurrent.ConcurrentTaskScheduler;
+import org.springframework.security.core.context.SecurityContext;
+import org.springframework.security.core.context.SecurityContextHolder;
 
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.ArgumentMatchers.isA;
+import static org.mockito.BDDMockito.willAnswer;
 import static org.mockito.Mockito.verify;
 
 /**
@@ -47,22 +52,30 @@ public class DelegatingSecurityContextTaskSchedulerTests {
 	@Mock
 	private TaskScheduler scheduler;
 
+	@Mock
+	private SecurityContext securityContext;
+
 	@Mock
 	private Runnable runnable;
 
 	@Mock
 	private Trigger trigger;
 
+	private SecurityContext originalSecurityContext;
+
 	private DelegatingSecurityContextTaskScheduler delegatingSecurityContextTaskScheduler;
 
 	@Before
 	public void setup() {
 		MockitoAnnotations.initMocks(this);
-		this.delegatingSecurityContextTaskScheduler = new DelegatingSecurityContextTaskScheduler(this.scheduler);
+		this.originalSecurityContext = SecurityContextHolder.createEmptyContext();
+		this.delegatingSecurityContextTaskScheduler = new DelegatingSecurityContextTaskScheduler(this.scheduler,
+				this.securityContext);
 	}
 
 	@After
 	public void cleanup() {
+		SecurityContextHolder.clearContext();
 		this.delegatingSecurityContextTaskScheduler = null;
 	}
 
@@ -71,6 +84,36 @@ public class DelegatingSecurityContextTaskSchedulerTests {
 		assertThatIllegalArgumentException().isThrownBy(() -> new DelegatingSecurityContextTaskScheduler(null));
 	}
 
+	@Test
+	public void testSchedulerCurrentSecurityContext() throws Exception {
+		willAnswer((invocation) -> {
+			assertThat(SecurityContextHolder.getContext()).isEqualTo(this.originalSecurityContext);
+			return null;
+		}).given(this.runnable).run();
+		TaskScheduler delegateTaskScheduler = new ConcurrentTaskScheduler();
+		this.delegatingSecurityContextTaskScheduler = new DelegatingSecurityContextTaskScheduler(delegateTaskScheduler);
+		assertWrapped(this.runnable);
+	}
+
+	@Test
+	public void testSchedulerExplicitSecurityContext() throws Exception {
+		willAnswer((invocation) -> {
+			assertThat(SecurityContextHolder.getContext()).isEqualTo(this.securityContext);
+			return null;
+		}).given(this.runnable).run();
+		TaskScheduler delegateTaskScheduler = new ConcurrentTaskScheduler();
+		this.delegatingSecurityContextTaskScheduler = new DelegatingSecurityContextTaskScheduler(delegateTaskScheduler,
+				this.securityContext);
+		assertWrapped(this.runnable);
+	}
+
+	private void assertWrapped(Runnable runnable) throws Exception {
+		ScheduledFuture<?> schedule = this.delegatingSecurityContextTaskScheduler.schedule(runnable, new Date());
+		schedule.get();
+		verify(this.runnable).run();
+		assertThat(SecurityContextHolder.getContext()).isEqualTo(this.originalSecurityContext);
+	}
+
 	@Test
 	public void testSchedulerWithRunnableAndTrigger() {
 		this.delegatingSecurityContextTaskScheduler.schedule(this.runnable, this.trigger);
@@ -87,7 +130,6 @@ public class DelegatingSecurityContextTaskSchedulerTests {
 	@Test
 	public void testScheduleAtFixedRateWithRunnableAndDate() {
 		Date date = new Date(1544751374L);
-		Duration duration = Duration.ofSeconds(4L);
 		this.delegatingSecurityContextTaskScheduler.scheduleAtFixedRate(this.runnable, date, 1000L);
 		verify(this.scheduler).scheduleAtFixedRate(isA(Runnable.class), isA(Date.class), eq(1000L));
 	}