shell.go 4.74 KB
Newer Older
ale's avatar
ale committed
1 2 3 4
package tabacco

import (
	"context"
5
	"fmt"
ale's avatar
ale committed
6
	"io"
7
	"io/ioutil"
ale's avatar
ale committed
8 9 10
	"log"
	"os"
	"os/exec"
11
	"path/filepath"
ale's avatar
ale committed
12
	"strconv"
13
	"strings"
14 15

	"git.autistici.org/ale/tabacco/jobs"
ale's avatar
ale committed
16 17
)

18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43
// Environment variables starting with any of the following strings
// are removed from the command execution environment.
var filteredEnvVars = []string{
	"LANG=",
	"LC_",
	"RESTIC_",
}

// Return a copy of os.Environ(), with filteredEnvVars removed.
func safeEnv() []string {
	var env []string
	for _, s := range os.Environ() {
		skip := false
		for _, pfx := range filteredEnvVars {
			if strings.HasPrefix(s, pfx) {
				skip = true
				break
			}
		}
		if !skip {
			env = append(env, s)
		}
	}
	return env
}

ale's avatar
ale committed
44 45 46 47 48 49 50
// Shell runs commands, with some options (a global dry-run flag
// preventing all executions, nice level, i/o class). As one may guess
// by the name, commands are run using the shell, so variable
// substitutions and other shell features are available.
type Shell struct {
	dryRun      bool
	niceLevel   int
ale's avatar
ale committed
51
	ioniceClass int
52
	env         []string
ale's avatar
ale committed
53 54 55 56 57 58
}

// NewShell creates a new Shell.
func NewShell(dryRun bool) *Shell {
	return &Shell{
		dryRun:      dryRun,
ale's avatar
ale committed
59 60
		niceLevel:   10,
		ioniceClass: 2,
61
		env:         safeEnv(),
ale's avatar
ale committed
62 63 64 65 66 67 68 69 70
	}
}

// SetNiceLevel sets the nice(1) level.
func (s *Shell) SetNiceLevel(n int) {
	s.niceLevel = n
}

// SetIOClass sets the ionice(1) i/o class.
ale's avatar
ale committed
71 72
func (s *Shell) SetIOClass(n int) {
	s.ioniceClass = n
ale's avatar
ale committed
73 74
}

75 76 77
// command builds an exec.Cmd and gets some parameters from the
// context - notably it sets log output and working directory to be in
// the working dir if the job has been wrapped in WithWorkDir().
ale's avatar
ale committed
78 79 80 81 82 83 84 85 86 87
func (s *Shell) command(ctx context.Context, arg string) *exec.Cmd {
	var args []string
	if s.dryRun {
		args = []string{"/bin/echo", arg}
	} else {
		args = []string{"/bin/sh", "-c", arg}
	}

	if s.niceLevel != 0 {
		args = append(
88
			[]string{"nice", "-n", strconv.Itoa(s.niceLevel)},
ale's avatar
ale committed
89 90 91
			args...,
		)
	}
ale's avatar
ale committed
92
	if s.ioniceClass != 0 {
ale's avatar
ale committed
93
		args = append(
94
			[]string{"ionice", "-c", strconv.Itoa(s.ioniceClass)},
ale's avatar
ale committed
95 96 97 98 99
			args...,
		)
	}

	c := exec.CommandContext(ctx, args[0], args[1:]...) // #nosec
100 101 102
	var env []string
	env = append(env, s.env...)
	c.Env = env
ale's avatar
ale committed
103
	c.Stderr = os.Stderr
104
	c.Dir = getWorkDir(ctx)
105 106 107

	log.Printf("sh: %s", arg)

ale's avatar
ale committed
108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153
	return c
}

// RunStdoutPipe runs a command with a function connected to its
// standard output via a pipe.
func (s *Shell) RunStdoutPipe(ctx context.Context, arg string, fn func(io.Reader) error) error {
	cmd := s.command(ctx, arg)
	log.Printf("stdout_pipe: %s", arg)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}
	if err := cmd.Start(); err != nil {
		return err
	}
	if err := fn(stdout); err != nil {
		return err
	}
	return cmd.Wait()
}

// RunStdinPipe runs a command with a function connected to its
// standard input via a pipe.
func (s *Shell) RunStdinPipe(ctx context.Context, arg string, fn func(io.Writer) error) error {
	cmd := s.command(ctx, arg)
	log.Printf("stdin_pipe: %s", arg)
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return err
	}
	if err := cmd.Start(); err != nil {
		return err
	}
	if err := fn(stdin); err != nil {
		return err
	}
	return cmd.Wait()
}

// Run a command.
func (s *Shell) Run(ctx context.Context, arg string) error {
	c := s.command(ctx, arg)
	c.Stdout = os.Stdout
	return c.Run()
}

154 155 156 157
// RunWithEnv runs a command with additional environment variables.
func (s *Shell) RunWithEnv(ctx context.Context, arg string, envMap map[string]string) error {
	c := s.command(ctx, arg)
	for k, v := range envMap {
158
		c.Env = append(c.Env, fmt.Sprintf("%s=%s", k, v))
159 160 161 162 163
	}
	c.Stdout = os.Stdout
	return c.Run()
}

ale's avatar
ale committed
164 165 166 167
// Output runs a command and returns the standard output.
func (s *Shell) Output(ctx context.Context, arg string) ([]byte, error) {
	return s.command(ctx, arg).Output()
}
168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220

// Manage a work directory (temporary scratch with lots of space).
type workdirManager struct {
	dir string
}

func newWorkdirManager(tmpdir string) (*workdirManager, error) {
	dir, err := ioutil.TempDir(tmpdir, "tabacco-")
	if err != nil {
		return nil, err
	}
	return &workdirManager{dir}, nil
}

func (w *workdirManager) Close() {
	os.RemoveAll(w.dir) // nolint
}

func (w *workdirManager) withWorkDir(j jobs.Job) jobs.Job {
	jobDir := filepath.Join(w.dir, "job-"+j.ID())
	return &workdirJob{
		Job: j,
		dir: jobDir,
	}
}

type workdirJob struct {
	jobs.Job
	dir string
}

type workdirKeyType int

var workdirKey workdirKeyType = 1

func (j *workdirJob) RunContext(ctx context.Context) error {
	if err := os.Mkdir(j.dir, 0700); err != nil {
		return err
	}

	ctx = context.WithValue(ctx, workdirKey, j.dir)
	return j.Job.RunContext(ctx)
}

// getWorkDir returns the current working directory, provided the
// WithWorkDir wrapper has been called.
func getWorkDir(ctx context.Context) string {
	dir, ok := ctx.Value(workdirKey).(string)
	if ok {
		return dir
	}
	return "/tmp"
}