<commit_before>// Copyright 2014 SteelSeries ApS.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This package implements a basic LISP interpretor for embedding in a go program for scripting.
// This file contains the bytearray primitive functions.

package golisp

import (
	"fmt"
	"unsafe"
)

func RegisterBytearrayPrimitives() {
	MakePrimitiveFunction("list-to-bytearray", 1, ListToBytesImpl)
	MakePrimitiveFunction("bytearray-to-list", 1, BytesToListImpl)
	MakePrimitiveFunction("replace-byte", 3, ReplaceByteImpl)
	MakePrimitiveFunction("replace-byte!", 3, ReplaceByteBangImpl)
	MakePrimitiveFunction("extract-byte", 2, ExtractByteImpl)
	MakePrimitiveFunction("append-bytes", -1, AppendBytesImpl)
	MakePrimitiveFunction("append-bytes!", -1, AppendBytesBangImpl)
	MakePrimitiveFunction("extract-bytes", 3, ExtractBytesImpl)
}

func ListToBytesImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	if NilP(Car(args)) {
		err = ProcessError("Argument to ListToByutes can not be nil.", env)
		return
	}
	list, err := Eval(Car(args), env)
	if err != nil {
		return
	}

	bytes := make([]byte, 0, int(Length(list)))
	for c := list; NotNilP(c); c = Cdr(c) {
		var n *Data
		n, err = Eval(Car(c), env)
		if !IntegerP(n) && !(ObjectP(n) && ObjectType(n) == "[]byte") {
			err = ProcessError(fmt.Sprintf("Byte arrays can only contain numbers, but found %v.", n), env)
			return
		}

		if IntegerP(n) {
			b := IntegerValue(n)
			if b > 255 {
				err = ProcessError(fmt.Sprintf("Byte arrays can only contain bytes, but found %d.", b), env)
				return
			}
			bytes = append(bytes, byte(b))
		} else {
			otherArrayBytes := *(*[]byte)(ObjectValue(n))
			bytes = append(bytes, otherArrayBytes...)
		}
	}
	return ObjectWithTypeAndValue("[]byte", unsafe.Pointer(&bytes)), nil
}

func BytesToListImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	dataByteObject, err := Eval(Car(args), env)
	if err != nil {
		panic(err)
	}
	if !ObjectP(dataByteObject) || ObjectType(dataByteObject) != "[]byte" {
		err = ProcessError(fmt.Sprintf("Bytearray object should return []byte but returned %s.", ObjectType(dataByteObject)), env)
		return
	}

	dataBytes := (*[]byte)(ObjectValue(dataByteObject))
	var bytes = make([]*Data, 0, len(*dataBytes))

	for _, b := range *dataBytes {
		bytes = append(bytes, IntegerWithValue(int64(b)))
	}

	result = ArrayToList(bytes)
	return
}

func internalReplaceByte(args *Data, env *SymbolTableFrame, makeCopy bool) (result *Data, err error) {

	if First(args) == nil {
		err = ProcessError("replace-byte requires a non-nil bytearray argument.", env)
		return
	}
	if Second(args) == nil {
		err = ProcessError("replace-byte requires a non-nil index argument.", env)
		return
	}
	if Third(args) == nil {
		err = ProcessError("replace-byte requires a non-nil value argument.", env)
		return
	}

	dataByteObject, err := Eval(First(args), env)
	if err != nil {
		panic(err)
	}
	if !ObjectP(dataByteObject) || ObjectType(dataByteObject) != "[]byte" {
		err = ProcessError(fmt.Sprintf("Bytearray object should return []byte but returned %s.", ObjectType(dataByteObject)), env)
		return
	}

	dataBytes := (*[]byte)(ObjectValue(dataByteObject))
	var newBytes *[]byte
	if makeCopy {
		temp := make([]byte, len(*dataBytes))
		newBytes = &temp
		copy(*newBytes, *dataBytes)
	} else {
		newBytes = dataBytes
	}

	indexObject, err := Eval(Second(args), env)
	if err != nil {
		panic(err)
	}
	if !IntegerP(indexObject) {
		panic(ProcessError("Bytearray index should be a number.", env))
	}
	index := int(IntegerValue(indexObject))

	if index >= len(*dataBytes) {
		err = ProcessError(fmt.Sprintf("replace-byte index was out of range. Was %d but bytearray has length of %d.", index, len(*dataBytes)), env)
		return
	}

	if WalkList(args, "add") == nil {
		err = ProcessError("replace-byte requires a non-nil value argument.", env)
		return
	}

	valueObject, err := Eval(Third(args), env)
	if err != nil {
		panic(err)
	}
	if !IntegerP(valueObject) {
		panic(ProcessError("Bytearray value should be a number.", env))
	}

	value := byte(IntegerValue(valueObject))

	if value > 255 {
		err = ProcessError(fmt.Sprintf("replace-byte value was not a byte. Was %d.", index), env)
		return
	}

	(*newBytes)[index] = value

	if makeCopy {
		result = ObjectWithTypeAndValue("[]byte", unsafe.Pointer(newBytes))
	} else {
		result = dataByteObject
	}
	return
}

func ReplaceByteImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return internalReplaceByte(args, env, true)
}

func ReplaceByteBangImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return internalReplaceByte(args, env, false)
}

func ExtractByteImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	if Car(args) == nil {
		err = ProcessError("extract-byte requires a non-nil bytearray argument.", env)
		return
	}
	if Cadr(args) == nil {
		err = ProcessError("extract-byte requires a non-nil index argument.", env)
		return
	}

	dataByteObject, err := Eval(Car(args), env)
	if err != nil {
		panic(err)
	}
	if !ObjectP(dataByteObject) || ObjectType(dataByteObject) != "[]byte" {
		panic(ProcessError(fmt.Sprintf("Bytearray object should return []byte but returned %s.", ObjectType(dataByteObject)), env))
	}

	dataBytes := (*[]byte)(ObjectValue(dataByteObject))

	indexObject, err := Eval(Cadr(args), env)
	if err != nil {
		panic(err)
	}
	if !IntegerP(indexObject) {
		panic(ProcessError("Bytearray index should be a number.", env))
	}
	index := int(IntegerValue(indexObject))

	if index >= len(*dataBytes) {
		err = ProcessError(fmt.Sprintf("extract-byte index was out of range. Was %d but bytearray has length of %d.", index, len(*dataBytes)), env)
		return
	}

	extractedValue := (*dataBytes)[index]
	result = IntegerWithValue(int64(extractedValue))
	return
}

func internalAppendBytes(args *Data, env *SymbolTableFrame) (newBytes *[]byte, err error) {
	if Car(args) == nil {
		err = ProcessError("append-bytes requires a non-nil bytearray argument.", env)
		return
	}
	if Cadr(args) == nil {
		err = ProcessError("append-bytes requires a non-nil list of bytes to append.", env)
		return
	}

	dataByteObject, err := Eval(Car(args), env)
	if err != nil {
		panic(err)
	}
	if !ObjectP(dataByteObject) || ObjectType(dataByteObject) != "[]byte" {
		panic(ProcessError(fmt.Sprintf("Bytearray object should return []byte but returned %s.", ObjectType(dataByteObject)), env))
	}

	dataBytes := (*[]byte)(ObjectValue(dataByteObject))

	var extraByteObj *Data
	var evaledArg *Data
	if NilP(Cddr(args)) {
		evaledArg, err = Eval(Cadr(args), env)
		if err != nil {
			return
		}
		if ObjectP(evaledArg) && ObjectType(evaledArg) == "[]byte" {
			extraByteObj = evaledArg
		} else if ListP(evaledArg) {
			extraByteObj, err = ListToBytesImpl(InternalMakeList(QuoteIt(evaledArg)), env)
		} else {
			extraByteObj, err = ListToBytesImpl(InternalMakeList(QuoteIt(Cdr(args))), env)
		}
	} else {
		extraByteObj, err = ListToBytesImpl(InternalMakeList(QuoteIt(Cdr(args))), env)
	}

	if err != nil {
		return
	}

	extraBytes := (*[]byte)(ObjectValue(extraByteObj))

	temp := make([]byte, len(*dataBytes)+len(*extraBytes))
	newBytes = &temp
	copy(*newBytes, *dataBytes)
	copy((*newBytes)[len(*dataBytes):], *extraBytes)

	return
}

func AppendBytesImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	newBytesPtr, err := internalAppendBytes(args, env)
	if err != nil {
		return
	}
	result = ObjectWithTypeAndValue("[]byte", unsafe.Pointer(newBytesPtr))
	return
}

func AppendBytesBangImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	newBytesPtr, err := internalAppendBytes(args, env)
	dataByteObject, _ := Eval(Car(args), env)
	objectData := BoxedObjectValue(dataByteObject)
	objectData.Obj = unsafe.Pointer(newBytesPtr)
	result = dataByteObject
	return
}

func ExtractBytesImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	if Car(args) == nil {
		err = ProcessError("extract-bytes requires a non-nil bytearray argument.", env)
		return
	}
	if Cadr(args) == nil {
		err = ProcessError("extract-bytes requires a non-nil index argument.", env)
		return
	}
	if Caddr(args) == nil {
		err = ProcessError("extract-bytes requires a non-nil number to extract argument.", env)
		return
	}

	dataByteObject, err := Eval(Car(args), env)
	if err != nil {
		panic(err)
	}
	if !ObjectP(dataByteObject) || ObjectType(dataByteObject) != "[]byte" {
		panic(ProcessError(fmt.Sprintf("Bytearray object should return []byte but returned %s.", ObjectType(dataByteObject)), env))
	}

	dataBytes := (*[]byte)(ObjectValue(dataByteObject))

	indexObject, err := Eval(Cadr(args), env)
	if err != nil {
		panic(err)
	}
	if !IntegerP(indexObject) {
		panic(ProcessError("Bytearray index should be a number.", env))
	}
	index := int(IntegerValue(indexObject))

	numToExtractObject, err := Eval(Caddr(args), env)
	if err != nil {
		panic(err)
	}
	if !IntegerP(numToExtractObject) {
		panic(ProcessError("Number to extract should be a number.", env))
	}
	numToExtract := int(IntegerValue(numToExtractObject))

	if index >= len(*dataBytes) {
		err = ProcessError(fmt.Sprintf("extract-bytes index was out of range. Was %d but bytearray has length of %d.", index, len(*dataBytes)), env)
		return
	}
	if index+numToExtract > len(*dataBytes) {
		err = ProcessError(fmt.Sprintf("extract-bytes final index was out of range.  Was %d but bytearray has length of %d.", index+numToExtract-1, len(*dataBytes)), env)
		return
	}

	result, err = DropImpl(InternalMakeList(indexObject, dataByteObject), Global)
	if err != nil {
		return
	}
	result, err = TakeImpl(InternalMakeList(numToExtractObject, result), Global)
	return
}
<commit_msg>Fix overzealous error panicing!<commit_after>// Copyright 2014 SteelSeries ApS.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This package implements a basic LISP interpretor for embedding in a go program for scripting.
// This file contains the bytearray primitive functions.

package golisp

import (
	"fmt"
	"unsafe"
)

func RegisterBytearrayPrimitives() {
	MakePrimitiveFunction("list-to-bytearray", 1, ListToBytesImpl)
	MakePrimitiveFunction("bytearray-to-list", 1, BytesToListImpl)
	MakePrimitiveFunction("replace-byte", 3, ReplaceByteImpl)
	MakePrimitiveFunction("replace-byte!", 3, ReplaceByteBangImpl)
	MakePrimitiveFunction("extract-byte", 2, ExtractByteImpl)
	MakePrimitiveFunction("append-bytes", -1, AppendBytesImpl)
	MakePrimitiveFunction("append-bytes!", -1, AppendBytesBangImpl)
	MakePrimitiveFunction("extract-bytes", 3, ExtractBytesImpl)
}

func ListToBytesImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	if NilP(Car(args)) {
		err = ProcessError("Argument to ListToByutes can not be nil.", env)
		return
	}
	list, err := Eval(Car(args), env)
	if err != nil {
		return
	}

	bytes := make([]byte, 0, int(Length(list)))
	for c := list; NotNilP(c); c = Cdr(c) {
		var n *Data
		n, err = Eval(Car(c), env)
		if !IntegerP(n) && !(ObjectP(n) && ObjectType(n) == "[]byte") {
			err = ProcessError(fmt.Sprintf("Byte arrays can only contain numbers, but found %v.", n), env)
			return
		}

		if IntegerP(n) {
			b := IntegerValue(n)
			if b > 255 {
				err = ProcessError(fmt.Sprintf("Byte arrays can only contain bytes, but found %d.", b), env)
				return
			}
			bytes = append(bytes, byte(b))
		} else {
			otherArrayBytes := *(*[]byte)(ObjectValue(n))
			bytes = append(bytes, otherArrayBytes...)
		}
	}
	return ObjectWithTypeAndValue("[]byte", unsafe.Pointer(&bytes)), nil
}

func BytesToListImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	dataByteObject, err := Eval(Car(args), env)
	if err != nil {
		panic(err)
	}
	if !ObjectP(dataByteObject) || ObjectType(dataByteObject) != "[]byte" {
		err = ProcessError(fmt.Sprintf("Bytearray object should return []byte but returned %s.", ObjectType(dataByteObject)), env)
		return
	}

	dataBytes := (*[]byte)(ObjectValue(dataByteObject))
	var bytes = make([]*Data, 0, len(*dataBytes))

	for _, b := range *dataBytes {
		bytes = append(bytes, IntegerWithValue(int64(b)))
	}

	result = ArrayToList(bytes)
	return
}

func internalReplaceByte(args *Data, env *SymbolTableFrame, makeCopy bool) (result *Data, err error) {

	if First(args) == nil {
		err = ProcessError("replace-byte requires a non-nil bytearray argument.", env)
		return
	}
	if Second(args) == nil {
		err = ProcessError("replace-byte requires a non-nil index argument.", env)
		return
	}
	if Third(args) == nil {
		err = ProcessError("replace-byte requires a non-nil value argument.", env)
		return
	}

	dataByteObject, err := Eval(First(args), env)
	if err != nil {
		panic(err)
	}
	if !ObjectP(dataByteObject) || ObjectType(dataByteObject) != "[]byte" {
		err = ProcessError(fmt.Sprintf("Bytearray object should return []byte but returned %s.", ObjectType(dataByteObject)), env)
		return
	}

	dataBytes := (*[]byte)(ObjectValue(dataByteObject))
	var newBytes *[]byte
	if makeCopy {
		temp := make([]byte, len(*dataBytes))
		newBytes = &temp
		copy(*newBytes, *dataBytes)
	} else {
		newBytes = dataBytes
	}

	indexObject, err := Eval(Second(args), env)
	if err != nil {
		panic(err)
	}
	if !IntegerP(indexObject) {
		err = ProcessError("Bytearray index should be a number.", env)
		return
	}
	index := int(IntegerValue(indexObject))

	if index >= len(*dataBytes) {
		err = ProcessError(fmt.Sprintf("replace-byte index was out of range. Was %d but bytearray has length of %d.", index, len(*dataBytes)), env)
		return
	}

	if WalkList(args, "add") == nil {
		err = ProcessError("replace-byte requires a non-nil value argument.", env)
		return
	}

	valueObject, err := Eval(Third(args), env)
	if err != nil {
		panic(err)
	}
	if !IntegerP(valueObject) {
		err = ProcessError("Bytearray value should be a number.", env)
		return
	}

	value := byte(IntegerValue(valueObject))

	if value > 255 {
		err = ProcessError(fmt.Sprintf("replace-byte value was not a byte. Was %d.", index), env)
		return
	}

	(*newBytes)[index] = value

	if makeCopy {
		result = ObjectWithTypeAndValue("[]byte", unsafe.Pointer(newBytes))
	} else {
		result = dataByteObject
	}
	return
}

func ReplaceByteImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return internalReplaceByte(args, env, true)
}

func ReplaceByteBangImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return internalReplaceByte(args, env, false)
}

func ExtractByteImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	if Car(args) == nil {
		err = ProcessError("extract-byte requires a non-nil bytearray argument.", env)
		return
	}
	if Cadr(args) == nil {
		err = ProcessError("extract-byte requires a non-nil index argument.", env)
		return
	}

	dataByteObject, err := Eval(Car(args), env)
	if err != nil {
		panic(err)
	}
	if !ObjectP(dataByteObject) || ObjectType(dataByteObject) != "[]byte" {
		err = ProcessError(fmt.Sprintf("Bytearray object should return []byte but returned %s.", ObjectType(dataByteObject)), env)
		return
	}

	dataBytes := (*[]byte)(ObjectValue(dataByteObject))

	indexObject, err := Eval(Cadr(args), env)
	if err != nil {
		panic(err)
	}
	if !IntegerP(indexObject) {
		err = ProcessError("Bytearray index should be a number.", env)
		return
	}
	index := int(IntegerValue(indexObject))

	if index >= len(*dataBytes) {
		err = ProcessError(fmt.Sprintf("extract-byte index was out of range. Was %d but bytearray has length of %d.", index, len(*dataBytes)), env)
		return
	}

	extractedValue := (*dataBytes)[index]
	result = IntegerWithValue(int64(extractedValue))
	return
}

func internalAppendBytes(args *Data, env *SymbolTableFrame) (newBytes *[]byte, err error) {
	if Car(args) == nil {
		err = ProcessError("append-bytes requires a non-nil bytearray argument.", env)
		return
	}
	if Cadr(args) == nil {
		err = ProcessError("append-bytes requires a non-nil list of bytes to append.", env)
		return
	}

	dataByteObject, err := Eval(Car(args), env)
	if err != nil {
		panic(err)
	}
	if !ObjectP(dataByteObject) || ObjectType(dataByteObject) != "[]byte" {
		err = ProcessError(fmt.Sprintf("Bytearray object should return []byte but returned %s.", ObjectType(dataByteObject)), env)
		return
	}

	dataBytes := (*[]byte)(ObjectValue(dataByteObject))

	var extraByteObj *Data
	var evaledArg *Data
	if NilP(Cddr(args)) {
		evaledArg, err = Eval(Cadr(args), env)
		if err != nil {
			return
		}
		if ObjectP(evaledArg) && ObjectType(evaledArg) == "[]byte" {
			extraByteObj = evaledArg
		} else if ListP(evaledArg) {
			extraByteObj, err = ListToBytesImpl(InternalMakeList(QuoteIt(evaledArg)), env)
		} else {
			extraByteObj, err = ListToBytesImpl(InternalMakeList(QuoteIt(Cdr(args))), env)
		}
	} else {
		extraByteObj, err = ListToBytesImpl(InternalMakeList(QuoteIt(Cdr(args))), env)
	}

	if err != nil {
		return
	}

	extraBytes := (*[]byte)(ObjectValue(extraByteObj))

	temp := make([]byte, len(*dataBytes)+len(*extraBytes))
	newBytes = &temp
	copy(*newBytes, *dataBytes)
	copy((*newBytes)[len(*dataBytes):], *extraBytes)

	return
}

func AppendBytesImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	newBytesPtr, err := internalAppendBytes(args, env)
	if err != nil {
		return
	}
	result = ObjectWithTypeAndValue("[]byte", unsafe.Pointer(newBytesPtr))
	return
}

func AppendBytesBangImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	newBytesPtr, err := internalAppendBytes(args, env)
	dataByteObject, _ := Eval(Car(args), env)
	objectData := BoxedObjectValue(dataByteObject)
	objectData.Obj = unsafe.Pointer(newBytesPtr)
	result = dataByteObject
	return
}

func ExtractBytesImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	if Car(args) == nil {
		err = ProcessError("extract-bytes requires a non-nil bytearray argument.", env)
		return
	}
	if Cadr(args) == nil {
		err = ProcessError("extract-bytes requires a non-nil index argument.", env)
		return
	}
	if Caddr(args) == nil {
		err = ProcessError("extract-bytes requires a non-nil number to extract argument.", env)
		return
	}

	dataByteObject, err := Eval(Car(args), env)
	if err != nil {
		panic(err)
	}
	if !ObjectP(dataByteObject) || ObjectType(dataByteObject) != "[]byte" {
		err = ProcessError(fmt.Sprintf("Bytearray object should return []byte but returned %s.", ObjectType(dataByteObject)), env)
		return
	}

	dataBytes := (*[]byte)(ObjectValue(dataByteObject))

	indexObject, err := Eval(Cadr(args), env)
	if err != nil {
		panic(err)
	}
	if !IntegerP(indexObject) {
		err = ProcessError("Bytearray index should be a number.", env)
		return
	}
	index := int(IntegerValue(indexObject))

	numToExtractObject, err := Eval(Caddr(args), env)
	if err != nil {
		panic(err)
	}
	if !IntegerP(numToExtractObject) {
		err = ProcessError("Number to extract should be a number.", env)
		return
	}
	numToExtract := int(IntegerValue(numToExtractObject))

	if index >= len(*dataBytes) {
		err = ProcessError(fmt.Sprintf("extract-bytes index was out of range. Was %d but bytearray has length of %d.", index, len(*dataBytes)), env)
		return
	}
	if index+numToExtract > len(*dataBytes) {
		err = ProcessError(fmt.Sprintf("extract-bytes final index was out of range.  Was %d but bytearray has length of %d.", index+numToExtract-1, len(*dataBytes)), env)
		return
	}

	result, err = DropImpl(InternalMakeList(indexObject, dataByteObject), Global)
	if err != nil {
		return
	}
	result, err = TakeImpl(InternalMakeList(numToExtractObject, result), Global)
	return
}
<|endoftext|>


<commit_before>package logrus_cloudwatchlogs

import (
	"bytes"
	"encoding/json"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/kardianos/osext"
)

type ProdFormatter struct {
	hostname                string
	appname                 string
	httpRequestKey          string
	httpRequestHeaderFilter []string
}

type ProdFormatterOption func(*ProdFormatter) error

// Hostname is a formatter option that specifies the hostname this
// program is running on. If this is not specified, the system's hostname
// will be used.
func Hostname(name string) ProdFormatterOption {
	return func(f *ProdFormatter) error {
		f.hostname = name
		return nil
	}
}

// AppName is a formatter option that specifies the name of the app.
// If this is not specified, the default is to use the executable name.
func AppName(name string) ProdFormatterOption {
	return func(f *ProdFormatter) error {
		f.appname = name
		return nil
	}
}

// HTTPRequest is a formatter option that allows you to indicate
// that a certain key will contain an *http.Request. If it does, it
// will be formatted in the output. You can provide an optional list
// of keys to filter out of the serialized header. This is useful so
// you don't include sensitive information (like the authorization header).
// Note: if you do not provide this and you pass in an *http.Request, it will
// fail because encoding/json cannot serialize *http.Request.
func HTTPRequest(key string, headerFilter ...string) ProdFormatterOption {
	return func(f *ProdFormatter) error {
		f.httpRequestKey = key
		f.httpRequestHeaderFilter = headerFilter
		return nil
	}
}

// NewProdFormatter creates a new cloudwatchlogs production formatter.
// This is opinionated and you can feel free to create your own.
func NewProdFormatter(options ...ProdFormatterOption) *ProdFormatter {
	f := &ProdFormatter{}

	for _, option := range options {
		option(f)
	}

	var err error
	if f.hostname == "" {
		if f.hostname, err = os.Hostname(); err != nil {
			f.hostname = "unknown"
		}
	}

	if f.appname == "" {
		if f.appname, err = osext.Executable(); err == nil {
			f.appname = filepath.Base(f.appname)
		} else {
			f.appname = "app"
		}
	}

	return f
}

// Format formats logrus.Entry in the form of:
// [timestamp] [jsondata]
func (f *ProdFormatter) Format(entry *logrus.Entry) ([]byte, error) {
	//tmp := make([]byte, 50)
	b := &bytes.Buffer{}

	// //b.WriteString(time.Now().Format("2006-01-02T15:04:05.999Z"))
	now := time.Now()
	// year, month, day := now.Date()
	// hour, minute, second := now.Clock()
	// nano := now.Nanosecond()
	// fourDigits(&tmp, 0, year)
	// tmp[4] = '-'
	// twoDigits(&tmp, 5, int(month))
	// tmp[7] = '-'
	// twoDigits(&tmp, 8, day)
	// tmp[10] = 'T'
	// twoDigits(&tmp, 11, hour)
	// tmp[13] = ':'
	// twoDigits(&tmp, 14, minute)
	// tmp[16] = ':'
	// twoDigits(&tmp, 17, second)
	// tmp[19] = '.'
	// threeDigits(&tmp, 20, nano)
	// tmp[23] = 'Z'
	// b.Write(tmp[:24])
	//
	// b.WriteRune(' ')

	// // This is so incredibly hacky. Needed until logrus implements
	// // this.
	// skip := 8
	// if len(entry.Data) == 0 {
	// 	skip = 9
	// }
	// file, line, function := fileInfo(skip)

	data := map[string]interface{}{
		"time":  now.Unix(), //b.String()[:24],
		"msg":   entry.Message,
		"level": entry.Level.String(),
		"host":  f.hostname,
		// "file":  file,
		// "line":  line,
		// "func":  function,
		"app": f.appname,
	}

	for k, v := range entry.Data {
		switch v := v.(type) {
		case error:
			// Otherwise errors are ignored by `encoding/json`
			// https://github.com/Sirupsen/logrus/issues/137
			data[k] = v.Error()
		default:
			data[k] = v
		}
	}

	if v, ok := data[f.httpRequestKey]; ok {
		if req, ok := v.(*http.Request); ok {
			header := make(map[string]interface{})
			for key, value := range req.Header {
				header[key] = value
			}
			for _, key := range f.httpRequestHeaderFilter {
				delete(header, key)
			}

			data[f.httpRequestKey] = map[string]interface{}{
				"method": req.Method,
				// We have to use RequestURI because URL may be
				// modified by routes.
				"url":         req.RequestURI, //t.URL.String(),
				"host":        req.Host,
				"remote_addr": req.RemoteAddr,
				"header":      header,
			}
		}
	}

	j, err := json.Marshal(data)
	if err != nil {
		return nil, err
	}

	b.Write(j)

	return b.Bytes(), nil
}

