<commit_before>// Proberk is an HTTP prober that periodically sends out HTTP requests to specified
// endpoints and reports if the returned results match the expectations. The results
// of the probe, including latency, are recored in metrics2.
package main

import (
	"crypto/md5"
	"crypto/tls"
	"flag"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/flynn/json5"
	"go.skia.org/infra/go/auth"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/httputils"
	"go.skia.org/infra/go/human"
	"go.skia.org/infra/go/metrics2"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
	"go.skia.org/infra/proberk/go/types"
)

// flags
var (
	config   = flag.String("config", "probersk.json5", "Prober config filename.")
	local    = flag.Bool("local", false, "True if running locally (as opposed to in production).")
	promPort = flag.String("prom_port", ":10110", "Metrics service address (e.g., ':10110')")
	runEvery = flag.Duration("run_every", 1*time.Minute, "How often to run the probes.")
	validate = flag.Bool("validate", false, "Validate the config file and then exit.")
)

var (
	// responseTesters is a mapping of names to functions that test response bodies.
	responseTesters = map[string]types.ResponseTester{
		"nonZeroContentLength":     nonZeroContentLength,
		"skfiddleJSONBad":          skfiddleJSONBad,
		"skfiddleJSONGood":         skfiddleJSONGood,
		"skfiddleJSONSecViolation": skfiddleJSONSecViolation,
		"validJSON":                validJSON,
	}

	// The hash of the config file contents when the app started.
	startHash = ""
)

const (
	DIAL_TIMEOUT               = time.Duration(5 * time.Second)
	REQUEST_TIMEOUT            = time.Duration(30 * time.Second)
	ISSUE_TRACKER_PERIOD       = 15 * time.Minute
	DEFAULT_SSL_VALID_DURATION = 10 * 24 * time.Hour // cert should be valid for at least 10 days
)

func readConfigFile(filename string) (types.Probes, error) {
	allProbes := types.Probes{}
	errs := []string{}
	file, err := os.Open(filename)
	if err != nil {
		return nil, fmt.Errorf("Failed to open config file: %s", err)
	}
	d := json5.NewDecoder(file)
	p := &types.Probes{}
	if err := d.Decode(p); err != nil {
		return nil, fmt.Errorf("Failed to decode JSON in config file: %s", err)
	}
	for k, v := range *p {
		v.Failure = map[string]metrics2.Int64Metric{}
		v.Latency = map[string]metrics2.Int64Metric{}
		if v.ResponseTestName != "" {
			if f, ok := responseTesters[v.ResponseTestName]; ok {
				v.ResponseTest = f
			} else {
				errs = append(errs, fmt.Sprintf("ResponseTestName Not Found %q", k))
			}
		}
		allProbes[k] = v
	}
	if len(errs) != 0 {
		return nil, fmt.Errorf("%s", strings.Join(errs, "\n  "))
	}
	return allProbes, nil
}

// In returns true if n is found in list.
func In(n int, list []int) bool {
	for _, x := range list {
		if x == n {
			return true
		}
	}
	return false
}

// nonZeroContentLength tests whether the Content-Length value is non-zero.
func nonZeroContentLength(r io.Reader, headers http.Header) bool {
	return headers.Get("Content-Length") != "0"
}

// validJSON tests whether the response contains valid JSON.
func validJSON(r io.Reader, headers http.Header) bool {
	var i interface{}
	return json5.NewDecoder(r).Decode(&i) == nil
}

type skfiddleResp struct {
	CompileErrors []interface{} `json:"compile_errors"`
	RuntimeError  string        `json:"runtime_error"`
}

// skfiddleJSONSecViolation tests that the compile failed with a runtime error (which includes security violations).
func skfiddleJSONSecViolation(r io.Reader, headers http.Header) bool {
	dec := json5.NewDecoder(r)
	s := skfiddleResp{
		CompileErrors: []interface{}{},
	}
	if err := dec.Decode(&s); err != nil {
		sklog.Warningf("Failed to decode skfiddle JSON: %#v %s", s, err)
		return false
	}
	sklog.Infof("%#v", s)
	return s.RuntimeError != ""
}

// skfiddleJSONGood tests that the compile completed w/o error.
func skfiddleJSONGood(r io.Reader, headers http.Header) bool {
	dec := json5.NewDecoder(r)
	s := skfiddleResp{
		CompileErrors: []interface{}{},
	}
	if err := dec.Decode(&s); err != nil {
		sklog.Warningf("Failed to decode skfiddle JSON: %#v %s", s, err)
		return false
	}
	sklog.Infof("%#v", s)
	return len(s.CompileErrors) == 0 && s.RuntimeError == ""
}

// skfiddleJSONBad tests that the compile completed w/error.
func skfiddleJSONBad(r io.Reader, headers http.Header) bool {
	dec := json5.NewDecoder(r)
	s := skfiddleResp{
		CompileErrors: []interface{}{},
	}
	if err := dec.Decode(&s); err != nil {
		sklog.Warningf("Failed to decode skfiddle JSON: %#v %s", s, err)
		return false
	}
	sklog.Infof("%#v", s)
	return len(s.CompileErrors) != 0
}

// decodeJSONObject reads a JSON object from r and returns the resulting object. Returns nil if the
// JSON is invalid or can't be decoded to a map[string]interface{}.
func decodeJSONObject(r io.Reader) map[string]interface{} {
	var obj map[string]interface{}
	if json5.NewDecoder(r).Decode(&obj) != nil {
		return nil
	}
	return obj
}

// hasKeys tests that the given decoded JSON object has at least the provided keys. If obj is nil,
// returns false.
func hasKeys(obj map[string]interface{}, keys []string) bool {
	if obj == nil {
		return false
	}
	for _, key := range keys {
		if _, ok := obj[key]; !ok {
			return false
		}
	}
	return true
}

func probeOneRound(cfg types.Probes, anonymousClient, authClient *http.Client) {
	var resp *http.Response
	var begin time.Time
	for name, probe := range cfg {
		c := anonymousClient
		if probe.Authenticated {
			c = authClient
		}
		for _, url := range probe.URLs {
			sklog.Infof("Probe: %s Starting fail value: %d", name, probe.Failure[url].Get())
			begin = time.Now()
			var err error
			if probe.Method == "GET" {
				resp, err = c.Get(url)
			} else if probe.Method == "HEAD" {
				resp, err = c.Head(url)
			} else if probe.Method == "POST" {
				resp, err = c.Post(url, probe.MimeType, strings.NewReader(probe.Body))
			} else if probe.Method == "SSL" {
				// SSL is a fictitious method that tests the SSL cert.
				// TODO(stephana): We should consider refactoring the prober into a
				// more generic prober beyond just HTTP related probes. SSL probing
				// would fit cleaner into such a framework.
				if err := probeSSL(probe, url); err != nil {
					sklog.Errorf("While testing %s we got SSL error: %s", url, err)
					probe.Failure[url].Update(1)
				} else {
					probe.Failure[url].Update(0)
				}
				continue
			} else {
				sklog.Errorf("Error: unknown method: %s", probe.Method)
				continue
			}
			d := time.Since(begin)
			probe.Latency[url].Update(d.Nanoseconds() / int64(time.Millisecond))
			if err != nil {
				errorOkay := false
				for _, expect := range probe.Expected {
					if strings.Contains(err.Error(), fmt.Sprintf("statuscode %d", expect)) {
						errorOkay = true
						break
					}
				}
				if !errorOkay {
					sklog.Warningf("Failed to make request: Name: %s URL: %s Error: %s", name, url, err)
					probe.Failure[url].Update(1)
					continue
				}
			}
			if resp != nil {
				responseTestResults := true
				if probe.ResponseTest != nil && resp.Body != nil {
					responseTestResults = probe.ResponseTest(resp.Body, resp.Header)
				}
				if resp.Body != nil {
					util.Close(resp.Body)
				}
				// TODO(jcgregorio) Save the last N responses and present them in a web UI.

				if !In(resp.StatusCode, probe.Expected) {
					sklog.Errorf("Got wrong status code: Name %s Got %d Want %v", name, resp.StatusCode, probe.Expected)
					probe.Failure[url].Update(1)
					continue
				}
				if !responseTestResults {
					sklog.Warningf("Response test failed: Name: %s %#v", name, probe)
					probe.Failure[url].Update(1)
					continue
				}
			}

			probe.Failure[url].Update(0)
		}
	}
}

// probeSSL inspects the SSL cert for the given URL and checks whether
// the time to expiration is below a certain number of days.
func probeSSL(probe *types.Probe, URL string) error {
	parsedURL, err := url.Parse(URL)
	targetAddr := parsedURL.Host
	if parsedURL.Port() == "" {
		targetAddr += ":443"
	}

	// TODO(stephana): Revisit whether we want to do host checking.

	// Don't verify the host, since we also want to test IP addresses and
	// this is not a generic SSL prober. We will catch a host mismatch during
	// manual testing.
	conf := &tls.Config{
		InsecureSkipVerify: true,
	}

	conn, err := tls.Dial("tcp", targetAddr, conf)
	if err != nil {
		return sklog.FmtErrorf("Error dialing %s: %s", targetAddr, err)
	}
	defer util.Close(conn)

	// Establish the connection.
	if err := conn.Handshake(); err != nil {
		return sklog.FmtErrorf("Handshake with %s failed: %s", targetAddr, err)
	}

	certs := conn.ConnectionState().PeerCertificates
	if len(certs) == 0 {
		return sklog.FmtErrorf("Unable to retrieve peer certificates for %s", targetAddr)
	}

	// Validate the certificate chain
	now := time.Now()
	for _, cert := range certs {
		// Make sure the cert is valid.
		if now.Before(cert.NotBefore) {
			return sklog.FmtErrorf("Certificate for %s is not valid yet", targetAddr)
		}

		// If the 'expected' value of the probe configuration contains a positive
		// integer we interpret it as the number of days in the future the cert
		// should be valid.
		minExpirationDelta := DEFAULT_SSL_VALID_DURATION
		if len(probe.Expected) > 0 && probe.Expected[0] > 0 {
			minExpirationDelta = time.Duration(probe.Expected[0]) * 24 * time.Hour
		}

		// Make sure the cert is not expired.
		delta := cert.NotAfter.Sub(now)
		if delta < minExpirationDelta {
			return sklog.FmtErrorf("Certificate for %s is expired or will expire in %s.", targetAddr, human.Duration(delta))
		}
	}

	return nil
}

func getHash() (string, error) {
	f, err := os.Open(*config)
	if err != nil {
		return "", fmt.Errorf("Failed to read config file while checking hash: %s", err)
	}
	defer util.Close(f)

	h := md5.New()
	if _, err := io.Copy(h, f); err != nil {
		return "", fmt.Errorf("Failed to copy bytes while checking hash: %s", err)
	}

	return fmt.Sprintf("%x", h.Sum(nil)), nil
}

func main() {
	common.InitWithMust(
		"probeserver",
		common.PrometheusOpt(promPort),
		common.MetricsLoggingOpt(),
	)
	var err error
	startHash, err = getHash()
	if err != nil {
		sklog.Fatalln("Failed to calculate hash of config file: ", err)
	}
	cfg, err := readConfigFile(*config)
	if *validate {
		if err != nil {
			fmt.Printf("Validation Failed:\n  %s\n", err)
			os.Exit(1)
		}
		fmt.Println("Validation Successful")
		os.Exit(0)
	}
	if err != nil {
		sklog.Fatalln("Failed to read config file: ", err)
	}

	liveness := metrics2.NewLiveness("probes")

	// Register counters for each probe.
	for name, probe := range cfg {
		for _, url := range probe.URLs {
			probe.Failure[url] = metrics2.GetInt64Metric("prober", map[string]string{"type": "failure", "probename": name, "url": url})
			probe.Latency[url] = metrics2.GetInt64Metric("prober", map[string]string{"type": "latency", "probename": name, "url": url})
		}
	}

	// Create a client that uses our dialer with a timeout.
	anonymousClient := httputils.NewConfiguredTimeoutClient(DIAL_TIMEOUT, REQUEST_TIMEOUT)
	anonymousClient.CheckRedirect = func(req *http.Request, via []*http.Request) error {
		return http.ErrUseLastResponse
	}
	ts, err := auth.NewDefaultTokenSource(*local, auth.SCOPE_USERINFO_EMAIL)
	if err != nil {
		sklog.Fatal(err)
	}
	authClient := httputils.DefaultClientConfig().WithTokenSource(ts).WithoutRetries().With2xxOnly().Client()
	authClient.CheckRedirect = func(req *http.Request, via []*http.Request) error {
		return http.ErrUseLastResponse
	}

	probeOneRound(cfg, anonymousClient, authClient)
	for range time.Tick(*runEvery) {
		probeOneRound(cfg, anonymousClient, authClient)
		liveness.Reset()

		currentHash, err := getHash()
		if err != nil {
			sklog.Errorf("Failed to verify hash of config file: %s", err)
			continue
		}
		if currentHash != startHash {
			fmt.Println("Restarting to pick up new config.")
			os.Exit(0)
		}
	}
}
<commit_msg>[proberk] Stop using ClientConfig.With2xxOnly<commit_after>// Proberk is an HTTP prober that periodically sends out HTTP requests to specified
// endpoints and reports if the returned results match the expectations. The results
// of the probe, including latency, are recored in metrics2.
package main

