<commit_before>/**
 * (C) Copyright 2014, Deft Labs
 *
 * 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 dlshared

import (
	"fmt"
	"strings"
	"reflect"
	"strconv"
	"net/http"
	"io/ioutil"
	"encoding/json"
	"labix.org/v2/mgo/bson"
	"github.com/gorilla/mux"
)

type HttpParamDataType int8
type HttpParamType int8

const(
	HttpIntParam = HttpParamDataType(0)
	HttpStringParam = HttpParamDataType(1)
	HttpFloatParam = HttpParamDataType(2)
	HttpBoolParam = HttpParamDataType(3) // Boolean types include: 1, t, T, TRUE, true, True, 0, f, F, FALSE, false
	HttpObjectIdParam = HttpParamDataType(4)

	// All of the param types only support single values (i.e., no slices). If multiple values are present, the
	// first is taken.
	HttpParamPost = HttpParamType(0)
	HttpParamJsonPost = HttpParamType(1) // When the content body is posted in json format. This only supports one level
	HttpParamQuery = HttpParamType(2)
	HttpParamHeader = HttpParamType(3)
	HttpParamPath = HttpParamType(4) // This must be declared as {someName} in the path mapping
)

type HttpContext struct {
	Response http.ResponseWriter
	Request *http.Request
	Params map[string]*HttpParam
	ErrorCodes []string
	Errors []error

	postJson map[string]interface{}

	body []byte
}

type HttpParam struct {
	Name string
	InvalidErrorCode string
	DataType HttpParamDataType
	Type HttpParamType
	Required bool
	MinLength int
	MaxLength int
	Post bool
	Value interface{}
	Raw string
	Valid bool
	Present bool // If value is present and parsed properly
}

// Make sure your params are present and valid before trying to access.
func (self *HttpParam) Int() int { return self.Value.(int) }

func (self *HttpParam) Float() float64 { return self.Value.(float64) }

func (self *HttpParam) String() string { return self.Value.(string) }

func (self *HttpParam) Bool() bool { return self.Value.(bool) }

func (self *HttpParam) ObjectId() *bson.ObjectId { return self.Value.(*bson.ObjectId) }

// Set a valid value for a param. Missing can be valid, but not present.
func (self *HttpParam) setPresentValue(value interface{}) {
	self.Present = true
	self.Value = value
}

// Validate the params. If any of the params are invalid, false is returned. You must call
// this first before calling the ErrorCodes []string. If not params are defined, this always
// returns "true". If there are raw data extraction errors, this is always false (e.g., body missing or incorrect).
func (self *HttpContext) ParamsAreValid() bool {

	if len(self.Errors) != 0 { return false }

	if len(self.Params) == 0 { return true }

	for _, param := range self.Params {
		switch param.DataType {
			case HttpIntParam: validateIntParam(self, param)
			case HttpStringParam: validateStringParam(self, param)
			case HttpFloatParam: validateFloatParam(self, param)
			case HttpBoolParam: validateBoolParam(self, param)
			case HttpObjectIdParam: validateObjectIdParam(self, param)
		}
	}

	return len(self.ErrorCodes) == 0
}

func (self *HttpContext) ParamInt(name string) int { return self.Params[name].Int() }

func (self *HttpContext) ParamFloat(name string) float64 { return self.Params[name].Float() }

func (self *HttpContext) ParamString(name string) string { return self.Params[name].String() }

func (self *HttpContext) HasParam(name string) bool {
	if val, found := self.Params[name]; !found || val == nil || val.Value == nil { return false }
	return true
}

func (self *HttpContext) ParamBool(name string) bool { return self.Params[name].Bool() }

func (self *HttpContext) ParamObjectId(name string) *bson.ObjectId { return self.Params[name].ObjectId() }

func (self *HttpContext) HasRawErrors() bool { return len(self.Errors) > 0 }

// This returns the param value as a string. If the param is missing or empty,
// the string will be len == 0.
func retrieveParamValue(ctx *HttpContext, param *HttpParam) string {
	switch param.Type {
		case HttpParamPost: return strings.TrimSpace(ctx.Request.PostFormValue(param.Name))
		case HttpParamJsonPost: return retrieveJsonParamValue(ctx, param)
		case HttpParamQuery: return strings.TrimSpace(ctx.Request.FormValue(param.Name))
		case HttpParamHeader: return strings.TrimSpace(ctx.Request.Header.Get(param.Name))
		case HttpParamPath: return strings.TrimSpace(mux.Vars(ctx.Request)[param.Name])
	}
	return nadaStr
}

func retrieveJsonParamValue(ctx *HttpContext, param *HttpParam) string {

	if len(ctx.Errors) > 0 {
		return nadaStr
	}

	// If this is the first access, read the body
	if len(ctx.body) == 0 {
		var err error
		ctx.body, err = ioutil.ReadAll(ctx.Request.Body)
		if err != nil {
			ctx.Errors = append(ctx.Errors, NewStackError("Error in raw data extraction - error: %v", err))
			return nadaStr
		}
	}

	if ctx.postJson == nil {
		var genJson interface{}
		err := json.Unmarshal(ctx.body, &genJson)
		if err != nil {
			ctx.Errors = append(ctx.Errors, NewStackError("Error in raw json data extraction - error: %v", err))
			return nadaStr
		}

		ctx.postJson = genJson.(map[string]interface{})
	}

	// Look for the value in the json. The json may hold the data in a variety
	// of formats. Convert back to a string to deal with the other data types :-(
	val, found := ctx.postJson[param.Name]
	if !found {
		return nadaStr
	}

	valType := reflect.TypeOf(val)

	if valType == nil { return nadaStr }

	switch valType.Kind() {
		case reflect.Invalid: return nadaStr
		case reflect.Bool: return fmt.Sprintf("%t", val.(bool))
		case reflect.Float64: return fmt.Sprintf("%g", val.(float64))
		case reflect.String: return val.(string)
		default: return nadaStr
	}

	return nadaStr
}

func appendInvalidErrorCode(ctx *HttpContext, param *HttpParam) {
	// Do not dupplicate error codes.
	for i := range ctx.ErrorCodes { if ctx.ErrorCodes[i] == param.InvalidErrorCode { return } }

	if len(param.InvalidErrorCode) == 0 {
		panic(fmt.Sprintf("We do not have an error code defined for param: %s - required: %t", param.Name, param.Required))
	}

	ctx.ErrorCodes = append(ctx.ErrorCodes, param.InvalidErrorCode)
	param.Valid = false
}

func validateIntParam(ctx *HttpContext, param *HttpParam) {
	param.Raw = retrieveParamValue(ctx, param)

	if len(param.Raw) == 0 && param.Required {
		appendInvalidErrorCode(ctx, param)
		return
	}

	if len(param.Raw) == 0 {
		return
	}

	if val, err := strconv.Atoi(param.Raw); err != nil {
		appendInvalidErrorCode(ctx, param)
	} else {
		param.setPresentValue(val)
	}
}

func validateStringParam(ctx *HttpContext, param *HttpParam) {

	param.Raw = retrieveParamValue(ctx, param)

	if len(param.Raw) == 0 && param.Required {
		appendInvalidErrorCode(ctx, param)
		return
	}

	if param.Required && param.MinLength > 0 && len(param.Raw) < param.MinLength {
		appendInvalidErrorCode(ctx, param)
		return
	}

	if param.Required && param.MaxLength > 0 && len(param.Raw) > param.MaxLength {
		appendInvalidErrorCode(ctx, param)
		return
	}

	param.setPresentValue(param.Raw)
}

func validateFloatParam(ctx *HttpContext, param *HttpParam) {
	param.Raw = retrieveParamValue(ctx, param)

	if len(param.Raw) == 0 && param.Required {
		appendInvalidErrorCode(ctx, param)
		return
	}

	if len(param.Raw) == 0 {
		return
	}

	if val, err := strconv.ParseFloat(param.Raw, 64); err != nil {
		appendInvalidErrorCode(ctx, param)
	} else {
		param.setPresentValue(val)
	}
}

func validateObjectIdParam(ctx *HttpContext, param *HttpParam) {

	param.Raw = retrieveParamValue(ctx, param)

	if len(param.Raw) == 0 && param.Required {
		appendInvalidErrorCode(ctx, param)
		return
	}

	if len(param.Raw) == 0 { return }

	if !bson.IsObjectIdHex(param.Raw) {
		appendInvalidErrorCode(ctx, param)
		return
	}

	value := bson.ObjectIdHex(param.Raw)
	param.setPresentValue(&value)
}

// Boolean types include: 1, t, T, TRUE, true, True, 0, f, F, FALSE, false
func validateBoolParam(ctx *HttpContext, param *HttpParam) {

	param.Raw = retrieveParamValue(ctx, param)

	if len(param.Raw) == 0 && param.Required {
		appendInvalidErrorCode(ctx, param)
		return
	}

	if len(param.Raw) == 0 { return }

	if val, err := strconv.ParseBool(param.Raw); err != nil {
		appendInvalidErrorCode(ctx, param)
	} else {
		param.setPresentValue(val)
	}
}

func (self *HttpContext) DefineIntParam(name, invalidErrorCode string, paramType HttpParamType, required bool) {
	self.Params[name] = &HttpParam{ Name: name, InvalidErrorCode: invalidErrorCode, DataType: HttpIntParam, Required: required, Type: paramType, Valid: true }
}

// Boolean types include: 1, t, T, TRUE, true, True, 0, f, F, FALSE, false
func (self *HttpContext) DefineBoolParam(name, invalidErrorCode string, paramType HttpParamType, required bool) {
	self.Params[name] = &HttpParam{ Name: name, InvalidErrorCode: invalidErrorCode, DataType: HttpBoolParam, Required: required, Type: paramType, Valid: true }
}

func (self *HttpContext) DefineFloatParam(name, invalidErrorCode string, paramType HttpParamType, required bool) {
	self.Params[name] = &HttpParam{ Name: name, InvalidErrorCode: invalidErrorCode, DataType: HttpFloatParam, Required: required, Type: paramType, Valid: true }
}

func (self *HttpContext) DefineObjectIdParam(name, invalidErrorCode string, paramType HttpParamType, required bool) {
	self.Params[name] = &HttpParam{ Name: name, InvalidErrorCode: invalidErrorCode, DataType: HttpObjectIdParam, Required: required, Type: paramType, Valid: true }
}

func (self *HttpContext) DefineStringParam(name, invalidErrorCode string, paramType HttpParamType, required bool, minLength, maxLength int) {
	self.Params[name] = &HttpParam{ Name: name, InvalidErrorCode: invalidErrorCode, DataType: HttpStringParam, Required: required, Type: paramType, Valid: true }
}

// Call this method to init the http context struct.
func NewHttpContext(response http.ResponseWriter, request *http.Request) *HttpContext {
	return &HttpContext{ Response: response, Request: request, Params: make(map[string]*HttpParam) }
}

// This method returns true if the http request method is a HTTP post. If the
// field missing or incorrect, false is returned. This method will panic if
// the request is nil.
func IsHttpMethodPost(request *http.Request) bool {
	if request == nil {
		panic("request param is nil")
	}
	return len(request.Method) > 0 && strings.ToUpper(request.Method) == HttpPostMethod
}

// Write an http ok response string. The content type is text/plain.
func WriteOkResponseString(response http.ResponseWriter, msg string) error {
	if response == nil {
		return NewStackError("response param is nil")
	}

	msgLength := len(msg)

	if msgLength == 0 {
		return NewStackError("Response message is an empty string")
	}

	response.Header().Set(ContentTypeHeader, ContentTypeTextPlain)

	written, err := response.Write([]byte(msg))

	if err != nil {
		return err
	}

	if written != msgLength {
		return NewStackError("Did not write full message - bytes written %d - expected %d", written, msgLength)
	}

	return nil
}

// Encode and write a json response. If there is a problem encoding an http 500 is sent and an
// error is returned. If there are problems writting the response an error is returned.
func JsonEncodeAndWriteResponse(response http.ResponseWriter, value interface{}) error {

	if value == nil {
		return NewStackError("Nil value passed")
	}

	rawJson, err := json.Marshal(value)
	if err != nil {
		http.Error(response, "Error", 500)
		return NewStackError("Unable to marshal json: %v", err)
	}

	response.Header().Set(ContentTypeHeader, ContentTypeJson)

	written, err := response.Write(rawJson)
	if err != nil {
		return NewStackError("Unable to write response: %v", err)
	}

	if written != len(rawJson) {
		return NewStackError("Unable to write full response - wrote: %d - expected: %d", written, len(rawJson))
	}

	return nil
}

<commit_msg>layout change<commit_after>/**
 * (C) Copyright 2014, Deft Labs
 *
 * 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 dlshared

import (
	"fmt"
	"strings"
	"reflect"
	"strconv"
	"net/http"
	"io/ioutil"
	"encoding/json"
	"labix.org/v2/mgo/bson"
	"github.com/gorilla/mux"
)

type HttpParamDataType int8
type HttpParamType int8

const(
	HttpIntParam = HttpParamDataType(0)
	HttpStringParam = HttpParamDataType(1)
	HttpFloatParam = HttpParamDataType(2)
	HttpBoolParam = HttpParamDataType(3) // Boolean types include: 1, t, T, TRUE, true, True, 0, f, F, FALSE, false
	HttpObjectIdParam = HttpParamDataType(4)

	// All of the param types only support single values (i.e., no slices). If multiple values are present, the
	// first is taken.
	HttpParamPost = HttpParamType(0)
	HttpParamJsonPost = HttpParamType(1) // When the content body is posted in json format. This only supports one level
	HttpParamQuery = HttpParamType(2)
	HttpParamHeader = HttpParamType(3)
	HttpParamPath = HttpParamType(4) // This must be declared as {someName} in the path mapping
)

type HttpContext struct {
	Response http.ResponseWriter
	Request *http.Request
	Params map[string]*HttpParam
	ErrorCodes []string
	Errors []error

	postJson map[string]interface{}

	body []byte
}

type HttpParam struct {
	Name string
	InvalidErrorCode string
	DataType HttpParamDataType
	Type HttpParamType
	Required bool
	MinLength int
	MaxLength int
	Post bool
	Value interface{}
	Raw string
	Valid bool
	Present bool // If value is present and parsed properly
}

// Make sure your params are present and valid before trying to access.
func (self *HttpParam) Int() int { return self.Value.(int) }

func (self *HttpParam) Float() float64 { return self.Value.(float64) }

func (self *HttpParam) String() string { return self.Value.(string) }

func (self *HttpParam) Bool() bool { return self.Value.(bool) }

func (self *HttpParam) ObjectId() *bson.ObjectId { return self.Value.(*bson.ObjectId) }

// Set a valid value for a param. Missing can be valid, but not present.
func (self *HttpParam) setPresentValue(value interface{}) {
	self.Present = true
	self.Value = value
}

// Validate the params. If any of the params are invalid, false is returned. You must call
// this first before calling the ErrorCodes []string. If not params are defined, this always
// returns "true". If there are raw data extraction errors, this is always false (e.g., body missing or incorrect).
func (self *HttpContext) ParamsAreValid() bool {

	if len(self.Errors) != 0 { return false }

	if len(self.Params) == 0 { return true }

	for _, param := range self.Params {
		switch param.DataType {
			case HttpIntParam: validateIntParam(self, param)
			case HttpStringParam: validateStringParam(self, param)
			case HttpFloatParam: validateFloatParam(self, param)
			case HttpBoolParam: validateBoolParam(self, param)
			case HttpObjectIdParam: validateObjectIdParam(self, param)
		}
	}

	return len(self.ErrorCodes) == 0
}

func (self *HttpContext) ParamInt(name string) int { return self.Params[name].Int() }

func (self *HttpContext) ParamFloat(name string) float64 { return self.Params[name].Float() }

func (self *HttpContext) ParamString(name string) string { return self.Params[name].String() }

func (self *HttpContext) HasParam(name string) bool {
	if val, found := self.Params[name]; !found || val == nil || val.Value == nil { return false }
	return true
}

func (self *HttpContext) ParamBool(name string) bool { return self.Params[name].Bool() }

func (self *HttpContext) ParamObjectId(name string) *bson.ObjectId { return self.Params[name].ObjectId() }

func (self *HttpContext) HasRawErrors() bool { return len(self.Errors) > 0 }

// This returns the param value as a string. If the param is missing or empty,
// the string will be len == 0.
func retrieveParamValue(ctx *HttpContext, param *HttpParam) string {
	switch param.Type {
		case HttpParamPost: return strings.TrimSpace(ctx.Request.PostFormValue(param.Name))
		case HttpParamJsonPost: return retrieveJsonParamValue(ctx, param)
		case HttpParamQuery: return strings.TrimSpace(ctx.Request.FormValue(param.Name))
		case HttpParamHeader: return strings.TrimSpace(ctx.Request.Header.Get(param.Name))
		case HttpParamPath: return strings.TrimSpace(mux.Vars(ctx.Request)[param.Name])
	}
	return nadaStr
}

func retrieveJsonParamValue(ctx *HttpContext, param *HttpParam) string {

	if len(ctx.Errors) > 0 {
		return nadaStr
	}

	// If this is the first access, read the body
	if len(ctx.body) == 0 {
		var err error
		ctx.body, err = ioutil.ReadAll(ctx.Request.Body)
		if err != nil {
			ctx.Errors = append(ctx.Errors, NewStackError("Error in raw data extraction - error: %v", err))
			return nadaStr
		}
	}

	if ctx.postJson == nil {
		var genJson interface{}
		err := json.Unmarshal(ctx.body, &genJson)
		if err != nil {
			ctx.Errors = append(ctx.Errors, NewStackError("Error in raw json data extraction - error: %v", err))
			return nadaStr
		}

		ctx.postJson = genJson.(map[string]interface{})
	}

	// Look for the value in the json. The json may hold the data in a variety
	// of formats. Convert back to a string to deal with the other data types :-(
	val, found := ctx.postJson[param.Name]
	if !found {
		return nadaStr
	}

	valType := reflect.TypeOf(val)

	if valType == nil { return nadaStr }

	switch valType.Kind() {
		case reflect.Invalid: return nadaStr
		case reflect.Bool: return fmt.Sprintf("%t", val.(bool))
		case reflect.Float64: return fmt.Sprintf("%g", val.(float64))
		case reflect.String: return val.(string)
		default: return nadaStr
	}

	return nadaStr
}

func appendInvalidErrorCode(ctx *HttpContext, param *HttpParam) {
	// Do not dupplicate error codes.
	for i := range ctx.ErrorCodes { if ctx.ErrorCodes[i] == param.InvalidErrorCode { return } }

	if len(param.InvalidErrorCode) == 0 {
		panic(fmt.Sprintf("We do not have an error code defined for param: %s - required: %t", param.Name, param.Required))
	}

	ctx.ErrorCodes = append(ctx.ErrorCodes, param.InvalidErrorCode)
	param.Valid = false
}

func validateIntParam(ctx *HttpContext, param *HttpParam) {

	param.Raw = retrieveParamValue(ctx, param)

	if len(param.Raw) == 0 && param.Required { appendInvalidErrorCode(ctx, param); return }

	if len(param.Raw) == 0 { return }

	if val, err := strconv.Atoi(param.Raw); err != nil { appendInvalidErrorCode(ctx, param)
	} else { param.setPresentValue(val) }
}

func validateStringParam(ctx *HttpContext, param *HttpParam) {

	param.Raw = retrieveParamValue(ctx, param)

	if len(param.Raw) == 0 && param.Required {
		appendInvalidErrorCode(ctx, param)
		return
	}

	if param.Required && param.MinLength > 0 && len(param.Raw) < param.MinLength {
		appendInvalidErrorCode(ctx, param)
		return
	}

	if param.Required && param.MaxLength > 0 && len(param.Raw) > param.MaxLength {
		appendInvalidErrorCode(ctx, param)
		return
	}

	param.setPresentValue(param.Raw)
}

func validateFloatParam(ctx *HttpContext, param *HttpParam) {
	param.Raw = retrieveParamValue(ctx, param)

	if len(param.Raw) == 0 && param.Required {
		appendInvalidErrorCode(ctx, param)
		return
	}

	if len(param.Raw) == 0 {
		return
	}

	if val, err := strconv.ParseFloat(param.Raw, 64); err != nil {
		appendInvalidErrorCode(ctx, param)
	} else {
		param.setPresentValue(val)
	}
}

func validateObjectIdParam(ctx *HttpContext, param *HttpParam) {

	param.Raw = retrieveParamValue(ctx, param)

	if len(param.Raw) == 0 && param.Required {
		appendInvalidErrorCode(ctx, param)
		return
	}

	if len(param.Raw) == 0 { return }

	if !bson.IsObjectIdHex(param.Raw) {
		appendInvalidErrorCode(ctx, param)
		return
	}

	value := bson.ObjectIdHex(param.Raw)
	param.setPresentValue(&value)
}

// Boolean types include: 1, t, T, TRUE, true, True, 0, f, F, FALSE, false
func validateBoolParam(ctx *HttpContext, param *HttpParam) {

	param.Raw = retrieveParamValue(ctx, param)

	if len(param.Raw) == 0 && param.Required {
		appendInvalidErrorCode(ctx, param)
		return
	}

	if len(param.Raw) == 0 { return }

	if val, err := strconv.ParseBool(param.Raw); err != nil {
		appendInvalidErrorCode(ctx, param)
	} else {
		param.setPresentValue(val)
	}
}

func (self *HttpContext) DefineIntParam(name, invalidErrorCode string, paramType HttpParamType, required bool) {
	self.Params[name] = &HttpParam{ Name: name, InvalidErrorCode: invalidErrorCode, DataType: HttpIntParam, Required: required, Type: paramType, Valid: true }
}

// Boolean types include: 1, t, T, TRUE, true, True, 0, f, F, FALSE, false
func (self *HttpContext) DefineBoolParam(name, invalidErrorCode string, paramType HttpParamType, required bool) {
	self.Params[name] = &HttpParam{ Name: name, InvalidErrorCode: invalidErrorCode, DataType: HttpBoolParam, Required: required, Type: paramType, Valid: true }
}

func (self *HttpContext) DefineFloatParam(name, invalidErrorCode string, paramType HttpParamType, required bool) {
	self.Params[name] = &HttpParam{ Name: name, InvalidErrorCode: invalidErrorCode, DataType: HttpFloatParam, Required: required, Type: paramType, Valid: true }
}

func (self *HttpContext) DefineObjectIdParam(name, invalidErrorCode string, paramType HttpParamType, required bool) {
	self.Params[name] = &HttpParam{ Name: name, InvalidErrorCode: invalidErrorCode, DataType: HttpObjectIdParam, Required: required, Type: paramType, Valid: true }
}

func (self *HttpContext) DefineStringParam(name, invalidErrorCode string, paramType HttpParamType, required bool, minLength, maxLength int) {
	self.Params[name] = &HttpParam{ Name: name, InvalidErrorCode: invalidErrorCode, DataType: HttpStringParam, Required: required, Type: paramType, Valid: true }
}

// Call this method to init the http context struct.
func NewHttpContext(response http.ResponseWriter, request *http.Request) *HttpContext {
	return &HttpContext{ Response: response, Request: request, Params: make(map[string]*HttpParam) }
}

// This method returns true if the http request method is a HTTP post. If the
// field missing or incorrect, false is returned. This method will panic if
// the request is nil.
func IsHttpMethodPost(request *http.Request) bool {
	if request == nil {
		panic("request param is nil")
	}
	return len(request.Method) > 0 && strings.ToUpper(request.Method) == HttpPostMethod
}

// Write an http ok response string. The content type is text/plain.
func WriteOkResponseString(response http.ResponseWriter, msg string) error {
	if response == nil {
		return NewStackError("response param is nil")
	}

	msgLength := len(msg)

	if msgLength == 0 {
		return NewStackError("Response message is an empty string")
	}

	response.Header().Set(ContentTypeHeader, ContentTypeTextPlain)

	written, err := response.Write([]byte(msg))

	if err != nil {
		return err
	}

	if written != msgLength {
		return NewStackError("Did not write full message - bytes written %d - expected %d", written, msgLength)
	}

	return nil
}

// Encode and write a json response. If there is a problem encoding an http 500 is sent and an
// error is returned. If there are problems writting the response an error is returned.
func JsonEncodeAndWriteResponse(response http.ResponseWriter, value interface{}) error {

	if value == nil {
		return NewStackError("Nil value passed")
	}

	rawJson, err := json.Marshal(value)
	if err != nil {
		http.Error(response, "Error", 500)
		return NewStackError("Unable to marshal json: %v", err)
	}

	response.Header().Set(ContentTypeHeader, ContentTypeJson)

	written, err := response.Write(rawJson)
	if err != nil {
		return NewStackError("Unable to write response: %v", err)
	}

	if written != len(rawJson) {
		return NewStackError("Unable to write full response - wrote: %d - expected: %d", written, len(rawJson))
	}

	return nil
}

<|endoftext|>


<commit_before>// HTTP server for EditInsite
package httpserver

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strings"

	"github.com/editinsite/editinsite/config"
	"github.com/editinsite/editinsite/projects"
)

func projectHandler(w http.ResponseWriter, r *http.Request) {
	if len(r.URL.Path) > len("/projects/") {
		fileHandler(w, r)
	} else {
		list := projects.SortByID()
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(list)
	}
}

func fileHandler(w http.ResponseWriter, r *http.Request) {
	path := r.URL.Path[len("/projects/"):]
	if strings.Contains(path, "..") {
		// for security
		http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
		return
	}

	projEnd := strings.IndexByte(path, '/')
	if projEnd == -1 {
		http.Redirect(w, r, r.URL.Path+"/", http.StatusMovedPermanently)
		return
	}
	projectID := path[0:projEnd]
	project := projects.Registry[projectID]

	filePath := path[projEnd:]
	isDir := filePath[len(filePath)-1] == '/'
	if !isDir {
		var err error
		isDir, err = project.IsDir(filePath)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		if isDir {
			http.Redirect(w, r, r.URL.Path+"/", http.StatusMovedPermanently)
			return
		}
	}
	if isDir {
		listHandler(w, r, project, filePath)
	} else {
		if r.Method == "POST" {
			saveHandler(w, r, project, filePath)
		} else {
			loadHandler(w, r, project, filePath)
		}
	}
}

func listHandler(w http.ResponseWriter, r *http.Request, p *projects.Workspace, subDir string) {
	list, err := p.Files(subDir)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	//w.Header().Set("Content-Length", strconv.Itoa(len(buffer)))
	json.NewEncoder(w).Encode(list)
}

func loadHandler(w http.ResponseWriter, r *http.Request, p *projects.Workspace, file string) {
	if err := p.ServeFile(w, r, file); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}

func saveHandler(w http.ResponseWriter, r *http.Request, p *projects.Workspace, file string) {
	err := p.SaveFile(file, r.Body)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}

func Start() error {
	port := fmt.Sprintf(":%d", config.Values.Port)
	http.HandleFunc("/projects/", projectHandler)
	http.Handle("/", http.FileServer(http.Dir("ui")))
	return http.ListenAndServe(port, nil)
}
<commit_msg>Show HTML directory listings<commit_after>// HTTP server for EditInsite
package httpserver

import (
	"encoding/json"
	"fmt"
	"html/template"
	"net/http"
	"strings"

	"github.com/editinsite/editinsite/config"
	"github.com/editinsite/editinsite/projects"
)

var dirListTemplate = template.Must(template.New("dir").Parse(`<!doctype html>

	  <meta charset="utf-8">
	  <title>Directory listing</title>

	  	{{range .}}
{{.}}
{{end}}

	`))

func projectHandler(w http.ResponseWriter, r *http.Request) {
	if len(r.URL.Path) > len("/projects/") {
		fileHandler(w, r)
	} else {
		list := projects.SortByID()
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(list)
	}
}

func fileHandler(w http.ResponseWriter, r *http.Request) {
	path := r.URL.Path[len("/projects/"):]
	if strings.Contains(path, "..") {
		// for security
		http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
		return
	}

	projEnd := strings.IndexByte(path, '/')
	if projEnd == -1 {
		http.Redirect(w, r, r.URL.Path+"/", http.StatusMovedPermanently)
		return
	}
	projectID := path[0:projEnd]
	project := projects.Registry[projectID]

	filePath := path[projEnd:]
	isDir := filePath[len(filePath)-1] == '/'
	if !isDir {
		var err error
		isDir, err = project.IsDir(filePath)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		if isDir {
			http.Redirect(w, r, r.URL.Path+"/", http.StatusMovedPermanently)
			return
		}
	}
	if isDir {
		listHandler(w, r, project, filePath)
	} else {
		if r.Method == "POST" {
			saveHandler(w, r, project, filePath)
		} else {
			loadHandler(w, r, project, filePath)
		}
	}
}

func listHandler(w http.ResponseWriter, r *http.Request, p *projects.Workspace, subDir string) {
	list, err := p.Files(subDir)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	accept := r.Header.Get("Accept")
	if strings.Contains(accept, "application/json") {
		w.Header().Set("Content-Type", "application/json")
		err = json.NewEncoder(w).Encode(list)
	} else {
		err = dirListTemplate.Execute(w, list)
	}
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}

func loadHandler(w http.ResponseWriter, r *http.Request, p *projects.Workspace, file string) {
	if err := p.ServeFile(w, r, file); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}

func saveHandler(w http.ResponseWriter, r *http.Request, p *projects.Workspace, file string) {
	err := p.SaveFile(file, r.Body)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}

func Start() error {
	port := fmt.Sprintf(":%d", config.Values.Port)
	http.HandleFunc("/projects/", projectHandler)
	http.Handle("/", http.FileServer(http.Dir("ui")))
	return http.ListenAndServe(port, nil)
}
<|endoftext|>


<commit_before>package grpcutil

import (
	"sync"
)

var bufPool = sync.Pool{
	New: func() interface{} {
		// This buffer size is:
		// 1. Reasonably smaller than the max gRPC size
		// 2. Small enough that having hundreds of these buffers won't
		// overwhelm the node
		// 3. Large enough for message-sending to be efficient
		return make([]byte, MaxMsgSize/10)
	},
}

// Get returns a buffer.  The buffer may or may not be freshly allocated.
func GetBuffer() []byte {
	return bufPool.Get().([]byte)
}

// Put returns the buffer to the pool.
func PutBuffer(buf []byte) {
	bufPool.Put(buf)
}
<commit_msg>Fix linter errors<commit_after>package grpcutil

import (
	"sync"
)

var bufPool = sync.Pool{
	New: func() interface{} {
		// This buffer size is:
		// 1. Reasonably smaller than the max gRPC size
		// 2. Small enough that having hundreds of these buffers won't
		// overwhelm the node
		// 3. Large enough for message-sending to be efficient
		return make([]byte, MaxMsgSize/10)
	},
}

// GetBuffer returns a buffer.  The buffer may or may not be freshly
// allocated.
func GetBuffer() []byte {
	return bufPool.Get().([]byte)
}

// PutBuffer returns the buffer to the pool.
func PutBuffer(buf []byte) {
	bufPool.Put(buf)
}
<|endoftext|>


<commit_before>package grpcutil

import (
	"strings"
	"sync"

	"github.com/pachyderm/pachyderm/src/client/pkg/tracing"
	"google.golang.org/grpc"
)

// Dialer defines a grpc.ClientConn connection dialer.
type Dialer interface {
	Dial(address string) (*grpc.ClientConn, error)
	CloseConns() error
}

// NewDialer creates a Dialer.
func NewDialer(opts ...grpc.DialOption) Dialer {
	return newDialer(opts...)
}

type dialer struct {
	opts []grpc.DialOption
	// A map from addresses to connections
	connMap map[string]*grpc.ClientConn
	lock    sync.Mutex
}

func newDialer(opts ...grpc.DialOption) *dialer {
	return &dialer{
		opts:    opts,
		connMap: make(map[string]*grpc.ClientConn),
	}
}

func (d *dialer) Dial(addr string) (*grpc.ClientConn, error) {
	d.lock.Lock()
	defer d.lock.Unlock()

	if conn, ok := d.connMap[addr]; ok {
		return conn, nil
	}
	opts := append(d.opts,
		grpc.WithUnaryInterceptor(tracing.UnaryClientInterceptor()),
		grpc.WithStreamInterceptor(tracing.StreamClientInterceptor()),
	)
	if strings.Index(addr, "dns:///") == -1 {
		addr = "dns:///" + addr
	}
	conn, err := grpc.Dial(addr, opts...)
	if err != nil {
		return nil, err
	}
	d.connMap[addr] = conn
	return conn, nil
}

func (d *dialer) CloseConns() error {
	d.lock.Lock()
	defer d.lock.Unlock()
	for addr, conn := range d.connMap {
		if err := conn.Close(); err != nil {
			return err
		}
		delete(d.connMap, addr)
	}
	return nil
}
<commit_msg>Fix lint issues with string.Index<commit_after>package grpcutil

import (
	"strings"
	"sync"

	"github.com/pachyderm/pachyderm/src/client/pkg/tracing"
	"google.golang.org/grpc"
)

// Dialer defines a grpc.ClientConn connection dialer.
type Dialer interface {
	Dial(address string) (*grpc.ClientConn, error)
	CloseConns() error
}

// NewDialer creates a Dialer.
func NewDialer(opts ...grpc.DialOption) Dialer {
	return newDialer(opts...)
}

type dialer struct {
	opts []grpc.DialOption
	// A map from addresses to connections
	connMap map[string]*grpc.ClientConn
	lock    sync.Mutex
}

func newDialer(opts ...grpc.DialOption) *dialer {
	return &dialer{
		opts:    opts,
		connMap: make(map[string]*grpc.ClientConn),
	}
}

func (d *dialer) Dial(addr string) (*grpc.ClientConn, error) {
	d.lock.Lock()
	defer d.lock.Unlock()

	if conn, ok := d.connMap[addr]; ok {
		return conn, nil
	}
	opts := append(d.opts,
		grpc.WithUnaryInterceptor(tracing.UnaryClientInterceptor()),
		grpc.WithStreamInterceptor(tracing.StreamClientInterceptor()),
	)
	if !strings.HasPrefix(addr, "dns:///") {
		addr = "dns:///" + addr
	}
	conn, err := grpc.Dial(addr, opts...)
	if err != nil {
		return nil, err
	}
	d.connMap[addr] = conn
	return conn, nil
}

func (d *dialer) CloseConns() error {
	d.lock.Lock()
	defer d.lock.Unlock()
	for addr, conn := range d.connMap {
		if err := conn.Close(); err != nil {
			return err
		}
		delete(d.connMap, addr)
	}
	return nil
}
<|endoftext|>


<commit_before>package grpcutil

import (
	"context"
	"fmt"
	"math"
	"net"
	"time"

	"golang.org/x/sync/errgroup"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/keepalive"

	"github.com/pachyderm/pachyderm/src/client/pkg/tls"
	"github.com/pachyderm/pachyderm/src/client/pkg/tracing"
	log "github.com/sirupsen/logrus"
)

// Server is a convenience wrapper to gRPC servers that simplifies their
// setup and execution
type Server struct {
	Server *grpc.Server
	eg     *errgroup.Group
}

// NewServer creates a new gRPC server, but does not start serving yet.
//
// If 'publicPortTLSAllowed' is set, grpcutil may enable TLS. This should be
// set for public ports that serve GRPC services to 3rd party clients. If set,
// the criterion for actually serving over TLS is: if a signed TLS cert and
// corresponding private key in 'TLSVolumePath', this will serve GRPC traffic
// over TLS. If either are missing this will serve GRPC traffic over
// unencrypted HTTP,
func NewServer(ctx context.Context, publicPortTLSAllowed bool) (*Server, error) {
	opts := []grpc.ServerOption{
		grpc.MaxConcurrentStreams(math.MaxUint32),
		grpc.MaxRecvMsgSize(MaxMsgSize),
		grpc.MaxSendMsgSize(MaxMsgSize),
		grpc.KeepaliveEnforcementPolicy(keepalive.EnforcementPolicy{
			MinTime:             5 * time.Second,
			PermitWithoutStream: true,
		}),
		grpc.UnaryInterceptor(tracing.UnaryServerInterceptor()),
		grpc.StreamInterceptor(tracing.StreamServerInterceptor()),
	}

	if publicPortTLSAllowed {
		// Validate environment
		certPath, keyPath, err := tls.GetCertPaths()
		if err != nil {
			log.Warnf("TLS disabled: %v", err)
		} else {
			// Read TLS cert and key
			transportCreds, err := credentials.NewServerTLSFromFile(certPath, keyPath)
			if err != nil {
				return nil, fmt.Errorf("couldn't build transport creds: %v", err)
			}
			opts = append(opts, grpc.Creds(transportCreds))
		}
	}

	server := grpc.NewServer(opts...)
	eg, ctx := errgroup.WithContext(ctx)

	eg.Go(func() error {
		<-ctx.Done()
		server.GracefulStop() // This also closes the listeners
		return nil
	})

	return &Server{
		Server: server,
		eg:     eg,
	}, nil
}

// ListenTCP causes the gRPC server to listen on a given TCP host and port
func (s *Server) ListenTCP(host string, port uint16) (net.Listener, error) {
	listener, err := net.Listen("tcp", fmt.Sprintf("%s:%d", host, port))
	if err != nil {
		return nil, err
	}

	s.eg.Go(func() error {
		return s.Server.Serve(listener)
	})

	return listener, nil
}

// ListenUDS causes the gRPC server to listen on a given Unix domain socket
// path
func (s *Server) ListenUDS(name string) (net.Listener, error) {
	listener, err := net.Listen("unix", name)
	if err != nil {
		return nil, err
	}

	s.eg.Go(func() error {
		return s.Server.Serve(listener)
	})

	return listener, nil
}

// Wait causes the gRPC server to wait until it finishes, returning any errors
// that happened
func (s *Server) Wait() error {
	return s.eg.Wait()
}
<commit_msg>Removed ListenUDS<commit_after>package grpcutil

import (
	"context"
	"fmt"
	"math"
	"net"
	"time"

	"golang.org/x/sync/errgroup"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/keepalive"

	"github.com/pachyderm/pachyderm/src/client/pkg/tls"
	"github.com/pachyderm/pachyderm/src/client/pkg/tracing"
	log "github.com/sirupsen/logrus"
)

// Server is a convenience wrapper to gRPC servers that simplifies their
// setup and execution
type Server struct {
	Server *grpc.Server
	eg     *errgroup.Group
}

// NewServer creates a new gRPC server, but does not start serving yet.
//
// If 'publicPortTLSAllowed' is set, grpcutil may enable TLS. This should be
// set for public ports that serve GRPC services to 3rd party clients. If set,
// the criterion for actually serving over TLS is: if a signed TLS cert and
// corresponding private key in 'TLSVolumePath', this will serve GRPC traffic
// over TLS. If either are missing this will serve GRPC traffic over
// unencrypted HTTP,
func NewServer(ctx context.Context, publicPortTLSAllowed bool) (*Server, error) {
	opts := []grpc.ServerOption{
		grpc.MaxConcurrentStreams(math.MaxUint32),
		grpc.MaxRecvMsgSize(MaxMsgSize),
		grpc.MaxSendMsgSize(MaxMsgSize),
		grpc.KeepaliveEnforcementPolicy(keepalive.EnforcementPolicy{
			MinTime:             5 * time.Second,
			PermitWithoutStream: true,
		}),
		grpc.UnaryInterceptor(tracing.UnaryServerInterceptor()),
		grpc.StreamInterceptor(tracing.StreamServerInterceptor()),
	}

	if publicPortTLSAllowed {
		// Validate environment
		certPath, keyPath, err := tls.GetCertPaths()
		if err != nil {
			log.Warnf("TLS disabled: %v", err)
		} else {
			// Read TLS cert and key
			transportCreds, err := credentials.NewServerTLSFromFile(certPath, keyPath)
			if err != nil {
				return nil, fmt.Errorf("couldn't build transport creds: %v", err)
			}
			opts = append(opts, grpc.Creds(transportCreds))
		}
	}

	server := grpc.NewServer(opts...)
	eg, ctx := errgroup.WithContext(ctx)

	eg.Go(func() error {
		<-ctx.Done()
		server.GracefulStop() // This also closes the listeners
		return nil
	})

	return &Server{
		Server: server,
		eg:     eg,
	}, nil
}

// ListenTCP causes the gRPC server to listen on a given TCP host and port
func (s *Server) ListenTCP(host string, port uint16) (net.Listener, error) {
	listener, err := net.Listen("tcp", fmt.Sprintf("%s:%d", host, port))
	if err != nil {
		return nil, err
	}

	s.eg.Go(func() error {
		return s.Server.Serve(listener)
	})

	return listener, nil
}

// Wait causes the gRPC server to wait until it finishes, returning any errors
// that happened
func (s *Server) Wait() error {
	return s.eg.Wait()
}
<|endoftext|>


<commit_before>package main

import (
	"testing"
	"time"

	mkr "github.com/mackerelio/mackerel-client-go"
)

func TestFormatJoinedAlert(t *testing.T) {

	const location = "UTC"

	loc, err := time.LoadLocation(location)
	if err != nil {
		loc = time.FixedZone(location, 0*60*60)
	}
	time.Local = loc

	a := &mkr.Alert{ID: "123", Type: "connectivity", Status: "critical", HostID: "1234", MonitorID: "12345", OpenedAt: 100}
	h := &mkr.Host{ID: "1234", Name: "foo", Roles: mkr.Roles{}, Status: "working"}
	m := &mkr.MonitorConnectivity{ID: "12345", Type: "connectivity"}
	as := alertSet{a, h, m}
	answer := "123 1970-01-01 00:01:40 critical connectivity foo working []"

	str := formatJoinedAlert(&as, false)
	if str != answer {
		t.Errorf("should be '%s' but '%s'", answer, str)
	}
}
<commit_msg>refactor TestFormatJoinedAlert in alerts_test<commit_after>package main

import (
	"testing"
	"time"

	mkr "github.com/mackerelio/mackerel-client-go"
)

func TestFormatJoinedAlert(t *testing.T) {

	const location = "UTC"

	loc, err := time.LoadLocation(location)
	if err != nil {
		loc = time.FixedZone(location, 0*60*60)
	}
	time.Local = loc

	testCases := []struct {
		alertSet *alertSet
		want     string
	}{
		{
			&alertSet{
				&mkr.Alert{ID: "2tZhm", Type: "connectivity", Status: "critical", HostID: "3XYyG", MonitorID: "5rXR3", OpenedAt: 100},
				&mkr.Host{ID: "3XYyG", Name: "foo", Roles: mkr.Roles{}, Status: "working"},
				&mkr.MonitorConnectivity{ID: "5rXR3", Type: "connectivity"},
			},
			"2tZhm 1970-01-01 00:01:40 critical connectivity foo working []",
		},
	}

	for _, testCase := range testCases {
		str := formatJoinedAlert(testCase.alertSet, false)
		if str != testCase.want {
			t.Errorf("should be '%s' but '%s'", testCase.want, str)
		}
	}
}
<|endoftext|>


<commit_before><commit_msg>Add size in collecting container logs<commit_after><|endoftext|>


<commit_before>package main

import (
	"bufio"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"time"
)

func (s *server) syncFromPeers() {
	for _, peer := range s.peers {
		go s.syncFromPeer(peer)
	}
}

func (srv *server) syncFromPeer(host string) {
	var ims time.Time
	var sleep time.Duration
	for {
		time.Sleep(sleep)
		req, err := http.NewRequest("GET", "http://"+host+"/webfist/bump", nil)
		if err != nil {
			panic("Bogus host: " + host + ": " + err.Error())
		}
		sleep = 30 * time.Second
		res, err := http.DefaultClient.Do(req)
		if err != nil {
			log.Printf("Fetch error from %s: %v", host, err)
			continue
		}
		sc := bufio.NewScanner(res.Body)
		for sc.Scan() {
			s := strings.Split(sc.Text(), " ")
			if len(s) != 2 {
				continue
			}
			modTime, err := time.Parse(time.RFC3339, s[0])
			if err != nil {
				continue
			}
			if modTime.After(ims) {
				ims = modTime
			}
			s = strings.Split(s[1], "-")
			if len(s) != 2 {
				continue
			}
			addrHexKey, encSHA1 := s[0], s[1]
			_, err = srv.storage.StatEncryptedEmail(addrHexKey, encSHA1)
			if err != nil {
				log.Printf("Need to fetch %s-%s", addrHexKey, encSHA1)
				res, err := http.Get("http://" + host + "/webfist/proof/" + addrHexKey + "-" + encSHA1)
				if err != nil {
					log.Printf("Error fetching %s-%s: %v", addrHexKey, encSHA1, err)
					continue
				}
				slurp, err := ioutil.ReadAll(io.LimitReader(res.Body, 100<<10))
				if err != nil {
					log.Printf("Error fetching %s-%s: %v", addrHexKey, encSHA1, err)
					continue
				}
				err = srv.storage.PutEncryptedEmail(addrHexKey, encSHA1, slurp)
				if err != nil {
					log.Printf("Error storing fetched %s-%s: %v", addrHexKey, encSHA1, err)
					continue
				}
			}
		}
		if err := sc.Err(); err != nil {
			log.Printf("Scan error: %v", err)
		}
		res.Body.Close()
	}
}
<commit_msg>sync logs<commit_after>package main

import (
	"bufio"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"time"
)

func (s *server) syncFromPeers() {
	for _, peer := range s.peers {
		go s.syncFromPeer(peer)
	}
}

func (srv *server) syncFromPeer(host string) {
	log.Printf("Starting sync from %v", host)
	var ims time.Time
	var sleep time.Duration
	for {
		time.Sleep(sleep)
		url := "http://"+host+"/webfist/bump"
		log.Printf("Doing request to %v", url)
		req, err := http.NewRequest("GET", url, nil)
		if err != nil {
			panic("Bogus host: " + host + ": " + err.Error())
		}
		sleep = 30 * time.Second
		res, err := http.DefaultClient.Do(req)
		if err != nil {
			log.Printf("Fetch error from %s: %v", host, err)
			continue
		}
		sc := bufio.NewScanner(res.Body)
		for sc.Scan() {
			s := strings.Split(sc.Text(), " ")
			if len(s) != 2 {
				continue
			}
			modTime, err := time.Parse(time.RFC3339, s[0])
			if err != nil {
				continue
			}
			if modTime.After(ims) {
				ims = modTime
			}
			s = strings.Split(s[1], "-")
			if len(s) != 2 {
				continue
			}
			addrHexKey, encSHA1 := s[0], s[1]
			_, err = srv.storage.StatEncryptedEmail(addrHexKey, encSHA1)
			if err != nil {
				log.Printf("Need to fetch %s-%s", addrHexKey, encSHA1)
				res, err := http.Get("http://" + host + "/webfist/proof/" + addrHexKey + "-" + encSHA1)
				if err != nil {
					log.Printf("Error fetching %s-%s: %v", addrHexKey, encSHA1, err)
					continue
				}
				slurp, err := ioutil.ReadAll(io.LimitReader(res.Body, 100<<10))
				if err != nil {
					log.Printf("Error fetching %s-%s: %v", addrHexKey, encSHA1, err)
					continue
				}
				err = srv.storage.PutEncryptedEmail(addrHexKey, encSHA1, slurp)
				if err != nil {
					log.Printf("Error storing fetched %s-%s: %v", addrHexKey, encSHA1, err)
					continue
				}
				log.Printf("Synced %s-%s from %s", addrHexKey, encSHA1, host)
			}
		}
		if err := sc.Err(); err != nil {
			log.Printf("Scan error: %v", err)
		}
		res.Body.Close()
	}
}
<|endoftext|>


<commit_before>package event

import (
	"fmt"
	"log"
	"testing"
)

var (
	numIndexes = 0
)

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

func newTestIndex() *Index {
	numIndexes += 1
	return NewIndex(fmt.Sprintf("test%d.db", numIndexes))
}

func newTestEvent(h, s, ss string, m float64) *Event {
	return &Event{
		Host:       h,
		Service:    s,
		SubService: ss,
		Metric:     m,
	}
}

func TestIndexPut(t *testing.T) {
	i := newTestIndex()
	defer i.Delete()

	e1 := newTestEvent("h", "s", "ss", 0.0)
	e2 := newTestEvent("h", "s", "ss", 0.0)

	i.PutEvent(e1)
	i.PutEvent(e2)

	e3 := i.GetEvent([]byte(e2.IndexName()))
	log.Println(e3)

}

func BenchmarkIndexPutJSON(b *testing.B) {
	index := newTestIndex()
	index.pool = NewEncodingPool(NewJsonEncoder, NewJsonDecoder, 4)
	defer index.Delete()
	events := make([]*Event, 1000)

	for i := 0; i < 1000; i++ {
		events[i] = newTestEvent("h", "s", "ss", float64(i))
	}

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		index.PutEvent(events[i%1000])
	}
}

func BenchmarkIndexPutMSGP(b *testing.B) {
	index := newTestIndex()
	index.pool = NewEncodingPool(NewMsgPackEncoder, NewMsgPackDecoder, 4)
	defer index.Delete()
	events := make([]*Event, 1000)

	for i := 0; i < 1000; i++ {
		events[i] = newTestEvent("h", "s", "ss", float64(i))
	}

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		index.PutEvent(events[i%1000])
	}
}
func TestDeleteIncidentById(t *testing.T) {
	i := newTestIndex()
	defer i.Delete()
	e := newTestEvent("h", "s", "ss", 1)
	in := NewIncident("DeleteIncident", i, e)

	in = i.GetIncident(in.Id)
	if in == nil {
		t.Fail()
	}

	i.DeleteIncidentById(in.Id)

	in = i.GetIncident(in.Id)
	if in != nil {
		t.Fail()
	}
}

func TestDeleteIncidentByEvent(t *testing.T) {
	i := newTestIndex()
	defer i.Delete()
	e := newTestEvent("h", "s", "ss", 1)
	in := NewIncident("DeleteIncident", i, e)

	in = i.GetIncident(in.Id)
	if in == nil {
		t.Fail()
	}

	i.DeleteIncidentByEvent(e)
	in = i.GetIncident(in.Id)
	if in != nil {
		t.Fail()
	}
}

func TestListIncidents(t *testing.T) {
	i := newTestIndex()
	defer i.Delete()
	e := newTestEvent("h", "s", "ss", 1)
	in := NewIncident("ListIncidents", i, e)

	ins := i.ListIncidents()

	if len(ins[0].EventName) != len(in.EventName) {
		t.Fail()
	}
}

func TestAddIncident(t *testing.T) {
	i := newTestIndex()
	defer i.Delete()
	e := newTestEvent("h", "s", "ss", 1)
	in := NewIncident("test", i, e)

	b := i.GetIncident(in.Id)

	if len(in.EventName) != len(b.EventName) {
		t.Fail()
	}
}
<commit_msg>added back test<commit_after>package event

import (
	"fmt"
	"log"
	"testing"
)

var (
	numIndexes = 0
)

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

func newTestIndex() *Index {
	numIndexes += 1
	return NewIndex(fmt.Sprintf("test%d.db", numIndexes))
}

func newTestEvent(h, s, ss string, m float64) *Event {
	return &Event{
		Host:       h,
		Service:    s,
		SubService: ss,
		Metric:     m,
	}
}

func TestIndexPut(t *testing.T) {
	i := newTestIndex()
	defer i.Delete()

	e1 := newTestEvent("h", "s", "ss", 0.0)
	e2 := newTestEvent("h", "s", "ss", 0.0)

	i.PutEvent(e1)
	i.PutEvent(e2)

	e3 := i.GetEvent([]byte(e2.IndexName()))

	if e3.LastEvent == nil {
		t.Fail()
	}

}

func BenchmarkIndexPutJSON(b *testing.B) {
	index := newTestIndex()
	index.pool = NewEncodingPool(NewJsonEncoder, NewJsonDecoder, 4)
	defer index.Delete()
	events := make([]*Event, 1000)

	for i := 0; i < 1000; i++ {
		events[i] = newTestEvent("h", "s", "ss", float64(i))
	}

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		index.PutEvent(events[i%1000])
	}
}

func BenchmarkIndexPutMSGP(b *testing.B) {
	index := newTestIndex()
	index.pool = NewEncodingPool(NewMsgPackEncoder, NewMsgPackDecoder, 4)
	defer index.Delete()
	events := make([]*Event, 1000)

	for i := 0; i < 1000; i++ {
		events[i] = newTestEvent("h", "s", "ss", float64(i))
	}

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		index.PutEvent(events[i%1000])
	}
}
func TestDeleteIncidentById(t *testing.T) {
	i := newTestIndex()
	defer i.Delete()
	e := newTestEvent("h", "s", "ss", 1)
	in := NewIncident("DeleteIncident", i, e)

	in = i.GetIncident(in.Id)
	if in == nil {
		t.Fail()
	}

	i.DeleteIncidentById(in.Id)

	in = i.GetIncident(in.Id)
	if in != nil {
		t.Fail()
	}
}

func TestDeleteIncidentByEvent(t *testing.T) {
	i := newTestIndex()
	defer i.Delete()
	e := newTestEvent("h", "s", "ss", 1)
	in := NewIncident("DeleteIncident", i, e)

	in = i.GetIncident(in.Id)
	if in == nil {
		t.Fail()
	}

	i.DeleteIncidentByEvent(e)
	in = i.GetIncident(in.Id)
	if in != nil {
		t.Fail()
	}
}

func TestListIncidents(t *testing.T) {
	i := newTestIndex()
	defer i.Delete()
	e := newTestEvent("h", "s", "ss", 1)
	in := NewIncident("ListIncidents", i, e)

	ins := i.ListIncidents()

	if len(ins[0].EventName) != len(in.EventName) {
		t.Fail()
	}
}

func TestAddIncident(t *testing.T) {
	i := newTestIndex()
	defer i.Delete()
	e := newTestEvent("h", "s", "ss", 1)
	in := NewIncident("test", i, e)

	b := i.GetIncident(in.Id)

	if len(in.EventName) != len(b.EventName) {
		t.Fail()
	}
}
<|endoftext|>


<commit_before>package gamekit

import (
	"sync"

	"github.com/sheenobu/go-gamekit/pair"
	"github.com/sheenobu/rxgen/rx"

	"github.com/veandco/go-sdl2/sdl"
)

// The WindowManager is responsible for handling a group of windows with different event handlers
type WindowManager struct {
	WindowCount     *rx.Uint32
	CurrentWindowID *rx.Uint32

	lock     sync.RWMutex
	windows  map[uint32]*Window
	handlers map[uint32]func(e sdl.Event)
}

// NewWindowManager builds a new window manager
func NewWindowManager() *WindowManager {
	return &WindowManager{
		CurrentWindowID: rx.NewUint32(0),
		WindowCount:     rx.NewUint32(0),

		windows:  make(map[uint32]*Window),
		handlers: make(map[uint32]func(e sdl.Event)),
	}
}

// NewWindow creates a new window, returning the window
func (wm *WindowManager) NewWindow(title string, w int, h int, extraFlags int) (*Window, error) {

	wm.lock.Lock()
	defer wm.lock.Unlock()

	window, renderer, err := sdl.CreateWindowAndRenderer(w, h, 0)

	if err != nil {
		return nil, err
	}

	win := &Window{
		ID:       window.GetID(),
		Size:     pair.NewRxInt32Pair(pair.Int32Pair{L: int32(w), R: int32(h)}),
		Window:   window,
		Renderer: renderer,
		Mouse:    NewMouse(),

		wm: wm,
	}

	wm.windows[window.GetID()] = win

	wm.handlers[window.GetID()] = func(e sdl.Event) {
		switch t := e.(type) {
		case *sdl.WindowEvent:
			if t.Event == sdl.WINDOWEVENT_CLOSE {
				window.Destroy()
			}
		}
	}

	wm.WindowCount.Set(wm.WindowCount.Get() + 1)

	return win, nil
}

// Destroy destroys the window
func (wm *WindowManager) Destroy(id uint32) {
	wm.lock.RLock()
	defer wm.lock.RUnlock()

	wm.windows[id].Window.Destroy()
	delete(wm.windows, id)
	delete(wm.handlers, id)
}

// SetHandler sets the event handler for the given window id
func (wm *WindowManager) SetHandler(id uint32, h func(e sdl.Event)) {
	wm.lock.Lock()
	defer wm.lock.Unlock()
	wm.handlers[id] = h
}

// DispatchEvents handles events and sends them to the required window event handlers.
func (wm *WindowManager) DispatchEvents() {
	if e := sdl.PollEvent(); e != nil {
		wm.lock.RLock()
		defer wm.lock.RUnlock()

		switch t := e.(type) {
		case *sdl.WindowEvent:
			switch t.Event {
			case sdl.WINDOWEVENT_RESIZED:
				// if we handle these events we actually get a crash????
				// wm.windows[t.WindowID].Size.Set(pair.Int32Pair{t.Data1, t.Data2})
				return
			case sdl.WINDOWEVENT_SIZE_CHANGED:
				wm.windows[t.WindowID].Size.Set(pair.Int32Pair{L: t.Data1, R: t.Data2})
				return
			case sdl.WINDOWEVENT_FOCUS_GAINED:
				wm.CurrentWindowID.Set(t.WindowID)
				return
			case sdl.WINDOWEVENT_FOCUS_LOST:
				wm.CurrentWindowID.Set(0)
				return
			case sdl.WINDOWEVENT_CLOSE:
				wm.handlers[t.WindowID](e)
				wm.WindowCount.Set(wm.WindowCount.Get() - 1)
				return
			}

		case *sdl.MouseMotionEvent:
			wm.windows[t.WindowID].Mouse.Position.Set(pair.Int32Pair{L: t.X, R: t.Y})
		case *sdl.MouseButtonEvent:

			mouse := wm.windows[t.WindowID].Mouse
			switch t.Button {
			case sdl.BUTTON_LEFT:
				go mouse.LeftButtonState.Set(t.State == sdl.PRESSED)
			case sdl.BUTTON_RIGHT:
				go mouse.RightButtonState.Set(t.State == sdl.PRESSED)
			}
		case *sdl.QuitEvent:
			return
		}

		cwd := wm.CurrentWindowID.Get()
		if cwd == 0 {
			return
		}

		wm.handlers[cwd](e)
	}
}
<commit_msg>windowManager - check current window ID /before/ dispatching window events<commit_after>package gamekit

import (
	"sync"

	"github.com/sheenobu/go-gamekit/pair"
	"github.com/sheenobu/rxgen/rx"

	"github.com/veandco/go-sdl2/sdl"
)

// The WindowManager is responsible for handling a group of windows with different event handlers
type WindowManager struct {
	WindowCount     *rx.Uint32
	CurrentWindowID *rx.Uint32

	lock     sync.RWMutex
	windows  map[uint32]*Window
	handlers map[uint32]func(e sdl.Event)
}

// NewWindowManager builds a new window manager
func NewWindowManager() *WindowManager {
	return &WindowManager{
		CurrentWindowID: rx.NewUint32(0),
		WindowCount:     rx.NewUint32(0),

		windows:  make(map[uint32]*Window),
		handlers: make(map[uint32]func(e sdl.Event)),
	}
}

// NewWindow creates a new window, returning the window
func (wm *WindowManager) NewWindow(title string, w int, h int, extraFlags int) (*Window, error) {

	wm.lock.Lock()
	defer wm.lock.Unlock()

	window, renderer, err := sdl.CreateWindowAndRenderer(w, h, 0)

	if err != nil {
		return nil, err
	}

	win := &Window{
		ID:       window.GetID(),
		Size:     pair.NewRxInt32Pair(pair.Int32Pair{L: int32(w), R: int32(h)}),
		Window:   window,
		Renderer: renderer,
		Mouse:    NewMouse(),

		wm: wm,
	}

	wm.windows[window.GetID()] = win

	wm.handlers[window.GetID()] = func(e sdl.Event) {
		switch t := e.(type) {
		case *sdl.WindowEvent:
			if t.Event == sdl.WINDOWEVENT_CLOSE {
				window.Destroy()
			}
		}
	}

	wm.WindowCount.Set(wm.WindowCount.Get() + 1)

	return win, nil
}

// Destroy destroys the window
func (wm *WindowManager) Destroy(id uint32) {
	wm.lock.RLock()
	defer wm.lock.RUnlock()

	wm.windows[id].Window.Destroy()
	delete(wm.windows, id)
	delete(wm.handlers, id)
}

// SetHandler sets the event handler for the given window id
func (wm *WindowManager) SetHandler(id uint32, h func(e sdl.Event)) {
	wm.lock.Lock()
	defer wm.lock.Unlock()
	wm.handlers[id] = h
}

// DispatchEvents handles events and sends them to the required window event handlers.
func (wm *WindowManager) DispatchEvents() {
	if e := sdl.PollEvent(); e != nil {
		wm.lock.RLock()
		defer wm.lock.RUnlock()

		cwd := wm.CurrentWindowID.Get()
		if cwd == 0 {
			return
		}

		switch t := e.(type) {
		case *sdl.WindowEvent:
			switch t.Event {
			case sdl.WINDOWEVENT_RESIZED:
				// if we handle these events we actually get a crash????
				// wm.windows[t.WindowID].Size.Set(pair.Int32Pair{t.Data1, t.Data2})
				return
			case sdl.WINDOWEVENT_SIZE_CHANGED:
				wm.windows[t.WindowID].Size.Set(pair.Int32Pair{L: t.Data1, R: t.Data2})
				return
			case sdl.WINDOWEVENT_FOCUS_GAINED:
				wm.CurrentWindowID.Set(t.WindowID)
				return
			case sdl.WINDOWEVENT_FOCUS_LOST:
				wm.CurrentWindowID.Set(0)
				return
			case sdl.WINDOWEVENT_CLOSE:
				wm.handlers[t.WindowID](e)
				wm.WindowCount.Set(wm.WindowCount.Get() - 1)
				return
			}

		case *sdl.MouseMotionEvent:
			wm.windows[t.WindowID].Mouse.Position.Set(pair.Int32Pair{L: t.X, R: t.Y})
		case *sdl.MouseButtonEvent:

			mouse := wm.windows[t.WindowID].Mouse
			switch t.Button {
			case sdl.BUTTON_LEFT:
				go mouse.LeftButtonState.Set(t.State == sdl.PRESSED)
			case sdl.BUTTON_RIGHT:
				go mouse.RightButtonState.Set(t.State == sdl.PRESSED)
			}
		case *sdl.QuitEvent:
			return
		}

		wm.handlers[cwd](e)
	}
}
<|endoftext|>


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

package worker

import (
	"fmt"
	"github.com/Matir/gobuster/client"
	"github.com/Matir/gobuster/logging"
	"github.com/Matir/gobuster/results"
	ss "github.com/Matir/gobuster/settings"
	"github.com/Matir/gobuster/util"
	"github.com/Matir/gobuster/workqueue"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"
)

type Stoppable interface {
	Stop()
}

type PageWorker interface {
	Eligible(*http.Response) bool
	Handle(*url.URL, io.Reader)
}

// Workers do the work of connecting to the server, issuing the request, and
// then optionally parsing the response.  Normally a pool of several workers
// will be used due to network latency.
type Worker struct {
	// client for connections
	client *client.Client
	// Channel for URLs to scan
	src <-chan *url.URL
	// Function to add future work
	adder workqueue.QueueAddFunc
	// Function to mark work done
	done workqueue.QueueDoneFunc
	// Channel for scan results
	rchan chan<- results.Result
	// Settings
	settings *ss.ScanSettings
	// HTML worker to parse page
	pageWorker PageWorker
	// Channel to trigger stopping
	stop chan bool
	// Request for redirection
	redir *http.Request
}

// Construct a worker with given settings.
func NewWorker(settings *ss.ScanSettings,
	factory client.ClientFactory,
	src <-chan *url.URL,
	adder workqueue.QueueAddFunc,
	done workqueue.QueueDoneFunc,
	rchan chan<- results.Result) *Worker {
	w := &Worker{
		client:   factory.Get(),
		settings: settings,
		src:      src,
		adder:    adder,
		done:     done,
		rchan:    rchan,
		stop:     make(chan bool),
	}

	// Install redirect handler
	redirHandler := func(req *http.Request, _ []*http.Request) error {
		w.redir = req
		return fmt.Errorf("Stop redirect.")
	}
	w.client.CheckRedirect = redirHandler

	return w
}

func (w *Worker) SetPageWorker(pw PageWorker) {
	w.pageWorker = pw
}

func (w *Worker) Run() {
	for true {
		select {
		case <-w.stop:
			return
		case task, ok := <-w.src:
			if !ok {
				return
			}
			w.HandleURL(task)
		}
	}
}

func (w *Worker) RunInBackground() {
	go w.Run()
}

func (w *Worker) Stop() {
	w.stop <- true
}

func (w *Worker) HandleURL(task *url.URL) {
	logging.Logf(logging.LogDebug, "Trying Raw URL (unmangled): %s", task.String())
	w.TryURL(task)
	if !util.URLIsDir(task) {
		w.TryMangleURL(task)
		for _, ext := range w.settings.Extensions {
			task := *task
			task.Path += "." + ext
			w.TryURL(&task)
			w.TryMangleURL(&task)
		}
	}
	// Mark as done
	w.done(1)
}

func (w *Worker) TryMangleURL(task *url.URL) {
	if !w.settings.Mangle {
		return
	}
	clone := *task
	spos := strings.LastIndex(clone.Path, "/")
	if spos == -1 {
		return
	}
	dirname := clone.Path[:spos]
	basename := clone.Path[spos+1:]
	for _, newname := range Mangle(basename) {
		clone := clone
		clone.Path = dirname + "/" + newname
		w.TryURL(&clone)
	}
}

func (w *Worker) TryURL(task *url.URL) {
	logging.Logf(logging.LogInfo, "Trying: %s", task.String())
	w.redir = nil
	if resp, err := w.client.RequestURL(task); err != nil && w.redir == nil {
		result := results.Result{URL: task, Error: err}
		if resp != nil {
			result.Code = resp.StatusCode
		}
		w.rchan <- result
	} else {
		defer resp.Body.Close()
		// Do we keep going?
		if util.URLIsDir(task) && w.KeepSpidering(resp.StatusCode) {
			logging.Logf(logging.LogDebug, "Referring %s back for spidering.", task.String())
			w.adder(task)
		}
		if w.redir != nil {
			logging.Logf(logging.LogDebug, "Referring redirect %s back.", w.redir.URL.String())
			w.adder(w.redir.URL)
		}
		if w.pageWorker != nil && w.pageWorker.Eligible(resp) {
			w.pageWorker.Handle(task, resp.Body)
		}
		var redir *url.URL
		if w.redir != nil {
			redir = w.redir.URL
		}
		w.rchan <- results.Result{
			URL:    task,
			Code:   resp.StatusCode,
			Redir:  redir,
			Length: resp.ContentLength,
		}
	}
	if w.settings.SleepTime != 0 {
		time.Sleep(w.settings.SleepTime)
	}
}

// Should we keep spidering from this code?
func (w *Worker) KeepSpidering(code int) bool {
	for _, v := range w.settings.SpiderCodes {
		if code == v {
			return true
		}
	}
	return false
}

// Starts a batch of workers based on the relevant settings.
func StartWorkers(settings *ss.ScanSettings,
	factory client.ClientFactory,
	src <-chan *url.URL,
	adder workqueue.QueueAddFunc,
	done workqueue.QueueDoneFunc,
	rchan chan<- results.Result) []*Worker {
	count := settings.Workers
	workers := make([]*Worker, count)
	for i := 0; i < count; i++ {
		workers[i] = NewWorker(settings, factory, src, adder, done, rchan)
		workers[i].RunInBackground()
	}
	return workers
}

// Set PageWorker for a bunch of Workers
func SetPageWorkers(batch []*Worker, pw PageWorker) {
	for _, w := range batch {
		w.SetPageWorker(pw)
	}
}

// Mangle a basename
func Mangle(basename string) []string {
	mangleRules := []string{
		".%s.swp", // VIM Swap File
		"%s~",     // Backup file
		"%s.bak",  // Backup file
		"%s.orig", //Backup file
	}
	res := make([]string, len(mangleRules))
	for i, rule := range mangleRules {
		res[i] = fmt.Sprintf(rule, basename)
	}
	return res
}
<commit_msg>Only mangle successful URLs.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package worker

import (
	"fmt"
	"github.com/Matir/gobuster/client"
	"github.com/Matir/gobuster/logging"
	"github.com/Matir/gobuster/results"
	ss "github.com/Matir/gobuster/settings"
	"github.com/Matir/gobuster/util"
	"github.com/Matir/gobuster/workqueue"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"
)

type Stoppable interface {
	Stop()
}

type PageWorker interface {
	Eligible(*http.Response) bool
	Handle(*url.URL, io.Reader)
}

// Workers do the work of connecting to the server, issuing the request, and
// then optionally parsing the response.  Normally a pool of several workers
// will be used due to network latency.
type Worker struct {
	// client for connections
	client *client.Client
	// Channel for URLs to scan
	src <-chan *url.URL
	// Function to add future work
	adder workqueue.QueueAddFunc
	// Function to mark work done
	done workqueue.QueueDoneFunc
	// Channel for scan results
	rchan chan<- results.Result
	// Settings
	settings *ss.ScanSettings
	// HTML worker to parse page
	pageWorker PageWorker
	// Channel to trigger stopping
	stop chan bool
	// Request for redirection
	redir *http.Request
}

// Construct a worker with given settings.
func NewWorker(settings *ss.ScanSettings,
	factory client.ClientFactory,
	src <-chan *url.URL,
	adder workqueue.QueueAddFunc,
	done workqueue.QueueDoneFunc,
	rchan chan<- results.Result) *Worker {
	w := &Worker{
		client:   factory.Get(),
		settings: settings,
		src:      src,
		adder:    adder,
		done:     done,
		rchan:    rchan,
		stop:     make(chan bool),
	}

	// Install redirect handler
	redirHandler := func(req *http.Request, _ []*http.Request) error {
		w.redir = req
		return fmt.Errorf("Stop redirect.")
	}
	w.client.CheckRedirect = redirHandler

	return w
}

func (w *Worker) SetPageWorker(pw PageWorker) {
	w.pageWorker = pw
}

func (w *Worker) Run() {
	for true {
		select {
		case <-w.stop:
			return
		case task, ok := <-w.src:
			if !ok {
				return
			}
			w.HandleURL(task)
		}
	}
}

func (w *Worker) RunInBackground() {
	go w.Run()
}

func (w *Worker) Stop() {
	w.stop <- true
}

func (w *Worker) HandleURL(task *url.URL) {
	logging.Logf(logging.LogDebug, "Trying Raw URL (unmangled): %s", task.String())
	withMangle := w.TryURL(task)
	if !util.URLIsDir(task) {
		if withMangle {
			w.TryMangleURL(task)
		}
		for _, ext := range w.settings.Extensions {
			task := *task
			task.Path += "." + ext
			if w.TryURL(&task) {
				w.TryMangleURL(&task)
			}
		}
	}
	// Mark as done
	w.done(1)
}

func (w *Worker) TryMangleURL(task *url.URL) {
	if !w.settings.Mangle {
		return
	}
	clone := *task
	spos := strings.LastIndex(clone.Path, "/")
	if spos == -1 {
		return
	}
	dirname := clone.Path[:spos]
	basename := clone.Path[spos+1:]
	for _, newname := range Mangle(basename) {
		clone := clone
		clone.Path = dirname + "/" + newname
		w.TryURL(&clone)
	}
}

func (w *Worker) TryURL(task *url.URL) bool {
	logging.Logf(logging.LogInfo, "Trying: %s", task.String())
	tryMangle := false
	w.redir = nil
	if resp, err := w.client.RequestURL(task); err != nil && w.redir == nil {
		result := results.Result{URL: task, Error: err}
		if resp != nil {
			result.Code = resp.StatusCode
		}
		w.rchan <- result
	} else {
		defer resp.Body.Close()
		// Do we keep going?
		if util.URLIsDir(task) && w.KeepSpidering(resp.StatusCode) {
			logging.Logf(logging.LogDebug, "Referring %s back for spidering.", task.String())
			w.adder(task)
		}
		if w.redir != nil {
			logging.Logf(logging.LogDebug, "Referring redirect %s back.", w.redir.URL.String())
			w.adder(w.redir.URL)
		}
		if w.pageWorker != nil && w.pageWorker.Eligible(resp) {
			w.pageWorker.Handle(task, resp.Body)
		}
		var redir *url.URL
		if w.redir != nil {
			redir = w.redir.URL
		}
		w.rchan <- results.Result{
			URL:    task,
			Code:   resp.StatusCode,
			Redir:  redir,
			Length: resp.ContentLength,
		}
		tryMangle = w.KeepSpidering(resp.StatusCode)
	}
	if w.settings.SleepTime != 0 {
		time.Sleep(w.settings.SleepTime)
	}
	return tryMangle
}

// Should we keep spidering from this code?
func (w *Worker) KeepSpidering(code int) bool {
	for _, v := range w.settings.SpiderCodes {
		if code == v {
			return true
		}
	}
	return false
}

// Starts a batch of workers based on the relevant settings.
func StartWorkers(settings *ss.ScanSettings,
	factory client.ClientFactory,
	src <-chan *url.URL,
	adder workqueue.QueueAddFunc,
	done workqueue.QueueDoneFunc,
	rchan chan<- results.Result) []*Worker {
	count := settings.Workers
	workers := make([]*Worker, count)
	for i := 0; i < count; i++ {
		workers[i] = NewWorker(settings, factory, src, adder, done, rchan)
		workers[i].RunInBackground()
	}
	return workers
}

// Set PageWorker for a bunch of Workers
func SetPageWorkers(batch []*Worker, pw PageWorker) {
	for _, w := range batch {
		w.SetPageWorker(pw)
	}
}

// Mangle a basename
func Mangle(basename string) []string {
	mangleRules := []string{
		".%s.swp", // VIM Swap File
		"%s~",     // Backup file
		"%s.bak",  // Backup file
		"%s.orig", //Backup file
	}
	res := make([]string, len(mangleRules))
	for i, rule := range mangleRules {
		res[i] = fmt.Sprintf(rule, basename)
	}
	return res
}
<|endoftext|>


<commit_before>// Copyright 2011 Xing Xing <mikespook@gmail.com> All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package worker

import (
    "time"
    "bytes"
    "bitbucket.org/mikespook/gearman-go/common"
)

const (
    Unlimited = 0
    OneByOne = 1

    Immediately = 0
)

var (
    ErrConnection = common.ErrConnection
)
// Job handler
type JobHandler func(*Job) error

type JobFunc func(*Job) ([]byte, error)

// The definition of the callback function.
type jobFunc struct {
    f JobFunc
    timeout uint32
}

// Map for added function.
type JobFuncs map[string]*jobFunc

/*
Worker side api for gearman

usage:
w = worker.New(worker.Unlimited)
w.AddFunction("foobar", foobar)
w.AddServer("127.0.0.1:4730")
w.Work() // Enter the worker's main loop

The definition of the callback function 'foobar' should suit for the type 'JobFunction'.
It looks like this:

func foobar(job *Job) (data []byte, err os.Error) {
    //sth. here
    //plaplapla...
    return
}
*/
type Worker struct {
    agents  []*agent
    funcs   JobFuncs
    in chan *Job
    running bool
    limit chan bool

    Id string
    // assign a ErrFunc to handle errors
    ErrHandler common.ErrorHandler
    JobHandler JobHandler
}

