<commit_before>package eval

import (
	"fmt"
	"monkey/ast"
	"monkey/object"
)

var (
	TRUE  = &object.Boolean{Value: true}
	FALSE = &object.Boolean{Value: false}
	NULL  = &object.Null{}
)

func Eval(node ast.Node, scope *object.Scope) object.Object {
	switch node := node.(type) {
	case *ast.Program:
		return evalProgram(node, scope)
	case *ast.CallExpression:
		f_scope := object.NewScope(scope)
		fn, ok := scope.Get(node.Function.String())
		if !ok {
			switch f := node.Function.(type) {
			case *ast.CallExpression:
				fn = &object.Function{Literal: f.Function.(*ast.FunctionLiteral), Scope: scope}
			case *ast.FunctionLiteral:
				fn = &object.Function{Literal: f, Scope: scope}
			}
			scope.Set(node.Function.String(), fn)
		}
		return evalFunctionCall(node, f_scope)
	case *ast.FunctionLiteral:
		return &object.Function{Literal: node, Scope: scope}
	case *ast.LetStatement:
		val := Eval(node.Value, scope)
		if val.Type() == object.ERROR_OBJ {
			return val
		}
		return scope.Set(node.Name.String(), val)
	case *ast.Identifier:
		if val, ok := scope.Get(node.String()); ok {
			return val
		}
		return &object.Error{Message: fmt.Sprintf("unknown identifier: %s", node.String())}
	case *ast.ExpressionStatement:
		return Eval(node.Expression, scope)
	case *ast.ReturnStatement:
		value := Eval(node.ReturnValue, scope)
		if value != nil {
			return &object.ReturnValue{Value: value}
		}
		return NULL
	case *ast.BlockStatement:
		return evalBlockStatements(node.Statements, scope)
	case *ast.InfixExpression:
		left := Eval(node.Left, scope)
		right := Eval(node.Right, scope)
		if left.Type() == object.ERROR_OBJ {
			return left
		} else if right.Type() == object.ERROR_OBJ {
			return right
		}
		return evalInfixExpression(node.Operator, left, right)
	case *ast.PrefixExpression:
		right := Eval(node.Right, scope)
		if right.Type() == object.ERROR_OBJ {
			return right
		}
		return evalPrefixExpression(node.Operator, right)
	case *ast.IfExpression:
		condition := Eval(node.Condition, scope)
		if isTrue(condition) {
			return evalBlockStatements(node.Consequence.Statements, scope)
		} else if node.Alternative != nil {
			return evalBlockStatements(node.Alternative.Statements, scope)
		}
		return NULL
	case *ast.IntegerLiteral:
		return &object.Integer{Value: node.Value}
	case *ast.Boolean:
		return nativeBoolToBooleanObject(node.Value)
	}
	return nil
}

func isTrue(obj object.Object) bool {
	switch obj {
	case TRUE:
		return true
	case FALSE:
		return false
	case NULL:
		return false
	default:
		return true
	}
}

func nativeBoolToBooleanObject(input bool) *object.Boolean {
	if input {
		return TRUE
	}
	return FALSE
}

func evalBlockStatements(block []ast.Statement, scope *object.Scope) object.Object {
	var results object.Object

	for _, statement := range block {
		results = Eval(statement, scope)
		if results != nil && results.Type() == object.RETURN_VALUE_OBJ {
			return results
		}
	}
	return results
}

func evalProgram(program *ast.Program, scope *object.Scope) object.Object {
	var results object.Object

	for _, statement := range program.Statements {
		results = Eval(statement, scope)
		switch s := results.(type) {
		case *object.ReturnValue:
			return s.Value
		case *object.Error:
			return s
		}
	}

	return results
}

func evalInfixExpression(operator string, left object.Object, right object.Object) object.Object {
	var errMsg string
	switch {
	case left.Type() == object.INTEGER_OBJ && right.Type() == object.INTEGER_OBJ:
		return evalIntInfixExpression(operator, left, right)
	case operator == "==":
		return nativeBoolToBooleanObject(left == right)
	case operator == "!=":
		return nativeBoolToBooleanObject(left != right)
	case left.Type() != right.Type():
		errMsg = fmt.Sprintf("type mismatch: %s %s %s", left.Type(), operator, right.Type())
	default:
		errMsg = fmt.Sprintf("unknown operator: %s %s %s", left.Type(), operator, right.Type())
	}
	if errMsg != "" {
		return &object.Error{Message: errMsg}
	}
	return NULL
}

func evalIntInfixExpression(operator string, left object.Object, right object.Object) object.Object {
	l := left.(*object.Integer)
	r := right.(*object.Integer)

	switch operator {
	case "+":
		return &object.Integer{Value: l.Value + r.Value}
	case "-":
		return &object.Integer{Value: l.Value - r.Value}
	case "*":
		return &object.Integer{Value: l.Value * r.Value}
	case "/":
		return &object.Integer{Value: l.Value / r.Value}
	case ">":
		return nativeBoolToBooleanObject(l.Value > r.Value)
	case "<":
		return nativeBoolToBooleanObject(l.Value < r.Value)
	case "==":
		return nativeBoolToBooleanObject(l.Value == r.Value)
	case "!=":
		return nativeBoolToBooleanObject(l.Value != r.Value)
	}
	return NULL
}

func evalPrefixExpression(operator string, right object.Object) object.Object {
	switch operator {
	case "!":
		return evalBangOperatorExpression(right)
	case "-":
		if i, ok := right.(*object.Integer); ok {
			i.Value = -i.Value
			return right
		}
		msg := fmt.Sprintf("unknown operator: %s%s", operator, right.Type())
		return &object.Error{Message: msg}
	default:
		return NULL
	}
}

func evalBangOperatorExpression(right object.Object) object.Object {
	switch right {
	case TRUE:
		return FALSE
	case FALSE:
		return TRUE
	case NULL:
		return TRUE
	default:
		return FALSE
	}
}

func evalFunctionCall(call *ast.CallExpression, scope *object.Scope) object.Object {
	v, ok := scope.Get(call.Function.String())
	if !ok {
		return &object.Error{Message: fmt.Sprintf("unknown identifier: %s", call.Function.String())}
	}
	fn := v.(*object.Function)
	fn.Scope = scope
	for i, v := range fn.Literal.Parameters {
		value := Eval(call.Arguments[i], fn.Scope)
		if value.Type() == object.RETURN_VALUE_OBJ {
			value = value.(*object.ReturnValue).Value
		}
		scope.Set(v.String(), value)
	}
	v = Eval(fn.Literal.Body, scope)
	if obj, ok := v.(*object.ReturnValue); ok {
		return obj.Value
	}
	return v
}
<commit_msg>removed redundanct type assertion<commit_after>package eval

import (
	"fmt"
	"monkey/ast"
	"monkey/object"
)

var (
	TRUE  = &object.Boolean{Value: true}
	FALSE = &object.Boolean{Value: false}
	NULL  = &object.Null{}
)

func Eval(node ast.Node, scope *object.Scope) object.Object {
	switch node := node.(type) {
	case *ast.Program:
		return evalProgram(node, scope)
	case *ast.CallExpression:
		f_scope := object.NewScope(scope)
		fn, ok := scope.Get(node.Function.String())
		if !ok {
			switch f := node.Function.(type) {
			case *ast.CallExpression:
				fn = &object.Function{Literal: f.Function.(*ast.FunctionLiteral), Scope: scope}
			case *ast.FunctionLiteral:
				fn = &object.Function{Literal: f, Scope: scope}
			}
			scope.Set(node.Function.String(), fn)
		}
		return evalFunctionCall(node, f_scope)
	case *ast.FunctionLiteral:
		return &object.Function{Literal: node, Scope: scope}
	case *ast.LetStatement:
		val := Eval(node.Value, scope)
		if val.Type() == object.ERROR_OBJ {
			return val
		}
		return scope.Set(node.Name.String(), val)
	case *ast.Identifier:
		if val, ok := scope.Get(node.String()); ok {
			return val
		}
		return &object.Error{Message: fmt.Sprintf("unknown identifier: %s", node.String())}
	case *ast.ExpressionStatement:
		return Eval(node.Expression, scope)
	case *ast.ReturnStatement:
		value := Eval(node.ReturnValue, scope)
		if value != nil {
			return &object.ReturnValue{Value: value}
		}
		return NULL
	case *ast.BlockStatement:
		return evalBlockStatements(node.Statements, scope)
	case *ast.InfixExpression:
		left := Eval(node.Left, scope)
		right := Eval(node.Right, scope)
		if left.Type() == object.ERROR_OBJ {
			return left
		} else if right.Type() == object.ERROR_OBJ {
			return right
		}
		return evalInfixExpression(node.Operator, left, right)
	case *ast.PrefixExpression:
		right := Eval(node.Right, scope)
		if right.Type() == object.ERROR_OBJ {
			return right
		}
		return evalPrefixExpression(node.Operator, right)
	case *ast.IfExpression:
		condition := Eval(node.Condition, scope)
		if isTrue(condition) {
			return evalBlockStatements(node.Consequence.Statements, scope)
		} else if node.Alternative != nil {
			return evalBlockStatements(node.Alternative.Statements, scope)
		}
		return NULL
	case *ast.IntegerLiteral:
		return &object.Integer{Value: node.Value}
	case *ast.Boolean:
		return nativeBoolToBooleanObject(node.Value)
	}
	return nil
}

func isTrue(obj object.Object) bool {
	switch obj {
	case TRUE:
		return true
	case FALSE:
		return false
	case NULL:
		return false
	default:
		return true
	}
}

func nativeBoolToBooleanObject(input bool) *object.Boolean {
	if input {
		return TRUE
	}
	return FALSE
}

func evalBlockStatements(block []ast.Statement, scope *object.Scope) object.Object {
	var results object.Object

	for _, statement := range block {
		results = Eval(statement, scope)
		if results != nil && results.Type() == object.RETURN_VALUE_OBJ {
			return results
		}
	}
	return results
}

func evalProgram(program *ast.Program, scope *object.Scope) object.Object {
	var results object.Object

	for _, statement := range program.Statements {
		results = Eval(statement, scope)
		switch s := results.(type) {
		case *object.ReturnValue:
			return s.Value
		case *object.Error:
			return s
		}
	}

	return results
}

func evalInfixExpression(operator string, left object.Object, right object.Object) object.Object {
	var errMsg string
	switch {
	case left.Type() == object.INTEGER_OBJ && right.Type() == object.INTEGER_OBJ:
		return evalIntInfixExpression(operator, left, right)
	case operator == "==":
		return nativeBoolToBooleanObject(left == right)
	case operator == "!=":
		return nativeBoolToBooleanObject(left != right)
	case left.Type() != right.Type():
		errMsg = fmt.Sprintf("type mismatch: %s %s %s", left.Type(), operator, right.Type())
	default:
		errMsg = fmt.Sprintf("unknown operator: %s %s %s", left.Type(), operator, right.Type())
	}
	if errMsg != "" {
		return &object.Error{Message: errMsg}
	}
	return NULL
}

func evalIntInfixExpression(operator string, left object.Object, right object.Object) object.Object {
	l := left.(*object.Integer)
	r := right.(*object.Integer)

	switch operator {
	case "+":
		return &object.Integer{Value: l.Value + r.Value}
	case "-":
		return &object.Integer{Value: l.Value - r.Value}
	case "*":
		return &object.Integer{Value: l.Value * r.Value}
	case "/":
		return &object.Integer{Value: l.Value / r.Value}
	case ">":
		return nativeBoolToBooleanObject(l.Value > r.Value)
	case "<":
		return nativeBoolToBooleanObject(l.Value < r.Value)
	case "==":
		return nativeBoolToBooleanObject(l.Value == r.Value)
	case "!=":
		return nativeBoolToBooleanObject(l.Value != r.Value)
	}
	return NULL
}

func evalPrefixExpression(operator string, right object.Object) object.Object {
	switch operator {
	case "!":
		return evalBangOperatorExpression(right)
	case "-":
		if i, ok := right.(*object.Integer); ok {
			i.Value = -i.Value
			return right
		}
		msg := fmt.Sprintf("unknown operator: %s%s", operator, right.Type())
		return &object.Error{Message: msg}
	default:
		return NULL
	}
}

func evalBangOperatorExpression(right object.Object) object.Object {
	switch right {
	case TRUE:
		return FALSE
	case FALSE:
		return TRUE
	case NULL:
		return TRUE
	default:
		return FALSE
	}
}

func evalFunctionCall(call *ast.CallExpression, scope *object.Scope) object.Object {
	v, ok := scope.Get(call.Function.String())
	if !ok {
		return &object.Error{Message: fmt.Sprintf("unknown identifier: %s", call.Function.String())}
	}
	fn := v.(*object.Function)
	fn.Scope = scope
	for i, v := range fn.Literal.Parameters {
		value := Eval(call.Arguments[i], fn.Scope)
		scope.Set(v.String(), value)
	}
	v = Eval(fn.Literal.Body, scope)
	if obj, ok := v.(*object.ReturnValue); ok {
		return obj.Value
	}
	return v
}
<|endoftext|>


<commit_before>package collector

import (
	"fmt"
	"strconv"
	"strings"
	"sync"
	"time"

	"gitlab.com/isard/isardvdi/stats/cfg"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/rs/zerolog"
	"golang.org/x/crypto/ssh"
)

