state.go

     1  // Copyright 2022 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package script
     6  
     7  import (
     8  	"bytes"
     9  	"context"
    10  	"fmt"
    11  	"internal/txtar"
    12  	"io"
    13  	"io/fs"
    14  	"os"
    15  	"os/exec"
    16  	"path/filepath"
    17  	"regexp"
    18  	"strings"
    19  )
    20  
    21  // A State encapsulates the current state of a running script engine,
    22  // including the script environment and any running background commands.
    23  type State struct {
    24  	engine *Engine // the engine currently executing the script, if any
    25  
    26  	ctx    context.Context
    27  	cancel context.CancelFunc
    28  	file   string
    29  	log    bytes.Buffer
    30  
    31  	workdir string            // initial working directory
    32  	pwd     string            // current working directory during execution
    33  	env     []string          // environment list (for os/exec)
    34  	envMap  map[string]string // environment mapping (matches env)
    35  	stdout  string            // standard output from last 'go' command; for 'stdout' command
    36  	stderr  string            // standard error from last 'go' command; for 'stderr' command
    37  
    38  	background []backgroundCmd
    39  }
    40  
    41  type backgroundCmd struct {
    42  	*command
    43  	wait WaitFunc
    44  }
    45  
    46  // NewState returns a new State permanently associated with ctx, with its
    47  // initial working directory in workdir and its initial environment set to
    48  // initialEnv (or os.Environ(), if initialEnv is nil).
    49  //
    50  // The new State also contains pseudo-environment-variables for
    51  // ${/} and ${:} (for the platform's path and list separators respectively),
    52  // but does not pass those to subprocesses.
    53  func NewState(ctx context.Context, workdir string, initialEnv []string) (*State, error) {
    54  	absWork, err := filepath.Abs(workdir)
    55  	if err != nil {
    56  		return nil, err
    57  	}
    58  
    59  	ctx, cancel := context.WithCancel(ctx)
    60  
    61  	// Make a fresh copy of the env slice to avoid aliasing bugs if we ever
    62  	// start modifying it in place; this also establishes the invariant that
    63  	// s.env contains no duplicates.
    64  	env := cleanEnv(initialEnv, absWork)
    65  
    66  	envMap := make(map[string]string, len(env))
    67  
    68  	// Add entries for ${:} and ${/} to make it easier to write platform-independent
    69  	// paths in scripts.
    70  	envMap["/"] = string(os.PathSeparator)
    71  	envMap[":"] = string(os.PathListSeparator)
    72  
    73  	for _, kv := range env {
    74  		if k, v, ok := strings.Cut(kv, "="); ok {
    75  			envMap[k] = v
    76  		}
    77  	}
    78  
    79  	s := &State{
    80  		ctx:     ctx,
    81  		cancel:  cancel,
    82  		workdir: absWork,
    83  		pwd:     absWork,
    84  		env:     env,
    85  		envMap:  envMap,
    86  	}
    87  	s.Setenv("PWD", absWork)
    88  	return s, nil
    89  }
    90  
    91  // CloseAndWait cancels the State's Context and waits for any background commands to
    92  // finish. If any remaining background command ended in an unexpected state,
    93  // Close returns a non-nil error.
    94  func (s *State) CloseAndWait(log io.Writer) error {
    95  	s.cancel()
    96  	wait, err := Wait().Run(s)
    97  	if wait != nil {
    98  		panic("script: internal error: Wait unexpectedly returns its own WaitFunc")
    99  	}
   100  	if flushErr := s.flushLog(log); err == nil {
   101  		err = flushErr
   102  	}
   103  	return err
   104  }
   105  
   106  // Chdir changes the State's working directory to the given path.
   107  func (s *State) Chdir(path string) error {
   108  	dir := s.Path(path)
   109  	if _, err := os.Stat(dir); err != nil {
   110  		return &fs.PathError{Op: "Chdir", Path: dir, Err: err}
   111  	}
   112  	s.pwd = dir
   113  	s.Setenv("PWD", dir)
   114  	return nil
   115  }
   116  
   117  // Context returns the Context with which the State was created.
   118  func (s *State) Context() context.Context {
   119  	return s.ctx
   120  }
   121  
   122  // Environ returns a copy of the current script environment,
   123  // in the form "key=value".
   124  func (s *State) Environ() []string {
   125  	return append([]string(nil), s.env...)
   126  }
   127  
   128  // ExpandEnv replaces ${var} or $var in the string according to the values of
   129  // the environment variables in s. References to undefined variables are
   130  // replaced by the empty string.
   131  func (s *State) ExpandEnv(str string, inRegexp bool) string {
   132  	return os.Expand(str, func(key string) string {
   133  		e := s.envMap[key]
   134  		if inRegexp {
   135  			// Quote to literal strings: we want paths like C:\work\go1.4 to remain
   136  			// paths rather than regular expressions.
   137  			e = regexp.QuoteMeta(e)
   138  		}
   139  		return e
   140  	})
   141  }
   142  
   143  // ExtractFiles extracts the files in ar to the state's current directory,
   144  // expanding any environment variables within each name.
   145  //
   146  // The files must reside within the working directory with which the State was
   147  // originally created.
   148  func (s *State) ExtractFiles(ar *txtar.Archive) error {
   149  	wd := s.workdir
   150  
   151  	// Add trailing separator to terminate wd.
   152  	// This prevents extracting to outside paths which prefix wd,
   153  	// e.g. extracting to /home/foobar when wd is /home/foo
   154  	if wd == "" {
   155  		panic("s.workdir is unexpectedly empty")
   156  	}
   157  	if !os.IsPathSeparator(wd[len(wd)-1]) {
   158  		wd += string(filepath.Separator)
   159  	}
   160  
   161  	for _, f := range ar.Files {
   162  		name := s.Path(s.ExpandEnv(f.Name, false))
   163  
   164  		if !strings.HasPrefix(name, wd) {
   165  			return fmt.Errorf("file %#q is outside working directory", f.Name)
   166  		}
   167  
   168  		if err := os.MkdirAll(filepath.Dir(name), 0777); err != nil {
   169  			return err
   170  		}
   171  		if err := os.WriteFile(name, f.Data, 0666); err != nil {
   172  			return err
   173  		}
   174  	}
   175  
   176  	return nil
   177  }
   178  
   179  // Getwd returns the directory in which to run the next script command.
   180  func (s *State) Getwd() string { return s.pwd }
   181  
   182  // Logf writes output to the script's log without updating its stdout or stderr
   183  // buffers. (The output log functions as a kind of meta-stderr.)
   184  func (s *State) Logf(format string, args ...any) {
   185  	fmt.Fprintf(&s.log, format, args...)
   186  }
   187  
   188  // flushLog writes the contents of the script's log to w and clears the log.
   189  func (s *State) flushLog(w io.Writer) error {
   190  	_, err := w.Write(s.log.Bytes())
   191  	s.log.Reset()
   192  	return err
   193  }
   194  
   195  // LookupEnv retrieves the value of the environment variable in s named by the key.
   196  func (s *State) LookupEnv(key string) (string, bool) {
   197  	v, ok := s.envMap[key]
   198  	return v, ok
   199  }
   200  
   201  // Path returns the absolute path in the host operating system for a
   202  // script-based (generally slash-separated and relative) path.
   203  func (s *State) Path(path string) string {
   204  	if filepath.IsAbs(path) {
   205  		return filepath.Clean(path)
   206  	}
   207  	return filepath.Join(s.pwd, path)
   208  }
   209  
   210  // Setenv sets the value of the environment variable in s named by the key.
   211  func (s *State) Setenv(key, value string) error {
   212  	s.env = cleanEnv(append(s.env, key+"="+value), s.pwd)
   213  	s.envMap[key] = value
   214  	return nil
   215  }
   216  
   217  // Stdout returns the stdout output of the last command run,
   218  // or the empty string if no command has been run.
   219  func (s *State) Stdout() string { return s.stdout }
   220  
   221  // Stderr returns the stderr output of the last command run,
   222  // or the empty string if no command has been run.
   223  func (s *State) Stderr() string { return s.stderr }
   224  
   225  // cleanEnv returns a copy of env with any duplicates removed in favor of
   226  // later values and any required system variables defined.
   227  //
   228  // If env is nil, cleanEnv copies the environment from os.Environ().
   229  func cleanEnv(env []string, pwd string) []string {
   230  	// There are some funky edge-cases in this logic, especially on Windows (with
   231  	// case-insensitive environment variables and variables with keys like "=C:").
   232  	// Rather than duplicating exec.dedupEnv here, cheat and use exec.Cmd directly.
   233  	cmd := &exec.Cmd{Env: env}
   234  	cmd.Dir = pwd
   235  	return cmd.Environ()
   236  }
   237
View as plain text