github.com/graybobo/golang.org-package-offline-cache@v0.0.0-20200626051047-6608995c132f/x/tools/playground/socket/socket.go (about)

     1  // Copyright 2012 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  // +build !appengine
     6  
     7  // Package socket implements an WebSocket-based playground backend.
     8  // Clients connect to a websocket handler and send run/kill commands, and
     9  // the server sends the output and exit status of the running processes.
    10  // Multiple clients running multiple processes may be served concurrently.
    11  // The wire format is JSON and is described by the Message type.
    12  //
    13  // This will not run on App Engine as WebSockets are not supported there.
    14  package socket // import "golang.org/x/tools/playground/socket"
    15  
    16  import (
    17  	"bytes"
    18  	"encoding/json"
    19  	"errors"
    20  	"go/parser"
    21  	"go/token"
    22  	"io"
    23  	"io/ioutil"
    24  	"log"
    25  	"net"
    26  	"net/http"
    27  	"net/url"
    28  	"os"
    29  	"os/exec"
    30  	"path/filepath"
    31  	"runtime"
    32  	"strconv"
    33  	"strings"
    34  	"time"
    35  	"unicode/utf8"
    36  
    37  	"golang.org/x/net/websocket"
    38  )
    39  
    40  // RunScripts specifies whether the socket handler should execute shell scripts
    41  // (snippets that start with a shebang).
    42  var RunScripts = true
    43  
    44  // Environ provides an environment when a binary, such as the go tool, is
    45  // invoked.
    46  var Environ func() []string = os.Environ
    47  
    48  const (
    49  	// The maximum number of messages to send per session (avoid flooding).
    50  	msgLimit = 1000
    51  
    52  	// Batch messages sent in this interval and send as a single message.
    53  	msgDelay = 10 * time.Millisecond
    54  )
    55  
    56  // Message is the wire format for the websocket connection to the browser.
    57  // It is used for both sending output messages and receiving commands, as
    58  // distinguished by the Kind field.
    59  type Message struct {
    60  	Id      string // client-provided unique id for the process
    61  	Kind    string // in: "run", "kill" out: "stdout", "stderr", "end"
    62  	Body    string
    63  	Options *Options `json:",omitempty"`
    64  }
    65  
    66  // Options specify additional message options.
    67  type Options struct {
    68  	Race bool // use -race flag when building code (for "run" only)
    69  }
    70  
    71  // NewHandler returns a websocket server which checks the origin of requests.
    72  func NewHandler(origin *url.URL) websocket.Server {
    73  	return websocket.Server{
    74  		Config:    websocket.Config{Origin: origin},
    75  		Handshake: handshake,
    76  		Handler:   websocket.Handler(socketHandler),
    77  	}
    78  }
    79  
    80  // handshake checks the origin of a request during the websocket handshake.
    81  func handshake(c *websocket.Config, req *http.Request) error {
    82  	o, err := websocket.Origin(c, req)
    83  	if err != nil {
    84  		log.Println("bad websocket origin:", err)
    85  		return websocket.ErrBadWebSocketOrigin
    86  	}
    87  	_, port, err := net.SplitHostPort(c.Origin.Host)
    88  	if err != nil {
    89  		log.Println("bad websocket origin:", err)
    90  		return websocket.ErrBadWebSocketOrigin
    91  	}
    92  	ok := c.Origin.Scheme == o.Scheme && (c.Origin.Host == o.Host || c.Origin.Host == net.JoinHostPort(o.Host, port))
    93  	if !ok {
    94  		log.Println("bad websocket origin:", o)
    95  		return websocket.ErrBadWebSocketOrigin
    96  	}
    97  	log.Println("accepting connection from:", req.RemoteAddr)
    98  	return nil
    99  }
   100  
   101  // socketHandler handles the websocket connection for a given present session.
   102  // It handles transcoding Messages to and from JSON format, and starting
   103  // and killing processes.
   104  func socketHandler(c *websocket.Conn) {
   105  	in, out := make(chan *Message), make(chan *Message)
   106  	errc := make(chan error, 1)
   107  
   108  	// Decode messages from client and send to the in channel.
   109  	go func() {
   110  		dec := json.NewDecoder(c)
   111  		for {
   112  			var m Message
   113  			if err := dec.Decode(&m); err != nil {
   114  				errc <- err
   115  				return
   116  			}
   117  			in <- &m
   118  		}
   119  	}()
   120  
   121  	// Receive messages from the out channel and encode to the client.
   122  	go func() {
   123  		enc := json.NewEncoder(c)
   124  		for m := range out {
   125  			if err := enc.Encode(m); err != nil {
   126  				errc <- err
   127  				return
   128  			}
   129  		}
   130  	}()
   131  	defer close(out)
   132  
   133  	// Start and kill processes and handle errors.
   134  	proc := make(map[string]*process)
   135  	for {
   136  		select {
   137  		case m := <-in:
   138  			switch m.Kind {
   139  			case "run":
   140  				log.Println("running snippet from:", c.Request().RemoteAddr)
   141  				proc[m.Id].Kill()
   142  				proc[m.Id] = startProcess(m.Id, m.Body, out, m.Options)
   143  			case "kill":
   144  				proc[m.Id].Kill()
   145  			}
   146  		case err := <-errc:
   147  			if err != io.EOF {
   148  				// A encode or decode has failed; bail.
   149  				log.Println(err)
   150  			}
   151  			// Shut down any running processes.
   152  			for _, p := range proc {
   153  				p.Kill()
   154  			}
   155  			return
   156  		}
   157  	}
   158  }
   159  
   160  // process represents a running process.
   161  type process struct {
   162  	out  chan<- *Message
   163  	done chan struct{} // closed when wait completes
   164  	run  *exec.Cmd
   165  	bin  string
   166  }
   167  
   168  // startProcess builds and runs the given program, sending its output
   169  // and end event as Messages on the provided channel.
   170  func startProcess(id, body string, dest chan<- *Message, opt *Options) *process {
   171  	var (
   172  		done = make(chan struct{})
   173  		out  = make(chan *Message)
   174  		p    = &process{out: out, done: done}
   175  	)
   176  	go func() {
   177  		defer close(done)
   178  		for m := range buffer(limiter(out, p)) {
   179  			m.Id = id
   180  			dest <- m
   181  		}
   182  	}()
   183  	var err error
   184  	if path, args := shebang(body); path != "" {
   185  		if RunScripts {
   186  			err = p.startProcess(path, args, body)
   187  		} else {
   188  			err = errors.New("script execution is not allowed")
   189  		}
   190  	} else {
   191  		err = p.start(body, opt)
   192  	}
   193  	if err != nil {
   194  		p.end(err)
   195  		return nil
   196  	}
   197  	go func() {
   198  		p.end(p.run.Wait())
   199  	}()
   200  	return p
   201  }
   202  
   203  // end sends an "end" message to the client, containing the process id and the
   204  // given error value. It also removes the binary, if present.
   205  func (p *process) end(err error) {
   206  	if p.bin != "" {
   207  		defer os.Remove(p.bin)
   208  	}
   209  	m := &Message{Kind: "end"}
   210  	if err != nil {
   211  		m.Body = err.Error()
   212  	}
   213  	p.out <- m
   214  	close(p.out)
   215  }
   216  
   217  // A killer provides a mechanism to terminate a process.
   218  // The Kill method returns only once the process has exited.
   219  type killer interface {
   220  	Kill()
   221  }
   222  
   223  // limiter returns a channel that wraps the given channel.
   224  // It receives Messages from the given channel and sends them to the returned
   225  // channel until it passes msgLimit messages, at which point it will kill the
   226  // process and pass only the "end" message.
   227  // When the given channel is closed, or when the "end" message is received,
   228  // it closes the returned channel.
   229  func limiter(in <-chan *Message, p killer) <-chan *Message {
   230  	out := make(chan *Message)
   231  	go func() {
   232  		defer close(out)
   233  		n := 0
   234  		for m := range in {
   235  			switch {
   236  			case n < msgLimit || m.Kind == "end":
   237  				out <- m
   238  				if m.Kind == "end" {
   239  					return
   240  				}
   241  			case n == msgLimit:
   242  				// Kill in a goroutine as Kill will not return
   243  				// until the process' output has been
   244  				// processed, and we're doing that in this loop.
   245  				go p.Kill()
   246  			default:
   247  				continue // don't increment
   248  			}
   249  			n++
   250  		}
   251  	}()
   252  	return out
   253  }
   254  
   255  // buffer returns a channel that wraps the given channel. It receives messages
   256  // from the given channel and sends them to the returned channel.
   257  // Message bodies are gathered over the period msgDelay and coalesced into a
   258  // single Message before they are passed on. Messages of the same kind are
   259  // coalesced; when a message of a different kind is received, any buffered
   260  // messages are flushed. When the given channel is closed, buffer flushes the
   261  // remaining buffered messages and closes the returned channel.
   262  func buffer(in <-chan *Message) <-chan *Message {
   263  	out := make(chan *Message)
   264  	go func() {
   265  		defer close(out)
   266  		var (
   267  			t     = time.NewTimer(msgDelay)
   268  			tc    <-chan time.Time
   269  			buf   []byte
   270  			kind  string
   271  			flush = func() {
   272  				if len(buf) == 0 {
   273  					return
   274  				}
   275  				out <- &Message{Kind: kind, Body: safeString(buf)}
   276  				buf = buf[:0] // recycle buffer
   277  				kind = ""
   278  			}
   279  		)
   280  		for {
   281  			select {
   282  			case m, ok := <-in:
   283  				if !ok {
   284  					flush()
   285  					return
   286  				}
   287  				if m.Kind == "end" {
   288  					flush()
   289  					out <- m
   290  					return
   291  				}
   292  				if kind != m.Kind {
   293  					flush()
   294  					kind = m.Kind
   295  					if tc == nil {
   296  						tc = t.C
   297  						t.Reset(msgDelay)
   298  					}
   299  				}
   300  				buf = append(buf, m.Body...)
   301  			case <-tc:
   302  				flush()
   303  				tc = nil
   304  			}
   305  		}
   306  	}()
   307  	return out
   308  }
   309  
   310  // Kill stops the process if it is running and waits for it to exit.
   311  func (p *process) Kill() {
   312  	if p == nil || p.run == nil {
   313  		return
   314  	}
   315  	p.run.Process.Kill()
   316  	<-p.done // block until process exits
   317  }
   318  
   319  // shebang looks for a shebang ('#!') at the beginning of the passed string.
   320  // If found, it returns the path and args after the shebang.
   321  // args includes the command as args[0].
   322  func shebang(body string) (path string, args []string) {
   323  	body = strings.TrimSpace(body)
   324  	if !strings.HasPrefix(body, "#!") {
   325  		return "", nil
   326  	}
   327  	if i := strings.Index(body, "\n"); i >= 0 {
   328  		body = body[:i]
   329  	}
   330  	fs := strings.Fields(body[2:])
   331  	return fs[0], fs
   332  }
   333  
   334  // startProcess starts a given program given its path and passing the given body
   335  // to the command standard input.
   336  func (p *process) startProcess(path string, args []string, body string) error {
   337  	cmd := &exec.Cmd{
   338  		Path:   path,
   339  		Args:   args,
   340  		Stdin:  strings.NewReader(body),
   341  		Stdout: &messageWriter{kind: "stdout", out: p.out},
   342  		Stderr: &messageWriter{kind: "stderr", out: p.out},
   343  	}
   344  	if err := cmd.Start(); err != nil {
   345  		return err
   346  	}
   347  	p.run = cmd
   348  	return nil
   349  }
   350  
   351  // start builds and starts the given program, sending its output to p.out,
   352  // and stores the running *exec.Cmd in the run field.
   353  func (p *process) start(body string, opt *Options) error {
   354  	// We "go build" and then exec the binary so that the
   355  	// resultant *exec.Cmd is a handle to the user's program
   356  	// (rather than the go tool process).
   357  	// This makes Kill work.
   358  
   359  	bin := filepath.Join(tmpdir, "compile"+strconv.Itoa(<-uniq))
   360  	src := bin + ".go"
   361  	if runtime.GOOS == "windows" {
   362  		bin += ".exe"
   363  	}
   364  
   365  	// write body to x.go
   366  	defer os.Remove(src)
   367  	err := ioutil.WriteFile(src, []byte(body), 0666)
   368  	if err != nil {
   369  		return err
   370  	}
   371  
   372  	// build x.go, creating x
   373  	p.bin = bin // to be removed by p.end
   374  	dir, file := filepath.Split(src)
   375  	args := []string{"go", "build", "-tags", "OMIT"}
   376  	if opt != nil && opt.Race {
   377  		p.out <- &Message{
   378  			Kind: "stderr",
   379  			Body: "Running with race detector.\n",
   380  		}
   381  		args = append(args, "-race")
   382  	}
   383  	args = append(args, "-o", bin, file)
   384  	cmd := p.cmd(dir, args...)
   385  	cmd.Stdout = cmd.Stderr // send compiler output to stderr
   386  	if err := cmd.Run(); err != nil {
   387  		return err
   388  	}
   389  
   390  	// run x
   391  	if isNacl() {
   392  		cmd, err = p.naclCmd(bin)
   393  		if err != nil {
   394  			return err
   395  		}
   396  	} else {
   397  		cmd = p.cmd("", bin)
   398  	}
   399  	if opt != nil && opt.Race {
   400  		cmd.Env = append(cmd.Env, "GOMAXPROCS=2")
   401  	}
   402  	if err := cmd.Start(); err != nil {
   403  		// If we failed to exec, that might be because they built
   404  		// a non-main package instead of an executable.
   405  		// Check and report that.
   406  		if name, err := packageName(body); err == nil && name != "main" {
   407  			return errors.New(`executable programs must use "package main"`)
   408  		}
   409  		return err
   410  	}
   411  	p.run = cmd
   412  	return nil
   413  }
   414  
   415  // cmd builds an *exec.Cmd that writes its standard output and error to the
   416  // process' output channel.
   417  func (p *process) cmd(dir string, args ...string) *exec.Cmd {
   418  	cmd := exec.Command(args[0], args[1:]...)
   419  	cmd.Dir = dir
   420  	cmd.Env = Environ()
   421  	cmd.Stdout = &messageWriter{kind: "stdout", out: p.out}
   422  	cmd.Stderr = &messageWriter{kind: "stderr", out: p.out}
   423  	return cmd
   424  }
   425  
   426  func isNacl() bool {
   427  	for _, v := range append(Environ(), os.Environ()...) {
   428  		if v == "GOOS=nacl" {
   429  			return true
   430  		}
   431  	}
   432  	return false
   433  }
   434  
   435  // naclCmd returns an *exec.Cmd that executes bin under native client.
   436  func (p *process) naclCmd(bin string) (*exec.Cmd, error) {
   437  	pwd, err := os.Getwd()
   438  	if err != nil {
   439  		return nil, err
   440  	}
   441  	var args []string
   442  	env := []string{
   443  		"NACLENV_GOOS=" + runtime.GOOS,
   444  		"NACLENV_GOROOT=/go",
   445  		"NACLENV_NACLPWD=" + strings.Replace(pwd, runtime.GOROOT(), "/go", 1),
   446  	}
   447  	switch runtime.GOARCH {
   448  	case "amd64":
   449  		env = append(env, "NACLENV_GOARCH=amd64p32")
   450  		args = []string{"sel_ldr_x86_64"}
   451  	case "386":
   452  		env = append(env, "NACLENV_GOARCH=386")
   453  		args = []string{"sel_ldr_x86_32"}
   454  	case "arm":
   455  		env = append(env, "NACLENV_GOARCH=arm")
   456  		selLdr, err := exec.LookPath("sel_ldr_arm")
   457  		if err != nil {
   458  			return nil, err
   459  		}
   460  		args = []string{"nacl_helper_bootstrap_arm", selLdr, "--reserved_at_zero=0xXXXXXXXXXXXXXXXX"}
   461  	default:
   462  		return nil, errors.New("native client does not support GOARCH=" + runtime.GOARCH)
   463  	}
   464  
   465  	cmd := p.cmd("", append(args, "-l", "/dev/null", "-S", "-e", bin)...)
   466  	cmd.Env = append(cmd.Env, env...)
   467  
   468  	return cmd, nil
   469  }
   470  
   471  func packageName(body string) (string, error) {
   472  	f, err := parser.ParseFile(token.NewFileSet(), "prog.go",
   473  		strings.NewReader(body), parser.PackageClauseOnly)
   474  	if err != nil {
   475  		return "", err
   476  	}
   477  	return f.Name.String(), nil
   478  }
   479  
   480  // messageWriter is an io.Writer that converts all writes to Message sends on
   481  // the out channel with the specified id and kind.
   482  type messageWriter struct {
   483  	kind string
   484  	out  chan<- *Message
   485  }
   486  
   487  func (w *messageWriter) Write(b []byte) (n int, err error) {
   488  	w.out <- &Message{Kind: w.kind, Body: safeString(b)}
   489  	return len(b), nil
   490  }
   491  
   492  // safeString returns b as a valid UTF-8 string.
   493  func safeString(b []byte) string {
   494  	if utf8.Valid(b) {
   495  		return string(b)
   496  	}
   497  	var buf bytes.Buffer
   498  	for len(b) > 0 {
   499  		r, size := utf8.DecodeRune(b)
   500  		b = b[size:]
   501  		buf.WriteRune(r)
   502  	}
   503  	return buf.String()
   504  }
   505  
   506  var tmpdir string
   507  
   508  func init() {
   509  	// find real path to temporary directory
   510  	var err error
   511  	tmpdir, err = filepath.EvalSymlinks(os.TempDir())
   512  	if err != nil {
   513  		log.Fatal(err)
   514  	}
   515  }
   516  
   517  var uniq = make(chan int) // a source of numbers for naming temporary files
   518  
   519  func init() {
   520  	go func() {
   521  		for i := 0; ; i++ {
   522  			uniq <- i
   523  		}
   524  	}()
   525  }