import (
	"crypto/md5"
	"crypto/tls"
	"flag"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/flynn/json5"
	"go.skia.org/infra/go/auth"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/httputils"
	"go.skia.org/infra/go/human"
	"go.skia.org/infra/go/metrics2"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
	"go.skia.org/infra/proberk/go/types"
)

// flags
var (
	config   = flag.String("config", "probersk.json5", "Prober config filename.")
	local    = flag.Bool("local", false, "True if running locally (as opposed to in production).")
	promPort = flag.String("prom_port", ":10110", "Metrics service address (e.g., ':10110')")
	runEvery = flag.Duration("run_every", 1*time.Minute, "How often to run the probes.")
	validate = flag.Bool("validate", false, "Validate the config file and then exit.")
)

var (
	// responseTesters is a mapping of names to functions that test response bodies.
	responseTesters = map[string]types.ResponseTester{
		"nonZeroContentLength":     nonZeroContentLength,
		"skfiddleJSONBad":          skfiddleJSONBad,
		"skfiddleJSONGood":         skfiddleJSONGood,
		"skfiddleJSONSecViolation": skfiddleJSONSecViolation,
		"validJSON":                validJSON,
	}

	// The hash of the config file contents when the app started.
	startHash = ""
)

const (
	DIAL_TIMEOUT               = time.Duration(5 * time.Second)
	REQUEST_TIMEOUT            = time.Duration(30 * time.Second)
	ISSUE_TRACKER_PERIOD       = 15 * time.Minute
	DEFAULT_SSL_VALID_DURATION = 10 * 24 * time.Hour // cert should be valid for at least 10 days
)

func readConfigFile(filename string) (types.Probes, error) {
	allProbes := types.Probes{}
	errs := []string{}
	file, err := os.Open(filename)
	if err != nil {
		return nil, fmt.Errorf("Failed to open config file: %s", err)
	}
	d := json5.NewDecoder(file)
	p := &types.Probes{}
	if err := d.Decode(p); err != nil {
		return nil, fmt.Errorf("Failed to decode JSON in config file: %s", err)
	}
	for k, v := range *p {
		v.Failure = map[string]metrics2.Int64Metric{}
		v.Latency = map[string]metrics2.Int64Metric{}
		if v.ResponseTestName != "" {
			if f, ok := responseTesters[v.ResponseTestName]; ok {
				v.ResponseTest = f
			} else {
				errs = append(errs, fmt.Sprintf("ResponseTestName Not Found %q", k))
			}
		}
		allProbes[k] = v
	}
	if len(errs) != 0 {
		return nil, fmt.Errorf("%s", strings.Join(errs, "\n  "))
	}
	return allProbes, nil
}

// In returns true if n is found in list.
func In(n int, list []int) bool {
	for _, x := range list {
		if x == n {
			return true
		}
	}
	return false
}

// nonZeroContentLength tests whether the Content-Length value is non-zero.
func nonZeroContentLength(r io.Reader, headers http.Header) bool {
	return headers.Get("Content-Length") != "0"
}

// validJSON tests whether the response contains valid JSON.
func validJSON(r io.Reader, headers http.Header) bool {
	var i interface{}
	return json5.NewDecoder(r).Decode(&i) == nil
}

type skfiddleResp struct {
	CompileErrors []interface{} `json:"compile_errors"`
	RuntimeError  string        `json:"runtime_error"`
}

// skfiddleJSONSecViolation tests that the compile failed with a runtime error (which includes security violations).
func skfiddleJSONSecViolation(r io.Reader, headers http.Header) bool {
	dec := json5.NewDecoder(r)
	s := skfiddleResp{
		CompileErrors: []interface{}{},
	}
	if err := dec.Decode(&s); err != nil {
		sklog.Warningf("Failed to decode skfiddle JSON: %#v %s", s, err)
		return false
	}
	sklog.Infof("%#v", s)
	return s.RuntimeError != ""
}

// skfiddleJSONGood tests that the compile completed w/o error.
func skfiddleJSONGood(r io.Reader, headers http.Header) bool {
	dec := json5.NewDecoder(r)
	s := skfiddleResp{
		CompileErrors: []interface{}{},
	}
	if err := dec.Decode(&s); err != nil {
		sklog.Warningf("Failed to decode skfiddle JSON: %#v %s", s, err)
		return false
	}
	sklog.Infof("%#v", s)
	return len(s.CompileErrors) == 0 && s.RuntimeError == ""
}

// skfiddleJSONBad tests that the compile completed w/error.
func skfiddleJSONBad(r io.Reader, headers http.Header) bool {
	dec := json5.NewDecoder(r)
	s := skfiddleResp{
		CompileErrors: []interface{}{},
	}
	if err := dec.Decode(&s); err != nil {
		sklog.Warningf("Failed to decode skfiddle JSON: %#v %s", s, err)
		return false
	}
	sklog.Infof("%#v", s)
	return len(s.CompileErrors) != 0
}

// decodeJSONObject reads a JSON object from r and returns the resulting object. Returns nil if the
// JSON is invalid or can't be decoded to a map[string]interface{}.
func decodeJSONObject(r io.Reader) map[string]interface{} {
	var obj map[string]interface{}
	if json5.NewDecoder(r).Decode(&obj) != nil {
		return nil
	}
	return obj
}

// hasKeys tests that the given decoded JSON object has at least the provided keys. If obj is nil,
// returns false.
func hasKeys(obj map[string]interface{}, keys []string) bool {
	if obj == nil {
		return false
	}
	for _, key := range keys {
		if _, ok := obj[key]; !ok {
			return false
		}
	}
	return true
}

func probeOneRound(cfg types.Probes, anonymousClient, authClient *http.Client) {
	var resp *http.Response
	var begin time.Time
	for name, probe := range cfg {
		c := anonymousClient
		if probe.Authenticated {
			c = authClient
		}
		for _, url := range probe.URLs {
			sklog.Infof("Probe: %s Starting fail value: %d", name, probe.Failure[url].Get())
			begin = time.Now()
			var err error
			if probe.Method == "GET" {
				resp, err = c.Get(url)
			} else if probe.Method == "HEAD" {
				resp, err = c.Head(url)
			} else if probe.Method == "POST" {
				resp, err = c.Post(url, probe.MimeType, strings.NewReader(probe.Body))
			} else if probe.Method == "SSL" {
				// SSL is a fictitious method that tests the SSL cert.
				// TODO(stephana): We should consider refactoring the prober into a
				// more generic prober beyond just HTTP related probes. SSL probing
				// would fit cleaner into such a framework.
				if err := probeSSL(probe, url); err != nil {
					sklog.Errorf("While testing %s we got SSL error: %s", url, err)
					probe.Failure[url].Update(1)
				} else {
					probe.Failure[url].Update(0)
				}
				continue
			} else {
				sklog.Errorf("Error: unknown method: %s", probe.Method)
				continue
			}
			d := time.Since(begin)
			probe.Latency[url].Update(d.Nanoseconds() / int64(time.Millisecond))
			if err != nil {
				sklog.Warningf("Failed to make request: Name: %s URL: %s Error: %s", name, url, err)
				probe.Failure[url].Update(1)
				continue
			}
			if resp != nil {
				responseTestResults := true
				if probe.ResponseTest != nil && resp.Body != nil {
					responseTestResults = probe.ResponseTest(resp.Body, resp.Header)
				}
				if resp.Body != nil {
					util.Close(resp.Body)
				}
				// TODO(jcgregorio) Save the last N responses and present them in a web UI.

				if !In(resp.StatusCode, probe.Expected) {
					sklog.Errorf("Got wrong status code: Name %s Got %d Want %v", name, resp.StatusCode, probe.Expected)
					probe.Failure[url].Update(1)
					continue
				}
				if !responseTestResults {
					sklog.Warningf("Response test failed: Name: %s %#v", name, probe)
					probe.Failure[url].Update(1)
					continue
				}
			}

			probe.Failure[url].Update(0)
		}
	}
}

// probeSSL inspects the SSL cert for the given URL and checks whether
// the time to expiration is below a certain number of days.
func probeSSL(probe *types.Probe, URL string) error {
	parsedURL, err := url.Parse(URL)
	targetAddr := parsedURL.Host
	if parsedURL.Port() == "" {
		targetAddr += ":443"
	}

	// TODO(stephana): Revisit whether we want to do host checking.

	// Don't verify the host, since we also want to test IP addresses and
	// this is not a generic SSL prober. We will catch a host mismatch during
	// manual testing.
	conf := &tls.Config{
		InsecureSkipVerify: true,
	}

	conn, err := tls.Dial("tcp", targetAddr, conf)
	if err != nil {
		return sklog.FmtErrorf("Error dialing %s: %s", targetAddr, err)
	}
	defer util.Close(conn)

	// Establish the connection.
	if err := conn.Handshake(); err != nil {
		return sklog.FmtErrorf("Handshake with %s failed: %s", targetAddr, err)
	}

	certs := conn.ConnectionState().PeerCertificates
	if len(certs) == 0 {
		return sklog.FmtErrorf("Unable to retrieve peer certificates for %s", targetAddr)
	}

	// Validate the certificate chain
	now := time.Now()
	for _, cert := range certs {
		// Make sure the cert is valid.
		if now.Before(cert.NotBefore) {
			return sklog.FmtErrorf("Certificate for %s is not valid yet", targetAddr)
		}

		// If the 'expected' value of the probe configuration contains a positive
		// integer we interpret it as the number of days in the future the cert
		// should be valid.
		minExpirationDelta := DEFAULT_SSL_VALID_DURATION
		if len(probe.Expected) > 0 && probe.Expected[0] > 0 {
			minExpirationDelta = time.Duration(probe.Expected[0]) * 24 * time.Hour
		}

		// Make sure the cert is not expired.
		delta := cert.NotAfter.Sub(now)
		if delta < minExpirationDelta {
			return sklog.FmtErrorf("Certificate for %s is expired or will expire in %s.", targetAddr, human.Duration(delta))
		}
	}

	return nil
}

func getHash() (string, error) {
	f, err := os.Open(*config)
	if err != nil {
		return "", fmt.Errorf("Failed to read config file while checking hash: %s", err)
	}
	defer util.Close(f)

	h := md5.New()
	if _, err := io.Copy(h, f); err != nil {
		return "", fmt.Errorf("Failed to copy bytes while checking hash: %s", err)
	}

	return fmt.Sprintf("%x", h.Sum(nil)), nil
}

func main() {
	common.InitWithMust(
		"probeserver",
		common.PrometheusOpt(promPort),
		common.MetricsLoggingOpt(),
	)
	var err error
	startHash, err = getHash()
	if err != nil {
		sklog.Fatalln("Failed to calculate hash of config file: ", err)
	}
	cfg, err := readConfigFile(*config)
	if *validate {
		if err != nil {
			fmt.Printf("Validation Failed:\n  %s\n", err)
			os.Exit(1)
		}
		fmt.Println("Validation Successful")
		os.Exit(0)
	}
	if err != nil {
		sklog.Fatalln("Failed to read config file: ", err)
	}

	liveness := metrics2.NewLiveness("probes")

	// Register counters for each probe.
	for name, probe := range cfg {
		for _, url := range probe.URLs {
			probe.Failure[url] = metrics2.GetInt64Metric("prober", map[string]string{"type": "failure", "probename": name, "url": url})
			probe.Latency[url] = metrics2.GetInt64Metric("prober", map[string]string{"type": "latency", "probename": name, "url": url})
		}
	}

	// Create a client that uses our dialer with a timeout.
	anonymousClient := httputils.NewConfiguredTimeoutClient(DIAL_TIMEOUT, REQUEST_TIMEOUT)
	anonymousClient.CheckRedirect = func(req *http.Request, via []*http.Request) error {
		return http.ErrUseLastResponse
	}
	ts, err := auth.NewDefaultTokenSource(*local, auth.SCOPE_USERINFO_EMAIL)
	if err != nil {
		sklog.Fatal(err)
	}
	authClient := httputils.DefaultClientConfig().WithTokenSource(ts).WithoutRetries().Client()
	authClient.CheckRedirect = func(req *http.Request, via []*http.Request) error {
		return http.ErrUseLastResponse
	}

	probeOneRound(cfg, anonymousClient, authClient)
	for range time.Tick(*runEvery) {
		probeOneRound(cfg, anonymousClient, authClient)
		liveness.Reset()

		currentHash, err := getHash()
		if err != nil {
			sklog.Errorf("Failed to verify hash of config file: %s", err)
			continue
		}
		if currentHash != startHash {
			fmt.Println("Restarting to pick up new config.")
			os.Exit(0)
		}
	}
}
<|endoftext|>


