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 #include "host/libs/websocket/websocket_server.h"
22
23 namespace cuttlefish {
24
25 namespace {
AppendData(const char * data,size_t len,std::string & buffer)26 void AppendData(const char* data, size_t len, std::string& buffer) {
27 auto ptr = reinterpret_cast<const uint8_t*>(data);
28 buffer.reserve(buffer.size() + len);
29 buffer.insert(buffer.end(), ptr, ptr + len);
30 }
31 } // namespace
32
WebSocketHandler(struct lws * wsi)33 WebSocketHandler::WebSocketHandler(struct lws* wsi) : wsi_(wsi) {}
34
EnqueueMessage(const uint8_t * data,size_t len,bool binary)35 void WebSocketHandler::EnqueueMessage(const uint8_t* data, size_t len,
36 bool binary) {
37 std::vector<uint8_t> buffer(LWS_PRE + len, 0);
38 std::copy(data, data + len, buffer.begin() + LWS_PRE);
39 buffer_queue_.emplace_front(std::move(buffer), binary);
40 lws_callback_on_writable(wsi_);
41 }
42
43 // Attempts to write what's left on a websocket buffer to the websocket,
44 // updating the buffer.
WriteWsBuffer(WebSocketHandler::WsBuffer & ws_buffer)45 void WebSocketHandler::WriteWsBuffer(WebSocketHandler::WsBuffer& ws_buffer) {
46 auto len = ws_buffer.data.size() - LWS_PRE;
47 // For http2 there must be LWS_PRE bytes at the end as well.
48 ws_buffer.data.resize(ws_buffer.data.size() + LWS_PRE);
49 auto flags = lws_write_ws_flags(
50 ws_buffer.binary ? LWS_WRITE_BINARY : LWS_WRITE_TEXT, true, true);
51 auto res = lws_write(wsi_, &ws_buffer.data[LWS_PRE], len,
52 static_cast<enum lws_write_protocol>(flags));
53 // lws_write will write all bytes of the provided buffer or enqueue the ones
54 // it couldn't write for later, but it guarantees it will consume the entire
55 // buffer, so we only need to check for error.
56 if (res < 0) {
57 // This shouldn't happen since this function is called in response to a
58 // LWS_CALLBACK_SERVER_WRITEABLE call.
59 LOG(ERROR) << "Failed to write data on the websocket";
60 }
61 }
62
OnWritable()63 bool WebSocketHandler::OnWritable() {
64 if (buffer_queue_.empty()) {
65 return close_;
66 }
67 WriteWsBuffer(buffer_queue_.back());
68 buffer_queue_.pop_back();
69
70 if (!buffer_queue_.empty()) {
71 lws_callback_on_writable(wsi_);
72 }
73 // Only close if there are no more queued writes
74 return buffer_queue_.empty() && close_;
75 }
76
Close()77 void WebSocketHandler::Close() {
78 close_ = true;
79 lws_callback_on_writable(wsi_);
80 }
81
DynHandler(struct lws * wsi)82 DynHandler::DynHandler(struct lws* wsi) : wsi_(wsi), out_buffer_(LWS_PRE, 0) {}
83
AppendDataOut(const std::string & data)84 void DynHandler::AppendDataOut(const std::string& data) {
85 AppendData(data.c_str(), data.size(), out_buffer_);
86 }
87
AppendDataIn(void * data,size_t len)88 void DynHandler::AppendDataIn(void* data, size_t len) {
89 AppendData(reinterpret_cast<char*>(data), len, in_buffer_);
90 }
91
OnWritable()92 int DynHandler::OnWritable() {
93 auto len = out_buffer_.size() - LWS_PRE;
94 // For http2 there must be LWS_PRE bytes at the end as well.
95 out_buffer_.resize(out_buffer_.size() + LWS_PRE);
96 auto res = lws_write(wsi_, reinterpret_cast<uint8_t*>(&out_buffer_[LWS_PRE]),
97 len, LWS_WRITE_HTTP_FINAL);
98 if (res != len) {
99 // This shouldn't happen since this function is called in response to a
100 // LWS_CALLBACK_SERVER_WRITEABLE call.
101 LOG(ERROR) << "Failed to write HTTP response";
102 }
103 return lws_http_transaction_completed(wsi_);
104 }
content_len() const105 size_t DynHandler::content_len() const { return out_buffer_.size() - LWS_PRE; }
106 } // namespace cuttlefish
107