<commit_before>package waveguide

// TODO: Change the error handling so users see non-technical error messages with at most an id that would correspond to something in the logs.

import (
	"fmt"
	"html"
	"io/ioutil"
	"net/http"
	"net/url"
	"sort"
	"strconv"
	"strings"
	"time"

	"golang.org/x/net/context"
	"google.golang.org/appengine"
	"google.golang.org/appengine/datastore"
	"google.golang.org/appengine/log"
	"google.golang.org/appengine/taskqueue"
	"google.golang.org/appengine/urlfetch"
)

func init() {
	http.HandleFunc("/", handle(root))
	http.HandleFunc("/update_all", handle(updateAll))
	http.HandleFunc("/update_one", handle(updateOne))
	http.HandleFunc("/show", handle(show))
	http.HandleFunc("/clear", handle(clear))
	http.HandleFunc("/delete_one", handle(deleteOne))
	http.HandleFunc("/coords", handle(coords))
	http.HandleFunc("/clear_coords", handle(clearCoords))
	http.HandleFunc("/map", handle(doMap))
}

func handle(f func(ctx context.Context, w http.ResponseWriter, r *http.Request) (int, error)) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		ctx := appengine.NewContext(r)
		status, err := f(ctx, w, r)
		if err != nil {
			log.Errorf(ctx, "%v", err)
			w.WriteHeader(status)
			fmt.Fprintf(w, "%v", err)
		}
	}
}

// root shows the main page.
func root(ctx context.Context, w http.ResponseWriter, r *http.Request) (int, error) {
	if r.URL.Path != "/" {
		return http.StatusNotFound, fmt.Errorf("%s not found", r.URL.Path)
	}
	if err := r.ParseForm(); err != nil {
		return http.StatusInternalServerError, err
	}
	sn := r.FormValue("n")
	if sn == "" {
		sn = "250"
	}
	n, err := strconv.Atoi(sn)
	if err != nil {
		return http.StatusBadRequest, fmt.Errorf("Bad value %q for n parameter. Should be an integer.", sn)
	}
	log.Infof(ctx, "Limiting to top %d spots", n)
	q := datastore.NewQuery("Spot").Order("-Cond.Rating").Limit(n)
	var spots []Spot
	_, err = q.GetAll(ctx, &spots)
	if err != nil {
		return http.StatusInternalServerError, fmt.Errorf("root: fetching spots: %v", err)
	}
	sort.Sort(ByRating(spots))
	data := struct {
		Spots []Spot
	}{
		Spots: spots,
	}
	err = tmpl.ExecuteTemplate(w, "root", data)
	if err != nil {
		return http.StatusInternalServerError, fmt.Errorf("root: template: %v", err)
	}
	return http.StatusOK, nil
}

// updateAll fetches the site map from msw, saves Spots based on it, and adds tasks to update all the Spots.
func updateAll(ctx context.Context, w http.ResponseWriter, r *http.Request) (int, error) {
	// Get form arg n: number of spots to update, or -1 for all of them
	if err := r.ParseForm(); err != nil {
		return http.StatusInternalServerError, err
	}
	sn := r.FormValue("n")
	if sn == "" {
		sn = "-1"
	}
	n, err := strconv.Atoi(sn)

	// Download magicseaweed.com/site-map.php
	client := urlfetch.Client(ctx)
	body, err := get(client, "http://magicseaweed.com/site-map.php")
	if err != nil {
		return http.StatusInternalServerError, err
	}
	reportPaths := reportRx.FindAll(body, n)

	// For each report url:
	for _, rp := range reportPaths {
		// Add a task to the queue to download the report and store its stars and wave
		// height to the db.
		srp := string(rp)
		su := "http://magicseaweed.com" + srp
		t := taskqueue.NewPOSTTask("/update_one", map[string][]string{"url": {su}})
		if _, err := taskqueue.Add(ctx, t, ""); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return http.StatusInternalServerError, err
		}
	}
	data := struct{ Message string }{fmt.Sprintf("Processed %d paths.\n", len(reportPaths))}
	tmpl.ExecuteTemplate(w, "action_response", data)
	return http.StatusOK, nil
}

// updateOne updates the conditions for a single surfing spot, given the
// path to the surf report for it.
func updateOne(ctx context.Context, w http.ResponseWriter, r *http.Request) (int, error) {
	if err := r.ParseForm(); err != nil {
		return http.StatusInternalServerError, err
	}
	su := r.FormValue("url")
	u, err := url.Parse(su)
	if err != nil {
		return http.StatusBadRequest, err
	}
	path := u.Path
	qual, err := fetchSpotQuality(ctx, su)
	if err != nil {
		return http.StatusInternalServerError, err
	}
	SetSpotQuality(ctx, path, qual)
	data := struct{ Message string }{fmt.Sprintf("Updated %s", path)}
	tmpl.ExecuteTemplate(w, "action_response", data)
	return http.StatusOK, nil
}

// show shows the data for a single spot given its url, for debugging.
func show(ctx context.Context, w http.ResponseWriter, r *http.Request) (int, error) {
	if err := r.ParseForm(); err != nil {
		return http.StatusInternalServerError, err
	}
	su := r.FormValue("url")
	u, err := url.Parse(su)
	if err != nil {
		return http.StatusBadRequest, err
	}
	path := u.Path
	key := SpotKey(ctx, path)
	var s Spot
	if err := datastore.Get(ctx, key, &s); err != nil {
		return http.StatusInternalServerError, err
	}
	fmt.Fprintf(w, "%+v", s)
	return http.StatusOK, nil
}

// clear removes all the Spots from datastore. They can easily be brought back with updateAll.
func clear(ctx context.Context, w http.ResponseWriter, _ *http.Request) (int, error) {
	// TODO: This may time out. Rewrite it to queue up a bunch of tasks to delete individual spots or batches of
	// them.
	q := datastore.NewQuery("Spot")
	var spots []Spot
	_, err := q.GetAll(ctx, &spots)
	if err != nil {
		return http.StatusInternalServerError, err
	}
	keys := make([]*datastore.Key, len(spots))
	for i, s := range spots {
		keys[i] = SpotKey(ctx, s.MswPath)
		t := taskqueue.NewPOSTTask("/delete_one", map[string][]string{"path": {s.MswPath}})
		_, err = taskqueue.Add(ctx, t, "")
		if err != nil {
			return http.StatusInternalServerError, err
		}
	}
	data := struct{ Message string }{fmt.Sprintf("Queued %d spots for deletion.\n", len(spots))}
	tmpl.ExecuteTemplate(w, "action_response", data)
	return http.StatusOK, nil
}

// deleteOne deletes one Spot, given its MSW path as the "path" form value.
func deleteOne(ctx context.Context, w http.ResponseWriter, r *http.Request) (int, error) {
	if err := r.ParseForm(); err != nil {
		return http.StatusInternalServerError, err
	}
	p := r.FormValue("path")
	key := SpotKey(ctx, p)
	err := datastore.Delete(ctx, key)
	if err != nil {
		return http.StatusInternalServerError, err
	}
	data := struct{ Message string }{fmt.Sprintf("Deleted Spot %q\n", p)}
	tmpl.ExecuteTemplate(w, "action_response", data)
	return http.StatusOK, nil
}

// coords sets the coordinates for a Spot.
func coords(ctx context.Context, w http.ResponseWriter, r *http.Request) (int, error) {
	if r.Method != "POST" {
		return http.StatusBadRequest, fmt.Errorf("Got %q, want POST.", r.Method)
	}
	if err := r.ParseForm(); err != nil {
		return http.StatusInternalServerError, err
	}
	p := r.FormValue("path")
	coordsStr := r.FormValue("coordinates")
	key := SpotKey(ctx, p)
	var s Spot
	err := datastore.Get(ctx, key, &s)
	if err != nil {
		return http.StatusInternalServerError, err
	}
	m := coordsRx.FindStringSubmatch(coordsStr)
	if len(m) != 3 {
		return http.StatusBadRequest, fmt.Errorf("Got %q, want match to %q", coordsStr, coordsRx)
	}
	latStr := m[1]
	lngStr := m[2]
	lat, err := strconv.ParseFloat(latStr, 32)
	lng, err := strconv.ParseFloat(lngStr, 32)
	s.Coordinates.Lat = lat
	s.Coordinates.Lng = lng
	_, err = datastore.Put(ctx, key, &s)
	if err != nil {
		return http.StatusInternalServerError, err
	}
	data := struct{ Message string }{"ok"}
	// TODO: check err here and elsewhere.
	tmpl.ExecuteTemplate(w, "action_response", data)
	return http.StatusOK, nil
}

// clearCoords clears the coordinates for a Spot.
// It's not done by a DELETE request to /coords because it has to be accessible from links on a page.
func clearCoords(ctx context.Context, w http.ResponseWriter, r *http.Request) (int, error) {
	if err := r.ParseForm(); err != nil {
		return http.StatusInternalServerError, err
	}
	p := r.FormValue("path")
	key := SpotKey(ctx, p)
	var s Spot
	err := datastore.Get(ctx, key, &s)
	if err != nil {
		return http.StatusInternalServerError, err
	}
	s.Coordinates = appengine.GeoPoint{}
	_, err = datastore.Put(ctx, key, &s)
	if err != nil {
		return http.StatusInternalServerError, err
	}
	data := struct{ Message string }{"ok"}
	tmpl.ExecuteTemplate(w, "action_response", data)
	return http.StatusOK, nil
}

func doMap(ctx context.Context, w http.ResponseWriter, r *http.Request) (int, error) {
	if err := r.ParseForm(); err != nil {
		return http.StatusInternalServerError, err
	}
	sn := r.FormValue("n")
	if sn == "" {
		sn = "250"
	}
	n, err := strconv.Atoi(sn)
	if err != nil {
		return http.StatusBadRequest, fmt.Errorf("map: Bad value %q for n parameter. Should be an integer.", sn)
	}
	log.Infof(ctx, "Limiting to top %d spots", n)
	q := datastore.NewQuery("Spot").Order("-Cond.Rating").Limit(n)
	var spots []Spot
	_, err = q.GetAll(ctx, &spots)
	if err != nil {
		return http.StatusInternalServerError, fmt.Errorf("map: Fetching spots: %v", err)
	}
	sort.Sort(ByRating(spots))
	err = tmpl.ExecuteTemplate(w, "map", spots)
	if err != nil {
		return http.StatusInternalServerError, fmt.Errorf("map: %v", err)
	}
	return http.StatusOK, nil
}

func get(client *http.Client, url string) ([]byte, error) {
	res, err := client.Get(url)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, fmt.Errorf("Failed to read body of %s", url)
	}
	return body, nil
}

func surfReportPathToName(srp string) (string, error) {
	s := srpTailRx.ReplaceAllString(srp, "")
	// TODO: Use PathUnescape once GAE supports Go 1.8
	s, err := url.QueryUnescape(s)
	if err != nil {
		return "", fmt.Errorf("Failed to unescape path in surfReportPathToName(%q). %v", srp, err)
	}
	s = html.UnescapeString(s)
	s = s[1:] // Remove leading /
	s = strings.Replace(s, "-", " ", -1)
	return s, nil
}

// fetchSpotQuality scrapes the surf report at the given url and returns a Quality struct
// summarizing the conditions.
func fetchSpotQuality(ctx context.Context, url string) (*Quality, error) {
	log.Infof(ctx, "Fetching %s", url)
	client := urlfetch.Client(ctx)
	body, err := get(client, url)
	if err != nil {
		return nil, err
	}
	rating := countTopStars(body)
	hMatch := heightRx.FindSubmatch(body)
	if len(hMatch) != 2 {
		return nil, fmt.Errorf("Wave height regex failed.")
	}
	height := fmt.Sprintf("%s ft", hMatch[1])
	q := &Quality{
		Rating:     rating,
		WaveHeight: height,
		TimeUnix:   time.Now().Unix(),
	}
	return q, nil
}

// countTopStars returns the number of stars in the first rating section on the page.
func countTopStars(body []byte) int {
	starSection := starSectionRx.Find(body)
	foundStars := starRx.FindAll(starSection, -1)
	return len(foundStars)
}
<commit_msg>Check errors on all template calls<commit_after>package waveguide

// TODO: Change the error handling so users see non-technical error messages with at most an id that would correspond to something in the logs.

import (
	"fmt"
	"html"
	"io/ioutil"
	"net/http"
	"net/url"
	"sort"
	"strconv"
	"strings"
	"time"

	"golang.org/x/net/context"
	"google.golang.org/appengine"
	"google.golang.org/appengine/datastore"
	"google.golang.org/appengine/log"
	"google.golang.org/appengine/taskqueue"
	"google.golang.org/appengine/urlfetch"
)

