LFG

internal/proxy/proxy.go

@@ -0,0 +1,238 @@
+package proxy
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/http/httputil"
+	"net/url"
+	"sync"
+	"time"
+
+	"github.com/iwasforcedtobehere/goRZ/internal/config"
+	"github.com/iwasforcedtobehere/goRZ/internal/logger"
+)
+
+// Server represents the reverse proxy server
+type Server struct {
+	config      *config.Config
+	logger      *logger.Logger
+	loadBalancer LoadBalancer
+	healthChecker *HealthChecker
+}
+
+// LoadBalancer defines the interface for load balancing strategies
+type LoadBalancer interface {
+	NextTarget() (*config.TargetConfig, error)
+	UpdateTargets(targets []config.TargetConfig)
+}
+
+// HealthChecker handles health checking for proxy targets
+type HealthChecker struct {
+	config     *config.Config
+	logger     *logger.Logger
+	targets    map[string]*config.TargetConfig
+	httpClient *http.Client
+	stopChan   chan struct{}
+	mu         sync.RWMutex
+}
+
+// NewServer creates a new reverse proxy server
+func NewServer(cfg *config.Config, logger *logger.Logger) (*Server, error) {
+	// Create load balancer based on configuration
+	var lb LoadBalancer
+	switch cfg.Proxy.LoadBalancer {
+	case "roundrobin":
+		lb = NewRoundRobinLoadBalancer(cfg.Proxy.Targets)
+	case "leastconn":
+		lb = NewLeastConnectionsLoadBalancer(cfg.Proxy.Targets)
+	case "random":
+		lb = NewRandomLoadBalancer(cfg.Proxy.Targets)
+	default:
+		return nil, fmt.Errorf("unsupported load balancer: %s", cfg.Proxy.LoadBalancer)
+	}
+
+	// Create health checker
+	healthChecker := NewHealthChecker(cfg, logger)
+
+	return &Server{
+		config:       cfg,
+		logger:       logger,
+		loadBalancer: lb,
+		healthChecker: healthChecker,
+	}, nil
+}
+
+// ServeHTTP handles HTTP requests
+func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	// Get next target based on load balancing strategy
+	target, err := s.loadBalancer.NextTarget()
+	if err != nil {
+		s.logger.Error("Failed to get target", logger.Error(err))
+		http.Error(w, "Service unavailable", http.StatusServiceUnavailable)
+		return
+	}
+
+	// Check if target is healthy
+	if !target.Healthy {
+		s.logger.Warn("Target is unhealthy", logger.String("target", target.Name))
+		http.Error(w, "Service unavailable", http.StatusServiceUnavailable)
+		return
+	}
+
+	// Create reverse proxy
+	targetURL, err := url.Parse(target.Address)
+	if err != nil {
+		s.logger.Error("Failed to parse target URL", 
+			logger.String("target", target.Name), 
+			logger.String("address", target.Address),
+			logger.Error(err))
+		http.Error(w, "Internal server error", http.StatusInternalServerError)
+		return
+	}
+
+	proxy := httputil.NewSingleHostReverseProxy(targetURL)
+
+	// Set up custom director to modify request
+	originalDirector := proxy.Director
+	proxy.Director = func(req *http.Request) {
+		originalDirector(req)
+		// Add custom headers
+		req.Header.Set("X-Forwarded-Host", req.Host)
+		req.Header.Set("X-Forwarded-Proto", "https")
+		if req.TLS != nil {
+			req.Header.Set("X-Forwarded-Ssl", "on")
+		}
+		// Add proxy identification
+		req.Header.Set("X-Proxy-Server", "goRZ")
+	}
+
+	// Set up error handler
+	proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
+		s.logger.Error("Proxy error", logger.Error(err))
+		http.Error(w, "Proxy error", http.StatusBadGateway)
+	}
+
+	// Set up modify response handler
+	proxy.ModifyResponse = func(resp *http.Response) error {
+		// Add custom headers to response
+		resp.Header.Set("X-Served-By", "goRZ")
+		resp.Header.Set("X-Target", target.Name)
+		return nil
+	}
+
+	// Serve the request
+	s.logger.Debug("Proxying request", 
+		logger.String("method", r.Method),
+		logger.String("path", r.URL.Path),
+		logger.String("target", target.Name))
+
+	proxy.ServeHTTP(w, r)
+}
+
+// Start starts the proxy server
+func (s *Server) Start() error {
+	// Start health checker
+	if err := s.healthChecker.Start(); err != nil {
+		return fmt.Errorf("failed to start health checker: %w", err)
+	}
+	return nil
+}
+
+// Stop stops the proxy server
+func (s *Server) Stop() {
+	s.healthChecker.Stop()
+}
+
+// NewHealthChecker creates a new health checker
+func NewHealthChecker(cfg *config.Config, logger *logger.Logger) *HealthChecker {
+	targets := make(map[string]*config.TargetConfig)
+	for i := range cfg.Proxy.Targets {
+		targets[cfg.Proxy.Targets[i].Name] = &cfg.Proxy.Targets[i]
+	}
+
+	return &HealthChecker{
+		config:     cfg,
+		logger:     logger,
+		targets:    targets,
+		httpClient: &http.Client{Timeout: 5 * time.Second},
+		stopChan:   make(chan struct{}),
+	}
+}
+
+// Start starts the health checker
+func (h *HealthChecker) Start() error {
+	// Initial health check
+	h.checkAllTargets()
+
+	// Start periodic health checks
+	go func() {
+		ticker := time.NewTicker(time.Duration(h.config.Proxy.HealthCheckInterval) * time.Second)
+		defer ticker.Stop()
+
+		for {
+			select {
+			case <-ticker.C:
+				h.checkAllTargets()
+			case <-h.stopChan:
+				return
+			}
+		}
+	}()
+
+	return nil
+}
+
+// Stop stops the health checker
+func (h *HealthChecker) Stop() {
+	close(h.stopChan)
+}
+
+// checkAllTargets checks the health of all targets
+func (h *HealthChecker) checkAllTargets() {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+
+	for name, target := range h.targets {
+		healthy := h.checkTargetHealth(target)
+		if target.Healthy != healthy {
+			h.logger.Info("Target health status changed", 
+				logger.String("target", name), 
+				logger.Bool("healthy", healthy))
+			target.Healthy = healthy
+		}
+	}
+}
+
+// checkTargetHealth checks the health of a single target
+func (h *HealthChecker) checkTargetHealth(target *config.TargetConfig) bool {
+	targetURL, err := url.Parse(target.Address)
+	if err != nil {
+		h.logger.Error("Failed to parse target URL for health check", 
+			logger.String("target", target.Name),
+			logger.Error(err))
+		return false
+	}
+
+	healthURL := *targetURL
+	healthURL.Path = h.config.Proxy.HealthCheckPath
+
+	req, err := http.NewRequest("GET", healthURL.String(), nil)
+	if err != nil {
+		h.logger.Error("Failed to create health check request", 
+			logger.String("target", target.Name),
+			logger.Error(err))
+		return false
+	}
+
+	resp, err := h.httpClient.Do(req)
+	if err != nil {
+		h.logger.Error("Health check request failed", 
+			logger.String("target", target.Name),
+			logger.Error(err))
+		return false
+	}
+	defer resp.Body.Close()
+
+	return resp.StatusCode == http.StatusOK
+}