affine_core 0.0.2

use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use sha2::{Digest, Sha256};

use super::{AccessGrant, QuotaUsage, evaluate_workspace_quota};

pub const STORAGE_RESERVATION_MINUTES: i32 = 60;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum StorageIntegrityError {
  MetadataMismatch,
  ChecksumMismatch,
}

pub fn validate_storage_object(
  resource: StorageResource,
  key: &str,
  body: &[u8],
  declared_size: Option<i64>,
  declared_mime: Option<&str>,
  content_length: i64,
  content_type: &str,
) -> Result<(), StorageIntegrityError> {
  if declared_size != Some(content_length)
    || declared_mime != Some(content_type)
    || i64::try_from(body.len()).ok() != declared_size
  {
    return Err(StorageIntegrityError::MetadataMismatch);
  }
  if resource == StorageResource::Blob && URL_SAFE_NO_PAD.encode(Sha256::digest(body)) != key.trim_end_matches('=') {
    return Err(StorageIntegrityError::ChecksumMismatch);
  }
  Ok(())
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum StorageResource {
  Blob,
  CommentAttachment,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct StorageReservationRequest {
  pub resource: StorageResource,
  pub size: i64,
  pub has_key: bool,
  pub has_doc_id: bool,
  pub has_name: bool,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct StorageAuthorizationFacts {
  pub workspace_upload: bool,
  pub doc_read: bool,
  pub doc_comment_create: bool,
}

impl StorageAuthorizationFacts {
  pub const fn allows_upload(self, resource: StorageResource) -> bool {
    match resource {
      StorageResource::Blob => self.workspace_upload,
      StorageResource::CommentAttachment => self.doc_read && self.doc_comment_create,
    }
  }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum StorageObjectFacts {
  NotChecked,
  Missing,
  Matching,
  MetadataMismatch,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum StorageLedgerFacts {
  Missing,
  Pending {
    live: bool,
    size_matches: bool,
    mime_matches: bool,
    owner_matches: bool,
  },
  Completed {
    size_matches: bool,
    mime_matches: bool,
    owner_matches: bool,
    object: StorageObjectFacts,
  },
  Deleted,
  InvalidStatus,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct StorageReservationFacts {
  pub authorization: StorageAuthorizationFacts,
  pub grant: AccessGrant,
  pub usage: QuotaUsage,
  pub ledger: StorageLedgerFacts,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum StorageMutationIntent {
  Create,
  Resume,
  ReplaceExpired,
  RepairMissingObject,
  AlreadyUploaded,
  InspectCompletedObject,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum StorageReservationDenial {
  Forbidden,
  WorkspaceReadonly,
  BlobLimitExceeded { limit: i64 },
  StorageLimitExceeded { limit: i64 },
  ArithmeticOverflow,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum StorageReservationPlan {
  Mutate {
    intent: StorageMutationIntent,
    requested_increment: i64,
  },
  Deny(StorageReservationDenial),
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum StorageReservationError {
  InvalidRequest,
  MissingAttachmentFields,
  Deleted,
  SizeMismatch,
  MimeMismatch,
  OwnerMismatch,
  InvalidLedgerStatus,
  ObjectMetadataMismatch,
}

pub fn plan_storage_reservation(
  request: StorageReservationRequest,
  facts: StorageReservationFacts,
) -> Result<StorageReservationPlan, StorageReservationError> {
  if request.size <= 0 || !request.has_key {
    return Err(StorageReservationError::InvalidRequest);
  }
  if request.resource == StorageResource::CommentAttachment && (!request.has_doc_id || !request.has_name) {
    return Err(StorageReservationError::MissingAttachmentFields);
  }
  if !facts.authorization.allows_upload(request.resource) {
    return Ok(StorageReservationPlan::Deny(StorageReservationDenial::Forbidden));
  }
  if request.size > facts.grant.limits.blob_limit {
    return Ok(StorageReservationPlan::Deny(
      StorageReservationDenial::BlobLimitExceeded {
        limit: facts.grant.limits.blob_limit,
      },
    ));
  }
  let workspace_readonly = !evaluate_workspace_quota(&facts.grant, facts.usage)
    .readonly_reasons
    .is_empty();

  let (intent, requested_increment) = match facts.ledger {
    StorageLedgerFacts::Missing => (StorageMutationIntent::Create, request.size),
    StorageLedgerFacts::Deleted => return Err(StorageReservationError::Deleted),
    StorageLedgerFacts::InvalidStatus => return Err(StorageReservationError::InvalidLedgerStatus),
    StorageLedgerFacts::Pending {
      live,
      size_matches,
      mime_matches,
      owner_matches,
    } => {
      validate_metadata(size_matches, mime_matches, owner_matches)?;
      if live {
        if workspace_readonly {
          return Ok(StorageReservationPlan::Deny(
            StorageReservationDenial::WorkspaceReadonly,
          ));
        }
        return Ok(StorageReservationPlan::Mutate {
          intent: StorageMutationIntent::Resume,
          requested_increment: 0,
        });
      }
      (StorageMutationIntent::ReplaceExpired, request.size)
    }
    StorageLedgerFacts::Completed {
      size_matches,
      mime_matches,
      owner_matches,
      object,
    } => {
      validate_metadata(size_matches, mime_matches, owner_matches)?;
      match object {
        StorageObjectFacts::NotChecked => {
          return Ok(StorageReservationPlan::Mutate {
            intent: StorageMutationIntent::InspectCompletedObject,
            requested_increment: 0,
          });
        }
        StorageObjectFacts::Matching => {
          return Ok(StorageReservationPlan::Mutate {
            intent: StorageMutationIntent::AlreadyUploaded,
            requested_increment: 0,
          });
        }
        StorageObjectFacts::Missing => (StorageMutationIntent::RepairMissingObject, 0),
        StorageObjectFacts::MetadataMismatch => return Err(StorageReservationError::ObjectMetadataMismatch),
      }
    }
  };
  let Some(next_storage_bytes) = facts.usage.storage_bytes.checked_add(requested_increment) else {
    return Ok(StorageReservationPlan::Deny(
      StorageReservationDenial::ArithmeticOverflow,
    ));
  };
  if workspace_readonly {
    return Ok(StorageReservationPlan::Deny(
      StorageReservationDenial::WorkspaceReadonly,
    ));
  }
  match next_storage_bytes {
    next if next <= facts.grant.limits.storage_quota => Ok(StorageReservationPlan::Mutate {
      intent,
      requested_increment,
    }),
    _ => Ok(StorageReservationPlan::Deny(
      StorageReservationDenial::StorageLimitExceeded {
        limit: facts.grant.limits.storage_quota,
      },
    )),
  }
}

fn validate_metadata(
  size_matches: bool,
  mime_matches: bool,
  owner_matches: bool,
) -> Result<(), StorageReservationError> {
  if !size_matches {
    Err(StorageReservationError::SizeMismatch)
  } else if !mime_matches {
    Err(StorageReservationError::MimeMismatch)
  } else if !owner_matches {
    Err(StorageReservationError::OwnerMismatch)
  } else {
    Ok(())
  }
}

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