<commit_before>/*
 * Copyright (c) 2014 Santiago Arias | Remy Jourde
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package models

import (
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"appengine"

	"github.com/santiaago/gonawin/helpers/log"
)

type Tday struct {
	Date    time.Time
	Matches []Tmatch
}

type Tphase struct {
	Name string
	Days []Tday
}

// ByDate implements sort.Interface for []Tday based on the date field.
type ByDate []Tday

func (a ByDate) Len() int           { return len(a) }
func (a ByDate) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByDate) Less(i, j int) bool { return a[i].Date.Before(a[j].Date) }

// Check if the match m passed as argument is the last match of a phase in a specific tournament.
// it returns a boolean and the index of the phase the match was found
func lastMatchOfPhase(c appengine.Context, m *Tmatch, phases *[]Tphase) (bool, int64) {

	for i, ph := range *phases {
		if n := len(ph.Days); n >= 1 {
			lastDay := ph.Days[n-1]
			if n = len(lastDay.Matches); n >= 1 {
				lastMatch := lastDay.Matches[n-1]
				if lastMatch.IdNumber == m.IdNumber {
					return true, int64(i)
				}
			}
		}
	}
	return false, int64(-1)
}

// UpdateNextPhase updates next phase in tournament.
//
func UpdateNextPhase(c appengine.Context, t *Tournament, currentphase *Tphase, nextphase *Tphase) error {

	// the array of phases that will be update.
	// it is an array as a phase can trigger an update in multiple phases, like semi-finals
	// trigger update of Third place and Finals
	var phases []*Tphase
	phases = append(phases, nextphase)
	// compute ranking of previous phase
	var mapOfTeams map[string]*Tteam
	mapOfTeams = make(map[string]*Tteam)

	if currentphase.Name == cFirstStage {
		// compute ranking of groups
		// get all groups.
		groups := Groups(c, t.GroupIds)
		for _, g := range groups {
			team1, argTeam1 := getFirstTeamInGroup(c, g)
			team2, _ := getSecondTeamInGroup(c, g, argTeam1)
			team3, _ := getNthTeamInGroup(c, g, 3)
			mapOfTeams["1"+g.Name] = team1
			mapOfTeams["2"+g.Name] = team2
			mapOfTeams["3"+g.Name] = team3
		}
	} else {
		// compute ranking just by match winners
		if currentphase.Name == cFinals || currentphase.Name == cThirdPlace {
			// nothing to do.
			return nil
		}

		currentmatches := GetMatchesByPhase(c, t, currentphase.Name)
		if currentphase.Name != cSemiFinals {
			log.Infof(c, "Update Next phase: current %v", currentphase.Name)
			log.Infof(c, "Update Next phase: next %v", nextphase.Name)

			for _, m := range currentmatches {
				// ToDo: handle penalties
				if m.Result1 >= m.Result2 {
					team1, _ := TTeamById(c, m.TeamId1)
					mapOfTeams["W"+strconv.Itoa(int(m.IdNumber))] = team1
					log.Infof(c, "Update Next phase: rule: W%v teams: %v", strconv.Itoa(int(m.IdNumber)), team1.Name)

				} else if m.Result1 < m.Result2 {
					team2, _ := TTeamById(c, m.TeamId2)
					mapOfTeams["W"+strconv.Itoa(int(m.IdNumber))] = team2
					log.Infof(c, "Update Next phase: rule: W%v teams: %v", strconv.Itoa(int(m.IdNumber)), team2.Name)
				}
			}
		} else {
			// append finals phases to array of phases to update.
			var finals Tphase
			finals.Name = cFinals
			phases = append(phases, &finals)

			for _, m := range currentmatches {
				// ToDo: handle penalties
				if m.Result1 >= m.Result2 {
					team1, _ := TTeamById(c, m.TeamId1)
					team2, _ := TTeamById(c, m.TeamId2)
					mapOfTeams["W"+strconv.Itoa(int(m.IdNumber))] = team1
					mapOfTeams["L"+strconv.Itoa(int(m.IdNumber))] = team2
					log.Infof(c, "Update Next phase: rule: W%v teams: %v", strconv.Itoa(int(m.IdNumber)), team1.Name)
					log.Infof(c, "Update Next phase: rule: L%v teams: %v", strconv.Itoa(int(m.IdNumber)), team2.Name)

				} else if m.Result1 < m.Result2 {
					team2, _ := TTeamById(c, m.TeamId2)
					team1, _ := TTeamById(c, m.TeamId1)
					mapOfTeams["W"+strconv.Itoa(int(m.IdNumber))] = team2
					mapOfTeams["L"+strconv.Itoa(int(m.IdNumber))] = team1
					log.Infof(c, "Update Next phase: rule: W%v teams: %v", strconv.Itoa(int(m.IdNumber)), team2.Name)
					log.Infof(c, "Update Next phase: rule: L%v teams: %v", strconv.Itoa(int(m.IdNumber)), team1.Name)
				}
			}

		}
	}

	// update phase (matches) with new teams
	for _, ph := range phases {
		matches := GetMatchesByPhase(c, t, ph.Name)
		for i, m := range matches {
			log.Infof(c, "Update Next phase: current rule: %v", m.Rule)
			rule := strings.Split(m.Rule, " ")
			if len(rule) != 2 {
				continue
			}

			if val, ok := mapOfTeams[rule[0]]; ok {
				log.Infof(c, "Update Next phase: match found: %v", val.Name)
				matches[i].TeamId1 = val.Id
			} else if len(rule[0]) == 3 {
				// try of get rule 3AD <=> the best between the 3A and 3D

				r1 := fmt.Sprintf("%s%s", rule[0][0], rule[0][1])
				r2 := fmt.Sprintf("%s%s", rule[0][0], rule[0][2])
				val1, ok1 := mapOfTeams[r1]
				val2, ok2 := mapOfTeams[r2]

				if ok1 && ok2 {
					// compare the two teams.
					groups := Groups(c, t.GroupIds)
					maxVal := getMaxTeam(groups, *val1, *val2)
					if maxVal != nil && maxVal.Id == val1.Id {
						log.Infof(c, "Update Next phase: match found: %v", maxVal.Name)
						matches[i].TeamId1 = maxVal.Id
					} else if rule[0][1] == 'B' && maxVal.Id == val2.Id {
						r3 := fmt.Sprintf("%s%s", rule[0][0], "A")
						val3, ok3 := mapOfTeams[r3]
						if ok3 {
							maxVal2 := getMaxTeam(groups, *val1, *val3)
							if maxVal2 != nil && maxVal2.Id == val1.Id {
								log.Infof(c, "Update Next phase: match found: %v", val1.Name)
								matches[i].TeamId1 = val1.Id
							} else if maxVal2.Id == val3.Id {
								log.Infof(c, "Update Next phase: match found: %v", maxVal.Name)
								matches[i].TeamId1 = maxVal.Id
							} else {
								return errors.New(fmt.Sprintf("Cannot parse rule in tournament =%d", t.Id))
							}
						}
					} else if rule[0][1] == 'A' && maxVal == val2 {
						r3 := fmt.Sprintf("%s%s", rule[0][0], "B")
						val3, ok3 := mapOfTeams[r3]
						if ok3 {
							maxVal2 := getMaxTeam(groups, *val1, *val3)
							if maxVal2 != nil && maxVal2.Id == val1.Id {
								log.Infof(c, "Update Next phase: match found: %v", val1.Name)
								matches[i].TeamId1 = val1.Id
							} else if maxVal2.Id == val3.Id {
								log.Infof(c, "Update Next phase: match found: %v", maxVal.Name)
								matches[i].TeamId1 = maxVal.Id
							} else {
								return errors.New(fmt.Sprintf("Cannot parse rule in tournament =%d", t.Id))
							}
						}
					} else {
						return errors.New(fmt.Sprintf("Cannot parse rule in tournament =%d", t.Id))
					}
				} else {
					return errors.New(fmt.Sprintf("Cannot parse rule in tournament =%d", t.Id))
				}

			} else {
				return errors.New(fmt.Sprintf("Cannot parse rule in tournament =%d", t.Id))
			}
			if val, ok := mapOfTeams[rule[1]]; ok {
				log.Infof(c, "Update Next phase: match found: %v", val.Name)
				matches[i].TeamId2 = val.Id
			} else if len(rule[0]) == 3 {
				// try of get rule 3AD <=> the best between the 3A and 3D

				r1 := fmt.Sprintf("%s%s", rule[1][0], rule[1][1])
				r2 := fmt.Sprintf("%s%s", rule[1][0], rule[1][2])
				val1, ok1 := mapOfTeams[r1]
				val2, ok2 := mapOfTeams[r2]

				if ok1 && ok2 {
					// compare the two teams.
					groups := Groups(c, t.GroupIds)
					maxVal := getMaxTeam(groups, *val1, *val2)
					if maxVal != nil {
						log.Infof(c, "Update Next phase: match found: %v", maxVal.Name)
						matches[i].TeamId1 = maxVal.Id
					} else {
						return errors.New(fmt.Sprintf("Cannot parse rule in tournament =%d", t.Id))
					}
				} else {
					return errors.New(fmt.Sprintf("Cannot parse rule in tournament =%d", t.Id))
				}
			} else {
				return errors.New(fmt.Sprintf("Cannot parse rule in tournament =%d", t.Id))
			}
			matches[i].Rule = ""
			matches[i].Ready = true
			matches[i].CanPredict = true
		}
		if err := UpdateMatches(c, matches); err != nil {
			log.Errorf(c, "Set Results: unable to set results on matches: %v", err)
			return err
		}
	}

	return nil
}
<commit_msg>forgot the other case.<commit_after>/*
 * Copyright (c) 2014 Santiago Arias | Remy Jourde
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package models

import (
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"appengine"

	"github.com/santiaago/gonawin/helpers/log"
)

type Tday struct {
	Date    time.Time
	Matches []Tmatch
}

type Tphase struct {
	Name string
	Days []Tday
}

// ByDate implements sort.Interface for []Tday based on the date field.
type ByDate []Tday

func (a ByDate) Len() int           { return len(a) }
func (a ByDate) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByDate) Less(i, j int) bool { return a[i].Date.Before(a[j].Date) }

// Check if the match m passed as argument is the last match of a phase in a specific tournament.
// it returns a boolean and the index of the phase the match was found
func lastMatchOfPhase(c appengine.Context, m *Tmatch, phases *[]Tphase) (bool, int64) {

	for i, ph := range *phases {
		if n := len(ph.Days); n >= 1 {
			lastDay := ph.Days[n-1]
			if n = len(lastDay.Matches); n >= 1 {
				lastMatch := lastDay.Matches[n-1]
				if lastMatch.IdNumber == m.IdNumber {
					return true, int64(i)
				}
			}
		}
	}
	return false, int64(-1)
}

// UpdateNextPhase updates next phase in tournament.
//
func UpdateNextPhase(c appengine.Context, t *Tournament, currentphase *Tphase, nextphase *Tphase) error {

	// the array of phases that will be update.
	// it is an array as a phase can trigger an update in multiple phases, like semi-finals
	// trigger update of Third place and Finals
	var phases []*Tphase
	phases = append(phases, nextphase)
	// compute ranking of previous phase
	var mapOfTeams map[string]*Tteam
	mapOfTeams = make(map[string]*Tteam)

	if currentphase.Name == cFirstStage {
		// compute ranking of groups
		// get all groups.
		groups := Groups(c, t.GroupIds)
		for _, g := range groups {
			team1, argTeam1 := getFirstTeamInGroup(c, g)
			team2, _ := getSecondTeamInGroup(c, g, argTeam1)
			team3, _ := getNthTeamInGroup(c, g, 3)
			mapOfTeams["1"+g.Name] = team1
			mapOfTeams["2"+g.Name] = team2
			mapOfTeams["3"+g.Name] = team3
		}
	} else {
		// compute ranking just by match winners
		if currentphase.Name == cFinals || currentphase.Name == cThirdPlace {
			// nothing to do.
			return nil
		}

		currentmatches := GetMatchesByPhase(c, t, currentphase.Name)
		if currentphase.Name != cSemiFinals {
			log.Infof(c, "Update Next phase: current %v", currentphase.Name)
			log.Infof(c, "Update Next phase: next %v", nextphase.Name)

			for _, m := range currentmatches {
				// ToDo: handle penalties
				if m.Result1 >= m.Result2 {
					team1, _ := TTeamById(c, m.TeamId1)
					mapOfTeams["W"+strconv.Itoa(int(m.IdNumber))] = team1
					log.Infof(c, "Update Next phase: rule: W%v teams: %v", strconv.Itoa(int(m.IdNumber)), team1.Name)

				} else if m.Result1 < m.Result2 {
					team2, _ := TTeamById(c, m.TeamId2)
					mapOfTeams["W"+strconv.Itoa(int(m.IdNumber))] = team2
					log.Infof(c, "Update Next phase: rule: W%v teams: %v", strconv.Itoa(int(m.IdNumber)), team2.Name)
				}
			}
		} else {
			// append finals phases to array of phases to update.
			var finals Tphase
			finals.Name = cFinals
			phases = append(phases, &finals)

			for _, m := range currentmatches {
				// ToDo: handle penalties
				if m.Result1 >= m.Result2 {
					team1, _ := TTeamById(c, m.TeamId1)
					team2, _ := TTeamById(c, m.TeamId2)
					mapOfTeams["W"+strconv.Itoa(int(m.IdNumber))] = team1
					mapOfTeams["L"+strconv.Itoa(int(m.IdNumber))] = team2
					log.Infof(c, "Update Next phase: rule: W%v teams: %v", strconv.Itoa(int(m.IdNumber)), team1.Name)
					log.Infof(c, "Update Next phase: rule: L%v teams: %v", strconv.Itoa(int(m.IdNumber)), team2.Name)

				} else if m.Result1 < m.Result2 {
					team2, _ := TTeamById(c, m.TeamId2)
					team1, _ := TTeamById(c, m.TeamId1)
					mapOfTeams["W"+strconv.Itoa(int(m.IdNumber))] = team2
					mapOfTeams["L"+strconv.Itoa(int(m.IdNumber))] = team1
					log.Infof(c, "Update Next phase: rule: W%v teams: %v", strconv.Itoa(int(m.IdNumber)), team2.Name)
					log.Infof(c, "Update Next phase: rule: L%v teams: %v", strconv.Itoa(int(m.IdNumber)), team1.Name)
				}
			}

		}
	}

	// update phase (matches) with new teams
	for _, ph := range phases {
		matches := GetMatchesByPhase(c, t, ph.Name)
		for i, m := range matches {
			log.Infof(c, "Update Next phase: current rule: %v", m.Rule)
			rule := strings.Split(m.Rule, " ")
			if len(rule) != 2 {
				continue
			}

			if val, ok := mapOfTeams[rule[0]]; ok {
				log.Infof(c, "Update Next phase: match found: %v", val.Name)
				matches[i].TeamId1 = val.Id
			} else if len(rule[0]) == 3 {
				// try of get rule 3AD <=> the best between the 3A and 3D

				r1 := fmt.Sprintf("%s%s", rule[0][0], rule[0][1])
				r2 := fmt.Sprintf("%s%s", rule[0][0], rule[0][2])
				val1, ok1 := mapOfTeams[r1]
				val2, ok2 := mapOfTeams[r2]

				if ok1 && ok2 {
					// compare the two teams.
					groups := Groups(c, t.GroupIds)
					maxVal := getMaxTeam(groups, *val1, *val2)
					if maxVal != nil && maxVal.Id == val1.Id {
						log.Infof(c, "Update Next phase: match found: %v", maxVal.Name)
						matches[i].TeamId1 = maxVal.Id
					} else if rule[0][1] == 'B' && maxVal.Id == val2.Id {
						r3 := fmt.Sprintf("%s%s", rule[0][0], "A")
						val3, ok3 := mapOfTeams[r3]
						if ok3 {
							maxVal2 := getMaxTeam(groups, *val1, *val3)
							if maxVal2 != nil && maxVal2.Id == val1.Id {
								log.Infof(c, "Update Next phase: match found: %v", val1.Name)
								matches[i].TeamId1 = val1.Id
							} else if maxVal2.Id == val3.Id {
								log.Infof(c, "Update Next phase: match found: %v", maxVal.Name)
								matches[i].TeamId1 = maxVal.Id
							} else {
								return errors.New(fmt.Sprintf("Cannot parse rule in tournament =%d", t.Id))
							}
						}
					} else if rule[0][1] == 'A' && maxVal == val2 {
						r3 := fmt.Sprintf("%s%s", rule[0][0], "B")
						val3, ok3 := mapOfTeams[r3]
						if ok3 {
							maxVal2 := getMaxTeam(groups, *val1, *val3)
							if maxVal2 != nil && maxVal2.Id == val1.Id {
								log.Infof(c, "Update Next phase: match found: %v", val1.Name)
								matches[i].TeamId1 = val1.Id
							} else if maxVal2.Id == val3.Id {
								log.Infof(c, "Update Next phase: match found: %v", maxVal.Name)
								matches[i].TeamId1 = maxVal.Id
							} else {
								return errors.New(fmt.Sprintf("Cannot parse rule in tournament =%d", t.Id))
							}
						}
					} else {
						return errors.New(fmt.Sprintf("Cannot parse rule in tournament =%d", t.Id))
					}
				} else {
					return errors.New(fmt.Sprintf("Cannot parse rule in tournament =%d", t.Id))
				}

			} else {
				return errors.New(fmt.Sprintf("Cannot parse rule in tournament =%d", t.Id))
			}
			if val, ok := mapOfTeams[rule[1]]; ok {
				log.Infof(c, "Update Next phase: match found: %v", val.Name)
				matches[i].TeamId2 = val.Id
			} else if len(rule[1]) == 3 {
				// try of get rule 3AD <=> the best between the 3A and 3D
				r1 := fmt.Sprintf("%s%s", rule[1][0], rule[1][1])
				r2 := fmt.Sprintf("%s%s", rule[1][0], rule[1][2])
				val1, ok1 := mapOfTeams[r1]
				val2, ok2 := mapOfTeams[r2]

				if ok1 && ok2 {
					// compare the two teams.
					groups := Groups(c, t.GroupIds)
					maxVal := getMaxTeam(groups, *val1, *val2)
					if maxVal != nil && maxVal.Id == val1.Id {
						log.Infof(c, "Update Next phase: match found: %v", maxVal.Name)
						matches[i].TeamId1 = maxVal.Id
					} else if rule[1][1] == 'B' && maxVal.Id == val2.Id {
						r3 := fmt.Sprintf("%s%s", rule[1][0], "A")
						val3, ok3 := mapOfTeams[r3]
						if ok3 {
							maxVal2 := getMaxTeam(groups, *val1, *val3)
							if maxVal2 != nil && maxVal2.Id == val1.Id {
								log.Infof(c, "Update Next phase: match found: %v", val1.Name)
								matches[i].TeamId1 = val1.Id
							} else if maxVal2.Id == val3.Id {
								log.Infof(c, "Update Next phase: match found: %v", maxVal.Name)
								matches[i].TeamId1 = maxVal.Id
							} else {
								return errors.New(fmt.Sprintf("Cannot parse rule in tournament =%d", t.Id))
							}
						}
					} else if rule[1][1] == 'A' && maxVal == val2 {
						r3 := fmt.Sprintf("%s%s", rule[1][0], "B")
						val3, ok3 := mapOfTeams[r3]
						if ok3 {
							maxVal2 := getMaxTeam(groups, *val1, *val3)
							if maxVal2 != nil && maxVal2.Id == val1.Id {
								log.Infof(c, "Update Next phase: match found: %v", val1.Name)
								matches[i].TeamId1 = val1.Id
							} else if maxVal2.Id == val3.Id {
								log.Infof(c, "Update Next phase: match found: %v", maxVal.Name)
								matches[i].TeamId1 = maxVal.Id
							} else {
								return errors.New(fmt.Sprintf("Cannot parse rule in tournament =%d", t.Id))
							}
						}
					} else {
						return errors.New(fmt.Sprintf("Cannot parse rule in tournament =%d", t.Id))
					}
				} else {
					return errors.New(fmt.Sprintf("Cannot parse rule in tournament =%d", t.Id))
				}
			} else {
				return errors.New(fmt.Sprintf("Cannot parse rule in tournament =%d", t.Id))
			}
			matches[i].Rule = ""
			matches[i].Ready = true
			matches[i].CanPredict = true
		}
		if err := UpdateMatches(c, matches); err != nil {
			log.Errorf(c, "Set Results: unable to set results on matches: %v", err)
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package es

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"path"
	"strings"
	"time"

	"github.com/grafana/grafana/pkg/components/simplejson"
	"github.com/grafana/grafana/pkg/log"
	"github.com/grafana/grafana/pkg/tsdb"

	"github.com/grafana/grafana/pkg/models"
	"golang.org/x/net/context/ctxhttp"
)

const loggerName = "tsdb.elasticsearch.client"

var (
	clientLog = log.New(loggerName)
)

// Client represents a client which can interact with elasticsearch api
type Client interface {
	GetVersion() int
	GetTimeField() string
	GetMinInterval(queryInterval string) (time.Duration, error)
	ExecuteMultisearch(r *MultiSearchRequest) (*MultiSearchResponse, error)
	MultiSearch() *MultiSearchRequestBuilder
}

// NewClient creates a new elasticsearch client
var NewClient = func(ctx context.Context, ds *models.DataSource, timeRange *tsdb.TimeRange) (Client, error) {
	version, err := ds.JsonData.Get("esVersion").Int()
	if err != nil {
		return nil, fmt.Errorf("eleasticsearch version is required, err=%v", err)
	}

	timeField, err := ds.JsonData.Get("timeField").String()
	if err != nil {
		return nil, fmt.Errorf("eleasticsearch time field name is required, err=%v", err)
	}

	indexInterval := ds.JsonData.Get("interval").MustString()
	ip, err := newIndexPattern(indexInterval, ds.Database)
	if err != nil {
		return nil, err
	}

	indices, err := ip.GetIndices(timeRange)
	if err != nil {
		return nil, err
	}

	bc := &baseClientImpl{
		ctx:       ctx,
		ds:        ds,
		version:   version,
		timeField: timeField,
		indices:   indices,
	}

	clientLog.Debug("Creating new client", "version", version, "timeField", timeField, "indices", strings.Join(indices, ", "))

	switch version {
	case 2:
		return newV2Client(bc)
	case 5:
		return newV5Client(bc)
	case 56:
		return newV56Client(bc)
	}

	return nil, fmt.Errorf("elasticsearch version=%d is not supported", version)
}

type baseClient interface {
	Client
	getSettings() *simplejson.Json
	executeBatchRequest(uriPath string, requests []*multiRequest) (*http.Response, error)
	executeRequest(method, uriPath string, body []byte) (*http.Response, error)
	createMultiSearchRequests(searchRequests []*SearchRequest) []*multiRequest
}

type baseClientImpl struct {
	ctx       context.Context
	ds        *models.DataSource
	version   int
	timeField string
	indices   []string
}

func (c *baseClientImpl) GetVersion() int {
	return c.version
}

func (c *baseClientImpl) GetTimeField() string {
	return c.timeField
}

func (c *baseClientImpl) GetMinInterval(queryInterval string) (time.Duration, error) {
	return tsdb.GetIntervalFrom(c.ds, simplejson.NewFromAny(map[string]string{
		"interval": queryInterval,
	}), 15*time.Second)
}

func (c *baseClientImpl) getSettings() *simplejson.Json {
	return c.ds.JsonData
}

type multiRequest struct {
	header map[string]interface{}
	body   interface{}
}

func (c *baseClientImpl) executeBatchRequest(uriPath string, requests []*multiRequest) (*http.Response, error) {
	payload := bytes.Buffer{}
	for _, r := range requests {
		reqHeader, err := json.Marshal(r.header)
		if err != nil {
			return nil, err
		}
		payload.WriteString(string(reqHeader) + "\n")

		reqBody, err := json.Marshal(r.body)
		if err != nil {
			return nil, err
		}
		payload.WriteString(string(reqBody) + "\n")
	}

	return c.executeRequest(http.MethodPost, uriPath, payload.Bytes())
}

func (c *baseClientImpl) executeRequest(method, uriPath string, body []byte) (*http.Response, error) {
	u, _ := url.Parse(c.ds.Url)
	u.Path = path.Join(u.Path, uriPath)

	var req *http.Request
	var err error
	if method == http.MethodPost {
		req, err = http.NewRequest(http.MethodPost, u.String(), bytes.NewBuffer(body))
	} else {
		req, err = http.NewRequest(http.MethodGet, u.String(), nil)
	}
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", "Grafana")
	req.Header.Set("Content-Type", "application/json")

	if c.ds.BasicAuth {
		clientLog.Debug("Request configured to use basic authentication")
		req.SetBasicAuth(c.ds.BasicAuthUser, c.ds.BasicAuthPassword)
	}

	if !c.ds.BasicAuth && c.ds.User != "" {
		clientLog.Debug("Request configured to use basic authentication")
		req.SetBasicAuth(c.ds.User, c.ds.Password)
	}

	httpClient, err := c.ds.GetHttpClient()
	if err != nil {
		return nil, err
	}

	if method == http.MethodPost {
		clientLog.Debug("Executing request", "url", req.URL.String(), "method", method)
	} else {
		clientLog.Debug("Executing request", "url", req.URL.String(), "method", method)
	}

	return ctxhttp.Do(c.ctx, httpClient, req)
}

func (c *baseClientImpl) ExecuteMultisearch(r *MultiSearchRequest) (*MultiSearchResponse, error) {
	multiRequests := c.createMultiSearchRequests(r.Requests)
	res, err := c.executeBatchRequest("_msearch", multiRequests)
	if err != nil {
		return nil, err
	}

	var msr MultiSearchResponse
	defer res.Body.Close()
	dec := json.NewDecoder(res.Body)
	err = dec.Decode(&msr)
	if err != nil {
		return nil, err
	}

	clientLog.Debug("Received multisearch response", "code", res.StatusCode, "status", res.Status, "content-length", res.ContentLength)

	msr.status = res.StatusCode

	return &msr, nil
}

func (c *baseClientImpl) createMultiSearchRequests(searchRequests []*SearchRequest) []*multiRequest {
	multiRequests := []*multiRequest{}

	for _, searchReq := range searchRequests {
		multiRequests = append(multiRequests, &multiRequest{
			header: map[string]interface{}{
				"search_type":        "query_then_fetch",
				"ignore_unavailable": true,
				"index":              strings.Join(c.indices, ","),
			},
			body: searchReq,
		})
	}

	return multiRequests
}

type v2Client struct {
	baseClient
}

func newV2Client(bc baseClient) (*v2Client, error) {
	c := v2Client{
		baseClient: bc,
	}

	return &c, nil
}

func (c *v2Client) createMultiSearchRequests(searchRequests []*SearchRequest) []*multiRequest {
	multiRequests := c.baseClient.createMultiSearchRequests(searchRequests)

	for _, mr := range multiRequests {
		mr.header["search_type"] = "count"
	}

	return multiRequests
}

type v5Client struct {
	baseClient
}

func newV5Client(bc baseClient) (*v5Client, error) {
	c := v5Client{
		baseClient: bc,
	}

	return &c, nil
}

type v56Client struct {
	*v5Client
	maxConcurrentShardRequests int
}

func newV56Client(bc baseClient) (*v56Client, error) {
	v5Client := v5Client{
		baseClient: bc,
	}
	maxConcurrentShardRequests := bc.getSettings().Get("maxConcurrentShardRequests").MustInt(256)

	c := v56Client{
		v5Client:                   &v5Client,
		maxConcurrentShardRequests: maxConcurrentShardRequests,
	}

	return &c, nil
}

func (c *v56Client) createMultiSearchRequests(searchRequests []*SearchRequest) []*multiRequest {
	multiRequests := c.v5Client.createMultiSearchRequests(searchRequests)

	for _, mr := range multiRequests {
		mr.header["max_concurrent_shard_requests"] = c.maxConcurrentShardRequests
	}

	return multiRequests
}

func (c *baseClientImpl) MultiSearch() *MultiSearchRequestBuilder {
	return NewMultiSearchRequestBuilder(c.GetVersion())
}
<commit_msg>elasticsearch: default interval fix<commit_after>package es

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"path"
	"strings"
	"time"

	"github.com/grafana/grafana/pkg/components/simplejson"
	"github.com/grafana/grafana/pkg/log"
	"github.com/grafana/grafana/pkg/tsdb"

	"github.com/grafana/grafana/pkg/models"
	"golang.org/x/net/context/ctxhttp"
)

const loggerName = "tsdb.elasticsearch.client"

var (
	clientLog = log.New(loggerName)
)

// Client represents a client which can interact with elasticsearch api
type Client interface {
	GetVersion() int
	GetTimeField() string
	GetMinInterval(queryInterval string) (time.Duration, error)
	ExecuteMultisearch(r *MultiSearchRequest) (*MultiSearchResponse, error)
	MultiSearch() *MultiSearchRequestBuilder
}

// NewClient creates a new elasticsearch client
var NewClient = func(ctx context.Context, ds *models.DataSource, timeRange *tsdb.TimeRange) (Client, error) {
	version, err := ds.JsonData.Get("esVersion").Int()
	if err != nil {
		return nil, fmt.Errorf("eleasticsearch version is required, err=%v", err)
	}

	timeField, err := ds.JsonData.Get("timeField").String()
	if err != nil {
		return nil, fmt.Errorf("eleasticsearch time field name is required, err=%v", err)
	}

	indexInterval := ds.JsonData.Get("interval").MustString()
	ip, err := newIndexPattern(indexInterval, ds.Database)
	if err != nil {
		return nil, err
	}

	indices, err := ip.GetIndices(timeRange)
	if err != nil {
		return nil, err
	}

	bc := &baseClientImpl{
		ctx:       ctx,
		ds:        ds,
		version:   version,
		timeField: timeField,
		indices:   indices,
	}

	clientLog.Debug("Creating new client", "version", version, "timeField", timeField, "indices", strings.Join(indices, ", "))

	switch version {
	case 2:
		return newV2Client(bc)
	case 5:
		return newV5Client(bc)
	case 56:
		return newV56Client(bc)
	}

	return nil, fmt.Errorf("elasticsearch version=%d is not supported", version)
}

type baseClient interface {
	Client
	getSettings() *simplejson.Json
	executeBatchRequest(uriPath string, requests []*multiRequest) (*http.Response, error)
	executeRequest(method, uriPath string, body []byte) (*http.Response, error)
	createMultiSearchRequests(searchRequests []*SearchRequest) []*multiRequest
}

type baseClientImpl struct {
	ctx       context.Context
	ds        *models.DataSource
	version   int
	timeField string
	indices   []string
}

func (c *baseClientImpl) GetVersion() int {
	return c.version
}

func (c *baseClientImpl) GetTimeField() string {
	return c.timeField
}

func (c *baseClientImpl) GetMinInterval(queryInterval string) (time.Duration, error) {
	return tsdb.GetIntervalFrom(c.ds, simplejson.NewFromAny(map[string]string{
		"interval": queryInterval,
	}), 5*time.Second)
}

func (c *baseClientImpl) getSettings() *simplejson.Json {
	return c.ds.JsonData
}

type multiRequest struct {
	header map[string]interface{}
	body   interface{}
}

func (c *baseClientImpl) executeBatchRequest(uriPath string, requests []*multiRequest) (*http.Response, error) {
	payload := bytes.Buffer{}
	for _, r := range requests {
		reqHeader, err := json.Marshal(r.header)
		if err != nil {
			return nil, err
		}
		payload.WriteString(string(reqHeader) + "\n")

		reqBody, err := json.Marshal(r.body)
		if err != nil {
			return nil, err
		}
		payload.WriteString(string(reqBody) + "\n")
	}

	return c.executeRequest(http.MethodPost, uriPath, payload.Bytes())
}

func (c *baseClientImpl) executeRequest(method, uriPath string, body []byte) (*http.Response, error) {
	u, _ := url.Parse(c.ds.Url)
	u.Path = path.Join(u.Path, uriPath)

	var req *http.Request
	var err error
	if method == http.MethodPost {
		req, err = http.NewRequest(http.MethodPost, u.String(), bytes.NewBuffer(body))
	} else {
		req, err = http.NewRequest(http.MethodGet, u.String(), nil)
	}
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", "Grafana")
	req.Header.Set("Content-Type", "application/json")

	if c.ds.BasicAuth {
		clientLog.Debug("Request configured to use basic authentication")
		req.SetBasicAuth(c.ds.BasicAuthUser, c.ds.BasicAuthPassword)
	}

	if !c.ds.BasicAuth && c.ds.User != "" {
		clientLog.Debug("Request configured to use basic authentication")
		req.SetBasicAuth(c.ds.User, c.ds.Password)
	}

	httpClient, err := c.ds.GetHttpClient()
	if err != nil {
		return nil, err
	}

	if method == http.MethodPost {
		clientLog.Debug("Executing request", "url", req.URL.String(), "method", method)
	} else {
		clientLog.Debug("Executing request", "url", req.URL.String(), "method", method)
	}

	return ctxhttp.Do(c.ctx, httpClient, req)
}

func (c *baseClientImpl) ExecuteMultisearch(r *MultiSearchRequest) (*MultiSearchResponse, error) {
	multiRequests := c.createMultiSearchRequests(r.Requests)
	res, err := c.executeBatchRequest("_msearch", multiRequests)
	if err != nil {
		return nil, err
	}

	var msr MultiSearchResponse
	defer res.Body.Close()
	dec := json.NewDecoder(res.Body)
	err = dec.Decode(&msr)
	if err != nil {
		return nil, err
	}

	clientLog.Debug("Received multisearch response", "code", res.StatusCode, "status", res.Status, "content-length", res.ContentLength)

	msr.status = res.StatusCode

	return &msr, nil
}

func (c *baseClientImpl) createMultiSearchRequests(searchRequests []*SearchRequest) []*multiRequest {
	multiRequests := []*multiRequest{}

	for _, searchReq := range searchRequests {
		multiRequests = append(multiRequests, &multiRequest{
			header: map[string]interface{}{
				"search_type":        "query_then_fetch",
				"ignore_unavailable": true,
				"index":              strings.Join(c.indices, ","),
			},
			body: searchReq,
		})
	}

	return multiRequests
}

type v2Client struct {
	baseClient
}

func newV2Client(bc baseClient) (*v2Client, error) {
	c := v2Client{
		baseClient: bc,
	}

	return &c, nil
}

func (c *v2Client) createMultiSearchRequests(searchRequests []*SearchRequest) []*multiRequest {
	multiRequests := c.baseClient.createMultiSearchRequests(searchRequests)

	for _, mr := range multiRequests {
		mr.header["search_type"] = "count"
	}

	return multiRequests
}

type v5Client struct {
	baseClient
}

func newV5Client(bc baseClient) (*v5Client, error) {
	c := v5Client{
		baseClient: bc,
	}

	return &c, nil
}

type v56Client struct {
	*v5Client
	maxConcurrentShardRequests int
}

func newV56Client(bc baseClient) (*v56Client, error) {
	v5Client := v5Client{
		baseClient: bc,
	}
	maxConcurrentShardRequests := bc.getSettings().Get("maxConcurrentShardRequests").MustInt(256)

	c := v56Client{
		v5Client:                   &v5Client,
		maxConcurrentShardRequests: maxConcurrentShardRequests,
	}

	return &c, nil
}

func (c *v56Client) createMultiSearchRequests(searchRequests []*SearchRequest) []*multiRequest {
	multiRequests := c.v5Client.createMultiSearchRequests(searchRequests)

	for _, mr := range multiRequests {
		mr.header["max_concurrent_shard_requests"] = c.maxConcurrentShardRequests
	}

	return multiRequests
}

func (c *baseClientImpl) MultiSearch() *MultiSearchRequestBuilder {
	return NewMultiSearchRequestBuilder(c.GetVersion())
}
<|endoftext|>


<commit_before>// Package base64 provides a byte slice type that marshals into json as
// a raw (no padding) base64url value.
package base64

import (
	"encoding/base64"
	"errors"
	"reflect"
)

// Value is a base64url encoded json object,
type Value []byte

// NewValue returns a pointer to a Value, cast from the given byte slice.
// This is a convenience function, handling the address creation that a direct
// cast would not allow.
func NewValue(b []byte) *Value {
	v := Value(b)
	return &v
}

// MarshalJSON returns the ba64url encoding of bv for JSON representation.
func (bv *Value) MarshalJSON() ([]byte, error) {
	return []byte("\"" + base64.RawURLEncoding.EncodeToString(*bv) + "\""), nil
}

// UnmarshalJSON sets bv to the bytes represented in the base64url encoding b.
func (bv *Value) UnmarshalJSON(b []byte) error {
	if len(b) < 2 || b[0] != byte('"') || b[len(b)-1] != byte('"') {
		return errors.New("value is not a string")
	}

	out := make([]byte, base64.RawURLEncoding.DecodedLen(len(b)-2))
	n, err := base64.RawURLEncoding.Decode(out, b[1:len(b)-1])
	if err != nil {
		return err
	}

	v := reflect.ValueOf(bv).Elem()
	v.SetBytes(out[:n])
	return nil
}
<commit_msg>Add String method to Base64 Package<commit_after>// Package base64 provides a byte slice type that marshals into json as
// a raw (no padding) base64url value.
package base64

import (
	"encoding/base64"
	"errors"
	"reflect"
)

// Value is a base64url encoded json object,
type Value []byte

// NewValue returns a pointer to a Value, cast from the given byte slice.
// This is a convenience function, handling the address creation that a direct
// cast would not allow.
func NewValue(b []byte) *Value {
	v := Value(b)
	return &v
}

// MarshalJSON returns the ba64url encoding of bv for JSON representation.
func (bv *Value) MarshalJSON() ([]byte, error) {
	return []byte("\"" + base64.RawURLEncoding.EncodeToString(*bv) + "\""), nil
}

func (bv *Value) String() string {
	return base64.RawURLEncoding.EncodeToString(*bv)
}

// UnmarshalJSON sets bv to the bytes represented in the base64url encoding b.
func (bv *Value) UnmarshalJSON(b []byte) error {
	if len(b) < 2 || b[0] != byte('"') || b[len(b)-1] != byte('"') {
		return errors.New("value is not a string")
	}

	out := make([]byte, base64.RawURLEncoding.DecodedLen(len(b)-2))
	n, err := base64.RawURLEncoding.Decode(out, b[1:len(b)-1])
	if err != nil {
		return err
	}

	v := reflect.ValueOf(bv).Elem()
	v.SetBytes(out[:n])
	return nil
}
<|endoftext|>


<commit_before>package api

import (
	"fmt"
	"math/big"
	"net/http"
	"strings"

	"github.com/NebulousLabs/Sia/types"
)

// WalletSiafundsBalance contains fields relating to the siafunds balance.
type WalletSiafundsBalance struct {
	SiafundBalance      types.Currency
	SiacoinClaimBalance types.Currency
}

// scanAmount scans a types.Currency.
func scanAmount(amountStr string) (amount types.Currency, err error) {
	// exponential format
	if strings.ContainsAny(amountStr, "Ee") {
		amountRat := new(big.Rat)
		_, err = fmt.Sscan(amountStr, amountRat)
		if err != nil {
			return
		}
		amount = types.NewCurrency(new(big.Int).Div(amountRat.Num(), amountRat.Denom()))
		return
	}

	// standard format
	_, err = fmt.Sscan(amountStr, &amount)
	return
}

// scanAddres scans a types.UnlockHash.
func scanAddress(addrStr string) (addr types.UnlockHash, err error) {
	err = addr.LoadString(addrStr)
	return
}

// walletAddressHandler handles the API request for a new address.
func (srv *Server) walletAddressHandler(w http.ResponseWriter, req *http.Request) {
	coinAddress, _, err := srv.wallet.CoinAddress(true) // true indicates that the address should be visible to the user
	if err != nil {
		writeError(w, "Failed to get a coin address", http.StatusInternalServerError)
		return
	}

	// Since coinAddress is not a struct, we define one here so that writeJSON
	// writes an object instead of a bare value. In addition, we transmit the
	// coinAddress as a hex-encoded string rather than a byte array.
	writeJSON(w, struct{ Address types.UnlockHash }{coinAddress})
}

// walletSendHandler handles the API call to send coins to another address.
func (srv *Server) walletSendHandler(w http.ResponseWriter, req *http.Request) {
	// Scan the amount.
	amount, err := scanAmount(req.FormValue("amount"))
	if err != nil {
		writeError(w, "Malformed amount", http.StatusBadRequest)
		return
	}

	// Scan the destination address.
	dest, err := scanAddress(req.FormValue("destination"))
	if err != nil {
		writeError(w, "Malformed coin address", http.StatusBadRequest)
		return
	}

	// Send the coins.
	_, err = srv.wallet.SendCoins(amount, dest)
	if err != nil {
		writeError(w, "Failed to create transaction: "+err.Error(), http.StatusInternalServerError)
		return
	}

	writeSuccess(w)
}

// walletSiafundsBalanceHandler handles the API call querying the balance of
// siafunds.
func (srv *Server) walletSiafundsBalanceHandler(w http.ResponseWriter, req *http.Request) {
	var wsb WalletSiafundsBalance
	wsb.SiafundBalance, wsb.SiacoinClaimBalance = srv.wallet.SiafundBalance()
	writeJSON(w, wsb)
}

// walletSiafundsSendHandler handles the API request to send siafunds.
func (srv *Server) walletSiafundsSendHandler(w http.ResponseWriter, req *http.Request) {
	// Scan the amount.
	amount, err := scanAmount(req.FormValue("amount"))
	if err != nil {
		writeError(w, "Malformed amount", http.StatusBadRequest)
		return
	}

	// Scan the destination address.
	dest, err := scanAddress(req.FormValue("destination"))
	if err != nil {
		writeError(w, "Malformed coin address", http.StatusBadRequest)
		return
	}

	// Scan the keyfile list.
	keyfiles := strings.Split(req.FormValue("keyfiles"), ",")
	if len(keyfiles) == 0 {
		writeError(w, "Missing keyfiles", http.StatusBadRequest)
		return
	}

	_, err = srv.wallet.SendSiagSiafunds(amount, dest, keyfiles)
	if err != nil {
		writeError(w, "Failed to send siafunds: "+err.Error(), http.StatusInternalServerError)
		return
	}

	writeSuccess(w)
}

// walletSiafundsWatchsiagaddressHandler handles the API request to watch a
// siag address.
func (srv *Server) walletSiafundsWatchsiagaddressHandler(w http.ResponseWriter, req *http.Request) {
	err := srv.wallet.WatchSiagSiafundAddress(req.FormValue("keyfile"))
	if err != nil {
		writeError(w, err.Error(), http.StatusBadRequest)
		return
	}
	writeSuccess(w)
}

// walletStatusHandler handles the API call querying the status of the wallet.
func (srv *Server) walletStatusHandler(w http.ResponseWriter, req *http.Request) {
	writeJSON(w, srv.wallet.Info())
}
<commit_msg>make scanAmount safer<commit_after>package api

import (
	"math/big"
	"net/http"
	"strings"

	"github.com/NebulousLabs/Sia/types"
)

// WalletSiafundsBalance contains fields relating to the siafunds balance.
type WalletSiafundsBalance struct {
	SiafundBalance      types.Currency
	SiacoinClaimBalance types.Currency
}

// scanAmount scans a types.Currency.
func scanAmount(amount string) (types.Currency, bool) {
	// use SetString manually to ensure that amount does not contain
	// multiple values, which would confuse fmt.Scan
	i, ok := new(big.Int).SetString(amount, 10)
	if !ok {
		return types.Currency{}, ok
	}
	return types.NewCurrency(i), true
}

// scanAddres scans a types.UnlockHash.
func scanAddress(addrStr string) (addr types.UnlockHash, err error) {
	err = addr.LoadString(addrStr)
	return
}

// walletAddressHandler handles the API request for a new address.
func (srv *Server) walletAddressHandler(w http.ResponseWriter, req *http.Request) {
	coinAddress, _, err := srv.wallet.CoinAddress(true) // true indicates that the address should be visible to the user
	if err != nil {
		writeError(w, "Failed to get a coin address", http.StatusInternalServerError)
		return
	}

	// Since coinAddress is not a struct, we define one here so that writeJSON
	// writes an object instead of a bare value. In addition, we transmit the
	// coinAddress as a hex-encoded string rather than a byte array.
	writeJSON(w, struct{ Address types.UnlockHash }{coinAddress})
}

// walletSendHandler handles the API call to send coins to another address.
func (srv *Server) walletSendHandler(w http.ResponseWriter, req *http.Request) {
	// Scan the amount.
	amount, ok := scanAmount(req.FormValue("amount"))
	if !ok {
		writeError(w, "Malformed amount", http.StatusBadRequest)
		return
	}

	// Scan the destination address.
	dest, err := scanAddress(req.FormValue("destination"))
	if err != nil {
		writeError(w, "Malformed coin address", http.StatusBadRequest)
		return
	}

	// Send the coins.
	_, err = srv.wallet.SendCoins(amount, dest)
	if err != nil {
		writeError(w, "Failed to create transaction: "+err.Error(), http.StatusInternalServerError)
		return
	}

	writeSuccess(w)
}

// walletSiafundsBalanceHandler handles the API call querying the balance of
// siafunds.
func (srv *Server) walletSiafundsBalanceHandler(w http.ResponseWriter, req *http.Request) {
	var wsb WalletSiafundsBalance
	wsb.SiafundBalance, wsb.SiacoinClaimBalance = srv.wallet.SiafundBalance()
	writeJSON(w, wsb)
}

// walletSiafundsSendHandler handles the API request to send siafunds.
func (srv *Server) walletSiafundsSendHandler(w http.ResponseWriter, req *http.Request) {
	// Scan the amount.
	amount, ok := scanAmount(req.FormValue("amount"))
	if !ok {
		writeError(w, "Malformed amount", http.StatusBadRequest)
		return
	}

	// Scan the destination address.
	dest, err := scanAddress(req.FormValue("destination"))
	if err != nil {
		writeError(w, "Malformed coin address", http.StatusBadRequest)
		return
	}

	// Scan the keyfile list.
	keyfiles := strings.Split(req.FormValue("keyfiles"), ",")
	if len(keyfiles) == 0 {
		writeError(w, "Missing keyfiles", http.StatusBadRequest)
		return
	}

	_, err = srv.wallet.SendSiagSiafunds(amount, dest, keyfiles)
	if err != nil {
		writeError(w, "Failed to send siafunds: "+err.Error(), http.StatusInternalServerError)
		return
	}

	writeSuccess(w)
}

// walletSiafundsWatchsiagaddressHandler handles the API request to watch a
// siag address.
func (srv *Server) walletSiafundsWatchsiagaddressHandler(w http.ResponseWriter, req *http.Request) {
	err := srv.wallet.WatchSiagSiafundAddress(req.FormValue("keyfile"))
	if err != nil {
		writeError(w, err.Error(), http.StatusBadRequest)
		return
	}
	writeSuccess(w)
}

// walletStatusHandler handles the API call querying the status of the wallet.
func (srv *Server) walletStatusHandler(w http.ResponseWriter, req *http.Request) {
	writeJSON(w, srv.wallet.Info())
}
<|endoftext|>


<commit_before>package paymill

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
)

const APIScheme string = "https"
const APIBase string = "api.paymill.com"
const APIVersion string = "v2"

type ApiClient struct {
	Token string
}

func NewApiClient(token string) (c *ApiClient) {
	if strings.Trim(token, " ") == "" {
		return nil
	}

	c = &ApiClient{
		Token: token,
	}
	return
}

func UrlFor(entity string, data url.Values) url.URL {
	url := url.URL{Scheme: APIScheme, Host: APIBase}
	path := fmt.Sprintf("/%s/%s", APIVersion, entity)

	url.Path = path
	url.RawQuery = data.Encode()

	return url
}

func (c *ApiClient) doRequest(resource string, method string, urlData url.Values, formData url.Values) (resp *http.Response, body []byte) {
	http_client := &http.Client{}

	// This can be wrapped in a method
	var req *http.Request
	url := UrlFor(resource, urlData)
	req, err := http.NewRequest(method, url.String(), strings.NewReader(formData.Encode()))
	if err != nil {
		panic(err)
	}

	if method == "POST" {
		req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	}
	req.SetBasicAuth(c.Token, "")

	resp, err = http_client.Do(req)
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()

	body, err = ioutil.ReadAll(resp.Body)

	if err != nil {
		panic(err)
	}

	return
}
<commit_msg>Load sha256, so we can interact with the new PAymill certificate.<commit_after>package paymill

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
)

import _ "crypto/sha512"

const APIScheme string = "https"
const APIBase string = "api.paymill.com"
const APIVersion string = "v2"

type ApiClient struct {
	Token string
}

func NewApiClient(token string) (c *ApiClient) {
	if strings.Trim(token, " ") == "" {
		return nil
	}

	c = &ApiClient{
		Token: token,
	}
	return
}

func UrlFor(entity string, data url.Values) url.URL {
	url := url.URL{Scheme: APIScheme, Host: APIBase}
	path := fmt.Sprintf("/%s/%s", APIVersion, entity)

	url.Path = path
	url.RawQuery = data.Encode()

	return url
}

func (c *ApiClient) doRequest(resource string, method string, urlData url.Values, formData url.Values) (resp *http.Response, body []byte) {
	http_client := &http.Client{}

	// This can be wrapped in a method
	var req *http.Request
	url := UrlFor(resource, urlData)
	req, err := http.NewRequest(method, url.String(), strings.NewReader(formData.Encode()))
	if err != nil {
		panic(err)
	}

	if method == "POST" {
		req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	}
	req.SetBasicAuth(c.Token, "")

	resp, err = http_client.Do(req)
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()

	body, err = ioutil.ReadAll(resp.Body)

	if err != nil {
		panic(err)
	}

	return
}
<|endoftext|>


<commit_before>// Package transport provides an net/http Transport in a single place.
// HTTP(S) timeout can be set with a flag.
// It is a separate package so libs and plugins use it instead of net/http.
// That way, in case HTTP(S) needs to go through a proxy or else, only this
// file needs to be edited.
// This also allows us to use our own TLS package, also extensible.
package transport

import (
	"flag"
	"net"
	"net/http"
	neturl "net/url"
	"time"

	"github.com/StalkR/goircbot/lib/tls"
)

var timeout = flag.Duration("http_timeout", 10*time.Second, "Timeout for HTTP(S) connections.")

// ByURL returns an transport given an URL for use with a client.
// It extracts the hostname from URL and give it to TLS for SNI.
func ByURL(url string) (*http.Transport, error) {
	u, err := neturl.Parse(url)
	if err != nil {
		return nil, err
	}
	return &http.Transport{
		Dial:            timeoutDialer(*timeout),
		TLSClientConfig: tls.Config(u.Host),
	}, nil
}

// Client returns a client given an URL.
func Client(url string) (*http.Client, error) {
	trans, err := ByURL(url)
	if err != nil {
		return nil, err
	}
	return &http.Client{Transport: trans}, nil
}

func timeoutDialer(d time.Duration) func(net, addr string) (net.Conn, error) {
	return func(netw, addr string) (net.Conn, error) {
		return net.DialTimeout(netw, addr, d)
	}
}
<commit_msg>lib/transport: force IPv4 on YouTube as they carelessly block wide IPv6 ranges<commit_after>// Package transport provides an net/http Transport in a single place.
// HTTP(S) timeout can be set with a flag.
// It is a separate package so libs and plugins use it instead of net/http.
// That way, in case HTTP(S) needs to go through a proxy or else, only this
// file needs to be edited.
// This also allows us to use our own TLS package, also extensible.
package transport

import (
	"flag"
	"net"
	"net/http"
	neturl "net/url"
	"strings"
	"time"

	"github.com/StalkR/goircbot/lib/tls"
)

var timeout = flag.Duration("http_timeout", 10*time.Second, "Timeout for HTTP(S) connections.")

// ByURL returns an transport given an URL for use with a client.
// It extracts the hostname from URL and give it to TLS for SNI.
func ByURL(url string) (*http.Transport, error) {
	u, err := neturl.Parse(url)
	if err != nil {
		return nil, err
	}
	return &http.Transport{
		Dial:            timeoutDialer(*timeout),
		TLSClientConfig: tls.Config(u.Host),
	}, nil
}

// Client returns a client given an URL.
func Client(url string) (*http.Client, error) {
	trans, err := ByURL(url)
	if err != nil {
		return nil, err
	}
	return &http.Client{Transport: trans}, nil
}

func timeoutDialer(d time.Duration) func(net, addr string) (net.Conn, error) {
	return func(netw, addr string) (net.Conn, error) {
		host := strings.Split(addr, ":")[0]
		// force IPv4 for YouTube as they carelessly block wide IPv6 ranges
		if host == "youtube.com" || host == "www.youtube.com" {
			netw = "tcp4"
		}
		return net.DialTimeout(netw, addr, d)
	}
}
<|endoftext|>


<commit_before>//go:build linux
// +build linux

package net

import (
	"errors"
	"fmt"
	"net"
	gonetURL "net/url"
	"os"
	"strings"

	bosherr "github.com/cloudfoundry/bosh-utils/errors"
	// NOTE: "cgroups is only intended to be used/compiled on linux based system"
	// see: https://github.com/containerd/cgroups/issues/19
	"github.com/containerd/cgroups"
	"github.com/opencontainers/runtime-spec/specs-go"

	"github.com/coreos/go-iptables/iptables"
)

const (
	/* "natsIsolationClassID" This is the integer value of the argument "0xb0540002", which is
	   b054:0002 . The major number (the left-hand side) is "BOSH", leet-ified.
	   The minor number (the right-hand side) is 2, indicating that this is the
	   second thing in our "BOSH" classid namespace.

	   _Hopefully_ noone uses a major number of "b054", and we avoid collisions _forever_!
	   If you need to select new classids for firewall rules or traffic control rules, keep
	   the major number "b054" for bosh stuff, unless there's a good reason to not.

	   The net_cls.classid structure is described in more detail here:
	   https://www.kernel.org/doc/Documentation/cgroup-v1/net_cls.txt
	*/
	natsIsolationClassID uint32 = 2958295042
)

