<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor: Aaron Meihm ameihm@mozilla.com [:alm]
// Contributor: Tristan Weir tweir@mozilla.com [:weir]

// runner-scribe is a mig-runner plugin that processes results coming from automated
// actions and forwards the results as vulnerability events to MozDef
package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"strings"
	"time"
	"log"

	"github.com/mozilla/gozdef"
	"github.com/mozilla/mig"
	scribemod "github.com/mozilla/mig/modules/scribe"
	"gopkg.in/gcfg.v1"
)

const sourceName = "runner-scribe"

// config represents the configuration used by runner-scribe, and is read in on
// initialization
//
// URL is mandatory
type config struct {
	MozDef struct {
		URL      string // URL to post events to MozDef
		UseProxy bool   // A switch to enable/disable the use of a system-configured proxy
	}
	api 		ServiceApi
}

type ServiceApiAsset struct {
	Id 				string `json:"id"`
	AssetType 		string `json:"asset_type"`
	AssetIdentifier string `json:"asset_identifier"`
	Team 			string `json:"team"`
	Operator 		string `json:"operator"`
	Zone 			string `json:"zone"`
	Timestamp 		string `json:"timestamp_utc"`
	Description 	string `json:"description"`
	Score 			int `json:"score"`
}

type ServiceApi struct {
	URL				string
	AuthEndpoint 	string
	ClientID 		string
	ClientSecret	string
	Token 			string // ephemeral token we generate to connect to ServiceAPI
}

type Auth0Token struct {
	AccessToken	string `json:"access_token"`
	Scope			string `json:"scope"`
	ExpiresIn		time.Duration `json:"expires_in"`
	TokenType		string `json:"token_type"`
}

const configPath string = "/etc/mig/runner-scribe.conf"

var conf config

func main() {
	defer func() {
		if e := recover(); e != nil {
			fmt.Fprintf(os.Stderr, "error: %v\n", e)
			os.Exit(1)
		}
	}()

	var (
		err     error
		results mig.RunnerResult
	)

	err = gcfg.ReadFileInto(&conf, configPath)
	if err != nil {
		log.Println(err)
	}

	// generate a realtime auth0 auth token
	conf.api.Token, err = GetAuthToken(conf.api)
	if err != nil {
		log.Println(err)
	}

	// load a searchable map of assets from ServiceAPI
	var serviceApiAssets = make(map[string]ServiceApiAsset)
	err = GetAssets(serviceApiAssets, conf.api)
	if err != nil {
		log.Println(err)
	}

	buf, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		log.Println(err)
	}
	err = json.Unmarshal(buf, &results)
	if err != nil {
		log.Println(err)
	}
	var items []gozdef.VulnEvent
	for _, x := range results.Commands {
		// Process the incoming commands, under normal circumstances we will have one
		// returned command per host. However, this function can handle cases where
		// more than one command applies to a given host. If data for a host already
		// exists in items, makeVulnerability should attempt to append this data to
		// the host rather than add a new item.
		var err error
		items, err = makeVulnerability(items, x, serviceApiAssets)
		if err != nil {
			log.Println(err)
		}
	}
	for _, y := range items {
		y.SourceName = sourceName
		err = sendVulnerability(y)
		if err != nil {
			log.Println(err)
		}
	}
}

func sendVulnerability(item gozdef.VulnEvent) (err error) {
	ac := gozdef.APIConf{
		URL:      conf.MozDef.URL,
		UseProxy: conf.MozDef.UseProxy,
	}
	pub, err := gozdef.InitAPI(ac)
	if err != nil {
		return
	}
	err = pub.Send(item)
	return
}

func makeVulnerability(initems []gozdef.VulnEvent, cmd mig.Command, serviceApiAssets map[string]ServiceApiAsset) (items []gozdef.VulnEvent, err error) {
	var (
		itemptr                       *gozdef.VulnEvent
		assethostname, assetipaddress string
		insertNew                     bool
		assetoperator, assetteam      string
	)
	items = initems

	assethostname = cmd.Agent.Name
	for _, x := range cmd.Agent.Env.Addresses {
		if !strings.Contains(x, ".") {
			continue
		}
		ipt, _, err := net.ParseCIDR(x)
		if err != nil {
			continue
		}
		assetipaddress = ipt.String()
		break
	}

	// First, see if we can locate a preexisting item for this asset
	for i := range items {
		if items[i].Asset.Hostname == assethostname &&
			items[i].Asset.IPAddress == assetipaddress {
			itemptr = &items[i]
			break
		}
	}
	if itemptr == nil {
		// Initialize a new event we will insert later
		newevent, err := gozdef.NewVulnEvent()
		if err != nil {
			return items, err
		}
		newevent.Description = "MIG vulnerability identification"
		newevent.Zone = "mig"
		newevent.Asset.Hostname = assethostname
		newevent.Asset.IPAddress = assetipaddress
		newevent.Asset.OS = cmd.Agent.Env.OS

		assetoperator, assetteam = LookupOperatorTeam(assethostname, serviceApiAssets)
		newevent.Asset.Owner.Operator = assetoperator
		newevent.Asset.Owner.Team = assetteam

		// if we didn't find an operator from ServiceAPI assets
		// set it based on the tag
		if len(cmd.Agent.Tags) != 0 && newevent.Asset.Owner.Operator == "" {
			if _, ok := cmd.Agent.Tags["operator"]; ok {
				newevent.Asset.Owner.Operator = cmd.Agent.Tags["operator"]
			}
		}
		// Apply a v2bkey to the event
		if newevent.Asset.Owner.V2Bkey == "" {
			if newevent.Asset.Owner.Operator != "" {
				newevent.Asset.Owner.V2Bkey = newevent.Asset.Owner.Operator
			}
			if newevent.Asset.Owner.Team != "" {
				newevent.Asset.Owner.V2Bkey += "-" + newevent.Asset.Owner.Team
			}
		}
		// Always set credentialed checks here
		newevent.CredentialedChecks = true
		insertNew = true
		itemptr = &newevent
	}

	for _, result := range cmd.Results {
		var el scribemod.ScribeElements
		err = result.GetElements(&el)
		if err != nil {
			return items, err
		}
		for _, x := range el.Results {
			if !x.MasterResult {
				// Result was false (vulnerability did not match)
				continue
			}
			newve := gozdef.VulnVuln{}
			newve.Name = x.TestName
			for _, y := range x.Tags {
				if y.Key == "severity" {
					newve.Risk = y.Value
				} else if y.Key == "link" {
					newve.Link = y.Value
				}
			}
			// If no risk value is set on the vulnerability, we just treat this as
			// informational and ignore it. This will apply to things like the result
			// from platform dependency checks associated with real vulnerability checks.
			if newve.Risk == "" {
				continue
			}
			newve.Risk = normalizeRisk(newve.Risk)
			newve.LikelihoodIndicator = likelihoodFromRisk(newve.Risk)
			if newve.CVSS == "" {
				newve.CVSS = cvssFromRisk(newve.Risk)
			}
			// Use the identifier for each true subresult in the
			// test as a proof section
			for _, y := range x.Results {
				if y.Result {
					newve.Packages = append(newve.Packages, y.Identifier)
				}
			}
			itemptr.Vuln = append(itemptr.Vuln, newve)
		}
	}
	if insertNew {
		items = append(items, *itemptr)
	}
	return
}

// given config for an API behind Auth0 (including client ID and Secret),
// return an Auth0 access token beginning with "Bearer "
// pattern from https://auth0.com/docs/api-auth/tutorials/client-credentials
func GetAuthToken(api ServiceApi) (string, error) {
	payload := strings.NewReader(fmt.Sprintf(`{
		"grant_type": "client_credentials",
		"client_id": "%s",
		"client_secret": "%s",
		"audience": "%s"
		}`, api.ClientID, api.ClientSecret, api.URL))

	req, _ := http.NewRequest("POST", api.AuthEndpoint, payload)
	req.Header.Add("content-type", "application/json")

	res, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", err
	}

	defer res.Body.Close()
	bodyJSON, _ := ioutil.ReadAll(res.Body)

	// unpack the JSON into an Auth0 token struct
	var body Auth0Token
	err = json.Unmarshal(bodyJSON, &body)
	if err != nil {
		return "", err
	}

	// serviceAPI expects the Access token in the form of "Bearer <token>"
	authToken := "Bearer " + body.AccessToken
	return authToken, err
}

// query a ServiceAPI instance for the set of all assets
// load them into a searchable map, keyed to asset hostname
// the ServiceAPI object must already be loaded with a Bearer token
func GetAssets(m map[string]ServiceApiAsset, api ServiceApi) (err error){

	// get json array of assets from serviceapi
	requestURL := api.URL + "api/v1/assets/"
	req, err := http.NewRequest(http.MethodGet, requestURL, nil)
	if err != nil {
		return err
	}
	req.Header.Add("accept", "application/json")
	req.Header.Add("Authorization", api.Token)
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}

	// unpack the HTTP request response
	defer res.Body.Close()
	body, readErr := ioutil.ReadAll(res.Body)
	if readErr != nil {
		return readErr
	}

	// because of the way that ServiceAPI returns the JSON content,
	// we need to Unmarshal it twice
	var allAssetsJson string
	err  = json.Unmarshal(body, &allAssetsJson)
	if err != nil {
		return err
	}

	// convert json into array of ServiceApiAsset objects
	var allAssets []ServiceApiAsset
	err  = json.Unmarshal([]byte(allAssetsJson), &allAssets)
	if err != nil {
		return err
	}

	// build a searchable map, keyed on AssetIdentifier (which is usually hostname)
	for _, tempAsset := range allAssets {
		m[tempAsset.AssetIdentifier] = tempAsset
	}

	return
}

// return the operator and team for a given hostname, provided they are in the map of
// ServiceApiAssets. If they are not in the map or if the values are not present,
// operator and/or team will return as an empty string ""
func LookupOperatorTeam(hostname string, m map[string]ServiceApiAsset) (operator string, team string) {
	operator = m[hostname].Operator
	team = m[hostname].Team

	return
}

// cvssFromRisk returns a synthesized CVSS score as a string given a risk label
func cvssFromRisk(risk string) string {
	switch risk {
	case "critical":
		return "10.0"
	case "high":
		return "8.0"
	case "medium":
		return "5.0"
	case "low":
		return "2.5"
	}
	return "0.0"
}

// likelihoodFromRisk returns a likelihood indicator value given a risk label
func likelihoodFromRisk(risk string) string {
	switch risk {
	case "high":
		return "high"
	case "medium":
		return "medium"
	case "low":
		return "low"
	case "critical":
		return "maximum"
	}
	return "unknown"
}

// normalizeRisk converts known risk labels into a standardized form, if we can't identify
// the value we just return it as is
func normalizeRisk(in string) string {
	switch strings.ToLower(in) {
	case "high":
		return "high"
	case "medium":
		return "medium"
	case "low":
		return "low"
	case "critical":
		return "critical"
	}
	return in
}
<commit_msg>Resolve more review comments<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor: Aaron Meihm ameihm@mozilla.com [:alm]
// Contributor: Tristan Weir tweir@mozilla.com [:weir]

// runner-scribe is a mig-runner plugin that processes results coming from automated
// actions and forwards the results as vulnerability events to MozDef
package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"strings"
	"time"
	"log"

	"github.com/mozilla/gozdef"
	"github.com/mozilla/mig"
	scribemod "github.com/mozilla/mig/modules/scribe"
	"gopkg.in/gcfg.v1"
)

const sourceName = "runner-scribe"

// config represents the configuration used by runner-scribe, and is read in on
// initialization
//
// URL is mandatory
type config struct {
	MozDef struct {
		URL      string // URL to post events to MozDef
		UseProxy bool   // A switch to enable/disable the use of a system-configured proxy
	}
	api 		ServiceApi
}

type ServiceApiAsset struct {
	Id 				string `json:"id"`
	AssetType 		string `json:"asset_type"`
	AssetIdentifier string `json:"asset_identifier"`
	Team 			string `json:"team"`
	Operator 		string `json:"operator"`
	Zone 			string `json:"zone"`
	Timestamp 		string `json:"timestamp_utc"`
	Description 	string `json:"description"`
	Score 			int `json:"score"`
}

type ServiceApi struct {
	URL				string
	AuthEndpoint 	string
	ClientID 		string
	ClientSecret	string
	Token 			string // ephemeral token we generate to connect to ServiceAPI
}

type Auth0Token struct {
	AccessToken	string `json:"access_token"`
	Scope			string `json:"scope"`
	ExpiresIn		time.Duration `json:"expires_in"`
	TokenType		string `json:"token_type"`
}

const configPath string = "/etc/mig/runner-scribe.conf"

var conf config

func main() {
	defer func() {
		if e := recover(); e != nil {
			fmt.Fprintf(os.Stderr, "error: %v\n", e)
			os.Exit(1)
		}
	}()

	var (
		err     error
		results mig.RunnerResult
	)

	err = gcfg.ReadFileInto(&conf, configPath)
	if err != nil {
		log.Println(err)
	}

	// generate a realtime auth0 auth token
	conf.api.Token, err = GetAuthToken(conf.api)
	if err != nil {
		log.Println(err)
	}

	// load a searchable map of assets from ServiceAPI
	var serviceApiAssets = make(map[string]ServiceApiAsset)
	err = GetAssets(serviceApiAssets, conf.api)
	if err != nil {
		log.Println(err)
	}

	buf, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		log.Println(err)
	}
	err = json.Unmarshal(buf, &results)
	if err != nil {
		log.Println(err)
	}
	var items []gozdef.VulnEvent
	for _, x := range results.Commands {
		// Process the incoming commands, under normal circumstances we will have one
		// returned command per host. However, this function can handle cases where
		// more than one command applies to a given host. If data for a host already
		// exists in items, makeVulnerability should attempt to append this data to
		// the host rather than add a new item.
		var err error
		items, err = makeVulnerability(items, x, serviceApiAssets)
		if err != nil {
			log.Println(err)
		}
	}
	for _, y := range items {
		y.SourceName = sourceName
		err = sendVulnerability(y)
		if err != nil {
			log.Println(err)
		}
	}
}

