<commit_before>// Implementation of several read-write locks in Go.
//
// Eli Bendersky [http://eli.thegreenplace.net]
// This code is in the public domain.
package main

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"

	"golang.org/x/sync/semaphore"
)

// RWLocker is an interface RW locks are implementing.
type RWLocker interface {
	RLock()
	RUnlock()
	WLock()
	WUnlock()
}

// DummyRWLock is a dummy no-op implementing the RWLocker interface.
// It's not locking anything, so it will demonstrate that tests catch failure.
type DummyRWLock struct {
}

func (*DummyRWLock) RLock()   {}
func (*DummyRWLock) RUnlock() {}
func (*DummyRWLock) WLock()   {}
func (*DummyRWLock) WUnlock() {}

// MutexAsRWLock is an adapter for sync.Mutex to implement the RWLocker
// interface.
// It uses a single mutex for both R and W locking, so it's not efficient, but
// should be correct.
type MutexAsRWLock struct {
	m sync.Mutex
}

func (l *MutexAsRWLock) RLock()   { l.m.Lock() }
func (l *MutexAsRWLock) RUnlock() { l.m.Unlock() }
func (l *MutexAsRWLock) WLock()   { l.m.Lock() }
func (l *MutexAsRWLock) WUnlock() { l.m.Unlock() }

// RWMutexAsRWLock is an adapter for sync.RWMutex to implement the RWLocker
// interface.
type RWMutexAsRWLock struct {
	rwm sync.RWMutex
}

func (l *RWMutexAsRWLock) RLock()   { l.rwm.RLock() }
func (l *RWMutexAsRWLock) RUnlock() { l.rwm.RUnlock() }
func (l *RWMutexAsRWLock) WLock()   { l.rwm.Lock() }
func (l *RWMutexAsRWLock) WUnlock() { l.rwm.Unlock() }

// ReaderCountRWLock uses a single mutex along with a count of readers. Writers
// will spin on the lock until they see there are no readers.
type ReaderCountRWLock struct {
	m           sync.Mutex
	readerCount int
}

func (l *ReaderCountRWLock) RLock() {
	l.m.Lock()
	l.readerCount++
	l.m.Unlock()
}

func (l *ReaderCountRWLock) RUnlock() {
	l.m.Lock()
	l.readerCount--
	if l.readerCount < 0 {
		panic("readerCount negative")
	}
	l.m.Unlock()
}

func (l *ReaderCountRWLock) WLock() {
	for {
		l.m.Lock()
		if l.readerCount > 0 {
			l.m.Unlock()
		} else {
			break
		}
	}
}

func (l *ReaderCountRWLock) WUnlock() {
	l.m.Unlock()
}

// ReaderCountCondRWLock is similar to ReaderCountRWLock, but uses a condition
// variable to wake up waiting writers (instead of spinning).
type ReaderCountCondRWLock struct {
	readerCount int
	c           *sync.Cond
}

// NewReaderCountCondRWLock creates a new ReaderCountRWLock.
func NewReaderCountCondRWLock() *ReaderCountCondRWLock {
	return &ReaderCountCondRWLock{0, sync.NewCond(new(sync.Mutex))}
}

func (l *ReaderCountCondRWLock) RLock() {
	l.c.L.Lock()
	l.readerCount++
	l.c.L.Unlock()
}

func (l *ReaderCountCondRWLock) RUnlock() {
	l.c.L.Lock()
	l.readerCount--
	if l.readerCount < 0 {
		panic("readerCount negative")
	} else if l.readerCount == 0 {
		l.c.Signal()
	}
	l.c.L.Unlock()
}

func (l *ReaderCountCondRWLock) WLock() {
	l.c.L.Lock()
	for l.readerCount > 0 {
		l.c.Wait()
	}
}

func (l *ReaderCountCondRWLock) WUnlock() {
	l.c.Signal()
	l.c.L.Unlock()
}

// SemaRWLock uses a semaphore; each reader acquires one unit and each writer
// acquires the full semaphore.
const maxWeight int64 = 1 << 30

type SemaRWLock struct {
	s *semaphore.Weighted
}

// NewSemaRWLock creates a new SemaRWLock.
func NewSemaRWLock() *SemaRWLock {
	return &SemaRWLock{semaphore.NewWeighted(maxWeight)}
}

func (l *SemaRWLock) RLock() {
	l.s.Acquire(context.Background(), 1)
}

func (l *SemaRWLock) RUnlock() {
	l.s.Release(1)
}

func (l *SemaRWLock) WLock() {
	l.s.Acquire(context.Background(), maxWeight)
}

func (l *SemaRWLock) WUnlock() {
	l.s.Release(maxWeight)
}

// WritePreferRWLock is a writer-preferring RW lock.
// (from https://en.wikipedia.org/wiki/Readers%E2%80%93writer_lock)
type WritePreferRWLock struct {
	readerCount int
	hasWriter   bool
	c           *sync.Cond
}

func NewWritePreferRWLock() *WritePreferRWLock {
	return &WritePreferRWLock{0, false, sync.NewCond(new(sync.Mutex))}
}

func (l *WritePreferRWLock) RLock() {
	l.c.L.Lock()
	// Wait as long as there are writers waiting to take the lock.
	for l.hasWriter {
		l.c.Wait()
	}
	l.readerCount++
	l.c.L.Unlock()
}

func (l *WritePreferRWLock) RUnlock() {
	l.c.L.Lock()
	l.readerCount--
	if l.readerCount == 0 {
		l.c.Signal()
	}
	l.c.L.Unlock()
}

func (l *WritePreferRWLock) WLock() {
	l.c.L.Lock()
	// Wait until there are no other writers contending for the lock, and mark
	// hasWriter = true.
	for l.hasWriter {
		l.c.Wait()
	}
	l.hasWriter = true
	for l.readerCount > 0 {
		l.c.Wait()
	}
	l.c.L.Unlock()
}

func (l *WritePreferRWLock) WUnlock() {
	l.c.L.Lock()
	l.hasWriter = false
	l.c.Broadcast()
	l.c.L.Unlock()
}

// WritePreferFastRWLock is a faster implementation of a RW lock that prefers
// writers. It's similar to the Go standard library implementation of RWMutex,
// but uses channels instead of runtime_Sem* primitives.
type WritePreferFastRWLock struct {
	// w guarantees mutual exclusion between writers.
	w sync.Mutex

	// These channels serve as semaphores for a writer to wait on readers and for
	// readers to wait for a writer, without spinning.
	writerWait chan struct{}
	readerWait chan struct{}

	// numPending marks the number of readers or writers that are still using the
	// lock. Readers increment it by one, but a writer subtracts maxReaders;
	// therefore, a negative value means a writer is currently using the lock.
	numPending int32

	// readersDeparting is used when a writer waits for all existing readers to
	// flush out before doing its thing. When a writer comes in to take a look
	// and guarantees it's the only writer, it posts -maxReaders onto numPending
	// to notify new readers not to enter. It then also posts the number of
	// pending readers onto readersDeparting. Each reader will decrement this
	// field before relinquishing the lock, and the writer waits for all of them
	// before proceeding.
	readersDeparting int32
}

const maxReaders int32 = 1 << 30

func NewWritePreferFastRWLock() *WritePreferFastRWLock {
	var l WritePreferFastRWLock
	l.writerWait = make(chan struct{})
	l.readerWait = make(chan struct{})
	return &l
}

func (l *WritePreferFastRWLock) RLock() {
	if atomic.AddInt32(&l.numPending, 1) < 0 {
		// A writer is pending, wait for it.
		<-l.readerWait
	}
}

func (l *WritePreferFastRWLock) RUnlock() {
	if r := atomic.AddInt32(&l.numPending, -1); r < 0 {
		// If numPending is now negative, it can be either because of an
		// error ...
		if r+1 == 0 || r+1 == -maxReaders {
			panic("RUnlock of unlocked RWLock")
		}
		// ... or because a writer is pending
		if atomic.AddInt32(&l.readersDeparting, -1) == 0 {
			// The last reader unblocks the writer
			l.writerWait <- struct{}{}
		}
	}
}

func (l *WritePreferFastRWLock) WLock() {
	l.w.Lock()
	// Announce to readers there is a pending writer by decrementing maxReaders
	// from numPending. r will hold the number of pending readers.
	r := atomic.AddInt32(&l.numPending, -maxReaders) + maxReaders
	// Wait for active readers
	if r != 0 && atomic.AddInt32(&l.readersDeparting, r) != 0 {
		<-l.writerWait
	}
}

func (l *WritePreferFastRWLock) WUnlock() {
	// Announce to readers there is no longer an active writer.
	r := atomic.AddInt32(&l.numPending, maxReaders)
	if r >= maxReaders {
		panic("WUnlock of unlocked RWLock")
	}
	// Unblock blocked readers, if any.
	for i := 0; i < int(r); i++ {
		l.readerWait <- struct{}{}
	}
	// Allow other writers to proceed.
	l.w.Unlock()
}

func main() {
	fmt.Println("Run tests")
}
<commit_msg>Replace Signal by Broadcast to be safer (see blog post footnote for details)<commit_after>// Implementation of several read-write locks in Go.
//
// Eli Bendersky [http://eli.thegreenplace.net]
// This code is in the public domain.
package main

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"

	"golang.org/x/sync/semaphore"
)

// RWLocker is an interface RW locks are implementing.
type RWLocker interface {
	RLock()
	RUnlock()
	WLock()
	WUnlock()
}

// DummyRWLock is a dummy no-op implementing the RWLocker interface.
// It's not locking anything, so it will demonstrate that tests catch failure.
type DummyRWLock struct {
}

func (*DummyRWLock) RLock()   {}
func (*DummyRWLock) RUnlock() {}
func (*DummyRWLock) WLock()   {}
func (*DummyRWLock) WUnlock() {}

// MutexAsRWLock is an adapter for sync.Mutex to implement the RWLocker
// interface.
// It uses a single mutex for both R and W locking, so it's not efficient, but
// should be correct.
type MutexAsRWLock struct {
	m sync.Mutex
}

func (l *MutexAsRWLock) RLock()   { l.m.Lock() }
func (l *MutexAsRWLock) RUnlock() { l.m.Unlock() }
func (l *MutexAsRWLock) WLock()   { l.m.Lock() }
func (l *MutexAsRWLock) WUnlock() { l.m.Unlock() }

// RWMutexAsRWLock is an adapter for sync.RWMutex to implement the RWLocker
// interface.
type RWMutexAsRWLock struct {
	rwm sync.RWMutex
}

func (l *RWMutexAsRWLock) RLock()   { l.rwm.RLock() }
func (l *RWMutexAsRWLock) RUnlock() { l.rwm.RUnlock() }
func (l *RWMutexAsRWLock) WLock()   { l.rwm.Lock() }
func (l *RWMutexAsRWLock) WUnlock() { l.rwm.Unlock() }

// ReaderCountRWLock uses a single mutex along with a count of readers. Writers
// will spin on the lock until they see there are no readers.
type ReaderCountRWLock struct {
	m           sync.Mutex
	readerCount int
}

func (l *ReaderCountRWLock) RLock() {
	l.m.Lock()
	l.readerCount++
	l.m.Unlock()
}

func (l *ReaderCountRWLock) RUnlock() {
	l.m.Lock()
	l.readerCount--
	if l.readerCount < 0 {
		panic("readerCount negative")
	}
	l.m.Unlock()
}

func (l *ReaderCountRWLock) WLock() {
	for {
		l.m.Lock()
		if l.readerCount > 0 {
			l.m.Unlock()
		} else {
			break
		}
	}
}

func (l *ReaderCountRWLock) WUnlock() {
	l.m.Unlock()
}

// ReaderCountCondRWLock is similar to ReaderCountRWLock, but uses a condition
// variable to wake up waiting writers (instead of spinning).
type ReaderCountCondRWLock struct {
	readerCount int
	c           *sync.Cond
}

// NewReaderCountCondRWLock creates a new ReaderCountRWLock.
func NewReaderCountCondRWLock() *ReaderCountCondRWLock {
	return &ReaderCountCondRWLock{0, sync.NewCond(new(sync.Mutex))}
}

func (l *ReaderCountCondRWLock) RLock() {
	l.c.L.Lock()
	l.readerCount++
	l.c.L.Unlock()
}

func (l *ReaderCountCondRWLock) RUnlock() {
	l.c.L.Lock()
	l.readerCount--
	if l.readerCount < 0 {
		panic("readerCount negative")
	} else if l.readerCount == 0 {
		l.c.Signal()
	}
	l.c.L.Unlock()
}

func (l *ReaderCountCondRWLock) WLock() {
	l.c.L.Lock()
	for l.readerCount > 0 {
		l.c.Wait()
	}
}

func (l *ReaderCountCondRWLock) WUnlock() {
	l.c.Signal()
	l.c.L.Unlock()
}

// SemaRWLock uses a semaphore; each reader acquires one unit and each writer
// acquires the full semaphore.
const maxWeight int64 = 1 << 30

type SemaRWLock struct {
	s *semaphore.Weighted
}

// NewSemaRWLock creates a new SemaRWLock.
func NewSemaRWLock() *SemaRWLock {
	return &SemaRWLock{semaphore.NewWeighted(maxWeight)}
}

func (l *SemaRWLock) RLock() {
	l.s.Acquire(context.Background(), 1)
}

func (l *SemaRWLock) RUnlock() {
	l.s.Release(1)
}

func (l *SemaRWLock) WLock() {
	l.s.Acquire(context.Background(), maxWeight)
}

func (l *SemaRWLock) WUnlock() {
	l.s.Release(maxWeight)
}

// WritePreferRWLock is a writer-preferring RW lock.
// (from https://en.wikipedia.org/wiki/Readers%E2%80%93writer_lock)
type WritePreferRWLock struct {
	readerCount int
	hasWriter   bool
	c           *sync.Cond
}

func NewWritePreferRWLock() *WritePreferRWLock {
	return &WritePreferRWLock{0, false, sync.NewCond(new(sync.Mutex))}
}

