nano-run/server/mode_bin.go

package server

import (
	"context"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"time"

	"nano-run/server/internal"
)

type markerResponse struct {
	dataSent bool
	res      http.ResponseWriter
}

func (m *markerResponse) Header() http.Header {
	return m.res.Header()
}

func (m *markerResponse) Write(bytes []byte) (int, error) {
	m.dataSent = true
	return m.res.Write(bytes)
}

func (m *markerResponse) WriteHeader(statusCode int) {
	m.dataSent = true
	m.res.WriteHeader(statusCode)
}

type binHandler struct {
	user            string
	command         string
	workDir         string
	shell           string
	environment     []string
	timeout         time.Duration
	gracefulTimeout time.Duration
}

func (bh *binHandler) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	marker := &markerResponse{res: writer}
	ctx := request.Context()
	cmd := exec.Command(bh.shell, "-c", bh.command) //nolint:gosec

	if bh.workDir == "" {
		tmpDir, err := ioutil.TempDir("", "")
		if err != nil {
			http.Error(writer, err.Error(), http.StatusInternalServerError)
			return
		}
		defer os.RemoveAll(tmpDir)
		cmd.Dir = tmpDir
	} else {
		cmd.Dir = bh.workDir
	}

	var env = bh.cloneEnv()
	for k, v := range request.Header {
		ke := strings.ToUpper(strings.Replace(k, "-", "_", -1))
		env = append(env, ke+"="+strings.Join(v, ","))
	}

	cmd.Stderr = os.Stderr
	cmd.Stdin = request.Body
	cmd.Stdout = marker
	cmd.Env = env
	internal.SetBinFlags(cmd)
	err := setUser(cmd, bh.user)
	if err != nil {
		writer.Header().Set("X-Return-Code", strconv.Itoa(cmd.ProcessState.ExitCode()))
		writer.WriteHeader(http.StatusInternalServerError)
		return
	}
	err = bh.run(ctx, cmd)

	if codeReset, ok := writer.(interface{ Status(status int) }); ok && err != nil {
		codeReset.Status(http.StatusBadGateway)
	}

	if err != nil {
		writer.Header().Set("X-Return-Code", strconv.Itoa(cmd.ProcessState.ExitCode()))
		writer.WriteHeader(http.StatusBadGateway)
	} else {
		writer.Header().Set("X-Return-Code", strconv.Itoa(cmd.ProcessState.ExitCode()))
		writer.WriteHeader(http.StatusNoContent)
	}
}

func (bh *binHandler) run(global context.Context, cmd *exec.Cmd) error {
	err := cmd.Start()
	if err != nil {
		return err
	}

	var (
		gracefulTimeout <-chan time.Time
		ctx             context.Context
	)

	var gracefulTimer *time.Ticker
	if bh.gracefulTimeout > 0 {
		gracefulTimer = time.NewTicker(bh.gracefulTimeout)
		defer gracefulTimer.Stop()
		gracefulTimeout = gracefulTimer.C
	}

	if bh.timeout > 0 {
		t, cancel := context.WithTimeout(global, bh.timeout)
		defer cancel()
		ctx = t
	} else {
		ctx = global
	}

	var process = make(chan error, 1)

	go func() {
		defer close(process)
		process <- cmd.Wait()
	}()

	for {
		select {
		case <-gracefulTimeout:
			if err := internal.IntSignal(cmd); err != nil {
				log.Println("failed send signal to process:", err)
			} else {
				log.Println("sent graceful shutdown to process")
			}
			gracefulTimer.Stop()
			gracefulTimeout = nil
		case <-ctx.Done():
			if err := internal.KillSignal(cmd); err != nil {
				log.Println("failed send kill to process:", err)
			} else {
				log.Println("sent kill to process")
			}
			return <-process
		case err := <-process:
			return err
		}
	}
}

func (bh *binHandler) cloneEnv() []string {
	var cp = make([]string, len(bh.environment))
	copy(cp, bh.environment)
	return cp
}