Skip to main content

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