Improve 2FA

* Store the auth state in the session so we don't have to query it every
time.
* Added some tests

Signed-off-by: Roeland Jago Douma <roeland@famdouma.nl>

lib/private/Authentication/TwoFactorAuth/Manager.php

@@ -28,8 +28,11 @@ use Exception;
 use OC;
 use OC\App\AppManager;
 use OC_App;
+use OC\Authentication\Exceptions\InvalidTokenException;
+use OC\Authentication\Token\IProvider as TokenProvider;
 use OCP\Activity\IManager;
 use OCP\AppFramework\QueryException;
+use OCP\AppFramework\Utility\ITimeFactory;
 use OCP\Authentication\TwoFactorAuth\IProvider;
 use OCP\IConfig;
 use OCP\ILogger;
@@ -39,6 +42,7 @@ use OCP\IUser;
 class Manager {
 
 	const SESSION_UID_KEY = 'two_factor_auth_uid';
+	const SESSION_UID_DONE = 'two_factor_auth_passed';
 	const BACKUP_CODES_APP_ID = 'twofactor_backupcodes';
 	const BACKUP_CODES_PROVIDER_ID = 'backup_codes';
 	const REMEMBER_LOGIN = 'two_factor_remember_login';
@@ -58,20 +62,35 @@ class Manager {
 	/** @var ILogger */
 	private $logger;
 
+	/** @var TokenProvider */
+	private $tokenProvider;
+
+	/** @var ITimeFactory */
+	private $timeFactory;
+
 	/**
 	 * @param AppManager $appManager
 	 * @param ISession $session
 	 * @param IConfig $config
 	 * @param IManager $activityManager
 	 * @param ILogger $logger
+	 * @param TokenProvider $tokenProvider
+	 * @param ITimeFactory $timeFactory
 	 */
-	public function __construct(AppManager $appManager, ISession $session, IConfig $config, IManager $activityManager,
+	public function __construct(AppManager $appManager,
-		ILogger $logger) {
+								ISession $session,
+								IConfig $config,
+								IManager $activityManager,
+								ILogger $logger,
+								TokenProvider $tokenProvider,
+								ITimeFactory $timeFactory) {
 		$this->appManager = $appManager;
 		$this->session = $session;
 		$this->config = $config;
 		$this->activityManager = $activityManager;
 		$this->logger = $logger;
+		$this->tokenProvider = $tokenProvider;
+		$this->timeFactory = $timeFactory;
 	}
 
 	/**
@@ -199,6 +218,13 @@ class Manager {
 			}
 			$this->session->remove(self::SESSION_UID_KEY);
 			$this->session->remove(self::REMEMBER_LOGIN);
+			$this->session->set(self::SESSION_UID_DONE, $user->getUID());
+
+			// Clear token from db
+			$sessionId = $this->session->getId();
+			$token = $this->tokenProvider->getToken($sessionId);
+			$tokenId = $token->getId();
+			$this->config->deleteUserValue($user->getUID(), 'login_token_2fa', $tokenId);
 
 			$this->publishEvent($user, 'twofactor_success', [
 				'provider' => $provider->getDisplayName(),
@@ -239,16 +265,50 @@ class Manager {
 	 * @return boolean
 	 */
 	public function needsSecondFactor(IUser $user = null) {
-		if (is_null($user) || !$this->session->exists(self::SESSION_UID_KEY)) {
+		if ($user === null) {
 			return false;
 		}
 
+		// First check if the session tells us we should do 2FA (99% case)
+		if (!$this->session->exists(self::SESSION_UID_KEY)) {
+
+			// Check if the session tells us it is 2FA authenticated already
+			if ($this->session->exists(self::SESSION_UID_DONE) &&
+				$this->session->get(self::SESSION_UID_DONE) === $user->getUID()) {
+				return false;
+			}
+
+			/*
+			 * If the session is expired check if we are not logged in by a token
+			 * that still needs 2FA auth
+			 */
+			try {
+				$sessionId = $this->session->getId();
+				$token = $this->tokenProvider->getToken($sessionId);
+				$tokenId = $token->getId();
+				$tokensNeeding2FA = $this->config->getUserKeys($user->getUID(), 'login_token_2fa');
+
+				if (!in_array($tokenId, $tokensNeeding2FA, true)) {
+					$this->session->set(self::SESSION_UID_DONE, $user->getUID());
+					return false;
+				}
+			} catch (InvalidTokenException $e) {
+				return true;
+			}
+		}
+
+
 		if (!$this->isTwoFactorAuthenticated($user)) {
 			// There is no second factor any more -> let the user pass
 			//   This prevents infinite redirect loops when a user is about
 			//   to solve the 2FA challenge, and the provider app is
 			//   disabled the same time
 			$this->session->remove(self::SESSION_UID_KEY);
+
+			$keys = $this->config->getUserKeys($user->getUID(), 'login_token_2fa');
+			foreach ($keys as $key) {
+				$this->config->deleteUserValue($user->getUID(), 'login_token_2fa', $key);
+			}
 			return false;
 		}
 
@@ -264,6 +324,10 @@ class Manager {
 	public function prepareTwoFactorLogin(IUser $user, $rememberMe) {
 		$this->session->set(self::SESSION_UID_KEY, $user->getUID());
 		$this->session->set(self::REMEMBER_LOGIN, $rememberMe);
+
+		$id = $this->session->getId();
+		$token = $this->tokenProvider->getToken($id);
+		$this->config->setUserValue($user->getUID(), 'login_token_2fa', $token->getId(), $this->timeFactory->getTime());
 	}
 
 }

lib/private/Server.php

@@ -102,6 +102,7 @@ use OC\Template\SCSSCacher;
 use OCA\Theming\ThemingDefaults;
 
 use OCP\App\IAppManager;
+use OCP\AppFramework\Utility\ITimeFactory;
 use OCP\Defaults;
 use OCA\Theming\Util;
 use OCP\Federation\ICloudIdManager;
@@ -379,7 +380,15 @@ class Server extends ServerContainer implements IServerContainer {
 		$this->registerAlias('UserSession', \OCP\IUserSession::class);
 
 		$this->registerService(\OC\Authentication\TwoFactorAuth\Manager::class, function (Server $c) {
-			return new \OC\Authentication\TwoFactorAuth\Manager($c->getAppManager(), $c->getSession(), $c->getConfig(), $c->getActivityManager(), $c->getLogger());
+			return new \OC\Authentication\TwoFactorAuth\Manager(
+				$c->getAppManager(),
+				$c->getSession(),
+				$c->getConfig(),
+				$c->getActivityManager(),
+				$c->getLogger(),
+				$c->query(\OC\Authentication\Token\IProvider::class),
+				$c->query(ITimeFactory::class)
+			);
 		});
 
 		$this->registerAlias(\OCP\INavigationManager::class, \OC\NavigationManager::class);

tests/lib/Authentication/TwoFactorAuth/ManagerTest.php

@@ -25,9 +25,11 @@ namespace Test\Authentication\TwoFactorAuth;
 use Exception;
 use OC;
 use OC\App\AppManager;
+use OC\Authentication\Token\IProvider as TokenProvider;
 use OC\Authentication\TwoFactorAuth\Manager;
 use OCP\Activity\IEvent;
 use OCP\Activity\IManager;
+use OCP\AppFramework\Utility\ITimeFactory;
 use OCP\Authentication\TwoFactorAuth\IProvider;
 use OCP\IConfig;
 use OCP\ILogger;
@@ -37,33 +39,39 @@ use Test\TestCase;
 
 class ManagerTest extends TestCase {
 
-	/** @var IUser|PHPUnit_Framework_MockObject_MockObject */
+	/** @var IUser|\PHPUnit_Framework_MockObject_MockObject */
 	private $user;
 
-	/** @var AppManager|PHPUnit_Framework_MockObject_MockObject */
+	/** @var AppManager|\PHPUnit_Framework_MockObject_MockObject */
 	private $appManager;
 
-	/** @var ISession|PHPUnit_Framework_MockObject_MockObject */
+	/** @var ISession|\PHPUnit_Framework_MockObject_MockObject */
 	private $session;
 
 	/** @var Manager */
 	private $manager;
 
-	/** @var IConfig|PHPUnit_Framework_MockObject_MockObject */
+	/** @var IConfig|\PHPUnit_Framework_MockObject_MockObject */
 	private $config;
 
-	/** @var IManager|PHPUnit_Framework_MockObject_MockObject */
+	/** @var IManager|\PHPUnit_Framework_MockObject_MockObject */
 	private $activityManager;
 
-	/** @var ILogger|PHPUnit_Framework_MockObject_MockObject */
+	/** @var ILogger|\PHPUnit_Framework_MockObject_MockObject */
 	private $logger;
 
-	/** @var IProvider|PHPUnit_Framework_MockObject_MockObject */
+	/** @var IProvider|\PHPUnit_Framework_MockObject_MockObject */
 	private $fakeProvider;
 
-	/** @var IProvider|PHPUnit_Framework_MockObject_MockObject */
+	/** @var IProvider|\PHPUnit_Framework_MockObject_MockObject */
 	private $backupProvider;
 
+	/** @var TokenProvider|\PHPUnit_Framework_MockObject_MockObject */
+	private $tokenProvider;
+
+	/** @var ITimeFactory|\PHPUnit_Framework_MockObject_MockObject */
+	private $timeFactory;
+
 	protected function setUp() {
 		parent::setUp();
 
@@ -73,9 +81,19 @@ class ManagerTest extends TestCase {
 		$this->config = $this->createMock(IConfig::class);
 		$this->activityManager = $this->createMock(IManager::class);
 		$this->logger = $this->createMock(ILogger::class);
+		$this->tokenProvider = $this->createMock(TokenProvider::class);
+		$this->timeFactory = $this->createMock(ITimeFactory::class);
 
-		$this->manager = $this->getMockBuilder('\OC\Authentication\TwoFactorAuth\Manager')
+		$this->manager = $this->getMockBuilder(Manager::class)
-			->setConstructorArgs([$this->appManager, $this->session, $this->config, $this->activityManager, $this->logger])
+			->setConstructorArgs([
+				$this->appManager,
+				$this->session,
+				$this->config,
+				$this->activityManager,
+				$this->logger,
+				$this->tokenProvider,
+				$this->timeFactory
+			])
 			->setMethods(['loadTwoFactorApp']) // Do not actually load the apps
 			->getMock();
 
@@ -242,6 +260,7 @@ class ManagerTest extends TestCase {
 			->method('verifyChallenge')
 			->with($this->user, $challenge)
 			->will($this->returnValue(true));
+
 		$this->session->expects($this->once())
 			->method('get')
 			->with('two_factor_remember_login')
@@ -252,12 +271,20 @@ class ManagerTest extends TestCase {
 		$this->session->expects($this->at(2))
 			->method('remove')
 			->with('two_factor_remember_login');
+		$this->session->expects($this->at(3))
+			->method('set')
+			->with(Manager::SESSION_UID_DONE, 'jos');
+		$this->session->method('getId')
+			->willReturn('mysessionid');
+
 		$this->activityManager->expects($this->once())
 			->method('generateEvent')
 			->willReturn($event);
+
 		$this->user->expects($this->any())
 			->method('getUID')
 			->willReturn('jos');
+
 		$event->expects($this->once())
 			->method('setApp')
 			->with($this->equalTo('core'))
@@ -284,6 +311,17 @@ class ManagerTest extends TestCase {
 			]))
 			->willReturnSelf();
 
+		$token = $this->createMock(OC\Authentication\Token\IToken::class);
+		$this->tokenProvider->method('getToken')
+			->with('mysessionid')
+			->willReturn($token);
+		$token->method('getId')
+			->willReturn(42);
+
+		$this->config->expects($this->once())
+			->method('deleteUserValue')
+			->with('jos', 'login_token_2fa', 42);
+
 		$this->assertTrue($this->manager->verifyChallenge('email', $this->user, $challenge));
 	}
 
@@ -348,12 +386,50 @@ class ManagerTest extends TestCase {
 
 	public function testNeedsSecondFactor() {
 		$user = $this->createMock(IUser::class);
-		$this->session->expects($this->once())
+		$this->session->expects($this->at(0))
 			->method('exists')
 			->with('two_factor_auth_uid')
 			->will($this->returnValue(false));
+		$this->session->expects($this->at(1))
+			->method('exists')
+			->with(Manager::SESSION_UID_DONE)
+			->willReturn(false);
 
-		$this->assertFalse($this->manager->needsSecondFactor($user));
+		$this->session->method('getId')
+			->willReturn('mysessionid');
+		$token = $this->createMock(OC\Authentication\Token\IToken::class);
+		$this->tokenProvider->method('getToken')
+			->with('mysessionid')
+			->willReturn($token);
+		$token->method('getId')
+			->willReturn(42);
+
+		$user->method('getUID')
+			->willReturn('user');
+		$this->config->method('getUserKeys')
+			->with('user', 'login_token_2fa')
+			->willReturn([
+				42
+			]);
+
+		$manager = $this->getMockBuilder(Manager::class)
+			->setConstructorArgs([
+				$this->appManager,
+				$this->session,
+				$this->config,
+				$this->activityManager,
+				$this->logger,
+				$this->tokenProvider,
+				$this->timeFactory
+			])
+			->setMethods(['loadTwoFactorApp','isTwoFactorAuthenticated']) // Do not actually load the apps
+			->getMock();
+
+		$manager->method('isTwoFactorAuthenticated')
+			->with($user)
+			->willReturn(true);
+
+		$this->assertTrue($manager->needsSecondFactor($user));
 	}
 
 	public function testNeedsSecondFactorUserIsNull() {
@@ -380,8 +456,7 @@ class ManagerTest extends TestCase {
 	}
 
 	public function testPrepareTwoFactorLogin() {
-		$this->user->expects($this->once())
+		$this->user->method('getUID')
-			->method('getUID')
 			->will($this->returnValue('ferdinand'));
 
 		$this->session->expects($this->at(0))
@@ -391,12 +466,27 @@ class ManagerTest extends TestCase {
 			->method('set')
 			->with('two_factor_remember_login', true);
 
+		$this->session->method('getId')
+			->willReturn('mysessionid');
+		$token = $this->createMock(OC\Authentication\Token\IToken::class);
+		$this->tokenProvider->method('getToken')
+			->with('mysessionid')
+			->willReturn($token);
+		$token->method('getId')
+			->willReturn(42);
+
+		$this->timeFactory->method('getTime')
+			->willReturn(1337);
+
+		$this->config->method('setUserValue')
+			->with('ferdinand', 'login_token_2fa', 42, 1337);
+
+
 		$this->manager->prepareTwoFactorLogin($this->user, true);
 	}
 
 	public function testPrepareTwoFactorLoginDontRemember() {
-		$this->user->expects($this->once())
+		$this->user->method('getUID')
-			->method('getUID')
 			->will($this->returnValue('ferdinand'));
 
 		$this->session->expects($this->at(0))
@@ -406,7 +496,89 @@ class ManagerTest extends TestCase {
 			->method('set')
 			->with('two_factor_remember_login', false);
 
+		$this->session->method('getId')
+			->willReturn('mysessionid');
+		$token = $this->createMock(OC\Authentication\Token\IToken::class);
+		$this->tokenProvider->method('getToken')
+			->with('mysessionid')
+			->willReturn($token);
+		$token->method('getId')
+			->willReturn(42);
+
+		$this->timeFactory->method('getTime')
+			->willReturn(1337);
+
+		$this->config->method('setUserValue')
+			->with('ferdinand', 'login_token_2fa', 42, 1337);
+
 		$this->manager->prepareTwoFactorLogin($this->user, false);
 	}
 
+	public function testNeedsSecondFactorSessionAuth() {
+		$user = $this->createMock(IUser::class);
+		$user->method('getUID')
+			->willReturn('user');
+
+		$this->session->method('exists')
+			->will($this->returnCallback(function($var) {
+				if ($var === Manager::SESSION_UID_KEY) {
+					return false;
+				}
+				return true;
+			}));
+		$this->session->expects($this->once())
+			->method('get')
+			->with(Manager::SESSION_UID_DONE)
+			->willReturn('user');
+
+		$this->assertFalse($this->manager->needsSecondFactor($user));
+	}
+
+	public function testNeedsSecondFactorSessionAuthFailDBPass() {
+		$user = $this->createMock(IUser::class);
+		$user->method('getUID')
+			->willReturn('user');
+
+		$this->session->method('exists')
+			->willReturn(false);
+		$this->session->method('getId')
+			->willReturn('mysessionid');
+
+		$token = $this->createMock(OC\Authentication\Token\IToken::class);
+		$token->method('getId')
+			->willReturn(40);
+
+		$this->tokenProvider->method('getToken')
+			->with('mysessionid')
+			->willReturn($token);
+
+		$this->config->method('getUserKeys')
+			->with('user', 'login_token_2fa')
+			->willReturn([
+				42, 43, 44
+			]);
+
+		$this->session->expects($this->once())
+			->method('set')
+			->with(Manager::SESSION_UID_DONE, 'user');
+
+		$this->assertFalse($this->manager->needsSecondFactor($user));
+	}
+
+	public function testNeedsSecondFactorInvalidToken() {
+		$user = $this->createMock(IUser::class);
+		$user->method('getUID')
+			->willReturn('user');
+
+		$this->session->method('exists')
+			->willReturn(false);
+		$this->session->method('getId')
+			->willReturn('mysessionid');
+
+		$this->tokenProvider->method('getToken')
+			->with('mysessionid')
+			->willThrowException(new OC\Authentication\Exceptions\InvalidTokenException());
+
+		$this->assertTrue($this->manager->needsSecondFactor($user));
+	}
 }