OpenShell/crates/openshell-server/src/policy_store.rs

// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::persistence::{
    DraftChunkRecord, PersistenceError, PersistenceResult, PolicyRecord, SetResourceVersion, Store,
};
use crate::storage_proto::{DraftChunkPayload, PolicyRevisionPayload};
use openshell_core::proto::{NetworkPolicyRule, Sandbox, SandboxPolicy as ProtoSandboxPolicy};
use prost::Message;
use std::collections::HashMap;

#[derive(Clone, PartialEq, Message)]
struct RawPolicyRevisionPayload {
    #[prost(bytes = "vec", optional, tag = "1")]
    policy: Option<Vec<u8>>,
    #[prost(string, tag = "2")]
    hash: String,
    #[prost(string, tag = "3")]
    load_error: String,
    #[prost(int64, tag = "4")]
    loaded_at_ms: i64,
    #[prost(map = "string, string", tag = "5")]
    provenance: HashMap<String, String>,
}

#[derive(Debug, Clone)]
pub struct AtomicPolicyRevisionWrite {
    pub id: String,
    pub sandbox_id: String,
    pub workspace: String,
    pub version: i64,
    pub policy_payload: Vec<u8>,
    pub policy_hash: String,
    pub provenance: HashMap<String, String>,
    pub expected_resource_version: u64,
    pub annotations: HashMap<String, String>,
    /// Populate the create-time baseline, or replace it while startup admission
    /// is blocked and no workload has consumed the static restrictions.
    pub backfill_policy: Option<ProtoSandboxPolicy>,
}

pub fn policy_record_for_atomic_write(
    write: &AtomicPolicyRevisionWrite,
    created_at_ms: i64,
) -> PolicyRecord {
    PolicyRecord {
        id: write.id.clone(),
        sandbox_id: write.sandbox_id.clone(),
        version: write.version,
        policy_payload: write.policy_payload.clone(),
        policy_hash: write.policy_hash.clone(),
        status: "pending".to_string(),
        load_error: None,
        created_at_ms,
        loaded_at_ms: None,
        provenance: write.provenance.clone(),
    }
}

/// Only a new sandbox with durable evidence of no prior activation can replace
/// static restrictions. Legacy records are conservative: absence is not false.
pub fn permits_initial_static_policy_repair(sandbox: &Sandbox) -> bool {
    sandbox.status.as_ref().is_some_and(|status| {
        status.configuration_activated == Some(false)
            && status
                .configuration_admission
                .as_ref()
                .is_some_and(|admission| {
                    matches!(
                        openshell_core::proto::ConfigurationAdmissionState::try_from(
                            admission.state
                        ),
                        Ok(openshell_core::proto::ConfigurationAdmissionState::Pending
                            | openshell_core::proto::ConfigurationAdmissionState::Rejected)
                    )
                })
    })
}

pub fn project_policy_revision_onto_sandbox(
    write: &AtomicPolicyRevisionWrite,
    payload: &[u8],
    current_resource_version: u64,
) -> PersistenceResult<(Sandbox, bool)> {
    if write.expected_resource_version != 0
        && write.expected_resource_version != current_resource_version
    {
        return Err(PersistenceError::Conflict {
            current_resource_version: Some(current_resource_version),
        });
    }

    let payload = crate::persistence::migrate_legacy_time_fields("sandbox", payload)?;
    let mut sandbox = Sandbox::decode(payload.as_slice())
        .map_err(|e| PersistenceError::Decode(format!("decode sandbox payload failed: {e}")))?;
    sandbox.set_resource_version(current_resource_version);

    let mut changed = false;
    let startup_blocked = permits_initial_static_policy_repair(&sandbox);
    if let Some(backfill_policy) = write.backfill_policy.as_ref() {
        let spec = sandbox
            .spec
            .as_mut()
            .ok_or_else(|| PersistenceError::Decode("sandbox payload missing spec".to_string()))?;
        match spec.policy.as_ref() {
            None => {
                spec.policy = Some(backfill_policy.clone());
                changed = true;
            }
            Some(current) if current == backfill_policy => {}
            Some(_) if startup_blocked => {
                spec.policy = Some(backfill_policy.clone());
                changed = true;
            }
            Some(_) => {
                return Err(PersistenceError::Conflict {
                    current_resource_version: Some(current_resource_version),
                });
            }
        }
    }

    if !write.annotations.is_empty() {
        let metadata = sandbox.metadata.as_mut().ok_or_else(|| {
            PersistenceError::Decode("sandbox payload missing metadata".to_string())
        })?;
        for (key, value) in &write.annotations {
            if metadata.annotations.get(key) != Some(value) {
                metadata.annotations.insert(key.clone(), value.clone());
                changed = true;
            }
        }
    }

    Ok((sandbox, changed))
}

