affine_core 0.0.5

AFFiNE primitive core.
Documentation
use super::{
  AccessContext, AccessDiagnostic, AccessGrant, Deployment, EntitlementFact, EntitlementSource, EntitlementStatus,
  LicenseError, LicenseVerifier, Limits, MAX_SEAT_QUANTITY, Plan, Rights, TargetType,
};

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct EntitlementInput {
  pub deployment: Deployment,
  pub target_type: TargetType,
  pub plan: Option<Plan>,
  pub quantity: Option<i32>,
  pub signed: bool,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ValidatedEntitlement {
  Catalog(PlanAccess),
  SignedLicense,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct PlanAccess {
  pub plan: Plan,
  pub rights: Rights,
  pub limits: Limits,
  pub quantity: Option<i32>,
}

impl From<PlanAccess> for AccessGrant {
  fn from(access: PlanAccess) -> Self {
    let mut limits = access.limits;
    if access.rights.unlimited_copilot {
      limits.copilot_action_limit = None;
    }
    Self {
      plan: if access.plan == Plan::Ai {
        Plan::Free
      } else {
        access.plan
      },
      rights: access.rights,
      limits,
      quantity: access.quantity,
      expires_at: None,
      diagnostics: Vec::new(),
    }
  }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum EntitlementInputError {
  InvalidQuantity,
  SignedFieldsConflict,
  SignedTargetMismatch,
  SelfHostedCommercialRequiresSignature,
  PlanTargetMismatch,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum EntitlementTransitionEffect {
  RevokePendingInvitations,
  DemoteNonOwnerAdmins,
}

#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct EntitlementTransitionPlan {
  pub effects: Vec<EntitlementTransitionEffect>,
}

pub fn validate_entitlement_input(input: EntitlementInput) -> Result<ValidatedEntitlement, EntitlementInputError> {
  validate_quantity(input.quantity)?;
  if input.signed {
    if input.plan.is_some() || input.quantity.is_some() {
      return Err(EntitlementInputError::SignedFieldsConflict);
    }
    if input.deployment != Deployment::SelfHosted || input.target_type != TargetType::Workspace {
      return Err(EntitlementInputError::SignedTargetMismatch);
    }
    return Ok(ValidatedEntitlement::SignedLicense);
  }

  let plan = input.plan.unwrap_or(match input.deployment {
    Deployment::Cloud => Plan::Free,
    Deployment::SelfHosted => Plan::SelfHostedFree,
  });
  if input.deployment == Deployment::SelfHosted && plan != Plan::SelfHostedFree {
    return Err(EntitlementInputError::SelfHostedCommercialRequiresSignature);
  }
  let eligible = match (input.deployment, input.target_type, plan) {
    (Deployment::Cloud, TargetType::User, Plan::Free | Plan::Pro | Plan::LifetimePro | Plan::Ai)
    | (Deployment::Cloud, TargetType::Workspace, Plan::Free | Plan::Team)
    | (Deployment::Cloud, TargetType::Instance, Plan::Free)
    | (Deployment::SelfHosted, _, Plan::SelfHostedFree) => true,
    _ => false,
  };
  if !eligible {
    return Err(EntitlementInputError::PlanTargetMismatch);
  }
  describe_plan(plan, input.quantity).map(ValidatedEntitlement::Catalog)
}

pub fn describe_plan(plan: Plan, quantity: Option<i32>) -> Result<PlanAccess, EntitlementInputError> {
  validate_quantity(quantity)?;
  let quantity = if matches!(plan, Plan::Team | Plan::SelfHostedTeam) {
    Some(quantity.unwrap_or(1))
  } else {
    None
  };
  Ok(PlanAccess {
    plan,
    rights: Rights::for_plan(plan),
    limits: Limits::for_plan(plan, quantity),
    quantity,
  })
}

pub fn plan_entitlement_transition(previous: &AccessGrant, next: &AccessGrant) -> EntitlementTransitionPlan {
  if previous.rights.commercial && !next.rights.commercial {
    EntitlementTransitionPlan {
      effects: vec![
        EntitlementTransitionEffect::RevokePendingInvitations,
        EntitlementTransitionEffect::DemoteNonOwnerAdmins,
      ],
    }
  } else {
    EntitlementTransitionPlan::default()
  }
}

fn validate_quantity(quantity: Option<i32>) -> Result<(), EntitlementInputError> {
  if quantity.is_some_and(|quantity| !(1..=MAX_SEAT_QUANTITY).contains(&quantity)) {
    return Err(EntitlementInputError::InvalidQuantity);
  }
  Ok(())
}

pub fn resolve_entitlements(context: &AccessContext<'_>, facts: &[EntitlementFact<'_>]) -> AccessGrant {
  let mut grant = AccessGrant::fallback(context.deployment);
  let mut base: Option<(
    EntitlementStatus,
    Plan,
    Option<i32>,
    Option<chrono::DateTime<chrono::Utc>>,
  )> = None;
  let mut unlimited_copilot = false;

  for fact in facts {
    let Some(source) = EntitlementSource::parse(fact.source) else {
      grant.diagnostics.push(AccessDiagnostic::InvalidSource);
      continue;
    };
    let Some(status) = EntitlementStatus::parse(fact.status) else {
      grant.diagnostics.push(AccessDiagnostic::InvalidStatus);
      continue;
    };
    if fact.starts_at.is_some_and(|starts_at| starts_at > context.now)
      || match status {
        EntitlementStatus::Active => fact.expires_at.is_some_and(|expires_at| context.now >= expires_at),
        EntitlementStatus::Grace => fact.grace_until.is_none_or(|grace_until| context.now >= grace_until),
        EntitlementStatus::Revoked | EntitlementStatus::Expired => true,
      }
    {
      grant.diagnostics.push(AccessDiagnostic::Inactive);
      continue;
    }

    let candidate = match context.deployment {
      Deployment::Cloud => cloud_candidate(context.target_type, source, fact),
      Deployment::SelfHosted => selfhost_candidate(context, source, fact),
    };
    let candidate = match candidate {
      Ok(candidate) => candidate,
      Err(diagnostic) => {
        grant.diagnostics.push(diagnostic);
        continue;
      }
    };
    let (plan, quantity, expires_at) = candidate;
    if plan == Plan::Ai {
      unlimited_copilot = true;
      grant.rights.copilot_byok = true;
      continue;
    }
    let should_replace = base.is_none_or(|(current_status, current_plan, _, _)| {
      (status.priority(), plan.priority()) > (current_status.priority(), current_plan.priority())
    });
    if should_replace {
      base = Some((status, plan, quantity, expires_at));
    }
  }

  if let Some((_, plan, quantity, expires_at)) = base {
    let access: AccessGrant = describe_plan(plan, quantity)
      .expect("resolved entitlement input is validated")
      .into();
    grant.plan = access.plan;
    grant.quantity = access.quantity;
    grant.expires_at = expires_at;
    grant.limits = access.limits;
    grant.rights.commercial = access.rights.commercial;
    grant.rights.copilot_byok = access.rights.copilot_byok;
  }
  grant.rights.unlimited_copilot = unlimited_copilot;
  if unlimited_copilot {
    grant.limits.copilot_action_limit = None;
  }
  grant
}

fn cloud_candidate(
  target_type: TargetType,
  source: EntitlementSource,
  fact: &EntitlementFact<'_>,
) -> Result<(Plan, Option<i32>, Option<chrono::DateTime<chrono::Utc>>), AccessDiagnostic> {
  if !matches!(
    source,
    EntitlementSource::CloudSubscription | EntitlementSource::AdminGrant
  ) {
    return Err(AccessDiagnostic::IneligibleTarget);
  }
  let plan = Plan::parse(fact.plan).ok_or(AccessDiagnostic::InvalidPlan)?;
  let validated = validate_entitlement_input(EntitlementInput {
    deployment: Deployment::Cloud,
    target_type,
    plan: Some(plan),
    quantity: fact.quantity,
    signed: false,
  })
  .map_err(|error| match error {
    EntitlementInputError::InvalidQuantity => AccessDiagnostic::InvalidQuantity,
    _ => AccessDiagnostic::IneligibleTarget,
  })?;
  let ValidatedEntitlement::Catalog(access) = validated else {
    unreachable!("cloud entitlement validation cannot return signed license")
  };
  Ok((access.plan, access.quantity, fact.expires_at))
}

fn selfhost_candidate(
  context: &AccessContext<'_>,
  source: EntitlementSource,
  fact: &EntitlementFact<'_>,
) -> Result<(Plan, Option<i32>, Option<chrono::DateTime<chrono::Utc>>), AccessDiagnostic> {
  if context.target_type != TargetType::Workspace || source != EntitlementSource::SelfHostedLicense {
    return Err(AccessDiagnostic::IneligibleTarget);
  }
  if Plan::parse(fact.plan) != Some(Plan::SelfHostedTeam) {
    return Err(AccessDiagnostic::InvalidPlan);
  }
  let payload = fact
    .signed_payload
    .ok_or(AccessDiagnostic::InvalidLicense(LicenseError::InvalidEnvelope))?;
  let public_key = context
    .license_public_key
    .ok_or(AccessDiagnostic::InvalidLicense(LicenseError::InvalidPublicKey))?;
  let claims = LicenseVerifier::verify(payload, public_key, context.workspace_id, context.now)
    .map_err(AccessDiagnostic::InvalidLicense)?;
  Ok((
    Plan::SelfHostedTeam,
    Some(claims.seat_quantity()),
    Some(claims.expires_at()),
  ))
}

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