<commit_before>package gokeepasslib

import (
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
)

//Constant enumerator for the inner random stream ID
const (
	NoStreamID    uint32 = 0
	ARC4StreamID         = 1
	SalsaStreamID        = 2
)

//Constants enumerator for compression flags
const (
	NoCompressionFlag   uint32 = 0
	GzipCompressionFlag        = 1
)

var AESCipherID = []byte{0x31, 0xC1, 0xF2, 0xE6, 0xBF, 0x71, 0x43, 0x50, 0xBE, 0x58, 0x05, 0x21, 0x6A, 0xFC, 0x5A, 0xFF}

var ErrEndOfHeaders = errors.New("gokeepasslib: header id was 0, end of headers")

type ErrUnknownHeaderID int

func (i ErrUnknownHeaderID) Error() string {
	return fmt.Sprintf("gokeepasslib: unknown header ID of %d", i)
}

// FileHeaders holds the header information of the Keepass File.
type FileHeaders struct {
	Comment             []byte // FieldID:  1
	CipherID            []byte // FieldID:  2
	CompressionFlags    uint32 // FieldID:  3
	MasterSeed          []byte // FieldID:  4
	TransformSeed       []byte // FieldID:  5
	TransformRounds     uint64 // FieldID:  6
	EncryptionIV        []byte // FieldID:  7
	ProtectedStreamKey  []byte // FieldID:  8
	StreamStartBytes    []byte // FieldID:  9
	InnerRandomStreamID uint32 // FieldID: 10
}

// NewFileHeaders creates a new FileHeaders with good defaults
func NewFileHeaders() *FileHeaders {
	masterSeed := make([]byte, 32)
	rand.Read(masterSeed)

	transformSeed := make([]byte, 32)
	rand.Read(transformSeed)

	encryptionIV := make([]byte, 16)
	rand.Read(encryptionIV)

	protectedStreamKey := make([]byte, 32)
	rand.Read(protectedStreamKey)

	streamStartBytes := make([]byte, 32)
	rand.Read(streamStartBytes)

	return &FileHeaders{
		CipherID:            []byte(AESCipherID),
		CompressionFlags:    GzipCompressionFlag,
		MasterSeed:          masterSeed,
		TransformSeed:       transformSeed,
		TransformRounds:     6000,
		EncryptionIV:        encryptionIV,
		ProtectedStreamKey:  protectedStreamKey,
		StreamStartBytes:    streamStartBytes,
		InnerRandomStreamID: SalsaStreamID,
	}
}

func (h FileHeaders) String() string {
	return fmt.Sprintf(
		"(1) Comment: %x\n"+
			"(2) CipherID: %x\n"+
			"(3) CompressionFlags: %x\n"+
			"(4) MasterSeed: %x\n"+
			"(5) TransformSeed: %x\n"+
			"(6) TransformRounds: %d\n"+
			"(7) EncryptionIV: %x\n"+
			"(8) ProtectedStreamKey: %x\n"+
			"(9) StreamStartBytes: %x\n"+
			"(10) InnerRandomStreamID: %d\n",
		h.Comment,
		h.CipherID,
		h.CompressionFlags,
		h.MasterSeed,
		h.TransformSeed,
		h.TransformRounds,
		h.EncryptionIV,
		h.ProtectedStreamKey,
		h.StreamStartBytes,
		h.InnerRandomStreamID,
	)
}

func (headers *FileHeaders) SetHeader(h Header) error {
	switch h.ID {
	case 0:
		return ErrEndOfHeaders
	case 1:
		headers.Comment = h.Data
	case 2:
		headers.CipherID = h.Data
	case 3:
		headers.CompressionFlags = binary.LittleEndian.Uint32(h.Data)
	case 4:
		headers.MasterSeed = h.Data
	case 5:
		headers.TransformSeed = h.Data
	case 6:
		headers.TransformRounds = binary.LittleEndian.Uint64(h.Data)
	case 7:
		headers.EncryptionIV = h.Data
	case 8:
		headers.ProtectedStreamKey = h.Data
	case 9:
		headers.StreamStartBytes = h.Data
	case 10:
		headers.InnerRandomStreamID = binary.LittleEndian.Uint32(h.Data)
	default:
		return ErrUnknownHeaderID(h.ID)
	}
	return nil
}

// ReadHeaders reads the headers from an io.Reader and
// creates a structure containing the parsed header information
func (h *FileHeaders) ReadFrom(r io.Reader) error {
	var header Header
	for {
		if err := header.ReadFrom(r); err != nil {
			return err
		}
		if err := h.SetHeader(header); err != nil {
			if err == ErrEndOfHeaders {
				return nil
			}
			return err
		}
	}
}

// WriteHeaders takes the contents of the corresponding FileHeaders struct
// and writes them to the given io.Writer
func (headers *FileHeaders) WriteTo(w io.Writer) error {
	var header Header
	var err error
	header = NewHeader(1,headers.Comment)
	if err = header.WriteTo(w); err != nil {
		return err
	}
	header = NewHeader(2, headers.CipherID)
	if err = header.WriteTo(w); err != nil {
		return err
	}
	header = NewHeader(3,make([]byte,4))
	binary.LittleEndian.PutUint32(header.Data, headers.CompressionFlags)
	if err = header.WriteTo(w); err != nil {
		return err
	}
	header = NewHeader(4, headers.MasterSeed)
	if err = header.WriteTo(w); err != nil {
		return err
	}
	header = NewHeader(5, headers.TransformSeed)
	if err = header.WriteTo(w); err != nil {
		return err
	}
	header = NewHeader(6,make([]byte,8))
	binary.LittleEndian.PutUint64(header.Data, headers.TransformRounds)
	if err = header.WriteTo(w); err != nil {
		return err
	}
	header = NewHeader(7, headers.EncryptionIV)
	if err = header.WriteTo(w); err != nil {
		return err
	}
	header = NewHeader(8, headers.ProtectedStreamKey)
	if err = header.WriteTo(w); err != nil {
		return err
	}
	header = NewHeader(9, headers.StreamStartBytes)
	if err = header.WriteTo(w); err != nil {
		return err
	}
	header = NewHeader(10,make([]byte,4))
	binary.LittleEndian.PutUint32(header.Data, headers.InnerRandomStreamID)
	if err = header.WriteTo(w); err != nil {
		return err
	}
	err = EndHeader.WriteTo(w)
	return err
}

var EndHeader = Header{0, 4, []byte{0x0d, 0x0a, 0x0d, 0x0a}}

type Header struct {
	ID     uint8
	Length uint16
	Data   []byte
}

// NewHeader creates a new header ,setting length automaticaly
func NewHeader(id int, data []byte) Header {
	return Header{
		ID:     uint8(id),
		Length: uint16(len(data)),
		Data:   data,
	}
}
func (h *Header) ReadFrom(r io.Reader) error {
	if err := binary.Read(r, binary.LittleEndian, &h.ID); err != nil {
		return err
	}
	if err := binary.Read(r, binary.LittleEndian, &h.Length); err != nil {
		return err
	}
	h.Data = make([]byte, h.Length)
	if err := binary.Read(r, binary.LittleEndian, h.Data); err != nil {
		return err
	}
	return nil
}
func (h Header) WriteTo(w io.Writer) error {
	if len(h.Data) > 0 {
		if err := binary.Write(w, binary.LittleEndian, h.ID); err != nil {
			return err
		}
		if err := binary.Write(w, binary.LittleEndian, h.Length); err != nil {
			return err
		}
		if err := binary.Write(w, binary.LittleEndian, h.Data); err != nil {
			return err
		}
	}
	return nil
}
func (h *Header) FixLength() {
	h.Length = uint16(len(h.Data))
}
func (h Header) String() string {
	return fmt.Sprintf("ID: %d, Length: %d, Data: %x", h.ID, h.Length, h.Data)
}
<commit_msg>fixed typo<commit_after>package gokeepasslib

import (
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
)

//Constant enumerator for the inner random stream ID
const (
	NoStreamID    uint32 = 0
	ARC4StreamID         = 1
	SalsaStreamID        = 2
)

//Constants enumerator for compression flags
const (
	NoCompressionFlag   uint32 = 0
	GzipCompressionFlag        = 1
)

var AESCipherID = []byte{0x31, 0xC1, 0xF2, 0xE6, 0xBF, 0x71, 0x43, 0x50, 0xBE, 0x58, 0x05, 0x21, 0x6A, 0xFC, 0x5A, 0xFF}

var ErrEndOfHeaders = errors.New("gokeepasslib: header id was 0, end of headers")

type ErrUnknownHeaderID int

func (i ErrUnknownHeaderID) Error() string {
	return fmt.Sprintf("gokeepasslib: unknown header ID of %d", i)
}

// FileHeaders holds the header information of the Keepass File.
type FileHeaders struct {
	Comment             []byte // FieldID:  1
	CipherID            []byte // FieldID:  2
	CompressionFlags    uint32 // FieldID:  3
	MasterSeed          []byte // FieldID:  4
	TransformSeed       []byte // FieldID:  5
	TransformRounds     uint64 // FieldID:  6
	EncryptionIV        []byte // FieldID:  7
	ProtectedStreamKey  []byte // FieldID:  8
	StreamStartBytes    []byte // FieldID:  9
	InnerRandomStreamID uint32 // FieldID: 10
}

// NewFileHeaders creates a new FileHeaders with good defaults
func NewFileHeaders() *FileHeaders {
	masterSeed := make([]byte, 32)
	rand.Read(masterSeed)

	transformSeed := make([]byte, 32)
	rand.Read(transformSeed)

	encryptionIV := make([]byte, 16)
	rand.Read(encryptionIV)

	protectedStreamKey := make([]byte, 32)
	rand.Read(protectedStreamKey)

	streamStartBytes := make([]byte, 32)
	rand.Read(streamStartBytes)

	return &FileHeaders{
		CipherID:            []byte(AESCipherID),
		CompressionFlags:    GzipCompressionFlag,
		MasterSeed:          masterSeed,
		TransformSeed:       transformSeed,
		TransformRounds:     6000,
		EncryptionIV:        encryptionIV,
		ProtectedStreamKey:  protectedStreamKey,
		StreamStartBytes:    streamStartBytes,
		InnerRandomStreamID: SalsaStreamID,
	}
}

func (h FileHeaders) String() string {
	return fmt.Sprintf(
		"(1) Comment: %x\n"+
			"(2) CipherID: %x\n"+
			"(3) CompressionFlags: %d\n"+
			"(4) MasterSeed: %x\n"+
			"(5) TransformSeed: %x\n"+
			"(6) TransformRounds: %d\n"+
			"(7) EncryptionIV: %x\n"+
			"(8) ProtectedStreamKey: %x\n"+
			"(9) StreamStartBytes: %x\n"+
			"(10) InnerRandomStreamID: %d\n",
		h.Comment,
		h.CipherID,
		h.CompressionFlags,
		h.MasterSeed,
		h.TransformSeed,
		h.TransformRounds,
		h.EncryptionIV,
		h.ProtectedStreamKey,
		h.StreamStartBytes,
		h.InnerRandomStreamID,
	)
}

func (headers *FileHeaders) SetHeader(h Header) error {
	switch h.ID {
	case 0:
		return ErrEndOfHeaders
	case 1:
		headers.Comment = h.Data
	case 2:
		headers.CipherID = h.Data
	case 3:
		headers.CompressionFlags = binary.LittleEndian.Uint32(h.Data)
	case 4:
		headers.MasterSeed = h.Data
	case 5:
		headers.TransformSeed = h.Data
	case 6:
		headers.TransformRounds = binary.LittleEndian.Uint64(h.Data)
	case 7:
		headers.EncryptionIV = h.Data
	case 8:
		headers.ProtectedStreamKey = h.Data
	case 9:
		headers.StreamStartBytes = h.Data
	case 10:
		headers.InnerRandomStreamID = binary.LittleEndian.Uint32(h.Data)
	default:
		return ErrUnknownHeaderID(h.ID)
	}
	return nil
}

// ReadHeaders reads the headers from an io.Reader and
// creates a structure containing the parsed header information
func (h *FileHeaders) ReadFrom(r io.Reader) error {
	var header Header
	for {
		if err := header.ReadFrom(r); err != nil {
			return err
		}
		if err := h.SetHeader(header); err != nil {
			if err == ErrEndOfHeaders {
				return nil
			}
			return err
		}
	}
}

// WriteTo takes the contents of the corresponding FileHeaders struct
// and writes them to the given io.Writer
func (headers *FileHeaders) WriteTo(w io.Writer) error {
	var header Header
	var err error
	header = NewHeader(1,headers.Comment)
	if err = header.WriteTo(w); err != nil {
		return err
	}
	header = NewHeader(2, headers.CipherID)
	if err = header.WriteTo(w); err != nil {
		return err
	}
	header = NewHeader(3,make([]byte,4))
	binary.LittleEndian.PutUint32(header.Data, headers.CompressionFlags)
	if err = header.WriteTo(w); err != nil {
		return err
	}
	header = NewHeader(4, headers.MasterSeed)
	if err = header.WriteTo(w); err != nil {
		return err
	}
	header = NewHeader(5, headers.TransformSeed)
	if err = header.WriteTo(w); err != nil {
		return err
	}
	header = NewHeader(6,make([]byte,8))
	binary.LittleEndian.PutUint64(header.Data, headers.TransformRounds)
	if err = header.WriteTo(w); err != nil {
		return err
	}
	header = NewHeader(7, headers.EncryptionIV)
	if err = header.WriteTo(w); err != nil {
		return err
	}
	header = NewHeader(8, headers.ProtectedStreamKey)
	if err = header.WriteTo(w); err != nil {
		return err
	}
	header = NewHeader(9, headers.StreamStartBytes)
	if err = header.WriteTo(w); err != nil {
		return err
	}
	header = NewHeader(10,make([]byte,4))
	binary.LittleEndian.PutUint32(header.Data, headers.InnerRandomStreamID)
	if err = header.WriteTo(w); err != nil {
		return err
	}
	err = EndHeader.WriteTo(w)
	return err
}

var EndHeader = Header{0, 4, []byte{0x0d, 0x0a, 0x0d, 0x0a}}

