From af8125591e75992c0d91ac1302668426ee5c1e63 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nicol=C3=A1s=20Venturo?= <nicolas.venturo@gmail.com>
Date: Thu, 1 Aug 2019 20:06:14 -0300
Subject: [PATCH] Add _upgradeRelayHub

---
 contracts/gsn/GSNContext.sol       | 11 +++++++++
 contracts/mocks/GSNContextMock.sol |  8 ++++++
 test/gsn/GSNContext.test.js        | 39 ++++++++++++++++++++++++++++--
 3 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/contracts/gsn/GSNContext.sol b/contracts/gsn/GSNContext.sol
index 1732270ce..480e3f04f 100644
--- a/contracts/gsn/GSNContext.sol
+++ b/contracts/gsn/GSNContext.sol
@@ -13,10 +13,21 @@ import "./Context.sol";
 contract GSNContext is Context {
     address private _relayHub = 0x537F27a04470242ff6b2c3ad247A05248d0d27CE;
 
+    event RelayHubUpgraded(address indexed oldRelayHub, address indexed newRelayHub);
+
     function _getRelayHub() internal view returns (address) {
         return _relayHub;
     }
 
+    function _upgradeRelayHub(address newRelayHub) internal {
+        require(newRelayHub != address(0), "GSNContext: new RelayHub is the zero address");
+        require(newRelayHub != _relayHub, "GSNContext: new RelayHub is the current one");
+
+        emit RelayHubUpgraded(_relayHub, newRelayHub);
+
+        _relayHub = newRelayHub;
+    }
+
     // Overrides for Context's functions: when called from RelayHub, sender and
     // data require some pre-processing: the actual sender is stored at the end
     // of the call data, which in turns means it needs to be removed from it
diff --git a/contracts/mocks/GSNContextMock.sol b/contracts/mocks/GSNContextMock.sol
index 1ec111ca6..08b6412ec 100644
--- a/contracts/mocks/GSNContextMock.sol
+++ b/contracts/mocks/GSNContextMock.sol
@@ -35,4 +35,12 @@ contract GSNContextMock is ContextMock, GSNContext, IRelayRecipient {
     function postRelayedCall(bytes calldata, bool, uint256, bytes32) external {
         // solhint-disable-previous-line no-empty-blocks
     }
+
+    function getRelayHub() public view returns (address) {
+        return _getRelayHub();
+    }
+
+    function upgradeRelayHub(address newRelayHub) public {
+        return _upgradeRelayHub(newRelayHub);
+    }
 }
diff --git a/test/gsn/GSNContext.test.js b/test/gsn/GSNContext.test.js
index 5c46f34a0..ea284aa24 100644
--- a/test/gsn/GSNContext.test.js
+++ b/test/gsn/GSNContext.test.js
@@ -1,4 +1,5 @@
-const { BN, expectEvent } = require('openzeppelin-test-helpers');
+const { BN, constants, expectEvent, expectRevert } = require('openzeppelin-test-helpers');
+const { ZERO_ADDRESS } = constants;
 const gsn = require('@openzeppelin/gsn-helpers');
 
 const GSNContextMock = artifacts.require('GSNContextMock');
@@ -6,12 +7,46 @@ const ContextMockCaller = artifacts.require('ContextMockCaller');
 
 const { shouldBehaveLikeRegularContext } = require('./Context.behavior');
 
-contract('GSNContext', function ([_, deployer, sender]) {
+contract('GSNContext', function ([_, deployer, sender, newRelayHub]) {
   beforeEach(async function () {
     this.context = await GSNContextMock.new();
     this.caller = await ContextMockCaller.new();
   });
 
+  describe('get/set RelayHub', function () {
+    const singletonRelayHub = '0x537F27a04470242ff6b2c3ad247A05248d0d27CE';
+
+    it('initially returns the singleton instance address', async function () {
+      expect(await this.context.getRelayHub()).to.equal(singletonRelayHub);
+    });
+
+    it('can be upgraded to a new RelayHub', async function () {
+      const { logs } = await this.context.upgradeRelayHub(newRelayHub);
+      expectEvent.inLogs(logs, 'RelayHubUpgraded', { oldRelayHub: singletonRelayHub, newRelayHub });
+    });
+
+    it('cannot upgrade to the same RelayHub', async function () {
+      await expectRevert(
+        this.context.upgradeRelayHub(singletonRelayHub),
+        'GSNContext: new RelayHub is the current one'
+      );
+    });
+
+    it('cannot upgrade to the zero address', async function () {
+      await expectRevert(this.context.upgradeRelayHub(ZERO_ADDRESS), 'GSNContext: new RelayHub is the zero address');
+    });
+
+    context('with new RelayHub', function () {
+      beforeEach(async function () {
+        await this.context.upgradeRelayHub(newRelayHub);
+      });
+
+      it('returns the new instance address', async function () {
+        expect(await this.context.getRelayHub()).to.equal(newRelayHub);
+      });
+    });
+  });
+
   context('when called directly', function () {
     shouldBehaveLikeRegularContext(sender);
   });