Implement TTY resize for hosts

common/xfer/controls.go

@@ -3,6 +3,7 @@ package xfer
 import (
 	"fmt"
 	"net/rpc"
+	"strconv"
 	"sync"
 )
 
@@ -54,6 +55,41 @@ func (c ControlHandlerFunc) Handle(req Request, res *Response) error {
 	return nil
 }
 
+// ResizeTTYControlWrapper extracts the arguments needed by the resize tty control handler
+func ResizeTTYControlWrapper(next func(pipeID string, height, width uint) Response) ControlHandlerFunc {
+	return func(req Request) Response {
+		var (
+			height, width uint64
+			err           error
+		)
+
+		pipeID, ok := req.ControlArgs["pipeID"]
+		if !ok {
+			return ResponseErrorf("Missing argument: pipeID")
+		}
+		heightS, ok := req.ControlArgs["height"]
+		if !ok {
+			return ResponseErrorf("Missing argument: height")
+		}
+		widthS, ok := req.ControlArgs["width"]
+		if !ok {
+			return ResponseErrorf("Missing argument: width")
+		}
+
+		height, err = strconv.ParseUint(heightS, 10, 32)
+		if err != nil {
+			return ResponseErrorf("Bad parameter: height (%q): %v", heightS, err)
+		}
+		width, err = strconv.ParseUint(widthS, 10, 32)
+		if err != nil {
+			return ResponseErrorf("Bad parameter: width (%q): %v", widthS, err)
+		}
+
+		return next(pipeID, uint(height), uint(width))
+
+	}
+}
+
 // ResponseErrorf creates a new Response with the given formatted error string.
 func ResponseErrorf(format string, a ...interface{}) Response {
 	return Response{

probe/docker/controls.go

@@ -1,8 +1,6 @@
 package docker
 
 import (
-	"strconv"
-
 	docker_client "github.com/fsouza/go-dockerclient"
 
 	log "github.com/Sirupsen/logrus"
@@ -164,34 +162,7 @@ func (r *registry) execContainer(containerID string, req xfer.Request) xfer.Resp
 	}
 }
 
-func (r *registry) resizeExecTTY(containerID string, req xfer.Request) xfer.Response {
-	var (
-		height, width int
-		err           error
-	)
-
-	pipeID, ok := req.ControlArgs["pipeID"]
-	if !ok {
-		return xfer.ResponseErrorf("Missing argument: pipeID")
-	}
-	heightS, ok := req.ControlArgs["height"]
-	if !ok {
-		return xfer.ResponseErrorf("Missing argument: height")
-	}
-	widthS, ok := req.ControlArgs["width"]
-	if !ok {
-		return xfer.ResponseErrorf("Missing argument: width")
-	}
-
-	height, err = strconv.Atoi(heightS)
-	if err != nil {
-		return xfer.ResponseErrorf("Bad parameter: height (%q): %v", heightS, err)
-	}
-	width, err = strconv.Atoi(widthS)
-	if err != nil {
-		return xfer.ResponseErrorf("Bad parameter: width (%q): %v", widthS, err)
-	}
-
+func (r *registry) resizeExecTTY(pipeID string, height, width uint) xfer.Response {
 	r.Lock()
 	execID, ok := r.pipeIDToexecID[pipeID]
 	r.Unlock()
@@ -200,7 +171,7 @@ func (r *registry) resizeExecTTY(containerID string, req xfer.Request) xfer.Resp
 		return xfer.ResponseErrorf("Unknown pipeID (%q)", pipeID)
 	}
 
-	if r.client.ResizeExecTTY(execID, int(height), int(width)); err != nil {
+	if err := r.client.ResizeExecTTY(execID, int(height), int(width)); err != nil {
 		return xfer.ResponseErrorf(
 			"Error setting terminal size (%d, %d) of pipe %s: %v",
 			height, width, pipeID, err)
@@ -229,7 +200,7 @@ func (r *registry) registerControls() {
 		RemoveContainer:  captureContainerID(r.removeContainer),
 		AttachContainer:  captureContainerID(r.attachContainer),
 		ExecContainer:    captureContainerID(r.execContainer),
-		ResizeExecTTY:    captureContainerID(r.resizeExecTTY),
+		ResizeExecTTY:    xfer.ResizeTTYControlWrapper(r.resizeExecTTY),
 	}
 	r.handlerRegistry.Batch(nil, controls)
 }

probe/host/controls.go

@@ -4,6 +4,7 @@ import (
 	"os/exec"
 
 	log "github.com/Sirupsen/logrus"
+	"github.com/docker/docker/pkg/term"
 	"github.com/kr/pty"
 
 	"github.com/weaveworks/scope/common/xfer"
@@ -12,15 +13,18 @@ import (
 
 // Control IDs used by the host integration.
 const (
-	ExecHost = "host_exec"
+	ExecHost      = "host_exec"
+	ResizeExecTTY = "host_resize_exec_tty"
 )
 
 func (r *Reporter) registerControls() {
 	r.handlerRegistry.Register(ExecHost, r.execHost)
+	r.handlerRegistry.Register(ResizeExecTTY, xfer.ResizeTTYControlWrapper(r.resizeExecTTY))
 }
 
 func (r *Reporter) deregisterControls() {
 	r.handlerRegistry.Rm(ExecHost)
+	r.handlerRegistry.Rm(ResizeExecTTY)
 }
 
 func (r *Reporter) execHost(req xfer.Request) xfer.Response {
@@ -35,6 +39,11 @@ func (r *Reporter) execHost(req xfer.Request) xfer.Response {
 	if err != nil {
 		return xfer.ResponseError(err)
 	}
+
+	r.Lock()
+	r.pipeIDToTTY[id] = ptyPipe.Fd()
+	r.Unlock()
+
 	pipe.OnClose(func() {
 		if err := cmd.Process.Kill(); err != nil {
 			log.Errorf("Error stopping host shell: %v", err)
@@ -42,6 +51,9 @@ func (r *Reporter) execHost(req xfer.Request) xfer.Response {
 		if err := ptyPipe.Close(); err != nil {
 			log.Errorf("Error closing host shell's pty: %v", err)
 		}
+		r.Lock()
+		delete(r.pipeIDToTTY, id)
+		r.Unlock()
 		log.Info("Host shell closed.")
 	})
 	go func() {
@@ -52,7 +64,32 @@ func (r *Reporter) execHost(req xfer.Request) xfer.Response {
 	}()
 
 	return xfer.Response{
-		Pipe:   id,
-		RawTTY: true,
+		Pipe:             id,
+		RawTTY:           true,
+		ResizeTTYControl: ResizeExecTTY,
 	}
 }
+
+func (r *Reporter) resizeExecTTY(pipeID string, height, width uint) xfer.Response {
+	r.Lock()
+	fd, ok := r.pipeIDToTTY[pipeID]
+	r.Unlock()
+
+	if !ok {
+		return xfer.ResponseErrorf("Unknown pipeID (%q)", pipeID)
+	}
+
+	size := term.Winsize{
+		Height: uint16(height),
+		Width:  uint16(width),
+	}
+
+	if err := term.SetWinsize(fd, &size); err != nil {
+		return xfer.ResponseErrorf(
+			"Error setting terminal size (%d, %d) of pipe %s: %v",
+			height, width, pipeID, err)
+	}
+
+	return xfer.Response{}
+
+}

probe/host/reporter.go

@@ -3,6 +3,7 @@ package host
 import (
 	"net"
 	"runtime"
+	"sync"
 	"time"
 
 	"github.com/weaveworks/scope/common/mtime"
@@ -52,6 +53,7 @@ var (
 
 // Reporter generates Reports containing the host topology.
 type Reporter struct {
+	sync.RWMutex
 	hostID          string
 	hostName        string
 	probeID         string
@@ -59,6 +61,7 @@ type Reporter struct {
 	pipes           controls.PipeClient
 	hostShellCmd    []string
 	handlerRegistry *controls.HandlerRegistry
+	pipeIDToTTY     map[string]uintptr
 }
 
 // NewReporter returns a Reporter which produces a report containing host
@@ -72,6 +75,7 @@ func NewReporter(hostID, hostName, probeID, version string, pipes controls.PipeC
 		version:         version,
 		hostShellCmd:    getHostShellCmd(),
 		handlerRegistry: handlerRegistry,
+		pipeIDToTTY:     map[string]uintptr{},
 	}
 	r.registerControls()
 	return r