type Socket struct {
	mux    *sync.Mutex
	domain string
	Log    *zerolog.Logger
	conn   *ssh.Client

	descScrapeDuration *prometheus.Desc
	descScrapeSuccess  *prometheus.Desc
	descSent           *prometheus.Desc
	descRecv           *prometheus.Desc
}

func NewSocket(mux *sync.Mutex, cfg cfg.Cfg, log *zerolog.Logger, conn *ssh.Client) *Socket {
	s := &Socket{
		mux:    mux,
		domain: cfg.Domain,
		Log:    log,
		conn:   conn,
	}

	s.descScrapeDuration = prometheus.NewDesc(
		prometheus.BuildFQName(namespace, s.String(), "scrape_duration_seconds"),
		"node_exporter: Duration of a collector scrape",
		[]string{},
		prometheus.Labels{
			"hypervisor": cfg.Domain,
		},
	)
	s.descScrapeSuccess = prometheus.NewDesc(
		prometheus.BuildFQName(namespace, s.String(), "scrape_duration_success"),
		"node_exporter: Whether a collector succeed",
		[]string{},
		prometheus.Labels{
			"hypervisor": cfg.Domain,
		},
	)
	s.descSent = prometheus.NewDesc(
		prometheus.BuildFQName(namespace, s.String(), "sent_bytes"),
		"Bytes sent",
		[]string{"port"},
		prometheus.Labels{
			"hypervisor": cfg.Domain,
		},
	)
	s.descRecv = prometheus.NewDesc(
		prometheus.BuildFQName(namespace, s.String(), "recv_bytes"),
		"Bytes recieved",
		[]string{"port"},
		prometheus.Labels{
			"hypervisor": cfg.Domain,
		},
	)

	return s
}

func (s *Socket) String() string {
	return "socket"
}

func (s *Socket) Describe(ch chan<- *prometheus.Desc) {
	ch <- s.descScrapeDuration
	ch <- s.descScrapeSuccess
	ch <- s.descSent
	ch <- s.descRecv
}

func (s *Socket) Collect(ch chan<- prometheus.Metric) {
	start := time.Now()

	success := 1
	viewers, err := s.collectViewers()
	if err != nil {
		s.Log.Info().Str("collector", s.String()).Err(err).Msg("collect viewers")
		success = 0
	}

	for src, v := range viewers {
		port := strconv.Itoa(src)
		ch <- prometheus.MustNewConstMetric(s.descSent, prometheus.GaugeValue, float64(v.Sent), port)
		ch <- prometheus.MustNewConstMetric(s.descRecv, prometheus.GaugeValue, float64(v.Recv), port)
	}

	duration := time.Since(start)

	ch <- prometheus.MustNewConstMetric(s.descScrapeDuration, prometheus.GaugeValue, duration.Seconds())
	ch <- prometheus.MustNewConstMetric(s.descScrapeSuccess, prometheus.GaugeValue, float64(success))

}

type viewer struct {
	DstPorts      []int
	PID           int
	SentByDstPort []int
	RecvByDstPort []int
	Sent          int
	Recv          int
}

func (s *Socket) collectViewers() (map[int]*viewer, error) {
	s.mux.Lock()
	defer s.mux.Unlock()

	sess, err := s.conn.NewSession()
	if err != nil {
		return nil, fmt.Errorf("create ssh session: %w", err)
	}
	defer sess.Close()

	// TODO: 5700 for websocket ports?
	b, err := sess.CombinedOutput(`ss -t state established -o state established -t -n -p -i "( sport > 5900 )" -O`)
	if err != nil {
		return nil, fmt.Errorf("collect viewers: %w", err)
	}

	viewers := map[int]*viewer{}
	lines := strings.Split(string(b), "\n")

	for i, line := range lines {
		if strings.Contains(line, "qemu") {
			split := strings.Split(line, ":")

			src, err := strconv.Atoi(strings.Split(split[1], " ")[0])
			if err != nil {
				return map[int]*viewer{}, fmt.Errorf("parse source port: %w", err)
			}

			dst, err := strconv.Atoi(strings.Split(split[2], " ")[0])
			if err != nil {
				return map[int]*viewer{}, fmt.Errorf("parse destination port: %w", err)
			}

			var (
				sent, recv int
			)

			bAcked := strings.Split(lines[i], "bytes_acked:")
			if len(bAcked) > 1 {
				sent, err = strconv.Atoi(strings.Split(bAcked[1], " ")[0])
				if err != nil {
					return map[int]*viewer{}, fmt.Errorf("parse bytes sent: %w", err)
				}

				recv, err = strconv.Atoi(strings.Split(strings.Split(lines[i], "bytes_received:")[1], " ")[0])
				if err != nil {
					return map[int]*viewer{}, fmt.Errorf("parse bytes received: %w", err)
				}
			}

			pid, err := strconv.Atoi(strings.Split(strings.Split(line, "pid=")[1], ",")[0])
			if err != nil {
				return map[int]*viewer{}, fmt.Errorf("parse pid: %w", err)
			}

			if v, ok := viewers[src]; !ok {
				viewers[src] = &viewer{
					DstPorts:      []int{},
					PID:           pid,
					SentByDstPort: []int{},
					RecvByDstPort: []int{},
					Sent:          0,
					Recv:          0,
				}

			} else {
				v.DstPorts = append(v.DstPorts, dst)
				v.SentByDstPort = append(v.SentByDstPort, sent)
				v.RecvByDstPort = append(v.RecvByDstPort, recv)
				v.Sent += sent
				v.Recv += recv
			}
		}
	}

	return viewers, nil
}
<commit_msg>fix(stats): fix the socket collector<commit_after>package collector

import (
	"fmt"
	"strconv"
	"strings"
	"sync"
	"time"

	"gitlab.com/isard/isardvdi/stats/cfg"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/rs/zerolog"
	"golang.org/x/crypto/ssh"
)

type Socket struct {
	mux    *sync.Mutex
	domain string
	Log    *zerolog.Logger
	conn   *ssh.Client

	descScrapeDuration *prometheus.Desc
	descScrapeSuccess  *prometheus.Desc
	descSent           *prometheus.Desc
	descRecv           *prometheus.Desc
}

func NewSocket(mux *sync.Mutex, cfg cfg.Cfg, log *zerolog.Logger, conn *ssh.Client) *Socket {
	s := &Socket{
		mux:    mux,
		domain: cfg.Domain,
		Log:    log,
		conn:   conn,
	}

	s.descScrapeDuration = prometheus.NewDesc(
		prometheus.BuildFQName(namespace, s.String(), "scrape_duration_seconds"),
		"node_exporter: Duration of a collector scrape",
		[]string{},
		prometheus.Labels{
			"hypervisor": cfg.Domain,
		},
	)
	s.descScrapeSuccess = prometheus.NewDesc(
		prometheus.BuildFQName(namespace, s.String(), "scrape_duration_success"),
		"node_exporter: Whether a collector succeed",
		[]string{},
		prometheus.Labels{
			"hypervisor": cfg.Domain,
		},
	)
	s.descSent = prometheus.NewDesc(
		prometheus.BuildFQName(namespace, s.String(), "sent_bytes"),
		"Bytes sent",
		[]string{"port"},
		prometheus.Labels{
			"hypervisor": cfg.Domain,
		},
	)
	s.descRecv = prometheus.NewDesc(
		prometheus.BuildFQName(namespace, s.String(), "recv_bytes"),
		"Bytes recieved",
		[]string{"port"},
		prometheus.Labels{
			"hypervisor": cfg.Domain,
		},
	)

	return s
}

func (s *Socket) String() string {
	return "socket"
}

func (s *Socket) Describe(ch chan<- *prometheus.Desc) {
	ch <- s.descScrapeDuration
	ch <- s.descScrapeSuccess
	ch <- s.descSent
	ch <- s.descRecv
}

func (s *Socket) Collect(ch chan<- prometheus.Metric) {
	start := time.Now()

	success := 1
	viewers, err := s.collectViewers()
	if err != nil {
		s.Log.Info().Str("collector", s.String()).Err(err).Msg("collect viewers")
		success = 0
	}

	for src, v := range viewers {
		port := strconv.Itoa(src)
		ch <- prometheus.MustNewConstMetric(s.descSent, prometheus.GaugeValue, float64(v.Sent), port)
		ch <- prometheus.MustNewConstMetric(s.descRecv, prometheus.GaugeValue, float64(v.Recv), port)
	}

	duration := time.Since(start)

	ch <- prometheus.MustNewConstMetric(s.descScrapeDuration, prometheus.GaugeValue, duration.Seconds())
	ch <- prometheus.MustNewConstMetric(s.descScrapeSuccess, prometheus.GaugeValue, float64(success))

}

type viewer struct {
	DstPorts      []int
	PID           int
	SentByDstPort []int
	RecvByDstPort []int
	Sent          int
	Recv          int
}

func (s *Socket) collectViewers() (map[int]*viewer, error) {
	s.mux.Lock()
	defer s.mux.Unlock()

	sess, err := s.conn.NewSession()
	if err != nil {
		return nil, fmt.Errorf("create ssh session: %w", err)
	}
	defer sess.Close()

	// TODO: 5700 for websocket ports?
	b, err := sess.CombinedOutput(`ss -t state established -o state established -t -n -p -i "( sport > 5900 )" -O`)
	if err != nil {
		return nil, fmt.Errorf("collect viewers: %w", err)
	}

	viewers := map[int]*viewer{}
	lines := strings.Split(string(b), "\n")

	for i, line := range lines {
		if strings.Contains(line, "qemu") {
			split := strings.Split(line, ":")

			src, err := strconv.Atoi(strings.Split(split[1], " ")[0])
			if err != nil {
				return map[int]*viewer{}, fmt.Errorf("parse source port: %w", err)
			}

			dst, err := strconv.Atoi(strings.Split(split[2], " ")[0])
			if err != nil {
				return map[int]*viewer{}, fmt.Errorf("parse destination port: %w", err)
			}

			var (
				sent, recv int
			)

			bAcked := strings.Split(lines[i], "bytes_acked:")
			if len(bAcked) > 1 {
				sent, err = strconv.Atoi(strings.Split(bAcked[1], " ")[0])
				if err != nil {
					return map[int]*viewer{}, fmt.Errorf("parse bytes sent: %w", err)
				}

				bRecv := strings.Split(lines[i], "bytes_received:")
				if len(bRecv) > 1 {
					recv, err = strconv.Atoi(strings.Split(bRecv[1], " ")[0])
					if err != nil {
						return map[int]*viewer{}, fmt.Errorf("parse bytes received: %w", err)
					}
				}
			}

			pid, err := strconv.Atoi(strings.Split(strings.Split(line, "pid=")[1], ",")[0])
			if err != nil {
				return map[int]*viewer{}, fmt.Errorf("parse pid: %w", err)
			}

			if v, ok := viewers[src]; !ok {
				viewers[src] = &viewer{
					DstPorts:      []int{},
					PID:           pid,
					SentByDstPort: []int{},
					RecvByDstPort: []int{},
					Sent:          0,
					Recv:          0,
				}

			} else {
				v.DstPorts = append(v.DstPorts, dst)
				v.SentByDstPort = append(v.SentByDstPort, sent)
				v.RecvByDstPort = append(v.RecvByDstPort, recv)
				v.Sent += sent
				v.Recv += recv
			}
		}
	}

	return viewers, nil
}
<|endoftext|>


<commit_before>package goexcel

import (
	"bytes"
	"encoding/csv"
	"errors"
	"net/http"
	"reflect"
	"strings"

	"google.golang.org/appengine"
	"google.golang.org/appengine/log"

	"golang.org/x/net/context"

	"github.com/news-ai/tabulae/models"
)

func csvGetCustomFields(r *http.Request, c context.Context, numberOfColumns int, headers []string) map[string]bool {
	customFields := make(map[string]bool, len(headers))

	for x := 0; x < numberOfColumns; x++ {
		columnName := headers[x]
		if !customOrNative(columnName) {
			customFields[columnName] = true
		}
	}
	return customFields
}

func csvRowToContact(r *http.Request, c context.Context, singleRow []string, headers []string) (models.Contact, error) {

}

func CsvToContactList(r *http.Request, file []byte, headers []string) ([]models.Contact, map[string]bool, error) {
	c := appengine.NewContext(r)

	readerFile := bytes.NewReader(file)
	incomingRecords := csv.NewReader(readerFile)
	records, err := incomingRecords.ReadAll()

	// Number of columns in sheet to compare
	numberOfColumns := len(records[0])
	if numberOfColumns != len(headers) {
		return []models.Contact{}, map[string]bool{}, errors.New("Number of headers does not match the ones for the sheet")
	}

	// Loop through all the rows
	// Extract information
	emptyContact := models.Contact{}
	contacts := []models.Contact{}

	for i := 0; i < len(records); i++ {
		contact, err := csvRowToContact(r, c, records[0], headers)
		if err != nil {
			return []models.Contact{}, map[string]bool{}, err
		}

		// To get rid of empty contacts. We don't want to create empty contacts.
		if !reflect.DeepEqual(emptyContact, contact) {
			contacts = append(contacts, contact)
		}
	}

	// Get custom fields
	customFields := getCustomFields(r, c, numberOfColumns, headers)

	return contacts, customFields, nil
}

