1
pragma solidity 0.8.20;
2

3
import {IPuzzle} from "lib/IPuzzle.sol";
4

5
import {IERC1822Proxiable} from "lib/openzeppelin-contracts/contracts/interfaces/draft-IERC1822.sol";
6
import {ERC1967Proxy} from "lib/openzeppelin-contracts/contracts/proxy/ERC1967/ERC1967Proxy.sol";
7
import {ERC1967Utils} from "lib/openzeppelin-contracts/contracts/proxy/ERC1967/ERC1967Utils.sol";
8
import {UUPSUpgradeable} from "lib/openzeppelin-contracts/contracts/proxy/utils/UUPSUpgradeable.sol";
9
import {OwnableUpgradeable} from "lib/openzeppelin-contracts-upgradeable/contracts/access/OwnableUpgradeable.sol";
10

11
contract SafeCurta is IPuzzle {
12
    mapping(uint => SafeChallenge) public factories;
13

14
    function name() external pure returns (string memory){
15
        return "ZSafe";
16
    }
17

18
    function generate(address _seed) public returns (uint256){
19
        return uint256(keccak256(abi.encode("Can you unlock the safe?", _seed)));
20
    }
21

22
    function verify(uint256 _start, uint256) external returns (bool) {
23
        return factories[_start].isUnlocked();
24
    }
25

26
    function deploy(uint256 _start, address owner) external returns (address) {
27
        bytes32 rng_seed = keccak256(abi.encodePacked(_start));
28
        factories[_start] = new SafeChallenge(owner, rng_seed);
29
        return address(factories[_start]);
30
    }
31

32
}
33

34
contract SafeChallenge {
35
    bytes32 public seed;
36
    SafeProxy public proxy;
37
    bool public isUnlocked;
38

39
    constructor(address owner, bytes32 _seed){
40
        //init both
41
        SafeProxy impl1 = new SafeSecret(); 
42
        SafeProxy impl2 = new SafeSecretAdmin(); 
43

44
        bytes32[] memory whitelist = new bytes32[](2);
45
        whitelist[0] = address(impl1).codehash;
46
        whitelist[1] = address(impl2).codehash;
47

48
        bytes memory init_data = abi.encodeCall(impl1.initialize, (owner, whitelist));
49

50
        address proxy_impl = address(new ERC1967Proxy(address(impl1), init_data));
51

52
        proxy = SafeProxy(proxy_impl);
53
        seed = _seed;
54
        isUnlocked = false;
55
    }
56

57

58
    function unlock(bytes32[3] calldata r, bytes32[3] calldata s) external {
59
        for(uint i = 0; i < 2; ++i){
60
            require(uint(r[i]) < uint(r[i+1]));
61
        }
62
        
63
        for(uint i = 0; i < 3; ++i){
64
            check(r[i], s[i]);
65
        }
66
        
67
        isUnlocked = true;
68
    }
69

70
    function check(bytes32 _r, bytes32 _s) internal {
71
        uint8 v = 27;
72
        address owner = proxy.owner();        
73

74
        //--------
75

76
        bytes32 message1_hash = keccak256(abi.encodePacked(seed, address(0xdead)));
77
        bytes32 r1 = transform_r1(_r);
78
        bytes32 s1 = transform_s1(_s);
79

80
        address signer = ecrecover(message1_hash, v, r1, s1);
81
        require(signer != address(0), "no sig match :<");
82
        require(signer == owner, "no owner match :<");
83

84
        //---------
85

86
        bytes32 message2_hash = keccak256(abi.encodePacked(seed, address(0xbeef)));
87
        bytes32 r2 = transform_r2(_r);
88
        bytes32 s2 = transform_s2(_s);
89

90
        address signer2 = ecrecover(message2_hash, v, r2, s2);
91
        require(signer2 != address(0), "no sig match :<");
92
        require(signer2 == owner, "no owner match :<");
93

94
        //--------
95

96
    }
97

98
    function transform_r1(bytes32 r) internal pure returns (bytes32) {
99
        return r;
100
    }
101

102
    function transform_s1(bytes32 s) internal view returns (bytes32) {
103
        return bytes32(uint256(s) ^ proxy.p2());
104
    }
105

106
    function transform_r2(bytes32 r) internal view returns (bytes32) {
107
        unchecked{
108
            return bytes32(uint256(r) + proxy.p1()); 
109
        }
110
    }
111

112
    function transform_s2(bytes32 s) internal view returns (bytes32) {
113
        return keccak256(abi.encodePacked(uint256(s) ^ proxy.p2(), seed));
114
    }
115
}
116

117
//Butchered implementation from https://github.com/OpenZeppelin/openzeppelin-contracts/blob/v5.0.0/contracts/proxy/utils/UUPSUpgradeable.sol with only the features I need
118
abstract contract SafeUpgradeable {
119
    mapping(bytes32 => bool) internal whitelist;
120

121
    address private immutable __self = address(this);
122

123
    modifier onlyProxy() {
124
        _checkProxy();
125
        _;
126
    }
127

128
    function upgradeToAndCall(address newImplementation, bytes memory data) public payable virtual onlyProxy {
129
        _authorizeUpgrade(newImplementation);
130
        _upgradeToAndCallSafe(newImplementation, data);
131
    }
132

133
    function _checkProxy() internal view virtual {
134
        if (
135
            address(this) == __self || 
136
            ERC1967Utils.getImplementation() != __self 
137
        ) {
138
            revert("No hacc");
139
        }
140
    }
141

142
    function _authorizeUpgrade(address newImplementation) internal {
143
        require(whitelist[newImplementation.codehash], "wtf no whitelisted no hacc pls");
144
    }
145

146
    function _upgradeToAndCallSafe(address newImplementation, bytes memory data) private {
147
            ERC1967Utils.upgradeToAndCall(newImplementation, data);
148
    }
149
}
150

151

152
abstract contract SafeProxy is OwnableUpgradeable, SafeUpgradeable {
153
    uint256 internal p1_secret;
154
    uint256 internal p2_secret;
155

156
    function initialize(address owner, bytes32[] calldata whitelisted_hashes) public initializer{
157

158
        for(uint i = 0; i < whitelisted_hashes.length; ++i){
159
            whitelist[whitelisted_hashes[i]] = true;
160
        }
161

162
        p1_secret = uint256(keccak256(abi.encodePacked(keccak256(abi.encode(uint256(blockhash(block.number)))))));
163
        p2_secret = uint256(keccak256(abi.encodePacked(keccak256(abi.encode(p1_secret)))));
164

165
        __Ownable_init(owner);
166
    }
167

168
    function p1() external view virtual returns (uint256);
169
    function p2() external view virtual returns (uint256);
170
}
171

172

173
contract SafeSecret is SafeProxy {
174
    function p1() external view virtual override returns (uint256){
175
        return p1_secret;
176
    }
177

178
    function p2() external view virtual override returns (uint256){
179
        return p2_secret;
180
    }
181
}
182

183
contract SafeSecretAdmin is SafeProxy {
184
    uint256 private offsetp1;
185
    uint256 private offsetp2;
186

187
    function p1() external view virtual override returns (uint256){
188
        unchecked{
189
            return p1_secret+offsetp1; 
190
        }
191
    }
192

193
    function p2() external view virtual override returns (uint256){
194
        unchecked{
195
            return p2_secret+offsetp2;
196
        }
197
    }
198

199
    function set_offset(uint256 _p1, uint256 _p2) external {
200
        offsetp1 = _p1;
201
        offsetp2 = _p2;
202
    }
203
}
204