Golang pty.Start函数代码示例

// run command from basedir and print output to stdout.
func runCmd(baseDir string, env []string, command string, args ...string) error {
	cmd := exec.Command(command, args...)
	if baseDir != "" {
		cmd.Dir = baseDir
	}
	cmd.Env = env

	// print command being run
	fmt.Println("$", strings.Join(cmd.Args, " "))

	tty, err := pty.Start(cmd)
	if err != nil {
		return err
	}

	scanner := bufio.NewScanner(tty)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		for scanner.Scan() {
			out := scanner.Text()
			fmt.Printf("%s\n", fmtOutput(out))
		}
	}()

	err = cmd.Wait()
	if err != nil {
		return err
	}

	wg.Wait()

	return nil
}

func (this *DockerServer) Run() {
	f, err := pty.Start(this.cmd)
	if err != nil {
		panic(err)
	}
	io.Copy(os.Stdout, f)
}

func establishPty(session *SessionConfig) error {
	defer trace.End(trace.Begin("initializing pty handling for session " + session.ID))

	// TODO: if we want to allow raw output to the log so that subsequent tty enabled
	// processing receives the control characters then we should be binding the PTY
	// during attach, and using the same path we have for non-tty here
	session.wait = &sync.WaitGroup{}

	var err error
	session.Pty, err = pty.Start(&session.Cmd)
	if session.Pty != nil {
		session.wait.Add(1)

		// TODO: do we need to ensure all reads have completed before calling Wait on the process?
		// it frees up all resources - does that mean it frees the output buffers?
		go func() {
			_, gerr := io.Copy(session.Outwriter, session.Pty)
			log.Debugf("PTY stdout copy: %s", gerr)

			session.wait.Done()
		}()
		go func() {
			_, gerr := io.Copy(session.Pty, session.Reader)
			log.Debugf("PTY stdin copy: %s", gerr)
		}()
	}

	return err
}

func (r *Reporter) execHost(req xfer.Request) xfer.Response {
	cmd := exec.Command(r.hostShellCmd[0], r.hostShellCmd[1:]...)
	cmd.Env = []string{"TERM=xterm"}
	ptyPipe, err := pty.Start(cmd)
	if err != nil {
		return xfer.ResponseError(err)
	}

	id, pipe, err := controls.NewPipeFromEnds(nil, ptyPipe, r.pipes, req.AppID)
	if err != nil {
		return xfer.ResponseError(err)
	}
	pipe.OnClose(func() {
		if err := cmd.Process.Kill(); err != nil {
			log.Errorf("Error stopping host shell: %v", err)
		}
		if err := ptyPipe.Close(); err != nil {
			log.Errorf("Error closing host shell's pty: %v", err)
		}
		log.Info("Host shell closed.")
	})
	go func() {
		if err := cmd.Wait(); err != nil {
			log.Errorf("Error waiting on host shell: %v", err)
		}
		pipe.Close()
	}()

	return xfer.Response{
		Pipe:   id,
		RawTTY: true,
	}
}

func (sshConfig *SSHConfig) Login() string {
	// beego.Debug(sshConfig.Agent.LoginName, sshConfig.Agent.Host)
	var err error
	sshConfig.WsPty.Cmd = exec.Command("ssh", "-o", "StrictHostKeyChecking=no", sshConfig.Agent.LoginName+"@"+sshConfig.Agent.Host)
	sshConfig.WsPty.Pty, err = pty.Start(sshConfig.WsPty.Cmd)
	if err != nil {
		beego.Error("Failed to start command: %s\n", err)
	}

	i := 0
	for {
		if i >= 10 {
			beego.Error("login error")
			sshConfig.WsPty.Pty.Close()
			return "login error"
		}
		buf := make([]byte, 1024)
		size, err := sshConfig.WsPty.Pty.Read(buf)
		if err != nil {
			beego.Error("login Read error")
		}

		if !strings.Contains(string([]byte(buf[:size])), "password") {
			i++
			continue
		}

		_, err = sshConfig.WsPty.Pty.Write([]byte(sshConfig.Agent.LoginPass + "\n"))
		if err != nil {
			beego.Error("login Write error")
		}
		return ""
	}
}

// Start will create start cmd and return a PTY wrapping it. The PTY implements
// ReadWriteCloser and must be closed when execution is done.
//
// Note: This will overwrite stdio for cmd.
func Start(cmd *exec.Cmd) (*PTY, error) {
	f, err := pty.Start(cmd)
	if err != nil {
		return nil, err
	}
	return &PTY{f}, nil
}

