Luis Paterrebuild branchf5dc3801 year agoHistory
317 lines · 8.5 KB · Go
RawDownload
package codex import (	"context"	"errors"	"fmt"	"net"	"net/http"	"strings"	"sync"	"time" 	log "github.com/sirupsen/logrus") // OAuthServer handles the local HTTP server for OAuth callbacks.// It listens for the authorization code response from the OAuth provider// and captures the necessary parameters to complete the authentication flow.type OAuthServer struct {	// server is the underlying HTTP server instance	server *http.Server	// port is the port number on which the server listens	port int	// resultChan is a channel for sending OAuth results	resultChan chan *OAuthResult	// errorChan is a channel for sending OAuth errors	errorChan chan error	// mu is a mutex for protecting server state	mu sync.Mutex	// running indicates whether the server is currently running	running bool} // OAuthResult contains the result of the OAuth callback.// It holds either the authorization code and state for successful authentication// or an error message if the authentication failed.type OAuthResult struct {	// Code is the authorization code received from the OAuth provider	Code string	// State is the state parameter used to prevent CSRF attacks	State string	// Error contains any error message if the OAuth flow failed	Error string} // NewOAuthServer creates a new OAuth callback server.// It initializes the server with the specified port and creates channels// for handling OAuth results and errors.//// Parameters://   - port: The port number on which the server should listen//// Returns://   - *OAuthServer: A new OAuthServer instancefunc NewOAuthServer(port int) *OAuthServer {	return &OAuthServer{		port:       port,		resultChan: make(chan *OAuthResult, 1),		errorChan:  make(chan error, 1),	}} // Start starts the OAuth callback server.// It sets up the HTTP handlers for the callback and success endpoints,// and begins listening on the specified port.//// Returns://   - error: An error if the server fails to startfunc (s *OAuthServer) Start() error {	s.mu.Lock()	defer s.mu.Unlock() 	if s.running {		return fmt.Errorf("server is already running")	} 	// Check if port is available	if !s.isPortAvailable() {		return fmt.Errorf("port %d is already in use", s.port)	} 	mux := http.NewServeMux()	mux.HandleFunc("/auth/callback", s.handleCallback)	mux.HandleFunc("/success", s.handleSuccess) 	s.server = &http.Server{		Addr:         fmt.Sprintf(":%d", s.port),		Handler:      mux,		ReadTimeout:  10 * time.Second,		WriteTimeout: 10 * time.Second,	} 	s.running = true 	// Start server in goroutine	go func() {		if err := s.server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {			s.errorChan <- fmt.Errorf("server failed to start: %w", err)		}	}() 	// Give server a moment to start	time.Sleep(100 * time.Millisecond) 	return nil} // Stop gracefully stops the OAuth callback server.// It performs a graceful shutdown of the HTTP server with a timeout.//// Parameters://   - ctx: The context for controlling the shutdown process//// Returns://   - error: An error if the server fails to stop gracefullyfunc (s *OAuthServer) Stop(ctx context.Context) error {	s.mu.Lock()	defer s.mu.Unlock() 	if !s.running || s.server == nil {		return nil	} 	log.Debug("Stopping OAuth callback server") 	// Create a context with timeout for shutdown	shutdownCtx, cancel := context.WithTimeout(ctx, 5*time.Second)	defer cancel() 	err := s.server.Shutdown(shutdownCtx)	s.running = false	s.server = nil 	return err} // WaitForCallback waits for the OAuth callback with a timeout.// It blocks until either an OAuth result is received, an error occurs,// or the specified timeout is reached.//// Parameters://   - timeout: The maximum time to wait for the callback//// Returns://   - *OAuthResult: The OAuth result if successful//   - error: An error if the callback times out or an error occursfunc (s *OAuthServer) WaitForCallback(timeout time.Duration) (*OAuthResult, error) {	select {	case result := <-s.resultChan:		return result, nil	case err := <-s.errorChan:		return nil, err	case <-time.After(timeout):		return nil, fmt.Errorf("timeout waiting for OAuth callback")	}} // handleCallback handles the OAuth callback endpoint.// It extracts the authorization code and state from the callback URL,// validates the parameters, and sends the result to the waiting channel.//// Parameters://   - w: The HTTP response writer//   - r: The HTTP requestfunc (s *OAuthServer) handleCallback(w http.ResponseWriter, r *http.Request) {	log.Debug("Received OAuth callback") 	// Validate request method	if r.Method != http.MethodGet {		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)		return	} 	// Extract parameters	query := r.URL.Query()	code := query.Get("code")	state := query.Get("state")	errorParam := query.Get("error") 	// Validate required parameters	if errorParam != "" {		log.Errorf("OAuth error received: %s", errorParam)		result := &OAuthResult{			Error: errorParam,		}		s.sendResult(result)		http.Error(w, fmt.Sprintf("OAuth error: %s", errorParam), http.StatusBadRequest)		return	} 	if code == "" {		log.Error("No authorization code received")		result := &OAuthResult{			Error: "no_code",		}		s.sendResult(result)		http.Error(w, "No authorization code received", http.StatusBadRequest)		return	} 	if state == "" {		log.Error("No state parameter received")		result := &OAuthResult{			Error: "no_state",		}		s.sendResult(result)		http.Error(w, "No state parameter received", http.StatusBadRequest)		return	} 	// Send successful result	result := &OAuthResult{		Code:  code,		State: state,	}	s.sendResult(result) 	// Redirect to success page	http.Redirect(w, r, "/success", http.StatusFound)} // handleSuccess handles the success page endpoint.// It serves a user-friendly HTML page indicating that authentication was successful.//// Parameters://   - w: The HTTP response writer//   - r: The HTTP requestfunc (s *OAuthServer) handleSuccess(w http.ResponseWriter, r *http.Request) {	log.Debug("Serving success page") 	w.Header().Set("Content-Type", "text/html; charset=utf-8")	w.WriteHeader(http.StatusOK) 	// Parse query parameters for customization	query := r.URL.Query()	setupRequired := query.Get("setup_required") == "true"	platformURL := query.Get("platform_url")	if platformURL == "" {		platformURL = "https://platform.openai.com"	} 	// Generate success page HTML with dynamic content	successHTML := s.generateSuccessHTML(setupRequired, platformURL) 	_, err := w.Write([]byte(successHTML))	if err != nil {		log.Errorf("Failed to write success page: %v", err)	}} // generateSuccessHTML creates the HTML content for the success page.// It customizes the page based on whether additional setup is required// and includes a link to the platform.//// Parameters://   - setupRequired: Whether additional setup is required after authentication//   - platformURL: The URL to the platform for additional setup//// Returns://   - string: The HTML content for the success pagefunc (s *OAuthServer) generateSuccessHTML(setupRequired bool, platformURL string) string {	html := LoginSuccessHtml 	// Replace platform URL placeholder	html = strings.Replace(html, "{{PLATFORM_URL}}", platformURL, -1) 	// Add setup notice if required	if setupRequired {		setupNotice := strings.Replace(SetupNoticeHtml, "{{PLATFORM_URL}}", platformURL, -1)		html = strings.Replace(html, "{{SETUP_NOTICE}}", setupNotice, 1)	} else {		html = strings.Replace(html, "{{SETUP_NOTICE}}", "", 1)	} 	return html} // sendResult sends the OAuth result to the waiting channel.// It ensures that the result is sent without blocking the handler.//// Parameters://   - result: The OAuth result to sendfunc (s *OAuthServer) sendResult(result *OAuthResult) {	select {	case s.resultChan <- result:		log.Debug("OAuth result sent to channel")	default:		log.Warn("OAuth result channel is full, result dropped")	}} // isPortAvailable checks if the specified port is available.// It attempts to listen on the port to determine availability.//// Returns://   - bool: True if the port is available, false otherwisefunc (s *OAuthServer) isPortAvailable() bool {	addr := fmt.Sprintf(":%d", s.port)	listener, err := net.Listen("tcp", addr)	if err != nil {		return false	}	defer func() {		_ = listener.Close()	}()	return true} // IsRunning returns whether the server is currently running.//// Returns://   - bool: True if the server is running, false otherwisefunc (s *OAuthServer) IsRunning() bool {	s.mu.Lock()	defer s.mu.Unlock()	return s.running}

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t