<commit_before>/*
Copyright 2014 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 nfs

import (
	"fmt"
	"os"
	"testing"

	"k8s.io/mount-utils"

	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/kubernetes/fake"
	utiltesting "k8s.io/client-go/util/testing"
	"k8s.io/kubernetes/pkg/volume"
	volumetest "k8s.io/kubernetes/pkg/volume/testing"
)

func TestCanSupport(t *testing.T) {
	tmpDir, err := utiltesting.MkTmpdir("nfs_test")
	if err != nil {
		t.Fatalf("error creating temp dir: %v", err)
	}
	defer os.RemoveAll(tmpDir)

	plugMgr := volume.VolumePluginMgr{}
	plugMgr.InitPlugins(ProbeVolumePlugins(volume.VolumeConfig{}), nil /* prober */, volumetest.NewFakeVolumeHost(t, tmpDir, nil, nil))
	plug, err := plugMgr.FindPluginByName("kubernetes.io/nfs")
	if err != nil {
		t.Errorf("Can't find the plugin by name")
	}
	if plug.GetPluginName() != "kubernetes.io/nfs" {
		t.Errorf("Wrong name: %s", plug.GetPluginName())
	}

	if !plug.CanSupport(&volume.Spec{Volume: &v1.Volume{VolumeSource: v1.VolumeSource{NFS: &v1.NFSVolumeSource{}}}}) {
		t.Errorf("Expected true")
	}
	if !plug.CanSupport(&volume.Spec{PersistentVolume: &v1.PersistentVolume{Spec: v1.PersistentVolumeSpec{PersistentVolumeSource: v1.PersistentVolumeSource{NFS: &v1.NFSVolumeSource{}}}}}) {
		t.Errorf("Expected true")
	}
	if plug.CanSupport(&volume.Spec{Volume: &v1.Volume{VolumeSource: v1.VolumeSource{}}}) {
		t.Errorf("Expected false")
	}
}

func TestGetAccessModes(t *testing.T) {
	tmpDir, err := utiltesting.MkTmpdir("nfs_test")
	if err != nil {
		t.Fatalf("error creating temp dir: %v", err)
	}
	defer os.RemoveAll(tmpDir)

	plugMgr := volume.VolumePluginMgr{}
	plugMgr.InitPlugins(ProbeVolumePlugins(volume.VolumeConfig{}), nil /* prober */, volumetest.NewFakeVolumeHost(t, tmpDir, nil, nil))

	plug, err := plugMgr.FindPersistentPluginByName("kubernetes.io/nfs")
	if err != nil {
		t.Errorf("Can't find the plugin by name")
	}
	if !volumetest.ContainsAccessMode(plug.GetAccessModes(), v1.ReadWriteOnce) || !volumetest.ContainsAccessMode(plug.GetAccessModes(), v1.ReadOnlyMany) || !volumetest.ContainsAccessMode(plug.GetAccessModes(), v1.ReadWriteMany) {
		t.Errorf("Expected three AccessModeTypes:  %s, %s, and %s", v1.ReadWriteOnce, v1.ReadOnlyMany, v1.ReadWriteMany)
	}
}

func TestRecycler(t *testing.T) {
	tmpDir, err := utiltesting.MkTmpdir("nfs_test")
	if err != nil {
		t.Fatalf("error creating temp dir: %v", err)
	}
	defer os.RemoveAll(tmpDir)

	plugMgr := volume.VolumePluginMgr{}
	plugMgr.InitPlugins([]volume.VolumePlugin{&nfsPlugin{nil, volume.VolumeConfig{}}}, nil, volumetest.NewFakeVolumeHost(t, tmpDir, nil, nil))

	spec := &volume.Spec{PersistentVolume: &v1.PersistentVolume{Spec: v1.PersistentVolumeSpec{PersistentVolumeSource: v1.PersistentVolumeSource{NFS: &v1.NFSVolumeSource{Path: "/foo"}}}}}
	_, pluginErr := plugMgr.FindRecyclablePluginBySpec(spec)
	if pluginErr != nil {
		t.Errorf("Can't find the plugin by name")
	}
}

func doTestPlugin(t *testing.T, spec *volume.Spec) {
	tmpDir, err := utiltesting.MkTmpdir("nfs_test")
	if err != nil {
		t.Fatalf("error creating temp dir: %v", err)
	}
	defer os.RemoveAll(tmpDir)

	plugMgr := volume.VolumePluginMgr{}
	plugMgr.InitPlugins(ProbeVolumePlugins(volume.VolumeConfig{}), nil /* prober */, volumetest.NewFakeVolumeHost(t, tmpDir, nil, nil))
	plug, err := plugMgr.FindPluginByName("kubernetes.io/nfs")
	if err != nil {
		t.Errorf("Can't find the plugin by name")
	}
	fake := mount.NewFakeMounter(nil)
	pod := &v1.Pod{ObjectMeta: metav1.ObjectMeta{UID: types.UID("poduid")}}
	mounter, err := plug.(*nfsPlugin).newMounterInternal(spec, pod, fake)
	if err != nil {
		t.Errorf("Failed to make a new Mounter: %v", err)
	}
	if mounter == nil {
		t.Errorf("Got a nil Mounter")
	}
	volumePath := mounter.GetPath()
	expectedPath := fmt.Sprintf("%s/pods/poduid/volumes/kubernetes.io~nfs/vol1", tmpDir)
	if volumePath != expectedPath {
		t.Errorf("Unexpected path, expected %q, got: %q", expectedPath, volumePath)
	}
	if err := mounter.SetUp(volume.MounterArgs{}); err != nil {
		t.Errorf("Expected success, got: %v", err)
	}
	if _, err := os.Stat(volumePath); err != nil {
		if os.IsNotExist(err) {
			t.Errorf("SetUp() failed, volume path not created: %s", volumePath)
		} else {
			t.Errorf("SetUp() failed: %v", err)
		}
	}
	if mounter.(*nfsMounter).readOnly {
		t.Errorf("The volume source should not be read-only and it is.")
	}
	mntDevs, err := fake.List()
	if err != nil {
		t.Errorf("fakeMounter.List() failed: %v", err)
	}
	if len(mntDevs) != 1 {
		t.Errorf("unexpected number of mounted devices. expected: %v, got %v", 1, len(mntDevs))
	} else {
		if mntDevs[0].Type != "nfs" {
			t.Errorf("unexpected type of mounted devices. expected: %v, got %v", "nfs", mntDevs[0].Type)
		}
		src, _, _ := getVolumeSource(spec)
		srvr := getServerFromSource(src)
		expectedDevice := fmt.Sprintf("%s:%s", srvr, src.Path)
		if mntDevs[0].Device != expectedDevice {
			t.Errorf("unexpected nfs device, expected %q, got: %q", expectedDevice, mntDevs[0].Device)
		}
	}
	log := fake.GetLog()
	if len(log) != 1 {
		t.Errorf("Mount was not called exactly one time. It was called %d times.", len(log))
	} else {
		if log[0].Action != mount.FakeActionMount {
			t.Errorf("Unexpected mounter action: %#v", log[0])
		}
	}
	fake.ResetLog()

	unmounter, err := plug.(*nfsPlugin).newUnmounterInternal("vol1", types.UID("poduid"), fake)
	if err != nil {
		t.Errorf("Failed to make a new Unmounter: %v", err)
	}
	if unmounter == nil {
		t.Errorf("Got a nil Unmounter")
	}
	if err := unmounter.TearDown(); err != nil {
		t.Errorf("Expected success, got: %v", err)
	}
	if _, err := os.Stat(volumePath); err == nil {
		t.Errorf("TearDown() failed, volume path still exists: %s", volumePath)
	} else if !os.IsNotExist(err) {
		t.Errorf("TearDown() failed: %v", err)
	}
	log = fake.GetLog()
	if len(log) != 1 {
		t.Errorf("Unmount was not called exactly one time. It was called %d times.", len(log))
	} else {
		if log[0].Action != mount.FakeActionUnmount {
			t.Errorf("Unexpected unmounter action: %#v", log[0])
		}
	}

	fake.ResetLog()
}

func TestPluginVolume(t *testing.T) {
	vol := &v1.Volume{
		Name:         "vol1",
		VolumeSource: v1.VolumeSource{NFS: &v1.NFSVolumeSource{Server: "localhost", Path: "/somepath", ReadOnly: false}},
	}
	doTestPlugin(t, volume.NewSpecFromVolume(vol))
}

func TestIPV6VolumeSource(t *testing.T) {
	vol := &v1.Volume{
		Name:         "vol1",
		VolumeSource: v1.VolumeSource{NFS: &v1.NFSVolumeSource{Server: "0:0:0:0:0:0:0:1", Path: "/somepath", ReadOnly: false}},
	}
	doTestPlugin(t, volume.NewSpecFromVolume(vol))
}

func TestIPV4VolumeSource(t *testing.T) {
	vol := &v1.Volume{
		Name:         "vol1",
		VolumeSource: v1.VolumeSource{NFS: &v1.NFSVolumeSource{Server: "127.0.0.1", Path: "/somepath", ReadOnly: false}},
	}
	doTestPlugin(t, volume.NewSpecFromVolume(vol))
}

func TestPluginPersistentVolume(t *testing.T) {
	vol := &v1.PersistentVolume{
		ObjectMeta: metav1.ObjectMeta{
			Name: "vol1",
		},
		Spec: v1.PersistentVolumeSpec{
			PersistentVolumeSource: v1.PersistentVolumeSource{
				NFS: &v1.NFSVolumeSource{Server: "localhost", Path: "/somepath", ReadOnly: false},
			},
		},
	}

	doTestPlugin(t, volume.NewSpecFromPersistentVolume(vol, false))
}

func TestPersistentClaimReadOnlyFlag(t *testing.T) {
	tmpDir, err := utiltesting.MkTmpdir("nfs_test")
	if err != nil {
		t.Fatalf("error creating temp dir: %v", err)
	}
	defer os.RemoveAll(tmpDir)

	pv := &v1.PersistentVolume{
		ObjectMeta: metav1.ObjectMeta{
			Name: "pvA",
		},
		Spec: v1.PersistentVolumeSpec{
			PersistentVolumeSource: v1.PersistentVolumeSource{
				NFS: &v1.NFSVolumeSource{},
			},
			ClaimRef: &v1.ObjectReference{
				Name: "claimA",
			},
		},
	}

	claim := &v1.PersistentVolumeClaim{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "claimA",
			Namespace: "nsA",
		},
		Spec: v1.PersistentVolumeClaimSpec{
			VolumeName: "pvA",
		},
		Status: v1.PersistentVolumeClaimStatus{
			Phase: v1.ClaimBound,
		},
	}

	client := fake.NewSimpleClientset(pv, claim)

	plugMgr := volume.VolumePluginMgr{}
	plugMgr.InitPlugins(ProbeVolumePlugins(volume.VolumeConfig{}), nil /* prober */, volumetest.NewFakeVolumeHost(t, tmpDir, client, nil))
	plug, _ := plugMgr.FindPluginByName(nfsPluginName)

	// readOnly bool is supplied by persistent-claim volume source when its mounter creates other volumes
	spec := volume.NewSpecFromPersistentVolume(pv, true)
	pod := &v1.Pod{ObjectMeta: metav1.ObjectMeta{UID: types.UID("poduid")}}
	mounter, _ := plug.NewMounter(spec, pod, volume.VolumeOptions{})
	if mounter == nil {
		t.Fatalf("Got a nil Mounter")
	}

	if !mounter.GetAttributes().ReadOnly {
		t.Errorf("Expected true for mounter.IsReadOnly")
	}
}
<commit_msg>add expectedDevice as arg for testing<commit_after>/*
Copyright 2014 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 nfs

import (
	"fmt"
	"os"
	"testing"

	"k8s.io/mount-utils"

	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/kubernetes/fake"
	utiltesting "k8s.io/client-go/util/testing"
	"k8s.io/kubernetes/pkg/volume"
	volumetest "k8s.io/kubernetes/pkg/volume/testing"
)

func TestCanSupport(t *testing.T) {
	tmpDir, err := utiltesting.MkTmpdir("nfs_test")
	if err != nil {
		t.Fatalf("error creating temp dir: %v", err)
	}
	defer os.RemoveAll(tmpDir)

	plugMgr := volume.VolumePluginMgr{}
	plugMgr.InitPlugins(ProbeVolumePlugins(volume.VolumeConfig{}), nil /* prober */, volumetest.NewFakeVolumeHost(t, tmpDir, nil, nil))
	plug, err := plugMgr.FindPluginByName("kubernetes.io/nfs")
	if err != nil {
		t.Errorf("Can't find the plugin by name")
	}
	if plug.GetPluginName() != "kubernetes.io/nfs" {
		t.Errorf("Wrong name: %s", plug.GetPluginName())
	}

	if !plug.CanSupport(&volume.Spec{Volume: &v1.Volume{VolumeSource: v1.VolumeSource{NFS: &v1.NFSVolumeSource{}}}}) {
		t.Errorf("Expected true")
	}
	if !plug.CanSupport(&volume.Spec{PersistentVolume: &v1.PersistentVolume{Spec: v1.PersistentVolumeSpec{PersistentVolumeSource: v1.PersistentVolumeSource{NFS: &v1.NFSVolumeSource{}}}}}) {
		t.Errorf("Expected true")
	}
	if plug.CanSupport(&volume.Spec{Volume: &v1.Volume{VolumeSource: v1.VolumeSource{}}}) {
		t.Errorf("Expected false")
	}
}

func TestGetAccessModes(t *testing.T) {
	tmpDir, err := utiltesting.MkTmpdir("nfs_test")
	if err != nil {
		t.Fatalf("error creating temp dir: %v", err)
	}
	defer os.RemoveAll(tmpDir)

	plugMgr := volume.VolumePluginMgr{}
	plugMgr.InitPlugins(ProbeVolumePlugins(volume.VolumeConfig{}), nil /* prober */, volumetest.NewFakeVolumeHost(t, tmpDir, nil, nil))

	plug, err := plugMgr.FindPersistentPluginByName("kubernetes.io/nfs")
	if err != nil {
		t.Errorf("Can't find the plugin by name")
	}
	if !volumetest.ContainsAccessMode(plug.GetAccessModes(), v1.ReadWriteOnce) || !volumetest.ContainsAccessMode(plug.GetAccessModes(), v1.ReadOnlyMany) || !volumetest.ContainsAccessMode(plug.GetAccessModes(), v1.ReadWriteMany) {
		t.Errorf("Expected three AccessModeTypes:  %s, %s, and %s", v1.ReadWriteOnce, v1.ReadOnlyMany, v1.ReadWriteMany)
	}
}