// SetupNatsFirewall will setup the outgoing cgroup based rule that prevents everything except the agent to open connections to the nats api
func SetupNatsFirewall(mbus string) error {
	// return early if
	// we get a https url for mbus. case for create-env
	// we get an empty string. case for http_metadata_service (responsible to extract the agent-settings.json from the metadata endpoint)
	// we find that v1cgroups are not mounted (warden stemcells)
	if mbus == "" || strings.HasPrefix(mbus, "https://") {
		return nil
	}
	_, err := cgroups.V1()
	if err != nil {
		if errors.Is(err, cgroups.ErrMountPointNotExist) {
			return nil // v1cgroups are not mounted (warden stemcells)
		}
		return bosherr.WrapError(err, "Error retrieving cgroups mount point")
	}
	mbusURL, err := gonetURL.Parse(mbus)
	if err != nil || mbusURL.Hostname() == "" {
		return bosherr.WrapError(err, "Error parsing MbusURL")
	}

	host, port, err := net.SplitHostPort(mbusURL.Host)
	if err != nil {
		return bosherr.WrapError(err, "Error getting Port")
	}
	ipt, err := iptables.New()
	if err != nil {
		return bosherr.WrapError(err, "Iptables Error")
	}
	exists, err := ipt.Exists("mangle", "POSTROUTING",
		"-d", "127.0.0.1",
		"-p", "tcp",
		"--dport", "2822",
		"-m", "cgroup",
		"--cgroup", fmt.Sprintf("%v", natsIsolationClassID),
		"-j", "ACCEPT",
	)
	if err != nil {
		return bosherr.WrapError(err, "Iptables Error checking for monit rule")
	}
	if !exists {
		err = ipt.Insert("mangle", "POSTROUTING", 1,
			"-d", "127.0.0.1",
			"-p", "tcp",
			"--dport", "2822",
			"-m", "cgroup",
			"--cgroup", fmt.Sprintf("%v", natsIsolationClassID),
			"-j", "ACCEPT",
		)
		if err != nil {
			return bosherr.WrapError(err, "Iptables Error inserting for monit rule")
		}
	}
	err = ipt.AppendUnique("mangle", "POSTROUTING",
		"-d", host,
		"-p", "tcp",
		"--dport", port,
		"-m", "cgroup",
		"--cgroup", fmt.Sprintf("%v", natsIsolationClassID),
		"-j", "ACCEPT",
	)
	if err != nil {
		return bosherr.WrapError(err, "Iptables Error inserting for agent ACCEPT rule")
	}
	err = ipt.AppendUnique("mangle", "POSTROUTING",
		"-d", host,
		"-p", "tcp",
		"--dport", port,
		"-j", "DROP",
	)
	if err != nil {
		return bosherr.WrapError(err, "Iptables Error inserting for non-agent DROP rule")
	}
	var isolationClassID = natsIsolationClassID
	natsAPICgroup, err := cgroups.New(cgroups.SingleSubsystem(cgroups.V1, cgroups.NetCLS), cgroups.StaticPath("/nats-api-access"), &specs.LinuxResources{
		Network: &specs.LinuxNetwork{
			ClassID: &isolationClassID,
		},
	})
	if err != nil {
		return bosherr.WrapError(err, "Error setting up cgroups for nats api access")
	}

	err = natsAPICgroup.AddProc(uint64(os.Getpid()), cgroups.NetCLS)
	return err
}
<commit_msg>Add V6 iptables support<commit_after>//go:build linux
// +build linux

