<commit_before>package cwl

import (
	"fmt"
	"sort"
	"strings"
)

// Input represents "CommandInputParameter".
// @see http://www.commonwl.org/v1.0/CommandLineTool.html#CommandInputParameter
type Input struct {
	ID             string          `json:"id"`
	Label          string          `json:"label"`
	Doc            string          `json:"doc"`
	Format         string          `json:"format"`
	Binding        *Binding        `json:"inputBinding"`
	Default        *InputDefault   `json:"default"`
	Types          []Type          `json:"type"`
	SecondaryFiles []SecondaryFile `json:"secondary_files"`
	// Input.Provided is what provided by parameters.(json|yaml)
	Provided *Provided `json:"-"`
	// Requirement ..
	RequiredType *Type
	Requirements Requirements
}

// New constructs "Input" struct from interface{}.
func (input Input) New(i interface{}) *Input {
	dest := &Input{}
	switch x := i.(type) {
	case map[string]interface{}:
		for key, v := range x {
			switch key {
			case "id":
				dest.ID = v.(string)
			case "type":
				dest.Types = Type{}.NewList(v)
			case "label":
				dest.Label = v.(string)
			case "doc":
				dest.Doc = v.(string)
			case "inputBinding":
				dest.Binding = Binding{}.New(v)
			case "default":
				dest.Default = InputDefault{}.New(v)
			case "format":
				dest.Format = v.(string)
			case "secondaryFiles":
				dest.SecondaryFiles = SecondaryFile{}.NewList(v)
			}
		}
		if dest.Default != nil {
			dest.Default.ID = dest.ID
		}
	case string:
		dest.Types = Type{}.NewList(x)
	case []interface{}:
		for _, v := range x {
			dest.Types = append(dest.Types, Type{}.New(v))
		}
	}
	return dest
}

// flatten
func (input *Input) flatten(typ Type, binding *Binding, prov interface{}) []string {
	flattened := []string{}
	switch typ.Type {
	case "int": // Array of Int
		tobejoined := []string{}
		for _, e := range input.Provided.Raw.([]interface{}) {
			tobejoined = append(tobejoined, fmt.Sprintf("%v", e))
		}
		flattened = append(flattened, strings.Join(tobejoined, input.Binding.Separator))
	case "File": // Array of Files
		switch arr := input.Provided.Raw.(type) {
		case []string:
			// TODO:
		case []interface{}:
			separated := []string{}
			for _, e := range arr {
				switch v := e.(type) {
				case map[interface{}]interface{}:
					if binding != nil && binding.Prefix != "" {
						separated = append(separated, binding.Prefix)
					}
					separated = append(separated, fmt.Sprintf("%v", v["location"]))
				default:
					// TODO:
				}
			}
			// In case it's Array of Files, unlike array of int,
			// it's NOT gonna be joined with .Binding.Separator.
			flattened = append(flattened, separated...)
		}
	case "string": // Array of string
		switch arr := prov.(type) {
		case []interface{}:
			separated := []string{}
			for _, e := range arr {
				switch v := e.(type) {
				case string:
					if binding != nil && binding.Prefix != "" {
						separated = append(separated, binding.Prefix)
					}
					separated = append(separated, fmt.Sprintf("%v", v))
				default:
					// TODO
				}
			}
			flattened = append(flattened, separated...)
		default:
			// TODO
		}
	case "array":
		switch arr := prov.(type) {
		case []interface{}:
			flattened = append(flattened, input.flatten(typ.Items[0], typ.Binding, arr[0])...)
		default:
			// TODO
		}
	default:
		if input.RequiredType != nil {
			flattened = append(flattened, input.flattenWithRequiredType()...)
		} else {
			// TODO
		}
	}
	return flattened
}

func (input *Input) flattenWithRequiredType() []string {
	flattened := []string{}
	key, needed := input.Types[0].NeedRequirement()
	if !needed {
		return flattened
	}
	if input.RequiredType.Name != key {
		return flattened
	}
	switch provided := input.Provided.Raw.(type) {
	case []interface{}:
		for _, e := range provided {
			switch v := e.(type) {
			case map[interface{}]interface{}:
				for _, field := range input.RequiredType.Fields {
					val, ok := v[field.Name]
					if !ok {
						continue
					}
					if field.Binding == nil {
						// Without thinking anything, just append it!!!
						flattened = append(flattened, fmt.Sprintf("%v", val))
					} else {
						if field.Binding.Prefix != "" {
							if field.Binding.Separate {
								flattened = append(flattened, field.Binding.Prefix, fmt.Sprintf("%v", val))
							} else {
								// TODO: Join if .Separator is given
								flattened = append(flattened, fmt.Sprintf("%s%v", field.Binding.Prefix, val))
							}
						} else {
							switch v2 := val.(type) {
							case []interface{}:
								for _, val2 := range v2 {
									switch v3 := val2.(type) {
									case []interface{}:
									case map[interface{}]interface{}:
										for _, types := range input.Requirements[0].SchemaDefRequirement.Types {
											val3array := []string{}
											var val3count int = 0
											sort.Sort(types.Fields)
											for _, fields := range types.Fields {
												for key3, val3 := range v3 {
													if fields.Name == key3 {
														for _, val3type := range fields.Types {
															if val3type.Type == "" {
															} else {
																switch val3type.Type {
																case "enum":
																	for _, symbol := range val3type.Symbols {
																		if symbol == val3 {
																			val3array = append(val3array, fmt.Sprintf("%v", val3))
																			val3count = val3count + 1
																		}
																	}
																case "int":
																	if fields.Binding.Prefix != "" {
																		val3array = append(val3array, fields.Binding.Prefix, fmt.Sprintf("%v", val3))
																		val3count = val3count + 1
																	} else {
																		val3array = append(val3array, fmt.Sprintf("%v", val3))
																		val3count = val3count + 1
																	}
																}
															}
														}
													}
												}
											}
											if len(v3) == val3count {
												flattened = append(flattened, val3array...)
											}
										}
									}
								}
							}
						}
					}
				}
			}
		}
	}
	return flattened
}

// Flatten ...
func (input *Input) Flatten() []string {
	if input.Provided == nil {
		// In case "input.Default == nil" should be validated by usage layer.
		if input.Default != nil {
			return input.Default.Flatten(input.Binding)
		} else {
			return []string{}
		}
	}
	flattened := []string{}
	if repr := input.Types[0]; len(input.Types) == 1 {
		switch repr.Type {
		case "array":
			flattened = append(flattened, input.flatten(repr.Items[0], repr.Binding, input.Provided.Raw)...)
		case "int":
			flattened = append(flattened, fmt.Sprintf("%v", input.Provided.Int))
		case "File":
			flattened = append(flattened, input.Provided.Entry.Location)
		default:
			flattened = append(flattened, fmt.Sprintf("%v", input.Provided))
		}
	}
	if input.Binding != nil && input.Binding.Prefix != "" {
		flattened = append([]string{input.Binding.Prefix}, flattened...)
	}

	return flattened
}

// Inputs represents "inputs" field in CWL.
type Inputs []*Input

// New constructs new "Inputs" struct.
func (ins Inputs) New(i interface{}) Inputs {
	dest := Inputs{}
	switch x := i.(type) {
	case []interface{}:
		for _, v := range x {
			dest = append(dest, Input{}.New(v))
		}
	case map[string]interface{}:
		for key, v := range x {
			input := Input{}.New(v)
			input.ID = key
			dest = append(dest, input)
		}
	}
	return dest
}

// Len for sorting.
func (ins Inputs) Len() int {
	return len(ins)
}

// Less for sorting.
func (ins Inputs) Less(i, j int) bool {
	prev, next := ins[i].Binding, ins[j].Binding
	switch [2]bool{prev == nil, next == nil} {
	case [2]bool{true, true}:
		return true
	case [2]bool{false, true}:
		return prev.Position < 0
	case [2]bool{true, false}:
		return next.Position > 0
	default:
		return prev.Position <= next.Position
	}
}

// Swap for sorting.
func (ins Inputs) Swap(i, j int) {
	ins[i], ins[j] = ins[j], ins[i]
}
<commit_msg>For test case 13<commit_after>package cwl

import (
	"fmt"
	"sort"
	"strings"

	"github.com/robertkrimen/otto"
)

// Input represents "CommandInputParameter".
// @see http://www.commonwl.org/v1.0/CommandLineTool.html#CommandInputParameter
type Input struct {
	ID             string          `json:"id"`
	Label          string          `json:"label"`
	Doc            string          `json:"doc"`
	Format         string          `json:"format"`
	Binding        *Binding        `json:"inputBinding"`
	Default        *InputDefault   `json:"default"`
	Types          []Type          `json:"type"`
	SecondaryFiles []SecondaryFile `json:"secondary_files"`
	// Input.Provided is what provided by parameters.(json|yaml)
	Provided *Provided `json:"-"`
	// Requirement ..
	RequiredType *Type
	Requirements Requirements
}

// New constructs "Input" struct from interface{}.
func (input Input) New(i interface{}) *Input {
	dest := &Input{}
	switch x := i.(type) {
	case map[string]interface{}:
		for key, v := range x {
			switch key {
			case "id":
				dest.ID = v.(string)
			case "type":
				dest.Types = Type{}.NewList(v)
			case "label":
				dest.Label = v.(string)
			case "doc":
				dest.Doc = v.(string)
			case "inputBinding":
				dest.Binding = Binding{}.New(v)
			case "default":
				dest.Default = InputDefault{}.New(v)
			case "format":
				dest.Format = v.(string)
			case "secondaryFiles":
				dest.SecondaryFiles = SecondaryFile{}.NewList(v)
			}
		}
		if dest.Default != nil {
			dest.Default.ID = dest.ID
		}
	case string:
		dest.Types = Type{}.NewList(x)
	case []interface{}:
		for _, v := range x {
			dest.Types = append(dest.Types, Type{}.New(v))
		}
	}
	return dest
}

// flatten
func (input *Input) flatten(typ Type, binding *Binding, prov interface{}) []string {
	flattened := []string{}
	switch typ.Type {
	case "int": // Array of Int
		tobejoined := []string{}
		for _, e := range input.Provided.Raw.([]interface{}) {
			tobejoined = append(tobejoined, fmt.Sprintf("%v", e))
		}
		flattened = append(flattened, strings.Join(tobejoined, input.Binding.Separator))
	case "File": // Array of Files
		switch arr := input.Provided.Raw.(type) {
		case []string:
			// TODO:
		case []interface{}:
			separated := []string{}
			for _, e := range arr {
				switch v := e.(type) {
				case map[interface{}]interface{}:
					if binding != nil && binding.Prefix != "" {
						separated = append(separated, binding.Prefix)
					}
					separated = append(separated, fmt.Sprintf("%v", v["location"]))
				default:
					// TODO:
				}
			}
			// In case it's Array of Files, unlike array of int,
			// it's NOT gonna be joined with .Binding.Separator.
			flattened = append(flattened, separated...)
		}
	case "string": // Array of string
		switch arr := prov.(type) {
		case []interface{}:
			separated := []string{}
			for _, e := range arr {
				switch v := e.(type) {
				case string:
					if binding != nil && binding.Prefix != "" {
						separated = append(separated, binding.Prefix)
					}
					separated = append(separated, fmt.Sprintf("%v", v))
				default:
					// TODO
				}
			}
			flattened = append(flattened, separated...)
		default:
			// TODO
		}
	case "array":
		switch arr := prov.(type) {
		case []interface{}:
			flattened = append(flattened, input.flatten(typ.Items[0], typ.Binding, arr[0])...)
		default:
			// TODO
		}
	default:
		if input.RequiredType != nil {
			flattened = append(flattened, input.flattenWithRequiredType()...)
		} else {
			// TODO
		}
	}
	return flattened
}

