diff --git a/cmd/proxy/main.go b/cmd/proxy/main.go new file mode 100644 index 0000000..f8b9289 --- /dev/null +++ b/cmd/proxy/main.go @@ -0,0 +1,127 @@ +package main + +import ( + "context" + "flag" + "fmt" + "log" + "net/http" + "os" + "os/signal" + "syscall" + "time" + + "github.com/user/go-proxy/internal/config" + "github.com/user/go-proxy/internal/proxy" + "github.com/user/go-proxy/internal/router" + "github.com/user/go-proxy/pkg/middleware" +) + +func main() { + // Parse command line flags + configPath := flag.String("config", "config.yaml", "Path to configuration file") + flag.Parse() + + // 1. Load configuration dynamically + cfg, err := config.LoadConfig(*configPath) + if err != nil { + log.Fatalf("Error loading config: %v", err) + os.Exit(1) + } + + fmt.Printf("Loaded configuration from %s\n", *configPath) + + // 2. Setup Router + r := router.NewRouter(cfg.HTTP) + + // 3. Setup Middleware Chain + handler := middleware.Chain( + r, + middleware.Recovery(), + middleware.Logger(), + middleware.RateLimit(10, 20), + ) + + // 4. Setup Config Watcher (Hot-Reload) + err = config.WatchConfig(*configPath, func(newCfg *config.Config) { + log.Println("Applying new configuration...") + r.UpdateConfig(newCfg.HTTP) + }) + if err != nil { + log.Printf("Warning: Failed to start config watcher: %v", err) + } + + // 5. Start TCP Proxy in a separate goroutine (if configured) + if cfg.TCP.Port > 0 && cfg.TCP.Backend != "" { + go func() { + if err := proxy.RunTCPProxy(cfg.TCP.Port, cfg.TCP.Backend); err != nil { + log.Printf("TCP proxy failed: %v", err) + } + }() + } + + // 6. Start UDP Proxy in a separate goroutine (if configured) + if cfg.UDP.Port > 0 && cfg.UDP.Backend != "" { + go func() { + if err := proxy.RunUDPProxy(cfg.UDP.Port, cfg.UDP.Backend); err != nil { + log.Printf("UDP proxy failed: %v", err) + } + }() + } + + // Prepare HTTP Servers + httpServer := &http.Server{ + Addr: fmt.Sprintf(":%d", cfg.HTTP.Port), + Handler: handler, + } + + var httpsServer *http.Server + if cfg.HTTP.TLS.Enabled { + httpsServer = &http.Server{ + Addr: fmt.Sprintf(":%d", cfg.HTTP.TLS.Port), + Handler: handler, + } + } + + // 6. Start HTTPS Server (if enabled) + if httpsServer != nil { + go func() { + log.Printf("Starting HTTPS proxy server on port %s\n", httpsServer.Addr) + if err := httpsServer.ListenAndServeTLS(cfg.HTTP.TLS.CertFile, cfg.HTTP.TLS.KeyFile); err != nil && err != http.ErrServerClosed { + log.Fatalf("HTTPS Server error: %v", err) + } + }() + } + + // 7. Start HTTP Server + go func() { + log.Printf("Starting HTTP proxy server on port %s\n", httpServer.Addr) + if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed { + log.Fatalf("HTTP Server error: %v", err) + } + }() + + // 8. Graceful Shutdown Waiter + quit := make(chan os.Signal, 1) + // kill (no param) default send syscall.SIGTERM + // kill -2 is syscall.SIGINT + // kill -9 is syscall.SIGKILL but can't be caught, so don't need to add it + signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM) + <-quit + log.Println("Shutdown signal received, draining traffic...") + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + if httpsServer != nil { + if err := httpsServer.Shutdown(ctx); err != nil { + log.Printf("HTTPS Server forced to shutdown: %v", err) + } + } + + if err := httpServer.Shutdown(ctx); err != nil { + log.Printf("HTTP Server forced to shutdown: %v", err) + } + + log.Println("Proxy server exiting gracefully") +} diff --git a/config.yaml b/config.yaml new file mode 100644 index 0000000..3686c44 --- /dev/null +++ b/config.yaml @@ -0,0 +1,34 @@ +http: + port: 18080 + tls: + enabled: false + port: 18443 + cert_file: "./cert.pem" + key_file: "./key.pem" + routers: + # Example 1: Route by API path with Load Balancing + - path: "/api/v1" + backends: + - "http://localhost:18081" + - "http://localhost:18082" + strip_prefix: true + + # Example 2: Route by Domain/Host + - host: "admin.local" + path: "/" + backend_url: "http://localhost:8082" + strip_prefix: false + + # Example 3: Default fallback route + - path: "/" + backend_url: "http://localhost:8083" + strip_prefix: false + +tcp: + port: 19000 + backend: "localhost:3306" + +udp: + port: 15353 + backend: "8.8.8.8:53" + diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..23b163e --- /dev/null +++ b/go.mod @@ -0,0 +1,10 @@ +module github.com/user/go-proxy + +go 1.26.6 + +require ( + github.com/fsnotify/fsnotify v1.10.1 // indirect + golang.org/x/sys v0.13.0 // indirect + golang.org/x/time v0.16.0 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect +) diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..7b71133 --- /dev/null +++ b/go.sum @@ -0,0 +1,9 @@ +github.com/fsnotify/fsnotify v1.10.1 h1:b0/UzAf9yR5rhf3RPm9gf3ehBPpf0oZKIjtpKrx59Ho= +github.com/fsnotify/fsnotify v1.10.1/go.mod h1:TLheqan6HD6GBK6PrDWyDPBaEV8LspOxvPSjC+bVfgo= +golang.org/x/sys v0.13.0 h1:Af8nKPmuFypiUBjVoU9V20FiaFXOcuZI21p0ycVYYGE= +golang.org/x/sys v0.13.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/time v0.16.0 h1:vMb6ptszcQMkcwiRTAuNNU50gom6++Q/6gY2hDM6VDE= +golang.org/x/time v0.16.0/go.mod h1:rVKOqvZeKvrDKTQiAHJ7wmwP0RzleSphoEA9RcdLA0s= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/internal/config/config.go b/internal/config/config.go new file mode 100644 index 0000000..86b1e89 --- /dev/null +++ b/internal/config/config.go @@ -0,0 +1,77 @@ +package config + +import ( + "encoding/json" + "fmt" + "os" + "path/filepath" + + "gopkg.in/yaml.v3" +) + +// Config represents the root configuration of the proxy. +type Config struct { + HTTP HTTPConfig `yaml:"http" json:"http"` + TCP TCPConfig `yaml:"tcp" json:"tcp"` + UDP UDPConfig `yaml:"udp" json:"udp"` +} + +type UDPConfig struct { + Port int `yaml:"port" json:"port"` + Backend string `yaml:"backend" json:"backend"` +} + +type HTTPConfig struct { + Port int `yaml:"port" json:"port"` + TLS TLSConfig `yaml:"tls" json:"tls"` + Routers []RouterRoute `yaml:"routers" json:"routers"` +} + +type TLSConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + Port int `yaml:"port" json:"port"` + CertFile string `yaml:"cert_file" json:"cert_file"` + KeyFile string `yaml:"key_file" json:"key_file"` +} + +type TCPConfig struct { + Port int `yaml:"port" json:"port"` + Backend string `yaml:"backend" json:"backend"` +} + +type RouterRoute struct { + Path string `yaml:"path" json:"path"` + Host string `yaml:"host" json:"host"` + Backends []string `yaml:"backends" json:"backends"` // Used for Load Balancing + BackendURL string `yaml:"backend_url" json:"backend_url"` // Legacy / Single backend + StripPrefix bool `yaml:"strip_prefix" json:"strip_prefix"` +} + +// LoadConfig reads a configuration file (YAML or JSON) and parses it into the Config struct. +func LoadConfig(path string) (*Config, error) { + file, err := os.Open(path) + if err != nil { + return nil, fmt.Errorf("failed to open config file: %w", err) + } + defer file.Close() + + var cfg Config + ext := filepath.Ext(path) + + switch ext { + case ".yaml", ".yml": + decoder := yaml.NewDecoder(file) + if err := decoder.Decode(&cfg); err != nil { + return nil, fmt.Errorf("failed to parse yaml config: %w", err) + } + case ".json": + decoder := json.NewDecoder(file) + if err := decoder.Decode(&cfg); err != nil { + return nil, fmt.Errorf("failed to parse json config: %w", err) + } + default: + return nil, fmt.Errorf("unsupported config format: %s", ext) + } + + return &cfg, nil +} diff --git a/internal/config/watcher.go b/internal/config/watcher.go new file mode 100644 index 0000000..6bac50b --- /dev/null +++ b/internal/config/watcher.go @@ -0,0 +1,55 @@ +package config + +import ( + "log" + "path/filepath" + + "github.com/fsnotify/fsnotify" +) + +// WatchConfig watches the config file for changes and triggers the callback +func WatchConfig(path string, onChange func(*Config)) error { + watcher, err := fsnotify.NewWatcher() + if err != nil { + return err + } + + absPath, err := filepath.Abs(path) + if err != nil { + return err + } + configDir := filepath.Dir(absPath) + + go func() { + defer watcher.Close() + for { + select { + case event, ok := <-watcher.Events: + if !ok { + return + } + // Many editors trigger Write or Rename when saving + if event.Has(fsnotify.Write) || event.Has(fsnotify.Create) { + if filepath.Base(event.Name) == filepath.Base(absPath) { + log.Printf("Config file changed: %s", event.Name) + newCfg, err := LoadConfig(absPath) + if err != nil { + log.Printf("Failed to reload config: %v", err) + continue + } + onChange(newCfg) + } + } + case err, ok := <-watcher.Errors: + if !ok { + return + } + log.Printf("Watcher error: %v", err) + } + } + }() + + // Watch the directory instead of the file to handle cases where + // editors use atomic saves (create new file, delete old, rename new) + return watcher.Add(configDir) +} diff --git a/internal/proxy/http.go b/internal/proxy/http.go new file mode 100644 index 0000000..4d8df48 --- /dev/null +++ b/internal/proxy/http.go @@ -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 +} diff --git a/internal/proxy/tcp.go b/internal/proxy/tcp.go new file mode 100644 index 0000000..253c1d0 --- /dev/null +++ b/internal/proxy/tcp.go @@ -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 +} diff --git a/internal/proxy/udp.go b/internal/proxy/udp.go new file mode 100644 index 0000000..b5a8bc5 --- /dev/null +++ b/internal/proxy/udp.go @@ -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]) + } +} diff --git a/internal/router/router.go b/internal/router/router.go new file mode 100644 index 0000000..271fd4f --- /dev/null +++ b/internal/router/router.go @@ -0,0 +1,90 @@ +package router + +import ( + "log" + "net/http" + "net/http/httputil" + "strings" + "sync" + + "github.com/user/go-proxy/internal/config" + "github.com/user/go-proxy/internal/proxy" +) + +// routeEntry stores the configuration and its pre-compiled ReverseProxy +type routeEntry struct { + config config.RouterRoute + proxy *httputil.ReverseProxy +} + +type Router struct { + mu sync.RWMutex + routes []routeEntry +} + +func NewRouter(cfg config.HTTPConfig) *Router { + r := &Router{} + r.UpdateConfig(cfg) + return r +} + +func (r *Router) UpdateConfig(cfg config.HTTPConfig) { + var newRoutes []routeEntry + + for _, routeCfg := range cfg.Routers { + // Collect backends (support both Backends array and single BackendURL) + backends := routeCfg.Backends + if len(backends) == 0 && routeCfg.BackendURL != "" { + backends = []string{routeCfg.BackendURL} + } + + if len(backends) == 0 { + log.Printf("Warning: Route %s has no backends configured, skipping", routeCfg.Path) + continue + } + + p, err := proxy.NewLoadBalancedHTTPProxy(backends) + if err != nil { + log.Printf("Error creating proxy for route %s: %v", routeCfg.Path, err) + continue + } + + newRoutes = append(newRoutes, routeEntry{ + config: routeCfg, + proxy: p, + }) + } + + r.mu.Lock() + defer r.mu.Unlock() + r.routes = newRoutes +} + +func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) { + r.mu.RLock() + routes := r.routes + r.mu.RUnlock() + + for _, route := range routes { + cfg := route.config + + if cfg.Host != "" && req.Host != cfg.Host { + continue + } + + if strings.HasPrefix(req.URL.Path, cfg.Path) { + if cfg.StripPrefix { + req.URL.Path = strings.TrimPrefix(req.URL.Path, cfg.Path) + if req.URL.Path == "" || req.URL.Path[0] != '/' { + req.URL.Path = "/" + req.URL.Path + } + } + + // Serve using the pre-compiled proxy + route.proxy.ServeHTTP(w, req) + return + } + } + + http.Error(w, "Not Found", http.StatusNotFound) +} diff --git a/pkg/middleware/logger.go b/pkg/middleware/logger.go new file mode 100644 index 0000000..1b92b7a --- /dev/null +++ b/pkg/middleware/logger.go @@ -0,0 +1,41 @@ +package middleware + +import ( + "log" + "net/http" + "time" +) + +// responseWriter is a custom wrapper to capture the status code +type responseWriter struct { + http.ResponseWriter + status int +} + +func (rw *responseWriter) WriteHeader(code int) { + rw.status = code + rw.ResponseWriter.WriteHeader(code) +} + +// Logger is a middleware that logs each incoming HTTP request +func Logger() Middleware { + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + start := time.Now() + + // Default status is 200 OK + rw := &responseWriter{ResponseWriter: w, status: http.StatusOK} + + next.ServeHTTP(rw, r) + + log.Printf( + "[%s] %s %s | %d | %v", + r.Method, + r.Host, + r.URL.Path, + rw.status, + time.Since(start), + ) + }) + } +} diff --git a/pkg/middleware/middleware.go b/pkg/middleware/middleware.go new file mode 100644 index 0000000..3223998 --- /dev/null +++ b/pkg/middleware/middleware.go @@ -0,0 +1,15 @@ +package middleware + +import "net/http" + +// Middleware defines a function that wraps an http.Handler +type Middleware func(http.Handler) http.Handler + +// Chain constructs a new http.Handler by chaining multiple Middlewares around a final http.Handler +func Chain(handler http.Handler, middlewares ...Middleware) http.Handler { + // Loop in reverse to preserve the order of execution + for i := len(middlewares) - 1; i >= 0; i-- { + handler = middlewares[i](handler) + } + return handler +} diff --git a/pkg/middleware/ratelimit.go b/pkg/middleware/ratelimit.go new file mode 100644 index 0000000..6d290d6 --- /dev/null +++ b/pkg/middleware/ratelimit.go @@ -0,0 +1,92 @@ +package middleware + +import ( + "net" + "net/http" + "sync" + "time" + + "golang.org/x/time/rate" +) + +// visitor defines the rate limiter for a single IP +type visitor struct { + limiter *rate.Limiter + lastSeen time.Time +} + +// IPRateLimiter holds a map of rate limiters per IP +type IPRateLimiter struct { + mu sync.Mutex + visitors map[string]*visitor + rate rate.Limit + burst int +} + +func NewIPRateLimiter(r rate.Limit, b int) *IPRateLimiter { + i := &IPRateLimiter{ + visitors: make(map[string]*visitor), + rate: r, + burst: b, + } + + // Background routine to clean up old visitors + go func() { + for { + time.Sleep(time.Minute) + i.cleanupVisitors() + } + }() + + return i +} + +func (i *IPRateLimiter) getLimiter(ip string) *rate.Limiter { + i.mu.Lock() + defer i.mu.Unlock() + + v, exists := i.visitors[ip] + if !exists { + limiter := rate.NewLimiter(i.rate, i.burst) + i.visitors[ip] = &visitor{ + limiter: limiter, + lastSeen: time.Now(), + } + return limiter + } + + v.lastSeen = time.Now() + return v.limiter +} + +func (i *IPRateLimiter) cleanupVisitors() { + i.mu.Lock() + defer i.mu.Unlock() + for ip, v := range i.visitors { + if time.Since(v.lastSeen) > 3*time.Minute { + delete(i.visitors, ip) + } + } +} + +// RateLimit is a middleware that limits requests per IP +func RateLimit(requestsPerSecond float64, burst int) Middleware { + limiter := NewIPRateLimiter(rate.Limit(requestsPerSecond), burst) + + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Extract IP + ip, _, err := net.SplitHostPort(r.RemoteAddr) + if err != nil { + ip = r.RemoteAddr + } + + if !limiter.getLimiter(ip).Allow() { + http.Error(w, "Too Many Requests", http.StatusTooManyRequests) + return + } + + next.ServeHTTP(w, r) + }) + } +} diff --git a/pkg/middleware/recovery.go b/pkg/middleware/recovery.go new file mode 100644 index 0000000..af8f75b --- /dev/null +++ b/pkg/middleware/recovery.go @@ -0,0 +1,22 @@ +package middleware + +import ( + "log" + "net/http" +) + +// Recovery is a middleware that recovers from any panics and writes a 500 if there was one +func Recovery() Middleware { + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + defer func() { + if err := recover(); err != nil { + log.Printf("PANIC RECOVERED: %v", err) + http.Error(w, "Internal Server Error", http.StatusInternalServerError) + } + }() + + next.ServeHTTP(w, r) + }) + } +} diff --git a/proxy.exe b/proxy.exe new file mode 100644 index 0000000..ca66528 Binary files /dev/null and b/proxy.exe differ diff --git a/tools/echoserver/main.go b/tools/echoserver/main.go new file mode 100644 index 0000000..4b1b9b1 --- /dev/null +++ b/tools/echoserver/main.go @@ -0,0 +1,36 @@ +package main + +import ( + "fmt" + "log" + "net/http" +) + +func main() { + ports := []string{"18081", "18082", "18083"} + + for _, port := range ports { + go startEchoServer(port) + } + + // Block main goroutine + select {} +} + +func startEchoServer(port string) { + mux := http.NewServeMux() + mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + msg := fmt.Sprintf("Response from Backend Server [Port: %s] | Requested Path: %s", port, r.URL.Path) + // Print to console so we can see which server handled it + log.Printf("[Port %s] Handled request for %s", port, r.URL.Path) + + w.Header().Set("Content-Type", "text/plain") + w.WriteHeader(http.StatusOK) + w.Write([]byte(msg + "\n")) + }) + + log.Printf("Starting Echo Server on :%s", port) + if err := http.ListenAndServe(":"+port, mux); err != nil { + log.Fatalf("Server on %s failed: %v", port, err) + } +}