Add support for owner screenshot HTTP GET API

common/network/CMakeLists.txt

@@ -1,6 +1,7 @@
 include_directories(${CMAKE_SOURCE_DIR}/common ${CMAKE_SOURCE_DIR}/unix/kasmvncpasswd)
 
 set(NETWORK_SOURCES
+  GetAPIMessager.cxx
   Socket.cxx
   TcpSocket.cxx
   websocket.c

common/network/GetAPI.h

@@ -0,0 +1,54 @@
+/* Copyright (C) 2021 Kasm
+ *
+ * This is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; either version 2 of the License, or
+ * (at your option) any later version.
+ *
+ * This software is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this software; if not, write to the Free Software
+ * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307,
+ * USA.
+ */
+
+#ifndef __NETWORK_GET_API_H__
+#define __NETWORK_GET_API_H__
+
+#include <pthread.h>
+#include <rfb/PixelBuffer.h>
+#include <rfb/PixelFormat.h>
+#include <stdint.h>
+#include <vector>
+
+namespace network {
+
+  class GetAPIMessager {
+  public:
+    GetAPIMessager();
+
+    // from main thread
+    void mainUpdateScreen(rfb::PixelBuffer *pb);
+
+    // from network threads
+    uint8_t *netGetScreenshot(uint16_t w, uint16_t h,
+                              const uint8_t q, const bool dedup,
+                              uint32_t &len, uint8_t *staging);
+  private:
+    pthread_mutex_t screenMutex;
+    rfb::ManagedPixelBuffer screenPb;
+    uint16_t screenW, screenH;
+    uint64_t screenHash;
+
+    std::vector<uint8_t> cachedJpeg;
+    uint16_t cachedW, cachedH;
+    uint8_t cachedQ;
+  };
+
+}
+
+#endif // __NETWORK_GET_API_H__

common/network/GetAPIMessager.cxx

@@ -0,0 +1,183 @@
+/* Copyright (C) 2021 Kasm
+ *
+ * This is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; either version 2 of the License, or
+ * (at your option) any later version.
+ *
+ * This software is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this software; if not, write to the Free Software
+ * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307,
+ * USA.
+ */
+
+#define __STDC_FORMAT_MACROS
+
+#include <inttypes.h>
+#include <network/GetAPI.h>
+#include <rfb/ConnParams.h>
+#include <rfb/LogWriter.h>
+#include <rfb/JpegCompressor.h>
+#include <rfb/xxhash.h>
+#include <stdio.h>
+
+using namespace network;
+using namespace rfb;
+
+static LogWriter vlog("GetAPIMessager");
+
+PixelBuffer *progressiveBilinearScale(const PixelBuffer *pb,
+					const uint16_t tgtw, const uint16_t tgth,
+					const float tgtdiff);
+
+struct TightJPEGConfiguration {
+    int quality;
+    int subsampling;
+};
+
+static const struct TightJPEGConfiguration conf[10] = {
+  {  15, subsample4X }, // 0
+  {  29, subsample4X }, // 1
+  {  41, subsample4X }, // 2
+  {  42, subsample2X }, // 3
+  {  62, subsample2X }, // 4
+  {  77, subsample2X }, // 5
+  {  79, subsampleNone }, // 6
+  {  86, subsampleNone }, // 7
+  {  92, subsampleNone }, // 8
+  { 100, subsampleNone }  // 9
+};
+
+GetAPIMessager::GetAPIMessager(): screenW(0), screenH(0), screenHash(0),
+					cachedW(0), cachedH(0), cachedQ(0) {
+
+	pthread_mutex_init(&screenMutex, NULL);
+}
+
+// from main thread
+void GetAPIMessager::mainUpdateScreen(rfb::PixelBuffer *pb) {
+	if (pthread_mutex_trylock(&screenMutex))
+		return;
+
+	int stride;
+	const rdr::U8 * const buf = pb->getBuffer(pb->getRect(), &stride);
+
+	if (pb->width() != screenW || pb->height() != screenH) {
+		screenHash = 0;
+		screenW = pb->width();
+		screenH = pb->height();
+		screenPb.setPF(pb->getPF());
+		screenPb.setSize(screenW, screenH);
+
+		cachedW = cachedH = cachedQ = 0;
+		cachedJpeg.clear();
+	}
+
+	const uint64_t newHash = XXH64(buf, pb->area() * 4, 0);
+	if (newHash != screenHash) {
+		cachedW = cachedH = cachedQ = 0;
+		cachedJpeg.clear();
+
+		screenHash = newHash;
+		rdr::U8 *rw = screenPb.getBufferRW(screenPb.getRect(), &stride);
+		memcpy(rw, buf, screenW * screenH * 4);
+		screenPb.commitBufferRW(screenPb.getRect());
+	}
+
+	pthread_mutex_unlock(&screenMutex);
+}
+
+// from network threads
+uint8_t *GetAPIMessager::netGetScreenshot(uint16_t w, uint16_t h,
+	const uint8_t q, const bool dedup,
+	uint32_t &len, uint8_t *staging) {
+
+	uint8_t *ret = NULL;
+	len = 0;
+
+	if (w > screenW)
+		w = screenW;
+	if (h > screenH)
+		h = screenH;
+
+	if (!w || !h || q > 9 || !staging)
+		return NULL;
+
+	if (pthread_mutex_lock(&screenMutex))
+		return NULL;
+
+	if (w == cachedW && h == cachedH && q == cachedQ) {
+		if (dedup) {
+			// Return the hash of the unchanged image
+			sprintf((char *) staging, "%" PRIx64, screenHash);
+			ret = staging;
+			len = 16;
+		} else {
+			// Return the cached image
+			len = cachedJpeg.size();
+			ret = staging;
+			memcpy(ret, &cachedJpeg[0], len);
+
+			vlog.info("Returning cached screenshot");
+		}
+	} else {
+		// Encode the new JPEG, cache it
+		JpegCompressor jc;
+		int quality, subsampling;
+
+		quality = conf[q].quality;
+		subsampling = conf[q].subsampling;
+
+		jc.clear();
+		int stride;
+
+		if (w != screenW || h != screenH) {
+			float xdiff = w / (float) screenW;
+			float ydiff = h / (float) screenH;
+			const float diff = xdiff < ydiff ? xdiff : ydiff;
+
+			const uint16_t neww = screenW * diff;
+			const uint16_t newh = screenH * diff;
+
+			const PixelBuffer *scaled = progressiveBilinearScale(&screenPb, neww, newh, diff);
+			const rdr::U8 * const buf = scaled->getBuffer(scaled->getRect(), &stride);
+
+			jc.compress(buf, stride, scaled->getRect(),
+					scaled->getPF(), quality, subsampling);
+
+			cachedJpeg.resize(jc.length());
+			memcpy(&cachedJpeg[0], jc.data(), jc.length());
+
+			delete scaled;
+
+			vlog.info("Returning scaled screenshot");
+		} else {
+			const rdr::U8 * const buf = screenPb.getBuffer(screenPb.getRect(), &stride);
+
+			jc.compress(buf, stride, screenPb.getRect(),
+					screenPb.getPF(), quality, subsampling);
+
+			cachedJpeg.resize(jc.length());
+			memcpy(&cachedJpeg[0], jc.data(), jc.length());
+
+			vlog.info("Returning normal screenshot");
+		}
+
+		cachedQ = q;
+		cachedW = w;
+		cachedH = h;
+
+		len = cachedJpeg.size();
+		ret = staging;
+		memcpy(ret, &cachedJpeg[0], len);
+	}
+
+	pthread_mutex_unlock(&screenMutex);
+
+	return ret;
+}

common/network/Socket.h

@@ -74,6 +74,8 @@ namespace network {
     virtual ~ConnectionFilter() {}
   };
 
+  class GetAPIMessager;
+
   class SocketListener {
   public:
     SocketListener(int fd);
@@ -93,6 +95,8 @@ namespace network {
     void setFilter(ConnectionFilter* f) {filter = f;}
     int getFd() {return fd;}
 
+    virtual GetAPIMessager *getMessager() { return NULL; }
+
   protected:
     SocketListener();
 

common/network/TcpSocket.cxx

@@ -42,6 +42,7 @@
 #include <wordexp.h>
 #include "websocket.h"
 
+#include <network/GetAPI.h>
 #include <network/TcpSocket.h>
 #include <rfb/LogWriter.h>
 #include <rfb/Configuration.h>
@@ -431,6 +432,14 @@ int TcpListener::getMyPort() {
 
 extern settings_t settings;
 
+static uint8_t *screenshotCb(void *messager, uint16_t w, uint16_t h, const uint8_t q,
+                             const uint8_t dedup,
+                             uint32_t *len, uint8_t *staging)
+{
+  GetAPIMessager *msgr = (GetAPIMessager *) messager;
+  return msgr->netGetScreenshot(w, h, q, dedup, *len, staging);
+}
+
 WebsocketListener::WebsocketListener(const struct sockaddr *listenaddr,
                          socklen_t listenaddrlen,
                          bool sslonly, const char *cert, const char *certkey,
@@ -515,6 +524,9 @@ WebsocketListener::WebsocketListener(const struct sockaddr *listenaddr,
 
   settings.listen_sock = sock;
 
+  settings.messager = messager = new GetAPIMessager;
+  settings.screenshotCb = screenshotCb;
+
   pthread_t tid;
   pthread_create(&tid, NULL, start_server, NULL);
 }

common/network/TcpSocket.h

@@ -100,8 +100,12 @@ namespace network {
 
     int internalSocket;
 
+    virtual GetAPIMessager *getMessager() { return messager; }
+
   protected:
     virtual Socket* createSocket(int fd);
+  private:
+    GetAPIMessager *messager;
   };
 
   void createLocalTcpListeners(std::list<SocketListener*> *listeners,

common/network/websocket.c

@@ -9,6 +9,7 @@
  */
 #define _GNU_SOURCE
 
+#include <ctype.h>
 #include <pthread.h>
 #include <unistd.h>
 #include <crypt.h>
@@ -83,6 +84,32 @@ int resolve_host(struct in_addr *sin_addr, const char *hostname)
     return 0;
 }
 
+static const char *parse_get(const char * const in, const char * const opt, unsigned *len) {
+	const char *start = in;
+	const char *end = strchrnul(start, '&');
+	const unsigned optlen = strlen(opt);
+	*len = 0;
+
+	while (1) {
+		if (!strncmp(start, opt, optlen)) {
+			const char *arg = strchr(start, '=');
+			if (!arg)
+				return "";
+			arg++;
+			*len = end - arg;
+			return arg;
+		}
+
+		if (!*end)
+			break;
+
+		end++;
+		start = end;
+		end = strchrnul(start, '&');
+	}
+
+	return "";
+}
 
 /*
  * SSL Wrapper Code
@@ -814,6 +841,116 @@ nope:
     ws_send(ws_ctx, buf, strlen(buf));
 }
 
+static uint8_t ownerapi(ws_ctx_t *ws_ctx, const char *in) {
+    char buf[4096], path[4096], fullpath[4096], args[4096] = "";
+    uint8_t ret = 0; // 0 = continue checking
+
+    if (strncmp(in, "GET ", 4)) {
+        wserr("non-GET request, rejecting\n");
+        return 0;
+    }
+    in += 4;
+    const char *end = strchr(in, ' ');
+    unsigned len = end - in;
+
+    if (len < 1 || len > 1024 || strstr(in, "../")) {
+        wserr("Request too long (%u) or attempted dir traversal attack, rejecting\n", len);
+        return 0;
+    }
+
+    end = memchr(in, '?', len);
+    if (end) {
+        len = end - in;
+        end++;
+
+        const char *argend = strchr(end, ' ');
+        strncpy(args, end, argend - end);
+        args[argend - end] = '\0';
+    }
+
+    memcpy(path, in, len);
+    path[len] = '\0';
+
+    wserr("Requested owner api '%s' with args '%s'\n", path, args);
+
+    #define entry(x) if (!strcmp(path, x))
+
+    const char *param;
+
+    entry("/api/get_screenshot") {
+        uint8_t q = 7, dedup = 0;
+        uint16_t w = 4096, h = 4096;
+
+        param = parse_get(args, "width", &len);
+        if (len && isdigit(param[0]))
+            w = atoi(param);
+
+        param = parse_get(args, "height", &len);
+        if (len && isdigit(param[0]))
+            h = atoi(param);
+
+        param = parse_get(args, "quality", &len);
+        if (len && isdigit(param[0]))
+            q = atoi(param);
+
+        param = parse_get(args, "deduplicate", &len);
+        if (len && isalpha(param[0])) {
+            if (!strncmp(param, "true", len))
+                dedup = 1;
+        }
+
+        uint8_t *staging = malloc(1024 * 1024 * 8);
+
+        settings.screenshotCb(settings.messager, w, h, q, dedup, &len, staging);
+
+        if (len == 16) {
+            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: %u\r\n"
+                     "\r\n", len);
+            ws_send(ws_ctx, buf, strlen(buf));
+            ws_send(ws_ctx, staging, len);
+
+            wserr("Screenshot hadn't changed and dedup was requested, sent hash\n");
+            ret = 1;
+        } else if (len) {
+            sprintf(buf, "HTTP/1.1 200 OK\r\n"
+                     "Server: KasmVNC/4.0\r\n"
+                     "Connection: close\r\n"
+                     "Content-type: image/jpeg\r\n"
+                     "Content-length: %u\r\n"
+                     "\r\n", len);
+            ws_send(ws_ctx, buf, strlen(buf));
+            ws_send(ws_ctx, staging, len);
+
+            wserr("Sent screenshot %u bytes\n", len);
+            ret = 1;
+        }
+
+        free(staging);
+
+        if (!len) {
+            wserr("Invalid params to screenshot\n");
+            goto nope;
+        }
+    }
+
+    #undef entry
+
+    return ret;
+nope:
+    sprintf(buf, "HTTP/1.1 400 Bad Request\r\n"
+                 "Server: KasmVNC/4.0\r\n"
+                 "Connection: close\r\n"
+                 "Content-type: text/plain\r\n"
+                 "\r\n"
+                 "400 Bad Request");
+    ws_send(ws_ctx, buf, strlen(buf));
+    return 1;
+}
+
 ws_ctx_t *do_handshake(int sock) {
     char handshake[4096], response[4096], sha1[29], trailer[17];
     char *scheme, *pre;
@@ -883,6 +1020,7 @@ ws_ctx_t *do_handshake(int sock) {
     }
 
     const char *colon;
+    unsigned char owner = 0;
     if ((colon = strchr(settings.basicauth, ':'))) {
         const char *hdr = strstr(handshake, "Authorization: Basic ");
         if (!hdr) {
@@ -938,6 +1076,9 @@ ws_ctx_t *do_handshake(int sock) {
                             snprintf(authbuf, 4096, "%s:%s", set->entries[i].user,
                                      set->entries[i].password);
                             authbuf[4095] = '\0';
+
+                            if (set->entries[i].owner)
+                                owner = 1;
                             break;
                         }
                     }
@@ -978,9 +1119,14 @@ ws_ctx_t *do_handshake(int sock) {
     if (!parse_handshake(ws_ctx, handshake)) {
         handler_emsg("Invalid WS request, maybe a HTTP one\n");
 
+        if (strstr(handshake, "/api/") && owner)
+            if (ownerapi(ws_ctx, handshake))
+                goto done;
+
         if (settings.httpdir && settings.httpdir[0])
             servefile(ws_ctx, handshake);
 
+done:
         free_ws_ctx(ws_ctx);
         return NULL;
     }

common/network/websocket.h

@@ -1,4 +1,5 @@
 #include <openssl/ssl.h>
+#include <stdint.h>
 
 #define BUFSIZE 65536
 #define DBUFSIZE (BUFSIZE * 3) / 4 - 20
@@ -74,6 +75,11 @@ typedef struct {
     const char *passwdfile;
     int ssl_only;
     const char *httpdir;
+
+    void *messager;
+    uint8_t *(*screenshotCb)(void *messager, uint16_t w, uint16_t h, const uint8_t q,
+                             const uint8_t dedup,
+                             uint32_t *len, uint8_t *staging);
 } settings_t;
 
 #ifdef __cplusplus

common/rfb/EncodeManager.cxx

@@ -966,7 +966,7 @@ static PixelBuffer *bilinearScale(const PixelBuffer *pb, const uint16_t w, const
   return newpb;
 }
 
-static PixelBuffer *progressiveBilinearScale(const PixelBuffer *pb,
+PixelBuffer *progressiveBilinearScale(const PixelBuffer *pb,
                                  const uint16_t tgtw, const uint16_t tgth,
                                  const float tgtdiff)
 {

common/rfb/VNCServerST.cxx

@@ -51,6 +51,8 @@
 #include <assert.h>
 #include <stdlib.h>
 
+#include <network/GetAPI.h>
+
 #include <rfb/ComparingUpdateTracker.h>
 #include <rfb/KeyRemapper.h>
 #include <rfb/ListConnInfo.h>
@@ -91,7 +93,7 @@ VNCServerST::VNCServerST(const char* name_, SDesktop* desktop_)
     renderedCursorInvalid(false),
     queryConnectionHandler(0), keyRemapper(&KeyRemapper::defInstance),
     lastConnectionTime(0), disableclients(false),
-    frameTimer(this)
+    frameTimer(this), apimessager(NULL)
 {
   lastUserInputTime = lastDisconnectTime = time(0);
   slog.debug("creating single-threaded server %s", name.buf);
@@ -709,6 +711,9 @@ void VNCServerST::writeUpdate()
     }
   }
 
+  if (apimessager)
+    apimessager->mainUpdateScreen(pb);
+
   for (ci = clients.begin(); ci != clients.end(); ci = ci_next) {
     ci_next = ci; ci_next++;
 

common/rfb/VNCServerST.h

@@ -43,6 +43,7 @@ namespace rfb {
   class ListConnInfo;
   class PixelBuffer;
   class KeyRemapper;
+  class network::GetAPIMessager;
 
   class VNCServerST : public VNCServer,
                       public Timer::Callback,
@@ -186,6 +187,8 @@ namespace rfb {
     bool getDisable() { return disableclients;};
     void setDisable(bool disable) { disableclients = disable;};
 
+    void setAPIMessager(network::GetAPIMessager *msgr) { apimessager = msgr; }
+
   protected:
 
     friend class VNCSConnectionST;
@@ -251,6 +254,8 @@ namespace rfb {
     Timer frameTimer;
 
     int inotifyfd;
+
+    network::GetAPIMessager *apimessager;
   };
 
 };

unix/xserver/hw/vnc/XserverDesktop.cc

@@ -86,6 +86,8 @@ XserverDesktop::XserverDesktop(int screenIndex_,
        i != listeners.end();
        i++) {
     vncSetNotifyFd((*i)->getFd(), screenIndex, true, false);
+    if ((*i)->getMessager())
+      server->setAPIMessager((*i)->getMessager());
   }
 }