// regression test for #12546
func (s *DockerSuite) TestExecInteractiveStdinClose(c *check.C) {
	testRequires(c, DaemonIsLinux)
	out, _ := dockerCmd(c, "run", "-itd", "busybox", "/bin/cat")
	contID := strings.TrimSpace(out)

	cmd := exec.Command(dockerBinary, "exec", "-i", contID, "echo", "-n", "hello")
	p, err := pty.Start(cmd)
	if err != nil {
		c.Fatal(err)
	}

	b := bytes.NewBuffer(nil)
	go io.Copy(b, p)

	ch := make(chan error)
	go func() { ch <- cmd.Wait() }()

	select {
	case err := <-ch:
		if err != nil {
			c.Errorf("cmd finished with error %v", err)
		}
		if output := b.String(); strings.TrimSpace(output) != "hello" {
			c.Fatalf("Unexpected output %s", output)
		}
	case <-time.After(1 * time.Second):
		c.Fatal("timed out running docker exec")
	}
}

func (app *App) handleWS(w http.ResponseWriter, r *http.Request) {
	log.Printf("New client connected: %s", r.RemoteAddr)

	if r.Method != "GET" {
		http.Error(w, "Method not allowed", 405)
		return
	}

	conn, err := app.upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Print("Failed to upgrade connection")
		return
	}

	cmd := exec.Command(app.options.Command[0], app.options.Command[1:]...)
	ptyIo, err := pty.Start(cmd)
	if err != nil {
		log.Print("Failed to execute command")
		return
	}
	log.Printf("Command is running for client %s with PID %d", r.RemoteAddr, cmd.Process.Pid)

	context := &clientContext{
		app:        app,
		request:    r,
		connection: conn,
		command:    cmd,
		pty:        ptyIo,
	}

	context.goHandleClient()
}

func ptyExec(test *testing.T, mux *pat.Router) {
	mux.Post("/exec", func(res http.ResponseWriter, req *http.Request) {
		test.Log("got exec")
		conn, rw, err := res.(http.Hijacker).Hijack()
		if err != nil {
			res.WriteHeader(http.StatusInternalServerError)
			res.Write([]byte(err.Error()))
			return
		}
		defer conn.Close()
		script := req.FormValue("cmd")
		if script == "" {
			test.Log("missing script")
			test.FailNow()
		}
		test.Log("executing", script)
		cmd := exec.Command("/bin/bash", "-c", script)
		file, err := pty.Start(cmd)
		if err != nil {
			test.Log(err)
			test.FailNow()
		}
		test.Log("copying from pty")
		go io.Copy(file, rw)
		io.Copy(conn, file)
		test.Log("finished running")
	})
}

/**
Expect:
 - all of the input to come back out, because terminals default to echo mode.
 - then the grep'd string should come out, because the command recieved it and matched.
 - the grep'd string should come out surrounded by the escape codes for color, since grep's auto mode should detect that we're in a terminal.
*/
func TestPtySanity(t *testing.T) {
	assert := assrt.NewAssert(t)

	c := exec.Command("grep", "--color=auto", "bar")
	f, err := pty.Start(c)
	if err != nil {
		t.Fatal(err)
	}

	go func() {
		f.Write([]byte("foo\nbar\nbaz\n"))
		/*
			All of the input must be written in a single call to prevent this test from occationally failing nondeterministically, because:
			grep operates stream-wise and will start printing output before it has all of its input,
			and 3/4ths of the output lines are actually from the terminal operating in echo mode on the same channel.
			So there's actually a race between the terminal itself (somewhere down in kernel land I suppose?) and the output of grep.
		*/
		f.Write([]byte{4}) // EOT
	}()

	outBuffer := new(bytes.Buffer)
	io.Copy(outBuffer, f)
	out := string(outBuffer.Bytes())

	expected := // I have no idea where the CR characters come from.
		"foo\r\n" +
			"bar\r\n" +
			"baz\r\n" +
			"[01;31m[Kbar[m[K\r\n"

	assert.Equal(
		expected,
		out,
	)
}