type Header struct {
	ID     uint8
	Length uint16
	Data   []byte
}

// NewHeader creates a new header ,setting length automaticaly
func NewHeader(id int, data []byte) Header {
	return Header{
		ID:     uint8(id),
		Length: uint16(len(data)),
		Data:   data,
	}
}
func (h *Header) ReadFrom(r io.Reader) error {
	if err := binary.Read(r, binary.LittleEndian, &h.ID); err != nil {
		return err
	}
	if err := binary.Read(r, binary.LittleEndian, &h.Length); err != nil {
		return err
	}
	h.Data = make([]byte, h.Length)
	if err := binary.Read(r, binary.LittleEndian, h.Data); err != nil {
		return err
	}
	return nil
}
func (h Header) WriteTo(w io.Writer) error {
	if len(h.Data) > 0 {
		if err := binary.Write(w, binary.LittleEndian, h.ID); err != nil {
			return err
		}
		if err := binary.Write(w, binary.LittleEndian, h.Length); err != nil {
			return err
		}
		if err := binary.Write(w, binary.LittleEndian, h.Data); err != nil {
			return err
		}
	}
	return nil
}
func (h *Header) FixLength() {
	h.Length = uint16(len(h.Data))
}
func (h Header) String() string {
	return fmt.Sprintf("ID: %d, Length: %d, Data: %x", h.ID, h.Length, h.Data)
}
<|endoftext|>


<commit_before>package context

import (
	"fmt"
	"net/http"
	"strconv"
)

// HttpError extends the error interface to include error information for HTTP responses
type HttpError interface {
	error
	Code() int
}

// httpError is the default, internal HttpError implementation
type httpError struct {
	code    int
	message string
}

func (err *httpError) Error() string {
	return err.message
}

func (err *httpError) String() string {
	return strconv.Itoa(err.code) + ":" + err.message
}

func (err *httpError) Code() int {
	return err.code
}

// NewHttpError creates a new HttpError object.  This object implements
// go's builtin error interface.
func NewHttpError(code int, message string) HttpError {
	return &httpError{code, message}
}

// WriteJsonError writes a standard JSON error to the response
func WriteJsonError(response http.ResponseWriter, code int, message string) error {
	response.Header().Set(ContentTypeHeader, JsonContentType)
	response.Header().Set(ContentTypeOptionsHeader, NoSniff)
	response.WriteHeader(code)
	_, err := response.Write(
		[]byte(
			fmt.Sprintf(`{"message": "%s"}`, message),
		),
	)

	return err
}

// WriteError handles writing errors, possibly from panic, in a standard way.
// This method permits a variety of types for the err value.
func WriteError(response http.ResponseWriter, err interface{}) error {
	switch value := err.(type) {
	case HttpError:
		return WriteJsonError(response, value.Code(), value.Error())

	case error:
		return WriteJsonError(response, http.StatusInternalServerError, value.Error())

	case int:
		response.Header().Set(ContentTypeOptionsHeader, NoSniff)
		response.WriteHeader(value)

	default:
		response.Header().Set(ContentTypeOptionsHeader, NoSniff)
		response.WriteHeader(http.StatusInternalServerError)
	}

	return nil
}
<commit_msg>Made things a bit more concise<commit_after>package context

import (
	"fmt"
	"net/http"
	"strconv"
)

// HttpError extends the error interface to include error information for HTTP responses
type HttpError interface {
	error
	Code() int
}

// httpError is the default, internal HttpError implementation
type httpError struct {
	code    int
	message string
}

func (err *httpError) Error() string {
	return err.message
}

func (err *httpError) String() string {
	return strconv.Itoa(err.code) + ":" + err.message
}

func (err *httpError) Code() int {
	return err.code
}

// NewHttpError creates a new HttpError object.  This object implements
// go's builtin error interface.
func NewHttpError(code int, message string) HttpError {
	return &httpError{code, message}
}

// WriteJsonError writes a standard JSON error to the response
func WriteJsonError(response http.ResponseWriter, code int, message string) error {
	response.Header().Set(ContentTypeHeader, JsonContentType)
	response.Header().Set(ContentTypeOptionsHeader, NoSniff)

	response.WriteHeader(code)
	_, err := fmt.Fprintf(response, `{"message": "%s"}`, message)
	return err
}

// WriteError handles writing errors, possibly from panic, in a standard way.
// This method permits a variety of types for the err value.
func WriteError(response http.ResponseWriter, err interface{}) error {
	switch value := err.(type) {
	case HttpError:
		return WriteJsonError(response, value.Code(), value.Error())

	case error:
		return WriteJsonError(response, http.StatusInternalServerError, value.Error())

	case int:
		response.Header().Set(ContentTypeOptionsHeader, NoSniff)
		response.WriteHeader(value)

	default:
		response.Header().Set(ContentTypeOptionsHeader, NoSniff)
		response.WriteHeader(http.StatusInternalServerError)
	}

	return nil
}
<|endoftext|>


<commit_before>package forest

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
)

// maskHeaderNames is used to prevent logging secrets.
var maskHeaderNames = []string{}

// MaskChar is used the create a masked header value.
var MaskChar = "*"

// MaskHeader is used to prevent logging secrets.
func MaskHeader(name string) {
	maskHeaderNames = append(maskHeaderNames, name)
}

// IsMaskedHeader return true if the name is part of (case-insensitive match) the MaskHeaderNames.
func IsMaskedHeader(name string) bool {
	for _, each := range maskHeaderNames {
		if strings.ToLower(each) == strings.ToLower(name) {
			return true
		}
	}
	return false
}

// Dump is a convenient method to log the full contents of a request and its response.
func Dump(t T, resp *http.Response) {
	// dump request
	var buffer bytes.Buffer
	buffer.WriteString("\n")
	buffer.WriteString(fmt.Sprintf("%v %v\n", resp.Request.Method, resp.Request.URL))
	for k, v := range resp.Request.Header {
		if IsMaskedHeader(k) {
			v = []string{strings.Repeat(MaskChar, len(v[0]))}
		}
		if len(k) > 0 {
			buffer.WriteString(fmt.Sprintf("%s : %v\n", k, strings.Join(v, ",")))
		}
	}
	// dump request payload, only available is there is a Body.
	if resp != nil && resp.Request != nil && resp.Request.Body != nil {
		rc, err := resp.Request.GetBody()
		body, err := ioutil.ReadAll(rc)
		if err != nil {
			buffer.WriteString(fmt.Sprintf("unable to read request body:%v", err))
		} else {
			if len(body) > 0 {
				buffer.WriteString("\n")
			}
			buffer.WriteString(string(body))
			buffer.WriteString("\n")
		}
	}
	// dump response payload
	if resp == nil {
		buffer.WriteString("-- no response --")
		Logf(t, buffer.String())
		return
	}
	buffer.WriteString(fmt.Sprintf("\n%s\n", resp.Status))
	for k, v := range resp.Header {
		if len(k) > 0 {
			buffer.WriteString(fmt.Sprintf("%s : %v\n", k, strings.Join(v, ",")))
		}
	}
	if resp.Body != nil {
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			if resp.StatusCode/100 == 3 {
				// redirect closes body ; nothing to read
				buffer.WriteString("\n")
			} else {
				buffer.WriteString(fmt.Sprintf("unable to read response body:%v", err))
			}
		} else {
			if len(body) > 0 {
				buffer.WriteString("\n")
			}
			buffer.WriteString(string(body))
		}
		resp.Body.Close()
		// put the body back for re-reads
		resp.Body = ioutil.NopCloser(bytes.NewReader(body))
	}
	buffer.WriteString("\n")
	Logf(t, buffer.String())
}

type skippeable interface {
	Skipf(string, ...interface{})
}

// SkipUnless will Skip the test unless the LABELS environment variable includes any of the provided labels.
//
//	LABELS=integration,nightly go test -v
//
func SkipUnless(t skippeable, labels ...string) {
	env := strings.Split(os.Getenv("LABELS"), ",")
	for _, each := range labels {
		for _, other := range env {
			if each == other {
				return
			}
		}
	}
	t.Skipf("skipped because provided LABELS=%v does not include any of %v", env, labels)
}

func headersString(h http.Header) string {
	if len(h) == 0 {
		return ""
	}
	masked := http.Header{}
	for k, v := range h {
		if IsMaskedHeader(k) {
			v = []string{strings.Repeat(MaskChar, len(v[0]))}
		}
		masked[k] = v
	}
	return fmt.Sprintf("%v", masked)
}

func copyHeaders(from, to http.Header) {
	for k, list := range from {
		for _, v := range list {
			to.Set(k, v)
		}
	}
}

func setFormData(config *RequestConfig, req *http.Request) {
	// set form data if available
	if len(config.FormData) > 0 {
		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
		req.Body = ioutil.NopCloser(strings.NewReader(config.FormData.Encode()))
	}
}

func setBasicAuth(config *RequestConfig, req *http.Request) {
	if len(config.User) > 0 {
		req.SetBasicAuth(config.User, config.Password)
	}
}

func ensureResponse(req *http.Request, resp *http.Response) *http.Response {
	if resp != nil {
		return resp
	}
	// wrap the request into an empty response ; Body must be nil
	wrap := new(http.Response)
	wrap.Request = req
	return wrap
}

func URLPathEncode(path string) string {
	buf := new(bytes.Buffer)
	for _, each := range path {
		switch each {
		case '+':
			buf.WriteString("%20")
		case ' ':
			buf.WriteString("%20")
		default:
			buf.WriteRune(each)
		}
	}
	return buf.String()
}
<commit_msg>fixes dump response that has Go format markers<commit_after>package forest

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
)

// maskHeaderNames is used to prevent logging secrets.
var maskHeaderNames = []string{}

// MaskChar is used the create a masked header value.
var MaskChar = "*"

// MaskHeader is used to prevent logging secrets.
func MaskHeader(name string) {
	maskHeaderNames = append(maskHeaderNames, name)
}

// IsMaskedHeader return true if the name is part of (case-insensitive match) the MaskHeaderNames.
func IsMaskedHeader(name string) bool {
	for _, each := range maskHeaderNames {
		if strings.ToLower(each) == strings.ToLower(name) {
			return true
		}
	}
	return false
}

// Dump is a convenient method to log the full contents of a request and its response.
func Dump(t T, resp *http.Response) {
	// dump request
	var buffer bytes.Buffer
	buffer.WriteString("\n")
	buffer.WriteString(fmt.Sprintf("%v %v\n", resp.Request.Method, resp.Request.URL))
	for k, v := range resp.Request.Header {
		if IsMaskedHeader(k) {
			v = []string{strings.Repeat(MaskChar, len(v[0]))}
		}
		if len(k) > 0 {
			buffer.WriteString(fmt.Sprintf("%s : %v\n", k, strings.Join(v, ",")))
		}
	}
	// dump request payload, only available is there is a Body.
	if resp != nil && resp.Request != nil && resp.Request.Body != nil {
		rc, err := resp.Request.GetBody()
		body, err := ioutil.ReadAll(rc)
		if err != nil {
			buffer.WriteString(fmt.Sprintf("unable to read request body:%v", err))
		} else {
			if len(body) > 0 {
				buffer.WriteString("\n")
			}
			buffer.WriteString(string(body))
			buffer.WriteString("\n")
		}
	}
	// dump response payload
	if resp == nil {
		buffer.WriteString("-- no response --")
		Logf(t, buffer.String())
		return
	}
	buffer.WriteString(fmt.Sprintf("\n%s\n", resp.Status))
	for k, v := range resp.Header {
		if len(k) > 0 {
			buffer.WriteString(fmt.Sprintf("%s : %v\n", k, strings.Join(v, ",")))
		}
	}
	if resp.Body != nil {
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			if resp.StatusCode/100 == 3 {
				// redirect closes body ; nothing to read
				buffer.WriteString("\n")
			} else {
				buffer.WriteString(fmt.Sprintf("unable to read response body:%v", err))
			}
		} else {
			if len(body) > 0 {
				buffer.WriteString("\n")
			}
			buffer.WriteString(string(body))
		}
		resp.Body.Close()
		// put the body back for re-reads
		resp.Body = ioutil.NopCloser(bytes.NewReader(body))
	}
	buffer.WriteString("\n")
	Logf(t, "%s", buffer.String())
}

type skippeable interface {
	Skipf(string, ...interface{})
}

// SkipUnless will Skip the test unless the LABELS environment variable includes any of the provided labels.
//
//	LABELS=integration,nightly go test -v
//
func SkipUnless(t skippeable, labels ...string) {
	env := strings.Split(os.Getenv("LABELS"), ",")
	for _, each := range labels {
		for _, other := range env {
			if each == other {
				return
			}
		}
	}
	t.Skipf("skipped because provided LABELS=%v does not include any of %v", env, labels)
}

func headersString(h http.Header) string {
	if len(h) == 0 {
		return ""
	}
	masked := http.Header{}
	for k, v := range h {
		if IsMaskedHeader(k) {
			v = []string{strings.Repeat(MaskChar, len(v[0]))}
		}
		masked[k] = v
	}
	return fmt.Sprintf("%v", masked)
}

func copyHeaders(from, to http.Header) {
	for k, list := range from {
		for _, v := range list {
			to.Set(k, v)
		}
	}
}

func setFormData(config *RequestConfig, req *http.Request) {
	// set form data if available
	if len(config.FormData) > 0 {
		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
		req.Body = ioutil.NopCloser(strings.NewReader(config.FormData.Encode()))
	}
}

func setBasicAuth(config *RequestConfig, req *http.Request) {
	if len(config.User) > 0 {
		req.SetBasicAuth(config.User, config.Password)
	}
}

func ensureResponse(req *http.Request, resp *http.Response) *http.Response {
	if resp != nil {
		return resp
	}
	// wrap the request into an empty response ; Body must be nil
	wrap := new(http.Response)
	wrap.Request = req
	return wrap
}

func URLPathEncode(path string) string {
	buf := new(bytes.Buffer)
	for _, each := range path {
		switch each {
		case '+':
			buf.WriteString("%20")
		case ' ':
			buf.WriteString("%20")
		default:
			buf.WriteRune(each)
		}
	}
	return buf.String()
}
<|endoftext|>


<commit_before>package test

