diff --git a/src/mod/dynamicproxy/proxyRequestHandler.go b/src/mod/dynamicproxy/proxyRequestHandler.go index 87578cc..348fe11 100644 --- a/src/mod/dynamicproxy/proxyRequestHandler.go +++ b/src/mod/dynamicproxy/proxyRequestHandler.go @@ -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, diff --git a/src/mod/dynamicproxy/rewrite/headervars.go b/src/mod/dynamicproxy/rewrite/headervars.go new file mode 100644 index 0000000..d57c9cc --- /dev/null +++ b/src/mod/dynamicproxy/rewrite/headervars.go @@ -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 +} diff --git a/src/mod/dynamicproxy/rewrite/headervars_test.go b/src/mod/dynamicproxy/rewrite/headervars_test.go new file mode 100644 index 0000000..8d52ce8 --- /dev/null +++ b/src/mod/dynamicproxy/rewrite/headervars_test.go @@ -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]) + } + } +} diff --git a/src/mod/dynamicproxy/rewrite/typedef.go b/src/mod/dynamicproxy/rewrite/typedef.go index 3462909..c4965ee 100644 --- a/src/mod/dynamicproxy/rewrite/typedef.go +++ b/src/mod/dynamicproxy/rewrite/typedef.go @@ -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 +} diff --git a/src/start.go b/src/start.go index b1764f0..fb3e05c 100644 --- a/src/start.go +++ b/src/start.go @@ -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