Automated Code Annotation with LLMs for Establishing TEE Boundaries

Automated Code Annotation with LLMs for Establishing TEE Boundaries

Speakers: Varun Gadey

Conference: NDSS Symposium 2026

YouTube: https://www.youtube.com/watch?v=TOTIK4sMqO0

Overview

Deciding which code should run inside a Trusted Execution Environment (TEE) versus untrusted space is a critical security decision that has traditionally required manual analysis -- a process that is time-consuming, error-prone, and does not scale to large codebases. This talk presents an automated tool that uses fine-tuned large language models to predict at the line level which code is security-sensitive and should be placed within TEE boundaries, achieving 97% precision and recall at the line level and zero false positives and false negatives at the function level.

The approach centers on a novel notion called "CryptiCS code" -- security-sensitive code defined around cryptographic operations and their data dependencies. The researchers built a manually annotated dataset from open-source GitHub projects, fine-tuned Google Gemma (2 billion parameters) with QLoRA (4-bit quantized Low-Rank Adaptation), and evaluated on both in-distribution test sets and unseen out-of-distribution codebases. The tool successfully identifies encryption/decryption logic, access control, secure communication, authentication, and data encryption code across diverse projects.

Background

▶ Watch: The problem: manually finding TEE boundaries doesn't scale (0:00)

Trusted Execution Environments like Intel SGX, ARM TrustZone, and AMD SEV provide hardware-based isolation for sensitive computations. The fundamental design question is code partitioning: which code needs the protection of the TEE and which can run in untrusted space? Getting this wrong in either direction has consequences -- including too much code bloats the Trusted Computing Base (TCB), increasing attack surface, while excluding sensitive code defeats the purpose of the TEE.

Current approaches fall into two camps. Manual code splitting (as in SOAP and DataShield) requires developers to identify security-sensitive code by hand, which is slow and doesn't scale. Whole-codebase migration (as in Graphene and SCONE) moves everything into the TEE, but this dramatically increases TCB size and the associated attack surface. Neither approach is satisfactory for large, evolving codebases.

The research question is whether LLMs can be fine-tuned to automatically identify security-sensitive code with sufficient accuracy to replace manual annotation, reducing TCB size while maintaining security guarantees.

Key Findings

▶ Watch: Dataset construction via snowballing from crypto library calls (4:00)

The researchers achieved strong classification performance:

Line-level results: Using the fine-tuned Gemma 2B model, predictions showed only 1-6 false positives and 452 false negatives across the entire test set, yielding approximately 97% precision, F1 score, and recall. This means the model correctly identifies the vast majority of security-sensitive lines while rarely flagging benign code.

Function-level results: When line-level predictions are resolved to function granularity, false positives and false negatives drop to zero. The model correctly identifies all security-sensitive functions and all non-sensitive functions without any false predictions.

Out-of-distribution generalization: Case studies on unseen codebases demonstrate that the model adapts to new projects. In the demonstrated example, a decryption code block with a TLS crypto library call, error handling, and plaintext output was correctly classified with all lines receiving probability scores above 0.5.

Model selection: Among Google Gemma, Meta LLaMA, and other models tested, Gemma 2B provided the best results, suggesting that smaller, well-fine-tuned models can outperform larger ones for this specialized task.

Technical Deep Dive

▶ Watch: Feature engineering: local context, data flows, and call graphs (6:00)

The tool's architecture has three stages:

Data Preprocessing: For each line of code, the system constructs an input sequence combining three feature types:

• Local features: Pre- and post-context lines forming a local context window around the target line
• Global features: Data flow dependencies computed from Code Property Graphs (using the Joern tool) and function call graphs, capturing long-range dependencies that indicate whether a line feeds into or receives from cryptographic operations
• Metadata features: Function and file-level tracing information to enable resolution from line-level to function-level predictions

CryptiCS Code Definition: The dataset is built by first identifying crypto library calls in open-source projects, then defining an initial set of security-sensitive functions around those calls, and performing bidirectional snowballing -- tracing both forward and backward through call graphs and data flows to find dependent lines and functions. Data flow paths to cryptographic sinks are also labeled as security-sensitive.

QLoRA Fine-tuning: The pre-trained LLM is frozen, and small trainable adapter matrices (linear layers) are added to each transformer block. 4-bit quantization reduces memory requirements significantly. Dropout prevents overfitting to the relatively small manually annotated dataset. Class weighting balances learning between security-sensitive and non-sensitive code classes, ensuring the model learns subtle patterns of both.

The classification threshold is 0.5 probability -- any line above this threshold is predicted as security-sensitive. The researchers note that boundary cases (0.49 vs 0.53) can be fine-tuned for specific codebases if needed.

Demo / Proof of Concept

▶ Watch: Results: 97% precision/recall at line level, zero errors at function level (10:00)

The case study demonstrates the tool on an unseen codebase containing encryption/decryption logic. A code block with five lines is analyzed:

• Line 1: TLS crypto library call (embed_tls decryption) -- probability 0.69
• Lines 2-4: Error handling for the crypto operation -- probabilities above 0.5, with the if return != 0 check at 0.93
• Line 5: Print decrypted plaintext -- probability above 0.5

All lines are correctly identified as security-sensitive. The Q&A addresses why probability scores vary within the block (e.g., the error handling line scoring higher than the crypto call itself) -- the researchers explain this reflects the model's global learning rather than per-codebase optimization, and the 0.5 threshold is the relevant decision boundary.

Defensive Implications

▶ Watch: Case study: correctly identifying TLS decryption code in unseen project (10:30)

This tool addresses a practical barrier to TEE adoption: the manual effort required to determine TEE boundaries. Organizations considering TEE deployment for sensitive workloads can use automated annotation to:

• Reduce TCB size by identifying exactly which code needs TEE protection rather than migrating entire applications
• Scale TEE adoption to large codebases where manual annotation is infeasible
• Maintain TEE boundaries as codebases evolve, by re-running the tool on new code
• Audit existing TEE deployments by comparing automated predictions against manual annotations to find gaps

The zero-false-positive, zero-false-negative performance at function level means the tool can reliably guide code partitioning decisions. The practical deployment model is to use the tool as a recommendation engine, with developers reviewing the automated annotations before finalizing TEE boundaries.

However, the CryptiCS code definition is inherently limited to cryptography-adjacent security sensitivity. Code that is security-sensitive for other reasons (e.g., business logic authorization, rate limiting, input validation) may not be captured unless it has data flow paths to cryptographic operations.

Key Takeaways

• Fine-tuned LLMs can predict security-sensitive code at line level with 97% precision/recall and zero false positives/negatives at function level
• The CryptiCS code notion defines security sensitivity around cryptographic operations and their data dependencies
• Google Gemma 2B with QLoRA fine-tuning outperformed larger models for this task
• The tool generalizes to unseen codebases in out-of-distribution case studies
• Combining local context, global data flow/call graph features, and metadata features enables accurate predictions
• The approach can scale TEE adoption by automating the most labor-intensive part of code partitioning

About the Speaker(s)

Varun Gadey is a researcher from the University of Duisburg-Essen, Germany, working under Professor Deit Franco. The research was conducted while the team was at the University of Wurzburg, Germany. The work was developed in collaboration with an industry partner, indicating practical deployment requirements beyond academic evaluation. The team has expertise in trusted execution environments, LLM fine-tuning, and code analysis.