<commit_before>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2019 Red Hat, Inc.
 *
 */

package tests_test

import (
	"fmt"
	"strings"
	"time"

	v1 "k8s.io/api/core/v1"
	"k8s.io/api/policy/v1beta1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"kubevirt.io/client-go/kubecli"
	"kubevirt.io/kubevirt/tests"
)

const (
	WaitSecondsBeforeDeploymentCheck     = 2 * time.Second
	DefaultStabilizationTimeoutInSeconds = 300
	DefaultPollIntervalInSeconds         = 3
	labelKey                             = "control-plane-test"
	labelValue                           = "selected"
)

var _ = Describe("[ref_id:2717]KubeVirt control plane resilience", func() {

	RegisterFailHandler(Fail)

	virtCli, err := kubecli.GetKubevirtClient()
	Expect(err).ToNot(HaveOccurred())
	deploymentsClient := virtCli.AppsV1().Deployments(tests.KubeVirtInstallNamespace)

	controlPlaneDeploymentNames := []string{"virt-api", "virt-controller"}

	Context("pod eviction", func() {

		var nodeNames []string
		var selectedNodeName string

		tests.FlagParse()

		getRunningReadyPods := func(podList *v1.PodList, podNames []string, nodeNames ...string) (pods []*v1.Pod) {
			pods = make([]*v1.Pod, 0)
			for _, pod := range podList.Items {
				if pod.Status.Phase != v1.PodRunning {
					continue
				}
				podReady := tests.PodReady(&pod)
				if podReady != v1.ConditionTrue {
					continue
				}
				for _, podName := range podNames {
					if strings.HasPrefix(pod.Name, podName) {
						if len(nodeNames) > 0 {
							for _, nodeName := range nodeNames {
								if pod.Spec.NodeName == nodeName {
									deepCopy := pod.DeepCopy()
									pods = append(pods, deepCopy)
								}
							}
						} else {
							deepCopy := pod.DeepCopy()
							pods = append(pods, deepCopy)
						}
					}
				}
			}
			return
		}

		getPodList := func() (podList *v1.PodList, err error) {
			podList, err = virtCli.CoreV1().Pods(tests.KubeVirtInstallNamespace).List(metav1.ListOptions{})
			return
		}

		getSelectedNode := func() (selectedNode *v1.Node, err error) {
			selectedNode, err = virtCli.CoreV1().Nodes().Get(selectedNodeName, metav1.GetOptions{})
			return
		}

		waitForDeploymentsToStabilize := func() (bool, error) {
			for _, deploymentName := range controlPlaneDeploymentNames {
				deployment, err := deploymentsClient.Get(deploymentName, metav1.GetOptions{})
				if err != nil {
					return false, err
				}

				if !(deployment.Status.UpdatedReplicas == *(deployment.Spec.Replicas) &&
					deployment.Status.Replicas == *(deployment.Spec.Replicas) &&
					deployment.Status.AvailableReplicas == *(deployment.Spec.Replicas)) {
					return false, err
				}
			}
			return true, nil
		}

		addLabelToSelectedNode := func() (bool, error) {
			selectedNode, err := getSelectedNode()
			if err != nil {
				return false, err
			}
			if selectedNode.Labels == nil {
				selectedNode.Labels = make(map[string]string)
			}
			selectedNode.Labels[labelKey] = labelValue
			_, err = virtCli.CoreV1().Nodes().Update(selectedNode)
			if err != nil {
				return false, fmt.Errorf("failed to update node: %v", err)
			}
			return true, nil
		}

		// Add nodeSelector to deployments so that they get scheduled to selectedNode
		addNodeSelectorToDeployments := func() (bool, error) {
			for _, deploymentName := range controlPlaneDeploymentNames {
				deployment, err := deploymentsClient.Get(deploymentName, metav1.GetOptions{})
				if err != nil {
					return false, err
				}

				labelMap := make(map[string]string)
				labelMap[labelKey] = labelValue
				if deployment.Spec.Template.Spec.NodeSelector == nil {
					deployment.Spec.Template.Spec.NodeSelector = make(map[string]string)
				}
				deployment.Spec.Template.Spec.NodeSelector[labelKey] = labelValue
				_, err = deploymentsClient.Update(deployment)
				if err != nil {
					return false, err
				}
			}
			return true, nil
		}

		checkControlPlanePodsHaveNodeSelector := func() (bool, error) {
			podList, err := getPodList()
			if err != nil {
				return false, err
			}
			runningControlPlanePods := getRunningReadyPods(podList, controlPlaneDeploymentNames)
			for _, pod := range runningControlPlanePods {
				if actualLabelValue, ok := pod.Spec.NodeSelector[labelKey]; ok {
					if actualLabelValue != labelValue {
						return false, fmt.Errorf("pod %s has node selector %s with value %s, expected was %s", pod.Name, labelKey, actualLabelValue, labelValue)
					}
				} else {
					return false, fmt.Errorf("pod %s has no node selector %s", pod.Name, labelKey)
				}
			}
			return true, nil
		}

		eventuallyWithTimeout := func(f func() (bool, error)) {
			Eventually(f,
				DefaultStabilizationTimeoutInSeconds, DefaultPollIntervalInSeconds,
			).Should(BeTrue())
		}

		setSelectedNodeUnschedulable := func() {
			selectedNode, err := getSelectedNode()
			Expect(err).ToNot(HaveOccurred())
			selectedNode.Spec.Unschedulable = true
			_, err = virtCli.CoreV1().Nodes().Update(selectedNode)
			Expect(err).ToNot(HaveOccurred())
		}

		BeforeEach(func() {
			tests.BeforeTestCleanup()

			nodes := tests.GetAllSchedulableNodes(virtCli).Items
			nodeNames = make([]string, len(nodes))
			for index, node := range nodes {
				nodeNames[index] = node.Name
			}

			// select one node from result for test, first node will do
			selectedNodeName = nodes[0].Name

			eventuallyWithTimeout(addLabelToSelectedNode)
			eventuallyWithTimeout(addNodeSelectorToDeployments)

			time.Sleep(WaitSecondsBeforeDeploymentCheck)

			eventuallyWithTimeout(checkControlPlanePodsHaveNodeSelector)
			eventuallyWithTimeout(waitForDeploymentsToStabilize)

			setSelectedNodeUnschedulable()
		})

		removeNodeSelectorFromDeployments := func() (bool, error) {
			for _, deploymentName := range controlPlaneDeploymentNames {
				deployment, err := deploymentsClient.Get(deploymentName, metav1.GetOptions{})
				if err != nil {
					return false, err
				}
				delete(deployment.Spec.Template.Spec.NodeSelector, labelKey)
				_, err = deploymentsClient.Update(deployment)
				if err != nil {
					return false, err
				}
			}
			return true, nil
		}

		// Clean up selectedNode: Remove label and make schedulable again
		cleanUpSelectedNode := func() (bool, error) {
			selectedNode, err := getSelectedNode()
			if err != nil {
				return false, err
			}
			selectedNode.Spec.Unschedulable = false
			delete(selectedNode.Labels, labelKey)
			_, err = virtCli.CoreV1().Nodes().Update(selectedNode)
			if err != nil {
				return false, err
			}
			return true, nil
		}

		checkControlPlanePodsDontHaveNodeSelector := func() (bool, error) {
			podList, err := getPodList()
			if err != nil {
				return false, err
			}
			runningControlPlanePods := getRunningReadyPods(podList, controlPlaneDeploymentNames)
			for _, pod := range runningControlPlanePods {
				if _, ok := pod.Spec.NodeSelector[labelKey]; ok {
					return false, fmt.Errorf("pod %s has still node selector %s", pod.Name, labelKey)
				}
			}
			return true, nil
		}

		AfterEach(func() {
			eventuallyWithTimeout(removeNodeSelectorFromDeployments)
			eventuallyWithTimeout(cleanUpSelectedNode)

			time.Sleep(WaitSecondsBeforeDeploymentCheck)

			eventuallyWithTimeout(checkControlPlanePodsDontHaveNodeSelector)
			eventuallyWithTimeout(waitForDeploymentsToStabilize)
		})

		When("evicting pods of control plane", func() {

			test := func(podName string) {
				By(fmt.Sprintf("Try to evict all pods %s from node %s\n", podName, selectedNodeName))
				podList, err := getPodList()
				Expect(err).ToNot(HaveOccurred())
				runningPods := getRunningReadyPods(podList, []string{podName})
				for index, pod := range runningPods {
					err = virtCli.CoreV1().Pods(tests.KubeVirtInstallNamespace).Evict(&v1beta1.Eviction{ObjectMeta: metav1.ObjectMeta{Name: pod.Name}})
					if index < len(runningPods)-1 {
						Expect(err).ToNot(HaveOccurred())
					}
				}
				Expect(err).To(HaveOccurred(), "no error occurred on evict of last pod")
			}

			It("[test_id:2830]last eviction should fail for virt-controller pods", func() { test("virt-controller") })
			It("[test_id:2799]last eviction should fail for virt-api pods", func() { test("virt-api") })

		})

	})

	Context("control plane components check", func() {

		When("control plane pods are running", func() {

			It("[test_id:2806]virt-controller and virt-api pods have a pod disruption budget", func() {

				By("check deployments")
				deployments, err := deploymentsClient.List(metav1.ListOptions{})
				Expect(err).ToNot(HaveOccurred())
				expectedDeployments := []string{"cdi-http-import-server", "virt-operator", "virt-api", "virt-controller"}
				for _, expectedDeployment := range expectedDeployments {
					found := false
					for _, deployment := range deployments.Items {
						if deployment.Name != expectedDeployment {
							continue
						}
						found = true
						break
					}
					if !found {
						Fail(fmt.Sprintf("deployment %s not found", expectedDeployment))
					}
				}

				By("check pod disruption budgets exist")
				podDisruptionBudgetList, err := virtCli.PolicyV1beta1().PodDisruptionBudgets(tests.KubeVirtInstallNamespace).List(metav1.ListOptions{})
				Expect(err).ToNot(HaveOccurred())
				for _, controlPlaneDeploymentName := range controlPlaneDeploymentNames {
					pdbName := controlPlaneDeploymentName + "-pdb"
					found := false
					for _, pdb := range podDisruptionBudgetList.Items {
						if pdb.Name != pdbName {
							continue
						}
						found = true
						break
					}
					if !found {
						Fail(fmt.Sprintf("pod disruption budget %s not found for control plane pod %s", pdbName, controlPlaneDeploymentName))
					}
				}
			})

		})

	})

})
<commit_msg>Remove "cdi-http-import-server" and "virt-operator" from expected deployments<commit_after>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2019 Red Hat, Inc.
 *
 */

package tests_test

import (
	"fmt"
	"strings"
	"time"

	v1 "k8s.io/api/core/v1"
	"k8s.io/api/policy/v1beta1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"kubevirt.io/client-go/kubecli"
	"kubevirt.io/kubevirt/tests"
)

const (
	WaitSecondsBeforeDeploymentCheck     = 2 * time.Second
	DefaultStabilizationTimeoutInSeconds = 300
	DefaultPollIntervalInSeconds         = 3
	labelKey                             = "control-plane-test"
	labelValue                           = "selected"
)

