feat: enhance terminal input handling and improve WebSocket connection safety

new-backend/executor/executor.go

@@ -21,13 +21,26 @@ import (
 	"github.com/ishikabhoyar/monaco/new-backend/models"
 )
 
+// SafeWebSocketConn wraps a WebSocket connection with a mutex for thread-safe writes
+type SafeWebSocketConn struct {
+	conn  *websocket.Conn
+	mutex sync.Mutex
+}
+
+// WriteJSON writes JSON to the WebSocket connection in a thread-safe manner
+func (s *SafeWebSocketConn) WriteJSON(v interface{}) error {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+	return s.conn.WriteJSON(v)
+}
+
 // CodeExecutor handles code execution for all languages
 type CodeExecutor struct {
 	config              *config.Config
 	execQueue           chan *models.CodeSubmission
 	submissions         map[string]*models.CodeSubmission
 	submissionsMutex    sync.RWMutex
-	terminalConnections map[string][]*websocket.Conn
+	terminalConnections map[string][]*SafeWebSocketConn
 	terminalMutex       sync.RWMutex
 	inputChannels       map[string]chan string
 	inputMutex          sync.RWMutex
@@ -39,7 +52,7 @@ func NewCodeExecutor(cfg *config.Config) *CodeExecutor {
 		config:              cfg,
 		execQueue:           make(chan *models.CodeSubmission, cfg.Executor.QueueCapacity),
 		submissions:         make(map[string]*models.CodeSubmission),
-		terminalConnections: make(map[string][]*websocket.Conn),
+		terminalConnections: make(map[string][]*SafeWebSocketConn),
 		inputChannels:       make(map[string]chan string),
 	}
 
@@ -87,7 +100,8 @@ func (e *CodeExecutor) RegisterTerminalConnection(submissionID string, conn *web
 	e.terminalMutex.Lock()
 	defer e.terminalMutex.Unlock()
 
-	e.terminalConnections[submissionID] = append(e.terminalConnections[submissionID], conn)
+	safeConn := &SafeWebSocketConn{conn: conn}
+	e.terminalConnections[submissionID] = append(e.terminalConnections[submissionID], safeConn)
 	
 	log.Printf("WebSocket connection registered for submission %s (total: %d)",
 		submissionID, len(e.terminalConnections[submissionID]))
@@ -103,13 +117,14 @@ func (e *CodeExecutor) UnregisterTerminalConnection(submissionID string, conn *w
 
 	connections := e.terminalConnections[submissionID]
 	for i, c := range connections {
-		if c == conn {
+		if c.conn == conn {
 			// Remove the connection
 			e.terminalConnections[submissionID] = append(connections[:i], connections[i+1:]...)
 			break
 		}
 	}
 
+
 	// Clean up if no more connections
 	if len(e.terminalConnections[submissionID]) == 0 {
 		delete(e.terminalConnections, submissionID)
@@ -326,7 +341,7 @@ func (e *CodeExecutor) executeJava(submission *models.CodeSubmission, tempDir st
 
 // executeC executes C code
 func (e *CodeExecutor) executeC(submission *models.CodeSubmission, tempDir string, langConfig config.LanguageConfig) {
-    // Write code to file
+    // Write code to file exactly as submitted by user
     codeFile := filepath.Join(tempDir, "code"+langConfig.FileExt)
     if err := os.WriteFile(codeFile, []byte(submission.Code), 0644); err != nil {
         submission.Status = "failed"
@@ -334,62 +349,6 @@ func (e *CodeExecutor) executeC(submission *models.CodeSubmission, tempDir strin
         return
     }
 
-    // Create a wrapper script that will include setbuf to disable buffering
-    wrapperCode := `#include <stdio.h>
-
-// Forward declaration of user's main function
-int user_main();
-
-int main() {
-    // Disable buffering completely for stdout
-    setbuf(stdout, NULL);
-    
-    // Call the user's code
-    return user_main();
-}
-
-// User's code begins here
-`
-    
-    // Modify the user's code to be a function called from our wrapper
-    modifiedCode := submission.Code
-    // Replace main function with our wrapper
-    mainRegex := regexp.MustCompile(`int\s+main\s*\([^)]*\)\s*{`)
-    if mainRegex.MatchString(modifiedCode) {
-        // Rename user's main to user_main
-        modifiedCode = mainRegex.ReplaceAllString(modifiedCode, "int user_main() {")
-        
-        // Combine wrapper with modified user code
-        finalCode := wrapperCode + modifiedCode
-        
-        // Write the final code with wrapper to file
-        if err := os.WriteFile(codeFile, []byte(finalCode), 0644); err != nil {
-            submission.Status = "failed"
-            submission.Output = "Failed to write code file: " + err.Error()
-            return
-        }
-    } else {
-        // If no main function found, create a minimal program that includes the user code
-        finalCode := `#include <stdio.h>
-
-int main() {
-    // Disable buffering completely for stdout
-    setbuf(stdout, NULL);
-    
-    // Execute the user's code
-    ` + submission.Code + `
-    
-    return 0;
-}
-`
-        // Write the final code to file
-        if err := os.WriteFile(codeFile, []byte(finalCode), 0644); err != nil {
-            submission.Status = "failed"
-            submission.Output = "Failed to write code file: " + err.Error()
-            return
-        }
-    }
-
     // Compile C code
     compileCmd := exec.Command(
         "docker", "run", "--rm",
@@ -406,7 +365,7 @@ int main() {
         return
     }
     
-    // Setup Docker run command
+    // Setup Docker run command with stdbuf to force line buffering
     cmd := exec.Command(
         "docker", "run", "--rm", "-i",
         "--network=none",
@@ -415,7 +374,7 @@ int main() {
         "--pids-limit=20",
         "-v", tempDir+":/code",
         langConfig.Image,
-        "/code/program",
+        "stdbuf", "-oL", "-eL", "/code/program",
     )
 
     // Execute the code with input handling
@@ -448,7 +407,7 @@ func (e *CodeExecutor) executeCpp(submission *models.CodeSubmission, tempDir str
 		return
 	}
 	
-	// Setup Docker run command
+	// Setup Docker run command with stdbuf to force line buffering
 	cmd := exec.Command(
 		"docker", "run", "--rm", "-i",
 		"--network=none",
@@ -457,7 +416,7 @@ func (e *CodeExecutor) executeCpp(submission *models.CodeSubmission, tempDir str
 		"--pids-limit=20",
 		"-v", tempDir+":/code",
 		langConfig.Image,
-		"/code/program",
+		"stdbuf", "-oL", "-eL", "/code/program",
 	)
 
 	// Execute the code with input handling