<commit_before>package nat

import (
	"errors"
	"fmt"
	"strconv"
	"strings"
	"sync"
	"time"

	nat "github.com/fd/go-nat"
	logging "github.com/ipfs/go-log"
	ma "github.com/jbenet/go-multiaddr"
	manet "github.com/jbenet/go-multiaddr-net"
	goprocess "github.com/jbenet/goprocess"
	periodic "github.com/jbenet/goprocess/periodic"
	notifier "github.com/whyrusleeping/go-notifier"
)

var (
	// ErrNoMapping signals no mapping exists for an address
	ErrNoMapping = errors.New("mapping not established")
)

var log = logging.Logger("nat")

// MappingDuration is a default port mapping duration.
// Port mappings are renewed every (MappingDuration / 3)
const MappingDuration = time.Second * 60

// CacheTime is the time a mapping will cache an external address for
const CacheTime = time.Second * 15

// DiscoverNAT looks for a NAT device in the network and
// returns an object that can manage port mappings.
func DiscoverNAT() *NAT {
	nat, err := nat.DiscoverGateway()
	if err != nil {
		log.Debug("DiscoverGateway error:", err)
		return nil
	}
	addr, err := nat.GetDeviceAddress()
	if err != nil {
		log.Debug("DiscoverGateway address error:", err)
	} else {
		log.Debug("DiscoverGateway address:", addr)
	}
	return newNAT(nat)
}

// NAT is an object that manages address port mappings in
// NATs (Network Address Translators). It is a long-running
// service that will periodically renew port mappings,
// and keep an up-to-date list of all the external addresses.
type NAT struct {
	nat  nat.NAT
	proc goprocess.Process // manages nat mappings lifecycle

	mappingmu sync.RWMutex // guards mappings
	mappings  map[*mapping]struct{}

	Notifier
}

func newNAT(realNAT nat.NAT) *NAT {
	return &NAT{
		nat:      realNAT,
		proc:     goprocess.WithParent(goprocess.Background()),
		mappings: make(map[*mapping]struct{}),
	}
}

// Close shuts down all port mappings. NAT can no longer be used.
func (nat *NAT) Close() error {
	return nat.proc.Close()
}

// Process returns the nat's life-cycle manager, for making it listen
// to close signals.
func (nat *NAT) Process() goprocess.Process {
	return nat.proc
}

// Notifier is an object that assists NAT in notifying listeners.
// It is implemented using thirdparty/notifier
type Notifier struct {
	n notifier.Notifier
}

func (n *Notifier) notifyAll(notify func(n Notifiee)) {
	n.n.NotifyAll(func(n notifier.Notifiee) {
		notify(n.(Notifiee))
	})
}

// Notify signs up notifiee to listen to NAT events.
func (n *Notifier) Notify(notifiee Notifiee) {
	n.n.Notify(n)
}

// StopNotify stops signaling events to notifiee.
func (n *Notifier) StopNotify(notifiee Notifiee) {
	n.n.StopNotify(notifiee)
}

// Notifiee is an interface objects must implement to listen to NAT events.
type Notifiee interface {

	// Called every time a successful mapping happens
	// Warning: the port mapping may have changed. If that is the
	// case, both MappingSuccess and MappingChanged are called.
	MappingSuccess(nat *NAT, m Mapping)

	// Called when mapping a port succeeds, but the mapping is
	// with a different port than an earlier success.
	MappingChanged(nat *NAT, m Mapping, oldport int, newport int)

	// Called when a port mapping fails. NAT will continue attempting after
	// the next period. To stop trying, use: mapping.Close(). After this failure,
	// mapping.ExternalPort() will be zero, and nat.ExternalAddrs() will not
	// return the address for this mapping. With luck, the next attempt will
	// succeed, without the client needing to do anything.
	MappingFailed(nat *NAT, m Mapping, oldport int, err error)
}

// Mapping represents a port mapping in a NAT.
type Mapping interface {
	// NAT returns the NAT object this Mapping belongs to.
	NAT() *NAT

	// Protocol returns the protocol of this port mapping. This is either
	// "tcp" or "udp" as no other protocols are likely to be NAT-supported.
	Protocol() string

	// InternalPort returns the internal device port. Mapping will continue to
	// try to map InternalPort() to an external facing port.
	InternalPort() int

	// ExternalPort returns the external facing port. If the mapping is not
	// established, port will be 0
	ExternalPort() int

	// InternalAddr returns the internal address.
	InternalAddr() ma.Multiaddr

	// ExternalAddr returns the external facing address. If the mapping is not
	// established, addr will be nil, and and ErrNoMapping will be returned.
	ExternalAddr() (addr ma.Multiaddr, err error)

	// Close closes the port mapping
	Close() error
}

// keeps republishing
type mapping struct {
	sync.Mutex // guards all fields

	nat       *NAT
	proto     string
	intport   int
	extport   int
	permanent bool
	intaddr   ma.Multiaddr
	proc      goprocess.Process

	cached    ma.Multiaddr
	cacheTime time.Time
	cacheLk   sync.Mutex
}

func (m *mapping) NAT() *NAT {
	m.Lock()
	defer m.Unlock()
	return m.nat
}

func (m *mapping) Protocol() string {
	m.Lock()
	defer m.Unlock()
	return m.proto
}

func (m *mapping) InternalPort() int {
	m.Lock()
	defer m.Unlock()
	return m.intport
}

func (m *mapping) ExternalPort() int {
	m.Lock()
	defer m.Unlock()
	return m.extport
}

func (m *mapping) setExternalPort(p int) {
	m.Lock()
	defer m.Unlock()
	m.extport = p
}

func (m *mapping) setPermanent(p bool) {
	m.Lock()
	defer m.Unlock()
	m.permanent = p
}

func (m *mapping) isPermanent() bool {
	m.Lock()
	defer m.Unlock()
	return m.permanent
}

func (m *mapping) InternalAddr() ma.Multiaddr {
	m.Lock()
	defer m.Unlock()
	return m.intaddr
}

func (m *mapping) ExternalAddr() (ma.Multiaddr, error) {
	m.cacheLk.Lock()
	ctime := m.cacheTime
	cval := m.cached
	m.cacheLk.Unlock()
	if time.Since(ctime) < CacheTime {
		return cval, nil
	}

	if m.ExternalPort() == 0 { // dont even try right now.
		return nil, ErrNoMapping
	}

	ip, err := m.nat.nat.GetExternalAddress()
	if err != nil {
		return nil, err
	}

	ipmaddr, err := manet.FromIP(ip)
	if err != nil {
		return nil, fmt.Errorf("error parsing ip")
	}

	// call m.ExternalPort again, as mapping may have changed under our feet. (tocttou)
	extport := m.ExternalPort()
	if extport == 0 {
		return nil, ErrNoMapping
	}

	tcp, err := ma.NewMultiaddr(fmt.Sprintf("/%s/%d", m.Protocol(), extport))
	if err != nil {
		return nil, err
	}

	maddr2 := ipmaddr.Encapsulate(tcp)

	m.cacheLk.Lock()
	m.cached = maddr2
	m.cacheTime = time.Now()
	m.cacheLk.Unlock()
	return maddr2, nil
}

func (m *mapping) Close() error {
	return m.proc.Close()
}

// Mappings returns a slice of all NAT mappings
func (nat *NAT) Mappings() []Mapping {
	nat.mappingmu.Lock()
	maps2 := make([]Mapping, 0, len(nat.mappings))
	for m := range nat.mappings {
		maps2 = append(maps2, m)
	}
	nat.mappingmu.Unlock()
	return maps2
}

func (nat *NAT) addMapping(m *mapping) {
	// make mapping automatically close when nat is closed.
	nat.proc.AddChild(m.proc)

	nat.mappingmu.Lock()
	nat.mappings[m] = struct{}{}
	nat.mappingmu.Unlock()
}

func (nat *NAT) rmMapping(m *mapping) {
	nat.mappingmu.Lock()
	delete(nat.mappings, m)
	nat.mappingmu.Unlock()
}

// NewMapping attemps to construct a mapping on protocol and internal port
// It will also periodically renew the mapping until the returned Mapping
// -- or its parent NAT -- is Closed.
//
// May not succeed, and mappings may change over time;
// NAT devices may not respect our port requests, and even lie.
// Clients should not store the mapped results, but rather always
// poll our object for the latest mappings.
func (nat *NAT) NewMapping(maddr ma.Multiaddr) (Mapping, error) {
	if nat == nil {
		return nil, fmt.Errorf("no nat available")
	}

	network, addr, err := manet.DialArgs(maddr)
	if err != nil {
		return nil, fmt.Errorf("DialArgs failed on addr:", maddr.String())
	}

	switch network {
	case "tcp", "tcp4", "tcp6":
		network = "tcp"
	case "udp", "udp4", "udp6":
		network = "udp"
	default:
		return nil, fmt.Errorf("transport not supported by NAT: %s", network)
	}

	intports := strings.Split(addr, ":")[1]
	intport, err := strconv.Atoi(intports)
	if err != nil {
		return nil, err
	}

	m := &mapping{
		nat:     nat,
		proto:   network,
		intport: intport,
		intaddr: maddr,
	}
	m.proc = goprocess.WithTeardown(func() error {
		nat.rmMapping(m)
		return nil
	})
	nat.addMapping(m)

	m.proc.AddChild(periodic.Every(MappingDuration/3, func(worker goprocess.Process) {
		nat.establishMapping(m)
	}))

	// do it once synchronously, so first mapping is done right away, and before exiting,
	// allowing users -- in the optimistic case -- to use results right after.
	nat.establishMapping(m)
	return m, nil
}

func (nat *NAT) establishMapping(m *mapping) {
	oldport := m.ExternalPort()
	if oldport != 0 && m.isPermanent() {
		// mapping was already established and it is permanent
		return
	}

	log.Debugf("Attempting port map: %s/%d", m.Protocol(), m.InternalPort())
	permanent := false

	newport, err := nat.nat.AddPortMapping(m.Protocol(), m.InternalPort(), "libp2p", MappingDuration)

	if err != nil {
		// Some hardware does not support mappings with timeout, so try that
		newport, err = nat.nat.AddPortMapping(m.Protocol(), m.InternalPort(), "libp2p", 0)
		permanent = (err == nil)
	}

	failure := func() {
		m.setExternalPort(0) // clear mapping
		// TODO: log.Event
		log.Warningf("failed to establish port mapping: %s", err)
		nat.Notifier.notifyAll(func(n Notifiee) {
			n.MappingFailed(nat, m, oldport, err)
		})

		// we do not close if the mapping failed,
		// because it may work again next time.
	}

	if err != nil || newport == 0 {
		failure()
		return
	}

	m.setPermanent(permanent)
	m.setExternalPort(newport)
	ext, err := m.ExternalAddr()
	if err != nil {
		log.Debugf("NAT Mapping addr error: %s %s", m.InternalAddr(), err)
		failure()
		return
	}

	log.Debugf("NAT Mapping: %s --> %s", m.InternalAddr(), ext)
	if oldport != 0 && newport != oldport {
		log.Debugf("failed to renew same port mapping: ch %d -> %d", oldport, newport)
		nat.Notifier.notifyAll(func(n Notifiee) {
			n.MappingChanged(nat, m, oldport, newport)
		})
	}

	nat.Notifier.notifyAll(func(n Notifiee) {
		n.MappingSuccess(nat, m)
	})
}

// PortMapAddrs attempts to open (and continue to keep open)
// port mappings for given addrs. This function blocks until
// all  addresses have been tried. This allows clients to
// retrieve results immediately after:
//
//   nat.PortMapAddrs(addrs)
//   mapped := nat.ExternalAddrs()
//
// Some may not succeed, and mappings may change over time;
// NAT devices may not respect our port requests, and even lie.
// Clients should not store the mapped results, but rather always
// poll our object for the latest mappings.
func (nat *NAT) PortMapAddrs(addrs []ma.Multiaddr) {
	// spin off addr mappings independently.
	var wg sync.WaitGroup
	for _, addr := range addrs {
		// do all of them concurrently
		wg.Add(1)
		go func() {
			defer wg.Done()
			nat.NewMapping(addr)
		}()
	}
	wg.Wait()
}

// MappedAddrs returns address mappings NAT believes have been
// successfully established. Unsuccessful mappings are nil. This is:
//
// 		map[internalAddr]externalAddr
//
// This set of mappings _may not_ be correct, as NAT devices are finicky.
// Consider this with _best effort_ semantics.
func (nat *NAT) MappedAddrs() map[ma.Multiaddr]ma.Multiaddr {

	mappings := nat.Mappings()
	addrmap := make(map[ma.Multiaddr]ma.Multiaddr, len(mappings))

	for _, m := range mappings {
		i := m.InternalAddr()
		e, err := m.ExternalAddr()
		if err != nil {
			addrmap[i] = nil
		} else {
			addrmap[i] = e
		}
	}
	return addrmap
}

// ExternalAddrs returns a list of addresses that NAT believes have
// been successfully established. Unsuccessful mappings are omitted,
// so nat.ExternalAddrs() may return less addresses than nat.InternalAddrs().
// To see which addresses are mapped, use nat.MappedAddrs().
//
// This set of mappings _may not_ be correct, as NAT devices are finicky.
// Consider this with _best effort_ semantics.
func (nat *NAT) ExternalAddrs() []ma.Multiaddr {
	mappings := nat.Mappings()
	addrs := make([]ma.Multiaddr, 0, len(mappings))
	for _, m := range mappings {
		a, err := m.ExternalAddr()
		if err != nil {
			continue // this mapping not currently successful.
		}
		addrs = append(addrs, a)
	}
	return addrs
}
<commit_msg>nat: do not shortcircuit permanent mappings<commit_after>package nat

import (
	"errors"
	"fmt"
	"strconv"
	"strings"
	"sync"
	"time"

	nat "github.com/fd/go-nat"
	logging "github.com/ipfs/go-log"
	ma "github.com/jbenet/go-multiaddr"
	manet "github.com/jbenet/go-multiaddr-net"
	goprocess "github.com/jbenet/goprocess"
	periodic "github.com/jbenet/goprocess/periodic"
	notifier "github.com/whyrusleeping/go-notifier"
)

var (
	// ErrNoMapping signals no mapping exists for an address
	ErrNoMapping = errors.New("mapping not established")
)

var log = logging.Logger("nat")

