Rewrite local association scenario for cleaner error handling and cle…#260
Open
Rewrite local association scenario for cleaner error handling and cle…#260
Conversation
WalkthroughLocalAssociationScenario was rewritten to implement IDisposable and to provide an asynchronous StartAndExecute flow with cancellation, WebSocket/backoff/key-exchange handling and cleanup. Three call sites now use Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 Pre-merge checks | ✅ 1 | ❌ 3❌ Failed checks (2 warnings, 1 inconclusive)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Tip Issue Planner is now in beta. Read the docs and try it out! Share your feedback on Discord. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Fix all issues with AI agents
In `@Runtime/codebase/SolanaMobileStack/LocalAssociationScenario.cs`:
- Around line 176-184: ConnectWithBackoffAsync can hang because it awaits
_wsConnected.Task but OnWsError only logs; update the WebSocket error handler
(OnWsError) to signal the waiting task (use _wsConnected.TrySetException with
the exception or _wsConnected.TrySetResult(false)) so ConnectWithBackoffAsync's
await unblocks and the retry/backoff logic proceeds; also ensure the close
handler (OnWsClose/OnClose) uses _wsConnected.TrySetResult(true/false)
consistently and guard against double-sets by using TrySet... methods when
touching _wsConnected to avoid race conditions.
- Around line 53-57: You created a linked CancellationTokenSource with
CancelAfter(_overallTimeout) but still assign and use the original ct, so the
overall timeout never cancels ConnectWithBackoffAsync or the action loop; change
the assignment to use the linked CTS token (set _cancellationToken = cts.Token)
and ensure all cancellation checks and ThrowIfCancellationRequested() calls use
_cancellationToken (and not the original ct); apply the same fix to the other
linked-CTS block (the one around lines with CreateLinkedTokenSource/CancelAfter
at the second site) so the overall timeout is honored across
ConnectWithBackoffAsync and the action loop.
- Around line 138-143: The Debug.Log in StartActivityForAssociation is emitting
sensitive association tokens (and there are similar logs around lines 250-257) —
remove or redact raw tokens/decrypted payloads and instead log only
non-sensitive metadata (e.g., port, token length, a stable prefix/suffix, or a
short hash). Update StartActivityForAssociation and the other logging sites to
stop printing associationToken or decrypted payloads; if a trace is required,
compute and log a one-way hash (SHA-256) or mask all but 4 characters of the
token and include contextual text like "Launched intent for port {port},
token_hash={hash}" or "token_masked=xxxx". Ensure no plaintext tokens or
payloads are written to Debug.Log anywhere in the class.
Comment on lines
+53
to
+57
| using var cts = CancellationTokenSource.CreateLinkedTokenSource(ct); | ||
| cts.CancelAfter(_overallTimeout); | ||
|
|
||
| _cancellationToken = ct; | ||
| _tcs = new TaskCompletionSource<Response<object>>(); |
There was a problem hiding this comment.
Overall timeout isn’t enforced across the connection/action flow.
You create a linked CTS with CancelAfter, but _cancellationToken and ThrowIfCancellationRequested use the original ct, so the 30s timeout won’t cancel ConnectWithBackoffAsync or the action loop when no external token is provided.
✅ Minimal fix to honor the overall timeout
- _cancellationToken = ct;
+ _cancellationToken = cts.Token;
_tcs = new TaskCompletionSource<Response<object>>();
@@
- ct.ThrowIfCancellationRequested();
+ _cancellationToken.ThrowIfCancellationRequested();Also applies to: 94-97
🤖 Prompt for AI Agents
In `@Runtime/codebase/SolanaMobileStack/LocalAssociationScenario.cs` around lines
53 - 57, You created a linked CancellationTokenSource with
CancelAfter(_overallTimeout) but still assign and use the original ct, so the
overall timeout never cancels ConnectWithBackoffAsync or the action loop; change
the assignment to use the linked CTS token (set _cancellationToken = cts.Token)
and ensure all cancellation checks and ThrowIfCancellationRequested() calls use
_cancellationToken (and not the original ct); apply the same fix to the other
linked-CTS block (the one around lines with CreateLinkedTokenSource/CancelAfter
at the second site) so the overall timeout is honored across
ConnectWithBackoffAsync and the action loop.
Comment on lines
+138
to
143
| private void StartActivityForAssociation(string associationToken, int port) | ||
| { | ||
| var intent = LocalAssociationIntentCreator.CreateAssociationIntent(associationToken, port); | ||
| _currentActivity.Call("startActivityForResult", intent, 0); | ||
| _currentActivity.Call("runOnUiThread", new AndroidJavaRunnable(TryConnectWs)); | ||
| _startAssociationTaskCompletionSource = new TaskCompletionSource<Response<object>>(); | ||
| return _startAssociationTaskCompletionSource.Task; | ||
| Debug.Log($"[MWA] Launched intent for port {port}, token {associationToken}"); | ||
| } |
There was a problem hiding this comment.
Avoid logging association tokens and decrypted payloads.
These logs can expose sensitive wallet association data and signed payload contents in device logs. Consider redaction or logging only metadata.
🔒️ Safer logging
- Debug.Log($"[MWA] Launched intent for port {port}, token {associationToken}");
+ Debug.Log($"[MWA] Launched intent for port {port}");
@@
- Debug.Log($"[MWA] Received: {json}");
+ Debug.Log($"[MWA] Received response (len={json.Length})");Also applies to: 250-257
🤖 Prompt for AI Agents
In `@Runtime/codebase/SolanaMobileStack/LocalAssociationScenario.cs` around lines
138 - 143, The Debug.Log in StartActivityForAssociation is emitting sensitive
association tokens (and there are similar logs around lines 250-257) — remove or
redact raw tokens/decrypted payloads and instead log only non-sensitive metadata
(e.g., port, token length, a stable prefix/suffix, or a short hash). Update
StartActivityForAssociation and the other logging sites to stop printing
associationToken or decrypted payloads; if a trace is required, compute and log
a one-way hash (SHA-256) or mask all but 4 characters of the token and include
contextual text like "Launched intent for port {port}, token_hash={hash}" or
"token_masked=xxxx". Ensure no plaintext tokens or payloads are written to
Debug.Log anywhere in the class.
Comment on lines
+176
to
+184
| var startTime = DateTime.UtcNow; | ||
| _wsConnected = new TaskCompletionSource<bool>(); | ||
|
|
||
| attempt++; | ||
| Debug.Log($"[MWA] Connect attempt {attempt}, state: {_webSocket.State}"); | ||
| _webSocket.Connect(); | ||
|
|
||
| var success = await _wsConnected.Task; | ||
| Debug.Log($"[MWA] Connect attempt {attempt} result, state: {_webSocket.State}"); |
There was a problem hiding this comment.
Connect loop can hang if OnError fires without OnClose.
ConnectWithBackoffAsync awaits _wsConnected.Task, but OnWsError only logs. If the socket errors without a close, the await never completes and the retry loop stalls.
🛠️ Signal connection failure on error
- private static void OnWsError(string message)
+ private void OnWsError(string message)
{
Debug.Log($"[MWA] WS Error: {message}");
+ if (_isConnecting)
+ _wsConnected?.TrySetResult(false);
+ else
+ _tcs?.TrySetException(new Exception($"WS error: {message}"));
}Also applies to: 229-232
🤖 Prompt for AI Agents
In `@Runtime/codebase/SolanaMobileStack/LocalAssociationScenario.cs` around lines
176 - 184, ConnectWithBackoffAsync can hang because it awaits _wsConnected.Task
but OnWsError only logs; update the WebSocket error handler (OnWsError) to
signal the waiting task (use _wsConnected.TrySetException with the exception or
_wsConnected.TrySetResult(false)) so ConnectWithBackoffAsync's await unblocks
and the retry/backoff logic proceeds; also ensure the close handler
(OnWsClose/OnClose) uses _wsConnected.TrySetResult(true/false) consistently and
guard against double-sets by using TrySet... methods when touching _wsConnected
to avoid race conditions.
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In `@Runtime/codebase/SessionWallet.cs`:
- Around line 69-85: GetSessionWallet(string publicKey, string privateKey,
PublicKey targetProgram) lacks the singleton guard and null checks: add the same
early return used in the other overloads (if (Instance != null) return
Instance;) before constructing a new SessionWallet, and validate Web3.Wallet and
its ActiveRpcClient (e.g., if (Web3.Wallet == null) throw/new error or handle;
if (_externalWallet.ActiveRpcClient == null) get node address another way or
throw) to avoid NullReferenceException; ensure you still set _externalWallet =
Web3.Wallet, create sessionAccount, compute SessionTokenPDA via
FindSessionToken, and populate SessionWallet fields only after these checks.
In `@Runtime/codebase/Web3.cs`:
- Around line 31-37: You removed Unity serialization for the RpcCluster field by
adding [NonSerialized] to rpcCluster, which prevents Inspector values from being
saved; restore Inspector configurability by removing [NonSerialized] from the
rpcCluster field (or, if you intend compile-time locking via the MAINNET symbol,
keep the field serialized but make it read-only in the Inspector instead), and
ensure the conditional defaulting using MAINNET still sets the initial value
(rpcCluster and RpcCluster) while preserving serialization or explicitly
document the breaking change if you must keep it non-serialized.
Comment on lines
+69
to
+85
| public static SessionWallet GetSessionWallet(string publicKey, string privateKey, PublicKey targetProgram) | ||
| { | ||
| _externalWallet = Web3.Wallet; | ||
|
|
||
| var sessionAccount = new Account(privateKey, publicKey); | ||
|
|
||
| // TODO: ActiveRpcClient can be null, get node address some other way | ||
| var sessionWallet = new SessionWallet(_externalWallet.RpcCluster, _externalWallet.ActiveRpcClient.NodeAddress.ToString()) | ||
| { | ||
| TargetProgram = targetProgram, | ||
| EncryptedKeystoreKey = $"{_externalWallet.Account.PublicKey}_SessionKeyStore", | ||
| SessionTokenPDA = FindSessionToken(targetProgram, sessionAccount, _externalWallet.Account), | ||
| Account = sessionAccount, | ||
| }; | ||
|
|
||
| return sessionWallet; | ||
| } |
There was a problem hiding this comment.
Missing singleton guard — will throw if Instance is already set.
The other GetSessionWallet overloads (lines 97 and 129) check if (Instance != null) return Instance; before constructing a new SessionWallet. This overload skips that check, so calling it when a session wallet already exists will hit the throw new Exception("SessionWallet already exists") in the constructor (line 37).
Additionally, Web3.Wallet (line 71) and _externalWallet.ActiveRpcClient (line 76) can both be null, which would cause a NullReferenceException.
Proposed fix
public static SessionWallet GetSessionWallet(string publicKey, string privateKey, PublicKey targetProgram)
{
+ if (Instance != null) return Instance;
+
_externalWallet = Web3.Wallet;
+ if (_externalWallet?.ActiveRpcClient == null)
+ throw new InvalidOperationException("External wallet or its RPC client is not initialized.");
var sessionAccount = new Account(privateKey, publicKey);
- // TODO: ActiveRpcClient can be null, get node address some other way
var sessionWallet = new SessionWallet(_externalWallet.RpcCluster, _externalWallet.ActiveRpcClient.NodeAddress.ToString())🤖 Prompt for AI Agents
In `@Runtime/codebase/SessionWallet.cs` around lines 69 - 85,
GetSessionWallet(string publicKey, string privateKey, PublicKey targetProgram)
lacks the singleton guard and null checks: add the same early return used in the
other overloads (if (Instance != null) return Instance;) before constructing a
new SessionWallet, and validate Web3.Wallet and its ActiveRpcClient (e.g., if
(Web3.Wallet == null) throw/new error or handle; if
(_externalWallet.ActiveRpcClient == null) get node address another way or throw)
to avoid NullReferenceException; ensure you still set _externalWallet =
Web3.Wallet, create sessionAccount, compute SessionTokenPDA via
FindSessionToken, and populate SessionWallet fields only after these checks.
Comment on lines
+31
to
+37
|
|
||
| #if MAINNET | ||
| [NonSerialized] public RpcCluster rpcCluster = RpcCluster.MainNet; | ||
| #else | ||
| [NonSerialized] public RpcCluster rpcCluster = RpcCluster.DevNet; | ||
| #endif | ||
|
|
There was a problem hiding this comment.
[NonSerialized] removes Unity Inspector configurability — likely a breaking change.
Adding [NonSerialized] to rpcCluster means Unity will no longer serialize this field, so any value previously set in the Inspector will be silently ignored. Users who relied on the Inspector to select their cluster will now always get the compile-time default with no visual feedback that the setting is gone.
If the intent is to lock the cluster at compile time via MAINNET, consider keeping the field serialized but read-only in the Inspector (e.g., show it disabled), or at minimum document this as a breaking change. If the [NonSerialized] was unintentional, remove it so Inspector-based configuration continues to work.
🤖 Prompt for AI Agents
In `@Runtime/codebase/Web3.cs` around lines 31 - 37, You removed Unity
serialization for the RpcCluster field by adding [NonSerialized] to rpcCluster,
which prevents Inspector values from being saved; restore Inspector
configurability by removing [NonSerialized] from the rpcCluster field (or, if
you intend compile-time locking via the MAINNET symbol, keep the field
serialized but make it read-only in the Inspector instead), and ensure the
conditional defaulting using MAINNET still sets the initial value (rpcCluster
and RpcCluster) while preserving serialization or explicitly document the
breaking change if you must keep it non-serialized.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Issues connecting to and disconnecting from wallet apps on android
Solution
Rewrote LocalAssociationScenario.cs
Summary by CodeRabbit
New Features
Improvements