// Get a new worker
func New(l int) (worker *Worker) {
    worker = &Worker{
        agents: make([]*agent, 0),
        funcs: make(JobFuncs),
        in:  make(chan *Job, common.QUEUE_SIZE),
    }
    if l != Unlimited {
        worker.limit = make(chan bool, l)
        for i := 0; i < l; i ++ {
            worker.limit <- true
        }
    }
    return
}

//
func (worker *Worker)err(e error) {
    if worker.ErrHandler != nil {
        worker.ErrHandler(e)
    }
}

// Add a server. The addr should be 'host:port' format.
// The connection is established at this time.
func (worker *Worker) AddServer(addr string) (err error) {
    // Create a new job server's client as a agent of server
    server, err := newAgent(addr, worker)
    if err != nil {
        return err
    }
    worker.agents = append(worker.agents, server)
    return
}

// Write a job to job server.
// Here, the job's mean is not the oraginal mean.
// Just looks like a network package for job's result or tell job server, there was a fail.
func (worker *Worker) broadcast(job *Job) {
    for _, v := range worker.agents {
        v.WriteJob(job)
    }
}

// Add a function.
// Plz added job servers first, then functions.
// The API will tell every connected job server that 'I can do this'
func (worker *Worker) AddFunc(funcname string,
f JobFunc, timeout uint32) (err error) {
    if _, ok := worker.funcs[funcname]; ok {
        return common.Errorf("The function already exists: %s", funcname)
    }
    worker.funcs[funcname] = &jobFunc{f: f, timeout: timeout}

    if worker.running {
        worker.addFunc(funcname, timeout)
    }
    return
}

