<commit_before>package resolver

import (
	"crypto/rand"
	"fmt"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"nimona.io/pkg/context"
	"nimona.io/pkg/crypto"
	"nimona.io/pkg/localpeer"
	"nimona.io/pkg/network"
	"nimona.io/pkg/object"
	"nimona.io/pkg/peer"
)

func TestResolver_TwoPeersCanFindEachOther(t *testing.T) {
	_, k0, kc0, n0, ctx0 := newPeer(t, "peer0")

	d0 := New(
		ctx0,
		n0,
	)

	ba := []*peer.Peer{
		{
			Addresses: n0.LocalPeer().GetAddresses(),
			Metadata: object.Metadata{
				Owner: kc0.GetPrimaryPeerKey().PublicKey(),
			},
		},
	}

	time.Sleep(time.Millisecond * 100)

	_, k1, _, n1, ctx1 := newPeer(t, "peer1")

	d1 := New(
		ctx1,
		n1,
		WithBoostrapPeers(ba),
	)

	time.Sleep(time.Millisecond * 100)

	ctx := context.New(
		context.WithCorrelationID("req1"),
		context.WithTimeout(time.Second),
	)

	peersChan, err := d1.Lookup(ctx, LookupByOwner(k0.PublicKey()))

	peers := gatherPeers(peersChan)
	require.NoError(t, err)
	require.Len(t, peers, 1)
	require.ElementsMatch(t, n0.LocalPeer().GetAddresses(), peers[0].Addresses)

	ctxR2 := context.New(
		context.WithCorrelationID("req2"),
		context.WithTimeout(time.Second),
	)
	peersChan, err = d0.Lookup(ctxR2, LookupByOwner(k1.PublicKey()))
	peers = gatherPeers(peersChan)
	require.NoError(t, err)
	require.Len(t, peers, 1)
	require.Equal(t, n1.LocalPeer().GetAddresses(), peers[0].Addresses)
}

func TestResolver_TwoPeersAndOneBootstrapCanFindEachOther(t *testing.T) {
	_, k0, kc0, n0, ctx0 := newPeer(t, "peer0")

	// bootstrap node
	New(ctx0, n0)

	time.Sleep(time.Millisecond * 100)

	_, k1, _, n1, ctx1 := newPeer(t, "peer1")
	_, k2, _, n2, ctx2 := newPeer(t, "peer2")

	// bootstrap address
	ba := []*peer.Peer{
		{
			Addresses: n0.LocalPeer().GetAddresses(),
			Metadata: object.Metadata{
				Owner: kc0.GetPrimaryPeerKey().PublicKey(),
			},
		},
	}

	// node 1
	d1 := New(
		ctx1,
		n1,
		WithBoostrapPeers(ba),
	)

	time.Sleep(time.Millisecond * 100)

	// node 2
	d2 := New(
		ctx2,
		n2,
		WithBoostrapPeers(ba),
	)

	time.Sleep(time.Millisecond * 100)

	// find bootstrap from node1
	ctx := context.New(
		context.WithCorrelationID("req1"),
		context.WithTimeout(time.Second*2),
	)
	peersChan, err := d1.Lookup(ctx, LookupByOwner(k0.PublicKey()))
	peers := gatherPeers(peersChan)
	require.NoError(t, err)
	require.Len(t, peers, 1)
	require.ElementsMatch(t, n0.LocalPeer().GetAddresses(), peers[0].Addresses)

	// find node 1 from node 2
	ctx = context.New(
		context.WithCorrelationID("req2"),
		context.WithTimeout(time.Second*2),
	)
	peersChan, err = d2.Lookup(ctx, LookupByOwner(k1.PublicKey()))
	peers = gatherPeers(peersChan)
	require.NoError(t, err)
	require.Len(t, peers, 1)
	require.ElementsMatch(t, n1.LocalPeer().GetAddresses(), peers[0].Addresses)

	// find node 2 from node 1
	ctx = context.New(
		context.WithCorrelationID("req3"),
		context.WithTimeout(time.Second*2),
	)

	peersChan, err = d1.Lookup(ctx, LookupByOwner(k2.PublicKey()))
	peers = gatherPeers(peersChan)
	require.NoError(t, err)
	require.Len(t, peers, 1)
	require.ElementsMatch(t, n2.LocalPeer().GetAddresses(), peers[0].Addresses)
}

func TestResolver_TwoPeersAndOneBootstrapCanProvide(t *testing.T) {
	_, k0, kc0, n0, ctx0 := newPeer(t, "peer0")
	_, k1, kc1, n1, ctx1 := newPeer(t, "peer1")
	_, k2, _, n2, ctx2 := newPeer(t, "peer2")

	// make peer 1 a provider
	token := make([]byte, 32)
	rand.Read(token) // nolint: errcheck
	ch := object.Hash("foo")
	kc1.PutContentHashes(ch)

	// print peer info
	fmt.Println("0", k0.PublicKey(), n0.LocalPeer().GetAddresses())
	fmt.Println("1", k1.PublicKey(), n1.LocalPeer().GetAddresses())
	fmt.Println("2", k2.PublicKey(), n2.LocalPeer().GetAddresses())

	// bootstrap peer
	d0 := New(
		ctx0,
		n0,
	)

	time.Sleep(time.Millisecond * 100)

	// bootstrap address
	ba := []*peer.Peer{
		{
			Addresses: n0.LocalPeer().GetAddresses(),
			Metadata: object.Metadata{
				Owner: kc0.GetPrimaryPeerKey().PublicKey(),
			},
		},
	}

	// peer 1
	New(
		ctx1,
		n1,
		WithBoostrapPeers(ba),
	)

	time.Sleep(time.Millisecond * 100)

	// peer 2
	d2 := New(
		ctx2,
		n2,
		WithBoostrapPeers(ba),
	)

	time.Sleep(time.Millisecond * 100)

	// find peer 1 from peer 2
	ctx := context.New(
		context.WithCorrelationID("req1"),
		context.WithTimeout(time.Second),
	)
	providersChan, err := d2.Lookup(ctx, LookupByContentHash(ch))
	require.NoError(t, err)
	providers := gatherPeers(providersChan)
	require.Len(t, providers, 1)
	require.Equal(t, k1.PublicKey(), providers[0].PublicKey())

	// find peer 1 from bootstrap
	ctx = context.New(
		context.WithCorrelationID("req2"),
		context.WithTimeout(time.Second*2),
	)
	providersChan, err = d0.Lookup(ctx, LookupByContentHash(ch))
	require.NoError(t, err)
	providers = gatherPeers(providersChan)
	require.Len(t, providers, 1)
	require.Equal(t, k1.PublicKey(), providers[0].PublicKey())
}

// nolint: gocritic
func newPeer(
	t *testing.T,
	name string,
) (
	crypto.PrivateKey,
	crypto.PrivateKey,
	localpeer.LocalPeer,
	network.Network,
	context.Context,
) {
	ctx := context.New(context.WithCorrelationID(name))

	// identity key
	opk, err := crypto.GenerateEd25519PrivateKey()
	assert.NoError(t, err)

	// peer key
	pk, err := crypto.GenerateEd25519PrivateKey()
	assert.NoError(t, err)

	// peer certificate
	c, err := object.NewCertificate(
		pk.PublicKey(),
		opk,
	)
	require.NoError(t, err)

	kc := localpeer.New()
	kc.PutPrimaryPeerKey(pk)
	kc.PutPrimaryIdentityKey(opk)
	kc.PutCertificate(c)

	n := network.New(
		ctx,
		network.WithLocalPeer(kc),
	)

	_, err = n.Listen(ctx, "127.0.0.1:0", network.ListenOnLocalIPs)
	require.NoError(t, err)

	return opk, pk, kc, n, ctx
}

func gatherPeers(p <-chan *peer.Peer) []*peer.Peer {
	ps := []*peer.Peer{}
	for p := range p {
		p := p
		ps = append(ps, p)
	}
	return peer.Unique(ps)
}
<commit_msg>fix(resolver): fix flaky TwoPeersCanFindEachOther test<commit_after>package resolver

import (
	"crypto/rand"
	"fmt"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"nimona.io/pkg/context"
	"nimona.io/pkg/crypto"
	"nimona.io/pkg/localpeer"
	"nimona.io/pkg/network"
	"nimona.io/pkg/object"
	"nimona.io/pkg/peer"
)

func TestResolver_TwoPeersCanFindEachOther(t *testing.T) {
	_, k0, kc0, n0, ctx0 := newPeer(t, "peer0")

	d0 := New(
		ctx0,
		n0,
	)

	ba := []*peer.Peer{
		{
			Addresses: n0.LocalPeer().GetAddresses(),
			Metadata: object.Metadata{
				Owner: kc0.GetPrimaryPeerKey().PublicKey(),
			},
		},
	}

	time.Sleep(time.Millisecond * 100)

	_, k1, _, n1, ctx1 := newPeer(t, "peer1")

	d1 := New(
		ctx1,
		n1,
		WithBoostrapPeers(ba),
	)

	time.Sleep(time.Millisecond * 100)

	ctx := context.New(
		context.WithCorrelationID("req1"),
		context.WithTimeout(time.Second),
	)

	peersChan, err := d1.Lookup(ctx, LookupByOwner(k0.PublicKey()))

	peers := gatherPeers(peersChan)
	require.NoError(t, err)
	require.Len(t, peers, 1)
	require.ElementsMatch(t, n0.LocalPeer().GetAddresses(), peers[0].Addresses)

	ctxR2 := context.New(
		context.WithCorrelationID("req2"),
		context.WithTimeout(time.Second),
	)
	peersChan, err = d0.Lookup(ctxR2, LookupByOwner(k1.PublicKey()))
	peers = gatherPeers(peersChan)
	require.NoError(t, err)
	require.Len(t, peers, 1)
	require.ElementsMatch(t, n1.LocalPeer().GetAddresses(), peers[0].Addresses)
}

func TestResolver_TwoPeersAndOneBootstrapCanFindEachOther(t *testing.T) {
	_, k0, kc0, n0, ctx0 := newPeer(t, "peer0")

	// bootstrap node
	New(ctx0, n0)

	time.Sleep(time.Millisecond * 100)

	_, k1, _, n1, ctx1 := newPeer(t, "peer1")
	_, k2, _, n2, ctx2 := newPeer(t, "peer2")

	// bootstrap address
	ba := []*peer.Peer{
		{
			Addresses: n0.LocalPeer().GetAddresses(),
			Metadata: object.Metadata{
				Owner: kc0.GetPrimaryPeerKey().PublicKey(),
			},
		},
	}

	// node 1
	d1 := New(
		ctx1,
		n1,
		WithBoostrapPeers(ba),
	)

	time.Sleep(time.Millisecond * 100)

	// node 2
	d2 := New(
		ctx2,
		n2,
		WithBoostrapPeers(ba),
	)

	time.Sleep(time.Millisecond * 100)

	// find bootstrap from node1
	ctx := context.New(
		context.WithCorrelationID("req1"),
		context.WithTimeout(time.Second*2),
	)
	peersChan, err := d1.Lookup(ctx, LookupByOwner(k0.PublicKey()))
	peers := gatherPeers(peersChan)
	require.NoError(t, err)
	require.Len(t, peers, 1)
	require.ElementsMatch(t, n0.LocalPeer().GetAddresses(), peers[0].Addresses)

	// find node 1 from node 2
	ctx = context.New(
		context.WithCorrelationID("req2"),
		context.WithTimeout(time.Second*2),
	)
	peersChan, err = d2.Lookup(ctx, LookupByOwner(k1.PublicKey()))
	peers = gatherPeers(peersChan)
	require.NoError(t, err)
	require.Len(t, peers, 1)
	require.ElementsMatch(t, n1.LocalPeer().GetAddresses(), peers[0].Addresses)

	// find node 2 from node 1
	ctx = context.New(
		context.WithCorrelationID("req3"),
		context.WithTimeout(time.Second*2),
	)

	peersChan, err = d1.Lookup(ctx, LookupByOwner(k2.PublicKey()))
	peers = gatherPeers(peersChan)
	require.NoError(t, err)
	require.Len(t, peers, 1)
	require.ElementsMatch(t, n2.LocalPeer().GetAddresses(), peers[0].Addresses)
}

