Adding a very simple CORS implementation (#93)

* Added a very simple CORS implementation

* Update header add location

---------

Co-authored-by: Sandeep Singh <sandeep@projectdiscovery.io>
Co-authored-by: NoF0rte <nof0rte@users.noreply.github.com>
Co-authored-by: shubhamrasal <shubhamdharmarasal@gmail.com>

Author: 3 people

internal/runner/options.go

@@ -35,6 +35,7 @@ type Options struct {
 	MaxFileSize     int
 	HTTP1Only       bool
 	MaxDumpBodySize int
+	CORS            bool
 	HTTPHeaders     HTTPHeaders
 }
 
@@ -64,6 +65,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.BoolVar(&options.CORS, "cors", false, "Enable Cross-Origin Resource Sharing (CORS)")
 	flag.Var(&options.HTTPHeaders, "header", "Add HTTP Response Header (name: value), can be used multiple times")
 	flag.Parse()
 

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),
+		CORS:              r.options.CORS,
 		HTTPHeaders:       r.options.HTTPHeaders,
 	})
 	if err != nil {

pkg/httpserver/corslayer.go

@@ -0,0 +1,28 @@
+package httpserver
+
+import (
+	"net/http"
+	"strings"
+)
+
+func (t *HTTPServer) corslayer(handler http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		headers := w.Header()
+		headers.Set("Access-Control-Allow-Origin", "*")
+		if r.Method != http.MethodOptions {
+			handler.ServeHTTP(w, r)
+			return
+		}
+
+		headers.Add("Vary", "Origin")
+		headers.Add("Vary", "Access-Control-Request-Method")
+		headers.Add("Vary", "Access-Control-Request-Headers")
+
+		reqMethod := r.Header.Get("Access-Control-Request-Method")
+		if reqMethod != "" {
+			headers.Set("Access-Control-Allow-Methods", strings.ToUpper(reqMethod))
+		}
+
+		w.WriteHeader(http.StatusOK)
+	})
+}

pkg/httpserver/httpserver.go

@@ -27,6 +27,7 @@ type Options struct {
 	HTTP1Only         bool
 	MaxFileSize       int // 50Mb
 	MaxDumpBodySize   int64
+	CORS              bool
 	HTTPHeaders       []HTTPHeader
 }
 
@@ -72,6 +73,10 @@ func New(options *Options) (*HTTPServer, error) {
 		addHandler(h.basicauthlayer)
 	}
 
+	if options.CORS {
+		addHandler(h.corslayer)
+	}
+
 	httpHandler = h.loglayer(httpHandler)
 	httpHandler = h.headerlayer(httpHandler, options.HTTPHeaders)