Add support for custom HTTP response headers (projectdiscovery#86)

DominicBreuker · ehsandeep · web-flow · commit 1f800063c031 · 2022-07-25T20:26:41.000+05:30
* add support for custom HTTP response headers

* update flag list in readme

* fix readme typo and adjust cases of words

Co-authored-by: Sandeep Singh &lt;sandeep@projectdiscovery.io&gt;
diff --git a/README.md b/README.md
@@ -69,6 +69,7 @@ This will display help for the tool. Here are all the switches it supports.
 | `-realm`         | Basic auth message                                      | `simplehttpserver -realm "insert the credentials"` |
 | `-version`       | Show version                                            | `simplehttpserver -version`                        |
 | `-silent`        | Show only results                                       | `simplehttpserver -silent`                         |
+| `-header`        | HTTP response header (can be used multiple times)       | `simplehttpserver -header 'X-Powered-By: Go'`      |
 
 ### Running simplehttpserver in the current folder  
 
diff --git a/internal/runner/options.go b/internal/runner/options.go
@@ -2,37 +2,40 @@ package runner
 
 import (
 	"flag"
+	"fmt"
 	"os"
 	"path/filepath"
 	"strings"
 
 	"github.com/projectdiscovery/gologger"
 	"github.com/projectdiscovery/gologger/levels"
+	"github.com/projectdiscovery/simplehttpserver/pkg/httpserver"
 )
 
 // Options of the tool
 type Options struct {
-	ListenAddress  string
-	Folder         string
-	BasicAuth      string
-	username       string
-	password       string
-	Realm          string
-	TLSCertificate string
-	TLSKey         string
-	TLSDomain      string
-	HTTPS          bool
-	Verbose        bool
-	EnableUpload   bool
-	EnableTCP      bool
-	RulesFile      string
-	TCPWithTLS     bool
-	Version        bool
-	Silent         bool
-	Sandbox        bool
-	MaxFileSize    int
-	HTTP1Only      bool
+	ListenAddress   string
+	Folder          string
+	BasicAuth       string
+	username        string
+	password        string
+	Realm           string
+	TLSCertificate  string
+	TLSKey          string
+	TLSDomain       string
+	HTTPS           bool
+	Verbose         bool
+	EnableUpload    bool
+	EnableTCP       bool
+	RulesFile       string
+	TCPWithTLS      bool
+	Version         bool
+	Silent          bool
+	Sandbox         bool
+	MaxFileSize     int
+	HTTP1Only       bool
 	MaxDumpBodySize int
+	HTTPHeaders     HTTPHeaders
 }
 
 // ParseOptions parses the command line options for application
@@ -61,6 +64,7 @@ func ParseOptions() *Options {
 	flag.BoolVar(&options.HTTP1Only, "http1", false, "Enable only HTTP1")
 	flag.IntVar(&options.MaxFileSize, "max-file-size", 50, "Max Upload File Size")
 	flag.IntVar(&options.MaxDumpBodySize, "max-dump-body-size", -1, "Max Dump Body Size")
+	flag.Var(&options.HTTPHeaders, "header", "Add HTTP Response Header (name: value), can be used multiple times")
 	flag.Parse()
 
 	// Read the inputs and configure the logging
@@ -113,3 +117,21 @@ func (options *Options) FolderAbsPath() string {
 	}
 	return abspath
 }
+
+// HTTPHeaders is a slice of HTTPHeader structs
+type HTTPHeaders []httpserver.HTTPHeader
+
+func (h *HTTPHeaders) String() string {
+	return fmt.Sprint(*h)
+}
+
+// Set sets a new header, which must be a string of the form 'name: value'
+func (h *HTTPHeaders) Set(value string) error {
+	tokens := strings.SplitN(value, ":", 2)
+	if len(tokens) != 2 {
+		return fmt.Errorf("header '%s' not in format 'name: value'", value)
+	}
+
+	*h = append(*h, httpserver.HTTPHeader{Name: tokens[0], Value: tokens[1]})
+	return nil
+}
diff --git a/internal/runner/runner.go b/internal/runner/runner.go
@@ -68,6 +68,7 @@ func New(options *Options) (*Runner, error) {
 		MaxFileSize:       r.options.MaxFileSize,
 		HTTP1Only:         r.options.HTTP1Only,
 		MaxDumpBodySize:   unit.ToMb(r.options.MaxDumpBodySize),
+		HTTPHeaders:       r.options.HTTPHeaders,
 	})
 	if err != nil {
 		return nil, err
diff --git a/pkg/httpserver/headerlayer.go b/pkg/httpserver/headerlayer.go
@@ -0,0 +1,20 @@
+package httpserver
+
+import (
+	"net/http"
+)
+
+// HTTPHeader represents an HTTP header
+type HTTPHeader struct {
+	Name  string
+	Value string
+}
+
+func (t *HTTPServer) headerlayer(handler http.Handler, headers []HTTPHeader) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		for _, header := range headers {
+			w.Header().Set(header.Name, header.Value)
+		}
+		handler.ServeHTTP(w, r)
+	})
+}
diff --git a/pkg/httpserver/httpserver.go b/pkg/httpserver/httpserver.go
@@ -27,6 +27,7 @@ type Options struct {
 	HTTP1Only         bool
 	MaxFileSize       int // 50Mb
 	MaxDumpBodySize   int64
+	HTTPHeaders       []HTTPHeader
 }
 
 // HTTPServer instance
@@ -72,6 +73,7 @@ func New(options *Options) (*HTTPServer, error) {
 	}
 
 	httpHandler = h.loglayer(httpHandler)
+	httpHandler = h.headerlayer(httpHandler, options.HTTPHeaders)
 
 	// add handler
 	h.layers = httpHandler

Original file line number	Diff line number	Diff line change
`@@ -27,6 +27,7 @@ type Options struct {`
`27`	`27`	`HTTP1Only bool`
`28`	`28`	`MaxFileSize int // 50Mb`
`29`	`29`	`MaxDumpBodySize int64`
	`30`	`+ HTTPHeaders []HTTPHeader`
`30`	`31`	`}`
`31`	`32`
`32`	`33`	`// HTTPServer instance`
`@@ -72,6 +73,7 @@ func New(options Options) (HTTPServer, error) {`
`72`	`73`	`}`
`73`	`74`
`74`	`75`	`httpHandler = h.loglayer(httpHandler)`
	`76`	`+ httpHandler = h.headerlayer(httpHandler, options.HTTPHeaders)`
`75`	`77`
`76`	`78`	`// add handler`
`77`	`79`	`h.layers = httpHandler`