This commit is contained in:
nora 2024-08-10 01:59:59 +02:00
parent be5b437548
commit aa96faa697

View file

@ -107,23 +107,8 @@ impl ServerConnection {
} }
impl ServerConnection { impl ServerConnection {
pub fn recv_bytes(&mut self, mut bytes: &[u8]) -> Result<()> { pub fn recv_bytes(&mut self, bytes: &[u8]) -> Result<()> {
while let Some(consumed) = self.recv_bytes_step(bytes)? { if let ServerState::ProtoExchange { received } = &mut self.state {
bytes = &bytes[consumed..];
if bytes.is_empty() {
break;
}
}
Ok(())
}
fn recv_bytes_step(&mut self, bytes: &[u8]) -> Result<Option<usize>> {
if !matches!(self.state, ServerState::ProtoExchange { .. }) {
self.packet_transport.recv_bytes(bytes);
}
let result = match &mut self.state {
ServerState::ProtoExchange { received } => {
// TODO: get rid of this allocation :(
received.extend_from_slice(bytes); received.extend_from_slice(bytes);
if received.windows(2).find(|win| win == b"\r\n").is_some() { if received.windows(2).find(|win| win == b"\r\n").is_some() {
// TODO: care that its SSH 2.0 instead of anythin anything else // TODO: care that its SSH 2.0 instead of anythin anything else
@ -134,8 +119,14 @@ impl ServerConnection {
client_identification, client_identification,
}; };
} }
None // This means that we must be called at least twice, which is fine I think.
return Ok(());
} }
self.packet_transport.recv_bytes(bytes)?;
match &mut self.state {
ServerState::ProtoExchange { .. } => unreachable!("handled above"),
ServerState::KeyExchangeInit { ServerState::KeyExchangeInit {
client_identification, client_identification,
} => match self.packet_transport.next_packet() { } => match self.packet_transport.next_packet() {
@ -219,10 +210,8 @@ impl ServerConnection {
client_kexinit: data.payload, client_kexinit: data.payload,
server_kexinit: server_kexinit_payload, server_kexinit: server_kexinit_payload,
}; };
None
} }
None => None, None => {},
}, },
ServerState::DhKeyInit { ServerState::DhKeyInit {
client_identification, client_identification,
@ -313,10 +302,8 @@ impl ServerConnection {
})); }));
self.state = ServerState::NewKeys; self.state = ServerState::NewKeys;
// TODO: set keys for transport // TODO: set keys for transport
None
} }
None => None, None => {},
}, },
ServerState::NewKeys => match self.packet_transport.next_packet() { ServerState::NewKeys => match self.packet_transport.next_packet() {
Some(data) => { Some(data) => {
@ -328,14 +315,12 @@ impl ServerConnection {
payload: vec![Packet::SSH_MSG_NEWKEYS], payload: vec![Packet::SSH_MSG_NEWKEYS],
})); }));
self.state = ServerState::ServiceRequest {}; self.state = ServerState::ServiceRequest {};
None
} }
None => None, None => {},
}, },
ServerState::ServiceRequest {} => None, ServerState::ServiceRequest {} => {},
}; }
Ok(result) Ok(())
} }
pub fn next_message_to_send(&mut self) -> Option<Msg> { pub fn next_message_to_send(&mut self) -> Option<Msg> {