1use std::borrow::Cow;
11use std::collections::BTreeMap;
12use std::fmt;
13use std::sync::{Arc, OnceLock};
14
15use async_trait::async_trait;
16use serde::{Deserialize, Serialize};
17use systemprompt_identifiers::UserId;
18
19#[derive(Debug, thiserror::Error)]
20#[error("invalid rule_type: {0}")]
21pub struct RuleTypeError(pub String);
22
23#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
37#[serde(transparent)]
38pub struct RuleType(Cow<'static, str>);
39
40impl RuleType {
41 pub const USER: Self = Self(Cow::Borrowed("user"));
42 pub const ROLE: Self = Self(Cow::Borrowed("role"));
43
44 pub fn extension(slug: impl Into<Cow<'static, str>>) -> Result<Self, RuleTypeError> {
45 let slug = slug.into();
46 if !is_extension_slug(slug.as_bytes()) {
47 return Err(RuleTypeError(slug.into_owned()));
48 }
49 Ok(Self(slug))
50 }
51
52 #[must_use]
53 pub const fn extension_static(slug: &'static str) -> Self {
54 assert!(
55 is_extension_slug(slug.as_bytes()),
56 "extension rule type must be [a-z0-9_], not edged by `_`, and not `user` or `role`"
57 );
58 Self(Cow::Borrowed(slug))
59 }
60
61 #[must_use]
62 pub fn as_str(&self) -> &str {
63 &self.0
64 }
65}
66
67const fn is_extension_slug(slug: &[u8]) -> bool {
68 if matches!(slug, [] | [b'_', ..] | [.., b'_'])
69 || bytes_eq(slug, b"user")
70 || bytes_eq(slug, b"role")
71 {
72 return false;
73 }
74 let mut rest = slug;
75 while let [byte, tail @ ..] = rest {
76 if !(byte.is_ascii_lowercase() || byte.is_ascii_digit() || *byte == b'_') {
77 return false;
78 }
79 rest = tail;
80 }
81 true
82}
83
84const fn bytes_eq(left: &[u8], right: &[u8]) -> bool {
85 if left.len() != right.len() {
86 return false;
87 }
88 let (mut left, mut right) = (left, right);
89 while let ([a, left_tail @ ..], [b, right_tail @ ..]) = (left, right) {
90 if *a != *b {
91 return false;
92 }
93 left = left_tail;
94 right = right_tail;
95 }
96 true
97}
98
99impl fmt::Display for RuleType {
100 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
101 f.write_str(self.as_str())
102 }
103}
104
105impl From<&str> for RuleType {
106 fn from(s: &str) -> Self {
107 match s {
108 "user" => Self::USER,
109 "role" => Self::ROLE,
110 other => Self(Cow::Owned(other.to_owned())),
111 }
112 }
113}
114
115#[cfg(feature = "sqlx")]
116impl sqlx::Type<sqlx::Postgres> for RuleType {
117 fn type_info() -> <sqlx::Postgres as sqlx::Database>::TypeInfo {
118 <str as sqlx::Type<sqlx::Postgres>>::type_info()
119 }
120
121 fn compatible(ty: &<sqlx::Postgres as sqlx::Database>::TypeInfo) -> bool {
122 <str as sqlx::Type<sqlx::Postgres>>::compatible(ty)
123 }
124}
125
126#[cfg(feature = "sqlx")]
127impl<'q> sqlx::Encode<'q, sqlx::Postgres> for RuleType {
128 fn encode_by_ref(
129 &self,
130 buf: &mut <sqlx::Postgres as sqlx::Database>::ArgumentBuffer,
131 ) -> Result<sqlx::encode::IsNull, sqlx::error::BoxDynError> {
132 <&str as sqlx::Encode<'q, sqlx::Postgres>>::encode(self.as_str(), buf)
133 }
134}
135
136#[cfg(feature = "sqlx")]
137impl<'r> sqlx::Decode<'r, sqlx::Postgres> for RuleType {
138 fn decode(
139 value: <sqlx::Postgres as sqlx::Database>::ValueRef<'r>,
140 ) -> Result<Self, sqlx::error::BoxDynError> {
141 let raw = <&str as sqlx::Decode<'r, sqlx::Postgres>>::decode(value)?;
142 Ok(Self::from(raw))
143 }
144}
145
146#[derive(Debug, Clone, PartialEq, Eq)]
151pub struct SubjectDimension {
152 pub rule_type: RuleType,
153 pub label: &'static str,
154 pub precedence: u16,
155}
156
157#[derive(Debug, Clone, Default, PartialEq, Eq)]
162pub struct SubjectAttributes(BTreeMap<RuleType, Vec<String>>);
163
164impl SubjectAttributes {
165 pub const EMPTY: Self = Self(BTreeMap::new());
166
167 #[must_use]
168 pub const fn new() -> Self {
169 Self::EMPTY
170 }
171
172 pub fn insert(&mut self, rule_type: RuleType, values: Vec<String>) {
173 self.0.insert(rule_type, values);
174 }
175
176 #[must_use]
177 pub fn values(&self, rule_type: &RuleType) -> &[String] {
178 self.0.get(rule_type).map_or(&[], Vec::as_slice)
179 }
180
181 #[must_use]
182 pub fn is_empty(&self) -> bool {
183 self.0.is_empty()
184 }
185}
186
187impl FromIterator<(RuleType, Vec<String>)> for SubjectAttributes {
188 fn from_iter<I: IntoIterator<Item = (RuleType, Vec<String>)>>(iter: I) -> Self {
189 Self(iter.into_iter().collect())
190 }
191}
192
193#[derive(Debug, thiserror::Error)]
194#[non_exhaustive]
195pub enum SubjectSourceError {
196 #[error("the subject attribute source is not ready")]
197 NotReady,
198 #[error("no subject dimension '{0}' is registered")]
199 UnknownDimension(String),
200 #[error("subject dimension '{dimension}' could not be read")]
201 Provider {
202 dimension: String,
203 #[source]
204 source: Box<dyn std::error::Error + Send + Sync + 'static>,
205 },
206}
207
208#[async_trait]
218pub trait SubjectAttributeSource: Send + Sync {
219 fn dimensions(&self) -> &[SubjectDimension];
220
221 async fn values_for(
222 &self,
223 user: &UserId,
224 dimension: &SubjectDimension,
225 ) -> Result<Vec<String>, SubjectSourceError>;
226
227 async fn resolve_all(&self, user: &UserId) -> Result<SubjectAttributes, SubjectSourceError>;
228
229 async fn known_values(
230 &self,
231 _dimension: &SubjectDimension,
232 ) -> Result<Option<Vec<String>>, SubjectSourceError> {
233 Ok(None)
234 }
235}
236
237#[derive(Debug, Clone, Copy, thiserror::Error)]
238#[error("the subject attribute source is already bound")]
239pub struct AlreadyBound;
240
241#[derive(Clone, Default)]
249pub struct SubjectSourceHandle(Arc<OnceLock<Arc<dyn SubjectAttributeSource>>>);
250
251impl fmt::Debug for SubjectSourceHandle {
252 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
253 f.debug_struct("SubjectSourceHandle")
254 .field("bound", &self.0.get().is_some())
255 .finish()
256 }
257}
258
259impl SubjectSourceHandle {
260 pub fn bind(&self, source: Arc<dyn SubjectAttributeSource>) -> Result<(), AlreadyBound> {
261 self.0.set(source).map_err(|_unbound| AlreadyBound)
262 }
263}
264
265#[async_trait]
266impl SubjectAttributeSource for SubjectSourceHandle {
267 fn dimensions(&self) -> &[SubjectDimension] {
268 self.0.get().map_or(&[], |source| source.dimensions())
269 }
270
271 async fn values_for(
272 &self,
273 user: &UserId,
274 dimension: &SubjectDimension,
275 ) -> Result<Vec<String>, SubjectSourceError> {
276 match self.0.get() {
277 Some(source) => source.values_for(user, dimension).await,
278 None => Err(SubjectSourceError::NotReady),
279 }
280 }
281
282 async fn resolve_all(&self, user: &UserId) -> Result<SubjectAttributes, SubjectSourceError> {
283 match self.0.get() {
284 Some(source) => source.resolve_all(user).await,
285 None => Err(SubjectSourceError::NotReady),
286 }
287 }
288
289 async fn known_values(
290 &self,
291 dimension: &SubjectDimension,
292 ) -> Result<Option<Vec<String>>, SubjectSourceError> {
293 match self.0.get() {
294 Some(source) => source.known_values(dimension).await,
295 None => Err(SubjectSourceError::NotReady),
296 }
297 }
298}