xref: /aosp_15_r20/tools/netsim/rust/daemon/src/transport/websocket.rs (revision cf78ab8cffb8fc9207af348f23af247fb04370a6)
1 // Copyright 2023 Google LLC
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //     https://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 use std::sync::{Arc, Mutex};
16 use std::{collections::HashMap, io::Cursor, net::TcpStream};
17 
18 use bytes::Bytes;
19 use http::Request;
20 use log::{error, info, warn};
21 use netsim_proto::common::ChipKind;
22 use netsim_proto::startup::DeviceInfo as ProtoDeviceInfo;
23 use tungstenite::{protocol::Role, Message, WebSocket};
24 
25 use crate::devices::chip;
26 use crate::devices::devices_handler::{add_chip, remove_chip};
27 use crate::http_server::server_response::ResponseWritable;
28 use crate::wireless;
29 use crate::wireless::packet::{register_transport, unregister_transport, Response};
30 
31 use super::h4;
32 
33 // This feature is enabled only for CMake builds
34 #[cfg(feature = "local_ssl")]
35 use crate::openssl;
36 
37 /// Generate Sec-Websocket-Accept value from given Sec-Websocket-Key value
generate_websocket_accept(websocket_key: String) -> String38 fn generate_websocket_accept(websocket_key: String) -> String {
39     let concat = websocket_key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
40     let hashed = openssl::sha::sha1(concat.as_bytes());
41     data_encoding::BASE64.encode(&hashed)
42 }
43 
44 /// Handler for websocket server connection
handle_websocket(request: &Request<Vec<u8>>, param: &str, writer: ResponseWritable)45 pub fn handle_websocket(request: &Request<Vec<u8>>, param: &str, writer: ResponseWritable) {
46     if param != "bt" {
47         writer.put_error(404, "netsim websocket currently supports bt, uri=/v1/websocket/bt");
48     }
49     let websocket_accept = match request.headers().get("Sec-Websocket-Key") {
50         Some(key) => match key.to_str() {
51             Ok(key_str) => generate_websocket_accept(key_str.to_string()),
52             Err(_) => {
53                 writer.put_error(
54                     404,
55                     "The HeaderValue of Sec-Websocket-Key cannot be converted to str",
56                 );
57                 return;
58             }
59         },
60         None => {
61             writer.put_error(404, "Missing Sec-Websocket-Key in header");
62             return;
63         }
64     };
65     writer.put_ok_switch_protocol(
66         "websocket",
67         vec![("Sec-WebSocket-Accept".to_string(), websocket_accept)],
68     )
69 }
70 
71 struct WebSocketTransport {
72     websocket_writer: Arc<Mutex<WebSocket<TcpStream>>>,
73 }
74 
75 impl Response for WebSocketTransport {
response(&mut self, packet: Bytes, packet_type: u8)76     fn response(&mut self, packet: Bytes, packet_type: u8) {
77         let mut buffer = Vec::new();
78         buffer.push(packet_type);
79         buffer.extend(packet);
80         if let Err(err) = self
81             .websocket_writer
82             .lock()
83             .expect("Failed to acquire lock on WebSocket")
84             .send(Message::Binary(buffer))
85         {
86             error!("{err}");
87         };
88     }
89 }
90 
91 /// Run websocket transport for packet flow in netsim
run_websocket_transport(stream: TcpStream, queries: HashMap<&str, &str>)92 pub fn run_websocket_transport(stream: TcpStream, queries: HashMap<&str, &str>) {
93     let chip_create_params = chip::CreateParams {
94         kind: ChipKind::BLUETOOTH,
95         address: queries.get("address").unwrap_or(&"").to_string(),
96         name: Some(format!("websocket-{}", stream.peer_addr().unwrap())),
97         manufacturer: "Google".to_string(),
98         product_name: "Google".to_string(),
99     };
100     #[cfg(not(test))]
101     let wireless_create_params =
102         wireless::CreateParam::Bluetooth(wireless::bluetooth::CreateParams {
103             address: chip_create_params.address.clone(),
104             bt_properties: None,
105         });
106     #[cfg(test)]
107     let wireless_create_params = wireless::CreateParam::Mock(wireless::mocked::CreateParams {
108         chip_kind: ChipKind::BLUETOOTH,
109     });
110     let device_info = ProtoDeviceInfo { kind: "WEBSOCKET".to_string(), ..Default::default() };
111     // Add Chip
112     let result = match add_chip(
113         &stream.peer_addr().unwrap().port().to_string(),
114         queries
115             .get("name")
116             .unwrap_or(&format!("websocket-device-{}", stream.peer_addr().unwrap()).as_str()),
117         &chip_create_params,
118         &wireless_create_params,
119         device_info,
120     ) {
121         Ok(chip_result) => chip_result,
122         Err(err) => {
123             warn!("{err}");
124             return;
125         }
126     };
127 
128     // Create websocket_writer to handle packet responses, write pong or close messages
129     let websocket_writer = Arc::new(Mutex::new(WebSocket::from_raw_socket(
130         stream.try_clone().unwrap(),
131         Role::Server,
132         None,
133     )));
134     // Websocket reader
135     let mut websocket_reader = WebSocket::from_raw_socket(stream, Role::Server, None);
136 
137     // Sending cloned websocket into packet dispatcher
138     register_transport(
139         result.chip_id,
140         Box::new(WebSocketTransport { websocket_writer: websocket_writer.clone() }),
141     );
142 
143     // Running Websocket server
144     loop {
145         let packet_msg =
146             match websocket_reader.read().map_err(|_| "Failed to read Websocket message") {
147                 Ok(message) => message,
148                 Err(err) => {
149                     error!("{err}");
150                     break;
151                 }
152             };
153         if packet_msg.is_binary() {
154             let mut cursor = Cursor::new(packet_msg.into_data());
155             match h4::read_h4_packet(&mut cursor) {
156                 Ok(packet) => {
157                     wireless::handle_request(result.chip_id, &packet.payload, packet.h4_type);
158                 }
159                 Err(error) => {
160                     error!(
161                         "netsimd: end websocket reader {}: {:?}",
162                         websocket_reader.get_ref().peer_addr().unwrap(),
163                         error
164                     );
165                     break;
166                 }
167             }
168         } else if packet_msg.is_ping() {
169             if let Err(err) = websocket_writer
170                 .lock()
171                 .expect("Failed to acquire lock on WebSocket")
172                 .send(Message::Pong(packet_msg.into_data()))
173             {
174                 error!("{err}");
175             }
176         } else if packet_msg.is_close() {
177             if let Message::Close(close_frame) = packet_msg {
178                 if let Err(err) = websocket_writer
179                     .lock()
180                     .expect("Failed to acquire lock on WebSocket")
181                     .close(close_frame)
182                     .map_err(|_| "Failed to close Websocket")
183                 {
184                     error!("{err}");
185                 }
186             }
187             break;
188         }
189     }
190 
191     // unregister before remove_chip because facade may re-use facade_id
192     // on an intertwining create_chip and the unregister here might remove
193     // the recently added chip creating a disconnected transport.
194     unregister_transport(result.chip_id);
195 
196     if let Err(err) = remove_chip(result.device_id, result.chip_id) {
197         warn!("{err}");
198     };
199     info!("Removed chip: device_id: {}, chip_id: {}", result.device_id, result.chip_id);
200 }
201