bat-cli 0.11.0

Blockchain Auditor Toolkit (BAT)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146

use std::collections::HashSet;
use std::fs;

pub struct FileClassification {
    pub context_accounts_names: HashSet<String>,
    pub solana_account_names: HashSet<String>,
    pub entrypoint_function_names: HashSet<String>,
}

pub fn classify_file_from_path(path: &str) -> FileClassification {
    let content = fs::read_to_string(path).unwrap_or_default();
    classify_file(&content)
}

pub fn classify_file(file_content: &str) -> FileClassification {
    let mut classification = FileClassification {
        context_accounts_names: HashSet::new(),
        solana_account_names: HashSet::new(),
        entrypoint_function_names: HashSet::new(),
    };

    let Ok(file) = syn::parse_file(file_content) else {
        return classification;
    };

    for item in &file.items {
        match item {
            syn::Item::Struct(item_struct) => {
                if has_derive_accounts(item_struct) {
                    classification
                        .context_accounts_names
                        .insert(item_struct.ident.to_string());
                } else if has_account_attribute(item_struct) {
                    classification
                        .solana_account_names
                        .insert(item_struct.ident.to_string());
                }
            }
            syn::Item::Mod(item_mod) => {
                if has_program_attribute(item_mod) {
                    extract_entrypoint_functions(item_mod, &mut classification);
                }
            }
            _ => {}
        }
    }

    classification
}

/// Finds the `Context<T>` type parameter for a given entrypoint function name
/// by parsing the `#[program]` module with syn. Returns `Some(T)` if found.
pub fn get_context_type_for_entrypoint(
    file_content: &str,
    entrypoint_name: &str,
) -> Option<String> {
    let file = syn::parse_file(file_content).ok()?;
    for item in &file.items {
        if let syn::Item::Mod(item_mod) = item {
            if !has_program_attribute(item_mod) {
                continue;
            }
            let Some((_, items)) = &item_mod.content else {
                continue;
            };
            for inner in items {
                if let syn::Item::Fn(item_fn) = inner {
                    if item_fn.sig.ident != entrypoint_name {
                        continue;
                    }
                    // Look for the Context<T> parameter
                    for arg in &item_fn.sig.inputs {
                        if let syn::FnArg::Typed(pat_type) = arg {
                            if let Some(ctx_type) = extract_context_type(&pat_type.ty) {
                                return Some(ctx_type);
                            }
                        }
                    }
                }
            }
        }
    }
    None
}

/// Extracts `T` from `Context<'_, T>` or `Context<T>` type.
fn extract_context_type(ty: &syn::Type) -> Option<String> {
    use quote::ToTokens;
    if let syn::Type::Path(type_path) = ty {
        let segment = type_path.path.segments.last()?;
        if segment.ident != "Context" {
            return None;
        }
        if let syn::PathArguments::AngleBracketed(args) = &segment.arguments {
            // The last generic argument that is a Type (not a lifetime) is T
            for arg in args.args.iter().rev() {
                if let syn::GenericArgument::Type(inner_ty) = arg {
                    let ty_str = inner_ty.to_token_stream().to_string();
                    return Some(
                        crate::batbelt::parser::function_parser::normalize_generic_type(&ty_str),
                    );
                }
            }
        }
    }
    None
}

fn has_derive_accounts(item: &syn::ItemStruct) -> bool {
    item.attrs.iter().any(|attr| {
        if !attr.path().is_ident("derive") {
            return false;
        }
        let Ok(nested) = attr.parse_args_with(
            syn::punctuated::Punctuated::<syn::Path, syn::Token![,]>::parse_terminated,
        ) else {
            return false;
        };
        nested.iter().any(|path| path.is_ident("Accounts"))
    })
}

fn has_account_attribute(item: &syn::ItemStruct) -> bool {
    item.attrs
        .iter()
        .any(|attr| attr.path().is_ident("account"))
}

fn has_program_attribute(item: &syn::ItemMod) -> bool {
    item.attrs
        .iter()
        .any(|attr| attr.path().is_ident("program"))
}

fn extract_entrypoint_functions(item_mod: &syn::ItemMod, classification: &mut FileClassification) {
    let Some((_, items)) = &item_mod.content else {
        return;
    };
    for item in items {
        if let syn::Item::Fn(item_fn) = item {
            classification
                .entrypoint_function_names
                .insert(item_fn.sig.ident.to_string());
        }
    }
}