Skip to content

Commit

Permalink
Merge pull request #297 from LATOKEN/keygen-spam-check
Browse files Browse the repository at this point in the history 
Keygen spam check
  • Loading branch information
@tbssajal
tbssajal authored Jan 13, 2023
2 parents c6d2d6e + c728bad commit e7bb48b
Show file tree
Hide file tree
Showing 7 changed files with 178 additions and 18 deletions.
12 changes: 9 additions & 3 deletions src/Lachain.Consensus/ThresholdKeygen/Data/State.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,11 +12,13 @@ public class State : IEquatable<State>
public Commitment? Commitment { get; set; }
public readonly Fr[] Values;
public readonly bool[] Acks;
public bool Confirmation;

public State(int n)
{
Values = new Fr[n];
Acks = new bool[n];
Confirmation = false;
}

public int ValueCount()
Expand Down Expand Up @@ -54,6 +56,7 @@ public byte[] ToBytes()
stream.Write(Acks.Length.ToBytes().ToArray());
stream.Write(values);
stream.Write(acks);
stream.Write(Confirmation ? new byte[1] {1} : new byte[1] {0});
return stream.ToArray();
}

Expand All @@ -66,15 +69,17 @@ public static State FromBytes(ReadOnlyMemory<byte> bytes)
.Batch(Fr.ByteSize)
.Select(x => Fr.FromBytes(x.ToArray()))
.ToArray();
var acks = bytes.Slice(8 + cLen + n * Fr.ByteSize).ToArray()
var acks = bytes.Slice(8 + cLen + n * Fr.ByteSize, n).ToArray()
.Select(b => b != 0)
.ToArray();
var confirmation = bytes.Slice(8 + cLen + n * Fr.ByteSize + n).ToArray();
var result = new State(values.Length) {Commitment = commitment};
for (var i = 0; i < n; ++i)
{
result.Values[i] = values[i];
result.Acks[i] = acks[i];
}
result.Confirmation = confirmation[0] == 1 ? true : false;

return result;
}
Expand All @@ -85,7 +90,8 @@ public bool Equals(State? other)
if (ReferenceEquals(this, other)) return true;
return Values.SequenceEqual(other.Values) &&
Acks.SequenceEqual(other.Acks) &&
Equals(Commitment, other.Commitment);
Equals(Commitment, other.Commitment) &&
Confirmation == other.Confirmation;
}

public override bool Equals(object? obj)
Expand All @@ -98,7 +104,7 @@ public override bool Equals(object? obj)

public override int GetHashCode()
{
return HashCode.Combine(Values, Acks, Commitment);
return HashCode.Combine(Values, Acks, Commitment, Confirmation);
}
}
}
8 changes: 7 additions & 1 deletion src/Lachain.Consensus/ThresholdKeygen/TrustlessKeygen.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -134,11 +134,17 @@ public bool HandleSendValue(int sender, ValueMessage message)
return true;
}

public bool HandleConfirm(PublicKey tpkeKey, PublicKeySet tsKeys)
public bool HandleConfirm(int sender, PublicKey tpkeKey, PublicKeySet tsKeys)
{
if (_keyGenStates[sender].Confirmation)
{
Logger.LogDebug($"Received confirmation from {sender} more than once");
throw new ArgumentException("Already handled this value");
}
var keyringHash = tpkeKey.ToBytes().Concat(tsKeys.ToBytes()).Keccak();
_confirmations.PutIfAbsent(keyringHash, 0);
_confirmations[keyringHash] += 1;
_keyGenStates[sender].Confirmation = true;
return _confirmations[keyringHash] == Players - Faulty;
}

Expand Down
83 changes: 83 additions & 0 deletions src/Lachain.Core/Blockchain/SystemContracts/GovernanceContract.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ public class GovernanceContract : ISystemContract