func TestResolver_TwoPeersAndOneBootstrapCanProvide(t *testing.T) {
	_, k0, kc0, n0, ctx0 := newPeer(t, "peer0")
	_, k1, kc1, n1, ctx1 := newPeer(t, "peer1")
	_, k2, _, n2, ctx2 := newPeer(t, "peer2")

	// make peer 1 a provider
	token := make([]byte, 32)
	rand.Read(token) // nolint: errcheck
	ch := object.Hash("foo")
	kc1.PutContentHashes(ch)

	// print peer info
	fmt.Println("0", k0.PublicKey(), n0.LocalPeer().GetAddresses())
	fmt.Println("1", k1.PublicKey(), n1.LocalPeer().GetAddresses())
	fmt.Println("2", k2.PublicKey(), n2.LocalPeer().GetAddresses())

	// bootstrap peer
	d0 := New(
		ctx0,
		n0,
	)

	time.Sleep(time.Millisecond * 100)

	// bootstrap address
	ba := []*peer.Peer{
		{
			Addresses: n0.LocalPeer().GetAddresses(),
			Metadata: object.Metadata{
				Owner: kc0.GetPrimaryPeerKey().PublicKey(),
			},
		},
	}

	// peer 1
	New(
		ctx1,
		n1,
		WithBoostrapPeers(ba),
	)

	time.Sleep(time.Millisecond * 100)

	// peer 2
	d2 := New(
		ctx2,
		n2,
		WithBoostrapPeers(ba),
	)

	time.Sleep(time.Millisecond * 100)

	// find peer 1 from peer 2
	ctx := context.New(
		context.WithCorrelationID("req1"),
		context.WithTimeout(time.Second),
	)
	providersChan, err := d2.Lookup(ctx, LookupByContentHash(ch))
	require.NoError(t, err)
	providers := gatherPeers(providersChan)
	require.Len(t, providers, 1)
	require.Equal(t, k1.PublicKey(), providers[0].PublicKey())

	// find peer 1 from bootstrap
	ctx = context.New(
		context.WithCorrelationID("req2"),
		context.WithTimeout(time.Second*2),
	)
	providersChan, err = d0.Lookup(ctx, LookupByContentHash(ch))
	require.NoError(t, err)
	providers = gatherPeers(providersChan)
	require.Len(t, providers, 1)
	require.Equal(t, k1.PublicKey(), providers[0].PublicKey())
}

// nolint: gocritic
func newPeer(
	t *testing.T,
	name string,
) (
	crypto.PrivateKey,
	crypto.PrivateKey,
	localpeer.LocalPeer,
	network.Network,
	context.Context,
) {
	ctx := context.New(context.WithCorrelationID(name))

	// identity key
	opk, err := crypto.GenerateEd25519PrivateKey()
	assert.NoError(t, err)

	// peer key
	pk, err := crypto.GenerateEd25519PrivateKey()
	assert.NoError(t, err)

	// peer certificate
	c, err := object.NewCertificate(
		pk.PublicKey(),
		opk,
	)
	require.NoError(t, err)

	kc := localpeer.New()
	kc.PutPrimaryPeerKey(pk)
	kc.PutPrimaryIdentityKey(opk)
	kc.PutCertificate(c)

	n := network.New(
		ctx,
		network.WithLocalPeer(kc),
	)

	_, err = n.Listen(ctx, "127.0.0.1:0", network.ListenOnLocalIPs)
	require.NoError(t, err)

	return opk, pk, kc, n, ctx
}

func gatherPeers(p <-chan *peer.Peer) []*peer.Peer {
	ps := []*peer.Peer{}
	for p := range p {
		p := p
		ps = append(ps, p)
	}
	return peer.Unique(ps)
}
<|endoftext|>