var _ = Describe("[ref_id:2717]KubeVirt control plane resilience", func() {

	RegisterFailHandler(Fail)

	virtCli, err := kubecli.GetKubevirtClient()
	Expect(err).ToNot(HaveOccurred())
	deploymentsClient := virtCli.AppsV1().Deployments(tests.KubeVirtInstallNamespace)

	controlPlaneDeploymentNames := []string{"virt-api", "virt-controller"}

	Context("pod eviction", func() {

		var nodeNames []string
		var selectedNodeName string

		tests.FlagParse()

		getRunningReadyPods := func(podList *v1.PodList, podNames []string, nodeNames ...string) (pods []*v1.Pod) {
			pods = make([]*v1.Pod, 0)
			for _, pod := range podList.Items {
				if pod.Status.Phase != v1.PodRunning {
					continue
				}
				podReady := tests.PodReady(&pod)
				if podReady != v1.ConditionTrue {
					continue
				}
				for _, podName := range podNames {
					if strings.HasPrefix(pod.Name, podName) {
						if len(nodeNames) > 0 {
							for _, nodeName := range nodeNames {
								if pod.Spec.NodeName == nodeName {
									deepCopy := pod.DeepCopy()
									pods = append(pods, deepCopy)
								}
							}
						} else {
							deepCopy := pod.DeepCopy()
							pods = append(pods, deepCopy)
						}
					}
				}
			}
			return
		}

		getPodList := func() (podList *v1.PodList, err error) {
			podList, err = virtCli.CoreV1().Pods(tests.KubeVirtInstallNamespace).List(metav1.ListOptions{})
			return
		}

		getSelectedNode := func() (selectedNode *v1.Node, err error) {
			selectedNode, err = virtCli.CoreV1().Nodes().Get(selectedNodeName, metav1.GetOptions{})
			return
		}

		waitForDeploymentsToStabilize := func() (bool, error) {
			for _, deploymentName := range controlPlaneDeploymentNames {
				deployment, err := deploymentsClient.Get(deploymentName, metav1.GetOptions{})
				if err != nil {
					return false, err
				}

				if !(deployment.Status.UpdatedReplicas == *(deployment.Spec.Replicas) &&
					deployment.Status.Replicas == *(deployment.Spec.Replicas) &&
					deployment.Status.AvailableReplicas == *(deployment.Spec.Replicas)) {
					return false, err
				}
			}
			return true, nil
		}

		addLabelToSelectedNode := func() (bool, error) {
			selectedNode, err := getSelectedNode()
			if err != nil {
				return false, err
			}
			if selectedNode.Labels == nil {
				selectedNode.Labels = make(map[string]string)
			}
			selectedNode.Labels[labelKey] = labelValue
			_, err = virtCli.CoreV1().Nodes().Update(selectedNode)
			if err != nil {
				return false, fmt.Errorf("failed to update node: %v", err)
			}
			return true, nil
		}

		// Add nodeSelector to deployments so that they get scheduled to selectedNode
		addNodeSelectorToDeployments := func() (bool, error) {
			for _, deploymentName := range controlPlaneDeploymentNames {
				deployment, err := deploymentsClient.Get(deploymentName, metav1.GetOptions{})
				if err != nil {
					return false, err
				}

				labelMap := make(map[string]string)
				labelMap[labelKey] = labelValue
				if deployment.Spec.Template.Spec.NodeSelector == nil {
					deployment.Spec.Template.Spec.NodeSelector = make(map[string]string)
				}
				deployment.Spec.Template.Spec.NodeSelector[labelKey] = labelValue
				_, err = deploymentsClient.Update(deployment)
				if err != nil {
					return false, err
				}
			}
			return true, nil
		}

		checkControlPlanePodsHaveNodeSelector := func() (bool, error) {
			podList, err := getPodList()
			if err != nil {
				return false, err
			}
			runningControlPlanePods := getRunningReadyPods(podList, controlPlaneDeploymentNames)
			for _, pod := range runningControlPlanePods {
				if actualLabelValue, ok := pod.Spec.NodeSelector[labelKey]; ok {
					if actualLabelValue != labelValue {
						return false, fmt.Errorf("pod %s has node selector %s with value %s, expected was %s", pod.Name, labelKey, actualLabelValue, labelValue)
					}
				} else {
					return false, fmt.Errorf("pod %s has no node selector %s", pod.Name, labelKey)
				}
			}
			return true, nil
		}

		eventuallyWithTimeout := func(f func() (bool, error)) {
			Eventually(f,
				DefaultStabilizationTimeoutInSeconds, DefaultPollIntervalInSeconds,
			).Should(BeTrue())
		}

		setSelectedNodeUnschedulable := func() {
			selectedNode, err := getSelectedNode()
			Expect(err).ToNot(HaveOccurred())
			selectedNode.Spec.Unschedulable = true
			_, err = virtCli.CoreV1().Nodes().Update(selectedNode)
			Expect(err).ToNot(HaveOccurred())
		}

		BeforeEach(func() {
			tests.BeforeTestCleanup()

			nodes := tests.GetAllSchedulableNodes(virtCli).Items
			nodeNames = make([]string, len(nodes))
			for index, node := range nodes {
				nodeNames[index] = node.Name
			}

			// select one node from result for test, first node will do
			selectedNodeName = nodes[0].Name

			eventuallyWithTimeout(addLabelToSelectedNode)
			eventuallyWithTimeout(addNodeSelectorToDeployments)

			time.Sleep(WaitSecondsBeforeDeploymentCheck)

			eventuallyWithTimeout(checkControlPlanePodsHaveNodeSelector)
			eventuallyWithTimeout(waitForDeploymentsToStabilize)

			setSelectedNodeUnschedulable()
		})

		removeNodeSelectorFromDeployments := func() (bool, error) {
			for _, deploymentName := range controlPlaneDeploymentNames {
				deployment, err := deploymentsClient.Get(deploymentName, metav1.GetOptions{})
				if err != nil {
					return false, err
				}
				delete(deployment.Spec.Template.Spec.NodeSelector, labelKey)
				_, err = deploymentsClient.Update(deployment)
				if err != nil {
					return false, err
				}
			}
			return true, nil
		}

		// Clean up selectedNode: Remove label and make schedulable again
		cleanUpSelectedNode := func() (bool, error) {
			selectedNode, err := getSelectedNode()
			if err != nil {
				return false, err
			}
			selectedNode.Spec.Unschedulable = false
			delete(selectedNode.Labels, labelKey)
			_, err = virtCli.CoreV1().Nodes().Update(selectedNode)
			if err != nil {
				return false, err
			}
			return true, nil
		}

		checkControlPlanePodsDontHaveNodeSelector := func() (bool, error) {
			podList, err := getPodList()
			if err != nil {
				return false, err
			}
			runningControlPlanePods := getRunningReadyPods(podList, controlPlaneDeploymentNames)
			for _, pod := range runningControlPlanePods {
				if _, ok := pod.Spec.NodeSelector[labelKey]; ok {
					return false, fmt.Errorf("pod %s has still node selector %s", pod.Name, labelKey)
				}
			}
			return true, nil
		}

		AfterEach(func() {
			eventuallyWithTimeout(removeNodeSelectorFromDeployments)
			eventuallyWithTimeout(cleanUpSelectedNode)

			time.Sleep(WaitSecondsBeforeDeploymentCheck)

			eventuallyWithTimeout(checkControlPlanePodsDontHaveNodeSelector)
			eventuallyWithTimeout(waitForDeploymentsToStabilize)
		})

		When("evicting pods of control plane", func() {

			test := func(podName string) {
				By(fmt.Sprintf("Try to evict all pods %s from node %s\n", podName, selectedNodeName))
				podList, err := getPodList()
				Expect(err).ToNot(HaveOccurred())
				runningPods := getRunningReadyPods(podList, []string{podName})
				for index, pod := range runningPods {
					err = virtCli.CoreV1().Pods(tests.KubeVirtInstallNamespace).Evict(&v1beta1.Eviction{ObjectMeta: metav1.ObjectMeta{Name: pod.Name}})
					if index < len(runningPods)-1 {
						Expect(err).ToNot(HaveOccurred())
					}
				}
				Expect(err).To(HaveOccurred(), "no error occurred on evict of last pod")
			}

			It("[test_id:2830]last eviction should fail for virt-controller pods", func() { test("virt-controller") })
			It("[test_id:2799]last eviction should fail for virt-api pods", func() { test("virt-api") })

		})

	})

	Context("control plane components check", func() {

		When("control plane pods are running", func() {

			It("[test_id:2806]virt-controller and virt-api pods have a pod disruption budget", func() {

				By("check deployments")
				deployments, err := deploymentsClient.List(metav1.ListOptions{})
				Expect(err).ToNot(HaveOccurred())
				expectedDeployments := []string{"virt-api", "virt-controller"}
				for _, expectedDeployment := range expectedDeployments {
					found := false
					for _, deployment := range deployments.Items {
						if deployment.Name != expectedDeployment {
							continue
						}
						found = true
						break
					}
					if !found {
						Fail(fmt.Sprintf("deployment %s not found", expectedDeployment))
					}
				}

				By("check pod disruption budgets exist")
				podDisruptionBudgetList, err := virtCli.PolicyV1beta1().PodDisruptionBudgets(tests.KubeVirtInstallNamespace).List(metav1.ListOptions{})
				Expect(err).ToNot(HaveOccurred())
				for _, controlPlaneDeploymentName := range controlPlaneDeploymentNames {
					pdbName := controlPlaneDeploymentName + "-pdb"
					found := false
					for _, pdb := range podDisruptionBudgetList.Items {
						if pdb.Name != pdbName {
							continue
						}
						found = true
						break
					}
					if !found {
						Fail(fmt.Sprintf("pod disruption budget %s not found for control plane pod %s", pdbName, controlPlaneDeploymentName))
					}
				}
			})

		})

	})

})
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
  "os/exec"
  "strings"
)

func main() {
  oid := latest()
  for _, filename := range changed(oid) {
    fmt.Println(filename)
  }
  os.Exit(1)
}

func changed(oid string) []string {
  lines := simpleExec("git", "diff", "--cached", "--name-only", oid)
  return strings.Split(lines, "\n")
}

func latest() string {
  if oid := simpleExec("git", "rev-parse", "--verify", "HEAD"); oid != "" {
    return oid
  }

  // Initial commit: diff against an empty tree object
  return "4b825dc642cb6eb9a060e54bf8d69288fbee4904"
}

func simpleExec(name string, arg ...string) string {
	output, err := exec.Command(name, arg...).Output()
  if _, ok := err.(*exec.ExitError); ok {
    return ""
  } else if err != nil {
    fmt.Printf("error running: %s %s\n", name, arg)
		panic(err)
	}

	return strings.Trim(string(output), " \n")
}
<commit_msg>really basic pre-commit hook that rejects commits with files > 5MB<commit_after>package main

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

const sizelimit = 5 * 1024 * 1024

func main() {
	wd, err := os.Getwd()
	if err != nil {
		panic(err)
	}

	oid := latest()
	bad := make(map[string]int64)
	for _, filename := range changed(oid) {
		check(wd, filename, bad)
	}

	if numbad := len(bad); numbad > 0 {
		fmt.Printf("%d bad file(s):\n", numbad)
		for name, size := range bad {
			fmt.Printf("%s %d\n", name, size)
		}
	}
}

func check(working, filename string, bad map[string]int64) {
	full := filepath.Join(working, filename)
	stat, err := os.Lstat(full)
	if err != nil {
		fmt.Println(err.Error())
		os.Exit(1)
	}

	if filesize := stat.Size(); filesize > sizelimit {
		bad[filename] = filesize
	}
}

func changed(oid string) []string {
	lines := simpleExec("git", "diff", "--cached", "--name-only", oid)
	return strings.Split(lines, "\n")
}

func latest() string {
	if oid := simpleExec("git", "rev-parse", "--verify", "HEAD"); oid != "" {
		return oid
	}

	// Initial commit: diff against an empty tree object
	return "4b825dc642cb6eb9a060e54bf8d69288fbee4904"
}

func simpleExec(name string, arg ...string) string {
	output, err := exec.Command(name, arg...).Output()
	if _, ok := err.(*exec.ExitError); ok {
		return ""
	} else if err != nil {
		fmt.Printf("error running: %s %s\n", name, arg)
		panic(err)
	}

	return strings.Trim(string(output), " \n")
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/tarm/goserial"
	"log"
	"io"
	"os"
	"encoding/json"
	"time"
)

var configFilePath string = "config.json"

var settings struct {
	PortName string `json:"portName"`
	BaudRate int `json:"baudRate"`
	PlayerType int `json:"playerType"`
	PollRateSeconds time.Duration `json:"pollRateSeconds"`
}

func main() {
	//banner print
	displayBanner()

	//load config jSON into struct
	loadConfig()

	//verify port name and baud rate set before continuing
	if settings.PortName != "" && settings.BaudRate > 0 {
		//config & open serial connection
		serialConf := &serial.Config{Name: settings.PortName, Baud: settings.BaudRate}

		ser, err := serial.OpenPort(serialConf)

		if err != nil {
			endEarly("Could not connect to serial port. Cannot continue.", err.Error())
		} else {
			infoMessage("Connected to serial port successfully.")
		}

		//start our X seconds timer
		startTicker(ser)
	} else {
		//something is not configured, back out.
		endEarly("No baudRate and / or portName specified in config file. Cannot continue.", "")
	}
}

//kicks off our ticker, fires the elapsed once to start
func startTicker(ser io.ReadWriteCloser) {
	tickerElapsed(ser)

	ticker := time.NewTicker(time.Second * settings.PollRateSeconds)
    for _ = range ticker.C {
    	tickerElapsed(ser)
	}
}

//determines which player to poll, then acts
func tickerElapsed(ser io.ReadWriteCloser) {
	infoMessage("Test timer elapsed")
}

//loads a jSON config file, parses it into a struct
func loadConfig() {
	configFile, err := os.Open(configFilePath)

	if err != nil {
		endEarly("Couldn't open config file. Cannot continue.", err.Error())
	} else{
		infoMessage("Opened config file successfully.")
	}

	jsonParser := json.NewDecoder(configFile)

	if err = jsonParser.Decode(&settings); err != nil {
		endEarly("Couldn't parse config file. Cannot continue.", err.Error())
	} else {
		infoMessage("Parsed config file successfully.")
	}
}

//prints a fatal error message, then waits for a keypress prior to exit
func endEarly(msg string, err string) {
	log.Fatal(msg, " FULL ERROR: ", err, "\n\nExiting immediately.\n\n")
}

//logs a standard string message
func infoMessage(msg string) {
	log.Print(msg)
}

//sends string data to the serial port passed in
func sendToSerial(ser io.ReadWriteCloser, msg string) {
	ser.Write([]byte(msg))
}