func CsvFileToExcelHeader(r *http.Request, file []byte) ([]Column, error) {
	c := appengine.NewContext(r)

	readerFile := bytes.NewReader(file)
	incomingRecords := csv.NewReader(readerFile)
	records, err := incomingRecords.ReadAll()
	if err != nil {
		log.Errorf(c, "%v", err)
	}

	log.Infof(c, "%v", records)

	numberOfRows := 15
	if len(records) < numberOfRows+1 {
		numberOfRows = len(records)
	}

	numberOfColumns := len(records[0])
	columns := make([]Column, numberOfColumns)

	for _, row := range records[0:numberOfRows] {
		for currentColumn, cell := range row {
			columns[currentColumn].Rows = append(columns[currentColumn].Rows, strings.Trim(cell, " "))
		}
	}

	return columns, nil
}
<commit_msg>Finishing uploading csv sheets<commit_after>package goexcel

import (
	"bytes"
	"encoding/csv"
	"errors"
	"net/http"
	"reflect"
	"strings"

	"google.golang.org/appengine"
	"google.golang.org/appengine/log"

	"golang.org/x/net/context"

	"github.com/news-ai/tabulae/models"
)

func csvGetCustomFields(r *http.Request, c context.Context, numberOfColumns int, headers []string) map[string]bool {
	customFields := make(map[string]bool, len(headers))

	for x := 0; x < numberOfColumns; x++ {
		columnName := headers[x]
		if !customOrNative(columnName) {
			customFields[columnName] = true
		}
	}
	return customFields
}

func csvRowToContact(r *http.Request, c context.Context, singleRow []string, headers []string) (models.Contact, error) {
	var (
		contact       models.Contact
		employers     []int64
		pastEmployers []int64
		customFields  []models.CustomContactField
	)

	for x := 0; x < len(singleRow); x++ {
		columnName := headers[x]
		cellName := singleRow[x]
		rowToContact(r, c, columnName, cellName, &contact, &employers, &pastEmployers, &customFields)
	}

	contact.CustomFields = customFields
	contact.Employers = employers
	contact.PastEmployers = pastEmployers
	return contact, nil
}

func CsvToContactList(r *http.Request, file []byte, headers []string) ([]models.Contact, map[string]bool, error) {
	c := appengine.NewContext(r)

	readerFile := bytes.NewReader(file)
	incomingRecords := csv.NewReader(readerFile)
	records, err := incomingRecords.ReadAll()
	if err != nil {
		log.Infof(c, "%v", err)
		return []models.Contact{}, map[string]bool{}, err
	}

	// Number of columns in sheet to compare
	numberOfColumns := len(records[0])
	if numberOfColumns != len(headers) {
		err := errors.New("Number of headers does not match the ones for the sheet")
		log.Infof(c, "%v", err)
		return []models.Contact{}, map[string]bool{}, err
	}

	// Loop through all the rows
	// Extract information
	emptyContact := models.Contact{}
	contacts := []models.Contact{}

	for i := 0; i < len(records); i++ {
		contact, err := csvRowToContact(r, c, records[0], headers)
		if err != nil {
			log.Infof(c, "%v", err)
			return []models.Contact{}, map[string]bool{}, err
		}

		// To get rid of empty contacts. We don't want to create empty contacts.
		if !reflect.DeepEqual(emptyContact, contact) {
			contacts = append(contacts, contact)
		}
	}

	// Get custom fields
	customFields := getCustomFields(r, c, numberOfColumns, headers)

	return contacts, customFields, nil
}

func CsvFileToExcelHeader(r *http.Request, file []byte) ([]Column, error) {
	c := appengine.NewContext(r)

	readerFile := bytes.NewReader(file)
	incomingRecords := csv.NewReader(readerFile)
	records, err := incomingRecords.ReadAll()
	if err != nil {
		log.Errorf(c, "%v", err)
		return []Column{}, err
	}

	numberOfRows := 15
	if len(records) < numberOfRows+1 {
		numberOfRows = len(records)
	}

	numberOfColumns := len(records[0])
	columns := make([]Column, numberOfColumns)

	for _, row := range records[0:numberOfRows] {
		for currentColumn, cell := range row {
			columns[currentColumn].Rows = append(columns[currentColumn].Rows, strings.Trim(cell, " "))
		}
	}

	return columns, nil
}
<|endoftext|>


<commit_before>package exec

import (
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"syscall"
)

// The fork executor is a simple executor that just runs the build script in a
// a temporary directory in a forked process.

type forkExecutor struct {
	Dir       string
	ScriptURL string
	Args      []string
	Env       []string
}

func (f *forkExecutor) SaveFile(relative string, data []byte) error {
	// TODO: ensure no leaks with ..
	filename := path.Join(f.Dir, relative)
	dir, _ := filepath.Split(filename)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return err
	}

	return ioutil.WriteFile(filename, data, 0755)
}

func (f *forkExecutor) Run(filename string, stdout chan<- []byte) error {
	var err error

	if err = f.run(filename, stdout); err != nil {
		close(stdout)
		return err
	}
	return nil
}

func (f *forkExecutor) Directory() string {
	return f.Dir
}

func (f *forkExecutor) Cleanup() error {
	return os.RemoveAll(f.Dir)
}

func (f *forkExecutor) run(filename string, ch chan<- []byte) error {
	var err error
	cmd := exec.Command(filename, f.Args...)
	cmd.Dir = f.Dir
	cmd.Stdin = nil
	cmd.Env = append(os.Environ(), f.Env...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}
	cmd.Stderr = cmd.Stdout
	if err := cmd.Start(); err != nil {
		return err
	}

	buf := make([]byte, 1024)
	for {
		n, err := stdout.Read(buf)
		if n != 0 {
			ch <- buf[:n]
		}
		if err != nil {
			if err != io.EOF {
				log.Printf("Error reading stdout: %v\n", err)
			}
			break
		}
	}

	return cmd.Wait()
}

// AsUnixStatusCode takes an error and resolves Unix status code from it
func AsUnixStatusCode(err error) int {
	var exitErr *exec.ExitError
	var status syscall.WaitStatus
	var ok bool
	if err == nil {
		return 0
	}
	if exitErr, ok = err.(*exec.ExitError); !ok {
		log.Printf("Got unexpected error %v\n", err)
		return -1
	}
	if status, ok = exitErr.Sys().(syscall.WaitStatus); !ok {
		log.Println("ExitError was not a WaitStatus. Is this Unix?")
		return -2
	}
	return status.ExitStatus()
}
<commit_msg>Ensure fork runner script is relative to path<commit_after>package exec

import (
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"syscall"
)

// The fork executor is a simple executor that just runs the build script in a
// a temporary directory in a forked process.

type forkExecutor struct {
	Dir       string
	ScriptURL string
	Args      []string
	Env       []string
}

func (f *forkExecutor) SaveFile(relative string, data []byte) error {
	// TODO: ensure no leaks with ..
	filename := path.Join(f.Dir, relative)
	dir, _ := filepath.Split(filename)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return err
	}

	return ioutil.WriteFile(filename, data, 0755)
}

func (f *forkExecutor) Run(relative string, stdout chan<- []byte) error {
	filename := path.Join(f.Dir, relative)
	if err := f.run(filename, stdout); err != nil {
		close(stdout)
		return err
	}
	return nil
}

func (f *forkExecutor) Directory() string {
	return f.Dir
}

func (f *forkExecutor) Cleanup() error {
	return os.RemoveAll(f.Dir)
}

func (f *forkExecutor) run(filename string, ch chan<- []byte) error {
	var err error
	cmd := exec.Command(filename, f.Args...)
	cmd.Dir = f.Dir
	cmd.Stdin = nil
	cmd.Env = append(os.Environ(), f.Env...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}
	cmd.Stderr = cmd.Stdout
	if err := cmd.Start(); err != nil {
		return err
	}

	buf := make([]byte, 1024)
	for {
		n, err := stdout.Read(buf)
		if n != 0 {
			ch <- buf[:n]
		}
		if err != nil {
			if err != io.EOF {
				log.Printf("Error reading stdout: %v\n", err)
			}
			break
		}
	}

	return cmd.Wait()
}

// AsUnixStatusCode takes an error and resolves Unix status code from it
func AsUnixStatusCode(err error) int {
	var exitErr *exec.ExitError
	var status syscall.WaitStatus
	var ok bool
	if err == nil {
		return 0
	}
	if exitErr, ok = err.(*exec.ExitError); !ok {
		log.Printf("Got unexpected error %v\n", err)
		return -1
	}
	if status, ok = exitErr.Sys().(syscall.WaitStatus); !ok {
		log.Println("ExitError was not a WaitStatus. Is this Unix?")
		return -2
	}
	return status.ExitStatus()
}
<|endoftext|>


<commit_before>package requesthelpers

import (
	"bytes"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"net/textproto"
	"path/filepath"
	"strings"
)

type ReqHelperFaxFile struct {
	metadata []byte
	filename string
	writer   *multipart.Writer
	headers  http.Header
	body     []byte
}

func NewReqHelperFaxFile(metadata []byte, filename string) ReqHelperFaxFile {
	fax := ReqHelperFaxFile{}
	fax.metadata = metadata
	fax.filename = filename
	fax.build()
	return fax
}

func (fax *ReqHelperFaxFile) build() {
	body := &bytes.Buffer{}
	fax.writer = multipart.NewWriter(body)
	fax.SetMetadata(fax.metadata)
	fax.SetFile(fax.filename)
	fax.writer.Close()
	fax.body = body.Bytes()
	fax.setHeaders()
}

func (fax *ReqHelperFaxFile) SetMetadata(metadata []byte) {
	// MIME HEADER
	metaHead := textproto.MIMEHeader{}
	metaHead.Add("Content-Type", "application/json")

	// MIME PART
	metaPart, _ := fax.writer.CreatePart(metaHead)
	metaPart.Write(metadata)
}

func (fax *ReqHelperFaxFile) SetFile(path string) error {
	// READ FILE
	bytes, err := ioutil.ReadFile(path)
	if err != nil {
		return err
	}

	// BODY HEAD
	_, filename := filepath.Split(path)

	bodyHead := textproto.MIMEHeader{}
	bodyHead.Add("Content-Type", "application/octet-stream")
	if len(filename) > 0 {
		bodyHead.Add("Content-Disposition", strings.Join([]string{`attachment; filename="`, filename, `"`}, ""))
	} else {
		bodyHead.Add("Content-Disposition", "attachment")
	}

	// BODY PART
	bodyPart, _ := fax.writer.CreatePart(bodyHead)
	bodyPart.Write(bytes)
	return nil
}

func (fax *ReqHelperFaxFile) setHeaders() {
	fax.headers = http.Header{}
	fax.headers.Add("Content-Type", strings.Join([]string{"multipart/mixed; boundary=", fax.writer.Boundary()}, ""))
}

func (fax *ReqHelperFaxFile) GetHeaders() http.Header {
	return fax.headers
}

func (fax *ReqHelperFaxFile) GetBody() []byte {
	return fax.body
}
<commit_msg>variable name update<commit_after>package requesthelpers

import (
	"bytes"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"net/textproto"
	"path/filepath"
	"strings"
)

type ReqHelperFaxFile struct {
	metadata []byte
	filename string
	writer   *multipart.Writer
	headers  http.Header
	body     []byte
}

func NewReqHelperFaxFile(metadata []byte, filename string) ReqHelperFaxFile {
	fax := ReqHelperFaxFile{}
	fax.metadata = metadata
	fax.filename = filename
	fax.build()
	return fax
}

func (fax *ReqHelperFaxFile) build() {
	body := &bytes.Buffer{}
	fax.writer = multipart.NewWriter(body)
	fax.SetMetadata(fax.metadata)
	fax.SetFile(fax.filename)
	fax.writer.Close()
	fax.body = body.Bytes()
	fax.setHeaders()
}

func (fax *ReqHelperFaxFile) SetMetadata(metadata []byte) {
	// MIME HEADER
	metaHead := textproto.MIMEHeader{}
	metaHead.Add("Content-Type", "application/json")

	// MIME PART
	metaPart, _ := fax.writer.CreatePart(metaHead)
	metaPart.Write(metadata)
}

func (fax *ReqHelperFaxFile) SetFile(path string) error {
	// READ FILE
	bytes, err := ioutil.ReadFile(path)
	if err != nil {
		return err
	}

	// FILE HEAD
	_, filename := filepath.Split(path)

	fileHead := textproto.MIMEHeader{}
	fileHead.Add("Content-Type", "application/octet-stream")
	if len(filename) > 0 {
		fileHead.Add("Content-Disposition", strings.Join([]string{`attachment; filename="`, filename, `"`}, ""))
	} else {
		fileHead.Add("Content-Disposition", "attachment")
	}

	// FILE PART
	filePart, _ := fax.writer.CreatePart(fileHead)
	filePart.Write(bytes)
	return nil
}

func (fax *ReqHelperFaxFile) setHeaders() {
	fax.headers = http.Header{}
	fax.headers.Add("Content-Type", strings.Join([]string{"multipart/mixed; boundary=", fax.writer.Boundary()}, ""))
}

