From 536ef3244d58a4bea2283b1ef38f20157823576d Mon Sep 17 00:00:00 2001
From: Alex Bilbie <hello@alexbilbie.com>
Date: Tue, 30 Sep 2014 22:28:49 +0100
Subject: [PATCH] Inject the session into the token type

---
 src/Grant/AuthCodeGrant.php          |  1 +
 src/Grant/ClientCredentialsGrant.php |  1 +
 src/Grant/PasswordGrant.php          |  1 +
 src/Grant/RefreshTokenGrant.php      |  1 +
 src/TokenType/AbstractTokenType.php  | 18 ++++++++++++++++++
 5 files changed, 22 insertions(+)

diff --git a/src/Grant/AuthCodeGrant.php b/src/Grant/AuthCodeGrant.php
index d73c8af9..1d930e36 100644
--- a/src/Grant/AuthCodeGrant.php
+++ b/src/Grant/AuthCodeGrant.php
@@ -230,6 +230,7 @@ class AuthCodeGrant extends AbstractGrant
            $accessToken->associateScope($scope);
         }
 
+        $this->server->getTokenType()->setSession($session);
         $this->server->getTokenType()->setParam('access_token', $accessToken->getId());
         $this->server->getTokenType()->setParam('expires_in', $this->getAccessTokenTTL());
 
diff --git a/src/Grant/ClientCredentialsGrant.php b/src/Grant/ClientCredentialsGrant.php
index 94e1094c..24c7b44c 100644
--- a/src/Grant/ClientCredentialsGrant.php
+++ b/src/Grant/ClientCredentialsGrant.php
@@ -112,6 +112,7 @@ class ClientCredentialsGrant extends AbstractGrant
         $accessToken->setSession($session);
         $accessToken->save($this->server->getStorage('access_token'));
 
+        $this->server->getTokenType()->setSession($session);
         $this->server->getTokenType()->setParam('access_token', $accessToken->getId());
         $this->server->getTokenType()->setParam('expires_in', $this->getAccessTokenTTL());
 
diff --git a/src/Grant/PasswordGrant.php b/src/Grant/PasswordGrant.php
index 4fb35126..cd778356 100644
--- a/src/Grant/PasswordGrant.php
+++ b/src/Grant/PasswordGrant.php
@@ -148,6 +148,7 @@ class PasswordGrant extends AbstractGrant
            $accessToken->associateScope($scope);
         }
 
+        $this->server->getTokenType()->setSession($session);
         $this->server->getTokenType()->setParam('access_token', $accessToken->getId());
         $this->server->getTokenType()->setParam('expires_in', $this->getAccessTokenTTL());
 
diff --git a/src/Grant/RefreshTokenGrant.php b/src/Grant/RefreshTokenGrant.php
index 1bfe9e68..b6a1b0ac 100644
--- a/src/Grant/RefreshTokenGrant.php
+++ b/src/Grant/RefreshTokenGrant.php
@@ -139,6 +139,7 @@ class RefreshTokenGrant extends AbstractGrant
         $oldAccessToken->expire($this->server->getStorage('access_token'));
         $newAccessToken->save($this->server->getStorage('access_token'));
 
+        $this->server->getTokenType()->setSession($session);
         $this->server->getTokenType()->setParam('access_token', $newAccessToken->getId());
         $this->server->getTokenType()->setParam('expires_in', $this->getAccessTokenTTL());
 
diff --git a/src/TokenType/AbstractTokenType.php b/src/TokenType/AbstractTokenType.php
index fb6cd6fb..7691ec9e 100644
--- a/src/TokenType/AbstractTokenType.php
+++ b/src/TokenType/AbstractTokenType.php
@@ -13,6 +13,7 @@ namespace League\OAuth2\Server\TokenType;
 
 use Symfony\Component\HttpFoundation\Request;
 use League\OAuth2\Server\AbstractServer;
+use League\OAuth2\Server\Entity\SessionEntity;
 
 abstract class AbstractTokenType
 {
@@ -28,6 +29,12 @@ abstract class AbstractTokenType
      */
     protected $server;
 
+    /**
+     * Server
+     * @var \League\OAuth2\Server\Entity\SessionEntity $session
+     */
+    protected $session;
+
     /**
      * Set the server
      * @param \League\OAuth2\Server\AbstractServer $server
@@ -38,6 +45,17 @@ abstract class AbstractTokenType
         return $this;
     }
 
+    /**
+     * Set the session entity
+     * @param \League\OAuth2\Server\Entity\SessionEntity $session
+     */
+    public function setSession(SessionEntity $session)
+    {
+        $this->session = $session;
+        return $this;
+    }
+
+    /**
      * Set a key/value response pair
      * @param string $key
      * @param mixed  $value