<commit_before>package server

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"sync"

	"github.com/Sirupsen/logrus"
	"github.com/docker/docker/pkg/registrar"
	"github.com/docker/docker/pkg/truncindex"
	"github.com/kubernetes-incubator/ocid/oci"
	"github.com/kubernetes-incubator/ocid/utils"
	rspec "github.com/opencontainers/runtime-spec/specs-go"
	"github.com/rajatchopra/ocicni"
)

const (
	runtimeAPIVersion = "v1alpha1"
	imageStore        = "/var/lib/ocid/images"
)

// Server implements the RuntimeService and ImageService
type Server struct {
	root         string
	runtime      *oci.Runtime
	sandboxDir   string
	stateLock    sync.Mutex
	state        *serverState
	netPlugin    ocicni.CNIPlugin
	podNameIndex *registrar.Registrar
	podIDIndex   *truncindex.TruncIndex
}

func (s *Server) loadSandbox(id string) error {
	config, err := ioutil.ReadFile(filepath.Join(s.sandboxDir, id, "config.json"))
	if err != nil {
		return err
	}
	var m rspec.Spec
	if err = json.Unmarshal(config, &m); err != nil {
		return err
	}
	labels := make(map[string]string)
	if err = json.Unmarshal([]byte(m.Annotations["ocid/labels"]), &labels); err != nil {
		return err
	}
	name := m.Annotations["ocid/name"]
	name, err = s.reservePodName(id, name)
	if err != nil {
		return err
	}
	s.addSandbox(&sandbox{
		id:         id,
		name:       name,
		logDir:     m.Annotations["ocid/log_path"],
		labels:     labels,
		containers: oci.NewMemoryStore(),
	})
	sandboxPath := filepath.Join(s.sandboxDir, id)
	scontainer, err := oci.NewContainer(m.Annotations["ocid/container_name"], sandboxPath, sandboxPath, labels, id, false)
	if err != nil {
		return err
	}
	s.addContainer(scontainer)
	if err = s.runtime.UpdateStatus(scontainer); err != nil {
		logrus.Warnf("error updating status for container %s: %v", scontainer, err)
	}
	if err = s.podIDIndex.Add(id); err != nil {
		return err
	}
	return nil
}

func (s *Server) restore() error {
	dir, err := ioutil.ReadDir(s.sandboxDir)
	if err != nil {
		return err
	}
	for _, v := range dir {
		if err := s.loadSandbox(v.Name()); err != nil {
			return err
		}
	}
	return nil
}

func (s *Server) reservePodName(id, name string) (string, error) {
	if err := s.podNameIndex.Reserve(name, id); err != nil {
		if err == registrar.ErrNameReserved {
			id, err := s.podNameIndex.Get(name)
			if err != nil {
				logrus.Warnf("name %s already reserved for %s", name, id)
				return "", err
			}
			return "", fmt.Errorf("conflict, name %s already reserved", name)
		}
		return "", fmt.Errorf("error reserving name %s", name)
	}
	return name, nil
}

// New creates a new Server with options provided
func New(runtimePath, root, sandboxDir, containerDir string) (*Server, error) {
	// TODO: This will go away later when we have wrapper process or systemd acting as
	// subreaper.
	if err := utils.SetSubreaper(1); err != nil {
		return nil, fmt.Errorf("failed to set server as subreaper: %v", err)
	}

	utils.StartReaper()

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

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

	r, err := oci.New(runtimePath, containerDir)
	if err != nil {
		return nil, err
	}
	sandboxes := make(map[string]*sandbox)
	containers := oci.NewMemoryStore()
	netPlugin, err := ocicni.InitCNI("")
	if err != nil {
		return nil, err
	}
	s := &Server{
		root:       root,
		runtime:    r,
		netPlugin:  netPlugin,
		sandboxDir: sandboxDir,
		state: &serverState{
			sandboxes:  sandboxes,
			containers: containers,
		},
	}
	s.podIDIndex = truncindex.NewTruncIndex([]string{})
	s.podNameIndex = registrar.NewRegistrar()
	if err := s.restore(); err != nil {
		logrus.Warnf("couldn't restore: %v", err)
	}
	logrus.Debugf("sandboxes: %v", s.state.sandboxes)
	logrus.Debugf("containers: %v", s.state.containers)
	return s, nil
}

type serverState struct {
	sandboxes  map[string]*sandbox
	containers oci.Store
}

func (s *Server) addSandbox(sb *sandbox) {
	s.stateLock.Lock()
	s.state.sandboxes[sb.id] = sb
	s.stateLock.Unlock()
}

func (s *Server) getSandbox(id string) *sandbox {
	s.stateLock.Lock()
	sb := s.state.sandboxes[id]
	s.stateLock.Unlock()
	return sb
}

func (s *Server) hasSandbox(id string) bool {
	s.stateLock.Lock()
	_, ok := s.state.sandboxes[id]
	s.stateLock.Unlock()
	return ok
}

func (s *Server) addContainer(c *oci.Container) {
	s.stateLock.Lock()
	sandbox := s.state.sandboxes[c.Sandbox()]
	// TODO(runcom): handle !ok above!!! otherwise it panics!
	sandbox.addContainer(c)
	s.state.containers.Add(c.Name(), c)
	s.stateLock.Unlock()
}

func (s *Server) getContainer(name string) *oci.Container {
	s.stateLock.Lock()
	c := s.state.containers.Get(name)
	s.stateLock.Unlock()
	return c
}

func (s *Server) removeContainer(c *oci.Container) {
	s.stateLock.Lock()
	sandbox := s.state.sandboxes[c.Sandbox()]
	sandbox.removeContainer(c)
	s.state.containers.Delete(c.Name())
	s.stateLock.Unlock()
}
<commit_msg>Add API to release Pod name when not required<commit_after>package server

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"sync"

	"github.com/Sirupsen/logrus"
	"github.com/docker/docker/pkg/registrar"
	"github.com/docker/docker/pkg/truncindex"
	"github.com/kubernetes-incubator/ocid/oci"
	"github.com/kubernetes-incubator/ocid/utils"
	rspec "github.com/opencontainers/runtime-spec/specs-go"
	"github.com/rajatchopra/ocicni"
)

const (
	runtimeAPIVersion = "v1alpha1"
	imageStore        = "/var/lib/ocid/images"
)

// Server implements the RuntimeService and ImageService
type Server struct {
	root         string
	runtime      *oci.Runtime
	sandboxDir   string
	stateLock    sync.Mutex
	state        *serverState
	netPlugin    ocicni.CNIPlugin
	podNameIndex *registrar.Registrar
	podIDIndex   *truncindex.TruncIndex
}

func (s *Server) loadSandbox(id string) error {
	config, err := ioutil.ReadFile(filepath.Join(s.sandboxDir, id, "config.json"))
	if err != nil {
		return err
	}
	var m rspec.Spec
	if err = json.Unmarshal(config, &m); err != nil {
		return err
	}
	labels := make(map[string]string)
	if err = json.Unmarshal([]byte(m.Annotations["ocid/labels"]), &labels); err != nil {
		return err
	}
	name := m.Annotations["ocid/name"]
	name, err = s.reservePodName(id, name)
	if err != nil {
		return err
	}
	s.addSandbox(&sandbox{
		id:         id,
		name:       name,
		logDir:     m.Annotations["ocid/log_path"],
		labels:     labels,
		containers: oci.NewMemoryStore(),
	})
	sandboxPath := filepath.Join(s.sandboxDir, id)
	scontainer, err := oci.NewContainer(m.Annotations["ocid/container_name"], sandboxPath, sandboxPath, labels, id, false)
	if err != nil {
		return err
	}
	s.addContainer(scontainer)
	if err = s.runtime.UpdateStatus(scontainer); err != nil {
		logrus.Warnf("error updating status for container %s: %v", scontainer, err)
	}
	if err = s.podIDIndex.Add(id); err != nil {
		return err
	}
	return nil
}

func (s *Server) restore() error {
	dir, err := ioutil.ReadDir(s.sandboxDir)
	if err != nil {
		return err
	}
	for _, v := range dir {
		if err := s.loadSandbox(v.Name()); err != nil {
			return err
		}
	}
	return nil
}

func (s *Server) reservePodName(id, name string) (string, error) {
	if err := s.podNameIndex.Reserve(name, id); err != nil {
		if err == registrar.ErrNameReserved {
			id, err := s.podNameIndex.Get(name)
			if err != nil {
				logrus.Warnf("name %s already reserved for %s", name, id)
				return "", err
			}
			return "", fmt.Errorf("conflict, name %s already reserved", name)
		}
		return "", fmt.Errorf("error reserving name %s", name)
	}
	return name, nil
}

func (s *Server) releasePodName(name string) {
	s.podNameIndex.Release(name)
}

// New creates a new Server with options provided
func New(runtimePath, root, sandboxDir, containerDir string) (*Server, error) {
	// TODO: This will go away later when we have wrapper process or systemd acting as
	// subreaper.
	if err := utils.SetSubreaper(1); err != nil {
		return nil, fmt.Errorf("failed to set server as subreaper: %v", err)
	}

	utils.StartReaper()

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

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

	r, err := oci.New(runtimePath, containerDir)
	if err != nil {
		return nil, err
	}
	sandboxes := make(map[string]*sandbox)
	containers := oci.NewMemoryStore()
	netPlugin, err := ocicni.InitCNI("")
	if err != nil {
		return nil, err
	}
	s := &Server{
		root:       root,
		runtime:    r,
		netPlugin:  netPlugin,
		sandboxDir: sandboxDir,
		state: &serverState{
			sandboxes:  sandboxes,
			containers: containers,
		},
	}
	s.podIDIndex = truncindex.NewTruncIndex([]string{})
	s.podNameIndex = registrar.NewRegistrar()
	if err := s.restore(); err != nil {
		logrus.Warnf("couldn't restore: %v", err)
	}
	logrus.Debugf("sandboxes: %v", s.state.sandboxes)
	logrus.Debugf("containers: %v", s.state.containers)
	return s, nil
}

type serverState struct {
	sandboxes  map[string]*sandbox
	containers oci.Store
}

func (s *Server) addSandbox(sb *sandbox) {
	s.stateLock.Lock()
	s.state.sandboxes[sb.id] = sb
	s.stateLock.Unlock()
}

func (s *Server) getSandbox(id string) *sandbox {
	s.stateLock.Lock()
	sb := s.state.sandboxes[id]
	s.stateLock.Unlock()
	return sb
}

func (s *Server) hasSandbox(id string) bool {
	s.stateLock.Lock()
	_, ok := s.state.sandboxes[id]
	s.stateLock.Unlock()
	return ok
}

func (s *Server) addContainer(c *oci.Container) {
	s.stateLock.Lock()
	sandbox := s.state.sandboxes[c.Sandbox()]
	// TODO(runcom): handle !ok above!!! otherwise it panics!
	sandbox.addContainer(c)
	s.state.containers.Add(c.Name(), c)
	s.stateLock.Unlock()
}

func (s *Server) getContainer(name string) *oci.Container {
	s.stateLock.Lock()
	c := s.state.containers.Get(name)
	s.stateLock.Unlock()
	return c
}

func (s *Server) removeContainer(c *oci.Container) {
	s.stateLock.Lock()
	sandbox := s.state.sandboxes[c.Sandbox()]
	sandbox.removeContainer(c)
	s.state.containers.Delete(c.Name())
	s.stateLock.Unlock()
}
<|endoftext|>


<commit_before>package queue

import (
	"bytes"
	"crypto/rand"
	"fmt"
	"github.com/FoundationDB/fdb-go/fdb"
	"github.com/FoundationDB/fdb-go/fdb/tuple"
	"github.com/happypancake/go-layers/subspace"
	"time"
)

type Queue struct {
	Subspace       subspace.Subspace
	HighContention bool
	conflictedPop  subspace.Subspace
	conflictedItem subspace.Subspace
	queueItem      subspace.Subspace
}

func New(sub subspace.Subspace, highContention bool) Queue {

	conflict := sub.Item(tuple.Tuple{"conflict"})
	pop := sub.Item(tuple.Tuple{"pop"})
	item := sub.Item(tuple.Tuple{"item"})

	return Queue{sub, highContention, pop, conflict, item}
}

func (queue *Queue) Clear(tr fdb.Transaction) {
	tr.ClearRange(queue.Subspace.FullRange())
}

func (queue *Queue) Peek(tr fdb.Transaction) (value []byte, ok bool) {
	if val, ok := queue.getFirstItem(tr); ok {
		return decodeValue(val.Value), true
	}
	return

}

func decodeValue(val []byte) []byte {
	if t, err := tuple.Unpack(val); err != nil {
		panic(err)
	} else {
		return t[0].([]byte)
	}

}
func encodeValue(value []byte) []byte {
	return tuple.Tuple{value}.Pack()
}