// MappingDuration is a default port mapping duration.
// Port mappings are renewed every (MappingDuration / 3)
const MappingDuration = time.Second * 60

// CacheTime is the time a mapping will cache an external address for
const CacheTime = time.Second * 15

// DiscoverNAT looks for a NAT device in the network and
// returns an object that can manage port mappings.
func DiscoverNAT() *NAT {
	nat, err := nat.DiscoverGateway()
	if err != nil {
		log.Debug("DiscoverGateway error:", err)
		return nil
	}
	addr, err := nat.GetDeviceAddress()
	if err != nil {
		log.Debug("DiscoverGateway address error:", err)
	} else {
		log.Debug("DiscoverGateway address:", addr)
	}
	return newNAT(nat)
}

// NAT is an object that manages address port mappings in
// NATs (Network Address Translators). It is a long-running
// service that will periodically renew port mappings,
// and keep an up-to-date list of all the external addresses.
type NAT struct {
	nat  nat.NAT
	proc goprocess.Process // manages nat mappings lifecycle

	mappingmu sync.RWMutex // guards mappings
	mappings  map[*mapping]struct{}

	Notifier
}

func newNAT(realNAT nat.NAT) *NAT {
	return &NAT{
		nat:      realNAT,
		proc:     goprocess.WithParent(goprocess.Background()),
		mappings: make(map[*mapping]struct{}),
	}
}

// Close shuts down all port mappings. NAT can no longer be used.
func (nat *NAT) Close() error {
	return nat.proc.Close()
}

// Process returns the nat's life-cycle manager, for making it listen
// to close signals.
func (nat *NAT) Process() goprocess.Process {
	return nat.proc
}

// Notifier is an object that assists NAT in notifying listeners.
// It is implemented using thirdparty/notifier
type Notifier struct {
	n notifier.Notifier
}

func (n *Notifier) notifyAll(notify func(n Notifiee)) {
	n.n.NotifyAll(func(n notifier.Notifiee) {
		notify(n.(Notifiee))
	})
}

// Notify signs up notifiee to listen to NAT events.
func (n *Notifier) Notify(notifiee Notifiee) {
	n.n.Notify(n)
}

// StopNotify stops signaling events to notifiee.
func (n *Notifier) StopNotify(notifiee Notifiee) {
	n.n.StopNotify(notifiee)
}

// Notifiee is an interface objects must implement to listen to NAT events.
type Notifiee interface {

	// Called every time a successful mapping happens
	// Warning: the port mapping may have changed. If that is the
	// case, both MappingSuccess and MappingChanged are called.
	MappingSuccess(nat *NAT, m Mapping)

	// Called when mapping a port succeeds, but the mapping is
	// with a different port than an earlier success.
	MappingChanged(nat *NAT, m Mapping, oldport int, newport int)

	// Called when a port mapping fails. NAT will continue attempting after
	// the next period. To stop trying, use: mapping.Close(). After this failure,
	// mapping.ExternalPort() will be zero, and nat.ExternalAddrs() will not
	// return the address for this mapping. With luck, the next attempt will
	// succeed, without the client needing to do anything.
	MappingFailed(nat *NAT, m Mapping, oldport int, err error)
}

// Mapping represents a port mapping in a NAT.
type Mapping interface {
	// NAT returns the NAT object this Mapping belongs to.
	NAT() *NAT

	// Protocol returns the protocol of this port mapping. This is either
	// "tcp" or "udp" as no other protocols are likely to be NAT-supported.
	Protocol() string

	// InternalPort returns the internal device port. Mapping will continue to
	// try to map InternalPort() to an external facing port.
	InternalPort() int

	// ExternalPort returns the external facing port. If the mapping is not
	// established, port will be 0
	ExternalPort() int

	// InternalAddr returns the internal address.
	InternalAddr() ma.Multiaddr

	// ExternalAddr returns the external facing address. If the mapping is not
	// established, addr will be nil, and and ErrNoMapping will be returned.
	ExternalAddr() (addr ma.Multiaddr, err error)

	// Close closes the port mapping
	Close() error
}

// keeps republishing
type mapping struct {
	sync.Mutex // guards all fields

	nat       *NAT
	proto     string
	intport   int
	extport   int
	permanent bool
	intaddr   ma.Multiaddr
	proc      goprocess.Process

	comment string

	cached    ma.Multiaddr
	cacheTime time.Time
	cacheLk   sync.Mutex
}

func (m *mapping) NAT() *NAT {
	m.Lock()
	defer m.Unlock()
	return m.nat
}

func (m *mapping) Protocol() string {
	m.Lock()
	defer m.Unlock()
	return m.proto
}

func (m *mapping) InternalPort() int {
	m.Lock()
	defer m.Unlock()
	return m.intport
}

func (m *mapping) ExternalPort() int {
	m.Lock()
	defer m.Unlock()
	return m.extport
}

func (m *mapping) setExternalPort(p int) {
	m.Lock()
	defer m.Unlock()
	m.extport = p
}

func (m *mapping) InternalAddr() ma.Multiaddr {
	m.Lock()
	defer m.Unlock()
	return m.intaddr
}

func (m *mapping) ExternalAddr() (ma.Multiaddr, error) {
	m.cacheLk.Lock()
	ctime := m.cacheTime
	cval := m.cached
	m.cacheLk.Unlock()
	if time.Since(ctime) < CacheTime {
		return cval, nil
	}

	if m.ExternalPort() == 0 { // dont even try right now.
		return nil, ErrNoMapping
	}

	ip, err := m.nat.nat.GetExternalAddress()
	if err != nil {
		return nil, err
	}

	ipmaddr, err := manet.FromIP(ip)
	if err != nil {
		return nil, fmt.Errorf("error parsing ip")
	}

	// call m.ExternalPort again, as mapping may have changed under our feet. (tocttou)
	extport := m.ExternalPort()
	if extport == 0 {
		return nil, ErrNoMapping
	}

	tcp, err := ma.NewMultiaddr(fmt.Sprintf("/%s/%d", m.Protocol(), extport))
	if err != nil {
		return nil, err
	}

	maddr2 := ipmaddr.Encapsulate(tcp)

	m.cacheLk.Lock()
	m.cached = maddr2
	m.cacheTime = time.Now()
	m.cacheLk.Unlock()
	return maddr2, nil
}

func (m *mapping) Close() error {
	return m.proc.Close()
}

// Mappings returns a slice of all NAT mappings
func (nat *NAT) Mappings() []Mapping {
	nat.mappingmu.Lock()
	maps2 := make([]Mapping, 0, len(nat.mappings))
	for m := range nat.mappings {
		maps2 = append(maps2, m)
	}
	nat.mappingmu.Unlock()
	return maps2
}

func (nat *NAT) addMapping(m *mapping) {
	// make mapping automatically close when nat is closed.
	nat.proc.AddChild(m.proc)

	nat.mappingmu.Lock()
	nat.mappings[m] = struct{}{}
	nat.mappingmu.Unlock()
}

func (nat *NAT) rmMapping(m *mapping) {
	nat.mappingmu.Lock()
	delete(nat.mappings, m)
	nat.mappingmu.Unlock()
}

// NewMapping attemps to construct a mapping on protocol and internal port
// It will also periodically renew the mapping until the returned Mapping
// -- or its parent NAT -- is Closed.
//
// May not succeed, and mappings may change over time;
// NAT devices may not respect our port requests, and even lie.
// Clients should not store the mapped results, but rather always
// poll our object for the latest mappings.
func (nat *NAT) NewMapping(maddr ma.Multiaddr) (Mapping, error) {
	if nat == nil {
		return nil, fmt.Errorf("no nat available")
	}

	network, addr, err := manet.DialArgs(maddr)
	if err != nil {
		return nil, fmt.Errorf("DialArgs failed on addr:", maddr.String())
	}

	switch network {
	case "tcp", "tcp4", "tcp6":
		network = "tcp"
	case "udp", "udp4", "udp6":
		network = "udp"
	default:
		return nil, fmt.Errorf("transport not supported by NAT: %s", network)
	}

	intports := strings.Split(addr, ":")[1]
	intport, err := strconv.Atoi(intports)
	if err != nil {
		return nil, err
	}

	m := &mapping{
		nat:     nat,
		proto:   network,
		intport: intport,
		intaddr: maddr,
	}
	m.proc = goprocess.WithTeardown(func() error {
		nat.rmMapping(m)
		return nil
	})
	nat.addMapping(m)

	m.proc.AddChild(periodic.Every(MappingDuration/3, func(worker goprocess.Process) {
		nat.establishMapping(m)
	}))

	// do it once synchronously, so first mapping is done right away, and before exiting,
	// allowing users -- in the optimistic case -- to use results right after.
	nat.establishMapping(m)
	return m, nil
}

func (nat *NAT) establishMapping(m *mapping) {
	oldport := m.ExternalPort()

	log.Debugf("Attempting port map: %s/%d", m.Protocol(), m.InternalPort())
	comment := "libp2p"
	if m.comment != "" {
		comment = "libp2p-" + m.comment
	}

	newport, err := nat.nat.AddPortMapping(m.Protocol(), m.InternalPort(), comment, MappingDuration)
	if err != nil {
		// Some hardware does not support mappings with timeout, so try that
		newport, err = nat.nat.AddPortMapping(m.Protocol(), m.InternalPort(), comment, 0)
	}

	failure := func() {
		m.setExternalPort(0) // clear mapping
		// TODO: log.Event
		log.Warningf("failed to establish port mapping: %s", err)
		nat.Notifier.notifyAll(func(n Notifiee) {
			n.MappingFailed(nat, m, oldport, err)
		})

		// we do not close if the mapping failed,
		// because it may work again next time.
	}

	if err != nil || newport == 0 {
		failure()
		return
	}

	m.setExternalPort(newport)
	ext, err := m.ExternalAddr()
	if err != nil {
		log.Debugf("NAT Mapping addr error: %s %s", m.InternalAddr(), err)
		failure()
		return
	}

	log.Debugf("NAT Mapping: %s --> %s", m.InternalAddr(), ext)
	if oldport != 0 && newport != oldport {
		log.Debugf("failed to renew same port mapping: ch %d -> %d", oldport, newport)
		nat.Notifier.notifyAll(func(n Notifiee) {
			n.MappingChanged(nat, m, oldport, newport)
		})
	}

	nat.Notifier.notifyAll(func(n Notifiee) {
		n.MappingSuccess(nat, m)
	})
}

// PortMapAddrs attempts to open (and continue to keep open)
// port mappings for given addrs. This function blocks until
// all  addresses have been tried. This allows clients to
// retrieve results immediately after:
//
//   nat.PortMapAddrs(addrs)
//   mapped := nat.ExternalAddrs()
//
// Some may not succeed, and mappings may change over time;
// NAT devices may not respect our port requests, and even lie.
// Clients should not store the mapped results, but rather always
// poll our object for the latest mappings.
func (nat *NAT) PortMapAddrs(addrs []ma.Multiaddr) {
	// spin off addr mappings independently.
	var wg sync.WaitGroup
	for _, addr := range addrs {
		// do all of them concurrently
		wg.Add(1)
		go func() {
			defer wg.Done()
			nat.NewMapping(addr)
		}()
	}
	wg.Wait()
}

// MappedAddrs returns address mappings NAT believes have been
// successfully established. Unsuccessful mappings are nil. This is:
//
// 		map[internalAddr]externalAddr
//
// This set of mappings _may not_ be correct, as NAT devices are finicky.
// Consider this with _best effort_ semantics.
func (nat *NAT) MappedAddrs() map[ma.Multiaddr]ma.Multiaddr {

	mappings := nat.Mappings()
	addrmap := make(map[ma.Multiaddr]ma.Multiaddr, len(mappings))

	for _, m := range mappings {
		i := m.InternalAddr()
		e, err := m.ExternalAddr()
		if err != nil {
			addrmap[i] = nil
		} else {
			addrmap[i] = e
		}
	}
	return addrmap
}

// ExternalAddrs returns a list of addresses that NAT believes have
// been successfully established. Unsuccessful mappings are omitted,
// so nat.ExternalAddrs() may return less addresses than nat.InternalAddrs().
// To see which addresses are mapped, use nat.MappedAddrs().
//
// This set of mappings _may not_ be correct, as NAT devices are finicky.
// Consider this with _best effort_ semantics.
func (nat *NAT) ExternalAddrs() []ma.Multiaddr {
	mappings := nat.Mappings()
	addrs := make([]ma.Multiaddr, 0, len(mappings))
	for _, m := range mappings {
		a, err := m.ExternalAddr()
		if err != nil {
			continue // this mapping not currently successful.
		}
		addrs = append(addrs, a)
	}
	return addrs
}
<|endoftext|>


<commit_before>// Copyright (c) 2010, Suryandaru Triandana. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This package implements the Rabbit encryption algorithm as defined in eSTREAM portfolio.
package rabbit

// This Go implementation is derived in part from the reference
// ANSI C implementation, which carries the following notice:
//
//	File name: rabbit.c
//
//	Source file for reference C version of the Rabbit stream cipher.
//
//	For further documentation, see "Rabbit Stream Cipher, Algorithm
//	Specification" which can be found at http://www.cryptico.com/.
//
//	This source code is for little-endian processors (e.g. x86).
//
//	Copyright (C) Cryptico ApS. All rights reserved.
//
//	YOU SHOULD CAREFULLY READ THIS LEGAL NOTICE BEFORE USING THIS SOFTWARE.
//
//	This software is developed by Cryptico ApS and/or its suppliers. It is
//	free for commercial and non-commercial use.
//
//	Cryptico ApS shall not in any way be liable for any use or export/import
//	of this software. The software is provided "as is" without any express or
//	implied warranty.
//
//	Cryptico, CryptiCore, the Cryptico logo and "Re-thinking encryption" are
//	either trademarks or registered trademarks of Cryptico ApS.

import (
	"os"
	"strconv"
)

// A Cipher is an instance of Rabbit encryption using a particular key.
type Cipher struct {
	x, c, cx, cc [8]uint32
	carry, ccarry bool
	r []byte
}

type KeySizeError struct {
	t, sz int
}

func (k *KeySizeError) String() string {
	switch(k.t) {
	case 1:
		return "crypto/rabbit: invalid key size " + strconv.Itoa(int(k.sz))
	case 2:
		return "crypto/rabbit: invalid iv size " + strconv.Itoa(int(k.sz))
	}
	return "crypto/rabbit: unknown key error type"
}