pub trait PolicyStoreExt {
    async fn put_policy_revision(
        &self,
        id: &str,
        sandbox_id: &str,
        workspace: &str,
        version: i64,
        payload: &[u8],
        hash: &str,
    ) -> PersistenceResult<()>;

    /// Atomically create a sandbox policy revision, project its annotations
    /// onto sandbox metadata, and optionally backfill `spec.policy`.
    async fn put_policy_revision_atomic(
        &self,
        write: &AtomicPolicyRevisionWrite,
    ) -> PersistenceResult<Sandbox>;

    async fn get_latest_policy(&self, sandbox_id: &str) -> PersistenceResult<Option<PolicyRecord>>;

    #[allow(dead_code)]
    async fn get_latest_loaded_policy(
        &self,
        sandbox_id: &str,
    ) -> PersistenceResult<Option<PolicyRecord>>;

    async fn get_policy_by_version(
        &self,
        sandbox_id: &str,
        version: i64,
    ) -> PersistenceResult<Option<PolicyRecord>>;

    #[allow(dead_code)]
    async fn list_policies(
        &self,
        sandbox_id: &str,
        limit: u32,
        offset: u32,
    ) -> PersistenceResult<Vec<PolicyRecord>>;

    async fn list_policies_before(
        &self,
        sandbox_id: &str,
        limit: u32,
        before_version: Option<i64>,
    ) -> PersistenceResult<Vec<PolicyRecord>>;

    async fn update_policy_status(
        &self,
        sandbox_id: &str,
        version: i64,
        status: &str,
        load_error: Option<&str>,
        loaded_at_ms: Option<i64>,
    ) -> PersistenceResult<bool>;

    async fn supersede_older_policies(
        &self,
        sandbox_id: &str,
        before_version: i64,
    ) -> PersistenceResult<u64>;

    /// Persist a draft chunk. When `dedup_key` is `Some`, duplicate inserts
    /// for the same `(sandbox, dedup_key)` fold into the existing row's
    /// `hit_count` instead of creating a second chunk — appropriate for
    /// observation-driven proposals from the mechanistic mapper. When
    /// `None`, the chunk is inserted unconditionally — appropriate for
    /// agent-authored proposals where each submission is an intentional
    /// act and the redraft loop relies on every proposal getting its own
    /// `chunk_id` even when the target endpoint is unchanged.
    ///
    /// Returns the **effective** row id. On a fresh insert that equals
    /// `chunk.id`; on dedup fold-in it is the existing row's id. Callers
    /// must use the returned id (not `chunk.id`) when reporting the chunk
    /// to clients — otherwise the response advertises an id that was never
    /// persisted.
    async fn put_draft_chunk(
        &self,
        chunk: &DraftChunkRecord,
        dedup_key: Option<&str>,
        workspace: &str,
    ) -> PersistenceResult<String>;

    async fn get_draft_chunk(&self, id: &str) -> PersistenceResult<Option<DraftChunkRecord>>;

    async fn list_draft_chunks(
        &self,
        sandbox_id: &str,
        status_filter: Option<&str>,
    ) -> PersistenceResult<Vec<DraftChunkRecord>>;