type KeyReader interface {
	GetKey(key fdb.Selectable) fdb.FutureKey
}

// to make private
func (queue *Queue) GetNextIndex(tr KeyReader, sub subspace.Subspace) int64 {

	r := sub.Range(tuple.Tuple{})

	key := tr.GetKey(fdb.LastLessThan(r.End)).GetOrPanic()

	if i := bytes.Compare(key, []byte(r.BeginKey())); i < 0 {
		return 0
	}

	if t, err := sub.Unpack(key); err != nil {
		panic("Failed to unpack key")
	} else {
		return t[0].(int64) + 1
	}
}

func (queue *Queue) GetNextQueueIndex(tr fdb.Transaction) int64 {
	return queue.GetNextIndex(tr.Snapshot(), queue.queueItem)
}
func (queue *Queue) Push(tr fdb.Transaction, value []byte) {
	snap := tr.Snapshot()
	index := queue.GetNextIndex(snap, queue.queueItem)
	queue.pushAt(tr, value, index)
}

// Pop the next item from the queue. Cannot be composed with other functions in a single transaction.
func (queue *Queue) Pop(db fdb.Database) (value []byte, ok bool) {

	if queue.HighContention {
		if result, ok := queue.popHighContention(db); ok {
			return decodeValue(result), true
		}
	} else {

		val, _ := db.Transact(func(tr fdb.Transaction) (interface{}, error) {
			if result, ok := queue.popSimple(tr); ok {
				return decodeValue(result), nil
			}
			return nil, nil

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

	}
	return
}

func (queue *Queue) pushAt(tr fdb.Transaction, value []byte, index int64) {
	key := queue.queueItem.Pack(tuple.Tuple{index, nextRandom()})
	val := encodeValue(value)

	tr.Set(fdb.Key(key), val)
}

// popSimple does not attempt to avoid conflicts
// if many clients are trying to pop simultaneously, only one will be able to succeed at a time.
func (queue *Queue) popSimple(tr fdb.Transaction) (value []byte, ok bool) {
	if kv, ok := queue.getFirstItem(tr); ok {
		tr.Clear(kv.Key)
		return kv.Value, true
	}

	return
}

func (queue *Queue) addConflictedPop(tr fdb.Transaction, forced bool) (val []byte, ok bool) {
	index := queue.GetNextIndex(tr.Snapshot(), queue.conflictedPop)

	if index == 0 && !forced {
		return nil, false
	}
	key := queue.conflictedPop.Pack(tuple.Tuple{index, nextRandom()})
	//read := tr.Get(key)
	tr.Set(fdb.Key(key), []byte(""))
	return key, true
}

func (queue *Queue) popSimpleOrRegisterWaitKey(tr fdb.Transaction) (value []byte, waitKey []byte) {

	// TODO: deal with FDB error in defer

	// Check if there are other people waiting to be popped. If so, we
	// cannot pop before them.
	if key, ok := queue.addConflictedPop(tr, false); ok {
		tr.Commit().BlockUntilReady()
		return nil, key
	} else {
		// No one else was waiting to be popped
		value, ok = queue.popSimple(tr)
		tr.Commit().BlockUntilReady()
		return value, nil
	}
}

// popHighContention attempts to avoid collisions by registering
// itself in a semi-ordered set of poppers if it doesn't initially succeed.
// It then enters a polling loop where it attempts to fulfill outstanding pops
// and then checks to see if it has been fulfilled.
func (queue *Queue) popHighContention(db fdb.Database) (value []byte, ok bool) {
	//panic("Not implemented")
	backoff := 0.01

	tr, err := db.CreateTransaction()

	if err != nil {
		panic(err)
	}

	value, waitKey := queue.popSimpleOrRegisterWaitKey(tr)
	if value != nil {
		return value, true
	}

	randId := queue.conflictedPop.MustUnpack(waitKey)[1].([]byte)
	// The result of the pop will be stored at this key once it has been fulfilled
	resultKey := queue.conflictedItemKey(randId)

	tr.Reset()

	for {
		for done := queue.fulfilConflictedPops(db); !done; {

		}

		tr.Reset()
		value := tr.Get(fdb.Key(waitKey))
		result := tr.Get(fdb.Key(resultKey))

		// If waitKey is present, then we have not been fulfilled
		if value.IsReady() {
			time.Sleep(time.Duration(backoff) * time.Second)
			backoff = backoff * 2
			if backoff > 1 {
				backoff = 1
			}
			continue
		}
		if !result.IsReady() {
			return nil, false
		}
		tr.Clear(fdb.Key(resultKey))
		tr.Commit().BlockUntilReady()

		return result.GetOrPanic(), true

	}

	return nil, false
}

func (queue *Queue) getWaitingPops(tr fdb.Transaction, numPops int) fdb.RangeResult {
	r := queue.conflictedPop.FullRange()
	return tr.GetRange(r, fdb.RangeOptions{Limit: numPops})
}

func (queue *Queue) getItems(tr fdb.Transaction, numPops int) fdb.RangeResult {
	r := queue.queueItem.FullRange()
	return tr.GetRange(r, fdb.RangeOptions{Limit: numPops})
}

func minLength(a, b []fdb.KeyValue) int {
	if al, bl := len(a), len(b); al < bl {
		return al
	} else {
		return bl
	}

}

func (queue *Queue) conflictedItemKey(subkey []byte) []byte {
	return queue.conflictedItem.Pack(tuple.Tuple{subkey})
}

func (queue *Queue) fulfilConflictedPops(db fdb.Database) bool {
	numPops := 100

	v, err := db.Transact(func(tr fdb.Transaction) (interface{}, error) {
		pops := queue.getWaitingPops(tr, numPops).GetSliceOrPanic()
		items := queue.getItems(tr, numPops).GetSliceOrPanic()

		min := minLength(pops, items)

		for i := 0; i < min; i++ {
			pop, k, v := pops[i], items[i].Key, items[i].Value

			key := queue.conflictedPop.MustUnpack(pop.Key)
			storageKey := queue.conflictedItemKey(key[0].([]byte))
			tr.Set(fdb.Key(storageKey), v)
			_ = tr.Get(k)
			_ = tr.Get(pop.Key)
			tr.Clear(pop.Key)
			tr.Clear(k)
		}

		for _, pop := range pops[min:] {
			_ = tr.Get(pop.Key)
			tr.Clear(pop.Key)
		}

		return len(pops) < numPops, nil

	})

	if err != nil {
		panic(err)
	}

	return v.(bool)
}

func nextRandom() []byte {
	b := make([]byte, 20)
	if _, err := rand.Read(b); err == nil {
		return b
	} else {
		fmt.Println("Panic", err)
		panic(err)
	}
}

func (queue *Queue) Empty(tr fdb.Transaction) bool {
	_, ok := queue.getFirstItem(tr)
	return ok == false
}

func (queue *Queue) getFirstItem(tr fdb.Transaction) (kv fdb.KeyValue, ok bool) {
	r := queue.queueItem.FullRange()
	opt := fdb.RangeOptions{Limit: 1}

	if kvs := tr.GetRange(r, opt).GetSliceOrPanic(); len(kvs) == 1 {
		return kvs[0], true
	}
	return
}
<commit_msg>Queues finally work in high contention mode<commit_after>package queue

import (
	"bytes"
	"crypto/rand"
	"fmt"
	"github.com/FoundationDB/fdb-go/fdb"
	"github.com/FoundationDB/fdb-go/fdb/tuple"
	"github.com/happypancake/go-layers/subspace"
	"time"
)

type Queue struct {
	Subspace       subspace.Subspace
	HighContention bool
	conflictedPop  subspace.Subspace // stores int64 index, randId []byte
	conflictedItem subspace.Subspace
	queueItem      subspace.Subspace
}

func New(sub subspace.Subspace, highContention bool) Queue {

	conflict := sub.Item(tuple.Tuple{"conflict"})
	pop := sub.Item(tuple.Tuple{"pop"})
	item := sub.Item(tuple.Tuple{"item"})

	return Queue{sub, highContention, pop, conflict, item}
}

func (queue *Queue) Clear(tr fdb.Transaction) {

	tr.ClearRange(queue.Subspace.FullRange())
}

func (queue *Queue) Peek(tr fdb.Transaction) (value []byte, ok bool) {
	if val, ok := queue.getFirstItem(tr); ok {
		return decodeValue(val.Value), true
	}
	return

}

func decodeValue(val []byte) []byte {
	if t, err := tuple.Unpack(val); err != nil {
		panic(err)
	} else {
		return t[0].([]byte)
	}

}
func encodeValue(value []byte) []byte {
	return tuple.Tuple{value}.Pack()
}

type KeyReader interface {
	GetKey(key fdb.Selectable) fdb.FutureKey
}

// to make private
func (queue *Queue) GetNextIndex(tr KeyReader, sub subspace.Subspace) int64 {

	r := sub.Range(tuple.Tuple{})

	key := tr.GetKey(fdb.LastLessThan(r.End)).GetOrPanic()

	if i := bytes.Compare(key, []byte(r.BeginKey())); i < 0 {
		return 0
	}

	if t, err := sub.Unpack(key); err != nil {
		panic("Failed to unpack key")
	} else {
		return t[0].(int64) + 1
	}
}

func (queue *Queue) GetNextQueueIndex(tr fdb.Transaction) int64 {
	return queue.GetNextIndex(tr.Snapshot(), queue.queueItem)
}
func (queue *Queue) Push(tr fdb.Transaction, value []byte) {
	snap := tr.Snapshot()
	index := queue.GetNextIndex(snap, queue.queueItem)
	queue.pushAt(tr, value, index)
}

// Pop the next item from the queue. Cannot be composed with other functions in a single transaction.
func (queue *Queue) Pop(db fdb.Database) (value []byte, ok bool) {

	if queue.HighContention {
		if result, ok := queue.popHighContention(db); ok {
			return decodeValue(result), true
		}
	} else {

		val, _ := db.Transact(func(tr fdb.Transaction) (interface{}, error) {
			if result, ok := queue.popSimple(tr); ok {
				return decodeValue(result), nil
			}
			return nil, nil

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

	}
	return
}

func (queue *Queue) pushAt(tr fdb.Transaction, value []byte, index int64) {
	key := queue.queueItem.Pack(tuple.Tuple{index, nextRandom()})
	val := encodeValue(value)

	tr.Set(fdb.Key(key), val)
}

// popSimple does not attempt to avoid conflicts
// if many clients are trying to pop simultaneously, only one will be able to succeed at a time.
func (queue *Queue) popSimple(tr fdb.Transaction) (value []byte, ok bool) {
	if kv, ok := queue.getFirstItem(tr); ok {
		tr.Clear(kv.Key)
		return kv.Value, true
	}

	return
}

func (queue *Queue) addConflictedPop(tr fdb.Transaction, forced bool) (val []byte) {
	index := queue.GetNextIndex(tr.Snapshot(), queue.conflictedPop)

	if (index == 0) && (!forced) {
		return nil
	}
	key := queue.conflictedPop.Pack(tuple.Tuple{index, nextRandom()})
	// why do we read no
	_ = tr.Get(fdb.Key(key))
	tr.Set(fdb.Key(key), []byte(""))
	return key
}

// popHighContention attempts to avoid collisions by registering
// itself in a semi-ordered set of poppers if it doesn't initially succeed.
// It then enters a polling loop where it attempts to fulfill outstanding pops
// and then checks to see if it has been fulfilled.
func (queue *Queue) popHighContention(db fdb.Database) (value []byte, ok bool) {
	//panic("Not implemented")
	backoff := 0.01

	tr, err := db.CreateTransaction()

	if err != nil {
		panic(err)
	}

	// Check if there are other people waiting to be popped. If so, we
	// cannot pop before them.

	waitKey := queue.addConflictedPop(tr, false)
	if waitKey == nil {
		value, ok := queue.popSimple(tr)

		// if we managed to commit without collisions
		if err := tr.Commit().GetWithError(); err == nil {

			return value, ok
		}
	}

	if err := tr.Commit().GetWithError(); err != nil {
		fmt.Println("Panic in #", err)
	}

	if waitKey == nil {
		waitKey = queue.addConflictedPop(tr, true)
	}

	randId := queue.conflictedPop.MustUnpack(waitKey)[1].([]byte)
	// The result of the pop will be stored at this key once it has been fulfilled
	resultKey := queue.conflictedItemKey(randId)

	tr.Reset()

	for {
		for done := queue.fulfilConflictedPops(db); !done; {

		}

		tr.Reset()
		value := tr.Get(fdb.Key(waitKey))
		result := tr.Get(fdb.Key(resultKey))

		// If waitKey is present, then we have not been fulfilled
		if value.IsReady() {
			time.Sleep(time.Duration(backoff) * time.Second)
			backoff = backoff * 2
			if backoff > 1 {
				backoff = 1
			}
			continue
		}
		if !result.IsReady() {
			return nil, false
		}
		tr.Clear(fdb.Key(resultKey))
		tr.Commit().BlockUntilReady()

		return result.GetOrPanic(), true

	}

	return nil, false
}

func (queue *Queue) getWaitingPops(tr fdb.Transaction, numPops int) fdb.RangeResult {
	r := queue.conflictedPop.FullRange()
	return tr.GetRange(r, fdb.RangeOptions{Limit: numPops})
}

func (queue *Queue) getItems(tr fdb.Transaction, numPops int) fdb.RangeResult {
	r := queue.queueItem.FullRange()
	return tr.GetRange(r, fdb.RangeOptions{Limit: numPops})
}

func minLength(a, b []fdb.KeyValue) int {
	if al, bl := len(a), len(b); al < bl {
		return al
	} else {
		return bl
	}

}

func (queue *Queue) conflictedItemKey(subkey []byte) []byte {
	return queue.conflictedItem.Pack(tuple.Tuple{subkey})
}

func (queue *Queue) fulfilConflictedPops(db fdb.Database) bool {
	numPops := 100

	v, err := db.Transact(func(tr fdb.Transaction) (interface{}, error) {
		pops := queue.getWaitingPops(tr, numPops).GetSliceOrPanic()
		items := queue.getItems(tr, numPops).GetSliceOrPanic()

		min := minLength(pops, items)

		for i := 0; i < min; i++ {
			pop, k, v := pops[i], items[i].Key, items[i].Value

			key := queue.conflictedPop.MustUnpack(pop.Key)
			storageKey := queue.conflictedItemKey(key[1].([]byte))
			tr.Set(fdb.Key(storageKey), v)
			_ = tr.Get(k)
			_ = tr.Get(pop.Key)
			tr.Clear(pop.Key)
			tr.Clear(k)
		}

		for _, pop := range pops[min:] {
			_ = tr.Get(pop.Key)
			tr.Clear(pop.Key)
		}

		return len(pops) < numPops, nil

	})

	if err != nil {
		panic(err)
	}

	return v.(bool)
}

func nextRandom() []byte {
	b := make([]byte, 20)
	if _, err := rand.Read(b); err == nil {
		return b
	} else {

		panic(err)
	}
}

func (queue *Queue) Empty(tr fdb.Transaction) bool {
	_, ok := queue.getFirstItem(tr)
	return ok == false
}

func (queue *Queue) getFirstItem(tr fdb.Transaction) (kv fdb.KeyValue, ok bool) {
	r := queue.queueItem.FullRange()
	opt := fdb.RangeOptions{Limit: 1}

	if kvs := tr.GetRange(r, opt).GetSliceOrPanic(); len(kvs) == 1 {
		return kvs[0], true
	}
	return
}
<|endoftext|>


<commit_before>/*
Copyright 2021 The Rook Authors. All rights reserved.

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

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

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

package v1

import (
	_ "embed"
	"fmt"

	secv1 "github.com/openshift/api/security/v1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// NewSecurityContextConstraints returns a new SecurityContextConstraints for Rook-Ceph to run on
// OpenShift.
func NewSecurityContextConstraints(name, namespace string) *secv1.SecurityContextConstraints {
	return &secv1.SecurityContextConstraints{
		TypeMeta: metav1.TypeMeta{
			APIVersion: "security.openshift.io/v1",
			Kind:       "SecurityContextConstraints",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: namespace,
		},
		AllowPrivilegedContainer: true,
		AllowHostDirVolumePlugin: true,
		ReadOnlyRootFilesystem:   false,
		AllowHostIPC:             true,
		AllowHostNetwork:         false,
		AllowHostPorts:           false,
		RequiredDropCapabilities: []corev1.Capability{},
		DefaultAddCapabilities:   []corev1.Capability{},
		RunAsUser: secv1.RunAsUserStrategyOptions{
			Type: secv1.RunAsUserStrategyRunAsAny,
		},
		SELinuxContext: secv1.SELinuxContextStrategyOptions{
			Type: secv1.SELinuxStrategyMustRunAs,
		},
		FSGroup: secv1.FSGroupStrategyOptions{
			Type: secv1.FSGroupStrategyMustRunAs,
		},
		SupplementalGroups: secv1.SupplementalGroupsStrategyOptions{
			Type: secv1.SupplementalGroupsStrategyRunAsAny,
		},
		Volumes: []secv1.FSType{
			secv1.FSTypeConfigMap,
			secv1.FSTypeDownwardAPI,
			secv1.FSTypeEmptyDir,
			secv1.FSTypeHostPath,
			secv1.FSTypePersistentVolumeClaim,
			secv1.FSProjected,
			secv1.FSTypeSecret,
		},
		Users: []string{
			fmt.Sprintf("system:serviceaccount:%s:rook-ceph-system", namespace),
			fmt.Sprintf("system:serviceaccount:%s:default", namespace),
			fmt.Sprintf("system:serviceaccount:%s:rook-ceph-mgr", namespace),
			fmt.Sprintf("system:serviceaccount:%s:rook-ceph-osd", namespace),
		},
	}
}
<commit_msg>core: remove unused imports<commit_after>/*
Copyright 2021 The Rook Authors. All rights reserved.

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

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

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

package v1

import (
	"fmt"

	secv1 "github.com/openshift/api/security/v1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// NewSecurityContextConstraints returns a new SecurityContextConstraints for Rook-Ceph to run on
// OpenShift.
func NewSecurityContextConstraints(name, namespace string) *secv1.SecurityContextConstraints {
	return &secv1.SecurityContextConstraints{
		TypeMeta: metav1.TypeMeta{
			APIVersion: "security.openshift.io/v1",
			Kind:       "SecurityContextConstraints",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: namespace,
		},
		AllowPrivilegedContainer: true,
		AllowHostDirVolumePlugin: true,
		ReadOnlyRootFilesystem:   false,
		AllowHostIPC:             true,
		AllowHostNetwork:         false,
		AllowHostPorts:           false,
		RequiredDropCapabilities: []corev1.Capability{},
		DefaultAddCapabilities:   []corev1.Capability{},
		RunAsUser: secv1.RunAsUserStrategyOptions{
			Type: secv1.RunAsUserStrategyRunAsAny,
		},
		SELinuxContext: secv1.SELinuxContextStrategyOptions{
			Type: secv1.SELinuxStrategyMustRunAs,
		},
		FSGroup: secv1.FSGroupStrategyOptions{
			Type: secv1.FSGroupStrategyMustRunAs,
		},
		SupplementalGroups: secv1.SupplementalGroupsStrategyOptions{
			Type: secv1.SupplementalGroupsStrategyRunAsAny,
		},
		Volumes: []secv1.FSType{
			secv1.FSTypeConfigMap,
			secv1.FSTypeDownwardAPI,
			secv1.FSTypeEmptyDir,
			secv1.FSTypeHostPath,
			secv1.FSTypePersistentVolumeClaim,
			secv1.FSProjected,
			secv1.FSTypeSecret,
		},
		Users: []string{
			fmt.Sprintf("system:serviceaccount:%s:rook-ceph-system", namespace),
			fmt.Sprintf("system:serviceaccount:%s:default", namespace),
			fmt.Sprintf("system:serviceaccount:%s:rook-ceph-mgr", namespace),
			fmt.Sprintf("system:serviceaccount:%s:rook-ceph-osd", namespace),
		},
	}
}
<|endoftext|>


<commit_before>package queue

import (
	"net/url"
	"sync"
	"sync/atomic"

	"github.com/gocolly/colly/v2"
)

const stop = true

// Storage is the interface of the queue's storage backend
type Storage interface {
	// Init initializes the storage
	Init() error
	// AddRequest adds a serialized request to the queue
	AddRequest([]byte) error
	// GetRequest pops the next request from the queue
	// or returns error if the queue is empty
	GetRequest() ([]byte, error)
	// QueueSize returns with the size of the queue
	QueueSize() (int, error)
}

// Queue is a request queue which uses a Collector to consume
// requests in multiple threads
type Queue struct {
	// Threads defines the number of consumer threads
	Threads           int
	storage           Storage
	activeThreadCount int32
	requestsOut       chan *colly.Request
}

// InMemoryQueueStorage is the default implementation of the Storage interface.
// InMemoryQueueStorage holds the request queue in memory.
type InMemoryQueueStorage struct {
	// MaxSize defines the capacity of the queue.
	// New requests are discarded if the queue size reaches MaxSize
	MaxSize int
	lock    *sync.RWMutex
	size    int
	first   *inMemoryQueueItem
	last    *inMemoryQueueItem
}

type inMemoryQueueItem struct {
	Request []byte
	Next    *inMemoryQueueItem
}

// New creates a new queue with a Storage specified in argument
// A standard InMemoryQueueStorage is used if Storage argument is nil.
func New(threads int, s Storage) (*Queue, error) {
	if s == nil {
		s = &InMemoryQueueStorage{MaxSize: 100000}
	}
	if err := s.Init(); err != nil {
		return nil, err
	}
	return &Queue{
		Threads:     threads,
		storage:     s,
		requestsOut: make(chan *colly.Request),
	}, nil
}

// IsEmpty returns true if the queue is empty
func (q *Queue) IsEmpty() bool {
	s, _ := q.Size()
	return s == 0
}

// AddURL adds a new URL to the queue
func (q *Queue) AddURL(URL string) error {
	u, err := url.Parse(URL)
	if err != nil {
		return err
	}
	r := &colly.Request{
		URL:    u,
		Method: "GET",
	}
	d, err := r.Marshal()
	if err != nil {
		return err
	}
	return q.storage.AddRequest(d)
}

// AddRequest adds a new Request to the queue
func (q *Queue) AddRequest(r *colly.Request) error {
	d, err := r.Marshal()
	if err != nil {
		return err
	}
	return q.storage.AddRequest(d)
}

// Size returns the size of the queue
func (q *Queue) Size() (int, error) {
	return q.storage.QueueSize()
}

// Run starts consumer threads and calls the Collector
// to perform requests. Run blocks while the queue has active requests
func (q *Queue) Run(c *colly.Collector) error {
	wg := &sync.WaitGroup{}
	for i := 0; i < q.Threads; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			atomic.AddInt32(&q.activeThreadCount, 1)
			for r := range q.requestsOut {
				r.Do()
			}
			atomic.AddInt32(&q.activeThreadCount, -1)
		}()
	}

	wg.Add(1)
	go func(c *colly.Collector, s Storage) {
		defer wg.Done()
		for {
			if q.IsEmpty() {
				if q.activeThreadCount == 0 {
					q.finish()
					break
				}
				continue
			}
			rb, err := s.GetRequest()
			if err != nil || rb == nil {
				//q.finish()
				break
			}
			t := make([]byte, len(rb))
			copy(t, rb)
			r, err := c.UnmarshalRequest(t[:])
			if err != nil || r == nil {
				continue
			}
			q.requestsOut <- r
		}
	}(c, q.storage)

	wg.Wait()
	return nil
}

func (q *Queue) finish() {
	close(q.requestsOut)
}

// Init implements Storage.Init() function
func (q *InMemoryQueueStorage) Init() error {
	q.lock = &sync.RWMutex{}
	return nil
}

// AddRequest implements Storage.AddRequest() function
func (q *InMemoryQueueStorage) AddRequest(r []byte) error {
	q.lock.Lock()
	defer q.lock.Unlock()
	// Discard URLs if size limit exceeded
	if q.MaxSize > 0 && q.size >= q.MaxSize {
		return nil
	}
	i := &inMemoryQueueItem{Request: r}
	if q.first == nil {
		q.first = i
	} else {
		q.last.Next = i
	}
	q.last = i
	q.size++
	return nil
}

// GetRequest implements Storage.GetRequest() function
func (q *InMemoryQueueStorage) GetRequest() ([]byte, error) {
	q.lock.Lock()
	defer q.lock.Unlock()
	if q.size == 0 {
		return nil, nil
	}
	r := q.first.Request
	q.first = q.first.Next
	q.size--
	return r, nil
}

// QueueSize implements Storage.QueueSize() function
func (q *InMemoryQueueStorage) QueueSize() (int, error) {
	q.lock.Lock()
	defer q.lock.Unlock()
	return q.size, nil
}
<commit_msg>Return ErrQueueFull when q.size >= q.MaxSize<commit_after>package queue

import (
	"net/url"
	"sync"
	"sync/atomic"

	"github.com/gocolly/colly/v2"
)

const stop = true

// Storage is the interface of the queue's storage backend
type Storage interface {
	// Init initializes the storage
	Init() error
	// AddRequest adds a serialized request to the queue
	AddRequest([]byte) error
	// GetRequest pops the next request from the queue
	// or returns error if the queue is empty
	GetRequest() ([]byte, error)
	// QueueSize returns with the size of the queue
	QueueSize() (int, error)
}

// Queue is a request queue which uses a Collector to consume
// requests in multiple threads
type Queue struct {
	// Threads defines the number of consumer threads
	Threads           int
	storage           Storage
	activeThreadCount int32
	requestsOut       chan *colly.Request
}

// InMemoryQueueStorage is the default implementation of the Storage interface.
// InMemoryQueueStorage holds the request queue in memory.
type InMemoryQueueStorage struct {
	// MaxSize defines the capacity of the queue.
	// New requests are discarded if the queue size reaches MaxSize
	MaxSize int
	lock    *sync.RWMutex
	size    int
	first   *inMemoryQueueItem
	last    *inMemoryQueueItem
}

type inMemoryQueueItem struct {
	Request []byte
	Next    *inMemoryQueueItem
}

// New creates a new queue with a Storage specified in argument
// A standard InMemoryQueueStorage is used if Storage argument is nil.
func New(threads int, s Storage) (*Queue, error) {
	if s == nil {
		s = &InMemoryQueueStorage{MaxSize: 100000}
	}
	if err := s.Init(); err != nil {
		return nil, err
	}
	return &Queue{
		Threads:     threads,
		storage:     s,
		requestsOut: make(chan *colly.Request),
	}, nil
}

// IsEmpty returns true if the queue is empty
func (q *Queue) IsEmpty() bool {
	s, _ := q.Size()
	return s == 0
}

// AddURL adds a new URL to the queue
func (q *Queue) AddURL(URL string) error {
	u, err := url.Parse(URL)
	if err != nil {
		return err
	}
	r := &colly.Request{
		URL:    u,
		Method: "GET",
	}
	d, err := r.Marshal()
	if err != nil {
		return err
	}
	return q.storage.AddRequest(d)
}

// AddRequest adds a new Request to the queue
func (q *Queue) AddRequest(r *colly.Request) error {
	d, err := r.Marshal()
	if err != nil {
		return err
	}
	return q.storage.AddRequest(d)
}

// Size returns the size of the queue
func (q *Queue) Size() (int, error) {
	return q.storage.QueueSize()
}

// Run starts consumer threads and calls the Collector
// to perform requests. Run blocks while the queue has active requests
func (q *Queue) Run(c *colly.Collector) error {
	wg := &sync.WaitGroup{}
	for i := 0; i < q.Threads; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			atomic.AddInt32(&q.activeThreadCount, 1)
			for r := range q.requestsOut {
				r.Do()
			}
			atomic.AddInt32(&q.activeThreadCount, -1)
		}()
	}

	wg.Add(1)
	go func(c *colly.Collector, s Storage) {
		defer wg.Done()
		for {
			if q.IsEmpty() {
				if q.activeThreadCount == 0 {
					q.finish()
					break
				}
				continue
			}
			rb, err := s.GetRequest()
			if err != nil || rb == nil {
				//q.finish()
				break
			}
			t := make([]byte, len(rb))
			copy(t, rb)
			r, err := c.UnmarshalRequest(t[:])
			if err != nil || r == nil {
				continue
			}
			q.requestsOut <- r
		}
	}(c, q.storage)

	wg.Wait()
	return nil
}

func (q *Queue) finish() {
	close(q.requestsOut)
}

// Init implements Storage.Init() function
func (q *InMemoryQueueStorage) Init() error {
	q.lock = &sync.RWMutex{}
	return nil
}

// AddRequest implements Storage.AddRequest() function
func (q *InMemoryQueueStorage) AddRequest(r []byte) error {
	q.lock.Lock()
	defer q.lock.Unlock()
	// Discard URLs if size limit exceeded
	if q.MaxSize > 0 && q.size >= q.MaxSize {
		return colly.ErrQueueFull
	}
	i := &inMemoryQueueItem{Request: r}
	if q.first == nil {
		q.first = i
	} else {
		q.last.Next = i
	}
	q.last = i
	q.size++
	return nil
}

// GetRequest implements Storage.GetRequest() function
func (q *InMemoryQueueStorage) GetRequest() ([]byte, error) {
	q.lock.Lock()
	defer q.lock.Unlock()
	if q.size == 0 {
		return nil, nil
	}
	r := q.first.Request
	q.first = q.first.Next
	q.size--
	return r, nil
}

// QueueSize implements Storage.QueueSize() function
func (q *InMemoryQueueStorage) QueueSize() (int, error) {
	q.lock.Lock()
	defer q.lock.Unlock()
	return q.size, nil
}
<|endoftext|>


<commit_before>// Package internal holds asset templates used by bootkube.
package internal

var (
	KubeConfigTemplate = []byte(`apiVersion: v1
kind: Config
clusters:
- name: local
  cluster:
    server: {{ .Server }}
    certificate-authority-data: {{ .CACert }}
users:
- name: kubelet
  user:
    client-certificate-data: {{ .KubeletCert}}
    client-key-data: {{ .KubeletKey }}
contexts:
- context:
    cluster: local
    user: kubelet
`)
	KubeletTemplate = []byte(`apiVersion: extensions/v1beta1
kind: DaemonSet
metadata:
  name: kubelet
  namespace: kube-system
  labels:
    k8s-app: kubelet
spec:
  template:
    metadata:
      labels:
        k8s-app: kubelet
    spec:
      containers:
      - name: kubelet
        image: quay.io/coreos/hyperkube:v1.5.1_coreos.0
        command:
        - /nsenter
        - --target=1
        - --mount
        - --wd=.
        - --
        - ./hyperkube
        - kubelet
        - --pod-manifest-path=/etc/kubernetes/manifests
        - --allow-privileged
        - --hostname-override=$(MY_POD_IP)
        - --cluster-dns=10.3.0.10
        - --cluster-domain=cluster.local
        - --kubeconfig=/etc/kubernetes/kubeconfig
        - --require-kubeconfig
        - --lock-file=/var/run/lock/kubelet.lock
        env:
          - name: MY_POD_IP
            valueFrom:
              fieldRef:
                fieldPath: status.podIP
        securityContext:
          privileged: true
        volumeMounts:
        - name: dev
          mountPath: /dev
        - name: run
          mountPath: /run
        - name: sys
          mountPath: /sys
          readOnly: true
        - name: etc-kubernetes
          mountPath: /etc/kubernetes
          readOnly: true
        - name: etc-ssl-certs
          mountPath: /etc/ssl/certs
          readOnly: true
        - name: var-lib-docker
          mountPath: /var/lib/docker
        - name: var-lib-kubelet
          mountPath: /var/lib/kubelet
        - name: var-lib-rkt
          mountPath: /var/lib/rkt
      hostNetwork: true
      hostPID: true
      volumes:
      - name: dev
        hostPath:
          path: /dev
      - name: run
        hostPath:
          path: /run
      - name: sys
        hostPath:
          path: /sys
      - name: etc-kubernetes
        hostPath:
          path: /etc/kubernetes
      - name: etc-ssl-certs
        hostPath:
          path: /usr/share/ca-certificates
      - name: var-lib-docker
        hostPath:
          path: /var/lib/docker
      - name: var-lib-kubelet
        hostPath:
          path: /var/lib/kubelet
      - name: var-lib-rkt
        hostPath:
          path: /var/lib/rkt
`)
	APIServerTemplate = []byte(`apiVersion: "extensions/v1beta1"
kind: DaemonSet
metadata:
  name: kube-apiserver
  namespace: kube-system
  labels:
    k8s-app: kube-apiserver
spec:
  template:
    metadata:
      labels:
        k8s-app: kube-apiserver
      annotations:
        checkpointer.alpha.coreos.com/checkpoint: "true"
    spec:
      nodeSelector:
        master: "true"
      hostNetwork: true
      containers:
      - name: kube-apiserver
        image: quay.io/coreos/hyperkube:v1.5.1_coreos.0
        command:
        - /hyperkube
        - apiserver
        - --bind-address=0.0.0.0
        - --secure-port=443
        - --insecure-port=8080
        - --advertise-address=$(MY_POD_IP)
        - --etcd-servers={{ range $i, $e := .EtcdServers }}{{ if $i }},{{end}}{{ $e }}{{end}}
        - --storage-backend={{.StorageBackend}}
        - --allow-privileged=true
        - --service-cluster-ip-range=10.3.0.0/24
        - --admission-control=NamespaceLifecycle,LimitRanger,ServiceAccount,ResourceQuota
        - --runtime-config=api/all=true
        - --tls-cert-file=/etc/kubernetes/secrets/apiserver.crt
        - --tls-private-key-file=/etc/kubernetes/secrets/apiserver.key
        - --service-account-key-file=/etc/kubernetes/secrets/service-account.pub
        - --client-ca-file=/etc/kubernetes/secrets/ca.crt
        - --cloud-provider={{ .CloudProvider  }}
        - --anonymous-auth=false
        env:
          - name: MY_POD_IP
            valueFrom:
              fieldRef:
                fieldPath: status.podIP
        volumeMounts:
        - mountPath: /etc/ssl/certs
          name: ssl-certs-host
          readOnly: true
        - mountPath: /etc/kubernetes/secrets
          name: secrets
          readOnly: true
      volumes:
      - name: ssl-certs-host
        hostPath:
          path: /usr/share/ca-certificates
      - name: secrets
        secret:
          secretName: kube-apiserver
`)
	CheckpointerTemplate = []byte(`apiVersion: "extensions/v1beta1"
kind: DaemonSet
metadata:
  name: checkpoint-installer
  namespace: kube-system
  labels:
    k8s-app: pod-checkpoint-installer
spec:
  template:
    metadata:
      labels:
        k8s-app: pod-checkpoint-installer
    spec:
      nodeSelector:
        master: "true"
      hostNetwork: true
      containers:
      - name: checkpoint-installer
        image: quay.io/coreos/pod-checkpointer:b4f0353cc12d95737628b8815625cc8e5cedb6fc
        command:
        - /checkpoint-installer.sh
        volumeMounts:
        - mountPath: /etc/kubernetes/manifests
          name: etc-k8s-manifests
      volumes:
      - name: etc-k8s-manifests
        hostPath:
          path: /etc/kubernetes/manifests
`)
	ControllerManagerTemplate = []byte(`apiVersion: extensions/v1beta1
kind: Deployment
metadata:
  name: kube-controller-manager
  namespace: kube-system
  labels:
    k8s-app: kube-controller-manager
spec:
  replicas: 2
  template:
    metadata:
      labels:
        k8s-app: kube-controller-manager
    spec:
      containers:
      - name: kube-controller-manager
        image: quay.io/coreos/hyperkube:v1.5.1_coreos.0
        command:
        - ./hyperkube
        - controller-manager
        - --root-ca-file=/etc/kubernetes/secrets/ca.crt
        - --service-account-private-key-file=/etc/kubernetes/secrets/service-account.key
        - --leader-elect=true
        - --cloud-provider={{ .CloudProvider  }}
        - --configure-cloud-routes=false
        volumeMounts:
        - name: secrets
          mountPath: /etc/kubernetes/secrets
          readOnly: true
        - name: ssl-host
          mountPath: /etc/ssl/certs
          readOnly: true
      volumes:
      - name: secrets
        secret:
          secretName: kube-controller-manager
      - name: ssl-host
        hostPath:
          path: /usr/share/ca-certificates
      dnsPolicy: Default # Don't use cluster DNS.
`)
	SchedulerTemplate = []byte(`apiVersion: extensions/v1beta1
kind: Deployment
metadata:
  name: kube-scheduler
  namespace: kube-system
  labels:
    k8s-app: kube-scheduler
spec:
  replicas: 2
  template:
    metadata:
      labels:
        k8s-app: kube-scheduler
    spec:
      containers:
      - name: kube-scheduler
        image: quay.io/coreos/hyperkube:v1.5.1_coreos.0
        command:
        - ./hyperkube
        - scheduler
        - --leader-elect=true
`)
	ProxyTemplate = []byte(`apiVersion: "extensions/v1beta1"
kind: DaemonSet
metadata:
  name: kube-proxy
  namespace: kube-system
  labels:
    k8s-app: kube-proxy
spec:
  template:
    metadata:
      labels:
        k8s-app: kube-proxy
    spec:
      hostNetwork: true
      containers:
      - name: kube-proxy
        image: quay.io/coreos/hyperkube:v1.5.1_coreos.0
        command:
        - /hyperkube
        - proxy
        - --kubeconfig=/etc/kubernetes/kubeconfig
        - --proxy-mode=iptables
        - --hostname-override=$(POD_IP)
        - --cluster-cidr=10.2.0.0/16
        env:
        - name: POD_IP
          valueFrom:
            fieldRef:
              fieldPath: status.podIP
        securityContext:
          privileged: true
        volumeMounts:
        - mountPath: /etc/ssl/certs
          name: ssl-certs-host
          readOnly: true
        - name: etc-kubernetes
          mountPath: /etc/kubernetes
          readOnly: true
      volumes:
      - hostPath:
          path: /usr/share/ca-certificates
        name: ssl-certs-host
      - name: etc-kubernetes
        hostPath:
          path: /etc/kubernetes
`)
	DNSDeploymentTemplate = []byte(`apiVersion: extensions/v1beta1
kind: Deployment
metadata:
  name: kube-dns
  namespace: kube-system
  labels:
    k8s-app: kube-dns
    kubernetes.io/cluster-service: "true"
spec:
  # replicas: not specified here:
  # 1. In order to make Addon Manager do not reconcile this replicas parameter.
  # 2. Default is 1.
  # 3. Will be tuned in real time if DNS horizontal auto-scaling is turned on.
  strategy:
    rollingUpdate:
      maxSurge: 10%
      maxUnavailable: 0
  selector:
    matchLabels:
      k8s-app: kube-dns
  template:
    metadata:
      labels:
        k8s-app: kube-dns
      annotations:
        scheduler.alpha.kubernetes.io/critical-pod: ''
        scheduler.alpha.kubernetes.io/tolerations: '[{"key":"CriticalAddonsOnly", "operator":"Exists"}]'
    spec:
      containers:
      - name: kubedns
        image: gcr.io/google_containers/kubedns-amd64:1.9
        resources:
          # TODO: Set memory limits when we've profiled the container for large
          # clusters, then set request = limit to keep this container in
          # guaranteed class. Currently, this container falls into the
          # "burstable" category so the kubelet doesn't backoff from restarting it.
          limits:
            memory: 170Mi
          requests:
            cpu: 100m
            memory: 70Mi
        livenessProbe:
          httpGet:
            path: /healthz-kubedns
            port: 8080
            scheme: HTTP
          initialDelaySeconds: 60
          timeoutSeconds: 5
          successThreshold: 1
          failureThreshold: 5
        readinessProbe:
          httpGet:
            path: /readiness
            port: 8081
            scheme: HTTP
          # we poll on pod startup for the Kubernetes master service and
          # only setup the /readiness HTTP server once that's available.
          initialDelaySeconds: 3
          timeoutSeconds: 5
        args:
        - --domain=cluster.local.
        - --dns-port=10053
        - --config-map=kube-dns
        # This should be set to v=2 only after the new image (cut from 1.5) has
        # been released, otherwise we will flood the logs.
        - --v=0
        env:
        - name: PROMETHEUS_PORT
          value: "10055"
        ports:
        - containerPort: 10053
          name: dns-local
          protocol: UDP
        - containerPort: 10053
          name: dns-tcp-local
          protocol: TCP
        - containerPort: 10055
          name: metrics
          protocol: TCP
      - name: dnsmasq
        image: gcr.io/google_containers/kube-dnsmasq-amd64:1.4
        livenessProbe:
          httpGet:
            path: /healthz-dnsmasq
            port: 8080
            scheme: HTTP
          initialDelaySeconds: 60
          timeoutSeconds: 5
          successThreshold: 1
          failureThreshold: 5
        args:
        - --cache-size=1000
        - --no-resolv
        - --server=127.0.0.1#10053
        - --log-facility=-
        ports:
        - containerPort: 53
          name: dns
          protocol: UDP
        - containerPort: 53
          name: dns-tcp
          protocol: TCP
        # see: https://github.com/kubernetes/kubernetes/issues/29055 for details
        resources:
          requests:
            cpu: 150m
            memory: 10Mi
      - name: dnsmasq-metrics
        image: gcr.io/google_containers/dnsmasq-metrics-amd64:1.0
        livenessProbe:
          httpGet:
            path: /metrics
            port: 10054
            scheme: HTTP
          initialDelaySeconds: 60
          timeoutSeconds: 5
          successThreshold: 1
          failureThreshold: 5
        args:
        - --v=2
        - --logtostderr
        ports:
        - containerPort: 10054
          name: metrics
          protocol: TCP
        resources:
          requests:
            memory: 10Mi
      - name: healthz
        image: gcr.io/google_containers/exechealthz-amd64:1.2
        resources:
          limits:
            memory: 50Mi
          requests:
            cpu: 10m
            # Note that this container shouldn't really need 50Mi of memory. The
            # limits are set higher than expected pending investigation on #29688.
            # The extra memory was stolen from the kubedns container to keep the
            # net memory requested by the pod constant.
            memory: 50Mi
        args:
        - --cmd=nslookup kubernetes.default.svc.cluster.local 127.0.0.1 >/dev/null
        - --url=/healthz-dnsmasq
        - --cmd=nslookup kubernetes.default.svc.cluster.local 127.0.0.1:10053 >/dev/null
        - --url=/healthz-kubedns
        - --port=8080
        - --quiet
        ports:
        - containerPort: 8080
          protocol: TCP
      dnsPolicy: Default  # Don't use cluster DNS.
`)
	DNSSvcTemplate = []byte(`apiVersion: v1
kind: Service
metadata:
  name: kube-dns
  namespace: kube-system
  labels:
    k8s-app: kube-dns
    kubernetes.io/cluster-service: "true"
    kubernetes.io/name: "KubeDNS"
spec:
  selector:
    k8s-app: kube-dns
  clusterIP: 10.3.0.10
  ports:
  - name: dns
    port: 53
    protocol: UDP
  - name: dns-tcp
    port: 53
    protocol: TCP
`)
	EtcdOperatorTemplate = []byte(`apiVersion: extensions/v1beta1
kind: Deployment
metadata:
  name: etcd-operator
  namespace: kube-system
  labels:
    k8s-app: etcd-operator
spec:
  replicas: 1
  template:
    metadata:
      labels:
        k8s-app: etcd-operator
    spec:
      containers:
      - name: etcd-operator
        image: quay.io/coreos/etcd-operator
        env:
        - name: MY_POD_NAMESPACE
          valueFrom:
            fieldRef:
              fieldPath: metadata.namespace
`)
	EtcdSvcTemplate = []byte(`apiVersion: v1
kind: Service
metadata:
  name: etcd-service
  namespace: kube-system
spec:
  selector:
    app: etcd
    etcd_cluster: etcd-cluster
  clusterIP: 10.3.0.15
  ports:
  - name: client
    port: 2379
    protocol: TCP
`)
)
<commit_msg>checkpointer: bump to new implementation<commit_after>// Package internal holds asset templates used by bootkube.
package internal

var (
	KubeConfigTemplate = []byte(`apiVersion: v1
kind: Config
clusters:
- name: local
  cluster:
    server: {{ .Server }}
    certificate-authority-data: {{ .CACert }}
users:
- name: kubelet
  user:
    client-certificate-data: {{ .KubeletCert}}
    client-key-data: {{ .KubeletKey }}
contexts:
- context:
    cluster: local
    user: kubelet
`)
	KubeletTemplate = []byte(`apiVersion: extensions/v1beta1
kind: DaemonSet
metadata:
  name: kubelet
  namespace: kube-system
  labels:
    k8s-app: kubelet
spec:
  template:
    metadata:
      labels:
        k8s-app: kubelet
    spec:
      containers:
      - name: kubelet
        image: quay.io/coreos/hyperkube:v1.5.1_coreos.0
        command:
        - /nsenter
        - --target=1
        - --mount
        - --wd=.
        - --
        - ./hyperkube
        - kubelet
        - --pod-manifest-path=/etc/kubernetes/manifests
        - --allow-privileged
        - --hostname-override=$(MY_POD_IP)
        - --cluster-dns=10.3.0.10
        - --cluster-domain=cluster.local
        - --kubeconfig=/etc/kubernetes/kubeconfig
        - --require-kubeconfig
        - --lock-file=/var/run/lock/kubelet.lock
        env:
          - name: MY_POD_IP
            valueFrom:
              fieldRef:
                fieldPath: status.podIP
        securityContext:
          privileged: true
        volumeMounts:
        - name: dev
          mountPath: /dev
        - name: run
          mountPath: /run
        - name: sys
          mountPath: /sys
          readOnly: true
        - name: etc-kubernetes
          mountPath: /etc/kubernetes
          readOnly: true
        - name: etc-ssl-certs
          mountPath: /etc/ssl/certs
          readOnly: true
        - name: var-lib-docker
          mountPath: /var/lib/docker
        - name: var-lib-kubelet
          mountPath: /var/lib/kubelet
        - name: var-lib-rkt
          mountPath: /var/lib/rkt
      hostNetwork: true
      hostPID: true
      volumes:
      - name: dev
        hostPath:
          path: /dev
      - name: run
        hostPath:
          path: /run
      - name: sys
        hostPath:
          path: /sys
      - name: etc-kubernetes
        hostPath:
          path: /etc/kubernetes
      - name: etc-ssl-certs
        hostPath:
          path: /usr/share/ca-certificates
      - name: var-lib-docker
        hostPath:
          path: /var/lib/docker
      - name: var-lib-kubelet
        hostPath:
          path: /var/lib/kubelet
      - name: var-lib-rkt
        hostPath:
          path: /var/lib/rkt
`)
	APIServerTemplate = []byte(`apiVersion: "extensions/v1beta1"
kind: DaemonSet
metadata:
  name: kube-apiserver
  namespace: kube-system
  labels:
    k8s-app: kube-apiserver
spec:
  template:
    metadata:
      labels:
        k8s-app: kube-apiserver
      annotations:
        checkpointer.alpha.coreos.com/checkpoint: "true"
    spec:
      nodeSelector:
        master: "true"
      hostNetwork: true
      containers:
      - name: kube-apiserver
        image: quay.io/coreos/hyperkube:v1.5.1_coreos.0
        command:
        - /hyperkube
        - apiserver
        - --bind-address=0.0.0.0
        - --secure-port=443
        - --insecure-port=8080
        - --advertise-address=$(MY_POD_IP)
        - --etcd-servers={{ range $i, $e := .EtcdServers }}{{ if $i }},{{end}}{{ $e }}{{end}}
        - --storage-backend={{.StorageBackend}}
        - --allow-privileged=true
        - --service-cluster-ip-range=10.3.0.0/24
        - --admission-control=NamespaceLifecycle,LimitRanger,ServiceAccount,ResourceQuota
        - --runtime-config=api/all=true
        - --tls-cert-file=/etc/kubernetes/secrets/apiserver.crt
        - --tls-private-key-file=/etc/kubernetes/secrets/apiserver.key
        - --service-account-key-file=/etc/kubernetes/secrets/service-account.pub
        - --client-ca-file=/etc/kubernetes/secrets/ca.crt
        - --cloud-provider={{ .CloudProvider  }}
        - --anonymous-auth=false
        env:
          - name: MY_POD_IP
            valueFrom:
              fieldRef:
                fieldPath: status.podIP
        volumeMounts:
        - mountPath: /etc/ssl/certs
          name: ssl-certs-host
          readOnly: true
        - mountPath: /etc/kubernetes/secrets
          name: secrets
          readOnly: true
      volumes:
      - name: ssl-certs-host
        hostPath:
          path: /usr/share/ca-certificates
      - name: secrets
        secret:
          secretName: kube-apiserver
`)
	CheckpointerTemplate = []byte(`apiVersion: "extensions/v1beta1"
kind: DaemonSet
metadata:
  name: checkpoint-installer
  namespace: kube-system
  labels:
    k8s-app: pod-checkpoint-installer
spec:
  template:
    metadata:
      labels:
        k8s-app: pod-checkpoint-installer
    spec:
      nodeSelector:
        master: "true"
      hostNetwork: true
      containers:
      - name: checkpoint-installer
        image: quay.io/coreos/pod-checkpointer:5f0ba7be7c958eac22cfd3d1e8e5b004bc6988a0
        command:
        - /checkpoint-installer.sh
        volumeMounts:
        - mountPath: /etc/kubernetes/manifests
          name: etc-k8s-manifests
      volumes:
      - name: etc-k8s-manifests
        hostPath:
          path: /etc/kubernetes/manifests
`)
	ControllerManagerTemplate = []byte(`apiVersion: extensions/v1beta1
kind: Deployment
metadata:
  name: kube-controller-manager
  namespace: kube-system
  labels:
    k8s-app: kube-controller-manager
spec:
  replicas: 2
  template:
    metadata:
      labels:
        k8s-app: kube-controller-manager
    spec:
      containers:
      - name: kube-controller-manager
        image: quay.io/coreos/hyperkube:v1.5.1_coreos.0
        command:
        - ./hyperkube
        - controller-manager
        - --root-ca-file=/etc/kubernetes/secrets/ca.crt
        - --service-account-private-key-file=/etc/kubernetes/secrets/service-account.key
        - --leader-elect=true
        - --cloud-provider={{ .CloudProvider  }}
        - --configure-cloud-routes=false
        volumeMounts:
        - name: secrets
          mountPath: /etc/kubernetes/secrets
          readOnly: true
        - name: ssl-host
          mountPath: /etc/ssl/certs
          readOnly: true
      volumes:
      - name: secrets
        secret:
          secretName: kube-controller-manager
      - name: ssl-host
        hostPath:
          path: /usr/share/ca-certificates
      dnsPolicy: Default # Don't use cluster DNS.
`)
	SchedulerTemplate = []byte(`apiVersion: extensions/v1beta1
kind: Deployment
metadata:
  name: kube-scheduler
  namespace: kube-system
  labels:
    k8s-app: kube-scheduler
spec:
  replicas: 2
  template:
    metadata:
      labels:
        k8s-app: kube-scheduler
    spec:
      containers:
      - name: kube-scheduler
        image: quay.io/coreos/hyperkube:v1.5.1_coreos.0
        command:
        - ./hyperkube
        - scheduler
        - --leader-elect=true
`)
	ProxyTemplate = []byte(`apiVersion: "extensions/v1beta1"
kind: DaemonSet
metadata:
  name: kube-proxy
  namespace: kube-system
  labels:
    k8s-app: kube-proxy
spec:
  template:
    metadata:
      labels:
        k8s-app: kube-proxy
    spec:
      hostNetwork: true
      containers:
      - name: kube-proxy
        image: quay.io/coreos/hyperkube:v1.5.1_coreos.0
        command:
        - /hyperkube
        - proxy
        - --kubeconfig=/etc/kubernetes/kubeconfig
        - --proxy-mode=iptables
        - --hostname-override=$(POD_IP)
        - --cluster-cidr=10.2.0.0/16
        env:
        - name: POD_IP
          valueFrom:
            fieldRef:
              fieldPath: status.podIP
        securityContext:
          privileged: true
        volumeMounts:
        - mountPath: /etc/ssl/certs
          name: ssl-certs-host
          readOnly: true
        - name: etc-kubernetes
          mountPath: /etc/kubernetes
          readOnly: true
      volumes:
      - hostPath:
          path: /usr/share/ca-certificates
        name: ssl-certs-host
      - name: etc-kubernetes
        hostPath:
          path: /etc/kubernetes
`)
	DNSDeploymentTemplate = []byte(`apiVersion: extensions/v1beta1
kind: Deployment
metadata:
  name: kube-dns
  namespace: kube-system
  labels:
    k8s-app: kube-dns
    kubernetes.io/cluster-service: "true"
spec:
  # replicas: not specified here:
  # 1. In order to make Addon Manager do not reconcile this replicas parameter.
  # 2. Default is 1.
  # 3. Will be tuned in real time if DNS horizontal auto-scaling is turned on.
  strategy:
    rollingUpdate:
      maxSurge: 10%
      maxUnavailable: 0
  selector:
    matchLabels:
      k8s-app: kube-dns
  template:
    metadata:
      labels:
        k8s-app: kube-dns
      annotations:
        scheduler.alpha.kubernetes.io/critical-pod: ''
        scheduler.alpha.kubernetes.io/tolerations: '[{"key":"CriticalAddonsOnly", "operator":"Exists"}]'
    spec:
      containers:
      - name: kubedns
        image: gcr.io/google_containers/kubedns-amd64:1.9
        resources:
          # TODO: Set memory limits when we've profiled the container for large
          # clusters, then set request = limit to keep this container in
          # guaranteed class. Currently, this container falls into the
          # "burstable" category so the kubelet doesn't backoff from restarting it.
          limits:
            memory: 170Mi
          requests:
            cpu: 100m
            memory: 70Mi
        livenessProbe:
          httpGet:
            path: /healthz-kubedns
            port: 8080
            scheme: HTTP
          initialDelaySeconds: 60
          timeoutSeconds: 5
          successThreshold: 1
          failureThreshold: 5
        readinessProbe:
          httpGet:
            path: /readiness
            port: 8081
            scheme: HTTP
          # we poll on pod startup for the Kubernetes master service and
          # only setup the /readiness HTTP server once that's available.
          initialDelaySeconds: 3
          timeoutSeconds: 5
        args:
        - --domain=cluster.local.
        - --dns-port=10053
        - --config-map=kube-dns
        # This should be set to v=2 only after the new image (cut from 1.5) has
        # been released, otherwise we will flood the logs.
        - --v=0
        env:
        - name: PROMETHEUS_PORT
          value: "10055"
        ports:
        - containerPort: 10053
          name: dns-local
          protocol: UDP
        - containerPort: 10053
          name: dns-tcp-local
          protocol: TCP
        - containerPort: 10055
          name: metrics
          protocol: TCP
      - name: dnsmasq
        image: gcr.io/google_containers/kube-dnsmasq-amd64:1.4
        livenessProbe:
          httpGet:
            path: /healthz-dnsmasq
            port: 8080
            scheme: HTTP
          initialDelaySeconds: 60
          timeoutSeconds: 5
          successThreshold: 1
          failureThreshold: 5
        args:
        - --cache-size=1000
        - --no-resolv
        - --server=127.0.0.1#10053
        - --log-facility=-
        ports:
        - containerPort: 53
          name: dns
          protocol: UDP
        - containerPort: 53
          name: dns-tcp
          protocol: TCP
        # see: https://github.com/kubernetes/kubernetes/issues/29055 for details
        resources:
          requests:
            cpu: 150m
            memory: 10Mi
      - name: dnsmasq-metrics
        image: gcr.io/google_containers/dnsmasq-metrics-amd64:1.0
        livenessProbe:
          httpGet:
            path: /metrics
            port: 10054
            scheme: HTTP
          initialDelaySeconds: 60
          timeoutSeconds: 5
          successThreshold: 1
          failureThreshold: 5
        args:
        - --v=2
        - --logtostderr
        ports:
        - containerPort: 10054
          name: metrics
          protocol: TCP
        resources:
          requests:
            memory: 10Mi
      - name: healthz
        image: gcr.io/google_containers/exechealthz-amd64:1.2
        resources:
          limits:
            memory: 50Mi
          requests:
            cpu: 10m
            # Note that this container shouldn't really need 50Mi of memory. The
            # limits are set higher than expected pending investigation on #29688.
            # The extra memory was stolen from the kubedns container to keep the
            # net memory requested by the pod constant.
            memory: 50Mi
        args:
        - --cmd=nslookup kubernetes.default.svc.cluster.local 127.0.0.1 >/dev/null
        - --url=/healthz-dnsmasq
        - --cmd=nslookup kubernetes.default.svc.cluster.local 127.0.0.1:10053 >/dev/null
        - --url=/healthz-kubedns
        - --port=8080
        - --quiet
        ports:
        - containerPort: 8080
          protocol: TCP
      dnsPolicy: Default  # Don't use cluster DNS.
`)
	DNSSvcTemplate = []byte(`apiVersion: v1
kind: Service
metadata:
  name: kube-dns
  namespace: kube-system
  labels:
    k8s-app: kube-dns
    kubernetes.io/cluster-service: "true"
    kubernetes.io/name: "KubeDNS"
spec:
  selector:
    k8s-app: kube-dns
  clusterIP: 10.3.0.10
  ports:
  - name: dns
    port: 53
    protocol: UDP
  - name: dns-tcp
    port: 53
    protocol: TCP
`)
	EtcdOperatorTemplate = []byte(`apiVersion: extensions/v1beta1
kind: Deployment
metadata:
  name: etcd-operator
  namespace: kube-system
  labels:
    k8s-app: etcd-operator
spec:
  replicas: 1
  template:
    metadata:
      labels:
        k8s-app: etcd-operator
    spec:
      containers:
      - name: etcd-operator
        image: quay.io/coreos/etcd-operator
        env:
        - name: MY_POD_NAMESPACE
          valueFrom:
            fieldRef:
              fieldPath: metadata.namespace
`)
	EtcdSvcTemplate = []byte(`apiVersion: v1
kind: Service
metadata:
  name: etcd-service
  namespace: kube-system
spec:
  selector:
    app: etcd
    etcd_cluster: etcd-cluster
  clusterIP: 10.3.0.15
  ports:
  - name: client
    port: 2379
    protocol: TCP
`)
)
<|endoftext|>


<commit_before>package server

import (
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"

	"github.com/khlieng/dispatch/Godeps/_workspace/src/github.com/gorilla/websocket"
	"github.com/khlieng/dispatch/Godeps/_workspace/src/github.com/spf13/viper"
	"github.com/khlieng/dispatch/Godeps/_workspace/src/golang.org/x/net/http2"

	"github.com/khlieng/dispatch/letsencrypt"
	"github.com/khlieng/dispatch/storage"
)

var (
	sessions     *sessionStore
	channelStore *storage.ChannelStore

	upgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}
)

func Run() {
	sessions = newSessionStore()
	channelStore = storage.NewChannelStore()

	reconnectIRC()
	initAuth()
	initFileServer()
	startHTTP()
}

func startHTTP() {
	port := viper.GetString("port")

	if viper.GetBool("https.enabled") {
		portHTTPS := viper.GetString("https.port")
		redirect := viper.GetBool("https.redirect")

		if redirect {
			log.Println("[HTTP] Listening on port", port, "(HTTPS Redirect)")
			go http.ListenAndServe(":"+port, createHTTPSRedirect(portHTTPS))
		}

		server := &http.Server{
			Addr:    ":" + portHTTPS,
			Handler: http.HandlerFunc(serve),
		}

		http2.ConfigureServer(server, nil)

		if certExists() {
			log.Println("[HTTPS] Listening on port", portHTTPS)
			server.ListenAndServeTLS(viper.GetString("https.cert"), viper.GetString("https.key"))
		} else if domain := viper.GetString("letsencrypt.domain"); domain != "" {
			dir := storage.Path.LetsEncrypt()
			email := viper.GetString("letsencrypt.email")
			lePort := viper.GetString("letsencrypt.port")

			if viper.GetBool("letsencrypt.proxy") && lePort != "" && (port != "80" || !redirect) {
				log.Println("[HTTP] Listening on port 80 (Let's Encrypt Proxy))")
				go http.ListenAndServe(":80", http.HandlerFunc(letsEncryptProxy))
			}

			letsEncrypt, err := letsencrypt.Run(dir, domain, email, ":"+lePort)
			if err != nil {
				log.Fatal(err)
			}

			server.TLSConfig.GetCertificate = letsEncrypt.GetCertificate

			log.Println("[HTTPS] Listening on port", portHTTPS)
			log.Fatal(listenAndServeTLS(server))
		} else {
			log.Fatal("Could not locate SSL certificate or private key")
		}
	} else {
		log.Println("[HTTP] Listening on port", port)
		log.Fatal(http.ListenAndServe(":"+port, http.HandlerFunc(serve)))
	}
}

func serve(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		w.WriteHeader(404)
		return
	}

	if r.URL.Path == "/ws" {
		session := handleAuth(w, r)
		if session == nil {
			log.Println("[Auth] No session")
			w.WriteHeader(500)
			return
		}

		upgradeWS(w, r, session)
	} else {
		serveFiles(w, r)
	}
}

func upgradeWS(w http.ResponseWriter, r *http.Request, session *Session) {
	conn, err := upgrader.Upgrade(w, r, w.Header())
	if err != nil {
		log.Println(err)
		return
	}

	newWSHandler(conn, session).run()
}

func createHTTPSRedirect(portHTTPS string) http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasPrefix(r.URL.Path, "/.well-known/acme-challenge") {
			letsEncryptProxy(w, r)
			return
		}

		host, _, err := net.SplitHostPort(r.Host)
		if err != nil {
			host = r.Host
		}

		u := url.URL{
			Scheme: "https",
			Host:   net.JoinHostPort(host, portHTTPS),
			Path:   r.RequestURI,
		}

		w.Header().Set("Location", u.String())
		w.WriteHeader(http.StatusMovedPermanently)
	})
}

func letsEncryptProxy(w http.ResponseWriter, r *http.Request) {
	host, _, err := net.SplitHostPort(r.Host)
	if err != nil {
		host = r.Host
	}

	upstream := &url.URL{
		Scheme: "http",
		Host:   net.JoinHostPort(host, viper.GetString("letsencrypt.port")),
	}

	httputil.NewSingleHostReverseProxy(upstream).ServeHTTP(w, r)
}
<commit_msg>Only allow websocket connections from the same origin<commit_after>package server

import (
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"

	"github.com/khlieng/dispatch/Godeps/_workspace/src/github.com/gorilla/websocket"
	"github.com/khlieng/dispatch/Godeps/_workspace/src/github.com/spf13/viper"
	"github.com/khlieng/dispatch/Godeps/_workspace/src/golang.org/x/net/http2"

	"github.com/khlieng/dispatch/letsencrypt"
	"github.com/khlieng/dispatch/storage"
)

var (
	sessions     *sessionStore
	channelStore *storage.ChannelStore

	upgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}
)

func Run() {
	sessions = newSessionStore()
	channelStore = storage.NewChannelStore()

	reconnectIRC()
	initAuth()
	initFileServer()
	startHTTP()
}

func startHTTP() {
	port := viper.GetString("port")

	if viper.GetBool("https.enabled") {
		portHTTPS := viper.GetString("https.port")
		redirect := viper.GetBool("https.redirect")

		if redirect {
			log.Println("[HTTP] Listening on port", port, "(HTTPS Redirect)")
			go http.ListenAndServe(":"+port, createHTTPSRedirect(portHTTPS))
		}

		server := &http.Server{
			Addr:    ":" + portHTTPS,
			Handler: http.HandlerFunc(serve),
		}

		http2.ConfigureServer(server, nil)

		if certExists() {
			log.Println("[HTTPS] Listening on port", portHTTPS)
			server.ListenAndServeTLS(viper.GetString("https.cert"), viper.GetString("https.key"))
		} else if domain := viper.GetString("letsencrypt.domain"); domain != "" {
			dir := storage.Path.LetsEncrypt()
			email := viper.GetString("letsencrypt.email")
			lePort := viper.GetString("letsencrypt.port")

			if viper.GetBool("letsencrypt.proxy") && lePort != "" && (port != "80" || !redirect) {
				log.Println("[HTTP] Listening on port 80 (Let's Encrypt Proxy))")
				go http.ListenAndServe(":80", http.HandlerFunc(letsEncryptProxy))
			}

			letsEncrypt, err := letsencrypt.Run(dir, domain, email, ":"+lePort)
			if err != nil {
				log.Fatal(err)
			}

			server.TLSConfig.GetCertificate = letsEncrypt.GetCertificate

			log.Println("[HTTPS] Listening on port", portHTTPS)
			log.Fatal(listenAndServeTLS(server))
		} else {
			log.Fatal("Could not locate SSL certificate or private key")
		}
	} else {
		log.Println("[HTTP] Listening on port", port)
		log.Fatal(http.ListenAndServe(":"+port, http.HandlerFunc(serve)))
	}
}

func serve(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		w.WriteHeader(404)
		return
	}

	if r.URL.Path == "/ws" {
		session := handleAuth(w, r)
		if session == nil {
			log.Println("[Auth] No session")
			w.WriteHeader(500)
			return
		}

		upgradeWS(w, r, session)
	} else {
		serveFiles(w, r)
	}
}

func upgradeWS(w http.ResponseWriter, r *http.Request, session *Session) {
	conn, err := upgrader.Upgrade(w, r, w.Header())
	if err != nil {
		log.Println(err)
		return
	}

	newWSHandler(conn, session).run()
}

func createHTTPSRedirect(portHTTPS string) http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasPrefix(r.URL.Path, "/.well-known/acme-challenge") {
			letsEncryptProxy(w, r)
			return
		}

		host, _, err := net.SplitHostPort(r.Host)
		if err != nil {
			host = r.Host
		}

		u := url.URL{
			Scheme: "https",
			Host:   net.JoinHostPort(host, portHTTPS),
			Path:   r.RequestURI,
		}

		w.Header().Set("Location", u.String())
		w.WriteHeader(http.StatusMovedPermanently)
	})
}

func letsEncryptProxy(w http.ResponseWriter, r *http.Request) {
	host, _, err := net.SplitHostPort(r.Host)
	if err != nil {
		host = r.Host
	}

	upstream := &url.URL{
		Scheme: "http",
		Host:   net.JoinHostPort(host, viper.GetString("letsencrypt.port")),
	}

	httputil.NewSingleHostReverseProxy(upstream).ServeHTTP(w, r)
}
<|endoftext|>


<commit_before>/*
Copyright 2019 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 cache

import (
	"context"
	"testing"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/latest"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/util"
	"github.com/GoogleContainerTools/skaffold/testutil"
)

type stubDependencyLister struct {
	dependencies []string
}

func (m *stubDependencyLister) DependenciesForArtifact(ctx context.Context, artifact *latest.Artifact) ([]string, error) {
	return m.dependencies, nil
}

var mockCacheHasher = func(s string) (string, error) {
	return s, nil
}

var fakeArtifactConfig = func(a *latest.Artifact) (string, error) {
	if a.ArtifactType.DockerArtifact != nil {
		return "docker/target=" + a.ArtifactType.DockerArtifact.Target, nil
	}
	return "other", nil
}

func TestGetHashForArtifact(t *testing.T) {
	tests := []struct {
		description  string
		dependencies []string
		artifact     *latest.Artifact
		expected     string
	}{
		{
			description:  "hash for artifact",
			dependencies: []string{"a", "b"},
			artifact:     &latest.Artifact{},
			expected:     "1caa15f7ce87536bddbac30a39768e8e3b212bf591f9b64926fa50c40b614c66",
		},
		{
			description:  "dependencies in different orders",
			dependencies: []string{"b", "a"},
			artifact:     &latest.Artifact{},
			expected:     "1caa15f7ce87536bddbac30a39768e8e3b212bf591f9b64926fa50c40b614c66",
		},
		{
			description: "no dependencies",
			artifact:    &latest.Artifact{},
			expected:    "53ebd85adc9b03923a7dacfe6002879af526ef6067d441419d6e62fb9bf608ab",
		},
		{
			description: "docker target",
			artifact: &latest.Artifact{
				ArtifactType: latest.ArtifactType{
					DockerArtifact: &latest.DockerArtifact{
						Target: "target",
					},
				},
			},
			expected: "f947b5aad32734914aa2dea0ec95bceff257037e6c2a529007183c3f21547eae",
		},
		{
			description: "different docker target",
			artifact: &latest.Artifact{
				ArtifactType: latest.ArtifactType{
					DockerArtifact: &latest.DockerArtifact{
						Target: "other",
					},
				},
			},
			expected: "09b366c764d0e39f942283cc081d5522b9dde52e725376661808054e3ed0177f",
		}, {
			description: "build args",
			artifact: &latest.Artifact{
				ArtifactType: latest.ArtifactType{
					DockerArtifact: &latest.DockerArtifact{
						BuildArgs: map[string]*string{"one": stringPointer("1"), "two": stringPointer("2")},
					},
				},
			},
			expected: "f5b610f4fea07461411b2ea0e2cddfd2ffc28d1baed49180f5d3acee5a18f9e7",
		}, {
			description: "build args in different order",
			artifact: &latest.Artifact{
				ArtifactType: latest.ArtifactType{
					DockerArtifact: &latest.DockerArtifact{
						BuildArgs: map[string]*string{"two": stringPointer("2"), "one": stringPointer("1")},
					},
				},
			},
			expected: "f5b610f4fea07461411b2ea0e2cddfd2ffc28d1baed49180f5d3acee5a18f9e7",
		}, {
			description: "different build args",
			artifact: &latest.Artifact{
				ArtifactType: latest.ArtifactType{
					DockerArtifact: &latest.DockerArtifact{
						BuildArgs: map[string]*string{"one": stringPointer("")},
					},
				},
			},
			expected: "961582bfb8d159de1129f89fa12852308afce65816b5f3c521ee57cf9ec524d7",
		},
	}
	for _, test := range tests {
		testutil.Run(t, test.description, func(t *testutil.T) {
			t.Override(&hashFunction, mockCacheHasher)
			t.Override(&artifactConfigFunction, fakeArtifactConfig)

			depLister := &stubDependencyLister{dependencies: test.dependencies}
			actual, err := getHashForArtifact(context.Background(), depLister, test.artifact)

			t.CheckNoError(err)
			t.CheckDeepEqual(test.expected, actual)
		})
	}
}

func TestArtifactConfig(t *testing.T) {
	testutil.Run(t, "", func(t *testutil.T) {
		config1, err := artifactConfig(&latest.Artifact{
			ArtifactType: latest.ArtifactType{
				DockerArtifact: &latest.DockerArtifact{
					Target: "target",
				},
			},
		})
		t.CheckNoError(err)

		config2, err := artifactConfig(&latest.Artifact{
			ArtifactType: latest.ArtifactType{
				DockerArtifact: &latest.DockerArtifact{
					Target: "other",
				},
			},
		})
		t.CheckNoError(err)

		if config1 == config2 {
			t.Errorf("configs should be different: [%s] [%s]", config1, config2)
		}
	})
}

func TestBuildArgsEnvSubstitution(t *testing.T) {
	testutil.Run(t, "", func(t *testutil.T) {
		original := util.OSEnviron
		defer func() { util.OSEnviron = original }()
		util.OSEnviron = func() []string {
			return []string{"FOO=bar"}
		}

		artifact := &latest.Artifact{
			ArtifactType: latest.ArtifactType{
				DockerArtifact: &latest.DockerArtifact{
					BuildArgs: map[string]*string{"env": stringPointer("${{.FOO}}")},
				},
			},
		}

		t.Override(&hashFunction, mockCacheHasher)
		t.Override(&artifactConfigFunction, fakeArtifactConfig)

		depLister := &stubDependencyLister{dependencies: []string{"dep"}}
		actual, err := getHashForArtifact(context.Background(), depLister, artifact)

		t.CheckNoError(err)
		t.CheckDeepEqual("afc018228939a12195994275d6a24ef8f0eb7fa5019ae0582b2daf9e4c353656", actual)

		// Make sure hash is different with a new env

		util.OSEnviron = func() []string {
			return []string{"FOO=baz"}
		}

		actual, err = getHashForArtifact(context.Background(), depLister, artifact)

		t.CheckNoError(err)
		t.CheckDeepEqual("f698ab606ce86ad1c6b7842890a6573060dfaa770e6df0913076f83f14ac32ae", actual)
	})
}

func TestCacheHasher(t *testing.T) {
	tests := []struct {
		description   string
		differentHash bool
		newFilename   string
		update        func(oldFile string, folder *testutil.TempDir)
	}{
		{
			description:   "change filename",
			differentHash: true,
			newFilename:   "newfoo",
			update: func(oldFile string, folder *testutil.TempDir) {
				folder.Rename(oldFile, "newfoo")
			},
		},
		{
			description:   "change file contents",
			differentHash: true,
			update: func(oldFile string, folder *testutil.TempDir) {
				folder.Write(oldFile, "newcontents")
			},
		},
		{
			description:   "change both",
			differentHash: true,
			newFilename:   "newfoo",
			update: func(oldFile string, folder *testutil.TempDir) {
				folder.Rename(oldFile, "newfoo")
				folder.Write(oldFile, "newcontents")
			},
		},
		{
			description:   "change nothing",
			differentHash: false,
			update:        func(oldFile string, folder *testutil.TempDir) {},
		},
	}
	for _, test := range tests {
		testutil.Run(t, test.description, func(t *testutil.T) {
			originalFile := "foo"
			originalContents := "contents"

			tmpDir := t.NewTempDir().
				Write(originalFile, originalContents)

			path := originalFile
			depLister := &stubDependencyLister{dependencies: []string{tmpDir.Path(originalFile)}}

			oldHash, err := getHashForArtifact(context.Background(), depLister, &latest.Artifact{})
			t.CheckNoError(err)

			test.update(originalFile, tmpDir)
			if test.newFilename != "" {
				path = test.newFilename
			}

			depLister = &stubDependencyLister{dependencies: []string{tmpDir.Path(path)}}
			newHash, err := getHashForArtifact(context.Background(), depLister, &latest.Artifact{})

			t.CheckNoError(err)
			t.CheckDeepEqual(false, test.differentHash && oldHash == newHash)
			t.CheckDeepEqual(false, !test.differentHash && oldHash != newHash)
		})
	}
}

func TestRetrieveBuildArgs(t *testing.T) {
	tests := []struct {
		description  string
		artifactType latest.ArtifactType
		expected     map[string]*string
	}{
		{
			description: "docker artifact with build args",
			artifactType: latest.ArtifactType{
				DockerArtifact: &latest.DockerArtifact{
					BuildArgs: map[string]*string{},
				},
			},
			expected: map[string]*string{},
		}, {
			description: "docker artifact without build args",
			artifactType: latest.ArtifactType{
				DockerArtifact: &latest.DockerArtifact{},
			},
		}, {
			description: "kaniko artifact with build args",
			artifactType: latest.ArtifactType{
				KanikoArtifact: &latest.KanikoArtifact{
					BuildArgs: map[string]*string{},
				},
			},
			expected: map[string]*string{},
		}, {
			description: "kaniko artifact without build args",
			artifactType: latest.ArtifactType{
				KanikoArtifact: &latest.KanikoArtifact{},
			},
		}, {
			description: "custom artifact, dockerfile dependency, with build args",
			artifactType: latest.ArtifactType{
				CustomArtifact: &latest.CustomArtifact{
					Dependencies: &latest.CustomDependencies{
						Dockerfile: &latest.DockerfileDependency{
							BuildArgs: map[string]*string{},
						},
					},
				},
			},
			expected: map[string]*string{},
		}, {
			description: "custom artifact, no dockerfile dependency",
			artifactType: latest.ArtifactType{
				CustomArtifact: &latest.CustomArtifact{
					Dependencies: &latest.CustomDependencies{},
				},
			},
		},
	}

	for _, test := range tests {
		t.Run(test.description, func(t *testing.T) {
			actual := retrieveBuildArgs(&latest.Artifact{
				ArtifactType: test.artifactType,
			})
			testutil.CheckDeepEqual(t, test.expected, actual)
		})
	}
}

func TestConvertBuildArgsToStringArray(t *testing.T) {
	tests := []struct {
		description string
		buildArgs   map[string]*string
		expected    []string
	}{
		{
			description: "regular key:value build args",
			buildArgs: map[string]*string{
				"one": stringPointer("1"),
				"two": stringPointer("2"),
			},
			expected: []string{"one=1", "two=2"},
		}, {
			description: "empty key:value build args",
			buildArgs: map[string]*string{
				"one": stringPointer(""),
				"two": stringPointer(""),
			},
			expected: []string{"one=", "two="},
		}, {
			description: "build args with nil vlaue",
			buildArgs: map[string]*string{
				"one": nil,
				"two": nil,
			},
			expected: []string{"one", "two"},
		},
	}

	for _, test := range tests {
		t.Run(test.description, func(t *testing.T) {
			actual := convertBuildArgsToStringArray(test.buildArgs)
			testutil.CheckDeepEqual(t, test.expected, actual)
		})
	}
}

func stringPointer(s string) *string {
	return &s
}
<commit_msg>separate build args test to validate hashes<commit_after>/*
Copyright 2019 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 cache

import (
	"context"
	"testing"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/latest"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/util"
	"github.com/GoogleContainerTools/skaffold/testutil"
)

type stubDependencyLister struct {
	dependencies []string
}

func (m *stubDependencyLister) DependenciesForArtifact(ctx context.Context, artifact *latest.Artifact) ([]string, error) {
	return m.dependencies, nil
}

var mockCacheHasher = func(s string) (string, error) {
	return s, nil
}

var fakeArtifactConfig = func(a *latest.Artifact) (string, error) {
	if a.ArtifactType.DockerArtifact != nil {
		return "docker/target=" + a.ArtifactType.DockerArtifact.Target, nil
	}
	return "other", nil
}

func TestGetHashForArtifact(t *testing.T) {
	tests := []struct {
		description  string
		dependencies []string
		artifact     *latest.Artifact
		expected     string
	}{
		{
			description:  "hash for artifact",
			dependencies: []string{"a", "b"},
			artifact:     &latest.Artifact{},
			expected:     "1caa15f7ce87536bddbac30a39768e8e3b212bf591f9b64926fa50c40b614c66",
		},
		{
			description:  "dependencies in different orders",
			dependencies: []string{"b", "a"},
			artifact:     &latest.Artifact{},
			expected:     "1caa15f7ce87536bddbac30a39768e8e3b212bf591f9b64926fa50c40b614c66",
		},
		{
			description: "no dependencies",
			artifact:    &latest.Artifact{},
			expected:    "53ebd85adc9b03923a7dacfe6002879af526ef6067d441419d6e62fb9bf608ab",
		},
		{
			description: "docker target",
			artifact: &latest.Artifact{
				ArtifactType: latest.ArtifactType{
					DockerArtifact: &latest.DockerArtifact{
						Target: "target",
					},
				},
			},
			expected: "f947b5aad32734914aa2dea0ec95bceff257037e6c2a529007183c3f21547eae",
		},
		{
			description: "different docker target",
			artifact: &latest.Artifact{
				ArtifactType: latest.ArtifactType{
					DockerArtifact: &latest.DockerArtifact{
						Target: "other",
					},
				},
			},
			expected: "09b366c764d0e39f942283cc081d5522b9dde52e725376661808054e3ed0177f",
		},
	}
	for _, test := range tests {
		testutil.Run(t, test.description, func(t *testutil.T) {
			t.Override(&hashFunction, mockCacheHasher)
			t.Override(&artifactConfigFunction, fakeArtifactConfig)

			depLister := &stubDependencyLister{dependencies: test.dependencies}
			actual, err := getHashForArtifact(context.Background(), depLister, test.artifact)

			t.CheckNoError(err)
			t.CheckDeepEqual(test.expected, actual)
		})
	}
}

func TestArtifactConfig(t *testing.T) {
	testutil.Run(t, "", func(t *testutil.T) {
		config1, err := artifactConfig(&latest.Artifact{
			ArtifactType: latest.ArtifactType{
				DockerArtifact: &latest.DockerArtifact{
					Target: "target",
				},
			},
		})
		t.CheckNoError(err)

		config2, err := artifactConfig(&latest.Artifact{
			ArtifactType: latest.ArtifactType{
				DockerArtifact: &latest.DockerArtifact{
					Target: "other",
				},
			},
		})
		t.CheckNoError(err)

		if config1 == config2 {
			t.Errorf("configs should be different: [%s] [%s]", config1, config2)
		}
	})
}

func TestBuildArgs(t *testing.T) {
	testutil.Run(t, "", func(t *testutil.T) {
		expected := "f5b610f4fea07461411b2ea0e2cddfd2ffc28d1baed49180f5d3acee5a18f9e7"

		artifact := &latest.Artifact{
			ArtifactType: latest.ArtifactType{
				DockerArtifact: &latest.DockerArtifact{
					BuildArgs: map[string]*string{"one": stringPointer("1"), "two": stringPointer("2")},
				},
			},
		}

		t.Override(&hashFunction, mockCacheHasher)
		t.Override(&artifactConfigFunction, fakeArtifactConfig)

		actual, err := getHashForArtifact(context.Background(), &stubDependencyLister{}, artifact)

		t.CheckNoError(err)
		t.CheckDeepEqual(expected, actual)

		// Change order of buildargs
		artifact.ArtifactType.DockerArtifact.BuildArgs = map[string]*string{"two": stringPointer("2"), "one": stringPointer("1")}
		actual, err = getHashForArtifact(context.Background(), &stubDependencyLister{}, artifact)

		t.CheckNoError(err)
		t.CheckDeepEqual(expected, actual)

		// Change build args, get different hash
		artifact.ArtifactType.DockerArtifact.BuildArgs = map[string]*string{"one": stringPointer("1")}
		actual, err = getHashForArtifact(context.Background(), &stubDependencyLister{}, artifact)

		t.CheckNoError(err)
		if actual == expected {
			t.Fatal("got same hash as different artifact; expected different hashes.")
		}
	})
}

func TestBuildArgsEnvSubstitution(t *testing.T) {
	testutil.Run(t, "", func(t *testutil.T) {
		original := util.OSEnviron
		defer func() { util.OSEnviron = original }()
		util.OSEnviron = func() []string {
			return []string{"FOO=bar"}
		}

		artifact := &latest.Artifact{
			ArtifactType: latest.ArtifactType{
				DockerArtifact: &latest.DockerArtifact{
					BuildArgs: map[string]*string{"env": stringPointer("${{.FOO}}")},
				},
			},
		}

		t.Override(&hashFunction, mockCacheHasher)
		t.Override(&artifactConfigFunction, fakeArtifactConfig)

		depLister := &stubDependencyLister{dependencies: []string{"dep"}}
		actual, err := getHashForArtifact(context.Background(), depLister, artifact)

		t.CheckNoError(err)
		t.CheckDeepEqual("afc018228939a12195994275d6a24ef8f0eb7fa5019ae0582b2daf9e4c353656", actual)

		// Make sure hash is different with a new env

		util.OSEnviron = func() []string {
			return []string{"FOO=baz"}
		}

		actual, err = getHashForArtifact(context.Background(), depLister, artifact)

		t.CheckNoError(err)
		t.CheckDeepEqual("f698ab606ce86ad1c6b7842890a6573060dfaa770e6df0913076f83f14ac32ae", actual)
	})
}

func TestCacheHasher(t *testing.T) {
	tests := []struct {
		description   string
		differentHash bool
		newFilename   string
		update        func(oldFile string, folder *testutil.TempDir)
	}{
		{
			description:   "change filename",
			differentHash: true,
			newFilename:   "newfoo",
			update: func(oldFile string, folder *testutil.TempDir) {
				folder.Rename(oldFile, "newfoo")
			},
		},
		{
			description:   "change file contents",
			differentHash: true,
			update: func(oldFile string, folder *testutil.TempDir) {
				folder.Write(oldFile, "newcontents")
			},
		},
		{
			description:   "change both",
			differentHash: true,
			newFilename:   "newfoo",
			update: func(oldFile string, folder *testutil.TempDir) {
				folder.Rename(oldFile, "newfoo")
				folder.Write(oldFile, "newcontents")
			},
		},
		{
			description:   "change nothing",
			differentHash: false,
			update:        func(oldFile string, folder *testutil.TempDir) {},
		},
	}
	for _, test := range tests {
		testutil.Run(t, test.description, func(t *testutil.T) {
			originalFile := "foo"
			originalContents := "contents"

			tmpDir := t.NewTempDir().
				Write(originalFile, originalContents)

			path := originalFile
			depLister := &stubDependencyLister{dependencies: []string{tmpDir.Path(originalFile)}}

			oldHash, err := getHashForArtifact(context.Background(), depLister, &latest.Artifact{})
			t.CheckNoError(err)

			test.update(originalFile, tmpDir)
			if test.newFilename != "" {
				path = test.newFilename
			}

			depLister = &stubDependencyLister{dependencies: []string{tmpDir.Path(path)}}
			newHash, err := getHashForArtifact(context.Background(), depLister, &latest.Artifact{})

			t.CheckNoError(err)
			t.CheckDeepEqual(false, test.differentHash && oldHash == newHash)
			t.CheckDeepEqual(false, !test.differentHash && oldHash != newHash)
		})
	}
}

func TestRetrieveBuildArgs(t *testing.T) {
	tests := []struct {
		description  string
		artifactType latest.ArtifactType
		expected     map[string]*string
	}{
		{
			description: "docker artifact with build args",
			artifactType: latest.ArtifactType{
				DockerArtifact: &latest.DockerArtifact{
					BuildArgs: map[string]*string{},
				},
			},
			expected: map[string]*string{},
		}, {
			description: "docker artifact without build args",
			artifactType: latest.ArtifactType{
				DockerArtifact: &latest.DockerArtifact{},
			},
		}, {
			description: "kaniko artifact with build args",
			artifactType: latest.ArtifactType{
				KanikoArtifact: &latest.KanikoArtifact{
					BuildArgs: map[string]*string{},
				},
			},
			expected: map[string]*string{},
		}, {
			description: "kaniko artifact without build args",
			artifactType: latest.ArtifactType{
				KanikoArtifact: &latest.KanikoArtifact{},
			},
		}, {
			description: "custom artifact, dockerfile dependency, with build args",
			artifactType: latest.ArtifactType{
				CustomArtifact: &latest.CustomArtifact{
					Dependencies: &latest.CustomDependencies{
						Dockerfile: &latest.DockerfileDependency{
							BuildArgs: map[string]*string{},
						},
					},
				},
			},
			expected: map[string]*string{},
		}, {
			description: "custom artifact, no dockerfile dependency",
			artifactType: latest.ArtifactType{
				CustomArtifact: &latest.CustomArtifact{
					Dependencies: &latest.CustomDependencies{},
				},
			},
		},
	}

	for _, test := range tests {
		t.Run(test.description, func(t *testing.T) {
			actual := retrieveBuildArgs(&latest.Artifact{
				ArtifactType: test.artifactType,
			})
			testutil.CheckDeepEqual(t, test.expected, actual)
		})
	}
}

func TestConvertBuildArgsToStringArray(t *testing.T) {
	tests := []struct {
		description string
		buildArgs   map[string]*string
		expected    []string
	}{
		{
			description: "regular key:value build args",
			buildArgs: map[string]*string{
				"one": stringPointer("1"),
				"two": stringPointer("2"),
			},
			expected: []string{"one=1", "two=2"},
		}, {
			description: "empty key:value build args",
			buildArgs: map[string]*string{
				"one": stringPointer(""),
				"two": stringPointer(""),
			},
			expected: []string{"one=", "two="},
		}, {
			description: "build args with nil vlaue",
			buildArgs: map[string]*string{
				"one": nil,
				"two": nil,
			},
			expected: []string{"one", "two"},
		},
	}

	for _, test := range tests {
		t.Run(test.description, func(t *testing.T) {
			actual := convertBuildArgsToStringArray(test.buildArgs)
			testutil.CheckDeepEqual(t, test.expected, actual)
		})
	}
}

func stringPointer(s string) *string {
	return &s
}
<|endoftext|>