//prints a welcome banner to the start of the app
func displayBanner() {
	fmt.Println("\n    _   __                 ____  __            _            ")
	fmt.Println("   / | / /___ _      __   / __ \\/ /___ ___  __(_)___  ____ _")
	fmt.Println("  /  |/ / __ \\ | /| / /  / /_/ / / __ `/ / / / / __ \\/ __ `/")
	fmt.Println(" / /|  / /_/ / |/ |/ /  / ____/ / /_/ / /_/ / / / / / /_/ / ")
	fmt.Println("/_/ |_/\\____/|__/|__/  /_/   /_/\\__,_/\\__, /_/_/ /_/\\__, /  ")
	fmt.Println("   __           _____           _    /____/        /____/   ")
	fmt.Println("  / /_____     / ___/___  _____(_)___ _/ /                  ")
	fmt.Println(" / __/ __ \\    \\__ \\/ _ \\/ ___/ / __ `/ /                   ")
	fmt.Println("/ /_/ /_/ /   ___/ /  __/ /  / / /_/ / /                    ")
	fmt.Println("\\__/\\____/   /____/\\___/_/  /_/\\__,_/_/                     \n")
}
<commit_msg>Basic web requests<commit_after>package main

import (
	"fmt"
	"github.com/tarm/goserial"
	"log"
	"io"
	"io/ioutil"
	"os"
	"encoding/json"
	"time"
	"net/http"
)

var configFilePath string = "config.json"

var settings struct {
	PortName string `json:"portName"`
	BaudRate int `json:"baudRate"`
	PlayerType int `json:"playerType"`
	PollRateSeconds time.Duration `json:"pollRateSeconds"`
	VlcWebUrl string `json:"vlcWebUrl"`
	VlcWebPassword string `json:"vlcWebPassword"`
}

func main() {
	//banner print
	displayBanner()

	//load config jSON into struct
	loadConfig()

	//verify port name and baud rate set before continuing
	if settings.PortName != "" && settings.BaudRate > 0 {
		//config & open serial connection
		serialConf := &serial.Config{Name: settings.PortName, Baud: settings.BaudRate}

		ser, err := serial.OpenPort(serialConf)

		if err != nil {
			endEarly("Could not connect to serial port. Cannot continue.", err.Error())
		} else {
			infoMessage("Connected to serial port successfully.")
		}

		//start our X seconds timer
		startTicker(ser)
	} else {
		//something is not configured, back out.
		endEarly("No baudRate and / or portName specified in config file. Cannot continue.", "")
	}
}

//kicks off our ticker, fires the elapsed once to start
func startTicker(ser io.ReadWriteCloser) {
	infoMessage("Ticker initialized.")

	tickerElapsed(ser)

	ticker := time.NewTicker(time.Second * settings.PollRateSeconds)
	for _ = range ticker.C {
		tickerElapsed(ser)
	}
}

//determines which player to poll, then acts
func tickerElapsed(ser io.ReadWriteCloser) {
	if settings.PlayerType == 1 {
		pollVlc(ser)
	} else {
		pollSpotify(ser)
	}
}

//poll VLC for now playing info
func pollVlc(ser io.ReadWriteCloser) {
	content, err := getResponseContent(settings.VlcWebUrl, settings.VlcWebPassword)

	if err != nil {
		infoMessage("Unable to get data from VLC. Double check your url and password.")
		infoMessage(err.Error())
	} else {
		infoMessage(string(content))
	}
}

//poll Spotify for now playing info
func pollSpotify(ser io.ReadWriteCloser) {
	infoMessage("Spotify polling not implemented yet.")
}

// Code adapted from http://www.codingcookies.com/2013/03/21/consuming-json-apis-with-go/
// This function fetch the content of a URL will return it as an
// array of bytes if retrieved successfully.
func getResponseContent(url string, password string) ([]byte, error) {
	// Build the request
	req, err := http.NewRequest("GET", url, nil)

	if password != "" {
		req.SetBasicAuth("", password)
	}

	if err != nil {
		return nil, err
	}

	// Send the request via a client
	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	// Defer the closing of the body
	defer resp.Body.Close()

	// Read the content into a byte array
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	// At this point we're done - simply return the bytes
	return body, nil
}

//loads a jSON config file, parses it into a struct
func loadConfig() {
	configFile, err := os.Open(configFilePath)

	if err != nil {
		endEarly("Couldn't open config file. Cannot continue.", err.Error())
	} else{
		infoMessage("Opened config file successfully.")
	}

	jsonParser := json.NewDecoder(configFile)

	if err = jsonParser.Decode(&settings); err != nil {
		endEarly("Couldn't parse config file. Cannot continue.", err.Error())
	} else {
		infoMessage("Parsed config file successfully.")
	}
}

//prints a fatal error message which in turn exits immediately
func endEarly(msg string, err string) {
	log.Fatal(msg, " FULL ERROR: ", err, "\n\nExiting immediately.\n\n")
}

//logs a standard string message
func infoMessage(msg string) {
	log.Print(msg)
}

//sends string data to the serial port passed in
func sendToSerial(ser io.ReadWriteCloser, msg string) {
	ser.Write([]byte(msg))

	infoMessage(msg)
}

//prints a welcome banner to the start of the app
func displayBanner() {
	fmt.Println("\n    _   __                 ____  __            _            ")
	fmt.Println("   / | / /___ _      __   / __ \\/ /___ ___  __(_)___  ____ _")
	fmt.Println("  /  |/ / __ \\ | /| / /  / /_/ / / __ `/ / / / / __ \\/ __ `/")
	fmt.Println(" / /|  / /_/ / |/ |/ /  / ____/ / /_/ / /_/ / / / / / /_/ / ")
	fmt.Println("/_/ |_/\\____/|__/|__/  /_/   /_/\\__,_/\\__, /_/_/ /_/\\__, /  ")
	fmt.Println("   __           _____           _    /____/        /____/   ")
	fmt.Println("  / /_____     / ___/___  _____(_)___ _/ /                  ")
	fmt.Println(" / __/ __ \\    \\__ \\/ _ \\/ ___/ / __ `/ /                   ")
	fmt.Println("/ /_/ /_/ /   ___/ /  __/ /  / / /_/ / /                    ")
	fmt.Println("\\__/\\____/   /____/\\___/_/  /_/\\__,_/_/                     \n")
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"strings"
	"time"
)

// WIDTH represents number of char allocated for time zone name.
const WIDTH = 16

// srv represents time server connection with chan for reading time.
type srv struct {
	name string
	addr string
	time chan string
	conn net.Conn
}

// servers represents list of time srv and methods to with all of them.
type servers []*srv

func (s servers) printTitle() {
	var buf bytes.Buffer

	for _, ts := range s {
		name := ts.name
		if len(name) > WIDTH {
			name = fmt.Sprintf("%s...", name[:WIDTH-3])
		}
		buf.WriteString(fmt.Sprintf("%*s|", WIDTH, name))
	}

	// Don't care about non-ascii names.
	rowLen := buf.Len()
	buf.WriteRune('\n')

	plusIndex := WIDTH
	for i := 0; i < rowLen; i++ {
		if i == plusIndex {
			buf.WriteRune('+')
			plusIndex += (WIDTH + 1)
		} else {
			buf.WriteRune('-')
		}
	}
	fmt.Println(buf.String())
}

func (s servers) printTime(sleep time.Duration) {
	for {
		time.Sleep(sleep)

		var down int
		var time string

		for _, ts := range s {
			t, ok := <-ts.time
			if !ok {
				t = "DISCONNECTED"
				down++
			}
			time += fmt.Sprintf("%*s|", WIDTH, t)
		}
		fmt.Printf("\r%s", time)
		if down == len(s) {
			fmt.Fprintln(os.Stderr, "\nall time servers are down! exiting...")
			os.Exit(1)
		}

	}
}

func (s servers) startFetching() {
	for _, ts := range s {
		go func(server *srv) {
			defer server.conn.Close()
			defer close(server.time)

			reader := bufio.NewReader(server.conn)

			for {
				line, _, err := reader.ReadLine()
				if err == io.EOF {
					break
				}
				if err != nil {
					log.Fatal(err)
				}
				server.time <- string(line)
			}
		}(ts)
	}
}

func (s servers) dialAll() {
	for _, ts := range s {
		conn, err := net.Dial("tcp", ts.addr)
		if err != nil {
			log.Fatal(err)
		}
		ts.conn = conn
	}
}

func main() {
	timeServers := make(servers, 0, len(os.Args)-1)

	for _, param := range os.Args[1:] {
		args := strings.Split(param, "=")
		timeServers = append(timeServers,
			&srv{name: args[0], addr: args[1], time: make(chan string)})
	}

	timeServers.dialAll()
	timeServers.startFetching()
	timeServers.printTitle()
	timeServers.printTime(time.Second)
}
<commit_msg>[8.2] remove temp variable name<commit_after>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"strings"
	"time"
)

// WIDTH represents number of char allocated for time zone name.
const WIDTH = 16

// srv represents time server connection with chan for reading time.
type srv struct {
	name string
	addr string
	time chan string
	conn net.Conn
}

// servers represents list of time srv and methods to with all of them.
type servers []*srv

func (s servers) printTitle() {
	var buf bytes.Buffer

	for _, ts := range s {
		if len(ts.name) > WIDTH {
			name = fmt.Sprintf("%s...", ts.name[:WIDTH-3])
		}
		buf.WriteString(fmt.Sprintf("%*s|", WIDTH, ts.name))
	}

	// Don't care about non-ascii names.
	rowLen := buf.Len()
	buf.WriteRune('\n')

	plusIndex := WIDTH
	for i := 0; i < rowLen; i++ {
		if i == plusIndex {
			buf.WriteRune('+')
			plusIndex += (WIDTH + 1)
		} else {
			buf.WriteRune('-')
		}
	}
	fmt.Println(buf.String())
}

func (s servers) printTime(sleep time.Duration) {
	for {
		time.Sleep(sleep)

		var down int
		var time string

		for _, ts := range s {
			t, ok := <-ts.time
			if !ok {
				t = "DISCONNECTED"
				down++
			}
			time += fmt.Sprintf("%*s|", WIDTH, t)
		}
		fmt.Printf("\r%s", time)
		if down == len(s) {
			fmt.Fprintln(os.Stderr, "\nall time servers are down! exiting...")
			os.Exit(1)
		}

	}
}

func (s servers) startFetching() {
	for _, ts := range s {
		go func(server *srv) {
			defer server.conn.Close()
			defer close(server.time)

			reader := bufio.NewReader(server.conn)

			for {
				line, _, err := reader.ReadLine()
				if err == io.EOF {
					break
				}
				if err != nil {
					log.Fatal(err)
				}
				server.time <- string(line)
			}
		}(ts)
	}
}

func (s servers) dialAll() {
	for _, ts := range s {
		conn, err := net.Dial("tcp", ts.addr)
		if err != nil {
			log.Fatal(err)
		}
		ts.conn = conn
	}
}

func main() {
	timeServers := make(servers, 0, len(os.Args)-1)

	for _, param := range os.Args[1:] {
		args := strings.Split(param, "=")
		timeServers = append(timeServers,
			&srv{name: args[0], addr: args[1], time: make(chan string)})
	}

	timeServers.dialAll()
	timeServers.startFetching()
	timeServers.printTitle()
	timeServers.printTime(time.Second)
}
<|endoftext|>


<commit_before>package schedule

import (
	"github.com/0xfoo/punchcard/git"
	"io/ioutil"
	"math/rand"
	"strconv"
	"strings"
	"time"
)

// RandomSchedule creates random commits over the past 365/366 days.
// These commits will be created in the location specified in the command.
func RandomSchedule(min, max int, location string) {
	git.Init(location)
	days := getDaysSinceDateMinusOneYear(time.Now())
	for day := range days {
		rnd := getRandomNumber(min, max)
		commits := RandomCommits(day, rnd)
		for commit := range commits {
			filename := strconv.Itoa(time.Now().Nanosecond())
			git.Add(location, filename)
			git.Commit(location, commit.message, commit.dateTime.String())
		}
	}
}

// RandomCommits returns a channel of random commits for a given day.
func RandomCommits(day time.Time, rnd int) chan Commit {
	commitChannel := make(chan Commit)
	go func() {
		for i := 0; i < rnd; i++ {
			commitChannel <- Commit{
				dateTime: getRandomTime(day),
				message:  getRandomCommitMessage(8),
			}
		}
		close(commitChannel)
	}()
	return commitChannel
}

// getRandomTime sets a random time on the given date.
func getRandomTime(day time.Time) time.Time {
	hours := time.Duration(getRandomNumber(0, 23)) * time.Hour
	minutes := time.Duration(getRandomNumber(0, 59)) * time.Minute
	seconds := time.Duration(getRandomNumber(0, 59)) * time.Second
	return day.Add(hours + minutes + seconds)
}

