A better PIR algorithm

In the previous PIR algorithm, the user had to send a query vector of the same size as the database itself. Can we do better than that? Yes!

Let's look at how to reduce the user's query size by making use of matrix-vector product and dot product.

How will our improved algorithm work?

Instead of representing the database as a vector of length n, we'll represent it as a matrix of dimension sqrt(n) by sqrt(n). As prior, everything in the database is unencrypted.

Let's say Alice wants to retrieve item_1,2 from the database. This time, Alice will send two query vectors to the server; each query vector is of length sqrt(n) so Alice's communication cost will be reduced from n to 2 · sqrt(n).

Information retrieval (without privacy)

What goes in the query vectors?

Query Vector #1: Used in the matrix-vector product. The query vector will have a 0 in every place, except for the column number Alice is interested in. Since Alice wants the 2nd column, the vector takes the following form: [0, 1, 0, ..., 0].
Query Vector #2: Used in the dot product. This query vector will have a 0 in every place, except for the row number Alice is interested in. Since Alice wants the 1st row, the vector takes the following form: [1, 0, ..., 0].

When we take the matrix-vector product of the database with query vector #1, we get a vector back. What's in this vector?

[item_1,2, item_2,2, item_3,2, ..., item_sqrt(n),2]

When we take the dot product of the previous result with query vector #2, we get the item Alice is interested in. Why? Well:

[item_1,2, item_2,2, item_3,2, ..., item_sqrt(n),2] · [1, 0, ..., 0]^t = item_1,2

Making the query private

Since Alice doesn't want to reveal to the server which item she's interested in, she encrypts her two query vectors with respect to her FHE public key.

Query Vector #1: [Enc(0), Enc(1), Enc(0), ..., Enc(0)]
Query Vector #2: [Enc(1), Enc(0), ..., Enc(0)]

When the server takes the matrix-vector product of the database with encrypted query vector #1, we get:

[Enc(item_1,2), Enc(item_2,2), Enc(item_3,2), ..., Enc(item_sqrt(n),2)]

When the server takes the dot product of the above vector with encrypted query vector #2, voila, we get Alice's desired item:

[Enc(item_1,2), Enc(item_2,2), Enc(item_3,2), ..., Enc(item_sqrt(n),2)] · [Enc(1), Enc(0), ..., Enc(0)]^t = Enc(item_1,2)

Program walkthrough

Our database will have 100 items in it. We'll use an array to represent our database.

Setup

#![allow(unused)]
fn main() {
use sunscreen::{
    fhe_program,
    types::{bfv::Signed, Cipher},
};

const SQRT_DATABASE_SIZE: usize = 10;

#[fhe_program(scheme = "bfv")]
/// This program takes a user's query and looks up the entry in the database.
/// Queries are arrays containing a single 1 element at the
/// desired item's index and 0s elsewhere.
fn lookup(
    col_query: [Cipher<Signed>; SQRT_DATABASE_SIZE],
    row_query: [Cipher<Signed>; SQRT_DATABASE_SIZE],
    database: [[Signed; SQRT_DATABASE_SIZE]; SQRT_DATABASE_SIZE],
) -> Cipher<Signed> {
    // Copy col_query just so we get an initialized object of the right
    // type in col.
    let mut col = [col_query[0]; SQRT_DATABASE_SIZE];

    // Perform matrix-vector multiplication with col_query to extract
    // Alice's desired column
    for i in 0..SQRT_DATABASE_SIZE {
        for j in 0..SQRT_DATABASE_SIZE {
            if j == 0 {
                col[i] = database[i][j] * col_query[j];
            } else {
                col[i] = col[i] + database[i][j] * col_query[j];
            }
        }
    }

    let mut sum = col[0] * row_query[0];

    // Dot product the result with the row query to get the result
    for i in 1..SQRT_DATABASE_SIZE {
        sum = sum + col[i] * row_query[i];
    }

    sum
}
}

We begin by importing the stuff we're going to use.

