affine_core 0.0.3

use serde::{Deserialize, Serialize};
use thiserror::Error;

use crate::blob_access::SourceIdentity;

pub const INVALIDATION_CHANNEL_V1: &str = "affine:backend-runtime:invalidation:v1";
const INVALIDATION_WIRE_VERSION: u32 = 1;

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub enum SubjectId {
  User(String),
  Workspace(String),
}

impl SubjectId {
  fn id(&self) -> &str {
    match self {
      Self::User(id) | Self::Workspace(id) => id,
    }
  }

  fn encode(&self) -> Result<String, InvalidationCodecError> {
    if self.id().is_empty() {
      return Err(InvalidationCodecError::InvalidSubject);
    }
    Ok(match self {
      Self::User(id) => format!("user:{id}"),
      Self::Workspace(id) => format!("workspace:{id}"),
    })
  }

  fn decode(value: String) -> Result<Self, InvalidationCodecError> {
    match value.split_once(':') {
      Some(("user", id)) if !id.is_empty() => Ok(Self::User(id.to_string())),
      Some(("workspace", id)) if !id.is_empty() => Ok(Self::Workspace(id.to_string())),
      _ => Err(InvalidationCodecError::InvalidSubject),
    }
  }
}

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub enum QuotaCacheKey {
  Entitlement(SubjectId),
  OwnerMapping(String),
  StorageUsage(SubjectId),
  SeatUsage(String),
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum InvalidationHintV1 {
  QuotaEntitlement { subject: SubjectId },
  QuotaOwnerMapping { workspace_id: String },
  QuotaStorageUsage { subject: SubjectId },
  QuotaSeatUsage { workspace_id: String },
  BlobSource { source: SourceIdentity },
}

impl InvalidationHintV1 {
  pub fn quota_key(&self) -> Option<QuotaCacheKey> {
    match self {
      Self::QuotaEntitlement { subject } => Some(QuotaCacheKey::Entitlement(subject.clone())),
      Self::QuotaOwnerMapping { workspace_id } => Some(QuotaCacheKey::OwnerMapping(workspace_id.clone())),
      Self::QuotaStorageUsage { subject } => Some(QuotaCacheKey::StorageUsage(subject.clone())),
      Self::QuotaSeatUsage { workspace_id } => Some(QuotaCacheKey::SeatUsage(workspace_id.clone())),
      Self::BlobSource { .. } => None,
    }
  }
}

#[derive(Debug, Error)]
pub enum InvalidationCodecError {
  #[error("invalid invalidation payload")]
  InvalidPayload(#[from] serde_json::Error),
  #[error("unsupported invalidation version")]
  UnsupportedVersion,
  #[error("invalid invalidation subject")]
  InvalidSubject,
  #[error("invalid blob source")]
  InvalidSource,
  #[error("invalidation id must not be empty")]
  EmptyId,
}

#[derive(Deserialize, Serialize)]
struct Envelope {
  version: u32,
  #[serde(flatten)]
  hint: WireHintV1,
}

#[derive(Deserialize, Serialize)]
#[serde(
  tag = "kind",
  rename_all = "camelCase",
  rename_all_fields = "camelCase",
  deny_unknown_fields
)]
enum WireHintV1 {
  QuotaEntitlement { subject: String },
  QuotaOwnerMapping { workspace_id: String },
  QuotaStorageUsage { subject: String },
  QuotaSeatUsage { workspace_id: String },
  BlobSource { source: SourceIdentity },
}

pub fn encode_invalidation_v1(hint: &InvalidationHintV1) -> Result<Vec<u8>, InvalidationCodecError> {
  let wire = match hint {
    InvalidationHintV1::QuotaEntitlement { subject } => WireHintV1::QuotaEntitlement {
      subject: subject.encode()?,
    },
    InvalidationHintV1::QuotaOwnerMapping { workspace_id } => {
      require_id(workspace_id)?;
      WireHintV1::QuotaOwnerMapping {
        workspace_id: workspace_id.clone(),
      }
    }
    InvalidationHintV1::QuotaStorageUsage { subject } => WireHintV1::QuotaStorageUsage {
      subject: subject.encode()?,
    },
    InvalidationHintV1::QuotaSeatUsage { workspace_id } => {
      require_id(workspace_id)?;
      WireHintV1::QuotaSeatUsage {
        workspace_id: workspace_id.clone(),
      }
    }
    InvalidationHintV1::BlobSource { source } => {
      source.validate().map_err(|_| InvalidationCodecError::InvalidSource)?;
      WireHintV1::BlobSource { source: source.clone() }
    }
  };
  Ok(serde_json::to_vec(&Envelope {
    version: INVALIDATION_WIRE_VERSION,
    hint: wire,
  })?)
}

pub fn decode_invalidation_v1(payload: &[u8]) -> Result<InvalidationHintV1, InvalidationCodecError> {
  let envelope: Envelope = serde_json::from_slice(payload)?;
  if envelope.version != INVALIDATION_WIRE_VERSION {
    return Err(InvalidationCodecError::UnsupportedVersion);
  }
  Ok(match envelope.hint {
    WireHintV1::QuotaEntitlement { subject } => InvalidationHintV1::QuotaEntitlement {
      subject: SubjectId::decode(subject)?,
    },
    WireHintV1::QuotaOwnerMapping { workspace_id } => {
      require_id(&workspace_id)?;
      InvalidationHintV1::QuotaOwnerMapping { workspace_id }
    }
    WireHintV1::QuotaStorageUsage { subject } => InvalidationHintV1::QuotaStorageUsage {
      subject: SubjectId::decode(subject)?,
    },
    WireHintV1::QuotaSeatUsage { workspace_id } => {
      require_id(&workspace_id)?;
      InvalidationHintV1::QuotaSeatUsage { workspace_id }
    }
    WireHintV1::BlobSource { source } => {
      source.validate().map_err(|_| InvalidationCodecError::InvalidSource)?;
      InvalidationHintV1::BlobSource { source }
    }
  })
}

fn require_id(id: &str) -> Result<(), InvalidationCodecError> {
  if id.is_empty() {
    Err(InvalidationCodecError::EmptyId)
  } else {
    Ok(())
  }
}

#[cfg(test)]
#[path = "tests/invalidation/tests.rs"]
mod tests;