func init() {
	http.HandleFunc("/", handle(root))
	http.HandleFunc("/update_all", handle(updateAll))
	http.HandleFunc("/update_one", handle(updateOne))
	http.HandleFunc("/show", handle(show))
	http.HandleFunc("/clear", handle(clear))
	http.HandleFunc("/delete_one", handle(deleteOne))
	http.HandleFunc("/coords", handle(coords))
	http.HandleFunc("/clear_coords", handle(clearCoords))
	http.HandleFunc("/map", handle(doMap))
}

func handle(f func(ctx context.Context, w http.ResponseWriter, r *http.Request) (int, error)) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		ctx := appengine.NewContext(r)
		status, err := f(ctx, w, r)
		if err != nil {
			log.Errorf(ctx, "%v", err)
			w.WriteHeader(status)
			fmt.Fprintf(w, "%v", err)
		}
	}
}

// root shows the main page.
func root(ctx context.Context, w http.ResponseWriter, r *http.Request) (int, error) {
	if r.URL.Path != "/" {
		return http.StatusNotFound, fmt.Errorf("%s not found", r.URL.Path)
	}
	if err := r.ParseForm(); err != nil {
		return http.StatusInternalServerError, err
	}
	sn := r.FormValue("n")
	if sn == "" {
		sn = "250"
	}
	n, err := strconv.Atoi(sn)
	if err != nil {
		return http.StatusBadRequest, fmt.Errorf("Bad value %q for n parameter. Should be an integer.", sn)
	}
	log.Infof(ctx, "Limiting to top %d spots", n)
	q := datastore.NewQuery("Spot").Order("-Cond.Rating").Limit(n)
	var spots []Spot
	_, err = q.GetAll(ctx, &spots)
	if err != nil {
		return http.StatusInternalServerError, fmt.Errorf("root: fetching spots: %v", err)
	}
	sort.Sort(ByRating(spots))
	data := struct {
		Spots []Spot
	}{
		Spots: spots,
	}
	err = tmpl.ExecuteTemplate(w, "root", data)
	if err != nil {
		return http.StatusInternalServerError, fmt.Errorf("root: template: %v", err)
	}
	return http.StatusOK, nil
}

// updateAll fetches the site map from msw, saves Spots based on it, and adds tasks to update all the Spots.
func updateAll(ctx context.Context, w http.ResponseWriter, r *http.Request) (int, error) {
	// Get form arg n: number of spots to update, or -1 for all of them
	if err := r.ParseForm(); err != nil {
		return http.StatusInternalServerError, err
	}
	sn := r.FormValue("n")
	if sn == "" {
		sn = "-1"
	}
	n, err := strconv.Atoi(sn)

	// Download magicseaweed.com/site-map.php
	client := urlfetch.Client(ctx)
	body, err := get(client, "http://magicseaweed.com/site-map.php")
	if err != nil {
		return http.StatusInternalServerError, err
	}
	reportPaths := reportRx.FindAll(body, n)

	// For each report url:
	for _, rp := range reportPaths {
		// Add a task to the queue to download the report and store its stars and wave
		// height to the db.
		srp := string(rp)
		su := "http://magicseaweed.com" + srp
		t := taskqueue.NewPOSTTask("/update_one", map[string][]string{"url": {su}})
		if _, err := taskqueue.Add(ctx, t, ""); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return http.StatusInternalServerError, err
		}
	}
	data := struct{ Message string }{fmt.Sprintf("Processed %d paths.\n", len(reportPaths))}
	err = tmpl.ExecuteTemplate(w, "action_response", data)
	if err != nil {
		return http.StatusInternalServerError, fmt.Errorf("update_all: %v", err)
	}
	return http.StatusOK, nil
}

// updateOne updates the conditions for a single surfing spot, given the
// path to the surf report for it.
func updateOne(ctx context.Context, w http.ResponseWriter, r *http.Request) (int, error) {
	if err := r.ParseForm(); err != nil {
		return http.StatusInternalServerError, err
	}
	su := r.FormValue("url")
	u, err := url.Parse(su)
	if err != nil {
		return http.StatusBadRequest, err
	}
	path := u.Path
	qual, err := fetchSpotQuality(ctx, su)
	if err != nil {
		return http.StatusInternalServerError, err
	}
	SetSpotQuality(ctx, path, qual)
	data := struct{ Message string }{fmt.Sprintf("Updated %s", path)}
	err = tmpl.ExecuteTemplate(w, "action_response", data)
	if err != nil {
		return http.StatusInternalServerError, fmt.Errorf("update_one: %v", err)
	}
	return http.StatusOK, nil
}

// show shows the data for a single spot given its url, for debugging.
func show(ctx context.Context, w http.ResponseWriter, r *http.Request) (int, error) {
	if err := r.ParseForm(); err != nil {
		return http.StatusInternalServerError, err
	}
	su := r.FormValue("url")
	u, err := url.Parse(su)
	if err != nil {
		return http.StatusBadRequest, err
	}
	path := u.Path
	key := SpotKey(ctx, path)
	var s Spot
	if err := datastore.Get(ctx, key, &s); err != nil {
		return http.StatusInternalServerError, err
	}
	fmt.Fprintf(w, "%+v", s)
	return http.StatusOK, nil
}

// clear removes all the Spots from datastore. They can easily be brought back with updateAll.
func clear(ctx context.Context, w http.ResponseWriter, _ *http.Request) (int, error) {
	// TODO: This may time out. Rewrite it to queue up a bunch of tasks to delete individual spots or batches of
	// them.
	q := datastore.NewQuery("Spot")
	var spots []Spot
	_, err := q.GetAll(ctx, &spots)
	if err != nil {
		return http.StatusInternalServerError, err
	}
	keys := make([]*datastore.Key, len(spots))
	for i, s := range spots {
		keys[i] = SpotKey(ctx, s.MswPath)
		t := taskqueue.NewPOSTTask("/delete_one", map[string][]string{"path": {s.MswPath}})
		_, err = taskqueue.Add(ctx, t, "")
		if err != nil {
			return http.StatusInternalServerError, err
		}
	}
	data := struct{ Message string }{fmt.Sprintf("Queued %d spots for deletion.\n", len(spots))}
	err = tmpl.ExecuteTemplate(w, "action_response", data)
	if err != nil {
		return http.StatusInternalServerError, fmt.Errorf("clear: %v", err)
	}
	return http.StatusOK, nil
}

// deleteOne deletes one Spot, given its MSW path as the "path" form value.
func deleteOne(ctx context.Context, w http.ResponseWriter, r *http.Request) (int, error) {
	if err := r.ParseForm(); err != nil {
		return http.StatusInternalServerError, err
	}
	p := r.FormValue("path")
	key := SpotKey(ctx, p)
	err := datastore.Delete(ctx, key)
	if err != nil {
		return http.StatusInternalServerError, err
	}
	data := struct{ Message string }{fmt.Sprintf("Deleted Spot %q\n", p)}
	err = tmpl.ExecuteTemplate(w, "action_response", data)
	if err != nil {
		return http.StatusInternalServerError, fmt.Errorf("delete_one: %v", err)
	}
	return http.StatusOK, nil
}

// coords sets the coordinates for a Spot.
func coords(ctx context.Context, w http.ResponseWriter, r *http.Request) (int, error) {
	if r.Method != "POST" {
		return http.StatusBadRequest, fmt.Errorf("Got %q, want POST.", r.Method)
	}
	if err := r.ParseForm(); err != nil {
		return http.StatusInternalServerError, err
	}
	p := r.FormValue("path")
	coordsStr := r.FormValue("coordinates")
	key := SpotKey(ctx, p)
	var s Spot
	err := datastore.Get(ctx, key, &s)
	if err != nil {
		return http.StatusInternalServerError, err
	}
	m := coordsRx.FindStringSubmatch(coordsStr)
	if len(m) != 3 {
		return http.StatusBadRequest, fmt.Errorf("Got %q, want match to %q", coordsStr, coordsRx)
	}
	latStr := m[1]
	lngStr := m[2]
	lat, err := strconv.ParseFloat(latStr, 32)
	lng, err := strconv.ParseFloat(lngStr, 32)
	s.Coordinates.Lat = lat
	s.Coordinates.Lng = lng
	_, err = datastore.Put(ctx, key, &s)
	if err != nil {
		return http.StatusInternalServerError, err
	}
	data := struct{ Message string }{"ok"}
	err = tmpl.ExecuteTemplate(w, "action_response", data)
	if err != nil {
		return http.StatusInternalServerError, fmt.Errorf("coords: %v", err)
	}
	return http.StatusOK, nil
}

// clearCoords clears the coordinates for a Spot.
// It's not done by a DELETE request to /coords because it has to be accessible from links on a page.
func clearCoords(ctx context.Context, w http.ResponseWriter, r *http.Request) (int, error) {
	if err := r.ParseForm(); err != nil {
		return http.StatusInternalServerError, err
	}
	p := r.FormValue("path")
	key := SpotKey(ctx, p)
	var s Spot
	err := datastore.Get(ctx, key, &s)
	if err != nil {
		return http.StatusInternalServerError, err
	}
	s.Coordinates = appengine.GeoPoint{}
	_, err = datastore.Put(ctx, key, &s)
	if err != nil {
		return http.StatusInternalServerError, err
	}
	data := struct{ Message string }{"ok"}
	err = tmpl.ExecuteTemplate(w, "action_response", data)
	if err != nil {
		return http.StatusInternalServerError, fmt.Errorf("clear_coords: %v", err)
	}
	return http.StatusOK, nil
}

func doMap(ctx context.Context, w http.ResponseWriter, r *http.Request) (int, error) {
	if err := r.ParseForm(); err != nil {
		return http.StatusInternalServerError, err
	}
	sn := r.FormValue("n")
	if sn == "" {
		sn = "250"
	}
	n, err := strconv.Atoi(sn)
	if err != nil {
		return http.StatusBadRequest, fmt.Errorf("map: Bad value %q for n parameter. Should be an integer.", sn)
	}
	log.Infof(ctx, "Limiting to top %d spots", n)
	q := datastore.NewQuery("Spot").Order("-Cond.Rating").Limit(n)
	var spots []Spot
	_, err = q.GetAll(ctx, &spots)
	if err != nil {
		return http.StatusInternalServerError, fmt.Errorf("map: Fetching spots: %v", err)
	}
	sort.Sort(ByRating(spots))
	err = tmpl.ExecuteTemplate(w, "map", spots)
	if err != nil {
		return http.StatusInternalServerError, fmt.Errorf("map: %v", err)
	}
	return http.StatusOK, nil
}

func get(client *http.Client, url string) ([]byte, error) {
	res, err := client.Get(url)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, fmt.Errorf("Failed to read body of %s", url)
	}
	return body, nil
}

func surfReportPathToName(srp string) (string, error) {
	s := srpTailRx.ReplaceAllString(srp, "")
	// TODO: Use PathUnescape once GAE supports Go 1.8
	s, err := url.QueryUnescape(s)
	if err != nil {
		return "", fmt.Errorf("Failed to unescape path in surfReportPathToName(%q). %v", srp, err)
	}
	s = html.UnescapeString(s)
	s = s[1:] // Remove leading /
	s = strings.Replace(s, "-", " ", -1)
	return s, nil
}

// fetchSpotQuality scrapes the surf report at the given url and returns a Quality struct
// summarizing the conditions.
func fetchSpotQuality(ctx context.Context, url string) (*Quality, error) {
	log.Infof(ctx, "Fetching %s", url)
	client := urlfetch.Client(ctx)
	body, err := get(client, url)
	if err != nil {
		return nil, err
	}
	rating := countTopStars(body)
	hMatch := heightRx.FindSubmatch(body)
	if len(hMatch) != 2 {
		return nil, fmt.Errorf("Wave height regex failed.")
	}
	height := fmt.Sprintf("%s ft", hMatch[1])
	q := &Quality{
		Rating:     rating,
		WaveHeight: height,
		TimeUnix:   time.Now().Unix(),
	}
	return q, nil
}

// countTopStars returns the number of stars in the first rating section on the page.
func countTopStars(body []byte) int {
	starSection := starSectionRx.Find(body)
	foundStars := starRx.FindAll(starSection, -1)
	return len(foundStars)
}
<|endoftext|>


<commit_before>// Package Wbzr provides tool to create a wooble in a given language and to package
// (wrap) some woobles.
package wbzr

import (
	"bytes"
	"io/ioutil"
	"text/template"

	"github.com/woobleio/wooblizer/wbzr/engine"
)

// ScriptLang are constants for implemented script languages.
type ScriptLang int

const (
	JSES5 ScriptLang = iota
)

type Wbzr struct {
	DomainsSec []string
	Scripts    []engine.Script

	lang     ScriptLang
	skeleton string
}

// New takes a script language which is used to inject and output a file.
func New(sl ScriptLang) *Wbzr {
	var skeleton string
	switch sl {
	case JSES5:
		skeleton = wbJses5
	default:
		panic("Language not supported")
	}

	return &Wbzr{
		nil,
		make([]engine.Script, 0),
		sl,
		skeleton,
	}
}

