1use std::collections::HashMap;
5
6use reifydb_auth::service::{AuthResponse, AuthService};
7use reifydb_core::{
8 actors::server::{
9 ServerAuthResponse, ServerLogoutResponse, ServerMessage, ServerResponse, ServerSubscribeResponse,
10 },
11 execution::ExecutionResult,
12};
13use reifydb_engine::{engine::StandardEngine, session::RetryStrategy};
14use reifydb_runtime::{
15 actor::{
16 context::Context,
17 reply::Reply,
18 traits::{Actor, Directive},
19 },
20 context::clock::Clock,
21};
22use reifydb_value::{
23 params::Params,
24 value::{duration::Duration, identity::IdentityId},
25};
26
27pub struct ServerActor {
28 engine: StandardEngine,
29 auth_service: AuthService,
30 clock: Clock,
31 retry: RetryStrategy,
32}
33
34impl ServerActor {
35 pub fn new(engine: StandardEngine, auth_service: AuthService, clock: Clock) -> Self {
36 Self {
37 engine,
38 auth_service,
39 clock,
40 retry: RetryStrategy::default(),
41 }
42 }
43
44 fn dispatch_execute(
45 &self,
46 identity: IdentityId,
47 rql: String,
48 params: Params,
49 reply: Reply<ServerResponse>,
50 execute: impl Fn(&StandardEngine, IdentityId, &str, Params) -> ExecutionResult,
51 ) {
52 let t = self.clock.instant();
53 let result = self
54 .retry
55 .execute(self.engine.rng(), &rql, || execute(&self.engine, identity, &rql, params.clone()));
56 if let Some(err) = result.error {
57 reply.send(ServerResponse::EngineError {
58 diagnostic: Box::new(err.diagnostic()),
59 rql,
60 });
61 } else {
62 reply.send(ServerResponse::Success {
63 frames: result.frames,
64 duration: Duration::from_std(t.elapsed()),
65 metrics: result.metrics,
66 });
67 }
68 }
69
70 #[inline]
71 fn handle_subscribe(&self, identity: IdentityId, rql: String, reply: Reply<ServerSubscribeResponse>) {
72 let t = self.clock.instant();
73 let result = self.engine.subscribe_as(identity, &rql, Params::None);
74 if let Some(err) = result.error {
75 reply.send(ServerSubscribeResponse::EngineError {
76 diagnostic: Box::new(err.diagnostic()),
77 rql,
78 });
79 } else {
80 reply.send(ServerSubscribeResponse::Subscribed {
81 frames: result.frames,
82 duration: Duration::from_std(t.elapsed()),
83 metrics: result.metrics,
84 });
85 }
86 }
87
88 #[inline]
89 fn handle_authenticate(
90 &self,
91 method: String,
92 credentials: HashMap<String, String>,
93 reply: Reply<ServerAuthResponse>,
94 ) {
95 match self.auth_service.authenticate(&method, credentials) {
96 Ok(AuthResponse::Authenticated {
97 identity,
98 token,
99 }) => reply.send(ServerAuthResponse::Authenticated {
100 identity,
101 token,
102 }),
103 Ok(AuthResponse::Challenge {
104 challenge_id,
105 payload,
106 }) => reply.send(ServerAuthResponse::Challenge {
107 challenge_id,
108 payload,
109 }),
110 Ok(AuthResponse::Failed {
111 reason,
112 }) => reply.send(ServerAuthResponse::Failed {
113 reason,
114 }),
115 Err(e) => reply.send(ServerAuthResponse::Error(e.to_string())),
116 }
117 }
118
119 #[inline]
120 fn handle_logout(&self, token: String, reply: Reply<ServerLogoutResponse>) {
121 if self.auth_service.revoke_token(&token) {
122 reply.send(ServerLogoutResponse::Ok);
123 } else {
124 reply.send(ServerLogoutResponse::InvalidToken);
125 }
126 }
127}
128
129impl Actor for ServerActor {
130 type State = ();
131 type Message = ServerMessage;
132
133 fn init(&self, _ctx: &Context<Self::Message>) -> Self::State {}
134
135 fn handle(&self, _state: &mut (), msg: ServerMessage, _ctx: &Context<ServerMessage>) -> Directive {
136 match msg {
137 ServerMessage::Query {
138 identity,
139 rql,
140 params,
141 reply,
142 } => {
143 self.dispatch_execute(identity, rql, params, reply, |e, id, s, p| e.query_as(id, s, p));
144 }
145 ServerMessage::Command {
146 identity,
147 rql,
148 params,
149 reply,
150 } => {
151 self.dispatch_execute(identity, rql, params, reply, |e, id, s, p| {
152 e.command_as(id, s, p)
153 });
154 }
155 ServerMessage::Admin {
156 identity,
157 rql,
158 params,
159 reply,
160 } => {
161 self.dispatch_execute(identity, rql, params, reply, |e, id, s, p| e.admin_as(id, s, p));
162 }
163 ServerMessage::Subscribe {
164 identity,
165 rql,
166 reply,
167 } => {
168 self.handle_subscribe(identity, rql, reply);
169 }
170 ServerMessage::Authenticate {
171 method,
172 credentials,
173 reply,
174 } => {
175 self.handle_authenticate(method, credentials, reply);
176 }
177 ServerMessage::Logout {
178 token,
179 reply,
180 } => {
181 self.handle_logout(token, reply);
182 }
183 }
184 Directive::Continue
185 }
186}