<commit_before>// Package nosurf implements an HTTP handler that
// mitigates Cross-Site Request Forgery Attacks.
package nosurf

import (
	"net/http"
	"regexp"
)

const (
	CookieName  = "csrf_token"
	HeaderName  = "X-CSRF-Token"
	FailureCode = 400
)

type CSRFHandler struct {
	// Handlers that CSRFHandler wraps.
	successHandler http.Handler
	failureHandler http.Handler

	// Slices of URLs that are exempt from CSRF checks.
	// They can be specified by...
	// ...an exact URL
	exemptURLs []string
	// ...a glob (as used by path.Match())
	exemptGlobs []string
	// ...a regexp.
	exemptRegexps []*regexp.Regexp

	// All of those will be matched against Request.URL.Path,
	// So they should begin with a leading slash
}

func defaultFailureHandler(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(FailureCode)
}

// Constructs a new `CSRFHandler` that calls
// the specified handler if the CSRF check succeeds.
func New(handler http.Handler) *CSRFHandler {
	csrf := &CSRFHandler{successHandler: handler,
		failureHandler: http.HandlerFunc(defaultFailureHandler),
		exemptURLs:     make([]string, 0),
		exemptGlobs:    make([]string, 0),
		exemptRegexps:  make([]*regexp.Regexp, 0),
	}

	return csrf
}

// Sets the handler to call in case the CSRF check
// fails. By default it's `defaultFailureHandler`.
func (h *CSRFHandler) SetFailureHandler(handler http.Handler) {
	h.failureHandler = handler
}
<commit_msg>Documentation updates<commit_after>// Package nosurf implements an HTTP handler that
// mitigates Cross-Site Request Forgery Attacks.
package nosurf

import (
	"net/http"
	"regexp"
)

const (
	// the name of CSRF cookie
	CookieName = "csrf_token"
	// the name of CSRF header
	HeaderName = "X-CSRF-Token"
	// the HTTP status code for the default failure handler
	FailureCode = 400
)

type CSRFHandler struct {
	// Handlers that CSRFHandler wraps.
	successHandler http.Handler
	failureHandler http.Handler

	// Slices of URLs that are exempt from CSRF checks.
	// They can be specified by...
	// ...an exact URL
	exemptURLs []string
	// ...a glob (as used by path.Match())
	exemptGlobs []string
	// ...a regexp.
	exemptRegexps []*regexp.Regexp

	// All of those will be matched against Request.URL.Path,
	// So they should begin with a leading slash
}

func defaultFailureHandler(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(FailureCode)
}

// Constructs a new CSRFHandler that calls
// the specified handler if the CSRF check succeeds.
func New(handler http.Handler) *CSRFHandler {
	csrf := &CSRFHandler{successHandler: handler,
		failureHandler: http.HandlerFunc(defaultFailureHandler),
		exemptURLs:     make([]string, 0),
		exemptGlobs:    make([]string, 0),
		exemptRegexps:  make([]*regexp.Regexp, 0),
	}

	return csrf
}

// Sets the handler to call in case the CSRF check
// fails. By default it's defaultFailureHandler.
func (h *CSRFHandler) SetFailureHandler(handler http.Handler) {
	h.failureHandler = handler
}
<|endoftext|>


<commit_before>package dochaincore

import (
	"bytes"
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"html/template"
	"net/http"
	"net/url"
	"os"
	"path"
	"strings"
	"sync"
	"time"
)

const indexPageHTML = `<!DOCTYPE html>

		<title>One-Click Chain Core DigitalOcean</title>
		<link rel="stylesheet" href="/static/style.css">

Install Chain Core on a DigitalOcean droplet. This installer creates a new 1gb droplet and a 100gb block storage volume on your DigitalOcean account. It installs Chain Core on the droplet using the attached volume for storage. The approximate cost on DigitalOcean is $20/month.

			Install Chain Core

`
const progressPageHTML = `
<!DOCTYPE html>

		<title>One-Click Chain Core DigitalOcean</title>
		<link rel="stylesheet" href="/static/style.css">

Initializing droplet…

Success! Chain Core has been installed on your DigitalOcean droplet. To access
				Chain Core's API and Dashboard, you'll need your client token:

URL:

Token:

				Open dashboard

`

func Handler(oauthClientID, oauthClientSecret, host string) http.Handler {
	h := &handler{
		oauthClientID:     oauthClientID,
		oauthClientSecret: oauthClientSecret,
		host:              host,
		progressTmpl:      template.Must(template.New("progresspage").Parse(progressPageHTML)),
		indexTmpl:         template.Must(template.New("index").Parse(indexPageHTML)),
		installs:          make(map[string]*install),
	}
	mux := http.NewServeMux()
	mux.HandleFunc("/status/", h.status)
	mux.HandleFunc("/grant", h.grant)
	mux.HandleFunc("/install/", h.progressPage)
	mux.Handle("/static/", http.StripPrefix("/static/", http.FileServer(http.Dir("static"))))
	mux.HandleFunc("/", h.index)
	return mux
}

type handler struct {
	oauthClientID     string
	oauthClientSecret string
	host              string
	progressTmpl      *template.Template
	indexTmpl         *template.Template

	installMu sync.Mutex
	installs  map[string]*install
}

func (h *handler) index(rw http.ResponseWriter, req *http.Request) {
	b := make([]byte, 16)
	_, err := rand.Read(b)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusInternalServerError)
		return
	}
	installID := hex.EncodeToString(b)

	h.installMu.Lock()
	h.installs[installID] = &install{Status: "pending auth"}
	h.installMu.Unlock()

	vals := make(url.Values)
	vals.Set("response_type", "code")
	vals.Set("client_id", h.oauthClientID)
	vals.Set("state", installID)
	vals.Set("scope", "read write")
	vals.Set("redirect_uri", h.host+"/grant")
	u := url.URL{
		Scheme:   "https",
		Host:     "cloud.digitalocean.com",
		Path:     "/v1/oauth/authorize",
		RawQuery: vals.Encode(),
	}

	tmplData := struct {
		InstallLink string
	}{
		InstallLink: u.String(),
	}
	h.indexTmpl.Execute(rw, tmplData)
}

func (h *handler) grant(rw http.ResponseWriter, req *http.Request) {
	code, state := req.FormValue("code"), req.FormValue("state")
	if code == "" || state == "" {
		http.Error(rw, "invalid oauth2 grant", http.StatusBadRequest)
		return
	}

	h.installMu.Lock()
	curr := h.installs[state]
	h.installMu.Unlock()
	if curr == nil {
		http.Error(rw, "invalid oauth2 state", http.StatusBadRequest)
		return
	}

	// Claim the code grant
	vals := make(url.Values)
	vals.Set("grant_type", "authorization_code")
	vals.Set("code", code)
	vals.Set("client_id", h.oauthClientID)
	vals.Set("client_secret", h.oauthClientSecret)
	vals.Set("redirect_uri", h.host+"/grant")
	u := url.URL{
		Scheme:   "https",
		Host:     "cloud.digitalocean.com",
		Path:     "/v1/oauth/token",
		RawQuery: vals.Encode(),
	}
	resp, err := http.Post(u.String(), "application/x-www-form-urlencoded", nil)
	if err != nil {
		http.Error(rw, "internal server error", http.StatusInternalServerError)
		return
	}
	defer resp.Body.Close()

	var decodedResponse struct {
		AccessToken  string `json:"access_token"`
		TokenType    string `json:"bearer"`
		ExpiresIn    int    `json:"expires_in"`
		RefreshToken string `json:"refresh_token"`
		Scope        string `json:"scope"`
		Info         struct {
			Name  string `json:"name"`
			Email string `json:"email"`
			UUID  string `json:"uuid"`
		} `json:"info"`
	}
	err = json.NewDecoder(resp.Body).Decode(&decodedResponse)
	if err != nil {
		http.Error(rw, "err decoding access token grant", http.StatusBadRequest)
		return
	}
	if decodedResponse.AccessToken == "" {
		http.Error(rw, "missing access token", http.StatusBadRequest)
		return
	}
	if decodedResponse.Scope != "read write" {
		http.Error(rw, "need read write OAuth scope", http.StatusBadRequest)
		return
	}

	curr.mu.Lock()
	curr.accessToken = decodedResponse.AccessToken
	curr.mu.Unlock()

	http.Redirect(rw, req, "/install/"+state, http.StatusFound)
}

func (h *handler) progressPage(rw http.ResponseWriter, req *http.Request) {
	id := path.Base(req.URL.Path)
	h.installMu.Lock()
	curr := h.installs[id]
	h.installMu.Unlock()

	if curr == nil {
		http.NotFound(rw, req)
		return
	}

	go curr.init(id)

	tmplData := struct {
		InstallID string
	}{
		InstallID: id,
	}
	err := h.progressTmpl.Execute(rw, tmplData)
	if err != nil {
		fmt.Fprintf(os.Stderr, "executing template: %s", err.Error())
	}
}

func (h *handler) status(rw http.ResponseWriter, req *http.Request) {
	id := path.Base(req.URL.Path)
	h.installMu.Lock()
	curr := h.installs[id]
	h.installMu.Unlock()

	if curr == nil {
		http.NotFound(rw, req)
		return
	}

	// Marshal to a buffer first so that a really slow request can't
	// keep curr.mu locked indefinitely.
	var buf bytes.Buffer
	curr.mu.Lock()
	_ = json.NewEncoder(&buf).Encode(curr)
	curr.mu.Unlock()

	rw.Header().Set("Content-Type", "application/json")
	rw.Write(buf.Bytes())
}

type install struct {
	mu          sync.Mutex
	Status      string `json:"status"`
	ClientToken string `json:"client_token"`
	IPAddress   string `json:"ip_address"`
	accessToken string
	c           *Core
}

func (i *install) setStatus(status string) {
	i.mu.Lock()
	defer i.mu.Unlock()
	i.Status = status
}

func (i *install) init(state string) {
	defer revoke(i.accessToken)

	// Set a 10 minute timeout for the installation. From beginning
	// to end it should only take a ~2 minutes, but make sure we
	// cleanup and revoke the access token even if it takes longer.
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
	defer cancel()

	var core *Core
	var err error

	defer func() {
		if err != nil {
			i.setStatus(err.Error())
		}
	}()

	// Start deploying and create the droplet.
	core, err = Deploy(ctx, i.accessToken, DropletName("chain-core-"+state[:6]))
	if err != nil {
		return
	}

	i.mu.Lock()
	i.IPAddress = core.IPv4Address
	i.c = core
	i.Status = "waiting for ssh"
	i.mu.Unlock()

	err = WaitForSSH(ctx, core)
	if err != nil {
		return
	}

	i.setStatus("waiting for http")
	err = WaitForHTTP(ctx, core)
	if err != nil {
		return
	}

	i.setStatus("creating client token")
	token, err := CreateClientToken(ctx, core)
	if err != nil {
		return
	}

	i.mu.Lock()
	i.Status = "done"
	i.ClientToken = token
	i.c = nil // garbage collect the SSH keys
	i.mu.Unlock()
}

func revoke(accessToken string) error {
	body := strings.NewReader(url.Values{"token": {accessToken}}.Encode())
	req, err := http.NewRequest("POST", "https://cloud.digitalocean.com/v1/oauth/revoke", body)
	if err != nil {
		return err
	}
	req.Header.Set("Authorization", "Bearer "+accessToken)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("revoke endpoint returned %d status code", resp.StatusCode)
	}
	return nil
}
<commit_msg>handler: add reminder about block storage volume<commit_after>package dochaincore

import (
	"bytes"
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"html/template"
	"net/http"
	"net/url"
	"os"
	"path"
	"strings"
	"sync"
	"time"
)

const indexPageHTML = `<!DOCTYPE html>

		<title>One-Click Chain Core DigitalOcean</title>
		<link rel="stylesheet" href="/static/style.css">

Install Chain Core on a DigitalOcean droplet. This installer creates a new 1gb droplet and a 100gb block storage volume on your DigitalOcean account. It installs Chain Core on the droplet using the attached volume for storage. The approximate cost on DigitalOcean is $20/month.

			Install Chain Core

`
const progressPageHTML = `
<!DOCTYPE html>

		<title>One-Click Chain Core DigitalOcean</title>
		<link rel="stylesheet" href="/static/style.css">

Initializing droplet…

Success! Chain Core has been installed on your DigitalOcean droplet. To access
				Chain Core's API and Dashboard, you'll need your client token:

URL:

Token:

				Open dashboard

When destroying the droplet, remember to also destroy its block storage volume.

`

func Handler(oauthClientID, oauthClientSecret, host string) http.Handler {
	h := &handler{
		oauthClientID:     oauthClientID,
		oauthClientSecret: oauthClientSecret,
		host:              host,
		progressTmpl:      template.Must(template.New("progresspage").Parse(progressPageHTML)),
		indexTmpl:         template.Must(template.New("index").Parse(indexPageHTML)),
		installs:          make(map[string]*install),
	}
	mux := http.NewServeMux()
	mux.HandleFunc("/status/", h.status)
	mux.HandleFunc("/grant", h.grant)
	mux.HandleFunc("/install/", h.progressPage)
	mux.Handle("/static/", http.StripPrefix("/static/", http.FileServer(http.Dir("static"))))
	mux.HandleFunc("/", h.index)
	return mux
}

type handler struct {
	oauthClientID     string
	oauthClientSecret string
	host              string
	progressTmpl      *template.Template
	indexTmpl         *template.Template

	installMu sync.Mutex
	installs  map[string]*install
}

