mz_repr/
role_id.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10use std::fmt;
11use std::mem::size_of;
12use std::str::FromStr;
13
14use anyhow::{Error, anyhow};
15use columnation::{Columnation, CopyRegion};
16use mz_lowertest::MzReflect;
17use mz_proto::{RustType, TryFromProtoError};
18use proptest_derive::Arbitrary;
19use serde::{Deserialize, Serialize};
20
21include!(concat!(env!("OUT_DIR"), "/mz_repr.role_id.rs"));
22
23const SYSTEM_CHAR: char = 's';
24const SYSTEM_BYTE: u8 = b's';
25const PREDEFINED_CHAR: char = 'g';
26const PREDEFINED_BYTE: u8 = b'g';
27const USER_CHAR: char = 'u';
28const USER_BYTE: u8 = b'u';
29const PUBLIC_CHAR: char = 'p';
30const PUBLIC_BYTE: u8 = b'p';
31
32/// The identifier for a role.
33#[derive(
34    Arbitrary,
35    Clone,
36    Copy,
37    Debug,
38    Eq,
39    PartialEq,
40    Ord,
41    PartialOrd,
42    Hash,
43    Serialize,
44    Deserialize,
45    MzReflect,
46)]
47pub enum RoleId {
48    System(u64),
49    /// Like system roles, these are roles built into the system. However, they are grantable to
50    /// users and provide access to certain, commonly needed, privileged capabilities and
51    /// information (modelled after <https://www.postgresql.org/docs/16/predefined-roles.html>).
52    Predefined(u64),
53    User(u64),
54    Public,
55}
56
57impl RoleId {
58    pub fn is_system(&self) -> bool {
59        matches!(self, Self::System(_))
60    }
61
62    pub fn is_user(&self) -> bool {
63        matches!(self, Self::User(_))
64    }
65
66    pub fn is_public(&self) -> bool {
67        matches!(self, Self::Public)
68    }
69
70    pub fn is_predefined(&self) -> bool {
71        matches!(self, Self::Predefined(_))
72    }
73
74    pub fn is_builtin(&self) -> bool {
75        self.is_public() || self.is_system() || self.is_predefined()
76    }
77
78    pub fn encode_binary(&self) -> Vec<u8> {
79        let mut res = Vec::with_capacity(Self::binary_size());
80        match self {
81            RoleId::System(id) => {
82                res.push(SYSTEM_BYTE);
83                res.extend_from_slice(&id.to_le_bytes());
84            }
85            RoleId::Predefined(id) => {
86                res.push(PREDEFINED_BYTE);
87                res.extend_from_slice(&id.to_le_bytes());
88            }
89            RoleId::User(id) => {
90                res.push(USER_BYTE);
91                res.extend_from_slice(&id.to_le_bytes());
92            }
93            RoleId::Public => {
94                res.push(PUBLIC_BYTE);
95                res.extend_from_slice(&0_u64.to_le_bytes());
96            }
97        }
98        res
99    }
100
101    pub fn decode_binary(raw: &[u8]) -> Result<RoleId, Error> {
102        if raw.len() != RoleId::binary_size() {
103            return Err(anyhow!(
104                "invalid binary size, expecting {}, found {}",
105                RoleId::binary_size(),
106                raw.len()
107            ));
108        }
109
110        let variant = raw[0];
111        let id = u64::from_le_bytes(raw[1..].try_into()?);
112
113        match variant {
114            SYSTEM_BYTE => Ok(RoleId::System(id)),
115            PREDEFINED_BYTE => Ok(RoleId::Predefined(id)),
116            USER_BYTE => Ok(RoleId::User(id)),
117            PUBLIC_BYTE => Ok(RoleId::Public),
118            _ => Err(anyhow!("unrecognized role id variant byte '{variant}'")),
119        }
120    }
121
122    pub const fn binary_size() -> usize {
123        1 + size_of::<u64>()
124    }
125}
126
127impl FromStr for RoleId {
128    type Err = Error;
129
130    fn from_str(s: &str) -> Result<Self, Self::Err> {
131        fn parse_u64(s: &str) -> Result<u64, Error> {
132            if s.len() < 2 {
133                return Err(anyhow!("couldn't parse role id '{s}'"));
134            }
135            s[1..]
136                .parse()
137                .map_err(|_| anyhow!("couldn't parse role id '{s}'"))
138        }
139
140        match s.chars().next() {
141            Some(SYSTEM_CHAR) => {
142                let val = parse_u64(s)?;
143                Ok(Self::System(val))
144            }
145            Some(PREDEFINED_CHAR) => {
146                let val = parse_u64(s)?;
147                Ok(Self::Predefined(val))
148            }
149            Some(USER_CHAR) => {
150                let val = parse_u64(s)?;
151                Ok(Self::User(val))
152            }
153            Some(PUBLIC_CHAR) => {
154                if s.len() == 1 {
155                    Ok(Self::Public)
156                } else {
157                    Err(anyhow!("couldn't parse role id '{s}'"))
158                }
159            }
160            _ => Err(anyhow!("couldn't parse role id '{s}'")),
161        }
162    }
163}
164
165impl fmt::Display for RoleId {
166    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
167        match self {
168            Self::System(id) => write!(f, "{SYSTEM_CHAR}{id}"),
169            Self::Predefined(id) => write!(f, "{PREDEFINED_CHAR}{id}"),
170            Self::User(id) => write!(f, "{USER_CHAR}{id}"),
171            Self::Public => write!(f, "{PUBLIC_CHAR}"),
172        }
173    }
174}
175
176impl RustType<ProtoRoleId> for RoleId {
177    fn into_proto(&self) -> ProtoRoleId {
178        use proto_role_id::Kind::*;
179        ProtoRoleId {
180            kind: Some(match self {
181                RoleId::System(x) => System(*x),
182                RoleId::Predefined(x) => Predefined(*x),
183                RoleId::User(x) => User(*x),
184                RoleId::Public => Public(()),
185            }),
186        }
187    }
188
189    fn from_proto(proto: ProtoRoleId) -> Result<Self, TryFromProtoError> {
190        use proto_role_id::Kind::*;
191        match proto.kind {
192            Some(System(x)) => Ok(RoleId::System(x)),
193            Some(Predefined(x)) => Ok(RoleId::Predefined(x)),
194            Some(User(x)) => Ok(RoleId::User(x)),
195            Some(Public(_)) => Ok(RoleId::Public),
196            None => Err(TryFromProtoError::missing_field("ProtoRoleId::kind")),
197        }
198    }
199}
200
201impl Columnation for RoleId {
202    type InnerRegion = CopyRegion<RoleId>;
203}
204
205#[mz_ore::test]
206fn test_role_id_parsing() {
207    let s = "s42";
208    let role_id: RoleId = s.parse().unwrap();
209    assert_eq!(RoleId::System(42), role_id);
210    assert_eq!(s, role_id.to_string());
211
212    let s = "g24";
213    let role_id: RoleId = s.parse().unwrap();
214    assert_eq!(RoleId::Predefined(24), role_id);
215    assert_eq!(s, role_id.to_string());
216
217    let s = "u666";
218    let role_id: RoleId = s.parse().unwrap();
219    assert_eq!(RoleId::User(666), role_id);
220    assert_eq!(s, role_id.to_string());
221
222    let s = "p";
223    let role_id: RoleId = s.parse().unwrap();
224    assert_eq!(RoleId::Public, role_id);
225    assert_eq!(s, role_id.to_string());
226
227    let s = "p23";
228    mz_ore::assert_err!(s.parse::<RoleId>());
229
230    let s = "d23";
231    mz_ore::assert_err!(s.parse::<RoleId>());
232
233    let s = "asfje90uf23i";
234    mz_ore::assert_err!(s.parse::<RoleId>());
235
236    let s = "";
237    mz_ore::assert_err!(s.parse::<RoleId>());
238}
239
240#[mz_ore::test]
241fn test_role_id_binary() {
242    let role_id = RoleId::System(42);
243    assert_eq!(
244        role_id,
245        RoleId::decode_binary(&role_id.encode_binary()).unwrap()
246    );
247
248    let role_id = RoleId::Predefined(24);
249    assert_eq!(
250        role_id,
251        RoleId::decode_binary(&role_id.encode_binary()).unwrap()
252    );
253
254    let role_id = RoleId::User(666);
255    assert_eq!(
256        role_id,
257        RoleId::decode_binary(&role_id.encode_binary()).unwrap()
258    );
259
260    let role_id = RoleId::Public;
261    assert_eq!(
262        role_id,
263        RoleId::decode_binary(&role_id.encode_binary()).unwrap()
264    );
265
266    mz_ore::assert_err!(RoleId::decode_binary(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 0]))
267}
268
269#[mz_ore::test]
270fn test_role_id_binary_size() {
271    assert_eq!(9, RoleId::binary_size());
272}