add ability to stream command output as a response

This commit is contained in:
lostsnow 2021-08-04 12:49:20 +08:00
parent 2a36f24269
commit afd20c569e
3 changed files with 44 additions and 9 deletions

View file

@ -570,6 +570,7 @@ type Hook struct {
ResponseMessage string `json:"response-message,omitempty"`
ResponseHeaders ResponseHeaders `json:"response-headers,omitempty"`
CaptureCommandOutput bool `json:"include-command-output-in-response,omitempty"`
StreamCommandOutput bool `json:"stream-command-output,omitempty"`
CaptureCommandOutputOnError bool `json:"include-command-output-in-response-on-error,omitempty"`
PassEnvironmentToCommand []Argument `json:"pass-environment-to-command,omitempty"`
PassArgumentsToCommand []Argument `json:"pass-arguments-to-command,omitempty"`

View file

@ -5,6 +5,7 @@ import (
"encoding/json"
"flag"
"fmt"
"io"
"io/ioutil"
"log"
"net"
@ -63,6 +64,19 @@ var (
pidFile *pidfile.PIDFile
)
type flushWriter struct {
f http.Flusher
w io.Writer
}
func (fw *flushWriter) Write(p []byte) (n int, err error) {
n, err = fw.w.Write(p)
if fw.f != nil {
fw.f.Flush()
}
return
}
func matchLoadedHook(id string) *hook.Hook {
for _, hooks := range loadedHooksFromFiles {
if hook := hooks.Match(id); hook != nil {
@ -504,8 +518,10 @@ func hookHandler(w http.ResponseWriter, r *http.Request) {
w.Header().Set(responseHeader.Name, responseHeader.Value)
}
if matchedHook.CaptureCommandOutput {
response, err := handleHook(matchedHook, req)
if matchedHook.StreamCommandOutput {
handleHook(matchedHook, req, w)
} else if matchedHook.CaptureCommandOutput {
response, err := handleHook(matchedHook, req, nil)
if err != nil {
w.WriteHeader(http.StatusInternalServerError)
@ -523,7 +539,7 @@ func hookHandler(w http.ResponseWriter, r *http.Request) {
fmt.Fprint(w, response)
}
} else {
go handleHook(matchedHook, req)
go handleHook(matchedHook, req, nil)
// Check if a success return code is configured for the hook
if matchedHook.SuccessHttpResponseCode != 0 {
@ -546,7 +562,7 @@ func hookHandler(w http.ResponseWriter, r *http.Request) {
fmt.Fprint(w, "Hook rules were not satisfied.")
}
func handleHook(h *hook.Hook, r *hook.Request) (string, error) {
func handleHook(h *hook.Hook, r *hook.Request, w http.ResponseWriter) (string, error) {
var errors []error
// check the command exists
@ -615,13 +631,31 @@ func handleHook(h *hook.Hook, r *hook.Request) (string, error) {
log.Printf("[%s] executing %s (%s) with arguments %q and environment %s using %s as cwd\n", r.ID, h.ExecuteCommand, cmd.Path, cmd.Args, envs, cmd.Dir)
out, err := cmd.CombinedOutput()
var out []byte
if w != nil {
log.Printf("[%s] command output will be streamed to response", r.ID)
// Implementation from https://play.golang.org/p/PpbPyXbtEs
// as described in https://stackoverflow.com/questions/19292113/not-buffered-http-responsewritter-in-golang
fw := flushWriter{w: w}
if f, ok := w.(http.Flusher); ok {
fw.f = f
}
cmd.Stderr = &fw
cmd.Stdout = &fw
if err := cmd.Run(); err != nil {
log.Printf("[%s] error occurred: %+v\n", r.ID, err)
}
} else {
out, err = cmd.CombinedOutput()
log.Printf("[%s] command output: %s\n", r.ID, out)
if err != nil {
log.Printf("[%s] error occurred: %+v\n", r.ID, err)
}
}
for i := range files {
if files[i].File != nil {

View file

@ -58,7 +58,7 @@ func TestStaticParams(t *testing.T) {
ID: "test",
Headers: spHeaders,
}
_, err = handleHook(spHook, r)
_, err = handleHook(spHook, r, nil)
if err != nil {
t.Fatalf("Unexpected error: %v\n", err)
}