<commit_before>package models

import (
	"crypto/rand"
	"errors"
	"log"
	"math/big"
	"time"

	r "github.com/dancannon/gorethink"
	"golang.org/x/crypto/bcrypt"
)

// User holds attribute for an Ease user.
type User struct {
	ID           string    `gorethink:"id,omitempty" json:"id"`
	Username     string    `gorethink:"username" json:"username"`
	PasswordHash string    `gorethink:"password_hash" json:"-"`
	APIToken     string    `gorethink:"api_token" json:"api_token"`
	LoginToken   string    `gorethink:"login_token" json:"login_token"`
	CreatedAt    time.Time `gorethink:"created_at" json:"created_at"`
}

// UserQuerier queries the user table and logs users in.
type UserQuerier struct {
	session *r.Session
}

// NewUserQuerier returns a new UserQuerier.
func NewUserQuerier(session *r.Session) *UserQuerier {
	return &UserQuerier{session: session}
}

// NewUser creates a new user with tokens and a hashed password.
func NewUser(username, password string) (*User, error) {
	user := &User{}
	user.CreatedAt = time.Now()
	user.Username = username
	apiToken, err := generateRandomString(30)
	if err != nil {
		log.Println("Error: Couldn't generate random API token.")
		log.Println(err)
		return nil, err
	}
	user.APIToken = apiToken
	loginToken, err := generateRandomString(30)
	if err != nil {
		log.Println("Error: Couldn't generate random login token.")
		log.Println(err)
		return nil, err
	}
	user.LoginToken = loginToken

	byteHash, err :=
		bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
	user.PasswordHash = string(byteHash)
	if err != nil {
		log.Println("Error: Couldn't hash password.")
		log.Println(err)
		return nil, err
	}
	return user, nil
}

// Save saves the given user and returns it.
// It verifies that the given username isn't already taken.
// Returns the updated user.
func (querier *UserQuerier) Save(user *User) (*User, error) {
	// Check that a user with the given username doesn't already exist.
	otherUser := querier.Find(user.Username)
	if otherUser != nil && user.ID != otherUser.ID {
		return nil, errors.New("A user with that name already exists")
	}

	// Upsert the user.
	res, err := r.DB("test").Table("users").Insert(
		user, r.InsertOpts{Conflict: "replace"},
	).RunWrite(querier.session)

	if err != nil {
		friendlyErr := errors.New("Couldn't save user")
		log.Println(friendlyErr)
		log.Println(err)
		return nil, friendlyErr
	}

	// Get the user's ID if one was generated.
	if user.ID == "" {
		user.ID = res.GeneratedKeys[0]
	}

	return user, nil
}

// Find finds the user with the given username.  Returns nil if none found.
func (querier *UserQuerier) Find(username string) *User {
	res, err := r.DB("test").Table("users").Filter(map[string]string{
		"username": username,
	}).Run(querier.session)
	if err != nil || res.IsNil() {
		return nil
	}
	var user *User
	err = res.One(&user)
	if err != nil {
		return nil
	}
	return user
}

// AttemptLogin attempts to login the user with the given username and password.
// Returns the user if successful, nil if failed.
func (querier *UserQuerier) AttemptLogin(username, password string) (*User, error) {
	user := querier.Find(username)
	if user == nil {
		return nil, errors.New("Couldn't find user with that username")
	}
	err :=
		bcrypt.CompareHashAndPassword([]byte(user.PasswordHash), []byte(password))
	if err != nil {
		return nil, errors.New("Password was invalid")
	}
	return user, nil
}

// Possible token chars.
var letters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")

// generateRandomString generates a random string of length n.
func generateRandomString(n int) (string, error) {
	b := make([]rune, n)
	for i := range b {
		randInt, err := rand.Int(rand.Reader, big.NewInt(int64(len(letters))))
		if err != nil {
			return "", err
		}
		b[i] = letters[int(randInt.Int64())]
	}
	return string(b), nil
}
<commit_msg>Remove conflict replacement.<commit_after>package models

import (
	"crypto/rand"
	"errors"
	"log"
	"math/big"
	"time"

	r "github.com/dancannon/gorethink"
	"golang.org/x/crypto/bcrypt"
)

// User holds attribute for an Ease user.
type User struct {
	ID           string    `gorethink:"id,omitempty" json:"id"`
	Username     string    `gorethink:"username" json:"username"`
	PasswordHash string    `gorethink:"password_hash" json:"-"`
	APIToken     string    `gorethink:"api_token" json:"api_token"`
	LoginToken   string    `gorethink:"login_token" json:"login_token"`
	CreatedAt    time.Time `gorethink:"created_at" json:"created_at"`
}

// UserQuerier queries the user table and logs users in.
type UserQuerier struct {
	session *r.Session
}

// NewUserQuerier returns a new UserQuerier.
func NewUserQuerier(session *r.Session) *UserQuerier {
	return &UserQuerier{session: session}
}

// NewUser creates a new user with tokens and a hashed password.
func NewUser(username, password string) (*User, error) {
	user := &User{}
	user.CreatedAt = time.Now()
	user.Username = username
	apiToken, err := generateRandomString(30)
	if err != nil {
		log.Println("Error: Couldn't generate random API token.")
		log.Println(err)
		return nil, err
	}
	user.APIToken = apiToken
	loginToken, err := generateRandomString(30)
	if err != nil {
		log.Println("Error: Couldn't generate random login token.")
		log.Println(err)
		return nil, err
	}
	user.LoginToken = loginToken

	byteHash, err :=
		bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
	user.PasswordHash = string(byteHash)
	if err != nil {
		log.Println("Error: Couldn't hash password.")
		log.Println(err)
		return nil, err
	}
	return user, nil
}

// Save saves the given user and returns it.
// It verifies that the given username isn't already taken.
// Returns the updated user.
func (querier *UserQuerier) Save(user *User) (*User, error) {
	// Check that a user with the given username doesn't already exist.
	otherUser := querier.Find(user.Username)
	if otherUser != nil && user.ID != otherUser.ID {
		return nil, errors.New("A user with that name already exists")
	}

	// Upsert the user.
	res, err := r.DB("test").Table("users").Insert(user).RunWrite(querier.session)

	if err != nil {
		friendlyErr := errors.New("Couldn't save user")
		log.Println(friendlyErr)
		log.Println(err)
		return nil, friendlyErr
	}

	// Get the user's ID if one was generated.
	if user.ID == "" {
		user.ID = res.GeneratedKeys[0]
	}

	return user, nil
}

// Find finds the user with the given username.  Returns nil if none found.
func (querier *UserQuerier) Find(username string) *User {
	res, err := r.DB("test").Table("users").Filter(map[string]string{
		"username": username,
	}).Run(querier.session)
	if err != nil || res.IsNil() {
		return nil
	}
	var user *User
	err = res.One(&user)
	if err != nil {
		return nil
	}
	return user
}

// AttemptLogin attempts to login the user with the given username and password.
// Returns the user if successful, nil if failed.
func (querier *UserQuerier) AttemptLogin(username, password string) (*User, error) {
	user := querier.Find(username)
	if user == nil {
		return nil, errors.New("Couldn't find user with that username")
	}
	err :=
		bcrypt.CompareHashAndPassword([]byte(user.PasswordHash), []byte(password))
	if err != nil {
		return nil, errors.New("Password was invalid")
	}
	return user, nil
}

// Possible token chars.
var letters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")

// generateRandomString generates a random string of length n.
func generateRandomString(n int) (string, error) {
	b := make([]rune, n)
	for i := range b {
		randInt, err := rand.Int(rand.Reader, big.NewInt(int64(len(letters))))
		if err != nil {
			return "", err
		}
		b[i] = letters[int(randInt.Int64())]
	}
	return string(b), nil
}
<|endoftext|>


<commit_before>package poke

import (
	"math"
)

const LIFE_MULTIPLIER = 100000
const LEVEL = 50.0
const DIED_PENALTY = 100

type BestMoveSetResult struct {
	PokemonId   int     `json:"id"`
	PokemonName string  `json:"name"`
	MoveSet     [4]Move `json:"move_set"`
	TotalKt     int     `json:"kt"`
}

func TypeEffectiveness(attType Type, defTypes []Type) float64 {
	mult := 1.0
	for _, t := range defTypes {
		mult *= EffectMap[attType.Name][t.Name]
	}
	return mult
}

// Analyze all pokemon. Returns its kt.
func BestMoveSet(pokemon Pokemon, list []Pokemon) (best [4]Move, bestKt int) {
	bestKt = 10000000000

	combinations := GenerateCombinations(len(pokemon.LearnableMoves), 4) // TODO: What if learns < 4
	for _, moveVector := range combinations {
		moveSet := intersectMoves(pokemon, moveVector)
		if thereExistStrictlyBetter(pokemon, moveSet) {
			continue
		}

		// For each 150 pokemon:  Fight and get 'effective damage' of best move
		totalKt := 0
		for _, enemy := range list {
			kt := spamMoveFight(pokemon, moveSet, enemy)
			totalKt += kt
		}

		// Maintain bestKt
		if totalKt < bestKt {
			best = moveSet
			bestKt = totalKt
		}
	}
	return best, bestKt
}

// True if there is an attack of same PorS with higher Power*Accuracy.
func thereExistStrictlyBetter(p Pokemon, moveSet [4]Move) bool {
	for _, m := range p.LearnableMoves {
		if isInMoveSet(m, moveSet) {
			continue
		}
		for _, mm := range moveSet {
			sameCategory := (m.isPhysical && mm.isPhysical) || (!m.isPhysical && !mm.isPhysical)
			sameType := m.Type.Name == mm.Type.Name
			otherIsBetter := expectedPower(m) > expectedPower(mm)
			if sameCategory && sameType && otherIsBetter {
				return true
			}
		}
	}
	return false
}

func isInMoveSet(m Move, moveSet [4]Move) bool {
	for _, mm := range moveSet {
		if m.Name == mm.Name {
			return true
		}
	}
	return false
}

func intersectMoves(p Pokemon, moveVector []bool) [4]Move {
	moves := [4]Move{}
	i := 0
	for j, b := range moveVector {
		if b {
			moves[i] = p.LearnableMoves[j]
			i++
		}
	}
	return moves
}

// Simulates a fight where pokemons spam the best available move in every turn.
// Returns kt (kill turn).  If enemy wins, kt = 10
func spamMoveFight(a Pokemon, moveSet [4]Move, b Pokemon) (kt int) {
	// Check expected damage from each move.
	_, aDamage := bestMove(a, moveSet[:], b)
	_, bDamage := bestMove(b, b.LearnableMoves, a) // out of all possible moves.

	aKillTurn := int(math.Ceil(float64(b.BaseStats.Hp*LIFE_MULTIPLIER)/aDamage) * 2) // * 2 since have to wait for opponent
	bKillTurn := int(math.Ceil(float64(a.BaseStats.Hp*LIFE_MULTIPLIER)/bDamage) * 2)
	if a.BaseStats.Speed >= b.BaseStats.Speed {
		aKillTurn -= 1
	} else {
		bKillTurn -= 1
	}

	if aKillTurn > bKillTurn {
		aKillTurn += DIED_PENALTY
	}
	return aKillTurn
}

func bestMove(a Pokemon, moves []Move, b Pokemon) (move Move, damage float64) {
	damage = -1.0
	for _, m := range moves {
		d := computeMoveDamage(a, m, b)

		if d > damage {
			move = m
			damage = d
		}
	}
	return move, damage
}