func getRandomCommitMessage(length int) string {
	content, _ := ioutil.ReadFile(COMMIT_MESSAGE_BASE)
	words := strings.Split(string(content), " ")
	return getRandomWords(words, getRandomNumber(1, length))
}

// getRandomNumber returns a number in the range of min and max.
func getRandomNumber(min, max int) int {
	if min == max {
		return min
	}
	return rand.Intn(max-min) + min
}

func getRandomWords(inWords []string, numWords int) string {
	outWords := make([]string, numWords)
	for i := 0; i < numWords; i++ {
		outWords = append(outWords, inWords[getRandomNumber(0, len(inWords))])
	}
	return strings.Join(outWords, " ")
}
<commit_msg>Move file creation into its own function.<commit_after>package schedule

import (
	"github.com/0xfoo/punchcard/git"
	"io/ioutil"
	"math/rand"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

// RandomSchedule creates random commits over the past 365/366 days.
// These commits will be created in the location specified in the command.
func RandomSchedule(min, max int, location string) {
	git.Init(location)
	days := getDaysSinceDateMinusOneYear(time.Now())
	for day := range days {
		rnd := getRandomNumber(min, max)
		commits := RandomCommits(day, rnd)
		for commit := range commits {
			filename := createFileInDir(location)
			git.Add(location, filename)
			git.Commit(location, commit.message, commit.dateTime.String())
		}
	}
}

// RandomCommits returns a channel of random commits for a given day.
func RandomCommits(day time.Time, rnd int) chan Commit {
	commitChannel := make(chan Commit)
	go func() {
		for i := 0; i < rnd; i++ {
			commitChannel <- Commit{
				dateTime: getRandomTime(day),
				message:  getRandomCommitMessage(8),
			}
		}
		close(commitChannel)
	}()
	return commitChannel
}

// getRandomTime sets a random time on the given date.
func getRandomTime(day time.Time) time.Time {
	hours := time.Duration(getRandomNumber(0, 23)) * time.Hour
	minutes := time.Duration(getRandomNumber(0, 59)) * time.Minute
	seconds := time.Duration(getRandomNumber(0, 59)) * time.Second
	return day.Add(hours + minutes + seconds)
}

// getRandomCommitMessage returns a commit message, no longer than length
func getRandomCommitMessage(length int) string {
	content, _ := ioutil.ReadFile(COMMIT_MESSAGE_BASE)
	words := strings.Split(string(content), " ")
	return getRandomWords(words, getRandomNumber(1, length))
}

// getRandomNumber returns a number in the range of min and max.
func getRandomNumber(min, max int) int {
	if min == max {
		return min
	}
	return rand.Intn(max-min) + min
}

// getRandomWords returns numWords random elements of the input []string
func getRandomWords(inWords []string, numWords int) string {
	outWords := make([]string, numWords)
	for i := 0; i < numWords; i++ {
		outWords = append(outWords, inWords[getRandomNumber(0, len(inWords))])
	}
	return strings.TrimSpace(strings.Join(outWords, " "))
}

// createFileWithTimeStamp creates a file with the current nano seconds as the
// file name, and returns this time stamp (i.e. filename)
func createFileInDir(dir string) string {
	filename := strconv.Itoa(time.Now().Nanosecond())
	file, _ := os.Create(filepath.Join(dir, filename))
	file.Close()
	return filename
}
<|endoftext|>


<commit_before>/* Licensed to the Apache Software Foundation (ASF) under one or more
contributor license agreements.  See the NOTICE file distributed with
this work for additional information regarding copyright ownership.
The ASF licenses this file to You 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 avro

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

	avro "github.com/elodina/go-avro"
)

const (
	GET_SCHEMA_BY_ID             = "/schemas/ids/%d"
	GET_SUBJECTS                 = "/subjects"
	GET_SUBJECT_VERSIONS         = "/subjects/%s/versions"
	GET_SPECIFIC_SUBJECT_VERSION = "/subjects/%s/versions/%s"
	REGISTER_NEW_SCHEMA          = "/subjects/%s/versions"
	CHECK_IS_REGISTERED          = "/subjects/%s"
	TEST_COMPATIBILITY           = "/compatibility/subjects/%s/versions/%s"
	CONFIG                       = "/config"
)

type SchemaRegistryClient interface {
	Register(subject string, schema avro.Schema) (int32, error)
	GetByID(id int32) (avro.Schema, error)
	GetLatestSchemaMetadata(subject string) (*SchemaMetadata, error)
	GetVersion(subject string, schema avro.Schema) (int32, error)
}

type SchemaMetadata struct {
	Id      int32
	Version int32
	Schema  string
}

type CompatibilityLevel string

const (
	BackwardCompatibilityLevel CompatibilityLevel = "BACKWARD"
	ForwardCompatibilityLevel  CompatibilityLevel = "FORWARD"
	FullCompatibilityLevel     CompatibilityLevel = "FULL"
	NoneCompatibilityLevel     CompatibilityLevel = "NONE"
)

const (
	SCHEMA_REGISTRY_V1_JSON               = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_V1_JSON_WEIGHTED      = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_MOST_SPECIFIC_DEFAULT = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_DEFAULT_JSON          = "application/vnd.schemaregistry+json"
	SCHEMA_REGISTRY_DEFAULT_JSON_WEIGHTED = "application/vnd.schemaregistry+json qs=0.9"
	JSON                                  = "application/json"
	JSON_WEIGHTED                         = "application/json qs=0.5"
	GENERIC_REQUEST                       = "application/octet-stream"
)

var PREFERRED_RESPONSE_TYPES = []string{SCHEMA_REGISTRY_V1_JSON, SCHEMA_REGISTRY_DEFAULT_JSON, JSON}

type ErrorMessage struct {
	Error_code int32
	Message    string
}

func (this *ErrorMessage) Error() string {
	return fmt.Sprintf("%s(error code: %d)", this.Message, this.Error_code)
}

type RegisterSchemaResponse struct {
	Id int32
}

type GetSchemaResponse struct {
	Schema string
}

type GetSubjectVersionResponse struct {
	Subject string
	Version int32
	Id      int32
	Schema  string
}

type CachedSchemaRegistryClient struct {
	registryURL  string
	schemaCache  map[string]map[avro.Schema]int32
	idCache      map[int32]avro.Schema
	versionCache map[string]map[avro.Schema]int32
}

func NewCachedSchemaRegistryClient(registryURL string) *CachedSchemaRegistryClient {
	return &CachedSchemaRegistryClient{
		registryURL:  registryURL,
		schemaCache:  make(map[string]map[avro.Schema]int32),
		idCache:      make(map[int32]avro.Schema),
		versionCache: make(map[string]map[avro.Schema]int32),
	}
}

func (this *CachedSchemaRegistryClient) Register(subject string, schema avro.Schema) (int32, error) {
	var schemaIdMap map[avro.Schema]int32
	var exists bool
	if schemaIdMap, exists = this.schemaCache[subject]; !exists {
		schemaIdMap = make(map[avro.Schema]int32)
		this.schemaCache[subject] = schemaIdMap
	}

	var id int32
	if id, exists = schemaIdMap[schema]; exists {
		return id, nil
	}

	request, err := this.newDefaultRequest("POST",
		fmt.Sprintf(REGISTER_NEW_SCHEMA, subject),
		strings.NewReader(fmt.Sprintf("{\"schema\": %s}", strconv.Quote(schema.String()))))
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return 0, err
	}

	if this.isOK(response) {
		decodedResponse := &RegisterSchemaResponse{}
		if this.handleSuccess(response, decodedResponse) != nil {
			return 0, err
		}

		schemaIdMap[schema] = decodedResponse.Id
		this.idCache[decodedResponse.Id] = schema

		return decodedResponse.Id, err
	} else {
		return 0, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetByID(id int32) (avro.Schema, error) {
	var schema avro.Schema
	var exists bool
	if schema, exists = this.idCache[id]; exists {
		return schema, nil
	}

	request, err := this.newDefaultRequest("GET", fmt.Sprintf(GET_SCHEMA_BY_ID, id), nil)
	if err != nil {
		return nil, err
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return nil, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSchemaResponse{}
		if this.handleSuccess(response, decodedResponse) != nil {
			return nil, err
		}
		schema, err := avro.ParseSchema(decodedResponse.Schema)
		this.idCache[id] = schema

		return schema, err
	} else {
		return nil, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetLatestSchemaMetadata(subject string) (*SchemaMetadata, error) {
	request, err := this.newDefaultRequest("GET", fmt.Sprintf(GET_SPECIFIC_SUBJECT_VERSION, subject, "latest"), nil)
	if err != nil {
		return nil, err
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return nil, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSubjectVersionResponse{}
		if this.handleSuccess(response, decodedResponse) != nil {
			return nil, err
		}

		return &SchemaMetadata{decodedResponse.Id, decodedResponse.Version, decodedResponse.Schema}, err
	} else {
		return nil, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetVersion(subject string, schema avro.Schema) (int32, error) {
	var schemaVersionMap map[avro.Schema]int32
	var exists bool
	if schemaVersionMap, exists = this.versionCache[subject]; !exists {
		schemaVersionMap = make(map[avro.Schema]int32)
		this.versionCache[subject] = schemaVersionMap
	}

	var version int32
	if version, exists = schemaVersionMap[schema]; exists {
		return version, nil
	}

	request, err := this.newDefaultRequest("POST",
		fmt.Sprintf(CHECK_IS_REGISTERED, subject),
		strings.NewReader(fmt.Sprintf("{\"schema\": %s}", strconv.Quote(schema.String()))))
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return 0, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSubjectVersionResponse{}
		if this.handleSuccess(response, decodedResponse) != nil {
			return 0, err
		}
		schemaVersionMap[schema] = decodedResponse.Version

		return decodedResponse.Version, err
	} else {
		return 0, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) newDefaultRequest(method string, uri string, reader io.Reader) (*http.Request, error) {
	url := fmt.Sprintf("%s%s", this.registryURL, uri)
	request, err := http.NewRequest(method, url, reader)
	if err != nil {
		return nil, err
	}
	request.Header.Set("Accept", SCHEMA_REGISTRY_V1_JSON)
	request.Header.Set("Content-Type", SCHEMA_REGISTRY_V1_JSON)
	return request, nil
}

func (this *CachedSchemaRegistryClient) isOK(response *http.Response) bool {
	return response.StatusCode >= 200 && response.StatusCode < 300
}

func (this *CachedSchemaRegistryClient) handleSuccess(response *http.Response, model interface{}) error {
	responseBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	return json.Unmarshal(responseBytes, model)
}

func (this *CachedSchemaRegistryClient) handleError(response *http.Response) error {
	registryError := &ErrorMessage{}
	responseBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	err = json.Unmarshal(responseBytes, registryError)
	if err != nil {
		return err
	}

	return registryError
}
<commit_msg>Synchronized schema cache access to avoid data races<commit_after>/* Licensed to the Apache Software Foundation (ASF) under one or more
contributor license agreements.  See the NOTICE file distributed with
this work for additional information regarding copyright ownership.
The ASF licenses this file to You 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 avro

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

	avro "github.com/elodina/go-avro"
	"sync"
)

const (
	GET_SCHEMA_BY_ID             = "/schemas/ids/%d"
	GET_SUBJECTS                 = "/subjects"
	GET_SUBJECT_VERSIONS         = "/subjects/%s/versions"
	GET_SPECIFIC_SUBJECT_VERSION = "/subjects/%s/versions/%s"
	REGISTER_NEW_SCHEMA          = "/subjects/%s/versions"
	CHECK_IS_REGISTERED          = "/subjects/%s"
	TEST_COMPATIBILITY           = "/compatibility/subjects/%s/versions/%s"
	CONFIG                       = "/config"
)

type SchemaRegistryClient interface {
	Register(subject string, schema avro.Schema) (int32, error)
	GetByID(id int32) (avro.Schema, error)
	GetLatestSchemaMetadata(subject string) (*SchemaMetadata, error)
	GetVersion(subject string, schema avro.Schema) (int32, error)
}

type SchemaMetadata struct {
	Id      int32
	Version int32
	Schema  string
}

type CompatibilityLevel string

const (
	BackwardCompatibilityLevel CompatibilityLevel = "BACKWARD"
	ForwardCompatibilityLevel  CompatibilityLevel = "FORWARD"
	FullCompatibilityLevel     CompatibilityLevel = "FULL"
	NoneCompatibilityLevel     CompatibilityLevel = "NONE"
)

const (
	SCHEMA_REGISTRY_V1_JSON               = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_V1_JSON_WEIGHTED      = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_MOST_SPECIFIC_DEFAULT = "application/vnd.schemaregistry.v1+json"
	SCHEMA_REGISTRY_DEFAULT_JSON          = "application/vnd.schemaregistry+json"
	SCHEMA_REGISTRY_DEFAULT_JSON_WEIGHTED = "application/vnd.schemaregistry+json qs=0.9"
	JSON                                  = "application/json"
	JSON_WEIGHTED                         = "application/json qs=0.5"
	GENERIC_REQUEST                       = "application/octet-stream"
)

var PREFERRED_RESPONSE_TYPES = []string{SCHEMA_REGISTRY_V1_JSON, SCHEMA_REGISTRY_DEFAULT_JSON, JSON}

type ErrorMessage struct {
	Error_code int32
	Message    string
}

func (this *ErrorMessage) Error() string {
	return fmt.Sprintf("%s(error code: %d)", this.Message, this.Error_code)
}

type RegisterSchemaResponse struct {
	Id int32
}

type GetSchemaResponse struct {
	Schema string
}

type GetSubjectVersionResponse struct {
	Subject string
	Version int32
	Id      int32
	Schema  string
}

type CachedSchemaRegistryClient struct {
	registryURL  string
	schemaCache  map[string]map[avro.Schema]int32
	idCache      map[int32]avro.Schema
	versionCache map[string]map[avro.Schema]int32
	lock         sync.RWMutex
}

func NewCachedSchemaRegistryClient(registryURL string) *CachedSchemaRegistryClient {
	return &CachedSchemaRegistryClient{
		registryURL:  registryURL,
		schemaCache:  make(map[string]map[avro.Schema]int32),
		idCache:      make(map[int32]avro.Schema),
		versionCache: make(map[string]map[avro.Schema]int32),
	}
}

func (this *CachedSchemaRegistryClient) Register(subject string, schema avro.Schema) (int32, error) {
	var schemaIdMap map[avro.Schema]int32
	var exists bool

	this.lock.RLock()
	if schemaIdMap, exists = this.schemaCache[subject]; exists {
		var id int32
		if id, exists = schemaIdMap[schema]; exists {
			return id, nil
		}
	}
	this.lock.RUnlock()

	this.lock.Lock()
	defer this.lock.Unlock()
	if schemaIdMap, exists = this.schemaCache[subject]; !exists {
		schemaIdMap = make(map[avro.Schema]int32)
		this.schemaCache[subject] = schemaIdMap
	}

	request, err := this.newDefaultRequest("POST",
		fmt.Sprintf(REGISTER_NEW_SCHEMA, subject),
		strings.NewReader(fmt.Sprintf("{\"schema\": %s}", strconv.Quote(schema.String()))))
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return 0, err
	}

	if this.isOK(response) {
		decodedResponse := &RegisterSchemaResponse{}
		if this.handleSuccess(response, decodedResponse) != nil {
			return 0, err
		}

		schemaIdMap[schema] = decodedResponse.Id
		this.idCache[decodedResponse.Id] = schema

		return decodedResponse.Id, err
	} else {
		return 0, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetByID(id int32) (avro.Schema, error) {
	var schema avro.Schema
	var exists bool
	if schema, exists = this.idCache[id]; exists {
		return schema, nil
	}

	request, err := this.newDefaultRequest("GET", fmt.Sprintf(GET_SCHEMA_BY_ID, id), nil)
	if err != nil {
		return nil, err
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return nil, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSchemaResponse{}
		if this.handleSuccess(response, decodedResponse) != nil {
			return nil, err
		}
		schema, err := avro.ParseSchema(decodedResponse.Schema)
		this.idCache[id] = schema

		return schema, err
	} else {
		return nil, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetLatestSchemaMetadata(subject string) (*SchemaMetadata, error) {
	request, err := this.newDefaultRequest("GET", fmt.Sprintf(GET_SPECIFIC_SUBJECT_VERSION, subject, "latest"), nil)
	if err != nil {
		return nil, err
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return nil, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSubjectVersionResponse{}
		if this.handleSuccess(response, decodedResponse) != nil {
			return nil, err
		}

		return &SchemaMetadata{decodedResponse.Id, decodedResponse.Version, decodedResponse.Schema}, err
	} else {
		return nil, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) GetVersion(subject string, schema avro.Schema) (int32, error) {
	var schemaVersionMap map[avro.Schema]int32
	var exists bool
	if schemaVersionMap, exists = this.versionCache[subject]; !exists {
		schemaVersionMap = make(map[avro.Schema]int32)
		this.versionCache[subject] = schemaVersionMap
	}

	var version int32
	if version, exists = schemaVersionMap[schema]; exists {
		return version, nil
	}

	request, err := this.newDefaultRequest("POST",
		fmt.Sprintf(CHECK_IS_REGISTERED, subject),
		strings.NewReader(fmt.Sprintf("{\"schema\": %s}", strconv.Quote(schema.String()))))
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return 0, err
	}

	if this.isOK(response) {
		decodedResponse := &GetSubjectVersionResponse{}
		if this.handleSuccess(response, decodedResponse) != nil {
			return 0, err
		}
		schemaVersionMap[schema] = decodedResponse.Version

		return decodedResponse.Version, err
	} else {
		return 0, this.handleError(response)
	}
}

func (this *CachedSchemaRegistryClient) newDefaultRequest(method string, uri string, reader io.Reader) (*http.Request, error) {
	url := fmt.Sprintf("%s%s", this.registryURL, uri)
	request, err := http.NewRequest(method, url, reader)
	if err != nil {
		return nil, err
	}
	request.Header.Set("Accept", SCHEMA_REGISTRY_V1_JSON)
	request.Header.Set("Content-Type", SCHEMA_REGISTRY_V1_JSON)
	return request, nil
}

func (this *CachedSchemaRegistryClient) isOK(response *http.Response) bool {
	return response.StatusCode >= 200 && response.StatusCode < 300
}

func (this *CachedSchemaRegistryClient) handleSuccess(response *http.Response, model interface{}) error {
	responseBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	return json.Unmarshal(responseBytes, model)
}

func (this *CachedSchemaRegistryClient) handleError(response *http.Response) error {
	registryError := &ErrorMessage{}
	responseBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	err = json.Unmarshal(responseBytes, registryError)
	if err != nil {
		return err
	}

	return registryError
}
<|endoftext|>


<commit_before>package idworker

import (
	"fmt"
	"sync"
	"time"
)

const (
	nano = 1000 * 1000
)

const (
	WorkerIdBits       = 5
	DatacenterIdBits   = 5
	MaxWorkerId        = -1 ^ (-1 << WorkerIdBits)
	MaxDatacenterId    = -1 ^ (-1 << DatacenterIdBits)
	SequenceBits       = 12
	WorkerIdShift      = SequenceBits
	DatacenterIdShift  = SequenceBits + WorkerIdBits
	TimestampLeftShift = SequenceBits + WorkerIdBits + DatacenterIdBits
	SequenceMask       = -1 ^ (-1 << SequenceBits)
)

var (
	Epoch uint64 = 1288834974657 /* tweet poch */
)

