Implement WebSocket support for terminal connections and enhance terminal UI

2025-03-30 00:13:32 +05:30
parent 305650925e
commit 3a75000e12
9 changed files with 590 additions and 225 deletions
--- a/backend/go.mod
+++ b/backend/go.mod
@@ -6,6 +6,7 @@ require github.com/stretchr/testify v1.9.0

 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/gorilla/websocket v1.5.3 // indirect
 	github.com/kr/pretty v0.3.0 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/rogpeppe/go-internal v1.12.0 // indirect
--- a/backend/go.sum
+++ b/backend/go.sum
@@ -1,6 +1,8 @@
 github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
+github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
 github.com/kr/pretty v0.3.0 h1:WgNl7dwNpEZ6jJ9k1snq4pZsg7DOEN8hP9Xw0Tsjwk0=
 github.com/kr/pretty v0.3.0/go.mod h1:640gp4NfQd8pI5XOwp5fnNeVWj67G7CFk/SaSQn7NBk=
--- a/backend/handler/handler.go
+++ b/backend/handler/handler.go
@@ -2,12 +2,15 @@ package handler

 import (
 	"encoding/json"
+	"fmt"
+	"log"
 	"net/http"
 	"sync"
 	"time"

 	"github.com/arnab-afk/monaco/model"
 	"github.com/arnab-afk/monaco/service"
+	"github.com/gorilla/websocket"
 )

 // Handler manages HTTP requests for code submissions
@@ -179,6 +182,63 @@ func (h *Handler) QueueStatsHandler(w http.ResponseWriter, r *http.Request) {
 	})
 }

+// ConnectTerminal connects a WebSocket to a running execution
+func (h *Handler) ConnectTerminal(conn *websocket.Conn, executionID string) {
+	// Get submission from storage
+	h.mu.Lock()
+	submission, found := h.submissions[executionID]
+	status := "not found"
+	if found {
+		status = submission.Status
+	}
+	h.mu.Unlock()
+
+	log.Printf("[WS-%s] Terminal connection request, submission status: %s", executionID, status)
+
+	if !found {
+		log.Printf("[WS-%s] Execution not found", executionID)
+		conn.WriteMessage(websocket.TextMessage, []byte("Execution not found"))
+		conn.Close()
+		return
+	}
+
+	// If execution is already completed, send stored output and close
+	if submission.Status == "completed" || submission.Status == "failed" {
+		log.Printf("[WS-%s] Execution already %s, sending stored output (length: %d)",
+			executionID, submission.Status, len(submission.Output))
+		conn.WriteMessage(websocket.TextMessage, []byte(submission.Output))
+		conn.Close()
+		return
+	}
+
+	log.Printf("[WS-%s] Registering connection for real-time updates, current status: %s",
+		executionID, submission.Status)
+
+	// Register this connection with the execution service for real-time updates
+	h.executionService.RegisterTerminalConnection(executionID, conn)
+
+	// Send initial connection confirmation
+	initialMsg := fmt.Sprintf("[System] Connected to process (ID: %s, Status: %s)\n",
+		executionID, submission.Status)
+	conn.WriteMessage(websocket.TextMessage, []byte(initialMsg))
+
+	// Handle incoming messages from the terminal (for stdin)
+	go func() {
+		for {
+			_, message, err := conn.ReadMessage()
+			if err != nil {
+				log.Printf("[WS-%s] Read error: %v", executionID, err)
+				h.executionService.UnregisterTerminalConnection(executionID, conn)
+				break
+			}
+
+			log.Printf("[WS-%s] Received input from client: %s", executionID, string(message))
+			// Send input to the execution if it's waiting for input
+			h.executionService.SendInput(executionID, string(message))
+		}
+	}()
+}
+
 // generateID creates a unique ID for submissions
 func (h *Handler) generateID() string {
 	return service.GenerateUUID()
--- a/backend/main.go
+++ b/backend/main.go
@@ -7,6 +7,7 @@ import (
 	"time"

 	"github.com/arnab-afk/monaco/handler"
+	"github.com/gorilla/websocket"
 )

 func main() {
@@ -46,7 +47,40 @@ func main() {
 		}
 	}

-	// Register handlers with logging and CORS middleware
+	// Configure WebSocket upgrader
+	upgrader := websocket.Upgrader{
+		ReadBufferSize:  1024,
+		WriteBufferSize: 1024,
+		// Allow connections from any origin
+		CheckOrigin: func(r *http.Request) bool {
+			return true
+		},
+	}
+
+	// WebSocket handler for terminal connection
+	http.HandleFunc("/ws/terminal", func(w http.ResponseWriter, r *http.Request) {
+		// Get execution ID from query parameters
+		executionID := r.URL.Query().Get("id")
+		if executionID == "" {
+			log.Println("[WS] Missing execution ID")
+			http.Error(w, "Missing execution ID", http.StatusBadRequest)
+			return
+		}
+
+		// Upgrade HTTP connection to WebSocket
+		conn, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			log.Printf("[WS] Failed to upgrade connection: %v", err)
+			return
+		}
+
+		log.Printf("[WS] Terminal connection established for execution ID: %s", executionID)
+
+		// Connect this WebSocket to the execution service for real-time streaming
+		h.ConnectTerminal(conn, executionID)
+	})
+
+	// Register REST API handlers with logging and CORS middleware
 	http.HandleFunc("/submit", corsMiddleware(loggingMiddleware(h.SubmitHandler)))
 	http.HandleFunc("/status", corsMiddleware(loggingMiddleware(h.StatusHandler)))
 	http.HandleFunc("/result", corsMiddleware(loggingMiddleware(h.ResultHandler)))