<commit_before>/*
Copyright 2017 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 api

import (
	"fmt"

	"k8s.io/api/core/v1"
	policyv1 "k8s.io/api/policy/v1beta1"
	"k8s.io/apimachinery/pkg/types"

	arbcorev1 "github.com/kubernetes-incubator/kube-arbitrator/pkg/apis/scheduling/v1alpha1"
	"github.com/kubernetes-incubator/kube-arbitrator/pkg/apis/utils"
)

type TaskID types.UID

type TaskInfo struct {
	UID TaskID
	Job JobID

	Name      string
	Namespace string

	Resreq *Resource

	NodeName string
	Status   TaskStatus
	Priority int32

	Pod *v1.Pod
}

func getJobID(pod *v1.Pod) JobID {
	if len(pod.Annotations) != 0 {
		if gn, found := pod.Annotations[arbcorev1.GroupNameAnnotationKey]; found && len(gn) != 0 {
			// Make sure Pod and PodGroup belong to the same namespace.
			jobID := fmt.Sprintf("%s/%s", pod.Namespace, gn)
			return JobID(jobID)
		}
	}
	return JobID(utils.GetController(pod))
}

func NewTaskInfo(pod *v1.Pod) *TaskInfo {
	req := EmptyResource()

	// TODO(k82cn): also includes initContainers' resource.
	for _, c := range pod.Spec.Containers {
		req.Add(NewResource(c.Resources.Requests))
	}

	ti := &TaskInfo{
		UID:       TaskID(pod.UID),
		Job:       getJobID(pod),
		Name:      pod.Name,
		Namespace: pod.Namespace,
		NodeName:  pod.Spec.NodeName,
		Status:    getTaskStatus(pod),
		Priority:  1,

		Pod:    pod,
		Resreq: req,
	}

	if pod.Spec.Priority != nil {
		ti.Priority = *pod.Spec.Priority
	}

	return ti
}

func (ti *TaskInfo) Clone() *TaskInfo {
	return &TaskInfo{
		UID:       ti.UID,
		Job:       ti.Job,
		Name:      ti.Name,
		Namespace: ti.Namespace,
		NodeName:  ti.NodeName,
		Status:    ti.Status,
		Priority:  ti.Priority,
		Pod:       ti.Pod,
		Resreq:    ti.Resreq.Clone(),
	}
}

func (ti TaskInfo) String() string {
	return fmt.Sprintf("Task (%v:%v/%v): job %v, status %v, pri %v, resreq %v",
		ti.UID, ti.Namespace, ti.Name, ti.Job, ti.Status, ti.Priority, ti.Resreq)
}

// JobID is the type of JobInfo's ID.
type JobID types.UID

type tasksMap map[TaskID]*TaskInfo

type JobInfo struct {
	UID JobID

	Name      string
	Namespace string

	Queue QueueID

	Priority int

	NodeSelector map[string]string
	MinAvailable int32

	// All tasks of the Job.
	TaskStatusIndex map[TaskStatus]tasksMap
	Tasks           tasksMap

	Allocated    *Resource
	TotalRequest *Resource

	PodGroup *arbcorev1.PodGroup

	// TODO(k82cn): keep backward compatbility, removed it when v1alpha1 finalized.
	PDB *policyv1.PodDisruptionBudget
}

func NewJobInfo(uid JobID) *JobInfo {
	return &JobInfo{
		UID: uid,

		MinAvailable: 0,
		NodeSelector: make(map[string]string),

		Allocated:    EmptyResource(),
		TotalRequest: EmptyResource(),

		TaskStatusIndex: map[TaskStatus]tasksMap{},
		Tasks:           tasksMap{},
	}
}

func (ji *JobInfo) UnsetPodGroup() {
	ji.PodGroup = nil
}

func (ji *JobInfo) SetPodGroup(pg *arbcorev1.PodGroup) {
	ji.Name = pg.Name
	ji.Namespace = pg.Namespace
	ji.MinAvailable = pg.Spec.NumMember

	if len(pg.Spec.Queue) == 0 {
		ji.Queue = QueueID(pg.Namespace)
	} else {
		ji.Queue = QueueID(pg.Spec.Queue)
	}

	ji.PodGroup = pg
}

func (ji *JobInfo) SetPDB(pdb *policyv1.PodDisruptionBudget) {
	ji.Name = pdb.Name
	ji.MinAvailable = pdb.Spec.MinAvailable.IntVal
	ji.Namespace = pdb.Namespace
	ji.Queue = QueueID(pdb.Namespace)

	ji.PDB = pbd
}

func (ji *JobInfo) UnsetPDB() {
	ji.PDB = nil
}

func (ji *JobInfo) GetTasks(statuses ...TaskStatus) []*TaskInfo {
	var res []*TaskInfo

	for _, status := range statuses {
		if tasks, found := ji.TaskStatusIndex[status]; found {
			for _, task := range tasks {
				res = append(res, task.Clone())
			}
		}
	}

	return res
}

func (ji *JobInfo) addTaskIndex(ti *TaskInfo) {
	if _, found := ji.TaskStatusIndex[ti.Status]; !found {
		ji.TaskStatusIndex[ti.Status] = tasksMap{}
	}

	ji.TaskStatusIndex[ti.Status][ti.UID] = ti
}

func (ji *JobInfo) AddTaskInfo(ti *TaskInfo) {
	ji.Tasks[ti.UID] = ti
	ji.addTaskIndex(ti)

	ji.TotalRequest.Add(ti.Resreq)

	if AllocatedStatus(ti.Status) {
		ji.Allocated.Add(ti.Resreq)
	}
}

func (ji *JobInfo) UpdateTaskStatus(task *TaskInfo, status TaskStatus) error {
	if err := validateStatusUpdate(task.Status, status); err != nil {
		return err
	}

	// Remove the task from the task list firstly
	ji.DeleteTaskInfo(task)

	// Update task's status to the target status
	task.Status = status
	ji.AddTaskInfo(task)

	return nil
}

func (ji *JobInfo) deleteTaskIndex(ti *TaskInfo) {
	if tasks, found := ji.TaskStatusIndex[ti.Status]; found {
		delete(tasks, ti.UID)

		if len(tasks) == 0 {
			delete(ji.TaskStatusIndex, ti.Status)
		}
	}
}

func (ji *JobInfo) DeleteTaskInfo(ti *TaskInfo) error {
	if task, found := ji.Tasks[ti.UID]; found {
		ji.TotalRequest.Sub(task.Resreq)

		if AllocatedStatus(task.Status) {
			ji.Allocated.Sub(task.Resreq)
		}

		delete(ji.Tasks, task.UID)

		ji.deleteTaskIndex(task)
		return nil
	}

	return fmt.Errorf("failed to find task <%v/%v> in job <%v/%v>",
		ti.Namespace, ti.Name, ji.Namespace, ji.Name)
}

func (ji *JobInfo) Clone() *JobInfo {
	info := &JobInfo{
		UID:       ji.UID,
		Name:      ji.Name,
		Namespace: ji.Namespace,
		Queue:     ji.Queue,

		MinAvailable: ji.MinAvailable,
		NodeSelector: map[string]string{},
		Allocated:    ji.Allocated.Clone(),
		TotalRequest: ji.TotalRequest.Clone(),

		PDB:      ji.PDB,
		PodGroup: ji.PodGroup,

		TaskStatusIndex: map[TaskStatus]tasksMap{},
		Tasks:           tasksMap{},
	}

	for k, v := range ji.NodeSelector {
		info.NodeSelector[k] = v
	}

	for _, task := range ji.Tasks {
		info.AddTaskInfo(task.Clone())
	}

	return info
}

func (ji JobInfo) String() string {
	res := ""

	i := 0
	for _, task := range ji.Tasks {
		res = res + fmt.Sprintf("\n\t %d: %v", i, task)
		i++
	}

	return fmt.Sprintf("Job (%v): name %v, minAvailable %d", ji.UID, ji.Name, ji.MinAvailable) + res
}
<commit_msg>fix typos<commit_after>/*
Copyright 2017 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 api

import (
	"fmt"

	"k8s.io/api/core/v1"
	policyv1 "k8s.io/api/policy/v1beta1"
	"k8s.io/apimachinery/pkg/types"

	arbcorev1 "github.com/kubernetes-incubator/kube-arbitrator/pkg/apis/scheduling/v1alpha1"
	"github.com/kubernetes-incubator/kube-arbitrator/pkg/apis/utils"
)

type TaskID types.UID

type TaskInfo struct {
	UID TaskID
	Job JobID

	Name      string
	Namespace string

	Resreq *Resource

	NodeName string
	Status   TaskStatus
	Priority int32

	Pod *v1.Pod
}

func getJobID(pod *v1.Pod) JobID {
	if len(pod.Annotations) != 0 {
		if gn, found := pod.Annotations[arbcorev1.GroupNameAnnotationKey]; found && len(gn) != 0 {
			// Make sure Pod and PodGroup belong to the same namespace.
			jobID := fmt.Sprintf("%s/%s", pod.Namespace, gn)
			return JobID(jobID)
		}
	}
	return JobID(utils.GetController(pod))
}

func NewTaskInfo(pod *v1.Pod) *TaskInfo {
	req := EmptyResource()

	// TODO(k82cn): also includes initContainers' resource.
	for _, c := range pod.Spec.Containers {
		req.Add(NewResource(c.Resources.Requests))
	}

	ti := &TaskInfo{
		UID:       TaskID(pod.UID),
		Job:       getJobID(pod),
		Name:      pod.Name,
		Namespace: pod.Namespace,
		NodeName:  pod.Spec.NodeName,
		Status:    getTaskStatus(pod),
		Priority:  1,

		Pod:    pod,
		Resreq: req,
	}

	if pod.Spec.Priority != nil {
		ti.Priority = *pod.Spec.Priority
	}

	return ti
}

func (ti *TaskInfo) Clone() *TaskInfo {
	return &TaskInfo{
		UID:       ti.UID,
		Job:       ti.Job,
		Name:      ti.Name,
		Namespace: ti.Namespace,
		NodeName:  ti.NodeName,
		Status:    ti.Status,
		Priority:  ti.Priority,
		Pod:       ti.Pod,
		Resreq:    ti.Resreq.Clone(),
	}
}

func (ti TaskInfo) String() string {
	return fmt.Sprintf("Task (%v:%v/%v): job %v, status %v, pri %v, resreq %v",
		ti.UID, ti.Namespace, ti.Name, ti.Job, ti.Status, ti.Priority, ti.Resreq)
}

// JobID is the type of JobInfo's ID.
type JobID types.UID

type tasksMap map[TaskID]*TaskInfo

type JobInfo struct {
	UID JobID

	Name      string
	Namespace string

	Queue QueueID

	Priority int

	NodeSelector map[string]string
	MinAvailable int32

	// All tasks of the Job.
	TaskStatusIndex map[TaskStatus]tasksMap
	Tasks           tasksMap

	Allocated    *Resource
	TotalRequest *Resource

	PodGroup *arbcorev1.PodGroup

	// TODO(k82cn): keep backward compatbility, removed it when v1alpha1 finalized.
	PDB *policyv1.PodDisruptionBudget
}

func NewJobInfo(uid JobID) *JobInfo {
	return &JobInfo{
		UID: uid,

		MinAvailable: 0,
		NodeSelector: make(map[string]string),

		Allocated:    EmptyResource(),
		TotalRequest: EmptyResource(),

		TaskStatusIndex: map[TaskStatus]tasksMap{},
		Tasks:           tasksMap{},
	}
}

func (ji *JobInfo) UnsetPodGroup() {
	ji.PodGroup = nil
}

func (ji *JobInfo) SetPodGroup(pg *arbcorev1.PodGroup) {
	ji.Name = pg.Name
	ji.Namespace = pg.Namespace
	ji.MinAvailable = pg.Spec.NumMember

	if len(pg.Spec.Queue) == 0 {
		ji.Queue = QueueID(pg.Namespace)
	} else {
		ji.Queue = QueueID(pg.Spec.Queue)
	}

	ji.PodGroup = pg
}

func (ji *JobInfo) SetPDB(pdb *policyv1.PodDisruptionBudget) {
	ji.Name = pdb.Name
	ji.MinAvailable = pdb.Spec.MinAvailable.IntVal
	ji.Namespace = pdb.Namespace
	ji.Queue = QueueID(pdb.Namespace)

	ji.PDB = pdb
}

func (ji *JobInfo) UnsetPDB() {
	ji.PDB = nil
}

func (ji *JobInfo) GetTasks(statuses ...TaskStatus) []*TaskInfo {
	var res []*TaskInfo

	for _, status := range statuses {
		if tasks, found := ji.TaskStatusIndex[status]; found {
			for _, task := range tasks {
				res = append(res, task.Clone())
			}
		}
	}

	return res
}

func (ji *JobInfo) addTaskIndex(ti *TaskInfo) {
	if _, found := ji.TaskStatusIndex[ti.Status]; !found {
		ji.TaskStatusIndex[ti.Status] = tasksMap{}
	}

	ji.TaskStatusIndex[ti.Status][ti.UID] = ti
}

func (ji *JobInfo) AddTaskInfo(ti *TaskInfo) {
	ji.Tasks[ti.UID] = ti
	ji.addTaskIndex(ti)

	ji.TotalRequest.Add(ti.Resreq)

	if AllocatedStatus(ti.Status) {
		ji.Allocated.Add(ti.Resreq)
	}
}

func (ji *JobInfo) UpdateTaskStatus(task *TaskInfo, status TaskStatus) error {
	if err := validateStatusUpdate(task.Status, status); err != nil {
		return err
	}

	// Remove the task from the task list firstly
	ji.DeleteTaskInfo(task)

	// Update task's status to the target status
	task.Status = status
	ji.AddTaskInfo(task)

	return nil
}

func (ji *JobInfo) deleteTaskIndex(ti *TaskInfo) {
	if tasks, found := ji.TaskStatusIndex[ti.Status]; found {
		delete(tasks, ti.UID)

		if len(tasks) == 0 {
			delete(ji.TaskStatusIndex, ti.Status)
		}
	}
}

func (ji *JobInfo) DeleteTaskInfo(ti *TaskInfo) error {
	if task, found := ji.Tasks[ti.UID]; found {
		ji.TotalRequest.Sub(task.Resreq)

		if AllocatedStatus(task.Status) {
			ji.Allocated.Sub(task.Resreq)
		}

		delete(ji.Tasks, task.UID)

		ji.deleteTaskIndex(task)
		return nil
	}

	return fmt.Errorf("failed to find task <%v/%v> in job <%v/%v>",
		ti.Namespace, ti.Name, ji.Namespace, ji.Name)
}

func (ji *JobInfo) Clone() *JobInfo {
	info := &JobInfo{
		UID:       ji.UID,
		Name:      ji.Name,
		Namespace: ji.Namespace,
		Queue:     ji.Queue,

		MinAvailable: ji.MinAvailable,
		NodeSelector: map[string]string{},
		Allocated:    ji.Allocated.Clone(),
		TotalRequest: ji.TotalRequest.Clone(),

		PDB:      ji.PDB,
		PodGroup: ji.PodGroup,

		TaskStatusIndex: map[TaskStatus]tasksMap{},
		Tasks:           tasksMap{},
	}

	for k, v := range ji.NodeSelector {
		info.NodeSelector[k] = v
	}

	for _, task := range ji.Tasks {
		info.AddTaskInfo(task.Clone())
	}

	return info
}

func (ji JobInfo) String() string {
	res := ""

	i := 0
	for _, task := range ji.Tasks {
		res = res + fmt.Sprintf("\n\t %d: %v", i, task)
		i++
	}

	return fmt.Sprintf("Job (%v): name %v, minAvailable %d", ji.UID, ji.Name, ji.MinAvailable) + res
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Skaffold 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 kaniko

import (
	"context"
	"fmt"
	"io"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/constants"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/docker"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/kubernetes"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/v1alpha2"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/util"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

func RunKanikoBuild(ctx context.Context, out io.Writer, artifact *v1alpha2.Artifact, cfg *v1alpha2.KanikoBuild) (string, error) {
	dockerfilePath := artifact.DockerArtifact.DockerfilePath

	initialTag := util.RandomID()
	tarName := "context.tar.gz" // TODO(r2d4): until this is configurable upstream
	if err := docker.UploadContextToGCS(ctx, dockerfilePath, artifact.Workspace, cfg.GCSBucket, tarName); err != nil {
		return "", errors.Wrap(err, "uploading tar to gcs")
	}

	client, err := kubernetes.GetClientset()
	if err != nil {
		return "", errors.Wrap(err, "")
	}

	imageList := kubernetes.NewImageList()
	imageList.Add(constants.DefaultKanikoImage)

	logger := kubernetes.NewLogAggregator(out, imageList, kubernetes.NewColorPicker([]*v1alpha2.Artifact{artifact}))
	if err := logger.Start(ctx); err != nil {
		return "", errors.Wrap(err, "starting log streamer")
	}
	imageDst := fmt.Sprintf("%s:%s", artifact.ImageName, initialTag)
	p, err := client.CoreV1().Pods(cfg.Namespace).Create(&v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "kaniko",
			Labels:    map[string]string{"skaffold-kaniko": "skaffold-kaniko"},
			Namespace: cfg.Namespace,
		},
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				{
					Name:            "kaniko",
					Image:           constants.DefaultKanikoImage,
					ImagePullPolicy: v1.PullIfNotPresent,
					Args: []string{
						fmt.Sprintf("--dockerfile=%s", dockerfilePath),
						fmt.Sprintf("--bucket=%s", cfg.GCSBucket),
						fmt.Sprintf("--destination=%s", imageDst),
						fmt.Sprintf("-v=%s", logrus.GetLevel().String()),
					},
					VolumeMounts: []v1.VolumeMount{
						{
							Name:      constants.DefaultKanikoSecretName,
							MountPath: "/secret",
						},
					},
					Env: []v1.EnvVar{
						{
							Name:  "GOOGLE_APPLICATION_CREDENTIALS",
							Value: "/secret/kaniko-secret",
						},
					},
				},
			},
			Volumes: []v1.Volume{
				{
					Name: constants.DefaultKanikoSecretName,
					VolumeSource: v1.VolumeSource{
						Secret: &v1.SecretVolumeSource{
							SecretName: cfg.PullSecretName,
						},
					},
				},
			},
			RestartPolicy: v1.RestartPolicyNever,
		},
	})

	if err != nil {
		return "", errors.Wrap(err, "creating kaniko pod")
	}

	defer func() {
		imageList.Remove(constants.DefaultKanikoImage)
		if err := client.CoreV1().Pods(cfg.Namespace).Delete(p.Name, &metav1.DeleteOptions{
			GracePeriodSeconds: new(int64),
		}); err != nil {
			logrus.Fatalf("deleting pod: %s", err)
		}
	}()

	if err := kubernetes.WaitForPodComplete(client.CoreV1().Pods(cfg.Namespace), p.Name); err != nil {
		return "", errors.Wrap(err, "waiting for pod to complete")
	}

	return imageDst, nil
}
<commit_msg>Generate pod name<commit_after>/*
Copyright 2018 The Skaffold 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 kaniko

import (
	"context"
	"fmt"
	"io"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/constants"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/docker"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/kubernetes"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/v1alpha2"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/util"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

func RunKanikoBuild(ctx context.Context, out io.Writer, artifact *v1alpha2.Artifact, cfg *v1alpha2.KanikoBuild) (string, error) {
	dockerfilePath := artifact.DockerArtifact.DockerfilePath

	initialTag := util.RandomID()
	tarName := "context.tar.gz" // TODO(r2d4): until this is configurable upstream
	if err := docker.UploadContextToGCS(ctx, dockerfilePath, artifact.Workspace, cfg.GCSBucket, tarName); err != nil {
		return "", errors.Wrap(err, "uploading tar to gcs")
	}

	client, err := kubernetes.GetClientset()
	if err != nil {
		return "", errors.Wrap(err, "")
	}

	imageList := kubernetes.NewImageList()
	imageList.Add(constants.DefaultKanikoImage)

	logger := kubernetes.NewLogAggregator(out, imageList, kubernetes.NewColorPicker([]*v1alpha2.Artifact{artifact}))
	if err := logger.Start(ctx); err != nil {
		return "", errors.Wrap(err, "starting log streamer")
	}
	imageDst := fmt.Sprintf("%s:%s", artifact.ImageName, initialTag)
	p, err := client.CoreV1().Pods(cfg.Namespace).Create(&v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			GenerateName: "kaniko",
			Labels:       map[string]string{"skaffold-kaniko": "skaffold-kaniko"},
			Namespace:    cfg.Namespace,
		},
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				{
					Name:            "kaniko",
					Image:           constants.DefaultKanikoImage,
					ImagePullPolicy: v1.PullIfNotPresent,
					Args: []string{
						fmt.Sprintf("--dockerfile=%s", dockerfilePath),
						fmt.Sprintf("--bucket=%s", cfg.GCSBucket),
						fmt.Sprintf("--destination=%s", imageDst),
						fmt.Sprintf("-v=%s", logrus.GetLevel().String()),
					},
					VolumeMounts: []v1.VolumeMount{
						{
							Name:      constants.DefaultKanikoSecretName,
							MountPath: "/secret",
						},
					},
					Env: []v1.EnvVar{
						{
							Name:  "GOOGLE_APPLICATION_CREDENTIALS",
							Value: "/secret/kaniko-secret",
						},
					},
				},
			},
			Volumes: []v1.Volume{
				{
					Name: constants.DefaultKanikoSecretName,
					VolumeSource: v1.VolumeSource{
						Secret: &v1.SecretVolumeSource{
							SecretName: cfg.PullSecretName,
						},
					},
				},
			},
			RestartPolicy: v1.RestartPolicyNever,
		},
	})

	if err != nil {
		return "", errors.Wrap(err, "creating kaniko pod")
	}

	defer func() {
		imageList.Remove(constants.DefaultKanikoImage)
		if err := client.CoreV1().Pods(cfg.Namespace).Delete(p.Name, &metav1.DeleteOptions{
			GracePeriodSeconds: new(int64),
		}); err != nil {
			logrus.Fatalf("deleting pod: %s", err)
		}
	}()

	if err := kubernetes.WaitForPodComplete(client.CoreV1().Pods(cfg.Namespace), p.Name); err != nil {
		return "", errors.Wrap(err, "waiting for pod to complete")
	}

	return imageDst, nil
}
<|endoftext|>


<commit_before>// Package proxy is a transparent proxy built on the go-micro/server
package proxy

<commit_msg>Add proxy interface<commit_after>// Package proxy is a transparent proxy built on the go-micro/server
package proxy

import (
	"context"

	"github.com/micro/go-micro/init"
	"github.com/micro/go-micro/server"
)

// Proxy can be used as a proxy server for go-micro services
type Proxy interface {
	init.Options
	// ServeRequest will serve a request
	ServeRequest(context.Context, Request, Response) error
	// run the proxy
	Run() error
}
<|endoftext|>


<commit_before>package proxy

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httputil"
	"os"
	"regexp"
	"sync"
	"syscall"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/docker/docker/pkg/tlsconfig"
	"github.com/gorilla/mux"
	"github.com/pkg/errors"
	"github.com/rancher/websocket-proxy/k8s"
	"github.com/rancher/websocket-proxy/proxy/apiinterceptor"
	"github.com/rancher/websocket-proxy/proxy/proxyprotocol"
	proxyTls "github.com/rancher/websocket-proxy/proxy/tls"
)

var slashRegex = regexp.MustCompile("[/]{2,}")

type Starter struct {
	BackendPaths       []string
	FrontendPaths      []string
	FrontendHTTPPaths  []string
	StatsPaths         []string
	CattleProxyPaths   []string
	CattleWSProxyPaths []string
	Config             *Config
}

func (s *Starter) StartProxy() error {
	switcher := NewSwitcher(s.Config)

	backendMultiplexers := make(map[string]*multiplexer)
	bpm := &backendProxyManager{
		multiplexers: backendMultiplexers,
		mu:           &sync.RWMutex{},
	}

	frontendHandler := switcher.Wrap(&FrontendHandler{
		backend:         bpm,
		parsedPublicKey: s.Config.PublicKey,
	})

	statsHandler := switcher.Wrap(&StatsHandler{
		backend:         bpm,
		parsedPublicKey: s.Config.PublicKey,
	})

	backendHandler := switcher.Wrap(&BackendHandler{
		proxyManager:    bpm,
		parsedPublicKey: s.Config.PublicKey,
	})

	frontendHTTPHandler := switcher.Wrap(&FrontendHTTPHandler{
		FrontendHandler: FrontendHandler{
			backend:         bpm,
			parsedPublicKey: s.Config.PublicKey,
		},
		HTTPSPorts:  s.Config.ProxyProtoHTTPSPorts,
		TokenLookup: NewTokenLookup(s.Config.CattleAddr),
	})

	cattleProxy, cattleWsProxy, err := newCattleProxies(s.Config)
	if err != nil {
		log.Fatalf("Couldn't create cattle proxies: %v", err)
	}

	router := mux.NewRouter()

	router.HandleFunc("/version", k8s.Version)
	router.HandleFunc("/swaggerapi/api/v1", k8s.Swagger)

	for _, p := range s.BackendPaths {
		router.Handle(p, backendHandler).Methods("GET")
	}
	for _, p := range s.FrontendPaths {
		router.Handle(p, frontendHandler).Methods("GET")
	}
	for _, p := range s.FrontendHTTPPaths {
		router.Handle(p, frontendHTTPHandler).Methods("GET", "POST", "PUT", "DELETE", "PATCH")
	}
	for _, p := range s.StatsPaths {
		router.Handle(p, statsHandler).Methods("GET")
	}

	if s.Config.CattleAddr != "" {
		for _, p := range s.CattleWSProxyPaths {
			router.Handle(p, cattleWsProxy)
		}

		for _, p := range s.CattleProxyPaths {
			router.Handle(p, cattleProxy)
		}
	}

	if s.Config.ParentPid != 0 {
		go func() {
			for {
				process, err := os.FindProcess(s.Config.ParentPid)
				if err != nil {
					log.Fatalf("Failed to find process: %s\n", err)
				} else {
					err := process.Signal(syscall.Signal(0))
					if err != nil {
						log.Fatal("Parent process went away. Shutting down.")
					}
				}
				time.Sleep(time.Millisecond * 250)
			}
		}()
	}

	pcRouter := &pathCleaner{
		router: router,
	}

	swarmHandler := &SwarmHandler{
		FrontendHandler: frontendHTTPHandler,
		DefaultHandler:  pcRouter,
	}

	server := &http.Server{
		Handler:   swarmHandler,
		Addr:      s.Config.ListenAddr,
		ConnState: proxyprotocol.StateCleanup,
	}

	listener, err := net.Listen("tcp", s.Config.ListenAddr)
	if err != nil {
		log.Fatalf("Couldn't create listener: %s\n", err)
	}

	listener = &proxyprotocol.Listener{listener}

	if s.Config.TLSListenAddr != "" {
		tlsConfig, err := s.setupTLS()
		if err != nil {
			return err
		}

		if s.Config.TLSListenAddr == s.Config.ListenAddr {
			listener = &proxyTls.SplitListener{
				Listener: listener,
				Config:   tlsConfig,
			}
		} else {
			tlsListener, err := net.Listen("tcp", s.Config.TLSListenAddr)
			if err != nil {
				return err
			}
			tlsListener = &proxyprotocol.Listener{tlsListener}
			go func() {
				defer listener.Close()
				log.Error(server.Serve(tls.NewListener(tlsListener, tlsConfig)))
			}()
		}
	}

	err = server.Serve(listener)
	return err
}

func (s *Starter) setupTLS() (*tls.Config, error) {
	if s.Config.CattleAccessKey == "" {
		return nil, fmt.Errorf("No access key supplied to download cert")
	}

	certs, err := s.Config.GetCerts()
	if err != nil {
		return nil, err
	}

	tlsCert, err := tls.X509KeyPair(certs.Cert, certs.Key)
	if err != nil {
		return nil, err
	}

	clientCas := x509.NewCertPool()
	if !clientCas.AppendCertsFromPEM(certs.CA) {
		return nil, err
	}

	tlsConfig := tlsconfig.ServerDefault()
	tlsConfig.ClientAuth = tls.VerifyClientCertIfGiven
	tlsConfig.ClientCAs = clientCas
	tlsConfig.Certificates = []tls.Certificate{tlsCert}

	return tlsConfig, nil
}

type pathCleaner struct {
	router *mux.Router
}

func (p *pathCleaner) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	if cleanedPath := p.cleanPath(req.URL.Path); cleanedPath != req.URL.Path {
		req.URL.Path = cleanedPath
		req.URL.Scheme = "http"
	}
	p.router.ServeHTTP(rw, req)
}

func (p *pathCleaner) cleanPath(path string) string {
	return slashRegex.ReplaceAllString(path, "/")
}

func newWSProxy(config *Config) http.Handler {
	cattleAddr := config.CattleAddr
	director := func(req *http.Request) {
		req.URL.Scheme = "http"
		req.URL.Host = cattleAddr
	}

	cattleProxy := &httputil.ReverseProxy{
		Director:      director,
		FlushInterval: time.Millisecond * 100,
	}

	reverseProxy := &proxyProtocolConverter{
		p:          cattleProxy,
		httpsPorts: config.ProxyProtoHTTPSPorts,
	}

	wsProxy := &cattleWSProxy{
		reverseProxy: reverseProxy,
		cattleAddr:   cattleAddr,
	}

	return wsProxy
}

func newCattleProxies(config *Config) (*proxyProtocolConverter, *cattleWSProxy, error) {
	cattleAddr := config.CattleAddr

	apiProxyHandler, err := apiinterceptor.NewInterceptor(config.APIInterceptorConfigFile, cattleAddr)
	if err != nil {
		return nil, nil, errors.Wrap(err, "Couldn't create API interceptor")
	}

	reverseProxy := &proxyProtocolConverter{
		httpsPorts: config.ProxyProtoHTTPSPorts,
		p:          apiProxyHandler,
	}

	wsProxy := &cattleWSProxy{
		reverseProxy: reverseProxy,
		cattleAddr:   cattleAddr,
	}

	return reverseProxy, wsProxy, nil
}

type proxyProtocolConverter struct {
	httpsPorts map[int]bool
	p          http.Handler
}

func (h *proxyProtocolConverter) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	proxyprotocol.AddHeaders(req, h.httpsPorts)
	h.p.ServeHTTP(rw, req)
}

type cattleWSProxy struct {
	reverseProxy *proxyProtocolConverter
	cattleAddr   string
}

func (h *cattleWSProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	if len(req.Header.Get("Upgrade")) > 0 {
		proxyprotocol.AddHeaders(req, h.reverseProxy.httpsPorts)
		h.serveWebsocket(rw, req)
	} else {
		h.reverseProxy.ServeHTTP(rw, req)
	}
}

func (h *cattleWSProxy) serveWebsocket(rw http.ResponseWriter, req *http.Request) {
	// Inspired by https://groups.google.com/forum/#!searchin/golang-nuts/httputil.ReverseProxy$20$2B$20websockets/golang-nuts/KBx9pDlvFOc/01vn1qUyVdwJ
	target := h.cattleAddr
	d, err := net.Dial("tcp", target)
	if err != nil {
		log.WithField("error", err).Error("Error dialing websocket backend.")
		http.Error(rw, "Unable to establish websocket connection: can't dial.", 500)
		return
	}
	hj, ok := rw.(http.Hijacker)
	if !ok {
		http.Error(rw, "Unable to establish websocket connection: no hijacker.", 500)
		return
	}
	nc, _, err := hj.Hijack()
	if err != nil {
		log.WithField("error", err).Error("Hijack error.")
		http.Error(rw, "Unable to establish websocket connection: can't hijack.", 500)
		return
	}
	defer nc.Close()
	defer d.Close()

	err = req.Write(d)
	if err != nil {
		log.WithField("error", err).Error("Error copying request to target.")
		return
	}

	errc := make(chan error, 2)
	cp := func(dst io.Writer, src io.Reader) {
		_, err := io.Copy(dst, src)
		errc <- err
	}
	go cp(d, nc)
	go cp(nc, d)
	<-errc
}
<commit_msg>Support HEAD in http proxy<commit_after>package proxy

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httputil"
	"os"
	"regexp"
	"sync"
	"syscall"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/docker/docker/pkg/tlsconfig"
	"github.com/gorilla/mux"
	"github.com/pkg/errors"
	"github.com/rancher/websocket-proxy/k8s"
	"github.com/rancher/websocket-proxy/proxy/apiinterceptor"
	"github.com/rancher/websocket-proxy/proxy/proxyprotocol"
	proxyTls "github.com/rancher/websocket-proxy/proxy/tls"
)

var slashRegex = regexp.MustCompile("[/]{2,}")

type Starter struct {
	BackendPaths       []string
	FrontendPaths      []string
	FrontendHTTPPaths  []string
	StatsPaths         []string
	CattleProxyPaths   []string
	CattleWSProxyPaths []string
	Config             *Config
}

func (s *Starter) StartProxy() error {
	switcher := NewSwitcher(s.Config)

	backendMultiplexers := make(map[string]*multiplexer)
	bpm := &backendProxyManager{
		multiplexers: backendMultiplexers,
		mu:           &sync.RWMutex{},
	}

	frontendHandler := switcher.Wrap(&FrontendHandler{
		backend:         bpm,
		parsedPublicKey: s.Config.PublicKey,
	})

	statsHandler := switcher.Wrap(&StatsHandler{
		backend:         bpm,
		parsedPublicKey: s.Config.PublicKey,
	})

	backendHandler := switcher.Wrap(&BackendHandler{
		proxyManager:    bpm,
		parsedPublicKey: s.Config.PublicKey,
	})

	frontendHTTPHandler := switcher.Wrap(&FrontendHTTPHandler{
		FrontendHandler: FrontendHandler{
			backend:         bpm,
			parsedPublicKey: s.Config.PublicKey,
		},
		HTTPSPorts:  s.Config.ProxyProtoHTTPSPorts,
		TokenLookup: NewTokenLookup(s.Config.CattleAddr),
	})

	cattleProxy, cattleWsProxy, err := newCattleProxies(s.Config)
	if err != nil {
		log.Fatalf("Couldn't create cattle proxies: %v", err)
	}

	router := mux.NewRouter()

	router.HandleFunc("/version", k8s.Version)
	router.HandleFunc("/swaggerapi/api/v1", k8s.Swagger)

	for _, p := range s.BackendPaths {
		router.Handle(p, backendHandler).Methods("GET")
	}
	for _, p := range s.FrontendPaths {
		router.Handle(p, frontendHandler).Methods("GET")
	}
	for _, p := range s.FrontendHTTPPaths {
		router.Handle(p, frontendHTTPHandler).Methods("GET", "POST", "PUT", "DELETE", "PATCH", "HEAD")
	}
	for _, p := range s.StatsPaths {
		router.Handle(p, statsHandler).Methods("GET")
	}

	if s.Config.CattleAddr != "" {
		for _, p := range s.CattleWSProxyPaths {
			router.Handle(p, cattleWsProxy)
		}

		for _, p := range s.CattleProxyPaths {
			router.Handle(p, cattleProxy)
		}
	}

	if s.Config.ParentPid != 0 {
		go func() {
			for {
				process, err := os.FindProcess(s.Config.ParentPid)
				if err != nil {
					log.Fatalf("Failed to find process: %s\n", err)
				} else {
					err := process.Signal(syscall.Signal(0))
					if err != nil {
						log.Fatal("Parent process went away. Shutting down.")
					}
				}
				time.Sleep(time.Millisecond * 250)
			}
		}()
	}

	pcRouter := &pathCleaner{
		router: router,
	}

	swarmHandler := &SwarmHandler{
		FrontendHandler: frontendHTTPHandler,
		DefaultHandler:  pcRouter,
	}

	server := &http.Server{
		Handler:   swarmHandler,
		Addr:      s.Config.ListenAddr,
		ConnState: proxyprotocol.StateCleanup,
	}

	listener, err := net.Listen("tcp", s.Config.ListenAddr)
	if err != nil {
		log.Fatalf("Couldn't create listener: %s\n", err)
	}

	listener = &proxyprotocol.Listener{listener}

	if s.Config.TLSListenAddr != "" {
		tlsConfig, err := s.setupTLS()
		if err != nil {
			return err
		}

		if s.Config.TLSListenAddr == s.Config.ListenAddr {
			listener = &proxyTls.SplitListener{
				Listener: listener,
				Config:   tlsConfig,
			}
		} else {
			tlsListener, err := net.Listen("tcp", s.Config.TLSListenAddr)
			if err != nil {
				return err
			}
			tlsListener = &proxyprotocol.Listener{tlsListener}
			go func() {
				defer listener.Close()
				log.Error(server.Serve(tls.NewListener(tlsListener, tlsConfig)))
			}()
		}
	}

	err = server.Serve(listener)
	return err
}

func (s *Starter) setupTLS() (*tls.Config, error) {
	if s.Config.CattleAccessKey == "" {
		return nil, fmt.Errorf("No access key supplied to download cert")
	}

	certs, err := s.Config.GetCerts()
	if err != nil {
		return nil, err
	}

	tlsCert, err := tls.X509KeyPair(certs.Cert, certs.Key)
	if err != nil {
		return nil, err
	}

	clientCas := x509.NewCertPool()
	if !clientCas.AppendCertsFromPEM(certs.CA) {
		return nil, err
	}

	tlsConfig := tlsconfig.ServerDefault()
	tlsConfig.ClientAuth = tls.VerifyClientCertIfGiven
	tlsConfig.ClientCAs = clientCas
	tlsConfig.Certificates = []tls.Certificate{tlsCert}

	return tlsConfig, nil
}

type pathCleaner struct {
	router *mux.Router
}

func (p *pathCleaner) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	if cleanedPath := p.cleanPath(req.URL.Path); cleanedPath != req.URL.Path {
		req.URL.Path = cleanedPath
		req.URL.Scheme = "http"
	}
	p.router.ServeHTTP(rw, req)
}

func (p *pathCleaner) cleanPath(path string) string {
	return slashRegex.ReplaceAllString(path, "/")
}

func newWSProxy(config *Config) http.Handler {
	cattleAddr := config.CattleAddr
	director := func(req *http.Request) {
		req.URL.Scheme = "http"
		req.URL.Host = cattleAddr
	}

	cattleProxy := &httputil.ReverseProxy{
		Director:      director,
		FlushInterval: time.Millisecond * 100,
	}

	reverseProxy := &proxyProtocolConverter{
		p:          cattleProxy,
		httpsPorts: config.ProxyProtoHTTPSPorts,
	}

	wsProxy := &cattleWSProxy{
		reverseProxy: reverseProxy,
		cattleAddr:   cattleAddr,
	}

	return wsProxy
}

func newCattleProxies(config *Config) (*proxyProtocolConverter, *cattleWSProxy, error) {
	cattleAddr := config.CattleAddr

	apiProxyHandler, err := apiinterceptor.NewInterceptor(config.APIInterceptorConfigFile, cattleAddr)
	if err != nil {
		return nil, nil, errors.Wrap(err, "Couldn't create API interceptor")
	}

	reverseProxy := &proxyProtocolConverter{
		httpsPorts: config.ProxyProtoHTTPSPorts,
		p:          apiProxyHandler,
	}

	wsProxy := &cattleWSProxy{
		reverseProxy: reverseProxy,
		cattleAddr:   cattleAddr,
	}

	return reverseProxy, wsProxy, nil
}

type proxyProtocolConverter struct {
	httpsPorts map[int]bool
	p          http.Handler
}

func (h *proxyProtocolConverter) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	proxyprotocol.AddHeaders(req, h.httpsPorts)
	h.p.ServeHTTP(rw, req)
}

type cattleWSProxy struct {
	reverseProxy *proxyProtocolConverter
	cattleAddr   string
}

func (h *cattleWSProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	if len(req.Header.Get("Upgrade")) > 0 {
		proxyprotocol.AddHeaders(req, h.reverseProxy.httpsPorts)
		h.serveWebsocket(rw, req)
	} else {
		h.reverseProxy.ServeHTTP(rw, req)
	}
}

func (h *cattleWSProxy) serveWebsocket(rw http.ResponseWriter, req *http.Request) {
	// Inspired by https://groups.google.com/forum/#!searchin/golang-nuts/httputil.ReverseProxy$20$2B$20websockets/golang-nuts/KBx9pDlvFOc/01vn1qUyVdwJ
	target := h.cattleAddr
	d, err := net.Dial("tcp", target)
	if err != nil {
		log.WithField("error", err).Error("Error dialing websocket backend.")
		http.Error(rw, "Unable to establish websocket connection: can't dial.", 500)
		return
	}
	hj, ok := rw.(http.Hijacker)
	if !ok {
		http.Error(rw, "Unable to establish websocket connection: no hijacker.", 500)
		return
	}
	nc, _, err := hj.Hijack()
	if err != nil {
		log.WithField("error", err).Error("Hijack error.")
		http.Error(rw, "Unable to establish websocket connection: can't hijack.", 500)
		return
	}
	defer nc.Close()
	defer d.Close()

	err = req.Write(d)
	if err != nil {
		log.WithField("error", err).Error("Error copying request to target.")
		return
	}

	errc := make(chan error, 2)
	cp := func(dst io.Writer, src io.Reader) {
		_, err := io.Copy(dst, src)
		errc <- err
	}
	go cp(d, nc)
	go cp(nc, d)
	<-errc
}
<|endoftext|>


<commit_before>package proxy

import (
	"bufio"
	"fmt"
	"log"
	"os"
	"reflect"
	"regexp"
	"sort"
	"strings"
	"text/tabwriter"

	"github.com/andystanton/proxybastard/util"
	"github.com/deckarep/golang-set"
)

func addToMap(frequencyMap map[string]int, value string) map[string]int {
	if len(strings.TrimSpace(value)) > 0 {
		if _, ok := frequencyMap[value]; !ok {
			frequencyMap[value] = 0
		}
		frequencyMap[value] = frequencyMap[value] + 1
	}
	return frequencyMap
}

func awaitInput(prompt string, pattern string, prefix string) string {
	var matched string
	var found bool
	fmt.Printf("%s%s\n", prefix, prompt)
	fmt.Printf("%s> ", prefix)
	for i := 0; i < 3; i++ {
		reader := bufio.NewReader(os.Stdin)
		text, _ := reader.ReadString('\n')
		rPattern := regexp.MustCompile(pattern)

		if result := strings.TrimSpace(text); rPattern.MatchString(result) {
			matched = result
			found = true
			break
		} else if i < 2 {
			log.Printf("%sThat doesn't look right - try again...\n", prefix)
			log.Println()
			fmt.Printf("%s%s\n", prefix, prompt)
			fmt.Printf("%s> ", prefix)
		} else {
			log.Println()
			log.Print("Three failed attempts - aborting!")
		}
	}
	if !found {
		os.Exit(1)
	}
	return matched
}

func awaitFileInput(prompt string, prefix string) string {
	var matched string
	var found bool
	fmt.Printf("%s%s\n", prefix, prompt)
	fmt.Printf("%s> ", prefix)
	for i := 0; i < 3; i++ {
		reader := bufio.NewReader(os.Stdin)
		text, _ := reader.ReadString('\n')
		if file := util.SanitisePath(strings.TrimSpace(text)); util.FileExists(file) {
			matched = file
			found = true
			break
		} else if i < 2 {
			log.Printf("%sThat doesn't look right - try again...\n", prefix)
			log.Println()
			fmt.Printf("%s%s\n", prefix, prompt)
			fmt.Printf("%s> ", prefix)
		} else {
			log.Println()
			log.Print("Three failed attempts - aborting!")
		}
	}
	if !found {
		os.Exit(1)
	}
	return matched
}

// Setup presents the user with setup options.
func Setup(version string, acceptDefaults bool) {
	suggestedConfiguration := suggestConfiguration()
	actualConfiguration := Configuration{}
	actualConfiguration.Version = version
	readyToWrite := false

	if acceptDefaults {
		actualConfiguration = suggestedConfiguration
		readyToWrite = true
	} else {
		httpProxySet := false
		if len(suggestedConfiguration.ProxyHost) > 0 {
			message := fmt.Sprintf("Use suggested http proxy %s:%s? [Yn]", suggestedConfiguration.ProxyHost, suggestedConfiguration.ProxyPort)
			input := awaitInput(message, "(y|n|^$)", "")
			httpProxySet = strings.EqualFold(input, "y") || strings.EqualFold(input, "")
			actualConfiguration.ProxyHost = suggestedConfiguration.ProxyHost
			actualConfiguration.ProxyPort = suggestedConfiguration.ProxyPort
			fmt.Println()
		}

		if !httpProxySet {
			proxyHostPattern := "(?:https?://)?(.+):(\\d+)"
			proxyHostRegexp := regexp.MustCompile(proxyHostPattern)
			matches := proxyHostRegexp.FindStringSubmatch(awaitInput("Please enter an http proxy e.g. http://proxybastard:1234", proxyHostPattern, ""))
			actualConfiguration.ProxyHost = fmt.Sprintf("http://%s", matches[1])
			actualConfiguration.ProxyPort = matches[2]
			httpProxySet = true
			fmt.Println()
		}

		socksProxySet := false
		if len(suggestedConfiguration.SOCKSProxyHost) > 0 {
			message := fmt.Sprintf("Use suggested SOCKS proxy %s:%s? [Yn]", suggestedConfiguration.SOCKSProxyHost, suggestedConfiguration.SOCKSProxyPort)
			input := awaitInput(message, "(y|n|^$)", "")
			socksProxySet = strings.EqualFold(input, "y") || strings.EqualFold(input, "")
			actualConfiguration.SOCKSProxyHost = suggestedConfiguration.SOCKSProxyHost
			actualConfiguration.SOCKSProxyPort = suggestedConfiguration.SOCKSProxyPort
			fmt.Println()
		}

		if !socksProxySet {
			socksHostPattern := "(?:(.+):(\\d+)|^$)"
			sockHostRegexp := regexp.MustCompile(socksHostPattern)
			matches := sockHostRegexp.FindStringSubmatch(awaitInput("Please enter a SOCKS proxy or press return for none e.g. socks.proxybastard:1234", socksHostPattern, ""))
			if len(matches) > 0 {
				socksProxySet = true
				actualConfiguration.SOCKSProxyHost = matches[1]
				actualConfiguration.SOCKSProxyPort = matches[2]
			}
			fmt.Println()
		}

		if suggestedConfiguration.Targets != nil {
			targetsField := reflect.Indirect(reflect.ValueOf(suggestedConfiguration.Targets))
			for i := 0; i < targetsField.NumField(); i++ {
				fieldName := targetsField.Type().Field(i).Name

				// valueForFieldRequired := false
				if !util.InterfaceIsZero(targetsField.Field(i).Interface()) {
					targetField := reflect.Indirect(reflect.ValueOf(targetsField.Field(i).Interface()))

					// step 1 - go simple and ask if user wants suggested config
					message := fmt.Sprintf("Found %s! Use suggested configuration? [Yn]", fieldName)
					input := awaitInput(message, "(y|n|^$)", "")
					configurationSet := strings.EqualFold(input, "y") || strings.EqualFold(input, "")

					if configurationSet {
						// great! put the suggested config in the actual config
						if actualConfiguration.Targets == nil {
							actualConfiguration.Targets = &TargetsConfiguration{}
						}
						actualField := reflect.Indirect(reflect.ValueOf(actualConfiguration.Targets)).FieldByName(fieldName)
						actualField.Set(reflect.ValueOf(targetsField.Field(i).Interface()))
					} else {
						fmt.Println()

						message := fmt.Sprintf("Enable %s? [Yn]", fieldName)
						input := awaitInput(message, "(y|n|^$)", "  ")
						if strings.EqualFold(input, "y") || strings.EqualFold(input, "") {
							if actualConfiguration.Targets == nil {
								actualConfiguration.Targets = &TargetsConfiguration{}
							}
							actualField := reflect.New(reflect.TypeOf(targetField.Interface())).Interface()
							reflect.Indirect(reflect.ValueOf(actualField)).FieldByName("Enabled").Set(reflect.ValueOf(true))
							fmt.Println()

							if util.ValueHasField(targetField, "Files") {
								fieldFiles := targetField.FieldByName("Files").Interface().([]string)
								message := fmt.Sprintf("Use suggested file %s? [Yn]", strings.Join(fieldFiles, ", "))
								input := awaitInput(message, "(y|n|^$)", "  ")
								if strings.EqualFold(input, "y") || strings.EqualFold(input, "") {
									reflect.Indirect(reflect.ValueOf(actualField)).FieldByName("Files").Set(reflect.ValueOf(fieldFiles))
								} else {
									fmt.Println()
									message := fmt.Sprintf("Enter the path to a valid %s configuration file", fieldName)
									input := awaitFileInput(message, "  ")
									reflect.Indirect(reflect.ValueOf(actualField)).FieldByName("Files").Set(reflect.ValueOf([]string{input}))
								}
							}

							if util.ValueHasMethod(targetField, "CustomPrompt") {
								customMethod := targetField.MethodByName("CustomPrompt")
								actualField = customMethod.Call([]reflect.Value{})[0].Interface()
							}

							reflect.Indirect(reflect.ValueOf(actualConfiguration.Targets)).FieldByName(fieldName).Set(reflect.ValueOf(actualField))
						}
					}
					fmt.Println()
				} else {
					fmt.Println("Suggested config does not contain " + fieldName)
				}
			}
		}

		w := tabwriter.NewWriter(os.Stdout, 20, 0, 1, ' ', 0)

		fmt.Fprintln(w, "Settings")
		fmt.Fprintln(w, "================================================================")
		fmt.Fprintln(w, "HTTP Proxy")
		fmt.Fprintf(w, " - Host\t : %s\n", actualConfiguration.ProxyHost)
		fmt.Fprintf(w, " - Port\t : %s\n", actualConfiguration.ProxyPort)

		if len(actualConfiguration.SOCKSProxyHost) > 0 {
			fmt.Fprintln(w)
			fmt.Fprintln(w, "SOCKS Proxy")
			fmt.Fprintf(w, " - Host\t : %s\n", actualConfiguration.SOCKSProxyHost)
			fmt.Fprintf(w, " - Port\t : %s\n", actualConfiguration.SOCKSProxyPort)
		}
		fmt.Fprintln(w, "================================================================")
		fmt.Fprintln(w)

		if actualConfiguration.Targets != nil {
			targetsField := reflect.Indirect(reflect.ValueOf(actualConfiguration.Targets))

			for i := 0; i < targetsField.NumField(); i++ {
				fieldName := targetsField.Type().Field(i).Name

				if !util.InterfaceIsZero(targetsField.Field(i).Interface()) {
					targetField := reflect.Indirect(reflect.ValueOf(targetsField.Field(i).Interface()))
					withConfig, _ := targetField.Interface().(WithConfig)
					fmt.Fprintln(w, fieldName)
					fmt.Fprintf(w, " - Enabled\t : %v\n", withConfig.isEnabled())

					if util.ValueHasField(targetField, "Files") {
						fieldFiles := targetField.FieldByName("Files").Interface().([]string)
						newFiles := []string{}
						for _, file := range fieldFiles {
							newFiles = append(newFiles, util.UnsanitisePath(file))
						}
						fmt.Fprintf(w, " - Files\t : %s\n", strings.Join(newFiles, ","))
					}

					if util.ValueHasMethod(targetField, "CustomFields") {
						customMethod := targetField.MethodByName("CustomFields")
						extraFields := customMethod.Call([]reflect.Value{})[0].Interface().(map[string]string)
						for k, v := range extraFields {
							fmt.Fprintf(w, " - %s\t : %s\n", k, v)
						}
					}
					fmt.Fprintln(w)
				}
			}
		}
		fmt.Fprintln(w, "================================================================")
		fmt.Fprintln(w)
		w.Flush()

		input := awaitInput("Write these settings to ~/.proxybastard/config.json? [Yn]", "(y|n|^$)", "")
		fmt.Println()
		readyToWrite = strings.EqualFold(input, "y") || strings.EqualFold(input, "")
	}

	// marshalled, err := json.MarshalIndent(actualConfiguration, "", "    ")
	// if err != nil {
	// 	log.Fatal(err)
	// }
	// fmt.Printf("%s\n\n", string(marshalled))
	if readyToWrite {
		fmt.Println("Done")
	} else {
		fmt.Println("kthx")
	}
}

func getHighestFrequency(frequencyMap map[string]int) string {
	var highestValue int
	var mostFrequentKey string

	keys := []string{}
	for key := range frequencyMap {
		keys = append(keys, key)
	}
	sort.Strings(keys)

	for _, key := range keys {
		if value := frequencyMap[key]; value > highestValue {
			highestValue = value
			mostFrequentKey = key
		}
	}
	return mostFrequentKey
}

func suggestConfiguration() Configuration {
	var suggestedConfiguration Configuration

	suggestedProxyHosts := make(map[string]int)
	suggestedProxyPorts := make(map[string]int)
	suggestedSOCKSProxyHosts := make(map[string]int)
	suggestedSOCKSProxyPorts := make(map[string]int)
	suggestedNonProxyHosts := mapset.NewSet()

	targetsField := reflect.Indirect(reflect.ValueOf(&TargetsConfiguration{}))
	for i := 0; i < targetsField.NumField(); i++ {

		configurationField := reflect.New(reflect.TypeOf(targetsField.Field(i).Interface()).Elem()).Interface()
		withConfig, hasConfig := configurationField.(WithConfig)
		if hasConfig {
			fieldName := targetsField.Type().Field(i).Name

			if suggestedItemConfiguration := withConfig.suggestConfiguration(); suggestedItemConfiguration != nil {
				if suggestedConfiguration.Targets == nil {
					suggestedConfiguration.Targets = &TargetsConfiguration{}
				}

				addToMap(suggestedProxyHosts, util.SanitiseHTTPProxyURL(suggestedItemConfiguration.ProxyHost))
				addToMap(suggestedProxyPorts, suggestedItemConfiguration.ProxyPort)
				addToMap(suggestedSOCKSProxyHosts, suggestedItemConfiguration.SOCKSProxyHost)
				addToMap(suggestedSOCKSProxyPorts, suggestedItemConfiguration.SOCKSProxyPort)

				for _, nonProxyHost := range suggestedItemConfiguration.NonProxyHosts {
					suggestedNonProxyHosts.Add(nonProxyHost)
				}

				targetsField := reflect.Indirect(reflect.ValueOf(suggestedConfiguration.Targets))
				targetsFieldSuggested := reflect.Indirect(reflect.ValueOf(suggestedItemConfiguration.Targets))

				targetsField.FieldByName(fieldName).Set(targetsFieldSuggested.FieldByName(fieldName))
			}
		}
	}

	suggestedConfiguration.ProxyHost = getHighestFrequency(suggestedProxyHosts)
	suggestedConfiguration.ProxyPort = getHighestFrequency(suggestedProxyPorts)
	suggestedConfiguration.SOCKSProxyHost = getHighestFrequency(suggestedSOCKSProxyHosts)
	suggestedConfiguration.SOCKSProxyPort = getHighestFrequency(suggestedSOCKSProxyPorts)
	for suggestedNonProxyHost := range suggestedNonProxyHosts.Iter() {
		suggestedConfiguration.NonProxyHosts = append(suggestedConfiguration.NonProxyHosts, suggestedNonProxyHost.(string))
	}

	return suggestedConfiguration
}
<commit_msg>write config to file<commit_after>package proxy

import (
	"bufio"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"reflect"
	"regexp"
	"sort"
	"strings"
	"text/tabwriter"

	"github.com/andystanton/proxybastard/util"
	"github.com/deckarep/golang-set"
)

func addToMap(frequencyMap map[string]int, value string) map[string]int {
	if len(strings.TrimSpace(value)) > 0 {
		if _, ok := frequencyMap[value]; !ok {
			frequencyMap[value] = 0
		}
		frequencyMap[value] = frequencyMap[value] + 1
	}
	return frequencyMap
}

func awaitInput(prompt string, pattern string, prefix string) string {
	var matched string
	var found bool
	fmt.Printf("%s%s\n", prefix, prompt)
	fmt.Printf("%s> ", prefix)
	for i := 0; i < 3; i++ {
		reader := bufio.NewReader(os.Stdin)
		text, _ := reader.ReadString('\n')
		rPattern := regexp.MustCompile(pattern)

		if result := strings.TrimSpace(text); rPattern.MatchString(result) {
			matched = result
			found = true
			break
		} else if i < 2 {
			log.Printf("%sThat doesn't look right - try again...\n", prefix)
			log.Println()
			fmt.Printf("%s%s\n", prefix, prompt)
			fmt.Printf("%s> ", prefix)
		} else {
			log.Println()
			log.Print("Three failed attempts - aborting!")
		}
	}
	if !found {
		os.Exit(1)
	}
	return matched
}

func awaitFileInput(prompt string, prefix string) string {
	var matched string
	var found bool
	fmt.Printf("%s%s\n", prefix, prompt)
	fmt.Printf("%s> ", prefix)
	for i := 0; i < 3; i++ {
		reader := bufio.NewReader(os.Stdin)
		text, _ := reader.ReadString('\n')
		if file := util.SanitisePath(strings.TrimSpace(text)); util.FileExists(file) {
			matched = file
			found = true
			break
		} else if i < 2 {
			log.Printf("%sThat doesn't look right - try again...\n", prefix)
			log.Println()
			fmt.Printf("%s%s\n", prefix, prompt)
			fmt.Printf("%s> ", prefix)
		} else {
			log.Println()
			log.Print("Three failed attempts - aborting!")
		}
	}
	if !found {
		os.Exit(1)
	}
	return matched
}

// Setup presents the user with setup options.
func Setup(version string, acceptDefaults bool) {
	suggestedConfiguration := suggestConfiguration()
	actualConfiguration := Configuration{}
	actualConfiguration.Version = version
	readyToWrite := false

	if acceptDefaults {
		actualConfiguration = suggestedConfiguration
		readyToWrite = true
	} else {
		httpProxySet := false
		if len(suggestedConfiguration.ProxyHost) > 0 {
			message := fmt.Sprintf("Use suggested http proxy %s:%s? [Yn]", suggestedConfiguration.ProxyHost, suggestedConfiguration.ProxyPort)
			input := awaitInput(message, "(y|n|^$)", "")
			httpProxySet = strings.EqualFold(input, "y") || strings.EqualFold(input, "")
			actualConfiguration.ProxyHost = suggestedConfiguration.ProxyHost
			actualConfiguration.ProxyPort = suggestedConfiguration.ProxyPort
			fmt.Println()
		}

		if !httpProxySet {
			proxyHostPattern := "(?:https?://)?(.+):(\\d+)"
			proxyHostRegexp := regexp.MustCompile(proxyHostPattern)
			matches := proxyHostRegexp.FindStringSubmatch(awaitInput("Please enter an http proxy e.g. http://proxybastard:1234", proxyHostPattern, ""))
			actualConfiguration.ProxyHost = fmt.Sprintf("http://%s", matches[1])
			actualConfiguration.ProxyPort = matches[2]
			httpProxySet = true
			fmt.Println()
		}

		socksProxySet := false
		if len(suggestedConfiguration.SOCKSProxyHost) > 0 {
			message := fmt.Sprintf("Use suggested SOCKS proxy %s:%s? [Yn]", suggestedConfiguration.SOCKSProxyHost, suggestedConfiguration.SOCKSProxyPort)
			input := awaitInput(message, "(y|n|^$)", "")
			socksProxySet = strings.EqualFold(input, "y") || strings.EqualFold(input, "")
			actualConfiguration.SOCKSProxyHost = suggestedConfiguration.SOCKSProxyHost
			actualConfiguration.SOCKSProxyPort = suggestedConfiguration.SOCKSProxyPort
			fmt.Println()
		}

		if !socksProxySet {
			socksHostPattern := "(?:(.+):(\\d+)|^$)"
			sockHostRegexp := regexp.MustCompile(socksHostPattern)
			matches := sockHostRegexp.FindStringSubmatch(awaitInput("Please enter a SOCKS proxy or press return for none e.g. socks.proxybastard:1234", socksHostPattern, ""))
			if len(matches) > 0 {
				socksProxySet = true
				actualConfiguration.SOCKSProxyHost = matches[1]
				actualConfiguration.SOCKSProxyPort = matches[2]
			}
			fmt.Println()
		}

		if suggestedConfiguration.Targets != nil {
			targetsField := reflect.Indirect(reflect.ValueOf(suggestedConfiguration.Targets))
			for i := 0; i < targetsField.NumField(); i++ {
				fieldName := targetsField.Type().Field(i).Name

				// valueForFieldRequired := false
				if !util.InterfaceIsZero(targetsField.Field(i).Interface()) {
					targetField := reflect.Indirect(reflect.ValueOf(targetsField.Field(i).Interface()))

					// step 1 - go simple and ask if user wants suggested config
					message := fmt.Sprintf("Found %s! Use suggested configuration? [Yn]", fieldName)
					input := awaitInput(message, "(y|n|^$)", "")
					configurationSet := strings.EqualFold(input, "y") || strings.EqualFold(input, "")

					if configurationSet {
						// great! put the suggested config in the actual config
						if actualConfiguration.Targets == nil {
							actualConfiguration.Targets = &TargetsConfiguration{}
						}
						actualField := reflect.Indirect(reflect.ValueOf(actualConfiguration.Targets)).FieldByName(fieldName)
						actualField.Set(reflect.ValueOf(targetsField.Field(i).Interface()))
					} else {
						fmt.Println()

						message := fmt.Sprintf("Enable %s? [Yn]", fieldName)
						input := awaitInput(message, "(y|n|^$)", "  ")
						if strings.EqualFold(input, "y") || strings.EqualFold(input, "") {
							if actualConfiguration.Targets == nil {
								actualConfiguration.Targets = &TargetsConfiguration{}
							}
							actualField := reflect.New(reflect.TypeOf(targetField.Interface())).Interface()
							reflect.Indirect(reflect.ValueOf(actualField)).FieldByName("Enabled").Set(reflect.ValueOf(true))
							fmt.Println()

							if util.ValueHasField(targetField, "Files") {
								fieldFiles := targetField.FieldByName("Files").Interface().([]string)
								message := fmt.Sprintf("Use suggested file %s? [Yn]", strings.Join(fieldFiles, ", "))
								input := awaitInput(message, "(y|n|^$)", "  ")
								if strings.EqualFold(input, "y") || strings.EqualFold(input, "") {
									reflect.Indirect(reflect.ValueOf(actualField)).FieldByName("Files").Set(reflect.ValueOf(fieldFiles))
								} else {
									fmt.Println()
									message := fmt.Sprintf("Enter the path to a valid %s configuration file", fieldName)
									input := awaitFileInput(message, "  ")
									reflect.Indirect(reflect.ValueOf(actualField)).FieldByName("Files").Set(reflect.ValueOf([]string{input}))
								}
							}

							if util.ValueHasMethod(targetField, "CustomPrompt") {
								customMethod := targetField.MethodByName("CustomPrompt")
								actualField = customMethod.Call([]reflect.Value{})[0].Interface()
							}

							reflect.Indirect(reflect.ValueOf(actualConfiguration.Targets)).FieldByName(fieldName).Set(reflect.ValueOf(actualField))
						}
					}
					fmt.Println()
				} else {
					fmt.Println("Suggested config does not contain " + fieldName)
				}
			}
		}

		w := tabwriter.NewWriter(os.Stdout, 20, 0, 1, ' ', 0)

		fmt.Fprintln(w, "Settings")
		fmt.Fprintln(w, "================================================================")
		fmt.Fprintln(w, "HTTP Proxy")
		fmt.Fprintf(w, " - Host\t : %s\n", actualConfiguration.ProxyHost)
		fmt.Fprintf(w, " - Port\t : %s\n", actualConfiguration.ProxyPort)

		if len(actualConfiguration.SOCKSProxyHost) > 0 {
			fmt.Fprintln(w)
			fmt.Fprintln(w, "SOCKS Proxy")
			fmt.Fprintf(w, " - Host\t : %s\n", actualConfiguration.SOCKSProxyHost)
			fmt.Fprintf(w, " - Port\t : %s\n", actualConfiguration.SOCKSProxyPort)
		}
		fmt.Fprintln(w, "================================================================")
		fmt.Fprintln(w)

		if actualConfiguration.Targets != nil {
			targetsField := reflect.Indirect(reflect.ValueOf(actualConfiguration.Targets))

			for i := 0; i < targetsField.NumField(); i++ {
				fieldName := targetsField.Type().Field(i).Name

				if !util.InterfaceIsZero(targetsField.Field(i).Interface()) {
					targetField := reflect.Indirect(reflect.ValueOf(targetsField.Field(i).Interface()))
					withConfig, _ := targetField.Interface().(WithConfig)
					fmt.Fprintln(w, fieldName)
					fmt.Fprintf(w, " - Enabled\t : %v\n", withConfig.isEnabled())

					if util.ValueHasField(targetField, "Files") {
						fieldFiles := targetField.FieldByName("Files").Interface().([]string)
						newFiles := []string{}
						for _, file := range fieldFiles {
							newFiles = append(newFiles, util.UnsanitisePath(file))
						}
						fmt.Fprintf(w, " - Files\t : %s\n", strings.Join(newFiles, ","))
					}

					if util.ValueHasMethod(targetField, "CustomFields") {
						customMethod := targetField.MethodByName("CustomFields")
						extraFields := customMethod.Call([]reflect.Value{})[0].Interface().(map[string]string)
						for k, v := range extraFields {
							fmt.Fprintf(w, " - %s\t : %s\n", k, v)
						}
					}
					fmt.Fprintln(w)
				}
			}
		}
		fmt.Fprintln(w, "================================================================")
		fmt.Fprintln(w)
		w.Flush()

		input := awaitInput("Write these settings to ~/.proxybastard/config.json? [Yn]", "(y|n|^$)", "")
		fmt.Println()
		readyToWrite = strings.EqualFold(input, "y") || strings.EqualFold(input, "")
	}

	marshalled, err := json.MarshalIndent(actualConfiguration, "", "    ")
	if err != nil {
		log.Fatal(err)
	}
	// fmt.Printf("%s\n\n", string(marshalled))
	if readyToWrite {
		util.WriteSliceToFile(util.SanitisePath("~/.proxybastard/config.json"), strings.Split(string(marshalled), "\n"))
		fmt.Println("Done")
	} else {
		fmt.Println("kthx")
	}
}

func getHighestFrequency(frequencyMap map[string]int) string {
	var highestValue int
	var mostFrequentKey string

	keys := []string{}
	for key := range frequencyMap {
		keys = append(keys, key)
	}
	sort.Strings(keys)

	for _, key := range keys {
		if value := frequencyMap[key]; value > highestValue {
			highestValue = value
			mostFrequentKey = key
		}
	}
	return mostFrequentKey
}

func suggestConfiguration() Configuration {
	var suggestedConfiguration Configuration

	suggestedProxyHosts := make(map[string]int)
	suggestedProxyPorts := make(map[string]int)
	suggestedSOCKSProxyHosts := make(map[string]int)
	suggestedSOCKSProxyPorts := make(map[string]int)
	suggestedNonProxyHosts := mapset.NewSet()

	targetsField := reflect.Indirect(reflect.ValueOf(&TargetsConfiguration{}))
	for i := 0; i < targetsField.NumField(); i++ {

		configurationField := reflect.New(reflect.TypeOf(targetsField.Field(i).Interface()).Elem()).Interface()
		withConfig, hasConfig := configurationField.(WithConfig)
		if hasConfig {
			fieldName := targetsField.Type().Field(i).Name

			if suggestedItemConfiguration := withConfig.suggestConfiguration(); suggestedItemConfiguration != nil {
				if suggestedConfiguration.Targets == nil {
					suggestedConfiguration.Targets = &TargetsConfiguration{}
				}

				addToMap(suggestedProxyHosts, util.SanitiseHTTPProxyURL(suggestedItemConfiguration.ProxyHost))
				addToMap(suggestedProxyPorts, suggestedItemConfiguration.ProxyPort)
				addToMap(suggestedSOCKSProxyHosts, suggestedItemConfiguration.SOCKSProxyHost)
				addToMap(suggestedSOCKSProxyPorts, suggestedItemConfiguration.SOCKSProxyPort)

				for _, nonProxyHost := range suggestedItemConfiguration.NonProxyHosts {
					suggestedNonProxyHosts.Add(nonProxyHost)
				}

				targetsField := reflect.Indirect(reflect.ValueOf(suggestedConfiguration.Targets))
				targetsFieldSuggested := reflect.Indirect(reflect.ValueOf(suggestedItemConfiguration.Targets))

				targetsField.FieldByName(fieldName).Set(targetsFieldSuggested.FieldByName(fieldName))
			}
		}
	}

	suggestedConfiguration.ProxyHost = getHighestFrequency(suggestedProxyHosts)
	suggestedConfiguration.ProxyPort = getHighestFrequency(suggestedProxyPorts)
	suggestedConfiguration.SOCKSProxyHost = getHighestFrequency(suggestedSOCKSProxyHosts)
	suggestedConfiguration.SOCKSProxyPort = getHighestFrequency(suggestedSOCKSProxyPorts)
	for suggestedNonProxyHost := range suggestedNonProxyHosts.Iter() {
		suggestedConfiguration.NonProxyHosts = append(suggestedConfiguration.NonProxyHosts, suggestedNonProxyHost.(string))
	}

	return suggestedConfiguration
}
<|endoftext|>


<commit_before>// Copyright 2015 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"

	"github.com/google/syzkaller/pkg/csource"
	"github.com/google/syzkaller/pkg/log"
	"github.com/google/syzkaller/pkg/mgrconfig"
	"github.com/google/syzkaller/pkg/osutil"
	"github.com/google/syzkaller/pkg/report"
	"github.com/google/syzkaller/pkg/repro"
	"github.com/google/syzkaller/prog"
	"github.com/google/syzkaller/vm"
)

var (
	flagConfig = flag.String("config", "", "manager configuration file (manager.cfg)")
	flagCount  = flag.Int("count", 0, "number of VMs to use (overrides config count param)")
	flagDebug  = flag.Bool("debug", false, "print debug output")
)

func main() {
	os.Args = append(append([]string{}, os.Args[0], "-v=10"), os.Args[1:]...)
	flag.Parse()
	if len(flag.Args()) != 1 || flagConfig == nil {
		log.Fatalf("usage: syz-repro -config=manager.cfg execution.log")
	}
	cfg, err := mgrconfig.LoadFile(*flagConfig)
	if err != nil {
		log.Fatalf("%v: %v", *flagConfig, err)
	}
	logFile := flag.Args()[0]
	data, err := ioutil.ReadFile(logFile)
	if err != nil {
		log.Fatalf("failed to open log file %v: %v", logFile, err)
	}
	if _, err := prog.GetTarget(cfg.TargetOS, cfg.TargetArch); err != nil {
		log.Fatalf("%v", err)
	}
	vmPool, err := vm.Create(cfg, *flagDebug)
	if err != nil {
		log.Fatalf("%v", err)
	}
	vmCount := vmPool.Count()
	if *flagCount > 0 && *flagCount < vmCount {
		vmCount = *flagCount
	}
	if vmCount > 4 {
		vmCount = 4
	}
	vmIndexes := make([]int, vmCount)
	for i := range vmIndexes {
		vmIndexes[i] = i
	}
	reporter, err := report.NewReporter(cfg)
	if err != nil {
		log.Fatalf("%v", err)
	}
	osutil.HandleInterrupts(vm.Shutdown)

	res, stats, err := repro.Run(data, cfg, reporter, vmPool, vmIndexes)
	if err != nil {
		log.Logf(0, "reproduction failed: %v", err)
	}
	if stats != nil {
		fmt.Printf("Extracting prog: %v\n", stats.ExtractProgTime)
		fmt.Printf("Minimizing prog: %v\n", stats.MinimizeProgTime)
		fmt.Printf("Simplifying prog options: %v\n", stats.SimplifyProgTime)
		fmt.Printf("Extracting C: %v\n", stats.ExtractCTime)
		fmt.Printf("Simplifying C: %v\n", stats.SimplifyCTime)
	}
	if res == nil {
		return
	}

	fmt.Printf("opts: %+v crepro: %v\n\n", res.Opts, res.CRepro)
	fmt.Printf("%s\n", res.Prog.Serialize())
	if res.CRepro {
		src, err := csource.Write(res.Prog, res.Opts)
		if err != nil {
			log.Fatalf("failed to generate C repro: %v", err)
		}
		if formatted, err := csource.Format(src); err == nil {
			src = formatted
		}
		fmt.Printf("%s\n", src)
	}
}
<commit_msg>Review<commit_after>// Copyright 2015 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"

	"github.com/google/syzkaller/pkg/csource"
	"github.com/google/syzkaller/pkg/log"
	"github.com/google/syzkaller/pkg/mgrconfig"
	"github.com/google/syzkaller/pkg/osutil"
	"github.com/google/syzkaller/pkg/report"
	"github.com/google/syzkaller/pkg/repro"
	"github.com/google/syzkaller/prog"
	"github.com/google/syzkaller/vm"
)

var (
	flagConfig = flag.String("config", "", "manager configuration file (manager.cfg)")
	flagCount  = flag.Int("count", 0, "number of VMs to use (overrides config count param)")
	flagDebug  = flag.Bool("debug", false, "print debug output")
)

func main() {
	os.Args = append(append([]string{}, os.Args[0], "-v=10"), os.Args[1:]...)
	flag.Parse()
	if len(flag.Args()) != 1 || *flagConfig == "" {
		log.Fatalf("usage: syz-repro -config=manager.cfg execution.log")
	}
	cfg, err := mgrconfig.LoadFile(*flagConfig)
	if err != nil {
		log.Fatalf("%v: %v", *flagConfig, err)
	}
	logFile := flag.Args()[0]
	data, err := ioutil.ReadFile(logFile)
	if err != nil {
		log.Fatalf("failed to open log file %v: %v", logFile, err)
	}
	if _, err := prog.GetTarget(cfg.TargetOS, cfg.TargetArch); err != nil {
		log.Fatalf("%v", err)
	}
	vmPool, err := vm.Create(cfg, *flagDebug)
	if err != nil {
		log.Fatalf("%v", err)
	}
	vmCount := vmPool.Count()
	if *flagCount > 0 && *flagCount < vmCount {
		vmCount = *flagCount
	}
	if vmCount > 4 {
		vmCount = 4
	}
	vmIndexes := make([]int, vmCount)
	for i := range vmIndexes {
		vmIndexes[i] = i
	}
	reporter, err := report.NewReporter(cfg)
	if err != nil {
		log.Fatalf("%v", err)
	}
	osutil.HandleInterrupts(vm.Shutdown)

	res, stats, err := repro.Run(data, cfg, reporter, vmPool, vmIndexes)
	if err != nil {
		log.Logf(0, "reproduction failed: %v", err)
	}
	if stats != nil {
		fmt.Printf("Extracting prog: %v\n", stats.ExtractProgTime)
		fmt.Printf("Minimizing prog: %v\n", stats.MinimizeProgTime)
		fmt.Printf("Simplifying prog options: %v\n", stats.SimplifyProgTime)
		fmt.Printf("Extracting C: %v\n", stats.ExtractCTime)
		fmt.Printf("Simplifying C: %v\n", stats.SimplifyCTime)
	}
	if res == nil {
		return
	}

	fmt.Printf("opts: %+v crepro: %v\n\n", res.Opts, res.CRepro)
	fmt.Printf("%s\n", res.Prog.Serialize())
	if res.CRepro {
		src, err := csource.Write(res.Prog, res.Opts)
		if err != nil {
			log.Fatalf("failed to generate C repro: %v", err)
		}
		if formatted, err := csource.Format(src); err == nil {
			src = formatted
		}
		fmt.Printf("%s\n", src)
	}
}
<|endoftext|>