myNodeServer.js

/**
 * Copyright (c) 2023 MERCENARIES.AI PTE. LTD.
 * All rights reserved.
 */
//@ts-check
const VERSION = '0.0.8e';

const express = require('express');
const http = require('http');
const session = require('express-session');
const axios = require('axios');
const { spawn } = require('child_process');
const { createProxyMiddleware } = require('http-proxy-middleware');

const app = express();

const CONTAINER_HOST = '127.0.0.1';
const OMNI_URL = 'http://127.0.0.1:1688'; // URL of the OMNITOOL service
const PROXY_PORT_OMNITOOL = 4444;
const CONTAINER_PORT_OMNITOOL = 1688;
const OMNITOOL_HEALTH_PING = '/api/v1/mercenaries/ping';


const OMNITOOL_SPACE_NAME = 'omnitool_test_3';
const OMNITOOL_SPACE_OWNER = 'manu-sapiens';
const HF_SPACE_URL = `https://huggingface.co/spaces/${OMNITOOL_SPACE_OWNER}/${OMNITOOL_SPACE_NAME}`;
const HF_SPACE_DUPLICATE_URL = `${HF_SPACE_URL}?duplicate=true`;

const DELAY_OMNITOOL_SET_TO_RUNNING = 2000; // 2 seconds
const CHECK_OMNI_INTERVAL = 60000; // 1 minute

// Global variable
global.OMNITOOL_RUNNING = false;
global.OMNITOOL_READY = false;
global.ALREADY_STARTING = false;
global.LOCAL_URL = "";
global.PING_URL = "";
global.PROTOCOL = "";
global.PROXY_STARTED = false;
global.logs = [];
global.OMNITOOL_PROXY = null;
global.CONNECTED_TO_MASTER = false;
global.SPACE_OWNER = "";


console.log(`************ Omnitool Proxy Server v${VERSION} ************`);

// HTML templates and constants
const LOG_CONTAINER_HTML = `<div id="logContainer"></div>`;
const COMMON_STYLES = `
    <style>
        #logContainer {
            height: 400px;
            overflow-y: scroll;
            background-color: black;
            color: lime;
            font-family: 'Courier New', Courier, monospace;
            padding: 10px;
            white-space: pre-wrap;
            border: 1px solid #ddd;
        }
        .highlight-button {
            animation: pulseAnimation 1s infinite;
            background-color: yellow;
            color: black;
            font-weight: bold;
        }    
        .button-like-link {
            display: inline-block;
            background-color: #f0f0f0; /* Button color */
            color: #000; /* Text color */
            padding: 10px 15px;
            text-align: center;
            text-decoration: none;
            border: 1px solid #ddd;
            border-radius: 5px;
            font-size: 16px;
            cursor: pointer;
        }
        .button-like-link:hover {
            background-color: #e9e9e9;
        }                     
        @keyframes pulseAnimation {
            0% { transform: scale(1); }
            50% { transform: scale(1.1); }
            100% { transform: scale(1); }
        }
    </style>`;

async function startOmnitoolServer()
{
    if (global.ALREADY_STARTING) return;
    global.ALREADY_STARTING = true;

    console.log('Starting Omnitool Server...');
    return new Promise((resolve, reject) =>
    {
        const omnitoolStartProcess = spawn('./omnitool_start.sh');
        omnitoolStartProcess.stdout.on('data', (data) =>
        {
            global.logs.push(data.toString());
            console.log(`omnitool stdout: ${data}`);

            if (data.toString().includes(`Server has started and is ready to accept connections`))
            {
                console.log('Omnitool server started successfully');
                setOmnitoolRunning(true);
            }

        });

        omnitoolStartProcess.stderr.on('data', (data) =>
        {
            console.error(`omnitool stderr: ${data}`);
            global.logs.push(data.toString());
        });

        omnitoolStartProcess.on('close', (code) =>
        {
            const message = `Omnitool server process exited with code: ${code}`;
            console.log(message);
            global.logs.push(message);

            global.ALREADY_STARTING = false;
            if (code === 0)
            {
                //@ts-ignore
                resolve();
            } else
            {
                reject(`Omnitool server did not start properly.`);
            }
        });
    });
}

