Golang remote.Cmd類代碼示例

func TestStart(t *testing.T) {
	address := newMockLineServer(t)
	parts := strings.Split(address, ":")

	r := &terraform.InstanceState{
		Ephemeral: terraform.EphemeralState{
			ConnInfo: map[string]string{
				"type":     "ssh",
				"user":     "user",
				"password": "pass",
				"host":     parts[0],
				"port":     parts[1],
				"timeout":  "30s",
			},
		},
	}

	c, err := New(r)
	if err != nil {
		t.Fatalf("error creating communicator: %s", err)
	}

	var cmd remote.Cmd
	stdout := new(bytes.Buffer)
	cmd.Command = "echo foo"
	cmd.Stdout = stdout

	err = c.Start(&cmd)
	if err != nil {
		t.Fatalf("error executing remote command: %s", err)
	}
}

// Start implementation of communicator.Communicator interface
func (c *MockCommunicator) Start(r *remote.Cmd) error {
	if !c.Commands[r.Command] {
		return fmt.Errorf("Command not found!")
	}

	r.SetExited(0)

	return nil
}

func runCommand(shell *winrm.Shell, cmd *winrm.Command, rc *remote.Cmd) {
	defer shell.Close()

	go io.Copy(rc.Stdout, cmd.Stdout)
	go io.Copy(rc.Stderr, cmd.Stderr)

	cmd.Wait()
	rc.SetExited(cmd.ExitCode())
}

// Start implementation of communicator.Communicator interface
func (c *Communicator) Start(cmd *remote.Cmd) error {
	session, err := c.newSession()
	if err != nil {
		return err
	}

	// Setup our session
	session.Stdin = cmd.Stdin
	session.Stdout = cmd.Stdout
	session.Stderr = cmd.Stderr

	if !c.config.noPty {
		// Request a PTY
		termModes := ssh.TerminalModes{
			ssh.ECHO:          0,     // do not echo
			ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
			ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
		}

		if err := session.RequestPty("xterm", 80, 40, termModes); err != nil {
			return err
		}
	}

	log.Printf("starting remote command: %s", cmd.Command)
	err = session.Start(cmd.Command + "\n")
	if err != nil {
		return err
	}

	// Start a goroutine to wait for the session to end and set the
	// exit boolean and status.
	go func() {
		defer session.Close()

		err := session.Wait()
		exitStatus := 0
		if err != nil {
			exitErr, ok := err.(*ssh.ExitError)
			if ok {
				exitStatus = exitErr.ExitStatus()
			}
		}

		log.Printf("remote command exited with '%d': %s", exitStatus, cmd.Command)
		cmd.SetExited(exitStatus)
	}()

	return nil
}

func runCommand(shell *winrm.Shell, cmd *winrm.Command, rc *remote.Cmd) {
	defer shell.Close()

	var wg sync.WaitGroup
	go func() {
		wg.Add(1)
		io.Copy(rc.Stdout, cmd.Stdout)
		wg.Done()
	}()
	go func() {
		wg.Add(1)
		io.Copy(rc.Stderr, cmd.Stderr)
		wg.Done()
	}()

	cmd.Wait()
	wg.Wait()
	rc.SetExited(cmd.ExitCode())
}

func TestStart_KeyFile(t *testing.T) {
	address := newMockLineServer(t)
	parts := strings.Split(address, ":")

	keyFile, err := ioutil.TempFile("", "tf")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	keyFilePath := keyFile.Name()
	keyFile.Write([]byte(testClientPrivateKey))
	keyFile.Close()
	defer os.Remove(keyFilePath)

	r := &terraform.InstanceState{
		Ephemeral: terraform.EphemeralState{
			ConnInfo: map[string]string{
				"type":     "ssh",
				"user":     "user",
				"key_file": keyFilePath,
				"host":     parts[0],
				"port":     parts[1],
				"timeout":  "30s",
			},
		},
	}

	c, err := New(r)
	if err != nil {
		t.Fatalf("error creating communicator: %s", err)
	}

	var cmd remote.Cmd
	stdout := new(bytes.Buffer)
	cmd.Command = "echo foo"
	cmd.Stdout = stdout

	err = c.Start(&cmd)
	if err != nil {
		t.Fatalf("error executing remote command: %s", err)
	}
}

// runScripts is used to copy and execute a set of scripts
func (p *ResourceProvisioner) runScripts(
	o terraform.UIOutput,
	comm communicator.Communicator,
	scripts []io.ReadCloser) error {
	// Wait and retry until we establish the connection
	err := retryFunc(comm.Timeout(), func() error {
		err := comm.Connect(o)
		return err
	})
	if err != nil {
		return err
	}
	defer comm.Disconnect()

	for _, script := range scripts {
		var cmd *remote.Cmd
		outR, outW := io.Pipe()
		errR, errW := io.Pipe()
		outDoneCh := make(chan struct{})
		errDoneCh := make(chan struct{})
		go p.copyOutput(o, outR, outDoneCh)
		go p.copyOutput(o, errR, errDoneCh)

		err = retryFunc(comm.Timeout(), func() error {
			remotePath := comm.ScriptPath()

			if err := comm.UploadScript(remotePath, script); err != nil {
				return fmt.Errorf("Failed to upload script: %v", err)
			}

			cmd = &remote.Cmd{
				Command: remotePath,
				Stdout:  outW,
				Stderr:  errW,
			}
			if err := comm.Start(cmd); err != nil {
				return fmt.Errorf("Error starting script: %v", err)
			}

			// Upload a blank follow up file in the same path to prevent residual
			// script contents from remaining on remote machine
			empty := bytes.NewReader([]byte(""))
			if err := comm.Upload(remotePath, empty); err != nil {
				return fmt.Errorf("Failed to upload empty follow up script: %v", err)
			}

			return nil
		})
		if err == nil {
			cmd.Wait()
			if cmd.ExitStatus != 0 {
				err = fmt.Errorf("Script exited with non-zero exit status: %d", cmd.ExitStatus)
			}
		}

		// Wait for output to clean up
		outW.Close()
		errW.Close()
		<-outDoneCh
		<-errDoneCh

		// If we have an error, return it out now that we've cleaned up
		if err != nil {
			return err
		}
	}

	return nil
}