innernet/server/src/api/user.rs

use crate::{
    api::inject_endpoints,
    db::{DatabaseCidr, DatabasePeer},
    form_body, with_session, with_unredeemed_session, Context, ServerError, Session,
    UnredeemedSession,
};
use hyper::StatusCode;
use shared::{EndpointContents, PeerContents, RedeemContents, State, REDEEM_TRANSITION_WAIT};
use warp::Filter;
use wgctrl::DeviceConfigBuilder;

pub fn routes(
    context: Context,
) -> impl Filter<Extract = (impl warp::Reply,), Error = warp::Rejection> + Clone {
    warp::path("user").and(
        routes::state(context.clone())
            .or(routes::redeem(context.clone()))
            .or(routes::override_endpoint(context.clone())),
    )
}

pub mod routes {
    use super::*;

    pub fn state(
        context: Context,
    ) -> impl Filter<Extract = (impl warp::Reply,), Error = warp::Rejection> + Clone {
        warp::path("state")
            .and(warp::path::end())
            .and(warp::get())
            .and(with_session(context))
            .and_then(handlers::state)
    }

    pub fn redeem(
        context: Context,
    ) -> impl Filter<Extract = (impl warp::Reply,), Error = warp::Rejection> + Clone {
        warp::path("redeem")
            .and(warp::path::end())
            .and(warp::post())
            .and(form_body())
            .and(with_unredeemed_session(context))
            .and_then(handlers::redeem)
    }

    pub fn override_endpoint(
        context: Context,
    ) -> impl Filter<Extract = (impl warp::Reply,), Error = warp::Rejection> + Clone {
        warp::path("endpoint")
            .and(warp::path::end())
            .and(warp::put())
            .and(form_body())
            .and(with_session(context))
            .and_then(handlers::endpoint)
    }
}

mod handlers {
    use super::*;

    /// Get the current state of the network, in the eyes of the current peer.
    ///
    /// This endpoint returns the visible CIDRs and Peers, providing all the necessary
    /// information for the peer to create connections to all of them.
    pub async fn state(session: Session) -> Result<impl warp::Reply, warp::Rejection> {
        let conn = session.context.db.lock();
        let selected_peer = DatabasePeer::get(&conn, session.peer.id)?;

        let cidrs: Vec<_> = DatabaseCidr::list(&conn)?;

        let mut peers: Vec<_> = selected_peer
            .get_all_allowed_peers(&conn)?
            .into_iter()
            .map(|p| p.inner)
            .collect();
        inject_endpoints(&session, &mut peers);

        Ok(warp::reply::json(&State { cidrs, peers }))
    }

    /// Redeems an invitation. An invitation includes a WireGuard keypair generated by either the server
    /// or a peer with admin rights.
    ///
    /// Redemption is the process of an invitee generating their own keypair and exchanging their temporary
    /// key with their permanent one.
    ///
    /// Until this API endpoint is called, the invited peer will not show up to other peers, and once
    /// it is called and succeeds, it cannot be called again.
    pub async fn redeem(
        form: RedeemContents,
        session: UnredeemedSession,
    ) -> Result<impl warp::Reply, warp::Rejection> {
        let conn = session.context.db.lock();
        let mut selected_peer = DatabasePeer::get(&conn, session.peer.id)?;

        let old_public_key = wgctrl::Key::from_base64(&selected_peer.public_key)
            .map_err(|_| ServerError::WireGuard)?;

        if selected_peer.is_redeemed {
            Ok(StatusCode::GONE)
        } else {
            selected_peer.redeem(&conn, &form.public_key)?;

            if cfg!(not(test)) {
                let interface = session.context.interface.clone();

                // If we were to modify the WireGuard interface immediately, the HTTP response wouldn't
                // get through. Instead, we need to wait a reasonable amount for the HTTP response to
                // flush, then update the interface.
                //
                // The client is also expected to wait the same amount of time after receiving a success
                // response from /redeem.
                //
                // This might be avoidable if we were able to run code after we were certain the response
                // had flushed over the TCP socket, but that isn't easily accessible from this high-level
                // web framework.
                tokio::task::spawn(async move {
                    tokio::time::sleep(*REDEEM_TRANSITION_WAIT).await;
                    log::info!(
                        "WireGuard: adding new peer {}, removing old pubkey {}",
                        &*selected_peer,
                        old_public_key.to_base64()
                    );
                    DeviceConfigBuilder::new()
                        .remove_peer_by_key(&old_public_key)
                        .add_peer((&*selected_peer).into())
                        .apply(&interface)
                        .map_err(|e| log::error!("{:?}", e))
                        .ok();
                });
            }
            Ok(StatusCode::NO_CONTENT)
        }
    }

