pkg/beam/examples/beamsh/beamsh.go

package main

import (
	"bufio"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"github.com/dotcloud/docker/pkg/beam"
	"github.com/dotcloud/docker/pkg/beam/data"
	"github.com/dotcloud/docker/pkg/dockerscript"
	"github.com/dotcloud/docker/pkg/term"
	"io"
	"net"
	"net/url"
	"os"
	"os/exec"
	"path"
	"strings"
	"sync"
)

func main() {
	devnull, err := Devnull()
	if err != nil {
		Fatal(err)
	}
	defer devnull.Close()
	if len(os.Args) == 1 {
		if term.IsTerminal(0) {
			input := bufio.NewScanner(os.Stdin)
			for {
				os.Stdout.Write([]byte("beamsh> "))
				if !input.Scan() {
					break
				}
				line := input.Text()
				if len(line) != 0 {
					cmd, err := dockerscript.Parse(strings.NewReader(line))
					if err != nil {
						fmt.Fprintf(os.Stderr, "error: %v\n", err)
						continue
					}
					if err := executeScript(devnull, cmd); err != nil {
						Fatal(err)
					}
				}
				if err := input.Err(); err == io.EOF {
					break
				} else if err != nil {
					Fatal(err)
				}
			}
		} else {
			script, err := dockerscript.Parse(os.Stdin)
			if err != nil {
				Fatal("parse error: %v\n", err)
			}
			if err := executeScript(devnull, script); err != nil {
				Fatal(err)
			}
		}
	} else {
		for _, scriptpath := range os.Args[1:] {
			f, err := os.Open(scriptpath)
			if err != nil {
				Fatal(err)
			}
			script, err := dockerscript.Parse(f)
			if err != nil {
				Fatal("parse error: %v\n", err)
			}
			if err := executeScript(devnull, script); err != nil {
				Fatal(err)
			}
		}
	}
}

func beamCopy(dst *net.UnixConn, src *net.UnixConn) (int, error) {
	var n int
	for {
		payload, attachment, err := beam.Receive(src)
		if err == io.EOF {
			return n, nil
		} else if err != nil {
			return n, err
		}
		if err := beam.Send(dst, payload, attachment); err != nil {
			if attachment != nil {
				attachment.Close()
			}
			return n, err
		}
		n++
	}
	panic("impossibru!")
	return n, nil
}

type Handler func([]string, *net.UnixConn, *net.UnixConn)

func Devnull() (*net.UnixConn, error) {
	priv, pub, err := beam.USocketPair()
	if err != nil {
		return nil, err
	}
	go func() {
		defer priv.Close()
		for {
			_, attachment, err := beam.Receive(priv)
			if err != nil {
				return
			}
			if attachment != nil {
				attachment.Close()
			}
		}
	}()
	return pub, nil
}

func scriptString(script []*dockerscript.Command) string {
	lines := make([]string, 0, len(script))
	for _, cmd := range script {
		line := strings.Join(cmd.Args, " ")
		if len(cmd.Children) > 0 {
			line += fmt.Sprintf(" { %s }", scriptString(cmd.Children))
		} else {
			line += " {}"
		}
		lines = append(lines, line)
	}
	return fmt.Sprintf("'%s'", strings.Join(lines, "; "))
}

func executeScript(client *net.UnixConn, script []*dockerscript.Command) error {
	Debugf("executeScript(%s)\n", scriptString(script))
	defer Debugf("executeScript(%s) DONE\n", scriptString(script))
	var background sync.WaitGroup
	defer background.Wait()
	for _, cmd := range script {
		if cmd.Background {
			background.Add(1)
			go func(client *net.UnixConn, cmd *dockerscript.Command) {
				executeCommand(client, cmd)
				background.Done()
			}(client, cmd)
		} else {
			if err := executeCommand(client, cmd); err != nil {
				return err
			}
		}
	}
	return nil
}

func dialer(addr string) (chan net.Conn, error) {
	u, err := url.Parse(addr)
	if err != nil {
		return nil, err
	}
	connections := make(chan net.Conn)
	go func() {
		defer close(connections)
		for {
			conn, err := net.Dial(u.Scheme, u.Host)
			if err != nil {
				return
			}
			connections <-conn
		}
	}()
	return connections, nil
}

