fix version

go.mod

@@ -5,7 +5,7 @@ go 1.24.0
 require (
 	github.com/creack/pty/v2 v2.0.1
 	golang.org/x/crypto v0.43.0
-	golang.org/x/term v0.37.0
+	golang.org/x/term v0.36.0
 )
 
 require golang.org/x/sys v0.37.0 // indirect

main.go

@@ -97,15 +97,32 @@ func handleChannel(channel ssh.Channel, requests <-chan *ssh.Request) {
 	defer channel.Close()
 	var ptmx *os.File
 	var termWidth, termHeight uint32 = 80, 24
-	var term string
+	var clientTerm string
+	var ptyAllocated bool // Track if PTY is allocated
+
 	for req := range requests {
 		switch req.Type {
+		case "env":
+			// Handle environment requests (e.g., client TERM propagation)
+			if len(req.Payload) >= 7 && string(req.Payload[:7]) == "TERM=\x00" {
+				termLen := int(req.Payload[7])
+				if len(req.Payload) >= 8+termLen {
+					clientTerm = string(req.Payload[8 : 8+termLen])
+					log.Println("Client TERM from env:", clientTerm)
+				}
+			}
+			req.Reply(true, nil)
 		case "pty-req":
+			// Allocate PTY early on pty-req, before shell/exec
+			if ptyAllocated {
+				req.Reply(false, nil)
+				continue
+			}
 			if len(req.Payload) >= 4 {
 				termLen := binary.BigEndian.Uint32(req.Payload[0:4])
 				if len(req.Payload) >= int(4+termLen+16) {
-					term = string(req.Payload[4 : 4+termLen])
-					log.Println("Client TERM:", term)
+					clientTerm = string(req.Payload[4 : 4+termLen])
+					log.Println("Client TERM from pty-req:", clientTerm)
 					cols := binary.BigEndian.Uint32(req.Payload[4+termLen : 4+termLen+4])
 					rows := binary.BigEndian.Uint32(req.Payload[4+termLen+4 : 4+termLen+8])
 					if cols > 0 {
@@ -116,50 +133,81 @@ func handleChannel(channel ssh.Channel, requests <-chan *ssh.Request) {
 					}
 				}
 			}
-			req.Reply(true, nil)
-		case "window-change":
-			width := binary.BigEndian.Uint32(req.Payload)
-			height := binary.BigEndian.Uint32(req.Payload[4:])
-			if width > 0 {
-				termWidth = width
+			// Set default TERM if not provided (TUI-compatible)
+			if clientTerm == "" {
+				clientTerm = "xterm-256color"
 			}
-			if height > 0 {
-				termHeight = height
-			}
-			if ptmx != nil {
-				pty.Setsize(ptmx, &pty.Winsize{Cols: uint16(termWidth), Rows: uint16(termHeight)})
-			}
-			req.Reply(true, nil)
-		case "shell":
-			req.Reply(true, nil)
 			command := os.Getenv("COMMAND")
 			if command == "" {
 				command = "/app/tui"
 			}
 			cmd := exec.Command(command)
-			envTerm := "TERM=xterm-256color"
-			if term != "" {
-				envTerm = "TERM=" + term
-			}
+			envTerm := "TERM=" + clientTerm
 			cmd.Env = []string{"PATH=/bin", envTerm}
 			cmd.Dir = "/"
 			var err error
 			ptmx, err = pty.StartWithSize(cmd, &pty.Winsize{Cols: uint16(termWidth), Rows: uint16(termHeight)})
 			if err != nil {
 				log.Println("PTY start error:", err)
+				req.Reply(false, nil)
 				return
 			}
-			if err := term.MakeRaw(int(ptmx.Fd())); err != nil {
+			// Make raw mode on master (server side)
+			if _, err := term.MakeRaw(int(ptmx.Fd())); err != nil {
 				log.Println("MakeRaw master error:", err)
 			}
+			// Note: Slave (cmd side) is already raw via pty.Start, but ensure via setsid if needed
+			ptyAllocated = true
+			// Start I/O bridging immediately
 			go func() {
-				defer ptmx.Close()
-				go io.Copy(channel, ptmx)
-				go io.Copy(ptmx, channel)
-				cmd.Wait()
-				channel.Close()
+				defer func() {
+					ptmx.Close()
+					channel.Close()
+				}()
+				// Bidirectional copy with error handling
+				done := make(chan error, 2)
+				go func() { done <- io.Copy(channel, ptmx) }()
+				go func() { done <- io.Copy(ptmx, channel) }()
+				<-done // Wait for one to finish
+				cmd.Process.Signal(os.Interrupt) // Graceful shutdown
+				<-done
 			}()
+			req.Reply(true, nil)
+			// Wait for cmd to finish after PTY setup
+			go func() {
+				cmd.Wait()
+				channel.SendRequest("exit-status", false, []byte{0}) // Send exit status
+			}()
+			return // PTY session started, no more requests
+		case "window-change":
+			// Handle resizes post-PTY allocation
+			if ptyAllocated && ptmx != nil {
+				width := binary.BigEndian.Uint32(req.Payload)
+				height := binary.BigEndian.Uint32(req.Payload[4:])
+				if width > 0 {
+					termWidth = width
+				}
+				if height > 0 {
+					termHeight = height
+				}
+				pty.Setsize(ptmx, &pty.Winsize{Cols: uint16(termWidth), Rows: uint16(termHeight)})
+			}
+			req.Reply(true, nil)
+		case "shell":
+			// Only handle if no PTY (fallback, but PTY should be allocated first)
+			if !ptyAllocated {
+				req.Reply(true, nil)
+				// Fallback to original shell logic if needed
+				log.Println("Shell requested without PTY")
+			} else {
+				req.Reply(false, nil)
+			}
 		case "exec":
+			// Handle exec separately if no PTY
+			if ptyAllocated {
+				req.Reply(false, nil)
+				continue
+			}
 			req.Reply(true, nil)
 			command := string(req.Payload[4:])
 			runCommand(channel, command)
@@ -171,6 +219,7 @@ func handleChannel(channel ssh.Channel, requests <-chan *ssh.Request) {
 }
 
 func runCommand(channel ssh.Channel, command string) {
+	defer channel.Close()
 	cmd := exec.Command("/bin/sh", "-c", command)
 	cmd.Env = []string{"PATH=/bin"}
 	cmd.Dir = "/"
@@ -193,8 +242,9 @@ func runCommand(channel ssh.Channel, command string) {
 		log.Println("Start error:", err)
 		return
 	}
-	go io.Copy(stdin, channel)
-	go io.Copy(channel, stdout)
-	go io.Copy(channel, stderr)
+	go func() { io.Copy(stdin, channel); stdin.Close() }()
+	go func() { io.Copy(channel, stdout); stdout.Close() }()
+	go func() { io.Copy(channel, stderr); stderr.Close() }()
 	cmd.Wait()
 }
+