wowlab_sentinel/http/services/
nodes.rs

1//! Node registration, token, and unlink use cases.
2
3use async_trait::async_trait;
4use uuid::Uuid;
5use wowlab_centrifuge::TokenError;
6use wowlab_common::{
7    NodePublicKey,
8    node_http::{
9        NodeRegistrationAccessRule, NodeRegistrationRequest, NodeRegistrationResponse,
10        NodeTokenResponse, NodeUnlinkResponse,
11    },
12};
13use wowlab_types::sensitive::Sensitive;
14
15use crate::ExposeSecret;
16
17const DEFAULT_CORE_COUNT: i32 = 4;
18
19#[derive(Debug, thiserror::Error)]
20pub(in crate::http) enum NodeOperationError {
21    #[error("invalid token")]
22    InvalidToken {
23        #[source]
24        source: Option<sqlx::Error>,
25    },
26    #[error("failed to create node")]
27    CreateNode {
28        #[source]
29        source: sqlx::Error,
30    },
31    #[error("token generation failed")]
32    TokenGeneration(#[source] TokenIssueError),
33    #[error("node not found or not claimed")]
34    NodeNotClaimed {
35        #[source]
36        source: Option<sqlx::Error>,
37    },
38    #[error("node not found")]
39    NodeNotFound,
40    #[error("database error")]
41    Database(#[source] sqlx::Error),
42}
43
44#[derive(Debug, thiserror::Error)]
45pub(in crate::http) enum TokenIssueError {
46    #[error(transparent)]
47    Centrifuge(#[from] TokenError),
48    #[cfg(test)]
49    #[error("{0}")]
50    Test(&'static str),
51}
52
53#[derive(Debug)]
54pub(in crate::http) struct NewNode<'a> {
55    public_key: &'a str,
56    user_id: Uuid,
57    name: &'a str,
58    total_cores: i32,
59    max_parallel: i32,
60    platform: &'a str,
61    version: Option<&'a str>,
62}
63
64#[async_trait]
65pub(in crate::http) trait NodeRepository: Send + Sync {
66    async fn user_by_claim(&self, token_claim: &str) -> Result<Option<Uuid>, sqlx::Error>;
67    async fn exists(&self, public_key: &str) -> Result<bool, sqlx::Error>;
68    async fn update_user(&self, public_key: &str, user_id: Uuid) -> Result<(), sqlx::Error>;
69    async fn insert(&self, node: NewNode<'_>) -> Result<(), sqlx::Error>;
70    async fn upsert_permission(
71        &self,
72        public_key: &str,
73        rule: &NodeRegistrationAccessRule,
74    ) -> Result<(), sqlx::Error>;
75    async fn is_claimed(&self, public_key: &str) -> Result<bool, sqlx::Error>;
76    async fn delete(&self, public_key: &str) -> Result<bool, sqlx::Error>;
77}
78
79pub(in crate::http) trait BeaconTokenIssuer: Send + Sync {
80    fn issue(&self, public_key: &NodePublicKey) -> Result<Sensitive<String>, TokenIssueError>;
81}
82
83pub(in crate::http) struct SqlNodeRepository<'a> {
84    pool: &'a sqlx::PgPool,
85}
86
87#[async_trait]
88impl NodeRepository for SqlNodeRepository<'_> {
89    async fn user_by_claim(&self, token_claim: &str) -> Result<Option<Uuid>, sqlx::Error> {
90        sqlx::query_file_scalar!("queries/nodes_get_user_by_token_claim.sql", token_claim)
91            .fetch_optional(self.pool)
92            .await
93    }
94
95    async fn exists(&self, public_key: &str) -> Result<bool, sqlx::Error> {
96        Ok(
97            sqlx::query_file_scalar!("queries/nodes_check_exists.sql", public_key)
98                .fetch_optional(self.pool)
99                .await?
100                .is_some(),
101        )
102    }
103
104    async fn update_user(&self, public_key: &str, user_id: Uuid) -> Result<(), sqlx::Error> {
105        sqlx::query_file!("queries/nodes_update_user.sql", user_id, public_key)
106            .execute(self.pool)
107            .await?;
108
109        Ok(())
110    }
111
112    async fn insert(&self, node: NewNode<'_>) -> Result<(), sqlx::Error> {
113        sqlx::query_file!(
114            "queries/nodes_insert.sql",
115            node.public_key,
116            node.user_id,
117            node.name,
118            node.total_cores,
119            node.max_parallel,
120            node.platform,
121            node.version,
122        )
123        .execute(self.pool)
124        .await?;
125
126        Ok(())
127    }
128
129    async fn upsert_permission(
130        &self,
131        public_key: &str,
132        rule: &NodeRegistrationAccessRule,
133    ) -> Result<(), sqlx::Error> {
134        let target_id = rule.target_id.as_deref().unwrap_or("");
135
136        sqlx::query_file!(
137            "queries/nodes_upsert_permission.sql",
138            public_key,
139            &rule.access_type,
140            target_id,
141        )
142        .execute(self.pool)
143        .await?;
144
145        Ok(())
146    }
147
148    async fn is_claimed(&self, public_key: &str) -> Result<bool, sqlx::Error> {
149        Ok(
150            sqlx::query_file_scalar!("queries/nodes_check_claimed.sql", public_key)
151                .fetch_optional(self.pool)
152                .await?
153                .unwrap_or(false),
154        )
155    }
156
157    async fn delete(&self, public_key: &str) -> Result<bool, sqlx::Error> {
158        let result = sqlx::query_file!("queries/nodes_delete.sql", public_key)
159            .execute(self.pool)
160            .await?;
161
162        Ok(result.rows_affected() > 0)
163    }
164}
165
166pub(in crate::http) struct CentrifugeTokenIssuer<'a> {
167    secret: &'a str,
168}
169
170impl BeaconTokenIssuer for CentrifugeTokenIssuer<'_> {
171    fn issue(&self, public_key: &NodePublicKey) -> Result<Sensitive<String>, TokenIssueError> {
172        Ok(wowlab_centrifuge::generate_token(
173            &public_key.to_string(),
174            self.secret,
175        )?)
176    }
177}
178
179pub(in crate::http) struct NodeOperations<R, T> {
180    repository: R,
181    token_issuer: T,
182}
183
184pub(in crate::http) fn node_operations(
185    state: &crate::state::ServerState,
186) -> NodeOperations<SqlNodeRepository<'_>, CentrifugeTokenIssuer<'_>> {
187    NodeOperations {
188        repository: SqlNodeRepository {
189            pool: state.dbs.get::<crate::http::HttpDb>(),
190        },
191        token_issuer: CentrifugeTokenIssuer {
192            secret: state.config.centrifugo_token_secret.expose_secret(),
193        },
194    }
195}
196
197impl<R, T> NodeOperations<R, T>
198where
199    R: NodeRepository,
200    T: BeaconTokenIssuer,
201{
202    pub(in crate::http) async fn register(
203        &self,
204        public_key: &NodePublicKey,
205        request: &NodeRegistrationRequest,
206    ) -> Result<NodeRegistrationResponse, NodeOperationError> {
207        let user_id = match self.repository.user_by_claim(&request.token_claim).await {
208            Ok(Some(user_id)) => user_id,
209            Ok(None) => return Err(NodeOperationError::InvalidToken { source: None }),
210            Err(source) => {
211                return Err(NodeOperationError::InvalidToken {
212                    source: Some(source),
213                });
214            }
215        };
216        let public_key_text = public_key.to_string();
217
218        if self
219            .repository
220            .exists(&public_key_text)
221            .await
222            .unwrap_or(false)
223        {
224            if let Err(error) = self.repository.update_user(&public_key_text, user_id).await {
225                tracing::error!(%error, public_key = %public_key, "Failed to update node user_id");
226            }
227
228            return Ok(NodeRegistrationResponse {
229                beacon_token: self
230                    .token_issuer
231                    .issue(public_key)
232                    .ok()
233                    .map(Sensitive::into_inner),
234            });
235        }
236
237        let total_cores = request.total_cores.unwrap_or(DEFAULT_CORE_COUNT);
238        let node = NewNode {
239            public_key: &public_key_text,
240            user_id,
241            name: request.hostname.as_deref().unwrap_or("WowLab Node"),
242            total_cores,
243            max_parallel: request.enabled_cores.unwrap_or(total_cores),
244            platform: request.platform.as_deref().unwrap_or("unknown"),
245            version: request.version.as_deref(),
246        };
247
248        if let Err(source) = self.repository.insert(node).await {
249            tracing::error!(error = %source, "Failed to insert node");
250
251            return Err(NodeOperationError::CreateNode { source });
252        }
253
254        let mut permission_failures = 0usize;
255
256        for rule in &request.access_rules {
257            if self
258                .repository
259                .upsert_permission(&public_key_text, rule)
260                .await
261                .is_err()
262            {
263                permission_failures += 1;
264            }
265        }
266
267        if permission_failures > 0 {
268            tracing::error!(permission_failures, "Failed to insert node permissions");
269        }
270
271        Ok(NodeRegistrationResponse {
272            beacon_token: self
273                .token_issuer
274                .issue(public_key)
275                .ok()
276                .map(Sensitive::into_inner),
277        })
278    }
279
280    pub(in crate::http) async fn refresh_token(
281        &self,
282        public_key: &NodePublicKey,
283    ) -> Result<NodeTokenResponse, NodeOperationError> {
284        let public_key_text = public_key.to_string();
285
286        match self.repository.is_claimed(&public_key_text).await {
287            Ok(true) => {}
288            Ok(false) => return Err(NodeOperationError::NodeNotClaimed { source: None }),
289            Err(source) => {
290                return Err(NodeOperationError::NodeNotClaimed {
291                    source: Some(source),
292                });
293            }
294        }
295
296        let token = match self.token_issuer.issue(public_key) {
297            Ok(token) => token,
298            Err(source) => {
299                tracing::error!("Failed to generate beacon token");
300
301                return Err(NodeOperationError::TokenGeneration(source));
302            }
303        };
304
305        Ok(NodeTokenResponse::new(token.into_inner()))
306    }
307
308    pub(in crate::http) async fn unlink(
309        &self,
310        public_key: &NodePublicKey,
311    ) -> Result<NodeUnlinkResponse, NodeOperationError> {
312        let public_key_text = public_key.to_string();
313
314        match self.repository.delete(&public_key_text).await {
315            Ok(true) => {
316                tracing::info!(public_key = %public_key, "Node unlinked");
317
318                Ok(NodeUnlinkResponse::success())
319            }
320            Ok(false) => Err(NodeOperationError::NodeNotFound),
321            Err(source) => {
322                tracing::error!(error = %source, "Failed to delete node");
323
324                Err(NodeOperationError::Database(source))
325            }
326        }
327    }
328}
329
330impl From<TokenIssueError> for NodeOperationError {
331    fn from(error: TokenIssueError) -> Self {
332        Self::TokenGeneration(error)
333    }
334}
335
336#[cfg(test)]
337mod tests;
wowlab_sentinel/http/services/nodes.rs

wowlab_sentinel/http/services/
nodes.rs
