Source code

Revision control

Copy as Markdown

Other Tools

/* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this
* file, You can obtain one at http://mozilla.org/MPL/2.0/. */
use crate::widestring;
use serde_json::Value;
use std::ffi::{c_void, CString, NulError};
use std::num::TryFromIntError;
use std::ptr;
use std::string::FromUtf8Error;
use thiserror::Error;
use url::{ParseError, Url};
use windows_sys::Win32::{Foundation::*, Networking::WinHttp::*};
fn check_win_api() -> u32 {
unsafe { GetLastError() }
}
// TODO: drop once windows-sys names its own import libraries (Bug 2071329).
#[link(name = "winhttp")]
unsafe extern "system" {}
#[derive(Error, Debug)]
pub enum Error {
#[error("WinHttpWebSocketCompleteUpgrade failed: {0}")]
CompleteUpgrade(u32),
#[error("Invalid number")]
IntParseError(#[from] TryFromIntError),
#[error("Invalid scheme. Not wss:// or ws://")]
InvalidScheme,
#[error("Invalid URL")]
InvalidURL(#[from] ParseError),
#[error("WinHttpConnect failed: {0}")]
NoConnection(u32),
#[error("No hostname")]
NoHostName,
#[error("WinHttpOpenRequest failed: {0}")]
NoRequest(u32),
#[error("WinHttpOpen failed: {0}")]
NoSession(u32),
#[error("No Websocket")]
NoWebSocket,
#[error("Unexpected null error")]
NulError(#[from] NulError),
#[error("WinHttpReceiveResponse failed: {0}")]
ReceiveResponse(u32),
#[error("WinHttpSendRequest failed: {0}")]
SendRequest(u32),
#[error("WinHttpSetOption failed: {0}")]
SetOption(u32),
#[error("Invalid UTF8")]
Utf8Error(#[from] FromUtf8Error),
#[error("WinHttpWebSocketReceive failed with error code {0}")]
WebSocketReceive(u32),
#[error("WinHttpWebSocketSend failed with error code {0}")]
WebSocketSend(u32),
}
pub enum Event {
Closed,
Notification,
Other(String),
Uaid(String),
}
pub struct PushConnection {
websocket: *mut c_void,
connection: *mut c_void,
session: *mut c_void,
}
impl PushConnection {
pub fn new() -> PushConnection {
PushConnection {
websocket: ptr::null_mut(),
connection: ptr::null_mut(),
session: ptr::null_mut(),
}
}
pub fn connect(&mut self, url_str: &str) -> Result<(), Error> {
// Close the connection if one already existed to avoid leaking
self.close();
let url = Url::parse(url_str)?;
if url.scheme() != "ws" && url.scheme() != "wss" {
return Err(Error::InvalidScheme);
}
let user_agent = widestring::WideString::new("Firefox background notification");
self.session = unsafe {
WinHttpOpen(
user_agent.pcwstr(),
WINHTTP_ACCESS_TYPE_AUTOMATIC_PROXY,
ptr::null(),
ptr::null(),
0,
)
};
if self.session == ptr::null_mut() {
return Err(Error::NoSession(check_win_api()));
}
let Some(host) = url.host_str() else {
return Err(Error::NoHostName);
};
let hostname = widestring::WideString::new(host);
let secure = url.scheme() == "wss";
// ws and wss are a known default.
let port = url.port_or_known_default().unwrap_or(if secure {
INTERNET_DEFAULT_HTTPS_PORT
} else {
INTERNET_DEFAULT_HTTP_PORT
});
self.connection = unsafe { WinHttpConnect(self.session, hostname.pcwstr(), port, 0) };
if self.connection == ptr::null_mut() {
return Err(Error::NoConnection(check_win_api()));
}
let verb = widestring::WideString::new("GET");
let path = widestring::WideString::new(&url[url::Position::BeforePath..]);
let flags = if secure { WINHTTP_FLAG_SECURE } else { 0 };
let request = unsafe {
WinHttpOpenRequest(
self.connection,
verb.pcwstr(),
path.pcwstr(),
ptr::null(),
ptr::null(),
ptr::null(),
flags,
)
};
if request == ptr::null_mut() {
return Err(Error::NoRequest(check_win_api()));
}
if 0 == unsafe {
WinHttpSetOption(
request,
WINHTTP_OPTION_UPGRADE_TO_WEB_SOCKET,
ptr::null(),
0,
)
} {
return Err(Error::SetOption(check_win_api()));
}
if 0 == unsafe { WinHttpSendRequest(request, ptr::null(), 0, ptr::null(), 0, 0, 0) } {
return Err(Error::SendRequest(check_win_api()));
}
if 0 == unsafe { WinHttpReceiveResponse(request, ptr::null_mut()) } {
return Err(Error::ReceiveResponse(check_win_api()));
}
self.websocket = unsafe { WinHttpWebSocketCompleteUpgrade(request, 0) };
if self.websocket == ptr::null_mut() {
return Err(Error::CompleteUpgrade(check_win_api()));
}
unsafe { WinHttpCloseHandle(request) };
Ok(())
}
pub fn send(&self, message: &str) -> Result<bool, Error> {
if self.websocket == ptr::null_mut() {
return Err(Error::NoWebSocket);
}
let message_c = CString::new(message)?;
let result = unsafe {
WinHttpWebSocketSend(
self.websocket,
WINHTTP_WEB_SOCKET_UTF8_MESSAGE_BUFFER_TYPE,
message_c.as_ptr() as *const c_void,
message.len().try_into()?,
)
};
if result != NO_ERROR {
return Err(Error::WebSocketSend(result));
}
return Ok(result == NO_ERROR);
}
pub fn send_hello(&self, uaid: &str) -> Result<bool, Error> {
let hello_msg = format!("{{\"messageType\":\"hello\",\"broadcasts\":{{}},\"use_webpush\":true,\"uaid\":\"{uaid}\"}}");
return self.send(hello_msg.as_str());
}
/*
pub fn send_keep_alive(&self) -> Result<bool, Error> {
return self.send("{}");
}
*/
fn parse_message(message: &str) -> Event {
let parsed_message = serde_json::from_str(message).unwrap_or_else(|_| Value::Null);
let message_type = match &parsed_message["messageType"] {
Value::String(mtype) => mtype,
_ => {
return Event::Other(message.to_owned());
}
};
return match message_type.as_str() {
"notification" => Event::Notification,
"hello" => {
let Some(uaid) = parsed_message["uaid"].as_str() else {
return Event::Other(message.to_owned());
};
Event::Uaid(uaid.to_owned())
}
_ => Event::Other(message.to_owned()),
};
}
pub fn wait_for_message(&self) -> Result<Event, Error> {
if self.websocket == ptr::null_mut() {
return Err(Error::NoWebSocket);
}
let mut bytes_read: u32 = 0;
let buffer_size: u32 = 8192;
let mut buffer_type: WINHTTP_WEB_SOCKET_BUFFER_TYPE = WINHTTP_WEB_SOCKET_CLOSE_BUFFER_TYPE;
let mut message_buffer: Vec<u8> = vec![];
let mut buffer: Vec<u8> = vec![0; buffer_size as usize];
loop {
let result = unsafe {
WinHttpWebSocketReceive(
self.websocket,
buffer.as_mut_ptr() as *mut c_void,
buffer_size,
&mut bytes_read as *mut u32,
&mut buffer_type as *mut WINHTTP_WEB_SOCKET_BUFFER_TYPE,
)
};
if result != NO_ERROR {
return Err(Error::WebSocketReceive(result));
}
if buffer_type == WINHTTP_WEB_SOCKET_CLOSE_BUFFER_TYPE {
return Ok(Event::Closed);
}
message_buffer.extend(&buffer[..bytes_read as usize]);
if buffer_type != WINHTTP_WEB_SOCKET_UTF8_MESSAGE_BUFFER_TYPE
&& buffer_type != WINHTTP_WEB_SOCKET_BINARY_MESSAGE_BUFFER_TYPE
{
continue;
}
let message = String::from_utf8(message_buffer)?;
println!("Message: {message}");
return Ok(Self::parse_message(&message));
}
}
pub fn close(&mut self) {
if self.websocket != ptr::null_mut() {
unsafe {
WinHttpWebSocketClose(
self.websocket,
WINHTTP_WEB_SOCKET_SUCCESS_CLOSE_STATUS as u16,
ptr::null(),
0,
);
WinHttpCloseHandle(self.websocket);
}
self.websocket = ptr::null_mut();
}
if self.connection != ptr::null_mut() {
unsafe {
WinHttpCloseHandle(self.connection);
self.connection = ptr::null_mut();
}
}
if self.session != ptr::null_mut() {
unsafe {
WinHttpCloseHandle(self.session);
self.session = ptr::null_mut();
}
}
}
}
impl Drop for PushConnection {
fn drop(&mut self) {
self.close();
}
}