// inner add function
func (worker *Worker) addFunc(funcname string, timeout uint32) {
    var datatype uint32
    var data []byte
    if timeout == 0 {
        datatype = common.CAN_DO
        data = []byte(funcname)
    } else {
        datatype = common.CAN_DO_TIMEOUT
        data = []byte(funcname + "\x00")
        t := common.Uint32ToBytes(timeout)
        data = append(data, t[:]...)
    }
    job := newJob(common.REQ, datatype, data)
    worker.broadcast(job)

}

// Remove a function.
// Tell job servers 'I can not do this now' at the same time.
func (worker *Worker) RemoveFunc(funcname string) (err error) {
    if _, ok := worker.funcs[funcname]; !ok {
        return common.Errorf("The function does not exist: %s", funcname)
    }
    delete(worker.funcs, funcname)
    if worker.running {
        worker.removeFunc(funcname)
    }
    return
}

// inner remove function
func (worker *Worker) removeFunc(funcname string) {
    job := newJob(common.REQ, common.CANT_DO, []byte(funcname))
    worker.broadcast(job)
}

// Main loop
func (worker *Worker) Work() {
    defer func() {
        worker.running = false
        for _, v := range worker.agents {
            v.Close()
        }
    }()
    for funcname, f := range worker.funcs {
        worker.addFunc(funcname, f.timeout)
    }
    worker.running = true
    for _, v := range worker.agents {
        go v.Work()
    }
    ok := true
    var job *Job
    for ok {
        if job, ok = <-worker.in; ok {
            go func() {
                defer job.Close()
                switch job.DataType {
                case common.ERROR:
                    _, err := common.GetError(job.Data)
                    worker.err(err)
                case common.JOB_ASSIGN, common.JOB_ASSIGN_UNIQ:
                    if err := worker.exec(job); err != nil {
                        worker.err(err)
                    }
                default:
                    worker.handleJob(job)
                }
            }()
        }
    }
}