func listener(addr string) (chan net.Conn, error) {
	u, err := url.Parse(addr)
	if err != nil {
		return nil, err
	}
	l, err := net.Listen(u.Scheme, u.Host)
	if err != nil {
		return nil, err
	}
	connections := make(chan net.Conn)
	go func() {
		defer close(connections)
		for {
			conn, err := l.Accept()
			if err != nil {
				return
			}
			Logf("new connection\n")
			connections<-conn
		}
	}()
	return connections, nil
}

func msgDesc(payload []byte, attachment *os.File) string {
	var filedesc string = "<nil>"
	if attachment != nil {
		filedesc = fmt.Sprintf("%d", attachment.Fd())
	}
	return fmt.Sprintf("'%s'[%s]", payload, filedesc)

}

func SendToConn(connections chan net.Conn, src *net.UnixConn) error {
	var tasks sync.WaitGroup
	defer tasks.Wait()
	for {
		payload, attachment, err := beam.Receive(src)
		if err == io.EOF {
			return nil
		} else if err != nil {
			return err
		}
		conn, ok := <-connections
		if !ok {
			break
		}
		Logf("Sending %s\n", msgDesc(payload, attachment))
		tasks.Add(1)
		go func(payload []byte, attachment *os.File, conn net.Conn) {
			defer tasks.Done()
			if _, err := conn.Write([]byte(data.EncodeString(string(payload)))); err != nil {
				return
			}
			if attachment == nil {
				conn.Close()
				return
			}
			var iotasks sync.WaitGroup
			iotasks.Add(2)
			go func(attachment *os.File, conn net.Conn) {
				defer iotasks.Done()
				Debugf("copying the connection to [%d]\n", attachment.Fd())
				io.Copy(attachment, conn)
				attachment.Close()
				Debugf("done copying the connection to [%d]\n", attachment.Fd())
			}(attachment, conn)
			go func(attachment *os.File, conn net.Conn) {
				defer iotasks.Done()
				Debugf("copying [%d] to the connection\n", attachment.Fd())
				io.Copy(conn, attachment)
				conn.Close()
				Debugf("done copying [%d] to the connection\n", attachment.Fd())
			}(attachment, conn)
			iotasks.Wait()
		}(payload, attachment, conn)
	}
	return nil
}

func bicopy(a, b io.ReadWriteCloser) {
	var iotasks sync.WaitGroup
	oneCopy := func(dst io.WriteCloser, src io.Reader) {
		defer iotasks.Done()
		io.Copy(dst, src)
		dst.Close()
	}
	iotasks.Add(2)
	go oneCopy(a, b)
	go oneCopy(b, a)
	iotasks.Wait()
}

func ReceiveFromConn(connections chan net.Conn, dst *net.UnixConn) error {
	for conn := range connections {
		err := func () error {
			Logf("parsing message from network...\n")
			defer Logf("done parsing message from network\n")
			buf := make([]byte, 4098)
			n, err := conn.Read(buf)
			if n == 0 {
				conn.Close()
				if err == io.EOF {
					return nil
				} else {
					return err
				}
			}
			Logf("decoding message from '%s'\n", buf[:n])
			header, skip, err := data.DecodeString(string(buf[:n]))
			if err != nil {
				conn.Close()
				return err
			}
			pub, priv, err := beam.SocketPair()
			if err != nil {
				return err
			}
			Logf("decoded message: %s\n", data.Message(header).Pretty())
			go func(skipped []byte, conn net.Conn, f *os.File) {
				// this closes both conn and f
				if len(skipped) > 0 {
					if _, err := f.Write(skipped); err != nil {
						Logf("ERROR: %v\n", err)
						f.Close()
						conn.Close()
						return
					}
				}
				bicopy(conn, f)
			}(buf[skip:n], conn, pub)
			if err := beam.Send(dst, []byte(header), priv); err != nil {
				return err
			}
			return nil
		}()
		if err != nil {
			Logf("Error reading from connection: %v\n", err)
		}
	}
	return nil
}