func (h *handler) index(rw http.ResponseWriter, req *http.Request) {
	b := make([]byte, 16)
	_, err := rand.Read(b)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusInternalServerError)
		return
	}
	installID := hex.EncodeToString(b)

	h.installMu.Lock()
	h.installs[installID] = &install{Status: "pending auth"}
	h.installMu.Unlock()

	vals := make(url.Values)
	vals.Set("response_type", "code")
	vals.Set("client_id", h.oauthClientID)
	vals.Set("state", installID)
	vals.Set("scope", "read write")
	vals.Set("redirect_uri", h.host+"/grant")
	u := url.URL{
		Scheme:   "https",
		Host:     "cloud.digitalocean.com",
		Path:     "/v1/oauth/authorize",
		RawQuery: vals.Encode(),
	}

	tmplData := struct {
		InstallLink string
	}{
		InstallLink: u.String(),
	}
	h.indexTmpl.Execute(rw, tmplData)
}

func (h *handler) grant(rw http.ResponseWriter, req *http.Request) {
	code, state := req.FormValue("code"), req.FormValue("state")
	if code == "" || state == "" {
		http.Error(rw, "invalid oauth2 grant", http.StatusBadRequest)
		return
	}

	h.installMu.Lock()
	curr := h.installs[state]
	h.installMu.Unlock()
	if curr == nil {
		http.Error(rw, "invalid oauth2 state", http.StatusBadRequest)
		return
	}

	// Claim the code grant
	vals := make(url.Values)
	vals.Set("grant_type", "authorization_code")
	vals.Set("code", code)
	vals.Set("client_id", h.oauthClientID)
	vals.Set("client_secret", h.oauthClientSecret)
	vals.Set("redirect_uri", h.host+"/grant")
	u := url.URL{
		Scheme:   "https",
		Host:     "cloud.digitalocean.com",
		Path:     "/v1/oauth/token",
		RawQuery: vals.Encode(),
	}
	resp, err := http.Post(u.String(), "application/x-www-form-urlencoded", nil)
	if err != nil {
		http.Error(rw, "internal server error", http.StatusInternalServerError)
		return
	}
	defer resp.Body.Close()

	var decodedResponse struct {
		AccessToken  string `json:"access_token"`
		TokenType    string `json:"bearer"`
		ExpiresIn    int    `json:"expires_in"`
		RefreshToken string `json:"refresh_token"`
		Scope        string `json:"scope"`
		Info         struct {
			Name  string `json:"name"`
			Email string `json:"email"`
			UUID  string `json:"uuid"`
		} `json:"info"`
	}
	err = json.NewDecoder(resp.Body).Decode(&decodedResponse)
	if err != nil {
		http.Error(rw, "err decoding access token grant", http.StatusBadRequest)
		return
	}
	if decodedResponse.AccessToken == "" {
		http.Error(rw, "missing access token", http.StatusBadRequest)
		return
	}
	if decodedResponse.Scope != "read write" {
		http.Error(rw, "need read write OAuth scope", http.StatusBadRequest)
		return
	}

	curr.mu.Lock()
	curr.accessToken = decodedResponse.AccessToken
	curr.mu.Unlock()

	http.Redirect(rw, req, "/install/"+state, http.StatusFound)
}

func (h *handler) progressPage(rw http.ResponseWriter, req *http.Request) {
	id := path.Base(req.URL.Path)
	h.installMu.Lock()
	curr := h.installs[id]
	h.installMu.Unlock()

	if curr == nil {
		http.NotFound(rw, req)
		return
	}

	go curr.init(id)

	tmplData := struct {
		InstallID string
	}{
		InstallID: id,
	}
	err := h.progressTmpl.Execute(rw, tmplData)
	if err != nil {
		fmt.Fprintf(os.Stderr, "executing template: %s", err.Error())
	}
}

func (h *handler) status(rw http.ResponseWriter, req *http.Request) {
	id := path.Base(req.URL.Path)
	h.installMu.Lock()
	curr := h.installs[id]
	h.installMu.Unlock()

	if curr == nil {
		http.NotFound(rw, req)
		return
	}

	// Marshal to a buffer first so that a really slow request can't
	// keep curr.mu locked indefinitely.
	var buf bytes.Buffer
	curr.mu.Lock()
	_ = json.NewEncoder(&buf).Encode(curr)
	curr.mu.Unlock()

	rw.Header().Set("Content-Type", "application/json")
	rw.Write(buf.Bytes())
}

type install struct {
	mu          sync.Mutex
	Status      string `json:"status"`
	ClientToken string `json:"client_token"`
	IPAddress   string `json:"ip_address"`
	accessToken string
	c           *Core
}

func (i *install) setStatus(status string) {
	i.mu.Lock()
	defer i.mu.Unlock()
	i.Status = status
}

func (i *install) init(state string) {
	defer revoke(i.accessToken)

	// Set a 10 minute timeout for the installation. From beginning
	// to end it should only take a ~2 minutes, but make sure we
	// cleanup and revoke the access token even if it takes longer.
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
	defer cancel()

	var core *Core
	var err error

	defer func() {
		if err != nil {
			i.setStatus(err.Error())
		}
	}()

	// Start deploying and create the droplet.
	core, err = Deploy(ctx, i.accessToken, DropletName("chain-core-"+state[:6]))
	if err != nil {
		return
	}

	i.mu.Lock()
	i.IPAddress = core.IPv4Address
	i.c = core
	i.Status = "waiting for ssh"
	i.mu.Unlock()

	err = WaitForSSH(ctx, core)
	if err != nil {
		return
	}

	i.setStatus("waiting for http")
	err = WaitForHTTP(ctx, core)
	if err != nil {
		return
	}

	i.setStatus("creating client token")
	token, err := CreateClientToken(ctx, core)
	if err != nil {
		return
	}

	i.mu.Lock()
	i.Status = "done"
	i.ClientToken = token
	i.c = nil // garbage collect the SSH keys
	i.mu.Unlock()
}

func revoke(accessToken string) error {
	body := strings.NewReader(url.Values{"token": {accessToken}}.Encode())
	req, err := http.NewRequest("POST", "https://cloud.digitalocean.com/v1/oauth/revoke", body)
	if err != nil {
		return err
	}
	req.Header.Set("Authorization", "Bearer "+accessToken)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("revoke endpoint returned %d status code", resp.StatusCode)
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"

	"github.com/BurntSushi/toml"
	"github.com/progrium/go-basher"
)

var Version string
var PluginPath string

func assert(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

func TomlGet(args []string) {
	bytes, err := ioutil.ReadAll(os.Stdin)
	assert(err)
	var t map[string]interface{}
	_, err = toml.Decode(string(bytes), &t)
	assert(err)
	fmt.Println(t[args[0]].(map[string]interface{})[args[1]])
}

func TomlExport(args []string) {
	plugin := args[0]
	bytes, err := ioutil.ReadAll(os.Stdin)
	assert(err)

	var c map[string]map[string]string
	_, err = toml.Decode(string(bytes), &c)
	assert(err)
	config := c[plugin]
	prefix := strings.ToUpper(strings.Replace(plugin, "-", "_", -1))

	var p map[string]map[string]interface{}
	_, err = toml.DecodeFile(PluginPath+"/available/"+plugin+"/plugin.toml", &p)
	assert(err)
	config_def := p["plugin"]["config"].(map[string]interface{})

	for key := range config_def {
		k := strings.ToUpper(strings.Replace(key, "-", "_", -1))
		fmt.Println("export CONFIG_" + k + "=\"${" + prefix + "_" + k + ":-\"" + config[key] + "\"}\"")
	}
}

func TomlSet(args []string) {
	bytes, err := ioutil.ReadAll(os.Stdin)
	assert(err)
	var t map[string]map[string]string
	_, err = toml.DecodeFile(args[0], &t)
	assert(err)
	if t[args[1]] == nil {
		t[args[1]] = make(map[string]string)
	}
	t[args[1]][args[2]] = string(bytes)
	f, err := os.Create(args[0])
	assert(err)
	assert(toml.NewEncoder(f).Encode(t))
	f.Close()
}

func main() {
	os.Setenv("PLUGN_VERSION", Version)
	if data, err := ioutil.ReadFile(".plugn"); err == nil {
		if path, err := filepath.Abs(string(data)); err == nil {
			os.Setenv("PLUGIN_PATH", path)
		}
	}
	if os.Getenv("PLUGIN_PATH") == "" {
		fmt.Println("!! PLUGIN_PATH is not set in environment")
		os.Exit(2)
	}
	PluginPath = os.Getenv("PLUGIN_PATH")
	if len(os.Args) > 1 && os.Args[1] == "gateway" {
		runGateway()
		return
	}
	basher.Application(map[string]func([]string){
		"toml-get":        TomlGet,
		"toml-set":        TomlSet,
		"toml-export":     TomlExport,
		"trigger-gateway": TriggerGateway,
		"reload-gateway":  ReloadGateway,
	}, []string{
		"bashenv/bash.bash",
		"bashenv/fn.bash",
		"bashenv/cmd.bash",
		"bashenv/plugn.bash",
	}, Asset, true)
}
<commit_msg>feat: allow specifying a BASH_BIN to override the embedded bash in go-basher<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"

	"github.com/BurntSushi/toml"
	"github.com/progrium/go-basher"
)

var Version string
var PluginPath string

func assert(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

func TomlGet(args []string) {
	bytes, err := ioutil.ReadAll(os.Stdin)
	assert(err)
	var t map[string]interface{}
	_, err = toml.Decode(string(bytes), &t)
	assert(err)
	fmt.Println(t[args[0]].(map[string]interface{})[args[1]])
}

func TomlExport(args []string) {
	plugin := args[0]
	bytes, err := ioutil.ReadAll(os.Stdin)
	assert(err)

	var c map[string]map[string]string
	_, err = toml.Decode(string(bytes), &c)
	assert(err)
	config := c[plugin]
	prefix := strings.ToUpper(strings.Replace(plugin, "-", "_", -1))

	var p map[string]map[string]interface{}
	_, err = toml.DecodeFile(PluginPath+"/available/"+plugin+"/plugin.toml", &p)
	assert(err)
	config_def := p["plugin"]["config"].(map[string]interface{})

	for key := range config_def {
		k := strings.ToUpper(strings.Replace(key, "-", "_", -1))
		fmt.Println("export CONFIG_" + k + "=\"${" + prefix + "_" + k + ":-\"" + config[key] + "\"}\"")
	}
}

func TomlSet(args []string) {
	bytes, err := ioutil.ReadAll(os.Stdin)
	assert(err)
	var t map[string]map[string]string
	_, err = toml.DecodeFile(args[0], &t)
	assert(err)
	if t[args[1]] == nil {
		t[args[1]] = make(map[string]string)
	}
	t[args[1]][args[2]] = string(bytes)
	f, err := os.Create(args[0])
	assert(err)
	assert(toml.NewEncoder(f).Encode(t))
	f.Close()
}

func main() {
	os.Setenv("PLUGN_VERSION", Version)
	if data, err := ioutil.ReadFile(".plugn"); err == nil {
		if path, err := filepath.Abs(string(data)); err == nil {
			os.Setenv("PLUGIN_PATH", path)
		}
	}
	if os.Getenv("PLUGIN_PATH") == "" {
		fmt.Println("!! PLUGIN_PATH is not set in environment")
		os.Exit(2)
	}
	PluginPath = os.Getenv("PLUGIN_PATH")
	if len(os.Args) > 1 && os.Args[1] == "gateway" {
		runGateway()
		return
	}
	funcs := map[string]func([]string){
		"toml-get":        TomlGet,
		"toml-set":        TomlSet,
		"toml-export":     TomlExport,
		"trigger-gateway": TriggerGateway,
		"reload-gateway":  ReloadGateway,
	}
	scripts := []string{
		"bashenv/bash.bash",
		"bashenv/fn.bash",
		"bashenv/cmd.bash",
		"bashenv/plugn.bash",
	}

	if os.Getenv("BASH_BIN") == "" {
		basher.Application(funcs, scripts, Asset, true)
	} else {
		basher.ApplicationWithPath(funcs, scripts, Asset, true, os.Getenv("BASH_BIN"))
	}
}
<|endoftext|>


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

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path"
	"strings"
	"sync/atomic"
	"time"
)

const (
	blobStoreUri = "blobs/{blobId}/signedurl"
)

type bundleManagerInterface interface {
	initializeBundleDownloading()
	queueDownloadRequest(*DataDeployment)
	enqueueRequest(*DownloadRequest)
	makeDownloadRequest(string) *DownloadRequest
	deleteBundlesFromDeployments([]DataDeployment)
	deleteBundleById(string)
	Close()
}

type bundleManager struct {
	blobServerUrl             string
	dbMan                     dbManagerInterface
	apiMan                    apiManagerInterface
	concurrentDownloads       int
	markDeploymentFailedAfter time.Duration
	bundleDownloadConnTimeout time.Duration
	bundleRetryDelay          time.Duration
	bundleCleanupDelay        time.Duration
	downloadQueue             chan *DownloadRequest
	isClosed                  *int32
	workers                   []*BundleDownloader
}

type blobServerResponse struct {
	id                       string `json:"id"`
	kind                     string `json:"kind"`
	self                     string `json:"self"`
	signedUrl                string `json:"signedurl"`
	signedUrlExpiryTimestamp string `json:"signedurlexpirytimestamp"`
}

func (bm *bundleManager) initializeBundleDownloading() {
	atomic.StoreInt32(bm.isClosed, 0)
	bm.workers = make([]*BundleDownloader, bm.concurrentDownloads)

	// create workers
	for i := 0; i < bm.concurrentDownloads; i++ {
		worker := BundleDownloader{
			id:       i + 1,
			workChan: make(chan *DownloadRequest),
			bm:       bm,
		}
		bm.workers[i] = &worker
		worker.Start()
	}
}

// download bundle blob and resource blob
// TODO do not download duplicate blobs
func (bm *bundleManager) queueDownloadRequest(dep *DataDeployment) {
	blobReq := bm.makeDownloadRequest(dep.BlobID)
	resourceReq := bm.makeDownloadRequest(dep.BlobResourceID)

	go func() {
		bm.enqueueRequest(blobReq)
		bm.enqueueRequest(resourceReq)
	}()
}