func (fax *ReqHelperFaxFile) GetHeaders() http.Header {
	return fax.headers
}

func (fax *ReqHelperFaxFile) GetBody() []byte {
	return fax.body
}
<|endoftext|>


<commit_before>package disgo

import (
	"bytes"
	"compress/zlib"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/slf4go/logger"
)

type shard struct {
	// Parent session
	session *Session

	// Websocket information used for identification and reconnection
	webSocket *websocket.Conn
	shard     int
	sessionID string
	sequence  int
	heartbeat int

	// Mutex locks, reconnect to make sure there is only 1 process reconnecting and concurrent read/write accesses on the socket
	reconnectLock       sync.Mutex
	readLock, writeLock sync.Mutex
	isReconnecting      bool

	// Channels to pass around messages
	closeMessage chan int
	stopListen   chan bool
	stopRead     chan bool
}

func connectShard(session *Session, shardNum int) (*shard, error) {
	conn, _, err := websocket.DefaultDialer.Dial(session.wsUrl, http.Header{})
	if err != nil {
		return nil, err
	}

	s := &shard{
		session:      session,
		webSocket:    conn,
		shard:        shardNum,
		closeMessage: make(chan int, 1),
		stopListen:   make(chan bool),
		stopRead:     make(chan bool),
	}

	if err = s.handshake(); err != nil {
		return nil, err
	}

	return s, nil
}

func (s *shard) handshake() error {
	s.webSocket.SetCloseHandler(s.onClose)

	helloFrame, err := s.readFrame()
	if err != nil {
		return err
	} else if helloFrame.Op != opHello {
		return errors.New("First frame sent from the Discord Gateway is not hello")
	}

	hello := helloPayload{}
	if err = json.Unmarshal(helloFrame.Data, &hello); err != nil {
		return err
	}

	logger.Debugf("Connected to Discord servers: %s", strings.Join(hello.Servers, ", "))
	logger.Debugf("Setting up a heartbeat interval of %d ms", hello.HeartbeatInterval)
	s.heartbeat = hello.HeartbeatInterval

	if err = s.identify(); err != nil {
		return err
	}

	go s.mainLoop()
	return nil
}

func (s *shard) identify() error {
	if s.sessionID != "" {
		logger.Debugf("Resuming connection starting at sequence %d.", s.sequence)
		s.sendFrame(&gatewayFrame{opResume, resumePayload{
			Token:     s.session.token,
			SessionID: s.sessionID,
			Sequence:  s.sequence,
		}})
	} else {
		logger.Debugf("Identifying to websocket")
		s.sendFrame(&gatewayFrame{opIdentify, identifyPayload{
			Token:          s.session.token,
			Compress:       true,
			LargeThreshold: 250,
			Shard:          [2]int{s.shard, cap(s.session.shards)},
			Properties: propertiesPayload{
				OS:      runtime.GOOS,
				Browser: "DisGo",
				Device:  "DisGo",
			},
		}})
	}

	for {
		frame, err := s.readFrame()
		if err != nil {
			return err
		}

		if frame.Op == opDispatch {
			switch frame.EventName {
			case "READY":
				ready := ReadyEvent{}
				if err := json.Unmarshal(frame.Data, &ready); err != nil {
					return err
				}

				s.sessionID = ready.SessionID

				fallthrough
			case "RESUMED":
				s.session.dispatchEvent(frame)
				return nil // Break out of the loop, we have what we want
			default:
				s.session.dispatchEvent(frame) // Nope, resume action, let's wait for more frames
			}
		} else if frame.Op == opInvalidSession {
			s.sessionID = "" // Invalidate session and retry
			s.identify()
			return nil
		} else {
			return fmt.Errorf("Unexpected opCode received from Discord: %d", frame.Op)
		}
	}
}