private readonly StorageVariable _nextValidators;
private readonly StorageMapping _confirmations;
private readonly StorageMapping? _confirmationsReceived;
private readonly StorageVariable _blockReward;
private readonly StorageVariable _playersCount;
private readonly StorageVariable _tsKeys;
Expand Down Expand Up @@ -110,6 +111,14 @@ public GovernanceContract(InvocationContext context)
new BigInteger(9).ToUInt256()
);
}
if (HardforkHeights.IsHardfork_15Active(context.Snapshot.Blocks.GetTotalBlockHeight()))
{
_confirmationsReceived = new StorageMapping(
ContractRegisterer.GovernanceContract,
context.Snapshot.Storage,
new BigInteger(10).ToUInt256()
);
}
}

public ContractStandard ContractStandard => ContractStandard.GovernanceContract;
Expand Down Expand Up @@ -292,6 +301,23 @@ public ExecutionStatus KeyGenConfirm(UInt256 cycle, byte[] tpkePublicKey, byte[]

frame.ReturnValue = new byte[] { };
frame.UseGas(GasMetering.KeygenConfirmCost);

if (HardforkHeights.IsHardfork_15Active(_context.Snapshot.Blocks.GetTotalBlockHeight()))
{
var senderPubKey = _context.Receipt.RecoverPublicKey(
HardforkHeights.IsHardfork_9Active(_context.Snapshot.Blocks.GetTotalBlockHeight())
);
var nextValidators = _nextValidators.Get()
.Batch(CryptoUtils.PublicKeyLength)
.Select(x => x.ToArray().ToPublicKey())
.ToArray();
if (!nextValidators.Contains(senderPubKey))
{
Logger.LogDebug($"non validator (public key: {senderPubKey.ToHex()}) sent MethodKeygenConfirm tx");
return ExecutionStatus.ExecutionHalted;
}
}

var players = thresholdSignaturePublicKeys.Length;
var faulty = (players - 1) / 3;

Expand All @@ -318,6 +344,13 @@ public ExecutionStatus KeyGenConfirm(UInt256 cycle, byte[] tpkePublicKey, byte[]
}

var gen = GetConsensusGeneration(frame);
var sender = MsgSender();
if (ConfirmationFromPlayer(sender, gen))
{
Logger.LogError($"GovernanceContract is halted in KeyGenConfirm: sender {sender.ToHex()} sent confirmation more than once");
return ExecutionStatus.ExecutionHalted;
}
ConfirmationReceivedFromPlayer(sender, gen);
var votes = GetConfirmations(keyringHash.ToBytes(), gen);
SetConfirmations(keyringHash.ToBytes(), gen, votes + 1);

Expand Down Expand Up @@ -347,6 +380,23 @@ public ExecutionStatus KeyGenConfirmWithVerification(UInt256 cycle, byte[] tpkeP

frame.ReturnValue = new byte[] { };
frame.UseGas(GasMetering.KeygenConfirmCost);

if (HardforkHeights.IsHardfork_15Active(_context.Snapshot.Blocks.GetTotalBlockHeight()))
{
var senderPubKey = _context.Receipt.RecoverPublicKey(
HardforkHeights.IsHardfork_9Active(_context.Snapshot.Blocks.GetTotalBlockHeight())
);
var nextValidators = _nextValidators.Get()
.Batch(CryptoUtils.PublicKeyLength)
.Select(x => x.ToArray().ToPublicKey())
.ToArray();
if (!nextValidators.Contains(senderPubKey))
{
Logger.LogDebug($"non validator (public key: {senderPubKey.ToHex()}) sent MethodKeygenConfirmWithVerification tx");
return ExecutionStatus.ExecutionHalted;
}
}

var players = thresholdSignaturePublicKeys.Length;
var faulty = (players - 1) / 3;

Expand Down Expand Up @@ -397,6 +447,15 @@ public ExecutionStatus KeyGenConfirmWithVerification(UInt256 cycle, byte[] tpkeP
}

var gen = GetConsensusGeneration(frame);
var sender = MsgSender();
if (ConfirmationFromPlayer(sender, gen))
{
Logger.LogError(
$"GovernanceContract is halted in KeyGenConfirmWithVerification: sender {sender.ToHex()} sent confirmation more than once"
);
return ExecutionStatus.ExecutionHalted;
}
ConfirmationReceivedFromPlayer(sender, gen);
var votes = GetConfirmations(keyringHash.ToBytes(), gen);
SetConfirmations(keyringHash.ToBytes(), gen, votes + 1);

Expand All @@ -416,6 +475,13 @@ public ExecutionStatus KeyGenConfirmWithVerification(UInt256 cycle, byte[] tpkeP
public ExecutionStatus FinishCycle(UInt256 cycle, SystemContractExecutionFrame frame)
{
Logger.LogDebug("FinishCycle()");

if (!MsgSender().IsZero())
{
Logger.LogError("!MsgSender().IsZero(): governance function FinishCycle() called by non-zero address");
return ExecutionStatus.ExecutionHalted;
}

var currentBlock = frame.InvocationContext.Receipt.Block;
if (GetBlockNumberInCycle(currentBlock) != 0)
{
Expand Down Expand Up @@ -576,6 +642,23 @@ private List<PublicKey> GetTpkeVerificationKeys()
.Select(x => PublicKey.FromBytes(x)).ToList();
}

private bool ConfirmationFromPlayer(UInt160 sender, ulong cycle)
{
if (_confirmationsReceived is null)
return false;
var value = _confirmationsReceived.GetValue(sender.ToBytes());
if (value.Length == 0 || value.AsReadOnlySpan().ToUInt64() != cycle)
return false;
return value.AsReadOnlySpan().Slice(8)[0] == 1;
}

private void ConfirmationReceivedFromPlayer(UInt160 sender, ulong cycle)
{
if (_confirmationsReceived is null)
return;
_confirmations.SetValue(sender.ToBytes(), cycle.ToBytes().Concat(new byte[1] {1}).ToArray());
}

private int GetConfirmations(IEnumerable<byte> key, ulong gen)
{
var votes = _confirmations.GetValue(key);
Expand Down
44 changes: 41 additions & 3 deletions src/Lachain.Core/Vault/KeyGenManager.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -287,7 +287,7 @@ private void BlockManagerOnSystemContractInvoked(object _, InvocationContext con
return;
}

if (keygen.HandleConfirm(tpkePublicKey, tsKeys))
if (keygen.HandleConfirm(sender, tpkePublicKey, tsKeys))
{
var keys = keygen.TryGetKeys() ?? throw new Exception();
Logger.LogTrace($"Generated keyring with public hash {keys.PublicPartHash().ToHex()}");
Expand Down Expand Up @@ -342,7 +342,7 @@ private void BlockManagerOnSystemContractInvoked(object _, InvocationContext con
return;
}

if (keygen.HandleConfirm(tpkePublicKey, tsKeys))
if (keygen.HandleConfirm(sender, tpkePublicKey, tsKeys))
{
var keys = keygen.TryGetKeys() ?? throw new Exception();
Logger.LogTrace($"Generated keyring with public hash {keys.PublicPartHash().ToHex()}");
Expand Down Expand Up @@ -525,7 +525,45 @@ public bool RescanBlockChainForKeys(IPublicConsensusKeySet publicKeysToSearch)
keygen.Faulty
);

if (!keygen.HandleConfirm(tpkePublicKey, tsKeys)) continue;
if (!keygen.HandleConfirm(sender, tpkePublicKey, tsKeys)) continue;
var keys = keygen.TryGetKeys() ?? throw new Exception();
Logger.LogTrace($"Generated keyring with public hash {keys.PublicPartHash().ToHex()}");
Logger.LogTrace($" - TPKE public key: {keys.TpkePublicKey.ToHex()}");
Logger.LogTrace(
$" - TS public key: {keys.ThresholdSignaturePrivateKey.GetPublicKeyShare().ToHex()}");
Logger.LogTrace(
$" - TS public key set: {string.Join(", ", keys.ThresholdSignaturePublicKeySet.Keys.Select(key => key.ToHex()))}"
);
var lastBlockInCurrentCycle =
(i / StakingContract.CycleDuration + 1) * StakingContract.CycleDuration;
_privateWallet.AddThresholdSignatureKeyAfterBlock(
lastBlockInCurrentCycle, keys.ThresholdSignaturePrivateKey
);
_privateWallet.AddTpkePrivateKeyAfterBlock(lastBlockInCurrentCycle, keys.TpkePrivateKey);
Logger.LogDebug("New keyring saved to wallet");
return true;
}
else if (signature == ContractEncoder.MethodSignatureAsInt(GovernanceInterface.MethodKeygenConfirmWithVerification))
{
Logger.LogDebug(
$"Detected call of GovernanceContract.{GovernanceInterface.MethodKeygenConfirmWithVerification}");
var sender = keygen.GetSenderByPublicKey(tx.RecoverPublicKey(HardforkHeights.IsHardfork_9Active(i)));
if (sender < 0)
{
Logger.LogWarning($"Skipping call because of invalid sender: {sender}");
continue;
}

var args = decoder.Decode(GovernanceInterface.MethodKeygenConfirmWithVerification);
var tpkePublicKey = PublicKey.FromBytes(args[1] as byte[] ?? throw new Exception());
var tsKeys = new PublicKeySet(
(args[2] as byte[][] ?? throw new Exception()).Select(x =>
Crypto.ThresholdSignature.PublicKey.FromBytes(x)
),
keygen.Faulty
);

if (!keygen.HandleConfirm(sender, tpkePublicKey, tsKeys)) continue;
var keys = keygen.TryGetKeys() ?? throw new Exception();
Logger.LogTrace($"Generated keyring with public hash {keys.PublicPartHash().ToHex()}");
Logger.LogTrace($" - TPKE public key: {keys.TpkePublicKey.ToHex()}");
Expand Down
33 changes: 24 additions & 9 deletions test/Lachain.CoreTest/Blockchain/SystemContracts/GovernanceContractTest.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
using System.Reflection;
using Lachain.Consensus.ThresholdKeygen;
using Lachain.Consensus.ThresholdKeygen.Data;
using Lachain.Core.Blockchain.Hardfork;
using Lachain.Core.Blockchain.Interface;
using Lachain.Core.Blockchain.SystemContracts;
using Lachain.Core.Blockchain.SystemContracts.ContractManager;
Expand All @@ -20,6 +21,7 @@
using Lachain.Core.DI.SimpleInjector;
using Lachain.Crypto;
using Lachain.Crypto.ECDSA;
using Lachain.Networking;
using Lachain.Proto;
using Lachain.Storage.State;
using Lachain.Utility.Containers;
Expand All @@ -34,6 +36,7 @@ namespace Lachain.CoreTest.Blockchain.SystemContracts
public class GovernanceContractTest
{
private IContainer? _container;
private IConfigManager? _configManager;
private static readonly ICrypto Crypto = CryptoProvider.GetCrypto();

[SetUp]
Expand All @@ -51,6 +54,18 @@ public void Setup()
containerBuilder.RegisterModule<StorageModule>();

_container = containerBuilder.Build();

_configManager = _container.Resolve<IConfigManager>();

// set chainId from config
if (TransactionUtils.ChainId(false) == 0)
{
var chainId = _configManager.GetConfig<NetworkConfig>("network")?.ChainId;
var newChainId = _configManager.GetConfig<NetworkConfig>("network")?.NewChainId;
TransactionUtils.SetChainId((int)chainId!, (int)newChainId!);
HardforkHeights.SetHardforkHeights(_configManager.GetConfig<HardforkConfig>("hardfork") ?? throw new InvalidOperationException());
StakingContract.Initialize(_configManager.GetConfig<NetworkConfig>("network")!);
}
}

[TearDown]
Expand All @@ -64,14 +79,15 @@ public void Teardown()
[Test]
public void Test_OneNodeCycle()
{

var stateManager = _container?.Resolve<IStateManager>();
var contractRegisterer = _container?.Resolve<IContractRegisterer>();
var tx = new TransactionReceipt();
var keyPair = new EcdsaKeyPair("0xD95D6DB65F3E2223703C5D8E205D98E3E6B470F067B0F94F6C6BF73D4301CE48"
.HexToBytes().ToPrivateKey());
var tx = TestUtils.GetCustomTransactionFromAddress(keyPair, "10", "0.0000000001", 0, false);
var sender = new BigInteger(0).ToUInt160();
var context = new InvocationContext(sender, stateManager!.LastApprovedSnapshot, tx);
var contract = new GovernanceContract(context);
var keyPair = new EcdsaKeyPair("0xD95D6DB65F3E2223703C5D8E205D98E3E6B470F067B0F94F6C6BF73D4301CE48"
.HexToBytes().ToPrivateKey());
byte[] pubKey = CryptoUtils.EncodeCompressed(keyPair.PublicKey);
ECDSAPublicKey[] allKeys = {keyPair.PublicKey};
var keygen = new TrustlessKeygen(keyPair, allKeys, 0, 0);
Expand Down Expand Up @@ -146,7 +162,7 @@ public void Test_OneNodeCycle()
Assert.AreEqual(ExecutionStatus.Ok, contract.KeyGenConfirm(cycle, keyring!.Value.TpkePublicKey.ToBytes(),
keyring!.Value.ThresholdSignaturePublicKeySet.Keys.Select(key => key.ToBytes()).ToArray(), frame));
// set keygen state
Assert.IsTrue(keygen.HandleConfirm(keyring!.Value.TpkePublicKey,
Assert.IsTrue(keygen.HandleConfirm(0, keyring!.Value.TpkePublicKey,
keyring!.Value.ThresholdSignaturePublicKeySet));
}
// check no validators in storage
Expand Down Expand Up @@ -174,7 +190,7 @@ public void Test_InvalidValidatorKey()
{
var stateManager = _container?.Resolve<IStateManager>();
var contractRegisterer = _container?.Resolve<IContractRegisterer>();
var tx = new TransactionReceipt();
var tx = TestUtils.GetRandomTransaction(false);
var sender = new BigInteger(0).ToUInt160();
var context = new InvocationContext(sender, stateManager!.LastApprovedSnapshot, tx);
var contract = new GovernanceContract(context);
Expand Down Expand Up @@ -232,10 +248,10 @@ public void Test_InvalidValidatorKey()
var call = contractRegisterer.DecodeContract(context, ContractRegisterer.GovernanceContract, input);
Assert.IsNotNull(call);
var frame = new SystemContractExecutionFrame(call!, context, input, 100_000_000);
Assert.AreEqual(ExecutionStatus.Ok, contract.KeyGenConfirm(cycle, keyring!.Value.TpkePublicKey.ToBytes(),
Assert.AreEqual(ExecutionStatus.ExecutionHalted, contract.KeyGenConfirm(cycle, keyring!.Value.TpkePublicKey.ToBytes(),
keyring!.Value.ThresholdSignaturePublicKeySet.Keys.Select(key => key.ToBytes()).ToArray(), frame));
// set keygen state
Assert.IsTrue(keygen.HandleConfirm(keyring!.Value.TpkePublicKey,
Assert.IsTrue(keygen.HandleConfirm(0, keyring!.Value.TpkePublicKey,
keyring!.Value.ThresholdSignaturePublicKeySet));
}

Expand All @@ -253,8 +269,7 @@ public void Test_InvalidValidatorKey()
Assert.AreEqual(ExecutionStatus.Ok, contract.FinishCycle(cycle, frame));
}

// check no validators in storage again
Assert.IsEmpty(context.Snapshot.Validators.GetValidatorsPublicKeys());
Assert.Throws<ConsensusStateNotPresentException>(()=>context.Snapshot.Validators.GetValidatorsPublicKeys());
}

private class QueueItem
Expand Down
Loading

0 comments on commit e7bb48b

Please sign in to comment.