func (bm *bundleManager) makeDownloadRequest(id string) *DownloadRequest {
	markFailedAt := time.Now().Add(bm.markDeploymentFailedAfter)
	retryIn := bm.bundleRetryDelay
	maxBackOff := 5 * time.Minute

	return &DownloadRequest{
		blobServerURL: bm.blobServerUrl,
		bm:            bm,
		blobId:        id,
		backoffFunc:   createBackoff(retryIn, maxBackOff),
		markFailedAt:  markFailedAt,
		connTimeout:   bm.bundleDownloadConnTimeout,
	}
}

// a blocking method to enqueue download requests
func (bm *bundleManager) enqueueRequest(r *DownloadRequest) {
	if atomic.LoadInt32(bm.isClosed) == 1 {
		return
	}
	/*
		defer func() {
			if r := recover(); r != nil {
				log.Warn("trying to enque requests to closed bundleManager")
			}
		}()
	*/
	bm.downloadQueue <- r
}

func (bm *bundleManager) Close() {
	atomic.StoreInt32(bm.isClosed, 1)
	close(bm.downloadQueue)
}

func (bm *bundleManager) deleteBundlesFromDeployments(deletedDeployments []DataDeployment) {
	for _, dep := range deletedDeployments {
		go bm.deleteBundleById(dep.BlobID)
		go bm.deleteBundleById(dep.BlobResourceID)
	}

	/*
		log.Debugf("will delete %d old bundles", len(deletedDeployments))
		go func() {
			// give clients a minute to avoid conflicts
			time.Sleep(bm.bundleCleanupDelay)
			for _, dep := range deletedDeployments {
				bundleFile := getBlobFilePath(dep.BlobID)
				log.Debugf("removing old bundle: %v", bundleFile)
				// TODO Remove from the Database table apid_blob_available
				safeDelete(bundleFile)
			}
		}()
	*/
}

// TODO add delete support
func (bm *bundleManager) deleteBundleById(blobId string) {

}

type DownloadRequest struct {
	bm            *bundleManager
	blobId        string
	backoffFunc   func()
	markFailedAt  time.Time
	connTimeout   time.Duration
	blobServerURL string
}

func (r *DownloadRequest) downloadBundle() error {

	log.Debugf("starting bundle download attempt for blobId=%s", r.blobId)

	if r.checkTimeout() {
		return &timeoutError{
			markFailedAt: r.markFailedAt,
		}
	}

	downloadedFile, err := downloadFromURI(r.blobServerURL, r.blobId, r.connTimeout)

	if err != nil {
		log.Errorf("Unable to download blob file blobId=%s err:%v", r.blobId, err)
		return err
	}

	log.Debugf("blod downloaded. blobid=%s  filepath=%s", r.blobId, downloadedFile)

	err = r.bm.dbMan.updateLocalFsLocation(r.blobId, downloadedFile)
	if err != nil {
		log.Errorf("updateLocalFsLocation failed: blobId=%s", r.blobId)
		return err
	}

	log.Debugf("bundle downloaded: blobId=%s", r.blobId)

	// TODO send changed deployments to subscribers (API call with "block")
	//r.bm.apiMan.addChangedDeployment(dep.ID)

	return nil
}

func (r *DownloadRequest) checkTimeout() bool {

	if !r.markFailedAt.IsZero() && time.Now().After(r.markFailedAt) {
		r.markFailedAt = time.Time{}
		log.Debugf("bundle download timeout. blobId=", r.blobId)
		// TODO notify gateway of this failure
		return true
	}
	return false
}

func getBlobFilePath(blobId string) string {
	return path.Join(bundlePath, base64.StdEncoding.EncodeToString([]byte(blobId)))
}

func getSignedURL(blobServerURL string, blobId string, bundleDownloadConnTimeout time.Duration) (string, error) {

	blobUri, err := url.Parse(blobServerURL)
	if err != nil {
		log.Panicf("bad url value for config %s: %s", blobUri, err)
	}

	blobUri.Path += strings.Replace(blobStoreUri, "{blobId}", blobId, 1)
	parameters := url.Values{}
	parameters.Add("action", "GET")
	blobUri.RawQuery = parameters.Encode()

	uri := blobUri.String()

	surl, err := getURIReader(uri, bundleDownloadConnTimeout)
	if err != nil {
		log.Errorf("Unable to get signed URL from BlobServer %s: %v", uri, err)
		return "", err
	}

	body, err := ioutil.ReadAll(surl)
	if err != nil {
		log.Errorf("Invalid response from BlobServer for {%s} error: {%v}", uri, err)
		return "", err
	}
	res := blobServerResponse{}
	err = json.Unmarshal(body, &res)
	if err != nil {
		log.Errorf("Invalid response from BlobServer for {%s} error: {%v}", uri, err)
		return "", err
	}

	return string(res.signedUrl), nil
}

// downloadFromURI involves retrieving the signed URL for the blob, and storing the resource locally
// after downloading the resource from GCS (via the signed URL)
func downloadFromURI(blobServerURL string, blobId string, bundleDownloadConnTimeout time.Duration) (tempFileName string, err error) {

	var tempFile *os.File
	log.Debugf("Downloading bundle: %s", blobId)

	uri, err := getSignedURL(blobServerURL, blobId, bundleDownloadConnTimeout)
	if err != nil {
		log.Errorf("Unable to get signed URL for blobId {%s}, error : {%v}", blobId, err)
		return
	}

	tempFile, err = ioutil.TempFile(bundlePath, "blob")
	if err != nil {
		log.Errorf("Unable to create temp file: %v", err)
		return
	}
	defer tempFile.Close()
	tempFileName = tempFile.Name()

	var confReader io.ReadCloser
	confReader, err = getURIReader(uri, bundleDownloadConnTimeout)
	if err != nil {
		log.Errorf("Unable to retrieve bundle %s: %v", uri, err)
		return
	}
	defer confReader.Close()

	_, err = io.Copy(tempFile, confReader)
	if err != nil {
		log.Errorf("Unable to write bundle %s: %v", tempFileName, err)
		return
	}

	log.Debugf("Bundle %s downloaded to: %s", uri, tempFileName)
	return
}

// retrieveBundle retrieves bundle data from a URI
func getURIReader(uriString string, bundleDownloadConnTimeout time.Duration) (io.ReadCloser, error) {

	client := http.Client{
		Timeout: bundleDownloadConnTimeout,
	}
	res, err := client.Get(uriString)
	if err != nil {
		return nil, err
	}
	if res.StatusCode != 200 {
		return nil, fmt.Errorf("GET uri %s failed with status %d", uriString, res.StatusCode)
	}
	return res.Body, nil
}

type BundleDownloader struct {
	id       int
	workChan chan *DownloadRequest
	bm       *bundleManager
}

func (w *BundleDownloader) Start() {
	go func() {
		log.Debugf("started bundle downloader %d", w.id)

		for req := range w.bm.downloadQueue {
			log.Debugf("starting download blobId=%s", req.blobId)
			err := req.downloadBundle()
			if err != nil {
				// timeout
				if _, ok := err.(*timeoutError); ok {
					continue
				}
				go func() {
					req.backoffFunc()
					w.bm.enqueueRequest(req)
				}()
			}
		}
		log.Debugf("bundle downloader %d stopped", w.id)
	}()
}

// simple doubling back-off
func createBackoff(retryIn, maxBackOff time.Duration) func() {
	return func() {
		log.Debugf("backoff called. will retry in %s.", retryIn)
		time.Sleep(retryIn)
		retryIn = retryIn * time.Duration(2)
		if retryIn > maxBackOff {
			retryIn = maxBackOff
		}
	}
}

type timeoutError struct {
	markFailedAt time.Time
}

func (e *timeoutError) Error() string {
	return fmt.Sprintf("Timeout. markFailedAt=%v", e.markFailedAt)
}
<commit_msg>update blob server uri path<commit_after>// Copyright 2017 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package apiGatewayConfDeploy

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path"
	"strings"
	"sync/atomic"
	"time"
)

const (
	blobStoreUri = "/blobs/{blobId}"
)

type bundleManagerInterface interface {
	initializeBundleDownloading()
	queueDownloadRequest(*DataDeployment)
	enqueueRequest(*DownloadRequest)
	makeDownloadRequest(string) *DownloadRequest
	deleteBundlesFromDeployments([]DataDeployment)
	deleteBundleById(string)
	Close()
}

type bundleManager struct {
	blobServerUrl             string
	dbMan                     dbManagerInterface
	apiMan                    apiManagerInterface
	concurrentDownloads       int
	markDeploymentFailedAfter time.Duration
	bundleDownloadConnTimeout time.Duration
	bundleRetryDelay          time.Duration
	bundleCleanupDelay        time.Duration
	downloadQueue             chan *DownloadRequest
	isClosed                  *int32
	workers                   []*BundleDownloader
}

type blobServerResponse struct {
	id                       string `json:"id"`
	kind                     string `json:"kind"`
	self                     string `json:"self"`
	signedUrl                string `json:"signedurl"`
	signedUrlExpiryTimestamp string `json:"signedurlexpirytimestamp"`
}

func (bm *bundleManager) initializeBundleDownloading() {
	atomic.StoreInt32(bm.isClosed, 0)
	bm.workers = make([]*BundleDownloader, bm.concurrentDownloads)

	// create workers
	for i := 0; i < bm.concurrentDownloads; i++ {
		worker := BundleDownloader{
			id:       i + 1,
			workChan: make(chan *DownloadRequest),
			bm:       bm,
		}
		bm.workers[i] = &worker
		worker.Start()
	}
}

// download bundle blob and resource blob
// TODO do not download duplicate blobs
func (bm *bundleManager) queueDownloadRequest(dep *DataDeployment) {
	blobReq := bm.makeDownloadRequest(dep.BlobID)
	resourceReq := bm.makeDownloadRequest(dep.BlobResourceID)

	go func() {
		bm.enqueueRequest(blobReq)
		bm.enqueueRequest(resourceReq)
	}()
}

func (bm *bundleManager) makeDownloadRequest(id string) *DownloadRequest {
	markFailedAt := time.Now().Add(bm.markDeploymentFailedAfter)
	retryIn := bm.bundleRetryDelay
	maxBackOff := 5 * time.Minute

	return &DownloadRequest{
		blobServerURL: bm.blobServerUrl,
		bm:            bm,
		blobId:        id,
		backoffFunc:   createBackoff(retryIn, maxBackOff),
		markFailedAt:  markFailedAt,
		connTimeout:   bm.bundleDownloadConnTimeout,
	}
}

// a blocking method to enqueue download requests
func (bm *bundleManager) enqueueRequest(r *DownloadRequest) {
	if atomic.LoadInt32(bm.isClosed) == 1 {
		return
	}
	/*
		defer func() {
			if r := recover(); r != nil {
				log.Warn("trying to enque requests to closed bundleManager")
			}
		}()
	*/
	bm.downloadQueue <- r
}

func (bm *bundleManager) Close() {
	atomic.StoreInt32(bm.isClosed, 1)
	close(bm.downloadQueue)
}

func (bm *bundleManager) deleteBundlesFromDeployments(deletedDeployments []DataDeployment) {
	for _, dep := range deletedDeployments {
		go bm.deleteBundleById(dep.BlobID)
		go bm.deleteBundleById(dep.BlobResourceID)
	}

	/*
		log.Debugf("will delete %d old bundles", len(deletedDeployments))
		go func() {
			// give clients a minute to avoid conflicts
			time.Sleep(bm.bundleCleanupDelay)
			for _, dep := range deletedDeployments {
				bundleFile := getBlobFilePath(dep.BlobID)
				log.Debugf("removing old bundle: %v", bundleFile)
				// TODO Remove from the Database table apid_blob_available
				safeDelete(bundleFile)
			}
		}()
	*/
}

// TODO add delete support
func (bm *bundleManager) deleteBundleById(blobId string) {

}

type DownloadRequest struct {
	bm            *bundleManager
	blobId        string
	backoffFunc   func()
	markFailedAt  time.Time
	connTimeout   time.Duration
	blobServerURL string
}

func (r *DownloadRequest) downloadBundle() error {

	log.Debugf("starting bundle download attempt for blobId=%s", r.blobId)

	if r.checkTimeout() {
		return &timeoutError{
			markFailedAt: r.markFailedAt,
		}
	}

	downloadedFile, err := downloadFromURI(r.blobServerURL, r.blobId, r.connTimeout)

	if err != nil {
		log.Errorf("Unable to download blob file blobId=%s err:%v", r.blobId, err)
		return err
	}

	log.Debugf("blod downloaded. blobid=%s  filepath=%s", r.blobId, downloadedFile)

	err = r.bm.dbMan.updateLocalFsLocation(r.blobId, downloadedFile)
	if err != nil {
		log.Errorf("updateLocalFsLocation failed: blobId=%s", r.blobId)
		return err
	}

	log.Debugf("bundle downloaded: blobId=%s", r.blobId)

	// TODO send changed deployments to subscribers (API call with "block")
	//r.bm.apiMan.addChangedDeployment(dep.ID)

	return nil
}

func (r *DownloadRequest) checkTimeout() bool {

	if !r.markFailedAt.IsZero() && time.Now().After(r.markFailedAt) {
		r.markFailedAt = time.Time{}
		log.Debugf("bundle download timeout. blobId=", r.blobId)
		// TODO notify gateway of this failure
		return true
	}
	return false
}

func getBlobFilePath(blobId string) string {
	return path.Join(bundlePath, base64.StdEncoding.EncodeToString([]byte(blobId)))
}

func getSignedURL(blobServerURL string, blobId string, bundleDownloadConnTimeout time.Duration) (string, error) {

	blobUri, err := url.Parse(blobServerURL)
	if err != nil {
		log.Panicf("bad url value for config %s: %s", blobUri, err)
	}

	blobUri.Path += strings.Replace(blobStoreUri, "{blobId}", blobId, 1)
	parameters := url.Values{}
	parameters.Add("action", "GET")
	blobUri.RawQuery = parameters.Encode()

	uri := blobUri.String()

	surl, err := getURIReader(uri, bundleDownloadConnTimeout)
	if err != nil {
		log.Errorf("Unable to get signed URL from BlobServer %s: %v", uri, err)
		return "", err
	}

	body, err := ioutil.ReadAll(surl)
	if err != nil {
		log.Errorf("Invalid response from BlobServer for {%s} error: {%v}", uri, err)
		return "", err
	}
	res := blobServerResponse{}
	err = json.Unmarshal(body, &res)
	if err != nil {
		log.Errorf("Invalid response from BlobServer for {%s} error: {%v}", uri, err)
		return "", err
	}

	return string(res.signedUrl), nil
}