func (input *Input) flattenWithRequiredType() []string {
	flattened := []string{}
	key, needed := input.Types[0].NeedRequirement()
	if !needed {
		return flattened
	}
	if input.RequiredType.Name != key {
		return flattened
	}
	switch provided := input.Provided.Raw.(type) {
	case []interface{}:
		for _, e := range provided {
			switch v := e.(type) {
			case map[interface{}]interface{}:
				for _, field := range input.RequiredType.Fields {
					val, ok := v[field.Name]
					if !ok {
						continue
					}
					if field.Binding == nil {
						// Without thinking anything, just append it!!!
						flattened = append(flattened, fmt.Sprintf("%v", val))
					} else {
						if field.Binding.Prefix != "" {
							if field.Binding.Separate {
								flattened = append(flattened, field.Binding.Prefix, fmt.Sprintf("%v", val))
							} else {
								// TODO: Join if .Separator is given
								flattened = append(flattened, fmt.Sprintf("%s%v", field.Binding.Prefix, val))
							}
						} else {
							switch v2 := val.(type) {
							case []interface{}:
								for _, val2 := range v2 {
									switch v3 := val2.(type) {
									case []interface{}:
									case map[interface{}]interface{}:
										for _, types := range input.Requirements[0].SchemaDefRequirement.Types {
											val3array := []string{}
											var val3count int = 0
											sort.Sort(types.Fields)
											for _, fields := range types.Fields {
												for key3, val3 := range v3 {
													if fields.Name == key3 {
														for _, val3type := range fields.Types {
															if val3type.Type == "" {
															} else {
																switch val3type.Type {
																case "enum":
																	for _, symbol := range val3type.Symbols {
																		if symbol == val3 {
																			val3array = append(val3array, fmt.Sprintf("%v", val3))
																			val3count = val3count + 1
																		}
																	}
																case "int":
																	if fields.Binding.Prefix != "" {
																		val3array = append(val3array, fields.Binding.Prefix, fmt.Sprintf("%v", val3))
																		val3count = val3count + 1
																	} else {
																		val3array = append(val3array, fmt.Sprintf("%v", val3))
																		val3count = val3count + 1
																	}
																}
															}
														}
													}
												}
											}
											if len(v3) == val3count {
												flattened = append(flattened, val3array...)
											}
										}
									}
								}
							}
						}
					}
				}
			}
		}
	}
	return flattened
}

// Flatten ...
func (input *Input) Flatten() []string {
	if input.Provided == nil {
		// In case "input.Default == nil" should be validated by usage layer.
		if input.Default != nil {
			return input.Default.Flatten(input.Binding)
		} else {
			return []string{}
		}
	}
	flattened := []string{}
	if repr := input.Types[0]; len(input.Types) == 1 {
		switch repr.Type {
		case "array":
			flattened = append(flattened, input.flatten(repr.Items[0], repr.Binding, input.Provided.Raw)...)
		case "int":
			flattened = append(flattened, fmt.Sprintf("%v", input.Provided.Int))
		case "File":
			flattened = append(flattened, input.Provided.Entry.Location)
		default:
			flattened = append(flattened, fmt.Sprintf("%v", input.Provided))
		}
	}
	if input.Binding != nil && input.Binding.Prefix != "" {
		flattened = append([]string{input.Binding.Prefix}, flattened...)
	}

	return flattened
}

// Inputs represents "inputs" field in CWL.
type Inputs []*Input

// New constructs new "Inputs" struct.
func (ins Inputs) New(i interface{}) Inputs {
	dest := Inputs{}
	switch x := i.(type) {
	case []interface{}:
		for _, v := range x {
			dest = append(dest, Input{}.New(v))
		}
	case map[string]interface{}:
		for key, v := range x {
			input := Input{}.New(v)
			input.ID = key
			dest = append(dest, input)
		}
	}
	return dest
}

// Len for sorting.
func (ins Inputs) Len() int {
	return len(ins)
}

// Less for sorting.
func (ins Inputs) Less(i, j int) bool {
	prev, next := ins[i].Binding, ins[j].Binding
	switch [2]bool{prev == nil, next == nil} {
	case [2]bool{true, true}:
		return true
	case [2]bool{false, true}:
		return prev.Position < 0
	case [2]bool{true, false}:
		return next.Position > 0
	default:
		return prev.Position <= next.Position
	}
}

// Swap for sorting.
func (ins Inputs) Swap(i, j int) {
	ins[i], ins[j] = ins[j], ins[i]
}

// CreateVM ...
func (ins Inputs) CreateVM(srcdir string) (*otto.Otto, error) {
	self := map[string]map[string]interface{}{}
	for _, i := range ins {
		if i.Default != nil {
			self[i.ID] = map[string]interface{}{
				"path": i.Default.Entry.Location,
			}
		}
		if i.Provided != nil {
			self[i.ID] = map[string]interface{}{
				"path": i.Provided.Entry.Location,
			}
		}
	}

	// No contents to load
	if len(self) == 0 {
		return nil, nil
	}

	vm := otto.New()
	if err := vm.Set("inputs", self); err != nil {
		return nil, err
	}
	return vm, nil
}
<|endoftext|>


<commit_before>//
// Fluentd Forwarder
//
// Copyright (C) 2014 Treasure Data, Inc.
//
// 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 fluentd_forwarder

import (
	"bufio"
	"errors"
	"fmt"
	logging "github.com/op/go-logging"
	"github.com/ugorji/go/codec"
	"io"
	"net"
	"reflect"
	"sync"
	"sync/atomic"
)

type forwardClient struct {
	input  *ForwardInput
	logger *logging.Logger
	conn   *net.TCPConn
	codec  *codec.MsgpackHandle
	dec    *codec.Decoder
}

type ForwardInput struct {
	entries        int64  // This variable must be on 64-bit alignment. Otherwise atomic.AddInt64 will cause a crash on ARM and x86-32
	port           Port
	logger         *logging.Logger
	bind           string
	listener       *net.TCPListener
	codec          *codec.MsgpackHandle
	clientsMtx     sync.Mutex
	clients        map[*net.TCPConn]*forwardClient
	wg             sync.WaitGroup
	acceptChan     chan *net.TCPConn
	shutdownChan   chan struct{}
	isShuttingDown uintptr
}

type EntryCountTopic struct{}

type ConnectionCountTopic struct{}

type ForwardInputFactory struct{}

func coerceInPlace(data map[string]interface{}) {
	for k, v := range data {
		switch v_ := v.(type) {
		case []byte:
			data[k] = string(v_) // XXX: byte => rune
		case map[string]interface{}:
			coerceInPlace(v_)
		}
	}
}

func (c *forwardClient) decodeRecordSet(tag []byte, entries []interface{}) (FluentRecordSet, error) {
	records := make([]TinyFluentRecord, len(entries))
	for i, _entry := range entries {
		entry, ok := _entry.([]interface{})
		if !ok {
			return FluentRecordSet{}, errors.New("Failed to decode recordSet")
		}
		timestamp, ok := entry[0].(uint64)
		if !ok {
			return FluentRecordSet{}, errors.New("Failed to decode timestamp field")
		}
		data, ok := entry[1].(map[string]interface{})
		if !ok {
			return FluentRecordSet{}, errors.New("Failed to decode data field")
		}
		coerceInPlace(data)
		records[i] = TinyFluentRecord{
			Timestamp: timestamp,
			Data:      data,
		}
	}
	return FluentRecordSet{
		Tag:     string(tag), // XXX: byte => rune
		Records: records,
	}, nil
}

func (c *forwardClient) decodeEntries() ([]FluentRecordSet, error) {
	v := []interface{}{nil, nil, nil}
	err := c.dec.Decode(&v)
	if err != nil {
		return nil, err
	}
	tag, ok := v[0].([]byte)
	if !ok {
		return nil, errors.New("Failed to decode tag field")
	}

	var retval []FluentRecordSet
	switch timestamp_or_entries := v[1].(type) {
	case uint64:
		timestamp := timestamp_or_entries
		data, ok := v[2].(map[string]interface{})
		if !ok {
			return nil, errors.New("Failed to decode data field")
		}
		coerceInPlace(data)
		retval = []FluentRecordSet{
			{
				Tag: string(tag), // XXX: byte => rune
				Records: []TinyFluentRecord{
					{
						Timestamp: timestamp,
						Data:      data,
					},
				},
			},
		}
	case float64:
		timestamp := uint64(timestamp_or_entries)
		data, ok := v[2].(map[string]interface{})
		if !ok {
			return nil, errors.New("Failed to decode data field")
		}
		retval = []FluentRecordSet{
			{
				Tag: string(tag), // XXX: byte => rune
				Records: []TinyFluentRecord{
					{
						Timestamp: timestamp,
						Data:      data,
					},
				},
			},
		}
	case []interface{}:
		if !ok {
			return nil, errors.New("Unexpected payload format")
		}
		recordSet, err := c.decodeRecordSet(tag, timestamp_or_entries)
		if err != nil {
			return nil, err
		}
		retval = []FluentRecordSet{recordSet}
	case []byte:
		entries := make([]interface{}, 0)
		err := codec.NewDecoderBytes(timestamp_or_entries, c.codec).Decode(&entries)
		if err != nil {
			return nil, err
		}
		recordSet, err := c.decodeRecordSet(tag, entries)
		if err != nil {
			return nil, err
		}
		retval = []FluentRecordSet{recordSet}
	default:
		return nil, errors.New(fmt.Sprintf("Unknown type: %t", timestamp_or_entries))
	}
	atomic.AddInt64(&c.input.entries, int64(len(retval)))
	return retval, nil
}

func (c *forwardClient) startHandling() {
	c.input.wg.Add(1)
	go func() {
		defer func() {
			err := c.conn.Close()
			if err != nil {
				c.logger.Debug("Close: %s", err.Error())
			}
			c.input.markDischarged(c)
			c.input.wg.Done()
		}()
		c.input.logger.Info("Started handling connection from %s", c.conn.RemoteAddr().String())
		for {
			recordSets, err := c.decodeEntries()
			if err != nil {
				err_, ok := err.(net.Error)
				if ok {
					if err_.Temporary() {
						c.logger.Info("Temporary failure: %s", err_.Error())
						continue
					}
				}
				if err == io.EOF {
					c.logger.Info("Client %s closed the connection", c.conn.RemoteAddr().String())
				} else {
					c.logger.Error(err.Error())
				}
				break
			}

			if len(recordSets) > 0 {
				err_ := c.input.port.Emit(recordSets)
				if err_ != nil {
					c.logger.Error(err_.Error())
					break
				}
			}
		}
		c.input.logger.Info("Ended handling connection from %s", c.conn.RemoteAddr().String())
	}()
}

func (c *forwardClient) shutdown() {
	err := c.conn.Close()
	if err != nil {
		c.input.logger.Info("Error during closing connection: %s", err.Error())
	}
}

func newForwardClient(input *ForwardInput, logger *logging.Logger, conn *net.TCPConn, _codec *codec.MsgpackHandle) *forwardClient {
	c := &forwardClient{
		input:  input,
		logger: logger,
		conn:   conn,
		codec:  _codec,
		dec:    codec.NewDecoder(bufio.NewReader(conn), _codec),
	}
	input.markCharged(c)
	return c
}

func (input *ForwardInput) spawnAcceptor() {
	input.logger.Notice("Spawning acceptor")
	input.wg.Add(1)
	go func() {
		defer func() {
			close(input.acceptChan)
			input.wg.Done()
		}()
		input.logger.Notice("Acceptor started")
		for {
			conn, err := input.listener.AcceptTCP()
			if err != nil {
				input.logger.Notice(err.Error())
				break
			}
			if conn != nil {
				input.logger.Notice("Connected from %s", conn.RemoteAddr().String())
				input.acceptChan <- conn
			} else {
				input.logger.Notice("AcceptTCP returned nil; something went wrong")
				break
			}
		}
		input.logger.Notice("Acceptor ended")
	}()
}

func (input *ForwardInput) spawnDaemon() {
	input.logger.Notice("Spawning daemon")
	input.wg.Add(1)
	go func() {
		defer func() {
			close(input.shutdownChan)
			input.wg.Done()
		}()
		input.logger.Notice("Daemon started")
	loop:
		for {
			select {
			case conn := <-input.acceptChan:
				if conn != nil {
					input.logger.Notice("Got conn from acceptChan")
					newForwardClient(input, input.logger, conn, input.codec).startHandling()
				}
			case <-input.shutdownChan:
				input.listener.Close()
				for _, client := range input.clients {
					client.shutdown()
				}
				break loop
			}
		}
		input.logger.Notice("Daemon ended")
	}()
}

func (input *ForwardInput) markCharged(c *forwardClient) {
	input.clientsMtx.Lock()
	defer input.clientsMtx.Unlock()
	input.clients[c.conn] = c
}

func (input *ForwardInput) markDischarged(c *forwardClient) {
	input.clientsMtx.Lock()
	defer input.clientsMtx.Unlock()
	delete(input.clients, c.conn)
}

func (input *ForwardInput) String() string {
	return "input"
}

func (input *ForwardInput) Start() {
	input.spawnAcceptor()
	input.spawnDaemon()
}

func (input *ForwardInput) WaitForShutdown() {
	input.wg.Wait()
}

func (input *ForwardInput) Stop() {
	if atomic.CompareAndSwapUintptr(&input.isShuttingDown, uintptr(0), uintptr(1)) {
		input.shutdownChan <- struct{}{}
	}
}

