docs: refactor getTokenFromWeb for OAuth2 callback handling

gmail/quickstart/quickstart.go

@@ -19,11 +19,16 @@ package main
 
 import (
 	"context"
+	"crypto/rand"
+	"encoding/base64"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"log"
+	"net"
 	"net/http"
 	"os"
+	"time"
 
 	"golang.org/x/oauth2"
 	"golang.org/x/oauth2/google"
@@ -45,18 +50,87 @@ func getClient(config *oauth2.Config) *http.Client {
 	return config.Client(context.Background(), tok)
 }
 
+// generateState creates a secure random string for the OAuth2 state parameter.
+func generateState() string {
+	b := make([]byte, 16)
+	rand.Read(b)
+	return base64.URLEncoding.EncodeToString(b)
+}
+
+// Request a token from the web, then returns the retrieved token.
 // Request a token from the web, then returns the retrieved token.
 func getTokenFromWeb(config *oauth2.Config) *oauth2.Token {
-	authURL := config.AuthCodeURL("state-token", oauth2.AccessTypeOffline)
-	fmt.Printf("Go to the following link in your browser then type the "+
-		"authorization code: \n%v\n", authURL)
+	listener, err := net.Listen("tcp", "localhost:0")
+	if err != nil {
+		log.Fatalf("Unable to start local listener: %v", err)
+	}
+
+	port := listener.Addr().(*net.TCPAddr).Port
+	config.RedirectURL = fmt.Sprintf("http://localhost:%d", port)
+
+	// Generate a dynamic, cryptographically secure state parameter
+	state := generateState()
+
+	authURL := config.AuthCodeURL(state, oauth2.AccessTypeOffline)
+	fmt.Printf("Go to the following link in your browser: \n%v\n", authURL)
+
+	// Channels to handle successful codes and errors
+	codeCh := make(chan string)
+	errCh := make(chan error)
+
+	m := http.NewServeMux()
+	server := &http.Server{Handler: m}
+
+	m.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+		if errStr := r.URL.Query().Get("error"); errStr != "" {
+			fmt.Fprintf(w, "Authentication error: %s. You may close this window.", errStr)
+			errCh <- errors.New(errStr)
+			return
+		}
+
+		returnedState := r.URL.Query().Get("state")
+		if returnedState != state {
+			fmt.Fprintf(w, "Security error: invalid state parameter. You may close this window.")
+			errCh <- errors.New("invalid state parameter")
+			return
+		}
+
+		code := r.URL.Query().Get("code")
+		if code != "" {
+			fmt.Fprintf(w, "Authentication successful! You may close this window.")
+			codeCh <- code
+		} else {
+			fmt.Fprintf(w, "Failed to get authorization code. You may close this window.")
+			errCh <- errors.New("authorization code missing")
+		}
+	})
+
+	go func() {
+		if err := server.Serve(listener); err != nil && err != http.ErrServerClosed {
+			log.Fatalf("Unable to start local web server: %v", err)
+		}
+	}()
 
 	var authCode string
-	if _, err := fmt.Scan(&authCode); err != nil {
-		log.Fatalf("Unable to read authorization code: %v", err)
+
+	// Wait for a successful code, an error, or a timeout
+	select {
+	case authCode = <-codeCh:
+		// Success case, proceed to shutdown and exchange
+	case callbackErr := <-errCh:
+		server.Shutdown(context.Background())
+		log.Fatalf("Authorization failed during callback: %v", callbackErr)
+	case <-time.After(3 * time.Minute):
+		// Timeout case: user took too long or closed the browser
+		server.Shutdown(context.Background())
+		log.Fatalf("Authorization timed out after 3 minutes. Please try again.")
 	}
 
-	tok, err := config.Exchange(context.TODO(), authCode)
+	// Shutdown the server gracefully upon success
+	server.Shutdown(context.Background())
+
+	// Exchange the authorization code for an access token using context.Background()
+	tok, err := config.Exchange(context.Background(), authCode)
 	if err != nil {
 		log.Fatalf("Unable to retrieve token from web: %v", err)
 	}