support for root key ids

- `BiscuitBuilder` now supports setting the root key id
- `Biscuit.from_bytes` and `Biscuit.from_base64` now take either a public key or a function returning a public key, allowing key selection based on the root key id.

biscuit_test.py

@@ -171,6 +171,37 @@ def test_authorizer_builder():
 allow if true;
 """
 
+def test_key_selection():
+    private_key = PrivateKey.from_hex("473b5189232f3f597b5c2f3f9b0d5e28b1ee4e7cce67ec6b7fbf5984157a6b97")
+    root = KeyPair.from_private_key(private_key)
+    other_root = KeyPair()
+
+    def choose_key(kid):
+        if kid is None:
+            return other_root.public_key
+        elif kid == 1:
+            return root.public_key
+        else:
+            raise Exception("Unknown key identifier") 
+
+    biscuit_builder0 = BiscuitBuilder("user({id})", { 'id': "1234" })
+    token0 = biscuit_builder0.build(other_root.private_key).to_base64()
+    biscuit_builder1 = BiscuitBuilder("user({id})", { 'id': "1234" })
+    biscuit_builder1.set_root_key_id(1)
+    token1 = biscuit_builder1.build(private_key).to_base64()
+    biscuit_builder2 = BiscuitBuilder("user({id})", { 'id': "1234" })
+    biscuit_builder2.set_root_key_id(2)
+    token2 = biscuit_builder2.build(private_key).to_base64()
+
+
+    Biscuit.from_base64(token0, choose_key)
+    Biscuit.from_base64(token0, other_root.public_key)
+    Biscuit.from_base64(token1, choose_key)
+    try:
+      Biscuit.from_base64(token2, choose_key)
+    except:
+      pass
+
 def test_complete_lifecycle():
     private_key = PrivateKey.from_hex("473b5189232f3f597b5c2f3f9b0d5e28b1ee4e7cce67ec6b7fbf5984157a6b97")
     root = KeyPair.from_private_key(private_key)

docs/basic-use.rst

@@ -57,6 +57,21 @@ Parse and authorize a biscuit token
 >>> authorizer.authorize()
 0
 
+In order to help with key rotation, biscuit tokens can optionally carry a root key identifier, helping the verifying party choose between several valid public keys.
+
+>>> def public_key_fn(kid):
+...   if kid is None:
+...     return PublicKey.from_hex("9e124fbb46ff99a87219aef4b09f4f6c3b7fd96b7bd279e38af3ef429a101c69")
+...   elif kid == 1:
+...     return PublicKey.from_hex("1d211ddaf521cc45b620431817ba4fe0457be467ba4d724ecf514db3070b53cc")
+...   else:
+...     raise Exception("unknown key identifier")
+>>> token = Biscuit.from_base64("CAESfQoTCgQxMjM0GAMiCQoHCAoSAxiACBIkCAASII5WVsvM52T91C12wnzButmyzmtGSX_rbM6hCSIJihX2GkDwAcVxTnY8aeMLm-i2R_VzTfIMQZya49ogXO2h2Fg2TJsDcG3udIki9il5PA05lKUwrfPNroS7Qg5e04AyLLcHIiIKII5rh75jrCrgE6Rzw6GVYczMn1IOo287uO4Ef5wp7obY", public_key_fn)
+>>> authorizer = Authorizer( """ time({now}); allow if user($u); """, { 'now': datetime.now(tz = timezone.utc)} )
+>>> authorizer.add_token(token)
+>>> authorizer.authorize()
+0
+
 Query an authorizer
 -------------------
 

src/lib.rs

@@ -1,5 +1,6 @@
 // There seem to be false positives with pyo3
 #![allow(clippy::borrow_deref_ref)]
+use ::biscuit_auth::RootKeyProvider;
 use chrono::DateTime;
 use chrono::TimeZone;
 use chrono::Utc;
@@ -41,6 +42,33 @@ create_exception!(
     pyo3::exceptions::PyException
 );
 
+struct PyKeyProvider {
+    py_value: PyObject,
+}
+
+impl RootKeyProvider for PyKeyProvider {
+    fn choose(&self, kid: Option<u32>) -> Result<PublicKey, error::Format> {
+        Python::with_gil(|py| {
+            if self.py_value.as_ref(py).is_callable() {
+                let result = self
+                    .py_value
+                    .call1(py, (kid,))
+                    .map_err(|_| error::Format::UnknownPublicKey)?;
+                let py_pk: PyPublicKey = result
+                    .extract(py)
+                    .map_err(|_| error::Format::UnknownPublicKey)?;
+                Ok(py_pk.0)
+            } else {
+                let py_pk: PyPublicKey = self
+                    .py_value
+                    .extract(py)
+                    .map_err(|_| error::Format::UnknownPublicKey)?;
+                Ok(py_pk.0)
+            }
+        })
+    }
+}
+
 /// Builder class allowing to create a biscuit from a datalog block
 ///
 /// :param source: a datalog snippet
@@ -141,6 +169,10 @@ impl PyBiscuitBuilder {
         self.0.merge(builder.0.clone())
     }
 
+    pub fn set_root_key_id(&mut self, root_key_id: u32) {
+        self.0.set_root_key_id(root_key_id)
+    }
+
     fn __repr__(&self) -> String {
         self.0.to_string()
     }
@@ -162,21 +194,27 @@ impl PyBiscuit {
 
     /// Deserializes a token from raw data
     ///
-    /// This will check the signature using the root key
+    /// This will check the signature using the provided root key (or function)
+    ///
+    /// :param data: a (url-safe) base64-encoded string
+    /// :param root: either a public key or a function taking an integer (or `None`) and returning an public key
     #[classmethod]
-    pub fn from_bytes(_: &PyType, data: &[u8], root: &PyPublicKey) -> PyResult<PyBiscuit> {
-        match Biscuit::from(data, root.0) {
+    pub fn from_bytes(_: &PyType, data: &[u8], root: PyObject) -> PyResult<PyBiscuit> {
+        match Biscuit::from(data, PyKeyProvider { py_value: root }) {
             Ok(biscuit) => Ok(PyBiscuit(biscuit)),
             Err(error) => Err(BiscuitValidationError::new_err(error.to_string())),
         }
     }
 
     /// Deserializes a token from URL safe base 64 data
     ///
-    /// This will check the signature using the root key
+    /// This will check the signature using the provided root key (or function)
+    ///
+    /// :param data: a (url-safe) base64-encoded string
+    /// :param root: either a public key or a function taking an integer (or `None`) and returning an public key
     #[classmethod]
-    pub fn from_base64(_: &PyType, data: &str, root: &PyPublicKey) -> PyResult<PyBiscuit> {
-        match Biscuit::from_base64(data, root.0) {
+    pub fn from_base64(_: &PyType, data: &str, root: PyObject) -> PyResult<PyBiscuit> {
+        match Biscuit::from_base64(data, PyKeyProvider { py_value: root }) {
             Ok(biscuit) => Ok(PyBiscuit(biscuit)),
             Err(error) => Err(BiscuitValidationError::new_err(error.to_string())),
         }