OpenAI Login with Subscription Support #2035

package main

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"strings"

	"golang.org/x/oauth2"
)

// Configurable settings
var (
	clientID     = os.Getenv("OAUTH_CLIENT_ID") // only client_id required
	redirectURL  = "http://localhost:8080/callback"
	authServerAuthURL  = "https://accounts.google.com/o/oauth2/v2/auth"
	authServerTokenURL = "https://oauth2.googleapis.com/token"
	scopes             = []string{"openid", "profile", "email"}
)

var (
	oauthConfig *oauth2.Config
	codeVerifier string
	tokenStore   *oauth2.Token
)

func main() {
	oauthConfig = &oauth2.Config{
		ClientID:    clientID,
		RedirectURL: redirectURL,
		Scopes:      scopes,
		Endpoint: oauth2.Endpoint{
			AuthURL:  authServerAuthURL,
			TokenURL: authServerTokenURL,
		},
	}

	http.HandleFunc("/", handleMain)
	http.HandleFunc("/login", handleLogin)
	http.HandleFunc("/callback", handleCallback)

	log.Println("Server started at http://localhost:8080")
	log.Fatal(http.ListenAndServe(":8080", nil))
}

func handleMain(w http.ResponseWriter, r *http.Request) {
	if tokenStore == nil {
		fmt.Fprintf(w, `<a href="/login">Login with OAuth (PKCE)</a>`)
		return
	}
	data, _ := json.MarshalIndent(tokenStore, "", "  ")
	fmt.Fprintf(w, "Access token stored:\n<pre>%s</pre>", string(data))
}

func handleLogin(w http.ResponseWriter, r *http.Request) {
	// Generate PKCE verifier & challenge
	codeVerifier = generateCodeVerifier()
	codeChallenge := generateCodeChallenge(codeVerifier)

	// Construct Auth URL with PKCE params
	authURL := oauthConfig.AuthCodeURL("state-token",
		oauth2.AccessTypeOffline,
		oauth2.SetAuthURLParam("code_challenge", codeChallenge),
		oauth2.SetAuthURLParam("code_challenge_method", "S256"),
	)

	http.Redirect(w, r, authURL, http.StatusFound)
}

func handleCallback(w http.ResponseWriter, r *http.Request) {
	code := r.URL.Query().Get("code")
	if code == "" {
		http.Error(w, "No code in request", http.StatusBadRequest)
		return
	}

	// Exchange code with PKCE verifier
	token, err := oauthConfig.Exchange(
		context.Background(),
		code,
		oauth2.SetAuthURLParam("code_verifier", codeVerifier),
	)
	if err != nil {
		http.Error(w, "Failed to exchange token: "+err.Error(), http.StatusInternalServerError)
		return
	}

	tokenStore = token
	data, _ := json.MarshalIndent(token, "", "  ")
	fmt.Fprintf(w, "OAuth2 PKCE token received:\n<pre>%s</pre>", string(data))
}

// --- Helpers for PKCE ---
func generateCodeVerifier() string {
	b := make([]byte, 32)
	if _, err := rand.Read(b); err != nil {
		panic(err)
	}
	return base64.RawURLEncoding.EncodeToString(b)
}

func generateCodeChallenge(verifier string) string {
	h := sha256.New()
	h.Write([]byte(verifier))
	return base64.RawURLEncoding.EncodeToString(h.Sum(nil))
}