import (
	"testing"

	"golang.org/x/net/websocket"

	"github.com/neptulon/neptulon"
	"github.com/neptulon/neptulon/middleware"
	"github.com/neptulon/randstr"
)

type echoMsg struct {
	Message string `json:"message"`
}

var (
	msg1 = "Lorem ipsum dolor sit amet, consectetur adipiscing elit."
	msg2 = "In sit amet lectus felis, at pellentesque turpis."
	msg3 = "Nunc urna enim, cursus varius aliquet ac, imperdiet eget tellus."
	msg4 = randstr.Get(45 * 1000)       // 0.45 MB
	msg5 = randstr.Get(5 * 1000 * 1000) // 5.0 MB
)

func TestEchoWithoutTestHelpers(t *testing.T) {
	s := neptulon.NewServer("127.0.0.1:3001")
	go s.Start()
	defer s.Close()

	s.Middleware(func(ctx *neptulon.ReqCtx) error {
		t.Log("Request received:", ctx.Method)
		ctx.Res = "response-wow!"
		return ctx.Next()
	})

	origin := "http://127.0.0.1"
	url := "ws://127.0.0.1:3001"
	ws, err := websocket.Dial(url, "", origin)
	if err != nil {
		t.Fatal(err)
	}
	if err := websocket.JSON.Send(ws, map[string]string{"id": "123", "method": "test"}); err != nil {
		t.Fatal(err)
	}
	var res interface{}
	if err := websocket.JSON.Receive(ws, &res); err != nil {
		t.Fatal(err)
	}
	t.Log("Got response:", res)

	if err := ws.Close(); err != nil {
		t.Fatal(err)
	}
}

func TestEcho(t *testing.T) {
	sh := NewServerHelper(t)
	rout := middleware.NewRouter()
	sh.Middleware(middleware.Logger)
	sh.Middleware(rout.Middleware)
	rout.Request("echo", middleware.Echo)
	defer sh.Start().CloseWait()

	ch := sh.GetConnHelper()
	ch.Middleware(middleware.Logger)
	defer ch.Connect().CloseWait()

	m := "Hello!"
	ch.SendRequest("echo", echoMsg{Message: m}, func(ctx *neptulon.ResCtx) error {
		var msg echoMsg
		if err := ctx.Result(&msg); err != nil {
			t.Fatal(err)
		}
		if msg.Message != m {
			t.Fatalf("expected: %v got: %v", m, msg.Message)
		}
		return nil
	})
}

func TestMessages(t *testing.T) {
	// todo: verify all message echoes from small to big
}

func TestBidirectional(t *testing.T) {
	// todo: test simultaneous read/writes
}

func TestTLS(t *testing.T) {
	// todo: ...
}
<commit_msg>wait for server to start in manual test<commit_after>package test

import (
	"testing"
	"time"

	"golang.org/x/net/websocket"

	"github.com/neptulon/neptulon"
	"github.com/neptulon/neptulon/middleware"
	"github.com/neptulon/randstr"
)

type echoMsg struct {
	Message string `json:"message"`
}

var (
	msg1 = "Lorem ipsum dolor sit amet, consectetur adipiscing elit."
	msg2 = "In sit amet lectus felis, at pellentesque turpis."
	msg3 = "Nunc urna enim, cursus varius aliquet ac, imperdiet eget tellus."
	msg4 = randstr.Get(45 * 1000)       // 0.45 MB
	msg5 = randstr.Get(5 * 1000 * 1000) // 5.0 MB
)

func TestEchoWithoutTestHelpers(t *testing.T) {
	s := neptulon.NewServer("127.0.0.1:3001")
	go s.Start()
	time.Sleep(time.Millisecond * 10)
	defer s.Close()

	s.Middleware(func(ctx *neptulon.ReqCtx) error {
		t.Log("Request received:", ctx.Method)
		ctx.Res = "response-wow!"
		return ctx.Next()
	})

	origin := "http://127.0.0.1"
	url := "ws://127.0.0.1:3001"
	ws, err := websocket.Dial(url, "", origin)
	if err != nil {
		t.Fatal(err)
	}
	if err := websocket.JSON.Send(ws, map[string]string{"id": "123", "method": "test"}); err != nil {
		t.Fatal(err)
	}
	var res interface{}
	if err := websocket.JSON.Receive(ws, &res); err != nil {
		t.Fatal(err)
	}
	t.Log("Got response:", res)

	if err := ws.Close(); err != nil {
		t.Fatal(err)
	}
}

func TestEcho(t *testing.T) {
	sh := NewServerHelper(t)
	rout := middleware.NewRouter()
	sh.Middleware(middleware.Logger)
	sh.Middleware(rout.Middleware)
	rout.Request("echo", middleware.Echo)
	defer sh.Start().CloseWait()

	ch := sh.GetConnHelper()
	ch.Middleware(middleware.Logger)
	defer ch.Connect().CloseWait()

	m := "Hello!"
	ch.SendRequest("echo", echoMsg{Message: m}, func(ctx *neptulon.ResCtx) error {
		var msg echoMsg
		if err := ctx.Result(&msg); err != nil {
			t.Fatal(err)
		}
		if msg.Message != m {
			t.Fatalf("expected: %v got: %v", m, msg.Message)
		}
		return nil
	})
}

func TestMessages(t *testing.T) {
	// todo: verify all message echoes from small to big
}

func TestBidirectional(t *testing.T) {
	// todo: test simultaneous read/writes
}

func TestTLS(t *testing.T) {
	// todo: ...
}
<|endoftext|>


<commit_before>package controller

import (
	"sync"
	"sync/atomic"
	"time"

	"golang.org/x/net/context"

	"github.com/giantswarm/inago/task"
)

// executeTaskAction executes the given function f to create a new task and block until
// it finishes. If the task fails, an error is raised.
func (c controller) executeTaskAction(f func(ctx context.Context, req Request) (*task.Task, error), ctx context.Context, req Request) error {
	taskObject, err := f(ctx, req)
	if err != nil {
		return maskAny(err)
	}
	closer := make(<-chan struct{})
	taskObject, err = c.WaitForTask(ctx, taskObject.ID, closer)
	if err != nil {
		return maskAny(err)
	}
	if task.HasFailedStatus(taskObject) {
		return maskAny(taskObject.Error)
	}
	return nil
}

func (c controller) getNumRunningSlices(ctx context.Context, req Request) (int, error) {
	c.Config.Logger.Debug(ctx, "controller: getting number of running slices")

	groupStatus, err := c.groupStatus(ctx, req)
	if IsUnitNotFound(err) {
		return 0, nil
	} else if err != nil {
		return 0, maskAny(err)
	}
	grouped, err := UnitStatusList(groupStatus).Group()
	if err != nil {
		return 0, maskAny(err)
	}

	c.Config.Logger.Debug(ctx, "controller: grouped list: %#v", grouped)

	var sliceIDs []string
	for _, us := range grouped {
		groupedStatuses := grouped.unitStatusesBySliceID(us.SliceID)
		if len(groupedStatuses) > 1 {
			// This group has an inconsistent state. Thus we do not consider it
			// running.
			continue
		}

		aggregator := Aggregator{
			Logger: c.Config.Logger,
		}
		ok, err := aggregator.UnitHasStatus(groupedStatuses[0], StatusRunning)
		if err != nil {
			return 0, maskAny(err)
		}
		if !ok {
			continue
		}

		sliceIDs = append(sliceIDs, us.SliceID)
	}

	c.Config.Logger.Debug(ctx, "controller: found %v running slices", len(sliceIDs))

	return len(sliceIDs), nil
}

func (c controller) isGroupRemovalAllowed(ctx context.Context, req Request, minAlive int, removeInProgress *int64) (bool, error) {
	c.Config.Logger.Debug(ctx, "controller: checking group removal allowed, req: %v", req)

	c.Config.Logger.Debug(
		ctx, "controller: removeInProgress: %v, minAlive: %v",
		*removeInProgress, minAlive,
	)

	if (minAlive - int(*removeInProgress)) > 0 {
		c.Config.Logger.Debug(ctx, "controller: group removal allowed ((minAlive - int(removeInProgress)) > 0)")
		return true, nil
	}

	c.Config.Logger.Debug(ctx, "controller: group removal not allowed ((minAlive - int(removeInProgress)) <= 0)")
	return false, nil
}

func (c controller) isGroupAdditionAllowed(ctx context.Context, req Request, maxGrowth int, additionInProgress *int64) (bool, error) {
	c.Config.Logger.Debug(ctx, "controller: checking group addition allowed, req: %v", req)

	c.Config.Logger.Debug(
		ctx, "controller: additionInProgress: %v, maxGrowth: %v",
		*additionInProgress, maxGrowth,
	)

	if (maxGrowth - int(*additionInProgress)) > 0 {
		c.Config.Logger.Debug(ctx, "controller: group addition allowed ((maxGrowth  - additionInProgress) > 0)")
		return true, nil
	}

	c.Config.Logger.Debug(ctx, "controller: group addition not allowed ((maxGrowth  - additionInProgress) <= 0)")
	return false, nil
}

func (c controller) addFirst(ctx context.Context, req Request, opts UpdateOptions) ([]string, error) {
	c.Config.Logger.Debug(ctx, "controller: running addFirst")

	c.Config.Logger.Debug(ctx, "controller: running add worker")
	newReq, err := c.runAddWorker(ctx, req, opts)
	if err != nil {
		return nil, maskAny(err)
	}

	c.Config.Logger.Debug(ctx, "controller: checking number of running slices")
	n, err := c.getNumRunningSlices(ctx, newReq)
	if err != nil {
		return nil, maskAny(err)
	}
	if n != len(newReq.SliceIDs) {
		return nil, maskAnyf(updateFailedError, "addFirst: slice not running: %d != %v", n, newReq.SliceIDs)
	}

	c.Config.Logger.Debug(ctx, "controller: running remove worker")
	err = c.runRemoveWorker(ctx, req)
	if err != nil {
		return nil, maskAny(err)
	}

	return newReq.SliceIDs, nil
}

func (c controller) runAddWorker(ctx context.Context, req Request, opts UpdateOptions) (Request, error) {
	c.Config.Logger.Info(ctx, "controller: adding units")

	// Create new random IDs.
	req.DesiredSlices = 1
	req.SliceIDs = nil
	newReq, err := c.ExtendWithRandomSliceIDs(ctx, req)
	if err != nil {
		return Request{}, maskAny(err)
	}

	// Submit.
	if err := c.executeTaskAction(c.Submit, ctx, newReq); err != nil {
		return Request{}, maskAny(err)
	}

	// Start.
	if err := c.executeTaskAction(c.Start, ctx, newReq); err != nil {
		return Request{}, maskAny(err)
	}

	time.Sleep(time.Duration(opts.ReadySecs) * time.Second)

	return newReq, nil
}

func (c controller) removeFirst(ctx context.Context, req Request, opts UpdateOptions) ([]string, error) {
	c.Config.Logger.Debug(ctx, "controller: running removeFirst")

	c.Config.Logger.Debug(ctx, "controller: running remove worker")
	err := c.runRemoveWorker(ctx, req)
	if err != nil {
		return nil, maskAny(err)
	}

	c.Config.Logger.Debug(ctx, "controller: running add worker")
	newReq, err := c.runAddWorker(ctx, req, opts)
	if err != nil {
		return nil, maskAny(err)
	}

	c.Config.Logger.Debug(ctx, "controller: checking number of running slices")
	n, err := c.getNumRunningSlices(ctx, newReq)
	if err != nil {
		return nil, maskAny(err)
	}
	if n != len(newReq.SliceIDs) {
		return nil, maskAnyf(updateFailedError, "removeFirst: slice not running: %d != %v", n, newReq.SliceIDs)
	}

	return newReq.SliceIDs, nil
}

func (c controller) runRemoveWorker(ctx context.Context, req Request) error {
	c.Config.Logger.Info(ctx, "controller: removing units")
	c.Config.Logger.Debug(ctx, "controller: executing stop action, req: %v", req)
	// Stop.
	if err := c.executeTaskAction(c.Stop, ctx, req); err != nil {
		return maskAny(err)
	}

	c.Config.Logger.Debug(ctx, "controller: executing destroy action, req: %v", req)
	// Destroy.
	if err := c.executeTaskAction(c.Destroy, ctx, req); err != nil {
		return maskAny(err)
	}
	return nil
}

// UpdateOptions represents the options defining the strategy of an update
// process. Lets have a look at how the update process of 3 group slices would
// look like using the given options.
//
//     MaxGrowth    1
//     MinAlive     2
//     ReadySecs    30
//
//     @1 (running)  ->  @1 (stopped/destroyed)
//     @2 (running)  ->  @2 (running)            ->  @2 (stopped/destroyed)
//     @3 (running)  ->  @3 (running)            ->  @3 (running)            ->  @3 (stopped/destroyed)
//                   ->  @1 (submitted/running)  ->  @1 (running)            ->  @1 (running)            ->  @1 (running)
//                                               ->  @2 (submitted/running)  ->  @2 (running)            ->  @2 (running)
//                                                                           ->  @3 (submitted/running)  ->  @3 (running)
//
type UpdateOptions struct {
	// MaxGrowth represents the number of groups allowed to be added at a given
	// time. No more than MaxGrowth groups will be added at the same time during
	// the update process.
	MaxGrowth int

	// MinAlive represents the number of groups required to stay healthy during
	// the update process. No more than MinAlive groups will be removed at the
	// same time during the update process.
	MinAlive int

	// ReadySecs represents the number of seconds required to wait before ending
	// the update process of one group and starting the update process of another
	// group. This is basically a cool down where the update process sleeps
	// before updating the next group.
	ReadySecs int
}

