gort.go

package gort

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/exec"
	"strconv"

	"github.com/go-chi/chi"
	"github.com/go-chi/chi/middleware"
	"github.com/idestis/gort/utils"
)

// Script sctruct will hold an entity to define which script we should run
type Script struct {
	// Executor defines the name of
	// the interpreter program to be executed
	Executor string `json:"executor"`
	// Script will refer to script name from the dist folder
	Script string `json:"script"`
	// EnvVars slice will hold all environment variables for the script
	EnvVars []string `json:"env_vars"`
	// Args slice will hold additional arguments for the command
	Args []string `json:"args"`
}

const (
	// defaultPort if not specified as environment variable
	defaultPort = 5000
	// defaultScriptsDir is a path to find scripts
	defaultScriptsDir = "./dist"
)

var (
	port       int
	scriptsDir string
	scripts    []string
)

// init is here with one reason gort need to be initialized first
func init() {
	port, _ = strconv.Atoi(os.Getenv("PORT"))
	if port == 0 {
		port = defaultPort
	}
	scriptsDir = os.Getenv("SCRIPTS_DIR")
	if scriptsDir == "" {
		scriptsDir = defaultScriptsDir
	}
	if _, err := os.Stat(scriptsDir); os.IsNotExist(err) {
		log.Panic(err)
	}
	scripts = utils.ScanScripts(scriptsDir)
}

func main() {
	r := chi.NewRouter()

	r.Use(middleware.RealIP)
	r.Use(middleware.Logger)
	r.Use(middleware.Recoverer)
	r.Use(middleware.Compress(9))

	r.Route("/v1", func(r chi.Router) {
		if os.Getenv("GORT_RATE_LIMIT") != "" {
			rl, _ := strconv.Atoi(os.Getenv("GORT_RATE_LIMIT"))
			log.Println("GORT_RATE_LIMIT was set globally for", rl)
			r.Use(middleware.Throttle(rl))
		}
		r.Post("/start", StartScriptHandler)                            // /v1/start
		r.Get("/list-dist", ListScriptsHandler)                         // /v1/list-dist
		r.Get("/health", func(w http.ResponseWriter, r *http.Request) { // /v1/health
			w.Write([]byte("OK"))
		})
	})

	r.NotFound(NotFoundHandler)
	log.Printf("Gort is started on port %d\n", port)
	http.ListenAndServe(fmt.Sprintf(":%d", port), r)

}

// ListScriptsHandler will return scripts list from the SCRIPTS_DIR
func ListScriptsHandler(w http.ResponseWriter, r *http.Request) {
	if len(scripts) == 0 {
		fmt.Fprintf(w, "%s seems like empty", scriptsDir)
		return
	}
	for _, script := range scripts {
		fmt.Fprintln(w, script)
	}
}

// StartScriptHandler will start requested script and print output to stdout
func StartScriptHandler(w http.ResponseWriter, r *http.Request) {
	var script Script
	err := json.NewDecoder(r.Body).Decode(&script)
	if err != nil {
		http.Error(w, "Not able to parse data as valid JSON", 422)
		return
	}

	if script.Executor == "" || script.Script == "" {
		http.Error(w, "Required parameters 'executor' and 'script' were not found in the payload", 400)
		return
	}

	_, err = exec.LookPath(script.Executor)
	if err != nil {
		http.Error(w, "Requested executor is not installed", 500)
		return
	}

	_, found := utils.Find(scripts, script.Script)
	if !found {
		http.Error(w, "Requested script is not found in the scripts directory", 501)
		return
	}
	command := []string{scriptsDir + "/" + script.Script}
	cmd := exec.Command(script.Executor, command...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	// Append arguments to command if they are present in request
	if len(script.Args) >= 1 {
		command = append(command, script.Args...)
		cmd = exec.Command(script.Executor, command...)
	}

	// Append arguments to command.Environment if they are present in request
	if len(script.EnvVars) >= 1 {
		cmd.Env = os.Environ()
		for _, envVar := range script.EnvVars {
			cmd.Env = append(cmd.Env, envVar)
		}
	}
	// Start and don't wait till execution ends
	cmd.Start()
	log.Println("Just ran subprocess of", script.Script, "with PID", cmd.Process.Pid)
	fmt.Fprintf(w, "The function will be executed in the background. Refer to container logs to see the output")
}

// NotFoundHandler will return custom error message
func NotFoundHandler(w http.ResponseWriter, r *http.Request) {
	http.Error(w, "This page does not exist!", 404)
}