func (l *WritePreferRWLock) RLock() {
	l.c.L.Lock()
	// Wait as long as there are writers waiting to take the lock.
	for l.hasWriter {
		l.c.Wait()
	}
	l.readerCount++
	l.c.L.Unlock()
}

func (l *WritePreferRWLock) RUnlock() {
	l.c.L.Lock()
	l.readerCount--
	if l.readerCount == 0 {
		l.c.Broadcast()
	}
	l.c.L.Unlock()
}

func (l *WritePreferRWLock) WLock() {
	l.c.L.Lock()
	// Wait until there are no other writers contending for the lock, and mark
	// hasWriter = true.
	for l.hasWriter {
		l.c.Wait()
	}
	l.hasWriter = true
	for l.readerCount > 0 {
		l.c.Wait()
	}
	l.c.L.Unlock()
}

func (l *WritePreferRWLock) WUnlock() {
	l.c.L.Lock()
	l.hasWriter = false
	l.c.Broadcast()
	l.c.L.Unlock()
}

// WritePreferFastRWLock is a faster implementation of a RW lock that prefers
// writers. It's similar to the Go standard library implementation of RWMutex,
// but uses channels instead of runtime_Sem* primitives.
type WritePreferFastRWLock struct {
	// w guarantees mutual exclusion between writers.
	w sync.Mutex

	// These channels serve as semaphores for a writer to wait on readers and for
	// readers to wait for a writer, without spinning.
	writerWait chan struct{}
	readerWait chan struct{}

	// numPending marks the number of readers or writers that are still using the
	// lock. Readers increment it by one, but a writer subtracts maxReaders;
	// therefore, a negative value means a writer is currently using the lock.
	numPending int32

	// readersDeparting is used when a writer waits for all existing readers to
	// flush out before doing its thing. When a writer comes in to take a look
	// and guarantees it's the only writer, it posts -maxReaders onto numPending
	// to notify new readers not to enter. It then also posts the number of
	// pending readers onto readersDeparting. Each reader will decrement this
	// field before relinquishing the lock, and the writer waits for all of them
	// before proceeding.
	readersDeparting int32
}

const maxReaders int32 = 1 << 30

func NewWritePreferFastRWLock() *WritePreferFastRWLock {
	var l WritePreferFastRWLock
	l.writerWait = make(chan struct{})
	l.readerWait = make(chan struct{})
	return &l
}

func (l *WritePreferFastRWLock) RLock() {
	if atomic.AddInt32(&l.numPending, 1) < 0 {
		// A writer is pending, wait for it.
		<-l.readerWait
	}
}

func (l *WritePreferFastRWLock) RUnlock() {
	if r := atomic.AddInt32(&l.numPending, -1); r < 0 {
		// If numPending is now negative, it can be either because of an
		// error ...
		if r+1 == 0 || r+1 == -maxReaders {
			panic("RUnlock of unlocked RWLock")
		}
		// ... or because a writer is pending
		if atomic.AddInt32(&l.readersDeparting, -1) == 0 {
			// The last reader unblocks the writer
			l.writerWait <- struct{}{}
		}
	}
}

func (l *WritePreferFastRWLock) WLock() {
	l.w.Lock()
	// Announce to readers there is a pending writer by decrementing maxReaders
	// from numPending. r will hold the number of pending readers.
	r := atomic.AddInt32(&l.numPending, -maxReaders) + maxReaders
	// Wait for active readers
	if r != 0 && atomic.AddInt32(&l.readersDeparting, r) != 0 {
		<-l.writerWait
	}
}

func (l *WritePreferFastRWLock) WUnlock() {
	// Announce to readers there is no longer an active writer.
	r := atomic.AddInt32(&l.numPending, maxReaders)
	if r >= maxReaders {
		panic("WUnlock of unlocked RWLock")
	}
	// Unblock blocked readers, if any.
	for i := 0; i < int(r); i++ {
		l.readerWait <- struct{}{}
	}
	// Allow other writers to proceed.
	l.w.Unlock()
}

func main() {
	fmt.Println("Run tests")
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"os/signal"
	"syscall"

	"github.com/coreos/go-iptables/iptables"
	"github.com/spf13/cobra"
	coreapi "k8s.io/api/core/v1"
	networkingv1 "k8s.io/api/networking/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/cache"

	"github.com/weaveworks/weave/common"
	"github.com/weaveworks/weave/net"
	"github.com/weaveworks/weave/net/ipset"
	"github.com/weaveworks/weave/npc"
	"github.com/weaveworks/weave/npc/metrics"
	"github.com/weaveworks/weave/npc/ulogd"
)

var (
	version     = "unreleased"
	metricsAddr string
	logLevel    string
	allowMcast  bool
	nodeName    string
	maxList     int
)

func handleError(err error) { common.CheckFatal(err) }

func makeController(getter cache.Getter, resource string,
	objType runtime.Object, handlers cache.ResourceEventHandlerFuncs) cache.Controller {
	listWatch := cache.NewListWatchFromClient(getter, resource, "", fields.Everything())
	_, controller := cache.NewInformer(listWatch, objType, 0, handlers)
	return controller
}

func resetIPTables(ipt *iptables.IPTables) error {
	// Flush chains first so there are no refs to extant ipsets
	if err := ipt.ClearChain(npc.TableFilter, npc.IngressChain); err != nil {
		return err
	}

	if err := ipt.ClearChain(npc.TableFilter, npc.DefaultChain); err != nil {
		return err
	}

	if err := ipt.ClearChain(npc.TableFilter, npc.MainChain); err != nil {
		return err
	}

	if err := ipt.ClearChain(npc.TableFilter, npc.EgressMarkChain); err != nil {
		return err
	}

	if err := ipt.ClearChain(npc.TableFilter, npc.EgressCustomChain); err != nil {
		return err
	}

	if err := ipt.ClearChain(npc.TableFilter, npc.EgressDefaultChain); err != nil {
		return err
	}

	// We do not clear npc.EgressChain here because otherwise, in the case of restarting
	// weave-npc process, all egress traffic is allowed for a short period of time.
	// The chain is created in createBaseRules.

	return nil
}

func resetIPSets(ips ipset.Interface) error {
	// Remove ipsets prefixed `weave-` only.

	sets, err := ips.List(npc.IpsetNamePrefix)
	if err != nil {
		common.Log.Errorf("Failed to retrieve list of ipsets")
		return err
	}

	common.Log.Debugf("Got list of ipsets: %v", sets)

	// Must remove references to ipsets by other ipsets before they're destroyed
	for _, s := range sets {
		common.Log.Debugf("Flushing ipset '%s'", string(s))
		if err := ips.Flush(s); err != nil {
			common.Log.Errorf("Failed to flush ipset '%s'", string(s))
			return err
		}
	}

	for _, s := range sets {
		// LocalIPset might be used by WEAVE-NPC-EGRESS chain which we do not
		// flush, so we cannot destroy it.
		if s == npc.LocalIpset {
			continue
		}
		common.Log.Debugf("Destroying ipset '%s'", string(s))
		if err := ips.Destroy(s); err != nil {
			common.Log.Errorf("Failed to destroy ipset '%s'", string(s))
			return err
		}
	}

	return nil
}

func createBaseRules(ipt *iptables.IPTables, ips ipset.Interface) error {
	// Configure main chain static rules
	if err := ipt.Append(npc.TableFilter, npc.MainChain,
		"-m", "state", "--state", "RELATED,ESTABLISHED", "-j", "ACCEPT"); err != nil {
		return err
	}

	if allowMcast {
		if err := ipt.Append(npc.TableFilter, npc.MainChain,
			"-d", "224.0.0.0/4", "-j", "ACCEPT"); err != nil {
			return err
		}
	}

	// If the destination address is not any of the local pods, let it through
	if err := ipt.Append(npc.TableFilter, npc.MainChain,
		"-m", "physdev", "--physdev-out=vethwe-bridge", "-j", "ACCEPT"); err != nil {
		return err
	}

	if err := ipt.Append(npc.TableFilter, npc.MainChain,
		"-m", "state", "--state", "NEW", "-j", string(npc.DefaultChain)); err != nil {
		return err
	}

	if err := ipt.Append(npc.TableFilter, npc.MainChain,
		"-m", "state", "--state", "NEW", "-j", string(npc.IngressChain)); err != nil {
		return err
	}

	if err := ipt.Append(npc.TableFilter, npc.EgressMarkChain,
		"-j", "MARK", "--set-xmark", npc.EgressMark); err != nil {
		return err
	}

	// Egress rules:
	//
	// -A WEAVE-NPC-EGRESS -m state --state RELATED,ESTABLISHED -j ACCEPT
	// -A WEAVE-NPC-EGRESS -m physdev --physdev-in=vethwe-bridge -j RETURN
	// -A WEAVE-NPC-EGRESS -d 224.0.0.0/4 -j RETURN
	// -A WEAVE-NPC-EGRESS -m state --state NEW -j WEAVE-NPC-EGRESS-DEFAULT
	// -A WEAVE-NPC-EGRESS -m state --state NEW -m mark ! --mark 0x40000/0x40000 -j WEAVE-NPC-EGRESS-CUSTOM
	// -A WEAVE-NPC-EGRESS -m state --state NEW -m mark ! --mark 0x40000/0x40000 -j NFLOG --nflog-group 86
	// -A WEAVE-NPC-EGRESS -m mark ! --mark 0x40000/0x40000 -j DROP
	//
	// -A WEAVE-NPC-EGRESS-CUSTOM <rulespec> -j MARK --set-xmark 0x40000/0x40000
	// -A WEAVE-NPC-EGRESS-CUSTOM <rulespec> -j RETURN
	//
	// -A WEAVE-NPC-EGRESS-DEFAULT <rulespec> -j MARK --set-xmark 0x40000/0x40000
	// -A WEAVE-NPC-EGRESS-DEFAULT <rulespec> -j RETURN
	//
	// For each rule we create two (mark and return). We cannot just accept
	// a packet if it matches any rule, as a packet might need to traverse
	// the ingress npc as well which happens later in the chain (in some cases
	// we cannot detect whether packet is ingress or egress, so we need to
	// check both chains).

	ruleSpecs := [][]string{
		{"-m", "state", "--state", "RELATED,ESTABLISHED", "-j", "ACCEPT"},
		{"-m", "physdev", "--physdev-in=vethwe-bridge", "-j", "RETURN"},
	}
	if allowMcast {
		ruleSpecs = append(ruleSpecs, []string{"-d", "224.0.0.0/4", "-j", "RETURN"})
	}
	ruleSpecs = append(ruleSpecs, [][]string{
		{"-m", "state", "--state", "NEW", "-j", string(npc.EgressDefaultChain)},
		{"-m", "state", "--state", "NEW", "-m", "mark", "!", "--mark", npc.EgressMark, "-j", string(npc.EgressCustomChain)},
		{"-m", "state", "--state", "NEW", "-m", "mark", "!", "--mark", npc.EgressMark, "-j", "NFLOG", "--nflog-group", "86"},
		{"-m", "mark", "!", "--mark", npc.EgressMark, "-j", "DROP"},
	}...)
	if err := net.AddChainWithRules(ipt, npc.TableFilter, npc.EgressChain, ruleSpecs); err != nil {
		return err
	}

	return nil
}

// Dummy way to check whether a given ipset exists.
// TODO(brb) Use "ipset -exist create <..>" for our purpose instead (for some reasons
// creating an ipset with -exist fails).
func ipsetExist(ips ipset.Interface, name ipset.Name) (bool, error) {
	sets, err := ips.List(string(name))
	if err != nil {
		return false, err
	}
	for _, s := range sets {
		if s == name {
			return true, nil
		}
	}
	return false, nil
}

func root(cmd *cobra.Command, args []string) {
	var npController cache.Controller

	common.SetLogLevel(logLevel)
	if nodeName == "" {
		// HOSTNAME is set by Kubernetes for pods in the host network namespace
		nodeName = os.Getenv("HOSTNAME")
	}
	if nodeName == "" {
		common.Log.Fatalf("Must set node name via --node-name or $HOSTNAME")
	}
	common.Log.Infof("Starting Weaveworks NPC %s; node name %q", version, nodeName)

	if err := metrics.Start(metricsAddr); err != nil {
		common.Log.Fatalf("Failed to start metrics: %v", err)
	}

	if err := ulogd.Start(); err != nil {
		common.Log.Fatalf("Failed to start ulogd: %v", err)
	}

	config, err := rest.InClusterConfig()
	handleError(err)

	client, err := kubernetes.NewForConfig(config)
	handleError(err)

	ipt, err := iptables.New()
	handleError(err)

	ips := ipset.New(common.LogLogger(), maxList)

	handleError(resetIPTables(ipt))
	handleError(resetIPSets(ips))
	handleError(createBaseRules(ipt, ips))

	npc := npc.New(nodeName, ipt, ips)

	nsController := makeController(client.Core().RESTClient(), "namespaces", &coreapi.Namespace{},
		cache.ResourceEventHandlerFuncs{
			AddFunc: func(obj interface{}) {
				handleError(npc.AddNamespace(obj.(*coreapi.Namespace)))
			},
			DeleteFunc: func(obj interface{}) {
				switch obj := obj.(type) {
				case *coreapi.Namespace:
					handleError(npc.DeleteNamespace(obj))
				case cache.DeletedFinalStateUnknown:
					// We know this object has gone away, but its final state is no longer
					// available from the API server. Instead we use the last copy of it
					// that we have, which is good enough for our cleanup.
					handleError(npc.DeleteNamespace(obj.Obj.(*coreapi.Namespace)))
				}
			},
			UpdateFunc: func(old, new interface{}) {
				handleError(npc.UpdateNamespace(old.(*coreapi.Namespace), new.(*coreapi.Namespace)))
			}})

	podController := makeController(client.Core().RESTClient(), "pods", &coreapi.Pod{},
		cache.ResourceEventHandlerFuncs{
			AddFunc: func(obj interface{}) {
				handleError(npc.AddPod(obj.(*coreapi.Pod)))
			},
			DeleteFunc: func(obj interface{}) {
				switch obj := obj.(type) {
				case *coreapi.Pod:
					handleError(npc.DeletePod(obj))
				case cache.DeletedFinalStateUnknown:
					// We know this object has gone away, but its final state is no longer
					// available from the API server. Instead we use the last copy of it
					// that we have, which is good enough for our cleanup.
					handleError(npc.DeletePod(obj.Obj.(*coreapi.Pod)))
				}
			},
			UpdateFunc: func(old, new interface{}) {
				handleError(npc.UpdatePod(old.(*coreapi.Pod), new.(*coreapi.Pod)))
			}})

	npHandlers := cache.ResourceEventHandlerFuncs{
		AddFunc: func(obj interface{}) {
			handleError(npc.AddNetworkPolicy(obj))
		},
		DeleteFunc: func(obj interface{}) {
			switch obj := obj.(type) {
			case cache.DeletedFinalStateUnknown:
				// We know this object has gone away, but its final state is no longer
				// available from the API server. Instead we use the last copy of it
				// that we have, which is good enough for our cleanup.
				handleError(npc.DeleteNetworkPolicy(obj.Obj))
			default:
				handleError(npc.DeleteNetworkPolicy(obj))
			}
		},
		UpdateFunc: func(old, new interface{}) {
			handleError(npc.UpdateNetworkPolicy(old, new))
		},
	}
	npController = makeController(client.NetworkingV1().RESTClient(), "networkpolicies", &networkingv1.NetworkPolicy{}, npHandlers)

	go nsController.Run(wait.NeverStop)
	go podController.Run(wait.NeverStop)
	go npController.Run(wait.NeverStop)

	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
	common.Log.Fatalf("Exiting: %v", <-signals)
}

