gbf_suite/
main.rs

1#![feature(backtrace_frames)]
2
3use std::{
4    backtrace::Backtrace,
5    env,
6    fs::{self},
7    path, process,
8    time::Instant,
9};
10
11use aws_upload::AwsUpload;
12use consts::{ExitCode, GBF_SUITE_INPUT_DIR_ENV_VAR};
13use dotenv::dotenv;
14use gbf_core::{
15    cfg_dot::{CfgDotConfig, DotRenderableGraph},
16    decompiler::{
17        ast::visitors::emit_context::EmitContextBuilder,
18        function_decompiler::{FunctionDecompilerBuilder, FunctionDecompilerErrorDetails},
19    },
20    module::ModuleBuilder,
21    utils::VERSION,
22};
23use gbf_result::{
24    GbfFunctionDao, GbfFunctionErrorDao, GbfGraphvizStructureAnalaysisDao, GbfModuleDao,
25    GbfSimplifiedBacktrace, GbfSimplifiedBacktraceFrame, GbfVersionDao,
26};
27use regex::Regex;
28use utils::hash_file;
29
30pub mod aws_upload;
31pub mod consts;
32pub mod gbf_result;
33pub mod utils;
34
35#[tokio::main]
36async fn main() {
37    // Load .env file if it exists
38    dotenv().ok();
39
40    let config_path = path::Path::new(env!("CARGO_MANIFEST_DIR")).join("logging_config.yaml");
41    log4rs::init_file(config_path, Default::default()).unwrap();
42
43    // Attempt to unwrap the environment variable. If it fails, exit with an error code.
44    let value = env::var(GBF_SUITE_INPUT_DIR_ENV_VAR).unwrap_or_else(|_| {
45        log::error!(
46            "Environment variable {} not found",
47            GBF_SUITE_INPUT_DIR_ENV_VAR
48        );
49        process::exit(ExitCode::EnvVarNotFound.into());
50    });
51
52    // Attempt to open the directory. If it fails, exit with an error code.
53    let dir = std::fs::read_dir(value).unwrap_or_else(|_| {
54        log::error!("Invalid directory");
55        process::exit(ExitCode::InvalidDir.into());
56    });
57
58    let uploader = aws_upload::AwsUpload::new().await;
59
60    // If GBF_VERSION is set, create option and pass it to `process_module`
61    let gbf_version_override = env::var("GBF_VERSION").ok();
62
63    // Iterate over the directory entries.
64    let time = Instant::now();
65    for entry in dir {
66        // Attempt to unwrap the entry. If it fails, log an error and continue.
67        let entry = match entry {
68            Ok(entry) => entry,
69            Err(err) => {
70                log::error!("Failed to get the directory entry: {}", err);
71                continue;
72            }
73        };
74
75        // If entry is a directory, skip it.
76        if entry.file_type().unwrap().is_dir() {
77            continue;
78        }
79
80        let result = process_module(
81            &uploader,
82            entry.path().as_ref(),
83            gbf_version_override.clone(),
84        )
85        .await;
86
87        if let Err(e) = result {
88            log::error!("Failed to process module: {}", e);
89            std::process::exit(ExitCode::UnexpectedError.into());
90        }
91    }
92    let total_time = time.elapsed();
93
94    let gbf_version = GbfVersionDao {
95        gbf_version: gbf_version_override.clone().unwrap_or(VERSION.to_string()),
96        total_time,
97        suite_timestamp: std::time::SystemTime::now()
98            .duration_since(std::time::UNIX_EPOCH)
99            .unwrap()
100            .as_secs(),
101    };
102
103    match uploader.upload_gbf_version(gbf_version).await {
104        Ok(_) => log::info!("Uploaded GBF version"),
105        Err(e) => {
106            log::error!("Failed to upload GBF version: {}", e);
107            std::process::exit(ExitCode::UnexpectedError.into());
108        }
109    }
110}
111
112async fn process_module(
113    uploader: &AwsUpload,
114    path: &path::Path,
115    gbf_version_override: Option<String>,
116) -> Result<(), Box<dyn std::error::Error>> {
117    let module_name = path
118        .file_name()
119        .ok_or("Failed to get file name")?
120        .to_str()
121        .ok_or("Failed to convert file name to string")?
122        .to_string();
123
124    let module_id = hash_file(path)?;
125
126    let reader = fs::File::open(path)?;
127
128    let time = Instant::now();
129    let module = match ModuleBuilder::new()
130        .name(module_name.clone())
131        .reader(Box::new(reader))
132        .build()
133    {
134        Ok(module) => module,
135        Err(e) => {
136            log::error!("Failed to build module {}: {:?}", module_name, e);
137            return Ok(());
138        }
139    };
140    let module_time = time.elapsed();
141
142    log::info!("Decompiling module {}", module_name.clone());
143
144    let mut module_dao = GbfModuleDao {
145        gbf_version: gbf_version_override.clone().unwrap_or(VERSION.to_string()),
146        module_id: module_id.to_string(),
147        file_name: module_name,
148        module_load_time: module_time,
149        decompile_success: true,
150    };
151
152    for func in module.iter() {
153        let func_basic_block_dot = func.render_dot(CfgDotConfig::default());
154        let func_basic_block_dot_key = uploader.upload_graphviz_dot(func_basic_block_dot).await?;
155
156        log::info!(
157            "Decompiling function {}",
158            func.id.name.clone().unwrap_or("entry".to_string())
159        );
160
161        let time = Instant::now();
162
163        let mut decompiler = FunctionDecompilerBuilder::new(func)
164            .structure_analysis_max_iterations(100)
165            .structure_debug_mode(true)
166            .build();
167
168        let res = decompiler.decompile(
169            EmitContextBuilder::default()
170                .include_ssa_versions(true)
171                .build(),
172        );
173        let function_time = time.elapsed();
174
175        let decompile_success = if res.is_err() {
176            let error = GbfFunctionErrorDao {
177                gbf_version: gbf_version_override.clone().unwrap_or(VERSION.to_string()),
178                module_id: module_id.to_string(),
179                function_address: func.id.address,
180                error_type: res.as_ref().unwrap_err().error_type().to_string(),
181                message: res.as_ref().unwrap_err().to_string(),
182                backtrace: process_backtrace(res.as_ref().unwrap_err().backtrace()),
183                context: res.as_ref().unwrap_err().context().clone(),
184            };
185            module_dao.decompile_success = false;
186            uploader.upload_gbf_function_error(error).await?;
187            false
188        } else {
189            true
190        };
191
192        let function_dao = GbfFunctionDao {
193            gbf_version: gbf_version_override.clone().unwrap_or(VERSION.to_string()),
194            module_id: module_id.to_string(),
195            function_address: func.id.address,
196            function_name: func.id.name.clone(),
197            decompile_success,
198            total_time: function_time,
199            dot_key: func_basic_block_dot_key,
200            decompile_result: res.ok().map(|r| r.to_string()),
201        };
202
203        uploader.upload_gbf_function(function_dao).await?;
204
205        log::info!(
206            "Decompiled function {} in {}ms",
207            func.id.name.clone().unwrap_or("entry".to_string()),
208            function_time.as_millis()
209        );
210
211        // Upload all of the intermediate dot files
212        for (i, dot) in decompiler
213            .get_structure_analysis_snapshots()?
214            .iter()
215            .enumerate()
216        {
217            let dot_key = uploader.upload_graphviz_dot(dot.clone()).await?;
218
219            let graphviz_dao = GbfGraphvizStructureAnalaysisDao {
220                gbf_version: gbf_version_override.clone().unwrap_or(VERSION.to_string()),
221                module_id: module_id.to_string(),
222                function_address: func.id.address,
223                structure_analysis_step: i,
224                dot_key,
225            };
226
227            uploader.upload_gbf_graphviz_dao(graphviz_dao).await?;
228        }
229    }
230
231    uploader.upload_gbf_module(module_dao).await?;
232
233    Ok(())
234}
235
236pub fn process_backtrace(backtrace: &Backtrace) -> GbfSimplifiedBacktrace {
237    // Filter frames to include only those starting with "gbf_core"
238    let include_fn_regex = Regex::new(r"^gbf_core::").unwrap();
239
240    let mut simplified_frames = Vec::new();
241
242    for frame in backtrace.frames() {
243        // Convert the frame to a string representation
244        let frame_str = format!("{:?}", frame);
245
246        // Extract the function name
247        let function_name = frame_str
248            .split_once("fn: \"")
249            .and_then(|(_, rest)| rest.split_once("\""))
250            .map(|(fn_name, _)| fn_name.to_string());
251
252        // Extract the file path
253        let file_path = frame_str
254            .split_once("file: \"")
255            .and_then(|(_, rest)| rest.split_once("\""))
256            .map(|(file, _)| file.to_string());
257
258        // Extract the line number
259        let line_number = frame_str
260            .split_once("line: ")
261            .and_then(|(_, rest)| rest.split_once("}"))
262            .and_then(|(line, _)| line.trim().parse::<u32>().ok());
263
264        // Only include frames that match the `gbf_core` prefix
265        if let Some(function_name) = function_name {
266            if include_fn_regex.is_match(&function_name) {
267                simplified_frames.push(GbfSimplifiedBacktraceFrame {
268                    function: function_name,
269                    file: file_path.unwrap(),
270                    line: line_number.unwrap(),
271                });
272            }
273        }
274    }
275
276    GbfSimplifiedBacktrace {
277        frames: simplified_frames,
278    }
279}
gbf_suite/main.rs

gbf_suite/
main.rs