// job handler
func (worker *Worker) handleJob(job *Job) {
    if worker.JobHandler != nil {
        if err := worker.JobHandler(job); err != nil {
            worker.err(err)
        }
    }
}

// Close.
func (worker *Worker) Close() {
    close(worker.in)
    if worker.limit != nil {
        close(worker.limit)
    }
}

// Send a something out, get the samething back.
func (worker *Worker) Echo(data []byte) {
    job := newJob(common.REQ, common.ECHO_REQ, data)
    worker.broadcast(job)
}

// Remove all of functions.
// Both from the worker or job servers.
func (worker *Worker) Reset() {
    job := newJob(common.REQ, common.RESET_ABILITIES, nil)
    worker.broadcast(job)
    worker.funcs = make(JobFuncs)
}

// Set the worker's unique id.
func (worker *Worker) SetId(id string) {
    worker.Id = id
    job := newJob(common.REQ, common.SET_CLIENT_ID, []byte(id))
    worker.broadcast(job)
}

// Execute the job. And send back the result.
func (worker *Worker) exec(job *Job) (err error) {
    defer func() {
        if r := recover(); r != nil {
            if e, ok := r.(error); ok {
                err = e
            } else {
                err = common.ErrUnknown
            }
        }
    } ()
    if worker.limit != nil {
        <-worker.limit
        defer func() {
            worker.limit <- true
        }()
    }
    var limit int
    if job.DataType == common.JOB_ASSIGN {
        limit = 3
    } else {
        limit = 4
    }
    jobdata := bytes.SplitN(job.Data, []byte{'\x00'}, limit)
    job.Handle = string(jobdata[0])
    funcname := string(jobdata[1])
    if job.DataType == common.JOB_ASSIGN {
        job.Data = jobdata[2]
    } else {
        job.UniqueId = string(jobdata[2])
        job.Data = jobdata[3]
    }
    f, ok := worker.funcs[funcname]
    if !ok {
        return common.Errorf("The function does not exist: %s", funcname)
    }
    var r *result
    if f.timeout == 0 {
        d, e := f.f(job)
        r = &result{data:d, err: e}
    } else {
        r = execTimeout(f.f, job, time.Duration(f.timeout) * time.Second)
    }
    var datatype uint32
    if r.err == nil {
        datatype = common.WORK_COMPLETE
    } else {
        if r.data == nil {
            datatype = common.WORK_FAIL
        } else {
            datatype = common.WORK_EXCEPTION
        }
        err = r.err
    }

    job.magicCode = common.REQ
    job.DataType = datatype
    job.Data = r.data
    job.agent.WriteJob(job)
    return
}

