pub

2026-09-15 01:38:15 +09:00
parent e748221831
commit 4e8a65f40d
16 changed files with 906 additions and 0 deletions
@@ -0,0 +1,132 @@
+package proxy
+
+import (
+	"fmt"
+	"log"
+	"net"
+	"net/http"
+	"net/http/httputil"
+	"net/url"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+// Backend represents a target server with its health state
+type Backend struct {
+	URL   *url.URL
+	Alive bool
+	mu    sync.RWMutex
+}
+
+func (b *Backend) SetAlive(alive bool) {
+	b.mu.Lock()
+	b.Alive = alive
+	b.mu.Unlock()
+}
+
+func (b *Backend) IsAlive() bool {
+	b.mu.RLock()
+	alive := b.Alive
+	b.mu.RUnlock()
+	return alive
+}
+
+// LoadBalancer holds a slice of backends and a counter for Round-Robin
+type LoadBalancer struct {
+	backends []*Backend
+	current  uint64
+}
+
+// Next returns the next healthy backend URL using Round-Robin
+func (lb *LoadBalancer) Next() *url.URL {
+	// Loop at most len(backends) times to find a healthy one
+	for i := 0; i < len(lb.backends); i++ {
+		next := atomic.AddUint64(&lb.current, 1)
+		idx := next % uint64(len(lb.backends))
+		b := lb.backends[idx]
+
+		if b.IsAlive() {
+			return b.URL
+		}
+	}
+	return nil // No healthy backends available
+}
+
+// StartHealthCheck begins a background routine to ping backends
+func (lb *LoadBalancer) StartHealthCheck() {
+	go func() {
+		for {
+			for _, b := range lb.backends {
+				go func(backend *Backend) {
+					// TCP Dial is a safe, universal health check for HTTP servers
+					host := backend.URL.Host
+					if backend.URL.Port() == "" {
+						if backend.URL.Scheme == "https" {
+							host += ":443"
+						} else {
+							host += ":80"
+						}
+					}
+					
+					conn, err := net.DialTimeout("tcp", host, 2*time.Second)
+					if err != nil {
+						if backend.IsAlive() {
+							log.Printf("[HealthCheck] Backend %s went DOWN: %v", backend.URL.String(), err)
+							backend.SetAlive(false)
+						}
+					} else {
+						conn.Close()
+						if !backend.IsAlive() {
+							log.Printf("[HealthCheck] Backend %s came UP", backend.URL.String())
+							backend.SetAlive(true)
+						}
+					}
+				}(b)
+			}
+			time.Sleep(5 * time.Second) // Check every 5 seconds
+		}
+	}()
+}
+
+// NewLoadBalancedHTTPProxy creates a ReverseProxy that load balances across multiple backends
+func NewLoadBalancedHTTPProxy(backends []string) (*httputil.ReverseProxy, error) {
+	if len(backends) == 0 {
+		return nil, fmt.Errorf("no backends provided")
+	}
+
+	lb := &LoadBalancer{}
+	for _, b := range backends {
+		u, err := url.Parse(b)
+		if err != nil {
+			return nil, fmt.Errorf("invalid backend url %s: %w", b, err)
+		}
+		// Initially mark all as alive
+		lb.backends = append(lb.backends, &Backend{URL: u, Alive: true})
+	}
+
+	// Start health checking in the background
+	lb.StartHealthCheck()
+
+	proxy := &httputil.ReverseProxy{
+		Director: func(req *http.Request) {
+			target := lb.Next()
+			if target == nil {
+				// We can't return an error directly from Director, but setting Scheme to empty
+				// will cause the RoundTripper to fail, returning 502 Bad Gateway by default.
+				return
+			}
+			req.URL.Scheme = target.Scheme
+			req.URL.Host = target.Host
+			// Note: We don't overwrite req.URL.Path here, as it's modified by the router
+			req.Header.Set("X-Forwarded-Host", req.Header.Get("Host"))
+		},
+		// Optional custom error handler when backends are unreachable
+		ErrorHandler: func(w http.ResponseWriter, r *http.Request, err error) {
+			log.Printf("Proxy error: %v", err)
+			http.Error(w, "502 Bad Gateway - No healthy backends available", http.StatusBadGateway)
+		},
+	}
+
+	return proxy, nil
+}
@@ -0,0 +1,59 @@
+package proxy
+
+import (
+	"fmt"
+	"io"
+	"log"
+	"net"
+)
+
+// RunTCPProxy starts a basic TCP reverse proxy listening on the given port
+// and forwarding all traffic to the target backend.
+func RunTCPProxy(port int, targetBackend string) error {
+	addr := fmt.Sprintf(":%d", port)
+	listener, err := net.Listen("tcp", addr)
+	if err != nil {
+		return err
+	}
+	defer listener.Close()
+
+	log.Printf("Starting TCP proxy on %s, forwarding to %s\n", addr, targetBackend)
+
+	for {
+		clientConn, err := listener.Accept()
+		if err != nil {
+			log.Printf("TCP Accept error: %v", err)
+			continue
+		}
+
+		go handleTCPConnection(clientConn, targetBackend)
+	}
+}
+
+func handleTCPConnection(clientConn net.Conn, targetBackend string) {
+	defer clientConn.Close()
+
+	backendConn, err := net.Dial("tcp", targetBackend)
+	if err != nil {
+		log.Printf("TCP Dial error to %s: %v", targetBackend, err)
+		return
+	}
+	defer backendConn.Close()
+
+	errChan := make(chan error, 2)
+
+	// Copy data from client to backend
+	go func() {
+		_, err := io.Copy(backendConn, clientConn)
+		errChan <- err
+	}()
+
+	// Copy data from backend to client
+	go func() {
+		_, err := io.Copy(clientConn, backendConn)
+		errChan <- err
+	}()
+
+	// Wait for one of the streams to finish (either EOF or error)
+	<-errChan
+}
@@ -0,0 +1,107 @@
+package proxy
+
+import (
+	"fmt"
+	"log"
+	"net"
+	"sync"
+	"time"
+)
+
+// udpSession tracks the state of a single client's UDP "connection"
+type udpSession struct {
+	clientAddr  *net.UDPAddr
+	backendConn *net.UDPConn
+	lastActive  time.Time
+}
+
+// RunUDPProxy starts a UDP reverse proxy, establishing a NAT-like session mapping for clients
+func RunUDPProxy(port int, targetBackend string) error {
+	addr := fmt.Sprintf(":%d", port)
+	proxyAddr, err := net.ResolveUDPAddr("udp", addr)
+	if err != nil {
+		return err
+	}
+
+	backendAddr, err := net.ResolveUDPAddr("udp", targetBackend)
+	if err != nil {
+		return err
+	}
+
+	proxyConn, err := net.ListenUDP("udp", proxyAddr)
+	if err != nil {
+		return err
+	}
+	defer proxyConn.Close()
+
+	log.Printf("Starting UDP proxy on %s, forwarding to %s\n", addr, targetBackend)
+
+	var mu sync.RWMutex
+	sessions := make(map[string]*udpSession)
+
+	// Background routine to clean up inactive sessions
+	go func() {
+		for {
+			time.Sleep(1 * time.Minute)
+			mu.Lock()
+			for k, s := range sessions {
+				// If a client hasn't sent or received packets in 3 minutes, close it
+				if time.Since(s.lastActive) > 3*time.Minute {
+					s.backendConn.Close()
+					delete(sessions, k)
+				}
+			}
+			mu.Unlock()
+		}
+	}()
+
+	buffer := make([]byte, 65535)
+	for {
+		n, clientAddr, err := proxyConn.ReadFromUDP(buffer)
+		if err != nil {
+			log.Printf("UDP Read error: %v", err)
+			continue
+		}
+
+		clientKey := clientAddr.String()
+		mu.RLock()
+		session, exists := sessions[clientKey]
+		mu.RUnlock()
+
+		if !exists {
+			// New client session detected, dial the backend
+			bConn, err := net.DialUDP("udp", nil, backendAddr)
+			if err != nil {
+				log.Printf("Failed to connect to UDP backend: %v", err)
+				continue
+			}
+
+			session = &udpSession{
+				clientAddr:  clientAddr,
+				backendConn: bConn,
+				lastActive:  time.Now(),
+			}
+
+			mu.Lock()
+			sessions[clientKey] = session
+			mu.Unlock()
+
+			// Start a goroutine to read from the backend and send back to this specific client
+			go func(s *udpSession) {
+				bBuf := make([]byte, 65535)
+				for {
+					bn, err := s.backendConn.Read(bBuf)
+					if err != nil {
+						break // Connection closed or error
+					}
+					s.lastActive = time.Now()
+					proxyConn.WriteToUDP(bBuf[:bn], s.clientAddr)
+				}
+			}(session)
+		}
+
+		// Forward the client's packet to the backend
+		session.lastActive = time.Now()
+		session.backendConn.Write(buffer[:n])
+	}
+}