type IdWorker struct {
	lastTimestamp uint64
	workerId      uint64
	dataCenterId  uint64
	sequence      uint64
	lock          sync.Mutex
}

func timeGen() uint64 {
	return uint64(time.Now().UnixNano() / nano)
}

func timestamp() uint64 {
	return uint64(time.Now().UnixNano() / nano)
}

func tillNextMillis(ts uint64) uint64 {
	i := timestamp()
	for i <= ts {
		i = timestamp()
	}
	return i
}

func (worker *IdWorker) Next() (uint64, error) {
	ts := timeGen()
	if ts < worker.lastTimestamp {
		err := fmt.Errorf("Clock is moving backwards. Rejecting requests until %d.", worker.lastTimestamp)
		return 1, err
	}

	worker.lock.Lock()
	defer worker.lock.Unlock()

	if worker.lastTimestamp == ts {

		worker.sequence = (worker.sequence + 1) & SequenceMask
		if worker.sequence == 0 {
			ts = tillNextMillis(ts)
		}
	} else {
		worker.sequence = 0
	}

	worker.lastTimestamp = ts

	id := ((worker.lastTimestamp - Epoch) << TimestampLeftShift) |
		(worker.dataCenterId << DatacenterIdShift) |
		(worker.workerId << WorkerIdShift) |
		worker.sequence

	return id, nil
}

func NewIdWorker(workerId uint64, datacenterId uint64) (*IdWorker, error) {
	if workerId > MaxWorkerId || workerId < 0 {
		return nil, fmt.Errorf("workerId can't be greater than %d or less than 0", workerId)
	}

	if datacenterId > MaxDatacenterId || datacenterId < 0 {
		return nil, fmt.Errorf("datacenterId can't be greater than %d or less than 0", datacenterId)
	}
	return &IdWorker{workerId: workerId, dataCenterId: datacenterId, lastTimestamp: 1, sequence: 0}, nil
}
<commit_msg>moved the lock in the Next function to mimic the synchronized block in the original Scala code<commit_after>package idworker

import (
	"fmt"
	"sync"
	"time"
)

const (
	nano = 1000 * 1000
)

const (
	WorkerIdBits       = 5
	DatacenterIdBits   = 5
	MaxWorkerId        = -1 ^ (-1 << WorkerIdBits)
	MaxDatacenterId    = -1 ^ (-1 << DatacenterIdBits)
	SequenceBits       = 12
	WorkerIdShift      = SequenceBits
	DatacenterIdShift  = SequenceBits + WorkerIdBits
	TimestampLeftShift = SequenceBits + WorkerIdBits + DatacenterIdBits
	SequenceMask       = -1 ^ (-1 << SequenceBits)
)

var (
	Epoch uint64 = 1288834974657 /* tweet poch */
)

type IdWorker struct {
	lastTimestamp uint64
	workerId      uint64
	dataCenterId  uint64
	sequence      uint64
	lock          sync.Mutex
}

func timeGen() uint64 {
	return uint64(time.Now().UnixNano() / nano)
}

func timestamp() uint64 {
	return uint64(time.Now().UnixNano() / nano)
}

func tillNextMillis(ts uint64) uint64 {
	i := timestamp()
	for i <= ts {
		i = timestamp()
	}
	return i
}

func (worker *IdWorker) Next() (uint64, error) {
	worker.lock.Lock()
	defer worker.lock.Unlock()

	ts := timeGen()
	if ts < worker.lastTimestamp {
		err := fmt.Errorf("Clock is moving backwards. Rejecting requests until %d.", worker.lastTimestamp)
		return 1, err
	}

	if worker.lastTimestamp == ts {
		worker.sequence = (worker.sequence + 1) & SequenceMask
		if worker.sequence == 0 {
			ts = tillNextMillis(ts)
		}
	} else {
		worker.sequence = 0
	}

	worker.lastTimestamp = ts

	id := ((worker.lastTimestamp - Epoch) << TimestampLeftShift) |
		(worker.dataCenterId << DatacenterIdShift) |
		(worker.workerId << WorkerIdShift) |
		worker.sequence

	return id, nil
}

func NewIdWorker(workerId uint64, datacenterId uint64) (*IdWorker, error) {
	if workerId > MaxWorkerId || workerId < 0 {
		return nil, fmt.Errorf("workerId can't be greater than %d or less than 0", workerId)
	}

	if datacenterId > MaxDatacenterId || datacenterId < 0 {
		return nil, fmt.Errorf("datacenterId can't be greater than %d or less than 0", datacenterId)
	}
	return &IdWorker{workerId: workerId, dataCenterId: datacenterId, lastTimestamp: 1, sequence: 0}, nil
}
<|endoftext|>


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

package db_service

import (
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"os"

	"code.cloudfoundry.org/lager"
	"github.com/go-sql-driver/mysql"
	"github.com/jinzhu/gorm"
	"github.com/spf13/viper"
)

const (
	caCertProp     = "db.ca.cert"
	clientCertProp = "db.client.cert"
	clientKeyProp  = "db.client.key"
	dbHostProp     = "db.host"
	dbUserProp     = "db.user"
	dbPassProp     = "db.password"
	dbPortProp     = "db.port"
	dbNameProp     = "db.name"
	dbTypeProp     = "db.type"
	dbPathProp     = "db.path"

	DbTypeMysql   = "mysql"
	DbTypeSqlite3 = "sqlite3"
)

func init() {
	viper.BindEnv(caCertProp, "CA_CERT")
	viper.BindEnv(clientCertProp, "CLIENT_CERT")
	viper.BindEnv(clientKeyProp, "CLIENT_KEY")

	viper.BindEnv(dbHostProp, "DB_HOST")
	viper.BindEnv(dbUserProp, "DB_USERNAME")
	viper.BindEnv(dbPassProp, "DB_PASSWORD")

	viper.BindEnv(dbPortProp, "DB_PORT")
	viper.SetDefault(dbPortProp, "3306")
	viper.BindEnv(dbNameProp, "DB_NAME")
	viper.SetDefault(dbNameProp, "servicebroker")

	viper.BindEnv(dbTypeProp, "DB_TYPE")
	viper.SetDefault(dbTypeProp, "mysql")

	viper.BindEnv(dbPathProp, "DB_PATH")
}

