From 5cd7bae50586a339403b7cf2bc0e18c56b6b703b Mon Sep 17 00:00:00 2001
From: Elio Bischof <eliobischof@gmail.com>
Date: Thu, 8 Apr 2021 09:54:39 +0200
Subject: [PATCH] fix: cli client (#92)

* fix: cli client

* fix: print shutdown error correctly
---
 pkg/client/rp/cli/cli.go | 15 ++++++++-------
 pkg/utils/http.go        |  8 ++++++--
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/pkg/client/rp/cli/cli.go b/pkg/client/rp/cli/cli.go
index 6cbb364..a00f0bd 100644
--- a/pkg/client/rp/cli/cli.go
+++ b/pkg/client/rp/cli/cli.go
@@ -13,13 +13,14 @@ const (
 	loginPath = "/login"
 )
 
-func CodeFlow(relyingParty rp.RelyingParty, callbackPath, port string, stateProvider func() string) *oidc.Tokens {
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
+func CodeFlow(ctx context.Context, relyingParty rp.RelyingParty, callbackPath, port string, stateProvider func() string) *oidc.Tokens {
+	codeflowCtx, codeflowCancel := context.WithCancel(ctx)
+	defer codeflowCancel()
+
+	tokenChan := make(chan *oidc.Tokens, 1)
 
-	var token *oidc.Tokens
 	callback := func(w http.ResponseWriter, r *http.Request, tokens *oidc.Tokens, state string) {
-		token = tokens
+		tokenChan <- tokens
 		msg := "<p><strong>Success!</strong></p>"
 		msg = msg + "<p>You are authenticated and can now return to the CLI.</p>"
 		w.Write([]byte(msg))
@@ -27,9 +28,9 @@ func CodeFlow(relyingParty rp.RelyingParty, callbackPath, port string, stateProv
 	http.Handle(loginPath, rp.AuthURLHandler(stateProvider, relyingParty))
 	http.Handle(callbackPath, rp.CodeExchangeHandler(callback, relyingParty))
 
-	utils.StartServer(ctx, port)
+	utils.StartServer(codeflowCtx, ":"+port)
 
 	utils.OpenBrowser("http://localhost:" + port + loginPath)
 
-	return token
+	return <-tokenChan
 }
diff --git a/pkg/utils/http.go b/pkg/utils/http.go
index 6632053..27f96f9 100644
--- a/pkg/utils/http.go
+++ b/pkg/utils/http.go
@@ -97,7 +97,11 @@ func StartServer(ctx context.Context, port string) {
 
 	go func() {
 		<-ctx.Done()
-		err := server.Shutdown(ctx)
-		log.Fatalf("Shutdown(): %v", err)
+		ctxShutdown, cancelShutdown := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancelShutdown()
+		err := server.Shutdown(ctxShutdown)
+		if err != nil {
+			log.Fatalf("Shutdown(): %v", err)
+		}
 	}()
 }