1 /*
2 * Copyright (C) 2021 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
17 #include "src/websocket_bridge/websocket_bridge.h"
18
19 #include <stdint.h>
20
21 #include <map>
22 #include <memory>
23 #include <vector>
24
25 #include "perfetto/ext/base/http/http_server.h"
26 #include "perfetto/ext/base/unix_socket.h"
27 #include "perfetto/ext/base/unix_task_runner.h"
28 #include "perfetto/ext/tracing/ipc/default_socket.h"
29
30 namespace perfetto {
31 namespace {
32
33 constexpr int kWebsocketPort = 8037;
34
35 struct Endpoint {
36 const char* uri;
37 const char* endpoint;
38 base::SockFamily family;
39 };
40
41 class WSBridge : public base::HttpRequestHandler,
42 public base::UnixSocket::EventListener {
43 public:
44 void Main(int argc, char** argv);
45
46 // base::HttpRequestHandler implementation.
47 void OnHttpRequest(const base::HttpRequest&) override;
48 void OnWebsocketMessage(const base::WebsocketMessage&) override;
49 void OnHttpConnectionClosed(base::HttpServerConnection*) override;
50
51 // base::UnixSocket::EventListener implementation.
52 void OnNewIncomingConnection(base::UnixSocket*,
53 std::unique_ptr<base::UnixSocket>) override;
54 void OnConnect(base::UnixSocket*, bool) override;
55 void OnDisconnect(base::UnixSocket*) override;
56 void OnDataAvailable(base::UnixSocket* self) override;
57
58 private:
59 base::HttpServerConnection* GetWebsocket(base::UnixSocket*);
60
61 base::UnixTaskRunner task_runner_;
62 std::vector<Endpoint> endpoints_;
63 std::map<base::HttpServerConnection*, std::unique_ptr<base::UnixSocket>>
64 conns_;
65 };
66
Main(int,char **)67 void WSBridge::Main(int, char**) {
68 #if PERFETTO_BUILDFLAG(PERFETTO_OS_WIN)
69 // On Windows traced used a TCP socket.
70 const auto kTracedFamily = base::SockFamily::kInet;
71 #else
72 const auto kTracedFamily = base::SockFamily::kUnix;
73 #endif
74 endpoints_.push_back({"/traced", GetConsumerSocket(), kTracedFamily});
75 endpoints_.push_back({"/adb", "127.0.0.1:5037", base::SockFamily::kInet});
76
77 base::HttpServer srv(&task_runner_, this);
78 srv.AddAllowedOrigin("http://localhost:10000");
79 srv.AddAllowedOrigin("http://127.0.0.1:10000");
80 srv.AddAllowedOrigin("https://ui.perfetto.dev");
81
82 srv.Start(kWebsocketPort);
83 PERFETTO_LOG("[WSBridge] Listening on 127.0.0.1:%d", kWebsocketPort);
84 task_runner_.Run();
85 }
86
OnHttpRequest(const base::HttpRequest & req)87 void WSBridge::OnHttpRequest(const base::HttpRequest& req) {
88 for (const auto& ep : endpoints_) {
89 if (req.uri != ep.uri || !req.is_websocket_handshake)
90 continue;
91
92 // Connect to the endpoint in blocking mode.
93 auto sock_raw =
94 base::UnixSocketRaw::CreateMayFail(ep.family, base::SockType::kStream);
95 if (!sock_raw) {
96 PERFETTO_PLOG("socket() failed");
97 req.conn->SendResponseAndClose("500 Server Error");
98 return;
99 }
100 PERFETTO_LOG("[WSBridge] New connection from \"%.*s\"",
101 static_cast<int>(req.origin.size()), req.origin.data());
102 sock_raw.SetTxTimeout(3000);
103 sock_raw.SetBlocking(true);
104
105 if (!sock_raw.Connect(ep.endpoint)) {
106 PERFETTO_ELOG("[WSBridge] Connection to %s failed", ep.endpoint);
107 req.conn->SendResponseAndClose("503 Service Unavailable");
108 return;
109 }
110 sock_raw.SetBlocking(false);
111
112 PERFETTO_DLOG("[WSBridge] Connected to %s", ep.endpoint);
113 conns_[req.conn] = base::UnixSocket::AdoptConnected(
114 sock_raw.ReleaseFd(), this, &task_runner_, ep.family,
115 base::SockType::kStream);
116
117 req.conn->UpgradeToWebsocket(req);
118 return;
119 } // for (endpoint)
120 req.conn->SendResponseAndClose("404 Not Found");
121 }
122
123 // Called when an inbound websocket message is received from the browser.
OnWebsocketMessage(const base::WebsocketMessage & msg)124 void WSBridge::OnWebsocketMessage(const base::WebsocketMessage& msg) {
125 auto it = conns_.find(msg.conn);
126 PERFETTO_CHECK(it != conns_.end());
127 // Pass through the websocket message onto the endpoint TCP socket.
128 base::UnixSocket& sock = *it->second;
129 sock.Send(msg.data.data(), msg.data.size());
130 }
131
132 // Called when a TCP message is received from the endpoint.
OnDataAvailable(base::UnixSocket * sock)133 void WSBridge::OnDataAvailable(base::UnixSocket* sock) {
134 base::HttpServerConnection* websocket = GetWebsocket(sock);
135 PERFETTO_CHECK(websocket);
136
137 char buf[8192];
138 auto rsize = sock->Receive(buf, sizeof(buf));
139 if (rsize > 0) {
140 websocket->SendWebsocketMessage(buf, static_cast<size_t>(rsize));
141 } else {
142 // Connection closed or errored.
143 sock->Shutdown(/*notify=*/true); // Will trigger OnDisconnect().
144 websocket->Close();
145 }
146 }
147
148 // Called when the browser terminates the websocket connection.
OnHttpConnectionClosed(base::HttpServerConnection * websocket)149 void WSBridge::OnHttpConnectionClosed(base::HttpServerConnection* websocket) {
150 PERFETTO_DLOG("[WSBridge] Websocket connection closed");
151 auto it = conns_.find(websocket);
152 if (it == conns_.end())
153 return; // Can happen if ADB closed first.
154 base::UnixSocket& sock = *it->second;
155 sock.Shutdown(/*notify=*/true);
156 conns_.erase(websocket);
157 }
158
OnDisconnect(base::UnixSocket * sock)159 void WSBridge::OnDisconnect(base::UnixSocket* sock) {
160 base::HttpServerConnection* websocket = GetWebsocket(sock);
161 if (!websocket)
162 return;
163 websocket->Close();
164 sock->Shutdown(/*notify=*/false);
165 conns_.erase(websocket);
166 PERFETTO_DLOG("[WSBridge] Socket connection closed");
167 }
168
GetWebsocket(base::UnixSocket * sock)169 base::HttpServerConnection* WSBridge::GetWebsocket(base::UnixSocket* sock) {
170 for (const auto& it : conns_) {
171 if (it.second.get() == sock) {
172 return it.first;
173 }
174 }
175 return nullptr;
176 }
177
OnConnect(base::UnixSocket *,bool)178 void WSBridge::OnConnect(base::UnixSocket*, bool) {}
OnNewIncomingConnection(base::UnixSocket *,std::unique_ptr<base::UnixSocket>)179 void WSBridge::OnNewIncomingConnection(base::UnixSocket*,
180 std::unique_ptr<base::UnixSocket>) {}
181
182 } // namespace
183
WebsocketBridgeMain(int argc,char ** argv)184 int PERFETTO_EXPORT_ENTRYPOINT WebsocketBridgeMain(int argc, char** argv) {
185 perfetto::WSBridge ws_bridge;
186 ws_bridge.Main(argc, argv);
187 return 0;
188 }
189
190 } // namespace perfetto
191