func (this *SSH_Config) Login() error {
	var err error
	this.Command = exec.Command("ssh", "-o", "StrictHostKeyChecking=no", this.Username+"@"+this.Ip)
	this.Pty, err = pty.Start(this.Command)
	if err != nil {
		return err
	}
	i := 0
	for {
		if i >= 10 {

			return errors.New("login error")
		}
		buf := make([]byte, 1024)
		size, err2 := this.Pty.Read(buf)
		if err2 != nil {
			return err
		}

		if !strings.Contains(string([]byte(buf[:size])), "password") {
			i++
			continue
		}
		this.Pty.Write([]byte(this.Password + "\n"))
		if err != nil {
			panic(err)
		}
		if err != nil {
			return errors.New("login error")
		}
		return nil
	}
}

func (s *RunVSuite) TestStartWithTty(c *check.C) {
	ctrName := "TestStartWithTty"
	spec := defaultTestSpec
	spec.Process.Terminal = true
	spec.Process.Args = []string{"sh"}
	c.Assert(s.addSpec(&spec), checker.IsNil)

	cmdArgs := []string{"--kernel", s.kernelPath, "--initrd", s.initrdPath}
	cmdArgs = append(cmdArgs, "start", "--bundle", s.bundlePath, ctrName)
	cmd := exec.Command(s.binaryPath, cmdArgs...)
	tty, err := pty.Start(cmd)
	c.Assert(err, checker.IsNil)
	defer tty.Close()

	_, err = tty.Write([]byte("uname && sleep 2 && exit\n"))
	c.Assert(err, checker.IsNil)

	chErr := make(chan error)
	go func() {
		chErr <- cmd.Wait()
	}()
	select {
	case err := <-chErr:
		c.Assert(err, checker.IsNil)
	case <-time.After(15 * time.Second):
		c.Fatal("timeout waiting for start to exit")
	}

	buf := make([]byte, 256)
	n, err := tty.Read(buf)
	c.Assert(err, checker.IsNil)
	c.Assert(bytes.Contains(buf, []byte("Linux")), checker.Equals, true, check.Commentf(string(buf[:n])))
}

func (s *DockerSuite) TestExecTTY(c *check.C) {
	testRequires(c, DaemonIsLinux, SameHostDaemon)
	dockerCmd(c, "run", "-d", "--name=test", "busybox", "sh", "-c", "echo hello > /foo && top")

	cmd := exec.Command(dockerBinary, "exec", "-it", "test", "sh")
	p, err := pty.Start(cmd)
	c.Assert(err, checker.IsNil)
	defer p.Close()

	_, err = p.Write([]byte("cat /foo && exit\n"))
	c.Assert(err, checker.IsNil)

	chErr := make(chan error)
	go func() {
		chErr <- cmd.Wait()
	}()
	select {
	case err := <-chErr:
		c.Assert(err, checker.IsNil)
	case <-time.After(3 * time.Second):
		c.Fatal("timeout waiting for exec to exit")
	}

	buf := make([]byte, 256)
	read, err := p.Read(buf)
	c.Assert(err, checker.IsNil)
	c.Assert(bytes.Contains(buf, []byte("hello")), checker.Equals, true, check.Commentf(string(buf[:read])))
}

func pipePty(cmd *exec.Cmd, handler *interactive.ShellHandler) error {
	// Start the shell as TTY
	f, err := pty.Start(cmd)
	if err == nil {
		// Connect pipes (close stderr as tty only has two streams)
		go ioext.CopyAndClose(f, handler.StdinPipe())
		go ioext.CopyAndClose(handler.StdoutPipe(), f)
		go handler.StderrPipe().Close()

		// Start communication
		handler.Communicate(func(cols, rows uint16) error {
			return setTTYSize(f, cols, rows)
		}, func() error {
			if cmd.Process != nil {
				return cmd.Process.Kill()
			}
			return nil
		})
	}
	// If pty wasn't supported for some reason we fall back to normal execution
	if err == pty.ErrUnsupported {
		return pipeCommand(cmd, handler)
	}
	if err != nil {
		handler.Communicate(nil, func() error {
			// If cmd.Start() failed, then we don't have a process, but we start
			// the communication flow anyways.
			if cmd.Process != nil {
				return cmd.Process.Kill()
			}
			return nil
		})
	}
	return err
}

func NewPty(cmd *exec.Cmd) (*Pty, error) {
	cmdPty, err := pty.Start(cmd)
	if err != nil {
		return nil, err
	}
	return &Pty{cmdPty}, nil
}