package net

import (
	"errors"
	"fmt"
	"net"
	gonetURL "net/url"
	"os"
	"strings"

	bosherr "github.com/cloudfoundry/bosh-utils/errors"
	// NOTE: "cgroups is only intended to be used/compiled on linux based system"
	// see: https://github.com/containerd/cgroups/issues/19
	"github.com/containerd/cgroups"
	"github.com/opencontainers/runtime-spec/specs-go"

	"github.com/coreos/go-iptables/iptables"
)

const (
	/* "natsIsolationClassID" This is the integer value of the argument "0xb0540002", which is
	   b054:0002 . The major number (the left-hand side) is "BOSH", leet-ified.
	   The minor number (the right-hand side) is 2, indicating that this is the
	   second thing in our "BOSH" classid namespace.

	   _Hopefully_ noone uses a major number of "b054", and we avoid collisions _forever_!
	   If you need to select new classids for firewall rules or traffic control rules, keep
	   the major number "b054" for bosh stuff, unless there's a good reason to not.

	   The net_cls.classid structure is described in more detail here:
	   https://www.kernel.org/doc/Documentation/cgroup-v1/net_cls.txt
	*/
	natsIsolationClassID uint32 = 2958295042
)

// SetupNatsFirewall will setup the outgoing cgroup based rule that prevents everything except the agent to open connections to the nats api
func SetupNatsFirewall(mbus string) error {
	// return early if
	// we get a https url for mbus. case for create-env
	// we get an empty string. case for http_metadata_service (responsible to extract the agent-settings.json from the metadata endpoint)
	// we find that v1cgroups are not mounted (warden stemcells)
	if mbus == "" || strings.HasPrefix(mbus, "https://") {
		return nil
	}
	_, err := cgroups.V1()
	if err != nil {
		if errors.Is(err, cgroups.ErrMountPointNotExist) {
			return nil // v1cgroups are not mounted (warden stemcells)
		}
		return bosherr.WrapError(err, "Error retrieving cgroups mount point")
	}
	mbusURL, err := gonetURL.Parse(mbus)
	if err != nil || mbusURL.Hostname() == "" {
		return bosherr.WrapError(err, "Error parsing MbusURL")
	}

	host, port, err := net.SplitHostPort(mbusURL.Host)
	if err != nil {
		return bosherr.WrapError(err, "Error getting Port")
	}

	// Run the lookup for Host as it could be potentially a Hostname | IPv4 | IPv6
	// the return for LookupIP will be a list of IP Addr and in case of the Input being an IP Addr,
	// it will only contain one element with the Input IP
	addr_array, err := net.LookupIP(host)
	if err != nil {
		return bosherr.WrapError(err, fmt.Sprintf("Error resolving mbus host: %v", host))
	}

	ipt, err := iptables.New()
	if err != nil {
		return bosherr.WrapError(err, "Creating Iptables Error")
	}
	// Even on a V6 VM, Monit will listen to only V4 loopback
	// First create Monit V4 rules for natsIsolationClassID
	exists, err := ipt.Exists("mangle", "POSTROUTING",
		"-d", "127.0.0.1",
		"-p", "tcp",
		"--dport", "2822",
		"-m", "cgroup",
		"--cgroup", fmt.Sprintf("%v", natsIsolationClassID),
		"-j", "ACCEPT",
	)
	if err != nil {
		return bosherr.WrapError(err, "Iptables Error checking for monit rule")
	}
	if !exists {
		err = ipt.Insert("mangle", "POSTROUTING", 1,
			"-d", "127.0.0.1",
			"-p", "tcp",
			"--dport", "2822",
			"-m", "cgroup",
			"--cgroup", fmt.Sprintf("%v", natsIsolationClassID),
			"-j", "ACCEPT",
		)
		if err != nil {
			return bosherr.WrapError(err, "Iptables Error inserting for monit rule")
		}
	}

	// For nats iptables rules we default to V4 unless below dns resolution gives us a V6 target
	ipVersion := iptables.ProtocolIPv4
	// Check if we're dealing with a V4 Target
	if addr_array[0].To4() == nil {
		ipVersion = iptables.ProtocolIPv6
	}
	ipt, err = iptables.NewWithProtocol(ipVersion)
	if err != nil {
		return bosherr.WrapError(err, "Creating Iptables Error")
	}

	err = ipt.AppendUnique("mangle", "POSTROUTING",
		"-d", host,
		"-p", "tcp",
		"--dport", port,
		"-m", "cgroup",
		"--cgroup", fmt.Sprintf("%v", natsIsolationClassID),
		"-j", "ACCEPT",
	)
	if err != nil {
		return bosherr.WrapError(err, "Iptables Error inserting for agent ACCEPT rule")
	}
	err = ipt.AppendUnique("mangle", "POSTROUTING",
		"-d", host,
		"-p", "tcp",
		"--dport", port,
		"-j", "DROP",
	)
	if err != nil {
		return bosherr.WrapError(err, "Iptables Error inserting for non-agent DROP rule")
	}

	var isolationClassID = natsIsolationClassID
	natsAPICgroup, err := cgroups.New(cgroups.SingleSubsystem(cgroups.V1, cgroups.NetCLS), cgroups.StaticPath("/nats-api-access"), &specs.LinuxResources{
		Network: &specs.LinuxNetwork{
			ClassID: &isolationClassID,
		},
	})
	if err != nil {
		return bosherr.WrapError(err, "Error setting up cgroups for nats api access")
	}

	err = natsAPICgroup.AddProc(uint64(os.Getpid()), cgroups.NetCLS)
	return err
}
<|endoftext|>