    /// Update a draft chunk's status, optionally recording a free-form
    /// `rejection_reason` for the reviewer's note. Pass `Some` only on the
    /// reject path; other status transitions pass `None` to leave any prior
    /// reason intact.
    async fn update_draft_chunk_status(
        &self,
        id: &str,
        status: &str,
        decided_at_ms: Option<i64>,
        rejection_reason: Option<&str>,
    ) -> PersistenceResult<bool>;

    async fn conditionally_reject_draft_chunk(
        &self,
        id: &str,
        decided_at_ms: i64,
        rejection_reason: &str,
    ) -> PersistenceResult<bool>;

    /// Replace the policy-dependent evaluation payload for a pending chunk.
    /// Status and observation counters remain owned by their dedicated
    /// columns and are not changed by this update.
    async fn update_draft_chunk_evaluation(
        &self,
        chunk: &DraftChunkRecord,
    ) -> PersistenceResult<bool>;

    async fn delete_draft_chunks(&self, sandbox_id: &str, status: &str) -> PersistenceResult<u64>;

    async fn get_draft_version(&self, sandbox_id: &str) -> PersistenceResult<i64>;
}

impl PolicyStoreExt for Store {
    async fn put_policy_revision(
        &self,
        id: &str,
        sandbox_id: &str,
        workspace: &str,
        version: i64,
        payload: &[u8],
        hash: &str,
    ) -> PersistenceResult<()> {
        match self {
            Self::Postgres(store) => {
                store
                    .put_policy_revision(id, sandbox_id, workspace, version, payload, hash)
                    .await
            }
            Self::Sqlite(store) => {
                store
                    .put_policy_revision(id, sandbox_id, workspace, version, payload, hash)
                    .await
            }
        }
    }

    async fn put_policy_revision_atomic(
        &self,
        write: &AtomicPolicyRevisionWrite,
    ) -> PersistenceResult<Sandbox> {
        match self {
            Self::Postgres(store) => store.put_policy_revision_atomic(write).await,
            Self::Sqlite(store) => store.put_policy_revision_atomic(write).await,
        }
    }

    async fn get_latest_policy(&self, sandbox_id: &str) -> PersistenceResult<Option<PolicyRecord>> {
        match self {
            Self::Postgres(store) => store.get_latest_policy(sandbox_id).await,
            Self::Sqlite(store) => store.get_latest_policy(sandbox_id).await,
        }
    }

    async fn get_latest_loaded_policy(
        &self,
        sandbox_id: &str,
    ) -> PersistenceResult<Option<PolicyRecord>> {
        match self {
            Self::Postgres(store) => store.get_latest_loaded_policy(sandbox_id).await,
            Self::Sqlite(store) => store.get_latest_loaded_policy(sandbox_id).await,
        }
    }

    async fn get_policy_by_version(
        &self,
        sandbox_id: &str,
        version: i64,
    ) -> PersistenceResult<Option<PolicyRecord>> {
        match self {
            Self::Postgres(store) => store.get_policy_by_version(sandbox_id, version).await,
            Self::Sqlite(store) => store.get_policy_by_version(sandbox_id, version).await,
        }
    }

    async fn list_policies(
        &self,
        sandbox_id: &str,
        limit: u32,
        offset: u32,
    ) -> PersistenceResult<Vec<PolicyRecord>> {
        match self {
            Self::Postgres(store) => store.list_policies(sandbox_id, limit, offset).await,
            Self::Sqlite(store) => store.list_policies(sandbox_id, limit, offset).await,
        }
    }

    async fn list_policies_before(
        &self,
        sandbox_id: &str,
        limit: u32,
        before_version: Option<i64>,
    ) -> PersistenceResult<Vec<PolicyRecord>> {
        match self {
            Self::Postgres(store) => {
                store
                    .list_policies_before(sandbox_id, limit, before_version)
                    .await
            }
            Self::Sqlite(store) => {
                store
                    .list_policies_before(sandbox_id, limit, before_version)
                    .await
            }
        }
    }

