use std::collections::HashMap;
use serde_json::{Value, json};
use crate::base64::base64url_decode;
use crate::crypto::aes::{aes128_cbc_decrypt, aes128_ecb_decrypt};
use crate::error::{MegaError, Result};
use crate::fs::node::{Node, NodeType};
use crate::session::Session;
const MAX_PENDING_NODES: usize = 4096;
impl Session {
fn finalize_refreshed_tree_cache_state(&mut self) -> Result<()> {
self.nodes_state_ready = true;
self.recompute_state_current();
self.action_packets_current =
self.state_current && (self.current_seqtag.is_none() || self.current_seqtag_seen);
self.persist_tree_cache_state()
}
pub async fn refresh(&mut self) -> Result<()> {
self.ensure_keys_attribute().await?;
self.reset_state_current_tracking();
let response = self
.api_mut()
.request(json!({"a": "f", "c": 1, "r": 1, "ca": 1}))
.await?;
if let Some(sn) = response.get("sn").and_then(|v| v.as_str()) {
self.scsn = Some(sn.to_string());
self.wsc_url = None;
self.sc_catchup = true;
self.current_seqtag = None;
self.current_seqtag_seen = false;
self.alerts_catchup_pending = true;
} else {
self.sc_catchup = false;
self.sc_batch_catchup_done = true;
}
if let Some(ok_array) = response.get("ok").and_then(|v| v.as_array()) {
self.parse_share_keys(ok_array);
}
if let Some(s_array) = response.get("s").and_then(|v| v.as_array()) {
self.ingest_outshares_from_fetch(s_array);
}
let public_links = response
.get("ph")
.and_then(|v| v.as_array())
.map(|arr| Self::parse_public_links(arr))
.unwrap_or_default();
let nodes_array = response
.get("f")
.and_then(|v| v.as_array())
.ok_or(MegaError::InvalidResponse)?;
for node_json in nodes_array {
if let Some(1) = node_json.get("t").and_then(|v| v.as_i64())
&& let (Some(handle), Some(kstr)) = (
node_json.get("h").and_then(|v| v.as_str()),
node_json.get("k").and_then(|v| v.as_str()),
)
{
for part in kstr.split('/') {
if let Some((key_handle, encrypted_key)) = part.split_once(':')
&& key_handle == self.user_handle
&& let Ok(enc) = base64url_decode(encrypted_key)
{
let dec = aes128_ecb_decrypt(&enc, self.master_key());
if dec.len() >= 16 {
let mut key = [0u8; 16];
key.copy_from_slice(&dec[..16]);
if self.key_manager.get_share_key_from_str(handle).is_none() {
self.key_manager.add_share_key_from_str(handle, &key);
}
}
}
}
}
}
self.pending_nodes.clear();
let mut nodes = Vec::new();
for node_json in nodes_array {
if let Some(mut node) = self.try_parse_or_stash(node_json) {
if let Some(link) = public_links.get(&node.handle) {
node.link = Some(link.clone());
}
nodes.push(node);
}
}
{
let handle_idx: HashMap<String, usize> = nodes
.iter()
.enumerate()
.map(|(i, n)| (n.handle.clone(), i))
.collect();
let parent_types: Vec<Option<NodeType>> = nodes
.iter()
.map(|n| {
n.parent_handle
.as_ref()
.and_then(|ph| handle_idx.get(ph))
.map(|&pidx| nodes[pidx].node_type)
})
.collect();
for i in 0..nodes.len() {
if self.outshares.contains_key(&nodes[i].handle)
|| self.pending_outshares.contains_key(&nodes[i].handle)
{
nodes[i].is_outshare = true;
}
if nodes[i].share_key.is_some() && nodes[i].node_type == NodeType::Folder {
let is_inshare = match parent_types[i] {
None => true,
Some(pt) => matches!(pt, NodeType::Root | NodeType::Network),
};
if is_inshare {
nodes[i].is_inshare = true;
}
}
if nodes[i].is_inshare {
let handle = nodes[i].handle.clone();
if let Some(node_json) = nodes_array
.iter()
.find(|j| j.get("h").and_then(|v| v.as_str()) == Some(&handle))
{
nodes[i].share_access = node_json
.get("r")
.and_then(|v| v.as_i64())
.map(|v| v as i32);
}
}
}
}
Self::build_node_paths(&mut nodes);
self.nodes = nodes;
self.drain_pending_nodes();
self.finalize_refreshed_tree_cache_state()?;
if self.clear_inuse_flags_for_missing_shares() {
let _ = self.persist_keys_with_retry().await;
}
Ok(())
}
fn parse_share_keys(&mut self, ok_array: &[Value]) {
for ok in ok_array {
if let (Some(h), Some(k)) = (
ok.get("h").and_then(|v| v.as_str()),
ok.get("k").and_then(|v| v.as_str()),
) {
if k.len() > 22 {
if let Ok(encrypted) = base64url_decode(k)
&& let Some(decrypted) = self.rsa_key().decrypt(&encrypted)
&& decrypted.len() >= 16
{
let mut key = [0u8; 16];
key.copy_from_slice(&decrypted[..16]);
if self.key_manager.get_share_key_from_str(h).is_none() {
self.key_manager.add_share_key_from_str(h, &key);
}
}
} else if let Ok(encrypted) = base64url_decode(k) {
let decrypted = aes128_ecb_decrypt(&encrypted, self.master_key());
if decrypted.len() >= 16 {
let mut key = [0u8; 16];
key.copy_from_slice(&decrypted[..16]);
if self.key_manager.get_share_key_from_str(h).is_none() {
self.key_manager.add_share_key_from_str(h, &key);
}
}
}
}
}
}
fn parse_public_links(ph_array: &[Value]) -> HashMap<String, String> {
let mut links = HashMap::new();
for ph in ph_array {
if let (Some(h), Some(ph_handle)) = (
ph.get("h").and_then(|v| v.as_str()),
ph.get("ph").and_then(|v| v.as_str()),
) {
links.insert(h.to_string(), ph_handle.to_string());
}
}
links
}
pub(crate) fn parse_node(&self, json: &Value) -> Option<Node> {
let handle = json.get("h")?.as_str()?.to_string();
let parent_handle = json
.get("p")
.and_then(|v| v.as_str())
.map(|s| s.to_string());
let node_type_int = json.get("t")?.as_i64()?;
let node_type = NodeType::from_i64(node_type_int)?;
let size = json.get("s").and_then(|v| v.as_u64()).unwrap_or(0);
let timestamp = json.get("ts").and_then(|v| v.as_i64()).unwrap_or(0);
let file_attr = json
.get("fa")
.and_then(|v| v.as_str())
.map(|s| s.to_string());
let (name, node_key, share_key, share_handle) = match node_type {
NodeType::Root => ("Root".to_string(), Vec::new(), None, None),
NodeType::Inbox => ("Inbox".to_string(), Vec::new(), None, None),
NodeType::Trash => ("Trash".to_string(), Vec::new(), None, None),
_ => {
let attrs_b64 = json.get("a")?.as_str()?;
let key_str = json.get("k")?.as_str()?;
let (node_key, used_share_key) = self.decrypt_node_key(key_str)?;
let sh = used_share_key.and_then(|_| {
key_str.split('/').find_map(|part| {
let (kh, _) = part.split_once(':')?;
if kh != self.user_handle && self.key_manager.contains_share_key(kh) {
Some(kh.to_string())
} else {
None
}
})
});
match self.decrypt_node_attrs(attrs_b64, &node_key) {
Some(name) => (name, node_key, used_share_key, sh),
None => return None,
}
}
};
Some(Node {
name,
handle,
parent_handle,
node_type,
size,
timestamp,
key: node_key,
path: None,
link: None,
file_attr,
share_key,
share_handle,
is_inshare: false,
is_outshare: false,
share_access: None,
})
}
pub(crate) fn try_parse_or_stash(&mut self, json: &Value) -> Option<Node> {
if let Some(node) = self.parse_node(json) {
return Some(node);
}
let handle = json.get("h").and_then(|v| v.as_str())?;
let node_type = json
.get("t")
.and_then(|v| v.as_i64())
.and_then(NodeType::from_i64)?;
if !matches!(node_type, NodeType::File | NodeType::Folder) {
return None;
}
let key_str = json.get("k").and_then(|v| v.as_str())?;
let has_recognized_handle = key_str.split('/').any(|part| {
part.split_once(':')
.map(|(h, _)| h == self.user_handle || self.key_manager.contains_share_key(h))
.unwrap_or(false)
});
if !has_recognized_handle {
self.stash_pending_node(json.clone(), handle);
} else if self.decrypt_node_key(key_str).is_some() {
self.stash_pending_node(json.clone(), handle);
}
None
}
fn stash_pending_node(&mut self, json: Value, handle: &str) {
if self.pending_nodes.len() >= MAX_PENDING_NODES {
tracing::warn!(
pending_count = self.pending_nodes.len(),
"pending_nodes queue at capacity ({MAX_PENDING_NODES}), dropping oldest entry"
);
self.pending_nodes.remove(0);
}
tracing::debug!(node_handle = handle, "stashing node with missing key");
self.pending_nodes.push(json);
}
pub(crate) fn drain_pending_nodes(&mut self) -> bool {
if self.pending_nodes.is_empty() {
return false;
}
let pending = std::mem::take(&mut self.pending_nodes);
let count = pending.len();
let mut recovered = 0usize;
for json in pending {
if let Some(node) = self.parse_node(&json) {
tracing::debug!(node_handle = %node.handle, "recovered node from pending queue");
self.upsert_node(node);
recovered += 1;
} else {
self.pending_nodes.push(json);
}
}
if recovered > 0 {
Self::build_node_paths(&mut self.nodes);
}
tracing::debug!(
recovered,
remaining = self.pending_nodes.len(),
total = count,
"drain_pending_nodes complete"
);
recovered > 0
}
fn decrypt_node_key(&self, key_str: &str) -> Option<(Vec<u8>, Option<[u8; 16]>)> {
for part in key_str.split('/') {
if let Some((key_handle, encrypted_key)) = part.split_once(':') {
let (decrypt_key_arr, used_share_key) = if key_handle == self.user_handle {
(Some(*self.master_key()), None)
} else if let Some(k) = self.key_manager.get_share_key_from_str(key_handle) {
(Some(k), Some(k))
} else {
(None, None)
};
if let Some(key) = decrypt_key_arr.as_ref()
&& let Ok(encrypted) = base64url_decode(encrypted_key)
{
let decrypted = aes128_ecb_decrypt(&encrypted, key);
return Some((decrypted, used_share_key));
}
}
}
None
}
pub(crate) fn decrypt_node_attrs(&self, attrs_b64: &str, node_key: &[u8]) -> Option<String> {
let encrypted = base64url_decode(attrs_b64).ok()?;
let aes_key: [u8; 16] = if node_key.len() >= 32 {
let mut key = [0u8; 16];
for i in 0..16 {
key[i] = node_key[i] ^ node_key[i + 16];
}
key
} else if node_key.len() >= 16 {
node_key[..16].try_into().ok()?
} else {
return None;
};
let decrypted = aes128_cbc_decrypt(&encrypted, &aes_key);
let text = String::from_utf8_lossy(&decrypted);
if !text.starts_with("MEGA") {
return None;
}
let json_str = text.trim_start_matches("MEGA").trim_end_matches('\0');
let attrs: Value = serde_json::from_str(json_str).ok()?;
attrs
.get("n")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
}
pub(crate) fn build_node_paths(nodes: &mut [Node]) {
let handle_map: HashMap<&str, usize> = nodes
.iter()
.enumerate()
.map(|(i, n)| (n.handle.as_str(), i))
.collect();
let paths: Vec<String> = (0..nodes.len())
.map(|i| Self::build_node_path(nodes, i, &handle_map, 0))
.collect();
for (i, path) in paths.into_iter().enumerate() {
nodes[i].path = Some(path);
}
}
fn build_node_path(
nodes: &[Node],
idx: usize,
handle_map: &HashMap<&str, usize>,
depth: usize,
) -> String {
if depth > 100 {
return format!("/{}", nodes[idx].name);
}
let node = &nodes[idx];
if matches!(
node.node_type,
NodeType::Root | NodeType::Inbox | NodeType::Trash | NodeType::Network
) {
return format!("/{}", node.name);
}
if let Some(parent_handle) = &node.parent_handle
&& let Some(&parent_idx) = handle_map.get(parent_handle.as_str())
{
let parent_path = Self::build_node_path(nodes, parent_idx, handle_map, depth + 1);
return format!("{}/{}", parent_path.trim_end_matches('/'), node.name);
}
format!("/{}", node.name)
}
}
#[cfg(test)]
mod tests {
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use serde_json::json;
use crate::base64::base64url_encode;
use crate::crypto::aes::{aes128_cbc_encrypt, aes128_ecb_encrypt};
use crate::fs::{Node, NodeType};
use crate::session::Session;
use crate::session::runtime::persistence::{MemoryPersistenceBackend, PersistenceRuntime};
use super::MAX_PENDING_NODES;
fn make_folder_node_json(
handle: &str,
name: &str,
key_handle: &str,
encrypt_key: &[u8; 16],
) -> serde_json::Value {
let node_key = [0xAA; 16];
let enc_key = aes128_ecb_encrypt(&node_key, encrypt_key);
let k_field = format!("{}:{}", key_handle, base64url_encode(&enc_key));
let attrs_plain = format!("MEGA{{\"n\":\"{name}\"}}");
let mut padded = attrs_plain.into_bytes();
let rem = padded.len() % 16;
if rem != 0 {
padded.resize(padded.len() + (16 - rem), 0);
}
let attrs_enc = aes128_cbc_encrypt(&padded, &node_key);
let a_field = base64url_encode(&attrs_enc);
json!({
"h": handle,
"p": "root",
"t": 1,
"ts": 12345,
"k": k_field,
"a": a_field,
})
}
#[test]
fn drain_recovers_stashed_node_after_share_key_added() {
let mut session = Session::test_dummy();
let share_key: [u8; 16] = [0x02; 16];
let node_json = make_folder_node_json("nodeA", "shared_folder", "shareXYZ", &share_key);
let result = session.try_parse_or_stash(&node_json);
assert!(result.is_none());
assert_eq!(session.pending_nodes.len(), 1);
assert!(session.nodes.is_empty());
session
.key_manager
.add_share_key_from_str("shareXYZ", &share_key);
assert!(session.drain_pending_nodes());
assert!(session.pending_nodes.is_empty());
assert_eq!(session.nodes.len(), 1);
assert_eq!(session.nodes[0].handle, "nodeA");
assert_eq!(session.nodes[0].name, "shared_folder");
}
#[test]
fn structurally_invalid_json_is_not_stashed() {
let mut session = Session::test_dummy();
let no_handle = json!({"t": 1, "k": "x:y", "a": "z"});
assert!(session.try_parse_or_stash(&no_handle).is_none());
assert!(session.pending_nodes.is_empty());
let no_type = json!({"h": "abc", "k": "x:y", "a": "z"});
assert!(session.try_parse_or_stash(&no_type).is_none());
assert!(session.pending_nodes.is_empty());
let bad_type = json!({"h": "abc", "t": 999, "k": "x:y", "a": "z"});
assert!(session.try_parse_or_stash(&bad_type).is_none());
assert!(session.pending_nodes.is_empty());
let no_key = json!({"h": "abc", "t": 0, "a": "z"});
assert!(session.try_parse_or_stash(&no_key).is_none());
assert!(session.pending_nodes.is_empty());
}
#[test]
fn corrupt_data_with_recognised_handle_is_not_stashed() {
let mut session = Session::test_dummy();
let corrupt = json!({
"h": "badnode",
"t": 0,
"p": "root",
"k": "myhandle:!!!invalid-base64!!!",
"a": "also-garbage",
});
assert!(session.try_parse_or_stash(&corrupt).is_none());
assert!(session.pending_nodes.is_empty());
}
#[test]
fn queue_cap_is_enforced() {
let mut session = Session::test_dummy();
for i in 0..MAX_PENDING_NODES + 10 {
let node = json!({
"h": format!("node{i}"),
"t": 1,
"p": "root",
"k": format!("unknown_handle{}:AAAA", i),
"a": "AAAA",
});
session.try_parse_or_stash(&node);
}
assert_eq!(session.pending_nodes.len(), MAX_PENDING_NODES);
let last_h = session
.pending_nodes
.last()
.and_then(|v| v.get("h"))
.and_then(|v| v.as_str())
.unwrap();
let expected = format!("node{}", MAX_PENDING_NODES + 9);
assert_eq!(last_h, expected);
}
#[test]
fn drain_with_no_pending_is_noop() {
let mut session = Session::test_dummy();
assert!(!session.drain_pending_nodes());
}
#[test]
fn root_nodes_parse_directly_without_stashing() {
let mut session = Session::test_dummy();
let root_json = json!({"h": "rootH", "t": 2, "ts": 0});
let node = session.try_parse_or_stash(&root_json);
assert!(node.is_some());
assert!(session.pending_nodes.is_empty());
assert_eq!(node.unwrap().name, "Root");
}
#[test]
fn finalize_refreshed_tree_cache_state_persists_coherent_snapshot() {
let persistence = PersistenceRuntime::new(Arc::new(MemoryPersistenceBackend::default()));
let mut session = Session::test_dummy().with_persistence_for_tests(persistence.clone());
session.scsn = Some("refresh-scsn".to_string());
session.nodes = vec![
Node {
name: "Root".to_string(),
handle: "root".to_string(),
parent_handle: None,
node_type: NodeType::Root,
size: 0,
timestamp: 0,
key: Vec::new(),
path: Some("/Root".to_string()),
link: None,
file_attr: None,
share_key: None,
share_handle: None,
is_inshare: false,
is_outshare: false,
share_access: None,
},
Node {
name: "docs".to_string(),
handle: "docs".to_string(),
parent_handle: Some("root".to_string()),
node_type: NodeType::Folder,
size: 0,
timestamp: 1,
key: vec![0x11; 16],
path: Some("/Root/docs".to_string()),
link: None,
file_attr: None,
share_key: None,
share_handle: None,
is_inshare: false,
is_outshare: true,
share_access: None,
},
];
session.pending_nodes = vec![json!({"h": "pending", "p": "docs", "t": 0})];
session.outshares =
HashMap::from([("docs".to_string(), HashSet::from(["EXP".to_string()]))]);
session
.finalize_refreshed_tree_cache_state()
.expect("refresh finalization should persist");
let stored = persistence
.load_engine_state(&session.persistence_scope())
.expect("load should succeed")
.expect("tree/cache snapshot should exist");
let tree = stored.tree.expect("refresh should persist tree snapshot");
assert_eq!(stored.sc.scsn.as_deref(), Some("refresh-scsn"));
assert!(session.nodes_state_ready);
assert_eq!(tree.nodes.len(), 2);
assert_eq!(tree.pending_nodes.len(), 1);
assert_eq!(
tree.outshares.get("docs"),
Some(&HashSet::from(["EXP".to_string()]))
);
}
#[test]
fn finalized_refresh_snapshot_restores_after_restart() {
let persistence = PersistenceRuntime::new(Arc::new(MemoryPersistenceBackend::default()));
let mut session = Session::test_dummy().with_persistence_for_tests(persistence.clone());
session.scsn = Some("refresh-scsn".to_string());
session.nodes = vec![
Node {
name: "Root".to_string(),
handle: "root".to_string(),
parent_handle: None,
node_type: NodeType::Root,
size: 0,
timestamp: 0,
key: Vec::new(),
path: Some("/Root".to_string()),
link: None,
file_attr: None,
share_key: None,
share_handle: None,
is_inshare: false,
is_outshare: false,
share_access: None,
},
Node {
name: "docs".to_string(),
handle: "docs".to_string(),
parent_handle: Some("root".to_string()),
node_type: NodeType::Folder,
size: 0,
timestamp: 1,
key: vec![0x11; 16],
path: Some("/Root/docs".to_string()),
link: None,
file_attr: None,
share_key: None,
share_handle: None,
is_inshare: false,
is_outshare: true,
share_access: None,
},
];
session.pending_nodes = vec![json!({"h": "pending", "p": "docs", "t": 0})];
session.outshares =
HashMap::from([("docs".to_string(), HashSet::from(["EXP".to_string()]))]);
session.pending_outshares = HashMap::from([(
"docs".to_string(),
HashSet::from(["pending-user".to_string()]),
)]);
session
.finalize_refreshed_tree_cache_state()
.expect("refresh finalization should persist");
let mut restored = Session::test_dummy().with_persistence_for_tests(persistence);
let loaded = restored
.restore_tree_cache_state()
.expect("restored refresh snapshot should load");
assert!(loaded);
assert_eq!(restored.scsn.as_deref(), Some("refresh-scsn"));
assert_eq!(restored.nodes.len(), 2);
assert_eq!(
restored
.nodes
.iter()
.find(|node| node.handle == "docs")
.and_then(Node::path),
Some("/Root/docs")
);
assert_eq!(restored.pending_nodes.len(), 1);
assert_eq!(
restored.outshares.get("docs"),
Some(&HashSet::from(["EXP".to_string()]))
);
assert_eq!(
restored.pending_outshares.get("docs"),
Some(&HashSet::from(["pending-user".to_string()]))
);
}
}