Skip to content

Commit a592366

Browse files
ricott1ricott1
committed
Move run and run_on_socket to Server trait
1 parent 0763767 commit a592366

File tree

5 files changed

+81
-78
lines changed

5 files changed

+81
-78
lines changed

‎russh/examples/echoserver.rs‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ use std::collections::HashMap;
22
use std::sync::Arc;
33

44
use async_trait::async_trait;
5+
use russh::server::Server as _;
56
use russh::server::{Msg, Session};
67
use russh::*;
78
use russh_keys::*;
@@ -25,9 +26,7 @@ async fn main() {
2526
clients: Arc::new(Mutex::new(HashMap::new())),
2627
id: 0,
2728
};
28-
russh::server::run(config, ("0.0.0.0", 2222), &mut sh)
29-
.await
30-
.unwrap();
29+
sh.run_on_address(config, ("0.0.0.0", 2222)).await.unwrap();
3130
}
3231

3332
#[derive(Clone)]

‎russh/examples/sftp_server.rs‎

Lines changed: 15 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,9 @@ use std::time::Duration;
55

66
use async_trait::async_trait;
77
use log::{error, info, LevelFilter};
8+
use russh::server::Server as _;
89
use russh::server::{Auth, Msg, Session};
10+
911
use russh::{Channel, ChannelId};
1012
use russh_keys::key::KeyPair;
1113
use russh_sftp::protocol::{File, FileAttributes, Handle, Name, Status, StatusCode, Version};
@@ -195,17 +197,17 @@ async fn main() {
195197

196198
let mut server = Server;
197199

198-
russh::server::run(
199-
Arc::new(config),
200-
(
201-
"0.0.0.0",
202-
std::env::var("PORT")
203-
.unwrap_or("22".to_string())
204-
.parse()
205-
.unwrap(),
206-
),
207-
&mut server,
208-
)
209-
.await
210-
.unwrap();
200+
server
201+
.run_on_address(
202+
Arc::new(config),
203+
(
204+
"0.0.0.0",
205+
std::env::var("PORT")
206+
.unwrap_or("22".to_string())
207+
.parse()
208+
.unwrap(),
209+
),
210+
)
211+
.await
212+
.unwrap();
211213
}

‎russh/examples/test.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ use std::sync::{Arc, Mutex};
33

44
use async_trait::async_trait;
55
use log::debug;
6+
use russh::server::Server as _;
67
use russh::server::{Auth, Msg, Session};
78
use russh::*;
89
use russh_keys::*;
@@ -22,7 +23,7 @@ async fn main() -> anyhow::Result<()> {
2223
};
2324
tokio::time::timeout(
2425
std::time::Duration::from_secs(60),
25-
russh::server::run(config, ("0.0.0.0", 2222), &mut sh),
26+
sh.run_on_address(config, ("0.0.0.0", 2222)),
2627
)
2728
.await
2829
.unwrap_or(Ok(()))?;

‎russh/src/server/mod.rs‎

Lines changed: 60 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,6 @@ use crate::*;
5151

5252
mod kex;
5353
mod session;
54-
pub use self::kex::*;
5554
pub use self::session::*;
5655
mod encrypted;
5756

@@ -525,80 +524,81 @@ pub trait Handler: Sized {
525524
}
526525
}
527526

527+
#[async_trait]
528528
/// Trait used to create new handlers when clients connect.
529529
pub trait Server {
530530
/// The type of handlers.
531-
type Handler: Handler + Send;
531+
type Handler: Handler + Send + 'static;
532532
/// Called when a new client connects.
533533
fn new_client(&mut self, peer_addr: Option<std::net::SocketAddr>) -> Self::Handler;
534534
/// Called when an active connection fails.
535535
fn handle_session_error(&mut self, _error: <Self::Handler as Handler>::Error) {}
536-
}
537536