func TestRecycler(t *testing.T) {
	tmpDir, err := utiltesting.MkTmpdir("nfs_test")
	if err != nil {
		t.Fatalf("error creating temp dir: %v", err)
	}
	defer os.RemoveAll(tmpDir)

	plugMgr := volume.VolumePluginMgr{}
	plugMgr.InitPlugins([]volume.VolumePlugin{&nfsPlugin{nil, volume.VolumeConfig{}}}, nil, volumetest.NewFakeVolumeHost(t, tmpDir, nil, nil))

	spec := &volume.Spec{PersistentVolume: &v1.PersistentVolume{Spec: v1.PersistentVolumeSpec{PersistentVolumeSource: v1.PersistentVolumeSource{NFS: &v1.NFSVolumeSource{Path: "/foo"}}}}}
	_, pluginErr := plugMgr.FindRecyclablePluginBySpec(spec)
	if pluginErr != nil {
		t.Errorf("Can't find the plugin by name")
	}
}

func doTestPlugin(t *testing.T, spec *volume.Spec, expectedDevice string) {
	tmpDir, err := utiltesting.MkTmpdir("nfs_test")
	if err != nil {
		t.Fatalf("error creating temp dir: %v", err)
	}
	defer os.RemoveAll(tmpDir)

	plugMgr := volume.VolumePluginMgr{}
	plugMgr.InitPlugins(ProbeVolumePlugins(volume.VolumeConfig{}), nil /* prober */, volumetest.NewFakeVolumeHost(t, tmpDir, nil, nil))
	plug, err := plugMgr.FindPluginByName("kubernetes.io/nfs")
	if err != nil {
		t.Errorf("Can't find the plugin by name")
	}
	fake := mount.NewFakeMounter(nil)
	pod := &v1.Pod{ObjectMeta: metav1.ObjectMeta{UID: types.UID("poduid")}}
	mounter, err := plug.(*nfsPlugin).newMounterInternal(spec, pod, fake)
	if err != nil {
		t.Errorf("Failed to make a new Mounter: %v", err)
	}
	if mounter == nil {
		t.Errorf("Got a nil Mounter")
	}
	volumePath := mounter.GetPath()
	expectedPath := fmt.Sprintf("%s/pods/poduid/volumes/kubernetes.io~nfs/vol1", tmpDir)
	if volumePath != expectedPath {
		t.Errorf("Unexpected path, expected %q, got: %q", expectedPath, volumePath)
	}
	if err := mounter.SetUp(volume.MounterArgs{}); err != nil {
		t.Errorf("Expected success, got: %v", err)
	}
	if _, err := os.Stat(volumePath); err != nil {
		if os.IsNotExist(err) {
			t.Errorf("SetUp() failed, volume path not created: %s", volumePath)
		} else {
			t.Errorf("SetUp() failed: %v", err)
		}
	}
	if mounter.(*nfsMounter).readOnly {
		t.Errorf("The volume source should not be read-only and it is.")
	}
	mntDevs, err := fake.List()
	if err != nil {
		t.Errorf("fakeMounter.List() failed: %v", err)
	}
	if len(mntDevs) != 1 {
		t.Errorf("unexpected number of mounted devices. expected: %v, got %v", 1, len(mntDevs))
	} else {
		if mntDevs[0].Type != "nfs" {
			t.Errorf("unexpected type of mounted devices. expected: %v, got %v", "nfs", mntDevs[0].Type)
		}
		if mntDevs[0].Device != expectedDevice {
			t.Errorf("unexpected nfs device, expected %q, got: %q", expectedDevice, mntDevs[0].Device)
		}
	}
	log := fake.GetLog()
	if len(log) != 1 {
		t.Errorf("Mount was not called exactly one time. It was called %d times.", len(log))
	} else {
		if log[0].Action != mount.FakeActionMount {
			t.Errorf("Unexpected mounter action: %#v", log[0])
		}
	}
	fake.ResetLog()

	unmounter, err := plug.(*nfsPlugin).newUnmounterInternal("vol1", types.UID("poduid"), fake)
	if err != nil {
		t.Errorf("Failed to make a new Unmounter: %v", err)
	}
	if unmounter == nil {
		t.Errorf("Got a nil Unmounter")
	}
	if err := unmounter.TearDown(); err != nil {
		t.Errorf("Expected success, got: %v", err)
	}
	if _, err := os.Stat(volumePath); err == nil {
		t.Errorf("TearDown() failed, volume path still exists: %s", volumePath)
	} else if !os.IsNotExist(err) {
		t.Errorf("TearDown() failed: %v", err)
	}
	log = fake.GetLog()
	if len(log) != 1 {
		t.Errorf("Unmount was not called exactly one time. It was called %d times.", len(log))
	} else {
		if log[0].Action != mount.FakeActionUnmount {
			t.Errorf("Unexpected unmounter action: %#v", log[0])
		}
	}

	fake.ResetLog()
}

func TestPluginVolume(t *testing.T) {
	vol := &v1.Volume{
		Name:         "vol1",
		VolumeSource: v1.VolumeSource{NFS: &v1.NFSVolumeSource{Server: "localhost", Path: "/somepath", ReadOnly: false}},
	}
	doTestPlugin(t, volume.NewSpecFromVolume(vol), "localhost:/somepath")
}

func TestIPV6VolumeSource(t *testing.T) {
	vol := &v1.Volume{
		Name:         "vol1",
		VolumeSource: v1.VolumeSource{NFS: &v1.NFSVolumeSource{Server: "0:0:0:0:0:0:0:1", Path: "/somepath", ReadOnly: false}},
	}
	doTestPlugin(t, volume.NewSpecFromVolume(vol), "[0:0:0:0:0:0:0:1]:/somepath")
}

func TestIPV4VolumeSource(t *testing.T) {
	vol := &v1.Volume{
		Name:         "vol1",
		VolumeSource: v1.VolumeSource{NFS: &v1.NFSVolumeSource{Server: "127.0.0.1", Path: "/somepath", ReadOnly: false}},
	}
	doTestPlugin(t, volume.NewSpecFromVolume(vol), "127.0.0.1:/somepath")
}

func TestPluginPersistentVolume(t *testing.T) {
	vol := &v1.PersistentVolume{
		ObjectMeta: metav1.ObjectMeta{
			Name: "vol1",
		},
		Spec: v1.PersistentVolumeSpec{
			PersistentVolumeSource: v1.PersistentVolumeSource{
				NFS: &v1.NFSVolumeSource{Server: "localhost", Path: "/somepath", ReadOnly: false},
			},
		},
	}

	doTestPlugin(t, volume.NewSpecFromPersistentVolume(vol, false), "localhost:/somepath")
}

func TestPersistentClaimReadOnlyFlag(t *testing.T) {
	tmpDir, err := utiltesting.MkTmpdir("nfs_test")
	if err != nil {
		t.Fatalf("error creating temp dir: %v", err)
	}
	defer os.RemoveAll(tmpDir)

	pv := &v1.PersistentVolume{
		ObjectMeta: metav1.ObjectMeta{
			Name: "pvA",
		},
		Spec: v1.PersistentVolumeSpec{
			PersistentVolumeSource: v1.PersistentVolumeSource{
				NFS: &v1.NFSVolumeSource{},
			},
			ClaimRef: &v1.ObjectReference{
				Name: "claimA",
			},
		},
	}

	claim := &v1.PersistentVolumeClaim{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "claimA",
			Namespace: "nsA",
		},
		Spec: v1.PersistentVolumeClaimSpec{
			VolumeName: "pvA",
		},
		Status: v1.PersistentVolumeClaimStatus{
			Phase: v1.ClaimBound,
		},
	}

	client := fake.NewSimpleClientset(pv, claim)

	plugMgr := volume.VolumePluginMgr{}
	plugMgr.InitPlugins(ProbeVolumePlugins(volume.VolumeConfig{}), nil /* prober */, volumetest.NewFakeVolumeHost(t, tmpDir, client, nil))
	plug, _ := plugMgr.FindPluginByName(nfsPluginName)

	// readOnly bool is supplied by persistent-claim volume source when its mounter creates other volumes
	spec := volume.NewSpecFromPersistentVolume(pv, true)
	pod := &v1.Pod{ObjectMeta: metav1.ObjectMeta{UID: types.UID("poduid")}}
	mounter, _ := plug.NewMounter(spec, pod, volume.VolumeOptions{})
	if mounter == nil {
		t.Fatalf("Got a nil Mounter")
	}

	if !mounter.GetAttributes().ReadOnly {
		t.Errorf("Expected true for mounter.IsReadOnly")
	}
}
<|endoftext|>


<commit_before>package sentryapi

import (
	"flag"
	"fmt"
	"os"
	"os/user"
	"strconv"
	"testing"
)

const (
	defaultHost = "localhost"
	defaultPort = "8038"
	component   = "solr"
	hostEnv     = "SENTRY_HOST"
	portEnv     = "SENTRY_PORT"
	userEnv     = "SENTRY_USER"
)

var (
	client        ClientAPI
	genericClient ClientAPI
)

func TestMain(m *testing.M) {
	// Read config from environment vars and create Hive and
	// generic clients.
	flag.Parse()
	host := os.Getenv(hostEnv)
	port := os.Getenv(portEnv)
	userName := os.Getenv(userEnv)

	if host == "" {
		host = defaultHost
	}
	if port == "" {
		port = defaultPort
	}
	if userName == "" {
		currentUser, _ := user.Current()
		userName = currentUser.Username
	}

	portVal, err := strconv.Atoi(port)
	if err != nil {
		panic("invalid port" + port)
	}

	// Create Hive protocol client
	client, err = GetClient(PolicyProtocol, host, portVal, "", userName)
	if err != nil {
		panic(err)
	}
	// Create generic protocol client
	genericClient, err = GetClient(PolicyProtocol, host, portVal, component, userName)
	if err != nil {
		panic(err)
	}
	os.Exit(m.Run())
}

func createRemoveRole(client ClientAPI, roleName string) error {
	if err := client.CreateRole(roleName); err != nil {
		return fmt.Errorf("failed to create role %s: %v", roleName, err)
	}
	if err := client.RemoveRole(roleName); err != nil {
		return fmt.Errorf("failed to delete role %s: %v", roleName, err)

	}
	return nil
}

func ExampleSentryClient_CreateRole() {
	roleName := "exampleRole"
	err := client.CreateRole(roleName)
	if err != nil {
		fmt.Println("failed to create role: ", err)
		os.Exit(1)
	}
	fmt.Println("Created role", roleName)
	err = client.RemoveRole(roleName)
	if err != nil {
		fmt.Println("failed to create role: ", err)
		os.Exit(1)
	}
	fmt.Println("Removed role", roleName)
	// Output:
	// Created role exampleRole
	// Removed role exampleRole
}

func TestSentryClient_CreateRole(t *testing.T) {
	if err := createRemoveRole(client, "sentryTestRole"); err != nil {
		t.Error(err)
	}
}

// Repeatedly create and remove a role
func BenchmarkSentryClient_CreateRole(b *testing.B) {
	for i := 0; i < b.N; i++ {
		if err := createRemoveRole(client, "testRole"); err != nil {
			b.Fail()
		}
	}
}

// Repeatedly create roles then repeatedly remove them
func BenchmarkSentryClient_CreateRoles(b *testing.B) {
	roleName := "testRole"
	for i := 0; i < b.N; i++ {
		name := fmt.Sprintf("%s_%d", roleName, i)
		err := client.CreateRole(name)
		if err != nil {
			b.Fail()
		}
	}
	for i := 0; i < b.N; i++ {
		name := fmt.Sprintf("%s_%d", roleName, i)
		err := client.RemoveRole(name)
		if err != nil {
			b.Errorf("can't remove %s: %v", name, err)
			continue
		}
	}
}

func TestGenericSentryClient_CreateRole(t *testing.T) {
	if err := createRemoveRole(genericClient,
		"sentryGenTestRole"); err != nil {
		t.Error(err)
	}
}

// Repeatedly create and remove a role
func BenchmarkGenericSentryClient_CreateRole(b *testing.B) {
	for i := 0; i < b.N; i++ {
		if err := createRemoveRole(genericClient, "testRole"); err != nil {
			b.Fail()
		}
	}
}

// Repeatedly create roles then repeatedly remove them
func BenchmarkGenericSentryClient_CreateRoles(b *testing.B) {
	roleName := "testRole"
	for i := 0; i < b.N; i++ {
		name := fmt.Sprintf("%s_%d", roleName, i)
		err := genericClient.CreateRole(name)
		if err != nil {
			b.Fail()
		}
	}
	for i := 0; i < b.N; i++ {
		name := fmt.Sprintf("%s_%d", roleName, i)
		err := genericClient.RemoveRole(name)
		if err != nil {
			b.Errorf("can't remove %s: %v", name, err)
			continue
		}
	}
}
<commit_msg>Add two more benchmarks<commit_after>package sentryapi

import (
	"flag"
	"fmt"
	"os"
	"os/user"
	"strconv"
	"testing"
)

const (
	defaultHost = "localhost"
	defaultPort = "8038"
	component   = "solr"
	hostEnv     = "SENTRY_HOST"
	portEnv     = "SENTRY_PORT"
	userEnv     = "SENTRY_USER"
)

var (
	client        ClientAPI
	genericClient ClientAPI
)