<commit_before>package ach

import (
	"testing"
)

func mockBatchCCDHeader() *BatchHeader {
	bh := NewBatchHeader()
	bh.ServiceClassCode = 220
	bh.StandardEntryClassCode = "CCD"
	bh.CompanyName = "Your Company, inc"
	bh.CompanyIdentification = "121042882"
	bh.CompanyEntryDescription = "Vndr Pay"
	bh.ODFIIdentification = "121042882"
	return bh
}

func mockCCDEntryDetail() *EntryDetail {
	entry := NewEntryDetail()
	entry.TransactionCode = 27
	entry.SetRDFI("231380104")
	entry.DFIAccountNumber = "744-5678-99"
	entry.Amount = 5000000
	entry.IdentificationNumber = "location #23"
	entry.SetReceivingCompany("Best Co. #23")
	entry.SetTraceNumber(mockBatchCCDHeader().ODFIIdentification, 1)
	entry.DiscretionaryData = "S"
	return entry
}

func mockBatchCCD() *BatchCCD {
	mockBatch := NewBatchCCD(mockBatchCCDHeader())
	mockBatch.AddEntry(mockCCDEntryDetail())
	mockBatch.GetEntries()[0].AddAddenda(mockAddenda05())
	if err := mockBatch.Create(); err != nil {
		panic(err)
	}
	return mockBatch
}