func main() {
	rootCmd := &cobra.Command{
		Use:   "weave-npc",
		Short: "Weaveworks Kubernetes Network Policy Controller",
		Run:   root}

	rootCmd.PersistentFlags().StringVar(&metricsAddr, "metrics-addr", ":6781", "metrics server bind address")
	rootCmd.PersistentFlags().StringVar(&logLevel, "log-level", "debug", "logging level (debug, info, warning, error)")
	rootCmd.PersistentFlags().BoolVar(&allowMcast, "allow-mcast", true, "allow all multicast traffic")
	rootCmd.PersistentFlags().StringVar(&nodeName, "node-name", "", "only generate rules that apply to this node")
	rootCmd.PersistentFlags().IntVar(&maxList, "max-list-size", 1024, "maximum size of ipset list (for namespaces)")

	handleError(rootCmd.Execute())
}
<commit_msg>addressing review comments<commit_after>package main

import (
	"os"
	"os/signal"
	"syscall"

	"github.com/coreos/go-iptables/iptables"
	"github.com/spf13/cobra"
	coreapi "k8s.io/api/core/v1"
	networkingv1 "k8s.io/api/networking/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/cache"

	"github.com/weaveworks/weave/common"
	"github.com/weaveworks/weave/net"
	"github.com/weaveworks/weave/net/ipset"
	"github.com/weaveworks/weave/npc"
	"github.com/weaveworks/weave/npc/metrics"
	"github.com/weaveworks/weave/npc/ulogd"
)

var (
	version     = "unreleased"
	metricsAddr string
	logLevel    string
	allowMcast  bool
	nodeName    string
	maxList     int
	bridgeName  string
)

func handleError(err error) { common.CheckFatal(err) }

func makeController(getter cache.Getter, resource string,
	objType runtime.Object, handlers cache.ResourceEventHandlerFuncs) cache.Controller {
	listWatch := cache.NewListWatchFromClient(getter, resource, "", fields.Everything())
	_, controller := cache.NewInformer(listWatch, objType, 0, handlers)
	return controller
}

func resetIPTables(ipt *iptables.IPTables) error {
	// Flush chains first so there are no refs to extant ipsets
	if err := ipt.ClearChain(npc.TableFilter, npc.IngressChain); err != nil {
		return err
	}

	if err := ipt.ClearChain(npc.TableFilter, npc.DefaultChain); err != nil {
		return err
	}

	if err := ipt.ClearChain(npc.TableFilter, npc.MainChain); err != nil {
		return err
	}

	if err := ipt.ClearChain(npc.TableFilter, npc.EgressMarkChain); err != nil {
		return err
	}

	if err := ipt.ClearChain(npc.TableFilter, npc.EgressCustomChain); err != nil {
		return err
	}

	if err := ipt.ClearChain(npc.TableFilter, npc.EgressDefaultChain); err != nil {
		return err
	}

	// We do not clear npc.EgressChain here because otherwise, in the case of restarting
	// weave-npc process, all egress traffic is allowed for a short period of time.
	// The chain is created in createBaseRules.

	return nil
}

func resetIPSets(ips ipset.Interface) error {
	// Remove ipsets prefixed `weave-` only.

	sets, err := ips.List(npc.IpsetNamePrefix)
	if err != nil {
		common.Log.Errorf("Failed to retrieve list of ipsets")
		return err
	}

	common.Log.Debugf("Got list of ipsets: %v", sets)

	// Must remove references to ipsets by other ipsets before they're destroyed
	for _, s := range sets {
		common.Log.Debugf("Flushing ipset '%s'", string(s))
		if err := ips.Flush(s); err != nil {
			common.Log.Errorf("Failed to flush ipset '%s'", string(s))
			return err
		}
	}

	for _, s := range sets {
		// LocalIPset might be used by WEAVE-NPC-EGRESS chain which we do not
		// flush, so we cannot destroy it.
		if s == npc.LocalIpset {
			continue
		}
		common.Log.Debugf("Destroying ipset '%s'", string(s))
		if err := ips.Destroy(s); err != nil {
			common.Log.Errorf("Failed to destroy ipset '%s'", string(s))
			return err
		}
	}

	return nil
}

func createBaseRules(ipt *iptables.IPTables, ips ipset.Interface) error {
	// Configure main chain static rules
	if err := ipt.Append(npc.TableFilter, npc.MainChain,
		"-m", "state", "--state", "RELATED,ESTABLISHED", "-j", "ACCEPT"); err != nil {
		return err
	}

	if allowMcast {
		if err := ipt.Append(npc.TableFilter, npc.MainChain,
			"-d", "224.0.0.0/4", "-j", "ACCEPT"); err != nil {
			return err
		}
	}

	// If the destination address is not any of the local pods, let it through
	if err := ipt.Append(npc.TableFilter, npc.MainChain,
		"-m", "physdev", "--physdev-out="+bridgeName, "-j", "ACCEPT"); err != nil {
		return err
	}

	if err := ipt.Append(npc.TableFilter, npc.MainChain,
		"-m", "state", "--state", "NEW", "-j", string(npc.DefaultChain)); err != nil {
		return err
	}

	if err := ipt.Append(npc.TableFilter, npc.MainChain,
		"-m", "state", "--state", "NEW", "-j", string(npc.IngressChain)); err != nil {
		return err
	}

	if err := ipt.Append(npc.TableFilter, npc.EgressMarkChain,
		"-j", "MARK", "--set-xmark", npc.EgressMark); err != nil {
		return err
	}

	// Egress rules:
	//
	// -A WEAVE-NPC-EGRESS -m state --state RELATED,ESTABLISHED -j ACCEPT
	// -A WEAVE-NPC-EGRESS -m physdev --physdev-in=vethwe-bridge -j RETURN
	// -A WEAVE-NPC-EGRESS -d 224.0.0.0/4 -j RETURN
	// -A WEAVE-NPC-EGRESS -m state --state NEW -j WEAVE-NPC-EGRESS-DEFAULT
	// -A WEAVE-NPC-EGRESS -m state --state NEW -m mark ! --mark 0x40000/0x40000 -j WEAVE-NPC-EGRESS-CUSTOM
	// -A WEAVE-NPC-EGRESS -m state --state NEW -m mark ! --mark 0x40000/0x40000 -j NFLOG --nflog-group 86
	// -A WEAVE-NPC-EGRESS -m mark ! --mark 0x40000/0x40000 -j DROP
	//
	// -A WEAVE-NPC-EGRESS-CUSTOM <rulespec> -j MARK --set-xmark 0x40000/0x40000
	// -A WEAVE-NPC-EGRESS-CUSTOM <rulespec> -j RETURN
	//
	// -A WEAVE-NPC-EGRESS-DEFAULT <rulespec> -j MARK --set-xmark 0x40000/0x40000
	// -A WEAVE-NPC-EGRESS-DEFAULT <rulespec> -j RETURN
	//
	// For each rule we create two (mark and return). We cannot just accept
	// a packet if it matches any rule, as a packet might need to traverse
	// the ingress npc as well which happens later in the chain (in some cases
	// we cannot detect whether packet is ingress or egress, so we need to
	// check both chains).

	ruleSpecs := [][]string{
		{"-m", "state", "--state", "RELATED,ESTABLISHED", "-j", "ACCEPT"},
		{"-m", "physdev", "--physdev-in=" + bridgeName, "-j", "RETURN"},
	}
	if allowMcast {
		ruleSpecs = append(ruleSpecs, []string{"-d", "224.0.0.0/4", "-j", "RETURN"})
	}
	ruleSpecs = append(ruleSpecs, [][]string{
		{"-m", "state", "--state", "NEW", "-j", string(npc.EgressDefaultChain)},
		{"-m", "state", "--state", "NEW", "-m", "mark", "!", "--mark", npc.EgressMark, "-j", string(npc.EgressCustomChain)},
		{"-m", "state", "--state", "NEW", "-m", "mark", "!", "--mark", npc.EgressMark, "-j", "NFLOG", "--nflog-group", "86"},
		{"-m", "mark", "!", "--mark", npc.EgressMark, "-j", "DROP"},
	}...)
	if err := net.AddChainWithRules(ipt, npc.TableFilter, npc.EgressChain, ruleSpecs); err != nil {
		return err
	}

	return nil
}

// Dummy way to check whether a given ipset exists.
// TODO(brb) Use "ipset -exist create <..>" for our purpose instead (for some reasons
// creating an ipset with -exist fails).
func ipsetExist(ips ipset.Interface, name ipset.Name) (bool, error) {
	sets, err := ips.List(string(name))
	if err != nil {
		return false, err
	}
	for _, s := range sets {
		if s == name {
			return true, nil
		}
	}
	return false, nil
}

func root(cmd *cobra.Command, args []string) {
	var npController cache.Controller

	common.SetLogLevel(logLevel)
	if nodeName == "" {
		// HOSTNAME is set by Kubernetes for pods in the host network namespace
		nodeName = os.Getenv("HOSTNAME")
	}
	if nodeName == "" {
		common.Log.Fatalf("Must set node name via --node-name or $HOSTNAME")
	}
	common.Log.Infof("Starting Weaveworks NPC %s; node name %q", version, nodeName)

	if err := metrics.Start(metricsAddr); err != nil {
		common.Log.Fatalf("Failed to start metrics: %v", err)
	}

	if err := ulogd.Start(); err != nil {
		common.Log.Fatalf("Failed to start ulogd: %v", err)
	}

	config, err := rest.InClusterConfig()
	handleError(err)

	client, err := kubernetes.NewForConfig(config)
	handleError(err)

	ipt, err := iptables.New()
	handleError(err)

	ips := ipset.New(common.LogLogger(), maxList)

	handleError(resetIPTables(ipt))
	handleError(resetIPSets(ips))
	handleError(createBaseRules(ipt, ips))

	npc := npc.New(nodeName, ipt, ips)

	nsController := makeController(client.Core().RESTClient(), "namespaces", &coreapi.Namespace{},
		cache.ResourceEventHandlerFuncs{
			AddFunc: func(obj interface{}) {
				handleError(npc.AddNamespace(obj.(*coreapi.Namespace)))
			},
			DeleteFunc: func(obj interface{}) {
				switch obj := obj.(type) {
				case *coreapi.Namespace:
					handleError(npc.DeleteNamespace(obj))
				case cache.DeletedFinalStateUnknown:
					// We know this object has gone away, but its final state is no longer
					// available from the API server. Instead we use the last copy of it
					// that we have, which is good enough for our cleanup.
					handleError(npc.DeleteNamespace(obj.Obj.(*coreapi.Namespace)))
				}
			},
			UpdateFunc: func(old, new interface{}) {
				handleError(npc.UpdateNamespace(old.(*coreapi.Namespace), new.(*coreapi.Namespace)))
			}})

	podController := makeController(client.Core().RESTClient(), "pods", &coreapi.Pod{},
		cache.ResourceEventHandlerFuncs{
			AddFunc: func(obj interface{}) {
				handleError(npc.AddPod(obj.(*coreapi.Pod)))
			},
			DeleteFunc: func(obj interface{}) {
				switch obj := obj.(type) {
				case *coreapi.Pod:
					handleError(npc.DeletePod(obj))
				case cache.DeletedFinalStateUnknown:
					// We know this object has gone away, but its final state is no longer
					// available from the API server. Instead we use the last copy of it
					// that we have, which is good enough for our cleanup.
					handleError(npc.DeletePod(obj.Obj.(*coreapi.Pod)))
				}
			},
			UpdateFunc: func(old, new interface{}) {
				handleError(npc.UpdatePod(old.(*coreapi.Pod), new.(*coreapi.Pod)))
			}})

	npHandlers := cache.ResourceEventHandlerFuncs{
		AddFunc: func(obj interface{}) {
			handleError(npc.AddNetworkPolicy(obj))
		},
		DeleteFunc: func(obj interface{}) {
			switch obj := obj.(type) {
			case cache.DeletedFinalStateUnknown:
				// We know this object has gone away, but its final state is no longer
				// available from the API server. Instead we use the last copy of it
				// that we have, which is good enough for our cleanup.
				handleError(npc.DeleteNetworkPolicy(obj.Obj))
			default:
				handleError(npc.DeleteNetworkPolicy(obj))
			}
		},
		UpdateFunc: func(old, new interface{}) {
			handleError(npc.UpdateNetworkPolicy(old, new))
		},
	}
	npController = makeController(client.NetworkingV1().RESTClient(), "networkpolicies", &networkingv1.NetworkPolicy{}, npHandlers)

	go nsController.Run(wait.NeverStop)
	go podController.Run(wait.NeverStop)
	go npController.Run(wait.NeverStop)

	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
	common.Log.Fatalf("Exiting: %v", <-signals)
}