// updateCurrentSliceIDs updates the list of current slice IDs,
// removing the slice that was modified, and adding any new slice IDs.
func (c controller) updateCurrentSliceIDs(ctx context.Context, currentSliceIDs []string, modifiedSliceIDs []string, newSliceIDs []string) []string {
	c.Config.Logger.Debug(ctx, "current slice IDs: %v", currentSliceIDs)
	c.Config.Logger.Debug(ctx, "modified slice IDs: %v", modifiedSliceIDs)
	c.Config.Logger.Debug(ctx, "new slice IDs: %v", newSliceIDs)

	// Remove any slice IDs that have been modified.
	for i, currentSliceID := range currentSliceIDs {
		for _, modifiedSliceID := range modifiedSliceIDs {
			if currentSliceID == modifiedSliceID {
				currentSliceIDs = append(
					currentSliceIDs[:i],
					currentSliceIDs[i+1:]...,
				)
			}
		}
	}

	// And add the new slice IDs.
	currentSliceIDs = append(currentSliceIDs, newSliceIDs...)

	c.Config.Logger.Debug(ctx, "updated slice IDs: %v", currentSliceIDs)

	return currentSliceIDs
}

func (c controller) UpdateWithStrategy(ctx context.Context, req Request, opts UpdateOptions) error {
	c.Config.Logger.Debug(ctx, "controller: running update for group '%v'", req.Group)

	fail := make(chan error, 1)
	numTotal := len(req.SliceIDs)

	done := make(chan struct{}, numTotal)

	var addInProgress, removeInProgress int64

	c.Config.Logger.Debug(ctx, "controller: checking if request is sliceable")
	if !req.isSliceable() {
		return maskAnyf(updateNotAllowedError, "cannot update unsliceable group")
	}

	c.Config.Logger.Debug(ctx, "controller: checking for slice ids in request")
	for _, sliceID := range req.SliceIDs {
		if sliceID == "" {
			return maskAnyf(updateNotAllowedError, "group misses slice ID")
		}
	}

	if numTotal < opts.MinAlive {
		return maskAnyf(updateNotAllowedError, "invalid min alive option")
	}

	// We need to track which slice IDs are currently in use.
	// This list is updated as slices are added and removed.
	currentSliceIDsMutex := sync.Mutex{}
	currentSliceIDs := []string{}
	for _, id := range req.SliceIDs {
		currentSliceIDs = append(currentSliceIDs, id)
	}

	for _, sliceID := range req.SliceIDs {
		newReq := req
		newReq.SliceIDs = []string{sliceID}

		for {
			c.Config.Logger.Debug(ctx, "controller: attempting to add slice: %v", sliceID)

			currentSliceReq := req
			// We try to add slices first. We are only allowed to increase to number
			// slices if the MaxGrowth value minus the additions that are currently
			// in progress is greater than zero.
			// => ((opts.MaxGrowth  - additionInProgress) > 0)
			// See also isGroupAdditionAllowed.
			c.Config.Logger.Debug(
				ctx, "controller: opts.MaxGrowth: %v, numTotal: %v, opts.MinAlive: %v, addInProgress: %v",
				opts.MaxGrowth, numTotal, opts.MinAlive, int(addInProgress),
			)

			c.Config.Logger.Debug(ctx, "controller: currentSliceIDs: %v", currentSliceIDs)
			currentSliceReq.SliceIDs = currentSliceIDs
			ok, err := c.isGroupAdditionAllowed(ctx, currentSliceReq, opts.MaxGrowth, &addInProgress)
			if err != nil {
				return maskAny(err)
			}
			if ok {
				ctx = context.WithValue(ctx, "slice ID", sliceID)
				// we increase the addInProgress counter before starting the goroutine
				// to avoid a race condition in the allowed calculation
				atomic.AddInt64(&addInProgress, 1)
				go func(ctx context.Context) {
					ctx = context.WithValue(ctx, "add slice", sliceID)
					c.Config.Logger.Debug(ctx, "controller: starting to add slice: %v", sliceID)

					newSliceIDs, err := c.addFirst(ctx, newReq, opts)
					if err != nil {
						fail <- maskAny(err)
						return
					}

					currentSliceIDsMutex.Lock()
					currentSliceIDs = c.updateCurrentSliceIDs(ctx, currentSliceIDs, newReq.SliceIDs, newSliceIDs)
					currentSliceIDsMutex.Unlock()

					atomic.AddInt64(&addInProgress, -1)
					done <- struct{}{}
				}(ctx)

				break
			}

			c.Config.Logger.Debug(ctx, "controller: attempting to remove slice: %v", sliceID)
			// remove
			// we are only allowed to remove if the number of minAlive slices
			// minus the ones, that are beeing removed right now is greater than 0
			//=> (minAlive - int(removeInProgress)) > 0)
			c.Config.Logger.Debug(
				ctx, "controller: opts.MinAlive: %v, removeInProgress: %v",
				opts.MinAlive, int(removeInProgress),
			)

			c.Config.Logger.Debug(ctx, "controller: currentSliceIDs: %v", currentSliceIDs)
			currentSliceReq.SliceIDs = currentSliceIDs
			ok, err = c.isGroupRemovalAllowed(ctx, currentSliceReq, opts.MinAlive, &removeInProgress)
			if err != nil {
				return maskAny(err)
			}
			if ok {
				// we increase the removeInProgress counter before starting the goroutine
				// to avoid a race condition in the allowed calculation
				atomic.AddInt64(&removeInProgress, 1)
				go func(ctx context.Context) {
					ctx = context.WithValue(ctx, "remove slice", sliceID)
					c.Config.Logger.Debug(ctx, "controller: starting to remove slice: %v", sliceID)

					newSliceIDs, err := c.removeFirst(ctx, newReq, opts)
					if err != nil {
						fail <- maskAny(err)
						return
					}

					currentSliceIDsMutex.Lock()
					currentSliceIDs = c.updateCurrentSliceIDs(ctx, currentSliceIDs, newReq.SliceIDs, newSliceIDs)
					currentSliceIDsMutex.Unlock()

					atomic.AddInt64(&removeInProgress, -1)
					done <- struct{}{}
				}(ctx)

				break
			}

			time.Sleep(c.WaitSleep)
		}
	}

	tc := 0

	for {
		select {
		case err := <-fail:
			return maskAny(err)
		case <-done:
			tc++
			if tc == numTotal {
				return nil
			}
		case <-time.After(c.WaitTimeout):
			return maskAny(waitTimeoutReachedError)
		}
	}
}
<commit_msg>show correct sliceID in logs when running goroutines (#243)<commit_after>package controller

import (
	"sync"
	"sync/atomic"
	"time"

	"golang.org/x/net/context"

	"github.com/giantswarm/inago/task"
)

// executeTaskAction executes the given function f to create a new task and block until
// it finishes. If the task fails, an error is raised.
func (c controller) executeTaskAction(f func(ctx context.Context, req Request) (*task.Task, error), ctx context.Context, req Request) error {
	taskObject, err := f(ctx, req)
	if err != nil {
		return maskAny(err)
	}
	closer := make(<-chan struct{})
	taskObject, err = c.WaitForTask(ctx, taskObject.ID, closer)
	if err != nil {
		return maskAny(err)
	}
	if task.HasFailedStatus(taskObject) {
		return maskAny(taskObject.Error)
	}
	return nil
}

func (c controller) getNumRunningSlices(ctx context.Context, req Request) (int, error) {
	c.Config.Logger.Debug(ctx, "controller: getting number of running slices")

	groupStatus, err := c.groupStatus(ctx, req)
	if IsUnitNotFound(err) {
		return 0, nil
	} else if err != nil {
		return 0, maskAny(err)
	}
	grouped, err := UnitStatusList(groupStatus).Group()
	if err != nil {
		return 0, maskAny(err)
	}

	c.Config.Logger.Debug(ctx, "controller: grouped list: %#v", grouped)

	var sliceIDs []string
	for _, us := range grouped {
		groupedStatuses := grouped.unitStatusesBySliceID(us.SliceID)
		if len(groupedStatuses) > 1 {
			// This group has an inconsistent state. Thus we do not consider it
			// running.
			continue
		}

		aggregator := Aggregator{
			Logger: c.Config.Logger,
		}
		ok, err := aggregator.UnitHasStatus(groupedStatuses[0], StatusRunning)
		if err != nil {
			return 0, maskAny(err)
		}
		if !ok {
			continue
		}

		sliceIDs = append(sliceIDs, us.SliceID)
	}

	c.Config.Logger.Debug(ctx, "controller: found %v running slices", len(sliceIDs))

	return len(sliceIDs), nil
}

func (c controller) isGroupRemovalAllowed(ctx context.Context, req Request, minAlive int, removeInProgress *int64) (bool, error) {
	c.Config.Logger.Debug(ctx, "controller: checking group removal allowed, req: %v", req)

	c.Config.Logger.Debug(
		ctx, "controller: removeInProgress: %v, minAlive: %v",
		*removeInProgress, minAlive,
	)

	if (minAlive - int(*removeInProgress)) > 0 {
		c.Config.Logger.Debug(ctx, "controller: group removal allowed ((minAlive - int(removeInProgress)) > 0)")
		return true, nil
	}

	c.Config.Logger.Debug(ctx, "controller: group removal not allowed ((minAlive - int(removeInProgress)) <= 0)")
	return false, nil
}

func (c controller) isGroupAdditionAllowed(ctx context.Context, req Request, maxGrowth int, additionInProgress *int64) (bool, error) {
	c.Config.Logger.Debug(ctx, "controller: checking group addition allowed, req: %v", req)

	c.Config.Logger.Debug(
		ctx, "controller: additionInProgress: %v, maxGrowth: %v",
		*additionInProgress, maxGrowth,
	)

	if (maxGrowth - int(*additionInProgress)) > 0 {
		c.Config.Logger.Debug(ctx, "controller: group addition allowed ((maxGrowth  - additionInProgress) > 0)")
		return true, nil
	}

	c.Config.Logger.Debug(ctx, "controller: group addition not allowed ((maxGrowth  - additionInProgress) <= 0)")
	return false, nil
}

func (c controller) addFirst(ctx context.Context, req Request, opts UpdateOptions) ([]string, error) {
	c.Config.Logger.Debug(ctx, "controller: running addFirst")

	c.Config.Logger.Debug(ctx, "controller: running add worker")
	newReq, err := c.runAddWorker(ctx, req, opts)
	if err != nil {
		return nil, maskAny(err)
	}

	c.Config.Logger.Debug(ctx, "controller: checking number of running slices")
	n, err := c.getNumRunningSlices(ctx, newReq)
	if err != nil {
		return nil, maskAny(err)
	}
	if n != len(newReq.SliceIDs) {
		return nil, maskAnyf(updateFailedError, "addFirst: slice not running: %d != %v", n, newReq.SliceIDs)
	}

	c.Config.Logger.Debug(ctx, "controller: running remove worker")
	err = c.runRemoveWorker(ctx, req)
	if err != nil {
		return nil, maskAny(err)
	}

	return newReq.SliceIDs, nil
}

func (c controller) runAddWorker(ctx context.Context, req Request, opts UpdateOptions) (Request, error) {
	c.Config.Logger.Info(ctx, "controller: adding units")

	// Create new random IDs.
	req.DesiredSlices = 1
	req.SliceIDs = nil
	newReq, err := c.ExtendWithRandomSliceIDs(ctx, req)
	if err != nil {
		return Request{}, maskAny(err)
	}

	// Submit.
	if err := c.executeTaskAction(c.Submit, ctx, newReq); err != nil {
		return Request{}, maskAny(err)
	}

	// Start.
	if err := c.executeTaskAction(c.Start, ctx, newReq); err != nil {
		return Request{}, maskAny(err)
	}

	time.Sleep(time.Duration(opts.ReadySecs) * time.Second)

	return newReq, nil
}

func (c controller) removeFirst(ctx context.Context, req Request, opts UpdateOptions) ([]string, error) {
	c.Config.Logger.Debug(ctx, "controller: running removeFirst")

	c.Config.Logger.Debug(ctx, "controller: running remove worker")
	err := c.runRemoveWorker(ctx, req)
	if err != nil {
		return nil, maskAny(err)
	}

	c.Config.Logger.Debug(ctx, "controller: running add worker")
	newReq, err := c.runAddWorker(ctx, req, opts)
	if err != nil {
		return nil, maskAny(err)
	}

	c.Config.Logger.Debug(ctx, "controller: checking number of running slices")
	n, err := c.getNumRunningSlices(ctx, newReq)
	if err != nil {
		return nil, maskAny(err)
	}
	if n != len(newReq.SliceIDs) {
		return nil, maskAnyf(updateFailedError, "removeFirst: slice not running: %d != %v", n, newReq.SliceIDs)
	}

	return newReq.SliceIDs, nil
}

func (c controller) runRemoveWorker(ctx context.Context, req Request) error {
	c.Config.Logger.Info(ctx, "controller: removing units")
	c.Config.Logger.Debug(ctx, "controller: executing stop action, req: %v", req)
	// Stop.
	if err := c.executeTaskAction(c.Stop, ctx, req); err != nil {
		return maskAny(err)
	}

	c.Config.Logger.Debug(ctx, "controller: executing destroy action, req: %v", req)
	// Destroy.
	if err := c.executeTaskAction(c.Destroy, ctx, req); err != nil {
		return maskAny(err)
	}
	return nil
}

// UpdateOptions represents the options defining the strategy of an update
// process. Lets have a look at how the update process of 3 group slices would
// look like using the given options.
//
//     MaxGrowth    1
//     MinAlive     2
//     ReadySecs    30
//
//     @1 (running)  ->  @1 (stopped/destroyed)
//     @2 (running)  ->  @2 (running)            ->  @2 (stopped/destroyed)
//     @3 (running)  ->  @3 (running)            ->  @3 (running)            ->  @3 (stopped/destroyed)
//                   ->  @1 (submitted/running)  ->  @1 (running)            ->  @1 (running)            ->  @1 (running)
//                                               ->  @2 (submitted/running)  ->  @2 (running)            ->  @2 (running)
//                                                                           ->  @3 (submitted/running)  ->  @3 (running)
//
type UpdateOptions struct {
	// MaxGrowth represents the number of groups allowed to be added at a given
	// time. No more than MaxGrowth groups will be added at the same time during
	// the update process.
	MaxGrowth int

	// MinAlive represents the number of groups required to stay healthy during
	// the update process. No more than MinAlive groups will be removed at the
	// same time during the update process.
	MinAlive int

	// ReadySecs represents the number of seconds required to wait before ending
	// the update process of one group and starting the update process of another
	// group. This is basically a cool down where the update process sleeps
	// before updating the next group.
	ReadySecs int
}

// updateCurrentSliceIDs updates the list of current slice IDs,
// removing the slice that was modified, and adding any new slice IDs.
func (c controller) updateCurrentSliceIDs(ctx context.Context, currentSliceIDs []string, modifiedSliceIDs []string, newSliceIDs []string) []string {
	c.Config.Logger.Debug(ctx, "current slice IDs: %v", currentSliceIDs)
	c.Config.Logger.Debug(ctx, "modified slice IDs: %v", modifiedSliceIDs)
	c.Config.Logger.Debug(ctx, "new slice IDs: %v", newSliceIDs)

	// Remove any slice IDs that have been modified.
	for i, currentSliceID := range currentSliceIDs {
		for _, modifiedSliceID := range modifiedSliceIDs {
			if currentSliceID == modifiedSliceID {
				currentSliceIDs = append(
					currentSliceIDs[:i],
					currentSliceIDs[i+1:]...,
				)
			}
		}
	}

	// And add the new slice IDs.
	currentSliceIDs = append(currentSliceIDs, newSliceIDs...)

	c.Config.Logger.Debug(ctx, "updated slice IDs: %v", currentSliceIDs)

	return currentSliceIDs
}

func (c controller) UpdateWithStrategy(ctx context.Context, req Request, opts UpdateOptions) error {
	c.Config.Logger.Debug(ctx, "controller: running update for group '%v'", req.Group)

	fail := make(chan error, 1)
	numTotal := len(req.SliceIDs)

	done := make(chan struct{}, numTotal)

	var addInProgress, removeInProgress int64

	c.Config.Logger.Debug(ctx, "controller: checking if request is sliceable")
	if !req.isSliceable() {
		return maskAnyf(updateNotAllowedError, "cannot update unsliceable group")
	}

	c.Config.Logger.Debug(ctx, "controller: checking for slice ids in request")
	for _, sliceID := range req.SliceIDs {
		if sliceID == "" {
			return maskAnyf(updateNotAllowedError, "group misses slice ID")
		}
	}

	if numTotal < opts.MinAlive {
		return maskAnyf(updateNotAllowedError, "invalid min alive option")
	}

	// We need to track which slice IDs are currently in use.
	// This list is updated as slices are added and removed.
	currentSliceIDsMutex := sync.Mutex{}
	currentSliceIDs := []string{}
	for _, id := range req.SliceIDs {
		currentSliceIDs = append(currentSliceIDs, id)
	}

	for _, sliceID := range req.SliceIDs {
		newReq := req
		newReq.SliceIDs = []string{sliceID}

		for {
			c.Config.Logger.Debug(ctx, "controller: attempting to add slice: %v", sliceID)

			currentSliceReq := req
			// We try to add slices first. We are only allowed to increase to number
			// slices if the MaxGrowth value minus the additions that are currently
			// in progress is greater than zero.
			// => ((opts.MaxGrowth  - additionInProgress) > 0)
			// See also isGroupAdditionAllowed.
			c.Config.Logger.Debug(
				ctx, "controller: opts.MaxGrowth: %v, numTotal: %v, opts.MinAlive: %v, addInProgress: %v",
				opts.MaxGrowth, numTotal, opts.MinAlive, int(addInProgress),
			)

			c.Config.Logger.Debug(ctx, "controller: currentSliceIDs: %v", currentSliceIDs)
			currentSliceReq.SliceIDs = currentSliceIDs
			ok, err := c.isGroupAdditionAllowed(ctx, currentSliceReq, opts.MaxGrowth, &addInProgress)
			if err != nil {
				return maskAny(err)
			}
			if ok {
				ctx = context.WithValue(ctx, "slice ID", sliceID)
				// we increase the addInProgress counter before starting the goroutine
				// to avoid a race condition in the allowed calculation
				atomic.AddInt64(&addInProgress, 1)
				go func(req Request) {
					ctx := context.WithValue(ctx, "add slice", req.SliceIDs)
					c.Config.Logger.Debug(ctx, "controller: starting to add slice: %v", req.SliceIDs)

					newSliceIDs, err := c.addFirst(ctx, req, opts)
					if err != nil {
						fail <- maskAny(err)
						return
					}

					currentSliceIDsMutex.Lock()
					currentSliceIDs = c.updateCurrentSliceIDs(ctx, currentSliceIDs, req.SliceIDs, newSliceIDs)
					currentSliceIDsMutex.Unlock()

					atomic.AddInt64(&addInProgress, -1)
					done <- struct{}{}
				}(newReq)

				break
			}

			c.Config.Logger.Debug(ctx, "controller: attempting to remove slice: %v", sliceID)
			// remove
			// we are only allowed to remove if the number of minAlive slices
			// minus the ones, that are beeing removed right now is greater than 0
			//=> (minAlive - int(removeInProgress)) > 0)
			c.Config.Logger.Debug(
				ctx, "controller: opts.MinAlive: %v, removeInProgress: %v",
				opts.MinAlive, int(removeInProgress),
			)

			c.Config.Logger.Debug(ctx, "controller: currentSliceIDs: %v", currentSliceIDs)
			currentSliceReq.SliceIDs = currentSliceIDs
			ok, err = c.isGroupRemovalAllowed(ctx, currentSliceReq, opts.MinAlive, &removeInProgress)
			if err != nil {
				return maskAny(err)
			}
			if ok {
				// we increase the removeInProgress counter before starting the goroutine
				// to avoid a race condition in the allowed calculation
				atomic.AddInt64(&removeInProgress, 1)
				go func(req Request) {
					ctx := context.WithValue(ctx, "remove slice", req.SliceIDs)
					c.Config.Logger.Debug(ctx, "controller: starting to remove slice: %v", req.SliceIDs)

					newSliceIDs, err := c.removeFirst(ctx, req, opts)
					if err != nil {
						fail <- maskAny(err)
						return
					}

					currentSliceIDsMutex.Lock()
					currentSliceIDs = c.updateCurrentSliceIDs(ctx, currentSliceIDs, req.SliceIDs, newSliceIDs)
					currentSliceIDsMutex.Unlock()

					atomic.AddInt64(&removeInProgress, -1)
					done <- struct{}{}
				}(newReq)

				break
			}

			time.Sleep(c.WaitSleep)
		}
	}

	tc := 0

	for {
		select {
		case err := <-fail:
			return maskAny(err)
		case <-done:
			tc++
			if tc == numTotal {
				return nil
			}
		case <-time.After(c.WaitTimeout):
			return maskAny(waitTimeoutReachedError)
		}
	}
}
<|endoftext|>


<commit_before>package worker

import (
	"encoding/json"
	"fmt"
	"time"

	gocontext "context"

	"github.com/Jeffail/tunny"
	"github.com/bitly/go-simplejson"
	"github.com/sirupsen/logrus"
	"github.com/streadway/amqp"
	"github.com/travis-ci/worker/backend"
	"github.com/travis-ci/worker/context"
	"github.com/travis-ci/worker/metrics"
)

// AMQPJobQueue is a JobQueue that uses AMQP
type AMQPJobQueue struct {
	conn  *amqp.Connection
	queue string

	stateUpdatePool *tunny.Pool

	DefaultLanguage, DefaultDist, DefaultGroup, DefaultOS string
}

// NewAMQPJobQueue creates a AMQPJobQueue backed by the given AMQP connections and
// connects to the AMQP queue with the given name. The queue will be declared
// in AMQP when this function is called, so an error could be raised if the
// queue already exists, but with different attributes than we expect.
func NewAMQPJobQueue(conn *amqp.Connection, queue string, stateUpdatePoolSize int) (*AMQPJobQueue, error) {
	channel, err := conn.Channel()
	if err != nil {
		return nil, err
	}

	_, err = channel.QueueDeclare(queue, true, false, false, false, nil)
	if err != nil {
		return nil, err
	}

	_, err = channel.QueueDeclare("reporting.jobs.builds", true, false, false, false, nil)
	if err != nil {
		return nil, err
	}

	err = channel.ExchangeDeclare("reporting", "topic", true, false, false, false, nil)
	if err != nil {
		return nil, err
	}

	_, err = channel.QueueDeclare("reporting.jobs.logs", true, false, false, false, nil)
	if err != nil {
		return nil, err
	}

	err = channel.QueueBind("reporting.jobs.logs", "reporting.jobs.logs", "reporting", false, nil)
	if err != nil {
		return nil, err
	}

	err = channel.Close()
	if err != nil {
		return nil, err
	}

	stateUpdatePool := newStateUpdatePool(conn, stateUpdatePoolSize)

	go reportPoolMetrics("state_update_pool", stateUpdatePool)

	return &AMQPJobQueue{
		conn:  conn,
		queue: queue,

		stateUpdatePool: stateUpdatePool,
	}, nil
}

func newStateUpdatePool(conn *amqp.Connection, poolSize int) *tunny.Pool {
	return tunny.New(poolSize, func() tunny.Worker {
		stateUpdateChan, err := conn.Channel()
		if err != nil {
			logrus.WithField("err", err).Panic("could not create state update amqp channel")
		}
		return &amqpStateUpdateWorker{
			stateUpdateChan: stateUpdateChan,
		}
	})
}

func reportPoolMetrics(poolName string, pool *tunny.Pool) {
	for {
		metrics.Gauge(fmt.Sprintf("travis.worker.%s.queue_length", poolName), pool.QueueLength())
		metrics.Gauge(fmt.Sprintf("travis.worker.%s.pool_size", poolName), pool.GetSize())
		time.Sleep(10 * time.Second)
	}
}

// Jobs creates a new consumer on the queue, and returns three channels. The
// first channel gets sent every BuildJob that we receive from AMQP. The
// stopChan is a channel that can be closed in order to stop the consumer.
func (q *AMQPJobQueue) Jobs(ctx gocontext.Context) (outChan <-chan Job, err error) {
	jobsChannel, err := q.conn.Channel()
	if err != nil {
		return
	}

	err = jobsChannel.Qos(1, 0, false)
	if err != nil {
		return
	}

	deliveries, err := jobsChannel.Consume(q.queue, "build-job-consumer", false, false, false, false, nil)
	if err != nil {
		return
	}

	logWriterChannel, err := q.conn.Channel()
	if err != nil {
		return
	}

	buildJobChan := make(chan Job)
	outChan = buildJobChan

	go func() {
		defer jobsChannel.Close()
		defer logWriterChannel.Close()
		defer close(buildJobChan)

		logger := context.LoggerFromContext(ctx).WithFields(logrus.Fields{
			"self": "amqp_job_queue",
			"inst": fmt.Sprintf("%p", q),
		})

		for {
			if ctx.Err() != nil {
				return
			}

			select {
			case <-ctx.Done():
				return
			case <-time.After(time.Second):
				continue
			case delivery, ok := <-deliveries:
				if !ok {
					logger.Info("job queue channel closed")
					return
				}

				buildJob := &amqpJob{
					payload:         &JobPayload{},
					startAttributes: &backend.StartAttributes{},
					stateUpdatePool: q.stateUpdatePool,
				}
				startAttrs := &jobPayloadStartAttrs{Config: &backend.StartAttributes{}}

				err := json.Unmarshal(delivery.Body, buildJob.payload)
				if err != nil {
					logger.WithField("err", err).Error("payload JSON parse error, attempting to ack+drop delivery")
					err := delivery.Ack(false)
					if err != nil {
						logger.WithField("err", err).WithField("delivery", delivery).Error("couldn't ack+drop delivery")
					}
					continue
				}

				logger.WithField("job_id", buildJob.payload.Job.ID).Info("received amqp delivery")

				err = json.Unmarshal(delivery.Body, &startAttrs)
				if err != nil {
					logger.WithField("err", err).Error("start attributes JSON parse error, attempting to ack+drop delivery")
					err := delivery.Ack(false)
					if err != nil {
						logger.WithField("err", err).WithField("delivery", delivery).Error("couldn't ack+drop delivery")
					}
					continue
				}

				buildJob.rawPayload, err = simplejson.NewJson(delivery.Body)
				if err != nil {
					logger.WithField("err", err).Error("raw payload JSON parse error, attempting to ack+drop delivery")
					err := delivery.Ack(false)
					if err != nil {
						logger.WithField("err", err).WithField("delivery", delivery).Error("couldn't ack+drop delivery")
					}
					continue
				}

				buildJob.startAttributes = startAttrs.Config
				buildJob.startAttributes.VMType = buildJob.payload.VMType
				buildJob.startAttributes.SetDefaults(q.DefaultLanguage, q.DefaultDist, q.DefaultGroup, q.DefaultOS, VMTypeDefault)
				buildJob.conn = q.conn
				buildJob.logWriterChan = logWriterChannel
				buildJob.delivery = delivery
				buildJob.stateCount = buildJob.payload.Meta.StateUpdateCount

				jobSendBegin := time.Now()
				select {
				case buildJobChan <- buildJob:
					metrics.TimeSince("travis.worker.job_queue.amqp.blocking_time", jobSendBegin)
					logger.WithFields(logrus.Fields{
						"source":           "amqp",
						"send_duration_ms": time.Since(jobSendBegin).Seconds() * 1e3,
					}).Info("sent job to output channel")
				case <-ctx.Done():
					delivery.Nack(false, true)
					return
				}
			}
		}
	}()

	return
}

// Name returns the name of this queue type, wow!
func (q *AMQPJobQueue) Name() string {
	return "amqp"
}

// Cleanup closes the underlying AMQP connection
func (q *AMQPJobQueue) Cleanup() error {
	q.stateUpdatePool.Close()
	return q.conn.Close()
}
<commit_msg>whoopsie<commit_after>package worker

import (
	"encoding/json"
	"fmt"
	"time"

	gocontext "context"

	"github.com/Jeffail/tunny"
	"github.com/bitly/go-simplejson"
	"github.com/sirupsen/logrus"
	"github.com/streadway/amqp"
	"github.com/travis-ci/worker/backend"
	"github.com/travis-ci/worker/context"
	"github.com/travis-ci/worker/metrics"
)

// AMQPJobQueue is a JobQueue that uses AMQP
type AMQPJobQueue struct {
	conn  *amqp.Connection
	queue string

	stateUpdatePool *tunny.Pool

	DefaultLanguage, DefaultDist, DefaultGroup, DefaultOS string
}

// NewAMQPJobQueue creates a AMQPJobQueue backed by the given AMQP connections and
// connects to the AMQP queue with the given name. The queue will be declared
// in AMQP when this function is called, so an error could be raised if the
// queue already exists, but with different attributes than we expect.
func NewAMQPJobQueue(conn *amqp.Connection, queue string, stateUpdatePoolSize int) (*AMQPJobQueue, error) {
	channel, err := conn.Channel()
	if err != nil {
		return nil, err
	}

	_, err = channel.QueueDeclare(queue, true, false, false, false, nil)
	if err != nil {
		return nil, err
	}

	_, err = channel.QueueDeclare("reporting.jobs.builds", true, false, false, false, nil)
	if err != nil {
		return nil, err
	}

	err = channel.ExchangeDeclare("reporting", "topic", true, false, false, false, nil)
	if err != nil {
		return nil, err
	}

	_, err = channel.QueueDeclare("reporting.jobs.logs", true, false, false, false, nil)
	if err != nil {
		return nil, err
	}

	err = channel.QueueBind("reporting.jobs.logs", "reporting.jobs.logs", "reporting", false, nil)
	if err != nil {
		return nil, err
	}

	err = channel.Close()
	if err != nil {
		return nil, err
	}

	stateUpdatePool := newStateUpdatePool(conn, stateUpdatePoolSize)

	go reportPoolMetrics("state_update_pool", stateUpdatePool)

	return &AMQPJobQueue{
		conn:  conn,
		queue: queue,

		stateUpdatePool: stateUpdatePool,
	}, nil
}

func newStateUpdatePool(conn *amqp.Connection, poolSize int) *tunny.Pool {
	return tunny.New(poolSize, func() tunny.Worker {
		stateUpdateChan, err := conn.Channel()
		if err != nil {
			logrus.WithField("err", err).Panic("could not create state update amqp channel")
		}
		return &amqpStateUpdateWorker{
			stateUpdateChan: stateUpdateChan,
		}
	})
}

func reportPoolMetrics(poolName string, pool *tunny.Pool) {
	for {
		metrics.Gauge(fmt.Sprintf("travis.worker.%s.queue_length", poolName), pool.QueueLength())
		metrics.Gauge(fmt.Sprintf("travis.worker.%s.pool_size", poolName), int64(pool.GetSize()))
		time.Sleep(10 * time.Second)
	}
}

// Jobs creates a new consumer on the queue, and returns three channels. The
// first channel gets sent every BuildJob that we receive from AMQP. The
// stopChan is a channel that can be closed in order to stop the consumer.
func (q *AMQPJobQueue) Jobs(ctx gocontext.Context) (outChan <-chan Job, err error) {
	jobsChannel, err := q.conn.Channel()
	if err != nil {
		return
	}

	err = jobsChannel.Qos(1, 0, false)
	if err != nil {
		return
	}

	deliveries, err := jobsChannel.Consume(q.queue, "build-job-consumer", false, false, false, false, nil)
	if err != nil {
		return
	}

	logWriterChannel, err := q.conn.Channel()
	if err != nil {
		return
	}

	buildJobChan := make(chan Job)
	outChan = buildJobChan

	go func() {
		defer jobsChannel.Close()
		defer logWriterChannel.Close()
		defer close(buildJobChan)

		logger := context.LoggerFromContext(ctx).WithFields(logrus.Fields{
			"self": "amqp_job_queue",
			"inst": fmt.Sprintf("%p", q),
		})

		for {
			if ctx.Err() != nil {
				return
			}

			select {
			case <-ctx.Done():
				return
			case <-time.After(time.Second):
				continue
			case delivery, ok := <-deliveries:
				if !ok {
					logger.Info("job queue channel closed")
					return
				}

				buildJob := &amqpJob{
					payload:         &JobPayload{},
					startAttributes: &backend.StartAttributes{},
					stateUpdatePool: q.stateUpdatePool,
				}
				startAttrs := &jobPayloadStartAttrs{Config: &backend.StartAttributes{}}

				err := json.Unmarshal(delivery.Body, buildJob.payload)
				if err != nil {
					logger.WithField("err", err).Error("payload JSON parse error, attempting to ack+drop delivery")
					err := delivery.Ack(false)
					if err != nil {
						logger.WithField("err", err).WithField("delivery", delivery).Error("couldn't ack+drop delivery")
					}
					continue
				}

				logger.WithField("job_id", buildJob.payload.Job.ID).Info("received amqp delivery")

				err = json.Unmarshal(delivery.Body, &startAttrs)
				if err != nil {
					logger.WithField("err", err).Error("start attributes JSON parse error, attempting to ack+drop delivery")
					err := delivery.Ack(false)
					if err != nil {
						logger.WithField("err", err).WithField("delivery", delivery).Error("couldn't ack+drop delivery")
					}
					continue
				}

				buildJob.rawPayload, err = simplejson.NewJson(delivery.Body)
				if err != nil {
					logger.WithField("err", err).Error("raw payload JSON parse error, attempting to ack+drop delivery")
					err := delivery.Ack(false)
					if err != nil {
						logger.WithField("err", err).WithField("delivery", delivery).Error("couldn't ack+drop delivery")
					}
					continue
				}

				buildJob.startAttributes = startAttrs.Config
				buildJob.startAttributes.VMType = buildJob.payload.VMType
				buildJob.startAttributes.SetDefaults(q.DefaultLanguage, q.DefaultDist, q.DefaultGroup, q.DefaultOS, VMTypeDefault)
				buildJob.conn = q.conn
				buildJob.logWriterChan = logWriterChannel
				buildJob.delivery = delivery
				buildJob.stateCount = buildJob.payload.Meta.StateUpdateCount

				jobSendBegin := time.Now()
				select {
				case buildJobChan <- buildJob:
					metrics.TimeSince("travis.worker.job_queue.amqp.blocking_time", jobSendBegin)
					logger.WithFields(logrus.Fields{
						"source":           "amqp",
						"send_duration_ms": time.Since(jobSendBegin).Seconds() * 1e3,
					}).Info("sent job to output channel")
				case <-ctx.Done():
					delivery.Nack(false, true)
					return
				}
			}
		}
	}()

	return
}

// Name returns the name of this queue type, wow!
func (q *AMQPJobQueue) Name() string {
	return "amqp"
}

// Cleanup closes the underlying AMQP connection
func (q *AMQPJobQueue) Cleanup() error {
	q.stateUpdatePool.Close()
	return q.conn.Close()
}
<|endoftext|>


<commit_before>package controllers

import (
	"fmt"
	"github.com/gin-gonic/gin"
	"net/http"

	"github.com/eirka/eirka-libs/audit"
	"github.com/eirka/eirka-libs/config"
	e "github.com/eirka/eirka-libs/errors"
	"github.com/eirka/eirka-libs/redis"
	"github.com/eirka/eirka-libs/user"

	"github.com/eirka/eirka-post/models"
	u "github.com/eirka/eirka-post/utils"
)

// Input from new reply form
type replyForm struct {
	Key     string `form:"askey" binding:"required"`
	Comment string `form:"comment"`
	Thread  uint   `form:"thread" binding:"required"`
}

// ReplyController handles the creation of new threads
func ReplyController(c *gin.Context) {
	var err error
	var rf replyForm
	req := c.Request

	// get userdata from session middleware
	userdata := c.MustGet("userdata").(user.User)

	// check size of content
	if req.ContentLength > int64(config.Settings.Limits.ImageMaxSize) {
		c.JSON(http.StatusExpectationFailed, gin.H{"error_message": e.ErrImageSize.Error()})
		c.Error(e.ErrImageSize)
		return
	}

	// set max bytes reader
	req.Body = http.MaxBytesReader(c.Writer, req.Body, int64(config.Settings.Limits.ImageMaxSize))

	err = c.Bind(&rf)
	if err != nil {
		c.JSON(e.ErrorMessage(e.ErrInvalidParam))
		c.Error(err)
		return
	}

	// Test for antispam key from Prim
	antispam := rf.Key
	if antispam != config.Settings.Antispam.AntispamKey {
		c.JSON(http.StatusBadRequest, gin.H{"error_message": e.ErrInvalidKey.Error()})
		c.Error(e.ErrInvalidKey)
		return
	}

	// Set parameters to ReplyModel
	m := models.ReplyModel{
		Uid:     userdata.Id,
		Ip:      c.ClientIP(),
		Comment: rf.Comment,
		Thread:  rf.Thread,
		Image:   true,
	}

	image := u.ImageType{}

	// Check if theres a file
	image.File, image.Header, err = req.FormFile("file")
	if err == http.ErrMissingFile {
		m.Image = false
	}

	// Validate input parameters
	err = m.ValidateInput()
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error_message": err.Error()})
		c.Error(err)
		return
	}

	// Check thread status
	err = m.Status()
	if err == e.ErrThreadClosed {
		c.JSON(http.StatusBadRequest, gin.H{"error_message": err.Error()})
		c.Error(err)
		return
	} else if err != nil {
		c.JSON(e.ErrorMessage(e.ErrInternalError))
		c.Error(err)
		return
	}

	if m.Comment != "" {

		// Check comment in SFS and Akismet
		check := u.CheckComment{
			Ip:      m.Ip,
			Ua:      req.UserAgent(),
			Referer: req.Referer(),
			Comment: m.Comment,
		}

		err = check.Get()
		if err != nil {
			c.JSON(http.StatusBadRequest, gin.H{"error_message": err.Error()})
			c.Error(err)
			return
		}

	}

	if m.Image {

		// process the uploaded file, this creates an md5
		err = image.ProcessFile()
		if err != nil {
			c.JSON(http.StatusBadRequest, gin.H{"error_message": err.Error()})
			c.Error(err)
			return
		}

		// Set MD5 from results
		m.MD5 = image.MD5

		// Initialize check duplicate
		duplicate := u.CheckDuplicate{
			Ib:  m.Ib,
			MD5: m.MD5,
		}

		// Check database for duplicate image hashes
		err = duplicate.Get()
		if err == e.ErrDuplicateImage {
			c.JSON(http.StatusBadRequest, gin.H{"error_message": err.Error(), "thread": duplicate.Thread, "post": duplicate.Post})
			c.Error(err)
			return
		} else if err != nil {
			c.JSON(e.ErrorMessage(e.ErrInternalError))
			c.Error(err)
			return
		}

		if image.Ext == ".webm" {

			// Save the webm to a file
			err = image.SaveWebM()
			if err != nil {
				c.JSON(http.StatusBadRequest, gin.H{"error_message": err.Error()})
				c.Error(err)
				return
			}

		} else {

			// Save the image to a file
			err = image.SaveImage()
			if err != nil {
				c.JSON(http.StatusBadRequest, gin.H{"error_message": err.Error()})
				c.Error(err)
				return
			}

		}

		m.OrigWidth = image.OrigWidth
		m.OrigHeight = image.OrigHeight
		m.ThumbWidth = image.ThumbWidth
		m.ThumbHeight = image.ThumbHeight
		m.Filename = image.Filename
		m.Thumbnail = image.Thumbnail

	}

	// Post data
	err = m.Post()
	if err != nil {
		c.JSON(e.ErrorMessage(e.ErrInternalError))
		c.Error(err)
		return
	}

	// Initialize cache handle
	cache := redis.RedisCache

	// Delete redis stuff
	index_key := fmt.Sprintf("%s:%d", "index", m.Ib)
	directory_key := fmt.Sprintf("%s:%d", "directory", m.Ib)
	thread_key := fmt.Sprintf("%s:%d:%d", "thread", m.Ib, m.Thread)
	image_key := fmt.Sprintf("%s:%d", "image", m.Ib)

	err = cache.Delete(index_key, directory_key, thread_key, image_key)
	if err != nil {
		c.JSON(e.ErrorMessage(e.ErrInternalError))
		c.Error(err)
		return
	}

	// get board domain and redirect to it
	redirect, err := u.Link(m.Ib, req.Referer())
	if err != nil {
		c.JSON(e.ErrorMessage(e.ErrInternalError))
		c.Error(err)
		return
	}

	c.Redirect(303, redirect)

	audit := audit.Audit{
		User:   userdata.Id,
		Ib:     m.Ib,
		Ip:     m.Ip,
		Action: audit.AuditReply,
		Info:   fmt.Sprintf("%d/%d", m.Thread, m.PostNum),
	}

	err = audit.Submit()
	if err != nil {
		c.Error(err)
	}

	return

}
<commit_msg>all database inserts should have struct data validated<commit_after>package controllers