func (worker *Worker) removeAgent(a *agent) {
    for k, v := range worker.agents {
        if v == a {
            worker.agents = append(worker.agents[:k], worker.agents[k + 1:] ...)
        }
    }
    if len(worker.agents) == 0 {
        worker.err(common.ErrNoActiveAgent)
    }
}

type result struct {
    data []byte
    err error
}

func execTimeout(f JobFunc, job *Job, timeout time.Duration) (r *result) {
    rslt := make(chan *result)
    defer close(rslt)
    go func() {
        defer func() {recover()}()
        d, e := f(job)
        rslt <- &result{data: d, err: e}
    }()
    select {
    case r = <-rslt:
    case <-time.After(timeout):
        go job.cancel()
        return &result{err:common.ErrExecTimeOut}
    }
    return r
}
<commit_msg>a wild pointer captrued<commit_after>// Copyright 2011 Xing Xing <mikespook@gmail.com> All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package worker

import (
    "time"
    "bytes"
    "bitbucket.org/mikespook/gearman-go/common"
)

const (
    Unlimited = 0
    OneByOne = 1

    Immediately = 0
)

var (
    ErrConnection = common.ErrConnection
)
// Job handler
type JobHandler func(*Job) error

type JobFunc func(*Job) ([]byte, error)

// The definition of the callback function.
type jobFunc struct {
    f JobFunc
    timeout uint32
}