// pulls db credentials from the environment, connects to the db, and returns the db connection
func SetupDb(logger lager.Logger) *gorm.DB {
	dbType := viper.GetString(dbTypeProp)
	var db *gorm.DB
	var err error
	switch dbType {
	default:
		logger.Error("Database Setup", fmt.Errorf("Invalid database type %q, valid types are: sqlite3 and mysql", dbType))
		os.Exit(1)
	case DbTypeMysql:
		db, err = setupMysqlDb(logger)
	case DbTypeSqlite3:
		db, err = setupSqlite3Db(logger)
	}

	if err != nil {
		logger.Error("Database Setup", err)
		os.Exit(1)
	}

	return db
}

func setupSqlite3Db(logger lager.Logger) (*gorm.DB, error) {
	dbPath := viper.GetString(dbPathProp)
	if dbPath == "" {
		return nil, fmt.Errorf("You must set a database path when using SQLite3 databases")
	}

	logger.Info("WARNING: DO NOT USE SQLITE3 IN PRODUCTION!")

	return gorm.Open(DbTypeSqlite3, dbPath)
}

func setupMysqlDb(logger lager.Logger) (*gorm.DB, error) {
	// connect to database
	dbHost := viper.GetString(dbHostProp)
	dbUsername := viper.GetString(dbUserProp)
	dbPassword := viper.GetString(dbPassProp)

	if dbPassword == "" || dbHost == "" || dbUsername == "" {
		return nil, errors.New("DB_HOST, DB_USERNAME and DB_PASSWORD are required environment variables.")
	}

	dbPort := viper.GetString(dbPortProp)
	dbName := viper.GetString(dbNameProp)

	tlsStr, err := generateTlsStringFromEnv()
	if err != nil {
		return nil, fmt.Errorf("Error generating TLS string from env: %s", err)
	}

	logger.Info("Connecting to MySQL Database", lager.Data{
		"host": dbHost,
		"port": dbPort,
		"name": dbName,
	})

	connStr := fmt.Sprintf("%v:%v@tcp(%v:%v)/%v?charset=utf8&parseTime=True&loc=Local%v", dbUsername, dbPassword, dbHost, dbPort, dbName, tlsStr)
	return gorm.Open(DbTypeMysql, connStr)
}

func generateTlsStringFromEnv() (string, error) {
	caCert := viper.GetString(caCertProp)
	clientCertStr := viper.GetString(clientCertProp)
	clientKeyStr := viper.GetString(clientKeyProp)
	tlsStr := "&tls=custom"

	// make sure ssl is set up for this connection
	if caCert != "" && clientCertStr != "" && clientKeyStr != "" {

		rootCertPool := x509.NewCertPool()

		if ok := rootCertPool.AppendCertsFromPEM([]byte(caCert)); !ok {
			return "", fmt.Errorf("Error appending cert: %s", errors.New(""))
		}
		clientCert := make([]tls.Certificate, 0, 1)

		certs, err := tls.X509KeyPair([]byte(clientCertStr), []byte(clientKeyStr))
		if err != nil {
			return "", fmt.Errorf("Error parsing cert pair: %s", err)
		}
		clientCert = append(clientCert, certs)
		mysql.RegisterTLSConfig("custom", &tls.Config{
			RootCAs:            rootCertPool,
			Certificates:       clientCert,
			InsecureSkipVerify: true,
		})
	} else {
		tlsStr = ""
	}

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

package db_service

import (
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"os"

	"code.cloudfoundry.org/lager"
	"github.com/go-sql-driver/mysql"
	"github.com/jinzhu/gorm"
	"github.com/spf13/viper"
)

const (
	caCertProp     = "db.ca.cert"
	clientCertProp = "db.client.cert"
	clientKeyProp  = "db.client.key"
	dbHostProp     = "db.host"
	dbUserProp     = "db.user"
	dbPassProp     = "db.password"
	dbPortProp     = "db.port"
	dbNameProp     = "db.name"
	dbTypeProp     = "db.type"
	dbPathProp     = "db.path"

	DbTypeMysql   = "mysql"
	DbTypeSqlite3 = "sqlite3"
)

func init() {
	viper.BindEnv(caCertProp, "CA_CERT")
	viper.BindEnv(clientCertProp, "CLIENT_CERT")
	viper.BindEnv(clientKeyProp, "CLIENT_KEY")

	viper.BindEnv(dbHostProp, "DB_HOST")
	viper.BindEnv(dbUserProp, "DB_USERNAME")
	viper.BindEnv(dbPassProp, "DB_PASSWORD")

	viper.BindEnv(dbPortProp, "DB_PORT")
	viper.SetDefault(dbPortProp, "3306")
	viper.BindEnv(dbNameProp, "DB_NAME")
	viper.SetDefault(dbNameProp, "servicebroker")

	viper.BindEnv(dbTypeProp, "DB_TYPE")
	viper.SetDefault(dbTypeProp, DbTypeMysql)

	viper.BindEnv(dbPathProp, "DB_PATH")
}

// pulls db credentials from the environment, connects to the db, and returns the db connection
func SetupDb(logger lager.Logger) *gorm.DB {
	dbType := viper.GetString(dbTypeProp)
	var db *gorm.DB
	var err error
	switch dbType {
	default:
		logger.Error("Database Setup", fmt.Errorf("Invalid database type %q, valid types are: sqlite3 and mysql", dbType))
		os.Exit(1)
	case DbTypeMysql:
		db, err = setupMysqlDb(logger)
	case DbTypeSqlite3:
		db, err = setupSqlite3Db(logger)
	}

	if err != nil {
		logger.Error("Database Setup", err)
		os.Exit(1)
	}

	return db
}

func setupSqlite3Db(logger lager.Logger) (*gorm.DB, error) {
	dbPath := viper.GetString(dbPathProp)
	if dbPath == "" {
		return nil, fmt.Errorf("You must set a database path when using SQLite3 databases")
	}

	logger.Info("WARNING: DO NOT USE SQLITE3 IN PRODUCTION!")

	return gorm.Open(DbTypeSqlite3, dbPath)
}

func setupMysqlDb(logger lager.Logger) (*gorm.DB, error) {
	// connect to database
	dbHost := viper.GetString(dbHostProp)
	dbUsername := viper.GetString(dbUserProp)
	dbPassword := viper.GetString(dbPassProp)

	if dbPassword == "" || dbHost == "" || dbUsername == "" {
		return nil, errors.New("DB_HOST, DB_USERNAME and DB_PASSWORD are required environment variables.")
	}

	dbPort := viper.GetString(dbPortProp)
	dbName := viper.GetString(dbNameProp)

	tlsStr, err := generateTlsStringFromEnv()
	if err != nil {
		return nil, fmt.Errorf("Error generating TLS string from env: %s", err)
	}

	logger.Info("Connecting to MySQL Database", lager.Data{
		"host": dbHost,
		"port": dbPort,
		"name": dbName,
	})

	connStr := fmt.Sprintf("%v:%v@tcp(%v:%v)/%v?charset=utf8&parseTime=True&loc=Local%v", dbUsername, dbPassword, dbHost, dbPort, dbName, tlsStr)
	return gorm.Open(DbTypeMysql, connStr)
}

func generateTlsStringFromEnv() (string, error) {
	caCert := viper.GetString(caCertProp)
	clientCertStr := viper.GetString(clientCertProp)
	clientKeyStr := viper.GetString(clientKeyProp)
	tlsStr := "&tls=custom"

	// make sure ssl is set up for this connection
	if caCert != "" && clientCertStr != "" && clientKeyStr != "" {

		rootCertPool := x509.NewCertPool()

		if ok := rootCertPool.AppendCertsFromPEM([]byte(caCert)); !ok {
			return "", fmt.Errorf("Error appending cert: %s", errors.New(""))
		}
		clientCert := make([]tls.Certificate, 0, 1)

		certs, err := tls.X509KeyPair([]byte(clientCertStr), []byte(clientKeyStr))
		if err != nil {
			return "", fmt.Errorf("Error parsing cert pair: %s", err)
		}
		clientCert = append(clientCert, certs)
		mysql.RegisterTLSConfig("custom", &tls.Config{
			RootCAs:            rootCertPool,
			Certificates:       clientCert,
			InsecureSkipVerify: true,
		})
	} else {
		tlsStr = ""
	}

	return tlsStr, nil
}
<|endoftext|>


<commit_before>package main

import "fmt"

func sum(a []int, ch chan int) {
    result := 0
    for _, value := range a {
        result += value
    }
    ch <- result
}

func Example() {
    var numbers [100000]int
    for i := range numbers {
        numbers[i] = i
    }
    ch := make(chan int, 2)
    go sum(numbers[:len(numbers)/2], ch)
    go sum(numbers[len(numbers)/2:], ch)
    fmt.Println(<-ch + <-ch)
    // Output: 4999950000
}
<commit_msg>Better typing on channel parameters<commit_after>package main

import "fmt"

func sum(a []int, ch chan<- int) {
    result := 0
    for _, value := range a {
        result += value
    }
    ch <- result
}

func Example() {
    var numbers [100000]int
    for i := range numbers {
        numbers[i] = i
    }
    ch := make(chan int, 2)
    go sum(numbers[:len(numbers)/2], ch)
    go sum(numbers[len(numbers)/2:], ch)
    fmt.Println(<-ch + <-ch)
    // Output: 4999950000
}
<|endoftext|>


<commit_before>package utils

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

	"github.com/keybase/client/go/protocol/chat1"
	"github.com/keybase/client/go/protocol/gregor1"
	"github.com/keybase/client/go/protocol/keybase1"
)

// parseDurationExtended is like time.ParseDuration, but adds "d" unit. "1d" is
// one day, defined as 24*time.Hour. Only whole days are supported for "d"
// unit, but it can be followed by smaller units, e.g., "1d1h".
func ParseDurationExtended(s string) (d time.Duration, err error) {
	p := strings.Index(s, "d")
	if p == -1 {
		// no "d" suffix
		return time.ParseDuration(s)
	}

	var days int
	if days, err = strconv.Atoi(s[:p]); err != nil {
		return time.Duration(0), err
	}
	d = time.Duration(days) * 24 * time.Hour

	if p < len(s)-1 {
		var dur time.Duration
		if dur, err = time.ParseDuration(s[p+1:]); err != nil {
			return time.Duration(0), err
		}
		d += dur
	}

	return d, nil
}

func ParseTimeFromRFC3339OrDurationFromPast(kbCtx KeybaseContext, s string) (t time.Time, err error) {
	var errt, errd error
	var d time.Duration

	if s == "" {
		return
	}

	if t, errt = time.Parse(time.RFC3339, s); errt == nil {
		return t, nil
	}
	if d, errd = ParseDurationExtended(s); errd == nil {
		return kbCtx.Clock().Now().Add(-d), nil
	}

	return time.Time{}, fmt.Errorf("given string is neither a valid time (%s) nor a valid duration (%v)", errt, errd)
}

// upper bounds takes higher priority
func Collar(lower int, ideal int, upper int) int {
	if ideal > upper {
		return upper
	}
	if ideal < lower {
		return lower
	}
	return ideal
}

func FilterByType(msgs []chat1.MessageUnboxed, query *chat1.GetThreadQuery) (res []chat1.MessageUnboxed) {
	if query != nil && len(query.MessageTypes) > 0 {
		typmap := make(map[chat1.MessageType]bool)
		for _, mt := range query.MessageTypes {
			typmap[mt] = true
		}
		for _, msg := range msgs {
			if _, ok := typmap[msg.GetMessageType()]; ok {
				res = append(res, msg)
			}
		}
	} else {
		res = msgs
	}
	return res
}

// AggRateLimitsP takes a list of rate limit responses and dedups them to the last one received
// of each category
func AggRateLimitsP(rlimits []*chat1.RateLimit) (res []chat1.RateLimit) {
	m := make(map[string]chat1.RateLimit)
	for _, l := range rlimits {
		if l != nil {
			m[l.Name] = *l
		}
	}
	for _, v := range m {
		res = append(res, v)
	}
	return res
}

func AggRateLimits(rlimits []chat1.RateLimit) (res []chat1.RateLimit) {
	m := make(map[string]chat1.RateLimit)
	for _, l := range rlimits {
		m[l.Name] = l
	}
	for _, v := range m {
		res = append(res, v)
	}
	return res
}

// Reorder participants based on the order in activeList.
// This never fails, worse comes to worst it just returns the split of tlfname.
func ReorderParticipants(udc *UserDeviceCache, uimap *UserInfoMapper, tlfname string, activeList []gregor1.UID) []string {
	usedUsers := make(map[string]bool)
	var users []string

	// Fill from the active list first.
	for _, uid := range activeList {
		kbUID := keybase1.UID(uid.String())
		user, err := udc.LookupUsername(uimap, kbUID)
		if err != nil {
			continue
		}
		if _, used := usedUsers[user]; !used {
			users = append(users, user)
			usedUsers[user] = true
		}
	}

	// Include participants even if they're not in the active list, in stable order.
	for _, user := range splitTlfName(tlfname) {
		if _, used := usedUsers[user]; !used {
			users = append(users, user)
			usedUsers[user] = true
		}
	}

	return users
}

