fix(auth): use a fresh mux and server per interactive auth attempt

pkg/auth_service.go

@@ -9,11 +9,14 @@ import (
 	"fmt"
 	"io"
 	"math/rand"
+	"net"
 	"net/http"
 	"net/url"
 	"os"
 	"strconv"
 	"strings"
+	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/pkg/browser"
@@ -22,18 +25,32 @@ import (
 )
 
 const (
-	httpAuthPort   = 10999
-	oAuthQoveryUrl = "https://auth.qovery.com/login?code_challenge_method=S256&scope=%s&client=%s&protocol=oauth2&response_type=%s&audience=%s&redirect_uri=%s&code_challenge=%s"
+	httpAuthPort     = 10999
+	qoveryConsoleUrl = "https://console.qovery.com"
+	oAuthQoveryUrl   = "https://auth.qovery.com/login?code_challenge_method=S256&scope=%s&client=%s&protocol=oauth2&response_type=%s&audience=%s&redirect_uri=%s&code_challenge=%s"
 )
 
 var (
+	// Delay before the callback server shuts down, so the browser gets its response.
+	authShutdownDelay = time.Second
+
 	oAuthUrlParamValueClient         = "MJ2SJpu12PxIzgmc5z5Y7N8m5MnaF7Y0"
 	oAuthUrlParamValueHeadlessClient = "f9drkTNpxsEw2VU2PVDrxhyT3vVuFT0Y"
 	oAuthUrlParamValueAudience       = "https://core.qovery.com"
 	oAuthUrlParamValueResponseType   = "code"
 	oAuthUrlParamValueScopes         = "offline_access openid profile email"
 	oAuthUrlParamValueRedirect       = "http://localhost:" + strconv.Itoa(httpAuthPort) + "/authorization"
 	oAuthTokenEndpoint               = "https://auth.qovery.com/oauth/token"
+
+	authSuccessMessage = "Authentication successful, you'll be redirected to Qovery console. If it's not the case, click on this link: "
+	authFailureMessage = "Authentication failed. Run 'qovery auth' again, or contact #support on https://discord.qovery.com."
+
+	// Upper bound for the graceful shutdown. It frees the port; it does not
+	// decide when the failure response is complete, see awaitFailureDelivered.
+	authShutdownTimeout = 5 * time.Second
+
+	// exitProcess is replaceable so tests can observe an authentication failure.
+	exitProcess = os.Exit
 )
 
 type TokensResponse struct {
@@ -51,8 +68,6 @@ type DeviceFlowParameters struct {
 }
 
 func DoRequestUserToAuthenticate(headless bool, skipVersionCheck bool) {
-	qoveryConsoleUrl := "https://console.qovery.com"
-
 	if !skipVersionCheck {
 		available, message, _ := CheckAvailableNewVersion()
 		if available {
@@ -70,66 +85,255 @@ func DoRequestUserToAuthenticate(headless bool, skipVersionCheck bool) {
 		utils.PrintlnError(errors.New("can not create authorization code challenge. Please contact the #support at 'https://discord.qovery.com'. "))
 		os.Exit(0)
 	}
+	// Listen before opening the browser so a busy port is reported instead of
+	// leaving the user waiting on a callback that can never arrive.
+	srv := newAuthorizationServer(verifier, storeTokens)
+	listener, err := net.Listen("tcp", srv.Addr)
+	if err != nil {
+		utils.PrintlnError(fmt.Errorf("can not listen on %s for the authentication callback, is another 'qovery auth' running? %w", srv.Addr, err))
+		os.Exit(1)
+	}
+
 	// TODO link to web auth
 	_ = browser.OpenURL(fmt.Sprintf(oAuthQoveryUrl, url.QueryEscape(oAuthUrlParamValueScopes), oAuthUrlParamValueClient, url.QueryEscape(oAuthUrlParamValueResponseType),
 		url.QueryEscape(oAuthUrlParamValueAudience), url.QueryEscape(oAuthUrlParamValueRedirect), challenge))
 
 	fmt.Println("\nOpening your browser, waiting for your authentication... ")
 
-	srv := &http.Server{Addr: fmt.Sprintf("localhost:%d", httpAuthPort)}
+	serveAuthorization(srv, listener)
+}
+
+// authorizationServer is the callback server of one authentication attempt.
+type authorizationServer struct {
+	*http.Server
+	failed atomic.Bool
+
+	// shutdownTimeout is read once at construction, so the package variable is
+	// not touched from server goroutines.
+	shutdownTimeout time.Duration
+
+	// failedConn is the connection that carries the failure response.
+	// delivered closes once that connection is idle again after the failing
+	// handler: the http server only gets there after the response is completely
+	// written. Closed and hijacked connections prove nothing about delivery.
+	failMu        sync.Mutex
+	failedConn    net.Conn
+	delivered     chan struct{}
+	deliveredOnce sync.Once
+}
+
+type connContextKey struct{}
+
+func (s *authorizationServer) markDelivered() {
+	s.deliveredOnce.Do(func() { close(s.delivered) })
+}
+
+// serveAuthorization serves until the attempt ends. Serve returns as soon as
+// Shutdown starts, and Shutdown closes the listener first, so the fixed port is
+// free again when this function returns. On failure the process exits here, not
+// in the handler, and only once the failure response is confirmed completely
+// written: the browser must not get it truncated. If that never gets confirmed
+// (connection lost mid-write), the CLI keeps waiting instead of exiting.
+func serveAuthorization(srv *authorizationServer, listener net.Listener) {
+	_ = srv.Serve(listener)
+	srv.shutdown()
+	if srv.failed.Load() {
+		srv.awaitFailureDelivered()
+		exitProcess(1)
+	}
+}
+
+// shutdown stops the server and waits, within authShutdownTimeout, for in-flight
+// requests to finish. It is safe to call more than once. A timeout here does not
+// mean the failure response is complete: see awaitFailureDelivered.
+func (s *authorizationServer) shutdown() {
+	ctx, cancel := context.WithTimeout(context.Background(), s.shutdownTimeout)
+	defer cancel()
+	if err := s.Shutdown(ctx); err != nil {
+		utils.PrintlnError(err)
+	}
+}
+
+// awaitFailureDelivered blocks until the failure response is confirmed
+// completely written. There is no timeout on purpose: exiting earlier could
+// truncate the response the user is about to read.
+func (s *authorizationServer) awaitFailureDelivered() {
+	<-s.delivered
+}
+
+// fail answers the browser with the failure message and ends the attempt. The
+// handler only returns: exiting here would kill the process before the response
+// is complete (a flushed response is still chunked and unterminated).
+func (s *authorizationServer) fail(writer http.ResponseWriter, request *http.Request, status int) {
+	// Without a tracked connection nothing can ever confirm delivery, so the
+	// attempt stays failed but neither exits nor shuts down.
+	conn, _ := request.Context().Value(connContextKey{}).(net.Conn)
+	s.failMu.Lock()
+	s.failedConn = conn
+	s.failMu.Unlock()
+
+	http.Error(writer, authFailureMessage, status)
+	utils.PrintlnError(errors.New("authentication unsuccessful. Try again later or contact #support on 'https://discord.qovery.com'. "))
+	if s.failed.CompareAndSwap(false, true) && conn != nil {
+		// Shutdown disables keep-alive, so a response finished during it closes
+		// the connection without ever going idle and delivery could not be
+		// confirmed. Shut down only once it is.
+		go func() {
+			s.awaitFailureDelivered()
+			s.shutdown()
+		}()
+	}
+}
+
+func storeTokens(tokens TokensResponse) {
+	expiredAt := time.Now().Local().Add(time.Duration(tokens.ExpiresIn-60) * time.Second)
+	_ = utils.SetAccessToken(utils.AccessToken(tokens.AccessToken), expiredAt, utils.RefreshToken(tokens.RefreshToken))
+}
+
+// authorizationPage is the callback page. The element is named statusElement
+// because a global status resolves to window.status, which is string-valued:
+// assigning the DOM element to it would store a string and lose the element.
+func authorizationPage(port int) string {
+	return fmt.Sprintf(`<p id="status">Authenticating...</p>
+<script type="text/javascript" charset="utf-8">
+	var statusElement = document.getElementById("status");
+	var code = new URLSearchParams(window.location.search).get("code") || "";
+	var xmlHttp = new XMLHttpRequest();
+	xmlHttp.open("GET", "http://localhost:%d/authorization/valid?code=" + encodeURIComponent(code), true);
+	xmlHttp.onload = function () {
+		statusElement.textContent = xmlHttp.responseText;
+		if (xmlHttp.status === 200) {
+			var link = document.createElement("a");
+			link.href = %q;
+			link.textContent = %q;
+			statusElement.appendChild(link);
+			window.setTimeout(function () { window.location = %q; }, 2000);
+		}
+	};
+	xmlHttp.onerror = function () {
+		statusElement.textContent = "Authentication failed, the Qovery CLI could not be reached. Run 'qovery auth' again.";
+	};
+	xmlHttp.send(null);
+</script>`, port, qoveryConsoleUrl, qoveryConsoleUrl, qoveryConsoleUrl)
+}
 
-	http.HandleFunc("/authorization", func(writer http.ResponseWriter, request *http.Request) {
-		js := fmt.Sprintf(`<script type="text/javascript" charset="utf-8">
-				var hash = window.location.search.split("=")[1].split("&")[0];
-				var xmlHttp = new XMLHttpRequest();
-				xmlHttp.open("GET", "http://localhost:%d/authorization/valid?code=" + hash, false);
-				xmlHttp.send(null);
-				xmlHttp.responseText;
-				window.setTimeout('window.location="`+qoveryConsoleUrl+`"; ',2000);
-             </script>`, httpAuthPort)
+// newAuthorizationServer builds a server with its own mux for one interactive
+// authentication attempt. The handlers must not be registered on
+// http.DefaultServeMux: a second attempt in the same process (re-auth after a
+// 401) would panic on the duplicate "/authorization" pattern. The PKCE verifier
+// is captured per attempt.
+func newAuthorizationServer(verifier string, onTokens func(TokensResponse)) *authorizationServer {
+	srv := &authorizationServer{
+		Server:    &http.Server{Addr: fmt.Sprintf("localhost:%d", httpAuthPort)},
+		delivered: make(chan struct{}),
+
+		shutdownTimeout: authShutdownTimeout,
+	}
+	srv.ConnContext = func(ctx context.Context, conn net.Conn) context.Context {
+		return context.WithValue(ctx, connContextKey{}, conn)
+	}
+	srv.ConnState = func(conn net.Conn, state http.ConnState) {
+		switch state {
+		case http.StateIdle:
+			srv.failMu.Lock()
+			isFailed := conn == srv.failedConn
+			srv.failMu.Unlock()
+			if isFailed {
+				srv.markDelivered()
+			}
+		case http.StateNew, http.StateActive, http.StateHijacked, http.StateClosed:
+		}
+	}
+	mux := http.NewServeMux()
 
-		_, _ = writer.Write([]byte(js))
-		_, _ = writer.Write([]byte("Authentication successful, you'll be redirected to Qovery console. If it's not the case, click on this link: <a href='" + qoveryConsoleUrl + "'>" + qoveryConsoleUrl + "</a>"))
+	// The page only shows what /authorization/valid answers: the success message
+	// is not part of it, so a failed token exchange can never look like a success.
+	mux.HandleFunc("/authorization", func(writer http.ResponseWriter, request *http.Request) {
+		writer.Header().Set("Content-Type", "text/html; charset=utf-8")
+		_, _ = writer.Write([]byte(authorizationPage(httpAuthPort)))
 	})
 
-	http.HandleFunc("/authorization/valid", func(writer http.ResponseWriter, request *http.Request) {
-		code := request.URL.Query()["code"][0]
-		res, err := http.PostForm(oAuthTokenEndpoint, url.Values{
-			"grant_type":    {"authorization_code"},
-			"client_id":     {oAuthUrlParamValueClient},
-			"code":          {code},
-			"redirect_uri":  {oAuthUrlParamValueRedirect},
-			"code_verifier": {verifier},
-		})
+	// The callback is one-shot. A repeat (page reload, double request) must not
+	// replay the consumed code: it gets the first success again, whatever it
+	// carries. The lock also makes a concurrent repeat wait for the first exchange.
+	var (
+		mu   sync.Mutex
+		done bool
+	)
+	mux.HandleFunc("/authorization/valid", func(writer http.ResponseWriter, request *http.Request) {
+		mu.Lock()
+		defer mu.Unlock()
+		if done {
+			_, _ = writer.Write([]byte(authSuccessMessage))
+			return
+		}
+
+		// The attempt already failed and is shutting down: do not exchange again.
+		if srv.failed.Load() {
+			http.Error(writer, authFailureMessage, http.StatusUnauthorized)
+			return
+		}
 
+		codes := request.URL.Query()["code"]
+		if len(codes) == 0 || codes[0] == "" {
+			srv.fail(writer, request, http.StatusBadRequest)
+			return
+		}
+		tokens, err := exchangeAuthorizationCode(verifier, codes[0])
 		if err != nil {
-			utils.PrintlnError(errors.New("authentication unsuccessful. Try again later or contact #support on 'https://discord.qovery.com'. "))
-			os.Exit(0)
-		} else {
-			defer func(Body io.ReadCloser) {
-				_ = Body.Close()
-			}(res.Body)
-
-			tokens := TokensResponse{}
-			err := json.NewDecoder(res.Body).Decode(&tokens)
-			if err != nil {
-				utils.PrintlnError(errors.New("authentication unsuccessful. Try again later or contact #support on 'https://discord.qovery.com'. "))
-				os.Exit(0)
-			}
-			expiredAt := time.Now().Local().Add(time.Duration(tokens.ExpiresIn-60) * time.Second)
-			_ = utils.SetAccessToken(utils.AccessToken(tokens.AccessToken), expiredAt, utils.RefreshToken(tokens.RefreshToken))
-			utils.PrintlnInfo("Success!")
+			// A rejected or stale code must not be stored nor reported as a success.
+			srv.fail(writer, request, http.StatusUnauthorized)
+			return
 		}
+		done = true
+		onTokens(tokens)
+		utils.PrintlnInfo("Success!")
+		_, _ = writer.Write([]byte(authSuccessMessage))
 
+		delay := authShutdownDelay
 		go func() {
-			time.Sleep(time.Second)
-			if err := srv.Shutdown(context.TODO()); err != nil {
-				utils.PrintlnError(err)
-			}
+			time.Sleep(delay)
+			srv.shutdown()
 		}()
 	})
 
-	_ = srv.ListenAndServe()
+	srv.Handler = mux
+	return srv
+}
+
+// exchangeAuthorizationCode trades the authorization code for tokens. It fails
+// on a non-200 answer or an empty access token, which is how the token endpoint
+// reports a rejected or already used code.
+func exchangeAuthorizationCode(verifier string, code string) (TokensResponse, error) {
+	res, err := http.PostForm(oAuthTokenEndpoint, url.Values{
+		"grant_type":    {"authorization_code"},
+		"client_id":     {oAuthUrlParamValueClient},
+		"code":          {code},
+		"redirect_uri":  {oAuthUrlParamValueRedirect},
+		"code_verifier": {verifier},
+	})
+	if err != nil {
+		return TokensResponse{}, err
+	}
+	defer func() { _ = res.Body.Close() }()
+
+	if res.StatusCode != http.StatusOK {
+		return TokensResponse{}, fmt.Errorf("token endpoint answered %d", res.StatusCode)
+	}
+	// Read the whole body: Unmarshal rejects trailing data after the first JSON
+	// value, which a streaming Decode would accept.
+	payload, err := io.ReadAll(res.Body)
+	if err != nil {
+		return TokensResponse{}, err
+	}
+	tokens := TokensResponse{}
+	if err := json.Unmarshal(payload, &tokens); err != nil {
+		return TokensResponse{}, err
+	}
+	if tokens.AccessToken == "" {
+		return TokensResponse{}, errors.New("token endpoint returned an empty access token")
+	}
+	return tokens, nil
 }
 
 func createCodeVerifier() string {
@@ -171,8 +375,7 @@ func runHeadlessFlow() {
 		tokens, err := getTokensWith(parameters)
 
 		if err == nil {
-			expiredAt := time.Now().Local().Add(time.Duration(tokens.ExpiresIn-60) * time.Second)
-			_ = utils.SetAccessToken(utils.AccessToken(tokens.AccessToken), expiredAt, utils.RefreshToken(tokens.RefreshToken))
+			storeTokens(tokens)
 			utils.PrintlnInfo("Success!")
 			return
 		}