import (
	"fmt"
	"github.com/gin-gonic/gin"
	"net/http"

	"github.com/eirka/eirka-libs/audit"
	"github.com/eirka/eirka-libs/config"
	e "github.com/eirka/eirka-libs/errors"
	"github.com/eirka/eirka-libs/redis"
	"github.com/eirka/eirka-libs/user"

	"github.com/eirka/eirka-post/models"
	u "github.com/eirka/eirka-post/utils"
)

// Input from new reply form
type replyForm struct {
	Key     string `form:"askey" binding:"required"`
	Comment string `form:"comment"`
	Thread  uint   `form:"thread" binding:"required"`
}

// ReplyController handles the creation of new threads
func ReplyController(c *gin.Context) {
	var err error
	var rf replyForm
	req := c.Request

	// get userdata from session middleware
	userdata := c.MustGet("userdata").(user.User)

	// check size of content
	if req.ContentLength > int64(config.Settings.Limits.ImageMaxSize) {
		c.JSON(http.StatusExpectationFailed, gin.H{"error_message": e.ErrImageSize.Error()})
		c.Error(e.ErrImageSize)
		return
	}

	// set max bytes reader
	req.Body = http.MaxBytesReader(c.Writer, req.Body, int64(config.Settings.Limits.ImageMaxSize))

	err = c.Bind(&rf)
	if err != nil {
		c.JSON(e.ErrorMessage(e.ErrInvalidParam))
		c.Error(err)
		return
	}

	// Test for antispam key from Prim
	antispam := rf.Key
	if antispam != config.Settings.Antispam.AntispamKey {
		c.JSON(http.StatusBadRequest, gin.H{"error_message": e.ErrInvalidKey.Error()})
		c.Error(e.ErrInvalidKey)
		return
	}

	// Set parameters to ReplyModel
	m := models.ReplyModel{
		Uid:     userdata.Id,
		Ip:      c.ClientIP(),
		Comment: rf.Comment,
		Thread:  rf.Thread,
		Image:   true,
	}

	image := u.ImageType{}

	// Check if theres a file
	image.File, image.Header, err = req.FormFile("file")
	if err == http.ErrMissingFile {
		m.Image = false
	}

	// Validate input parameters
	err = m.ValidateInput()
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error_message": err.Error()})
		c.Error(err)
		return
	}

	// Check thread status
	err = m.Status()
	if err == e.ErrThreadClosed {
		c.JSON(http.StatusBadRequest, gin.H{"error_message": err.Error()})
		c.Error(err)
		return
	} else if err != nil {
		c.JSON(e.ErrorMessage(e.ErrInternalError))
		c.Error(err)
		return
	}

	if m.Comment != "" {

		// Check comment in SFS and Akismet
		check := u.CheckComment{
			Ip:      m.Ip,
			Ua:      req.UserAgent(),
			Referer: req.Referer(),
			Comment: m.Comment,
		}

		err = check.Get()
		if err != nil {
			c.JSON(http.StatusBadRequest, gin.H{"error_message": err.Error()})
			c.Error(err)
			return
		}

	}

	if m.Image {

		// process the uploaded file, this creates an md5
		err = image.ProcessFile()
		if err != nil {
			c.JSON(http.StatusBadRequest, gin.H{"error_message": err.Error()})
			c.Error(err)
			return
		}

		// Set MD5 from results
		m.MD5 = image.MD5

		// Initialize check duplicate
		duplicate := u.CheckDuplicate{
			Ib:  m.Ib,
			MD5: m.MD5,
		}

		// Check database for duplicate image hashes
		err = duplicate.Get()
		if err == e.ErrDuplicateImage {
			c.JSON(http.StatusBadRequest, gin.H{"error_message": err.Error(), "thread": duplicate.Thread, "post": duplicate.Post})
			c.Error(err)
			return
		} else if err != nil {
			c.JSON(e.ErrorMessage(e.ErrInternalError))
			c.Error(err)
			return
		}

		if image.Ext == ".webm" {

			// Save the webm to a file
			err = image.SaveWebM()
			if err != nil {
				c.JSON(http.StatusBadRequest, gin.H{"error_message": err.Error()})
				c.Error(err)
				return
			}

		} else {

			// Save the image to a file
			err = image.SaveImage()
			if err != nil {
				c.JSON(http.StatusBadRequest, gin.H{"error_message": err.Error()})
				c.Error(err)
				return
			}

		}

		m.OrigWidth = image.OrigWidth
		m.OrigHeight = image.OrigHeight
		m.ThumbWidth = image.ThumbWidth
		m.ThumbHeight = image.ThumbHeight
		m.Filename = image.Filename
		m.Thumbnail = image.Thumbnail

	}

	// Post data
	err = m.Post()
	if err != nil {
		c.JSON(e.ErrorMessage(e.ErrInternalError))
		c.Error(err)
		return
	}

	// Initialize cache handle
	cache := redis.RedisCache

	// Delete redis stuff
	index_key := fmt.Sprintf("%s:%d", "index", m.Ib)
	directory_key := fmt.Sprintf("%s:%d", "directory", m.Ib)
	thread_key := fmt.Sprintf("%s:%d:%d", "thread", m.Ib, m.Thread)
	image_key := fmt.Sprintf("%s:%d", "image", m.Ib)

	err = cache.Delete(index_key, directory_key, thread_key, image_key)
	if err != nil {
		c.JSON(e.ErrorMessage(e.ErrInternalError))
		c.Error(err)
		return
	}

	// get board domain and redirect to it
	redirect, err := u.Link(m.Ib, req.Referer())
	if err != nil {
		c.JSON(e.ErrorMessage(e.ErrInternalError))
		c.Error(err)
		return
	}

	c.Redirect(303, redirect)

	audit := audit.Audit{
		User:   userdata.Id,
		Ib:     m.Ib,
		Ip:     m.Ip,
		Action: audit.AuditReply,
		Info:   fmt.Sprintf("%d/%d", m.Thread),
	}

	err = audit.Submit()
	if err != nil {
		c.Error(err)
	}

	return

}
<|endoftext|>


