From 81f2496b5eb263f58ca01ea17dbf0cac14ef5f0c Mon Sep 17 00:00:00 2001
From: Guillaume Tardif <guillaume.tardif@docker.com>
Date: Thu, 3 Sep 2020 15:35:33 +0200
Subject: [PATCH] Allow Ctrl+C to cancel CLI when using Azure Device Code Flow
 login

Signed-off-by: Guillaume Tardif <guillaume.tardif@docker.com>
---
 aci/login/login.go | 30 ++++++++++++++++++++++++------
 1 file changed, 24 insertions(+), 6 deletions(-)

diff --git a/aci/login/login.go b/aci/login/login.go
index 775ad065..433b4114 100644
--- a/aci/login/login.go
+++ b/aci/login/login.go
@@ -176,18 +176,20 @@ func (login *AzureLoginService) Login(ctx context.Context, requestedTenantID str
 		return errors.Wrap(errdefs.ErrLoginFailed, "empty redirect URL")
 	}
 
+	deviceCodeFlowCh := make(chan deviceCodeFlowResponse, 1)
 	if err = login.apiHelper.openAzureLoginPage(redirectURL); err != nil {
-		fmt.Println("Could not automatically open a browser, falling back to Azure device code flow authentication")
-		token, err := login.apiHelper.getDeviceCodeFlowToken()
-		if err != nil {
-			return errors.Wrapf(errdefs.ErrLoginFailed, "could not get token using device code flow: %s", err)
-		}
-		return login.getTenantAndValidateLogin(token.AccessToken, token.RefreshToken, requestedTenantID)
+		login.startDeviceCodeFlow(deviceCodeFlowCh)
 	}
 
 	select {
 	case <-ctx.Done():
 		return ctx.Err()
+	case dcft := <-deviceCodeFlowCh:
+		if dcft.err != nil {
+			return errors.Wrapf(errdefs.ErrLoginFailed, "could not get token using device code flow: %s", err)
+		}
+		token := dcft.token
+		return login.getTenantAndValidateLogin(token.AccessToken, token.RefreshToken, requestedTenantID)
 	case q := <-queryCh:
 		if q.err != nil {
 			return errors.Wrapf(errdefs.ErrLoginFailed, "unhandled local login server error: %s", err)
@@ -211,6 +213,22 @@ func (login *AzureLoginService) Login(ctx context.Context, requestedTenantID str
 	}
 }
 
+type deviceCodeFlowResponse struct {
+	token adal.Token
+	err   error
+}
+
+func (login *AzureLoginService) startDeviceCodeFlow(deviceCodeFlowCh chan deviceCodeFlowResponse) {
+	fmt.Println("Could not automatically open a browser, falling back to Azure device code flow authentication")
+	go func() {
+		token, err := login.apiHelper.getDeviceCodeFlowToken()
+		if err != nil {
+			deviceCodeFlowCh <- deviceCodeFlowResponse{err: err}
+		}
+		deviceCodeFlowCh <- deviceCodeFlowResponse{token: token}
+	}()
+}
+
 func getTenantID(tenantValues []tenantValue, requestedTenantID string) (string, error) {
 	if requestedTenantID == "" {
 		if len(tenantValues) < 1 {