<commit_before>// Interface to both live and offline pcap parsing.
package pcap

/*
#cgo LDFLAGS: -lpcap
#include <stdlib.h>
#include <pcap.h>

// Workaround for not knowing how to cast to const u_char**
int hack_pcap_next_ex(pcap_t *p, struct pcap_pkthdr **pkt_header,
                      u_char **pkt_data) {
    return pcap_next_ex(p, pkt_header, (const u_char **)pkt_data);
}
*/
import "C"
import (
	"errors"
	"net"
	"syscall"
	"time"
	"unsafe"
)

type Pcap struct {
	cptr *C.pcap_t
}

type Stat struct {
	PacketsReceived  uint32
	PacketsDropped   uint32
	PacketsIfDropped uint32
}

type Interface struct {
	Name        string
	Description string
	Addresses   []IFAddress
	// TODO: add more elements
}

type IFAddress struct {
	IP      net.IP
	Netmask net.IPMask
	// TODO: add broadcast + PtP dst ?
}

func (p *Pcap) Next() (pkt *Packet) {
	rv, _ := p.NextEx()
	return rv
}

// Openlive opens a device and returns a *Pcap handler
func Openlive(device string, snaplen int32, promisc bool, timeout_ms int32) (handle *Pcap, err error) {
	var buf *C.char
	buf = (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	h := new(Pcap)
	var pro int32
	if promisc {
		pro = 1
	}

	dev := C.CString(device)
	defer C.free(unsafe.Pointer(dev))

	h.cptr = C.pcap_open_live(dev, C.int(snaplen), C.int(pro), C.int(timeout_ms), buf)
	if nil == h.cptr {
		handle = nil
		err = errors.New(C.GoString(buf))
	} else {
		handle = h
	}
	C.free(unsafe.Pointer(buf))
	return
}

func Openoffline(file string) (handle *Pcap, err error) {
	var buf *C.char
	buf = (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	h := new(Pcap)

	cf := C.CString(file)
	defer C.free(unsafe.Pointer(cf))

	h.cptr = C.pcap_open_offline(cf, buf)
	if nil == h.cptr {
		handle = nil
		err = errors.New(C.GoString(buf))
	} else {
		handle = h
	}
	C.free(unsafe.Pointer(buf))
	return
}

func (p *Pcap) NextEx() (pkt *Packet, result int32) {
	var pkthdr_ptr *C.struct_pcap_pkthdr
	var pkthdr C.struct_pcap_pkthdr

	var buf_ptr *C.u_char
	var buf unsafe.Pointer
	result = int32(C.hack_pcap_next_ex(p.cptr, &pkthdr_ptr, &buf_ptr))

	buf = unsafe.Pointer(buf_ptr)
	pkthdr = *pkthdr_ptr

	if nil == buf {
		return
	}
	pkt = new(Packet)
	pkt.Time = time.Unix(int64(pkthdr.ts.tv_sec), int64(pkthdr.ts.tv_usec))
	pkt.Caplen = uint32(pkthdr.caplen)
	pkt.Len = uint32(pkthdr.len)
	pkt.Data = make([]byte, pkthdr.caplen)

	for i := uint32(0); i < pkt.Caplen; i++ {
		pkt.Data[i] = *(*byte)(unsafe.Pointer(uintptr(buf) + uintptr(i)))
	}
	return
}

func (p *Pcap) Close() {
	C.pcap_close(p.cptr)
}

func (p *Pcap) Geterror() error {
	return errors.New(C.GoString(C.pcap_geterr(p.cptr)))
}

func (p *Pcap) Getstats() (stat *Stat, err error) {
	var cstats _Ctype_struct_pcap_stat
	if -1 == C.pcap_stats(p.cptr, &cstats) {
		return nil, p.Geterror()
	}
	stats := new(Stat)
	stats.PacketsReceived = uint32(cstats.ps_recv)
	stats.PacketsDropped = uint32(cstats.ps_drop)
	stats.PacketsIfDropped = uint32(cstats.ps_ifdrop)

	return stats, nil
}

func (p *Pcap) Setfilter(expr string) (err error) {
	var bpf _Ctype_struct_bpf_program
	cexpr := C.CString(expr)
	defer C.free(unsafe.Pointer(cexpr))

	if -1 == C.pcap_compile(p.cptr, &bpf, cexpr, 1, 0) {
		return p.Geterror()
	}

	if -1 == C.pcap_setfilter(p.cptr, &bpf) {
		C.pcap_freecode(&bpf)
		return p.Geterror()
	}
	C.pcap_freecode(&bpf)
	return nil
}

func Version() string {
	return C.GoString(C.pcap_lib_version())
}

func (p *Pcap) Datalink() int {
	return int(C.pcap_datalink(p.cptr))
}

func (p *Pcap) Setdatalink(dlt int) error {
	if -1 == C.pcap_set_datalink(p.cptr, C.int(dlt)) {
		return p.Geterror()
	}
	return nil
}

func DatalinkValueToName(dlt int) string {
	if name := C.pcap_datalink_val_to_name(C.int(dlt)); name != nil {
		return C.GoString(name)
	}
	return ""
}

func DatalinkValueToDescription(dlt int) string {
	if desc := C.pcap_datalink_val_to_description(C.int(dlt)); desc != nil {
		return C.GoString(desc)
	}
	return ""
}

func Findalldevs() (ifs []Interface, err error) {
	var buf *C.char
	buf = (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	defer C.free(unsafe.Pointer(buf))
	var alldevsp *C.pcap_if_t

	if -1 == C.pcap_findalldevs((**C.pcap_if_t)(&alldevsp), buf) {
		return nil, errors.New(C.GoString(buf))
	}
	defer C.pcap_freealldevs((*C.pcap_if_t)(alldevsp))
	dev := alldevsp
	var i uint32
	for i = 0; dev != nil; dev = (*C.pcap_if_t)(dev.next) {
		i++
	}
	ifs = make([]Interface, i)
	dev = alldevsp
	for j := uint32(0); dev != nil; dev = (*C.pcap_if_t)(dev.next) {
		var iface Interface
		iface.Name = C.GoString(dev.name)
		iface.Description = C.GoString(dev.description)
		iface.Addresses = findalladdresses(dev.addresses)
		// TODO: add more elements
		ifs[j] = iface
		j++
	}
	return
}

func findalladdresses(addresses *_Ctype_struct_pcap_addr) (retval []IFAddress) {
	// TODO - make it support more than IPv4 and IPv6?
	retval = make([]IFAddress, 0, 1)
	for curaddr := addresses; curaddr != nil; curaddr = (*_Ctype_struct_pcap_addr)(curaddr.next) {
		var a IFAddress
		var err error
		if a.IP, err = sockaddr_to_IP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr))); err != nil {
			continue
		}
		if a.Netmask, err = sockaddr_to_IP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr))); err != nil {
			continue
		}
		retval = append(retval, a)
	}
	return
}

func sockaddr_to_IP(rsa *syscall.RawSockaddr) (IP []byte, err error) {
	switch rsa.Family {
	case syscall.AF_INET:
		pp := (*syscall.RawSockaddrInet4)(unsafe.Pointer(rsa))
		IP = make([]byte, 4)
		for i := 0; i < len(IP); i++ {
			IP[i] = pp.Addr[i]
		}
		return
	case syscall.AF_INET6:
		pp := (*syscall.RawSockaddrInet6)(unsafe.Pointer(rsa))
		IP = make([]byte, 16)
		for i := 0; i < len(IP); i++ {
			IP[i] = pp.Addr[i]
		}
		return
	}
	err = errors.New("Unsupported address type")
	return
}

func (p *Pcap) Inject(data []byte) (err error) {
	buf := (*C.char)(C.malloc((C.size_t)(len(data))))

	for i := 0; i < len(data); i++ {
		*(*byte)(unsafe.Pointer(uintptr(unsafe.Pointer(buf)) + uintptr(i))) = data[i]
	}

	if -1 == C.pcap_inject(p.cptr, unsafe.Pointer(buf), (C.size_t)(len(data))) {
		err = p.Geterror()
	}
	C.free(unsafe.Pointer(buf))
	return
}
<commit_msg>Fixed precision of timestmaps.<commit_after>// Interface to both live and offline pcap parsing.
package pcap

/*
#cgo LDFLAGS: -lpcap
#include <stdlib.h>
#include <pcap.h>

// Workaround for not knowing how to cast to const u_char**
int hack_pcap_next_ex(pcap_t *p, struct pcap_pkthdr **pkt_header,
                      u_char **pkt_data) {
    return pcap_next_ex(p, pkt_header, (const u_char **)pkt_data);
}
*/
import "C"
import (
	"errors"
	"net"
	"syscall"
	"time"
	"unsafe"
)

type Pcap struct {
	cptr *C.pcap_t
}

type Stat struct {
	PacketsReceived  uint32
	PacketsDropped   uint32
	PacketsIfDropped uint32
}

type Interface struct {
	Name        string
	Description string
	Addresses   []IFAddress
	// TODO: add more elements
}

type IFAddress struct {
	IP      net.IP
	Netmask net.IPMask
	// TODO: add broadcast + PtP dst ?
}

func (p *Pcap) Next() (pkt *Packet) {
	rv, _ := p.NextEx()
	return rv
}

// Openlive opens a device and returns a *Pcap handler
func Openlive(device string, snaplen int32, promisc bool, timeout_ms int32) (handle *Pcap, err error) {
	var buf *C.char
	buf = (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	h := new(Pcap)
	var pro int32
	if promisc {
		pro = 1
	}

	dev := C.CString(device)
	defer C.free(unsafe.Pointer(dev))

	h.cptr = C.pcap_open_live(dev, C.int(snaplen), C.int(pro), C.int(timeout_ms), buf)
	if nil == h.cptr {
		handle = nil
		err = errors.New(C.GoString(buf))
	} else {
		handle = h
	}
	C.free(unsafe.Pointer(buf))
	return
}

func Openoffline(file string) (handle *Pcap, err error) {
	var buf *C.char
	buf = (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	h := new(Pcap)

	cf := C.CString(file)
	defer C.free(unsafe.Pointer(cf))

	h.cptr = C.pcap_open_offline(cf, buf)
	if nil == h.cptr {
		handle = nil
		err = errors.New(C.GoString(buf))
	} else {
		handle = h
	}
	C.free(unsafe.Pointer(buf))
	return
}

func (p *Pcap) NextEx() (pkt *Packet, result int32) {
	var pkthdr_ptr *C.struct_pcap_pkthdr
	var pkthdr C.struct_pcap_pkthdr

	var buf_ptr *C.u_char
	var buf unsafe.Pointer
	result = int32(C.hack_pcap_next_ex(p.cptr, &pkthdr_ptr, &buf_ptr))

	buf = unsafe.Pointer(buf_ptr)
	pkthdr = *pkthdr_ptr

	if nil == buf {
		return
	}
	pkt = new(Packet)
	pkt.Time = time.Unix(int64(pkthdr.ts.tv_sec), int64(pkthdr.ts.tv_usec)*1000)
	pkt.Caplen = uint32(pkthdr.caplen)
	pkt.Len = uint32(pkthdr.len)
	pkt.Data = make([]byte, pkthdr.caplen)

	for i := uint32(0); i < pkt.Caplen; i++ {
		pkt.Data[i] = *(*byte)(unsafe.Pointer(uintptr(buf) + uintptr(i)))
	}
	return
}

func (p *Pcap) Close() {
	C.pcap_close(p.cptr)
}

func (p *Pcap) Geterror() error {
	return errors.New(C.GoString(C.pcap_geterr(p.cptr)))
}

func (p *Pcap) Getstats() (stat *Stat, err error) {
	var cstats _Ctype_struct_pcap_stat
	if -1 == C.pcap_stats(p.cptr, &cstats) {
		return nil, p.Geterror()
	}
	stats := new(Stat)
	stats.PacketsReceived = uint32(cstats.ps_recv)
	stats.PacketsDropped = uint32(cstats.ps_drop)
	stats.PacketsIfDropped = uint32(cstats.ps_ifdrop)

	return stats, nil
}

func (p *Pcap) Setfilter(expr string) (err error) {
	var bpf _Ctype_struct_bpf_program
	cexpr := C.CString(expr)
	defer C.free(unsafe.Pointer(cexpr))

	if -1 == C.pcap_compile(p.cptr, &bpf, cexpr, 1, 0) {
		return p.Geterror()
	}

	if -1 == C.pcap_setfilter(p.cptr, &bpf) {
		C.pcap_freecode(&bpf)
		return p.Geterror()
	}
	C.pcap_freecode(&bpf)
	return nil
}

func Version() string {
	return C.GoString(C.pcap_lib_version())
}

func (p *Pcap) Datalink() int {
	return int(C.pcap_datalink(p.cptr))
}

func (p *Pcap) Setdatalink(dlt int) error {
	if -1 == C.pcap_set_datalink(p.cptr, C.int(dlt)) {
		return p.Geterror()
	}
	return nil
}

func DatalinkValueToName(dlt int) string {
	if name := C.pcap_datalink_val_to_name(C.int(dlt)); name != nil {
		return C.GoString(name)
	}
	return ""
}

func DatalinkValueToDescription(dlt int) string {
	if desc := C.pcap_datalink_val_to_description(C.int(dlt)); desc != nil {
		return C.GoString(desc)
	}
	return ""
}

func Findalldevs() (ifs []Interface, err error) {
	var buf *C.char
	buf = (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	defer C.free(unsafe.Pointer(buf))
	var alldevsp *C.pcap_if_t

	if -1 == C.pcap_findalldevs((**C.pcap_if_t)(&alldevsp), buf) {
		return nil, errors.New(C.GoString(buf))
	}
	defer C.pcap_freealldevs((*C.pcap_if_t)(alldevsp))
	dev := alldevsp
	var i uint32
	for i = 0; dev != nil; dev = (*C.pcap_if_t)(dev.next) {
		i++
	}
	ifs = make([]Interface, i)
	dev = alldevsp
	for j := uint32(0); dev != nil; dev = (*C.pcap_if_t)(dev.next) {
		var iface Interface
		iface.Name = C.GoString(dev.name)
		iface.Description = C.GoString(dev.description)
		iface.Addresses = findalladdresses(dev.addresses)
		// TODO: add more elements
		ifs[j] = iface
		j++
	}
	return
}

func findalladdresses(addresses *_Ctype_struct_pcap_addr) (retval []IFAddress) {
	// TODO - make it support more than IPv4 and IPv6?
	retval = make([]IFAddress, 0, 1)
	for curaddr := addresses; curaddr != nil; curaddr = (*_Ctype_struct_pcap_addr)(curaddr.next) {
		var a IFAddress
		var err error
		if a.IP, err = sockaddr_to_IP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr))); err != nil {
			continue
		}
		if a.Netmask, err = sockaddr_to_IP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr))); err != nil {
			continue
		}
		retval = append(retval, a)
	}
	return
}