func computeMoveDamage(a Pokemon, move Move, b Pokemon) float64 {
	var A float64
	var D float64
	if move.isPhysical {
		A = float64(a.BaseStats.Attack)
		D = float64(b.BaseStats.Defense)
	} else {
		A = float64(a.BaseStats.SpecialAttack)
		D = float64(b.BaseStats.SpecialDefense)
	}

	stab := 1.0 // same-type attack bonus
	for _, t := range a.Types {
		if move.Type == t {
			stab = 1.5
		}
	}
	levelTerm := ((2.0 * LEVEL) / 5.0) + 2.0
	modifier := TypeEffectiveness(move.Type, b.Types) * stab
	damage := (((levelTerm * expectedPower(move) * A / D) / 50) + 2) * modifier

	if damage <= 0.0 {
		damage = 1.0
	}
	return damage
}

func expectedPower(move Move) float64 {
	return float64(move.Power) * float64(move.Accuracy) / 100.0
}

// PROBLEM: When you fight; do you select best attack at the moment, or
// start with a plan and stick to it?  Maybe a combination: start with
// a plan and adjust appropriately (re-compute)?
<commit_msg>ThereExistStrictlyBetter Heuristic 2.  Before generating combs<commit_after>package poke

import (
	"math"
)

const LIFE_MULTIPLIER = 100000
const LEVEL = 50.0
const DIED_PENALTY = 100

type BestMoveSetResult struct {
	PokemonId   int     `json:"id"`
	PokemonName string  `json:"name"`
	MoveSet     [4]Move `json:"move_set"`
	TotalKt     int     `json:"kt"`
}

func TypeEffectiveness(attType Type, defTypes []Type) float64 {
	mult := 1.0
	for _, t := range defTypes {
		mult *= EffectMap[attType.Name][t.Name]
	}
	return mult
}

// Analyze all pokemon. Returns its kt.
func BestMoveSet(pokemon Pokemon, list []Pokemon) (best [4]Move, bestKt int) {
	bestKt = 10000000000

	candidates := learnableCandidates(pokemon.LearnableMoves)
	combinations := GenerateCombinations(len(candidates), 4) // TODO: What if learns < 4
	for _, moveVector := range combinations {
		moveSet := intersectMoves(candidates, moveVector)

		// For each 150 pokemon:  Fight and get 'effective damage' of best move
		totalKt := 0
		for _, enemy := range list {
			kt := spamMoveFight(pokemon, moveSet, enemy)
			totalKt += kt
		}

		// Maintain bestKt
		if totalKt < bestKt {
			best = moveSet
			bestKt = totalKt
		}
	}
	return best, bestKt
}

// Skip moves where there exist a strictly better option
func learnableCandidates(moves []Move) []Move {
	res := []Move{}
	for _, m := range moves {
		thereExistBetter := false
		for _, mm := range moves {
			if m == mm {
				continue
			}
			sameCategory := (m.isPhysical && mm.isPhysical) || (!m.isPhysical && !mm.isPhysical)
			sameType := m.Type.Name == mm.Type.Name
			otherIsBetter := expectedPower(m) < expectedPower(mm)
			if sameCategory && sameType && otherIsBetter {
				thereExistBetter = true
				break
			}
		}
		if !thereExistBetter {
			res = append(res, m)
		}
	}
	return res
}

func intersectMoves(moveSpace []Move, moveVector []bool) [4]Move {
	moves := [4]Move{}
	i := 0
	for j, b := range moveVector {
		if b {
			moves[i] = moveSpace[j]
			i++
		}
	}
	return moves
}

// Simulates a fight where pokemons spam the best available move in every turn.
// Returns kt (kill turn).  If enemy wins, kt = 10
func spamMoveFight(a Pokemon, moveSet [4]Move, b Pokemon) (kt int) {
	// Check expected damage from each move.
	_, aDamage := bestMove(a, moveSet[:], b)
	_, bDamage := bestMove(b, b.LearnableMoves, a) // out of all possible moves.

	aKillTurn := int(math.Ceil(float64(b.BaseStats.Hp*LIFE_MULTIPLIER)/aDamage) * 2) // * 2 since have to wait for opponent
	bKillTurn := int(math.Ceil(float64(a.BaseStats.Hp*LIFE_MULTIPLIER)/bDamage) * 2)
	if a.BaseStats.Speed >= b.BaseStats.Speed {
		aKillTurn -= 1
	} else {
		bKillTurn -= 1
	}

	if aKillTurn > bKillTurn {
		aKillTurn += DIED_PENALTY
	}
	return aKillTurn
}

func bestMove(a Pokemon, moves []Move, b Pokemon) (move Move, damage float64) {
	damage = -1.0
	for _, m := range moves {
		d := computeMoveDamage(a, m, b)

		if d > damage {
			move = m
			damage = d
		}
	}
	return move, damage
}

func computeMoveDamage(a Pokemon, move Move, b Pokemon) float64 {
	var A float64
	var D float64
	if move.isPhysical {
		A = float64(a.BaseStats.Attack)
		D = float64(b.BaseStats.Defense)
	} else {
		A = float64(a.BaseStats.SpecialAttack)
		D = float64(b.BaseStats.SpecialDefense)
	}

	stab := 1.0 // same-type attack bonus
	for _, t := range a.Types {
		if move.Type == t {
			stab = 1.5
		}
	}
	levelTerm := ((2.0 * LEVEL) / 5.0) + 2.0
	modifier := TypeEffectiveness(move.Type, b.Types) * stab
	damage := (((levelTerm * expectedPower(move) * A / D) / 50) + 2) * modifier

	if damage <= 0.0 {
		damage = 1.0
	}
	return damage
}

func expectedPower(move Move) float64 {
	return float64(move.Power) * float64(move.Accuracy) / 100.0
}

// PROBLEM: When you fight; do you select best attack at the moment, or
// start with a plan and stick to it?  Maybe a combination: start with
// a plan and adjust appropriately (re-compute)?
<|endoftext|>


<commit_before>/*
 * This file is part of the libvirt-go project
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 *
 * Copyright (C) 2017 Red Hat, Inc.
 *
 */

package libvirt

/*
#cgo pkg-config: libvirt
#include <libvirt/libvirt.h>
#include <assert.h>
#include "stream_compat.h"

int virStreamRecvFlagsCompat(virStreamPtr st,
			     char *data,
			     size_t nbytes,
			     unsigned int flags)
{
#if LIBVIR_VERSION_NUMBER < 3002000
    assert(0); // Caller should have checked version
#else
    return virStreamRecvFlags(st, data, nbytes, flags);
#endif
}

int virStreamSendHoleCompat(virStreamPtr st,
			    long long length,
			    unsigned int flags)
{
#if LIBVIR_VERSION_NUMBER < 3002000
    assert(0); // Caller should have checked version
#else
    return virStreamSendHole(st, length, flags);
#endif
}

int virStreamRecvHoleCompat(virStreamPtr st,
			    long long *length,
			    unsigned int flags)
{
#if LIBVIR_VERSION_NUMBER < 3002000
    assert(0); // Caller should have checked version
#else
    return virStreamRecvHole(st, length, flags);
#endif
}

*/
import "C"
<commit_msg>Fix version number check for stream funcs<commit_after>/*
 * This file is part of the libvirt-go project
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 *
 * Copyright (C) 2017 Red Hat, Inc.
 *
 */

package libvirt

/*
#cgo pkg-config: libvirt
#include <libvirt/libvirt.h>
#include <assert.h>
#include "stream_compat.h"

int virStreamRecvFlagsCompat(virStreamPtr st,
			     char *data,
			     size_t nbytes,
			     unsigned int flags)
{
#if LIBVIR_VERSION_NUMBER < 3004000
    assert(0); // Caller should have checked version
#else
    return virStreamRecvFlags(st, data, nbytes, flags);
#endif
}

int virStreamSendHoleCompat(virStreamPtr st,
			    long long length,
			    unsigned int flags)
{
#if LIBVIR_VERSION_NUMBER < 3004000
    assert(0); // Caller should have checked version
#else
    return virStreamSendHole(st, length, flags);
#endif
}

int virStreamRecvHoleCompat(virStreamPtr st,
			    long long *length,
			    unsigned int flags)
{
#if LIBVIR_VERSION_NUMBER < 3004000
    assert(0); // Caller should have checked version
#else
    return virStreamRecvHole(st, length, flags);
#endif
}

*/
import "C"
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"strings"

	"github.com/spf13/cobra"
)

var environment []string
var volumes []string

var checkSh = []byte(`#!/bin/bash
type docker >/dev/null 2>&1 || { echo >&2 "Can't find docker command."; exit 1; }

docker_path=""

if [ $OSTYPE == "linux-gnu" ] ; then
  docker_path=$(which docker)
  libltdl_path=$(ldd /usr/bin/docker | grep libltdl | awk '{print $3}')
  libltdl_path="--volume $libltdl_path:/usr/lib/$(basename $libltdl_path)"
elif  [[ $OSTYPE == *"darwin"* ]]; then
  docker_path="/usr/bin/docker"
  libltdl_path=""
else
  echo "Unrecognized OS: $OSTYPE"
  exit 1
fi

docker run --rm \
  $@ \
  $libltdl_path \
  --volume $PWD:$PWD \
  --volume /tmp:/tmp \
  --volume $docker_path:/usr/bin/docker \
  --volume /var/run/docker.sock:/var/run/docker.sock \
  --workdir $PWD \
  -e SKIP='/root/clone.sh' \
  ivotron/popperci
`)

var checkCmd = &cobra.Command{
	Use:   "check",
	Short: "Check integrity of an experiment",
	Long:  ``,
	Run: func(cmd *cobra.Command, args []string) {
		err := ioutil.WriteFile("/tmp/poppercheck", checkSh, 0755)
		if err != nil {
			log.Fatalln("Error writing shell command to /tmp")
		}
		envVols := ""
		if len(environment) > 0 {
			envVols += "-e " + strings.Join(environment, " -e ")
		}
		if len(volumes) > 0 {
			envVols += "-v " + strings.Join(volumes, " -v ")
		}
		fmt.Printf("%v\n", envVols)
		//sh.Command("/tmp/poppercheck " + envVols).Run()

	},
}

func init() {
	RootCmd.AddCommand(checkCmd)

	checkCmd.Flags().StringSliceVarP(&environment, "environment", "e", []string{}, "Environment variables to be defined inside the test container.")
	checkCmd.Flags().StringSliceVarP(&volumes, "volume", "v", []string{}, "Volumes to be passed to the test container.")
}
<commit_msg>Adds a --skip flag to 'popper check'<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"strings"

	sh "github.com/codeskyblue/go-sh"
	"github.com/spf13/cobra"
)

var environment []string
var volumes []string
var skip string
var checksh = `#!/bin/bash
set -e
type docker >/dev/null 2>&1 || { echo >&2 "Can't find docker command."; exit 1; }

docker_path=""

if [ $OSTYPE == "linux-gnu" ] ; then
  docker_path=$(which docker)
  libltdl_path=$(ldd $docker_path | grep libltdl | awk '{print $3}')
	if [ -n "$libltdl_path" ] ; then
    libltdl_path="--volume $libltdl_path:/usr/lib/$(basename $libltdl_path)"
	fi
elif  [[ $OSTYPE == *"darwin"* ]]; then
  docker_path="/usr/bin/docker"
  libltdl_path=""
else
  echo "Unrecognized OS: $OSTYPE"
  exit 1
fi

echo "Starting Popper check"
docker pull ivotron/popperci-experimenter &> /dev/null
docker run --rm %s \
  $libltdl_path \
  --volume $PWD:$PWD \
  --volume $docker_path:/usr/bin/docker \
  --volume /var/run/docker.sock:/var/run/docker.sock \
  --workdir $PWD \
  ivotron/popperci-experimenter %s
`

func writePopperCheckScript() {
	env := ""
	if len(environment) > 0 {
		env += " -e " + strings.Join(environment, " -e ")
	}
	if len(volumes) > 0 {
		env += " -v " + strings.Join(volumes, " -v ")
	}
	content := fmt.Sprintf(checksh, env, "--skip "+skip)
	err := ioutil.WriteFile("/tmp/poppercheck", []byte(content), 0755)
	if err != nil {
		log.Fatalln("Error writing bash script to /tmp")
	}
}

var checkCmd = &cobra.Command{
	Use:   "check",
	Short: "Check integrity of an experiment",
	Long:  ``,
	Run: func(cmd *cobra.Command, args []string) {
		writePopperCheckScript()
		if err := sh.Command("/tmp/poppercheck").Run(); err != nil {
			log.Fatalln(err)
		}
	},
}

