beam/examples/beamsh: use beam.Router to simplify the 'stdio' command

Docker-DCO-1.1-Signed-off-by: Solomon Hykes <solomon@docker.com> (github: shykes)

pkg/beam/examples/beamsh/beamsh.go

@@ -407,44 +407,23 @@ func GetHandler(name string) Handler {
 			defer stderr.Close()
 			var tasks sync.WaitGroup
 			defer tasks.Wait()
-			for {
+
-				payload, attachment, err := in.Receive()
+			r := beam.NewRouter(out)
-				if err != nil {
+			r.NewRoute().HasAttachment().KeyStartsWith("cmd", "log").Handler(func(payload []byte, attachment *os.File) error {
-					return
+				tasks.Add(1)
-				}
+				go func() {
-				cmd := data.Message(payload).Get("cmd")
+					defer tasks.Done()
-				if attachment != nil && len(cmd) > 0 && cmd[0] == "log" {
+					defer attachment.Close()
-					w, err := beam.SendPipe(out, payload)
+					io.Copy(os.Stdout, attachment)
-					if err != nil {
+					attachment.Close()
-						attachment.Close()
+				}()
-						fmt.Fprintf(stderr, "sendpipe: %v\n", err)
+				return nil
-						return
+			}).Tee(out)
-					}
+
-					tasks.Add(1)
+			if _, err := beam.Copy(r, in); err != nil {
-					go func(payload []byte, attachment *os.File, sink *os.File) {
+				Fatal(err)
-						defer tasks.Done()
+				fmt.Fprintf(stderr, "%v\n", err)
-						defer attachment.Close()
+				return
-						defer sink.Close()
-						cmd := data.Message(payload).Get("cmd")
-						if cmd == nil || len(cmd) == 0 {
-							return
-						}
-						if cmd[0] != "log" {
-							return
-						}
-						var output io.Writer
-						if len(cmd) == 1 || cmd[1] == "stdout" {
-							output = os.Stdout
-						} else if cmd[1] == "stderr" {
-							output = os.Stderr
-						}
-						io.Copy(io.MultiWriter(output, sink), attachment)
-					}(payload, attachment, w)
-				} else {
-					if err := out.Send(payload, attachment); err != nil {
-						return
-					}
-				}
 			}
 		}
 	} else if name == "echo" {