func sockaddr_to_IP(rsa *syscall.RawSockaddr) (IP []byte, err error) {
	switch rsa.Family {
	case syscall.AF_INET:
		pp := (*syscall.RawSockaddrInet4)(unsafe.Pointer(rsa))
		IP = make([]byte, 4)
		for i := 0; i < len(IP); i++ {
			IP[i] = pp.Addr[i]
		}
		return
	case syscall.AF_INET6:
		pp := (*syscall.RawSockaddrInet6)(unsafe.Pointer(rsa))
		IP = make([]byte, 16)
		for i := 0; i < len(IP); i++ {
			IP[i] = pp.Addr[i]
		}
		return
	}
	err = errors.New("Unsupported address type")
	return
}

func (p *Pcap) Inject(data []byte) (err error) {
	buf := (*C.char)(C.malloc((C.size_t)(len(data))))

	for i := 0; i < len(data); i++ {
		*(*byte)(unsafe.Pointer(uintptr(unsafe.Pointer(buf)) + uintptr(i))) = data[i]
	}

	if -1 == C.pcap_inject(p.cptr, unsafe.Pointer(buf), (C.size_t)(len(data))) {
		err = p.Geterror()
	}
	C.free(unsafe.Pointer(buf))
	return
}
<|endoftext|>


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

package main

import (
	"fmt"
	"launchpad.net/tomb"
	"log"
	"net"
	"sort"
)

// PeerTuple represents a single IP+port pair of a peer
type PeerTuple struct {
	IP   net.IP
	Port uint16
}

type Peer struct {
	peer PeerTuple
}

type PeerManager struct {
	peersCh <-chan PeerTuple
	statsCh chan Stats
	connsCh <-chan net.Conn
	peers	map[string]PeerInfo
	t       tomb.Tomb
}

type PeerInfo struct {
	peerId          string
	isActive        bool // The peer is connected and unchoked
	availablePieces []bool
	activeRequests  map[int]struct{}
	qtyPiecesNeeded int                 // The quantity of pieces that this peer has that we haven't yet downloaded.
	requestPieceCh  chan<- RequestPiece // Other end is Peer. Used to tell the peer to request a particular piece.
	cancelPieceCh   chan<- CancelPiece  // Other end is Peer. Used to tell the peer to cancel a particular piece.
	havePieceCh		chan<- HavePiece 	// Other end is Peer. Used to tell the peer that we have a new piece.
}

type SortedPeers []PeerInfo

func (sp SortedPeers) Less(i, j int) bool {
	return sp[i].qtyPiecesNeeded <= sp[j].qtyPiecesNeeded
}

func (sp SortedPeers) Swap(i, j int) {
	tmp := sp[i]
	sp[i] = sp[j]
	sp[j] = tmp
}

func (sp SortedPeers) Len() int {
	return len(sp)
}

func sortedPeersByQtyPiecesNeeded(peers map[string]PeerInfo) SortedPeers {
	peerInfoSlice := make(SortedPeers, 0)

	for _, peerInfo := range peers {
		peerInfoSlice = append(peerInfoSlice, peerInfo)
	}
	sort.Sort(peerInfoSlice)

	return peerInfoSlice
}

func NewPeerManager(peersCh chan PeerTuple, statsCh chan Stats, connsCh chan net.Conn) *PeerManager {
	pm := new(PeerManager)
	pm.peersCh = peersCh
	pm.statsCh = statsCh
	pm.connsCh = connsCh
	return pm
}

func NewPeerInfo(quantityOfPieces int) *PeerInfo {
	pi := new(PeerInfo)
	pi.availablePieces = make([]bool, quantityOfPieces)
	pi.activeRequests = make(map[int]struct{})

	// FIXME Not finished. Need to hook these channels into the Peer struct
	pi.requestPieceCh = make(chan<- RequestPiece)
	pi.cancelPieceCh = make(chan<- CancelPiece)
	return pi
}

/*
func NewPeer() {
}

func NewPeerConn(net.Conn) {
}

func NewPeerTuple() {
}
*/

func (pm *PeerManager) Stop() error {
	log.Println("PeerManager : Stop : Stopping")
	pm.t.Kill(nil)
	return pm.t.Wait()
}

func (pm *PeerManager) Run() {
	log.Println("PeerManager : Run : Started")
	defer pm.t.Done()
	defer log.Println("PeerManager : Run : Completed")

	for {
		select {
		case peer := <-pm.peersCh:
			_, ok := peers[peer]
			if ok {
				// peer already exists
				log.Printf("Peer %s:%d already in map\n", peer.IP.String, peer.Port)
			} else {
				peers[peer] = "foo"
			}
			fmt.Println(peer)
		case conn := <-pm.connsCh:
			fmt.Println(conn)
		case <-pm.t.Dying():
			return
		}
	}
}
<commit_msg>Start writing code to connect to other peers<commit_after>// Copyright 2013 Jari Takkala. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"fmt"
	"launchpad.net/tomb"
	"log"
	"net"
	"sort"
//	"syscall"
)

// PeerTuple represents a single IP+port pair of a peer
type PeerTuple struct {
	IP   net.IP
	Port uint16
}

type Peer struct {
	peer PeerTuple
	conn net.Conn
}

type PeerManager struct {
	peersCh <-chan PeerTuple
	statsCh chan Stats
	connsCh <-chan net.Conn
	peers	map[string]*Peer
	t       tomb.Tomb
}

type PeerInfo struct {
	peerId          string
	isActive        bool // The peer is connected and unchoked
	availablePieces []bool
	activeRequests  map[int]struct{}
	qtyPiecesNeeded int                 // The quantity of pieces that this peer has that we haven't yet downloaded.
	requestPieceCh  chan<- RequestPiece // Other end is Peer. Used to tell the peer to request a particular piece.
	cancelPieceCh   chan<- CancelPiece  // Other end is Peer. Used to tell the peer to cancel a particular piece.
	havePieceCh		chan<- HavePiece 	// Other end is Peer. Used to tell the peer that we have a new piece.
}

type SortedPeers []PeerInfo

func (sp SortedPeers) Less(i, j int) bool {
	return sp[i].qtyPiecesNeeded <= sp[j].qtyPiecesNeeded
}

func (sp SortedPeers) Swap(i, j int) {
	tmp := sp[i]
	sp[i] = sp[j]
	sp[j] = tmp
}

func (sp SortedPeers) Len() int {
	return len(sp)
}

func sortedPeersByQtyPiecesNeeded(peers map[string]PeerInfo) SortedPeers {
	peerInfoSlice := make(SortedPeers, 0)

	for _, peerInfo := range peers {
		peerInfoSlice = append(peerInfoSlice, peerInfo)
	}
	sort.Sort(peerInfoSlice)

	return peerInfoSlice
}

func NewPeerManager(peersCh chan PeerTuple, statsCh chan Stats, connsCh chan net.Conn) *PeerManager {
	pm := new(PeerManager)
	pm.peersCh = peersCh
	pm.statsCh = statsCh
	pm.connsCh = connsCh
	pm.peers = make(map[string]*Peer)
	return pm
}

func NewPeerInfo(quantityOfPieces int) *PeerInfo {
	pi := new(PeerInfo)
	pi.availablePieces = make([]bool, quantityOfPieces)
	pi.activeRequests = make(map[int]struct{})

	// FIXME Not finished. Need to hook these channels into the Peer struct
	pi.requestPieceCh = make(chan<- RequestPiece)
	pi.cancelPieceCh = make(chan<- CancelPiece)
	return pi
}

func NewPeer(peerTuple PeerTuple) *Peer {
	peer := new(Peer)
	raddr := net.TCPAddr { peerTuple.IP, int(peerTuple.Port), "" }
	/*
	go func() {
	conn, err := net.DialTCP("tcp4", nil, &raddr)
	if err != nil {
		if e, ok := err.(*net.OpError); ok {
			if e.Err == syscall.ECONNREFUSED {
				fmt.Printf("%#q\n", e)
				return
			}
		}
		log.Fatal(err)
	}
	fmt.Println(conn)
	}()
	*/
	fmt.Println(raddr)

	return peer
}

/*
func NewPeerConn(net.Conn) {
}

func NewPeerTuple() {
}
*/

func (pm *PeerManager) Stop() error {
	log.Println("PeerManager : Stop : Stopping")
	pm.t.Kill(nil)
	return pm.t.Wait()
}

func (pm *PeerManager) Run() {
	log.Println("PeerManager : Run : Started")
	defer pm.t.Done()
	defer log.Println("PeerManager : Run : Completed")

	for {
		select {
		case peer := <-pm.peersCh:
			peerID := fmt.Sprintf("%s:%d", peer.IP.String(), peer.Port)
			_, ok := pm.peers[peerID]
			if ok {
				// Peer already exists
				log.Printf("Peer %s already in map\n", peerID)
			} else {
				pm.peers[peerID] = NewPeer(peer)
			}
		case conn := <-pm.connsCh:
			fmt.Println(conn)
		case <-pm.t.Dying():
			return
		}
	}
}
<|endoftext|>


<commit_before>package gontpd

import (
	"crypto/md5"
	"log"
	"net"
	"time"

	"github.com/beevik/ntp"
)

const (
	replyNum       = 4
	goodFilter     = 2
	invalidStratum = 16
	maxPoll        = 16
	minPoll        = 5
)

type peer struct {
	addr       net.IP
	reply      [replyNum]*ntp.Response
	offset     time.Duration
	delay      time.Duration
	err        time.Duration
	refId      uint32
	stratum    uint8
	trustLevel uint8
	good       bool
	enable     bool
}

func newPeer(addr net.IP) (p *peer) {
	log.Printf("new peer:%s", addr.String())
	p = &peer{
		addr:       addr,
		trustLevel: minPoll,
		enable:     true,
	}
	p.refId = makeSendRefId(addr)
	return
}

func (p *peer) update() {

	goodCount := 0
	p.good = false
	ts := 2 * time.Second

	for i := 0; i < replyNum; i++ {
		time.Sleep(ts)
		resp, err := ntp.Query(p.addr.String())

		if resp != nil {
			switch resp.KissCode {
			case "RATE":
				ts += time.Second
			case "DENY":
				p.enable = false
				return
			}
		}
		if err != nil {
			log.Printf("%s update failed %s", p.addr.String(), err)
			p.reply[i] = &ntp.Response{Stratum: invalidStratum}
			continue
		}

		goodCount += 1
		p.reply[i] = resp
	}

	p.good = goodCount > goodFilter

	if debug {
		log.Printf("%s is good=%v", p.addr, p.good)
	}

	if p.good {
		if p.trustLevel < maxPoll {
			p.trustLevel += 1
		}
	} else {
		if p.trustLevel > minPoll {
			p.trustLevel -= 1
		}
	}
}

func makeSendRefId(ip net.IP) (id uint32) {

	if len(ip) > 10 && ip[11] == 255 {
		// ipv4
		id = uint32(ip[12])<<24 + uint32(ip[13])<<16 + uint32(ip[14])<<
			8 + uint32(ip[15])
	} else {
		h := md5.New()
		hr := h.Sum(ip)
		// 255.b2.b3.b4 for ipv6 hash
		// https://support.ntp.org/bin/view/Dev/UpdatingTheRefidFormat
		id = uint32(255)<<24 + uint32(hr[1])<<16 + uint32(hr[2])<<8 + uint32(hr[3])
	}
	return
}
<commit_msg>disable peer while can't be reach<commit_after>package gontpd