// testBatchCCDHeader creates a batch CCD header
func testBatchCCDHeader(t testing.TB) {
	batch, _ := NewBatch(mockBatchCCDHeader())
	_, ok := batch.(*BatchCCD)
	if !ok {
		t.Error("Expecting BatchCCD")
	}
}

// TestBatchCCDHeader tests creating a batch CCD header
func TestBatchCCDHeader(t *testing.T) {
	testBatchCCDHeader(t)
}

// BenchmarkBatchCCDHeader benchmark creating a batch CCD header
func BenchmarkBatchCCDHeader(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		testBatchCCDHeader(b)
	}
}

// testBatchCCDAddendumCount batch control CCD can only have one addendum per entry detail
func testBatchCCDAddendumCount(t testing.TB) {
	mockBatch := mockBatchCCD()
	// Adding a second addenda to the mock entry
	mockBatch.GetEntries()[0].AddAddenda(mockAddenda05())
	if err := mockBatch.Validate(); err != nil {
		if e, ok := err.(*BatchError); ok {
			if e.FieldName != "EntryAddendaCount" {
				t.Errorf("%T: %s", err, err)
			}
		} else {
			t.Errorf("%T: %s", err, err)
		}
	}
}

// TestBatchCCDAddendumCount tests batch control CCD can only have one addendum per entry detail
func TestBatchCCDAddendumCount(t *testing.T) {
	testBatchCCDAddendumCount(t)
}

