Skip to content

⚡ Optimize find program key by adding static buffers and reusing the same SHA256 while finding an off-curve key #62

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.

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

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
24 changes: 2 additions & 22 deletions src/Solana.Unity.Dex/Orca/Address/AddressUtils.cs
Original file line number Diff line number Diff line change
Expand Up @@ -20,28 +20,8 @@ public static class AddressUtils
/// <returns>A Pda (program-derived address) off the curve, or null.</returns>
public static Pda FindProgramAddress(IEnumerable<byte[]> seeds, PublicKey programId)
{
byte nonce = 255;

while (nonce != 0)
{
PublicKey address;
List<byte[]> seedsWithNonce = new List<byte[]>(seeds);
seedsWithNonce.Add(new byte[] { nonce });

//try to generate the address
bool created = PublicKey.TryCreateProgramAddress(new List<byte[]>(seedsWithNonce), programId, out address);

//if succeeded, return
if (created)
{
return new Pda(address, nonce);
}

//decrease the nonce and retry if failed
nonce--;
}

return null;
return PublicKey.TryFindProgramAddress(seeds, programId, out PublicKey pubkey, out byte bump)
? new Pda(pubkey, bump) : null;
}

/// <summary>
Expand Down
89 changes: 62 additions & 27 deletions src/Solana.Unity.Wallet/PublicKey.cs
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
using System.Linq;
using System.Security.Cryptography;
using System.Text;
using System.Threading;

namespace Solana.Unity.Wallet
{
Expand Down Expand Up @@ -256,44 +257,61 @@
/// <summary>
/// The bytes of the `ProgramDerivedAddress` string.
/// </summary>
private static readonly byte[] ProgramDerivedAddressBytes = Encoding.UTF8.GetBytes("ProgramDerivedAddress");
private static readonly byte[] ProgramDerivedAddressBytes = "ProgramDerivedAddress"u8.ToArray();

/// <summary>
/// A thread-local buffer for all Pda Seeds taht we can reuse
/// </summary>
private static readonly ThreadLocal<List<byte[]>> PdaSeedsBuffer = new(() => []);

Check failure on line 265 in src/Solana.Unity.Wallet/PublicKey.cs

View workflow job for this annotation

GitHub Actions / build

Invalid expression term '['

Check failure on line 265 in src/Solana.Unity.Wallet/PublicKey.cs

View workflow job for this annotation

GitHub Actions / build

Syntax error; value expected

/// <summary>
/// A thread local buffer for the bump array that we can reuse
/// </summary>
private static readonly ThreadLocal<byte[]> BumpArray = new(() => new byte[1]);

/// <summary>
/// A thread local sha256 instance that we can reuse
/// </summary>
private static readonly ThreadLocal<SHA256> Sha256 = new(SHA256.Create);

/// <summary>
/// Derives a program address.
/// </summary>
/// <param name="seeds">The address seeds.</param>
/// <param name="programId">The program Id.</param>
/// <param name="publicKey">The derived public key, returned as inline out.</param>
/// <returns>true if it could derive the program address for the given seeds, otherwise false..</returns>
/// <returns>true if it could derive the program address for the given seeds, otherwise false.</returns>
/// <exception cref="ArgumentException">Throws exception when one of the seeds has an invalid length.</exception>
public static bool TryCreateProgramAddress(ICollection<byte[]> seeds, PublicKey programId, out PublicKey publicKey)
{
MemoryStream buffer = new(PublicKeyLength * seeds.Count + ProgramDerivedAddressBytes.Length + programId.KeyBytes.Length);

SHA256 sha256 = Sha256.Value;
sha256.Initialize();

foreach (byte[] seed in seeds)
{
if (seed.Length > PublicKeyLength)
{
throw new ArgumentException("max seed length exceeded", nameof(seeds));
}
buffer.Write(seed,0, seed.Length);

sha256.TransformBlock(seed,0, seed.Length, null, 0);
}

buffer.Write(programId.KeyBytes, 0, programId.KeyBytes.Length);
buffer.Write(ProgramDerivedAddressBytes, 0, ProgramDerivedAddressBytes.Length);

SHA256 sha256 = SHA256.Create();
byte[] hash = sha256.ComputeHash(new ReadOnlySpan<byte>(buffer.GetBuffer(), 0, (int)buffer.Length).ToArray());

sha256.TransformBlock(programId.KeyBytes,0, programId.KeyBytes.Length, null,0);
sha256.TransformBlock(ProgramDerivedAddressBytes, 0, ProgramDerivedAddressBytes.Length, null, 0);
sha256.TransformFinalBlock([], 0, 0);

Check failure on line 302 in src/Solana.Unity.Wallet/PublicKey.cs

View workflow job for this annotation

GitHub Actions / build

) expected

Check failure on line 302 in src/Solana.Unity.Wallet/PublicKey.cs

View workflow job for this annotation

GitHub Actions / build

Syntax error; value expected

Check failure on line 302 in src/Solana.Unity.Wallet/PublicKey.cs

View workflow job for this annotation

GitHub Actions / build

; expected

Check failure on line 302 in src/Solana.Unity.Wallet/PublicKey.cs

View workflow job for this annotation

GitHub Actions / build

} expected