import (
	"crypto/md5"
	"log"
	"net"
	"time"

	"github.com/beevik/ntp"
)

const (
	replyNum       = 4
	goodFilter     = 2
	invalidStratum = 16
	maxPoll        = 16
	minPoll        = 5
)

type peer struct {
	addr       net.IP
	reply      [replyNum]*ntp.Response
	offset     time.Duration
	delay      time.Duration
	err        time.Duration
	refId      uint32
	stratum    uint8
	trustLevel uint8
	good       bool
	enable     bool
}

func newPeer(addr net.IP) (p *peer) {
	log.Printf("new peer:%s", addr.String())
	p = &peer{
		addr:       addr,
		trustLevel: minPoll,
		enable:     true,
	}
	p.refId = makeSendRefId(addr)
	return
}

func (p *peer) update() {

	goodCount := 0
	p.good = false
	ts := 2 * time.Second

	for i := 0; i < replyNum; i++ {
		time.Sleep(ts)
		resp, err := ntp.Query(p.addr.String())
		if resp != nil && resp.Stratum == 0 {
			switch resp.KissCode {
			case "RATE":
				ts += time.Second
			case "DENY":
				p.enable = false
				return
			}
		}

		if err != nil {
			log.Printf("%s update failed %s", p.addr.String(), err)
			p.reply[i] = &ntp.Response{Stratum: invalidStratum}
			if nerr, ok := err.(net.Error); ok {
				if !nerr.Temporary() {
					log.Printf("%s can't be reach, disabled", p.addr.String())
					p.enable = false
				}
			}
			continue
		}

		goodCount += 1
		p.reply[i] = resp
	}

	p.good = goodCount > goodFilter

	if debug {
		log.Printf("%s is good=%v", p.addr, p.good)
	}

	if p.good {
		if p.trustLevel < maxPoll {
			p.trustLevel += 1
		}
	} else {
		if p.trustLevel > minPoll {
			p.trustLevel -= 1
		}
	}
}

func makeSendRefId(ip net.IP) (id uint32) {

	if len(ip) > 10 && ip[11] == 255 {
		// ipv4
		id = uint32(ip[12])<<24 + uint32(ip[13])<<16 + uint32(ip[14])<<
			8 + uint32(ip[15])
	} else {
		h := md5.New()
		hr := h.Sum(ip)
		// 255.b2.b3.b4 for ipv6 hash
		// https://support.ntp.org/bin/view/Dev/UpdatingTheRefidFormat
		id = uint32(255)<<24 + uint32(hr[1])<<16 + uint32(hr[2])<<8 + uint32(hr[3])
	}
	return
}
<|endoftext|>


<commit_before>package raft

import (
	"errors"
	"sync"
	"time"
)

//------------------------------------------------------------------------------
//
// Typedefs
//
//------------------------------------------------------------------------------

// A peer is a reference to another server involved in the consensus protocol.
type Peer struct {
	server         *Server
	name           string
	prevLogIndex   uint64
	mutex          sync.Mutex
	heartbeatTimer *Timer
}

type FlushResponse struct {
	term    uint64
	success bool
	err     error
	peer    *Peer
}

//------------------------------------------------------------------------------
//
// Constructor
//
//------------------------------------------------------------------------------

// Creates a new peer.
func NewPeer(server *Server, name string, heartbeatTimeout time.Duration) *Peer {
	p := &Peer{
		server:         server,
		name:           name,
		heartbeatTimer: NewTimer(heartbeatTimeout, heartbeatTimeout),
	}

	return p
}

//------------------------------------------------------------------------------
//
// Accessors
//
//------------------------------------------------------------------------------

// Retrieves the name of the peer.
func (p *Peer) Name() string {
	return p.name
}

// Retrieves the heartbeat timeout.
func (p *Peer) HeartbeatTimeout() time.Duration {
	return p.heartbeatTimer.MinDuration()
}

// Sets the heartbeat timeout.
func (p *Peer) SetHeartbeatTimeout(duration time.Duration) {
	p.heartbeatTimer.SetDuration(duration)
}

func (p *Peer) StartHeartbeat() {
	go p.heartbeat()
}

//------------------------------------------------------------------------------
//
// Methods
//
//------------------------------------------------------------------------------

//--------------------------------------
// State
//--------------------------------------

// Stops the peer entirely.
func (p *Peer) stop() {
	p.mutex.Lock()
	defer p.mutex.Unlock()
	p.heartbeatTimer.Stop()
}

//--------------------------------------
// Copying
//--------------------------------------

// Clones the state of the peer. The clone is not attached to a server and
// the heartbeat timer will not exist.
func (p *Peer) clone() *Peer {
	p.mutex.Lock()
	defer p.mutex.Unlock()
	return &Peer{
		name:         p.name,
		prevLogIndex: p.prevLogIndex,
	}
}

//--------------------------------------
// Flush
//--------------------------------------

// Sends an AppendEntries RPC but does not obtain a lock
// on the server.
func (p *Peer) flush() (uint64, bool, error) {
	// We need to hold the log lock to create AppendEntriesRequest
	// avoid snapshot to delete the desired entries before AEQ()
	req := p.server.createAppendEntriesRequest(p.prevLogIndex)

	if req != nil {
		return p.sendFlushRequest(req)
	} else {
		req := p.server.createSnapshotRequest()
		return p.sendSnapshotRequest(req)
	}

}

// send VoteRequest Request
func (p *Peer) sendVoteRequest(req *RequestVoteRequest, c chan *RequestVoteResponse){
	req.peer = p
	debugln(p.server.Name(), "Send Vote Request to ", p.Name())
	if resp, _ := p.server.transporter.SendVoteRequest(p.server, p, req); resp != nil {
		resp.peer = p
		c <- resp
	}
}

// send Snapshot Request
func (p *Peer) sendSnapshotRequest(req *SnapshotRequest) (uint64, bool, error) {
	// Ignore any null requests.
	if req == nil {
		return 0, false, errors.New("raft.Peer: Request required")
	}

	// Generate an snapshot request based on the state of the server and
	// log. Send the request through the user-provided handler and process the
	// result.
	resp, err := p.server.transporter.SendSnapshotRequest(p.server, p, req)

	if resp == nil {
		return 0, false, err
	}

	// If successful then update the previous log index. If it was
	// unsuccessful then decrement the previous log index and we'll try again
	// next time.
	if resp.Success {
		p.prevLogIndex = req.LastIndex

	} else {
		panic(resp)
	}

	return resp.Term, resp.Success, err
}

// Flushes a request through the server's transport.
func (p *Peer) sendFlushRequest(req *AppendEntriesRequest) (uint64, bool, error) {
	// Ignore any null requests.
	if req == nil {
		return 0, false, errors.New("raft.Peer: Request required")
	}

	// Generate an AppendEntries request based on the state of the server and
	// log. Send the request through the user-provided handler and process the
	// result.
	//debugln("flush to ", p.Name())
	debugln("[HeartBeat] Leader ", p.server.Name(), " to ", p.Name(), " ", len(req.Entries), " ", time.Now())

	if p.server.State() != Leader {
		return 0, false, errors.New("Not leader anymore")
	}

	resp, err := p.server.transporter.SendAppendEntriesRequest(p.server, p, req)

	//debugln("receive flush response from ", p.Name())

	if resp == nil {
		return 0, false, err
	}

	// If successful then update the previous log index. If it was
	// unsuccessful then decrement the previous log index and we'll try again
	// next time.
	if resp.Success {
		if len(req.Entries) > 0 {
			p.prevLogIndex = req.Entries[len(req.Entries)-1].Index
			debugln("Peer ", p.Name(), "'s' log update to ", p.prevLogIndex)
		}
	} else {

		if p.server.State() != Leader {
			return 0, false, errors.New("Not leader anymore")
		}

		if resp.Term > p.server.currentTerm {
			return resp.Term, false, errors.New("Step down")
		}
		// Decrement the previous log index down until we find a match. Don't
		// let it go below where the peer's commit index is though. That's a
		// problem.
		if p.prevLogIndex > 0 {
			p.prevLogIndex--
		}
		if resp.CommitIndex > p.prevLogIndex {
			debugln("%v %v %v %v", resp.CommitIndex, p.prevLogIndex,
				p.server.currentTerm, resp.Term)
			panic("commitedIndex is greater than prevLogIndex")
		}
	}

	return resp.Term, resp.Success, err
}

//--------------------------------------
// Heartbeat
//--------------------------------------

// Listens to the heartbeat timeout and flushes an AppendEntries RPC.
func (p *Peer) heartbeat() {
	for {

		// (1) timeout/fire happens, flush the peer
		// (2) stopped, return

		if p.heartbeatTimer.Start() {

			var f FlushResponse

			f.peer = p

			f.term, f.success, f.err = p.flush()

			// if the peer successfully appended the log entry
			// we will tell the commit center
			if f.success {
				if p.prevLogIndex > p.server.log.CommitIndex() {
					debugln("[Heartbeat] Peer", p.Name(), "send to commit center")
					p.server.response <- f
					debugln("[Heartbeat] Peer", p.Name(), "back from commit center")
				}

			} else {
				// shutdown the heartbeat
				if f.term > p.server.currentTerm {
					debugln("[Heartbeat] SetpDown!")
					select {
					case p.server.stepDown <- f.term:
						return
					default:
						return
					}
				}
			}

		} else {
			// shutdown
			return
		}
	}
}
<commit_msg>timeout for transfere. This helps when the node is far left behind. The server will not wait for sending a large package before send the other heartbeat<commit_after>package raft

import (
	"errors"
	"sync"
	"time"
)

//------------------------------------------------------------------------------
//
// Typedefs
//
//------------------------------------------------------------------------------

// A peer is a reference to another server involved in the consensus protocol.
type Peer struct {
	server         *Server
	name           string
	prevLogIndex   uint64
	mutex          sync.Mutex
	heartbeatTimer *Timer
}

type FlushResponse struct {
	term    uint64
	success bool
	err     error
	peer    *Peer
}

//------------------------------------------------------------------------------
//
// Constructor
//
//------------------------------------------------------------------------------

// Creates a new peer.
func NewPeer(server *Server, name string, heartbeatTimeout time.Duration) *Peer {
	p := &Peer{
		server:         server,
		name:           name,
		heartbeatTimer: NewTimer(heartbeatTimeout, heartbeatTimeout),
	}

	return p
}

//------------------------------------------------------------------------------
//
// Accessors
//
//------------------------------------------------------------------------------

// Retrieves the name of the peer.
func (p *Peer) Name() string {
	return p.name
}

// Retrieves the heartbeat timeout.
func (p *Peer) HeartbeatTimeout() time.Duration {
	return p.heartbeatTimer.MinDuration()
}

// Sets the heartbeat timeout.
func (p *Peer) SetHeartbeatTimeout(duration time.Duration) {
	p.heartbeatTimer.SetDuration(duration)
}

func (p *Peer) StartHeartbeat() {
	go p.heartbeat()
}

//------------------------------------------------------------------------------
//
// Methods
//
//------------------------------------------------------------------------------

//--------------------------------------
// State
//--------------------------------------

// Stops the peer entirely.
func (p *Peer) stop() {
	p.mutex.Lock()
	defer p.mutex.Unlock()
	p.heartbeatTimer.Stop()
}

//--------------------------------------
// Copying
//--------------------------------------

// Clones the state of the peer. The clone is not attached to a server and
// the heartbeat timer will not exist.
func (p *Peer) clone() *Peer {
	p.mutex.Lock()
	defer p.mutex.Unlock()
	return &Peer{
		name:         p.name,
		prevLogIndex: p.prevLogIndex,
	}
}

//--------------------------------------
// Flush
//--------------------------------------

// Sends an AppendEntries RPC but does not obtain a lock
// on the server.
func (p *Peer) flush() (uint64, bool, error) {
	// We need to hold the log lock to create AppendEntriesRequest
	// avoid snapshot to delete the desired entries before AEQ()
	req := p.server.createAppendEntriesRequest(p.prevLogIndex)

	if req != nil {
		return p.sendFlushRequest(req)
	} else {
		req := p.server.createSnapshotRequest()
		return p.sendSnapshotRequest(req)
	}

}

// send VoteRequest Request
func (p *Peer) sendVoteRequest(req *RequestVoteRequest, c chan *RequestVoteResponse) {
	req.peer = p
	debugln(p.server.Name(), "Send Vote Request to ", p.Name())
	if resp, _ := p.server.transporter.SendVoteRequest(p.server, p, req); resp != nil {
		resp.peer = p
		c <- resp
	}
}

// send Snapshot Request
func (p *Peer) sendSnapshotRequest(req *SnapshotRequest) (uint64, bool, error) {
	// Ignore any null requests.
	if req == nil {
		return 0, false, errors.New("raft.Peer: Request required")
	}

	// Generate an snapshot request based on the state of the server and
	// log. Send the request through the user-provided handler and process the
	// result.
	resp, err := p.server.transporter.SendSnapshotRequest(p.server, p, req)

	if resp == nil {
		return 0, false, err
	}

	// If successful then update the previous log index. If it was
	// unsuccessful then decrement the previous log index and we'll try again
	// next time.
	if resp.Success {
		p.prevLogIndex = req.LastIndex

	} else {
		panic(resp)
	}

	return resp.Term, resp.Success, err
}