// Split a tlf name into its users.
// Does not validate the usernames.
func splitTlfName(tlfname string) []string {
	writerSep := ","
	readerSep := "#"
	extensionSep := " "

	// Strip off the suffix
	s2 := strings.Split(tlfname, extensionSep)
	tlfname = s2[0]
	// Replace "#" with ","
	tlfname = strings.Replace(tlfname, writerSep, readerSep, -1)
	// Split on ","
	return strings.Split(tlfname, readerSep)
}
<commit_msg>Only allow users from tlfname<commit_after>package utils

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

	"github.com/keybase/client/go/protocol/chat1"
	"github.com/keybase/client/go/protocol/gregor1"
	"github.com/keybase/client/go/protocol/keybase1"
)

// parseDurationExtended is like time.ParseDuration, but adds "d" unit. "1d" is
// one day, defined as 24*time.Hour. Only whole days are supported for "d"
// unit, but it can be followed by smaller units, e.g., "1d1h".
func ParseDurationExtended(s string) (d time.Duration, err error) {
	p := strings.Index(s, "d")
	if p == -1 {
		// no "d" suffix
		return time.ParseDuration(s)
	}

	var days int
	if days, err = strconv.Atoi(s[:p]); err != nil {
		return time.Duration(0), err
	}
	d = time.Duration(days) * 24 * time.Hour

	if p < len(s)-1 {
		var dur time.Duration
		if dur, err = time.ParseDuration(s[p+1:]); err != nil {
			return time.Duration(0), err
		}
		d += dur
	}

	return d, nil
}

func ParseTimeFromRFC3339OrDurationFromPast(kbCtx KeybaseContext, s string) (t time.Time, err error) {
	var errt, errd error
	var d time.Duration

	if s == "" {
		return
	}

	if t, errt = time.Parse(time.RFC3339, s); errt == nil {
		return t, nil
	}
	if d, errd = ParseDurationExtended(s); errd == nil {
		return kbCtx.Clock().Now().Add(-d), nil
	}

	return time.Time{}, fmt.Errorf("given string is neither a valid time (%s) nor a valid duration (%v)", errt, errd)
}

// upper bounds takes higher priority
func Collar(lower int, ideal int, upper int) int {
	if ideal > upper {
		return upper
	}
	if ideal < lower {
		return lower
	}
	return ideal
}

func FilterByType(msgs []chat1.MessageUnboxed, query *chat1.GetThreadQuery) (res []chat1.MessageUnboxed) {
	if query != nil && len(query.MessageTypes) > 0 {
		typmap := make(map[chat1.MessageType]bool)
		for _, mt := range query.MessageTypes {
			typmap[mt] = true
		}
		for _, msg := range msgs {
			if _, ok := typmap[msg.GetMessageType()]; ok {
				res = append(res, msg)
			}
		}
	} else {
		res = msgs
	}
	return res
}

// AggRateLimitsP takes a list of rate limit responses and dedups them to the last one received
// of each category
func AggRateLimitsP(rlimits []*chat1.RateLimit) (res []chat1.RateLimit) {
	m := make(map[string]chat1.RateLimit)
	for _, l := range rlimits {
		if l != nil {
			m[l.Name] = *l
		}
	}
	for _, v := range m {
		res = append(res, v)
	}
	return res
}

func AggRateLimits(rlimits []chat1.RateLimit) (res []chat1.RateLimit) {
	m := make(map[string]chat1.RateLimit)
	for _, l := range rlimits {
		m[l.Name] = l
	}
	for _, v := range m {
		res = append(res, v)
	}
	return res
}

// Reorder participants based on the order in activeList.
// Only allows usernames from tlfname in the output.
// This never fails, worse comes to worst it just returns the split of tlfname.
func ReorderParticipants(udc *UserDeviceCache, uimap *UserInfoMapper, tlfname string, activeList []gregor1.UID) []string {
	tlfnameList := splitTlfName(tlfname)
	allowedUsers := make(map[string]bool)
	var users []string

	// Allow all users from tlfname.
	for _, user := range tlfnameList {
		allowedUsers[user] = true
	}

	// Fill from the active list first.
	for _, uid := range activeList {
		kbUID := keybase1.UID(uid.String())
		user, err := udc.LookupUsername(uimap, kbUID)
		if err != nil {
			continue
		}
		if allowed, _ := allowedUsers[user]; allowed {
			users = append(users, user)
			// Allow only one occurrence.
			allowedUsers[user] = false
		}
	}

	// Include participants even if they weren't in the active list, in stable order.
	for _, user := range tlfnameList {
		if allowed, _ := allowedUsers[user]; allowed {
			users = append(users, user)
			allowedUsers[user] = false
		}
	}

	return users
}

// Split a tlf name into its users.
// Does not validate the usernames.
func splitTlfName(tlfname string) []string {
	writerSep := ","
	readerSep := "#"
	extensionSep := " "

	// Strip off the suffix
	s2 := strings.Split(tlfname, extensionSep)
	tlfname = s2[0]
	// Replace "#" with ","
	tlfname = strings.Replace(tlfname, writerSep, readerSep, -1)
	// Split on ","
	return strings.Split(tlfname, readerSep)
}
<|endoftext|>


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

package ascii85

import (
	"bytes"
	"io"
	"io/ioutil"
	"testing"
)

type testpair struct {
	decoded, encoded string
}

var pairs = []testpair{
	// Wikipedia example
	{
		"Man is distinguished, not only by his reason, but by this singular passion from " +
			"other animals, which is a lust of the mind, that by a perseverance of delight in " +
			"the continued and indefatigable generation of knowledge, exceeds the short " +
			"vehemence of any carnal pleasure.",
		"9jqo^BlbD-BleB1DJ+*+F(f,q/0JhKF<GL>Cj@.4Gp$d7F!,L7@<6@)/0JDEF<G%<+EV:2F!,\n" +
			"O<DJ+*.@<*K0@<6L(Df-\\0Ec5e;DffZ(EZee.Bl.9pF\"AGXBPCsi+DGm>@3BB/F*&OCAfu2/AKY\n" +
			"i(DIb:@FD,*)+C]U=@3BN#EcYf8ATD3s@q?d$AftVqCh[NqF<G:8+EV:.+Cf>-FD5W8ARlolDIa\n" +
			"l(DId<j@<?3r@:F%a+D58'ATD4$Bl@l3De:,-DJs`8ARoFb/0JMK@qB4^F!,R<AKZ&-DfTqBG%G\n" +
			">uD.RTpAKYo'+CT/5+Cei#DII?(E,9)oF*2M7/c\n",
	},
	// Special case when shortening !!!!! to z.
	{
		"\000\000\000\000",
		"z",
	},
}

var bigtest = pairs[len(pairs)-1]

func testEqual(t *testing.T, msg string, args ...interface{}) bool {
	if args[len(args)-2] != args[len(args)-1] {
		t.Errorf(msg, args...)
		return false
	}
	return true
}

func strip85(s string) string {
	t := make([]byte, len(s))
	w := 0
	for r := 0; r < len(s); r++ {
		c := s[r]
		if c > ' ' {
			t[w] = c
			w++
		}
	}
	return string(t[0:w])
}

func TestEncode(t *testing.T) {
	for _, p := range pairs {
		buf := make([]byte, MaxEncodedLen(len(p.decoded)))
		n := Encode(buf, []byte(p.decoded))
		buf = buf[0:n]
		testEqual(t, "Encode(%q) = %q, want %q", p.decoded, strip85(string(buf)), strip85(p.encoded))
	}
}

func TestEncoder(t *testing.T) {
	for _, p := range pairs {
		bb := &bytes.Buffer{}
		encoder := NewEncoder(bb)
		encoder.Write([]byte(p.decoded))
		encoder.Close()
		testEqual(t, "Encode(%q) = %q, want %q", p.decoded, strip85(bb.String()), strip85(p.encoded))
	}
}

func TestEncoderBuffering(t *testing.T) {
	input := []byte(bigtest.decoded)
	for bs := 1; bs <= 12; bs++ {
		bb := &bytes.Buffer{}
		encoder := NewEncoder(bb)
		for pos := 0; pos < len(input); pos += bs {
			end := pos + bs
			if end > len(input) {
				end = len(input)
			}
			n, err := encoder.Write(input[pos:end])
			testEqual(t, "Write(%q) gave error %v, want %v", input[pos:end], err, error(nil))
			testEqual(t, "Write(%q) gave length %v, want %v", input[pos:end], n, end-pos)
		}
		err := encoder.Close()
		testEqual(t, "Close gave error %v, want %v", err, error(nil))
		testEqual(t, "Encoding/%d of %q = %q, want %q", bs, bigtest.decoded, strip85(bb.String()), strip85(bigtest.encoded))
	}
}

func TestDecode(t *testing.T) {
	for _, p := range pairs {
		dbuf := make([]byte, 4*len(p.encoded))
		ndst, nsrc, err := Decode(dbuf, []byte(p.encoded), true)
		testEqual(t, "Decode(%q) = error %v, want %v", p.encoded, err, error(nil))
		testEqual(t, "Decode(%q) = nsrc %v, want %v", p.encoded, nsrc, len(p.encoded))
		testEqual(t, "Decode(%q) = ndst %v, want %v", p.encoded, ndst, len(p.decoded))
		testEqual(t, "Decode(%q) = %q, want %q", p.encoded, string(dbuf[0:ndst]), p.decoded)
	}
}

func TestDecoder(t *testing.T) {
	for _, p := range pairs {
		decoder := NewDecoder(bytes.NewBufferString(p.encoded))
		dbuf, err := ioutil.ReadAll(decoder)
		if err != nil {
			t.Fatal("Read failed", err)
		}
		testEqual(t, "Read from %q = length %v, want %v", p.encoded, len(dbuf), len(p.decoded))
		testEqual(t, "Decoding of %q = %q, want %q", p.encoded, string(dbuf), p.decoded)
		if err != nil {
			testEqual(t, "Read from %q = %v, want %v", p.encoded, err, io.EOF)
		}
	}
}

func TestDecoderBuffering(t *testing.T) {
	for bs := 1; bs <= 12; bs++ {
		decoder := NewDecoder(bytes.NewBufferString(bigtest.encoded))
		buf := make([]byte, len(bigtest.decoded)+12)
		var total int
		for total = 0; total < len(bigtest.decoded); {
			n, err := decoder.Read(buf[total : total+bs])
			testEqual(t, "Read from %q at pos %d = %d, %v, want _, %v", bigtest.encoded, total, n, err, error(nil))
			total += n
		}
		testEqual(t, "Decoding/%d of %q = %q, want %q", bs, bigtest.encoded, string(buf[0:total]), bigtest.decoded)
	}
}

func TestDecodeCorrupt(t *testing.T) {
	type corrupt struct {
		e string
		p int
	}
	examples := []corrupt{
		{"v", 0},
		{"!z!!!!!!!!!", 1},
	}

	for _, e := range examples {
		dbuf := make([]byte, 4*len(e.e))
		_, _, err := Decode(dbuf, []byte(e.e), true)
		switch err := err.(type) {
		case CorruptInputError:
			testEqual(t, "Corruption in %q at offset %v, want %v", e.e, int(err), e.p)
		default:
			t.Error("Decoder failed to detect corruption in", e)
		}
	}
}

func TestBig(t *testing.T) {
	n := 3*1000 + 1
	raw := make([]byte, n)
	const alpha = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
	for i := 0; i < n; i++ {
		raw[i] = alpha[i%len(alpha)]
	}
	encoded := new(bytes.Buffer)
	w := NewEncoder(encoded)
	nn, err := w.Write(raw)
	if nn != n || err != nil {
		t.Fatalf("Encoder.Write(raw) = %d, %v want %d, nil", nn, err, n)
	}
	err = w.Close()
	if err != nil {
		t.Fatalf("Encoder.Close() = %v want nil", err)
	}
	decoded, err := ioutil.ReadAll(NewDecoder(encoded))
	if err != nil {
		t.Fatalf("io.ReadAll(NewDecoder(...)): %v", err)
	}

	if !bytes.Equal(raw, decoded) {
		var i int
		for i = 0; i < len(decoded) && i < len(raw); i++ {
			if decoded[i] != raw[i] {
				break
			}
		}
		t.Errorf("Decode(Encode(%d-byte string)) failed at offset %d", n, i)
	}
}
<commit_msg>encoding/ascii85: add empty string case for Encode test<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ascii85

import (
	"bytes"
	"io"
	"io/ioutil"
	"testing"
)

type testpair struct {
	decoded, encoded string
}

