1/*
2 * Copyright (c) 2023 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#include <atomic>
17#include <thread>
18#include "CommandLineInterface.h"
19#include "PreviewerEngineLog.h"
20#include "WebSocketServer.h"
21
22lws* WebSocketServer::webSocket = nullptr;
23std::atomic<bool> WebSocketServer::interrupted = false;
24WebSocketServer::WebSocketState WebSocketServer::webSocketWritable = WebSocketState::INIT;
25uint8_t* WebSocketServer::firstImageBuffer = nullptr;
26uint64_t WebSocketServer::firstImagebufferSize = 0;
27
28WebSocketServer::WebSocketServer() : serverThread(nullptr), serverPort(0)
29{
30    protocols[0] = {"ws", WebSocketServer::ProtocolCallback, 0, MAX_PAYLOAD_SIZE};
31    protocols[1] = {NULL, NULL, 0, 0};
32}
33
34WebSocketServer::~WebSocketServer() {}
35
36WebSocketServer& WebSocketServer::GetInstance()
37{
38    static WebSocketServer server;
39    return server;
40}
41
42void WebSocketServer::SetServerPort(int port)
43{
44    serverPort = port;
45}
46
47int WebSocketServer::ProtocolCallback(struct lws* wsi,
48                                      enum lws_callback_reasons reason,
49                                      void* user,
50                                      void* in,
51                                      size_t len)
52{
53    switch (reason) {
54        case LWS_CALLBACK_PROTOCOL_INIT:
55            ILOG("Engine Websocket protocol init");
56            break;
57        case LWS_CALLBACK_ESTABLISHED:
58            ILOG("Websocket client connect");
59            webSocket = wsi;
60            lws_callback_on_writable(wsi);
61            break;
62        case LWS_CALLBACK_RECEIVE:
63            break;
64        case LWS_CALLBACK_SERVER_WRITEABLE:
65            ILOG("Engine websocket server writeable");
66            if (firstImagebufferSize > 0 && webSocketWritable == WebSocketState::UNWRITEABLE) {
67                ILOG("Send last image after websocket reconnected");
68                std::lock_guard<std::mutex> guard(WebSocketServer::GetInstance().mutex);
69                lws_write(wsi,
70                          firstImageBuffer + LWS_PRE,
71                          firstImagebufferSize,
72                          LWS_WRITE_BINARY);
73            }
74            webSocketWritable = WebSocketState::WRITEABLE;
75            break;
76        case LWS_CALLBACK_CLOSED:
77            ILOG("Websocket client connection closed");
78            webSocketWritable = WebSocketState::UNWRITEABLE;
79            break;
80        default:
81            break;
82    }
83    return 0;
84}
85
86void WebSocketServer::SignalHandler(int sig)
87{
88    interrupted = true;
89}
90
91void WebSocketServer::StartWebsocketListening()
92{
93    const auto sig = signal(SIGINT, SignalHandler);
94    if (sig == SIG_ERR) {
95        ELOG("StartWebsocketListening failed");
96        return;
97    }
98    ILOG("Begin to start websocket listening!");
99    struct lws_context_creation_info contextInfo = {0};
100    contextInfo.port = serverPort;
101    contextInfo.iface = serverHostname;
102    contextInfo.protocols = protocols;
103    contextInfo.ip_limit_wsi = websocketMaxConn;
104    contextInfo.options  = LWS_SERVER_OPTION_VALIDATE_UTF8;
105    struct lws_context* context = lws_create_context(&contextInfo);
106    if (context == nullptr) {
107        ELOG("WebSocketServer::StartWebsocketListening context memory allocation failed");
108        return;
109    }
110    while (!interrupted) {
111        if (lws_service(context, WEBSOCKET_SERVER_TIMEOUT)) {
112            interrupted = true;
113        }
114    }
115    lws_context_destroy(context);
116}
117
118void WebSocketServer::Run()
119{
120    serverThread = std::make_unique<std::thread>([this]() {
121        this->StartWebsocketListening();
122    });
123    if (serverThread == nullptr) {
124        ELOG("WebSocketServer::Start serverThread memory allocation failed");
125    }
126    serverThread->detach();
127}
128
129size_t WebSocketServer::WriteData(unsigned char* data, size_t length)
130{
131    while (webSocketWritable != WebSocketState::WRITEABLE) {
132        std::this_thread::sleep_for(std::chrono::milliseconds(1));
133    }
134    if (webSocket != nullptr && webSocketWritable == WebSocketState::WRITEABLE) {
135        return lws_write(webSocket, data, length, LWS_WRITE_BINARY);
136    }
137    return 0;
138}
139