// Flushes a request through the server's transport.
func (p *Peer) sendFlushRequest(req *AppendEntriesRequest) (uint64, bool, error) {
	// Ignore any null requests.
	if req == nil {
		return 0, false, errors.New("raft.Peer: Request required")
	}

	// Generate an AppendEntries request based on the state of the server and
	// log. Send the request through the user-provided handler and process the
	// result.
	//debugln("flush to ", p.Name())
	debugln("[HeartBeat] Leader ", p.server.Name(), " to ", p.Name(), " ", len(req.Entries), " ", time.Now())

	respChan := make(chan *AppendEntriesResponse, 2)

	go func() {
		tranResp, _ := p.server.transporter.SendAppendEntriesRequest(p.server, p, req)
		respChan <- tranResp
	}()

	go func() {
		<-time.After(p.server.heartbeatTimeout)
		respChan <- nil
	}()

	resp := <-respChan

	debugln("receive flush response from ", p.Name())

	if resp == nil {
		return 0, false, errors.New("Network problem")
	}

	// If successful then update the previous log index. If it was
	// unsuccessful then decrement the previous log index and we'll try again
	// next time.
	if resp.Success {
		if len(req.Entries) > 0 {
			p.prevLogIndex = req.Entries[len(req.Entries)-1].Index
			debugln("Peer ", p.Name(), "'s' log update to ", p.prevLogIndex)
		}
	} else {

		if resp.Term > p.server.currentTerm {
			return resp.Term, false, errors.New("Step down")
		}

		// we may miss a response from peer
		if resp.CommitIndex > p.prevLogIndex {
			p.prevLogIndex = resp.CommitIndex
		} else if p.prevLogIndex > 0 {
			// Decrement the previous log index down until we find a match. Don't
			// let it go below where the peer's commit index is though. That's a
			// problem.
			p.prevLogIndex--
		}

	}
	return resp.Term, resp.Success, nil
}

//--------------------------------------
// Heartbeat
//--------------------------------------

// Listens to the heartbeat timeout and flushes an AppendEntries RPC.
func (p *Peer) heartbeat() {
	for {

		// (1) timeout/fire happens, flush the peer
		// (2) stopped, return

		if p.heartbeatTimer.Start() {

			var f FlushResponse

			f.peer = p

			f.term, f.success, f.err = p.flush()

			// if the peer successfully appended the log entry
			// we will tell the commit center
			if f.success {
				if p.prevLogIndex > p.server.log.CommitIndex() {
					debugln("[Heartbeat] Peer", p.Name(), "send to commit center")
					p.server.response <- f
					debugln("[Heartbeat] Peer", p.Name(), "back from commit center")
				}

			} else {
				// shutdown the heartbeat
				if f.term > p.server.currentTerm {
					debugln("[Heartbeat] SetpDown!")
					select {
					case p.server.stepDown <- f.term:
						return
					default:
						return
					}
				}
			}

		} else {
			// shutdown
			return
		}
	}
}
<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"strings"

	"github.com/jbenet/go-ipfs/commands"
)

// Parse parses the input commandline string (cmd, flags, and args).
// returns the corresponding command Request object.
func Parse(input []string, root *commands.Command) (commands.Request, error) {
	path, input := parsePath(input, root)
	opts, args, err := parseOptions(input)
	if err != nil {
		return nil, err
	}

	return commands.NewRequest(path, opts, args), nil
}

// parsePath gets the command path from the command line input
func parsePath(input []string, root *commands.Command) ([]string, []string) {
	cmd := root
	i := 0

	for _, blob := range input {
		if strings.HasPrefix(blob, "-") {
			break
		}

		cmd := cmd.Subcommand(blob)
		if cmd == nil {
			break
		}

		i++
	}

	return input[:i], input[i:]
}

// parseOptions parses the raw string values of the given options
// returns the parsed options as strings, along with the CLI args
func parseOptions(input []string) (map[string]interface{}, []string, error) {
	opts := make(map[string]interface{})
	args := []string{}

	for i := 0; i < len(input); i++ {
		blob := input[i]

		if strings.HasPrefix(blob, "-") {
			name := blob[1:]
			value := ""

			// support single and double dash
			if strings.HasPrefix(name, "-") {
				name = name[1:]
			}

			if strings.Contains(name, "=") {
				split := strings.SplitN(name, "=", 2)
				name = split[0]
				value = split[1]
			}

			if _, ok := opts[name]; ok {
				return nil, nil, fmt.Errorf("Duplicate values for option '%s'", name)
			}

			opts[name] = value

		} else {
			args = append(args, blob)
		}
	}

	return opts, args, nil
}
<commit_msg>commands/cli: Made Parse return component fields instead of a Request<commit_after>package cli

import (
	"fmt"
	"strings"

	"github.com/jbenet/go-ipfs/commands"
)

// Parse parses the input commandline string (cmd, flags, and args).
// returns the corresponding command Request object.
func Parse(input []string, root *commands.Command) ([]string, map[string]interface{}, []string, error) {
	path, input := parsePath(input, root)
	opts, args, err := parseOptions(input)
	if err != nil {
		return nil, nil, nil, err
	}

	return path, opts, args, nil
}

// parsePath gets the command path from the command line input
func parsePath(input []string, root *commands.Command) ([]string, []string) {
	cmd := root
	i := 0

	for _, blob := range input {
		if strings.HasPrefix(blob, "-") {
			break
		}

		cmd := cmd.Subcommand(blob)
		if cmd == nil {
			break
		}

		i++
	}

	return input[:i], input[i:]
}

// parseOptions parses the raw string values of the given options
// returns the parsed options as strings, along with the CLI args
func parseOptions(input []string) (map[string]interface{}, []string, error) {
	opts := make(map[string]interface{})
	args := []string{}

	for i := 0; i < len(input); i++ {
		blob := input[i]

		if strings.HasPrefix(blob, "-") {
			name := blob[1:]
			value := ""

			// support single and double dash
			if strings.HasPrefix(name, "-") {
				name = name[1:]
			}

			if strings.Contains(name, "=") {
				split := strings.SplitN(name, "=", 2)
				name = split[0]
				value = split[1]
			}

			if _, ok := opts[name]; ok {
				return nil, nil, fmt.Errorf("Duplicate values for option '%s'", name)
			}

			opts[name] = value

		} else {
			args = append(args, blob)
		}
	}

	return opts, args, nil
}
<|endoftext|>


<commit_before>//  Copyright (c) 2015 Rackspace
//
//  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 conf

import (
	"fmt"
	"strconv"
	"strings"
)

type Policy struct {
	Index      int
	Type       string
	Name       string
	Aliases    []string
	Default    bool
	Deprecated bool
	Config     map[string]string
}

func (p Policy) GetDbPartPower() (uint, error) {
	dbPartPower := 0
	if p.Config["db_part_power"] != "" {
		dbPartPowerInt64, err := strconv.ParseInt(p.Config["db_part_power"], 10, 64)
		if err != nil {
			return 0, fmt.Errorf("Could not parse db_part_power value %q: %s", p.Config["db_part_power"], err)
		}
		dbPartPower = int(dbPartPowerInt64)
	}
	if dbPartPower < 1 {
		dbPartPower = 5
	}
	return uint(dbPartPower), nil
}
func (p Policy) GetDbSubDirs() (int, error) {
	subdirs := 0
	if p.Config["subdirs"] != "" {
		subdirsInt64, err := strconv.ParseInt(p.Config["subdirs"], 10, 64)
		if err != nil {
			return 0, fmt.Errorf("Could not parse subdirs value %q: %s", p.Config["subdirs"], err)
		}
		subdirs = int(subdirsInt64)
	}
	if subdirs < 1 {
		subdirs = 32
	}
	return subdirs, nil
}

type PolicyList map[int]*Policy

func (p PolicyList) Default() int {
	for _, v := range p {
		if v.Default {
			return v.Index
		}
	}
	return 0
}

func (p PolicyList) NameLookup(s string) *Policy {
	// a map would be faster, but this only happens on container PUT, so how fast does it need to be?
	s = strings.ToUpper(strings.TrimSpace(s))
	for _, v := range p {
		if s == strings.ToUpper(v.Name) {
			return v
		}
		for _, a := range v.Aliases {
			if s == strings.ToUpper(a) {
				return v
			}
		}
	}
	return nil
}

func (p PolicyList) GetPolicyInfo() []map[string]interface{} {
	policyInfo := []map[string]interface{}{}
	for _, v := range p {
		if v.Deprecated {
			continue
		}
		pol := map[string]interface{}{}
		pol["name"] = v.Name
		if v.Default {
			pol["default"] = v.Default
		}
		pol["aliases"] = strings.Join(v.Aliases, ", ")
		policyInfo = append(policyInfo, pol)
	}
	return policyInfo
}

// LoadPolicies loads policies, probably from /etc/swift/swift.conf
func GetPolicies() (PolicyList, error) {
	policies := map[int]*Policy{0: {
		Index:      0,
		Type:       "replication",
		Name:       "Policy-0",
		Aliases:    nil,
		Default:    false,
		Deprecated: false,
	}}
	for _, loc := range configLocations {
		if conf, e := LoadConfig(loc); e == nil {
			for key := range conf.File {
				var policyIndex int
				if c, err := fmt.Sscanf(key, "storage-policy:%d", &policyIndex); err == nil && c == 1 {
					name := conf.GetDefault(key, "name", fmt.Sprintf("Policy-%d", policyIndex))
					aliases := []string{name}
					aliasList := conf.GetDefault(key, "aliases", "")
					for _, alias := range strings.Split(aliasList, ",") {
						alias = strings.Trim(alias, " ")
						if alias != "" && alias != name {
							aliases = append(aliases, alias)
						}
					}
					policies[policyIndex] = &Policy{
						Index:      policyIndex,
						Type:       conf.GetDefault(key, "policy_type", "replication"),
						Name:       name,
						Aliases:    aliases,
						Deprecated: conf.GetBool(key, "deprecated", false),
						Default:    conf.GetBool(key, "default", false),
						Config:     map[string]string(conf.File[key]),
					}
				}
			}
			break
		}
	}
	defaultFound := false
	for _, policy := range policies {
		if policy.Default {
			defaultFound = true
		}
	}
	if !defaultFound {
		policies[0].Default = true
	}
	return PolicyList(policies), nil
}
<commit_msg>HUM-858 : lower default db_part_power<commit_after>//  Copyright (c) 2015 Rackspace
//
//  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 conf

import (
	"fmt"
	"strconv"
	"strings"
)

type Policy struct {
	Index      int
	Type       string
	Name       string
	Aliases    []string
	Default    bool
	Deprecated bool
	Config     map[string]string
}

func (p Policy) GetDbPartPower() (uint, error) {
	dbPartPower := 0
	if p.Config["db_part_power"] != "" {
		dbPartPowerInt64, err := strconv.ParseInt(p.Config["db_part_power"], 10, 64)
		if err != nil {
			return 0, fmt.Errorf("Could not parse db_part_power value %q: %s", p.Config["db_part_power"], err)
		}
		dbPartPower = int(dbPartPowerInt64)
	}
	if dbPartPower < 1 {
		dbPartPower = 3
	}
	return uint(dbPartPower), nil
}
func (p Policy) GetDbSubDirs() (int, error) {
	subdirs := 0
	if p.Config["subdirs"] != "" {
		subdirsInt64, err := strconv.ParseInt(p.Config["subdirs"], 10, 64)
		if err != nil {
			return 0, fmt.Errorf("Could not parse subdirs value %q: %s", p.Config["subdirs"], err)
		}
		subdirs = int(subdirsInt64)
	}
	if subdirs < 1 {
		subdirs = 32
	}
	return subdirs, nil
}

type PolicyList map[int]*Policy

func (p PolicyList) Default() int {
	for _, v := range p {
		if v.Default {
			return v.Index
		}
	}
	return 0
}

func (p PolicyList) NameLookup(s string) *Policy {
	// a map would be faster, but this only happens on container PUT, so how fast does it need to be?
	s = strings.ToUpper(strings.TrimSpace(s))
	for _, v := range p {
		if s == strings.ToUpper(v.Name) {
			return v
		}
		for _, a := range v.Aliases {
			if s == strings.ToUpper(a) {
				return v
			}
		}
	}
	return nil
}

func (p PolicyList) GetPolicyInfo() []map[string]interface{} {
	policyInfo := []map[string]interface{}{}
	for _, v := range p {
		if v.Deprecated {
			continue
		}
		pol := map[string]interface{}{}
		pol["name"] = v.Name
		if v.Default {
			pol["default"] = v.Default
		}
		pol["aliases"] = strings.Join(v.Aliases, ", ")
		policyInfo = append(policyInfo, pol)
	}
	return policyInfo
}

