axum-tws 0.1.0

A high-performance WebSocket backend for axum powered by tokio-websockets.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101

#![doc = include_str!("../README.md")]
#![forbid(unsafe_code)]

pub mod upgrade;
pub mod websocket;

use std::fmt::Display;

use axum_core::body::Body;
use axum_core::response::IntoResponse;
use axum_core::response::Response;
use http::StatusCode;

pub use tokio_websockets::*;

pub use crate::{upgrade::WebSocketUpgrade, websocket::WebSocket};

#[derive(Debug)]
pub enum WebSocketError {
    ConnectionNotUpgradeable,
    Internal(tokio_websockets::Error),
    InvalidConnectionHeader,
    InvalidUpgradeHeader,
    InvalidWebSocketVersionHeader,
    MethodNotGet,
    UpgradeFailed(hyper::Error),
}

impl Display for WebSocketError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            WebSocketError::ConnectionNotUpgradeable => {
                write!(f, "connection is not upgradeable")
            }
            WebSocketError::Internal(e) => {
                write!(f, "internal server error: {}", e)
            }
            WebSocketError::InvalidConnectionHeader => {
                write!(f, "invalid `Connection` header")
            }
            WebSocketError::InvalidUpgradeHeader => {
                write!(f, "invalid `Upgrade` header")
            }
            WebSocketError::InvalidWebSocketVersionHeader => {
                write!(f, "invalid `Sec-WebSocket-Version` header")
            }
            WebSocketError::MethodNotGet => {
                write!(f, "http request method must be `GET`")
            }
            WebSocketError::UpgradeFailed(e) => {
                write!(f, "upgrade failed: {}", e)
            }
        }
    }
}

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

impl IntoResponse for WebSocketError {
    fn into_response(self) -> Response<Body> {
        let status = match self {
            WebSocketError::ConnectionNotUpgradeable => StatusCode::UPGRADE_REQUIRED,

            // Request headers are invalid or missing.
            WebSocketError::InvalidConnectionHeader
            | WebSocketError::InvalidUpgradeHeader
            | WebSocketError::InvalidWebSocketVersionHeader => StatusCode::BAD_REQUEST,

            // Invalid request method.
            WebSocketError::MethodNotGet => StatusCode::METHOD_NOT_ALLOWED,

            // All other errors will be treated as internal server errors.
            _ => StatusCode::INTERNAL_SERVER_ERROR,
        };

        Response::builder()
            .status(status)
            .body(Body::empty())
            .unwrap()
    }
}

impl From<tokio_websockets::Error> for WebSocketError {
    fn from(e: tokio_websockets::Error) -> Self {
        WebSocketError::Internal(e)
    }
}

impl From<hyper::Error> for WebSocketError {
    fn from(e: hyper::Error) -> Self {
        WebSocketError::UpgradeFailed(e)
    }
}