Merge pull request #58 from ydb-platform/fix-token-refresh

ilyakharev · web-flow · commit 83b09252468d · 2023-05-23T18:28:22.000+03:00
Fixed refresh token when it expired
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
@@ -1,4 +1,4 @@
-name: CI test connection
+name: CI tests
 
 on:
   pull_request:
@@ -9,7 +9,7 @@ permissions:
   contents: read
 
 jobs:
-  build:
+  tests:
     services:
       ydb:
         image: cr.yandex/yc/yandex-docker-local-ydb:stable-22-5
@@ -23,7 +23,9 @@ jobs:
           YDB_USE_IN_MEMORY_PDISKS: true
         options: '-h localhost'
     runs-on: ubuntu-latest
-
+    strategy:
+      matrix:
+        php-versions: [ '7.4' ]
     steps:
       - uses: actions/checkout@v3
         name: Checkout
@@ -38,7 +40,7 @@ jobs:
         id: php
         with:
           extensions: grpc
-          php-version: 7.4
+          php-version: ${{ matrix.php-versions }}
 
       - run: composer validate --strict
         name: Validate composer.json and composer.lock
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,4 @@
+* fixed refresh token when it expired
 * fixed retry at BAD_SESSION
 * added credentials authentication
 * added CI test
diff --git a/composer.json b/composer.json
@@ -20,7 +20,7 @@
         "psr/log": "^1|^2|^3"
     },
     "require-dev": {
-        "phpunit/phpunit": "9.*"
+        "phpunit/phpunit": ">= 6.0, <10.0"
     },
     "autoload": {
         "psr-4": {
diff --git a/src/Auth/TokenInfo.php b/src/Auth/TokenInfo.php
@@ -4,6 +4,7 @@
 
 class TokenInfo
 {
+    const _PRIVATE_REFRESH_RATIO = 0.1;
     /**
      * @var string
      */
@@ -13,10 +14,13 @@ class TokenInfo
      */
     protected $expiresAt;
 
+    private $refreshAt;
+
     public function __construct(string $token, int $expiresAt)
     {
         $this->token = $token;
         $this->expiresAt = $expiresAt;
+        $this->refreshAt = time() + round(TokenInfo::_PRIVATE_REFRESH_RATIO*($this->expiresAt-time()),0);
     }
 
     /**
@@ -34,4 +38,9 @@ public function getToken(): string
     {
         return $this->token;
     }
+
+    public function getRefreshAt(): int
+    {
+        return $this->refreshAt;
+    }
 }
diff --git a/src/Discovery.php b/src/Discovery.php
@@ -30,6 +30,10 @@ class Discovery
      * @var LoggerInterface
      */
     protected $logger;
+    /**
+     * @var Iam
+     */
+    protected $credentials;
 
     /**
      * @param Ydb $ydb
@@ -43,6 +47,8 @@ public function __construct(Ydb $ydb, LoggerInterface $logger = null)
 
         $this->meta = $ydb->meta();
 
+        $this->credentials = $ydb->iam();
+
         $this->database = $ydb->database();
 
         $this->logger = $logger;
@@ -82,4 +88,4 @@ protected function request($method, array $data = [])
     {
         return $this->doRequest('Discovery', $method, $data);
     }
-}
+}
diff --git a/src/Iam.php b/src/Iam.php
@@ -45,6 +45,11 @@ class Iam implements IamTokenContract
      */
     protected $logger;
 
+    /**
+     * @var int
+     */
+    protected $refresh_at;
+
     /**
      * @param array $config
      * @param LoggerInterface|null $logger
@@ -96,9 +101,11 @@ public function newToken()
         $tokenInfo = $this->config('credentials')->getTokenInfo();
         $this->iam_token = $tokenInfo->getToken();
         $this->expires_at = $tokenInfo->getExpiresAt();
+        $this->refresh_at = $tokenInfo->getRefreshAt();
         $this->saveToken((object)[
             "iamToken" => $tokenInfo->getToken(),
             "expiresAt" => $tokenInfo->getExpiresAt(),
+            "refreshAt" => $tokenInfo->getRefreshAt()
         ]);
         return $tokenInfo->getToken();
     }
@@ -347,11 +354,14 @@ protected function loadToken()
     {
         if ($this->iam_token)
         {
-            if ($this->expires_at > time())
-            {
-                return $this->iam_token;
+            if ($this->refresh_at <= time()){
+                try {
+                    return $this->newToken();
+                } catch (\Exception $e){
+                    return $this->iam_token;
+                }
             }
-            return $this->newToken();
+            return $this->iam_token;
         }
 
         return $this->loadTokenFromFile();
@@ -372,6 +382,7 @@ protected function loadTokenFromFile()
             {
                 $this->iam_token = $token->iamToken;
                 $this->expires_at = $token->expiresAt;
+                $this->refresh_at = $token->refreshAt ?? time();
                 $this->logger()->info('YDB: Reused IAM token [...' . substr($this->iam_token, -6) . '].');
                 return $token->iamToken;
             }
@@ -390,10 +401,12 @@ protected function saveToken($token)
 
         $this->iam_token = $token->iamToken;
         $this->expires_at = $this->convertExpiresAt($token->expiresAt ?? '');
+        $this->refresh_at = $token->refreshAt;
 
         file_put_contents($tokenFile, json_encode([
             'iamToken' => $this->iam_token,
             'expiresAt' => $this->expires_at,
+            'refreshAt' => $this->refresh_at
         ]));
     }
 
diff --git a/src/Operations.php b/src/Operations.php
@@ -25,6 +25,10 @@ class Operations
      * @var LoggerInterface
      */
     protected $logger;
+    /**
+     * @var Iam
+     */
+    protected $credentials;
 
     /**
      * @param Ydb $ydb
@@ -38,6 +42,8 @@ public function __construct(Ydb $ydb, LoggerInterface $logger = null)
 
         $this->meta = $ydb->meta();
 
+        $this->credentials = $ydb->iam();
+
         $this->logger = $logger;
     }
 
diff --git a/src/Scheme.php b/src/Scheme.php
@@ -45,6 +45,8 @@ public function __construct(Ydb $ydb, LoggerInterface $logger = null)
 
         $this->meta = $ydb->meta();
 
+        $this->credentials = $ydb->iam();
+
         $this->database = $ydb->database();
 
         $this->logger = $logger;
diff --git a/src/Scripting.php b/src/Scripting.php
@@ -25,6 +25,10 @@ class Scripting
      * @var LoggerInterface
      */
     protected $logger;
+    /**
+     * @var Iam
+     */
+    protected $credentials;
 
     /**
      * @param Ydb $ydb
@@ -38,6 +42,8 @@ public function __construct(Ydb $ydb, LoggerInterface $logger = null)
 
         $this->meta = $ydb->meta();
 
+        $this->credentials = $ydb->iam();
+
         $this->logger = $logger;
     }
 
diff --git a/src/Session.php b/src/Session.php
@@ -82,6 +82,7 @@ public function __construct(Table $table, $session_id)
 
         $this->client = $table->client();
         $this->meta = $table->meta();
+        $this->credentials = $table->credentials();
         $this->path = $table->path();
         $this->logger = $table->getLogger();
 
diff --git a/src/Table.php b/src/Table.php
@@ -41,6 +41,10 @@ class Table
      * @var LoggerInterface
      */
     protected $logger;
+    /**
+     * @var Iam
+     */
+    protected $credentials;
 
     /**
      * @param Ydb $ydb
@@ -54,6 +58,8 @@ public function __construct(Ydb $ydb, LoggerInterface $logger = null)
 
         $this->meta = $ydb->meta();
 
+        $this->credentials = $ydb->iam();
+
         $this->path = $ydb->database();
 
         $this->logger = $logger;
@@ -284,6 +290,14 @@ public function createTable($table, $columns, $primary_key = 'id', $indexes = []
         return $this->session()->createTable($table, $columns, $primary_key, $indexes);
     }
 
+    /**
+     * @return Iam
+     */
+    public function credentials(): Iam
+    {
+        return $this->credentials;
+    }
+
     /**
      * Proxy to Session::copyTable.
      *
@@ -407,4 +421,4 @@ protected function streamRequest($method, array $data = [])
         return $this->doStreamRequest('Table', $method, $data);
     }
 
-}
+}
diff --git a/src/Traits/RequestTrait.php b/src/Traits/RequestTrait.php
@@ -41,6 +41,8 @@ trait RequestTrait
      */
     protected function doRequest($service, $method, array $data = [])
     {
+        $this->meta['x-ydb-auth-ticket'] = [$this->credentials->token()];
+
         $this->saveLastRequest($service, $method, $data);
 
         $requestClass = '\\Ydb\\' . $service . '\\' . $method . 'Request';
diff --git a/tests/RefreshTokenTest.php b/tests/RefreshTokenTest.php
@@ -0,0 +1,113 @@
+<?php
+
+namespace YdbPlatform\Ydb\Test;
+
+use PHPUnit\Framework\TestCase;
+use YdbPlatform\Ydb\Auth\Auth;
+use YdbPlatform\Ydb\Auth\TokenInfo;
+use YdbPlatform\Ydb\Ydb;
+
+class FakeCredentials extends Auth {
+
+    /**
+     * @var int
+     */
+    protected $counter;
+
+    protected $tokenLiveTime;
+
+    public function __construct(&$counter, &$tokenLiveTime)
+    {
+        $this->counter = &$counter;
+        $this->tokenLiveTime = &$tokenLiveTime;
+    }
+
+    public function getTokenInfo(): TokenInfo
+    {
+        $this->counter++;
+        if ($this->counter==2){
+            throw new \Exception("Some error");
+        }
+        return new TokenInfo(time()+$this->tokenLiveTime,time()+$this->tokenLiveTime);
+    }
+
+    public function getName(): string
+    {
+        return "FakeCredentials";
+    }
+}
+
+class MetaGetter extends \YdbPlatform\Ydb\Session{
+    public static function getMeta(\YdbPlatform\Ydb\Session $session){
+        return $session->meta;
+    }
+}
+
+class RefreshTokenTest extends TestCase
+{
+    public function test(){
+
+        $counter = 0;
+
+        $TOKEN_LIVE_TIME = 10;
+
+        $config = [
+
+            // Database path
+            'database'    => '/local',
+
+            // Database endpoint
+            'endpoint'    => 'localhost:2136',
+
+            // Auto discovery (dedicated server only)
+            'discovery'   => false,
+
+            // IAM config
+            'iam_config'  => [
+                'insecure' => true,
+            ],
+            'credentials' => new FakeCredentials($counter, $TOKEN_LIVE_TIME)
+        ];
+        $ydb = new Ydb($config);
+        $table = $ydb->table();
+        $session = $table->session();
+        $token = MetaGetter::getMeta($session)["x-ydb-auth-ticket"][0];
+        self::assertEquals(
+            1,
+            $counter
+        );
+
+        // Check that the token will not be updated until a refresh time
+        $session->query('select 1 as res');
+        self::assertEquals(
+            1,
+            $counter
+        );
+        self::assertEquals(
+            $token,
+            MetaGetter::getMeta($session)["x-ydb-auth-ticket"][0]
+        );
+        // Check that sdk used old token when failed refreshing
+        usleep(TokenInfo::_PRIVATE_REFRESH_RATIO*$TOKEN_LIVE_TIME*1000*1000); // waiting 10% from token live time
+        $session->query('select 1 as res');
+        self::assertEquals(
+            2,
+            $counter
+        );
+        self::assertEquals(
+            $token,
+            MetaGetter::getMeta($session)["x-ydb-auth-ticket"][0]
+        );
+
+        // Check that token refreshed
+        $session->query('select 1 as res');
+        self::assertEquals(
+            3,
+            $counter
+        );
+        self::assertNotEquals(
+            $token,
+            MetaGetter::getMeta($session)["x-ydb-auth-ticket"][0]
+        );
+    }
+}
diff --git a/tests/RetryOnBadSessionTest.php b/tests/RetryOnBadSessionTest.php

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,4 @@`
	`1`	`+* fixed refresh token when it expired`
`1`	`2`	`* fixed retry at BAD_SESSION`
`2`	`3`	`* added credentials authentication`
`3`	`4`	`* added CI test`
Original file line number	Diff line number	Diff line change
`@@ -4,6 +4,7 @@`
`4`	`4`
`5`	`5`	`class TokenInfo`
`6`	`6`	`{`
	`7`	`+ const _PRIVATE_REFRESH_RATIO = 0.1;`
`7`	`8`	`/**`
`8`	`9`	`* @var string`
`9`	`10`	`*/`
`@@ -13,10 +14,13 @@ class TokenInfo`
`13`	`14`	`*/`
`14`	`15`	`protected $expiresAt;`
`15`	`16`
	`17`	`+ private $refreshAt;`
	`18`	`+`
`16`	`19`	`public function __construct(string $token, int $expiresAt)`
`17`	`20`	`{`
`18`	`21`	`$this->token = $token;`
`19`	`22`	`$this->expiresAt = $expiresAt;`
	`23`	`+ $this->refreshAt = time() + round(TokenInfo::_PRIVATE_REFRESH_RATIO*($this->expiresAt-time()),0);`
`20`	`24`	`}`
`21`	`25`
`22`	`26`	`/**`
`@@ -34,4 +38,9 @@ public function getToken(): string`
`34`	`38`	`{`
`35`	`39`	`return $this->token;`
`36`	`40`	`}`
	`41`	`+`
	`42`	`+ public function getRefreshAt(): int`
	`43`	`+ {`
	`44`	`+ return $this->refreshAt;`
	`45`	`+ }`
`37`	`46`	`}`