From a4e310a7f8a1b82fef76f4040bc1ec1fabcbffbb Mon Sep 17 00:00:00 2001
From: allburov <allburov@gmail.com>
Date: Thu, 30 Nov 2023 14:36:27 +0700
Subject: [PATCH] [core] Get screenshot return only buffer

---
 src/api/screenshot.controller.ts | 14 +++++++-------
 src/core/abc/session.abc.ts      |  2 +-
 src/core/session.noweb.core.ts   |  2 +-
 src/core/session.venom.core.ts   |  2 +-
 src/core/session.webjs.core.ts   |  7 +++++--
 5 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/src/api/screenshot.controller.ts b/src/api/screenshot.controller.ts
index c4240639..6a625909 100644
--- a/src/api/screenshot.controller.ts
+++ b/src/api/screenshot.controller.ts
@@ -1,4 +1,4 @@
-import { Controller, Get, Query, Res } from '@nestjs/common';
+import { Controller, Get, Query, Res, StreamableFile } from '@nestjs/common';
 import { ApiSecurity, ApiTags } from '@nestjs/swagger';
 import { Response } from 'express';
 import { Readable } from 'stream';
@@ -13,17 +13,17 @@ export class ScreenshotController {
   constructor(private manager: SessionManager) {}
 
   @Get('/screenshot')
-  async screenshot(@Res() res: Response, @Query() sessionQuery: SessionQuery) {
+  async screenshot(
+    @Res({ passthrough: true }) res: Response,
+    @Query() sessionQuery: SessionQuery,
+  ) {
     const whatsappService = this.manager.getSession(sessionQuery.session);
     const buffer = await whatsappService.getScreenshot();
-    const stream = new Readable();
-    stream.push(buffer);
-    stream.push(null);
-
+    const file = new StreamableFile(buffer);
     res.set({
       'Content-Type': 'image/png',
       'Content-Length': buffer.length,
     });
-    stream.pipe(res);
+    return file;
   }
 }
diff --git a/src/core/abc/session.abc.ts b/src/core/abc/session.abc.ts
index 6f2b515f..6aa5b091 100644
--- a/src/core/abc/session.abc.ts
+++ b/src/core/abc/session.abc.ts
@@ -206,7 +206,7 @@ export abstract class WhatsappSession {
     throw new NotImplementedByEngineError();
   }
 
-  abstract getScreenshot(): Promise<Buffer | string>;
+  abstract getScreenshot(): Promise<Buffer>;
 
   public getSessionMeInfo(): Promise<MeInfo | null> {
     throw new NotImplementedByEngineError();
diff --git a/src/core/session.noweb.core.ts b/src/core/session.noweb.core.ts
index 98448bff..757f9308 100644
--- a/src/core/session.noweb.core.ts
+++ b/src/core/session.noweb.core.ts
@@ -288,7 +288,7 @@ export class WhatsappSessionNoWebCore extends WhatsappSession {
     return { code: codeRepr };
   }
 
-  async getScreenshot(): Promise<Buffer | string> {
+  async getScreenshot(): Promise<Buffer> {
     if (this.status === WAHASessionStatus.STARTING) {
       throw new UnprocessableEntityException(
         `The session is starting, please try again after few seconds`,
diff --git a/src/core/session.venom.core.ts b/src/core/session.venom.core.ts
index 7bbbdc79..31b05ef5 100644
--- a/src/core/session.venom.core.ts
+++ b/src/core/session.venom.core.ts
@@ -126,7 +126,7 @@ export class WhatsappSessionVenomCore extends WhatsappSession {
   /**
    * START - Methods for API
    */
-  getScreenshot(): Promise<Buffer | string> {
+  getScreenshot(): Promise<Buffer> {
     if (this.status === WAHASessionStatus.STARTING) {
       throw new UnprocessableEntityException(
         `The session is starting, please try again after few seconds`,
diff --git a/src/core/session.webjs.core.ts b/src/core/session.webjs.core.ts
index aebce269..74e298f0 100644
--- a/src/core/session.webjs.core.ts
+++ b/src/core/session.webjs.core.ts
@@ -177,13 +177,16 @@ export class WhatsappSessionWebJSCore extends WhatsappSession {
     return Promise.resolve(this.qr.get());
   }
 
-  async getScreenshot(): Promise<Buffer | string> {
+  async getScreenshot(): Promise<Buffer> {
     if (this.status === WAHASessionStatus.FAILED) {
       throw new UnprocessableEntityException(
         `The session under FAILED status. Please try to restart it.`,
       );
     }
-    return await this.whatsapp.pupPage.screenshot();
+    const screenshot = await this.whatsapp.pupPage.screenshot({
+      encoding: 'binary',
+    });
+    return screenshot as Buffer;
   }
 
   async checkNumberStatus(