func init() {
	RootCmd.AddCommand(checkCmd)

	checkCmd.Flags().StringSliceVarP(&environment, "environment", "e", []string{}, "Environment variables to be defined inside the test container.")
	checkCmd.Flags().StringSliceVarP(&volumes, "volume", "v", []string{}, "Volumes to be passed to the test container.")
	checkCmd.Flags().StringVarP(&skip, "skip", "s", "", "Comma-separated list of stages to skip.")
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/chzyer/readline"
	"github.com/golang/protobuf/ptypes"
	"github.com/google/shlex"
	"github.com/relab/gorums/examples/storage/proto"
)

var help = `
This interface allows you to run RPCs and quorum calls against the Storage
Servers interactively. The following commands can be used:

rpc	[node index] [operation]	Executes an RPC on the given node.
qc 	[operation]             	Executes a quorum call on all nodes.
cfg	[config] [opertation]   	Executes a quorum call on a given configuration.

The following operations are supported:

read 	[key]        	Read a value
write	[key] [value]	Write a value

Examples:

> rpc 0 write foo bar
The command performs the 'write' RPC on node 0, and sets 'foo' = 'bar'

> qc read foo
The command performs the 'read' quorum call, and returns the value of 'foo'

> cfg 1:3 write foo bar
The command performs the write qourum call on node 1 and 2

> cfg 0,2 write foo 'bar baz'
The command performs the write quorum call on node 0 and 2
`

var prompt = "> "

type repl struct {
	mgr *proto.Manager
	cfg *proto.Configuration
}

func Repl(mgr *proto.Manager, defaultCfg *proto.Configuration) {
	r := &repl{mgr, defaultCfg}
	rl, err := readline.New(prompt)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to create readline instance: %v\n", err)
		os.Exit(1)
	}
	defer rl.Close()

	fmt.Println(help)
	fmt.Print(prompt)

	for {
		l, err := rl.Readline()
		if errors.Is(err, io.EOF) {
			return
		}
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to read line: %v\n", err)
			os.Exit(1)
		}
		args, err := shlex.Split(l)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to split command: %v\n", err)
			os.Exit(1)
		}
		if len(args) < 1 {
			continue
		}

		switch args[0] {
		case "exit":
			return
		case "help":
			fmt.Println(help)
		case "rpc":
			r.rpc(args[1:])
		case "qc":
			r.qc(args[1:])
		case "cfg":
			r.qcCfg(args[1:])
		default:
			fmt.Printf("Unknown command '%s'. Type 'help' to see available commands.\n", args[0])
		}
	}
}

func (r repl) rpc(args []string) {
	if len(args) < 2 {
		fmt.Println("'rpc' requires a node index and an operation.")
		return
	}

	index, err := strconv.Atoi(args[0])
	if err != nil {
		fmt.Printf("Invalid id '%s'. node index must be numeric.\n", args[0])
		return
	}

	if index < 0 || index >= r.cfg.Size() {
		fmt.Printf("Invalid index. Must be between 0 and %d.\n", r.cfg.Size())
		return
	}

	node := r.cfg.Nodes()[index]

	switch args[1] {
	case "read":
		r.readRPC(args[2:], node)
	case "write":
		r.writeRPC(args[2:], node)
	}
}

func (r repl) qc(args []string) {
	if len(args) < 1 {
		fmt.Println("'qc' requires an operation.")
		return
	}

	switch args[0] {
	case "read":
		r.readQC(args[1:], r.cfg)
	case "write":
		r.writeQC(args[1:], r.cfg)
	}
}

func (r repl) qcCfg(args []string) {
	if len(args) < 2 {
		fmt.Println("'cfg' requires a configuration and an operation.")
		return
	}
	cfg := r.parseConfiguration(args[0])
	if cfg == nil {
		return
	}
	switch args[1] {
	case "read":
		r.readQC(args[2:], cfg)
	case "write":
		r.writeQC(args[2:], cfg)
	}
}

