bugfix: added extra check to update context with assigned issuer

contracts/src/SchemaRegistry.sol

@@ -15,13 +15,17 @@ contract SchemaRegistry is OwnableUpgradeable {
   IRouter public router;
   /// @dev The list of Schemas, accessed by their ID
   mapping(bytes32 id => Schema schema) private schemas;
+  /// @dev The list of Schemas, assigned to the respective issuers
+  mapping(bytes32 id => address issuer) private schemasIssuers;
   /// @dev The list of Schema IDs
   bytes32[] public schemaIds;
 
   /// @notice Error thrown when an invalid Router address is given
   error RouterInvalid();
   /// @notice Error thrown when a non-issuer tries to call a method that can only be called by an issuer
   error OnlyIssuer();
+  /// @notice Error thrown when a non-assigned issuer tries to call a method that can only be called by an assigned issuer
+  error OnlyAssignedIssuer();
   /// @notice Error thrown when an identical Schema was already registered
   error SchemaAlreadyExists();
   /// @notice Error thrown when attempting to add a Schema without a name
@@ -30,6 +34,8 @@ contract SchemaRegistry is OwnableUpgradeable {
   error SchemaStringMissing();
   /// @notice Error thrown when attempting to get a Schema that is not registered
   error SchemaNotRegistered();
+  /// @notice Error thrown when an invalid Issuer address is given
+  error IssuerInvalid();
 
   /// @notice Event emitted when a Schema is created and registered
   event SchemaCreated(bytes32 indexed id, string name, string description, string context, string schemaString);
@@ -65,6 +71,16 @@ contract SchemaRegistry is OwnableUpgradeable {
     router = IRouter(_router);
   }
 
+  /**
+   * @notice Updates schema issuer, adds it if not present for historical schemas
+   * @dev Updates issuer for the given schemaId in the mapping schemaIssuers and adds the issuer if no record found.
+   */
+  function updateSchemaIssuer(bytes32 _schemaId, address _issuer) public onlyOwner {
+    if (!isRegistered(_schemaId)) revert SchemaNotRegistered();
+    if (_issuer == address(0)) revert IssuerInvalid();
+    schemasIssuers[_schemaId] = _issuer;
+  }
+
   /**
    * Generate an ID for a given schema
    * @param schema the string defining a schema
@@ -103,6 +119,7 @@ contract SchemaRegistry is OwnableUpgradeable {
 
     schemas[schemaId] = Schema(name, description, context, schemaString);
     schemaIds.push(schemaId);
+    schemasIssuers[schemaId] = msg.sender;
     emit SchemaCreated(schemaId, name, description, context, schemaString);
   }
 
@@ -114,6 +131,7 @@ contract SchemaRegistry is OwnableUpgradeable {
    */
   function updateContext(bytes32 schemaId, string memory context) public onlyIssuers(msg.sender) {
     if (!isRegistered(schemaId)) revert SchemaNotRegistered();
+    if (schemasIssuers[schemaId] != msg.sender) revert OnlyAssignedIssuer();
     schemas[schemaId].context = context;
   }
 

contracts/test/SchemaRegistry.t.sol

@@ -17,6 +17,7 @@ contract SchemaRegistryTest is Test {
   string private expectedContext = "Context";
   string private expectedString = "this is a schema";
   address private user = makeAddr("user");
+  address private unassignedUser = makeAddr("unassignedUser");
 
   event SchemaCreated(bytes32 indexed id, string name, string description, string context, string schemaString);
   event Initialized(uint8 version);
@@ -32,6 +33,7 @@ contract SchemaRegistryTest is Test {
     portalRegistryAddress = address(portalRegistryMock);
     router.updatePortalRegistry(portalRegistryAddress);
     portalRegistryMock.setIssuer(user);
+    portalRegistryMock.setIssuer(unassignedUser);
   }
 
   function test_initialize_ContractAlreadyInitialized() public {
@@ -56,6 +58,27 @@ contract SchemaRegistryTest is Test {
     testSchemaRegistry.updateRouter(address(0));
   }
 
+  function test_updateSchemaIssuer() public {
+    vm.prank(user);
+    schemaRegistry.createSchema(expectedName, expectedDescription, expectedContext, expectedString);
+    vm.prank(address(0));
+    schemaRegistry.updateSchemaIssuer(expectedId, address(2));
+  }
+
+  function test_updateSchemaIssuer_SchemaNotRegistered() public {
+    vm.expectRevert(SchemaRegistry.SchemaNotRegistered.selector);
+    vm.prank(address(0));
+    schemaRegistry.updateSchemaIssuer(expectedId, address(0));
+  }
+
+  function test_updateSchemaIssuer_IssuerInvalid() public {
+    vm.prank(user);
+    schemaRegistry.createSchema(expectedName, expectedDescription, expectedContext, expectedString);
+    vm.prank(address(0));
+    vm.expectRevert(SchemaRegistry.IssuerInvalid.selector);
+    schemaRegistry.updateSchemaIssuer(expectedId, address(0));
+  }
+
   function test_getIdFromSchemaString() public {
     bytes32 id = schemaRegistry.getIdFromSchemaString(expectedString);
     assertEq(id, expectedId);
@@ -132,6 +155,15 @@ contract SchemaRegistryTest is Test {
     vm.stopPrank();
   }
 
+  function test_updateContext_OnlyAssignedIssuer() public {
+    vm.startPrank(user);
+    schemaRegistry.createSchema(expectedName, expectedDescription, expectedContext, expectedString);
+    vm.expectRevert(SchemaRegistry.OnlyAssignedIssuer.selector);
+    vm.startPrank(unassignedUser);
+    schemaRegistry.updateContext(expectedId, "New context");
+    vm.stopPrank();
+  }
+
   function test_updateContext_withEmptyContext() public {
     vm.startPrank(user);
     schemaRegistry.createSchema(expectedName, expectedDescription, expectedContext, expectedString);