// Get returns an injected source.
func (wb *Wbzr) Get(name string) (engine.Script, error) {
	for _, sc := range wb.Scripts {
		if sc.GetName() == name {
			return sc, nil
		}
	}
	return nil, ErrUniqueName
}

// Inject injects a source code to be wooblized. It takes a name which must be
// unique. Src can be empty, it'll create a default object
func (wb *Wbzr) Inject(src string, name string) (engine.Script, error) {
	if _, err := wb.Get(name); err == nil {
		return nil, ErrUniqueName
	}
	var sc engine.Script
	var err error

	switch wb.lang {
	case JSES5:
		sc, err = engine.NewJSES5(src, name)
	}

	if err != nil {
		return nil, err
	}

	wb.Scripts = append(wb.Scripts, sc)
	return sc, nil
}

// InjectFile injects a source from a file.
func (wb *Wbzr) InjectFile(path string, name string) (engine.Script, error) {
	c, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}

	return wb.Inject(string(c[:]), name)
}

// Secure set some domains to protect the script and make it works only for specific domains
func (wb *Wbzr) Secure(domains ...string) {
	wb.DomainsSec = domains
}

func (wb *Wbzr) SecureAndWrap(domains ...string) (*bytes.Buffer, error) {
	wb.Secure(domains...)
	return wb.Wrap()
}

// Wrap packages some woobles (all the woobles injected in the Wbzr)
// and build a file which contains the wooble library.
func (wb *Wbzr) Wrap() (*bytes.Buffer, error) {
	for _, sc := range wb.Scripts {
		if _, err := sc.Build(); err != nil {
			return nil, err
		}
	}

	tmpl := template.Must(template.New("WbJSES5").Parse(wbJses5))

	var out bytes.Buffer
	if err := tmpl.Execute(&out, wb); err != nil {
		return nil, err
	}

	return &out, nil
}

// WooblyJSES5 is a Wooble creation template for JSES5
var WooblyJSES5 = `woobly = {
	doc: document,
  attribute: "a value (optionnal)",
  _init: function() {
    // Creation code at runtime
  },
  method: function(a, b) {
    // a method (optionnal)
  }
}`

var wbJses5 = `
function Wb(id) {
	{{if .DomainsSec}}
	var ah = [{{range $i, $o := .DomainsSec}}"{{$o}}"{{if not $i}},{{end}}{{end}}];
  var xx = ah.indexOf(window.location.hostname);
  if(ah.indexOf(window.location.hostname) == -1) {
  	console.log("Wooble error : domain restricted");
    return;
  }
	{{end}}

	if(window === this) {
  	return new Wb(id);
  }

  var cs = {
  	{{range $i, $o := .Scripts}}"{{$o.GetName}}":{{"{"}}{{$o.GetSource}}{{"}"}}{{if not $i}},{{end}}{{end}}
  }

  var c = cs[id];
  if(typeof c == 'undefined') {
  	console.log("Wooble error : creation", id, "not found");
    return undefined;
  }

  this.init = function (target) {
    if(document.querySelector(target) == null) {
    	console.log("Wooble error : Element", target, "not found in the document");
      return;
    }

    if("_buildDoc" in c) c._buildDoc(target);
    if("_buildStyle" in c) c._buildStyle();
    if("_init" in c) c._init();

		return c;
  }

  this.get = function() {
  	return c;
  }

  return this;
}
`
<commit_msg>Adds template indication<commit_after>// Package Wbzr provides tool to create a wooble in a given language and to package
// (wrap) some woobles.
package wbzr

import (
	"bytes"
	"io/ioutil"
	"text/template"

	"github.com/woobleio/wooblizer/wbzr/engine"
)

// ScriptLang are constants for implemented script languages.
type ScriptLang int

const (
	JSES5 ScriptLang = iota
)

type Wbzr struct {
	DomainsSec []string
	Scripts    []engine.Script

	lang     ScriptLang
	skeleton string
}

// New takes a script language which is used to inject and output a file.
func New(sl ScriptLang) *Wbzr {
	var skeleton string
	switch sl {
	case JSES5:
		skeleton = wbJses5
	default:
		panic("Language not supported")
	}

	return &Wbzr{
		nil,
		make([]engine.Script, 0),
		sl,
		skeleton,
	}
}

// Get returns an injected source.
func (wb *Wbzr) Get(name string) (engine.Script, error) {
	for _, sc := range wb.Scripts {
		if sc.GetName() == name {
			return sc, nil
		}
	}
	return nil, ErrUniqueName
}

// Inject injects a source code to be wooblized. It takes a name which must be
// unique. Src can be empty, it'll create a default object
func (wb *Wbzr) Inject(src string, name string) (engine.Script, error) {
	if _, err := wb.Get(name); err == nil {
		return nil, ErrUniqueName
	}
	var sc engine.Script
	var err error

	switch wb.lang {
	case JSES5:
		sc, err = engine.NewJSES5(src, name)
	}

	if err != nil {
		return nil, err
	}

	wb.Scripts = append(wb.Scripts, sc)
	return sc, nil
}

// InjectFile injects a source from a file.
func (wb *Wbzr) InjectFile(path string, name string) (engine.Script, error) {
	c, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}

	return wb.Inject(string(c[:]), name)
}

// Secure set some domains to protect the script and make it works only for specific domains
func (wb *Wbzr) Secure(domains ...string) {
	wb.DomainsSec = domains
}

func (wb *Wbzr) SecureAndWrap(domains ...string) (*bytes.Buffer, error) {
	wb.Secure(domains...)
	return wb.Wrap()
}

// Wrap packages some woobles (all the woobles injected in the Wbzr)
// and build a file which contains the wooble library.
func (wb *Wbzr) Wrap() (*bytes.Buffer, error) {
	for _, sc := range wb.Scripts {
		if _, err := sc.Build(); err != nil {
			return nil, err
		}
	}

	tmpl := template.Must(template.New("WbJSES5").Parse(wbJses5))

	var out bytes.Buffer
	if err := tmpl.Execute(&out, wb); err != nil {
		return nil, err
	}

	return &out, nil
}

// WooblyJSES5 is a Wooble creation template for JSES5
var WooblyJSES5 = `woobly = {
	// Use this attribute to use this creation document (such as this._doc.stuff), otherwise
	// you may use 'document' (document.stuff) to use creation's parent when a user uses your creation
	doc: document,
  attribute: "a value (optionnal)",
  _init: function() {
    // Creation code at runtime
  },
  method: function(a, b) {
    // a method (optionnal)
  }
}`

var wbJses5 = `
function Wb(id) {
	{{if .DomainsSec}}
	var ah = [{{range $i, $o := .DomainsSec}}"{{$o}}"{{if not $i}},{{end}}{{end}}];
  var xx = ah.indexOf(window.location.hostname);
  if(ah.indexOf(window.location.hostname) == -1) {
  	console.log("Wooble error : domain restricted");
    return;
  }
	{{end}}

	if(window === this) {
  	return new Wb(id);
  }

  var cs = {
  	{{range $i, $o := .Scripts}}"{{$o.GetName}}":{{"{"}}{{$o.GetSource}}{{"}"}}{{if not $i}},{{end}}{{end}}
  }

  var c = cs[id];
  if(typeof c == 'undefined') {
  	console.log("Wooble error : creation", id, "not found");
    return undefined;
  }

  this.init = function (target) {
    if(document.querySelector(target) == null) {
    	console.log("Wooble error : Element", target, "not found in the document");
      return;
    }

    if("_buildDoc" in c) c._buildDoc(target);
    if("_buildStyle" in c) c._buildStyle();
    if("_init" in c) c._init();

		return c;
  }

  this.get = function() {
  	return c;
  }

  return this;
}
`
<|endoftext|>


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

// Package prometheus provides bindings to the Prometheus HTTP API:
// http://prometheus.io/docs/querying/api/
package prometheus

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/prometheus/common/model"
	"golang.org/x/net/context"
	"golang.org/x/net/context/ctxhttp"
)

const (
	statusAPIError = 422
	apiPrefix      = "/api/v1"

	epQuery       = "/query"
	epQueryRange  = "/query_range"
	epLabelValues = "/label/:name/values"
	epSeries      = "/series"
)

type ErrorType string

const (
	// The different API error types.
	ErrBadData     ErrorType = "bad_data"
	ErrTimeout               = "timeout"
	ErrCanceled              = "canceled"
	ErrExec                  = "execution"
	ErrBadResponse           = "bad_response"
)

// Error is an error returned by the API.
type Error struct {
	Type ErrorType
	Msg  string
}

func (e *Error) Error() string {
	return fmt.Sprintf("%s: %s", e.Type, e.Msg)
}

// CancelableTransport is like net.Transport but provides
// per-request cancelation functionality.
type CancelableTransport interface {
	http.RoundTripper
	CancelRequest(req *http.Request)
}

var DefaultTransport CancelableTransport = &http.Transport{
	Proxy: http.ProxyFromEnvironment,
	Dial: (&net.Dialer{
		Timeout:   30 * time.Second,
		KeepAlive: 30 * time.Second,
	}).Dial,
	TLSHandshakeTimeout: 10 * time.Second,
}

// Config defines configuration parameters for a new client.
type Config struct {
	// The address of the Prometheus to connect to.
	Address string

	// Transport is used by the Client to drive HTTP requests. If not
	// provided, DefaultTransport will be used.
	Transport CancelableTransport
}

func (cfg *Config) transport() CancelableTransport {
	if cfg.Transport == nil {
		return DefaultTransport
	}
	return cfg.Transport
}

type Client interface {
	url(ep string, args map[string]string) *url.URL
	do(context.Context, *http.Request) (*http.Response, []byte, error)
}

// New returns a new Client.
func New(cfg Config) (Client, error) {
	u, err := url.Parse(cfg.Address)
	if err != nil {
		return nil, err
	}
	u.Path = strings.TrimRight(u.Path, "/") + apiPrefix

	return &httpClient{
		endpoint:  u,
		transport: cfg.transport(),
	}, nil
}

type httpClient struct {
	endpoint  *url.URL
	transport CancelableTransport
}

func (c *httpClient) url(ep string, args map[string]string) *url.URL {
	p := path.Join(c.endpoint.Path, ep)

	for arg, val := range args {
		arg = ":" + arg
		p = strings.Replace(p, arg, val, -1)
	}

	u := *c.endpoint
	u.Path = p

	return &u
}

func (c *httpClient) do(ctx context.Context, req *http.Request) (*http.Response, []byte, error) {
	resp, err := ctxhttp.Do(ctx, &http.Client{Transport: c.transport}, req)

	defer func() {
		if resp != nil {
			resp.Body.Close()
		}
	}()

	if err != nil {
		return nil, nil, err
	}

	var body []byte
	done := make(chan struct{})
	go func() {
		body, err = ioutil.ReadAll(resp.Body)
		close(done)
	}()

	select {
	case <-ctx.Done():
		err = resp.Body.Close()
		<-done
		if err == nil {
			err = ctx.Err()
		}
	case <-done:
	}

	return resp, body, err
}

// apiClient wraps a regular client and processes successful API responses.
// Successful also includes responses that errored at the API level.
type apiClient struct {
	Client
}

type apiResponse struct {
	Status    string          `json:"status"`
	Data      json.RawMessage `json:"data"`
	ErrorType ErrorType       `json:"errorType"`
	Error     string          `json:"error"`
}

func (c apiClient) do(ctx context.Context, req *http.Request) (*http.Response, []byte, error) {
	resp, body, err := c.Client.do(ctx, req)
	if err != nil {
		return resp, body, err
	}

	code := resp.StatusCode

	if code/100 != 2 && code != statusAPIError {
		return resp, body, &Error{
			Type: ErrBadResponse,
			Msg:  fmt.Sprintf("bad response code %d", resp.StatusCode),
		}
	}

	var result apiResponse

	if err = json.Unmarshal(body, &result); err != nil {
		return resp, body, &Error{
			Type: ErrBadResponse,
			Msg:  err.Error(),
		}
	}

	if (code == statusAPIError) != (result.Status == "error") {
		err = &Error{
			Type: ErrBadResponse,
			Msg:  "inconsistent body for response code",
		}
	}

	if code == statusAPIError && result.Status == "error" {
		err = &Error{
			Type: result.ErrorType,
			Msg:  result.Error,
		}
	}

	return resp, []byte(result.Data), err
}

// Range represents a sliced time range.
type Range struct {
	// The boundaries of the time range.
	Start, End time.Time
	// The maximum time between two slices within the boundaries.
	Step time.Duration
}

// queryResult contains result data for a query.
type queryResult struct {
	Type   model.ValueType `json:"resultType"`
	Result interface{}     `json:"result"`

	// The decoded value.
	v model.Value
}

