Release v0.16.0

Misc corrections

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "sn0int"
-version = "0.15.0"
+version = "0.16.0"
 description = "Semi-automatic OSINT framework and package manager"
 authors = ["kpcyrd <git@rxv.cc>"]
 license = "GPL-3.0"
@@ -40,10 +40,11 @@ dirs = "2.0"
 url = "2.0"
 percent-encoding = "2.1"
 #chrootable-https = { path = "../chrootable-https" }
-chrootable-https = "0.13"
+chrootable-https = "0.14"
 rustls = { version="0.16", features=["dangerous_configuration"] }
 webpki = "0.21"
 webpki-roots = "0.18"
+ct-logs = "0.6"
 pem = "0.7"
 base64 = "0.11"
 data-encoding = "2.1.2"
@@ -74,6 +75,7 @@ bytesize = "1.0"
 ipnetwork = "0.15"
 strum = "0.17"
 strum_macros = "0.17"
+tungstenite = { version = "0.9", default-features = false }
 
 digest = "0.8.0"
 bs58 = "0.3"

README.md

@@ -203,6 +203,7 @@ For everything else please have a look at the [detailed list][1].
   - [sn0int_version](https://sn0int.readthedocs.io/en/latest/reference.html#sn0int-version)
   - [sock_connect](https://sn0int.readthedocs.io/en/latest/reference.html#sock-connect)
   - [sock_upgrade_tls](https://sn0int.readthedocs.io/en/latest/reference.html#sock-upgrade-tls)
+  - [sock_options](https://sn0int.readthedocs.io/en/latest/reference.html#sock-options)
   - [sock_send](https://sn0int.readthedocs.io/en/latest/reference.html#sock-send)
   - [sock_recv](https://sn0int.readthedocs.io/en/latest/reference.html#sock-recv)
   - [sock_sendline](https://sn0int.readthedocs.io/en/latest/reference.html#sock-sendline)
@@ -229,6 +230,14 @@ For everything else please have a look at the [detailed list][1].
   - [utf8_decode](https://sn0int.readthedocs.io/en/latest/reference.html#utf8-decode)
   - [warn](https://sn0int.readthedocs.io/en/latest/reference.html#warn)
   - [warn_once](https://sn0int.readthedocs.io/en/latest/reference.html#warn-once)
+  - [ws_connect](https://sn0int.readthedocs.io/en/latest/reference.html#ws-connect)
+  - [ws_options](https://sn0int.readthedocs.io/en/latest/reference.html#ws-options)
+  - [ws_recv_text](https://sn0int.readthedocs.io/en/latest/reference.html#ws-recv-text)
+  - [ws_recv_binary](https://sn0int.readthedocs.io/en/latest/reference.html#ws-recv-binary)
+  - [ws_recv_json](https://sn0int.readthedocs.io/en/latest/reference.html#ws-recv-json)
+  - [ws_send_text](https://sn0int.readthedocs.io/en/latest/reference.html#ws-send-text)
+  - [ws_send_binary](https://sn0int.readthedocs.io/en/latest/reference.html#ws-send-binary)
+  - [ws_send_json](https://sn0int.readthedocs.io/en/latest/reference.html#ws-send-json)
   - [x509_parse_pem](https://sn0int.readthedocs.io/en/latest/reference.html#x509-parse-pem)
   - [xml_decode](https://sn0int.readthedocs.io/en/latest/reference.html#xml-decode)
   - [xml_named](https://sn0int.readthedocs.io/en/latest/reference.html#xml-named)

docs/reference.rst

@@ -834,6 +834,13 @@ The following options are available:
 ``proxy``
   Use a socks5 proxy in the format ``127.0.0.1:9050``. This option only works
   if it doesn't conflict with the global proxy settings.
+``connect_timeout``
+  Abort tcp connection attempts after ``n`` seconds.
+``read_timeout``
+  Abort read attempts after ``n`` seconds. This can be used to wake up
+  connections periodically.
+``write_timeout``
+  Abort write attempts after ``n`` seconds.
 
 .. code-block:: lua
 
@@ -865,6 +872,23 @@ discarded when using sock_connect_ directly with ``tls=true``.
 
     info(tls)
 
+sock_options
+------------
+
+Update options of an existing connection:
+
+``read_timeout``
+  Abort read attempts after ``n`` seconds. This can be used to wake up
+  connections periodically.
+``write_timeout``
+  Abort write attempts after ``n`` seconds.
+
+.. code-block:: lua
+
+    sock_options(sock, {
+        read_timeout=3,
+    })
+
 sock_send
 ---------
 
@@ -1141,6 +1165,111 @@ a ``run`` execution.
     warn_once('ohai')
     warn_once('ohai')
 
+ws_connect
+----------
+
+Create a websocket connection. The url format is ``ws://example.com/asdf``,
+``wss://`` is also supported.
+
+The following options are available:
+
+``headers``
+  A map of additional headers that should be set for the request.
+``proxy``
+  Use a socks5 proxy in the format ``127.0.0.1:9050``. This option only works
+  if it doesn't conflict with the global proxy settings.
+``connect_timeout``
+  Abort tcp connection attempts after ``n`` seconds.
+``read_timeout``
+  Abort read attempts after ``n`` seconds. This can be used to wake up
+  connections periodically.
+``write_timeout``
+  Abort write attempts after ``n`` seconds.
+
+.. code-block:: lua
+
+    sock = ws_connect("wss://example.com/asdf", {})
+
+ws_options
+----------
+
+Update options of an existing connection:
+
+``read_timeout``
+  Abort read attempts after ``n`` seconds. This can be used to wake up
+  connections periodically.
+``write_timeout``
+  Abort write attempts after ``n`` seconds.
+
+.. code-block:: lua
+
+    ws_options(sock, {
+        read_timeout=3,
+    })
+
+ws_recv_text
+------------
+
+Wait until the server sends a text frame. A binary frame is considered an
+error. Ping requests are answered automatically.
+
+.. code-block:: lua
+
+    msg = ws_recv_text(sock)
+
+ws_recv_binary
+--------------
+
+Wait until the server sends a binary frame. A text frame is considered an
+error. Ping requests are answered automatically.
+
+.. code-block:: lua
+
+    msg = ws_recv_binary(sock)
+
+ws_recv_json
+------------
+
+Identical to ws_send_text_ but automatically runs json_decode_ on the
+response.
+
+.. code-block:: lua
+
+    msg = ws_recv_json(sock)
+
+ws_send_text
+------------
+
+Send a text frame on the websocket connection.
+
+.. code-block:: lua
+
+    ws_send_text(sock, "ohai!")
+
+ws_send_binary
+--------------
+
+Send a binary frame on the websocket connection.
+
+.. code-block:: lua
+
+    ws_send_binary(sock, "\x00\x01\x02")
+
+ws_send_json
+------------
+
+Encode the object as json string and send it as a text frame on the websocket
+connection.
+
+.. code-block:: lua
+
+    ws_send_text(sock, {
+        foo="ohai!",
+        x={
+            y={1,3,3,7},
+        },
+    })
+
 x509_parse_pem
 --------------
 

modules/harness/google-tls.lua

@@ -0,0 +1,29 @@
+-- Description: Test various tls functions
+-- Version: 0.1.0
+-- License: GPL-3.0
+
+function run()
+    info('sending https request to google.com')
+    session = http_mksession()
+    req = http_request(session, 'GET', 'https://google.com/', {})
+    r = http_send(req)
+    if last_err() then return end
+    debug(r)
+
+    info('creating tls socket to google.com')
+    sock = sock_connect('google.com', 443, {
+        tls=true,
+    })
+    if last_err() then return end
+    debug(sock)
+
+    info('creating socket to google.com, wrapping afterwards')
+    sock = sock_connect('google.com', 443, {})
+    if last_err() then return end
+    tls = sock_upgrade_tls(sock, {
+        sni_value='google.com',
+    })
+    if last_err() then return end
+    debug(sock)
+    debug(tls)
+end

modules/harness/socket-ping.lua

@@ -0,0 +1,33 @@
+-- Description: Connect somewhere and send a ping every 3s
+-- Version: 0.1.0
+-- License: GPL-3.0
+
+INTERVAL = 3
+
+function run()
+    local sock = sock_connect('127.0.0.1', 4444, {
+        read_timeout=INTERVAL,
+    })
+    if last_err() then return end
+
+    local last_ping = time_unix()
+    while true do
+        local now = time_unix()
+        local sleep = last_ping + INTERVAL - now
+
+        if sleep <= 0 then
+            sock_send(sock, sn0int_time() .. ' ping\n')
+            last_ping = now
+            sleep = INTERVAL
+        end
+
+        sock_options(sock, {
+            read_timeout=sleep,
+        })
+        if last_err() then return end
+
+        local buf = sock_recv(sock)
+        if last_err() then return end
+        info(buf)
+    end
+end

modules/harness/ws-ping.lua

@@ -0,0 +1,33 @@
+-- Description: Connect somewhere and send a ping every 3s
+-- Version: 0.1.0
+-- License: GPL-3.0
+
+INTERVAL = 3
+
+function run()
+    local sock = ws_connect('ws://127.0.0.1:8080', {
+        read_timeout=INTERVAL,
+    })
+    if last_err() then return end
+
+    local last_ping = time_unix()
+    while true do
+        local now = time_unix()
+        local sleep = last_ping + INTERVAL - now
+
+        if sleep <= 0 then
+            ws_send_text(sock, sn0int_time() .. ' ping\n')
+            last_ping = now
+            sleep = INTERVAL
+        end
+
+        ws_options(sock, {
+            read_timeout=sleep,
+        })
+        if last_err() then return end
+
+        local buf = ws_recv_text(sock)
+        if last_err() then return end
+        info(buf)
+    end
+end

modules/harness/ws.lua

@@ -0,0 +1,23 @@
+-- Description: Create an echo websocket connection
+-- Version: 0.1.0
+-- License: GPL-3.0
+
+function run()
+    local target = 'ws://echo.websocket.org'
+
+    info('connecting to ' .. target)
+    local sock = ws_connect(target, {})
+    if last_err() then return end
+
+    info('sending')
+    ws_send_text(sock, 'ohai wurld')
+    if last_err() then return end
+
+    info('recieving')
+    local msg = ws_recv_text(sock)
+    if last_err() then return end
+
+    if msg ~= 'ohai wurld' then
+        return 'echo failed, got: ' .. msg
+    end
+end

modules/harness/wss.lua

@@ -0,0 +1,30 @@
+-- Description: Create an encrypted websocket connection
+-- Version: 0.1.0
+-- License: GPL-3.0
+
+function run()
+    -- local target = 'ws://echo.websocket.org' -- doesn't support proper ciphers
+    local target = 'wss://rocket.events.ccc.de/sockjs/258/whi0yr1y/websocket'
+
+    info('connecting to ' .. target)
+    local sock = ws_connect(target, {})
+    if last_err() then return end
+
+    info('recieving 1/2')
+    local msg = ws_recv_text(sock)
+    if last_err() then return end
+
+    if msg ~= 'o' then
+        return 'recieve failed, got ' .. msg
+    end
+
+    info('recieving 2/2')
+    local msg = ws_recv_text(sock)
+    if last_err() then return end
+
+    if msg ~= 'a["{\\"server_id\\":\\"0\\"}"]' then
+        return 'recieve failed, got ' .. msg
+    end
+
+    info('handshake succeeded')
+end

src/db/mod.rs

@@ -121,6 +121,8 @@ impl Database {
             .context("Failed to enable write ahead log")?;
         db.execute("PRAGMA foreign_keys = ON")
             .context("Failed to enforce foreign keys")?;
+        db.execute("PRAGMA synchronous = NORMAL")
+            .context("Failed to enforce foreign keys")?;
 
         let autonoscope = RuleSet::load(&db)?;
 

src/engine/ctx.rs

@@ -12,6 +12,7 @@ use crate::lazy::Lazy;
 use crate::runtime;
 use crate::sockets::{Socket, SocketOptions, TlsData};
 use crate::web::{HttpSession, HttpRequest, RequestOptions};
+use crate::websockets::{WebSocket, WebSocketOptions};
 use crate::worker::{Event, LogEvent, DatabaseEvent, StdioEvent, RatelimitEvent};
 use crate::ratelimits::RatelimitResponse;
 use chrootable_https::{self, Resolver};
@@ -168,6 +169,10 @@ pub trait State {
 
     fn sock_upgrade_tls(&self, id: &str, options: &SocketOptions) -> Result<TlsData>;
 
+    fn ws_connect(&self, url: url::Url, options: &WebSocketOptions) -> Result<String>;
+
+    fn get_ws(&self, id: &str)-> Arc<Mutex<WebSocket>>;
+
     fn http(&self, proxy: &Option<SocketAddr>) -> Result<Arc<chrootable_https::Client<Resolver>>>;
 
     fn http_mksession(&self) -> String;
@@ -194,6 +199,7 @@ pub struct LuaState {
     error: Mutex<Option<Error>>,
     logger: Arc<Mutex<Box<dyn Reporter>>>,
     socket_sessions: Mutex<HashMap<String, Arc<Mutex<Socket>>>>,
+    ws_sessions: Mutex<HashMap<String, Arc<Mutex<WebSocket>>>>,
     blobs: Mutex<HashMap<String, Arc<Blob>>>,
     http_sessions: Mutex<HashMap<String, HttpSession>>,
     http_clients: Mutex<HashMap<String, Arc<chrootable_https::Client<Resolver>>>>,
@@ -301,13 +307,13 @@ impl State for LuaState {
 
     fn get_sock(&self, id: &str)-> Arc<Mutex<Socket>> {
         let mtx = self.socket_sessions.lock().unwrap();
-        let sock = mtx.get(id).expect("invalid session reference"); // TODO
+        let sock = mtx.get(id).expect("Invalid socket reference"); // TODO
         sock.clone()
     }
 
     fn sock_upgrade_tls(&self, id: &str, options: &SocketOptions) -> Result<TlsData> {
         let mut mtx = self.socket_sessions.lock().unwrap();
-        let sock = mtx.remove(id).expect("invalid session reference"); // TODO
+        let sock = mtx.remove(id).expect("Invalid socket reference"); // TODO
 
         let sock = Arc::try_unwrap(sock).unwrap();
         let sock = sock.into_inner().unwrap();
@@ -319,6 +325,22 @@ impl State for LuaState {
         Ok(tls)
     }
 
+    fn ws_connect(&self, url: url::Url, options: &WebSocketOptions) -> Result<String> {
+        let mut mtx = self.ws_sessions.lock().unwrap();
+        let id = self.random_id();
+
+        let sock = WebSocket::connect(&self.dns_config, url, options)?;
+        mtx.insert(id.clone(), Arc::new(Mutex::new(sock)));
+
+        Ok(id)
+    }
+
+    fn get_ws(&self, id: &str)-> Arc<Mutex<WebSocket>> {
+        let mtx = self.ws_sessions.lock().unwrap();
+        let sock = mtx.get(id).expect("Invalid ws reference"); // TODO
+        sock.clone()
+    }
+
     fn http(&self, proxy: &Option<SocketAddr>) -> Result<Arc<chrootable_https::Client<Resolver>>> {
         let proxy = self.resolve_proxy_options(proxy)?;
 
@@ -354,7 +376,7 @@ impl State for LuaState {
 
     fn http_request(&self, session_id: &str, method: String, url: String, options: RequestOptions) -> HttpRequest {
         let mtx = self.http_sessions.lock().unwrap();
-        let session = mtx.get(session_id).expect("invalid session reference"); // TODO
+        let session = mtx.get(session_id).expect("Invalid session reference"); // TODO
 
         HttpRequest::new(&session, method, url, options)
     }
@@ -413,6 +435,7 @@ pub fn ctx<'a>(env: Environment, logger: Arc<Mutex<Box<dyn Reporter>>>) -> (hlua
         error: Mutex::new(None),
         logger,
         socket_sessions: Mutex::new(HashMap::new()),
+        ws_sessions: Mutex::new(HashMap::new()),
         blobs: Mutex::new(HashMap::new()),
         http_sessions: Mutex::new(HashMap::new()),
         http_clients: Mutex::new(HashMap::new()),
@@ -497,6 +520,7 @@ pub fn ctx<'a>(env: Environment, logger: Arc<Mutex<Box<dyn Reporter>>>) -> (hlua
     runtime::sn0int_version(&mut lua, state.clone());
     runtime::sock_connect(&mut lua, state.clone());
     runtime::sock_upgrade_tls(&mut lua, state.clone());
+    runtime::sock_options(&mut lua, state.clone());
     runtime::sock_send(&mut lua, state.clone());
     runtime::sock_recv(&mut lua, state.clone());
     runtime::sock_sendline(&mut lua, state.clone());
@@ -523,6 +547,14 @@ pub fn ctx<'a>(env: Environment, logger: Arc<Mutex<Box<dyn Reporter>>>) -> (hlua
     runtime::utf8_decode(&mut lua, state.clone());
     runtime::warn(&mut lua, state.clone());
     runtime::warn_once(&mut lua, state.clone());
+    runtime::ws_connect(&mut lua, state.clone());
+    runtime::ws_options(&mut lua, state.clone());
+    runtime::ws_recv_text(&mut lua, state.clone());
+    runtime::ws_recv_binary(&mut lua, state.clone());
+    runtime::ws_recv_json(&mut lua, state.clone());
+    runtime::ws_send_text(&mut lua, state.clone());
+    runtime::ws_send_binary(&mut lua, state.clone());
+    runtime::ws_send_json(&mut lua, state.clone());
     runtime::x509_parse_pem(&mut lua, state.clone());
     runtime::xml_decode(&mut lua, state.clone());
     runtime::xml_named(&mut lua, state.clone());

src/engine/isolation.rs

@@ -245,10 +245,7 @@ pub fn run_worker(geoip: Option<MaxmindReader>, asn: Option<MaxmindReader>, psl:
     let mut reporter = Arc::try_unwrap(mtx).expect("Failed to consume Arc")
                         .into_inner().expect("Failed to consume Mutex");
 
-    let event = match result {
-        Ok(_) => ExitEvent::Ok,
-        Err(err) => ExitEvent::Err(err.to_string()),
-    };
+    let event = result.into();
     reporter.send(&Event::Exit(event))?;
 
     Ok(())

src/lib.rs

@@ -46,6 +46,7 @@ pub mod term;
 pub mod update;
 pub mod utils;
 pub mod web;
+pub mod websockets;
 pub mod worker;
 pub mod workspaces;
 pub mod xml;

src/models/breach_email.rs

@@ -118,11 +118,16 @@ impl BreachEmail {
 pub struct PrintableBreachEmail {
     breach: String,
     email: String,
+    password: Option<String>,
 }
 
 impl fmt::Display for PrintableBreachEmail {
     fn fmt(&self, w: &mut fmt::Formatter) -> fmt::Result {
-        write!(w, "{:?} -> {:?}", self.breach, self.email)
+        write!(w, "{:?} -> {:?}", self.breach, self.email)?;
+        if let Some(password) = &self.password {
+            write!(w, " ({:?})", password)?;
+        }
+        Ok(())
     }
 }
 
@@ -133,6 +138,7 @@ impl Printable<PrintableBreachEmail> for BreachEmail {
         Ok(PrintableBreachEmail {
             breach: breach.value.to_string(),
             email: email.value.to_string(),
+            password: self.password.clone(),
         })
     }
 }
@@ -163,6 +169,7 @@ impl Printable<PrintableBreachEmail> for NewBreachEmail {
         Ok(PrintableBreachEmail {
             breach: breach.value.to_string(),
             email: email.value.to_string(),
+            password: self.password.clone(),
         })
     }
 }

src/runtime/mod.rs

@@ -32,5 +32,6 @@ import_fns!(sock);
 import_fns!(stdio);
 import_fns!(url);
 import_fns!(utf8);
+import_fns!(websockets);
 import_fns!(x509);
 import_fns!(xml);

src/runtime/sock.rs

@@ -10,7 +10,7 @@ use std::sync::Arc;
 pub fn sock_connect(lua: &mut hlua::Lua, state: Arc<dyn State>) {
     lua.set("sock_connect", hlua::function3(move |host: String, port: u16, options: AnyLuaValue| -> Result<String> {
         let options = SocketOptions::try_from(options)
-            .context("invalid socket options")
+            .context("Invalid socket options")
             .map_err(|err| state.set_error(Error::from(err)))?;
 
         state.sock_connect(&host, port, &options)
@@ -21,7 +21,7 @@ pub fn sock_connect(lua: &mut hlua::Lua, state: Arc<dyn State>) {
 pub fn sock_upgrade_tls(lua: &mut hlua::Lua, state: Arc<dyn State>) {
     lua.set("sock_upgrade_tls", hlua::function2(move |sock: String, options: AnyLuaValue| -> Result<AnyLuaValue> {
         let options = SocketOptions::try_from(options)
-            .context("invalid socket options")
+            .context("Invalid socket options")
             .map_err(|err| state.set_error(Error::from(err)))?;
 
         let tls = state.sock_upgrade_tls(&sock, &options)
@@ -31,6 +31,22 @@ pub fn sock_upgrade_tls(lua: &mut hlua::Lua, state: Arc<dyn State>) {
     }))
 }
 
+pub fn sock_options(lua: &mut hlua::Lua, state: Arc<dyn State>) {
+    lua.set("sock_options", hlua::function2(move |sock: String, options: AnyLuaValue| -> Result<()> {
+        let options = SocketOptions::try_from(options)
+            .context("Invalid socket options")
+            .map_err(|err| state.set_error(Error::from(err)))?;
+
+        let sock = state.get_sock(&sock);
+        let sock = sock.lock().unwrap();
+
+        sock.options(&options)
+            .map_err(|err| state.set_error(err))?;
+
+        Ok(())
+    }))
+}
+
 pub fn sock_send(lua: &mut hlua::Lua, state: Arc<dyn State>) {
     lua.set("sock_send", hlua::function2(move |sock: String, bytes: AnyLuaValue| -> Result<()> {
         let sock = state.get_sock(&sock);

src/runtime/websockets.rs

@@ -0,0 +1,127 @@
+use crate::errors::*;
+
+use crate::engine::ctx::State;
+use crate::engine::structs::{byte_array, lua_bytes};
+use crate::hlua::{self, AnyLuaValue};
+use crate::websockets::WebSocketOptions;
+use std::sync::Arc;
+use crate::json;
+use url::Url;
+
+pub fn ws_connect(lua: &mut hlua::Lua, state: Arc<dyn State>) {
+    lua.set("ws_connect", hlua::function2(move |url: String, options: AnyLuaValue| -> Result<String> {
+        let options = WebSocketOptions::try_from(options)
+            .context("Invalid websocket options")
+            .map_err(|err| state.set_error(Error::from(err)))?;
+
+        let url = Url::parse(&url)
+            .context("Failed to parse url")
+            .map_err(|err| state.set_error(Error::from(err)))?;
+
+        state.ws_connect(url, &options)
+            .map_err(|err| state.set_error(err))
+    }))
+}
+
+pub fn ws_options(lua: &mut hlua::Lua, state: Arc<dyn State>) {
+    lua.set("ws_options", hlua::function2(move |sock: String, options: AnyLuaValue| -> Result<()> {
+        let options = WebSocketOptions::try_from(options)
+            .context("Invalid websocket options")
+            .map_err(|err| state.set_error(Error::from(err)))?;
+
+        let sock = state.get_ws(&sock);
+        let sock = sock.lock().unwrap();
+
+        sock.options(&options)
+            .map_err(|err| state.set_error(err))?;
+
+        Ok(())
+    }))
+}
+
+pub fn ws_recv_text(lua: &mut hlua::Lua, state: Arc<dyn State>) {
+    lua.set("ws_recv_text", hlua::function1(move |sock: String| -> Result<Option<String>> {
+        let sock = state.get_ws(&sock);
+        let mut sock = sock.lock().unwrap();
+
+        let text = sock.read_text()
+            .map_err(|err| state.set_error(err))?;
+
+        Ok(text)
+    }))
+}
+
+pub fn ws_recv_binary(lua: &mut hlua::Lua, state: Arc<dyn State>) {
+    lua.set("ws_recv_binary", hlua::function1(move |sock: String| -> Result<Option<AnyLuaValue>> {
+        let sock = state.get_ws(&sock);
+        let mut sock = sock.lock().unwrap();
+
+        let bytes = sock.read_binary()
+            .map_err(|err| state.set_error(err))?
+            .map(|bytes| lua_bytes(&bytes));
+
+        Ok(bytes)
+    }))
+}
+
+pub fn ws_recv_json(lua: &mut hlua::Lua, state: Arc<dyn State>) {
+    lua.set("ws_recv_json", hlua::function1(move |sock: String| -> Result<Option<AnyLuaValue>> {
+        let sock = state.get_ws(&sock);
+        let mut sock = sock.lock().unwrap();
+
+        let json = sock.read_text()
+            .map_err(|err| state.set_error(err))?;
+
+        let json = if let Some(json) = json {
+            let json = json::decode(&json)
+                .map_err(|err| state.set_error(err))?;
+            Some(json)
+        } else {
+            None
+        };
+
+        Ok(json)
+    }))
+}
+
+pub fn ws_send_text(lua: &mut hlua::Lua, state: Arc<dyn State>) {
+    lua.set("ws_send_text", hlua::function2(move |sock: String, text: String| -> Result<()> {
+        let sock = state.get_ws(&sock);
+        let mut sock = sock.lock().unwrap();
+
+        sock.write_text(text)
+            .map_err(|err| state.set_error(err))?;
+
+        Ok(())
+    }))
+}
+
+pub fn ws_send_binary(lua: &mut hlua::Lua, state: Arc<dyn State>) {
+    lua.set("ws_send_binary", hlua::function2(move |sock: String, bytes: AnyLuaValue| -> Result<()> {
+        let sock = state.get_ws(&sock);
+        let mut sock = sock.lock().unwrap();
+
+        let bytes = byte_array(bytes)
+            .map_err(|err| state.set_error(err))?;
+
+        sock.write_binary(bytes)
+            .map_err(|err| state.set_error(err))?;
+
+        Ok(())
+    }))
+}
+
+pub fn ws_send_json(lua: &mut hlua::Lua, state: Arc<dyn State>) {
+    lua.set("ws_send_json", hlua::function2(move |sock: String, json: AnyLuaValue| -> Result<()> {
+        let sock = state.get_ws(&sock);
+        let mut sock = sock.lock().unwrap();
+
+        let json = json::encode(json)
+            .map_err(|err| state.set_error(err))?;
+
+        sock.write_text(json)
+            .map_err(|err| state.set_error(err))?;
+
+        Ok(())
+    }))
+}

src/sockets/mod.rs

@@ -16,6 +16,7 @@ use std::io::BufRead;
 use std::net::SocketAddr;
 use std::net::TcpStream;
 use std::net::{IpAddr, Ipv4Addr};
+use std::time::Duration;
 
 mod tls;
 pub use self::tls::TlsData;
@@ -39,14 +40,21 @@ fn unwrap_socket(socket: tokio::net::TcpStream) -> Result<TcpStream> {
 pub struct SocketOptions {
     #[serde(default)]
     pub tls: bool,
-    sni_value: Option<String>,
+    pub sni_value: Option<String>,
     #[serde(default)]
-    disable_tls_verify: bool,
+    pub disable_tls_verify: bool,
     pub proxy: Option<SocketAddr>,
 
     // TODO: enable_sni (default to true)
+    // TODO: sni_name
     // TODO: cacert
-    // TODO: timeout
+
+    #[serde(default)]
+    pub connect_timeout: u64,
+    #[serde(default)]
+    pub read_timeout: u64,
+    #[serde(default)]
+    pub write_timeout: u64,
 }
 
 impl SocketOptions {
@@ -57,17 +65,97 @@ impl SocketOptions {
     }
 }
 
-#[derive(Debug)]
-pub struct Socket {
-    stream: BufStream<Stream>,
-    newline: String,
+impl SocketOptions {
+    pub fn apply(&self, stream: &Stream) -> Result<()> {
+        let socket = match stream {
+            Stream::Tcp(s) => s,
+            Stream::Tls(s) => s.get_ref(),
+        };
+        self.apply_tcp(socket)
+    }
+
+    pub fn apply_tcp(&self, socket: &TcpStream) -> Result<()> {
+        let read_timeout = self.read_timeout;
+        if read_timeout > 0 {
+            socket.set_read_timeout(Some(Duration::from_secs(read_timeout)))?;
+        }
+
+        let write_timeout = self.write_timeout;
+        if write_timeout > 0 {
+            socket.set_write_timeout(Some(Duration::from_secs(write_timeout)))?;
+        }
+
+        Ok(())
+    }
 }
 
-enum Stream {
+pub enum Stream {
     Tcp(TcpStream),
     Tls(rustls::StreamOwned<rustls::ClientSession, TcpStream>),
 }
 
+impl Stream {
+    pub fn connect_stream<R: DnsResolver>(resolver: &R, host: &str, port: u16, options: &SocketOptions) -> Result<Stream> {
+        let addrs = match host.parse::<IpAddr>() {
+            Ok(addr) => vec![addr],
+            Err(_) => resolver.resolve(host, RecordType::A)
+                .wait_for_response()?
+                .success()?,
+        };
+
+        let mut errors = Vec::new();
+
+        for addr in addrs {
+            match Stream::connect_addr(host, (addr, port).into(), &options) {
+                Ok(socket) => {
+                    return Ok(socket);
+                },
+                Err(err) => errors.push((addr, err)),
+            }
+        }
+
+        if errors.is_empty() {
+            bail!("no dns records found");
+        } else {
+            bail!("couldn't connect: {:?}", errors);
+        }
+    }
+
+    fn connect_addr(host: &str, addr: SocketAddr, options: &SocketOptions) -> Result<Stream> {
+        debug!("connecting to {}", addr);
+
+        let connect_timeout = options.connect_timeout;
+        let socket = if connect_timeout > 0 {
+            TcpStream::connect_timeout(&addr, Duration::from_secs(connect_timeout))?
+        } else {
+            TcpStream::connect(&addr)?
+        };
+        debug!("successfully connected to {:?}", addr);
+
+        options.apply_tcp(&socket)?;
+
+        tls::wrap_if_enabled(socket, host, options)
+    }
+
+    pub fn connect_socks5_stream(proxy: &SocketAddr, host: &str, port: u16, options: &SocketOptions) -> Result<Stream> {
+        debug!("connecting to {:?}:{:?} with socks5 on {:?}", host, port, proxy);
+
+        let addr = match host.parse::<Ipv4Addr>() {
+            Ok(ipaddr) => ProxyDest::Ipv4Addr(ipaddr),
+            _ => ProxyDest::Domain(host.to_string()),
+        };
+
+        let fut = socks5::connect(proxy, addr, port);
+
+        let mut rt = Runtime::new()?;
+        let socket = rt.block_on(fut)?;
+
+        let socket = unwrap_socket(socket)?;
+
+        tls::wrap_if_enabled(socket, host, options)
+    }
+}
+
 impl fmt::Debug for Stream {
     fn fmt(&self, w: &mut fmt::Formatter) -> fmt::Result {
         match self {
@@ -102,6 +190,12 @@ impl Write for Stream {
     }
 }
 
+#[derive(Debug)]
+pub struct Socket {
+    stream: BufStream<Stream>,
+    newline: String,
+}
+
 impl Socket {
     fn new(stream: Stream) -> Socket {
         let stream = BufStream::new(stream);
@@ -112,60 +206,31 @@ impl Socket {
     }
 
     pub fn connect<R: DnsResolver>(resolver: &R, host: &str, port: u16, options: &SocketOptions) -> Result<Socket> {
-        let addrs = match host.parse::<IpAddr>() {
-            Ok(addr) => vec![addr],
-            Err(_) => resolver.resolve(host, RecordType::A)
-                .wait_for_response()?
-                .success()?,
-        };
-
-        let mut errors = Vec::new();
-
-        for addr in addrs {
-            debug!("connecting to {}:{}", addr, port);
-            match TcpStream::connect((addr, port)) {
-                Ok(socket) => {
-                    debug!("successfully connected to {:?}", addr);
-                    return tls::wrap_if_enabled(socket, host, options);
-                },
-                Err(err) => errors.push((addr, err)),
-            }
-        }
-
-        if errors.is_empty() {
-            bail!("no dns records found");
-        } else {
-            bail!("couldn't connect: {:?}", errors);
-        }
+        let stream = Stream::connect_stream(resolver, host, port, options)?;
+        Ok(Socket::new(stream))
     }
 
     pub fn connect_socks5(proxy: &SocketAddr, host: &str, port: u16, options: &SocketOptions) -> Result<Socket> {
-        debug!("connecting to {:?}:{:?} with socks5 on {:?}", host, port, proxy);
-
-        let addr = match host.parse::<Ipv4Addr>() {
-            Ok(ipaddr) => ProxyDest::Ipv4Addr(ipaddr),
-            _ => ProxyDest::Domain(host.to_string()),
-        };
-
-        let fut = socks5::connect(proxy, addr, port);
-
-        let mut rt = Runtime::new()?;
-        let socket = rt.block_on(fut)?;
-
-        let socket = unwrap_socket(socket)?;
-
-        tls::wrap_if_enabled(socket, host, options)
+        let stream = Stream::connect_socks5_stream(proxy, host, port, options)?;
+        Ok(Socket::new(stream))
     }
 
     pub fn upgrade_to_tls(self, options: &SocketOptions) -> Result<(Socket, TlsData)> {
         let stream = self.stream.into_inner()?;
 
-        match stream {
-            Stream::Tcp(stream) => tls::wrap(stream, "", options),
-            _ => bail!("Only tcp streams can be upgraded"),
+        if let Stream::Tcp(stream) = stream {
+            let (stream, tls) = tls::wrap(stream, "", options)?;
+            let socket = Socket::new(stream);
+            Ok((socket, tls))
+        } else {
+            bail!("Only tcp streams can be upgraded")
         }
     }
 
+    pub fn options(&self, options: &SocketOptions) -> Result<()> {
+        options.apply(self.stream.get_ref())
+    }
+
     pub fn send(&mut self, data: &[u8]) -> Result<()> {
         match str::from_utf8(&data) {
             Ok(data) => debug!("send: {:?}", data),
@@ -178,7 +243,12 @@ impl Socket {
 
     pub fn recv(&mut self) -> Result<Vec<u8>> {
         let mut buf = [0; 4096];
-        let n = self.stream.read(&mut buf)?;
+        let n = match self.stream.read(&mut buf) {
+            Ok(n) if n == 0 => bail!("Connection closed"),
+            Ok(n) => n,
+            Err(err) if err.kind() == io::ErrorKind::WouldBlock => 0,
+            Err(err) => return Err(err.into()),
+        };
         let data = buf[..n].to_vec();
         match str::from_utf8(&data) {
             Ok(data) => debug!("recv: {:?}", data),

src/sockets/tls.rs

@@ -2,14 +2,14 @@ use crate::errors::*;
 
 use crate::hlua::AnyLuaValue;
 use crate::json::LuaJsonValue;
-use rustls::{self, ClientConfig, Session, ClientSession, RootCertStore};
+use rustls::{self, ClientConfig, Session, ClientSession};
 
 use std::str;
 use std::result;
 use std::sync::Arc;
 use std::net::TcpStream;
 
-use super::{Socket, Stream, SocketOptions};
+use super::{Stream, SocketOptions};
 
 
 #[derive(Debug, Serialize)]
@@ -26,22 +26,22 @@ impl TlsData {
     }
 }
 
-pub fn wrap_if_enabled(stream: TcpStream, host: &str, options: &SocketOptions) -> Result<Socket> {
+pub fn wrap_if_enabled(stream: TcpStream, host: &str, options: &SocketOptions) -> Result<Stream> {
     if !options.tls {
         let stream = Stream::Tcp(stream);
-        return Ok(Socket::new(stream));
+        Ok(stream)
     } else {
         let (socket, _) = wrap(stream, host, options)?;
         Ok(socket)
     }
 }
 
-pub fn wrap(stream: TcpStream, host: &str, options: &SocketOptions) -> Result<(Socket, TlsData)> {
-    let mut anchors = RootCertStore::empty();
-    anchors.add_server_trust_anchors(&webpki_roots::TLS_SERVER_ROOTS);
-
+pub fn wrap(stream: TcpStream, host: &str, options: &SocketOptions) -> Result<(Stream, TlsData)> {
     let mut config = ClientConfig::new();
-    config.root_store = anchors;
+    config
+        .root_store
+        .add_server_trust_anchors(&webpki_roots::TLS_SERVER_ROOTS);
+    config.ct_logs = Some(&ct_logs::LOGS);
 
     if options.disable_tls_verify {
         info!("tls verification has been disabled");
@@ -73,14 +73,16 @@ fn get_dns_name(config: &mut ClientConfig, host: &str) -> webpki::DNSName {
     }
 }
 
-fn setup(mut stream: TcpStream, mut session: ClientSession) -> Result<(Socket, TlsData)> {
+fn setup(mut stream: TcpStream, mut session: ClientSession) -> Result<(Stream, TlsData)> {
     info!("starting tls handshake");
     if session.is_handshaking() {
-        session.complete_io(&mut stream)?;
+        session.complete_io(&mut stream)
+            .context("Failed to read reply to tls client hello")?;
     }
 
     if session.wants_write() {
-        session.complete_io(&mut stream)?;
+        session.complete_io(&mut stream)
+            .context("wants_write->complete_io failed")?;
     }
 
     let mut tls = TlsData {
@@ -106,7 +108,7 @@ fn setup(mut stream: TcpStream, mut session: ClientSession) -> Result<(Socket, T
     info!("successfully established tls connection");
     let stream = rustls::StreamOwned::new(session, stream);
     let stream = Stream::Tls(stream);
-    Ok((Socket::new(stream), tls))
+    Ok((stream, tls))
 }
 
 pub struct NoCertificateVerification {}

src/web.rs

@@ -239,6 +239,10 @@ impl HttpRequest {
         }
         resp.insert("headers", headers);
 
+        if let Some(ipaddr) = res.ipaddr {
+            resp.insert_str("ipaddr", ipaddr.to_string());
+        }
+
         if self.into_blob {
             let blob = Blob::create(res.body);
             let id = state.register_blob(blob);

src/websockets.rs

@@ -0,0 +1,153 @@
+use chrootable_https::DnsResolver;
+use crate::errors::*;
+use crate::hlua::AnyLuaValue;
+use crate::json::LuaJsonValue;
+use crate::sockets::{Stream, SocketOptions};
+use std::borrow::Cow;
+use std::collections::HashMap;
+use std::net::SocketAddr;
+use std::io;
+use tungstenite::handshake::client::Request;
+use tungstenite::protocol::{self, Message};
+use url::Url;
+
+#[derive(Debug, Default, Deserialize)]
+pub struct WebSocketOptions {
+    pub headers: Option<HashMap<String, String>>,
+    pub proxy: Option<SocketAddr>,
+
+    #[serde(default)]
+    pub connect_timeout: u64,
+    #[serde(default)]
+    pub read_timeout: u64,
+    #[serde(default)]
+    pub write_timeout: u64,
+}
+
+impl WebSocketOptions {
+    pub fn try_from(x: AnyLuaValue) -> Result<WebSocketOptions> {
+        let x = LuaJsonValue::from(x);
+        let x = serde_json::from_value(x.into())?;
+        Ok(x)
+    }
+}
+
+pub enum Event {
+    Text(String),
+    Binary(Vec<u8>),
+    Close,
+    Timeout,
+}
+
+pub struct WebSocket {
+    sock: protocol::WebSocket<Stream>,
+}
+
+impl WebSocket {
+    pub fn negotiate(stream: Stream, url: Url, headers: Option<&HashMap<String, String>>) -> Result<WebSocket> {
+        let extra_headers = headers.map(|headers| {
+            headers.iter()
+                .map(|(k, v)| (Cow::Borrowed(k.as_str()), Cow::Borrowed(v.as_str())))
+                .collect()
+        });
+
+        let (sock, _resp) = tungstenite::client::client(
+            Request {
+                url,
+                extra_headers,
+            },
+            stream,
+        )?;
+        Ok(WebSocket {
+            sock,
+        })
+    }
+
+    pub fn connect<R: DnsResolver>(resolver: &R, url: Url, options: &WebSocketOptions) -> Result<WebSocket> {
+        let tls = match url.scheme() {
+            "ws" => false,
+            "wss" => true,
+            _ => bail!("Invalid websocket protocol"),
+        };
+
+        let host = url.host_str()
+            .ok_or_else(|| format_err!("Missing host in url"))?;
+
+        let port = match (url.port(), tls) {
+            (Some(port), _) => port,
+            (None, true) => 443,
+            (None, false) => 80,
+        };
+
+        let stream = Stream::connect_stream(resolver, host, port, &SocketOptions {
+            tls,
+            sni_value: None,
+            disable_tls_verify: false,
+            proxy: options.proxy,
+
+            connect_timeout: options.connect_timeout,
+            read_timeout: options.read_timeout,
+            write_timeout: options.write_timeout,
+        })?;
+        Self::negotiate(stream, url, options.headers.as_ref())
+    }
+
+    pub fn options(&self, options: &WebSocketOptions) -> Result<()> {
+        let mut o = SocketOptions::default();
+        o.read_timeout = options.read_timeout;
+        o.write_timeout = options.write_timeout;
+        o.apply(self.sock.get_ref())
+    }
+
+    fn read_msg(&mut self) -> Result<Event> {
+        loop {
+            let msg = match self.sock.read_message() {
+                Ok(Message::Text(body)) => Event::Text(body),
+                Ok(Message::Binary(body)) => Event::Binary(body),
+                Ok(Message::Ping(ping)) => {
+                    self.sock.write_message(Message::Pong(ping))?;
+                    continue;
+                },
+                Ok(Message::Pong(_)) => continue, // this should never happen
+                Ok(Message::Close(_)) => Event::Close,
+                Err(tungstenite::Error::ConnectionClosed) => Event::Close,
+                Err(tungstenite::Error::AlreadyClosed) => Event::Close,
+                Err(tungstenite::Error::Io(err)) if err.kind() == io::ErrorKind::WouldBlock => Event::Timeout,
+                Err(err) => return Err(err.into()),
+            };
+            return Ok(msg);
+        }
+    }
+
+    pub fn read_text(&mut self) -> Result<Option<String>> {
+        match self.read_msg()? {
+            Event::Text(text) => Ok(Some(text)),
+            Event::Binary(_) => bail!("Unexpected message type: binary"),
+            Event::Close => bail!("Connection closed"),
+            Event::Timeout => Ok(None),
+        }
+    }
+
+    pub fn read_binary(&mut self) -> Result<Option<Vec<u8>>> {
+        match self.read_msg()? {
+            Event::Text(_) => bail!("Unexpected message type: text"),
+            Event::Binary(body) => Ok(Some(body)),
+            Event::Close => bail!("Connection closed"),
+            Event::Timeout => Ok(None),
+        }
+    }
+
+    fn write_msg(&mut self, msg: Message) -> Result<()> {
+        self.sock.write_message(msg)?;
+        self.sock.write_pending()?;
+        Ok(())
+    }
+
+    pub fn write_text(&mut self, text: String) -> Result<()> {
+        self.write_msg(Message::Text(text))
+    }
+
+    pub fn write_binary(&mut self, binary: Vec<u8>) -> Result<()> {
+        self.write_msg(Message::Binary(binary))
+    }
+}

src/worker.rs

@@ -100,6 +100,21 @@ pub enum ExitEvent {
     SetupFailed(String),
 }
 
+impl From<Result<()>> for ExitEvent {
+    fn from(result: Result<()>) -> ExitEvent {
+        match result {
+            Ok(_) => ExitEvent::Ok,
+            Err(err) => {
+                let err = err.iter_chain()
+                    .map(|e| e.to_string())
+                    .collect::<Vec<_>>()
+                    .join(": ");
+                ExitEvent::Err(err.to_string())
+            },
+        }
+    }
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 pub enum LogEvent {
     Info(String),
@@ -531,18 +546,8 @@ pub fn spawn_multi<T: Task, F>(tasks: Vec<T>, mut done_fn: F, threads: usize) ->
 
             tx.send(Event2::Start);
 
-            let exit = match task.run(&tx) {
-                Ok(_) => ExitEvent::Ok,
-                Err(err) => {
-                    let err = err.iter_chain()
-                        .map(|e| e.to_string())
-                        .collect::<Vec<_>>()
-                        .join(": ");
-                    ExitEvent::Err(err.to_string())
-                },
-            };
-
-            tx.send(Event2::Exit(exit));
+            let exit = task.run(&tx);
+            tx.send(Event2::Exit(exit.into()));
         });
         expected += 1;
     }