func TestMain(m *testing.M) {
	// Read config from environment vars and create Hive and
	// generic clients.
	flag.Parse()
	host := os.Getenv(hostEnv)
	port := os.Getenv(portEnv)
	userName := os.Getenv(userEnv)

	if host == "" {
		host = defaultHost
	}
	if port == "" {
		port = defaultPort
	}
	if userName == "" {
		currentUser, _ := user.Current()
		userName = currentUser.Username
	}

	portVal, err := strconv.Atoi(port)
	if err != nil {
		panic("invalid port" + port)
	}

	// Create Hive protocol client
	client, err = GetClient(PolicyProtocol, host, portVal, "", userName)
	if err != nil {
		panic(err)
	}
	// Create generic protocol client
	genericClient, err = GetClient(PolicyProtocol, host, portVal, component, userName)
	if err != nil {
		panic(err)
	}
	os.Exit(m.Run())
}

func createRemoveRole(client ClientAPI, roleName string) error {
	if err := client.CreateRole(roleName); err != nil {
		return fmt.Errorf("failed to create role %s: %v", roleName, err)
	}
	if err := client.RemoveRole(roleName); err != nil {
		return fmt.Errorf("failed to delete role %s: %v", roleName, err)

	}
	return nil
}

func ExampleSentryClient_CreateRole() {
	roleName := "exampleRole"
	err := client.CreateRole(roleName)
	if err != nil {
		fmt.Println("failed to create role: ", err)
		os.Exit(1)
	}
	fmt.Println("Created role", roleName)
	err = client.RemoveRole(roleName)
	if err != nil {
		fmt.Println("failed to create role: ", err)
		os.Exit(1)
	}
	fmt.Println("Removed role", roleName)
	// Output:
	// Created role exampleRole
	// Removed role exampleRole
}

func TestSentryClient_CreateRole(t *testing.T) {
	if err := createRemoveRole(client, "sentryTestRole"); err != nil {
		t.Error(err)
	}
}

// Repeatedly create and remove a role
func BenchmarkSentryClient_CreateRole(b *testing.B) {
	for i := 0; i < b.N; i++ {
		if err := createRemoveRole(client, "testRole"); err != nil {
			b.Fail()
		}
	}
}

// Repeatedly create roles then repeatedly remove them
func BenchmarkSentryClient_CreateRoles(b *testing.B) {
	roleName := "testRole"
	for i := 0; i < b.N; i++ {
		name := fmt.Sprintf("%s_%d", roleName, i)
		err := client.CreateRole(name)
		if err != nil {
			b.Fail()
		}
	}
	for i := 0; i < b.N; i++ {
		name := fmt.Sprintf("%s_%d", roleName, i)
		err := client.RemoveRole(name)
		if err != nil {
			b.Errorf("can't remove %s: %v", name, err)
			continue
		}
	}
}

func TestGenericSentryClient_CreateRole(t *testing.T) {
	if err := createRemoveRole(genericClient,
		"sentryGenTestRole"); err != nil {
		t.Error(err)
	}
}

// Repeatedly create and remove a role
func BenchmarkGenericSentryClient_CreateRole(b *testing.B) {
	for i := 0; i < b.N; i++ {
		if err := createRemoveRole(genericClient, "testRole"); err != nil {
			b.Fail()
		}
	}
}

// Repeatedly create roles then repeatedly remove them
func BenchmarkGenericSentryClient_CreateRoles(b *testing.B) {
	roleName := "testRole"
	for i := 0; i < b.N; i++ {
		name := fmt.Sprintf("%s_%d", roleName, i)
		err := genericClient.CreateRole(name)
		if err != nil {
			b.Fail()
		}
	}
	for i := 0; i < b.N; i++ {
		name := fmt.Sprintf("%s_%d", roleName, i)
		err := genericClient.RemoveRole(name)
		if err != nil {
			b.Errorf("can't remove %s: %v", name, err)
			continue
		}
	}
}

func BenchmarkSentryClient_ListPrivilegesByRole(b *testing.B) {
	roleName := "testRole"
	err := client.CreateRole(roleName)
	if err != nil {
		b.Fail()
	}

	for i := 0; i < b.N; i++ {
		_, err := client.ListPrivilegesByRole(roleName)
		if err != nil {
			b.Errorf("can't list privileges")
			break
		}
	}

	err = client.RemoveRole(roleName)
	if err != nil {
		b.Fail()
	}
}

func BenchmarkSentryClient_GrantAndRevokePrivilege(b *testing.B) {
	roleName := "testRole"
	err := client.CreateRole(roleName)
	if err != nil {
		b.Fail()
	}
	priv := Privilege{
		Server:      "some_server",
		Database:    "some_database",
		Table:       "some_table",
		Column:      "some_column",
		URI:         "some_uri",
		Action:      "all",
		GrantOption: true,
	}

	for i := 0; i < b.N; i++ {
		err := client.GrantPrivilege(roleName, &priv)
		if err != nil {
			b.Errorf("can't grant privilege: %v", err)
			break
		}
		err = client.RevokePrivilege(roleName, &priv)
		if err != nil {
			b.Errorf("can't revoke privilege: %v", err)
			break
		}
	}

	err = client.RemoveRole(roleName)
	if err != nil {
		b.Fail()
	}
}
<|endoftext|>


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

package models

import (
	"fmt"

	"github.com/Unknwon/com"

	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/markdown"
	"code.gitea.io/gitea/modules/setting"
)

func (issue *Issue) mailSubject() string {
	return fmt.Sprintf("[%s] %s (#%d)", issue.Repo.Name, issue.Title, issue.Index)
}

// mailIssueCommentToParticipants can be used for only for comment.
// This function sends two list of emails:
// 1. Repository watchers and users who are participated in comments.
// 2. Users who are not in 1. but get mentioned in current issue/comment.
func mailIssueCommentToParticipants(issue *Issue, doer *User, comment *Comment, mentions []string) error {

	names, tos, err := prepareMailToParticipants(issue, doer)

	if(err != nil) {
		return fmt.Errorf("PrepareMailToParticipants: %v", err)
	}

	SendIssueCommentMail(issue, doer, comment, tos)

	// Mail mentioned people and exclude watchers.
	names = append(names, doer.Name)
	tos = make([]string, 0, len(mentions)) // list of user names.
	for i := range mentions {
		if com.IsSliceContainsStr(names, mentions[i]) {
			continue
		}

		tos = append(tos, mentions[i])
	}
	SendIssueMentionMail(issue, doer, comment, GetUserEmailsByNames(tos))

	return nil
}

func mailIssueActionToParticipants(issue *Issue, doer *User, mentions []string) error {
	names, tos, err := prepareMailToParticipants(issue, doer)

	if(err != nil) {
		return fmt.Errorf("PrepareMailToParticipants: %v", err)
	}

	SendIssueActionMail(issue, doer, tos)

	// Mail mentioned people and exclude watchers.
	names = append(names, doer.Name)
	tos = make([]string, 0, len(mentions)) // list of user names.
	for i := range mentions {
		if com.IsSliceContainsStr(names, mentions[i]) {
			continue
		}

		tos = append(tos, mentions[i])
	}
	SendIssueMentionInActionMail(issue, doer, GetUserEmailsByNames(tos))

	return nil
}

// prepareMailToParticipants creates the tos and names list for an issue and the issue's creator.
func prepareMailToParticipants(issue *Issue, doer *User) ([]string, []string, error)  {
	if !setting.Service.EnableNotifyMail {
		return nil, nil, nil
	}

	watchers, err := GetWatchers(issue.RepoID)
	if err != nil {
		return nil, nil, err
	}
	participants, err := GetParticipantsByIssueID(issue.ID)
	if err != nil {
		return nil, nil, err
	}

	// In case the issue poster is not watching the repository,
	// even if we have duplicated in watchers, can be safely filtered out.
	if issue.PosterID != doer.ID {
		participants = append(participants, issue.Poster)
	}

	tos := make([]string, 0, len(watchers)) // List of email addresses.
	names := make([]string, 0, len(watchers))
	for i := range watchers {
		if watchers[i].UserID == doer.ID {
			continue
		}

		to, err := GetUserByID(watchers[i].UserID)
		if err != nil {
			return nil, nil, err
		}
		if to.IsOrganization() {
			continue
		}

		tos = append(tos, to.Email)
		names = append(names, to.Name)
	}
	for i := range participants {
		if participants[i].ID == doer.ID {
			continue
		} else if com.IsSliceContainsStr(names, participants[i].Name) {
			continue
		}

		tos = append(tos, participants[i].Email)
		names = append(names, participants[i].Name)
	}
	return tos, names, nil
}

// MailParticipants sends new issue thread created emails to repository watchers
// and mentioned people.
func (issue *Issue) MailParticipants() (err error) {
	mentions := markdown.FindAllMentions(issue.Content)
	if err = UpdateIssueMentions(x, issue.ID, mentions); err != nil {
		return fmt.Errorf("UpdateIssueMentions [%d]: %v", issue.ID, err)
	}

	if err = mailIssueActionToParticipants(issue, issue.Poster, mentions); err != nil {
		log.Error(4, "mailIssueCommentToParticipants: %v", err)
	}

	return nil
}

<commit_msg>Added comment's hashtag to url for mail notifications. Added explanation to return statement.<commit_after>// Copyright 2016 The Gogs Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package models

import (
	"fmt"

	"github.com/Unknwon/com"

	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/markdown"
	"code.gitea.io/gitea/modules/setting"
)

func (issue *Issue) mailSubject() string {
	return fmt.Sprintf("[%s] %s (#%d)", issue.Repo.Name, issue.Title, issue.Index)
}

// mailIssueCommentToParticipants can be used for only for comment.
// This function sends two list of emails:
// 1. Repository watchers and users who are participated in comments.
// 2. Users who are not in 1. but get mentioned in current issue/comment.
func mailIssueCommentToParticipants(issue *Issue, doer *User, comment *Comment, mentions []string) error {

	names, tos, err := prepareMailToParticipants(issue, doer)

	if(err != nil) {
		return fmt.Errorf("PrepareMailToParticipants: %v", err)
	}

	SendIssueCommentMail(issue, doer, comment, tos)

	// Mail mentioned people and exclude watchers.
	names = append(names, doer.Name)
	tos = make([]string, 0, len(mentions)) // list of user names.
	for i := range mentions {
		if com.IsSliceContainsStr(names, mentions[i]) {
			continue
		}

		tos = append(tos, mentions[i])
	}
	SendIssueMentionMail(issue, doer, comment, GetUserEmailsByNames(tos))

	return nil
}

func mailIssueActionToParticipants(issue *Issue, doer *User, mentions []string) error {
	names, tos, err := prepareMailToParticipants(issue, doer)

	if(err != nil) {
		return fmt.Errorf("PrepareMailToParticipants: %v", err)
	}

	SendIssueActionMail(issue, doer, tos)

	// Mail mentioned people and exclude watchers.
	names = append(names, doer.Name)
	tos = make([]string, 0, len(mentions)) // list of user names.
	for i := range mentions {
		if com.IsSliceContainsStr(names, mentions[i]) {
			continue
		}

		tos = append(tos, mentions[i])
	}
	SendIssueMentionInActionMail(issue, doer, GetUserEmailsByNames(tos))

	return nil
}

// prepareMailToParticipants creates the tos and names list for an issue and the issue's creator.
func prepareMailToParticipants(issue *Issue, doer *User) (tos []string, names []string, error error)  {
	if !setting.Service.EnableNotifyMail {
		return nil, nil, nil
	}

	watchers, err := GetWatchers(issue.RepoID)
	if err != nil {
		return nil, nil, err
	}
	participants, err := GetParticipantsByIssueID(issue.ID)
	if err != nil {
		return nil, nil, err
	}

	// In case the issue poster is not watching the repository,
	// even if we have duplicated in watchers, can be safely filtered out.
	if issue.PosterID != doer.ID {
		participants = append(participants, issue.Poster)
	}

	tos := make([]string, 0, len(watchers)) // List of email addresses.
	names := make([]string, 0, len(watchers))
	for i := range watchers {
		if watchers[i].UserID == doer.ID {
			continue
		}

		to, err := GetUserByID(watchers[i].UserID)
		if err != nil {
			return nil, nil, err
		}
		if to.IsOrganization() {
			continue
		}

		tos = append(tos, to.Email)
		names = append(names, to.Name)
	}
	for i := range participants {
		if participants[i].ID == doer.ID {
			continue
		} else if com.IsSliceContainsStr(names, participants[i].Name) {
			continue
		}

		tos = append(tos, participants[i].Email)
		names = append(names, participants[i].Name)
	}
	return tos, names, nil
}

// MailParticipants sends new issue thread created emails to repository watchers
// and mentioned people.
func (issue *Issue) MailParticipants() (err error) {
	mentions := markdown.FindAllMentions(issue.Content)
	if err = UpdateIssueMentions(x, issue.ID, mentions); err != nil {
		return fmt.Errorf("UpdateIssueMentions [%d]: %v", issue.ID, err)
	}

	if err = mailIssueActionToParticipants(issue, issue.Poster, mentions); err != nil {
		log.Error(4, "mailIssueCommentToParticipants: %v", err)
	}

	return nil
}

<|endoftext|>


<commit_before>package kubernetes

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

	"github.com/coredns/caddy"
	"github.com/coredns/coredns/core/dnsserver"
	"github.com/coredns/coredns/plugin"
	"github.com/coredns/coredns/plugin/pkg/dnsutil"
	clog "github.com/coredns/coredns/plugin/pkg/log"
	"github.com/coredns/coredns/plugin/pkg/upstream"

	"github.com/miekg/dns"
	meta "k8s.io/apimachinery/pkg/apis/meta/v1"
	_ "k8s.io/client-go/plugin/pkg/client/auth/gcp"       // pull this in here, because we want it excluded if plugin.cfg doesn't have k8s
	_ "k8s.io/client-go/plugin/pkg/client/auth/oidc"      // pull this in here, because we want it excluded if plugin.cfg doesn't have k8s
	_ "k8s.io/client-go/plugin/pkg/client/auth/openstack" // pull this in here, because we want it excluded if plugin.cfg doesn't have k8s
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/klog/v2"
)

const pluginName = "kubernetes"

var log = clog.NewWithPlugin(pluginName)

func init() { plugin.Register(pluginName, setup) }

