Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Create KeyBuilder for native contracts #1748

Merged
merged 1 commit into from
Jul 4, 2020
Merged

Create KeyBuilder for native contracts #1748

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 52 additions & 0 deletions src/neo/SmartContract/Native/KeyBuilder.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
using Neo.IO;
using Neo.Ledger;
using System;
using System.IO;

namespace Neo.SmartContract.Native
{
internal class KeyBuilder
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
internal class KeyBuilder
internal class KeyBuilder : IDisposable

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't need IDisposable.

{
private readonly int id;
private readonly MemoryStream stream = new MemoryStream();

public KeyBuilder(int id, byte prefix)
{
this.id = id;
this.stream.WriteByte(prefix);
}

public KeyBuilder Add(ReadOnlySpan<byte> key)
{
stream.Write(key);
return this;
}

public KeyBuilder Add(ISerializable key)
{
using (BinaryWriter writer = new BinaryWriter(stream, Utility.StrictUTF8, true))
{
key.Serialize(writer);
writer.Flush();
}
return this;
}

unsafe public KeyBuilder Add<T>(T key) where T : unmanaged
{
return Add(new ReadOnlySpan<byte>(&key, sizeof(T)));
}

public static implicit operator StorageKey(KeyBuilder builder)
{
using (builder.stream)
{
return new StorageKey
{
Id = builder.id,
Key = builder.stream.ToArray()
};
}
}
}
}
17 changes: 2 additions & 15 deletions src/neo/SmartContract/Native/NativeContract.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
using Neo.IO;
using Neo.Ledger;
using Neo.SmartContract.Manifest;
using Neo.SmartContract.Native.Tokens;
using Neo.VM;
Expand Down Expand Up @@ -79,21 +78,9 @@ protected NativeContract()
contractsHashDictionary.Add(Hash, this);
}

protected StorageKey CreateStorageKey(byte prefix, byte[] key = null)
private protected KeyBuilder CreateStorageKey(byte prefix)
{
StorageKey storageKey = new StorageKey
{
Id = Id,
Key = new byte[sizeof(byte) + (key?.Length ?? 0)]
};
storageKey.Key[0] = prefix;
key?.CopyTo(storageKey.Key.AsSpan(1));
return storageKey;
}

internal protected StorageKey CreateStorageKey(byte prefix, ISerializable key)
{
return CreateStorageKey(prefix, key.ToArray());
return new KeyBuilder(Id, prefix);
}

