Skip to content

Commit

Permalink
re-write bls light account contract
Browse files Browse the repository at this point in the history
  • Loading branch information
@fanhousanbu
fanhousanbu committed Oct 4, 2024
1 parent b96ba11 commit cfe3a9e
Show file tree
Hide file tree
Showing 3 changed files with 190 additions and 90 deletions.
2 changes: 1 addition & 1 deletion cache/solidity-files-cache.json
View file

Large diffs are not rendered by default.

235 changes: 165 additions & 70 deletions src/BLSLightAccount.sol
View file
Original file line number Diff line number Diff line change
@@ -1,93 +1,134 @@
// SPDX-License-Identifier: GPL-3.0
pragma solidity ^0.8.23;

//the key components of this new BLSLightAccount contract:
// We inherit from LightAccount to maintain all original functionality.
// We introduce a new storage struct BLSAccountStorage to store the BLS public key.
// We override the initialize function to include setting the BLS public key.
// We add a setBlsPublicKey function to allow updating the BLS public key.
// We override the _validateSignature function to handle BLS signatures.
// We introduce a placeholder _validateBLSSignature function that needs to be implemented.

// TODO: Implement the BLS signature validation logic in _validateBLSSignature.
// Done: Update the SignatureType enum in BaseLightAccount to include the BLS type.
// Possibly integrate with a BLS library for the necessary cryptographic operations.

import "@openzeppelin/contracts/utils/cryptography/MessageHashUtils.sol";
import {LightAccount} from "../lib/light-account/src/LightAccount.sol";
import {IEntryPoint} from "../lib/light-account/lib/account-abstraction/contracts/interfaces/IEntryPoint.sol";
import {PackedUserOperation} from "../lib/light-account/lib/account-abstraction/contracts/interfaces/PackedUserOperation.sol";

