Refactor request type handlers into a stateful struct.

Author: shazow

execd.go

@@ -7,20 +7,17 @@ import (
 	"errors"
 	"flag"
 	"fmt"
-	"io"
 	"io/ioutil"
 	"log"
 	"net"
 	"os"
 	"os/exec"
 	"path/filepath"
 	"strings"
-	"sync"
 	"syscall"
 
-	"code.google.com/p/go.crypto/ssh"
 	"github.com/flynn/go-shlex"
-	"github.com/kr/pty"
+	"golang.org/x/crypto/ssh"
 )
 
 var host = flag.String("h", "", "host ip to listen on")
@@ -50,27 +47,6 @@ func exitStatus(err error) (exitStatusMsg, error) {
 	return exitStatusMsg{0}, nil
 }
 
-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, nil, err
-	}
-	go func() {
-		io.Copy(stdout, cmdPty)
-		wg.Done()
-	}()
-	go func() {
-		io.Copy(cmdPty, stdin)
-		wg.Done()
-	}()
-
-	return cmdPty, &wg, nil
-}
-
 func addKey(conf *ssh.ServerConfig, block *pem.Block) (err error) {
 	var key interface{}
 	switch block.Type {
@@ -223,135 +199,3 @@ func handleConn(conn net.Conn, conf *ssh.ServerConfig, execHandler []string) {
 		go handleChannel(sshConn, ch, execHandler)
 	}
 }
-
-func handleChannel(conn *ssh.ServerConn, newChan ssh.NewChannel, execHandler []string) {
-	ch, reqs, err := newChan.Accept()
-	if err != nil {
-		log.Println("newChan.Accept failed:", err)
-		return
-	}
-
-	assert := func(at string, err error) bool {
-		if err != nil {
-			log.Printf("%s failed: %s", at, err)
-			ch.Stderr().Write([]byte("Internal error.\n"))
-			return true
-		}
-		return false
-	}
-
-	var stdout, stderr io.Writer
-	if *debug {
-		stdout = io.MultiWriter(ch, os.Stdout)
-		stderr = io.MultiWriter(ch.Stderr(), os.Stdout)
-	} else {
-		stdout = ch
-		stderr = ch.Stderr()
-	}
-
-	var ptyShell *os.File
-
-	for req := range reqs {
-		switch req.Type {
-		case "exec":
-			if req.WantReply {
-				req.Reply(true, nil)
-			}
-
-			cmdline := string(req.Payload[4:])
-			var cmd *exec.Cmd
-			if *shell {
-				shellcmd := flag.Arg(1) + " " + cmdline
-				cmd = exec.Command(os.Getenv("SHELL"), "-c", shellcmd)
-			} else {
-				cmdargs, err := shlex.Split(cmdline)
-				if assert("shlex.Split", err) {
-					ch.Close()
-					return
-				}
-				cmd = exec.Command(execHandler[0], append(execHandler[1:], cmdargs...)...)
-			}
-			if *env {
-				cmd.Env = os.Environ()
-			} else {
-				cmd.Env = []string{}
-			}
-			if conn.Permissions != nil {
-				// Using Permissions.Extensions as a way to get state from PublicKeyCallback
-				if conn.Permissions.Extensions["environ"] != "" {
-					cmd.Env = append(cmd.Env, strings.Split(conn.Permissions.Extensions["environ"], "\n")...)
-				}
-				cmd.Env = append(cmd.Env, "USER="+conn.Permissions.Extensions["user"])
-			}
-			cmd.Env = append(cmd.Env, "SSH_ORIGINAL_COMMAND="+cmdline)
-
-			// cmd.Wait closes the stdin when it's done, so we need to proxy it through a pipe
-			stdinPipe, err := cmd.StdinPipe()
-			if assert("cmd.StdinPipe", err) {
-				ch.Close()
-				return
-			}
-			go io.Copy(stdinPipe, ch)
-
-			cmd.Stdout = stdout
-			cmd.Stderr = stderr
-
-			go func() {
-				status, err := exitStatus(cmd.Run())
-				if !assert("exec run", err) {
-					_, err := ch.SendRequest("exit-status", false, ssh.Marshal(&status))
-					assert("exec exit", err)
-				}
-				ch.Close()
-			}()
-		case "pty-req":
-			width, height, okSize := parsePtyRequest(req.Payload)
-
-			var cmd *exec.Cmd
-			if *shell {
-				cmd = exec.Command(os.Getenv("SHELL"))
-			} else {
-				cmd = exec.Command(execHandler[0], execHandler[1:]...)
-			}
-			if *env {
-				cmd.Env = os.Environ()
-			} else {
-				cmd.Env = []string{}
-			}
-			if conn.Permissions != nil {
-				// Using Permissions.Extensions as a way to get state from PublicKeyCallback
-				if conn.Permissions.Extensions["environ"] != "" {
-					cmd.Env = append(cmd.Env, strings.Split(conn.Permissions.Extensions["environ"], "\n")...)
-				}
-				cmd.Env = append(cmd.Env, "USER="+conn.Permissions.Extensions["user"])
-			}
-			ptyShell, _, err := attachShell(cmd, stdout, ch)
-			if assert("attachShell", err) {
-				ch.Close()
-				return
-			}
-			if okSize {
-				setWinsize(ptyShell.Fd(), width, height)
-				req.Reply(true, nil)
-			}
-
-			go func() {
-				status, err := exitStatus(cmd.Wait())
-				if !assert("pty run", err) {
-					_, err := ch.SendRequest("exit-status", false, ssh.Marshal(&status))
-					assert("pty exit", err)
-				}
-				ch.Close()
-			}()
-		case "window-change":
-			width, height, okSize := parsePtyRequest(req.Payload)
-			if okSize {
-				setWinsize(ptyShell.Fd(), width, height)
-			}
-		}
-
-		if req.WantReply {
-			req.Reply(true, nil)
-		}
-	}
-}