JavaScript is disabled, refresh for a better experience. ambee/giterated

ambee/giterated

Git repository hosting, collaboration, and discovery for the Fediverse.

Fixes

View files Patch

Amber - ⁨2⁩ years ago

parent: tbd commit: ⁨25c3410⁩

⁨giterated-daemon/src/connection/wrapper.rs⁩ - ⁨6679⁩ bytes
Raw
1 use std::{
2 collections::HashMap,
3 net::SocketAddr,
4 sync::{
5 atomic::{AtomicBool, Ordering},
6 Arc,
7 },
8 };
9
10 use anyhow::Error;
11 use futures_util::{SinkExt, StreamExt};
12 use giterated_models::{
13 messages::error::ConnectionError,
14 model::{authenticated::Authenticated, instance::Instance},
15 };
16 use rsa::RsaPublicKey;
17 use serde::Serialize;
18 use serde_json::Value;
19 use tokio::{
20 net::TcpStream,
21 sync::{Mutex, RwLock},
22 };
23 use tokio_tungstenite::{tungstenite::Message, WebSocketStream};
24
25 use crate::{
26 authentication::AuthenticationTokenGranter,
27 backend::{RepositoryBackend, UserBackend},
28 connection::forwarded::wrap_forwarded,
29 federation::connections::InstanceConnections,
30 message::NetworkMessage,
31 };
32
33 use super::{
34 authentication::authentication_handle, handshake::handshake_handle,
35 repository::repository_handle, user::user_handle, Connections,
36 };
37
38 pub async fn connection_wrapper(
39 socket: WebSocketStream<TcpStream>,
40 connections: Arc<Mutex<Connections>>,
41 repository_backend: Arc<Mutex<dyn RepositoryBackend + Send>>,
42 user_backend: Arc<Mutex<dyn UserBackend + Send>>,
43 auth_granter: Arc<Mutex<AuthenticationTokenGranter>>,
44 addr: SocketAddr,
45 instance: impl ToOwned<Owned = Instance>,
46 instance_connections: Arc<Mutex<InstanceConnections>>,
47 ) {
48 let connection_state = ConnectionState {
49 socket: Arc::new(Mutex::new(socket)),
50 connections,
51 repository_backend,
52 user_backend,
53 auth_granter,
54 addr,
55 instance: instance.to_owned(),
56 handshaked: Arc::new(AtomicBool::new(false)),
57 cached_keys: Arc::default(),
58 };
59
60 let mut handshaked = false;
61
62 loop {
63 let mut socket = connection_state.socket.lock().await;
64 let message = socket.next().await;
65 drop(socket);
66
67 match message {
68 Some(Ok(message)) => {
69 let payload = match message {
70 Message::Binary(payload) => payload,
71 Message::Ping(_) => {
72 let mut socket = connection_state.socket.lock().await;
73 let _ = socket.send(Message::Pong(vec![])).await;
74 drop(socket);
75 continue;
76 }
77 Message::Close(_) => return,
78 _ => continue,
79 };
80
81 info!(
82 "Received payload: {}",
83 std::str::from_utf8(&payload).unwrap()
84 );
85
86 let message = NetworkMessage(payload.clone());
87
88 if !handshaked {
89 if handshake_handle(&message, &connection_state).await.is_ok() {
90 if connection_state.handshaked.load(Ordering::SeqCst) {
91 handshaked = true;
92 }
93 }
94 } else {
95 let raw = serde_json::from_slice::<Authenticated<Value>>(&payload).unwrap();
96
97 if let Some(target_instance) = &raw.target_instance {
98 // Forward request
99 info!("Forwarding message to {}", target_instance.url);
100 let mut instance_connections = instance_connections.lock().await;
101 let pool = instance_connections.get_or_open(&target_instance).unwrap();
102 let pool_clone = pool.clone();
103 drop(pool);
104
105 let result = wrap_forwarded(&pool_clone, raw).await;
106
107 let mut socket = connection_state.socket.lock().await;
108 let _ = socket.send(result).await;
109
110 continue;
111 }
112
113 let message_type = &raw.message_type;
114
115 info!("Handling message with type: {}", message_type);
116
117 match authentication_handle(message_type, &message, &connection_state).await {
118 Err(e) => {
119 let _ = connection_state.send(ConnectionError(e.to_string())).await;
120 }
121 Ok(true) => continue,
122 Ok(false) => {}
123 }
124
125 match repository_handle(message_type, &message, &connection_state).await {
126 Err(e) => {
127 let _ = connection_state.send(ConnectionError(e.to_string())).await;
128 }
129 Ok(true) => continue,
130 Ok(false) => {}
131 }
132
133 match user_handle(message_type, &message, &connection_state).await {
134 Err(e) => {
135 let _ = connection_state.send(ConnectionError(e.to_string())).await;
136 }
137 Ok(true) => continue,
138 Ok(false) => {}
139 }
140
141 match authentication_handle(message_type, &message, &connection_state).await {
142 Err(e) => {
143 let _ = connection_state.send(ConnectionError(e.to_string())).await;
144 }
145 Ok(true) => continue,
146 Ok(false) => {}
147 }
148
149 error!(
150 "Message completely unhandled: {}",
151 std::str::from_utf8(&payload).unwrap()
152 );
153 }
154 }
155 Some(Err(e)) => {
156 error!("Closing connection for {:?} for {}", e, addr);
157 return;
158 }
159 _ => {
160 info!("Unhandled");
161 continue;
162 }
163 }
164 }
165 }
166
167 #[derive(Clone)]
168 pub struct ConnectionState {
169 socket: Arc<Mutex<WebSocketStream<TcpStream>>>,
170 pub connections: Arc<Mutex<Connections>>,
171 pub repository_backend: Arc<Mutex<dyn RepositoryBackend + Send>>,
172 pub user_backend: Arc<Mutex<dyn UserBackend + Send>>,
173 pub auth_granter: Arc<Mutex<AuthenticationTokenGranter>>,
174 pub addr: SocketAddr,
175 pub instance: Instance,
176 pub handshaked: Arc<AtomicBool>,
177 pub cached_keys: Arc<RwLock<HashMap<Instance, RsaPublicKey>>>,
178 }
179
180 impl ConnectionState {
181 pub async fn send<T: Serialize>(&self, message: T) -> Result<(), Error> {
182 let payload = serde_json::to_string(&message)?;
183 info!("Sending payload: {}", &payload);
184 self.socket
185 .lock()
186 .await
187 .send(Message::Binary(payload.into_bytes()))
188 .await?;
189
190 Ok(())
191 }
192 }
193