common.go 8.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392
  1. package common
  2. //
  3. // gopsutil is a port of psutil(http://pythonhosted.org/psutil/).
  4. // This covers these architectures.
  5. // - linux (amd64, arm)
  6. // - freebsd (amd64)
  7. // - windows (amd64)
  8. import (
  9. "bufio"
  10. "bytes"
  11. "context"
  12. "errors"
  13. "fmt"
  14. "io/ioutil"
  15. "net/url"
  16. "os"
  17. "os/exec"
  18. "path"
  19. "path/filepath"
  20. "reflect"
  21. "runtime"
  22. "strconv"
  23. "strings"
  24. "time"
  25. )
  26. var (
  27. Timeout = 3 * time.Second
  28. ErrTimeout = errors.New("command timed out")
  29. )
  30. type Invoker interface {
  31. Command(string, ...string) ([]byte, error)
  32. CommandWithContext(context.Context, string, ...string) ([]byte, error)
  33. }
  34. type Invoke struct{}
  35. func (i Invoke) Command(name string, arg ...string) ([]byte, error) {
  36. ctx, cancel := context.WithTimeout(context.Background(), Timeout)
  37. defer cancel()
  38. return i.CommandWithContext(ctx, name, arg...)
  39. }
  40. func (i Invoke) CommandWithContext(ctx context.Context, name string, arg ...string) ([]byte, error) {
  41. cmd := exec.CommandContext(ctx, name, arg...)
  42. var buf bytes.Buffer
  43. cmd.Stdout = &buf
  44. cmd.Stderr = &buf
  45. if err := cmd.Start(); err != nil {
  46. return buf.Bytes(), err
  47. }
  48. if err := cmd.Wait(); err != nil {
  49. return buf.Bytes(), err
  50. }
  51. return buf.Bytes(), nil
  52. }
  53. type FakeInvoke struct {
  54. Suffix string // Suffix species expected file name suffix such as "fail"
  55. Error error // If Error specfied, return the error.
  56. }
  57. // Command in FakeInvoke returns from expected file if exists.
  58. func (i FakeInvoke) Command(name string, arg ...string) ([]byte, error) {
  59. if i.Error != nil {
  60. return []byte{}, i.Error
  61. }
  62. arch := runtime.GOOS
  63. commandName := filepath.Base(name)
  64. fname := strings.Join(append([]string{commandName}, arg...), "")
  65. fname = url.QueryEscape(fname)
  66. fpath := path.Join("testdata", arch, fname)
  67. if i.Suffix != "" {
  68. fpath += "_" + i.Suffix
  69. }
  70. if PathExists(fpath) {
  71. return ioutil.ReadFile(fpath)
  72. }
  73. return []byte{}, fmt.Errorf("could not find testdata: %s", fpath)
  74. }
  75. func (i FakeInvoke) CommandWithContext(ctx context.Context, name string, arg ...string) ([]byte, error) {
  76. return i.Command(name, arg...)
  77. }
  78. var ErrNotImplementedError = errors.New("not implemented yet")
  79. // ReadLines reads contents from a file and splits them by new lines.
  80. // A convenience wrapper to ReadLinesOffsetN(filename, 0, -1).
  81. func ReadLines(filename string) ([]string, error) {
  82. return ReadLinesOffsetN(filename, 0, -1)
  83. }
  84. // ReadLines reads contents from file and splits them by new line.
  85. // The offset tells at which line number to start.
  86. // The count determines the number of lines to read (starting from offset):
  87. // n >= 0: at most n lines
  88. // n < 0: whole file
  89. func ReadLinesOffsetN(filename string, offset uint, n int) ([]string, error) {
  90. f, err := os.Open(filename)
  91. if err != nil {
  92. return []string{""}, err
  93. }
  94. defer f.Close()
  95. var ret []string
  96. r := bufio.NewReader(f)
  97. for i := 0; i < n+int(offset) || n < 0; i++ {
  98. line, err := r.ReadString('\n')
  99. if err != nil {
  100. break
  101. }
  102. if i < int(offset) {
  103. continue
  104. }
  105. ret = append(ret, strings.Trim(line, "\n"))
  106. }
  107. return ret, nil
  108. }
  109. func IntToString(orig []int8) string {
  110. ret := make([]byte, len(orig))
  111. size := -1
  112. for i, o := range orig {
  113. if o == 0 {
  114. size = i
  115. break
  116. }
  117. ret[i] = byte(o)
  118. }
  119. if size == -1 {
  120. size = len(orig)
  121. }
  122. return string(ret[0:size])
  123. }
  124. func UintToString(orig []uint8) string {
  125. ret := make([]byte, len(orig))
  126. size := -1
  127. for i, o := range orig {
  128. if o == 0 {
  129. size = i
  130. break
  131. }
  132. ret[i] = byte(o)
  133. }
  134. if size == -1 {
  135. size = len(orig)
  136. }
  137. return string(ret[0:size])
  138. }
  139. func ByteToString(orig []byte) string {
  140. n := -1
  141. l := -1
  142. for i, b := range orig {
  143. // skip left side null
  144. if l == -1 && b == 0 {
  145. continue
  146. }
  147. if l == -1 {
  148. l = i
  149. }
  150. if b == 0 {
  151. break
  152. }
  153. n = i + 1
  154. }
  155. if n == -1 {
  156. return string(orig)
  157. }
  158. return string(orig[l:n])
  159. }
  160. // ReadInts reads contents from single line file and returns them as []int32.
  161. func ReadInts(filename string) ([]int64, error) {
  162. f, err := os.Open(filename)
  163. if err != nil {
  164. return []int64{}, err
  165. }
  166. defer f.Close()
  167. var ret []int64
  168. r := bufio.NewReader(f)
  169. // The int files that this is concerned with should only be one liners.
  170. line, err := r.ReadString('\n')
  171. if err != nil {
  172. return []int64{}, err
  173. }
  174. i, err := strconv.ParseInt(strings.Trim(line, "\n"), 10, 32)
  175. if err != nil {
  176. return []int64{}, err
  177. }
  178. ret = append(ret, i)
  179. return ret, nil
  180. }
  181. // Parse to int32 without error
  182. func mustParseInt32(val string) int32 {
  183. vv, _ := strconv.ParseInt(val, 10, 32)
  184. return int32(vv)
  185. }
  186. // Parse to uint64 without error
  187. func mustParseUint64(val string) uint64 {
  188. vv, _ := strconv.ParseInt(val, 10, 64)
  189. return uint64(vv)
  190. }
  191. // Parse to Float64 without error
  192. func mustParseFloat64(val string) float64 {
  193. vv, _ := strconv.ParseFloat(val, 64)
  194. return vv
  195. }
  196. // StringsHas checks the target string slice contains src or not
  197. func StringsHas(target []string, src string) bool {
  198. for _, t := range target {
  199. if strings.TrimSpace(t) == src {
  200. return true
  201. }
  202. }
  203. return false
  204. }
  205. // StringsContains checks the src in any string of the target string slice
  206. func StringsContains(target []string, src string) bool {
  207. for _, t := range target {
  208. if strings.Contains(t, src) {
  209. return true
  210. }
  211. }
  212. return false
  213. }
  214. // IntContains checks the src in any int of the target int slice.
  215. func IntContains(target []int, src int) bool {
  216. for _, t := range target {
  217. if src == t {
  218. return true
  219. }
  220. }
  221. return false
  222. }
  223. // get struct attributes.
  224. // This method is used only for debugging platform dependent code.
  225. func attributes(m interface{}) map[string]reflect.Type {
  226. typ := reflect.TypeOf(m)
  227. if typ.Kind() == reflect.Ptr {
  228. typ = typ.Elem()
  229. }
  230. attrs := make(map[string]reflect.Type)
  231. if typ.Kind() != reflect.Struct {
  232. return nil
  233. }
  234. for i := 0; i < typ.NumField(); i++ {
  235. p := typ.Field(i)
  236. if !p.Anonymous {
  237. attrs[p.Name] = p.Type
  238. }
  239. }
  240. return attrs
  241. }
  242. func PathExists(filename string) bool {
  243. if _, err := os.Stat(filename); err == nil {
  244. return true
  245. }
  246. return false
  247. }
  248. //GetEnv retrieves the environment variable key. If it does not exist it returns the default.
  249. func GetEnv(key string, dfault string, combineWith ...string) string {
  250. value := os.Getenv(key)
  251. if value == "" {
  252. value = dfault
  253. }
  254. switch len(combineWith) {
  255. case 0:
  256. return value
  257. case 1:
  258. return filepath.Join(value, combineWith[0])
  259. default:
  260. all := make([]string, len(combineWith)+1)
  261. all[0] = value
  262. copy(all[1:], combineWith)
  263. return filepath.Join(all...)
  264. }
  265. panic("invalid switch case")
  266. }
  267. func HostProc(combineWith ...string) string {
  268. return GetEnv("HOST_PROC", "/proc", combineWith...)
  269. }
  270. func HostSys(combineWith ...string) string {
  271. return GetEnv("HOST_SYS", "/sys", combineWith...)
  272. }
  273. func HostEtc(combineWith ...string) string {
  274. return GetEnv("HOST_ETC", "/etc", combineWith...)
  275. }
  276. func HostVar(combineWith ...string) string {
  277. return GetEnv("HOST_VAR", "/var", combineWith...)
  278. }
  279. func HostRun(combineWith ...string) string {
  280. return GetEnv("HOST_RUN", "/run", combineWith...)
  281. }
  282. // https://gist.github.com/kylelemons/1525278
  283. func Pipeline(cmds ...*exec.Cmd) ([]byte, []byte, error) {
  284. // Require at least one command
  285. if len(cmds) < 1 {
  286. return nil, nil, nil
  287. }
  288. // Collect the output from the command(s)
  289. var output bytes.Buffer
  290. var stderr bytes.Buffer
  291. last := len(cmds) - 1
  292. for i, cmd := range cmds[:last] {
  293. var err error
  294. // Connect each command's stdin to the previous command's stdout
  295. if cmds[i+1].Stdin, err = cmd.StdoutPipe(); err != nil {
  296. return nil, nil, err
  297. }
  298. // Connect each command's stderr to a buffer
  299. cmd.Stderr = &stderr
  300. }
  301. // Connect the output and error for the last command
  302. cmds[last].Stdout, cmds[last].Stderr = &output, &stderr
  303. // Start each command
  304. for _, cmd := range cmds {
  305. if err := cmd.Start(); err != nil {
  306. return output.Bytes(), stderr.Bytes(), err
  307. }
  308. }
  309. // Wait for each command to complete
  310. for _, cmd := range cmds {
  311. if err := cmd.Wait(); err != nil {
  312. return output.Bytes(), stderr.Bytes(), err
  313. }
  314. }
  315. // Return the pipeline output and the collected standard error
  316. return output.Bytes(), stderr.Bytes(), nil
  317. }
  318. // getSysctrlEnv sets LC_ALL=C in a list of env vars for use when running
  319. // sysctl commands (see DoSysctrl).
  320. func getSysctrlEnv(env []string) []string {
  321. foundLC := false
  322. for i, line := range env {
  323. if strings.HasPrefix(line, "LC_ALL") {
  324. env[i] = "LC_ALL=C"
  325. foundLC = true
  326. }
  327. }
  328. if !foundLC {
  329. env = append(env, "LC_ALL=C")
  330. }
  331. return env
  332. }