// BenchmarkBatchCCDAddendumCoun benchmarks batch control CCD can only have one addendum per entry detail
func BenchmarkBatchCCDAddendumCount(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		testBatchCCDAddendumCount(b)
	}
}

// testBatchCCDReceivingCompanyName validates Receiving company / Individual name is a mandatory field
func testBatchCCDReceivingCompanyName(t testing.TB) {
	mockBatch := mockBatchCCD()
	// modify the Individual name / receiving company to nothing
	mockBatch.GetEntries()[0].SetReceivingCompany("")
	if err := mockBatch.Validate(); err != nil {
		if e, ok := err.(*BatchError); ok {
			if e.FieldName != "IndividualName" {
				t.Errorf("%T: %s", err, err)
			}
		} else {
			t.Errorf("%T: %s", err, err)
		}
	}
}

// TestBatchCCDReceivingCompanyName tests validating receiving company / Individual name is a mandatory field
func TestBatchCCDReceivingCompanyName(t *testing.T) {
	testBatchCCDReceivingCompanyName(t)
}

// BenchmarkBatchCCDReceivingCompanyName benchmarks validating receiving company / Individual name is a mandatory field
func BenchmarkBatchCCDReceivingCompanyName(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		testBatchCCDReceivingCompanyName(b)
	}
}

// testBatchCCDAddendaTypeCode verifies addenda type code is 05
func testBatchCCDAddendaTypeCode(t testing.TB) {
	mockBatch := mockBatchCCD()
	mockBatch.GetEntries()[0].Addendum[0].(*Addenda05).typeCode = "07"
	if err := mockBatch.Validate(); err != nil {
		if e, ok := err.(*BatchError); ok {
			if e.FieldName != "TypeCode" {
				t.Errorf("%T: %s", err, err)
			}
		} else {
			t.Errorf("%T: %s", err, err)
		}
	}
}

// TestBatchCCDAddendaTypeCode tests verifying addenda type code is 05
func TestBatchCCDAddendaTypeCode(t *testing.T) {
	testBatchCCDAddendaTypeCode(t)
}

// BenchmarkBatchCCDAddendaTypeCod benchmarks verifying addenda type code is 05
func BenchmarkBatchCCDAddendaTypeCode(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		testBatchCCDAddendaTypeCode(b)
	}
}

// testBatchCCDSEC verifies that the standard entry class code is CCD for batchCCD
func testBatchCCDSEC(t testing.TB) {
	mockBatch := mockBatchCCD()
	mockBatch.header.StandardEntryClassCode = "RCK"
	if err := mockBatch.Validate(); err != nil {
		if e, ok := err.(*BatchError); ok {
			if e.FieldName != "StandardEntryClassCode" {
				t.Errorf("%T: %s", err, err)
			}
		} else {
			t.Errorf("%T: %s", err, err)
		}
	}
}

// TestBatchCCDSEC test verifying that the standard entry class code is CCD for batchCCD
func TestBatchCCDSEC(t *testing.T) {
	testBatchCCDSEC(t)
}

// BenchmarkBatchCCDSEC benchmarks verifying that the standard entry class code is CCD for batchCCD
func BenchmarkBatchCCDSEC(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		testBatchCCDSEC(b)
	}
}

// testBatchCCDAddendaCount verifies batch CCD addenda count
func testBatchCCDAddendaCount(t testing.TB) {
	mockBatch := mockBatchCCD()
	mockBatch.GetEntries()[0].AddAddenda(mockAddenda05())
	mockBatch.Create()
	if err := mockBatch.Validate(); err != nil {
		if e, ok := err.(*BatchError); ok {
			if e.FieldName != "AddendaCount" {
				t.Errorf("%T: %s", err, err)
			}
		} else {
			t.Errorf("%T: %s", err, err)
		}
	}
}

// TestBatchCCDAddendaCount tests verifying batch CCD addenda count
func TestBatchCCDAddendaCount(t *testing.T) {
	testBatchCCDAddendaCount(t)
}

// BenchmarkBatchCCDAddendaCount benchmarks verifying batch CCD addenda count
func BenchmarkBatchCCDAddendaCount(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		testBatchCCDAddendaCount(b)
	}
}

// testBatchCCDCreate creates a batch CCD
func testBatchCCDCreate(t testing.TB) {
	mockBatch := mockBatchCCD()
	// Batch Header information is required to Create a batch.
	mockBatch.GetHeader().ServiceClassCode = 0
	mockBatch.Create()
	if err := mockBatch.Validate(); err != nil {
		if e, ok := err.(*BatchError); ok {
			if e.FieldName != "ServiceClassCode" {
				t.Errorf("%T: %s", err, err)
			}
		} else {
			t.Errorf("%T: %s", err, err)
		}
	}
}

// TestBatchCCDCreate Test creating a batch CCD
func TestBatchCCDCreate(t *testing.T) {
	testBatchCCDCreate(t)
}

// BenchmarkBatchCCDCreate benchmark creating a batch CCD
func BenchmarkBatchCCDCreate(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		testBatchCCDCreate(b)
	}
}
<commit_msg>#166 batchCCD_test benchmark<commit_after>package ach

import (
	"testing"
)

// mockBatchCCDHeader creates a CCD batch header
func mockBatchCCDHeader() *BatchHeader {
	bh := NewBatchHeader()
	bh.ServiceClassCode = 220
	bh.StandardEntryClassCode = "CCD"
	bh.CompanyName = "Your Company, inc"
	bh.CompanyIdentification = "121042882"
	bh.CompanyEntryDescription = "Vndr Pay"
	bh.ODFIIdentification = "121042882"
	return bh
}

// mockCCDEntryDetail creates a CCD entry detail
func mockCCDEntryDetail() *EntryDetail {
	entry := NewEntryDetail()
	entry.TransactionCode = 27
	entry.SetRDFI("231380104")
	entry.DFIAccountNumber = "744-5678-99"
	entry.Amount = 5000000
	entry.IdentificationNumber = "location #23"
	entry.SetReceivingCompany("Best Co. #23")
	entry.SetTraceNumber(mockBatchCCDHeader().ODFIIdentification, 1)
	entry.DiscretionaryData = "S"
	return entry
}

// mockBatchCCD creates a CCD batch
func mockBatchCCD() *BatchCCD {
	mockBatch := NewBatchCCD(mockBatchCCDHeader())
	mockBatch.AddEntry(mockCCDEntryDetail())
	mockBatch.GetEntries()[0].AddAddenda(mockAddenda05())
	if err := mockBatch.Create(); err != nil {
		panic(err)
	}
	return mockBatch
}

// testBatchCCDHeader creates a CCD batch header
func testBatchCCDHeader(t testing.TB) {
	batch, _ := NewBatch(mockBatchCCDHeader())
	_, ok := batch.(*BatchCCD)
	if !ok {
		t.Error("Expecting BatchCCD")
	}
}

// TestBatchCCDHeader tests creating a CCD batch header
func TestBatchCCDHeader(t *testing.T) {
	testBatchCCDHeader(t)
}

// BenchmarkBatchCCDHeader benchmark creating a CCD batch header
func BenchmarkBatchCCDHeader(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		testBatchCCDHeader(b)
	}
}

// testBatchCCDAddendumCount batch control CCD can only have one addendum per entry detail
func testBatchCCDAddendumCount(t testing.TB) {
	mockBatch := mockBatchCCD()
	// Adding a second addenda to the mock entry
	mockBatch.GetEntries()[0].AddAddenda(mockAddenda05())
	if err := mockBatch.Validate(); err != nil {
		if e, ok := err.(*BatchError); ok {
			if e.FieldName != "EntryAddendaCount" {
				t.Errorf("%T: %s", err, err)
			}
		} else {
			t.Errorf("%T: %s", err, err)
		}
	}
}

// TestBatchCCDAddendumCount tests batch control CCD can only have one addendum per entry detail
func TestBatchCCDAddendumCount(t *testing.T) {
	testBatchCCDAddendumCount(t)
}

// BenchmarkBatchCCDAddendumCoun benchmarks batch control CCD can only have one addendum per entry detail
func BenchmarkBatchCCDAddendumCount(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		testBatchCCDAddendumCount(b)
	}
}

// testBatchCCDReceivingCompanyName validates Receiving company / Individual name is a mandatory field
func testBatchCCDReceivingCompanyName(t testing.TB) {
	mockBatch := mockBatchCCD()
	// modify the Individual name / receiving company to nothing
	mockBatch.GetEntries()[0].SetReceivingCompany("")
	if err := mockBatch.Validate(); err != nil {
		if e, ok := err.(*BatchError); ok {
			if e.FieldName != "IndividualName" {
				t.Errorf("%T: %s", err, err)
			}
		} else {
			t.Errorf("%T: %s", err, err)
		}
	}
}

// TestBatchCCDReceivingCompanyName tests validating receiving company / Individual name is a mandatory field
func TestBatchCCDReceivingCompanyName(t *testing.T) {
	testBatchCCDReceivingCompanyName(t)
}

// BenchmarkBatchCCDReceivingCompanyName benchmarks validating receiving company / Individual name is a mandatory field
func BenchmarkBatchCCDReceivingCompanyName(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		testBatchCCDReceivingCompanyName(b)
	}
}

// testBatchCCDAddendaTypeCode validates addenda type code is 05
func testBatchCCDAddendaTypeCode(t testing.TB) {
	mockBatch := mockBatchCCD()
	mockBatch.GetEntries()[0].Addendum[0].(*Addenda05).typeCode = "07"
	if err := mockBatch.Validate(); err != nil {
		if e, ok := err.(*BatchError); ok {
			if e.FieldName != "TypeCode" {
				t.Errorf("%T: %s", err, err)
			}
		} else {
			t.Errorf("%T: %s", err, err)
		}
	}
}

// TestBatchCCDAddendaTypeCode tests validating addenda type code is 05
func TestBatchCCDAddendaTypeCode(t *testing.T) {
	testBatchCCDAddendaTypeCode(t)
}

// BenchmarkBatchCCDAddendaTypeCod benchmarks validating addenda type code is 05
func BenchmarkBatchCCDAddendaTypeCode(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		testBatchCCDAddendaTypeCode(b)
	}
}