538-
/// Run a server on a specified `tokio::net::TcpListener`. Useful when dropping
539-
/// privileges immediately after socket binding, for example.
540-
pub async fn run_on_socket<H: Server + Send + 'static>(
541-
config: Arc<Config>,
542-
socket: &TcpListener,
543-
server: &mut H,
544-
) -> Result<(), std::io::Error> {
545-
if config.maximum_packet_size > 65535 {
546-
error!(
547-
"Maximum packet size ({:?}) should not larger than a TCP packet (65535)",
548-
config.maximum_packet_size
549-
);
550-
}
551-
552-
let (error_tx, mut error_rx) = tokio::sync::mpsc::unbounded_channel();
553-
554-
loop {
555-
tokio::select! {
556-
accept_result = socket.accept() => {
557-
match accept_result {
558-
Ok((socket, _)) => {
559-
let config = config.clone();
560-
let handler = server.new_client(socket.peer_addr().ok());
561-
let error_tx = error_tx.clone();
562-
tokio::spawn(async move {
563-
let session = match run_stream(config, socket, handler).await {
564-
Ok(s) => s,
565-
Err(e) => {
566-
debug!("Connection setup failed");
567-
let _ = error_tx.send(e);
568-
return
569-
}
570-
};
571-
match session.await {
572-
Ok(_) => debug!("Connection closed"),
573-
Err(e) => {
574-
debug!("Connection closed with error");
575-
let _ = error_tx.send(e);
537+
/// Run a server on a specified `tokio::net::TcpListener`. Useful when dropping
538+
/// privileges immediately after socket binding, for example.
539+
async fn run_on_socket(
540+
&mut self,
541+
config: Arc<Config>,
542+
socket: &TcpListener,
543+
) -> Result<(), std::io::Error> {
544+
if config.maximum_packet_size > 65535 {
545+
error!(
546+
"Maximum packet size ({:?}) should not larger than a TCP packet (65535)",
547+
config.maximum_packet_size
548+
);
549+
}
550+
551+
let (error_tx, mut error_rx) = tokio::sync::mpsc::unbounded_channel();
552+
553+
loop {
554+
tokio::select! {
555+
accept_result = socket.accept() => {
556+
match accept_result {
557+
Ok((socket, _)) => {
558+
let config = config.clone();
559+
let handler = self.new_client(socket.peer_addr().ok());
560+
let error_tx = error_tx.clone();
561+
tokio::spawn(async move {
562+
let session = match run_stream(config, socket, handler).await {
563+
Ok(s) => s,
564+
Err(e) => {
565+
debug!("Connection setup failed");
566+
let _ = error_tx.send(e);
567+
return
568+
}
569+
};
570+
match session.await {
571+
Ok(_) => debug!("Connection closed"),
572+
Err(e) => {
573+
debug!("Connection closed with error");
574+
let _ = error_tx.send(e);
575+
}
576576
}
577-
}
578-
});
577+
});
578+
}
579+
_ => break,
579580
}
580-
_ => break,
581+
},
582+
Some(error) = error_rx.recv() => {
583+
self.handle_session_error(error);
581584
}
582-
},
583-
Some(error) = error_rx.recv() => {
584-
server.handle_session_error(error);
585585
}
586586
}
587-
}
588587

589-
Ok(())
590-
}
588+
Ok(())
589+
}
591590

592-
/// Run a server.
593-
/// Create a new `Connection` from the server's configuration, a
594-
/// stream and a [`Handler`](trait.Handler.html).
595-
pub async fn run<H: Server + Send + 'static, A: ToSocketAddrs>(
596-
config: Arc<Config>,
597-
addrs: A,
598-
server: &mut H,
599-
) -> Result<(), std::io::Error> {
600-
let socket = TcpListener::bind(addrs).await?;
601-
run_on_socket(config, &socket, server).await
591+
/// Run a server.
592+
/// Create a new `Connection` from the server's configuration, a
593+
/// stream and a [`Handler`](trait.Handler.html).
594+
async fn run_on_address<A: ToSocketAddrs + Send>(
595+
&mut self,
596+
config: Arc<Config>,
597+
addrs: A,
598+
) -> Result<(), std::io::Error> {
599+
let socket = TcpListener::bind(addrs).await?;
600+
self.run_on_socket(config, &socket).await
601+
}
602602
}
603603

604604
use std::cell::RefCell;

‎russh/tests/test_data_stream.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ use std::net::{SocketAddr, TcpListener, TcpStream};
22
use std::sync::Arc;
33

44
use rand::RngCore;
5+
use russh::server::Server as _;
56
use russh::server::{self, Auth, Msg, Session};
67
use russh::{client, Channel};
78
use russh_keys::key;
@@ -90,7 +91,7 @@ impl Server {
9091
});
9192
let mut sh = Server {};
9293

93-
russh::server::run(config, addr, &mut sh).await.unwrap();
94+
sh.run_on_address(config, addr).await.unwrap();
9495
}
9596
}
9697

0 commit comments

Comments
 (0)