go-copilot-proxy/internal/cmd/serve.go

package cmd

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/rs/zerolog/log"
	"github.com/spf13/cobra"

	"gitlab.com/technofab/go-copilot-proxy/internal/auth"
	"gitlab.com/technofab/go-copilot-proxy/internal/server"
)

var serveCmd = &cobra.Command{
	Use:   "serve",
	Short: "Starts the proxy server",
	Long:  "Starts the HTTP proxy server that forwards requests to the GitHub Copilot API with automatic token management.",
	Run:   runServe,
}

func init() {
	serveCmd.Flags().String("host", "127.0.0.1", "Host to bind the server to")
	serveCmd.Flags().Int("port", 8080, "Port to run the server on")
}

func runServe(cmd *cobra.Command, args []string) {
	accessToken := os.Getenv("GO_COPILOT_PROXY_TOKEN")
	if accessToken == "" {
		log.Fatal().Msg("GO_COPILOT_PROXY_TOKEN environment variable is not set. Please set it to a secure token to protect your proxy.")
	}
	log.Debug().Msg("Proxy access token is configured.")

	ctx, cancel := context.WithCancel(context.Background())
	var wg sync.WaitGroup

	copilotAuth, err := auth.NewCopilotAuth()
	if err != nil {
		log.Fatal().Err(err).Msg("Failed to initialize Copilot Auth")
	}

	if err := copilotAuth.Start(ctx, &wg); err != nil {
		log.Fatal().Err(err).Msg("Failed to start Copilot auth background tasks. Did you run the 'auth' command first?")
	}

	r := server.NewRouter(accessToken, copilotAuth)

	host, _ := cmd.Flags().GetString("host")
	port, _ := cmd.Flags().GetInt("port")
	addr := fmt.Sprintf("%s:%d", host, port)
	httpServer := &http.Server{
		Addr:    addr,
		Handler: r,
	}

	go func() {
		log.Info().Str("address", addr).Msg("Starting server...")
		if err := httpServer.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
			log.Fatal().Err(err).Msg("Server failed to start")
		}
	}()

	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit
	log.Info().Msg("Shutting down server...")

	cancel()

	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer shutdownCancel()

	if err := httpServer.Shutdown(shutdownCtx); err != nil {
		log.Error().Err(err).Msg("Server forced to shutdown")
	}

	log.Debug().Msg("Waiting for background tasks to complete...")
	wg.Wait()
	log.Info().Msg("Server exiting.")
}