Rusteze

Inspiration

A combination of factors led us to this project. We learned about exploitable cybersecurity vulnerabilities in class, and we saw that the White House was encouraging migration towards memory safe languages, giving us a clear value proposition if we could find a cost effective and time efficient way to perform migrations. We also wanted an excuse to learn Rust! And so, we decided to build Rusteze.

What it does

Rust-eze transfers your C project into memory safe Rust code. It does this by traversing your directory and using Gemini pro to translate files to Rust code. Then we can iteratively make calls to fix and improve the resulting Rust code until it compiles. Rust will not compile code unless it is sure that it is memory safe. By changing the code to working Rust, it closes memory safety vulnerabilities and eliminates memory leaks.

How we built it

We built a VS-Code extension in typescript that detects the current workspace directory and then calls a script which walks the directory, scanning for c files and prompts Gemini to rewrite them. For large files, we chunk these calls, providing full context first and then asking for code in smaller pieces. This was necessary to take advantage of Gemini’s large input context window and smaller output token limit. The script also has functionality to test files for error output and reprompt Gemini with error information, allowing Gemini to rewrite Rust code it previously generated to fix its mistakes. Furthermore, we have a set of unit tests to validate Gemini’s performance - these tests cross compare with the original code for correct output to ensure the generated Rust code is correct, and does not break against attackers while the C code does break against the same attack.

Challenges we ran into

We ran into challenges when getting the extension to work, especially when we wanted to run it on a different machine. To solve this we had to learn the organizational structure of VS-Code extensions. Another challenge we ran into was the non-deterministic nature of the LLM output. This made it difficult because the output requires the code to be runnable and have the same behavior as the input code. Our solution to this was to automatically revise and improve the code when the first changes are unsuccessful. One final challenge we ran into was writing large rust files. The output token limits on Gemini Pro made us get clever with how we designed this.

Accomplishments that we're proud of

We were able to successfully develop a pipeline that works with any size directory or file. We are also proud of how we implemented the retry loop to improve the rust code iteratively. On top of this, it was exciting to implement this as a VS-Code extension. It made it very seamless to use which we were very proud of.

What we learned

We learned that creating a robust, accurate infrastructure using an LLM is very difficult. LLM’ are designed to be nondeterministic which makes relying on a consistent, objective output more difficult. Because our use case relies on generating more consistent output that is correct according to the Rust compiler, we have had to work with the LLM more to generate output more aligned to our goals. We also learned a little more about where the capabilities of LLMs begin and end.

The biggest thing we learned as a group is the problem facing a large amount of code in our nation. The memory unsafe nature of C and C++ leads to large security vulnerabilities in code. This is a real problem and we got to explore how memory unsafe code can be taken advantage of. We also learned more in depth about these memory issues, and how exploits can be taken against existing C code to exploit these vulnerabilities. Moving forward, we will make sure to develop code that protects against such attacks.

What's next for Rust-eze

We would like to be able to use Gemini 1.5 Pro (we require more prompts than we are allowed from Gemini 1.5’s rate limit) in the future to allow a larger token output, allowing us to generate larger C files at once, and use the better code generation features of Gemini 1.5. We would also like to automate the generation of a complete, consolidated Rust codebase including the Cargo.toml and necessary packages for the created Rust codebase. As mentioned before, due to the nature of LLMs, it is difficult to rely on them for consistency. We would like to explore generating multiple versions of each code sample, which we can then consolidate into a single result, which is ideally more correct than the samples it is made from, reducing the variance in the results we generate.