11import { Injectable , Logger } from '@nestjs/common' ;
22import path from 'path' ;
3- import { writeFileSync , readFileSync , unlink , mkdirSync , existsSync } from 'fs' ;
3+ import { writeFileSync , readFileSync , unlink , mkdirSync , existsSync , createWriteStream , unlinkSync , statSync , readdir } from 'fs' ;
44import { PNG , PNGWithMetadata } from 'pngjs' ;
55
6+ import {
7+ S3Client ,
8+ PutObjectCommand ,
9+ DeleteObjectCommand ,
10+ GetObjectCommand ,
11+ } from "@aws-sdk/client-s3" ;
12+ import { Readable } from 'stream' ;
13+
614export const IMAGE_PATH = 'imageUploads/' ;
715
816@Injectable ( )
917export class StaticService {
18+
1019 private readonly logger : Logger = new Logger ( StaticService . name ) ;
1120
21+ private readonly USE_AWS_S3_BUCKET = process . env . USE_AWS_S3_BUCKET ?. trim ( ) . toLowerCase ( ) ;
22+ private readonly AWS_ACCESS_KEY_ID = process . env . AWS_ACCESS_KEY_ID ;
23+ private readonly AWS_SECRET_ACCESS_KEY = process . env . AWS_SECRET_ACCESS_KEY
24+ private readonly AWS_REGION = process . env . AWS_REGION
25+ private readonly AWS_S3_BUCKET_NAME = process . env . AWS_S3_BUCKET_NAME
26+ private readonly MAX_DISK_USAGE : number = parseInt ( process . env . MAX_TEMP_STORAGE_FOR_S3_DOWNLOAD ) * 1 * 1024 * 1024 ;
27+ private s3Client : S3Client ;
28+
29+ private readonly DELETE_INTERVAL = 60 * 60 * 1000 ; // Check for deletions every hour
30+ private deletionQueue : { filePath : string ; size : number } [ ] = [ ] ;
31+
32+ constructor ( ) {
33+ setInterval ( ( ) => this . cleanupQueuedFiles ( ) , this . DELETE_INTERVAL ) ;
34+ this . cleanUpAllFiles ( ) ;
35+ }
36+
37+ getEnvBoolean ( value : string , defaultValue : boolean = false ) : boolean {
38+ if ( value === 'true' || value === '1' ) return true ;
39+ if ( value === 'false' || value === '0' ) return false ;
40+ return defaultValue ;
41+ }
42+
43+ isAWSDefined ( ) : boolean {
44+ const areAWSVariablesValid =
45+ this . getEnvBoolean ( this . USE_AWS_S3_BUCKET ) &&
46+ this . AWS_ACCESS_KEY_ID ?. trim ( ) . length > 1 &&
47+ this . AWS_SECRET_ACCESS_KEY ?. trim ( ) . length > 1 &&
48+ this . AWS_S3_BUCKET_NAME ?. trim ( ) . length > 1 ;
49+
50+ if ( areAWSVariablesValid && ! this . s3Client ) {
51+ this . s3Client = new S3Client ( {
52+ credentials : {
53+ accessKeyId : this . AWS_ACCESS_KEY_ID ,
54+ secretAccessKey : this . AWS_SECRET_ACCESS_KEY ,
55+ } ,
56+ region : this . AWS_REGION ,
57+ } ) ;
58+ this . logger . log ( 'Using AWS S3 bucket.' ) ;
59+ }
60+ return areAWSVariablesValid ;
61+ }
62+
1263 generateNewImage ( type : 'screenshot' | 'diff' | 'baseline' ) : { imageName : string ; imagePath : string } {
1364 const imageName = `${ Date . now ( ) } ${ type } ;
1465 return {
@@ -22,37 +73,152 @@ export class StaticService {
2273 return path . resolve ( IMAGE_PATH , imageName ) ;
2374 }
2475
25- saveImage ( type : 'screenshot' | 'diff' | 'baseline' , imageBuffer : Buffer ) : string {
76+ async saveImage ( type : 'screenshot' | 'diff' | 'baseline' , imageBuffer : Buffer ) : Promise < string > {
77+ const { imageName, imagePath } = this . generateNewImage ( type ) ;
2678 try {
27- new PNG ( ) . parse ( imageBuffer ) ;
79+ if ( this . isAWSDefined ( ) ) {
80+ await this . s3Client . send (
81+ new PutObjectCommand ( {
82+ Bucket : this . AWS_S3_BUCKET_NAME ,
83+ Key : imageName ,
84+ ContentType : 'image/png' ,
85+ Body : imageBuffer ,
86+ } ) ,
87+ ) ;
88+ return imageName ;
89+ } else {
90+ new PNG ( ) . parse ( imageBuffer ) ;
91+ writeFileSync ( imagePath , imageBuffer ) ;
92+ return imageName ;
93+ }
2894 } catch ( ex ) {
29- throw new Error ( 'Cannot parse image as PNG file' ) ;
95+ throw new Error ( 'Cannot parse image as PNG file: ' + ex ) ;
3096 }
97+ }
3198
32- const { imageName, imagePath } = this . generateNewImage ( type ) ;
33- writeFileSync ( imagePath , imageBuffer ) ;
34- return imageName ;
99+ scheduleFileDeletion ( fileName : string ) : void {
100+ const filePath = this . getImagePath ( fileName ) ;
101+ const fileSize = statSync ( filePath ) . size ;
102+ this . deletionQueue . push ( { filePath, size : fileSize } ) ;
103+ this . logger . log ( `Scheduled deletion for file: ${ filePath } ${ fileSize } ) ;
35104 }
36105
37- getImage ( imageName : string ) : PNGWithMetadata {
38- if ( ! imageName ) return ;
106+ checkDiskUsageAndClean ( ) : void {
107+ const totalDiskUsage = this . getTotalDiskUsage ( ) ;
108+ if ( totalDiskUsage > this . MAX_DISK_USAGE ) {
109+ this . logger . log ( `Disk usage exceeded. Triggering cleanup.` ) ;
110+ this . cleanupQueuedFiles ( ) ;
111+ }
112+ }
113+
114+ private getTotalDiskUsage ( ) : number {
115+ return this . deletionQueue . reduce ( ( total , file ) => total + file . size , 0 ) ;
116+ }
117+
118+ private cleanUpAllFiles ( ) {
119+ readdir ( path . resolve ( IMAGE_PATH ) , ( err , files ) => {
120+ if ( err ) throw err ;
121+ for ( const file of files ) {
122+ unlink ( path . join ( path . resolve ( IMAGE_PATH ) , file ) , ( err ) => {
123+ if ( err ) throw err ;
124+ } ) ;
125+ }
126+ } ) ;
127+
128+ }
129+
130+ private cleanupQueuedFiles ( ) : void {
131+ const totalDiskUsage = this . getTotalDiskUsage ( ) ;
132+ this . logger . log ( `Cleaning up files. Total disk usage: ${ totalDiskUsage } ) ;
133+
134+ // Delete files until the total disk usage is within the limit
135+ let currentUsage = totalDiskUsage ;
136+
137+ if ( currentUsage > this . MAX_DISK_USAGE ) {
138+
139+ // Sort files by the earliest (oldest) first for deletion
140+ this . deletionQueue . sort ( ( a , b ) => statSync ( a . filePath ) . birthtimeMs - statSync ( b . filePath ) . birthtimeMs ) ;
141+
142+ // Reduce the size by half
143+ while ( currentUsage > ( this . MAX_DISK_USAGE / 2 ) && this . deletionQueue . length > 0 ) {
144+ const fileToDelete = this . deletionQueue . shift ( ) ;
145+ if ( fileToDelete ) {
146+ try {
147+ unlinkSync ( fileToDelete . filePath ) ;
148+ currentUsage -= fileToDelete . size ;
149+ this . logger . log ( `Deleted file: ${ fileToDelete . filePath } ) ;
150+ } catch ( error ) {
151+ this . logger . log ( `Failed to delete file: ${ fileToDelete . filePath } ${ error . message } ) ;
152+ }
153+ }
154+ }
155+ }
156+ }
157+
158+ doesFileExist ( fileName : string ) {
159+ return existsSync ( this . getImagePath ( fileName ) ) ;
160+ }
161+
162+ async getImage ( fileName : string ) : Promise < PNGWithMetadata > {
163+ if ( ! fileName ) return null ;
39164 try {
40- return PNG . sync . read ( readFileSync ( this . getImagePath ( imageName ) ) ) ;
165+ if ( ! this . doesFileExist ( fileName ) && this . isAWSDefined ( ) ) {
166+ const localFileStream = await this . saveFileToServerFromS3 ( fileName ) ;
167+ await new Promise < void > ( ( resolve , reject ) => {
168+ localFileStream . on ( 'finish' , ( ) => {
169+ this . scheduleFileDeletion ( fileName ) ;
170+ this . checkDiskUsageAndClean ( ) ;
171+ resolve ( ) ;
172+ } ) ;
173+ localFileStream . on ( 'error' , ( error ) => {
174+ this . logger . error ( 'Error writing file:' , error ) ;
175+ reject ( error ) ;
176+ } ) ;
177+ } ) ;
178+ }
179+ return PNG . sync . read ( readFileSync ( this . getImagePath ( fileName ) ) ) ;
41180 } catch ( ex ) {
42- this . logger . error ( `Cannot get image: ${ imageName } ${ ex } ) ;
181+ this . logger . error ( `Error from read : Cannot get image: ${ fileName } ${ ex } ) ;
182+ }
183+ }
184+
185+ async saveFileToServerFromS3 ( fileName : string ) {
186+ if ( this . isAWSDefined ( ) ) {
187+ const command = new GetObjectCommand ( { Bucket : this . AWS_S3_BUCKET_NAME , Key : fileName } ) ;
188+ const s3Response = await this . s3Client . send ( command ) ;
189+ const fileStream = s3Response . Body as Readable ;
190+ const localFileStream = createWriteStream ( this . getImagePath ( fileName ) ) ;
191+ fileStream . pipe ( localFileStream ) ;
192+ this . logger . log ( `File saved locally at ${ this . getImagePath ( fileName ) } ) ;
193+ return localFileStream ;
194+ } else {
195+ throw Error ( 'Error connecting to AWS' ) ;
43196 }
44197 }
45198
46199 async deleteImage ( imageName : string ) : Promise < boolean > {
47- if ( ! imageName ) return ;
48- return new Promise ( ( resolvePromise ) => {
49- unlink ( this . getImagePath ( imageName ) , ( err ) => {
50- if ( err ) {
51- this . logger . error ( err ) ;
52- }
53- resolvePromise ( true ) ;
200+ if ( ! imageName ) return false ;
201+ if ( this . isAWSDefined ( ) ) {
202+ try {
203+ await this . s3Client . send (
204+ new DeleteObjectCommand ( { Bucket : this . AWS_S3_BUCKET_NAME , Key : imageName } ) ,
205+ ) ;
206+ return true ;
207+ } catch ( error ) {
208+ this . logger . log ( `Failed to delete image ${ imageName } , error ) ;
209+ return false ; // Return `false` if an error occurs.
210+ }
211+ } else {
212+ return new Promise ( ( resolvePromise ) => {
213+ unlink ( this . getImagePath ( imageName ) , ( err ) => {
214+ if ( err ) {
215+ this . logger . error ( err ) ;
216+ resolvePromise ( false ) ;
217+ }
218+ resolvePromise ( true ) ;
219+ } ) ;
54220 } ) ;
55- } ) ;
221+ }
56222 }
57223
58224 private ensureDirectoryExistence ( dir : string ) {
0 commit comments