Skip to content

Commit 085dd79

Browse files
author
Hang SU
committed
console: Refactor backend handling with pattern matching
Rework conditional logic using match expressions for backend type dispatch. No functional changes, this patch is a preparation for new backend type: uds. Signed-off-by: Hang SU <[email protected]>
1 parent 5a71380 commit 085dd79

File tree

2 files changed

+132
-88
lines changed

2 files changed

+132
-88
lines changed

‎vhost-device-console/src/backend.rs

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -74,14 +74,19 @@ impl VuConsoleConfig {
7474
}
7575

7676
pub fn generate_vm_socks(&self) -> Vec<String> {
77-
let tcp_port_base = self.tcp_port.clone();
78-
79-
let make_tcp_port = |i: u32| -> String {
80-
let port_num: u32 = tcp_port_base.clone().parse().unwrap();
81-
"127.0.0.1:".to_owned() + &(port_num + i).to_string()
82-
};
83-
84-
(0..self.socket_count).map(make_tcp_port).collect()
77+
match self.backend {
78+
// if type is Nested, result will be dropped.
79+
BackendType::Nested => {
80+
vec![String::new()]
81+
}
82+
83+
BackendType::Network => {
84+
let port_base: u32 = self.tcp_port.parse().unwrap();
85+
let make_tcp_port =
86+
|i: u32| -> String { "127.0.0.1:".to_owned() + &(port_base + i).to_string() };
87+
(0..self.socket_count).map(make_tcp_port).collect()
88+
}
89+
}
8590
}
8691
}
8792

‎vhost-device-console/src/vhu_console.rs

Lines changed: 119 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -174,19 +174,22 @@ impl VhostUserConsoleBackend {
174174
}
175175

