Close shell properly on exit, and return exit code. Update winsize.

Author: shazow

execd.go

@@ -87,14 +87,14 @@ func attachCmd(cmd *exec.Cmd, stdout io.Writer, stderr io.Writer, stdin io.Reade
 	return &wg, nil
 }
 
-func attachShell(cmd *exec.Cmd, stdout io.Writer, stdin io.Reader) (*sync.WaitGroup, error) {
+func attachShell(cmd *exec.Cmd, stdout io.Writer, stdin io.Reader) (*os.File, *sync.WaitGroup, error) {
 	var wg sync.WaitGroup
 	wg.Add(2)
 
 	// Note that pty merges stdout and stderr.
 	cmdPty, err := pty.Start(cmd)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	go func() {
 		io.Copy(stdout, cmdPty)
@@ -105,7 +105,7 @@ func attachShell(cmd *exec.Cmd, stdout io.Writer, stdin io.Reader) (*sync.WaitGr
 		wg.Done()
 	}()
 
-	return &wg, nil
+	return cmdPty, &wg, nil
 }
 
 func addKey(conf *ssh.ServerConfig, block *pem.Block) (err error) {
@@ -270,7 +270,6 @@ func handleChannel(conn *ssh.ServerConn, newChan ssh.NewChannel, execHandler []s
 		log.Println("newChan.Accept failed:", err)
 		return
 	}
-	defer ch.Close()
 
 	assert := func(at string, err error) bool {
 		if err != nil {
@@ -290,11 +289,13 @@ func handleChannel(conn *ssh.ServerConn, newChan ssh.NewChannel, execHandler []s
 		stderr = ch.Stderr()
 	}
 
-	// FIXME: We currently bail out as soon as a req is matched. Technically ssh
-	// can send multiple requests out of band, which we might want to handle?
+	var ptyShell *os.File
+
 	for req := range reqs {
 		switch req.Type {
 		case "exec":
+			defer ch.Close()
+
 			if req.WantReply {
 				req.Reply(true, nil)
 			}
@@ -338,9 +339,8 @@ func handleChannel(conn *ssh.ServerConn, newChan ssh.NewChannel, execHandler []s
 			assert("sendExit", err)
 			return
 		case "pty-req":
-			// TODO: Handle "shell" which technically comes after "pty-req".
-			// TODO: Handle "window-change" for pty
-			// TODO: Parse payload to set initial window dimensions
+			width, height, okSize := parsePtyRequest(req.Payload)
+
 			var cmd *exec.Cmd
 			if *shell {
 				cmd = exec.Command(os.Getenv("SHELL"))
@@ -357,17 +357,29 @@ func handleChannel(conn *ssh.ServerConn, newChan ssh.NewChannel, execHandler []s
 				}
 				cmd.Env = append(cmd.Env, "USER="+conn.Permissions.Extensions["user"])
 			}
-			done, err := attachShell(cmd, stdout, ch)
+			ptyShell, _, err := attachShell(cmd, stdout, ch)
 			if assert("attachShell", err) {
+				ch.Close()
 				return
 			}
-			done.Wait()
-			_, err = exitStatus(cmd.Wait())
-			if assert("exitStatus", err) {
-				return
+			if okSize {
+				setWinsize(ptyShell.Fd(), width, height)
+				req.Reply(true, nil)
+			}
+
+			go func() {
+				status, err := exitStatus(cmd.Wait())
+				if !assert("exitStatus", err) {
+					_, err := ch.SendRequest("exit-status", false, ssh.Marshal(&status))
+					assert("sendExit", err)
+				}
+				ch.Close()
+			}()
+		case "window-change":
+			width, height, okSize := parsePtyRequest(req.Payload)
+			if okSize {
+				setWinsize(ptyShell.Fd(), width, height)
 			}
-			// FIXME: Do we want to send anything back before closing?
-			return
 		}
 
 		if req.WantReply {

pty.go

@@ -0,0 +1,89 @@
+package main
+
+import (
+	"encoding/binary"
+	"syscall"
+	"unsafe"
+)
+
+//// winsize is borrowed from https://github.com/creack/termios/blob/master/win/win.go
+
+// winsize stores the Heighty and Width of a terminal.
+type winsize struct {
+	Height uint16
+	Width  uint16
+	x      uint16 // unused
+	y      uint16 // unused
+}
+
+func setWinsize(fd uintptr, width int, height int) {
+	ws := &winsize{Width: uint16(width), Height: uint16(height)}
+	syscall.Syscall(syscall.SYS_IOCTL, fd, uintptr(syscall.TIOCSWINSZ), uintptr(unsafe.Pointer(ws)))
+}
+
+//// Helpers below are borrowed from go.crypto circa 2011:
+
+// parsePtyRequest parses the payload of the pty-req message and extracts the
+// dimensions of the terminal. See RFC 4254, section 6.2.
+func parsePtyRequest(s []byte) (width, height int, ok bool) {
+	_, s, ok = parseString(s)
+	if !ok {
+		return
+	}
+	width32, s, ok := parseUint32(s)
+	if !ok {
+		return
+	}
+	height32, _, ok := parseUint32(s)
+	width = int(width32)
+	height = int(height32)
+	if width < 1 {
+		ok = false
+	}
+	if height < 1 {
+		ok = false
+	}
+	return
+}
+
+func parseWinchRequest(s []byte) (width, height int, ok bool) {
+	width32, s, ok := parseUint32(s)
+	if !ok {
+		return
+	}
+	height32, s, ok := parseUint32(s)
+	if !ok {
+		return
+	}
+
+	width = int(width32)
+	height = int(height32)
+	if width < 1 {
+		ok = false
+	}
+	if height < 1 {
+		ok = false
+	}
+	return
+}
+
+func parseString(in []byte) (out string, rest []byte, ok bool) {
+	if len(in) < 4 {
+		return
+	}
+	length := binary.BigEndian.Uint32(in)
+	if uint32(len(in)) < 4+length {
+		return
+	}
+	out = string(in[4 : 4+length])
+	rest = in[4+length:]
+	ok = true
+	return
+}
+
+func parseUint32(in []byte) (uint32, []byte, bool) {
+	if len(in) < 4 {
+		return 0, nil, false
+	}
+	return binary.BigEndian.Uint32(in), in[4:], true
+}