<commit_before>package horenso

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"time"

	"github.com/Songmu/timestamper"
	"github.com/Songmu/wrapcommander"
	"github.com/jessevdk/go-flags"
	"github.com/kballard/go-shellquote"
	"github.com/lestrrat-go/strftime"
	"golang.org/x/sync/errgroup"
	"golang.org/x/text/transform"
)

type horenso struct {
	Reporter       []string `short:"r" long:"reporter" value-name:"/path/to/reporter.pl" description:"handler for reporting the result of the job"`
	Noticer        []string `short:"n" long:"noticer" value-name:"'ruby /path/to/noticer.rb'" description:"handler for noticing the start of the job"`
	TimeStamp      bool     `short:"T" long:"timestamp" description:"add timestamp to merged output"`
	Tag            string   `short:"t" long:"tag" value-name:"job-name" description:"tag of the job"`
	OverrideStatus bool     `short:"o" long:"override-status" description:"override command exit status, always exit 0"`
	Verbose        []bool   `short:"v" long:"verbose" description:"verbose output. it can be stacked like -vv for more detailed log"`
	Logfile        string   `short:"l" long:"log" value-name:"/path/to/logfile" description:"logfile path. The strftime format like '%Y%m%d.log' is available."`
	Config         string   `short:"c" long:"config" value-name:"/path/to/config.yaml" description:"config file"`

	outStream, errStream io.Writer
}