    async fn update_policy_status(
        &self,
        sandbox_id: &str,
        version: i64,
        status: &str,
        load_error: Option<&str>,
        loaded_at_ms: Option<i64>,
    ) -> PersistenceResult<bool> {
        match self {
            Self::Postgres(store) => {
                store
                    .update_policy_status(sandbox_id, version, status, load_error, loaded_at_ms)
                    .await
            }
            Self::Sqlite(store) => {
                store
                    .update_policy_status(sandbox_id, version, status, load_error, loaded_at_ms)
                    .await
            }
        }
    }

    async fn supersede_older_policies(
        &self,
        sandbox_id: &str,
        before_version: i64,
    ) -> PersistenceResult<u64> {
        match self {
            Self::Postgres(store) => {
                store
                    .supersede_older_policies(sandbox_id, before_version)
                    .await
            }
            Self::Sqlite(store) => {
                store
                    .supersede_older_policies(sandbox_id, before_version)
                    .await
            }
        }
    }

    async fn put_draft_chunk(
        &self,
        chunk: &DraftChunkRecord,
        dedup_key: Option<&str>,
        workspace: &str,
    ) -> PersistenceResult<String> {
        match self {
            Self::Postgres(store) => store.put_draft_chunk(chunk, dedup_key, workspace).await,
            Self::Sqlite(store) => store.put_draft_chunk(chunk, dedup_key, workspace).await,
        }
    }

    async fn get_draft_chunk(&self, id: &str) -> PersistenceResult<Option<DraftChunkRecord>> {
        match self {
            Self::Postgres(store) => store.get_draft_chunk(id).await,
            Self::Sqlite(store) => store.get_draft_chunk(id).await,
        }
    }

    async fn list_draft_chunks(
        &self,
        sandbox_id: &str,
        status_filter: Option<&str>,
    ) -> PersistenceResult<Vec<DraftChunkRecord>> {
        match self {
            Self::Postgres(store) => store.list_draft_chunks(sandbox_id, status_filter).await,
            Self::Sqlite(store) => store.list_draft_chunks(sandbox_id, status_filter).await,
        }
    }

    async fn update_draft_chunk_status(
        &self,
        id: &str,
        status: &str,
        decided_at_ms: Option<i64>,
        rejection_reason: Option<&str>,
    ) -> PersistenceResult<bool> {
        match self {
            Self::Postgres(store) => {
                store
                    .update_draft_chunk_status(id, status, decided_at_ms, rejection_reason)
                    .await
            }
            Self::Sqlite(store) => {
                store
                    .update_draft_chunk_status(id, status, decided_at_ms, rejection_reason)
                    .await
            }
        }
    }

    async fn conditionally_reject_draft_chunk(
        &self,
        id: &str,
        decided_at_ms: i64,
        rejection_reason: &str,
    ) -> PersistenceResult<bool> {
        match self {
            Self::Postgres(store) => {
                store
                    .conditionally_reject_draft_chunk(id, decided_at_ms, rejection_reason)
                    .await
            }
            Self::Sqlite(store) => {
                store
                    .conditionally_reject_draft_chunk(id, decided_at_ms, rejection_reason)
                    .await
            }
        }
    }

    async fn update_draft_chunk_evaluation(
        &self,
        chunk: &DraftChunkRecord,
    ) -> PersistenceResult<bool> {
        match self {
            Self::Postgres(store) => store.update_draft_chunk_evaluation(chunk).await,
            Self::Sqlite(store) => store.update_draft_chunk_evaluation(chunk).await,
        }
    }

    async fn delete_draft_chunks(&self, sandbox_id: &str, status: &str) -> PersistenceResult<u64> {
        match self {
            Self::Postgres(store) => store.delete_draft_chunks(sandbox_id, status).await,
            Self::Sqlite(store) => store.delete_draft_chunks(sandbox_id, status).await,
        }
    }

    async fn get_draft_version(&self, sandbox_id: &str) -> PersistenceResult<i64> {
        match self {
            Self::Postgres(store) => store.get_draft_version(sandbox_id).await,
            Self::Sqlite(store) => store.get_draft_version(sandbox_id).await,
        }
    }
}

