diff --git a/gmail/quickstart/quickstart.go b/gmail/quickstart/quickstart.go index 295a9b1..74184aa 100644 --- a/gmail/quickstart/quickstart.go +++ b/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) }