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.

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

⚡ Improve Session Wallet #173

Merged
merged 1 commit into from
Sep 8, 2023
Merged
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
56 changes: 33 additions & 23 deletions Runtime/codebase/SessionWallet.cs
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ public class SessionWallet : InGameWallet
public PublicKey SessionTokenPDA { get; protected set; }

public static SessionWallet Instance;
private static WalletBase _externalWallet;

private SessionWallet(RpcCluster rpcCluster = RpcCluster.DevNet,
string customRpcUri = null, string customStreamingRpcUri = null,
Expand Down Expand Up @@ -58,25 +59,28 @@ private static PublicKey FindSessionToken(PublicKey TargetProgram, Account Accou
sessionSigner: Account.PublicKey
);
}


public void SignInitSessionTx(Transaction tx)
{
tx.PartialSign(new[] { _externalWallet.Account, Account });
}


/// <summary>
/// Creates a new SessionWallet instance and logs in with the provided password if a session wallet exists, otherwise creates a new account and logs in.
/// </summary>
/// <param name="targetProgram">The target program to interact with.</param>
/// <param name="password">The password to decrypt the session keystore.</param>
/// <param name="rpcCluster">The Solana RPC cluster to connect to.</param>
/// <param name="customRpcUri">A custom URI to connect to the Solana RPC cluster.</param>
/// <param name="customStreamingRpcUri">A custom URI to connect to the Solana streaming RPC cluster.</param>
/// <param name="autoConnectOnStartup">Whether to automatically connect to the Solana RPC cluster on startup.</param>
/// <param name="externalWallet">The external wallet</param>
/// <returns>A SessionWallet instance.</returns>
public static async Task<SessionWallet> GetSessionWallet(PublicKey targetProgram, string password, RpcCluster rpcCluster = RpcCluster.DevNet,
string customRpcUri = null, string customStreamingRpcUri = null,
bool autoConnectOnStartup = false)
public static async Task<SessionWallet> GetSessionWallet(PublicKey targetProgram, string password, WalletBase externalWallet = null)
{
if(Instance != null) return Instance;
SessionWallet sessionWallet = new SessionWallet(rpcCluster, customRpcUri, customStreamingRpcUri, autoConnectOnStartup);
externalWallet ??= Web3.Wallet;
_externalWallet = externalWallet;
SessionWallet sessionWallet = new SessionWallet(externalWallet.RpcCluster, externalWallet.ActiveRpcClient.NodeAddress.ToString());
sessionWallet.TargetProgram = targetProgram;
sessionWallet.EncryptedKeystoreKey = $"{Web3.Account.PublicKey}_SessionKeyStore";
sessionWallet.EncryptedKeystoreKey = $"{_externalWallet.Account.PublicKey}_SessionKeyStore";
var derivedPassword = DeriveSessionPassword(password);

if (sessionWallet.HasSessionWallet())
Expand All @@ -89,10 +93,10 @@ public static async Task<SessionWallet> GetSessionWallet(PublicKey targetProgram
sessionWallet.DeleteSessionWallet();
sessionWallet.Logout();
Instance = null;
return await GetSessionWallet(targetProgram, password, rpcCluster, customRpcUri, customStreamingRpcUri, autoConnectOnStartup);
return await GetSessionWallet(targetProgram, password, externalWallet);
}

sessionWallet.SessionTokenPDA = FindSessionToken(targetProgram, sessionWallet.Account, Web3.Account);
sessionWallet.SessionTokenPDA = FindSessionToken(targetProgram, sessionWallet.Account, _externalWallet.Account);

Debug.Log(sessionWallet.SessionTokenPDA);

Expand All @@ -109,14 +113,14 @@ public static async Task<SessionWallet> GetSessionWallet(PublicKey targetProgram
}

Debug.Log("Session Token is invalid");
await sessionWallet.PrepareLogout();
await sessionWallet.CloseSession();
sessionWallet.Logout();
Instance = null;
return await GetSessionWallet(targetProgram, password, rpcCluster, customRpcUri, customStreamingRpcUri, autoConnectOnStartup);
return await GetSessionWallet(targetProgram, password, externalWallet);
}

sessionWallet.Account = await sessionWallet.CreateAccount(password: derivedPassword);
sessionWallet.SessionTokenPDA = FindSessionToken(targetProgram, sessionWallet.Account, Web3.Account);
sessionWallet.SessionTokenPDA = FindSessionToken(targetProgram, sessionWallet.Account, _externalWallet.Account);
return sessionWallet;
}

Expand All @@ -133,7 +137,7 @@ public TransactionInstruction CreateSessionIX(bool topUp, long sessionValidity)
{
SessionToken = SessionTokenPDA,
SessionSigner = Account.PublicKey,
Authority = Web3.Account,
Authority = _externalWallet.Account,
TargetProgram = TargetProgram,
SystemProgram = SystemProgram.ProgramIdKey,
};
Expand All @@ -155,7 +159,7 @@ public TransactionInstruction RevokeSessionIX()
{
SessionToken = SessionTokenPDA,
// Only the authority of the session token can receive the refund
Authority = Web3.Account,
Authority = _externalWallet.Account,
SystemProgram = SystemProgram.ProgramIdKey,
};

Expand Down Expand Up @@ -185,8 +189,15 @@ public async Task<bool> IsSessionTokenValid()
return SessionToken.Deserialize(Convert.FromBase64String(sessionTokenData)).ValidUntil > DateTimeOffset.UtcNow.ToUnixTimeSeconds();
}

public async Task<PublicKey> Authority()
{
var sessionTokenData = (await ActiveRpcClient.GetAccountInfoAsync(SessionTokenPDA)).Result.Value.Data[0];
if (sessionTokenData == null) return null;
return SessionToken.Deserialize(Convert.FromBase64String(sessionTokenData)).Authority;
}

private static string DeriveSessionPassword(string password) {
var rawData = Web3.Account.PublicKey.Key + password + Application.platform;
var rawData = _externalWallet.Account.PublicKey.Key + password + Application.platform;
using SHA256 sha256Hash = SHA256.Create();
// ComputeHash - returns byte array
var bytes = sha256Hash.ComputeHash(Encoding.UTF8.GetBytes(rawData));
Expand All @@ -207,15 +218,15 @@ private void DeleteSessionWallet()
/// NOTE: You must call PrepareLogout before calling Logout to ensure that the session token account is revoked and the refund is issued.
/// </summary>
/// <returns>A task that represents the asynchronous operation.</returns>
public async Task PrepareLogout()
public async Task CloseSession(Commitment commitment = Commitment.Confirmed)
{
Debug.Log("Preparing Logout");
// Revoke Session
var tx = new Transaction()
{
FeePayer = Account,
Instructions = new List<TransactionInstruction>(),
RecentBlockHash = await Web3.BlockHash()
RecentBlockHash = await GetBlockHash(commitment)
};

// Get balance and calculate refund
Expand All @@ -226,9 +237,8 @@ public async Task PrepareLogout()

tx.Add(RevokeSessionIX());
// Issue Refund
tx.Add(SystemProgram.Transfer(Account.PublicKey, Web3.Account.PublicKey, (ulong)refund));
var rest = await SignAndSendTransaction(tx);
Debug.Log("Session refund transaction: " + rest.RawRpcResponse);
tx.Add(SystemProgram.Transfer(Account.PublicKey, _externalWallet.Account.PublicKey, (ulong)refund));
var rest = await SignAndSendTransaction(tx, commitment: commitment);
DeleteSessionWallet();
}
}
Expand Down