func (s *shard) mainLoop() {
	logger.Debugf("Starting main loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))
	defer logger.Debugf("Exiting main loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))

	heartbeat := time.NewTicker(time.Duration(s.heartbeat) * time.Millisecond)
	defer heartbeat.Stop()
	sentHeartBeat := false

	reader := make(chan *receivedFrame, 1)
	go s.readWebSocket(reader)

	for {
		select {
		case <-heartbeat.C:
			if !sentHeartBeat {
				go s.sendFrame(&gatewayFrame{opHeartbeat, s.sequence})
				sentHeartBeat = true
			} else {
				go s.disconnect(websocket.CloseAbnormalClosure, "Did not respond to previous heartbeat")
			}
		case <-s.stopListen:
			return
		case frame := <-reader:
			switch frame.Op {
			case opHeartbeat:
				s.sendFrame(&gatewayFrame{Op: opHeartbeatAck})
			case opHeartbeatAck:
				sentHeartBeat = false
			case opReconnect:
				s.disconnect(websocket.CloseNormalClosure, "op Reconnect")
			case opDispatch:
				s.session.dispatchEvent(frame)
			default:
				logger.Errorf("Unexpected opCode received: %d", frame.Op)
			}
		}
	}
}

func (s *shard) sendFrame(frame *gatewayFrame) {
	// Resume will only be sent within the reconnect lock, we don't relock to prevent a deadlock
	if frame.Op != opResume {
		s.reconnectLock.Lock()
		defer s.reconnectLock.Unlock()
	}

	s.writeLock.Lock()
	defer s.writeLock.Unlock()

	logger.Debugf("Sending frame with opCode: %d", frame.Op)
	s.webSocket.WriteJSON(frame)
}

func (s *shard) readWebSocket(reader chan *receivedFrame) {
	logger.Debugf("Starting read loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))
	defer logger.Debugf("Exiting read loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))

	for {
		select {
		case <-s.stopRead:
			return
		default:
			frame, err := s.readFrame()
			if err != nil {
				if !s.session.isShuttingDown() {
					if closeErr, isClose := err.(*websocket.CloseError); isClose {
						s.closeMessage <- closeErr.Code
					}
					logger.ErrorE(err)
					go s.disconnect(websocket.CloseAbnormalClosure, err.Error())
				}
				<-s.stopRead // Wait for the connection to be closed
				return
			}

			reader <- frame
		}
	}

}

func (s *shard) readFrame() (*receivedFrame, error) {
	s.readLock.Lock()
	defer s.readLock.Unlock()

	logger.Tracef("Shard.readFrame() called")
	msgType, msg, err := s.webSocket.ReadMessage()
	if err != nil {
		return nil, err
	}

	var reader io.Reader
	reader = bytes.NewBuffer(msg)

	if msgType == websocket.BinaryMessage {
		zReader, err := zlib.NewReader(reader)
		if err != nil {
			return nil, err
		}

		defer zReader.Close()

		reader = zReader
	}

	frame := receivedFrame{Sequence: -1}
	json.NewDecoder(reader).Decode(&frame)

	logger.Debugf("Received frame with opCode: %d", frame.Op)

	if frame.Sequence != -1 {
		logger.Tracef("Last sequence received set to %d.", frame.Sequence)
		s.sequence = frame.Sequence
	}
	return &frame, nil
}

func (s *shard) startReconnect() {
	if !s.isReconnecting {
		s.reconnectLock.Lock()
		s.isReconnecting = true
		go s.reconnect()
	}
}

func (s *shard) reconnect() {
	if !s.session.isShuttingDown() {
		logger.Noticef("Reconnecting shard [%d/%d]", s.shard+1, cap(s.session.shards))
		conn, _, err := websocket.DefaultDialer.Dial(s.session.wsUrl, http.Header{})

		if err == nil {
			s.webSocket = conn
			err = s.handshake()

			if err != nil {
				conn.Close()
			}
		}

		if err != nil {
			logger.ErrorE(err)
			time.Sleep(1 * time.Second)
			go s.reconnect()
		} else {
			s.isReconnecting = false
			s.reconnectLock.Unlock()
		}
	}
}

func (s *shard) onClose(code int, text string) error {
	logger.Debugf("Received Close Frame from Discord. Code: %d. Text: %s", code, text)
	s.closeMessage <- code
	s.startReconnect()
	return nil
}

func (s *shard) disconnect(code int, text string) {
	logger.Tracef("Shard.disconnect() called")
	s.stopListen <- true

	s.writeLock.Lock()
	err := s.webSocket.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(code, text))
	s.writeLock.Unlock()
	if err != nil {
		logger.ErrorE(err)
	}

	select {
	case code := <-s.closeMessage:
		logger.Debugf("Discord connection closed with code %d", code)
	case <-time.After(1 * time.Second):
		logger.Warn("Discord did not reply to the close message, force-closing connection.")
	}

	s.webSocket.Close()
	s.stopRead <- true

	s.startReconnect()
}
<commit_msg>Locked s.isReconnecting<commit_after>package disgo

import (
	"bytes"
	"compress/zlib"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/slf4go/logger"
)

type shard struct {
	// Parent session
	session *Session

	// Websocket information used for identification and reconnection
	webSocket *websocket.Conn
	shard     int
	sessionID string
	sequence  int
	heartbeat int

	// Mutex locks, reconnect to make sure there is only 1 process reconnecting and concurrent read/write accesses on the socket
	readLock, writeLock sync.Mutex
	reconnectLock       sync.RWMutex
	isReconnecting      bool

	// Channels to pass around messages
	closeMessage chan int
	stopListen   chan bool
	stopRead     chan bool
}

func connectShard(session *Session, shardNum int) (*shard, error) {
	conn, _, err := websocket.DefaultDialer.Dial(session.wsUrl, http.Header{})
	if err != nil {
		return nil, err
	}

	s := &shard{
		session:      session,
		webSocket:    conn,
		shard:        shardNum,
		closeMessage: make(chan int, 1),
		stopListen:   make(chan bool),
		stopRead:     make(chan bool),
	}

	if err = s.handshake(); err != nil {
		return nil, err
	}

	return s, nil
}

func (s *shard) handshake() error {
	s.webSocket.SetCloseHandler(s.onClose)

	helloFrame, err := s.readFrame()
	if err != nil {
		return err
	} else if helloFrame.Op != opHello {
		return errors.New("First frame sent from the Discord Gateway is not hello")
	}

	hello := helloPayload{}
	if err = json.Unmarshal(helloFrame.Data, &hello); err != nil {
		return err
	}

	logger.Debugf("Connected to Discord servers: %s", strings.Join(hello.Servers, ", "))
	logger.Debugf("Setting up a heartbeat interval of %d ms", hello.HeartbeatInterval)
	s.heartbeat = hello.HeartbeatInterval

	if err = s.identify(); err != nil {
		return err
	}

	go s.mainLoop()
	return nil
}

func (s *shard) identify() error {
	if s.sessionID != "" {
		logger.Debugf("Resuming connection starting at sequence %d.", s.sequence)
		s.sendFrame(&gatewayFrame{opResume, resumePayload{
			Token:     s.session.token,
			SessionID: s.sessionID,
			Sequence:  s.sequence,
		}})
	} else {
		logger.Debugf("Identifying to websocket")
		s.sendFrame(&gatewayFrame{opIdentify, identifyPayload{
			Token:          s.session.token,
			Compress:       true,
			LargeThreshold: 250,
			Shard:          [2]int{s.shard, cap(s.session.shards)},
			Properties: propertiesPayload{
				OS:      runtime.GOOS,
				Browser: "DisGo",
				Device:  "DisGo",
			},
		}})
	}

	for {
		frame, err := s.readFrame()
		if err != nil {
			return err
		}

		if frame.Op == opDispatch {
			switch frame.EventName {
			case "READY":
				ready := ReadyEvent{}
				if err := json.Unmarshal(frame.Data, &ready); err != nil {
					return err
				}

				s.sessionID = ready.SessionID

				fallthrough
			case "RESUMED":
				s.session.dispatchEvent(frame)
				return nil // Break out of the loop, we have what we want
			default:
				s.session.dispatchEvent(frame) // Nope, resume action, let's wait for more frames
			}
		} else if frame.Op == opInvalidSession {
			s.sessionID = "" // Invalidate session and retry
			s.identify()
			return nil
		} else {
			return fmt.Errorf("Unexpected opCode received from Discord: %d", frame.Op)
		}
	}
}

func (s *shard) mainLoop() {
	logger.Debugf("Starting main loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))
	defer logger.Debugf("Exiting main loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))

	heartbeat := time.NewTicker(time.Duration(s.heartbeat) * time.Millisecond)
	defer heartbeat.Stop()
	sentHeartBeat := false

	reader := make(chan *receivedFrame, 1)
	go s.readWebSocket(reader)

	for {
		select {
		case <-heartbeat.C:
			if !sentHeartBeat {
				go s.sendFrame(&gatewayFrame{opHeartbeat, s.sequence})
				sentHeartBeat = true
			} else {
				go s.disconnect(websocket.CloseAbnormalClosure, "Did not respond to previous heartbeat")
			}
		case <-s.stopListen:
			return
		case frame := <-reader:
			switch frame.Op {
			case opHeartbeat:
				s.sendFrame(&gatewayFrame{Op: opHeartbeatAck})
			case opHeartbeatAck:
				sentHeartBeat = false
			case opReconnect:
				s.disconnect(websocket.CloseNormalClosure, "op Reconnect")
			case opDispatch:
				s.session.dispatchEvent(frame)
			default:
				logger.Errorf("Unexpected opCode received: %d", frame.Op)
			}
		}
	}
}

func (s *shard) sendFrame(frame *gatewayFrame) {
	// Resume will only be sent within the reconnect lock, we don't relock to prevent a deadlock
	if frame.Op != opResume {
		s.reconnectLock.Lock()
		defer s.reconnectLock.Unlock()
	}

	s.writeLock.Lock()
	defer s.writeLock.Unlock()

	logger.Debugf("Sending frame with opCode: %d", frame.Op)
	s.webSocket.WriteJSON(frame)
}

func (s *shard) readWebSocket(reader chan *receivedFrame) {
	logger.Debugf("Starting read loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))
	defer logger.Debugf("Exiting read loop for shard [%d/%d]", s.shard+1, cap(s.session.shards))

	for {
		select {
		case <-s.stopRead:
			return
		default:
			frame, err := s.readFrame()
			if err != nil {
				if !s.session.isShuttingDown() {
					if closeErr, isClose := err.(*websocket.CloseError); isClose {
						s.closeMessage <- closeErr.Code
					}
					logger.ErrorE(err)
					go s.disconnect(websocket.CloseAbnormalClosure, err.Error())
				}
				<-s.stopRead // Wait for the connection to be closed
				return
			}

			reader <- frame
		}
	}

}

func (s *shard) readFrame() (*receivedFrame, error) {
	s.readLock.Lock()
	defer s.readLock.Unlock()

	logger.Tracef("Shard.readFrame() called")
	msgType, msg, err := s.webSocket.ReadMessage()
	if err != nil {
		return nil, err
	}

	var reader io.Reader
	reader = bytes.NewBuffer(msg)

	if msgType == websocket.BinaryMessage {
		zReader, err := zlib.NewReader(reader)
		if err != nil {
			return nil, err
		}

		defer zReader.Close()

		reader = zReader
	}

	frame := receivedFrame{Sequence: -1}
	json.NewDecoder(reader).Decode(&frame)

	logger.Debugf("Received frame with opCode: %d", frame.Op)

	if frame.Sequence != -1 {
		logger.Tracef("Last sequence received set to %d.", frame.Sequence)
		s.sequence = frame.Sequence
	}
	return &frame, nil
}

func (s *shard) startReconnect() {
	s.reconnectLock.RLock()
	if s.isReconnecting {
		return
	}
	s.reconnectLock.RUnlock()

	s.reconnectLock.Lock()
	s.isReconnecting = true
	go s.reconnect()
}

func (s *shard) reconnect() {
	if !s.session.isShuttingDown() {
		logger.Noticef("Reconnecting shard [%d/%d]", s.shard+1, cap(s.session.shards))
		conn, _, err := websocket.DefaultDialer.Dial(s.session.wsUrl, http.Header{})

		if err == nil {
			s.webSocket = conn
			err = s.handshake()

			if err != nil {
				conn.Close()
			}
		}

		if err != nil {
			logger.ErrorE(err)
			time.Sleep(1 * time.Second)
			go s.reconnect()
		} else {
			s.isReconnecting = false
			s.reconnectLock.Unlock()
		}
	}
}

func (s *shard) onClose(code int, text string) error {
	logger.Debugf("Received Close Frame from Discord. Code: %d. Text: %s", code, text)
	s.closeMessage <- code
	s.startReconnect()
	return nil
}

func (s *shard) disconnect(code int, text string) {
	logger.Tracef("Shard.disconnect() called")
	s.stopListen <- true

	s.writeLock.Lock()
	err := s.webSocket.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(code, text))
	s.writeLock.Unlock()
	if err != nil {
		logger.ErrorE(err)
	}

	select {
	case code := <-s.closeMessage:
		logger.Debugf("Discord connection closed with code %d", code)
	case <-time.After(1 * time.Second):
		logger.Warn("Discord did not reply to the close message, force-closing connection.")
	}

	s.webSocket.Close()
	s.stopRead <- true

	s.startReconnect()
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"sync"
)

// DirStore is a KVStore backed by a directory.
type DirStore struct {
	Dir string
}

func (ds DirStore) path(key []byte) string {
	return filepath.Join(ds.Dir, hex.EncodeToString(key))
}

// Set sets key-value pair.
func (ds DirStore) Set(key, value []byte) error {
	return ioutil.WriteFile(ds.path(key), value, 0600)
}

// Get gets a value with the given key.
// If the KV pair is not set, returns ErrNotExist.
func (ds DirStore) Get(key []byte) ([]byte, error) {
	dat, err := ioutil.ReadFile(ds.path(key))
	if err != nil {
		if os.IsNotExist(err) {
			err = ErrNotExist
		}
		return nil, err
	}
	return dat, err
}

// MemStore is an in-memory KVStore.
type MemStore struct {
	m sync.Map
}

// Set sets key-value pair.
func (ms *MemStore) Set(key, value []byte) error {
	ms.m.Store(hex.EncodeToString(key), value)
	return nil
}

// Get gets a value with the given key.
// If the KV pair is not set, returns ErrNotExist.
func (ms *MemStore) Get(key []byte) ([]byte, error) {
	dat, ok := ms.m.Load(hex.EncodeToString(key))
	if !ok {
		return nil, ErrNotExist
	}
	return dat.([]byte), nil
}

// ErrNotExist is an error indicating that a KV pair is not set.
var ErrNotExist = errors.New("kv pair does not exist")

// KVStore is an interface for a key-value store.
type KVStore interface {
	// Set sets key-value pair.
	Set(key, value []byte) error

	// Get gets a value with the given key.
	// If the KV pair is not set, returns ErrNotExist.
	Get(key []byte) ([]byte, error)
}

// Code is a struct containing code with metadata.
type Code struct {
	Code     string `json:"code"`
	Language string `json:"language"`
}

// CodeStore is a code storage system using a KVStore.
type CodeStore struct {
	KV KVStore
}

// Get retrieves a Code struct from the store.
func (cs CodeStore) Get(key string) (Code, error) {
	// decode key
	k, err := hex.DecodeString(key)
	if err != nil {
		return Code{}, err
	}

	// get code from store
	dat, err := cs.KV.Get(k)
	if err != nil {
		return Code{}, err
	}

	// unmarshal code
	var c Code
	err = json.Unmarshal(dat, &c)
	if err != nil {
		return Code{}, err
	}

	return c, nil
}

// Store stores Code into tha KVStore.
func (cs CodeStore) Store(c Code) (string, error) {
	// encode Code
	dat, err := json.Marshal(&c)
	if err != nil {
		return "", err
	}

	// hash Code to generate key
	hash := sha256.Sum256(dat)

	// save in KVStore
	err = cs.KV.Set(hash[:], dat)
	if err != nil {
		return "", err
	}

	// encode hash key into text format
	return hex.EncodeToString(hash[:]), nil
}

func main() {
	var driver string
	var dir string
	flag.StringVar(&driver, "driver", "mem", "driver for key-value store")
	flag.StringVar(&dir, "dir", "", "directory to use for dir driver")
	flag.Parse()

	// initialize KVStore
	var kv KVStore
	switch driver {
	case "mem":
		kv = new(MemStore)
	case "dir":
		kv = DirStore{dir}
	default:
		panic(fmt.Errorf("unrecognized driver %s", driver))
	}

	cs := CodeStore{kv}

	// store Code
	http.HandleFunc("/store", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			http.Error(w, "method not supported", http.StatusMethodNotAllowed)
			return
		}

		var c Code
		err := json.NewDecoder(r.Body).Decode(&c)
		if err != nil {
			http.Error(w, fmt.Sprintf("failed to decode request: %s", err.Error()), http.StatusBadRequest)
			return
		}

		key, err := cs.Store(c)
		if err != nil {
			http.Error(w, fmt.Sprintf("Failed to store: %s", err.Error()), http.StatusInternalServerError)
			return
		}

		// write back key
		w.Header().Add("Content-Type", "text/plain")
		w.Write([]byte(key))
	})

	// load Code
	http.HandleFunc("/load", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodGet {
			http.Error(w, "method not supported", http.StatusMethodNotAllowed)
			return
		}

		key := r.URL.Query().Get("key")
		c, err := cs.Get(key)
		if err != nil {
			http.Error(w, fmt.Sprintf("failed to load: %s", err.Error()), http.StatusInternalServerError)
			return
		}
		w.Header().Add("Content-Type", "application/json")
		json.NewEncoder(w).Encode(c)
	})

	panic(http.ListenAndServe(":80", nil))
}
<commit_msg>implement etag-based caching in store<commit_after>package main

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"sync"
)

// DirStore is a KVStore backed by a directory.
type DirStore struct {
	Dir string
}

func (ds DirStore) path(key []byte) string {
	return filepath.Join(ds.Dir, hex.EncodeToString(key))
}

// Set sets key-value pair.
func (ds DirStore) Set(key, value []byte) error {
	return ioutil.WriteFile(ds.path(key), value, 0600)
}

// Get gets a value with the given key.
// If the KV pair is not set, returns ErrNotExist.
func (ds DirStore) Get(key []byte) ([]byte, error) {
	dat, err := ioutil.ReadFile(ds.path(key))
	if err != nil {
		if os.IsNotExist(err) {
			err = ErrNotExist
		}
		return nil, err
	}
	return dat, err
}

// MemStore is an in-memory KVStore.
type MemStore struct {
	m sync.Map
}

// Set sets key-value pair.
func (ms *MemStore) Set(key, value []byte) error {
	ms.m.Store(hex.EncodeToString(key), value)
	return nil
}

// Get gets a value with the given key.
// If the KV pair is not set, returns ErrNotExist.
func (ms *MemStore) Get(key []byte) ([]byte, error) {
	dat, ok := ms.m.Load(hex.EncodeToString(key))
	if !ok {
		return nil, ErrNotExist
	}
	return dat.([]byte), nil
}

// ErrNotExist is an error indicating that a KV pair is not set.
var ErrNotExist = errors.New("kv pair does not exist")

// KVStore is an interface for a key-value store.
type KVStore interface {
	// Set sets key-value pair.
	Set(key, value []byte) error

	// Get gets a value with the given key.
	// If the KV pair is not set, returns ErrNotExist.
	Get(key []byte) ([]byte, error)
}

// Code is a struct containing code with metadata.
type Code struct {
	Code     string `json:"code"`
	Language string `json:"language"`
}

// CodeStore is a code storage system using a KVStore.
type CodeStore struct {
	KV KVStore
}

// Get retrieves a Code struct from the store.
func (cs CodeStore) Get(key string) (Code, error) {
	// decode key
	k, err := hex.DecodeString(key)
	if err != nil {
		return Code{}, err
	}

	// get code from store
	dat, err := cs.KV.Get(k)
	if err != nil {
		return Code{}, err
	}

	// unmarshal code
	var c Code
	err = json.Unmarshal(dat, &c)
	if err != nil {
		return Code{}, err
	}

	return c, nil
}

// Store stores Code into tha KVStore.
func (cs CodeStore) Store(c Code) (string, error) {
	// encode Code
	dat, err := json.Marshal(&c)
	if err != nil {
		return "", err
	}

	// hash Code to generate key
	hash := sha256.Sum256(dat)

	// save in KVStore
	err = cs.KV.Set(hash[:], dat)
	if err != nil {
		return "", err
	}

	// encode hash key into text format
	return hex.EncodeToString(hash[:]), nil
}

func main() {
	var driver string
	var dir string
	flag.StringVar(&driver, "driver", "mem", "driver for key-value store")
	flag.StringVar(&dir, "dir", "", "directory to use for dir driver")
	flag.Parse()

	// initialize KVStore
	var kv KVStore
	switch driver {
	case "mem":
		kv = new(MemStore)
	case "dir":
		kv = DirStore{dir}
	default:
		panic(fmt.Errorf("unrecognized driver %s", driver))
	}

	cs := CodeStore{kv}

	// store Code
	http.HandleFunc("/store", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			http.Error(w, "method not supported", http.StatusMethodNotAllowed)
			return
		}

		var c Code
		err := json.NewDecoder(r.Body).Decode(&c)
		if err != nil {
			http.Error(w, fmt.Sprintf("failed to decode request: %s", err.Error()), http.StatusBadRequest)
			return
		}

		key, err := cs.Store(c)
		if err != nil {
			http.Error(w, fmt.Sprintf("Failed to store: %s", err.Error()), http.StatusInternalServerError)
			return
		}

		// write back key
		w.Header().Add("Content-Type", "text/plain")
		w.Write([]byte(key))
	})

	// load Code
	http.HandleFunc("/load", func(w http.ResponseWriter, r *http.Request) {
		// check method
		if r.Method != http.MethodGet {
			http.Error(w, "method not supported", http.StatusMethodNotAllowed)
			return
		}

		// get key
		key := r.URL.Query().Get("key")

		// handle ETag caching
		if etag := r.Header.Get("If-None-Match"); etag != "" && etag == key {
			w.WriteHeader(http.StatusNotModified)
			return
		}

		// run KV lookup
		c, err := cs.Get(key)
		if err != nil {
			http.Error(w, fmt.Sprintf("failed to load: %s", err.Error()), http.StatusInternalServerError)
			return
		}

		// set headers
		w.Header().Add("ETag", key)
		w.Header().Add("Content-Type", "application/json")

		// send response
		json.NewEncoder(w).Encode(c)
	})

	panic(http.ListenAndServe(":80", nil))
}
<|endoftext|>


