charcuterie/src/main.rs

#[macro_use]
extern crate serde_json;

use async_std::task;
use async_std::net::{SocketAddr, TcpListener, TcpStream};
use async_tungstenite::tungstenite::protocol::Message;
use futures::pin_mut;
use futures::prelude::*;
use futures::channel::mpsc::{unbounded, UnboundedSender};
use glob::glob;
use handlebars::Handlebars;
use log::*;
use rodio::Source;
use serde::Deserialize;
use tide::{Body, Request, StatusCode};

use std::collections::HashMap;
use std::fs::File;
use std::io::BufReader;
use std::sync::{Arc, Mutex};


type Tx = UnboundedSender<Message>;
type PeerMap = Arc<Mutex<HashMap<SocketAddr, Tx>>>;


/**
 * Simple struct to deserialize some query parameters
 */
#[derive(Debug, Deserialize)]
struct Query {
    admin: String,
}

async fn index(req: Request<()>) -> Result<Body, tide::Error> {
    let mut hb = Handlebars::new();

    if let Err(failure) = hb.register_templates_directory(".hbs", "views") {
        error!("Failed to render: {:?}", failure);
        return Err(tide::Error::from_str(StatusCode::InternalServerError,  "Could not render templates"))
    }

    let mut admin = false;
    if let Ok(query) = req.query::<Query>() {
        info!("admin: {:?}", query);
        admin = query.admin == "rtyler";
    }

    let mut sounds = vec![];

    for sound in glob("sounds/*.wav").expect("Failed to glob sounds") {
        if let Ok(sound) = sound {
            if let Some(name) = sound.as_path().file_stem() {
                // factory whistle is a special admin only sound ^_^
                if name != "factorywhistle" {
                    sounds.push(name.to_os_string().into_string().unwrap());
                }
            }
        }
    }
    info!("sounds: {:?}", sounds);

    let view = hb.render("index",
        &json!({"sounds": sounds, "admin" : admin}))
        .expect("Failed to render");
    let mut body = Body::from_string(view);
    body.set_mime("text/html");
    Ok(body)
}

async fn play(req: Request<()>) -> tide::Result {
    if let Ok(sound) = req.param::<String>("sound") {
        let device = rodio::default_output_device().unwrap();
        let file = File::open(format!("sounds/{}.wav", sound)).unwrap();
        let source = rodio::Decoder::new(BufReader::new(file)).unwrap();
        rodio::play_raw(&device, source.convert_samples());
        Ok(tide::Redirect::new("/").into())
    }
    else {
        Err(tide::Error::from_str(StatusCode::NotFound,  "Could not find sound"))
    }
}

async fn handle_websocket(peer_map: PeerMap, raw_stream: TcpStream, addr: SocketAddr) {
    println!("Incoming TCP connection from: {}", addr);

    let ws_stream = async_tungstenite::accept_async(raw_stream)
        .await
        .expect("Error during the websocket handshake occurred");
    println!("WebSocket connection established: {}", addr);

    // Insert the write part of this peer to the peer map.
    let (tx, rx) = unbounded();
    peer_map.lock().unwrap().insert(addr, tx);

    let (outgoing, incoming) = ws_stream.split();

    let broadcast_incoming = incoming
        .try_filter(|msg| {
            // Broadcasting a Close message from one client
            // will close the other clients.
            future::ready(!msg.is_close())
        })
        .try_for_each(|msg| {
            info!(
                "Received a message from {}: {}",
                addr,
                msg.to_text().unwrap()
            );
            let peers = peer_map.lock().unwrap();

            if let Ok(json) = serde_json::from_str::<serde_json::Value>(msg.to_text().unwrap()) {
                info!("JSON: {:?}", json);
            }

            // We want to broadcast the message to everyone except ourselves.
            let broadcast_recipients = peers
                .iter()
                .filter(|(peer_addr, _)| peer_addr != &&addr)
                .map(|(_, ws_sink)| ws_sink);

            for recp in broadcast_recipients {
                recp.unbounded_send(msg.clone()).unwrap();
            }

            future::ok(())
        });

    let receive_from_others = rx.map(Ok).forward(outgoing);

    pin_mut!(broadcast_incoming, receive_from_others);
    future::select(broadcast_incoming, receive_from_others).await;

    println!("{} disconnected", &addr);
    peer_map.lock().unwrap().remove(&addr);
}

#[async_std::main]
async fn main() -> Result<(), tide::Error> {
    pretty_env_logger::init();

    let mut app = tide::new();
    app.at("/play/:sound").post(play);
    app.at("/").get(index);
    app.at("/assets").serve_dir("assets/");

    task::spawn(async move {
        let addr = "0.0.0.0:9078";
        let state = PeerMap::new(Mutex::new(HashMap::new()));
        // Create the event loop and TCP listener we'll accept connections on.
        let try_socket = TcpListener::bind(&addr).await;
        let listener = try_socket.expect("Failed to bind");
        info!("Listening on: {}", addr);

        while let Ok((stream, addr)) = listener.accept().await {
            task::spawn(handle_websocket(state.clone(), stream, addr));
        }
    });

    Ok(app.listen("0.0.0.0:9077").await?)
}