func rotl(v, n uint32) uint32 {
	return v<<n | v>>(32-n)
}

func rabbitCalcG(x uint32) uint32 {
	var a, b uint32
	a = x&0xFFFF;
	b = x>>16;
	return ((((a*a)>>17 + a*b)>>15) + b*b)^(x*x)
}

func booltoi(b bool) uint32 {
	if b {
		return 1
	}
	return 0
}

func (c *Cipher) rabbitNext() {
	var c0, c1, c2, c3, c4, c5, c6, c7 uint32

	c0, c1, c2, c3 = c.c[0], c.c[1], c.c[2], c.c[3]
	c4, c5, c6, c7 = c.c[4], c.c[5], c.c[6], c.c[7]

	c0 = c0 + 0x4D34D34D + booltoi(c.carry)
	c1 = c1 + 0xD34D34D3 + booltoi(c0 < c.c[0])
	c2 = c2 + 0x34D34D34 + booltoi(c1 < c.c[1])
	c3 = c3 + 0x4D34D34D + booltoi(c2 < c.c[2])
	c4 = c4 + 0xD34D34D3 + booltoi(c3 < c.c[3])
	c5 = c5 + 0x34D34D34 + booltoi(c4 < c.c[4])
	c6 = c6 + 0x4D34D34D + booltoi(c5 < c.c[5])
	c7 = c7 + 0xD34D34D3 + booltoi(c6 < c.c[6])
	c.carry = (c7 < c.c[7])

	g0 := rabbitCalcG(c.x[0] + c0)
	g1 := rabbitCalcG(c.x[1] + c1)
	g2 := rabbitCalcG(c.x[2] + c2)
	g3 := rabbitCalcG(c.x[3] + c3)
	g4 := rabbitCalcG(c.x[4] + c4)
	g5 := rabbitCalcG(c.x[5] + c5)
	g6 := rabbitCalcG(c.x[6] + c6)
	g7 := rabbitCalcG(c.x[7] + c7)

	c.x[0] = g0 + rotl(g7,16) + rotl(g6, 16)
	c.x[1] = g1 + rotl(g0, 8) + g7
	c.x[2] = g2 + rotl(g1,16) + rotl(g0, 16)
	c.x[3] = g3 + rotl(g2, 8) + g1
	c.x[4] = g4 + rotl(g3,16) + rotl(g2, 16)
	c.x[5] = g5 + rotl(g4, 8) + g3
	c.x[6] = g6 + rotl(g5,16) + rotl(g4, 16)
	c.x[7] = g7 + rotl(g6, 8) + g5

	c.c[0], c.c[1], c.c[2], c.c[3] = c0, c1, c2, c3
	c.c[4], c.c[5], c.c[6], c.c[7] = c4, c5, c6, c7
}

func (c *Cipher) rabbitGen(buf *[16]byte) {
	c.rabbitNext()
	var d0, d1, d2, d3 uint32
	d0 = c.x[0] ^ (c.x[5]>>16 ^ c.x[3]<<16)
	d1 = c.x[2] ^ (c.x[7]>>16 ^ c.x[5]<<16)
	d2 = c.x[4] ^ (c.x[1]>>16 ^ c.x[7]<<16)
	d3 = c.x[6] ^ (c.x[3]>>16 ^ c.x[1]<<16)
	buf[ 0], buf[ 1], buf[ 2], buf[ 3] = byte(d0), byte(d0>>8), byte(d0>>16), byte(d0>>24)
	buf[ 4], buf[ 5], buf[ 6], buf[ 7] = byte(d1), byte(d1>>8), byte(d1>>16), byte(d1>>24)
	buf[ 8], buf[ 9], buf[10], buf[11] = byte(d2), byte(d2>>8), byte(d2>>16), byte(d2>>24)
	buf[12], buf[13], buf[14], buf[15] = byte(d3), byte(d3>>8), byte(d3>>16), byte(d3>>24)
}

// NewCipher creates and returns a Cipher.
// Rabbit key, must be 16 bytes.
func NewCipher(key []byte) (*Cipher, os.Error) {
	k := len(key)
	if k != 16 {
		return nil, &KeySizeError{1, k}
	}
	var c Cipher

	var k0, k1, k2, k3 uint32
	k0 = uint32(key[ 0]) | uint32(key[ 1])<<8 | uint32(key[ 2])<<16 | uint32(key[ 3])<<24
	k1 = uint32(key[ 4]) | uint32(key[ 5])<<8 | uint32(key[ 6])<<16 | uint32(key[ 7])<<24
	k2 = uint32(key[ 8]) | uint32(key[ 9])<<8 | uint32(key[10])<<16 | uint32(key[11])<<24
	k3 = uint32(key[12]) | uint32(key[13])<<8 | uint32(key[14])<<16 | uint32(key[15])<<24

	c.x[0] = k0
	c.x[2] = k1
	c.x[4] = k2
	c.x[6] = k3
	c.x[1] = k3<<16 | k2>>16
	c.x[3] = k0<<16 | k3>>16
	c.x[5] = k1<<16 | k0>>16
	c.x[7] = k2<<16 | k1>>16

	c.c[0] = rotl(k2, 16)
	c.c[2] = rotl(k3, 16)
	c.c[4] = rotl(k0, 16)
	c.c[6] = rotl(k1, 16)
	c.c[1] = (k0&0xFFFF0000) | (k1&0xFFFF)
	c.c[3] = (k1&0xFFFF0000) | (k2&0xFFFF)
	c.c[5] = (k2&0xFFFF0000) | (k3&0xFFFF)
	c.c[7] = (k3&0xFFFF0000) | (k0&0xFFFF)

	c.carry = false

	for i := 0; i < 4; i++ {
		c.rabbitNext()
	}

	for i := range c.c {
		c.c[i] ^= c.x[(i+4)&0x7]
	}

	for i := range c.c {
		c.cx[i] = c.x[i]
		c.cc[i] = c.c[i]
	}
	c.ccarry = c.carry

	return &c, nil
}

// SetupIV will setup Initialization vector.
// Rabbit iv, must be 8 bytes.
func (c *Cipher) SetupIV(iv []byte) os.Error {
	k := len(iv)
	if k != 8 {
		return &KeySizeError{2, k}
	}

	var d0, d1, d2, d3 uint32
	d0 = uint32(iv[0]) | uint32(iv[1])<<8 | uint32(iv[2])<<16 | uint32(iv[3])<<24
	d2 = uint32(iv[4]) | uint32(iv[5])<<8 | uint32(iv[6])<<16 | uint32(iv[7])<<24
	d1 = d0>>16 | (d2&0xFFFF0000)
	d3 = d2<<16 | (d0&0x0000FFFF)

	c.c[0] = c.cc[0] ^ d0
	c.c[1] = c.cc[1] ^ d1
	c.c[2] = c.cc[2] ^ d2
	c.c[3] = c.cc[3] ^ d3
	c.c[4] = c.cc[4] ^ d0
	c.c[5] = c.cc[5] ^ d1
	c.c[6] = c.cc[6] ^ d2
	c.c[7] = c.cc[7] ^ d3

	for i := range c.x {
		c.x[i] = c.cx[i]
	}
	c.carry = c.ccarry

	for i := 0; i < 4; i++ {
		c.rabbitNext()
	}

	return nil
}

// ProcessStream will encrypt or decrypt given buffer.
func (c *Cipher) ProcessStream(buf []byte) {
	l := len(buf)
	i := 0
	if m := len(c.r); m > 0 {
		for ; i < m && i < l; i++ {
			buf[i] ^= c.r[i]
		}
		c.r = nil
	}
	for i < l {
		c.rabbitNext()

		if n := l - i; n >= 16 {
			o0 := c.x[0] ^ (c.x[5]>>16 ^ c.x[3]<<16)
			o1 := c.x[2] ^ (c.x[7]>>16 ^ c.x[5]<<16)
			o2 := c.x[4] ^ (c.x[1]>>16 ^ c.x[7]<<16)
			o3 := c.x[6] ^ (c.x[3]>>16 ^ c.x[1]<<16)
			buf[i + 0] ^= byte(o0     )
			buf[i + 1] ^= byte(o0 >> 8)
			buf[i + 2] ^= byte(o0 >>16)
			buf[i + 3] ^= byte(o0 >>24)
			buf[i + 4] ^= byte(o1     )
			buf[i + 5] ^= byte(o1 >> 8)
			buf[i + 6] ^= byte(o1 >>16)
			buf[i + 7] ^= byte(o1 >>24)
			buf[i + 8] ^= byte(o2     )
			buf[i + 9] ^= byte(o2 >> 8)
			buf[i +10] ^= byte(o2 >>16)
			buf[i +11] ^= byte(o2 >>24)
			buf[i +12] ^= byte(o3     )
			buf[i +13] ^= byte(o3 >> 8)
			buf[i +14] ^= byte(o3 >>16)
			buf[i +15] ^= byte(o3 >>24)
			i += 16
		} else {
			for b, j, z, f := buf, 0, c.x[0] ^ (c.x[5]>>16 ^ c.x[3]<<16), false; j < 4; j++ {
				for k := uint32(0); k < 4; k++ {
					b[i] ^= byte(z>>(k*8))
					if i++; f == false && i >= l {
						l = (3 - j)*4 + (3 - int(k))
						if l == 0 { return }
						c.r = make([]byte, l)
						b, i, f = c.r, 0, true
					}
				}
				switch(j) {
				case  0: z = c.x[2] ^ (c.x[7]>>16 ^ c.x[5]<<16)
				case  1: z = c.x[4] ^ (c.x[1]>>16 ^ c.x[7]<<16)
				case  2: z = c.x[6] ^ (c.x[3]>>16 ^ c.x[1]<<16)
				}
			}
		}
	}
}

// ResetCipher reset cipher round to original state. Initialization vector will be erased.
func (c *Cipher) ResetCipher() {
	for i := range c.c {
		c.c[i] = c.cc[i]
	}
	for i := range c.x {
		c.x[i] = c.cx[i]
	}
	c.carry = c.ccarry
}

// Reset zeros the key data so that it will no longer appear in the
// process's memory.
func (c *Cipher) Reset() {
	for i := range c.x {
		c.x[i], c.c[i], c.cx[i], c.cc[i] = 0, 0, 0, 0
	}
	c.carry, c.carry = false, false
}

<commit_msg>crypto/rabbit: fix SetupIV and ResetCipher that not reset cipher correctly crypto/rabbit: fix os.Error implementation crypto/rabbit: remove unused rabbitGen method<commit_after>// Copyright (c) 2010, Suryandaru Triandana. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This package implements the Rabbit encryption algorithm as defined in eSTREAM portfolio.
package rabbit

// This Go implementation is derived in part from the reference
// ANSI C implementation, which carries the following notice:
//
//	File name: rabbit.c
//
//	Source file for reference C version of the Rabbit stream cipher.
//
//	For further documentation, see "Rabbit Stream Cipher, Algorithm
//	Specification" which can be found at http://www.cryptico.com/.
//
//	This source code is for little-endian processors (e.g. x86).
//
//	Copyright (C) Cryptico ApS. All rights reserved.
//
//	YOU SHOULD CAREFULLY READ THIS LEGAL NOTICE BEFORE USING THIS SOFTWARE.
//
//	This software is developed by Cryptico ApS and/or its suppliers. It is
//	free for commercial and non-commercial use.
//
//	Cryptico ApS shall not in any way be liable for any use or export/import
//	of this software. The software is provided "as is" without any express or
//	implied warranty.
//
//	Cryptico, CryptiCore, the Cryptico logo and "Re-thinking encryption" are
//	either trademarks or registered trademarks of Cryptico ApS.

import (
	"os"
	"strconv"
)

// A Cipher is an instance of Rabbit encryption using a particular key.
type Cipher struct {
	x, c, cx, cc [8]uint32
	carry, ccarry bool
	r []byte
}

type KeySizeError int

func (k KeySizeError) String() string {
	return "crypto/rabbit: invalid key size " + strconv.Itoa(int(k))
}

type IVSizeError int

func (k IVSizeError) String() string {
	return "crypto/rabbit: invalid iv size " + strconv.Itoa(int(k))
}

func rotl(v, n uint32) uint32 {
	return v<<n | v>>(32-n)
}

func rabbitCalcG(x uint32) uint32 {
	var a, b uint32
	a = x&0xFFFF;
	b = x>>16;
	return ((((a*a)>>17 + a*b)>>15) + b*b)^(x*x)
}

func booltoi(b bool) uint32 {
	if b {
		return 1
	}
	return 0
}

func (c *Cipher) rabbitNext() {
	var c0, c1, c2, c3, c4, c5, c6, c7 uint32

	c0, c1, c2, c3 = c.c[0], c.c[1], c.c[2], c.c[3]
	c4, c5, c6, c7 = c.c[4], c.c[5], c.c[6], c.c[7]

	c0 = c0 + 0x4D34D34D + booltoi(c.carry)
	c1 = c1 + 0xD34D34D3 + booltoi(c0 < c.c[0])
	c2 = c2 + 0x34D34D34 + booltoi(c1 < c.c[1])
	c3 = c3 + 0x4D34D34D + booltoi(c2 < c.c[2])
	c4 = c4 + 0xD34D34D3 + booltoi(c3 < c.c[3])
	c5 = c5 + 0x34D34D34 + booltoi(c4 < c.c[4])
	c6 = c6 + 0x4D34D34D + booltoi(c5 < c.c[5])
	c7 = c7 + 0xD34D34D3 + booltoi(c6 < c.c[6])
	c.carry = (c7 < c.c[7])

	g0 := rabbitCalcG(c.x[0] + c0)
	g1 := rabbitCalcG(c.x[1] + c1)
	g2 := rabbitCalcG(c.x[2] + c2)
	g3 := rabbitCalcG(c.x[3] + c3)
	g4 := rabbitCalcG(c.x[4] + c4)
	g5 := rabbitCalcG(c.x[5] + c5)
	g6 := rabbitCalcG(c.x[6] + c6)
	g7 := rabbitCalcG(c.x[7] + c7)

	c.x[0] = g0 + rotl(g7,16) + rotl(g6, 16)
	c.x[1] = g1 + rotl(g0, 8) + g7
	c.x[2] = g2 + rotl(g1,16) + rotl(g0, 16)
	c.x[3] = g3 + rotl(g2, 8) + g1
	c.x[4] = g4 + rotl(g3,16) + rotl(g2, 16)
	c.x[5] = g5 + rotl(g4, 8) + g3
	c.x[6] = g6 + rotl(g5,16) + rotl(g4, 16)
	c.x[7] = g7 + rotl(g6, 8) + g5

	c.c[0], c.c[1], c.c[2], c.c[3] = c0, c1, c2, c3
	c.c[4], c.c[5], c.c[6], c.c[7] = c4, c5, c6, c7
}