func fileInfo(skip int) (string, int, string) {
	function := "<???>"
	pc, file, line, ok := runtime.Caller(skip)
	if !ok {
		file = "<???>"
		line = 1
	} else {
		slash := strings.LastIndex(file, "/")
		if slash >= 0 {
			file = file[slash+1:]
		}

		me := runtime.FuncForPC(pc)
		if me != nil {
			function = me.Name()
		}
	}

	return file, line, function
}

const ddigits = `0001020304050607080910111213141516171819` +
	`2021222324252627282930313233343536373839` +
	`4041424344454647484950515253545556575859` +
	`6061626364656667686970717273747576777879` +
	`8081828384858687888990919293949596979899`

// itoa converts an integer d to its ascii representation
// i is the deintation index in buf
// algorithm from https://www.facebook.com/notes/facebook-engineering/three-optimization-tips-for-c/10151361643253920
func itoa(buf *[]byte, i, d int) int {
	j := len(*buf)

	for d >= 100 {
		// Integer division is slow, so we do it by 2
		index := (d % 100) * 2
		d /= 100
		j--
		(*buf)[j] = ddigits[index+1]
		j--
		(*buf)[j] = ddigits[index]
	}

	if d < 10 {
		j--
		(*buf)[j] = byte(int('0') + d)
		return copy((*buf)[i:], (*buf)[j:])
	}

	index := d * 2
	j--
	(*buf)[j] = ddigits[index+1]
	j--
	(*buf)[j] = ddigits[index]

	return copy((*buf)[i:], (*buf)[j:])
}

const digits = "0123456789"

// twoDigits converts an integer d to its ascii representation
// i is the destination index in buf
func twoDigits(buf *[]byte, i, d int) {
	(*buf)[i+1] = digits[d%10]
	d /= 10
	(*buf)[i] = digits[d%10]
}

func threeDigits(buf *[]byte, i, d int) {
	(*buf)[i+2] = digits[d%10]
	d /= 10
	(*buf)[i+1] = digits[d%10]
	d /= 10
	(*buf)[i] = digits[d%10]
}

func fourDigits(buf *[]byte, i, d int) {
	(*buf)[i+3] = digits[d%10]
	d /= 10
	(*buf)[i+2] = digits[d%10]
	d /= 10
	(*buf)[i+1] = digits[d%10]
	d /= 10
	(*buf)[i] = digits[d%10]
}
<commit_msg>Added Marshaler interface<commit_after>package logrus_cloudwatchlogs

import (
	"bytes"
	"encoding/json"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/kardianos/osext"
)

type ProdFormatter struct {
	hostname                string
	appname                 string
	httpRequestKey          string
	httpRequestHeaderFilter []string
}

type ProdFormatterOption func(*ProdFormatter) error

//Marshaler is an interface any type can implement to change its output in our production logs.
type Marshaler interface {
	MarshalLog() map[string]interface{}
}

// Hostname is a formatter option that specifies the hostname this
// program is running on. If this is not specified, the system's hostname
// will be used.
func Hostname(name string) ProdFormatterOption {
	return func(f *ProdFormatter) error {
		f.hostname = name
		return nil
	}
}

// AppName is a formatter option that specifies the name of the app.
// If this is not specified, the default is to use the executable name.
func AppName(name string) ProdFormatterOption {
	return func(f *ProdFormatter) error {
		f.appname = name
		return nil
	}
}

// HTTPRequest is a formatter option that allows you to indicate
// that a certain key will contain an *http.Request. If it does, it
// will be formatted in the output. You can provide an optional list
// of keys to filter out of the serialized header. This is useful so
// you don't include sensitive information (like the authorization header).
// Note: if you do not provide this and you pass in an *http.Request, it will
// fail because encoding/json cannot serialize *http.Request.
func HTTPRequest(key string, headerFilter ...string) ProdFormatterOption {
	return func(f *ProdFormatter) error {
		f.httpRequestKey = key
		f.httpRequestHeaderFilter = headerFilter
		return nil
	}
}

// NewProdFormatter creates a new cloudwatchlogs production formatter.
// This is opinionated and you can feel free to create your own.
func NewProdFormatter(options ...ProdFormatterOption) *ProdFormatter {
	f := &ProdFormatter{}

	for _, option := range options {
		option(f)
	}

	var err error
	if f.hostname == "" {
		if f.hostname, err = os.Hostname(); err != nil {
			f.hostname = "unknown"
		}
	}

	if f.appname == "" {
		if f.appname, err = osext.Executable(); err == nil {
			f.appname = filepath.Base(f.appname)
		} else {
			f.appname = "app"
		}
	}

	return f
}

// Format formats logrus.Entry in the form of:
// [timestamp] [jsondata]
func (f *ProdFormatter) Format(entry *logrus.Entry) ([]byte, error) {
	//tmp := make([]byte, 50)
	b := &bytes.Buffer{}

	// //b.WriteString(time.Now().Format("2006-01-02T15:04:05.999Z"))
	now := time.Now()
	// year, month, day := now.Date()
	// hour, minute, second := now.Clock()
	// nano := now.Nanosecond()
	// fourDigits(&tmp, 0, year)
	// tmp[4] = '-'
	// twoDigits(&tmp, 5, int(month))
	// tmp[7] = '-'
	// twoDigits(&tmp, 8, day)
	// tmp[10] = 'T'
	// twoDigits(&tmp, 11, hour)
	// tmp[13] = ':'
	// twoDigits(&tmp, 14, minute)
	// tmp[16] = ':'
	// twoDigits(&tmp, 17, second)
	// tmp[19] = '.'
	// threeDigits(&tmp, 20, nano)
	// tmp[23] = 'Z'
	// b.Write(tmp[:24])
	//
	// b.WriteRune(' ')

	// // This is so incredibly hacky. Needed until logrus implements
	// // this.
	// skip := 8
	// if len(entry.Data) == 0 {
	// 	skip = 9
	// }
	// file, line, function := fileInfo(skip)

	data := map[string]interface{}{
		"time":  now.Unix(), //b.String()[:24],
		"msg":   entry.Message,
		"level": entry.Level.String(),
		"host":  f.hostname,
		// "file":  file,
		// "line":  line,
		// "func":  function,
		"app": f.appname,
	}

	for k, v := range entry.Data {
		switch v := v.(type) {
		case error:
			// Otherwise errors are ignored by `encoding/json`
			// https://github.com/Sirupsen/logrus/issues/137
			data[k] = v.Error()
		case marshaler:
			data[k] = v.MarshalLog()
		default:
			data[k] = v
		}
	}

	if v, ok := data[f.httpRequestKey]; ok {
		if req, ok := v.(*http.Request); ok {
			header := make(map[string]interface{})
			for key, value := range req.Header {
				header[key] = value
			}
			for _, key := range f.httpRequestHeaderFilter {
				delete(header, key)
			}

			data[f.httpRequestKey] = map[string]interface{}{
				"method": req.Method,
				// We have to use RequestURI because URL may be
				// modified by routes.
				"url":         req.RequestURI, //t.URL.String(),
				"host":        req.Host,
				"remote_addr": req.RemoteAddr,
				"header":      header,
			}
		}
	}

	j, err := json.Marshal(data)
	if err != nil {
		return nil, err
	}

	b.Write(j)

	return b.Bytes(), nil
}

func fileInfo(skip int) (string, int, string) {
	function := "<???>"
	pc, file, line, ok := runtime.Caller(skip)
	if !ok {
		file = "<???>"
		line = 1
	} else {
		slash := strings.LastIndex(file, "/")
		if slash >= 0 {
			file = file[slash+1:]
		}

		me := runtime.FuncForPC(pc)
		if me != nil {
			function = me.Name()
		}
	}

	return file, line, function
}

const ddigits = `0001020304050607080910111213141516171819` +
	`2021222324252627282930313233343536373839` +
	`4041424344454647484950515253545556575859` +
	`6061626364656667686970717273747576777879` +
	`8081828384858687888990919293949596979899`

// itoa converts an integer d to its ascii representation
// i is the deintation index in buf
// algorithm from https://www.facebook.com/notes/facebook-engineering/three-optimization-tips-for-c/10151361643253920
func itoa(buf *[]byte, i, d int) int {
	j := len(*buf)

	for d >= 100 {
		// Integer division is slow, so we do it by 2
		index := (d % 100) * 2
		d /= 100
		j--
		(*buf)[j] = ddigits[index+1]
		j--
		(*buf)[j] = ddigits[index]
	}

	if d < 10 {
		j--
		(*buf)[j] = byte(int('0') + d)
		return copy((*buf)[i:], (*buf)[j:])
	}

	index := d * 2
	j--
	(*buf)[j] = ddigits[index+1]
	j--
	(*buf)[j] = ddigits[index]

	return copy((*buf)[i:], (*buf)[j:])
}

const digits = "0123456789"

// twoDigits converts an integer d to its ascii representation
// i is the destination index in buf
func twoDigits(buf *[]byte, i, d int) {
	(*buf)[i+1] = digits[d%10]
	d /= 10
	(*buf)[i] = digits[d%10]
}

func threeDigits(buf *[]byte, i, d int) {
	(*buf)[i+2] = digits[d%10]
	d /= 10
	(*buf)[i+1] = digits[d%10]
	d /= 10
	(*buf)[i] = digits[d%10]
}

func fourDigits(buf *[]byte, i, d int) {
	(*buf)[i+3] = digits[d%10]
	d /= 10
	(*buf)[i+2] = digits[d%10]
	d /= 10
	(*buf)[i+1] = digits[d%10]
	d /= 10
	(*buf)[i] = digits[d%10]
}
<|endoftext|>


<commit_before>// Copyright 2015 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package prog

import (
	"bytes"
	"math/rand"
	"testing"
	"time"
)

func initTest(t *testing.T) (rand.Source, int) {
	iters := 10000
	if testing.Short() {
		iters = 100
	}
	seed := int64(time.Now().UnixNano())
	rs := rand.NewSource(seed)
	t.Logf("seed=%v", seed)
	return rs, iters
}

func TestGeneration(t *testing.T) {
	rs, iters := initTest(t)
	for i := 0; i < iters; i++ {
		Generate(rs, 20, nil)
	}
}

func TestSerialize(t *testing.T) {
	rs, iters := initTest(t)
	for i := 0; i < iters; i++ {
		p := Generate(rs, 10, nil)
		data := p.Serialize()
		p1, err := Deserialize(data)
		if err != nil {
			t.Fatalf("failed to deserialize program: %v\n%s", err, data)
		}
		data1 := p1.Serialize()
		if len(p.Calls) != len(p1.Calls) {
			t.Fatalf("different number of calls")
		}
		if !bytes.Equal(data, data1) {
			t.Fatalf("program changed after serialize/deserialize\noriginal:\n%s\n\nnew:\n%s\n", data, data1)
		}
	}
}
<commit_msg>prog: test that Deserialize does not return nil prog<commit_after>// Copyright 2015 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package prog

import (
	"bytes"
	"math/rand"
	"testing"
	"time"
)

func initTest(t *testing.T) (rand.Source, int) {
	iters := 10000
	if testing.Short() {
		iters = 100
	}
	seed := int64(time.Now().UnixNano())
	rs := rand.NewSource(seed)
	t.Logf("seed=%v", seed)
	return rs, iters
}

func TestGeneration(t *testing.T) {
	rs, iters := initTest(t)
	for i := 0; i < iters; i++ {
		Generate(rs, 20, nil)
	}
}

func TestSerialize(t *testing.T) {
	rs, iters := initTest(t)
	for i := 0; i < iters; i++ {
		p := Generate(rs, 10, nil)
		data := p.Serialize()
		p1, err := Deserialize(data)
		if err != nil {
			t.Fatalf("failed to deserialize program: %v\n%s", err, data)
		}
		if p1 == nil {
			t.Fatalf("deserialized nil program:\n%s", data)
		}
		data1 := p1.Serialize()
		if len(p.Calls) != len(p1.Calls) {
			t.Fatalf("different number of calls")
		}
		if !bytes.Equal(data, data1) {
			t.Fatalf("program changed after serialize/deserialize\noriginal:\n%s\n\nnew:\n%s\n", data, data1)
		}
	}
}
<|endoftext|>


<commit_before>package kite

import (
	"flag"
	"fmt"
	logging "github.com/op/go-logging"
	"koding/newkite/dnode"
	"koding/newkite/dnode/rpc"
	"koding/newkite/protocol"
	"koding/newkite/utils"
	stdlog "log"
	"net"
	"net/http"
	"os"
	"time"
)

var log = logging.MustGetLogger("Kite")

// GetLogger returns a new logger which is used within the application itsel
// (in main package).
func GetLogger() *logging.Logger {
	return log
}