pub fn policy_payload_from_record(record: &PolicyRecord) -> PersistenceResult<Vec<u8>> {
    ProtoSandboxPolicy::decode(record.policy_payload.as_slice())
        .map_err(|e| PersistenceError::Decode(format!("decode policy payload failed: {e}")))?;
    Ok(RawPolicyRevisionPayload {
        policy: Some(record.policy_payload.clone()),
        hash: record.policy_hash.clone(),
        load_error: record.load_error.clone().unwrap_or_default(),
        loaded_at_ms: record.loaded_at_ms.unwrap_or(0),
        provenance: record.provenance.clone(),
    }
    .encode_to_vec())
}

pub fn policy_record_from_parts(
    id: String,
    sandbox_id: String,
    version: i64,
    status: String,
    payload: &[u8],
    created_at_ms: i64,
) -> PersistenceResult<PolicyRecord> {
    let raw_wrapper = RawPolicyRevisionPayload::decode(payload)
        .map_err(|e| PersistenceError::Decode(format!("decode raw policy wrapper failed: {e}")))?;
    let wrapper = PolicyRevisionPayload::decode(payload)
        .map_err(|e| PersistenceError::Decode(format!("decode policy wrapper failed: {e}")))?;
    wrapper
        .policy
        .ok_or_else(|| PersistenceError::Decode("policy wrapper missing policy".to_string()))?;
    let policy_payload = raw_wrapper
        .policy
        .ok_or_else(|| PersistenceError::Decode("policy wrapper missing policy".to_string()))?;
    Ok(PolicyRecord {
        id,
        sandbox_id,
        version,
        policy_payload,
        policy_hash: wrapper.hash,
        status,
        load_error: if wrapper.load_error.is_empty() {
            None
        } else {
            Some(wrapper.load_error)
        },
        created_at_ms,
        loaded_at_ms: (wrapper.loaded_at_ms > 0).then_some(wrapper.loaded_at_ms),
        provenance: wrapper.provenance,
    })
}

/// Observation-mode dedup key: `host|port|binary`. Used by the mechanistic
/// mapper path where N denials targeting the same endpoint should fold into
/// one chunk instead of N near-identical chunks. Agent-authored proposals
/// pass `None` for the `dedup_key` argument to `put_draft_chunk` so each
/// proposal lands as its own chunk regardless of target — the redraft loop
/// depends on this.
pub fn observation_dedup_key(chunk: &DraftChunkRecord) -> String {
    format!("{}|{}|{}", chunk.host, chunk.port, chunk.binary)
}

pub fn draft_chunk_payload_from_record(chunk: &DraftChunkRecord) -> PersistenceResult<Vec<u8>> {
    let proposed_rule = if chunk.proposed_rule.is_empty() {
        None
    } else {
        Some(
            NetworkPolicyRule::decode(chunk.proposed_rule.as_slice())
                .map_err(|e| PersistenceError::Decode(format!("decode draft rule failed: {e}")))?,
        )
    };
    Ok(DraftChunkPayload {
        rule_name: chunk.rule_name.clone(),
        proposed_rule,
        rationale: chunk.rationale.clone(),
        security_notes: chunk.security_notes.clone(),
        #[allow(clippy::cast_possible_truncation)] // f64->f32 for confidence scores
        confidence: chunk.confidence as f32,
        decided_at_ms: chunk.decided_at_ms.unwrap_or(0),
        host: chunk.host.clone(),
        port: chunk.port,
        binary: chunk.binary.clone(),
        draft_version: chunk.draft_version,
        validation_result: chunk.validation_result.clone(),
        rejection_reason: chunk.rejection_reason.clone(),
        application_error: chunk.application_error.clone(),
        review_token: chunk.review_token.clone(),
        current_effective_policy_hash: chunk.current_effective_policy_hash.clone(),
        candidate_effective_policy_hash: chunk.candidate_effective_policy_hash.clone(),
        current_effective_policy: chunk.current_effective_policy.clone(),
        candidate_effective_policy: chunk.candidate_effective_policy.clone(),
    }
    .encode_to_vec())
}

