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

// This file implements sysSocket and accept for platforms that
// provide a fast path for setting SetNonblock and CloseOnExec.

// +build freebsd linux

package net

import "syscall"

// Wrapper around the socket system call that marks the returned file
// descriptor as nonblocking and close-on-exec.
func sysSocket(family, sotype, proto int) (int, error) {
	s, err := syscall.Socket(family, sotype|syscall.SOCK_NONBLOCK|syscall.SOCK_CLOEXEC, proto)
	// On Linux the SOCK_NONBLOCK and SOCK_CLOEXEC flags were
	// introduced in 2.6.27 kernel and on FreeBSD both flags were
	// introduced in 10 kernel. If we get an EINVAL error on Linux
	// or EPROTONOSUPPORT error on FreeBSD, fall back to using
	// socket without them.
	if err == nil || (err != syscall.EPROTONOSUPPORT && err != syscall.EINVAL) {
		return s, err
	}

	// See ../syscall/exec_unix.go for description of ForkLock.
	syscall.ForkLock.RLock()
	s, err = syscall.Socket(family, sotype, proto)
	if err == nil {
		syscall.CloseOnExec(s)
	}
	syscall.ForkLock.RUnlock()
	if err != nil {
		return -1, err
	}
	if err = syscall.SetNonblock(s, true); err != nil {
		syscall.Close(s)
		return -1, err
	}
	return s, nil
}

// Wrapper around the accept system call that marks the returned file
// descriptor as nonblocking and close-on-exec.
func accept(s int) (int, syscall.Sockaddr, error) {
	ns, sa, err := syscall.Accept4(s, syscall.SOCK_NONBLOCK|syscall.SOCK_CLOEXEC)
	// On Linux the accept4 system call was introduced in 2.6.28
	// kernel and on FreeBSD it was introduced in 10 kernel. If we
	// get an ENOSYS error on both Linux and FreeBSD, or EINVAL
	// error on Linux, fall back to using accept.
	if err == nil || (err != syscall.ENOSYS && err != syscall.EINVAL) {
		return ns, sa, err
	}

	// See ../syscall/exec_unix.go for description of ForkLock.
	// It is probably okay to hold the lock across syscall.Accept
	// because we have put fd.sysfd into non-blocking mode.
	// However, a call to the File method will put it back into
	// blocking mode. We can't take that risk, so no use of ForkLock here.
	ns, sa, err = syscall.Accept(s)
	if err == nil {
		syscall.CloseOnExec(ns)
	}
	if err != nil {
		return -1, nil, err
	}
	if err = syscall.SetNonblock(ns, true); err != nil {
		syscall.Close(ns)
		return -1, nil, err
	}
	return ns, sa, nil
}
<commit_msg>net: accept a few more errors in Accept4 wrapper<commit_after>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file implements sysSocket and accept for platforms that
// provide a fast path for setting SetNonblock and CloseOnExec.

// +build freebsd linux

package net

import "syscall"

// Wrapper around the socket system call that marks the returned file
// descriptor as nonblocking and close-on-exec.
func sysSocket(family, sotype, proto int) (int, error) {
	s, err := syscall.Socket(family, sotype|syscall.SOCK_NONBLOCK|syscall.SOCK_CLOEXEC, proto)
	// On Linux the SOCK_NONBLOCK and SOCK_CLOEXEC flags were
	// introduced in 2.6.27 kernel and on FreeBSD both flags were
	// introduced in 10 kernel. If we get an EINVAL error on Linux
	// or EPROTONOSUPPORT error on FreeBSD, fall back to using
	// socket without them.
	if err == nil || (err != syscall.EPROTONOSUPPORT && err != syscall.EINVAL) {
		return s, err
	}

	// See ../syscall/exec_unix.go for description of ForkLock.
	syscall.ForkLock.RLock()
	s, err = syscall.Socket(family, sotype, proto)
	if err == nil {
		syscall.CloseOnExec(s)
	}
	syscall.ForkLock.RUnlock()
	if err != nil {
		return -1, err
	}
	if err = syscall.SetNonblock(s, true); err != nil {
		syscall.Close(s)
		return -1, err
	}
	return s, nil
}

// Wrapper around the accept system call that marks the returned file
// descriptor as nonblocking and close-on-exec.
func accept(s int) (int, syscall.Sockaddr, error) {
	ns, sa, err := syscall.Accept4(s, syscall.SOCK_NONBLOCK|syscall.SOCK_CLOEXEC)
	// On Linux the accept4 system call was introduced in 2.6.28
	// kernel and on FreeBSD it was introduced in 10 kernel. If we
	// get an ENOSYS error on both Linux and FreeBSD, or EINVAL
	// error on Linux, fall back to using accept.
	switch err {
	default: // nil and errors other than the ones listed
		return ns, sa, err
	case syscall.ENOSYS: // syscall missing
	case syscall.EINVAL: // some Linux use this instead of ENOSYS
	case syscall.EACCES: // some Linux use this instead of ENOSYS
	case syscall.EFAULT: // some Linux use this instead of ENOSYS
	}

	// See ../syscall/exec_unix.go for description of ForkLock.
	// It is probably okay to hold the lock across syscall.Accept
	// because we have put fd.sysfd into non-blocking mode.
	// However, a call to the File method will put it back into
	// blocking mode. We can't take that risk, so no use of ForkLock here.
	ns, sa, err = syscall.Accept(s)
	if err == nil {
		syscall.CloseOnExec(ns)
	}
	if err != nil {
		return -1, nil, err
	}
	if err = syscall.SetNonblock(ns, true); err != nil {
		syscall.Close(ns)
		return -1, nil, err
	}
	return ns, sa, nil
}
<|endoftext|>


<commit_before>package herd

import (
	"fmt"
	"github.com/Symantec/Dominator/lib/format"
	"io"
	"sort"
	"time"
)

var timeFormat string = "02 Jan 2006 15:04:05.99 MST"

type rDuration time.Duration

func (d rDuration) selector(sub *Sub) bool {
	if time.Since(sub.lastReachableTime) <= time.Duration(d) {
		return true
	}
	return false
}

