Added support for custom header variables

- Added support for using nginx-like variables in custom headers
- Supported variables includes: $host, $remote_addr, $request_uri, $request_method, $content_length, $content_type, $uri, $args, $scheme, $query_string, $http_user_agent and $http_referer
- Added test case for custom header variable rewriter
This commit is contained in:
Toby Chui 2024-10-27 14:47:01 +08:00
parent f98e04a9fc
commit 172c5afa60
5 changed files with 262 additions and 3 deletions

View File

@ -159,9 +159,12 @@ func (h *ProxyHandler) hostRequest(w http.ResponseWriter, r *http.Request, targe
r.URL, _ = url.Parse(originalHostHeader)
}
//Populate the user-defined headers with the values from the request
rewrittenUserDefinedHeaders := rewrite.PopulateRequestHeaderVariables(r, target.UserDefinedHeaders)
//Build downstream and upstream header rules
upstreamHeaders, downstreamHeaders := rewrite.SplitUpDownStreamHeaders(&rewrite.HeaderRewriteOptions{
UserDefinedHeaders: target.UserDefinedHeaders,
UserDefinedHeaders: rewrittenUserDefinedHeaders,
HSTSMaxAge: target.HSTSMaxAge,
HSTSIncludeSubdomains: target.ContainsWildcardName(true),
EnablePermissionPolicyHeader: target.EnablePermissionPolicyHeader,
@ -234,9 +237,12 @@ func (h *ProxyHandler) vdirRequest(w http.ResponseWriter, r *http.Request, targe
r.URL, _ = url.Parse(originalHostHeader)
}
//Populate the user-defined headers with the values from the request
rewrittenUserDefinedHeaders := rewrite.PopulateRequestHeaderVariables(r, target.parent.UserDefinedHeaders)
//Build downstream and upstream header rules, use the parent (subdomain) endpoint's headers
upstreamHeaders, downstreamHeaders := rewrite.SplitUpDownStreamHeaders(&rewrite.HeaderRewriteOptions{
UserDefinedHeaders: target.parent.UserDefinedHeaders,
UserDefinedHeaders: rewrittenUserDefinedHeaders,
HSTSMaxAge: target.parent.HSTSMaxAge,
HSTSIncludeSubdomains: target.parent.ContainsWildcardName(true),
EnablePermissionPolicyHeader: target.parent.EnablePermissionPolicyHeader,

View File

@ -0,0 +1,63 @@
package rewrite
import (
"fmt"
"net/http"
"strings"
)
// GetHeaderVariableValuesFromRequest returns a map of header variables and their values
// note that variables behavior is not exactly identical to nginx variables
func GetHeaderVariableValuesFromRequest(r *http.Request) map[string]string {
vars := make(map[string]string)
// Request-specific variables
vars["$host"] = r.Host
vars["$remote_addr"] = r.RemoteAddr
vars["$request_uri"] = r.RequestURI
vars["$request_method"] = r.Method
vars["$content_length"] = fmt.Sprintf("%d", r.ContentLength)
vars["$content_type"] = r.Header.Get("Content-Type")
// Parsed URI elements
vars["$uri"] = r.URL.Path
vars["$args"] = r.URL.RawQuery
vars["$scheme"] = r.URL.Scheme
vars["$query_string"] = r.URL.RawQuery
// User agent and referer
vars["$http_user_agent"] = r.UserAgent()
vars["$http_referer"] = r.Referer()
return vars
}
// CustomHeadersIncludeDynamicVariables checks if the user-defined headers contain dynamic variables
// use for early exit when processing the headers
func CustomHeadersIncludeDynamicVariables(userDefinedHeaders []*UserDefinedHeader) bool {
for _, header := range userDefinedHeaders {
if strings.Contains(header.Value, "$") {
return true
}
}
return false
}
// PopulateRequestHeaderVariables populates the user-defined headers with the values from the request
func PopulateRequestHeaderVariables(r *http.Request, userDefinedHeaders []*UserDefinedHeader) []*UserDefinedHeader {
if !CustomHeadersIncludeDynamicVariables(userDefinedHeaders) {
// Early exit if there are no dynamic variables
return userDefinedHeaders
}
vars := GetHeaderVariableValuesFromRequest(r)
populatedHeaders := []*UserDefinedHeader{}
// Populate the user-defined headers with the values from the request
for _, header := range userDefinedHeaders {
thisHeader := header.Copy()
for key, value := range vars {
thisHeader.Value = strings.ReplaceAll(thisHeader.Value, key, value)
}
populatedHeaders = append(populatedHeaders, thisHeader)
}
return populatedHeaders
}

View File

@ -0,0 +1,172 @@
package rewrite
import (
"net/http/httptest"
"testing"
)
func TestGetHeaderVariableValuesFromRequest(t *testing.T) {
// Create a sample request
req := httptest.NewRequest("GET", "https://example.com/test?foo=bar", nil)
req.Host = "example.com"
req.RemoteAddr = "192.168.1.1:12345"
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", "TestAgent")
req.Header.Set("Referer", "https://referer.com")
// Call the function
vars := GetHeaderVariableValuesFromRequest(req)
// Expected results
expected := map[string]string{
"$host": "example.com",
"$remote_addr": "192.168.1.1:12345",
"$request_uri": "https://example.com/test?foo=bar",
"$request_method": "GET",
"$content_length": "0", // ContentLength is 0 because there's no body in the request
"$content_type": "application/json",
"$uri": "/test",
"$args": "foo=bar",
"$scheme": "https",
"$query_string": "foo=bar",
"$http_user_agent": "TestAgent",
"$http_referer": "https://referer.com",
}
// Check each expected variable
for key, expectedValue := range expected {
if vars[key] != expectedValue {
t.Errorf("Expected %s to be %s, but got %s", key, expectedValue, vars[key])
}
}
}
func TestCustomHeadersIncludeDynamicVariables(t *testing.T) {
tests := []struct {
name string
headers []*UserDefinedHeader
expectedHasVar bool
}{
{
name: "No headers",
headers: []*UserDefinedHeader{},
expectedHasVar: false,
},
{
name: "Headers without dynamic variables",
headers: []*UserDefinedHeader{
{
Direction: HeaderDirection_ZoraxyToUpstream,
Key: "X-Custom-Header",
Value: "staticValue",
IsRemove: false,
},
{
Direction: HeaderDirection_ZoraxyToDownstream,
Key: "X-Another-Header",
Value: "staticValue",
IsRemove: false,
},
},
expectedHasVar: false,
},
{
name: "Headers with one dynamic variable",
headers: []*UserDefinedHeader{
{
Direction: HeaderDirection_ZoraxyToUpstream,
Key: "X-Custom-Header",
Value: "$dynamicValue",
IsRemove: false,
},
},
expectedHasVar: true,
},
{
name: "Headers with multiple dynamic variables",
headers: []*UserDefinedHeader{
{
Direction: HeaderDirection_ZoraxyToUpstream,
Key: "X-Custom-Header",
Value: "$dynamicValue1",
IsRemove: false,
},
{
Direction: HeaderDirection_ZoraxyToDownstream,
Key: "X-Another-Header",
Value: "$dynamicValue2",
IsRemove: false,
},
},
expectedHasVar: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
hasVar := CustomHeadersIncludeDynamicVariables(tt.headers)
if hasVar != tt.expectedHasVar {
t.Errorf("Expected %v, but got %v", tt.expectedHasVar, hasVar)
}
})
}
}
func TestPopulateRequestHeaderVariables(t *testing.T) {
// Create a sample request with specific values
req := httptest.NewRequest("GET", "https://example.com/test?foo=bar", nil)
req.Host = "example.com"
req.RemoteAddr = "192.168.1.1:12345"
req.Header.Set("User-Agent", "TestAgent")
req.Header.Set("Referer", "https://referer.com")
// Define user-defined headers with dynamic variables
userDefinedHeaders := []*UserDefinedHeader{
{
Direction: HeaderDirection_ZoraxyToUpstream,
Key: "X-Forwarded-Host",
Value: "$host",
},
{
Direction: HeaderDirection_ZoraxyToDownstream,
Key: "X-Client-IP",
Value: "$remote_addr",
},
{
Direction: HeaderDirection_ZoraxyToDownstream,
Key: "X-Custom-Header",
Value: "$request_uri",
},
}
// Call the function with the test data
resultHeaders := PopulateRequestHeaderVariables(req, userDefinedHeaders)
// Expected results after variable substitution
expectedHeaders := []*UserDefinedHeader{
{
Direction: HeaderDirection_ZoraxyToUpstream,
Key: "X-Forwarded-Host",
Value: "example.com",
},
{
Direction: HeaderDirection_ZoraxyToDownstream,
Key: "X-Client-IP",
Value: "192.168.1.1:12345",
},
{
Direction: HeaderDirection_ZoraxyToDownstream,
Key: "X-Custom-Header",
Value: "https://example.com/test?foo=bar",
},
}
// Validate results
for i, expected := range expectedHeaders {
if resultHeaders[i].Direction != expected.Direction ||
resultHeaders[i].Key != expected.Key ||
resultHeaders[i].Value != expected.Value {
t.Errorf("Expected header %v, but got %v", expected, resultHeaders[i])
}
}
}

View File

@ -1,6 +1,10 @@
package rewrite
import "imuslab.com/zoraxy/mod/dynamicproxy/permissionpolicy"
import (
"encoding/json"
"imuslab.com/zoraxy/mod/dynamicproxy/permissionpolicy"
)
/*
typdef.go
@ -32,3 +36,16 @@ type HeaderRewriteOptions struct {
EnablePermissionPolicyHeader bool //Enable injection of permission policy header
PermissionPolicy *permissionpolicy.PermissionsPolicy //Permission policy header
}
// Utilities for header rewrite
func (h *UserDefinedHeader) GetDirection() HeaderDirection {
return h.Direction
}
// Copy eturns a deep copy of the UserDefinedHeader
func (h *UserDefinedHeader) Copy() *UserDefinedHeader {
result := UserDefinedHeader{}
js, _ := json.Marshal(h)
json.Unmarshal(js, &result)
return &result
}

View File

@ -318,6 +318,7 @@ func startupSequence() {
SystemWideLogger.PrintAndLog("warning", "Invalid start flag combination: docker=true && runtime.GOOS == windows. Running in docker UX development mode.", nil)
}
DockerUXOptimizer = dockerux.NewDockerOptimizer(*runningInDocker, SystemWideLogger)
}
// This sequence start after everything is initialized