// LoadPolicies loads policies, probably from /etc/swift/swift.conf
func GetPolicies() (PolicyList, error) {
	policies := map[int]*Policy{0: {
		Index:      0,
		Type:       "replication",
		Name:       "Policy-0",
		Aliases:    nil,
		Default:    false,
		Deprecated: false,
	}}
	for _, loc := range configLocations {
		if conf, e := LoadConfig(loc); e == nil {
			for key := range conf.File {
				var policyIndex int
				if c, err := fmt.Sscanf(key, "storage-policy:%d", &policyIndex); err == nil && c == 1 {
					name := conf.GetDefault(key, "name", fmt.Sprintf("Policy-%d", policyIndex))
					aliases := []string{name}
					aliasList := conf.GetDefault(key, "aliases", "")
					for _, alias := range strings.Split(aliasList, ",") {
						alias = strings.Trim(alias, " ")
						if alias != "" && alias != name {
							aliases = append(aliases, alias)
						}
					}
					policies[policyIndex] = &Policy{
						Index:      policyIndex,
						Type:       conf.GetDefault(key, "policy_type", "replication"),
						Name:       name,
						Aliases:    aliases,
						Deprecated: conf.GetBool(key, "deprecated", false),
						Default:    conf.GetBool(key, "default", false),
						Config:     map[string]string(conf.File[key]),
					}
				}
			}
			break
		}
	}
	defaultFound := false
	for _, policy := range policies {
		if policy.Default {
			defaultFound = true
		}
	}
	if !defaultFound {
		policies[0].Default = true
	}
	return PolicyList(policies), nil
}
<|endoftext|>


<commit_before>package service_test

import (
	"errors"
	"fmt"

	testapi "github.com/cloudfoundry/cli/cf/api/fakes"
	"github.com/cloudfoundry/cli/cf/command_registry"
	"github.com/cloudfoundry/cli/cf/configuration/core_config"
	cferrors "github.com/cloudfoundry/cli/cf/errors"
	"github.com/cloudfoundry/cli/cf/models"
	testcmd "github.com/cloudfoundry/cli/testhelpers/commands"
	testconfig "github.com/cloudfoundry/cli/testhelpers/configuration"
	"github.com/cloudfoundry/cli/testhelpers/maker"
	testreq "github.com/cloudfoundry/cli/testhelpers/requirements"
	testterm "github.com/cloudfoundry/cli/testhelpers/terminal"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	. "github.com/cloudfoundry/cli/testhelpers/matchers"
)

var _ = Describe("purge-service-offering command", func() {
	var (
		requirementsFactory *testreq.FakeReqFactory
		config              core_config.Repository
		ui                  *testterm.FakeUI
		serviceRepo         *testapi.FakeServiceRepository
		deps                command_registry.Dependency
	)

	updateCommandDependency := func(pluginCall bool) {
		deps.Ui = ui
		deps.RepoLocator = deps.RepoLocator.SetServiceRepository(serviceRepo)
		deps.Config = config
		command_registry.Commands.SetCommand(command_registry.Commands.FindCommand("purge-service-offering").SetDependency(deps, pluginCall))
	}

	runCommand := func(args []string) bool {
		return testcmd.RunCliCommand("purge-service-offering", args, requirementsFactory, updateCommandDependency, false)
	}

	BeforeEach(func() {
		ui = &testterm.FakeUI{}
		config = testconfig.NewRepositoryWithDefaults()
		serviceRepo = &testapi.FakeServiceRepository{}
		requirementsFactory = &testreq.FakeReqFactory{LoginSuccess: true, TargetedSpaceSuccess: true}
	})

	Describe("requirements", func() {
		It("fails when not logged in", func() {
			requirementsFactory.LoginSuccess = false

			passed := runCommand([]string{"-f", "whatever"})

			Expect(passed).To(BeFalse())
		})

		It("fails when called without exactly one arg", func() {
			requirementsFactory.LoginSuccess = true

			passed := runCommand([]string{})

			Expect(passed).To(BeFalse())
			Expect(ui.Outputs).To(ContainSubstrings(
				[]string{"Incorrect Usage", "Requires", "argument"},
			))
		})
	})

	It("works when given -p and a provider name", func() {
		offering := maker.NewServiceOffering("the-service-name")
		serviceRepo.FindServiceOfferingByLabelAndProviderReturns(offering, nil)

		ui.Inputs = []string{"yes"}

		runCommand([]string{"-p", "the-provider", "the-service-name"})

		name, provider := serviceRepo.FindServiceOfferingByLabelAndProviderArgsForCall(0)
		Expect(name).To(Equal("the-service-name"))
		Expect(provider).To(Equal("the-provider"))

		Expect(serviceRepo.PurgeServiceOfferingArgsForCall(0)).To(Equal(offering))
	})

	It("works when not given a provider", func() {
		offering := maker.NewServiceOffering("the-service-name")
		serviceRepo.FindServiceOfferingByLabelAndProviderReturns(offering, nil)

		ui.Inputs = []string{"yes"}

		runCommand([]string{"the-service-name"})

		Expect(ui.Outputs).To(ContainSubstrings([]string{"WARNING"}))
		Expect(ui.Prompts).To(ContainSubstrings([]string{"Really purge service", "the-service-name"}))
		Expect(ui.Outputs).To(ContainSubstrings([]string{"Purging service the-service-name..."}))

		name, provider := serviceRepo.FindServiceOfferingByLabelAndProviderArgsForCall(0)
		Expect(name).To(Equal("the-service-name"))
		Expect(provider).To(BeEmpty())
		Expect(serviceRepo.PurgeServiceOfferingArgsForCall(0)).To(Equal(offering))

		Expect(ui.Outputs).To(ContainSubstrings([]string{"OK"}))
	})

	It("exits when the user does not acknowledge the confirmation", func() {
		ui.Inputs = []string{"no"}

		runCommand([]string{"the-service-name"})

		Expect(serviceRepo.FindServiceOfferingByLabelAndProviderCallCount()).To(Equal(1))
		Expect(serviceRepo.PurgeServiceOfferingCallCount()).To(BeZero())
	})

	It("does not prompt with confirmation when -f is passed", func() {
		offering := maker.NewServiceOffering("the-service-name")
		serviceRepo.FindServiceOfferingByLabelAndProviderReturns(offering, nil)

		runCommand(
			[]string{"-f", "the-service-name"},
		)

		Expect(len(ui.Prompts)).To(Equal(0))
		Expect(serviceRepo.PurgeServiceOfferingCallCount()).To(Equal(1))
	})

	It("fails with an error message when the request fails", func() {
		serviceRepo.FindServiceOfferingByLabelAndProviderReturns(models.ServiceOffering{}, errors.New("oh no!"))

		runCommand(
			[]string{"-f", "-p", "the-provider", "the-service-name"},
		)

		Expect(ui.Outputs).To(ContainSubstrings(
			[]string{"FAILED"},
			[]string{"oh no!"},
		))

		Expect(serviceRepo.PurgeServiceOfferingCallCount()).To(BeZero())
	})

	It("fails with an error message when the purging request fails", func() {
		serviceRepo.PurgeServiceOfferingReturns(fmt.Errorf("crumpets insufficiently buttered"))

		runCommand(
			[]string{"-f", "-p", "the-provider", "the-service-name"},
		)

		Expect(ui.Outputs).To(ContainSubstrings(
			[]string{"FAILED"},
			[]string{"crumpets insufficiently buttered"},
		))
	})

	It("indicates when a service doesn't exist", func() {
		serviceRepo.FindServiceOfferingByLabelAndProviderReturns(models.ServiceOffering{}, cferrors.NewModelNotFoundError("Service Offering", ""))

		ui.Inputs = []string{"yes"}

		runCommand(
			[]string{"-p", "the-provider", "the-service-name"},
		)

		Expect(ui.Outputs).To(ContainSubstrings([]string{"Service offering", "does not exist"}))
		Expect(ui.Outputs).ToNot(ContainSubstrings([]string{"WARNING"}))
		Expect(ui.Outputs).ToNot(ContainSubstrings([]string{"Ok"}))

		Expect(serviceRepo.PurgeServiceOfferingCallCount()).To(BeZero())
	})
})
<commit_msg>Rewrite purge_service_offering_test<commit_after>package service_test

import (
	"github.com/cloudfoundry/cli/cf/command_registry"
	"github.com/cloudfoundry/cli/cf/commands/service"
	"github.com/cloudfoundry/cli/cf/configuration/core_config"
	"github.com/cloudfoundry/cli/cf/errors"
	"github.com/cloudfoundry/cli/cf/models"
	"github.com/cloudfoundry/cli/cf/requirements"
	"github.com/cloudfoundry/cli/flags"

	testapi "github.com/cloudfoundry/cli/cf/api/fakes"
	fakerequirements "github.com/cloudfoundry/cli/cf/requirements/fakes"
	testconfig "github.com/cloudfoundry/cli/testhelpers/configuration"
	testterm "github.com/cloudfoundry/cli/testhelpers/terminal"

	. "github.com/cloudfoundry/cli/testhelpers/matchers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("PurgeServiceOffering", func() {
	var (
		ui          *testterm.FakeUI
		configRepo  core_config.Repository
		serviceRepo *testapi.FakeServiceRepository

		cmd         command_registry.Command
		deps        command_registry.Dependency
		factory     *fakerequirements.FakeFactory
		flagContext flags.FlagContext

		loginRequirement requirements.Requirement
	)

	BeforeEach(func() {
		ui = &testterm.FakeUI{}
		configRepo = testconfig.NewRepositoryWithDefaults()
		serviceRepo = &testapi.FakeServiceRepository{}
		repoLocator := deps.RepoLocator.SetServiceRepository(serviceRepo)

		deps = command_registry.Dependency{
			Ui:          ui,
			Config:      configRepo,
			RepoLocator: repoLocator,
		}

		cmd = &service.PurgeServiceOffering{}
		cmd.SetDependency(deps, false)

		flagContext = flags.NewFlagContext(cmd.MetaData().Flags)
		factory = &fakerequirements.FakeFactory{}

		loginRequirement = &passingRequirement{Name: "login-requirement"}
		factory.NewLoginRequirementReturns(loginRequirement)
	})

	Describe("Requirements", func() {
		Context("when not provided exactly one arg", func() {
			BeforeEach(func() {
				flagContext.Parse("service", "extra-arg")
			})

			It("fails with usage", func() {
				Expect(func() { cmd.Requirements(factory, flagContext) }).To(Panic())
				Expect(ui.Outputs).To(ContainSubstrings(
					[]string{"FAILED"},
					[]string{"Incorrect Usage. Requires an argument"},
				))
			})
		})

		Context("when provided exactly one arg", func() {
			BeforeEach(func() {
				flagContext.Parse("service")
			})

			It("returns a LoginRequirement", func() {
				actualRequirements, err := cmd.Requirements(factory, flagContext)
				Expect(err).NotTo(HaveOccurred())
				Expect(factory.NewLoginRequirementCallCount()).To(Equal(1))
				Expect(actualRequirements).To(ContainElement(loginRequirement))
			})
		})
	})

	Describe("Execute", func() {
		BeforeEach(func() {
			err := flagContext.Parse("service")
			Expect(err).NotTo(HaveOccurred())
			_, err = cmd.Requirements(factory, flagContext)
			Expect(err).NotTo(HaveOccurred())
		})

		It("tries to find the service offering by label and provider", func() {
			ui.Inputs = []string{"n"}
			cmd.Execute(flagContext)
			Expect(serviceRepo.FindServiceOfferingByLabelAndProviderCallCount()).To(Equal(1))
		})

		Context("when finding the service offering succeeds", func() {
			BeforeEach(func() {
				serviceOffering := models.ServiceOffering{}
				serviceOffering.Guid = "service-offering-guid"
				serviceRepo.FindServiceOfferingByLabelAndProviderReturns(serviceOffering, nil)
			})

			It("asks the user to confirm", func() {
				ui.Inputs = []string{"n"}
				cmd.Execute(flagContext)
				Expect(ui.Outputs).To(ContainSubstrings([]string{"WARNING"}))
				Expect(ui.Prompts).To(ContainSubstrings([]string{"Really purge service offering service from Cloud Foundry?"}))
			})

			Context("when the user confirms", func() {
				BeforeEach(func() {
					ui.Inputs = []string{"y"}
				})

				It("tries to purge the service offering", func() {
					cmd.Execute(flagContext)
					Expect(serviceRepo.PurgeServiceOfferingCallCount()).To(Equal(1))
				})

				Context("when purging succeeds", func() {
					BeforeEach(func() {
						serviceRepo.PurgeServiceOfferingReturns(nil)
					})

					It("says OK", func() {
						cmd.Execute(flagContext)
						Expect(ui.Outputs).To(ContainSubstrings([]string{"OK"}))
					})
				})
			})

			Context("when the user does not confirm", func() {
				BeforeEach(func() {
					ui.Inputs = []string{"n"}
				})

				It("does not try to purge the service offering", func() {
					cmd.Execute(flagContext)
					Expect(serviceRepo.PurgeServiceOfferingCallCount()).To(BeZero())
				})
			})
		})

		Context("when finding the service offering fails with an error other than 404", func() {
			BeforeEach(func() {
				serviceRepo.FindServiceOfferingByLabelAndProviderReturns(models.ServiceOffering{}, errors.New("find-err"))
			})

			It("fails with error", func() {
				Expect(func() { cmd.Execute(flagContext) }).To(Panic())
				Expect(ui.Outputs).To(ContainSubstrings([]string{"FAILED"}))
			})
		})

		Context("when finding the service offering fails with 404 not found", func() {
			BeforeEach(func() {
				serviceRepo.FindServiceOfferingByLabelAndProviderReturns(
					models.ServiceOffering{},
					errors.NewModelNotFoundError("model-type", "find-err"),
				)
			})

			It("warns the user", func() {
				cmd.Execute(flagContext)
				Expect(ui.Outputs).To(ContainSubstrings(
					[]string{"Service offering does not exist"},
				))
			})
		})
	})
})
<|endoftext|>


