wmhttp 0.3.1

a http server for rust
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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

// Copyright 2022 - 2023 Wenmeng See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
//
// Author: tickbh
// -----
// Created Date: 2023/09/14 09:42:25

use crate::{ws::WsCodec, ProtResult};

use algorithm::buf::{Binary, Bt};
use std::{
    pin::Pin,
    task::{ready, Context, Poll},
};
use tokio::io::{AsyncRead, AsyncWrite};

pub struct WsStateHandshake {
    /// 当前握手状态
    state: WsHandshaking,
    /// 是否为客户端
    is_client: bool,
}

/// 握手状态
enum WsHandshaking {
    /// 还未进行握手, 确定http2协议则开始握手
    Wait,
    /// 协议升级信息写入
    Flushing(Flush),
    /// 已完成握手, 不可重复握手
    Done,
}

/// Flush a Sink
struct Flush(Binary);

impl WsStateHandshake {
    pub fn new_server() -> WsStateHandshake {
        WsStateHandshake {
            state: WsHandshaking::Wait,
            is_client: false,
        }
    }

    // pub fn new_client() -> WsStateHandshake {
    //     WsStateHandshake {
    //         state: WsHandshaking::Wait,
    //         is_client: true,
    //     }
    // }

    pub fn poll_handle<T>(
        &mut self,
        cx: &mut Context<'_>,
        codec: &mut WsCodec<T>,
    ) -> Poll<ProtResult<()>>
    where
        T: AsyncRead + AsyncWrite + Unpin,
    {
        loop {
            match &mut self.state {
                WsHandshaking::Wait => {
                    return Poll::Pending;
                }
                WsHandshaking::Flushing(flush) => {
                    match ready!(flush.poll_handle(cx, codec)) {
                        Ok(_) => {
                            tracing::trace!(flush.poll = %"Ready");
                            self.state = WsHandshaking::Done;
                            continue;
                        }
                        Err(e) => return Poll::Ready(Err(e)),
                    };
                }
                WsHandshaking::Done => {
                    return Poll::Ready(Ok(()));
                }
            }
        }
    }

    pub fn set_handshake_status(&mut self, binary: Binary, is_client: bool) {
        self.is_client = is_client;
        if binary.is_empty() {
            self.state = WsHandshaking::Done;
        } else {
            self.state = WsHandshaking::Flushing(Flush(binary));
        }
    }
}

impl Flush {
    pub fn poll_handle<T>(
        &mut self,
        cx: &mut Context<'_>,
        codec: &mut WsCodec<T>,
    ) -> Poll<ProtResult<()>>
    where
        T: AsyncRead + AsyncWrite + Unpin,
    {
        if !self.0.has_remaining() {
            return Poll::Ready(Ok(()));
        }

        loop {
            match ready!(Pin::new(codec.get_mut()).poll_write(cx, self.0.chunk())) {
                Ok(n) => {
                    self.0.advance(n);
                }
                Err(e) => return Poll::Ready(Err(e.into())),
            }
            if !self.0.has_remaining() {
                return Poll::Ready(Ok(()));
            }
        }
    }
}

unsafe impl Send for WsStateHandshake {}

unsafe impl Sync for WsStateHandshake {}