    /// Redeems an invitation. An invitation includes a WireGuard keypair generated by either the server
    /// or a peer with admin rights.
    ///
    /// Redemption is the process of an invitee generating their own keypair and exchanging their temporary
    /// key with their permanent one.
    ///
    /// Until this API endpoint is called, the invited peer will not show up to other peers, and once
    /// it is called and succeeds, it cannot be called again.
    pub async fn endpoint(
        contents: EndpointContents,
        session: Session,
    ) -> Result<impl warp::Reply, warp::Rejection> {
        let conn = session.context.db.lock();
        let mut selected_peer = DatabasePeer::get(&conn, session.peer.id)?;
        selected_peer.update(
            &conn,
            PeerContents {
                endpoint: contents.into(),
                ..selected_peer.contents.clone()
            },
        )?;

        Ok(StatusCode::NO_CONTENT)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{db::DatabaseAssociation, test};
    use anyhow::Result;
    use shared::{AssociationContents, CidrContents, EndpointContents};
    use warp::http::StatusCode;

    #[tokio::test]
    async fn test_get_state_from_developer1() -> Result<()> {
        let server = test::Server::new()?;
        let filter = crate::routes(server.context());
        let res = server
            .request_from_ip(test::DEVELOPER1_PEER_IP)
            .path("/v1/user/state")
            .reply(&filter)
            .await;

        assert_eq!(res.status(), StatusCode::OK);

        let State { peers, .. } = serde_json::from_slice(&res.body())?;
        let mut peer_names = peers.iter().map(|p| &p.contents.name).collect::<Vec<_>>();
        peer_names.sort();
        // Developers should see only peers in infra CIDR and developer CIDR.
        assert_eq!(
            &["developer1", "developer2", "innernet-server"],
            &peer_names[..]
        );

        Ok(())
    }

    #[tokio::test]
    async fn test_override_endpoint() -> Result<()> {
        let server = test::Server::new()?;
        let filter = crate::routes(server.context());
        assert_eq!(
            server
                .put_request_from_ip(test::DEVELOPER1_PEER_IP)
                .path("/v1/user/endpoint")
                .body(serde_json::to_string(&EndpointContents::Set(
                    "1.1.1.1:51820".parse()?
                ))?)
                .reply(&filter)
                .await
                .status(),
            StatusCode::NO_CONTENT
        );

        println!("{}", serde_json::to_string(&EndpointContents::Unset)?);
        assert_eq!(
            server
                .put_request_from_ip(test::DEVELOPER1_PEER_IP)
                .path("/v1/user/endpoint")
                .body(serde_json::to_string(&EndpointContents::Unset)?)
                .reply(&filter)
                .await
                .status(),
            StatusCode::NO_CONTENT
        );

        assert_eq!(
            server
                .put_request_from_ip(test::DEVELOPER1_PEER_IP)
                .path("/v1/user/endpoint")
                .body("endpoint=blah")
                .reply(&filter)
                .await
                .status(),
            StatusCode::BAD_REQUEST
        );

        Ok(())
    }

    #[tokio::test]
    async fn test_list_peers_from_unknown_ip() -> Result<()> {
        let server = test::Server::new()?;
        let filter = crate::routes(server.context());

        // Request comes from an unknown IP.
        let res = server
            .request_from_ip("10.80.80.80")
            .path("/v1/user/state")
            .reply(&filter)
            .await;

        assert_eq!(res.status(), StatusCode::UNAUTHORIZED);

        Ok(())
    }

    #[tokio::test]
    async fn test_list_peers_for_developer_subcidr() -> Result<()> {
        let server = test::Server::new()?;
        let filter = crate::routes(server.context());
        {
            let db = server.db.lock();
            let cidr = DatabaseCidr::create(
                &db,
                CidrContents {
                    name: "experiment cidr".to_string(),
                    cidr: test::EXPERIMENTAL_CIDR.parse()?,
                    parent: Some(test::ROOT_CIDR_ID),
                },
            )?;
            let subcidr = DatabaseCidr::create(
                &db,
                CidrContents {
                    name: "experiment subcidr".to_string(),
                    cidr: test::EXPERIMENTAL_SUBCIDR.parse()?,
                    parent: Some(cidr.id),
                },
            )?;
            DatabasePeer::create(
                &db,
                test::peer_contents(
                    "experiment-peer",
                    test::EXPERIMENT_SUBCIDR_PEER_IP,
                    subcidr.id,
                    false,
                )?,
            )?;

            // Add a peering between the developer's CIDR and the experimental *parent* cidr.
            DatabaseAssociation::create(
                &db,
                AssociationContents {
                    cidr_id_1: test::DEVELOPER_CIDR_ID,
                    cidr_id_2: cidr.id,
                },
            )?;
            DatabaseAssociation::create(
                &db,
                AssociationContents {
                    cidr_id_1: test::INFRA_CIDR_ID,
                    cidr_id_2: cidr.id,
                },
            )?;
        }

        for ip in &[test::DEVELOPER1_PEER_IP, test::EXPERIMENT_SUBCIDR_PEER_IP] {
            let res = server
                .request_from_ip(ip)
                .path("/v1/user/state")
                .reply(&filter)
                .await;
            assert_eq!(res.status(), StatusCode::OK);
            let State { peers, .. } = serde_json::from_slice(&res.body())?;
            let mut peer_names = peers.iter().map(|p| &p.contents.name).collect::<Vec<_>>();
            peer_names.sort();
            // Developers should see only peers in infra CIDR and developer CIDR.
            assert_eq!(
                &[
                    "developer1",
                    "developer2",
                    "experiment-peer",
                    "innernet-server"
                ],
                &peer_names[..]
            );
        }

        Ok(())
    }

    #[tokio::test]
    async fn test_redeem() -> Result<()> {
        let server = test::Server::new()?;

        let experimental_cidr = DatabaseCidr::create(
            &server.db().lock(),
            CidrContents {
                name: "experimental".to_string(),
                cidr: test::EXPERIMENTAL_CIDR.parse()?,
                parent: Some(test::ROOT_CIDR_ID),
            },
        )?;

        let mut peer_contents = test::peer_contents(
            "experiment-peer",
            test::EXPERIMENT_SUBCIDR_PEER_IP,
            experimental_cidr.id,
            false,
        )?;
        peer_contents.is_redeemed = false;
        let _experiment_peer = DatabasePeer::create(&server.db().lock(), peer_contents)?;

        let filter = crate::routes(server.context());

        // Step 1: Ensure that before redeeming, other endpoints aren't yet accessible.
        let res = server
            .request_from_ip(test::EXPERIMENT_SUBCIDR_PEER_IP)
            .path("/v1/user/state")
            .reply(&filter)
            .await;
        assert_eq!(res.status(), StatusCode::UNAUTHORIZED);

        // Step 2: Ensure that redemption works.
        let body = RedeemContents {
            public_key: "YBVIgpfLbi/knrMCTEb0L6eVy0daiZnJJQkxBK9s+2I=".into(),
        };
        let res = server
            .post_request_from_ip(test::EXPERIMENT_SUBCIDR_PEER_IP)
            .path("/v1/user/redeem")
            .body(serde_json::to_string(&body)?)
            .reply(&filter)
            .await;
        assert!(res.status().is_success());

        // Step 3: Ensure that a second attempt at redemption DOESN'T work.
        let res = server
            .post_request_from_ip(test::EXPERIMENT_SUBCIDR_PEER_IP)
            .path("/v1/user/redeem")
            .body(serde_json::to_string(&body)?)
            .reply(&filter)
            .await;
        assert!(res.status().is_client_error());

        // Step 3: Ensure that after redemption, fetching state works.
        let res = server
            .request_from_ip(test::EXPERIMENT_SUBCIDR_PEER_IP)
            .path("/v1/user/state")
            .reply(&filter)
            .await;
        assert_eq!(res.status(), StatusCode::OK);
        Ok(())
    }
}