<commit_before>package bongo

import (
	"os"
	"path/filepath"
	"strings"

	"github.com/bongo-contrib/models"

	"github.com/bongo-contrib/utils"
)

var supportedExtensions = []string{".md", ".MD", "..markdown"}

//DefaultLoader  is the default FileLoader implementation
type DefaultLoader struct{}

// NewLoader returns default FileLoader implementation.
func NewLoader() *DefaultLoader {
	return &DefaultLoader{}
}

// Load loads files found in the base path for processing.
func (d DefaultLoader) Load(base string) ([]string, error) {
	return func(base string) ([]string, error) {
		var rst []string
		err := filepath.Walk(base, func(path string, info os.FileInfo, err error) error {
			switch {
			case err != nil:
				return err
			case info.IsDir():
				return nil
			case !utils.HasExt(path, supportedExtensions...):
				return nil
			case strings.Contains(path, models.OutputDir):
				return nil
			}
			rst = append(rst, path)
			return nil
		})
		if err != nil {
			return nil, err
		}
		return rst, nil

	}(base)

}
<commit_msg>Remove unnecessary dependencies<commit_after>package bongo

import (
	"os"
	"path/filepath"
	"strings"
)

var supportedExtensions = []string{".md", ".MD", "..markdown"}

//DefaultLoader  is the default FileLoader implementation
type DefaultLoader struct{}

// NewLoader returns default FileLoader implementation.
func NewLoader() *DefaultLoader {
	return &DefaultLoader{}
}

// Load loads files found in the base path for processing.
func (d DefaultLoader) Load(base string) ([]string, error) {
	return func(base string) ([]string, error) {
		var rst []string
		err := filepath.Walk(base, func(path string, info os.FileInfo, err error) error {
			switch {
			case err != nil:
				return err
			case info.IsDir():
				return nil
			case !HasExt(path, supportedExtensions...):
				return nil
			case strings.Contains(path, OutputDir):
				return nil
			}
			rst = append(rst, path)
			return nil
		})
		if err != nil {
			return nil, err
		}
		return rst, nil

	}(base)

}
<|endoftext|>


<commit_before>package kubernetesprovider

import (
	"context"

	detector "github.com/rancher/kubernetes-provider-detector"
	v3 "github.com/rancher/rancher/pkg/apis/management.cattle.io/v3"
	v32 "github.com/rancher/rancher/pkg/generated/controllers/management.cattle.io/v3"
	"github.com/rancher/rancher/pkg/wrangler"
	"k8s.io/client-go/kubernetes"
)

type handler struct {
	ctx                context.Context
	clusters           v32.ClusterClient
	localClusterClient kubernetes.Interface
	mcm                wrangler.MultiClusterManager
}

func Register(ctx context.Context,
	clusters v32.ClusterController,
	localClusterClient kubernetes.Interface,
	mcm wrangler.MultiClusterManager,
) {
	h := &handler{
		ctx:                ctx,
		clusters:           clusters,
		localClusterClient: localClusterClient,
		mcm:                mcm,
	}
	clusters.OnChange(ctx, "kubernetes-provider", h.OnChange)
}

func (h *handler) OnChange(key string, cluster *v3.Cluster) (*v3.Cluster, error) {
	if cluster == nil || cluster.Status.Provider != "" {
		return cluster, nil
	}

	if !v3.ClusterConditionReady.IsTrue(cluster) {
		return cluster, nil
	}

	var client kubernetes.Interface
	if cluster.Spec.Internal {
		client = h.localClusterClient
	} else if k8s, err := h.mcm.K8sClient(cluster.Name); err != nil {
		// ignore error. If we can't get a client just ignore it. The cluster probably isn't happy
		// yet and we will get an update later when it is.
		return nil, nil
	} else if k8s != nil {
		client = k8s
	}

	if client == nil {
		return cluster, nil
	}

	provider, err := detector.DetectProvider(h.ctx, client)
	if err == detector.ErrUnknownProvider {
		return cluster, nil
	} else if err != nil {
		return cluster, err
	}
	cluster = cluster.DeepCopy()
	cluster.Status.Provider = provider
	return h.clusters.Update(cluster)
}
<commit_msg>Don't return if cluster is windows type but provider is not windows<commit_after>package kubernetesprovider

import (
	"context"

	detector "github.com/rancher/kubernetes-provider-detector"
	v3 "github.com/rancher/rancher/pkg/apis/management.cattle.io/v3"
	v32 "github.com/rancher/rancher/pkg/generated/controllers/management.cattle.io/v3"
	"github.com/rancher/rancher/pkg/wrangler"
	"k8s.io/client-go/kubernetes"
)

type handler struct {
	ctx                context.Context
	clusters           v32.ClusterClient
	localClusterClient kubernetes.Interface
	mcm                wrangler.MultiClusterManager
}

func Register(ctx context.Context,
	clusters v32.ClusterController,
	localClusterClient kubernetes.Interface,
	mcm wrangler.MultiClusterManager,
) {
	h := &handler{
		ctx:                ctx,
		clusters:           clusters,
		localClusterClient: localClusterClient,
		mcm:                mcm,
	}
	clusters.OnChange(ctx, "kubernetes-provider", h.OnChange)
}

func (h *handler) OnChange(key string, cluster *v3.Cluster) (*v3.Cluster, error) {
	if cluster == nil || cluster.Status.Provider != "" {
		// if cluster has windows enabled and provider is not rke.windows, continue to detect if it has windows nodes
		if cluster == nil || !cluster.Spec.WindowsPreferedCluster || cluster.Status.Provider == "rke.windows" {
			return cluster, nil
		}
	}

	if !v3.ClusterConditionReady.IsTrue(cluster) {
		return cluster, nil
	}

	var client kubernetes.Interface
	if cluster.Spec.Internal {
		client = h.localClusterClient
	} else if k8s, err := h.mcm.K8sClient(cluster.Name); err != nil {
		// ignore error. If we can't get a client just ignore it. The cluster probably isn't happy
		// yet and we will get an update later when it is.
		return nil, nil
	} else if k8s != nil {
		client = k8s
	}

	if client == nil {
		return cluster, nil
	}

	provider, err := detector.DetectProvider(h.ctx, client)
	if err == detector.ErrUnknownProvider {
		return cluster, nil
	} else if err != nil {
		return cluster, err
	}
	cluster = cluster.DeepCopy()
	cluster.Status.Provider = provider
	return h.clusters.Update(cluster)
}
<|endoftext|>


<commit_before>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package tests_test

import (
	"context"
	"fmt"
	"sync"
	"time"

	"kubevirt.io/kubevirt/tests/framework/checks"

	appsv1 "k8s.io/api/apps/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/informers"
	"k8s.io/client-go/tools/cache"
	"k8s.io/client-go/util/retry"

	v1 "kubevirt.io/api/core/v1"
	"kubevirt.io/client-go/kubecli"

	"kubevirt.io/kubevirt/tests/flags"
	"kubevirt.io/kubevirt/tests/util"

	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
)

var _ = Describe("[Serial][sig-operator]virt-handler canary upgrade", Serial, func() {

	var err error
	var originalKV *v1.KubeVirt
	var virtCli kubecli.KubevirtClient
	var dsInformer cache.SharedIndexInformer
	var stopCh chan struct{}
	var lastObservedEvent string

	const (
		e2eCanaryTestAnnotation = "e2e-canary-test"
		canaryTestNodeTimeout   = 60
	)

	BeforeEach(func() {
		if !checks.HasAtLeastTwoNodes() {
			Skip("this test requires at least 2 nodes")
		}

		virtCli, err = kubecli.GetKubevirtClient()
		Expect(err).ToNot(HaveOccurred())

		originalKV = util.GetCurrentKv(virtCli).DeepCopy()

		stopCh = make(chan struct{})
		lastObservedEvent = ""

		informerFactory := informers.NewSharedInformerFactoryWithOptions(virtCli, 0, informers.WithNamespace(flags.KubeVirtInstallNamespace), informers.WithTweakListOptions(func(opts *metav1.ListOptions) {
			opts.LabelSelector = "kubevirt.io=virt-handler"
		}))
		dsInformer = informerFactory.Apps().V1().DaemonSets().Informer()
	})

	AfterEach(func() {
		close(stopCh)

		retry.RetryOnConflict(retry.DefaultRetry, func() error {
			_, err := virtCli.KubeVirt(flags.KubeVirtInstallNamespace).Update(originalKV)
			return err
		})

		Eventually(func() bool {
			ds, err := virtCli.AppsV1().DaemonSets(originalKV.Namespace).Get(context.Background(), "virt-handler", metav1.GetOptions{})
			Expect(err).ToNot(HaveOccurred())
			return ds.Status.DesiredNumberScheduled == ds.Status.NumberReady && ds.Spec.UpdateStrategy.RollingUpdate.MaxUnavailable.IntValue() == 1
		}, 60*time.Second, 1*time.Second).Should(BeTrue(), "waiting for virt-handler to be ready")
	})

	getVirtHandler := func() *appsv1.DaemonSet {
		daemonSet, err := virtCli.AppsV1().DaemonSets(flags.KubeVirtInstallNamespace).Get(context.Background(), "virt-handler", metav1.GetOptions{})
		ExpectWithOffset(1, err).NotTo(HaveOccurred())
		return daemonSet
	}

	updateVirtHandler := func() error {
		kv := util.GetCurrentKv(virtCli)

		patch := fmt.Sprintf(`{"spec": { "template": {"metadata": {"annotations": {"%s": "test"}}}}}`,
			e2eCanaryTestAnnotation)
		kv.Spec.CustomizeComponents = v1.CustomizeComponents{
			Patches: []v1.CustomizeComponentsPatch{
				{
					ResourceName: "virt-handler",
					ResourceType: "DaemonSet",
					Type:         v1.StrategicMergePatchType,
					Patch:        patch,
				},
			},
		}
		return retry.RetryOnConflict(retry.DefaultRetry, func() error {
			_, err := virtCli.KubeVirt(flags.KubeVirtInstallNamespace).Update(kv)
			return err
		})
	}

	isVirtHandlerUpdated := func(ds *appsv1.DaemonSet) bool {
		_, exists := ds.Spec.Template.Annotations[e2eCanaryTestAnnotation]
		return exists
	}

	removeExpectedAtHead := func(eventsQueue []string, expectedEvent string) []string {
		defer GinkgoRecover()

		if expectedEvent == lastObservedEvent {
			return eventsQueue
		}
		if len(eventsQueue) == 0 {
			return eventsQueue
		}

		head := eventsQueue[0]
		ExpectWithOffset(1, head).To(Equal(expectedEvent), fmt.Sprintf("was expecting event %s but got %s instead", expectedEvent, head))
		lastObservedEvent = expectedEvent
		return eventsQueue[1:]
	}

	processDsEvent := func(ds *appsv1.DaemonSet, eventsQueue []string) []string {
		update := ds.Spec.UpdateStrategy.RollingUpdate
		if update == nil {
			return eventsQueue
		}
		maxUnavailable := update.MaxUnavailable
		if maxUnavailable == nil {
			return eventsQueue
		}

		if maxUnavailable.String() == "10%" {
			eventsQueue = removeExpectedAtHead(eventsQueue, "maxUnavailable=10%")
		}
		if maxUnavailable.IntValue() == 1 {
			eventsQueue = removeExpectedAtHead(eventsQueue, "maxUnavailable=1")
		}
		if ds.Status.DesiredNumberScheduled == ds.Status.NumberReady {
			pods, err := virtCli.CoreV1().Pods(flags.KubeVirtInstallNamespace).List(context.Background(), metav1.ListOptions{LabelSelector: "kubevirt.io=virt-handler"})
			Expect(err).ToNot(HaveOccurred())

			var updatedPods int32
			for i := range pods.Items {
				if _, exists := pods.Items[i].Annotations[e2eCanaryTestAnnotation]; exists {
					updatedPods++
				}
			}

			if updatedPods > 0 && updatedPods == ds.Status.DesiredNumberScheduled {
				eventsQueue = removeExpectedAtHead(eventsQueue, "virt-handler=ready")
			}
		}
		return eventsQueue
	}

	It("should successfully upgrade virt-handler", func() {
		var expectedEventsLock sync.Mutex
		expectedEvents := []string{
			"maxUnavailable=1",
			"maxUnavailable=10%",
			"virt-handler=ready",
			"maxUnavailable=1",
		}

		ds, err := virtCli.AppsV1().DaemonSets(flags.KubeVirtInstallNamespace).Get(context.Background(), "virt-handler", metav1.GetOptions{})
		Expect(err).ToNot(HaveOccurred())
		nodesToUpdate := ds.Status.DesiredNumberScheduled
		Expect(nodesToUpdate > 0).To(BeTrue())

		go dsInformer.Run(stopCh)
		cache.WaitForCacheSync(stopCh, dsInformer.HasSynced)

		dsInformer.AddEventHandler(cache.ResourceEventHandlerFuncs{
			UpdateFunc: func(old, curr interface{}) {
				ds := curr.(*appsv1.DaemonSet)
				expectedEventsLock.Lock()
				defer expectedEventsLock.Unlock()
				expectedEvents = processDsEvent(ds, expectedEvents)
			},
		})

		err = updateVirtHandler()
		Expect(err).ToNot(HaveOccurred())

		updateTimeout := time.Duration(canaryTestNodeTimeout * nodesToUpdate)
		Eventually(func() bool {
			expectedEventsLock.Lock()
			defer expectedEventsLock.Unlock()
			return len(expectedEvents) == 0
		}, updateTimeout*time.Second, 1*time.Second).Should(BeTrue(), fmt.Sprintf("events %v were expected but did not occur", expectedEvents))

		Expect(isVirtHandlerUpdated(getVirtHandler())).To(BeTrue())
	})
})
<commit_msg>Quarantine canary upgrade test<commit_after>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package tests_test

