...

Source file src/golang.conradwood.net/go-easyops/linux/linux.go

Documentation: golang.conradwood.net/go-easyops/linux

     1  /*
     2  Package linux provides methods to execute commands on linux
     3  */
     4  package linux
     5  
     6  import (
     7  	"context"
     8  	"flag"
     9  	"fmt"
    10  	"golang.conradwood.net/go-easyops/auth"
    11  	"golang.conradwood.net/go-easyops/ctx"
    12  	"golang.conradwood.net/go-easyops/errors"
    13  	"io"
    14  	"os"
    15  	"os/exec"
    16  	"strings"
    17  	"sync"
    18  	"time"
    19  )
    20  
    21  const (
    22  	add_serialised_context = false
    23  )
    24  
    25  var (
    26  	cmdLock    sync.Mutex
    27  	curCmd     string
    28  	LogExe     = flag.Bool("ge_debug_exe", false, "debug execution of third party binaries")
    29  	maxRuntime = flag.Duration("ge_default_max_runtime_exe", time.Duration(5)*time.Second, "mĚ€ax_runtime for external binaries")
    30  )
    31  
    32  type linux struct {
    33  	Runtime          time.Duration
    34  	AllowConcurrency bool
    35  	ctx              context.Context
    36  	context_set      bool // if user-supplied context
    37  	envs             []string
    38  	lastcmd          []string
    39  	runforever       bool
    40  	extraFiles       []*os.File
    41  }
    42  
    43  type Linux interface {
    44  	SafelyExecute(cmd []string, stdin io.Reader) (string, error)
    45  	SafelyExecuteWithDir(cmd []string, dir string, stdin io.Reader) (string, error)
    46  	MyIP() string
    47  	SetMaxRuntime(time.Duration)
    48  	SetRunForever() // incompatible with setmaxruntime
    49  	SetAllowConcurrency(bool)
    50  	SetEnvironment([]string)
    51  	AddFileDescriptor(fd int)
    52  }
    53  
    54  func NewWithContext(ctx context.Context) Linux {
    55  	l := New()
    56  	ln := l.(*linux)
    57  	ln.context_set = true
    58  	ln.ctx = ctx
    59  	return l
    60  }
    61  func New() Linux {
    62  	res := &linux{
    63  		Runtime:          *maxRuntime,
    64  		AllowConcurrency: false,
    65  	}
    66  	res.recalc_context_from_timeout()
    67  	return res
    68  }
    69  
    70  func (l *linux) recalc_context_from_timeout() {
    71  	if l.runforever {
    72  		l.ctx = context.Background()
    73  		return
    74  	}
    75  	cb := ctx.NewContextBuilder()
    76  	cb.WithTimeout(l.Runtime)
    77  	l.ctx = cb.ContextWithAutoCancel()
    78  }
    79  
    80  // execute a command...
    81  // print stdout/err (so it ends up in the logs)
    82  // also we add a timeout - if program hangs we return an error
    83  // rather than 'hanging' forever
    84  // and we use a low-level lock to avoid calling binaries at the same time
    85  func (l *linux) SafelyExecute(cmd []string, stdin io.Reader) (string, error) {
    86  	return l.SafelyExecuteWithDir(cmd, "", stdin)
    87  }
    88  
    89  /*
    90  execute a command within a working directory
    91  */
    92  func (l *linux) SafelyExecuteWithDir(cmd []string, dir string, stdin io.Reader) (string, error) {
    93  	// avoid possible segfaults (afterall it's called 'safely...')
    94  	if len(cmd) == 0 {
    95  		return "", errors.Errorf("no command specified for execute.")
    96  	}
    97  	l.lastcmd = cmd
    98  	if !l.AllowConcurrency {
    99  		if curCmd != "" {
   100  			if *LogExe {
   101  				fmt.Printf("Waiting for %s to complete...\n", curCmd)
   102  			}
   103  		}
   104  		cmdLock.Lock()
   105  		defer cmdLock.Unlock()
   106  	}
   107  	curCmd = cmd[0]
   108  	if curCmd == "sudo" {
   109  		if len(curCmd) < 2 {
   110  			return "", errors.Errorf("sudo without parameters not allowed")
   111  		}
   112  		curCmd = cmd[1]
   113  	}
   114  	// execute
   115  	if *LogExe {
   116  		fmt.Printf("[go-easyops] preparing to execute below command:\n%s\n", l.ComWithParas())
   117  	}
   118  	c := exec.CommandContext(l.ctx, cmd[0], cmd[1:]...)
   119  	if dir != "" {
   120  		c.Dir = dir
   121  	}
   122  	if stdin != nil {
   123  		c.Stdin = stdin
   124  	}
   125  	// set environment
   126  	c.ExtraFiles = l.extraFiles
   127  	c.Env = os.Environ()
   128  	l.env(c)
   129  	output, err := l.syncExecute(c, l.Runtime, !l.runforever)
   130  	if *LogExe {
   131  		printOutput(l.ComName(), output)
   132  	}
   133  	curCmd = ""
   134  	if err != nil {
   135  		fmt.Printf("[go-easyops] ---- %s -----\n%s\n---- end output----\n", strings.Join(cmd, " "), output)
   136  		return output, errors.Wrap(err)
   137  	}
   138  	return output, nil
   139  }
   140  
   141  // execute with timeout.
   142  // sends SIGKILL to process on timeout and returns error
   143  func (l *linux) syncExecute(c *exec.Cmd, timeout time.Duration, hastimeout bool) (string, error) {
   144  	running := false
   145  	killed := false
   146  	if hastimeout {
   147  		timer1 := time.NewTimer(timeout)
   148  		go func() {
   149  			<-timer1.C
   150  			if running {
   151  				if c.Process == nil {
   152  					fmt.Printf("[go-easyops] no process to kill after %0.2fs\n", timeout.Seconds())
   153  					return
   154  				}
   155  				if !running {
   156  					return
   157  				}
   158  				c.Process.Kill()
   159  				killed = true
   160  				if *LogExe {
   161  					fmt.Printf("[go-easyops] process killed after %0.2fs\n", timeout.Seconds())
   162  				}
   163  			}
   164  		}()
   165  	}
   166  	// racecondition - timer might expire between
   167  	// setting flag and starting process.
   168  	// (if timer is really short)
   169  	running = true
   170  	if *LogExe {
   171  		fmt.Printf("[go-easyops] executing command %s (timeout=%0.2fs)\n", l.ComName(), timeout.Seconds())
   172  	}
   173  	b, err := c.CombinedOutput()
   174  	if *LogExe {
   175  		fmt.Printf("[go-easyops] process terminated\n")
   176  	}
   177  	running = false
   178  	if killed {
   179  		err = fmt.Errorf("Process killed after %0.2f seconds", timeout.Seconds())
   180  	}
   181  	return string(b), err
   182  }
   183  
   184  func printOutput(cmd string, output string) {
   185  	fmt.Printf("====BEGIN OUTPUT OF %s====\n", cmd)
   186  	fmt.Printf("%s\n", output)
   187  	fmt.Printf("====END OUTPUT OF %s====\n", cmd)
   188  }
   189  func (l *linux) SetEnvironment(sx []string) {
   190  	l.envs = sx
   191  }
   192  func (l *linux) SetRunForever() {
   193  	l.runforever = true
   194  	if !l.context_set {
   195  		l.recalc_context_from_timeout()
   196  	}
   197  }
   198  func (l *linux) SetMaxRuntime(d time.Duration) {
   199  	l.runforever = false
   200  	l.Runtime = d
   201  	if !l.context_set {
   202  		l.recalc_context_from_timeout()
   203  	}
   204  }
   205  func (l *linux) SetAllowConcurrency(b bool) {
   206  	l.AllowConcurrency = b
   207  }
   208  
   209  // normally all filedescriptors except stdin, stdout and stderr are closed.
   210  // here we can give an additional one to pass on to the child
   211  func (l *linux) AddFileDescriptor(fd int) {
   212  	l.extraFiles = append(l.extraFiles, os.NewFile(uintptr(fd), "addfile"))
   213  }
   214  
   215  // add context to environment
   216  func (l *linux) env(c *exec.Cmd) error {
   217  	if l.context_set {
   218  		nc, err := auth.SerialiseContextToString(l.ctx)
   219  		if err != nil {
   220  			return err
   221  		}
   222  		ncs := fmt.Sprintf("GE_CTX=%s", nc)
   223  		for i, e := range c.Env {
   224  			if strings.HasPrefix(e, "GE_CTX=") {
   225  				c.Env[i] = ncs
   226  				return nil
   227  			}
   228  		}
   229  		c.Env = append(c.Env, ncs)
   230  	}
   231  	for _, e := range l.envs {
   232  		c.Env = append(c.Env, e)
   233  	}
   234  	return nil
   235  }
   236  
   237  func (l *linux) ComWithParas() string {
   238  	if len(l.lastcmd) == 0 {
   239  		return "<no command executed>"
   240  	}
   241  	return strings.Join(l.lastcmd, " ")
   242  }
   243  func (l *linux) ComName() string {
   244  	if len(l.lastcmd) == 0 || l.lastcmd[0] == "" {
   245  		return "<no command executed>"
   246  	}
   247  	s := l.lastcmd[0]
   248  	if strings.Contains(s, "sudo") && len(l.lastcmd) > 1 {
   249  		return "sudo " + l.lastcmd[1]
   250  	}
   251  	return l.lastcmd[0]
   252  }
   253  

View as plain text