func setup(c *caddy.Controller) error {
	klog.InitFlags(nil)
	klog.SetOutput(os.Stdout)

	k, err := kubernetesParse(c)
	if err != nil {
		return plugin.Error(pluginName, err)
	}

	onStart, onShut, err := k.InitKubeCache(context.Background())
	if err != nil {
		return plugin.Error(pluginName, err)
	}
	if onStart != nil {
		c.OnStartup(onStart)
	}
	if onShut != nil {
		c.OnShutdown(onShut)
	}

	dnsserver.GetConfig(c).AddPlugin(func(next plugin.Handler) plugin.Handler {
		k.Next = next
		return k
	})

	// get locally bound addresses
	c.OnStartup(func() error {
		k.localIPs = boundIPs(c)
		return nil
	})

	return nil
}

func kubernetesParse(c *caddy.Controller) (*Kubernetes, error) {
	var (
		k8s *Kubernetes
		err error
	)

	i := 0
	for c.Next() {
		if i > 0 {
			return nil, plugin.ErrOnce
		}
		i++

		k8s, err = ParseStanza(c)
		if err != nil {
			return k8s, err
		}
	}
	return k8s, nil
}

// ParseStanza parses a kubernetes stanza
func ParseStanza(c *caddy.Controller) (*Kubernetes, error) {

	k8s := New([]string{""})
	k8s.autoPathSearch = searchFromResolvConf()

	opts := dnsControlOpts{
		initEndpointsCache: true,
		ignoreEmptyService: false,
	}
	k8s.opts = opts

	k8s.Zones = plugin.OriginsFromArgsOrServerBlock(c.RemainingArgs(), c.ServerBlockKeys)

	k8s.primaryZoneIndex = -1
	for i, z := range k8s.Zones {
		if dnsutil.IsReverse(z) > 0 {
			continue
		}
		k8s.primaryZoneIndex = i
		break
	}

	if k8s.primaryZoneIndex == -1 {
		return nil, errors.New("non-reverse zone name must be used")
	}

	k8s.Upstream = upstream.New()

	for c.NextBlock() {
		switch c.Val() {
		case "endpoint_pod_names":
			args := c.RemainingArgs()
			if len(args) > 0 {
				return nil, c.ArgErr()
			}
			k8s.endpointNameMode = true
			continue
		case "pods":
			args := c.RemainingArgs()
			if len(args) == 1 {
				switch args[0] {
				case podModeDisabled, podModeInsecure, podModeVerified:
					k8s.podMode = args[0]
				default:
					return nil, fmt.Errorf("wrong value for pods: %s,  must be one of: disabled, verified, insecure", args[0])
				}
				continue
			}
			return nil, c.ArgErr()
		case "namespaces":
			args := c.RemainingArgs()
			if len(args) > 0 {
				for _, a := range args {
					k8s.Namespaces[a] = struct{}{}
				}
				continue
			}
			return nil, c.ArgErr()
		case "endpoint":
			args := c.RemainingArgs()
			if len(args) > 0 {
				// Multiple endpoints are deprecated but still could be specified,
				// only the first one be used, though
				k8s.APIServerList = args
				if len(args) > 1 {
					log.Warningf("Multiple endpoints have been deprecated, only the first specified endpoint '%s' is used", args[0])
				}
				continue
			}
			return nil, c.ArgErr()
		case "tls": // cert key cacertfile
			args := c.RemainingArgs()
			if len(args) == 3 {
				k8s.APIClientCert, k8s.APIClientKey, k8s.APICertAuth = args[0], args[1], args[2]
				continue
			}
			return nil, c.ArgErr()
		case "labels":
			args := c.RemainingArgs()
			if len(args) > 0 {
				labelSelectorString := strings.Join(args, " ")
				ls, err := meta.ParseToLabelSelector(labelSelectorString)
				if err != nil {
					return nil, fmt.Errorf("unable to parse label selector value: '%v': %v", labelSelectorString, err)
				}
				k8s.opts.labelSelector = ls
				continue
			}
			return nil, c.ArgErr()
		case "namespace_labels":
			args := c.RemainingArgs()
			if len(args) > 0 {
				namespaceLabelSelectorString := strings.Join(args, " ")
				nls, err := meta.ParseToLabelSelector(namespaceLabelSelectorString)
				if err != nil {
					return nil, fmt.Errorf("unable to parse namespace_label selector value: '%v': %v", namespaceLabelSelectorString, err)
				}
				k8s.opts.namespaceLabelSelector = nls
				continue
			}
			return nil, c.ArgErr()
		case "fallthrough":
			k8s.Fall.SetZonesFromArgs(c.RemainingArgs())
		case "ttl":
			args := c.RemainingArgs()
			if len(args) == 0 {
				return nil, c.ArgErr()
			}
			t, err := strconv.Atoi(args[0])
			if err != nil {
				return nil, err
			}
			if t < 0 || t > 3600 {
				return nil, c.Errf("ttl must be in range [0, 3600]: %d", t)
			}
			k8s.ttl = uint32(t)
		case "noendpoints":
			if len(c.RemainingArgs()) != 0 {
				return nil, c.ArgErr()
			}
			k8s.opts.initEndpointsCache = false
		case "ignore":
			args := c.RemainingArgs()
			if len(args) > 0 {
				ignore := args[0]
				if ignore == "empty_service" {
					k8s.opts.ignoreEmptyService = true
					continue
				} else {
					return nil, fmt.Errorf("unable to parse ignore value: '%v'", ignore)
				}
			}
		case "kubeconfig":
			args := c.RemainingArgs()
			if len(args) != 1 && len(args) != 2 {
				return nil, c.ArgErr()
			}
			overrides := &clientcmd.ConfigOverrides{}
			if len(args) == 2 {
				overrides.CurrentContext = args[1]
			}
			config := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(
				&clientcmd.ClientConfigLoadingRules{ExplicitPath: args[0]},
				overrides,
			)
			k8s.ClientConfig = config
		default:
			return nil, c.Errf("unknown property '%s'", c.Val())
		}
	}

	if len(k8s.Namespaces) != 0 && k8s.opts.namespaceLabelSelector != nil {
		return nil, c.Errf("namespaces and namespace_labels cannot both be set")
	}

	return k8s, nil
}

func searchFromResolvConf() []string {
	rc, err := dns.ClientConfigFromFile("/etc/resolv.conf")
	if err != nil {
		return nil
	}
	plugin.Zones(rc.Search).Normalize()
	return rc.Search
}
<commit_msg>plugin/kubernetes: fix reload panic (#4881)<commit_after>package kubernetes

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

	"github.com/coredns/caddy"
	"github.com/coredns/coredns/core/dnsserver"
	"github.com/coredns/coredns/plugin"
	"github.com/coredns/coredns/plugin/pkg/dnsutil"
	clog "github.com/coredns/coredns/plugin/pkg/log"
	"github.com/coredns/coredns/plugin/pkg/upstream"

	"github.com/miekg/dns"
	meta "k8s.io/apimachinery/pkg/apis/meta/v1"
	_ "k8s.io/client-go/plugin/pkg/client/auth/gcp"       // pull this in here, because we want it excluded if plugin.cfg doesn't have k8s
	_ "k8s.io/client-go/plugin/pkg/client/auth/oidc"      // pull this in here, because we want it excluded if plugin.cfg doesn't have k8s
	_ "k8s.io/client-go/plugin/pkg/client/auth/openstack" // pull this in here, because we want it excluded if plugin.cfg doesn't have k8s
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/klog/v2"
)

const pluginName = "kubernetes"

var log = clog.NewWithPlugin(pluginName)

func init() { plugin.Register(pluginName, setup) }

func setup(c *caddy.Controller) error {
	// Do not call klog.InitFlags(nil) here.  It will cause reload to panic.
	klog.SetOutput(os.Stdout)

	k, err := kubernetesParse(c)
	if err != nil {
		return plugin.Error(pluginName, err)
	}

	onStart, onShut, err := k.InitKubeCache(context.Background())
	if err != nil {
		return plugin.Error(pluginName, err)
	}
	if onStart != nil {
		c.OnStartup(onStart)
	}
	if onShut != nil {
		c.OnShutdown(onShut)
	}

	dnsserver.GetConfig(c).AddPlugin(func(next plugin.Handler) plugin.Handler {
		k.Next = next
		return k
	})

	// get locally bound addresses
	c.OnStartup(func() error {
		k.localIPs = boundIPs(c)
		return nil
	})

	return nil
}

func kubernetesParse(c *caddy.Controller) (*Kubernetes, error) {
	var (
		k8s *Kubernetes
		err error
	)

	i := 0
	for c.Next() {
		if i > 0 {
			return nil, plugin.ErrOnce
		}
		i++

		k8s, err = ParseStanza(c)
		if err != nil {
			return k8s, err
		}
	}
	return k8s, nil
}

// ParseStanza parses a kubernetes stanza
func ParseStanza(c *caddy.Controller) (*Kubernetes, error) {

	k8s := New([]string{""})
	k8s.autoPathSearch = searchFromResolvConf()

	opts := dnsControlOpts{
		initEndpointsCache: true,
		ignoreEmptyService: false,
	}
	k8s.opts = opts

	k8s.Zones = plugin.OriginsFromArgsOrServerBlock(c.RemainingArgs(), c.ServerBlockKeys)

	k8s.primaryZoneIndex = -1
	for i, z := range k8s.Zones {
		if dnsutil.IsReverse(z) > 0 {
			continue
		}
		k8s.primaryZoneIndex = i
		break
	}

	if k8s.primaryZoneIndex == -1 {
		return nil, errors.New("non-reverse zone name must be used")
	}

	k8s.Upstream = upstream.New()

	for c.NextBlock() {
		switch c.Val() {
		case "endpoint_pod_names":
			args := c.RemainingArgs()
			if len(args) > 0 {
				return nil, c.ArgErr()
			}
			k8s.endpointNameMode = true
			continue
		case "pods":
			args := c.RemainingArgs()
			if len(args) == 1 {
				switch args[0] {
				case podModeDisabled, podModeInsecure, podModeVerified:
					k8s.podMode = args[0]
				default:
					return nil, fmt.Errorf("wrong value for pods: %s,  must be one of: disabled, verified, insecure", args[0])
				}
				continue
			}
			return nil, c.ArgErr()
		case "namespaces":
			args := c.RemainingArgs()
			if len(args) > 0 {
				for _, a := range args {
					k8s.Namespaces[a] = struct{}{}
				}
				continue
			}
			return nil, c.ArgErr()
		case "endpoint":
			args := c.RemainingArgs()
			if len(args) > 0 {
				// Multiple endpoints are deprecated but still could be specified,
				// only the first one be used, though
				k8s.APIServerList = args
				if len(args) > 1 {
					log.Warningf("Multiple endpoints have been deprecated, only the first specified endpoint '%s' is used", args[0])
				}
				continue
			}
			return nil, c.ArgErr()
		case "tls": // cert key cacertfile
			args := c.RemainingArgs()
			if len(args) == 3 {
				k8s.APIClientCert, k8s.APIClientKey, k8s.APICertAuth = args[0], args[1], args[2]
				continue
			}
			return nil, c.ArgErr()
		case "labels":
			args := c.RemainingArgs()
			if len(args) > 0 {
				labelSelectorString := strings.Join(args, " ")
				ls, err := meta.ParseToLabelSelector(labelSelectorString)
				if err != nil {
					return nil, fmt.Errorf("unable to parse label selector value: '%v': %v", labelSelectorString, err)
				}
				k8s.opts.labelSelector = ls
				continue
			}
			return nil, c.ArgErr()
		case "namespace_labels":
			args := c.RemainingArgs()
			if len(args) > 0 {
				namespaceLabelSelectorString := strings.Join(args, " ")
				nls, err := meta.ParseToLabelSelector(namespaceLabelSelectorString)
				if err != nil {
					return nil, fmt.Errorf("unable to parse namespace_label selector value: '%v': %v", namespaceLabelSelectorString, err)
				}
				k8s.opts.namespaceLabelSelector = nls
				continue
			}
			return nil, c.ArgErr()
		case "fallthrough":
			k8s.Fall.SetZonesFromArgs(c.RemainingArgs())
		case "ttl":
			args := c.RemainingArgs()
			if len(args) == 0 {
				return nil, c.ArgErr()
			}
			t, err := strconv.Atoi(args[0])
			if err != nil {
				return nil, err
			}
			if t < 0 || t > 3600 {
				return nil, c.Errf("ttl must be in range [0, 3600]: %d", t)
			}
			k8s.ttl = uint32(t)
		case "noendpoints":
			if len(c.RemainingArgs()) != 0 {
				return nil, c.ArgErr()
			}
			k8s.opts.initEndpointsCache = false
		case "ignore":
			args := c.RemainingArgs()
			if len(args) > 0 {
				ignore := args[0]
				if ignore == "empty_service" {
					k8s.opts.ignoreEmptyService = true
					continue
				} else {
					return nil, fmt.Errorf("unable to parse ignore value: '%v'", ignore)
				}
			}
		case "kubeconfig":
			args := c.RemainingArgs()
			if len(args) != 1 && len(args) != 2 {
				return nil, c.ArgErr()
			}
			overrides := &clientcmd.ConfigOverrides{}
			if len(args) == 2 {
				overrides.CurrentContext = args[1]
			}
			config := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(
				&clientcmd.ClientConfigLoadingRules{ExplicitPath: args[0]},
				overrides,
			)
			k8s.ClientConfig = config
		default:
			return nil, c.Errf("unknown property '%s'", c.Val())
		}
	}

	if len(k8s.Namespaces) != 0 && k8s.opts.namespaceLabelSelector != nil {
		return nil, c.Errf("namespaces and namespace_labels cannot both be set")
	}

	return k8s, nil
}

func searchFromResolvConf() []string {
	rc, err := dns.ClientConfigFromFile("/etc/resolv.conf")
	if err != nil {
		return nil
	}
	plugin.Zones(rc.Search).Normalize()
	return rc.Search
}
<|endoftext|>


<commit_before>package hmutil

