3 files changed, 62 insertions, 15 deletions

diff --git a/Csrf/TokenStorage/NativeSessionTokenStorage.php b/Csrf/TokenStorage/NativeSessionTokenStorage.php
index 8956743..c01967c 100644
--- a/Csrf/TokenStorage/NativeSessionTokenStorage.php
+++ b/Csrf/TokenStorage/NativeSessionTokenStorage.php
@@ -11,6 +11,8 @@
 
 namespace Symfony\Component\Security\Csrf\TokenStorage;
 
+use Symfony\Component\Security\Csrf\Exception\TokenNotFoundException;
+
 /**
  * Token storage that uses PHP's native session handling.
  *
@@ -49,17 +51,17 @@ class NativeSessionTokenStorage implements TokenStorageInterface
     /**
      * {@inheritdoc}
      */
-    public function getToken($tokenId, $default = null)
+    public function getToken($tokenId)
     {
         if (!$this->sessionStarted) {
             $this->startSession();
         }
 
-        if (isset($_SESSION[$this->namespace][$tokenId])) {
-            return $_SESSION[$this->namespace][$tokenId];
+        if (!isset($_SESSION[$this->namespace][$tokenId])) {
+            throw new TokenNotFoundException('The CSRF token with ID '.$tokenId.' does not exist.');
         }
 
-        return $default;
+        return (string) $_SESSION[$this->namespace][$tokenId];
     }
 
     /**
@@ -71,7 +73,7 @@ class NativeSessionTokenStorage implements TokenStorageInterface
             $this->startSession();
         }
 
-        $_SESSION[$this->namespace][$tokenId] = $token;
+        $_SESSION[$this->namespace][$tokenId] = (string) $token;
     }
 
     /**
@@ -86,6 +88,24 @@ class NativeSessionTokenStorage implements TokenStorageInterface
         return isset($_SESSION[$this->namespace][$tokenId]);
     }
 
+    /**
+     * {@inheritdoc}
+     */
+    public function removeToken($tokenId)
+    {
+        if (!$this->sessionStarted) {
+            $this->startSession();
+        }
+
+        $token = isset($_SESSION[$this->namespace][$tokenId])
+            ? $_SESSION[$this->namespace][$tokenId]
+            : null;
+
+        unset($_SESSION[$this->namespace][$tokenId]);
+
+        return $token;
+    }
+
     private function startSession()
     {
         if (version_compare(PHP_VERSION, '5.4', '>=')) {
diff --git a/Csrf/TokenStorage/SessionTokenStorage.php b/Csrf/TokenStorage/SessionTokenStorage.php
index 3878e4c..f08eb96 100644
--- a/Csrf/TokenStorage/SessionTokenStorage.php
+++ b/Csrf/TokenStorage/SessionTokenStorage.php
@@ -12,6 +12,7 @@
 namespace Symfony\Component\Security\Csrf\TokenStorage;
 
 use Symfony\Component\HttpFoundation\Session\SessionInterface;
+use Symfony\Component\Security\Csrf\Exception\TokenNotFoundException;
 
 /**
  * Token storage that uses a Symfony2 Session object.
@@ -54,13 +55,17 @@ class SessionTokenStorage implements TokenStorageInterface
     /**
      * {@inheritdoc}
      */
-    public function getToken($tokenId, $default = null)
+    public function getToken($tokenId)
     {
         if (!$this->session->isStarted()) {
             $this->session->start();
         }
 
-        return $this->session->get($this->namespace . '/' . $tokenId, $default);
+        if (!$this->session->has($this->namespace.'/'.$tokenId)) {
+            throw new TokenNotFoundException('The CSRF token with ID '.$tokenId.' does not exist.');
+        }
+
+        return (string) $this->session->get($this->namespace.'/'.$tokenId);
     }
 
     /**
@@ -72,7 +77,7 @@ class SessionTokenStorage implements TokenStorageInterface
             $this->session->start();
         }
 
-        $this->session->set($this->namespace . '/' . $tokenId, $token);
+        $this->session->set($this->namespace.'/'.$tokenId, (string) $token);
     }
 
     /**
@@ -84,6 +89,18 @@ class SessionTokenStorage implements TokenStorageInterface
             $this->session->start();
         }
 
-        return $this->session->has($this->namespace . '/' . $tokenId);
+        return $this->session->has($this->namespace.'/'.$tokenId);
+    }
+
+    /**
+     * {@inheritdoc}
+     */
+    public function removeToken($tokenId)
+    {
+        if (!$this->session->isStarted()) {
+            $this->session->start();
+        }
+
+        return $this->session->remove($this->namespace.'/'.$tokenId);
     }
 }
diff --git a/Csrf/TokenStorage/TokenStorageInterface.php b/Csrf/TokenStorage/TokenStorageInterface.php
index 7dba9e5..3fb3191 100644
--- a/Csrf/TokenStorage/TokenStorageInterface.php
+++ b/Csrf/TokenStorage/TokenStorageInterface.php
@@ -23,27 +23,37 @@ interface TokenStorageInterface
      * Reads a stored CSRF token.
      *
      * @param string $tokenId The token ID
-     * @param mixed  $default The value to be returned if no token is set
      *
-     * @return mixed The stored token or the default value, if no token is set
+     * @return string The stored token
+     *
+     * @throws \Symfony\Component\Security\Csrf\Exception\TokenNotFoundException If the token ID does not exist
      */
-    public function getToken($tokenId, $default = null);
+    public function getToken($tokenId);
 
     /**
      * Stores a CSRF token.
      *
      * @param string $tokenId The token ID
-     * @param mixed  $token   The CSRF token
+     * @param string $token   The CSRF token
      */
     public function setToken($tokenId, $token);
 
     /**
+     * Removes a CSRF token.
+     *
+     * @param string $tokenId The token ID
+     *
+     * @return string|null Returns the removed token if one existed, NULL
+     *                     otherwise
+     */
+    public function removeToken($tokenId);
+
+    /**
      * Checks whether a token with the given token ID exists.
      *
      * @param string $tokenId The token ID
      *
-     * @return Boolean Returns true if a token is stored for the given token ID,
-     *                 false otherwise.
+     * @return Boolean Whether a token exists with the given ID
      */
     public function hasToken($tokenId);
 }