func sendVulnerability(item gozdef.VulnEvent) (err error) {
	ac := gozdef.APIConf{
		URL:      conf.MozDef.URL,
		UseProxy: conf.MozDef.UseProxy,
	}
	pub, err := gozdef.InitAPI(ac)
	if err != nil {
		return
	}
	err = pub.Send(item)
	return
}

func makeVulnerability(initems []gozdef.VulnEvent, cmd mig.Command, serviceApiAssets map[string]ServiceApiAsset) (items []gozdef.VulnEvent, err error) {
	var (
		itemptr                       *gozdef.VulnEvent
		assethostname, assetipaddress string
		insertNew                     bool
		assetoperator, assetteam      string
	)
	items = initems

	assethostname = cmd.Agent.Name
	for _, x := range cmd.Agent.Env.Addresses {
		if !strings.Contains(x, ".") {
			continue
		}
		ipt, _, err := net.ParseCIDR(x)
		if err != nil {
			continue
		}
		assetipaddress = ipt.String()
		break
	}

	// First, see if we can locate a preexisting item for this asset
	for i := range items {
		if items[i].Asset.Hostname == assethostname &&
			items[i].Asset.IPAddress == assetipaddress {
			itemptr = &items[i]
			break
		}
	}
	if itemptr == nil {
		// Initialize a new event we will insert later
		newevent, err := gozdef.NewVulnEvent()
		if err != nil {
			return items, err
		}
		newevent.Description = "MIG vulnerability identification"
		newevent.Zone = "mig"
		newevent.Asset.Hostname = assethostname
		newevent.Asset.IPAddress = assetipaddress
		newevent.Asset.OS = cmd.Agent.Env.OS

		assetoperator, assetteam = LookupOperatorTeam(assethostname, serviceApiAssets)
		newevent.Asset.Owner.Operator = assetoperator
		newevent.Asset.Owner.Team = assetteam

		// if we didn't find an operator from ServiceAPI assets
		// set it based on the tag
		if len(cmd.Agent.Tags) != 0 && newevent.Asset.Owner.Operator == "" {
			if _, ok := cmd.Agent.Tags["operator"]; ok {
				newevent.Asset.Owner.Operator = cmd.Agent.Tags["operator"]
			}
		}
		// Apply a v2bkey to the event
		if newevent.Asset.Owner.V2Bkey == "" {
			if newevent.Asset.Owner.Operator != "" {
				newevent.Asset.Owner.V2Bkey = newevent.Asset.Owner.Operator
			}
			if newevent.Asset.Owner.Team != "" {
				newevent.Asset.Owner.V2Bkey += "-" + newevent.Asset.Owner.Team
			}
		}
		// Always set credentialed checks here
		newevent.CredentialedChecks = true
		insertNew = true
		itemptr = &newevent
	}

	for _, result := range cmd.Results {
		var el scribemod.ScribeElements
		err = result.GetElements(&el)
		if err != nil {
			return items, err
		}
		for _, x := range el.Results {
			if !x.MasterResult {
				// Result was false (vulnerability did not match)
				continue
			}
			newve := gozdef.VulnVuln{}
			newve.Name = x.TestName
			for _, y := range x.Tags {
				if y.Key == "severity" {
					newve.Risk = y.Value
				} else if y.Key == "link" {
					newve.Link = y.Value
				}
			}
			// If no risk value is set on the vulnerability, we just treat this as
			// informational and ignore it. This will apply to things like the result
			// from platform dependency checks associated with real vulnerability checks.
			if newve.Risk == "" {
				continue
			}
			newve.Risk = normalizeRisk(newve.Risk)
			newve.LikelihoodIndicator = likelihoodFromRisk(newve.Risk)
			if newve.CVSS == "" {
				newve.CVSS = cvssFromRisk(newve.Risk)
			}
			// Use the identifier for each true subresult in the
			// test as a proof section
			for _, y := range x.Results {
				if y.Result {
					newve.Packages = append(newve.Packages, y.Identifier)
				}
			}
			itemptr.Vuln = append(itemptr.Vuln, newve)
		}
	}
	if insertNew {
		items = append(items, *itemptr)
	}
	return
}

// given config for an API behind Auth0 (including client ID and Secret),
// return an Auth0 access token beginning with "Bearer "
// pattern from https://auth0.com/docs/api-auth/tutorials/client-credentials
func GetAuthToken(api ServiceApi) (string, error) {
	payload := strings.NewReader(fmt.Sprintf(`{
		"grant_type": "client_credentials",
		"client_id": "%s",
		"client_secret": "%s",
		"audience": "%s"
		}`, api.ClientID, api.ClientSecret, api.URL))

	req, err := http.NewRequest("POST", api.AuthEndpoint, payload)
	if err != nil {
		return "", err
	}

	req.Header.Add("content-type", "application/json")

	res, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", err
	}

	defer res.Body.Close()
	bodyJSON, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return "", err
	}

	// unpack the JSON into an Auth0 token struct
	var body Auth0Token
	err = json.Unmarshal(bodyJSON, &body)
	if err != nil {
		return "", err
	}

	// serviceAPI expects the Access token in the form of "Bearer <token>"
	authToken := "Bearer " + body.AccessToken
	return authToken, err
}

// query a ServiceAPI instance for the set of all assets
// load them into a searchable map, keyed to asset hostname
// the ServiceAPI object must already be loaded with a Bearer token
func GetAssets(m map[string]ServiceApiAsset, api ServiceApi) (error){

	// get json array of assets from serviceapi
	requestURL := api.URL + "api/v1/assets/"
	req, err := http.NewRequest(http.MethodGet, requestURL, nil)
	if err != nil {
		return err
	}
	req.Header.Add("accept", "application/json")
	req.Header.Add("Authorization", api.Token)
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}

	// unpack the HTTP request response
	defer res.Body.Close()
	body, readErr := ioutil.ReadAll(res.Body)
	if readErr != nil {
		return readErr
	}

	// because of the way that ServiceAPI returns the JSON content,
	// we need to Unmarshal it twice
	var allAssetsJson string
	err  = json.Unmarshal(body, &allAssetsJson)
	if err != nil {
		return err
	}

	// convert json into array of ServiceApiAsset objects
	var allAssets []ServiceApiAsset
	err  = json.Unmarshal([]byte(allAssetsJson), &allAssets)
	if err != nil {
		return err
	}

	// build a searchable map, keyed on AssetIdentifier (which is usually hostname)
	for _, tempAsset := range allAssets {
		m[tempAsset.AssetIdentifier] = tempAsset
	}

	return err
}

// return the operator and team for a given hostname, provided they are in the map of
// ServiceApiAssets. If they are not in the map or if the values are not present,
// operator and/or team will return as an empty string ""
func LookupOperatorTeam(hostname string, m map[string]ServiceApiAsset) (operator string, team string) {
	operator = m[hostname].Operator
	team = m[hostname].Team

	return operator, team
}

// cvssFromRisk returns a synthesized CVSS score as a string given a risk label
func cvssFromRisk(risk string) string {
	switch risk {
	case "critical":
		return "10.0"
	case "high":
		return "8.0"
	case "medium":
		return "5.0"
	case "low":
		return "2.5"
	}
	return "0.0"
}

// likelihoodFromRisk returns a likelihood indicator value given a risk label
func likelihoodFromRisk(risk string) string {
	switch risk {
	case "high":
		return "high"
	case "medium":
		return "medium"
	case "low":
		return "low"
	case "critical":
		return "maximum"
	}
	return "unknown"
}

// normalizeRisk converts known risk labels into a standardized form, if we can't identify
// the value we just return it as is
func normalizeRisk(in string) string {
	switch strings.ToLower(in) {
	case "high":
		return "high"
	case "medium":
		return "medium"
	case "low":
		return "low"
	case "critical":
		return "critical"
	}
	return in
}
<|endoftext|>


<commit_before>package rivet

import (
	"net/http"
)

/**
ResponseWriter 提供状态支持
*/
type ResponseWriter interface {
	http.ResponseWriter
	http.Flusher
	// Status 返回调用 WriteHeader 设定的值, 初始值为 0
	Status() int
	// Size 返回调用 Write 写入的字节数, 初始值为 0
	Size() int
	// Written 返回 Status()!=0 || Size()!=0 的结果
	Written() bool
}

/**
NewResponseWriterFakeFlusher 返回 ResponseWriter 实例, 可能是伪 http.Flusher.
如果 rw 已经实现了 ResponseWriter 接口, 返回 rw.(ResponseWriter).
否则返回 &ResponseWriteFakeFlusher 伪 http.Flusher 实例.
*/
func NewResponseWriterFakeFlusher(rw http.ResponseWriter) ResponseWriter {

	if res, ok := rw.(ResponseWriter); ok {
		return res
	}
	return &ResponseWriteFakeFlusher{rw, 0, 0}
}

/**
ResponseWriteFakeFlusher 实现了 http.ResponseWriter 接口和伪 http.Flusher 接口.
Flush() 是个方法, 是否支持 Flusher 取决于原 http.ResponseWriter 实例.
*/
type ResponseWriteFakeFlusher struct {
	http.ResponseWriter
	status int
	size   int
}

// Flush() 是个伪方法, 是否支持 Flusher 取决于原 http.ResponseWriter 实例.
func (rw *ResponseWriteFakeFlusher) Flush() {
	flusher, ok := rw.ResponseWriter.(http.Flusher)
	if ok {
		flusher.Flush()
	}
}

func (rw *ResponseWriteFakeFlusher) WriteHeader(s int) {
	rw.ResponseWriter.WriteHeader(s)
	rw.status = s
}

func (rw *ResponseWriteFakeFlusher) Write(b []byte) (int, error) {
	if !rw.Written() {
		rw.WriteHeader(http.StatusOK)
	}
	size, err := rw.ResponseWriter.Write(b)
	rw.size += size
	return size, err
}

func (rw *ResponseWriteFakeFlusher) Status() int {
	return rw.status
}

func (rw *ResponseWriteFakeFlusher) Size() int {
	return rw.size
}

func (rw *ResponseWriteFakeFlusher) Written() bool {
	return rw.status != 0 || rw.size != 0
}
<commit_msg>增加注释, 重申 ResponseWriter.Written 的语义.<commit_after>package rivet

import (
	"net/http"
)

/**
ResponseWriter 提供状态支持
*/
type ResponseWriter interface {
	http.ResponseWriter
	http.Flusher
	// Status 返回调用 WriteHeader 设定的值, 初始值为 0
	Status() int
	// Size 返回调用 Write 写入的字节数, 初始值为 0
	Size() int

	/**
	Written 返回是否已经写入了内容.
	包括两种情况, WriteHeader 和 Write. 实现是如何判断的可能有差异.
	*/
	Written() bool
}

/**
NewResponseWriterFakeFlusher 返回 ResponseWriter 实例, 可能是伪 http.Flusher.
如果 rw 已经实现了 ResponseWriter 接口, 返回 rw.(ResponseWriter).
否则返回 &ResponseWriteFakeFlusher 伪 http.Flusher 实例.
*/
func NewResponseWriterFakeFlusher(rw http.ResponseWriter) ResponseWriter {

	if res, ok := rw.(ResponseWriter); ok {
		return res
	}
	return &ResponseWriteFakeFlusher{rw, 0, 0}
}

/**
ResponseWriteFakeFlusher 实现了 http.ResponseWriter 接口和伪 http.Flusher 接口.
Flush() 是个方法, 是否支持 Flusher 取决于原 http.ResponseWriter 实例.
*/
type ResponseWriteFakeFlusher struct {
	http.ResponseWriter
	status int
	size   int
}

// Flush() 是个伪方法, 是否支持 Flusher 取决于原 http.ResponseWriter 实例.
func (rw *ResponseWriteFakeFlusher) Flush() {
	flusher, ok := rw.ResponseWriter.(http.Flusher)
	if ok {
		flusher.Flush()
	}
}

/**
WriteHeader 向相应发送状态码 s.
*/
func (rw *ResponseWriteFakeFlusher) WriteHeader(s int) {
	rw.ResponseWriter.WriteHeader(s)
	rw.status = s
}

/**
Write 向相应写入 b, 返回本次写入的字节和发生的错误.
*/
func (rw *ResponseWriteFakeFlusher) Write(b []byte) (int, error) {
	if !rw.Written() {
		rw.WriteHeader(http.StatusOK)
	}
	size, err := rw.ResponseWriter.Write(b)
	rw.size += size
	return size, err
}

/**
Status 返回通过 WriteHeader 设置的值.
*/
func (rw *ResponseWriteFakeFlusher) Status() int {
	return rw.status
}

/**
Size 返回通过 Write 的总字节数.
*/
func (rw *ResponseWriteFakeFlusher) Size() int {
	return rw.size
}

/**
Written 返回  Status()!=0 || Size()!=0 的结果
*/
func (rw *ResponseWriteFakeFlusher) Written() bool {
	return rw.status != 0 || rw.size != 0
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"time"
)

func randSeq(n int64, step int64) string {
	ret := ""
	fmt_str := fmt.Sprintf("%%0%dx", step)
	t := time.Now()
	for i := int64(0); i < n; i += step {
		ret += fmt.Sprintf(fmt_str, t)
	}

	return ret[:n]
}