func main() {
	rootCmd := &cobra.Command{
		Use:   "weave-npc",
		Short: "Weaveworks Kubernetes Network Policy Controller",
		Run:   root}

	rootCmd.PersistentFlags().StringVar(&metricsAddr, "metrics-addr", ":6781", "metrics server bind address")
	rootCmd.PersistentFlags().StringVar(&logLevel, "log-level", "debug", "logging level (debug, info, warning, error)")
	rootCmd.PersistentFlags().BoolVar(&allowMcast, "allow-mcast", true, "allow all multicast traffic")
	rootCmd.PersistentFlags().StringVar(&nodeName, "node-name", "", "only generate rules that apply to this node")
	rootCmd.PersistentFlags().IntVar(&maxList, "max-list-size", 1024, "maximum size of ipset list (for namespaces)")
	rootCmd.PersistentFlags().StringVar(&bridgeName, "bridge-name", "vethwe-bridge", "name of the brige on which packets are received and sent")

	handleError(rootCmd.Execute())
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"os/signal"
	"syscall"

	log "github.com/Sirupsen/logrus"
	"github.com/coreos/go-iptables/iptables"
	"github.com/spf13/cobra"
	"k8s.io/client-go/kubernetes"
	coreapi "k8s.io/client-go/pkg/api/v1"
	extnapi "k8s.io/client-go/pkg/apis/extensions/v1beta1"
	"k8s.io/client-go/pkg/fields"
	"k8s.io/client-go/pkg/runtime"
	"k8s.io/client-go/pkg/util/wait"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/cache"

	"github.com/weaveworks/weave/npc"
	"github.com/weaveworks/weave/npc/ipset"
	"github.com/weaveworks/weave/npc/metrics"
	"github.com/weaveworks/weave/npc/ulogd"
)

var (
	version     = "(unreleased)"
	metricsAddr string
)

