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