// downloadFromURI involves retrieving the signed URL for the blob, and storing the resource locally
// after downloading the resource from GCS (via the signed URL)
func downloadFromURI(blobServerURL string, blobId string, bundleDownloadConnTimeout time.Duration) (tempFileName string, err error) {

	var tempFile *os.File
	log.Debugf("Downloading bundle: %s", blobId)

	uri, err := getSignedURL(blobServerURL, blobId, bundleDownloadConnTimeout)
	if err != nil {
		log.Errorf("Unable to get signed URL for blobId {%s}, error : {%v}", blobId, err)
		return
	}

	tempFile, err = ioutil.TempFile(bundlePath, "blob")
	if err != nil {
		log.Errorf("Unable to create temp file: %v", err)
		return
	}
	defer tempFile.Close()
	tempFileName = tempFile.Name()

	var confReader io.ReadCloser
	confReader, err = getURIReader(uri, bundleDownloadConnTimeout)
	if err != nil {
		log.Errorf("Unable to retrieve bundle %s: %v", uri, err)
		return
	}
	defer confReader.Close()

	_, err = io.Copy(tempFile, confReader)
	if err != nil {
		log.Errorf("Unable to write bundle %s: %v", tempFileName, err)
		return
	}

	log.Debugf("Bundle %s downloaded to: %s", uri, tempFileName)
	return
}

// retrieveBundle retrieves bundle data from a URI
func getURIReader(uriString string, bundleDownloadConnTimeout time.Duration) (io.ReadCloser, error) {

	client := http.Client{
		Timeout: bundleDownloadConnTimeout,
	}
	res, err := client.Get(uriString)
	if err != nil {
		return nil, err
	}
	if res.StatusCode != 200 {
		return nil, fmt.Errorf("GET uri %s failed with status %d", uriString, res.StatusCode)
	}
	return res.Body, nil
}

type BundleDownloader struct {
	id       int
	workChan chan *DownloadRequest
	bm       *bundleManager
}

func (w *BundleDownloader) Start() {
	go func() {
		log.Debugf("started bundle downloader %d", w.id)

		for req := range w.bm.downloadQueue {
			log.Debugf("starting download blobId=%s", req.blobId)
			err := req.downloadBundle()
			if err != nil {
				// timeout
				if _, ok := err.(*timeoutError); ok {
					continue
				}
				go func() {
					req.backoffFunc()
					w.bm.enqueueRequest(req)
				}()
			}
		}
		log.Debugf("bundle downloader %d stopped", w.id)
	}()
}

// simple doubling back-off
func createBackoff(retryIn, maxBackOff time.Duration) func() {
	return func() {
		log.Debugf("backoff called. will retry in %s.", retryIn)
		time.Sleep(retryIn)
		retryIn = retryIn * time.Duration(2)
		if retryIn > maxBackOff {
			retryIn = maxBackOff
		}
	}
}

type timeoutError struct {
	markFailedAt time.Time
}

func (e *timeoutError) Error() string {
	return fmt.Sprintf("Timeout. markFailedAt=%v", e.markFailedAt)
}
<|endoftext|>


<commit_before>// Copyright 2018 Marc-Antoine Ruel. All rights reserved.
// Use of this source code is governed under the Apache License, Version 2.0
// that can be found in the LICENSE file.

// Package gohci defines the configuration schemas for 'gohci.yml' and
// '.gohci.yml'.
//
// '.gohci.yml' is found in the repository and defines the checks to run.
//
// 'gohci.yml' is found on the worker itself and defines the http port, webhook
// secret and OAuth2 access token.
package gohci

// WorkerConfig is a worker configuration.
//
// It is found as `gohci.yml` in the gohci-worker working directory.
type WorkerConfig struct {
	// TCP port number for the HTTP server.
	Port int
	// WebHookSecret is the shared secret that keeps people on the internet from
	// running tasks on your worker.
	//
	// gohci-worker generates a good secret by default.
	//
	// See https://developer.github.com/webhooks/ for more information.
	WebHookSecret string
	// Oauth2AccessToken is the OAuth2 Access Token to be able to create gist and
	// update commit status.
	//
	// https://github.com/settings/tokens, check "repo:status" and "gist"
	Oauth2AccessToken string
	// Display name to use in the status report on Github.
	//
	// Defaults to the machine hostname.
	Name string
}

// Check is a single command to run.
type Check struct {
	Cmd []string // Command to run.
	Env []string // Optional environment variables to use.
}

// WorkerProjectConfig is the project configuration via ".gohci.yml" for a
// specific worker.
type WorkerProjectConfig struct {
	// Name is the worker which this config belongs to.
	//
	// If empty, this is the default configuration to use.
	Name string
	// Checks are the commands to run to test the repository. They are run one
	// after the other from the repository's root.
	Checks []Check
}

// ProjectConfig is a configuration file found in a project as ".gohci.yml" in
// the root directory of the repository.
type ProjectConfig struct {
	Version int                   // Current 1
	Workers []WorkerProjectConfig //
}
<commit_msg>gohci: rename to ProjectWorkerConfig so it is ordered<commit_after>// Copyright 2018 Marc-Antoine Ruel. All rights reserved.
// Use of this source code is governed under the Apache License, Version 2.0
// that can be found in the LICENSE file.

// Package gohci defines the configuration schemas for 'gohci.yml' and
// '.gohci.yml'.
//
// '.gohci.yml' is found in the repository and defines the checks to run.
//
// 'gohci.yml' is found on the worker itself and defines the http port, webhook
// secret and OAuth2 access token.
package gohci

// WorkerConfig is a worker configuration.
//
// It is found as `gohci.yml` in the gohci-worker working directory.
type WorkerConfig struct {
	// TCP port number for the HTTP server.
	Port int
	// WebHookSecret is the shared secret that keeps people on the internet from
	// running tasks on your worker.
	//
	// gohci-worker generates a good secret by default.
	//
	// See https://developer.github.com/webhooks/ for more information.
	WebHookSecret string
	// Oauth2AccessToken is the OAuth2 Access Token to be able to create gist and
	// update commit status.
	//
	// https://github.com/settings/tokens, check "repo:status" and "gist"
	Oauth2AccessToken string
	// Display name to use in the status report on Github.
	//
	// Defaults to the machine hostname.
	Name string
}

// Check is a single command to run.
type Check struct {
	Cmd []string // Command to run.
	Env []string // Optional environment variables to use.
}

// ProjectWorkerConfig is the project configuration via ".gohci.yml" for a
// specific worker.
type ProjectWorkerConfig struct {
	// Name is the worker which this config belongs to.
	//
	// If empty, this is the default configuration to use.
	Name string
	// Checks are the commands to run to test the repository. They are run one
	// after the other from the repository's root.
	Checks []Check
}

// ProjectConfig is a configuration file found in a project as ".gohci.yml" in
// the root directory of the repository.
type ProjectConfig struct {
	Version int                   // Current 1
	Workers []ProjectWorkerConfig //
}
<|endoftext|>


<commit_before>package gohex

import (
	"bufio"
	"encoding/hex"
	"io"
	"sort"
)

// Constants definitions of IntelHex record types
const (
	_DATA_RECORD    byte = 0 // Record with data bytes
	_EOF_RECORD     byte = 1 // Record with end of file indicator
	_ADDRESS_RECORD byte = 4 // Record with extended linear address
	_START_RECORD   byte = 5 // Record with start linear address
)

// Structure with binary data segment fields
type DataSegment struct {
	Address uint32 // Starting address of data segment
	Data    []byte // Data segment bytes
}

// Helper type for data segments sorting operations
type sortByAddress []*DataSegment

func (segs sortByAddress) Len() int           { return len(segs) }
func (segs sortByAddress) Swap(i, j int)      { segs[i], segs[j] = segs[j], segs[i] }
func (segs sortByAddress) Less(i, j int) bool { return segs[i].Address < segs[j].Address }

// Main structure with private fields of IntelHex parser
type Memory struct {
	dataSegments     []*DataSegment // Slice with pointers to DataSegments
	startAddress     uint32         // Start linear address
	extendedAddress  uint32         // Extended linear address
	eofFlag          bool           // End of file record exist flag
	startFlag        bool           // Start address record exist flag
	lineNum          uint           // Parser input line number
	firstAddressFlag bool           // Dump first address line
}

// Constructor of Memory structure
func NewMemory() *Memory {
	m := new(Memory)
	m.Clear()
	return m
}

// Method to getting start address from IntelHex data
func (m *Memory) GetStartAddress() (adr uint32, ok bool) {
	if m.startFlag {
		return m.startAddress, true
	}
	return 0, false
}

// Method to setting start address to IntelHex data
func (m *Memory) SetStartAddress(adr uint32) {
	m.startAddress = adr
	m.startFlag = true
}

// Method to getting data segments address from IntelHex data
func (m *Memory) GetDataSegments() []DataSegment {
	segs := []DataSegment{}
	for _, s := range m.dataSegments {
		segs = append(segs, *s)
	}
	return segs
}

// Method to clear memory structure
func (m *Memory) Clear() {
	m.startAddress = 0
	m.extendedAddress = 0
	m.lineNum = 0
	m.dataSegments = []*DataSegment{}
	m.startFlag = false
	m.eofFlag = false
	m.firstAddressFlag = false
}

func (seg *DataSegment) IsOverlap(adr uint32, size uint32) bool {
	if ((adr >= seg.Address) && (adr < seg.Address+uint32(len(seg.Data)))) ||
		((adr < seg.Address) && (adr+size) > seg.Address) {
		return true
	}
	return false
}

func (m *Memory) removeSegment(index int) {
	size := len(m.dataSegments)

	if size == 0 {
		return
	} else if size == 1 {
		m.dataSegments = []*DataSegment{}
	} else {
		if index == 0 {
			m.dataSegments = m.dataSegments[1:]
		} else if index == size-1 {
			m.dataSegments = m.dataSegments[:index]
		} else {
			m.dataSegments = append(m.dataSegments[:index], m.dataSegments[index+1:]...)
		}
	}
}

func (m *Memory) findDataSegment(adr uint32) (seg *DataSegment, offset uint32, index int) {
	for i, s := range m.dataSegments {
		if s.IsOverlap(adr, 1) == true {
			return s, adr - s.Address, i
		}
	}
	return nil, 0, 0
}

// Method to add binary data to memory (auto segmented and sorted)
func (m *Memory) AddBinary(adr uint32, bytes []byte) error {
	var segBefore *DataSegment = nil
	var segAfter *DataSegment = nil
	var segAfterIndex int
	for i, s := range m.dataSegments {
		if s.IsOverlap(adr, uint32(len(bytes))) == true {
			return newParseError(_DATA_ERROR, "data segments overlap", m.lineNum)
		}

		if adr == s.Address+uint32(len(s.Data)) {
			segBefore = s
		}
		if adr+uint32(len(bytes)) == s.Address {
			segAfter, segAfterIndex = s, i
		}
	}

	if segBefore != nil && segAfter != nil {
		segBefore.Data = append(segBefore.Data, bytes...)
		segBefore.Data = append(segBefore.Data, segAfter.Data...)
		m.dataSegments = append(m.dataSegments[:segAfterIndex], m.dataSegments[segAfterIndex+1:]...)

	} else if segBefore != nil && segAfter == nil {
		segBefore.Data = append(segBefore.Data, bytes...)
	} else if segBefore == nil && segAfter != nil {
		segAfter.Address = adr
		segAfter.Data = append(bytes, segAfter.Data...)
	} else {
		m.dataSegments = append(m.dataSegments, &DataSegment{Address: adr, Data: bytes})
	}
	sort.Sort(sortByAddress(m.dataSegments))
	return nil
}

// Method to set binary data to memory (data overlapped will change, auto segmented and sorted)
func (m *Memory) SetBinary(adr uint32, bytes []byte) {
	for a, b := range bytes {
		currentAdr := adr + uint32(a)
		seg, offset, _ := m.findDataSegment(currentAdr)

		if seg != nil {
			seg.Data[offset] = b
		} else {
			m.AddBinary(currentAdr, []byte{b})
		}
	}
}

// Method to remove binary data from memory (auto segmented and sorted)
func (m *Memory) RemoveBinary(adr uint32, size uint32) {
	adrEnd := adr + size
	for currentAdr := adr; currentAdr < adrEnd; currentAdr++ {
		seg, offset, index := m.findDataSegment(currentAdr)

		if seg == nil {
			continue
		}

		if offset == 0 {
			seg.Address += 1
			if len(seg.Data) > 1 {
				seg.Data = seg.Data[1:]
			} else {
				m.removeSegment(index)
			}
		} else if offset == uint32(len(seg.Data)-1) {
			if len(seg.Data) > 1 {
				seg.Data = seg.Data[:offset]
			} else {
				m.removeSegment(index)
			}
		} else {
			newSeg := DataSegment{Address: seg.Address + offset + 1, Data: seg.Data[offset+1:]}
			seg.Data = seg.Data[:offset]
			m.dataSegments = append(m.dataSegments, &newSeg)
		}
	}
	sort.Sort(sortByAddress(m.dataSegments))
}