// Kite defines a single process that enables distributed service messaging
// amongst the peers it is connected. A Kite process acts as a Client and as a
// Server. That means it can receive request, process them, but it also can
// make request to other kites. A Kite can be anything. It can be simple Image
// processing kite (which would process data), it could be a Chat kite that
// enables peer-to-peer chat. For examples we have FileSystem kite that expose
// the file system to a client, which in order build the filetree.
type Kite struct {
	protocol.Kite

	// KodingKey is used for authenticate to Kontrol.
	KodingKey string

	// Points to the Kontrol instance if enabled
	Kontrol *Kontrol

	// Wheter we want to connect to Kontrol on startup, true by default.
	KontrolEnabled bool

	// Wheter we want to register our Kite to Kontrol, true by default.
	RegisterToKontrol bool

	// method map for exported methods
	handlers map[string]HandlerFunc

	// Dnode rpc server
	server *rpc.Server

	// Handlers to call when a Kite opens a connection to this Kite.
	onConnectHandlers []func(*RemoteKite)

	// Contains different functions for authenticating user from request.
	// Keys are the authentication types (options.authentication.type).
	Authenticators map[string]func(*CallOptions) error
}

// New creates, initialize and then returns a new Kite instance. It accept
// three  arguments. options is a config struct that needs to be filled with
// several informations like Name, Port, IP and so on.
func New(options *protocol.Options) *Kite {
	var err error
	if options == nil {
		options, err = utils.ReadKiteOptions("manifest.json")
		if err != nil {
			log.Fatal("error: could not read config file", err)
		}
	}

	// some simple validations for config
	if options.Kitename == "" {
		log.Fatal("ERROR: options.Kitename field is not set")
	}

	if options.Region == "" {
		log.Fatal("ERROR: options.Region field is not set")
	}

	if options.Environment == "" {
		log.Fatal("ERROR: options.Environment field is not set")
	}

	hostname, _ := os.Hostname()
	kiteID := utils.GenerateUUID()

	kodingKey, err := utils.GetKodingKey()
	if err != nil {
		log.Fatal("Couldn't find koding.key. Please run 'kd register'.")
	}

	port := options.Port
	if options.Port == "" {
		port = "0" // OS binds to an automatic port
	}

	if options.KontrolAddr == "" {
		options.KontrolAddr = "127.0.0.1:4000" // local fallback address
	}

	k := &Kite{
		Kite: protocol.Kite{
			Name:        options.Kitename,
			Username:    options.Username,
			ID:          kiteID,
			Version:     options.Version,
			Hostname:    hostname,
			Port:        port,
			Environment: options.Environment,
			Region:      options.Region,

			// PublicIP will be set by Kontrol after registering if it is not set.
			PublicIP: options.PublicIP,
		},
		KodingKey:         kodingKey,
		server:            rpc.NewServer(),
		KontrolEnabled:    true,
		RegisterToKontrol: true,
		Authenticators:    make(map[string]func(*CallOptions) error),
		handlers:          make(map[string]HandlerFunc),
	}
	k.Kontrol = k.NewKontrol(options.KontrolAddr)

	// Every kite should be able to authenticate the user from token.
	k.Authenticators["token"] = k.AuthenticateFromToken
	// A kite accepts requests from Kontrol.
	k.Authenticators["kodingKey"] = k.AuthenticateFromKodingKey

	// Register our internal methods
	k.HandleFunc("status", new(Status).Info)
	k.HandleFunc("heartbeat", k.handleHeartbeat)
	k.HandleFunc("log", k.handleLog)

	return k
}

func (k *Kite) HandleFunc(method string, handler HandlerFunc) {
	k.server.HandleFunc(method, func(msg *dnode.Message, tr dnode.Transport) {
		request, responseCallback, err := k.parseRequest(msg, tr)
		if err != nil {
			log.Notice("Did not understand request: %s", err)
			return
		}

		result, err := handler(request)
		if responseCallback == nil {
			return
		}

		if err != nil {
			err = responseCallback(err.Error(), result)
		} else {
			err = responseCallback(nil, result)
		}

		if err != nil {
			log.Error(err.Error())
		}
	})
}

// Run is a blocking method. It runs the kite server and then accepts requests
// asynchronously. It can be started in a goroutine if you wish to use kite as a
// client too.
func (k *Kite) Run() {
	k.parseVersionFlag()
	k.setupLogging()

	err := k.listenAndServe()
	if err != nil {
		log.Fatal(err)
	}
}

func (k *Kite) handleHeartbeat(r *Request) (interface{}, error) {
	args, err := r.Args.Array()
	if err != nil {
		return nil, err
	}

	if len(args) != 2 {
		return nil, fmt.Errorf("Invalid args: %s", string(r.Args.Raw))
	}

	seconds, ok := args[0].(float64)
	if !ok {
		return nil, fmt.Errorf("Invalid interval: %s", args[0])
	}

	ping, ok := args[1].(dnode.Function)
	if !ok {
		return nil, fmt.Errorf("Invalid callback: %s", args[1])
	}

	go func() {
		for {
			time.Sleep(time.Duration(seconds) * time.Second)
			if ping() != nil {
				return
			}
		}
	}()

	return nil, nil
}

// handleLog prints a log message to stdout.
func (k *Kite) handleLog(r *Request) (interface{}, error) {
	s, err := r.Args.String()
	if err != nil {
		return nil, err
	}

	log.Info(fmt.Sprintf("%s: %s", r.RemoteKite.Name, s))
	return nil, nil
}

// setupLogging is used to setup the logging format, destination and level.
func (k *Kite) setupLogging() {
	log.Module = k.Name
	logging.SetFormatter(logging.MustStringFormatter("%{level:-8s} ▶ %{message}"))

	stderrBackend := logging.NewLogBackend(os.Stderr, "", stdlog.LstdFlags)
	stderrBackend.Color = true

	syslogBackend, _ := logging.NewSyslogBackend(k.Name)
	logging.SetBackend(stderrBackend, syslogBackend)

	// Set logging level. Default level is INFO.
	level := logging.INFO
	if k.hasDebugFlag() {
		level = logging.DEBUG
	}
	logging.SetLevel(level, log.Module)
}

// If the user wants to call flag.Parse() the flag must be defined in advance.
var _ = flag.Bool("version", false, "show version")
var _ = flag.Bool("debug", false, "print debug logs")

// parseVersionFlag prints the version number of the kite and exits with 0
// if "-version" flag is enabled.
// We did not use the "flag" package because it causes trouble if the user
// also calls "flag.Parse()" in his code. flag.Parse() can be called only once.
func (k *Kite) parseVersionFlag() {
	for _, flag := range os.Args {
		if flag == "-version" {
			fmt.Println(k.Version)
			os.Exit(0)
		}
	}
}

// hasDebugFlag returns true if -debug flag is present in os.Args.
func (k *Kite) hasDebugFlag() bool {
	for _, flag := range os.Args {
		if flag == "-debug" {
			return true
		}
	}
	return false
}

// listenAndServe starts our rpc server with the given addr.
func (k *Kite) listenAndServe() error {
	listener, err := net.Listen("tcp4", ":"+k.Port)
	if err != nil {
		return err
	}

	log.Info("Listening: %s", listener.Addr().String())

	// Port is known here if "0" is used as port number
	_, k.Port, _ = net.SplitHostPort(listener.Addr().String())

	// We must connect to Kontrol after starting to listen on port
	if k.KontrolEnabled {
		if k.RegisterToKontrol {
			k.Kontrol.RemoteKite.OnConnect(k.registerToKontrol)
		}

		k.Kontrol.DialForever()
	}

	return http.Serve(listener, k.server)
}

func (k *Kite) registerToKontrol() {
	err := k.Kontrol.Register()
	if err != nil {
		log.Fatalf("Cannot register to Kontrol: %s", err)
	}

	log.Info("Registered to Kontrol successfully")
}

// OnConnect registers a function to run when a Kite connects to this Kite.
func (k *Kite) OnConnect(handler func(*RemoteKite)) {
	k.onConnectHandlers = append(k.onConnectHandlers, handler)
}

// OnDisconnect registers a function to run when a connected Kite is disconnected.
func (k *Kite) OnDisconnect(handler func(*RemoteKite)) {
	k.server.OnDisconnect(func(c *rpc.Client) {
		if r, ok := c.Properties()["remoteKite"]; ok {
			handler(r.(*RemoteKite))
		}
	})
}

// notifyRemoteKiteConnected runs the registered handlers with OnConnect().
func (k *Kite) notifyRemoteKiteConnected(r *RemoteKite) {
	// Print log messages on connect and disconnect.
	log.Info("Client has connected: %s", r.Addr())
	k.OnDisconnect(func(r *RemoteKite) {
		log.Info("Client has disconnected: %s", r.Addr())
	})

	// Run handlers.
	for _, handler := range k.onConnectHandlers {
		go handler(r)
	}
}
<commit_msg>kite: rename status method to systemInfo<commit_after>package kite

import (
	"flag"
	"fmt"
	logging "github.com/op/go-logging"
	"koding/newkite/dnode"
	"koding/newkite/dnode/rpc"
	"koding/newkite/protocol"
	"koding/newkite/utils"
	stdlog "log"
	"net"
	"net/http"
	"os"
	"time"
)

var log = logging.MustGetLogger("Kite")

// GetLogger returns a new logger which is used within the application itsel
// (in main package).
func GetLogger() *logging.Logger {
	return log
}

// Kite defines a single process that enables distributed service messaging
// amongst the peers it is connected. A Kite process acts as a Client and as a
// Server. That means it can receive request, process them, but it also can
// make request to other kites. A Kite can be anything. It can be simple Image
// processing kite (which would process data), it could be a Chat kite that
// enables peer-to-peer chat. For examples we have FileSystem kite that expose
// the file system to a client, which in order build the filetree.
type Kite struct {
	protocol.Kite

	// KodingKey is used for authenticate to Kontrol.
	KodingKey string

	// Points to the Kontrol instance if enabled
	Kontrol *Kontrol

	// Wheter we want to connect to Kontrol on startup, true by default.
	KontrolEnabled bool

	// Wheter we want to register our Kite to Kontrol, true by default.
	RegisterToKontrol bool

	// method map for exported methods
	handlers map[string]HandlerFunc

	// Dnode rpc server
	server *rpc.Server

	// Handlers to call when a Kite opens a connection to this Kite.
	onConnectHandlers []func(*RemoteKite)

	// Contains different functions for authenticating user from request.
	// Keys are the authentication types (options.authentication.type).
	Authenticators map[string]func(*CallOptions) error
}

// New creates, initialize and then returns a new Kite instance. It accept
// three  arguments. options is a config struct that needs to be filled with
// several informations like Name, Port, IP and so on.
func New(options *protocol.Options) *Kite {
	var err error
	if options == nil {
		options, err = utils.ReadKiteOptions("manifest.json")
		if err != nil {
			log.Fatal("error: could not read config file", err)
		}
	}

	// some simple validations for config
	if options.Kitename == "" {
		log.Fatal("ERROR: options.Kitename field is not set")
	}

	if options.Region == "" {
		log.Fatal("ERROR: options.Region field is not set")
	}

	if options.Environment == "" {
		log.Fatal("ERROR: options.Environment field is not set")
	}

	hostname, _ := os.Hostname()
	kiteID := utils.GenerateUUID()

	kodingKey, err := utils.GetKodingKey()
	if err != nil {
		log.Fatal("Couldn't find koding.key. Please run 'kd register'.")
	}

	port := options.Port
	if options.Port == "" {
		port = "0" // OS binds to an automatic port
	}

	if options.KontrolAddr == "" {
		options.KontrolAddr = "127.0.0.1:4000" // local fallback address
	}

	k := &Kite{
		Kite: protocol.Kite{
			Name:        options.Kitename,
			Username:    options.Username,
			ID:          kiteID,
			Version:     options.Version,
			Hostname:    hostname,
			Port:        port,
			Environment: options.Environment,
			Region:      options.Region,

			// PublicIP will be set by Kontrol after registering if it is not set.
			PublicIP: options.PublicIP,
		},
		KodingKey:         kodingKey,
		server:            rpc.NewServer(),
		KontrolEnabled:    true,
		RegisterToKontrol: true,
		Authenticators:    make(map[string]func(*CallOptions) error),
		handlers:          make(map[string]HandlerFunc),
	}
	k.Kontrol = k.NewKontrol(options.KontrolAddr)

	// Every kite should be able to authenticate the user from token.
	k.Authenticators["token"] = k.AuthenticateFromToken
	// A kite accepts requests from Kontrol.
	k.Authenticators["kodingKey"] = k.AuthenticateFromKodingKey

	// Register our internal methods
	k.HandleFunc("systemInfo", new(Status).Info)
	k.HandleFunc("heartbeat", k.handleHeartbeat)
	k.HandleFunc("log", k.handleLog)

	return k
}

func (k *Kite) HandleFunc(method string, handler HandlerFunc) {
	k.server.HandleFunc(method, func(msg *dnode.Message, tr dnode.Transport) {
		request, responseCallback, err := k.parseRequest(msg, tr)
		if err != nil {
			log.Notice("Did not understand request: %s", err)
			return
		}

		result, err := handler(request)
		if responseCallback == nil {
			return
		}

		if err != nil {
			err = responseCallback(err.Error(), result)
		} else {
			err = responseCallback(nil, result)
		}

		if err != nil {
			log.Error(err.Error())
		}
	})
}

// Run is a blocking method. It runs the kite server and then accepts requests
// asynchronously. It can be started in a goroutine if you wish to use kite as a
// client too.
func (k *Kite) Run() {
	k.parseVersionFlag()
	k.setupLogging()

	err := k.listenAndServe()
	if err != nil {
		log.Fatal(err)
	}
}