function setOmnitoolRunning(set_running)
{
    if (set_running === true)
    {
        if (global.OMNITOOL_READY === false) 
        {
            global.OMNITOOL_READY = true;
            console.log('Omnitool server is READY! ');
            setTimeout(() =>
            {
                console.log('Omnitool server is RUNNING! ');
                global.OMNITOOL_RUNNING = true;
            }, DELAY_OMNITOOL_SET_TO_RUNNING); // Delay by 2 second
        }

    }
    else
    {
        global.OMNITOOL_READY = false;
        global.OMNITOOL_RUNNING = false;
    }
}

// Function to check the status of the external service
async function checkOmnitoolStatus()
{
    console.log("Checking external service");
    try
    {
        //@ts-ignore
        const response = await axios.get('http://127.0.0.1:1688/api/v1/mercenaries/ping');
        if (response.data && response.data.ping === 'pong' && Object.keys(response.data.payload).length === 0)
        {
            setOmnitoolRunning(true);
        }
        else
        {
            setOmnitoolRunning(false);
        }
    } catch (error)
    {
        console.error('Cannot access OMNITOOL. It is probably not running...', error.message);
        setOmnitoolRunning(false);
    }
}

async function handleGetStartOmnitoolServer(req, res)
{
    console.log(`Omnitool Server:ALREADY_STARTING =  ${global.ALREADY_STARTING}`);
    if (global.ALREADY_STARTING) 
    {
        res.writeHead(200, { 'Content-Type': 'text/html' });
        res.end("Omnitool server already starting");
        return;
    }
    try
    {
        await startOmnitoolServer();
        res.writeHead(200, { 'Content-Type': 'text/html' });
        res.end("Omnitool server started successfully");
    }
    catch (error)
    {
        console.error(error);
        global.ALREADY_STARTING = false;
        res.writeHead(500, { 'Content-Type': 'text/html' });
        res.end(`Error starting Omnitool server: ${error}`);
    }
}

async function handleGetOmnitoolLogs(req, res)
{
    res.writeHead(200, { 'Content-Type': 'application/json' });
    const reply = { logs: global.logs, ready: global.OMNITOOL_READY };
    res.end(JSON.stringify(reply));
    return;
}

async function handleGetBurstIframe(req, res)
{
    res.writeHead(200, { 'Content-Type': 'application/json' });
    req.session.isVisited = true;

    const reply = { burst_iframe: true };
    res.end(JSON.stringify(reply));
    return;
}

async function proxyRequest(req, res)
{

    console.log('Proxying request');
    if (global.PROXY_STARTED) return;
    global.PROXY_STARTED = true;

    // Proxy logic...
    const options = { hostname: CONTAINER_HOST, port: CONTAINER_PORT_OMNITOOL, path: req.url, method: req.method, headers: req.headers, };
    const proxy = http.request(options, (proxyRes) =>
    {
        res.writeHead(proxyRes.statusCode, proxyRes.headers);
        proxyRes.pipe(res, { end: true });
    });
    req.pipe(proxy, { end: true });
    return;
}