--- a/backend/service/execution.go
+++ b/backend/service/execution.go
@@ -1,6 +1,8 @@
 package service

 import (
+	"bytes"
+	"context"
 	"fmt"
 	"io"
 	"log"
@@ -8,25 +10,101 @@ import (
 	"os/exec"
 	"path/filepath"
 	"regexp"
-	"strings"
 	"sync"
 	"time"

 	"github.com/arnab-afk/monaco/model"
 	"github.com/arnab-afk/monaco/queue"
+	"github.com/gorilla/websocket"
 )

 // ExecutionService handles code execution for multiple languages
 type ExecutionService struct {
-	mu    sync.Mutex
-	queue *queue.JobQueue
+	mu                  sync.Mutex
+	queue               *queue.JobQueue
+	terminalConnections map[string][]*websocket.Conn // Map of executionID to WebSocket connections
+	execInputChannels   map[string]chan string       // Map of executionID to input channels
 }

 // NewExecutionService creates a new execution service
 func NewExecutionService() *ExecutionService {
 	log.Println("Initializing execution service with 3 concurrent workers")
 	return &ExecutionService{
-		queue: queue.NewJobQueue(35), // 3 concurrent executions max
+		queue:               queue.NewJobQueue(3), // 3 concurrent executions max
+		terminalConnections: make(map[string][]*websocket.Conn),
+		execInputChannels:   make(map[string]chan string),
+	}
+}
+
+// RegisterTerminalConnection registers a WebSocket connection for an execution
+func (s *ExecutionService) RegisterTerminalConnection(executionID string, conn *websocket.Conn) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if _, exists := s.terminalConnections[executionID]; !exists {
+		s.terminalConnections[executionID] = make([]*websocket.Conn, 0)
+	}
+	s.terminalConnections[executionID] = append(s.terminalConnections[executionID], conn)
+	log.Printf("[WS-%s] Terminal connection registered, total connections: %d",
+		executionID, len(s.terminalConnections[executionID]))
+}
+
+// UnregisterTerminalConnection removes a WebSocket connection
+func (s *ExecutionService) UnregisterTerminalConnection(executionID string, conn *websocket.Conn) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	connections, exists := s.terminalConnections[executionID]
+	if !exists {
+		return
+	}
+
+	// Remove the specific connection
+	for i, c := range connections {
+		if c == conn {
+			s.terminalConnections[executionID] = append(connections[:i], connections[i+1:]...)
+			break
+		}
+	}
+
+	// If no more connections, clean up
+	if len(s.terminalConnections[executionID]) == 0 {
+		delete(s.terminalConnections, executionID)
+	}
+
+	log.Printf("[WS-%s] Terminal connection unregistered", executionID)
+}
+
+// SendOutputToTerminals sends output to all connected terminals for an execution
+func (s *ExecutionService) SendOutputToTerminals(executionID string, output string) {
+	s.mu.Lock()
+	connections := s.terminalConnections[executionID]
+	s.mu.Unlock()
+
+	for _, conn := range connections {
+		if err := conn.WriteMessage(websocket.TextMessage, []byte(output)); err != nil {
+			log.Printf("[WS-%s] Error sending to terminal: %v", executionID, err)
+			// Unregister this connection on error
+			s.UnregisterTerminalConnection(executionID, conn)
+		}
+	}
+}
+
+// SendInput sends user input to a running process
+func (s *ExecutionService) SendInput(executionID string, input string) {
+	s.mu.Lock()
+	inputChan, exists := s.execInputChannels[executionID]
+	s.mu.Unlock()
+
+	if exists {
+		select {
+		case inputChan <- input:
+			log.Printf("[WS-%s] Sent input to execution: %s", executionID, input)
+		default:
+			log.Printf("[WS-%s] Execution not ready for input", executionID)
+		}
+	} else {
+		log.Printf("[WS-%s] No input channel for execution", executionID)
 	}
 }