<commit_before>package lib

import (
	"time"

	"code.google.com/p/go-uuid/uuid"

	"github.com/mitchellh/multistep"
	"github.com/travis-ci/worker/lib/backend"
	"golang.org/x/net/context"
)

// InstanceStartTimeout is the timeout for starting an instance and waiting for
// it to be available.
var InstanceStartTimeout = 300 * time.Second

// JobTimeout is the maximum time a job can take before it times out.
var JobTimeout = 50 * time.Minute

// A Processor will process build jobs on a channel, one by one, until it is
// told to shut down or the channel of build jobs closes.
type Processor struct {
	processorUUID uuid.UUID

	ctx           context.Context
	buildJobsChan <-chan Job
	provider      backend.Provider
	generator     BuildScriptGenerator

	graceful  chan struct{}
	terminate context.CancelFunc
}

// NewProcessor creates a new processor that will run the build jobs on the
// given channel using the given provider and getting build scripts from the
// generator.
func NewProcessor(ctx context.Context, buildJobsChan <-chan Job, provider backend.Provider, generator BuildScriptGenerator) *Processor {
	processorUUID := uuid.NewRandom()

	ctx, cancel := context.WithCancel(contextFromProcessor(ctx, processorUUID.String()))

	return &Processor{
		processorUUID: processorUUID,

		ctx:           contextFromProcessor(ctx, processorUUID.String()),
		buildJobsChan: buildJobsChan,
		provider:      provider,
		generator:     generator,

		graceful:  make(chan struct{}),
		terminate: cancel,
	}
}

// Run starts the processor. This method will not return until the processor is
// terminated, either by calling the GracefulShutdown or Terminate methods, or
// if the build jobs channel is closed.
func (p *Processor) Run() {
	for {
		select {
		case <-p.graceful:
			LoggerFromContext(p.ctx).Info("processor is done, terminating")
			return
		case buildJob, ok := <-p.buildJobsChan:
			if !ok {
				return
			}
			p.process(contextFromJob(p.ctx, buildJob), buildJob)
		}
	}
}

// GracefulShutdown tells the processor to finish the job it is currently
// processing, but not pick up any new jobs. This method will return
// immediately, the processor is done when Run() returns.
func (p *Processor) GracefulShutdown() {
	LoggerFromContext(p.ctx).Info("processor initiating graceful shutdown")
	close(p.graceful)
}

// Terminate tells the processor to stop working on the current job as soon as
// possible.
func (p *Processor) Terminate() {
	p.terminate()
}

func (p *Processor) process(ctx context.Context, buildJob Job) {
	state := new(multistep.BasicStateBag)
	state.Put("buildJob", buildJob)
	state.Put("ctx", ctx)

	steps := []multistep.Step{
		&stepGenerateScript{generator: p.generator},
		&stepStartInstance{provider: p.provider},
		&stepUploadScript{},
		&stepRunScript{},
	}

	runner := &multistep.BasicRunner{Steps: steps}

	runner.Run(state)
}
<commit_msg>Log when job starts and ends<commit_after>package lib

import (
	"time"

	"code.google.com/p/go-uuid/uuid"

	"github.com/mitchellh/multistep"
	"github.com/travis-ci/worker/lib/backend"
	"golang.org/x/net/context"
)

// InstanceStartTimeout is the timeout for starting an instance and waiting for
// it to be available.
var InstanceStartTimeout = 300 * time.Second

// JobTimeout is the maximum time a job can take before it times out.
var JobTimeout = 50 * time.Minute

// A Processor will process build jobs on a channel, one by one, until it is
// told to shut down or the channel of build jobs closes.
type Processor struct {
	processorUUID uuid.UUID

	ctx           context.Context
	buildJobsChan <-chan Job
	provider      backend.Provider
	generator     BuildScriptGenerator

	graceful  chan struct{}
	terminate context.CancelFunc
}

// NewProcessor creates a new processor that will run the build jobs on the
// given channel using the given provider and getting build scripts from the
// generator.
func NewProcessor(ctx context.Context, buildJobsChan <-chan Job, provider backend.Provider, generator BuildScriptGenerator) *Processor {
	processorUUID := uuid.NewRandom()

	ctx, cancel := context.WithCancel(contextFromProcessor(ctx, processorUUID.String()))

	return &Processor{
		processorUUID: processorUUID,

		ctx:           contextFromProcessor(ctx, processorUUID.String()),
		buildJobsChan: buildJobsChan,
		provider:      provider,
		generator:     generator,

		graceful:  make(chan struct{}),
		terminate: cancel,
	}
}

// Run starts the processor. This method will not return until the processor is
// terminated, either by calling the GracefulShutdown or Terminate methods, or
// if the build jobs channel is closed.
func (p *Processor) Run() {
	for {
		select {
		case <-p.graceful:
			LoggerFromContext(p.ctx).Info("processor is done, terminating")
			return
		case buildJob, ok := <-p.buildJobsChan:
			if !ok {
				return
			}
			p.process(contextFromJob(p.ctx, buildJob), buildJob)
		}
	}
}

// GracefulShutdown tells the processor to finish the job it is currently
// processing, but not pick up any new jobs. This method will return
// immediately, the processor is done when Run() returns.
func (p *Processor) GracefulShutdown() {
	LoggerFromContext(p.ctx).Info("processor initiating graceful shutdown")
	close(p.graceful)
}

// Terminate tells the processor to stop working on the current job as soon as
// possible.
func (p *Processor) Terminate() {
	p.terminate()
}

func (p *Processor) process(ctx context.Context, buildJob Job) {
	state := new(multistep.BasicStateBag)
	state.Put("buildJob", buildJob)
	state.Put("ctx", ctx)

	steps := []multistep.Step{
		&stepGenerateScript{generator: p.generator},
		&stepStartInstance{provider: p.provider},
		&stepUploadScript{},
		&stepRunScript{},
	}

	runner := &multistep.BasicRunner{Steps: steps}

	LoggerFromContext(p.ctx).Info("starting job")
	runner.Run(state)
	LoggerFromContext(p.ctx).Info("finished job")
}
<|endoftext|>


<commit_before>package lib

import (
	"time"

	"code.google.com/p/go-uuid/uuid"

	"github.com/mitchellh/multistep"
	"github.com/travis-ci/worker/lib/backend"
	"golang.org/x/net/context"
)

// InstanceStartTimeout is the timeout for starting an instance and waiting for
// it to be available.
var InstanceStartTimeout = 300 * time.Second

// JobTimeout is the maximum time a job can take before it times out.
var JobTimeout = 50 * time.Minute

// A Processor will process build jobs on a channel, one by one, until it is
// told to shut down or the channel of build jobs closes.
type Processor struct {
	processorUUID uuid.UUID

	ctx           context.Context
	buildJobsChan <-chan Job
	provider      backend.Provider
	generator     BuildScriptGenerator

	graceful  chan struct{}
	terminate context.CancelFunc
}

// NewProcessor creates a new processor that will run the build jobs on the
// given channel using the given provider and getting build scripts from the
// generator.
func NewProcessor(ctx context.Context, buildJobsChan <-chan Job, provider backend.Provider, generator BuildScriptGenerator) *Processor {
	processorUUID := uuid.NewRandom()

	ctx, cancel := context.WithCancel(contextFromProcessor(ctx, processorUUID.String()))

	return &Processor{
		processorUUID: processorUUID,

		ctx:           contextFromProcessor(ctx, processorUUID.String()),
		buildJobsChan: buildJobsChan,
		provider:      provider,
		generator:     generator,

		graceful:  make(chan struct{}),
		terminate: cancel,
	}
}

// Run starts the processor. This method will not return until the processor is
// terminated, either by calling the GracefulShutdown or Terminate methods, or
// if the build jobs channel is closed.
func (p *Processor) Run() {
	for {
		select {
		case <-p.graceful:
			LoggerFromContext(p.ctx).Info("processor is done, terminating")
			return
		case buildJob, ok := <-p.buildJobsChan:
			if !ok {
				return
			}
			p.process(contextFromJob(p.ctx, buildJob), buildJob)
		}
	}
}

// GracefulShutdown tells the processor to finish the job it is currently
// processing, but not pick up any new jobs. This method will return
// immediately, the processor is done when Run() returns.
func (p *Processor) GracefulShutdown() {
	LoggerFromContext(p.ctx).Info("processor initiating graceful shutdown")
	close(p.graceful)
}

// Terminate tells the processor to stop working on the current job as soon as
// possible.
func (p *Processor) Terminate() {
	p.terminate()
}

func (p *Processor) process(ctx context.Context, buildJob Job) {
	state := new(multistep.BasicStateBag)
	state.Put("buildJob", buildJob)
	state.Put("ctx", ctx)

	steps := []multistep.Step{
		&stepGenerateScript{generator: p.generator},
		&stepStartInstance{provider: p.provider},
		&stepUploadScript{},
		&stepRunScript{},
	}

	runner := &multistep.BasicRunner{Steps: steps}

	LoggerFromContext(p.ctx).Info("starting job")
	runner.Run(state)
	LoggerFromContext(p.ctx).Info("finished job")
}
<commit_msg>processor: use the right context for job start/finish logs<commit_after>package lib

import (
	"time"

	"code.google.com/p/go-uuid/uuid"

	"github.com/mitchellh/multistep"
	"github.com/travis-ci/worker/lib/backend"
	"golang.org/x/net/context"
)

// InstanceStartTimeout is the timeout for starting an instance and waiting for
// it to be available.
var InstanceStartTimeout = 300 * time.Second

// JobTimeout is the maximum time a job can take before it times out.
var JobTimeout = 50 * time.Minute

// A Processor will process build jobs on a channel, one by one, until it is
// told to shut down or the channel of build jobs closes.
type Processor struct {
	processorUUID uuid.UUID

	ctx           context.Context
	buildJobsChan <-chan Job
	provider      backend.Provider
	generator     BuildScriptGenerator

	graceful  chan struct{}
	terminate context.CancelFunc
}

// NewProcessor creates a new processor that will run the build jobs on the
// given channel using the given provider and getting build scripts from the
// generator.
func NewProcessor(ctx context.Context, buildJobsChan <-chan Job, provider backend.Provider, generator BuildScriptGenerator) *Processor {
	processorUUID := uuid.NewRandom()

	ctx, cancel := context.WithCancel(contextFromProcessor(ctx, processorUUID.String()))

	return &Processor{
		processorUUID: processorUUID,

		ctx:           contextFromProcessor(ctx, processorUUID.String()),
		buildJobsChan: buildJobsChan,
		provider:      provider,
		generator:     generator,

		graceful:  make(chan struct{}),
		terminate: cancel,
	}
}

// Run starts the processor. This method will not return until the processor is
// terminated, either by calling the GracefulShutdown or Terminate methods, or
// if the build jobs channel is closed.
func (p *Processor) Run() {
	for {
		select {
		case <-p.graceful:
			LoggerFromContext(p.ctx).Info("processor is done, terminating")
			return
		case buildJob, ok := <-p.buildJobsChan:
			if !ok {
				return
			}
			p.process(contextFromJob(p.ctx, buildJob), buildJob)
		}
	}
}

// GracefulShutdown tells the processor to finish the job it is currently
// processing, but not pick up any new jobs. This method will return
// immediately, the processor is done when Run() returns.
func (p *Processor) GracefulShutdown() {
	LoggerFromContext(p.ctx).Info("processor initiating graceful shutdown")
	close(p.graceful)
}

// Terminate tells the processor to stop working on the current job as soon as
// possible.
func (p *Processor) Terminate() {
	p.terminate()
}

func (p *Processor) process(ctx context.Context, buildJob Job) {
	state := new(multistep.BasicStateBag)
	state.Put("buildJob", buildJob)
	state.Put("ctx", ctx)

	steps := []multistep.Step{
		&stepGenerateScript{generator: p.generator},
		&stepStartInstance{provider: p.provider},
		&stepUploadScript{},
		&stepRunScript{},
	}

	runner := &multistep.BasicRunner{Steps: steps}

	LoggerFromContext(ctx).Info("starting job")
	runner.Run(state)
	LoggerFromContext(ctx).Info("finished job")
}
<|endoftext|>


<commit_before>// Copyright 2019 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package submission

import (
	"context"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"sync"
	"time"

	"github.com/google/certificate-transparency-go/client"
	"github.com/google/certificate-transparency-go/ctpolicy"
	"github.com/google/certificate-transparency-go/jsonclient"
	"github.com/google/certificate-transparency-go/loglist2"
	"github.com/google/certificate-transparency-go/trillian/ctfe"
	"github.com/google/certificate-transparency-go/x509"
	"github.com/google/trillian/monitoring"

	ct "github.com/google/certificate-transparency-go"
)

var (
	// Metrics are per-log, per-endpoint and some per-response-status code.
	distOnce      sync.Once
	reqsCounter   monitoring.Counter   // logurl, ep => value
	rspsCounter   monitoring.Counter   // logurl, ep, sc => value
	logRspLatency monitoring.Histogram // logurl, ep => value
)

