From 5efb2050c02d381749c8bec5e83cdc0ecd0c6a6d Mon Sep 17 00:00:00 2001
From: Eric Mann <eric@eamann.com>
Date: Tue, 21 Nov 2017 20:23:53 -0800
Subject: [PATCH] Clean up the secret padding to fit better with the hashing
 spec.

---
 php/functions.php              | 41 ++++++++++++++++++++++++++++++----
 test/phpunit/CoreTest.php      | 41 ++++++++++++++++++++++++++--------
 test/phpunit/ReferenceTest.php | 18 +++++++++++++++
 3 files changed, 87 insertions(+), 13 deletions(-)

diff --git a/php/functions.php b/php/functions.php
index c1e88e6..6f9bd83 100644
--- a/php/functions.php
+++ b/php/functions.php
@@ -53,6 +53,31 @@ function is_valid_authcode($key, $authcode)
     return false;
 }
 
+/**
+ * Pad a short secret with bytes from the same until it's the correct length
+ * for hashing.
+ *
+ * @param string $secret Secret key to pad
+ * @param int    $length Byte length of the desired padded secret
+ *
+ * @throws \InvalidArgumentException If the secret or length are invalid
+ *
+ * @return string
+ */
+function pad_secret($secret, $length)
+{
+    if (empty($secret)) {
+        throw new \InvalidArgumentException('Secret must be non-empty!');
+    }
+
+    $length = intval($length);
+    if ($length <= 0) {
+        throw new \InvalidArgumentException('Padding length must be non-zero');
+    }
+
+    return str_pad($secret, $length, $secret, STR_PAD_RIGHT);
+}
+
 /**
  * Calculate a valid code given the shared secret key
  *
@@ -72,10 +97,18 @@ function calc_totp($key, $step_count = false, $digits = 6, $hash = 'sha1', $time
 
     $secret = Encoding::base32DecodeUpper((string)$key);
 
-    if ($hash === 'sha256') {
-        $secret .= substr($secret, 0, 12);
-    } else if ($hash === 'sha512') {
-        $secret .= $secret . $secret . substr($secret, 0, 4);
+    switch($hash) {
+        case 'sha1':
+            $secret = pad_secret($secret, 20);
+            break;
+        case 'sha256':
+            $secret = pad_secret($secret, 32);
+            break;
+        case 'sha512':
+            $secret = pad_secret($secret, 64);
+            break;
+        default:
+            throw new \InvalidArgumentException('Invalid hash type specified!');
     }
 
     if (false === $step_count) {
diff --git a/test/phpunit/CoreTest.php b/test/phpunit/CoreTest.php
index 7b442b9..e9659ab 100644
--- a/test/phpunit/CoreTest.php
+++ b/test/phpunit/CoreTest.php
@@ -1,19 +1,23 @@
 <?php
+
 namespace EAMann\TOTP;
 
 use PHPUnit\Framework\TestCase;
 
-class CoreTest extends TestCase {
-    public function test_key_generation_uniqueness() {
+class CoreTest extends TestCase
+{
+    public function test_key_generation_uniqueness()
+    {
         $first = generate_key();
         $second = generate_key();
 
-        $this->assertNotEquals((string) $first, (string) $second);
+        $this->assertNotEquals((string)$first, (string)$second);
     }
 
-    public function test_key_validation() {
-        if ( PHP_INT_SIZE === 4 ) {
-            $this->markTestSkipped( 'calc_totp requires 64-bit PHP' );
+    public function test_key_validation()
+    {
+        if (PHP_INT_SIZE === 4) {
+            $this->markTestSkipped('calc_totp requires 64-bit PHP');
         }
 
         $key = new Key();
@@ -22,13 +26,32 @@ public function test_key_validation() {
         $this->assertTrue(is_valid_authcode($key, $otp));
     }
 
-    public function test_invalid_otp() {
-        if ( PHP_INT_SIZE === 4 ) {
-            $this->markTestSkipped( 'calc_totp requires 64-bit PHP' );
+    public function test_invalid_otp()
+    {
+        if (PHP_INT_SIZE === 4) {
+            $this->markTestSkipped('calc_totp requires 64-bit PHP');
         }
 
         $key = new Key();
 
         $this->assertFalse(is_valid_authcode($key, '000000'));
     }
+
+    public function pad_short_secrets()
+    {
+        $secret = 'abc';
+
+        $this->assertEquals('abc', pad_secret($secret, 3));
+        $this->assertEquals('abcabc', pad_secret($secret, 6));
+        $this->assertEquals('abcabcabcabc', pad_secret($secret, 12));
+        $this->assertEquals('abcab', pad_secret($secret, 5));
+    }
+
+    public function test_invalid_hash()
+    {
+        $this->expectException(\InvalidArgumentException::class);
+
+        $key = new Key();
+        calc_totp($key, false, 6, 'md5');
+    }
 }
\ No newline at end of file
diff --git a/test/phpunit/ReferenceTest.php b/test/phpunit/ReferenceTest.php
index 851b6ec..11ddce0 100644
--- a/test/phpunit/ReferenceTest.php
+++ b/test/phpunit/ReferenceTest.php
@@ -48,6 +48,12 @@ public function test_sha1() {
 
             $this->assertEquals($vector[0], calc_totp($token, false, 8, $hash, self::$step));
         }
+
+        foreach (self::$vectors as $time => $vector) {
+            self::$now = (int) $time;
+
+            $this->assertEquals(substr($vector[0], 2), calc_totp($token, false, 6, $hash, self::$step));
+        }
     }
 
 
@@ -60,6 +66,12 @@ public function test_sha256() {
 
             $this->assertEquals($vector[1], calc_totp($token, false, 8, $hash, self::$step));
         }
+
+        foreach (self::$vectors as $time => $vector) {
+            self::$now = (int) $time;
+
+            $this->assertEquals(substr($vector[1], 2), calc_totp($token, false, 6, $hash, self::$step));
+        }
     }
 
 
@@ -72,5 +84,11 @@ public function test_sha512() {
 
             $this->assertEquals($vector[2], calc_totp($token, false, 8, $hash, self::$step));
         }
+
+        foreach (self::$vectors as $time => $vector) {
+            self::$now = (int) $time;
+
+            $this->assertEquals(substr($vector[2], 2), calc_totp($token, false, 6, $hash, self::$step));
+        }
     }
 }
\ No newline at end of file