var pairs = []testpair{
	// Encode returns 0 when len(src) is 0
	{
		"",
		"",
	},
	// Wikipedia example
	{
		"Man is distinguished, not only by his reason, but by this singular passion from " +
			"other animals, which is a lust of the mind, that by a perseverance of delight in " +
			"the continued and indefatigable generation of knowledge, exceeds the short " +
			"vehemence of any carnal pleasure.",
		"9jqo^BlbD-BleB1DJ+*+F(f,q/0JhKF<GL>Cj@.4Gp$d7F!,L7@<6@)/0JDEF<G%<+EV:2F!,\n" +
			"O<DJ+*.@<*K0@<6L(Df-\\0Ec5e;DffZ(EZee.Bl.9pF\"AGXBPCsi+DGm>@3BB/F*&OCAfu2/AKY\n" +
			"i(DIb:@FD,*)+C]U=@3BN#EcYf8ATD3s@q?d$AftVqCh[NqF<G:8+EV:.+Cf>-FD5W8ARlolDIa\n" +
			"l(DId<j@<?3r@:F%a+D58'ATD4$Bl@l3De:,-DJs`8ARoFb/0JMK@qB4^F!,R<AKZ&-DfTqBG%G\n" +
			">uD.RTpAKYo'+CT/5+Cei#DII?(E,9)oF*2M7/c\n",
	},
	// Special case when shortening !!!!! to z.
	{
		"\000\000\000\000",
		"z",
	},
}

var bigtest = pairs[len(pairs)-1]

func testEqual(t *testing.T, msg string, args ...interface{}) bool {
	if args[len(args)-2] != args[len(args)-1] {
		t.Errorf(msg, args...)
		return false
	}
	return true
}

func strip85(s string) string {
	t := make([]byte, len(s))
	w := 0
	for r := 0; r < len(s); r++ {
		c := s[r]
		if c > ' ' {
			t[w] = c
			w++
		}
	}
	return string(t[0:w])
}

func TestEncode(t *testing.T) {
	for _, p := range pairs {
		buf := make([]byte, MaxEncodedLen(len(p.decoded)))
		n := Encode(buf, []byte(p.decoded))
		buf = buf[0:n]
		testEqual(t, "Encode(%q) = %q, want %q", p.decoded, strip85(string(buf)), strip85(p.encoded))
	}
}

func TestEncoder(t *testing.T) {
	for _, p := range pairs {
		bb := &bytes.Buffer{}
		encoder := NewEncoder(bb)
		encoder.Write([]byte(p.decoded))
		encoder.Close()
		testEqual(t, "Encode(%q) = %q, want %q", p.decoded, strip85(bb.String()), strip85(p.encoded))
	}
}

func TestEncoderBuffering(t *testing.T) {
	input := []byte(bigtest.decoded)
	for bs := 1; bs <= 12; bs++ {
		bb := &bytes.Buffer{}
		encoder := NewEncoder(bb)
		for pos := 0; pos < len(input); pos += bs {
			end := pos + bs
			if end > len(input) {
				end = len(input)
			}
			n, err := encoder.Write(input[pos:end])
			testEqual(t, "Write(%q) gave error %v, want %v", input[pos:end], err, error(nil))
			testEqual(t, "Write(%q) gave length %v, want %v", input[pos:end], n, end-pos)
		}
		err := encoder.Close()
		testEqual(t, "Close gave error %v, want %v", err, error(nil))
		testEqual(t, "Encoding/%d of %q = %q, want %q", bs, bigtest.decoded, strip85(bb.String()), strip85(bigtest.encoded))
	}
}

func TestDecode(t *testing.T) {
	for _, p := range pairs {
		dbuf := make([]byte, 4*len(p.encoded))
		ndst, nsrc, err := Decode(dbuf, []byte(p.encoded), true)
		testEqual(t, "Decode(%q) = error %v, want %v", p.encoded, err, error(nil))
		testEqual(t, "Decode(%q) = nsrc %v, want %v", p.encoded, nsrc, len(p.encoded))
		testEqual(t, "Decode(%q) = ndst %v, want %v", p.encoded, ndst, len(p.decoded))
		testEqual(t, "Decode(%q) = %q, want %q", p.encoded, string(dbuf[0:ndst]), p.decoded)
	}
}

func TestDecoder(t *testing.T) {
	for _, p := range pairs {
		decoder := NewDecoder(bytes.NewBufferString(p.encoded))
		dbuf, err := ioutil.ReadAll(decoder)
		if err != nil {
			t.Fatal("Read failed", err)
		}
		testEqual(t, "Read from %q = length %v, want %v", p.encoded, len(dbuf), len(p.decoded))
		testEqual(t, "Decoding of %q = %q, want %q", p.encoded, string(dbuf), p.decoded)
		if err != nil {
			testEqual(t, "Read from %q = %v, want %v", p.encoded, err, io.EOF)
		}
	}
}

func TestDecoderBuffering(t *testing.T) {
	for bs := 1; bs <= 12; bs++ {
		decoder := NewDecoder(bytes.NewBufferString(bigtest.encoded))
		buf := make([]byte, len(bigtest.decoded)+12)
		var total int
		for total = 0; total < len(bigtest.decoded); {
			n, err := decoder.Read(buf[total : total+bs])
			testEqual(t, "Read from %q at pos %d = %d, %v, want _, %v", bigtest.encoded, total, n, err, error(nil))
			total += n
		}
		testEqual(t, "Decoding/%d of %q = %q, want %q", bs, bigtest.encoded, string(buf[0:total]), bigtest.decoded)
	}
}

func TestDecodeCorrupt(t *testing.T) {
	type corrupt struct {
		e string
		p int
	}
	examples := []corrupt{
		{"v", 0},
		{"!z!!!!!!!!!", 1},
	}

	for _, e := range examples {
		dbuf := make([]byte, 4*len(e.e))
		_, _, err := Decode(dbuf, []byte(e.e), true)
		switch err := err.(type) {
		case CorruptInputError:
			testEqual(t, "Corruption in %q at offset %v, want %v", e.e, int(err), e.p)
		default:
			t.Error("Decoder failed to detect corruption in", e)
		}
	}
}

func TestBig(t *testing.T) {
	n := 3*1000 + 1
	raw := make([]byte, n)
	const alpha = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
	for i := 0; i < n; i++ {
		raw[i] = alpha[i%len(alpha)]
	}
	encoded := new(bytes.Buffer)
	w := NewEncoder(encoded)
	nn, err := w.Write(raw)
	if nn != n || err != nil {
		t.Fatalf("Encoder.Write(raw) = %d, %v want %d, nil", nn, err, n)
	}
	err = w.Close()
	if err != nil {
		t.Fatalf("Encoder.Close() = %v want nil", err)
	}
	decoded, err := ioutil.ReadAll(NewDecoder(encoded))
	if err != nil {
		t.Fatalf("io.ReadAll(NewDecoder(...)): %v", err)
	}

	if !bytes.Equal(raw, decoded) {
		var i int
		for i = 0; i < len(decoded) && i < len(raw); i++ {
			if decoded[i] != raw[i] {
				break
			}
		}
		t.Errorf("Decode(Encode(%d-byte string)) failed at offset %d", n, i)
	}
}
<|endoftext|>


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

// The suffixarray package implements substring search in logarithmic time
// using an in-memory suffix array.
//
// Example use:
//
//	// create index for some data
//	index := suffixarray.New(data)
//
//	// lookup byte slice s
//	offsets1 := index.Lookup(s, -1) // the list of all indices where s occurs in data
//	offsets2 := index.Lookup(s, 3)  // the list of at most 3 indices where s occurs in data
//
package suffixarray

import (
	"bytes"
	"container/vector"
	"sort"
)

// BUG(gri): For larger data (10MB) which contains very long (say 100000)
// contiguous sequences of identical bytes, index creation time will be extremely slow.

// TODO(gri): Use a more sophisticated algorithm to create the suffix array.

// Index implements a suffix array for fast substring search.
type Index struct {
	data []byte
	sa   []int // suffix array for data
}

// New creates a new Index for data.
// Index creation time is approximately O(N*log(N)) for N = len(data).
//
func New(data []byte) *Index {
	sa := make([]int, len(data))
	for i, _ := range sa {
		sa[i] = i
	}
	x := &index{data, sa}
	sort.Sort(x)
	return (*Index)(x)
}

func (x *Index) at(i int) []byte {
	return x.data[x.sa[i]:]
}

// Binary search according to "A Method of Programming", E.W. Dijkstra.
func (x *Index) search(s []byte) int {
	i, j := 0, len(x.sa)
	// i < j for non-empty x
	for i+1 < j {
		// 0 <= i < j <= len(x.sa) && (x.at(i) <= s < x.at(j) || (s is not in x))
		h := i + (j-i)/2 // i < h < j
		if bytes.Compare(x.at(h), s) <= 0 {
			i = h
		} else { // s < x.at(h)
			j = h
		}
	}
	// i+1 == j for non-empty x
	return i
}

// Lookup returns an unsorted list of at most n indices where the byte string s
// occurs in the indexed data. If n < 0, all occurrences are returned.
// The result is nil if s is empty, s is not found, or n == 0.
// Lookup time is O((log(N) + len(res))*len(s)) where N is the
// size of the indexed data.
//
func (x *Index) Lookup(s []byte, n int) []int {
	var res vector.IntVector

	if len(s) > 0 && n != 0 {
		// find matching suffix index i
		i := x.search(s)
		// x.at(i) <= s < x.at(i+1)

		// ignore the first suffix if it is < s
		if i < len(x.sa) && bytes.Compare(x.at(i), s) < 0 {
			i++
		}

		// collect the following suffixes with matching prefixes
		for (n < 0 || len(res) < n) && i < len(x.sa) && bytes.HasPrefix(x.at(i), s) {
			res.Push(x.sa[i])
			i++
		}
	}

	return res
}

// index is like Index; it is only used to hide the sort.Interface methods
type index struct {
	data []byte
	sa   []int
}

// index implements sort.Interface

func (x *index) Len() int           { return len(x.sa) }
func (x *index) Less(i, j int) bool { return bytes.Compare(x.at(i), x.at(j)) < 0 }
func (x *index) Swap(i, j int)      { x.sa[i], x.sa[j] = x.sa[j], x.sa[i] }
func (a *index) at(i int) []byte    { return a.data[a.sa[i]:] }
<commit_msg>suffixarray: cleanup per suggestion from Roger Peppe<commit_after>// Copyright 2010 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The suffixarray package implements substring search in logarithmic time
// using an in-memory suffix array.
//
// Example use:
//
//	// create index for some data
//	index := suffixarray.New(data)
//
//	// lookup byte slice s
//	offsets1 := index.Lookup(s, -1) // the list of all indices where s occurs in data
//	offsets2 := index.Lookup(s, 3)  // the list of at most 3 indices where s occurs in data
//
package suffixarray

import (
	"bytes"
	"container/vector"
	"sort"
)

// BUG(gri): For larger data (10MB) which contains very long (say 100000)
// contiguous sequences of identical bytes, index creation time will be extremely slow.

// TODO(gri): Use a more sophisticated algorithm to create the suffix array.

// Index implements a suffix array for fast substring search.
type Index struct {
	data []byte
	sa   []int // suffix array for data
}

// New creates a new Index for data.
// Index creation time is approximately O(N*log(N)) for N = len(data).
//
func New(data []byte) *Index {
	sa := make([]int, len(data))
	for i, _ := range sa {
		sa[i] = i
	}
	x := &Index{data, sa}
	sort.Sort((*index)(x))
	return x
}

func (x *Index) at(i int) []byte {
	return x.data[x.sa[i]:]
}

// Binary search according to "A Method of Programming", E.W. Dijkstra.
func (x *Index) search(s []byte) int {
	i, j := 0, len(x.sa)
	// i < j for non-empty x
	for i+1 < j {
		// 0 <= i < j <= len(x.sa) && (x.at(i) <= s < x.at(j) || (s is not in x))
		h := i + (j-i)/2 // i < h < j
		if bytes.Compare(x.at(h), s) <= 0 {
			i = h
		} else { // s < x.at(h)
			j = h
		}
	}
	// i+1 == j for non-empty x
	return i
}

// Lookup returns an unsorted list of at most n indices where the byte string s
// occurs in the indexed data. If n < 0, all occurrences are returned.
// The result is nil if s is empty, s is not found, or n == 0.
// Lookup time is O((log(N) + len(result))*len(s)) where N is the
// size of the indexed data.
//
func (x *Index) Lookup(s []byte, n int) []int {
	var res vector.IntVector

	if len(s) > 0 && n != 0 {
		// find matching suffix index i
		i := x.search(s)
		// x.at(i) <= s < x.at(i+1)

		// ignore the first suffix if it is < s
		if i < len(x.sa) && bytes.Compare(x.at(i), s) < 0 {
			i++
		}

		// collect the following suffixes with matching prefixes
		for (n < 0 || len(res) < n) && i < len(x.sa) && bytes.HasPrefix(x.at(i), s) {
			res.Push(x.sa[i])
			i++
		}
	}

	return res
}

// index is used to hide the sort.Interface
type index Index

func (x *index) Len() int           { return len(x.sa) }
func (x *index) Less(i, j int) bool { return bytes.Compare(x.at(i), x.at(j)) < 0 }
func (x *index) Swap(i, j int)      { x.sa[i], x.sa[j] = x.sa[j], x.sa[i] }
func (a *index) at(i int) []byte    { return a.data[a.sa[i]:] }
<|endoftext|>