remember redirect_url when solving the 2FA challenge

2 files changed, 22 insertions, 9 deletions

diff --git a/tests/Core/Controller/TwoFactorChallengeControllerTest.php b/tests/Core/Controller/TwoFactorChallengeControllerTest.php
index aa1c7d39cfa..2da6dcd52ac 100644
--- a/tests/Core/Controller/TwoFactorChallengeControllerTest.php
+++ b/tests/Core/Controller/TwoFactorChallengeControllerTest.php
@@ -69,9 +69,10 @@ class TwoFactorChallengeControllerTest extends TestCase {
 
 		$expected = new \OCP\AppFramework\Http\TemplateResponse('core', 'twofactorselectchallenge', [
 			'providers' => $providers,
-			], 'guest');
+			'redirect_url' => '/some/url',
+		], 'guest');
 
-		$this->assertEquals($expected, $this->controller->selectChallenge());
+		$this->assertEquals($expected, $this->controller->selectChallenge('/some/url'));
 	}
 
 	public function testShowChallenge() {
@@ -112,7 +113,7 @@ class TwoFactorChallengeControllerTest extends TestCase {
 			'template' => '<html/>',
 			], 'guest');
 
-		$this->assertEquals($expected, $this->controller->showChallenge('myprovider'));
+		$this->assertEquals($expected, $this->controller->showChallenge('myprovider', '/re/dir/ect/url'));
 	}
 
 	public function testShowInvalidChallenge() {
@@ -132,7 +133,7 @@ class TwoFactorChallengeControllerTest extends TestCase {
 
 		$expected = new \OCP\AppFramework\Http\RedirectResponse('select/challenge/url');
 
-		$this->assertEquals($expected, $this->controller->showChallenge('myprovider'));
+		$this->assertEquals($expected, $this->controller->showChallenge('myprovider', 'redirect/url'));
 	}
 
 	public function testSolveChallenge() {
@@ -207,6 +208,7 @@ class TwoFactorChallengeControllerTest extends TestCase {
 			->method('linkToRoute')
 			->with('core.TwoFactorChallenge.showChallenge', [
 				'challengeProviderId' => 'myprovider',
+				'redirect_url' => '/url',
 			])
 			->will($this->returnValue('files/index/url'));
 		$provider->expects($this->once())
@@ -214,7 +216,7 @@ class TwoFactorChallengeControllerTest extends TestCase {
 			->will($this->returnValue('myprovider'));
 
 		$expected = new \OCP\AppFramework\Http\RedirectResponse('files/index/url');
-		$this->assertEquals($expected, $this->controller->solveChallenge('myprovider', 'token'));
+		$this->assertEquals($expected, $this->controller->solveChallenge('myprovider', 'token', '/url'));
 	}
 
 }
diff --git a/tests/Core/Middleware/TwoFactorMiddlewareTest.php b/tests/Core/Middleware/TwoFactorMiddlewareTest.php
index 248793bf987..6b8f4928928 100644
--- a/tests/Core/Middleware/TwoFactorMiddlewareTest.php
+++ b/tests/Core/Middleware/TwoFactorMiddlewareTest.php
@@ -23,6 +23,7 @@
 namespace Test\Core\Middleware;
 
 use OC\Core\Middleware\TwoFactorMiddleware;
+use OC\AppFramework\Http\Request;
 use Test\TestCase;
 
 class TwoFactorMiddlewareTest extends TestCase {
@@ -32,6 +33,7 @@ class TwoFactorMiddlewareTest extends TestCase {
 	private $session;
 	private $urlGenerator;
 	private $reflector;
+	private $request;
 
 	/** @var TwoFactorMiddleware */
 	private $middleware;
@@ -48,8 +50,17 @@ class TwoFactorMiddlewareTest extends TestCase {
 		$this->session = $this->getMock('\OCP\ISession');
 		$this->urlGenerator = $this->getMock('\OCP\IURLGenerator');
 		$this->reflector = $this->getMock('\OCP\AppFramework\Utility\IControllerMethodReflector');
-
-		$this->middleware = new TwoFactorMiddleware($this->twoFactorManager, $this->userSession, $this->session, $this->urlGenerator, $this->reflector);
+		$this->request = new Request(
+			[
+				'server' => [
+					'REQUEST_URI' => 'test/url'
+				]
+			],
+			$this->getMock('\OCP\Security\ISecureRandom'),
+			$this->getMock('\OCP\IConfig')
+		);
+
+		$this->middleware = new TwoFactorMiddleware($this->twoFactorManager, $this->userSession, $this->session, $this->urlGenerator, $this->reflector, $this->request);
 	}
 
 	public function testBeforeControllerNotLoggedIn() {
@@ -162,8 +173,8 @@ class TwoFactorMiddlewareTest extends TestCase {
 		$this->urlGenerator->expects($this->once())
 			->method('linkToRoute')
 			->with('core.TwoFactorChallenge.selectChallenge')
-			->will($this->returnValue('redirect/url'));
-		$expected = new \OCP\AppFramework\Http\RedirectResponse('redirect/url');
+			->will($this->returnValue('test/url'));
+		$expected = new \OCP\AppFramework\Http\RedirectResponse('test/url');
 
 		$this->assertEquals($expected, $this->middleware->afterException(null, 'index', $ex));
 	}