Check failure on line 302 in src/Solana.Unity.Wallet/PublicKey.cs

View workflow job for this annotation

GitHub Actions / build

; expected

Check failure on line 302 in src/Solana.Unity.Wallet/PublicKey.cs

View workflow job for this annotation

GitHub Actions / build

} expected

Check failure on line 302 in src/Solana.Unity.Wallet/PublicKey.cs

View workflow job for this annotation

GitHub Actions / build

; expected

Check failure on line 302 in src/Solana.Unity.Wallet/PublicKey.cs

View workflow job for this annotation

GitHub Actions / build

} expected

byte[] hash = sha256.Hash!;
if (hash.IsOnCurve())
{
publicKey = null;
return false;
}
publicKey = new(hash);

publicKey = new PublicKey(hash);
return true;
}

/// <summary>
/// Attempts to find a program address for the passed seeds and program Id.
/// </summary>
Expand All @@ -304,28 +322,45 @@
/// <returns>True whenever the address for a nonce was found, otherwise false.</returns>
public static bool TryFindProgramAddress(IEnumerable<byte[]> seeds, PublicKey programId, out PublicKey address, out byte bump)
{
byte seedBump = 255;
List<byte[]> buffer = seeds.ToList();
var bumpArray = new byte[1];
buffer.Add(bumpArray);

while (seedBump != 0)
List<byte[]> pdaSeedsBuffer = PdaSeedsBuffer.Value;
pdaSeedsBuffer.Clear();
pdaSeedsBuffer.AddRange(seeds);

if (pdaSeedsBuffer.Any(seed => seed.Length > PublicKeyLength))
{
bumpArray[0] = seedBump;
bool success = TryCreateProgramAddress(buffer, programId, out PublicKey derivedAddress);
throw new ArgumentException("max seed length exceeded", nameof(seeds));
}

byte[] bumpArray = BumpArray.Value;
SHA256 sha256 = Sha256.Value;

for (bump = 255; ; bump--)
{
sha256.Initialize();

if (success)
foreach (byte[] seed in pdaSeedsBuffer)
{
sha256.TransformBlock(seed, 0, seed.Length, null, 0);
}

bumpArray[0] = bump;
sha256.TransformBlock(bumpArray, 0, 1, null, 0);
sha256.TransformBlock(programId.KeyBytes, 0, programId.KeyBytes.Length, null, 0);
sha256.TransformBlock(ProgramDerivedAddressBytes, 0, ProgramDerivedAddressBytes.Length, null, 0);
sha256.TransformFinalBlock([], 0, 0);

byte[] hash = sha256.Hash!;
if (!hash.IsOnCurve())
{
address = derivedAddress;
bump = seedBump;
address = new PublicKey(hash);
return true;
}

seedBump--;
if (bump == 0)
break;
}

address = null;
bump = 0;
address = null!;
return false;
}

Expand Down
Loading