actix-web/actix-http/src/ws/client/connect.rs

//! Http client request
use std::str;

#[cfg(feature = "cookies")]
use cookie::Cookie;
use http::header::{HeaderName, HeaderValue};
use http::{Error as HttpError, HttpTryFrom, Uri};

use super::ClientError;
use crate::header::IntoHeaderValue;
use crate::message::RequestHead;

/// `WebSocket` connection
pub struct Connect {
    pub(super) head: RequestHead,
    pub(super) err: Option<ClientError>,
    pub(super) http_err: Option<HttpError>,
    pub(super) origin: Option<HeaderValue>,
    pub(super) protocols: Option<String>,
    pub(super) max_size: usize,
    pub(super) server_mode: bool,
}

impl Connect {
    /// Create new websocket connection
    pub fn new<S: AsRef<str>>(uri: S) -> Connect {
        let mut cl = Connect {
            head: RequestHead::default(),
            err: None,
            http_err: None,
            origin: None,
            protocols: None,
            max_size: 65_536,
            server_mode: false,
        };

        match Uri::try_from(uri.as_ref()) {
            Ok(uri) => cl.head.uri = uri,
            Err(e) => cl.http_err = Some(e.into()),
        }

        cl
    }

    /// Set supported websocket protocols
    pub fn protocols<U, V>(mut self, protos: U) -> Self
    where
        U: IntoIterator<Item = V> + 'static,
        V: AsRef<str>,
    {
        let mut protos = protos
            .into_iter()
            .fold(String::new(), |acc, s| acc + s.as_ref() + ",");
        protos.pop();
        self.protocols = Some(protos);
        self
    }

    // #[cfg(feature = "cookies")]
    // /// Set cookie for handshake request
    // pub fn cookie(mut self, cookie: Cookie) -> Self {
    //    self.request.cookie(cookie);
    //    self
    // }

    /// Set request Origin
    pub fn origin<V>(mut self, origin: V) -> Self
    where
        HeaderValue: HttpTryFrom<V>,
    {
        match HeaderValue::try_from(origin) {
            Ok(value) => self.origin = Some(value),
            Err(e) => self.http_err = Some(e.into()),
        }
        self
    }

    /// Set max frame size
    ///
    /// By default max size is set to 64kb
    pub fn max_frame_size(mut self, size: usize) -> Self {
        self.max_size = size;
        self
    }

    /// Disable payload masking. By default ws client masks frame payload.
    pub fn server_mode(mut self) -> Self {
        self.server_mode = true;
        self
    }

    /// Set request header
    pub fn header<K, V>(mut self, key: K, value: V) -> Self
    where
        HeaderName: HttpTryFrom<K>,
        V: IntoHeaderValue,
    {
        match HeaderName::try_from(key) {
            Ok(key) => match value.try_into() {
                Ok(value) => {
                    self.head.headers.append(key, value);
                }
                Err(e) => self.http_err = Some(e.into()),
            },
            Err(e) => self.http_err = Some(e.into()),
        }
        self
    }
}