cmd/gomote: add gomote rdp subcommand to open an RDP proxy to a Windows buildlet

Fixes golang/go#26090

Change-Id: Ib58fb7767384778b67d0d62d34c1132d70d75c23
Reviewed-on: https://go-review.googlesource.com/c/build/+/207378
Run-TryBot: Brad Fitzpatrick <[email protected]>
TryBot-Result: Gobot Gobot <[email protected]>
Reviewed-by: Bryan C. Mills <[email protected]>

Author: bradfitz

buildlet/buildletclient.go

@@ -188,7 +188,7 @@ func (c *Client) String() string {
 
 // RemoteName returns the name of this client's buildlet on the
 // coordinator. If this buildlet isn't a remote buildlet created via
-// a buildlet, this returns the empty string.
+// gomote, this returns the empty string.
 func (c *Client) RemoteName() string {
 	return c.remoteBuildlet
 }
@@ -258,6 +258,37 @@ func (c *Client) do(req *http.Request) (*http.Response, error) {
 	return c.httpClient.Do(req)
 }
 
+// ProxyTCP connects to the given port on the remote buildlet.
+// The buildlet client must currently be a gomote client (RemoteName != "")
+// and the target type must be a VM type running on GCE. This was primarily
+// created for RDP to Windows machines, but it might get reused for other
+// purposes in the future.
+func (c *Client) ProxyTCP(port int) (io.ReadWriteCloser, error) {
+	if c.RemoteName() == "" {
+		return nil, errors.New("ProxyTCP currently only supports gomote-created buildlets")
+	}
+	req, err := http.NewRequest("POST", c.URL()+"/tcpproxy", nil)
+	if err != nil {
+		return nil, err
+	}
+	req.Header.Add("X-Target-Port", fmt.Sprint(port))
+	res, err := c.do(req)
+	if err != nil {
+		return nil, err
+	}
+	if res.StatusCode != http.StatusSwitchingProtocols {
+		slurp, _ := ioutil.ReadAll(io.LimitReader(res.Body, 4<<10))
+		res.Body.Close()
+		return nil, fmt.Errorf("wanted 101 Switching Protocols; unexpected response: %v, %q", res.Status, slurp)
+	}
+	rwc, ok := res.Body.(io.ReadWriteCloser)
+	if !ok {
+		res.Body.Close()
+		return nil, fmt.Errorf("tcpproxy response was not a Writer")
+	}
+	return rwc, nil
+}
+
 // ProxyRoundTripper returns a RoundTripper that sends HTTP requests directly
 // through to the underlying buildlet, adding auth and X-Buildlet-Proxy headers
 // as necessary. This is really only intended for use by the coordinator.

cmd/gomote/gomote.go

@@ -34,6 +34,7 @@ To list the subcommands, run "gomote" without arguments:
     put14      put Go 1.4 in place
     puttar     extract a tar.gz to a buildlet
     rm         delete files or directories
+    rdp        RDP (Remote Desktop Protocol) to a Windows buildlet
     run        run a command on a buildlet
     ssh        ssh to a buildlet
 
@@ -149,6 +150,7 @@ func registerCommands() {
 	registerCommand("put", "put files on a buildlet", put)
 	registerCommand("put14", "put Go 1.4 in place", put14)
 	registerCommand("puttar", "extract a tar.gz to a buildlet", putTar)
+	registerCommand("rdp", "RDP (Remote Desktop Protocol) to a Windows buildlet", rdp)
 	registerCommand("rm", "delete files or directories", rm)
 	registerCommand("run", "run a command on a buildlet", run)
 	registerCommand("ssh", "ssh to a buildlet", ssh)

cmd/gomote/rdp.go

@@ -0,0 +1,88 @@
+// Copyright 2019 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package main
+
+import (
+	"context"
+	"errors"
+	"flag"
+	"fmt"
+	"io"
+	"log"
+	"net"
+	"os"
+
+	"golang.org/x/build/buildlet"
+	"golang.org/x/sync/errgroup"
+)
+
+const rdpPort = 3389
+
+func rdp(args []string) error {
+	fs := flag.NewFlagSet("rdp", flag.ContinueOnError)
+	fs.Usage = func() {
+		fmt.Fprintln(os.Stderr, "rdp usage: gomote rdp [--listen=...] <instance>")
+		fs.PrintDefaults()
+		os.Exit(1)
+	}
+	var listen string
+	fs.StringVar(&listen, "listen", "localhost:"+fmt.Sprint(rdpPort), "local address to listen on")
+	fs.Parse(args)
+	if fs.NArg() != 1 {
+		fs.Usage()
+	}
+	name := fs.Arg(0)
+	bc, err := remoteClient(name)
+	if err != nil {
+		return err
+	}
+
+	ln, err := net.Listen("tcp", listen)
+	if err != nil {
+		return err
+	}
+	log.Printf("Listening on %v to proxy RDP.", ln.Addr())
+	for {
+		c, err := ln.Accept()
+		if err != nil {
+			return err
+		}
+		go handleRDPConn(bc, c)
+	}
+}
+
+func handleRDPConn(bc *buildlet.Client, c net.Conn) {
+	const Lmsgprefix = 64 // new in Go 1.14, harmless before
+	log := log.New(os.Stderr, c.RemoteAddr().String()+": ", log.LstdFlags|Lmsgprefix)
+	log.Printf("accepted connection, dialing buildlet via coordinator proxy...")
+	rwc, err := bc.ProxyTCP(rdpPort)
+	if err != nil {
+		c.Close()
+		log.Printf("failed to connect to buildlet via coordinator: %v", err)
+		return
+	}
+
+	log.Printf("connected to buildlet; proxying data.")
+
+	grp, ctx := errgroup.WithContext(context.Background())
+	grp.Go(func() error {
+		_, err := io.Copy(rwc, c)
+		if err == nil {
+			return errors.New("local client closed")
+		}
+		return fmt.Errorf("error copying from local to remote: %v", err)
+	})
+	grp.Go(func() error {
+		_, err := io.Copy(c, rwc)
+		if err == nil {
+			return errors.New("remote server closed")
+		}
+		return fmt.Errorf("error copying from remote to local: %v", err)
+	})
+	<-ctx.Done()
+	rwc.Close()
+	c.Close()
+	log.Printf("closing RDP connection: %v", grp.Wait())
+}