func (qr *queryResult) UnmarshalJSON(b []byte) error {
	v := struct {
		Type   model.ValueType `json:"resultType"`
		Result json.RawMessage `json:"result"`
	}{}

	err := json.Unmarshal(b, &v)
	if err != nil {
		return err
	}

	switch v.Type {
	case model.ValScalar:
		var sv model.Scalar
		err = json.Unmarshal(v.Result, &sv)
		qr.v = &sv

	case model.ValVector:
		var vv model.Vector
		err = json.Unmarshal(v.Result, &vv)
		qr.v = vv

	case model.ValMatrix:
		var mv model.Matrix
		err = json.Unmarshal(v.Result, &mv)
		qr.v = mv

	default:
		err = fmt.Errorf("unexpected value type %q", v.Type)
	}
	return err
}

// QueryAPI provides bindings the Prometheus's query API.
type QueryAPI interface {
	// Query performs a query for the given time.
	Query(ctx context.Context, query string, ts time.Time) (model.Value, error)
	// Query performs a query for the given range.
	QueryRange(ctx context.Context, query string, r Range) (model.Value, error)
}

// NewQueryAPI returns a new QueryAPI for the client.
func NewQueryAPI(c Client) QueryAPI {
	return &httpQueryAPI{client: apiClient{c}}
}

type httpQueryAPI struct {
	client Client
}

func (h *httpQueryAPI) Query(ctx context.Context, query string, ts time.Time) (model.Value, error) {
	u := h.client.url(epQuery, nil)
	q := u.Query()

	q.Set("query", query)
	q.Set("time", ts.Format(time.RFC3339Nano))

	u.RawQuery = q.Encode()

	req, _ := http.NewRequest("GET", u.String(), nil)

	_, body, err := h.client.do(ctx, req)
	if err != nil {
		return nil, err
	}

	var qres queryResult
	err = json.Unmarshal(body, &qres)

	return model.Value(qres.v), err
}

func (h *httpQueryAPI) QueryRange(ctx context.Context, query string, r Range) (model.Value, error) {
	u := h.client.url(epQueryRange, nil)
	q := u.Query()

	var (
		start = r.Start.Format(time.RFC3339Nano)
		end   = r.End.Format(time.RFC3339Nano)
		step  = strconv.FormatFloat(r.Step.Seconds(), 'f', 3, 64)
	)

	q.Set("query", query)
	q.Set("start", start)
	q.Set("end", end)
	q.Set("step", step)

	u.RawQuery = q.Encode()

	req, _ := http.NewRequest("GET", u.String(), nil)

	_, body, err := h.client.do(ctx, req)
	if err != nil {
		return nil, err
	}

	var qres queryResult
	err = json.Unmarshal(body, &qres)

	return model.Value(qres.v), err
}
<commit_msg>api: document goroutine safeness<commit_after>// Copyright 2015 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package prometheus provides bindings to the Prometheus HTTP API:
// http://prometheus.io/docs/querying/api/
package prometheus

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/prometheus/common/model"
	"golang.org/x/net/context"
	"golang.org/x/net/context/ctxhttp"
)

const (
	statusAPIError = 422
	apiPrefix      = "/api/v1"

	epQuery       = "/query"
	epQueryRange  = "/query_range"
	epLabelValues = "/label/:name/values"
	epSeries      = "/series"
)

type ErrorType string

const (
	// The different API error types.
	ErrBadData     ErrorType = "bad_data"
	ErrTimeout               = "timeout"
	ErrCanceled              = "canceled"
	ErrExec                  = "execution"
	ErrBadResponse           = "bad_response"
)

// Error is an error returned by the API.
type Error struct {
	Type ErrorType
	Msg  string
}

func (e *Error) Error() string {
	return fmt.Sprintf("%s: %s", e.Type, e.Msg)
}

// CancelableTransport is like net.Transport but provides
// per-request cancelation functionality.
type CancelableTransport interface {
	http.RoundTripper
	CancelRequest(req *http.Request)
}

var DefaultTransport CancelableTransport = &http.Transport{
	Proxy: http.ProxyFromEnvironment,
	Dial: (&net.Dialer{
		Timeout:   30 * time.Second,
		KeepAlive: 30 * time.Second,
	}).Dial,
	TLSHandshakeTimeout: 10 * time.Second,
}

// Config defines configuration parameters for a new client.
type Config struct {
	// The address of the Prometheus to connect to.
	Address string

	// Transport is used by the Client to drive HTTP requests. If not
	// provided, DefaultTransport will be used.
	Transport CancelableTransport
}

func (cfg *Config) transport() CancelableTransport {
	if cfg.Transport == nil {
		return DefaultTransport
	}
	return cfg.Transport
}

type Client interface {
	url(ep string, args map[string]string) *url.URL
	do(context.Context, *http.Request) (*http.Response, []byte, error)
}

// New returns a new Client.
//
// It is safe to use the returned Client from multiple goroutines.
func New(cfg Config) (Client, error) {
	u, err := url.Parse(cfg.Address)
	if err != nil {
		return nil, err
	}
	u.Path = strings.TrimRight(u.Path, "/") + apiPrefix

	return &httpClient{
		endpoint:  u,
		transport: cfg.transport(),
	}, nil
}

type httpClient struct {
	endpoint  *url.URL
	transport CancelableTransport
}

func (c *httpClient) url(ep string, args map[string]string) *url.URL {
	p := path.Join(c.endpoint.Path, ep)

	for arg, val := range args {
		arg = ":" + arg
		p = strings.Replace(p, arg, val, -1)
	}

	u := *c.endpoint
	u.Path = p

	return &u
}

func (c *httpClient) do(ctx context.Context, req *http.Request) (*http.Response, []byte, error) {
	resp, err := ctxhttp.Do(ctx, &http.Client{Transport: c.transport}, req)

	defer func() {
		if resp != nil {
			resp.Body.Close()
		}
	}()

	if err != nil {
		return nil, nil, err
	}

	var body []byte
	done := make(chan struct{})
	go func() {
		body, err = ioutil.ReadAll(resp.Body)
		close(done)
	}()

	select {
	case <-ctx.Done():
		err = resp.Body.Close()
		<-done
		if err == nil {
			err = ctx.Err()
		}
	case <-done:
	}

	return resp, body, err
}

// apiClient wraps a regular client and processes successful API responses.
// Successful also includes responses that errored at the API level.
type apiClient struct {
	Client
}

type apiResponse struct {
	Status    string          `json:"status"`
	Data      json.RawMessage `json:"data"`
	ErrorType ErrorType       `json:"errorType"`
	Error     string          `json:"error"`
}

func (c apiClient) do(ctx context.Context, req *http.Request) (*http.Response, []byte, error) {
	resp, body, err := c.Client.do(ctx, req)
	if err != nil {
		return resp, body, err
	}

	code := resp.StatusCode

	if code/100 != 2 && code != statusAPIError {
		return resp, body, &Error{
			Type: ErrBadResponse,
			Msg:  fmt.Sprintf("bad response code %d", resp.StatusCode),
		}
	}

	var result apiResponse

	if err = json.Unmarshal(body, &result); err != nil {
		return resp, body, &Error{
			Type: ErrBadResponse,
			Msg:  err.Error(),
		}
	}

	if (code == statusAPIError) != (result.Status == "error") {
		err = &Error{
			Type: ErrBadResponse,
			Msg:  "inconsistent body for response code",
		}
	}

	if code == statusAPIError && result.Status == "error" {
		err = &Error{
			Type: result.ErrorType,
			Msg:  result.Error,
		}
	}

	return resp, []byte(result.Data), err
}

// Range represents a sliced time range.
type Range struct {
	// The boundaries of the time range.
	Start, End time.Time
	// The maximum time between two slices within the boundaries.
	Step time.Duration
}

// queryResult contains result data for a query.
type queryResult struct {
	Type   model.ValueType `json:"resultType"`
	Result interface{}     `json:"result"`

	// The decoded value.
	v model.Value
}

func (qr *queryResult) UnmarshalJSON(b []byte) error {
	v := struct {
		Type   model.ValueType `json:"resultType"`
		Result json.RawMessage `json:"result"`
	}{}

	err := json.Unmarshal(b, &v)
	if err != nil {
		return err
	}

	switch v.Type {
	case model.ValScalar:
		var sv model.Scalar
		err = json.Unmarshal(v.Result, &sv)
		qr.v = &sv

	case model.ValVector:
		var vv model.Vector
		err = json.Unmarshal(v.Result, &vv)
		qr.v = vv

	case model.ValMatrix:
		var mv model.Matrix
		err = json.Unmarshal(v.Result, &mv)
		qr.v = mv

	default:
		err = fmt.Errorf("unexpected value type %q", v.Type)
	}
	return err
}

// QueryAPI provides bindings the Prometheus's query API.
type QueryAPI interface {
	// Query performs a query for the given time.
	Query(ctx context.Context, query string, ts time.Time) (model.Value, error)
	// Query performs a query for the given range.
	QueryRange(ctx context.Context, query string, r Range) (model.Value, error)
}

// NewQueryAPI returns a new QueryAPI for the client.
//
// It is safe to use the returned QueryAPI from multiple goroutines.
func NewQueryAPI(c Client) QueryAPI {
	return &httpQueryAPI{client: apiClient{c}}
}

type httpQueryAPI struct {
	client Client
}

func (h *httpQueryAPI) Query(ctx context.Context, query string, ts time.Time) (model.Value, error) {
	u := h.client.url(epQuery, nil)
	q := u.Query()

	q.Set("query", query)
	q.Set("time", ts.Format(time.RFC3339Nano))

	u.RawQuery = q.Encode()

	req, _ := http.NewRequest("GET", u.String(), nil)

	_, body, err := h.client.do(ctx, req)
	if err != nil {
		return nil, err
	}

	var qres queryResult
	err = json.Unmarshal(body, &qres)

	return model.Value(qres.v), err
}

func (h *httpQueryAPI) QueryRange(ctx context.Context, query string, r Range) (model.Value, error) {
	u := h.client.url(epQueryRange, nil)
	q := u.Query()

	var (
		start = r.Start.Format(time.RFC3339Nano)
		end   = r.End.Format(time.RFC3339Nano)
		step  = strconv.FormatFloat(r.Step.Seconds(), 'f', 3, 64)
	)

	q.Set("query", query)
	q.Set("start", start)
	q.Set("end", end)
	q.Set("step", step)

	u.RawQuery = q.Encode()

	req, _ := http.NewRequest("GET", u.String(), nil)

	_, body, err := h.client.do(ctx, req)
	if err != nil {
		return nil, err
	}

	var qres queryResult
	err = json.Unmarshal(body, &qres)

	return model.Value(qres.v), err
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"time"

	"github.com/mkideal/cli"
)

var sockFile = filepath.Join(os.Getenv("HOME"), ".rpc.sock")

func main() {
	if err := cli.Root(root,
		cli.Tree(help),
		cli.Tree(daemon),
		cli.Tree(api,
			cli.Tree(ping),
		),
	).Run(os.Args[1:]); err != nil {
		fmt.Println(err)
	}
}

//------
// root
//------
var root = &cli.Command{
	Fn: func(ctx *cli.Context) error {
		ctx.WriteUsage()
		return nil
	},
}

//------
// help
//------
var help = &cli.Command{
	Name:        "help",
	Desc:        "display help",
	CanSubRoute: true,
	HTTPRouters: []string{"/v1/help"},
	HTTPMethods: []string{"GET"},

	Fn: cli.HelpCommandFn,
}

//--------
// daemon
//--------
type daemonT struct {
	cli.Helper
	Port uint16 `cli:"p,port" usage:"http port" dft:"8080"`
}

func (t *daemonT) Validate(ctx *cli.Context) error {
	if t.Port == 0 {
		return fmt.Errorf("please don't use 0 as http port")
	}
	return nil
}

var daemon = &cli.Command{
	Name: "daemon",
	Desc: "startup app as daemon",
	Argv: func() interface{} { return new(daemonT) },
	Fn: func(ctx *cli.Context) error {
		argv := ctx.Argv().(*daemonT)
		if argv.Help {
			ctx.WriteUsage()
			return nil
		}

		cli.EnableDebug()

		addr := fmt.Sprintf(":%d", argv.Port)
		r := ctx.Command().Root()
		if err := r.RegisterHTTP(ctx); err != nil {
			return err
		}

		listeners := make([]net.Listener, 0, 2)

		// http listener
		if httpListener, err := net.Listen("tcp", addr); err != nil {
			return err
		} else if ln, ok := httpListener.(*net.TCPListener); ok {
			listeners = append(listeners, tcpKeepAliveListener{ln})
		}

		// unix listener
		if err := os.Remove(sockFile); err != nil {
			return err
		}
		if unixListener, err := net.Listen("unix", sockFile); err != nil {
			return err
		} else {
			listeners = append(listeners, unixListener)
		}

		ctx.String("listeners size: %d\n", len(listeners))

		return r.Serve(listeners...)
	},
}

// http client use unix sock
var httpc = &http.Client{
	Transport: &http.Transport{
		Dial: func(_, _ string) (net.Conn, error) {
			return net.Dial("unix", sockFile)
		},
	},
}

type tcpKeepAliveListener struct {
	*net.TCPListener
}