import (
	// "log"
	// "os/exec"
	"strings"
	"fmt"
	"os"
	"path/filepath"
	"bytes"
	"io/ioutil"
	"io"
	"crypto/aes"
	"crypto/cipher"
	"compress/gzip"
	"archive/tar"

	"github.com/rlmcpherson/s3gof3r"
	"github.com/laher/scp-go/scp"
	"github.com/dutchcoders/goftp"
)

// func System(cmd string) ([]byte, error) {
// 	log.Println(cmd)
// 	return exec.Command("sh", "-c", cmd).CombinedOutput()
// }

func ReplaceVars(str string, replacements map[string]string) string {
	for from, to := range replacements {
		str = strings.Replace(str, from, to, -1)
	}
	return str
}

func find(prefix string, arr []string, elem string) int {
	for i, e := range arr {
		if filepath.Join(prefix, e) == elem {
			return i
		}
	}
	return -1
}

func removeFromSlice(prefix string, arr []string, elem string) ([]string, bool) {
	index := find(prefix, arr, elem)

	if index == -1 {
		return arr, false
	} else {
		return append(arr[:index], arr[index+1:]...), true
	}
}

func removeExcludedFiles(dir string, files, excludeFiles []string) {
	var changed bool
	var deleted []string = []string{}

	for _, e := range excludeFiles {
		files, changed = removeFromSlice(dir, files, e)
		if changed {
			deleted = append(deleted, e)
		}
	}
	for _, e := range deleted {
		removeFromSlice("", excludeFiles, e)
	}
}

func Tar(dir string, files, excludeFiles []string, tw *tar.Writer) error {
	if len(files) == 1 && files[0] == "*" {
		files = []string{}
		d, _err := ioutil.ReadDir(dir)

		if _err != nil {
			return _err
		}

		for _, e := range d {
			files = append(files, e.Name())
		}
	}

	// for _, e := range excludeFiles {
	// 	fmt.Println(filpath.Join(dir, e))
	// 	files = removeFromSlice(files, e)
	// 	excludeFiles = removeFromSlice(excludeFiles, e)
	// }

	removeExcludedFiles(dir, files, excludeFiles)

	for _, file := range files {
		f, _err := os.Open(filepath.Join(dir, file))
		if _err != nil {
			continue
		}

		s, _err := f.Stat()

		if s.IsDir() {
			Tar(filepath.Join(dir, file), []string{"*"}, excludeFiles, tw)
		} else {
			if _err != nil {
				return _err
			}

			header := &tar.Header{
				Name: filepath.Join(dir, s.Name()),
				Size: s.Size(),
				Mode: 0777,
			}

			if _err = tw.WriteHeader(header); _err != nil {
				return _err
			}

			buffer := make([]byte, s.Size())
			buffer, _err = ioutil.ReadFile(filepath.Join(dir, s.Name())) // file.Read(buffer)
			if _err != nil {
				return _err
			}

			if _, _err = tw.Write(buffer); _err != nil {
				return _err
			}
		}
	}

	return nil
}

// func Gzip(pathToFile string, buffer *bytes.Buffer) {
// 	gzFile, _err := os.Create(pathToFile)
// 	handleError(_err)

// 	gzWriter := gzip.NewWriter(gzFile)
// 	_, _err = gzWriter.Write(buffer.Bytes())
// 	handleError(_err)

// 	gzWriter.Close()
// }

func Gzip(buffer *bytes.Buffer) (bytes.Buffer, error) {
	var gzFile bytes.Buffer
	gzWriter := gzip.NewWriter(&gzFile)
	_, _err := gzWriter.Write(buffer.Bytes())
	gzWriter.Flush()
	gzWriter.Close()
	if _err != nil {
		return bytes.Buffer{}, _err
	}

	return gzFile, nil
}

func Encode(buffer *bytes.Buffer, encodeKey []byte) (bytes.Buffer, error) {
	var outbuffer bytes.Buffer

	block, _err := aes.NewCipher(encodeKey)
	if _err != nil {
		return bytes.Buffer{}, _err
	}

	var iv [aes.BlockSize]byte
	stream := cipher.NewOFB(block, iv[:])

	writer := &cipher.StreamWriter{S: stream, W: &outbuffer}
	defer writer.Close()
	if _, _err = io.Copy(writer, buffer); _err != nil {
		return bytes.Buffer{}, _err
	}

	return outbuffer, nil
}

func WriteToFile(path string, buffer bytes.Buffer) error {
	file, _err := os.Create(path)
	defer file.Close()
	if _err != nil {
		return _err
	}
	file.Write(buffer.Bytes())
	return nil
}

func PackAndCompress(dir string, files, excludeFiles []string, outputFile string, key []byte, encrypt bool) error {
	outdir, _ := filepath.Split(outputFile)
	_err := os.MkdirAll(outdir, 0777)
	if _err != nil {
		return _err
	}

	var tarFile bytes.Buffer

	tarWriter := tar.NewWriter(&tarFile)
	_err = Tar(dir, files, excludeFiles, tarWriter)

	if _err != nil {
		return _err
	}

	tarWriter.Close()
	gzipedBuffer, _err := Gzip(&tarFile)

	if _err != nil {
		return _err
	}

	if encrypt {
		encryptedBuffer, _err := Encode(&gzipedBuffer, key)
		if _err != nil {
			return _err
		}
		return WriteToFile(outputFile + ".encrypted", encryptedBuffer)
	} else {
		return WriteToFile(outputFile, gzipedBuffer)
	}
}

func handleError(e error) {
	if e != nil {
		panic(e)
	}
}

func UploadToS3(k s3gof3r.Keys, bucketName string, pathToFile string) error {
	s3 := s3gof3r.New("", k)
	b := s3.Bucket(bucketName)

	file, err := os.Open(pathToFile)
	if err != nil {
		return err
	}

	stats, _ := file.Stat()

	w, err := b.PutWriter(stats.Name(), nil, nil)
	if err != nil {
		return err
	}

	if _, err = io.Copy(w, file); err != nil { // Copy into S3
		return err
	}

	if err = w.Close(); err != nil {
		return err
	}

	return nil
}

func DownloadFromS3(k s3gof3r.Keys, bucketName, filename string) (bytes.Buffer, error) {
	s3 := s3gof3r.New("", k)
	b := s3.Bucket(bucketName)

	r, _, err := b.GetReader(filename, nil)
	if err != nil {
		return bytes.Buffer{}, err
	}

	// outputFile, err := os.Create(filepath.Join(outputDirPath, filename))
	// if err != nil {
	// 	return nil, err
	// }

	var output bytes.Buffer

	if _, err = io.Copy(&output, r); err != nil {
		return bytes.Buffer{}, err
	}

	err = r.Close()
	if err != nil {
		return bytes.Buffer{}, err
	}

	return output, nil
}

func SSHUploader(port int, id_rsa, host, user, srcFile, dstFile string) *scp.SecureCopier {
	return scp.NewSecureCopier(port, false, false, false, true, false, false, id_rsa, "", "", srcFile, host, user, dstFile)
}

func SSHDownloader(port int, id_rsa, host, user, srcFile, dstFile string) *scp.SecureCopier {
	return scp.NewSecureCopier(port, false, false, false, true, false, false, id_rsa, host, user, srcFile, "", "", dstFile)
}

func SSHExec(ssh *scp.SecureCopier) (bytes.Buffer, error) {
	var o, e, result bytes.Buffer
	err, _ := ssh.Exec(nil, &o, &e, &result)
	return result, err
}

func prepareFTPConn(port int, host, login, password string) (*goftp.FTP, error) {
	var ftp *goftp.FTP
	var err error

	if ftp, err = goftp.Connect(fmt.Sprintf("%s:%d", host, port)); err != nil {
		return nil, err
	}

	if err = ftp.Login(login, password); err != nil {
		return nil, err
	}

	return ftp, nil
}

func FTPUpload(port int, host, login, password, pathToFile, pathToRemoteFile string) error {
	ftp, err := prepareFTPConn(port, host, login, password)
	if err != nil {
		return err
	}

	file, err := os.Open(pathToFile)
	if err != nil {
		return err
	}

	if err = ftp.Stor(pathToRemoteFile, file); err != nil {
		return err
	}

	return ftp.Quit()
}

func FTPDownload(port int, host, login, password, pathToRemoteFile string) (bytes.Buffer, error) {
	ftp, err := prepareFTPConn(port, host, login, password)
	if err != nil {
		return bytes.Buffer{}, err
	}

	var buffer bytes.Buffer

	_, err = ftp.Retr(pathToRemoteFile, func(r io.Reader) error {
		// file, _err := os.Create(pathToFile)
		// if _err != nil {
		// 	return _err
		// }

		if _, _err := io.Copy(&buffer, r); _err != nil {
			return _err
		}

		return nil
	})

	if err != nil {
		return bytes.Buffer{}, err
	}

	if err = ftp.Quit(); err != nil {
		return buffer, err
	}

	return buffer, nil
}

func Decode(buffer *bytes.Buffer, key []byte) (bytes.Buffer, error) {
	block, err := aes.NewCipher(key)
  if err != nil {
    return bytes.Buffer{}, err
  }

  var iv [aes.BlockSize]byte
  stream := cipher.NewOFB(block, iv[:])

  var outbuffer bytes.Buffer
  reader := &cipher.StreamReader{S: stream, R: buffer}

  if _, err = io.Copy(&outbuffer, reader); err != nil {
    return bytes.Buffer{}, err
  }

  return outbuffer, nil
}

/*
func ErrString(err error) (s *string) {
	if err != nil {
		tmp := err.Error()
		s = &tmp
	}
	return
}
*/
<commit_msg>Change scp lib<commit_after>package hmutil

import (
	// "log"
	// "os/exec"
	"strings"
	"fmt"
	"os"
	"path/filepath"
	"bytes"
	"io/ioutil"
	"io"
	"crypto/aes"
	"crypto/cipher"
	"compress/gzip"
	"archive/tar"

	"github.com/rlmcpherson/s3gof3r"
	"github.com/Twizty/scp-go/scp"
	"github.com/dutchcoders/goftp"
)

// func System(cmd string) ([]byte, error) {
// 	log.Println(cmd)
// 	return exec.Command("sh", "-c", cmd).CombinedOutput()
// }

func ReplaceVars(str string, replacements map[string]string) string {
	for from, to := range replacements {
		str = strings.Replace(str, from, to, -1)
	}
	return str
}

func find(prefix string, arr []string, elem string) int {
	for i, e := range arr {
		if filepath.Join(prefix, e) == elem {
			return i
		}
	}
	return -1
}

func removeFromSlice(prefix string, arr []string, elem string) ([]string, bool) {
	index := find(prefix, arr, elem)

	if index == -1 {
		return arr, false
	} else {
		return append(arr[:index], arr[index+1:]...), true
	}
}

func removeExcludedFiles(dir string, files, excludeFiles []string) {
	var changed bool
	var deleted []string = []string{}

	for _, e := range excludeFiles {
		files, changed = removeFromSlice(dir, files, e)
		if changed {
			deleted = append(deleted, e)
		}
	}
	for _, e := range deleted {
		removeFromSlice("", excludeFiles, e)
	}
}

func Tar(dir string, files, excludeFiles []string, tw *tar.Writer) error {
	if len(files) == 1 && files[0] == "*" {
		files = []string{}
		d, _err := ioutil.ReadDir(dir)

		if _err != nil {
			return _err
		}

		for _, e := range d {
			files = append(files, e.Name())
		}
	}

	// for _, e := range excludeFiles {
	// 	fmt.Println(filpath.Join(dir, e))
	// 	files = removeFromSlice(files, e)
	// 	excludeFiles = removeFromSlice(excludeFiles, e)
	// }

	removeExcludedFiles(dir, files, excludeFiles)

	for _, file := range files {
		f, _err := os.Open(filepath.Join(dir, file))
		if _err != nil {
			continue
		}

		s, _err := f.Stat()

		if s.IsDir() {
			Tar(filepath.Join(dir, file), []string{"*"}, excludeFiles, tw)
		} else {
			if _err != nil {
				return _err
			}

			header := &tar.Header{
				Name: filepath.Join(dir, s.Name()),
				Size: s.Size(),
				Mode: 0777,
			}

			if _err = tw.WriteHeader(header); _err != nil {
				return _err
			}

			buffer := make([]byte, s.Size())
			buffer, _err = ioutil.ReadFile(filepath.Join(dir, s.Name())) // file.Read(buffer)
			if _err != nil {
				return _err
			}

			if _, _err = tw.Write(buffer); _err != nil {
				return _err
			}
		}
	}

	return nil
}

// func Gzip(pathToFile string, buffer *bytes.Buffer) {
// 	gzFile, _err := os.Create(pathToFile)
// 	handleError(_err)

// 	gzWriter := gzip.NewWriter(gzFile)
// 	_, _err = gzWriter.Write(buffer.Bytes())
// 	handleError(_err)

// 	gzWriter.Close()
// }

func Gzip(buffer *bytes.Buffer) (bytes.Buffer, error) {
	var gzFile bytes.Buffer
	gzWriter := gzip.NewWriter(&gzFile)
	_, _err := gzWriter.Write(buffer.Bytes())
	gzWriter.Flush()
	gzWriter.Close()
	if _err != nil {
		return bytes.Buffer{}, _err
	}

	return gzFile, nil
}

func Encode(buffer *bytes.Buffer, encodeKey []byte) (bytes.Buffer, error) {
	var outbuffer bytes.Buffer

	block, _err := aes.NewCipher(encodeKey)
	if _err != nil {
		return bytes.Buffer{}, _err
	}

	var iv [aes.BlockSize]byte
	stream := cipher.NewOFB(block, iv[:])

	writer := &cipher.StreamWriter{S: stream, W: &outbuffer}
	defer writer.Close()
	if _, _err = io.Copy(writer, buffer); _err != nil {
		return bytes.Buffer{}, _err
	}

	return outbuffer, nil
}

func WriteToFile(path string, buffer bytes.Buffer) error {
	file, _err := os.Create(path)
	defer file.Close()
	if _err != nil {
		return _err
	}
	file.Write(buffer.Bytes())
	return nil
}