func (m *Memory) parseIntelHexRecord(bytes []byte) error {
	if len(bytes) < 5 {
		return newParseError(_DATA_ERROR, "not enought data bytes", m.lineNum)
	}
	err := checkSum(bytes)
	if err != nil {
		return newParseError(_CHECKSUM_ERROR, err.Error(), m.lineNum)
	}
	err = checkRecordSize(bytes)
	if err != nil {
		return newParseError(_DATA_ERROR, err.Error(), m.lineNum)
	}
	switch record_type := bytes[3]; record_type {
	case _DATA_RECORD:
		a, data := getDataLine(bytes)
		adr := uint32(a) + m.extendedAddress
		err = m.AddBinary(adr, data)
		if err != nil {
			return err
		}
	case _EOF_RECORD:
		err = checkEOF(bytes)
		if err != nil {
			return newParseError(_RECORD_ERROR, err.Error(), m.lineNum)
		}
		m.eofFlag = true
	case _ADDRESS_RECORD:
		m.extendedAddress, err = getExtendedAddress(bytes)
		if err != nil {
			return newParseError(_RECORD_ERROR, err.Error(), m.lineNum)
		}
	case _START_RECORD:
		if m.startFlag == true {
			return newParseError(_DATA_ERROR, "multiple start address lines", m.lineNum)
		}
		m.startAddress, err = getStartAddress(bytes)
		if err != nil {
			return newParseError(_RECORD_ERROR, err.Error(), m.lineNum)
		}
		m.startFlag = true
	}
	return nil
}

func (m *Memory) parseIntelHexLine(line string) error {
	if len(line) == 0 {
		return nil
	}
	if line[0] != ':' {
		return newParseError(_SYNTAX_ERROR, "no colon char on the first line character", m.lineNum)
	}
	bytes, err := hex.DecodeString(line[1:])
	if err != nil {
		return newParseError(_SYNTAX_ERROR, err.Error(), m.lineNum)
	}
	return m.parseIntelHexRecord(bytes)
}

// Method to parsing IntelHex data and add into memory
func (m *Memory) ParseIntelHex(reader io.Reader) error {
	scanner := bufio.NewScanner(reader)
	m.Clear()
	for scanner.Scan() {
		m.lineNum++
		line := scanner.Text()
		err := m.parseIntelHexLine(line)
		if err != nil {
			return err
		}
	}
	if err := scanner.Err(); err != nil {
		return newParseError(_SYNTAX_ERROR, err.Error(), m.lineNum)
	}
	if m.eofFlag == false {
		return newParseError(_DATA_ERROR, "no end of file line", m.lineNum)
	}

	return nil
}

func (m *Memory) dumpDataSegment(writer io.Writer, s *DataSegment, lineLength byte) error {
	lineAdr := s.Address
	lineData := []byte{}
	for byteAdr := s.Address; byteAdr < s.Address+uint32(len(s.Data)); byteAdr++ {
		if ((byteAdr & 0xFFFF0000) != m.extendedAddress) || (m.firstAddressFlag == false) {
			m.firstAddressFlag = true
			if len(lineData) != 0 {
				err := writeDataLine(writer, &lineAdr, byteAdr, &lineData)
				if err != nil {
					return err
				}
			}
			m.extendedAddress = (byteAdr & 0xFFFF0000)
			writeExtendedAddressLine(writer, m.extendedAddress)
		}
		if len(lineData) >= int(lineLength) {
			err := writeDataLine(writer, &lineAdr, byteAdr, &lineData)
			if err != nil {
				return err
			}
		}
		lineData = append(lineData, s.Data[byteAdr-s.Address])
	}

	if len(lineData) != 0 {
		return writeDataLine(writer, &lineAdr, 0, &lineData)
	}
	return nil
}

// Method to dumping IntelHex data previously loaded into memory
func (m *Memory) DumpIntelHex(writer io.Writer, lineLength byte) error {
	if m.startFlag {
		err := writeStartAddressLine(writer, m.startAddress)
		if err != nil {
			return err
		}
	}

	m.firstAddressFlag = false
	m.extendedAddress = 0
	for _, s := range m.dataSegments {
		err := m.dumpDataSegment(writer, s, lineLength)
		if err != nil {
			return err
		}
	}

	return writeEofLine(writer)
}

// Method to load binary data previously loaded into memory
func (m *Memory) ToBinary(address uint32, size uint32, padding byte) []byte {
	data := make([]byte, size)

	i := uint32(0)
	for i < size {
		ok := false
		for _, s := range m.dataSegments {
			if (address >= s.Address) && (address < s.Address+uint32(len(s.Data))) {
				data[i] = s.Data[address-s.Address]
				i++
				address++
				ok = true
				break
			}
		}
		if ok == false {
			data[i] = padding
			i++
			address++
		}
	}

	return data
}
<commit_msg>IsOverlap method hidded<commit_after>package gohex

import (
	"bufio"
	"encoding/hex"
	"io"
	"sort"
)

// Constants definitions of IntelHex record types
const (
	_DATA_RECORD    byte = 0 // Record with data bytes
	_EOF_RECORD     byte = 1 // Record with end of file indicator
	_ADDRESS_RECORD byte = 4 // Record with extended linear address
	_START_RECORD   byte = 5 // Record with start linear address
)

// Structure with binary data segment fields
type DataSegment struct {
	Address uint32 // Starting address of data segment
	Data    []byte // Data segment bytes
}

// Helper type for data segments sorting operations
type sortByAddress []*DataSegment

func (segs sortByAddress) Len() int           { return len(segs) }
func (segs sortByAddress) Swap(i, j int)      { segs[i], segs[j] = segs[j], segs[i] }
func (segs sortByAddress) Less(i, j int) bool { return segs[i].Address < segs[j].Address }

// Main structure with private fields of IntelHex parser
type Memory struct {
	dataSegments     []*DataSegment // Slice with pointers to DataSegments
	startAddress     uint32         // Start linear address
	extendedAddress  uint32         // Extended linear address
	eofFlag          bool           // End of file record exist flag
	startFlag        bool           // Start address record exist flag
	lineNum          uint           // Parser input line number
	firstAddressFlag bool           // Dump first address line
}

// Constructor of Memory structure
func NewMemory() *Memory {
	m := new(Memory)
	m.Clear()
	return m
}

// Method to getting start address from IntelHex data
func (m *Memory) GetStartAddress() (adr uint32, ok bool) {
	if m.startFlag {
		return m.startAddress, true
	}
	return 0, false
}

// Method to setting start address to IntelHex data
func (m *Memory) SetStartAddress(adr uint32) {
	m.startAddress = adr
	m.startFlag = true
}

// Method to getting data segments address from IntelHex data
func (m *Memory) GetDataSegments() []DataSegment {
	segs := []DataSegment{}
	for _, s := range m.dataSegments {
		segs = append(segs, *s)
	}
	return segs
}

// Method to clear memory structure
func (m *Memory) Clear() {
	m.startAddress = 0
	m.extendedAddress = 0
	m.lineNum = 0
	m.dataSegments = []*DataSegment{}
	m.startFlag = false
	m.eofFlag = false
	m.firstAddressFlag = false
}

func (seg *DataSegment) isOverlap(adr uint32, size uint32) bool {
	if ((adr >= seg.Address) && (adr < seg.Address+uint32(len(seg.Data)))) ||
		((adr < seg.Address) && (adr+size) > seg.Address) {
		return true
	}
	return false
}

func (m *Memory) removeSegment(index int) {
	size := len(m.dataSegments)

	if size == 0 {
		return
	} else if size == 1 {
		m.dataSegments = []*DataSegment{}
	} else {
		if index == 0 {
			m.dataSegments = m.dataSegments[1:]
		} else if index == size-1 {
			m.dataSegments = m.dataSegments[:index]
		} else {
			m.dataSegments = append(m.dataSegments[:index], m.dataSegments[index+1:]...)
		}
	}
}

func (m *Memory) findDataSegment(adr uint32) (seg *DataSegment, offset uint32, index int) {
	for i, s := range m.dataSegments {
		if s.isOverlap(adr, 1) == true {
			return s, adr - s.Address, i
		}
	}
	return nil, 0, 0
}

// Method to add binary data to memory (auto segmented and sorted)
func (m *Memory) AddBinary(adr uint32, bytes []byte) error {
	var segBefore *DataSegment = nil
	var segAfter *DataSegment = nil
	var segAfterIndex int
	for i, s := range m.dataSegments {
		if s.isOverlap(adr, uint32(len(bytes))) == true {
			return newParseError(_DATA_ERROR, "data segments overlap", m.lineNum)
		}

		if adr == s.Address+uint32(len(s.Data)) {
			segBefore = s
		}
		if adr+uint32(len(bytes)) == s.Address {
			segAfter, segAfterIndex = s, i
		}
	}

	if segBefore != nil && segAfter != nil {
		segBefore.Data = append(segBefore.Data, bytes...)
		segBefore.Data = append(segBefore.Data, segAfter.Data...)
		m.dataSegments = append(m.dataSegments[:segAfterIndex], m.dataSegments[segAfterIndex+1:]...)

	} else if segBefore != nil && segAfter == nil {
		segBefore.Data = append(segBefore.Data, bytes...)
	} else if segBefore == nil && segAfter != nil {
		segAfter.Address = adr
		segAfter.Data = append(bytes, segAfter.Data...)
	} else {
		m.dataSegments = append(m.dataSegments, &DataSegment{Address: adr, Data: bytes})
	}
	sort.Sort(sortByAddress(m.dataSegments))
	return nil
}

// Method to set binary data to memory (data overlapped will change, auto segmented and sorted)
func (m *Memory) SetBinary(adr uint32, bytes []byte) {
	for a, b := range bytes {
		currentAdr := adr + uint32(a)
		seg, offset, _ := m.findDataSegment(currentAdr)

		if seg != nil {
			seg.Data[offset] = b
		} else {
			m.AddBinary(currentAdr, []byte{b})
		}
	}
}

// Method to remove binary data from memory (auto segmented and sorted)
func (m *Memory) RemoveBinary(adr uint32, size uint32) {
	adrEnd := adr + size
	for currentAdr := adr; currentAdr < adrEnd; currentAdr++ {
		seg, offset, index := m.findDataSegment(currentAdr)

		if seg == nil {
			continue
		}

		if offset == 0 {
			seg.Address += 1
			if len(seg.Data) > 1 {
				seg.Data = seg.Data[1:]
			} else {
				m.removeSegment(index)
			}
		} else if offset == uint32(len(seg.Data)-1) {
			if len(seg.Data) > 1 {
				seg.Data = seg.Data[:offset]
			} else {
				m.removeSegment(index)
			}
		} else {
			newSeg := DataSegment{Address: seg.Address + offset + 1, Data: seg.Data[offset+1:]}
			seg.Data = seg.Data[:offset]
			m.dataSegments = append(m.dataSegments, &newSeg)
		}
	}
	sort.Sort(sortByAddress(m.dataSegments))
}

func (m *Memory) parseIntelHexRecord(bytes []byte) error {
	if len(bytes) < 5 {
		return newParseError(_DATA_ERROR, "not enought data bytes", m.lineNum)
	}
	err := checkSum(bytes)
	if err != nil {
		return newParseError(_CHECKSUM_ERROR, err.Error(), m.lineNum)
	}
	err = checkRecordSize(bytes)
	if err != nil {
		return newParseError(_DATA_ERROR, err.Error(), m.lineNum)
	}
	switch record_type := bytes[3]; record_type {
	case _DATA_RECORD:
		a, data := getDataLine(bytes)
		adr := uint32(a) + m.extendedAddress
		err = m.AddBinary(adr, data)
		if err != nil {
			return err
		}
	case _EOF_RECORD:
		err = checkEOF(bytes)
		if err != nil {
			return newParseError(_RECORD_ERROR, err.Error(), m.lineNum)
		}
		m.eofFlag = true
	case _ADDRESS_RECORD:
		m.extendedAddress, err = getExtendedAddress(bytes)
		if err != nil {
			return newParseError(_RECORD_ERROR, err.Error(), m.lineNum)
		}
	case _START_RECORD:
		if m.startFlag == true {
			return newParseError(_DATA_ERROR, "multiple start address lines", m.lineNum)
		}
		m.startAddress, err = getStartAddress(bytes)
		if err != nil {
			return newParseError(_RECORD_ERROR, err.Error(), m.lineNum)
		}
		m.startFlag = true
	}
	return nil
}

func (m *Memory) parseIntelHexLine(line string) error {
	if len(line) == 0 {
		return nil
	}
	if line[0] != ':' {
		return newParseError(_SYNTAX_ERROR, "no colon char on the first line character", m.lineNum)
	}
	bytes, err := hex.DecodeString(line[1:])
	if err != nil {
		return newParseError(_SYNTAX_ERROR, err.Error(), m.lineNum)
	}
	return m.parseIntelHexRecord(bytes)
}

// Method to parsing IntelHex data and add into memory
func (m *Memory) ParseIntelHex(reader io.Reader) error {
	scanner := bufio.NewScanner(reader)
	m.Clear()
	for scanner.Scan() {
		m.lineNum++
		line := scanner.Text()
		err := m.parseIntelHexLine(line)
		if err != nil {
			return err
		}
	}
	if err := scanner.Err(); err != nil {
		return newParseError(_SYNTAX_ERROR, err.Error(), m.lineNum)
	}
	if m.eofFlag == false {
		return newParseError(_DATA_ERROR, "no end of file line", m.lineNum)
	}

	return nil
}

func (m *Memory) dumpDataSegment(writer io.Writer, s *DataSegment, lineLength byte) error {
	lineAdr := s.Address
	lineData := []byte{}
	for byteAdr := s.Address; byteAdr < s.Address+uint32(len(s.Data)); byteAdr++ {
		if ((byteAdr & 0xFFFF0000) != m.extendedAddress) || (m.firstAddressFlag == false) {
			m.firstAddressFlag = true
			if len(lineData) != 0 {
				err := writeDataLine(writer, &lineAdr, byteAdr, &lineData)
				if err != nil {
					return err
				}
			}
			m.extendedAddress = (byteAdr & 0xFFFF0000)
			writeExtendedAddressLine(writer, m.extendedAddress)
		}
		if len(lineData) >= int(lineLength) {
			err := writeDataLine(writer, &lineAdr, byteAdr, &lineData)
			if err != nil {
				return err
			}
		}
		lineData = append(lineData, s.Data[byteAdr-s.Address])
	}

	if len(lineData) != 0 {
		return writeDataLine(writer, &lineAdr, 0, &lineData)
	}
	return nil
}

