intake/core/action.go

package core

import (
	"bufio"
	"context"
	"database/sql/driver"
	"encoding/json"
	"errors"
	"io"
	"log"
	"os"
	"os/exec"
	"strings"
	"time"
)

// Type alias for storing string array as jsonb
type argList []string

func (a argList) Value() (driver.Value, error) {
	return json.Marshal(a)
}

func (a *argList) Scan(value interface{}) error {
	return json.Unmarshal([]byte(value.(string)), a)
}

func AddAction(db *DB, source string, name string, argv []string) error {
	_, err := db.Exec(`
		insert into actions (source, name, argv)
		values (?, ?, jsonb(?))
	`, source, name, argList(argv))
	return err
}

func UpdateAction(db *DB, source string, name string, argv []string) error {
	_, err := db.Exec(`
		update actions
		set argv = jsonb(?)
		where source = ? and name = ?
	`, argList(argv), source, name)
	return err
}

func GetActionsForSource(db *DB, source string) ([]string, error) {
	rows, err := db.Query(`
		select name
		from actions
		where source = ?
	`, source)
	if err != nil {
		return nil, err
	}
	var names []string
	for rows.Next() {
		var name string
		err = rows.Scan(&name)
		if err != nil {
			return nil, err
		}
		names = append(names, name)
	}
	return names, nil
}

func GetArgvForAction(db *DB, source string, name string) ([]string, error) {
	rows := db.QueryRow(`
		select json(argv)
		from actions
		where source = ? and name = ?
	`, source, name)
	var argv argList
	err := rows.Scan(&argv)
	if err != nil {
		return nil, err
	}
	return argv, nil
}

func DeleteAction(db *DB, source string, name string) error {
	_, err := db.Exec(`
		delete from actions
		where source = ? and name = ?
	`, source, name)
	return err
}

func readStdout(stdout io.ReadCloser, source string, items chan Item, cparse chan bool) {
	var item Item
	parseError := false
	scanout := bufio.NewScanner(stdout)
	for scanout.Scan() {
		data := scanout.Bytes()
		err := json.Unmarshal(data, &item)
		if err != nil || item.Id == "" {
			log.Printf("[%s: stdout] %s\n", source, strings.TrimSpace(string(data)))
			parseError = true
		} else {
			item.Active = true   // These fields aren't up to
			item.Created = 0     // the action to set and
			item.Source = source // shouldn't be overrideable
			log.Printf("[%s:   item] %s\n", source, item.Id)
			items <- item
		}
	}
	// Only send the parsing result at the end, to block main until stdout is drained
	cparse <- parseError
	close(items)
}

func readStderr(stderr io.ReadCloser, source string, done chan bool) {
	scanerr := bufio.NewScanner(stderr)
	for scanerr.Scan() {
		text := strings.TrimSpace(scanerr.Text())
		log.Printf("[%s: stderr] %s\n", source, text)
	}
	done <- true
}

func writeStdin(stdin io.WriteCloser, text string) {
	defer stdin.Close()
	io.WriteString(stdin, text)
}

func Execute(
	source string,
	argv []string,
	env []string,
	input string,
	timeout time.Duration,
) ([]Item, error) {
	log.Printf("Executing %v", argv)

	if len(argv) == 0 {
		return nil, errors.New("empty argv")
	}

	env = append(env, "STATE_PATH=")

	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	cmd := exec.CommandContext(ctx, argv[0], argv[1:]...)
	cmd.Env = append(os.Environ(), env...)
	cmd.WaitDelay = time.Second * 5

	// Open pipes to the command
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, err
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil, err
	}

	cout := make(chan Item)
	cparse := make(chan bool)
	cerr := make(chan bool)

	// Sink routine for items produced
	var items []Item
	go func() {
		for item := range cout {
			items = append(items, item)
		}
	}()

	// Routines handling the process i/o
	go writeStdin(stdin, input)
	go readStdout(stdout, source, cout, cparse)
	go readStderr(stderr, source, cerr)

	// Kick off the command
	err = cmd.Start()
	if err != nil {
		return nil, err
	}

	// Block until std{out,err} close
	<-cerr
	parseError := <-cparse

	err = cmd.Wait()
	if ctx.Err() == context.DeadlineExceeded {
		log.Printf("Timed out after %v\n", timeout)
		return nil, err
	} else if exiterr, ok := err.(*exec.ExitError); ok {
		log.Printf("error: %s failed with exit code %d\n", argv[0], exiterr.ExitCode())
		return nil, err
	} else if err != nil {
		log.Printf("error: %s failed with error: %s\n", argv[0], err)
		return nil, err
	}

	if parseError {
		log.Printf("error: could not parse item\n")
		return nil, errors.New("invalid JSON")
	}

	return items, nil
}