pub fn draft_chunk_record_from_parts(
    id: String,
    sandbox_id: String,
    status: String,
    hit_count: i64,
    payload: &[u8],
    created_at_ms: i64,
    updated_at_ms: i64,
) -> PersistenceResult<DraftChunkRecord> {
    let wrapper = DraftChunkPayload::decode(payload)
        .map_err(|e| PersistenceError::Decode(format!("decode draft chunk wrapper failed: {e}")))?;
    let proposed_rule = wrapper
        .proposed_rule
        .map(|rule| rule.encode_to_vec())
        .unwrap_or_default();
    Ok(DraftChunkRecord {
        id,
        sandbox_id,
        draft_version: wrapper.draft_version,
        status,
        rule_name: wrapper.rule_name,
        proposed_rule,
        rationale: wrapper.rationale,
        security_notes: wrapper.security_notes,
        confidence: f64::from(wrapper.confidence),
        created_at_ms,
        decided_at_ms: (wrapper.decided_at_ms > 0).then_some(wrapper.decided_at_ms),
        host: wrapper.host,
        port: wrapper.port,
        binary: wrapper.binary,
        hit_count: i32::try_from(hit_count).unwrap_or(i32::MAX),
        first_seen_ms: created_at_ms,
        last_seen_ms: updated_at_ms,
        validation_result: wrapper.validation_result,
        rejection_reason: wrapper.rejection_reason,
        application_error: wrapper.application_error,
        review_token: wrapper.review_token,
        current_effective_policy_hash: wrapper.current_effective_policy_hash,
        candidate_effective_policy_hash: wrapper.candidate_effective_policy_hash,
        current_effective_policy: wrapper.current_effective_policy,
        candidate_effective_policy: wrapper.candidate_effective_policy,
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use openshell_core::proto::{
        ConfigurationAdmissionState as Admission, SandboxConfigurationAdmission, SandboxSpec,
        SandboxStatus,
    };

    #[test]
    fn static_projection_requires_durable_evidence_of_no_previous_activation() {
        let baseline = openshell_policy::restrictive_default_policy();
        let mut replacement = baseline.clone();
        replacement
            .filesystem
            .as_mut()
            .unwrap()
            .read_write
            .push("/new-static-path".to_string());
        let write = AtomicPolicyRevisionWrite {
            id: "revision".to_string(),
            sandbox_id: "sandbox".to_string(),
            workspace: "default".to_string(),
            version: 2,
            policy_payload: replacement.encode_to_vec(),
            policy_hash: String::new(),
            provenance: HashMap::new(),
            expected_resource_version: 0,
            annotations: HashMap::new(),
            backfill_policy: Some(replacement.clone()),
        };
        for activated in [None, Some(false), Some(true)] {
            for state in [Admission::Pending, Admission::Rejected, Admission::Accepted] {
                let sandbox = Sandbox {
                    spec: Some(SandboxSpec {
                        policy: Some(baseline.clone()),
                        ..Default::default()
                    }),
                    status: Some(SandboxStatus {
                        configuration_activated: activated,
                        configuration_admission: Some(SandboxConfigurationAdmission {
                            state: state.into(),
                            ..Default::default()
                        }),
                        ..Default::default()
                    }),
                    ..Default::default()
                };
                let result =
                    project_policy_revision_onto_sandbox(&write, &sandbox.encode_to_vec(), 1);
                if activated == Some(false) && state != Admission::Accepted {
                    let (projected, changed) = result.unwrap();
                    assert!(changed);
                    assert_eq!(projected.spec.unwrap().policy, Some(replacement.clone()));
                } else {
                    assert!(
                        matches!(result, Err(PersistenceError::Conflict { .. })),
                        "{activated:?} {state:?}"
                    );
                }
            }
        }
    }
}