// Method to dumping IntelHex data previously loaded into memory
func (m *Memory) DumpIntelHex(writer io.Writer, lineLength byte) error {
	if m.startFlag {
		err := writeStartAddressLine(writer, m.startAddress)
		if err != nil {
			return err
		}
	}

	m.firstAddressFlag = false
	m.extendedAddress = 0
	for _, s := range m.dataSegments {
		err := m.dumpDataSegment(writer, s, lineLength)
		if err != nil {
			return err
		}
	}

	return writeEofLine(writer)
}

// Method to load binary data previously loaded into memory
func (m *Memory) ToBinary(address uint32, size uint32, padding byte) []byte {
	data := make([]byte, size)

	i := uint32(0)
	for i < size {
		ok := false
		for _, s := range m.dataSegments {
			if (address >= s.Address) && (address < s.Address+uint32(len(s.Data))) {
				data[i] = s.Data[address-s.Address]
				i++
				address++
				ok = true
				break
			}
		}
		if ok == false {
			data[i] = padding
			i++
			address++
		}
	}

	return data
}
<|endoftext|>


<commit_before>package gohll

//**
// HLL++ Implemintation by Micha Gorelick
// paper -- http://im.micha.gd/1dc0z0S
//**

import (
	"errors"
	"github.com/mynameisfiber/gohll/mmh3"
	"math"
)

// Defined the constants used to identify spase vs normal mode HLL
const (
	SPARSE byte = iota
	NORMAL
)

var (
	ErrInvalidP             = errors.New("invalid value of P, must be 4<=p<=25")
	ErrSameP                = errors.New("both HLL instances must have the same value of P")
	ErrErrorRateOutOfBounds = errors.New("error rate must be 0.26>=errorRate>=0.00025390625")
)

// MMH3Hash is the default hasher and uses murmurhash to return a uint64
func MMH3Hash(value string) uint64 {
	h1, _ := mmh3.Hash128(value)
	return h1
}

type HLL struct {
	P uint8

	Hasher func(string) uint64

	m1 uint
	m2 uint

	alpha  float64
	format byte

	tempSet    *tempSet
	sparseList *sparseList

	registers []uint8
}

// NewHLLByError creates a new HLL object with error rate given by `errorRate`.
// The error must be between 26% and 0.0253%
func NewHLLByError(errorRate float64) (*HLL, error) {
	if errorRate < 0.00025390625 || errorRate > 0.26 {
		return nil, ErrErrorRateOutOfBounds
	}
	p := uint8(math.Ceil(math.Log2(math.Pow(1.04/errorRate, 2))))
	return NewHLL(p)
}

// NewHLL creates a new HLL object given a normal mode precision between 4 and
// 25
func NewHLL(p uint8) (*HLL, error) {
	if p < 4 || p > 25 {
		return nil, ErrInvalidP
	}

	m1 := uint(1 << p)
	m2 := uint(1 << 25)

	var alpha float64
	switch m1 {
	case 16:
		alpha = 0.673
	case 32:
		alpha = 0.697
	case 64:
		alpha = 0.709
	default:
		alpha = 0.7213 / (1 + 1.079/float64(m1))
	}

	format := SPARSE

	// Since HLL.registers is a uint8 slice and the SparseList is a uint32
	// slice, we switch from sparse to normal with the sparse list is |m1/4| in
	// size (ie: the same size as the registers would be.
	sparseList := newSparseList(p, int(m1/4))
	tempSet := make(tempSet, 0, int(m1/16))

	return &HLL{
		P:          p,
		Hasher:     MMH3Hash,
		m1:         m1,
		m2:         m2,
		alpha:      alpha,
		format:     format,
		tempSet:    &tempSet,
		sparseList: sparseList,
	}, nil
}

// Add will add the given string value to the HLL using the currently set
// Hasher function
func (h *HLL) Add(value string) {
	hash := h.Hasher(value)
	switch h.format {
	case NORMAL:
		h.addNormal(hash)
	case SPARSE:
		h.addSparse(hash)
	}
}

func (h *HLL) addNormal(hash uint64) {
	index := sliceUint64(hash, 63, 64-h.P)
	w := sliceUint64(hash, 63-h.P, 0) << h.P
	rho := leadingBitUint64(w) + 1
	if h.registers[index] < rho {
		h.registers[index] = rho
	}
}

func (h *HLL) addSparse(hash uint64) {
	k := encodeHash(hash, h.P)
	h.tempSet = h.tempSet.Append(k)
	if h.tempSet.Full() {
		h.mergeSparse()
		h.checkModeChange()
	}
}

func (h *HLL) mergeSparse() {
	h.sparseList.Merge(h.tempSet)
	h.tempSet.Clear()
}

func (h *HLL) checkModeChange() {
	if h.sparseList.Full() {
		h.ToNormal()
	}
}

// ToNormal will convert the current HLL to normal mode, maintaining any data
// already inserted into the structure, if it is in sparse mode
func (h *HLL) ToNormal() {
	if h.format != SPARSE {
		return
	}
	h.format = NORMAL
	h.registers = make([]uint8, h.m1)
	for _, value := range h.sparseList.Data {
		index, rho := decodeHash(value, h.P)
		if h.registers[index] < rho {
			h.registers[index] = rho
		}
	}
	for _, value := range *(h.tempSet) {
		index, rho := decodeHash(value, h.P)
		if h.registers[index] < rho {
			h.registers[index] = rho
		}
	}
	h.tempSet.Clear()
	h.sparseList.Clear()
}

// Cardinality returns the estimated cardinality of the current HLL object
func (h *HLL) Cardinality() float64 {
	var cardinality float64
	switch h.format {
	case NORMAL:
		cardinality = h.cardinalityNormal()
	case SPARSE:
		cardinality = h.cardinalitySparse()
	}
	return cardinality
}

func (h *HLL) cardinalityNormal() float64 {
	var V int
	Ebottom := 0.0
	for _, value := range h.registers {
		Ebottom += math.Pow(2, -1.0*float64(value))
		if value == 0 {
			V++
		}
	}

	return h.cardinalityNormalCorrected(Ebottom, V)
}

func (h *HLL) cardinalityNormalCorrected(Ebottom float64, V int) float64 {
	E := h.alpha * float64(h.m1*h.m1) / Ebottom
	var Eprime float64
	if E < 5*float64(h.m1) {
		Eprime = E - estimateBias(E, h.P)
	} else {
		Eprime = E
	}

	var H float64
	if V != 0 {
		H = linearCounting(h.m1, V)
	} else {
		H = Eprime
	}

	if H <= threshold(h.P) {
		return H
	}
	return Eprime
}

func (h *HLL) cardinalitySparse() float64 {
	h.mergeSparse()
	return linearCounting(h.m2, int(h.m2)-h.sparseList.Len())
}

// Union will merge all data in another HLL object into this one.
func (h *HLL) Union(other *HLL) error {
	if h.P != other.P {
		return ErrSameP
	}
	if other.format == NORMAL {
		if h.format == SPARSE {
			h.ToNormal()
		}
		for i := uint(0); i < h.m1; i++ {
			if other.registers[i] > h.registers[i] {
				h.registers[i] = other.registers[i]
			}
		}
	} else if h.format == NORMAL && other.format == SPARSE {
		other.mergeSparse()
		for _, value := range other.sparseList.Data {
			index, rho := decodeHash(value, h.P)
			if h.registers[index] < rho {
				h.registers[index] = rho
			}
		}
	} else if h.format == SPARSE && other.format == SPARSE {
		h.mergeSparse()
		other.mergeSparse()
		h.sparseList.Merge(other.sparseList)
		h.checkModeChange()
	}
	return nil
}

// CardinalityIntersection returns the estimated cardinality of the
// intersection between this HLL object and another one.  That is, it returns
// an estimate of the number of unique items that occur in both this and the
// other HLL object.  This is done with the Inclusion–exclusion principle and
// does not satisfy the error guarantee.
func (h *HLL) CardinalityIntersection(other *HLL) (float64, error) {
	if h.P != other.P {
		return 0.0, ErrSameP
	}
	A := h.Cardinality()
	B := other.Cardinality()
	AuB, _ := h.CardinalityUnion(other)
	return A + B - AuB, nil
}

// CardinalityUnion returns the estimated cardinality of the union between this
// and another HLL object.  This result would be the same as first taking the
// union between this and the other object and then calling Cardinality.
// However, by calling this function we are not making any changes to the HLL
// object.
func (h *HLL) CardinalityUnion(other *HLL) (float64, error) {
	if h.P != other.P {
		return 0.0, ErrSameP
	}
	cardinality := 0.0
	if h.format == NORMAL && other.format == NORMAL {
		cardinality = h.cardinalityUnionNN(other)
	} else if h.format == NORMAL && other.format == SPARSE {
		cardinality = h.cardinalityUnionNS(other)
	} else if h.format == SPARSE && other.format == NORMAL {
		cardinality, _ = other.CardinalityUnion(h)
	} else if h.format == SPARSE && other.format == SPARSE {
		cardinality = h.cardinalityUnionSS(other)
	}
	return cardinality, nil
}

func (h *HLL) cardinalityUnionNN(other *HLL) float64 {
	var V int
	Ebottom := 0.0
	for i, value := range h.registers {
		if other.registers[i] > value {
			value = other.registers[i]
		}
		Ebottom += math.Pow(2, -1.0*float64(value))
		if value == 0 {
			V++
		}
	}
	return h.cardinalityNormalCorrected(Ebottom, V)
}

func (h *HLL) cardinalityUnionNS(other *HLL) float64 {
	var V int
	other.mergeSparse()
	registerOther := make([]uint8, h.m1)
	for _, value := range other.sparseList.Data {
		index, rho := decodeHash(value, other.P)
		if registerOther[index] < rho {
			registerOther[index] = rho
		}
	}
	Ebottom := 0.0
	for i, value := range h.registers {
		if registerOther[i] > value {
			value = registerOther[i]
		}
		Ebottom += math.Pow(2, -1.0*float64(value))
		if value == 0 {
			V++
		}
	}
	registerOther = registerOther[:0]
	return h.cardinalityNormalCorrected(Ebottom, V)
}

func (h *HLL) cardinalityUnionSS(other *HLL) float64 {
	h.mergeSparse()
	other.mergeSparse()
	if h.sparseList.Len() == 0 {
		return other.Cardinality()
	} else if other.sparseList.Len() == 0 {
		return h.Cardinality()
	}
	var i, j, V int
	var idxH, idxOther uint32
	for i < h.sparseList.Len()-1 || j < other.sparseList.Len()-1 {
		if i < h.sparseList.Len() {
			idxH = getIndexSparse(h.sparseList.Get(i))
		}
		if j < other.sparseList.Len() {
			idxOther = getIndexSparse(other.sparseList.Get(j))
		}
		V++
		if idxH < idxOther {
			i++
		} else if idxH > idxOther {
			j++
		} else {
			i++
			j++
		}
	}
	return linearCounting(h.m2, int(h.m2)-V)
}
<commit_msg>docstring<commit_after>package gohll

//**
// HLL++ Implemintation by Micha Gorelick
// paper -- http://im.micha.gd/1dc0z0S
//**

import (
	"errors"
	"github.com/mynameisfiber/gohll/mmh3"
	"math"
)

// Defined the constants used to identify spase vs normal mode HLL
const (
	SPARSE byte = iota
	NORMAL
)

var (
	// ErrInvalidP is returned if an invalid precision is requested
	ErrInvalidP = errors.New("invalid value of P, must be 4<=p<=25")

	// ErrSameP is returned if an operation is requested between two HLL
	// objects with different precisions
	ErrSameP = errors.New("both HLL instances must have the same value of P")

	// ErrErrorRateOutOfBounds is returned if an invalid error rate is
	// requested
	ErrErrorRateOutOfBounds = errors.New("error rate must be 0.26>=errorRate>=0.00025390625")
)

// MMH3Hash is the default hasher and uses murmurhash to return a uint64
// NOTE: This hashing function will clobber the original hash
func MMH3Hash(value string) uint64 {
	h1, _ := mmh3.Hash128(value)
	return h1
}

// HLL is the structure holding the HLL registers and maintains state.  State
// includes:
// - Whether we are in normal or spase mode
// - Register values
// - Desired precision
// - Reference to the hashing function used
type HLL struct {
	P uint8

	Hasher func(string) uint64

	m1 uint
	m2 uint

	alpha  float64
	format byte

	tempSet    *tempSet
	sparseList *sparseList

	registers []uint8
}

// NewHLLByError creates a new HLL object with error rate given by `errorRate`.
// The error must be between 26% and 0.0253%
func NewHLLByError(errorRate float64) (*HLL, error) {
	if errorRate < 0.00025390625 || errorRate > 0.26 {
		return nil, ErrErrorRateOutOfBounds
	}
	p := uint8(math.Ceil(math.Log2(math.Pow(1.04/errorRate, 2))))
	return NewHLL(p)
}

// NewHLL creates a new HLL object given a normal mode precision between 4 and
// 25
func NewHLL(p uint8) (*HLL, error) {
	if p < 4 || p > 25 {
		return nil, ErrInvalidP
	}

	m1 := uint(1 << p)
	m2 := uint(1 << 25)

	var alpha float64
	switch m1 {
	case 16:
		alpha = 0.673
	case 32:
		alpha = 0.697
	case 64:
		alpha = 0.709
	default:
		alpha = 0.7213 / (1 + 1.079/float64(m1))
	}

	format := SPARSE

	// Since HLL.registers is a uint8 slice and the SparseList is a uint32
	// slice, we switch from sparse to normal with the sparse list is |m1/4| in
	// size (ie: the same size as the registers would be.
	sparseList := newSparseList(p, int(m1/4))
	tempSet := make(tempSet, 0, int(m1/16))

	return &HLL{
		P:          p,
		Hasher:     MMH3Hash,
		m1:         m1,
		m2:         m2,
		alpha:      alpha,
		format:     format,
		tempSet:    &tempSet,
		sparseList: sparseList,
	}, nil
}