func NewForwardInput(logger *logging.Logger, bind string, port Port) (*ForwardInput, error) {
	_codec := codec.MsgpackHandle{}
	_codec.MapType = reflect.TypeOf(map[string]interface{}(nil))
	_codec.RawToString = false
	addr, err := net.ResolveTCPAddr("tcp", bind)
	if err != nil {
		logger.Error("%s", err.Error())
		return nil, err
	}
	listener, err := net.ListenTCP("tcp", addr)
	if err != nil {
		logger.Error("%s", err.Error())
		return nil, err
	}
	return &ForwardInput{
		port:           port,
		logger:         logger,
		bind:           bind,
		listener:       listener,
		codec:          &_codec,
		clients:        make(map[*net.TCPConn]*forwardClient),
		clientsMtx:     sync.Mutex{},
		entries:        0,
		wg:             sync.WaitGroup{},
		acceptChan:     make(chan *net.TCPConn),
		shutdownChan:   make(chan struct{}),
		isShuttingDown: uintptr(0),
	}, nil
}
<commit_msg>Support PackedForward messages correctly<commit_after>//
// Fluentd Forwarder
//
// Copyright (C) 2014 Treasure Data, Inc.
//
// 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 fluentd_forwarder

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	logging "github.com/op/go-logging"
	"github.com/ugorji/go/codec"
	"io"
	"net"
	"reflect"
	"sync"
	"sync/atomic"
)

type forwardClient struct {
	input  *ForwardInput
	logger *logging.Logger
	conn   *net.TCPConn
	codec  *codec.MsgpackHandle
	dec    *codec.Decoder
}

type ForwardInput struct {
	entries        int64 // This variable must be on 64-bit alignment. Otherwise atomic.AddInt64 will cause a crash on ARM and x86-32
	port           Port
	logger         *logging.Logger
	bind           string
	listener       *net.TCPListener
	codec          *codec.MsgpackHandle
	clientsMtx     sync.Mutex
	clients        map[*net.TCPConn]*forwardClient
	wg             sync.WaitGroup
	acceptChan     chan *net.TCPConn
	shutdownChan   chan struct{}
	isShuttingDown uintptr
}

type EntryCountTopic struct{}

type ConnectionCountTopic struct{}

type ForwardInputFactory struct{}

func coerceInPlace(data map[string]interface{}) {
	for k, v := range data {
		switch v_ := v.(type) {
		case []byte:
			data[k] = string(v_) // XXX: byte => rune
		case map[string]interface{}:
			coerceInPlace(v_)
		}
	}
}

func (c *forwardClient) decodeRecordSet(tag []byte, entries []interface{}) (FluentRecordSet, error) {
	records := make([]TinyFluentRecord, len(entries))
	for i, _entry := range entries {
		entry, ok := _entry.([]interface{})
		if !ok {
			return FluentRecordSet{}, errors.New("Failed to decode recordSet")
		}
		timestamp, ok := entry[0].(uint64)
		if !ok {
			return FluentRecordSet{}, errors.New("Failed to decode timestamp field")
		}
		data, ok := entry[1].(map[string]interface{})
		if !ok {
			return FluentRecordSet{}, errors.New("Failed to decode data field")
		}
		coerceInPlace(data)
		records[i] = TinyFluentRecord{
			Timestamp: timestamp,
			Data:      data,
		}
	}
	return FluentRecordSet{
		Tag:     string(tag), // XXX: byte => rune
		Records: records,
	}, nil
}

func (c *forwardClient) decodeEntries() ([]FluentRecordSet, error) {
	v := []interface{}{nil, nil, nil}
	err := c.dec.Decode(&v)
	if err != nil {
		return nil, err
	}
	tag, ok := v[0].([]byte)
	if !ok {
		return nil, errors.New("Failed to decode tag field")
	}

	var retval []FluentRecordSet
	switch timestamp_or_entries := v[1].(type) {
	case uint64:
		timestamp := timestamp_or_entries
		data, ok := v[2].(map[string]interface{})
		if !ok {
			return nil, errors.New("Failed to decode data field")
		}
		coerceInPlace(data)
		retval = []FluentRecordSet{
			{
				Tag: string(tag), // XXX: byte => rune
				Records: []TinyFluentRecord{
					{
						Timestamp: timestamp,
						Data:      data,
					},
				},
			},
		}
	case float64:
		timestamp := uint64(timestamp_or_entries)
		data, ok := v[2].(map[string]interface{})
		if !ok {
			return nil, errors.New("Failed to decode data field")
		}
		retval = []FluentRecordSet{
			{
				Tag: string(tag), // XXX: byte => rune
				Records: []TinyFluentRecord{
					{
						Timestamp: timestamp,
						Data:      data,
					},
				},
			},
		}
	case []interface{}:
		if !ok {
			return nil, errors.New("Unexpected payload format")
		}
		recordSet, err := c.decodeRecordSet(tag, timestamp_or_entries)
		if err != nil {
			return nil, err
		}
		retval = []FluentRecordSet{recordSet}
	case []byte:
		entries := make([]interface{}, 0)
		reader := bytes.NewReader(timestamp_or_entries)
		dec := codec.NewDecoder(reader, c.codec)
		for reader.Len() > 0 { // codec.Decoder doesn't return EOF.
			entry := []interface{}{}
			if err != dec.Decode(&entry) {
				if err == io.EOF { // in case codec.Decoder changes its behavior
					break
				}
				return nil, err
			}
			entries = append(entries, entry)
		}
		recordSet, err := c.decodeRecordSet(tag, entries)
		if err != nil {
			return nil, err
		}
		retval = []FluentRecordSet{recordSet}
	default:
		return nil, errors.New(fmt.Sprintf("Unknown type: %t", timestamp_or_entries))
	}
	atomic.AddInt64(&c.input.entries, int64(len(retval)))
	return retval, nil
}

func (c *forwardClient) startHandling() {
	c.input.wg.Add(1)
	go func() {
		defer func() {
			err := c.conn.Close()
			if err != nil {
				c.logger.Debug("Close: %s", err.Error())
			}
			c.input.markDischarged(c)
			c.input.wg.Done()
		}()
		c.input.logger.Info("Started handling connection from %s", c.conn.RemoteAddr().String())
		for {
			recordSets, err := c.decodeEntries()
			if err != nil {
				err_, ok := err.(net.Error)
				if ok {
					if err_.Temporary() {
						c.logger.Info("Temporary failure: %s", err_.Error())
						continue
					}
				}
				if err == io.EOF {
					c.logger.Info("Client %s closed the connection", c.conn.RemoteAddr().String())
				} else {
					c.logger.Error(err.Error())
				}
				break
			}

			if len(recordSets) > 0 {
				err_ := c.input.port.Emit(recordSets)
				if err_ != nil {
					c.logger.Error(err_.Error())
					break
				}
			}
		}
		c.input.logger.Info("Ended handling connection from %s", c.conn.RemoteAddr().String())
	}()
}

func (c *forwardClient) shutdown() {
	err := c.conn.Close()
	if err != nil {
		c.input.logger.Info("Error during closing connection: %s", err.Error())
	}
}

func newForwardClient(input *ForwardInput, logger *logging.Logger, conn *net.TCPConn, _codec *codec.MsgpackHandle) *forwardClient {
	c := &forwardClient{
		input:  input,
		logger: logger,
		conn:   conn,
		codec:  _codec,
		dec:    codec.NewDecoder(bufio.NewReader(conn), _codec),
	}
	input.markCharged(c)
	return c
}

func (input *ForwardInput) spawnAcceptor() {
	input.logger.Notice("Spawning acceptor")
	input.wg.Add(1)
	go func() {
		defer func() {
			close(input.acceptChan)
			input.wg.Done()
		}()
		input.logger.Notice("Acceptor started")
		for {
			conn, err := input.listener.AcceptTCP()
			if err != nil {
				input.logger.Notice(err.Error())
				break
			}
			if conn != nil {
				input.logger.Notice("Connected from %s", conn.RemoteAddr().String())
				input.acceptChan <- conn
			} else {
				input.logger.Notice("AcceptTCP returned nil; something went wrong")
				break
			}
		}
		input.logger.Notice("Acceptor ended")
	}()
}

func (input *ForwardInput) spawnDaemon() {
	input.logger.Notice("Spawning daemon")
	input.wg.Add(1)
	go func() {
		defer func() {
			close(input.shutdownChan)
			input.wg.Done()
		}()
		input.logger.Notice("Daemon started")
	loop:
		for {
			select {
			case conn := <-input.acceptChan:
				if conn != nil {
					input.logger.Notice("Got conn from acceptChan")
					newForwardClient(input, input.logger, conn, input.codec).startHandling()
				}
			case <-input.shutdownChan:
				input.listener.Close()
				for _, client := range input.clients {
					client.shutdown()
				}
				break loop
			}
		}
		input.logger.Notice("Daemon ended")
	}()
}

func (input *ForwardInput) markCharged(c *forwardClient) {
	input.clientsMtx.Lock()
	defer input.clientsMtx.Unlock()
	input.clients[c.conn] = c
}

func (input *ForwardInput) markDischarged(c *forwardClient) {
	input.clientsMtx.Lock()
	defer input.clientsMtx.Unlock()
	delete(input.clients, c.conn)
}

func (input *ForwardInput) String() string {
	return "input"
}

func (input *ForwardInput) Start() {
	input.spawnAcceptor()
	input.spawnDaemon()
}

func (input *ForwardInput) WaitForShutdown() {
	input.wg.Wait()
}

func (input *ForwardInput) Stop() {
	if atomic.CompareAndSwapUintptr(&input.isShuttingDown, uintptr(0), uintptr(1)) {
		input.shutdownChan <- struct{}{}
	}
}

func NewForwardInput(logger *logging.Logger, bind string, port Port) (*ForwardInput, error) {
	_codec := codec.MsgpackHandle{}
	_codec.MapType = reflect.TypeOf(map[string]interface{}(nil))
	_codec.RawToString = false
	addr, err := net.ResolveTCPAddr("tcp", bind)
	if err != nil {
		logger.Error("%s", err.Error())
		return nil, err
	}
	listener, err := net.ListenTCP("tcp", addr)
	if err != nil {
		logger.Error("%s", err.Error())
		return nil, err
	}
	return &ForwardInput{
		port:           port,
		logger:         logger,
		bind:           bind,
		listener:       listener,
		codec:          &_codec,
		clients:        make(map[*net.TCPConn]*forwardClient),
		clientsMtx:     sync.Mutex{},
		entries:        0,
		wg:             sync.WaitGroup{},
		acceptChan:     make(chan *net.TCPConn),
		shutdownChan:   make(chan struct{}),
		isShuttingDown: uintptr(0),
	}, nil
}
<|endoftext|>


<commit_before>package survey

import (
	"os"

	"github.com/AlecAivazis/survey/core"
	"github.com/AlecAivazis/survey/terminal"
)

// Input is a regular text input that prints each character the user types on the screen
// and accepts the input with the enter key.
type Input struct {
	core.Renderer
	Message string
	Default string
	Help    string
}

// data available to the templates when processing
type InputTemplateData struct {
	Input
	Answer     string
	ShowAnswer bool
	ShowHelp   bool
}

// Templates with Color formatting. See Documentation: https://github.com/mgutz/ansi#style-format
var InputQuestionTemplate = `
{{- if .ShowHelp }}{{- color "cyan"}}{{ HelpIcon }} {{ .Help }}{{color "reset"}}{{"\n"}}{{end}}
{{- color "green+hb"}}{{ QuestionIcon }} {{color "reset"}}
{{- color "default+hb"}}{{ .Message }} {{color "reset"}}
{{- if .ShowAnswer}}
  {{- color "cyan"}}{{.Answer}}{{color "reset"}}{{"\n"}}
{{- else }}
  {{- if and .Help (not .ShowHelp)}}{{color "cyan"}}[{{ HelpInputRune }} for help]{{color "reset"}} {{end}}
  {{- if .Default}}{{color "white"}}({{.Default}}) {{color "reset"}}{{end}}
{{- end}}`

func (i *Input) Prompt() (interface{}, error) {
	// render the template
	err := i.Render(
		InputQuestionTemplate,
		InputTemplateData{Input: *i},
	)
	if err != nil {
		return "", err
	}

	// start reading runes from the standard in
	rr := terminal.NewRuneReader(os.Stdin)
	rr.SetTermMode()
	defer rr.RestoreTermMode()

	line := []rune{}
	// get the next line
	for {
		line, err = rr.ReadLine(0)
		if err != nil {
			return string(line), err
		}
		// terminal will echo the \n so we need to jump back up one row
		terminal.CursorPreviousLine(1)

		if string(line) == string(core.HelpInputRune) && i.Help != "" {
			err = i.Render(
				InputQuestionTemplate,
				InputTemplateData{Input: *i, ShowHelp: true},
			)
			if err != nil {
				return "", err
			}
			continue
		}
		break
	}

	// if the line is empty
	if line == nil || len(line) == 0 {
		// use the default value
		return i.Default, err
	}

	// we're done
	return string(line), err
}

