RDSViewer.jsx

Go to the documentation of this file.

/**
 * RDS (Remote Desktop) Viewer Component
 * 
 * Integrates with existing WebSocket RDS infrastructure to display remote desktop
 * sessions using HTML canvas rendering with mouse/keyboard input.
 * 
 * Features:
 * - Real-time screen streaming via WebSocket
 * - Mouse click and movement
 * - Keyboard input
 * - Clipboard sync (bidirectional)
 * - Connection status indicators
 * - Auto-reconnect on disconnect
 */
 
import { useState, useEffect, useRef } from 'react';
import {
    Box,
    Paper,
    Alert,
    Typography,
    IconButton,
    Tooltip,
    CircularProgress,
    TextField,
    Button,
    Stack
} from '@mui/material';
import {
    Fullscreen as FullscreenIcon,
    FullscreenExit as FullscreenExitIcon,
    Refresh as RefreshIcon,
    ContentCopy as CopyIcon,
    ContentPaste as PasteIcon
} from '@mui/icons-material';
import { apiFetch } from '../lib/api';
 
export default function RDSViewer({ agentId, agentUuid }) {
    const [status, setStatus] = useState('idle'); // idle, starting, connecting, connected, error, disconnected
    const [error, setError] = useState(null);
    const [sessionId, setSessionId] = useState(null);
    const [fullscreen, setFullscreen] = useState(false);
    const [clipboard, setClipboard] = useState('');
    const [agentClipboard, setAgentClipboard] = useState('');
 
    const canvasRef = useRef(null);
    const wsRef = useRef(null);
    const containerRef = useRef(null);
    const contextRef = useRef(null);
    const imageRef = useRef(new Image());
 
    // Start RDS session
    const startSession = async () => {
        try {
            setStatus('starting');
            setError(null);
 
            // Request session from backend
            const response = await apiFetch(`/rds/start/${agentUuid}`, {
                method: 'POST'
            });
 
            setSessionId(response.sessionId);
            setStatus('connecting');
 
            // Connect to viewer WebSocket
            connectWebSocket(response.sessionId);
 
        } catch (err) {
            console.error('[RDS] Failed to start session:', err);
            setError(err.message || 'Failed to start remote desktop session');
            setStatus('error');
        }
    };
 
    // Connect to WebSocket viewer
    const connectWebSocket = (sid) => {
        const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
        const host = window.location.host;
        const wsUrl = `${protocol}//${host}/api/rds/view/${sid}?session=${sid}`;
 
        console.log('[RDS] Connecting to:', wsUrl);
 
        const ws = new WebSocket(wsUrl);
        ws.binaryType = 'arraybuffer';
        wsRef.current = ws;
 
        ws.onopen = () => {
            console.log('[RDS] WebSocket connected');
            setStatus('connected');
        };
 
        ws.onmessage = (event) => {
            // Handle JSON messages (clipboard, status)
            if (typeof event.data === 'string') {
                try {
                    const msg = JSON.parse(event.data);
                    handleMessage(msg);
                } catch (err) {
                    console.error('[RDS] Failed to parse message:', err);
                }
                return;
            }
 
            // Handle binary screen data (JPEG frames)
            if (event.data instanceof ArrayBuffer) {
                const blob = new Blob([event.data], { type: 'image/jpeg' });
                const url = URL.createObjectURL(blob);
 
                imageRef.current.onload = () => {
                    if (canvasRef.current && contextRef.current) {
                        const canvas = canvasRef.current;
                        const ctx = contextRef.current;
 
                        // Auto-resize canvas to match image
                        if (canvas.width !== imageRef.current.width || canvas.height !== imageRef.current.height) {
                            canvas.width = imageRef.current.width;
                            canvas.height = imageRef.current.height;
                        }
 
                        ctx.drawImage(imageRef.current, 0, 0);
                    }
                    URL.revokeObjectURL(url);
                };
 
                imageRef.current.src = url;
            }
        };
 
        ws.onerror = (err) => {
            console.error('[RDS] WebSocket error:', err);
            setError('WebSocket connection error');
            setStatus('error');
        };
 
        ws.onclose = () => {
            console.log('[RDS] WebSocket closed');
            setStatus('disconnected');
        };
    };
 
    // Handle JSON messages (clipboard, status, etc.)
    const handleMessage = (msg) => {
        if (msg.type === 'clipboard') {
            setAgentClipboard(msg.data || '');
        } else if (msg.type === 'status') {
            console.log('[RDS] Status update:', msg);
        }
    };
 
    // Send mouse event
    const handleMouse = (event) => {
        if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN || !sessionId) return;
 
        const canvas = canvasRef.current;
        const rect = canvas.getBoundingClientRect();
        const x = Math.floor((event.clientX - rect.left) * (canvas.width / rect.width));
        const y = Math.floor((event.clientY - rect.top) * (canvas.height / rect.height));
 
        let button = 0;
        if (event.type === 'mousedown' || event.type === 'mouseup' || event.type === 'click') {
            button = event.button;
        }
 
        wsRef.current.send(JSON.stringify({
            type: 'mouse',
            sessionId: sessionId,
            x: x,
            y: y,
            button: button,
            action: event.type
        }));
    };
 
    // Send keyboard event
    const handleKeyboard = (event) => {
        if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN || !sessionId) return;
 
        wsRef.current.send(JSON.stringify({
            type: 'keyboard',
            sessionId: sessionId,
            key: event.key,
            code: event.code,
            action: event.type,
            ctrlKey: event.ctrlKey,
            shiftKey: event.shiftKey,
            altKey: event.altKey,
            metaKey: event.metaKey
        }));
    };
 
    // Send clipboard to agent
    const sendClipboard = () => {
        if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN || !sessionId) return;
 
        wsRef.current.send(JSON.stringify({
            type: 'clipboard',
            sessionId: sessionId,
            data: clipboard
        }));
 
        console.log('[RDS] Sent clipboard to agent');
    };
 
    // Copy agent clipboard to local
    const copyAgentClipboard = async () => {
        try {
            await navigator.clipboard.writeText(agentClipboard);
            console.log('[RDS] Copied agent clipboard to local');
        } catch (err) {
            console.error('[RDS] Failed to copy to clipboard:', err);
        }
    };
 
    // Toggle fullscreen
    const toggleFullscreen = () => {
        if (!fullscreen) {
            containerRef.current?.requestFullscreen();
            setFullscreen(true);
        } else {
            document.exitFullscreen();
            setFullscreen(false);
        }
    };
 
    // Stop session
    const stopSession = async () => {
        if (wsRef.current) {
            wsRef.current.close();
            wsRef.current = null;
        }
 
        if (sessionId) {
            try {
                await apiFetch(`/rds/stop/${agentUuid}`, {
                    method: 'POST'
                });
            } catch (err) {
                console.error('[RDS] Failed to stop session:', err);
            }
            setSessionId(null);
        }
 
        setStatus('idle');
    };
 
    // Initialize canvas context
    useEffect(() => {
        if (canvasRef.current) {
            contextRef.current = canvasRef.current.getContext('2d');
        }
    }, []);
 
    // Cleanup on unmount
    useEffect(() => {
        return () => {
            stopSession();
        };
    }, []);
 
    const statusColors = {
        idle: 'default',
        starting: 'info',
        connecting: 'info',
        connected: 'success',
        error: 'error',
        disconnected: 'warning'
    };
 
    const statusLabels = {
        idle: 'Not connected',
        starting: 'Starting session...',
        connecting: 'Connecting...',
        connected: 'Connected',
        error: 'Error',
        disconnected: 'Disconnected'
    };
 
    return (
        <Box sx={{ height: '100%', display: 'flex', flexDirection: 'column' }}>
            {/* Status Bar */}
            <Alert severity={statusColors[status]} sx={{ mb: 2 }}>
                Status: {statusLabels[status]}
                {status === 'starting' && <CircularProgress size={16} sx={{ ml: 1 }} />}
            </Alert>
 
            {error && (
                <Alert severity="error" sx={{ mb: 2 }}>
                    {error}
                </Alert>
            )}
 
            {/* Control Bar */}
            <Paper sx={{ p: 1, mb: 2, display: 'flex', gap: 1, alignItems: 'center' }}>
                {status === 'idle' && (
                    <Button variant="contained" onClick={startSession}>
                        Start Remote Desktop
                    </Button>
                )}
 
                {status !== 'idle' && status !== 'starting' && (
                    <>
                        <Button variant="outlined" color="error" onClick={stopSession}>
                            Stop Session
                        </Button>
 
                        <Tooltip title="Refresh">
                            <IconButton onClick={() => { stopSession(); setTimeout(startSession, 1000); }}>
                                <RefreshIcon />
                            </IconButton>
                        </Tooltip>
 
                        <Tooltip title={fullscreen ? "Exit Fullscreen" : "Fullscreen"}>
                            <IconButton onClick={toggleFullscreen}>
                                {fullscreen ? <FullscreenExitIcon /> : <FullscreenIcon />}
                            </IconButton>
                        </Tooltip>
                    </>
                )}
            </Paper>
 
            {/* Viewer Canvas */}
            <Paper
                ref={containerRef}
                sx={{
                    flex: 1,
                    overflow: 'auto',
                    bgcolor: '#1e1e1e',
                    display: 'flex',
                    alignItems: 'center',
                    justifyContent: 'center',
                    position: 'relative'
                }}
            >
                {status === 'idle' && (
                    <Typography color="text.secondary">
                        Click "Start Remote Desktop" to begin
                    </Typography>
                )}
 
                {(status === 'starting' || status === 'connecting') && (
                    <Box sx={{ textAlign: 'center' }}>
                        <CircularProgress />
                        <Typography color="text.secondary" sx={{ mt: 2 }}>
                            {statusLabels[status]}
                        </Typography>
                    </Box>
                )}
 
                {status === 'connected' && (
                    <canvas
                        ref={canvasRef}
                        onMouseDown={handleMouse}
                        onMouseUp={handleMouse}
                        onMouseMove={handleMouse}
                        onClick={handleMouse}
                        onKeyDown={handleKeyboard}
                        onKeyUp={handleKeyboard}
                        tabIndex={0}
                        style={{
                            maxWidth: '100%',
                            maxHeight: '100%',
                            cursor: 'crosshair',
                            outline: 'none'
                        }}
                    />
                )}
            </Paper>
 
            {/* Clipboard Controls */}
            {status === 'connected' && (
                <Paper sx={{ p: 2, mt: 2 }}>
                    <Typography variant="h6" gutterBottom>
                        Clipboard Sync
                    </Typography>
 
                    <Stack spacing={2}>
                        {/* Send to Agent */}
                        <Box>
                            <Typography variant="subtitle2" gutterBottom>
                                Send to Remote Desktop
                            </Typography>
                            <Stack direction="row" spacing={1}>
                                <TextField
                                    size="small"
                                    fullWidth
                                    multiline
                                    rows={2}
                                    value={clipboard}
                                    onChange={(e) => setClipboard(e.target.value)}
                                    placeholder="Type text to send to remote desktop"
                                />
                                <Button
                                    variant="contained"
                                    startIcon={<PasteIcon />}
                                    onClick={sendClipboard}
                                    disabled={!clipboard}
                                >
                                    Send
                                </Button>
                            </Stack>
                        </Box>
 
                        {/* Receive from Agent */}
                        <Box>
                            <Typography variant="subtitle2" gutterBottom>
                                Received from Remote Desktop
                            </Typography>
                            <Stack direction="row" spacing={1}>
                                <TextField
                                    size="small"
                                    fullWidth
                                    multiline
                                    rows={2}
                                    value={agentClipboard}
                                    InputProps={{ readOnly: true }}
                                    placeholder="Clipboard content from remote desktop will appear here"
                                />
                                <Button
                                    variant="outlined"
                                    startIcon={<CopyIcon />}
                                    onClick={copyAgentClipboard}
                                    disabled={!agentClipboard}
                                >
                                    Copy
                                </Button>
                            </Stack>
                        </Box>
                    </Stack>
                </Paper>
            )}
        </Box>
    );
}