// Report is represents the result of the command
type Report struct {
	Command     string     `json:"command"`
	CommandArgs []string   `json:"commandArgs"`
	Tag         string     `json:"tag,omitempty"`
	Output      string     `json:"output"`
	Stdout      string     `json:"stdout"`
	Stderr      string     `json:"stderr"`
	ExitCode    int        `json:"exitCode"`
	Signaled    bool       `json:"signaled"`
	Result      string     `json:"result"`
	Hostname    string     `json:"hostname"`
	Pid         int        `json:"pid,omitempty"`
	StartAt     *time.Time `json:"startAt,omitempty"`
	EndAt       *time.Time `json:"endAt,omitempty"`
	SystemTime  float64    `json:"systemTime,omitempty"`
	UserTime    float64    `json:"userTime,omitempty"`
}

func (ho *horenso) openLog() (io.WriteCloser, error) {
	logfile, err := strftime.Format(ho.Logfile, time.Now())
	if err != nil {
		return nil, fmt.Errorf("failed to parse log file format %q: %s", ho.Logfile, err)
	}
	f, err := os.OpenFile(logfile, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0644)
	if err != nil {
		return nil, fmt.Errorf("failed to open log file %q: %s", logfile, err)
	}
	return f, nil
}

func (ho *horenso) loadConfig() error {
	conf := ho.Config
	if conf == "" {
		conf = os.Getenv("HORENSO_CONFIG")
	}
	if conf == "" {
		return nil
	}
	c, err := loadConfig(conf)
	if err != nil {
		return err
	}
	ho.Reporter = append(ho.Reporter, c.Reporter...)
	ho.Noticer = append(ho.Noticer, c.Noticer...)
	if !ho.TimeStamp {
		ho.TimeStamp = c.Timestamp
	}
	if ho.Tag == "" {
		ho.Tag = c.Tag
	}
	if !ho.OverrideStatus {
		ho.OverrideStatus = c.OverrideStatus
	}
	if ho.Logfile == "" {
		ho.Logfile = c.Logfile
	}
	return nil
}

func (ho *horenso) run(args []string) (Report, error) {
	log.SetPrefix("[horenso] ")
	log.SetFlags(0)
	log.SetOutput(ho.errStream)

	if err := ho.loadConfig(); err != nil {
		ho.logf(warn, "failed to load config: %s", err)
	}

	hostname, _ := os.Hostname()
	r := Report{
		Command:     shellquote.Join(args...),
		CommandArgs: args,
		Tag:         ho.Tag,
		ExitCode:    -1,
		Hostname:    hostname,
	}
	cmd := exec.Command(args[0], args[1:]...)

	stdoutPipe, err := cmd.StdoutPipe()
	if err != nil {
		return ho.failReport(r, err.Error()), err
	}
	defer stdoutPipe.Close()

	stderrPipe, err := cmd.StderrPipe()
	if err != nil {
		return ho.failReport(r, err.Error()), err
	}
	defer stderrPipe.Close()

	var bufStdout bytes.Buffer
	var bufStderr bytes.Buffer
	var bufMerged bytes.Buffer

	var wtr io.Writer = &bufMerged
	if ho.Logfile != "" {
		if f, err := ho.openLog(); err != nil {
			ho.log(warn, err.Error())
		} else {
			defer f.Close()
			wtr = io.MultiWriter(wtr, f)
		}
	}
	if ho.TimeStamp {
		wc := transform.NewWriter(wtr, timestamper.New())
		defer wc.Close()
		wtr = wc
	}
	stdoutPipe2 := io.TeeReader(stdoutPipe, io.MultiWriter(&bufStdout, wtr))
	stderrPipe2 := io.TeeReader(stderrPipe, io.MultiWriter(&bufStderr, wtr))

	ho.logf(info, "starting executiton of the command %q", r.Command)
	r.StartAt = now()
	err = cmd.Start()
	if err != nil {
		return ho.failReport(r, err.Error()), err
	}
	if cmd.Process != nil {
		r.Pid = cmd.Process.Pid
	}
	done := make(chan error)
	go func(r Report) {
		done <- ho.runNoticer(r)
	}(r)

	eg := &errgroup.Group{}
	eg.Go(func() error {
		defer stdoutPipe.Close()
		_, err := io.Copy(ho.outStream, stdoutPipe2)
		return err
	})
	eg.Go(func() error {
		defer stderrPipe.Close()
		_, err := io.Copy(ho.errStream, stderrPipe2)
		return err
	})
	if err := eg.Wait(); err != nil {
		ho.logf(warn, "something went wrong while executing the command: %s", err)
	}
	err = cmd.Wait()
	r.EndAt = now()
	es := wrapcommander.ResolveExitStatus(err)
	r.ExitCode = es.ExitCode()
	r.Signaled = es.Signaled()
	r.Result = fmt.Sprintf("command exited with code: %d", r.ExitCode)
	if r.Signaled {
		r.Result = fmt.Sprintf("command died with signal: %d", r.ExitCode&127)
	}
	ho.logf(info, "the command %q finished: %s", r.Command, r.Result)
	r.Stdout = bufStdout.String()
	r.Stderr = bufStderr.String()
	r.Output = bufMerged.String()
	if p := cmd.ProcessState; p != nil {
		r.UserTime = float64(p.UserTime()) / float64(time.Second)
		r.SystemTime = float64(p.SystemTime()) / float64(time.Second)
	}
	ho.runReporter(r)
	<-done
	ho.logf(info, "all processes are completed for the job %q", r.Command)
	return r, nil
}

func now() *time.Time {
	now := time.Now()
	return &now
}

func parseArgs(args []string) (*flags.Parser, *horenso, []string, error) {
	ho := &horenso{}
	p := flags.NewParser(ho, flags.Default)
	p.Usage = fmt.Sprintf(`--reporter /path/to/reporter.pl -- /path/to/job [...]

Version: %s (rev: %s/%s)`, version, revision, runtime.Version())
	rest, err := p.ParseArgs(args)
	ho.outStream = os.Stdout
	ho.errStream = os.Stderr
	return p, ho, rest, err
}

// Run the horenso
func Run(args []string) int {
	p, ho, cmdArgs, err := parseArgs(args)
	if err != nil || len(cmdArgs) < 1 {
		if ferr, ok := err.(*flags.Error); !ok || ferr.Type != flags.ErrHelp {
			p.WriteHelp(ho.errStream)
		}
		return 2
	}
	r, err := ho.run(cmdArgs)
	if err != nil {
		return wrapcommander.ResolveExitCode(err)
	}
	if ho.OverrideStatus {
		return 0
	}
	return r.ExitCode
}

func (ho *horenso) failReport(r Report, errStr string) Report {
	r.Result = fmt.Sprintf("failed to execute the command: %s", errStr)
	ho.logf(warn, "failed to execute the command %q: %s", r.Command, errStr)
	done := make(chan error)
	go func() {
		done <- ho.runNoticer(r)
	}()
	ho.runReporter(r)
	<-done
	return r
}

func (ho *horenso) appendOut(base, out string) string {
	out = strings.TrimSpace(out)
	if out == "" {
		return base
	}
	if !strings.HasSuffix(base, "\n") {
		base += "\n"
	}
	const indent = "  "
	return base + indent + strings.Replace("Output:\n"+out, "\n", "\n"+indent, -1)
}

func (ho *horenso) runHandler(cmdStr string, json []byte) error {
	ho.logf(info, "starting to run the handler %q", cmdStr)
	args, err := shellquote.Split(cmdStr)
	if err != nil || len(args) < 1 {
		ho.logf(warn, "failed to run the handler %q: invalid handler arguments", cmdStr)
		return fmt.Errorf("invalid handler: %q", cmdStr)
	}
	cmd := exec.Command(args[0], args[1:]...)
	stdinPipe, _ := cmd.StdinPipe()
	var b bytes.Buffer
	cmd.Stdout = &b
	cmd.Stderr = &b
	if err := cmd.Start(); err != nil {
		stdinPipe.Close()
		logoutput := fmt.Sprintf("failed to run the handler %q: %s", cmdStr, err)
		ho.log(warn, ho.appendOut(logoutput, b.String()))
		return err
	}
	stdinPipe.Write(json)
	stdinPipe.Close()
	err = cmd.Wait()
	if err != nil || ho.logLevel() >= info {
		var logoutput string
		lv := info
		if err != nil {
			lv = warn
			logoutput = fmt.Sprintf("failed to run the handler %q: %s", cmdStr, err)
		} else {
			logoutput = fmt.Sprintf("finished to run the handler %q", cmdStr)
		}
		ho.log(lv, ho.appendOut(logoutput, b.String()))
	}
	return err
}