func (ln tcpKeepAliveListener) Accept() (c net.Conn, err error) {
	tc, err := ln.AcceptTCP()
	if err != nil {
		return
	}
	tc.SetKeepAlive(true)
	tc.SetKeepAlivePeriod(1 * time.Minute)
	return tc, nil
}

//-----
// api
//-----
var api = &cli.Command{
	Name: "api",
	Desc: "display all api",
	Fn: func(ctx *cli.Context) error {
		cmd := ctx.Command().Root()
		if cmd.IsClient() {
			return cmd.RPC(httpc, ctx)
		}
		ctx.String("Commands:\n")
		ctx.String("    ping\n")
		return nil
	},
}

//------
// ping
//------
var ping = &cli.Command{
	Name: "ping",
	Desc: "ping server",
	Fn: func(ctx *cli.Context) error {
		cmd := ctx.Command().Root()
		if cmd.IsClient() {
			for {
				if err := cmd.RPC(httpc, ctx); err != nil {
					return err
				}
				time.Sleep(time.Millisecond * 1000)
			}
		}
		ctx.String(ctx.Color().Bold(time.Now().Format(time.RFC3339)) + " pong\n")
		return nil
	},
}
<commit_msg>remove color print for example rpc<commit_after>package main

import (
	"fmt"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"time"

	"github.com/mkideal/cli"
)

var sockFile = filepath.Join(os.Getenv("HOME"), ".rpc.sock")

func main() {
	if err := cli.Root(root,
		cli.Tree(help),
		cli.Tree(daemon),
		cli.Tree(api,
			cli.Tree(ping),
		),
	).Run(os.Args[1:]); err != nil {
		fmt.Println(err)
	}
}

//------
// root
//------
var root = &cli.Command{
	Fn: func(ctx *cli.Context) error {
		ctx.WriteUsage()
		return nil
	},
}

//------
// help
//------
var help = &cli.Command{
	Name:        "help",
	Desc:        "display help",
	CanSubRoute: true,
	HTTPRouters: []string{"/v1/help"},
	HTTPMethods: []string{"GET"},

	Fn: cli.HelpCommandFn,
}

//--------
// daemon
//--------
type daemonT struct {
	cli.Helper
	Port uint16 `cli:"p,port" usage:"http port" dft:"8080"`
}

func (t *daemonT) Validate(ctx *cli.Context) error {
	if t.Port == 0 {
		return fmt.Errorf("please don't use 0 as http port")
	}
	return nil
}

var daemon = &cli.Command{
	Name: "daemon",
	Desc: "startup app as daemon",
	Argv: func() interface{} { return new(daemonT) },
	Fn: func(ctx *cli.Context) error {
		argv := ctx.Argv().(*daemonT)
		if argv.Help {
			ctx.WriteUsage()
			return nil
		}

		cli.EnableDebug()

		addr := fmt.Sprintf(":%d", argv.Port)
		r := ctx.Command().Root()
		if err := r.RegisterHTTP(ctx); err != nil {
			return err
		}

		listeners := make([]net.Listener, 0, 2)

		// http listener
		if httpListener, err := net.Listen("tcp", addr); err != nil {
			return err
		} else if ln, ok := httpListener.(*net.TCPListener); ok {
			listeners = append(listeners, tcpKeepAliveListener{ln})
		}

		// unix listener
		if err := os.Remove(sockFile); err != nil {
			return err
		}
		if unixListener, err := net.Listen("unix", sockFile); err != nil {
			return err
		} else {
			listeners = append(listeners, unixListener)
		}

		ctx.String("listeners size: %d\n", len(listeners))

		return r.Serve(listeners...)
	},
}

// http client use unix sock
var httpc = &http.Client{
	Transport: &http.Transport{
		Dial: func(_, _ string) (net.Conn, error) {
			return net.Dial("unix", sockFile)
		},
	},
}

type tcpKeepAliveListener struct {
	*net.TCPListener
}

func (ln tcpKeepAliveListener) Accept() (c net.Conn, err error) {
	tc, err := ln.AcceptTCP()
	if err != nil {
		return
	}
	tc.SetKeepAlive(true)
	tc.SetKeepAlivePeriod(1 * time.Minute)
	return tc, nil
}

//-----
// api
//-----
var api = &cli.Command{
	Name: "api",
	Desc: "display all api",
	Fn: func(ctx *cli.Context) error {
		cmd := ctx.Command().Root()
		if cmd.IsClient() {
			return cmd.RPC(httpc, ctx)
		}
		ctx.String("Commands:\n")
		ctx.String("    ping\n")
		return nil
	},
}

//------
// ping
//------
var ping = &cli.Command{
	Name: "ping",
	Desc: "ping server",
	Fn: func(ctx *cli.Context) error {
		cmd := ctx.Command().Root()
		if cmd.IsClient() {
			for {
				if err := cmd.RPC(httpc, ctx); err != nil {
					return err
				}
				time.Sleep(time.Millisecond * 1000)
			}
		}
		ctx.String(time.Now().Format(time.RFC3339) + " pong\n")
		return nil
	},
}
<|endoftext|>


<commit_before>package qemu

import (
	"context"
	"fmt"
	"log"

	"github.com/hashicorp/packer/common/net"
	"github.com/hashicorp/packer/helper/multistep"
	"github.com/hashicorp/packer/packer"
)

// This step adds a NAT port forwarding definition so that SSH is available
// on the guest machine.
//
// Uses:
//
// Produces:
type stepForwardSSH struct {
	l *net.Listener
}

func (s *stepForwardSSH) Run(ctx context.Context, state multistep.StateBag) multistep.StepAction {
	config := state.Get("config").(*Config)
	ui := state.Get("ui").(packer.Ui)

	log.Printf("Looking for available communicator (SSH, WinRM, etc) port between %d and %d", config.SSHHostPortMin, config.SSHHostPortMax)
	var err error
	s.l, err = net.ListenRangeConfig{
		Addr:    config.VNCBindAddress,
		Min:     config.VNCPortMin,
		Max:     config.VNCPortMax,
		Network: "tcp",
	}.Listen(ctx)
	if err != nil {
		err := fmt.Errorf("Error finding port: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}
	s.l.Listener.Close() // free port, but don't unlock lock file
	sshHostPort := s.l.Port
	ui.Say(fmt.Sprintf("Found port for communicator (SSH, WinRM, etc): %d.", sshHostPort))

	// Save the port we're using so that future steps can use it
	state.Put("sshHostPort", sshHostPort)

	return multistep.ActionContinue
}

func (s *stepForwardSSH) Cleanup(state multistep.StateBag) {
	if s.l != nil {
		err := s.l.Close()
		if err != nil {
			log.Printf("failed to unlock port lockfile: %v", err)
		}
	}
}
<commit_msg>fix copypasta mistake switching ssh port mix/max for vnc port min/max<commit_after>package qemu

import (
	"context"
	"fmt"
	"log"

	"github.com/hashicorp/packer/common/net"
	"github.com/hashicorp/packer/helper/multistep"
	"github.com/hashicorp/packer/packer"
)

// This step adds a NAT port forwarding definition so that SSH is available
// on the guest machine.
//
// Uses:
//
// Produces:
type stepForwardSSH struct {
	l *net.Listener
}

func (s *stepForwardSSH) Run(ctx context.Context, state multistep.StateBag) multistep.StepAction {
	config := state.Get("config").(*Config)
	ui := state.Get("ui").(packer.Ui)

	log.Printf("Looking for available communicator (SSH, WinRM, etc) port between %d and %d", config.SSHHostPortMin, config.SSHHostPortMax)
	var err error
	s.l, err = net.ListenRangeConfig{
		Addr:    config.VNCBindAddress,
		Min:     config.SSHHostPortMin,
		Max:     config.SSHHostPortMax,
		Network: "tcp",
	}.Listen(ctx)
	if err != nil {
		err := fmt.Errorf("Error finding port: %s", err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}
	s.l.Listener.Close() // free port, but don't unlock lock file
	sshHostPort := s.l.Port
	ui.Say(fmt.Sprintf("Found port for communicator (SSH, WinRM, etc): %d.", sshHostPort))

	// Save the port we're using so that future steps can use it
	state.Put("sshHostPort", sshHostPort)

	return multistep.ActionContinue
}

func (s *stepForwardSSH) Cleanup(state multistep.StateBag) {
	if s.l != nil {
		err := s.l.Close()
		if err != nil {
			log.Printf("failed to unlock port lockfile: %v", err)
		}
	}
}
<|endoftext|>


<commit_before>package internet

import (
	"net"
	"syscall"
	"golang.org/x/sys/unix"
)

const (
	// For incoming connections.
	TCP_FASTOPEN = 23
	// For out-going connections.
	TCP_FASTOPEN_CONNECT = 30
)

func bindAddr(fd uintptr, ip []byte, port uint32) error {
	if err := syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, syscall.SO_REUSEADDR, 1); err != nil {
		return newError("failed to set resuse_addr").Base(err).AtWarning()
	}

	if err := syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, unix.SO_REUSEPORT, 1); err != nil {
		return newError("failed to set resuse_port").Base(err).AtWarning()
	}

	var sockaddr syscall.Sockaddr

	switch len(ip) {
	case net.IPv4len:
		a4 := &syscall.SockaddrInet4{
			Port: int(port),
		}
		copy(a4.Addr[:], ip)
		sockaddr = a4
	case net.IPv6len:
		a6 := &syscall.SockaddrInet6{
			Port: int(port),
		}
		copy(a6.Addr[:], ip)
		sockaddr = a6
	default:
		return newError("unexpected length of ip")
	}

	return syscall.Bind(int(fd), sockaddr)
}

func applyOutboundSocketOptions(network string, address string, fd uintptr, config *SocketConfig) error {
	if config.Mark != 0 {
		if err := syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, syscall.SO_MARK, int(config.Mark)); err != nil {
			return newError("failed to set SO_MARK").Base(err)
		}
	}

	if isTCPSocket(network) {
		switch config.Tfo {
		case SocketConfig_Enable:
			if err := syscall.SetsockoptInt(int(fd), syscall.SOL_TCP, TCP_FASTOPEN_CONNECT, 1); err != nil {
				return newError("failed to set TCP_FASTOPEN_CONNECT=1").Base(err)
			}
		case SocketConfig_Disable:
			if err := syscall.SetsockoptInt(int(fd), syscall.SOL_TCP, TCP_FASTOPEN_CONNECT, 0); err != nil {
				return newError("failed to set TCP_FASTOPEN_CONNECT=0").Base(err)
			}
		}
	}

	if config.Tproxy.IsEnabled() {
		if err := syscall.SetsockoptInt(int(fd), syscall.SOL_IP, syscall.IP_TRANSPARENT, 1); err != nil {
			return newError("failed to set IP_TRANSPARENT").Base(err)
		}
	}

	return nil
}

func applyInboundSocketOptions(network string, fd uintptr, config *SocketConfig) error {
	if config.Mark != 0 {
		if err := syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, syscall.SO_MARK, int(config.Mark)); err != nil {
			return newError("failed to set SO_MARK").Base(err)
		}
	}
	if isTCPSocket(network) {
		switch config.Tfo {
		case SocketConfig_Enable:
			if err := syscall.SetsockoptInt(int(fd), syscall.SOL_TCP, TCP_FASTOPEN, 1); err != nil {
				return newError("failed to set TCP_FASTOPEN=1").Base(err)
			}
		case SocketConfig_Disable:
			if err := syscall.SetsockoptInt(int(fd), syscall.SOL_TCP, TCP_FASTOPEN, 0); err != nil {
				return newError("failed to set TCP_FASTOPEN=0").Base(err)
			}
		}
	}

	if config.Tproxy.IsEnabled() {
		if err := syscall.SetsockoptInt(int(fd), syscall.SOL_IP, syscall.IP_TRANSPARENT, 1); err != nil {
			return newError("failed to set IP_TRANSPARENT").Base(err)
		}
	}

	if config.ReceiveOriginalDestAddress && isUDPSocket(network) {
		if err := syscall.SetsockoptInt(int(fd), syscall.SOL_IP, syscall.IP_RECVORIGDSTADDR, 1); err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>Add SO_REUSEPORT to inbound<commit_after>package internet

import (
	"net"
	"syscall"
	"golang.org/x/sys/unix"
)

const (
	// For incoming connections.
	TCP_FASTOPEN = 23
	// For out-going connections.
	TCP_FASTOPEN_CONNECT = 30
)

func bindAddr(fd uintptr, ip []byte, port uint32) error {
	if err := syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, syscall.SO_REUSEADDR, 1); err != nil {
		return newError("failed to set resuse_addr").Base(err).AtWarning()
	}

	if err := syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, unix.SO_REUSEPORT, 1); err != nil {
		return newError("failed to set resuse_port").Base(err).AtWarning()
	}

	var sockaddr syscall.Sockaddr

	switch len(ip) {
	case net.IPv4len:
		a4 := &syscall.SockaddrInet4{
			Port: int(port),
		}
		copy(a4.Addr[:], ip)
		sockaddr = a4
	case net.IPv6len:
		a6 := &syscall.SockaddrInet6{
			Port: int(port),
		}
		copy(a6.Addr[:], ip)
		sockaddr = a6
	default:
		return newError("unexpected length of ip")
	}

	return syscall.Bind(int(fd), sockaddr)
}

