<commit_before>/*
Copyright 2019 The Kubernetes Authors All rights reserved.

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

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

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

package e2e

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path"
	"regexp"
	"sort"
	"strings"
	"testing"

	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus/testutil/promlint"
	dto "github.com/prometheus/client_model/go"

	ksmFramework "k8s.io/kube-state-metrics/v2/tests/e2e/framework"
)

var framework *ksmFramework.Framework

func TestMain(m *testing.M) {
	ksmHTTPMetricsURL := flag.String(
		"ksm-http-metrics-url",
		"",
		"url to access the kube-state-metrics service",
	)
	ksmTelemetryURL := flag.String(
		"ksm-telemetry-url",
		"",
		"url to access the kube-state-metrics telemetry endpoint",
	)
	flag.Parse()

	var (
		err      error
		exitCode int
	)

	if framework, err = ksmFramework.New(*ksmHTTPMetricsURL, *ksmTelemetryURL); err != nil {
		log.Fatalf("failed to setup framework: %v\n", err)
	}

	exitCode = m.Run()

	os.Exit(exitCode)
}

func TestIsHealthz(t *testing.T) {
	ok, err := framework.KsmClient.IsHealthz()
	if err != nil {
		t.Fatalf("kube-state-metrics healthz check failed: %v", err)
	}

	if ok == false {
		t.Fatal("kube-state-metrics is unhealthy")
	}
}

func TestLintMetrics(t *testing.T) {
	buf := &bytes.Buffer{}

	err := framework.KsmClient.Metrics(buf)
	if err != nil {
		t.Fatalf("failed to get metrics from kube-state-metrics: %v", err)
	}

	l := promlint.New(buf)
	problems, err := l.Lint()
	if err != nil {
		t.Fatalf("failed to lint: %v", err)
	}

	if len(problems) != 0 {
		t.Fatalf("the problems encountered in Lint are: %v", problems)
	}
}

func TestDocumentation(t *testing.T) {
	labelsDocumentation, err := getLabelsDocumentation()
	if err != nil {
		t.Fatal("Cannot get labels documentation", err)
	}

	metricFamilies, err := framework.ParseMetrics(framework.KsmClient.Metrics)
	if err != nil {
		t.Fatal("Failed to get or decode metrics", err)
	}

	for _, metricFamily := range metricFamilies {
		metric := metricFamily.GetName()

		acceptedLabelNames, ok := labelsDocumentation[metric]
		if !ok {
			t.Errorf("Metric %s not found in documentation.", metric)
			continue
		}
		for _, m := range metricFamily.Metric {
			for _, l := range m.Label {
				labelName := l.GetName()
				labelNameMatched := false
				for _, labelPattern := range acceptedLabelNames {
					re, err := regexp.Compile(labelPattern)
					if err != nil {
						t.Errorf("Cannot compile pattern %s: %w", labelPattern, err)
						continue
					}
					if re.MatchString(labelName) {
						labelNameMatched = true
						break
					}
				}
				if !labelNameMatched {
					t.Errorf("Label %s not found in documentation. Documented labels for metric %s are: %s",
						labelName, metric, strings.Join(acceptedLabelNames, ", "))
				}
			}
		}
	}
}

// getLabelsDocumentation is a helper function that gets metric mabels documentation.
// It returns a map where keys are metric names, and values are slices of label names,
// and an error in case of failure.
// By convention, UPPER_CASE parts in label names denotes wilcard patterns, used for dynamic labels.
func getLabelsDocumentation() (map[string][]string, error) {
	documentedMetrics := map[string][]string{}

	docPath := "../../docs/"
	docFiles, err := os.ReadDir(docPath)
	if err != nil {
		return nil, errors.Wrap(err, "failed to read documentation directory")
	}

	// Match file names such as daemonset-metrics.md
	fileRe := regexp.MustCompile(`^([a-z]*)-metrics.md$`)
	// Match doc lines such as | kube_node_created | Gauge | `node`=<node-address>| STABLE |
	lineRe := regexp.MustCompile(`^\| *(kube_[a-z_]+) *\| *[a-zA-Z]+ *\|(.*)\| *[A-Z]+`)
	// Match label names in label documentation
	labelsRe := regexp.MustCompile("`([a-zA-Z_][a-zA-Z0-9_]*)`")
	// Match wildcard patterns for dynamic labels such as label_CRONJOB_LABEL
	patternRe := regexp.MustCompile(`_[A-Z_]+`)

	for _, file := range docFiles {
		if file.IsDir() || !fileRe.MatchString(file.Name()) {
			continue
		}

		filePath := path.Join(docPath, file.Name())
		f, err := os.Open(filePath)
		if err != nil {
			return nil, errors.Wrapf(err, "cannot read file %s", filePath)
		}
		scanner := bufio.NewScanner(f)
		for scanner.Scan() {
			params := lineRe.FindStringSubmatch(scanner.Text())
			if len(params) != 3 {
				continue
			}
			metric := params[1]
			labelsDoc := params[2]

			labels := labelsRe.FindAllStringSubmatch(labelsDoc, -1)
			labelPatterns := make([]string, len(labels))
			for i, l := range labels {
				if len(l) <= 1 {
					return nil, errors.Errorf("Label documentation %s did not match regex", labelsDoc)
				}
				labelPatterns[i] = patternRe.ReplaceAllString(l[1], "_.*")
			}

			documentedMetrics[metric] = labelPatterns
		}
	}
	return documentedMetrics, nil
}

func TestKubeStateMetricsErrorMetrics(t *testing.T) {
	metricFamilies, err := framework.ParseMetrics(framework.KsmClient.TelemetryMetrics)
	if err != nil {
		t.Fatal("Failed to get or decode telemetry metrics", err)
	}

	// This map's keys are the metrics expected in kube-state-metrics telemetry.
	// Its values are booleans, set to true when the metric is found.
	foundMetricFamily := map[string]bool{
		"kube_state_metrics_list_total":  false,
		"kube_state_metrics_watch_total": false,
	}

	for _, metricFamily := range metricFamilies {
		name := metricFamily.GetName()
		if _, expectedMetric := foundMetricFamily[name]; expectedMetric {
			foundMetricFamily[name] = true

			for _, m := range metricFamily.Metric {
				if hasLabelError(m) && m.GetCounter().GetValue() > 0 {
					t.Errorf("Metric %s in telemetry shows a list/watch error", prettyPrintCounter(name, m))
				}
			}
		}
	}

	for metricFamily, found := range foundMetricFamily {
		if !found {
			t.Errorf("Metric family %s was not found in telemetry metrics", metricFamily)
		}
	}
}

func hasLabelError(metric *dto.Metric) bool {
	for _, l := range metric.Label {
		if l.GetName() == "result" && l.GetValue() == "error" {
			return true
		}
	}
	return false
}

func prettyPrintCounter(name string, metric *dto.Metric) string {
	labelStrings := []string{}
	for _, l := range metric.Label {
		labelStrings = append(labelStrings, fmt.Sprintf(`%s="%s"`, l.GetName(), l.GetValue()))
	}
	return fmt.Sprintf("%s{%s} %d", name, strings.Join(labelStrings, ","), int(metric.GetCounter().GetValue()))
}

func TestDefaultCollectorMetricsAvailable(t *testing.T) {
	buf := &bytes.Buffer{}

	err := framework.KsmClient.Metrics(buf)
	if err != nil {
		t.Fatalf("failed to get metrics from kube-state-metrics: %v", err)
	}

	resources := map[string]struct{}{}
	files, err := ioutil.ReadDir("../../internal/store/")
	if err != nil {
		t.Fatalf("failed to read dir to get all resouces name: %v", err)
	}

	re := regexp.MustCompile(`^([a-z]*).go$`)
	for _, file := range files {
		params := re.FindStringSubmatch(file.Name())
		if len(params) != 2 {
			continue
		}
		if params[1] == "builder" || params[1] == "utils" || params[1] == "testutils" || params[1] == "verticalpodautoscaler" {
			continue
		}
		resources[params[1]] = struct{}{}
	}

	re = regexp.MustCompile(`^kube_([a-z]*)_`)
	scanner := bufio.NewScanner(buf)
	for scanner.Scan() {
		params := re.FindStringSubmatch(scanner.Text())
		if len(params) != 2 {
			continue
		}
		delete(resources, params[1])
	}

	err = scanner.Err()
	if err != nil {
		t.Fatalf("failed to scan metrics: %v", err)
	}

	if len(resources) != 0 {
		s := []string{}
		for k := range resources {
			s = append(s, k)
		}
		sort.Strings(s)
		t.Fatalf("failed to find metrics of resources: %s", strings.Join(s, ", "))
	}
}
<commit_msg>e2e TestDefaultCollectorMetricsAvailable improvements:<commit_after>/*
Copyright 2019 The Kubernetes Authors All rights reserved.

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

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

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

package e2e

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"log"
	"os"
	"path"
	"regexp"
	"sort"
	"strings"
	"testing"

	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus/testutil/promlint"
	dto "github.com/prometheus/client_model/go"

	ksmFramework "k8s.io/kube-state-metrics/v2/tests/e2e/framework"
)

var framework *ksmFramework.Framework

func TestMain(m *testing.M) {
	ksmHTTPMetricsURL := flag.String(
		"ksm-http-metrics-url",
		"",
		"url to access the kube-state-metrics service",
	)
	ksmTelemetryURL := flag.String(
		"ksm-telemetry-url",
		"",
		"url to access the kube-state-metrics telemetry endpoint",
	)
	flag.Parse()

	var (
		err      error
		exitCode int
	)

	if framework, err = ksmFramework.New(*ksmHTTPMetricsURL, *ksmTelemetryURL); err != nil {
		log.Fatalf("failed to setup framework: %v\n", err)
	}

	exitCode = m.Run()

	os.Exit(exitCode)
}

func TestIsHealthz(t *testing.T) {
	ok, err := framework.KsmClient.IsHealthz()
	if err != nil {
		t.Fatalf("kube-state-metrics healthz check failed: %v", err)
	}

	if ok == false {
		t.Fatal("kube-state-metrics is unhealthy")
	}
}

func TestLintMetrics(t *testing.T) {
	buf := &bytes.Buffer{}

	err := framework.KsmClient.Metrics(buf)
	if err != nil {
		t.Fatalf("failed to get metrics from kube-state-metrics: %v", err)
	}

	l := promlint.New(buf)
	problems, err := l.Lint()
	if err != nil {
		t.Fatalf("failed to lint: %v", err)
	}

	if len(problems) != 0 {
		t.Fatalf("the problems encountered in Lint are: %v", problems)
	}
}

func TestDocumentation(t *testing.T) {
	labelsDocumentation, err := getLabelsDocumentation()
	if err != nil {
		t.Fatal("Cannot get labels documentation", err)
	}

	metricFamilies, err := framework.ParseMetrics(framework.KsmClient.Metrics)
	if err != nil {
		t.Fatal("Failed to get or decode metrics", err)
	}

	for _, metricFamily := range metricFamilies {
		metric := metricFamily.GetName()

		acceptedLabelNames, ok := labelsDocumentation[metric]
		if !ok {
			t.Errorf("Metric %s not found in documentation.", metric)
			continue
		}
		for _, m := range metricFamily.Metric {
			for _, l := range m.Label {
				labelName := l.GetName()
				labelNameMatched := false
				for _, labelPattern := range acceptedLabelNames {
					re, err := regexp.Compile(labelPattern)
					if err != nil {
						t.Errorf("Cannot compile pattern %s: %w", labelPattern, err)
						continue
					}
					if re.MatchString(labelName) {
						labelNameMatched = true
						break
					}
				}
				if !labelNameMatched {
					t.Errorf("Label %s not found in documentation. Documented labels for metric %s are: %s",
						labelName, metric, strings.Join(acceptedLabelNames, ", "))
				}
			}
		}
	}
}

// getLabelsDocumentation is a helper function that gets metric mabels documentation.
// It returns a map where keys are metric names, and values are slices of label names,
// and an error in case of failure.
// By convention, UPPER_CASE parts in label names denotes wilcard patterns, used for dynamic labels.
func getLabelsDocumentation() (map[string][]string, error) {
	documentedMetrics := map[string][]string{}

	docPath := "../../docs/"
	docFiles, err := os.ReadDir(docPath)
	if err != nil {
		return nil, errors.Wrap(err, "failed to read documentation directory")
	}

	// Match file names such as daemonset-metrics.md
	fileRe := regexp.MustCompile(`^([a-z]*)-metrics.md$`)
	// Match doc lines such as | kube_node_created | Gauge | `node`=<node-address>| STABLE |
	lineRe := regexp.MustCompile(`^\| *(kube_[a-z_]+) *\| *[a-zA-Z]+ *\|(.*)\| *[A-Z]+`)
	// Match label names in label documentation
	labelsRe := regexp.MustCompile("`([a-zA-Z_][a-zA-Z0-9_]*)`")
	// Match wildcard patterns for dynamic labels such as label_CRONJOB_LABEL
	patternRe := regexp.MustCompile(`_[A-Z_]+`)

	for _, file := range docFiles {
		if file.IsDir() || !fileRe.MatchString(file.Name()) {
			continue
		}

		filePath := path.Join(docPath, file.Name())
		f, err := os.Open(filePath)
		if err != nil {
			return nil, errors.Wrapf(err, "cannot read file %s", filePath)
		}
		scanner := bufio.NewScanner(f)
		for scanner.Scan() {
			params := lineRe.FindStringSubmatch(scanner.Text())
			if len(params) != 3 {
				continue
			}
			metric := params[1]
			labelsDoc := params[2]

			labels := labelsRe.FindAllStringSubmatch(labelsDoc, -1)
			labelPatterns := make([]string, len(labels))
			for i, l := range labels {
				if len(l) <= 1 {
					return nil, errors.Errorf("Label documentation %s did not match regex", labelsDoc)
				}
				labelPatterns[i] = patternRe.ReplaceAllString(l[1], "_.*")
			}

			documentedMetrics[metric] = labelPatterns
		}
	}
	return documentedMetrics, nil
}

func TestKubeStateMetricsErrorMetrics(t *testing.T) {
	metricFamilies, err := framework.ParseMetrics(framework.KsmClient.TelemetryMetrics)
	if err != nil {
		t.Fatal("Failed to get or decode telemetry metrics", err)
	}

	// This map's keys are the metrics expected in kube-state-metrics telemetry.
	// Its values are booleans, set to true when the metric is found.
	foundMetricFamily := map[string]bool{
		"kube_state_metrics_list_total":  false,
		"kube_state_metrics_watch_total": false,
	}

	for _, metricFamily := range metricFamilies {
		name := metricFamily.GetName()
		if _, expectedMetric := foundMetricFamily[name]; expectedMetric {
			foundMetricFamily[name] = true

			for _, m := range metricFamily.Metric {
				if hasLabelError(m) && m.GetCounter().GetValue() > 0 {
					t.Errorf("Metric %s in telemetry shows a list/watch error", prettyPrintCounter(name, m))
				}
			}
		}
	}

	for metricFamily, found := range foundMetricFamily {
		if !found {
			t.Errorf("Metric family %s was not found in telemetry metrics", metricFamily)
		}
	}
}

func hasLabelError(metric *dto.Metric) bool {
	for _, l := range metric.Label {
		if l.GetName() == "result" && l.GetValue() == "error" {
			return true
		}
	}
	return false
}

func prettyPrintCounter(name string, metric *dto.Metric) string {
	labelStrings := []string{}
	for _, l := range metric.Label {
		labelStrings = append(labelStrings, fmt.Sprintf(`%s="%s"`, l.GetName(), l.GetValue()))
	}
	return fmt.Sprintf("%s{%s} %d", name, strings.Join(labelStrings, ","), int(metric.GetCounter().GetValue()))
}

func TestDefaultCollectorMetricsAvailable(t *testing.T) {
	buf := &bytes.Buffer{}

	err := framework.KsmClient.Metrics(buf)
	if err != nil {
		t.Fatalf("failed to get metrics from kube-state-metrics: %v", err)
	}

	resources := map[string]struct{}{}
	files, err := os.ReadDir("../../internal/store/")
	if err != nil {
		t.Fatalf("failed to read dir to get all resouces name: %v", err)
	}

	re := regexp.MustCompile(`^([a-z]+).go$`)
	for _, file := range files {
		params := re.FindStringSubmatch(file.Name())
		if len(params) != 2 {
			continue
		}
		if params[1] == "builder" || params[1] == "utils" || params[1] == "testutils" {
			// Non resource file
			continue
		}
		if params[1] == "verticalpodautoscaler" {
			// Resource disabled by default
			continue
		}
		resources[params[1]] = struct{}{}
	}

	re = regexp.MustCompile(`^kube_([a-z]+)_`)
	scanner := bufio.NewScanner(buf)
	for scanner.Scan() {
		params := re.FindStringSubmatch(scanner.Text())
		if len(params) != 2 {
			continue
		}
		delete(resources, params[1])
	}

	err = scanner.Err()
	if err != nil {
		t.Fatalf("failed to scan metrics: %v", err)
	}

	if len(resources) != 0 {
		s := []string{}
		for k := range resources {
			s = append(s, k)
		}
		sort.Strings(s)
		t.Fatalf("failed to find metrics of resources: %s", strings.Join(s, ", "))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"github.com/ansible-semaphore/semaphore/db"
	"github.com/ansible-semaphore/semaphore/db/sql"
	"github.com/ansible-semaphore/semaphore/util"
	"github.com/snikch/goodman/transaction"
	"math/rand"
	"os"
	"time"
)

var tablesShouldBeTruncated = [...]string{
	"access_key",
	"event",
	"user__token",
	"project",
	"task__output",
	"task",
	"session",
	"project__environment",
	"project__inventory",
	"project__repository",
	"project__template",
	"project__schedule",
	"project__user",
	"user",
	"project__view",
}

// Test Runner User
func addTestRunnerUser() {
	uid := getUUID()
	testRunnerUser = &db.User{
		Username: "ITU-" + uid,
		Name:     "ITU-" + uid,
		Email:    uid + "@semaphore.test",
		Created:  db.GetParsedTime(time.Now()),
		Admin:    true,
	}

	dbConnect()
	defer store.Close()

	truncateAll()

	newUser, err := store.CreateUserWithoutPassword(*testRunnerUser)

	if err != nil {
		panic(err)
	}

	testRunnerUser.ID = newUser.ID

	addToken(adminToken, testRunnerUser.ID)
}

func truncateAll() {
	tx, err := store.Sql().Begin()
	if err != nil {
		panic(err)
	}

	_, err = tx.Exec("SET FOREIGN_KEY_CHECKS = 0")
	if err == nil {
		for _, tableName := range tablesShouldBeTruncated {
			tx.Exec("TRUNCATE TABLE " + tableName)
		}
		tx.Exec("SET FOREIGN_KEY_CHECKS = 1")
	}

	if err := tx.Commit(); err != nil {
		panic(err)
	}
}

func removeTestRunnerUser(transactions []*transaction.Transaction) {
	dbConnect()
	defer store.Close()
	_ = store.DeleteAPIToken(testRunnerUser.ID, adminToken)
	_ = store.DeleteUser(testRunnerUser.ID)
}

// Parameter Substitution
func setupObjectsAndPaths(t *transaction.Transaction) {
	alterRequestPath(t)
	alterRequestBody(t)
}

// Object Lifecycle
func addUserProjectRelation(pid int, user int) {
	_, err := store.CreateProjectUser(db.ProjectUser{
		ProjectID: pid,
		UserID:    user,
		Admin:     true,
	})
	if err != nil {
		panic(err)
	}
}

func deleteUserProjectRelation(pid int, user int) {
	err := store.DeleteProjectUser(pid, user)
	if err != nil {
		panic(err)
	}
}

func addAccessKey(pid *int) *db.AccessKey {
	uid := getUUID()
	secret := "5up3r53cr3t\n"

	key, err := store.CreateAccessKey(db.AccessKey{
		Name:      "ITK-" + uid,
		Type:      "ssh",
		Secret:    &secret,
		ProjectID: pid,
	})

	if err != nil {
		panic(err)
	}
	return &key
}

func addProject() *db.Project {
	uid := getUUID()
	project := db.Project{
		Name:    "ITP-" + uid,
		Created: time.Now(),
	}
	project, err := store.CreateProject(project)
	if err != nil {
		panic(err)
	}
	return &project
}

func addUser() *db.User {
	uid := getUUID()
	user := db.User{
		Created:  time.Now(),
		Username: "ITU-" + uid,
		Email:    "test@semaphore." + uid,
		Name:     "ITU-" + uid,
	}

	user, err := store.CreateUserWithoutPassword(user)

	if err != nil {
		panic(err)
	}
	return &user
}

func addView() *db.View {
	view, err := store.CreateView(db.View{
		ProjectID: userProject.ID,
		Title:     "Test",
		Position:  1,
	})

	if err != nil {
		panic(err)
	}

	return &view
}

func addSchedule() *db.Schedule {
	schedule, err := store.CreateSchedule(db.Schedule{
		TemplateID: int(templateID),
		CronFormat: "* * * 1 *",
		ProjectID:  userProject.ID,
	})

	if err != nil {
		panic(err)
	}

	return &schedule
}

func addTask() *db.Task {
	t := db.Task{
		ProjectID:  userProject.ID,
		TemplateID: int(templateID),
		Status:     "testing",
		UserID:     &userPathTestUser.ID,
		Created:    db.GetParsedTime(time.Now()),
	}
	_, err := store.CreateTask(t)
	if err != nil {
		fmt.Println("error during insertion of task:")
		if j, err := json.Marshal(t); err == nil {
			fmt.Println(string(j))
		} else {
			fmt.Println("can not stringify task object")
		}
		panic(err)
	}
	return &t
}

// Token Handling
func addToken(tok string, user int) {
	_, err := store.CreateAPIToken(db.APIToken{
		ID:      tok,
		Created: time.Now(),
		UserID:  user,
		Expired: false,
	})
	if err != nil {
		panic(err)
	}
}

// HELPERS
var r *rand.Rand
var randSetup = false

func getUUID() string {
	if !randSetup {
		r = rand.New(rand.NewSource(time.Now().UnixNano()))
		randSetup = true
	}
	return randomString(8)
}
func randomString(strlen int) string {
	const chars = "abcdefghijklmnopqrstuvwxyz0123456789"
	result := ""
	for i := 0; i < strlen; i++ {
		index := r.Intn(len(chars))
		result += chars[index : index+1]
	}
	return result
}

func loadConfig() {
	cwd, _ := os.Getwd()
	file, _ := os.Open(cwd + "/.dredd/config.json")
	if err := json.NewDecoder(file).Decode(&util.Config); err != nil {
		fmt.Println("Could not decode configuration!")
		panic(err)
	}
}

var store sql.SqlDb

func dbConnect() {
	store = sql.SqlDb{}

	if err := store.Connect(); err != nil {
		panic(err)
	}
}

func stringInSlice(a string, list []string) (int, bool) {
	for k, b := range list {
		if b == a {
			return k, true
		}
	}
	return 0, false
}

func printError(err error) {
	if err != nil {
		//fmt.Println(err)
		panic(err)
	}
}
<commit_msg>test: fix task creation<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"github.com/ansible-semaphore/semaphore/db"
	"github.com/ansible-semaphore/semaphore/db/sql"
	"github.com/ansible-semaphore/semaphore/util"
	"github.com/snikch/goodman/transaction"
	"math/rand"
	"os"
	"time"
)

var tablesShouldBeTruncated = [...]string{
	"access_key",
	"event",
	"user__token",
	"project",
	"task__output",
	"task",
	"session",
	"project__environment",
	"project__inventory",
	"project__repository",
	"project__template",
	"project__schedule",
	"project__user",
	"user",
	"project__view",
}

// Test Runner User
func addTestRunnerUser() {
	uid := getUUID()
	testRunnerUser = &db.User{
		Username: "ITU-" + uid,
		Name:     "ITU-" + uid,
		Email:    uid + "@semaphore.test",
		Created:  db.GetParsedTime(time.Now()),
		Admin:    true,
	}

	dbConnect()
	defer store.Close()

	truncateAll()

	newUser, err := store.CreateUserWithoutPassword(*testRunnerUser)

	if err != nil {
		panic(err)
	}

	testRunnerUser.ID = newUser.ID

	addToken(adminToken, testRunnerUser.ID)
}

func truncateAll() {
	tx, err := store.Sql().Begin()
	if err != nil {
		panic(err)
	}

	_, err = tx.Exec("SET FOREIGN_KEY_CHECKS = 0")
	if err == nil {
		for _, tableName := range tablesShouldBeTruncated {
			tx.Exec("TRUNCATE TABLE " + tableName)
		}
		tx.Exec("SET FOREIGN_KEY_CHECKS = 1")
	}

	if err := tx.Commit(); err != nil {
		panic(err)
	}
}

func removeTestRunnerUser(transactions []*transaction.Transaction) {
	dbConnect()
	defer store.Close()
	_ = store.DeleteAPIToken(testRunnerUser.ID, adminToken)
	_ = store.DeleteUser(testRunnerUser.ID)
}

// Parameter Substitution
func setupObjectsAndPaths(t *transaction.Transaction) {
	alterRequestPath(t)
	alterRequestBody(t)
}

// Object Lifecycle
func addUserProjectRelation(pid int, user int) {
	_, err := store.CreateProjectUser(db.ProjectUser{
		ProjectID: pid,
		UserID:    user,
		Admin:     true,
	})
	if err != nil {
		panic(err)
	}
}

func deleteUserProjectRelation(pid int, user int) {
	err := store.DeleteProjectUser(pid, user)
	if err != nil {
		panic(err)
	}
}

func addAccessKey(pid *int) *db.AccessKey {
	uid := getUUID()
	secret := "5up3r53cr3t\n"

	key, err := store.CreateAccessKey(db.AccessKey{
		Name:      "ITK-" + uid,
		Type:      "ssh",
		Secret:    &secret,
		ProjectID: pid,
	})

	if err != nil {
		panic(err)
	}
	return &key
}

func addProject() *db.Project {
	uid := getUUID()
	project := db.Project{
		Name:    "ITP-" + uid,
		Created: time.Now(),
	}
	project, err := store.CreateProject(project)
	if err != nil {
		panic(err)
	}
	return &project
}

func addUser() *db.User {
	uid := getUUID()
	user := db.User{
		Created:  time.Now(),
		Username: "ITU-" + uid,
		Email:    "test@semaphore." + uid,
		Name:     "ITU-" + uid,
	}

	user, err := store.CreateUserWithoutPassword(user)

	if err != nil {
		panic(err)
	}
	return &user
}

func addView() *db.View {
	view, err := store.CreateView(db.View{
		ProjectID: userProject.ID,
		Title:     "Test",
		Position:  1,
	})

	if err != nil {
		panic(err)
	}

	return &view
}

func addSchedule() *db.Schedule {
	schedule, err := store.CreateSchedule(db.Schedule{
		TemplateID: int(templateID),
		CronFormat: "* * * 1 *",
		ProjectID:  userProject.ID,
	})

	if err != nil {
		panic(err)
	}

	return &schedule
}

func addTask() *db.Task {
	t := db.Task{
		ProjectID:  userProject.ID,
		TemplateID: int(templateID),
		Status:     "testing",
		UserID:     &userPathTestUser.ID,
		Created:    db.GetParsedTime(time.Now()),
	}
	t, err := store.CreateTask(t)
	if err != nil {
		fmt.Println("error during insertion of task:")
		if j, err := json.Marshal(t); err == nil {
			fmt.Println(string(j))
		} else {
			fmt.Println("can not stringify task object")
		}
		panic(err)
	}
	return &t
}

// Token Handling
func addToken(tok string, user int) {
	_, err := store.CreateAPIToken(db.APIToken{
		ID:      tok,
		Created: time.Now(),
		UserID:  user,
		Expired: false,
	})
	if err != nil {
		panic(err)
	}
}

// HELPERS
var r *rand.Rand
var randSetup = false

func getUUID() string {
	if !randSetup {
		r = rand.New(rand.NewSource(time.Now().UnixNano()))
		randSetup = true
	}
	return randomString(8)
}
func randomString(strlen int) string {
	const chars = "abcdefghijklmnopqrstuvwxyz0123456789"
	result := ""
	for i := 0; i < strlen; i++ {
		index := r.Intn(len(chars))
		result += chars[index : index+1]
	}
	return result
}

func loadConfig() {
	cwd, _ := os.Getwd()
	file, _ := os.Open(cwd + "/.dredd/config.json")
	if err := json.NewDecoder(file).Decode(&util.Config); err != nil {
		fmt.Println("Could not decode configuration!")
		panic(err)
	}
}

var store sql.SqlDb

func dbConnect() {
	store = sql.SqlDb{}

	if err := store.Connect(); err != nil {
		panic(err)
	}
}

func stringInSlice(a string, list []string) (int, bool) {
	for k, b := range list {
		if b == a {
			return k, true
		}
	}
	return 0, false
}

func printError(err error) {
	if err != nil {
		//fmt.Println(err)
		panic(err)
	}
}
<|endoftext|>


<commit_before>package file

import (
	"bufio"
	"cred-alert/sniff"
	"fmt"
	"os"

	"github.com/pivotal-golang/lager"
)

type fileScanner struct {
	path         string
	bufioScanner *bufio.Scanner
	lineNumber   int
}

func NewFileScanner(file *os.File) *fileScanner {
	bufioScanner := bufio.NewScanner(file)

	return &fileScanner{
		path:         file.Name(),
		bufioScanner: bufioScanner,
	}
}

func (s *fileScanner) Scan(logger lager.Logger) bool {
	logger = logger.Session("file-scanner")

	success := s.bufioScanner.Scan()

	if err := s.bufioScanner.Err(); err != nil {
		fmt.Println(err)
		return false
	}

	if success {
		s.lineNumber++
	}
	return success
}

func (s *fileScanner) Line() *sniff.Line {
	return &sniff.Line{
		Content:    s.bufioScanner.Text(),
		LineNumber: s.lineNumber,
		Path:       s.path,
	}
}
<commit_msg>[#121957927] Use logger to log bufio errors<commit_after>package file

import (
	"bufio"
	"cred-alert/sniff"
	"os"

	"github.com/pivotal-golang/lager"
)

type fileScanner struct {
	path         string
	bufioScanner *bufio.Scanner
	lineNumber   int
}

func NewFileScanner(file *os.File) *fileScanner {
	bufioScanner := bufio.NewScanner(file)

	return &fileScanner{
		path:         file.Name(),
		bufioScanner: bufioScanner,
	}
}

func (s *fileScanner) Scan(logger lager.Logger) bool {
	logger = logger.Session("file-scanner")

	success := s.bufioScanner.Scan()

	if err := s.bufioScanner.Err(); err != nil {
		logger.Error("bufio-error", err)
		return false
	}

	if success {
		s.lineNumber++
	}
	return success
}

func (s *fileScanner) Line() *sniff.Line {
	return &sniff.Line{
		Content:    s.bufioScanner.Text(),
		LineNumber: s.lineNumber,
		Path:       s.path,
	}
}
<|endoftext|>


<commit_before>// Copyright 2022 Google LLC
//
// 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
//
//     https://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 tests

import (
	"fmt"
	"testing"
	"time"

	"github.com/google/go-cmp/cmp"
	"github.com/googleapis/cloud-bigtable-clients-test/testproxypb"
	"github.com/stretchr/testify/assert"
	btpb "google.golang.org/genproto/googleapis/bigtable/v2"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/keepalive"
	"google.golang.org/grpc/metadata"
)

// dummyChunkData returns a chunkData object with hardcoded family name and qualifier.
func dummyChunkData(rowKey string, value string, status RowStatus) chunkData {
	return chunkData{
		rowKey: []byte(rowKey), familyName: "f", qualifier: "col", value: value, status: status}
}

// TestReadRows_Generic_Headers tests that ReadRows request has client and resource info in the
// header.
func TestReadRows_Generic_Headers(t *testing.T) {
	// 0. Common variables
	tableName := buildTableName("table")

	// 1. Instantiate the mock server
	// Don't call mockReadRowsFn() as the behavior is to record metadata of the request.
	mdRecords := make(chan metadata.MD, 1)
	server := initMockServer(t)
	server.ReadRowsFn = func(req *btpb.ReadRowsRequest, srv btpb.Bigtable_ReadRowsServer) error {
		md, _ := metadata.FromIncomingContext(srv.Context())
		mdRecords <- md
		return nil
	}

	// 2. Build the request to test proxy
	req := testproxypb.ReadRowsRequest{
		ClientId: t.Name(),
		Request:  &btpb.ReadRowsRequest{TableName: tableName},
	}

	// 3. Perform the operation via test proxy
	doReadRowsOp(t, server, &req, nil)

	// 4. Check the request headers in the metadata
	md := <-mdRecords
	assert.NotEmpty(t, md["x-goog-api-client"])
	assert.Contains(t, md["x-goog-request-params"][0], tableName)
}

// TestReadRows_NoRetry_OutOfOrderError tests that client will fail on receiving out of order row keys.
func TestReadRows_NoRetry_OutOfOrderError(t *testing.T) {
	// 1. Instantiate the mock server
	action := &readRowsAction{
		chunks: []chunkData{
			dummyChunkData("row-01", "v1", Commit),
			// The following two rows are in bad order
			dummyChunkData("row-07", "v7", Commit),
			dummyChunkData("row-03", "v3", Commit),
		},
	}
	server := initMockServer(t)
	server.ReadRowsFn = mockReadRowsFnSimple(nil, action)

	// 2. Build the request to test proxyk
	req := testproxypb.ReadRowsRequest{
		ClientId: t.Name(),
		Request:  &btpb.ReadRowsRequest{TableName: buildTableName("table")},
	}

	// 3. Perform the operation via test proxy
	res := doReadRowsOp(t, server, &req, nil)

	// 4. Check the response (C++ and Java clients have different error messages)
	assert.Contains(t, res.GetStatus().GetMessage(), "increasing")
	t.Logf("The full error message is: %s", res.GetStatus().GetMessage())
}

// TestReadRows_NoRetry_ErrorAfterLastRow tests that when receiving a transient error after receiving
// the last row, the read will still finish successfully.
func TestReadRows_NoRetry_ErrorAfterLastRow(t *testing.T) {
	// 1. Instantiate the mock server
	sequence := []*readRowsAction{
		&readRowsAction{
			chunks: []chunkData{
				dummyChunkData("row-01", "v1", Commit)}},
		&readRowsAction{rpcError: codes.DeadlineExceeded}, // Error after returning the requested row
		&readRowsAction{
			chunks: []chunkData{
				dummyChunkData("row-05", "v5", Commit)}},
	}
	server := initMockServer(t)
	server.ReadRowsFn = mockReadRowsFn(nil, sequence)

	// 2. Build the request to test proxy
	req := testproxypb.ReadRowsRequest{
		ClientId: t.Name(),
		Request: &btpb.ReadRowsRequest{
			TableName: buildTableName("table"),
			RowsLimit: 1,
		},
	}

	// 3. Perform the operation via test proxy
	res := doReadRowsOp(t, server, &req, nil)

	// 4. Verify that the read succeeds
	checkResultOkStatus(t, res)
	assert.Equal(t, 1, len(res.GetRow()))
	assert.Equal(t, "row-01", string(res.Row[0].Key))
}

// TestReadRows_Retry_PausedScan tests that client will transparently resume the scan when a stream
// is paused.
func TestReadRows_Retry_PausedScan(t *testing.T) {
	// 1. Instantiate the mock server
	recorder := make(chan *readRowsReqRecord, 2)
	sequence := []*readRowsAction{
		&readRowsAction{
			chunks: []chunkData{
				dummyChunkData("row-01", "v1", Commit)}},
		&readRowsAction{rpcError: codes.Aborted}, // close the stream by aborting it
		&readRowsAction{
			chunks: []chunkData{
				dummyChunkData("row-05", "v5", Commit)}},
	}
	server := initMockServer(t)
	server.ReadRowsFn = mockReadRowsFn(recorder, sequence)

	// 2. Build the request to test proxy
	req := testproxypb.ReadRowsRequest{
		ClientId: t.Name(),
		Request:  &btpb.ReadRowsRequest{TableName: buildTableName("table")},
	}

	// 3. Perform the operation via test proxy
	res := doReadRowsOp(t, server, &req, nil)

	// 4a. Verify that two rows were read successfully
	checkResultOkStatus(t, res)
	assert.Equal(t, 2, len(res.GetRow()))
	assert.Equal(t, "row-01", string(res.Row[0].Key))
	assert.Equal(t, "row-05", string(res.Row[1].Key))

	// 4b. Verify that client sent the retry request properly
	loggedReq := <-recorder
	loggedRetry := <-recorder
	assert.Empty(t, loggedReq.req.GetRows().GetRowRanges())
	assert.True(t, cmp.Equal(loggedRetry.req.GetRows().GetRowRanges()[0].StartKey, &btpb.RowRange_StartKeyOpen{StartKeyOpen: []byte("row-01")}))
}

// TestReadRows_Retry_LastScannedRow tests that client will resume from last scan row key.
func TestReadRows_Retry_LastScannedRow(t *testing.T) {
	// 1. Instantiate the mock server
	recorder := make(chan *readRowsReqRecord, 2)
	sequence := []*readRowsAction{
		&readRowsAction{
			chunks: []chunkData{
				dummyChunkData("abar", "v_a", Commit)}},
		&readRowsAction{
			chunks: []chunkData{
				dummyChunkData("qfoo", "v_q", Drop)}}, // Chunkless response due to Drop
		&readRowsAction{rpcError: codes.DeadlineExceeded}, // Server-side DeadlineExceeded should be retry-able.
		&readRowsAction{
			chunks: []chunkData{
				dummyChunkData("zbar", "v_z", Commit)}},
	}
	server := initMockServer(t)
	server.ReadRowsFn = mockReadRowsFn(recorder, sequence)

	// 2. Build the request to test proxy
	req := testproxypb.ReadRowsRequest{
		ClientId: t.Name(),
		Request:  &btpb.ReadRowsRequest{TableName: buildTableName("table")},
	}

	// 3. Perform the operation via test proxy
	res := doReadRowsOp(t, server, &req, nil)

	// 4a. Verify that rows aabar and zzbar were read successfully (qqfoo doesn't match the filter)
	checkResultOkStatus(t, res)
	assert.Equal(t, 2, len(res.GetRow()))
	assert.Equal(t, "abar", string(res.Row[0].Key))
	assert.Equal(t, "zbar", string(res.Row[1].Key))

	// 4b. Verify that client sent the retry request properly
	loggedReq := <-recorder
	loggedRetry := <-recorder
	assert.Empty(t, loggedReq.req.GetRows().GetRowRanges())
	assert.True(t, cmp.Equal(loggedRetry.req.GetRows().GetRowRanges()[0].StartKey, &btpb.RowRange_StartKeyOpen{StartKeyOpen: []byte("qfoo")}))
}

// TestReadRows_Generic_MultiStreams tests that client can have multiple concurrent streams.
func TestReadRows_Generic_MultiStreams(t *testing.T) {
	// 0. Common variable
	rowKeys := [][]string{
		[]string{"op0-row-a", "op0-row-b"},
		[]string{"op1-row-a", "op1-row-b"},
		[]string{"op2-row-a", "op2-row-b"},
		[]string{"op3-row-a", "op3-row-b"},
		[]string{"op4-row-a", "op4-row-b"},
	}
	concurrency := len(rowKeys)
	const requestRecorderCapacity = 10

	// 1. Instantiate the mock server
	recorder := make(chan *readRowsReqRecord, requestRecorderCapacity)
	actions := make([]*readRowsAction, concurrency)
	for i := 0; i < concurrency; i++ {
		// Each request will get a different response.
		actions[i] = &readRowsAction{
			chunks: []chunkData{
				dummyChunkData(rowKeys[i][0], fmt.Sprintf("value%d-a", i), Commit),
				dummyChunkData(rowKeys[i][1], fmt.Sprintf("value%d-b", i), Commit),
			},
			delayStr: "2s",
		}
	}
	server := initMockServer(t)
	server.ReadRowsFn = mockReadRowsFnSimple(recorder, actions...)

	// 2. Build the requests to test proxy
	reqs := make([]*testproxypb.ReadRowsRequest, concurrency)
	for i := 0; i < concurrency; i++ {
		reqs[i] = &testproxypb.ReadRowsRequest{
			ClientId: t.Name(),
			Request: &btpb.ReadRowsRequest{
				TableName: buildTableName("table"),
				Rows: &btpb.RowSet{
					RowKeys: [][]byte{[]byte(rowKeys[i][0]), []byte(rowKeys[i][1])},
				},
			},
		}
	}

	// 3. Perform the operations via test proxy
	results := doReadRowsOps(t, server, reqs, nil)

	// 4a. Check that all the requests succeeded
	assert.Equal(t, concurrency, len(results))
	checkResultOkStatus(t, results...)

	// 4b. Check that the timestamps of requests should be very close
	assert.Equal(t, concurrency, len(recorder))
	checkRequestsAreWithin(t, 1000, recorder)

	// 4c. Check the row keys in the results.
	for i := 0; i < concurrency; i++ {
		assert.Equal(t, rowKeys[i][0], string(results[i].Row[0].Key))
		assert.Equal(t, rowKeys[i][1], string(results[i].Row[1].Key))
	}
}

// TestReadRows_Retry_StreamReset tests that client will retry on stream reset.
func TestReadRows_Retry_StreamReset(t *testing.T) {
	// 0. Common variable
	const maxConnAge = 4 * time.Second
	const maxConnAgeGrace = time.Second

	// 1. Instantiate the mock server
	recorder := make(chan *readRowsReqRecord, 3)
	sequence := []*readRowsAction{
		&readRowsAction{
			chunks: []chunkData{
				dummyChunkData("abar", "v_a", Commit)}},
		&readRowsAction{
			chunks: []chunkData{
				dummyChunkData("qbar", "v_q", Commit)},
			delayStr: "10s"}, // Stream resets before sending chunks.
		&readRowsAction{
			chunks: []chunkData{
				dummyChunkData("qbar", "v_q", Commit)}},
		&readRowsAction{
			chunks: []chunkData{
				dummyChunkData("zbar", "v_z", Commit)}},
	}
	serverOpt := grpc.KeepaliveParams(
		keepalive.ServerParameters{
			MaxConnectionAge:      maxConnAge,
			MaxConnectionAgeGrace: maxConnAgeGrace,
		})
	server := initMockServer(t, serverOpt)
	server.ReadRowsFn = mockReadRowsFn(recorder, sequence)

	// 2. Build the request to test proxy
	req := testproxypb.ReadRowsRequest{
		ClientId: t.Name(),
		Request:  &btpb.ReadRowsRequest{TableName: buildTableName("table")},
	}

	// 3. Perform the operation via test proxy
	res := doReadRowsOp(t, server, &req, nil)

	// 4a. Verify that rows were read successfully
	checkResultOkStatus(t, res)
	assert.Equal(t, 3, len(res.GetRow()))
	assert.Equal(t, "abar", string(res.Row[0].Key))
	assert.Equal(t, "qbar", string(res.Row[1].Key))
	assert.Equal(t, "zbar", string(res.Row[2].Key))

	// 4b. Verify that client sent the only retry request properly
	assert.Equal(t, 2, len(recorder))
	loggedReq := <-recorder
	loggedRetry := <-recorder
	assert.Empty(t, loggedReq.req.GetRows().GetRowRanges())
	assert.True(t, cmp.Equal(loggedRetry.req.GetRows().GetRowRanges()[0].StartKey, &btpb.RowRange_StartKeyOpen{StartKeyOpen: []byte("abar")}))
}

// TestReadRows_Generic_EmptyTableNoRows tests that reads on an empty table returns 0 rows.
func TestReadRows_NoRetry_EmptyTableNoRows(t *testing.T) {
	// 1. Instantiate the mock server
	recorder := make(chan *readRowsReqRecord, 3)
	action := &readRowsAction{
		chunks: []chunkData{}}
	server := initMockServer(t)
	server.ReadRowsFn = mockReadRowsFnSimple(recorder, action)

	// 2. Build the request to test proxy
	req := testproxypb.ReadRowsRequest{
		ClientId: t.Name(),
		Request:  &btpb.ReadRowsRequest{TableName: buildTableName("table")},
	}

	// 3. Perform the operation via test proxy
	res := doReadRowsOp(t, server, &req, nil)
	assert.Len(t, res.Row, 0)
}
<commit_msg>Update readrows_test.go<commit_after>// Copyright 2022 Google LLC
//
// 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
//
//     https://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 tests

import (
	"fmt"
	"testing"
	"time"

	"github.com/google/go-cmp/cmp"
	"github.com/googleapis/cloud-bigtable-clients-test/testproxypb"
	"github.com/stretchr/testify/assert"
	btpb "google.golang.org/genproto/googleapis/bigtable/v2"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/keepalive"
	"google.golang.org/grpc/metadata"
)

// dummyChunkData returns a chunkData object with hardcoded family name and qualifier.
func dummyChunkData(rowKey string, value string, status RowStatus) chunkData {
	return chunkData{
		rowKey: []byte(rowKey), familyName: "f", qualifier: "col", value: value, status: status}
}

// TestReadRows_Generic_Headers tests that ReadRows request has client and resource info in the
// header.
func TestReadRows_Generic_Headers(t *testing.T) {
	// 0. Common variables
	tableName := buildTableName("table")

	// 1. Instantiate the mock server
	// Don't call mockReadRowsFn() as the behavior is to record metadata of the request.
	mdRecords := make(chan metadata.MD, 1)
	server := initMockServer(t)
	server.ReadRowsFn = func(req *btpb.ReadRowsRequest, srv btpb.Bigtable_ReadRowsServer) error {
		md, _ := metadata.FromIncomingContext(srv.Context())
		mdRecords <- md
		return nil
	}

	// 2. Build the request to test proxy
	req := testproxypb.ReadRowsRequest{
		ClientId: t.Name(),
		Request:  &btpb.ReadRowsRequest{TableName: tableName},
	}

	// 3. Perform the operation via test proxy
	doReadRowsOp(t, server, &req, nil)

	// 4. Check the request headers in the metadata
	md := <-mdRecords
	assert.NotEmpty(t, md["x-goog-api-client"])
	assert.Contains(t, md["x-goog-request-params"][0], tableName)
}

// TestReadRows_NoRetry_OutOfOrderError tests that client will fail on receiving out of order row keys.
func TestReadRows_NoRetry_OutOfOrderError(t *testing.T) {
	// 1. Instantiate the mock server
	action := &readRowsAction{
		chunks: []chunkData{
			dummyChunkData("row-01", "v1", Commit),
			// The following two rows are in bad order
			dummyChunkData("row-07", "v7", Commit),
			dummyChunkData("row-03", "v3", Commit),
		},
	}
	server := initMockServer(t)
	server.ReadRowsFn = mockReadRowsFnSimple(nil, action)

	// 2. Build the request to test proxyk
	req := testproxypb.ReadRowsRequest{
		ClientId: t.Name(),
		Request:  &btpb.ReadRowsRequest{TableName: buildTableName("table")},
	}

	// 3. Perform the operation via test proxy
	res := doReadRowsOp(t, server, &req, nil)

	// 4. Check the response (C++ and Java clients have different error messages)
	assert.Contains(t, res.GetStatus().GetMessage(), "increasing")
	t.Logf("The full error message is: %s", res.GetStatus().GetMessage())
}

// TestReadRows_NoRetry_ErrorAfterLastRow tests that when receiving a transient error after receiving
// the last row, the read will still finish successfully.
func TestReadRows_NoRetry_ErrorAfterLastRow(t *testing.T) {
	// 1. Instantiate the mock server
	sequence := []*readRowsAction{
		&readRowsAction{
			chunks: []chunkData{
				dummyChunkData("row-01", "v1", Commit)}},
		&readRowsAction{rpcError: codes.DeadlineExceeded}, // Error after returning the requested row
		&readRowsAction{
			chunks: []chunkData{
				dummyChunkData("row-05", "v5", Commit)}},
	}
	server := initMockServer(t)
	server.ReadRowsFn = mockReadRowsFn(nil, sequence)

	// 2. Build the request to test proxy
	req := testproxypb.ReadRowsRequest{
		ClientId: t.Name(),
		Request: &btpb.ReadRowsRequest{
			TableName: buildTableName("table"),
			RowsLimit: 1,
		},
	}

	// 3. Perform the operation via test proxy
	res := doReadRowsOp(t, server, &req, nil)

	// 4. Verify that the read succeeds
	checkResultOkStatus(t, res)
	assert.Equal(t, 1, len(res.GetRow()))
	assert.Equal(t, "row-01", string(res.Row[0].Key))
}

// TestReadRows_Retry_PausedScan tests that client will transparently resume the scan when a stream
// is paused.
func TestReadRows_Retry_PausedScan(t *testing.T) {
	// 1. Instantiate the mock server
	recorder := make(chan *readRowsReqRecord, 2)
	sequence := []*readRowsAction{
		&readRowsAction{
			chunks: []chunkData{
				dummyChunkData("row-01", "v1", Commit)}},
		&readRowsAction{rpcError: codes.Aborted}, // close the stream by aborting it
		&readRowsAction{
			chunks: []chunkData{
				dummyChunkData("row-05", "v5", Commit)}},
	}
	server := initMockServer(t)
	server.ReadRowsFn = mockReadRowsFn(recorder, sequence)

	// 2. Build the request to test proxy
	req := testproxypb.ReadRowsRequest{
		ClientId: t.Name(),
		Request:  &btpb.ReadRowsRequest{TableName: buildTableName("table")},
	}

	// 3. Perform the operation via test proxy
	res := doReadRowsOp(t, server, &req, nil)

	// 4a. Verify that two rows were read successfully
	checkResultOkStatus(t, res)
	assert.Equal(t, 2, len(res.GetRow()))
	assert.Equal(t, "row-01", string(res.Row[0].Key))
	assert.Equal(t, "row-05", string(res.Row[1].Key))

	// 4b. Verify that client sent the retry request properly
	loggedReq := <-recorder
	loggedRetry := <-recorder
	assert.Empty(t, loggedReq.req.GetRows().GetRowRanges())
	assert.True(t, cmp.Equal(loggedRetry.req.GetRows().GetRowRanges()[0].StartKey, &btpb.RowRange_StartKeyOpen{StartKeyOpen: []byte("row-01")}))
}

// TestReadRows_Retry_LastScannedRow tests that client will resume from last scan row key.
func TestReadRows_Retry_LastScannedRow(t *testing.T) {
	// 1. Instantiate the mock server
	recorder := make(chan *readRowsReqRecord, 2)
	sequence := []*readRowsAction{
		&readRowsAction{
			chunks: []chunkData{
				dummyChunkData("abar", "v_a", Commit)}},
		&readRowsAction{
			chunks: []chunkData{
				dummyChunkData("qfoo", "v_q", Drop)}}, // Chunkless response due to Drop
		&readRowsAction{rpcError: codes.DeadlineExceeded}, // Server-side DeadlineExceeded should be retry-able.
		&readRowsAction{
			chunks: []chunkData{
				dummyChunkData("zbar", "v_z", Commit)}},
	}
	server := initMockServer(t)
	server.ReadRowsFn = mockReadRowsFn(recorder, sequence)

	// 2. Build the request to test proxy
	req := testproxypb.ReadRowsRequest{
		ClientId: t.Name(),
		Request:  &btpb.ReadRowsRequest{TableName: buildTableName("table")},
	}

	// 3. Perform the operation via test proxy
	res := doReadRowsOp(t, server, &req, nil)

	// 4a. Verify that rows aabar and zzbar were read successfully (qqfoo doesn't match the filter)
	checkResultOkStatus(t, res)
	assert.Equal(t, 2, len(res.GetRow()))
	assert.Equal(t, "abar", string(res.Row[0].Key))
	assert.Equal(t, "zbar", string(res.Row[1].Key))

	// 4b. Verify that client sent the retry request properly
	loggedReq := <-recorder
	loggedRetry := <-recorder
	assert.Empty(t, loggedReq.req.GetRows().GetRowRanges())
	assert.True(t, cmp.Equal(loggedRetry.req.GetRows().GetRowRanges()[0].StartKey, &btpb.RowRange_StartKeyOpen{StartKeyOpen: []byte("qfoo")}))
}

// TestReadRows_Generic_MultiStreams tests that client can have multiple concurrent streams.
func TestReadRows_Generic_MultiStreams(t *testing.T) {
	// 0. Common variable
	rowKeys := [][]string{
		[]string{"op0-row-a", "op0-row-b"},
		[]string{"op1-row-a", "op1-row-b"},
		[]string{"op2-row-a", "op2-row-b"},
		[]string{"op3-row-a", "op3-row-b"},
		[]string{"op4-row-a", "op4-row-b"},
	}
	concurrency := len(rowKeys)
	const requestRecorderCapacity = 10

	// 1. Instantiate the mock server
	recorder := make(chan *readRowsReqRecord, requestRecorderCapacity)
	actions := make([]*readRowsAction, concurrency)
	for i := 0; i < concurrency; i++ {
		// Each request will get a different response.
		actions[i] = &readRowsAction{
			chunks: []chunkData{
				dummyChunkData(rowKeys[i][0], fmt.Sprintf("value%d-a", i), Commit),
				dummyChunkData(rowKeys[i][1], fmt.Sprintf("value%d-b", i), Commit),
			},
			delayStr: "2s",
		}
	}
	server := initMockServer(t)
	server.ReadRowsFn = mockReadRowsFnSimple(recorder, actions...)

	// 2. Build the requests to test proxy
	reqs := make([]*testproxypb.ReadRowsRequest, concurrency)
	for i := 0; i < concurrency; i++ {
		reqs[i] = &testproxypb.ReadRowsRequest{
			ClientId: t.Name(),
			Request: &btpb.ReadRowsRequest{
				TableName: buildTableName("table"),
				Rows: &btpb.RowSet{
					RowKeys: [][]byte{[]byte(rowKeys[i][0]), []byte(rowKeys[i][1])},
				},
			},
		}
	}

	// 3. Perform the operations via test proxy
	results := doReadRowsOps(t, server, reqs, nil)

	// 4a. Check that all the requests succeeded
	assert.Equal(t, concurrency, len(results))
	checkResultOkStatus(t, results...)

	// 4b. Check that the timestamps of requests should be very close
	assert.Equal(t, concurrency, len(recorder))
	checkRequestsAreWithin(t, 1000, recorder)

	// 4c. Check the row keys in the results.
	for i := 0; i < concurrency; i++ {
		assert.Equal(t, rowKeys[i][0], string(results[i].Row[0].Key))
		assert.Equal(t, rowKeys[i][1], string(results[i].Row[1].Key))
	}
}

// TestReadRows_Retry_StreamReset tests that client will retry on stream reset.
func TestReadRows_Retry_StreamReset(t *testing.T) {
	// 0. Common variable
	const maxConnAge = 4 * time.Second
	const maxConnAgeGrace = time.Second

	// 1. Instantiate the mock server
	recorder := make(chan *readRowsReqRecord, 3)
	sequence := []*readRowsAction{
		&readRowsAction{
			chunks: []chunkData{
				dummyChunkData("abar", "v_a", Commit)}},
		&readRowsAction{
			chunks: []chunkData{
				dummyChunkData("qbar", "v_q", Commit)},
			delayStr: "10s"}, // Stream resets before sending chunks.
		&readRowsAction{
			chunks: []chunkData{
				dummyChunkData("qbar", "v_q", Commit)}},
		&readRowsAction{
			chunks: []chunkData{
				dummyChunkData("zbar", "v_z", Commit)}},
	}
	serverOpt := grpc.KeepaliveParams(
		keepalive.ServerParameters{
			MaxConnectionAge:      maxConnAge,
			MaxConnectionAgeGrace: maxConnAgeGrace,
		})
	server := initMockServer(t, serverOpt)
	server.ReadRowsFn = mockReadRowsFn(recorder, sequence)

	// 2. Build the request to test proxy
	req := testproxypb.ReadRowsRequest{
		ClientId: t.Name(),
		Request:  &btpb.ReadRowsRequest{TableName: buildTableName("table")},
	}

	// 3. Perform the operation via test proxy
	res := doReadRowsOp(t, server, &req, nil)

	// 4a. Verify that rows were read successfully
	checkResultOkStatus(t, res)
	assert.Equal(t, 3, len(res.GetRow()))
	assert.Equal(t, "abar", string(res.Row[0].Key))
	assert.Equal(t, "qbar", string(res.Row[1].Key))
	assert.Equal(t, "zbar", string(res.Row[2].Key))

	// 4b. Verify that client sent the only retry request properly
	assert.Equal(t, 2, len(recorder))
	loggedReq := <-recorder
	loggedRetry := <-recorder
	assert.Empty(t, loggedReq.req.GetRows().GetRowRanges())
	assert.True(t, cmp.Equal(loggedRetry.req.GetRows().GetRowRanges()[0].StartKey, &btpb.RowRange_StartKeyOpen{StartKeyOpen: []byte("abar")}))
}

// TestReadRows_NoRetry_EmptyTableNoRows tests that reads on an empty table returns 0 rows.
func TestReadRows_NoRetry_EmptyTableNoRows(t *testing.T) {
	// 1. Instantiate the mock server
	recorder := make(chan *readRowsReqRecord, 3)
	action := &readRowsAction{
		chunks: []chunkData{}}
	server := initMockServer(t)
	server.ReadRowsFn = mockReadRowsFnSimple(recorder, action)

	// 2. Build the request to test proxy
	req := testproxypb.ReadRowsRequest{
		ClientId: t.Name(),
		Request:  &btpb.ReadRowsRequest{TableName: buildTableName("table")},
	}

	// 3. Perform the operation via test proxy
	res := doReadRowsOp(t, server, &req, nil)
	assert.Len(t, res.Row, 0)
}
<|endoftext|>


<commit_before>// main simply contains the primary web serving code that allows peers to
// register and unregister as give mode peers running within the Lantern
// network
package main

import (
	"crypto/tls"
	"fmt"
	"net/http"
	"strconv"
	"strings"

	"github.com/getlantern/keyman"
	"github.com/getlantern/tlsdefaults"
)

const (
	PKFile   = "pk.pem"
	CertFile = "cert.pem"
)

func startHttp() {
	http.HandleFunc("/register", register)
	http.HandleFunc("/unregister", unregister)
	laddr := fmt.Sprintf(":%d", *port)

	tlsConfig := tlsdefaults.Server()
	_, _, err := keyman.StoredPKAndCert(PKFile, CertFile, "Lantern", "localhost")
	if err != nil {
		log.Fatalf("Unable to initialize private key and certificate: %v", err)
	}
	cert, err := tls.LoadX509KeyPair(CertFile, PKFile)
	if err != nil {
		log.Fatalf("Unable to load certificate and key from %s and %s: %s", CertFile, PKFile, err)
	}
	tlsConfig.Certificates = []tls.Certificate{cert}

	log.Debugf("About to listen at %v", laddr)
	l, err := tls.Listen("tcp", laddr, tlsConfig)
	if err != nil {
		log.Fatalf("Unable to listen for tls connections at %s: %s", laddr, err)
	}

	log.Debug("About to serve")
	err = http.Serve(l, nil)
	if err != nil {
		log.Fatalf("Unable to serve: %s", err)
	}
}

// register is the entry point for peers registering themselves with the service.
// If peers are successfully vetted, they'll be added to the DNS round robin.
func register(resp http.ResponseWriter, req *http.Request) {
	name, ip, port, err := getHostInfo(req)
	if err == nil && port != 443 {
		err = fmt.Errorf("Port %d not supported, only port 443 is supported", port)
	}
	if err != nil {
		resp.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(resp, err.Error())
		return
	}

	h := getOrCreateHost(name, ip)
	online, connectionRefused, timedOut := h.status()
	if timedOut {
		log.Debugf("%v timed out waiting for status, returning 500 error", h)
		resp.WriteHeader(500)
		fmt.Fprintf(resp, "Timed out waiting for status")
		return
	}

	if online {
		resp.WriteHeader(200)
		fmt.Fprintln(resp, "Connectivity to proxy confirmed")
		return
	}

	// Note this may not work across platforms, but the intent
	// is to tell the client if the connection was flat out
	// refused as opposed to timed out in order to allow them
	// to configure their router if possible.
	if connectionRefused {
		// 417 response code.
		resp.WriteHeader(http.StatusExpectationFailed)
		fmt.Fprintln(resp, "No connectivity to proxy - connection refused")
	} else {
		// 408 response code.
		resp.WriteHeader(http.StatusRequestTimeout)
		fmt.Fprintln(resp, "No connectivity to proxy - test request timed out")
	}
}

// unregister is the HTTP endpoint for removing peers from DNS. Peers are
// unregistered based on their ip (not their name).
func unregister(resp http.ResponseWriter, req *http.Request) {
	_, ip, _, err := getHostInfo(req)
	if err != nil {
		resp.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(resp, err.Error())
		return
	}

	h := getHostByIp(ip)
	msg := "Host not registered"
	if h != nil {
		h.unregister()
		msg = "Host unregistered"
	}
	resp.WriteHeader(200)
	fmt.Fprintln(resp, msg)
}

func getHostInfo(req *http.Request) (name string, ip string, port int, err error) {
	name = req.FormValue("name")
	if name == "" {
		err = fmt.Errorf("Please specify a name")
		return
	}
	ip = clientIpFor(req, name)
	if ip == "" {
		err = fmt.Errorf("Unable to determine IP address")
		return
	}
	portString := req.FormValue("port")

	if portString != "" {
		port, err = strconv.Atoi(portString)
		if err != nil {
			err = fmt.Errorf("Received invalid port for %v - %v: %v", name, ip, portString)
		}
	}

	return
}

func clientIpFor(req *http.Request, name string) string {
	// Client requested their info
	clientIp := req.Header.Get("X-Peerscanner-Forwarded-For")
	if clientIp == "" {
		clientIp = req.Header.Get("X-Forwarded-For")
	}
	if clientIp == "" && isFallback(name) {
		// Use direct IP for fallbacks
		clientIp = strings.Split(req.RemoteAddr, ":")[0]
	}
	// clientIp may contain multiple ips, use the first
	ips := strings.Split(clientIp, ",")
	ip := strings.TrimSpace(ips[0])
	// TODO: need a more robust way to determine when a non-fallback host looks
	// like a fallback.
	if !isFallback(name) && strings.HasPrefix(ip, "128.199") {
		log.Errorf("Found fallback ip %v for non-fallback host %v", ip, name)
		return ""
	} else if isFallback(name) && !strings.HasPrefix(ip, "128.199") {
		log.Errorf("Found non-fallback ip %v for fallback host %v", ip, name)
		return ""
	}
	return ip
}
<commit_msg>Including more ip prefixes as fallbacks, closes getlantern/lantern#1739<commit_after>// main simply contains the primary web serving code that allows peers to
// register and unregister as give mode peers running within the Lantern
// network
package main

import (
	"crypto/tls"
	"fmt"
	"net/http"
	"strconv"
	"strings"

	"github.com/getlantern/keyman"
	"github.com/getlantern/tlsdefaults"
)

const (
	PKFile   = "pk.pem"
	CertFile = "cert.pem"
)

func startHttp() {
	http.HandleFunc("/register", register)
	http.HandleFunc("/unregister", unregister)
	laddr := fmt.Sprintf(":%d", *port)

	tlsConfig := tlsdefaults.Server()
	_, _, err := keyman.StoredPKAndCert(PKFile, CertFile, "Lantern", "localhost")
	if err != nil {
		log.Fatalf("Unable to initialize private key and certificate: %v", err)
	}
	cert, err := tls.LoadX509KeyPair(CertFile, PKFile)
	if err != nil {
		log.Fatalf("Unable to load certificate and key from %s and %s: %s", CertFile, PKFile, err)
	}
	tlsConfig.Certificates = []tls.Certificate{cert}

	log.Debugf("About to listen at %v", laddr)
	l, err := tls.Listen("tcp", laddr, tlsConfig)
	if err != nil {
		log.Fatalf("Unable to listen for tls connections at %s: %s", laddr, err)
	}

	log.Debug("About to serve")
	err = http.Serve(l, nil)
	if err != nil {
		log.Fatalf("Unable to serve: %s", err)
	}
}

// register is the entry point for peers registering themselves with the service.
// If peers are successfully vetted, they'll be added to the DNS round robin.
func register(resp http.ResponseWriter, req *http.Request) {
	name, ip, port, err := getHostInfo(req)
	if err == nil && port != 443 {
		err = fmt.Errorf("Port %d not supported, only port 443 is supported", port)
	}
	if err != nil {
		resp.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(resp, err.Error())
		return
	}

	h := getOrCreateHost(name, ip)
	online, connectionRefused, timedOut := h.status()
	if timedOut {
		log.Debugf("%v timed out waiting for status, returning 500 error", h)
		resp.WriteHeader(500)
		fmt.Fprintf(resp, "Timed out waiting for status")
		return
	}

	if online {
		resp.WriteHeader(200)
		fmt.Fprintln(resp, "Connectivity to proxy confirmed")
		return
	}

	// Note this may not work across platforms, but the intent
	// is to tell the client if the connection was flat out
	// refused as opposed to timed out in order to allow them
	// to configure their router if possible.
	if connectionRefused {
		// 417 response code.
		resp.WriteHeader(http.StatusExpectationFailed)
		fmt.Fprintln(resp, "No connectivity to proxy - connection refused")
	} else {
		// 408 response code.
		resp.WriteHeader(http.StatusRequestTimeout)
		fmt.Fprintln(resp, "No connectivity to proxy - test request timed out")
	}
}

// unregister is the HTTP endpoint for removing peers from DNS. Peers are
// unregistered based on their ip (not their name).
func unregister(resp http.ResponseWriter, req *http.Request) {
	_, ip, _, err := getHostInfo(req)
	if err != nil {
		resp.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(resp, err.Error())
		return
	}

	h := getHostByIp(ip)
	msg := "Host not registered"
	if h != nil {
		h.unregister()
		msg = "Host unregistered"
	}
	resp.WriteHeader(200)
	fmt.Fprintln(resp, msg)
}

func getHostInfo(req *http.Request) (name string, ip string, port int, err error) {
	name = req.FormValue("name")
	if name == "" {
		err = fmt.Errorf("Please specify a name")
		return
	}
	ip = clientIpFor(req, name)
	if ip == "" {
		err = fmt.Errorf("Unable to determine IP address")
		return
	}
	portString := req.FormValue("port")

	if portString != "" {
		port, err = strconv.Atoi(portString)
		if err != nil {
			err = fmt.Errorf("Received invalid port for %v - %v: %v", name, ip, portString)
		}
	}

	return
}

func clientIpFor(req *http.Request, name string) string {
	// Client requested their info
	clientIp := req.Header.Get("X-Peerscanner-Forwarded-For")
	if clientIp == "" {
		clientIp = req.Header.Get("X-Forwarded-For")
	}
	if clientIp == "" && isFallback(name) {
		// Use direct IP for fallbacks
		clientIp = strings.Split(req.RemoteAddr, ":")[0]
	}
	// clientIp may contain multiple ips, use the first
	ips := strings.Split(clientIp, ",")
	ip := strings.TrimSpace(ips[0])
	// TODO: need a more robust way to determine when a non-fallback host looks
	// like a fallback.
	hasFallbackIp := isFallbackIp(ip)
	if !isFallback(name) && hasFallbackIp {
		log.Errorf("Found fallback ip %v for non-fallback host %v", ip, name)
		return ""
	} else if isFallback(name) && !hasFallbackIp {
		log.Errorf("Found non-fallback ip %v for fallback host %v", ip, name)
		return ""
	}
	return ip
}

func isFallbackIp(ip string) bool {
	for _, prefix := range fallbackIPPrefixes {
		if strings.HasPrefix(ip, prefix) {
			return true
		}
	}

	return false
}

var fallbackIPPrefixes = []string{
	"128.199",
	"178.62",
	"188.166",
}
<|endoftext|>


<commit_before>// main simply contains the primary web serving code that allows peers to
// register and unregister as give mode peers running within the Lantern
// network
package main

import (
	"crypto/tls"
	"fmt"
	"net/http"
	"strconv"
	"strings"

	"github.com/getlantern/keyman"
	"github.com/getlantern/tlsdefaults"
)

const (
	PKFile   = "pk.pem"
	CertFile = "cert.pem"
)

const (
	cloudflareBit = 1 << iota
	cloudfrontBit = 1 << iota
)

func startHttp() {
	http.HandleFunc("/register", register)
	http.HandleFunc("/unregister", unregister)
	laddr := fmt.Sprintf(":%d", *port)

	tlsConfig := tlsdefaults.Server()
	_, _, err := keyman.StoredPKAndCert(PKFile, CertFile, "Lantern", "localhost")
	if err != nil {
		log.Fatalf("Unable to initialize private key and certificate: %v", err)
	}
	cert, err := tls.LoadX509KeyPair(CertFile, PKFile)
	if err != nil {
		log.Fatalf("Unable to load certificate and key from %s and %s: %s", CertFile, PKFile, err)
	}
	tlsConfig.Certificates = []tls.Certificate{cert}

	log.Debugf("About to listen at %v", laddr)
	l, err := tls.Listen("tcp", laddr, tlsConfig)
	if err != nil {
		log.Fatalf("Unable to listen for tls connections at %s: %s", laddr, err)
	}

	log.Debug("About to serve")
	err = http.Serve(l, nil)
	if err != nil {
		log.Fatalf("Unable to serve: %s", err)
	}
}

// register is the entry point for peers registering themselves with the service.
// If peers are successfully vetted, they'll be added to the DNS round robin.
func register(resp http.ResponseWriter, req *http.Request) {
	name, ip, port, supportedFronts, err := getHostInfo(req)
	if err == nil && !(port == "80" || port == "443") {
		err = fmt.Errorf("Port %d not supported, only ports 80 and 443 are supported", port)
	}
	if err != nil {
		resp.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(resp, err.Error())
		return
	}
	if isPeer(name) {
		log.Debugf("Not adding peer %v because we're not using peers at the moment", name)
		resp.WriteHeader(200)
		fmt.Fprintln(resp, "Peers disabled at the moment")
		return
	}
	online := true
	connectionRefused := false
	timedOut := false

	h := getOrCreateHost(name, ip, port)
	online, connectionRefused, timedOut = h.status()
	if online {
		resp.WriteHeader(200)
		fmt.Fprintln(resp, "Connectivity to proxy confirmed")
		if (supportedFronts & cloudfrontBit) == cloudfrontBit {
			h.initCloudfront()
		}
		fstr := "frontfqdns: {cloudflare: " + name + "." + *cfldomain
		if h.cfrDist != nil {
			fstr += ", cloudfront: " + h.cfrDist.Domain
		}
		fstr += "}"
		fmt.Fprintln(resp, fstr)

		return
	}
	if timedOut {
		log.Debugf("%v timed out waiting for status, returning 500 error", h)
		resp.WriteHeader(500)
		fmt.Fprintf(resp, "Timed out waiting for status")
		return
	}

	// Note this may not work across platforms, but the intent
	// is to tell the client if the connection was flat out
	// refused as opposed to timed out in order to allow them
	// to configure their router if possible.
	if connectionRefused {
		// 417 response code.
		resp.WriteHeader(http.StatusExpectationFailed)
		fmt.Fprintln(resp, "No connectivity to proxy - connection refused")
	} else {
		// 408 response code.
		resp.WriteHeader(http.StatusRequestTimeout)
		fmt.Fprintln(resp, "No connectivity to proxy - test request timed out")
	}
}

// unregister is the HTTP endpoint for removing peers from DNS. Peers are
// unregistered based on their ip (not their name).
func unregister(resp http.ResponseWriter, req *http.Request) {
	_, ip, _, _, err := getHostInfo(req)
	if err != nil {
		resp.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(resp, err.Error())
		return
	}

	h := getHostByIp(ip)
	msg := "Host not registered"
	if h != nil {
		h.unregister()
		msg = "Host unregistered"
	}
	resp.WriteHeader(200)
	fmt.Fprintln(resp, msg)
}

func getHostInfo(req *http.Request) (name string, ip string, port string, supportedFronts int, err error) {
	err = req.ParseForm()
	if err != nil {
		err = fmt.Errorf("Couldn't parse form: %v", err)
		return
	}
	name = getSingleFormValue(req, "name")
	if name == "" {
		err = fmt.Errorf("Please specify a name")
		return
	}
	ip = clientIpFor(req, name)
	if ip == "" {
		err = fmt.Errorf("Unable to determine IP address")
		return
	}
	port = getSingleFormValue(req, "port")
	if port != "" {
		_, err = strconv.Atoi(port)
		if err != nil {
			err = fmt.Errorf("Received invalid port for %v - %v: %v", name, ip, port)
			return
		}
	}
	fronts := req.Form["fronts"]
	if len(fronts) == 0 {
		// backwards compatibility
		fronts = []string{"cloudflare"}
	}
	for _, front := range fronts {
		switch front {
		case "cloudflare":
			supportedFronts |= cloudflareBit
		case "cloudfront":
			supportedFronts |= cloudfrontBit
		default:
			// Ignore these for forward compatibility.
			log.Debugf("Unrecognized front: %v", front)
		}
	}
	return
}

func clientIpFor(req *http.Request, name string) string {
	// Client requested their info
	clientIp := req.Header.Get("X-Peerscanner-Forwarded-For")
	if clientIp == "" {
		clientIp = req.Header.Get("X-Forwarded-For")
	}
	if clientIp == "" && isFallback(name) {
		// Use direct IP for fallbacks
		clientIp = strings.Split(req.RemoteAddr, ":")[0]
	}
	// clientIp may contain multiple ips, use the first
	ips := strings.Split(clientIp, ",")
	ip := strings.TrimSpace(ips[0])
	// TODO: need a more robust way to determine when a non-fallback host looks
	// like a fallback.
	hasFallbackIp := isFallbackIp(ip)
	if !isFallback(name) && hasFallbackIp {
		log.Errorf("Found fallback ip %v for non-fallback host %v", ip, name)
		return ""
	} else if isFallback(name) && !hasFallbackIp {
		log.Errorf("Found non-fallback ip %v for fallback host %v", ip, name)
		return ""
	}
	return ip
}

func isFallbackIp(ip string) bool {
	for _, prefix := range fallbackIPPrefixes {
		if strings.HasPrefix(ip, prefix) {
			return true
		}
	}

	return false
}

var fallbackIPPrefixes = []string{
	"128.199",
	"178.62",
	"188.166",
}

func getSingleFormValue(req *http.Request, name string) string {
	ls := req.Form[name]
	if len(ls) == 0 {
		return ""
	}
	if len(ls) > 1 {
		// But we still allow it for robustness.
		log.Errorf("More than one '%v' provided in form: %v", name, ls)
	}
	return ls[0]
}
<commit_msg>add JP IPs seen so far<commit_after>// main simply contains the primary web serving code that allows peers to
// register and unregister as give mode peers running within the Lantern
// network
package main

import (
	"crypto/tls"
	"fmt"
	"net/http"
	"strconv"
	"strings"

	"github.com/getlantern/keyman"
	"github.com/getlantern/tlsdefaults"
)

const (
	PKFile   = "pk.pem"
	CertFile = "cert.pem"
)

const (
	cloudflareBit = 1 << iota
	cloudfrontBit = 1 << iota
)

func startHttp() {
	http.HandleFunc("/register", register)
	http.HandleFunc("/unregister", unregister)
	laddr := fmt.Sprintf(":%d", *port)

	tlsConfig := tlsdefaults.Server()
	_, _, err := keyman.StoredPKAndCert(PKFile, CertFile, "Lantern", "localhost")
	if err != nil {
		log.Fatalf("Unable to initialize private key and certificate: %v", err)
	}
	cert, err := tls.LoadX509KeyPair(CertFile, PKFile)
	if err != nil {
		log.Fatalf("Unable to load certificate and key from %s and %s: %s", CertFile, PKFile, err)
	}
	tlsConfig.Certificates = []tls.Certificate{cert}

	log.Debugf("About to listen at %v", laddr)
	l, err := tls.Listen("tcp", laddr, tlsConfig)
	if err != nil {
		log.Fatalf("Unable to listen for tls connections at %s: %s", laddr, err)
	}

	log.Debug("About to serve")
	err = http.Serve(l, nil)
	if err != nil {
		log.Fatalf("Unable to serve: %s", err)
	}
}

// register is the entry point for peers registering themselves with the service.
// If peers are successfully vetted, they'll be added to the DNS round robin.
func register(resp http.ResponseWriter, req *http.Request) {
	name, ip, port, supportedFronts, err := getHostInfo(req)
	if err == nil && !(port == "80" || port == "443") {
		err = fmt.Errorf("Port %d not supported, only ports 80 and 443 are supported", port)
	}
	if err != nil {
		resp.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(resp, err.Error())
		return
	}
	if isPeer(name) {
		log.Debugf("Not adding peer %v because we're not using peers at the moment", name)
		resp.WriteHeader(200)
		fmt.Fprintln(resp, "Peers disabled at the moment")
		return
	}
	online := true
	connectionRefused := false
	timedOut := false

	h := getOrCreateHost(name, ip, port)
	online, connectionRefused, timedOut = h.status()
	if online {
		resp.WriteHeader(200)
		fmt.Fprintln(resp, "Connectivity to proxy confirmed")
		if (supportedFronts & cloudfrontBit) == cloudfrontBit {
			h.initCloudfront()
		}
		fstr := "frontfqdns: {cloudflare: " + name + "." + *cfldomain
		if h.cfrDist != nil {
			fstr += ", cloudfront: " + h.cfrDist.Domain
		}
		fstr += "}"
		fmt.Fprintln(resp, fstr)

		return
	}
	if timedOut {
		log.Debugf("%v timed out waiting for status, returning 500 error", h)
		resp.WriteHeader(500)
		fmt.Fprintf(resp, "Timed out waiting for status")
		return
	}

	// Note this may not work across platforms, but the intent
	// is to tell the client if the connection was flat out
	// refused as opposed to timed out in order to allow them
	// to configure their router if possible.
	if connectionRefused {
		// 417 response code.
		resp.WriteHeader(http.StatusExpectationFailed)
		fmt.Fprintln(resp, "No connectivity to proxy - connection refused")
	} else {
		// 408 response code.
		resp.WriteHeader(http.StatusRequestTimeout)
		fmt.Fprintln(resp, "No connectivity to proxy - test request timed out")
	}
}

// unregister is the HTTP endpoint for removing peers from DNS. Peers are
// unregistered based on their ip (not their name).
func unregister(resp http.ResponseWriter, req *http.Request) {
	_, ip, _, _, err := getHostInfo(req)
	if err != nil {
		resp.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(resp, err.Error())
		return
	}

	h := getHostByIp(ip)
	msg := "Host not registered"
	if h != nil {
		h.unregister()
		msg = "Host unregistered"
	}
	resp.WriteHeader(200)
	fmt.Fprintln(resp, msg)
}

func getHostInfo(req *http.Request) (name string, ip string, port string, supportedFronts int, err error) {
	err = req.ParseForm()
	if err != nil {
		err = fmt.Errorf("Couldn't parse form: %v", err)
		return
	}
	name = getSingleFormValue(req, "name")
	if name == "" {
		err = fmt.Errorf("Please specify a name")
		return
	}
	ip = clientIpFor(req, name)
	if ip == "" {
		err = fmt.Errorf("Unable to determine IP address")
		return
	}
	port = getSingleFormValue(req, "port")
	if port != "" {
		_, err = strconv.Atoi(port)
		if err != nil {
			err = fmt.Errorf("Received invalid port for %v - %v: %v", name, ip, port)
			return
		}
	}
	fronts := req.Form["fronts"]
	if len(fronts) == 0 {
		// backwards compatibility
		fronts = []string{"cloudflare"}
	}
	for _, front := range fronts {
		switch front {
		case "cloudflare":
			supportedFronts |= cloudflareBit
		case "cloudfront":
			supportedFronts |= cloudfrontBit
		default:
			// Ignore these for forward compatibility.
			log.Debugf("Unrecognized front: %v", front)
		}
	}
	return
}

func clientIpFor(req *http.Request, name string) string {
	// Client requested their info
	clientIp := req.Header.Get("X-Peerscanner-Forwarded-For")
	if clientIp == "" {
		clientIp = req.Header.Get("X-Forwarded-For")
	}
	if clientIp == "" && isFallback(name) {
		// Use direct IP for fallbacks
		clientIp = strings.Split(req.RemoteAddr, ":")[0]
	}
	// clientIp may contain multiple ips, use the first
	ips := strings.Split(clientIp, ",")
	ip := strings.TrimSpace(ips[0])
	// TODO: need a more robust way to determine when a non-fallback host looks
	// like a fallback.
	hasFallbackIp := isFallbackIp(ip)
	if !isFallback(name) && hasFallbackIp {
		log.Errorf("Found fallback ip %v for non-fallback host %v", ip, name)
		return ""
	} else if isFallback(name) && !hasFallbackIp {
		log.Errorf("Found non-fallback ip %v for fallback host %v", ip, name)
		return ""
	}
	return ip
}

func isFallbackIp(ip string) bool {
	for _, prefix := range fallbackIPPrefixes {
		if strings.HasPrefix(ip, prefix) {
			return true
		}
	}

	return false
}

var fallbackIPPrefixes = []string{
	"128.199",
	"178.62",
	"188.166",
	"104.156.",
	"45.63.",
}

func getSingleFormValue(req *http.Request, name string) string {
	ls := req.Form[name]
	if len(ls) == 0 {
		return ""
	}
	if len(ls) > 1 {
		// But we still allow it for robustness.
		log.Errorf("More than one '%v' provided in form: %v", name, ls)
	}
	return ls[0]
}
<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
	"log"
	"os"

	"github.com/jlaffaye/ftp"
)

func retrieveFile(file string, serverID string) {
	fname := "files/" + file
	c, err := ftp.Connect(serverID)
	if err != nil {
		log.Fatal(err)
	}

	defer c.Quit()

	c.Login("test", "1234")
	if err != nil {
		log.Fatal(err)
	}
	defer c.Logout()
	r, err := c.Retr(fname)
	if err != nil {
		log.Fatal(err)
	}

	defer r.Close()

	buf, err := ioutil.ReadAll(r)
	if err != nil {
		log.Fatal(err)
	}

	err = ioutil.WriteFile(os.Args[2], buf, 0644)
	if err != nil {
		log.Fatal(err)
	}
}
func main() {
	serverID := "127.0.0.1:" + os.Args[1]
	fname := "files/" + os.Args[2]
	c, err := ftp.Connect(serverID)
	if err != nil {
		log.Fatal(err)
	}

	defer c.Quit()

	c.Login("test", "1234")
	if err != nil {
		log.Fatal(err)
	}
	defer c.Logout()
	r, err := c.Loc(fname)
	if err != nil {
		log.Fatal(err)
	}

	defer r.Close()

	buf, err := ioutil.ReadAll(r)
	if err != nil {
		log.Fatal(err)
	}

	// fmt.Println(string(buf))
	serverID = "127.0.0.1:" + string(buf)
	retrieveFile(os.Args[2], serverID)

}
<commit_msg>removing duplicate code<commit_after><|endoftext|>


<commit_before>package miniprofiler_gae

import (
	"appengine"
	"appengine/memcache"
	"appengine/user"
	"fmt"
	"github.com/mjibson/MiniProfiler/go/miniprofiler"
	"github.com/mjibson/appstats"
	"net/http"
)

func init() {
	miniprofiler.Enable = EnableIfAdminOrDev
	miniprofiler.Get = GetMemcache
	miniprofiler.Store = StoreMemcache
	miniprofiler.MachineName = Instance
}

func EnableIfAdminOrDev(r *http.Request) bool {
	if appengine.IsDevAppServer() {
		return true
	}
	c := appengine.NewContext(r)
	u := user.Current(c)
	return u.Admin
}

func Instance() string {
	if i := appengine.InstanceID(); i != "" {
		return i[len(i)-8:]
	}
	return miniprofiler.Hostname()
}

func StoreMemcache(r *http.Request, p *miniprofiler.Profile) {
	item := &memcache.Item{
		Key:   mp_key(string(p.Id)),
		Value: p.Json(),
	}
	c := appengine.NewContext(r)
	memcache.Set(c, item)
}

func GetMemcache(r *http.Request, id string) *miniprofiler.Profile {
	c := appengine.NewContext(r)
	item, err := memcache.Get(c, mp_key(id))
	if err != nil {
		return nil
	}
	return miniprofiler.ProfileFromJson(item.Value)
}

type Context struct {
	appstats.Context
	P *miniprofiler.Profile
}

func NewHandler(f func(Context, http.ResponseWriter, *http.Request)) appstats.Handler {
	return appstats.NewHandler(func(c appengine.Context, w http.ResponseWriter, r *http.Request) {
		pc := Context{
			Context: c.(appstats.Context),
		}

		if miniprofiler.Enabled(r) {
			pc.P = miniprofiler.NewProfile(w, r, miniprofiler.FuncName(f))
			f(pc, w, r)

			for _, v := range pc.Context.Stats.RPCStats {
				pc.P.Root.AddCustomTiming("RPC", &miniprofiler.CustomTiming{
					StartMilliseconds:    float64(v.Offset.Nanoseconds()) / 1000000,
					DurationMilliseconds: float64(v.Duration.Nanoseconds()) / 1000000,
				})
			}

			pc.P.CustomLink = pc.URL()
			pc.P.CustomLinkName = "appstats"
			pc.P.Finalize()
		} else {
			f(pc, w, r)
		}
	})
}

func mp_key(id string) string {
	return fmt.Sprintf("mini-profiler-results:%s", id)
}
<commit_msg>Attach RPCs to correct child<commit_after>package miniprofiler_gae

import (
	"appengine"
	"appengine/memcache"
	"appengine/user"
	"appengine_internal"
	"fmt"
	"github.com/mjibson/MiniProfiler/go/miniprofiler"
	"github.com/mjibson/appstats"
	"net/http"
)

func init() {
	miniprofiler.Enable = EnableIfAdminOrDev
	miniprofiler.Get = GetMemcache
	miniprofiler.Store = StoreMemcache
	miniprofiler.MachineName = Instance
}

func EnableIfAdminOrDev(r *http.Request) bool {
	if appengine.IsDevAppServer() {
		return true
	}
	c := appengine.NewContext(r)
	u := user.Current(c)
	return u.Admin
}

func Instance() string {
	if i := appengine.InstanceID(); i != "" {
		return i[len(i)-8:]
	}
	return miniprofiler.Hostname()
}

func StoreMemcache(r *http.Request, p *miniprofiler.Profile) {
	item := &memcache.Item{
		Key:   mp_key(string(p.Id)),
		Value: p.Json(),
	}
	c := appengine.NewContext(r)
	memcache.Set(c, item)
}

func GetMemcache(r *http.Request, id string) *miniprofiler.Profile {
	c := appengine.NewContext(r)
	item, err := memcache.Get(c, mp_key(id))
	if err != nil {
		return nil
	}
	return miniprofiler.ProfileFromJson(item.Value)
}

type Context struct {
	appstats.Context
	P *miniprofiler.Profile
}

func (c Context) Call(service, method string, in, out appengine_internal.ProtoMessage, opts *appengine_internal.CallOptions) error {
	err := c.Context.Call(service, method, in, out, opts)
	v := c.Context.Stats.RPCStats[len(c.Context.Stats.RPCStats)-1]
	c.P.AddCustomTiming("RPC", &miniprofiler.CustomTiming{
		StartMilliseconds:    float64(v.Offset.Nanoseconds()) / 1000000,
		DurationMilliseconds: float64(v.Duration.Nanoseconds()) / 1000000,
	})
	return err
}

func NewHandler(f func(Context, http.ResponseWriter, *http.Request)) appstats.Handler {
	return appstats.NewHandler(func(c appengine.Context, w http.ResponseWriter, r *http.Request) {
		pc := Context{
			Context: c.(appstats.Context),
		}

		if miniprofiler.Enabled(r) {
			pc.P = miniprofiler.NewProfile(w, r, miniprofiler.FuncName(f))
			f(pc, w, r)
			pc.P.CustomLink = pc.URL()
			pc.P.CustomLinkName = "appstats"
			pc.P.Finalize()
		} else {
			f(pc, w, r)
		}
	})
}

func mp_key(id string) string {
	return fmt.Sprintf("mini-profiler-results:%s", id)
}
<|endoftext|>


<commit_before>package bongo

import (
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"strings"

	"github.com/jinzhu/gorm"
)

// Fetch fetches the data from db by given parameters(fields of the struct)
func (b *Bongo) Fetch(i Modellable) error {
	if i.GetId() == 0 {
		return fmt.Errorf("Id is not set for %s", i.TableName())
	}

	if err := b.DB.Table(i.TableName()).
		Find(i).
		Error; err != nil {
		return err
	}

	return nil
}

// ById Fetches data from db by it's id
func (b *Bongo) ById(i Modellable, id int64) error {
	if err := b.DB.
		Table(i.TableName()).
		Where("id = ?", id).
		Find(i).
		Error; err != nil {
		return err
	}

	return nil
}

// Creates a new record with the given struct and its fields
func (b *Bongo) Create(i Modellable) error {
	if err := b.DB.Save(i).Error; err != nil {
		return err
	}

	return nil
}

// Update updates all fields of a struct with assigned data
func (b *Bongo) Update(i Modellable) error {
	if i.GetId() == 0 {
		return fmt.Errorf("id is not set for %s", i.TableName())
	}

	// Update and Create is using the Save method, so they are
	// same functions but GORM handles, AfterCreate and AfterUpdate
	// in correct manner
	return b.Create(i)
}

// Delete deletes the data by it's id, it doesnt take any other fields
// into consideration
func (b *Bongo) Delete(i Modellable) error {
	if i.GetId() == 0 {
		return fmt.Errorf("id is not set for %s", i.TableName())
	}

	if err := b.DB.Delete(i).Error; err != nil {
		return err
	}

	return nil
}

// fetches records by their ids and returns results in the same order
// as the ids; if no records in db we don't return error
func (b *Bongo) FetchByIds(i Modellable, data interface{}, ids []int64) error {
	if len(ids) == 0 {
		return nil
	}

	orderByQuery := ""
	comma := ""
	for _, id := range ids {
		orderByQuery = orderByQuery + comma + " id = " + strconv.FormatInt(id, 10) + " desc"
		comma = ","
	}

	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.TableName())

	query = query.Order(orderByQuery)

	query = query.Where(ids)

	query = query.Find(data)

	// supress not found errors
	return CheckErr(query)

}

func (b *Bongo) UpdatePartial(i Modellable, set map[string]interface{}) error {
	if i.GetId() == 0 {
		return fmt.Errorf("id is not set for %s", i.TableName())
	}

	// init query
	query := b.DB

	query = query.Table(i.TableName())

	query = query.Where(i.GetId())

	if err := query.Update(set).Error; err != nil {
		return err
	}

	if err := b.Fetch(i); err != nil {
		return err
	}

	b.AfterUpdate(i)
	return nil
}

// selector, set
func (b *Bongo) UpdateMulti(i Modellable, rest ...map[string]interface{}) error {
	var set, selector map[string]interface{}

	switch len(rest) {
	case 1:
		set = rest[0]
		selector = nil
	case 2:
		selector = rest[0]
		set = rest[1]
	default:
		return errors.New("update partial parameter list is wrong")
	}

	query := b.DB.Table(i.TableName())

	//add selector
	query = addWhere(query, selector)

	if err := query.Updates(set).Error; err != nil {
		return err
	}

	return nil
}

func (b *Bongo) Count(i Modellable, where ...interface{}) (int, error) {
	var count int

	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.TableName())

	// add query
	query = query.Where(where[0], where[1:len(where)]...)

	return count, query.Count(&count).Error
}

func (b *Bongo) CountWithQuery(i Modellable, q *Query) (int, error) {
	query := b.BuildQuery(i, q)
	var count int
	return count, query.Count(&count).Error
}

type Scope func(d *gorm.DB) *gorm.DB

type Query struct {
	Selector   map[string]interface{}
	Sort       map[string]string
	Pluck      string
	Pagination Pagination
	Scopes     []Scope
}

func (q *Query) AddScope(scope Scope) {
	if q.Scopes == nil {
		q.Scopes = make([]Scope, 0)
	}

	q.Scopes = append(q.Scopes, scope)
}

type Pagination struct {
	Limit int
	Skip  int
}

func NewPagination(limit int, skip int) *Pagination {
	return &Pagination{
		Limit: limit,
		Skip:  skip,
	}
}

func NewQS(selector map[string]interface{}) *Query {
	return &Query{
		Selector: selector,
	}
}

// selector, sort, limit, pluck,
func (b *Bongo) Some(i Modellable, data interface{}, q *Query) error {
	err := b.executeQuery(i, data, q)
	if err == gorm.RecordNotFound {
		return nil
	}
	return err
}

func (b *Bongo) One(i Modellable, data interface{}, q *Query) error {
	q.Pagination.Limit = 1
	return b.executeQuery(i, data, q)
}

func (b *Bongo) BuildQuery(i Modellable, q *Query) *gorm.DB {
	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.TableName())

	// add sort options
	query = addSort(query, q.Sort)

	query = addSkip(query, q.Pagination.Skip)

	query = addLimit(query, q.Pagination.Limit)

	// add selector
	query = addWhere(query, q.Selector)

	// put scopes
	if q.Scopes != nil && len(q.Scopes) > 0 {
		for _, scope := range q.Scopes {
			query = query.Scopes(scope)
		}
	}

	return query
}

func (b *Bongo) executeQuery(i Modellable, data interface{}, q *Query) error {
	// init query
	query := b.BuildQuery(i, q)

	var err error
	// TODO refactor this part
	if q.Pluck != "" {
		if strings.Contains(q.Pluck, ",") {
			// add pluck data
			query = addPluck(query, q.Pluck)

			err = query.Find(data).Error
		} else {
			err = query.Pluck(q.Pluck, data).Error
		}
	} else {
		err = query.Find(data).Error
	}

	return err
}

func (b *Bongo) PublishEvent(eventName string, i Modellable) error {
	data, err := json.Marshal(i)
	if err != nil {
		b.log.Error("Error while marshalling for publish %s", err)
		return err
	}

	err = b.Broker.Publish(i.TableName()+"_"+eventName, data)
	if err != nil {
		b.log.Error("Error while publishing %s", err)
		return err
	}

	return nil
}

func (b *Bongo) AfterCreate(i Modellable) {
	b.PublishEvent("created", i)
}

func (b *Bongo) AfterUpdate(i Modellable) {
	b.PublishEvent("updated", i)
}

func (b *Bongo) AfterDelete(i Modellable) {
	b.PublishEvent("deleted", i)
}

// addSort injects sort parameters into query
func addSort(query *gorm.DB, options map[string]string) *gorm.DB {

	if options == nil {
		return query
	}

	if len(options) == 0 {
		return query
	}

	var opts []string
	for key, val := range options {
		opts = append(opts, fmt.Sprintf("%s %v", key, val))
	}
	return query.Order(strings.Join(opts, ","))
}

// addPluck basically adds select statement for
// only required fields
func addPluck(query *gorm.DB, plucked string) *gorm.DB {
	if plucked == "" {
		return query
	}

	return query.Select(plucked)
}

// addWhere adds where query
func addWhere(query *gorm.DB, selector map[string]interface{}) *gorm.DB {
	if selector == nil {
		return query
	}

	// instead sending one selector, do chaining here
	return query.Where(selector)
}

// addSkip adds skip parameter into sql query
func addSkip(query *gorm.DB, skip int) *gorm.DB {
	if skip > 0 {
		return query.Offset(skip)
	}

	return query
}

// addLimit adds limit into query if set
func addLimit(query *gorm.DB, limit int) *gorm.DB {
	// if limit is minus or 0 ignore
	if limit > 0 {
		return query.Limit(limit)
	}

	return query
}

// CheckErr checks error exitence and returns
// if found, but this function suppress RecordNotFound errors
func CheckErr(res *gorm.DB) error {
	if res == nil {
		return nil
	}

	if res.Error == nil {
		return nil
	}

	if res.Error == gorm.RecordNotFound {
		return nil
	}

	return res.Error
}
<commit_msg>Social: minor fix<commit_after>package bongo

import (
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"strings"

	"github.com/jinzhu/gorm"
)

// Fetch fetches the data from db by given parameters(fields of the struct)
func (b *Bongo) Fetch(i Modellable) error {
	if i.GetId() == 0 {
		return fmt.Errorf("Id is not set for %s", i.TableName())
	}

	if err := b.DB.Table(i.TableName()).
		Find(i).
		Error; err != nil {
		return err
	}

	return nil
}

// ById Fetches data from db by it's id
func (b *Bongo) ById(i Modellable, id int64) error {
	if err := b.DB.
		Table(i.TableName()).
		Where("id = ?", id).
		Find(i).
		Error; err != nil {
		return err
	}

	return nil
}

// Creates a new record with the given struct and its fields
func (b *Bongo) Create(i Modellable) error {
	if err := b.DB.Save(i).Error; err != nil {
		return err
	}

	return nil
}

// Update updates all fields of a struct with assigned data
func (b *Bongo) Update(i Modellable) error {
	if i.GetId() == 0 {
		return fmt.Errorf("id is not set for %s", i.TableName())
	}

	// Update and Create is using the Save method, so they are
	// same functions but GORM handles, AfterCreate and AfterUpdate
	// in correct manner
	return b.Create(i)
}

// Delete deletes the data by it's id, it doesnt take any other fields
// into consideration
func (b *Bongo) Delete(i Modellable) error {
	if i.GetId() == 0 {
		return fmt.Errorf("id is not set for %s", i.TableName())
	}

	if err := b.DB.Delete(i).Error; err != nil {
		return err
	}

	return nil
}

// FetchByIds fetches records by their ids and returns results in the same order
// as the ids; if no records in db we don't return error
func (b *Bongo) FetchByIds(i Modellable, data interface{}, ids []int64) error {
	if len(ids) == 0 {
		return nil
	}

	orderByQuery := ""
	comma := ""
	for _, id := range ids {
		orderByQuery = orderByQuery + comma + " id = " + strconv.FormatInt(id, 10) + " desc"
		comma = ","
	}

	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.TableName())

	query = query.Order(orderByQuery)

	query = query.Where(ids)

	query = query.Find(data)

	// supress not found errors
	return CheckErr(query)

}

func (b *Bongo) UpdatePartial(i Modellable, set map[string]interface{}) error {
	if i.GetId() == 0 {
		return fmt.Errorf("id is not set for %s", i.TableName())
	}

	// init query
	query := b.DB

	query = query.Table(i.TableName())

	query = query.Where(i.GetId())

	if err := query.Update(set).Error; err != nil {
		return err
	}

	if err := b.Fetch(i); err != nil {
		return err
	}

	b.AfterUpdate(i)
	return nil
}

// selector, set
func (b *Bongo) UpdateMulti(i Modellable, rest ...map[string]interface{}) error {
	var set, selector map[string]interface{}

	switch len(rest) {
	case 1:
		set = rest[0]
		selector = nil
	case 2:
		selector = rest[0]
		set = rest[1]
	default:
		return errors.New("update partial parameter list is wrong")
	}

	query := b.DB.Table(i.TableName())

	//add selector
	query = addWhere(query, selector)

	if err := query.Updates(set).Error; err != nil {
		return err
	}

	return nil
}

func (b *Bongo) Count(i Modellable, where ...interface{}) (int, error) {
	var count int

	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.TableName())

	// add query
	query = query.Where(where[0], where[1:len(where)]...)

	return count, query.Count(&count).Error
}

func (b *Bongo) CountWithQuery(i Modellable, q *Query) (int, error) {
	query := b.BuildQuery(i, q)
	var count int
	return count, query.Count(&count).Error
}

type Scope func(d *gorm.DB) *gorm.DB

type Query struct {
	Selector   map[string]interface{}
	Sort       map[string]string
	Pluck      string
	Pagination Pagination
	Scopes     []Scope
}

func (q *Query) AddScope(scope Scope) {
	if q.Scopes == nil {
		q.Scopes = make([]Scope, 0)
	}

	q.Scopes = append(q.Scopes, scope)
}

type Pagination struct {
	Limit int
	Skip  int
}

func NewPagination(limit int, skip int) *Pagination {
	return &Pagination{
		Limit: limit,
		Skip:  skip,
	}
}

func NewQS(selector map[string]interface{}) *Query {
	return &Query{
		Selector: selector,
	}
}

// selector, sort, limit, pluck,
func (b *Bongo) Some(i Modellable, data interface{}, q *Query) error {
	err := b.executeQuery(i, data, q)
	if err == gorm.RecordNotFound {
		return nil
	}
	return err
}

func (b *Bongo) One(i Modellable, data interface{}, q *Query) error {
	q.Pagination.Limit = 1
	return b.executeQuery(i, data, q)
}

func (b *Bongo) BuildQuery(i Modellable, q *Query) *gorm.DB {
	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.TableName())

	// add sort options
	query = addSort(query, q.Sort)

	query = addSkip(query, q.Pagination.Skip)

	query = addLimit(query, q.Pagination.Limit)

	// add selector
	query = addWhere(query, q.Selector)

	// put scopes
	if q.Scopes != nil && len(q.Scopes) > 0 {
		for _, scope := range q.Scopes {
			query = query.Scopes(scope)
		}
	}

	return query
}

func (b *Bongo) executeQuery(i Modellable, data interface{}, q *Query) error {
	// init query
	query := b.BuildQuery(i, q)

	var err error
	// TODO refactor this part
	if q.Pluck != "" {
		if strings.Contains(q.Pluck, ",") {
			// add pluck data
			query = addPluck(query, q.Pluck)

			err = query.Find(data).Error
		} else {
			err = query.Pluck(q.Pluck, data).Error
		}
	} else {
		err = query.Find(data).Error
	}

	return err
}

func (b *Bongo) PublishEvent(eventName string, i Modellable) error {
	data, err := json.Marshal(i)
	if err != nil {
		b.log.Error("Error while marshalling for publish %s", err)
		return err
	}

	err = b.Broker.Publish(i.TableName()+"_"+eventName, data)
	if err != nil {
		b.log.Error("Error while publishing %s", err)
		return err
	}

	return nil
}

func (b *Bongo) AfterCreate(i Modellable) {
	b.PublishEvent("created", i)
}

func (b *Bongo) AfterUpdate(i Modellable) {
	b.PublishEvent("updated", i)
}

func (b *Bongo) AfterDelete(i Modellable) {
	b.PublishEvent("deleted", i)
}

// addSort injects sort parameters into query
func addSort(query *gorm.DB, options map[string]string) *gorm.DB {

	if options == nil {
		return query
	}

	if len(options) == 0 {
		return query
	}

	var opts []string
	for key, val := range options {
		opts = append(opts, fmt.Sprintf("%s %v", key, val))
	}
	return query.Order(strings.Join(opts, ","))
}

// addPluck basically adds select statement for
// only required fields
func addPluck(query *gorm.DB, plucked string) *gorm.DB {
	if plucked == "" {
		return query
	}

	return query.Select(plucked)
}

// addWhere adds where query
func addWhere(query *gorm.DB, selector map[string]interface{}) *gorm.DB {
	if selector == nil {
		return query
	}

	// instead sending one selector, do chaining here
	return query.Where(selector)
}

// addSkip adds skip parameter into sql query
func addSkip(query *gorm.DB, skip int) *gorm.DB {
	if skip > 0 {
		return query.Offset(skip)
	}

	return query
}

// addLimit adds limit into query if set
func addLimit(query *gorm.DB, limit int) *gorm.DB {
	// if limit is minus or 0 ignore
	if limit > 0 {
		return query.Limit(limit)
	}

	return query
}

// CheckErr checks error exitence and returns
// if found, but this function suppress RecordNotFound errors
func CheckErr(res *gorm.DB) error {
	if res == nil {
		return nil
	}

	if res.Error == nil {
		return nil
	}

	if res.Error == gorm.RecordNotFound {
		return nil
	}

	return res.Error
}
<|endoftext|>