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