import (
	"context"
	"fmt"
	"sync"
	"time"

	"kubevirt.io/kubevirt/tests/framework/checks"

	appsv1 "k8s.io/api/apps/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/informers"
	"k8s.io/client-go/tools/cache"
	"k8s.io/client-go/util/retry"

	v1 "kubevirt.io/api/core/v1"
	"kubevirt.io/client-go/kubecli"

	"kubevirt.io/kubevirt/tests/flags"
	"kubevirt.io/kubevirt/tests/util"

	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
)

var _ = Describe("[Serial][sig-operator]virt-handler canary upgrade", Serial, func() {

	var err error
	var originalKV *v1.KubeVirt
	var virtCli kubecli.KubevirtClient
	var dsInformer cache.SharedIndexInformer
	var stopCh chan struct{}
	var lastObservedEvent string

	const (
		e2eCanaryTestAnnotation = "e2e-canary-test"
		canaryTestNodeTimeout   = 60
	)

	BeforeEach(func() {
		if !checks.HasAtLeastTwoNodes() {
			Skip("this test requires at least 2 nodes")
		}

		virtCli, err = kubecli.GetKubevirtClient()
		Expect(err).ToNot(HaveOccurred())

		originalKV = util.GetCurrentKv(virtCli).DeepCopy()

		stopCh = make(chan struct{})
		lastObservedEvent = ""

		informerFactory := informers.NewSharedInformerFactoryWithOptions(virtCli, 0, informers.WithNamespace(flags.KubeVirtInstallNamespace), informers.WithTweakListOptions(func(opts *metav1.ListOptions) {
			opts.LabelSelector = "kubevirt.io=virt-handler"
		}))
		dsInformer = informerFactory.Apps().V1().DaemonSets().Informer()
	})

	AfterEach(func() {
		close(stopCh)

		retry.RetryOnConflict(retry.DefaultRetry, func() error {
			_, err := virtCli.KubeVirt(flags.KubeVirtInstallNamespace).Update(originalKV)
			return err
		})

		Eventually(func() bool {
			ds, err := virtCli.AppsV1().DaemonSets(originalKV.Namespace).Get(context.Background(), "virt-handler", metav1.GetOptions{})
			Expect(err).ToNot(HaveOccurred())
			return ds.Status.DesiredNumberScheduled == ds.Status.NumberReady && ds.Spec.UpdateStrategy.RollingUpdate.MaxUnavailable.IntValue() == 1
		}, 60*time.Second, 1*time.Second).Should(BeTrue(), "waiting for virt-handler to be ready")
	})

	getVirtHandler := func() *appsv1.DaemonSet {
		daemonSet, err := virtCli.AppsV1().DaemonSets(flags.KubeVirtInstallNamespace).Get(context.Background(), "virt-handler", metav1.GetOptions{})
		ExpectWithOffset(1, err).NotTo(HaveOccurred())
		return daemonSet
	}

	updateVirtHandler := func() error {
		kv := util.GetCurrentKv(virtCli)

		patch := fmt.Sprintf(`{"spec": { "template": {"metadata": {"annotations": {"%s": "test"}}}}}`,
			e2eCanaryTestAnnotation)
		kv.Spec.CustomizeComponents = v1.CustomizeComponents{
			Patches: []v1.CustomizeComponentsPatch{
				{
					ResourceName: "virt-handler",
					ResourceType: "DaemonSet",
					Type:         v1.StrategicMergePatchType,
					Patch:        patch,
				},
			},
		}
		return retry.RetryOnConflict(retry.DefaultRetry, func() error {
			_, err := virtCli.KubeVirt(flags.KubeVirtInstallNamespace).Update(kv)
			return err
		})
	}

	isVirtHandlerUpdated := func(ds *appsv1.DaemonSet) bool {
		_, exists := ds.Spec.Template.Annotations[e2eCanaryTestAnnotation]
		return exists
	}

	removeExpectedAtHead := func(eventsQueue []string, expectedEvent string) []string {
		defer GinkgoRecover()

		if expectedEvent == lastObservedEvent {
			return eventsQueue
		}
		if len(eventsQueue) == 0 {
			return eventsQueue
		}

		head := eventsQueue[0]
		ExpectWithOffset(1, head).To(Equal(expectedEvent), fmt.Sprintf("was expecting event %s but got %s instead", expectedEvent, head))
		lastObservedEvent = expectedEvent
		return eventsQueue[1:]
	}

	processDsEvent := func(ds *appsv1.DaemonSet, eventsQueue []string) []string {
		update := ds.Spec.UpdateStrategy.RollingUpdate
		if update == nil {
			return eventsQueue
		}
		maxUnavailable := update.MaxUnavailable
		if maxUnavailable == nil {
			return eventsQueue
		}

		if maxUnavailable.String() == "10%" {
			eventsQueue = removeExpectedAtHead(eventsQueue, "maxUnavailable=10%")
		}
		if maxUnavailable.IntValue() == 1 {
			eventsQueue = removeExpectedAtHead(eventsQueue, "maxUnavailable=1")
		}
		if ds.Status.DesiredNumberScheduled == ds.Status.NumberReady {
			pods, err := virtCli.CoreV1().Pods(flags.KubeVirtInstallNamespace).List(context.Background(), metav1.ListOptions{LabelSelector: "kubevirt.io=virt-handler"})
			Expect(err).ToNot(HaveOccurred())

			var updatedPods int32
			for i := range pods.Items {
				if _, exists := pods.Items[i].Annotations[e2eCanaryTestAnnotation]; exists {
					updatedPods++
				}
			}

			if updatedPods > 0 && updatedPods == ds.Status.DesiredNumberScheduled {
				eventsQueue = removeExpectedAtHead(eventsQueue, "virt-handler=ready")
			}
		}
		return eventsQueue
	}

	It("[QUARANTINE]should successfully upgrade virt-handler", func() {
		var expectedEventsLock sync.Mutex
		expectedEvents := []string{
			"maxUnavailable=1",
			"maxUnavailable=10%",
			"virt-handler=ready",
			"maxUnavailable=1",
		}

		ds, err := virtCli.AppsV1().DaemonSets(flags.KubeVirtInstallNamespace).Get(context.Background(), "virt-handler", metav1.GetOptions{})
		Expect(err).ToNot(HaveOccurred())
		nodesToUpdate := ds.Status.DesiredNumberScheduled
		Expect(nodesToUpdate > 0).To(BeTrue())

		go dsInformer.Run(stopCh)
		cache.WaitForCacheSync(stopCh, dsInformer.HasSynced)

		dsInformer.AddEventHandler(cache.ResourceEventHandlerFuncs{
			UpdateFunc: func(old, curr interface{}) {
				ds := curr.(*appsv1.DaemonSet)
				expectedEventsLock.Lock()
				defer expectedEventsLock.Unlock()
				expectedEvents = processDsEvent(ds, expectedEvents)
			},
		})

		err = updateVirtHandler()
		Expect(err).ToNot(HaveOccurred())

		updateTimeout := time.Duration(canaryTestNodeTimeout * nodesToUpdate)
		Eventually(func() bool {
			expectedEventsLock.Lock()
			defer expectedEventsLock.Unlock()
			return len(expectedEvents) == 0
		}, updateTimeout*time.Second, 1*time.Second).Should(BeTrue(), fmt.Sprintf("events %v were expected but did not occur", expectedEvents))

		Expect(isVirtHandlerUpdated(getVirtHandler())).To(BeTrue())
	})
})
<|endoftext|>


<commit_before>package envy

import (
	"errors"
	"fmt"
	"log"
	"os"
	"reflect"
	"strconv"
	"strings"
)

var (
	// Error used when an invalid reference is provided to the Load function
	ErrInvalidConfigType = errors.New("give me a struct")

	// Basic config error
	ErrConfigInvalid = errors.New("config is invalid")
)

var logger = log.New(os.Stderr, "[goenvy] ", log.LstdFlags|log.Lshortfile)

// interface that reads config from somewhere
type EnvironmentReader interface {
	// Method reads the environment from the source
	//
	// Returns: map[string]string of environment keys to values
	Read() map[string]string
}

// Loads directly from the environment
func Load(spec interface{}) error {
	osEnv := &OsEnvironmentReader{}
	return Load(osEnv)
}

// Loads config from the provided EnvironmentReader
func LoadFromEnv(reader EnvironmentReader, configSpec interface{}) error {
	source := reader.Read()

	// Find the value of the provided configSpec
	// It must be a struct of some kind in order for the values
	// to be set.
	s := reflect.ValueOf(configSpec).Elem()
	if s.Kind() != reflect.Struct {
		return ErrInvalidConfigType
	}

	// create a list of all errors
	errors := make([]error, 0)

	// iterate over all fields in the struct
	typeOfSpec := s.Type()
	for i := 0; i < s.NumField(); i++ {
		// reference to the value of the field (used for assignment)
		fieldValue := s.Field(i)
		// reference to the type of the field
		// (used to determine the name and any relevant struct tags)
		fieldType := typeOfSpec.Field(i)

		// Only uppercase values can be set (limitation of reflection)
		if fieldValue.CanSet() {
			fieldName := fieldType.Name

			// always assume uppercase key names
			key := strings.ToUpper(fieldName)

			// string used for outputting useful error messages
			example := fieldType.Tag.Get("example")

			// retrieve the value from the source, UPCASED
			// if this value is not available, track the error and continue with
			// the other options
			value, ok := source[key]
			if !ok {
				err := fmt.Errorf("Config not found: key=%s; example=%q", key, key, example)
				errors = append(errors, err)
				continue
			}

			// populate the struct values based on what type it is
			switch fieldValue.Kind() {
			case reflect.String:
				fieldValue.SetString(value)
			case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
				intValue, err := strconv.Atoi(value)
				if err != nil {
					err := fmt.Errorf("invalid value for int name=%s, value=%s; example=%q", key, value, key, example)
					errors = append(errors, err)
					continue
				}
				fieldValue.SetInt(int64(intValue))
			case reflect.Bool:
				boolValue, err := strconv.ParseBool(value)
				if err != nil {
					err := fmt.Errorf("invalid value for bool name=%s, value=%s; example=%q", key, value, key, example)
					errors = append(errors, err)
					continue
				}
				fieldValue.SetBool(boolValue)
			}
		}
	}

	if len(errors) > 0 {
		for _, err := range errors {
			logger.Println(err)
		}
		return ErrConfigInvalid
	}

	return nil
}

// Default EnvironmentReader
type OsEnvironmentReader struct{}

// Reads values from the os.Environ slice and returns the result
// as a map[string]string
func (o *OsEnvironmentReader) Read() map[string]string {
	result := make(map[string]string)
	for _, envVar := range os.Environ() {
		parts := strings.SplitN(envVar, "=", 2)
		result[parts[0]] = parts[1]
	}

	return result
}
<commit_msg>chang e<commit_after>package envy

import (
	"errors"
	"fmt"
	"log"
	"os"
	"reflect"
	"strconv"
	"strings"
)

var (
	// Error used when an invalid reference is provided to the Load function
	ErrInvalidConfigType = errors.New("a struct value is required for loading config")

	// Basic config error
	ErrConfigInvalid = errors.New("config is invalid")
)

var logger = log.New(os.Stderr, "[goenvy] ", log.LstdFlags|log.Lshortfile)

// interface that reads config from somewhere
type EnvironmentReader interface {
	// Method reads the environment from the source
	//
	// Returns: map[string]string of environment keys to values
	Read() map[string]string
}

// Loads directly from the environment
func Load(spec interface{}) error {
	osEnv := &OsEnvironmentReader{}
	return Load(osEnv)
}

