Skip to content

Commit

Permalink
reset proxy handler
Browse files Browse the repository at this point in the history
  • Loading branch information
@allnash
allnash committed Nov 20, 2024
1 parent e71f793 commit 3910eaf
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 130 deletions.
27 changes: 6 additions & 21 deletions fast-server/config/config.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,18 +14,11 @@ import (

var ProductionConfigPath = "/etc/fast/config.yaml"

const (
LoadBalanceMethodRoundRobin = "round_robin"
LoadBalanceMethodLeastConn = "least_conn"
)

type ProxyConfig struct {
Host string `yaml:"host,omitempty"`
Hosts []string `yaml:"hosts,omitempty"`
Port int `yaml:"port"`
Protocol string `yaml:"protocol,omitempty"`
InsecureSkipVerify bool `yaml:"insecure_skip_verify,omitempty"`
LoadBalanceMethod string `yaml:"load_balance_method,omitempty"`
Host string `yaml:"host"`
Port int `yaml:"port"`
Protocol string `yaml:"protocol,omitempty"`
InsecureSkipVerify bool `yaml:"insecure_skip_verify,omitempty"`
}

func (p *ProxyConfig) setDefaults() {
Expand All @@ -36,23 +29,15 @@ func (p *ProxyConfig) setDefaults() {
}

func (p *ProxyConfig) validate() error {
if p.Host == "" && len(p.Hosts) == 0 {
return fmt.Errorf("proxy host or hosts must be specified")
if p.Host == "" {
return fmt.Errorf("proxy host cannot be empty")
}
if p.Port <= 0 || p.Port > 65535 {
return fmt.Errorf("invalid proxy port: %d", p.Port)
}
if p.Protocol != "http" && p.Protocol != "https" {
return fmt.Errorf("invalid proxy protocol: %s", p.Protocol)
}
if p.LoadBalanceMethod != "" &&
p.LoadBalanceMethod != LoadBalanceMethodRoundRobin &&
p.LoadBalanceMethod != LoadBalanceMethodLeastConn {
return fmt.Errorf("invalid load balance method: %s, must be one of: %s, %s",
p.LoadBalanceMethod,
LoadBalanceMethodRoundRobin,
LoadBalanceMethodLeastConn)
}
return nil
}

Expand Down
116 changes: 7 additions & 109 deletions fast-server/handlers/proxy_handler.go
View file
Original file line number Diff line number Diff line change
@@ -1,73 +1,18 @@
package handlers

import (
"bytes"
"crypto/tls"
"fast/config"
"fmt"
"github.com/labstack/echo/v4"
"github.com/labstack/echo/v4/middleware"
"io"
"math"
"net/http"
"net/url"
"sort"
"strings"
"sync"
"sync/atomic"
"time"
)

var (
loadBalancerCounter uint32
activeConnections sync.Map
)

const (
LoadBalanceMethodRoundRobin = "round_robin"
LoadBalanceMethodLeastConn = "least_conn"
)

func getNextHost(proxyConfig config.ProxyConfig) string {
// If only single host is specified, return it
if len(proxyConfig.Hosts) == 0 {
return proxyConfig.Host
}

// Choose load balancing method
switch proxyConfig.LoadBalanceMethod {
case LoadBalanceMethodLeastConn:
return getLeastConnHost(proxyConfig.Hosts)
case LoadBalanceMethodRoundRobin:
return getRoundRobinHost(proxyConfig.Hosts)
default:
// If no method specified or unknown, default to round_robin
return getRoundRobinHost(proxyConfig.Hosts)
}
}

func getRoundRobinHost(hosts []string) string {
next := atomic.AddUint32(&loadBalancerCounter, 1)
return hosts[next%uint32(len(hosts))]
}

func getLeastConnHost(hosts []string) string {
var selectedHost string
minConns := int64(math.MaxInt64)

for _, host := range hosts {
// Initialize with 0 if not exists
connsValue, _ := activeConnections.LoadOrStore(host, int64(0))
currentConns := connsValue.(int64)
if currentConns < minConns {
minConns = currentConns
selectedHost = host
}
}

return selectedHost
}

func getMatchingLocation(path string, locations []config.Location) *config.Location {
// Sort locations by path length in descending order (longest first)
sortedLocations := make([]config.Location, len(locations))
Expand Down Expand Up @@ -97,33 +42,16 @@ func HandleProxy(c echo.Context, domain config.Domain) error {
return echo.ErrNotFound
}

c.Logger().Debugf("Matched location path: %s", location.Path)
c.Logger().Infof("Matched location path: %s", location.Path)

// Use location's proxy config
proxyConfig := location.Proxy

// Get next host using load balancing
selectedHost := getNextHost(proxyConfig)

// Increment active connections for the selected host if using least_conn
if proxyConfig.LoadBalanceMethod == LoadBalanceMethodLeastConn {
connsValue, _ := activeConnections.LoadOrStore(selectedHost, int64(0))
currentConns := connsValue.(int64)
activeConnections.Store(selectedHost, currentConns+1)
defer func() {
connsValue, _ := activeConnections.Load(selectedHost)
currentConns := connsValue.(int64)
activeConnections.Store(selectedHost, currentConns-1)
}()
c.Logger().Debugf("Active connections for host %s: %d", selectedHost, currentConns+1)
}

targetScheme := proxyConfig.Protocol
if targetScheme == "" {
targetScheme = "http"
}

target, err := url.Parse(fmt.Sprintf("%s://%s:%d", targetScheme, selectedHost, proxyConfig.Port))
target, err := url.Parse(fmt.Sprintf("%s://%s:%d", targetScheme, proxyConfig.Host, proxyConfig.Port))
if err != nil {
c.Logger().Errorf("Error parsing proxy URL: %v", err)
return echo.ErrInternalServerError
Expand All @@ -143,17 +71,9 @@ func HandleProxy(c echo.Context, domain config.Domain) error {
}
}

c.Logger().Debugf("Using rewrite map: %v", rewriteMap)

// More detailed logging for load balancing
c.Logger().Infof("[LoadBalancer] %s request: %s → %s (host: %s, method: %s, scheme: %s)",
proxyConfig.LoadBalanceMethod,
originalHost+requestPath,
target.String(),
selectedHost,
c.Request().Method,
targetScheme,
)
c.Logger().Infof("Using rewrite map: %v", rewriteMap)
c.Logger().Infof("Proxying %s request from %s to %s%s",
c.Request().Method, originalHost, target.String(), requestPath)

transport := &http.Transport{
Proxy: http.ProxyFromEnvironment,
Expand Down Expand Up @@ -193,39 +113,17 @@ func HandleProxy(c echo.Context, domain config.Domain) error {
Rewrite: rewriteMap,
Transport: transport,
ModifyResponse: func(res *http.Response) error {
// Read and log the response body
body, err := io.ReadAll(res.Body)
if err != nil {
c.Logger().Errorf("[LoadBalancer] Failed to read response body: %v", err)
return err
}
err = res.Body.Close()
if err != nil {
c.Logger().Errorf("[LoadBalancer] Failed to read response body: %v", err)
return err
}
// Create new ReadCloser for the body
res.Body = io.NopCloser(bytes.NewBuffer(body))

c.Logger().Infof("[LoadBalancer] Response from %s: %d %s, Content-Type: %s, Body: %s",
selectedHost,
res.StatusCode,
http.StatusText(res.StatusCode),
res.Header.Get("Content-Type"),
string(body),
)
c.Logger().Infof("Proxy response status: %d for path: %s", res.StatusCode, requestPath)
return nil
},
})

// Set proxy headers with more detailed tracking
// Set proxy headers
c.Request().Header.Set("X-Forwarded-Host", originalHost)
c.Request().Header.Set("X-Real-IP", c.RealIP())
c.Request().Header.Set("X-Forwarded-For", c.RealIP())
c.Request().Header.Set("X-Forwarded-Proto", sourceScheme)
c.Request().Header.Set("X-Original-URI", requestPath)
c.Request().Header.Set("X-Load-Balanced-Host", selectedHost)
c.Request().Header.Set("X-Load-Balance-Method", proxyConfig.LoadBalanceMethod)
c.Request().Host = originalHost

return proxyMiddleware(func(c echo.Context) error {
Expand Down

0 comments on commit 3910eaf

Please sign in to comment.