func (herd *Herd) writeHtml(writer io.Writer) {
	if herd.updatesDisabledReason != "" {
		herd.writeDisableStatus(writer)
		fmt.Fprintln(writer, "
")
	}
	numSubs := herd.countSelectedSubs(nil)
	fmt.Fprintf(writer, "Time since current cycle start: %s
\n",
		time.Since(herd.currentScanStartTime))
	if numSubs < 1 {
		fmt.Fprintf(writer, "Duration of previous cycle: %s
\n",
			herd.previousScanDuration)
	} else {
		fmt.Fprintf(writer, "Duration of previous cycle: %s (%s/sub)
\n",
			herd.previousScanDuration,
			herd.previousScanDuration/time.Duration(numSubs))
	}
	fmt.Fprintf(writer, "Image server: %s
\n",
		herd.imageManager, herd.imageManager)
	if herd.defaultImageName != "" {
		fmt.Fprintf(writer,
			"Default image: %s
\n",
			herd.imageManager, herd.defaultImageName, herd.defaultImageName)
	}
	fmt.Fprintf(writer,
		"Number of subs: %d
\n",
		numSubs)
	numSubs = herd.countSelectedSubs(selectAliveSub)
	fmt.Fprintf(writer,
		"Number of alive subs: %d
\n",
		numSubs)
	fmt.Fprint(writer, "Number of reachable subs in last: ")
	herd.writeReachableSubsLink(writer, time.Minute, "1 min", "1m", true)
	herd.writeReachableSubsLink(writer, time.Minute*10, "10 min", "10m", true)
	herd.writeReachableSubsLink(writer, time.Hour, "1 hour", "1h", true)
	herd.writeReachableSubsLink(writer, time.Hour*24, "1 day", "1d", true)
	herd.writeReachableSubsLink(writer, time.Hour*24*7, "1 week", "1w", false)
	numSubs = herd.countSelectedSubs(selectDeviantSub)
	fmt.Fprintf(writer,
		"Number of deviant subs: %d
\n",
		numSubs)
	numSubs = herd.countSelectedSubs(selectCompliantSub)
	fmt.Fprintf(writer,
		"Number of compliant subs: %d
\n",
		numSubs)
	subs := herd.getSelectedSubs(nil)
	connectDurations := getConnectDurations(subs)
	shortPollDurations := getPollDurations(subs, false)
	fullPollDurations := getPollDurations(subs, true)
	showDurationStats(writer, connectDurations, "Connect")
	showDurationStats(writer, shortPollDurations, "Short poll")
	showDurationStats(writer, fullPollDurations, "Full poll")
	// TODO(rgooch): Figure out a way of restoring this information.
	//fmt.Fprintf(writer, "Connection slots: %d out of %d
\n",
	//len(herd.connectionSemaphore), cap(herd.connectionSemaphore))
	fmt.Fprintf(writer, "Poll slots: %d out of %d
\n",
		len(herd.pollSemaphore), cap(herd.pollSemaphore))
	stats := herd.cpuSharer.GetStatistics()
	timeSinceLastIdleEvent := time.Since(stats.LastIdleEvent)
	fmt.Fprintf(writer,
		"CPU slots: %d out of %d; idle events: %d, last: %s, time since: %s
\n",
		stats.NumCpuRunning, stats.NumCpu, stats.NumIdleEvents,
		stats.LastIdleEvent.Format(timeFormat),
		format.Duration(timeSinceLastIdleEvent))
}

func (herd *Herd) writeDisableStatus(writer io.Writer) {
	timeString := herd.updatesDisabledTime.Format(timeFormat)
	if herd.updatesDisabledBy == "" {
		fmt.Fprintf(writer,
			"Updates disabled because: %s at %s\n",
			herd.updatesDisabledReason, timeString)
	} else {
		fmt.Fprintf(writer,
			"Updates disabled by: %s because: %s at %s",
			herd.updatesDisabledBy, herd.updatesDisabledReason, timeString)
	}
}

func (herd *Herd) writeReachableSubsLink(writer io.Writer,
	duration time.Duration, durationString string, query string,
	moreToCome bool) {
	numSubs := herd.countSelectedSubs(rDuration(duration).selector)
	fmt.Fprintf(writer, "%s",
		query, durationString)
	fmt.Fprintf(writer, "(%d)",
		query, numSubs)
	if moreToCome {
		fmt.Fprint(writer, ", ")
	} else {
		fmt.Fprintln(writer, "
")
	}
}

func selectAliveSub(sub *Sub) bool {
	switch sub.publishedStatus {
	case statusUnknown:
		return false
	case statusConnecting:
		return false
	case statusDNSError:
		return false
	case statusConnectionRefused:
		return false
	case statusNoRouteToHost:
		return false
	case statusConnectTimeout:
		return false
	case statusFailedToConnect:
		return false
	case statusFailedToPoll:
		return false
	}
	return true
}

func selectDeviantSub(sub *Sub) bool {
	switch sub.publishedStatus {
	case statusComputingUpdate:
		return true
	case statusSendingUpdate:
		return true
	case statusUpdatesDisabled:
		return true
	case statusUpdating:
		return true
	case statusUpdateDenied:
		return true
	case statusFailedToUpdate:
		return true
	}
	return false
}

func selectCompliantSub(sub *Sub) bool {
	if sub.publishedStatus == statusSynced {
		return true
	}
	return false
}

func getConnectDurations(subs []*Sub) []time.Duration {
	durations := make([]time.Duration, 0, len(subs))
	for _, sub := range subs {
		if sub.lastConnectDuration > 0 {
			durations = append(durations, sub.lastConnectDuration)
		}
	}
	sort.Sort(durationList(durations))
	return durations
}

func getPollDurations(subs []*Sub, full bool) []time.Duration {
	durations := make([]time.Duration, 0, len(subs))
	for _, sub := range subs {
		var duration time.Duration
		if full {
			duration = sub.lastFullPollDuration
		} else {
			duration = sub.lastShortPollDuration
		}
		if duration > 0 {
			durations = append(durations, duration)
		}
	}
	sort.Sort(durationList(durations))
	return durations
}

type durationList []time.Duration

func (dl durationList) Len() int {
	return len(dl)
}

func (dl durationList) Less(i, j int) bool {
	return dl[i] < dl[j]
}

func (dl durationList) Swap(i, j int) {
	dl[i], dl[j] = dl[j], dl[i]
}

func showDurationStats(writer io.Writer, durations []time.Duration,
	durationType string) {
	if len(durations) < 1 {
		return
	}
	var avgDuration time.Duration
	for _, duration := range durations {
		avgDuration += duration
	}
	avgDuration /= time.Duration(len(durations))
	medDuration := durations[len(durations)/2]
	unit := "ns"
	scale := 1.0
	switch {
	case medDuration > 1e9:
		unit = "s"
		scale = 1e-9
	case medDuration > 1e6:
		unit = "ms"
		scale = 1e-6
	case medDuration > 1e3:
		unit = "µs"
		scale = 1e-3
	}
	fmt.Fprintf(writer,
		"%s durations: %.3f/%.3f/%.3f/%.3f %s (avg/med/min/max)
\n",
		durationType,
		float64(avgDuration)*scale, float64(medDuration)*scale,
		float64(durations[0])*scale, float64(durations[len(durations)-1])*scale,
		unit)
}
<commit_msg>Show more cpusharer statistics in dominator status page.<commit_after>package herd

import (
	"fmt"
	"github.com/Symantec/Dominator/lib/format"
	"io"
	"sort"
	"time"
)

var timeFormat string = "02 Jan 2006 15:04:05.99 MST"

type rDuration time.Duration

func (d rDuration) selector(sub *Sub) bool {
	if time.Since(sub.lastReachableTime) <= time.Duration(d) {
		return true
	}
	return false
}

func (herd *Herd) writeHtml(writer io.Writer) {
	if herd.updatesDisabledReason != "" {
		herd.writeDisableStatus(writer)
		fmt.Fprintln(writer, "
")
	}
	numSubs := herd.countSelectedSubs(nil)
	fmt.Fprintf(writer, "Time since current cycle start: %s
\n",
		time.Since(herd.currentScanStartTime))
	if numSubs < 1 {
		fmt.Fprintf(writer, "Duration of previous cycle: %s
\n",
			herd.previousScanDuration)
	} else {
		fmt.Fprintf(writer, "Duration of previous cycle: %s (%s/sub)
\n",
			herd.previousScanDuration,
			herd.previousScanDuration/time.Duration(numSubs))
	}
	fmt.Fprintf(writer, "Image server: %s
\n",
		herd.imageManager, herd.imageManager)
	if herd.defaultImageName != "" {
		fmt.Fprintf(writer,
			"Default image: %s
\n",
			herd.imageManager, herd.defaultImageName, herd.defaultImageName)
	}
	fmt.Fprintf(writer,
		"Number of subs: %d
\n",
		numSubs)
	numSubs = herd.countSelectedSubs(selectAliveSub)
	fmt.Fprintf(writer,
		"Number of alive subs: %d
\n",
		numSubs)
	fmt.Fprint(writer, "Number of reachable subs in last: ")
	herd.writeReachableSubsLink(writer, time.Minute, "1 min", "1m", true)
	herd.writeReachableSubsLink(writer, time.Minute*10, "10 min", "10m", true)
	herd.writeReachableSubsLink(writer, time.Hour, "1 hour", "1h", true)
	herd.writeReachableSubsLink(writer, time.Hour*24, "1 day", "1d", true)
	herd.writeReachableSubsLink(writer, time.Hour*24*7, "1 week", "1w", false)
	numSubs = herd.countSelectedSubs(selectDeviantSub)
	fmt.Fprintf(writer,
		"Number of deviant subs: %d
\n",
		numSubs)
	numSubs = herd.countSelectedSubs(selectCompliantSub)
	fmt.Fprintf(writer,
		"Number of compliant subs: %d
\n",
		numSubs)
	subs := herd.getSelectedSubs(nil)
	connectDurations := getConnectDurations(subs)
	shortPollDurations := getPollDurations(subs, false)
	fullPollDurations := getPollDurations(subs, true)
	showDurationStats(writer, connectDurations, "Connect")
	showDurationStats(writer, shortPollDurations, "Short poll")
	showDurationStats(writer, fullPollDurations, "Full poll")
	// TODO(rgooch): Figure out a way of restoring this information.
	//fmt.Fprintf(writer, "Connection slots: %d out of %d
\n",
	//len(herd.connectionSemaphore), cap(herd.connectionSemaphore))
	fmt.Fprintf(writer, "Poll slots: %d out of %d
\n",
		len(herd.pollSemaphore), cap(herd.pollSemaphore))
	stats := herd.cpuSharer.GetStatistics()
	timeSinceLastIdleEvent := time.Since(stats.LastIdleEvent)
	fmt.Fprintf(writer,
		"CPU slots: %d out of %d; idle events: %d, last: %s, time since: %s, last acquire: %s, last yield: %s
\n",
		stats.NumCpuRunning, stats.NumCpu, stats.NumIdleEvents,
		stats.LastIdleEvent.Format(timeFormat),
		format.Duration(timeSinceLastIdleEvent),
		format.Duration(time.Since(stats.LastAcquireEvent)),
		format.Duration(time.Since(stats.LastYieldEvent)))
}

func (herd *Herd) writeDisableStatus(writer io.Writer) {
	timeString := herd.updatesDisabledTime.Format(timeFormat)
	if herd.updatesDisabledBy == "" {
		fmt.Fprintf(writer,
			"Updates disabled because: %s at %s\n",
			herd.updatesDisabledReason, timeString)
	} else {
		fmt.Fprintf(writer,
			"Updates disabled by: %s because: %s at %s",
			herd.updatesDisabledBy, herd.updatesDisabledReason, timeString)
	}
}

func (herd *Herd) writeReachableSubsLink(writer io.Writer,
	duration time.Duration, durationString string, query string,
	moreToCome bool) {
	numSubs := herd.countSelectedSubs(rDuration(duration).selector)
	fmt.Fprintf(writer, "%s",
		query, durationString)
	fmt.Fprintf(writer, "(%d)",
		query, numSubs)
	if moreToCome {
		fmt.Fprint(writer, ", ")
	} else {
		fmt.Fprintln(writer, "
")
	}
}

func selectAliveSub(sub *Sub) bool {
	switch sub.publishedStatus {
	case statusUnknown:
		return false
	case statusConnecting:
		return false
	case statusDNSError:
		return false
	case statusConnectionRefused:
		return false
	case statusNoRouteToHost:
		return false
	case statusConnectTimeout:
		return false
	case statusFailedToConnect:
		return false
	case statusFailedToPoll:
		return false
	}
	return true
}

func selectDeviantSub(sub *Sub) bool {
	switch sub.publishedStatus {
	case statusComputingUpdate:
		return true
	case statusSendingUpdate:
		return true
	case statusUpdatesDisabled:
		return true
	case statusUpdating:
		return true
	case statusUpdateDenied:
		return true
	case statusFailedToUpdate:
		return true
	}
	return false
}

func selectCompliantSub(sub *Sub) bool {
	if sub.publishedStatus == statusSynced {
		return true
	}
	return false
}

func getConnectDurations(subs []*Sub) []time.Duration {
	durations := make([]time.Duration, 0, len(subs))
	for _, sub := range subs {
		if sub.lastConnectDuration > 0 {
			durations = append(durations, sub.lastConnectDuration)
		}
	}
	sort.Sort(durationList(durations))
	return durations
}

func getPollDurations(subs []*Sub, full bool) []time.Duration {
	durations := make([]time.Duration, 0, len(subs))
	for _, sub := range subs {
		var duration time.Duration
		if full {
			duration = sub.lastFullPollDuration
		} else {
			duration = sub.lastShortPollDuration
		}
		if duration > 0 {
			durations = append(durations, duration)
		}
	}
	sort.Sort(durationList(durations))
	return durations
}

type durationList []time.Duration

func (dl durationList) Len() int {
	return len(dl)
}

func (dl durationList) Less(i, j int) bool {
	return dl[i] < dl[j]
}

func (dl durationList) Swap(i, j int) {
	dl[i], dl[j] = dl[j], dl[i]
}

func showDurationStats(writer io.Writer, durations []time.Duration,
	durationType string) {
	if len(durations) < 1 {
		return
	}
	var avgDuration time.Duration
	for _, duration := range durations {
		avgDuration += duration
	}
	avgDuration /= time.Duration(len(durations))
	medDuration := durations[len(durations)/2]
	unit := "ns"
	scale := 1.0
	switch {
	case medDuration > 1e9:
		unit = "s"
		scale = 1e-9
	case medDuration > 1e6:
		unit = "ms"
		scale = 1e-6
	case medDuration > 1e3:
		unit = "µs"
		scale = 1e-3
	}
	fmt.Fprintf(writer,
		"%s durations: %.3f/%.3f/%.3f/%.3f %s (avg/med/min/max)
\n",
		durationType,
		float64(avgDuration)*scale, float64(medDuration)*scale,
		float64(durations[0])*scale, float64(durations[len(durations)-1])*scale,
		unit)
}
<|endoftext|>


<commit_before>package libkb

import (
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"sync"
	"time"

	"github.com/keybase/client/go/kbcrypto"
	"github.com/keybase/client/go/msgpack"
	"github.com/keybase/client/go/protocol/keybase1"
	context "golang.org/x/net/context"
)

//
// NIST = "Non-Interactive Session Token"
//
// If a client has an unlocked device key, it's able to sign a statement
// of its own creation, and present it to the server as a session key.
// Or, to save bandwidth, can just present the hash of a previously-accepted
// NIST token.
//
// Use NIST tokens rather than the prior generation of session tokens so that
// we're more responsive when we come back from backgrounding, etc.
//

// If we're within 26 hours of expiration, generate a new NIST;
const nistExpirationMargin = 26 * time.Hour // I.e., half of the lifetime
const nistLifetime = 52 * time.Hour         // A little longer than 2 days.
const nistSessionIDLength = 16
const nistShortHashLen = 19
const nistWebAuthTokenLifetime = 24 * time.Hour // website tokens expire in a day

type nistType int
type nistMode int
type sessionVersion int

const (
	nistVersion             sessionVersion = 34
	nistVersionWebAuthToken sessionVersion = 35
	nistModeSignature       nistMode       = 1
	nistModeHash            nistMode       = 2
	nistClient              nistType       = 0
	nistWebAuthToken        nistType       = 1
)

func (t nistType) sessionVersion() sessionVersion {
	if t == nistClient {
		return nistVersion
	}
	return nistVersionWebAuthToken
}

func (t nistType) lifetime() time.Duration {
	if t == nistClient {
		return nistLifetime
	}
	return nistWebAuthTokenLifetime
}

func (t nistType) signaturePrefix() kbcrypto.SignaturePrefix {
	if t == nistClient {
		return kbcrypto.SignaturePrefixNIST
	}
	return kbcrypto.SignaturePrefixNISTWebAuthToken
}

func (t nistType) encoding() *base64.Encoding {
	if t == nistClient {
		return base64.StdEncoding
	}
	return base64.RawURLEncoding
}

type NISTFactory struct {
	Contextified
	sync.Mutex
	uid      keybase1.UID
	deviceID keybase1.DeviceID
	key      GenericKey // cached secret signing key
	nist     *NIST
}

type NISTToken struct {
	b        []byte
	nistType nistType
}

func (n NISTToken) Bytes() []byte  { return n.b }
func (n NISTToken) String() string { return n.nistType.encoding().EncodeToString(n.Bytes()) }
func (n NISTToken) Hash() []byte {
	tmp := sha256.Sum256(n.Bytes())
	return tmp[:]
}
func (n NISTToken) ShortHash() []byte {
	return n.Hash()[0:nistShortHashLen]
}

type NIST struct {
	Contextified
	sync.RWMutex
	expiresAt time.Time
	failed    bool
	succeeded bool
	long      *NISTToken
	short     *NISTToken
	nistType  nistType
}

func NewNISTFactory(g *GlobalContext, uid keybase1.UID, deviceID keybase1.DeviceID, key GenericKey) *NISTFactory {
	return &NISTFactory{
		Contextified: NewContextified(g),
		uid:          uid,
		deviceID:     deviceID,
		key:          key,
	}
}

func (f *NISTFactory) UID() keybase1.UID {
	if f == nil {
		return keybase1.UID("")
	}

	f.Lock()
	defer f.Unlock()
	return f.uid
}

func (f *NISTFactory) NIST(ctx context.Context) (ret *NIST, err error) {
	if f == nil {
		return nil, nil
	}

	f.Lock()
	defer f.Unlock()

	makeNew := true

	if f.nist == nil {
		f.G().Log.CDebugf(ctx, "| NISTFactory#NIST: nil NIST, making new one")
	} else if f.nist.DidFail() {
		f.G().Log.CDebugf(ctx, "| NISTFactory#NIST: NIST previously failed, so we'll make a new one")
	} else {
		valid, until := f.nist.IsStillValid()
		if valid {
			f.G().Log.CDebugf(ctx, "| NISTFactory#NIST: returning existing NIST (expires conservatively in %s, expiresAt: %s)", until, f.nist.expiresAt)
			makeNew = false
		} else {
			f.G().Log.CDebugf(ctx, "| NISTFactory#NIST: NIST expired (conservatively) %s ago, making a new one (expiresAt: %s)", -until, f.nist.expiresAt)
		}
	}

	if makeNew {
		ret = newNIST(f.G())
		err = ret.generate(ctx, f.uid, f.deviceID, f.key, nistClient)
		if err != nil {
			return nil, err
		}
		f.nist = ret
		f.G().Log.CDebugf(ctx, "| NISTFactory#NIST: Installing new NIST; expiresAt: %s", f.nist.expiresAt)
	}

	return f.nist, nil
}

func (f *NISTFactory) GenerateWebAuthToken(ctx context.Context) (ret *NIST, err error) {
	ret = newNIST(f.G())
	err = ret.generate(ctx, f.uid, f.deviceID, f.key, nistWebAuthToken)
	return ret, err
}

func durationToSeconds(d time.Duration) int64 {
	return int64(d / time.Second)
}

func (n *NIST) IsStillValid() (bool, time.Duration) {
	n.RLock()
	defer n.RUnlock()
	now := ForceWallClock(n.G().Clock().Now())
	diff := n.expiresAt.Sub(now) - nistExpirationMargin
	return (diff > 0), diff
}

func (n *NIST) IsExpired() bool {
	isValid, _ := n.IsStillValid()
	return !isValid
}

func (n *NIST) DidFail() bool {
	n.RLock()
	defer n.RUnlock()
	return n.failed
}

func (n *NIST) MarkFailure() {
	n.Lock()
	defer n.Unlock()
	n.failed = true
}

func (n *NIST) MarkSuccess() {
	n.Lock()
	defer n.Unlock()
	n.succeeded = true
}

func newNIST(g *GlobalContext) *NIST {
	return &NIST{Contextified: NewContextified(g)}
}

type nistPayload struct {
	_struct   bool `codec:",toarray"` //nolint
	Version   sessionVersion
	Mode      nistMode
	Hostname  string
	UID       []byte
	DeviceID  []byte
	KID       []byte
	Generated int64
	Lifetime  int64
	SessionID []byte
}

type nistSig struct {
	_struct bool `codec:",toarray"` //nolint
	Version sessionVersion
	Mode    nistMode
	Sig     []byte
	Payload nistPayloadShort
}

type nistPayloadShort struct {
	_struct   bool `codec:",toarray"` //nolint
	UID       []byte
	DeviceID  []byte
	Generated int64
	Lifetime  int64
	SessionID []byte
}

type nistHash struct {
	_struct bool `codec:",toarray"` //nolint
	Version sessionVersion
	Mode    nistMode
	Hash    []byte
}

func (h nistSig) pack(t nistType) (*NISTToken, error) {
	b, err := msgpack.Encode(h)
	if err != nil {
		return nil, err
	}
	return &NISTToken{b: b, nistType: t}, nil
}

func (h nistHash) pack(t nistType) (*NISTToken, error) {
	b, err := msgpack.Encode(h)
	if err != nil {
		return nil, err
	}
	return &NISTToken{b: b, nistType: t}, nil
}

func (n nistPayload) abbreviate() nistPayloadShort {
	return nistPayloadShort{
		UID:       n.UID,
		DeviceID:  n.DeviceID,
		Generated: n.Generated,
		Lifetime:  n.Lifetime,
		SessionID: n.SessionID,
	}
}

func (n *NIST) generate(ctx context.Context, uid keybase1.UID, deviceID keybase1.DeviceID, key GenericKey, typ nistType) (err error) {
	defer n.G().CTrace(ctx, "NIST#generate", func() error { return err })()

	n.Lock()
	defer n.Unlock()

	naclKey, ok := (key).(NaclSigningKeyPair)
	if !ok {
		return errors.New("cannot generate a NIST without a NaCl key")
	}

	var generated time.Time

	// For some tests we ignore the clock in n.G().Clock() and just use the standard
	// time.Now() clock, because otherwise, the server would start to reject our
	// NISTs.
	if n.G().Env.UseTimeClockForNISTs() {
		generated = time.Now()
	} else {
		generated = n.G().Clock().Now()
	}

	lifetime := typ.lifetime()
	expires := generated.Add(lifetime)
	version := typ.sessionVersion()

	payload := nistPayload{
		Version:   version,
		Mode:      nistModeSignature,
		Hostname:  CanonicalHost,
		UID:       uid.ToBytes(),
		Generated: generated.Unix(),
		Lifetime:  durationToSeconds(lifetime),
		KID:       key.GetBinaryKID(),
	}
	n.G().Log.CDebugf(ctx, "NIST: uid=%s; kid=%s; deviceID=%s", uid, key.GetKID().String(), deviceID)
	payload.DeviceID, err = hex.DecodeString(string(deviceID))
	if err != nil {
		return err
	}
	payload.SessionID, err = RandBytes(nistSessionIDLength)
	if err != nil {
		return err
	}
	var sigInfo kbcrypto.NaclSigInfo
	var payloadPacked []byte
	payloadPacked, err = msgpack.Encode(payload)
	if err != nil {
		return err
	}
	sigInfo, err = naclKey.SignV2(payloadPacked, typ.signaturePrefix())
	if err != nil {
		return err
	}

	var longTmp, shortTmp *NISTToken

	longTmp, err = (nistSig{
		Version: version,
		Mode:    nistModeSignature,
		Sig:     sigInfo.Sig[:],
		Payload: payload.abbreviate(),
	}).pack(typ)
	if err != nil {
		return err
	}

	shortTmp, err = (nistHash{
		Version: version,
		Mode:    nistModeHash,
		Hash:    longTmp.ShortHash(),
	}).pack(typ)
	if err != nil {
		return err
	}

	n.long = longTmp
	n.short = shortTmp
	n.expiresAt = ForceWallClock(expires)
	n.nistType = typ

	return nil
}

func (n *NIST) Token() *NISTToken {
	n.RLock()
	defer n.RUnlock()
	if n.succeeded {
		return n.short
	}
	return n.long
}
<commit_msg>Pass old valid nist for new one (#21754)<commit_after>package libkb

import (
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"sync"
	"time"

	"github.com/keybase/client/go/kbcrypto"
	"github.com/keybase/client/go/msgpack"
	"github.com/keybase/client/go/protocol/keybase1"
	context "golang.org/x/net/context"
)

//
// NIST = "Non-Interactive Session Token"
//
// If a client has an unlocked device key, it's able to sign a statement
// of its own creation, and present it to the server as a session key.
// Or, to save bandwidth, can just present the hash of a previously-accepted
// NIST token.
//
// Use NIST tokens rather than the prior generation of session tokens so that
// we're more responsive when we come back from backgrounding, etc.
//

// If we're within 26 hours of expiration, generate a new NIST;
const nistExpirationMargin = 26 * time.Hour // I.e., half of the lifetime
const nistLifetime = 52 * time.Hour         // A little longer than 2 days.
const nistSessionIDLength = 16
const nistShortHashLen = 19
const nistWebAuthTokenLifetime = 24 * time.Hour // website tokens expire in a day

type nistType int
type nistMode int
type sessionVersion int

const (
	nistVersion             sessionVersion = 34
	nistVersionWebAuthToken sessionVersion = 35
	nistModeSignature       nistMode       = 1
	nistModeHash            nistMode       = 2
	nistClient              nistType       = 0
	nistWebAuthToken        nistType       = 1
)

func (t nistType) sessionVersion() sessionVersion {
	if t == nistClient {
		return nistVersion
	}
	return nistVersionWebAuthToken
}

func (t nistType) lifetime() time.Duration {
	if t == nistClient {
		return nistLifetime
	}
	return nistWebAuthTokenLifetime
}

func (t nistType) signaturePrefix() kbcrypto.SignaturePrefix {
	if t == nistClient {
		return kbcrypto.SignaturePrefixNIST
	}
	return kbcrypto.SignaturePrefixNISTWebAuthToken
}

func (t nistType) encoding() *base64.Encoding {
	if t == nistClient {
		return base64.StdEncoding
	}
	return base64.RawURLEncoding
}

type NISTFactory struct {
	Contextified
	sync.Mutex
	uid                keybase1.UID
	deviceID           keybase1.DeviceID
	key                GenericKey // cached secret signing key
	nist               *NIST
	lastSuccessfulNIST *NIST
}

type NISTToken struct {
	b        []byte
	nistType nistType
}

func (n NISTToken) Bytes() []byte  { return n.b }
func (n NISTToken) String() string { return n.nistType.encoding().EncodeToString(n.Bytes()) }
func (n NISTToken) Hash() []byte {
	tmp := sha256.Sum256(n.Bytes())
	return tmp[:]
}
func (n NISTToken) ShortHash() []byte {
	return n.Hash()[0:nistShortHashLen]
}

type NIST struct {
	Contextified
	sync.RWMutex
	expiresAt time.Time
	failed    bool
	succeeded bool
	long      *NISTToken
	short     *NISTToken
	nistType  nistType
}

func NewNISTFactory(g *GlobalContext, uid keybase1.UID, deviceID keybase1.DeviceID, key GenericKey) *NISTFactory {
	return &NISTFactory{
		Contextified: NewContextified(g),
		uid:          uid,
		deviceID:     deviceID,
		key:          key,
	}
}

func (f *NISTFactory) UID() keybase1.UID {
	if f == nil {
		return keybase1.UID("")
	}

	f.Lock()
	defer f.Unlock()
	return f.uid
}

func (f *NISTFactory) NIST(ctx context.Context) (ret *NIST, err error) {
	if f == nil {
		return nil, nil
	}

	f.Lock()
	defer f.Unlock()

	makeNew := true

	if f.nist == nil {
		f.G().Log.CDebugf(ctx, "| NISTFactory#NIST: nil NIST, making new one")
	} else if f.nist.DidFail() {
		f.G().Log.CDebugf(ctx, "| NISTFactory#NIST: NIST previously failed, so we'll make a new one")
	} else {
		if f.nist.DidSucceed() {
			f.lastSuccessfulNIST = f.nist
		}

		valid, until := f.nist.IsStillValid()
		if valid {
			f.G().Log.CDebugf(ctx, "| NISTFactory#NIST: returning existing NIST (expires conservatively in %s, expiresAt: %s)", until, f.nist.expiresAt)
			makeNew = false
		} else {
			f.G().Log.CDebugf(ctx, "| NISTFactory#NIST: NIST expired (conservatively) %s ago, making a new one (expiresAt: %s)", -until, f.nist.expiresAt)
		}
	}

	if makeNew {
		ret = newNIST(f.G())
		var lastSuccessfulNISTShortHash []byte
		if f.lastSuccessfulNIST != nil {
			lastSuccessfulNISTShortHash = f.lastSuccessfulNIST.long.ShortHash()
		}
		err = ret.generate(ctx, f.uid, f.deviceID, f.key, nistClient, lastSuccessfulNISTShortHash)
		if err != nil {
			return nil, err
		}
		f.nist = ret
		f.G().Log.CDebugf(ctx, "| NISTFactory#NIST: Installing new NIST; expiresAt: %s", f.nist.expiresAt)
	}

	return f.nist, nil
}

func (f *NISTFactory) GenerateWebAuthToken(ctx context.Context) (ret *NIST, err error) {
	ret = newNIST(f.G())
	err = ret.generate(ctx, f.uid, f.deviceID, f.key, nistWebAuthToken, nil)
	return ret, err
}

func durationToSeconds(d time.Duration) int64 {
	return int64(d / time.Second)
}

func (n *NIST) IsStillValid() (bool, time.Duration) {
	n.RLock()
	defer n.RUnlock()
	now := ForceWallClock(n.G().Clock().Now())
	diff := n.expiresAt.Sub(now) - nistExpirationMargin
	return (diff > 0), diff
}

func (n *NIST) IsExpired() bool {
	isValid, _ := n.IsStillValid()
	return !isValid
}

func (n *NIST) DidSucceed() bool {
	n.RLock()
	defer n.RUnlock()
	return n.succeeded
}

func (n *NIST) DidFail() bool {
	n.RLock()
	defer n.RUnlock()
	return n.failed
}

func (n *NIST) MarkFailure() {
	n.Lock()
	defer n.Unlock()
	n.failed = true
}

func (n *NIST) MarkSuccess() {
	n.Lock()
	defer n.Unlock()
	n.succeeded = true
}

func newNIST(g *GlobalContext) *NIST {
	return &NIST{Contextified: NewContextified(g)}
}

type nistPayload struct {
	_struct   bool `codec:",toarray"` //nolint
	Version   sessionVersion
	Mode      nistMode
	Hostname  string
	UID       []byte
	DeviceID  []byte
	KID       []byte
	Generated int64
	Lifetime  int64
	SessionID []byte
}

type nistSig struct {
	_struct bool `codec:",toarray"` //nolint
	Version sessionVersion
	Mode    nistMode
	Sig     []byte
	Payload nistPayloadShort
}

type nistPayloadShort struct {
	_struct     bool `codec:",toarray"` //nolint
	UID         []byte
	DeviceID    []byte
	Generated   int64
	Lifetime    int64
	SessionID   []byte
	OldNISTHash []byte
}

type nistHash struct {
	_struct bool `codec:",toarray"` //nolint
	Version sessionVersion
	Mode    nistMode
	Hash    []byte
}

func (h nistSig) pack(t nistType) (*NISTToken, error) {
	b, err := msgpack.Encode(h)
	if err != nil {
		return nil, err
	}
	return &NISTToken{b: b, nistType: t}, nil
}

func (h nistHash) pack(t nistType) (*NISTToken, error) {
	b, err := msgpack.Encode(h)
	if err != nil {
		return nil, err
	}
	return &NISTToken{b: b, nistType: t}, nil
}

func (n nistPayload) abbreviate(lastSuccessfulNISTShortHash []byte) nistPayloadShort {
	short := nistPayloadShort{
		UID:         n.UID,
		DeviceID:    n.DeviceID,
		Generated:   n.Generated,
		Lifetime:    n.Lifetime,
		SessionID:   n.SessionID,
		OldNISTHash: lastSuccessfulNISTShortHash,
	}
	return short
}

func (n *NIST) generate(ctx context.Context, uid keybase1.UID, deviceID keybase1.DeviceID, key GenericKey, typ nistType, lastSuccessfulShortHash []byte) (err error) {
	defer n.G().CTrace(ctx, "NIST#generate", func() error { return err })()

	n.Lock()
	defer n.Unlock()

	naclKey, ok := (key).(NaclSigningKeyPair)
	if !ok {
		return errors.New("cannot generate a NIST without a NaCl key")
	}

	var generated time.Time

	// For some tests we ignore the clock in n.G().Clock() and just use the standard
	// time.Now() clock, because otherwise, the server would start to reject our
	// NISTs.
	if n.G().Env.UseTimeClockForNISTs() {
		generated = time.Now()
	} else {
		generated = n.G().Clock().Now()
	}

	lifetime := typ.lifetime()
	expires := generated.Add(lifetime)
	version := typ.sessionVersion()

	payload := nistPayload{
		Version:   version,
		Mode:      nistModeSignature,
		Hostname:  CanonicalHost,
		UID:       uid.ToBytes(),
		Generated: generated.Unix(),
		Lifetime:  durationToSeconds(lifetime),
		KID:       key.GetBinaryKID(),
	}
	n.G().Log.CDebugf(ctx, "NIST: uid=%s; kid=%s; deviceID=%s", uid, key.GetKID().String(), deviceID)
	payload.DeviceID, err = hex.DecodeString(string(deviceID))
	if err != nil {
		return err
	}
	payload.SessionID, err = RandBytes(nistSessionIDLength)
	if err != nil {
		return err
	}
	var sigInfo kbcrypto.NaclSigInfo
	var payloadPacked []byte
	payloadPacked, err = msgpack.Encode(payload)
	if err != nil {
		return err
	}
	sigInfo, err = naclKey.SignV2(payloadPacked, typ.signaturePrefix())
	if err != nil {
		return err
	}

	var longTmp, shortTmp *NISTToken

	long := nistSig{
		Version: version,
		Mode:    nistModeSignature,
		Sig:     sigInfo.Sig[:],
		Payload: payload.abbreviate(lastSuccessfulShortHash),
	}

	longTmp, err = (long).pack(typ)
	if err != nil {
		return err
	}

	shortTmp, err = (nistHash{
		Version: version,
		Mode:    nistModeHash,
		Hash:    longTmp.ShortHash(),
	}).pack(typ)
	if err != nil {
		return err
	}

	n.long = longTmp
	n.short = shortTmp
	n.expiresAt = ForceWallClock(expires)
	n.nistType = typ

	return nil
}

func (n *NIST) Token() *NISTToken {
	n.RLock()
	defer n.RUnlock()
	if n.succeeded {
		return n.short
	}
	return n.long
}
<|endoftext|>


<commit_before>package prompt

import (
	"errors"
	"strings"
	"syscall"
	"unsafe"
)

const (
	CREDUI_FLAGS_ALWAYS_SHOW_UI      = 0x00080
	CREDUI_FLAGS_GENERIC_CREDENTIALS = 0x40000
	CREDUI_FLAGS_KEEP_USERNAME       = 0x100000
)

type creduiInfoA struct {
	cbSize         uint32
	hwndParent     uintptr
	pszMessageText *uint16
	pszCaptionText *uint16
	hbmBanner      uintptr
}

func WindowsMfaPrompt(mfaSerial string) (string, error) {
	info := &creduiInfoA{
		hwndParent:     0,
		pszCaptionText: syscall.StringToUTF16Ptr("Enter token for aws-vault"),
		pszMessageText: syscall.StringToUTF16Ptr(mfaPromptMessage(mfaSerial)),
		hbmBanner:      0,
	}
	info.cbSize = uint32(unsafe.Sizeof(*info))
	passwordBuf := make([]uint16, 64)
	save := false
	flags := CREDUI_FLAGS_ALWAYS_SHOW_UI | CREDUI_FLAGS_KEEP_USERNAME | CREDUI_FLAGS_GENERIC_CREDENTIALS
	shortSerial := strings.ReplaceAll(strings.ReplaceAll(mfaSerial, "arn:aws:iam::", ""), ":mfa", "")

	ret, _, _ := syscall.NewLazyDLL("credui.dll").NewProc("CredUIPromptForCredentialsW").Call(
		uintptr(unsafe.Pointer(info)),
		uintptr(unsafe.Pointer(syscall.StringBytePtr("aws-vault"))),
		0,
		0,
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(shortSerial))),
		uintptr(len(shortSerial)+1),
		uintptr(unsafe.Pointer(&passwordBuf[0])),
		64,
		uintptr(unsafe.Pointer(&save)),
		uintptr(flags),
	)
	if ret != 0 {
		return "", errors.New("credui: call to CredUIPromptForCredentialsW failed")
	}

	return strings.TrimSpace(syscall.UTF16ToString(passwordBuf)), nil
}

func init() {
	Methods["credui"] = WindowsMfaPrompt
}
<commit_msg>Rename function<commit_after>package prompt

import (
	"errors"
	"strings"
	"syscall"
	"unsafe"
)

const (
	CREDUI_FLAGS_ALWAYS_SHOW_UI      = 0x00080
	CREDUI_FLAGS_GENERIC_CREDENTIALS = 0x40000
	CREDUI_FLAGS_KEEP_USERNAME       = 0x100000
)

type creduiInfoA struct {
	cbSize         uint32
	hwndParent     uintptr
	pszMessageText *uint16
	pszCaptionText *uint16
	hbmBanner      uintptr
}

func CredUiPrompt(mfaSerial string) (string, error) {
	info := &creduiInfoA{
		hwndParent:     0,
		pszCaptionText: syscall.StringToUTF16Ptr("Enter token for aws-vault"),
		pszMessageText: syscall.StringToUTF16Ptr(mfaPromptMessage(mfaSerial)),
		hbmBanner:      0,
	}
	info.cbSize = uint32(unsafe.Sizeof(*info))
	passwordBuf := make([]uint16, 64)
	save := false
	flags := CREDUI_FLAGS_ALWAYS_SHOW_UI | CREDUI_FLAGS_KEEP_USERNAME | CREDUI_FLAGS_GENERIC_CREDENTIALS
	shortSerial := strings.ReplaceAll(strings.ReplaceAll(mfaSerial, "arn:aws:iam::", ""), ":mfa", "")

	ret, _, _ := syscall.NewLazyDLL("credui.dll").NewProc("CredUIPromptForCredentialsW").Call(
		uintptr(unsafe.Pointer(info)),
		uintptr(unsafe.Pointer(syscall.StringBytePtr("aws-vault"))),
		0,
		0,
		uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(shortSerial))),
		uintptr(len(shortSerial)+1),
		uintptr(unsafe.Pointer(&passwordBuf[0])),
		64,
		uintptr(unsafe.Pointer(&save)),
		uintptr(flags),
	)
	if ret != 0 {
		return "", errors.New("credui: call to CredUIPromptForCredentialsW failed")
	}

	return strings.TrimSpace(syscall.UTF16ToString(passwordBuf)), nil
}

func init() {
	Methods["credui"] = CredUiPrompt
}
<|endoftext|>


<commit_before>// Mgmt
// Copyright (C) 2013-2016+ James Shubin and the project contributors
// Written by James Shubin <james@shubin.ca> and the project contributors
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package main

import (
	"bytes"
	"encoding/base64"
	"encoding/gob"
	"fmt"
	"log"
)

//go:generate stringer -type=resState -output=resstate_stringer.go
type resState int

const (
	resStateNil resState = iota
	resStateWatching
	resStateEvent // an event has happened, but we haven't poked yet
	resStateCheckApply
	resStatePoking
)

// a unique identifier for a resource, namely it's name, and the kind ("type")
type ResUUID interface {
	GetName() string
	Kind() string
	IFF(ResUUID) bool

	Reversed() bool // true means this resource happens before the generator
}

type BaseUUID struct {
	name string // name and kind are the values of where this is coming from
	kind string

	reversed *bool // piggyback edge information here
}

type AutoEdge interface {
	Next() []ResUUID  // call to get list of edges to add
	Test([]bool) bool // call until false
}

type MetaParams struct {
	AutoEdge  bool `yaml:"autoedge"`  // metaparam, should we generate auto edges? // XXX should default to true
	AutoGroup bool `yaml:"autogroup"` // metaparam, should we auto group? // XXX should default to true
}

// this interface is everything that is common to all resources
// everything here only needs to be implemented once, in the BaseRes
type Base interface {
	GetName() string // can't be named "Name()" because of struct field
	SetName(string)
	Kind() string
	GetMeta() MetaParams
	AssociateData(Converger)
	IsWatching() bool
	SetWatching(bool)
	GetState() resState
	SetState(resState)
	SendEvent(eventName, bool, bool) bool
	ReadEvent(*Event) (bool, bool) // TODO: optional here?
	GroupCmp(Res) bool             // TODO: is there a better name for this?
	GroupRes(Res) error            // group resource (arg) into self
	IsGrouped() bool               // am I grouped?
	SetGrouped(bool)               // set grouped bool
	GetGroup() []Res               // return everyone grouped inside me
	SetGroup([]Res)
}

// this is the minimum interface you need to implement to make a new resource
type Res interface {
	Base // include everything from the Base interface
	Init()
	//Validate() bool    // TODO: this might one day be added
	GetUUIDs() []ResUUID // most resources only return one
	Watch(chan Event)    // send on channel to signal process() events
	CheckApply(bool) (bool, error)
	AutoEdges() AutoEdge
	Compare(Res) bool
	CollectPattern(string) // XXX: temporary until Res collection is more advanced
}

type BaseRes struct {
	Name      string     `yaml:"name"`
	Meta      MetaParams `yaml:"meta"` // struct of all the metaparams
	kind      string
	events    chan Event
	converger Converger // converged tracking
	state     resState
	watching  bool  // is Watch() loop running ?
	isStateOK bool  // whether the state is okay based on events or not
	isGrouped bool  // am i contained within a group?
	grouped   []Res // list of any grouped resources
}

// wraps the IFF method when used with a list of UUID's
func UUIDExistsInUUIDs(uuid ResUUID, uuids []ResUUID) bool {
	for _, u := range uuids {
		if uuid.IFF(u) {
			return true
		}
	}
	return false
}

func (obj *BaseUUID) GetName() string {
	return obj.name
}

func (obj *BaseUUID) Kind() string {
	return obj.kind
}

// if and only if they are equivalent, return true
// if they are not equivalent, return false
// most resource will want to override this method, since it does the important
// work of actually discerning if two resources are identical in function
func (obj *BaseUUID) IFF(uuid ResUUID) bool {
	res, ok := uuid.(*BaseUUID)
	if !ok {
		return false
	}
	return obj.name == res.name
}

func (obj *BaseUUID) Reversed() bool {
	if obj.reversed == nil {
		log.Fatal("Programming error!")
	}
	return *obj.reversed
}

// initialize structures like channels if created without New constructor
func (obj *BaseRes) Init() {
	obj.events = make(chan Event) // unbuffered chan size to avoid stale events
}

// this method gets used by all the resources
func (obj *BaseRes) GetName() string {
	return obj.Name
}

func (obj *BaseRes) SetName(name string) {
	obj.Name = name
}

// return the kind of resource this is
func (obj *BaseRes) Kind() string {
	return obj.kind
}

func (obj *BaseRes) GetMeta() MetaParams {
	return obj.Meta
}

// AssociateData associates some data with the object in question
func (obj *BaseRes) AssociateData(converger Converger) {
	obj.converger = converger
}

// is the Watch() function running?
func (obj *BaseRes) IsWatching() bool {
	return obj.watching
}

// store status of if the Watch() function is running
func (obj *BaseRes) SetWatching(b bool) {
	obj.watching = b
}

func (obj *BaseRes) GetState() resState {
	return obj.state
}

func (obj *BaseRes) SetState(state resState) {
	if DEBUG {
		log.Printf("%v[%v]: State: %v -> %v", obj.Kind(), obj.GetName(), obj.GetState(), state)
	}
	obj.state = state
}

// push an event into the message queue for a particular vertex
func (obj *BaseRes) SendEvent(event eventName, sync bool, activity bool) bool {
	// TODO: isn't this race-y ?
	if !obj.IsWatching() { // element has already exited
		return false // if we don't return, we'll block on the send
	}
	if !sync {
		obj.events <- Event{event, nil, "", activity}
		return true
	}

	resp := make(chan bool)
	obj.events <- Event{event, resp, "", activity}
	for {
		value := <-resp
		// wait until true value
		if value {
			return true
		}
	}
}

// process events when a select gets one, this handles the pause code too!
// the return values specify if we should exit and poke respectively
func (obj *BaseRes) ReadEvent(event *Event) (exit, poke bool) {
	event.ACK()
	switch event.Name {
	case eventStart:
		return false, true

	case eventPoke:
		return false, true

	case eventBackPoke:
		return false, true // forward poking in response to a back poke!

	case eventExit:
		return true, false

	case eventPause:
		// wait for next event to continue
		select {
		case e := <-obj.events:
			e.ACK()
			if e.Name == eventExit {
				return true, false
			} else if e.Name == eventStart { // eventContinue
				return false, false // don't poke on unpause!
			} else {
				// if we get a poke event here, it's a bug!
				log.Fatalf("%v[%v]: Unknown event: %v, while paused!", obj.Kind(), obj.GetName(), e)
			}
		}

	default:
		log.Fatal("Unknown event: ", event)
	}
	return true, false // required to keep the stupid go compiler happy
}

func (obj *BaseRes) GroupRes(res Res) error {
	if l := len(res.GetGroup()); l > 0 {
		return fmt.Errorf("Res: %v already contains %d grouped resources!", res, l)
	}
	if res.IsGrouped() {
		return fmt.Errorf("Res: %v is already grouped!", res)
	}

	obj.grouped = append(obj.grouped, res)
	res.SetGrouped(true) // i am contained _in_ a group
	return nil
}

func (obj *BaseRes) IsGrouped() bool { // am I grouped?
	return obj.isGrouped
}

func (obj *BaseRes) SetGrouped(b bool) {
	obj.isGrouped = b
}

func (obj *BaseRes) GetGroup() []Res { // return everyone grouped inside me
	return obj.grouped
}

func (obj *BaseRes) SetGroup(g []Res) {
	obj.grouped = g
}

func (obj *BaseRes) CollectPattern(pattern string) {
	// XXX: default method is empty
}

// ResToB64 encodes a resource to a base64 encoded string (after serialization)
func ResToB64(res Res) (string, error) {
	b := bytes.Buffer{}
	e := gob.NewEncoder(&b)
	err := e.Encode(&res) // pass with &
	if err != nil {
		return "", fmt.Errorf("Gob failed to encode: %v", err)
	}
	return base64.StdEncoding.EncodeToString(b.Bytes()), nil
}

// B64ToRes decodes a resource from a base64 encoded string (after deserialization)
func B64ToRes(str string) (Res, error) {
	var output interface{}
	bb, err := base64.StdEncoding.DecodeString(str)
	if err != nil {
		return nil, fmt.Errorf("Base64 failed to decode: %v", err)
	}
	b := bytes.NewBuffer(bb)
	d := gob.NewDecoder(b)
	err = d.Decode(&output) // pass with &
	if err != nil {
		return nil, fmt.Errorf("Gob failed to decode: %v", err)
	}
	res, ok := output.(Res)
	if !ok {
		return nil, fmt.Errorf("Output %v is not a Res", res)

	}
	return res, nil
}
<commit_msg>Add a base method for GroupCmp<commit_after>// Mgmt
// Copyright (C) 2013-2016+ James Shubin and the project contributors
// Written by James Shubin <james@shubin.ca> and the project contributors
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package main

import (
	"bytes"
	"encoding/base64"
	"encoding/gob"
	"fmt"
	"log"
)

//go:generate stringer -type=resState -output=resstate_stringer.go
type resState int

const (
	resStateNil resState = iota
	resStateWatching
	resStateEvent // an event has happened, but we haven't poked yet
	resStateCheckApply
	resStatePoking
)

// a unique identifier for a resource, namely it's name, and the kind ("type")
type ResUUID interface {
	GetName() string
	Kind() string
	IFF(ResUUID) bool

	Reversed() bool // true means this resource happens before the generator
}

type BaseUUID struct {
	name string // name and kind are the values of where this is coming from
	kind string

	reversed *bool // piggyback edge information here
}

type AutoEdge interface {
	Next() []ResUUID  // call to get list of edges to add
	Test([]bool) bool // call until false
}

type MetaParams struct {
	AutoEdge  bool `yaml:"autoedge"`  // metaparam, should we generate auto edges? // XXX should default to true
	AutoGroup bool `yaml:"autogroup"` // metaparam, should we auto group? // XXX should default to true
}

// this interface is everything that is common to all resources
// everything here only needs to be implemented once, in the BaseRes
type Base interface {
	GetName() string // can't be named "Name()" because of struct field
	SetName(string)
	Kind() string
	GetMeta() MetaParams
	AssociateData(Converger)
	IsWatching() bool
	SetWatching(bool)
	GetState() resState
	SetState(resState)
	SendEvent(eventName, bool, bool) bool
	ReadEvent(*Event) (bool, bool) // TODO: optional here?
	GroupCmp(Res) bool             // TODO: is there a better name for this?
	GroupRes(Res) error            // group resource (arg) into self
	IsGrouped() bool               // am I grouped?
	SetGrouped(bool)               // set grouped bool
	GetGroup() []Res               // return everyone grouped inside me
	SetGroup([]Res)
}

// this is the minimum interface you need to implement to make a new resource
type Res interface {
	Base // include everything from the Base interface
	Init()
	//Validate() bool    // TODO: this might one day be added
	GetUUIDs() []ResUUID // most resources only return one
	Watch(chan Event)    // send on channel to signal process() events
	CheckApply(bool) (bool, error)
	AutoEdges() AutoEdge
	Compare(Res) bool
	CollectPattern(string) // XXX: temporary until Res collection is more advanced
}

type BaseRes struct {
	Name      string     `yaml:"name"`
	Meta      MetaParams `yaml:"meta"` // struct of all the metaparams
	kind      string
	events    chan Event
	converger Converger // converged tracking
	state     resState
	watching  bool  // is Watch() loop running ?
	isStateOK bool  // whether the state is okay based on events or not
	isGrouped bool  // am i contained within a group?
	grouped   []Res // list of any grouped resources
}

// wraps the IFF method when used with a list of UUID's
func UUIDExistsInUUIDs(uuid ResUUID, uuids []ResUUID) bool {
	for _, u := range uuids {
		if uuid.IFF(u) {
			return true
		}
	}
	return false
}

func (obj *BaseUUID) GetName() string {
	return obj.name
}

func (obj *BaseUUID) Kind() string {
	return obj.kind
}

// if and only if they are equivalent, return true
// if they are not equivalent, return false
// most resource will want to override this method, since it does the important
// work of actually discerning if two resources are identical in function
func (obj *BaseUUID) IFF(uuid ResUUID) bool {
	res, ok := uuid.(*BaseUUID)
	if !ok {
		return false
	}
	return obj.name == res.name
}

func (obj *BaseUUID) Reversed() bool {
	if obj.reversed == nil {
		log.Fatal("Programming error!")
	}
	return *obj.reversed
}

// initialize structures like channels if created without New constructor
func (obj *BaseRes) Init() {
	obj.events = make(chan Event) // unbuffered chan size to avoid stale events
}

// this method gets used by all the resources
func (obj *BaseRes) GetName() string {
	return obj.Name
}

func (obj *BaseRes) SetName(name string) {
	obj.Name = name
}

// return the kind of resource this is
func (obj *BaseRes) Kind() string {
	return obj.kind
}

func (obj *BaseRes) GetMeta() MetaParams {
	return obj.Meta
}

// AssociateData associates some data with the object in question
func (obj *BaseRes) AssociateData(converger Converger) {
	obj.converger = converger
}

// is the Watch() function running?
func (obj *BaseRes) IsWatching() bool {
	return obj.watching
}

// store status of if the Watch() function is running
func (obj *BaseRes) SetWatching(b bool) {
	obj.watching = b
}

func (obj *BaseRes) GetState() resState {
	return obj.state
}

func (obj *BaseRes) SetState(state resState) {
	if DEBUG {
		log.Printf("%v[%v]: State: %v -> %v", obj.Kind(), obj.GetName(), obj.GetState(), state)
	}
	obj.state = state
}

// push an event into the message queue for a particular vertex
func (obj *BaseRes) SendEvent(event eventName, sync bool, activity bool) bool {
	// TODO: isn't this race-y ?
	if !obj.IsWatching() { // element has already exited
		return false // if we don't return, we'll block on the send
	}
	if !sync {
		obj.events <- Event{event, nil, "", activity}
		return true
	}

	resp := make(chan bool)
	obj.events <- Event{event, resp, "", activity}
	for {
		value := <-resp
		// wait until true value
		if value {
			return true
		}
	}
}

// process events when a select gets one, this handles the pause code too!
// the return values specify if we should exit and poke respectively
func (obj *BaseRes) ReadEvent(event *Event) (exit, poke bool) {
	event.ACK()
	switch event.Name {
	case eventStart:
		return false, true

	case eventPoke:
		return false, true

	case eventBackPoke:
		return false, true // forward poking in response to a back poke!

	case eventExit:
		return true, false

	case eventPause:
		// wait for next event to continue
		select {
		case e := <-obj.events:
			e.ACK()
			if e.Name == eventExit {
				return true, false
			} else if e.Name == eventStart { // eventContinue
				return false, false // don't poke on unpause!
			} else {
				// if we get a poke event here, it's a bug!
				log.Fatalf("%v[%v]: Unknown event: %v, while paused!", obj.Kind(), obj.GetName(), e)
			}
		}

	default:
		log.Fatal("Unknown event: ", event)
	}
	return true, false // required to keep the stupid go compiler happy
}

// GroupCmp compares two resources and decides if they're suitable for grouping
// You'll probably want to override this method when implementing a resource...
func (obj *BaseRes) GroupCmp(res Res) bool {
	return false // base implementation assumes false, override me!
}

func (obj *BaseRes) GroupRes(res Res) error {
	if l := len(res.GetGroup()); l > 0 {
		return fmt.Errorf("Res: %v already contains %d grouped resources!", res, l)
	}
	if res.IsGrouped() {
		return fmt.Errorf("Res: %v is already grouped!", res)
	}

	obj.grouped = append(obj.grouped, res)
	res.SetGrouped(true) // i am contained _in_ a group
	return nil
}

func (obj *BaseRes) IsGrouped() bool { // am I grouped?
	return obj.isGrouped
}

func (obj *BaseRes) SetGrouped(b bool) {
	obj.isGrouped = b
}

func (obj *BaseRes) GetGroup() []Res { // return everyone grouped inside me
	return obj.grouped
}

func (obj *BaseRes) SetGroup(g []Res) {
	obj.grouped = g
}

func (obj *BaseRes) CollectPattern(pattern string) {
	// XXX: default method is empty
}

// ResToB64 encodes a resource to a base64 encoded string (after serialization)
func ResToB64(res Res) (string, error) {
	b := bytes.Buffer{}
	e := gob.NewEncoder(&b)
	err := e.Encode(&res) // pass with &
	if err != nil {
		return "", fmt.Errorf("Gob failed to encode: %v", err)
	}
	return base64.StdEncoding.EncodeToString(b.Bytes()), nil
}

// B64ToRes decodes a resource from a base64 encoded string (after deserialization)
func B64ToRes(str string) (Res, error) {
	var output interface{}
	bb, err := base64.StdEncoding.DecodeString(str)
	if err != nil {
		return nil, fmt.Errorf("Base64 failed to decode: %v", err)
	}
	b := bytes.NewBuffer(bb)
	d := gob.NewDecoder(b)
	err = d.Decode(&output) // pass with &
	if err != nil {
		return nil, fmt.Errorf("Gob failed to decode: %v", err)
	}
	res, ok := output.(Res)
	if !ok {
		return nil, fmt.Errorf("Output %v is not a Res", res)

	}
	return res, nil
}
<|endoftext|>


<commit_before>package webgo

import (
	"encoding/json"
	"fmt"
	"html/template"
	"net/http"
)

//ErrorData used to render the error page
type ErrorData struct {
	ErrCode        int
	ErrDescription string
}

type dOutput struct {
	Data   interface{} `json:"data"`
	Status int         `json:"status"`
}

type errOutput struct {
	Errors interface{} `json:"errors"`
	Status int         `json:"status"`
}

const (
	//HeaderContentType is the key for mentioning the response header content type
	HeaderContentType = "Content-Type"
	//JSONContentType is the MIME type when the response is JSON
	JSONContentType = "application/json"
	//HTMLContentType is the MIME type when the response is HTML
	HTMLContentType = "text/html; charset=UTF-8"

	//ErrInternalServer to send when there's an internal server error
	ErrInternalServer = "Internal server error."
)

//SendResponse is used to respond to any request (JSON response) based on the code, data etc.
func SendResponse(w http.ResponseWriter, data interface{}, rCode int) {
	w.Header().Set(HeaderContentType, JSONContentType)

	w.WriteHeader(rCode)

	// Encode data to json and send response
	if err := json.NewEncoder(w).Encode(dOutput{data, rCode}); err != nil {
		/*
			In case of encoding error, send "internal server error" after
			logging the actual error
		*/
		Log.Println(err)
		R500(w, struct {
			errors []string
		}{
			[]string{ErrInternalServer},
		})
	}
}

//SendError is used to respond to any request with an error
func SendError(w http.ResponseWriter, data interface{}, rCode int) {
	w.Header().Set(HeaderContentType, JSONContentType)

	w.WriteHeader(rCode)

	if err := json.NewEncoder(w).Encode(errOutput{data, rCode}); err != nil {
		/*
			In case of encoding error, send "internal server error" after
			logging the actual error
		*/
		Log.Println(err)
		R500(w, struct {
			errors []string
		}{
			[]string{ErrInternalServer},
		})
	}
}

// Render is used for rendering templates (HTML)
func Render(w http.ResponseWriter, data interface{}, rCode int, tpl *template.Template) {
	// In case of HTML response, setting appropriate header type for text/HTML response
	w.Header().Set(HeaderContentType, HTMLContentType)

	w.WriteHeader(rCode)

	// Rendering an HTML template with appropriate data
	tpl.Execute(w, data)
}

// Render404 - used to render a 404 page
func Render404(w http.ResponseWriter, tpl *template.Template) {
	Render(w, ErrorData{
		404,
		"Sorry, the URL you requested was not found on this server... Or you're lost :-/",
	},
		404,
		tpl,
	)
}

// R200 - Successful/OK response
func R200(w http.ResponseWriter, data interface{}) {
	SendResponse(w, data, 200)
}

// R201 - New item created
func R201(w http.ResponseWriter, data interface{}) {
	SendResponse(w, data, 201)
}

// R204 - empty, no content
func R204(w http.ResponseWriter) {
	w.WriteHeader(204)
	fmt.Fprint(w)
}

// R302 - Temporary redirect
func R302(w http.ResponseWriter, data interface{}) {
	SendResponse(w, data, 302)
}

// R400 - Invalid request, any incorrect/erraneous value in the request body
func R400(w http.ResponseWriter, data interface{}) {
	SendError(w, data, 400)
}

// R403 - Unauthorized access
func R403(w http.ResponseWriter, data interface{}) {
	SendError(w, data, 403)
}

// R404 - Resource not found
func R404(w http.ResponseWriter, data interface{}) {
	SendError(w, data, 404)
}

// R406 - Unacceptable header. For any error related to values set in header
func R406(w http.ResponseWriter, data interface{}) {
	SendError(w, data, 406)
}

// R451 - Resource taken down because of a legal request
func R451(w http.ResponseWriter, data interface{}) {
	SendError(w, data, 451)
}

// R500 - Internal server error
func R500(w http.ResponseWriter, data interface{}) {
	SendError(w, data, 500)
}
<commit_msg>encoded payload struct is now sent by pointer<commit_after>package webgo

import (
	"encoding/json"
	"fmt"
	"html/template"
	"net/http"
)

//ErrorData used to render the error page
type ErrorData struct {
	ErrCode        int
	ErrDescription string
}

type dOutput struct {
	Data   interface{} `json:"data"`
	Status int         `json:"status"`
}

type errOutput struct {
	Errors interface{} `json:"errors"`
	Status int         `json:"status"`
}

const (
	//HeaderContentType is the key for mentioning the response header content type
	HeaderContentType = "Content-Type"
	//JSONContentType is the MIME type when the response is JSON
	JSONContentType = "application/json"
	//HTMLContentType is the MIME type when the response is HTML
	HTMLContentType = "text/html; charset=UTF-8"

	//ErrInternalServer to send when there's an internal server error
	ErrInternalServer = "Internal server error."
)

//SendResponse is used to respond to any request (JSON response) based on the code, data etc.
func SendResponse(w http.ResponseWriter, data interface{}, rCode int) {
	w.Header().Set(HeaderContentType, JSONContentType)

	w.WriteHeader(rCode)

	// Encode data to json and send response
	if err := json.NewEncoder(w).Encode(&dOutput{data, rCode}); err != nil {
		/*
			In case of encoding error, send "internal server error" after
			logging the actual error
		*/
		Log.Println(err)
		R500(w, struct {
			errors []string
		}{
			[]string{ErrInternalServer},
		})
	}
}

//SendError is used to respond to any request with an error
func SendError(w http.ResponseWriter, data interface{}, rCode int) {
	w.Header().Set(HeaderContentType, JSONContentType)

	w.WriteHeader(rCode)

	if err := json.NewEncoder(w).Encode(&errOutput{data, rCode}); err != nil {
		/*
			In case of encoding error, send "internal server error" after
			logging the actual error
		*/
		Log.Println(err)
		R500(w, struct {
			errors []string
		}{
			[]string{ErrInternalServer},
		})
	}
}

// Render is used for rendering templates (HTML)
func Render(w http.ResponseWriter, data interface{}, rCode int, tpl *template.Template) {
	// In case of HTML response, setting appropriate header type for text/HTML response
	w.Header().Set(HeaderContentType, HTMLContentType)

	w.WriteHeader(rCode)

	// Rendering an HTML template with appropriate data
	tpl.Execute(w, data)
}

// Render404 - used to render a 404 page
func Render404(w http.ResponseWriter, tpl *template.Template) {
	Render(w, ErrorData{
		404,
		"Sorry, the URL you requested was not found on this server... Or you're lost :-/",
	},
		404,
		tpl,
	)
}

// R200 - Successful/OK response
func R200(w http.ResponseWriter, data interface{}) {
	SendResponse(w, data, 200)
}

// R201 - New item created
func R201(w http.ResponseWriter, data interface{}) {
	SendResponse(w, data, 201)
}

// R204 - empty, no content
func R204(w http.ResponseWriter) {
	w.WriteHeader(204)
	fmt.Fprint(w)
}

// R302 - Temporary redirect
func R302(w http.ResponseWriter, data interface{}) {
	SendResponse(w, data, 302)
}

// R400 - Invalid request, any incorrect/erraneous value in the request body
func R400(w http.ResponseWriter, data interface{}) {
	SendError(w, data, 400)
}

// R403 - Unauthorized access
func R403(w http.ResponseWriter, data interface{}) {
	SendError(w, data, 403)
}

// R404 - Resource not found
func R404(w http.ResponseWriter, data interface{}) {
	SendError(w, data, 404)
}

// R406 - Unacceptable header. For any error related to values set in header
func R406(w http.ResponseWriter, data interface{}) {
	SendError(w, data, 406)
}

// R451 - Resource taken down because of a legal request
func R451(w http.ResponseWriter, data interface{}) {
	SendError(w, data, 451)
}

// R500 - Internal server error
func R500(w http.ResponseWriter, data interface{}) {
	SendError(w, data, 500)
}
<|endoftext|>


<commit_before>//  Copyright (c) 2014 Couchbase, Inc.
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the
//  License. You may obtain a copy of the License at
//    http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing,
//  software distributed under the License is distributed on an "AS
//  IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
//  express or implied. See the License for the specific language
//  governing permissions and limitations under the License.

package main

import (
	"bytes"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"reflect"
	"testing"
)

func TestNewManagerRESTRouter(t *testing.T) {
	emptyDir, _ := ioutil.TempDir("./tmp", "test")
	defer os.RemoveAll(emptyDir)

	ring, err := NewMsgRing(nil, 1)

	cfg := NewCfgMem()
	mgr := NewManager(VERSION, cfg, NewUUID(), nil, "", 1, ":1000",
		emptyDir, "some-datasource", nil)
	r, err := NewManagerRESTRouter(mgr, emptyDir, ring)
	if r == nil || err != nil {
		t.Errorf("expected no errors")
	}

	mgr = NewManager(VERSION, cfg, NewUUID(), []string{"queryer", "anotherTag"},
		"", 1, ":1000", emptyDir, "some-datasource", nil)
	r, err = NewManagerRESTRouter(mgr, emptyDir, ring)
	if r == nil || err != nil {
		t.Errorf("expected no errors")
	}
}

func TestHandlers(t *testing.T) {
	emptyDir, _ := ioutil.TempDir("./tmp", "test")
	defer os.RemoveAll(emptyDir)

	cfg := NewCfgMem()
	meh := &TestMEH{}
	mgr := NewManager(VERSION, cfg, NewUUID(),
		nil, "", 1, ":1000", emptyDir, "some-datasource", meh)
	mgr.Start("wanted")
	mgr.Kick("test-start-kick")

	mr, _ := NewMsgRing(os.Stderr, 1000)

	router, err := NewManagerRESTRouter(mgr, "static", mr)
	if err != nil || router == nil {
		t.Errorf("no mux router")
	}

	tests := []struct {
		Desc          string
		Handler       http.Handler
		Path          string
		Method        string
		Params        url.Values
		Body          []byte
		Status        int
		ResponseBody  []byte
		ResponseMatch map[string]bool
	}{
		{
			Desc:         "log",
			Handler:      NewGetLogHandler(mr),
			Path:         "/api/log",
			Method:       "GET",
			Params:       nil,
			Body:         nil,
			Status:       http.StatusOK,
			ResponseBody: []byte(`{"messages":[]}`),
		},
		{
			Desc:         "list empty indexes",
			Handler:      NewListIndexHandler(mgr),
			Path:         "/api/index",
			Method:       "GET",
			Params:       nil,
			Body:         nil,
			Status:       http.StatusOK,
			ResponseBody: []byte(`{"status":"ok","indexDefs":null}`),
		},
		{
			Desc:         "try to get a nonexistent index",
			Handler:      NewGetIndexHandler(mgr),
			Path:         "/api/index/NOT-AN-INDEX",
			Method:       "GET",
			Params:       nil,
			Body:         nil,
			Status:       400,
			ResponseBody: []byte(`not an index`),
		},
		{
			Desc:    "try to create a default index with bad server",
			Handler: NewCreateIndexHandler(mgr),
			Path:    "/api/index/index-on-a-bad-server",
			Method:  "PUT",
			Params:  nil,
			Body:    nil,
			Status:  500,
			ResponseMatch: map[string]bool{
				`failed to connect`: true,
			},
		},
		{
			Desc:    "try to delete a nonexistent index when no indexes",
			Handler: NewDeleteIndexHandler(mgr),
			Path:    "/api/index/NOT-AN-INDEX",
			Method:  "DELETE",
			Params:  nil,
			Body:    nil,
			Status:  400,
			ResponseMatch: map[string]bool{
				`indexes do not exist`: true,
			},
		},
		{
			Desc:    "try to count a nonexistent index when no indexes",
			Handler: NewCountHandler(mgr),
			Path:    "/api/index/NOT-AN-INDEX/count",
			Method:  "GET",
			Params:  nil,
			Body:    nil,
			Status:  400,
			ResponseMatch: map[string]bool{
				`could not get indexDefs`: true,
			},
		},
		{
			Desc:    "try to query a nonexistent index when no indexes",
			Handler: NewQueryHandler(mgr),
			Path:    "/api/index/NOT-AN-INDEX/query",
			Method:  "POST",
			Params:  nil,
			Body:    nil,
			Status:  400,
			ResponseMatch: map[string]bool{
				`could not get indexDefs`: true,
			},
		},
	}

	for _, test := range tests {
		record := httptest.NewRecorder()
		req := &http.Request{
			Method: test.Method,
			URL:    &url.URL{Path: test.Path},
			Form:   test.Params,
			Body:   ioutil.NopCloser(bytes.NewBuffer(test.Body)),
		}
		router.ServeHTTP(record, req)
		if got, want := record.Code, test.Status; got != want {
			t.Errorf("%s: response code = %d, want %d", test.Desc, got, want)
			t.Errorf("%s: response body = %s", test.Desc, record.Body)
		}

		got := bytes.TrimRight(record.Body.Bytes(), "\n")
		if test.ResponseBody != nil {
			if !reflect.DeepEqual(got, test.ResponseBody) {
				t.Errorf("%s: expected: '%s', got: '%s'",
					test.Desc, test.ResponseBody, got)
			}
		}
		for pattern, shouldMatch := range test.ResponseMatch {
			didMatch := bytes.Contains(got, []byte(pattern))
			if didMatch != shouldMatch {
				t.Errorf("%s: expected match %t for pattern %s, got %t",
					test.Desc, shouldMatch, pattern, didMatch)
				t.Errorf("%s: response body was: %s", test.Desc, got)
			}
		}
	}
}
<commit_msg>first test of /api/cfg REST endpoint<commit_after>//  Copyright (c) 2014 Couchbase, Inc.
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the
//  License. You may obtain a copy of the License at
//    http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing,
//  software distributed under the License is distributed on an "AS
//  IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
//  express or implied. See the License for the specific language
//  governing permissions and limitations under the License.

package main

import (
	"bytes"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"reflect"
	"testing"
)

func TestNewManagerRESTRouter(t *testing.T) {
	emptyDir, _ := ioutil.TempDir("./tmp", "test")
	defer os.RemoveAll(emptyDir)

	ring, err := NewMsgRing(nil, 1)

	cfg := NewCfgMem()
	mgr := NewManager(VERSION, cfg, NewUUID(), nil, "", 1, ":1000",
		emptyDir, "some-datasource", nil)
	r, err := NewManagerRESTRouter(mgr, emptyDir, ring)
	if r == nil || err != nil {
		t.Errorf("expected no errors")
	}

	mgr = NewManager(VERSION, cfg, NewUUID(), []string{"queryer", "anotherTag"},
		"", 1, ":1000", emptyDir, "some-datasource", nil)
	r, err = NewManagerRESTRouter(mgr, emptyDir, ring)
	if r == nil || err != nil {
		t.Errorf("expected no errors")
	}
}

func TestHandlers(t *testing.T) {
	emptyDir, _ := ioutil.TempDir("./tmp", "test")
	defer os.RemoveAll(emptyDir)

	cfg := NewCfgMem()
	meh := &TestMEH{}
	mgr := NewManager(VERSION, cfg, NewUUID(),
		nil, "", 1, ":1000", emptyDir, "some-datasource", meh)
	mgr.Start("wanted")
	mgr.Kick("test-start-kick")

	mr, _ := NewMsgRing(os.Stderr, 1000)

	router, err := NewManagerRESTRouter(mgr, "static", mr)
	if err != nil || router == nil {
		t.Errorf("no mux router")
	}

	tests := []struct {
		Desc          string
		Handler       http.Handler
		Path          string
		Method        string
		Params        url.Values
		Body          []byte
		Status        int
		ResponseBody  []byte
		ResponseMatch map[string]bool
	}{
		{
			Desc:         "log",
			Handler:      NewGetLogHandler(mr),
			Path:         "/api/log",
			Method:       "GET",
			Params:       nil,
			Body:         nil,
			Status:       http.StatusOK,
			ResponseBody: []byte(`{"messages":[]}`),
		},
		{
			Desc:    "cfg",
			Handler: NewCfgGetHandler(mgr),
			Path:    "/api/cfg",
			Method:  "GET",
			Params:  nil,
			Body:    nil,
			Status:  http.StatusOK,
			ResponseMatch: map[string]bool{
				`"status":"ok"`: true,
				`"indexDefs":null`: true,
				`"nodeDefsKnown":{`: true,
				`"nodeDefsWanted":{`: true,
				`"planPIndexes":null`: true,
			},
		},
		{
			Desc:         "list empty indexes",
			Handler:      NewListIndexHandler(mgr),
			Path:         "/api/index",
			Method:       "GET",
			Params:       nil,
			Body:         nil,
			Status:       http.StatusOK,
			ResponseBody: []byte(`{"status":"ok","indexDefs":null}`),
		},
		{
			Desc:         "try to get a nonexistent index",
			Handler:      NewGetIndexHandler(mgr),
			Path:         "/api/index/NOT-AN-INDEX",
			Method:       "GET",
			Params:       nil,
			Body:         nil,
			Status:       400,
			ResponseBody: []byte(`not an index`),
		},
		{
			Desc:    "try to create a default index with bad server",
			Handler: NewCreateIndexHandler(mgr),
			Path:    "/api/index/index-on-a-bad-server",
			Method:  "PUT",
			Params:  nil,
			Body:    nil,
			Status:  500,
			ResponseMatch: map[string]bool{
				`failed to connect`: true,
			},
		},
		{
			Desc:    "try to delete a nonexistent index when no indexes",
			Handler: NewDeleteIndexHandler(mgr),
			Path:    "/api/index/NOT-AN-INDEX",
			Method:  "DELETE",
			Params:  nil,
			Body:    nil,
			Status:  400,
			ResponseMatch: map[string]bool{
				`indexes do not exist`: true,
			},
		},
		{
			Desc:    "try to count a nonexistent index when no indexes",
			Handler: NewCountHandler(mgr),
			Path:    "/api/index/NOT-AN-INDEX/count",
			Method:  "GET",
			Params:  nil,
			Body:    nil,
			Status:  400,
			ResponseMatch: map[string]bool{
				`could not get indexDefs`: true,
			},
		},
		{
			Desc:    "try to query a nonexistent index when no indexes",
			Handler: NewQueryHandler(mgr),
			Path:    "/api/index/NOT-AN-INDEX/query",
			Method:  "POST",
			Params:  nil,
			Body:    nil,
			Status:  400,
			ResponseMatch: map[string]bool{
				`could not get indexDefs`: true,
			},
		},
	}

	for _, test := range tests {
		record := httptest.NewRecorder()
		req := &http.Request{
			Method: test.Method,
			URL:    &url.URL{Path: test.Path},
			Form:   test.Params,
			Body:   ioutil.NopCloser(bytes.NewBuffer(test.Body)),
		}
		router.ServeHTTP(record, req)
		if got, want := record.Code, test.Status; got != want {
			t.Errorf("%s: response code = %d, want %d", test.Desc, got, want)
			t.Errorf("%s: response body = %s", test.Desc, record.Body)
		}

		got := bytes.TrimRight(record.Body.Bytes(), "\n")
		if test.ResponseBody != nil {
			if !reflect.DeepEqual(got, test.ResponseBody) {
				t.Errorf("%s: expected: '%s', got: '%s'",
					test.Desc, test.ResponseBody, got)
			}
		}
		for pattern, shouldMatch := range test.ResponseMatch {
			didMatch := bytes.Contains(got, []byte(pattern))
			if didMatch != shouldMatch {
				t.Errorf("%s: expected match %t for pattern %s, got %t",
					test.Desc, shouldMatch, pattern, didMatch)
				t.Errorf("%s: response body was: %s", test.Desc, got)
			}
		}
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package provider

import (
	"strings"
	"testing"

	"github.com/kubernetes-incubator/external-dns/endpoint"
	"github.com/kubernetes-incubator/external-dns/plan"
)

type fakeETCDClient struct {
	services map[string]*Service
}

func (c fakeETCDClient) GetServices(prefix string) ([]*Service, error) {
	var result []*Service
	for key, value := range c.services {
		if strings.HasPrefix(key, prefix) {
			value.Key = key
			result = append(result, value)
		}
	}
	return result, nil
}

func (c fakeETCDClient) SaveService(service *Service) error {
	c.services[service.Key] = service
	return nil
}

func (c fakeETCDClient) DeleteService(key string) error {
	delete(c.services, key)
	return nil
}

func TestAServiceTranslation(t *testing.T) {
	expectedTarget := "1.2.3.4"
	expectedDNSName := "example.com"
	expectedRecordType := endpoint.RecordTypeA

	client := fakeETCDClient{
		map[string]*Service{
			"/skydns/com/example": {Host: expectedTarget},
		},
	}
	provider := coreDNSProvider{client: client}
	endpoints, err := provider.Records()
	if err != nil {
		t.Fatal(err)
	}
	if len(endpoints) != 1 {
		t.Fatalf("got unexpected number of endpoints: %d", len(endpoints))
	}
	if endpoints[0].DNSName != expectedDNSName {
		t.Errorf("got unexpected DNS name: %s != %s", endpoints[0].DNSName, expectedDNSName)
	}
	if endpoints[0].Target != expectedTarget {
		t.Errorf("got unexpected DNS target: %s != %s", endpoints[0].Target, expectedTarget)
	}
	if endpoints[0].RecordType != expectedRecordType {
		t.Errorf("got unexpected DNS record type: %s != %s", endpoints[0].RecordType, expectedRecordType)
	}
}

func TestCNAMEServiceTranslation(t *testing.T) {
	expectedTarget := "example.net"
	expectedDNSName := "example.com"
	expectedRecordType := endpoint.RecordTypeCNAME

	client := fakeETCDClient{
		map[string]*Service{
			"/skydns/com/example": {Host: expectedTarget},
		},
	}
	provider := coreDNSProvider{client: client}
	endpoints, err := provider.Records()
	if err != nil {
		t.Fatal(err)
	}
	if len(endpoints) != 1 {
		t.Fatalf("got unexpected number of endpoints: %d", len(endpoints))
	}
	if endpoints[0].DNSName != expectedDNSName {
		t.Errorf("got unexpected DNS name: %s != %s", endpoints[0].DNSName, expectedDNSName)
	}
	if endpoints[0].Target != expectedTarget {
		t.Errorf("got unexpected DNS target: %s != %s", endpoints[0].Target, expectedTarget)
	}
	if endpoints[0].RecordType != expectedRecordType {
		t.Errorf("got unexpected DNS record type: %s != %s", endpoints[0].RecordType, expectedRecordType)
	}
}

func TestTXTServiceTranslation(t *testing.T) {
	expectedTarget := "string"
	expectedDNSName := "example.com"
	expectedRecordType := endpoint.RecordTypeTXT

	client := fakeETCDClient{
		map[string]*Service{
			"/skydns/com/example": {Text: expectedTarget},
		},
	}
	provider := coreDNSProvider{client: client}
	endpoints, err := provider.Records()
	if err != nil {
		t.Fatal(err)
	}
	if len(endpoints) != 1 {
		t.Fatalf("got unexpected number of endpoints: %d", len(endpoints))
	}
	if endpoints[0].DNSName != expectedDNSName {
		t.Errorf("got unexpected DNS name: %s != %s", endpoints[0].DNSName, expectedDNSName)
	}
	if endpoints[0].Target != expectedTarget {
		t.Errorf("got unexpected DNS target: %s != %s", endpoints[0].Target, expectedTarget)
	}
	if endpoints[0].RecordType != expectedRecordType {
		t.Errorf("got unexpected DNS record type: %s != %s", endpoints[0].RecordType, expectedRecordType)
	}
}

func TestAWithTXTServiceTranslation(t *testing.T) {
	expectedTargets := map[string]string{
		endpoint.RecordTypeA:   "1.2.3.4",
		endpoint.RecordTypeTXT: "string",
	}
	expectedDNSName := "example.com"

	client := fakeETCDClient{
		map[string]*Service{
			"/skydns/com/example": {Host: "1.2.3.4", Text: "string"},
		},
	}
	provider := coreDNSProvider{client: client}
	endpoints, err := provider.Records()
	if err != nil {
		t.Fatal(err)
	}
	if len(endpoints) != len(expectedTargets) {
		t.Fatalf("got unexpected number of endpoints: %d", len(endpoints))
	}

	for _, ep := range endpoints {
		expectedTarget := expectedTargets[ep.RecordType]
		if expectedTarget == "" {
			t.Errorf("got unexpected DNS record type: %s", ep.RecordType)
			continue
		}
		delete(expectedTargets, ep.RecordType)

		if ep.DNSName != expectedDNSName {
			t.Errorf("got unexpected DNS name: %s != %s", ep.DNSName, expectedDNSName)
		}

		if ep.Target != expectedTarget {
			t.Errorf("got unexpected DNS target: %s != %s", ep.Target, expectedTarget)
		}
	}
}

func TestCNAMEWithTXTServiceTranslation(t *testing.T) {
	expectedTargets := map[string]string{
		endpoint.RecordTypeCNAME: "example.net",
		endpoint.RecordTypeTXT:   "string",
	}
	expectedDNSName := "example.com"

	client := fakeETCDClient{
		map[string]*Service{
			"/skydns/com/example": {Host: "example.net", Text: "string"},
		},
	}
	provider := coreDNSProvider{client: client}
	endpoints, err := provider.Records()
	if err != nil {
		t.Fatal(err)
	}
	if len(endpoints) != len(expectedTargets) {
		t.Fatalf("got unexpected number of endpoints: %d", len(endpoints))
	}

	for _, ep := range endpoints {
		expectedTarget := expectedTargets[ep.RecordType]
		if expectedTarget == "" {
			t.Errorf("got unexpected DNS record type: %s", ep.RecordType)
			continue
		}
		delete(expectedTargets, ep.RecordType)

		if ep.DNSName != expectedDNSName {
			t.Errorf("got unexpected DNS name: %s != %s", ep.DNSName, expectedDNSName)
		}

		if ep.Target != expectedTarget {
			t.Errorf("got unexpected DNS target: %s != %s", ep.Target, expectedTarget)
		}
	}
}

func TestCoreDNSApplyChanges(t *testing.T) {
	client := fakeETCDClient{
		map[string]*Service{},
	}
	coredns := coreDNSProvider{client: client}

	changes1 := &plan.Changes{
		Create: []*endpoint.Endpoint{
			endpoint.NewEndpoint("domain1.local", "5.5.5.5", endpoint.RecordTypeA),
			endpoint.NewEndpoint("domain1.local", "string1", endpoint.RecordTypeTXT),
			endpoint.NewEndpoint("domain2.local", "site.local", endpoint.RecordTypeCNAME),
		},
	}
	coredns.ApplyChanges(changes1)

	expectedServices1 := map[string]*Service{
		"/skydns/local/domain1": {Host: "5.5.5.5", Text: "string1"},
		"/skydns/local/domain2": {Host: "site.local"},
	}
	validateServices(client.services, expectedServices1, t, 1)

	updatedEp := endpoint.NewEndpoint("domain1.local", "6.6.6.6", endpoint.RecordTypeA)
	updatedEp.Labels["originalText"] = "string1"
	changes2 := &plan.Changes{
		Create: []*endpoint.Endpoint{
			endpoint.NewEndpoint("domain3.local", "7.7.7.7", endpoint.RecordTypeA),
		},
		UpdateNew: []*endpoint.Endpoint{
			endpoint.NewEndpoint("domain1.local", "6.6.6.6", "A"),
		},
	}
	applyServiceChanges(coredns, changes2)

	expectedServices2 := map[string]*Service{
		"/skydns/local/domain1": {Host: "6.6.6.6", Text: "string1"},
		"/skydns/local/domain2": {Host: "site.local"},
		"/skydns/local/domain3": {Host: "7.7.7.7"},
	}
	validateServices(client.services, expectedServices2, t, 2)

	changes3 := &plan.Changes{
		Delete: []*endpoint.Endpoint{
			endpoint.NewEndpoint("domain1.local", "6.6.6.6", endpoint.RecordTypeA),
			endpoint.NewEndpoint("domain1.local", "string", endpoint.RecordTypeTXT),
			endpoint.NewEndpoint("domain3.local", "7.7.7.7", endpoint.RecordTypeA),
		},
	}

	applyServiceChanges(coredns, changes3)

	expectedServices3 := map[string]*Service{
		"/skydns/local/domain2": {Host: "site.local"},
	}
	validateServices(client.services, expectedServices3, t, 3)
}

func applyServiceChanges(provider coreDNSProvider, changes *plan.Changes) {
	records, _ := provider.Records()
	for _, col := range [][]*endpoint.Endpoint{changes.Create, changes.UpdateNew, changes.Delete} {
		for _, record := range col {
			for _, existingRecord := range records {
				if existingRecord.DNSName == record.DNSName && existingRecord.RecordType == record.RecordType {
					record.MergeLabels(existingRecord.Labels)
				}
			}
		}
	}
	provider.ApplyChanges(changes)
}

func validateServices(services, expectedServices map[string]*Service, t *testing.T, step int) {
	if len(services) != len(expectedServices) {
		t.Errorf("wrong number of records on step %d: %d != %d", step, len(services), len(expectedServices))
	}
	for key, value := range services {
		keyParts := strings.Split(key, "/")
		expectedKey := strings.Join(keyParts[:len(keyParts)-value.TargetStrip], "/")
		expectedService := expectedServices[expectedKey]
		if expectedService == nil {
			t.Errorf("unexpected service %s", key)
			continue
		}
		delete(expectedServices, key)
		if value.Host != expectedService.Host {
			t.Errorf("wrong host for service %s: %s != %s on step %d", key, value.Host, expectedService.Host, step)
		}
		if value.Text != expectedService.Text {
			t.Errorf("wrong text for service %s: %s != %s on step %d", key, value.Text, expectedService.Text, step)
		}
	}
}
<commit_msg>fix: make coredns provider compile<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package provider

import (
	"strings"
	"testing"

	"github.com/kubernetes-incubator/external-dns/endpoint"
	"github.com/kubernetes-incubator/external-dns/plan"
)

type fakeETCDClient struct {
	services map[string]*Service
}

func (c fakeETCDClient) GetServices(prefix string) ([]*Service, error) {
	var result []*Service
	for key, value := range c.services {
		if strings.HasPrefix(key, prefix) {
			value.Key = key
			result = append(result, value)
		}
	}
	return result, nil
}

func (c fakeETCDClient) SaveService(service *Service) error {
	c.services[service.Key] = service
	return nil
}

func (c fakeETCDClient) DeleteService(key string) error {
	delete(c.services, key)
	return nil
}

func TestAServiceTranslation(t *testing.T) {
	expectedTarget := "1.2.3.4"
	expectedDNSName := "example.com"
	expectedRecordType := endpoint.RecordTypeA

	client := fakeETCDClient{
		map[string]*Service{
			"/skydns/com/example": {Host: expectedTarget},
		},
	}
	provider := coreDNSProvider{client: client}
	endpoints, err := provider.Records()
	if err != nil {
		t.Fatal(err)
	}
	if len(endpoints) != 1 {
		t.Fatalf("got unexpected number of endpoints: %d", len(endpoints))
	}
	if endpoints[0].DNSName != expectedDNSName {
		t.Errorf("got unexpected DNS name: %s != %s", endpoints[0].DNSName, expectedDNSName)
	}
	if endpoints[0].Target != expectedTarget {
		t.Errorf("got unexpected DNS target: %s != %s", endpoints[0].Target, expectedTarget)
	}
	if endpoints[0].RecordType != expectedRecordType {
		t.Errorf("got unexpected DNS record type: %s != %s", endpoints[0].RecordType, expectedRecordType)
	}
}

func TestCNAMEServiceTranslation(t *testing.T) {
	expectedTarget := "example.net"
	expectedDNSName := "example.com"
	expectedRecordType := endpoint.RecordTypeCNAME

	client := fakeETCDClient{
		map[string]*Service{
			"/skydns/com/example": {Host: expectedTarget},
		},
	}
	provider := coreDNSProvider{client: client}
	endpoints, err := provider.Records()
	if err != nil {
		t.Fatal(err)
	}
	if len(endpoints) != 1 {
		t.Fatalf("got unexpected number of endpoints: %d", len(endpoints))
	}
	if endpoints[0].DNSName != expectedDNSName {
		t.Errorf("got unexpected DNS name: %s != %s", endpoints[0].DNSName, expectedDNSName)
	}
	if endpoints[0].Target != expectedTarget {
		t.Errorf("got unexpected DNS target: %s != %s", endpoints[0].Target, expectedTarget)
	}
	if endpoints[0].RecordType != expectedRecordType {
		t.Errorf("got unexpected DNS record type: %s != %s", endpoints[0].RecordType, expectedRecordType)
	}
}

func TestTXTServiceTranslation(t *testing.T) {
	expectedTarget := "string"
	expectedDNSName := "example.com"
	expectedRecordType := endpoint.RecordTypeTXT

	client := fakeETCDClient{
		map[string]*Service{
			"/skydns/com/example": {Text: expectedTarget},
		},
	}
	provider := coreDNSProvider{client: client}
	endpoints, err := provider.Records()
	if err != nil {
		t.Fatal(err)
	}
	if len(endpoints) != 1 {
		t.Fatalf("got unexpected number of endpoints: %d", len(endpoints))
	}
	if endpoints[0].DNSName != expectedDNSName {
		t.Errorf("got unexpected DNS name: %s != %s", endpoints[0].DNSName, expectedDNSName)
	}
	if endpoints[0].Target != expectedTarget {
		t.Errorf("got unexpected DNS target: %s != %s", endpoints[0].Target, expectedTarget)
	}
	if endpoints[0].RecordType != expectedRecordType {
		t.Errorf("got unexpected DNS record type: %s != %s", endpoints[0].RecordType, expectedRecordType)
	}
}

func TestAWithTXTServiceTranslation(t *testing.T) {
	expectedTargets := map[string]string{
		endpoint.RecordTypeA:   "1.2.3.4",
		endpoint.RecordTypeTXT: "string",
	}
	expectedDNSName := "example.com"

	client := fakeETCDClient{
		map[string]*Service{
			"/skydns/com/example": {Host: "1.2.3.4", Text: "string"},
		},
	}
	provider := coreDNSProvider{client: client}
	endpoints, err := provider.Records()
	if err != nil {
		t.Fatal(err)
	}
	if len(endpoints) != len(expectedTargets) {
		t.Fatalf("got unexpected number of endpoints: %d", len(endpoints))
	}

	for _, ep := range endpoints {
		expectedTarget := expectedTargets[ep.RecordType]
		if expectedTarget == "" {
			t.Errorf("got unexpected DNS record type: %s", ep.RecordType)
			continue
		}
		delete(expectedTargets, ep.RecordType)

		if ep.DNSName != expectedDNSName {
			t.Errorf("got unexpected DNS name: %s != %s", ep.DNSName, expectedDNSName)
		}

		if ep.Target != expectedTarget {
			t.Errorf("got unexpected DNS target: %s != %s", ep.Target, expectedTarget)
		}
	}
}

func TestCNAMEWithTXTServiceTranslation(t *testing.T) {
	expectedTargets := map[string]string{
		endpoint.RecordTypeCNAME: "example.net",
		endpoint.RecordTypeTXT:   "string",
	}
	expectedDNSName := "example.com"

	client := fakeETCDClient{
		map[string]*Service{
			"/skydns/com/example": {Host: "example.net", Text: "string"},
		},
	}
	provider := coreDNSProvider{client: client}
	endpoints, err := provider.Records()
	if err != nil {
		t.Fatal(err)
	}
	if len(endpoints) != len(expectedTargets) {
		t.Fatalf("got unexpected number of endpoints: %d", len(endpoints))
	}

	for _, ep := range endpoints {
		expectedTarget := expectedTargets[ep.RecordType]
		if expectedTarget == "" {
			t.Errorf("got unexpected DNS record type: %s", ep.RecordType)
			continue
		}
		delete(expectedTargets, ep.RecordType)

		if ep.DNSName != expectedDNSName {
			t.Errorf("got unexpected DNS name: %s != %s", ep.DNSName, expectedDNSName)
		}

		if ep.Target != expectedTarget {
			t.Errorf("got unexpected DNS target: %s != %s", ep.Target, expectedTarget)
		}
	}
}

func TestCoreDNSApplyChanges(t *testing.T) {
	client := fakeETCDClient{
		map[string]*Service{},
	}
	coredns := coreDNSProvider{client: client}

	changes1 := &plan.Changes{
		Create: []*endpoint.Endpoint{
			endpoint.NewEndpoint("domain1.local", "5.5.5.5", endpoint.RecordTypeA),
			endpoint.NewEndpoint("domain1.local", "string1", endpoint.RecordTypeTXT),
			endpoint.NewEndpoint("domain2.local", "site.local", endpoint.RecordTypeCNAME),
		},
	}
	coredns.ApplyChanges(changes1)

	expectedServices1 := map[string]*Service{
		"/skydns/local/domain1": {Host: "5.5.5.5", Text: "string1"},
		"/skydns/local/domain2": {Host: "site.local"},
	}
	validateServices(client.services, expectedServices1, t, 1)

	updatedEp := endpoint.NewEndpoint("domain1.local", "6.6.6.6", endpoint.RecordTypeA)
	updatedEp.Labels["originalText"] = "string1"
	changes2 := &plan.Changes{
		Create: []*endpoint.Endpoint{
			endpoint.NewEndpoint("domain3.local", "7.7.7.7", endpoint.RecordTypeA),
		},
		UpdateNew: []*endpoint.Endpoint{
			endpoint.NewEndpoint("domain1.local", "6.6.6.6", "A"),
		},
	}
	applyServiceChanges(coredns, changes2)

	expectedServices2 := map[string]*Service{
		"/skydns/local/domain1": {Host: "6.6.6.6", Text: "string1"},
		"/skydns/local/domain2": {Host: "site.local"},
		"/skydns/local/domain3": {Host: "7.7.7.7"},
	}
	validateServices(client.services, expectedServices2, t, 2)

	changes3 := &plan.Changes{
		Delete: []*endpoint.Endpoint{
			endpoint.NewEndpoint("domain1.local", "6.6.6.6", endpoint.RecordTypeA),
			endpoint.NewEndpoint("domain1.local", "string", endpoint.RecordTypeTXT),
			endpoint.NewEndpoint("domain3.local", "7.7.7.7", endpoint.RecordTypeA),
		},
	}

	applyServiceChanges(coredns, changes3)

	expectedServices3 := map[string]*Service{
		"/skydns/local/domain2": {Host: "site.local"},
	}
	validateServices(client.services, expectedServices3, t, 3)
}

func applyServiceChanges(provider coreDNSProvider, changes *plan.Changes) {
	records, _ := provider.Records()
	for _, col := range [][]*endpoint.Endpoint{changes.Create, changes.UpdateNew, changes.Delete} {
		for _, record := range col {
			for _, existingRecord := range records {
				if existingRecord.DNSName == record.DNSName && existingRecord.RecordType == record.RecordType {
					mergeLabels(record, existingRecord.Labels)
				}
			}
		}
	}
	provider.ApplyChanges(changes)
}

func validateServices(services, expectedServices map[string]*Service, t *testing.T, step int) {
	if len(services) != len(expectedServices) {
		t.Errorf("wrong number of records on step %d: %d != %d", step, len(services), len(expectedServices))
	}
	for key, value := range services {
		keyParts := strings.Split(key, "/")
		expectedKey := strings.Join(keyParts[:len(keyParts)-value.TargetStrip], "/")
		expectedService := expectedServices[expectedKey]
		if expectedService == nil {
			t.Errorf("unexpected service %s", key)
			continue
		}
		delete(expectedServices, key)
		if value.Host != expectedService.Host {
			t.Errorf("wrong host for service %s: %s != %s on step %d", key, value.Host, expectedService.Host, step)
		}
		if value.Text != expectedService.Text {
			t.Errorf("wrong text for service %s: %s != %s on step %d", key, value.Text, expectedService.Text, step)
		}
	}
}

// mergeLabels adds keys to labels if not defined for the endpoint
func mergeLabels(e *endpoint.Endpoint, labels map[string]string) {
	for k, v := range labels {
		if e.Labels[k] == "" {
			e.Labels[k] = v
		}
	}
}
<|endoftext|>