1/*
2 * Copyright (c) 2022 Huawei Device Co., Ltd.
3 * Licensed under the Apache License, Version 2.0 (the "License");
4 * you may not use this file except in compliance with the License.
5 * You may obtain a copy of the License at
6 *
7 *     http://www.apache.org/licenses/LICENSE-2.0
8 *
9 * Unless required by applicable law or agreed to in writing, software
10 * distributed under the License is distributed on an "AS IS" BASIS,
11 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 * See the License for the specific language governing permissions and
13 * limitations under the License.
14 */
15
16#ifndef ARKCOMPILER_TOOLCHAIN_WEBSOCKET_SERVER_WEBSOCKET_SERVER_H
17#define ARKCOMPILER_TOOLCHAIN_WEBSOCKET_SERVER_WEBSOCKET_SERVER_H
18
19#include "http.h"
20#include "websocket_base.h"
21
22#include <functional>
23
24namespace OHOS::ArkCompiler::Toolchain {
25class WebSocketServer final : public WebSocketBase {
26public:
27    using ValidateConnectionCallback = std::function<bool(const HttpRequest&)>;
28    using OpenConnectionCallback = std::function<void()>;
29
30public:
31    ~WebSocketServer() noexcept override;
32
33    /**
34     * @brief Accept new posix-socket connection.
35     * Safe to call concurrently with `Close`.
36     * Must not be called concurrently with `SendReply` or `Decode`, as it might lead to race condition.
37     */
38    bool AcceptNewConnection();
39
40    /**
41     * @brief Initialize server posix-socket.
42     * Non thread safe.
43     * On success, `serverFd_` is initialized and `serverUp_` evaluates true.
44     * @param port server TCP socket port number.
45     * @param timeoutLimit timeout in seconds for server socket. If zero, timeout is not set.
46     * @returns true on success, false otherwise.
47     */
48    bool InitTcpWebSocket(int port, uint32_t timeoutLimit = 0);
49
50#if !defined(WINDOWS_PLATFORM)
51    /**
52     * @brief Initialize server unix-socket.
53     * Non thread safe.
54     * On success, `serverFd_` is initialized and `serverUp_` evaluates true.
55     * @param sockName server socket name.
56     * @param timeoutLimit timeout in seconds for server socket. If zero, timeout is not set.
57     * @returns true on success, false otherwise.
58     */
59    bool InitUnixWebSocket(const std::string& sockName, uint32_t timeoutLimit = 0);
60
61    /**
62     * @brief Initialize connection with unix-socket.
63     * Non thread safe.
64     * On success, `serverFd_` stays uninitialized, but `serverUp_` evaluates true.
65     * Note that this mode supports only a single connection,
66     * which must be accepted by calling `ConnectUnixWebSocketBySocketpair`.
67     * @param socketfd connection socket file descriptor, must be correctly opened before calling the method.
68     * @returns true on success, false otherwise.
69     */
70    bool InitUnixWebSocket(int socketfd);
71
72    /**
73     * @brief Accept new unix-socket connection.
74     * Safe to call concurrently with `Close`.
75     */
76    bool ConnectUnixWebSocketBySocketpair();
77#endif  // WINDOWS_PLATFORM
78
79    /**
80     * @brief Set callback for calling after received HTTP handshake request.
81     * Non thread safe.
82     */
83    void SetValidateConnectionCallback(ValidateConnectionCallback cb);
84
85    /**
86     * @brief Set callback for calling after accepted new connection.
87     * Non thread safe.
88     */
89    void SetOpenConnectionCallback(OpenConnectionCallback cb);
90
91    /**
92     * @brief Close server endpoint and connection sockets.
93     * Safe to call concurrently with:
94     * `AcceptNewConnection`, `ConnectUnixWebSocketBySocketpair` `SendReply`, `Decode`, `CloseConnection`.
95     */
96    void Close();
97
98private:
99    bool BindAndListenTcpWebSocket(int port);
100
101    bool ValidateIncomingFrame(const WebSocketFrame& wsFrame) const override;
102    std::string CreateFrame(bool isLast, FrameType frameType) const override;
103    std::string CreateFrame(bool isLast, FrameType frameType, const std::string& payload) const override;
104    std::string CreateFrame(bool isLast, FrameType frameType, std::string&& payload) const override;
105    bool DecodeMessage(WebSocketFrame& wsFrame) const override;
106
107    bool HttpHandShake();
108    bool ProtocolUpgrade(const HttpRequest& req);
109    bool ResponseInvalidHandShake() const;
110
111    /**
112     * @brief Runs user-provided callback and performs transition from `CONNECTING` to `OPEN` state.
113     */
114    void OnNewConnection();
115
116    void CloseServerSocket();
117
118    /**
119     * @brief Performs transition from `CLOSED` to `CONNECTING` state if the server is up.
120     * @returns true on success, false otherwise.
121     */
122    bool MoveToConnectingState();
123
124    /**
125     * @brief Wait until concurrent `CONNECTING` state transition ends.
126     * There might be a concurrent call to `AcceptNewConnection`;
127     * in this case, it must finish and move into either:
128     * - `CLOSED` (due to failed `accept` or handshake),
129     * - `OPEN` (which can then concurrently transition to either `CLOSING` or `CLOSE`).
130     * @param connection previously loaded connection state.
131     * @returns updated connection state.
132     */
133    ConnectionState WaitConnectingStateEnds(ConnectionState connection);
134
135private:
136    // Server initialization status.
137    // Note that there could be alive connections after `serverUp_` switched to false,
138    // but they must terminate soon after. Users must track this with callbacks.
139    std::atomic_bool serverUp_ {false};
140
141    int32_t serverFd_ {-1};
142
143    // Callbacks used during different stages of connection lifecycle.
144    // E.g. validation callback is executed during handshake
145    // and used to indicate whether the incoming connection should be accepted.
146    ValidateConnectionCallback validateCb_;
147    OpenConnectionCallback openCb_;
148
149    static constexpr std::string_view BAD_REQUEST_RESPONSE = "HTTP/1.1 400 Bad Request\r\n\r\n";
150    static constexpr int NET_SUCCESS = 1;
151};
152} // namespace OHOS::ArkCompiler::Toolchain
153
154#endif // ARKCOMPILER_TOOLCHAIN_WEBSOCKET_SERVER_WEBSOCKET_SERVER_H
155