@@ -110,48 +188,123 @@ func (s *ExecutionService) executeLanguageSpecific(submission *model.CodeSubmiss
 func (s *ExecutionService) executeWithInput(cmd *exec.Cmd, input string, timeout time.Duration, submissionID string) ([]byte, error) {
 	log.Printf("[TIMEOUT-%s] Setting execution timeout: %v", submissionID, timeout)

-	// Set up input pipe if input is provided
-	if input != "" {
-		stdin, err := cmd.StdinPipe()
-		if err != nil {
-			log.Printf("[ERROR-%s] Failed to create stdin pipe: %v", submissionID, err)
-			return nil, err
-		}
-
-		// Write input in a goroutine to avoid blocking
-		go func() {
-			defer stdin.Close()
-			io.WriteString(stdin, input)
-		}()
-
-		log.Printf("[INPUT-%s] Providing input to process", submissionID)
+	// Create pipes for stdin, stdout, and stderr
+	stdin, stdinErr := cmd.StdinPipe()
+	if stdinErr != nil {
+		return nil, fmt.Errorf("failed to create stdin pipe: %v", stdinErr)
 	}

-	done := make(chan struct{})
-	var output []byte
-	var err error
+	stdout, stdoutErr := cmd.StdoutPipe()
+	if stdoutErr != nil {
+		return nil, fmt.Errorf("failed to create stdout pipe: %v", stdoutErr)
+	}

-	go func() {
-		log.Printf("[EXEC-%s] Starting command execution: %v", submissionID, cmd.Args)
-		output, err = cmd.CombinedOutput()
-		close(done)
+	stderr, stderrErr := cmd.StderrPipe()
+	if stderrErr != nil {
+		return nil, fmt.Errorf("failed to create stderr pipe: %v", stderrErr)
+	}
+
+	// Create an input channel and register it
+	inputChan := make(chan string, 10)
+	s.mu.Lock()
+	s.execInputChannels[submissionID] = inputChan
+	s.mu.Unlock()
+
+	// Clean up the input channel when done
+	defer func() {
+		s.mu.Lock()
+		delete(s.execInputChannels, submissionID)
+		s.mu.Unlock()
+		close(inputChan)
 	}()

+	// Start the command
+	if err := cmd.Start(); err != nil {
+		return nil, fmt.Errorf("failed to start process: %v", err)
+	}
+
+	// Create a buffer to collect all output
+	var outputBuffer bytes.Buffer
+
+	// Handle stdout in a goroutine
+	go func() {
+		buffer := make([]byte, 1024)
+		for {
+			n, err := stdout.Read(buffer)
+			if n > 0 {
+				data := buffer[:n]
+				outputBuffer.Write(data)
+				// Send real-time output to connected terminals
+				s.SendOutputToTerminals(submissionID, string(data))
+			}
+			if err != nil {
+				break
+			}
+		}
+	}()
+
+	// Handle stderr in a goroutine
+	go func() {
+		buffer := make([]byte, 1024)
+		for {
+			n, err := stderr.Read(buffer)
+			if n > 0 {
+				data := buffer[:n]
+				outputBuffer.Write(data)
+				// Send real-time output to connected terminals
+				s.SendOutputToTerminals(submissionID, string(data))
+			}
+			if err != nil {
+				break
+			}
+		}
+	}()
+
+	// Write initial input if provided
+	if input != "" {
+		io.WriteString(stdin, input+"\n")
+	}
+
+	// Process is in a separate context, but it needs to be killed if timeout occurs
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// Handle additional input from WebSocket in a goroutine
+	go func() {
+		for {
+			select {
+			case additionalInput, ok := <-inputChan:
+				if !ok {
+					return
+				}
+				log.Printf("[INPUT-%s] Received input from WebSocket: %s", submissionID, additionalInput)
+				io.WriteString(stdin, additionalInput+"\n")
+			case <-ctx.Done():
+				return
+			}
+		}
+	}()
+
+	// Wait for the command to complete with timeout
+	done := make(chan error, 1)
+	go func() {
+		done <- cmd.Wait()
+	}()
+
+	// Wait for completion or timeout
 	select {
 	case <-time.After(timeout):
+		cancel() // Stop the input handler
 		log.Printf("[TIMEOUT-%s] Execution timed out after %v seconds", submissionID, timeout.Seconds())
 		if err := cmd.Process.Kill(); err != nil {
 			log.Printf("[TIMEOUT-%s] Failed to kill process: %v", submissionID, err)
-			return nil, fmt.Errorf("timeout reached but failed to kill process: %v", err)
 		}
-		return nil, fmt.Errorf("execution timed out after %v seconds", timeout.Seconds())
-	case <-done:
-		if err != nil {
-			log.Printf("[EXEC-%s] Command execution failed: %v", submissionID, err)
-		} else {
-			log.Printf("[EXEC-%s] Command execution completed successfully", submissionID)
-		}
-		return output, err
+		s.SendOutputToTerminals(submissionID, fmt.Sprintf("\n[System] Process killed after timeout of %v seconds", timeout.Seconds()))
+		return outputBuffer.Bytes(), fmt.Errorf("execution timed out after %v seconds", timeout.Seconds())
+	case err := <-done:
+		cancel() // Stop the input handler
+		s.SendOutputToTerminals(submissionID, "\n[System] Process completed")
+		return outputBuffer.Bytes(), err
 	}
 }

@@ -169,15 +322,9 @@ func (s *ExecutionService) executePython(submission *model.CodeSubmission) {
 		"python:3.9", "python", "-c", submission.Code)

 	log.Printf("[PYTHON-%s] Executing Python code with timeout: 10s", submission.ID)
-	var output []byte
-	var err error

-	if submission.Input != "" {
-		cmd.Stdin = strings.NewReader(submission.Input)
-		output, err = cmd.CombinedOutput()
-	} else {
-		output, err = s.executeWithTimeout(cmd, 10*time.Second, submission.ID)
-	}
+	// Use the enhanced executeWithInput method for all executions
+	output, err := s.executeWithInput(cmd, submission.Input, 100*time.Second, submission.ID)

 	elapsed := time.Since(startTime)
 	log.Printf("[PYTHON-%s] Python execution completed in %v", submission.ID, elapsed)
@@ -255,7 +402,7 @@ func (s *ExecutionService) executeJava(submission *model.CodeSubmission) {

 	log.Printf("[JAVA-%s] Compilation successful", submission.ID)

-	// Now run the compiled class
+	// Now run the compiled class with the enhanced executeWithInput method
 	runCmd := exec.Command("docker", "run", "--rm", "-i",
 		"--network=none",       // No network access
 		"--memory=400m",        // Memory limit
@@ -267,17 +414,8 @@ func (s *ExecutionService) executeJava(submission *model.CodeSubmission) {
 		"-Xverify:none", "-Xms64m", "-Xmx256m",
 		"-cp", "/code", className)

-	// Add input if provided
-	var output []byte
-
-	if submission.Input != "" {
-		log.Printf("[JAVA-%s] Executing Java code with input", submission.ID)
-		runCmd.Stdin = strings.NewReader(submission.Input)
-		output, err = runCmd.CombinedOutput()
-	} else {
-		log.Printf("[JAVA-%s] Executing Java code without input", submission.ID)
-		output, err = s.executeWithTimeout(runCmd, 15*time.Second, submission.ID)
-	}
+	log.Printf("[JAVA-%s] Executing Java code", submission.ID)
+	output, err := s.executeWithInput(runCmd, submission.Input, 15*time.Second, submission.ID)

 	elapsed := time.Since(startTime)
 	log.Printf("[JAVA-%s] Java execution completed in %v", submission.ID, elapsed)
@@ -327,7 +465,7 @@ func (s *ExecutionService) executeC(submission *model.CodeSubmission) {

 	log.Printf("[C-%s] Compilation successful", submission.ID)

-	// Run C executable
+	// Run C executable using executeWithInput to support WebSockets
 	runCmd := exec.Command("docker", "run", "--rm", "-i",
 		"--network=none",       // No network access
 		"--memory=100m",        // Memory limit
@@ -336,17 +474,8 @@ func (s *ExecutionService) executeC(submission *model.CodeSubmission) {
 		"-v", tempDir+":/code", // Mount code directory
 		"gcc:latest", "/code/solution")

-	// Add input if provided
-	var output []byte
-	// Don't redeclare err here - use the existing variable
-	if submission.Input != "" {
-		log.Printf("[C-%s] Executing C code with input", submission.ID)
-		runCmd.Stdin = strings.NewReader(submission.Input)
-		output, err = runCmd.CombinedOutput() // Use the existing err variable
-	} else {
-		log.Printf("[C-%s] Executing C code without input", submission.ID)
-		output, err = s.executeWithTimeout(runCmd, 10*time.Second, submission.ID) // Use the existing err variable
-	}
+	log.Printf("[C-%s] Executing C code", submission.ID)
+	output, err := s.executeWithInput(runCmd, submission.Input, 30*time.Second, submission.ID)

 	elapsed := time.Since(startTime)
 	log.Printf("[C-%s] C execution completed in %v", submission.ID, elapsed)
@@ -396,7 +525,7 @@ func (s *ExecutionService) executeCpp(submission *model.CodeSubmission) {

 	log.Printf("[CPP-%s] Compilation successful", submission.ID)

-	// Run C++ executable
+	// Run C++ executable using executeWithInput to support WebSockets
 	runCmd := exec.Command("docker", "run", "--rm", "-i",
 		"--network=none",       // No network access
 		"--memory=100m",        // Memory limit
@@ -405,16 +534,8 @@ func (s *ExecutionService) executeCpp(submission *model.CodeSubmission) {
 		"-v", tempDir+":/code", // Mount code directory
 		"gcc:latest", "/code/solution")

-	// Add input if provided
-	var output []byte
-	if submission.Input != "" {
-		log.Printf("[CPP-%s] Executing C++ code with input", submission.ID)
-		runCmd.Stdin = strings.NewReader(submission.Input)
-		output, err = runCmd.CombinedOutput()
-	} else {
-		log.Printf("[CPP-%s] Executing C++ code without input", submission.ID)
-		output, err = s.executeWithTimeout(runCmd, 10*time.Second, submission.ID)
-	}
+	log.Printf("[CPP-%s] Executing C++ code", submission.ID)
+	output, err := s.executeWithInput(runCmd, submission.Input, 100*time.Second, submission.ID)

 	elapsed := time.Since(startTime)
 	log.Printf("[CPP-%s] C++ execution completed in %v", submission.ID, elapsed)
--- a/backend/tmp/main.exe
+++ b/backend/tmp/main.exe