.config/coc/extensions/coc-go-data/tools/pkg/mod/golang.org/x/tools@v0.1.1-0.20210319172145-bda8f5cee399/playground/socket/socket.go

   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 }