// Add will add the given string value to the HLL using the currently set
// Hasher function
func (h *HLL) Add(value string) {
	hash := h.Hasher(value)
	switch h.format {
	case NORMAL:
		h.addNormal(hash)
	case SPARSE:
		h.addSparse(hash)
	}
}

func (h *HLL) addNormal(hash uint64) {
	index := sliceUint64(hash, 63, 64-h.P)
	w := sliceUint64(hash, 63-h.P, 0) << h.P
	rho := leadingBitUint64(w) + 1
	if h.registers[index] < rho {
		h.registers[index] = rho
	}
}

func (h *HLL) addSparse(hash uint64) {
	k := encodeHash(hash, h.P)
	h.tempSet = h.tempSet.Append(k)
	if h.tempSet.Full() {
		h.mergeSparse()
		h.checkModeChange()
	}
}

func (h *HLL) mergeSparse() {
	h.sparseList.Merge(h.tempSet)
	h.tempSet.Clear()
}

func (h *HLL) checkModeChange() {
	if h.sparseList.Full() {
		h.ToNormal()
	}
}

// ToNormal will convert the current HLL to normal mode, maintaining any data
// already inserted into the structure, if it is in sparse mode
func (h *HLL) ToNormal() {
	if h.format != SPARSE {
		return
	}
	h.format = NORMAL
	h.registers = make([]uint8, h.m1)
	for _, value := range h.sparseList.Data {
		index, rho := decodeHash(value, h.P)
		if h.registers[index] < rho {
			h.registers[index] = rho
		}
	}
	for _, value := range *(h.tempSet) {
		index, rho := decodeHash(value, h.P)
		if h.registers[index] < rho {
			h.registers[index] = rho
		}
	}
	h.tempSet.Clear()
	h.sparseList.Clear()
}

// Cardinality returns the estimated cardinality of the current HLL object
func (h *HLL) Cardinality() float64 {
	var cardinality float64
	switch h.format {
	case NORMAL:
		cardinality = h.cardinalityNormal()
	case SPARSE:
		cardinality = h.cardinalitySparse()
	}
	return cardinality
}

func (h *HLL) cardinalityNormal() float64 {
	var V int
	Ebottom := 0.0
	for _, value := range h.registers {
		Ebottom += math.Pow(2, -1.0*float64(value))
		if value == 0 {
			V++
		}
	}

	return h.cardinalityNormalCorrected(Ebottom, V)
}

func (h *HLL) cardinalityNormalCorrected(Ebottom float64, V int) float64 {
	E := h.alpha * float64(h.m1*h.m1) / Ebottom
	var Eprime float64
	if E < 5*float64(h.m1) {
		Eprime = E - estimateBias(E, h.P)
	} else {
		Eprime = E
	}

	var H float64
	if V != 0 {
		H = linearCounting(h.m1, V)
	} else {
		H = Eprime
	}

	if H <= threshold(h.P) {
		return H
	}
	return Eprime
}

func (h *HLL) cardinalitySparse() float64 {
	h.mergeSparse()
	return linearCounting(h.m2, int(h.m2)-h.sparseList.Len())
}

// Union will merge all data in another HLL object into this one.
func (h *HLL) Union(other *HLL) error {
	if h.P != other.P {
		return ErrSameP
	}
	if other.format == NORMAL {
		if h.format == SPARSE {
			h.ToNormal()
		}
		for i := uint(0); i < h.m1; i++ {
			if other.registers[i] > h.registers[i] {
				h.registers[i] = other.registers[i]
			}
		}
	} else if h.format == NORMAL && other.format == SPARSE {
		other.mergeSparse()
		for _, value := range other.sparseList.Data {
			index, rho := decodeHash(value, h.P)
			if h.registers[index] < rho {
				h.registers[index] = rho
			}
		}
	} else if h.format == SPARSE && other.format == SPARSE {
		h.mergeSparse()
		other.mergeSparse()
		h.sparseList.Merge(other.sparseList)
		h.checkModeChange()
	}
	return nil
}

// CardinalityIntersection returns the estimated cardinality of the
// intersection between this HLL object and another one.  That is, it returns
// an estimate of the number of unique items that occur in both this and the
// other HLL object.  This is done with the Inclusion–exclusion principle and
// does not satisfy the error guarantee.
func (h *HLL) CardinalityIntersection(other *HLL) (float64, error) {
	if h.P != other.P {
		return 0.0, ErrSameP
	}
	A := h.Cardinality()
	B := other.Cardinality()
	AuB, _ := h.CardinalityUnion(other)
	return A + B - AuB, nil
}

// CardinalityUnion returns the estimated cardinality of the union between this
// and another HLL object.  This result would be the same as first taking the
// union between this and the other object and then calling Cardinality.
// However, by calling this function we are not making any changes to the HLL
// object.
func (h *HLL) CardinalityUnion(other *HLL) (float64, error) {
	if h.P != other.P {
		return 0.0, ErrSameP
	}
	cardinality := 0.0
	if h.format == NORMAL && other.format == NORMAL {
		cardinality = h.cardinalityUnionNN(other)
	} else if h.format == NORMAL && other.format == SPARSE {
		cardinality = h.cardinalityUnionNS(other)
	} else if h.format == SPARSE && other.format == NORMAL {
		cardinality, _ = other.CardinalityUnion(h)
	} else if h.format == SPARSE && other.format == SPARSE {
		cardinality = h.cardinalityUnionSS(other)
	}
	return cardinality, nil
}

func (h *HLL) cardinalityUnionNN(other *HLL) float64 {
	var V int
	Ebottom := 0.0
	for i, value := range h.registers {
		if other.registers[i] > value {
			value = other.registers[i]
		}
		Ebottom += math.Pow(2, -1.0*float64(value))
		if value == 0 {
			V++
		}
	}
	return h.cardinalityNormalCorrected(Ebottom, V)
}

func (h *HLL) cardinalityUnionNS(other *HLL) float64 {
	var V int
	other.mergeSparse()
	registerOther := make([]uint8, h.m1)
	for _, value := range other.sparseList.Data {
		index, rho := decodeHash(value, other.P)
		if registerOther[index] < rho {
			registerOther[index] = rho
		}
	}
	Ebottom := 0.0
	for i, value := range h.registers {
		if registerOther[i] > value {
			value = registerOther[i]
		}
		Ebottom += math.Pow(2, -1.0*float64(value))
		if value == 0 {
			V++
		}
	}
	registerOther = registerOther[:0]
	return h.cardinalityNormalCorrected(Ebottom, V)
}

func (h *HLL) cardinalityUnionSS(other *HLL) float64 {
	h.mergeSparse()
	other.mergeSparse()
	if h.sparseList.Len() == 0 {
		return other.Cardinality()
	} else if other.sparseList.Len() == 0 {
		return h.Cardinality()
	}
	var i, j, V int
	var idxH, idxOther uint32
	for i < h.sparseList.Len()-1 || j < other.sparseList.Len()-1 {
		if i < h.sparseList.Len() {
			idxH = getIndexSparse(h.sparseList.Get(i))
		}
		if j < other.sparseList.Len() {
			idxOther = getIndexSparse(other.sparseList.Get(j))
		}
		V++
		if idxH < idxOther {
			i++
		} else if idxH > idxOther {
			j++
		} else {
			i++
			j++
		}
	}
	return linearCounting(h.m2, int(h.m2)-V)
}
<|endoftext|>


<commit_before>// Package udp implements a BitTorrent tracker via the UDP protocol as
// described in BEP 15.
package udp

import (
	"bytes"
	"context"
	"encoding/binary"
	"math/rand"
	"net"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"

	"github.com/chihaya/chihaya/bittorrent"
	"github.com/chihaya/chihaya/frontend"
	"github.com/chihaya/chihaya/frontend/udp/bytepool"
	"github.com/chihaya/chihaya/pkg/log"
	"github.com/chihaya/chihaya/pkg/stop"
	"github.com/chihaya/chihaya/pkg/timecache"
)

var allowedGeneratedPrivateKeyRunes = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890")

func init() {
	prometheus.MustRegister(promResponseDurationMilliseconds)
}

var promResponseDurationMilliseconds = prometheus.NewHistogramVec(
	prometheus.HistogramOpts{
		Name:    "chihaya_udp_response_duration_milliseconds",
		Help:    "The duration of time it takes to receive and write a response to an API request",
		Buckets: prometheus.ExponentialBuckets(9.375, 2, 10),
	},
	[]string{"action", "address_family", "error"},
)

// recordResponseDuration records the duration of time to respond to a UDP
// Request in milliseconds .
func recordResponseDuration(action string, af *bittorrent.AddressFamily, err error, duration time.Duration) {
	var errString string
	if err != nil {
		if _, ok := err.(bittorrent.ClientError); ok {
			errString = err.Error()
		} else {
			errString = "internal error"
		}
	}

	var afString string
	if af == nil {
		afString = "Unknown"
	} else if *af == bittorrent.IPv4 {
		afString = "IPv4"
	} else if *af == bittorrent.IPv6 {
		afString = "IPv6"
	}

	promResponseDurationMilliseconds.
		WithLabelValues(action, afString, errString).
		Observe(float64(duration.Nanoseconds()) / float64(time.Millisecond))
}

// Config represents all of the configurable options for a UDP BitTorrent
// Tracker.
type Config struct {
	Addr                string        `yaml:"addr"`
	PrivateKey          string        `yaml:"private_key"`
	MaxClockSkew        time.Duration `yaml:"max_clock_skew"`
	EnableRequestTiming bool          `yaml:"enable_request_timing"`
	ParseOptions        `yaml:",inline"`
}

// LogFields renders the current config as a set of Logrus fields.
func (cfg Config) LogFields() log.Fields {
	return log.Fields{
		"addr":                cfg.Addr,
		"privateKey":          cfg.PrivateKey,
		"maxClockSkew":        cfg.MaxClockSkew,
		"enableRequestTiming": cfg.EnableRequestTiming,
		"allowIPSpoofing":     cfg.AllowIPSpoofing,
		"maxNumWant":          cfg.MaxNumWant,
		"defaultNumWant":      cfg.DefaultNumWant,
		"maxScrapeInfoHashes": cfg.MaxScrapeInfoHashes,
	}
}

// Frontend holds the state of a UDP BitTorrent Frontend.
type Frontend struct {
	socket  *net.UDPConn
	closing chan struct{}
	wg      sync.WaitGroup

	logic frontend.TrackerLogic
	Config
}

// NewFrontend creates a new instance of an UDP Frontend that asynchronously
// serves requests.
func NewFrontend(logic frontend.TrackerLogic, cfg Config) (*Frontend, error) {
	// Generate a private key if one isn't provided by the user.
	if cfg.PrivateKey == "" {
		rand.Seed(time.Now().UnixNano())
		pkeyRunes := make([]rune, 64)
		for i := range pkeyRunes {
			pkeyRunes[i] = allowedGeneratedPrivateKeyRunes[rand.Intn(len(allowedGeneratedPrivateKeyRunes))]
		}
		cfg.PrivateKey = string(pkeyRunes)

		log.Warn("UDP private key was not provided, using generated key", log.Fields{"key": cfg.PrivateKey})
	}

	f := &Frontend{
		closing: make(chan struct{}),
		logic:   logic,
		Config:  cfg,
	}

	go func() {
		if err := f.listenAndServe(); err != nil {
			log.Fatal("failed while serving udp", log.Err(err))
		}
	}()

	return f, nil
}

// Stop provides a thread-safe way to shutdown a currently running Frontend.
func (t *Frontend) Stop() <-chan error {
	select {
	case <-t.closing:
		return stop.AlreadyStopped
	default:
	}

	c := make(chan error)
	go func() {
		close(t.closing)
		t.socket.SetReadDeadline(time.Now())
		t.wg.Wait()
		if err := t.socket.Close(); err != nil {
			c <- err
		} else {
			close(c)
		}
	}()

	return c
}

// listenAndServe blocks while listening and serving UDP BitTorrent requests
// until Stop() is called or an error is returned.
func (t *Frontend) listenAndServe() error {
	udpAddr, err := net.ResolveUDPAddr("udp", t.Addr)
	if err != nil {
		return err
	}

	t.socket, err = net.ListenUDP("udp", udpAddr)
	if err != nil {
		return err
	}

	pool := bytepool.New(2048)

	t.wg.Add(1)
	defer t.wg.Done()

	for {
		// Check to see if we need to shutdown.
		select {
		case <-t.closing:
			log.Debug("udp listenAndServe() received shutdown signal")
			return nil
		default:
		}

		// Read a UDP packet into a reusable buffer.
		buffer := pool.Get()
		n, addr, err := t.socket.ReadFromUDP(buffer)
		if err != nil {
			pool.Put(buffer)
			if netErr, ok := err.(net.Error); ok && netErr.Temporary() {
				// A temporary failure is not fatal; just pretend it never happened.
				continue
			}
			return err
		}

		// We got nothin'
		if n == 0 {
			pool.Put(buffer)
			continue
		}

		t.wg.Add(1)
		go func() {
			defer t.wg.Done()
			defer pool.Put(buffer)

			if ip := addr.IP.To4(); ip != nil {
				addr.IP = ip
			}

			// Handle the request.
			var start time.Time
			if t.EnableRequestTiming {
				start = time.Now()
			}
			action, af, err := t.handleRequest(
				// Make sure the IP is copied, not referenced.
				Request{buffer[:n], append([]byte{}, addr.IP...)},
				ResponseWriter{t.socket, addr},
			)
			if t.EnableRequestTiming {
				recordResponseDuration(action, af, err, time.Since(start))
			} else {
				recordResponseDuration(action, af, err, time.Duration(0))
			}
		}()
	}
}