// distInitMetrics initializes all the exported metrics.
func distInitMetrics(mf monitoring.MetricFactory) {
	reqsCounter = mf.NewCounter("http_reqs", "Number of requests", "logurl", "ep")
	rspsCounter = mf.NewCounter("http_rsps", "Number of responses", "logurl", "ep", "httpstatus")
	logRspLatency = mf.NewHistogram("http_log_latency", "Latency of responses in seconds", "logurl", "ep")
}

const (
	// GetRootsTimeout timeout used for external requests within root-updates.
	getRootsTimeout = time.Second * 10
)

// pendingLogsPolicy is policy stub used for spreading submissions across
// Pending and Qualified Logs.
type pendingLogsPolicy struct {
}

func (stubP pendingLogsPolicy) LogsByGroup(cert *x509.Certificate, approved *loglist2.LogList) (ctpolicy.LogPolicyData, error) {
	baseGroup, err := ctpolicy.BaseGroupFor(approved, 1)
	groups := ctpolicy.LogPolicyData{baseGroup.Name: baseGroup}
	return groups, err
}

func (stubP pendingLogsPolicy) Name() string {
	return "Pending/Qualified Policy"
}

// Distributor operates policy-based submission across Logs.
type Distributor struct {
	ll                 *loglist2.LogList
	usableLl           *loglist2.LogList
	pendingQualifiedLl *loglist2.LogList

	mu sync.RWMutex

	// helper structs produced out of ll during init.
	logClients map[string]client.AddLogClient
	logRoots   loglist2.LogRoots
	rootPool   *ctfe.PEMCertPool

	rootDataFull bool

	policy            ctpolicy.CTPolicy
	pendingLogsPolicy ctpolicy.CTPolicy
}

// RefreshRoots requests roots from Logs and updates local copy.
// Returns error map keyed by log-URL for any Log experiencing roots retrieval
// problems
// If at least one root was successfully parsed for a log, log roots set gets
// the update.
func (d *Distributor) RefreshRoots(ctx context.Context) map[string]error {
	type RootsResult struct {
		LogURL string
		Roots  *ctfe.PEMCertPool
		Err    error
	}
	ch := make(chan RootsResult, len(d.logClients))

	rctx, cancel := context.WithTimeout(ctx, getRootsTimeout)
	defer cancel()

	for logURL, lc := range d.logClients {
		go func(logURL string, lc client.AddLogClient) {
			res := RootsResult{LogURL: logURL}

			roots, err := lc.GetAcceptedRoots(rctx)
			if err != nil {
				res.Err = fmt.Errorf("roots refresh for %s: couldn't collect roots. %s", logURL, err)
				ch <- res
				return
			}
			res.Roots = ctfe.NewPEMCertPool()
			for _, r := range roots {
				parsed, err := x509.ParseCertificate(r.Data)
				if x509.IsFatal(err) {
					errS := fmt.Errorf("roots refresh for %s: unable to parse root cert: %s", logURL, err)
					if res.Err != nil {
						res.Err = fmt.Errorf("%s\n%s", res.Err, errS)
					} else {
						res.Err = errS
					}
					continue
				}
				res.Roots.AddCert(parsed)
			}
			ch <- res
		}(logURL, lc)
	}

	// Collect get-roots results for every Log-client.
	freshRoots := make(loglist2.LogRoots)
	errors := make(map[string]error)
	for range d.logClients {
		r := <-ch
		// update roots
		if r.Err != nil {
			errors[r.LogURL] = r.Err
		}
		// Roots get update even if some returned roots couldn't get parsed.
		if r.Roots != nil {
			freshRoots[r.LogURL] = r.Roots
		}
	}

	d.mu.Lock()
	defer d.mu.Unlock()

	d.logRoots = freshRoots
	d.rootDataFull = len(d.logRoots) == len(d.logClients)
	// Merge individual root-pools into a unified one
	d.rootPool = ctfe.NewPEMCertPool()
	for _, pool := range d.logRoots {
		for _, c := range pool.RawCertificates() {
			d.rootPool.AddCert(c)
		}
	}

	return errors
}

// IsRootDataFull returns true if root certificates have been obtained for all Logs.
func (d *Distributor) isRootDataFull() bool {
	d.mu.RLock()
	defer d.mu.RUnlock()
	return d.rootDataFull
}

// incRspsCounter extracts HTTP status code and increments corresponding rspsCounter.
func incRspsCounter(logURL string, endpoint string, rspErr error) {
	status := http.StatusOK
	if rspErr != nil {
		status = http.StatusBadRequest // default to this if status code unavailable
		if err, ok := rspErr.(client.RspError); ok {
			status = err.StatusCode
		}
	}
	rspsCounter.Inc(logURL, endpoint, strconv.Itoa(status))
}

// SubmitToLog implements Submitter interface.
func (d *Distributor) SubmitToLog(ctx context.Context, logURL string, chain []ct.ASN1Cert, asPreChain bool) (*ct.SignedCertificateTimestamp, error) {
	lc, ok := d.logClients[logURL]
	if !ok {
		return nil, fmt.Errorf("no client registered for Log with URL %q", logURL)
	}

	// endpoint used for metrics
	endpoint := string(ctfe.AddChainName)
	if asPreChain {
		endpoint = string(ctfe.AddPreChainName)
	}

	defer func(start time.Time) {
		logRspLatency.Observe(time.Since(start).Seconds(), logURL, endpoint)
	}(time.Now())
	reqsCounter.Inc(logURL, endpoint)
	addChain := lc.AddChain
	if asPreChain {
		addChain = lc.AddPreChain
	}
	sct, err := addChain(ctx, chain)
	incRspsCounter(logURL, endpoint, err)
	return sct, err
}

// parseRawChain reads cert chain from bytes into x509.Certificate format.
func parseRawChain(rawChain [][]byte) ([]*x509.Certificate, error) {
	parsedChain := make([]*x509.Certificate, 0, len(rawChain))
	for _, certBytes := range rawChain {
		cert, err := x509.ParseCertificate(certBytes)
		if x509.IsFatal(err) {
			return nil, fmt.Errorf("distributor unable to parse cert-chain %v", err)
		}
		parsedChain = append(parsedChain, cert)
	}
	return parsedChain, nil
}

// addSomeChain is helper calling one of AddChain or AddPreChain based
// on asPreChain param.
func (d *Distributor) addSomeChain(ctx context.Context, rawChain [][]byte, loadPendingLogs bool, asPreChain bool) ([]*AssignedSCT, error) {
	if len(rawChain) == 0 {
		return nil, fmt.Errorf("distributor unable to process empty chain")
	}

	var parsedChain []*x509.Certificate
	var compatibleLogs loglist2.LogList

	d.mu.RLock()
	vOpts := ctfe.NewCertValidationOpts(d.rootPool, time.Time{}, false, false, nil, nil, false, nil)
	rootedChain, err := ctfe.ValidateChain(rawChain, vOpts)

	if err == nil {
		parsedChain = rootedChain
		compatibleLogs = d.usableLl.Compatible(rootedChain[0], rootedChain[len(rootedChain)-1], d.logRoots)
	} else if d.isRootDataFull() {
		// Could not verify the chain while root info for logs is complete.
		d.mu.RUnlock()
		return nil, fmt.Errorf("distributor unable to process cert-chain: %v", err)
	} else {
		// Chain might be rooted to the Log which has no root-info yet.
		parsedChain, err = parseRawChain(rawChain)
		if err != nil {
			return nil, fmt.Errorf("distributor unable to parse cert-chain: %v", err)
		}
		compatibleLogs = d.usableLl.Compatible(parsedChain[0], nil, d.logRoots)
	}
	d.mu.RUnlock()

	// Distinguish between precerts and certificates.
	isPrecert, err := ctfe.IsPrecertificate(parsedChain[0])
	if err != nil {
		return nil, fmt.Errorf("distributor unable to check certificate %v: \n%v", parsedChain[0], err)
	}
	if isPrecert != asPreChain {
		var methodType, inputType string
		if asPreChain {
			methodType = "pre-"
		}
		if isPrecert {
			inputType = "pre-"
		}
		return nil, fmt.Errorf("add-%schain method expected %scertificate, got %scertificate", methodType, methodType, inputType)
	}

	// Set up policy structs.
	groups, err := d.policy.LogsByGroup(parsedChain[0], &compatibleLogs)
	if err != nil {
		return nil, fmt.Errorf("distributor does not have enough compatible Logs to comply with the policy: %v", err)
	}
	chain := make([]ct.ASN1Cert, len(parsedChain))
	for i, c := range parsedChain {
		chain[i] = ct.ASN1Cert{Data: c.Raw}
	}
	if loadPendingLogs {
		go func() {
			pendingGroup, err := d.pendingLogsPolicy.LogsByGroup(parsedChain[0], d.pendingQualifiedLl)
			if err != nil {
				return
			}
			GetSCTs(ctx, d, chain, asPreChain, pendingGroup)
		}()
	}
	return GetSCTs(ctx, d, chain, asPreChain, groups)
}

// AddPreChain runs add-pre-chain calls across subset of logs according to
// Distributor's policy. May emit both SCTs array and error when SCTs
// collected do not satisfy the policy.
func (d *Distributor) AddPreChain(ctx context.Context, rawChain [][]byte, loadPendingLogs bool) ([]*AssignedSCT, error) {
	return d.addSomeChain(ctx, rawChain, loadPendingLogs, true)
}

// AddChain runs add-chain calls across subset of logs according to
// Distributor's policy. May emit both SCTs array and error when SCTs
// collected do not satisfy the policy.
func (d *Distributor) AddChain(ctx context.Context, rawChain [][]byte, loadPendingLogs bool) ([]*AssignedSCT, error) {
	return d.addSomeChain(ctx, rawChain, loadPendingLogs, false)
}

// LogClientBuilder builds client-interface instance for a given Log.
type LogClientBuilder func(*loglist2.Log) (client.AddLogClient, error)

// BuildLogClient is default (non-mock) LogClientBuilder.
func BuildLogClient(log *loglist2.Log) (client.AddLogClient, error) {
	u, err := url.Parse(log.URL)
	if err != nil {
		return nil, err
	}
	if u.Scheme == "" {
		u.Scheme = "https"
	}
	hc := &http.Client{Timeout: time.Second * 10}
	return client.New(u.String(), hc, jsonclient.Options{PublicKeyDER: log.Key})
}

// NewDistributor creates and inits a Distributor instance.
// The Distributor will asynchronously fetch the latest roots from all of the
// logs when active. Call Run() to fetch roots and init regular updates to keep
// the local copy of the roots up-to-date.
func NewDistributor(ll *loglist2.LogList, plc ctpolicy.CTPolicy, lcBuilder LogClientBuilder, mf monitoring.MetricFactory) (*Distributor, error) {
	var d Distributor
	// Divide Logs by statuses.
	d.ll = ll
	usableStat := []loglist2.LogStatus{loglist2.UsableLogStatus}
	active := ll.SelectByStatus(usableStat)
	d.usableLl = &active
	pendingQualifiedStat := []loglist2.LogStatus{
		loglist2.PendingLogStatus, loglist2.QualifiedLogStatus}
	pending := ll.SelectByStatus(pendingQualifiedStat)
	d.pendingQualifiedLl = &pending

	d.policy = plc
	d.pendingLogsPolicy = pendingLogsPolicy{}
	d.logClients = make(map[string]client.AddLogClient)
	d.logRoots = make(loglist2.LogRoots)
	d.rootPool = ctfe.NewPEMCertPool()

	// Build clients for each of the Logs. Also build log-to-id map.
	if err := d.buildLogClients(lcBuilder, d.usableLl); err != nil {
		return nil, err
	}
	d.buildLogClients(lcBuilder, d.pendingQualifiedLl)

	if mf == nil {
		mf = monitoring.InertMetricFactory{}
	}
	distOnce.Do(func() { distInitMetrics(mf) })
	return &d, nil
}

