<commit_before>/*
Copyright 2016 The Kubernetes 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 main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"path"
	"syscall"

	"github.com/golang/glog"
	"github.com/kubernetes-incubator/external-storage/lib/controller"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/strategicpatch"
	"k8s.io/apimachinery/pkg/util/uuid"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"

	"github.com/dustin/go-humanize"
	"github.com/virtuozzo/goploop-cli"
	"github.com/virtuozzo/ploop-flexvol/vstorage"
)

const (
	parentProvisionerAnn = "vzFSParentProvisioner"
	vzShareAnn           = "vzShare"
)

type vzFSProvisioner struct {
	// Kubernetes Client. Use to retrieve secrets with Virtuozzo Storage credentials
	client kubernetes.Interface
}

func newVzFSProvisioner(client kubernetes.Interface) controller.Provisioner {
	return &vzFSProvisioner{
		client: client,
	}
}

var _ controller.Provisioner = &vzFSProvisioner{}

const provisionerDir = "/export/virtuozzo-provisioner/"
const mountDir = provisionerDir + "mnt/"

func prepareVstorage(options map[string]string, clusterName string, clusterPassword string) error {
	mount := mountDir + clusterName
	mounted, _ := vstorage.IsVstorage(mount)
	if mounted {
		return nil
	}

	if err := os.MkdirAll(mount, 0755); err != nil {
		return err
	}

	v := vstorage.Vstorage{clusterName}
	p, _ := v.Mountpoint()
	if p != "" {
		return syscall.Mount(p, mount, "", syscall.MS_BIND, "")
	}

	if err := v.Auth(clusterPassword); err != nil {
		return err
	}
	if err := v.Mount(mount); err != nil {
		return err
	}

	return nil
}

func createPloop(mount string, options map[string]string) error {
	var (
		volumePath, volumeID, size string
	)

	for k, v := range options {
		switch k {
		case "volumePath":
			volumePath = v
		case "volumeID":
			volumeID = v
		case "size":
			size = v
		case "vzsReplicas":
		case "vzsFailureDomain":
		case "vzsEncoding":
		case "vzsTier":
		case "kubernetes.io/readwrite":
		case "kubernetes.io/fsType":
		default:
		}
	}

	if volumePath == "" {
		return fmt.Errorf("volumePath isn't specified")
	}

	if volumeID == "" {
		return fmt.Errorf("volumeID isn't specified")
	}

	if size == "" {
		return fmt.Errorf("size isn't specified")
	}

	// get a human readable size from the map
	bytes, _ := humanize.ParseBytes(size)

	// ploop driver takes kilobytes, so convert it
	volumeSize := bytes / 1024

	// create ploop deltas path
	if err := os.MkdirAll(path.Join(mount, options["deltaPath"]), 0755); err != nil {
		return err
	}

	ploopPath := path.Join(mount, options["volumePath"], options["volumeID"])
	deltaPath := path.Join(mount, options["deltasPath"], options["volumeID"])
	// Create the ploop volume
	_, err := ploop.PloopVolumeCreate(ploopPath, volumeSize, deltaPath)
	if err != nil {
		return err
	}

	for k, v := range options {
		attr := ""
		switch k {
		case "vzsReplicas":
			attr = "replicas"
		case "vzsTier":
			attr = "tier"
		case "vzsEncoding":
			attr = "encoding"
		case "vzsFailureDomain":
			attr = "failure-domain"
		}
		if attr != "" {
			cmd := "vstorage"
			args := []string{"set-attr", "-R", ploopPath,
				fmt.Sprintf("%s=%s", attr, v)}
			err = exec.Command(cmd, args...).Run()
		}

		if err != nil {
			os.RemoveAll(ploopPath)
			return fmt.Errorf("Unable to set %s to %s: %v", attr, v, err)
		}
	}

	return nil
}

func (p *vzFSProvisioner) patchSecret(oldSecret, newSecret *v1.Secret) error {
	oldData, err := json.Marshal(oldSecret)
	if err != nil {
		glog.Errorf("failed to marshal secret %s: %v", newSecret.Name, err)
		return err
	}
	newData, err := json.Marshal(newSecret)
	if err != nil {
		glog.Errorf("failed to marshal secret patch %s: %v", newSecret.Name, err)
		return err
	}

	patchBytes, err := strategicpatch.CreateTwoWayMergePatch(oldData, newData, v1.Secret{})
	if err != nil {
		glog.Errorf("failed to create patch for secret %s: %v", newSecret.Name, err)
		return err
	}
	glog.Infof("Secret %s patch: %s", newSecret.Name, string(patchBytes))

	_, err = p.client.Core().Secrets(newSecret.ObjectMeta.Namespace).Patch(newSecret.Name, types.StrategicMergePatchType, patchBytes)
	return err
}

// Provision creates a storage asset and returns a PV object representing it.
func (p *vzFSProvisioner) Provision(options controller.VolumeOptions) (*v1.PersistentVolume, error) {
	modes := options.PVC.Spec.AccessModes
	if len(modes) == 0 {
		// if AccessModes field is absent, ReadWriteOnce is used by default
		modes = append(modes, v1.ReadWriteOnce)
	} else {
		if len(modes) != 1 || modes[0] != v1.ReadWriteOnce {
			return nil, fmt.Errorf("Virtuozzo flexvolume provisioner supports only ReadWriteOnce access mode")
		}
	}
	capacity := options.PVC.Spec.Resources.Requests[v1.ResourceName(v1.ResourceStorage)]
	bytes := capacity.Value()

	if options.PVC.Spec.Selector != nil {
		return nil, fmt.Errorf("claim Selector is not supported")
	}
	share := fmt.Sprintf("kubernetes-dynamic-pvc-%s", options.PVC.UID)

	glog.Infof("Add %s %s", share, humanize.Bytes(uint64(bytes)))

	storageClassOptions := map[string]string{}
	for k, v := range options.Parameters {
		storageClassOptions[k] = v
	}

	storageClassOptions["volumeID"] = share
	storageClassOptions["size"] = fmt.Sprintf("%d", bytes)
	secretName := storageClassOptions["secretName"]
	delete(storageClassOptions, "secretName")

	secret, err := p.client.Core().Secrets(options.PVC.Namespace).Get(secretName, metav1.GetOptions{})
	if err != nil {
		return nil, err
	}

	name := string(secret.Data["clusterName"][:len(secret.Data["clusterName"])])
	password := string(secret.Data["clusterPassword"][:len(secret.Data["clusterPassword"])])
	if err := prepareVstorage(storageClassOptions, name, password); err != nil {
		return nil, err
	}

	if err := createPloop(mountDir+name, storageClassOptions); err != nil {
		return nil, err
	}

	u := uuid.NewUUID()
	storageClassOptions["clusterName"] = name
	pv := &v1.PersistentVolume{
		ObjectMeta: metav1.ObjectMeta{
			Name: options.PVName,
			UID:  u,
			Annotations: map[string]string{
				parentProvisionerAnn: *provisionerID,
				vzShareAnn:           share,
			},
		},
		Spec: v1.PersistentVolumeSpec{
			PersistentVolumeReclaimPolicy: options.PersistentVolumeReclaimPolicy,
			AccessModes:                   modes,
			Capacity: v1.ResourceList{
				v1.ResourceName(v1.ResourceStorage): options.PVC.Spec.Resources.Requests[v1.ResourceName(v1.ResourceStorage)],
			},
			PersistentVolumeSource: v1.PersistentVolumeSource{
				FlexVolume: &v1.FlexVolumeSource{
					Driver:    "virtuozzo/ploop",
					SecretRef: &v1.LocalObjectReference{Name: secretName},
					Options:   storageClassOptions,
				},
			},
		},
	}

	newSecret := *secret
	finalizer := fmt.Sprintf("virtuozzo.com/%s-pv", u)
	idx := -1
	for i, f := range newSecret.Finalizers {
		if f == finalizer {
			idx = i
			break
		}
	}
	if idx == -1 {
		newSecret.Finalizers = append(newSecret.Finalizers, finalizer)
		if err = p.patchSecret(secret, &newSecret); err != nil {
			glog.Errorf("Failed to update finalizers in secret: %s", secretName)
			return nil, err
		}
	}
	glog.Infof("successfully created virtuozzo storage share: %s", share)
	return pv, nil
}

// Delete removes the storage asset that was created by Provision represented
// by the given PV.
func (p *vzFSProvisioner) Delete(volume *v1.PersistentVolume) error {
	ann, ok := volume.Annotations[parentProvisionerAnn]
	if !ok {
		return errors.New("Parent provisioner name annotation not found on PV")
	}
	if ann != *provisionerID {
		return &controller.IgnoredError{"parent provisioner name annotation on PV does not match ours"}
	}
	share, ok := volume.Annotations[vzShareAnn]
	if !ok {
		return errors.New("vz share annotation not found on PV")
	}

	secretName := volume.Spec.PersistentVolumeSource.FlexVolume.SecretRef.Name
	options := volume.Spec.PersistentVolumeSource.FlexVolume.Options

	secret, err := p.client.Core().Secrets(volume.Spec.ClaimRef.Namespace).Get(secretName, metav1.GetOptions{})
	if err != nil {
		return err
	}

	name := string(secret.Data["clusterName"][:len(secret.Data["clusterName"])])
	password := string(secret.Data["clusterPassword"][:len(secret.Data["clusterPassword"])])
	mount := mountDir + name
	if err := prepareVstorage(options, name, password); err != nil {
		return err
	}

	ploopPath := path.Join(mount, options["volumePath"], options["volumeID"])
	vol, err := ploop.PloopVolumeOpen(ploopPath)
	if err != nil {
		return err
	}
	glog.Infof("Delete: %s", ploopPath)
	if err = vol.Delete(); err != nil {
		return err
	}

	newSecret := *secret
	finalizer := fmt.Sprintf("virtuozzo.com/%s-pv", volume.UID)
	idx := -1
	for i, f := range newSecret.Finalizers {
		if f == finalizer {
			idx = i
			break
		}
	}
	if idx != -1 {
		newSecret.Finalizers = append(newSecret.Finalizers[:idx], newSecret.Finalizers[idx+1:]...)
		if err = p.patchSecret(secret, &newSecret); err != nil {
			glog.Errorf("Failed to update finalizers in secret: %s", secretName)
			return err
		}
	}

	glog.Infof("successfully delete virtuozzo storage share: %s", share)

	return nil
}

var (
	master          = flag.String("master", "", "Master URL")
	kubeconfig      = flag.String("kubeconfig", "", "Absolute path to the kubeconfig")
	provisionerID   = flag.String("id", "", "Unique provisioner id")
	provisionerName = flag.String("name", "virtuozzo.com/virtuozzo-storage", "Unique provisioner name")
)

func main() {
	flag.Parse()
	flag.Set("logtostderr", "true")
	if *provisionerID == "" {
		glog.Fatalf("You should provide unique provisioner name!")
	}

	var config *rest.Config
	var err error
	if *master != "" || *kubeconfig != "" {
		config, err = clientcmd.BuildConfigFromFlags(*master, *kubeconfig)
	} else {
		config, err = rest.InClusterConfig()
	}

	if err != nil {
		glog.Fatalf("Failed to create config: %v", err)
	}
	clientset, err := kubernetes.NewForConfig(config)
	if err != nil {
		glog.Fatalf("Failed to create client: %v", err)
	}

	// The controller needs to know what the server version is because out-of-tree
	// provisioners aren't officially supported until 1.5
	serverVersion, err := clientset.Discovery().ServerVersion()
	if err != nil {
		glog.Fatalf("Error getting server version: %v", err)
	}

	// Create the provisioner: it implements the Provisioner interface expected by
	// the controller
	vzFSProvisioner := newVzFSProvisioner(clientset)

	// Start the provision controller which will dynamically provision Virtuozzo Storage PVs
	pc := controller.NewProvisionController(clientset,
		*provisionerName,
		vzFSProvisioner,
		serverVersion.GitVersion,
	)

	pc.Run(wait.NeverStop)
}
<commit_msg>fix typo in creating deltas path dir<commit_after>/*
Copyright 2016 The Kubernetes 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 main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"path"
	"syscall"

	"github.com/golang/glog"
	"github.com/kubernetes-incubator/external-storage/lib/controller"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/strategicpatch"
	"k8s.io/apimachinery/pkg/util/uuid"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"

	"github.com/dustin/go-humanize"
	"github.com/virtuozzo/goploop-cli"
	"github.com/virtuozzo/ploop-flexvol/vstorage"
)

const (
	parentProvisionerAnn = "vzFSParentProvisioner"
	vzShareAnn           = "vzShare"
)

type vzFSProvisioner struct {
	// Kubernetes Client. Use to retrieve secrets with Virtuozzo Storage credentials
	client kubernetes.Interface
}

func newVzFSProvisioner(client kubernetes.Interface) controller.Provisioner {
	return &vzFSProvisioner{
		client: client,
	}
}

var _ controller.Provisioner = &vzFSProvisioner{}

const provisionerDir = "/export/virtuozzo-provisioner/"
const mountDir = provisionerDir + "mnt/"

func prepareVstorage(options map[string]string, clusterName string, clusterPassword string) error {
	mount := mountDir + clusterName
	mounted, _ := vstorage.IsVstorage(mount)
	if mounted {
		return nil
	}

	if err := os.MkdirAll(mount, 0755); err != nil {
		return err
	}

	v := vstorage.Vstorage{clusterName}
	p, _ := v.Mountpoint()
	if p != "" {
		return syscall.Mount(p, mount, "", syscall.MS_BIND, "")
	}

	if err := v.Auth(clusterPassword); err != nil {
		return err
	}
	if err := v.Mount(mount); err != nil {
		return err
	}

	return nil
}

func createPloop(mount string, options map[string]string) error {
	var (
		volumePath, volumeID, size string
	)

	for k, v := range options {
		switch k {
		case "volumePath":
			volumePath = v
		case "volumeID":
			volumeID = v
		case "size":
			size = v
		case "vzsReplicas":
		case "vzsFailureDomain":
		case "vzsEncoding":
		case "vzsTier":
		case "kubernetes.io/readwrite":
		case "kubernetes.io/fsType":
		default:
		}
	}

	if volumePath == "" {
		return fmt.Errorf("volumePath isn't specified")
	}

	if volumeID == "" {
		return fmt.Errorf("volumeID isn't specified")
	}

	if size == "" {
		return fmt.Errorf("size isn't specified")
	}

	// get a human readable size from the map
	bytes, _ := humanize.ParseBytes(size)

	// ploop driver takes kilobytes, so convert it
	volumeSize := bytes / 1024

	// create ploop deltas path
	if err := os.MkdirAll(path.Join(mount, options["deltasPath"]), 0755); err != nil {
		return err
	}

	ploopPath := path.Join(mount, options["volumePath"], options["volumeID"])
	deltaPath := path.Join(mount, options["deltasPath"], options["volumeID"])
	// Create the ploop volume
	_, err := ploop.PloopVolumeCreate(ploopPath, volumeSize, deltaPath)
	if err != nil {
		return err
	}

	for k, v := range options {
		attr := ""
		switch k {
		case "vzsReplicas":
			attr = "replicas"
		case "vzsTier":
			attr = "tier"
		case "vzsEncoding":
			attr = "encoding"
		case "vzsFailureDomain":
			attr = "failure-domain"
		}
		if attr != "" {
			cmd := "vstorage"
			args := []string{"set-attr", "-R", ploopPath,
				fmt.Sprintf("%s=%s", attr, v)}
			err = exec.Command(cmd, args...).Run()
		}

		if err != nil {
			os.RemoveAll(ploopPath)
			return fmt.Errorf("Unable to set %s to %s: %v", attr, v, err)
		}
	}

	return nil
}

func (p *vzFSProvisioner) patchSecret(oldSecret, newSecret *v1.Secret) error {
	oldData, err := json.Marshal(oldSecret)
	if err != nil {
		glog.Errorf("failed to marshal secret %s: %v", newSecret.Name, err)
		return err
	}
	newData, err := json.Marshal(newSecret)
	if err != nil {
		glog.Errorf("failed to marshal secret patch %s: %v", newSecret.Name, err)
		return err
	}

	patchBytes, err := strategicpatch.CreateTwoWayMergePatch(oldData, newData, v1.Secret{})
	if err != nil {
		glog.Errorf("failed to create patch for secret %s: %v", newSecret.Name, err)
		return err
	}
	glog.Infof("Secret %s patch: %s", newSecret.Name, string(patchBytes))

	_, err = p.client.Core().Secrets(newSecret.ObjectMeta.Namespace).Patch(newSecret.Name, types.StrategicMergePatchType, patchBytes)
	return err
}

// Provision creates a storage asset and returns a PV object representing it.
func (p *vzFSProvisioner) Provision(options controller.VolumeOptions) (*v1.PersistentVolume, error) {
	modes := options.PVC.Spec.AccessModes
	if len(modes) == 0 {
		// if AccessModes field is absent, ReadWriteOnce is used by default
		modes = append(modes, v1.ReadWriteOnce)
	} else {
		if len(modes) != 1 || modes[0] != v1.ReadWriteOnce {
			return nil, fmt.Errorf("Virtuozzo flexvolume provisioner supports only ReadWriteOnce access mode")
		}
	}
	capacity := options.PVC.Spec.Resources.Requests[v1.ResourceName(v1.ResourceStorage)]
	bytes := capacity.Value()

	if options.PVC.Spec.Selector != nil {
		return nil, fmt.Errorf("claim Selector is not supported")
	}
	share := fmt.Sprintf("kubernetes-dynamic-pvc-%s", options.PVC.UID)

	glog.Infof("Add %s %s", share, humanize.Bytes(uint64(bytes)))

	storageClassOptions := map[string]string{}
	for k, v := range options.Parameters {
		storageClassOptions[k] = v
	}

	storageClassOptions["volumeID"] = share
	storageClassOptions["size"] = fmt.Sprintf("%d", bytes)
	secretName := storageClassOptions["secretName"]
	delete(storageClassOptions, "secretName")

	secret, err := p.client.Core().Secrets(options.PVC.Namespace).Get(secretName, metav1.GetOptions{})
	if err != nil {
		return nil, err
	}

	name := string(secret.Data["clusterName"][:len(secret.Data["clusterName"])])
	password := string(secret.Data["clusterPassword"][:len(secret.Data["clusterPassword"])])
	if err := prepareVstorage(storageClassOptions, name, password); err != nil {
		return nil, err
	}

	if err := createPloop(mountDir+name, storageClassOptions); err != nil {
		return nil, err
	}

	u := uuid.NewUUID()
	storageClassOptions["clusterName"] = name
	pv := &v1.PersistentVolume{
		ObjectMeta: metav1.ObjectMeta{
			Name: options.PVName,
			UID:  u,
			Annotations: map[string]string{
				parentProvisionerAnn: *provisionerID,
				vzShareAnn:           share,
			},
		},
		Spec: v1.PersistentVolumeSpec{
			PersistentVolumeReclaimPolicy: options.PersistentVolumeReclaimPolicy,
			AccessModes:                   modes,
			Capacity: v1.ResourceList{
				v1.ResourceName(v1.ResourceStorage): options.PVC.Spec.Resources.Requests[v1.ResourceName(v1.ResourceStorage)],
			},
			PersistentVolumeSource: v1.PersistentVolumeSource{
				FlexVolume: &v1.FlexVolumeSource{
					Driver:    "virtuozzo/ploop",
					SecretRef: &v1.LocalObjectReference{Name: secretName},
					Options:   storageClassOptions,
				},
			},
		},
	}

	newSecret := *secret
	finalizer := fmt.Sprintf("virtuozzo.com/%s-pv", u)
	idx := -1
	for i, f := range newSecret.Finalizers {
		if f == finalizer {
			idx = i
			break
		}
	}
	if idx == -1 {
		newSecret.Finalizers = append(newSecret.Finalizers, finalizer)
		if err = p.patchSecret(secret, &newSecret); err != nil {
			glog.Errorf("Failed to update finalizers in secret: %s", secretName)
			return nil, err
		}
	}
	glog.Infof("successfully created virtuozzo storage share: %s", share)
	return pv, nil
}

// Delete removes the storage asset that was created by Provision represented
// by the given PV.
func (p *vzFSProvisioner) Delete(volume *v1.PersistentVolume) error {
	ann, ok := volume.Annotations[parentProvisionerAnn]
	if !ok {
		return errors.New("Parent provisioner name annotation not found on PV")
	}
	if ann != *provisionerID {
		return &controller.IgnoredError{"parent provisioner name annotation on PV does not match ours"}
	}
	share, ok := volume.Annotations[vzShareAnn]
	if !ok {
		return errors.New("vz share annotation not found on PV")
	}

	secretName := volume.Spec.PersistentVolumeSource.FlexVolume.SecretRef.Name
	options := volume.Spec.PersistentVolumeSource.FlexVolume.Options

	secret, err := p.client.Core().Secrets(volume.Spec.ClaimRef.Namespace).Get(secretName, metav1.GetOptions{})
	if err != nil {
		return err
	}

	name := string(secret.Data["clusterName"][:len(secret.Data["clusterName"])])
	password := string(secret.Data["clusterPassword"][:len(secret.Data["clusterPassword"])])
	mount := mountDir + name
	if err := prepareVstorage(options, name, password); err != nil {
		return err
	}

	ploopPath := path.Join(mount, options["volumePath"], options["volumeID"])
	vol, err := ploop.PloopVolumeOpen(ploopPath)
	if err != nil {
		return err
	}
	glog.Infof("Delete: %s", ploopPath)
	if err = vol.Delete(); err != nil {
		return err
	}

	newSecret := *secret
	finalizer := fmt.Sprintf("virtuozzo.com/%s-pv", volume.UID)
	idx := -1
	for i, f := range newSecret.Finalizers {
		if f == finalizer {
			idx = i
			break
		}
	}
	if idx != -1 {
		newSecret.Finalizers = append(newSecret.Finalizers[:idx], newSecret.Finalizers[idx+1:]...)
		if err = p.patchSecret(secret, &newSecret); err != nil {
			glog.Errorf("Failed to update finalizers in secret: %s", secretName)
			return err
		}
	}

	glog.Infof("successfully delete virtuozzo storage share: %s", share)

	return nil
}

var (
	master          = flag.String("master", "", "Master URL")
	kubeconfig      = flag.String("kubeconfig", "", "Absolute path to the kubeconfig")
	provisionerID   = flag.String("id", "", "Unique provisioner id")
	provisionerName = flag.String("name", "virtuozzo.com/virtuozzo-storage", "Unique provisioner name")
)

func main() {
	flag.Parse()
	flag.Set("logtostderr", "true")
	if *provisionerID == "" {
		glog.Fatalf("You should provide unique provisioner name!")
	}

	var config *rest.Config
	var err error
	if *master != "" || *kubeconfig != "" {
		config, err = clientcmd.BuildConfigFromFlags(*master, *kubeconfig)
	} else {
		config, err = rest.InClusterConfig()
	}

	if err != nil {
		glog.Fatalf("Failed to create config: %v", err)
	}
	clientset, err := kubernetes.NewForConfig(config)
	if err != nil {
		glog.Fatalf("Failed to create client: %v", err)
	}

	// The controller needs to know what the server version is because out-of-tree
	// provisioners aren't officially supported until 1.5
	serverVersion, err := clientset.Discovery().ServerVersion()
	if err != nil {
		glog.Fatalf("Error getting server version: %v", err)
	}

	// Create the provisioner: it implements the Provisioner interface expected by
	// the controller
	vzFSProvisioner := newVzFSProvisioner(clientset)

	// Start the provision controller which will dynamically provision Virtuozzo Storage PVs
	pc := controller.NewProvisionController(clientset,
		*provisionerName,
		vzFSProvisioner,
		serverVersion.GitVersion,
	)

	pc.Run(wait.NeverStop)
}
<|endoftext|>


<commit_before>package lint

import (
	"errors"
	"strconv"
	"strings"

	"github.com/golang/protobuf/v2/reflect/protoreflect"
	descriptorpb "github.com/golang/protobuf/v2/types/descriptor"
)

var (
	// ErrLocationNotFound is the returned error when a location is not found.
	ErrLocationNotFound = errors.New("source: location not found")
	// ErrCommentsNotFound is the returned error when a comments is not found.
	ErrCommentsNotFound = errors.New("source: comments not found")
	// ErrSourceInfoNotAvailable is the returned error when creating a source
	// but the source information is not available.
	ErrSourceInfoNotAvailable = errors.New("source: source information is not available")
)

// Comments describes a collection of comments associate with an element,
// which contains leading, trailing, and leading-detached comments, in a
// source code file.
type Comments struct {
	LeadingComments         string
	TrailingComments        string
	LeadingDetachedComments []string
}

const sep = ","

// locPath represents a path in the SourceCodeInfo_Location,
// and this serves as a map key.
// It's a string representation of a slice because slices
// cannot be map keys.
// Representation: integers separated by commas. No spaces.
// Example: [4, 3, 2, 7] --> "4,3,2,7"
// See descriptor.proto for more explantion of semantics.
type locPath string

// newLocPath return a locPath from a list of index.
func newLocPath(p ...int) locPath {
	a := []string{}
	for _, i := range p {
		a = append(a, strconv.Itoa(i))
	}
	return locPath(strings.Join(a, sep))
}

// buildLocPathMap creates a map of locPath to *descriptorpb.SourceCodeInfo_Location
// from *descriptorpb.SourceCodeInfo.
func buildLocPathMap(sci *descriptorpb.SourceCodeInfo) map[locPath]*descriptorpb.SourceCodeInfo_Location {
	m := make(map[locPath]*descriptorpb.SourceCodeInfo_Location)
	if sci == nil {
		return m
	}

	for _, loc := range sci.GetLocation() {
		var path []int
		for _, v := range loc.GetPath() {
			path = append(path, int(v))
		}
		m[newLocPath(path...)] = loc
	}
	return m
}

// DescriptorSource represents a map of locPath to *descriptorpb.SourceCodeInfo_Location.
type DescriptorSource struct {
	m map[locPath]*descriptorpb.SourceCodeInfo_Location
}

// NewDescriptorSource creates a new DescriptorSource from a FileDescriptorProto.
// If source code information is not available, returns (nil, ErrSourceInfoNotAvailable).
func NewDescriptorSource(f *descriptorpb.FileDescriptorProto) (DescriptorSource, error) {
	if f.GetSourceCodeInfo() == nil {
		return DescriptorSource{}, ErrSourceInfoNotAvailable
	}
	return DescriptorSource{m: buildLocPathMap(f.GetSourceCodeInfo())}, nil
}

// findLocationByPath returns a `Location` if found in the map,
// and (nil, ErrLocationNotFound) if not found.
func (s DescriptorSource) findLocationByPath(path []int) (Location, error) {
	l := s.m[newLocPath(path...)]
	if l == nil {
		return Location{}, ErrLocationNotFound
	}
	return newLocationFromSpan(l.GetSpan()), nil
}

// findCommentsByPath returns a `Comments` for the path. If not found, returns
// (nil, ErrCommentsNotFound).
func (s DescriptorSource) findCommentsByPath(path []int) (Comments, error) {
	l := s.m[newLocPath(path...)]
	if l == nil {
		return Comments{}, ErrCommentsNotFound
	}
	return Comments{
		LeadingComments:         l.GetLeadingComments(),
		TrailingComments:        l.GetTrailingComments(),
		LeadingDetachedComments: l.GetLeadingDetachedComments(),
	}, nil
}

func newLocationFromSpan(span []int32) Location {
	if len(span) == 4 {
		return Location{
			Start: Position{
				Line:   int(span[0]),
				Column: int(span[1]),
			},
			End: Position{
				Line:   int(span[2]),
				Column: int(span[3]),
			},
		}
	}

	if len(span) == 3 {
		return Location{
			Start: Position{
				Line:   int(span[0]),
				Column: int(span[1]),
			},
			End: Position{
				Line:   int(span[0]),
				Column: int(span[2]),
			},
		}
	}

	return Location{}
}

// SyntaxLocation returns the location of the syntax definition.
func (s DescriptorSource) SyntaxLocation() (Location, error) {
	return s.findLocationByPath([]int{syntaxTag})
}

// SyntaxComments returns the comments of the syntax definition.
func (s DescriptorSource) SyntaxComments() (Comments, error) {
	return s.findCommentsByPath([]int{syntaxTag})
}

// DescriptorLocation returns a `Location` for the given descriptor.
// If not found, returns (nil, ErrLocationNotFound).
func (s DescriptorSource) DescriptorLocation(d protoreflect.Descriptor) (Location, error) {
	return s.findLocationByPath(getPath(d))
}

func getPath(d protoreflect.Descriptor) []int {
	path := []int{}
	for p := d; !isFileDescriptor(p); p, _ = p.Parent() {
		path = append(path, p.Index(), getDescriptorTag(p))
	}
	reverseInts(path)
	return path
}

const syntaxTag = 12

var enumTagInFile = 5
var enumTagInMessage = 4
var enumValueTag = 2
var fieldTag = 2
var extensionTagInFile = 7
var extensionTagInMessage = 6
var messageTagInFile = 4
var nestedMessageTag = 3
var oneofTag = 8
var serviceTag = 6
var methodTag = 2

func getDescriptorTag(d protoreflect.Descriptor) int {
	switch d.(type) {
	case protoreflect.EnumDescriptor:
		if isTopLevelDescriptor(d) {
			return enumTagInFile
		}
		return enumTagInMessage
	case protoreflect.EnumValueDescriptor:
		return enumValueTag
	case protoreflect.FieldDescriptor:
		if isFieldExtension(d) {
			if isTopLevelDescriptor(d) {
				return extensionTagInFile
			}
			return extensionTagInMessage
		}
		return fieldTag
	case protoreflect.MessageDescriptor:
		if isTopLevelDescriptor(d) {
			return messageTagInFile
		}
		return nestedMessageTag
	case protoreflect.MethodDescriptor:
		return methodTag
	case protoreflect.OneofDescriptor:
		return oneofTag
	case protoreflect.ServiceDescriptor:
		return serviceTag
	default:
		return 0
	}
}

func isFieldExtension(d protoreflect.Descriptor) bool {
	f, ok := d.(protoreflect.FieldDescriptor)
	return ok && f.ExtendedType() != nil
}

func isFileDescriptor(d protoreflect.Descriptor) bool {
	_, ok := d.(protoreflect.FileDescriptor)
	return ok
}

func isTopLevelDescriptor(d protoreflect.Descriptor) bool {
	p, _ := d.Parent()
	_, ok := p.(protoreflect.FileDescriptor)
	return ok
}

// DescriptorComments returns a `Comments` for the given descriptor.
// If not found, returns (nil, ErrCommentsNotFound).
func (s DescriptorSource) DescriptorComments(d protoreflect.Descriptor) (Comments, error) {
	return s.findCommentsByPath(getPath(d))
}

func reverseInts(a []int) {
	for left, right := 0, len(a)-1; left < right; left, right = left+1, right-1 {
		a[left], a[right] = a[right], a[left]
	}
}

// IsRuleDisabled check if a rule is disabled for a descriptor
// in the comments.
func (s DescriptorSource) IsRuleDisabled(id RuleID, d protoreflect.Descriptor) bool {
	comments, err := s.DescriptorComments(d)
	if err != nil {
		return true
	}

	commentsToCheck := []string{
		comments.LeadingComments,
		comments.TrailingComments,
	}
	commentsToCheck = append(commentsToCheck, s.fileComments().LeadingDetachedComments...)

	return stringsContains(commentsToCheck, ruleDisablingComment(id))
}

func stringsContains(comments []string, s string) bool {
	for _, c := range comments {
		if strings.Contains(c, s) {
			return true
		}
	}
	return false
}

func ruleDisablingComment(id RuleID) string {
	name := id.Set + "." + id.Name
	if id.Set == "" || id.Set == "core" {
		name = id.Name
	}
	return "(-- api-linter: " + name + "=disabled --)"
}

func (s DescriptorSource) fileComments() Comments {
	comments, _ := s.SyntaxComments()
	return comments
}
<commit_msg>Added an error for invalid spans, returning it if a span is provided with invalid number of elements.<commit_after>package lint

import (
	"errors"
	"strconv"
	"strings"

	"github.com/golang/protobuf/v2/reflect/protoreflect"
	descriptorpb "github.com/golang/protobuf/v2/types/descriptor"
)

var (
	// ErrLocationNotFound is the returned error when a location is not found.
	ErrLocationNotFound = errors.New("source: location not found")
	// ErrCommentsNotFound is the returned error when a comments is not found.
	ErrCommentsNotFound = errors.New("source: comments not found")
	// ErrSourceInfoNotAvailable is the returned error when creating a source
	// but the source information is not available.
	ErrSourceInfoNotAvailable = errors.New("source: source information is not available")
	// ErrInvalidSpan is the returned error when a span is provided that cannot be mapped to
	// a location (i.e. it does not have either 3 or 4 elmeents)
	ErrInvalidSpan = errors.New("source: invalid span")
)

// Comments describes a collection of comments associate with an element,
// which contains leading, trailing, and leading-detached comments, in a
// source code file.
type Comments struct {
	LeadingComments         string
	TrailingComments        string
	LeadingDetachedComments []string
}

const sep = ","

// locPath represents a path in the SourceCodeInfo_Location,
// and this serves as a map key.
// It's a string representation of a slice because slices
// cannot be map keys.
// Representation: integers separated by commas. No spaces.
// Example: [4, 3, 2, 7] --> "4,3,2,7"
// See descriptor.proto for more explantion of semantics.
type locPath string

// newLocPath return a locPath from a list of index.
func newLocPath(p ...int) locPath {
	a := []string{}
	for _, i := range p {
		a = append(a, strconv.Itoa(i))
	}
	return locPath(strings.Join(a, sep))
}

// buildLocPathMap creates a map of locPath to *descriptorpb.SourceCodeInfo_Location
// from *descriptorpb.SourceCodeInfo.
func buildLocPathMap(sci *descriptorpb.SourceCodeInfo) map[locPath]*descriptorpb.SourceCodeInfo_Location {
	m := make(map[locPath]*descriptorpb.SourceCodeInfo_Location)
	if sci == nil {
		return m
	}

	for _, loc := range sci.GetLocation() {
		var path []int
		for _, v := range loc.GetPath() {
			path = append(path, int(v))
		}
		m[newLocPath(path...)] = loc
	}
	return m
}

// DescriptorSource represents a map of locPath to *descriptorpb.SourceCodeInfo_Location.
type DescriptorSource struct {
	m map[locPath]*descriptorpb.SourceCodeInfo_Location
}

// NewDescriptorSource creates a new DescriptorSource from a FileDescriptorProto.
// If source code information is not available, returns (nil, ErrSourceInfoNotAvailable).
func NewDescriptorSource(f *descriptorpb.FileDescriptorProto) (DescriptorSource, error) {
	if f.GetSourceCodeInfo() == nil {
		return DescriptorSource{}, ErrSourceInfoNotAvailable
	}
	return DescriptorSource{m: buildLocPathMap(f.GetSourceCodeInfo())}, nil
}

// findLocationByPath returns a `Location` if found in the map,
// and (nil, ErrLocationNotFound) if not found.
func (s DescriptorSource) findLocationByPath(path []int) (Location, error) {
	l := s.m[newLocPath(path...)]
	if l == nil {
		return Location{}, ErrLocationNotFound
	}
	return newLocationFromSpan(l.GetSpan())
}

// findCommentsByPath returns a `Comments` for the path. If not found, returns
// (nil, ErrCommentsNotFound).
func (s DescriptorSource) findCommentsByPath(path []int) (Comments, error) {
	l := s.m[newLocPath(path...)]
	if l == nil {
		return Comments{}, ErrCommentsNotFound
	}
	return Comments{
		LeadingComments:         l.GetLeadingComments(),
		TrailingComments:        l.GetTrailingComments(),
		LeadingDetachedComments: l.GetLeadingDetachedComments(),
	}, nil
}

func newLocationFromSpan(span []int32) (Location, error) {
	if len(span) == 4 {
		return Location{
			Start: Position{
				Line:   int(span[0]),
				Column: int(span[1]),
			},
			End: Position{
				Line:   int(span[2]),
				Column: int(span[3]),
			},
		}, nil
	}

	if len(span) == 3 {
		return Location{
			Start: Position{
				Line:   int(span[0]),
				Column: int(span[1]),
			},
			End: Position{
				Line:   int(span[0]),
				Column: int(span[2]),
			},
		}, nil
	}

	return Location{}, ErrInvalidSpan
}

// SyntaxLocation returns the location of the syntax definition.
func (s DescriptorSource) SyntaxLocation() (Location, error) {
	return s.findLocationByPath([]int{syntaxTag})
}

// SyntaxComments returns the comments of the syntax definition.
func (s DescriptorSource) SyntaxComments() (Comments, error) {
	return s.findCommentsByPath([]int{syntaxTag})
}

// DescriptorLocation returns a `Location` for the given descriptor.
// If not found, returns (nil, ErrLocationNotFound).
func (s DescriptorSource) DescriptorLocation(d protoreflect.Descriptor) (Location, error) {
	return s.findLocationByPath(getPath(d))
}

func getPath(d protoreflect.Descriptor) []int {
	path := []int{}
	for p := d; !isFileDescriptor(p); p, _ = p.Parent() {
		path = append(path, p.Index(), getDescriptorTag(p))
	}
	reverseInts(path)
	return path
}

const syntaxTag = 12

var enumTagInFile = 5
var enumTagInMessage = 4
var enumValueTag = 2
var fieldTag = 2
var extensionTagInFile = 7
var extensionTagInMessage = 6
var messageTagInFile = 4
var nestedMessageTag = 3
var oneofTag = 8
var serviceTag = 6
var methodTag = 2

func getDescriptorTag(d protoreflect.Descriptor) int {
	switch d.(type) {
	case protoreflect.EnumDescriptor:
		if isTopLevelDescriptor(d) {
			return enumTagInFile
		}
		return enumTagInMessage
	case protoreflect.EnumValueDescriptor:
		return enumValueTag
	case protoreflect.FieldDescriptor:
		if isFieldExtension(d) {
			if isTopLevelDescriptor(d) {
				return extensionTagInFile
			}
			return extensionTagInMessage
		}
		return fieldTag
	case protoreflect.MessageDescriptor:
		if isTopLevelDescriptor(d) {
			return messageTagInFile
		}
		return nestedMessageTag
	case protoreflect.MethodDescriptor:
		return methodTag
	case protoreflect.OneofDescriptor:
		return oneofTag
	case protoreflect.ServiceDescriptor:
		return serviceTag
	default:
		return 0
	}
}

func isFieldExtension(d protoreflect.Descriptor) bool {
	f, ok := d.(protoreflect.FieldDescriptor)
	return ok && f.ExtendedType() != nil
}

func isFileDescriptor(d protoreflect.Descriptor) bool {
	_, ok := d.(protoreflect.FileDescriptor)
	return ok
}

func isTopLevelDescriptor(d protoreflect.Descriptor) bool {
	p, _ := d.Parent()
	_, ok := p.(protoreflect.FileDescriptor)
	return ok
}

// DescriptorComments returns a `Comments` for the given descriptor.
// If not found, returns (nil, ErrCommentsNotFound).
func (s DescriptorSource) DescriptorComments(d protoreflect.Descriptor) (Comments, error) {
	return s.findCommentsByPath(getPath(d))
}

func reverseInts(a []int) {
	for left, right := 0, len(a)-1; left < right; left, right = left+1, right-1 {
		a[left], a[right] = a[right], a[left]
	}
}

// IsRuleDisabled check if a rule is disabled for a descriptor
// in the comments.
func (s DescriptorSource) IsRuleDisabled(id RuleID, d protoreflect.Descriptor) bool {
	comments, err := s.DescriptorComments(d)
	if err != nil {
		return true
	}

	commentsToCheck := []string{
		comments.LeadingComments,
		comments.TrailingComments,
	}
	commentsToCheck = append(commentsToCheck, s.fileComments().LeadingDetachedComments...)

	return stringsContains(commentsToCheck, ruleDisablingComment(id))
}

func stringsContains(comments []string, s string) bool {
	for _, c := range comments {
		if strings.Contains(c, s) {
			return true
		}
	}
	return false
}

func ruleDisablingComment(id RuleID) string {
	name := id.Set + "." + id.Name
	if id.Set == "" || id.Set == "core" {
		name = id.Name
	}
	return "(-- api-linter: " + name + "=disabled --)"
}

func (s DescriptorSource) fileComments() Comments {
	comments, _ := s.SyntaxComments()
	return comments
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"
	"os"
	"strconv"
	"strings"

	"github.com/gorilla/mux"
	consul "github.com/hashicorp/consul/api"
	nomad "github.com/hashicorp/nomad/api"
	"github.com/jippi/hashi-ui/backend/config"
	consul_helper "github.com/jippi/hashi-ui/backend/consul/helper"
	nomad_helper "github.com/jippi/hashi-ui/backend/nomad/helper"
	newrelic "github.com/newrelic/go-agent"
	log "github.com/sirupsen/logrus"
)

var (
	GitCommit string // Filled in by the compiler
	AppConfig *config.Config
)

func startLogging(logLevel string) {
	level, err := log.ParseLevel(logLevel)
	if err != nil {
		fmt.Printf("%s (%s)", err, logLevel)
		os.Exit(1)
	}

	log.SetLevel(level)
}

func main() {
	cfg := config.DefaultConfig()
	cfg.Parse()

	AppConfig = cfg

	startLogging(cfg.LogLevel)

	if os.Getenv("LOG_FORMAT") == "json" {
		log.SetFormatter(&logrus.JSONFormatter{
			FieldMap: logrus.FieldMap{
				logrus.FieldKeyTime:  "@timestamp",
				logrus.FieldKeyLevel: "@level",
				logrus.FieldKeyMsg:   "@message",
			},
		})
	}

	log.Infof("----------------------------------------------------------------------------------")
	log.Infof("|                                 HASHI UI                                       |")
	log.Infof("----------------------------------------------------------------------------------")
	if !cfg.HttpsEnable {
		log.Infof("| listen-address            : http://%-43s |", cfg.ListenAddress)
	} else {
		log.Infof("| listen-address            : https://%-43s  |", cfg.ListenAddress)
	}
	log.Infof("| server-certificate        : %-50s |", cfg.ServerCert)
	log.Infof("| server-key                : %-50s |", cfg.ServerKey)
	log.Infof("| site-title                : %-50s |", cfg.SiteTitle)
	log.Infof("| proxy-address             : %-50s |", cfg.ProxyAddress)
	log.Infof("| log-level                 : %-50s |", cfg.LogLevel)
	log.Infof("| new-relic-enable          : %-50v |", cfg.NewRelicEnable)
	log.Infof("| new-relic-app-name        : %-50s |", cfg.NewRelicAppName)

	if cfg.ThrottleUpdateDuration != nil {
		log.Infof("| throttle-update-duration  : %-50s |", cfg.ThrottleUpdateDuration)
	} else {
		log.Infof("| throttle-update-duration  : %-50s |", "0s")
	}

	// Nomad
	log.Infof("| nomad-enable              : %-50t |", cfg.NomadEnable)
	if cfg.NomadReadOnly {
		log.Infof("| nomad-read-only           : %-50s |", "Yes")
	} else {
		log.Infof("| nomad-read-only           : %-50s |", "No (Hashi-UI can change Nomad state)")
	}
	log.Infof("| nomad-address             : %-50s |", cfg.NomadAddress)
	log.Infof("| nomad-acl-token           : %-50s |", cfg.NomadACLToken)
	log.Infof("| nomad-ca-cert             : %-50s |", cfg.NomadCACert)
	log.Infof("| nomad-client-cert         : %-50s |", cfg.NomadClientCert)
	log.Infof("| nomad-client-key          : %-50s |", cfg.NomadClientKey)
	log.Infof("| nomad-skip-verify         : %-50t |", cfg.NomadSkipVerify)
	log.Infof("| nomad-hide-env-data       : %-50v |", cfg.NomadHideEnvData)
	if cfg.NomadSkipVerify {
		log.Infof("| nomad-skip-verify         : %-50s |", "Yes")
	} else {
		log.Infof("| nomad-skip-verify         : %-50s |", "No")
	}
	if cfg.NomadAllowStale {
		log.Infof("| nomad-allow-stale         : %-50s |", "Yes")
	} else {
		log.Infof("| nomad-allow-stale         : %-50s |", "No")
	}
	log.Infof("| nomad-color               : %-50s |", cfg.NomadColor)

	// Consul
	log.Infof("| consul-enable             : %-50t |", cfg.ConsulEnable)
	if cfg.ConsulReadOnly {
		log.Infof("| consul-read-only          : %-50s |", "Yes")
	} else {
		log.Infof("| consul-read-only          : %-50s |", "No (Hashi-UI can change Consul state)")
	}
	log.Infof("| consul-address            : %-50s |", cfg.ConsulAddress)
	log.Infof("| consul-acl-token          : %-50s |", cfg.ConsulACLToken)
	log.Infof("| consul-color              : %-50s |", cfg.ConsulColor)

	log.Infof("----------------------------------------------------------------------------------")
	log.Infof("")

	if !cfg.NomadEnable && !cfg.ConsulEnable {
		log.Fatal("Please enable at least Consul (--consul-enable) or Nomad (--nomad-enable)")
	}

	myAssetFS := assetFS()
	router := mux.NewRouter()

	var app newrelic.Application

	if cfg.NewRelicEnable {
		config := newrelic.NewConfig(cfg.NewRelicAppName, cfg.NewRelicLicenseKey)
		var err error
		app, err = newrelic.NewApplication(config)
		if err != nil {
			log.Fatalf("Could not create NewRelic application: %v", err)
		}
	}

	// create clients
	var nomadClient *nomad.Client
	var consulClient *consul.Client

	if cfg.NomadEnable {
		var err error

		log.Info("Connecting to Nomad ...")
		nomadClient, err = nomad_helper.NewRegionClient(cfg, "")
		if err != nil {
			log.Fatalf("Unable to create Nomad client: %s", err)
		}
		if _, err := nomadClient.Status().Leader(); err != nil {
			log.Fatalf("Unable to communicate with Nomad: %s", err)
		}
		log.Info("done!")
	}

	if cfg.ConsulEnable {
		var err error

		log.Info("Connecting to Consul ...")
		consulClient, err = consul_helper.NewDatacenterClient(cfg, "")
		if err != nil {
			log.Fatalf("Unable to create Consul client: %s", err)
		}
		if _, err := consulClient.Status().Leader(); err != nil {
			log.Fatalf("Unable to communicate with Consul: %s", err)
		}
		log.Info("done!")
	}

	// setup http handlers

	if cfg.NomadEnable {
		router.HandleFunc(newrelic.WrapHandleFunc(app, "/", func(w http.ResponseWriter, r *http.Request) {
			log.Infof("Redirecting / to /nomad")
			w.Write([]byte("<script>document.location.href='" + cfg.ProxyPath + "nomad'</script>"))
			return
		}))

		router.HandleFunc(newrelic.WrapHandleFunc(app, "/ws/nomad", NomadHandler(cfg, nomadClient, consulClient)))
		router.HandleFunc(newrelic.WrapHandleFunc(app, "/ws/nomad/{region}", NomadHandler(cfg, nomadClient, consulClient)))
		router.HandleFunc(newrelic.WrapHandleFunc(app, "/api/nomad/{region}", NomadAPIHandler(cfg, nomadClient, consulClient)))
		router.HandleFunc(newrelic.WrapHandleFunc(app, "/nomad/{region}/download/{path:.*}", NomadDownloadFile(cfg)))
	}

	if cfg.ConsulEnable {
		if !cfg.NomadEnable {
			router.HandleFunc(newrelic.WrapHandleFunc(app, "/", func(w http.ResponseWriter, r *http.Request) {
				log.Infof("Redirecting / to /consul")
				http.Redirect(w, r, cfg.ProxyPath+"consul", 302)
			}))
		}

		router.HandleFunc(newrelic.WrapHandleFunc(app, "/ws/consul", ConsulHandler(cfg, nomadClient, consulClient)))
		router.HandleFunc(newrelic.WrapHandleFunc(app, "/ws/consul/{region}", ConsulHandler(cfg, nomadClient, consulClient)))
	}

	router.HandleFunc(newrelic.WrapHandleFunc(app, "/_status", StatusHandler(cfg, nomadClient, consulClient)))

	router.PathPrefix("/").HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		responseFile := "/index.html"

		if idx := strings.Index(r.URL.Path, "static/"); idx != -1 {
			responseFile = r.URL.Path[idx:]
		}

		if idx := strings.Index(r.URL.Path, "favicon.png"); idx != -1 {
			responseFile = "/favicon.png"
		}

		if idx := strings.Index(r.URL.Path, "config.js"); idx != -1 {
			response := make([]string, 0)
			response = append(response, "window.HASHI_DEV=false")
			response = append(response, fmt.Sprintf("window.GIT_HASH='%s'", GitCommit))

			response = append(response, fmt.Sprintf("window.CONSUL_ENABLED=%s", strconv.FormatBool(cfg.ConsulEnable)))
			response = append(response, fmt.Sprintf("window.CONSUL_READ_ONLY=%s", strconv.FormatBool(cfg.ConsulReadOnly)))

			response = append(response, fmt.Sprintf("window.NOMAD_ENABLED=%s", strconv.FormatBool(cfg.NomadEnable)))
			response = append(response, fmt.Sprintf("window.NOMAD_READ_ONLY=%s", strconv.FormatBool(cfg.NomadReadOnly)))

			response = append(response, fmt.Sprintf("window.NOMAD_COLOR='%s'", cfg.NomadColor))
			response = append(response, fmt.Sprintf("window.CONSUL_COLOR='%s'", cfg.ConsulColor))
			response = append(response, fmt.Sprintf("window.SITE_TITLE='%s'", cfg.SiteTitle))

			enabledServices := make([]string, 0)
			if cfg.ConsulEnable {
				enabledServices = append(enabledServices, "'consul'")
			}
			if cfg.NomadEnable {
				enabledServices = append(enabledServices, "'nomad'")
			}

			response = append(response, fmt.Sprintf("window.ENABLED_SERVICES=[%s]", strings.Join(enabledServices, ",")))
			response = append(response, fmt.Sprintf("window.HASHI_PATH_PREFIX='%s';", cfg.ProxyPath))
			response = append(response, "window.HASHI_ENDPOINT_PORT=window.HASHI_ENDPOINT_PORT || document.location.port")

			w.Header().Set("Content-Type", "application/javascript")

			// never ever cache config.js
			w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
			w.Header().Set("Pragma", "no-cache")
			w.Header().Set("Expires", "0")

			w.Write([]byte(strings.Join(response, "\n")))
			return
		}

		if bs, err := myAssetFS.Open(responseFile); err != nil {
			log.Errorf("%s: %s", responseFile, err)
		} else {
			stat, err := bs.Stat()
			if err != nil {
				w.WriteHeader(http.StatusNotFound)
				log.Errorf("Failed to stat %s: %s", responseFile, err)
			} else {
				http.ServeContent(w, r, responseFile[1:], stat.ModTime(), bs)
			}
		}
	})

	log.Infof("Listening ...")
	var err error
	if cfg.HttpsEnable {
		if cfg.ServerCert == "" || cfg.ServerKey == "" {
			log.Fatal("Using https protocol but server certificate or key were not specified.")
		}
		err = http.ListenAndServeTLS(cfg.ListenAddress, cfg.ServerCert, cfg.ServerKey, router)
	} else {
		err = http.ListenAndServe(cfg.ListenAddress, router)
	}
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>fixing references to logrus import<commit_after>package main

import (
	"fmt"
	"net/http"
	"os"
	"strconv"
	"strings"

	"github.com/gorilla/mux"
	consul "github.com/hashicorp/consul/api"
	nomad "github.com/hashicorp/nomad/api"
	"github.com/jippi/hashi-ui/backend/config"
	consul_helper "github.com/jippi/hashi-ui/backend/consul/helper"
	nomad_helper "github.com/jippi/hashi-ui/backend/nomad/helper"
	newrelic "github.com/newrelic/go-agent"
	log "github.com/sirupsen/logrus"
)

var (
	GitCommit string // Filled in by the compiler
	AppConfig *config.Config
)

func startLogging(logLevel string) {
	level, err := log.ParseLevel(logLevel)
	if err != nil {
		fmt.Printf("%s (%s)", err, logLevel)
		os.Exit(1)
	}

	log.SetLevel(level)
}

func main() {
	cfg := config.DefaultConfig()
	cfg.Parse()

	AppConfig = cfg

	startLogging(cfg.LogLevel)

	if os.Getenv("LOG_FORMAT") == "json" {
		log.SetFormatter(&log.JSONFormatter{
			FieldMap: log.FieldMap{
				log.FieldKeyTime:  "@timestamp",
				log.FieldKeyLevel: "@level",
				log.FieldKeyMsg:   "@message",
			},
		})
	}

	log.Infof("----------------------------------------------------------------------------------")
	log.Infof("|                                 HASHI UI                                       |")
	log.Infof("----------------------------------------------------------------------------------")
	if !cfg.HttpsEnable {
		log.Infof("| listen-address            : http://%-43s |", cfg.ListenAddress)
	} else {
		log.Infof("| listen-address            : https://%-43s  |", cfg.ListenAddress)
	}
	log.Infof("| server-certificate        : %-50s |", cfg.ServerCert)
	log.Infof("| server-key                : %-50s |", cfg.ServerKey)
	log.Infof("| site-title                : %-50s |", cfg.SiteTitle)
	log.Infof("| proxy-address             : %-50s |", cfg.ProxyAddress)
	log.Infof("| log-level                 : %-50s |", cfg.LogLevel)
	log.Infof("| new-relic-enable          : %-50v |", cfg.NewRelicEnable)
	log.Infof("| new-relic-app-name        : %-50s |", cfg.NewRelicAppName)

	if cfg.ThrottleUpdateDuration != nil {
		log.Infof("| throttle-update-duration  : %-50s |", cfg.ThrottleUpdateDuration)
	} else {
		log.Infof("| throttle-update-duration  : %-50s |", "0s")
	}

	// Nomad
	log.Infof("| nomad-enable              : %-50t |", cfg.NomadEnable)
	if cfg.NomadReadOnly {
		log.Infof("| nomad-read-only           : %-50s |", "Yes")
	} else {
		log.Infof("| nomad-read-only           : %-50s |", "No (Hashi-UI can change Nomad state)")
	}
	log.Infof("| nomad-address             : %-50s |", cfg.NomadAddress)
	log.Infof("| nomad-acl-token           : %-50s |", cfg.NomadACLToken)
	log.Infof("| nomad-ca-cert             : %-50s |", cfg.NomadCACert)
	log.Infof("| nomad-client-cert         : %-50s |", cfg.NomadClientCert)
	log.Infof("| nomad-client-key          : %-50s |", cfg.NomadClientKey)
	log.Infof("| nomad-skip-verify         : %-50t |", cfg.NomadSkipVerify)
	log.Infof("| nomad-hide-env-data       : %-50v |", cfg.NomadHideEnvData)
	if cfg.NomadSkipVerify {
		log.Infof("| nomad-skip-verify         : %-50s |", "Yes")
	} else {
		log.Infof("| nomad-skip-verify         : %-50s |", "No")
	}
	if cfg.NomadAllowStale {
		log.Infof("| nomad-allow-stale         : %-50s |", "Yes")
	} else {
		log.Infof("| nomad-allow-stale         : %-50s |", "No")
	}
	log.Infof("| nomad-color               : %-50s |", cfg.NomadColor)

	// Consul
	log.Infof("| consul-enable             : %-50t |", cfg.ConsulEnable)
	if cfg.ConsulReadOnly {
		log.Infof("| consul-read-only          : %-50s |", "Yes")
	} else {
		log.Infof("| consul-read-only          : %-50s |", "No (Hashi-UI can change Consul state)")
	}
	log.Infof("| consul-address            : %-50s |", cfg.ConsulAddress)
	log.Infof("| consul-acl-token          : %-50s |", cfg.ConsulACLToken)
	log.Infof("| consul-color              : %-50s |", cfg.ConsulColor)

	log.Infof("----------------------------------------------------------------------------------")
	log.Infof("")

	if !cfg.NomadEnable && !cfg.ConsulEnable {
		log.Fatal("Please enable at least Consul (--consul-enable) or Nomad (--nomad-enable)")
	}

	myAssetFS := assetFS()
	router := mux.NewRouter()

	var app newrelic.Application

	if cfg.NewRelicEnable {
		config := newrelic.NewConfig(cfg.NewRelicAppName, cfg.NewRelicLicenseKey)
		var err error
		app, err = newrelic.NewApplication(config)
		if err != nil {
			log.Fatalf("Could not create NewRelic application: %v", err)
		}
	}

	// create clients
	var nomadClient *nomad.Client
	var consulClient *consul.Client

	if cfg.NomadEnable {
		var err error

		log.Info("Connecting to Nomad ...")
		nomadClient, err = nomad_helper.NewRegionClient(cfg, "")
		if err != nil {
			log.Fatalf("Unable to create Nomad client: %s", err)
		}
		if _, err := nomadClient.Status().Leader(); err != nil {
			log.Fatalf("Unable to communicate with Nomad: %s", err)
		}
		log.Info("done!")
	}

	if cfg.ConsulEnable {
		var err error

		log.Info("Connecting to Consul ...")
		consulClient, err = consul_helper.NewDatacenterClient(cfg, "")
		if err != nil {
			log.Fatalf("Unable to create Consul client: %s", err)
		}
		if _, err := consulClient.Status().Leader(); err != nil {
			log.Fatalf("Unable to communicate with Consul: %s", err)
		}
		log.Info("done!")
	}

	// setup http handlers

	if cfg.NomadEnable {
		router.HandleFunc(newrelic.WrapHandleFunc(app, "/", func(w http.ResponseWriter, r *http.Request) {
			log.Infof("Redirecting / to /nomad")
			w.Write([]byte("<script>document.location.href='" + cfg.ProxyPath + "nomad'</script>"))
			return
		}))

		router.HandleFunc(newrelic.WrapHandleFunc(app, "/ws/nomad", NomadHandler(cfg, nomadClient, consulClient)))
		router.HandleFunc(newrelic.WrapHandleFunc(app, "/ws/nomad/{region}", NomadHandler(cfg, nomadClient, consulClient)))
		router.HandleFunc(newrelic.WrapHandleFunc(app, "/api/nomad/{region}", NomadAPIHandler(cfg, nomadClient, consulClient)))
		router.HandleFunc(newrelic.WrapHandleFunc(app, "/nomad/{region}/download/{path:.*}", NomadDownloadFile(cfg)))
	}

	if cfg.ConsulEnable {
		if !cfg.NomadEnable {
			router.HandleFunc(newrelic.WrapHandleFunc(app, "/", func(w http.ResponseWriter, r *http.Request) {
				log.Infof("Redirecting / to /consul")
				http.Redirect(w, r, cfg.ProxyPath+"consul", 302)
			}))
		}

		router.HandleFunc(newrelic.WrapHandleFunc(app, "/ws/consul", ConsulHandler(cfg, nomadClient, consulClient)))
		router.HandleFunc(newrelic.WrapHandleFunc(app, "/ws/consul/{region}", ConsulHandler(cfg, nomadClient, consulClient)))
	}

	router.HandleFunc(newrelic.WrapHandleFunc(app, "/_status", StatusHandler(cfg, nomadClient, consulClient)))

	router.PathPrefix("/").HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		responseFile := "/index.html"

		if idx := strings.Index(r.URL.Path, "static/"); idx != -1 {
			responseFile = r.URL.Path[idx:]
		}

		if idx := strings.Index(r.URL.Path, "favicon.png"); idx != -1 {
			responseFile = "/favicon.png"
		}

		if idx := strings.Index(r.URL.Path, "config.js"); idx != -1 {
			response := make([]string, 0)
			response = append(response, "window.HASHI_DEV=false")
			response = append(response, fmt.Sprintf("window.GIT_HASH='%s'", GitCommit))

			response = append(response, fmt.Sprintf("window.CONSUL_ENABLED=%s", strconv.FormatBool(cfg.ConsulEnable)))
			response = append(response, fmt.Sprintf("window.CONSUL_READ_ONLY=%s", strconv.FormatBool(cfg.ConsulReadOnly)))

			response = append(response, fmt.Sprintf("window.NOMAD_ENABLED=%s", strconv.FormatBool(cfg.NomadEnable)))
			response = append(response, fmt.Sprintf("window.NOMAD_READ_ONLY=%s", strconv.FormatBool(cfg.NomadReadOnly)))

			response = append(response, fmt.Sprintf("window.NOMAD_COLOR='%s'", cfg.NomadColor))
			response = append(response, fmt.Sprintf("window.CONSUL_COLOR='%s'", cfg.ConsulColor))
			response = append(response, fmt.Sprintf("window.SITE_TITLE='%s'", cfg.SiteTitle))

			enabledServices := make([]string, 0)
			if cfg.ConsulEnable {
				enabledServices = append(enabledServices, "'consul'")
			}
			if cfg.NomadEnable {
				enabledServices = append(enabledServices, "'nomad'")
			}

			response = append(response, fmt.Sprintf("window.ENABLED_SERVICES=[%s]", strings.Join(enabledServices, ",")))
			response = append(response, fmt.Sprintf("window.HASHI_PATH_PREFIX='%s';", cfg.ProxyPath))
			response = append(response, "window.HASHI_ENDPOINT_PORT=window.HASHI_ENDPOINT_PORT || document.location.port")

			w.Header().Set("Content-Type", "application/javascript")

			// never ever cache config.js
			w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
			w.Header().Set("Pragma", "no-cache")
			w.Header().Set("Expires", "0")

			w.Write([]byte(strings.Join(response, "\n")))
			return
		}

		if bs, err := myAssetFS.Open(responseFile); err != nil {
			log.Errorf("%s: %s", responseFile, err)
		} else {
			stat, err := bs.Stat()
			if err != nil {
				w.WriteHeader(http.StatusNotFound)
				log.Errorf("Failed to stat %s: %s", responseFile, err)
			} else {
				http.ServeContent(w, r, responseFile[1:], stat.ModTime(), bs)
			}
		}
	})

	log.Infof("Listening ...")
	var err error
	if cfg.HttpsEnable {
		if cfg.ServerCert == "" || cfg.ServerKey == "" {
			log.Fatal("Using https protocol but server certificate or key were not specified.")
		}
		err = http.ListenAndServeTLS(cfg.ListenAddress, cfg.ServerCert, cfg.ServerKey, router)
	} else {
		err = http.ListenAndServe(cfg.ListenAddress, router)
	}
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>// Backend server - main part of LogVoyage service.
// It accepts connections from "Client", parses string and pushes it to ElasticSearch index
package main

import (
	"encoding/json"
	"errors"
	"flag"
	"log"
	"runtime"
	"strings"
	"time"

	"github.com/firstrow/logvoyage/common"
	"github.com/firstrow/tcp_server"
)

var (
	defaultHost = ""
	defaultPort = "27077"
)

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	log.Print("Initializing server")

	host := flag.String("host", defaultHost, "Host to open server. Set to `localhost` to accept only local connections.")
	port := flag.String("port", defaultPort, "Port to accept new connections. Default value: "+defaultPort)
	flag.Parse()

	server := tcp_server.New(*host + ":" + *port)
	server.OnNewClient(func(c *tcp_server.Client) {
		log.Print("New client")
	})

	// Receives new message and send it to Elastic server
	server.OnNewMessage(func(c *tcp_server.Client, message string) {
		// Send data to elastic
		record := &common.LogRecord{
			Message:  message,
			Datetime: time.Now().UTC(),
		}

		indexName, err := getIndexName(record)
		if err == nil {
			record.Message = common.RemoveApiKey(record.Message)
			record.Message = strings.TrimSpace(record.Message)

			toElastic(indexName, record)
		}
	})

	server.OnClientConnectionClosed(func(c *tcp_server.Client, err error) {
		log.Print("Client disconnected")
	})
	server.Listen()
}

func getIndexName(record *common.LogRecord) (string, error) {
	key, err := common.ExtractApiKey(record.Message)
	if err != nil {
		log.Println(err.Error())
		return "", err
	}

	user := common.FindUserByApiKey(key)
	if user == nil {
		log.Println("User not found")
		return "", errors.New("Error. User not found")
	}

	return user.GetIndexName(), nil
}

func toElastic(indexName string, record *common.LogRecord) {
	j, err := json.Marshal(record)
	if err != nil {
		log.Print("Error encoding message to JSON")
	} else {
		common.SendToElastic(indexName+"/logs", "POST", j)
	}
}
<commit_msg>Refactred essage saving<commit_after>// Backend server - main part of LogVoyage service.
// It accepts connections from "Client", parses string and pushes it to ElasticSearch index
package main

import (
	"encoding/json"
	"errors"
	"flag"
	"log"
	"runtime"
	"strings"
	"time"

	"github.com/firstrow/logvoyage/common"
	"github.com/firstrow/tcp_server"
)

var (
	defaultHost = ""
	defaultPort = "27077"
)

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	log.Print("Initializing server")

	host := flag.String("host", defaultHost, "Host to open server. Set to `localhost` to accept only local connections.")
	port := flag.String("port", defaultPort, "Port to accept new connections. Default value: "+defaultPort)
	flag.Parse()

	server := tcp_server.New(*host + ":" + *port)
	server.OnNewClient(func(c *tcp_server.Client) {
		log.Print("New client")
	})

	// Receives new message and send it to Elastic server
	server.OnNewMessage(func(c *tcp_server.Client, message string) {
		indexName, err := getIndexName(message)
		if err != nil {
			// TODO: Log error
		} else {
			message = common.RemoveApiKey(message)
			message = strings.TrimSpace(message)

			// Send data to elastic
			record := &common.LogRecord{
				Message:  message,
				Datetime: time.Now().UTC(),
			}
			toElastic(indexName, record)
		}
	})

	server.OnClientConnectionClosed(func(c *tcp_server.Client, err error) {
		log.Print("Client disconnected")
	})
	server.Listen()
}

// Get users index name by apiKey
func getIndexName(message string) (string, error) {
	key, err := common.ExtractApiKey(message)
	if err != nil {
		log.Println(err.Error())
		return "", err
	}

	user := common.FindUserByApiKey(key)
	if user == nil {
		log.Println("User not found")
		return "", errors.New("Error. User not found")
	}

	return user.GetIndexName(), nil
}

func toElastic(indexName string, record *common.LogRecord) {
	j, err := json.Marshal(record)
	if err != nil {
		log.Print("Error encoding message to JSON")
	} else {
		common.SendToElastic(indexName+"/logs", "POST", j)
	}
}
<|endoftext|>


<commit_before>package backends

import (
  "time"
  "fmt"
  "strings"
  "github.com/docker/libswarm/beam"

  "launchpad.net/goamz/aws"
  "launchpad.net/goamz/ec2"
)

type ec2Config struct {
  securityGroup string
  instanceType  string
  keypair        string
  zone           string
  ami            string
  tag            string
  region         aws.Region
}

type ec2Client struct {
  config   *ec2Config
  ec2Conn  *ec2.EC2
  Server   *beam.Server
  instance *ec2.Instance
  dockerInstance *beam.Object
}

func (c *ec2Client) get(ctx *beam.Message) error {
  output, err := c.dockerInstance.Get()
  if (err != nil) {
    return err
  }
  ctx.Ret.Send(&beam.Message{Verb: beam.Set, Args: []string{output}})
  return nil
}

func (c *ec2Client) start(ctx *beam.Message) error {
  if instance, err := c.findInstance(); err != nil {
    return err
  } else if instance != nil {
    fmt.Printf("Found existing instance: %s\n", instance.InstanceId)
    c.instance = instance
  } else {
    if err := c.startInstance(); err != nil {
      return err
    }

    if err := c.tagtInstance(); err != nil {
      return err
    }
  }

  c.initDockerClientInstance(c.instance)
  ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: c.Server})
  return nil
}

func (c *ec2Client) log(ctx *beam.Message) error {
  ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: c.Server})
  return nil
}

func (c *ec2Client) spawn(ctx *beam.Message) error {
  out, err := c.dockerInstance.Spawn(ctx.Args...)
  if err != nil {
    return err
  }
  ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: out})
  return nil
}

func (c *ec2Client) ls(ctx *beam.Message) error {
  output, err := c.dockerInstance.Ls()
  if (err != nil) {
    return err
  }
  ctx.Ret.Send(&beam.Message{Verb: beam.Set, Args: output})
  return nil
}

func (c *ec2Client) error(ctx *beam.Message) error {
  ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: c.Server})
  return nil
}

func (c *ec2Client) stop(ctx *beam.Message) error {
  ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: c.Server})
  return nil
}

func (c *ec2Client) attach(ctx *beam.Message) error {
	if ctx.Args[0] == "" {
    ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: c.Server})
		<-make(chan struct{})
	} else {
    _, out, err := c.dockerInstance.Attach(ctx.Args[0])
    if err != nil {
      return err
    }
    ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: out})
  }

  return nil
}

func defaultConfigValues() (config *ec2Config) {
  config               = new(ec2Config)
  config.region        = aws.USEast
  config.instanceType  = "t1.micro"
  return config
}

func newConfig(args []string) (config *ec2Config, err error) {
  // TODO (aaron): fail fast on incorrect number of args

  var optValPair []string
  var opt, val string

  config = defaultConfigValues()

  for _, value := range args {
    optValPair = strings.Split(value, "=")
    opt, val   = optValPair[0], optValPair[1]

    switch opt {
      case "--region":
        config.region = convertToRegion(val)
      case "--zone":
        config.zone = val
      case "--tag":
        config.tag = val
      case "--ami":
        config.ami = val
      case "--keypair":
        config.keypair = val
      case "--security_group":
        config.securityGroup = val
      case "--instance_type":
        config.instanceType = val
      default:
        fmt.Printf("Unrecognizable option: %s value: %s", opt, val)
    }
  }
  return config, nil
}

func convertToRegion(input string) aws.Region {
  // TODO (aaron): fill in more regions
  switch input {
    case "us-east-1":
      return  aws.USEast
    case "us-west-1":
      return  aws.USWest
    case "us-west-2":
      return  aws.USWest2
    case "eu-west-1":
      return  aws.EUWest
    default:
      fmt.Println("Unrecognizable region, default to: us-east-1")
      return aws.USEast
  }
}

func awsInit(config *ec2Config)  (ec2Conn *ec2.EC2, err error) {
  auth, err := aws.EnvAuth()

  if err != nil {
    return nil, err
  }

  return ec2.New(auth, config.region), nil
}

func (c *ec2Client) findInstance() (instance *ec2.Instance, err error) {
  filter := ec2.NewFilter()
  filter.Add("tag:Name", c.config.tag)
  resp, err := c.ec2Conn.Instances([]string{}, filter)

  if err != nil {
    return nil, err
  } else {
    if resp.Reservations == nil {
      return nil, nil
    }

    instance := resp.Reservations[0].Instances[0]

    if (instance.State.Name == "running" || instance.State.Name == "pending") {
      return &instance, nil
    }

    return nil, nil
  }
}

func (c *ec2Client) tagtInstance() error {
  ec2Tags := []ec2.Tag{ec2.Tag{"Name", c.config.tag}}
  if _, err := c.ec2Conn.CreateTags([]string{c.instance.InstanceId}, ec2Tags); err != nil {
    return err
  }
  return nil
}

func (c *ec2Client) startInstance() error {
  // TODO (aaron): make sure to wait for cloud-init to finish before
  // executing docker commands
  options := ec2.RunInstances{
    ImageId:      c.config.ami,
    InstanceType: c.config.instanceType,
    KeyName:      c.config.keypair,
    UserData:     []byte("#include https://get.docker.io"),
  }

  resp, err := c.ec2Conn.RunInstances(&options)
  if err != nil {
    return err
  }

  // TODO (aaron): this really could be multiple instances, not just 1
  i := resp.Instances[0]

  for i.State.Name != "running" {
    time.Sleep(3 * time.Second)
    fmt.Printf("Waiting for instance to come up.  Current State: %s\n",
               i.State.Name)

    resp, err := c.ec2Conn.Instances([]string{i.InstanceId}, ec2.NewFilter())

    if err != nil {
      return err
    }

    i = resp.Reservations[0].Instances[0]
  }

  c.instance = &i

  fmt.Printf("Instance up and running - id: %s\n", i.InstanceId)
  return nil
}

func (c *ec2Client) initDockerClientInstance(instance *ec2.Instance) error {
  dockerClient := DockerClientWithConfig(&DockerClientConfig{
		Scheme:          "http",
		URLHost:         instance.IPAddress,
	})

  dockerBackend := beam.Obj(dockerClient)
  url := fmt.Sprintf("tcp://%s:4243", instance.IPAddress)
  dockerInstance, err := dockerBackend.Spawn(url)
  c.dockerInstance = dockerInstance

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

func Ec2() beam.Sender {
  backend := beam.NewServer()
  backend.OnSpawn(beam.Handler(func(ctx *beam.Message) error {
    var config, err = newConfig(ctx.Args)

    if (err != nil) {
      return err
    }

    ec2Conn, err := awsInit(config)
    if (err != nil) {
      return err
    }

    client := &ec2Client{config, ec2Conn, beam.NewServer(), nil, nil}
    client.Server.OnSpawn(beam.Handler(client.spawn))
    client.Server.OnStart(beam.Handler(client.start))
    client.Server.OnStop(beam.Handler(client.stop))
    client.Server.OnAttach(beam.Handler(client.attach))
    client.Server.OnLog(beam.Handler(client.log))
    client.Server.OnError(beam.Handler(client.error))
    client.Server.OnLs(beam.Handler(client.ls))
    client.Server.OnGet(beam.Handler(client.get))

    _, err = ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: client.Server})

    return err
  }))
  return backend
}
<commit_msg>allow more aws options to be configurable<commit_after>package backends

import (
  "time"
  "fmt"
  "strings"
  "github.com/docker/libswarm/beam"

  "launchpad.net/goamz/aws"
  "launchpad.net/goamz/ec2"
)

type ec2Config struct {
  securityGroup string
  instanceType  string
  keypair        string
  zone           string
  ami            string
  tag            string
  region         aws.Region
}

type ec2Client struct {
  config   *ec2Config
  ec2Conn  *ec2.EC2
  Server   *beam.Server
  instance *ec2.Instance
  dockerInstance *beam.Object
}

func (c *ec2Client) get(ctx *beam.Message) error {
  output, err := c.dockerInstance.Get()
  if (err != nil) {
    return err
  }
  ctx.Ret.Send(&beam.Message{Verb: beam.Set, Args: []string{output}})
  return nil
}

func (c *ec2Client) start(ctx *beam.Message) error {
  if instance, err := c.findInstance(); err != nil {
    return err
  } else if instance != nil {
    fmt.Printf("Found existing instance: %s\n", instance.InstanceId)
    c.instance = instance
  } else {
    if err := c.startInstance(); err != nil {
      return err
    }

    if err := c.tagtInstance(); err != nil {
      return err
    }
  }

  c.initDockerClientInstance(c.instance)
  ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: c.Server})
  return nil
}

func (c *ec2Client) log(ctx *beam.Message) error {
  ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: c.Server})
  return nil
}

func (c *ec2Client) spawn(ctx *beam.Message) error {
  out, err := c.dockerInstance.Spawn(ctx.Args...)
  if err != nil {
    return err
  }
  ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: out})
  return nil
}

func (c *ec2Client) ls(ctx *beam.Message) error {
  output, err := c.dockerInstance.Ls()
  if (err != nil) {
    return err
  }
  ctx.Ret.Send(&beam.Message{Verb: beam.Set, Args: output})
  return nil
}

func (c *ec2Client) error(ctx *beam.Message) error {
  ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: c.Server})
  return nil
}

func (c *ec2Client) stop(ctx *beam.Message) error {
  ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: c.Server})
  return nil
}

func (c *ec2Client) attach(ctx *beam.Message) error {
	if ctx.Args[0] == "" {
    ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: c.Server})
		<-make(chan struct{})
	} else {
    _, out, err := c.dockerInstance.Attach(ctx.Args[0])
    if err != nil {
      return err
    }
    ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: out})
  }

  return nil
}

func defaultConfigValues() (config *ec2Config) {
  config               = new(ec2Config)
  config.region        = aws.USEast
  config.instanceType  = "t1.micro"
  config.zone          = "us-east-1a"
  return config
}

func newConfig(args []string) (config *ec2Config, err error) {
  // TODO (aaron): fail fast on incorrect number of args

  var optValPair []string
  var opt, val string

  config = defaultConfigValues()

  for _, value := range args {
    optValPair = strings.Split(value, "=")
    opt, val   = optValPair[0], optValPair[1]

    switch opt {
      case "--region":
        config.region = convertToRegion(val)
      case "--zone":
        config.zone = val
      case "--tag":
        config.tag = val
      case "--ami":
        config.ami = val
      case "--keypair":
        config.keypair = val
      case "--security_group":
        config.securityGroup = val
      case "--instance_type":
        config.instanceType = val
      default:
        fmt.Printf("Unrecognizable option: %s value: %s", opt, val)
    }
  }
  return config, nil
}

func convertToRegion(input string) aws.Region {
  // TODO (aaron): fill in more regions
  switch input {
    case "us-east-1":
      return  aws.USEast
    case "us-west-1":
      return  aws.USWest
    case "us-west-2":
      return  aws.USWest2
    case "eu-west-1":
      return  aws.EUWest
    default:
      fmt.Println("Unrecognizable region, default to: us-east-1")
      return aws.USEast
  }
}

func awsInit(config *ec2Config)  (ec2Conn *ec2.EC2, err error) {
  auth, err := aws.EnvAuth()

  if err != nil {
    return nil, err
  }

  return ec2.New(auth, config.region), nil
}

func (c *ec2Client) findInstance() (instance *ec2.Instance, err error) {
  filter := ec2.NewFilter()
  filter.Add("tag:Name", c.config.tag)
  resp, err := c.ec2Conn.Instances([]string{}, filter)

  if err != nil {
    return nil, err
  } else {
    if resp.Reservations == nil {
      return nil, nil
    }

    instance := resp.Reservations[0].Instances[0]

    if (instance.State.Name == "running" || instance.State.Name == "pending") {
      return &instance, nil
    }

    return nil, nil
  }
}

func (c *ec2Client) tagtInstance() error {
  ec2Tags := []ec2.Tag{ec2.Tag{"Name", c.config.tag}}
  if _, err := c.ec2Conn.CreateTags([]string{c.instance.InstanceId}, ec2Tags); err != nil {
    return err
  }
  return nil
}

func (c *ec2Client) startInstance() error {
  // TODO (aaron): make sure to wait for cloud-init to finish before
  // executing docker commands
  options := ec2.RunInstances{
    ImageId:        c.config.ami,
    InstanceType:   c.config.instanceType,
    KeyName:        c.config.keypair,
    AvailZone:      c.config.zone,
    // TODO: allow more than one sg in the future
    SecurityGroups: []ec2.SecurityGroup{ec2.SecurityGroup{Name: c.config.securityGroup}},
    UserData:       []byte(userdata),
  }

  resp, err := c.ec2Conn.RunInstances(&options)
  if err != nil {
    return err
  }

  // TODO (aaron): this really could be multiple instances, not just 1
  i := resp.Instances[0]

  for i.State.Name != "running" {
    time.Sleep(3 * time.Second)
    fmt.Printf("Waiting for instance to come up.  Current State: %s\n",
               i.State.Name)

    resp, err := c.ec2Conn.Instances([]string{i.InstanceId}, ec2.NewFilter())

    if err != nil {
      return err
    }

    i = resp.Reservations[0].Instances[0]
  }

  c.instance = &i

  fmt.Printf("Instance up and running - id: %s\n", i.InstanceId)
  return nil
}

func (c *ec2Client) initDockerClientInstance(instance *ec2.Instance) error {
  dockerClient := DockerClientWithConfig(&DockerClientConfig{
		Scheme:          "http",
		URLHost:         instance.IPAddress,
	})

  dockerBackend := beam.Obj(dockerClient)
  url := fmt.Sprintf("tcp://%s:4243", instance.IPAddress)
  dockerInstance, err := dockerBackend.Spawn(url)
  c.dockerInstance = dockerInstance

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

func Ec2() beam.Sender {
  backend := beam.NewServer()
  backend.OnSpawn(beam.Handler(func(ctx *beam.Message) error {
    var config, err = newConfig(ctx.Args)

    if (err != nil) {
      return err
    }

    ec2Conn, err := awsInit(config)
    if (err != nil) {
      return err
    }

    client := &ec2Client{config, ec2Conn, beam.NewServer(), nil, nil}
    client.Server.OnSpawn(beam.Handler(client.spawn))
    client.Server.OnStart(beam.Handler(client.start))
    client.Server.OnStop(beam.Handler(client.stop))
    client.Server.OnAttach(beam.Handler(client.attach))
    client.Server.OnLog(beam.Handler(client.log))
    client.Server.OnError(beam.Handler(client.error))
    client.Server.OnLs(beam.Handler(client.ls))
    client.Server.OnGet(beam.Handler(client.get))

    _, err = ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: client.Server})

    return err
  }))
  return backend
}

// TODO (aaron): load this externally
const userdata = `
#!/bin/bash
yum install -y docker
cat << EOF > /etc/sysconfig/docker
other_args="-H tcp://127.0.0.1:4243"
EOF
service docker start
`
<|endoftext|>


<commit_before>package backends

import (
  "time"
  "fmt"
  "strings"
  "github.com/docker/libswarm/beam"
)

type ec2Config struct {
  access_key     string
  secret_key     string
  security_group string
  instance_type  string
  keypair        string
  region         string
  zone           string
  ami            string
  tag            string
}

type ec2Client struct {
  config *ec2Config
  Server *beam.Server
}

func (c *ec2Client) get(ctx *beam.Message) error {
  fmt.Println("*** ec2 Get ***")
  body := `{}`
  ctx.Ret.Send(&beam.Message{Verb: beam.Set, Args: []string{body}})
  return nil
}

func (c *ec2Client) start(ctx *beam.Message) error {
  fmt.Println("*** ec2 OnStart ***")
  ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: c.Server})
  return nil
}

func (c *ec2Client) log(ctx *beam.Message) error {
  fmt.Println("*** ec2 OnLog ***")
  ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: c.Server})
  return nil
}

func (c *ec2Client) spawn(ctx *beam.Message) error {
  fmt.Println("*** ec2 OnSpawn ***")
  ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: c.Server})
  return nil
}

func (c *ec2Client) ls(ctx *beam.Message) error {
  fmt.Println("*** ec2 OnLs ***")
  ctx.Ret.Send(&beam.Message{Verb: beam.Set, Args: []string{}})
  return nil
}

func (c *ec2Client) error(ctx *beam.Message) error {
  fmt.Println("*** ec2 OnError ***")
  ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: c.Server})
  return nil
}

func (c *ec2Client) stop(ctx *beam.Message) error {
  fmt.Println("*** ec2 OnStop ***")
  ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: c.Server})
  return nil
}

func (c *ec2Client) attach(ctx *beam.Message) error {
  fmt.Println("*** ec2 OnAttach ***")
  ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: c.Server})

  for {
     time.Sleep(1 * time.Second)
     (&beam.Object{ctx.Ret}).Log("ec2: heartbeat")
  }
  return nil
}

func newConfig(args []string) (config *ec2Config, err error) {
  // TODO (aaron): fail fast on incorrect number of args

  var optValPair []string
  var opt, val string

  config = new(ec2Config)
  for _, value := range args {
    optValPair = strings.Split(value, "=")
    opt = optValPair[0]
    val = optValPair[1]

    switch opt {
      case "--region":
        config.region = val
      case "--zone":
        config.zone = val
      case "--tag":
        config.tag = val
      case "--ami":
        config.ami = val
      case "--keypair":
        config.keypair = val
      case "--security_group":
        config.security_group = val
      case "--instance_type":
        config.instance_type = val
      default:
        fmt.Printf("Unrecognizable option: %s value: %s", opt, val)
    }
  }
  return config, nil
}

func Ec2() beam.Sender {
  backend := beam.NewServer()
  backend.OnSpawn(beam.Handler(func(ctx *beam.Message) error {
    fmt.Println("*** init ***")

    var config, err = newConfig(ctx.Args)

    if (err != nil) {
      return err
    }

    client := &ec2Client{config, beam.NewServer()}
    client.Server.OnSpawn(beam.Handler(client.spawn))
    client.Server.OnStart(beam.Handler(client.start))
    client.Server.OnStop(beam.Handler(client.stop))
    client.Server.OnAttach(beam.Handler(client.attach))
    client.Server.OnLog(beam.Handler(client.log))
    client.Server.OnError(beam.Handler(client.error))
    client.Server.OnLs(beam.Handler(client.ls))
    client.Server.OnGet(beam.Handler(client.get))
    _, err = ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: client.Server})
    return err
  }))
  return backend
}
<commit_msg>starting up ec2 instance<commit_after>package backends

import (
  "time"
  "fmt"
  "strings"
  "github.com/docker/libswarm/beam"

  "launchpad.net/goamz/aws"
  "launchpad.net/goamz/ec2"
)

type ec2Config struct {
  securityGroup string
  instanceType  string
  keypair        string
  zone           string
  ami            string
  tag            string
  region         aws.Region
}

type ec2Client struct {
  config  *ec2Config
  ec2Conn *ec2.EC2
  Server  *beam.Server
}

func (c *ec2Client) get(ctx *beam.Message) error {
  fmt.Println("*** ec2 Get ***")
  body := `{}`
  ctx.Ret.Send(&beam.Message{Verb: beam.Set, Args: []string{body}})
  return nil
}

func (c *ec2Client) start(ctx *beam.Message) error {
  fmt.Println("*** ec2 OnStart ***")
  err := c.startInstance()

  if (err != nil) {
    return err
  }

  ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: c.Server})
  return nil
}

func (c *ec2Client) log(ctx *beam.Message) error {
  fmt.Println("*** ec2 OnLog ***")
  ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: c.Server})
  return nil
}

func (c *ec2Client) spawn(ctx *beam.Message) error {
  fmt.Println("*** ec2 OnSpawn ***")
  ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: c.Server})
  return nil
}

func (c *ec2Client) ls(ctx *beam.Message) error {
  fmt.Println("*** ec2 OnLs ***")
  ctx.Ret.Send(&beam.Message{Verb: beam.Set, Args: []string{}})
  return nil
}

func (c *ec2Client) error(ctx *beam.Message) error {
  fmt.Println("*** ec2 OnError ***")
  ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: c.Server})
  return nil
}

func (c *ec2Client) stop(ctx *beam.Message) error {
  fmt.Println("*** ec2 OnStop ***")
  ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: c.Server})
  return nil
}

func (c *ec2Client) attach(ctx *beam.Message) error {
  fmt.Println("*** ec2 OnAttach ***")
  ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: c.Server})

  for {
     time.Sleep(1 * time.Second)
     (&beam.Object{ctx.Ret}).Log("ec2: heartbeat")
  }
  return nil
}

func defaultConfigValues() (config *ec2Config) {
  config               = new(ec2Config)
  config.region        = aws.USEast
  config.instanceType  = "t1.micro"
  return config
}

func newConfig(args []string) (config *ec2Config, err error) {
  // TODO (aaron): fail fast on incorrect number of args

  var optValPair []string
  var opt, val string

  config = defaultConfigValues()

  for _, value := range args {
    optValPair = strings.Split(value, "=")
    opt, val   = optValPair[0], optValPair[1]

    switch opt {
      case "--region":
        config.region = convertToRegion(val)
      case "--zone":
        config.zone = val
      case "--tag":
        config.tag = val
      case "--ami":
        config.ami = val
      case "--keypair":
        config.keypair = val
      case "--security_group":
        config.securityGroup = val
      case "--instance_type":
        config.instanceType = val
      default:
        fmt.Printf("Unrecognizable option: %s value: %s", opt, val)
    }
  }
  return config, nil
}

func convertToRegion(input string) aws.Region {
  // TODO (aaron): fill in more regions
  switch input {
    case "us-east-1":
      return  aws.USEast
    case "us-west-1":
      return  aws.USWest
    case "us-west-2":
      return  aws.USWest2
    case "eu-west-1":
      return  aws.EUWest
    default:
      fmt.Println("Unrecognizable region, default to: us-east-1")
      return aws.USEast
  }
}

func awsInit(config *ec2Config)  (ec2Conn *ec2.EC2, err error) {
  auth, err := aws.EnvAuth()

  if err != nil {
    return nil, err
  }

  return ec2.New(auth, config.region), nil
}

func (c *ec2Client) tagtInstance() error {
  return nil
}

func (c *ec2Client) startInstance() error {
  options := ec2.RunInstances{
    ImageId:      c.config.ami,
    InstanceType: c.config.instanceType,
    KeyName:      c.config.keypair,
  }

  resp, err := c.ec2Conn.RunInstances(&options)
  if err != nil {
    return err
  }

  i := resp.Instances[0]
  for i.State.Name != "running" {
    time.Sleep(3 * time.Second)
    fmt.Printf("Waiting for instance to come up.  Current State: %s\n",
               i.State.Name)

    resp, err := c.ec2Conn.Instances([]string{i.InstanceId}, ec2.NewFilter())

    if err != nil {
      return err
    }

    i = resp.Reservations[0].Instances[0]
  }

  fmt.Printf("Instance up and running - id: %s\n", i.InstanceId)
  return nil
}

func Ec2() beam.Sender {
  backend := beam.NewServer()
  backend.OnSpawn(beam.Handler(func(ctx *beam.Message) error {
    fmt.Println("*** init ***")

    var config, err = newConfig(ctx.Args)

    if (err != nil) {
      return err
    }

    ec2Conn, err := awsInit(config)
    if (err != nil) {
      return err
    }

    client := &ec2Client{config, ec2Conn, beam.NewServer()}
    client.Server.OnSpawn(beam.Handler(client.spawn))
    client.Server.OnStart(beam.Handler(client.start))
    client.Server.OnStop(beam.Handler(client.stop))
    client.Server.OnAttach(beam.Handler(client.attach))
    client.Server.OnLog(beam.Handler(client.log))
    client.Server.OnError(beam.Handler(client.error))
    client.Server.OnLs(beam.Handler(client.ls))
    client.Server.OnGet(beam.Handler(client.get))

    _, err = ctx.Ret.Send(&beam.Message{Verb: beam.Ack, Ret: client.Server})

    return err
  }))
  return backend
}
<|endoftext|>


<commit_before>package handler

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/rs/zerolog/hlog"

	"github.com/gilcrest/go-api-basic/domain/errs"
	"github.com/pkg/errors"

	"github.com/gilcrest/go-api-basic/domain/user"

	qt "github.com/frankban/quicktest"

	"github.com/gilcrest/go-api-basic/domain/auth"
)

func TestJSONContentTypeHandler(t *testing.T) {
	req, err := http.NewRequest("GET", "/ping", nil)
	if err != nil {
		t.Fatalf("http.NewRequest() error = %v", err)
	}

	testJSONContentTypeHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		contentType := w.Header().Get("Content-Type")
		if contentType != "application/json" {
			t.Fatalf("Content-Type %s is invalid", contentType)
		}
	})

	rr := httptest.NewRecorder()

	handlers := JSONContentTypeHandler(testJSONContentTypeHandler)
	handlers.ServeHTTP(rr, req)
}

func TestAccessTokenHandler(t *testing.T) {
	t.Run("typical", func(t *testing.T) {
		c := qt.New(t)

		req, err := http.NewRequest("GET", "/ping", nil)
		if err != nil {
			t.Fatalf("http.NewRequest() error = %v", err)
		}
		req.Header.Add("Authorization", "Bearer abcdef123")

		testAccessTokenHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			token, err := auth.FromRequest(r)
			if err != nil {
				t.Fatalf("auth.FromRequest() error = %v", err)
			}
			wantToken := auth.AccessToken{
				Token:     "abcdef123",
				TokenType: "Bearer",
			}
			t.Log(token)
			c.Assert(token, qt.Equals, wantToken)
		})

		rr := httptest.NewRecorder()

		handlers := AccessTokenHandler(testAccessTokenHandler)
		handlers.ServeHTTP(rr, req)

		// If there is any issues with the Access Token, the body
		// should be empty and the status code should be 401
		c.Assert(rr.Code, qt.Equals, http.StatusOK)
	})

	t.Run("no token", func(t *testing.T) {
		c := qt.New(t)

		req, err := http.NewRequest("GET", "/ping", nil)
		if err != nil {
			t.Fatalf("http.NewRequest() error = %v", err)
		}

		testAccessTokenHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			c.Fatal("handler should not make it here")
		})

		rr := httptest.NewRecorder()

		handlers := AccessTokenHandler(testAccessTokenHandler)
		handlers.ServeHTTP(rr, req)

		body, err := ioutil.ReadAll(rr.Body)
		if err != nil {
			t.Fatalf("ioutil.ReadAll() error = %v", err)
		}

		// If there is any issues with the Access Token, the status
		// code should be 401, and the body should be empty
		c.Assert(rr.Code, qt.Equals, http.StatusUnauthorized)
		c.Assert(string(body), qt.Equals, "")
	})
}

func TestNewStandardResponse(t *testing.T) {
	t.Run("no request id", func(t *testing.T) {
		c := qt.New(t)

		req, err := http.NewRequest("GET", "/ping", nil)
		if err != nil {
			t.Fatalf("http.NewRequest() error = %v", err)
		}
		u := user.User{Email: "otto.maddox711@gmail.com"}
		_, err = NewStandardResponse(req, u)
		wantErr := errs.E(errors.New("request ID not properly set to request context"))

		c.Assert(errs.Match(err, wantErr), qt.Equals, true)

	})

	t.Run("typical", func(t *testing.T) {
		c := qt.New(t)

		req, err := http.NewRequest("GET", "/ping", nil)
		if err != nil {
			t.Fatalf("http.NewRequest() error = %v", err)
		}
		rr := httptest.NewRecorder()

		var requestID string
		testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			rID, ok := hlog.IDFromRequest(r)
			if !ok {
				t.Fatal("Request ID not set to request context")
			}
			requestID = rID.String()
			req = r
		})

		h := hlog.RequestIDHandler("request_id", "Request-Id")(testHandler)
		h.ServeHTTP(rr, req)

		u := user.User{Email: "otto.maddox711@gmail.com"}
		sr, err := NewStandardResponse(req, u)
		if err != nil {
			t.Fatalf("NewStandardResponse() error = %v", err)
		}

		wantStandardResponse := &StandardResponse{Path: "/ping", RequestID: requestID, Data: u}

		c.Assert(sr, qt.DeepEquals, wantStandardResponse)

	})
}

func TestDecoderErr(t *testing.T) {
	t.Run("typical", func(t *testing.T) {
		c := qt.New(t)

		type testBody struct {
			Director string `json:"director"`
			Writer   string `json:"writer"`
		}

		requestBody := []byte(`{
				"director": "Alex Cox",
				"writer": "Alex Cox"
			}`)

		r, err := http.NewRequest(http.MethodPost, "/fake", bytes.NewBuffer(requestBody))
		if err != nil {
			t.Fatalf("http.NewRequest() error = %v", err)
		}

		// Decode JSON HTTP request body into a Decoder type
		// and unmarshal that into the testBody struct. DecoderErr
		// wraps errors from Decode when body is nil, json is malformed
		// or any other error
		wantBody := new(testBody)
		err = DecoderErr(json.NewDecoder(r.Body).Decode(&wantBody))
		defer r.Body.Close()
		c.Assert(err, qt.IsNil)
	})

	t.Run("malformed JSON", func(t *testing.T) {
		c := qt.New(t)

		type testBody struct {
			Director string `json:"director"`
			Writer   string `json:"writer"`
		}

		// removed trailing curly bracket
		requestBody := []byte(`{
				"director": "Alex Cox",
				"writer": "Alex Cox"`)

		r, err := http.NewRequest(http.MethodPost, "/fake", bytes.NewBuffer(requestBody))
		if err != nil {
			t.Fatalf("http.NewRequest() error = %v", err)
		}

		// Decode JSON HTTP request body into a Decoder type
		// and unmarshal that into the testBody struct. DecoderErr
		// wraps errors from Decode when body is nil, JSON is malformed
		// or any other error
		wantBody := new(testBody)
		err = DecoderErr(json.NewDecoder(r.Body).Decode(&wantBody))
		defer r.Body.Close()

		wantErr := errs.E(errs.InvalidRequest, errors.New("Malformed JSON"))
		c.Assert(errs.Match(err, wantErr), qt.IsTrue)
	})

	t.Run("empty request body", func(t *testing.T) {
		c := qt.New(t)

		type testBody struct {
			Director string `json:"director"`
			Writer   string `json:"writer"`
		}

		// empty body
		requestBody := []byte("")

		r, err := http.NewRequest(http.MethodPost, "/fake", bytes.NewBuffer(requestBody))
		if err != nil {
			t.Fatalf("http.NewRequest() error = %v", err)
		}

		// Decode JSON HTTP request body into a Decoder type
		// and unmarshal that into the testBody struct. DecoderErr
		// wraps errors from Decode when body is nil, JSON is malformed
		// or any other error
		wantBody := new(testBody)
		err = DecoderErr(json.NewDecoder(r.Body).Decode(&wantBody))
		defer r.Body.Close()

		wantErr := errs.E(errs.InvalidRequest, errors.New("Request Body cannot be empty"))
		c.Assert(errs.Match(err, wantErr), qt.IsTrue)
	})

	t.Run("invalid request body", func(t *testing.T) {
		c := qt.New(t)

		type testBody struct {
			Director string `json:"director"`
			Writer   string `json:"writer"`
		}

		// has unknown field
		requestBody := []byte(`{
				"director": "Alex Cox",
				"writer": "Alex Cox",
                "unknown_field": "I should fail"
			}`)

		r, err := http.NewRequest(http.MethodPost, "/fake", bytes.NewBuffer(requestBody))
		if err != nil {
			t.Fatalf("http.NewRequest() error = %v", err)
		}

		// force an error with DisallowUnknownFields
		wantBody := new(testBody)
		decoder := json.NewDecoder(r.Body)
		decoder.DisallowUnknownFields()
		err = DecoderErr(decoder.Decode(&wantBody))
		defer r.Body.Close()

		// check to make sure I have an error
		c.Assert(err != nil, qt.Equals, true)
	})
}
<commit_msg>remove log statement<commit_after>package handler

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/rs/zerolog/hlog"

	"github.com/gilcrest/go-api-basic/domain/errs"
	"github.com/pkg/errors"

	"github.com/gilcrest/go-api-basic/domain/user"

	qt "github.com/frankban/quicktest"

	"github.com/gilcrest/go-api-basic/domain/auth"
)

func TestJSONContentTypeHandler(t *testing.T) {
	req, err := http.NewRequest("GET", "/ping", nil)
	if err != nil {
		t.Fatalf("http.NewRequest() error = %v", err)
	}

	testJSONContentTypeHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		contentType := w.Header().Get("Content-Type")
		if contentType != "application/json" {
			t.Fatalf("Content-Type %s is invalid", contentType)
		}
	})

	rr := httptest.NewRecorder()

	handlers := JSONContentTypeHandler(testJSONContentTypeHandler)
	handlers.ServeHTTP(rr, req)
}

func TestAccessTokenHandler(t *testing.T) {
	t.Run("typical", func(t *testing.T) {
		c := qt.New(t)

		req, err := http.NewRequest("GET", "/ping", nil)
		if err != nil {
			t.Fatalf("http.NewRequest() error = %v", err)
		}
		req.Header.Add("Authorization", "Bearer abcdef123")

		testAccessTokenHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			token, err := auth.FromRequest(r)
			if err != nil {
				t.Fatalf("auth.FromRequest() error = %v", err)
			}
			wantToken := auth.AccessToken{
				Token:     "abcdef123",
				TokenType: "Bearer",
			}
			c.Assert(token, qt.Equals, wantToken)
		})

		rr := httptest.NewRecorder()

		handlers := AccessTokenHandler(testAccessTokenHandler)
		handlers.ServeHTTP(rr, req)

		// If there is any issues with the Access Token, the body
		// should be empty and the status code should be 401
		c.Assert(rr.Code, qt.Equals, http.StatusOK)
	})

	t.Run("no token", func(t *testing.T) {
		c := qt.New(t)

		req, err := http.NewRequest("GET", "/ping", nil)
		if err != nil {
			t.Fatalf("http.NewRequest() error = %v", err)
		}

		testAccessTokenHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			c.Fatal("handler should not make it here")
		})

		rr := httptest.NewRecorder()

		handlers := AccessTokenHandler(testAccessTokenHandler)
		handlers.ServeHTTP(rr, req)

		body, err := ioutil.ReadAll(rr.Body)
		if err != nil {
			t.Fatalf("ioutil.ReadAll() error = %v", err)
		}

		// If there is any issues with the Access Token, the status
		// code should be 401, and the body should be empty
		c.Assert(rr.Code, qt.Equals, http.StatusUnauthorized)
		c.Assert(string(body), qt.Equals, "")
	})
}

func TestNewStandardResponse(t *testing.T) {
	t.Run("no request id", func(t *testing.T) {
		c := qt.New(t)

		req, err := http.NewRequest("GET", "/ping", nil)
		if err != nil {
			t.Fatalf("http.NewRequest() error = %v", err)
		}
		u := user.User{Email: "otto.maddox711@gmail.com"}
		_, err = NewStandardResponse(req, u)
		wantErr := errs.E(errors.New("request ID not properly set to request context"))

		c.Assert(errs.Match(err, wantErr), qt.Equals, true)

	})

	t.Run("typical", func(t *testing.T) {
		c := qt.New(t)

		req, err := http.NewRequest("GET", "/ping", nil)
		if err != nil {
			t.Fatalf("http.NewRequest() error = %v", err)
		}
		rr := httptest.NewRecorder()

		var requestID string
		testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			rID, ok := hlog.IDFromRequest(r)
			if !ok {
				t.Fatal("Request ID not set to request context")
			}
			requestID = rID.String()
			req = r
		})

		h := hlog.RequestIDHandler("request_id", "Request-Id")(testHandler)
		h.ServeHTTP(rr, req)

		u := user.User{Email: "otto.maddox711@gmail.com"}
		sr, err := NewStandardResponse(req, u)
		if err != nil {
			t.Fatalf("NewStandardResponse() error = %v", err)
		}

		wantStandardResponse := &StandardResponse{Path: "/ping", RequestID: requestID, Data: u}

		c.Assert(sr, qt.DeepEquals, wantStandardResponse)

	})
}

func TestDecoderErr(t *testing.T) {
	t.Run("typical", func(t *testing.T) {
		c := qt.New(t)

		type testBody struct {
			Director string `json:"director"`
			Writer   string `json:"writer"`
		}

		requestBody := []byte(`{
				"director": "Alex Cox",
				"writer": "Alex Cox"
			}`)

		r, err := http.NewRequest(http.MethodPost, "/fake", bytes.NewBuffer(requestBody))
		if err != nil {
			t.Fatalf("http.NewRequest() error = %v", err)
		}

		// Decode JSON HTTP request body into a Decoder type
		// and unmarshal that into the testBody struct. DecoderErr
		// wraps errors from Decode when body is nil, json is malformed
		// or any other error
		wantBody := new(testBody)
		err = DecoderErr(json.NewDecoder(r.Body).Decode(&wantBody))
		defer r.Body.Close()
		c.Assert(err, qt.IsNil)
	})

	t.Run("malformed JSON", func(t *testing.T) {
		c := qt.New(t)

		type testBody struct {
			Director string `json:"director"`
			Writer   string `json:"writer"`
		}

		// removed trailing curly bracket
		requestBody := []byte(`{
				"director": "Alex Cox",
				"writer": "Alex Cox"`)

		r, err := http.NewRequest(http.MethodPost, "/fake", bytes.NewBuffer(requestBody))
		if err != nil {
			t.Fatalf("http.NewRequest() error = %v", err)
		}

		// Decode JSON HTTP request body into a Decoder type
		// and unmarshal that into the testBody struct. DecoderErr
		// wraps errors from Decode when body is nil, JSON is malformed
		// or any other error
		wantBody := new(testBody)
		err = DecoderErr(json.NewDecoder(r.Body).Decode(&wantBody))
		defer r.Body.Close()

		wantErr := errs.E(errs.InvalidRequest, errors.New("Malformed JSON"))
		c.Assert(errs.Match(err, wantErr), qt.IsTrue)
	})

	t.Run("empty request body", func(t *testing.T) {
		c := qt.New(t)

		type testBody struct {
			Director string `json:"director"`
			Writer   string `json:"writer"`
		}

		// empty body
		requestBody := []byte("")

		r, err := http.NewRequest(http.MethodPost, "/fake", bytes.NewBuffer(requestBody))
		if err != nil {
			t.Fatalf("http.NewRequest() error = %v", err)
		}

		// Decode JSON HTTP request body into a Decoder type
		// and unmarshal that into the testBody struct. DecoderErr
		// wraps errors from Decode when body is nil, JSON is malformed
		// or any other error
		wantBody := new(testBody)
		err = DecoderErr(json.NewDecoder(r.Body).Decode(&wantBody))
		defer r.Body.Close()

		wantErr := errs.E(errs.InvalidRequest, errors.New("Request Body cannot be empty"))
		c.Assert(errs.Match(err, wantErr), qt.IsTrue)
	})

	t.Run("invalid request body", func(t *testing.T) {
		c := qt.New(t)

		type testBody struct {
			Director string `json:"director"`
			Writer   string `json:"writer"`
		}

		// has unknown field
		requestBody := []byte(`{
				"director": "Alex Cox",
				"writer": "Alex Cox",
                "unknown_field": "I should fail"
			}`)

		r, err := http.NewRequest(http.MethodPost, "/fake", bytes.NewBuffer(requestBody))
		if err != nil {
			t.Fatalf("http.NewRequest() error = %v", err)
		}

		// force an error with DisallowUnknownFields
		wantBody := new(testBody)
		decoder := json.NewDecoder(r.Body)
		decoder.DisallowUnknownFields()
		err = DecoderErr(decoder.Decode(&wantBody))
		defer r.Body.Close()

		// check to make sure I have an error
		c.Assert(err != nil, qt.Equals, true)
	})
}
<|endoftext|>