function getButtonsString(buttonsHTML)
{
    // Add a button for the Hugging Face API call
    const huggingFaceButton = '<button onclick="getHuggingFaceUserName()">Get Hugging Face User Name</button>';


    // Message to display based on CONNECTED_TO_MASTER
    let connectionStatusMessage;
    if (global.CONNECTED_TO_MASTER)
    {
        connectionStatusMessage = '<div style="color: red;">You are now connected DIRECTLY to the REFERENCE Omnitool Space on Huggingface. It is HIGHLY recommended that you duplicate this space and make it private to secure your keys, recipes and outputs.</div>';
    } else
    {
        connectionStatusMessage = `<div style="color: green;">You are now connected to the space of ${global.SPACE_OWNER}. If this is NOT YOU, it is HIGHLY recommended that you duplicate this space and make it private to secure your keys, recipes and outputs.</div><div>In addition, consider subscribing to the 20GB local storage option from Huggingface to preserve your keys and work between reboot of this server.</div>`;
    }

    return `
        <html>
            <head>
                <title>Proxy Server</title>
                ${COMMON_STYLES}
            </head>
            <body>
                <h1>Proxy Server</h1>
                <p>Version: ${VERSION}</p>
                <p>Local URL: ${global.LOCAL_URL}</p>
                <p>Space Owner: ${global.SPACE_OWNER}</p>
                <p>Connected to master: ${global.CONNECTED_TO_MASTER}</p>
                <p>Omni URL: ${OMNI_URL}</p>
                <p>Omni Running: ${global.OMNITOOL_RUNNING}</p>
                <p>Omni Ready: ${global.OMNITOOL_READY}</p>
               
                ${connectionStatusMessage} <!-- Display the connection status message -->
                ${buttonsHTML}
                ${huggingFaceButton}
                ${LOG_CONTAINER_HTML}
                <script>
                    // Script to check Hugging Face login status and control 'Clone Repo' button
                
                    function startServer() {
                        document.getElementById('startServerButton').classList.remove
                        ('highlight-button');
                        document.getElementById('startServerButton').disabled = true;


                        fetch('/start-omnitool-server')
                            .then(response => response.json())
                            .then(data => startLogPolling());
  
                        
                    }

                    function exitIframe() {
                        fetch('/burst-iframe')
                            .then(response => response.json())
                            .then(data => window.open('/', '_blank'));
                    }

                    function startLogPolling() {

                        const interval = setInterval(() => {
                            fetch('/omnitool-logs')
                                .then(response => response.json())
                                .then(data => {
                                    const logContainer = document.getElementById('logContainer');
                                    logContainer.innerText = data.logs.join("");
                                    if (data.ready) {
                                        clearInterval(interval);
                                        document.getElementById('exitIframeButton').classList.add('highlight-button');


                                    }
                                    else {
                                        scrollToBottom(logContainer);
                                    }
                                });
                        }, 500);
                    }

                    function scrollToBottom(element) {
                        element.scrollTop = element.scrollHeight;
                    }
                    
                    // Function to handle the Hugging Face API call
                    function getHuggingFaceUserName() {
                        fetch('https://huggingface.co/api/whoami-v2')
                            .then(response => response.json())
                            .then(data => {
                                if(data && data.name) {
                                    alert('Hugging Face User Name: ' + data.name);
                                } else {
                                    alert('No name field in the response.');
                                }
                            })
                            .catch(error => {
                                console.error('Error fetching Hugging Face user name:', error);
                                alert('Error fetching data.');
                            });
                    }

                    startLogPolling();
                </script>
            </body>
        </html>`;
}


async function handleGetRoot(req, res)
{
    setGlobals(req);
    
    console.log(`req.session.isVisited = ${req.session.isVisited}`);
    console.log(`global.OMNITOOL_RUNNING = ${global.OMNITOOL_RUNNING}`);

    if (!req.session.isVisited || !global.OMNITOOL_RUNNING)
    {
        await checkOmnitoolStatus();
    }

    if (!global.OMNITOOL_RUNNING)
    {
        console.log('Omnitool server is not running');
        let startButtonClass = '';
        let cloneButtonClass = '';
        if (!global.ALREADY_STARTING) startButtonClass = 'highlight-button';
        const gotoButtonClass = '';

        let buttonsHTML = `
        <button id="startServerButton" class="${startButtonClass}" onclick="startServer()">Start Omnitool Server</button>
        <button id="exitIframeButton" class="${gotoButtonClass}" onclick="exitIframe()">GOTO OMNITOOL</button>
        <a href="${HF_SPACE_DUPLICATE_URL}" target="_blank" class="button-like-link" id="duplicateRepoButton">DUPLICATE Repo</a>`;

        const html = getButtonsString(buttonsHTML);
        res.writeHead(200, { 'Content-Type': 'text/html' });
        res.end(html);

        return;
    }

    if (!req.session.isVisited)
    {
        console.log('First time visitor');
        const gotoButtonClass = 'highlight-button';
        const buttonsHTML = `<button id="exitIframeButton" class="${gotoButtonClass}" onclick="exitIframe()">GOTO OMNITOOL</button>`;
        const html = getButtonsString(buttonsHTML);
        res.writeHead(200, { 'Content-Type': 'text/html' });
        res.end(html);

        return;
    }

    proxyRequest(req, res);
}

