Consolidate agent to single 'bind' command

Author: sreya

internal/cmd/bind.go

@@ -0,0 +1,148 @@
+package cmd
+
+import (
+	"context"
+	"os"
+	"regexp"
+	"strings"
+	"time"
+
+	"cdr.dev/slog/sloggers/sloghuman"
+	"github.com/spf13/pflag"
+	"golang.org/x/xerrors"
+
+	"go.coder.com/cli"
+	"go.coder.com/cloud-agent/internal/client"
+	"go.coder.com/cloud-agent/internal/config"
+	"go.coder.com/cloud-agent/internal/ideproxy"
+	"go.coder.com/flog"
+)
+
+var (
+	DefaultCloudURL       = "https://cloud.coder.com"
+	codeserverPasswordEnv = "CODESERVER_PASSWORD"
+	codeserverAddrEnv     = "CODESERVER_ADDR"
+)
+
+var codeServerNameRx = regexp.MustCompile("^[a-z][a-z0-9_]{0,50}$")
+
+type bindCmd struct {
+	cloudURL string
+}
+
+func (c *bindCmd) Spec() cli.CommandSpec {
+	return cli.CommandSpec{
+		Name:  "bind",
+		Usage: "[NAME]",
+		Desc:  "Bind a server to Coder Cloud. A name will be generated from the hostname if one is not provided.",
+	}
+}
+
+func (c *bindCmd) RegisterFlags(fl *pflag.FlagSet) {
+	fl.StringVar(&c.cloudURL, "cloud-url", DefaultCloudURL, "The Coder Cloud URL to connect to.")
+}
+
+func (c *bindCmd) Run(fl *pflag.FlagSet) {
+	var (
+		err error
+		ctx = context.Background()
+	)
+
+	name := fl.Arg(0)
+	if name == "" {
+		// Generate a name based on the hostname if one is not provided.
+		name, err = genServerName()
+		if err != nil {
+			flog.Fatal("Failed to generate server name: %v", err.Error())
+		}
+	}
+
+	if !codeServerNameRx.MatchString(name) {
+		flog.Fatal("Name must conform to regex %s", codeServerNameRx.String())
+	}
+
+	cli, err := client.FromEnv()
+	if xerrors.Is(err, os.ErrNotExist) {
+		cli, err = loginClient(c.cloudURL, name)
+	}
+	if err != nil {
+		flog.Fatal("Failed to login: %v", err)
+	}
+
+	// Register the server with Coder Cloud. This is an idempotent
+	// operation.
+	cs, err := cli.RegisterCodeServer(name)
+	if err != nil {
+		flog.Fatal("Failed to register server: %v", err)
+	}
+
+	// Get the Access URL for the user.
+	url, err := cli.AccessURL(cs.ID)
+	if err != nil {
+		flog.Fatal("Failed to query server: %v", err)
+	}
+
+	token, err := config.SessionToken.Read()
+	if err != nil {
+		flog.Fatal("Failed to read session token: %v", err)
+	}
+
+	agent := &ideproxy.Agent{
+		Log:                sloghuman.Make(os.Stderr),
+		CodeServerID:       cs.ID,
+		SessionToken:       token,
+		CloudProxyURL:      c.cloudURL,
+		CodeServerAddr:     os.Getenv(codeserverAddrEnv),
+		CodeServerPassword: os.Getenv(codeserverPasswordEnv),
+	}
+
+	proxy := func() {
+		err = agent.Proxy(ctx)
+		if err != nil {
+			flog.Error("Connection to Coder-Cloud disrupted, re-establishing connection: %v", err.Error())
+		}
+	}
+
+	flog.Info("Proxying code-server to Coder Cloud, you can access your IDE at %v", url)
+
+	proxy()
+
+	// Avoid a super tight loop.
+	ticker := time.NewTicker(time.Second)
+	for range ticker.C {
+		proxy()
+	}
+}
+
+func loginClient(url, serverName string) (*client.Client, error) {
+	token, err := client.Login(url, serverName)
+	if err != nil {
+		return nil, xerrors.Errorf("unable to login: %w", err)
+	}
+
+	err = config.SessionToken.Write(token)
+	if err != nil {
+		return nil, xerrors.Errorf("write session token to file: %w", err)
+	}
+
+	err = config.URL.Write(url)
+	if err != nil {
+		return nil, xerrors.Errorf("write coder-cloud url to file: %w", err)
+	}
+
+	return client.FromEnv()
+}
+
+func genServerName() (string, error) {
+	hostname, err := os.Hostname()
+	if err != nil {
+		xerrors.Errorf("get hostname: %w", err)
+	}
+
+	hostname = strings.ToLower(hostname)
+
+	// Only use the first token.
+	hostname = strings.SplitN(hostname, ".", 1)[0]
+	// '-' are not allowed, convert them to '_'.
+	return strings.Replace(hostname, "-", "_", -1), nil
+}

internal/cmd/cmd.go

@@ -27,8 +27,7 @@ func (c *rootCmd) Spec() cli.CommandSpec {
 
 func (c *rootCmd) Subcommands() []cli.Command {
 	return []cli.Command{
-		&linkCmd{},
-		&proxyCmd{},
+		&bindCmd{},
 		&versionCmd{},
 	}
 }

internal/cmd/link.go

@@ -5,8 +5,6 @@ import "golang.org/x/xerrors"
 var (
 	// SessionToken is the file containing the session token.
 	SessionToken File = "session"
-	// ServerID is the file containing the server ID.
-	ServerID File = "server"
 	// URL is the file containing the url to a Coder-Cloud instance.
 	URL File = "url"
 )