// testBatchCCDSEC validates that the standard entry class code is CCD for batchCCD
func testBatchCCDSEC(t testing.TB) {
	mockBatch := mockBatchCCD()
	mockBatch.header.StandardEntryClassCode = "RCK"
	if err := mockBatch.Validate(); err != nil {
		if e, ok := err.(*BatchError); ok {
			if e.FieldName != "StandardEntryClassCode" {
				t.Errorf("%T: %s", err, err)
			}
		} else {
			t.Errorf("%T: %s", err, err)
		}
	}
}

// TestBatchCCDSEC tests validating that the standard entry class code is CCD for batchCCD
func TestBatchCCDSEC(t *testing.T) {
	testBatchCCDSEC(t)
}

// BenchmarkBatchCCDSEC benchmarks validating that the standard entry class code is CCD for batchCCD
func BenchmarkBatchCCDSEC(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		testBatchCCDSEC(b)
	}
}

// testBatchCCDAddendaCount validates batch CCD addenda count
func testBatchCCDAddendaCount(t testing.TB) {
	mockBatch := mockBatchCCD()
	mockBatch.GetEntries()[0].AddAddenda(mockAddenda05())
	mockBatch.Create()
	if err := mockBatch.Validate(); err != nil {
		if e, ok := err.(*BatchError); ok {
			if e.FieldName != "AddendaCount" {
				t.Errorf("%T: %s", err, err)
			}
		} else {
			t.Errorf("%T: %s", err, err)
		}
	}
}

// TestBatchCCDAddendaCount tests validating batch CCD addenda count
func TestBatchCCDAddendaCount(t *testing.T) {
	testBatchCCDAddendaCount(t)
}

// BenchmarkBatchCCDAddendaCount benchmarks validating batch CCD addenda count
func BenchmarkBatchCCDAddendaCount(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		testBatchCCDAddendaCount(b)
	}
}

// testBatchCCDCreate creates a batch CCD
func testBatchCCDCreate(t testing.TB) {
	mockBatch := mockBatchCCD()
	// Batch Header information is required to Create a batch.
	mockBatch.GetHeader().ServiceClassCode = 0
	mockBatch.Create()
	if err := mockBatch.Validate(); err != nil {
		if e, ok := err.(*BatchError); ok {
			if e.FieldName != "ServiceClassCode" {
				t.Errorf("%T: %s", err, err)
			}
		} else {
			t.Errorf("%T: %s", err, err)
		}
	}
}

// TestBatchCCDCreate Test creating a batch CCD
func TestBatchCCDCreate(t *testing.T) {
	testBatchCCDCreate(t)
}

// BenchmarkBatchCCDCreate benchmark creating a batch CCD
func BenchmarkBatchCCDCreate(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		testBatchCCDCreate(b)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	pb "github.com/viru/berrybot/proto"

	log "github.com/Sirupsen/logrus"
	"github.com/kidoman/embd"
	_ "github.com/kidoman/embd/host/rpi" // This loads the RPi driver
	"google.golang.org/grpc"
)

// server is used to implement hellowrld.GreeterServer.
type server struct {
	front, rear *echo
	driver      *driver
}

type echo struct {
	name       string
	echo       embd.DigitalPin
	trig       embd.DigitalPin
	quit, done chan bool
	dist       int64
	last       time.Time
	enabled    bool
}

func newEcho(name string, trigPin, echoPin int) (*echo, error) {
	var e echo
	e.name = name
	e.quit = make(chan bool)
	e.done = make(chan bool)
	var err error
	e.trig, err = embd.NewDigitalPin(trigPin)
	if err != nil {
		return nil, fmt.Errorf("can't init trigger pin: %v", err)
	}
	e.echo, err = embd.NewDigitalPin(echoPin)
	if err != nil {
		return nil, fmt.Errorf("can't init echo pin: %v", err)
	}

	// Set direction.
	if err := e.trig.SetDirection(embd.Out); err != nil {
		return nil, fmt.Errorf("can't set trigger direction: %v", err)
	}
	if err := e.echo.SetDirection(embd.In); err != nil {
		return nil, fmt.Errorf("can't set echo direction: %v", err)
	}
	return &e, nil
}

func (e *echo) runDistancer() {
	if err := e.trig.Write(embd.Low); err != nil {
		log.Warnf("can't set trigger to low: %v", err)
	}
	time.Sleep(time.Second * 1)
	tick := time.NewTicker(time.Millisecond * 500)
	defer tick.Stop()
	for {
		select {
		case <-e.quit:
			e.done <- true
			return
		case <-tick.C:
			if !e.enabled {
				continue
			}
			log.Infof("%s: measuring...", e.name)
			if err := e.trig.Write(embd.High); err != nil {
				log.Warnf("can't set trigger to high: %v", err)
			}
			time.Sleep(time.Microsecond * 10)
			if err := e.trig.Write(embd.Low); err != nil {
				log.Warnf("can't set trigger to low: %v", err)
			}
			dur, err := e.echo.TimePulse(embd.High)
			if err != nil {
				log.Warnf("can't time pulse: %v", err)
			}
			log.Infof("%s: distance: %dcm", e.name, dur.Nanoseconds()/1000*34/1000/2)
			e.dist = dur.Nanoseconds() / 1000 * 34 / 1000 / 2
		}
	}
}

func (e *echo) close() {
	e.quit <- true
	<-e.done
	e.echo.Close()
	e.trig.Close()
}

type driver struct {
	left, right *engine
	mu          sync.Mutex
	moving      bool
	last        time.Time
}

func (d *driver) safetyStop() {
	ticker := time.NewTicker(time.Second)
	for range ticker.C {
		d.mu.Lock()
		if d.moving && d.last.Add(time.Second).Before(time.Now()) {
			d.mu.Unlock()
			d.stop()
			log.Warn("Emergency stop!")
			continue
		}
		d.mu.Unlock()
	}
}

func (d *driver) setMoving(moving bool) {
	d.mu.Lock()
	d.last = time.Now()
	d.moving = moving
	d.mu.Unlock()
}

func (d *driver) stop() {
	d.left.pwr.Write(embd.Low)
	d.right.pwr.Write(embd.Low)
	d.setMoving(false)
}

func (d *driver) forward() {
	d.left.pwr.Write(embd.High)
	d.left.fwd.Write(embd.High)
	d.right.pwr.Write(embd.High)
	d.right.fwd.Write(embd.High)
	d.setMoving(true)
}

func (d *driver) backward() {
	d.left.pwr.Write(embd.High)
	d.left.fwd.Write(embd.Low)
	d.right.pwr.Write(embd.High)
	d.right.fwd.Write(embd.Low)
	d.setMoving(true)
}

func (d *driver) sharpRight() {
	d.left.pwr.Write(embd.High)
	d.left.fwd.Write(embd.High)
	d.right.pwr.Write(embd.High)
	d.right.fwd.Write(embd.Low)
	d.setMoving(true)
}

func (d *driver) sharpLeft() {
	d.left.pwr.Write(embd.High)
	d.left.fwd.Write(embd.Low)
	d.right.pwr.Write(embd.High)
	d.right.fwd.Write(embd.High)
	d.setMoving(true)
}

func (d *driver) fwdRight() {
	d.left.pwr.Write(embd.High)
	d.left.fwd.Write(embd.High)
	d.right.pwr.Write(embd.Low)
	d.right.fwd.Write(embd.High)
	d.setMoving(true)
}

func (d *driver) fwdLeft() {
	d.left.pwr.Write(embd.Low)
	d.left.fwd.Write(embd.High)
	d.right.pwr.Write(embd.High)
	d.right.fwd.Write(embd.High)
	d.setMoving(true)
}

func (d *driver) backRight() {
	d.left.pwr.Write(embd.High)
	d.left.fwd.Write(embd.Low)
	d.right.pwr.Write(embd.Low)
	d.right.fwd.Write(embd.Low)
	d.setMoving(true)
}

func (d *driver) backLeft() {
	d.left.pwr.Write(embd.Low)
	d.left.fwd.Write(embd.Low)
	d.right.pwr.Write(embd.High)
	d.right.fwd.Write(embd.Low)
	d.setMoving(true)
}

const (
	safeStraightDist = 20
	safeTurningDist  = 10
)

func (s *server) drive(dir *pb.Direction) {
	switch {
	case dir.Dy > -5 && dir.Dy < 5 && dir.Dx > -5 && dir.Dx < 5:
		s.front.enabled = false
		s.rear.enabled = false
		s.driver.stop()
	case dir.Dy > 5 && dir.Dx > -5 && dir.Dx < 5:
		s.front.enabled = true
		s.driver.forward()
	case dir.Dy < -5 && dir.Dx > -5 && dir.Dx < 5:
		s.rear.enabled = true
		s.driver.backward()
	case dir.Dx > 5 && dir.Dy > -5 && dir.Dy < 5:
		s.driver.sharpRight()
	case dir.Dx < -5 && dir.Dy > -5 && dir.Dy < 5:
		s.driver.sharpLeft()
	case dir.Dx > 5 && dir.Dy > 5:
		s.driver.fwdRight()
	case dir.Dx < -5 && dir.Dy > 5:
		s.driver.fwdLeft()
	case dir.Dx > 5 && dir.Dy < -5:
		s.driver.backRight()
	case dir.Dx < -5 && dir.Dy < -5:
		s.driver.backLeft()
	}
}

type engine struct {
	fwd, pwr embd.DigitalPin
}

func newEngine(pwrPin, fwdPin int) (*engine, error) {
	var e engine
	var err error
	e.pwr, err = embd.NewDigitalPin(pwrPin)
	if err != nil {
		return nil, fmt.Errorf("can't init power pin: %v", err)
	}
	e.fwd, err = embd.NewDigitalPin(fwdPin)
	if err != nil {
		return nil, fmt.Errorf("can't init forward pin: %v", err)
	}

	// Set direction.
	if err := e.pwr.SetDirection(embd.Out); err != nil {
		return nil, fmt.Errorf("can't set power direction: %v", err)
	}
	if err := e.fwd.SetDirection(embd.Out); err != nil {
		return nil, fmt.Errorf("can't set forward direction: %v", err)
	}
	return &e, nil
}

func (e *engine) close() {
	e.pwr.Close()
	e.fwd.Close()
}

const (
	sensorUnknown = iota
	sensorFront
	sensorRear
)

func (s *server) Drive(stream pb.Driver_DriveServer) error {
	waitc := make(chan struct{})
	go func() {
		for {
			d, err := stream.Recv()
			if err != nil {
				log.Warnf("ERR from client: %v", err)
				close(waitc)
				return
			}
			s.drive(d)
		}
	}()

	for {
		select {
		case <-time.After(time.Second):
			if !s.front.enabled && !s.rear.enabled {
				continue
			}
			var speed int32
			if s.driver.moving {
				speed = 100
			}
			if err := stream.Send(&pb.Telemetry{Speed: speed, DistFront: int32(s.front.dist), DistRear: int32(s.rear.dist)}); err != nil {
				log.Errorf("can't send telemetry: %v", err)
				return err
			}
			log.Info("Sending telemetry!")
		case <-waitc:
			log.Info("got ERR from client, closing sending loop")
			return nil
		}
	}
}

var grpcPort = flag.String("grpc-port", "31337", "gRPC listen port")

func main() {
	flag.Parse()

	go http.ListenAndServe(":9191", nil)

	// Initialize GPIO.
	var err error
	if err = embd.InitGPIO(); err != nil {
		log.Fatalf("Can't init GPIO: %v", err)
	}
	defer embd.CloseGPIO()
	front, err := newEcho("front", 9, 10)
	if err != nil {
		log.Fatalf("Can't init front echo: %v", err)
	}
	defer front.close()
	rear, err := newEcho("rear", 19, 20)
	if err != nil {
		log.Fatalf("Can't init rear echo: %v", err)
	}
	defer rear.close()
	go front.runDistancer()
	go rear.runDistancer()

	left, err := newEngine(23, 4)
	if err != nil {
		log.Fatalf("Can't init left engine: %v", err)
	}
	defer left.close()
	right, err := newEngine(24, 17)
	if err != nil {
		log.Fatalf("Can't init right engine: %v", err)
	}
	defer right.close()

	// Listen for GRPC connections.
	lis, err := net.Listen("tcp", ":"+*grpcPort)
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}
	defer lis.Close()

	drv := driver{left: left, right: right}
	go drv.safetyStop()

	srv := server{front: front, rear: rear, driver: &drv}
	s := grpc.NewServer()
	pb.RegisterDriverServer(s, &srv)

	// Open broadcast connection.
	bcast, err := net.ListenPacket("udp", ":0")
	if err != nil {
		log.Fatal(err)
	}
	defer bcast.Close()

	broadcastAddr := "255.255.255.255:8032"
	dst, err := net.ResolveUDPAddr("udp", broadcastAddr)
	if err != nil {
		log.Fatal(err)
	}

	go func() {
		log.Infof("Starting to broadcast our port %s on %s", *grpcPort, broadcastAddr)
		for {
			if _, err := bcast.WriteTo([]byte(*grpcPort), dst); err != nil {
				log.Warn(err)
			}
			time.Sleep(time.Second)
		}
	}()

	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGTERM, syscall.SIGINT)
	go func() {
		sig := <-c
		log.Infof("Got %s, trying to shutdown gracefully", sig.String())
		front.close()
		rear.close()
		left.close()
		right.close()
		embd.CloseGPIO()
		lis.Close()
		bcast.Close()
		os.Exit(0)
	}()

	// Start serving GRPC.
	log.Fatal(s.Serve(lis))
}
<commit_msg>Initial PWM implementation for motors<commit_after>package main

