common.go 7.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379
  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. // ReadFile reads contents from a file.
  80. func ReadFile(filename string) (string, error) {
  81. content, err := ioutil.ReadFile(filename)
  82. if err != nil {
  83. return "", err
  84. }
  85. return string(content), nil
  86. }
  87. // ReadLines reads contents from a file and splits them by new lines.
  88. // A convenience wrapper to ReadLinesOffsetN(filename, 0, -1).
  89. func ReadLines(filename string) ([]string, error) {
  90. return ReadLinesOffsetN(filename, 0, -1)
  91. }
  92. // ReadLinesOffsetN reads contents from file and splits them by new line.
  93. // The offset tells at which line number to start.
  94. // The count determines the number of lines to read (starting from offset):
  95. // n >= 0: at most n lines
  96. // n < 0: whole file
  97. func ReadLinesOffsetN(filename string, offset uint, n int) ([]string, error) {
  98. f, err := os.Open(filename)
  99. if err != nil {
  100. return []string{""}, err
  101. }
  102. defer f.Close()
  103. var ret []string
  104. r := bufio.NewReader(f)
  105. for i := 0; i < n+int(offset) || n < 0; i++ {
  106. line, err := r.ReadString('\n')
  107. if err != nil {
  108. break
  109. }
  110. if i < int(offset) {
  111. continue
  112. }
  113. ret = append(ret, strings.Trim(line, "\n"))
  114. }
  115. return ret, nil
  116. }
  117. func IntToString(orig []int8) string {
  118. ret := make([]byte, len(orig))
  119. size := -1
  120. for i, o := range orig {
  121. if o == 0 {
  122. size = i
  123. break
  124. }
  125. ret[i] = byte(o)
  126. }
  127. if size == -1 {
  128. size = len(orig)
  129. }
  130. return string(ret[0:size])
  131. }
  132. func UintToString(orig []uint8) string {
  133. ret := make([]byte, len(orig))
  134. size := -1
  135. for i, o := range orig {
  136. if o == 0 {
  137. size = i
  138. break
  139. }
  140. ret[i] = o
  141. }
  142. if size == -1 {
  143. size = len(orig)
  144. }
  145. return string(ret[0:size])
  146. }
  147. func ByteToString(orig []byte) string {
  148. n := -1
  149. l := -1
  150. for i, b := range orig {
  151. // skip left side null
  152. if l == -1 && b == 0 {
  153. continue
  154. }
  155. if l == -1 {
  156. l = i
  157. }
  158. if b == 0 {
  159. break
  160. }
  161. n = i + 1
  162. }
  163. if n == -1 {
  164. return string(orig)
  165. }
  166. return string(orig[l:n])
  167. }
  168. // ReadInts reads contents from single line file and returns them as []int32.
  169. func ReadInts(filename string) ([]int64, error) {
  170. f, err := os.Open(filename)
  171. if err != nil {
  172. return []int64{}, err
  173. }
  174. defer f.Close()
  175. var ret []int64
  176. r := bufio.NewReader(f)
  177. // The int files that this is concerned with should only be one liners.
  178. line, err := r.ReadString('\n')
  179. if err != nil {
  180. return []int64{}, err
  181. }
  182. i, err := strconv.ParseInt(strings.Trim(line, "\n"), 10, 32)
  183. if err != nil {
  184. return []int64{}, err
  185. }
  186. ret = append(ret, i)
  187. return ret, nil
  188. }
  189. // HexToUint32 parses Hex to uint32 without error.
  190. func HexToUint32(hex string) uint32 {
  191. vv, _ := strconv.ParseUint(hex, 16, 32)
  192. return uint32(vv)
  193. }
  194. // mustParseInt32 parses to int32 without error.
  195. func mustParseInt32(val string) int32 {
  196. vv, _ := strconv.ParseInt(val, 10, 32)
  197. return int32(vv)
  198. }
  199. // mustParseUint64 parses to uint64 without error.
  200. func mustParseUint64(val string) uint64 {
  201. vv, _ := strconv.ParseInt(val, 10, 64)
  202. return uint64(vv)
  203. }
  204. // mustParseFloat64 parses to Float64 without error.
  205. func mustParseFloat64(val string) float64 {
  206. vv, _ := strconv.ParseFloat(val, 64)
  207. return vv
  208. }
  209. // StringsHas checks the target string slice contains src or not.
  210. func StringsHas(target []string, src string) bool {
  211. for _, t := range target {
  212. if strings.TrimSpace(t) == src {
  213. return true
  214. }
  215. }
  216. return false
  217. }
  218. // StringsContains checks the src in any string of the target string slice.
  219. func StringsContains(target []string, src string) bool {
  220. for _, t := range target {
  221. if strings.Contains(t, src) {
  222. return true
  223. }
  224. }
  225. return false
  226. }
  227. // IntContains checks the src in any int of the target int slice.
  228. func IntContains(target []int, src int) bool {
  229. for _, t := range target {
  230. if src == t {
  231. return true
  232. }
  233. }
  234. return false
  235. }
  236. // get struct attributes.
  237. // This method is used only for debugging platform dependent code.
  238. func attributes(m interface{}) map[string]reflect.Type {
  239. typ := reflect.TypeOf(m)
  240. if typ.Kind() == reflect.Ptr {
  241. typ = typ.Elem()
  242. }
  243. attrs := make(map[string]reflect.Type)
  244. if typ.Kind() != reflect.Struct {
  245. return nil
  246. }
  247. for i := 0; i < typ.NumField(); i++ {
  248. p := typ.Field(i)
  249. if !p.Anonymous {
  250. attrs[p.Name] = p.Type
  251. }
  252. }
  253. return attrs
  254. }
  255. func PathExists(filename string) bool {
  256. if _, err := os.Stat(filename); err == nil {
  257. return true
  258. }
  259. return false
  260. }
  261. // GetEnv retrieves the environment variable key. If it does not exist it returns the default.
  262. func GetEnv(key string, dfault string, combineWith ...string) string {
  263. value := os.Getenv(key)
  264. if value == "" {
  265. value = dfault
  266. }
  267. switch len(combineWith) {
  268. case 0:
  269. return value
  270. case 1:
  271. return filepath.Join(value, combineWith[0])
  272. default:
  273. all := make([]string, len(combineWith)+1)
  274. all[0] = value
  275. copy(all[1:], combineWith)
  276. return filepath.Join(all...)
  277. }
  278. }
  279. func HostProc(combineWith ...string) string {
  280. return GetEnv("HOST_PROC", "/proc", combineWith...)
  281. }
  282. func HostSys(combineWith ...string) string {
  283. return GetEnv("HOST_SYS", "/sys", combineWith...)
  284. }
  285. func HostEtc(combineWith ...string) string {
  286. return GetEnv("HOST_ETC", "/etc", combineWith...)
  287. }
  288. func HostVar(combineWith ...string) string {
  289. return GetEnv("HOST_VAR", "/var", combineWith...)
  290. }
  291. func HostRun(combineWith ...string) string {
  292. return GetEnv("HOST_RUN", "/run", combineWith...)
  293. }
  294. func HostDev(combineWith ...string) string {
  295. return GetEnv("HOST_DEV", "/dev", combineWith...)
  296. }
  297. // MockEnv set environment variable and return revert function.
  298. // MockEnv should be used testing only.
  299. func MockEnv(key string, value string) func() {
  300. original := os.Getenv(key)
  301. os.Setenv(key, value)
  302. return func() {
  303. os.Setenv(key, original)
  304. }
  305. }
  306. // getSysctrlEnv sets LC_ALL=C in a list of env vars for use when running
  307. // sysctl commands (see DoSysctrl).
  308. func getSysctrlEnv(env []string) []string {
  309. foundLC := false
  310. for i, line := range env {
  311. if strings.HasPrefix(line, "LC_ALL") {
  312. env[i] = "LC_ALL=C"
  313. foundLC = true
  314. }
  315. }
  316. if !foundLC {
  317. env = append(env, "LC_ALL=C")
  318. }
  319. return env
  320. }