1 //
2 // Copyright (C) 2020 The Android Open Source Project
3 //
4 // Licensed under the Apache License, Version 2.0 (the "License");
5 // you may not use this file except in compliance with the License.
6 // You may obtain a copy of the License at
7 //
8 // http://www.apache.org/licenses/LICENSE-2.0
9 //
10 // Unless required by applicable law or agreed to in writing, software
11 // distributed under the License is distributed on an "AS IS" BASIS,
12 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 // See the License for the specific language governing permissions and
14 // limitations under the License.
15
16 #include "host/libs/websocket/websocket_handler.h"
17
18 #include <android-base/logging.h>
19 #include <libwebsockets.h>
20
21 namespace cuttlefish {
22
WebSocketHandler(struct lws * wsi)23 WebSocketHandler::WebSocketHandler(struct lws* wsi) : wsi_(wsi) {}
24
EnqueueMessage(const uint8_t * data,size_t len,bool binary)25 void WebSocketHandler::EnqueueMessage(const uint8_t* data, size_t len,
26 bool binary) {
27 std::vector<uint8_t> buffer(LWS_PRE + len, 0);
28 std::copy(data, data + len, buffer.begin() + LWS_PRE);
29 buffer_queue_.emplace_front(std::move(buffer), binary);
30 lws_callback_on_writable(wsi_);
31 }
32
33 // Attempts to write what's left on a websocket buffer to the websocket,
34 // updating the buffer.
WriteWsBuffer(WebSocketHandler::WsBuffer & ws_buffer)35 void WebSocketHandler::WriteWsBuffer(WebSocketHandler::WsBuffer& ws_buffer) {
36 auto len = ws_buffer.data.size() - LWS_PRE;
37 auto flags = lws_write_ws_flags(
38 ws_buffer.binary ? LWS_WRITE_BINARY : LWS_WRITE_TEXT, true, true);
39 auto res = lws_write(wsi_, &ws_buffer.data[LWS_PRE], len,
40 static_cast<enum lws_write_protocol>(flags));
41 // lws_write will write all bytes of the provided buffer or enqueue the ones
42 // it couldn't write for later, but it guarantees it will consume the entire
43 // buffer, so we only need to check for error.
44 if (res < 0) {
45 // This shouldn't happen since this function is called in response to a
46 // LWS_CALLBACK_SERVER_WRITEABLE call.
47 LOG(FATAL) << "Failed to write data on the websocket";
48 }
49 }
50
OnWritable()51 bool WebSocketHandler::OnWritable() {
52 if (buffer_queue_.empty()) {
53 return close_;
54 }
55 WriteWsBuffer(buffer_queue_.back());
56 buffer_queue_.pop_back();
57
58 if (!buffer_queue_.empty()) {
59 lws_callback_on_writable(wsi_);
60 }
61 // Only close if there are no more queued writes
62 return buffer_queue_.empty() && close_;
63 }
64
Close()65 void WebSocketHandler::Close() {
66 close_ = true;
67 lws_callback_on_writable(wsi_);
68 }
69
70 } // namespace cuttlefish
71