func applyOutboundSocketOptions(network string, address string, fd uintptr, config *SocketConfig) error {
	if config.Mark != 0 {
		if err := syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, syscall.SO_MARK, int(config.Mark)); err != nil {
			return newError("failed to set SO_MARK").Base(err)
		}
	}

	if isTCPSocket(network) {
		switch config.Tfo {
		case SocketConfig_Enable:
			if err := syscall.SetsockoptInt(int(fd), syscall.SOL_TCP, TCP_FASTOPEN_CONNECT, 1); err != nil {
				return newError("failed to set TCP_FASTOPEN_CONNECT=1").Base(err)
			}
		case SocketConfig_Disable:
			if err := syscall.SetsockoptInt(int(fd), syscall.SOL_TCP, TCP_FASTOPEN_CONNECT, 0); err != nil {
				return newError("failed to set TCP_FASTOPEN_CONNECT=0").Base(err)
			}
		}
	}

	if config.Tproxy.IsEnabled() {
		if err := syscall.SetsockoptInt(int(fd), syscall.SOL_IP, syscall.IP_TRANSPARENT, 1); err != nil {
			return newError("failed to set IP_TRANSPARENT").Base(err)
		}
	}

	return nil
}

func applyInboundSocketOptions(network string, fd uintptr, config *SocketConfig) error {
	if config.Mark != 0 {
		if err := syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, syscall.SO_MARK, int(config.Mark)); err != nil {
			return newError("failed to set SO_MARK").Base(err)
		}
	}
	if isTCPSocket(network) {
		switch config.Tfo {
		case SocketConfig_Enable:
			if err := syscall.SetsockoptInt(int(fd), syscall.SOL_TCP, TCP_FASTOPEN, 1); err != nil {
				return newError("failed to set TCP_FASTOPEN=1").Base(err)
			}
		case SocketConfig_Disable:
			if err := syscall.SetsockoptInt(int(fd), syscall.SOL_TCP, TCP_FASTOPEN, 0); err != nil {
				return newError("failed to set TCP_FASTOPEN=0").Base(err)
			}
		}
	}

	if config.Tproxy.IsEnabled() {
		if err := syscall.SetsockoptInt(int(fd), syscall.SOL_IP, syscall.IP_TRANSPARENT, 1); err != nil {
			return newError("failed to set IP_TRANSPARENT").Base(err)
		}
	}

	if config.ReceiveOriginalDestAddress && isUDPSocket(network) {
		if err := syscall.SetsockoptInt(int(fd), syscall.SOL_IP, syscall.IP_RECVORIGDSTADDR, 1); err != nil {
			return err
		}
	}

	if err := syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, unix.SO_REUSEPORT, 1); err != nil {
		return newError("failed to set SO_REUSEPORT").Base(err).AtWarning()
	}

	return nil
}
<|endoftext|>


<commit_before>// +build !confonly

package websocket

import (
	"context"
	"crypto/tls"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"v2ray.com/core/common"
	"v2ray.com/core/common/net"
	http_proto "v2ray.com/core/common/protocol/http"
	"v2ray.com/core/common/session"
	"v2ray.com/core/transport/internet"
	v2tls "v2ray.com/core/transport/internet/tls"
)

type requestHandler struct {
	path string
	ln   *Listener
}

var upgrader = &websocket.Upgrader{
	ReadBufferSize:   4 * 1024,
	WriteBufferSize:  4 * 1024,
	HandshakeTimeout: time.Second * 4,
}

func (h *requestHandler) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	if request.URL.Path != h.path {
		writer.WriteHeader(http.StatusNotFound)
		return
	}
	conn, err := upgrader.Upgrade(writer, request, nil)
	if err != nil {
		newError("failed to convert to WebSocket connection").Base(err).WriteToLog()
		return
	}

	forwardedAddrs := http_proto.ParseXForwardedFor(request.Header)
	remoteAddr := conn.RemoteAddr()
	if len(forwardedAddrs) > 0 && forwardedAddrs[0].Family().IsIP() {
		remoteAddr.(*net.TCPAddr).IP = forwardedAddrs[0].IP()
	}

	h.ln.addConn(newConnection(conn, remoteAddr))
}

type Listener struct {
	sync.Mutex
	server   http.Server
	listener net.Listener
	config   *Config
	addConn  internet.ConnHandler
}

func ListenWS(ctx context.Context, address net.Address, port net.Port, streamSettings *internet.MemoryStreamConfig, addConn internet.ConnHandler) (internet.Listener, error) {
	wsSettings := streamSettings.ProtocolSettings.(*Config)

	var tlsConfig *tls.Config
	if config := v2tls.ConfigFromStreamSettings(streamSettings); config != nil {
		tlsConfig = config.GetTLSConfig()
	}

	listener, err := listenTCP(ctx, address, port, tlsConfig, streamSettings.SocketSettings)
	if err != nil {
		return nil, err
	}

	l := &Listener{
		config:   wsSettings,
		addConn:  addConn,
		listener: listener,
	}

	l.server = http.Server{
		Handler: &requestHandler{
			path: wsSettings.GetNormalizedPath(),
			ln:   l,
		},
		ReadHeaderTimeout: time.Second * 4,
		MaxHeaderBytes:    2048,
	}

	go func() {
		if err := l.server.Serve(l.listener); err != nil {
			newError("failed to serve http for WebSocket").Base(err).AtWarning().WriteToLog(session.ExportIDToError(ctx))
		}
	}()

	return l, err
}

func listenTCP(ctx context.Context, address net.Address, port net.Port, tlsConfig *tls.Config, sockopt *internet.SocketConfig) (net.Listener, error) {
	listener, err := internet.ListenSystem(ctx, &net.TCPAddr{
		IP:   address.IP(),
		Port: int(port),
	}, sockopt)
	if err != nil {
		return nil, newError("failed to listen TCP on", address, ":", port).Base(err)
	}

	if tlsConfig != nil {
		return tls.NewListener(listener, tlsConfig), nil
	}

	return listener, nil
}

// Addr implements net.Listener.Addr().
func (ln *Listener) Addr() net.Addr {
	return ln.listener.Addr()
}

// Close implements net.Listener.Close().
func (ln *Listener) Close() error {
	return ln.listener.Close()
}

func init() {
	common.Must(internet.RegisterTransportListener(protocolName, ListenWS))
}
<commit_msg>Allow the use of Browser Bridge<commit_after>// +build !confonly

package websocket

import (
	"context"
	"crypto/tls"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"v2ray.com/core/common"
	"v2ray.com/core/common/net"
	http_proto "v2ray.com/core/common/protocol/http"
	"v2ray.com/core/common/session"
	"v2ray.com/core/transport/internet"
	v2tls "v2ray.com/core/transport/internet/tls"
)

type requestHandler struct {
	path string
	ln   *Listener
}

var upgrader = &websocket.Upgrader{
	ReadBufferSize:   4 * 1024,
	WriteBufferSize:  4 * 1024,
	HandshakeTimeout: time.Second * 4,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

func (h *requestHandler) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	if request.URL.Path != h.path {
		writer.WriteHeader(http.StatusNotFound)
		return
	}
	conn, err := upgrader.Upgrade(writer, request, nil)
	if err != nil {
		newError("failed to convert to WebSocket connection").Base(err).WriteToLog()
		return
	}

	forwardedAddrs := http_proto.ParseXForwardedFor(request.Header)
	remoteAddr := conn.RemoteAddr()
	if len(forwardedAddrs) > 0 && forwardedAddrs[0].Family().IsIP() {
		remoteAddr.(*net.TCPAddr).IP = forwardedAddrs[0].IP()
	}

	h.ln.addConn(newConnection(conn, remoteAddr))
}

type Listener struct {
	sync.Mutex
	server   http.Server
	listener net.Listener
	config   *Config
	addConn  internet.ConnHandler
}

func ListenWS(ctx context.Context, address net.Address, port net.Port, streamSettings *internet.MemoryStreamConfig, addConn internet.ConnHandler) (internet.Listener, error) {
	wsSettings := streamSettings.ProtocolSettings.(*Config)

	var tlsConfig *tls.Config
	if config := v2tls.ConfigFromStreamSettings(streamSettings); config != nil {
		tlsConfig = config.GetTLSConfig()
	}

	listener, err := listenTCP(ctx, address, port, tlsConfig, streamSettings.SocketSettings)
	if err != nil {
		return nil, err
	}

	l := &Listener{
		config:   wsSettings,
		addConn:  addConn,
		listener: listener,
	}

	l.server = http.Server{
		Handler: &requestHandler{
			path: wsSettings.GetNormalizedPath(),
			ln:   l,
		},
		ReadHeaderTimeout: time.Second * 4,
		MaxHeaderBytes:    2048,
	}

	go func() {
		if err := l.server.Serve(l.listener); err != nil {
			newError("failed to serve http for WebSocket").Base(err).AtWarning().WriteToLog(session.ExportIDToError(ctx))
		}
	}()

	return l, err
}

func listenTCP(ctx context.Context, address net.Address, port net.Port, tlsConfig *tls.Config, sockopt *internet.SocketConfig) (net.Listener, error) {
	listener, err := internet.ListenSystem(ctx, &net.TCPAddr{
		IP:   address.IP(),
		Port: int(port),
	}, sockopt)
	if err != nil {
		return nil, newError("failed to listen TCP on", address, ":", port).Base(err)
	}

	if tlsConfig != nil {
		return tls.NewListener(listener, tlsConfig), nil
	}

	return listener, nil
}

// Addr implements net.Listener.Addr().
func (ln *Listener) Addr() net.Addr {
	return ln.listener.Addr()
}

// Close implements net.Listener.Close().
func (ln *Listener) Close() error {
	return ln.listener.Close()
}

func init() {
	common.Must(internet.RegisterTransportListener(protocolName, ListenWS))
}
<|endoftext|>


<commit_before>// Copyright 2015 ThoughtWorks, Inc.

// This file is part of Gauge.

// Gauge is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Gauge is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Gauge.  If not, see <http://www.gnu.org/licenses/>.

package execution

import (
	"strconv"
	"time"

	"fmt"

	"os"

	"strings"

	"github.com/getgauge/gauge/api"
	"github.com/getgauge/gauge/config"
	"github.com/getgauge/gauge/execution/event"
	"github.com/getgauge/gauge/execution/rerun"
	"github.com/getgauge/gauge/execution/result"
	"github.com/getgauge/gauge/gauge"
	"github.com/getgauge/gauge/logger"
	"github.com/getgauge/gauge/manifest"
	"github.com/getgauge/gauge/plugin"
	"github.com/getgauge/gauge/plugin/install"
	"github.com/getgauge/gauge/reporter"
	"github.com/getgauge/gauge/runner"
	"github.com/getgauge/gauge/validation"
)

var NumberOfExecutionStreams int
var InParallel bool

type execution interface {
	run() *result.SuiteResult
}

type executionInfo struct {
	manifest        *manifest.Manifest
	specs           *gauge.SpecCollection
	runner          runner.Runner
	pluginHandler   *plugin.Handler
	errMaps         *validation.ValidationErrMaps
	inParallel      bool
	numberOfStreams int
	stream          int
}

func newExecutionInfo(s *gauge.SpecCollection, r runner.Runner, ph *plugin.Handler, e *validation.ValidationErrMaps, p bool, stream int) *executionInfo {
	m, err := manifest.ProjectManifest()
	if err != nil {
		logger.Fatalf(err.Error())
	}
	return &executionInfo{
		manifest:        m,
		specs:           s,
		runner:          r,
		pluginHandler:   ph,
		errMaps:         e,
		inParallel:      p,
		numberOfStreams: NumberOfExecutionStreams,
		stream:          stream,
	}
}

func ExecuteSpecs(specDirs []string) int {
	err := validateFlags()
	if err != nil {
		logger.Fatalf(err.Error())
	}
	if config.CheckUpdates() {
		i := &install.UpdateFacade{}
		i.BufferUpdateDetails()
		defer i.PrintUpdateBuffer()
	}

	res := validation.ValidateSpecs(specDirs)
	if len(res.Errs) > 0 {
		os.Exit(1)
	}
	if res.SpecCollection.Size() < 1 {
		logger.Info("No specifications found in %s.", strings.Join(specDirs, ", "))
		res.Runner.Kill()
		os.Exit(0)
	}
	event.InitRegistry()
	reporter.ListenExecutionEvents()
	rerun.ListenFailedScenarios()
	ei := newExecutionInfo(res.SpecCollection, res.Runner, nil, res.ErrMap, InParallel, 0)
	e := newExecution(ei)
	return printExecutionStatus(e.run(), res.ErrMap)
}