func (k *Kite) handleHeartbeat(r *Request) (interface{}, error) {
	args, err := r.Args.Array()
	if err != nil {
		return nil, err
	}

	if len(args) != 2 {
		return nil, fmt.Errorf("Invalid args: %s", string(r.Args.Raw))
	}

	seconds, ok := args[0].(float64)
	if !ok {
		return nil, fmt.Errorf("Invalid interval: %s", args[0])
	}

	ping, ok := args[1].(dnode.Function)
	if !ok {
		return nil, fmt.Errorf("Invalid callback: %s", args[1])
	}

	go func() {
		for {
			time.Sleep(time.Duration(seconds) * time.Second)
			if ping() != nil {
				return
			}
		}
	}()

	return nil, nil
}

// handleLog prints a log message to stdout.
func (k *Kite) handleLog(r *Request) (interface{}, error) {
	s, err := r.Args.String()
	if err != nil {
		return nil, err
	}

	log.Info(fmt.Sprintf("%s: %s", r.RemoteKite.Name, s))
	return nil, nil
}

// setupLogging is used to setup the logging format, destination and level.
func (k *Kite) setupLogging() {
	log.Module = k.Name
	logging.SetFormatter(logging.MustStringFormatter("%{level:-8s} ▶ %{message}"))

	stderrBackend := logging.NewLogBackend(os.Stderr, "", stdlog.LstdFlags)
	stderrBackend.Color = true

	syslogBackend, _ := logging.NewSyslogBackend(k.Name)
	logging.SetBackend(stderrBackend, syslogBackend)

	// Set logging level. Default level is INFO.
	level := logging.INFO
	if k.hasDebugFlag() {
		level = logging.DEBUG
	}
	logging.SetLevel(level, log.Module)
}

// If the user wants to call flag.Parse() the flag must be defined in advance.
var _ = flag.Bool("version", false, "show version")
var _ = flag.Bool("debug", false, "print debug logs")

// parseVersionFlag prints the version number of the kite and exits with 0
// if "-version" flag is enabled.
// We did not use the "flag" package because it causes trouble if the user
// also calls "flag.Parse()" in his code. flag.Parse() can be called only once.
func (k *Kite) parseVersionFlag() {
	for _, flag := range os.Args {
		if flag == "-version" {
			fmt.Println(k.Version)
			os.Exit(0)
		}
	}
}

// hasDebugFlag returns true if -debug flag is present in os.Args.
func (k *Kite) hasDebugFlag() bool {
	for _, flag := range os.Args {
		if flag == "-debug" {
			return true
		}
	}
	return false
}

// listenAndServe starts our rpc server with the given addr.
func (k *Kite) listenAndServe() error {
	listener, err := net.Listen("tcp4", ":"+k.Port)
	if err != nil {
		return err
	}

	log.Info("Listening: %s", listener.Addr().String())

	// Port is known here if "0" is used as port number
	_, k.Port, _ = net.SplitHostPort(listener.Addr().String())

	// We must connect to Kontrol after starting to listen on port
	if k.KontrolEnabled {
		if k.RegisterToKontrol {
			k.Kontrol.RemoteKite.OnConnect(k.registerToKontrol)
		}

		k.Kontrol.DialForever()
	}

	return http.Serve(listener, k.server)
}

func (k *Kite) registerToKontrol() {
	err := k.Kontrol.Register()
	if err != nil {
		log.Fatalf("Cannot register to Kontrol: %s", err)
	}

	log.Info("Registered to Kontrol successfully")
}

// OnConnect registers a function to run when a Kite connects to this Kite.
func (k *Kite) OnConnect(handler func(*RemoteKite)) {
	k.onConnectHandlers = append(k.onConnectHandlers, handler)
}

// OnDisconnect registers a function to run when a connected Kite is disconnected.
func (k *Kite) OnDisconnect(handler func(*RemoteKite)) {
	k.server.OnDisconnect(func(c *rpc.Client) {
		if r, ok := c.Properties()["remoteKite"]; ok {
			handler(r.(*RemoteKite))
		}
	})
}

