socket.go

Documentation: golang.org/x/tools/playground/socket

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