mirror of https://github.com/hashicorp/consul
387 lines
8.7 KiB
Go
387 lines
8.7 KiB
Go
package common
|
|
|
|
//
|
|
// gopsutil is a port of psutil(http://pythonhosted.org/psutil/).
|
|
// This covers these architectures.
|
|
// - linux (amd64, arm)
|
|
// - freebsd (amd64)
|
|
// - windows (amd64)
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"errors"
|
|
"io/ioutil"
|
|
"log"
|
|
"net/url"
|
|
"os"
|
|
"os/exec"
|
|
"path"
|
|
"path/filepath"
|
|
"reflect"
|
|
"runtime"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
var (
|
|
Timeout = 3 * time.Second
|
|
ErrTimeout = errors.New("Command timed out.")
|
|
)
|
|
|
|
type Invoker interface {
|
|
Command(string, ...string) ([]byte, error)
|
|
}
|
|
|
|
type Invoke struct{}
|
|
|
|
func (i Invoke) Command(name string, arg ...string) ([]byte, error) {
|
|
cmd := exec.Command(name, arg...)
|
|
return CombinedOutputTimeout(cmd, Timeout)
|
|
}
|
|
|
|
type FakeInvoke struct {
|
|
CommandExpectedDir string // CommandExpectedDir specifies dir which includes expected outputs.
|
|
Suffix string // Suffix species expected file name suffix such as "fail"
|
|
Error error // If Error specfied, return the error.
|
|
}
|
|
|
|
// Command in FakeInvoke returns from expected file if exists.
|
|
func (i FakeInvoke) Command(name string, arg ...string) ([]byte, error) {
|
|
if i.Error != nil {
|
|
return []byte{}, i.Error
|
|
}
|
|
|
|
arch := runtime.GOOS
|
|
|
|
fname := strings.Join(append([]string{name}, arg...), "")
|
|
fname = url.QueryEscape(fname)
|
|
var dir string
|
|
if i.CommandExpectedDir == "" {
|
|
dir = "expected"
|
|
} else {
|
|
dir = i.CommandExpectedDir
|
|
}
|
|
fpath := path.Join(dir, arch, fname)
|
|
if i.Suffix != "" {
|
|
fpath += "_" + i.Suffix
|
|
}
|
|
if PathExists(fpath) {
|
|
return ioutil.ReadFile(fpath)
|
|
}
|
|
return exec.Command(name, arg...).Output()
|
|
}
|
|
|
|
var ErrNotImplementedError = errors.New("not implemented yet")
|
|
|
|
// ReadLines reads contents from a file and splits them by new lines.
|
|
// A convenience wrapper to ReadLinesOffsetN(filename, 0, -1).
|
|
func ReadLines(filename string) ([]string, error) {
|
|
return ReadLinesOffsetN(filename, 0, -1)
|
|
}
|
|
|
|
// ReadLines reads contents from file and splits them by new line.
|
|
// The offset tells at which line number to start.
|
|
// The count determines the number of lines to read (starting from offset):
|
|
// n >= 0: at most n lines
|
|
// n < 0: whole file
|
|
func ReadLinesOffsetN(filename string, offset uint, n int) ([]string, error) {
|
|
f, err := os.Open(filename)
|
|
if err != nil {
|
|
return []string{""}, err
|
|
}
|
|
defer f.Close()
|
|
|
|
var ret []string
|
|
|
|
r := bufio.NewReader(f)
|
|
for i := 0; i < n+int(offset) || n < 0; i++ {
|
|
line, err := r.ReadString('\n')
|
|
if err != nil {
|
|
break
|
|
}
|
|
if i < int(offset) {
|
|
continue
|
|
}
|
|
ret = append(ret, strings.Trim(line, "\n"))
|
|
}
|
|
|
|
return ret, nil
|
|
}
|
|
|
|
func IntToString(orig []int8) string {
|
|
ret := make([]byte, len(orig))
|
|
size := -1
|
|
for i, o := range orig {
|
|
if o == 0 {
|
|
size = i
|
|
break
|
|
}
|
|
ret[i] = byte(o)
|
|
}
|
|
if size == -1 {
|
|
size = len(orig)
|
|
}
|
|
|
|
return string(ret[0:size])
|
|
}
|
|
|
|
func UintToString(orig []uint8) string {
|
|
ret := make([]byte, len(orig))
|
|
size := -1
|
|
for i, o := range orig {
|
|
if o == 0 {
|
|
size = i
|
|
break
|
|
}
|
|
ret[i] = byte(o)
|
|
}
|
|
if size == -1 {
|
|
size = len(orig)
|
|
}
|
|
|
|
return string(ret[0:size])
|
|
}
|
|
|
|
func ByteToString(orig []byte) string {
|
|
n := -1
|
|
l := -1
|
|
for i, b := range orig {
|
|
// skip left side null
|
|
if l == -1 && b == 0 {
|
|
continue
|
|
}
|
|
if l == -1 {
|
|
l = i
|
|
}
|
|
|
|
if b == 0 {
|
|
break
|
|
}
|
|
n = i + 1
|
|
}
|
|
if n == -1 {
|
|
return string(orig)
|
|
}
|
|
return string(orig[l:n])
|
|
}
|
|
|
|
// ReadInts reads contents from single line file and returns them as []int32.
|
|
func ReadInts(filename string) ([]int64, error) {
|
|
f, err := os.Open(filename)
|
|
if err != nil {
|
|
return []int64{}, err
|
|
}
|
|
defer f.Close()
|
|
|
|
var ret []int64
|
|
|
|
r := bufio.NewReader(f)
|
|
|
|
// The int files that this is concerned with should only be one liners.
|
|
line, err := r.ReadString('\n')
|
|
if err != nil {
|
|
return []int64{}, err
|
|
}
|
|
|
|
i, err := strconv.ParseInt(strings.Trim(line, "\n"), 10, 32)
|
|
if err != nil {
|
|
return []int64{}, err
|
|
}
|
|
ret = append(ret, i)
|
|
|
|
return ret, nil
|
|
}
|
|
|
|
// Parse to int32 without error
|
|
func mustParseInt32(val string) int32 {
|
|
vv, _ := strconv.ParseInt(val, 10, 32)
|
|
return int32(vv)
|
|
}
|
|
|
|
// Parse to uint64 without error
|
|
func mustParseUint64(val string) uint64 {
|
|
vv, _ := strconv.ParseInt(val, 10, 64)
|
|
return uint64(vv)
|
|
}
|
|
|
|
// Parse to Float64 without error
|
|
func mustParseFloat64(val string) float64 {
|
|
vv, _ := strconv.ParseFloat(val, 64)
|
|
return vv
|
|
}
|
|
|
|
// StringsHas checks the target string slice contains src or not
|
|
func StringsHas(target []string, src string) bool {
|
|
for _, t := range target {
|
|
if strings.TrimSpace(t) == src {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
// StringsContains checks the src in any string of the target string slice
|
|
func StringsContains(target []string, src string) bool {
|
|
for _, t := range target {
|
|
if strings.Contains(t, src) {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
// IntContains checks the src in any int of the target int slice.
|
|
func IntContains(target []int, src int) bool {
|
|
for _, t := range target {
|
|
if src == t {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
// get struct attributes.
|
|
// This method is used only for debugging platform dependent code.
|
|
func attributes(m interface{}) map[string]reflect.Type {
|
|
typ := reflect.TypeOf(m)
|
|
if typ.Kind() == reflect.Ptr {
|
|
typ = typ.Elem()
|
|
}
|
|
|
|
attrs := make(map[string]reflect.Type)
|
|
if typ.Kind() != reflect.Struct {
|
|
return nil
|
|
}
|
|
|
|
for i := 0; i < typ.NumField(); i++ {
|
|
p := typ.Field(i)
|
|
if !p.Anonymous {
|
|
attrs[p.Name] = p.Type
|
|
}
|
|
}
|
|
|
|
return attrs
|
|
}
|
|
|
|
func PathExists(filename string) bool {
|
|
if _, err := os.Stat(filename); err == nil {
|
|
return true
|
|
}
|
|
return false
|
|
}
|
|
|
|
//GetEnv retrieves the environment variable key. If it does not exist it returns the default.
|
|
func GetEnv(key string, dfault string, combineWith ...string) string {
|
|
value := os.Getenv(key)
|
|
if value == "" {
|
|
value = dfault
|
|
}
|
|
|
|
switch len(combineWith) {
|
|
case 0:
|
|
return value
|
|
case 1:
|
|
return filepath.Join(value, combineWith[0])
|
|
default:
|
|
all := make([]string, len(combineWith)+1)
|
|
all[0] = value
|
|
copy(all[1:], combineWith)
|
|
return filepath.Join(all...)
|
|
}
|
|
panic("invalid switch case")
|
|
}
|
|
|
|
func HostProc(combineWith ...string) string {
|
|
return GetEnv("HOST_PROC", "/proc", combineWith...)
|
|
}
|
|
|
|
func HostSys(combineWith ...string) string {
|
|
return GetEnv("HOST_SYS", "/sys", combineWith...)
|
|
}
|
|
|
|
func HostEtc(combineWith ...string) string {
|
|
return GetEnv("HOST_ETC", "/etc", combineWith...)
|
|
}
|
|
|
|
// CombinedOutputTimeout runs the given command with the given timeout and
|
|
// returns the combined output of stdout and stderr.
|
|
// If the command times out, it attempts to kill the process.
|
|
// copied from https://github.com/influxdata/telegraf
|
|
func CombinedOutputTimeout(c *exec.Cmd, timeout time.Duration) ([]byte, error) {
|
|
var b bytes.Buffer
|
|
c.Stdout = &b
|
|
c.Stderr = &b
|
|
if err := c.Start(); err != nil {
|
|
return nil, err
|
|
}
|
|
err := WaitTimeout(c, timeout)
|
|
return b.Bytes(), err
|
|
}
|
|
|
|
// WaitTimeout waits for the given command to finish with a timeout.
|
|
// It assumes the command has already been started.
|
|
// If the command times out, it attempts to kill the process.
|
|
// copied from https://github.com/influxdata/telegraf
|
|
func WaitTimeout(c *exec.Cmd, timeout time.Duration) error {
|
|
timer := time.NewTimer(timeout)
|
|
done := make(chan error)
|
|
go func() { done <- c.Wait() }()
|
|
select {
|
|
case err := <-done:
|
|
timer.Stop()
|
|
return err
|
|
case <-timer.C:
|
|
if err := c.Process.Kill(); err != nil {
|
|
log.Printf("FATAL error killing process: %s", err)
|
|
return err
|
|
}
|
|
// wait for the command to return after killing it
|
|
<-done
|
|
return ErrTimeout
|
|
}
|
|
}
|
|
|
|
// https://gist.github.com/kylelemons/1525278
|
|
func Pipeline(cmds ...*exec.Cmd) ([]byte, []byte, error) {
|
|
// Require at least one command
|
|
if len(cmds) < 1 {
|
|
return nil, nil, nil
|
|
}
|
|
|
|
// Collect the output from the command(s)
|
|
var output bytes.Buffer
|
|
var stderr bytes.Buffer
|
|
|
|
last := len(cmds) - 1
|
|
for i, cmd := range cmds[:last] {
|
|
var err error
|
|
// Connect each command's stdin to the previous command's stdout
|
|
if cmds[i+1].Stdin, err = cmd.StdoutPipe(); err != nil {
|
|
return nil, nil, err
|
|
}
|
|
// Connect each command's stderr to a buffer
|
|
cmd.Stderr = &stderr
|
|
}
|
|
|
|
// Connect the output and error for the last command
|
|
cmds[last].Stdout, cmds[last].Stderr = &output, &stderr
|
|
|
|
// Start each command
|
|
for _, cmd := range cmds {
|
|
if err := cmd.Start(); err != nil {
|
|
return output.Bytes(), stderr.Bytes(), err
|
|
}
|
|
}
|
|
|
|
// Wait for each command to complete
|
|
for _, cmd := range cmds {
|
|
if err := cmd.Wait(); err != nil {
|
|
return output.Bytes(), stderr.Bytes(), err
|
|
}
|
|
}
|
|
|
|
// Return the pipeline output and the collected standard error
|
|
return output.Bytes(), stderr.Bytes(), nil
|
|
}
|