Kordophone/kordophone/src/api/http_client.rs

extern crate hyper;
extern crate serde;

use std::{path::PathBuf, str};

use crate::api::AuthenticationStore;
use crate::api::event_socket::EventSocket;
use hyper::{Body, Client, Method, Request, Uri};

use async_trait::async_trait;
use serde::{de::DeserializeOwned, Deserialize, Serialize};

use tokio::net::TcpStream;

use futures_util::{StreamExt, TryStreamExt};
use futures_util::stream::{SplitStream, SplitSink, Stream};
use futures_util::stream::BoxStream;

use tokio_tungstenite::connect_async;
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};

use crate::{
    model::{
        Conversation, 
        ConversationID, 
        JwtToken, 
        Message, 
        MessageID, 
        UpdateItem, 
        Event, 
        OutgoingMessage,
    }, 

    APIInterface
};

type HttpClient = Client<hyper::client::HttpConnector>;

pub struct HTTPAPIClient<K: AuthenticationStore + Send + Sync> {
    pub base_url: Uri,
    pub auth_store: K,
    client: HttpClient,
}

#[derive(Clone, Serialize, Deserialize, Debug)]
pub struct Credentials {
    pub username: String,
    pub password: String,
}

#[derive(Debug)]
pub enum Error {
    ClientError(String),
    HTTPError(hyper::Error),
    SerdeError(serde_json::Error),
    DecodeError,
}

impl std::error::Error for Error {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Error::HTTPError(ref err) => Some(err),
            _ => None,
        }
    }
}

impl std::fmt::Display for Error {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{:?}", self)
    }
}

impl From <hyper::Error> for Error {
    fn from(err: hyper::Error) -> Error {
        Error::HTTPError(err)
    }
}

impl From <serde_json::Error> for Error {
    fn from(err: serde_json::Error) -> Error {
        Error::SerdeError(err)
    }
}

impl From <tungstenite::Error> for Error {
    fn from(err: tungstenite::Error) -> Error {
        Error::ClientError(err.to_string())
    }
}

trait AuthBuilder {
    fn with_auth(self, token: &Option<JwtToken>) -> Self;
}

impl AuthBuilder for hyper::http::request::Builder {
    fn with_auth(self, token: &Option<JwtToken>) -> Self {
        if let Some(token) = &token {
            self.header("Authorization", token.to_header_value())
        } else { self }
    }
}

#[cfg(test)]
#[allow(dead_code)]
trait AuthSetting {
    fn authenticate(&mut self, token: &Option<JwtToken>);
}

#[cfg(test)]
impl<B> AuthSetting for hyper::http::Request<B> {
    fn authenticate(&mut self, token: &Option<JwtToken>) {
        if let Some(token) = &token {
            self.headers_mut().insert("Authorization", token.to_header_value());
        }
    }
}

type WebsocketSink = SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, tungstenite::Message>;
type WebsocketStream = SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>;

pub struct WebsocketEventSocket {
    _sink: WebsocketSink,
    stream: WebsocketStream,
}

impl WebsocketEventSocket {
    pub fn new(socket: WebSocketStream<MaybeTlsStream<TcpStream>>) -> Self {
        let (sink, stream) = socket.split();
        Self { _sink: sink, stream }
    }
}

impl WebsocketEventSocket {
    fn raw_update_stream(self) -> impl Stream<Item = Result<Vec<UpdateItem>, Error>> {
        self.stream
            .map_err(Error::from)
            .try_filter_map(|msg| async move {
                match msg {
                    tungstenite::Message::Text(text) => {
                        serde_json::from_str::<Vec<UpdateItem>>(&text)
                            .map(Some)
                            .map_err(Error::from)
                    }
                    _ => Ok(None)
                }
            })
    }
}

#[async_trait]
impl EventSocket for WebsocketEventSocket {
    type Error = Error;
    type EventStream = BoxStream<'static, Result<Event, Error>>;
    type UpdateStream = BoxStream<'static, Result<Vec<UpdateItem>, Error>>;

    async fn events(self) -> Self::EventStream {
        use futures_util::stream::iter;

        self.raw_update_stream()
            .map_ok(|updates| iter(updates.into_iter().map(|update| Ok(Event::from(update)))))
            .try_flatten()
            .boxed()
    }

    async fn raw_updates(self) -> Self::UpdateStream {
        self.raw_update_stream().boxed()
    }
}

#[async_trait]
impl<K: AuthenticationStore + Send + Sync> APIInterface for HTTPAPIClient<K> {
    type Error = Error;

    async fn get_version(&mut self) -> Result<String, Self::Error> {
        let version: String = self.request("version", Method::GET).await?;
        Ok(version)
    }

