#1219 Make 2FA tokens one-use only (#1848)

- Make TotpAuthenticator keep track of the used codes and reject any known ones

src/main/java/fr/xephi/authme/security/totp/GenerateTotpService.java

@@ -59,7 +59,7 @@ public class GenerateTotpService implements HasCleanup {
      */
     public boolean isTotpCodeCorrectForGeneratedTotpKey(Player player, String totpCode) {
         TotpGenerationResult totpDetails = totpKeys.get(player.getName().toLowerCase());
-        return totpDetails != null && totpAuthenticator.checkCode(totpDetails.getTotpKey(), totpCode);
+        return totpDetails != null && totpAuthenticator.checkCode(player.getName(), totpDetails.getTotpKey(), totpCode);
     }
 
     @Override

src/main/java/fr/xephi/authme/security/totp/TotpAuthenticator.java

@@ -1,23 +1,31 @@
 package fr.xephi.authme.security.totp;
 
+import com.google.common.collect.HashBasedTable;
+import com.google.common.collect.Table;
 import com.google.common.primitives.Ints;
 import com.warrenstrange.googleauth.GoogleAuthenticator;
 import com.warrenstrange.googleauth.GoogleAuthenticatorKey;
 import com.warrenstrange.googleauth.GoogleAuthenticatorQRGenerator;
 import com.warrenstrange.googleauth.IGoogleAuthenticator;
 import fr.xephi.authme.data.auth.PlayerAuth;
+import fr.xephi.authme.initialization.HasCleanup;
 import fr.xephi.authme.service.BukkitService;
 import org.bukkit.entity.Player;
 
 import javax.inject.Inject;
 
+import static fr.xephi.authme.util.Utils.MILLIS_PER_MINUTE;
+
 /**
  * Provides TOTP functions (wrapping a third-party TOTP implementation).
  */
-public class TotpAuthenticator {
+public class TotpAuthenticator implements HasCleanup {
+
+    private static final int CODE_RETENTION_MINUTES = 5;
 
     private final IGoogleAuthenticator authenticator;
     private final BukkitService bukkitService;
+    private final Table<String, Integer, Long> usedCodes = HashBasedTable.create();
 
     @Inject
     TotpAuthenticator(BukkitService bukkitService) {
@@ -33,19 +41,26 @@ public class TotpAuthenticator {
     }
 
     public boolean checkCode(PlayerAuth auth, String totpCode) {
-        return checkCode(auth.getTotpKey(), totpCode);
+        return checkCode(auth.getNickname(), auth.getTotpKey(), totpCode);
     }
 
     /**
      * Returns whether the given input code matches for the provided TOTP key.
      *
+     * @param playerName the player name
      * @param totpKey the key to check with
      * @param inputCode the input code to verify
      * @return true if code is valid, false otherwise
      */
-    public boolean checkCode(String totpKey, String inputCode) {
+    public boolean checkCode(String playerName, String totpKey, String inputCode) {
+        String nameLower = playerName.toLowerCase();
         Integer totpCode = Ints.tryParse(inputCode);
-        return totpCode != null && authenticator.authorize(totpKey, totpCode);
+        if (totpCode != null && !usedCodes.contains(nameLower, totpCode)
+            && authenticator.authorize(totpKey, totpCode)) {
+            usedCodes.put(nameLower, totpCode, System.currentTimeMillis());
+            return true;
+        }
+        return false;
     }
 
     public TotpGenerationResult generateTotpKey(Player player) {
@@ -55,6 +70,12 @@ public class TotpAuthenticator {
         return new TotpGenerationResult(credentials.getKey(), qrCodeUrl);
     }
 
+    @Override
+    public void performCleanup() {
+        long threshold = System.currentTimeMillis() - CODE_RETENTION_MINUTES * MILLIS_PER_MINUTE;
+        usedCodes.values().removeIf(value -> value < threshold);
+    }
+
     public static final class TotpGenerationResult {
         private final String totpKey;
         private final String authenticatorQrCodeUrl;

src/test/java/fr/xephi/authme/security/totp/GenerateTotpServiceTest.java

@@ -70,7 +70,7 @@ public class GenerateTotpServiceTest {
         given(totpAuthenticator.generateTotpKey(player)).willReturn(givenGenerationResult);
         generateTotpService.generateTotpKey(player);
         String validCode = "928374";
-        given(totpAuthenticator.checkCode(generatedKey, validCode)).willReturn(true);
+        given(totpAuthenticator.checkCode("Aria", generatedKey, validCode)).willReturn(true);
 
         // when
         boolean invalidCodeResult = generateTotpService.isTotpCodeCorrectForGeneratedTotpKey(player, "000000");
@@ -81,8 +81,8 @@ public class GenerateTotpServiceTest {
         assertThat(invalidCodeResult, equalTo(false));
         assertThat(validCodeResult, equalTo(true));
         assertThat(unknownPlayerResult, equalTo(false));
-        verify(totpAuthenticator).checkCode(generatedKey, "000000");
+        verify(totpAuthenticator).checkCode("Aria", generatedKey, "000000");
-        verify(totpAuthenticator).checkCode(generatedKey, validCode);
+        verify(totpAuthenticator).checkCode("Aria", generatedKey, validCode);
     }
 
     @Test

src/test/java/fr/xephi/authme/security/totp/TotpAuthenticatorTest.java

@@ -1,9 +1,12 @@
 package fr.xephi.authme.security.totp;
 
+import com.google.common.collect.Table;
 import com.warrenstrange.googleauth.IGoogleAuthenticator;
+import fr.xephi.authme.ReflectionTestUtils;
 import fr.xephi.authme.data.auth.PlayerAuth;
 import fr.xephi.authme.security.totp.TotpAuthenticator.TotpGenerationResult;
 import fr.xephi.authme.service.BukkitService;
+import fr.xephi.authme.util.Utils;
 import org.bukkit.entity.Player;
 import org.junit.Before;
 import org.junit.Test;
@@ -63,24 +66,26 @@ public class TotpAuthenticatorTest {
     }
 
     @Test
-    public void shouldCheckCode() {
+    public void shouldCheckCodeAndDeclareItValidOnlyOnce() {
         // given
         String secret = "the_secret";
         int code = 21398;
         given(googleAuthenticator.authorize(secret, code)).willReturn(true);
 
         // when
-        boolean result = totpAuthenticator.checkCode(secret, Integer.toString(code));
+        boolean result1 = totpAuthenticator.checkCode("pl", secret, Integer.toString(code));
+        boolean result2 = totpAuthenticator.checkCode("pl", secret, Integer.toString(code));
 
         // then
-        assertThat(result, equalTo(true));
+        assertThat(result1, equalTo(true));
+        assertThat(result2, equalTo(false));
         verify(googleAuthenticator).authorize(secret, code);
     }
 
     @Test
     public void shouldHandleInvalidNumberInput() {
         // given / when
-        boolean result = totpAuthenticator.checkCode("Some_Secret", "123ZZ");
+        boolean result = totpAuthenticator.checkCode("foo", "Some_Secret", "123ZZ");
 
         // then
         assertThat(result, equalTo(false));
@@ -96,7 +101,7 @@ public class TotpAuthenticatorTest {
             .totpKey(totpKey)
             .build();
         String inputCode = "408435";
-        given(totpAuthenticator.checkCode(totpKey, inputCode)).willReturn(true);
+        given(totpAuthenticator.checkCode("Maya", totpKey, inputCode)).willReturn(true);
 
         // when
         boolean result = totpAuthenticator.checkCode(auth, inputCode);
@@ -106,6 +111,23 @@ public class TotpAuthenticatorTest {
         verify(googleAuthenticator).authorize(totpKey, 408435);
     }
 
+    @Test
+    public void shouldRemoveOldEntries() {
+        // given
+        Table<String, Integer, Long> usedCodes = ReflectionTestUtils.getFieldValue(
+            TotpAuthenticator.class, totpAuthenticator, "usedCodes");
+        usedCodes.put("bobby", 414213, System.currentTimeMillis());
+        usedCodes.put("charlie", 732050, System.currentTimeMillis() - 6 * Utils.MILLIS_PER_MINUTE);
+        usedCodes.put("bobby", 236067, System.currentTimeMillis() - 9 * Utils.MILLIS_PER_MINUTE);
+
+        // when
+        totpAuthenticator.performCleanup();
+
+        // then
+        assertThat(usedCodes.size(), equalTo(1));
+        assertThat(usedCodes.contains("bobby", 414213), equalTo(true));
+    }
+
     private final class TotpAuthenticatorTestImpl extends TotpAuthenticator {
 
         TotpAuthenticatorTestImpl(BukkitService bukkitService) {