We declare our lookup function as an FHE program with the appropriate attribute (#[fhe_program(scheme = "bfv")]).

We'll represent our database as an array of length SQRT_DATABASE_SIZE (in this case = 10) with entries that are arrays of length SQRT_DATABASE_SIZE (also = 10). Since we want to write into the database, we'll need to initialize the array in the FHE program's function body.

lookup implements the matrix-vector and dot product formulas explained above to retrieve Alice's item privately. It takes in an unencrypted database along with Alice's encrypted query "vectors" (actually arrays). Recall that we have two queries— col_query will be used to select the appropriate column of the database whereas row_query will be used to select the appropriate row of the database.

Alice

#![allow(unused)]
fn main() {
use sunscreen::{
    FheRuntime, PrivateKey, PublicKey,
};

/// Alice is a party that wants to look up a value in the database without
/// revealing what she looked up.
struct Alice {
    /// Alice's public key
    pub public_key: PublicKey,

    /// Alice's private key
    private_key: PrivateKey,

    /// Alice's runtime
    runtime: FheRuntime,
}
}

Alice wants to retrieve an item from the database privately. She'll need a public/private key pair to do this (since she needs to encrypt her query with respect to her public key).

#![allow(unused)]
fn main() {
use sunscreen::{
    fhe_program,
    types::{bfv::Signed, Cipher},
    Ciphertext, CompiledFheProgram, Compiler, Error, FheProgramInput, FheRuntime,
    Params, PrivateKey, PublicKey,
};

const SQRT_DATABASE_SIZE: usize = 10;

/// Alice is a party that wants to look up a value in the database without
/// revealing what she looked up.
struct Alice {
    /// Alice's public key
    pub public_key: PublicKey,

    /// Alice's private key
    private_key: PrivateKey,

    /// Alice's runtime
    runtime: FheRuntime,
}

impl Alice {
    pub fn setup(params: &Params) -> Result<Alice, Error> {
        let runtime = FheRuntime::new(params)?;

        let (public_key, private_key) = runtime.generate_keys()?;

        Ok(Alice {
            public_key,
            private_key,
            runtime,
        })
    }

    pub fn create_query(&self, index: usize) -> Result<(Ciphertext, Ciphertext), Error> {
        let col = index % SQRT_DATABASE_SIZE;
        let row = index / SQRT_DATABASE_SIZE;

        let mut col_query = [Signed::from(0); SQRT_DATABASE_SIZE];
        let mut row_query = [Signed::from(0); SQRT_DATABASE_SIZE];
        col_query[col] = Signed::from(1);
        row_query[row] = Signed::from(1);

        Ok((
            self.runtime.encrypt(col_query, &self.public_key)?,
            self.runtime.encrypt(row_query, &self.public_key)?,
        ))
    }

    pub fn check_response(&self, value: Ciphertext) -> Result<(), Error> {
        let value: Signed = self.runtime.decrypt(&value, &self.private_key)?;

        let value: i64 = value.into();

        println!("Alice received {}", value);

        Ok(())
    }
}
}

Alice will need to construct a runtime. Once that's done, she can generate her public/private key pair.

create_query will allow Alice to create and encrypt her two query "vectors" (i.e. arrays). Alice will pass in an index which contains her desired column and row indices. Notice the 1's place of the index will allow Alice to select her desired column #, whereas the 10's place will allow Alice to select her desired row # (e.g. if index is 85, this denotes Alice is interested in entry located in the 5th column, 8th row).

We won't use this until the very end but check_response allows Alice to decrypt the server's response by passing in the ciphertext she received (value) along with her private_key.

Server

Let's look at the server next.

#![allow(unused)]
fn main() {
use sunscreen::{
    fhe_program,
    types::{bfv::Signed, Cipher},
    Ciphertext, CompiledFheProgram, Compiler, Error, FheProgramInput, FheRuntime,
    Params, PrivateKey, PublicKey,
};

/// This is the server that processes Alice's query.
struct Server {
    /// The compiled database query program
    pub compiled_lookup: CompiledFheProgram,

    /// The server's runtime
    runtime: FheRuntime,
}
}

Recall that the server is responsible for retrieving Alice's item from the database; thus, it will have to run the compiled lookup program (compiled_lookup).

#![allow(unused)]
fn main() {
use sunscreen::{
    fhe_program,
    types::{bfv::Signed, Cipher},
    Ciphertext, CompiledFheProgram, Compiler, Error, FheProgramInput, FheRuntime,
    Params, PrivateKey, PublicKey,
};

#[fhe_program(scheme = "bfv")]
/// This program takes a user's query and looks up the entry in the database.
/// Queries are arrays containing a single 1 element at the
/// desired item's index and 0s elsewhere.
fn lookup(
    col_query: [Cipher<Signed>; SQRT_DATABASE_SIZE],
    row_query: [Cipher<Signed>; SQRT_DATABASE_SIZE],
    database: [[Signed; SQRT_DATABASE_SIZE]; SQRT_DATABASE_SIZE],
) -> Cipher<Signed> {
    // Copy col_query just so we get an initialized object of the right
    // type in col.
    let mut col = [col_query[0]; SQRT_DATABASE_SIZE];

    // Perform matrix-vector multiplication with col_query to extract
    // Alice's desired column
    for i in 0..SQRT_DATABASE_SIZE {
        for j in 0..SQRT_DATABASE_SIZE {
            if j == 0 {
                col[i] = database[i][j] * col_query[j];
            } else {
                col[i] = col[i] + database[i][j] * col_query[j];
            }
        }
    }

    let mut sum = col[0] * row_query[0];

    // Dot product the result with the row query to get the result
    for i in 1..SQRT_DATABASE_SIZE {
        sum = sum + col[i] * row_query[i];
    }

    sum
}

const SQRT_DATABASE_SIZE: usize = 10;

/// This is the server that processes Alice's query.
struct Server {
    /// The compiled database query program
    pub compiled_lookup: CompiledFheProgram,

    /// The server's runtime
    runtime: FheRuntime,
}

impl Server {
    pub fn setup() -> Result<Server, Error> {
        let app = Compiler::new()
            .fhe_program(lookup)
            .compile()?;

        let runtime = FheRuntime::new(app.params())?;

        Ok(Server {
            compiled_lookup: app.get_fhe_program(lookup).unwrap().clone(),
            runtime,
        })
    }

    pub fn run_query(
        &self,
        col_query: Ciphertext,
        row_query: Ciphertext,
        public_key: &PublicKey,
    ) -> Result<Ciphertext, Error> {
        // Our database will consist of values between 400 and 500.
        let mut database = [[Signed::from(0); SQRT_DATABASE_SIZE]; SQRT_DATABASE_SIZE];
        let mut val = Signed::from(400);

        for i in 0..SQRT_DATABASE_SIZE {
            for j in 0..SQRT_DATABASE_SIZE {
                database[i][j] = val;
                val = val + 1;
            }
        }

        let args: Vec<FheProgramInput> = vec![col_query.into(), row_query.into(), database.into()];

        let results = self.runtime.run(&self.compiled_lookup, args, public_key)?;

        Ok(results[0].clone())
    }
}
}

The server constructs a runtime so that it can run the compiled FHE program compiled_lookup.

Using run_query, the server can return an (encrypted) response to Alice's query.

The items in the database will be the integers from 400 to 499, stored in ascending order. Recall that lookup takes in three arguments---the encrypted query for the column index (col_query), the encrypted query for the row index (row_query), and the unencrypted database. Unfortunately, we'll need to do some type conversion for the database as Sunscreen's compiler needs entries of the Signed type not i64 for its programs.

Additionally, to run FHE programs, we need to pass in arguments as a vec. Thus, we create a vec called args that contains our encrypted queries and unencrypted database (which now has Signed entries rather than i64 entries in it).

Once all that's done, the server can run the FHE program by passing in the compiled_lookup program, the arguments to the program args (now contained in a vec), and Alice's public_key.

Retrieving the item privately

use sunscreen::{
    fhe_program,
    types::{bfv::Signed, Cipher},
    Ciphertext, CompiledFheProgram, Compiler, Error, FheProgramInput, FheRuntime,
    Params, PrivateKey, PublicKey,
};

const SQRT_DATABASE_SIZE: usize = 10;

#[fhe_program(scheme = "bfv")]
/// This program takes a user's query and looks up the entry in the database.
/// Queries are arrays containing a single 1 element at the
/// desired item's index and 0s elsewhere.
fn lookup(
    col_query: [Cipher<Signed>; SQRT_DATABASE_SIZE],
    row_query: [Cipher<Signed>; SQRT_DATABASE_SIZE],
    database: [[Signed; SQRT_DATABASE_SIZE]; SQRT_DATABASE_SIZE],
) -> Cipher<Signed> {
    // Copy col_query just so we get an initialized object of the right
    // type in col.
    let mut col = [col_query[0]; SQRT_DATABASE_SIZE];

    // Perform matrix-vector multiplication with col_query to extract
    // Alice's desired column
    for i in 0..SQRT_DATABASE_SIZE {
        for j in 0..SQRT_DATABASE_SIZE {
            if j == 0 {
                col[i] = database[i][j] * col_query[j];
            } else {
                col[i] = col[i] + database[i][j] * col_query[j];
            }
        }
    }

    let mut sum = col[0] * row_query[0];

    // Dot product the result with the row query to get the result
    for i in 1..SQRT_DATABASE_SIZE {
        sum = sum + col[i] * row_query[i];
    }

    sum
}

/// Alice is a party that wants to look up a value in the database without
/// revealing what she looked up.
struct Alice {
    /// Alice's public key
    pub public_key: PublicKey,

    /// Alice's private key
    private_key: PrivateKey,

    /// Alice's runtime
    runtime: FheRuntime,
}

impl Alice {
    pub fn setup(params: &Params) -> Result<Alice, Error> {
        let runtime = FheRuntime::new(params)?;

        let (public_key, private_key) = runtime.generate_keys()?;

        Ok(Alice {
            public_key,
            private_key,
            runtime,
        })
    }

    pub fn create_query(&self, index: usize) -> Result<(Ciphertext, Ciphertext), Error> {
        let col = index % SQRT_DATABASE_SIZE;
        let row = index / SQRT_DATABASE_SIZE;

        let mut col_query = [Signed::from(0); SQRT_DATABASE_SIZE];
        let mut row_query = [Signed::from(0); SQRT_DATABASE_SIZE];
        col_query[col] = Signed::from(1);
        row_query[row] = Signed::from(1);

        Ok((
            self.runtime.encrypt(col_query, &self.public_key)?,
            self.runtime.encrypt(row_query, &self.public_key)?,
        ))
    }

    pub fn check_response(&self, value: Ciphertext) -> Result<(), Error> {
        let value: Signed = self.runtime.decrypt(&value, &self.private_key)?;

        let value: i64 = value.into();

        println!("Alice received {}", value);

        Ok(())
    }
}

/// This is the server that processes Alice's query.
struct Server {
    /// The compiled database query program
    pub compiled_lookup: CompiledFheProgram,

    /// The server's runtime
    runtime: FheRuntime,
}

impl Server {
    pub fn setup() -> Result<Server, Error> {
        let app = Compiler::new()
            .fhe_program(lookup)
            .compile()?;

        let runtime = FheRuntime::new(app.params())?;

        Ok(Server {
            compiled_lookup: app.get_fhe_program(lookup).unwrap().clone(),
            runtime,
        })
    }

    pub fn run_query(
        &self,
        col_query: Ciphertext,
        row_query: Ciphertext,
        public_key: &PublicKey,
    ) -> Result<Ciphertext, Error> {
        // Our database will consist of values between 400 and 500.
        let mut database = [[Signed::from(0); SQRT_DATABASE_SIZE]; SQRT_DATABASE_SIZE];
        let mut val = Signed::from(400);

        for i in 0..SQRT_DATABASE_SIZE {
            for j in 0..SQRT_DATABASE_SIZE {
                database[i][j] = val;
                val = val + 1;
            }
        }

        let args: Vec<FheProgramInput> = vec![col_query.into(), row_query.into(), database.into()];

        let results = self.runtime.run(&self.compiled_lookup, args, public_key)?;

        Ok(results[0].clone())
    }
}
fn main() -> Result<(), Error> {
    // Set up the database
    let server = Server::setup()?;

    // Alice sets herself up. The FHE scheme parameters are public to the
    // protocol, so Alice has them.
    let alice = Alice::setup(&server.compiled_lookup.metadata.params)?;

    let (col_query, row_query) = alice.create_query(94)?;

    let response = server.run_query(col_query, row_query, &alice.public_key)?;

    alice.check_response(response)?;

    Ok(())
}

We set up the server first and then Alice (notice that Alice relies on parameters generated from the server's setup). Both of them must use the same set of FHE scheme parameters for compatibility. In deployment, these values would likely be fixed at the protocol level.

Alice would like to privately retrieve the item at the 94th "position" (this will mean the entry located in the 4th row, 9th column) from the database so she calls create_query. create_query encrypts her query value of 94 properly (i.e. creates col_query and row_query which has encryptions of 0s and 1s in the appropriate places).

The server calls run_query to privately retrieve Alice's desired item to her. It passes in Alice's encrypted queries along with Alice's public key (alice.public_key).

Finally, Alice can decrypt to check what item she received via check_response.

Performance

The entire program (not including compilation time) takes ~376 ms on an Intel Xeon @ 3.0 GHz (with 8 cores and 16 GB RAM) and ~716 ms on a Macbook Air M1.

What's missing?

If you're interested in using PIR in a real application, there are much better PIR algorithms out there (e.g. SealPIR, Spiral) that are faster and compress the query vector further.

Additionally, in PIR, we assume that the user knows the index of the item she wants to retrieve. For many applications though, this might not be the case. Oblivious message detection allows the user to detect which item she's interested in and can be combined with PIR.

Sunscreen BFV Documentation