function extractSpaceOwner(url, spaceName)
{
    const regex = new RegExp(`^(?:https?:\\/\\/)?([^\\/]*?)-${spaceName}\\.hf\\.space\\/?$`);
    const match = url.match(regex);
    return match ? match[1] : null;
}

function setGlobals(req)
{
    console.log(`Setting global using req.headers['host'] = ${req.headers['host']}`);
    global.LOCAL_URL = req.headers['host'];
    if (req.protocol === "https") global.PROTOCOL = 'https'; else global.PROTOCOL = 'http';
    console.log(`global.LOCAL_URL = ${global.LOCAL_URL}\nglobal.PROTOCOL = ${global.PROTOCOL}`);

    if (!global.LOCAL_URL) throw new Error('No host header found in request' + JSON.stringify(req.headers, null, 2));

    const hostname = global.LOCAL_URL.split(':')[0];
    const newUrl = `${global.PROTOCOL}://${hostname}`;//:${PROXY_PORT_OMNITOOL}`;
    global.PING_URL = `${newUrl}${OMNITOOL_HEALTH_PING}`;
    console.log(`global.PING_URL = ${global.PING_URL}`);

    // New logic to set CONNECTED_TO_MASTER
    global.CONNECTED_TO_MASTER = false;
    global.SPACE_OWNER = extractSpaceOwner(global.LOCAL_URL, OMNITOOL_SPACE_NAME);
 console.log(`global.SPACE_OWNER = ${global.SPACE_OWNER}`);

    if (global.SPACE_OWNER === OMNITOOL_SPACE_OWNER) { global.CONNECTED_TO_MASTER = true; }
    console.log(`global.CONNECTED_TO_MASTER = ${global.CONNECTED_TO_MASTER}`);

}

async function startServer()
{
    // Start server
    http.createServer(app).listen(PROXY_PORT_OMNITOOL, () =>
    {
        console.log(`Server running on port ${PROXY_PORT_OMNITOOL}`);
    });

}

function omnitoolProxyMiddleware(req, res, next)
{
    if (global.OMNITOOL_RUNNING && req.session.isVisited)
    {
        // Proxy all requests to OMNITOOL when conditions are met
        return createProxyMiddleware({
            target: OMNI_URL,
            changeOrigin: true,
            ws: true // if you need WebSocket support
        })(req, res, next);
    } else
    {
        // Continue with normal processing for other cases
        next();
    }
}


async function main(app)
{
    await checkOmnitoolStatus();

    // Configure session middleware
    app.use(session({
        secret: 'your-secret-key',
        resave: false,
        saveUninitialized: true
    }));
    try
    {
        await startServer();
    }
    catch (error)
    {
        console.error(`There was an error starting the server: ${error}`);
        return;
    }


    app.use(omnitoolProxyMiddleware);

    app.get('/favicon.ico', (req, res) => res.status(204));
    app.get('/start-omnitool-server', (req, res) => handleGetStartOmnitoolServer(req, res));
    app.get('/omnitool-logs', (req, res) => handleGetOmnitoolLogs(req, res));
    app.get('/burst-iframe', (req, res) => handleGetBurstIframe(req, res));
    app.get('/', (req, res) => handleGetRoot(req, res));

    setInterval(async () =>
    {
        await checkOmnitoolStatus();
    }, CHECK_OMNI_INTERVAL);

}

// Define the proxy middleware outside the function
global.OMNITOOL_PROXY = createProxyMiddleware({
    target: OMNI_URL,
    changeOrigin: true,
    ws: true // if you need WebSocket support
});

function omnitoolProxyMiddleware(req, res, next)
{
    if (global.OMNITOOL_RUNNING && req.session.isVisited)
    {
        // Use the predefined proxy middleware
        return global.OMNITOOL_PROXY(req, res, next);
    } else
    {
        // Continue with normal processing for other cases
        next();
    }
}
main(app);