dashcaddy/dashcaddy-api/routes/exec.js

const { WebSocketServer } = require('ws');
const Docker = require('dockerode');
const url = require('url');

const docker = new Docker();

/**
 * Attach WebSocket server for container exec/shell
 * Route: ws://host/ws/exec/:containerId
 * @param {http.Server} server - The HTTP server instance
 * @param {Object} log - Logger
 */
module.exports = function attachExecWS(server, log, authManager) {
  const wss = new WebSocketServer({ noServer: true });

  // Authenticate WebSocket upgrade request before accepting it
  server.on('upgrade', async (req, socket, head) => {
    const parsed = url.parse(req.url, true);
    const match = parsed.pathname.match(/^\/ws\/exec\/([a-zA-Z0-9_.-]+)$/);
    if (!match) return; // Not our route — let other handlers deal with it

    const containerId = decodeURIComponent(match[1]);

    // Validate container ID format to prevent injection
    if (!/^[a-zA-Z0-9][a-zA-Z0-9_.-]{0,127}$/.test(containerId)) {
      log.warn('exec', 'Invalid container ID in WebSocket path', { containerId });
      socket.write('HTTP/1.1 400 Bad Request\r\n\r\n');
      socket.destroy();
      return;
    }

    // Check auth — require valid JWT or API key from the upgrade request
    let auth = null;
    const authHeader = req.headers.authorization;
    if (authHeader && authHeader.startsWith('Bearer ')) {
      const token = authHeader.substring(7);
      const payload = await authManager.verifyJWT(token);
      if (payload) {
        auth = { type: 'jwt', userId: payload.userId, scope: payload.scope || [] };
      }
    } else {
      const apiKey = req.headers['x-api-key'];
      if (apiKey) {
        const keyData = await authManager.verifyAPIKey(apiKey);
        if (keyData) {
          auth = { type: 'apikey', keyId: keyData.keyId, scope: keyData.scopes || [] };
        }
      }
    }

    if (!auth) {
      log.warn('exec', 'Unauthenticated WebSocket exec attempt', { containerId, ip: req.socket.remoteAddress });
      socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
      socket.destroy();
      return;
    }

    // Auth passed — proceed with WebSocket upgrade
    wss.handleUpgrade(req, socket, head, (ws) => {
      handleExec(ws, containerId, log, auth);
    });
  });

  return wss;
};

async function handleExec(ws, containerId, log, auth) {
  let execStream = null;
  let execInstance = null;

  try {
    const container = docker.getContainer(containerId);
    // Verify container exists and is running
    const info = await container.inspect();
    if (!info.State.Running) {
      ws.send(JSON.stringify({ type: 'error', message: 'Container is not running' }));
      ws.close();
      return;
    }

    log.info('exec', 'Authenticated exec session started', {
      containerId,
      authType: auth.type,
      authId: auth.type === 'jwt' ? auth.userId : auth.keyId
    });

    // Detect available shell
    let shell = '/bin/sh';
    try {
      const bashCheck = await container.exec({ Cmd: ['which', 'bash'], AttachStdout: true });
      const bashStream = await bashCheck.start();
      const chunks = [];
      await new Promise((resolve) => {
        bashStream.on('data', (chunk) => chunks.push(chunk));
        bashStream.on('end', resolve);
      });
      if (chunks.length > 0 && Buffer.concat(chunks).toString().includes('/bash')) {
        shell = '/bin/bash';
      }
    } catch (_) {
      // Fall back to /bin/sh when bash detection fails
    }

    execInstance = await container.exec({
      Cmd: [shell],
      AttachStdin: true,
      AttachStdout: true,
      AttachStderr: true,
      Tty: true,
    });

    execStream = await execInstance.start({ hijack: true, stdin: true, Tty: true });

    ws.send(JSON.stringify({ type: 'connected', shell, containerId }));

    // Docker → WebSocket
    execStream.on('data', (chunk) => {
      if (ws.readyState === ws.OPEN) {
        ws.send(chunk);
      }
    });

    execStream.on('end', () => {
      if (ws.readyState === ws.OPEN) {
        ws.send(JSON.stringify({ type: 'exit' }));
        ws.close();
      }
    });

    // WebSocket → Docker
    ws.on('message', (data) => {
      if (!execStream.writable) return;
      try {
        // Check for control messages (JSON)
        const str = data.toString();
        if (str.startsWith('{"type":')) {
          const msg = JSON.parse(str);
          if (msg.type === 'resize' && execInstance && msg.cols && msg.rows) {
            execInstance.resize({ h: msg.rows, w: msg.cols }).catch(() => {});
            return;
          }
        }
      } catch (_) {
        // Treat message as raw terminal input if JSON parsing fails
      }
      // Regular terminal input
      execStream.write(data);
    });

    ws.on('close', () => {
      if (execStream) {
        try { execStream.destroy(); } catch (_) {
          // Ignore stream teardown errors on socket close
        }
      }
    });

    ws.on('error', (err) => {
      log.warn('exec', 'WebSocket error', { containerId, error: err.message });
      if (execStream) {
        try { execStream.destroy(); } catch (_) {
          // Ignore stream teardown errors after websocket errors
        }
      }
    });

  } catch (err) {
    log.error('exec', 'Failed to start exec session', { containerId, error: err.message });
    if (ws.readyState === ws.OPEN) {
      ws.send(JSON.stringify({ type: 'error', message: err.message }));
      ws.close();
    }
  }
}