2021-01-03 04:40:44 +00:00
// code originally provided by tzlil
2020-11-17 14:52:12 +00:00
2020-12-03 16:30:33 +00:00
require ( "dotenv" ) . config ( ) ;
2020-11-17 14:52:12 +00:00
const os = require ( "os" ) ;
2021-01-10 01:50:29 +00:00
const { run } = require ( "../utils/image-runner.js" ) ;
2020-11-17 14:52:12 +00:00
const net = require ( "net" ) ;
2021-01-18 20:11:28 +00:00
const http = require ( "http" ) ;
2020-08-31 22:15:34 +00:00
2020-12-26 02:27:45 +00:00
const start = process . hrtime ( ) ;
2021-01-08 18:08:10 +00:00
const log = ( msg , jobNum ) => {
console . log ( ` [ ${ process . hrtime ( start ) [ 1 ] / 1000000 } ${ jobNum !== undefined ? ` : ${ jobNum } ` : "" } ] \t ${ msg } ` ) ;
2020-11-17 14:52:12 +00:00
} ;
2020-11-05 21:40:18 +00:00
2020-11-17 14:52:12 +00:00
const jobs = { } ;
// Should look like UUID : { addr : "someaddr", port: someport msg: "request" }
const queue = [ ] ;
// Array of UUIDs
2020-08-31 22:15:34 +00:00
2021-01-06 22:10:31 +00:00
const { v4 : uuidv4 } = require ( "uuid" ) ;
const MAX _JOBS = process . env . JOBS !== "" && process . env . JOBS !== undefined ? parseInt ( process . env . JOBS ) : os . cpus ( ) . length * 4 ; // Completely arbitrary, should usually be some multiple of your amount of cores
let jobAmount = 0 ;
2021-01-18 20:11:28 +00:00
const acceptJob = async ( uuid , sock ) => {
2021-01-06 22:10:31 +00:00
jobAmount ++ ;
queue . shift ( ) ;
try {
await runJob ( {
2020-11-17 14:52:12 +00:00
uuid : uuid ,
msg : jobs [ uuid ] . msg ,
2021-01-08 18:08:10 +00:00
num : jobs [ uuid ] . num
2021-01-18 20:11:28 +00:00
} , sock ) ;
2021-01-06 22:10:31 +00:00
jobAmount -- ;
if ( queue . length > 0 ) {
2021-02-20 03:01:41 +00:00
acceptJob ( queue [ 0 ] , sock ) ;
2021-01-06 22:10:31 +00:00
}
log ( ` Job ${ uuid } has finished ` ) ;
} catch ( err ) {
console . error ( ` Error on job ${ uuid } : ` , err ) ;
jobAmount -- ;
if ( queue . length > 0 ) {
2021-02-20 03:01:41 +00:00
acceptJob ( queue [ 0 ] , sock ) ;
2021-01-06 22:10:31 +00:00
}
delete jobs [ uuid ] ;
2021-01-18 20:11:28 +00:00
sock . write ( Buffer . concat ( [ Buffer . from ( [ 0x2 ] ) , Buffer . from ( uuid ) , Buffer . from ( err . toString ( ) ) ] ) ) ;
2021-01-06 22:10:31 +00:00
}
} ;
2020-11-17 14:52:12 +00:00
2021-01-18 20:11:28 +00:00
const httpServer = http . createServer ( ( req , res ) => {
if ( req . method !== "GET" ) {
res . statusCode = 405 ;
return res . end ( "405 Method Not Allowed" ) ;
}
const reqUrl = new URL ( req . url , ` http:// ${ req . headers . host } ` ) ;
if ( reqUrl . pathname === "/status" ) {
log ( ` Sending server status to ${ req . socket . remoteAddress } : ${ req . socket . remotePort } via HTTP ` ) ;
return res . end ( Buffer . from ( ( MAX _JOBS - jobAmount ) . toString ( ) ) ) ;
} else if ( reqUrl . pathname === "/image" ) {
if ( ! reqUrl . searchParams . has ( "id" ) ) {
res . statusCode = 400 ;
return res . end ( "400 Bad Request" ) ;
2020-11-17 14:52:12 +00:00
}
2021-01-18 20:11:28 +00:00
const id = reqUrl . searchParams . get ( "id" ) ;
if ( ! jobs [ id ] ) {
res . statusCode = 410 ;
return res . end ( "410 Gone" ) ;
}
log ( ` Sending image data for job ${ id } to ${ req . socket . remoteAddress } : ${ req . socket . remotePort } via HTTP ` ) ;
res . setHeader ( "ext" , jobs [ id ] . ext ) ;
return res . end ( jobs [ id ] . data , ( err ) => {
if ( err ) console . error ( err ) ;
delete jobs [ id ] ;
} ) ;
2021-01-06 22:10:31 +00:00
} else {
2021-01-18 20:11:28 +00:00
res . statusCode = 404 ;
return res . end ( "404 Not Found" ) ;
2021-01-06 22:10:31 +00:00
}
} ) ;
2021-01-18 20:11:28 +00:00
httpServer . on ( "error" , ( e ) => {
console . error ( "An HTTP error occurred: " , e ) ;
} ) ;
httpServer . listen ( 8081 , ( ) => {
log ( "HTTP listening on port 8081" ) ;
} ) ;
const server = net . createServer ( ( sock ) => { // Create a TCP socket/server to listen to requests
log ( ` TCP client ${ sock . remoteAddress } : ${ sock . remotePort } has connected ` ) ;
sock . on ( "error" , ( e ) => {
console . error ( e ) ;
} ) ;
sock . on ( "data" , ( msg ) => {
const opcode = msg . readUint8 ( 0 ) ;
const req = msg . toString ( ) . slice ( 1 , msg . length ) ;
console . log ( req ) ;
// 0x00 == Cancel job
// 0x01 == Queue job
if ( opcode == 0x00 ) {
delete queue [ queue . indexOf ( req ) - 1 ] ;
delete jobs [ req ] ;
} else if ( opcode == 0x01 ) {
const length = parseInt ( req . slice ( 0 , 1 ) ) ;
const num = req . slice ( 1 , length + 1 ) ;
const obj = req . slice ( length + 1 ) ;
const job = { addr : sock . remoteAddress , port : sock . remotePort , msg : obj , num : jobAmount } ;
const uuid = uuidv4 ( ) ;
jobs [ uuid ] = job ;
queue . push ( uuid ) ;
const newBuffer = Buffer . concat ( [ Buffer . from ( [ 0x00 ] ) , Buffer . from ( uuid ) , Buffer . from ( num ) ] ) ;
sock . write ( newBuffer ) ;
if ( jobAmount < MAX _JOBS ) {
log ( ` Got TCP request for job ${ job . msg } with id ${ uuid } ` , job . num ) ;
acceptJob ( uuid , sock ) ;
} else {
log ( ` Got TCP request for job ${ job . msg } with id ${ uuid } , queued in position ${ queue . indexOf ( uuid ) } ` , job . num ) ;
}
} else {
log ( "Could not parse TCP message" ) ;
}
} ) ;
sock . on ( "end" , ( ) => {
log ( ` TCP client ${ sock . remoteAddress } : ${ sock . remotePort } has disconnected ` ) ;
} ) ;
} ) ;
server . on ( "error" , ( e ) => {
console . error ( "A TCP error occurred: " , e ) ;
2021-01-06 22:10:31 +00:00
} ) ;
2021-01-18 20:11:28 +00:00
server . listen ( 8080 , ( ) => {
log ( "TCP listening on port 8080" ) ;
} ) ;
2021-01-06 22:10:31 +00:00
2021-01-18 20:11:28 +00:00
const runJob = ( job , sock ) => {
2021-03-16 18:24:21 +00:00
return new Promise ( ( resolve , reject ) => {
2021-01-08 18:08:10 +00:00
log ( ` Job ${ job . uuid } starting... ` , job . num ) ;
const object = JSON . parse ( job . msg ) ;
2021-01-10 01:50:29 +00:00
// If the image has a path, it must also have a type
if ( object . path && ! object . type ) {
reject ( new TypeError ( "Unknown image type" ) ) ;
2021-01-08 18:08:10 +00:00
}
2021-01-06 22:10:31 +00:00
2021-01-08 18:08:10 +00:00
log ( ` Job ${ job . uuid } started ` , job . num ) ;
2021-03-16 18:24:21 +00:00
run ( object ) . then ( ( data ) => {
2021-01-18 20:11:28 +00:00
log ( ` Sending result of job ${ job . uuid } back to the bot ` , job . num ) ;
2021-03-16 18:24:21 +00:00
jobs [ job . uuid ] . data = data . buffer ;
jobs [ job . uuid ] . ext = data . fileExtension ;
2021-01-18 20:11:28 +00:00
sock . write ( Buffer . concat ( [ Buffer . from ( [ 0x1 ] ) , Buffer . from ( job . uuid ) ] ) , ( e ) => {
if ( e ) return reject ( e ) ;
return resolve ( ) ;
2020-11-17 14:52:12 +00:00
} ) ;
2021-03-18 14:29:03 +00:00
return ;
2021-03-16 18:24:21 +00:00
} ) . catch ( e => {
2021-01-18 20:11:28 +00:00
reject ( e ) ;
2021-03-16 18:24:21 +00:00
} ) ;
2020-11-05 21:40:18 +00:00
} ) ;
2021-01-06 22:10:31 +00:00
} ;