func (_ repl) readRPC(args []string, node *proto.Node) {
	if len(args) < 1 {
		fmt.Println("Read requires a key to read.")
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	resp, err := node.ReadRPC(ctx, &proto.ReadRequest{Key: args[0]})
	cancel()
	if err != nil {
		fmt.Printf("Read RPC finished with error: %v\n", err)
		return
	}
	if !resp.GetOK() {
		fmt.Printf("%s was not found\n", args[0])
		return
	}
	fmt.Printf("%s = %s\n", args[0], resp.GetValue())
}

func (_ repl) writeRPC(args []string, node *proto.Node) {
	if len(args) < 2 {
		fmt.Println("Write requires a key and a value to write.")
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	resp, err := node.WriteRPC(ctx, &proto.WriteRequest{Key: args[0], Value: args[1], Time: ptypes.TimestampNow()})
	cancel()
	if err != nil {
		fmt.Printf("Write RPC finished with error: %v\n", err)
		return
	}
	if !resp.GetNew() {
		fmt.Printf("Failed to update %s: timestamp too old.\n", args[0])
		return
	}
	fmt.Println("Write OK")
}

func (_ repl) readQC(args []string, cfg *proto.Configuration) {
	if len(args) < 1 {
		fmt.Println("Read requires a key to read.")
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	resp, err := cfg.ReadQC(ctx, &proto.ReadRequest{Key: args[0]})
	cancel()
	if err != nil {
		fmt.Printf("Read RPC finished with error: %v\n", err)
		return
	}
	if !resp.GetOK() {
		fmt.Printf("%s was not found\n", args[0])
		return
	}
	fmt.Printf("%s = %s\n", args[0], resp.GetValue())
}

func (_ repl) writeQC(args []string, cfg *proto.Configuration) {
	if len(args) < 2 {
		fmt.Println("Write requires a key and a value to write.")
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	resp, err := cfg.WriteQC(ctx, &proto.WriteRequest{Key: args[0], Value: args[1], Time: ptypes.TimestampNow()})
	cancel()
	if err != nil {
		fmt.Printf("Write RPC finished with error: %v\n", err)
		return
	}
	if !resp.GetNew() {
		fmt.Printf("Failed to update %s: timestamp too old.\n", args[0])
		return
	}
	fmt.Println("Write OK")
}

func (r repl) parseConfiguration(cfgStr string) (cfg *proto.Configuration) {
	// configuration using range syntax
	if i := strings.Index(cfgStr, ":"); i > -1 {
		var start, stop int
		var err error
		numNodes, _ := r.mgr.Size()
		if i == 0 {
			start = 0
		} else {
			start, err = strconv.Atoi(cfgStr[:i])
			if err != nil {
				fmt.Printf("Failed to parse configuration: %v\n", err)
				return nil
			}
		}
		if i == len(cfgStr)-1 {
			stop = numNodes
		} else {
			stop, err = strconv.Atoi(cfgStr[i+1:])
			if err != nil {
				fmt.Printf("Failed to parse configuration: %v\n", err)
				return nil
			}
		}
		if start >= stop || start < 0 || stop >= numNodes {
			fmt.Println("Invalid configuration.")
			return nil
		}
		cfg, err = r.mgr.NewConfiguration(r.mgr.NodeIDs()[start:stop], &qspec{cfgSize: stop - start})
		if err != nil {
			fmt.Printf("Failed to create configuration: %v\n", err)
			return nil
		}
		return cfg
	}
	// configuration using list of indices
	if indices := strings.Split(cfgStr, ","); len(indices) > 0 {
		selectedNodes := make([]uint32, 0, len(indices))
		nodeIDs := r.mgr.NodeIDs()
		for _, index := range indices {
			i, err := strconv.Atoi(index)
			if err != nil {
				fmt.Printf("Failed to parse configuration: %v\n", err)
				return nil
			}
			if i < 0 || i >= len(nodeIDs) {
				fmt.Println("Invalid configuration.")
				return nil
			}
			selectedNodes = append(selectedNodes, nodeIDs[i])
		}
		cfg, err := r.mgr.NewConfiguration(selectedNodes, &qspec{cfgSize: len(selectedNodes)})
		if err != nil {
			fmt.Printf("Failed to create configuration: %v\n", err)
			return nil
		}
		return cfg
	}
	fmt.Println("Invalid configuration.")
	return nil
}
<commit_msg>examples/storage: added a command to print all nodes<commit_after>package main

import (
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/chzyer/readline"
	"github.com/golang/protobuf/ptypes"
	"github.com/google/shlex"
	"github.com/relab/gorums/examples/storage/proto"
)

var help = `
This interface allows you to run RPCs and quorum calls against the Storage
Servers interactively. Take a look at the files 'client.go' and 'server.go'
for the source code of the RPC handlers and quorum functions.
The following commands can be used:

help 	                        	Show this text
exit 	                        	Exit the program
nodes	                        	Print a list of the available nodes
rpc  	[node index] [operation]	Executes an RPC on the given node.
qc   	[operation]             	Executes a quorum call on all nodes.
cfg  	[config] [opertation]   	Executes a quorum call on a configuration.

The following operations are supported:

read 	[key]        	Read a value
write	[key] [value]	Write a value

Examples:

> rpc 0 write foo bar
The command performs the 'write' RPC on node 0, and sets 'foo' = 'bar'

> qc read foo
The command performs the 'read' quorum call, and returns the value of 'foo'

> cfg 1:3 write foo bar
The command performs the write qourum call on node 1 and 2

> cfg 0,2 write foo 'bar baz'
The command performs the write quorum call on node 0 and 2
`

var prompt = "> "

type repl struct {
	mgr *proto.Manager
	cfg *proto.Configuration
}

func Repl(mgr *proto.Manager, defaultCfg *proto.Configuration) {
	r := &repl{mgr, defaultCfg}
	rl, err := readline.New(prompt)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to create readline instance: %v\n", err)
		os.Exit(1)
	}
	defer rl.Close()

	fmt.Println(help)
	fmt.Print(prompt)

	for {
		l, err := rl.Readline()
		if errors.Is(err, io.EOF) {
			return
		}
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to read line: %v\n", err)
			os.Exit(1)
		}
		args, err := shlex.Split(l)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to split command: %v\n", err)
			os.Exit(1)
		}
		if len(args) < 1 {
			continue
		}

		switch args[0] {
		case "exit":
			return
		case "help":
			fmt.Println(help)
		case "rpc":
			r.rpc(args[1:])
		case "qc":
			r.qc(args[1:])
		case "cfg":
			r.qcCfg(args[1:])
		case "nodes":
			fmt.Println("Nodes: ")
			for i, n := range mgr.Nodes() {
				fmt.Printf("%d: %s\n", i, n.Address())
			}
		default:
			fmt.Printf("Unknown command '%s'. Type 'help' to see available commands.\n", args[0])
		}
	}
}

func (r repl) rpc(args []string) {
	if len(args) < 2 {
		fmt.Println("'rpc' requires a node index and an operation.")
		return
	}

	index, err := strconv.Atoi(args[0])
	if err != nil {
		fmt.Printf("Invalid id '%s'. node index must be numeric.\n", args[0])
		return
	}

	if index < 0 || index >= r.cfg.Size() {
		fmt.Printf("Invalid index. Must be between 0 and %d.\n", r.cfg.Size())
		return
	}

	node := r.cfg.Nodes()[index]

	switch args[1] {
	case "read":
		r.readRPC(args[2:], node)
	case "write":
		r.writeRPC(args[2:], node)
	}
}

func (r repl) qc(args []string) {
	if len(args) < 1 {
		fmt.Println("'qc' requires an operation.")
		return
	}

	switch args[0] {
	case "read":
		r.readQC(args[1:], r.cfg)
	case "write":
		r.writeQC(args[1:], r.cfg)
	}
}

func (r repl) qcCfg(args []string) {
	if len(args) < 2 {
		fmt.Println("'cfg' requires a configuration and an operation.")
		return
	}
	cfg := r.parseConfiguration(args[0])
	if cfg == nil {
		return
	}
	switch args[1] {
	case "read":
		r.readQC(args[2:], cfg)
	case "write":
		r.writeQC(args[2:], cfg)
	}
}

func (_ repl) readRPC(args []string, node *proto.Node) {
	if len(args) < 1 {
		fmt.Println("Read requires a key to read.")
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	resp, err := node.ReadRPC(ctx, &proto.ReadRequest{Key: args[0]})
	cancel()
	if err != nil {
		fmt.Printf("Read RPC finished with error: %v\n", err)
		return
	}
	if !resp.GetOK() {
		fmt.Printf("%s was not found\n", args[0])
		return
	}
	fmt.Printf("%s = %s\n", args[0], resp.GetValue())
}

func (_ repl) writeRPC(args []string, node *proto.Node) {
	if len(args) < 2 {
		fmt.Println("Write requires a key and a value to write.")
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	resp, err := node.WriteRPC(ctx, &proto.WriteRequest{Key: args[0], Value: args[1], Time: ptypes.TimestampNow()})
	cancel()
	if err != nil {
		fmt.Printf("Write RPC finished with error: %v\n", err)
		return
	}
	if !resp.GetNew() {
		fmt.Printf("Failed to update %s: timestamp too old.\n", args[0])
		return
	}
	fmt.Println("Write OK")
}

func (_ repl) readQC(args []string, cfg *proto.Configuration) {
	if len(args) < 1 {
		fmt.Println("Read requires a key to read.")
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	resp, err := cfg.ReadQC(ctx, &proto.ReadRequest{Key: args[0]})
	cancel()
	if err != nil {
		fmt.Printf("Read RPC finished with error: %v\n", err)
		return
	}
	if !resp.GetOK() {
		fmt.Printf("%s was not found\n", args[0])
		return
	}
	fmt.Printf("%s = %s\n", args[0], resp.GetValue())
}

func (_ repl) writeQC(args []string, cfg *proto.Configuration) {
	if len(args) < 2 {
		fmt.Println("Write requires a key and a value to write.")
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	resp, err := cfg.WriteQC(ctx, &proto.WriteRequest{Key: args[0], Value: args[1], Time: ptypes.TimestampNow()})
	cancel()
	if err != nil {
		fmt.Printf("Write RPC finished with error: %v\n", err)
		return
	}
	if !resp.GetNew() {
		fmt.Printf("Failed to update %s: timestamp too old.\n", args[0])
		return
	}
	fmt.Println("Write OK")
}

func (r repl) parseConfiguration(cfgStr string) (cfg *proto.Configuration) {
	// configuration using range syntax
	if i := strings.Index(cfgStr, ":"); i > -1 {
		var start, stop int
		var err error
		numNodes, _ := r.mgr.Size()
		if i == 0 {
			start = 0
		} else {
			start, err = strconv.Atoi(cfgStr[:i])
			if err != nil {
				fmt.Printf("Failed to parse configuration: %v\n", err)
				return nil
			}
		}
		if i == len(cfgStr)-1 {
			stop = numNodes
		} else {
			stop, err = strconv.Atoi(cfgStr[i+1:])
			if err != nil {
				fmt.Printf("Failed to parse configuration: %v\n", err)
				return nil
			}
		}
		if start >= stop || start < 0 || stop >= numNodes {
			fmt.Println("Invalid configuration.")
			return nil
		}
		cfg, err = r.mgr.NewConfiguration(r.mgr.NodeIDs()[start:stop], &qspec{cfgSize: stop - start})
		if err != nil {
			fmt.Printf("Failed to create configuration: %v\n", err)
			return nil
		}
		return cfg
	}
	// configuration using list of indices
	if indices := strings.Split(cfgStr, ","); len(indices) > 0 {
		selectedNodes := make([]uint32, 0, len(indices))
		nodeIDs := r.mgr.NodeIDs()
		for _, index := range indices {
			i, err := strconv.Atoi(index)
			if err != nil {
				fmt.Printf("Failed to parse configuration: %v\n", err)
				return nil
			}
			if i < 0 || i >= len(nodeIDs) {
				fmt.Println("Invalid configuration.")
				return nil
			}
			selectedNodes = append(selectedNodes, nodeIDs[i])
		}
		cfg, err := r.mgr.NewConfiguration(selectedNodes, &qspec{cfgSize: len(selectedNodes)})
		if err != nil {
			fmt.Printf("Failed to create configuration: %v\n", err)
			return nil
		}
		return cfg
	}
	fmt.Println("Invalid configuration.")
	return nil
}
<|endoftext|>


<commit_before>package httpd

import (
	"fmt"
	"io"
	"net"
	"net/http"

	"github.com/Symantec/Dominator/lib/log"
)

type HtmlWriter interface {
	WriteHtml(writer io.Writer)
}

var htmlWriters []HtmlWriter

func StartServer(portNum uint, logger log.DebugLogger) error {
	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", portNum))
	if err != nil {
		return err
	}
	http.HandleFunc("/", statusHandler)
	go http.Serve(listener, nil)
	return nil
}

func AddHtmlWriter(htmlWriter HtmlWriter) {
	htmlWriters = append(htmlWriters, htmlWriter)
}
<commit_msg>Add reverse listener support to subd.<commit_after>package httpd

import (
	"io"
	"net/http"

	"github.com/Symantec/Dominator/lib/log"
	"github.com/Symantec/Dominator/lib/net/reverseconnection"
	"github.com/Symantec/tricorder/go/tricorder"
)

type HtmlWriter interface {
	WriteHtml(writer io.Writer)
}

var htmlWriters []HtmlWriter

func StartServer(portNum uint, logger log.DebugLogger) error {
	listener, err := reverseconnection.Listen("tcp", portNum, logger)
	if err != nil {
		return err
	}
	if err := listener.RequestConnections("Dominator"); err != nil {
		return err
	}
	err = listener.RequestConnections(tricorder.CollectorServiceName)
	if err != nil {
		return err
	}
	http.HandleFunc("/", statusHandler)
	go http.Serve(listener, nil)
	return nil
}

func AddHtmlWriter(htmlWriter HtmlWriter) {
	htmlWriters = append(htmlWriters, htmlWriter)
}
<|endoftext|>


<commit_before>package ovn

import (
	"fmt"

	util "github.com/openvswitch/ovn-kubernetes/go-controller/pkg/util"
	"github.com/sirupsen/logrus"
	kapi "k8s.io/api/core/v1"
)

// AddEndpoints adds endpoints and creates corresponding resources in OVN
func (ovn *Controller) AddEndpoints(ep *kapi.Endpoints) error {
	// get service
	// TODO: cache the service
	svc, err := ovn.kube.GetService(ep.Namespace, ep.Name)
	if err != nil {
		// This is not necessarily an error. For e.g when there are endpoints
		// without a corresponding service.
		logrus.Debugf("no service found for endpoint %s in namespace %s",
			ep.Name, ep.Namespace)
		return nil
	}
	tcpPortMap := make(map[int32]([]string))
	udpPortMap := make(map[int32]([]string))
	for _, s := range ep.Subsets {
		for _, ip := range s.Addresses {
			for _, port := range s.Ports {
				var ips []string
				var portMap map[int32]([]string)
				var ok bool
				if port.Protocol == kapi.ProtocolUDP {
					portMap = udpPortMap
				} else if port.Protocol == kapi.ProtocolTCP {
					portMap = tcpPortMap
				}
				if ips, ok = portMap[port.Port]; !ok {
					ips = make([]string, 0)
				}
				ips = append(ips, ip.IP)
				portMap[port.Port] = ips
			}
		}
	}

	logrus.Debugf("Tcp table: %v\nUdp table: %v", tcpPortMap, udpPortMap)

	for targetPort, ips := range tcpPortMap {
		for _, svcPort := range svc.Spec.Ports {
			if svcPort.Protocol == kapi.ProtocolTCP && svcPort.TargetPort.IntVal == targetPort {
				if svc.Spec.Type == kapi.ServiceTypeNodePort && ovn.nodePortEnable {
					logrus.Debugf("Creating Gateways IP for NodePort: %d, %d, %v", svcPort.NodePort, targetPort, ips)
					err = ovn.createGatewaysVIP(string(svcPort.Protocol), svcPort.NodePort, targetPort, ips)
					if err != nil {
						logrus.Errorf("Error in creating Node Port for svc %s, node port: %d - %v\n", svc.Name, svcPort.NodePort, err)
						continue
					}
				}
				if svc.Spec.Type == kapi.ServiceTypeClusterIP || svc.Spec.Type == kapi.ServiceTypeNodePort {
					err = ovn.createLoadBalancerVIP(ovn.getLoadBalancer(svcPort.Protocol), svc.Spec.ClusterIP, svcPort.Port, ips, targetPort)
					if err != nil {
						logrus.Errorf("Error in creating Cluster IP for svc %s, target port: %d - %v\n", svc.Name, targetPort, err)
						return err
					}
				}
			}
		}
	}
	for targetPort, ips := range udpPortMap {
		for _, svcPort := range svc.Spec.Ports {
			if svcPort.Protocol == kapi.ProtocolUDP && svcPort.TargetPort.IntVal == targetPort {
				if svc.Spec.Type == kapi.ServiceTypeNodePort && ovn.nodePortEnable {
					err = ovn.createGatewaysVIP(string(svcPort.Protocol), svcPort.NodePort, targetPort, ips)
					if err != nil {
						logrus.Errorf("Error in creating Node Port for svc %s, node port: %d - %v\n", svc.Name, svcPort.NodePort, err)
						continue
					}
				} else if svc.Spec.Type == kapi.ServiceTypeNodePort || svc.Spec.Type == kapi.ServiceTypeClusterIP {
					err = ovn.createLoadBalancerVIP(ovn.getLoadBalancer(svcPort.Protocol), svc.Spec.ClusterIP, svcPort.Port, ips, targetPort)
					if err != nil {
						logrus.Errorf("Error in creating Cluster IP for svc %s, target port: %d - %v\n", svc.Name, targetPort, err)
						return err
					}
				}
			}
		}
	}
	return nil
}

func (ovn *Controller) deleteEndpoints(ep *kapi.Endpoints) error {
	svc, err := ovn.kube.GetService(ep.Namespace, ep.Name)
	if err != nil {
		// This is not necessarily an error. For e.g when a service is deleted,
		// you will get endpoint delete event and the call to fetch service
		// will fail.
		logrus.Debugf("no service found for endpoint %s in namespace %s",
			ep.Name, ep.Namespace)
		return nil
	}
	for _, svcPort := range svc.Spec.Ports {
		lb := ovn.getLoadBalancer(svcPort.Protocol)
		key := fmt.Sprintf("\"%s:%d\"", svc.Spec.ClusterIP, svcPort.Port)
		_, stderr, err := util.RunOVNNbctlUnix("remove", "load_balancer", lb,
			"vips", key)
		if err != nil {
			logrus.Errorf("Error in deleting endpoints for lb %s, "+
				"stderr: %q (%v)", lb, stderr, err)
		}
	}
	return nil
}
<commit_msg>support named ports for services<commit_after>package ovn

import (
	"fmt"

	util "github.com/openvswitch/ovn-kubernetes/go-controller/pkg/util"
	"github.com/sirupsen/logrus"
	kapi "k8s.io/api/core/v1"
)

type lbEndpoints struct {
	IPs  []string
	Port int32
}

// AddEndpoints adds endpoints and creates corresponding resources in OVN
func (ovn *Controller) AddEndpoints(ep *kapi.Endpoints) error {
	// get service
	// TODO: cache the service
	svc, err := ovn.kube.GetService(ep.Namespace, ep.Name)
	if err != nil {
		// This is not necessarily an error. For e.g when there are endpoints
		// without a corresponding service.
		logrus.Debugf("no service found for endpoint %s in namespace %s",
			ep.Name, ep.Namespace)
		return nil
	}
	tcpPortMap := make(map[string]lbEndpoints)
	udpPortMap := make(map[string]lbEndpoints)
	for _, s := range ep.Subsets {
		for _, ip := range s.Addresses {
			for _, port := range s.Ports {
				var ips []string
				var portMap map[string]lbEndpoints
				if port.Protocol == kapi.ProtocolUDP {
					portMap = udpPortMap
				} else if port.Protocol == kapi.ProtocolTCP {
					portMap = tcpPortMap
				}
				if lbEps, ok := portMap[port.Name]; ok {
					ips = lbEps.IPs
				} else {
					ips = make([]string, 0)
				}
				ips = append(ips, ip.IP)
				portMap[port.Name] = lbEndpoints{IPs: ips, Port: port.Port}
			}
		}
	}

	logrus.Debugf("Tcp table: %v\nUdp table: %v", tcpPortMap, udpPortMap)

	for svcPortName, lbEps := range tcpPortMap {
		ips := lbEps.IPs
		targetPort := lbEps.Port
		for _, svcPort := range svc.Spec.Ports {
			if svcPort.Protocol == kapi.ProtocolTCP && svcPort.Name == svcPortName {
				if svc.Spec.Type == kapi.ServiceTypeNodePort && ovn.nodePortEnable {
					logrus.Debugf("Creating Gateways IP for NodePort: %d, %v", svcPort.NodePort, ips)
					err = ovn.createGatewaysVIP(string(svcPort.Protocol), svcPort.NodePort, targetPort, ips)
					if err != nil {
						logrus.Errorf("Error in creating Node Port for svc %s, node port: %d - %v\n", svc.Name, svcPort.NodePort, err)
						continue
					}
				}
				if svc.Spec.Type == kapi.ServiceTypeClusterIP || svc.Spec.Type == kapi.ServiceTypeNodePort {
					err = ovn.createLoadBalancerVIP(ovn.getLoadBalancer(svcPort.Protocol), svc.Spec.ClusterIP, svcPort.Port, ips, targetPort)
					if err != nil {
						logrus.Errorf("Error in creating Cluster IP for svc %s, target port: %d - %v\n", svc.Name, targetPort, err)
						return err
					}
				}
			}
		}
	}
	for svcPortName, lbEps := range udpPortMap {
		ips := lbEps.IPs
		targetPort := lbEps.Port
		for _, svcPort := range svc.Spec.Ports {
			if svcPort.Protocol == kapi.ProtocolUDP && svcPort.Name == svcPortName {
				if svc.Spec.Type == kapi.ServiceTypeNodePort && ovn.nodePortEnable {
					err = ovn.createGatewaysVIP(string(svcPort.Protocol), svcPort.NodePort, targetPort, ips)
					if err != nil {
						logrus.Errorf("Error in creating Node Port for svc %s, node port: %d - %v\n", svc.Name, svcPort.NodePort, err)
						continue
					}
				} else if svc.Spec.Type == kapi.ServiceTypeNodePort || svc.Spec.Type == kapi.ServiceTypeClusterIP {
					err = ovn.createLoadBalancerVIP(ovn.getLoadBalancer(svcPort.Protocol), svc.Spec.ClusterIP, svcPort.Port, ips, targetPort)
					if err != nil {
						logrus.Errorf("Error in creating Cluster IP for svc %s, target port: %d - %v\n", svc.Name, targetPort, err)
						return err
					}
				}
			}
		}
	}
	return nil
}

func (ovn *Controller) deleteEndpoints(ep *kapi.Endpoints) error {
	svc, err := ovn.kube.GetService(ep.Namespace, ep.Name)
	if err != nil {
		// This is not necessarily an error. For e.g when a service is deleted,
		// you will get endpoint delete event and the call to fetch service
		// will fail.
		logrus.Debugf("no service found for endpoint %s in namespace %s",
			ep.Name, ep.Namespace)
		return nil
	}
	for _, svcPort := range svc.Spec.Ports {
		lb := ovn.getLoadBalancer(svcPort.Protocol)
		key := fmt.Sprintf("\"%s:%d\"", svc.Spec.ClusterIP, svcPort.Port)
		_, stderr, err := util.RunOVNNbctlUnix("remove", "load_balancer", lb,
			"vips", key)
		if err != nil {
			logrus.Errorf("Error in deleting endpoints for lb %s, "+
				"stderr: %q (%v)", lb, stderr, err)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package arn

import (
	"errors"
	"fmt"
	"reflect"
	"strings"

	"github.com/aerogo/aero"
	"github.com/aerogo/api"
	"github.com/aerogo/markdown"
	"github.com/animenotifier/arn/autocorrect"
)

// Force interface implementations
var (
	_ Postable          = (*Post)(nil)
	_ Likeable          = (*Post)(nil)
	_ LikeEventReceiver = (*Post)(nil)
	_ PostParent        = (*Post)(nil)
	_ fmt.Stringer      = (*Post)(nil)
	_ api.Newable       = (*Post)(nil)
	_ api.Editable      = (*Post)(nil)
	_ api.Actionable    = (*Post)(nil)
	_ api.Deletable     = (*Post)(nil)
)

// Actions
func init() {
	API.RegisterActions("Post", []*api.Action{
		// Like post
		LikeAction(),

		// Unlike post
		UnlikeAction(),
	})
}

// Authorize returns an error if the given API POST request is not authorized.
func (post *Post) Authorize(ctx *aero.Context, action string) error {
	if !ctx.HasSession() {
		return errors.New("Neither logged in nor in session")
	}

	if action == "edit" {
		user := GetUserFromContext(ctx)

		if post.CreatedBy != user.ID && user.Role != "admin" {
			return errors.New("Can't edit the posts of other users")
		}
	}

	return nil
}

// Create sets the data for a new post with data we received from the API request.
func (post *Post) Create(ctx *aero.Context) error {
	data, err := ctx.Request().Body().JSONObject()

	if err != nil {
		return err
	}

	user := GetUserFromContext(ctx)

	if user == nil {
		return errors.New("Not logged in")
	}

	post.ID = GenerateID("Post")
	post.Text, _ = data["text"].(string)
	post.CreatedBy = user.ID
	post.ParentID, _ = data["parentId"].(string)
	post.ParentType, _ = data["parentType"].(string)
	post.Likes = []string{}
	post.Created = DateTimeUTC()
	post.Edited = ""

	// Check parent type
	if !DB.HasType(post.ParentType) {
		return errors.New("Invalid parent type: " + post.ParentType)
	}

	// Post-process text
	post.Text = autocorrect.PostText(post.Text)

	if len(post.Text) < 5 {
		return errors.New("Text too short: Should be at least 5 characters")
	}

	// Tags
	tags, _ := data["tags"].([]interface{})
	post.Tags = make([]string, len(tags))

	for i := range post.Tags {
		post.Tags[i] = tags[i].(string)
	}

	// Thread
	parent := post.Parent()

	if parent == nil {
		return errors.New(post.ParentType + " does not exist")
	}

	// Is the parent locked?
	if IsLocked(parent) {
		return errors.New(post.ParentType + " is locked")
	}

	// Append to posts
	parent.AddPost(post.ID)

	// Save the parent thread
	parent.Save()

	// Send notification to the author of the parent post
	go func() {
		notifyUser := parent.Creator()

		// Don't notify the author himself
		if notifyUser.ID == post.CreatedBy {
			return
		}

		message := ""

		if post.ParentType == "Post" {
			message = fmt.Sprintf("%s replied to your comment in \"%s\".", user.Nick, parent.(*Post).Parent().TitleByUser(notifyUser))
		} else {
			message = fmt.Sprintf("%s replied in the %s \"%s\".", user.Nick, strings.ToLower(post.ParentType), parent.TitleByUser(notifyUser))
		}

		notifyUser.SendNotification(&PushNotification{
			Title:   user.Nick + " replied",
			Message: message,
			Icon:    "https:" + user.AvatarLink("large"),
			Link:    post.Link(),
			Type:    NotificationTypeForumReply,
		})
	}()

	// Write log entry
	logEntry := NewEditLogEntry(user.ID, "create", "Post", post.ID, "", "", "")
	logEntry.Save()

	return nil
}

// Edit saves a log entry for the edit.
func (post *Post) Edit(ctx *aero.Context, key string, value reflect.Value, newValue reflect.Value) (bool, error) {
	consumed := false
	user := GetUserFromContext(ctx)

	switch key {
	case "ParentID":
		newParentID := newValue.String()
		newParent, err := GetPost(newParentID)

		if err != nil {
			return false, err
		}

		post.SetParent(newParent)
		consumed = true
	}

	// Write log entry
	logEntry := NewEditLogEntry(user.ID, "edit", "Post", post.ID, key, fmt.Sprint(value.Interface()), fmt.Sprint(newValue.Interface()))
	logEntry.Save()

	return consumed, nil
}

// OnAppend saves a log entry.
func (post *Post) OnAppend(ctx *aero.Context, key string, index int, obj interface{}) {
	user := GetUserFromContext(ctx)
	logEntry := NewEditLogEntry(user.ID, "arrayAppend", "Post", post.ID, fmt.Sprintf("%s[%d]", key, index), "", fmt.Sprint(obj))
	logEntry.Save()
}

// OnRemove saves a log entry.
func (post *Post) OnRemove(ctx *aero.Context, key string, index int, obj interface{}) {
	user := GetUserFromContext(ctx)
	logEntry := NewEditLogEntry(user.ID, "arrayRemove", "Post", post.ID, fmt.Sprintf("%s[%d]", key, index), fmt.Sprint(obj), "")
	logEntry.Save()
}

// DeleteInContext deletes the post in the given context.
func (post *Post) DeleteInContext(ctx *aero.Context) error {
	user := GetUserFromContext(ctx)

	// Write log entry
	logEntry := NewEditLogEntry(user.ID, "delete", "Post", post.ID, "", fmt.Sprint(post), "")
	logEntry.Save()

	return post.Delete()
}

// Delete deletes the post from the database.
func (post *Post) Delete() error {
	// Remove child posts first
	for _, post := range post.Posts() {
		post.Delete()
	}

	parent := post.Parent()

	if parent == nil {
		return fmt.Errorf("Invalid %s parent ID: %s", post.ParentType, post.ParentID)
	}

	// Remove the reference of the post in the thread that contains it
	if !parent.RemovePost(post.ID) {
		return fmt.Errorf("This post does not exist in the %s", strings.ToLower(post.ParentType))
	}

	parent.Save()
	DB.Delete("Post", post.ID)
	return nil
}

// AfterEdit updates the date it has been edited.
func (post *Post) AfterEdit(ctx *aero.Context) error {
	post.Edited = DateTimeUTC()
	post.html = markdown.Render(post.Text)
	return nil
}

// Save saves the post object in the database.
func (post *Post) Save() {
	DB.Set("Post", post.ID, post)
}
<commit_msg>Allow thread IDs in parent editing<commit_after>package arn

import (
	"errors"
	"fmt"
	"reflect"
	"strings"

	"github.com/aerogo/aero"
	"github.com/aerogo/api"
	"github.com/aerogo/markdown"
	"github.com/animenotifier/arn/autocorrect"
)

// Force interface implementations
var (
	_ Postable          = (*Post)(nil)
	_ Likeable          = (*Post)(nil)
	_ LikeEventReceiver = (*Post)(nil)
	_ PostParent        = (*Post)(nil)
	_ fmt.Stringer      = (*Post)(nil)
	_ api.Newable       = (*Post)(nil)
	_ api.Editable      = (*Post)(nil)
	_ api.Actionable    = (*Post)(nil)
	_ api.Deletable     = (*Post)(nil)
)

// Actions
func init() {
	API.RegisterActions("Post", []*api.Action{
		// Like post
		LikeAction(),

		// Unlike post
		UnlikeAction(),
	})
}

// Authorize returns an error if the given API POST request is not authorized.
func (post *Post) Authorize(ctx *aero.Context, action string) error {
	if !ctx.HasSession() {
		return errors.New("Neither logged in nor in session")
	}

	if action == "edit" {
		user := GetUserFromContext(ctx)

		if post.CreatedBy != user.ID && user.Role != "admin" {
			return errors.New("Can't edit the posts of other users")
		}
	}

	return nil
}

// Create sets the data for a new post with data we received from the API request.
func (post *Post) Create(ctx *aero.Context) error {
	data, err := ctx.Request().Body().JSONObject()

	if err != nil {
		return err
	}

	user := GetUserFromContext(ctx)

	if user == nil {
		return errors.New("Not logged in")
	}

	post.ID = GenerateID("Post")
	post.Text, _ = data["text"].(string)
	post.CreatedBy = user.ID
	post.ParentID, _ = data["parentId"].(string)
	post.ParentType, _ = data["parentType"].(string)
	post.Likes = []string{}
	post.Created = DateTimeUTC()
	post.Edited = ""

	// Check parent type
	if !DB.HasType(post.ParentType) {
		return errors.New("Invalid parent type: " + post.ParentType)
	}

	// Post-process text
	post.Text = autocorrect.PostText(post.Text)

	if len(post.Text) < 5 {
		return errors.New("Text too short: Should be at least 5 characters")
	}

	// Tags
	tags, _ := data["tags"].([]interface{})
	post.Tags = make([]string, len(tags))

	for i := range post.Tags {
		post.Tags[i] = tags[i].(string)
	}

	// Thread
	parent := post.Parent()

	if parent == nil {
		return errors.New(post.ParentType + " does not exist")
	}

	// Is the parent locked?
	if IsLocked(parent) {
		return errors.New(post.ParentType + " is locked")
	}

	// Append to posts
	parent.AddPost(post.ID)

	// Save the parent thread
	parent.Save()

	// Send notification to the author of the parent post
	go func() {
		notifyUser := parent.Creator()

		// Don't notify the author himself
		if notifyUser.ID == post.CreatedBy {
			return
		}

		message := ""

		if post.ParentType == "Post" {
			message = fmt.Sprintf("%s replied to your comment in \"%s\".", user.Nick, parent.(*Post).Parent().TitleByUser(notifyUser))
		} else {
			message = fmt.Sprintf("%s replied in the %s \"%s\".", user.Nick, strings.ToLower(post.ParentType), parent.TitleByUser(notifyUser))
		}

		notifyUser.SendNotification(&PushNotification{
			Title:   user.Nick + " replied",
			Message: message,
			Icon:    "https:" + user.AvatarLink("large"),
			Link:    post.Link(),
			Type:    NotificationTypeForumReply,
		})
	}()

	// Write log entry
	logEntry := NewEditLogEntry(user.ID, "create", "Post", post.ID, "", "", "")
	logEntry.Save()

	return nil
}

// Edit saves a log entry for the edit.
func (post *Post) Edit(ctx *aero.Context, key string, value reflect.Value, newValue reflect.Value) (bool, error) {
	consumed := false
	user := GetUserFromContext(ctx)

	switch key {
	case "ParentID":
		var newParent PostParent
		newParentID := newValue.String()
		newParent, err := GetPost(newParentID)

		if err != nil {
			newParent, err = GetThread(newParentID)

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

		post.SetParent(newParent)
		consumed = true
	}

	// Write log entry
	logEntry := NewEditLogEntry(user.ID, "edit", "Post", post.ID, key, fmt.Sprint(value.Interface()), fmt.Sprint(newValue.Interface()))
	logEntry.Save()

	return consumed, nil
}

// OnAppend saves a log entry.
func (post *Post) OnAppend(ctx *aero.Context, key string, index int, obj interface{}) {
	user := GetUserFromContext(ctx)
	logEntry := NewEditLogEntry(user.ID, "arrayAppend", "Post", post.ID, fmt.Sprintf("%s[%d]", key, index), "", fmt.Sprint(obj))
	logEntry.Save()
}

// OnRemove saves a log entry.
func (post *Post) OnRemove(ctx *aero.Context, key string, index int, obj interface{}) {
	user := GetUserFromContext(ctx)
	logEntry := NewEditLogEntry(user.ID, "arrayRemove", "Post", post.ID, fmt.Sprintf("%s[%d]", key, index), fmt.Sprint(obj), "")
	logEntry.Save()
}

// DeleteInContext deletes the post in the given context.
func (post *Post) DeleteInContext(ctx *aero.Context) error {
	user := GetUserFromContext(ctx)

	// Write log entry
	logEntry := NewEditLogEntry(user.ID, "delete", "Post", post.ID, "", fmt.Sprint(post), "")
	logEntry.Save()

	return post.Delete()
}

// Delete deletes the post from the database.
func (post *Post) Delete() error {
	// Remove child posts first
	for _, post := range post.Posts() {
		post.Delete()
	}

	parent := post.Parent()

	if parent == nil {
		return fmt.Errorf("Invalid %s parent ID: %s", post.ParentType, post.ParentID)
	}

	// Remove the reference of the post in the thread that contains it
	if !parent.RemovePost(post.ID) {
		return fmt.Errorf("This post does not exist in the %s", strings.ToLower(post.ParentType))
	}

	parent.Save()
	DB.Delete("Post", post.ID)
	return nil
}

// AfterEdit updates the date it has been edited.
func (post *Post) AfterEdit(ctx *aero.Context) error {
	post.Edited = DateTimeUTC()
	post.html = markdown.Render(post.Text)
	return nil
}

// Save saves the post object in the database.
func (post *Post) Save() {
	DB.Set("Post", post.ID, post)
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Google Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package sqlparser

import (
	"bytes"
	"io"
	"strings"
	"testing"
)

// TestParseNextValid concatenates all the valid SQL test cases and check it can read
// them as one long string.
func TestParseNextValid(t *testing.T) {
	var sql bytes.Buffer
	for _, tcase := range validSQL {
		sql.WriteString(strings.TrimSuffix(tcase.input, ";"))
		sql.WriteRune(';')
	}

	tokens := NewStringTokenizer(sql.String())
	for i, tcase := range validSQL {
		input := tcase.input + ";"
		want := tcase.output
		if want == "" {
			want = tcase.input
		}

		tree, err := ParseNext(tokens)
		if err != nil {
			t.Fatalf("[%d] ParseNext(%q) err: %q, want nil", i, input, err)
			continue
		}

		if got := String(tree); got != want {
			t.Fatalf("[%d] ParseNext(%q) = %q, want %q", i, input, got, want)
		}
	}

	// Read once more and it should be EOF.
	if tree, err := ParseNext(tokens); err != io.EOF {
		t.Errorf("ParseNext(tokens) = (%q, %v) want io.EOF", String(tree), err)
	}
}

// TestParseNextErrors tests all the error cases, and ensures a valid
// SQL statement can be passed afterwards.
func TestParseNextErrors(t *testing.T) {
	for _, tcase := range invalidSQL {
		if tcase.skipMulti {
			// Skip tests which leave unclosed strings, or comments.
			continue
		}

		sql := tcase.input + " ; select 1 from t"
		tokens := NewStringTokenizer(sql)

		// The first statement should be an error
		_, err := ParseNext(tokens)
		if err == nil || err.Error() != tcase.output {
			t.Fatalf("[0] ParseNext(%q) err: %q, want %q", sql, err, tcase.output)
			continue
		}

		// The second should be valid
		tree, err := ParseNext(tokens)
		if err != nil {
			t.Fatalf("[1] ParseNext(%q) err: %q, want nil", sql, err)
			continue
		}

		want := "select 1 from t"
		if got := String(tree); got != want {
			t.Fatalf("[1] ParseNext(%q) = %q, want %q", sql, got, want)
		}

		// Read once more and it should be EOF.
		if tree, err := ParseNext(tokens); err != io.EOF {
			t.Errorf("ParseNext(tokens) = (%q, %v) want io.EOF", String(tree), err)
		}
	}
}

// TestParseNextEdgeCases tests various ParseNext edge cases.
func TestParseNextEdgeCases(t *testing.T) {
	tests := []struct {
		name  string
		input string
		want  []string
	}{{
		name:  "Trailing ;",
		input: "select 1 from a; update a set b = 2;",
		want:  []string{"select 1 from a", "update a set b = 2"},
	}, {
		name:  "No trailing ;",
		input: "select 1 from a; update a set b = 2",
		want:  []string{"select 1 from a", "update a set b = 2"},
	}, {
		name:  "Trailing whitespace",
		input: "select 1 from a; update a set b = 2    ",
		want:  []string{"select 1 from a", "update a set b = 2"},
	}, {
		name:  "Trailing whitespace and ;",
		input: "select 1 from a; update a set b = 2   ;   ",
		want:  []string{"select 1 from a", "update a set b = 2"},
	}, {
		name:  "Handle ForceEOF statements",
		input: "set character set utf8; select 1 from a",
		want:  []string{"set ", "select 1 from a"},
	}, {
		name:  "Semicolin inside a string",
		input: "set character set ';'; select 1 from a",
		want:  []string{"set ", "select 1 from a"},
	}, {
		name:  "Partial DDL",
		input: "create table a; select 1 from a",
		want:  []string{"create table a", "select 1 from a"},
	}, {
		name:  "Partial DDL",
		input: "create table a ignore me this is garbage; select 1 from a",
		want:  []string{"create table a", "select 1 from a"},
	}}

	for _, test := range tests {
		tokens := NewStringTokenizer(test.input)

		for i, want := range test.want {
			tree, err := ParseNext(tokens)
			if err != nil {
				t.Fatalf("[%d] ParseNext(%q) err = %q, want nil", i, test.input, err)
				continue
			}

			if got := String(tree); got != want {
				t.Fatalf("[%d] ParseNext(%q) = %q, want %q", i, test.input, got, want)
			}
		}

		// Read once more and it should be EOF.
		if tree, err := ParseNext(tokens); err != io.EOF {
			t.Errorf("ParseNext(%q) = (%q, %v) want io.EOF", test.input, String(tree), err)
		}

		// And again, once more should be EOF.
		if tree, err := ParseNext(tokens); err != io.EOF {
			t.Errorf("ParseNext(%q) = (%q, %v) want io.EOF", test.input, String(tree), err)
		}
	}
}
<commit_msg>Fixed a typo!<commit_after>/*
Copyright 2017 Google Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package sqlparser

import (
	"bytes"
	"io"
	"strings"
	"testing"
)

// TestParseNextValid concatenates all the valid SQL test cases and check it can read
// them as one long string.
func TestParseNextValid(t *testing.T) {
	var sql bytes.Buffer
	for _, tcase := range validSQL {
		sql.WriteString(strings.TrimSuffix(tcase.input, ";"))
		sql.WriteRune(';')
	}

	tokens := NewStringTokenizer(sql.String())
	for i, tcase := range validSQL {
		input := tcase.input + ";"
		want := tcase.output
		if want == "" {
			want = tcase.input
		}

		tree, err := ParseNext(tokens)
		if err != nil {
			t.Fatalf("[%d] ParseNext(%q) err: %q, want nil", i, input, err)
			continue
		}

		if got := String(tree); got != want {
			t.Fatalf("[%d] ParseNext(%q) = %q, want %q", i, input, got, want)
		}
	}

	// Read once more and it should be EOF.
	if tree, err := ParseNext(tokens); err != io.EOF {
		t.Errorf("ParseNext(tokens) = (%q, %v) want io.EOF", String(tree), err)
	}
}

// TestParseNextErrors tests all the error cases, and ensures a valid
// SQL statement can be passed afterwards.
func TestParseNextErrors(t *testing.T) {
	for _, tcase := range invalidSQL {
		if tcase.excludeMulti {
			// Skip tests which leave unclosed strings, or comments.
			continue
		}

		sql := tcase.input + " ; select 1 from t"
		tokens := NewStringTokenizer(sql)

		// The first statement should be an error
		_, err := ParseNext(tokens)
		if err == nil || err.Error() != tcase.output {
			t.Fatalf("[0] ParseNext(%q) err: %q, want %q", sql, err, tcase.output)
			continue
		}

		// The second should be valid
		tree, err := ParseNext(tokens)
		if err != nil {
			t.Fatalf("[1] ParseNext(%q) err: %q, want nil", sql, err)
			continue
		}

		want := "select 1 from t"
		if got := String(tree); got != want {
			t.Fatalf("[1] ParseNext(%q) = %q, want %q", sql, got, want)
		}

		// Read once more and it should be EOF.
		if tree, err := ParseNext(tokens); err != io.EOF {
			t.Errorf("ParseNext(tokens) = (%q, %v) want io.EOF", String(tree), err)
		}
	}
}

// TestParseNextEdgeCases tests various ParseNext edge cases.
func TestParseNextEdgeCases(t *testing.T) {
	tests := []struct {
		name  string
		input string
		want  []string
	}{{
		name:  "Trailing ;",
		input: "select 1 from a; update a set b = 2;",
		want:  []string{"select 1 from a", "update a set b = 2"},
	}, {
		name:  "No trailing ;",
		input: "select 1 from a; update a set b = 2",
		want:  []string{"select 1 from a", "update a set b = 2"},
	}, {
		name:  "Trailing whitespace",
		input: "select 1 from a; update a set b = 2    ",
		want:  []string{"select 1 from a", "update a set b = 2"},
	}, {
		name:  "Trailing whitespace and ;",
		input: "select 1 from a; update a set b = 2   ;   ",
		want:  []string{"select 1 from a", "update a set b = 2"},
	}, {
		name:  "Handle ForceEOF statements",
		input: "set character set utf8; select 1 from a",
		want:  []string{"set ", "select 1 from a"},
	}, {
		name:  "Semicolin inside a string",
		input: "set character set ';'; select 1 from a",
		want:  []string{"set ", "select 1 from a"},
	}, {
		name:  "Partial DDL",
		input: "create table a; select 1 from a",
		want:  []string{"create table a", "select 1 from a"},
	}, {
		name:  "Partial DDL",
		input: "create table a ignore me this is garbage; select 1 from a",
		want:  []string{"create table a", "select 1 from a"},
	}}

	for _, test := range tests {
		tokens := NewStringTokenizer(test.input)

		for i, want := range test.want {
			tree, err := ParseNext(tokens)
			if err != nil {
				t.Fatalf("[%d] ParseNext(%q) err = %q, want nil", i, test.input, err)
				continue
			}

			if got := String(tree); got != want {
				t.Fatalf("[%d] ParseNext(%q) = %q, want %q", i, test.input, got, want)
			}
		}

		// Read once more and it should be EOF.
		if tree, err := ParseNext(tokens); err != io.EOF {
			t.Errorf("ParseNext(%q) = (%q, %v) want io.EOF", test.input, String(tree), err)
		}

		// And again, once more should be EOF.
		if tree, err := ParseNext(tokens); err != io.EOF {
			t.Errorf("ParseNext(%q) = (%q, %v) want io.EOF", test.input, String(tree), err)
		}
	}
}
<|endoftext|>