func (i *Input) Cleanup(val interface{}) error {
	return i.Render(
		InputQuestionTemplate,
		InputTemplateData{Input: *i, Answer: val.(string), ShowAnswer: true},
	)
}
<commit_msg>[#81] added response type doc for input prompt<commit_after>package survey

import (
	"os"

	"github.com/AlecAivazis/survey/core"
	"github.com/AlecAivazis/survey/terminal"
)

// Input is a regular text input that prints each character the user types on the screen
// and accepts the input with the enter key. Response type is a string.
type Input struct {
	core.Renderer
	Message string
	Default string
	Help    string
}

// data available to the templates when processing
type InputTemplateData struct {
	Input
	Answer     string
	ShowAnswer bool
	ShowHelp   bool
}

// Templates with Color formatting. See Documentation: https://github.com/mgutz/ansi#style-format
var InputQuestionTemplate = `
{{- if .ShowHelp }}{{- color "cyan"}}{{ HelpIcon }} {{ .Help }}{{color "reset"}}{{"\n"}}{{end}}
{{- color "green+hb"}}{{ QuestionIcon }} {{color "reset"}}
{{- color "default+hb"}}{{ .Message }} {{color "reset"}}
{{- if .ShowAnswer}}
  {{- color "cyan"}}{{.Answer}}{{color "reset"}}{{"\n"}}
{{- else }}
  {{- if and .Help (not .ShowHelp)}}{{color "cyan"}}[{{ HelpInputRune }} for help]{{color "reset"}} {{end}}
  {{- if .Default}}{{color "white"}}({{.Default}}) {{color "reset"}}{{end}}
{{- end}}`

func (i *Input) Prompt() (interface{}, error) {
	// render the template
	err := i.Render(
		InputQuestionTemplate,
		InputTemplateData{Input: *i},
	)
	if err != nil {
		return "", err
	}

	// start reading runes from the standard in
	rr := terminal.NewRuneReader(os.Stdin)
	rr.SetTermMode()
	defer rr.RestoreTermMode()

	line := []rune{}
	// get the next line
	for {
		line, err = rr.ReadLine(0)
		if err != nil {
			return string(line), err
		}
		// terminal will echo the \n so we need to jump back up one row
		terminal.CursorPreviousLine(1)

		if string(line) == string(core.HelpInputRune) && i.Help != "" {
			err = i.Render(
				InputQuestionTemplate,
				InputTemplateData{Input: *i, ShowHelp: true},
			)
			if err != nil {
				return "", err
			}
			continue
		}
		break
	}

	// if the line is empty
	if line == nil || len(line) == 0 {
		// use the default value
		return i.Default, err
	}

	// we're done
	return string(line), err
}

func (i *Input) Cleanup(val interface{}) error {
	return i.Render(
		InputQuestionTemplate,
		InputTemplateData{Input: *i, Answer: val.(string), ShowAnswer: true},
	)
}
<|endoftext|>


<commit_before>package template

import (
	"strings"
	"testing"

	nomad "github.com/hashicorp/nomad/api"
)

const (
	testJobName 			= "levantExample"
	testJobNameOverwrite 	= "levantExampleOverwrite"
	testJobNameOverwrite2 	= "levantExampleOverwrite2"
	testDCName  			= "dc13"
)

func TestTemplater_RenderTemplate(t *testing.T) {

	var job *nomad.Job
	var err error

	// Start with an empty passed var args map.
	fVars := make(map[string]string)

	// Test basic TF template render.
	job, err = RenderJob("test-fixtures/single_templated.nomad", []string{"test-fixtures/test.tf"}, "", &fVars)
	if err != nil {
		t.Fatal(err)
	}
	if *job.Name != testJobName {
		t.Fatalf("expected %s but got %v", testJobName, *job.Name)
	}

	// Test basic YAML template render.
	job, err = RenderJob("test-fixtures/single_templated.nomad", []string{"test-fixtures/test.yaml"}, "", &fVars)
	if err != nil {
		t.Fatal(err)
	}
	if *job.Name != testJobName {
		t.Fatalf("expected %s but got %v", testJobName, *job.Name)
	}

	// Test multiple var-files
	job, err = RenderJob("test-fixtures/single_templated.nomad", []string{"test-fixtures/test.yaml","test-fixtures/test-overwrite.yaml"}, "", &fVars)
	if err != nil {
		t.Fatal(err)
	}
	if *job.Name != testJobNameOverwrite {
		t.Fatalf("expected %s but got %v", testJobNameOverwrite, *job.Name)
	}

	// Test multiple var-files of different types
	job, err = RenderJob("test-fixtures/single_templated.nomad", []string{"test-fixtures/test.tf","test-fixtures/test-overwrite.yaml"}, "", &fVars)
	if err != nil {
		t.Fatal(err)
	}
	if *job.Name != testJobNameOverwrite {
		t.Fatalf("expected %s but got %v", testJobNameOverwrite, *job.Name)
	}

	// Test multiple var-files with var-args
	fVars["job_name"] = testJobNameOverwrite2
	job, err = RenderJob("test-fixtures/single_templated.nomad", []string{"test-fixtures/test.tf","test-fixtures/test-overwrite.yaml"}, "", &fVars)
	if err != nil {
		t.Fatal(err)
	}
	if *job.Name != testJobNameOverwrite2 {
		t.Fatalf("expected %s but got %v", testJobNameOverwrite2, *job.Name)
	}

	// Test empty var-args and empty variable file render.
	job, err = RenderJob("test-fixtures/none_templated.nomad", []string{}, "", &fVars)
	if err != nil {
		t.Fatal(err)
	}
	if *job.Name != testJobName {
		t.Fatalf("expected %s but got %v", testJobName, *job.Name)
	}

	// Test var-args only render.
	delete(fVars, "job_name")
	fVars["job_name"] = testJobName
	job, err = RenderJob("test-fixtures/single_templated.nomad", []string{}, "", &fVars)
	if err != nil {
		t.Fatal(err)
	}
	if *job.Name != testJobName {
		t.Fatalf("expected %s but got %v", testJobName, *job.Name)
	}

	// Test var-args and variables file render.
	delete(fVars, "job_name")
	fVars["datacentre"] = testDCName
	job, err = RenderJob("test-fixtures/multi_templated.nomad", []string{"test-fixtures/test.yaml"}, "", &fVars)
	if err != nil {
		t.Fatal(err)
	}
	if *job.Name != testJobName {
		t.Fatalf("expected %s but got %v", testJobName, *job.Name)
	}
	if job.Datacenters[0] != testDCName {
		t.Fatalf("expected %s but got %v", testDCName, job.Datacenters[0])
	}

	// Test var-args only render.
	fVars["job_name"] = testJobName
	_, err = RenderTemplate("test-fixtures/missing_var.nomad", []string{}, "", &fVars)
	if err == nil {
		t.Fatal("expected err to not be nil")
	}
	if !strings.Contains(err.Error(), "binary_url") {
		t.Fatal("expected err to mention missing var (binary_url)")
	}
}
<commit_msg>#179 Support multiple var-files - formatting<commit_after>package template

import (
	"strings"
	"testing"

	nomad "github.com/hashicorp/nomad/api"
)

const (
	testJobName           = "levantExample"
	testJobNameOverwrite  = "levantExampleOverwrite"
	testJobNameOverwrite2 = "levantExampleOverwrite2"
	testDCName            = "dc13"
)

func TestTemplater_RenderTemplate(t *testing.T) {

	var job *nomad.Job
	var err error

	// Start with an empty passed var args map.
	fVars := make(map[string]string)

	// Test basic TF template render.
	job, err = RenderJob("test-fixtures/single_templated.nomad", []string{"test-fixtures/test.tf"}, "", &fVars)
	if err != nil {
		t.Fatal(err)
	}
	if *job.Name != testJobName {
		t.Fatalf("expected %s but got %v", testJobName, *job.Name)
	}

	// Test basic YAML template render.
	job, err = RenderJob("test-fixtures/single_templated.nomad", []string{"test-fixtures/test.yaml"}, "", &fVars)
	if err != nil {
		t.Fatal(err)
	}
	if *job.Name != testJobName {
		t.Fatalf("expected %s but got %v", testJobName, *job.Name)
	}

	// Test multiple var-files
	job, err = RenderJob("test-fixtures/single_templated.nomad", []string{"test-fixtures/test.yaml", "test-fixtures/test-overwrite.yaml"}, "", &fVars)
	if err != nil {
		t.Fatal(err)
	}
	if *job.Name != testJobNameOverwrite {
		t.Fatalf("expected %s but got %v", testJobNameOverwrite, *job.Name)
	}

	// Test multiple var-files of different types
	job, err = RenderJob("test-fixtures/single_templated.nomad", []string{"test-fixtures/test.tf", "test-fixtures/test-overwrite.yaml"}, "", &fVars)
	if err != nil {
		t.Fatal(err)
	}
	if *job.Name != testJobNameOverwrite {
		t.Fatalf("expected %s but got %v", testJobNameOverwrite, *job.Name)
	}

	// Test multiple var-files with var-args
	fVars["job_name"] = testJobNameOverwrite2
	job, err = RenderJob("test-fixtures/single_templated.nomad", []string{"test-fixtures/test.tf", "test-fixtures/test-overwrite.yaml"}, "", &fVars)
	if err != nil {
		t.Fatal(err)
	}
	if *job.Name != testJobNameOverwrite2 {
		t.Fatalf("expected %s but got %v", testJobNameOverwrite2, *job.Name)
	}

	// Test empty var-args and empty variable file render.
	job, err = RenderJob("test-fixtures/none_templated.nomad", []string{}, "", &fVars)
	if err != nil {
		t.Fatal(err)
	}
	if *job.Name != testJobName {
		t.Fatalf("expected %s but got %v", testJobName, *job.Name)
	}

	// Test var-args only render.
	delete(fVars, "job_name")
	fVars["job_name"] = testJobName
	job, err = RenderJob("test-fixtures/single_templated.nomad", []string{}, "", &fVars)
	if err != nil {
		t.Fatal(err)
	}
	if *job.Name != testJobName {
		t.Fatalf("expected %s but got %v", testJobName, *job.Name)
	}

	// Test var-args and variables file render.
	delete(fVars, "job_name")
	fVars["datacentre"] = testDCName
	job, err = RenderJob("test-fixtures/multi_templated.nomad", []string{"test-fixtures/test.yaml"}, "", &fVars)
	if err != nil {
		t.Fatal(err)
	}
	if *job.Name != testJobName {
		t.Fatalf("expected %s but got %v", testJobName, *job.Name)
	}
	if job.Datacenters[0] != testDCName {
		t.Fatalf("expected %s but got %v", testDCName, job.Datacenters[0])
	}

	// Test var-args only render.
	fVars["job_name"] = testJobName
	_, err = RenderTemplate("test-fixtures/missing_var.nomad", []string{}, "", &fVars)
	if err == nil {
		t.Fatal("expected err to not be nil")
	}
	if !strings.Contains(err.Error(), "binary_url") {
		t.Fatal("expected err to mention missing var (binary_url)")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"time"

	"github.com/Shopify/sarama"
	"github.com/redBorder/rbforwarder"
	"github.com/wvanbergen/kafka/consumergroup"
	"github.com/wvanbergen/kazoo-go"
)

// KafkaConsumer get messages from multiple kafka topics
type KafkaConsumer struct {
	backend       *rbforwarder.RBForwarder     // The backend to send messages
	consumerGroup *consumergroup.ConsumerGroup // The main kafka consumer
	Config        KafkaConfig                  // Cofiguration after the parsing
}

// KafkaConfig stores the configuration for the Kafka source
type KafkaConfig struct {
	topics              []string // Topics where listen for messages
	brokers             []string // Brokers to connect
	consumergroup       string   // ID for the consumer
	consumerGroupConfig *consumergroup.Config
}

// Start starts reading messages from kafka and pushing them to the pipeline
func (k *KafkaConsumer) Start() {
	var err error

	// Join the consumer group
	k.consumerGroup, err = consumergroup.JoinConsumerGroup(
		k.Config.consumergroup,
		k.Config.topics,
		k.Config.brokers,
		k.Config.consumerGroupConfig,
	)
	if err != nil {
		logger.Fatal(err)
	}

	// Check for errors
	go func() {
		for err := range k.consumerGroup.Errors() {
			logger.Error(err)
		}
	}()

	var offset uint64
	var eventsReported uint64
	var eventsSent uint64

	// Start processing reports
	done := make(chan struct{})
	go func() {
		for report := range k.backend.GetOrderedReports() {
			message := report.Metadata["sarama_message"].(*sarama.ConsumerMessage)

			if offset != report.ID {
				logger.Fatalf("Unexpected offset. Expected %d, found %d.",
					offset, report.ID)
			}

			if report.StatusCode != 0 {
				logger.
					WithField("ID", report.ID).
					WithField("STATUS", report.Status).
					WithField("OFFSET", message.Offset).
					Errorf("REPORT")
			} else {
				logger.
					WithField("ID", report.ID).
					WithField("STATUS", report.Status).
					WithField("OFFSET", message.Offset).
					Debugf("REPORT")
			}

			if err := k.consumerGroup.CommitUpto(message); err != nil {
				logger.Fatal(err)
			}
			offset++
			eventsReported++
		}

		done <- struct{}{}
	}()

	// Start consuming messages
	for message := range k.consumerGroup.Messages() {
		msg, err := k.backend.TakeMessage()
		if err != nil {
			break
		}

		if _, err := msg.InputBuffer.Write(message.Value); err != nil {
			logger.Error(err)
		}

		msg.Metadata["sarama_message"] = message
		msg.Metadata["topic"] = message.Topic
		if err := msg.Produce(); err != nil {
			break
		}

		eventsSent++
	}

	logger.Info("Consumer terminated")

	<-done
	logger.Infof("TOTAL SENT MESSAGES: %d", eventsSent)
	logger.Infof("TOTAL REPORTS: %d", eventsReported)

	return
}

// Close closes the connection with Kafka
func (k *KafkaConsumer) Close() {
	if err := k.consumerGroup.Close(); err != nil {
		logger.Errorf("Error closing the consumer: %s", err)
	}
}

// ParseKafkaConfig reads the configuration from the YAML config file and store it
// on the instance
func (k *KafkaConsumer) ParseKafkaConfig(config map[string]interface{}) {

	// Create the config
	k.Config.consumerGroupConfig = consumergroup.NewConfig()

	k.Config.consumerGroupConfig.Offsets.ProcessingTimeout = 5 * time.Second

	// Parse the brokers addresses
	if val, ok := config["begining"].(bool); ok {
		k.Config.consumerGroupConfig.Offsets.ResetOffsets = val
	}

	// Parse the brokers addresses
	if config["broker"] != nil {
		k.Config.brokers, k.Config.consumerGroupConfig.Zookeeper.Chroot =
			kazoo.ParseConnectionString(config["broker"].(string))
	}

	// Parse topics
	if config["topics"] != nil {
		topics := config["topics"].([]interface{})

		for _, topic := range topics {
			k.Config.topics = append(k.Config.topics, topic.(string))
		}
	}

	// Parse consumergroup
	if config["consumergroup"] != nil {
		k.Config.consumergroup = config["consumergroup"].(string)
	}
}
<commit_msg>Modified offset processing time<commit_after>package main

import (
	"time"

	"github.com/Shopify/sarama"
	"github.com/redBorder/rbforwarder"
	"github.com/wvanbergen/kafka/consumergroup"
	"github.com/wvanbergen/kazoo-go"
)

// KafkaConsumer get messages from multiple kafka topics
type KafkaConsumer struct {
	backend       *rbforwarder.RBForwarder     // The backend to send messages
	consumerGroup *consumergroup.ConsumerGroup // The main kafka consumer
	Config        KafkaConfig                  // Cofiguration after the parsing
}

// KafkaConfig stores the configuration for the Kafka source
type KafkaConfig struct {
	topics              []string // Topics where listen for messages
	brokers             []string // Brokers to connect
	consumergroup       string   // ID for the consumer
	consumerGroupConfig *consumergroup.Config
}

// Start starts reading messages from kafka and pushing them to the pipeline
func (k *KafkaConsumer) Start() {
	var err error

	// Join the consumer group
	k.consumerGroup, err = consumergroup.JoinConsumerGroup(
		k.Config.consumergroup,
		k.Config.topics,
		k.Config.brokers,
		k.Config.consumerGroupConfig,
	)
	if err != nil {
		logger.Fatal(err)
	}

	// Check for errors
	go func() {
		for err := range k.consumerGroup.Errors() {
			logger.Error(err)
		}
	}()

	var offset uint64
	var eventsReported uint64
	var eventsSent uint64

	// Start processing reports
	done := make(chan struct{})
	go func() {
		for report := range k.backend.GetOrderedReports() {
			message := report.Metadata["sarama_message"].(*sarama.ConsumerMessage)

			if offset != report.ID {
				logger.Fatalf("Unexpected offset. Expected %d, found %d.",
					offset, report.ID)
			}

			if report.StatusCode != 0 {
				logger.
					WithField("ID", report.ID).
					WithField("STATUS", report.Status).
					WithField("OFFSET", message.Offset).
					Errorf("REPORT")
			} else {
				logger.
					WithField("ID", report.ID).
					WithField("STATUS", report.Status).
					WithField("OFFSET", message.Offset).
					Debugf("REPORT")
			}

			if err := k.consumerGroup.CommitUpto(message); err != nil {
				logger.Fatal(err)
			}
			offset++
			eventsReported++
		}

		done <- struct{}{}
	}()

	// Start consuming messages
	for message := range k.consumerGroup.Messages() {
		msg, err := k.backend.TakeMessage()
		if err != nil {
			break
		}

		if _, err := msg.InputBuffer.Write(message.Value); err != nil {
			logger.Error(err)
		}

		msg.Metadata["sarama_message"] = message
		msg.Metadata["topic"] = message.Topic
		if err := msg.Produce(); err != nil {
			break
		}

		eventsSent++
	}

	logger.Info("Consumer terminated")

	<-done
	logger.Infof("TOTAL SENT MESSAGES: %d", eventsSent)
	logger.Infof("TOTAL REPORTS: %d", eventsReported)

	return
}

// Close closes the connection with Kafka
func (k *KafkaConsumer) Close() {
	if err := k.consumerGroup.Close(); err != nil {
		logger.Errorf("Error closing the consumer: %s", err)
	}
}

// ParseKafkaConfig reads the configuration from the YAML config file and store it
// on the instance
func (k *KafkaConsumer) ParseKafkaConfig(config map[string]interface{}) {

	// Create the config
	k.Config.consumerGroupConfig = consumergroup.NewConfig()

	k.Config.consumerGroupConfig.Offsets.ProcessingTimeout = 1 * time.Second

	// Parse the brokers addresses
	if val, ok := config["begining"].(bool); ok {
		k.Config.consumerGroupConfig.Offsets.ResetOffsets = val
	}

	// Parse the brokers addresses
	if config["broker"] != nil {
		k.Config.brokers, k.Config.consumerGroupConfig.Zookeeper.Chroot =
			kazoo.ParseConnectionString(config["broker"].(string))
	}

	// Parse topics
	if config["topics"] != nil {
		topics := config["topics"].([]interface{})

		for _, topic := range topics {
			k.Config.topics = append(k.Config.topics, topic.(string))
		}
	}

	// Parse consumergroup
	if config["consumergroup"] != nil {
		k.Config.consumergroup = config["consumergroup"].(string)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	ab "github.com/gernest/authboss"
	_ "github.com/gernest/authboss/auth"
	_ "github.com/gernest/authboss/register"
	_ "github.com/gernest/authboss/remember"
	"github.com/gorilla/mux"
	"github.com/justinas/nosurf"
	"github.com/monoculum/formam"
	"html/template"
	"log"
	"net/http"
	"os"
	"time"
	"github.com/boltdb/bolt"
	"github.com/gorilla/sessions"
)

const (
	VERSION = "0.0.1"
)

func init() {
	log.SetFlags(log.Lshortfile)
}

var funcs = template.FuncMap{
	"formatDate": func(date time.Time) string {
		return date.Format("2006/01/02 03:04pm")
	},
	"yield": func() string { return "" },
}

type Config struct {
	AccountsBucket  string
	TemplatesBucket string
	SessionBucket   string
	SessionName string
	DefaultTemplate string

	SessionDB string
	MainDB    string

	Secretts  []string
}

type Kesho struct {
	AccountsBucket  string
	TemplatesBucket string
	SessionBucket   string

	Store       Storage
	Assets      *Assets
	Templ       *KTemplate
	SessStore   *BStore
	SessionName string

	DefaultTemplate string

	SessionDB string
	MainDB    string
	Secretts  []string

	Cfg *Config
}

func NewKesho(cfg *Config) *Kesho {
	return &Kesho{Cfg:cfg}
}

func (k *Kesho)Initialize() {
	k.initDefaults()
}

func (k *Kesho)initDefaults() {
	k.AccountsBucket="accounts"
	k.TemplatesBucket="templates"
	k.SessionBucket="sessions"
	k.DefaultTemplate="web"
	k.SessionName="kesho_"

	kStore := NewStorage("mainstore.db", 0660)
	ass := &Assets{Bucket:"assets", Store:kStore}

	db, err := bolt.Open("sessions.db", 0600, nil)
	if err!=nil {
		log.Println(err)
	}
	secrets := []string{
		"892252c6eade0b4ebf32d94aaed79d20",
		"9451243db34445f4dbf86e0b13bec94d",
	}
	opts := &sessions.Options{MaxAge: 86400 * 30, Path: "/"}
	ss, err := NewBStoreFromDB(db, "sessions", 100, opts, []byte(secrets[0]), []byte(secrets[1]))
	if err != nil {
		log.Println(err)
	}
	k.Store=kStore
	k.Assets=ass
	k.Templ=&KTemplate{Store:kStore, Bucket:k.TemplatesBucket, Assets:ass}
	k.SessStore=ss
}

// Our HomePage
func (k *Kesho) HomePage(w http.ResponseWriter, r *http.Request) {
	k.RenderDefaultView(w, "index.html", nil)
	return
}

// Accounts
func (k *Kesho) AccountHome(w http.ResponseWriter, r *http.Request) {
	currentUsr, err := ab.CurrentUser(w, r)
	if err != nil {
		k.ServerProblem(w, err.Error())
		return
	}
	if currentUsr == nil {
		http.Redirect(w, r, "/auth/login", http.StatusFound)
		return
	}
	usr := currentUsr.(*Account)

	data := NewHtmlData()
	data.Set("Title", "Account")
	data.SetUser(usr)
	data.StausLogged()
	k.RenderDefaultView(w, "accounts/index.html", data.Data())
}

// Posts
func (k *Kesho) PostCreate(w http.ResponseWriter, r *http.Request) {
	log.Println("Creating new post")
	currentUsr, err := ab.CurrentUser(w, r)
	if err != nil {
		k.ServerProblem(w, err.Error())
		return
	}
	if currentUsr == nil {
		http.Redirect(w, r, "/auth/login", http.StatusFound)
		return
	}
	usr := currentUsr.(*Account)
	data := NewHtmlData()

	r.ParseForm()
	post := new(Post)
	if err = formam.Decode(r.Form, post); err != nil {
		k.ServerProblem(w, err.Error())
		return
	}
	post.Account = usr
	err = post.Create()
	if err != nil {
		k.ServerProblem(w, err.Error())
		return
	}
	data.SetSafe("Title", usr.UserName)
	data.Set("post", post)
	data.FlashSuccess(post.Title + " imepokelewa na kutangazwa")
	k.RenderDefaultView(w, "accounts/index.html", data.Data())
}

func (k *Kesho) PostUpdate(w http.ResponseWriter, r *http.Request) {
	currentUsr, err := ab.CurrentUser(w, r)
	if err != nil {
		k.ServerProblem(w, err.Error())
		return
	}
	if currentUsr == nil {
		http.Redirect(w, r, "/auth/login", http.StatusFound)
		return
	}

	vars := mux.Vars(r)
	postSlug := vars["slug"]
	if postSlug == "" {
		k.NotFound(w, r)
		return
	}
	usr := currentUsr.(*Account)
	post := new(Post)
	post.Account = usr
	post.Slug = postSlug
	err = post.Get()
	if err != nil {
		k.ServerProblem(w, err.Error())
		return
	}
	data := NewHtmlData()
	data.Set("Title", post.Title)
	data.SetUser(usr)
	data.StausLogged()
	data.Set("post", post)
	if r.Method == "POST" {
		r.ParseForm()
		if err = formam.Decode(r.Form, post); err != nil {
			k.ServerProblem(w, err.Error())
			return
		}
		if err = post.Update(); err != nil {
			k.ServerProblem(w, err.Error())
			return
		}
		data.FlashSuccess("Updated " + post.Title)
		k.RenderDefaultView(w, "accounts/index.html", data.Data())
		return
	}
	k.RenderDefaultView(w, "post/update.html", data.Data())
}

func (k *Kesho) PostDelete(w http.ResponseWriter, r *http.Request) {}

func (k *Kesho) PostView(w http.ResponseWriter, r *http.Request) {}

// Version
func (k *Kesho) Version(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte(VERSION))
	return
}

// Views
func (k *Kesho) ViewHome(w http.ResponseWriter, r *http.Request) {}
func (k *Kesho) ViewPost(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	uname := vars["username"]
	if uname == "" {
		k.NotFound(w, r)
		return
	}
	pslug := vars["slug"]
	if pslug == "" {
		k.NotFound(w, r)
		return
	}
	user := NewAccount(k.AccountsBucket, k.Store)
	user.UserName = uname
	if err := user.Get(); err != nil {
		k.ServerProblem(w, err.Error())
		return
	}
	post := new(Post)
	post.Slug = pslug
	post.Account = user
	if err := post.Get(); err != nil {
		k.ServerProblem(w, err.Error())
		return
	}
	data := NewHtmlData()
	data.Set("user", user)
	data.Set("post", post)
	k.RenderDefaultView(w, "post/post.html", data.Data())
}
func (k *Kesho) ViewSubHome(w http.ResponseWriter, r *http.Request) {}
func (k *Kesho) ViewSubPost(w http.ResponseWriter, r *http.Request) {}

func (k *Kesho) RenderDefaultView(w http.ResponseWriter, name string, data interface{}) {
	out := new(bytes.Buffer)
	err := k.Templ.Render(out, k.DefaultTemplate, name, data)

	if err != nil {
		log.Println(err, k.DefaultTemplate, name)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	w.Write(out.Bytes())
}

func (k *Kesho) Setup() {
	database := AccountAuth{k.Store, k.AccountsBucket, "tokens_"}
	ab.Cfg.Storer = database
	ab.Cfg.OAuth2Storer = database
	ab.Cfg.MountPath = "/auth"
	ab.Cfg.ViewsPath = ""
	ab.Cfg.ResponseTmpl = k.Templ.AuthTempl
	ab.Cfg.LogWriter = os.Stdout
	ab.Cfg.RootURL = `http://localhost:8080`

	ab.Cfg.LayoutDataMaker = k.AuthlayoutData
	ab.Cfg.XSRFName = "csrf_token"
	ab.Cfg.XSRFMaker = func(_ http.ResponseWriter, r *http.Request) string {
		return nosurf.Token(r)
	}
	ab.Cfg.PrimaryID = "user_name"

	ab.Cfg.Policies = []ab.Validator{
		ab.Rules{
			FieldName:       "user_name",
			Required:        true,
			MinLength:       5,
			MaxLength:       10,
			AllowWhitespace: false,
		},
		ab.Rules{
			FieldName:       "password",
			Required:        true,
			MinLength:       8,
			MaxLength:       20,
			AllowWhitespace: false,
		},
	}
	ab.Cfg.ConfirmFields = []string{"password", "confirm_password"}

	ab.Cfg.CookieStoreMaker = k.NewSessionStorer
	ab.Cfg.SessionStoreMaker = k.NewSessionStorer

	if err := ab.Init(); err != nil {
		log.Fatal(err)
	}

	if err := k.Templ.LoadToDB(k.DefaultTemplate); err != nil {
		log.Println(err)
	}
}

func (k Kesho) Run() {
	var (
		httpPort = "8080"
	)
	k.Initialize()
	k.Setup()
	log.Printf("Kesho is running at localhost:%s \n", httpPort)
	addr := fmt.Sprintf(":%s", httpPort)
	log.Fatal(http.ListenAndServe(addr, k.Routes()))

}
func (k *Kesho)Cleanup() error {
	os.Remove(k.MainDB)
	return os.Remove(k.SessionDB)
}<commit_msg>fix sessions<commit_after>package main

import (
	"bytes"
	"fmt"
	"github.com/boltdb/bolt"
	ab "github.com/gernest/authboss"
	_ "github.com/gernest/authboss/auth"
	_ "github.com/gernest/authboss/register"
	_ "github.com/gernest/authboss/remember"
	"github.com/gorilla/mux"
	"github.com/gorilla/sessions"
	"github.com/justinas/nosurf"
	"github.com/monoculum/formam"
	"html/template"
	"log"
	"net/http"
	"os"
	"time"
)

const (
	VERSION = "0.0.1"
)

func init() {
	log.SetFlags(log.Lshortfile)
}

var funcs = template.FuncMap{
	"formatDate": func(date time.Time) string {
		return date.Format("2006/01/02 03:04pm")
	},
	"yield": func() string { return "" },
}

type Config struct {
	AccountsBucket  string
	TemplatesBucket string
	SessionBucket   string
	SessionName     string
	DefaultTemplate string

	SessionDB string
	MainDB    string

	Secretts []string
}

type Kesho struct {
	AccountsBucket  string
	TemplatesBucket string
	SessionBucket   string

	Store       Storage
	Assets      *Assets
	Templ       *KTemplate
	SessStore   *BStore
	SessionName string

	DefaultTemplate string

	SessionDB string
	MainDB    string
	Secretts  []string

	Cfg *Config
}

func NewKesho(cfg *Config) *Kesho {
	return &Kesho{Cfg: cfg}
}

func (k *Kesho) Initialize() {
	k.initDefaults()
}

func (k *Kesho) initDefaults() {
	k.AccountsBucket = "accounts"
	k.TemplatesBucket = "templates"
	k.SessionBucket = "sessions"
	k.DefaultTemplate = "web"
	k.SessionName = "kesho_"

	kStore := NewStorage("mainstore.db", 0660)
	ass := &Assets{Bucket: "assets", Store: kStore}

	db, err := bolt.Open("sessions.db", 0600, nil)
	if err != nil {
		log.Println(err)
	}
	secrets := []string{
		"892252c6eade0b4ebf32d94aaed79d20",
		"9451243db34445f4dbf86e0b13bec94d",
	}
	opts := &sessions.Options{MaxAge: 86400 * 30, Path: "/"}
	ss, err := NewBStoreFromDB(db, "sessions", 100, opts, []byte(secrets[0]), []byte(secrets[1]))
	if err != nil {
		log.Println(err)
	}
	k.Store = kStore
	k.Assets = ass
	k.Templ = &KTemplate{Store: kStore, Bucket: k.TemplatesBucket, Assets: ass}
	k.SessStore = ss
}

// Our HomePage
func (k *Kesho) HomePage(w http.ResponseWriter, r *http.Request) {
	k.RenderDefaultView(w, "index.html", nil)
	return
}

// Accounts
func (k *Kesho) AccountHome(w http.ResponseWriter, r *http.Request) {
	currentUsr, err := ab.CurrentUser(w, r)
	if err != nil {
		k.ServerProblem(w, err.Error())
		return
	}
	if currentUsr == nil {
		http.Redirect(w, r, "/auth/login", http.StatusFound)
		return
	}
	usr := currentUsr.(*Account)

	data := NewHtmlData()
	data.Set("Title", "Account")
	data.SetUser(usr)
	data.StausLogged()
	k.RenderDefaultView(w, "accounts/index.html", data.Data())
}

// Posts
func (k *Kesho) PostCreate(w http.ResponseWriter, r *http.Request) {
	log.Println("Creating new post")
	currentUsr, err := ab.CurrentUser(w, r)
	if err != nil {
		k.ServerProblem(w, err.Error())
		return
	}
	if currentUsr == nil {
		http.Redirect(w, r, "/auth/login", http.StatusFound)
		return
	}
	usr := currentUsr.(*Account)
	data := NewHtmlData()

	r.ParseForm()
	post := new(Post)
	if err = formam.Decode(r.Form, post); err != nil {
		k.ServerProblem(w, err.Error())
		return
	}
	post.Account = usr
	err = post.Create()
	if err != nil {
		k.ServerProblem(w, err.Error())
		return
	}
	data.SetSafe("Title", usr.UserName)
	data.Set("post", post)
	data.FlashSuccess(post.Title + " imepokelewa na kutangazwa")
	k.RenderDefaultView(w, "accounts/index.html", data.Data())
}

func (k *Kesho) PostUpdate(w http.ResponseWriter, r *http.Request) {
	currentUsr, err := ab.CurrentUser(w, r)
	if err != nil {
		k.ServerProblem(w, err.Error())
		return
	}
	if currentUsr == nil {
		http.Redirect(w, r, "/auth/login", http.StatusFound)
		return
	}

	vars := mux.Vars(r)
	postSlug := vars["slug"]
	if postSlug == "" {
		k.NotFound(w, r)
		return
	}
	usr := currentUsr.(*Account)
	post := new(Post)
	post.Account = usr
	post.Slug = postSlug
	err = post.Get()
	if err != nil {
		k.ServerProblem(w, err.Error())
		return
	}
	data := NewHtmlData()
	data.Set("Title", post.Title)
	data.SetUser(usr)
	data.StausLogged()
	data.Set("post", post)
	if r.Method == "POST" {
		r.ParseForm()
		if err = formam.Decode(r.Form, post); err != nil {
			k.ServerProblem(w, err.Error())
			return
		}
		if err = post.Update(); err != nil {
			k.ServerProblem(w, err.Error())
			return
		}
		data.FlashSuccess("Updated " + post.Title)
		k.RenderDefaultView(w, "accounts/index.html", data.Data())
		return
	}
	k.RenderDefaultView(w, "post/update.html", data.Data())
}

func (k *Kesho) PostDelete(w http.ResponseWriter, r *http.Request) {}

func (k *Kesho) PostView(w http.ResponseWriter, r *http.Request) {}

// Version
func (k *Kesho) Version(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte(VERSION))
	return
}

// Views
func (k *Kesho) ViewHome(w http.ResponseWriter, r *http.Request) {}
func (k *Kesho) ViewPost(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	uname := vars["username"]
	if uname == "" {
		k.NotFound(w, r)
		return
	}
	pslug := vars["slug"]
	if pslug == "" {
		k.NotFound(w, r)
		return
	}
	user := NewAccount(k.AccountsBucket, k.Store)
	user.UserName = uname
	if err := user.Get(); err != nil {
		k.ServerProblem(w, err.Error())
		return
	}
	post := new(Post)
	post.Slug = pslug
	post.Account = user
	if err := post.Get(); err != nil {
		k.ServerProblem(w, err.Error())
		return
	}
	data := NewHtmlData()
	data.Set("user", user)
	data.Set("post", post)
	k.RenderDefaultView(w, "post/post.html", data.Data())
}
func (k *Kesho) ViewSubHome(w http.ResponseWriter, r *http.Request) {}
func (k *Kesho) ViewSubPost(w http.ResponseWriter, r *http.Request) {}

func (k *Kesho) RenderDefaultView(w http.ResponseWriter, name string, data interface{}) {
	out := new(bytes.Buffer)
	err := k.Templ.Render(out, k.DefaultTemplate, name, data)

	if err != nil {
		log.Println(err, k.DefaultTemplate, name)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	w.Write(out.Bytes())
}

func (k *Kesho) Setup() {
	database := AccountAuth{k.Store, k.AccountsBucket, "tokens_"}
	ab.Cfg.Storer = database
	ab.Cfg.OAuth2Storer = database
	ab.Cfg.MountPath = "/auth"
	ab.Cfg.ViewsPath = ""
	ab.Cfg.ResponseTmpl = k.Templ.AuthTempl
	ab.Cfg.LogWriter = os.Stdout
	ab.Cfg.RootURL = `http://localhost:8080`

	ab.Cfg.LayoutDataMaker = k.AuthlayoutData
	ab.Cfg.XSRFName = "csrf_token"
	ab.Cfg.XSRFMaker = func(_ http.ResponseWriter, r *http.Request) string {
		return nosurf.Token(r)
	}
	ab.Cfg.PrimaryID = "user_name"

	ab.Cfg.Policies = []ab.Validator{
		ab.Rules{
			FieldName:       "user_name",
			Required:        true,
			MinLength:       5,
			MaxLength:       10,
			AllowWhitespace: false,
		},
		ab.Rules{
			FieldName:       "password",
			Required:        true,
			MinLength:       8,
			MaxLength:       20,
			AllowWhitespace: false,
		},
	}
	ab.Cfg.ConfirmFields = []string{"password", "confirm_password"}

	ab.Cfg.CookieStoreMaker = k.NewSessionStorer
	ab.Cfg.SessionStoreMaker = k.NewSessionStorer

	if err := ab.Init(); err != nil {
		log.Fatal(err)
	}

	if err := k.Templ.LoadToDB(k.DefaultTemplate); err != nil {
		log.Println(err)
	}
}

func (k Kesho) Run() {
	var (
		httpPort = "8080"
	)
	k.Initialize()
	k.Setup()
	log.Printf("Kesho is running at localhost:%s \n", httpPort)
	addr := fmt.Sprintf(":%s", httpPort)
	log.Fatal(http.ListenAndServe(addr, k.Routes()))

}
func (k *Kesho) Cleanup() error {
	os.Remove(k.MainDB)
	return os.Remove(k.SessionDB)
}
<|endoftext|>


<commit_before>package logutils

import (
	"bytes"
	"io"
	"sync"
)

type LogLevel string

// LevelFilter is an io.Writer that can be used with a logger that
// will filter out log messages that aren't at least a certain level.
//
// Once the filter is in use somewhere, it is not safe to modify
// the structure.
type LevelFilter struct {
	// Levels is the list of log levels, in increasing order of
	// severity. Example might be: {"DEBUG", "WARN", "ERROR"}.
	Levels []LogLevel

	// MinLevel is the minimum level allowed through
	MinLevel LogLevel

	// The underlying io.Writer where log messages that pass the filter
	// will be set.
	Writer io.Writer

	badLevels map[LogLevel]struct{}
	once      sync.Once
}

// Check will check a given line if it would be included in the level
// filter.
func (f *LevelFilter) Check(line []byte) bool {
	f.once.Do(f.init)

	// Check for a log level
	var level LogLevel
	x := bytes.IndexByte(line, '[')
	if x >= 0 {
		y := bytes.IndexByte(line[x:], ']')
		if y >= 0 {
			level = LogLevel(line[x+1 : y])
		}
	}

	_, ok := f.badLevels[level]
	return !ok
}

func (f *LevelFilter) Write(p []byte) (n int, err error) {
	// Note in general that io.Writer can receive any byte sequence
	// to write, but the "log" package always guarantees that we only
	// get a single line. We use that as a slight optimization within
	// this method, assuming we're dealing with a single, complete line
	// of log data.

	if !f.Check(p) {
		return len(p), nil
	}

	return f.Writer.Write(p)
}

func (f *LevelFilter) init() {
	f.badLevels = make(map[LogLevel]struct{})
	for _, level := range f.Levels {
		if level == f.MinLevel {
			break
		}

		f.badLevels[level] = struct{}{}
	}
}
<commit_msg>Fix slice out of bounds error<commit_after>package logutils

import (
	"bytes"
	"io"
	"sync"
)

type LogLevel string

// LevelFilter is an io.Writer that can be used with a logger that
// will filter out log messages that aren't at least a certain level.
//
// Once the filter is in use somewhere, it is not safe to modify
// the structure.
type LevelFilter struct {
	// Levels is the list of log levels, in increasing order of
	// severity. Example might be: {"DEBUG", "WARN", "ERROR"}.
	Levels []LogLevel

	// MinLevel is the minimum level allowed through
	MinLevel LogLevel

	// The underlying io.Writer where log messages that pass the filter
	// will be set.
	Writer io.Writer

	badLevels map[LogLevel]struct{}
	once      sync.Once
}

// Check will check a given line if it would be included in the level
// filter.
func (f *LevelFilter) Check(line []byte) bool {
	f.once.Do(f.init)

	// Check for a log level
	var level LogLevel
	x := bytes.IndexByte(line, '[')
	if x >= 0 {
		y := bytes.IndexByte(line[x:], ']')
		if y >= 0 {
			level = LogLevel(line[x+1 : x+y])
		}
	}

	_, ok := f.badLevels[level]
	return !ok
}

func (f *LevelFilter) Write(p []byte) (n int, err error) {
	// Note in general that io.Writer can receive any byte sequence
	// to write, but the "log" package always guarantees that we only
	// get a single line. We use that as a slight optimization within
	// this method, assuming we're dealing with a single, complete line
	// of log data.

	if !f.Check(p) {
		return len(p), nil
	}

	return f.Writer.Write(p)
}

func (f *LevelFilter) init() {
	f.badLevels = make(map[LogLevel]struct{})
	for _, level := range f.Levels {
		if level == f.MinLevel {
			break
		}

		f.badLevels[level] = struct{}{}
	}
}
<|endoftext|>


<commit_before>package log

type Level int8

const (
	FATAL Level = iota
	ERROR
	INFO
	WARN
	DEBUG
	TRACE
)

var StringToLevels = map[string]Level{
	"TRACE": TRACE,
	"DEBUG": DEBUG,
	"WARN":  WARN,
	"INFO":  INFO,
	"ERROR": ERROR,
	"FATAL": FATAL,
}

var LevelsToString = map[Level]string{
	TRACE: "TRACE",
	DEBUG: "DEBUG",
	WARN:  "WARN",
	INFO:  "INFO",
	ERROR: "ERROR",
	FATAL: "FATAL",
}
<commit_msg>swap WARN INFO level<commit_after>package log

type Level int8

const (
	FATAL Level = iota
	ERROR
	WARN
	INFO
	DEBUG
	TRACE
)

var StringToLevels = map[string]Level{
	"TRACE": TRACE,
	"DEBUG": DEBUG,
	"INFO":  INFO,
	"WARN":  WARN,
	"ERROR": ERROR,
	"FATAL": FATAL,
}

var LevelsToString = map[Level]string{
	TRACE: "TRACE",
	DEBUG: "DEBUG",
	INFO:  "INFO",
	WARN:  "WARN",
	ERROR: "ERROR",
	FATAL: "FATAL",
}
<|endoftext|>


<commit_before>package center

import (
	"strings"
	"time"

	"github.com/cozy/cozy-stack/pkg/apps"
	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/couchdb/mango"
	"github.com/cozy/cozy-stack/pkg/instance"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/notification"
	"github.com/cozy/cozy-stack/pkg/oauth"
	"github.com/cozy/cozy-stack/pkg/permissions"
	"github.com/cozy/cozy-stack/pkg/workers/mails"
	"github.com/cozy/cozy-stack/pkg/workers/push"
	multierror "github.com/hashicorp/go-multierror"
)

// Push creates and send a new notification in database. This method verifies
// the permissions associated with this creation in order to check that it is
// granted to create a notification and to extract its source.
func Push(inst *instance.Instance, perm *permissions.Permission, n *notification.Notification) error {
	if n.Title == "" {
		return ErrBadNotification
	}

	var p *notification.Properties
	switch perm.Type {
	// Applications and services have TypeWebapp permissions
	case permissions.TypeWebapp:
		slug := strings.TrimPrefix(perm.SourceID, consts.Apps+"/")
		man, err := apps.GetWebappBySlug(inst, slug)
		if err != nil {
			return err
		}
		p = man.Notifications[n.Category]
	default:
		return ErrUnauthorized
	}

	// XXX: for retro-compatibility, we do not yet block applications from
	// sending notification from unknown category.
	if p != nil && p.Stateful {
		l, err := findLastNotification(inst, n.Source())
		if err != nil {
			return err
		}
		// when the state is the same for the last notification from this source,
		// we do not bother sending or creating a new notification.
		if l != nil && l.State == n.State {
			return nil
		}
	}

	preferredChannels := n.PreferredChannels
	if len(preferredChannels) == 0 {
		preferredChannels = []string{"mail"}
	}

	n.NID = ""
	n.NRev = ""
	n.SourceID = n.Source()
	n.CreatedAt = time.Now()
	n.PreferredChannels = nil

	if err := couchdb.CreateDoc(inst, n); err != nil {
		return err
	}

	var errm error
	for _, channel := range preferredChannels {
		switch channel {
		case "mobile":
			if p != nil {
				if err := sendPush(inst, p.Collapsible, n); err != nil {
					errm = multierror.Append(errm, err)
				}
			}
		case "mail":
			if err := sendMail(inst, n); err != nil {
				errm = multierror.Append(errm, err)
			}
		}
	}
	return errm
}

func findLastNotification(inst *instance.Instance, source string) (*notification.Notification, error) {
	var notifs []*notification.Notification
	req := &couchdb.FindRequest{
		UseIndex: "by-source-id",
		Selector: mango.Equal("source", source),
		Sort: mango.SortBy{
			{Field: "source_id", Direction: mango.Desc},
			{Field: "created_at", Direction: mango.Desc},
		},
		Limit: 1,
	}
	err := couchdb.FindDocs(inst, consts.Notifications, req, &notifs)
	if err != nil {
		return nil, err
	}
	if len(notifs) == 0 {
		return nil, nil
	}
	return notifs[0], nil
}

func sendPush(inst *instance.Instance, collapsible bool, n *notification.Notification) error {
	clients, err := oauth.GetNotifiables(inst)
	if err != nil {
		return err
	}
	source := n.Source()
	var errm error
	for _, c := range clients {
		push := push.Message{
			NotificationID: n.ID(),
			Source:         source,
			Title:          n.Title,
			Message:        n.Message,
			Priority:       n.Priority,
			Sound:          n.Sound,
			Collapsible:    collapsible,
			Platform:       c.NotificationPlatform,
			DeviceToken:    c.NotificationDeviceToken,
		}
		msg, err := jobs.NewMessage(&push)
		if err != nil {
			return err
		}
		_, err = jobs.System().PushJob(&jobs.JobRequest{
			Domain:     inst.Domain,
			WorkerType: "push",
			Message:    msg,
		})
		if err != nil {
			errm = multierror.Append(errm, err)
		}
	}
	return errm
}

func sendMail(inst *instance.Instance, n *notification.Notification) error {
	parts := []*mails.Part{
		{Body: n.Content, Type: "text/plain"},
	}
	if n.ContentHTML != "" {
		parts = append(parts, &mails.Part{Body: n.ContentHTML, Type: "text/html"})
	}
	mail := mails.Options{
		Mode:    mails.ModeNoReply,
		Subject: n.Title,
		Parts:   parts,
	}
	msg, err := jobs.NewMessage(&mail)
	if err != nil {
		return err
	}
	_, err = jobs.System().PushJob(&jobs.JobRequest{
		Domain:     inst.Domain,
		WorkerType: "sendmail",
		Message:    msg,
	})
	return err
}
<commit_msg>Fix: passing data from notification center to push worker<commit_after>package center

import (
	"strings"
	"time"

	"github.com/cozy/cozy-stack/pkg/apps"
	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/couchdb/mango"
	"github.com/cozy/cozy-stack/pkg/instance"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/notification"
	"github.com/cozy/cozy-stack/pkg/oauth"
	"github.com/cozy/cozy-stack/pkg/permissions"
	"github.com/cozy/cozy-stack/pkg/workers/mails"
	"github.com/cozy/cozy-stack/pkg/workers/push"
	multierror "github.com/hashicorp/go-multierror"
)

// Push creates and send a new notification in database. This method verifies
// the permissions associated with this creation in order to check that it is
// granted to create a notification and to extract its source.
func Push(inst *instance.Instance, perm *permissions.Permission, n *notification.Notification) error {
	if n.Title == "" {
		return ErrBadNotification
	}

	var p *notification.Properties
	switch perm.Type {
	// Applications and services have TypeWebapp permissions
	case permissions.TypeWebapp:
		slug := strings.TrimPrefix(perm.SourceID, consts.Apps+"/")
		man, err := apps.GetWebappBySlug(inst, slug)
		if err != nil {
			return err
		}
		p = man.Notifications[n.Category]
	default:
		return ErrUnauthorized
	}

	// XXX: for retro-compatibility, we do not yet block applications from
	// sending notification from unknown category.
	if p != nil && p.Stateful {
		l, err := findLastNotification(inst, n.Source())
		if err != nil {
			return err
		}
		// when the state is the same for the last notification from this source,
		// we do not bother sending or creating a new notification.
		if l != nil && l.State == n.State {
			return nil
		}
	}

	preferredChannels := n.PreferredChannels
	if len(preferredChannels) == 0 {
		preferredChannels = []string{"mail"}
	}

	n.NID = ""
	n.NRev = ""
	n.SourceID = n.Source()
	n.CreatedAt = time.Now()
	n.PreferredChannels = nil

	if err := couchdb.CreateDoc(inst, n); err != nil {
		return err
	}

	var errm error
	for _, channel := range preferredChannels {
		switch channel {
		case "mobile":
			if p != nil {
				if err := sendPush(inst, p.Collapsible, n); err != nil {
					errm = multierror.Append(errm, err)
				}
			}
		case "mail":
			if err := sendMail(inst, n); err != nil {
				errm = multierror.Append(errm, err)
			}
		}
	}
	return errm
}

func findLastNotification(inst *instance.Instance, source string) (*notification.Notification, error) {
	var notifs []*notification.Notification
	req := &couchdb.FindRequest{
		UseIndex: "by-source-id",
		Selector: mango.Equal("source", source),
		Sort: mango.SortBy{
			{Field: "source_id", Direction: mango.Desc},
			{Field: "created_at", Direction: mango.Desc},
		},
		Limit: 1,
	}
	err := couchdb.FindDocs(inst, consts.Notifications, req, &notifs)
	if err != nil {
		return nil, err
	}
	if len(notifs) == 0 {
		return nil, nil
	}
	return notifs[0], nil
}

func sendPush(inst *instance.Instance, collapsible bool, n *notification.Notification) error {
	clients, err := oauth.GetNotifiables(inst)
	if err != nil {
		return err
	}
	source := n.Source()
	var errm error
	for _, c := range clients {
		push := push.Message{
			NotificationID: n.ID(),
			Source:         source,
			Title:          n.Title,
			Message:        n.Message,
			Priority:       n.Priority,
			Sound:          n.Sound,
			Data:           n.Data,
			Collapsible:    collapsible,
			Platform:       c.NotificationPlatform,
			DeviceToken:    c.NotificationDeviceToken,
		}
		msg, err := jobs.NewMessage(&push)
		if err != nil {
			return err
		}
		_, err = jobs.System().PushJob(&jobs.JobRequest{
			Domain:     inst.Domain,
			WorkerType: "push",
			Message:    msg,
		})
		if err != nil {
			errm = multierror.Append(errm, err)
		}
	}
	return errm
}

func sendMail(inst *instance.Instance, n *notification.Notification) error {
	parts := []*mails.Part{
		{Body: n.Content, Type: "text/plain"},
	}
	if n.ContentHTML != "" {
		parts = append(parts, &mails.Part{Body: n.ContentHTML, Type: "text/html"})
	}
	mail := mails.Options{
		Mode:    mails.ModeNoReply,
		Subject: n.Title,
		Parts:   parts,
	}
	msg, err := jobs.NewMessage(&mail)
	if err != nil {
		return err
	}
	_, err = jobs.System().PushJob(&jobs.JobRequest{
		Domain:     inst.Domain,
		WorkerType: "sendmail",
		Message:    msg,
	})
	return err
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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 priorities

import (
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	appslisters "k8s.io/client-go/listers/apps/v1"
	corelisters "k8s.io/client-go/listers/core/v1"
	"k8s.io/klog"
	schedulerlisters "k8s.io/kubernetes/pkg/scheduler/listers"
	schedulernodeinfo "k8s.io/kubernetes/pkg/scheduler/nodeinfo"
)

// MetadataFactory is a factory to produce PriorityMetadata.
type MetadataFactory struct {
	serviceLister         corelisters.ServiceLister
	controllerLister      corelisters.ReplicationControllerLister
	replicaSetLister      appslisters.ReplicaSetLister
	statefulSetLister     appslisters.StatefulSetLister
	hardPodAffinityWeight int32
}

// NewMetadataFactory creates a MetadataFactory.
func NewMetadataFactory(
	serviceLister corelisters.ServiceLister,
	controllerLister corelisters.ReplicationControllerLister,
	replicaSetLister appslisters.ReplicaSetLister,
	statefulSetLister appslisters.StatefulSetLister,
	hardPodAffinityWeight int32,
) MetadataProducer {
	factory := &MetadataFactory{
		serviceLister:         serviceLister,
		controllerLister:      controllerLister,
		replicaSetLister:      replicaSetLister,
		statefulSetLister:     statefulSetLister,
		hardPodAffinityWeight: hardPodAffinityWeight,
	}
	return factory.PriorityMetadata
}

// priorityMetadata is a type that is passed as metadata for priority functions
type priorityMetadata struct {
	podLimits               *schedulernodeinfo.Resource
	podTolerations          []v1.Toleration
	affinity                *v1.Affinity
	podSelector             labels.Selector
	controllerRef           *metav1.OwnerReference
	podFirstServiceSelector labels.Selector
	totalNumNodes           int
	podTopologySpreadMap    *podTopologySpreadMap
}

// PriorityMetadata is a MetadataProducer.  Node info can be nil.
func (pmf *MetadataFactory) PriorityMetadata(
	pod *v1.Pod,
	filteredNodes []*v1.Node,
	sharedLister schedulerlisters.SharedLister,
) interface{} {
	// If we cannot compute metadata, just return nil
	if pod == nil {
		return nil
	}
	totalNumNodes := 0
	var allNodes []*schedulernodeinfo.NodeInfo
	if sharedLister != nil {
		if l, err := sharedLister.NodeInfos().List(); err == nil {
			totalNumNodes = len(l)
			allNodes = l
		}
	}
	tpSpreadMap, err := buildPodTopologySpreadMap(pod, filteredNodes, allNodes)
	if err != nil {
		klog.Errorf("Error building podTopologySpreadMap: %v", err)
		return nil
	}
	return &priorityMetadata{
		podLimits:               getResourceLimits(pod),
		podTolerations:          getAllTolerationPreferNoSchedule(pod.Spec.Tolerations),
		affinity:                pod.Spec.Affinity,
		podSelector:             getSelector(pod, pmf.serviceLister, pmf.controllerLister, pmf.replicaSetLister, pmf.statefulSetLister),
		controllerRef:           metav1.GetControllerOf(pod),
		podFirstServiceSelector: getFirstServiceSelector(pod, pmf.serviceLister),
		totalNumNodes:           totalNumNodes,
		podTopologySpreadMap:    tpSpreadMap,
	}
}

// getFirstServiceSelector returns one selector of services the given pod.
func getFirstServiceSelector(pod *v1.Pod, sl corelisters.ServiceLister) (firstServiceSelector labels.Selector) {
	if services, err := schedulerlisters.GetPodServices(sl, pod); err == nil && len(services) > 0 {
		return labels.SelectorFromSet(services[0].Spec.Selector)
	}
	return nil
}

// getSelector returns a selector for the services, RCs, RSs, and SSs matching the given pod.
func getSelector(pod *v1.Pod, sl corelisters.ServiceLister, cl corelisters.ReplicationControllerLister, rsl appslisters.ReplicaSetLister, ssl appslisters.StatefulSetLister) labels.Selector {
	labelSet := make(labels.Set)
	// Since services, RCs, RSs and SSs match the pod, they won't have conflicting
	// labels. Merging is safe.

	if services, err := schedulerlisters.GetPodServices(sl, pod); err == nil {
		for _, service := range services {
			labelSet = labels.Merge(labelSet, service.Spec.Selector)
		}
	}

	if rcs, err := cl.GetPodControllers(pod); err == nil {
		for _, rc := range rcs {
			labelSet = labels.Merge(labelSet, rc.Spec.Selector)
		}
	}

	selector := labels.NewSelector()
	if len(labelSet) != 0 {
		selector = labelSet.AsSelector()
	}

	if rss, err := rsl.GetPodReplicaSets(pod); err == nil {
		for _, rs := range rss {
			if other, err := metav1.LabelSelectorAsSelector(rs.Spec.Selector); err == nil {
				if r, ok := other.Requirements(); ok {
					selector = selector.Add(r...)
				}
			}
		}
	}

	if sss, err := ssl.GetPodStatefulSets(pod); err == nil {
		for _, ss := range sss {
			if other, err := metav1.LabelSelectorAsSelector(ss.Spec.Selector); err == nil {
				if r, ok := other.Requirements(); ok {
					selector = selector.Add(r...)
				}
			}
		}
	}

	return selector
}
<commit_msg>fix: remove totalNumNodes from priority metadata<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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 priorities

import (
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	appslisters "k8s.io/client-go/listers/apps/v1"
	corelisters "k8s.io/client-go/listers/core/v1"
	"k8s.io/klog"
	schedulerlisters "k8s.io/kubernetes/pkg/scheduler/listers"
	schedulernodeinfo "k8s.io/kubernetes/pkg/scheduler/nodeinfo"
)

// MetadataFactory is a factory to produce PriorityMetadata.
type MetadataFactory struct {
	serviceLister         corelisters.ServiceLister
	controllerLister      corelisters.ReplicationControllerLister
	replicaSetLister      appslisters.ReplicaSetLister
	statefulSetLister     appslisters.StatefulSetLister
	hardPodAffinityWeight int32
}

// NewMetadataFactory creates a MetadataFactory.
func NewMetadataFactory(
	serviceLister corelisters.ServiceLister,
	controllerLister corelisters.ReplicationControllerLister,
	replicaSetLister appslisters.ReplicaSetLister,
	statefulSetLister appslisters.StatefulSetLister,
	hardPodAffinityWeight int32,
) MetadataProducer {
	factory := &MetadataFactory{
		serviceLister:         serviceLister,
		controllerLister:      controllerLister,
		replicaSetLister:      replicaSetLister,
		statefulSetLister:     statefulSetLister,
		hardPodAffinityWeight: hardPodAffinityWeight,
	}
	return factory.PriorityMetadata
}

// priorityMetadata is a type that is passed as metadata for priority functions
type priorityMetadata struct {
	podLimits               *schedulernodeinfo.Resource
	podTolerations          []v1.Toleration
	affinity                *v1.Affinity
	podSelector             labels.Selector
	controllerRef           *metav1.OwnerReference
	podFirstServiceSelector labels.Selector
	podTopologySpreadMap    *podTopologySpreadMap
}

// PriorityMetadata is a MetadataProducer.  Node info can be nil.
func (pmf *MetadataFactory) PriorityMetadata(
	pod *v1.Pod,
	filteredNodes []*v1.Node,
	sharedLister schedulerlisters.SharedLister,
) interface{} {
	// If we cannot compute metadata, just return nil
	if pod == nil {
		return nil
	}
	var allNodes []*schedulernodeinfo.NodeInfo
	if sharedLister != nil {
		if l, err := sharedLister.NodeInfos().List(); err == nil {
			allNodes = l
		}
	}
	tpSpreadMap, err := buildPodTopologySpreadMap(pod, filteredNodes, allNodes)
	if err != nil {
		klog.Errorf("Error building podTopologySpreadMap: %v", err)
		return nil
	}
	return &priorityMetadata{
		podLimits:               getResourceLimits(pod),
		podTolerations:          getAllTolerationPreferNoSchedule(pod.Spec.Tolerations),
		affinity:                pod.Spec.Affinity,
		podSelector:             getSelector(pod, pmf.serviceLister, pmf.controllerLister, pmf.replicaSetLister, pmf.statefulSetLister),
		controllerRef:           metav1.GetControllerOf(pod),
		podFirstServiceSelector: getFirstServiceSelector(pod, pmf.serviceLister),
		podTopologySpreadMap:    tpSpreadMap,
	}
}

// getFirstServiceSelector returns one selector of services the given pod.
func getFirstServiceSelector(pod *v1.Pod, sl corelisters.ServiceLister) (firstServiceSelector labels.Selector) {
	if services, err := schedulerlisters.GetPodServices(sl, pod); err == nil && len(services) > 0 {
		return labels.SelectorFromSet(services[0].Spec.Selector)
	}
	return nil
}

// getSelector returns a selector for the services, RCs, RSs, and SSs matching the given pod.
func getSelector(pod *v1.Pod, sl corelisters.ServiceLister, cl corelisters.ReplicationControllerLister, rsl appslisters.ReplicaSetLister, ssl appslisters.StatefulSetLister) labels.Selector {
	labelSet := make(labels.Set)
	// Since services, RCs, RSs and SSs match the pod, they won't have conflicting
	// labels. Merging is safe.

	if services, err := schedulerlisters.GetPodServices(sl, pod); err == nil {
		for _, service := range services {
			labelSet = labels.Merge(labelSet, service.Spec.Selector)
		}
	}

	if rcs, err := cl.GetPodControllers(pod); err == nil {
		for _, rc := range rcs {
			labelSet = labels.Merge(labelSet, rc.Spec.Selector)
		}
	}

	selector := labels.NewSelector()
	if len(labelSet) != 0 {
		selector = labelSet.AsSelector()
	}

	if rss, err := rsl.GetPodReplicaSets(pod); err == nil {
		for _, rs := range rss {
			if other, err := metav1.LabelSelectorAsSelector(rs.Spec.Selector); err == nil {
				if r, ok := other.Requirements(); ok {
					selector = selector.Add(r...)
				}
			}
		}
	}

	if sss, err := ssl.GetPodStatefulSets(pod); err == nil {
		for _, ss := range sss {
			if other, err := metav1.LabelSelectorAsSelector(ss.Spec.Selector); err == nil {
				if r, ok := other.Requirements(); ok {
					selector = selector.Add(r...)
				}
			}
		}
	}

	return selector
}
<|endoftext|>