// NewCipher creates and returns a Cipher.
// Rabbit key, must be 16 bytes.
func NewCipher(key []byte) (*Cipher, os.Error) {
	k := len(key)
	if k != 16 {
		return nil, KeySizeError(k)
	}
	var c Cipher

	var k0, k1, k2, k3 uint32
	k0 = uint32(key[ 0]) | uint32(key[ 1])<<8 | uint32(key[ 2])<<16 | uint32(key[ 3])<<24
	k1 = uint32(key[ 4]) | uint32(key[ 5])<<8 | uint32(key[ 6])<<16 | uint32(key[ 7])<<24
	k2 = uint32(key[ 8]) | uint32(key[ 9])<<8 | uint32(key[10])<<16 | uint32(key[11])<<24
	k3 = uint32(key[12]) | uint32(key[13])<<8 | uint32(key[14])<<16 | uint32(key[15])<<24

	c.x[0] = k0
	c.x[2] = k1
	c.x[4] = k2
	c.x[6] = k3
	c.x[1] = k3<<16 | k2>>16
	c.x[3] = k0<<16 | k3>>16
	c.x[5] = k1<<16 | k0>>16
	c.x[7] = k2<<16 | k1>>16

	c.c[0] = rotl(k2, 16)
	c.c[2] = rotl(k3, 16)
	c.c[4] = rotl(k0, 16)
	c.c[6] = rotl(k1, 16)
	c.c[1] = (k0&0xFFFF0000) | (k1&0xFFFF)
	c.c[3] = (k1&0xFFFF0000) | (k2&0xFFFF)
	c.c[5] = (k2&0xFFFF0000) | (k3&0xFFFF)
	c.c[7] = (k3&0xFFFF0000) | (k0&0xFFFF)

	c.carry = false

	for i := 0; i < 4; i++ {
		c.rabbitNext()
	}

	for i := range c.c {
		c.c[i] ^= c.x[(i+4)&0x7]
	}

	for i := range c.c {
		c.cx[i] = c.x[i]
		c.cc[i] = c.c[i]
	}
	c.ccarry = c.carry

	return &c, nil
}

// SetupIV will setup Initialization vector.
// Rabbit iv, must be 8 bytes.
func (c *Cipher) SetupIV(iv []byte) os.Error {
	k := len(iv)
	if k != 8 {
		return IVSizeError(k)
	}

	var d0, d1, d2, d3 uint32
	d0 = uint32(iv[0]) | uint32(iv[1])<<8 | uint32(iv[2])<<16 | uint32(iv[3])<<24
	d2 = uint32(iv[4]) | uint32(iv[5])<<8 | uint32(iv[6])<<16 | uint32(iv[7])<<24
	d1 = d0>>16 | (d2&0xFFFF0000)
	d3 = d2<<16 | (d0&0x0000FFFF)

	c.c[0] = c.cc[0] ^ d0
	c.c[1] = c.cc[1] ^ d1
	c.c[2] = c.cc[2] ^ d2
	c.c[3] = c.cc[3] ^ d3
	c.c[4] = c.cc[4] ^ d0
	c.c[5] = c.cc[5] ^ d1
	c.c[6] = c.cc[6] ^ d2
	c.c[7] = c.cc[7] ^ d3

	for i := range c.x {
		c.x[i] = c.cx[i]
	}
	c.carry = c.ccarry
	c.r = nil

	for i := 0; i < 4; i++ {
		c.rabbitNext()
	}

	return nil
}

// ProcessStream will encrypt or decrypt given buffer.
func (c *Cipher) ProcessStream(buf []byte) {
	l := len(buf)
	i := 0
	if m := len(c.r); m > 0 {
		for ; i < m && i < l; i++ {
			buf[i] ^= c.r[i]
		}
		c.r = nil
	}
	for i < l {
		c.rabbitNext()

		if n := l - i; n >= 16 {
			o0 := c.x[0] ^ (c.x[5]>>16 ^ c.x[3]<<16)
			o1 := c.x[2] ^ (c.x[7]>>16 ^ c.x[5]<<16)
			o2 := c.x[4] ^ (c.x[1]>>16 ^ c.x[7]<<16)
			o3 := c.x[6] ^ (c.x[3]>>16 ^ c.x[1]<<16)
			buf[i + 0] ^= byte(o0     )
			buf[i + 1] ^= byte(o0 >> 8)
			buf[i + 2] ^= byte(o0 >>16)
			buf[i + 3] ^= byte(o0 >>24)
			buf[i + 4] ^= byte(o1     )
			buf[i + 5] ^= byte(o1 >> 8)
			buf[i + 6] ^= byte(o1 >>16)
			buf[i + 7] ^= byte(o1 >>24)
			buf[i + 8] ^= byte(o2     )
			buf[i + 9] ^= byte(o2 >> 8)
			buf[i +10] ^= byte(o2 >>16)
			buf[i +11] ^= byte(o2 >>24)
			buf[i +12] ^= byte(o3     )
			buf[i +13] ^= byte(o3 >> 8)
			buf[i +14] ^= byte(o3 >>16)
			buf[i +15] ^= byte(o3 >>24)
			i += 16
		} else {
			for b, j, z, f := buf, 0, c.x[0] ^ (c.x[5]>>16 ^ c.x[3]<<16), false; j < 4; j++ {
				for k := uint32(0); k < 4; k++ {
					b[i] ^= byte(z>>(k*8))
					if i++; f == false && i >= l {
						l = (3 - j)*4 + (3 - int(k))
						if l == 0 { return }
						c.r = make([]byte, l)
						b, i, f = c.r, 0, true
					}
				}
				switch(j) {
				case  0: z = c.x[2] ^ (c.x[7]>>16 ^ c.x[5]<<16)
				case  1: z = c.x[4] ^ (c.x[1]>>16 ^ c.x[7]<<16)
				case  2: z = c.x[6] ^ (c.x[3]>>16 ^ c.x[1]<<16)
				}
			}
		}
	}
}

// ResetCipher reset cipher round to original state. Initialization vector will be erased.
func (c *Cipher) ResetCipher() {
	for i := range c.c {
		c.c[i] = c.cc[i]
	}
	for i := range c.x {
		c.x[i] = c.cx[i]
	}
	c.carry = c.ccarry
	c.r = nil
}

// Reset zeros the key data so that it will no longer appear in the
// process's memory.
func (c *Cipher) Reset() {
	for i := range c.x {
		c.x[i], c.c[i], c.cx[i], c.cc[i] = 0, 0, 0, 0
	}
	c.carry, c.carry = false, false
}

<|endoftext|>


<commit_before>package page_runner

import (
	"fmt"
	"github.com/jesusslim/slimmysql"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
)

type PageRunner struct {
	id        string
	waitGroup sync.WaitGroup
	sql       *slimmysql.Sql
	path      string    //扫描路径
	suffix    string    //文件后缀
	baseUrl   string    //http基础url
	ignore    []string  //忽略的目录/文件名
	rep       []string  //需替换为空的字符串
	extError  []string  //其他错误标示
	c         chan bool //channel 建议根据服务器情况调整合适大小
	result    map[string]map[string]interface{}
	cookie    string //cookie for sessionid
	maxTimes  int    //最大尝试访问次数
}

func NewPageRunner(id string, path string, suffix string, baseUrl string, ignore []string, rep []string, extErr []string, sql *slimmysql.Sql, channelLength int, cookie string, maxTimes int) *PageRunner {
	if !strings.HasSuffix(baseUrl, "/") {
		baseUrl = baseUrl + "/"
	}
	if suffix == "" {
		suffix = "index.html"
	}
	if channelLength == 0 {
		channelLength = 1
	}
	if maxTimes == 0 {
		maxTimes = 20
	}
	return &PageRunner{
		sql:      sql,
		path:     path,
		suffix:   suffix,
		baseUrl:  baseUrl,
		ignore:   ignore,
		extError: extErr,
		rep:      rep,
		c:        make(chan bool, channelLength),
		result:   make(map[string]map[string]interface{}),
		cookie:   cookie,
		maxTimes: maxTimes,
	}
}

//for thinkphp
func NewPageRunnerTP(id string, path, baseUrl string, rp []string, extErr []string, sql *slimmysql.Sql, channelLength int, cookie string, maxTimes int) *PageRunner {
	return NewPageRunner(
		id,
		path,
		".html",
		baseUrl,
		[]string{
			"Widget",
		},
		append(rp, "View/index.html", "View/"),
		extErr,
		sql,
		channelLength,
		cookie,
		maxTimes)
}

type UrlModel struct {
	module     string
	controller string
	action     string
	url        string
	times      int
	lastErr    string
}

func NewUrlModel(baseUrl, subUrl string) *UrlModel {
	mdl := &UrlModel{
		url:   baseUrl + subUrl,
		times: 0,
	}
	subs := strings.Split(subUrl, "/")
	l := len(subs)
	if l > 0 {
		mdl.module = subs[0]
	}
	if l > 1 {
		mdl.controller = subs[1]
	}
	if l > 2 {
		mdl.action = subs[2]
	}
	return mdl
}

func (this *PageRunner) fetchUrl(id string, url *UrlModel) {
	this.c <- true
	start := time.Now()
	var err error
	var resp *http.Response
	if this.cookie == "" {
		//1.easy
		resp, err = http.Get(url.url)
	} else {
		//2.with cookie
		client := &http.Client{}
		req, _ := http.NewRequest("GET", url.url, nil)
		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
		req.Header.Set("Cookie", this.cookie)
		resp, err = client.Do(req)
	}
	if err != nil {
		<-this.c
		fmt.Println("ERROR:" + err.Error())
		if url.times == this.maxTimes-1 {
			//结束 give up
			this.result[id] = map[string]interface{}{
				"duration":    0,
				"url":         url.url,
				"status":      -1,
				"create_time": time.Now().Unix(),
				"same":        0,
				"task_id":     this.id,
				"is_err":      1,
				"err":         url.lastErr,
				"module":      url.module,
				"controller":  url.controller,
				"action":      url.action,
				"times":       url.times + 1,
			}
			this.waitGroup.Done()
		} else {
			url.times++
			url.lastErr = err.Error()
			this.fetchUrl(id, url)
		}
	} else {
		duration := time.Since(start).Seconds() * 1000
		defer resp.Body.Close()

		_, ok := this.result[id]

		if !ok {
			status := resp.StatusCode
			create_time := time.Now().Unix()
			header := resp.Header
			last_url := resp.Request.URL.String()
			same := 0
			is_err := 0
			if status != 200 {
				is_err = 1
			}
			if strings.EqualFold(url.url, last_url) {
				same = 1
			}
			// fmt.Println(header.Get("Date"))
			server := header.Get("Server")
			//xby := header.Get("X-Powered-By")
			if same == 0 {
				for _, e := range this.extError {
					if strings.Contains(last_url, e) {
						is_err = 1
						break
					}
				}
			}

			this.result[id] = map[string]interface{}{
				"duration":    int(duration),
				"url":         url.url,
				"status":      status,
				"create_time": create_time,
				"server":      server,
				"last_url":    last_url,
				"same":        same,
				"task_id":     this.id,
				"is_err":      is_err,
				"module":      url.module,
				"controller":  url.controller,
				"action":      url.action,
				"err":         url.lastErr,
				"times":       url.times + 1,
			}
		}

		<-this.c
		this.waitGroup.Done()
	}
}

func (this *PageRunner) walkDir() (map[string]*UrlModel, error) {
	files := map[string]*UrlModel{}
	suffix := strings.ToUpper(this.suffix)
	id := 0
	err := filepath.Walk(this.path, func(filename string, info os.FileInfo, err_inside error) error {
		if info.IsDir() {
			return nil
		}
		if strings.HasSuffix(strings.ToUpper(info.Name()), suffix) {
			for _, ig := range this.ignore {
				if strings.Contains(filename, ig) {
					return nil
				}
			}
			for _, rp := range this.rep {
				filename = strings.Replace(filename, rp, "", -1)
			}
			if strings.HasPrefix(filename, "/") {
				filename = filename[1:]
			}
			files[strconv.Itoa(id)] = NewUrlModel(this.baseUrl, filename)
			id++
		}
		return nil
	})
	return files, err
}

func (this *PageRunner) Run() {
	urls, err := this.walkDir()
	if err != nil {
		fmt.Println("ERROR:", err.Error())
	} else {
		num := len(urls)
		id, _ := this.sql.Table("task").Add(map[string]interface{}{
			"num":         num,
			"title":       "Test",
			"create_time": time.Now().Unix(),
		})
		this.id = string(strconv.FormatInt(id, 10))
		for k, v := range urls {
			this.waitGroup.Add(1)
			go this.fetchUrl(k, v)
		}
		this.waitGroup.Wait()
		fmt.Println("DataOk")
		success := len(this.result)
		this.sql.Clear().Table("task").Where("id = "+this.id).SetInc("success", success)
		this.sql = this.sql.Clear()
		for _, v := range this.result {
			this.sql.Table("url").Add(v)
		}
		fmt.Println("Finished")

	}
}

//example
// func main() {
// 	slimmysql.RegisterConnectionDefault(false, "127.0.0.1", "3307", "test", "root", "root", "", false)
// 	sql, _ := slimmysql.NewSqlInstanceDefault()
// 	baseUrl := "http://localhost:8888/teenager/Student"
// 	path := "/Applications/MAMP/htdocs/teenager/Application/Student"
// 	extErr := []string{
// 		"404",
// 		"error",
// 	}
// 	runner := NewPageRunnerTP("0", path, baseUrl,[]string{"/Applications/MAMP/htdocs/teenager/Application/"}, extErr, sql, 30, "PHPSESSID=c3146dcc95ba4e5992441718296aef1d", 50)
// 	runner.Run()
// }
<commit_msg>响应时间debug<commit_after>package page_runner

import (
	"fmt"
	"github.com/jesusslim/slimmysql"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
)