/// @title A LightAccount with added BLS signature support
/// @dev Extends LightAccount to support BLS signatures for user operations
contract BLSLightAccount is LightAccount {
// BLS-specific storage
struct BLSAccountStorage {
bytes32 blsPublicKey;
import {ECDSA} from "@openzeppelin/contracts/utils/cryptography/ECDSA.sol";
import {MessageHashUtils} from "@openzeppelin/contracts/utils/cryptography/MessageHashUtils.sol";
import {SignatureChecker} from "@openzeppelin/contracts/utils/cryptography/SignatureChecker.sol";
import {SIG_VALIDATION_FAILED} from "account-abstraction/core/Helpers.sol";
import {IEntryPoint} from "account-abstraction/interfaces/IEntryPoint.sol";
import {PackedUserOperation} from "account-abstraction/interfaces/PackedUserOperation.sol";

import {BaseLightAccount} from "light-account/src/common/BaseLightAccount.sol";
import {CustomSlotInitializable} from "light-account/src/common/CustomSlotInitializable.sol";

/// @title A simple ERC-4337 compatible smart contract account with a designated owner account.
/// @dev Like eth-infinitism's SimpleAccount, but with the following changes:
///
/// 1. Instead of the default storage slots, uses namespaced storage to avoid clashes when switching implementations.
///
/// 2. Ownership can be transferred via `transferOwnership`, similar to the behavior of an `Ownable` contract. This is
/// a simple single-step operation, so care must be taken to ensure that the ownership is being transferred to the
/// correct address.
///
/// 3. Supports [ERC-1271](https://eips.ethereum.org/EIPS/eip-1271) signature validation for both validating the
/// signature on user operations and in exposing its own `isValidSignature` method. This only works when the owner of
/// LightAccount also support ERC-1271.
///
/// ERC-4337's bundler validation rules limit the types of contracts that can be used as owners to validate user
/// operation signatures. For example, the contract's `isValidSignature` function may not use any forbidden opcodes
/// such as `TIMESTAMP` or `NUMBER`, and the contract may not be an ERC-1967 proxy as it accesses a constant
/// implementation slot not associated with the account, violating storage access rules. This also means that the
/// owner of a LightAccount may not be another LightAccount if you want to send user operations through a bundler.
///
/// 4. Event `SimpleAccountInitialized` renamed to `LightAccountInitialized`.
///
/// 5. Uses custom errors.
contract BLSLightAccount is BaseLightAccount, CustomSlotInitializable {
using ECDSA for bytes32;
using MessageHashUtils for bytes32;

/// @dev The version used for namespaced storage is not linked to the release version of the contract. Storage
/// versions will be updated only when storage layout changes are made.
/// keccak256(abi.encode(uint256(keccak256("blslight_account_v1.storage")) - 1)) & ~bytes32(uint256(0xff));
bytes32 internal constant _STORAGE_POSITION = 0x99f82d893aaaf0bfbe72b23177e6b97cf253cbd01abd0d2f793f75ec7cfc2e00;
/// @dev keccak256(abi.encode(uint256(keccak256("blslight_account_v1.initializable")) - 1)) & ~bytes32(uint256(0xff));
bytes32 internal constant _INITIALIZABLE_STORAGE_POSITION =
0x6c0a8bfcf5680463be0c279c880b798fd37cd3a6f41eba9aa29d36c2fe945100;

struct BLSLightAccountStorage {
address owner;
}

/// @dev keccak256(abi.encode(uint256(keccak256("bls_light_account_v1.storage")) - 1)) & ~bytes32(uint256(0xff));
bytes32 private constant _BLS_STORAGE_POSITION = 0x3e3e3e3e3e3e3e3e3e3e3e3e3e3e3e3e3e3e3e3e3e3e3e3e3e3e3e3e3e3e3e00;
address public _owner;
/// @notice Emitted when this account is first initialized.
/// @param entryPoint The entry point.
/// @param owner The initial owner.
event BLSLightAccountInitialized(IEntryPoint indexed entryPoint, address indexed owner);

event BLSPublicKeySet(bytes32 indexed blsPublicKey);
/// @notice Emitted when this account's owner changes. Also emitted once at initialization, with a
/// `previousOwner` of 0.
/// @param previousOwner The previous owner.
/// @param newOwner The new owner.
event OwnershipTransferred(address indexed previousOwner, address indexed newOwner);

error InvalidBLSSignature();
/// @dev The new owner is not a valid owner (e.g., `address(0)`, the account itself, or the current owner).
error InvalidOwner(address owner);

constructor(IEntryPoint entryPoint_) LightAccount(entryPoint_) {}
constructor(IEntryPoint entryPoint_) CustomSlotInitializable(_INITIALIZABLE_STORAGE_POSITION) {
_ENTRY_POINT = entryPoint_;
// _disableInitializers();
}

/// @notice Initializes the BLSLightAccount with an owner and a BLS public key
/// @param owner_ The initial owner of the account
/// @param blsPublicKey_ The BLS public key for this account
function initialize(address owner_, bytes32 blsPublicKey_) external initializer {
// super._initialize(owner_);
require(_owner == address(0), "Already initialized");
_owner = owner_;
_setBlsPublicKey(blsPublicKey_);
/// @notice Called once as part of initialization, either during initial deployment or when first upgrading to
/// this contract.
/// @dev The `_ENTRY_POINT` member is immutable, to reduce gas consumption. To update the entry point address, a new
/// implementation of LightAccount must be deployed with the new entry point address, and then `upgradeToAndCall`
/// must be called to upgrade the implementation.
/// @param owner_ The initial owner of the account.
function initialize(address owner_) external initializer {
_initialize(owner_);
}

/// @notice Sets or updates the BLS public key for this account
/// @param blsPublicKey_ The new BLS public key
function setBlsPublicKey(bytes32 blsPublicKey_) external onlyAuthorized {
_setBlsPublicKey(blsPublicKey_);
/// @notice Transfers ownership of the contract to a new account (`newOwner`). Can only be called by the current
/// owner or from the entry point via a user operation signed by the current owner.
/// @param newOwner The new owner.
function transferOwnership(address newOwner) external virtual onlyAuthorized {
if (newOwner == address(0) || newOwner == address(this)) {
revert InvalidOwner(newOwner);
}
_transferOwnership(newOwner);
}

/// @notice Return the current owner of this account.
/// @return The current owner.
function owner() public view returns (address) {
return _getStorage().owner;
}

function _initialize(address owner_) internal virtual {
if (owner_ == address(0)) {
revert InvalidOwner(address(0));
}
_getStorage().owner = owner_;
emit BLSLightAccountInitialized(_ENTRY_POINT, owner_);
emit OwnershipTransferred(address(0), owner_);
}

/// @notice Returns the current BLS public key for this account
function blsPublicKey() public view returns (bytes32) {
return _getBLSStorage().blsPublicKey;
function _transferOwnership(address newOwner) internal virtual {
BLSLightAccountStorage storage _storage = _getStorage();
address oldOwner = _storage.owner;
if (newOwner == oldOwner) {
revert InvalidOwner(newOwner);
}
_storage.owner = newOwner;
emit OwnershipTransferred(oldOwner, newOwner);
}

/// @dev Overrides the _validateSignature function to support BLS signatures
/// @dev Implement template method of BaseAccount.
/// Uses a modified version of `SignatureChecker.isValidSignatureNow` in which the digest is wrapped with an
/// "Ethereum Signed Message" envelope for the EOA-owner case but not in the ERC-1271 contract-owner case.
function _validateSignature(PackedUserOperation calldata userOp, bytes32 userOpHash)
internal
virtual
override
returns (uint256 validationData)
{
// if (userOp.signature.length >= 1) {
// uint8 signatureType = uint8(userOp.signature[0]);
// if (signatureType == uint8(SignatureType.BLS)) {
// // BLS signature validation
// return _validateBLSSignature(userOp, userOpHash);
// }
// }
// If not a BLS signature, fall back to the original validation
return super._validateSignature(userOp, userOpHash);
}

/// @dev Validates a BLS signature
function _validateBLSSignature(PackedUserOperation calldata userOp, bytes32 userOpHash)
internal
view
returns (uint256)
{
if (userOp.signature.length < 1) {
revert InvalidSignatureType();
}
uint8 signatureType = uint8(userOp.signature[0]);
if (signatureType == uint8(SignatureType.EOA)) {
// EOA signature
bytes32 signedHash = MessageHashUtils.toEthSignedMessageHash(userOpHash);
bytes32 signedHash = userOpHash.toEthSignedMessageHash();
bytes memory signature = userOp.signature[1:];
return _successToValidationData(_isValidEOAOwnerSignature(signedHash, signature));
} else if (signatureType == uint8(SignatureType.CONTRACT)) {
Expand All @@ -98,15 +139,69 @@ contract BLSLightAccount is LightAccount {
revert InvalidSignatureType();
}

function _setBlsPublicKey(bytes32 blsPublicKey_) internal {
_getBLSStorage().blsPublicKey = blsPublicKey_;
emit BLSPublicKeySet(blsPublicKey_);
/// @notice Check if the signature is a valid by the EOA owner for the given digest.
/// @dev Only supports 65-byte signatures, and uses the digest directly. Reverts if the signature is malformed.
/// @param digest The digest to be checked.
/// @param signature The signature to be checked.
/// @return True if the signature is valid and by the owner, false otherwise.
function _isValidEOAOwnerSignature(bytes32 digest, bytes memory signature) internal view returns (bool) {
address recovered = digest.recover(signature);
return recovered == owner();
}

/// @notice Check if the signature is a valid ERC-1271 signature by a contract owner for the given digest.
/// @param digest The digest to be checked.
/// @param signature The signature to be checked.
/// @return True if the signature is valid and by an owner, false otherwise.
function _isValidContractOwnerSignatureNow(bytes32 digest, bytes memory signature) internal view returns (bool) {
return SignatureChecker.isValidERC1271SignatureNow(owner(), digest, signature);
}

/// @dev The signature is valid if it is signed by the owner's private key (if the owner is an EOA) or if it is a
/// valid ERC-1271 signature from the owner (if the owner is a contract). Reverts if the signature is malformed.
/// Note that unlike the signature validation used in `validateUserOp`, this does **not** wrap the hash in an
/// "Ethereum Signed Message" envelope before checking the signature in the EOA-owner case.
function _isValidSignature(bytes32 replaySafeHash, bytes calldata signature)
internal
view
virtual
override
returns (bool)
{
if (signature.length < 1) {
revert InvalidSignatureType();
}
uint8 signatureType = uint8(signature[0]);
if (signatureType == uint8(SignatureType.EOA)) {
// EOA signature
return _isValidEOAOwnerSignature(replaySafeHash, signature[1:]);
} else if (signatureType == uint8(SignatureType.CONTRACT)) {
// Contract signature without address
return _isValidContractOwnerSignatureNow(replaySafeHash, signature[1:]);
}
revert InvalidSignatureType();
}

function _domainNameAndVersion()
internal
view
virtual
override
returns (string memory name, string memory version)
{
name = "BLSLightAccount";
// Set to the major version of the GitHub release at which the contract was last updated.
version = "1";
}

function _isFromOwner() internal view virtual override returns (bool) {
return msg.sender == owner();
}

function _getBLSStorage() internal pure returns (BLSAccountStorage storage blsStorageStruct) {
bytes32 position = _BLS_STORAGE_POSITION;
function _getStorage() internal pure returns (BLSLightAccountStorage storage storageStruct) {
bytes32 position = _STORAGE_POSITION;
assembly ("memory-safe") {
blsStorageStruct.slot := position
storageStruct.slot := position
}
}
}
}
43 changes: 24 additions & 19 deletions test/BLSLightAccount.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,32 +2,37 @@
pragma solidity ^0.8.0;

import "forge-std/Test.sol";
import {EntryPoint} from "../lib/light-account/lib/account-abstraction/contracts/core/EntryPoint.sol";
import {BLSLightAccount, IEntryPoint} from "../src/BLSLightAccount.sol";
import {LightAccountFactory} from "../lib/light-account/src/LightAccountFactory.sol";
import {LightAccount} from "../lib/light-account/src/LightAccount.sol";
import {EntryPoint} from "light-account/lib/account-abstraction/contracts/core/EntryPoint.sol";
import {BLSLightAccount, IEntryPoint, CustomSlotInitializable} from "../src/BLSLightAccount.sol";

contract BLSLightAccountTest is Test {
address owner = address(this);
IEntryPoint entryPoint;
LightAccountFactory factory;
bytes32 blsPublicKey = keccak256(bytes("blsPublicKey"));
// BLSLightAccount account;
LightAccount public account;
address public eoaAddress;
uint256 public constant EOA_PRIVATE_KEY = 1;
BLSLightAccount account;
EntryPoint entryPoint;

function setUp() public {
eoaAddress = vm.addr(EOA_PRIVATE_KEY);
entryPoint = new EntryPoint();
factory = new LightAccountFactory(address(this), entryPoint);
account = factory.createAccount(eoaAddress, 1);
account = new BLSLightAccount(entryPoint);
}

function testInitialize() public {
assertEq(entryPoint.getNonce(owner, 0), account.getNonce());
assertNotEq(account.owner(), owner);
// account.initialize(owner, blsPublicKey);
// account.setBlsPublicKey(blsPublicKey);
address owner = address(this);

// 初始化合约
account.initialize(owner);

// 检查初始化后的状态
assertEq(account.owner(), owner);
}

function testInitializeRevert() public {
address owner = address(this);

// 第一次初始化合约
account.initialize(owner);

// 尝试再次初始化合约，应该失败
bytes memory errorSelector = abi.encodeWithSelector(CustomSlotInitializable.InvalidInitialization.selector);
vm.expectRevert(errorSelector);
account.initialize(owner);
}
}

0 comments on commit cfe3a9e

Please sign in to comment.