type doc struct {
	Email        string `json:"email,omitempty"`
	Name         string `json:"name,omitempty"`
	Alt_email    string `json:"alt_email,omitempty"`
	City         string `json:"city,omitempty"`
	Coins        int64  `json:"coins,omitempty"`
	Achievements int64  `json:"achievements,omitempty"`
	Body         string `json:"body,omitempty"`
}

// Here's a basic interface to get Value.
type get_doc_value interface {
	newValue(i int64) *doc
}

type basic_doc struct {
}

type fixed_size_doc struct {
}

type multi_item_doc struct {
}

// newKey generates a new key with a common prefix and variable suffix based on
// sequential document identifier.
func newKey(i int64) string {
	return fmt.Sprintf("%s-%012x", prefix, i)
}

// newValue generates a new set of field(s) based on sequential document
// identifier. Current time is used in order to randomize document fields.
func (basic basic_doc) newValue(i int64) *doc {
	t := time.Now().UnixNano()

	return &doc{
		Email: fmt.Sprintf("%x@%09x.%d", t, i, i%10),
	}
}

func (basic fixed_size_doc) newValue(i int64) *doc {

	return &doc{
		City: fmt.Sprintf("%s", randSeq(size, size)),
	}
}

func (basic multi_item_doc) newValue(i int64) *doc {
	t := time.Now().UnixNano()

	return &doc{
		Email:        fmt.Sprintf("%012x@%09x.%d", t, i, i%10),
		Name:         fmt.Sprintf("%07x %08x", t, t),
		City:         fmt.Sprintf("%09x", t),
		Coins:        t,
		Achievements: t % prime,
		Body:         fmt.Sprintf("%s", randSeq(size-60, 300)),
	}
}
<commit_msg>Remove unused struct field<commit_after>package main

import (
	"fmt"
	"time"
)

func randSeq(n int64, step int64) string {
	ret := ""
	fmt_str := fmt.Sprintf("%%0%dx", step)
	t := time.Now()
	for i := int64(0); i < n; i += step {
		ret += fmt.Sprintf(fmt_str, t)
	}

	return ret[:n]
}

type doc struct {
	Email        string `json:"email,omitempty"`
	Name         string `json:"name,omitempty"`
	City         string `json:"city,omitempty"`
	Coins        int64  `json:"coins,omitempty"`
	Achievements int64  `json:"achievements,omitempty"`
	Body         string `json:"body,omitempty"`
}

// Here's a basic interface to get Value.
type get_doc_value interface {
	newValue(i int64) *doc
}

type basic_doc struct {
}

type fixed_size_doc struct {
}

type multi_item_doc struct {
}

// newKey generates a new key with a common prefix and variable suffix based on
// sequential document identifier.
func newKey(i int64) string {
	return fmt.Sprintf("%s-%012x", prefix, i)
}

// newValue generates a new set of field(s) based on sequential document
// identifier. Current time is used in order to randomize document fields.
func (basic basic_doc) newValue(i int64) *doc {
	t := time.Now().UnixNano()

	return &doc{
		Email: fmt.Sprintf("%x@%09x.%d", t, i, i%10),
	}
}

func (basic fixed_size_doc) newValue(i int64) *doc {
	return &doc{
		City: fmt.Sprintf("%s", randSeq(size, size)),
	}
}

func (basic multi_item_doc) newValue(i int64) *doc {
	t := time.Now().UnixNano()

	return &doc{
		Email:        fmt.Sprintf("%012x@%09x.%d", t, i, i%10),
		Name:         fmt.Sprintf("%07x %08x", t, t),
		City:         fmt.Sprintf("%09x", t),
		Coins:        t,
		Achievements: t % prime,
		Body:         fmt.Sprintf("%s", randSeq(size-60, 300)),
	}
}
<|endoftext|>


<commit_before>package gontpd

import (
	"encoding/binary"
	"log"
	"net"
	"time"
)

type Worker struct {
	conn  *net.UDPConn
	cfg   *Config
	state []byte
}

func (w *Worker) run(i uint) {
	log.Printf("worker:%d online", i)
	var (
		buf    []byte
		n      int
		remote *net.UDPAddr
		err    error

		rcvTime   time.Time
		referTime uint64
	)

	buf = make([]byte, 48, 48)
	// BCE
	_ = buf[47]
	_ = w.state[originTimeStamp-1]

	for {
		n, remote, err = w.conn.ReadFromUDP(buf)
		rcvTime = time.Now()
		if err != nil {
			continue
		}
		if n < 48 {
			continue
		}
		if !isValidNTPRequest(buf) {
			continue
		}
		if w.cfg.InACL(remote.IP) {
			continue
		}
		referTime = binary.BigEndian.Uint64(buf[transmitTimeStamp:])
		copy(buf, w.state)
		binary.BigEndian.PutUint64(buf[originTimeStamp:], referTime)
		binary.BigEndian.PutUint64(buf[receiveTimeStamp:], toNtpTime(rcvTime))
		binary.BigEndian.PutUint64(buf[transmitTimeStamp:], toNtpTime(time.Now()))
		_, err = w.conn.WriteToUDP(buf, remote)
		if err != nil {
			log.Println(err)
		}
	}
}

func isValidNTPRequest(p []byte) (r bool) {
	switch p[0] &^ 0xf8 {
	case 3: // modeClient
	default:
		return
	}

	return true
}

var ntpEpoch = time.Date(1900, 1, 1, 0, 0, 0, 0, time.UTC)

func toNtpTime(t time.Time) uint64 {
	nsec := uint64(t.Sub(ntpEpoch))
	sec := nsec / nanoPerSec
	frac := (nsec - sec*nanoPerSec) << 32 / nanoPerSec
	return uint64(sec<<32 | frac)
}
<commit_msg>worker drop port 0<commit_after>package gontpd

import (
	"encoding/binary"
	"log"
	"net"
	"time"
)

type Worker struct {
	conn  *net.UDPConn
	cfg   *Config
	state []byte
}

func (w *Worker) run(i uint) {
	log.Printf("worker:%d online", i)
	var (
		buf    []byte
		n      int
		remote *net.UDPAddr
		err    error

		rcvTime   time.Time
		referTime uint64
	)

	buf = make([]byte, 48, 48)
	// BCE
	_ = buf[47]
	_ = w.state[originTimeStamp-1]

	for {
		n, remote, err = w.conn.ReadFromUDP(buf)
		rcvTime = time.Now()
		if err != nil || remote.Port == 0 {
			continue
		}
		if n < 48 {
			continue
		}
		if !isValidNTPRequest(buf) {
			continue
		}
		if w.cfg.InACL(remote.IP) {
			continue
		}
		referTime = binary.BigEndian.Uint64(buf[transmitTimeStamp:])
		copy(buf, w.state)
		binary.BigEndian.PutUint64(buf[originTimeStamp:], referTime)
		binary.BigEndian.PutUint64(buf[receiveTimeStamp:], toNtpTime(rcvTime))
		binary.BigEndian.PutUint64(buf[transmitTimeStamp:], toNtpTime(time.Now()))
		_, err = w.conn.WriteToUDP(buf, remote)
		if err != nil {
			log.Println(err)
		}
	}
}

func isValidNTPRequest(p []byte) (r bool) {
	switch p[0] &^ 0xf8 {
	case 3: // modeClient
	default:
		return
	}

	return true
}

var ntpEpoch = time.Date(1900, 1, 1, 0, 0, 0, 0, time.UTC)

func toNtpTime(t time.Time) uint64 {
	nsec := uint64(t.Sub(ntpEpoch))
	sec := nsec / nanoPerSec
	frac := (nsec - sec*nanoPerSec) << 32 / nanoPerSec
	return uint64(sec<<32 | frac)
}
<|endoftext|>


<commit_before>// This worker does not model the way qless does it.
// I more or less modeled it after my own needs.
package goqless

import (
	// "encoding/json"
	"fmt"
	// "github.com/garyburd/redigo/redis"
	"errors"
	"log"
	"reflect"
	"strconv"
	"sync"
	"time"
)

func init() {
	log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile)
}

type JobFunc func(*Job) error
type JobCallback func(*Job) error

type Worker struct {
	Interval int // in time.Duration

	funcs map[string]JobFunc
	queue *Queue
	// events *Events

	cli *Client
}

func NewWorker(cli *Client, queueName string, interval int) *Worker {
	w := &Worker{
		Interval: interval,
		funcs:    make(map[string]JobFunc),
		// events: c.Events(),
		cli: cli,
	}

	w.queue = cli.Queue(queueName)

	return w
}

func heartbeatStart(job *Job, done chan bool, heartbeat int, clientLock sync.Mutex) {
	tick := time.NewTicker(time.Duration(heartbeat) * time.Duration(time.Second))
	for {
		select {
		case <-done:
			tick.Stop()
			return
		case <-tick.C:
			clientLock.Lock()
			success, err := job.Heartbeat()
			clientLock.Unlock()
			if err != nil {
				log.Printf("warning, slow, send heartbeat***jid:%v, queue:%v, success:%v, error:%v",
					job.Jid, job.Queue, success, err)
			} else {
				log.Printf("warning, slow, send heartbeat***jid:%v, queue:%v, success:%v",
					job.Jid, job.Queue, success)
			}
		}
	}
}

func (w *Worker) Start() error {
	// log.Println("worker Start")
	var clientLock sync.Mutex

	heartbeatStr, err := w.cli.GetConfig("heartbeat")
	heartbeat, err := strconv.Atoi(heartbeatStr)
	//log.Println("heartbeatStr:", heartbeat)
	if err != nil {
		heartbeat = 60
		log.Println(err)
	}

	heartbeat /= 2

	err = func(q *Queue) error {
		for {
			clientLock.Lock()
			jobs, err := q.Pop(1) //we may pop more if fast enough
			clientLock.Unlock()

			if err != nil {
				log.Println(err)
				return err
			}

			if len(jobs) == 0 {
				time.Sleep(time.Duration(w.Interval) * time.Millisecond)
				continue
			}

			for i := 0; i < len(jobs); i++ {
				if len(jobs[i].History) > 2 {
					log.Printf("warning, multiple processed exist %+v\n", jobs[i])
				}
				done := make(chan bool)
				//todo: using seprate connection to send heartbeat
				go heartbeatStart(jobs[i], done, heartbeat, clientLock)
				f, ok := w.funcs[jobs[i].Klass]
				if !ok { //we got a job that not belongs to us
					done <- false
					log.Fatalln("got a message not belongs to us, queue", q.Name, jobs[i])
					continue
				}

				err := f(jobs[i])
				if err != nil {
					// TODO: probably do something with this
					log.Println("error: job failed, id", jobs[i].Jid, "queue", jobs[i].Queue, err.Error())
					clientLock.Lock()
					success, err := jobs[i].Fail("fail", err.Error())
					clientLock.Unlock()
					done <- false
					if err != nil {
						log.Printf("fail job:%+v success:%v, error:%v",
							jobs[i].Jid, success, err)
						return err
					}
				} else {
					clientLock.Lock()
					success, err := jobs[i].CompleteWithNoData()
					clientLock.Unlock()
					done <- true
					if err != nil {
						log.Printf("complete job:%+v success:%v, error:%v",
							jobs[i].Jid, success, err)
						return err
					}
					//log.Printf("===job:%+v", jobs[0])
				}
			}
		}
		return nil
	}(w.queue)

	return err
}

func (w *Worker) AddFunc(name string, f JobFunc) error {
	if _, ok := w.funcs[name]; ok {
		return fmt.Errorf("function \"%s\" already exists", name)
	}

	w.funcs[name] = f
	return nil
}

// Adds all the methods in the passed interface as job functions.
// Job names are in the form of: name.methodname
func (w *Worker) AddService(name string, rcvr interface{}) error {
	typ := reflect.TypeOf(rcvr)
	val := reflect.ValueOf(rcvr)
	for i := 0; i < typ.NumMethod(); i++ {
		method := typ.Method(i)
		w.AddFunc(name+"."+method.Name, func(job *Job) error {
			ret := method.Func.Call([]reflect.Value{val, reflect.ValueOf(job)})
			if len(ret) > 0 {
				if err, ok := ret[0].Interface().(error); ok {
					return err
				}
			} else {
				errStr := "reflect len less than zero." + strconv.Itoa(len(ret))
				return errors.New(errStr)
			}

			return nil
		})
	}

	return nil
}
<commit_msg>retry complete the job when complete failed<commit_after>// This worker does not model the way qless does it.
// I more or less modeled it after my own needs.
package goqless

import (
	// "encoding/json"
	"fmt"
	// "github.com/garyburd/redigo/redis"
	"errors"
	"log"
	"reflect"
	"strconv"
	"strings"
	"sync"
	"time"
)

func init() {
	log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile)
}

type JobFunc func(*Job) error
type JobCallback func(*Job) error

type Worker struct {
	sync.Mutex
	Interval int // in time.Duration

	funcs     map[string]JobFunc
	queue     *Queue
	queueAddr string
	queueName string

	cli *Client
}

func NewWorker(queueAddr string, queueName string, interval int) (*Worker, error) {
	ipport := strings.Split(queueAddr, ":")
	client, err := Dial(ipport[0], ipport[1])
	if err != nil {
		log.Println("Dial err:", err)
		return nil, errors.New(err.Error())
	}

	w := &Worker{
		Interval: interval,
		funcs:    make(map[string]JobFunc),
		// events: c.Events(),
		cli:       client,
		queueName: queueName,
		queueAddr: queueAddr,
	}

	w.queue = w.cli.Queue(queueName)

	return w, nil
}