type PageRunner struct {
	id        string
	waitGroup sync.WaitGroup
	sql       *slimmysql.Sql
	path      string    //扫描路径
	suffix    string    //文件后缀
	baseUrl   string    //http基础url
	ignore    []string  //忽略的目录/文件名
	rep       []string  //需替换为空的字符串
	extError  []string  //其他错误标示
	c         chan bool //channel 建议根据服务器情况调整合适大小
	result    map[string]map[string]interface{}
	cookie    string //cookie for sessionid
	maxTimes  int    //最大尝试访问次数
}

func NewPageRunner(id string, path string, suffix string, baseUrl string, ignore []string, rep []string, extErr []string, sql *slimmysql.Sql, channelLength int, cookie string, maxTimes int) *PageRunner {
	if !strings.HasSuffix(baseUrl, "/") {
		baseUrl = baseUrl + "/"
	}
	if suffix == "" {
		suffix = "index.html"
	}
	if channelLength == 0 {
		channelLength = 1
	}
	if maxTimes == 0 {
		maxTimes = 20
	}
	return &PageRunner{
		sql:      sql,
		path:     path,
		suffix:   suffix,
		baseUrl:  baseUrl,
		ignore:   ignore,
		extError: extErr,
		rep:      rep,
		c:        make(chan bool, channelLength),
		result:   make(map[string]map[string]interface{}),
		cookie:   cookie,
		maxTimes: maxTimes,
	}
}

//for thinkphp
func NewPageRunnerTP(id string, path, baseUrl string, rp []string, extErr []string, sql *slimmysql.Sql, channelLength int, cookie string, maxTimes int) *PageRunner {
	return NewPageRunner(
		id,
		path,
		".html",
		baseUrl,
		[]string{
			"Widget",
		},
		append(rp, "View/index.html", "View/"),
		extErr,
		sql,
		channelLength,
		cookie,
		maxTimes)
}

type UrlModel struct {
	module     string
	controller string
	action     string
	url        string
	times      int
	lastErr    string
}

func NewUrlModel(baseUrl, subUrl string) *UrlModel {
	mdl := &UrlModel{
		url:   baseUrl + subUrl,
		times: 0,
	}
	subs := strings.Split(subUrl, "/")
	l := len(subs)
	if l > 0 {
		mdl.module = subs[0]
	}
	if l > 1 {
		mdl.controller = subs[1]
	}
	if l > 2 {
		mdl.action = subs[2]
	}
	return mdl
}

func (this *PageRunner) fetchUrl(id string, url *UrlModel) {
	this.c <- true
	var start time.Time
	var err error
	var resp *http.Response
	if this.cookie == "" {
		//1.easy
		start = time.Now()
		resp, err = http.Get(url.url)
	} else {
		//2.with cookie
		client := &http.Client{}
		req, _ := http.NewRequest("GET", url.url, nil)
		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
		req.Header.Set("Cookie", this.cookie)
		start := time.Now()
		resp, err = client.Do(req)
	}
	duration := time.Since(start).Seconds() * 1000
	if err != nil {
		<-this.c
		fmt.Println("ERROR:" + err.Error())
		if url.times == this.maxTimes-1 {
			//结束 give up
			this.result[id] = map[string]interface{}{
				"duration":    0,
				"url":         url.url,
				"status":      -1,
				"create_time": time.Now().Unix(),
				"same":        0,
				"task_id":     this.id,
				"is_err":      1,
				"err":         url.lastErr,
				"module":      url.module,
				"controller":  url.controller,
				"action":      url.action,
				"times":       url.times + 1,
			}
			this.waitGroup.Done()
		} else {
			url.times++
			url.lastErr = err.Error()
			this.fetchUrl(id, url)
		}
	} else {
		defer resp.Body.Close()

		_, ok := this.result[id]

		if !ok {
			status := resp.StatusCode
			create_time := time.Now().Unix()
			header := resp.Header
			last_url := resp.Request.URL.String()
			same := 0
			is_err := 0
			if status != 200 {
				is_err = 1
			}
			if strings.EqualFold(url.url, last_url) {
				same = 1
			}
			// fmt.Println(header.Get("Date"))
			server := header.Get("Server")
			//xby := header.Get("X-Powered-By")
			if same == 0 {
				for _, e := range this.extError {
					if strings.Contains(last_url, e) {
						is_err = 1
						break
					}
				}
			}

			this.result[id] = map[string]interface{}{
				"duration":    int(duration),
				"url":         url.url,
				"status":      status,
				"create_time": create_time,
				"server":      server,
				"last_url":    last_url,
				"same":        same,
				"task_id":     this.id,
				"is_err":      is_err,
				"module":      url.module,
				"controller":  url.controller,
				"action":      url.action,
				"err":         url.lastErr,
				"times":       url.times + 1,
			}
		}

		<-this.c
		this.waitGroup.Done()
	}
}

func (this *PageRunner) walkDir() (map[string]*UrlModel, error) {
	files := map[string]*UrlModel{}
	suffix := strings.ToUpper(this.suffix)
	id := 0
	err := filepath.Walk(this.path, func(filename string, info os.FileInfo, err_inside error) error {
		if info.IsDir() {
			return nil
		}
		if strings.HasSuffix(strings.ToUpper(info.Name()), suffix) {
			for _, ig := range this.ignore {
				if strings.Contains(filename, ig) {
					return nil
				}
			}
			for _, rp := range this.rep {
				filename = strings.Replace(filename, rp, "", -1)
			}
			if strings.HasPrefix(filename, "/") {
				filename = filename[1:]
			}
			files[strconv.Itoa(id)] = NewUrlModel(this.baseUrl, filename)
			id++
		}
		return nil
	})
	return files, err
}

func (this *PageRunner) Run() {
	urls, err := this.walkDir()
	if err != nil {
		fmt.Println("ERROR:", err.Error())
	} else {
		num := len(urls)
		id, _ := this.sql.Table("task").Add(map[string]interface{}{
			"num":         num,
			"title":       "Test",
			"create_time": time.Now().Unix(),
		})
		this.id = string(strconv.FormatInt(id, 10))
		for k, v := range urls {
			this.waitGroup.Add(1)
			go this.fetchUrl(k, v)
		}
		this.waitGroup.Wait()
		fmt.Println("DataOk")
		success := len(this.result)
		this.sql.Clear().Table("task").Where("id = "+this.id).SetInc("success", success)
		this.sql = this.sql.Clear()
		for _, v := range this.result {
			this.sql.Table("url").Add(v)
		}
		fmt.Println("Finished")

	}
}

//example
// func main() {
// 	slimmysql.RegisterConnectionDefault(false, "127.0.0.1", "3307", "test", "root", "root", "", false)
// 	sql, _ := slimmysql.NewSqlInstanceDefault()
// 	baseUrl := "http://localhost:8888/teenager/Student"
// 	path := "/Applications/MAMP/htdocs/teenager/Application/Student"
// 	extErr := []string{
// 		"404",
// 		"error",
// 	}
// 	runner := NewPageRunnerTP("0", path, baseUrl,[]string{"/Applications/MAMP/htdocs/teenager/Application/"}, extErr, sql, 30, "PHPSESSID=c3146dcc95ba4e5992441718296aef1d", 50)
// 	runner.Run()
// }
<|endoftext|>


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

package main

import (
	"strings"

	"github.com/golang/protobuf/protoc-gen-go/descriptor"
)

// packageDescriptor describes a package, which is a composition of proto files.
type packageDescriptor struct {
	baseDesc
	files       []*fileDescriptor
	name        string
	title       string
	overview    string
	location    string
	frontMatter []string
}

const (
	title       = "$title: "
	overview    = "$overview: "
	location    = "$location: "
	frontMatter = "$front_matter: "
)

func newPackageDescriptor(name string, desc []*descriptor.FileDescriptorProto) *packageDescriptor {
	p := &packageDescriptor{
		name: name,
	}

	for _, fd := range desc {
		f := newFileDescriptor(fd, p)
		p.files = append(p.files, f)

		loc := f.find(newPathVector(packagePath))

		if loc != nil {
			if p.loc == nil {
				p.loc = loc
			}

			if loc.LeadingDetachedComments != nil {
				for _, para := range loc.LeadingDetachedComments {
					lines := strings.Split(para, "\n")
					for _, l := range lines {
						l = strings.Trim(l, " ")

						if strings.HasPrefix(l, title) {
							p.title = l[len(title):]
						} else if strings.HasPrefix(l, overview) {
							p.overview = l[len(overview):]
						} else if strings.HasPrefix(l, location) {
							p.location = l[len(location):]
						} else if strings.HasPrefix(l, frontMatter) {
							p.frontMatter = append(p.frontMatter, l[len(frontMatter):])
						}
					}
				}
			}
		}
	}

	return p
}
<commit_msg>Fix handling of package-level comments in a multi-file package.<commit_after>// Copyright 2018 Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this currentFile except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"strings"

	"github.com/golang/protobuf/protoc-gen-go/descriptor"
)

// packageDescriptor describes a package, which is a composition of proto files.
type packageDescriptor struct {
	baseDesc
	files       []*fileDescriptor
	name        string
	title       string
	overview    string
	location    string
	frontMatter []string
}

const (
	title       = "$title: "
	overview    = "$overview: "
	location    = "$location: "
	frontMatter = "$front_matter: "
)

func newPackageDescriptor(name string, desc []*descriptor.FileDescriptorProto) *packageDescriptor {
	p := &packageDescriptor{
		name: name,
	}

	for _, fd := range desc {
		f := newFileDescriptor(fd, p)
		p.files = append(p.files, f)

		loc := f.find(newPathVector(packagePath))

		if loc != nil {
			if p.loc == nil {
				if loc.GetLeadingComments() != "" || loc.GetTrailingComments() != "" {
					p.loc = loc
				}
			}

			if loc.LeadingDetachedComments != nil {
				for _, para := range loc.LeadingDetachedComments {
					lines := strings.Split(para, "\n")
					for _, l := range lines {
						l = strings.Trim(l, " ")

						if strings.HasPrefix(l, title) {
							p.title = l[len(title):]
						} else if strings.HasPrefix(l, overview) {
							p.overview = l[len(overview):]
						} else if strings.HasPrefix(l, location) {
							p.location = l[len(location):]
						} else if strings.HasPrefix(l, frontMatter) {
							p.frontMatter = append(p.frontMatter, l[len(frontMatter):])
						}
					}
				}
			}
		}
	}

	return p
}
<|endoftext|>


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

package virtualips

import (
	"errors"
	"fmt"
	"path"
	"time"

	"github.com/control-center/serviced/coordinator/client"
	"github.com/control-center/serviced/domain/pool"
	"github.com/control-center/serviced/zzk"
	"github.com/zenoss/glog"
)

const (
	zkVirtualIP            = "/virtualIPs"
	virtualInterfacePrefix = ":zvip"
	maxRetries             = 2
	waitTimeout            = 30 * time.Second
)

var (
	ErrInvalidVirtualIP = errors.New("invalid virtual ip")
)

func vippath(nodes ...string) string {
	p := append([]string{zkVirtualIP}, nodes...)
	return path.Join(p...)
}

type VirtualIPNode struct {
	*pool.VirtualIP
	version interface{}
}

// ID implements zzk.Node
func (node *VirtualIPNode) GetID() string {
	return node.IP
}

// Create implements zzk.Node
func (node *VirtualIPNode) Create(conn client.Connection) error {
	return AddVirtualIP(conn, node.VirtualIP)
}

// Update implements zzk.Node
func (node *VirtualIPNode) Update(conn client.Connection) error {
	return nil
}

func (node *VirtualIPNode) Version() interface{}           { return node.version }
func (node *VirtualIPNode) SetVersion(version interface{}) { node.version = version }

// VirtualIPHandler is the handler interface for virtual ip bindings on the host
type VirtualIPHandler interface {
	BindVirtualIP(*pool.VirtualIP, string) error
	UnbindVirtualIP(*pool.VirtualIP) error
	VirtualInterfaceMap(string) (map[string]*pool.VirtualIP, error)
}

// VirtualIPListener is the listener object for watching the zk object for
// virtual IP nodes
type VirtualIPListener struct {
	conn    client.Connection
	handler VirtualIPHandler
	hostID  string

	index chan int
	ips   map[string]chan bool
	retry map[string]int
}

// NewVirtualIPListener instantiates a new VirtualIPListener object
func NewVirtualIPListener(handler VirtualIPHandler, hostID string) *VirtualIPListener {
	l := &VirtualIPListener{
		handler: handler,
		hostID:  hostID,
		index:   make(chan int),
		ips:     make(map[string]chan bool),
	}

	// Index generator for bind interface
	go func(start int) {
		for {
			l.index <- start
			start++
		}
	}(0)

	return l
}

// GetConnection implements zzk.Listener
func (l *VirtualIPListener) SetConnection(conn client.Connection) { l.conn = conn }

// GetPath implements zzk.Listener
func (l *VirtualIPListener) GetPath(nodes ...string) string {
	return vippath(nodes...)
}

// Ready removes all virtual IPs that may be present
func (l *VirtualIPListener) Ready() error {
	vmap, err := l.handler.VirtualInterfaceMap(virtualInterfacePrefix)
	if err != nil {
		return err
	}

	for _, vip := range vmap {
		if err := l.handler.UnbindVirtualIP(vip); err != nil {
			return err
		}
	}
	return nil
}

// Done implements zzk.Listener
func (l *VirtualIPListener) Done() {}

// PostProcess implements zzk.Listener
func (l *VirtualIPListener) PostProcess(p map[string]struct{}) {}