//	1) Find a handler for the command (if no handler, fail)
//	2) Attach new in & out pair to the handler
//	3) [in the background] Copy handler output to our own output
//	4) [in the background] Run the handler
//	5) Recursively executeScript() all children commands and wait for them to complete
//	6) Wait for handler to return and (shortly afterwards) output copy to complete
//	7) 
func executeCommand(client *net.UnixConn, cmd *dockerscript.Command) error {
	Debugf("executeCommand(%s)\n", strings.Join(cmd.Args, " "))
	defer Debugf("executeCommand(%s) DONE\n", strings.Join(cmd.Args, " "))
	handler := GetHandler(cmd.Args[0])
	if handler == nil {
		return fmt.Errorf("no such command: %s", cmd.Args[0])
	}
	inPub, inPriv, err := beam.USocketPair()
	if err != nil {
		return err
	}
	// Don't close inPub here. We close it to signify the end of input once
	// all children are completed (guaranteeing that no more input will be sent
	// by children).
	// Otherwise we get a deadlock.
	defer inPriv.Close()
	outPub, outPriv, err := beam.USocketPair()
	if err != nil {
		return err
	}
	defer outPub.Close()
	// don't close outPriv here. It must be closed after the handler is called,
	// but before the copy tasks associated with it completes.
	// Otherwise we get a deadlock.
	var tasks sync.WaitGroup
	tasks.Add(2)
	go func() {
		handler(cmd.Args, inPriv, outPriv)
		// FIXME: do we need to outPriv.sync before closing it?
		Debugf("[%s] handler returned, closing output\n", strings.Join(cmd.Args, " "))
		outPriv.Close()
		tasks.Done()
	}()
	go func() {
		Debugf("[%s] copy start...\n", strings.Join(cmd.Args, " "))
		n, err := beamCopy(client, outPub)
		if err != nil {
			Fatal(err)
		}
		Debugf("[%s] copied %d messages\n", strings.Join(cmd.Args, " "), n)
		Debugf("[%s] copy done\n", strings.Join(cmd.Args, " "))
		tasks.Done()
	}()
	// depth-first execution of children commands
	// executeScript() blocks until all commands are completed
	executeScript(inPub, cmd.Children)
	inPub.Close()
	Debugf("[%s] waiting for handler and output copy to complete...\n", strings.Join(cmd.Args, " "))
	tasks.Wait()
	Debugf("[%s] handler and output copy complete!\n", strings.Join(cmd.Args, " "))
	return nil
}

func randomId() string {
	id := make([]byte, 4)
	io.ReadFull(rand.Reader, id)
	return hex.EncodeToString(id)
}

func sendWPipe(conn *net.UnixConn, payload []byte) (*os.File, error) {
	r, w, err := os.Pipe()
	if err != nil {
		return nil, err
	}
	if err := beam.Send(conn, payload, r); err != nil {
		r.Close()
		w.Close()
		return nil, err
	}
	return w, nil
}