func heartbeatStart(job *Job, done chan bool, heartbeat int, l sync.Locker) {
	tick := time.NewTicker(time.Duration(heartbeat) * time.Duration(time.Second))
	for {
		select {
		case <-done:
			tick.Stop()
			return
		case <-tick.C:
			l.Lock()
			success, err := job.Heartbeat()
			l.Unlock()
			if err != nil {
				log.Printf("warning, slow, send heartbeat***jid:%v, queue:%v, success:%v, error:%v",
					job.Jid, job.Queue, success, err)
			} else {
				log.Printf("warning, slow, send heartbeat***jid:%v, queue:%v, success:%v",
					job.Jid, job.Queue, success)
			}
		}
	}
}

//try our best to complete the job
func (w *Worker) tryCompleteJob(job *Job) error {
	for i := 0; i < 3; i++ {
		time.Sleep(time.Second)
		log.Println("tryCompleteJob", job.Jid)
		ipport := strings.Split(w.queueAddr, ":")
		client, err := Dial(ipport[0], ipport[1])
		if err != nil {
			log.Println("Dial err:", err)
			continue
		}

		job.SetClient(client)
		if _, err := job.CompleteWithNoData(); err != nil {
			client.Close()
			log.Println("tryCompleteJob", job.Jid, err)
		} else {
			client.Close()
			return nil
		}
	}

	return errors.New("tryCompleteJob " + job.Jid + "failed")
}

func (w *Worker) Start() error {
	// log.Println("worker Start")

	defer func() {
		w.cli.Close()
		w.cli = nil
	}()

	heartbeatStr, err := w.cli.GetConfig("heartbeat")
	heartbeat, err := strconv.Atoi(heartbeatStr)
	//log.Println("heartbeatStr:", heartbeat)
	if err != nil {
		heartbeat = 60
		log.Println(err)
	}

	heartbeat /= 2

	err = func(q *Queue) error {
		for {
			w.Lock()
			jobs, err := q.Pop(1) //we may pop more if fast enough
			w.Unlock()

			if err != nil {
				log.Println(err)
				return err
			}

			if len(jobs) == 0 {
				time.Sleep(time.Duration(w.Interval) * time.Millisecond)
				continue
			}

			for i := 0; i < len(jobs); i++ {
				if len(jobs[i].History) > 2 {
					log.Printf("warning, multiple processed exist %+v\n", jobs[i])
				}
				done := make(chan bool)
				//todo: using seprate connection to send heartbeat
				go heartbeatStart(jobs[i], done, heartbeat, w)
				f, ok := w.funcs[jobs[i].Klass]
				if !ok { //we got a job that not belongs to us
					done <- false
					log.Fatalln("got a message not belongs to us, queue", q.Name, jobs[i])
					continue
				}

				err := f(jobs[i])
				if err != nil {
					// TODO: probably do something with this
					log.Println("error: job failed, id", jobs[i].Jid, "queue", jobs[i].Queue, err.Error())
					w.Lock()
					success, err := jobs[i].Fail("fail", err.Error())
					w.Unlock()
					done <- false
					if err != nil {
						log.Printf("fail job:%+v success:%v, error:%v",
							jobs[i].Jid, success, err)
						return err
					}
				} else {
					w.Lock()
					success, err := jobs[i].CompleteWithNoData()
					w.Unlock()
					done <- true
					if err != nil {
						err = w.tryCompleteJob(jobs[i])
						if err != nil {
							log.Printf("fail job:%+v success:%v, error:%v",
								jobs[i].Jid, success, err)
						} else {
							log.Println("retry complete job ", jobs[i].Jid, "ok")
						}
						return errors.New("restart")
					}
				}
			}
		}
		return nil
	}(w.queue)

	return err
}

func (w *Worker) AddFunc(name string, f JobFunc) error {
	if _, ok := w.funcs[name]; ok {
		return fmt.Errorf("function \"%s\" already exists", name)
	}

	w.funcs[name] = f
	return nil
}

// Adds all the methods in the passed interface as job functions.
// Job names are in the form of: name.methodname
func (w *Worker) AddService(name string, rcvr interface{}) error {
	typ := reflect.TypeOf(rcvr)
	val := reflect.ValueOf(rcvr)
	for i := 0; i < typ.NumMethod(); i++ {
		method := typ.Method(i)
		w.AddFunc(name+"."+method.Name, func(job *Job) error {
			ret := method.Func.Call([]reflect.Value{val, reflect.ValueOf(job)})
			if len(ret) > 0 {
				if err, ok := ret[0].Interface().(error); ok {
					return err
				}
			} else {
				errStr := "reflect len less than zero." + strconv.Itoa(len(ret))
				return errors.New(errStr)
			}

			return nil
		})
	}

	return nil
}
<|endoftext|>


<commit_before>package resque

import (
	"encoding/json"
	"sync"

	"github.com/go-redis/redis"
)

// Worker primitive to implement a consumer.
type Worker struct {
	// JobClass type of job being of executed.
	JobClass JobClass
	// Done channel to request the worker to stop execution.
	Done <-chan bool
}

// Run worker until an error occurs or the worker is requested to quit.
func (w Worker) Run() error {
	pubsub, err := w.JobClass.Queue.Subscribe()
	if err != nil {
		return err
	}

	chan_in := make(chan Job, 1)
	chan_err := make(chan error, 1)
	chan_exit := make(chan bool, 1)

	var wg sync.WaitGroup

	defer func() {
		wg.Wait()
		close(chan_in)
		close(chan_err)
		close(chan_exit)
		pubsub.Close()
	}()

	go consume(pubsub, chan_in, chan_err, chan_exit, wg)
	chan_exit <- false

	for {
		select {
		case job := <-chan_in:
			w.JobClass.Performer.Perform(job.Args...)

			// ready for the next message
			chan_exit <- false

		case err = chan_err:
			chan_exit <- true
			return err

		case <-w.Done:
			chan_exit <- true
			return nil
		}
	}
}

// consume retrieves a message from a Redis subscription,
// transforms it into a Job and
// send to the input channel.
// In case of error during this process,
// it will send the error message to an error channel.
func consume(pubsub *redis.PubSub, chan_in <-chan Job, chan_err <-chan error, chan_exit <-chan bool, wg sync.WaitGroup) {
	defer wg.Done()

	var job Job

	for {
		select {
		case quit := <-chan_exit:
			if quit {
				return
			}

			for {
				msg, err := pubsub.ReceiveMessage()
				if err != nil {
					chan_err <- err
					continue
				}

				job = Job{}
				err = json.Unmarshal([]byte(msg.Payload), &job)
				if err != nil {
					chan_err <- err
					continue
				}

				chan_in <- job
			}
		}
	}
}
<commit_msg>worker.go: added license header.<commit_after>//
// Copyright 2017 Pedro Salgado
//
// 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 resque

import (
	"encoding/json"
	"sync"

	"github.com/go-redis/redis"
)

// Worker primitive to implement a consumer.
type Worker struct {
	// JobClass type of job being of executed.
	JobClass JobClass
	// Done channel to request the worker to stop execution.
	Done <-chan bool
}

// Run worker until an error occurs or the worker is requested to quit.
func (w Worker) Run() error {
	pubsub, err := w.JobClass.Queue.Subscribe()
	if err != nil {
		return err
	}

	chan_in := make(chan Job, 1)
	chan_err := make(chan error, 1)
	chan_exit := make(chan bool, 1)

	var wg sync.WaitGroup

	defer func() {
		wg.Wait()
		close(chan_in)
		close(chan_err)
		close(chan_exit)
		pubsub.Close()
	}()

	go consume(pubsub, chan_in, chan_err, chan_exit, wg)
	chan_exit <- false

	for {
		select {
		case job := <-chan_in:
			w.JobClass.Performer.Perform(job.Args...)

			// ready for the next message
			chan_exit <- false

		case err = chan_err:
			chan_exit <- true
			return err

		case <-w.Done:
			chan_exit <- true
			return nil
		}
	}
}

// consume retrieves a message from a Redis subscription,
// transforms it into a Job and
// send to the input channel.
// In case of error during this process,
// it will send the error message to an error channel.
func consume(pubsub *redis.PubSub, chan_in <-chan Job, chan_err <-chan error, chan_exit <-chan bool, wg sync.WaitGroup) {
	defer wg.Done()

	var job Job

	for {
		select {
		case quit := <-chan_exit:
			if quit {
				return
			}

			for {
				msg, err := pubsub.ReceiveMessage()
				if err != nil {
					chan_err <- err
					continue
				}

				job = Job{}
				err = json.Unmarshal([]byte(msg.Payload), &job)
				if err != nil {
					chan_err <- err
					continue
				}

				chan_in <- job
			}
		}
	}
}
<|endoftext|>


<commit_before>// TODO: example in example_test?
// TODO: goroutine leak test?

package fakestdio

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"sync"
)

// FakeStdio can be used to fake stdin and capture stdout.
// Between creating a new FakeStdio and calling ReadAndRestore on it,
// code reading os.Stdin in the process will get the contents of stdinText
// passed to New. Output to os.Stdout will be captured and returned from
// ReadAndRestore.
type FakeStdio struct {
	origStdout   *os.File
	stdoutReader *os.File

	outBuf *bytes.Buffer
	outWg  *sync.WaitGroup

	origStdin   *os.File
	stdinWriter *os.File
}

func New(stdinText string) (*FakeStdio, error) {
	stdinReader, stdinWriter, err := os.Pipe()
	if err != nil {
		return nil, err
	}

	stdoutReader, stdoutWriter, err := os.Pipe()
	if err != nil {
		return nil, err
	}

	// Fake stdin and write input text.
	origStdin := os.Stdin
	os.Stdin = stdinReader

	_, err = stdinWriter.Write([]byte(stdinText))
	if err != nil {
		stdinWriter.Close()
		os.Stdin = origStdin
		return nil, err
	}

	// Fake stdout.
	origStdout := os.Stdout
	os.Stdout = stdoutWriter

	var wg sync.WaitGroup
	var outBuf bytes.Buffer

	// This goroutine continuously reads stdout into buf.
	wg.Add(1)
	go func() {
		defer wg.Done()
		buf := make([]byte, 2048)
		for {
			n, err := stdoutReader.Read(buf)
			if err != nil && err != io.EOF {
				log.Println("read error", err)
				return
			}
			if n == 0 {
				return
			}
			outBuf.Write(buf[:n])
		}
	}()

	return &FakeStdio{
		origStdout:   origStdout,
		stdoutReader: stdoutReader,
		outBuf:       &outBuf,
		outWg:        &wg,
		origStdin:    origStdin,
		stdinWriter:  stdinWriter,
	}, nil
}

// CloseStdin closes the fake stdin. This may be necessary if the process has
// logic for reading stdin until EOF; otherwise such code would block forever.
func (sf *FakeStdio) CloseStdin() {
	if sf.stdinWriter != nil {
		sf.stdinWriter.Close()
		sf.stdinWriter = nil
	}
}

// ReadAndRestore collects all captured stdout and returns it; it also restores
// os.Stdin and os.Stdout to their original values.
func (sf *FakeStdio) ReadAndRestore() ([]byte, error) {
	if sf.stdoutReader == nil {
		return nil, fmt.Errorf("ReadAndRestore from closed FakeStdio")
	}

	// TODO: explain
	os.Stdout.Close()
	sf.outWg.Wait()

	os.Stdout = sf.origStdout
	os.Stdin = sf.origStdin

	if sf.stdoutReader != nil {
		sf.stdoutReader.Close()
		sf.stdoutReader = nil
	}

	if sf.stdinWriter != nil {
		sf.stdinWriter.Close()
		sf.stdinWriter = nil
	}

	return sf.outBuf.Bytes(), nil
}
<commit_msg>Don't copy through buf. Use Buffer.ReadFrom to read directly from reader<commit_after>// TODO: example in example_test?
// TODO: goroutine leak test?

package fakestdio

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"sync"
)

// FakeStdio can be used to fake stdin and capture stdout.
// Between creating a new FakeStdio and calling ReadAndRestore on it,
// code reading os.Stdin in the process will get the contents of stdinText
// passed to New. Output to os.Stdout will be captured and returned from
// ReadAndRestore.
type FakeStdio struct {
	origStdout   *os.File
	stdoutReader *os.File

	outBuf *bytes.Buffer
	outWg  *sync.WaitGroup

	origStdin   *os.File
	stdinWriter *os.File
}

func New(stdinText string) (*FakeStdio, error) {
	stdinReader, stdinWriter, err := os.Pipe()
	if err != nil {
		return nil, err
	}

	stdoutReader, stdoutWriter, err := os.Pipe()
	if err != nil {
		return nil, err
	}

	// Fake stdin and write input text.
	origStdin := os.Stdin
	os.Stdin = stdinReader

	_, err = stdinWriter.Write([]byte(stdinText))
	if err != nil {
		stdinWriter.Close()
		os.Stdin = origStdin
		return nil, err
	}

	// Fake stdout.
	origStdout := os.Stdout
	os.Stdout = stdoutWriter

	var wg sync.WaitGroup
	var outBuf bytes.Buffer

	// This goroutine continuously reads stdout into outBuf.
	wg.Add(1)
	go func() {
		defer wg.Done()
		for {
			n, err := outBuf.ReadFrom(stdoutReader)
			if err != nil && err != io.EOF {
				log.Println("read error", err)
				return
			}
			if n == 0 {
				return
			}
		}
	}()

	return &FakeStdio{
		origStdout:   origStdout,
		stdoutReader: stdoutReader,
		outBuf:       &outBuf,
		outWg:        &wg,
		origStdin:    origStdin,
		stdinWriter:  stdinWriter,
	}, nil
}