func PackAndCompress(dir string, files, excludeFiles []string, outputFile string, key []byte, encrypt bool) error {
	outdir, _ := filepath.Split(outputFile)
	_err := os.MkdirAll(outdir, 0777)
	if _err != nil {
		return _err
	}

	var tarFile bytes.Buffer

	tarWriter := tar.NewWriter(&tarFile)
	_err = Tar(dir, files, excludeFiles, tarWriter)

	if _err != nil {
		return _err
	}

	tarWriter.Close()
	gzipedBuffer, _err := Gzip(&tarFile)

	if _err != nil {
		return _err
	}

	if encrypt {
		encryptedBuffer, _err := Encode(&gzipedBuffer, key)
		if _err != nil {
			return _err
		}
		return WriteToFile(outputFile + ".encrypted", encryptedBuffer)
	} else {
		return WriteToFile(outputFile, gzipedBuffer)
	}
}

func handleError(e error) {
	if e != nil {
		panic(e)
	}
}

func UploadToS3(k s3gof3r.Keys, bucketName string, pathToFile string) error {
	s3 := s3gof3r.New("", k)
	b := s3.Bucket(bucketName)

	file, err := os.Open(pathToFile)
	if err != nil {
		return err
	}

	stats, _ := file.Stat()

	w, err := b.PutWriter(stats.Name(), nil, nil)
	if err != nil {
		return err
	}

	if _, err = io.Copy(w, file); err != nil { // Copy into S3
		return err
	}

	if err = w.Close(); err != nil {
		return err
	}

	return nil
}

func DownloadFromS3(k s3gof3r.Keys, bucketName, filename string) (bytes.Buffer, error) {
	s3 := s3gof3r.New("", k)
	b := s3.Bucket(bucketName)

	r, _, err := b.GetReader(filename, nil)
	if err != nil {
		return bytes.Buffer{}, err
	}

	// outputFile, err := os.Create(filepath.Join(outputDirPath, filename))
	// if err != nil {
	// 	return nil, err
	// }

	var output bytes.Buffer

	if _, err = io.Copy(&output, r); err != nil {
		return bytes.Buffer{}, err
	}

	err = r.Close()
	if err != nil {
		return bytes.Buffer{}, err
	}

	return output, nil
}

func SSHUploader(port int, id_rsa, host, user, srcFile, dstFile string) *scp.SecureCopier {
	return scp.NewSecureCopier(port, false, false, false, true, false, false, id_rsa, "", "", srcFile, host, user, dstFile)
}

func SSHDownloader(port int, id_rsa, host, user, srcFile, dstFile string) *scp.SecureCopier {
	return scp.NewSecureCopier(port, false, false, false, true, false, false, id_rsa, host, user, srcFile, "", "", dstFile)
}

func SSHExec(ssh *scp.SecureCopier) (bytes.Buffer, error) {
	var o, e, result bytes.Buffer
	err, _ := ssh.Exec(nil, &o, &e, &result)
	return result, err
}

func prepareFTPConn(port int, host, login, password string) (*goftp.FTP, error) {
	var ftp *goftp.FTP
	var err error

	if ftp, err = goftp.Connect(fmt.Sprintf("%s:%d", host, port)); err != nil {
		return nil, err
	}

	if err = ftp.Login(login, password); err != nil {
		return nil, err
	}

	return ftp, nil
}

func FTPUpload(port int, host, login, password, pathToFile, pathToRemoteFile string) error {
	ftp, err := prepareFTPConn(port, host, login, password)
	if err != nil {
		return err
	}

	file, err := os.Open(pathToFile)
	if err != nil {
		return err
	}

	if err = ftp.Stor(pathToRemoteFile, file); err != nil {
		return err
	}

	return ftp.Quit()
}

func FTPDownload(port int, host, login, password, pathToRemoteFile string) (bytes.Buffer, error) {
	ftp, err := prepareFTPConn(port, host, login, password)
	if err != nil {
		return bytes.Buffer{}, err
	}

	var buffer bytes.Buffer

	_, err = ftp.Retr(pathToRemoteFile, func(r io.Reader) error {
		// file, _err := os.Create(pathToFile)
		// if _err != nil {
		// 	return _err
		// }

		if _, _err := io.Copy(&buffer, r); _err != nil {
			return _err
		}

		return nil
	})

	if err != nil {
		return bytes.Buffer{}, err
	}

	if err = ftp.Quit(); err != nil {
		return buffer, err
	}

	return buffer, nil
}

func Decode(buffer *bytes.Buffer, key []byte) (bytes.Buffer, error) {
	block, err := aes.NewCipher(key)
  if err != nil {
    return bytes.Buffer{}, err
  }

  var iv [aes.BlockSize]byte
  stream := cipher.NewOFB(block, iv[:])

  var outbuffer bytes.Buffer
  reader := &cipher.StreamReader{S: stream, R: buffer}

  if _, err = io.Copy(&outbuffer, reader); err != nil {
    return bytes.Buffer{}, err
  }

  return outbuffer, nil
}

/*
func ErrString(err error) (s *string) {
	if err != nil {
		tmp := err.Error()
		s = &tmp
	}
	return
}
*/
<|endoftext|>


<commit_before>package topiccsm

import (
	"fmt"
	"sync"
	"time"

	"github.com/mailgun/kafka-pixy/actor"
	"github.com/mailgun/kafka-pixy/config"
	"github.com/mailgun/kafka-pixy/consumer"
	"github.com/mailgun/kafka-pixy/consumer/dispatcher"
)

var requestTimeoutRs = dispatcher.Response{Err: consumer.ErrRequestTimeout}

// T implements a consumer request dispatch tier responsible for a particular
// topic. It receives requests on the `Requests()` channel and replies with
// messages received on `Messages()` channel. If there has been no message
// received for `Config.Consumer.LongPollingTimeout` then a timeout error is
// sent to the requests' reply channel.
//
// implements `dispatcher.Tier`.
// implements `multiplexer.Out`.
type T struct {
	actDesc            *actor.Descriptor
	childSpec          dispatcher.ChildSpec
	cfg                *config.Proxy
	group              string
	topic              string
	expireTimer        *time.Timer
	nilOrExpireTimerCh <-chan time.Time
	lifespanCh         chan<- *T
	messagesCh         chan consumer.Message
	wg                 sync.WaitGroup
}

// Creates a topic consumer instance. It should be explicitly started in
// accordance with the `dispatcher.Tier` contract.
func Spawn(parentActDesc *actor.Descriptor, group string, childSpec dispatcher.ChildSpec, cfg *config.Proxy, lifespanCh chan<- *T) *T {
	topic := string(childSpec.Key())
	actDesc := parentActDesc.NewChild(fmt.Sprintf("%s", topic))
	actDesc.AddLogField("kafka.group", group)
	actDesc.AddLogField("kafka.topic", topic)
	tc := T{
		actDesc:    actDesc,
		childSpec:  childSpec,
		cfg:        cfg,
		group:      group,
		topic:      topic,
		lifespanCh: lifespanCh,

		// Messages channel must be non-buffered. Otherwise we might end up
		// buffering a message from a partition that no longer belongs to this
		// consumer group member.
		messagesCh: make(chan consumer.Message),
	}
	actor.Spawn(tc.actDesc, &tc.wg, tc.run)
	return &tc
}

// Topic returns the topic name this topic consumer is responsible for.
func (tc *T) Topic() string {
	return tc.topic
}

// implements `multiplexer.Out`
func (tc *T) Messages() chan<- consumer.Message {
	return tc.messagesCh
}

func (tc *T) run() {
	defer tc.childSpec.Dispose()
	tc.lifespanCh <- tc
	defer func() {
		tc.lifespanCh <- tc
	}()

	for {
		select {
		case consumeReq, ok := <-tc.childSpec.Requests():
			if !ok {
				tc.actDesc.Log().Info("Signaled to shutdown")
				return
			}
			tc.stopExpireTimer()

			requestAge := time.Now().UTC().Sub(consumeReq.Timestamp)
			ttl := tc.cfg.Consumer.LongPollingTimeout - requestAge
			// The request has been waiting in the buffer for too long. If we
			// reply with a fetched message, then there is a good chance that the
			// client won't receive it due to the client HTTP timeout. Therefore
			// we reject the request to avoid message loss.
			if ttl <= 0 {
				consumeReq.ResponseCh <- requestTimeoutRs
				continue
			}
			select {
			case msg := <-tc.messagesCh:
				msg.EventsCh <- consumer.Event{consumer.EvOffered, msg.Offset}
				consumeReq.ResponseCh <- dispatcher.Response{Msg: msg}
			case <-time.After(ttl):
				consumeReq.ResponseCh <- requestTimeoutRs
			}
		case <-tc.nilOrExpireTimerCh:
			tc.actDesc.Log().Info("Topic registration expired")
			return
		default:
			tc.ensureExpireTimer()
		}
	}
}

func (tc *T) String() string {
	return tc.actDesc.String()
}

func (tc *T) stopExpireTimer() {
	if tc.expireTimer == nil {
		return
	}
	if !tc.expireTimer.Stop() {
		<-tc.expireTimer.C
	}
	tc.nilOrExpireTimerCh = nil
}

func (tc *T) ensureExpireTimer() {
	if tc.nilOrExpireTimerCh != nil {
		return
	}
	if tc.expireTimer == nil {
		tc.expireTimer = time.NewTimer(tc.cfg.Consumer.SubscriptionTimeout)
	} else {
		tc.expireTimer.Reset(tc.cfg.Consumer.SubscriptionTimeout)
	}
	tc.nilOrExpireTimerCh = tc.expireTimer.C
}
<commit_msg>Get rid of busy loop from topic consumer<commit_after>package topiccsm

import (
	"fmt"
	"sync"
	"time"

	"github.com/mailgun/kafka-pixy/actor"
	"github.com/mailgun/kafka-pixy/config"
	"github.com/mailgun/kafka-pixy/consumer"
	"github.com/mailgun/kafka-pixy/consumer/dispatcher"
)

var requestTimeoutRs = dispatcher.Response{Err: consumer.ErrRequestTimeout}

// T implements a consumer request dispatch tier responsible for a particular
// topic. It receives requests on the `Requests()` channel and replies with
// messages received on `Messages()` channel. If there has been no message
// received for `Config.Consumer.LongPollingTimeout` then a timeout error is
// sent to the requests' reply channel.
//
// implements `multiplexer.Out`.
type T struct {
	actDesc    *actor.Descriptor
	childSpec  dispatcher.ChildSpec
	cfg        *config.Proxy
	group      string
	topic      string
	lifespanCh chan<- *T
	messagesCh chan consumer.Message
	wg         sync.WaitGroup
}

// Creates a topic consumer instance. It should be explicitly started in
// accordance with the `dispatcher.Tier` contract.
func Spawn(parentActDesc *actor.Descriptor, group string, childSpec dispatcher.ChildSpec, cfg *config.Proxy, lifespanCh chan<- *T) *T {
	topic := string(childSpec.Key())
	actDesc := parentActDesc.NewChild(fmt.Sprintf("%s", topic))
	actDesc.AddLogField("kafka.group", group)
	actDesc.AddLogField("kafka.topic", topic)
	tc := T{
		actDesc:    actDesc,
		childSpec:  childSpec,
		cfg:        cfg,
		group:      group,
		topic:      topic,
		lifespanCh: lifespanCh,

		// Messages channel must be non-buffered. Otherwise we might end up
		// buffering a message from a partition that no longer belongs to this
		// consumer group member.
		messagesCh: make(chan consumer.Message),
	}
	actor.Spawn(tc.actDesc, &tc.wg, tc.run)
	return &tc
}

// Topic returns the topic name this topic consumer is responsible for.
func (tc *T) Topic() string {
	return tc.topic
}

// implements `multiplexer.Out`
func (tc *T) Messages() chan<- consumer.Message {
	return tc.messagesCh
}

func (tc *T) run() {
	defer tc.childSpec.Dispose()
	tc.lifespanCh <- tc
	defer func() {
		tc.lifespanCh <- tc
	}()

	expireTimer := time.NewTimer(tc.cfg.Consumer.SubscriptionTimeout)
	defer expireTimer.Stop()
	var lastRequestTime time.Time
	for {
		select {
		case consumeRq, ok := <-tc.childSpec.Requests():
			if !ok {
				tc.actDesc.Log().Info("Signaled to shutdown")
				return
			}
			lastRequestTime = time.Now().UTC()
			requestAge := lastRequestTime.Sub(consumeRq.Timestamp)
			requestTTL := tc.cfg.Consumer.LongPollingTimeout - requestAge
			// The request has been waiting in the buffer for too long. If we
			// reply with a fetched message, then there is a good chance that the
			// client won't receive it due to the client HTTP timeout. Therefore
			// we reject the request to avoid message loss.
			if requestTTL <= 0 {
				consumeRq.ResponseCh <- requestTimeoutRs
				continue
			}
			select {
			case msg := <-tc.messagesCh:
				msg.EventsCh <- consumer.Event{consumer.EvOffered, msg.Offset}
				consumeRq.ResponseCh <- dispatcher.Response{Msg: msg}
			case <-time.After(requestTTL):
				consumeRq.ResponseCh <- requestTimeoutRs
			}
		case <-expireTimer.C:
			now := time.Now().UTC()
			sinceLastRequest := now.Sub(lastRequestTime)
			subscriptionTTL := tc.cfg.Consumer.SubscriptionTimeout - sinceLastRequest
			if subscriptionTTL <= 0 {
				tc.actDesc.Log().Info("Topic subscription expired")
				return
			}
			expireTimer.Reset(subscriptionTTL)
		}
	}
}

func (tc *T) String() string {
	return tc.actDesc.String()
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"log"
	"os"
	"strconv"

	"gopkg.in/jcelliott/turnpike.v2"
)