func handleError(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

func makeController(getter cache.Getter, resource string,
	objType runtime.Object, handlers cache.ResourceEventHandlerFuncs) *cache.Controller {
	listWatch := cache.NewListWatchFromClient(getter, resource, "", fields.Everything())
	_, controller := cache.NewInformer(listWatch, objType, 0, handlers)
	return controller
}

func resetIPTables(ipt *iptables.IPTables) error {
	// Flush chains first so there are no refs to extant ipsets
	if err := ipt.ClearChain(npc.TableFilter, npc.IngressChain); err != nil {
		return err
	}

	if err := ipt.ClearChain(npc.TableFilter, npc.DefaultChain); err != nil {
		return err
	}

	if err := ipt.ClearChain(npc.TableFilter, npc.MainChain); err != nil {
		return err
	}

	// Configure main chain static rules
	if err := ipt.Append(npc.TableFilter, npc.MainChain,
		"-m", "state", "--state", "RELATED,ESTABLISHED", "-j", "ACCEPT"); err != nil {
		return err
	}

	if err := ipt.Append(npc.TableFilter, npc.MainChain,
		"-m", "state", "--state", "NEW", "-j", string(npc.DefaultChain)); err != nil {
		return err
	}

	if err := ipt.Append(npc.TableFilter, npc.MainChain,
		"-m", "state", "--state", "NEW", "-j", string(npc.IngressChain)); err != nil {
		return err
	}

	return nil
}

func resetIPSets(ips ipset.Interface) error {
	// TODO should restrict ipset operations to the `weave-` prefix:

	if err := ips.FlushAll(); err != nil {
		return err
	}

	if err := ips.DestroyAll(); err != nil {
		return err
	}

	return nil
}

func root(cmd *cobra.Command, args []string) {
	log.Infof("Starting Weaveworks NPC %s", version)

	if err := metrics.Start(metricsAddr); err != nil {
		log.Fatalf("Failed to start metrics: %v", err)
	}

	if err := ulogd.Start(); err != nil {
		log.Fatalf("Failed to start ulogd: %v", err)
	}

	config, err := rest.InClusterConfig()
	if err != nil {
		log.Fatal(err)
	}

	client, err := kubernetes.NewForConfig(config)
	if err != nil {
		log.Fatal(err)
	}

	ipt, err := iptables.New()
	if err != nil {
		log.Fatal(err)
	}

	ips := ipset.New()

	handleError(resetIPTables(ipt))
	handleError(resetIPSets(ips))

	npc := npc.New(ipt, ips)

	nsController := makeController(client.Core().RESTClient(), "namespaces", &coreapi.Namespace{},
		cache.ResourceEventHandlerFuncs{
			AddFunc: func(obj interface{}) {
				handleError(npc.AddNamespace(obj.(*coreapi.Namespace)))
			},
			DeleteFunc: func(obj interface{}) {
				handleError(npc.DeleteNamespace(obj.(*coreapi.Namespace)))
			},
			UpdateFunc: func(old, new interface{}) {
				handleError(npc.UpdateNamespace(old.(*coreapi.Namespace), new.(*coreapi.Namespace)))
			}})

	podController := makeController(client.Core().RESTClient(), "pods", &coreapi.Pod{},
		cache.ResourceEventHandlerFuncs{
			AddFunc: func(obj interface{}) {
				handleError(npc.AddPod(obj.(*coreapi.Pod)))
			},
			DeleteFunc: func(obj interface{}) {
				handleError(npc.DeletePod(obj.(*coreapi.Pod)))
			},
			UpdateFunc: func(old, new interface{}) {
				handleError(npc.UpdatePod(old.(*coreapi.Pod), new.(*coreapi.Pod)))
			}})

	npController := makeController(client.Extensions().RESTClient(), "networkpolicies", &extnapi.NetworkPolicy{},
		cache.ResourceEventHandlerFuncs{
			AddFunc: func(obj interface{}) {
				handleError(npc.AddNetworkPolicy(obj.(*extnapi.NetworkPolicy)))
			},
			DeleteFunc: func(obj interface{}) {
				handleError(npc.DeleteNetworkPolicy(obj.(*extnapi.NetworkPolicy)))
			},
			UpdateFunc: func(old, new interface{}) {
				handleError(npc.UpdateNetworkPolicy(old.(*extnapi.NetworkPolicy), new.(*extnapi.NetworkPolicy)))
			}})

	go nsController.Run(wait.NeverStop)
	go podController.Run(wait.NeverStop)
	go npController.Run(wait.NeverStop)

	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
	log.Fatalf("Exiting: %v", <-signals)
}

func main() {
	rootCmd := &cobra.Command{
		Use:   "weave-npc",
		Short: "Weaveworks Kubernetes Network Policy Controller",
		Run:   root}

	rootCmd.PersistentFlags().StringVar(&metricsAddr, "metrics-addr", ":8686", "metrics server bind address")

	if err := rootCmd.Execute(); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Change default NPC metrics port<commit_after>package main

import (
	"os"
	"os/signal"
	"syscall"

	log "github.com/Sirupsen/logrus"
	"github.com/coreos/go-iptables/iptables"
	"github.com/spf13/cobra"
	"k8s.io/client-go/kubernetes"
	coreapi "k8s.io/client-go/pkg/api/v1"
	extnapi "k8s.io/client-go/pkg/apis/extensions/v1beta1"
	"k8s.io/client-go/pkg/fields"
	"k8s.io/client-go/pkg/runtime"
	"k8s.io/client-go/pkg/util/wait"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/cache"

	"github.com/weaveworks/weave/npc"
	"github.com/weaveworks/weave/npc/ipset"
	"github.com/weaveworks/weave/npc/metrics"
	"github.com/weaveworks/weave/npc/ulogd"
)

var (
	version     = "(unreleased)"
	metricsAddr string
)

func handleError(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

func makeController(getter cache.Getter, resource string,
	objType runtime.Object, handlers cache.ResourceEventHandlerFuncs) *cache.Controller {
	listWatch := cache.NewListWatchFromClient(getter, resource, "", fields.Everything())
	_, controller := cache.NewInformer(listWatch, objType, 0, handlers)
	return controller
}

func resetIPTables(ipt *iptables.IPTables) error {
	// Flush chains first so there are no refs to extant ipsets
	if err := ipt.ClearChain(npc.TableFilter, npc.IngressChain); err != nil {
		return err
	}

	if err := ipt.ClearChain(npc.TableFilter, npc.DefaultChain); err != nil {
		return err
	}

	if err := ipt.ClearChain(npc.TableFilter, npc.MainChain); err != nil {
		return err
	}

	// Configure main chain static rules
	if err := ipt.Append(npc.TableFilter, npc.MainChain,
		"-m", "state", "--state", "RELATED,ESTABLISHED", "-j", "ACCEPT"); err != nil {
		return err
	}

	if err := ipt.Append(npc.TableFilter, npc.MainChain,
		"-m", "state", "--state", "NEW", "-j", string(npc.DefaultChain)); err != nil {
		return err
	}

	if err := ipt.Append(npc.TableFilter, npc.MainChain,
		"-m", "state", "--state", "NEW", "-j", string(npc.IngressChain)); err != nil {
		return err
	}

	return nil
}

func resetIPSets(ips ipset.Interface) error {
	// TODO should restrict ipset operations to the `weave-` prefix:

	if err := ips.FlushAll(); err != nil {
		return err
	}

	if err := ips.DestroyAll(); err != nil {
		return err
	}

	return nil
}

func root(cmd *cobra.Command, args []string) {
	log.Infof("Starting Weaveworks NPC %s", version)

	if err := metrics.Start(metricsAddr); err != nil {
		log.Fatalf("Failed to start metrics: %v", err)
	}

	if err := ulogd.Start(); err != nil {
		log.Fatalf("Failed to start ulogd: %v", err)
	}

	config, err := rest.InClusterConfig()
	if err != nil {
		log.Fatal(err)
	}

	client, err := kubernetes.NewForConfig(config)
	if err != nil {
		log.Fatal(err)
	}

	ipt, err := iptables.New()
	if err != nil {
		log.Fatal(err)
	}

	ips := ipset.New()

	handleError(resetIPTables(ipt))
	handleError(resetIPSets(ips))

	npc := npc.New(ipt, ips)

	nsController := makeController(client.Core().RESTClient(), "namespaces", &coreapi.Namespace{},
		cache.ResourceEventHandlerFuncs{
			AddFunc: func(obj interface{}) {
				handleError(npc.AddNamespace(obj.(*coreapi.Namespace)))
			},
			DeleteFunc: func(obj interface{}) {
				handleError(npc.DeleteNamespace(obj.(*coreapi.Namespace)))
			},
			UpdateFunc: func(old, new interface{}) {
				handleError(npc.UpdateNamespace(old.(*coreapi.Namespace), new.(*coreapi.Namespace)))
			}})

	podController := makeController(client.Core().RESTClient(), "pods", &coreapi.Pod{},
		cache.ResourceEventHandlerFuncs{
			AddFunc: func(obj interface{}) {
				handleError(npc.AddPod(obj.(*coreapi.Pod)))
			},
			DeleteFunc: func(obj interface{}) {
				handleError(npc.DeletePod(obj.(*coreapi.Pod)))
			},
			UpdateFunc: func(old, new interface{}) {
				handleError(npc.UpdatePod(old.(*coreapi.Pod), new.(*coreapi.Pod)))
			}})

	npController := makeController(client.Extensions().RESTClient(), "networkpolicies", &extnapi.NetworkPolicy{},
		cache.ResourceEventHandlerFuncs{
			AddFunc: func(obj interface{}) {
				handleError(npc.AddNetworkPolicy(obj.(*extnapi.NetworkPolicy)))
			},
			DeleteFunc: func(obj interface{}) {
				handleError(npc.DeleteNetworkPolicy(obj.(*extnapi.NetworkPolicy)))
			},
			UpdateFunc: func(old, new interface{}) {
				handleError(npc.UpdateNetworkPolicy(old.(*extnapi.NetworkPolicy), new.(*extnapi.NetworkPolicy)))
			}})

	go nsController.Run(wait.NeverStop)
	go podController.Run(wait.NeverStop)
	go npController.Run(wait.NeverStop)

	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
	log.Fatalf("Exiting: %v", <-signals)
}

func main() {
	rootCmd := &cobra.Command{
		Use:   "weave-npc",
		Short: "Weaveworks Kubernetes Network Policy Controller",
		Run:   root}

	rootCmd.PersistentFlags().StringVar(&metricsAddr, "metrics-addr", ":6781", "metrics server bind address")

	if err := rootCmd.Execute(); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>// +build linux

package overlay

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/containerd/containerd"
	"github.com/containerd/containerd/fs"
	"github.com/containerd/containerd/log"
	"github.com/containerd/containerd/plugin"
	"github.com/containerd/containerd/snapshot"
	"github.com/containerd/containerd/snapshot/storage"
	"github.com/pkg/errors"
)

func init() {
	plugin.Register("snapshot-overlay", &plugin.Registration{
		Type: plugin.SnapshotPlugin,
		Init: func(ic *plugin.InitContext) (interface{}, error) {
			return NewSnapshotter(filepath.Join(ic.Root, "snapshot", "overlay"))
		},
	})
}

type snapshotter struct {
	root string
	ms   *storage.MetaStore
}

type activeSnapshot struct {
	id       string
	name     string
	parentID interface{}
	readonly bool
}

// NewSnapshotter returns a Snapshotter which uses overlayfs. The overlayfs
// diffs are stored under the provided root. A metadata file is stored under
// the root.
func NewSnapshotter(root string) (snapshot.Snapshotter, error) {
	if err := os.MkdirAll(root, 0700); err != nil {
		return nil, err
	}
	ms, err := storage.NewMetaStore(filepath.Join(root, "metadata.db"))
	if err != nil {
		return nil, err
	}

	if err := os.Mkdir(filepath.Join(root, "snapshots"), 0700); err != nil && !os.IsExist(err) {
		return nil, err
	}

	return &snapshotter{
		root: root,
		ms:   ms,
	}, nil
}

// Stat returns the info for an active or committed snapshot by name or
// key.
//
// Should be used for parent resolution, existence checks and to discern
// the kind of snapshot.
func (o *snapshotter) Stat(ctx context.Context, key string) (snapshot.Info, error) {
	ctx, t, err := o.ms.TransactionContext(ctx, false)
	if err != nil {
		return snapshot.Info{}, err
	}
	defer t.Rollback()
	_, info, _, err := storage.GetInfo(ctx, key)
	if err != nil {
		return snapshot.Info{}, err
	}

	return info, nil
}

// Usage returns the resources taken by the snapshot identified by key.
//
// For active snapshots, this will scan the usage of the overlay "diff" (aka
// "upper") directory and may take some time.
//
// For committed snapshots, the value is returned from the metadata database.
func (o *snapshotter) Usage(ctx context.Context, key string) (snapshot.Usage, error) {
	ctx, t, err := o.ms.TransactionContext(ctx, false)
	if err != nil {
		return snapshot.Usage{}, err
	}
	id, info, usage, err := storage.GetInfo(ctx, key)
	if err != nil {
		return snapshot.Usage{}, err
	}

	upperPath := o.upperPath(id)
	t.Rollback() // transaction no longer needed at this point.

	if info.Kind == snapshot.KindActive {
		du, err := fs.DiskUsage(upperPath)
		if err != nil {
			// TODO(stevvooe): Consider not reporting an error in this case.
			return snapshot.Usage{}, err
		}

		usage = snapshot.Usage(du)
	}

	return usage, nil
}

func (o *snapshotter) Prepare(ctx context.Context, key, parent string) ([]containerd.Mount, error) {
	return o.createActive(ctx, key, parent, false)
}

func (o *snapshotter) View(ctx context.Context, key, parent string) ([]containerd.Mount, error) {
	return o.createActive(ctx, key, parent, true)
}

// Mounts returns the mounts for the transaction identified by key. Can be
// called on an read-write or readonly transaction.
//
// This can be used to recover mounts after calling View or Prepare.
func (o *snapshotter) Mounts(ctx context.Context, key string) ([]containerd.Mount, error) {
	ctx, t, err := o.ms.TransactionContext(ctx, false)
	if err != nil {
		return nil, err
	}
	active, err := storage.GetActive(ctx, key)
	t.Rollback()
	if err != nil {
		return nil, errors.Wrap(err, "failed to get active mount")
	}
	return o.mounts(active), nil
}

func (o *snapshotter) Commit(ctx context.Context, name, key string) error {
	ctx, t, err := o.ms.TransactionContext(ctx, true)
	if err != nil {
		return err
	}

	// grab the existing id
	id, _, _, err := storage.GetInfo(ctx, key)
	if err != nil {
		return err
	}

	usage, err := fs.DiskUsage(o.upperPath(id))
	if err != nil {
		return err
	}

	if _, err := storage.CommitActive(ctx, key, name, snapshot.Usage(usage)); err != nil {
		if rerr := t.Rollback(); rerr != nil {
			log.G(ctx).WithError(rerr).Warn("Failure rolling back transaction")
		}
		return errors.Wrap(err, "failed to commit snapshot")
	}
	return t.Commit()
}

// Remove abandons the transaction identified by key. All resources
// associated with the key will be removed.
func (o *snapshotter) Remove(ctx context.Context, key string) (err error) {
	ctx, t, err := o.ms.TransactionContext(ctx, true)
	if err != nil {
		return err
	}
	defer func() {
		if err != nil && t != nil {
			if rerr := t.Rollback(); rerr != nil {
				log.G(ctx).WithError(rerr).Warn("Failure rolling back transaction")
			}
		}
	}()

	id, _, err := storage.Remove(ctx, key)
	if err != nil {
		return errors.Wrap(err, "failed to remove")
	}

	path := filepath.Join(o.root, "snapshots", id)
	renamed := filepath.Join(o.root, "snapshots", "rm-"+id)
	if err := os.Rename(path, renamed); err != nil {
		return errors.Wrap(err, "failed to rename")
	}

	err = t.Commit()
	t = nil
	if err != nil {
		if err1 := os.Rename(renamed, path); err1 != nil {
			// May cause inconsistent data on disk
			log.G(ctx).WithError(err1).WithField("path", renamed).Errorf("Failed to rename after failed commit")
		}
		return errors.Wrap(err, "failed to commit")
	}
	if err := os.RemoveAll(renamed); err != nil {
		// Must be cleaned up, any "rm-*" could be removed if no active transactions
		log.G(ctx).WithError(err).WithField("path", renamed).Warnf("Failed to remove root filesystem")
	}

	return nil
}

// Walk the committed snapshots.
func (o *snapshotter) Walk(ctx context.Context, fn func(context.Context, snapshot.Info) error) error {
	ctx, t, err := o.ms.TransactionContext(ctx, false)
	if err != nil {
		return err
	}
	defer t.Rollback()
	return storage.WalkInfo(ctx, fn)
}

func (o *snapshotter) createActive(ctx context.Context, key, parent string, readonly bool) ([]containerd.Mount, error) {
	var (
		path        string
		snapshotDir = filepath.Join(o.root, "snapshots")
	)

	td, err := ioutil.TempDir(snapshotDir, "new-")
	if err != nil {
		return nil, errors.Wrap(err, "failed to create temp dir")
	}
	defer func() {
		if err != nil {
			if td != "" {
				if err1 := os.RemoveAll(td); err1 != nil {
					err = errors.Wrapf(err, "remove failed: %v", err1)
				}
			}
			if path != "" {
				if err1 := os.RemoveAll(path); err1 != nil {
					err = errors.Wrapf(err, "failed to remove path: %v", err1)
				}
			}
		}
	}()

	if err = os.MkdirAll(filepath.Join(td, "fs"), 0711); err != nil {
		return nil, err
	}
	if !readonly {
		if err = os.MkdirAll(filepath.Join(td, "work"), 0700); err != nil {
			return nil, err
		}
	}

	ctx, t, err := o.ms.TransactionContext(ctx, true)
	if err != nil {
		return nil, err
	}

	active, err := storage.CreateActive(ctx, key, parent, readonly)
	if err != nil {
		if rerr := t.Rollback(); rerr != nil {
			log.G(ctx).WithError(rerr).Warn("Failure rolling back transaction")
		}
		return nil, errors.Wrap(err, "failed to create active")
	}

	path = filepath.Join(snapshotDir, active.ID)
	if err = os.Rename(td, path); err != nil {
		if rerr := t.Rollback(); rerr != nil {
			log.G(ctx).WithError(rerr).Warn("Failure rolling back transaction")
		}
		return nil, errors.Wrap(err, "failed to rename")
	}
	td = ""

	if err = t.Commit(); err != nil {
		return nil, errors.Wrap(err, "commit failed")
	}

	return o.mounts(active), nil
}

func (o *snapshotter) mounts(active storage.Active) []containerd.Mount {
	if len(active.ParentIDs) == 0 {
		// if we only have one layer/no parents then just return a bind mount as overlay
		// will not work
		roFlag := "rw"
		if active.Readonly {
			roFlag = "ro"
		}

		return []containerd.Mount{
			{
				Source: o.upperPath(active.ID),
				Type:   "bind",
				Options: []string{
					roFlag,
					"rbind",
				},
			},
		}
	}
	var options []string

	if !active.Readonly {
		options = append(options,
			fmt.Sprintf("workdir=%s", o.workPath(active.ID)),
			fmt.Sprintf("upperdir=%s", o.upperPath(active.ID)),
		)
	} else if len(active.ParentIDs) == 1 {
		return []containerd.Mount{
			{
				Source: o.upperPath(active.ParentIDs[0]),
				Type:   "bind",
				Options: []string{
					"ro",
					"rbind",
				},
			},
		}
	}

	parentPaths := make([]string, len(active.ParentIDs))
	for i := range active.ParentIDs {
		parentPaths[i] = o.upperPath(active.ParentIDs[i])
	}

	options = append(options, fmt.Sprintf("lowerdir=%s", strings.Join(parentPaths, ":")))
	return []containerd.Mount{
		{
			Type:    "overlay",
			Source:  "overlay",
			Options: options,
		},
	}

}

func (o *snapshotter) upperPath(id string) string {
	return filepath.Join(o.root, "snapshots", id, "fs")
}

func (o *snapshotter) workPath(id string) string {
	return filepath.Join(o.root, "snapshots", id, "work")
}
<commit_msg>overlay: abort if backing fs lacks support for d_type<commit_after>// +build linux

package overlay

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/containerd/containerd"
	"github.com/containerd/containerd/fs"
	"github.com/containerd/containerd/fs/fsutils"
	"github.com/containerd/containerd/log"
	"github.com/containerd/containerd/plugin"
	"github.com/containerd/containerd/snapshot"
	"github.com/containerd/containerd/snapshot/storage"
	"github.com/pkg/errors"
)

func init() {
	plugin.Register("snapshot-overlay", &plugin.Registration{
		Type: plugin.SnapshotPlugin,
		Init: func(ic *plugin.InitContext) (interface{}, error) {
			return NewSnapshotter(filepath.Join(ic.Root, "snapshot", "overlay"))
		},
	})
}

type snapshotter struct {
	root string
	ms   *storage.MetaStore
}

type activeSnapshot struct {
	id       string
	name     string
	parentID interface{}
	readonly bool
}

// NewSnapshotter returns a Snapshotter which uses overlayfs. The overlayfs
// diffs are stored under the provided root. A metadata file is stored under
// the root.
func NewSnapshotter(root string) (snapshot.Snapshotter, error) {
	if err := os.MkdirAll(root, 0700); err != nil {
		return nil, err
	}
	supportsDType, err := fsutils.SupportsDType(root)
	if err != nil {
		return nil, err
	}
	if !supportsDType {
		return nil, fmt.Errorf("%s does not support d_type. If the backing filesystem is xfs, please reformat with ftype=1 to enable d_type support.", root)
	}
	ms, err := storage.NewMetaStore(filepath.Join(root, "metadata.db"))
	if err != nil {
		return nil, err
	}

	if err := os.Mkdir(filepath.Join(root, "snapshots"), 0700); err != nil && !os.IsExist(err) {
		return nil, err
	}

	return &snapshotter{
		root: root,
		ms:   ms,
	}, nil
}

// Stat returns the info for an active or committed snapshot by name or
// key.
//
// Should be used for parent resolution, existence checks and to discern
// the kind of snapshot.
func (o *snapshotter) Stat(ctx context.Context, key string) (snapshot.Info, error) {
	ctx, t, err := o.ms.TransactionContext(ctx, false)
	if err != nil {
		return snapshot.Info{}, err
	}
	defer t.Rollback()
	_, info, _, err := storage.GetInfo(ctx, key)
	if err != nil {
		return snapshot.Info{}, err
	}

	return info, nil
}

// Usage returns the resources taken by the snapshot identified by key.
//
// For active snapshots, this will scan the usage of the overlay "diff" (aka
// "upper") directory and may take some time.
//
// For committed snapshots, the value is returned from the metadata database.
func (o *snapshotter) Usage(ctx context.Context, key string) (snapshot.Usage, error) {
	ctx, t, err := o.ms.TransactionContext(ctx, false)
	if err != nil {
		return snapshot.Usage{}, err
	}
	id, info, usage, err := storage.GetInfo(ctx, key)
	if err != nil {
		return snapshot.Usage{}, err
	}

	upperPath := o.upperPath(id)
	t.Rollback() // transaction no longer needed at this point.

	if info.Kind == snapshot.KindActive {
		du, err := fs.DiskUsage(upperPath)
		if err != nil {
			// TODO(stevvooe): Consider not reporting an error in this case.
			return snapshot.Usage{}, err
		}

		usage = snapshot.Usage(du)
	}

	return usage, nil
}

func (o *snapshotter) Prepare(ctx context.Context, key, parent string) ([]containerd.Mount, error) {
	return o.createActive(ctx, key, parent, false)
}

func (o *snapshotter) View(ctx context.Context, key, parent string) ([]containerd.Mount, error) {
	return o.createActive(ctx, key, parent, true)
}

// Mounts returns the mounts for the transaction identified by key. Can be
// called on an read-write or readonly transaction.
//
// This can be used to recover mounts after calling View or Prepare.
func (o *snapshotter) Mounts(ctx context.Context, key string) ([]containerd.Mount, error) {
	ctx, t, err := o.ms.TransactionContext(ctx, false)
	if err != nil {
		return nil, err
	}
	active, err := storage.GetActive(ctx, key)
	t.Rollback()
	if err != nil {
		return nil, errors.Wrap(err, "failed to get active mount")
	}
	return o.mounts(active), nil
}

func (o *snapshotter) Commit(ctx context.Context, name, key string) error {
	ctx, t, err := o.ms.TransactionContext(ctx, true)
	if err != nil {
		return err
	}

	// grab the existing id
	id, _, _, err := storage.GetInfo(ctx, key)
	if err != nil {
		return err
	}

	usage, err := fs.DiskUsage(o.upperPath(id))
	if err != nil {
		return err
	}

	if _, err := storage.CommitActive(ctx, key, name, snapshot.Usage(usage)); err != nil {
		if rerr := t.Rollback(); rerr != nil {
			log.G(ctx).WithError(rerr).Warn("Failure rolling back transaction")
		}
		return errors.Wrap(err, "failed to commit snapshot")
	}
	return t.Commit()
}

// Remove abandons the transaction identified by key. All resources
// associated with the key will be removed.
func (o *snapshotter) Remove(ctx context.Context, key string) (err error) {
	ctx, t, err := o.ms.TransactionContext(ctx, true)
	if err != nil {
		return err
	}
	defer func() {
		if err != nil && t != nil {
			if rerr := t.Rollback(); rerr != nil {
				log.G(ctx).WithError(rerr).Warn("Failure rolling back transaction")
			}
		}
	}()

	id, _, err := storage.Remove(ctx, key)
	if err != nil {
		return errors.Wrap(err, "failed to remove")
	}

	path := filepath.Join(o.root, "snapshots", id)
	renamed := filepath.Join(o.root, "snapshots", "rm-"+id)
	if err := os.Rename(path, renamed); err != nil {
		return errors.Wrap(err, "failed to rename")
	}

	err = t.Commit()
	t = nil
	if err != nil {
		if err1 := os.Rename(renamed, path); err1 != nil {
			// May cause inconsistent data on disk
			log.G(ctx).WithError(err1).WithField("path", renamed).Errorf("Failed to rename after failed commit")
		}
		return errors.Wrap(err, "failed to commit")
	}
	if err := os.RemoveAll(renamed); err != nil {
		// Must be cleaned up, any "rm-*" could be removed if no active transactions
		log.G(ctx).WithError(err).WithField("path", renamed).Warnf("Failed to remove root filesystem")
	}

	return nil
}

// Walk the committed snapshots.
func (o *snapshotter) Walk(ctx context.Context, fn func(context.Context, snapshot.Info) error) error {
	ctx, t, err := o.ms.TransactionContext(ctx, false)
	if err != nil {
		return err
	}
	defer t.Rollback()
	return storage.WalkInfo(ctx, fn)
}

func (o *snapshotter) createActive(ctx context.Context, key, parent string, readonly bool) ([]containerd.Mount, error) {
	var (
		path        string
		snapshotDir = filepath.Join(o.root, "snapshots")
	)

	td, err := ioutil.TempDir(snapshotDir, "new-")
	if err != nil {
		return nil, errors.Wrap(err, "failed to create temp dir")
	}
	defer func() {
		if err != nil {
			if td != "" {
				if err1 := os.RemoveAll(td); err1 != nil {
					err = errors.Wrapf(err, "remove failed: %v", err1)
				}
			}
			if path != "" {
				if err1 := os.RemoveAll(path); err1 != nil {
					err = errors.Wrapf(err, "failed to remove path: %v", err1)
				}
			}
		}
	}()

	if err = os.MkdirAll(filepath.Join(td, "fs"), 0711); err != nil {
		return nil, err
	}
	if !readonly {
		if err = os.MkdirAll(filepath.Join(td, "work"), 0700); err != nil {
			return nil, err
		}
	}

	ctx, t, err := o.ms.TransactionContext(ctx, true)
	if err != nil {
		return nil, err
	}

	active, err := storage.CreateActive(ctx, key, parent, readonly)
	if err != nil {
		if rerr := t.Rollback(); rerr != nil {
			log.G(ctx).WithError(rerr).Warn("Failure rolling back transaction")
		}
		return nil, errors.Wrap(err, "failed to create active")
	}

	path = filepath.Join(snapshotDir, active.ID)
	if err = os.Rename(td, path); err != nil {
		if rerr := t.Rollback(); rerr != nil {
			log.G(ctx).WithError(rerr).Warn("Failure rolling back transaction")
		}
		return nil, errors.Wrap(err, "failed to rename")
	}
	td = ""

	if err = t.Commit(); err != nil {
		return nil, errors.Wrap(err, "commit failed")
	}

	return o.mounts(active), nil
}

func (o *snapshotter) mounts(active storage.Active) []containerd.Mount {
	if len(active.ParentIDs) == 0 {
		// if we only have one layer/no parents then just return a bind mount as overlay
		// will not work
		roFlag := "rw"
		if active.Readonly {
			roFlag = "ro"
		}

		return []containerd.Mount{
			{
				Source: o.upperPath(active.ID),
				Type:   "bind",
				Options: []string{
					roFlag,
					"rbind",
				},
			},
		}
	}
	var options []string

	if !active.Readonly {
		options = append(options,
			fmt.Sprintf("workdir=%s", o.workPath(active.ID)),
			fmt.Sprintf("upperdir=%s", o.upperPath(active.ID)),
		)
	} else if len(active.ParentIDs) == 1 {
		return []containerd.Mount{
			{
				Source: o.upperPath(active.ParentIDs[0]),
				Type:   "bind",
				Options: []string{
					"ro",
					"rbind",
				},
			},
		}
	}

	parentPaths := make([]string, len(active.ParentIDs))
	for i := range active.ParentIDs {
		parentPaths[i] = o.upperPath(active.ParentIDs[i])
	}

	options = append(options, fmt.Sprintf("lowerdir=%s", strings.Join(parentPaths, ":")))
	return []containerd.Mount{
		{
			Type:    "overlay",
			Source:  "overlay",
			Options: options,
		},
	}

}

func (o *snapshotter) upperPath(id string) string {
	return filepath.Join(o.root, "snapshots", id, "fs")
}

func (o *snapshotter) workPath(id string) string {
	return filepath.Join(o.root, "snapshots", id, "work")
}
<|endoftext|>


<commit_before>package testing_test

import (
	"bufio"
	"bytes"
	"context"
	"github.com/influxdata/flux/execute/executetest"
	"strings"
	"testing"

	"github.com/influxdata/flux"
	"github.com/influxdata/flux/ast"
	"github.com/influxdata/flux/execute"
	"github.com/influxdata/flux/lang"
	"github.com/influxdata/flux/parser"
	"github.com/influxdata/flux/stdlib"

	platform "github.com/influxdata/influxdb"
	"github.com/influxdata/influxdb/cmd/influxd/launcher"
	influxdbcontext "github.com/influxdata/influxdb/context"
	"github.com/influxdata/influxdb/mock"
	"github.com/influxdata/influxdb/query"
	itesting "github.com/influxdata/influxdb/query/stdlib/testing"

	_ "github.com/influxdata/influxdb/query/stdlib" // Import the stdlib
)

// Default context.
var ctx = influxdbcontext.SetAuthorizer(context.Background(), &mock.Authorization{})

func init() {
	flux.FinalizeBuiltIns()
}

func TestFluxEndToEnd(t *testing.T) {
	runEndToEnd(t, stdlib.FluxTestPackages)
}
func BenchmarkFluxEndToEnd(b *testing.B) {
	benchEndToEnd(b, stdlib.FluxTestPackages)
}

func runEndToEnd(t *testing.T, pkgs []*ast.Package) {
	l := launcher.RunTestLauncherOrFail(t, ctx)
	l.SetupOrFail(t)
	defer l.ShutdownOrFail(t, ctx)
	for _, pkg := range pkgs {
		pkg := pkg.Copy().(*ast.Package)
		name := strings.TrimSuffix(pkg.Files[0].Name, "_test.flux")
		t.Run(name, func(t *testing.T) {
			if reason, ok := itesting.FluxEndToEndSkipList[name]; ok {
				t.Skip(reason)
			}
			c := lang.ASTCompiler{AST: pkg}
			testFlux(t, l, pkg, c)
		})
	}
}

func benchEndToEnd(b *testing.B, pkgs []*ast.Package) {
	l := launcher.RunTestLauncherOrFail(b, ctx)
	l.SetupOrFail(b)
	defer l.ShutdownOrFail(b, ctx)
	for _, pkg := range pkgs {
		pkg := pkg.Copy().(*ast.Package)
		name := pkg.Files[0].Name
		b.Run(name, func(b *testing.B) {
			if reason, ok := itesting.FluxEndToEndSkipList[strings.TrimSuffix(name, ".flux")]; ok {
				b.Skip(reason)
			}
			b.ResetTimer()
			b.ReportAllocs()
			for i := 0; i < b.N; i++ {
				c := lang.ASTCompiler{AST: pkg}
				testFlux(b, l, pkg, c)
			}
		})
	}
}

var optionsSource = `
import "testing"
import c "csv"

// Options bucket and org are defined dynamically per test

option testing.loadStorage = (csv) => {
	c.from(csv: csv) |> to(bucket: bucket, org: org)
	return from(bucket: bucket)
}
`
var optionsAST *ast.File

func init() {
	pkg := parser.ParseSource(optionsSource)
	if ast.Check(pkg) > 0 {
		panic(ast.GetError(pkg))
	}
	optionsAST = pkg.Files[0]
}

func testFlux(t testing.TB, l *launcher.TestLauncher, pkg *ast.Package, c flux.Compiler) {

	// Query server to ensure write persists.

	b := &platform.Bucket{
		OrgID:           l.Org.ID,
		Name:            t.Name(),
		RetentionPeriod: 0,
	}

	s := l.BucketService()
	if err := s.CreateBucket(context.Background(), b); err != nil {
		t.Fatal(err)
	}

	// Define bucket and org options
	bucketOpt := &ast.OptionStatement{
		Assignment: &ast.VariableAssignment{
			ID:   &ast.Identifier{Name: "bucket"},
			Init: &ast.StringLiteral{Value: b.Name},
		},
	}

	orgOpt := &ast.OptionStatement{
		Assignment: &ast.VariableAssignment{
			ID:   &ast.Identifier{Name: "org"},
			Init: &ast.StringLiteral{Value: l.Org.Name},
		},
	}
	options := optionsAST.Copy().(*ast.File)
	options.Body = append([]ast.Statement{bucketOpt, orgOpt}, options.Body...)

	// Add options to pkg
	pkg.Files = append(pkg.Files, options)

	// Add testing.inspect call to ensure the data is loaded
	inspectCalls := stdlib.TestingInspectCalls(pkg)
	pkg.Files = append(pkg.Files, inspectCalls)

	req := &query.Request{
		OrganizationID: l.Org.ID,
		Compiler:       c,
	}

	program, err := c.Compile(context.Background())
	if p, ok := program.(lang.DependenciesAwareProgram); ok {
		p.SetExecutorDependencies(executetest.NewTestExecuteDependencies())
	}

	if r, err := l.FluxQueryService().Query(ctx, req); err != nil {
		t.Fatal(err)
	} else {
		for r.More() {
			v := r.Next()
			if err := v.Tables().Do(func(tbl flux.Table) error {
				return tbl.Do(func(reader flux.ColReader) error {
					return nil
				})
			}); err != nil {
				t.Error(err)
			}
		}
	}

	// quirk: our execution engine doesn't guarantee the order of execution for disconnected DAGS
	// so that our function-with-side effects call to `to` may run _after_ the test instead of before.
	// running twice makes sure that `to` happens at least once before we run the test.
	// this time we use a call to `run` so that the assertion error is triggered
	runCalls := stdlib.TestingRunCalls(pkg)
	pkg.Files[len(pkg.Files)-1] = runCalls
	r, err := l.FluxQueryService().Query(ctx, req)
	if err != nil {
		t.Fatal(err)
	}

	for r.More() {
		v := r.Next()
		if err := v.Tables().Do(func(tbl flux.Table) error {
			return tbl.Do(func(reader flux.ColReader) error {
				return nil
			})
		}); err != nil {
			t.Error(err)
		}
	}
	if err := r.Err(); err != nil {
		t.Error(err)
		// Replace the testing.run calls with testing.inspect calls.
		pkg.Files[len(pkg.Files)-1] = inspectCalls
		r, err := l.FluxQueryService().Query(ctx, req)
		if err != nil {
			t.Fatal(err)
		}
		var out bytes.Buffer
		defer func() {
			if t.Failed() {
				scanner := bufio.NewScanner(&out)
				for scanner.Scan() {
					t.Log(scanner.Text())
				}
			}
		}()
		for r.More() {
			v := r.Next()
			err := execute.FormatResult(&out, v)
			if err != nil {
				t.Error(err)
			}
		}
		if err := r.Err(); err != nil {
			t.Error(err)
		}
	}
}
<commit_msg>test(query): roll back some changes to end_to_end_test.go<commit_after>package testing_test

import (
	"bufio"
	"bytes"
	"context"
	"strings"
	"testing"

	"github.com/influxdata/flux"
	"github.com/influxdata/flux/ast"
	"github.com/influxdata/flux/execute"
	"github.com/influxdata/flux/lang"
	"github.com/influxdata/flux/parser"
	"github.com/influxdata/flux/stdlib"

	platform "github.com/influxdata/influxdb"
	"github.com/influxdata/influxdb/cmd/influxd/launcher"
	influxdbcontext "github.com/influxdata/influxdb/context"
	"github.com/influxdata/influxdb/mock"
	"github.com/influxdata/influxdb/query"
	itesting "github.com/influxdata/influxdb/query/stdlib/testing"

	_ "github.com/influxdata/influxdb/query/stdlib" // Import the stdlib
)

// Default context.
var ctx = influxdbcontext.SetAuthorizer(context.Background(), &mock.Authorization{})

func init() {
	flux.FinalizeBuiltIns()
}

func TestFluxEndToEnd(t *testing.T) {
	runEndToEnd(t, stdlib.FluxTestPackages)
}
func BenchmarkFluxEndToEnd(b *testing.B) {
	benchEndToEnd(b, stdlib.FluxTestPackages)
}

func runEndToEnd(t *testing.T, pkgs []*ast.Package) {
	l := launcher.RunTestLauncherOrFail(t, ctx)
	l.SetupOrFail(t)
	defer l.ShutdownOrFail(t, ctx)
	for _, pkg := range pkgs {
		pkg := pkg.Copy().(*ast.Package)
		name := strings.TrimSuffix(pkg.Files[0].Name, "_test.flux")
		t.Run(name, func(t *testing.T) {
			if reason, ok := itesting.FluxEndToEndSkipList[name]; ok {
				t.Skip(reason)
			}
			testFlux(t, l, pkg)
		})
	}
}

func benchEndToEnd(b *testing.B, pkgs []*ast.Package) {
	l := launcher.RunTestLauncherOrFail(b, ctx)
	l.SetupOrFail(b)
	defer l.ShutdownOrFail(b, ctx)
	for _, pkg := range pkgs {
		pkg := pkg.Copy().(*ast.Package)
		name := pkg.Files[0].Name
		b.Run(name, func(b *testing.B) {
			if reason, ok := itesting.FluxEndToEndSkipList[strings.TrimSuffix(name, ".flux")]; ok {
				b.Skip(reason)
			}
			b.ResetTimer()
			b.ReportAllocs()
			for i := 0; i < b.N; i++ {
				testFlux(b, l, pkg)
			}
		})
	}
}

var optionsSource = `
import "testing"
import c "csv"

// Options bucket and org are defined dynamically per test

option testing.loadStorage = (csv) => {
	c.from(csv: csv) |> to(bucket: bucket, org: org)
	return from(bucket: bucket)
}
`
var optionsAST *ast.File

func init() {
	pkg := parser.ParseSource(optionsSource)
	if ast.Check(pkg) > 0 {
		panic(ast.GetError(pkg))
	}
	optionsAST = pkg.Files[0]
}

func testFlux(t testing.TB, l *launcher.TestLauncher, pkg *ast.Package) {

	// Query server to ensure write persists.

	b := &platform.Bucket{
		OrgID:           l.Org.ID,
		Name:            t.Name(),
		RetentionPeriod: 0,
	}

	s := l.BucketService()
	if err := s.CreateBucket(context.Background(), b); err != nil {
		t.Fatal(err)
	}

	// Define bucket and org options
	bucketOpt := &ast.OptionStatement{
		Assignment: &ast.VariableAssignment{
			ID:   &ast.Identifier{Name: "bucket"},
			Init: &ast.StringLiteral{Value: b.Name},
		},
	}
	orgOpt := &ast.OptionStatement{
		Assignment: &ast.VariableAssignment{
			ID:   &ast.Identifier{Name: "org"},
			Init: &ast.StringLiteral{Value: l.Org.Name},
		},
	}
	options := optionsAST.Copy().(*ast.File)
	options.Body = append([]ast.Statement{bucketOpt, orgOpt}, options.Body...)

	// Add options to pkg
	pkg.Files = append(pkg.Files, options)

	// Add testing.inspect call to ensure the data is loaded
	inspectCalls := stdlib.TestingInspectCalls(pkg)
	pkg.Files = append(pkg.Files, inspectCalls)

	req := &query.Request{
		OrganizationID: l.Org.ID,
		Compiler:       lang.ASTCompiler{AST: pkg},
	}
	if r, err := l.FluxQueryService().Query(ctx, req); err != nil {
		t.Fatal(err)
	} else {
		for r.More() {
			v := r.Next()
			if err := v.Tables().Do(func(tbl flux.Table) error {
				return tbl.Do(func(reader flux.ColReader) error {
					return nil
				})
			}); err != nil {
				t.Error(err)
			}
		}
	}

	// quirk: our execution engine doesn't guarantee the order of execution for disconnected DAGS
	// so that our function-with-side effects call to `to` may run _after_ the test instead of before.
	// running twice makes sure that `to` happens at least once before we run the test.
	// this time we use a call to `run` so that the assertion error is triggered
	runCalls := stdlib.TestingRunCalls(pkg)
	pkg.Files[len(pkg.Files)-1] = runCalls
	r, err := l.FluxQueryService().Query(ctx, req)
	if err != nil {
		t.Fatal(err)
	}

	for r.More() {
		v := r.Next()
		if err := v.Tables().Do(func(tbl flux.Table) error {
			return tbl.Do(func(reader flux.ColReader) error {
				return nil
			})
		}); err != nil {
			t.Error(err)
		}
	}
	if err := r.Err(); err != nil {
		t.Error(err)
		// Replace the testing.run calls with testing.inspect calls.
		pkg.Files[len(pkg.Files)-1] = inspectCalls
		r, err := l.FluxQueryService().Query(ctx, req)
		if err != nil {
			t.Fatal(err)
		}
		var out bytes.Buffer
		defer func() {
			if t.Failed() {
				scanner := bufio.NewScanner(&out)
				for scanner.Scan() {
					t.Log(scanner.Text())
				}
			}
		}()
		for r.More() {
			v := r.Next()
			err := execute.FormatResult(&out, v)
			if err != nil {
				t.Error(err)
			}
		}
		if err := r.Err(); err != nil {
			t.Error(err)
		}
	}
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"regexp"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	elasticsearch "github.com/aws/aws-sdk-go/service/elasticsearchservice"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsElasticSearchDomain() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsElasticSearchDomainCreate,
		Read:   resourceAwsElasticSearchDomainRead,
		Update: resourceAwsElasticSearchDomainUpdate,
		Delete: resourceAwsElasticSearchDomainDelete,

		Schema: map[string]*schema.Schema{
			"access_policies": &schema.Schema{
				Type:      schema.TypeString,
				StateFunc: normalizeJson,
				Optional:  true,
			},
			"advanced_options": &schema.Schema{
				Type:     schema.TypeMap,
				Optional: true,
				Computed: true,
			},
			"domain_name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
				ValidateFunc: func(v interface{}, k string) (ws []string, errors []error) {
					value := v.(string)
					if !regexp.MustCompile(`^[0-9A-Za-z]+`).MatchString(value) {
						errors = append(errors, fmt.Errorf(
							"%q must start with a letter or number", k))
					}
					if !regexp.MustCompile(`^[0-9A-Za-z][0-9a-z-]+$`).MatchString(value) {
						errors = append(errors, fmt.Errorf(
							"%q can only contain lowercase characters, numbers and hyphens", k))
					}
					return
				},
			},
			"arn": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"domain_id": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"endpoint": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"ebs_options": &schema.Schema{
				Type:     schema.TypeList,
				Optional: true,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"ebs_enabled": &schema.Schema{
							Type:     schema.TypeBool,
							Required: true,
						},
						"iops": &schema.Schema{
							Type:     schema.TypeInt,
							Optional: true,
						},
						"volume_size": &schema.Schema{
							Type:     schema.TypeInt,
							Optional: true,
						},
						"volume_type": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
						},
					},
				},
			},
			"cluster_config": &schema.Schema{
				Type:     schema.TypeList,
				Optional: true,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"dedicated_master_count": &schema.Schema{
							Type:     schema.TypeInt,
							Optional: true,
						},
						"dedicated_master_enabled": &schema.Schema{
							Type:     schema.TypeBool,
							Optional: true,
							Default:  false,
						},
						"dedicated_master_type": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
						},
						"instance_count": &schema.Schema{
							Type:     schema.TypeInt,
							Optional: true,
							Default:  1,
						},
						"instance_type": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
							Default:  "m3.medium.elasticsearch",
						},
						"zone_awareness_enabled": &schema.Schema{
							Type:     schema.TypeBool,
							Optional: true,
						},
					},
				},
			},
			"snapshot_options": &schema.Schema{
				Type:     schema.TypeList,
				Optional: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"automated_snapshot_start_hour": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},
					},
				},
			},
		},
	}
}

func resourceAwsElasticSearchDomainCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).esconn

	input := elasticsearch.CreateElasticsearchDomainInput{
		DomainName: aws.String(d.Get("domain_name").(string)),
	}

	if v, ok := d.GetOk("access_policies"); ok {
		input.AccessPolicies = aws.String(v.(string))
	}

	if v, ok := d.GetOk("advanced_options"); ok {
		input.AdvancedOptions = stringMapToPointers(v.(map[string]interface{}))
	}

	if v, ok := d.GetOk("ebs_options"); ok {
		options := v.([]interface{})

		if len(options) > 1 {
			return fmt.Errorf("Only a single ebs_options block is expected")
		} else if len(options) == 1 {
			if options[0] == nil {
				return fmt.Errorf("At least one field is expected inside ebs_options")
			}

			s := options[0].(map[string]interface{})
			input.EBSOptions = expandESEBSOptions(s)
		}
	}

	if v, ok := d.GetOk("cluster_config"); ok {
		config := v.([]interface{})

		if len(config) > 1 {
			return fmt.Errorf("Only a single cluster_config block is expected")
		} else if len(config) == 1 {
			if config[0] == nil {
				return fmt.Errorf("At least one field is expected inside cluster_config")
			}
			m := config[0].(map[string]interface{})
			input.ElasticsearchClusterConfig = expandESClusterConfig(m)
		}
	}

	if v, ok := d.GetOk("snapshot_options"); ok {
		options := v.([]interface{})

		if len(options) > 1 {
			return fmt.Errorf("Only a single snapshot_options block is expected")
		} else if len(options) == 1 {
			if options[0] == nil {
				return fmt.Errorf("At least one field is expected inside snapshot_options")
			}

			o := options[0].(map[string]interface{})

			snapshotOptions := elasticsearch.SnapshotOptions{
				AutomatedSnapshotStartHour: aws.Int64(int64(o["automated_snapshot_start_hour"].(int))),
			}

			input.SnapshotOptions = &snapshotOptions
		}
	}

	log.Printf("[DEBUG] Creating ElasticSearch domain: %s", input)
	out, err := conn.CreateElasticsearchDomain(&input)
	if err != nil {
		return err
	}

	d.SetId(*out.DomainStatus.ARN)

	log.Printf("[DEBUG] Waiting for ElasticSearch domain %q to be created", d.Id())
	err = resource.Retry(15*time.Minute, func() error {
		out, err := conn.DescribeElasticsearchDomain(&elasticsearch.DescribeElasticsearchDomainInput{
			DomainName: aws.String(d.Get("domain_name").(string)),
		})
		if err != nil {
			return resource.RetryError{Err: err}
		}

		if !*out.DomainStatus.Processing && out.DomainStatus.Endpoint != nil {
			return nil
		}

		return fmt.Errorf("%q: Timeout while waiting for the domain to be created", d.Id())
	})
	if err != nil {
		return err
	}

	log.Printf("[DEBUG] ElasticSearch domain %q created", d.Id())

	return resourceAwsElasticSearchDomainRead(d, meta)
}

func resourceAwsElasticSearchDomainRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).esconn

	out, err := conn.DescribeElasticsearchDomain(&elasticsearch.DescribeElasticsearchDomainInput{
		DomainName: aws.String(d.Get("domain_name").(string)),
	})
	if err != nil {
		return err
	}

	log.Printf("[DEBUG] Received ElasticSearch domain: %s", out)

	ds := out.DomainStatus

	if ds.AccessPolicies != nil && *ds.AccessPolicies != "" {
		d.Set("access_policies", normalizeJson(*ds.AccessPolicies))
	}
	err = d.Set("advanced_options", pointersMapToStringList(ds.AdvancedOptions))
	if err != nil {
		return err
	}
	d.Set("domain_id", *ds.DomainId)
	d.Set("domain_name", *ds.DomainName)
	if ds.Endpoint != nil {
		d.Set("endpoint", *ds.Endpoint)
	}

	err = d.Set("ebs_options", flattenESEBSOptions(ds.EBSOptions))
	if err != nil {
		return err
	}
	err = d.Set("cluster_config", flattenESClusterConfig(ds.ElasticsearchClusterConfig))
	if err != nil {
		return err
	}
	if ds.SnapshotOptions != nil {
		d.Set("snapshot_options", map[string]interface{}{
			"automated_snapshot_start_hour": *ds.SnapshotOptions.AutomatedSnapshotStartHour,
		})
	}

	d.Set("arn", *ds.ARN)

	return nil
}

func resourceAwsElasticSearchDomainUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).esconn

	input := elasticsearch.UpdateElasticsearchDomainConfigInput{
		DomainName: aws.String(d.Get("domain_name").(string)),
	}

	if d.HasChange("access_policies") {
		input.AccessPolicies = aws.String(d.Get("access_policies").(string))
	}

	if d.HasChange("advanced_options") {
		input.AdvancedOptions = stringMapToPointers(d.Get("advanced_options").(map[string]interface{}))
	}

	if d.HasChange("ebs_options") {
		options := d.Get("ebs_options").([]interface{})

		if len(options) > 1 {
			return fmt.Errorf("Only a single ebs_options block is expected")
		} else if len(options) == 1 {
			s := options[0].(map[string]interface{})
			input.EBSOptions = expandESEBSOptions(s)
		}
	}

	if d.HasChange("cluster_config") {
		config := d.Get("cluster_config").([]interface{})

		if len(config) > 1 {
			return fmt.Errorf("Only a single cluster_config block is expected")
		} else if len(config) == 1 {
			m := config[0].(map[string]interface{})
			input.ElasticsearchClusterConfig = expandESClusterConfig(m)
		}
	}

	if d.HasChange("snapshot_options") {
		options := d.Get("snapshot_options").([]interface{})

		if len(options) > 1 {
			return fmt.Errorf("Only a single snapshot_options block is expected")
		} else if len(options) == 1 {
			o := options[0].(map[string]interface{})

			snapshotOptions := elasticsearch.SnapshotOptions{
				AutomatedSnapshotStartHour: aws.Int64(int64(o["automated_snapshot_start_hour"].(int))),
			}

			input.SnapshotOptions = &snapshotOptions
		}
	}

	_, err := conn.UpdateElasticsearchDomainConfig(&input)
	if err != nil {
		return err
	}

	err = resource.Retry(25*time.Minute, func() error {
		out, err := conn.DescribeElasticsearchDomain(&elasticsearch.DescribeElasticsearchDomainInput{
			DomainName: aws.String(d.Get("domain_name").(string)),
		})
		if err != nil {
			return resource.RetryError{Err: err}
		}

		if *out.DomainStatus.Processing == false {
			return nil
		}

		return fmt.Errorf("%q: Timeout while waiting for changes to be processed", d.Id())
	})
	if err != nil {
		return err
	}

	return resourceAwsElasticSearchDomainRead(d, meta)
}

func resourceAwsElasticSearchDomainDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).esconn

	log.Printf("[DEBUG] Deleting ElasticSearch domain: %q", d.Get("domain_name").(string))
	_, err := conn.DeleteElasticsearchDomain(&elasticsearch.DeleteElasticsearchDomainInput{
		DomainName: aws.String(d.Get("domain_name").(string)),
	})
	if err != nil {
		return err
	}

	log.Printf("[DEBUG] Waiting for ElasticSearch domain %q to be deleted", d.Get("domain_name").(string))
	err = resource.Retry(15*time.Minute, func() error {
		out, err := conn.DescribeElasticsearchDomain(&elasticsearch.DescribeElasticsearchDomainInput{
			DomainName: aws.String(d.Get("domain_name").(string)),
		})

		if err != nil {
			awsErr, ok := err.(awserr.Error)
			if !ok {
				return resource.RetryError{Err: err}
			}

			if awsErr.Code() == "ResourceNotFoundException" {
				return nil
			}

			return resource.RetryError{Err: awsErr}
		}

		if !*out.DomainStatus.Processing {
			return nil
		}

		return fmt.Errorf("%q: Timeout while waiting for the domain to be deleted", d.Id())
	})

	d.SetId("")

	return err
}
<commit_msg>Update validation on domain name<commit_after>package aws

import (
	"fmt"
	"log"
	"regexp"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	elasticsearch "github.com/aws/aws-sdk-go/service/elasticsearchservice"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsElasticSearchDomain() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsElasticSearchDomainCreate,
		Read:   resourceAwsElasticSearchDomainRead,
		Update: resourceAwsElasticSearchDomainUpdate,
		Delete: resourceAwsElasticSearchDomainDelete,

		Schema: map[string]*schema.Schema{
			"access_policies": &schema.Schema{
				Type:      schema.TypeString,
				StateFunc: normalizeJson,
				Optional:  true,
			},
			"advanced_options": &schema.Schema{
				Type:     schema.TypeMap,
				Optional: true,
				Computed: true,
			},
			"domain_name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
				ValidateFunc: func(v interface{}, k string) (ws []string, errors []error) {
					value := v.(string)
					if !regexp.MustCompile(`^[a-z][0-9a-z\-]{2,27}$`).MatchString(value) {
						errors = append(errors, fmt.Errorf(
							"%q must start with a lowercase alphabet and be at least 3 and no more than 28 characters long. Valid characters are a-z (lowercase letters), 0-9, and - (hyphen).", k))
					}
					return
				},
			},
			"arn": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"domain_id": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"endpoint": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"ebs_options": &schema.Schema{
				Type:     schema.TypeList,
				Optional: true,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"ebs_enabled": &schema.Schema{
							Type:     schema.TypeBool,
							Required: true,
						},
						"iops": &schema.Schema{
							Type:     schema.TypeInt,
							Optional: true,
						},
						"volume_size": &schema.Schema{
							Type:     schema.TypeInt,
							Optional: true,
						},
						"volume_type": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
						},
					},
				},
			},
			"cluster_config": &schema.Schema{
				Type:     schema.TypeList,
				Optional: true,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"dedicated_master_count": &schema.Schema{
							Type:     schema.TypeInt,
							Optional: true,
						},
						"dedicated_master_enabled": &schema.Schema{
							Type:     schema.TypeBool,
							Optional: true,
							Default:  false,
						},
						"dedicated_master_type": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
						},
						"instance_count": &schema.Schema{
							Type:     schema.TypeInt,
							Optional: true,
							Default:  1,
						},
						"instance_type": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
							Default:  "m3.medium.elasticsearch",
						},
						"zone_awareness_enabled": &schema.Schema{
							Type:     schema.TypeBool,
							Optional: true,
						},
					},
				},
			},
			"snapshot_options": &schema.Schema{
				Type:     schema.TypeList,
				Optional: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"automated_snapshot_start_hour": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},
					},
				},
			},
		},
	}
}

func resourceAwsElasticSearchDomainCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).esconn

	input := elasticsearch.CreateElasticsearchDomainInput{
		DomainName: aws.String(d.Get("domain_name").(string)),
	}

	if v, ok := d.GetOk("access_policies"); ok {
		input.AccessPolicies = aws.String(v.(string))
	}

	if v, ok := d.GetOk("advanced_options"); ok {
		input.AdvancedOptions = stringMapToPointers(v.(map[string]interface{}))
	}

	if v, ok := d.GetOk("ebs_options"); ok {
		options := v.([]interface{})

		if len(options) > 1 {
			return fmt.Errorf("Only a single ebs_options block is expected")
		} else if len(options) == 1 {
			if options[0] == nil {
				return fmt.Errorf("At least one field is expected inside ebs_options")
			}

			s := options[0].(map[string]interface{})
			input.EBSOptions = expandESEBSOptions(s)
		}
	}

	if v, ok := d.GetOk("cluster_config"); ok {
		config := v.([]interface{})

		if len(config) > 1 {
			return fmt.Errorf("Only a single cluster_config block is expected")
		} else if len(config) == 1 {
			if config[0] == nil {
				return fmt.Errorf("At least one field is expected inside cluster_config")
			}
			m := config[0].(map[string]interface{})
			input.ElasticsearchClusterConfig = expandESClusterConfig(m)
		}
	}

	if v, ok := d.GetOk("snapshot_options"); ok {
		options := v.([]interface{})

		if len(options) > 1 {
			return fmt.Errorf("Only a single snapshot_options block is expected")
		} else if len(options) == 1 {
			if options[0] == nil {
				return fmt.Errorf("At least one field is expected inside snapshot_options")
			}

			o := options[0].(map[string]interface{})

			snapshotOptions := elasticsearch.SnapshotOptions{
				AutomatedSnapshotStartHour: aws.Int64(int64(o["automated_snapshot_start_hour"].(int))),
			}

			input.SnapshotOptions = &snapshotOptions
		}
	}

	log.Printf("[DEBUG] Creating ElasticSearch domain: %s", input)
	out, err := conn.CreateElasticsearchDomain(&input)
	if err != nil {
		return err
	}

	d.SetId(*out.DomainStatus.ARN)

	log.Printf("[DEBUG] Waiting for ElasticSearch domain %q to be created", d.Id())
	err = resource.Retry(15*time.Minute, func() error {
		out, err := conn.DescribeElasticsearchDomain(&elasticsearch.DescribeElasticsearchDomainInput{
			DomainName: aws.String(d.Get("domain_name").(string)),
		})
		if err != nil {
			return resource.RetryError{Err: err}
		}

		if !*out.DomainStatus.Processing && out.DomainStatus.Endpoint != nil {
			return nil
		}

		return fmt.Errorf("%q: Timeout while waiting for the domain to be created", d.Id())
	})
	if err != nil {
		return err
	}

	log.Printf("[DEBUG] ElasticSearch domain %q created", d.Id())

	return resourceAwsElasticSearchDomainRead(d, meta)
}

func resourceAwsElasticSearchDomainRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).esconn

	out, err := conn.DescribeElasticsearchDomain(&elasticsearch.DescribeElasticsearchDomainInput{
		DomainName: aws.String(d.Get("domain_name").(string)),
	})
	if err != nil {
		return err
	}

	log.Printf("[DEBUG] Received ElasticSearch domain: %s", out)

	ds := out.DomainStatus

	if ds.AccessPolicies != nil && *ds.AccessPolicies != "" {
		d.Set("access_policies", normalizeJson(*ds.AccessPolicies))
	}
	err = d.Set("advanced_options", pointersMapToStringList(ds.AdvancedOptions))
	if err != nil {
		return err
	}
	d.Set("domain_id", *ds.DomainId)
	d.Set("domain_name", *ds.DomainName)
	if ds.Endpoint != nil {
		d.Set("endpoint", *ds.Endpoint)
	}

	err = d.Set("ebs_options", flattenESEBSOptions(ds.EBSOptions))
	if err != nil {
		return err
	}
	err = d.Set("cluster_config", flattenESClusterConfig(ds.ElasticsearchClusterConfig))
	if err != nil {
		return err
	}
	if ds.SnapshotOptions != nil {
		d.Set("snapshot_options", map[string]interface{}{
			"automated_snapshot_start_hour": *ds.SnapshotOptions.AutomatedSnapshotStartHour,
		})
	}

	d.Set("arn", *ds.ARN)

	return nil
}

func resourceAwsElasticSearchDomainUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).esconn

	input := elasticsearch.UpdateElasticsearchDomainConfigInput{
		DomainName: aws.String(d.Get("domain_name").(string)),
	}

	if d.HasChange("access_policies") {
		input.AccessPolicies = aws.String(d.Get("access_policies").(string))
	}

	if d.HasChange("advanced_options") {
		input.AdvancedOptions = stringMapToPointers(d.Get("advanced_options").(map[string]interface{}))
	}

	if d.HasChange("ebs_options") {
		options := d.Get("ebs_options").([]interface{})

		if len(options) > 1 {
			return fmt.Errorf("Only a single ebs_options block is expected")
		} else if len(options) == 1 {
			s := options[0].(map[string]interface{})
			input.EBSOptions = expandESEBSOptions(s)
		}
	}

	if d.HasChange("cluster_config") {
		config := d.Get("cluster_config").([]interface{})

		if len(config) > 1 {
			return fmt.Errorf("Only a single cluster_config block is expected")
		} else if len(config) == 1 {
			m := config[0].(map[string]interface{})
			input.ElasticsearchClusterConfig = expandESClusterConfig(m)
		}
	}

	if d.HasChange("snapshot_options") {
		options := d.Get("snapshot_options").([]interface{})

		if len(options) > 1 {
			return fmt.Errorf("Only a single snapshot_options block is expected")
		} else if len(options) == 1 {
			o := options[0].(map[string]interface{})

			snapshotOptions := elasticsearch.SnapshotOptions{
				AutomatedSnapshotStartHour: aws.Int64(int64(o["automated_snapshot_start_hour"].(int))),
			}

			input.SnapshotOptions = &snapshotOptions
		}
	}

	_, err := conn.UpdateElasticsearchDomainConfig(&input)
	if err != nil {
		return err
	}

	err = resource.Retry(25*time.Minute, func() error {
		out, err := conn.DescribeElasticsearchDomain(&elasticsearch.DescribeElasticsearchDomainInput{
			DomainName: aws.String(d.Get("domain_name").(string)),
		})
		if err != nil {
			return resource.RetryError{Err: err}
		}

		if *out.DomainStatus.Processing == false {
			return nil
		}

		return fmt.Errorf("%q: Timeout while waiting for changes to be processed", d.Id())
	})
	if err != nil {
		return err
	}

	return resourceAwsElasticSearchDomainRead(d, meta)
}

func resourceAwsElasticSearchDomainDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).esconn

	log.Printf("[DEBUG] Deleting ElasticSearch domain: %q", d.Get("domain_name").(string))
	_, err := conn.DeleteElasticsearchDomain(&elasticsearch.DeleteElasticsearchDomainInput{
		DomainName: aws.String(d.Get("domain_name").(string)),
	})
	if err != nil {
		return err
	}

	log.Printf("[DEBUG] Waiting for ElasticSearch domain %q to be deleted", d.Get("domain_name").(string))
	err = resource.Retry(15*time.Minute, func() error {
		out, err := conn.DescribeElasticsearchDomain(&elasticsearch.DescribeElasticsearchDomainInput{
			DomainName: aws.String(d.Get("domain_name").(string)),
		})

		if err != nil {
			awsErr, ok := err.(awserr.Error)
			if !ok {
				return resource.RetryError{Err: err}
			}

			if awsErr.Code() == "ResourceNotFoundException" {
				return nil
			}

			return resource.RetryError{Err: awsErr}
		}

		if !*out.DomainStatus.Processing {
			return nil
		}

		return fmt.Errorf("%q: Timeout while waiting for the domain to be deleted", d.Id())
	})

	d.SetId("")

	return err
}
<|endoftext|>


<commit_before>package config

import (
	"log"
	"github.com/alehano/gobootstrap/sys/cmd"
	"github.com/spf13/cobra"
	"golang.org/x/crypto/bcrypt"
)

func init() {

	cmd.RootCmd.AddCommand(&cobra.Command{
		Use:   "version",
		Short: "Show version",
		Run: func(cmd *cobra.Command, args []string) {
			log.Printf("Version: %s\n", Version)
		},
	})

	cmd.RootCmd.AddCommand(&cobra.Command{
		Use:   "admin_pwd [password]",
		Short: "Generate Admin password hash",
		Long: "Generate Admin password hash (BCrypt) for using in config file.",
		RunE: func(cmd *cobra.Command, args []string) error {
			if len(args) == 1 {
				hash, err := bcrypt.GenerateFromPassword([]byte(args[0]), bcrypt.DefaultCost)
				if err != nil {
					return err
				} else {
					log.Printf("Password Hash: %s\n", hash)
				}
			} else {
				log.Println("Error: password string wasn't provided")
			}
			return nil
		},
	})

}
<commit_msg>Config dump<commit_after>package config

import (
	"log"
	"github.com/alehano/gobootstrap/sys/cmd"
	"github.com/spf13/cobra"
	"golang.org/x/crypto/bcrypt"
	"gopkg.in/yaml.v2"
)

func init() {

	cmd.RootCmd.AddCommand(&cobra.Command{
		Use:   "version",
		Short: "Show version",
		Run: func(cmd *cobra.Command, args []string) {
			log.Printf("Version: %s\n", Version)
		},
	})

	cmd.RootCmd.AddCommand(&cobra.Command{
		Use:   "dumpconfig",
		Short: "Print config",
		Run: func(cmd *cobra.Command, args []string) {
			out, err := yaml.Marshal(Get())
			if err != nil {
				log.Println(err)
			} else {
				log.Println(string(out))
			}
		},
	})

	cmd.RootCmd.AddCommand(&cobra.Command{
		Use:   "admin_pwd [password]",
		Short: "Generate Admin password hash",
		Long: "Generate Admin password hash (BCrypt) for using in config file.",
		RunE: func(cmd *cobra.Command, args []string) error {
			if len(args) == 1 {
				hash, err := bcrypt.GenerateFromPassword([]byte(args[0]), bcrypt.DefaultCost)
				if err != nil {
					return err
				} else {
					log.Printf("Password Hash: %s\n", hash)
				}
			} else {
				log.Println("Error: password string wasn't provided")
			}
			return nil
		},
	})

}
<|endoftext|>