command.go

Documentation: sigs.k8s.io/release-utils/command

     1  /*
     2  Copyright 2019 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package command
    18  
    19  import (
    20  	"bytes"
    21  	"fmt"
    22  	"io"
    23  	"os"
    24  	"os/exec"
    25  	"regexp"
    26  	"strings"
    27  	"sync"
    28  	"syscall"
    29  
    30  	"github.com/sirupsen/logrus"
    31  )
    32  
    33  // A generic command abstraction
    34  type Command struct {
    35  	cmds                         []*command
    36  	stdErrWriters, stdOutWriters []io.Writer
    37  	env                          []string
    38  	verbose                      bool
    39  	filter                       *filter
    40  }
    41  
    42  // The internal command representation
    43  type command struct {
    44  	*exec.Cmd
    45  	pipeWriter *io.PipeWriter
    46  }
    47  
    48  // filter is the internally used struct for filtering command output.
    49  type filter struct {
    50  	regex      *regexp.Regexp
    51  	replaceAll string
    52  }
    53  
    54  // A generic command exit status
    55  type Status struct {
    56  	waitStatus syscall.WaitStatus
    57  	*Stream
    58  }
    59  
    60  // Stream combines standard output and error
    61  type Stream struct { //nolint: errname
    62  	stdOut string
    63  	stdErr string
    64  }
    65  
    66  // Commands is an abstraction over multiple Command structures
    67  type Commands []*Command
    68  
    69  // New creates a new command from the provided arguments.
    70  func New(cmd string, args ...string) *Command {
    71  	return NewWithWorkDir("", cmd, args...)
    72  }
    73  
    74  // NewWithWorkDir creates a new command from the provided workDir and the command
    75  // arguments.
    76  func NewWithWorkDir(workDir, cmd string, args ...string) *Command {
    77  	return &Command{
    78  		cmds: []*command{{
    79  			Cmd:        cmdWithDir(workDir, cmd, args...),
    80  			pipeWriter: nil,
    81  		}},
    82  		stdErrWriters: []io.Writer{},
    83  		stdOutWriters: []io.Writer{},
    84  		verbose:       false,
    85  	}
    86  }
    87  
    88  func cmdWithDir(dir, cmd string, args ...string) *exec.Cmd {
    89  	c := exec.Command(cmd, args...)
    90  	c.Dir = dir
    91  	return c
    92  }
    93  
    94  // Pipe creates a new command where the previous should be piped to
    95  func (c *Command) Pipe(cmd string, args ...string) *Command {
    96  	pipeCmd := cmdWithDir(c.cmds[0].Dir, cmd, args...)
    97  
    98  	reader, writer := io.Pipe()
    99  	c.cmds[len(c.cmds)-1].Stdout = writer
   100  	pipeCmd.Stdin = reader
   101  
   102  	c.cmds = append(c.cmds, &command{
   103  		Cmd:        pipeCmd,
   104  		pipeWriter: writer,
   105  	})
   106  	return c
   107  }
   108  
   109  // Env specifies the environment added to the command. Each entry is of the
   110  // form "key=value". The environment of the current process is being preserved,
   111  // while it is possible to overwrite already existing environment variables.
   112  func (c *Command) Env(env ...string) *Command {
   113  	c.env = append(c.env, env...)
   114  	return c
   115  }
   116  
   117  // Verbose enables verbose output aka printing the command before executing it.
   118  func (c *Command) Verbose() *Command {
   119  	c.verbose = true
   120  	return c
   121  }
   122  
   123  // isVerbose returns true if the command is in verbose mode, either set locally
   124  // or global
   125  func (c *Command) isVerbose() bool {
   126  	return GetGlobalVerbose() || c.verbose
   127  }
   128  
   129  // Add a command with the same working directory as well as verbosity mode.
   130  // Returns a new Commands instance.
   131  func (c *Command) Add(cmd string, args ...string) Commands {
   132  	addCmd := NewWithWorkDir(c.cmds[0].Dir, cmd, args...)
   133  	addCmd.verbose = c.verbose
   134  	addCmd.filter = c.filter
   135  	return Commands{c, addCmd}
   136  }
   137  
   138  // AddWriter can be used to add an additional output (stdout) and error
   139  // (stderr) writer to the command, for example when having the need to log to
   140  // files.
   141  func (c *Command) AddWriter(writer io.Writer) *Command {
   142  	c.AddOutputWriter(writer)
   143  	c.AddErrorWriter(writer)
   144  	return c
   145  }
   146  
   147  // AddErrorWriter can be used to add an additional error (stderr) writer to the
   148  // command, for example when having the need to log to files.
   149  func (c *Command) AddErrorWriter(writer io.Writer) *Command {
   150  	c.stdErrWriters = append(c.stdErrWriters, writer)
   151  	return c
   152  }
   153  
   154  // AddOutputWriter can be used to add an additional output (stdout) writer to
   155  // the command, for example when having the need to log to files.
   156  func (c *Command) AddOutputWriter(writer io.Writer) *Command {
   157  	c.stdOutWriters = append(c.stdOutWriters, writer)
   158  	return c
   159  }
   160  
   161  // Filter adds an output filter regular expression to the command. Every output
   162  // will then be replaced with the string provided by replaceAll.
   163  func (c *Command) Filter(regex, replaceAll string) (*Command, error) {
   164  	filterRegex, err := regexp.Compile(regex)
   165  	if err != nil {
   166  		return nil, fmt.Errorf("compile regular expression: %w", err)
   167  	}
   168  	c.filter = &filter{
   169  		regex:      filterRegex,
   170  		replaceAll: replaceAll,
   171  	}
   172  	return c, nil
   173  }
   174  
   175  // Run starts the command and waits for it to finish. It returns an error if
   176  // the command execution was not possible at all, otherwise the Status.
   177  // This method prints the commands output during execution
   178  func (c *Command) Run() (res *Status, err error) {
   179  	return c.run(true)
   180  }
   181  
   182  // RunSuccessOutput starts the command and waits for it to finish. It returns
   183  // an error if the command execution was not successful, otherwise its output.
   184  func (c *Command) RunSuccessOutput() (output *Stream, err error) {
   185  	res, err := c.run(true)
   186  	if err != nil {
   187  		return nil, err
   188  	}
   189  	if !res.Success() {
   190  		return nil, fmt.Errorf("command %v did not succeed: %v", c.String(), res.Error())
   191  	}
   192  	return res.Stream, nil
   193  }
   194  
   195  // RunSuccess starts the command and waits for it to finish. It returns an
   196  // error if the command execution was not successful.
   197  func (c *Command) RunSuccess() error {
   198  	_, err := c.RunSuccessOutput() //nolint: errcheck
   199  	return err
   200  }
   201  
   202  // String returns a string representation of the full command
   203  func (c *Command) String() string {
   204  	str := []string{}
   205  	for _, x := range c.cmds {
   206  		// Note: the following logic can be replaced with x.String(), which was
   207  		// implemented in go1.13
   208  		b := new(strings.Builder)
   209  		b.WriteString(x.Path)
   210  		for _, a := range x.Args[1:] {
   211  			b.WriteByte(' ')
   212  			b.WriteString(a)
   213  		}
   214  		str = append(str, b.String())
   215  	}
   216  	return strings.Join(str, " | ")
   217  }
   218  
   219  // Run starts the command and waits for it to finish. It returns an error if
   220  // the command execution was not possible at all, otherwise the Status.
   221  // This method does not print the output of the command during its execution.
   222  func (c *Command) RunSilent() (res *Status, err error) {
   223  	return c.run(false)
   224  }
   225  
   226  // RunSilentSuccessOutput starts the command and waits for it to finish. It
   227  // returns an error if the command execution was not successful, otherwise its
   228  // output. This method does not print the output of the command during its
   229  // execution.
   230  func (c *Command) RunSilentSuccessOutput() (output *Stream, err error) {
   231  	res, err := c.run(false)
   232  	if err != nil {
   233  		return nil, err
   234  	}
   235  	if !res.Success() {
   236  		return nil, fmt.Errorf("command %v did not succeed: %w", c.String(), res)
   237  	}
   238  	return res.Stream, nil
   239  }
   240  
   241  // RunSilentSuccess starts the command and waits for it to finish. It returns
   242  // an error if the command execution was not successful. This method does not
   243  // print the output of the command during its execution.
   244  func (c *Command) RunSilentSuccess() error {
   245  	_, err := c.RunSilentSuccessOutput() //nolint: errcheck
   246  	return err
   247  }
   248  
   249  // run is the internal run method
   250  func (c *Command) run(printOutput bool) (res *Status, err error) {
   251  	var runErr error
   252  	stdOutBuffer := &bytes.Buffer{}
   253  	stdErrBuffer := &bytes.Buffer{}
   254  	status := &Status{Stream: &Stream{}}
   255  
   256  	type done struct {
   257  		stdout error
   258  		stderr error
   259  	}
   260  	doneChan := make(chan done, 1)
   261  
   262  	var stdOutWriter io.Writer
   263  	for i, cmd := range c.cmds {
   264  		// Last command handling
   265  		if i+1 == len(c.cmds) {
   266  			stdout, err := cmd.StdoutPipe()
   267  			if err != nil {
   268  				return nil, err
   269  			}
   270  			stderr, err := cmd.StderrPipe()
   271  			if err != nil {
   272  				return nil, err
   273  			}
   274  
   275  			var stdErrWriter io.Writer
   276  			if printOutput {
   277  				stdOutWriter = io.MultiWriter(append(
   278  					[]io.Writer{os.Stdout, stdOutBuffer}, c.stdOutWriters...,
   279  				)...)
   280  				stdErrWriter = io.MultiWriter(append(
   281  					[]io.Writer{os.Stderr, stdErrBuffer}, c.stdErrWriters...,
   282  				)...)
   283  			} else {
   284  				stdOutWriter = stdOutBuffer
   285  				stdErrWriter = stdErrBuffer
   286  			}
   287  			go func() {
   288  				var stdoutErr, stderrErr error
   289  				wg := sync.WaitGroup{}
   290  
   291  				wg.Add(2)
   292  
   293  				filterCopy := func(read io.ReadCloser, write io.Writer) (err error) {
   294  					if c.filter != nil {
   295  						builder := &strings.Builder{}
   296  						_, err = io.Copy(builder, read)
   297  						if err != nil {
   298  							return err
   299  						}
   300  						str := c.filter.regex.ReplaceAllString(
   301  							builder.String(), c.filter.replaceAll,
   302  						)
   303  						_, err = io.Copy(write, strings.NewReader(str))
   304  					} else {
   305  						_, err = io.Copy(write, read)
   306  					}
   307  					return err
   308  				}
   309  
   310  				go func() {
   311  					stdoutErr = filterCopy(stdout, stdOutWriter)
   312  					wg.Done()
   313  				}()
   314  
   315  				go func() {
   316  					stderrErr = filterCopy(stderr, stdErrWriter)
   317  					wg.Done()
   318  				}()
   319  
   320  				wg.Wait()
   321  				doneChan <- done{stdoutErr, stderrErr}
   322  			}()
   323  		}
   324  
   325  		if c.isVerbose() {
   326  			logrus.Infof("+ %s", c.String())
   327  		}
   328  
   329  		cmd.Env = append(os.Environ(), c.env...)
   330  
   331  		if err := cmd.Start(); err != nil {
   332  			return nil, err
   333  		}
   334  
   335  		if i > 0 {
   336  			if err := c.cmds[i-1].Wait(); err != nil {
   337  				return nil, err
   338  			}
   339  		}
   340  
   341  		if cmd.pipeWriter != nil {
   342  			if err := cmd.pipeWriter.Close(); err != nil {
   343  				return nil, err
   344  			}
   345  		}
   346  
   347  		// Wait for last command in the pipe to finish
   348  		if i+1 == len(c.cmds) {
   349  			err := <-doneChan
   350  			if err.stdout != nil && strings.Contains(err.stdout.Error(), os.ErrClosed.Error()) {
   351  				return nil, fmt.Errorf("unable to copy stdout: %w", err.stdout)
   352  			}
   353  			if err.stderr != nil && strings.Contains(err.stderr.Error(), os.ErrClosed.Error()) {
   354  				return nil, fmt.Errorf("unable to copy stderr: %w", err.stderr)
   355  			}
   356  
   357  			runErr = cmd.Wait()
   358  		}
   359  	}
   360  
   361  	status.stdOut = stdOutBuffer.String()
   362  	status.stdErr = stdErrBuffer.String()
   363  
   364  	if exitErr, ok := runErr.(*exec.ExitError); ok {
   365  		if waitStatus, ok := exitErr.Sys().(syscall.WaitStatus); ok {
   366  			status.waitStatus = waitStatus
   367  			return status, nil
   368  		}
   369  	}
   370  
   371  	return status, runErr
   372  }
   373  
   374  // Success returns if a Status was successful
   375  func (s *Status) Success() bool {
   376  	return s.waitStatus.ExitStatus() == 0
   377  }
   378  
   379  // ExitCode returns the exit status of the command status
   380  func (s *Status) ExitCode() int {
   381  	return s.waitStatus.ExitStatus()
   382  }
   383  
   384  // Output returns stdout of the command status
   385  func (s *Stream) Output() string {
   386  	return s.stdOut
   387  }
   388  
   389  // OutputTrimNL returns stdout of the command status with newlines trimmed
   390  // Use only when output is expected to be a single "word", like a version string.
   391  func (s *Stream) OutputTrimNL() string {
   392  	return strings.TrimSpace(s.stdOut)
   393  }
   394  
   395  // Error returns the stderr of the command status
   396  func (s *Stream) Error() string {
   397  	return s.stdErr
   398  }
   399  
   400  // Execute is a convenience function which creates a new Command, executes it
   401  // and evaluates its status.
   402  func Execute(cmd string, args ...string) error {
   403  	status, err := New(cmd, args...).Run()
   404  	if err != nil {
   405  		return fmt.Errorf("command %q is not executable: %w", cmd, err)
   406  	}
   407  	if !status.Success() {
   408  		return fmt.Errorf(
   409  			"command %q did not exit successful (%d)",
   410  			cmd, status.ExitCode(),
   411  		)
   412  	}
   413  	return nil
   414  }
   415  
   416  // Available verifies that the specified `commands` are available within the
   417  // current `$PATH` environment and returns true if so. The function does not
   418  // check for duplicates nor if the provided slice is empty.
   419  func Available(commands ...string) (ok bool) {
   420  	ok = true
   421  	for _, command := range commands {
   422  		if _, err := exec.LookPath(command); err != nil {
   423  			logrus.Warnf("Unable to %v", err)
   424  			ok = false
   425  		}
   426  	}
   427  	return ok
   428  }
   429  
   430  // Add adds another command with the same working directory as well as
   431  // verbosity mode to the Commands.
   432  func (c Commands) Add(cmd string, args ...string) Commands {
   433  	addCmd := NewWithWorkDir(c[0].cmds[0].Dir, cmd, args...)
   434  	addCmd.verbose = c[0].verbose
   435  	addCmd.filter = c[0].filter
   436  	return append(c, addCmd)
   437  }
   438  
   439  // Run executes all commands sequentially and abort if any of those fails.
   440  func (c Commands) Run() (*Status, error) {
   441  	res := &Status{Stream: &Stream{}}
   442  	for _, cmd := range c {
   443  		output, err := cmd.RunSuccessOutput()
   444  		if err != nil {
   445  			return nil, fmt.Errorf("running command %q: %w", cmd.String(), err)
   446  		}
   447  		res.stdOut += "\n" + output.stdOut
   448  		res.stdErr += "\n" + output.stdErr
   449  	}
   450  	return res, nil
   451  }
   452
View as plain text