alice.pds.demo.boltless.dev
This repository has no description
did:plc:3fwecdnvtcscjnrx2p4n7a…
1use std::collections::{HashMap, HashSet};
2use std::net::{IpAddr, SocketAddr};
3use std::sync::Arc;
4
5use knot_runtime::{Clock, HttpTransport};
6use knot_types::OfferedKey;
7use russh::keys::ssh_key;
8use russh::server::{Auth, Handler, Msg, Server, Session};
9use russh::{Channel, ChannelId};
10use tokio_util::task::TaskTracker;
11
12use crate::SshState;
13use crate::exec::run_exec;
14
15pub(crate) struct KnotSshServer<H, C> {
16 pub(crate) state: Arc<SshState<H, C>>,
17 pub(crate) tracker: TaskTracker,
18}
19
20impl<H: HttpTransport, C: Clock> Server for KnotSshServer<H, C> {
21 type Handler = KnotSession<H, C>;
22
23 fn new_client(&mut self, peer: Option<SocketAddr>) -> Self::Handler {
24 KnotSession::new(
25 Arc::clone(&self.state),
26 self.tracker.clone(),
27 peer.map(|addr| addr.ip()),
28 )
29 }
30}
31
32pub(crate) struct KnotSession<H, C> {
33 state: Arc<SshState<H, C>>,
34 tracker: TaskTracker,
35 key: Option<OfferedKey>,
36 channels: HashMap<ChannelId, Channel<Msg>>,
37 protocols: HashSet<ChannelId>,
38 peer: Option<IpAddr>,
39}
40
41impl<H, C> KnotSession<H, C> {
42 fn new(state: Arc<SshState<H, C>>, tracker: TaskTracker, peer: Option<IpAddr>) -> Self {
43 Self {
44 state,
45 tracker,
46 key: None,
47 channels: HashMap::new(),
48 protocols: HashSet::new(),
49 peer,
50 }
51 }
52}
53
54impl<H: HttpTransport, C: Clock> Handler for KnotSession<H, C> {
55 type Error = russh::Error;
56
57 async fn auth_publickey(
58 &mut self,
59 _user: &str,
60 public_key: &ssh_key::PublicKey,
61 ) -> Result<Auth, Self::Error> {
62 let reject = Auth::Reject {
63 proceed_with_methods: None,
64 partial_success: false,
65 };
66 let Ok(blob) = public_key.to_bytes() else {
67 return Ok(reject);
68 };
69 let key = OfferedKey::from_bytes(blob);
70 if self
71 .state
72 .roster
73 .recognizes_fresh(&key, &self.state.index, &self.state.atproto)
74 .await
75 {
76 self.key = Some(key);
77 Ok(Auth::Accept)
78 } else {
79 Ok(reject)
80 }
81 }
82
83 async fn channel_open_session(
84 &mut self,
85 channel: Channel<Msg>,
86 _session: &mut Session,
87 ) -> Result<bool, Self::Error> {
88 self.channels.insert(channel.id(), channel);
89 Ok(true)
90 }
91
92 async fn env_request(
93 &mut self,
94 channel: ChannelId,
95 variable_name: &str,
96 variable_value: &str,
97 _session: &mut Session,
98 ) -> Result<(), Self::Error> {
99 if variable_name == "GIT_PROTOCOL"
100 && variable_value
101 .split(':')
102 .any(|token| token.trim() == "version=2")
103 {
104 self.protocols.insert(channel);
105 }
106 Ok(())
107 }
108
109 async fn channel_eof(
110 &mut self,
111 channel: ChannelId,
112 _session: &mut Session,
113 ) -> Result<(), Self::Error> {
114 self.protocols.remove(&channel);
115 Ok(())
116 }
117
118 async fn channel_close(
119 &mut self,
120 channel: ChannelId,
121 _session: &mut Session,
122 ) -> Result<(), Self::Error> {
123 self.channels.remove(&channel);
124 self.protocols.remove(&channel);
125 Ok(())
126 }
127
128 #[allow(clippy::too_many_arguments)]
129 async fn pty_request(
130 &mut self,
131 channel: ChannelId,
132 _term: &str,
133 _col_width: u32,
134 _row_height: u32,
135 _pix_width: u32,
136 _pix_height: u32,
137 _modes: &[(russh::Pty, u32)],
138 session: &mut Session,
139 ) -> Result<(), Self::Error> {
140 session.channel_success(channel)?;
141 Ok(())
142 }
143
144 async fn shell_request(
145 &mut self,
146 channel: ChannelId,
147 session: &mut Session,
148 ) -> Result<(), Self::Error> {
149 let Some(handle) = self.channels.remove(&channel) else {
150 session.channel_failure(channel)?;
151 return Ok(());
152 };
153 session.channel_success(channel)?;
154 let state = Arc::clone(&self.state);
155 let key = self.key.clone();
156 self.tracker.spawn(async move {
157 crate::exec::run_greeting(state, key, handle).await;
158 });
159 Ok(())
160 }
161
162 async fn exec_request(
163 &mut self,
164 channel: ChannelId,
165 data: &[u8],
166 session: &mut Session,
167 ) -> Result<(), Self::Error> {
168 let Some(handle) = self.channels.remove(&channel) else {
169 session.channel_failure(channel)?;
170 return Ok(());
171 };
172 session.channel_success(channel)?;
173 let protocol_v2 = self.protocols.remove(&channel);
174 let state = Arc::clone(&self.state);
175 let key = self.key.clone();
176 let peer = self.peer;
177 let command = data.to_vec();
178 self.tracker.spawn(async move {
179 run_exec(state, key, handle, &command, protocol_v2, peer).await;
180 });
181 Ok(())
182 }
183}