Type Level Security: The future of secure AI code generation?

Intro

With code being written (& generated) faster than ever before, there is the unfortunate side effect that security vulnerabilities are also coming faster than ever before. Asking your LLM not to include security vulnerabilities in its code doesn't always work. It is becoming clear that the way software is built today, manually or with assistance, is insufficient when it comes to reliably, consistently, and provably writing secure code.

Rust's rise in popularity has shown that it is possible to completely remove entire vulnerability classes in a reasonable and ergonomic way. Rust has effectively eliminated all (most) memory corruption vulnerabilities at compile time. Why, then, should we limit ourselves to this class of vulnerabilities?

In this post, I will discuss how it can be possible to write code in such a way that web application security vulnerabilities are uncompilable (or un-type-checkable), how having secure-by-design libraries or well-placed library wrappers can completely stop entire classes of vulnerabilities from being written, either manually or by an LLM. I will show code patterns in both Python and Rust that could be used to eliminate vulnerability classes.

Why types

When used properly, type systems can be incredibly powerful tools. They allow you to codify the invariants of a system, that is, all of the rules, properties, and relationships of every piece of data in your program. Rather than leaving a helpful comment or performing run-time assertions (that may only be hit when your first customer tries to do something important), you can ensure that every single caller of your add The function is only passing two int types, in every case, across the whole codebase.

Anecdotally, code I write with extremely strict types contains considerably fewer runtime bugs than code I write without, as I am forced to reason about every parameter, every piece of data, and every input. Codified appropriately, I make considerably fewer mistakes, and my code then only has business logic bugs rather than 'oops, I passed a string to an integer addition function'.

Many security vulnerabilities are just a specific type of bug, and if you take my above assertion to be true, then many such bugs should be solvable using a sufficiently flexible type system.

Trusted types

This post is, in part, inspired by the Trusted Types web API. This API effectively mitigates most forms of DOM XSS by ensuring that XSS injection sinks accept only known-safe values, such as those sanitized by an appropriate XSS sanitizer. This API can be further locked down using CSP to require the use of the Trusted Types API, throwing a TypeError if it is not used.

A similar mechanism is in use by the Linux Kernel with the __user macro, which ensures that pointers from userland are handled appropriately.

In the following sections, we will look at the generalization of this technique and how you can apply it to arbitrary security vulnerability classes.

Solving IDOR

Insecure Direct Object Reference (IDOR) is a pernicious vulnerability stemming from a lack of authentication and/or authorization checks when performing API actions. The classic example of this is an API endpoint that takes an incrementing user ID parameter and returns user data, but when changing the user ID value it returns data from a different user, without appropriate authorization checks.

Statistically, you, the reader, likely use Python, so we will first explore this vulnerability and how to solve it in Python using only type hints.

In Python

We start with a vulnerable example:

1import sqlite3
2
3from fastapi import FastAPI, Depends
4
5app = FastAPI()
6
7DATABASE = "app.db"
8
9def get_db():
10    conn = sqlite3.connect(DATABASE)
11    conn.row_factory = sqlite3.Row
12    try:
13        yield conn
14    finally:
15        conn.close()
16
17@app.get("/balance/{user_id}")
18def get_balance(user_id: int, db: sqlite3.Connection = Depends(get_db)):
19    cursor = db.execute(f"SELECT balance FROM accounts WHERE user_id = {user_id}")
20    row = cursor.fetchone()
21    return {"uid": user_id, "balance": row["balance"]}

This is a pretty standard-looking API call: you pass in an integer user ID, query the database, and return the result as the response. But whoops! We forgot to add any authentication or authorization checks; any user can request any other user's balance (assuming they have the ID), a classic IDOR. We receive the vulnerability report, we pay the researcher their bounty, and we update the function thus:

1@app.get("/balance/{user_id}")
2def get_balance(user_id: int, db: sqlite3.Connection = Depends(get_db)):
3    if not (check_authentication() and check_authorization()): return 403
4    ...

Perfect, the user data is safe. But that's a frightfully easy mistake to make. Forgetting one auth check on one API endpoint can lead to significant impacts. Now, let us look at the same endpoint, but using the type system to make sure this vulnerability doesn't happen again:

1import sqlite3
2
3from fastapi import FastAPI, Depends
4from typing import Never, NewType, Final
5from collections.abc import Generator
6from contextlib import contextmanager
7
8app = FastAPI()
9
10DATABASE = "app.db"
11
12def get_db():
13...
14
15UserID = NewType("UserID", int)
16
17class AuthenticationProvider:
18    def do_authn(self) -> "AuthenticationGuard":
19        # do something to check the auth
20        return AuthenticationGuard(guard=AuthenticationGuard._INSTANTIATION_TOKEN)  # pyright: ignore[reportPrivateUsage, reportArgumentType]
21
22class AuthenticationGuard:
23    _INSTANTIATION_TOKEN: Final = object()
24
25    def __init__(self, *, guard: Never):
26        # could also do some runtime checks to ensure caller
27        assert guard is self._INSTANTIATION_TOKEN
28
29class UncheckedUserID:
30    __value: UserID
31
32    def __init__(self, user_id: UserID):
33        self.__value = user_id
34
35    @contextmanager
36    def ensure_authn(self, guard: AuthenticationGuard) -> Generator[UserID, None, None]:
37        assert guard
38        yield self.__value
39
40
41@app.get("/balance/{user_id}")
42def get_balance(unauth_user_id: UncheckedUserID = Depends(UncheckedUserID), db: sqlite3.Connection = Depends(get_db)):
43    authguard = AuthenticationProvider().do_authn()
44    with unauth_user_id.ensure_authn(authguard) as user_id:
45        cursor = db.execute(f"SELECT balance FROM accounts WHERE user_id = {user_id}")
46        row = cursor.fetchone()
47        return {"uid": user_id, "balance": row["balance"]}