// Spawn implements zzk.Listener
func (l *VirtualIPListener) Spawn(shutdown <-chan interface{}, ip string) {
	// Check if this ip has exceeded the number of retries for this host
	if l.retry[ip] > maxRetries {
		glog.Warningf("Throttling acquisition of %s for %s", ip, l.hostID)
		select {
		case <-time.After(waitTimeout):
		case <-shutdown:
			return
		}
	}

	glog.V(2).Infof("Host %s waiting to acquire virtual ip %s", l.hostID, ip)
	// Try to take lead on the path
	leader := zzk.NewHostLeader(l.conn, l.hostID, "", l.GetPath(ip))
	_, err := leader.TakeLead()
	if err != nil {
		glog.Errorf("Error while trying to acquire a lock for %s: %s", ip, err)
		return
	}
	defer leader.ReleaseLead()

	select {
	case <-shutdown:
		return
	default:
	}

	// Check if the path still exists
	if exists, err := zzk.PathExists(l.conn, l.GetPath(ip)); err != nil {
		glog.Errorf("Error while checking ip %s: %s", ip, err)
		return
	} else if !exists {
		return
	}

	index := l.getIndex()
	for {
		var vip pool.VirtualIP
		event, err := l.conn.GetW(l.GetPath(ip), &VirtualIPNode{VirtualIP: &vip})
		if err != nil {
			glog.Errorf("Could not load virtual ip %s: %s", ip, err)
			return
		}

		glog.V(2).Infof("Host %s binding to %s", l.hostID, ip)
		rebind, err := l.bind(&vip, index)
		if err != nil {
			glog.Errorf("Could not bind to virtual ip %s: %s", ip, err)
			l.retry[ip]++
			return
		}

		if l.retry[ip] > 0 {
			l.retry[ip]--
		}

		select {
		case e := <-event:
			// If the virtual ip is changed, you need to update the bindings
			if err := l.unbind(ip); err != nil {
				glog.Errorf("Could not unbind to virtual ip %s: %s", ip, err)
				return
			}
			if e.Type == client.EventNodeDeleted {
				return
			}
			glog.V(4).Infof("virtual ip listener for %s receieved event: %v", ip, e)
		case <-rebind:
			// If the primary virtual IP is removed, all other virtual IPs on
			// that subnet are removed.  This is in place to restore the
			// virtual IPs that were removed soley by the removal of the
			// primary virtual IP.
			glog.V(2).Infof("Host %s rebinding to %s", l.hostID, ip)
		case <-shutdown:
			if err := l.unbind(ip); err != nil {
				glog.Errorf("Could not unbind to virtual ip %s: %s", ip, err)
				return
			}
		}
	}
}

func (l *VirtualIPListener) getIndex() int {
	return <-l.index
}

func (l *VirtualIPListener) reset() {
	for _, ipChan := range l.ips {
		ipChan <- true
	}
}

func (l *VirtualIPListener) get(ip string) <-chan bool {
	l.ips[ip] = make(chan bool, 1)
	return l.ips[ip]
}

func (l *VirtualIPListener) bind(vip *pool.VirtualIP, index int) (<-chan bool, error) {
	vmap, err := l.handler.VirtualInterfaceMap(virtualInterfacePrefix)
	if err != nil {
		return nil, err
	}

	if _, ok := vmap[vip.IP]; !ok {
		if vip.BindInterface == "" {
			return nil, ErrInvalidVirtualIP
		}
		vname := fmt.Sprintf("%s%s%d", vip.BindInterface, virtualInterfacePrefix, index)
		if err := l.handler.BindVirtualIP(vip, vname); err != nil {
			return nil, err
		}
	}

	return l.get(vip.IP), nil
}

func (l *VirtualIPListener) unbind(ip string) error {
	defer l.reset()
	vmap, err := l.handler.VirtualInterfaceMap(virtualInterfacePrefix)
	if err != nil {
		return err
	}

	if vip, ok := vmap[ip]; ok {
		return l.handler.UnbindVirtualIP(vip)
	}

	return nil
}

func SyncVirtualIPs(conn client.Connection, virtualIPs []pool.VirtualIP) error {
	nodes := make([]zzk.Node, len(virtualIPs))
	for i := range virtualIPs {
		nodes[i] = &VirtualIPNode{VirtualIP: &virtualIPs[i]}
	}
	return zzk.Sync(conn, nodes, vippath())
}

func AddVirtualIP(conn client.Connection, virtualIP *pool.VirtualIP) error {
	var node VirtualIPNode
	path := vippath(virtualIP.IP)

	glog.V(1).Infof("Adding virtual ip to zookeeper: %s", path)
	if err := conn.Create(path, &node); err != nil {
		return err
	}
	node.VirtualIP = virtualIP
	return conn.Set(path, &node)
}

func RemoveVirtualIP(conn client.Connection, ip string) error {
	glog.V(1).Infof("Removing virtual ip from zookeeper: %s", vippath(ip))
	err := conn.Delete(vippath(ip))
	if err == nil || err == client.ErrNoNode {
		return nil
	}
	return err
}

func GetHostID(conn client.Connection, ip string) (string, error) {
	leader := zzk.NewHostLeader(conn, "", "", vippath(ip))
	return zzk.GetHostID(leader)
}
<commit_msg>fixes CC-284: ensure that retry sentinel has good initial state<commit_after>// Copyright 2014 The Serviced Authors.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package virtualips

import (
	"errors"
	"fmt"
	"path"
	"time"

	"github.com/control-center/serviced/coordinator/client"
	"github.com/control-center/serviced/domain/pool"
	"github.com/control-center/serviced/zzk"
	"github.com/zenoss/glog"
)

const (
	zkVirtualIP            = "/virtualIPs"
	virtualInterfacePrefix = ":zvip"
	maxRetries             = 2
	waitTimeout            = 30 * time.Second
)

var (
	ErrInvalidVirtualIP = errors.New("invalid virtual ip")
)

func vippath(nodes ...string) string {
	p := append([]string{zkVirtualIP}, nodes...)
	return path.Join(p...)
}

type VirtualIPNode struct {
	*pool.VirtualIP
	version interface{}
}

// ID implements zzk.Node
func (node *VirtualIPNode) GetID() string {
	return node.IP
}

// Create implements zzk.Node
func (node *VirtualIPNode) Create(conn client.Connection) error {
	return AddVirtualIP(conn, node.VirtualIP)
}

// Update implements zzk.Node
func (node *VirtualIPNode) Update(conn client.Connection) error {
	return nil
}

func (node *VirtualIPNode) Version() interface{}           { return node.version }
func (node *VirtualIPNode) SetVersion(version interface{}) { node.version = version }

// VirtualIPHandler is the handler interface for virtual ip bindings on the host
type VirtualIPHandler interface {
	BindVirtualIP(*pool.VirtualIP, string) error
	UnbindVirtualIP(*pool.VirtualIP) error
	VirtualInterfaceMap(string) (map[string]*pool.VirtualIP, error)
}

// VirtualIPListener is the listener object for watching the zk object for
// virtual IP nodes
type VirtualIPListener struct {
	conn    client.Connection
	handler VirtualIPHandler
	hostID  string

	index chan int
	ips   map[string]chan bool
	retry map[string]int
}

// NewVirtualIPListener instantiates a new VirtualIPListener object
func NewVirtualIPListener(handler VirtualIPHandler, hostID string) *VirtualIPListener {
	l := &VirtualIPListener{
		handler: handler,
		hostID:  hostID,
		index:   make(chan int),
		ips:     make(map[string]chan bool),
	}

	// Index generator for bind interface
	go func(start int) {
		for {
			l.index <- start
			start++
		}
	}(0)

	return l
}

// GetConnection implements zzk.Listener
func (l *VirtualIPListener) SetConnection(conn client.Connection) { l.conn = conn }

// GetPath implements zzk.Listener
func (l *VirtualIPListener) GetPath(nodes ...string) string {
	return vippath(nodes...)
}

// Ready removes all virtual IPs that may be present
func (l *VirtualIPListener) Ready() error {
	vmap, err := l.handler.VirtualInterfaceMap(virtualInterfacePrefix)
	if err != nil {
		return err
	}

	for _, vip := range vmap {
		if err := l.handler.UnbindVirtualIP(vip); err != nil {
			return err
		}
	}
	return nil
}

// Done implements zzk.Listener
func (l *VirtualIPListener) Done() {}

// PostProcess implements zzk.Listener
func (l *VirtualIPListener) PostProcess(p map[string]struct{}) {}

// Spawn implements zzk.Listener
func (l *VirtualIPListener) Spawn(shutdown <-chan interface{}, ip string) {
	// ensure that the retry sentinel has good initial state
	if l.retry == nil {
		l.retry = make(map[string]int)
	}
	if _, ok := l.retry[ip]; !ok {
		l.retry[ip] = maxRetries
	}

	// Check if this ip has exceeded the number of retries for this host
	if l.retry[ip] > maxRetries {
		glog.Warningf("Throttling acquisition of %s for %s", ip, l.hostID)
		select {
		case <-time.After(waitTimeout):
		case <-shutdown:
			return
		}
	}

	glog.V(2).Infof("Host %s waiting to acquire virtual ip %s", l.hostID, ip)
	// Try to take lead on the path
	leader := zzk.NewHostLeader(l.conn, l.hostID, "", l.GetPath(ip))
	_, err := leader.TakeLead()
	if err != nil {
		glog.Errorf("Error while trying to acquire a lock for %s: %s", ip, err)
		return
	}
	defer leader.ReleaseLead()

	select {
	case <-shutdown:
		return
	default:
	}

	// Check if the path still exists
	if exists, err := zzk.PathExists(l.conn, l.GetPath(ip)); err != nil {
		glog.Errorf("Error while checking ip %s: %s", ip, err)
		return
	} else if !exists {
		return
	}

	index := l.getIndex()
	for {
		var vip pool.VirtualIP
		event, err := l.conn.GetW(l.GetPath(ip), &VirtualIPNode{VirtualIP: &vip})
		if err != nil {
			glog.Errorf("Could not load virtual ip %s: %s", ip, err)
			return
		}

		glog.V(2).Infof("Host %s binding to %s", l.hostID, ip)
		rebind, err := l.bind(&vip, index)
		if err != nil {
			glog.Errorf("Could not bind to virtual ip %s: %s", ip, err)
			l.retry[ip]++
			return
		}

		if l.retry[ip] > 0 {
			l.retry[ip]--
		}

		select {
		case e := <-event:
			// If the virtual ip is changed, you need to update the bindings
			if err := l.unbind(ip); err != nil {
				glog.Errorf("Could not unbind to virtual ip %s: %s", ip, err)
				return
			}
			if e.Type == client.EventNodeDeleted {
				return
			}
			glog.V(4).Infof("virtual ip listener for %s receieved event: %v", ip, e)
		case <-rebind:
			// If the primary virtual IP is removed, all other virtual IPs on
			// that subnet are removed.  This is in place to restore the
			// virtual IPs that were removed soley by the removal of the
			// primary virtual IP.
			glog.V(2).Infof("Host %s rebinding to %s", l.hostID, ip)
		case <-shutdown:
			if err := l.unbind(ip); err != nil {
				glog.Errorf("Could not unbind to virtual ip %s: %s", ip, err)
				return
			}
		}
	}
}

func (l *VirtualIPListener) getIndex() int {
	return <-l.index
}

func (l *VirtualIPListener) reset() {
	for _, ipChan := range l.ips {
		ipChan <- true
	}
}

func (l *VirtualIPListener) get(ip string) <-chan bool {
	l.ips[ip] = make(chan bool, 1)
	return l.ips[ip]
}

func (l *VirtualIPListener) bind(vip *pool.VirtualIP, index int) (<-chan bool, error) {
	vmap, err := l.handler.VirtualInterfaceMap(virtualInterfacePrefix)
	if err != nil {
		return nil, err
	}

	if _, ok := vmap[vip.IP]; !ok {
		if vip.BindInterface == "" {
			return nil, ErrInvalidVirtualIP
		}
		vname := fmt.Sprintf("%s%s%d", vip.BindInterface, virtualInterfacePrefix, index)
		if err := l.handler.BindVirtualIP(vip, vname); err != nil {
			return nil, err
		}
	}

	return l.get(vip.IP), nil
}

func (l *VirtualIPListener) unbind(ip string) error {
	defer l.reset()
	vmap, err := l.handler.VirtualInterfaceMap(virtualInterfacePrefix)
	if err != nil {
		return err
	}

	if vip, ok := vmap[ip]; ok {
		return l.handler.UnbindVirtualIP(vip)
	}

	return nil
}

func SyncVirtualIPs(conn client.Connection, virtualIPs []pool.VirtualIP) error {
	nodes := make([]zzk.Node, len(virtualIPs))
	for i := range virtualIPs {
		nodes[i] = &VirtualIPNode{VirtualIP: &virtualIPs[i]}
	}
	return zzk.Sync(conn, nodes, vippath())
}

func AddVirtualIP(conn client.Connection, virtualIP *pool.VirtualIP) error {
	var node VirtualIPNode
	path := vippath(virtualIP.IP)

	glog.V(1).Infof("Adding virtual ip to zookeeper: %s", path)
	if err := conn.Create(path, &node); err != nil {
		return err
	}
	node.VirtualIP = virtualIP
	return conn.Set(path, &node)
}

func RemoveVirtualIP(conn client.Connection, ip string) error {
	glog.V(1).Infof("Removing virtual ip from zookeeper: %s", vippath(ip))
	err := conn.Delete(vippath(ip))
	if err == nil || err == client.ErrNoNode {
		return nil
	}
	return err
}

func GetHostID(conn client.Connection, ip string) (string, error) {
	leader := zzk.NewHostLeader(conn, "", "", vippath(ip))
	return zzk.GetHostID(leader)
}
<|endoftext|>


<commit_before>package util

import (
	"errors"
	"net"
)

// NoDefaultIP means no suitable non-loopback address can be found.
var NoDefaultIP = errors.New("no suitable IP address")

// DefaultLocalIP4 returns an IPv4 address that this host can be reached
// on. Will return NoDefaultIP if no suitable address can be found.
func DefaultLocalIP4() (net.IP, error) {
	devices, err := net.Interfaces()
	if err != nil {
		return nil, err
	}
	for _, dev := range devices {
		if (dev.Flags&net.FlagUp != 0) && (dev.Flags&net.FlagLoopback == 0) {
			addrs, err := dev.Addrs()
			if err != nil {
				continue
			}
			for i := range addrs {
				if ip, ok := addrs[i].(*net.IPNet); ok {
					if ip.IP.To4() != nil {
						return ip.IP, nil
					}
				}
			}
		}
	}
	return nil, NoDefaultIP
}
<commit_msg>NoDefaultIP should be ErrNoDefaultIP<commit_after>package util

import (
	"errors"
	"net"
)

// NoDefaultIP means no suitable non-loopback address can be found.
var ErrorNoDefaultIP = errors.New("no suitable IP address")