func newExecution(executionInfo *executionInfo) execution {
	if executionInfo.inParallel {
		return newParallelExecution(executionInfo)
	}
	return newSimpleExecution(executionInfo)
}

func startAPI() runner.Runner {
	sc := api.StartAPI()
	select {
	case runner := <-sc.RunnerChan:
		return runner
	case err := <-sc.ErrorChan:
		logger.Fatalf("Failed to start gauge API: %s", err.Error())
	}
	return nil
}

func printExecutionStatus(suiteResult *result.SuiteResult, errMap *validation.ValidationErrMaps) int {
	nSkippedScenarios := len(errMap.ScenarioErrs)
	nSkippedSpecs := suiteResult.SpecsSkippedCount
	var nExecutedSpecs int
	if len(suiteResult.SpecResults) != 0 {
		nExecutedSpecs = len(suiteResult.SpecResults) - nSkippedSpecs
	}
	nFailedSpecs := suiteResult.SpecsFailedCount
	nPassedSpecs := nExecutedSpecs - nFailedSpecs

	nExecutedScenarios := 0
	nFailedScenarios := 0
	nPassedScenarios := 0
	for _, specResult := range suiteResult.SpecResults {
		nExecutedScenarios += specResult.ScenarioCount
		nFailedScenarios += specResult.ScenarioFailedCount
	}
	nExecutedScenarios -= nSkippedScenarios
	nPassedScenarios = nExecutedScenarios - nFailedScenarios

	if nExecutedScenarios < 0 {
		nExecutedScenarios = 0
	}

	if nPassedScenarios < 0 {
		nPassedScenarios = 0
	}

	logger.Info("Specifications:\t%d executed\t%d passed\t%d failed\t%d skipped", nExecutedSpecs, nPassedSpecs, nFailedSpecs, nSkippedSpecs)
	logger.Info("Scenarios:\t%d executed\t%d passed\t%d failed\t%d skipped", nExecutedScenarios, nPassedScenarios, nFailedScenarios, nSkippedScenarios)
	logger.Info("\nTotal time taken: %s", time.Millisecond*time.Duration(suiteResult.ExecutionTime))

	if suiteResult.IsFailed || (nSkippedSpecs+nSkippedScenarios) > 0 {
		return 1
	}
	return 0
}

func validateFlags() error {
	if !InParallel {
		return nil
	}
	if NumberOfExecutionStreams < 1 {
		return fmt.Errorf("Invalid input(%s) to --n flag.", strconv.Itoa(NumberOfExecutionStreams))
	}
	if !isValidStrategy(Strategy) {
		return fmt.Errorf("Invalid input(%s) to --strategy flag.", Strategy)
	}
	return nil
}
<commit_msg>Removing unused func<commit_after>// Copyright 2015 ThoughtWorks, Inc.

// This file is part of Gauge.

// Gauge is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Gauge is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Gauge.  If not, see <http://www.gnu.org/licenses/>.

package execution

import (
	"strconv"
	"time"

	"fmt"

	"os"

	"strings"

	"github.com/getgauge/gauge/config"
	"github.com/getgauge/gauge/execution/event"
	"github.com/getgauge/gauge/execution/rerun"
	"github.com/getgauge/gauge/execution/result"
	"github.com/getgauge/gauge/gauge"
	"github.com/getgauge/gauge/logger"
	"github.com/getgauge/gauge/manifest"
	"github.com/getgauge/gauge/plugin"
	"github.com/getgauge/gauge/plugin/install"
	"github.com/getgauge/gauge/reporter"
	"github.com/getgauge/gauge/runner"
	"github.com/getgauge/gauge/validation"
)

var NumberOfExecutionStreams int
var InParallel bool

type execution interface {
	run() *result.SuiteResult
}

type executionInfo struct {
	manifest        *manifest.Manifest
	specs           *gauge.SpecCollection
	runner          runner.Runner
	pluginHandler   *plugin.Handler
	errMaps         *validation.ValidationErrMaps
	inParallel      bool
	numberOfStreams int
	stream          int
}

func newExecutionInfo(s *gauge.SpecCollection, r runner.Runner, ph *plugin.Handler, e *validation.ValidationErrMaps, p bool, stream int) *executionInfo {
	m, err := manifest.ProjectManifest()
	if err != nil {
		logger.Fatalf(err.Error())
	}
	return &executionInfo{
		manifest:        m,
		specs:           s,
		runner:          r,
		pluginHandler:   ph,
		errMaps:         e,
		inParallel:      p,
		numberOfStreams: NumberOfExecutionStreams,
		stream:          stream,
	}
}

func ExecuteSpecs(specDirs []string) int {
	err := validateFlags()
	if err != nil {
		logger.Fatalf(err.Error())
	}
	if config.CheckUpdates() {
		i := &install.UpdateFacade{}
		i.BufferUpdateDetails()
		defer i.PrintUpdateBuffer()
	}

	res := validation.ValidateSpecs(specDirs)
	if len(res.Errs) > 0 {
		os.Exit(1)
	}
	if res.SpecCollection.Size() < 1 {
		logger.Info("No specifications found in %s.", strings.Join(specDirs, ", "))
		res.Runner.Kill()
		os.Exit(0)
	}
	event.InitRegistry()
	reporter.ListenExecutionEvents()
	rerun.ListenFailedScenarios()
	ei := newExecutionInfo(res.SpecCollection, res.Runner, nil, res.ErrMap, InParallel, 0)
	e := newExecution(ei)
	return printExecutionStatus(e.run(), res.ErrMap)
}

func newExecution(executionInfo *executionInfo) execution {
	if executionInfo.inParallel {
		return newParallelExecution(executionInfo)
	}
	return newSimpleExecution(executionInfo)
}

func printExecutionStatus(suiteResult *result.SuiteResult, errMap *validation.ValidationErrMaps) int {
	nSkippedScenarios := len(errMap.ScenarioErrs)
	nSkippedSpecs := suiteResult.SpecsSkippedCount
	var nExecutedSpecs int
	if len(suiteResult.SpecResults) != 0 {
		nExecutedSpecs = len(suiteResult.SpecResults) - nSkippedSpecs
	}
	nFailedSpecs := suiteResult.SpecsFailedCount
	nPassedSpecs := nExecutedSpecs - nFailedSpecs

	nExecutedScenarios := 0
	nFailedScenarios := 0
	nPassedScenarios := 0
	for _, specResult := range suiteResult.SpecResults {
		nExecutedScenarios += specResult.ScenarioCount
		nFailedScenarios += specResult.ScenarioFailedCount
	}
	nExecutedScenarios -= nSkippedScenarios
	nPassedScenarios = nExecutedScenarios - nFailedScenarios

	if nExecutedScenarios < 0 {
		nExecutedScenarios = 0
	}

	if nPassedScenarios < 0 {
		nPassedScenarios = 0
	}

	logger.Info("Specifications:\t%d executed\t%d passed\t%d failed\t%d skipped", nExecutedSpecs, nPassedSpecs, nFailedSpecs, nSkippedSpecs)
	logger.Info("Scenarios:\t%d executed\t%d passed\t%d failed\t%d skipped", nExecutedScenarios, nPassedScenarios, nFailedScenarios, nSkippedScenarios)
	logger.Info("\nTotal time taken: %s", time.Millisecond*time.Duration(suiteResult.ExecutionTime))

	if suiteResult.IsFailed || (nSkippedSpecs+nSkippedScenarios) > 0 {
		return 1
	}
	return 0
}

func validateFlags() error {
	if !InParallel {
		return nil
	}
	if NumberOfExecutionStreams < 1 {
		return fmt.Errorf("Invalid input(%s) to --n flag.", strconv.Itoa(NumberOfExecutionStreams))
	}
	if !isValidStrategy(Strategy) {
		return fmt.Errorf("Invalid input(%s) to --strategy flag.", Strategy)
	}
	return nil
}
<|endoftext|>


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

/*
Package hmac implements the Keyed-Hash Message Authentication Code (HMAC) as
defined in U.S. Federal Information Processing Standards Publication 198.
An HMAC is a cryptographic hash that uses a key to sign a message.
The receiver verifies the hash by recomputing it using the same key.

Receivers should be careful to use Equal to compare MACs in order to avoid
timing side-channels:

	// CheckMAC reports whether messageMAC is a valid HMAC tag for message.
	func CheckMAC(message, messageMAC, key []byte) bool {
		mac := hmac.New(sha256.New, key)
		mac.Write(message)
		expectedMAC := mac.Sum(nil)
		return hmac.Equal(messageMAC, expectedMAC)
	}
*/
package hmac

import (
	"crypto/subtle"
	"hash"
)

// FIPS 198-1:
// https://csrc.nist.gov/publications/fips/fips198-1/FIPS-198-1_final.pdf

// key is zero padded to the block size of the hash function
// ipad = 0x36 byte repeated for key length
// opad = 0x5c byte repeated for key length
// hmac = H([key ^ opad] H([key ^ ipad] text))

type hmac struct {
	size         int
	blocksize    int
	opad, ipad   []byte
	outer, inner hash.Hash
}

func (h *hmac) Sum(in []byte) []byte {
	origLen := len(in)
	in = h.inner.Sum(in)
	h.outer.Reset()
	h.outer.Write(h.opad)
	h.outer.Write(in[origLen:])
	return h.outer.Sum(in[:origLen])
}

func (h *hmac) Write(p []byte) (n int, err error) {
	return h.inner.Write(p)
}

func (h *hmac) Size() int { return h.size }

func (h *hmac) BlockSize() int { return h.blocksize }

func (h *hmac) Reset() {
	h.inner.Reset()
	h.inner.Write(h.ipad)
}

// New returns a new HMAC hash using the given hash.Hash type and key.
// Note that unlike other hash implementations in the standard library,
// the returned Hash does not implement encoding.BinaryMarshaler
// or encoding.BinaryUnmarshaler.
func New(h func() hash.Hash, key []byte) hash.Hash {
	hm := new(hmac)
	hm.outer = h()
	hm.inner = h()
	hm.size = hm.inner.Size()
	hm.blocksize = hm.inner.BlockSize()
	hm.ipad = make([]byte, hm.blocksize)
	hm.opad = make([]byte, hm.blocksize)
	if len(key) > hm.blocksize {
		// If key is too big, hash it.
		hm.outer.Write(key)
		key = hm.outer.Sum(nil)
	}
	copy(hm.ipad, key)
	copy(hm.opad, key)
	for i := range hm.ipad {
		hm.ipad[i] ^= 0x36
	}
	for i := range hm.opad {
		hm.opad[i] ^= 0x5c
	}
	hm.inner.Write(hm.ipad)
	return hm
}

// Equal compares two MACs for equality without leaking timing information.
func Equal(mac1, mac2 []byte) bool {
	// We don't have to be constant time if the lengths of the MACs are
	// different as that suggests that a completely different hash function
	// was used.
	return subtle.ConstantTimeCompare(mac1, mac2) == 1
}
<commit_msg>Hmac: re-vendor<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
Package hmac implements the Keyed-Hash Message Authentication Code (HMAC) as
defined in U.S. Federal Information Processing Standards Publication 198.
An HMAC is a cryptographic hash that uses a key to sign a message.
The receiver verifies the hash by recomputing it using the same key.

Receivers should be careful to use Equal to compare MACs in order to avoid
timing side-channels:

	// ValidMAC reports whether messageMAC is a valid HMAC tag for message.
	func ValidMAC(message, messageMAC, key []byte) bool {
		mac := hmac.New(sha256.New, key)
		mac.Write(message)
		expectedMAC := mac.Sum(nil)
		return hmac.Equal(messageMAC, expectedMAC)
	}
*/
package hmac

import (
	"crypto/subtle"
	"hash"
)

// FIPS 198-1:
// https://csrc.nist.gov/publications/fips/fips198-1/FIPS-198-1_final.pdf

// key is zero padded to the block size of the hash function
// ipad = 0x36 byte repeated for key length
// opad = 0x5c byte repeated for key length
// hmac = H([key ^ opad] H([key ^ ipad] text))

type hmac struct {
	size         int
	blocksize    int
	opad, ipad   []byte
	outer, inner hash.Hash
}

func (h *hmac) Sum(in []byte) []byte {
	origLen := len(in)
	in = h.inner.Sum(in)
	h.outer.Reset()
	h.outer.Write(h.opad)
	h.outer.Write(in[origLen:])
	return h.outer.Sum(in[:origLen])
}

func (h *hmac) Write(p []byte) (n int, err error) {
	return h.inner.Write(p)
}

func (h *hmac) Size() int { return h.size }

func (h *hmac) BlockSize() int { return h.blocksize }

func (h *hmac) Reset() {
	h.inner.Reset()
	h.inner.Write(h.ipad)
}