// CloseStdin closes the fake stdin. This may be necessary if the process has
// logic for reading stdin until EOF; otherwise such code would block forever.
func (sf *FakeStdio) CloseStdin() {
	if sf.stdinWriter != nil {
		sf.stdinWriter.Close()
		sf.stdinWriter = nil
	}
}

// ReadAndRestore collects all captured stdout and returns it; it also restores
// os.Stdin and os.Stdout to their original values.
func (sf *FakeStdio) ReadAndRestore() ([]byte, error) {
	if sf.stdoutReader == nil {
		return nil, fmt.Errorf("ReadAndRestore from closed FakeStdio")
	}

	// TODO: explain
	os.Stdout.Close()
	sf.outWg.Wait()

	os.Stdout = sf.origStdout
	os.Stdin = sf.origStdin

	if sf.stdoutReader != nil {
		sf.stdoutReader.Close()
		sf.stdoutReader = nil
	}

	if sf.stdinWriter != nil {
		sf.stdinWriter.Close()
		sf.stdinWriter = nil
	}

	return sf.outBuf.Bytes(), nil
}
<|endoftext|>


<commit_before>package httpsched

import (
	"context"
	"log"
	"net/http"
	"net/url"
	"time"

	"github.com/mesos/mesos-go/api/v1/lib"
	"github.com/mesos/mesos-go/api/v1/lib/backoff"
	mesosclient "github.com/mesos/mesos-go/api/v1/lib/client"
	"github.com/mesos/mesos-go/api/v1/lib/encoding"
	"github.com/mesos/mesos-go/api/v1/lib/httpcli"
	"github.com/mesos/mesos-go/api/v1/lib/httpcli/apierrors"
	"github.com/mesos/mesos-go/api/v1/lib/scheduler"
	"github.com/mesos/mesos-go/api/v1/lib/scheduler/calls"
)

var (
	errNotHTTPCli  = httpcli.ProtocolError("expected an httpcli.Response object, found something else instead")
	errBadLocation = httpcli.ProtocolError("failed to build new Mesos service endpoint URL from Location header")

	DefaultRedirectSettings = RedirectSettings{
		MaxAttempts:      9,
		MaxBackoffPeriod: 13 * time.Second,
		MinBackoffPeriod: 500 * time.Millisecond,
	}
)

type (
	RedirectSettings struct {
		MaxAttempts      int           // per httpDo invocation
		MaxBackoffPeriod time.Duration // should be more than minBackoffPeriod
		MinBackoffPeriod time.Duration // should be less than maxBackoffPeriod
	}

	client struct {
		*httpcli.Client
		redirect RedirectSettings
	}

	// Caller is the public interface a framework scheduler's should consume
	Caller interface {
		calls.Caller
		// httpDo is intentionally package-private; clients of this package may extend a Caller
		// generated by this package by overriding the Call func but may not customize httpDo.
		httpDo(context.Context, encoding.Marshaler, ...httpcli.RequestOpt) (mesos.Response, error)
	}

	callerInternal interface {
		Caller
		// WithTemporary configures the Client with the temporary option and returns the results of
		// invoking f(). Changes made to the Client by the temporary option are reverted before this
		// func returns.
		WithTemporary(opt httpcli.Opt, f func() error) error
	}

	// Option is a functional configuration option type
	Option func(*client) Option

	callerTemporary struct {
		callerInternal                      // delegate actually does the work
		requestOpts    []httpcli.RequestOpt // requestOpts are temporary per-request options
		opt            httpcli.Opt          // opt is a temporary client option
	}
)

func (ct *callerTemporary) httpDo(ctx context.Context, m encoding.Marshaler, opt ...httpcli.RequestOpt) (resp mesos.Response, err error) {
	ct.callerInternal.WithTemporary(ct.opt, func() error {
		if len(opt) == 0 {
			opt = ct.requestOpts
		} else if len(ct.requestOpts) > 0 {
			opt = append(opt[:], ct.requestOpts...)
		}
		resp, err = ct.callerInternal.httpDo(ctx, m, opt...)
		return nil
	})
	return
}

func (ct *callerTemporary) Call(ctx context.Context, call *scheduler.Call) (resp mesos.Response, err error) {
	ct.callerInternal.WithTemporary(ct.opt, func() error {
		resp, err = ct.callerInternal.Call(ctx, call)
		return nil
	})
	return
}

// MaxRedirects is a functional option that sets the maximum number of per-call HTTP redirects for a scheduler client
func MaxRedirects(mr int) Option {
	return func(c *client) Option {
		old := c.redirect.MaxAttempts
		c.redirect.MaxAttempts = mr
		return MaxRedirects(old)
	}
}

// NewCaller returns a scheduler API Client in the form of a Caller. Concurrent invocations
// of Call upon the returned caller are safely executed in a serial fashion. It is expected that
// there are no other users of the given Client since its state may be modified by this impl.
func NewCaller(cl *httpcli.Client, opts ...Option) calls.Caller {
	result := &client{Client: cl, redirect: DefaultRedirectSettings}
	cl.With(result.redirectHandler())
	for _, o := range opts {
		if o != nil {
			o(result)
		}
	}
	return &state{
		client: result,
		fn:     disconnectedFn,
	}
}

// httpDo decorates the inherited behavior w/ support for HTTP redirection to follow Mesos leadership changes.
// NOTE: this implementation will change the state of the client upon Mesos leadership changes.
func (cli *client) httpDo(ctx context.Context, m encoding.Marshaler, opt ...httpcli.RequestOpt) (resp mesos.Response, err error) {
	var (
		done            chan struct{} // avoid allocating these chans unless we actually need to redirect
		redirectBackoff <-chan struct{}
		getBackoff      = func() <-chan struct{} {
			if redirectBackoff == nil {
				done = make(chan struct{})
				redirectBackoff = backoff.Notifier(cli.redirect.MinBackoffPeriod, cli.redirect.MaxBackoffPeriod, done)
			}
			return redirectBackoff
		}
	)
	defer func() {
		if done != nil {
			close(done)
		}
	}()
	opt = append(opt, httpcli.Context(ctx))
	for attempt := 0; ; attempt++ {
		resp, err = cli.Client.Do(m, opt...)
		redirectErr, ok := err.(*mesosRedirectionError)
		if !ok {
			return resp, err
		}
		if attempt < cli.redirect.MaxAttempts {
			log.Println("redirecting to " + redirectErr.newURL)
			cli.With(httpcli.Endpoint(redirectErr.newURL))
			select {
			case <-getBackoff():
			case <-ctx.Done():
				return nil, ctx.Err()
			}
			continue
		}
		return
	}
}

// Call implements Client
func (cli *client) Call(ctx context.Context, call *scheduler.Call) (mesos.Response, error) {
	return cli.httpDo(ctx, call)
}

type mesosRedirectionError struct{ newURL string }

func (mre *mesosRedirectionError) Error() string {
	return "mesos server sent redirect to: " + mre.newURL
}

// redirectHandler returns a config options that decorates the default response handling routine;
// it transforms normal Mesos redirect "errors" into mesosRedirectionErrors by parsing the Location
// header and computing the address of the next endpoint that should be used to replay the failed
// HTTP request.
func (cli *client) redirectHandler() httpcli.Opt {
	return httpcli.HandleResponse(func(hres *http.Response, rc mesosclient.ResponseClass, err error) (mesos.Response, error) {
		resp, err := cli.HandleResponse(hres, rc, err) // default response handler
		if err == nil || !apierrors.CodeNotLeader.Matches(err) {
			return resp, err
		}
		// TODO(jdef) for now, we're tightly coupled to the httpcli package's Response type
		res, ok := resp.(*httpcli.Response)
		if !ok {
			if resp != nil {
				resp.Close()
			}
			return nil, errNotHTTPCli
		}
		log.Println("master changed?")
		location, ok := buildNewEndpoint(res.Header.Get("Location"), cli.Endpoint())
		if !ok {
			return nil, errBadLocation
		}
		res.Close()
		return nil, &mesosRedirectionError{location}
	})
}

