Update node contract

.mockery.yaml

@@ -1,5 +1,4 @@
 with-expecter: True
-testonly: True
 dir: pkg/mocks
 mockname: "Mock{{.InterfaceName}}"
 outpkg: mocks

contracts/src/Nodes.sol

@@ -1,91 +1,162 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.13;
 
+import "@openzeppelin/contracts/token/ERC721/ERC721.sol";
 import "@openzeppelin/contracts/access/Ownable.sol";
 
-contract Nodes is Ownable {
-    constructor() Ownable(msg.sender) {}
+/**
+A NFT contract for XMTP Node Operators.
 
+The deployer of this contract is responsible for minting NFTs and assinging them to node operators.
+
+All nodes on the network periodically check this contract to determine which nodes they should connect to.
+ */
+contract Nodes is ERC721, Ownable {
+    constructor() ERC721("XMTP Node Operator", "XMTP") Ownable(msg.sender) {}
+
+    // uint16 counter so that we cannot create more than 65k IDs
+    // The ERC721 standard expects the tokenID to be uint256 for standard methods unfortunately
+    uint16 private _nodeIdCounter;
+
+    // A node, as stored in the internal mapping
     struct Node {
+        bytes signingKeyPub;
         string httpAddress;
-        uint256 originatorId;
         bool isHealthy;
-        // Maybe we want a TLS cert separate from the public key for MTLS authenticated connections?
     }
 
-    event NodeUpdate(
-        bytes publicKey,
-        string httpAddress,
-        uint256 originatorId,
-        bool isHealthy
-    );
+    struct NodeWithId {
+        uint16 nodeId;
+        Node node;
+    }
 
-    // List of public keys
-    bytes[] publicKeys;
+    event NodeUpdated(uint256 nodeId, Node node);
 
-    // Mapping of publicKey to node
-    mapping(bytes => Node) public nodes;
+    // Mapping of token ID to Node
+    mapping(uint256 => Node) private _nodes;
 
     /**
-    Add a node to the network
+    Mint a new node NFT and store the metadata in the smart contract
      */
     function addNode(
-        bytes calldata publicKey,
+        address to,
+        bytes calldata signingKeyPub,
+        string calldata httpAddress
+    ) public onlyOwner returns (uint16) {
+        uint16 nodeId = _nodeIdCounter;
+        _mint(to, nodeId);
+        _nodes[nodeId] = Node(signingKeyPub, httpAddress, true);
+        _emitNodeUpdate(nodeId);
+        _nodeIdCounter++;
+
+        return nodeId;
+    }
+
+    /**
+    Override the built in transferFrom function to block NFT owners from transferring
+    node ownership.
+
+    NFT owners are only allowed to update their HTTP address and MTLS cert.
+     */
+    function transferFrom(
+        address from,
+        address to,
+        uint256 tokenId
+    ) public override {
+        require(
+            _msgSender() == owner(),
+            "Only the contract owner can transfer Node ownership"
+        );
+        super.transferFrom(from, to, tokenId);
+    }
+
+    /**
+    Allow a NFT holder to update the HTTP address of their node
+     */
+    function updateHttpAddress(
+        uint256 tokenId,
         string calldata httpAddress
-    ) public onlyOwner {
+    ) public {
         require(
-            bytes(nodes[publicKey].httpAddress).length == 0,
-            "Node already exists"
+            _msgSender() == ownerOf(tokenId),
+            "Only the owner of the Node NFT can update its http address"
         );
+        _nodes[tokenId].httpAddress = httpAddress;
+        _emitNodeUpdate(tokenId);
+    }
 
-        require(bytes(httpAddress).length != 0, "HTTP address is required");
+    /**
+    The contract owner may update the health status of the node.
 
-        nodes[publicKey] = Node({
-            httpAddress: httpAddress,
-            originatorId: publicKeys.length + 1,
-            isHealthy: true
-        });
+    No one else is allowed to call this function.
+     */
+    function updateHealth(uint256 tokenId, bool isHealthy) public onlyOwner {
+        // Make sure that the token exists
+        _requireOwned(tokenId);
+        _nodes[tokenId].isHealthy = isHealthy;
+        _emitNodeUpdate(tokenId);
+    }
 
-        publicKeys.push(publicKey);
+    /**
+    Get a list of healthy nodes with their ID and metadata
+     */
+    function healthyNodes() public view returns (NodeWithId[] memory) {
+        uint16 totalNodeCount = _nodeIdCounter;
+        uint256 healthyCount = 0;
 
-        emit NodeUpdate(publicKey, httpAddress, publicKeys.length, true);
+        // First, count the number of healthy nodes
+        for (uint256 i = 0; i < totalNodeCount; i++) {
+            if (_nodeExists(i) && _nodes[i].isHealthy) {
+                healthyCount++;
+            }
+        }
+
+        // Create an array to store healthy nodes
+        NodeWithId[] memory healthyNodesList = new NodeWithId[](healthyCount);
+        uint256 currentIndex = 0;
+
+        // Populate the array with healthy nodes
+        for (uint16 i = 0; i < totalNodeCount; i++) {
+            if (_nodeExists(i) && _nodes[i].isHealthy) {
+                healthyNodesList[currentIndex] = NodeWithId({
+                    nodeId: i,
+                    node: _nodes[i]
+                });
+                currentIndex++;
+            }
+        }
+
+        return healthyNodesList;
     }
 
     /**
-    The contract owner can use this function to mark a node as unhealthy
-    triggering all other nodes to stop replicating to/from this node
+    Get all nodes regardless of their health status
      */
-    function markNodeUnhealthy(bytes calldata publicKey) public onlyOwner {
-        require(
-            bytes(nodes[publicKey].httpAddress).length != 0,
-            "Node does not exist"
-        );
-        nodes[publicKey].isHealthy = false;
+    function allNodes() public view returns (NodeWithId[] memory) {
+        uint16 totalNodeCount = _nodeIdCounter;
+        NodeWithId[] memory allNodesList = new NodeWithId[](totalNodeCount);
 
-        emit NodeUpdate(
-            publicKey,
-            nodes[publicKey].httpAddress,
-            nodes[publicKey].originatorId,
-            false
-        );
+        for (uint16 i = 0; i < totalNodeCount; i++) {
+            allNodesList[i] = NodeWithId({nodeId: i, node: _nodes[i]});
+        }
+
+        return allNodesList;
     }
 
     /**
-    The contract owner can use this function to mark a node as healthy
-    triggering all other nodes to 
+    Get a node's metadata by ID
      */
-    function markNodeHealthy(bytes calldata publicKey) public onlyOwner {
-        require(
-            bytes(nodes[publicKey].httpAddress).length != 0,
-            "Node does not exist"
-        );
-        nodes[publicKey].isHealthy = true;
+    function getNode(uint256 tokenId) public view returns (Node memory) {
+        _requireOwned(tokenId);
+        return _nodes[tokenId];
+    }
 
-        emit NodeUpdate(
-            publicKey,
-            nodes[publicKey].httpAddress,
-            nodes[publicKey].originatorId,
-            true
-        );
+    function _emitNodeUpdate(uint256 tokenId) private {
+        emit NodeUpdated(tokenId, _nodes[tokenId]);
+    }
+
+    function _nodeExists(uint256 tokenId) private view returns (bool) {
+        address owner = _ownerOf(tokenId);
+        return owner != address(0);
     }
 }

contracts/test/GroupMessage.t.sol

@@ -4,7 +4,7 @@ pragma solidity ^0.8.13;
 import {Test, console} from "forge-std/Test.sol";
 import {GroupMessages} from "../src/GroupMessages.sol";
 
-contract CounterTest is Test {
+contract GroupMessagesTest is Test {
     GroupMessages public groupMessages;
 
     function setUp() public {

contracts/test/Nodes.sol

@@ -0,0 +1,159 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.13;
+
+import {Test, console} from "forge-std/Test.sol";
+import {Nodes} from "../src/Nodes.sol";
+import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
+
+contract NodesTest is Test {
+    Nodes public nodes;
+
+    function setUp() public {
+        nodes = new Nodes();
+    }
+
+    function _genBytes(uint32 length) internal pure returns (bytes memory) {
+        bytes memory message = new bytes(length);
+        for (uint256 i = 0; i < length; i++) {
+            message[i] = bytes1(uint8(i % 256));
+        }
+
+        return message;
+    }
+
+    function _genString(uint32 length) internal pure returns (string memory) {
+        return string(_genBytes(length));
+    }
+
+    function _randomNode(
+        bool isHealthy
+    ) internal pure returns (Nodes.Node memory) {
+        return
+            Nodes.Node({
+                signingKeyPub: _genBytes(32),
+                httpAddress: _genString(32),
+                isHealthy: isHealthy
+            });
+    }
+
+    function test_canAddNode() public {
+        Nodes.Node memory node = _randomNode(true);
+
+        address operatorAddress = vm.randomAddress();
+
+        uint16 nodeId = nodes.addNode(
+            operatorAddress,
+            node.signingKeyPub,
+            node.httpAddress
+        );
+
+        vm.assertEq(nodes.ownerOf(nodeId), operatorAddress);
+        vm.assertEq(nodes.getNode(nodeId).signingKeyPub, node.signingKeyPub);
+        vm.assertEq(nodes.getNode(nodeId).httpAddress, node.httpAddress);
+        vm.assertEq(nodes.getNode(nodeId).isHealthy, true);
+    }
+
+    function test_canAddMultiple() public {
+        Nodes.Node memory node1 = _randomNode(true);
+        Nodes.Node memory node2 = _randomNode(true);
+        Nodes.Node memory node3 = _randomNode(true);
+
+        address operator1 = vm.randomAddress();
+        address operator2 = vm.randomAddress();
+        address operator3 = vm.randomAddress();
+
+        uint16 node1Id = nodes.addNode(
+            operator1,
+            node1.signingKeyPub,
+            node1.httpAddress
+        );
+        nodes.addNode(operator2, node2.signingKeyPub, node2.httpAddress);
+        nodes.addNode(operator3, node3.signingKeyPub, node3.httpAddress);
+
+        Nodes.NodeWithId[] memory healthyNodes = nodes.healthyNodes();
+        vm.assertTrue(healthyNodes.length == 3);
+
+        nodes.updateHealth(node1Id, false);
+        healthyNodes = nodes.healthyNodes();
+        vm.assertTrue(healthyNodes.length == 2);
+    }
+
+    function test_canMarkUnhealthy() public {
+        Nodes.Node memory node = _randomNode(true);
+        address operator = vm.randomAddress();
+
+        uint16 nodeId = nodes.addNode(
+            operator,
+            node.signingKeyPub,
+            node.httpAddress
+        );
+
+        nodes.updateHealth(nodeId, false);
+
+        vm.assertEq(nodes.getNode(nodeId).isHealthy, false);
+        vm.assertEq(nodes.healthyNodes().length, 0);
+    }
+
+    function testFail_ownerCannotUpdateHealth() public {
+        vm.expectRevert(Ownable.OwnableUnauthorizedAccount.selector);
+        Nodes.Node memory node = _randomNode(true);
+        address operator = vm.randomAddress();
+
+        uint16 nodeId = nodes.addNode(
+            operator,
+            node.signingKeyPub,
+            node.httpAddress
+        );
+
+        vm.prank(operator);
+        nodes.updateHealth(nodeId, false);
+    }
+
+    function testFail_ownerCannotTransfer() public {
+        Nodes.Node memory node = _randomNode(true);
+        address operator = vm.randomAddress();
+
+        uint16 nodeId = nodes.addNode(
+            operator,
+            node.signingKeyPub,
+            node.httpAddress
+        );
+
+        vm.prank(operator);
+        nodes.safeTransferFrom(operator, vm.randomAddress(), uint256(nodeId));
+    }
+
+    function test_canChangeHttpAddress() public {
+        Nodes.Node memory node = _randomNode(true);
+        address operator = vm.randomAddress();
+
+        uint16 nodeId = nodes.addNode(
+            operator,
+            node.signingKeyPub,
+            node.httpAddress
+        );
+
+        vm.prank(operator);
+        nodes.updateHttpAddress(nodeId, "new-http-address");
+
+        vm.assertEq(nodes.getNode(nodeId).httpAddress, "new-http-address");
+    }
+
+    function testFail_cannotChangeOtherHttpAddress() public {
+        vm.expectRevert(
+            "Only the owner of the Node NFT can update its http address"
+        );
+
+        Nodes.Node memory node = _randomNode(true);
+        address operator = vm.randomAddress();
+
+        uint16 nodeId = nodes.addNode(
+            operator,
+            node.signingKeyPub,
+            node.httpAddress
+        );
+
+        vm.prank(vm.randomAddress());
+        nodes.updateHttpAddress(nodeId, "new-http-address");
+    }
+}