176176
pub fn assign_input_method(&mut self, vm_sock: String) -> Result<()> {
177-
if self.controller.read().unwrap().backend == BackendType::Nested {
178-
// Enable raw mode for local terminal if backend is nested
179-
enable_raw_mode().expect("Raw mode error");
180-
181-
let stdin_fd = io::stdin().as_raw_fd();
182-
let stdin: Box<dyn Read + Send + Sync> = Box::new(io::stdin());
183-
self.stdin = Some(stdin);
177+
let backend_type = self.controller.read().unwrap().backend;
178+
match backend_type {
179+
BackendType::Nested => {
180+
// Enable raw mode for local terminal if backend is nested
181+
enable_raw_mode().expect("Raw mode error");
182+
let stdin_fd = io::stdin().as_raw_fd();
183+
let stdin: Box<dyn Read + Send + Sync> = Box::new(io::stdin());
184+
self.stdin = Some(stdin);
185+
Self::epoll_register(self.epoll_fd.as_raw_fd(), stdin_fd, epoll::Events::EPOLLIN)
186+
.map_err(|_| Error::EpollAdd)?;
187+
}
184188

185-
Self::epoll_register(self.epoll_fd.as_raw_fd(), stdin_fd, epoll::Events::EPOLLIN)
186-
.map_err(|_| Error::EpollAdd)?;
187-
} else {
188-
let listener = TcpListener::bind(vm_sock).expect("Failed bind tcp address");
189-
self.listener = Some(listener);
189+
BackendType::Network => {
190+
let tcp_listener = TcpListener::bind(vm_sock).expect("Failed bind tcp address");
191+
self.tcp_listener = Some(tcp_listener);
192+
}
190193
}
191194
Ok(())
192195
}
@@ -264,14 +267,19 @@ impl VhostUserConsoleBackend {
264267
}
265268

266269
let my_string = String::from_utf8(tx_data).unwrap();
267-
if self.controller.read().unwrap().backend == BackendType::Nested {
268-
print!("{}", my_string);
269-
io::stdout().flush().unwrap();
270-
} else {
271-
self.output_queue
272-
.add(my_string)
273-
.map_err(|_| Error::RxCtrlQueueAddFailed)?;
274-
self.write_tcp_stream();
270+
let backend_type = self.controller.read().unwrap().backend;
271+
match backend_type {
272+
BackendType::Nested => {
273+
print!("{}", my_string);
274+
io::stdout().flush().unwrap();
275+
}
276+
277+
BackendType::Network => {
278+
self.output_queue
279+
.add(my_string)
280+
.map_err(|_| Error::RxCtrlQueueAddFailed)?;
281+
self.write_stream();
282+
}
275283
}
276284

277285
vring
@@ -529,19 +537,24 @@ impl VhostUserConsoleBackend {
529537
.register_listener(epoll_fd, EventSet::IN, u64::from(QueueEvents::KEY_EFD))
530538
.unwrap();
531539

532-
if self.controller.read().unwrap().backend == BackendType::Network {
533-
let listener_fd = self
534-
.tcp_listener
535-
.as_ref()
536-
.expect("Failed get tcp listener ref")
537-
.as_raw_fd();
538-
vring_worker
539-
.register_listener(
540-
listener_fd,
541-
EventSet::IN,
542-
u64::from(QueueEvents::LISTENER_EFD),
543-
)
544-
.unwrap();
540+
let backend_type = self.controller.read().unwrap().backend;
541+
match backend_type {
542+
BackendType::Nested => {}
543+
544+
BackendType::Network => {
545+
let listener_fd = self
546+
.tcp_listener
547+
.as_ref()
548+
.expect("Failed get tcp listener ref")
549+
.as_raw_fd();
550+
vring_worker
551+
.register_listener(
552+
listener_fd,
553+
EventSet::IN,
554+
u64::from(QueueEvents::LISTENER_EFD),
555+
)
556+
.unwrap();
557+
}
545558
}
546559
}
547560

@@ -570,42 +583,50 @@ impl VhostUserConsoleBackend {
570583
Ok(())
571584
}
572585

586+
/// Sets up a new stream connection with proper epoll registration
587+
fn handle_stream_connection<T>(&mut self, stream: T, addr_desc: String)
588+
where
589+
T: ReadWrite + Send + Sync + AsRawFd + 'static,
590+
{
591+
println!("New connection on: {}", addr_desc);
592+
let stream_raw_fd = stream.as_raw_fd();
593+
self.stream_fd = Some(stream_raw_fd);
594+
595+
if let Err(err) = Self::epoll_register(
596+
self.epoll_fd.as_raw_fd(),
597+
stream_raw_fd,
598+
epoll::Events::EPOLLIN,
599+
) {
600+
warn!("Failed to register with epoll: {:?}", err);
601+
}
602+
603+
self.stream = Some(Box::new(stream));
604+
self.write_stream();
605+
}
606+
573607
fn create_new_stream_thread(&mut self) {
574608
// Accept only one incoming connection
575-
if let Some(stream) = self
576-
.tcp_listener
577-
.as_ref()
578-
.expect("Failed get tcp listener ref")
579-
.incoming()
580-
.next()
581-
{
582-
match stream {
583-
Ok(stream) => {
584-
let local_addr = self
585-
.tcp_listener
586-
.as_ref()
587-
.expect("No listener")
588-
.local_addr()
589-
.unwrap();
590-
println!("New connection on: {}", local_addr);
591-
let stream_raw_fd = stream.as_raw_fd();
592-
self.stream_fd = Some(stream_raw_fd);
593-
if let Err(err) = Self::epoll_register(
594-
self.epoll_fd.as_raw_fd(),
595-
stream_raw_fd,
596-
epoll::Events::EPOLLIN,
597-
) {
598-
warn!("Failed to register with epoll: {:?}", err);
599-
}
609+
let backend_type = self.controller.read().unwrap().backend;
610+
match backend_type {
611+
BackendType::Nested => {}
600612

601-
let stream: Box<dyn ReadWrite + Send + Sync> = Box::new(stream);
602-
self.stream = Some(stream);
603-
self.write_stream();
604-
}
605-
Err(e) => {
606-
eprintln!("Stream error: {}", e);
613+
BackendType::Network => match self
614+
.tcp_listener
615+
.as_ref()
616+
.expect("No tcp listener")
617+
.incoming()
618+
.next()
619+
{
620+
Some(Ok(tcp_stream)) => {
621+
let addr_desc = tcp_stream
622+
.peer_addr()
623+
.map(|addr| format!("TCP {}", addr))
624+
.unwrap_or_else(|_| "unknown TCP peer".to_string());
625+
self.handle_stream_connection(tcp_stream, addr_desc);
607626
}
608-
}
627+
Some(Err(e)) => eprintln!("TCP stream error: {}", e),
628+
None => {}
629+
},
609630
}
610631
}
611632

@@ -635,13 +656,23 @@ impl VhostUserConsoleBackend {
635656
match self.stream.as_mut().expect("No stream").read(&mut buffer) {
636657
Ok(bytes_read) => {
637658
if bytes_read == 0 {
638-
let local_addr = self
639-
.tcp_listener
640-
.as_ref()
641-
.expect("No listener")
642-
.local_addr()
643-
.unwrap();
644-
println!("Close connection on: {}", local_addr);
659+
let backend_type = self.controller.read().unwrap().backend;
660+
661+
// Get connection address for logging
662+
let conn_addr = match backend_type {
663+
BackendType::Nested => String::new(),
664+
BackendType::Network => self
665+
.tcp_listener
666+
.as_ref()
667+
.and_then(|l| l.local_addr().ok())
668+
.map(|addr| addr.to_string())
669+
.unwrap_or_else(|| "unknown TCP address".to_string()),
670+
};
671+
672+
if !conn_addr.is_empty() {
673+
println!("Close connection on: {}", conn_addr);
674+
}
675+
645676
if let Err(err) = Self::epoll_unregister(
646677
self.epoll_fd.as_raw_fd(),
647678
self.stream_fd.expect("No stream fd"),
@@ -692,9 +723,13 @@ impl VhostUserConsoleBackend {
692723
}
693724

694725
pub fn prepare_exit(&self) {
695-
/* For the nested backend */
696-
if self.controller.read().unwrap().backend == BackendType::Nested {
697-
disable_raw_mode().expect("Raw mode error");
726+
let backend_type = self.controller.read().unwrap().backend;
727+
match backend_type {
728+
/* For the nested backend */
729+
BackendType::Nested => {
730+
disable_raw_mode().expect("Cannot reset terminal");
731+
}
732+
BackendType::Network => {}
698733
}
699734
}
700735
}
@@ -782,11 +817,15 @@ impl VhostUserBackendMut for VhostUserConsoleBackend {
782817
}
783818

784819
if device_event == QueueEvents::KEY_EFD {
785-
if self.controller.read().unwrap().backend == BackendType::Nested {
786-
return self.read_char_thread();
787-
} else {
788-
self.read_stream();
789-
return Ok(());
820+
let backend_type = self.controller.read().unwrap().backend;
821+
match backend_type {
822+
BackendType::Nested => {
823+
return self.read_char_thread();
824+
}
825+
BackendType::Network => {
826+
self.read_stream();
827+
return Ok(());
828+
}
790829
}
791830
}
792831

0 commit comments

Comments
 (0)