func GetHandler(name string) Handler {
	if name == "devnull" {
		return func(args []string, in *net.UnixConn, out *net.UnixConn) {
			for {
				_, attachment, err := beam.Receive(in)
				if err != nil {
					return
				}
				if attachment != nil {
					attachment.Close()
				}
			}
		}
	} else if name == "stdio" {
		return func(args []string, in *net.UnixConn, out *net.UnixConn) {
			var tasks sync.WaitGroup
			defer tasks.Wait()
			for {
				payload, attachment, err := beam.Receive(in)
				if err != nil {
					return
				}
				tasks.Add(1)
				go func(payload []byte, attachment *os.File) {
					defer tasks.Done()
					if attachment == nil {
						return
					}
					defer attachment.Close()
					cmd := data.Message(payload).Get("cmd")
					if cmd == nil || len(cmd) == 0 {
						return
					}
					if cmd[0] != "log" {
						return
					}
					if len(cmd) == 1 || cmd[1] == "stdout" {
						io.Copy(os.Stdout, attachment)
					} else if cmd[1] == "stderr" {
						io.Copy(os.Stderr, attachment)
					}
				}(payload, attachment)
			}
		}
	} else if name == "echo" {
		return func(args []string, in *net.UnixConn, out *net.UnixConn) {
			stdout, err := sendWPipe(out, data.Empty().Set("cmd", "log", "stdout").Bytes())
			if err != nil {
				return
			}
			fmt.Fprintln(stdout, strings.Join(args[1:], " "))
			stdout.Close()
		}
	} else if name == "pass" {
		return func(args []string, in *net.UnixConn, out *net.UnixConn) {
			for {
				payload, attachment, err := beam.Receive(in)
				if err != nil {
					return
				}
				if err := beam.Send(out, payload, attachment); err != nil {
					if attachment != nil {
						attachment.Close()
					}
					return
				}
			}
		}
	} else if name == "in" {
		return func(args []string, in *net.UnixConn, out *net.UnixConn) {
			os.Chdir(args[1])
			GetHandler("pass")([]string{"pass"}, in, out)
		}
	} else if name == "exec" {
		return func(args []string, in *net.UnixConn, out *net.UnixConn) {
			cmd := exec.Command(args[1], args[2:]...)
			outR, outW, err := os.Pipe()
			if err != nil {
				return
			}
			cmd.Stdout = outW
			errR, errW, err := os.Pipe()
			if err != nil {
				return
			}
			cmd.Stderr = errW
			cmd.Stdin = os.Stdin
			beam.Send(out, data.Empty().Set("cmd", "log", "stdout").Bytes(), outR)
			beam.Send(out, data.Empty().Set("cmd", "log", "stderr").Bytes(), errR)
			execErr := cmd.Run()
			var status string
			if execErr != nil {
				status = execErr.Error()
			} else {
				status = "ok"
			}
			beam.Send(out, data.Empty().Set("status", status).Set("cmd", args...).Bytes(), nil)
			outW.Close()
			errW.Close()
		}
	} else if name == "trace" {
		return func(args []string, in *net.UnixConn, out *net.UnixConn) {
			for {
				p, a, err := beam.Receive(in)
				if err != nil {
					return
				}
				var msg string
				if pretty := data.Message(string(p)).Pretty(); pretty != "" {
					msg = pretty
				} else {
					msg = string(p)
				}
				if a != nil {
					msg = fmt.Sprintf("%s [%d]", msg, a.Fd())
				}
				fmt.Printf("===> %s\n", msg)
				beam.Send(out, p, a)
			}
		}
	} else if name == "emit" {
		return func(args []string, in *net.UnixConn, out *net.UnixConn) {
			beam.Send(out, data.Parse(args[1:]).Bytes(), nil)
		}
	} else if name == "print" {
		return func(args []string, in *net.UnixConn, out *net.UnixConn) {
			for {
				_, a, err := beam.Receive(in)
				if err != nil {
					return
				}
				if a != nil {
					io.Copy(os.Stdout, a)
				}
			}
		}
	} else if name == "multiprint" {
		return func(args []string, in *net.UnixConn, out *net.UnixConn) {
			var tasks sync.WaitGroup
			for {
				payload, a, err := beam.Receive(in)
				if err != nil {
					return
				}
				if a != nil {
					tasks.Add(1)
					go func(payload []byte, attachment *os.File) {
						defer tasks.Done()
						msg := data.Message(string(payload))
						input := bufio.NewScanner(attachment)
						for input.Scan() {
							fmt.Printf("[%s] %s\n", msg.Pretty(), input.Text())
						}
					}(payload, a)
				}
			}
			tasks.Wait()
		}
	} else if name == "listen" {
		return func(args []string, in *net.UnixConn, out *net.UnixConn) {
			if len(args) != 2 {
				beam.Send(out, data.Empty().Set("status", "1").Set("message", "wrong number of arguments").Bytes(), nil)
				return
			}
			u, err := url.Parse(args[1])
			if err != nil {
				beam.Send(out, data.Empty().Set("status", "1").Set("message", err.Error()).Bytes(), nil)
				return
			}
			l, err := net.Listen(u.Scheme, u.Host)
			if err != nil {
				beam.Send(out, data.Empty().Set("status", "1").Set("message", err.Error()).Bytes(), nil)
				return
			}
			for {
				conn, err := l.Accept()
				if err != nil {
					beam.Send(out, data.Empty().Set("status", "1").Set("message", err.Error()).Bytes(), nil)
					return
				}
				f, err := connToFile(conn)
				if err != nil {
					conn.Close()
					continue
				}
				beam.Send(out, data.Empty().Set("type", "socket").Set("remoteaddr", conn.RemoteAddr().String()).Bytes(), f)
			}
		}
	} else if name == "beamsend" {
		return func(args []string, in *net.UnixConn, out *net.UnixConn) {
			if len(args) < 2 {
				if err := beam.Send(out, data.Empty().Set("status", "1").Set("message", "wrong number of arguments").Bytes(), nil); err != nil {
					Fatal(err)
				}
				return
			}
			var connector func(string) (chan net.Conn, error)
			connector = dialer
			connections, err := connector(args[1])
			if err != nil {
				beam.Send(out, data.Empty().Set("status", "1").Set("message", err.Error()).Bytes(), nil)
				return
			}
			// Copy in to conn
			SendToConn(connections, in)
		}
	} else if name == "beamreceive" {
		return func(args []string, in *net.UnixConn, out *net.UnixConn) {
			if len(args) != 2 {
				if err := beam.Send(out, data.Empty().Set("status", "1").Set("message", "wrong number of arguments").Bytes(), nil); err != nil {
					Fatal(err)
				}
				return
			}
			var connector func(string) (chan net.Conn, error)
			connector = listener
			connections, err := connector(args[1])
			if err != nil {
				beam.Send(out, data.Empty().Set("status", "1").Set("message", err.Error()).Bytes(), nil)
				return
			}
			// Copy in to conn
			ReceiveFromConn(connections, out)
		}
	} else if name == "connect" {
		return func(args []string, in *net.UnixConn, out *net.UnixConn) {
			if len(args) != 2 {
				beam.Send(out, data.Empty().Set("status", "1").Set("message", "wrong number of arguments").Bytes(), nil)
				return
			}
			u, err := url.Parse(args[1])
			if err != nil {
				beam.Send(out, data.Empty().Set("status", "1").Set("message", err.Error()).Bytes(), nil)
				return
			}
			var tasks sync.WaitGroup
			for {
				_, attachment, err := beam.Receive(in)
				if err != nil {
					break
				}
				if attachment == nil {
					continue
				}
				Logf("connecting to %s/%s\n", u.Scheme, u.Host)
				conn, err := net.Dial(u.Scheme, u.Host)
				if err != nil {
					beam.Send(out, data.Empty().Set("cmd", "msg", "connect error: " + err.Error()).Bytes(), nil)
					return
				}
				beam.Send(out, data.Empty().Set("cmd", "msg", "connection established").Bytes(), nil)
				tasks.Add(1)
				go func(attachment *os.File, conn net.Conn) {
					defer tasks.Done()
					// even when successful, conn.File() returns a duplicate,
					// so we must close the original
					var iotasks sync.WaitGroup
					iotasks.Add(2)
					go func(attachment *os.File, conn net.Conn) {
						defer iotasks.Done()
						io.Copy(attachment, conn)
					}(attachment, conn)
					go func(attachment *os.File, conn net.Conn) {
						defer iotasks.Done()
						io.Copy(conn, attachment)
					}(attachment, conn)
					iotasks.Wait()
					conn.Close()
					attachment.Close()
				}(attachment, conn)
			}
			tasks.Wait()
		}
	} else if name == "openfile" {
		return func(args []string, in *net.UnixConn, out *net.UnixConn) {
			for _, name := range args {
				f, err := os.Open(name)
				if err != nil {
					continue
				}
				if err := beam.Send(out, data.Empty().Set("path", name).Set("type", "file").Bytes(), f); err != nil {
					f.Close()
				}
			}
		}
	}
	return nil
}

func connToFile(conn net.Conn) (f *os.File, err error) {
	if connWithFile, ok := conn.(interface { File() (*os.File, error) }); !ok {
		return nil, fmt.Errorf("no file descriptor available")
	} else {
		f, err = connWithFile.File()
		if err != nil {
			return nil, err
		}
	}
	return f, err
}

type Msg struct {
	payload		[]byte
	attachment	*os.File
}

func Logf(msg string, args ...interface{}) (int, error) {
	if len(msg) == 0 || msg[len(msg) - 1] != '\n' {
		msg = msg + "\n"
	}
	msg = fmt.Sprintf("[%v] [%v] %s", os.Getpid(), path.Base(os.Args[0]), msg)
	return fmt.Printf(msg, args...)
}

func Debugf(msg string, args ...interface{}) {
	if os.Getenv("BEAMDEBUG") != "" {
		Logf(msg, args...)
	}
}

func Fatalf(msg string, args ...interface{})  {
	Logf(msg, args)
	os.Exit(1)
}

func Fatal(args ...interface{}) {
	Fatalf("%v", args[0])
}