    async fn get_conversations(&mut self) -> Result<Vec<Conversation>, Self::Error> {
        let conversations: Vec<Conversation> = self.request("conversations", Method::GET).await?;
        Ok(conversations)
    }

    async fn authenticate(&mut self, credentials: Credentials) -> Result<JwtToken, Self::Error> {
        #[derive(Deserialize, Debug)]
        struct AuthResponse {
            jwt: String,
        }

        log::debug!("Authenticating with username: {:?}", credentials.username);

        let body = || -> Body { serde_json::to_string(&credentials).unwrap().into() };
        let token: AuthResponse = self.request_with_body_retry("authenticate", Method::POST, body, false).await?;
        let token = JwtToken::new(&token.jwt).map_err(|_| Error::DecodeError)?;

        log::debug!("Saving token: {:?}", token);
        self.auth_store.set_token(token.clone()).await;

        Ok(token)
    }

    async fn get_messages(
        &mut self, 
        conversation_id: &ConversationID,
        limit: Option<u32>,
        before: Option<MessageID>,
        after: Option<MessageID>,
    ) -> Result<Vec<Message>, Self::Error> {
        let mut endpoint = format!("messages?guid={}", conversation_id);
        
        if let Some(limit_val) = limit {
            endpoint.push_str(&format!("&limit={}", limit_val));
        }
        
        if let Some(before_id) = before {
            endpoint.push_str(&format!("&beforeMessageGUID={}", before_id));
        }
        
        if let Some(after_id) = after {
            endpoint.push_str(&format!("&afterMessageGUID={}", after_id));
        }

        let messages: Vec<Message> = self.request(&endpoint, Method::GET).await?;
        Ok(messages)
    }

    async fn send_message(
        &mut self, 
        outgoing_message: &OutgoingMessage,
    ) -> Result<Message, Self::Error> {
        let message: Message = self.request_with_body(
            "sendMessage", 
            Method::POST, 
            || serde_json::to_string(&outgoing_message).unwrap().into() 
        ).await?;

        Ok(message)
    }

    async fn open_event_socket(&mut self) -> Result<WebsocketEventSocket, Self::Error> {
        use tungstenite::http::StatusCode;
        use tungstenite::handshake::client::Request as TungsteniteRequest;
        use tungstenite::handshake::client::generate_key;

        let uri = self.uri_for_endpoint("updates", Some(self.websocket_scheme()));

        log::debug!("Connecting to websocket: {:?}", uri);

        let auth = self.auth_store.get_token().await;
        let host = uri.authority().unwrap().host();
        let mut request = TungsteniteRequest::builder()
            .header("Host", host)
            .header("Connection", "Upgrade")
            .header("Upgrade", "websocket")
            .header("Sec-WebSocket-Version", "13")
            .header("Sec-WebSocket-Key", generate_key())
            .uri(uri.to_string())
            .body(())
            .expect("Unable to build websocket request");

        log::debug!("Websocket request: {:?}", request);

        if let Some(token) = &auth {
            let header_value = token.to_header_value().to_str().unwrap().parse().unwrap(); // ugh 
            request.headers_mut().insert("Authorization", header_value);
        }

        let (socket, response) = connect_async(request).await.map_err(Error::from)?;
        log::debug!("Websocket connected: {:?}", response.status());

        if response.status() != StatusCode::SWITCHING_PROTOCOLS {
            return Err(Error::ClientError("Websocket connection failed".into()));
        }

        Ok(WebsocketEventSocket::new(socket))
    }
}

impl<K: AuthenticationStore + Send + Sync> HTTPAPIClient<K> {
    pub fn new(base_url: Uri, auth_store: K) -> HTTPAPIClient<K> {
        HTTPAPIClient { 
            base_url,
            auth_store,
            client: Client::new(),
        }
    }

    fn uri_for_endpoint(&self, endpoint: &str, scheme: Option<&str>) -> Uri {
        let mut parts = self.base_url.clone().into_parts();
        let root_path: PathBuf = parts.path_and_query.unwrap().path().into();
        let path = root_path.join(endpoint);
        parts.path_and_query = Some(path.to_str().unwrap().parse().unwrap());

        if let Some(scheme) = scheme {
            parts.scheme = Some(scheme.parse().unwrap());
        }

        Uri::try_from(parts).unwrap()
    }

    fn websocket_scheme(&self) -> &str {
        if self.base_url.scheme().unwrap() == "https" {
            "wss"
        } else {
            "ws"
        }
    }

    async fn request<T: DeserializeOwned>(&mut self, endpoint: &str, method: Method) -> Result<T, Error> {
        self.request_with_body(endpoint, method, Body::empty).await
    }