// Loads config from the provided EnvironmentReader
func LoadFromEnv(reader EnvironmentReader, configSpec interface{}) error {
	source := reader.Read()

	// Find the value of the provided configSpec
	// It must be a struct of some kind in order for the values
	// to be set.
	s := reflect.ValueOf(configSpec).Elem()
	if s.Kind() != reflect.Struct {
		return ErrInvalidConfigType
	}

	// create a list of all errors
	errors := make([]error, 0)

	// iterate over all fields in the struct
	typeOfSpec := s.Type()
	for i := 0; i < s.NumField(); i++ {
		// reference to the value of the field (used for assignment)
		fieldValue := s.Field(i)
		// reference to the type of the field
		// (used to determine the name and any relevant struct tags)
		fieldType := typeOfSpec.Field(i)

		// Only uppercase values can be set (limitation of reflection)
		if fieldValue.CanSet() {
			fieldName := fieldType.Name

			// always assume uppercase key names
			key := strings.ToUpper(fieldName)

			// string used for outputting useful error messages
			example := fieldType.Tag.Get("example")

			// retrieve the value from the source, UPCASED
			// if this value is not available, track the error and continue with
			// the other options
			value, ok := source[key]
			if !ok {
				err := fmt.Errorf("Config not found: key=%s; example=%q", key, key, example)
				errors = append(errors, err)
				continue
			}

			// populate the struct values based on what type it is
			switch fieldValue.Kind() {
			case reflect.String:
				fieldValue.SetString(value)
			case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
				intValue, err := strconv.Atoi(value)
				if err != nil {
					err := fmt.Errorf("invalid value for int name=%s, value=%s; example=%q", key, value, key, example)
					errors = append(errors, err)
					continue
				}
				fieldValue.SetInt(int64(intValue))
			case reflect.Bool:
				boolValue, err := strconv.ParseBool(value)
				if err != nil {
					err := fmt.Errorf("invalid value for bool name=%s, value=%s; example=%q", key, value, key, example)
					errors = append(errors, err)
					continue
				}
				fieldValue.SetBool(boolValue)
			}
		}
	}

	if len(errors) > 0 {
		for _, err := range errors {
			logger.Println(err)
		}
		return ErrConfigInvalid
	}

	return nil
}

// Default EnvironmentReader
type OsEnvironmentReader struct{}

// Reads values from the os.Environ slice and returns the result
// as a map[string]string
func (o *OsEnvironmentReader) Read() map[string]string {
	result := make(map[string]string)
	for _, envVar := range os.Environ() {
		parts := strings.SplitN(envVar, "=", 2)
		result[parts[0]] = parts[1]
	}

	return result
}
<|endoftext|>


<commit_before>/*
Package epub generates valid EPUB 3.0 files with additional EPUB 2.0 table of
contents (as seen here: https://github.com/bmaupin/epub-samples) for maximum
compatibility.

Basic usage:

	// Create a new EPUB
	e := epub.NewEpub("My title")

	// Set the author
	e.SetAuthor("Hingle McCringleberry")

	// Add a section
	section1Content := `
Section 1

This is a paragraph.
`
	e.AddSection("Section 1", section1Content, "section0001.xhtml")

	// The section file name is optional
	section2Content := `
Section 2

This is a paragraph.
`
	e.AddSection("Section 2", section2Content, "")

	// Add an image from a local file
	e.AddImage("testdata/gophercolor16x16.png", "go-gopher.png")

	// Add an image from a URL. The image filename is also optional
	e.AddImage("https://golang.org/doc/gopher/gophercolor16x16.png", "")

	// Write the EPUB
	err = e.Write("My EPUB.epub")
	if err != nil {
		// handle error
	}

*/
package epub

import (
	"fmt"
	"path/filepath"

	"github.com/satori/go.uuid"
)

const (
	defaultEpubLang   = "en"
	imageFileFormat   = "image%04d%s"
	sectionFileFormat = "section%04d.xhtml"
	urnUUIDPrefix     = "urn:uuid:"
)

// Epub implements an EPUB file.
type Epub struct {
	author   string
	images   map[string]string // Images added to the EPUB
	lang     string            // Language
	pkg      *pkg              // The package file (package.opf)
	sections map[string]xhtml  // Sections (chapters)
	title    string
	toc      *toc // Table of contents
	uuid     string
}

// NewEpub returns a new Epub.
func NewEpub(title string) *Epub {
	e := &Epub{}
	e.images = make(map[string]string)
	e.sections = make(map[string]xhtml)
	e.pkg = newPackage()
	e.toc = newToc()
	// Set minimal required attributes
	e.SetLang(defaultEpubLang)
	e.SetTitle(title)
	e.SetUUID(uuid.NewV4().String())

	return e
}

// AddImage adds an image to the EPUB and returns a relative path that can be
// used in the content of a section.
//
// The image source should either be a URL or a path to a local file; in either
// case, the image will be retrieved and stored in the EPUB.
//
// The image filename will be used when storing the image in the EPUB and must
// be unique. The image filename is optional; if no filename is provided, one
// will be generated.
func (e *Epub) AddImage(imageSource string, imageFilename string) (string, error) {
	// Generate an image filename if one isn't provided
	if imageFilename == "" {
		imageFilename = fmt.Sprintf(imageFileFormat, len(e.images)+1, filepath.Ext(imageSource))
	}

	if _, ok := e.images[imageFilename]; ok {
		return "", fmt.Errorf("Image filename %s already used", imageFilename)
	}

	e.images[imageFilename] = imageSource

	return filepath.Join(
		"..",
		imageFolderName,
		imageFilename,
	), nil
}

// AddSection adds a new section (chapter, etc) to the EPUB and returns a
// relative path to the section that can be used from another section (for
// links).
//
// The title will be used for the table of contents.
//
// The content must be valid XHTML that will go between the  tags of the
// section XHTML file. The content will not be validated.
//
// The section filename will be used when storing the image in the EPUB and must
// be unique.The section filename is optional; if no filename is provided, one
// will be generated.
//
// The section will be shown in the table of contents in the same order it was
// added to the EPUB.
func (e *Epub) AddSection(sectionTitle string, sectionContent string, sectionFilename string) (string, error) {
	// Generate a section filename if one isn't provided
	if sectionFilename == "" {
		sectionFilename = fmt.Sprintf(sectionFileFormat, len(e.sections)+1)
	}

	if _, ok := e.sections[sectionFilename]; ok {
		return "", fmt.Errorf("Section filename %s already used", sectionFilename)
	}

	x := newXhtml(sectionContent)
	x.setTitle(sectionTitle)

	e.sections[sectionFilename] = *x

	return sectionFilename, nil
}

// Author returns the author of the EPUB.
func (e *Epub) Author() string {
	return e.author
}

// Lang returns the language of the EPUB.
func (e *Epub) Lang() string {
	return e.lang
}

// SetAuthor sets the author of the EPUB.
func (e *Epub) SetAuthor(author string) {
	e.author = author
	e.pkg.setAuthor(author)
}

// SetLang sets the language of the EPUB.
func (e *Epub) SetLang(lang string) {
	e.lang = lang
	e.pkg.setLang(lang)
}

// SetTitle sets the title of the EPUB.
func (e *Epub) SetTitle(title string) {
	e.title = title
	e.pkg.setTitle(title)
	e.toc.setTitle(title)
}

// SetUUID sets the UUID of the EPUB. A UUID will be automatically be generated
// for you when the NewEpub method is run.
func (e *Epub) SetUUID(uuid string) {
	e.uuid = uuid
	e.pkg.setUUID(urnUUIDPrefix + uuid)
	e.toc.setUUID(urnUUIDPrefix + uuid)
}

// Title returns the title of the EPUB.
func (e *Epub) Title() string {
	return e.title
}

// UUID returns the UUID of the EPUB.
func (e *Epub) UUID() string {
	return e.uuid
}
<commit_msg>Minimize the points and types of failures in epub.go (http://bitsquid.blogspot.ca/2012/02/sensible-error-handling-part-2.html)<commit_after>/*
Package epub generates valid EPUB 3.0 files with additional EPUB 2.0 table of
contents (as seen here: https://github.com/bmaupin/epub-samples) for maximum
compatibility.

Basic usage:

	// Create a new EPUB
	e := epub.NewEpub("My title")

	// Set the author
	e.SetAuthor("Hingle McCringleberry")

	// Add a section
	section1Content := `
Section 1

This is a paragraph.
`
	e.AddSection("Section 1", section1Content, "section0001.xhtml")

	// The section file name is optional
	section2Content := `
Section 2

This is a paragraph.
`
	e.AddSection("Section 2", section2Content, "")

	// Add an image from a local file
	e.AddImage("testdata/gophercolor16x16.png", "go-gopher.png")

	// Add an image from a URL. The image filename is also optional
	e.AddImage("https://golang.org/doc/gopher/gophercolor16x16.png", "")

	// Write the EPUB
	err = e.Write("My EPUB.epub")
	if err != nil {
		// handle error
	}

*/
package epub

import (
	"errors"
	"fmt"
	"path/filepath"

	"github.com/satori/go.uuid"
)

var (
	// ErrFilenameAlreadyUsed is thrown by AddImage and AddSection if the same
	// filename is used more than once
	ErrFilenameAlreadyUsed = errors.New("Filename already used")
)

const (
	defaultEpubLang   = "en"
	imageFileFormat   = "image%04d%s"
	sectionFileFormat = "section%04d.xhtml"
	urnUUIDPrefix     = "urn:uuid:"
)

// Epub implements an EPUB file.
type Epub struct {
	author   string
	images   map[string]string // Images added to the EPUB
	lang     string            // Language
	pkg      *pkg              // The package file (package.opf)
	sections map[string]xhtml  // Sections (chapters)
	title    string
	toc      *toc // Table of contents
	uuid     string
}

// NewEpub returns a new Epub.
func NewEpub(title string) *Epub {
	e := &Epub{}
	e.images = make(map[string]string)
	e.sections = make(map[string]xhtml)
	e.pkg = newPackage()
	e.toc = newToc()
	// Set minimal required attributes
	e.SetLang(defaultEpubLang)
	e.SetTitle(title)
	e.SetUUID(uuid.NewV4().String())

	return e
}

// AddImage adds an image to the EPUB and returns a relative path that can be
// used in the content of a section.
//
// The image source should either be a URL or a path to a local file; in either
// case, the image will be retrieved and stored in the EPUB.
//
// The image filename will be used when storing the image in the EPUB and must
// be unique. If the same filename is used more than once,
// ErrFilenameAlreadyUsed will be returned. The image filename is optional; if
// no filename is provided, one will be generated.
func (e *Epub) AddImage(imageSource string, imageFilename string) (string, error) {
	// Generate an image filename if one isn't provided
	if imageFilename == "" {
		imageFilename = fmt.Sprintf(imageFileFormat, len(e.images)+1, filepath.Ext(imageSource))
	}

	if _, ok := e.images[imageFilename]; ok {
		return "", ErrFilenameAlreadyUsed
	}

	e.images[imageFilename] = imageSource

	return filepath.Join(
		"..",
		imageFolderName,
		imageFilename,
	), nil
}

// AddSection adds a new section (chapter, etc) to the EPUB and returns a
// relative path to the section that can be used from another section (for
// links).
//
// The title will be used for the table of contents.
//
// The content must be valid XHTML that will go between the  tags of the
// section XHTML file. The content will not be validated.
//
// The section filename will be used when storing the image in the EPUB and must
// be unique. If the same filename is used more than once,
// ErrFilenameAlreadyUsed will be returned. The section filename is optional;
// if no filename is provided, one will be generated.
//
// The section will be shown in the table of contents in the same order it was
// added to the EPUB.
func (e *Epub) AddSection(sectionTitle string, sectionContent string, sectionFilename string) (string, error) {
	// Generate a section filename if one isn't provided
	if sectionFilename == "" {
		sectionFilename = fmt.Sprintf(sectionFileFormat, len(e.sections)+1)
	}

	if _, ok := e.sections[sectionFilename]; ok {
		return "", ErrFilenameAlreadyUsed
	}

	x := newXhtml(sectionContent)
	x.setTitle(sectionTitle)

	e.sections[sectionFilename] = *x

	return sectionFilename, nil
}

// Author returns the author of the EPUB.
func (e *Epub) Author() string {
	return e.author
}

// Lang returns the language of the EPUB.
func (e *Epub) Lang() string {
	return e.lang
}

// SetAuthor sets the author of the EPUB.
func (e *Epub) SetAuthor(author string) {
	e.author = author
	e.pkg.setAuthor(author)
}

// SetLang sets the language of the EPUB.
func (e *Epub) SetLang(lang string) {
	e.lang = lang
	e.pkg.setLang(lang)
}

// SetTitle sets the title of the EPUB.
func (e *Epub) SetTitle(title string) {
	e.title = title
	e.pkg.setTitle(title)
	e.toc.setTitle(title)
}

// SetUUID sets the UUID of the EPUB. A UUID will be automatically be generated
// for you when the NewEpub method is run.
func (e *Epub) SetUUID(uuid string) {
	e.uuid = uuid
	e.pkg.setUUID(urnUUIDPrefix + uuid)
	e.toc.setUUID(urnUUIDPrefix + uuid)
}

// Title returns the title of the EPUB.
func (e *Epub) Title() string {
	return e.title
}

// UUID returns the UUID of the EPUB.
func (e *Epub) UUID() string {
	return e.uuid
}
<|endoftext|>