// Request represents a UDP payload received by a Tracker.
type Request struct {
	Packet []byte
	IP     net.IP
}

// ResponseWriter implements the ability to respond to a Request via the
// io.Writer interface.
type ResponseWriter struct {
	socket *net.UDPConn
	addr   *net.UDPAddr
}

// Write implements the io.Writer interface for a ResponseWriter.
func (w ResponseWriter) Write(b []byte) (int, error) {
	w.socket.WriteToUDP(b, w.addr)
	return len(b), nil
}

// handleRequest parses and responds to a UDP Request.
func (t *Frontend) handleRequest(r Request, w ResponseWriter) (actionName string, af *bittorrent.AddressFamily, err error) {
	if len(r.Packet) < 16 {
		// Malformed, no client packets are less than 16 bytes.
		// We explicitly return nothing in case this is a DoS attempt.
		err = errMalformedPacket
		return
	}

	// Parse the headers of the UDP packet.
	connID := r.Packet[0:8]
	actionID := binary.BigEndian.Uint32(r.Packet[8:12])
	txID := r.Packet[12:16]

	// If this isn't requesting a new connection ID and the connection ID is
	// invalid, then fail.
	if actionID != connectActionID && !ValidConnectionID(connID, r.IP, timecache.Now(), t.MaxClockSkew, t.PrivateKey) {
		err = errBadConnectionID
		WriteError(w, txID, err)
		return
	}

	// Handle the requested action.
	switch actionID {
	case connectActionID:
		actionName = "connect"

		if !bytes.Equal(connID, initialConnectionID) {
			err = errMalformedPacket
			return
		}

		WriteConnectionID(w, txID, NewConnectionID(r.IP, timecache.Now(), t.PrivateKey))

	case announceActionID, announceV6ActionID:
		actionName = "announce"

		var req *bittorrent.AnnounceRequest
		req, err = ParseAnnounce(r, actionID == announceV6ActionID, t.ParseOptions)
		if err != nil {
			WriteError(w, txID, err)
			return
		}
		af = new(bittorrent.AddressFamily)
		*af = req.IP.AddressFamily

		var ctx context.Context
		var resp *bittorrent.AnnounceResponse
		ctx, resp, err = t.logic.HandleAnnounce(context.Background(), req)
		if err != nil {
			WriteError(w, txID, err)
			return
		}

		WriteAnnounce(w, txID, resp, actionID == announceV6ActionID)

		go t.logic.AfterAnnounce(ctx, req, resp)

	case scrapeActionID:
		actionName = "scrape"

		var req *bittorrent.ScrapeRequest
		req, err = ParseScrape(r, t.ParseOptions)
		if err != nil {
			WriteError(w, txID, err)
			return
		}

		if r.IP.To4() != nil {
			req.AddressFamily = bittorrent.IPv4
		} else if len(r.IP) == net.IPv6len { // implies r.IP.To4() == nil
			req.AddressFamily = bittorrent.IPv6
		} else {
			log.Error("udp: invalid IP: neither v4 nor v6", log.Fields{"IP": r.IP})
			WriteError(w, txID, bittorrent.ErrInvalidIP)
			return
		}
		af = new(bittorrent.AddressFamily)
		*af = req.AddressFamily

		var ctx context.Context
		var resp *bittorrent.ScrapeResponse
		ctx, resp, err = t.logic.HandleScrape(context.Background(), req)
		if err != nil {
			WriteError(w, txID, err)
			return
		}

		WriteScrape(w, txID, resp)

		go t.logic.AfterScrape(ctx, req, resp)

	default:
		err = errUnknownAction
		WriteError(w, txID, err)
	}

	return
}
<commit_msg>udp: set address family for connect metrics<commit_after>// Package udp implements a BitTorrent tracker via the UDP protocol as
// described in BEP 15.
package udp

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"math/rand"
	"net"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"

	"github.com/chihaya/chihaya/bittorrent"
	"github.com/chihaya/chihaya/frontend"
	"github.com/chihaya/chihaya/frontend/udp/bytepool"
	"github.com/chihaya/chihaya/pkg/log"
	"github.com/chihaya/chihaya/pkg/stop"
	"github.com/chihaya/chihaya/pkg/timecache"
)

var allowedGeneratedPrivateKeyRunes = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890")

func init() {
	prometheus.MustRegister(promResponseDurationMilliseconds)
}

var promResponseDurationMilliseconds = prometheus.NewHistogramVec(
	prometheus.HistogramOpts{
		Name:    "chihaya_udp_response_duration_milliseconds",
		Help:    "The duration of time it takes to receive and write a response to an API request",
		Buckets: prometheus.ExponentialBuckets(9.375, 2, 10),
	},
	[]string{"action", "address_family", "error"},
)

// recordResponseDuration records the duration of time to respond to a UDP
// Request in milliseconds .
func recordResponseDuration(action string, af *bittorrent.AddressFamily, err error, duration time.Duration) {
	var errString string
	if err != nil {
		if _, ok := err.(bittorrent.ClientError); ok {
			errString = err.Error()
		} else {
			errString = "internal error"
		}
	}

	var afString string
	if af == nil {
		afString = "Unknown"
	} else if *af == bittorrent.IPv4 {
		afString = "IPv4"
	} else if *af == bittorrent.IPv6 {
		afString = "IPv6"
	}

	promResponseDurationMilliseconds.
		WithLabelValues(action, afString, errString).
		Observe(float64(duration.Nanoseconds()) / float64(time.Millisecond))
}

// Config represents all of the configurable options for a UDP BitTorrent
// Tracker.
type Config struct {
	Addr                string        `yaml:"addr"`
	PrivateKey          string        `yaml:"private_key"`
	MaxClockSkew        time.Duration `yaml:"max_clock_skew"`
	EnableRequestTiming bool          `yaml:"enable_request_timing"`
	ParseOptions        `yaml:",inline"`
}

// LogFields renders the current config as a set of Logrus fields.
func (cfg Config) LogFields() log.Fields {
	return log.Fields{
		"addr":                cfg.Addr,
		"privateKey":          cfg.PrivateKey,
		"maxClockSkew":        cfg.MaxClockSkew,
		"enableRequestTiming": cfg.EnableRequestTiming,
		"allowIPSpoofing":     cfg.AllowIPSpoofing,
		"maxNumWant":          cfg.MaxNumWant,
		"defaultNumWant":      cfg.DefaultNumWant,
		"maxScrapeInfoHashes": cfg.MaxScrapeInfoHashes,
	}
}

// Frontend holds the state of a UDP BitTorrent Frontend.
type Frontend struct {
	socket  *net.UDPConn
	closing chan struct{}
	wg      sync.WaitGroup

	logic frontend.TrackerLogic
	Config
}

// NewFrontend creates a new instance of an UDP Frontend that asynchronously
// serves requests.
func NewFrontend(logic frontend.TrackerLogic, cfg Config) (*Frontend, error) {
	// Generate a private key if one isn't provided by the user.
	if cfg.PrivateKey == "" {
		rand.Seed(time.Now().UnixNano())
		pkeyRunes := make([]rune, 64)
		for i := range pkeyRunes {
			pkeyRunes[i] = allowedGeneratedPrivateKeyRunes[rand.Intn(len(allowedGeneratedPrivateKeyRunes))]
		}
		cfg.PrivateKey = string(pkeyRunes)

		log.Warn("UDP private key was not provided, using generated key", log.Fields{"key": cfg.PrivateKey})
	}

	f := &Frontend{
		closing: make(chan struct{}),
		logic:   logic,
		Config:  cfg,
	}

	go func() {
		if err := f.listenAndServe(); err != nil {
			log.Fatal("failed while serving udp", log.Err(err))
		}
	}()

	return f, nil
}

// Stop provides a thread-safe way to shutdown a currently running Frontend.
func (t *Frontend) Stop() <-chan error {
	select {
	case <-t.closing:
		return stop.AlreadyStopped
	default:
	}

	c := make(chan error)
	go func() {
		close(t.closing)
		t.socket.SetReadDeadline(time.Now())
		t.wg.Wait()
		if err := t.socket.Close(); err != nil {
			c <- err
		} else {
			close(c)
		}
	}()

	return c
}

// listenAndServe blocks while listening and serving UDP BitTorrent requests
// until Stop() is called or an error is returned.
func (t *Frontend) listenAndServe() error {
	udpAddr, err := net.ResolveUDPAddr("udp", t.Addr)
	if err != nil {
		return err
	}

	t.socket, err = net.ListenUDP("udp", udpAddr)
	if err != nil {
		return err
	}

	pool := bytepool.New(2048)

	t.wg.Add(1)
	defer t.wg.Done()

	for {
		// Check to see if we need to shutdown.
		select {
		case <-t.closing:
			log.Debug("udp listenAndServe() received shutdown signal")
			return nil
		default:
		}

		// Read a UDP packet into a reusable buffer.
		buffer := pool.Get()
		n, addr, err := t.socket.ReadFromUDP(buffer)
		if err != nil {
			pool.Put(buffer)
			if netErr, ok := err.(net.Error); ok && netErr.Temporary() {
				// A temporary failure is not fatal; just pretend it never happened.
				continue
			}
			return err
		}

		// We got nothin'
		if n == 0 {
			pool.Put(buffer)
			continue
		}

		t.wg.Add(1)
		go func() {
			defer t.wg.Done()
			defer pool.Put(buffer)

			if ip := addr.IP.To4(); ip != nil {
				addr.IP = ip
			}

			// Handle the request.
			var start time.Time
			if t.EnableRequestTiming {
				start = time.Now()
			}
			action, af, err := t.handleRequest(
				// Make sure the IP is copied, not referenced.
				Request{buffer[:n], append([]byte{}, addr.IP...)},
				ResponseWriter{t.socket, addr},
			)
			if t.EnableRequestTiming {
				recordResponseDuration(action, af, err, time.Since(start))
			} else {
				recordResponseDuration(action, af, err, time.Duration(0))
			}
		}()
	}
}

// Request represents a UDP payload received by a Tracker.
type Request struct {
	Packet []byte
	IP     net.IP
}

// ResponseWriter implements the ability to respond to a Request via the
// io.Writer interface.
type ResponseWriter struct {
	socket *net.UDPConn
	addr   *net.UDPAddr
}

// Write implements the io.Writer interface for a ResponseWriter.
func (w ResponseWriter) Write(b []byte) (int, error) {
	w.socket.WriteToUDP(b, w.addr)
	return len(b), nil
}

// handleRequest parses and responds to a UDP Request.
func (t *Frontend) handleRequest(r Request, w ResponseWriter) (actionName string, af *bittorrent.AddressFamily, err error) {
	if len(r.Packet) < 16 {
		// Malformed, no client packets are less than 16 bytes.
		// We explicitly return nothing in case this is a DoS attempt.
		err = errMalformedPacket
		return
	}

	// Parse the headers of the UDP packet.
	connID := r.Packet[0:8]
	actionID := binary.BigEndian.Uint32(r.Packet[8:12])
	txID := r.Packet[12:16]

	// If this isn't requesting a new connection ID and the connection ID is
	// invalid, then fail.
	if actionID != connectActionID && !ValidConnectionID(connID, r.IP, timecache.Now(), t.MaxClockSkew, t.PrivateKey) {
		err = errBadConnectionID
		WriteError(w, txID, err)
		return
	}

	// Handle the requested action.
	switch actionID {
	case connectActionID:
		actionName = "connect"

		if !bytes.Equal(connID, initialConnectionID) {
			err = errMalformedPacket
			return
		}

		af = new(bittorrent.AddressFamily)
		if r.IP.To4() != nil {
			*af = bittorrent.IPv4
		} else if len(r.IP) == net.IPv6len { // implies r.IP.To4() == nil
			*af = bittorrent.IPv6
		} else {
			// Should never happen - we got the IP straight from the UDP packet.
			panic(fmt.Sprintf("udp: invalid IP: neither v4 nor v6, IP: %#v", r.IP))
		}

		WriteConnectionID(w, txID, NewConnectionID(r.IP, timecache.Now(), t.PrivateKey))

	case announceActionID, announceV6ActionID:
		actionName = "announce"

		var req *bittorrent.AnnounceRequest
		req, err = ParseAnnounce(r, actionID == announceV6ActionID, t.ParseOptions)
		if err != nil {
			WriteError(w, txID, err)
			return
		}
		af = new(bittorrent.AddressFamily)
		*af = req.IP.AddressFamily

		var ctx context.Context
		var resp *bittorrent.AnnounceResponse
		ctx, resp, err = t.logic.HandleAnnounce(context.Background(), req)
		if err != nil {
			WriteError(w, txID, err)
			return
		}

		WriteAnnounce(w, txID, resp, actionID == announceV6ActionID)

		go t.logic.AfterAnnounce(ctx, req, resp)

	case scrapeActionID:
		actionName = "scrape"

		var req *bittorrent.ScrapeRequest
		req, err = ParseScrape(r, t.ParseOptions)
		if err != nil {
			WriteError(w, txID, err)
			return
		}

		if r.IP.To4() != nil {
			req.AddressFamily = bittorrent.IPv4
		} else if len(r.IP) == net.IPv6len { // implies r.IP.To4() == nil
			req.AddressFamily = bittorrent.IPv6
		} else {
			// Should never happen - we got the IP straight from the UDP packet.
			panic(fmt.Sprintf("udp: invalid IP: neither v4 nor v6, IP: %#v", r.IP))
		}
		af = new(bittorrent.AddressFamily)
		*af = req.AddressFamily

		var ctx context.Context
		var resp *bittorrent.ScrapeResponse
		ctx, resp, err = t.logic.HandleScrape(context.Background(), req)
		if err != nil {
			WriteError(w, txID, err)
			return
		}

		WriteScrape(w, txID, resp)

		go t.logic.AfterScrape(ctx, req, resp)

	default:
		err = errUnknownAction
		WriteError(w, txID, err)
	}

	return
}
<|endoftext|>