Add --check option and fix nits

cli/download/download.go

@@ -55,8 +55,8 @@ type newlineStdoutCloser struct {
 }
 
 func (n *newlineStdoutCloser) Close() error {
-	if isatty.IsTerminal(n.Fd()) {
-		n.Write([]byte{'\n'})
+	if isatty.IsTerminal(n.File.Fd()) {
+		n.File.Write([]byte{'\n'})
 	}
 	return n.File.Close()
 }
@@ -77,7 +77,7 @@ func printProtoToOutput(msg proto.Message, output io.Writer) error {
 	return err
 }
 
-func DownloadFile(uri string) error {
+func downloadFile(uri string) error {
 	if !strings.HasPrefix(uri, "/blobs") {
 		uri = "/blobs/" + uri
 	}
@@ -146,7 +146,7 @@ func HandleDownload(args []string) (int, error) {
 	}
 
 	uri := flags.Args()[0]
-	if err := DownloadFile(uri); err != nil {
+	if err := downloadFile(uri); err != nil {
 		log.Print(err)
 		return 1, nil
 	}

cli/login/login.go

@@ -9,7 +9,6 @@ import (
 	"os/exec"
 	"runtime"
 	"strings"
-	"sync"
 
 	"github.com/buildbuddy-io/buildbuddy/cli/arg"
 	"github.com/buildbuddy-io/buildbuddy/cli/log"
@@ -24,61 +23,73 @@ import (
 )
 
 const (
-	ApiKeyRepoSetting = "api-key"
+	apiKeyRepoSetting = "api-key"
 	apiKeyHeader      = "remote_header=x-buildbuddy-api-key"
 )
 
 var (
-	loginFlagSet = flag.NewFlagSet("login", flag.ContinueOnError)
+	flags = flag.NewFlagSet("login", flag.ContinueOnError)
 
-	apiTarget = loginFlagSet.String("target", "grpcs://remote.buildbuddy.io", "BuildBuddy gRPC target")
-	loginURL  = loginFlagSet.String("login_url", "https://app.buildbuddy.io/settings/cli-login", "URL for user to login")
+	check         = flags.Bool("check", false, "Just check whether logged in. Exits with code 0 if logged in, code 1 if not logged in, or 2 if there is an error.")
+	allowExisting = flags.Bool("allow_existing", false, "Don't force re-login if the current credentials are valid.")
+
+	loginURL  = flags.String("url", "https://app.buildbuddy.io/settings/cli-login", "Web URL for user to login")
+	apiTarget = flags.String("target", "grpcs://remote.buildbuddy.io", "BuildBuddy gRPC target")
+
+	usage = `
+bb ` + flags.Name() + ` [--check] [--allow_existing]
+
+Logs into BuildBuddy, saving your personal API key to .git/config.
+`
 )
 
-var getApiClient = sync.OnceValues(func() (bbspb.BuildBuddyServiceClient, error) {
+func authenticate(apiKey string) error {
 	conn, err := grpc_client.DialSimple(*apiTarget)
 	if err != nil {
-		return nil, fmt.Errorf("unable to dial BuildBuddy target: %v", err)
+		return fmt.Errorf("dial %s: %w", *apiTarget, err)
 	}
-	return bbspb.NewBuildBuddyServiceClient(conn), nil
-})
+	defer conn.Close()
+	client := bbspb.NewBuildBuddyServiceClient(conn)
 
-func isValidApiKey(apiKey string) bool {
-	bbsClient, err := getApiClient()
-	if err != nil {
-		log.Debugf("unable to get buildbuddy service client: %v", err)
-		return false
-	}
 	ctx := context.Background()
 	ctx = metadata.AppendToOutgoingContext(ctx, "x-buildbuddy-api-key", apiKey)
 
-	_, err = bbsClient.GetUser(ctx, &uspb.GetUserRequest{})
-	if err != nil && status.IsNotFoundError(err) && strings.Contains(err.Error(), "user not found") {
-		// Org-level API key is used
-		return true
-	}
+	_, err = client.GetUser(ctx, &uspb.GetUserRequest{})
 	if err != nil {
-		log.Debugf("fail to get user: %v", err)
-		return false
+		return err
 	}
-	return true
-}
-
-func hasValidApiKey() bool {
-	apiKey, err := storage.ReadRepoConfig(ApiKeyRepoSetting)
-	if err != nil {
-		log.Debugf("unable to read git repo config: %v", err)
-		return false
-	}
-	log.Printf("Found an API key in your .git/config. Verifying it's validity.")
-
-	return isValidApiKey(apiKey)
+	return nil
 }
 
 func HandleLogin(args []string) (exitCode int, err error) {
-	if hasValidApiKey() {
-		log.Printf("Current API key is valid. Skipping login.")
-		return 0, nil
+	if err := arg.ParseFlagSet(flags, args); err != nil {
+		if err == flag.ErrHelp {
+			log.Print(usage)
+			return 1, nil
+		}
+		return -1, err
+	}
+
+	if *check || *allowExisting {
+		apiKey, err := storage.ReadRepoConfig(apiKeyRepoSetting)
+		if err != nil {
+			return -1, fmt.Errorf("read .git/config: %w", err)
+		}
+		code := 0
+		if err := authenticate(apiKey); err != nil {
+			log.Printf("Failed to authenticate: %s", err)
+			if !status.IsUnauthenticatedError(err) {
+				code = 2
+			} else {
+				code = 1
+			}
+		}
+		if *check {
+			return code, nil
+		}
+		if *allowExisting && code == 0 {
+			return 0, nil
+		}
 	}
 
 	log.Printf("Press Enter to open %s in your browser...", *loginURL)
@@ -101,11 +112,12 @@ func HandleLogin(args []string) (exitCode int, err error) {
 	if apiKey == "" {
 		return -1, fmt.Errorf("invalid input: API key is empty")
 	}
-	if !isValidApiKey(apiKey) {
-		return -1, fmt.Errorf("invalid input: API key is not valid")
+
+	if err := authenticate(apiKey); err != nil {
+		return -1, fmt.Errorf("authenticate API key: %w", err)
 	}
 
-	if err := storage.WriteRepoConfig(ApiKeyRepoSetting, apiKey); err != nil {
+	if err := storage.WriteRepoConfig(apiKeyRepoSetting, apiKey); err != nil {
 		return -1, fmt.Errorf("failed to write API key to local .git/config: %s", err)
 	}
 
@@ -116,7 +128,7 @@ func HandleLogin(args []string) (exitCode int, err error) {
 }
 
 func HandleLogout(args []string) (exitCode int, err error) {
-	if err := storage.WriteRepoConfig(ApiKeyRepoSetting, ""); err != nil {
+	if err := storage.WriteRepoConfig(apiKeyRepoSetting, ""); err != nil {
 		return -1, fmt.Errorf("failed to clear api key from local .git/config: %s", err)
 	}
 
@@ -143,13 +155,13 @@ func ConfigureAPIKey(args []string) ([]string, error) {
 		return args, nil
 	}
 
-	apiKey, err := storage.ReadRepoConfig(ApiKeyRepoSetting)
+	apiKey, err := storage.ReadRepoConfig(apiKeyRepoSetting)
 	if err != nil {
 		// If we're not in a git repo, we'll fail to read the repo-specific
 		// config.
 		// Making this fatal would be inconvenient for new workspaces,
 		// so just log a debug message and move on.
-		log.Debugf("failed to configure API key from .git/config: %s", err)
+		log.Debugf("Failed to read API key from .git/config: %s", err)
 		return args, nil
 	}
 	if apiKey == "" {

cli/storage/storage.go

@@ -51,7 +51,7 @@ func CacheDir() (string, error) {
 	return cacheDir, nil
 }
 
-var RepoRootPath = sync.OnceValues(func() (string, error) {
+var repoRootPath = sync.OnceValues(func() (string, error) {
 	dir, err := exec.Command("git", "rev-parse", "--show-toplevel").CombinedOutput()
 	if err != nil {
 		return "", fmt.Errorf("failed to run git rev-parse: %s", err)
@@ -61,7 +61,7 @@ var RepoRootPath = sync.OnceValues(func() (string, error) {
 
 // ReadRepoConfig reads a repository-local configuration setting.
 func ReadRepoConfig(key string) (string, error) {
-	dir, err := RepoRootPath()
+	dir, err := repoRootPath()
 	if err != nil {
 		return "", err
 	}
@@ -89,7 +89,7 @@ func ReadRepoConfig(key string) (string, error) {
 
 // WriteRepoConfig writes a repository-local configuration setting.
 func WriteRepoConfig(key, value string) error {
-	dir, err := RepoRootPath()
+	dir, err := repoRootPath()
 	if err != nil {
 		return err
 	}