Add unit test for the message body and headers

Adds unit tests for passing of the message body and the message headers.

cmd/weblug/config.go

@@ -21,7 +21,8 @@ type ConfigSettings struct {
 	GID            int         `yaml:"gid"`           // Custom group ID or 0, if not being used
 	ReadTimeout    int         `yaml:"readtimeout"`   // Timeout for reading the whole request
 	WriteTimeout   int         `yaml:"writetimeout"`  // Timeout for writing the whole response
-	MaxHeaderBytes int         `yaml:"maxheadersize"` // Maximum size of the receive body
+	MaxHeaderBytes int         `yaml:"maxheadersize"` // Maximum size of the receive header
+	MaxBodySize    int64       `yaml:"maxbodysize"`   // Maximum size of the receive body
 	TLS            TLSSettings `yaml:"tls"`
 }
 

cmd/weblug/hook.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"bytes"
 	"fmt"
 	"net"
 	"os/exec"
@@ -24,6 +25,8 @@ type Hook struct {
 	BlockedAddresses []string          `yaml:"blocked"`    // Addresses that are explicitly blocked
 	HttpBasicAuth    BasicAuth         `yaml:"basic_auth"` // Optional requires http basic auth
 	Env              map[string]string `yaml:"env"`        // Optional environment variables
+
+	maxBodySize int64 // Maximum allowed body size for this hook
 }
 
 type BasicAuth struct {
@@ -79,10 +82,12 @@ func cmdSplit(command string) []string {
 	return ret
 }
 
-// Run executes the given command and return it's return code. It also respects the given concurrency number and will block until resources are free
-func (hook *Hook) Run() error {
+// Run executes the given command and return it's return code.
+// Will pass the given input string to the command
+// It also respects the given concurrency number and will block until resources are free
+func (hook *Hook) Run(buffer []byte) ([]byte, error) {
 	if hook.Command == "" {
-		return nil
+		return make([]byte, 0), nil
 	}
 
 	split := cmdSplit(hook.Command)
@@ -102,12 +107,8 @@ func (hook *Hook) Run() error {
 			cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", k, v))
 		}
 	}
-	if hook.Output {
-		buf, ret := cmd.Output()
-		fmt.Println(string(buf))
-		return ret
-	}
-	return cmd.Run()
+	cmd.Stdin = bytes.NewReader(buffer)
+	return cmd.Output()
 }
 
 func isAddressInList(addr string, addrList []string) (bool, error) {

cmd/weblug/weblug.go

@@ -4,8 +4,10 @@
 package main
 
 import (
+	"bytes"
 	"crypto/tls"
 	"fmt"
+	"io"
 	"log"
 	"net"
 	"net/http"
@@ -224,6 +226,10 @@ func RegisterHandlers(cf Config, mux *http.ServeMux) error {
 		if hook.Concurrency < 1 {
 			hook.Concurrency = 1
 		}
+		// allow hooks to have individual maxBodySize arguments.
+		if cf.Settings.MaxBodySize > 0 && hook.maxBodySize == 0 {
+			hook.maxBodySize = cf.Settings.MaxBodySize
+		}
 		mux.HandleFunc(hook.Route, createHandler(hook))
 	}
 	return nil
@@ -299,13 +305,37 @@ func createHandler(hook Hook) Handler {
 			return
 		}
 
+		// Input buffer used to pass the headers to the process
+		var buffer bytes.Buffer
+		for k, v := range r.Header {
+			buffer.WriteString(fmt.Sprintf("%s:%s\n", k, strings.Join(v, ",")))
+		}
+		buffer.WriteString("\n")
+		// Receive body only if configured. By default the body is ignored.
+		if hook.maxBodySize > 0 {
+			body := make([]byte, hook.maxBodySize)
+			n, err := r.Body.Read(body)
+			if err != nil && err != io.EOF {
+				log.Printf("receive body failed: %s", err)
+				w.Header().Add("Content-Type", "application/json")
+				w.WriteHeader(500)
+				fmt.Fprintf(w, "{\"status\":\"fail\",\"reason\":\"receive failure\"}")
+				return
+			}
+			buffer.Write(body[:n])
+		}
+
 		if hook.Background { // Execute command in background
 			w.Header().Add("Content-Type", "application/json")
 			w.WriteHeader(200)
 			fmt.Fprintf(w, "{\"status\":\"ok\"}")
 			go func() {
 				defer hook.Unlock()
-				if err := hook.Run(); err != nil {
+				buffer, err := hook.Run(buffer.Bytes())
+				if hook.Output {
+					fmt.Println(string(buffer))
+				}
+				if err != nil {
 					log.Printf("Hook \"%s\" failed: %s", hook.Name, err)
 				} else {
 					log.Printf("Hook \"%s\" completed", hook.Name)
@@ -313,7 +343,11 @@ func createHandler(hook Hook) Handler {
 			}()
 		} else {
 			defer hook.Unlock()
-			if err := hook.Run(); err != nil {
+			buffer, err := hook.Run(buffer.Bytes())
+			if hook.Output {
+				fmt.Println(string(buffer))
+			}
+			if err != nil {
 				log.Printf("ERR: \"%s\" exec failure: %s", hook.Name, err)
 				w.Header().Add("Content-Type", "application/json")
 				w.WriteHeader(500)

cmd/weblug/weblug_test.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"bytes"
 	"context"
 	"crypto/rand"
 	"crypto/rsa"
@@ -250,6 +251,117 @@ func TestTLSWebserver(t *testing.T) {
 	}
 }
 
+// Tests the run hook commands
+func TestRunHook(t *testing.T) {
+	testText := "hello Test"
+	hook := Hook{Name: "hook", Command: "cat"}
+
+	buffer, err := hook.Run([]byte(testText))
+	if err != nil {
+		t.Fatalf("running test hook failed: %s", err)
+	}
+	ret := string(buffer)
+	if ret != testText {
+		t.Error("returned string mismatch")
+	}
+}
+
+// Tests passing the request header and body
+func TestHeaderAndBody(t *testing.T) {
+	// Create temp file
+	tempFile, err := os.CreateTemp("", "test_header_body_*")
+	if err != nil {
+		panic(err)
+	}
+	defer func() {
+		os.Remove(tempFile.Name())
+	}()
+
+	// Create test webserver with receive hook, that passes all headers and the body to the temp file
+
+	var cf Config
+
+	bodyIncluded := "this is the request body\nit is awesome"
+	bodyIgnored := "this part of the body should be ignored\nIt is hopefully not present"
+	bodyText := fmt.Sprintf("%s\n%s", bodyIncluded, bodyIgnored)
+
+	cf.Settings.BindAddress = "127.0.0.1:2088"
+	cf.Settings.MaxBodySize = int64(len(bodyIncluded))
+	cf.Hooks = make([]Hook, 0)
+	cf.Hooks = append(cf.Hooks, Hook{Name: "hook", Command: fmt.Sprintf("tee %s", tempFile.Name()), Route: "/header_and_body"})
+
+	listener, err := CreateListener(cf)
+	if err != nil {
+		t.Fatalf("error creating listener: %s", err)
+		return
+	}
+	server := CreateWebserver(cf)
+	mux := http.NewServeMux()
+	server.Handler = mux
+	if err := RegisterHandlers(cf, mux); err != nil {
+		t.Fatalf("error registering handlers: %s", err)
+		return
+	}
+	go server.Serve(listener)
+	defer func() {
+		if err := server.Shutdown(context.Background()); err != nil {
+			t.Fatalf("error while server shutdown: %s", err)
+			return
+		}
+	}()
+
+	// Create http request with custom headers and a message body
+	client := &http.Client{}
+	req, err := http.NewRequest("GET", fmt.Sprintf("http://%s/header_and_body", cf.Settings.BindAddress), nil)
+	if err != nil {
+		panic(err)
+	}
+	headers := make(map[string]string, 0)
+	headers["Header1"] = "value1"
+	headers["Header2"] = "value2"
+	headers["Header3"] = "value3"
+	headers["Content-Type"] = "this is the content type"
+	for k, v := range headers {
+		req.Header.Set(k, v)
+	}
+	req.Body = io.NopCloser(bytes.NewReader([]byte(bodyText)))
+	res, err := client.Do(req)
+	if err != nil {
+		t.Fatalf("http request error: %s", err)
+	}
+	if res.StatusCode != http.StatusOK {
+		t.Fatalf("http request failed: %d != %d", res.StatusCode, http.StatusOK)
+	}
+
+	// Assert that the headers and the body is in the test file
+	buf, err := os.ReadFile(tempFile.Name())
+	if err != nil {
+		panic(err)
+	}
+	contents := string(buf)
+	assertHeader := func(key string, value string) {
+		if !strings.Contains(contents, key) {
+			t.Fatalf("Header %s is not present", key)
+		}
+		if !strings.Contains(contents, fmt.Sprintf("%s:%s\n", key, value)) {
+			t.Fatalf("Header %s has not the right value", key)
+		}
+	}
+	for k, v := range headers {
+		assertHeader(k, v)
+	}
+
+	// Assert the message body got passed as well
+	if !strings.Contains(contents, bodyIncluded) {
+		t.Fatal("Message body was not passed")
+	}
+
+	// Assert the messaeg body got cropped
+	if strings.Contains(contents, bodyIgnored) {
+		t.Fatal("Cut-off after max-body size didn't happen")
+	}
+}
+
 func generateKeypair(keyfile string, certfile string, hostnames []string) error {
 	key, err := rsa.GenerateKey(rand.Reader, 2048)
 	if err != nil {

weblug.yml

@@ -11,6 +11,7 @@ settings:
   readtimeout: 10            # if set, maximum number of seconds to receive the full request
   writetimeout: 10           # if set, maximum number of seconds to send the full response
   maxheadersize: 4096        # maximum header size
+  maxbodysize: 0             # maximum size of the body before cropping. Setting to 0 will ignore the http request body. Default: 0
   # Enable TLS
   # Note that it is not recommended to run weblug on the open internet without using a reverse proxy
   tls: