Add possibility to adding additional request and response headers

README.md

@@ -18,13 +18,13 @@ These capabilities make Protty a useful tool for a variety of purposes, such as
 The following command will start a proxy on port 8080, and after starting, all traffic from port 8080 will be redirected to a remote host located at https://example.com
 
 ```shell
-docker run -p8080:80 -e REMOTE_URI=https://example.com:443 mgerasimchuk/protty:v0.2.0
+docker run -p8080:80 -e REMOTE_URI=https://example.com:443 mgerasimchuk/protty:v0.3.0
 ```
 
 ## Running options and runtime configuration
 
 ```
-» ~  docker run -p8080:80 -it mgerasimchuk/protty:v0.2.0 /bin/sh -c 'protty start --help'  
+» ~  docker run -p8080:80 -it mgerasimchuk/protty:v0.3.0 /bin/sh -c 'protty start --help'  
 Start the proxy
 
 Usage:
@@ -43,6 +43,9 @@ Examples:
   # Start the proxy with a specific remote URI and specific throttle rate limit 
   protty start --remote-uri https://www.githubstatus.com --throttle-rate-limit 2
 
+  # Start the proxy with a specific additional request headers
+  protty start --additional-request-headers 'Authorization: Bearer authtoken-with:any:symbols' --additional-request-headers 'X-Another-One: another-value'
+
   # Start the proxy with a specific SED expression for response transformation
   protty start --transform-response-body-sed 's|old|new|g'
 
@@ -60,8 +63,10 @@ Flags:
       --local-port int                            Verbosity level (panic, fatal, error, warn, info, debug, trace) | Env variable alias: LOCAL_PORT | Request header alias: X-PROTTY-LOCAL-PORT (default 80)
       --remote-uri string                         Listening port for the proxy | Env variable alias: REMOTE_URI | Request header alias: X-PROTTY-REMOTE-URI (default "https://example.com:443")
       --throttle-rate-limit float                 How many requests can be send to the remote resource per second | Env variable alias: THROTTLE_RATE_LIMIT | Request header alias: X-PROTTY-THROTTLE-RATE-LIMIT
+      --additional-request-headers stringArray    Array of additional request headers in format Header: Value | Env variable alias: ADDITIONAL_REQUEST_HEADERS | Request header alias: X-PROTTY-ADDITIONAL-REQUEST-HEADERS
       --transform-request-body-sed stringArray    Pipeline of SED expressions for request body transformation | Env variable alias: TRANSFORM_REQUEST_BODY_SED | Request header alias: X-PROTTY-TRANSFORM-REQUEST-BODY-SED
       --transform-request-body-jq stringArray     Pipeline of JQ expressions for request body transformation | Env variable alias: TRANSFORM_REQUEST_BODY_JQ | Request header alias: X-PROTTY-TRANSFORM-REQUEST-BODY-JQ
+      --additional-response-headers stringArray   Array of additional response headers in format Header: Value | Env variable alias: ADDITIONAL_RESPONSE_HEADERS | Request header alias: X-PROTTY-ADDITIONAL-RESPONSE-HEADERS
       --transform-response-body-sed stringArray   Pipeline of SED expressions for response body transformation | Env variable alias: TRANSFORM_RESPONSE_BODY_SED | Request header alias: X-PROTTY-TRANSFORM-RESPONSE-BODY-SED
       --transform-response-body-jq stringArray    Pipeline of JQ expressions for response body transformation | Env variable alias: TRANSFORM_RESPONSE_BODY_JQ | Request header alias: X-PROTTY-TRANSFORM-RESPONSE-BODY-JQ
   -h, --help                                      help for start

internal/adapter/cli/start.go

@@ -36,8 +36,10 @@ func NewStartCommand(cfg *config.StartCommandConfig, reverseProxySvc *service.Re
 	startCommand.cobraCmd.Flags().IntVar(buildFlagArgs(&cfg.LocalPort))
 	startCommand.cobraCmd.Flags().StringVar(buildFlagArgs(&cfg.RemoteURI))
 	startCommand.cobraCmd.Flags().Float64Var(buildFlagArgs(&cfg.ThrottleRateLimit))
+	startCommand.cobraCmd.Flags().StringArrayVar(buildFlagArgs(&cfg.AdditionalRequestHeaders))
 	startCommand.cobraCmd.Flags().StringArrayVar(buildFlagArgs(&cfg.TransformRequestBodySED))
 	startCommand.cobraCmd.Flags().StringArrayVar(buildFlagArgs(&cfg.TransformRequestBodyJQ))
+	startCommand.cobraCmd.Flags().StringArrayVar(buildFlagArgs(&cfg.AdditionalResponseHeaders))
 	startCommand.cobraCmd.Flags().StringArrayVar(buildFlagArgs(&cfg.TransformResponseBodySED))
 	startCommand.cobraCmd.Flags().StringArrayVar(buildFlagArgs(&cfg.TransformResponseBodyJQ))
 
@@ -77,6 +79,9 @@ func (c *StartCommand) getExamples() string {
   # Start the proxy with a specific remote URI and specific throttle rate limit 
   {{ .Cmd.CommandPath }} --{{ .Cfg.RemoteURI.GetFlagName }} https://www.githubstatus.com --{{ .Cfg.ThrottleRateLimit.GetFlagName }} 2
 
+  # Start the proxy with a specific additional request headers
+  {{ .Cmd.CommandPath }} --{{ .Cfg.AdditionalRequestHeaders.GetFlagName }} 'Authorization: Bearer authtoken-with:any:symbols' --{{ .Cfg.AdditionalRequestHeaders.GetFlagName }} 'X-Another-One: another-value'
+
   # Start the proxy with a specific SED expression for response transformation
   {{ .Cmd.CommandPath }} --{{ .Cfg.TransformResponseBodySED.GetFlagName }} 's|old|new|g'
 

internal/infrastructure/config/start_command_config.go

@@ -15,14 +15,16 @@ import (
 )
 
 type StartCommandConfig struct {
-	LogLevel                 Option[string]   `default:"debug" description:"On which host, the throttle rate limit should be applied"`
-	LocalPort                Option[int]      `default:"80" description:"Verbosity level (panic, fatal, error, warn, info, debug, trace)"`
-	RemoteURI                Option[string]   `default:"https://example.com:443" description:"Listening port for the proxy"`
-	ThrottleRateLimit        Option[float64]  `description:"How many requests can be send to the remote resource per second"`
-	TransformRequestBodySED  Option[[]string] `description:"Pipeline of SED expressions for request body transformation"`
-	TransformRequestBodyJQ   Option[[]string] `description:"Pipeline of JQ expressions for request body transformation"`
-	TransformResponseBodySED Option[[]string] `description:"Pipeline of SED expressions for response body transformation"`
-	TransformResponseBodyJQ  Option[[]string] `description:"Pipeline of JQ expressions for response body transformation"`
+	LogLevel                  Option[string]   `default:"debug" description:"On which host, the throttle rate limit should be applied"`
+	LocalPort                 Option[int]      `default:"80" description:"Verbosity level (panic, fatal, error, warn, info, debug, trace)"`
+	RemoteURI                 Option[string]   `default:"https://example.com:443" description:"Listening port for the proxy"`
+	ThrottleRateLimit         Option[float64]  `description:"How many requests can be send to the remote resource per second"`
+	AdditionalRequestHeaders  Option[[]string] `description:"Array of additional request headers in format Header: Value"`
+	TransformRequestBodySED   Option[[]string] `description:"Pipeline of SED expressions for request body transformation"`
+	TransformRequestBodyJQ    Option[[]string] `description:"Pipeline of JQ expressions for request body transformation"`
+	AdditionalResponseHeaders Option[[]string] `description:"Array of additional response headers in format Header: Value"`
+	TransformResponseBodySED  Option[[]string] `description:"Pipeline of SED expressions for response body transformation"`
+	TransformResponseBodyJQ   Option[[]string] `description:"Pipeline of JQ expressions for response body transformation"`
 }
 
 func GetStartCommandConfig() *StartCommandConfig {

internal/infrastructure/service/reverse_proxy.go

@@ -61,6 +61,16 @@ func (s *ReverseProxyService) modifyRequest(cfg config.StartCommandConfig, req *
 	// Deleting encoding to keep availability for changing response
 	req.Header.Del("Accept-Encoding")
 
+	// Add request headers
+	for _, header := range cfg.AdditionalRequestHeaders.Value {
+		kv := strings.SplitN(header, ": ", 2)
+		if len(kv) != 2 {
+			s.logger.Errorf("%s: %s: %s - %+v", util.GetCurrentFuncName(), util.GetFuncName(strings.SplitN), "returns not 2 values", kv)
+			continue
+		}
+		req.Header.Add(kv[0], kv[1])
+	}
+
 	sourceRequestBody, err := ioutil.ReadAll(req.Body)
 	if err != nil {
 		s.logger.Errorf("%s: %s: %s", util.GetCurrentFuncName(), util.GetFuncName(ioutil.ReadAll), err)
@@ -151,6 +161,17 @@ func (s *ReverseProxyService) getModifyResponseFunc(cfg config.StartCommandConfi
 		resp.Body = ioutil.NopCloser(bytes.NewBuffer(modifiedResponseBody))
 		resp.Header["Content-Length"] = []string{strconv.Itoa(len(modifiedResponseBody))}
 		resp.ContentLength = int64(len(modifiedResponseBody))
+
+		// Add response headers
+		for _, header := range cfg.AdditionalResponseHeaders.Value {
+			kv := strings.SplitN(header, ": ", 2)
+			if len(kv) != 2 {
+				s.logger.Errorf("%s: %s: %s - %+v", util.GetCurrentFuncName(), util.GetFuncName(strings.SplitN), "returns not 2 values", kv)
+				continue
+			}
+			resp.Header.Add(kv[0], kv[1])
+		}
+
 		return nil
 	}
 }