// buildLogClients builds clients for every Log provided and adds them into
// Distributor internals.
func (d *Distributor) buildLogClients(lcBuilder LogClientBuilder, ll *loglist2.LogList) error {
	for _, op := range ll.Operators {
		for _, log := range op.Logs {
			lc, err := lcBuilder(log)
			if err != nil {
				return fmt.Errorf("failed to create log client for %s: %v", log.URL, err)
			}
			d.logClients[log.URL] = lc
		}
	}
	return nil
}
<commit_msg>Enforce http-error code on invalid SCT response for the metric (#607)<commit_after>// Copyright 2019 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package submission

import (
	"context"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"sync"
	"time"

	"github.com/google/certificate-transparency-go/client"
	"github.com/google/certificate-transparency-go/ctpolicy"
	"github.com/google/certificate-transparency-go/jsonclient"
	"github.com/google/certificate-transparency-go/loglist2"
	"github.com/google/certificate-transparency-go/trillian/ctfe"
	"github.com/google/certificate-transparency-go/x509"
	"github.com/google/trillian/monitoring"

	ct "github.com/google/certificate-transparency-go"
)

var (
	// Metrics are per-log, per-endpoint and some per-response-status code.
	distOnce      sync.Once
	reqsCounter   monitoring.Counter   // logurl, ep => value
	rspsCounter   monitoring.Counter   // logurl, ep, sc => value
	errCounter    monitoring.Counter   // logurl, ep, status => value
	logRspLatency monitoring.Histogram // logurl, ep => value
)

// distInitMetrics initializes all the exported metrics.
func distInitMetrics(mf monitoring.MetricFactory) {
	reqsCounter = mf.NewCounter("http_reqs", "Number of requests", "logurl", "ep")
	rspsCounter = mf.NewCounter("http_rsps", "Number of responses", "logurl", "ep", "httpstatus")
	errCounter = mf.NewCounter("err_count", "Number of errors", "logurl", "ep", "errtype")
	logRspLatency = mf.NewHistogram("http_log_latency", "Latency of responses in seconds", "logurl", "ep")
}

const (
	// GetRootsTimeout timeout used for external requests within root-updates.
	getRootsTimeout = time.Second * 10
)

// pendingLogsPolicy is policy stub used for spreading submissions across
// Pending and Qualified Logs.
type pendingLogsPolicy struct {
}

func (stubP pendingLogsPolicy) LogsByGroup(cert *x509.Certificate, approved *loglist2.LogList) (ctpolicy.LogPolicyData, error) {
	baseGroup, err := ctpolicy.BaseGroupFor(approved, 1)
	groups := ctpolicy.LogPolicyData{baseGroup.Name: baseGroup}
	return groups, err
}

func (stubP pendingLogsPolicy) Name() string {
	return "Pending/Qualified Policy"
}

// Distributor operates policy-based submission across Logs.
type Distributor struct {
	ll                 *loglist2.LogList
	usableLl           *loglist2.LogList
	pendingQualifiedLl *loglist2.LogList

	mu sync.RWMutex

	// helper structs produced out of ll during init.
	logClients map[string]client.AddLogClient
	logRoots   loglist2.LogRoots
	rootPool   *ctfe.PEMCertPool

	rootDataFull bool

	policy            ctpolicy.CTPolicy
	pendingLogsPolicy ctpolicy.CTPolicy
}

// RefreshRoots requests roots from Logs and updates local copy.
// Returns error map keyed by log-URL for any Log experiencing roots retrieval
// problems
// If at least one root was successfully parsed for a log, log roots set gets
// the update.
func (d *Distributor) RefreshRoots(ctx context.Context) map[string]error {
	type RootsResult struct {
		LogURL string
		Roots  *ctfe.PEMCertPool
		Err    error
	}
	ch := make(chan RootsResult, len(d.logClients))

	rctx, cancel := context.WithTimeout(ctx, getRootsTimeout)
	defer cancel()

	for logURL, lc := range d.logClients {
		go func(logURL string, lc client.AddLogClient) {
			res := RootsResult{LogURL: logURL}

			roots, err := lc.GetAcceptedRoots(rctx)
			if err != nil {
				res.Err = fmt.Errorf("roots refresh for %s: couldn't collect roots. %s", logURL, err)
				ch <- res
				return
			}
			res.Roots = ctfe.NewPEMCertPool()
			for _, r := range roots {
				parsed, err := x509.ParseCertificate(r.Data)
				if x509.IsFatal(err) {
					errS := fmt.Errorf("roots refresh for %s: unable to parse root cert: %s", logURL, err)
					if res.Err != nil {
						res.Err = fmt.Errorf("%s\n%s", res.Err, errS)
					} else {
						res.Err = errS
					}
					continue
				}
				res.Roots.AddCert(parsed)
			}
			ch <- res
		}(logURL, lc)
	}

	// Collect get-roots results for every Log-client.
	freshRoots := make(loglist2.LogRoots)
	errors := make(map[string]error)
	for range d.logClients {
		r := <-ch
		// update roots
		if r.Err != nil {
			errors[r.LogURL] = r.Err
		}
		// Roots get update even if some returned roots couldn't get parsed.
		if r.Roots != nil {
			freshRoots[r.LogURL] = r.Roots
		}
	}

	d.mu.Lock()
	defer d.mu.Unlock()

	d.logRoots = freshRoots
	d.rootDataFull = len(d.logRoots) == len(d.logClients)
	// Merge individual root-pools into a unified one
	d.rootPool = ctfe.NewPEMCertPool()
	for _, pool := range d.logRoots {
		for _, c := range pool.RawCertificates() {
			d.rootPool.AddCert(c)
		}
	}

	return errors
}

// IsRootDataFull returns true if root certificates have been obtained for all Logs.
func (d *Distributor) isRootDataFull() bool {
	d.mu.RLock()
	defer d.mu.RUnlock()
	return d.rootDataFull
}

// incRspsCounter extracts HTTP status code and increments corresponding rspsCounter.
func incRspsCounter(logURL string, endpoint string, rspErr error) {
	status := http.StatusOK
	if rspErr != nil {
		status = http.StatusBadRequest // default to this if status code unavailable
		if err, ok := rspErr.(client.RspError); ok {
			status = err.StatusCode
		}
	}
	rspsCounter.Inc(logURL, endpoint, strconv.Itoa(status))
}

// incErrCounter increments corresponding errCounter if any error occurred during
// submission to a Log.
func incErrCounter(logURL string, endpoint string, rspErr error) {
	if rspErr == nil {
		return
	}
	err, ok := rspErr.(client.RspError)
	switch {
	case !ok:
		errCounter.Inc(logURL, endpoint, "unknown_error")
	case err.Err != nil && err.StatusCode == http.StatusOK:
		errCounter.Inc(logURL, endpoint, "invalid_sct")
	case err.Err != nil: // err.StatusCode != http.StatusOK.
		errCounter.Inc(logURL, endpoint, "connection_error")
	}
}

// SubmitToLog implements Submitter interface.
func (d *Distributor) SubmitToLog(ctx context.Context, logURL string, chain []ct.ASN1Cert, asPreChain bool) (*ct.SignedCertificateTimestamp, error) {
	lc, ok := d.logClients[logURL]
	if !ok {
		return nil, fmt.Errorf("no client registered for Log with URL %q", logURL)
	}

	// endpoint used for metrics
	endpoint := string(ctfe.AddChainName)
	if asPreChain {
		endpoint = string(ctfe.AddPreChainName)
	}

	defer func(start time.Time) {
		logRspLatency.Observe(time.Since(start).Seconds(), logURL, endpoint)
	}(time.Now())
	reqsCounter.Inc(logURL, endpoint)
	addChain := lc.AddChain
	if asPreChain {
		addChain = lc.AddPreChain
	}
	sct, err := addChain(ctx, chain)
	incRspsCounter(logURL, endpoint, err)
	incErrCounter(logURL, endpoint, err)
	return sct, err
}

// parseRawChain reads cert chain from bytes into x509.Certificate format.
func parseRawChain(rawChain [][]byte) ([]*x509.Certificate, error) {
	parsedChain := make([]*x509.Certificate, 0, len(rawChain))
	for _, certBytes := range rawChain {
		cert, err := x509.ParseCertificate(certBytes)
		if x509.IsFatal(err) {
			return nil, fmt.Errorf("distributor unable to parse cert-chain %v", err)
		}
		parsedChain = append(parsedChain, cert)
	}
	return parsedChain, nil
}

// addSomeChain is helper calling one of AddChain or AddPreChain based
// on asPreChain param.
func (d *Distributor) addSomeChain(ctx context.Context, rawChain [][]byte, loadPendingLogs bool, asPreChain bool) ([]*AssignedSCT, error) {
	if len(rawChain) == 0 {
		return nil, fmt.Errorf("distributor unable to process empty chain")
	}

	var parsedChain []*x509.Certificate
	var compatibleLogs loglist2.LogList

	d.mu.RLock()
	vOpts := ctfe.NewCertValidationOpts(d.rootPool, time.Time{}, false, false, nil, nil, false, nil)
	rootedChain, err := ctfe.ValidateChain(rawChain, vOpts)

	if err == nil {
		parsedChain = rootedChain
		compatibleLogs = d.usableLl.Compatible(rootedChain[0], rootedChain[len(rootedChain)-1], d.logRoots)
	} else if d.isRootDataFull() {
		// Could not verify the chain while root info for logs is complete.
		d.mu.RUnlock()
		return nil, fmt.Errorf("distributor unable to process cert-chain: %v", err)
	} else {
		// Chain might be rooted to the Log which has no root-info yet.
		parsedChain, err = parseRawChain(rawChain)
		if err != nil {
			return nil, fmt.Errorf("distributor unable to parse cert-chain: %v", err)
		}
		compatibleLogs = d.usableLl.Compatible(parsedChain[0], nil, d.logRoots)
	}
	d.mu.RUnlock()

	// Distinguish between precerts and certificates.
	isPrecert, err := ctfe.IsPrecertificate(parsedChain[0])
	if err != nil {
		return nil, fmt.Errorf("distributor unable to check certificate %v: \n%v", parsedChain[0], err)
	}
	if isPrecert != asPreChain {
		var methodType, inputType string
		if asPreChain {
			methodType = "pre-"
		}
		if isPrecert {
			inputType = "pre-"
		}
		return nil, fmt.Errorf("add-%schain method expected %scertificate, got %scertificate", methodType, methodType, inputType)
	}

	// Set up policy structs.
	groups, err := d.policy.LogsByGroup(parsedChain[0], &compatibleLogs)
	if err != nil {
		return nil, fmt.Errorf("distributor does not have enough compatible Logs to comply with the policy: %v", err)
	}
	chain := make([]ct.ASN1Cert, len(parsedChain))
	for i, c := range parsedChain {
		chain[i] = ct.ASN1Cert{Data: c.Raw}
	}
	if loadPendingLogs {
		go func() {
			pendingGroup, err := d.pendingLogsPolicy.LogsByGroup(parsedChain[0], d.pendingQualifiedLl)
			if err != nil {
				return
			}
			GetSCTs(ctx, d, chain, asPreChain, pendingGroup)
		}()
	}
	return GetSCTs(ctx, d, chain, asPreChain, groups)
}

// AddPreChain runs add-pre-chain calls across subset of logs according to
// Distributor's policy. May emit both SCTs array and error when SCTs
// collected do not satisfy the policy.
func (d *Distributor) AddPreChain(ctx context.Context, rawChain [][]byte, loadPendingLogs bool) ([]*AssignedSCT, error) {
	return d.addSomeChain(ctx, rawChain, loadPendingLogs, true)
}

// AddChain runs add-chain calls across subset of logs according to
// Distributor's policy. May emit both SCTs array and error when SCTs
// collected do not satisfy the policy.
func (d *Distributor) AddChain(ctx context.Context, rawChain [][]byte, loadPendingLogs bool) ([]*AssignedSCT, error) {
	return d.addSomeChain(ctx, rawChain, loadPendingLogs, false)
}

// LogClientBuilder builds client-interface instance for a given Log.
type LogClientBuilder func(*loglist2.Log) (client.AddLogClient, error)

// BuildLogClient is default (non-mock) LogClientBuilder.
func BuildLogClient(log *loglist2.Log) (client.AddLogClient, error) {
	u, err := url.Parse(log.URL)
	if err != nil {
		return nil, err
	}
	if u.Scheme == "" {
		u.Scheme = "https"
	}
	hc := &http.Client{Timeout: time.Second * 10}
	return client.New(u.String(), hc, jsonclient.Options{PublicKeyDER: log.Key})
}

// NewDistributor creates and inits a Distributor instance.
// The Distributor will asynchronously fetch the latest roots from all of the
// logs when active. Call Run() to fetch roots and init regular updates to keep
// the local copy of the roots up-to-date.
func NewDistributor(ll *loglist2.LogList, plc ctpolicy.CTPolicy, lcBuilder LogClientBuilder, mf monitoring.MetricFactory) (*Distributor, error) {
	var d Distributor
	// Divide Logs by statuses.
	d.ll = ll
	usableStat := []loglist2.LogStatus{loglist2.UsableLogStatus}
	active := ll.SelectByStatus(usableStat)
	d.usableLl = &active
	pendingQualifiedStat := []loglist2.LogStatus{
		loglist2.PendingLogStatus, loglist2.QualifiedLogStatus}
	pending := ll.SelectByStatus(pendingQualifiedStat)
	d.pendingQualifiedLl = &pending

	d.policy = plc
	d.pendingLogsPolicy = pendingLogsPolicy{}
	d.logClients = make(map[string]client.AddLogClient)
	d.logRoots = make(loglist2.LogRoots)
	d.rootPool = ctfe.NewPEMCertPool()

	// Build clients for each of the Logs. Also build log-to-id map.
	if err := d.buildLogClients(lcBuilder, d.usableLl); err != nil {
		return nil, err
	}
	d.buildLogClients(lcBuilder, d.pendingQualifiedLl)

	if mf == nil {
		mf = monitoring.InertMetricFactory{}
	}
	distOnce.Do(func() { distInitMetrics(mf) })
	return &d, nil
}

// buildLogClients builds clients for every Log provided and adds them into
// Distributor internals.
func (d *Distributor) buildLogClients(lcBuilder LogClientBuilder, ll *loglist2.LogList) error {
	for _, op := range ll.Operators {
		for _, log := range op.Logs {
			lc, err := lcBuilder(log)
			if err != nil {
				return fmt.Errorf("failed to create log client for %s: %v", log.URL, err)
			}
			d.logClients[log.URL] = lc
		}
	}
	return nil
}
<|endoftext|>