public static NativeContract GetContract(UInt160 hash)
Expand Down
14 changes: 7 additions & 7 deletions src/neo/SmartContract/Native/Tokens/NeoToken.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ protected override void OnBalanceChanging(ApplicationEngine engine, UInt160 acco
if (amount.IsZero) return;
if (state.VoteTo != null)
{
StorageItem storage_validator = engine.Snapshot.Storages.GetAndChange(CreateStorageKey(Prefix_Candidate, state.VoteTo.ToArray()));
StorageItem storage_validator = engine.Snapshot.Storages.GetAndChange(CreateStorageKey(Prefix_Candidate).Add(state.VoteTo));
CandidateState state_validator = storage_validator.GetInteroperable<CandidateState>();
state_validator.Votes += amount;
}
Expand Down Expand Up @@ -112,7 +112,7 @@ protected override void OnPersist(ApplicationEngine engine)
[ContractMethod(0_03000000, CallFlags.AllowStates)]
public BigInteger UnclaimedGas(StoreView snapshot, UInt160 account, uint end)
{
StorageItem storage = snapshot.Storages.TryGet(CreateAccountKey(account));
StorageItem storage = snapshot.Storages.TryGet(CreateStorageKey(Prefix_Account).Add(account));
if (storage is null) return BigInteger.Zero;
NeoAccountState state = storage.GetInteroperable<NeoAccountState>();
return CalculateBonus(state.Balance, state.BalanceHeight, end);
Expand All @@ -129,7 +129,7 @@ private bool RegisterCandidate(ApplicationEngine engine, ECPoint pubkey)

private void RegisterCandidateInternal(StoreView snapshot, ECPoint pubkey)
{
StorageKey key = CreateStorageKey(Prefix_Candidate, pubkey);
StorageKey key = CreateStorageKey(Prefix_Candidate).Add(pubkey);
StorageItem item = snapshot.Storages.GetAndChange(key, () => new StorageItem(new CandidateState()));
CandidateState state = item.GetInteroperable<CandidateState>();
state.Registered = true;
Expand All @@ -140,7 +140,7 @@ private bool UnregisterCandidate(ApplicationEngine engine, ECPoint pubkey)
{
if (!engine.CheckWitnessInternal(Contract.CreateSignatureRedeemScript(pubkey).ToScriptHash()))
return false;
StorageKey key = CreateStorageKey(Prefix_Candidate, pubkey);
StorageKey key = CreateStorageKey(Prefix_Candidate).Add(pubkey);
if (engine.Snapshot.Storages.TryGet(key) is null) return true;
StorageItem item = engine.Snapshot.Storages.GetAndChange(key);
CandidateState state = item.GetInteroperable<CandidateState>();
Expand All @@ -160,13 +160,13 @@ private bool Vote(ApplicationEngine engine, UInt160 account, ECPoint voteTo)

private bool VoteInternal(StoreView snapshot, UInt160 account, ECPoint voteTo)
{
StorageKey key_account = CreateAccountKey(account);
StorageKey key_account = CreateStorageKey(Prefix_Account).Add(account);
if (snapshot.Storages.TryGet(key_account) is null) return false;
StorageItem storage_account = snapshot.Storages.GetAndChange(key_account);
NeoAccountState state_account = storage_account.GetInteroperable<NeoAccountState>();
if (state_account.VoteTo != null)
{
StorageKey key = CreateStorageKey(Prefix_Candidate, state_account.VoteTo.ToArray());
StorageKey key = CreateStorageKey(Prefix_Candidate).Add(state_account.VoteTo);
StorageItem storage_validator = snapshot.Storages.GetAndChange(key);
CandidateState state_validator = storage_validator.GetInteroperable<CandidateState>();
state_validator.Votes -= state_account.Balance;
Expand All @@ -176,7 +176,7 @@ private bool VoteInternal(StoreView snapshot, UInt160 account, ECPoint voteTo)
state_account.VoteTo = voteTo;
if (voteTo != null)
{
StorageKey key = CreateStorageKey(Prefix_Candidate, voteTo.ToArray());
StorageKey key = CreateStorageKey(Prefix_Candidate).Add(voteTo);
if (snapshot.Storages.TryGet(key) is null) return false;
StorageItem storage_validator = snapshot.Storages.GetAndChange(key);
CandidateState state_validator = storage_validator.GetInteroperable<CandidateState>();
Expand Down
15 changes: 5 additions & 10 deletions src/neo/SmartContract/Native/Tokens/Nep5Token.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,16 +58,11 @@ protected Nep5Token()
Manifest.Abi.Events = events.ToArray();
}

protected StorageKey CreateAccountKey(UInt160 account)
{
return CreateStorageKey(Prefix_Account, account);
}

internal protected virtual void Mint(ApplicationEngine engine, UInt160 account, BigInteger amount)
{
if (amount.Sign < 0) throw new ArgumentOutOfRangeException(nameof(amount));
if (amount.IsZero) return;
StorageItem storage = engine.Snapshot.Storages.GetAndChange(CreateAccountKey(account), () => new StorageItem(new TState()));
StorageItem storage = engine.Snapshot.Storages.GetAndChange(CreateStorageKey(Prefix_Account).Add(account), () => new StorageItem(new TState()));
TState state = storage.GetInteroperable<TState>();
OnBalanceChanging(engine, account, state, amount);
state.Balance += amount;
Expand All @@ -85,7 +80,7 @@ internal protected virtual void Burn(ApplicationEngine engine, UInt160 account,
{
if (amount.Sign < 0) throw new ArgumentOutOfRangeException(nameof(amount));
if (amount.IsZero) return;
StorageKey key = CreateAccountKey(account);
StorageKey key = CreateStorageKey(Prefix_Account).Add(account);
StorageItem storage = engine.Snapshot.Storages.GetAndChange(key);
TState state = storage.GetInteroperable<TState>();
if (state.Balance < amount) throw new InvalidOperationException();
Expand All @@ -112,7 +107,7 @@ public virtual BigInteger TotalSupply(StoreView snapshot)
[ContractMethod(0_01000000, CallFlags.AllowStates)]
public virtual BigInteger BalanceOf(StoreView snapshot, UInt160 account)
{
StorageItem storage = snapshot.Storages.TryGet(CreateAccountKey(account));
StorageItem storage = snapshot.Storages.TryGet(CreateStorageKey(Prefix_Account).Add(account));
if (storage is null) return BigInteger.Zero;
return storage.GetInteroperable<TState>().Balance;
}
Expand All @@ -125,7 +120,7 @@ protected virtual bool Transfer(ApplicationEngine engine, UInt160 from, UInt160
return false;
ContractState contract_to = engine.Snapshot.Contracts.TryGet(to);
if (contract_to?.Payable == false) return false;
StorageKey key_from = CreateAccountKey(from);
StorageKey key_from = CreateStorageKey(Prefix_Account).Add(from);
StorageItem storage_from = engine.Snapshot.Storages.GetAndChange(key_from);
if (amount.IsZero)
{
Expand All @@ -151,7 +146,7 @@ protected virtual bool Transfer(ApplicationEngine engine, UInt160 from, UInt160
engine.Snapshot.Storages.Delete(key_from);
else
state_from.Balance -= amount;
StorageKey key_to = CreateAccountKey(to);
StorageKey key_to = CreateStorageKey(Prefix_Account).Add(to);
StorageItem storage_to = engine.Snapshot.Storages.GetAndChange(key_to, () => new StorageItem(new TState()));
TState state_to = storage_to.GetInteroperable<TState>();
OnBalanceChanging(engine, to, state_to, amount);
Expand Down
10 changes: 10 additions & 0 deletions tests/neo.UnitTests/TestUtils.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
using Neo.Network.P2P.Payloads;
using Neo.SmartContract;
using Neo.SmartContract.Manifest;
using Neo.SmartContract.Native;
using Neo.VM;
using Neo.Wallets.NEP6;
using System;
Expand Down Expand Up @@ -55,6 +56,15 @@ public static ContractManifest CreateManifest(UInt160 hash, string method, Contr
return manifest;
}

public static StorageKey CreateStorageKey(this NativeContract contract, byte prefix, ISerializable key)
{
return new StorageKey
{
Id = contract.Id,
Key = key.ToArray().Prepend(prefix).ToArray()
};
}

public static byte[] GetByteArray(int length, byte firstByte)
{
byte[] array = new byte[length];
Expand Down