From b3e7dcae3beedd2a54563421898ca2d762145d5e Mon Sep 17 00:00:00 2001
From: Marcos Lilljedahl <marcosnils@gmail.com>
Date: Sat, 8 Oct 2016 14:26:25 +0200
Subject: [PATCH] Make session replication

---
 cookoo/multi.go    | 70 ++++++++++++++++++++++++++++++++++++++++++++++
 handlers/exec.go   | 51 ++++++++++++++++++++++-----------
 services/docker.go |  2 +-
 types/session.go   | 12 ++++++--
 4 files changed, 114 insertions(+), 21 deletions(-)
 create mode 100644 cookoo/multi.go

diff --git a/cookoo/multi.go b/cookoo/multi.go
new file mode 100644
index 0000000..84075c9
--- /dev/null
+++ b/cookoo/multi.go
@@ -0,0 +1,70 @@
+package cookoo
+
+import (
+	"fmt"
+	"io"
+	"os"
+)
+
+// MultiWriter enables you to have a writer that passes on the writing to one
+// of more Writers where the write is duplicated to each Writer. MultiWriter
+// is similar to the multiWriter that is part of Go. The difference is
+// this MultiWriter allows you to manager the Writers attached to it via CRUD
+// operations. To do this you will need to mock the type. For example,
+// mw := NewMultiWriter()
+// mw.(*MultiWriter).AddWriter("foo", foo)
+type MultiWriter struct {
+	writers map[string]io.Writer
+}
+
+// Write sends the bytes to each of the attached writers to be written.
+func (t *MultiWriter) Write(p []byte) (n int, err error) {
+	for name, w := range t.writers {
+		n, err = w.Write(p)
+		if err != nil {
+			// One broken logger should not stop the others.
+			fmt.Fprintf(os.Stderr, "Error logging to '%s': %s", name, err)
+			continue
+		}
+		if n < len(p) {
+			// One broken logger should not stop the others.
+			err = io.ErrShortWrite
+			fmt.Fprintf(os.Stderr, "Short write logging to '%s': Expected to write %d (%V), wrote %d", name, len(p), w, n)
+			continue
+		}
+	}
+	return len(p), err
+}
+
+// Init initializes the MultiWriter.
+func (t *MultiWriter) Init() *MultiWriter {
+	t.writers = make(map[string]io.Writer)
+	return t
+}
+
+// Writer retrieves a given io.Writer given its name.
+func (t *MultiWriter) Writer(name string) (io.Writer, bool) {
+	value, found := t.writers[name]
+	return value, found
+}
+
+// Writers retrieves a map of all io.Writers keyed by name.
+func (t *MultiWriter) Writers() map[string]io.Writer {
+	return t.writers
+}
+
+// AddWriter adds an io.Writer with an associated name.
+func (t *MultiWriter) AddWriter(name string, writer io.Writer) {
+	t.writers[name] = writer
+}
+
+// RemoveWriter removes an io.Writer given a name.
+func (t *MultiWriter) RemoveWriter(name string) {
+	delete(t.writers, name)
+}
+
+// NewMultiWriter returns an initialized MultiWriter.
+func NewMultiWriter() io.Writer {
+	w := new(MultiWriter).Init()
+	return w
+}
diff --git a/handlers/exec.go b/handlers/exec.go
index 5518e3a..3ab42b6 100644
--- a/handlers/exec.go
+++ b/handlers/exec.go
@@ -6,8 +6,10 @@ import (
 	"golang.org/x/net/context"
 	"golang.org/x/net/websocket"
 
+	"github.com/franela/play-with-docker/cookoo"
 	"github.com/franela/play-with-docker/services"
 	"github.com/go-zoo/bone"
+	"github.com/twinj/uuid"
 )
 
 // Echo the data received on the WebSocket.
@@ -20,26 +22,41 @@ func Exec(ws *websocket.Conn) {
 	session := services.GetSession(sessionId)
 	instance := services.GetInstance(session, instanceId)
 
-	if instance.ExecId == "" {
-		execId, err := services.CreateExecConnection(instance.Name, ctx)
+	if instance.Stdout == nil {
+		id, err := services.CreateExecConnection(instance.Name, ctx)
 		if err != nil {
 			return
 		}
-		instance.ExecId = execId
-	}
-	conn, err := services.AttachExecConnection(instance.ExecId, ctx)
-	if err != nil {
-		return
+		conn, err := services.AttachExecConnection(id, ctx)
+		if err != nil {
+			return
+		}
+
+		instance.Conn = conn
+		instance.Stdout = &cookoo.MultiWriter{}
+		instance.Stdout.Init()
+		u1 := uuid.NewV4()
+		instance.Stdout.AddWriter(u1.String(), ws)
+		go func() {
+			io.Copy(instance.Stdout, instance.Conn.Reader)
+		}()
+		defer conn.Close()
+		go func() {
+			io.Copy(instance.Conn.Conn, ws)
+		}()
+		select {
+		case <-ctx.Done():
+		}
+	} else {
+		u1 := uuid.NewV4()
+		instance.Stdout.AddWriter(u1.String(), ws)
+
+		go func() {
+			io.Copy(instance.Conn.Conn, ws)
+		}()
+		select {
+		case <-ctx.Done():
+		}
 	}
 
-	defer conn.Close()
-	go func() {
-		io.Copy(ws, conn.Reader)
-	}()
-	go func() {
-		io.Copy(conn.Conn, ws)
-	}()
-	select {
-	case <-ctx.Done():
-	}
 }
diff --git a/services/docker.go b/services/docker.go
index 7f232d6..9734c4e 100644
--- a/services/docker.go
+++ b/services/docker.go
@@ -60,7 +60,7 @@ func CreateExecConnection(id string, ctx context.Context) (string, error) {
 }
 
 func AttachExecConnection(execId string, ctx context.Context) (*types.HijackedResponse, error) {
-	conf := types.ExecConfig{Tty: true, AttachStdin: true, AttachStderr: true, AttachStdout: true, Cmd: []string{"sh"}}
+	conf := types.ExecConfig{Tty: true, AttachStdin: true, AttachStderr: true, AttachStdout: true}
 	conn, err := c.ContainerExecAttach(ctx, execId, conf)
 
 	if err != nil {
diff --git a/types/session.go b/types/session.go
index ad49c22..61cb120 100644
--- a/types/session.go
+++ b/types/session.go
@@ -1,12 +1,18 @@
 package types
 
+import (
+	"github.com/docker/docker/api/types"
+	"github.com/franela/play-with-docker/cookoo"
+)
+
 type Session struct {
 	Id        string               `json:"id"`
 	Instances map[string]*Instance `json:"instances"`
 }
 
 type Instance struct {
-	Name   string `json:"name"`
-	IP     string `json:"ip"`
-	ExecId string `json:"-"`
+	Name   string                  `json:"name"`
+	IP     string                  `json:"ip"`
+	Stdout *cookoo.MultiWriter     `json:"-"`
+	Conn   *types.HijackedResponse `json:"-"`
 }