diff --git a/core/contracts/staking/TokenStaking.sol b/core/contracts/staking/TokenStaking.sol
index defd83e1a..aa61fc59d 100644
--- a/core/contracts/staking/TokenStaking.sol
+++ b/core/contracts/staking/TokenStaking.sol
@@ -1,16 +1,20 @@
 // SPDX-License-Identifier: GPL-3.0-only
 
-pragma solidity ^0.8.9;
+pragma solidity ^0.8.20;
+
+import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 
 /// @title TokenStaking
 /// @notice A token staking contract for a specified standard ERC20 token. A
 ///         holder of the specified token can stake its tokens to this contract
 ///         and recover the stake after undelegation period is over.
 contract TokenStaking {
-    // TODO: use IERC20 contract as type
-    address internal immutable token;
+    using SafeERC20 for IERC20;
+    
+    IERC20 internal immutable token;
 
-    constructor(address _token) {
+    constructor(IERC20 _token) {
         require(
             address(_token) != address(0),
             "Token can not be the zero address"
diff --git a/core/hardhat.config.ts b/core/hardhat.config.ts
index 687d555c4..0e91f14aa 100644
--- a/core/hardhat.config.ts
+++ b/core/hardhat.config.ts
@@ -7,7 +7,7 @@ const config: HardhatUserConfig = {
   solidity: {
     compilers: [
       {
-        version: "0.8.17",
+        version: "0.8.20",
         settings: {
           optimizer: {
             enabled: true,
diff --git a/core/package.json b/core/package.json
index 983dd5827..fb677e3c7 100644
--- a/core/package.json
+++ b/core/package.json
@@ -56,5 +56,8 @@
     "ts-node": ">=8.0.0",
     "typechain": "^8.1.0",
     "typescript": ">=4.5.0"
+  },
+  "dependencies": {
+    "@openzeppelin/contracts": "^5.0.0"
   }
 }
diff --git a/core/yarn.lock b/core/yarn.lock
index e2d8e3206..f61f4d6f0 100644
--- a/core/yarn.lock
+++ b/core/yarn.lock
@@ -879,6 +879,11 @@
   resolved "https://registry.yarnpkg.com/@nomiclabs/hardhat-waffle/-/hardhat-waffle-2.0.6.tgz#d11cb063a5f61a77806053e54009c40ddee49a54"
   integrity sha512-+Wz0hwmJGSI17B+BhU/qFRZ1l6/xMW82QGXE/Gi+WTmwgJrQefuBs1lIf7hzQ1hLk6hpkvb/zwcNkpVKRYTQYg==
 
+"@openzeppelin/contracts@^5.0.0":
+  version "5.0.0"
+  resolved "https://registry.yarnpkg.com/@openzeppelin/contracts/-/contracts-5.0.0.tgz#ee0e4b4564f101a5c4ee398cd4d73c0bd92b289c"
+  integrity sha512-bv2sdS6LKqVVMLI5+zqnNrNU/CA+6z6CmwFXm/MzmOPBRSO5reEJN7z0Gbzvs0/bv/MZZXNklubpwy3v2+azsw==
+
 "@openzeppelin/defender-admin-client@^1.48.0":
   version "1.49.0"
   resolved "https://registry.yarnpkg.com/@openzeppelin/defender-admin-client/-/defender-admin-client-1.49.0.tgz#ed07318ccba10ac8a8a33cf594fc18b7ab5889f9"