510 lines
13 KiB
Go
510 lines
13 KiB
Go
package main
|
|
|
|
import (
|
|
"bufio"
|
|
"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"
|
|
"path"
|
|
"strings"
|
|
"sync"
|
|
"flag"
|
|
)
|
|
|
|
var rootPlugins = []string{
|
|
"stdio",
|
|
}
|
|
|
|
var (
|
|
flX bool
|
|
)
|
|
|
|
func main() {
|
|
flag.BoolVar(&flX, "x", false, "print commands as they are being executed")
|
|
flag.Parse()
|
|
if flag.NArg() == 0{
|
|
if term.IsTerminal(0) {
|
|
// No arguments, stdin is terminal --> interactive mode
|
|
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 := executeRootScript(cmd); err != nil {
|
|
Fatal(err)
|
|
}
|
|
}
|
|
if err := input.Err(); err == io.EOF {
|
|
break
|
|
} else if err != nil {
|
|
Fatal(err)
|
|
}
|
|
}
|
|
} else {
|
|
// No arguments, stdin not terminal --> batch mode
|
|
script, err := dockerscript.Parse(os.Stdin)
|
|
if err != nil {
|
|
Fatal("parse error: %v\n", err)
|
|
}
|
|
if err := executeRootScript(script); err != nil {
|
|
Fatal(err)
|
|
}
|
|
}
|
|
} else {
|
|
// 1+ arguments: parse them as script files
|
|
for _, scriptpath := range flag.Args() {
|
|
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 := executeRootScript(script); err != nil {
|
|
Fatal(err)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func executeRootScript(script []*dockerscript.Command) error {
|
|
if len(rootPlugins) > 0 {
|
|
// If there are root plugins, wrap the script inside them
|
|
var (
|
|
rootCmd *dockerscript.Command
|
|
lastCmd *dockerscript.Command
|
|
)
|
|
for _, plugin := range rootPlugins {
|
|
pluginCmd := &dockerscript.Command{
|
|
Args: []string{plugin},
|
|
}
|
|
if rootCmd == nil {
|
|
rootCmd = pluginCmd
|
|
} else {
|
|
lastCmd.Children = []*dockerscript.Command{pluginCmd}
|
|
}
|
|
lastCmd = pluginCmd
|
|
}
|
|
lastCmd.Children = script
|
|
script = []*dockerscript.Command{rootCmd}
|
|
}
|
|
handlers, err := Handlers()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer handlers.Close()
|
|
if err := executeScript(handlers, script); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func executeScript(out beam.Sender, 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(out beam.Sender, cmd *dockerscript.Command) {
|
|
executeCommand(out, cmd)
|
|
background.Done()
|
|
}(out, cmd)
|
|
} else {
|
|
if err := executeCommand(out, cmd); err != nil {
|
|
return 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) Profit
|
|
func executeCommand(out beam.Sender, cmd *dockerscript.Command) error {
|
|
if flX {
|
|
fmt.Printf("+ %v\n", strings.Replace(strings.TrimRight(cmd.String(), "\n"), "\n", "\n+ ", -1))
|
|
}
|
|
Debugf("executeCommand(%s)\n", strings.Join(cmd.Args, " "))
|
|
defer Debugf("executeCommand(%s) DONE\n", strings.Join(cmd.Args, " "))
|
|
if len(cmd.Args) == 0 {
|
|
return fmt.Errorf("empty command")
|
|
}
|
|
Debugf("[executeCommand] sending job '%s'\n", strings.Join(cmd.Args, " "))
|
|
job, err := beam.SendConn(out, data.Empty().Set("cmd", cmd.Args...).Bytes())
|
|
if err != nil {
|
|
return fmt.Errorf("%v\n", err)
|
|
}
|
|
var tasks sync.WaitGroup
|
|
tasks.Add(1)
|
|
Debugf("[executeCommand] spawning background copy of the output of '%s'\n", strings.Join(cmd.Args, " "))
|
|
go func() {
|
|
if out != nil {
|
|
Debugf("[executeCommand] background copy of the output of '%s'\n", strings.Join(cmd.Args, " "))
|
|
n, err := beam.Copy(out, job)
|
|
if err != nil {
|
|
Fatalf("[executeCommand] [%s] error during background copy: %v\n", strings.Join(cmd.Args, " "), err)
|
|
}
|
|
Debugf("[executeCommand] background copy done of the output of '%s': copied %d messages\n", strings.Join(cmd.Args, " "), n)
|
|
}
|
|
tasks.Done()
|
|
}()
|
|
// depth-first execution of children commands
|
|
// executeScript() blocks until all commands are completed
|
|
Debugf("[executeCommand] recursively running children of '%s'\n", strings.Join(cmd.Args, " "))
|
|
executeScript(job, cmd.Children)
|
|
Debugf("[executeCommand] DONE recursively running children of '%s'\n", strings.Join(cmd.Args, " "))
|
|
job.CloseWrite()
|
|
Debugf("[executeCommand] closing the input of '%s' (all children are completed)\n", strings.Join(cmd.Args, " "))
|
|
Debugf("[executeCommand] waiting for background copy of '%s' to complete...\n", strings.Join(cmd.Args, " "))
|
|
tasks.Wait()
|
|
Debugf("[executeCommand] background copy of '%s' complete! This means the job completed.\n", strings.Join(cmd.Args, " "))
|
|
return nil
|
|
}
|
|
|
|
|
|
type Handler func([]string, io.Writer, io.Writer, beam.Receiver, beam.Sender)
|
|
|
|
|
|
func Handlers() (*beam.UnixConn, error) {
|
|
var tasks sync.WaitGroup
|
|
pub, priv, err := beam.USocketPair()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
go func() {
|
|
defer func() {
|
|
Debugf("[handlers] closewrite() on endpoint\n")
|
|
// FIXME: this is not yet necessary but will be once
|
|
// there is synchronization over standard beam messages
|
|
priv.CloseWrite()
|
|
Debugf("[handlers] done closewrite() on endpoint\n")
|
|
}()
|
|
for {
|
|
Debugf("[handlers] waiting for next job...\n")
|
|
payload, conn, err := beam.ReceiveConn(priv)
|
|
Debugf("[handlers] ReceiveConn() returned %v\n", err)
|
|
if err != nil {
|
|
return
|
|
}
|
|
tasks.Add(1)
|
|
go func(payload []byte, conn *beam.UnixConn) {
|
|
defer tasks.Done()
|
|
defer func() {
|
|
Debugf("[handlers] '%s' closewrite\n", payload)
|
|
conn.CloseWrite()
|
|
Debugf("[handlers] '%s' done closewrite\n", payload)
|
|
}()
|
|
cmd := data.Message(payload).Get("cmd")
|
|
Debugf("[handlers] received %s\n", strings.Join(cmd, " "))
|
|
if len(cmd) == 0 {
|
|
return
|
|
}
|
|
handler := GetHandler(cmd[0])
|
|
if handler == nil {
|
|
return
|
|
}
|
|
stdout, err := beam.SendPipe(conn, data.Empty().Set("cmd", "log", "stdout").Set("fromcmd", cmd...).Bytes())
|
|
if err != nil {
|
|
return
|
|
}
|
|
defer stdout.Close()
|
|
stderr, err := beam.SendPipe(conn, data.Empty().Set("cmd", "log", "stderr").Set("fromcmd", cmd...).Bytes())
|
|
if err != nil {
|
|
return
|
|
}
|
|
defer stderr.Close()
|
|
Debugf("[handlers] calling %s\n", strings.Join(cmd, " "))
|
|
handler(cmd, stdout, stderr, beam.Receiver(conn), beam.Sender(conn))
|
|
Debugf("[handlers] returned: %s\n", strings.Join(cmd, " "))
|
|
}(payload, conn)
|
|
}
|
|
Debugf("[handlers] waiting for all tasks\n")
|
|
tasks.Wait()
|
|
Debugf("[handlers] all tasks returned\n")
|
|
}()
|
|
return pub, nil
|
|
}
|
|
|
|
func GetHandler(name string) Handler {
|
|
if name == "logger" {
|
|
return CmdLogger
|
|
} else if name == "render" {
|
|
return CmdRender
|
|
} else if name == "devnull" {
|
|
return CmdDevnull
|
|
} else if name == "prompt" {
|
|
return CmdPrompt
|
|
} else if name == "stdio" {
|
|
return CmdStdio
|
|
} else if name == "echo" {
|
|
return CmdEcho
|
|
} else if name == "pass" {
|
|
return CmdPass
|
|
} else if name == "in" {
|
|
return CmdIn
|
|
} else if name == "exec" {
|
|
return CmdExec
|
|
} else if name == "trace" {
|
|
return CmdTrace
|
|
} else if name == "emit" {
|
|
return CmdEmit
|
|
} else if name == "print" {
|
|
return CmdPrint
|
|
} else if name == "multiprint" {
|
|
return CmdMultiprint
|
|
} else if name == "listen" {
|
|
return CmdListen
|
|
} else if name == "beamsend" {
|
|
return CmdBeamsend
|
|
} else if name == "beamreceive" {
|
|
return CmdBeamreceive
|
|
} else if name == "connect" {
|
|
return CmdConnect
|
|
} else if name == "openfile" {
|
|
return CmdOpenfile
|
|
}
|
|
return nil
|
|
}
|
|
|
|
|
|
// VARIOUS HELPER FUNCTIONS:
|
|
|
|
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])
|
|
}
|
|
|
|
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 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 SendToConn(connections chan net.Conn, src beam.Receiver) error {
|
|
var tasks sync.WaitGroup
|
|
defer tasks.Wait()
|
|
for {
|
|
payload, attachment, err := src.Receive()
|
|
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 msgDesc(payload []byte, attachment *os.File) string {
|
|
return beam.MsgDesc(payload, attachment)
|
|
}
|
|
|
|
func ReceiveFromConn(connections chan net.Conn, dst beam.Sender) 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 := dst.Send([]byte(header), priv); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}()
|
|
if err != nil {
|
|
Logf("Error reading from connection: %v\n", err)
|
|
}
|
|
}
|
|
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()
|
|
}
|
|
|