geth/crates/geth-keychain/src/lib.rs

use geth_types::{AgentId, AuthOpId, DeviceId, KeyId, NodeId, UnixMillis, UserId};
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, BTreeSet};

pub const KEYCHAIN_SIGNATURE_NAMESPACE: &str = "geth.keychain.v1@geth.local";

pub type SignedKeychainOp = geth_codec::SignedEnvelope<KeychainOp, KeyId>;

pub fn keychain_signing_payload(op: &KeychainOp) -> Result<Vec<u8>, geth_codec::CodecError> {
    geth_codec::signing_payload(KEYCHAIN_SIGNATURE_NAMESPACE, op)
}

pub fn keychain_signing_payload_hash(
    op: &KeychainOp,
) -> Result<geth_types::BlobHash, geth_codec::CodecError> {
    geth_codec::signing_payload_hash(KEYCHAIN_SIGNATURE_NAMESPACE, op)
}

#[must_use]
pub fn signed_keychain_op(op: KeychainOp, signer: KeyId, signature: Vec<u8>) -> SignedKeychainOp {
    geth_codec::SignedEnvelope::new(KEYCHAIN_SIGNATURE_NAMESPACE, op, signer, signature)
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct KeychainOp {
    pub id: geth_types::AuthOpId,
    pub created_at: UnixMillis,
    pub kind: KeychainOpKind,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct KeychainOpSignature {
    pub op_id: AuthOpId,
    pub signer: KeyId,
    pub namespace: String,
    pub signature: Vec<u8>,
    pub created_at: UnixMillis,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "kebab-case")]
pub enum KeychainOpKind {
    KeychainInit,
    AdminKeyAdd {
        key: KeyId,
    },
    AdminKeyRevoke {
        key: KeyId,
    },
    UserAdd {
        user: UserId,
        name: String,
    },
    UserRename {
        user: UserId,
        name: String,
    },
    UserRevoke {
        user: UserId,
    },
    DeviceAdd {
        device: DeviceId,
        user: UserId,
    },
    DeviceRevoke {
        device: DeviceId,
    },
    DeviceKeyAdd {
        device: DeviceId,
        key: KeyId,
    },
    DeviceKeyRevoke {
        device: DeviceId,
        key: KeyId,
    },
    NodeAdd {
        node: NodeId,
        device: DeviceId,
        name: String,
    },
    NodeRename {
        node: NodeId,
        name: String,
    },
    NodeRevoke {
        node: NodeId,
    },
    NodeEndpointAdd {
        node: NodeId,
        endpoint: String,
    },
    NodeEndpointRevoke {
        node: NodeId,
        endpoint: String,
    },
    AgentBind {
        agent: AgentId,
        node: NodeId,
    },
}

#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct KeychainView {
    pub initialized: bool,
    pub admin_keys: Vec<KeyId>,
    pub users: BTreeMap<UserId, UserRecord>,
    pub devices: BTreeMap<DeviceId, DeviceRecord>,
    pub nodes: BTreeMap<NodeId, NodeRecord>,
    pub agents: BTreeMap<AgentId, NodeId>,
    pub endpoints: BTreeMap<String, NodeId>,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct UserRecord {
    pub id: UserId,
    pub name: String,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct DeviceRecord {
    pub id: DeviceId,
    pub user: UserId,
    pub keys: Vec<KeyId>,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct NodeRecord {
    pub id: NodeId,
    pub device: DeviceId,
    pub name: String,
    pub endpoints: Vec<String>,
}

pub fn reduce_keychain_ops(ops: &[KeychainOp]) -> KeychainView {
    let mut initialized = false;
    let mut admin_keys = BTreeSet::new();
    let mut users = BTreeMap::<UserId, UserRecord>::new();
    let mut revoked_users = BTreeSet::new();
    let mut devices = BTreeMap::<DeviceId, DeviceRecord>::new();
    let mut revoked_devices = BTreeSet::new();
    let mut nodes = BTreeMap::<NodeId, NodeRecord>::new();
    let mut revoked_nodes = BTreeSet::new();
    let mut agent_bindings = BTreeMap::<AgentId, NodeId>::new();

    for op in ops {
        match &op.kind {
            KeychainOpKind::KeychainInit => initialized = true,
            KeychainOpKind::AdminKeyAdd { key } => {
                admin_keys.insert(key.clone());
            }
            KeychainOpKind::AdminKeyRevoke { key } => {
                admin_keys.remove(key);
            }
            KeychainOpKind::UserAdd { user, name } => {
                revoked_users.remove(user);
                users.entry(user.clone()).or_insert_with(|| UserRecord {
                    id: user.clone(),
                    name: name.clone(),
                });
            }
            KeychainOpKind::UserRename { user, name } => {
                if let Some(record) = users.get_mut(user) {
                    record.name.clone_from(name);
                }
            }
            KeychainOpKind::UserRevoke { user } => {
                revoked_users.insert(user.clone());
            }
            KeychainOpKind::DeviceAdd { device, user } => {
                revoked_devices.remove(device);
                devices
                    .entry(device.clone())
                    .or_insert_with(|| DeviceRecord {
                        id: device.clone(),
                        user: user.clone(),
                        keys: Vec::new(),
                    });
            }
            KeychainOpKind::DeviceRevoke { device } => {
                revoked_devices.insert(device.clone());
            }
            KeychainOpKind::DeviceKeyAdd { device, key } => {
                if let Some(record) = devices.get_mut(device) {
                    if !record.keys.contains(key) {
                        record.keys.push(key.clone());
                        record.keys.sort();
                    }
                }
            }
            KeychainOpKind::DeviceKeyRevoke { device, key } => {
                if let Some(record) = devices.get_mut(device) {
                    record.keys.retain(|item| item != key);
                }
            }
            KeychainOpKind::NodeAdd { node, device, name } => {
                revoked_nodes.remove(node);
                nodes.entry(node.clone()).or_insert_with(|| NodeRecord {
                    id: node.clone(),
                    device: device.clone(),
                    name: name.clone(),
                    endpoints: Vec::new(),
                });
            }
            KeychainOpKind::NodeRename { node, name } => {
                if let Some(record) = nodes.get_mut(node) {
                    record.name.clone_from(name);
                }
            }
            KeychainOpKind::NodeRevoke { node } => {
                revoked_nodes.insert(node.clone());
            }
            KeychainOpKind::NodeEndpointAdd { node, endpoint } => {
                if let Some(record) = nodes.get_mut(node) {
                    if !record.endpoints.contains(endpoint) {
                        record.endpoints.push(endpoint.clone());
                        record.endpoints.sort();
                    }
                }
            }
            KeychainOpKind::NodeEndpointRevoke { node, endpoint } => {
                if let Some(record) = nodes.get_mut(node) {
                    record.endpoints.retain(|item| item != endpoint);
                }
            }
            KeychainOpKind::AgentBind { agent, node } => {
                agent_bindings.insert(agent.clone(), node.clone());
            }
        }
    }

    users.retain(|user, _| !revoked_users.contains(user));
    devices.retain(|device, record| {
        !revoked_devices.contains(device) && users.contains_key(&record.user)
    });
    nodes.retain(|node, record| {
        !revoked_nodes.contains(node) && devices.contains_key(&record.device)
    });
    agent_bindings.retain(|_, node| nodes.contains_key(node));

    let mut endpoints = BTreeMap::new();
    for (node, record) in &nodes {
        for endpoint in &record.endpoints {
            endpoints.insert(endpoint.clone(), node.clone());
        }
    }

    KeychainView {
        initialized,
        admin_keys: admin_keys.into_iter().collect(),
        users,
        devices,
        nodes,
        agents: agent_bindings,
        endpoints,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn keychain_structs_roundtrip() {
        let op = KeychainOp {
            id: "op:1".into(),
            created_at: UnixMillis(1),
            kind: KeychainOpKind::UserAdd {
                user: "user:eric".into(),
                name: "Eric".to_owned(),
            },
        };
        let json = serde_json::to_string(&op).expect("json");
        let decoded: KeychainOp = serde_json::from_str(&json).expect("decode");
        assert_eq!(decoded, op);
    }

    #[test]
    fn keychain_signing_payload_is_canonical_and_namespaced() {
        let op = KeychainOp {
            id: "op:1".into(),
            created_at: UnixMillis(1),
            kind: KeychainOpKind::AdminKeyAdd {
                key: "key:admin".into(),
            },
        };
        assert_eq!(
            keychain_signing_payload(&op).expect("payload"),
            keychain_signing_payload(&op).expect("payload again")
        );
        assert_ne!(
            keychain_signing_payload_hash(&op).expect("hash"),
            geth_codec::hash_canonical(&op).expect("raw op hash")
        );

        let signed = signed_keychain_op(op.clone(), "key:admin".into(), vec![1, 2, 3]);
        assert_eq!(signed.namespace(), KEYCHAIN_SIGNATURE_NAMESPACE);
        assert_eq!(signed.payload(), &op);
    }

    fn op(sequence: i64, kind: KeychainOpKind) -> KeychainOp {
        KeychainOp {
            id: format!("op:{sequence}").into(),
            created_at: UnixMillis(sequence),
            kind,
        }
    }

    #[test]
    fn reducer_tracks_active_keychain_view() {
        let ops = vec![
            op(1, KeychainOpKind::KeychainInit),
            op(
                2,
                KeychainOpKind::AdminKeyAdd {
                    key: "key:admin-a".into(),
                },
            ),
            op(
                3,
                KeychainOpKind::AdminKeyAdd {
                    key: "key:admin-b".into(),
                },
            ),
            op(
                4,
                KeychainOpKind::AdminKeyRevoke {
                    key: "key:admin-b".into(),
                },
            ),
            op(
                5,
                KeychainOpKind::UserAdd {
                    user: "user:eric".into(),
                    name: "Eric".to_owned(),
                },
            ),
            op(
                6,
                KeychainOpKind::UserRename {
                    user: "user:eric".into(),
                    name: "Eric Updated".to_owned(),
                },
            ),
            op(
                7,
                KeychainOpKind::DeviceAdd {
                    device: "device:laptop".into(),
                    user: "user:eric".into(),
                },
            ),
            op(
                8,
                KeychainOpKind::DeviceKeyAdd {
                    device: "device:laptop".into(),
                    key: "key:device-a".into(),
                },
            ),
            op(
                9,
                KeychainOpKind::DeviceKeyAdd {
                    device: "device:laptop".into(),
                    key: "key:device-b".into(),
                },
            ),
            op(
                10,
                KeychainOpKind::DeviceKeyRevoke {
                    device: "device:laptop".into(),
                    key: "key:device-b".into(),
                },
            ),
            op(
                11,
                KeychainOpKind::NodeAdd {
                    node: "node:laptop".into(),
                    device: "device:laptop".into(),
                    name: "laptop".to_owned(),
                },
            ),
            op(
                12,
                KeychainOpKind::NodeRename {
                    node: "node:laptop".into(),
                    name: "work-laptop".to_owned(),
                },
            ),
            op(
                13,
                KeychainOpKind::NodeEndpointAdd {
                    node: "node:laptop".into(),
                    endpoint: "endpoint:old".to_owned(),
                },
            ),
            op(
                14,
                KeychainOpKind::NodeEndpointAdd {
                    node: "node:laptop".into(),
                    endpoint: "endpoint:new".to_owned(),
                },
            ),
            op(
                15,
                KeychainOpKind::NodeEndpointRevoke {
                    node: "node:laptop".into(),
                    endpoint: "endpoint:old".to_owned(),
                },
            ),
            op(
                16,
                KeychainOpKind::AgentBind {
                    agent: "agent:daemon".into(),
                    node: "node:laptop".into(),
                },
            ),
        ];

        let view = reduce_keychain_ops(&ops);

        assert!(view.initialized);
        assert_eq!(view.admin_keys, vec![KeyId::from("key:admin-a")]);
        assert_eq!(
            view.users.get(&UserId::from("user:eric")),
            Some(&UserRecord {
                id: "user:eric".into(),
                name: "Eric Updated".to_owned(),
            })
        );
        assert_eq!(
            view.devices.get(&DeviceId::from("device:laptop")),
            Some(&DeviceRecord {
                id: "device:laptop".into(),
                user: "user:eric".into(),
                keys: vec!["key:device-a".into()],
            })
        );
        assert_eq!(
            view.nodes.get(&NodeId::from("node:laptop")),
            Some(&NodeRecord {
                id: "node:laptop".into(),
                device: "device:laptop".into(),
                name: "work-laptop".to_owned(),
                endpoints: vec!["endpoint:new".to_owned()],
            })
        );
        assert_eq!(
            view.agents.get(&AgentId::from("agent:daemon")),
            Some(&NodeId::from("node:laptop"))
        );
        assert_eq!(
            view.endpoints.get("endpoint:new"),
            Some(&NodeId::from("node:laptop"))
        );
        assert!(!view.endpoints.contains_key("endpoint:old"));
    }

    #[test]
    fn reducer_excludes_revoked_identity_subtrees() {
        let ops = vec![
            op(
                1,
                KeychainOpKind::UserAdd {
                    user: "user:eric".into(),
                    name: "Eric".to_owned(),
                },
            ),
            op(
                2,
                KeychainOpKind::DeviceAdd {
                    device: "device:laptop".into(),
                    user: "user:eric".into(),
                },
            ),
            op(
                3,
                KeychainOpKind::NodeAdd {
                    node: "node:laptop".into(),
                    device: "device:laptop".into(),
                    name: "laptop".to_owned(),
                },
            ),
            op(
                4,
                KeychainOpKind::NodeEndpointAdd {
                    node: "node:laptop".into(),
                    endpoint: "endpoint:live".to_owned(),
                },
            ),
            op(
                5,
                KeychainOpKind::AgentBind {
                    agent: "agent:daemon".into(),
                    node: "node:laptop".into(),
                },
            ),
            op(
                6,
                KeychainOpKind::UserRevoke {
                    user: "user:eric".into(),
                },
            ),
        ];

        let view = reduce_keychain_ops(&ops);

        assert!(view.users.is_empty());
        assert!(view.devices.is_empty());
        assert!(view.nodes.is_empty());
        assert!(view.agents.is_empty());
        assert!(view.endpoints.is_empty());
    }
}