func main() {
	turnpike.Debug()
	c, err := turnpike.NewWebsocketClient(turnpike.JSON, "ws://localhost:8000/", nil)
	if err != nil {
		log.Fatal(err)
	}
	_, err = c.JoinRealm("turnpike.examples", nil)
	if err != nil {
		log.Fatal(err)
	}

	quit := make(chan bool)
	c.Subscribe("alarm.ring", func([]interface{}, map[string]interface{}) {
		fmt.Println("The alarm rang!")
		c.Close()
		quit <- true
	})
	fmt.Print("Enter the timer duration: ")
	scanner := bufio.NewScanner(os.Stdin)
	scanner.Scan()
	if err := scanner.Err(); err != nil {
		log.Fatalln("reading stdin:", err)
	}
	text := scanner.Text()
	if duration, err := strconv.Atoi(text); err != nil {
		log.Fatalln("invalid integer input:", err)
	} else {
		if _, err := c.Call("alarm.set", []interface{}{duration}, nil); err != nil {
			log.Fatalln("error setting alarm:", err)
		}
	}
	<-quit
}
<commit_msg>Fix Call to pass in options<commit_after>package main

import (
	"bufio"
	"fmt"
	"log"
	"os"
	"strconv"

	"gopkg.in/jcelliott/turnpike.v2"
)

func main() {
	turnpike.Debug()
	c, err := turnpike.NewWebsocketClient(turnpike.JSON, "ws://localhost:8000/", nil)
	if err != nil {
		log.Fatal(err)
	}
	_, err = c.JoinRealm("turnpike.examples", nil)
	if err != nil {
		log.Fatal(err)
	}

	quit := make(chan bool)
	c.Subscribe("alarm.ring", func([]interface{}, map[string]interface{}) {
		fmt.Println("The alarm rang!")
		c.Close()
		quit <- true
	})
	fmt.Print("Enter the timer duration: ")
	scanner := bufio.NewScanner(os.Stdin)
	scanner.Scan()
	if err := scanner.Err(); err != nil {
		log.Fatalln("reading stdin:", err)
	}
	text := scanner.Text()
	if duration, err := strconv.Atoi(text); err != nil {
		log.Fatalln("invalid integer input:", err)
	} else {
		if _, err := c.Call("alarm.set", nil, []interface{}{duration}, nil); err != nil {
			log.Fatalln("error setting alarm:", err)
		}
	}
	<-quit
}
<|endoftext|>


<commit_before>package fluidbackup

import "fmt"

/* Represents a data store provider. Is the entity that actually
stores information. */
type FluidServer struct {
}

func (fv *FluidServer) PrintHelloWorld() {
	fmt.Println("Hellow World")
}

// A temporary store data file that is archaic
func (fv *FluidServer) StoreData(data []byte) {

}
<commit_msg>Update filestore.go to reflect name.<commit_after>package fluidbackup

import "fmt"

/* Represents a data store provider. Is the entity that actually
stores information. */
type FileStore struct {

}

// A temporary store data file that is archaic
func (this *FileStore) StoreData(path string, data []byte) {

}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"github.com/minio/minio-go"
	"github.com/minio/minio-go/pkg/policy"
	"github.com/pborman/uuid"
	"github.com/rcrowley/go-metrics"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"path"
	"syscall"
)

var (
	rootDir    string
	listenAddr string
	s3         S3
)

func init() {
	flag.StringVar(&rootDir, "root", "", "http root directory (default: .)")
	flag.StringVar(&listenAddr, "listen", "", "address to listen to (default: :8080)")
	flag.StringVar(&s3.endpoint, "s3-endpoint", "", "s3 endpoint")
	flag.StringVar(&s3.region, "s3-region", "", "s3 region")
	flag.StringVar(&s3.bucket, "s3-bucket", "", "s3 bucket")
	flag.StringVar(&s3.accessKey, "s3-access-key", "", "s3 access key")
	flag.StringVar(&s3.secretKey, "s3-secret-key", "", "s3 secret key")
}

func main() {
	flag.Parse()

	if len(rootDir) == 0 {
		listenAddr = os.Getenv("PUT_ROOT")
	}

	if len(rootDir) == 0 {
		rootDir = "."
	}

	if len(listenAddr) == 0 {
		listenAddr = os.Getenv("PUT_LISTEN")
	}

	if len(listenAddr) == 0 {
		listenAddr = ":8080"
	}

	if len(s3.endpoint) == 0 {
		s3.endpoint = os.Getenv("PUT_S3_ENDPOINT")
	}

	if len(s3.region) == 0 {
		s3.region = os.Getenv("PUT_S3_REGION")
	}

	if len(s3.bucket) == 0 {
		s3.bucket = os.Getenv("PUT_S3_BUCKET")
	}

	if len(s3.accessKey) == 0 {
		s3.accessKey = os.Getenv("PUT_S3_ACCESS_KEY")
	}

	if len(s3.secretKey) == 0 {
		s3.secretKey = os.Getenv("PUT_S3_SECRET_KEY")
	}
	l, err := listen(listenAddr)

	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	defer l.Close()

	metrics := NewMetrics()

	mux := http.NewServeMux()
	mux.Handle("/", NewRequestHandler(rootDir, &s3, metrics))
	mux.Handle("/health", NewHealthHandler(&s3))
	mux.Handle("/metrics", NewMetricsHandler(metrics, &s3))

	go http.Serve(l, mux)

	sig := make(chan os.Signal)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM, syscall.SIGKILL)
	<-sig
}

func listen(addr string) (net.Listener, error) {
	if len(addr) > 0 && addr[0] != '/' {
		return net.Listen("tcp", addr)
	}
	return net.Listen("unix", addr)
}

type S3 struct {
	minio     *minio.Client
	url       *url.URL
	endpoint  string
	bucket    string
	region    string
	secretKey string
	accessKey string
}

func (s3 *S3) Put(name, contentType string, body io.Reader) (string, error) {
	client, err := s3.client()

	if err != nil {
		return "", err
	}
	meta := &minio.PutObjectOptions{ContentType: contentType}

	if _, err := client.PutObject(s3.bucket, name, body, -1, meta); err != nil {
		return "", err
	}
	s3.url.Path = path.Join(s3.bucket, name)
	return s3.url.String(), nil
}

func (s3 *S3) HealthCheck() error {
	client, err := s3.client()

	if err != nil {
		return err
	}
	ok, err := client.BucketExists(s3.bucket)

	if err != nil {
		return err
	}

	if !ok {
		return fmt.Errorf("bucket unavailable: %s", s3.bucket)
	}
	return nil
}

func (s3 *S3) Metrics() (count int64, size int64, err error) {
	client, err := s3.client()

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

	for obj := range client.ListObjects(s3.bucket, "", false, make(chan struct{})) {
		count += 1
		size += obj.Size
	}
	return count, size, err
}

func (s3 *S3) client() (client *minio.Client, err error) {
	if s3.minio != nil {
		return s3.minio, nil
	}
	s3.url, err = url.Parse(s3.endpoint)

	if err != nil {
		return nil, err
	}
	client, err = minio.New(s3.url.Host, s3.accessKey, s3.secretKey, s3.url.Scheme == "https")

	if err != nil {
		return nil, err
	}
	ok, err := client.BucketExists(s3.bucket)

	if err != nil {
		return nil, err
	}

	if !ok {
		if err := client.MakeBucket(s3.bucket, s3.region); err != nil {
			return nil, err
		}
	}

	if err := client.SetBucketPolicy(s3.bucket, "", policy.BucketPolicyReadOnly); err != nil {
		return nil, err
	}

	s3.minio = client
	return client, nil
}

type RequestHandler struct {
	s3         *S3
	metrics    Metrics
	fileServer http.Handler
}

func NewRequestHandler(rootDir string, s3 *S3, metrics Metrics) *RequestHandler {
	return &RequestHandler{
		s3:         s3,
		metrics:    metrics,
		fileServer: http.FileServer(http.Dir(rootDir)),
	}
}

func (rh *RequestHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method == "GET" {
		rh.fileServer.ServeHTTP(w, r)
		return
	}
	rh.metrics.Total.Inc(1)

	if r.Method != "POST" && r.Method != "PUT" {
		rh.metrics.Invalid.Inc(1)
		w.Header().Set("Allow", "GET, POST, PUT")
		http.Error(w, "method not allowed", 405)
		return
	}
	objectUrl, err := rh.s3.Put(uuid.New(), r.Header.Get("Content-Type"), r.Body)

	if err != nil {
		rh.metrics.Failed.Inc(1)
		fmt.Println("could not store object:", err)
		http.Error(w, "internal server error", 500)
		return
	}

	if _, err := fmt.Fprintln(w, objectUrl); err != nil {
		rh.metrics.Failed.Inc(1)
		fmt.Println("could respond with object url:", err)
		return
	}
	rh.metrics.Success.Inc(1)
}

type HealthHandler struct {
	s3 *S3
}

type HealthPayload struct {
	Status string `json:"status"`
}

func NewHealthHandler(s3 *S3) *HealthHandler {
	return &HealthHandler{s3: s3}
}

func (hh *HealthHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	payload := &HealthPayload{Status: "ok"}

	if err := hh.s3.HealthCheck(); err != nil {
		payload.Status = fmt.Sprintf("critical: %s", err)
	}
	w.Header().Set("Content-Type", "application/json")

	if err := json.NewEncoder(w).Encode(payload); err != nil {
		fmt.Println("could not respond with health payload:", err)
	}
}

type Metrics struct {
	Total   metrics.Counter
	Invalid metrics.Counter
	Failed  metrics.Counter
	Success metrics.Counter
}

func NewMetrics() Metrics {
	return Metrics{
		Total:   metrics.NewCounter(),
		Invalid: metrics.NewCounter(),
		Failed:  metrics.NewCounter(),
		Success: metrics.NewCounter(),
	}
}

type MetricsHandler struct {
	metrics Metrics
	s3      *S3
}

type MetricsPayload struct {
	RequestsTotal   int64 `json:"requests_total"`
	RequestsInvalid int64 `json:"requests_invalid"`
	RequestsFailed  int64 `json:"requests_failed"`
	RequestsSuccess int64 `json:"requests_success"`
	S3Objects       int64 `json:"s3_objects"`
	S3Usage         int64 `json:"s3_usage"`
}

func NewMetricsHandler(metrics Metrics, s3 *S3) *MetricsHandler {
	return &MetricsHandler{metrics: metrics, s3: s3}
}

func (mh *MetricsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	s3Objects, s3Usage, _ := mh.s3.Metrics()

	payload := MetricsPayload{
		RequestsTotal:   mh.metrics.Total.Count(),
		RequestsInvalid: mh.metrics.Invalid.Count(),
		RequestsFailed:  mh.metrics.Failed.Count(),
		RequestsSuccess: mh.metrics.Success.Count(),
		S3Objects:       s3Objects,
		S3Usage:         s3Usage,
	}
	w.Header().Set("Content-Type", "application/json")

	if err := json.NewEncoder(w).Encode(payload); err != nil {
		fmt.Println("could not respond with metrics payload:", err)
	}
}
<commit_msg>minio lib upgrade<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"github.com/minio/minio-go"
	"github.com/minio/minio-go/pkg/policy"
	"github.com/pborman/uuid"
	"github.com/rcrowley/go-metrics"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"path"
	"syscall"
)

var (
	rootDir    string
	listenAddr string
	s3         S3
)

func init() {
	flag.StringVar(&rootDir, "root", "", "http root directory (default: .)")
	flag.StringVar(&listenAddr, "listen", "", "address to listen to (default: :8080)")
	flag.StringVar(&s3.endpoint, "s3-endpoint", "", "s3 endpoint")
	flag.StringVar(&s3.region, "s3-region", "", "s3 region")
	flag.StringVar(&s3.bucket, "s3-bucket", "", "s3 bucket")
	flag.StringVar(&s3.accessKey, "s3-access-key", "", "s3 access key")
	flag.StringVar(&s3.secretKey, "s3-secret-key", "", "s3 secret key")
}

func main() {
	flag.Parse()

	if len(rootDir) == 0 {
		listenAddr = os.Getenv("PUT_ROOT")
	}

	if len(rootDir) == 0 {
		rootDir = "."
	}

	if len(listenAddr) == 0 {
		listenAddr = os.Getenv("PUT_LISTEN")
	}

	if len(listenAddr) == 0 {
		listenAddr = ":8080"
	}

	if len(s3.endpoint) == 0 {
		s3.endpoint = os.Getenv("PUT_S3_ENDPOINT")
	}

	if len(s3.region) == 0 {
		s3.region = os.Getenv("PUT_S3_REGION")
	}

	if len(s3.bucket) == 0 {
		s3.bucket = os.Getenv("PUT_S3_BUCKET")
	}

	if len(s3.accessKey) == 0 {
		s3.accessKey = os.Getenv("PUT_S3_ACCESS_KEY")
	}

	if len(s3.secretKey) == 0 {
		s3.secretKey = os.Getenv("PUT_S3_SECRET_KEY")
	}
	l, err := listen(listenAddr)

	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	defer l.Close()

	metrics := NewMetrics()

	mux := http.NewServeMux()
	mux.Handle("/", NewRequestHandler(rootDir, &s3, metrics))
	mux.Handle("/health", NewHealthHandler(&s3))
	mux.Handle("/metrics", NewMetricsHandler(metrics, &s3))

	go http.Serve(l, mux)

	sig := make(chan os.Signal)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM, syscall.SIGKILL)
	<-sig
}

func listen(addr string) (net.Listener, error) {
	if len(addr) > 0 && addr[0] != '/' {
		return net.Listen("tcp", addr)
	}
	return net.Listen("unix", addr)
}

type S3 struct {
	minio     *minio.Client
	url       *url.URL
	endpoint  string
	bucket    string
	region    string
	secretKey string
	accessKey string
}

func (s3 *S3) Put(name, contentType string, body io.Reader) (string, error) {
	client, err := s3.client()

	if err != nil {
		return "", err
	}
	meta := minio.PutObjectOptions{ContentType: contentType}

	if _, err := client.PutObject(s3.bucket, name, body, -1, meta); err != nil {
		return "", err
	}
	s3.url.Path = path.Join(s3.bucket, name)
	return s3.url.String(), nil
}

func (s3 *S3) HealthCheck() error {
	client, err := s3.client()

	if err != nil {
		return err
	}
	ok, err := client.BucketExists(s3.bucket)

	if err != nil {
		return err
	}

	if !ok {
		return fmt.Errorf("bucket unavailable: %s", s3.bucket)
	}
	return nil
}