// Map for added function.
type JobFuncs map[string]*jobFunc

/*
Worker side api for gearman

usage:
w = worker.New(worker.Unlimited)
w.AddFunction("foobar", foobar)
w.AddServer("127.0.0.1:4730")
w.Work() // Enter the worker's main loop

The definition of the callback function 'foobar' should suit for the type 'JobFunction'.
It looks like this:

func foobar(job *Job) (data []byte, err os.Error) {
    //sth. here
    //plaplapla...
    return
}
*/
type Worker struct {
    agents  []*agent
    funcs   JobFuncs
    in chan *Job
    running bool
    limit chan bool

    Id string
    // assign a ErrFunc to handle errors
    ErrHandler common.ErrorHandler
    JobHandler JobHandler
}

// Get a new worker
func New(l int) (worker *Worker) {
    worker = &Worker{
        agents: make([]*agent, 0),
        funcs: make(JobFuncs),
        in:  make(chan *Job, common.QUEUE_SIZE),
    }
    if l != Unlimited {
        worker.limit = make(chan bool, l)
        for i := 0; i < l; i ++ {
            worker.limit <- true
        }
    }
    return
}

//
func (worker *Worker)err(e error) {
    if worker.ErrHandler != nil {
        worker.ErrHandler(e)
    }
}