// DefaultLocalIP4 returns an IPv4 address that this host can be reached
// on. Will return NoDefaultIP if no suitable address can be found.
func DefaultLocalIP4() (net.IP, error) {
	devices, err := net.Interfaces()
	if err != nil {
		return nil, err
	}
	for _, dev := range devices {
		if (dev.Flags&net.FlagUp != 0) && (dev.Flags&net.FlagLoopback == 0) {
			addrs, err := dev.Addrs()
			if err != nil {
				continue
			}
			for i := range addrs {
				if ip, ok := addrs[i].(*net.IPNet); ok {
					if ip.IP.To4() != nil {
						return ip.IP, nil
					}
				}
			}
		}
	}
	return nil, ErrorNoDefaultIP
}
<|endoftext|>


<commit_before>package md_test

import (
	"html"
	"strings"
	"testing"

	"github.com/google/go-cmp/cmp"
	. "src.elv.sh/pkg/md"
	"src.elv.sh/pkg/testutil"
)

var supplementalFmtCases = []testCase{
	{
		Section:  "Fenced code blocks",
		Name:     "Tilde fence with info starting with tilde",
		Markdown: "~~~ ~`\n" + "~~~",
	},
	{
		Section:  "Emphasis and strong emphasis",
		Name:     "Space at start of content",
		Markdown: "* x*",
	},
	{
		Section:  "Emphasis and strong emphasis",
		Name:     "Space at end of content",
		Markdown: "*x *",
	},
	{
		Section:  "Emphasis and strong emphasis",
		Name:     "Emphasis opener after word before punctuation",
		Markdown: "A*!*",
	},
	{
		Section:  "Emphasis and strong emphasis",
		Name:     "Emphasis closer after punctuation before word",
		Markdown: "*!*A",
	},
	{
		Section:  "Emphasis and strong emphasis",
		Name:     "Space-only content",
		Markdown: "* *",
	},
	{
		Section:  "Links",
		Name:     "Exclamation mark before link",
		Markdown: `\![a](b)`,
	},
	{
		Section:  "Links",
		Name:     "Link title with both single and double quotes",
		Markdown: `[a](b ('"))`,
	},
	{
		Section:  "Links",
		Name:     "Link title with fewer double quotes than single quotes and parens",
		Markdown: `[a](b "\"''()")`,
	},
	{
		Section:  "Links",
		Name:     "Link title with fewer single quotes than double quotes and parens",
		Markdown: `[a](b '\'""()')`,
	},
	{
		Section:  "Links",
		Name:     "Link title with fewer parens than single and double quotes",
		Markdown: `[a](b (\(''""))`,
	},
	{
		Section:  "Soft line breaks",
		Name:     "Space at start of line",
		Markdown: " foo",
	},
	{
		Section:  "Soft line breaks",
		Name:     "Space at end of line",
		Markdown: "foo ",
	},
}

var fmtTestCases = concat(htmlTestCases, supplementalFmtCases)

func TestFmtPreservesHTMLRender(t *testing.T) {
	testutil.Set(t, &UnescapeHTML, html.UnescapeString)
	for _, tc := range fmtTestCases {
		t.Run(tc.testName(), func(t *testing.T) {
			testFmtPreservesHTMLRender(t, tc.Markdown)
		})
	}
}

func FuzzFmtPreservesHTMLRender(f *testing.F) {
	for _, tc := range fmtTestCases {
		f.Add(tc.Markdown)
	}
	f.Fuzz(testFmtPreservesHTMLRender)
}

func testFmtPreservesHTMLRender(t *testing.T, original string) {
	t.Helper()
	formatted := formatAndSkipIfUnsupported(t, original)
	formattedRender := RenderString(formatted, &HTMLCodec{})
	originalRender := RenderString(original, &HTMLCodec{})
	if formattedRender != originalRender {
		t.Errorf("original:\n%s\nformatted:\n%s\n"+
			"markdown diff (-original +formatted):\n%s"+
			"HTML diff (-original +formatted):\n%s"+
			"ops diff (-original +formatted):\n%s",
			hr+"\n"+original+hr, hr+"\n"+formatted+hr,
			cmp.Diff(original, formatted),
			cmp.Diff(originalRender, formattedRender),
			cmp.Diff(RenderString(original, &TraceCodec{}), RenderString(formatted, &TraceCodec{})))
	}
}

func TestFmtIsIdempotent(t *testing.T) {
	testutil.Set(t, &UnescapeHTML, html.UnescapeString)
	for _, tc := range fmtTestCases {
		t.Run(tc.testName(), func(t *testing.T) {
			testFmtIsIdempotent(t, tc.Markdown)
		})
	}
}

func FuzzFmtIsIdempotent(f *testing.F) {
	for _, tc := range fmtTestCases {
		f.Add(tc.Markdown)
	}
	f.Fuzz(testFmtIsIdempotent)
}

func testFmtIsIdempotent(t *testing.T, original string) {
	formatted1 := formatAndSkipIfUnsupported(t, original)
	formatted2 := RenderString(formatted1, &FmtCodec{})
	if formatted1 != formatted2 {
		t.Errorf("original:\n%s\nformatted1:\n%s\nformatted2:\n%s\n"+
			"diff (-formatted1 +formatted2):\n%s",
			hr+"\n"+original+hr, hr+"\n"+formatted1+hr, hr+"\n"+formatted2+hr,
			cmp.Diff(formatted1, formatted2))
	}
}

func formatAndSkipIfUnsupported(t *testing.T, original string) string {
	t.Helper()
	if strings.Contains(original, "\t") {
		t.Skipf("input contains tab")
	}
	codec := &FmtCodec{}
	formatted := RenderString(original, codec)
	if u := codec.Unsupported(); u != nil {
		t.Skipf("input uses unsupported feature: %v", u)
	}
	return formatted
}
<commit_msg>pkg/md/fmt_test.go: Remove test for the idempotency of FmtCodec.<commit_after>package md_test

import (
	"html"
	"strings"
	"testing"

	"github.com/google/go-cmp/cmp"
	. "src.elv.sh/pkg/md"
	"src.elv.sh/pkg/testutil"
)

var supplementalFmtCases = []testCase{
	{
		Section:  "Fenced code blocks",
		Name:     "Tilde fence with info starting with tilde",
		Markdown: "~~~ ~`\n" + "~~~",
	},
	{
		Section:  "Emphasis and strong emphasis",
		Name:     "Space at start of content",
		Markdown: "* x*",
	},
	{
		Section:  "Emphasis and strong emphasis",
		Name:     "Space at end of content",
		Markdown: "*x *",
	},
	{
		Section:  "Emphasis and strong emphasis",
		Name:     "Emphasis opener after word before punctuation",
		Markdown: "A*!*",
	},
	{
		Section:  "Emphasis and strong emphasis",
		Name:     "Emphasis closer after punctuation before word",
		Markdown: "*!*A",
	},
	{
		Section:  "Emphasis and strong emphasis",
		Name:     "Space-only content",
		Markdown: "* *",
	},
	{
		Section:  "Links",
		Name:     "Exclamation mark before link",
		Markdown: `\![a](b)`,
	},
	{
		Section:  "Links",
		Name:     "Link title with both single and double quotes",
		Markdown: `[a](b ('"))`,
	},
	{
		Section:  "Links",
		Name:     "Link title with fewer double quotes than single quotes and parens",
		Markdown: `[a](b "\"''()")`,
	},
	{
		Section:  "Links",
		Name:     "Link title with fewer single quotes than double quotes and parens",
		Markdown: `[a](b '\'""()')`,
	},
	{
		Section:  "Links",
		Name:     "Link title with fewer parens than single and double quotes",
		Markdown: `[a](b (\(''""))`,
	},
	{
		Section:  "Soft line breaks",
		Name:     "Space at start of line",
		Markdown: " foo",
	},
	{
		Section:  "Soft line breaks",
		Name:     "Space at end of line",
		Markdown: "foo ",
	},
}

var fmtTestCases = concat(htmlTestCases, supplementalFmtCases)

func TestFmtPreservesHTMLRender(t *testing.T) {
	testutil.Set(t, &UnescapeHTML, html.UnescapeString)
	for _, tc := range fmtTestCases {
		t.Run(tc.testName(), func(t *testing.T) {
			testFmtPreservesHTMLRender(t, tc.Markdown)
		})
	}
}

func FuzzFmtPreservesHTMLRender(f *testing.F) {
	for _, tc := range fmtTestCases {
		f.Add(tc.Markdown)
	}
	f.Fuzz(testFmtPreservesHTMLRender)
}

func testFmtPreservesHTMLRender(t *testing.T, original string) {
	t.Helper()
	formatted := formatAndSkipIfUnsupported(t, original)
	formattedRender := RenderString(formatted, &HTMLCodec{})
	originalRender := RenderString(original, &HTMLCodec{})
	if formattedRender != originalRender {
		t.Errorf("original:\n%s\nformatted:\n%s\n"+
			"markdown diff (-original +formatted):\n%s"+
			"HTML diff (-original +formatted):\n%s"+
			"ops diff (-original +formatted):\n%s",
			hr+"\n"+original+hr, hr+"\n"+formatted+hr,
			cmp.Diff(original, formatted),
			cmp.Diff(originalRender, formattedRender),
			cmp.Diff(RenderString(original, &TraceCodec{}), RenderString(formatted, &TraceCodec{})))
	}
}

func formatAndSkipIfUnsupported(t *testing.T, original string) string {
	t.Helper()
	if strings.Contains(original, "\t") {
		t.Skipf("input contains tab")
	}
	codec := &FmtCodec{}
	formatted := RenderString(original, codec)
	if u := codec.Unsupported(); u != nil {
		t.Skipf("input uses unsupported feature: %v", u)
	}
	return formatted
}
<|endoftext|>


<commit_before>package object

import (
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"math"
	"reflect"
	"sort"
	"strings"

	"nimona.io/internal/encoding/base58"
	"nimona.io/pkg/errors"
)

// ObjectHash consistently hashes a map.
// It is based on Ben Laurie's object hash, but using the same type hints
// as TJSON instead.

const (
	EmptyHash = Hash("")
)

// SortHashes sorts a slice of Hashes in increasing order, and also returns it.
// The return part is mostly for allowing this to be used as a helper method in
// tests.
func SortHashes(a []Hash) []Hash {
	sort.Sort(HashSlice(a))
	return a
}

type HashSlice []Hash

func (p HashSlice) Len() int           { return len(p) }
func (p HashSlice) Less(i, j int) bool { return p[i] < p[j] }
func (p HashSlice) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
func (p HashSlice) Sort()              { sort.Sort(p) }

func (h Hash) IsEmpty() bool {
	return string(h) == ""
}

func NewHash(o *Object) (Hash, error) {
	m, err := objectToMap(o)
	if err != nil {
		return "", err
	}
	return hashMap(m)
}

func hintsFromKey(k string) []TypeHint {
	ps := strings.Split(k, ":")
	if len(ps) == 1 {
		return nil
	}
	hs := []TypeHint{}
	for _, sh := range ps[1] {
		hs = append(hs, GetTypeHint(string(sh)))
	}
	return hs
}

func hintFromKey(k string) string {
	ps := strings.Split(k, ":")
	if len(ps) == 1 {
		return ""
	}
	if ps[1] == HintObject.String() {
		return HintRef.String()
	}
	if ps[1] == HintArray.String()+HintObject.String() {
		return HintArray.String() + HintRef.String()
	}
	return ps[1]
}

func hashMap(m map[string]interface{}) (Hash, error) {
	if len(m) == 0 {
		return EmptyHash, nil
	}
	b := []byte{}
	ks := []string{}
	for k := range m {
		if len(k) > 0 && k[0] == '_' {
			continue
		}
		ks = append(ks, k)
	}
	sort.Strings(ks)
	for _, k := range ks {
		v := m[k]
		if v == nil {
			continue
		}
		ts := hintsFromKey(k)
		hv, err := hashValueAs(k, v, ts...)
		if err != nil {
			return EmptyHash, err
		}
		if hv.IsEmpty() {
			continue
		}
		// hash the key (including the hint)
		// but if we are dealing with an object, replace its hint with r
		ck := k
		if strings.HasSuffix(k, ":"+HintObject.String()) {
			ck = k[:len(k)-2] + ":" + HintRef.String()
		}
		// TODO consider not replacing hints for maps
		if strings.HasSuffix(k, ":"+HintMap.String()) {
			ck = k[:len(k)-2] + ":" + HintRef.String()
		}
		if strings.HasSuffix(k, ":a"+HintObject.String()) {
			ck = k[:len(k)-3] + ":a" + HintRef.String()
		}
		// TODO consider not replacing hints for maps
		if strings.HasSuffix(k, ":a"+HintMap.String()) {
			ck = k[:len(k)-3] + ":a" + HintRef.String()
		}
		hk, _ := hash("", []byte(ck)) // nolint: errcheck
		b = append(b, hk...)
		b = append(b, hv...)
	}
	if len(b) == 0 {
		return EmptyHash, nil
	}
	return hash(HintRef, b)
}

func hash(p TypeHint, v []byte) (Hash, error) {
	d := sha256.Sum256(append([]byte(p), v...))
	return Hash("oh1." + base58.Encode(d[:])), nil
}