func (ho *horenso) runHandlers(handlers []string, json []byte) error {
	eg := &errgroup.Group{}
	for _, handler := range handlers {
		h := handler
		eg.Go(func() error {
			return ho.runHandler(h, json)
		})
	}
	return eg.Wait()
}

func (ho *horenso) runNoticer(r Report) error {
	if len(ho.Noticer) < 1 {
		return nil
	}
	ho.logf(info, "starting to run the noticers")
	defer ho.logf(info, "finished to run the noticers")
	json, _ := json.Marshal(r)
	return ho.runHandlers(ho.Noticer, json)
}

func (ho *horenso) runReporter(r Report) error {
	ho.logf(info, "starting to run the reporters")
	defer ho.logf(info, "finished to run the reporters")
	json, _ := json.Marshal(r)
	return ho.runHandlers(ho.Reporter, json)
}
<commit_msg>Add splitHandlerCmdStr method<commit_after>package horenso

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"time"

	"github.com/Songmu/timestamper"
	"github.com/Songmu/wrapcommander"
	"github.com/jessevdk/go-flags"
	"github.com/kballard/go-shellquote"
	"github.com/lestrrat-go/strftime"
	"golang.org/x/sync/errgroup"
	"golang.org/x/text/transform"
)

type horenso struct {
	Reporter       []string `short:"r" long:"reporter" value-name:"/path/to/reporter.pl" description:"handler for reporting the result of the job"`
	Noticer        []string `short:"n" long:"noticer" value-name:"'ruby /path/to/noticer.rb'" description:"handler for noticing the start of the job"`
	TimeStamp      bool     `short:"T" long:"timestamp" description:"add timestamp to merged output"`
	Tag            string   `short:"t" long:"tag" value-name:"job-name" description:"tag of the job"`
	OverrideStatus bool     `short:"o" long:"override-status" description:"override command exit status, always exit 0"`
	Verbose        []bool   `short:"v" long:"verbose" description:"verbose output. it can be stacked like -vv for more detailed log"`
	Logfile        string   `short:"l" long:"log" value-name:"/path/to/logfile" description:"logfile path. The strftime format like '%Y%m%d.log' is available."`
	Config         string   `short:"c" long:"config" value-name:"/path/to/config.yaml" description:"config file"`

	outStream, errStream io.Writer
}

// Report is represents the result of the command
type Report struct {
	Command     string     `json:"command"`
	CommandArgs []string   `json:"commandArgs"`
	Tag         string     `json:"tag,omitempty"`
	Output      string     `json:"output"`
	Stdout      string     `json:"stdout"`
	Stderr      string     `json:"stderr"`
	ExitCode    int        `json:"exitCode"`
	Signaled    bool       `json:"signaled"`
	Result      string     `json:"result"`
	Hostname    string     `json:"hostname"`
	Pid         int        `json:"pid,omitempty"`
	StartAt     *time.Time `json:"startAt,omitempty"`
	EndAt       *time.Time `json:"endAt,omitempty"`
	SystemTime  float64    `json:"systemTime,omitempty"`
	UserTime    float64    `json:"userTime,omitempty"`
}

func (ho *horenso) openLog() (io.WriteCloser, error) {
	logfile, err := strftime.Format(ho.Logfile, time.Now())
	if err != nil {
		return nil, fmt.Errorf("failed to parse log file format %q: %s", ho.Logfile, err)
	}
	f, err := os.OpenFile(logfile, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0644)
	if err != nil {
		return nil, fmt.Errorf("failed to open log file %q: %s", logfile, err)
	}
	return f, nil
}

func (ho *horenso) loadConfig() error {
	conf := ho.Config
	if conf == "" {
		conf = os.Getenv("HORENSO_CONFIG")
	}
	if conf == "" {
		return nil
	}
	c, err := loadConfig(conf)
	if err != nil {
		return err
	}
	ho.Reporter = append(ho.Reporter, c.Reporter...)
	ho.Noticer = append(ho.Noticer, c.Noticer...)
	if !ho.TimeStamp {
		ho.TimeStamp = c.Timestamp
	}
	if ho.Tag == "" {
		ho.Tag = c.Tag
	}
	if !ho.OverrideStatus {
		ho.OverrideStatus = c.OverrideStatus
	}
	if ho.Logfile == "" {
		ho.Logfile = c.Logfile
	}
	return nil
}

func (ho *horenso) run(args []string) (Report, error) {
	log.SetPrefix("[horenso] ")
	log.SetFlags(0)
	log.SetOutput(ho.errStream)

	if err := ho.loadConfig(); err != nil {
		ho.logf(warn, "failed to load config: %s", err)
	}

	hostname, _ := os.Hostname()
	r := Report{
		Command:     shellquote.Join(args...),
		CommandArgs: args,
		Tag:         ho.Tag,
		ExitCode:    -1,
		Hostname:    hostname,
	}
	cmd := exec.Command(args[0], args[1:]...)

	stdoutPipe, err := cmd.StdoutPipe()
	if err != nil {
		return ho.failReport(r, err.Error()), err
	}
	defer stdoutPipe.Close()

	stderrPipe, err := cmd.StderrPipe()
	if err != nil {
		return ho.failReport(r, err.Error()), err
	}
	defer stderrPipe.Close()

	var bufStdout bytes.Buffer
	var bufStderr bytes.Buffer
	var bufMerged bytes.Buffer

	var wtr io.Writer = &bufMerged
	if ho.Logfile != "" {
		if f, err := ho.openLog(); err != nil {
			ho.log(warn, err.Error())
		} else {
			defer f.Close()
			wtr = io.MultiWriter(wtr, f)
		}
	}
	if ho.TimeStamp {
		wc := transform.NewWriter(wtr, timestamper.New())
		defer wc.Close()
		wtr = wc
	}
	stdoutPipe2 := io.TeeReader(stdoutPipe, io.MultiWriter(&bufStdout, wtr))
	stderrPipe2 := io.TeeReader(stderrPipe, io.MultiWriter(&bufStderr, wtr))

	ho.logf(info, "starting executiton of the command %q", r.Command)
	r.StartAt = now()
	err = cmd.Start()
	if err != nil {
		return ho.failReport(r, err.Error()), err
	}
	if cmd.Process != nil {
		r.Pid = cmd.Process.Pid
	}
	done := make(chan error)
	go func(r Report) {
		done <- ho.runNoticer(r)
	}(r)

	eg := &errgroup.Group{}
	eg.Go(func() error {
		defer stdoutPipe.Close()
		_, err := io.Copy(ho.outStream, stdoutPipe2)
		return err
	})
	eg.Go(func() error {
		defer stderrPipe.Close()
		_, err := io.Copy(ho.errStream, stderrPipe2)
		return err
	})
	if err := eg.Wait(); err != nil {
		ho.logf(warn, "something went wrong while executing the command: %s", err)
	}
	err = cmd.Wait()
	r.EndAt = now()
	es := wrapcommander.ResolveExitStatus(err)
	r.ExitCode = es.ExitCode()
	r.Signaled = es.Signaled()
	r.Result = fmt.Sprintf("command exited with code: %d", r.ExitCode)
	if r.Signaled {
		r.Result = fmt.Sprintf("command died with signal: %d", r.ExitCode&127)
	}
	ho.logf(info, "the command %q finished: %s", r.Command, r.Result)
	r.Stdout = bufStdout.String()
	r.Stderr = bufStderr.String()
	r.Output = bufMerged.String()
	if p := cmd.ProcessState; p != nil {
		r.UserTime = float64(p.UserTime()) / float64(time.Second)
		r.SystemTime = float64(p.SystemTime()) / float64(time.Second)
	}
	ho.runReporter(r)
	<-done
	ho.logf(info, "all processes are completed for the job %q", r.Command)
	return r, nil
}

func now() *time.Time {
	now := time.Now()
	return &now
}

func parseArgs(args []string) (*flags.Parser, *horenso, []string, error) {
	ho := &horenso{}
	p := flags.NewParser(ho, flags.Default)
	p.Usage = fmt.Sprintf(`--reporter /path/to/reporter.pl -- /path/to/job [...]

Version: %s (rev: %s/%s)`, version, revision, runtime.Version())
	rest, err := p.ParseArgs(args)
	ho.outStream = os.Stdout
	ho.errStream = os.Stderr
	return p, ho, rest, err
}

// Run the horenso
func Run(args []string) int {
	p, ho, cmdArgs, err := parseArgs(args)
	if err != nil || len(cmdArgs) < 1 {
		if ferr, ok := err.(*flags.Error); !ok || ferr.Type != flags.ErrHelp {
			p.WriteHelp(ho.errStream)
		}
		return 2
	}
	r, err := ho.run(cmdArgs)
	if err != nil {
		return wrapcommander.ResolveExitCode(err)
	}
	if ho.OverrideStatus {
		return 0
	}
	return r.ExitCode
}

func (ho *horenso) failReport(r Report, errStr string) Report {
	r.Result = fmt.Sprintf("failed to execute the command: %s", errStr)
	ho.logf(warn, "failed to execute the command %q: %s", r.Command, errStr)
	done := make(chan error)
	go func() {
		done <- ho.runNoticer(r)
	}()
	ho.runReporter(r)
	<-done
	return r
}

func (ho *horenso) appendOut(base, out string) string {
	out = strings.TrimSpace(out)
	if out == "" {
		return base
	}
	if !strings.HasSuffix(base, "\n") {
		base += "\n"
	}
	const indent = "  "
	return base + indent + strings.Replace("Output:\n"+out, "\n", "\n"+indent, -1)
}

func (ho *horenso) splitHandlerCmdStr(cmdStr string) ([]string, error) {
	args, err := shellquote.Split(cmdStr)
	return args, err
}

func (ho *horenso) runHandler(cmdStr string, json []byte) error {
	ho.logf(info, "starting to run the handler %q", cmdStr)
	args, err := ho.splitHandlerCmdStr(cmdStr)
	if err != nil || len(args) < 1 {
		ho.logf(warn, "failed to run the handler %q: invalid handler arguments", cmdStr)
		return fmt.Errorf("invalid handler: %q", cmdStr)
	}
	cmd := exec.Command(args[0], args[1:]...)
	stdinPipe, _ := cmd.StdinPipe()
	var b bytes.Buffer
	cmd.Stdout = &b
	cmd.Stderr = &b
	if err := cmd.Start(); err != nil {
		stdinPipe.Close()
		logoutput := fmt.Sprintf("failed to run the handler %q: %s", cmdStr, err)
		ho.log(warn, ho.appendOut(logoutput, b.String()))
		return err
	}
	stdinPipe.Write(json)
	stdinPipe.Close()
	err = cmd.Wait()
	if err != nil || ho.logLevel() >= info {
		var logoutput string
		lv := info
		if err != nil {
			lv = warn
			logoutput = fmt.Sprintf("failed to run the handler %q: %s", cmdStr, err)
		} else {
			logoutput = fmt.Sprintf("finished to run the handler %q", cmdStr)
		}
		ho.log(lv, ho.appendOut(logoutput, b.String()))
	}
	return err
}

func (ho *horenso) runHandlers(handlers []string, json []byte) error {
	eg := &errgroup.Group{}
	for _, handler := range handlers {
		h := handler
		eg.Go(func() error {
			return ho.runHandler(h, json)
		})
	}
	return eg.Wait()
}

func (ho *horenso) runNoticer(r Report) error {
	if len(ho.Noticer) < 1 {
		return nil
	}
	ho.logf(info, "starting to run the noticers")
	defer ho.logf(info, "finished to run the noticers")
	json, _ := json.Marshal(r)
	return ho.runHandlers(ho.Noticer, json)
}

func (ho *horenso) runReporter(r Report) error {
	ho.logf(info, "starting to run the reporters")
	defer ho.logf(info, "finished to run the reporters")
	json, _ := json.Marshal(r)
	return ho.runHandlers(ho.Reporter, json)
}
<|endoftext|>


<commit_before>package apps

import (
	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry-incubator/cf-test-helpers/generator"
	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry-incubator/cf-test-helpers/helpers"
	"github.com/cloudfoundry/cf-acceptance-tests/helpers/app_helpers"
	"github.com/cloudfoundry/cf-acceptance-tests/helpers/assets"

	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/ginkgo"
	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/gomega"
	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/gomega/gexec"
)

var _ = Describe("User services", func() {
	var testConfig = helpers.LoadConfig()
	var appName string

	Describe("When a user-defined service is bound", func() {
		var serviceName string

		BeforeEach(func() {
			appName = generator.PrefixedRandomName("CATS-APP-")
			serviceName = generator.PrefixedRandomName("CUPS-")

			Eventually(cf.Cf(
				"push",
				appName,
				"--no-start",
				"-b", testConfig.RubyBuildpackName,
				"-m", DEFAULT_MEMORY_LIMIT,
				"-p", assets.NewAssets().RubySimple,
				"-d", testConfig.AppsDomain), DEFAULT_TIMEOUT).Should(Exit(0), "Failed to push app")

			app_helpers.SetBackend(appName)
			Expect(cf.Cf("start", appName).Wait(CF_PUSH_TIMEOUT)).To(Exit(0))

			Eventually(cf.Cf("cups", serviceName, "-l", "does this matter"), DEFAULT_TIMEOUT).Should(Exit(0), "Failed to create syslog drain service")
			Eventually(cf.Cf("bind-service", appName, serviceName), DEFAULT_TIMEOUT).Should(Exit(0), "Failed to bind service")
		})

		// No AfterEach because the delete occurs in the test itself

		It("can be deleted", func() {
			// AppReport here instead of AfterEach because the delete occurs in the test itself
			app_helpers.AppReport(appName, DEFAULT_TIMEOUT)

			Eventually(cf.Cf("delete", appName, "-f", "-r"), DEFAULT_TIMEOUT).Should(Exit(0), "Failed to delete app")
			if serviceName != "" {
				Eventually(cf.Cf("delete-service", serviceName, "-f"), DEFAULT_TIMEOUT).Should(Exit(0), "Failed to delete service")
			}
		})
	})
})
<commit_msg>Remove cups regression test<commit_after><|endoftext|>