// notifyRemoteKiteConnected runs the registered handlers with OnConnect().
func (k *Kite) notifyRemoteKiteConnected(r *RemoteKite) {
	// Print log messages on connect and disconnect.
	log.Info("Client has connected: %s", r.Addr())
	k.OnDisconnect(func(r *RemoteKite) {
		log.Info("Client has disconnected: %s", r.Addr())
	})

	// Run handlers.
	for _, handler := range k.onConnectHandlers {
		go handler(r)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"syscall"
)

var daemonName = "npserver"

func main() {

	initFlags()

	// check permissions
	if os.Getuid() != 0 && os.Geteuid() != 0 {
		fmt.Printf("npserver-daemon should be run as root, have uid=%d and euid=%d\n", os.Getuid(), os.Geteuid())
		os.Exit(1)
	}

	// check start || sto
	if flags.Start == flags.Stop {
		fmt.Println("need --start or --stop flag")
	}

	if flags.Start {
		startDaemon()
	}
	if flags.Stop {
		stopDaemon()
	}

	// all good :)
}

func startDaemon() {
	// setup args for daemon call
	args := []string{
		fmt.Sprintf("--name=%s", daemonName),
		"--noconfig",
		"--errlog=/var/log/npserver-daemon.log",
		"--output=/var/log/npserver.log",
		fmt.Sprintf("--pidfile=%s", flags.PIDFile),
		"--unsafe",
		"--",
		"/usr/local/bin/npserver",
	}

	// append extra args to args
	args = append(args, extraArgs...)

	// start process
	proc, err := os.StartProcess("/usr/bin/daemon", args, &os.ProcAttr{
		Files: []*os.File{os.Stdin, os.Stdout, os.Stderr},
		Sys: &syscall.SysProcAttr{
			Credential: &syscall.Credential{
				Uid: uint32(os.Geteuid()),
				Gid: uint32(os.Getegid()),
			},
		},
	})
	if err != nil {
		fmt.Printf("os/exec returned an error: '%s'\n", err)
		os.Exit(1)
	}

	// wait for daemon to be ready
	_, err = proc.Wait()
	if err != nil {
		fmt.Printf("proc.Wait() failed. %s\n", err)
		os.Exit(1)
	}
}

func stopDaemon() {
	// stopping arguments
	args := []string{
		fmt.Sprintf("--name=%s", daemonName),
		fmt.Sprintf("--pidfile=%s", flags.PIDFile),
		"--stop",
	}

	// start process
	proc, err := os.StartProcess("/usr/bin/daemon", args, &os.ProcAttr{
		Files: []*os.File{os.Stdin, os.Stdout, os.Stderr},
		Sys: &syscall.SysProcAttr{
			Credential: &syscall.Credential{
				Uid: uint32(os.Geteuid()),
				Gid: uint32(os.Getegid()),
			},
		},
	})
	if err != nil {
		fmt.Printf("os/exec returned an error: '%s'\n", err)
		os.Exit(1)
	}

	// wait for daemon to be ready
	_, err = proc.Wait()
	if err != nil {
		fmt.Printf("proc.Wait() failed. %s\n", err)
		os.Exit(1)
	}
	// // open pid file (if available)
	// pidFile, err := os.Open(flags.PIDFile)
	// if err != nil {
	// 	if os.IsNotExist(err) {
	// 		fmt.Printf("it looks like npserver is not running")
	// 		os.Exit(0)
	// 	}
	// 	fmt.Printf("error on opening pidfile: %s\n", err)
	// 	os.Exit(1)
	// }

	// // read all file contents
	// pidFileContents, err := ioutil.ReadAll(pidFile)
	// pidFile.Close()
	// if err != nil {
	// 	fmt.Printf("error reading pidfile contents: %s\n", err)
	// 	os.Exit(1)
	// }

	// pidFileContentsString := string(pidFileContents)

	// // strip eventual whitespace
	// pidFileContentsString = strings.TrimRight(pidFileContentsString, " \r\n\t")

	// // convert pid string to pid int
	// pid, err := strconv.Atoi(pidFileContentsString)
	// if err != nil {
	// 	fmt.Printf("error parsing pidfile contents: %s\n", err)
	// 	os.Exit(1)
	// }

	// // lookup process
	// proc, err := os.FindProcess(pid)
	// if err != nil {
	// 	fmt.Printf("error finding process with pid %d: %s\n", pid, err)
	// 	os.Exit(1)
	// }

	// // signal process to stop
	// err = proc.Signal(os.Interrupt)
	// if err != nil {
	// 	fmt.Printf("error sending interrupt signal to npserver: %s\n", err)
	// 	os.Exit(1)
	// }

	// // wait until process is done
	// state, err := proc.Wait()
	// if err != nil {
	// 	if err.Error() != "wait: no child processes" {
	// 		fmt.Printf("error waiting for process to stop: %s\n", err)
	// 		os.Exit(1)
	// 	}
	// } else {
	// 	if !state.Exited() || !state.Success() {
	// 		fmt.Printf("npserver process exited badly")
	// 		os.Exit(1)
	// 	}
	// }

	// // remove pid file
	// err = os.Remove(flags.PIDFile)
	// if err != nil {
	// 	fmt.Printf("error removing pid file: %s\n", err)
	// }
}
<commit_msg>first --stop then --name?<commit_after>package main

import (
	"fmt"
	"os"
	"syscall"
)

var daemonName = "npserver"

func main() {

	initFlags()

	// check permissions
	if os.Getuid() != 0 && os.Geteuid() != 0 {
		fmt.Printf("npserver-daemon should be run as root, have uid=%d and euid=%d\n", os.Getuid(), os.Geteuid())
		os.Exit(1)
	}

	// check start || sto
	if flags.Start == flags.Stop {
		fmt.Println("need --start or --stop flag")
	}

	if flags.Start {
		startDaemon()
	}
	if flags.Stop {
		stopDaemon()
	}

	// all good :)
}

func startDaemon() {
	// setup args for daemon call
	args := []string{
		fmt.Sprintf("--name=%s", daemonName),
		"--noconfig",
		"--errlog=/var/log/npserver-daemon.log",
		"--output=/var/log/npserver.log",
		fmt.Sprintf("--pidfile=%s", flags.PIDFile),
		"--unsafe",
		"--",
		"/usr/local/bin/npserver",
	}

	// append extra args to args
	args = append(args, extraArgs...)

	// start process
	proc, err := os.StartProcess("/usr/bin/daemon", args, &os.ProcAttr{
		Files: []*os.File{os.Stdin, os.Stdout, os.Stderr},
		Sys: &syscall.SysProcAttr{
			Credential: &syscall.Credential{
				Uid: uint32(os.Geteuid()),
				Gid: uint32(os.Getegid()),
			},
		},
	})
	if err != nil {
		fmt.Printf("os/exec returned an error: '%s'\n", err)
		os.Exit(1)
	}

	// wait for daemon to be ready
	_, err = proc.Wait()
	if err != nil {
		fmt.Printf("proc.Wait() failed. %s\n", err)
		os.Exit(1)
	}
}

func stopDaemon() {
	// stopping arguments
	args := []string{
		"--stop",
		fmt.Sprintf("--name=%s", daemonName),
		fmt.Sprintf("--pidfile=%s", flags.PIDFile),
	}

	// start process
	proc, err := os.StartProcess("/usr/bin/daemon", args, &os.ProcAttr{
		Files: []*os.File{os.Stdin, os.Stdout, os.Stderr},
		Sys: &syscall.SysProcAttr{
			Credential: &syscall.Credential{
				Uid: uint32(os.Geteuid()),
				Gid: uint32(os.Getegid()),
			},
		},
	})
	if err != nil {
		fmt.Printf("os/exec returned an error: '%s'\n", err)
		os.Exit(1)
	}

	// wait for daemon to be ready
	_, err = proc.Wait()
	if err != nil {
		fmt.Printf("proc.Wait() failed. %s\n", err)
		os.Exit(1)
	}
	// // open pid file (if available)
	// pidFile, err := os.Open(flags.PIDFile)
	// if err != nil {
	// 	if os.IsNotExist(err) {
	// 		fmt.Printf("it looks like npserver is not running")
	// 		os.Exit(0)
	// 	}
	// 	fmt.Printf("error on opening pidfile: %s\n", err)
	// 	os.Exit(1)
	// }

	// // read all file contents
	// pidFileContents, err := ioutil.ReadAll(pidFile)
	// pidFile.Close()
	// if err != nil {
	// 	fmt.Printf("error reading pidfile contents: %s\n", err)
	// 	os.Exit(1)
	// }

	// pidFileContentsString := string(pidFileContents)

	// // strip eventual whitespace
	// pidFileContentsString = strings.TrimRight(pidFileContentsString, " \r\n\t")

	// // convert pid string to pid int
	// pid, err := strconv.Atoi(pidFileContentsString)
	// if err != nil {
	// 	fmt.Printf("error parsing pidfile contents: %s\n", err)
	// 	os.Exit(1)
	// }

	// // lookup process
	// proc, err := os.FindProcess(pid)
	// if err != nil {
	// 	fmt.Printf("error finding process with pid %d: %s\n", pid, err)
	// 	os.Exit(1)
	// }

	// // signal process to stop
	// err = proc.Signal(os.Interrupt)
	// if err != nil {
	// 	fmt.Printf("error sending interrupt signal to npserver: %s\n", err)
	// 	os.Exit(1)
	// }

	// // wait until process is done
	// state, err := proc.Wait()
	// if err != nil {
	// 	if err.Error() != "wait: no child processes" {
	// 		fmt.Printf("error waiting for process to stop: %s\n", err)
	// 		os.Exit(1)
	// 	}
	// } else {
	// 	if !state.Exited() || !state.Success() {
	// 		fmt.Printf("npserver process exited badly")
	// 		os.Exit(1)
	// 	}
	// }

	// // remove pid file
	// err = os.Remove(flags.PIDFile)
	// if err != nil {
	// 	fmt.Printf("error removing pid file: %s\n", err)
	// }
}
<|endoftext|>


<commit_before>BuildDB.go<commit_msg>Delete BuildDB.go<commit_after><|endoftext|>


<commit_before>// Copyright 2017 Mathew Robinson <mrobinson@praelatus.io>. All rights reserved.
// Use of this source code is governed by the AGPLv3 license that can be found in
// the LICENSE file.

package mongo

import (
	"time"

	"github.com/praelatus/praelatus/models"
	"github.com/praelatus/praelatus/repo"
	mgo "gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

type projectRepo struct {
	conn *mgo.Session
}

func (p projectRepo) coll() *mgo.Collection {
	return p.conn.DB(dbName).C(projects)
}

func permWithID(u *models.User, uid string) bson.M {
	base := permQuery(u)
	return bson.M{
		"$and": []bson.M{
			{
				"_id": uid,
			},
			base,
		},
	}
}

func (p projectRepo) Get(u *models.User, uid string) (models.Project, error) {
	var project models.Project
	q := permWithID(u, uid)
	err := p.coll().Find(q).One(&project)
	return project, mongoErr(err)
}

func (p projectRepo) Update(u *models.User, uid string, updated models.Project) error {
	q := permWithID(u, uid)
	return mongoErr(p.coll().Update(q, updated))
}

func (p projectRepo) Create(u *models.User, project models.Project) (models.Project, error) {
	if u == nil || !u.IsAdmin {
		return models.Project{}, repo.ErrAdminRequired
	}

	project.CreatedDate = time.Now()
	return project, mongoErr(p.coll().Insert(project))
}

func (p projectRepo) Delete(u *models.User, uid string) error {
	if u == nil || !u.IsAdmin {
		return repo.ErrAdminRequired
	}

	return p.coll().RemoveId(uid)
}

func (p projectRepo) Search(u *models.User, query string) ([]models.Project, error) {
	base := permQuery(u)
	q := base

	if query != "" {
		q = bson.M{
			"$and": []bson.M{
				base,
				{
					"$or": []bson.M{
						{
							"name": bson.M{
								"$regex":   query,
								"$options": "i",
							},
						},
						{
							"key": bson.M{
								"$regex":   query,
								"$options": "i",
							},
						},
						{
							"lead": bson.M{
								"$regex":   query,
								"$options": "i",
							},
						},
					},
				},
			},
		}
	}

	var projects []models.Project
	err := p.coll().Find(q).All(&projects)
	return projects, mongoErr(err)
}
<commit_msg>implement has lead for mongo repo<commit_after>// Copyright 2017 Mathew Robinson <mrobinson@praelatus.io>. All rights reserved.
// Use of this source code is governed by the AGPLv3 license that can be found in
// the LICENSE file.

package mongo

import (
	"time"

	"github.com/praelatus/praelatus/models"
	"github.com/praelatus/praelatus/repo"
	mgo "gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

type projectRepo struct {
	conn *mgo.Session
}

func (p projectRepo) coll() *mgo.Collection {
	return p.conn.DB(dbName).C(projects)
}

func permWithID(u *models.User, uid string) bson.M {
	base := permQuery(u)
	return bson.M{
		"$and": []bson.M{
			{
				"_id": uid,
			},
			base,
		},
	}
}

func (p projectRepo) Get(u *models.User, uid string) (models.Project, error) {
	var project models.Project
	q := permWithID(u, uid)
	err := p.coll().Find(q).One(&project)
	return project, mongoErr(err)
}

func (p projectRepo) Update(u *models.User, uid string, updated models.Project) error {
	q := permWithID(u, uid)
	return mongoErr(p.coll().Update(q, updated))
}

func (p projectRepo) Create(u *models.User, project models.Project) (models.Project, error) {
	if u == nil || !u.IsAdmin {
		return models.Project{}, repo.ErrAdminRequired
	}

	project.CreatedDate = time.Now()
	return project, mongoErr(p.coll().Insert(project))
}

func (p projectRepo) Delete(u *models.User, uid string) error {
	if u == nil || !u.IsAdmin {
		return repo.ErrAdminRequired
	}

	return p.coll().RemoveId(uid)
}

func (p projectRepo) Search(u *models.User, query string) ([]models.Project, error) {
	base := permQuery(u)
	q := base

	if query != "" {
		q = bson.M{
			"$and": []bson.M{
				base,
				{
					"$or": []bson.M{
						{
							"name": bson.M{
								"$regex":   query,
								"$options": "i",
							},
						},
						{
							"key": bson.M{
								"$regex":   query,
								"$options": "i",
							},
						},
						{
							"lead": bson.M{
								"$regex":   query,
								"$options": "i",
							},
						},
					},
				},
			},
		}
	}

	var projects []models.Project
	err := p.coll().Find(q).All(&projects)
	return projects, mongoErr(err)
}

func (p projectRepo) HasLead(u *models.User, lead models.User) ([]models.Project, error) {
	base := permQuery(u)
	q := bson.M{
		"$and": []bson.M{
			base,
			{
				"lead": lead.Username,
			},
		},
	}

	var projects []models.Project
	err := p.coll().Find(q).All(&projects)
	return projects, mongoErr(err)
}
<|endoftext|>


<commit_before>package iptables

import (
	"bufio"
	"bytes"
	"errors"
	"log"
	"os/exec"
	"regexp"
	"strconv"
	"sync"

	"github.com/deckarep/golang-set"
	"github.com/luizbafilho/fusis/api/types"
	"github.com/luizbafilho/fusis/config"
	"github.com/luizbafilho/fusis/net"
	"github.com/luizbafilho/fusis/state"
)

var (
	ErrIptablesNotFound = errors.New("Iptables not found")
	ErrIptablesSnat     = errors.New("Iptables: error when inserting SNAT rule")
)

// Defines iptables actions
var (
	ADD = "-A"
	DEL = "-D"
)

type Syncer interface {
	Sync(state state.State) error
}

type IptablesMngr struct {
	sync.Mutex
	config *config.BalancerConfig
	path   string
}

type SnatRule struct {
	vaddr    string
	vport    string
	toSource string
}

func New(config *config.BalancerConfig) (*IptablesMngr, error) {
	path, err := exec.LookPath("iptables")
	if err != nil {
		return nil, ErrIptablesNotFound
	}

	return &IptablesMngr{
		config: config,
		path:   path,
	}, nil
}

func (i IptablesMngr) Sync(s state.State) error {
	i.Lock()
	defer i.Unlock()

	stateSet, err := i.getStateRulesSet(s)
	if err != nil {
		return err
	}

	kernelSet, err := i.getKernelRulesSet()
	if err != nil {
		return err
	}

	rulesToAdd := stateSet.Difference(kernelSet)
	rulesToRemove := kernelSet.Difference(stateSet)

	for r := range rulesToAdd.Iter() {
		rule := r.(SnatRule)
		err := i.addRule(rule)
		if err != nil {
			return err
		}
	}

	for r := range rulesToRemove.Iter() {
		rule := r.(SnatRule)
		err := i.removeRule(rule)
		if err != nil {
			return err
		}
	}

	return nil
}

func (i IptablesMngr) getKernelRulesSet() (mapset.Set, error) {
	kernelRules, err := i.getSnatRules()
	if err != nil {
		return nil, err
	}

	kernelSet := mapset.NewSet()
	for _, r := range kernelRules {
		kernelSet.Add(r)
	}

	return kernelSet, nil

}

func (i IptablesMngr) getStateRulesSet(s state.State) (mapset.Set, error) {
	stateRules, err := i.convertServicesToSnatRules(i.getNatServices(s))
	if err != nil {
		return nil, err
	}

	stateSet := mapset.NewSet()
	for _, r := range stateRules {
		stateSet.Add(r)
	}

	return stateSet, nil
}

func (i IptablesMngr) convertServicesToSnatRules(svcs []types.Service) ([]SnatRule, error) {
	rules := []SnatRule{}
	for _, s := range svcs {

		r, err := i.serviceToSnatRule(s)
		if err != nil {
			return rules, err
		}

		rules = append(rules, *r)
	}

	return rules, nil
}

func (i IptablesMngr) getNatServices(s state.State) []types.Service {
	natServices := []types.Service{}

	for _, svc := range s.GetServices() {
		if svc.IsNat() {
			natServices = append(natServices, svc)
		}
	}

	return natServices
}

func (i IptablesMngr) serviceToSnatRule(svc types.Service) (*SnatRule, error) {
	privateIp, err := net.GetIpByInterface(i.config.Interfaces.Outbound)
	if err != nil {
		return nil, err
	}

	rule := &SnatRule{
		vaddr:    svc.Address,
		vport:    strconv.Itoa(int(svc.Port)),
		toSource: privateIp,
	}

	return rule, nil
}

func (i IptablesMngr) addRule(r SnatRule) error {
	return i.execIptablesCommand(ADD, r)
}

func (i IptablesMngr) removeRule(r SnatRule) error {
	return i.execIptablesCommand(DEL, r)
}

func (i IptablesMngr) execIptablesCommand(action string, r SnatRule) error {
	cmd := exec.Command(i.path, "--wait", "-t", "nat", action, "POSTROUTING", "-m", "ipvs", "--vaddr", r.vaddr+"/32", "--vport", r.vport, "-j", "SNAT", "--to-source", r.toSource)
	err := cmd.Run()
	if err != nil {
		return err
	}

	return nil
}

func (i IptablesMngr) getSnatRules() ([]SnatRule, error) {
	out, err := exec.Command(i.path, "--wait", "--list", "-t", "nat").Output()
	if err != nil {
		log.Fatal(err)
	}

	r, _ := regexp.Compile(`vaddr\s([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})\svport\s(\d+)\sto:([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})`)
	scanner := bufio.NewScanner(bytes.NewReader(out))

	rules := []SnatRule{}
	for scanner.Scan() {
		matches := r.FindStringSubmatch(scanner.Text())
		if len(matches) == 0 {
			continue
		}

		rules = append(rules, SnatRule{
			vaddr:    matches[1],
			vport:    matches[2],
			toSource: matches[3],
		})
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return rules, nil
}
<commit_msg>Use logrus on iptables package<commit_after>package iptables

import (
	"bufio"
	"bytes"
	"errors"
	"os/exec"
	"regexp"
	"strconv"
	"sync"

	log "github.com/Sirupsen/logrus"
	"github.com/deckarep/golang-set"
	"github.com/luizbafilho/fusis/api/types"
	"github.com/luizbafilho/fusis/config"
	"github.com/luizbafilho/fusis/net"
	"github.com/luizbafilho/fusis/state"
)

var (
	ErrIptablesNotFound = errors.New("Iptables not found")
	ErrIptablesSnat     = errors.New("Iptables: error when inserting SNAT rule")
)

// Defines iptables actions
var (
	ADD = "-A"
	DEL = "-D"
)

type Syncer interface {
	Sync(state state.State) error
}

type IptablesMngr struct {
	sync.Mutex
	config *config.BalancerConfig
	path   string
}

type SnatRule struct {
	vaddr    string
	vport    string
	toSource string
}

func New(config *config.BalancerConfig) (*IptablesMngr, error) {
	path, err := exec.LookPath("iptables")
	if err != nil {
		return nil, ErrIptablesNotFound
	}

	return &IptablesMngr{
		config: config,
		path:   path,
	}, nil
}

func (i IptablesMngr) Sync(s state.State) error {
	i.Lock()
	defer i.Unlock()

	stateSet, err := i.getStateRulesSet(s)
	if err != nil {
		return err
	}

	kernelSet, err := i.getKernelRulesSet()
	if err != nil {
		return err
	}

	rulesToAdd := stateSet.Difference(kernelSet)
	rulesToRemove := kernelSet.Difference(stateSet)

	for r := range rulesToAdd.Iter() {
		rule := r.(SnatRule)
		err := i.addRule(rule)
		if err != nil {
			return err
		}
	}

	for r := range rulesToRemove.Iter() {
		rule := r.(SnatRule)
		err := i.removeRule(rule)
		if err != nil {
			return err
		}
	}

	return nil
}

func (i IptablesMngr) getKernelRulesSet() (mapset.Set, error) {
	kernelRules, err := i.getSnatRules()
	if err != nil {
		return nil, err
	}

	kernelSet := mapset.NewSet()
	for _, r := range kernelRules {
		kernelSet.Add(r)
	}

	return kernelSet, nil

}

func (i IptablesMngr) getStateRulesSet(s state.State) (mapset.Set, error) {
	stateRules, err := i.convertServicesToSnatRules(i.getNatServices(s))
	if err != nil {
		return nil, err
	}

	stateSet := mapset.NewSet()
	for _, r := range stateRules {
		stateSet.Add(r)
	}

	return stateSet, nil
}

func (i IptablesMngr) convertServicesToSnatRules(svcs []types.Service) ([]SnatRule, error) {
	rules := []SnatRule{}
	for _, s := range svcs {

		r, err := i.serviceToSnatRule(s)
		if err != nil {
			return rules, err
		}

		rules = append(rules, *r)
	}

	return rules, nil
}

func (i IptablesMngr) getNatServices(s state.State) []types.Service {
	natServices := []types.Service{}

	for _, svc := range s.GetServices() {
		if svc.IsNat() {
			natServices = append(natServices, svc)
		}
	}

	return natServices
}

func (i IptablesMngr) serviceToSnatRule(svc types.Service) (*SnatRule, error) {
	privateIp, err := net.GetIpByInterface(i.config.Interfaces.Outbound)
	if err != nil {
		return nil, err
	}

	rule := &SnatRule{
		vaddr:    svc.Address,
		vport:    strconv.Itoa(int(svc.Port)),
		toSource: privateIp,
	}

	return rule, nil
}

func (i IptablesMngr) addRule(r SnatRule) error {
	return i.execIptablesCommand(ADD, r)
}

func (i IptablesMngr) removeRule(r SnatRule) error {
	return i.execIptablesCommand(DEL, r)
}

func (i IptablesMngr) execIptablesCommand(action string, r SnatRule) error {
	cmd := exec.Command(i.path, "--wait", "-t", "nat", action, "POSTROUTING", "-m", "ipvs", "--vaddr", r.vaddr+"/32", "--vport", r.vport, "-j", "SNAT", "--to-source", r.toSource)
	err := cmd.Run()
	if err != nil {
		return err
	}

	return nil
}

func (i IptablesMngr) getSnatRules() ([]SnatRule, error) {
	out, err := exec.Command(i.path, "--wait", "--list", "-t", "nat").Output()
	if err != nil {
		log.Fatal("Error executing iptables", err)
	}

	r, _ := regexp.Compile(`vaddr\s([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})\svport\s(\d+)\sto:([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})`)
	scanner := bufio.NewScanner(bytes.NewReader(out))

	rules := []SnatRule{}
	for scanner.Scan() {
		matches := r.FindStringSubmatch(scanner.Text())
		if len(matches) == 0 {
			continue
		}

		rules = append(rules, SnatRule{
			vaddr:    matches[1],
			vport:    matches[2],
			toSource: matches[3],
		})
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return rules, nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/base64"
	"fmt"
	"math/big"
	"os"
	"time"

	"encoding/json"

	"github.com/mhe/gabi"
	"github.com/privacybydesign/irmago"
)

func main() {
	if len(os.Args) != 3 {
		fmt.Println("Usage: irmago metadata_attribute_in_decimal path_to_irma_configuration")
	}

	metaint, ok := new(big.Int).SetString(os.Args[1], 10)
	if !ok {
		fmt.Println("Could not parse argument as decimal integer:")
		os.Exit(1)
	}

	confpath := os.Args[2]
	conf, err := irma.NewConfiguration(confpath, "")
	if err != nil {
		fmt.Println("Failed to parse irma_configuration:", err)
		os.Exit(1)
	}
	err = conf.ParseFolder()
	if err != nil {
		fmt.Println("Failed to parse irma_configuration:", err)
		os.Exit(1)
	}

	meta := irma.MetadataFromInt(metaint, conf)
	typ := meta.CredentialType()
	var key *gabi.PublicKey

	if typ == nil {
		fmt.Println("Unknown credential type, hash:", base64.StdEncoding.EncodeToString(meta.CredentialTypeHash()))
	} else {
		fmt.Println("Identifier      :", typ.Identifier())
		key, err = meta.PublicKey()
		if err != nil {
			fmt.Println("Failed to parse public key", err)
		}
	}
	fmt.Println("Signed          :", meta.SigningDate().String())
	fmt.Println("Expires         :", meta.Expiry().String())
	fmt.Println("IsValid         :", meta.IsValid())
	fmt.Println("Version         :", meta.Version())
	fmt.Println("KeyCounter      :", meta.KeyCounter())
	if key != nil {
		fmt.Println("KeyExpires      :", time.Unix(key.ExpiryDate, 0))
		fmt.Println("KeyModulusBitlen:", key.N.BitLen())
	}

	fmt.Println()
	fmt.Println("CredentialType  :", prettyprint(typ))
}

func prettyprint(ob interface{}) string {
	b, err := json.MarshalIndent(ob, "", "  ")
	if err != nil {
		fmt.Println("error:", err)
	}
	return string(b)
}
<commit_msg>Swap irma_configuration and bigint params in irmameta tool<commit_after>package main

import (
	"encoding/base64"
	"fmt"
	"math/big"
	"os"
	"time"

	"encoding/json"

	"github.com/mhe/gabi"
	"github.com/privacybydesign/irmago"
)

func main() {
	if len(os.Args) != 3 {
		fmt.Println("Usage: irmago metadata_attribute_in_decimal path_to_irma_configuration")
	}

	metaint, ok := new(big.Int).SetString(os.Args[2], 10)
	if !ok {
		fmt.Println("Could not parse argument as decimal integer:")
		os.Exit(1)
	}

	confpath := os.Args[1]
	conf, err := irma.NewConfiguration(confpath, "")
	if err != nil {
		fmt.Println("Failed to parse irma_configuration:", err)
		os.Exit(1)
	}
	err = conf.ParseFolder()
	if err != nil {
		fmt.Println("Failed to parse irma_configuration:", err)
		os.Exit(1)
	}

	meta := irma.MetadataFromInt(metaint, conf)
	typ := meta.CredentialType()
	var key *gabi.PublicKey

	if typ == nil {
		fmt.Println("Unknown credential type, hash:", base64.StdEncoding.EncodeToString(meta.CredentialTypeHash()))
	} else {
		fmt.Println("Identifier      :", typ.Identifier())
		key, err = meta.PublicKey()
		if err != nil {
			fmt.Println("Failed to parse public key", err)
		}
	}
	fmt.Println("Signed          :", meta.SigningDate().String())
	fmt.Println("Expires         :", meta.Expiry().String())
	fmt.Println("IsValid         :", meta.IsValid())
	fmt.Println("Version         :", meta.Version())
	fmt.Println("KeyCounter      :", meta.KeyCounter())
	if key != nil {
		fmt.Println("KeyExpires      :", time.Unix(key.ExpiryDate, 0))
		fmt.Println("KeyModulusBitlen:", key.N.BitLen())
	}

	fmt.Println()
	fmt.Println("CredentialType  :", prettyprint(typ))
}

func prettyprint(ob interface{}) string {
	b, err := json.MarshalIndent(ob, "", "  ")
	if err != nil {
		fmt.Println("error:", err)
	}
	return string(b)
}
<|endoftext|>


<commit_before>// Copyright Istio 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 cmd

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/spf13/cobra"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/client-go/dynamic"

	"istio.io/istio/istioctl/pkg/clioptions"
	"istio.io/istio/istioctl/pkg/util/handlers"
	"istio.io/istio/pilot/pkg/xds"
	"istio.io/istio/pkg/config"
	"istio.io/istio/pkg/config/schema/collection"
	"istio.io/istio/pkg/config/schema/collections"
	"istio.io/istio/pkg/kube"
)

var (
	forFlag      string
	nameflag     string
	threshold    float32
	timeout      time.Duration
	generation   string
	verbose      bool
	targetSchema collection.Schema
	clientGetter func(string, string) (dynamic.Interface, error)
)

const pollInterval = time.Second

// waitCmd represents the wait command
func waitCmd() *cobra.Command {
	var opts clioptions.ControlPlaneOptions
	cmd := &cobra.Command{
		Use:   "wait [flags] <type> <name>[.<namespace>]",
		Short: "Wait for an Istio resource",
		Long:  `Waits for the specified condition to be true of an Istio resource.`,
		Example: `  # Wait until the bookinfo virtual service has been distributed to all proxies in the mesh
  istioctl experimental wait --for=distribution virtualservice bookinfo.default

  # Wait until 99% of the proxies receive the distribution, timing out after 5 minutes
  istioctl experimental wait --for=distribution --threshold=.99 --timeout=300 virtualservice bookinfo.default
`,
		RunE: func(cmd *cobra.Command, args []string) error {
			printVerbosef(cmd, "kubeconfig %s", kubeconfig)
			printVerbosef(cmd, "ctx %s", configContext)
			if forFlag == "delete" {
				return errors.New("wait for delete is not yet implemented")
			} else if forFlag != "distribution" {
				return fmt.Errorf("--for must be 'delete' or 'distribution', got: %s", forFlag)
			}
			var w *watcher
			ctx, cancel := context.WithTimeout(context.Background(), timeout)
			defer cancel()
			if generation == "" {
				w = getAndWatchResource(ctx) // setup version getter from kubernetes
			} else {
				w = withContext(ctx)
				w.Go(func(result chan string) error {
					result <- generation
					return nil
				})
			}
			// wait for all deployed versions to be contained in generations
			t := time.NewTicker(pollInterval)
			printVerbosef(cmd, "getting first version from chan")
			firstVersion, err := w.BlockingRead()
			if err != nil {
				return fmt.Errorf("unable to retrieve Kubernetes resource %s: %v", "", err)
			}
			generations := []string{firstVersion}
			targetResource := config.Key(
				targetSchema.Resource().Group(), targetSchema.Resource().Version(), targetSchema.Resource().Kind(),
				nameflag, namespace)
			for {
				// run the check here as soon as we start
				// because tickers won't run immediately
				present, notpresent, sdcnum, err := poll(cmd, generations, targetResource, opts)
				printVerbosef(cmd, "Received poll result: %d/%d", present, present+notpresent)
				if err != nil {
					return err
				} else if float32(present)/float32(present+notpresent) >= threshold {
					_, _ = fmt.Fprintf(cmd.OutOrStdout(), "Resource %s present on %d out of %d configurations for totally %d sidecars\n",
						targetResource, present, present+notpresent, sdcnum)
					return nil
				}
				select {
				case newVersion := <-w.resultsChan:
					printVerbosef(cmd, "received new target version: %s", newVersion)
					generations = append(generations, newVersion)
				case <-t.C:
					printVerbosef(cmd, "tick")
					continue
				case err = <-w.errorChan:
					return fmt.Errorf("unable to retrieve Kubernetes resource2 %s: %v", "", err)
				case <-ctx.Done():
					printVerbosef(cmd, "timeout")
					// I think this means the timeout has happened:
					t.Stop()
					return fmt.Errorf("timeout expired before resource %s became effective on all sidecars",
						targetResource)
				}
			}
		},
		Args: func(cmd *cobra.Command, args []string) error {
			if err := cobra.ExactArgs(2)(cmd, args); err != nil {
				return err
			}
			nameflag, namespace = handlers.InferPodInfo(args[1], handlers.HandleNamespace(namespace, defaultNamespace))
			return validateType(args[0])
		},
	}
	cmd.PersistentFlags().StringVar(&forFlag, "for", "distribution",
		"Wait condition, must be 'distribution' or 'delete'")
	cmd.PersistentFlags().DurationVar(&timeout, "timeout", time.Second*30,
		"The duration to wait before failing")
	cmd.PersistentFlags().Float32Var(&threshold, "threshold", 1,
		"The ratio of distribution required for success")
	cmd.PersistentFlags().StringVar(&generation, "generation", "",
		"Wait for a specific generation of config to become current, rather than using whatever is latest in "+
			"Kubernetes")
	cmd.PersistentFlags().BoolVarP(&verbose, "verbose", "v", false, "enables verbose output")
	_ = cmd.PersistentFlags().MarkHidden("verbose")
	opts.AttachControlPlaneFlags(cmd)
	return cmd
}

func printVerbosef(cmd *cobra.Command, template string, args ...interface{}) {
	if verbose {
		_, _ = fmt.Fprintf(cmd.OutOrStdout(), template+"\n", args...)
	}
}

func validateType(kind string) error {
	originalKind := kind

	// Remove any dashes.
	kind = strings.ReplaceAll(kind, "-", "")

	for _, s := range collections.Pilot.All() {
		if strings.EqualFold(kind, s.Resource().Kind()) {
			targetSchema = s
			return nil
		}
	}
	return fmt.Errorf("type %s is not recognized", originalKind)
}

func countVersions(versionCount map[string]int, configVersion string) {
	if count, ok := versionCount[configVersion]; ok {
		versionCount[configVersion] = count + 1
	} else {
		versionCount[configVersion] = 1
	}
}

func poll(cmd *cobra.Command,
	acceptedVersions []string,
	targetResource string,
	opts clioptions.ControlPlaneOptions) (present, notpresent, sdcnum int, err error) {
	kubeClient, err := kubeClientWithRevision(kubeconfig, configContext, opts.Revision)
	if err != nil {
		return 0, 0, 0, err
	}
	path := fmt.Sprintf("/debug/config_distribution?resource=%s", targetResource)
	pilotResponses, err := kubeClient.AllDiscoveryDo(context.TODO(), istioNamespace, path)
	if err != nil {
		return 0, 0, 0, fmt.Errorf("unable to query pilot for distribution "+
			"(are you using pilot version >= 1.4 with config distribution tracking on): %s", err)
	}
	sdcnum = 0
	versionCount := make(map[string]int)
	for _, response := range pilotResponses {
		var configVersions []xds.SyncedVersions
		err = json.Unmarshal(response, &configVersions)
		if err != nil {
			return 0, 0, 0, err
		}
		printVerbosef(cmd, "sync status: %v", configVersions)
		sdcnum += len(configVersions)
		for _, configVersion := range configVersions {
			countVersions(versionCount, configVersion.ClusterVersion)
			countVersions(versionCount, configVersion.RouteVersion)
			countVersions(versionCount, configVersion.ListenerVersion)
		}
	}

	for version, count := range versionCount {
		if contains(acceptedVersions, version) {
			present += count
		} else {
			notpresent += count
		}
	}
	return present, notpresent, sdcnum, nil
}

func init() {
	clientGetter = func(kubeconfig, context string) (dynamic.Interface, error) {
		config, err := kube.DefaultRestConfig(kubeconfig, context)
		if err != nil {
			return nil, err
		}
		cfg := dynamic.ConfigFor(config)
		dclient, err := dynamic.NewForConfig(cfg)
		if err != nil {
			return nil, err
		}
		return dclient, nil
	}
}

// getAndWatchResource ensures that Generations always contains
// the current generation of the targetResource, adding new versions
// as they are created.
func getAndWatchResource(ictx context.Context) *watcher {
	g := withContext(ictx)
	// copy nameflag to avoid race
	nf := nameflag
	g.Go(func(result chan string) error {
		// retrieve latest generation from Kubernetes
		dclient, err := clientGetter(kubeconfig, configContext)
		if err != nil {
			return err
		}
		collectionParts := strings.Split(targetSchema.Name().String(), "/")
		group := targetSchema.Resource().Group()
		version := targetSchema.Resource().Version()
		resource := collectionParts[3]
		r := dclient.Resource(schema.GroupVersionResource{Group: group, Version: version, Resource: resource}).Namespace(namespace)
		watch, err := r.Watch(context.TODO(), metav1.ListOptions{FieldSelector: "metadata.name=" + nf})
		if err != nil {
			return err
		}
		for w := range watch.ResultChan() {
			o := w.Object.(metav1.Object)
			if o.GetName() == nf {
				result <- strconv.FormatInt(o.GetGeneration(), 10)
			}
			select {
			case <-ictx.Done():
				return ictx.Err()
			default:
				continue
			}
		}

		return nil
	})
	return g
}

type watcher struct {
	resultsChan chan string
	errorChan   chan error
	ctx         context.Context
}

func withContext(ctx context.Context) *watcher {
	return &watcher{
		resultsChan: make(chan string, 1),
		errorChan:   make(chan error, 1),
		ctx:         ctx,
	}
}

func (w *watcher) Go(f func(chan string) error) {
	go func() {
		if err := f(w.resultsChan); err != nil {
			w.errorChan <- err
		}
	}()
}

func (w *watcher) BlockingRead() (string, error) {
	select {
	case err := <-w.errorChan:
		return "", err
	case res := <-w.resultsChan:
		return res, nil
	case <-w.ctx.Done():
		return "", w.ctx.Err()
	}
}
<commit_msg>fixed istioctl wait bug for integ-distroless-k8s-tests (#33107)<commit_after>// Copyright Istio 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 cmd

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/spf13/cobra"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/client-go/dynamic"

	"istio.io/istio/istioctl/pkg/clioptions"
	"istio.io/istio/istioctl/pkg/util/handlers"
	"istio.io/istio/pilot/pkg/xds"
	"istio.io/istio/pkg/config"
	"istio.io/istio/pkg/config/schema/collection"
	"istio.io/istio/pkg/config/schema/collections"
	"istio.io/istio/pkg/kube"
)

var (
	forFlag      string
	nameflag     string
	threshold    float32
	timeout      time.Duration
	generation   string
	verbose      bool
	targetSchema collection.Schema
	clientGetter func(string, string) (dynamic.Interface, error)
)

const pollInterval = time.Second

// waitCmd represents the wait command
func waitCmd() *cobra.Command {
	var opts clioptions.ControlPlaneOptions
	cmd := &cobra.Command{
		Use:   "wait [flags] <type> <name>[.<namespace>]",
		Short: "Wait for an Istio resource",
		Long:  `Waits for the specified condition to be true of an Istio resource.`,
		Example: `  # Wait until the bookinfo virtual service has been distributed to all proxies in the mesh
  istioctl experimental wait --for=distribution virtualservice bookinfo.default

  # Wait until 99% of the proxies receive the distribution, timing out after 5 minutes
  istioctl experimental wait --for=distribution --threshold=.99 --timeout=300 virtualservice bookinfo.default
`,
		RunE: func(cmd *cobra.Command, args []string) error {
			printVerbosef(cmd, "kubeconfig %s", kubeconfig)
			printVerbosef(cmd, "ctx %s", configContext)
			if forFlag == "delete" {
				return errors.New("wait for delete is not yet implemented")
			} else if forFlag != "distribution" {
				return fmt.Errorf("--for must be 'delete' or 'distribution', got: %s", forFlag)
			}
			var w *watcher
			ctx, cancel := context.WithTimeout(context.Background(), timeout)
			defer cancel()
			if generation == "" {
				w = getAndWatchResource(ctx) // setup version getter from kubernetes
			} else {
				w = withContext(ctx)
				w.Go(func(result chan string) error {
					result <- generation
					return nil
				})
			}
			// wait for all deployed versions to be contained in generations
			t := time.NewTicker(pollInterval)
			printVerbosef(cmd, "getting first version from chan")
			firstVersion, err := w.BlockingRead()
			if err != nil {
				return fmt.Errorf("unable to retrieve Kubernetes resource %s: %v", "", err)
			}
			generations := []string{firstVersion}
			targetResource := config.Key(
				targetSchema.Resource().Group(), targetSchema.Resource().Version(), targetSchema.Resource().Kind(),
				nameflag, namespace)
			for {
				// run the check here as soon as we start
				// because tickers won't run immediately
				present, notpresent, sdcnum, err := poll(cmd, generations, targetResource, opts)
				printVerbosef(cmd, "Received poll result: %d/%d", present, present+notpresent)
				if err != nil {
					return err
				} else if float32(present)/float32(present+notpresent) >= threshold {
					_, _ = fmt.Fprintf(cmd.OutOrStdout(), "Resource %s present on %d out of %d configurations for totally %d sidecars\n",
						targetResource, present, present+notpresent, sdcnum)
					return nil
				}
				select {
				case newVersion := <-w.resultsChan:
					printVerbosef(cmd, "received new target version: %s", newVersion)
					generations = append(generations, newVersion)
				case <-t.C:
					printVerbosef(cmd, "tick")
					continue
				case err = <-w.errorChan:
					return fmt.Errorf("unable to retrieve Kubernetes resource2 %s: %v", "", err)
				case <-ctx.Done():
					printVerbosef(cmd, "timeout")
					// I think this means the timeout has happened:
					t.Stop()
					return fmt.Errorf("timeout expired before resource %s became effective on all sidecars",
						targetResource)
				}
			}
		},
		Args: func(cmd *cobra.Command, args []string) error {
			if err := cobra.ExactArgs(2)(cmd, args); err != nil {
				return err
			}
			nameflag, namespace = handlers.InferPodInfo(args[1], handlers.HandleNamespace(namespace, defaultNamespace))
			return validateType(args[0])
		},
	}
	cmd.PersistentFlags().StringVar(&forFlag, "for", "distribution",
		"Wait condition, must be 'distribution' or 'delete'")
	cmd.PersistentFlags().DurationVar(&timeout, "timeout", time.Second*30,
		"The duration to wait before failing")
	cmd.PersistentFlags().Float32Var(&threshold, "threshold", 1,
		"The ratio of distribution required for success")
	cmd.PersistentFlags().StringVar(&generation, "generation", "",
		"Wait for a specific generation of config to become current, rather than using whatever is latest in "+
			"Kubernetes")
	cmd.PersistentFlags().BoolVarP(&verbose, "verbose", "v", false, "enables verbose output")
	_ = cmd.PersistentFlags().MarkHidden("verbose")
	opts.AttachControlPlaneFlags(cmd)
	return cmd
}

func printVerbosef(cmd *cobra.Command, template string, args ...interface{}) {
	if verbose {
		_, _ = fmt.Fprintf(cmd.OutOrStdout(), template+"\n", args...)
	}
}

func validateType(kind string) error {
	originalKind := kind

	// Remove any dashes.
	kind = strings.ReplaceAll(kind, "-", "")

	for _, s := range collections.Pilot.All() {
		if strings.EqualFold(kind, s.Resource().Kind()) {
			targetSchema = s
			return nil
		}
	}
	return fmt.Errorf("type %s is not recognized", originalKind)
}

func countVersions(versionCount map[string]int, configVersion string) {
	if count, ok := versionCount[configVersion]; ok {
		versionCount[configVersion] = count + 1
	} else {
		versionCount[configVersion] = 1
	}
}

func poll(cmd *cobra.Command,
	acceptedVersions []string,
	targetResource string,
	opts clioptions.ControlPlaneOptions) (present, notpresent, sdcnum int, err error) {
	kubeClient, err := kubeClientWithRevision(kubeconfig, configContext, opts.Revision)
	if err != nil {
		return 0, 0, 0, err
	}
	path := fmt.Sprintf("/debug/config_distribution?resource=%s", targetResource)
	pilotResponses, err := kubeClient.AllDiscoveryDo(context.TODO(), istioNamespace, path)
	if err != nil {
		return 0, 0, 0, fmt.Errorf("unable to query pilot for distribution "+
			"(are you using pilot version >= 1.4 with config distribution tracking on): %s", err)
	}
	sdcnum = 0
	versionCount := make(map[string]int)
	for _, response := range pilotResponses {
		var configVersions []xds.SyncedVersions
		err = json.Unmarshal(response, &configVersions)
		if err != nil {
			return 0, 0, 0, err
		}
		printVerbosef(cmd, "sync status: %v", configVersions)
		sdcnum += len(configVersions)
		for _, configVersion := range configVersions {
			countVersions(versionCount, configVersion.ClusterVersion)
			countVersions(versionCount, configVersion.RouteVersion)
			countVersions(versionCount, configVersion.ListenerVersion)
		}
	}

	for version, count := range versionCount {
		if contains(acceptedVersions, version) {
			present += count
		} else {
			notpresent += count
		}
	}
	return present, notpresent, sdcnum, nil
}

func init() {
	clientGetter = func(kubeconfig, context string) (dynamic.Interface, error) {
		config, err := kube.DefaultRestConfig(kubeconfig, context)
		if err != nil {
			return nil, err
		}
		cfg := dynamic.ConfigFor(config)
		dclient, err := dynamic.NewForConfig(cfg)
		if err != nil {
			return nil, err
		}
		return dclient, nil
	}
}

// getAndWatchResource ensures that Generations always contains
// the current generation of the targetResource, adding new versions
// as they are created.
func getAndWatchResource(ictx context.Context) *watcher {
	g := withContext(ictx)
	// copy nameflag to avoid race
	nf := nameflag
	g.Go(func(result chan string) error {
		// retrieve latest generation from Kubernetes
		dclient, err := clientGetter(kubeconfig, configContext)
		if err != nil {
			return err
		}
		collectionParts := strings.Split(targetSchema.Name().String(), "/")
		group := targetSchema.Resource().Group()
		version := targetSchema.Resource().Version()
		resource := collectionParts[3]
		r := dclient.Resource(schema.GroupVersionResource{Group: group, Version: version, Resource: resource}).Namespace(namespace)
		watch, err := r.Watch(context.TODO(), metav1.ListOptions{FieldSelector: "metadata.name=" + nf})
		if err != nil {
			return err
		}
		for w := range watch.ResultChan() {
			o, ok := w.Object.(metav1.Object)
			if !ok {
				continue
			}
			if o.GetName() == nf {
				result <- strconv.FormatInt(o.GetGeneration(), 10)
			}
			select {
			case <-ictx.Done():
				return ictx.Err()
			default:
				continue
			}
		}

		return nil
	})
	return g
}

type watcher struct {
	resultsChan chan string
	errorChan   chan error
	ctx         context.Context
}

func withContext(ctx context.Context) *watcher {
	return &watcher{
		resultsChan: make(chan string, 1),
		errorChan:   make(chan error, 1),
		ctx:         ctx,
	}
}

func (w *watcher) Go(f func(chan string) error) {
	go func() {
		if err := f(w.resultsChan); err != nil {
			w.errorChan <- err
		}
	}()
}

func (w *watcher) BlockingRead() (string, error) {
	select {
	case err := <-w.errorChan:
		return "", err
	case res := <-w.resultsChan:
		return res, nil
	case <-w.ctx.Done():
		return "", w.ctx.Err()
	}
}
<|endoftext|>