triton_distributed/transports/nats/
slug.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3//
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at
7//
8// http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16use serde::de::{self, Deserializer, Visitor};
17use serde::{Deserialize, Serialize};
18use std::fmt;
19
20const REPLACEMENT_CHAR: char = '_';
21
22/// URL and NATS friendly string.
23/// Only a-z, 0-9, - and _.
24#[derive(Serialize, Clone, Debug, Eq, PartialEq)]
25pub struct Slug(String);
26
27impl Slug {
28    fn new(s: String) -> Slug {
29        // remove any leading REPLACEMENT_CHAR
30        let s = s.trim_start_matches(REPLACEMENT_CHAR).to_string();
31        Slug(s)
32    }
33
34    /// Create [`Slug`] from a string.
35    pub fn from_string(s: impl AsRef<str>) -> Slug {
36        Slug::slugify_unique(s.as_ref())
37    }
38
39    // /// Turn the string into a valid slug, replacing any not-web-or-nats-safe characters with '-'
40    // fn slugify(s: &str) -> Slug {
41    //     let out = s
42    //         .to_lowercase()
43    //         .chars()
44    //         .map(|c| {
45    //             let is_valid = c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-' || c == '_';
46    //             if is_valid {
47    //                 c
48    //             } else {
49    //                 REPLACEMENT_CHAR
50    //             }
51    //         })
52    //         .collect::<String>();
53    //     Slug::new(out)
54    // }
55
56    /// Like slugify but also add a four byte hash on the end, in case two different strings slug
57    /// to the same thing.
58    fn slugify_unique(s: &str) -> Slug {
59        let out = s
60            .to_lowercase()
61            .chars()
62            .map(|c| {
63                let is_valid = c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-' || c == '_';
64                if is_valid {
65                    c
66                } else {
67                    REPLACEMENT_CHAR
68                }
69            })
70            .collect::<String>();
71        let hash = blake3::hash(s.as_bytes()).to_string();
72        let out = format!("{out}-{}", &hash[(hash.len() - 8)..]);
73        Slug::new(out)
74    }
75}
76
77impl fmt::Display for Slug {
78    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
79        write!(f, "{}", self.0)
80    }
81}
82
83#[derive(Debug)]
84pub struct InvalidSlugError(char);
85
86impl fmt::Display for InvalidSlugError {
87    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
88        write!(
89            f,
90            "Invalid char '{}'. String can only contain a-z, 0-9, - and _.",
91            self.0
92        )
93    }
94}
95
96impl std::error::Error for InvalidSlugError {}
97
98impl TryFrom<&str> for Slug {
99    type Error = InvalidSlugError;
100
101    fn try_from(s: &str) -> Result<Self, Self::Error> {
102        s.to_string().try_into()
103    }
104}
105
106impl TryFrom<String> for Slug {
107    type Error = InvalidSlugError;
108
109    fn try_from(s: String) -> Result<Self, Self::Error> {
110        let is_invalid =
111            |c: &char| !c.is_ascii_lowercase() && !c.is_ascii_digit() && *c != '-' && *c != '_';
112        match s.chars().find(is_invalid) {
113            None => Ok(Slug(s)),
114            Some(c) => Err(InvalidSlugError(c)),
115        }
116    }
117}
118
119impl<'de> Deserialize<'de> for Slug {
120    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
121    where
122        D: Deserializer<'de>,
123    {
124        struct SlugVisitor;
125
126        impl Visitor<'_> for SlugVisitor {
127            type Value = Slug;
128
129            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
130                formatter
131                    .write_str("a valid slug string containing only characters a-z, 0-9, - and _.")
132            }
133
134            fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
135            where
136                E: de::Error,
137            {
138                Slug::try_from(v).map_err(de::Error::custom)
139            }
140
141            fn visit_string<E>(self, v: String) -> Result<Self::Value, E>
142            where
143                E: de::Error,
144            {
145                Slug::try_from(v.as_ref()).map_err(de::Error::custom)
146            }
147        }
148
149        deserializer.deserialize_string(SlugVisitor)
150    }
151}
152
153impl AsRef<str> for Slug {
154    fn as_ref(&self) -> &str {
155        &self.0
156    }
157}
158
159impl PartialEq<str> for Slug {
160    fn eq(&self, other: &str) -> bool {
161        self.0 == other
162    }
163}