remember redirect_url when solving the 2FA challenge

4 files changed, 38 insertions, 8 deletions

diff --git a/core/Controller/LoginController.php b/core/Controller/LoginController.php
index 39e1019abe7..c64f58ae2cc 100644
--- a/core/Controller/LoginController.php
+++ b/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'));
 		}
 
diff --git a/core/Controller/TwoFactorChallengeController.php b/core/Controller/TwoFactorChallengeController.php
index 16e1ee9c0c7..499898de3bc 100644
--- a/core/Controller/TwoFactorChallengeController.php
+++ b/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,
+		]));
 	}
 
 }
diff --git a/core/Middleware/TwoFactorMiddleware.php b/core/Middleware/TwoFactorMiddleware.php
index 495c4889c20..aa82897ad46 100644
--- a/core/Middleware/TwoFactorMiddleware.php
+++ b/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'));
diff --git a/core/templates/twofactorselectchallenge.php b/core/templates/twofactorselectchallenge.php
index 6db8c69d7ac..14d599aab3e 100644
--- a/core/templates/twofactorselectchallenge.php
+++ b/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>