func hashValueAs(k string, o interface{}, ts ...TypeHint) (Hash, error) {
	if o == nil {
		return EmptyHash, nil
	}

	v := reflect.ValueOf(o)
	t := reflect.TypeOf(o)

	if len(ts) == 0 {
		return EmptyHash, nil
	}

	switch ts[0] {
	case HintArray:
		if v.Len() == 0 {
			return EmptyHash, nil
		}
		vs := []interface{}{}
		for i := 0; i < v.Len(); i++ {
			vs = append(vs, v.Index(i).Interface())
		}
		b := []byte{}
		for _, iv := range vs {
			ivv, err := hashValueAs(k, iv, ts[1:]...)
			if err != nil {
				return EmptyHash, err
			}
			b = append(b, ivv...)
		}
		// TODO(geoah) hint SHOULD NOT be array, but array + inner hint
		return hash(TypeHint(hintFromKey(k)), b)
	case HintString:
		s, ok := o.(string)
		if !ok {
			if ss, ok := o.(interface{ String() string }); ok {
				s = ss.String()
			} else {
				return EmptyHash, nil
			}
		}
		if s == "" {
			return EmptyHash, nil
		}
		return hash(HintString, []byte(s))
	case HintData:
		switch t.Kind() {
		case reflect.String:
			d, err := base64.StdEncoding.DecodeString(o.(string))
			if err != nil {
				panic(err)
			}
			return hash(HintData, d)
		case reflect.Slice:
			bo := make([]byte, v.Len())
			switch t.Elem().Kind() {
			case reflect.Int,
				reflect.Int8,
				reflect.Int16,
				reflect.Int32,
				reflect.Int64:
				for i := 0; i < v.Len(); i++ {
					bo[i] = uint8(v.Index(i).Int())
				}
				return hash(HintData, bo)
			case reflect.Uint,
				reflect.Uint8,
				reflect.Uint16,
				reflect.Uint32,
				reflect.Uint64:
				for i := 0; i < v.Len(); i++ {
					bo[i] = uint8(v.Index(i).Uint())
				}
				return hash(HintData, bo)
			case reflect.Interface:
				for i := 0; i < v.Len(); i++ {
					iv := v.Index(i).Interface()
					switch ivv := iv.(type) {
					case uint8:
						bo[i] = ivv
					case uint64:
						bo[i] = uint8(ivv)
					case float64:
						bo[i] = uint8(ivv)
					default:
						panic("data should be some sort of number array, was " +
							t.Elem().Kind().String())
					}
				}
				return hash(HintData, bo)
			default:
				panic("data should be some sort of number array, was " +
					t.Elem().Kind().String())
			}
		}
		return hash(HintData, o.([]byte))
	case HintMap, HintObject:
		switch v := o.(type) {
		case map[string]interface{}:
			if len(v) == 0 {
				return EmptyHash, nil
			}
			h, err := hashMap(v)
			if err != nil {
				return EmptyHash, err
			}
			return h, nil
		case *Object:
			m := v.ToMap()
			if len(m) == 0 {
				return EmptyHash, nil
			}
			h, err := hashMap(m)
			if err != nil {
				return EmptyHash, err
			}
			return h, nil
		default:
			panic("hashing only supports map[string]interface{}")
		}
	case HintFloat:
		switch t.Kind() {
		case reflect.Float32,
			reflect.Float64:
			nf, err := hashFloat(v.Float())
			if err != nil {
				return EmptyHash, err
			}
			return nf, nil
		case reflect.Int,
			reflect.Int8,
			reflect.Int16,
			reflect.Int32,
			reflect.Int64:
			nf, err := hashFloat(float64(v.Int()))
			if err != nil {
				return EmptyHash, err
			}
			return nf, nil
		case reflect.Uint,
			reflect.Uint8,
			reflect.Uint16,
			reflect.Uint32,
			reflect.Uint64:
			nf, err := hashFloat(float64(v.Uint()))
			if err != nil {
				return EmptyHash, err
			}
			return nf, nil
		}
	case HintInt:
		switch t.Kind() {
		case reflect.Float32,
			reflect.Float64:
			return hash(HintInt, []byte(fmt.Sprintf("%d", int64(v.Float()))))
		case reflect.Int,
			reflect.Int8,
			reflect.Int16,
			reflect.Int32,
			reflect.Int64:
			return hash(HintInt, []byte(fmt.Sprintf("%d", v.Int())))
		case reflect.Uint,
			reflect.Uint8,
			reflect.Uint16,
			reflect.Uint32,
			reflect.Uint64:
			return hash(HintInt, []byte(fmt.Sprintf("%d", int64(v.Uint()))))
		}
	case HintUint:
		return hash(HintUint, []byte(fmt.Sprintf("%d", o)))
	case HintBool:
		if v.Bool() {
			return hash(HintBool, []byte{1})
		}
		return hash(HintBool, []byte{0})
	case HintRef:
		switch h := o.(type) {
		case string:
			return Hash(h), nil
		case Hash:
			return h, nil
		}
		return EmptyHash, errors.Error("unsupported type for ref")
	}
	panic(
		fmt.Sprintf("hash: unsupported type %s (%s) for key %s",
			string(ts[0]),
			t.Kind().String(),
			k,
		),
	)
}

// replacing ben's implementation with something less custom, based on:
// * https://github.com/benlaurie/objecthash
// * https://play.golang.org/p/3xraud43pi
// examples of same results in other languages
// * ruby: `[7.30363941192626953125].pack('G').unpack('B*').first`
// * js: `http://weitz.de/ieee`
func hashFloat(f float64) (Hash, error) {
	nf := ""
	switch {
	case math.IsInf(f, 1):
		nf = "Infinity"
	case math.IsInf(f, -1):
		nf = "-Infinity"
	case math.IsNaN(f):
		nf = "NaN"
	default:
		nf = fmt.Sprintf("%x", math.Float64bits(f))
	}

	return hash(HintFloat, []byte(nf))
}
<commit_msg>fix(object): fix uint hashing<commit_after>package object

import (
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"math"
	"reflect"
	"sort"
	"strings"

	"nimona.io/internal/encoding/base58"
	"nimona.io/pkg/errors"
)

// ObjectHash consistently hashes a map.
// It is based on Ben Laurie's object hash, but using the same type hints
// as TJSON instead.

const (
	EmptyHash = Hash("")
)

// SortHashes sorts a slice of Hashes in increasing order, and also returns it.
// The return part is mostly for allowing this to be used as a helper method in
// tests.
func SortHashes(a []Hash) []Hash {
	sort.Sort(HashSlice(a))
	return a
}

type HashSlice []Hash

func (p HashSlice) Len() int           { return len(p) }
func (p HashSlice) Less(i, j int) bool { return p[i] < p[j] }
func (p HashSlice) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
func (p HashSlice) Sort()              { sort.Sort(p) }

func (h Hash) IsEmpty() bool {
	return string(h) == ""
}

func NewHash(o *Object) (Hash, error) {
	m, err := objectToMap(o)
	if err != nil {
		return "", err
	}
	return hashMap(m)
}

func hintsFromKey(k string) []TypeHint {
	ps := strings.Split(k, ":")
	if len(ps) == 1 {
		return nil
	}
	hs := []TypeHint{}
	for _, sh := range ps[1] {
		hs = append(hs, GetTypeHint(string(sh)))
	}
	return hs
}

func hintFromKey(k string) string {
	ps := strings.Split(k, ":")
	if len(ps) == 1 {
		return ""
	}
	if ps[1] == HintObject.String() {
		return HintRef.String()
	}
	if ps[1] == HintArray.String()+HintObject.String() {
		return HintArray.String() + HintRef.String()
	}
	return ps[1]
}

func hashMap(m map[string]interface{}) (Hash, error) {
	if len(m) == 0 {
		return EmptyHash, nil
	}
	b := []byte{}
	ks := []string{}
	for k := range m {
		if len(k) > 0 && k[0] == '_' {
			continue
		}
		ks = append(ks, k)
	}
	sort.Strings(ks)
	for _, k := range ks {
		v := m[k]
		if v == nil {
			continue
		}
		ts := hintsFromKey(k)
		hv, err := hashValueAs(k, v, ts...)
		if err != nil {
			return EmptyHash, err
		}
		if hv.IsEmpty() {
			continue
		}
		// hash the key (including the hint)
		// but if we are dealing with an object, replace its hint with r
		ck := k
		if strings.HasSuffix(k, ":"+HintObject.String()) {
			ck = k[:len(k)-2] + ":" + HintRef.String()
		}
		// TODO consider not replacing hints for maps
		if strings.HasSuffix(k, ":"+HintMap.String()) {
			ck = k[:len(k)-2] + ":" + HintRef.String()
		}
		if strings.HasSuffix(k, ":a"+HintObject.String()) {
			ck = k[:len(k)-3] + ":a" + HintRef.String()
		}
		// TODO consider not replacing hints for maps
		if strings.HasSuffix(k, ":a"+HintMap.String()) {
			ck = k[:len(k)-3] + ":a" + HintRef.String()
		}
		hk, _ := hash("", []byte(ck)) // nolint: errcheck
		b = append(b, hk...)
		b = append(b, hv...)
	}
	if len(b) == 0 {
		return EmptyHash, nil
	}
	return hash(HintRef, b)
}

func hash(p TypeHint, v []byte) (Hash, error) {
	d := sha256.Sum256(append([]byte(p), v...))
	return Hash("oh1." + base58.Encode(d[:])), nil
}

func hashValueAs(k string, o interface{}, ts ...TypeHint) (Hash, error) {
	if o == nil {
		return EmptyHash, nil
	}

	v := reflect.ValueOf(o)
	t := reflect.TypeOf(o)

	if len(ts) == 0 {
		return EmptyHash, nil
	}

	switch ts[0] {
	case HintArray:
		if v.Len() == 0 {
			return EmptyHash, nil
		}
		vs := []interface{}{}
		for i := 0; i < v.Len(); i++ {
			vs = append(vs, v.Index(i).Interface())
		}
		b := []byte{}
		for _, iv := range vs {
			ivv, err := hashValueAs(k, iv, ts[1:]...)
			if err != nil {
				return EmptyHash, err
			}
			b = append(b, ivv...)
		}
		// TODO(geoah) hint SHOULD NOT be array, but array + inner hint
		return hash(TypeHint(hintFromKey(k)), b)
	case HintString:
		s, ok := o.(string)
		if !ok {
			if ss, ok := o.(interface{ String() string }); ok {
				s = ss.String()
			} else {
				return EmptyHash, nil
			}
		}
		if s == "" {
			return EmptyHash, nil
		}
		return hash(HintString, []byte(s))
	case HintData:
		switch t.Kind() {
		case reflect.String:
			d, err := base64.StdEncoding.DecodeString(o.(string))
			if err != nil {
				panic(err)
			}
			return hash(HintData, d)
		case reflect.Slice:
			bo := make([]byte, v.Len())
			switch t.Elem().Kind() {
			case reflect.Int,
				reflect.Int8,
				reflect.Int16,
				reflect.Int32,
				reflect.Int64:
				for i := 0; i < v.Len(); i++ {
					bo[i] = uint8(v.Index(i).Int())
				}
				return hash(HintData, bo)
			case reflect.Uint,
				reflect.Uint8,
				reflect.Uint16,
				reflect.Uint32,
				reflect.Uint64:
				for i := 0; i < v.Len(); i++ {
					bo[i] = uint8(v.Index(i).Uint())
				}
				return hash(HintData, bo)
			case reflect.Interface:
				for i := 0; i < v.Len(); i++ {
					iv := v.Index(i).Interface()
					switch ivv := iv.(type) {
					case uint8:
						bo[i] = ivv
					case uint64:
						bo[i] = uint8(ivv)
					case float64:
						bo[i] = uint8(ivv)
					default:
						panic("data should be some sort of number array, was " +
							t.Elem().Kind().String())
					}
				}
				return hash(HintData, bo)
			default:
				panic("data should be some sort of number array, was " +
					t.Elem().Kind().String())
			}
		}
		return hash(HintData, o.([]byte))
	case HintMap, HintObject:
		switch v := o.(type) {
		case map[string]interface{}:
			if len(v) == 0 {
				return EmptyHash, nil
			}
			h, err := hashMap(v)
			if err != nil {
				return EmptyHash, err
			}
			return h, nil
		case *Object:
			m := v.ToMap()
			if len(m) == 0 {
				return EmptyHash, nil
			}
			h, err := hashMap(m)
			if err != nil {
				return EmptyHash, err
			}
			return h, nil
		default:
			panic("hashing only supports map[string]interface{}")
		}
	case HintFloat:
		switch t.Kind() {
		case reflect.Float32,
			reflect.Float64:
			nf, err := hashFloat(v.Float())
			if err != nil {
				return EmptyHash, err
			}
			return nf, nil
		case reflect.Int,
			reflect.Int8,
			reflect.Int16,
			reflect.Int32,
			reflect.Int64:
			nf, err := hashFloat(float64(v.Int()))
			if err != nil {
				return EmptyHash, err
			}
			return nf, nil
		case reflect.Uint,
			reflect.Uint8,
			reflect.Uint16,
			reflect.Uint32,
			reflect.Uint64:
			nf, err := hashFloat(float64(v.Uint()))
			if err != nil {
				return EmptyHash, err
			}
			return nf, nil
		}
	case HintInt:
		switch t.Kind() {
		case reflect.Float32,
			reflect.Float64:
			return hash(HintInt, []byte(fmt.Sprintf("%d", int64(v.Float()))))
		case reflect.Int,
			reflect.Int8,
			reflect.Int16,
			reflect.Int32,
			reflect.Int64:
			return hash(HintInt, []byte(fmt.Sprintf("%d", v.Int())))
		case reflect.Uint,
			reflect.Uint8,
			reflect.Uint16,
			reflect.Uint32,
			reflect.Uint64:
			return hash(HintInt, []byte(fmt.Sprintf("%d", int64(v.Uint()))))
		}
		panic("unknown type for int")
	case HintUint:
		switch t.Kind() {
		case reflect.Float32,
			reflect.Float64:
			return hash(HintInt, []byte(fmt.Sprintf("%d", uint64(v.Float()))))
		case reflect.Int,
			reflect.Int8,
			reflect.Int16,
			reflect.Int32,
			reflect.Int64:
			return hash(HintInt, []byte(fmt.Sprintf("%d", uint64(v.Int()))))
		case reflect.Uint,
			reflect.Uint8,
			reflect.Uint16,
			reflect.Uint32,
			reflect.Uint64:
			return hash(HintInt, []byte(fmt.Sprintf("%d", uint64(v.Uint()))))
		}
		panic("unknown type for uint")
	case HintBool:
		if v.Bool() {
			return hash(HintBool, []byte{1})
		}
		return hash(HintBool, []byte{0})
	case HintRef:
		switch h := o.(type) {
		case string:
			return Hash(h), nil
		case Hash:
			return h, nil
		}
		return EmptyHash, errors.Error("unsupported type for ref")
	}
	panic(
		fmt.Sprintf("hash: unsupported type %s (%s) for key %s",
			string(ts[0]),
			t.Kind().String(),
			k,
		),
	)
}

// replacing ben's implementation with something less custom, based on:
// * https://github.com/benlaurie/objecthash
// * https://play.golang.org/p/3xraud43pi
// examples of same results in other languages
// * ruby: `[7.30363941192626953125].pack('G').unpack('B*').first`
// * js: `http://weitz.de/ieee`
func hashFloat(f float64) (Hash, error) {
	nf := ""
	switch {
	case math.IsInf(f, 1):
		nf = "Infinity"
	case math.IsInf(f, -1):
		nf = "-Infinity"
	case math.IsNaN(f):
		nf = "NaN"
	default:
		nf = fmt.Sprintf("%x", math.Float64bits(f))
	}

	return hash(HintFloat, []byte(nf))
}
<|endoftext|>