Implement /api/get_bottleneck_stats

common/network/GetAPI.h

@@ -24,6 +24,8 @@
 #include <rfb/PixelBuffer.h>
 #include <rfb/PixelFormat.h>
 #include <stdint.h>
+#include <map>
+#include <string>
 #include <vector>
 
 namespace network {
@@ -34,6 +36,7 @@ namespace network {
 
     // from main thread
     void mainUpdateScreen(rfb::PixelBuffer *pb);
+    void mainUpdateBottleneckStats(const char userid[], const char stats[]);
 
     // from network threads
     uint8_t *netGetScreenshot(uint16_t w, uint16_t h,
@@ -42,6 +45,7 @@ namespace network {
     uint8_t netAddUser(const char name[], const char pw[], const bool write);
     uint8_t netRemoveUser(const char name[]);
     uint8_t netGiveControlTo(const char name[]);
+    void netGetBottleneckStats(char *buf, uint32_t len);
 
     enum USER_ACTION {
       //USER_ADD, - handled locally for interactivity
@@ -68,6 +72,9 @@ namespace network {
     std::vector<uint8_t> cachedJpeg;
     uint16_t cachedW, cachedH;
     uint8_t cachedQ;
+
+    std::map<std::string, std::string> bottleneckStats;
+    pthread_mutex_t statMutex;
   };
 
 }

common/network/GetAPIMessager.cxx

@@ -60,6 +60,7 @@ GetAPIMessager::GetAPIMessager(const char *passwdfile_): passwdfile(passwdfile_)
 
 	pthread_mutex_init(&screenMutex, NULL);
 	pthread_mutex_init(&userMutex, NULL);
+	pthread_mutex_init(&statMutex, NULL);
 }
 
 // from main thread
@@ -95,6 +96,15 @@ void GetAPIMessager::mainUpdateScreen(rfb::PixelBuffer *pb) {
 	pthread_mutex_unlock(&screenMutex);
 }
 
+void GetAPIMessager::mainUpdateBottleneckStats(const char userid[], const char stats[]) {
+	if (pthread_mutex_trylock(&statMutex))
+		return;
+
+	bottleneckStats[userid] = stats;
+
+	pthread_mutex_unlock(&statMutex);
+}
+
 // from network threads
 uint8_t *GetAPIMessager::netGetScreenshot(uint16_t w, uint16_t h,
 	const uint8_t q, const bool dedup,
@@ -286,3 +296,70 @@ uint8_t GetAPIMessager::netGiveControlTo(const char name[]) {
 
 	return 1;
 }
+
+void GetAPIMessager::netGetBottleneckStats(char *buf, uint32_t len) {
+/*
+{
+    "username.1": {
+        "192.168.100.2:14908": [ 100, 100, 100, 100 ],
+        "192.168.100.3:14918": [ 100, 100, 100, 100 ]
+    },
+    "username.2": {
+        "192.168.100.5:14904": [ 100, 100, 100, 100 ]
+    }
+}
+*/
+	std::map<std::string, std::string>::const_iterator it;
+	const char *prev = NULL;
+	FILE *f;
+
+	if (pthread_mutex_lock(&statMutex)) {
+		buf[0] = 0;
+		return;
+	}
+
+	// Conservative estimate
+	if (len < bottleneckStats.size() * 60) {
+		buf[0] = 0;
+		goto out;
+	}
+
+	f = fmemopen(buf, len, "w");
+
+	fprintf(f, "{\n");
+
+	for (it = bottleneckStats.begin(); it != bottleneckStats.end(); it++) {
+		// user@127.0.0.1_1627311208.791752::websocket
+		const char *id = it->first.c_str();
+		const char *data = it->second.c_str();
+
+		const char *at = strchr(id, '@');
+		if (!at)
+			continue;
+
+		const unsigned userlen = at - id;
+		if (prev && !strncmp(prev, id, userlen)) {
+			// Same user
+			fprintf(f, ",\n\t\t\"%s\": %s", at + 1, data);
+		} else {
+			// New one
+			if (prev) {
+				fprintf(f, "\n\t},\n");
+			}
+			fprintf(f, "\t\"%.*s\": {\n", userlen, id);
+			fprintf(f, "\t\t\"%s\": %s", at + 1, data);
+		}
+
+		prev = id;
+	}
+
+	if (!bottleneckStats.size())
+		fprintf(f, "}\n");
+	else
+		fprintf(f, "\n\t}\n}\n");
+
+	fclose(f);
+
+out:
+	pthread_mutex_unlock(&statMutex);
+}

common/network/TcpSocket.cxx

@@ -459,6 +459,13 @@ static uint8_t givecontrolCb(void *messager, const char name[])
   return msgr->netGiveControlTo(name);
 }
 
+static void bottleneckStatsCb(void *messager, char *buf, uint32_t len)
+{
+  GetAPIMessager *msgr = (GetAPIMessager *) messager;
+  msgr->netGetBottleneckStats(buf, len);
+}
+
+
 WebsocketListener::WebsocketListener(const struct sockaddr *listenaddr,
                          socklen_t listenaddrlen,
                          bool sslonly, const char *cert, const char *certkey,
@@ -548,6 +555,7 @@ WebsocketListener::WebsocketListener(const struct sockaddr *listenaddr,
   settings.adduserCb = adduserCb;
   settings.removeCb = removeCb;
   settings.givecontrolCb = givecontrolCb;
+  settings.bottleneckStatsCb = bottleneckStatsCb;
 
   pthread_t tid;
   pthread_create(&tid, NULL, start_server, NULL);

common/network/websocket.c

@@ -1074,6 +1074,21 @@ static uint8_t ownerapi(ws_ctx_t *ws_ctx, const char *in) {
 
         wserr("Passed give_control request to main thread\n");
         ret = 1;
+    } else entry("/api/get_bottleneck_stats") {
+        char statbuf[4096];
+        settings.bottleneckStatsCb(settings.messager, statbuf, 4096);
+
+        sprintf(buf, "HTTP/1.1 200 OK\r\n"
+                 "Server: KasmVNC/4.0\r\n"
+                 "Connection: close\r\n"
+                 "Content-type: text/plain\r\n"
+                 "Content-length: %lu\r\n"
+                 "\r\n", strlen(statbuf));
+        ws_send(ws_ctx, buf, strlen(buf));
+        ws_send(ws_ctx, statbuf, strlen(statbuf));
+
+        wserr("Sent bottleneck stats to API caller\n");
+        ret = 1;
     }
 
     #undef entry

common/network/websocket.h

@@ -84,6 +84,7 @@ typedef struct {
                           const uint8_t write);
     uint8_t (*removeCb)(void *messager, const char name[]);
     uint8_t (*givecontrolCb)(void *messager, const char name[]);
+    void (*bottleneckStatsCb)(void *messager, char *buf, uint32_t len);
 } settings_t;
 
 #ifdef __cplusplus

common/rfb/SMsgHandler.h

@@ -63,7 +63,7 @@ namespace rfb {
                                         const size_t* lengths,
                                         const rdr::U8* const* data);
 
-    virtual void sendStats() = 0;
+    virtual void sendStats(const bool toClient = true) = 0;
 
     virtual bool canChangeKasmSettings() const = 0;
 

common/rfb/VNCSConnectionST.cxx

@@ -17,6 +17,7 @@
  * USA.
  */
 
+#include <network/GetAPI.h>
 #include <network/TcpSocket.h>
 
 #include <rfb/ComparingUpdateTracker.h>
@@ -1459,7 +1460,7 @@ static void pruneStatList(std::list<struct timeval> &list, const struct timeval
   }
 }
 
-void VNCSConnectionST::sendStats() {
+void VNCSConnectionST::sendStats(const bool toClient) {
   char buf[1024];
   struct timeval now;
 
@@ -1498,8 +1499,12 @@ void VNCSConnectionST::sendStats() {
 
   #undef ten
 
-  vlog.info("Sending client stats:\n%s\n", buf);
-  writer()->writeStats(buf, strlen(buf));
+  if (toClient) {
+    vlog.info("Sending client stats:\n%s\n", buf);
+    writer()->writeStats(buf, strlen(buf));
+  } else if (server->apimessager) {
+    server->apimessager->mainUpdateBottleneckStats(peerEndpoint.buf, buf);
+  }
 }
 
 // setCursor() is called whenever the cursor has changed shape or pixel format.

common/rfb/VNCSConnectionST.h

@@ -164,6 +164,8 @@ namespace rfb {
     void setStatus(int status);
     int getStatus();
 
+    virtual void sendStats(const bool toClient = true);
+
   private:
     // SConnection callbacks
 
@@ -191,7 +193,6 @@ namespace rfb {
     virtual void supportsContinuousUpdates();
     virtual void supportsLEDState();
 
-    virtual void sendStats();
     virtual bool canChangeKasmSettings() const {
         return (accessRights & (AccessPtrEvents | AccessKeyEvents)) ==
                (AccessPtrEvents | AccessKeyEvents);

common/rfb/VNCServerST.cxx

@@ -997,6 +997,9 @@ void VNCServerST::writeUpdate()
     (*ci)->add_copypassed(ui.copypassed);
     (*ci)->add_changed(ui.changed);
     (*ci)->writeFramebufferUpdateOrClose();
+
+    if (apimessager)
+      (*ci)->sendStats(false);
   }
 }