// Add a server. The addr should be 'host:port' format.
// The connection is established at this time.
func (worker *Worker) AddServer(addr string) (err error) {
    // Create a new job server's client as a agent of server
    server, err := newAgent(addr, worker)
    if err != nil {
        return err
    }
    worker.agents = append(worker.agents, server)
    return
}

// Write a job to job server.
// Here, the job's mean is not the oraginal mean.
// Just looks like a network package for job's result or tell job server, there was a fail.
func (worker *Worker) broadcast(job *Job) {
    for _, v := range worker.agents {
        v.WriteJob(job)
    }
}

// Add a function.
// Plz added job servers first, then functions.
// The API will tell every connected job server that 'I can do this'
func (worker *Worker) AddFunc(funcname string,
f JobFunc, timeout uint32) (err error) {
    if _, ok := worker.funcs[funcname]; ok {
        return common.Errorf("The function already exists: %s", funcname)
    }
    worker.funcs[funcname] = &jobFunc{f: f, timeout: timeout}

    if worker.running {
        worker.addFunc(funcname, timeout)
    }
    return
}

// inner add function
func (worker *Worker) addFunc(funcname string, timeout uint32) {
    var datatype uint32
    var data []byte
    if timeout == 0 {
        datatype = common.CAN_DO
        data = []byte(funcname)
    } else {
        datatype = common.CAN_DO_TIMEOUT
        data = []byte(funcname + "\x00")
        t := common.Uint32ToBytes(timeout)
        data = append(data, t[:]...)
    }
    job := newJob(common.REQ, datatype, data)
    worker.broadcast(job)

}

// Remove a function.
// Tell job servers 'I can not do this now' at the same time.
func (worker *Worker) RemoveFunc(funcname string) (err error) {
    if _, ok := worker.funcs[funcname]; !ok {
        return common.Errorf("The function does not exist: %s", funcname)
    }
    delete(worker.funcs, funcname)
    if worker.running {
        worker.removeFunc(funcname)
    }
    return
}

// inner remove function
func (worker *Worker) removeFunc(funcname string) {
    job := newJob(common.REQ, common.CANT_DO, []byte(funcname))
    worker.broadcast(job)
}

// Main loop
func (worker *Worker) Work() {
    defer func() {
        worker.running = false
        for _, v := range worker.agents {
            v.Close()
        }
    }()
    for funcname, f := range worker.funcs {
        worker.addFunc(funcname, f.timeout)
    }
    worker.running = true
    for _, v := range worker.agents {
        go v.Work()
    }
    ok := true
    for ok {
        var job *Job
        if job, ok = <-worker.in; ok {
            go func() {
                defer job.Close()
                switch job.DataType {
                case common.ERROR:
                    _, err := common.GetError(job.Data)
                    worker.err(err)
                case common.JOB_ASSIGN, common.JOB_ASSIGN_UNIQ:
                    if err := worker.exec(job); err != nil {
                        worker.err(err)
                    }
                default:
                    worker.handleJob(job)
                }
            }()
        }
    }
}

// job handler
func (worker *Worker) handleJob(job *Job) {
    if worker.JobHandler != nil {
        if err := worker.JobHandler(job); err != nil {
            worker.err(err)
        }
    }
}

// Close.
func (worker *Worker) Close() {
    close(worker.in)
    if worker.limit != nil {
        close(worker.limit)
    }
}

// Send a something out, get the samething back.
func (worker *Worker) Echo(data []byte) {
    job := newJob(common.REQ, common.ECHO_REQ, data)
    worker.broadcast(job)
}