There is a fair bit more code here, but in the grand scheme of things, it would be swallowed up by your authentication and authorization code anyway.

The main change in this new code is that we never pass around basic types (like the user ID as an int in the previous examples). Input data is abstracted away into an opaque class, and you can only access it by proving that you've already done the authentication and authorization steps. By ensuring that your UncheckedUserID class cannot do anything useful (like being used in a SQL expression), your type checker will reliably error if you've forgotten to perform the authentication check to extract the 'real' user_id value.

Python, being as dynamic as it is, means you could technically bypass the check and access to the internal value without providing a valid AuthenticationGuard, but you'd hope this would be caught in PR review as undesirable code. Other languages, such as Rust, can provide much stronger guarantees, meaning that even with ugly code, you cannot directly access the internal value of the wrapper class. Obviously, the code is still running, so you can take extreme measures like dumping the instance's memory, but if you just want to ensure your auth is consistent, why would you do this at all?

In Rust

Visibility specifiers in Rust provide strong control over which code can access the wrapped user_id value, making this an even more reliable way to ensure authentication and authorization checks. In the example below, we use Axum with an extractor to ensure that the input value is correctly deserialized into our wrapper safety class.

1use axum::{response::IntoResponse, routing::get, Json, Router};
2
3mod unchecked_input {
4    use axum::extract::{FromRequestParts, Path};
5    #[derive(serde::Deserialize, serde::Serialize, Debug)]
6    pub struct UserID(u32);
7
8    #[derive(FromRequestParts)]
9    pub struct UncheckedUserID {
10        #[from_request(via(Path))]
11        value: UserID,
12    }
13
14    impl UncheckedUserID {
15        pub const fn ensure_authn(self, _: super::authentication::AuthenticationGuard) -> UserID {
16            self.value
17        }
18    }
19}
20
21mod authentication {
22    use std::marker::PhantomData;
23
24    pub struct AuthenticationGuard {
25        _internal: PhantomData<()>,
26    }
27
28    pub struct AuthenticationProvider {}
29    impl AuthenticationProvider {
30        pub const fn do_auth() -> Option<AuthenticationGuard> {
31            // do something to check the auth
32            Some(AuthenticationGuard {
33                _internal: PhantomData,
34            })
35        }
36    }
37}
38
39async fn get_balance(user_id: unchecked_input::UncheckedUserID) -> impl IntoResponse {
40    match authentication::AuthenticationProvider::do_auth() {
41        Some(guard) => {
42            let user_id = user_id.ensure_authn(guard);
43            // go and do some SQL
44            Json(serde_json::json!({"uid": user_id, "balance": 0.0}))
45        }
46        None => Json(serde_json::json!({"error": "auth failure"})),
47    }
48}
49
50#[tokio::main]
51async fn main() {
52    let app = Router::new().route("/balance/{user_id}", get(get_balance));
53
54    let listener = tokio::net::TcpListener::bind("127.0.0.1:8000")
55        .await
56        .unwrap();
57    axum::serve(listener, app).await.unwrap();
58}

Practicality

Obviously, this method for mitigating security vulnerabilities requires significant buy-in, especially in a more established software project. Coding infrastructure would be necessary to ensure these patterns are consistently followed. The two main methods to achieve this would be wrapper libraries or custom linter rules.

Take the Python case, for example. As an organization, you could require the MyOrgFastAPI library, a lightweight wrapper over FastAPI, and ensure that all API endpoints use custom types that implement these security requirements. No endpoint should be allowed to take a str argument; it must be some form of UncheckedString.

Alternatively, if you don't want to maintain a custom wrapper library, these rules could be implemented at the linter level. Whilst this does not provide feedback to a developer as early, it can still provide the same safety net by banning basic types.

Whichever way you achieve it, it would ensure that every single case is handled properly. This naturally extends the code pattern into protecting against all sorts of vulnerabilities:

• An UncheckedString must be sanitized appropriately to get out a raw string for returning to the user, mitigating XSS

• An UncheckedString cannot be concatenated onto an SQL query, mitigating classic SQL vulnerabilities

• An UncheckedString cannot be concatenated onto a command string, mitigating command injection vulnerabilities.

The list goes on.

Such patterns would also apply equally to human developers and AI agents. Asking an AI agent to ensure authentication everywhere may not be 100% reliable, but enforcing authentication at the compilation or linting stage can be extremely powerful.

Conclusion

Implementing security like this at the type level can be an incredibly powerful tool for both human developers and AI agents equally. Whilst there may be some initial overhead, especially on existing software projects, the benefits can be significant, as entire classes of vulnerabilities can be wiped out.

How you handle your data can be very project-specific, but thinking about it early can pay dividends for secure code.