func (s3 *S3) Metrics() (count int64, size int64, err error) {
	client, err := s3.client()

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

	for obj := range client.ListObjects(s3.bucket, "", false, make(chan struct{})) {
		count += 1
		size += obj.Size
	}
	return count, size, err
}

func (s3 *S3) client() (client *minio.Client, err error) {
	if s3.minio != nil {
		return s3.minio, nil
	}
	s3.url, err = url.Parse(s3.endpoint)

	if err != nil {
		return nil, err
	}
	client, err = minio.New(s3.url.Host, s3.accessKey, s3.secretKey, s3.url.Scheme == "https")

	if err != nil {
		return nil, err
	}
	ok, err := client.BucketExists(s3.bucket)

	if err != nil {
		return nil, err
	}

	if !ok {
		if err := client.MakeBucket(s3.bucket, s3.region); err != nil {
			return nil, err
		}
	}

	if err := client.SetBucketPolicy(s3.bucket, "", policy.BucketPolicyReadOnly); err != nil {
		return nil, err
	}

	s3.minio = client
	return client, nil
}

type RequestHandler struct {
	s3         *S3
	metrics    Metrics
	fileServer http.Handler
}

func NewRequestHandler(rootDir string, s3 *S3, metrics Metrics) *RequestHandler {
	return &RequestHandler{
		s3:         s3,
		metrics:    metrics,
		fileServer: http.FileServer(http.Dir(rootDir)),
	}
}

func (rh *RequestHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method == "GET" {
		rh.fileServer.ServeHTTP(w, r)
		return
	}
	rh.metrics.Total.Inc(1)

	if r.Method != "POST" && r.Method != "PUT" {
		rh.metrics.Invalid.Inc(1)
		w.Header().Set("Allow", "GET, POST, PUT")
		http.Error(w, "method not allowed", 405)
		return
	}
	objectUrl, err := rh.s3.Put(uuid.New(), r.Header.Get("Content-Type"), r.Body)

	if err != nil {
		rh.metrics.Failed.Inc(1)
		fmt.Println("could not store object:", err)
		http.Error(w, "internal server error", 500)
		return
	}

	if _, err := fmt.Fprintln(w, objectUrl); err != nil {
		rh.metrics.Failed.Inc(1)
		fmt.Println("could respond with object url:", err)
		return
	}
	rh.metrics.Success.Inc(1)
}

type HealthHandler struct {
	s3 *S3
}

type HealthPayload struct {
	Status string `json:"status"`
}

func NewHealthHandler(s3 *S3) *HealthHandler {
	return &HealthHandler{s3: s3}
}

func (hh *HealthHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	payload := &HealthPayload{Status: "ok"}

	if err := hh.s3.HealthCheck(); err != nil {
		payload.Status = fmt.Sprintf("critical: %s", err)
	}
	w.Header().Set("Content-Type", "application/json")

	if err := json.NewEncoder(w).Encode(payload); err != nil {
		fmt.Println("could not respond with health payload:", err)
	}
}

type Metrics struct {
	Total   metrics.Counter
	Invalid metrics.Counter
	Failed  metrics.Counter
	Success metrics.Counter
}

func NewMetrics() Metrics {
	return Metrics{
		Total:   metrics.NewCounter(),
		Invalid: metrics.NewCounter(),
		Failed:  metrics.NewCounter(),
		Success: metrics.NewCounter(),
	}
}

type MetricsHandler struct {
	metrics Metrics
	s3      *S3
}

type MetricsPayload struct {
	RequestsTotal   int64 `json:"requests_total"`
	RequestsInvalid int64 `json:"requests_invalid"`
	RequestsFailed  int64 `json:"requests_failed"`
	RequestsSuccess int64 `json:"requests_success"`
	S3Objects       int64 `json:"s3_objects"`
	S3Usage         int64 `json:"s3_usage"`
}

func NewMetricsHandler(metrics Metrics, s3 *S3) *MetricsHandler {
	return &MetricsHandler{metrics: metrics, s3: s3}
}

func (mh *MetricsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	s3Objects, s3Usage, _ := mh.s3.Metrics()

	payload := MetricsPayload{
		RequestsTotal:   mh.metrics.Total.Count(),
		RequestsInvalid: mh.metrics.Invalid.Count(),
		RequestsFailed:  mh.metrics.Failed.Count(),
		RequestsSuccess: mh.metrics.Success.Count(),
		S3Objects:       s3Objects,
		S3Usage:         s3Usage,
	}
	w.Header().Set("Content-Type", "application/json")

	if err := json.NewEncoder(w).Encode(payload); err != nil {
		fmt.Println("could not respond with metrics payload:", err)
	}
}
<|endoftext|>


<commit_before>package lfs

import (
	"bytes"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/github/git-lfs/vendor/_nuts/github.com/ThomsonReutersEikon/go-ntlm/ntlm"
)

func (c *Configuration) ntlmClientSession(creds Creds) (ntlm.ClientSession, error) {
	if c.ntlmSession != nil {
		return c.ntlmSession, nil
	}
	splits := strings.Split(creds["username"], "\\")

	if len(splits) != 2 {
		errorMessage := fmt.Sprintf("Your user name must be of the form DOMAIN\\user. It is currently %s", creds["username"], "string")
		return nil, errors.New(errorMessage)
	}

	session, err := ntlm.CreateClientSession(ntlm.Version2, ntlm.ConnectionOrientedMode)

	if err != nil {
		return nil, err
	}

	session.SetUserInfo(splits[1], creds["password"], strings.ToUpper(splits[0]))
	c.ntlmSession = session
	return session, nil
}

func DoNTLMRequest(request *http.Request, retry bool) (*http.Response, error) {
	handReq, err := cloneRequest(request)
	if err != nil {
		return nil, err
	}

	res, err := Config.HttpClient().Do(handReq)
	if err != nil && res == nil {
		return nil, err
	}

	//If the status is 401 then we need to re-authenticate, otherwise it was successful
	if res.StatusCode == 401 {

		creds, err := getCredsForAPI(request)
		if err != nil {
			return nil, err
		}

		negotiateReq, err := cloneRequest(request)
		if err != nil {
			return nil, err
		}

		challengeMessage, err := negotiate(negotiateReq, getNegotiateMessage())
		if err != nil {
			return nil, err
		}

		challengeReq, err := cloneRequest(request)
		if err != nil {
			return nil, err
		}

		res, err := challenge(challengeReq, challengeMessage, creds)
		if err != nil {
			return nil, err
		}

		//If the status is 401 then we need to re-authenticate
		if res.StatusCode == 401 && retry == true {
			return DoNTLMRequest(challengeReq, false)
		}

		saveCredentials(creds, res)

		return res, nil
	}
	return res, nil
}

func negotiate(request *http.Request, message string) ([]byte, error) {
	request.Header.Add("Authorization", message)
	res, err := Config.HttpClient().Do(request)

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

	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()

	ret, err := parseChallengeResponse(res)
	if err != nil {
		return nil, err
	}

	return ret, nil
}

func challenge(request *http.Request, challengeBytes []byte, creds Creds) (*http.Response, error) {
	challenge, err := ntlm.ParseChallengeMessage(challengeBytes)
	if err != nil {
		return nil, err
	}

	session, err := Config.ntlmClientSession(creds)
	if err != nil {
		return nil, err
	}

	session.ProcessChallengeMessage(challenge)
	authenticate, err := session.GenerateAuthenticateMessage()
	if err != nil {
		return nil, err
	}

	authMsg := base64.StdEncoding.EncodeToString(authenticate.Bytes())
	request.Header.Add("Authorization", "NTLM "+authMsg)
	return Config.HttpClient().Do(request)
}

func parseChallengeResponse(response *http.Response) ([]byte, error) {
	header := response.Header.Get("Www-Authenticate")
	if len(header) < 6 {
		return nil, fmt.Errorf("Invalid NTLM challenge response: %q", header)
	}

	//parse out the "NTLM " at the beginning of the response
	challenge := header[5:]
	val, err := base64.StdEncoding.DecodeString(challenge)

	if err != nil {
		return nil, err
	}
	return []byte(val), nil
}

func cloneRequest(request *http.Request) (*http.Request, error) {
	var clonedReq *http.Request
	var err error

	if request.Body != nil {
		//If we have a body (POST/PUT etc.)
		//We need to do some magic to copy the request without closing the body stream

		buf, err := ioutil.ReadAll(request.Body)

		if err != nil {
			return nil, err
		}

		cloneReqBody := ioutil.NopCloser(bytes.NewBuffer(buf))
		request.Body = ioutil.NopCloser(bytes.NewBuffer(buf))
		clonedReq, err = http.NewRequest(request.Method, request.URL.String(), cloneReqBody)
		if err != nil {
			return nil, err
		}

	} else {
		clonedReq, err = http.NewRequest(request.Method, request.URL.String(), nil)
		if err != nil {
			return nil, err
		}
	}

	for k, v := range request.Header {
		clonedReq.Header.Add(k, v[0])
	}

	clonedReq.ContentLength = request.ContentLength

	return clonedReq, nil
}

func getNegotiateMessage() string {
	return "NTLM TlRMTVNTUAABAAAAB7IIogwADAAzAAAACwALACgAAAAKAAAoAAAAD1dJTExISS1NQUlOTk9SVEhBTUVSSUNB"
}

<commit_msg>ラララララ ラー ウウウ フフフ<commit_after>package lfs

import (
	"bytes"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/github/git-lfs/vendor/_nuts/github.com/ThomsonReutersEikon/go-ntlm/ntlm"
)

func (c *Configuration) ntlmClientSession(creds Creds) (ntlm.ClientSession, error) {
	if c.ntlmSession != nil {
		return c.ntlmSession, nil
	}
	splits := strings.Split(creds["username"], "\\")

	if len(splits) != 2 {
		errorMessage := fmt.Sprintf("Your user name must be of the form DOMAIN\\user. It is currently %s", creds["username"], "string")
		return nil, errors.New(errorMessage)
	}

	session, err := ntlm.CreateClientSession(ntlm.Version2, ntlm.ConnectionOrientedMode)

	if err != nil {
		return nil, err
	}

	session.SetUserInfo(splits[1], creds["password"], strings.ToUpper(splits[0]))
	c.ntlmSession = session
	return session, nil
}

func DoNTLMRequest(request *http.Request, retry bool) (*http.Response, error) {
	handReq, err := cloneRequest(request)
	if err != nil {
		return nil, err
	}

	res, err := Config.HttpClient().Do(handReq)
	if err != nil && res == nil {
		return nil, err
	}

	//If the status is 401 then we need to re-authenticate, otherwise it was successful
	if res.StatusCode == 401 {

		creds, err := getCredsForAPI(request)
		if err != nil {
			return nil, err
		}

		negotiateReq, err := cloneRequest(request)
		if err != nil {
			return nil, err
		}

		challengeMessage, err := negotiate(negotiateReq, ntlmNegotiateMessage)
		if err != nil {
			return nil, err
		}

		challengeReq, err := cloneRequest(request)
		if err != nil {
			return nil, err
		}

		res, err := challenge(challengeReq, challengeMessage, creds)
		if err != nil {
			return nil, err
		}

		//If the status is 401 then we need to re-authenticate
		if res.StatusCode == 401 && retry == true {
			return DoNTLMRequest(challengeReq, false)
		}

		saveCredentials(creds, res)

		return res, nil
	}
	return res, nil
}

func negotiate(request *http.Request, message string) ([]byte, error) {
	request.Header.Add("Authorization", message)
	res, err := Config.HttpClient().Do(request)

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

	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()

	ret, err := parseChallengeResponse(res)
	if err != nil {
		return nil, err
	}

	return ret, nil
}

func challenge(request *http.Request, challengeBytes []byte, creds Creds) (*http.Response, error) {
	challenge, err := ntlm.ParseChallengeMessage(challengeBytes)
	if err != nil {
		return nil, err
	}

	session, err := Config.ntlmClientSession(creds)
	if err != nil {
		return nil, err
	}

	session.ProcessChallengeMessage(challenge)
	authenticate, err := session.GenerateAuthenticateMessage()
	if err != nil {
		return nil, err
	}

	authMsg := base64.StdEncoding.EncodeToString(authenticate.Bytes())
	request.Header.Add("Authorization", "NTLM "+authMsg)
	return Config.HttpClient().Do(request)
}

func parseChallengeResponse(response *http.Response) ([]byte, error) {
	header := response.Header.Get("Www-Authenticate")
	if len(header) < 6 {
		return nil, fmt.Errorf("Invalid NTLM challenge response: %q", header)
	}

	//parse out the "NTLM " at the beginning of the response
	challenge := header[5:]
	val, err := base64.StdEncoding.DecodeString(challenge)

	if err != nil {
		return nil, err
	}
	return []byte(val), nil
}

func cloneRequest(request *http.Request) (*http.Request, error) {
	var clonedReq *http.Request
	var err error

	if request.Body != nil {
		//If we have a body (POST/PUT etc.)
		//We need to do some magic to copy the request without closing the body stream

		buf, err := ioutil.ReadAll(request.Body)

		if err != nil {
			return nil, err
		}

		cloneReqBody := ioutil.NopCloser(bytes.NewBuffer(buf))
		request.Body = ioutil.NopCloser(bytes.NewBuffer(buf))
		clonedReq, err = http.NewRequest(request.Method, request.URL.String(), cloneReqBody)
		if err != nil {
			return nil, err
		}

	} else {
		clonedReq, err = http.NewRequest(request.Method, request.URL.String(), nil)
		if err != nil {
			return nil, err
		}
	}

	for k, v := range request.Header {
		clonedReq.Header.Add(k, v[0])
	}

	clonedReq.ContentLength = request.ContentLength

	return clonedReq, nil
}

const ntlmNegotiateMessage = "NTLM TlRMTVNTUAABAAAAB7IIogwADAAzAAAACwALACgAAAAKAAAoAAAAD1dJTExISS1NQUlOTk9SVEhBTUVSSUNB"
<|endoftext|>