Initial agent

coder · kylecarbs · Feb 17, 2022 · Feb 15, 2022 · Feb 16, 2022 · Feb 16, 2022
commit 54133fc08e0a6d3079bc2a168b33e768372e3a87
diff --git a/.vscode/settings.json b/.vscode/settings.json
@@ -33,6 +33,7 @@
     "drpcserver",
     "fatih",
     "goleak",
+    "gossh",
     "hashicorp",
     "httpmw",
     "isatty",
@@ -54,6 +55,7 @@
     "retrier",
     "sdkproto",
     "stretchr",
+    "tcpip",
     "tfexec",
     "tfstate",
     "unconvert",

diff --git a/agent/agent.go b/agent/agent.go
@@ -0,0 +1,140 @@
+package agent
+
+import (
+	"context"
+	"errors"
+	"io"
+	"sync"
+	"time"
+
+	"cdr.dev/slog"
+	"github.com/coder/coder/peer"
+	"github.com/coder/coder/peerbroker"
+	"github.com/coder/retry"
+
+	"github.com/gliderlabs/ssh"
+	gossh "golang.org/x/crypto/ssh"
+)
+
+type Options struct {
+	Logger slog.Logger
+}
+
+type Dialer func(ctx context.Context) (*peerbroker.Listener, error)
+
+func Server(dialer Dialer, options *Options) io.Closer {
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	s := &server{
+		clientDialer: dialer,
+		options:      options,
+		closeCancel:  cancelFunc,
+	}
+	s.init(ctx)
+	return s
+}
+
+type server struct {
+	clientDialer Dialer
+	options      *Options
+
+	closeCancel context.CancelFunc
+	closeMutex  sync.Mutex
+	closed      chan struct{}
+	closeError  error
+
+	sshServer *ssh.Server
+}
+
+func (s *server) init(ctx context.Context) {
+	forwardHandler := &ssh.ForwardedTCPHandler{}
+	s.sshServer = &ssh.Server{
+		LocalPortForwardingCallback: func(ctx ssh.Context, destinationHost string, destinationPort uint32) bool {
+			return false
+		},
+		ReversePortForwardingCallback: func(ctx ssh.Context, bindHost string, bindPort uint32) bool {
+			return false
+		},
+		PtyCallback: func(ctx ssh.Context, pty ssh.Pty) bool {
+			return false
+		},
+		ServerConfigCallback: func(ctx ssh.Context) *gossh.ServerConfig {
+			return &gossh.ServerConfig{
+				Config: gossh.Config{
+					// "arcfour" is the fastest SSH cipher. We prioritize throughput
+					// over encryption here, because the WebRTC connection is already
+					// encrypted. If possible, we'd disable encryption entirely here.
+					Ciphers: []string{"arcfour"},
+				},
+				NoClientAuth: true,
+			}
+		},
+		RequestHandlers: map[string]ssh.RequestHandler{
+			"tcpip-forward":        forwardHandler.HandleSSHRequest,
+			"cancel-tcpip-forward": forwardHandler.HandleSSHRequest,
+		},
+	}
+
+	go s.run(ctx)
+}
+
+func (s *server) run(ctx context.Context) {
+	var peerListener *peerbroker.Listener
+	var err error
+	// An exponential back-off occurs when the connection is failing to dial.
+	// This is to prevent server spam in case of a coderd outage.
+	for retrier := retry.New(50*time.Millisecond, 10*time.Second); retrier.Wait(ctx); {
+		peerListener, err = s.clientDialer(ctx)
+		if err != nil {
+			if errors.Is(err, context.Canceled) {
+				return
+			}
+			if s.isClosed() {
+				return
+			}
+			s.options.Logger.Warn(context.Background(), "failed to dial", slog.Error(err))
+			continue
+		}
+		s.options.Logger.Debug(context.Background(), "connected")
+		break
+	}
+
+	for {
+		conn, err := peerListener.Accept()
+		if err != nil {
+			// This is closed!
+			return
+		}
+		go s.handle(ctx, conn)
+	}
+}
+
+func (s *server) handle(ctx context.Context, conn *peer.Conn) {
+	for {
+		channel, err := conn.Accept(ctx)
+		if err != nil {
+			// TODO: Log here!
+			return
+		}
+
+		switch channel.Protocol() {
+		case "ssh":
+			s.sshServer.HandleConn(channel.NetConn())
+		case "proxy":
+			// Proxy the port provided.
+		}
+	}
+}
+
+// isClosed returns whether the API is closed or not.
+func (s *server) isClosed() bool {
+	select {
+	case <-s.closed:
+		return true
+	default:
+		return false
+	}
+}
+
+func (s *server) Close() error {
+	return nil
+}
diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -0,0 +1,61 @@
+package agent_test
+
+import (
+	"context"
+	"net"
+	"os"
+	"testing"
+
+	"cdr.dev/slog/sloggers/slogtest"
+	"github.com/coder/coder/agent"
+	"github.com/coder/coder/peer"
+	"github.com/coder/coder/peerbroker"
+	"github.com/coder/coder/peerbroker/proto"
+	"github.com/coder/coder/provisionersdk"
+	"github.com/pion/webrtc/v3"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/crypto/ssh"
+)
+
+func TestAgent(t *testing.T) {
+	t.Run("asd", func(t *testing.T) {
+		ctx := context.Background()
+		client, server := provisionersdk.TransportPipe()
+		defer client.Close()
+		defer server.Close()
+		closer := agent.Server(func(ctx context.Context) (*peerbroker.Listener, error) {
+			return peerbroker.Listen(server, &peer.ConnOptions{
+				Logger: slogtest.Make(t, nil),
+			})
+		}, &agent.Options{
+			Logger: slogtest.Make(t, nil),
+		})
+		defer closer.Close()
+		api := proto.NewDRPCPeerBrokerClient(provisionersdk.Conn(client))
+		stream, err := api.NegotiateConnection(ctx)
+		require.NoError(t, err)
+		conn, err := peerbroker.Dial(stream, []webrtc.ICEServer{}, &peer.ConnOptions{
+			Logger: slogtest.Make(t, nil),
+		})
+		require.NoError(t, err)
+		defer conn.Close()
+		channel, err := conn.Dial(ctx, "example", &peer.ChannelOptions{
+			Protocol: "ssh",
+		})
+		require.NoError(t, err)
+		sshConn, channels, requests, err := ssh.NewClientConn(channel.NetConn(), "localhost:22", &ssh.ClientConfig{
+			User: "kyle",
+			HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
+				return nil
+			},
+		})
+		require.NoError(t, err)
+		sshClient := ssh.NewClient(sshConn, channels, requests)
+		session, err := sshClient.NewSession()
+		require.NoError(t, err)
+		session.Stdout = os.Stdout
+		session.Stderr = os.Stderr
+		err = session.Run("echo test")
+		require.NoError(t, err)
+	})
+}
diff --git a/go.mod b/go.mod
@@ -16,10 +16,11 @@ replace github.com/chzyer/readline => github.com/kylecarbs/readline v0.0.0-20220
 
 require (
 	cdr.dev/slog v1.4.1
-	github.com/Netflix/go-expect v0.0.0-20220104043353-73e0943537d2
 	github.com/briandowns/spinner v1.18.1
 	github.com/coder/retry v1.3.0
+	github.com/creack/pty v1.1.17
 	github.com/fatih/color v1.13.0
+	github.com/gliderlabs/ssh v0.3.3
 	github.com/go-chi/chi/v5 v5.0.7
 	github.com/go-chi/render v1.0.1
 	github.com/go-playground/validator/v10 v10.10.0
@@ -50,6 +51,7 @@ require (
 	go.uber.org/goleak v1.1.12
 	golang.org/x/crypto v0.0.0-20220131195533-30dcbda58838
 	golang.org/x/oauth2 v0.0.0-20211104180415-d3ed0bb246c8
+	golang.org/x/sys v0.0.0-20220114195835-da31bd327af9
 	golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1
 	google.golang.org/protobuf v1.27.1
 	nhooyr.io/websocket v1.8.7
@@ -63,11 +65,11 @@ require (
 	github.com/Nvveen/Gotty v0.0.0-20120604004816-cd527374f1e5 // indirect
 	github.com/agext/levenshtein v1.2.3 // indirect
 	github.com/alecthomas/chroma v0.10.0 // indirect
+	github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be // indirect
 	github.com/apparentlymart/go-textseg/v13 v13.0.0 // indirect
 	github.com/cenkalti/backoff/v4 v4.1.2 // indirect
 	github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e // indirect
 	github.com/containerd/continuity v0.2.2 // indirect
-	github.com/creack/pty v1.1.17 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/dhui/dktest v0.3.9 // indirect
 	github.com/dlclark/regexp2 v1.4.0 // indirect
@@ -124,7 +126,6 @@ require (
 	github.com/zeebo/errs v1.2.2 // indirect
 	go.opencensus.io v0.23.0 // indirect
 	golang.org/x/net v0.0.0-20220127200216-cd36cc0744dd // indirect
-	golang.org/x/sys v0.0.0-20220114195835-da31bd327af9 // indirect
 	golang.org/x/term v0.0.0-20210927222741-03fcf44c2211 // indirect
 	golang.org/x/text v0.3.7 // indirect
 	google.golang.org/appengine v1.6.7 // indirect

diff --git a/go.sum b/go.sum
@@ -103,8 +103,6 @@ github.com/Microsoft/hcsshim v0.8.23/go.mod h1:4zegtUJth7lAvFyc6cH2gGQ5B3OFQim01
 github.com/Microsoft/hcsshim/test v0.0.0-20201218223536-d3e5debf77da/go.mod h1:5hlzMzRKMLyo42nCZ9oml8AdTlq/0cvIaBv6tK1RehU=
 github.com/Microsoft/hcsshim/test v0.0.0-20210227013316-43a75bb4edd3/go.mod h1:mw7qgWloBUl75W/gVH3cQszUg1+gUITj7D6NY7ywVnY=
 github.com/NYTimes/gziphandler v0.0.0-20170623195520-56545f4a5d46/go.mod h1:3wb06e3pkSAbeQ52E9H9iFoQsEEwGN64994WTCIhntQ=
-github.com/Netflix/go-expect v0.0.0-20220104043353-73e0943537d2 h1:+vx7roKuyA63nhn5WAunQHLTznkw5W8b1Xc0dNjp83s=
-github.com/Netflix/go-expect v0.0.0-20220104043353-73e0943537d2/go.mod h1:HBCaDeC1lPdgDeDbhX8XFpy1jqjK0IBG8W5K+xYqA0w=
 github.com/Nvveen/Gotty v0.0.0-20120604004816-cd527374f1e5 h1:TngWCqHvy9oXAN6lEVMRuU21PR1EtLVZJmdB18Gu3Rw=
 github.com/Nvveen/Gotty v0.0.0-20120604004816-cd527374f1e5/go.mod h1:lmUJ/7eu/Q8D7ML55dXQrVaamCz2vxCfdQBasLZfHKk=
 github.com/OneOfOne/xxhash v1.2.2/go.mod h1:HSdplMjZKSmBqAxg5vPj2TmRDmfkzw+cTzAElWljhcU=
@@ -134,6 +132,8 @@ github.com/alecthomas/units v0.0.0-20190717042225-c3de453c63f4/go.mod h1:ybxpYRF
 github.com/alexflint/go-filemutex v0.0.0-20171022225611-72bdc8eae2ae/go.mod h1:CgnQgUtFrFz9mxFNtED3jI5tLDjKlOM+oUF/sTk6ps0=
 github.com/andybalholm/crlf v0.0.0-20171020200849-670099aa064f/go.mod h1:k8feO4+kXDxro6ErPXBRTJ/ro2mf0SsFG8s7doP9kJE=
 github.com/anmitsu/go-shlex v0.0.0-20161002113705-648efa622239/go.mod h1:2FmKhYUyUczH0OGQWaF5ceTx0UBShxjsH6f8oGKYe2c=
+github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be h1:9AeTilPcZAjCFIImctFaOjnTIavg87rW78vTPkQqLI8=
+github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be/go.mod h1:ySMOLuWl6zY27l47sB3qLNK6tF2fkHG55UZxx8oIVo4=
 github.com/antihax/optional v1.0.0/go.mod h1:uupD/76wgC+ih3iEmQUL+0Ugr19nfwCT1kdvxnR2qWY=
 github.com/apache/arrow/go/arrow v0.0.0-20210818145353-234c94e4ce64/go.mod h1:2qMFB56yOP3KzkB3PbYZ4AlUFg3a88F67TIx5lB/WwY=
 github.com/apache/arrow/go/arrow v0.0.0-20211013220434-5962184e7a30/go.mod h1:Q7yQnSMnLvcXlZ8RV+jwz/6y1rQTqbX6C82SndT52Zs=
@@ -443,6 +443,8 @@ github.com/gin-contrib/sse v0.1.0/go.mod h1:RHrZQHXnP2xjPF+u1gW/2HnVO7nvIa9PG3Gm
 github.com/gin-gonic/gin v1.6.3 h1:ahKqKTFpO5KTPHxWZjEdPScmYaGtLo8Y4DMHoEsnp14=
 github.com/gin-gonic/gin v1.6.3/go.mod h1:75u5sXoLsGZoRN5Sgbi1eraJ4GU3++wFwWzhwvtwp4M=
 github.com/gliderlabs/ssh v0.2.2/go.mod h1:U7qILu1NlMHj9FlMhZLlkCdDnU1DBEAqr0aevW3Awn0=
+github.com/gliderlabs/ssh v0.3.3 h1:mBQ8NiOgDkINJrZtoizkC3nDNYgSaWtxyem6S2XHBtA=
+github.com/gliderlabs/ssh v0.3.3/go.mod h1:ZSS+CUoKHDrqVakTfTWUlKSr9MtMFkC4UvtQKD7O914=
 github.com/go-chi/chi/v5 v5.0.7 h1:rDTPXLDHGATaeHvVlLcR4Qe0zftYethFucbjVQ1PxU8=
 github.com/go-chi/chi/v5 v5.0.7/go.mod h1:DslCQbL2OYiznFReuXYUmQ2hGd1aDpCnlMNITLSKoi8=
 github.com/go-chi/render v1.0.1 h1:4/5tis2cKaNdnv9zFLfXzcquC9HbeZgCnxGnKrltBS8=

diff --git a/peer/channel.go b/peer/channel.go
@@ -2,8 +2,10 @@ package peer
 
 import (
 	"context"
+	"fmt"
 	"io"
 	"net"
+	"runtime/debug"
 	"sync"
 	"time"
 
@@ -186,6 +188,7 @@ func (c *Channel) Read(bytes []byte) (int, error) {
 		if c.isClosed() {
 			return 0, c.closeError
 		}
+		debug.PrintStack()
 		// An EOF always occurs when the connection is closed.
 		// Alternative close errors will occur first if an unexpected
 		// close has occurred.
@@ -233,6 +236,8 @@ func (c *Channel) Write(bytes []byte) (n int, err error) {
 	// See: https://github.com/pion/sctp/issues/181
 	time.Sleep(time.Microsecond)
 
+	fmt.Printf("Writing %d\n", len(bytes))
+
 	return c.rwc.Write(bytes)
 }
 
@@ -246,6 +251,11 @@ func (c *Channel) Label() string {
 	return c.dc.Label()
 }
 
+// Protocol returns the protocol of the underlying DataChannel.
+func (c *Channel) Protocol() string {
+	return c.dc.Protocol()
+}
+
 // NetConn wraps the DataChannel in a struct fulfilling net.Conn.
 // Read, Write, and Close operations can still be used on the *Channel struct.
 func (c *Channel) NetConn() net.Conn {

diff --git a/peer/conn_test.go b/peer/conn_test.go
@@ -267,6 +267,23 @@ func TestConn(t *testing.T) {
 		_, err := client.Ping()
 		require.NoError(t, err)
 	})
+
+	t.Run("ShortBuffer", func(t *testing.T) {
+		t.Parallel()
+		client, server, _ := createPair(t)
+		exchange(client, server)
+		go func() {
+			channel, err := client.Dial(context.Background(), "test", nil)
+			require.NoError(t, err)
+			_, err = channel.Write([]byte{'1', '2'})
+			require.NoError(t, err)
+		}()
+
+		channel, err := server.Accept(context.Background())
+		require.NoError(t, err)
+		_, err = channel.Read(make([]byte, 1))
+		require.NoError(t, err)
+	})
 }
 
 func createPair(t *testing.T) (client *peer.Conn, server *peer.Conn, wan *vnet.Router) {