// New returns a new HMAC hash using the given hash.Hash type and key.
// Note that unlike other hash implementations in the standard library,
// the returned Hash does not implement encoding.BinaryMarshaler
// or encoding.BinaryUnmarshaler.
func New(h func() hash.Hash, key []byte) hash.Hash {
	hm := new(hmac)
	hm.outer = h()
	hm.inner = h()
	hm.size = hm.inner.Size()
	hm.blocksize = hm.inner.BlockSize()
	hm.ipad = make([]byte, hm.blocksize)
	hm.opad = make([]byte, hm.blocksize)
	if len(key) > hm.blocksize {
		// If key is too big, hash it.
		hm.outer.Write(key)
		key = hm.outer.Sum(nil)
	}
	copy(hm.ipad, key)
	copy(hm.opad, key)
	for i := range hm.ipad {
		hm.ipad[i] ^= 0x36
	}
	for i := range hm.opad {
		hm.opad[i] ^= 0x5c
	}
	hm.inner.Write(hm.ipad)
	return hm
}

// Equal compares two MACs for equality without leaking timing information.
func Equal(mac1, mac2 []byte) bool {
	// We don't have to be constant time if the lengths of the MACs are
	// different as that suggests that a completely different hash function
	// was used.
	return subtle.ConstantTimeCompare(mac1, mac2) == 1
}
<|endoftext|>


<commit_before>package store

import (
	"context"
	"database/sql"

	"github.com/jbub/pgbouncer_exporter/internal/domain"

	"github.com/jmoiron/sqlx"
)

type stat struct {
	Database          string `db:"database"`
	TotalRequests     int64  `db:"total_requests"`
	TotalReceived     int64  `db:"total_received"`
	TotalSent         int64  `db:"total_sent"`
	TotalQueryTime    int64  `db:"total_query_time"`
	TotalXactCount    int64  `db:"total_xact_count"`
	TotalXactTime     int64  `db:"total_xact_time"`
	TotalQueryCount   int64  `db:"total_query_count"`
	TotalWaitTime     int64  `db:"total_wait_time"`
	AverageRequests   int64  `db:"avg_req"`
	AverageReceived   int64  `db:"avg_recv"`
	AverageSent       int64  `db:"avg_sent"`
	AverageQuery      int64  `db:"avg_query"`
	AverageQueryCount int64  `db:"avg_query_count"`
	AverageQueryTime  int64  `db:"avg_query_time"`
	AverageXactTime   int64  `db:"avg_xact_time"`
	AverageXactCount  int64  `db:"avg_xact_count"`
	AverageWaitTime   int64  `db:"avg_wait_time"`
}

type pool struct {
	Database     string         `db:"database"`
	User         string         `db:"user"`
	Active       int64          `db:"cl_active"`
	Waiting      int64          `db:"cl_waiting"`
	ServerActive int64          `db:"sv_active"`
	ServerIdle   int64          `db:"sv_idle"`
	ServerUsed   int64          `db:"sv_used"`
	ServerTested int64          `db:"sv_tested"`
	ServerLogin  int64          `db:"sv_login"`
	MaxWait      int64          `db:"maxwait"`
	MaxWaitUs    int64          `db:"maxwait_us"`
	PoolMode     sql.NullString `db:"pool_mode"`
}

type database struct {
	Name               string         `db:"name"`
	Host               sql.NullString `db:"host"`
	Port               int64          `db:"port"`
	Database           string         `db:"database"`
	ForceUser          sql.NullString `db:"force_user"`
	PoolSize           int64          `db:"pool_size"`
	ReservePool        int64          `db:"reserve_pool"`
	PoolMode           sql.NullString `db:"pool_mode"`
	MaxConnections     int64          `db:"max_connections"`
	CurrentConnections int64          `db:"current_connections"`
	Paused             int64          `db:"paused"`
	Disabled           int64          `db:"disabled"`
}

type list struct {
	List  string `db:"list"`
	Items int64  `db:"items"`
}

// NewSQLStore returns new SQLStore.
func NewSQLStore(dataSource string) (*SQLStore, error) {
	db, err := sqlx.Connect("postgres", dataSource)
	if err != nil {
		return nil, err
	}
	return &SQLStore{db: db}, nil
}

// SQLStore is a sql based Store implementation.
type SQLStore struct {
	db *sqlx.DB
}

// GetStats returns stats.
func (s *SQLStore) GetStats(ctx context.Context) ([]domain.Stat, error) {
	var stats []stat
	if err := s.db.SelectContext(ctx, &stats, "SHOW STATS"); err != nil {
		return nil, err
	}
	var result []domain.Stat
	for _, row := range stats {
		result = append(result, domain.Stat(row))
	}
	return result, nil
}

// GetPools returns pools.
func (s *SQLStore) GetPools(ctx context.Context) ([]domain.Pool, error) {
	var pools []pool
	if err := s.db.SelectContext(ctx, &pools, "SHOW POOLS"); err != nil {
		return nil, err
	}
	var result []domain.Pool
	for _, row := range pools {
		result = append(result, domain.Pool{
			Database:     row.Database,
			User:         row.User,
			Active:       row.Active,
			Waiting:      row.Waiting,
			ServerActive: row.ServerActive,
			ServerIdle:   row.ServerIdle,
			ServerUsed:   row.ServerUsed,
			ServerTested: row.ServerTested,
			ServerLogin:  row.ServerLogin,
			MaxWait:      row.MaxWait,
			MaxWaitUs:    row.MaxWaitUs,
			PoolMode:     row.PoolMode.String,
		})
	}
	return result, nil
}

// GetDatabases returns databases.
func (s *SQLStore) GetDatabases(ctx context.Context) ([]domain.Database, error) {
	var databases []database
	if err := s.db.SelectContext(ctx, &databases, "SHOW DATABASES"); err != nil {
		return nil, err
	}
	var result []domain.Database
	for _, row := range databases {
		result = append(result, domain.Database{
			Name:               row.Name,
			Host:               row.Host.String,
			Port:               row.Port,
			Database:           row.Database,
			ForceUser:          row.ForceUser.String,
			PoolSize:           row.PoolSize,
			ReservePool:        row.ReservePool,
			PoolMode:           row.PoolMode.String,
			MaxConnections:     row.MaxConnections,
			CurrentConnections: row.CurrentConnections,
			Paused:             row.Paused,
			Disabled:           row.Disabled,
		})
	}
	return result, nil
}

// GetLists returns lists.
func (s *SQLStore) GetLists(ctx context.Context) ([]domain.List, error) {
	var lists []list
	if err := s.db.SelectContext(ctx, &lists, "SHOW LISTS"); err != nil {
		return nil, err
	}
	var result []domain.List
	for _, row := range lists {
		result = append(result, domain.List(row))
	}
	return result, nil
}

// Check checks the health of the store.
func (s *SQLStore) Check(ctx context.Context) error {
	// we cant use db.Ping because it is making a ";" sql query which pgbouncer does not support
	rows, err := s.db.QueryContext(ctx, "SHOW VERSION")
	if err != nil {
		return err
	}
	return rows.Close()
}

// Close closes the store.
func (s *SQLStore) Close() error {
	return s.db.Close()
}
<commit_msg>Use sqlx.Open instead of sqlx.Connect to skip calling Ping.<commit_after>package store

import (
	"context"
	"database/sql"

	"github.com/jbub/pgbouncer_exporter/internal/domain"

	"github.com/jmoiron/sqlx"
)

type stat struct {
	Database          string `db:"database"`
	TotalRequests     int64  `db:"total_requests"`
	TotalReceived     int64  `db:"total_received"`
	TotalSent         int64  `db:"total_sent"`
	TotalQueryTime    int64  `db:"total_query_time"`
	TotalXactCount    int64  `db:"total_xact_count"`
	TotalXactTime     int64  `db:"total_xact_time"`
	TotalQueryCount   int64  `db:"total_query_count"`
	TotalWaitTime     int64  `db:"total_wait_time"`
	AverageRequests   int64  `db:"avg_req"`
	AverageReceived   int64  `db:"avg_recv"`
	AverageSent       int64  `db:"avg_sent"`
	AverageQuery      int64  `db:"avg_query"`
	AverageQueryCount int64  `db:"avg_query_count"`
	AverageQueryTime  int64  `db:"avg_query_time"`
	AverageXactTime   int64  `db:"avg_xact_time"`
	AverageXactCount  int64  `db:"avg_xact_count"`
	AverageWaitTime   int64  `db:"avg_wait_time"`
}

type pool struct {
	Database     string         `db:"database"`
	User         string         `db:"user"`
	Active       int64          `db:"cl_active"`
	Waiting      int64          `db:"cl_waiting"`
	ServerActive int64          `db:"sv_active"`
	ServerIdle   int64          `db:"sv_idle"`
	ServerUsed   int64          `db:"sv_used"`
	ServerTested int64          `db:"sv_tested"`
	ServerLogin  int64          `db:"sv_login"`
	MaxWait      int64          `db:"maxwait"`
	MaxWaitUs    int64          `db:"maxwait_us"`
	PoolMode     sql.NullString `db:"pool_mode"`
}

type database struct {
	Name               string         `db:"name"`
	Host               sql.NullString `db:"host"`
	Port               int64          `db:"port"`
	Database           string         `db:"database"`
	ForceUser          sql.NullString `db:"force_user"`
	PoolSize           int64          `db:"pool_size"`
	ReservePool        int64          `db:"reserve_pool"`
	PoolMode           sql.NullString `db:"pool_mode"`
	MaxConnections     int64          `db:"max_connections"`
	CurrentConnections int64          `db:"current_connections"`
	Paused             int64          `db:"paused"`
	Disabled           int64          `db:"disabled"`
}

type list struct {
	List  string `db:"list"`
	Items int64  `db:"items"`
}

// NewSQLStore returns new SQLStore.
func NewSQLStore(dataSource string) (*SQLStore, error) {
	db, err := sqlx.Open("postgres", dataSource)
	if err != nil {
		return nil, err
	}
	return &SQLStore{db: db}, nil
}

// SQLStore is a sql based Store implementation.
type SQLStore struct {
	db *sqlx.DB
}

// GetStats returns stats.
func (s *SQLStore) GetStats(ctx context.Context) ([]domain.Stat, error) {
	var stats []stat
	if err := s.db.SelectContext(ctx, &stats, "SHOW STATS"); err != nil {
		return nil, err
	}
	var result []domain.Stat
	for _, row := range stats {
		result = append(result, domain.Stat(row))
	}
	return result, nil
}

// GetPools returns pools.
func (s *SQLStore) GetPools(ctx context.Context) ([]domain.Pool, error) {
	var pools []pool
	if err := s.db.SelectContext(ctx, &pools, "SHOW POOLS"); err != nil {
		return nil, err
	}
	var result []domain.Pool
	for _, row := range pools {
		result = append(result, domain.Pool{
			Database:     row.Database,
			User:         row.User,
			Active:       row.Active,
			Waiting:      row.Waiting,
			ServerActive: row.ServerActive,
			ServerIdle:   row.ServerIdle,
			ServerUsed:   row.ServerUsed,
			ServerTested: row.ServerTested,
			ServerLogin:  row.ServerLogin,
			MaxWait:      row.MaxWait,
			MaxWaitUs:    row.MaxWaitUs,
			PoolMode:     row.PoolMode.String,
		})
	}
	return result, nil
}

// GetDatabases returns databases.
func (s *SQLStore) GetDatabases(ctx context.Context) ([]domain.Database, error) {
	var databases []database
	if err := s.db.SelectContext(ctx, &databases, "SHOW DATABASES"); err != nil {
		return nil, err
	}
	var result []domain.Database
	for _, row := range databases {
		result = append(result, domain.Database{
			Name:               row.Name,
			Host:               row.Host.String,
			Port:               row.Port,
			Database:           row.Database,
			ForceUser:          row.ForceUser.String,
			PoolSize:           row.PoolSize,
			ReservePool:        row.ReservePool,
			PoolMode:           row.PoolMode.String,
			MaxConnections:     row.MaxConnections,
			CurrentConnections: row.CurrentConnections,
			Paused:             row.Paused,
			Disabled:           row.Disabled,
		})
	}
	return result, nil
}

// GetLists returns lists.
func (s *SQLStore) GetLists(ctx context.Context) ([]domain.List, error) {
	var lists []list
	if err := s.db.SelectContext(ctx, &lists, "SHOW LISTS"); err != nil {
		return nil, err
	}
	var result []domain.List
	for _, row := range lists {
		result = append(result, domain.List(row))
	}
	return result, nil
}

// Check checks the health of the store.
func (s *SQLStore) Check(ctx context.Context) error {
	// we cant use db.Ping because it is making a ";" sql query which pgbouncer does not support
	rows, err := s.db.QueryContext(ctx, "SHOW VERSION")
	if err != nil {
		return err
	}
	return rows.Close()
}

// Close closes the store.
func (s *SQLStore) Close() error {
	return s.db.Close()
}
<|endoftext|>