    async fn request_with_body<T>(&mut self, endpoint: &str, method: Method, body_fn: impl Fn() -> Body) -> Result<T, Error> 
        where T: DeserializeOwned
    {
        self.request_with_body_retry(endpoint, method, body_fn, true).await
    }

    async fn request_with_body_retry<T>(
        &mut self, 
        endpoint: &str, 
        method: Method, 
        body_fn: impl Fn() -> Body, 
        retry_auth: bool) -> Result<T, Error>
    where 
        T: DeserializeOwned, 
    {
        use hyper::StatusCode;

        let uri = self.uri_for_endpoint(endpoint, None);
        log::debug!("Requesting {:?} {:?}", method, uri);

        let build_request = move |auth: &Option<JwtToken>| {
            let body = body_fn();
            Request::builder()
                .method(&method)
                .uri(&uri)
                .with_auth(auth)
                .body(body)
                .expect("Unable to build request")
        };

        let token = self.auth_store.get_token().await;
        let request = build_request(&token);
        let mut response = self.client.request(request).await?;

        log::debug!("-> Response: {:}", response.status());

        match response.status() {
            StatusCode::OK => { /* cool */ },

            // 401: Unauthorized. Token may have expired or is invalid. Attempt to renew. 
            StatusCode::UNAUTHORIZED => {
                if !retry_auth {
                    return Err(Error::ClientError("Unauthorized".into()));
                }

                if let Some(credentials) = &self.auth_store.get_credentials().await {
                    log::debug!("Renewing token using credentials: u: {:?}", credentials.username);
                    let new_token = self.authenticate(credentials.clone()).await?;

                    let request = build_request(&Some(new_token));
                    response = self.client.request(request).await?;
                } else {
                    return Err(Error::ClientError("Unauthorized, no credentials provided".into()));
                }
            },

            // Other errors: bubble up. 
            _ => {
                let message = format!("Request failed ({:})", response.status());
                return Err(Error::ClientError(message));    
            }
        }

        // Read and parse response body
        let body = hyper::body::to_bytes(response.into_body()).await?;
        let parsed: T = match serde_json::from_slice(&body) {
            Ok(result) => Ok(result),
            Err(json_err) => {
                log::error!("Error deserializing JSON: {:?}", json_err);
                log::error!("Body: {:?}", String::from_utf8_lossy(&body));

                // If JSON deserialization fails, try to interpret it as plain text
                // Unfortunately the server does return things like this...
                let s = str::from_utf8(&body).map_err(|_| Error::DecodeError)?;
                serde_plain::from_str(s).map_err(|_| json_err)
            }
        }?;

        Ok(parsed)
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::api::InMemoryAuthenticationStore;
    
    #[cfg(test)]
    fn local_mock_client() -> HTTPAPIClient<InMemoryAuthenticationStore> {
        let base_url = "http://localhost:5738".parse().unwrap();
        let credentials = Credentials {
            username: "test".to_string(),
            password: "test".to_string(),
        };

        HTTPAPIClient::new(base_url, InMemoryAuthenticationStore::new(Some(credentials)))
    }

    #[cfg(test)]
    async fn mock_client_is_reachable() -> bool {
        let mut client = local_mock_client();
        let version = client.get_version().await;

        match version {
            Ok(_) => true,
            Err(e) => {
                log::error!("Mock client error: {:?}", e);
                false 
            }
        }
    }

    #[tokio::test]
    async fn test_version() {
        if !mock_client_is_reachable().await {
            log::warn!("Skipping http_client tests (mock server not reachable)");
            return;
        }

        let mut client = local_mock_client();
        let version = client.get_version().await.unwrap();
        assert!(version.starts_with("KordophoneMock-"));
    }

    #[tokio::test]
    async fn test_conversations() {
        if !mock_client_is_reachable().await {
            log::warn!("Skipping http_client tests (mock server not reachable)");
            return;
        }

        let mut client = local_mock_client();
        let conversations = client.get_conversations().await.unwrap();
        assert!(!conversations.is_empty());
    }

    #[tokio::test]
    async fn test_messages() {
        if !mock_client_is_reachable().await {
            log::warn!("Skipping http_client tests (mock server not reachable)");
            return;
        }

        let mut client = local_mock_client();
        let conversations = client.get_conversations().await.unwrap();
        let conversation = conversations.first().unwrap();
        let messages = client.get_messages(&conversation.guid, None, None, None).await.unwrap();
        assert!(!messages.is_empty());
    }

    #[tokio::test]
    async fn test_updates() {
        if !mock_client_is_reachable().await {
            log::warn!("Skipping http_client tests (mock server not reachable)");
            return;
        }

        let mut client = local_mock_client();

        // We just want to see if the connection is established, we won't wait for any events
        let _ = client.open_event_socket().await.unwrap();
        assert!(true);
    }
}