import (
	"flag"
	"fmt"
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	pb "github.com/viru/berrybot/proto"

	log "github.com/Sirupsen/logrus"
	"github.com/kidoman/embd"
	_ "github.com/kidoman/embd/host/rpi" // This loads the RPi driver
	"google.golang.org/grpc"
)

// server is used to implement hellowrld.GreeterServer.
type server struct {
	front, rear *echo
	driver      *driver
}

type echo struct {
	name       string
	echo       embd.DigitalPin
	trig       embd.DigitalPin
	quit, done chan bool
	dist       int64
	last       time.Time
	enabled    bool
}

func newEcho(name string, trigPin, echoPin int) (*echo, error) {
	var e echo
	e.name = name
	e.quit = make(chan bool)
	e.done = make(chan bool)
	var err error
	e.trig, err = embd.NewDigitalPin(trigPin)
	if err != nil {
		return nil, fmt.Errorf("can't init trigger pin: %v", err)
	}
	e.echo, err = embd.NewDigitalPin(echoPin)
	if err != nil {
		return nil, fmt.Errorf("can't init echo pin: %v", err)
	}

	// Set direction.
	if err := e.trig.SetDirection(embd.Out); err != nil {
		return nil, fmt.Errorf("can't set trigger direction: %v", err)
	}
	if err := e.echo.SetDirection(embd.In); err != nil {
		return nil, fmt.Errorf("can't set echo direction: %v", err)
	}
	return &e, nil
}

func (e *echo) measure() {
	log.Infof("%s: measuring...", e.name)
	if err := e.trig.Write(embd.High); err != nil {
		log.Warnf("can't set trigger to high: %v", err)
	}
	time.Sleep(time.Microsecond * 10)
	if err := e.trig.Write(embd.Low); err != nil {
		log.Warnf("can't set trigger to low: %v", err)
	}
	dur, err := e.echo.TimePulse(embd.High)
	if err != nil {
		log.Warnf("can't time pulse: %v", err)
	}
	log.Infof("%s: distance: %dcm", e.name, dur.Nanoseconds()/1000*34/1000/2)
	e.dist = dur.Nanoseconds() / 1000 * 34 / 1000 / 2
}

const (
	defaultFastDur = time.Millisecond * 250
	defaultSlowDur = time.Second
)

// Goroutine measuring distance in an infinite loop. If distancer is enabled (bot is driving)
// then measuring on fast timer, otherwise only once per second to save CPU cycles.
func (e *echo) runDistancer() {
	if err := e.trig.Write(embd.Low); err != nil {
		log.Warnf("can't set trigger to low: %v", err)
	}
	time.Sleep(time.Second * 1) // Settle time needed after initial activation.
	fast := time.NewTicker(defaultFastDur)
	defer fast.Stop()
	slow := time.NewTicker(defaultSlowDur)
	defer slow.Stop()
	for {
		select {
		case <-e.quit:
			e.done <- true
			return
		case <-slow.C:
			if !e.enabled {
				e.measure()
			}
		case <-fast.C:
			if e.enabled {
				e.measure()
			}
		}
	}
}

func (e *echo) close() {
	e.quit <- true
	<-e.done
	e.echo.Close()
	e.trig.Close()
}

type driver struct {
	left, right *engine
	mu          sync.Mutex
	moving      bool
	last        time.Time
}

func (d *driver) safetyStop() {
	ticker := time.NewTicker(time.Second)
	for range ticker.C {
		d.mu.Lock()
		if d.moving && d.last.Add(time.Second).Before(time.Now()) {
			d.mu.Unlock()
			d.stop()
			log.Warn("Emergency stop!")
			continue
		}
		d.mu.Unlock()
	}
}

func (d *driver) setMoving(moving bool) {
	d.mu.Lock()
	d.last = time.Now()
	d.moving = moving
	d.mu.Unlock()
}

func (d *driver) stop() {
	d.left.pwr = 0
	d.right.pwr = 0
	d.setMoving(false)
}

func (d *driver) forward(pwr int32) {
	d.left.pwr = pwr
	d.left.fwdPin.Write(embd.High)
	d.right.pwr = pwr
	d.right.fwdPin.Write(embd.High)
	d.setMoving(true)
}

func (d *driver) backward(pwr int32) {
	d.left.pwr = pwr
	d.left.fwdPin.Write(embd.Low)
	d.right.pwr = pwr
	d.right.fwdPin.Write(embd.Low)
	d.setMoving(true)
}

func (d *driver) sharpRight(pwr int32) {
	d.left.pwr = pwr
	d.left.fwdPin.Write(embd.High)
	d.right.pwr = pwr
	d.right.fwdPin.Write(embd.Low)
	d.setMoving(true)
}

func (d *driver) sharpLeft(pwr int32) {
	d.left.pwr = pwr
	d.left.fwdPin.Write(embd.Low)
	d.right.pwr = pwr
	d.right.fwdPin.Write(embd.High)
	d.setMoving(true)
}

func (d *driver) fwdRight() {
	d.left.pwr = 100
	d.left.fwdPin.Write(embd.High)
	d.right.pwr = 50
	d.right.fwdPin.Write(embd.High)
	d.setMoving(true)
}

func (d *driver) fwdLeft() {
	d.left.pwr = 50
	d.left.fwdPin.Write(embd.High)
	d.right.pwr = 100
	d.right.fwdPin.Write(embd.High)
	d.setMoving(true)
}

func (d *driver) backRight() {
	d.left.pwr = 100
	d.left.fwdPin.Write(embd.Low)
	d.right.pwr = 50
	d.right.fwdPin.Write(embd.Low)
	d.setMoving(true)
}

func (d *driver) backLeft() {
	d.left.pwr = 50
	d.left.fwdPin.Write(embd.Low)
	d.right.pwr = 100
	d.right.fwdPin.Write(embd.Low)
	d.setMoving(true)
}

const (
	safeStraightDist = 20
	safeTurningDist  = 10
)

func (s *server) drive(dir *pb.Direction) {
	switch {
	case dir.Dy > -5 && dir.Dy < 5 && dir.Dx > -5 && dir.Dx < 5:
		s.front.enabled = false
		s.rear.enabled = false
		s.driver.stop()
	case dir.Dy > 25 && dir.Dx > -25 && dir.Dx < 25:
		s.front.enabled = true
		s.driver.forward(dir.Dy)
	case dir.Dy < -25 && dir.Dx > -25 && dir.Dx < 25:
		s.rear.enabled = true
		s.driver.backward(-dir.Dy)
	case dir.Dx > 25 && dir.Dy > -25 && dir.Dy < 25:
		s.driver.sharpRight(dir.Dx)
	case dir.Dx < -25 && dir.Dy > -25 && dir.Dy < 25:
		s.driver.sharpLeft(-dir.Dx)
	case dir.Dx > 25 && dir.Dy > 25:
		s.driver.fwdRight()
	case dir.Dx < -25 && dir.Dy > 25:
		s.driver.fwdLeft()
	case dir.Dx > 25 && dir.Dy < -25:
		s.driver.backRight()
	case dir.Dx < -25 && dir.Dy < -25:
		s.driver.backLeft()
	}
}

type engine struct {
	fwdPin, pwrPin embd.DigitalPin
	pwr            int32
}

func newEngine(pwrPin, fwdPin int) (*engine, error) {
	var e engine
	var err error
	e.pwrPin, err = embd.NewDigitalPin(pwrPin)
	if err != nil {
		return nil, fmt.Errorf("can't init power pin: %v", err)
	}
	e.fwdPin, err = embd.NewDigitalPin(fwdPin)
	if err != nil {
		return nil, fmt.Errorf("can't init forward pin: %v", err)
	}

	// Set direction.
	if err := e.pwrPin.SetDirection(embd.Out); err != nil {
		return nil, fmt.Errorf("can't set power direction: %v", err)
	}
	if err := e.fwdPin.SetDirection(embd.Out); err != nil {
		return nil, fmt.Errorf("can't set forward direction: %v", err)
	}
	go e.startPWM()
	return &e, nil
}

func (e *engine) close() {
	e.pwrPin.Close()
	e.fwdPin.Close()
}

func (e *engine) startPWM() {
	ticker := time.NewTicker(time.Millisecond * 100)
	flap := embd.Low
	for range ticker.C {
		if e.pwr < 50 {
			e.pwrPin.Write(embd.Low)
		} else if e.pwr < 75 {
			e.pwrPin.Write(flap)
			if flap == embd.Low {
				flap = embd.High
			} else {
				flap = embd.Low
			}
		} else {
			e.pwrPin.Write(embd.High)
		}
	}
}

const (
	sensorUnknown = iota
	sensorFront
	sensorRear
)

func (s *server) Drive(stream pb.Driver_DriveServer) error {
	waitc := make(chan struct{})
	go func() {
		for {
			d, err := stream.Recv()
			if err != nil {
				log.Warnf("ERR from client: %v", err)
				close(waitc)
				return
			}
			s.drive(d)
		}
	}()

	for {
		select {
		case <-time.After(time.Millisecond * 500):
			var speed int32
			if s.driver.moving {
				speed = 100
			}
			if err := stream.Send(&pb.Telemetry{Speed: speed, DistFront: int32(s.front.dist), DistRear: int32(s.rear.dist)}); err != nil {
				log.Errorf("can't send telemetry: %v", err)
				return err
			}
			log.Info("Sending telemetry!")
		case <-waitc:
			log.Info("got ERR from client, closing sending loop")
			return nil
		}
	}
}

var grpcPort = flag.String("grpc-port", "31337", "gRPC listen port")

func main() {
	flag.Parse()

	go http.ListenAndServe(":9191", nil)

	// Initialize GPIO.
	var err error
	if err = embd.InitGPIO(); err != nil {
		log.Fatalf("Can't init GPIO: %v", err)
	}
	defer embd.CloseGPIO()
	front, err := newEcho("front", 9, 10)
	if err != nil {
		log.Fatalf("Can't init front echo: %v", err)
	}
	defer front.close()
	rear, err := newEcho("rear", 19, 20)
	if err != nil {
		log.Fatalf("Can't init rear echo: %v", err)
	}
	defer rear.close()
	go front.runDistancer()
	go rear.runDistancer()

	left, err := newEngine(23, 4)
	if err != nil {
		log.Fatalf("Can't init left engine: %v", err)
	}
	defer left.close()
	right, err := newEngine(24, 17)
	if err != nil {
		log.Fatalf("Can't init right engine: %v", err)
	}
	defer right.close()

	// Listen for GRPC connections.
	lis, err := net.Listen("tcp", ":"+*grpcPort)
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}
	defer lis.Close()

	drv := driver{left: left, right: right}
	go drv.safetyStop()

	srv := server{front: front, rear: rear, driver: &drv}
	s := grpc.NewServer()
	pb.RegisterDriverServer(s, &srv)

	// Open broadcast connection.
	bcast, err := net.ListenPacket("udp", ":0")
	if err != nil {
		log.Fatal(err)
	}
	defer bcast.Close()

	broadcastAddr := "255.255.255.255:8032"
	dst, err := net.ResolveUDPAddr("udp", broadcastAddr)
	if err != nil {
		log.Fatal(err)
	}

	go func() {
		log.Infof("Starting to broadcast our port %s on %s", *grpcPort, broadcastAddr)
		for {
			if _, err := bcast.WriteTo([]byte(*grpcPort), dst); err != nil {
				log.Warn(err)
			}
			time.Sleep(time.Second)
		}
	}()

	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGTERM, syscall.SIGINT)
	go func() {
		sig := <-c
		log.Infof("Got %s, trying to shutdown gracefully", sig.String())
		front.close()
		rear.close()
		left.close()
		right.close()
		embd.CloseGPIO()
		lis.Close()
		bcast.Close()
		os.Exit(0)
	}()

	// Start serving GRPC.
	log.Fatal(s.Serve(lis))
}
<|endoftext|>