// Remove all of functions.
// Both from the worker or job servers.
func (worker *Worker) Reset() {
    job := newJob(common.REQ, common.RESET_ABILITIES, nil)
    worker.broadcast(job)
    worker.funcs = make(JobFuncs)
}

// Set the worker's unique id.
func (worker *Worker) SetId(id string) {
    worker.Id = id
    job := newJob(common.REQ, common.SET_CLIENT_ID, []byte(id))
    worker.broadcast(job)
}

// Execute the job. And send back the result.
func (worker *Worker) exec(job *Job) (err error) {
    defer func() {
        if r := recover(); r != nil {
            if e, ok := r.(error); ok {
                err = e
            } else {
                err = common.ErrUnknown
            }
        }
    } ()
    if worker.limit != nil {
        <-worker.limit
        defer func() {
            worker.limit <- true
        }()
    }
    var limit int
    if job.DataType == common.JOB_ASSIGN {
        limit = 3
    } else {
        limit = 4
    }
    jobdata := bytes.SplitN(job.Data, []byte{'\x00'}, limit)
    job.Handle = string(jobdata[0])
    funcname := string(jobdata[1])
    if job.DataType == common.JOB_ASSIGN {
        job.Data = jobdata[2]
    } else {
        job.UniqueId = string(jobdata[2])
        job.Data = jobdata[3]
    }
    f, ok := worker.funcs[funcname]
    if !ok {
        return common.Errorf("The function does not exist: %s", funcname)
    }
    var r *result
    if f.timeout == 0 {
        d, e := f.f(job)
        r = &result{data:d, err: e}
    } else {
        r = execTimeout(f.f, job, time.Duration(f.timeout) * time.Second)
    }
    var datatype uint32
    if r.err == nil {
        datatype = common.WORK_COMPLETE
    } else {
        if r.data == nil {
            datatype = common.WORK_FAIL
        } else {
            datatype = common.WORK_EXCEPTION
        }
        err = r.err
    }

    job.magicCode = common.REQ
    job.DataType = datatype
    job.Data = r.data
    job.agent.WriteJob(job)
    return
}

func (worker *Worker) removeAgent(a *agent) {
    for k, v := range worker.agents {
        if v == a {
            worker.agents = append(worker.agents[:k], worker.agents[k + 1:] ...)
        }
    }
    if len(worker.agents) == 0 {
        worker.err(common.ErrNoActiveAgent)
    }
}

type result struct {
    data []byte
    err error
}

func execTimeout(f JobFunc, job *Job, timeout time.Duration) (r *result) {
    rslt := make(chan *result)
    defer close(rslt)
    go func() {
        defer func() {recover()}()
        d, e := f(job)
        rslt <- &result{data: d, err: e}
    }()
    select {
    case r = <-rslt:
    case <-time.After(timeout):
        go job.cancel()
        return &result{err:common.ErrExecTimeOut}
    }
    return r
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/metakeule/refl"
	"reflect"
)

type Rating struct {
	NumParticipants int
	Points          int
}

type Game struct {
	Rating
	Name string
}

func (ø *Game) Play(players int) {
	p("playing " + ø.Name + " with " + fmt.Sprintf("%d", players) + " players")
}

func p(i interface{}) {
	fmt.Println(i)
}

type A1 struct {
	Inner B1
}

type A2 struct {
	Inner B2
}

type B1 struct{}
type B2 struct{}

func SetInner(a interface{}, b interface{}) {
	refl.PolySetField(a, "Inner", b)
}

type WithFunc struct {
	Fn func(*B1)
}

func main() {
	g := Game{Rating{2, 300}, "Game of life"}
	p(refl.Type(g))              // => "Game"
	p(refl.Kind(g))              // => "struct"
	refl.Call(&g, "Play", 2)     // => "playing Game of life with 2 players"
	p(refl.GetField(&g, "Name")) // "Game of life"
	refl.SetField(&g.Rating, "Points", 400)
	p(g.Points) // 400
	refl.Set(&g.Rating, Rating{5, 800})
	p(g.Points) // 800

	a1, a2 := &A1{}, &A2{}
	SetInner(a1, B1{})
	SetInner(a2, B2{})
	fmt.Println(refl.Inspect(a1)) // &main.A1{Inner:main.B1{}} ()
	fmt.Println(refl.Inspect(a2)) // &main.A2{Inner:main.B2{}} ()

	w := WithFunc{func(b *B1) { refl.P(b) }}
	f := refl.GetField(&w, "Fn").(func(*B1))
	refl.CallFn(f, &B1{})
	f(&B1{})
}
<commit_msg>fix import<commit_after>package main

import (
	"fmt"
	"github.com/metakeule/refl"
)

type Rating struct {
	NumParticipants int
	Points          int
}

type Game struct {
	Rating
	Name string
}

func (ø *Game) Play(players int) {
	p("playing " + ø.Name + " with " + fmt.Sprintf("%d", players) + " players")
}

func p(i interface{}) {
	fmt.Println(i)
}

type A1 struct {
	Inner B1
}

type A2 struct {
	Inner B2
}

type B1 struct{}
type B2 struct{}

func SetInner(a interface{}, b interface{}) {
	refl.PolySetField(a, "Inner", b)
}

type WithFunc struct {
	Fn func(*B1)
}

func main() {
	g := Game{Rating{2, 300}, "Game of life"}
	p(refl.Type(g))              // => "Game"
	p(refl.Kind(g))              // => "struct"
	refl.Call(&g, "Play", 2)     // => "playing Game of life with 2 players"
	p(refl.GetField(&g, "Name")) // "Game of life"
	refl.SetField(&g.Rating, "Points", 400)
	p(g.Points) // 400
	refl.Set(&g.Rating, Rating{5, 800})
	p(g.Points) // 800

	a1, a2 := &A1{}, &A2{}
	SetInner(a1, B1{})
	SetInner(a2, B2{})
	fmt.Println(refl.Inspect(a1)) // &main.A1{Inner:main.B1{}} ()
	fmt.Println(refl.Inspect(a2)) // &main.A2{Inner:main.B2{}} ()

	w := WithFunc{func(b *B1) { refl.P(b) }}
	f := refl.GetField(&w, "Fn").(func(*B1))
	refl.CallFn(f, &B1{})
	f(&B1{})
}
<|endoftext|>


<commit_before>package worker

import (
	"log"
	"sync"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/sqs"
	"fmt"
)

// HandlerFunc is used to define the Handler that is run on for each message
type HandlerFunc func(msg *sqs.Message) error

func (f HandlerFunc) HandleMessage(msg *sqs.Message) error {
	return f(msg)
}

// Handler interface
type Handler interface {
	HandleMessage(msg *sqs.Message) error
}

type InvalidEventError struct {
	event string
	msg string
}
func (e *InvalidEventError) Error() string {
	return fmt.Sprintf("[Invalid Event: %s] %s", e.event, e.msg)
}

func NewInvalidEventError(event, msg string) *InvalidEventError {
	return &InvalidEventError{event: event, msg: msg}
}

// Exported Variables
var (
	// what is the queue url we are connecting to, Defaults to empty
	QueueURL string = ""
	// The maximum number of messages to return. Amazon SQS never returns more messages
	// than this value (however, fewer messages might be returned). Valid values
	// are 1 to 10. Default is 10.
	MaxNumberOfMessage int64 = 10
	// The duration (in seconds) for which the call waits for a message to arrive
	// in the queue before returning. If a message is available, the call returns
	// sooner than WaitTimeSeconds.
	WaitTimeSecond int64 = 20
)

// Start starts the polling and will continue polling till the application is forcibly stopped
func Start(svc *sqs.SQS, h Handler) {
	for {
		log.Println("worker: Start polling")
		params := &sqs.ReceiveMessageInput{
			QueueUrl: aws.String(QueueURL), // Required
			MaxNumberOfMessages: aws.Int64(MaxNumberOfMessage),
			MessageAttributeNames: []*string{
				aws.String("All"), // Required
			},
			WaitTimeSeconds:         aws.Int64(WaitTimeSecond),
		}

		resp, err := svc.ReceiveMessage(params)
		if err != nil {
			log.Println(err)
			continue
		}
		if len(resp.Messages) > 0 {
			run(svc, h, resp.Messages)
		}
	}
}

// poll launches goroutine per received message and wait for all message to be processed
func run(svc *sqs.SQS, h Handler, messages []*sqs.Message) {
	numMessages := len(messages)
	log.Printf("worker: Received %d messages", numMessages)

	var wg sync.WaitGroup
	wg.Add(numMessages)
	for i := range messages {
		go func(m *sqs.Message) {
			// launch goroutine
			log.Println("worker: Spawned worker goroutine")
			defer wg.Done()
			if err := handleMessage(svc, m, h); err != nil {
				log.Println(err)
			}
		}(messages[i])
	}

	wg.Wait()
}

func handleMessage(svc *sqs.SQS, m *sqs.Message, h Handler) error {
	var err error
	err = h.HandleMessage(m)
	if _, ok := err.(InvalidEventError); ok {
		log.Println(err.Error())
	} else {
		return err
	}

	params := &sqs.DeleteMessageInput{
		QueueUrl:      aws.String(QueueURL), // Required
		ReceiptHandle: m.ReceiptHandle, // Required
	}
	_, err = svc.DeleteMessage(params)
	if err != nil {
		return err
	}
	log.Printf("worker: deleted message from queue: %s", aws.StringValue(m.ReceiptHandle))

	return nil
}<commit_msg>Don't pass back a pointer.<commit_after>package worker

import (
	"log"
	"sync"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/sqs"
	"fmt"
)

// HandlerFunc is used to define the Handler that is run on for each message
type HandlerFunc func(msg *sqs.Message) error

func (f HandlerFunc) HandleMessage(msg *sqs.Message) error {
	return f(msg)
}

// Handler interface
type Handler interface {
	HandleMessage(msg *sqs.Message) error
}

type InvalidEventError struct {
	event string
	msg string
}
func (e InvalidEventError) Error() string {
	return fmt.Sprintf("[Invalid Event: %s] %s", e.event, e.msg)
}

func NewInvalidEventError(event, msg string) InvalidEventError {
	return InvalidEventError{event: event, msg: msg}
}

// Exported Variables
var (
	// what is the queue url we are connecting to, Defaults to empty
	QueueURL string = ""
	// The maximum number of messages to return. Amazon SQS never returns more messages
	// than this value (however, fewer messages might be returned). Valid values
	// are 1 to 10. Default is 10.
	MaxNumberOfMessage int64 = 10
	// The duration (in seconds) for which the call waits for a message to arrive
	// in the queue before returning. If a message is available, the call returns
	// sooner than WaitTimeSeconds.
	WaitTimeSecond int64 = 20
)

// Start starts the polling and will continue polling till the application is forcibly stopped
func Start(svc *sqs.SQS, h Handler) {
	for {
		log.Println("worker: Start polling")
		params := &sqs.ReceiveMessageInput{
			QueueUrl: aws.String(QueueURL), // Required
			MaxNumberOfMessages: aws.Int64(MaxNumberOfMessage),
			MessageAttributeNames: []*string{
				aws.String("All"), // Required
			},
			WaitTimeSeconds:         aws.Int64(WaitTimeSecond),
		}

		resp, err := svc.ReceiveMessage(params)
		if err != nil {
			log.Println(err)
			continue
		}
		if len(resp.Messages) > 0 {
			run(svc, h, resp.Messages)
		}
	}
}

// poll launches goroutine per received message and wait for all message to be processed
func run(svc *sqs.SQS, h Handler, messages []*sqs.Message) {
	numMessages := len(messages)
	log.Printf("worker: Received %d messages", numMessages)

	var wg sync.WaitGroup
	wg.Add(numMessages)
	for i := range messages {
		go func(m *sqs.Message) {
			// launch goroutine
			log.Println("worker: Spawned worker goroutine")
			defer wg.Done()
			if err := handleMessage(svc, m, h); err != nil {
				log.Println(err)
			}
		}(messages[i])
	}

	wg.Wait()
}

func handleMessage(svc *sqs.SQS, m *sqs.Message, h Handler) error {
	var err error
	err = h.HandleMessage(m)
	if _, ok := err.(InvalidEventError); ok {
		log.Println(err.Error())
	} else {
		return err
	}

	params := &sqs.DeleteMessageInput{
		QueueUrl:      aws.String(QueueURL), // Required
		ReceiptHandle: m.ReceiptHandle, // Required
	}
	_, err = svc.DeleteMessage(params)
	if err != nil {
		return err
	}
	log.Printf("worker: deleted message from queue: %s", aws.StringValue(m.ReceiptHandle))

	return nil
}<|endoftext|>