From 73393c9b5e1e06fc1ab3ee5424f033eefa8df5d0 Mon Sep 17 00:00:00 2001
From: Galaxy <30950645+GalaxySciTech@users.noreply.github.com>
Date: Fri, 12 Jan 2024 20:57:16 +0400
Subject: [PATCH 1/2] Update ProxyGateway.sol

---
 contracts/proxy/ProxyGateway.sol | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/contracts/proxy/ProxyGateway.sol b/contracts/proxy/ProxyGateway.sol
index 6115e02..01d0ced 100644
--- a/contracts/proxy/ProxyGateway.sol
+++ b/contracts/proxy/ProxyGateway.sol
@@ -7,6 +7,8 @@ import {ICheckpoint} from "../interfaces/ICheckpoint.sol";
 contract ProxyGateway is ProxyAdmin {
     // 0 full | 1 lite
     mapping(uint256 => TransparentUpgradeableProxy) public cscProxies;
+    // 0 full | 1 lite
+    mapping(uint256 => uint256) public cscVersion;
 
     event CreateProxy(TransparentUpgradeableProxy proxy);
 
@@ -50,6 +52,7 @@ contract ProxyGateway is ProxyAdmin {
             initEpoch
         );
         cscProxies[0] = createProxy(full, data);
+        cscVersion[0]++;
         return cscProxies[0];
     }
 
@@ -77,6 +80,7 @@ contract ProxyGateway is ProxyAdmin {
             initEpoch
         );
         cscProxies[1] = createProxy(lite, data);
+        cscVersion[0]++;
         return cscProxies[1];
     }
 }

From 9a7096ccc2cf266f3fd2d1ac4ab169befc2a4d6f Mon Sep 17 00:00:00 2001
From: Galaxy <30950645+GalaxySciTech@users.noreply.github.com>
Date: Fri, 12 Jan 2024 21:07:23 +0400
Subject: [PATCH 2/2] update

---
 contracts/proxy/ProxyGateway.sol | 28 +++++++++++++++++++++++-----
 test/ProxyGateway.js             |  7 +++++++
 2 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/contracts/proxy/ProxyGateway.sol b/contracts/proxy/ProxyGateway.sol
index 01d0ced..8d8d9ef 100644
--- a/contracts/proxy/ProxyGateway.sol
+++ b/contracts/proxy/ProxyGateway.sol
@@ -1,14 +1,15 @@
 // SPDX-License-Identifier: MIT
 pragma solidity =0.8.19;
 
-import {ProxyAdmin, TransparentUpgradeableProxy} from "@openzeppelin/contracts/proxy/transparent/ProxyAdmin.sol";
+import {ProxyAdmin, TransparentUpgradeableProxy, ITransparentUpgradeableProxy} from "@openzeppelin/contracts/proxy/transparent/ProxyAdmin.sol";
 import {ICheckpoint} from "../interfaces/ICheckpoint.sol";
 
 contract ProxyGateway is ProxyAdmin {
     // 0 full | 1 lite
     mapping(uint256 => TransparentUpgradeableProxy) public cscProxies;
-    // 0 full | 1 lite
-    mapping(uint256 => uint256) public cscVersion;
+
+    //proxy => version
+    mapping(ITransparentUpgradeableProxy => uint256) public version;
 
     event CreateProxy(TransparentUpgradeableProxy proxy);
 
@@ -52,7 +53,7 @@ contract ProxyGateway is ProxyAdmin {
             initEpoch
         );
         cscProxies[0] = createProxy(full, data);
-        cscVersion[0]++;
+
         return cscProxies[0];
     }
 
@@ -80,7 +81,24 @@ contract ProxyGateway is ProxyAdmin {
             initEpoch
         );
         cscProxies[1] = createProxy(lite, data);
-        cscVersion[0]++;
+
         return cscProxies[1];
     }
+
+    function upgrade(
+        ITransparentUpgradeableProxy proxy,
+        address implementation
+    ) public override {
+        super.upgrade(proxy, implementation);
+        version[proxy]++;
+    }
+
+    function upgradeAndCall(
+        ITransparentUpgradeableProxy proxy,
+        address implementation,
+        bytes memory data
+    ) public payable override {
+        super.upgradeAndCall(proxy, implementation, data);
+        version[proxy]++;
+    }
 }
diff --git a/test/ProxyGateway.js b/test/ProxyGateway.js
index 9921ebd..dc17e7f 100644
--- a/test/ProxyGateway.js
+++ b/test/ProxyGateway.js
@@ -91,12 +91,14 @@ describe("proxyGateway", () => {
         900
       );
       const fullProxyAddress = await proxyGateway.cscProxies(0);
+      const fullVersion = await proxyGateway.version(fullProxyAddress);
       const fullProxy = full.attach(fullProxyAddress);
       const fullMode = await fullProxy.MODE();
       const fullGap = await fullProxy.INIT_GAP();
       const fullResult = await fullProxy.getHeaderByNumber(1);
 
       expect(full.address).to.not.eq(fullProxy.address);
+      expect(fullVersion).to.eq(0);
       expect(fullMode).to.eq("full");
       expect(fullGap).to.eq(450);
 
@@ -105,10 +107,15 @@ describe("proxyGateway", () => {
 
       await proxyGateway.upgrade(fullProxyAddress, proxyTest.address);
       const afterUpgradeResult = await fullProxy.getHeaderByNumber(1);
+      const fullVersionAfterUpdate = await proxyGateway.version(
+        fullProxyAddress
+      );
 
       expect(fullResult.hash).to.eq(
         "0x684d18e0081cbe82cab66173647eaf2b078413da5f79a1082a5228314c23ae15"
       );
+
+      expect(fullVersionAfterUpdate).to.eq(1);
       expect(fullResult.number).to.eq(1);
 
       expect(afterUpgradeResult.hash).to.eq(