remember redirect_url when solving the 2FA challenge

core/Controller/LoginController.php

@@ -197,6 +197,11 @@ class LoginController extends Controller {
 
 		if ($this->twoFactorManager->isTwoFactorAuthenticated($loginResult)) {
 			$this->twoFactorManager->prepareTwoFactorLogin($loginResult);
+			if (!is_null($redirect_url)) {
+				return new RedirectResponse($this->urlGenerator->linkToRoute('core.TwoFactorChallenge.selectChallenge', [
+					'redirect_url' => $redirect_url
+				]));
+			}
 			return new RedirectResponse($this->urlGenerator->linkToRoute('core.TwoFactorChallenge.selectChallenge'));
 		}
 

core/Controller/TwoFactorChallengeController.php

@@ -65,14 +65,16 @@ class TwoFactorChallengeController extends Controller {
 	 * @NoAdminRequired
 	 * @NoCSRFRequired
 	 *
+	 * @param string $redirect_url
 	 * @return TemplateResponse
 	 */
-	public function selectChallenge() {
+	public function selectChallenge($redirect_url) {
 		$user = $this->userSession->getUser();
 		$providers = $this->twoFactorManager->getProviders($user);
 
 		$data = [
 			'providers' => $providers,
+			'redirect_url' => $redirect_url,
 		];
 		return new TemplateResponse($this->appName, 'twofactorselectchallenge', $data, 'guest');
 	}
@@ -83,9 +85,10 @@ class TwoFactorChallengeController extends Controller {
 	 * @UseSession
 	 *
 	 * @param string $challengeProviderId
+	 * @param string $redirect_url
 	 * @return TemplateResponse
 	 */
-	public function showChallenge($challengeProviderId) {
+	public function showChallenge($challengeProviderId, $redirect_url) {
 		$user = $this->userSession->getUser();
 		$provider = $this->twoFactorManager->getProvider($user, $challengeProviderId);
 		if (is_null($provider)) {
@@ -98,10 +101,12 @@ class TwoFactorChallengeController extends Controller {
 		} else {
 			$error = false;
 		}
+		$tmpl = $provider->getTemplate($user);
+		$tmpl->assign('redirect_url', $redirect_url);
 		$data = [
 			'error' => $error,
 			'provider' => $provider,
-			'template' => $provider->getTemplate($user)->fetchPage(),
+			'template' => $tmpl->fetchPage(),
 		];
 		return new TemplateResponse($this->appName, 'twofactorshowchallenge', $data, 'guest');
 	}
@@ -113,9 +118,10 @@ class TwoFactorChallengeController extends Controller {
 	 *
 	 * @param string $challengeProviderId
 	 * @param string $challenge
+	 * @param string $redirect_url
 	 * @return RedirectResponse
 	 */
-	public function solveChallenge($challengeProviderId, $challenge) {
+	public function solveChallenge($challengeProviderId, $challenge, $redirect_url = null) {
 		$user = $this->userSession->getUser();
 		$provider = $this->twoFactorManager->getProvider($user, $challengeProviderId);
 		if (is_null($provider)) {
@@ -123,11 +129,17 @@ class TwoFactorChallengeController extends Controller {
 		}
 
 		if ($this->twoFactorManager->verifyChallenge($challengeProviderId, $user, $challenge)) {
+			if (!is_null($redirect_url)) {
+				return new RedirectResponse($this->urlGenerator->getAbsoluteURL(urldecode($redirect_url)));
+			}
 			return new RedirectResponse($this->urlGenerator->linkToRoute('files.view.index'));
 		}
 
 		$this->session->set('two_factor_auth_error', true);
-		return new RedirectResponse($this->urlGenerator->linkToRoute('core.TwoFactorChallenge.showChallenge', ['challengeProviderId' => $provider->getId()]));
+		return new RedirectResponse($this->urlGenerator->linkToRoute('core.TwoFactorChallenge.showChallenge', [
+			'challengeProviderId' => $provider->getId(),
+			'redirect_url' => $redirect_url,
+		]));
 	}
 
 }

core/Middleware/TwoFactorMiddleware.php

@@ -1,4 +1,5 @@
 <?php
+
 /**
  * @author Christoph Wurst <christoph@owncloud.com>
  *
@@ -31,6 +32,7 @@ use OCP\AppFramework\Controller;
 use OCP\AppFramework\Http\RedirectResponse;
 use OCP\AppFramework\Middleware;
 use OCP\AppFramework\Utility\IControllerMethodReflector;
+use OCP\IRequest;
 use OCP\ISession;
 use OCP\IURLGenerator;
 
@@ -51,6 +53,9 @@ class TwoFactorMiddleware extends Middleware {
 	/** @var IControllerMethodReflector */
 	private $reflector;
 
+	/** @var IRequest */
+	private $request;
+
 	/**
 	 * @param Manager $twoFactorManager
 	 * @param Session $userSession
@@ -58,12 +63,13 @@ class TwoFactorMiddleware extends Middleware {
 	 * @param IURLGenerator $urlGenerator
 	 */
 	public function __construct(Manager $twoFactorManager, Session $userSession, ISession $session,
-		IURLGenerator $urlGenerator, IControllerMethodReflector $reflector) {
+		IURLGenerator $urlGenerator, IControllerMethodReflector $reflector, IRequest $request) {
 		$this->twoFactorManager = $twoFactorManager;
 		$this->userSession = $userSession;
 		$this->session = $session;
 		$this->urlGenerator = $urlGenerator;
 		$this->reflector = $reflector;
+		$this->request = $request;
 	}
 
 	/**
@@ -110,7 +116,9 @@ class TwoFactorMiddleware extends Middleware {
 
 	public function afterException($controller, $methodName, Exception $exception) {
 		if ($exception instanceof TwoFactorAuthRequiredException) {
-			return new RedirectResponse($this->urlGenerator->linkToRoute('core.TwoFactorChallenge.selectChallenge'));
+			return new RedirectResponse($this->urlGenerator->linkToRoute('core.TwoFactorChallenge.selectChallenge', [
+					'redirect_url' => urlencode($this->request->server['REQUEST_URI']),
+			]));
 		}
 		if ($exception instanceof UserAlreadyLoggedInException) {
 			return new RedirectResponse($this->urlGenerator->linkToRoute('files.view.index'));

core/templates/twofactorselectchallenge.php

@@ -7,7 +7,12 @@
 <?php foreach ($_['providers'] as $provider): ?>
 	<li>
 		<a class="two-factor-provider"
-		   href="<?php p(\OC::$server->getURLGenerator()->linkToRoute('core.TwoFactorChallenge.showChallenge', ['challengeProviderId' => $provider->getId()])) ?>">
+		   href="<?php p(\OC::$server->getURLGenerator()->linkToRoute('core.TwoFactorChallenge.showChallenge',
+							[
+								'challengeProviderId' => $provider->getId(),
+								'redirect_url' => $_['redirect_url'],
+							]
+						)) ?>">
 			<?php p($provider->getDescription()) ?>
 		</a>
 	</li>

lib/private/AppFramework/DependencyInjection/DIContainer.php

@@ -369,7 +369,8 @@ class DIContainer extends SimpleContainer implements IAppContainer {
 			$session = $app->getServer()->getSession();
 			$urlGenerator = $app->getServer()->getURLGenerator();
 			$reflector = $c['ControllerMethodReflector'];
-			return new TwoFactorMiddleware($twoFactorManager, $userSession, $session, $urlGenerator, $reflector);
+			$request = $app->getServer()->getRequest();
+			return new TwoFactorMiddleware($twoFactorManager, $userSession, $session, $urlGenerator, $reflector, $request);
 		});
 
 		$middleWares = &$this->middleWares;

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'));
 	}
 
 }

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));
 	}