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};
17#[cfg(any(test, feature = "proptest"))]
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    Clone,
35    Copy,
36    Debug,
37    Eq,
38    PartialEq,
39    Ord,
40    PartialOrd,
41    Hash,
42    Serialize,
43    Deserialize,
44    MzReflect
45)]
46#[cfg_attr(any(test, feature = "proptest"), derive(Arbitrary))]
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
201#[mz_ore::test]
202fn test_role_id_parsing() {
203    let s = "s42";
204    let role_id: RoleId = s.parse().unwrap();
205    assert_eq!(RoleId::System(42), role_id);
206    assert_eq!(s, role_id.to_string());
207
208    let s = "g24";
209    let role_id: RoleId = s.parse().unwrap();
210    assert_eq!(RoleId::Predefined(24), role_id);
211    assert_eq!(s, role_id.to_string());
212
213    let s = "u666";
214    let role_id: RoleId = s.parse().unwrap();
215    assert_eq!(RoleId::User(666), role_id);
216    assert_eq!(s, role_id.to_string());
217
218    let s = "p";
219    let role_id: RoleId = s.parse().unwrap();
220    assert_eq!(RoleId::Public, role_id);
221    assert_eq!(s, role_id.to_string());
222
223    let s = "p23";
224    mz_ore::assert_err!(s.parse::<RoleId>());
225
226    let s = "d23";
227    mz_ore::assert_err!(s.parse::<RoleId>());
228
229    let s = "asfje90uf23i";
230    mz_ore::assert_err!(s.parse::<RoleId>());
231
232    let s = "";
233    mz_ore::assert_err!(s.parse::<RoleId>());
234}
235
236#[mz_ore::test]
237fn test_role_id_binary() {
238    let role_id = RoleId::System(42);
239    assert_eq!(
240        role_id,
241        RoleId::decode_binary(&role_id.encode_binary()).unwrap()
242    );
243
244    let role_id = RoleId::Predefined(24);
245    assert_eq!(
246        role_id,
247        RoleId::decode_binary(&role_id.encode_binary()).unwrap()
248    );
249
250    let role_id = RoleId::User(666);
251    assert_eq!(
252        role_id,
253        RoleId::decode_binary(&role_id.encode_binary()).unwrap()
254    );
255
256    let role_id = RoleId::Public;
257    assert_eq!(
258        role_id,
259        RoleId::decode_binary(&role_id.encode_binary()).unwrap()
260    );
261
262    mz_ore::assert_err!(RoleId::decode_binary(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 0]))
263}
264
265#[mz_ore::test]
266fn test_role_id_binary_size() {
267    assert_eq!(9, RoleId::binary_size());
268}