func buildNewEndpoint(location, currentEndpoint string) (string, bool) {
	// TODO(jdef) refactor this
	// mesos v0.29 will actually send back fully-formed URLs in the Location header
	if location == "" {
		return "", false
	}
	// current format appears to be //x.y.z.w:port
	hostport, parseErr := url.Parse(location)
	if parseErr != nil || hostport.Host == "" {
		return "", false
	}
	current, parseErr := url.Parse(currentEndpoint)
	if parseErr != nil {
		return "", false
	}
	current.Host = hostport.Host
	return current.String(), true
}
<commit_msg>mute httpsched debug logging (#334)<commit_after>package httpsched

import (
	"context"
	"log"
	"net/http"
	"net/url"
	"time"

	"github.com/mesos/mesos-go/api/v1/lib"
	"github.com/mesos/mesos-go/api/v1/lib/backoff"
	mesosclient "github.com/mesos/mesos-go/api/v1/lib/client"
	"github.com/mesos/mesos-go/api/v1/lib/encoding"
	"github.com/mesos/mesos-go/api/v1/lib/httpcli"
	"github.com/mesos/mesos-go/api/v1/lib/httpcli/apierrors"
	"github.com/mesos/mesos-go/api/v1/lib/scheduler"
	"github.com/mesos/mesos-go/api/v1/lib/scheduler/calls"
)

var (
	errNotHTTPCli  = httpcli.ProtocolError("expected an httpcli.Response object, found something else instead")
	errBadLocation = httpcli.ProtocolError("failed to build new Mesos service endpoint URL from Location header")

	DefaultRedirectSettings = RedirectSettings{
		MaxAttempts:      9,
		MaxBackoffPeriod: 13 * time.Second,
		MinBackoffPeriod: 500 * time.Millisecond,
	}
)

type (
	RedirectSettings struct {
		MaxAttempts      int           // per httpDo invocation
		MaxBackoffPeriod time.Duration // should be more than minBackoffPeriod
		MinBackoffPeriod time.Duration // should be less than maxBackoffPeriod
	}

	client struct {
		*httpcli.Client
		redirect RedirectSettings
	}

	// Caller is the public interface a framework scheduler's should consume
	Caller interface {
		calls.Caller
		// httpDo is intentionally package-private; clients of this package may extend a Caller
		// generated by this package by overriding the Call func but may not customize httpDo.
		httpDo(context.Context, encoding.Marshaler, ...httpcli.RequestOpt) (mesos.Response, error)
	}

	callerInternal interface {
		Caller
		// WithTemporary configures the Client with the temporary option and returns the results of
		// invoking f(). Changes made to the Client by the temporary option are reverted before this
		// func returns.
		WithTemporary(opt httpcli.Opt, f func() error) error
	}

	// Option is a functional configuration option type
	Option func(*client) Option

	callerTemporary struct {
		callerInternal                      // delegate actually does the work
		requestOpts    []httpcli.RequestOpt // requestOpts are temporary per-request options
		opt            httpcli.Opt          // opt is a temporary client option
	}
)

func (ct *callerTemporary) httpDo(ctx context.Context, m encoding.Marshaler, opt ...httpcli.RequestOpt) (resp mesos.Response, err error) {
	ct.callerInternal.WithTemporary(ct.opt, func() error {
		if len(opt) == 0 {
			opt = ct.requestOpts
		} else if len(ct.requestOpts) > 0 {
			opt = append(opt[:], ct.requestOpts...)
		}
		resp, err = ct.callerInternal.httpDo(ctx, m, opt...)
		return nil
	})
	return
}

func (ct *callerTemporary) Call(ctx context.Context, call *scheduler.Call) (resp mesos.Response, err error) {
	ct.callerInternal.WithTemporary(ct.opt, func() error {
		resp, err = ct.callerInternal.Call(ctx, call)
		return nil
	})
	return
}

// MaxRedirects is a functional option that sets the maximum number of per-call HTTP redirects for a scheduler client
func MaxRedirects(mr int) Option {
	return func(c *client) Option {
		old := c.redirect.MaxAttempts
		c.redirect.MaxAttempts = mr
		return MaxRedirects(old)
	}
}

// NewCaller returns a scheduler API Client in the form of a Caller. Concurrent invocations
// of Call upon the returned caller are safely executed in a serial fashion. It is expected that
// there are no other users of the given Client since its state may be modified by this impl.
func NewCaller(cl *httpcli.Client, opts ...Option) calls.Caller {
	result := &client{Client: cl, redirect: DefaultRedirectSettings}
	cl.With(result.redirectHandler())
	for _, o := range opts {
		if o != nil {
			o(result)
		}
	}
	return &state{
		client: result,
		fn:     disconnectedFn,
	}
}

// httpDo decorates the inherited behavior w/ support for HTTP redirection to follow Mesos leadership changes.
// NOTE: this implementation will change the state of the client upon Mesos leadership changes.
func (cli *client) httpDo(ctx context.Context, m encoding.Marshaler, opt ...httpcli.RequestOpt) (resp mesos.Response, err error) {
	var (
		done            chan struct{} // avoid allocating these chans unless we actually need to redirect
		redirectBackoff <-chan struct{}
		getBackoff      = func() <-chan struct{} {
			if redirectBackoff == nil {
				done = make(chan struct{})
				redirectBackoff = backoff.Notifier(cli.redirect.MinBackoffPeriod, cli.redirect.MaxBackoffPeriod, done)
			}
			return redirectBackoff
		}
	)
	defer func() {
		if done != nil {
			close(done)
		}
	}()
	opt = append(opt, httpcli.Context(ctx))
	for attempt := 0; ; attempt++ {
		resp, err = cli.Client.Do(m, opt...)
		redirectErr, ok := err.(*mesosRedirectionError)
		if !ok {
			return resp, err
		}
		if attempt < cli.redirect.MaxAttempts {
			if debug {
				log.Println("redirecting to " + redirectErr.newURL)
			}
			cli.With(httpcli.Endpoint(redirectErr.newURL))
			select {
			case <-getBackoff():
			case <-ctx.Done():
				return nil, ctx.Err()
			}
			continue
		}
		return
	}
}

// Call implements Client
func (cli *client) Call(ctx context.Context, call *scheduler.Call) (mesos.Response, error) {
	return cli.httpDo(ctx, call)
}

type mesosRedirectionError struct{ newURL string }

func (mre *mesosRedirectionError) Error() string {
	return "mesos server sent redirect to: " + mre.newURL
}

// redirectHandler returns a config options that decorates the default response handling routine;
// it transforms normal Mesos redirect "errors" into mesosRedirectionErrors by parsing the Location
// header and computing the address of the next endpoint that should be used to replay the failed
// HTTP request.
func (cli *client) redirectHandler() httpcli.Opt {
	return httpcli.HandleResponse(func(hres *http.Response, rc mesosclient.ResponseClass, err error) (mesos.Response, error) {
		resp, err := cli.HandleResponse(hres, rc, err) // default response handler
		if err == nil || !apierrors.CodeNotLeader.Matches(err) {
			return resp, err
		}
		// TODO(jdef) for now, we're tightly coupled to the httpcli package's Response type
		res, ok := resp.(*httpcli.Response)
		if !ok {
			if resp != nil {
				resp.Close()
			}
			return nil, errNotHTTPCli
		}
		if debug {
			log.Println("master changed?")
		}
		location, ok := buildNewEndpoint(res.Header.Get("Location"), cli.Endpoint())
		if !ok {
			return nil, errBadLocation
		}
		res.Close()
		return nil, &mesosRedirectionError{location}
	})
}

func buildNewEndpoint(location, currentEndpoint string) (string, bool) {
	// TODO(jdef) refactor this
	// mesos v0.29 will actually send back fully-formed URLs in the Location header
	if location == "" {
		return "", false
	}
	// current format appears to be //x.y.z.w:port
	hostport, parseErr := url.Parse(location)
	if parseErr != nil || hostport.Host == "" {
		return "", false
	}
	current, parseErr := url.Parse(currentEndpoint)
	if parseErr != nil {
		return "", false
	}
	current.Host = hostport.Host
	return current.String(), true
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors 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 mungers

import (
	"fmt"
	"regexp"
	"time"

	"k8s.io/contrib/mungegithub/features"
	"k8s.io/contrib/mungegithub/github"

	"github.com/golang/glog"
	githubapi "github.com/google/go-github/github"
	"github.com/spf13/cobra"
)

const (
	day              = time.Hour * 24
	keepOpenLabel    = "keep-open"
	stalePullRequest = 90 * day // Close the PR if no human interaction for `stalePullRequest`
	closingComment   = `This PR hasn't been active in %s. Please re-open if necessary

You can add 'keep-open' label to prevent this from happening again.`
)

var (
	closingCommentRE = regexp.MustCompile(`This PR hasn't been active in \d+ days?\. Please re-open if necessary

You can add 'keep-open' label to prevent this from happening again\.`)
)

// CloseStalePR will ask the Bot to close any PullRequest that didn't
// have any human interactions in `stalePullRequest` duration.
//
// This is done by checking both review and issue comments, and by
// ignoring comments done with a bot name. We also consider re-open on the PR.
type CloseStalePR struct{}

func init() {
	s := CloseStalePR{}
	RegisterMungerOrDie(s)
	RegisterStaleComments(s)
}

// Name is the name usable in --pr-mungers
func (CloseStalePR) Name() string { return "close-stale-pr" }

// RequiredFeatures is a slice of 'features' that must be provided
func (CloseStalePR) RequiredFeatures() []string { return []string{} }

// Initialize will initialize the munger
func (CloseStalePR) Initialize(config *github.Config, features *features.Features) error {
	return nil
}

// EachLoop is called at the start of every munge loop
func (CloseStalePR) EachLoop() error { return nil }

// AddFlags will add any request flags to the cobra `cmd`
func (CloseStalePR) AddFlags(cmd *cobra.Command, config *github.Config) {}

func findLastHumanPullRequestUpdate(obj *github.MungeObject) (*time.Time, error) {
	pr, err := obj.GetPR()
	if err != nil {
		return nil, err
	}

	comments, err := obj.ListReviewComments()
	if err != nil {
		return nil, err
	}

	lastHuman := pr.CreatedAt
	for i := range comments {
		comment := comments[i]
		if comment.User == nil || comment.User.Login == nil || comment.CreatedAt == nil || comment.Body == nil {
			continue
		}
		if *comment.User.Login == botName || *comment.User.Login == jenkinsBotName {
			continue
		}
		if lastHuman.Before(*comment.UpdatedAt) {
			lastHuman = comment.UpdatedAt
		}
	}

	return lastHuman, nil
}

func findLastHumanIssueUpdate(obj *github.MungeObject) (*time.Time, error) {
	lastHuman := obj.Issue.CreatedAt

	comments, err := obj.ListComments()
	if err != nil {
		return nil, err
	}

	for i := range comments {
		comment := comments[i]
		if !validComment(comment) {
			continue
		}
		if mergeBotComment(comment) || jenkinsBotComment(comment) {
			continue
		}
		if lastHuman.Before(*comment.UpdatedAt) {
			lastHuman = comment.UpdatedAt
		}
	}

	return lastHuman, nil
}

func findLastInterestingEventUpdate(obj *github.MungeObject) (*time.Time, error) {
	lastInteresting := obj.Issue.CreatedAt

	events, err := obj.GetEvents()
	if err != nil {
		return nil, err
	}

	for i := range events {
		event := events[i]
		if event.Event == nil || *event.Event != "reopened" {
			continue
		}

		if lastInteresting.Before(*event.CreatedAt) {
			lastInteresting = event.CreatedAt
		}
	}

	return lastInteresting, nil
}

func findLastModificationTime(obj *github.MungeObject) (*time.Time, error) {
	lastHumanIssue, err := findLastHumanIssueUpdate(obj)
	if err != nil {
		return nil, err
	}
	lastHumanPR, err := findLastHumanPullRequestUpdate(obj)
	if err != nil {
		return nil, err
	}
	lastInterestingEvent, err := findLastInterestingEventUpdate(obj)
	if err != nil {
		return nil, err
	}

	lastModif := lastHumanPR
	if lastHumanIssue.After(*lastModif) {
		lastModif = lastHumanIssue
	}
	if lastInterestingEvent.After(*lastModif) {
		lastModif = lastInterestingEvent
	}

	return lastModif, nil
}

func durationToDays(duration time.Duration) string {
	days := duration / day
	dayString := "days"
	if days == 1 || days == -1 {
		dayString = "day"
	}
	return fmt.Sprintf("%d %s", days, dayString)
}

func closePullRequest(obj *github.MungeObject, inactiveFor time.Duration) {
	obj.WriteComment(
		fmt.Sprintf(closingComment, durationToDays(inactiveFor)))
	obj.ClosePR()
}

// Munge is the workhorse that will actually close the PRs
func (CloseStalePR) Munge(obj *github.MungeObject) {
	if !obj.IsPR() {
		return
	}

	if obj.HasLabel(keepOpenLabel) {
		return
	}

	lastModif, err := findLastModificationTime(obj)
	if err != nil {
		glog.Errorf("Failed to find last modification: %v", err)
		return
	}

	inactiveFor := time.Since(*lastModif)
	if inactiveFor > stalePullRequest {
		closePullRequest(obj, inactiveFor)
	} else {
		// Pull-request is active. Do nothing
	}
}

func (CloseStalePR) isStaleComment(obj *github.MungeObject, comment githubapi.IssueComment) bool {
	if !mergeBotComment(comment) {
		return false
	}

	if !closingCommentRE.MatchString(*comment.Body) {
		return false
	}

	return true
}

// StaleComments returns a slice of stale comments
func (s CloseStalePR) StaleComments(obj *github.MungeObject, comments []githubapi.IssueComment) []githubapi.IssueComment {
	return forEachCommentTest(obj, comments, s.isStaleComment)
}
<commit_msg>Warn before closing stale pull-request<commit_after>/*
Copyright 2016 The Kubernetes Authors 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 mungers

import (
	"fmt"
	"regexp"
	"time"

	"k8s.io/contrib/mungegithub/features"
	"k8s.io/contrib/mungegithub/github"

	"github.com/golang/glog"
	githubapi "github.com/google/go-github/github"
	"github.com/spf13/cobra"
)

const (
	day              = time.Hour * 24
	keepOpenLabel    = "keep-open"
	stalePullRequest = 90 * day // Close the PR if no human interaction for `stalePullRequest`
	startWarning     = 60 * day
	remindWarning    = 30 * day
	closingComment   = `This PR hasn't been active in %s. Feel free to reopen.

You can add 'keep-open' label to prevent this from happening again.`
	warningComment = `This PR hasn't been active in %s. Will be closed in %s.

You can add 'keep-open' label to prevent this from happening.`
)

var (
	closingCommentRE = regexp.MustCompile(`This PR hasn't been active in \d+ days?\. Feel free to reopen.

You can add 'keep-open' label to prevent this from happening again\.`)
	warningCommentRE = regexp.MustCompile(`This PR hasn't been active in \d+ days?\. Will be closed in \d+ days?\.

You can add 'keep-open' label to prevent this from happening\.`)
)

// CloseStalePR will ask the Bot to close any PullRequest that didn't
// have any human interactions in `stalePullRequest` duration.
//
// This is done by checking both review and issue comments, and by
// ignoring comments done with a bot name. We also consider re-open on the PR.
type CloseStalePR struct{}

func init() {
	s := CloseStalePR{}
	RegisterMungerOrDie(s)
	RegisterStaleComments(s)
}

// Name is the name usable in --pr-mungers
func (CloseStalePR) Name() string { return "close-stale-pr" }

// RequiredFeatures is a slice of 'features' that must be provided
func (CloseStalePR) RequiredFeatures() []string { return []string{} }

// Initialize will initialize the munger
func (CloseStalePR) Initialize(config *github.Config, features *features.Features) error {
	return nil
}

// EachLoop is called at the start of every munge loop
func (CloseStalePR) EachLoop() error { return nil }

// AddFlags will add any request flags to the cobra `cmd`
func (CloseStalePR) AddFlags(cmd *cobra.Command, config *github.Config) {}

func findLastHumanPullRequestUpdate(obj *github.MungeObject) (*time.Time, error) {
	pr, err := obj.GetPR()
	if err != nil {
		return nil, err
	}

	comments, err := obj.ListReviewComments()
	if err != nil {
		return nil, err
	}

	lastHuman := pr.CreatedAt
	for i := range comments {
		comment := comments[i]
		if comment.User == nil || comment.User.Login == nil || comment.CreatedAt == nil || comment.Body == nil {
			continue
		}
		if *comment.User.Login == botName || *comment.User.Login == jenkinsBotName {
			continue
		}
		if lastHuman.Before(*comment.UpdatedAt) {
			lastHuman = comment.UpdatedAt
		}
	}

	return lastHuman, nil
}

func findLastHumanIssueUpdate(obj *github.MungeObject) (*time.Time, error) {
	lastHuman := obj.Issue.CreatedAt

	comments, err := obj.ListComments()
	if err != nil {
		return nil, err
	}

	for i := range comments {
		comment := comments[i]
		if !validComment(comment) {
			continue
		}
		if mergeBotComment(comment) || jenkinsBotComment(comment) {
			continue
		}
		if lastHuman.Before(*comment.UpdatedAt) {
			lastHuman = comment.UpdatedAt
		}
	}

	return lastHuman, nil
}

func findLastInterestingEventUpdate(obj *github.MungeObject) (*time.Time, error) {
	lastInteresting := obj.Issue.CreatedAt

	events, err := obj.GetEvents()
	if err != nil {
		return nil, err
	}

	for i := range events {
		event := events[i]
		if event.Event == nil || *event.Event != "reopened" {
			continue
		}

		if lastInteresting.Before(*event.CreatedAt) {
			lastInteresting = event.CreatedAt
		}
	}

	return lastInteresting, nil
}

func findLastModificationTime(obj *github.MungeObject) (*time.Time, error) {
	lastHumanIssue, err := findLastHumanIssueUpdate(obj)
	if err != nil {
		return nil, err
	}
	lastHumanPR, err := findLastHumanPullRequestUpdate(obj)
	if err != nil {
		return nil, err
	}
	lastInterestingEvent, err := findLastInterestingEventUpdate(obj)
	if err != nil {
		return nil, err
	}

	lastModif := lastHumanPR
	if lastHumanIssue.After(*lastModif) {
		lastModif = lastHumanIssue
	}
	if lastInterestingEvent.After(*lastModif) {
		lastModif = lastInterestingEvent
	}

	return lastModif, nil
}

func findLatestWarningComment(obj *github.MungeObject) *githubapi.IssueComment {
	var lastFoundComment *githubapi.IssueComment

	comments, err := obj.ListComments()
	if err != nil {
		return nil
	}

	for i := range comments {
		comment := comments[i]
		if !validComment(comment) {
			continue
		}
		if !mergeBotComment(comment) {
			continue
		}

		if !warningCommentRE.MatchString(*comment.Body) {
			continue
		}

		if lastFoundComment == nil || lastFoundComment.CreatedAt.Before(*comment.UpdatedAt) {
			if lastFoundComment != nil {
				obj.DeleteComment(lastFoundComment)
			}
			lastFoundComment = &comment
		}
	}

	return lastFoundComment
}

func durationToDays(duration time.Duration) string {
	days := duration / day
	dayString := "days"
	if days == 1 || days == -1 {
		dayString = "day"
	}
	return fmt.Sprintf("%d %s", days, dayString)
}

func closePullRequest(obj *github.MungeObject, inactiveFor time.Duration) {
	comment := findLatestWarningComment(obj)
	if comment != nil {
		obj.DeleteComment(comment)
	}

	obj.WriteComment(fmt.Sprintf(closingComment, durationToDays(inactiveFor)))
	obj.ClosePR()
}

func postWarningComment(obj *github.MungeObject, inactiveFor time.Duration, closeIn time.Duration) {
	obj.WriteComment(fmt.Sprintf(
		warningComment,
		durationToDays(inactiveFor),
		durationToDays(closeIn)))
}

func checkAndWarn(obj *github.MungeObject, inactiveFor time.Duration, closeIn time.Duration) {
	if closeIn < day {
		// We are going to close the PR in less than a day. Too late to warn
		return
	}
	comment := findLatestWarningComment(obj)
	if comment == nil {
		// We don't already have the comment. Post it
		postWarningComment(obj, inactiveFor, closeIn)
	} else if time.Since(*comment.UpdatedAt) > remindWarning {
		// It's time to warn again
		obj.DeleteComment(comment)
		postWarningComment(obj, inactiveFor, closeIn)
	} else {
		// We already have a warning, and it's not expired. Do nothing
	}
}

// Munge is the workhorse that will actually close the PRs
func (CloseStalePR) Munge(obj *github.MungeObject) {
	if !obj.IsPR() {
		return
	}

	if obj.HasLabel(keepOpenLabel) {
		return
	}

	lastModif, err := findLastModificationTime(obj)
	if err != nil {
		glog.Errorf("Failed to find last modification: %v", err)
		return
	}

	closeIn := -time.Since(lastModif.Add(stalePullRequest))
	inactiveFor := time.Since(*lastModif)
	if closeIn <= 0 {
		closePullRequest(obj, inactiveFor)
	} else if closeIn <= startWarning {
		checkAndWarn(obj, inactiveFor, closeIn)
	} else {
		// Pull-request is active. Do nothing
	}
}

func (CloseStalePR) isStaleComment(obj *github.MungeObject, comment githubapi.IssueComment) bool {
	if !mergeBotComment(comment) {
		return false
	}

	if !closingCommentRE.MatchString(*comment.Body) {
		return false
	}

	return true
}

// StaleComments returns a slice of stale comments
func (s CloseStalePR) StaleComments(obj *github.MungeObject, comments []githubapi.IssueComment) []githubapi.IssueComment {
	return forEachCommentTest(obj, comments, s.isStaleComment)
}
<|endoftext|>


<commit_before>package gocb

import (
	"encoding/json"
)

type Transcoder interface {
	Decode([]byte, uint32, interface{}) error
	Encode(interface{}) ([]byte, uint32, error)
}

type DefaultTranscoder struct {
}

func (t DefaultTranscoder) Decode(bytes []byte, flags uint32, out interface{}) error {
	// Check for legacy flags
	if flags&cfMask == 0 {
		// Legacy Flags
		if flags == lfJson {
			// Legacy JSON
			flags = cfFmtJson
		} else {
			return clientError{"Unexpected legacy flags value"}
		}
	}

	// Make sure compression is disabled
	if flags&cfCmprMask != cfCmprNone {
		return clientError{"Unexpected value compression"}
	}

	// Normal types of decoding
	if flags&cfFmtMask == cfFmtBinary {
		switch typedOut := out.(type) {
		case *[]byte:
			*typedOut = bytes
			return nil
		case *interface{}:
			*typedOut = bytes
			return nil
		default:
			return clientError{"You must encode binary in a byte array or interface"}
		}
	} else if flags&cfFmtMask == cfFmtString {
		switch typedOut := out.(type) {
		case *string:
			*typedOut = string(bytes)
			return nil
		case *interface{}:
			*typedOut = string(bytes)
			return nil
		default:
			return clientError{"You must encode a string in a string or interface"}
		}
	} else if flags&cfFmtMask == cfFmtJson {
		err := json.Unmarshal(bytes, &out)
		if err != nil {
			return err
		}
		return nil
	} else {
		return clientError{"Unexpected flags value"}
	}
}

func (t DefaultTranscoder) Encode(value interface{}) ([]byte, uint32, error) {
	var bytes []byte
	var flags uint32
	var err error

	switch value.(type) {
	case []byte:
		bytes = value.([]byte)
		flags = cfFmtBinary
	case *[]byte:
		bytes = *value.(*[]byte)
		flags = cfFmtBinary
	case string:
		bytes = []byte(value.(string))
		flags = cfFmtString
	case *string:
		bytes = []byte(*value.(*string))
		flags = cfFmtString
	default:
		bytes, err = json.Marshal(value)
		if err != nil {
			return nil, 0, err
		}
		flags = cfFmtJson
	}

	// No compression supported currently

	return bytes, flags, nil
}
<commit_msg>Fix bug when passing *interface to storage methods.<commit_after>package gocb

import (
	"encoding/json"
)

type Transcoder interface {
	Decode([]byte, uint32, interface{}) error
	Encode(interface{}) ([]byte, uint32, error)
}

type DefaultTranscoder struct {
}

func (t DefaultTranscoder) Decode(bytes []byte, flags uint32, out interface{}) error {
	// Check for legacy flags
	if flags&cfMask == 0 {
		// Legacy Flags
		if flags == lfJson {
			// Legacy JSON
			flags = cfFmtJson
		} else {
			return clientError{"Unexpected legacy flags value"}
		}
	}

	// Make sure compression is disabled
	if flags&cfCmprMask != cfCmprNone {
		return clientError{"Unexpected value compression"}
	}

	// Normal types of decoding
	if flags&cfFmtMask == cfFmtBinary {
		switch typedOut := out.(type) {
		case *[]byte:
			*typedOut = bytes
			return nil
		case *interface{}:
			*typedOut = bytes
			return nil
		default:
			return clientError{"You must encode binary in a byte array or interface"}
		}
	} else if flags&cfFmtMask == cfFmtString {
		switch typedOut := out.(type) {
		case *string:
			*typedOut = string(bytes)
			return nil
		case *interface{}:
			*typedOut = string(bytes)
			return nil
		default:
			return clientError{"You must encode a string in a string or interface"}
		}
	} else if flags&cfFmtMask == cfFmtJson {
		err := json.Unmarshal(bytes, &out)
		if err != nil {
			return err
		}
		return nil
	} else {
		return clientError{"Unexpected flags value"}
	}
}

func (t DefaultTranscoder) Encode(value interface{}) ([]byte, uint32, error) {
	var bytes []byte
	var flags uint32
	var err error

	switch value.(type) {
	case []byte:
		bytes = value.([]byte)
		flags = cfFmtBinary
	case *[]byte:
		bytes = *value.(*[]byte)
		flags = cfFmtBinary
	case string:
		bytes = []byte(value.(string))
		flags = cfFmtString
	case *string:
		bytes = []byte(*value.(*string))
		flags = cfFmtString
	case *interface{}:
		return t.Encode(*value.(*interface{}))
	default:
		bytes, err = json.Marshal(value)
		if err != nil {
			return nil, 0, err
		}
		flags = cfFmtJson
	}

	// No compression supported currently

	return bytes, flags, nil
}
<|endoftext|>


<commit_before>package trousseau

import (
	"github.com/oleiade/serrure/openpgp"
	"os"
	"path"
)

// Declare encryption types
type CryptoType int

const (
	SYMMETRIC_ENCRYPTION  CryptoType = 0
	ASYMMETRIC_ENCRYPTION CryptoType = 1

	SYMMETRIC_ENCRYPTION_REPR  string = "symmetric"
	ASYMMETRIC_ENCRYPTION_REPR string = "asymmetric"
)

// Declare available encryption algorithms
type CryptoAlgorithm int

const (
	GPG_ENCRYPTION     CryptoAlgorithm = 0
	AES_256_ENCRYPTION CryptoAlgorithm = 1

	GPG_ENCRYPTION_REPR     string = "gpg"
	AES_256_ENCRYPTION_REPR string = "aes256"
)

// Gnupg variables
var GnupgHome = path.Join(os.Getenv("HOME"), ".gnupg")
var GnupgPubring func() string = func() string { return path.Join(GnupgHome, "pubring.gpg") }
var GnupgSecring func() string = func() string { return path.Join(GnupgHome, "secring.gpg") }

// DecryptAsymmetricPGP decrypts an OpenPGP message using GnuPG.
func DecryptAsymmetricPGP(encryptedData []byte, passphrase string) ([]byte, error) {
	// Decrypt store data
	decryptionKeys, err := openpgp.ReadSecRing(GnupgSecring())
	if err != nil {
		return nil, err
	}

	plainData, err := openpgp.Decrypt(encryptedData, decryptionKeys, passphrase)
	if err != nil {
		return nil, err
	}

	return plainData, nil
}

func EncryptAsymmetricPGP(plainData []byte, recipients []string) ([]byte, error) {
	encryptionKeys, err := openpgp.ReadPubRing(GnupgPubring(), recipients)
	if err != nil {
		return nil, err
	}

	encData, err := openpgp.Encrypt(plainData, encryptionKeys)
	if err != nil {
		return nil, err
	}

	return encData, nil
}
<commit_msg>Update DecrypticAsymmetricPGP function to use serrure<commit_after>package trousseau

import (
	"github.com/oleiade/serrure/openpgp"
	"os"
	"path"
)

// Declare encryption types
type CryptoType int

const (
	SYMMETRIC_ENCRYPTION  CryptoType = 0
	ASYMMETRIC_ENCRYPTION CryptoType = 1

	SYMMETRIC_ENCRYPTION_REPR  string = "symmetric"
	ASYMMETRIC_ENCRYPTION_REPR string = "asymmetric"
)

// Declare available encryption algorithms
type CryptoAlgorithm int

const (
	GPG_ENCRYPTION     CryptoAlgorithm = 0
	AES_256_ENCRYPTION CryptoAlgorithm = 1

	GPG_ENCRYPTION_REPR     string = "gpg"
	AES_256_ENCRYPTION_REPR string = "aes256"
)

// Gnupg variables
var GnupgHome = path.Join(os.Getenv("HOME"), ".gnupg")
var GnupgPubring func() string = func() string { return path.Join(GnupgHome, "pubring.gpg") }
var GnupgSecring func() string = func() string { return path.Join(GnupgHome, "secring.gpg") }

// DecryptAsymmetricPGP decrypts an OpenPGP message using GnuPG.
func DecryptAsymmetricPGP(ed []byte, passphrase string) ([]byte, error) {
	d, err := openpgp.NewOpenPGPDecrypter(GnupgSecring(), passphrase)
	if err != nil {
		return nil, err
	}
	// Decrypt store data
	decryptionKeys, err := openpgp.ReadSecRing(GnupgSecring())
	if err != nil {
		return nil, err
	}

	pd, err := d.Decrypt(ed)
	if err != nil {
		return nil, err
	}

	return pd, nil
}

func EncryptAsymmetricPGP(plainData []byte, recipients []string) ([]byte, error) {
	encryptionKeys, err := openpgp.ReadPubRing(GnupgPubring(), recipients)
	if err != nil {
		return nil, err
	}

	encData, err := openpgp.Encrypt(plainData, encryptionKeys)
	if err != nil {
		return nil, err
	}

	return encData, nil
}
<|endoftext|>


<commit_before>package hawk

import (
	"crypto/hmac"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/base64"
	"hash"
	"net"
	"net/url"
	"strconv"
	"strings"
)

const headerVersion = 1

type AuthType int

const (
	Header AuthType = iota
	Response
	Bewit
)

type Mac struct {
	Type       AuthType
	Credential *Credential
	Uri        string
	Method     string
	HostPort   string
	Option     *Option
}

type TsMac struct {
	TimeStamp  int64
	Credential *Credential
}

type PayloadHash struct {
	ContentType string
	Payload     string
	Alg         Alg
}

// String returns a base64 encoded message authentication code.
func (m *Mac) String() (string, error) {
	digest, err := m.digest()
	return base64.StdEncoding.EncodeToString(digest), err
}

func (m *Mac) digest() ([]byte, error) {
	s := getHash(m.Credential.Alg)

	mac := hmac.New(s, []byte(m.Credential.Key))
	ns, err := m.normalized()
	if err != nil {
		return nil, err
	}
	mac.Write([]byte(ns))

	return mac.Sum(nil), nil
}

func (m *Mac) normalized() (string, error) {
	return normalized(m.Type, m.Uri, m.Method, m.HostPort, m.Option)
}

// String returns a base64 encoded message authentication code for timestamp
func (tm *TsMac) String() string {
	digest := tm.digest()
	return base64.StdEncoding.EncodeToString(digest)
}

func (tm *TsMac) digest() []byte {
	s := getHash(tm.Credential.Alg)

	mac := hmac.New(s, []byte(tm.Credential.Key))
	ns := "hawk." + strconv.Itoa(headerVersion) + ".ts" + "\n" + strconv.FormatInt(tm.TimeStamp, 10) + "\n"
	mac.Write([]byte(ns))

	return mac.Sum(nil)
}

// String returns a base64 encoded hash value of payload
func (h *PayloadHash) String() string {
	hash := h.hash()
	return base64.StdEncoding.EncodeToString(hash)
}

func sanitizeContentType(contentType string) string {
	return strings.TrimSpace(strings.ToLower(strings.Split(contentType, ";")[0]))
}

func (h *PayloadHash) hash() []byte {
	s := getHash(h.Alg)()

	ns := "hawk." + strconv.Itoa(headerVersion) + ".payload" + "\n" + sanitizeContentType(h.ContentType) + "\n" + h.Payload + "\n"
	s.Write([]byte(ns))

	return s.Sum(nil)
}

func normalized(authType AuthType, uri, method, customHost string, option *Option) (string, error) {
	u, err := url.Parse(uri)
	if err != nil {
		return "", err
	}

	var h string
	if customHost != "" {
		h = customHost
	} else {
		h = u.Host
	}

	host, port, _ := net.SplitHostPort(h)
	if port == "" {
		switch u.Scheme {
		case "http":
			port = "80"
		case "https":
			port = "443"
		}
	}
	if host == "" {
		if customHost != "" {
			host = customHost
		} else {
			host = u.Host
		}
	}

	path := u.Path
	if u.Query().Encode() != "" {
		path = path + "?" + u.RawQuery
	}

	header := "hawk" + "." + strconv.Itoa(headerVersion) + "." + strings.ToLower(authType.String())

	ext := ""
	if option.Ext != "" {
		ext = strings.Replace(option.Ext, "\\", "\\\\", -1)
		ext = strings.Replace(ext, "\n", "\\n", -1)
	}

	ns := header + "\n" +
		strconv.FormatInt(option.TimeStamp, 10) + "\n" +
		option.Nonce + "\n" +
		strings.ToUpper(method) + "\n" +
		path + "\n" +
		strings.ToLower(host) + "\n" +
		port + "\n" +
		option.Hash + "\n" +
		ext + "\n"

	if option.App != "" {
		ns = ns + option.App + "\n"
		ns = ns + option.Dlg + "\n"
	}

	return ns, nil
}

func getHash(alg Alg) func() hash.Hash {
	switch alg {
	case SHA256:
		return sha256.New
	case SHA512:
		return sha512.New
	default:
		return sha256.New
	}
}
<commit_msg>#12 Fix MAC generation for encoded URIs with special chars<commit_after>package hawk

import (
	"crypto/hmac"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/base64"
	"hash"
	"net"
	"net/url"
	"strconv"
	"strings"
)

const headerVersion = 1

type AuthType int

const (
	Header AuthType = iota
	Response
	Bewit
)

type Mac struct {
	Type       AuthType
	Credential *Credential
	Uri        string
	Method     string
	HostPort   string
	Option     *Option
}

type TsMac struct {
	TimeStamp  int64
	Credential *Credential
}

type PayloadHash struct {
	ContentType string
	Payload     string
	Alg         Alg
}

// String returns a base64 encoded message authentication code.
func (m *Mac) String() (string, error) {
	digest, err := m.digest()
	return base64.StdEncoding.EncodeToString(digest), err
}

func (m *Mac) digest() ([]byte, error) {
	s := getHash(m.Credential.Alg)

	mac := hmac.New(s, []byte(m.Credential.Key))
	ns, err := m.normalized()
	if err != nil {
		return nil, err
	}
	mac.Write([]byte(ns))

	return mac.Sum(nil), nil
}

func (m *Mac) normalized() (string, error) {
	return normalized(m.Type, m.Uri, m.Method, m.HostPort, m.Option)
}

// String returns a base64 encoded message authentication code for timestamp
func (tm *TsMac) String() string {
	digest := tm.digest()
	return base64.StdEncoding.EncodeToString(digest)
}

func (tm *TsMac) digest() []byte {
	s := getHash(tm.Credential.Alg)

	mac := hmac.New(s, []byte(tm.Credential.Key))
	ns := "hawk." + strconv.Itoa(headerVersion) + ".ts" + "\n" + strconv.FormatInt(tm.TimeStamp, 10) + "\n"
	mac.Write([]byte(ns))

	return mac.Sum(nil)
}

// String returns a base64 encoded hash value of payload
func (h *PayloadHash) String() string {
	hash := h.hash()
	return base64.StdEncoding.EncodeToString(hash)
}

func sanitizeContentType(contentType string) string {
	return strings.TrimSpace(strings.ToLower(strings.Split(contentType, ";")[0]))
}

func (h *PayloadHash) hash() []byte {
	s := getHash(h.Alg)()

	ns := "hawk." + strconv.Itoa(headerVersion) + ".payload" + "\n" + sanitizeContentType(h.ContentType) + "\n" + h.Payload + "\n"
	s.Write([]byte(ns))

	return s.Sum(nil)
}

func normalized(authType AuthType, uri, method, customHost string, option *Option) (string, error) {
	u, err := url.Parse(uri)
	if err != nil {
		return "", err
	}

	var h string
	if customHost != "" {
		h = customHost
	} else {
		h = u.Host
	}

	host, port, _ := net.SplitHostPort(h)
	if port == "" {
		switch u.Scheme {
		case "http":
			port = "80"
		case "https":
			port = "443"
		}
	}
	if host == "" {
		if customHost != "" {
			host = customHost
		} else {
			host = u.Host
		}
	}

	path := u.RawPath
	if u.Query().Encode() != "" {
		path = path + "?" + u.RawQuery
	}

	header := "hawk" + "." + strconv.Itoa(headerVersion) + "." + strings.ToLower(authType.String())

	ext := ""
	if option.Ext != "" {
		ext = strings.Replace(option.Ext, "\\", "\\\\", -1)
		ext = strings.Replace(ext, "\n", "\\n", -1)
	}

	ns := header + "\n" +
		strconv.FormatInt(option.TimeStamp, 10) + "\n" +
		option.Nonce + "\n" +
		strings.ToUpper(method) + "\n" +
		path + "\n" +
		strings.ToLower(host) + "\n" +
		port + "\n" +
		option.Hash + "\n" +
		ext + "\n"

	if option.App != "" {
		ns = ns + option.App + "\n"
		ns = ns + option.Dlg + "\n"
	}

	return ns, nil
}

func getHash(alg Alg) func() hash.Hash {
	switch alg {
	case SHA256:
		return sha256.New
	case SHA512:
		return sha512.New
	default:
		return sha256.New
	}
}
<|endoftext|>


<commit_before>package aranGO

import (
	"encoding/json"
	"errors"
	"reflect"
	"time"
)

type Cursor struct {
	db *Database `json:"-"`
	Id string    `json:"Id"`

	Index  int           `json:"-"`
	Result []interface{} `json:"result"`
	More   bool          `json:"hasMore"`
	Amount int           `json:"count"`
	Data   Extra         `json:"extra"`

	Err    bool   `json:"error"`
	ErrMsg string `json:"errorMessage"`
	Code   int    `json:"code"`
	max    int
	Time   time.Duration `json:"time"`
}

func NewCursor(db *Database) *Cursor {
	var c Cursor
	if db == nil {
		return nil
	}
	c.db = db
	return &c
}

// Delete cursor in server and free RAM
func (c *Cursor) Delete() error {
	if c.Id == "" {
		return errors.New("Invalid cursor to delete")
	}
	res, err := c.db.send("cursor", c.Id, "DELETE", nil, c, c)
	if err != nil {
		return nil
	}

	switch res.Status() {
	case 202:
		return nil
	case 404:
		return errors.New("Cursor does not exist")
	default:
		return nil
	}

}

func (c *Cursor) FetchBatch(r interface{}) error {
	kind := reflect.ValueOf(r).Elem().Kind()
	if kind != reflect.Slice && kind != reflect.Array {
		return errors.New("Container must be Slice of array kind")
	}
	b, err := json.Marshal(c.Result)
	if err != nil {
		return err
	}
	err = json.Unmarshal(b, r)
	if err != nil {
		return err
	}

	// fetch next batch
	if c.HasMore() {
		res, err := c.db.send("cursor", c.Id, "PUT", nil, c, c)

		if res.Status() == 200 {
			return nil
		}

		if err != nil {
			return err
		}
	}

	return nil
}

// Iterates over cursor, returns false when no more values into batch, fetch next batch if necesary.
func (c *Cursor) FetchOne(r interface{}) bool {
	var max int = len(c.Result) - 1

	if c.Index < max {
		b, err := json.Marshal(c.Result[c.Index])
		err = json.Unmarshal(b, r)
		c.Index++ // move to next value into result
		if err != nil {
			return false
		} else {
			return true
		}
	}
	// Last document
	if c.Index == max {
		b, err := json.Marshal(c.Result[c.Index])
		err = json.Unmarshal(b, r)
		if err != nil {
			return false
		} else {
			if c.More {
				//fetch rest from server
				res, _ := c.db.send("cursor", c.Id, "PUT", nil, c, c)

				if err != nil {
					return false
				}

				if res.Status() == 200 {
					c.Index = 0
					return true
				} else {
					return false
				}

			} else {
				// last doc
				return false
			}
		}
	}

	return false
}

// move cursor index by 1
func (c *Cursor) Next(r interface{}) bool {
	if c.Index == c.max {
		return false
	} else {
		c.Index++
		if c.Index == c.max {
			return true
		} else {
			return false
		}
	}
}

type Extra struct {
	FullCount int `json:"fullCount"`
}

func (c Cursor) Count() int {
	return c.Amount
}

func (c *Cursor) FullCount() int {
	return c.Data.FullCount
}

func (c Cursor) HasMore() bool {
	return c.More
}

func (c Cursor) Error() bool {
	return c.Err
}

func (c Cursor) ErrCode() int {
	return c.Code
}
<commit_msg>fix FetchOne<commit_after>package aranGO

import (
	"encoding/json"
	"errors"
	"log"
	"reflect"
	"time"
)

type Cursor struct {
	db *Database `json:"-"`
	Id string    `json:"Id"`

	Index  int           `json:"-"`
	Result []interface{} `json:"result"`
	More   bool          `json:"hasMore"`
	Amount int           `json:"count"`
	Data   Extra         `json:"extra"`

	Err    bool   `json:"error"`
	ErrMsg string `json:"errorMessage"`
	Code   int    `json:"code"`
	max    int
	Time   time.Duration `json:"time"`
}

func NewCursor(db *Database) *Cursor {
	var c Cursor
	if db == nil {
		return nil
	}
	c.db = db
	return &c
}

// Delete cursor in server and free RAM
func (c *Cursor) Delete() error {
	if c.Id == "" {
		return errors.New("Invalid cursor to delete")
	}
	res, err := c.db.send("cursor", c.Id, "DELETE", nil, c, c)
	if err != nil {
		return nil
	}

	switch res.Status() {
	case 202:
		return nil
	case 404:
		return errors.New("Cursor does not exist")
	default:
		return nil
	}

}

func (c *Cursor) FetchBatch(r interface{}) error {
	kind := reflect.ValueOf(r).Elem().Kind()
	if kind != reflect.Slice && kind != reflect.Array {
		return errors.New("Container must be Slice of array kind")
	}
	b, err := json.Marshal(c.Result)
	if err != nil {
		return err
	}
	err = json.Unmarshal(b, r)
	if err != nil {
		return err
	}

	// fetch next batch
	if c.HasMore() {
		res, err := c.db.send("cursor", c.Id, "PUT", nil, c, c)

		if res.Status() == 200 {
			return nil
		}

		if err != nil {
			return err
		}
	}

	return nil
}

// Iterates over cursor, returns false when no more values into batch, fetch next batch if necesary.
func (c *Cursor) FetchOne(r interface{}) bool {
	var max int = len(c.Result) - 1
	if c.Index >= max {
		b, err := json.Marshal(c.Result[c.Index])
		err = json.Unmarshal(b, r)
		if err != nil {
			return false
		} else {
			if c.More {
				//fetch rest from server
				res, _ := c.db.send("cursor", c.Id, "PUT", nil, c, c)

				if err != nil {
					return false
				}

				if res.Status() == 200 {
					c.Index = 0
					return true
				} else {
					return false
				}

			} else {
				// last doc
				return false
			}
		}
	} else {
		b, err := json.Marshal(c.Result[c.Index])
		err = json.Unmarshal(b, r)
		c.Index++ // move to next value into result
		if c.Index == max {
			log.Println("index--->", c.Index)
		}
		if err != nil {
			log.Println(err)
			return false
		} else {
			return true
		}
	}
}

// move cursor index by 1
func (c *Cursor) Next(r interface{}) bool {
	if c.Index == c.max {
		return false
	} else {
		c.Index++
		if c.Index == c.max {
			return true
		} else {
			return false
		}
	}
}

type Extra struct {
	FullCount int `json:"fullCount"`
}

func (c Cursor) Count() int {
	return c.Amount
}

func (c *Cursor) FullCount() int {
	return c.Data.FullCount
}

func (c Cursor) HasMore() bool {
	return c.More
}

func (c Cursor) Error() bool {
	return c.Err
}

func (c Cursor) ErrCode() int {
	return c.Code
}
<|endoftext|>


<commit_before>// Copyright 2016 The Goutils Author. All Rights Reserved.
//
// -------------------------------------------------------------------

package goutils

func ArrayStringContain(a []string, s string) bool {

	for _, v := range a {
		if v == s {
			return true
		}
	}

	return false
}

func ArrayIntContain(a []int, s int) bool {

	for _, v := range a {
		if v == s {
			return true
		}
	}

	return false
}
<commit_msg>array remove<commit_after>// Copyright 2016 The Goutils Author. All Rights Reserved.
//
// -------------------------------------------------------------------

package goutils

func ArrayStringContain(a []string, s string) bool {

	for _, v := range a {
		if v == s {
			return true
		}
	}

	return false
}

func arrayStringRemove(a []string, s string) []string {

	for i, v := range a {

		if v == s {
			a = append(a[:i], a[i+1:]...)
		}
	}

	return a
}

func ArrayIntContain(a []int, s int) bool {

	for _, v := range a {
		if v == s {
			return true
		}
	}

	return false
}

func arrayIntRemove(a []int, s int) []int {

	for i, v := range a {

		if v == s {
			a = append(a[:i], a[i+1:]...)
		}
	}

	return a
}
<|endoftext|>