feat: option to rewrite the host header

README.md

@@ -106,6 +106,15 @@ You can enable dashboard to see the status of the agents with
 `--management-dashboard-enable=true` flag. If enabled, it is available at the 
 management server address under `/dashboard` path.
 
+#### Rewriting the `Host` Header
+
+> [!NOTE]
+> Available since v0.8.0
+
+In some cases (see: [#20](https://github.com/distantmagic/paddler/issues/20)), you might want to rewrite the `Host` header.
+
+In such cases, you can use the `--rewrite-host-header` flag. If used, Paddler will use the `external` host provided by agents instead of the balancer host when forwarding the requests.
+
 ## Feature Highlights
 
 ### Aggregated Health Status

cmd/Balancer.go

@@ -48,6 +48,8 @@ func (self *Balancer) Action(cliContext *cli.Context) error {
 		return err
 	}
 
+	self.Logger.Debug("rewrite-host-header", "enabled", self.LoadBalancerConfiguration.RewriteHostHeader)
+
 	bufferedRequestsStats := &loadbalancer.BufferedRequestsStats{}
 
 	managementServer := &management.Server{
@@ -61,6 +63,7 @@ func (self *Balancer) Action(cliContext *cli.Context) error {
 		RespondToRegisterTarget: &management.RespondToRegisterTarget{
 			LoadBalancerTargetRegistrar: &loadbalancer.LoadBalancerTargetRegistrar{
 				HttpClient:                   http.DefaultClient,
+				LoadBalancerConfiguration:    self.LoadBalancerConfiguration,
 				LoadBalancerTargetCollection: loadBalancer.LoadBalancerTargetCollection,
 				Logger:                       self.Logger,
 			},

loadbalancer/CreateLlamaCppTargetReverseProxy.go

@@ -0,0 +1,36 @@
+package loadbalancer
+
+import (
+	"net/http"
+	"net/http/httputil"
+
+	"github.com/hashicorp/go-hclog"
+)
+
+func CreateLlamaCppTargetReverseProxy(
+	logger hclog.Logger,
+	loadBalancerConfiguration *LoadBalancerConfiguration,
+	targetConfiguration *LlamaCppTargetConfiguration,
+) *httputil.ReverseProxy {
+	reverseProxy := httputil.NewSingleHostReverseProxy(
+		targetConfiguration.LlamaCppConfiguration.HttpAddress.GetBaseUrl(),
+	)
+
+	reverseProxy.ErrorLog = logger.StandardLogger(&hclog.StandardLoggerOptions{
+		InferLevels: true,
+	})
+
+	if !loadBalancerConfiguration.RewriteHostHeader {
+		return reverseProxy
+	}
+
+	originalDirector := reverseProxy.Director
+
+	reverseProxy.Director = func(req *http.Request) {
+		originalDirector(req)
+
+		req.Host = targetConfiguration.LlamaCppConfiguration.HttpAddress.GetHostWithPort()
+	}
+
+	return reverseProxy
+}

loadbalancer/LoadBalancerConfiguration.go

@@ -5,4 +5,5 @@ import "time"
 type LoadBalancerConfiguration struct {
 	RequestBufferSize    uint
 	RequestBufferTimeout time.Duration
+	RewriteHostHeader    bool
 }

loadbalancer/LoadBalancerTargetCollection_bench_test.go

@@ -18,6 +18,7 @@ func BenchmarkLoadBalancerTargetCollection(b *testing.B) {
 
 	loadBalancerTargetRegistrar := &LoadBalancerTargetRegistrar{
 		HttpClient:                   http.DefaultClient,
+		LoadBalancerConfiguration:    &LoadBalancerConfiguration{},
 		LoadBalancerTargetCollection: NewLoadBalancerTargetCollection(llamaCppHealthStatusAggregate),
 		Logger:                       hclog.NewNullLogger(),
 	}

loadbalancer/LoadBalancerTargetRegistrar.go

@@ -2,7 +2,6 @@ package loadbalancer
 
 import (
 	"net/http"
-	"net/http/httputil"
 	"time"
 
 	"github.com/distantmagic/paddler/llamacpp"
@@ -11,6 +10,7 @@ import (
 
 type LoadBalancerTargetRegistrar struct {
 	HttpClient                   *http.Client
+	LoadBalancerConfiguration    *LoadBalancerConfiguration
 	LoadBalancerTargetCollection *LoadBalancerTargetCollection
 	Logger                       hclog.Logger
 }
@@ -43,14 +43,6 @@ func (self *LoadBalancerTargetRegistrar) registerTarget(
 		"host", targetConfiguration.LlamaCppConfiguration.HttpAddress.GetHostWithPort(),
 	)
 
-	reverseProxy := httputil.NewSingleHostReverseProxy(
-		targetConfiguration.LlamaCppConfiguration.HttpAddress.GetBaseUrl(),
-	)
-
-	reverseProxy.ErrorLog = self.Logger.Named("ReverseProxy").StandardLogger(&hclog.StandardLoggerOptions{
-		InferLevels: true,
-	})
-
 	self.LoadBalancerTargetCollection.RegisterTarget(&LlamaCppTarget{
 		LastUpdate: time.Now(),
 		LlamaCppClient: &llamacpp.LlamaCppClient{
@@ -60,6 +52,10 @@ func (self *LoadBalancerTargetRegistrar) registerTarget(
 		LlamaCppSlotsAggregatedStatus: llamaCppSlotsAggregatedStatus,
 		LlamaCppTargetConfiguration:   targetConfiguration,
 		RemainingTicksUntilRemoved:    3,
-		ReverseProxy:                  reverseProxy,
+		ReverseProxy: CreateLlamaCppTargetReverseProxy(
+			self.Logger.Named("ReverseProxy"),
+			self.LoadBalancerConfiguration,
+			targetConfiguration,
+		),
 	})
 }

loadbalancer/LoadBalancerTargetRegistrar_test.go

@@ -19,6 +19,7 @@ func TestTargetOrderIsPreserved(t *testing.T) {
 
 	loadBalancerTargetRegistrar := &LoadBalancerTargetRegistrar{
 		HttpClient:                   http.DefaultClient,
+		LoadBalancerConfiguration:    &LoadBalancerConfiguration{},
 		LoadBalancerTargetCollection: NewLoadBalancerTargetCollection(llamaCppHealthStatusAggregate),
 		Logger:                       hclog.NewNullLogger(),
 	}

loadbalancer/LoadBalancerTemporalManager_bench_test.go

@@ -22,6 +22,7 @@ func BenchmarkLoadBalancerTemporalManager(b *testing.B) {
 
 	loadBalancerTargetRegistrar := &LoadBalancerTargetRegistrar{
 		HttpClient:                   http.DefaultClient,
+		LoadBalancerConfiguration:    &LoadBalancerConfiguration{},
 		LoadBalancerTargetCollection: loadBalancerTargetCollection,
 		Logger:                       hclog.NewNullLogger(),
 	}

main.go

@@ -282,6 +282,11 @@ func main() {
 						Value:       60 * time.Second,
 						Destination: &balancer.LoadBalancerConfiguration.RequestBufferTimeout,
 					},
+					&cli.BoolFlag{
+						Name:        "rewrite-host-header",
+						Value:       false,
+						Destination: &balancer.LoadBalancerConfiguration.RewriteHostHeader,
+					},
 				},
 			},
 			// {