<commit_before>// Package mdns is a multicast dns registry
package registry

import (
	"net"
	"strings"
	"sync"
	"time"

	"github.com/micro/mdns"
	hash "github.com/mitchellh/hashstructure"
)

type mdnsTxt struct {
	Service   string
	Version   string
	Endpoints []*Endpoint
	Metadata  map[string]string
}

type mdnsEntry struct {
	hash uint64
	id   string
	node *mdns.Server
}

type mdnsRegistry struct {
	opts Options

	sync.Mutex
	services map[string][]*mdnsEntry
}

func newRegistry(opts ...Option) Registry {
	options := Options{
		Timeout: time.Millisecond * 100,
	}

	return &mdnsRegistry{
		opts:     options,
		services: make(map[string][]*mdnsEntry),
	}
}

func (m *mdnsRegistry) Init(opts ...Option) error {
	for _, o := range opts {
		o(&m.opts)
	}
	return nil
}

func (m *mdnsRegistry) Options() Options {
	return m.opts
}

func (m *mdnsRegistry) Register(service *Service, opts ...RegisterOption) error {
	m.Lock()
	defer m.Unlock()

	entries, ok := m.services[service.Name]
	// first entry, create wildcard used for list queries
	if !ok {
		s, err := mdns.NewMDNSService(
			service.Name,
			"_services",
			"",
			"",
			9999,
			[]net.IP{net.ParseIP("0.0.0.0")},
			nil,
		)
		if err != nil {
			return err
		}

		srv, err := mdns.NewServer(&mdns.Config{Zone: &mdns.DNSSDService{s}})
		if err != nil {
			return err
		}

		// append the wildcard entry
		entries = append(entries, &mdnsEntry{id: "*", node: srv})
	}

	var gerr error

	for _, node := range service.Nodes {
		// create hash of service; uint64
		h, err := hash.Hash(node, nil)
		if err != nil {
			gerr = err
			continue
		}

		var seen bool
		var e *mdnsEntry

		for _, entry := range entries {
			if node.Id == entry.id {
				seen = true
				e = entry
				break
			}
		}

		// already registered, continue
		if seen && e.hash == h {
			continue
			// hash doesn't match, shutdown
		} else if seen {
			e.node.Shutdown()
			// doesn't exist
		} else {
			e = &mdnsEntry{hash: h}
		}

		txt, err := encode(&mdnsTxt{
			Service:   service.Name,
			Version:   service.Version,
			Endpoints: service.Endpoints,
			Metadata:  node.Metadata,
		})

		if err != nil {
			gerr = err
			continue
		}

		// we got here, new node
		s, err := mdns.NewMDNSService(
			node.Id,
			service.Name,
			"",
			"",
			node.Port,
			[]net.IP{net.ParseIP(node.Address)},
			txt,
		)
		if err != nil {
			gerr = err
			continue
		}

		srv, err := mdns.NewServer(&mdns.Config{Zone: s})
		if err != nil {
			gerr = err
			continue
		}

		e.id = node.Id
		e.node = srv
		entries = append(entries, e)
	}

	// save
	m.services[service.Name] = entries

	return gerr
}

func (m *mdnsRegistry) Deregister(service *Service) error {
	m.Lock()
	defer m.Unlock()

	var newEntries []*mdnsEntry

	// loop existing entries, check if any match, shutdown those that do
	for _, entry := range m.services[service.Name] {
		var remove bool

		for _, node := range service.Nodes {
			if node.Id == entry.id {
				entry.node.Shutdown()
				remove = true
				break
			}
		}

		// keep it?
		if !remove {
			newEntries = append(newEntries, entry)
		}
	}

	// last entry is the wildcard for list queries. Remove it.
	if len(newEntries) == 1 && newEntries[0].id == "*" {
		newEntries[0].node.Shutdown()
		delete(m.services, service.Name)
	} else {
		m.services[service.Name] = newEntries
	}

	return nil
}

func (m *mdnsRegistry) GetService(service string) ([]*Service, error) {
	p := mdns.DefaultParams(service)
	p.Timeout = m.opts.Timeout
	entryCh := make(chan *mdns.ServiceEntry, 10)
	p.Entries = entryCh

	exit := make(chan bool)
	defer close(exit)

	serviceMap := make(map[string]*Service)

	go func() {
		for {
			select {
			case e := <-entryCh:
				// list record so skip
				if p.Service == "_services" {
					continue
				}

				if e.TTL == 0 {
					continue
				}

				txt, err := decode(e.InfoFields)
				if err != nil {
					continue
				}

				if txt.Service != service {
					continue
				}

				s, ok := serviceMap[txt.Version]
				if !ok {
					s = &Service{
						Name:      txt.Service,
						Version:   txt.Version,
						Endpoints: txt.Endpoints,
					}
				}

				s.Nodes = append(s.Nodes, &Node{
					Id:       strings.TrimSuffix(e.Name, "."+p.Service+"."+p.Domain+"."),
					Address:  e.AddrV4.String(),
					Port:     e.Port,
					Metadata: txt.Metadata,
				})

				serviceMap[txt.Version] = s
			case <-exit:
				return
			}
		}
	}()

	if err := mdns.Query(p); err != nil {
		return nil, err
	}

	// create list and return
	var services []*Service

	for _, service := range serviceMap {
		services = append(services, service)
	}

	return services, nil
}

func (m *mdnsRegistry) ListServices() ([]*Service, error) {
	p := mdns.DefaultParams("_services")
	p.Timeout = m.opts.Timeout
	entryCh := make(chan *mdns.ServiceEntry, 10)
	p.Entries = entryCh

	exit := make(chan bool)
	defer close(exit)

	serviceMap := make(map[string]bool)
	var services []*Service

	go func() {
		for {
			select {
			case e := <-entryCh:
				if e.TTL == 0 {
					continue
				}

				name := strings.TrimSuffix(e.Name, "."+p.Service+"."+p.Domain+".")
				if !serviceMap[name] {
					serviceMap[name] = true
					services = append(services, &Service{Name: name})
				}
			case <-exit:
				return
			}
		}
	}()

	if err := mdns.Query(p); err != nil {
		return nil, err
	}

	return services, nil
}

func (m *mdnsRegistry) Watch(opts ...WatchOption) (Watcher, error) {
	var wo WatchOptions
	for _, o := range opts {
		o(&wo)
	}

	md := &mdnsWatcher{
		wo:   wo,
		ch:   make(chan *mdns.ServiceEntry, 32),
		exit: make(chan struct{}),
	}

	go func() {
		if err := mdns.Listen(md.ch, md.exit); err != nil {
			md.Stop()
		}
	}()

	return md, nil
}

func (m *mdnsRegistry) String() string {
	return "mdns"
}

// NewRegistry returns a new default registry which is mdns
func NewRegistry(opts ...Option) Registry {
	return newRegistry(opts...)
}
<commit_msg>update mdns to remove race condition<commit_after>// Package mdns is a multicast dns registry
package registry

import (
	"context"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/micro/mdns"
	hash "github.com/mitchellh/hashstructure"
)

type mdnsTxt struct {
	Service   string
	Version   string
	Endpoints []*Endpoint
	Metadata  map[string]string
}

type mdnsEntry struct {
	hash uint64
	id   string
	node *mdns.Server
}

type mdnsRegistry struct {
	opts Options

	sync.Mutex
	services map[string][]*mdnsEntry
}

func newRegistry(opts ...Option) Registry {
	options := Options{
		Timeout: time.Millisecond * 100,
	}

	return &mdnsRegistry{
		opts:     options,
		services: make(map[string][]*mdnsEntry),
	}
}

func (m *mdnsRegistry) Init(opts ...Option) error {
	for _, o := range opts {
		o(&m.opts)
	}
	return nil
}

func (m *mdnsRegistry) Options() Options {
	return m.opts
}

func (m *mdnsRegistry) Register(service *Service, opts ...RegisterOption) error {
	m.Lock()
	defer m.Unlock()

	entries, ok := m.services[service.Name]
	// first entry, create wildcard used for list queries
	if !ok {
		s, err := mdns.NewMDNSService(
			service.Name,
			"_services",
			"",
			"",
			9999,
			[]net.IP{net.ParseIP("0.0.0.0")},
			nil,
		)
		if err != nil {
			return err
		}

		srv, err := mdns.NewServer(&mdns.Config{Zone: &mdns.DNSSDService{s}})
		if err != nil {
			return err
		}

		// append the wildcard entry
		entries = append(entries, &mdnsEntry{id: "*", node: srv})
	}

	var gerr error

	for _, node := range service.Nodes {
		// create hash of service; uint64
		h, err := hash.Hash(node, nil)
		if err != nil {
			gerr = err
			continue
		}

		var seen bool
		var e *mdnsEntry

		for _, entry := range entries {
			if node.Id == entry.id {
				seen = true
				e = entry
				break
			}
		}

		// already registered, continue
		if seen && e.hash == h {
			continue
			// hash doesn't match, shutdown
		} else if seen {
			e.node.Shutdown()
			// doesn't exist
		} else {
			e = &mdnsEntry{hash: h}
		}

		txt, err := encode(&mdnsTxt{
			Service:   service.Name,
			Version:   service.Version,
			Endpoints: service.Endpoints,
			Metadata:  node.Metadata,
		})

		if err != nil {
			gerr = err
			continue
		}

		// we got here, new node
		s, err := mdns.NewMDNSService(
			node.Id,
			service.Name,
			"",
			"",
			node.Port,
			[]net.IP{net.ParseIP(node.Address)},
			txt,
		)
		if err != nil {
			gerr = err
			continue
		}

		srv, err := mdns.NewServer(&mdns.Config{Zone: s})
		if err != nil {
			gerr = err
			continue
		}

		e.id = node.Id
		e.node = srv
		entries = append(entries, e)
	}

	// save
	m.services[service.Name] = entries

	return gerr
}

func (m *mdnsRegistry) Deregister(service *Service) error {
	m.Lock()
	defer m.Unlock()

	var newEntries []*mdnsEntry

	// loop existing entries, check if any match, shutdown those that do
	for _, entry := range m.services[service.Name] {
		var remove bool

		for _, node := range service.Nodes {
			if node.Id == entry.id {
				entry.node.Shutdown()
				remove = true
				break
			}
		}

		// keep it?
		if !remove {
			newEntries = append(newEntries, entry)
		}
	}

	// last entry is the wildcard for list queries. Remove it.
	if len(newEntries) == 1 && newEntries[0].id == "*" {
		newEntries[0].node.Shutdown()
		delete(m.services, service.Name)
	} else {
		m.services[service.Name] = newEntries
	}

	return nil
}

func (m *mdnsRegistry) GetService(service string) ([]*Service, error) {
	serviceMap := make(map[string]*Service)
	entries := make(chan *mdns.ServiceEntry, 10)
	done := make(chan bool)

	p := mdns.DefaultParams(service)
	// set context with timeout
	p.Context, _ = context.WithTimeout(context.Background(), m.opts.Timeout)
	// set entries channel
	p.Entries = entries

	go func() {
		for {
			select {
			case e := <-entries:
				// list record so skip
				if p.Service == "_services" {
					continue
				}

				if e.TTL == 0 {
					continue
				}

				txt, err := decode(e.InfoFields)
				if err != nil {
					continue
				}

				if txt.Service != service {
					continue
				}

				s, ok := serviceMap[txt.Version]
				if !ok {
					s = &Service{
						Name:      txt.Service,
						Version:   txt.Version,
						Endpoints: txt.Endpoints,
					}
				}

				s.Nodes = append(s.Nodes, &Node{
					Id:       strings.TrimSuffix(e.Name, "."+p.Service+"."+p.Domain+"."),
					Address:  e.AddrV4.String(),
					Port:     e.Port,
					Metadata: txt.Metadata,
				})

				serviceMap[txt.Version] = s
			case <-p.Context.Done():
				close(done)
				return
			}
		}
	}()

	// execute the query
	if err := mdns.Query(p); err != nil {
		return nil, err
	}

	// wait for completion
	<-done

	// create list and return
	var services []*Service

	for _, service := range serviceMap {
		services = append(services, service)
	}

	return services, nil
}

func (m *mdnsRegistry) ListServices() ([]*Service, error) {
	serviceMap := make(map[string]bool)
	entries := make(chan *mdns.ServiceEntry, 10)
	done := make(chan bool)

	p := mdns.DefaultParams("_services")
	// set context with timeout
	p.Context, _ = context.WithTimeout(context.Background(), m.opts.Timeout)
	// set entries channel
	p.Entries = entries

	var services []*Service

	go func() {
		for {
			select {
			case e := <-entries:
				if e.TTL == 0 {
					continue
				}

				name := strings.TrimSuffix(e.Name, "."+p.Service+"."+p.Domain+".")
				if !serviceMap[name] {
					serviceMap[name] = true
					services = append(services, &Service{Name: name})
				}
			case <-p.Context.Done():
				close(done)
				return
			}
		}
	}()

	// execute query
	if err := mdns.Query(p); err != nil {
		return nil, err
	}

	// wait till done
	<-done

	return services, nil
}

func (m *mdnsRegistry) Watch(opts ...WatchOption) (Watcher, error) {
	var wo WatchOptions
	for _, o := range opts {
		o(&wo)
	}

	md := &mdnsWatcher{
		wo:   wo,
		ch:   make(chan *mdns.ServiceEntry, 32),
		exit: make(chan struct{}),
	}

	go func() {
		if err := mdns.Listen(md.ch, md.exit); err != nil {
			md.Stop()
		}
	}()

	return md, nil
}

func (m *mdnsRegistry) String() string {
	return "mdns"
}

// NewRegistry returns a new default registry which is mdns
func NewRegistry(opts ...Option) Registry {
	return newRegistry(opts...)
}
<|endoftext|>


<commit_before>package pstoremem

import (
	"context"
	"sort"
	"sync"
	"time"

	logging "github.com/ipfs/go-log"
	peer "github.com/libp2p/go-libp2p-peer"
	ma "github.com/multiformats/go-multiaddr"

	pstore "github.com/libp2p/go-libp2p-peerstore"
	addr "github.com/libp2p/go-libp2p-peerstore/addr"
)

var log = logging.Logger("peerstore")

type expiringAddr struct {
	Addr    ma.Multiaddr
	TTL     time.Duration
	Expires time.Time
}

func (e *expiringAddr) ExpiredBy(t time.Time) bool {
	return t.After(e.Expires)
}

var _ pstore.AddrBook = (*memoryAddrBook)(nil)

// memoryAddrBook manages addresses.
type memoryAddrBook struct {
	addrmu sync.RWMutex
	// Use pointers to save memory. Maps always leave some fraction of their
	// space unused. storing the *values* directly in the map will
	// drastically increase the space waste. In our case, by 6x.
	addrs map[peer.ID]map[string]*expiringAddr

	nextGC time.Time
	ctx    context.Context
	cancel func()

	subManager *AddrSubManager
}

func NewAddrBook() pstore.AddrBook {
	ctx, cancel := context.WithCancel(context.Background())

	ab := &memoryAddrBook{
		addrs:      make(map[peer.ID]map[string]*expiringAddr),
		subManager: NewAddrSubManager(),
		ctx:        ctx,
		cancel:     cancel,
	}

	go ab.background()
	return ab
}

// background periodically schedules a gc
func (mab *memoryAddrBook) background() {
	ticker := time.NewTicker(1 * time.Hour)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			mab.addrmu.Lock()
			mab.gc()
			mab.addrmu.Unlock()

		case <-mab.ctx.Done():
			return
		}
	}
}

// gc garbage collects the in-memory address book. The caller *must* hold the addrmu lock.
func (mab *memoryAddrBook) gc() {
	now := time.Now()
	if !now.After(mab.nextGC) {
		return
	}
	for p, amap := range mab.addrs {
		for k, addr := range amap {
			if addr.ExpiredBy(now) {
				delete(amap, k)
			}
		}
		if len(amap) == 0 {
			delete(mab.addrs, p)
		}
	}
	mab.nextGC = time.Now().Add(pstore.AddressTTL)
}

func (mab *memoryAddrBook) PeersWithAddrs() peer.IDSlice {
	mab.addrmu.RLock()
	defer mab.addrmu.RUnlock()

	pids := make(peer.IDSlice, 0, len(mab.addrs))
	for pid := range mab.addrs {
		pids = append(pids, pid)
	}
	return pids
}

// AddAddr calls AddAddrs(p, []ma.Multiaddr{addr}, ttl)
func (mab *memoryAddrBook) AddAddr(p peer.ID, addr ma.Multiaddr, ttl time.Duration) {
	mab.AddAddrs(p, []ma.Multiaddr{addr}, ttl)
}

// AddAddrs gives memoryAddrBook addresses to use, with a given ttl
// (time-to-live), after which the address is no longer valid.
// If the manager has a longer TTL, the operation is a no-op for that address
func (mab *memoryAddrBook) AddAddrs(p peer.ID, addrs []ma.Multiaddr, ttl time.Duration) {
	mab.addrmu.Lock()
	defer mab.addrmu.Unlock()

	// if ttl is zero, exit. nothing to do.
	if ttl <= 0 {
		return
	}

	amap := mab.addrs[p]
	if amap == nil {
		amap = make(map[string]*expiringAddr, len(addrs))
		mab.addrs[p] = amap
	}
	exp := time.Now().Add(ttl)
	for _, addr := range addrs {
		if addr == nil {
			log.Warningf("was passed nil multiaddr for %s", p)
			continue
		}
		addrstr := string(addr.Bytes())
		a, found := amap[addrstr]
		if !found || exp.After(a.Expires) {
			amap[addrstr] = &expiringAddr{Addr: addr, Expires: exp, TTL: ttl}

			mab.subManager.BroadcastAddr(p, addr)
		}
	}
	mab.gc()
}

// SetAddr calls mgr.SetAddrs(p, addr, ttl)
func (mab *memoryAddrBook) SetAddr(p peer.ID, addr ma.Multiaddr, ttl time.Duration) {
	mab.SetAddrs(p, []ma.Multiaddr{addr}, ttl)
}

// SetAddrs sets the ttl on addresses. This clears any TTL there previously.
// This is used when we receive the best estimate of the validity of an address.
func (mab *memoryAddrBook) SetAddrs(p peer.ID, addrs []ma.Multiaddr, ttl time.Duration) {
	mab.addrmu.Lock()
	defer mab.addrmu.Unlock()

	amap := mab.addrs[p]
	if amap == nil {
		amap = make(map[string]*expiringAddr, len(addrs))
		mab.addrs[p] = amap
	}

	exp := time.Now().Add(ttl)
	for _, addr := range addrs {
		if addr == nil {
			log.Warningf("was passed nil multiaddr for %s", p)
			continue
		}
		// re-set all of them for new ttl.
		addrstr := string(addr.Bytes())

		if ttl > 0 {
			amap[addrstr] = &expiringAddr{Addr: addr, Expires: exp, TTL: ttl}

			mab.subManager.BroadcastAddr(p, addr)
		} else {
			delete(amap, addrstr)
		}
	}
	mab.gc()
}

// UpdateAddrs updates the addresses associated with the given peer that have
// the given oldTTL to have the given newTTL.
func (mab *memoryAddrBook) UpdateAddrs(p peer.ID, oldTTL time.Duration, newTTL time.Duration) {
	mab.addrmu.Lock()
	defer mab.addrmu.Unlock()

	amap, found := mab.addrs[p]
	if !found {
		return
	}

	exp := time.Now().Add(newTTL)
	for k, addr := range amap {
		if oldTTL == addr.TTL {
			addr.TTL = newTTL
			addr.Expires = exp
			amap[k] = addr
		}
	}
	mab.gc()
}

// Addresses returns all known (and valid) addresses for a given
func (mab *memoryAddrBook) Addrs(p peer.ID) []ma.Multiaddr {
	mab.addrmu.RLock()
	defer mab.addrmu.RUnlock()

	amap, found := mab.addrs[p]
	if !found {
		return nil
	}

	now := time.Now()
	good := make([]ma.Multiaddr, 0, len(amap))
	for _, m := range amap {
		if !m.ExpiredBy(now) {
			good = append(good, m.Addr)
		}
	}

	return good
}

// ClearAddrs removes all previously stored addresses
func (mab *memoryAddrBook) ClearAddrs(p peer.ID) {
	mab.addrmu.Lock()
	defer mab.addrmu.Unlock()

	delete(mab.addrs, p)
}

// AddrStream returns a channel on which all new addresses discovered for a
// given peer ID will be published.
func (mab *memoryAddrBook) AddrStream(ctx context.Context, p peer.ID) <-chan ma.Multiaddr {
	mab.addrmu.RLock()
	defer mab.addrmu.RUnlock()

	baseaddrslice := mab.addrs[p]
	initial := make([]ma.Multiaddr, 0, len(baseaddrslice))
	for _, a := range baseaddrslice {
		initial = append(initial, a.Addr)
	}

	return mab.subManager.AddrStream(ctx, p, initial)
}

type addrSub struct {
	pubch  chan ma.Multiaddr
	lk     sync.Mutex
	buffer []ma.Multiaddr
	ctx    context.Context
}

func (s *addrSub) pubAddr(a ma.Multiaddr) {
	select {
	case s.pubch <- a:
	case <-s.ctx.Done():
	}
}

// An abstracted, pub-sub manager for address streams. Extracted from
// memoryAddrBook in order to support additional implementations.
type AddrSubManager struct {
	mu   sync.RWMutex
	subs map[peer.ID][]*addrSub
}

// NewAddrSubManager initializes an AddrSubManager.
func NewAddrSubManager() *AddrSubManager {
	return &AddrSubManager{
		subs: make(map[peer.ID][]*addrSub),
	}
}

// Used internally by the address stream coroutine to remove a subscription
// from the manager.
func (mgr *AddrSubManager) removeSub(p peer.ID, s *addrSub) {
	mgr.mu.Lock()
	defer mgr.mu.Unlock()

	subs := mgr.subs[p]
	if len(subs) == 1 {
		if subs[0] != s {
			return
		}
		delete(mgr.subs, p)
		return
	}

	for i, v := range subs {
		if v == s {
			subs[i] = subs[len(subs)-1]
			subs[len(subs)-1] = nil
			mgr.subs[p] = subs[:len(subs)-1]
			return
		}
	}
}

// BroadcastAddr broadcasts a new address to all subscribed streams.
func (mgr *AddrSubManager) BroadcastAddr(p peer.ID, addr ma.Multiaddr) {
	mgr.mu.RLock()
	defer mgr.mu.RUnlock()

	if subs, ok := mgr.subs[p]; ok {
		for _, sub := range subs {
			sub.pubAddr(addr)
		}
	}
}

// AddrStream creates a new subscription for a given peer ID, pre-populating the
// channel with any addresses we might already have on file.
func (mgr *AddrSubManager) AddrStream(ctx context.Context, p peer.ID, initial []ma.Multiaddr) <-chan ma.Multiaddr {
	sub := &addrSub{pubch: make(chan ma.Multiaddr), ctx: ctx}
	out := make(chan ma.Multiaddr)

	mgr.mu.Lock()
	if _, ok := mgr.subs[p]; ok {
		mgr.subs[p] = append(mgr.subs[p], sub)
	} else {
		mgr.subs[p] = []*addrSub{sub}
	}
	mgr.mu.Unlock()

	sort.Sort(addr.AddrList(initial))

	go func(buffer []ma.Multiaddr) {
		defer close(out)

		sent := make(map[string]bool, len(buffer))
		var outch chan ma.Multiaddr

		for _, a := range buffer {
			sent[string(a.Bytes())] = true
		}

		var next ma.Multiaddr
		if len(buffer) > 0 {
			next = buffer[0]
			buffer = buffer[1:]
			outch = out
		}

		for {
			select {
			case outch <- next:
				if len(buffer) > 0 {
					next = buffer[0]
					buffer = buffer[1:]
				} else {
					outch = nil
					next = nil
				}
			case naddr := <-sub.pubch:
				if sent[string(naddr.Bytes())] {
					continue
				}

				sent[string(naddr.Bytes())] = true
				if next == nil {
					next = naddr
					outch = out
				} else {
					buffer = append(buffer, naddr)
				}
			case <-ctx.Done():
				mgr.removeSub(p, sub)
				return
			}
		}

	}(initial)

	return out
}
<commit_msg>add Close method to memory address book, which cancels the gc process<commit_after>package pstoremem

import (
	"context"
	"sort"
	"sync"
	"time"

	logging "github.com/ipfs/go-log"
	peer "github.com/libp2p/go-libp2p-peer"
	ma "github.com/multiformats/go-multiaddr"

	pstore "github.com/libp2p/go-libp2p-peerstore"
	addr "github.com/libp2p/go-libp2p-peerstore/addr"
)

var log = logging.Logger("peerstore")

type expiringAddr struct {
	Addr    ma.Multiaddr
	TTL     time.Duration
	Expires time.Time
}

func (e *expiringAddr) ExpiredBy(t time.Time) bool {
	return t.After(e.Expires)
}

var _ pstore.AddrBook = (*memoryAddrBook)(nil)

// memoryAddrBook manages addresses.
type memoryAddrBook struct {
	addrmu sync.RWMutex
	// Use pointers to save memory. Maps always leave some fraction of their
	// space unused. storing the *values* directly in the map will
	// drastically increase the space waste. In our case, by 6x.
	addrs map[peer.ID]map[string]*expiringAddr

	nextGC time.Time
	ctx    context.Context
	cancel func()

	subManager *AddrSubManager
}

func NewAddrBook() pstore.AddrBook {
	ctx, cancel := context.WithCancel(context.Background())

	ab := &memoryAddrBook{
		addrs:      make(map[peer.ID]map[string]*expiringAddr),
		subManager: NewAddrSubManager(),
		ctx:        ctx,
		cancel:     cancel,
	}

	go ab.background()
	return ab
}

// background periodically schedules a gc
func (mab *memoryAddrBook) background() {
	ticker := time.NewTicker(1 * time.Hour)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			mab.addrmu.Lock()
			mab.gc()
			mab.addrmu.Unlock()

		case <-mab.ctx.Done():
			return
		}
	}
}

func (mab *memoryAddrBook) Close() error {
	mab.cancel()
	return nil
}

// gc garbage collects the in-memory address book. The caller *must* hold the addrmu lock.
func (mab *memoryAddrBook) gc() {
	now := time.Now()
	if !now.After(mab.nextGC) {
		return
	}
	for p, amap := range mab.addrs {
		for k, addr := range amap {
			if addr.ExpiredBy(now) {
				delete(amap, k)
			}
		}
		if len(amap) == 0 {
			delete(mab.addrs, p)
		}
	}
	mab.nextGC = time.Now().Add(pstore.AddressTTL)
}

func (mab *memoryAddrBook) PeersWithAddrs() peer.IDSlice {
	mab.addrmu.RLock()
	defer mab.addrmu.RUnlock()

	pids := make(peer.IDSlice, 0, len(mab.addrs))
	for pid := range mab.addrs {
		pids = append(pids, pid)
	}
	return pids
}

// AddAddr calls AddAddrs(p, []ma.Multiaddr{addr}, ttl)
func (mab *memoryAddrBook) AddAddr(p peer.ID, addr ma.Multiaddr, ttl time.Duration) {
	mab.AddAddrs(p, []ma.Multiaddr{addr}, ttl)
}

// AddAddrs gives memoryAddrBook addresses to use, with a given ttl
// (time-to-live), after which the address is no longer valid.
// If the manager has a longer TTL, the operation is a no-op for that address
func (mab *memoryAddrBook) AddAddrs(p peer.ID, addrs []ma.Multiaddr, ttl time.Duration) {
	mab.addrmu.Lock()
	defer mab.addrmu.Unlock()

	// if ttl is zero, exit. nothing to do.
	if ttl <= 0 {
		return
	}

	amap := mab.addrs[p]
	if amap == nil {
		amap = make(map[string]*expiringAddr, len(addrs))
		mab.addrs[p] = amap
	}
	exp := time.Now().Add(ttl)
	for _, addr := range addrs {
		if addr == nil {
			log.Warningf("was passed nil multiaddr for %s", p)
			continue
		}
		addrstr := string(addr.Bytes())
		a, found := amap[addrstr]
		if !found || exp.After(a.Expires) {
			amap[addrstr] = &expiringAddr{Addr: addr, Expires: exp, TTL: ttl}

			mab.subManager.BroadcastAddr(p, addr)
		}
	}
	mab.gc()
}

// SetAddr calls mgr.SetAddrs(p, addr, ttl)
func (mab *memoryAddrBook) SetAddr(p peer.ID, addr ma.Multiaddr, ttl time.Duration) {
	mab.SetAddrs(p, []ma.Multiaddr{addr}, ttl)
}

// SetAddrs sets the ttl on addresses. This clears any TTL there previously.
// This is used when we receive the best estimate of the validity of an address.
func (mab *memoryAddrBook) SetAddrs(p peer.ID, addrs []ma.Multiaddr, ttl time.Duration) {
	mab.addrmu.Lock()
	defer mab.addrmu.Unlock()

	amap := mab.addrs[p]
	if amap == nil {
		amap = make(map[string]*expiringAddr, len(addrs))
		mab.addrs[p] = amap
	}

	exp := time.Now().Add(ttl)
	for _, addr := range addrs {
		if addr == nil {
			log.Warningf("was passed nil multiaddr for %s", p)
			continue
		}
		// re-set all of them for new ttl.
		addrstr := string(addr.Bytes())

		if ttl > 0 {
			amap[addrstr] = &expiringAddr{Addr: addr, Expires: exp, TTL: ttl}

			mab.subManager.BroadcastAddr(p, addr)
		} else {
			delete(amap, addrstr)
		}
	}
	mab.gc()
}

// UpdateAddrs updates the addresses associated with the given peer that have
// the given oldTTL to have the given newTTL.
func (mab *memoryAddrBook) UpdateAddrs(p peer.ID, oldTTL time.Duration, newTTL time.Duration) {
	mab.addrmu.Lock()
	defer mab.addrmu.Unlock()

	amap, found := mab.addrs[p]
	if !found {
		return
	}

	exp := time.Now().Add(newTTL)
	for k, addr := range amap {
		if oldTTL == addr.TTL {
			addr.TTL = newTTL
			addr.Expires = exp
			amap[k] = addr
		}
	}
	mab.gc()
}

// Addresses returns all known (and valid) addresses for a given
func (mab *memoryAddrBook) Addrs(p peer.ID) []ma.Multiaddr {
	mab.addrmu.RLock()
	defer mab.addrmu.RUnlock()

	amap, found := mab.addrs[p]
	if !found {
		return nil
	}

	now := time.Now()
	good := make([]ma.Multiaddr, 0, len(amap))
	for _, m := range amap {
		if !m.ExpiredBy(now) {
			good = append(good, m.Addr)
		}
	}

	return good
}

// ClearAddrs removes all previously stored addresses
func (mab *memoryAddrBook) ClearAddrs(p peer.ID) {
	mab.addrmu.Lock()
	defer mab.addrmu.Unlock()

	delete(mab.addrs, p)
}

// AddrStream returns a channel on which all new addresses discovered for a
// given peer ID will be published.
func (mab *memoryAddrBook) AddrStream(ctx context.Context, p peer.ID) <-chan ma.Multiaddr {
	mab.addrmu.RLock()
	defer mab.addrmu.RUnlock()

	baseaddrslice := mab.addrs[p]
	initial := make([]ma.Multiaddr, 0, len(baseaddrslice))
	for _, a := range baseaddrslice {
		initial = append(initial, a.Addr)
	}

	return mab.subManager.AddrStream(ctx, p, initial)
}

type addrSub struct {
	pubch  chan ma.Multiaddr
	lk     sync.Mutex
	buffer []ma.Multiaddr
	ctx    context.Context
}

func (s *addrSub) pubAddr(a ma.Multiaddr) {
	select {
	case s.pubch <- a:
	case <-s.ctx.Done():
	}
}

// An abstracted, pub-sub manager for address streams. Extracted from
// memoryAddrBook in order to support additional implementations.
type AddrSubManager struct {
	mu   sync.RWMutex
	subs map[peer.ID][]*addrSub
}

// NewAddrSubManager initializes an AddrSubManager.
func NewAddrSubManager() *AddrSubManager {
	return &AddrSubManager{
		subs: make(map[peer.ID][]*addrSub),
	}
}

// Used internally by the address stream coroutine to remove a subscription
// from the manager.
func (mgr *AddrSubManager) removeSub(p peer.ID, s *addrSub) {
	mgr.mu.Lock()
	defer mgr.mu.Unlock()

	subs := mgr.subs[p]
	if len(subs) == 1 {
		if subs[0] != s {
			return
		}
		delete(mgr.subs, p)
		return
	}

	for i, v := range subs {
		if v == s {
			subs[i] = subs[len(subs)-1]
			subs[len(subs)-1] = nil
			mgr.subs[p] = subs[:len(subs)-1]
			return
		}
	}
}

// BroadcastAddr broadcasts a new address to all subscribed streams.
func (mgr *AddrSubManager) BroadcastAddr(p peer.ID, addr ma.Multiaddr) {
	mgr.mu.RLock()
	defer mgr.mu.RUnlock()

	if subs, ok := mgr.subs[p]; ok {
		for _, sub := range subs {
			sub.pubAddr(addr)
		}
	}
}

// AddrStream creates a new subscription for a given peer ID, pre-populating the
// channel with any addresses we might already have on file.
func (mgr *AddrSubManager) AddrStream(ctx context.Context, p peer.ID, initial []ma.Multiaddr) <-chan ma.Multiaddr {
	sub := &addrSub{pubch: make(chan ma.Multiaddr), ctx: ctx}
	out := make(chan ma.Multiaddr)

	mgr.mu.Lock()
	if _, ok := mgr.subs[p]; ok {
		mgr.subs[p] = append(mgr.subs[p], sub)
	} else {
		mgr.subs[p] = []*addrSub{sub}
	}
	mgr.mu.Unlock()

	sort.Sort(addr.AddrList(initial))

	go func(buffer []ma.Multiaddr) {
		defer close(out)

		sent := make(map[string]bool, len(buffer))
		var outch chan ma.Multiaddr

		for _, a := range buffer {
			sent[string(a.Bytes())] = true
		}

		var next ma.Multiaddr
		if len(buffer) > 0 {
			next = buffer[0]
			buffer = buffer[1:]
			outch = out
		}

		for {
			select {
			case outch <- next:
				if len(buffer) > 0 {
					next = buffer[0]
					buffer = buffer[1:]
				} else {
					outch = nil
					next = nil
				}
			case naddr := <-sub.pubch:
				if sent[string(naddr.Bytes())] {
					continue
				}

				sent[string(naddr.Bytes())] = true
				if next == nil {
					next = naddr
					outch = out
				} else {
					buffer = append(buffer, naddr)
				}
			case <-ctx.Done():
				mgr.removeSub(p, sub)
				return
			}
		}

	}(initial)

	return out
}
<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
	"log"
	"net/http"
	"strings"

	"github.com/golang/protobuf/proto"
)

// smtpdNewClientHandler handles /smtpdnewclient
// If smtpd client IP is 127.0.0.1 connection is rejected
func smtpdNewClientHandler(w http.ResponseWriter, r *http.Request) {
	// read request body, ie a protobuf serialized SmtpdNewClientMsg
	data, err := ioutil.ReadAll(r.Body)
	if err != nil {
		// for the poc simply return HTTP 500 on error
		w.WriteHeader(500)
		w.Write([]byte(err.Error()))
	}

	newClientMsg := &SmtpdNewClientMsg{}
	err = proto.Unmarshal(data, newClientMsg)
	if err != nil {
		w.WriteHeader(422)
		w.Write([]byte(err.Error()))
		return
	}

	// init response
	smtpResponse := &SmtpdResponse{
		SmtpCode: proto.Int32(0),
		SmtpMsg:  proto.String(""),
	}

	// test client IP (ip:port)
	ipPort := strings.Split(newClientMsg.GetRemoteIp(), ":")
	//  bad format ?
	if len(ipPort) != 2 {
		w.WriteHeader(422)
		w.Write([]byte("422 - Bad remote IP format. Expected ip:port. Got: " + newClientMsg.GetRemoteIp()))
		return
	}

	if ipPort[0] == "127.0.0.1" {
		// return SMTP permFail
		smtpResponse.SmtpCode = proto.Int32(570)
		smtpResponse.SmtpMsg = proto.String("sorry you are not allowed to speak to me")
		// Close connection
		smtpResponse.CloseConnection = proto.Bool(true)
	}
	data, err = proto.Marshal(smtpResponse)
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte(err.Error()))
	}
	w.Write(data)
}

func main() {
	http.HandleFunc("/smtpdnewclient", smtpdNewClientHandler)
	log.Fatalln(http.ListenAndServe(":3333", nil))
}
<commit_msg>missing return<commit_after>package main

import (
	"io/ioutil"
	"log"
	"net/http"
	"strings"

	"github.com/golang/protobuf/proto"
)

// smtpdNewClientHandler handles /smtpdnewclient
// If smtpd client IP is 127.0.0.1 connection is rejected
func smtpdNewClientHandler(w http.ResponseWriter, r *http.Request) {
	// read request body, ie a protobuf serialized SmtpdNewClientMsg
	data, err := ioutil.ReadAll(r.Body)
	if err != nil {
		// for the poc simply return HTTP 500 on error
		w.WriteHeader(500)
		w.Write([]byte(err.Error()))
	}

	newClientMsg := &SmtpdNewClientMsg{}
	err = proto.Unmarshal(data, newClientMsg)
	if err != nil {
		w.WriteHeader(422)
		w.Write([]byte(err.Error()))
		return
	}

	// init response
	smtpResponse := &SmtpdResponse{
		SmtpCode: proto.Int32(0),
		SmtpMsg:  proto.String(""),
	}

	// test client IP (ip:port)
	ipPort := strings.Split(newClientMsg.GetRemoteIp(), ":")
	//  bad format ?
	if len(ipPort) != 2 {
		w.WriteHeader(422)
		w.Write([]byte("422 - Bad remote IP format. Expected ip:port. Got: " + newClientMsg.GetRemoteIp()))
		return
	}

	if ipPort[0] == "127.0.0.1" {
		// return SMTP permFail
		smtpResponse.SmtpCode = proto.Int32(570)
		smtpResponse.SmtpMsg = proto.String("sorry you are not allowed to speak to me")
		// Close connection
		smtpResponse.CloseConnection = proto.Bool(true)
	}
	data, err = proto.Marshal(smtpResponse)
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte(err.Error()))
		return
	}
	w.Write(data)
}

func main() {
	http.HandleFunc("/smtpdnewclient", smtpdNewClientHandler)
	log.Fatalln(http.ListenAndServe(":3333", nil))
}
<|endoftext|>


<commit_before>package registry

import (
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"testing"

	"github.com/docker/docker/utils"
)

var (
	token = []string{"fake-token"}
)

const (
	imageID = "42d718c941f5c532ac049bf0b0ab53f0062f09a03afd4aa4a02c098e46032b9d"
	REPO    = "foo42/bar"
)

func spawnTestRegistrySession(t *testing.T) *Session {
	authConfig := &AuthConfig{}
	endpoint, err := NewEndpoint(makeURL("/v1/"), false)
	if err != nil {
		t.Fatal(err)
	}
	r, err := NewSession(authConfig, utils.NewHTTPRequestFactory(), endpoint, true)
	if err != nil {
		t.Fatal(err)
	}
	return r
}

func TestPingRegistryEndpoint(t *testing.T) {
	ep, err := NewEndpoint(makeURL("/v1/"), false)
	if err != nil {
		t.Fatal(err)
	}
	regInfo, err := ep.Ping()
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, regInfo.Standalone, true, "Expected standalone to be true (default)")
}

func TestGetRemoteHistory(t *testing.T) {
	r := spawnTestRegistrySession(t)
	hist, err := r.GetRemoteHistory(imageID, makeURL("/v1/"), token)
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, len(hist), 2, "Expected 2 images in history")
	assertEqual(t, hist[0], imageID, "Expected "+imageID+"as first ancestry")
	assertEqual(t, hist[1], "77dbf71da1d00e3fbddc480176eac8994025630c6590d11cfc8fe1209c2a1d20",
		"Unexpected second ancestry")
}

func TestLookupRemoteImage(t *testing.T) {
	r := spawnTestRegistrySession(t)
	found := r.LookupRemoteImage(imageID, makeURL("/v1/"), token)
	assertEqual(t, found, true, "Expected remote lookup to succeed")
	found = r.LookupRemoteImage("abcdef", makeURL("/v1/"), token)
	assertEqual(t, found, false, "Expected remote lookup to fail")
}

func TestGetRemoteImageJSON(t *testing.T) {
	r := spawnTestRegistrySession(t)
	json, size, err := r.GetRemoteImageJSON(imageID, makeURL("/v1/"), token)
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, size, 154, "Expected size 154")
	if len(json) <= 0 {
		t.Fatal("Expected non-empty json")
	}

	_, _, err = r.GetRemoteImageJSON("abcdef", makeURL("/v1/"), token)
	if err == nil {
		t.Fatal("Expected image not found error")
	}
}

func TestGetRemoteImageLayer(t *testing.T) {
	r := spawnTestRegistrySession(t)
	data, err := r.GetRemoteImageLayer(imageID, makeURL("/v1/"), token, 0)
	if err != nil {
		t.Fatal(err)
	}
	if data == nil {
		t.Fatal("Expected non-nil data result")
	}

	_, err = r.GetRemoteImageLayer("abcdef", makeURL("/v1/"), token, 0)
	if err == nil {
		t.Fatal("Expected image not found error")
	}
}

func TestGetRemoteTags(t *testing.T) {
	r := spawnTestRegistrySession(t)
	tags, err := r.GetRemoteTags([]string{makeURL("/v1/")}, REPO, token)
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, len(tags), 1, "Expected one tag")
	assertEqual(t, tags["latest"], imageID, "Expected tag latest to map to "+imageID)

	_, err = r.GetRemoteTags([]string{makeURL("/v1/")}, "foo42/baz", token)
	if err == nil {
		t.Fatal("Expected error when fetching tags for bogus repo")
	}
}

func TestGetRepositoryData(t *testing.T) {
	r := spawnTestRegistrySession(t)
	parsedURL, err := url.Parse(makeURL("/v1/"))
	if err != nil {
		t.Fatal(err)
	}
	host := "http://" + parsedURL.Host + "/v1/"
	data, err := r.GetRepositoryData("foo42/bar")
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, len(data.ImgList), 2, "Expected 2 images in ImgList")
	assertEqual(t, len(data.Endpoints), 2,
		fmt.Sprintf("Expected 2 endpoints in Endpoints, found %d instead", len(data.Endpoints)))
	assertEqual(t, data.Endpoints[0], host,
		fmt.Sprintf("Expected first endpoint to be %s but found %s instead", host, data.Endpoints[0]))
	assertEqual(t, data.Endpoints[1], "http://test.example.com/v1/",
		fmt.Sprintf("Expected first endpoint to be http://test.example.com/v1/ but found %s instead", data.Endpoints[1]))

}

func TestPushImageJSONRegistry(t *testing.T) {
	r := spawnTestRegistrySession(t)
	imgData := &ImgData{
		ID:       "77dbf71da1d00e3fbddc480176eac8994025630c6590d11cfc8fe1209c2a1d20",
		Checksum: "sha256:1ac330d56e05eef6d438586545ceff7550d3bdcb6b19961f12c5ba714ee1bb37",
	}

	err := r.PushImageJSONRegistry(imgData, []byte{0x42, 0xdf, 0x0}, makeURL("/v1/"), token)
	if err != nil {
		t.Fatal(err)
	}
}

func TestPushImageLayerRegistry(t *testing.T) {
	r := spawnTestRegistrySession(t)
	layer := strings.NewReader("")
	_, _, err := r.PushImageLayerRegistry(imageID, layer, makeURL("/v1/"), token, []byte{})
	if err != nil {
		t.Fatal(err)
	}
}

func TestResolveRepositoryName(t *testing.T) {
	_, _, err := ResolveRepositoryName("https://github.com/docker/docker")
	assertEqual(t, err, ErrInvalidRepositoryName, "Expected error invalid repo name")
	ep, repo, err := ResolveRepositoryName("fooo/bar")
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, ep, IndexServerAddress(), "Expected endpoint to be index server address")
	assertEqual(t, repo, "fooo/bar", "Expected resolved repo to be foo/bar")

	u := makeURL("")[7:]
	ep, repo, err = ResolveRepositoryName(u + "/private/moonbase")
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, ep, u, "Expected endpoint to be "+u)
	assertEqual(t, repo, "private/moonbase", "Expected endpoint to be private/moonbase")

	ep, repo, err = ResolveRepositoryName("ubuntu-12.04-base")
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, ep, IndexServerAddress(), "Expected endpoint to be "+IndexServerAddress())
	assertEqual(t, repo, "ubuntu-12.04-base", "Expected endpoint to be ubuntu-12.04-base")
}

func TestPushRegistryTag(t *testing.T) {
	r := spawnTestRegistrySession(t)
	err := r.PushRegistryTag("foo42/bar", imageID, "stable", makeURL("/v1/"), token)
	if err != nil {
		t.Fatal(err)
	}
}

func TestPushImageJSONIndex(t *testing.T) {
	r := spawnTestRegistrySession(t)
	imgData := []*ImgData{
		{
			ID:       "77dbf71da1d00e3fbddc480176eac8994025630c6590d11cfc8fe1209c2a1d20",
			Checksum: "sha256:1ac330d56e05eef6d438586545ceff7550d3bdcb6b19961f12c5ba714ee1bb37",
		},
		{
			ID:       "42d718c941f5c532ac049bf0b0ab53f0062f09a03afd4aa4a02c098e46032b9d",
			Checksum: "sha256:bea7bf2e4bacd479344b737328db47b18880d09096e6674165533aa994f5e9f2",
		},
	}
	repoData, err := r.PushImageJSONIndex("foo42/bar", imgData, false, nil)
	if err != nil {
		t.Fatal(err)
	}
	if repoData == nil {
		t.Fatal("Expected RepositoryData object")
	}
	repoData, err = r.PushImageJSONIndex("foo42/bar", imgData, true, []string{r.indexEndpoint.String()})
	if err != nil {
		t.Fatal(err)
	}
	if repoData == nil {
		t.Fatal("Expected RepositoryData object")
	}
}

func TestSearchRepositories(t *testing.T) {
	r := spawnTestRegistrySession(t)
	results, err := r.SearchRepositories("fakequery")
	if err != nil {
		t.Fatal(err)
	}
	if results == nil {
		t.Fatal("Expected non-nil SearchResults object")
	}
	assertEqual(t, results.NumResults, 1, "Expected 1 search results")
	assertEqual(t, results.Query, "fakequery", "Expected 'fakequery' as query")
	assertEqual(t, results.Results[0].StarCount, 42, "Expected 'fakeimage' a ot hae 42 stars")
}

func TestValidRepositoryName(t *testing.T) {
	if err := validateRepositoryName("docker/docker"); err != nil {
		t.Fatal(err)
	}
	// Support 64-byte non-hexadecimal names (hexadecimal names are forbidden)
	if err := validateRepositoryName("thisisthesongthatneverendsitgoesonandonandonthisisthesongthatnev"); err != nil {
		t.Fatal(err)
	}
	if err := validateRepositoryName("docker/Docker"); err == nil {
		t.Log("Repository name should be invalid")
		t.Fail()
	}
	if err := validateRepositoryName("docker///docker"); err == nil {
		t.Log("Repository name should be invalid")
		t.Fail()
	}
	if err := validateRepositoryName("1a3f5e7d9c1b3a5f7e9d1c3b5a7f9e1d3c5b7a9f1e3d5d7c9b1a3f5e7d9c1b3a"); err == nil {
		t.Log("Repository name should be invalid, 64-byte hexadecimal names forbidden")
		t.Fail()
	}
}

func TestTrustedLocation(t *testing.T) {
	for _, url := range []string{"http://example.com", "https://example.com:7777", "http://docker.io", "http://test.docker.com", "https://fakedocker.com"} {
		req, _ := http.NewRequest("GET", url, nil)
		if trustedLocation(req) == true {
			t.Fatalf("'%s' shouldn't be detected as a trusted location", url)
		}
	}

	for _, url := range []string{"https://docker.io", "https://test.docker.com:80"} {
		req, _ := http.NewRequest("GET", url, nil)
		if trustedLocation(req) == false {
			t.Fatalf("'%s' should be detected as a trusted location", url)
		}
	}
}

func TestAddRequiredHeadersToRedirectedRequests(t *testing.T) {
	for _, urls := range [][]string{
		{"http://docker.io", "https://docker.com"},
		{"https://foo.docker.io:7777", "http://bar.docker.com"},
		{"https://foo.docker.io", "https://example.com"},
	} {
		reqFrom, _ := http.NewRequest("GET", urls[0], nil)
		reqFrom.Header.Add("Content-Type", "application/json")
		reqFrom.Header.Add("Authorization", "super_secret")
		reqTo, _ := http.NewRequest("GET", urls[1], nil)

		AddRequiredHeadersToRedirectedRequests(reqTo, []*http.Request{reqFrom})

		if len(reqTo.Header) != 1 {
			t.Fatalf("Expected 1 headers, got %d", len(reqTo.Header))
		}

		if reqTo.Header.Get("Content-Type") != "application/json" {
			t.Fatal("'Content-Type' should be 'application/json'")
		}

		if reqTo.Header.Get("Authorization") != "" {
			t.Fatal("'Authorization' should be empty")
		}
	}

	for _, urls := range [][]string{
		{"https://docker.io", "https://docker.com"},
		{"https://foo.docker.io:7777", "https://bar.docker.com"},
	} {
		reqFrom, _ := http.NewRequest("GET", urls[0], nil)
		reqFrom.Header.Add("Content-Type", "application/json")
		reqFrom.Header.Add("Authorization", "super_secret")
		reqTo, _ := http.NewRequest("GET", urls[1], nil)

		AddRequiredHeadersToRedirectedRequests(reqTo, []*http.Request{reqFrom})

		if len(reqTo.Header) != 2 {
			t.Fatalf("Expected 2 headers, got %d", len(reqTo.Header))
		}

		if reqTo.Header.Get("Content-Type") != "application/json" {
			t.Fatal("'Content-Type' should be 'application/json'")
		}

		if reqTo.Header.Get("Authorization") != "super_secret" {
			t.Fatal("'Authorization' should be 'super_secret'")
		}
	}
}
<commit_msg>registry: add tests for IsSecure<commit_after>package registry

import (
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"testing"

	"github.com/docker/docker/utils"
)

var (
	token = []string{"fake-token"}
)

const (
	imageID = "42d718c941f5c532ac049bf0b0ab53f0062f09a03afd4aa4a02c098e46032b9d"
	REPO    = "foo42/bar"
)

func spawnTestRegistrySession(t *testing.T) *Session {
	authConfig := &AuthConfig{}
	endpoint, err := NewEndpoint(makeURL("/v1/"), false)
	if err != nil {
		t.Fatal(err)
	}
	r, err := NewSession(authConfig, utils.NewHTTPRequestFactory(), endpoint, true)
	if err != nil {
		t.Fatal(err)
	}
	return r
}

func TestPingRegistryEndpoint(t *testing.T) {
	ep, err := NewEndpoint(makeURL("/v1/"), false)
	if err != nil {
		t.Fatal(err)
	}
	regInfo, err := ep.Ping()
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, regInfo.Standalone, true, "Expected standalone to be true (default)")
}

func TestGetRemoteHistory(t *testing.T) {
	r := spawnTestRegistrySession(t)
	hist, err := r.GetRemoteHistory(imageID, makeURL("/v1/"), token)
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, len(hist), 2, "Expected 2 images in history")
	assertEqual(t, hist[0], imageID, "Expected "+imageID+"as first ancestry")
	assertEqual(t, hist[1], "77dbf71da1d00e3fbddc480176eac8994025630c6590d11cfc8fe1209c2a1d20",
		"Unexpected second ancestry")
}

func TestLookupRemoteImage(t *testing.T) {
	r := spawnTestRegistrySession(t)
	found := r.LookupRemoteImage(imageID, makeURL("/v1/"), token)
	assertEqual(t, found, true, "Expected remote lookup to succeed")
	found = r.LookupRemoteImage("abcdef", makeURL("/v1/"), token)
	assertEqual(t, found, false, "Expected remote lookup to fail")
}

func TestGetRemoteImageJSON(t *testing.T) {
	r := spawnTestRegistrySession(t)
	json, size, err := r.GetRemoteImageJSON(imageID, makeURL("/v1/"), token)
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, size, 154, "Expected size 154")
	if len(json) <= 0 {
		t.Fatal("Expected non-empty json")
	}

	_, _, err = r.GetRemoteImageJSON("abcdef", makeURL("/v1/"), token)
	if err == nil {
		t.Fatal("Expected image not found error")
	}
}

func TestGetRemoteImageLayer(t *testing.T) {
	r := spawnTestRegistrySession(t)
	data, err := r.GetRemoteImageLayer(imageID, makeURL("/v1/"), token, 0)
	if err != nil {
		t.Fatal(err)
	}
	if data == nil {
		t.Fatal("Expected non-nil data result")
	}

	_, err = r.GetRemoteImageLayer("abcdef", makeURL("/v1/"), token, 0)
	if err == nil {
		t.Fatal("Expected image not found error")
	}
}

func TestGetRemoteTags(t *testing.T) {
	r := spawnTestRegistrySession(t)
	tags, err := r.GetRemoteTags([]string{makeURL("/v1/")}, REPO, token)
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, len(tags), 1, "Expected one tag")
	assertEqual(t, tags["latest"], imageID, "Expected tag latest to map to "+imageID)

	_, err = r.GetRemoteTags([]string{makeURL("/v1/")}, "foo42/baz", token)
	if err == nil {
		t.Fatal("Expected error when fetching tags for bogus repo")
	}
}

func TestGetRepositoryData(t *testing.T) {
	r := spawnTestRegistrySession(t)
	parsedURL, err := url.Parse(makeURL("/v1/"))
	if err != nil {
		t.Fatal(err)
	}
	host := "http://" + parsedURL.Host + "/v1/"
	data, err := r.GetRepositoryData("foo42/bar")
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, len(data.ImgList), 2, "Expected 2 images in ImgList")
	assertEqual(t, len(data.Endpoints), 2,
		fmt.Sprintf("Expected 2 endpoints in Endpoints, found %d instead", len(data.Endpoints)))
	assertEqual(t, data.Endpoints[0], host,
		fmt.Sprintf("Expected first endpoint to be %s but found %s instead", host, data.Endpoints[0]))
	assertEqual(t, data.Endpoints[1], "http://test.example.com/v1/",
		fmt.Sprintf("Expected first endpoint to be http://test.example.com/v1/ but found %s instead", data.Endpoints[1]))

}

func TestPushImageJSONRegistry(t *testing.T) {
	r := spawnTestRegistrySession(t)
	imgData := &ImgData{
		ID:       "77dbf71da1d00e3fbddc480176eac8994025630c6590d11cfc8fe1209c2a1d20",
		Checksum: "sha256:1ac330d56e05eef6d438586545ceff7550d3bdcb6b19961f12c5ba714ee1bb37",
	}

	err := r.PushImageJSONRegistry(imgData, []byte{0x42, 0xdf, 0x0}, makeURL("/v1/"), token)
	if err != nil {
		t.Fatal(err)
	}
}

func TestPushImageLayerRegistry(t *testing.T) {
	r := spawnTestRegistrySession(t)
	layer := strings.NewReader("")
	_, _, err := r.PushImageLayerRegistry(imageID, layer, makeURL("/v1/"), token, []byte{})
	if err != nil {
		t.Fatal(err)
	}
}

func TestResolveRepositoryName(t *testing.T) {
	_, _, err := ResolveRepositoryName("https://github.com/docker/docker")
	assertEqual(t, err, ErrInvalidRepositoryName, "Expected error invalid repo name")
	ep, repo, err := ResolveRepositoryName("fooo/bar")
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, ep, IndexServerAddress(), "Expected endpoint to be index server address")
	assertEqual(t, repo, "fooo/bar", "Expected resolved repo to be foo/bar")

	u := makeURL("")[7:]
	ep, repo, err = ResolveRepositoryName(u + "/private/moonbase")
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, ep, u, "Expected endpoint to be "+u)
	assertEqual(t, repo, "private/moonbase", "Expected endpoint to be private/moonbase")

	ep, repo, err = ResolveRepositoryName("ubuntu-12.04-base")
	if err != nil {
		t.Fatal(err)
	}
	assertEqual(t, ep, IndexServerAddress(), "Expected endpoint to be "+IndexServerAddress())
	assertEqual(t, repo, "ubuntu-12.04-base", "Expected endpoint to be ubuntu-12.04-base")
}

func TestPushRegistryTag(t *testing.T) {
	r := spawnTestRegistrySession(t)
	err := r.PushRegistryTag("foo42/bar", imageID, "stable", makeURL("/v1/"), token)
	if err != nil {
		t.Fatal(err)
	}
}

func TestPushImageJSONIndex(t *testing.T) {
	r := spawnTestRegistrySession(t)
	imgData := []*ImgData{
		{
			ID:       "77dbf71da1d00e3fbddc480176eac8994025630c6590d11cfc8fe1209c2a1d20",
			Checksum: "sha256:1ac330d56e05eef6d438586545ceff7550d3bdcb6b19961f12c5ba714ee1bb37",
		},
		{
			ID:       "42d718c941f5c532ac049bf0b0ab53f0062f09a03afd4aa4a02c098e46032b9d",
			Checksum: "sha256:bea7bf2e4bacd479344b737328db47b18880d09096e6674165533aa994f5e9f2",
		},
	}
	repoData, err := r.PushImageJSONIndex("foo42/bar", imgData, false, nil)
	if err != nil {
		t.Fatal(err)
	}
	if repoData == nil {
		t.Fatal("Expected RepositoryData object")
	}
	repoData, err = r.PushImageJSONIndex("foo42/bar", imgData, true, []string{r.indexEndpoint.String()})
	if err != nil {
		t.Fatal(err)
	}
	if repoData == nil {
		t.Fatal("Expected RepositoryData object")
	}
}

func TestSearchRepositories(t *testing.T) {
	r := spawnTestRegistrySession(t)
	results, err := r.SearchRepositories("fakequery")
	if err != nil {
		t.Fatal(err)
	}
	if results == nil {
		t.Fatal("Expected non-nil SearchResults object")
	}
	assertEqual(t, results.NumResults, 1, "Expected 1 search results")
	assertEqual(t, results.Query, "fakequery", "Expected 'fakequery' as query")
	assertEqual(t, results.Results[0].StarCount, 42, "Expected 'fakeimage' a ot hae 42 stars")
}

func TestValidRepositoryName(t *testing.T) {
	if err := validateRepositoryName("docker/docker"); err != nil {
		t.Fatal(err)
	}
	// Support 64-byte non-hexadecimal names (hexadecimal names are forbidden)
	if err := validateRepositoryName("thisisthesongthatneverendsitgoesonandonandonthisisthesongthatnev"); err != nil {
		t.Fatal(err)
	}
	if err := validateRepositoryName("docker/Docker"); err == nil {
		t.Log("Repository name should be invalid")
		t.Fail()
	}
	if err := validateRepositoryName("docker///docker"); err == nil {
		t.Log("Repository name should be invalid")
		t.Fail()
	}
	if err := validateRepositoryName("1a3f5e7d9c1b3a5f7e9d1c3b5a7f9e1d3c5b7a9f1e3d5d7c9b1a3f5e7d9c1b3a"); err == nil {
		t.Log("Repository name should be invalid, 64-byte hexadecimal names forbidden")
		t.Fail()
	}
}

func TestTrustedLocation(t *testing.T) {
	for _, url := range []string{"http://example.com", "https://example.com:7777", "http://docker.io", "http://test.docker.com", "https://fakedocker.com"} {
		req, _ := http.NewRequest("GET", url, nil)
		if trustedLocation(req) == true {
			t.Fatalf("'%s' shouldn't be detected as a trusted location", url)
		}
	}

	for _, url := range []string{"https://docker.io", "https://test.docker.com:80"} {
		req, _ := http.NewRequest("GET", url, nil)
		if trustedLocation(req) == false {
			t.Fatalf("'%s' should be detected as a trusted location", url)
		}
	}
}

func TestAddRequiredHeadersToRedirectedRequests(t *testing.T) {
	for _, urls := range [][]string{
		{"http://docker.io", "https://docker.com"},
		{"https://foo.docker.io:7777", "http://bar.docker.com"},
		{"https://foo.docker.io", "https://example.com"},
	} {
		reqFrom, _ := http.NewRequest("GET", urls[0], nil)
		reqFrom.Header.Add("Content-Type", "application/json")
		reqFrom.Header.Add("Authorization", "super_secret")
		reqTo, _ := http.NewRequest("GET", urls[1], nil)

		AddRequiredHeadersToRedirectedRequests(reqTo, []*http.Request{reqFrom})

		if len(reqTo.Header) != 1 {
			t.Fatalf("Expected 1 headers, got %d", len(reqTo.Header))
		}

		if reqTo.Header.Get("Content-Type") != "application/json" {
			t.Fatal("'Content-Type' should be 'application/json'")
		}

		if reqTo.Header.Get("Authorization") != "" {
			t.Fatal("'Authorization' should be empty")
		}
	}

	for _, urls := range [][]string{
		{"https://docker.io", "https://docker.com"},
		{"https://foo.docker.io:7777", "https://bar.docker.com"},
	} {
		reqFrom, _ := http.NewRequest("GET", urls[0], nil)
		reqFrom.Header.Add("Content-Type", "application/json")
		reqFrom.Header.Add("Authorization", "super_secret")
		reqTo, _ := http.NewRequest("GET", urls[1], nil)

		AddRequiredHeadersToRedirectedRequests(reqTo, []*http.Request{reqFrom})

		if len(reqTo.Header) != 2 {
			t.Fatalf("Expected 2 headers, got %d", len(reqTo.Header))
		}

		if reqTo.Header.Get("Content-Type") != "application/json" {
			t.Fatal("'Content-Type' should be 'application/json'")
		}

		if reqTo.Header.Get("Authorization") != "super_secret" {
			t.Fatal("'Authorization' should be 'super_secret'")
		}
	}
}

func TestIsSecure(t *testing.T) {
	tests := []struct {
		addr               string
		insecureRegistries []string
		expected           bool
	}{
		{"example.com", []string{}, true},
		{"example.com", []string{"example.com"}, false},
		{"localhost", []string{"localhost:5000"}, true},
		{"localhost:5000", []string{"localhost:5000"}, false},
		{"localhost", []string{"example.com"}, true},
		{"127.0.0.1:5000", []string{"127.0.0.1:5000"}, false},
	}
	for _, tt := range tests {
		if sec := IsSecure(tt.addr, tt.insecureRegistries); sec != tt.expected {
			t.Errorf("IsSecure failed for %q %v, expected %v got %v", tt.addr, tt.insecureRegistries, tt.expected, sec)
		}
	}
}
<|endoftext|>


<commit_before>package rcmgr

import (
	"encoding/json"
	"fmt"
	"io"

	"github.com/libp2p/go-libp2p-core/peer"
	"github.com/libp2p/go-libp2p-core/protocol"

	"github.com/pbnjay/memory"
)

type limitConfig struct {
	// if true, then a dynamic limit is used
	Dynamic bool
	// either Memory is set for fixed memory limit
	Memory int64
	// or the following 3 fields for computed memory limits
	MinMemory      int64
	MaxMemory      int64
	MemoryFraction float64

	StreamsInbound  int
	StreamsOutbound int
	Streams         int

	ConnsInbound  int
	ConnsOutbound int
	Conns         int

	FD int
}

func (cfg *limitConfig) toLimit(base BaseLimit, mem MemoryLimit) (Limit, error) {
	if cfg == nil {
		m := mem.GetMemory(int64(memory.TotalMemory()))
		return &StaticLimit{
			Memory:    m,
			BaseLimit: base,
		}, nil
	}

	if cfg.Streams > 0 {
		base.Streams = cfg.Streams
	}
	if cfg.StreamsInbound > 0 {
		base.StreamsInbound = cfg.StreamsInbound
	}
	if cfg.StreamsOutbound > 0 {
		base.StreamsOutbound = cfg.StreamsOutbound
	}
	if cfg.Conns > 0 {
		base.Conns = cfg.Conns
	}
	if cfg.ConnsInbound > 0 {
		base.ConnsInbound = cfg.ConnsInbound
	}
	if cfg.ConnsOutbound > 0 {
		base.ConnsOutbound = cfg.ConnsOutbound
	}
	if cfg.FD > 0 {
		base.FD = cfg.FD
	}

	switch {
	case cfg.Memory > 0:
		return &StaticLimit{
			Memory:    cfg.Memory,
			BaseLimit: base,
		}, nil

	case cfg.Dynamic:
		if cfg.MemoryFraction < 0 {
			return nil, fmt.Errorf("negative memory fraction: %f", cfg.MemoryFraction)
		}
		if cfg.MemoryFraction > 0 {
			mem.MemoryFraction = cfg.MemoryFraction
		}
		if cfg.MinMemory > 0 {
			mem.MinMemory = cfg.MinMemory
		}
		if cfg.MaxMemory > 0 {
			mem.MaxMemory = cfg.MaxMemory
		}

		return &DynamicLimit{
			MemoryLimit: mem,
			BaseLimit:   base,
		}, nil

	default:
		if cfg.MemoryFraction < 0 {
			return nil, fmt.Errorf("negative memory fraction: %f", cfg.MemoryFraction)
		}
		if cfg.MemoryFraction > 0 {
			mem.MemoryFraction = cfg.MemoryFraction
		}
		if cfg.MinMemory > 0 {
			mem.MinMemory = cfg.MinMemory
		}
		if cfg.MaxMemory > 0 {
			mem.MaxMemory = cfg.MaxMemory
		}

		m := mem.GetMemory(int64(memory.TotalMemory()))
		return &StaticLimit{
			Memory:    m,
			BaseLimit: base,
		}, nil
	}
}

func (cfg *limitConfig) toLimitFixed(base BaseLimit, mem int64) (Limit, error) {
	if cfg == nil {
		return &StaticLimit{
			Memory:    mem,
			BaseLimit: base,
		}, nil
	}

	if cfg.Streams > 0 {
		base.Streams = cfg.Streams
	}
	if cfg.StreamsInbound > 0 {
		base.StreamsInbound = cfg.StreamsInbound
	}
	if cfg.StreamsOutbound > 0 {
		base.StreamsOutbound = cfg.StreamsOutbound
	}
	if cfg.Conns > 0 {
		base.Conns = cfg.Conns
	}
	if cfg.ConnsInbound > 0 {
		base.ConnsInbound = cfg.ConnsInbound
	}
	if cfg.ConnsOutbound > 0 {
		base.ConnsOutbound = cfg.ConnsOutbound
	}
	if cfg.FD > 0 {
		base.FD = cfg.FD
	}

	switch {
	case cfg.Memory > 0:
		return &StaticLimit{
			Memory:    cfg.Memory,
			BaseLimit: base,
		}, nil

	case cfg.Dynamic:
		return nil, fmt.Errorf("cannot specify dynamic limit for fixed memory limit")

	default:
		if cfg.MemoryFraction > 0 || cfg.MinMemory > 0 || cfg.MaxMemory > 0 {
			return nil, fmt.Errorf("cannot specify dynamic range for fixed memory limit")
		}
		return &StaticLimit{
			Memory:    mem,
			BaseLimit: base,
		}, nil
	}
}

type limiterConfig struct {
	System    *limitConfig
	Transient *limitConfig

	ServiceDefault     *limitConfig
	ServicePeerDefault *limitConfig
	Service            map[string]limitConfig
	ServicePeer        map[string]limitConfig

	ProtocolDefault     *limitConfig
	ProtocolPeerDefault *limitConfig
	Protocol            map[string]limitConfig
	ProtocolPeer        map[string]limitConfig

	PeerDefault *limitConfig
	Peer        map[string]limitConfig

	Conn   *limitConfig
	Stream *limitConfig
}

// NewDefaultLimiterFromJSON creates a new limiter by parsing a json configuration,
// using the default limits for fallback.
func NewDefaultLimiterFromJSON(in io.Reader) (*BasicLimiter, error) {
	return NewLimiterFromJSON(in, DefaultLimits)
}

// NewLimiterFromJSON creates a new limiter by parsing a json configuration.
func NewLimiterFromJSON(in io.Reader, defaults DefaultLimitConfig) (*BasicLimiter, error) {
	jin := json.NewDecoder(in)

	var cfg limiterConfig

	if err := jin.Decode(&cfg); err != nil {
		return nil, err
	}

	limiter := new(BasicLimiter)
	var err error

	limiter.SystemLimits, err = cfg.System.toLimit(defaults.SystemBaseLimit, defaults.SystemMemory)
	if err != nil {
		return nil, fmt.Errorf("invalid system limit: %w", err)
	}

	limiter.TransientLimits, err = cfg.Transient.toLimit(defaults.TransientBaseLimit, defaults.TransientMemory)
	if err != nil {
		return nil, fmt.Errorf("invalid transient limit: %w", err)
	}

	limiter.DefaultServiceLimits, err = cfg.ServiceDefault.toLimit(defaults.ServiceBaseLimit, defaults.ServiceMemory)
	if err != nil {
		return nil, fmt.Errorf("invlaid default service limit: %w", err)
	}

	limiter.DefaultServicePeerLimits, err = cfg.ServicePeerDefault.toLimit(defaults.ServicePeerBaseLimit, defaults.ServicePeerMemory)
	if err != nil {
		return nil, fmt.Errorf("invlaid default service peer limit: %w", err)
	}

	if len(cfg.Service) > 0 {
		limiter.ServiceLimits = make(map[string]Limit, len(cfg.Service))
		for svc, cfgLimit := range cfg.Service {
			limiter.ServiceLimits[svc], err = cfgLimit.toLimit(defaults.ServiceBaseLimit, defaults.ServiceMemory)
			if err != nil {
				return nil, fmt.Errorf("invalid service limit for %s: %w", svc, err)
			}
		}
	}

	if len(cfg.ServicePeer) > 0 {
		limiter.ServicePeerLimits = make(map[string]Limit, len(cfg.ServicePeer))
		for svc, cfgLimit := range cfg.ServicePeer {
			limiter.ServicePeerLimits[svc], err = cfgLimit.toLimit(defaults.ServicePeerBaseLimit, defaults.ServicePeerMemory)
			if err != nil {
				return nil, fmt.Errorf("invalid service peer limit for %s: %w", svc, err)
			}
		}
	}

	limiter.DefaultProtocolLimits, err = cfg.ProtocolDefault.toLimit(defaults.ProtocolBaseLimit, defaults.ProtocolMemory)
	if err != nil {
		return nil, fmt.Errorf("invlaid default protocol limit: %w", err)
	}

	limiter.DefaultProtocolPeerLimits, err = cfg.ProtocolPeerDefault.toLimit(defaults.ProtocolPeerBaseLimit, defaults.ProtocolPeerMemory)
	if err != nil {
		return nil, fmt.Errorf("invlaid default protocol peer limit: %w", err)
	}

	if len(cfg.Protocol) > 0 {
		limiter.ProtocolLimits = make(map[protocol.ID]Limit, len(cfg.Protocol))
		for p, cfgLimit := range cfg.Protocol {
			limiter.ProtocolLimits[protocol.ID(p)], err = cfgLimit.toLimit(defaults.ProtocolBaseLimit, defaults.ProtocolMemory)
			if err != nil {
				return nil, fmt.Errorf("invalid service limit for %s: %w", p, err)
			}
		}
	}

	if len(cfg.ProtocolPeer) > 0 {
		limiter.ProtocolPeerLimits = make(map[protocol.ID]Limit, len(cfg.ProtocolPeer))
		for p, cfgLimit := range cfg.ProtocolPeer {
			limiter.ProtocolPeerLimits[protocol.ID(p)], err = cfgLimit.toLimit(defaults.ProtocolPeerBaseLimit, defaults.ProtocolPeerMemory)
			if err != nil {
				return nil, fmt.Errorf("invalid service peer limit for %s: %w", p, err)
			}
		}
	}

	limiter.DefaultPeerLimits, err = cfg.PeerDefault.toLimit(defaults.PeerBaseLimit, defaults.PeerMemory)
	if err != nil {
		return nil, fmt.Errorf("invalid peer limit: %w", err)
	}

	if len(cfg.Peer) > 0 {
		limiter.PeerLimits = make(map[peer.ID]Limit, len(cfg.Peer))
		for p, cfgLimit := range cfg.Peer {
			pid, err := peer.Decode(p)
			if err != nil {
				return nil, fmt.Errorf("invalid peer ID %s: %w", p, err)
			}
			limiter.PeerLimits[pid], err = cfgLimit.toLimit(defaults.PeerBaseLimit, defaults.PeerMemory)
			if err != nil {
				return nil, fmt.Errorf("invalid peer limit for %s: %w", p, err)
			}
		}
	}

	limiter.ConnLimits, err = cfg.Conn.toLimitFixed(defaults.ConnBaseLimit, defaults.ConnMemory)
	if err != nil {
		return nil, fmt.Errorf("invalid conn limit: %w", err)
	}

	limiter.StreamLimits, err = cfg.Stream.toLimitFixed(defaults.StreamBaseLimit, defaults.StreamMemory)
	if err != nil {
		return nil, fmt.Errorf("invalid stream limit: %w", err)
	}

	return limiter, nil
}
<commit_msg>feat: export basic limiter config<commit_after>package rcmgr

import (
	"encoding/json"
	"fmt"
	"io"

	"github.com/libp2p/go-libp2p-core/peer"
	"github.com/libp2p/go-libp2p-core/protocol"

	"github.com/pbnjay/memory"
)

type BasicLimitConfig struct {
	// if true, then a dynamic limit is used
	Dynamic bool `json:",omitempty"`
	// either Memory is set for fixed memory limit
	Memory int64 `json:",omitempty"`
	// or the following 3 fields for computed memory limits
	MinMemory      int64   `json:",omitempty"`
	MaxMemory      int64   `json:",omitempty"`
	MemoryFraction float64 `json:",omitempty"`

	StreamsInbound  int
	StreamsOutbound int
	Streams         int

	ConnsInbound  int
	ConnsOutbound int
	Conns         int

	FD int
}

func (cfg *BasicLimitConfig) toLimit(base BaseLimit, mem MemoryLimit) (Limit, error) {
	if cfg == nil {
		m := mem.GetMemory(int64(memory.TotalMemory()))
		return &StaticLimit{
			Memory:    m,
			BaseLimit: base,
		}, nil
	}

	if cfg.Streams > 0 {
		base.Streams = cfg.Streams
	}
	if cfg.StreamsInbound > 0 {
		base.StreamsInbound = cfg.StreamsInbound
	}
	if cfg.StreamsOutbound > 0 {
		base.StreamsOutbound = cfg.StreamsOutbound
	}
	if cfg.Conns > 0 {
		base.Conns = cfg.Conns
	}
	if cfg.ConnsInbound > 0 {
		base.ConnsInbound = cfg.ConnsInbound
	}
	if cfg.ConnsOutbound > 0 {
		base.ConnsOutbound = cfg.ConnsOutbound
	}
	if cfg.FD > 0 {
		base.FD = cfg.FD
	}

	switch {
	case cfg.Memory > 0:
		return &StaticLimit{
			Memory:    cfg.Memory,
			BaseLimit: base,
		}, nil

	case cfg.Dynamic:
		if cfg.MemoryFraction < 0 {
			return nil, fmt.Errorf("negative memory fraction: %f", cfg.MemoryFraction)
		}
		if cfg.MemoryFraction > 0 {
			mem.MemoryFraction = cfg.MemoryFraction
		}
		if cfg.MinMemory > 0 {
			mem.MinMemory = cfg.MinMemory
		}
		if cfg.MaxMemory > 0 {
			mem.MaxMemory = cfg.MaxMemory
		}

		return &DynamicLimit{
			MemoryLimit: mem,
			BaseLimit:   base,
		}, nil

	default:
		if cfg.MemoryFraction < 0 {
			return nil, fmt.Errorf("negative memory fraction: %f", cfg.MemoryFraction)
		}
		if cfg.MemoryFraction > 0 {
			mem.MemoryFraction = cfg.MemoryFraction
		}
		if cfg.MinMemory > 0 {
			mem.MinMemory = cfg.MinMemory
		}
		if cfg.MaxMemory > 0 {
			mem.MaxMemory = cfg.MaxMemory
		}

		m := mem.GetMemory(int64(memory.TotalMemory()))
		return &StaticLimit{
			Memory:    m,
			BaseLimit: base,
		}, nil
	}
}

func (cfg *BasicLimitConfig) toLimitFixed(base BaseLimit, mem int64) (Limit, error) {
	if cfg == nil {
		return &StaticLimit{
			Memory:    mem,
			BaseLimit: base,
		}, nil
	}

	if cfg.Streams > 0 {
		base.Streams = cfg.Streams
	}
	if cfg.StreamsInbound > 0 {
		base.StreamsInbound = cfg.StreamsInbound
	}
	if cfg.StreamsOutbound > 0 {
		base.StreamsOutbound = cfg.StreamsOutbound
	}
	if cfg.Conns > 0 {
		base.Conns = cfg.Conns
	}
	if cfg.ConnsInbound > 0 {
		base.ConnsInbound = cfg.ConnsInbound
	}
	if cfg.ConnsOutbound > 0 {
		base.ConnsOutbound = cfg.ConnsOutbound
	}
	if cfg.FD > 0 {
		base.FD = cfg.FD
	}

	switch {
	case cfg.Memory > 0:
		return &StaticLimit{
			Memory:    cfg.Memory,
			BaseLimit: base,
		}, nil

	case cfg.Dynamic:
		return nil, fmt.Errorf("cannot specify dynamic limit for fixed memory limit")

	default:
		if cfg.MemoryFraction > 0 || cfg.MinMemory > 0 || cfg.MaxMemory > 0 {
			return nil, fmt.Errorf("cannot specify dynamic range for fixed memory limit")
		}
		return &StaticLimit{
			Memory:    mem,
			BaseLimit: base,
		}, nil
	}
}

type BasicLimiterConfig struct {
	System    *BasicLimitConfig `json:",omitempty"`
	Transient *BasicLimitConfig `json:",omitempty"`

	ServiceDefault     *BasicLimitConfig           `json:",omitempty"`
	ServicePeerDefault *BasicLimitConfig           `json:",omitempty"`
	Service            map[string]BasicLimitConfig `json:",omitempty"`
	ServicePeer        map[string]BasicLimitConfig `json:",omitempty"`

	ProtocolDefault     *BasicLimitConfig           `json:",omitempty"`
	ProtocolPeerDefault *BasicLimitConfig           `json:",omitempty"`
	Protocol            map[string]BasicLimitConfig `json:",omitempty"`
	ProtocolPeer        map[string]BasicLimitConfig `json:",omitempty"`

	PeerDefault *BasicLimitConfig           `json:",omitempty"`
	Peer        map[string]BasicLimitConfig `json:",omitempty"`

	Conn   *BasicLimitConfig `json:",omitempty"`
	Stream *BasicLimitConfig `json:",omitempty"`
}

// NewDefaultLimiterFromJSON creates a new limiter by parsing a json configuration,
// using the default limits for fallback.
func NewDefaultLimiterFromJSON(in io.Reader) (*BasicLimiter, error) {
	return NewLimiterFromJSON(in, DefaultLimits)
}

// NewLimiterFromJSON creates a new limiter by parsing a json configuration.
func NewLimiterFromJSON(in io.Reader, defaults DefaultLimitConfig) (*BasicLimiter, error) {
	jin := json.NewDecoder(in)

	var cfg BasicLimiterConfig

	if err := jin.Decode(&cfg); err != nil {
		return nil, err
	}

	return NewLimiter(cfg, defaults)
}

func NewLimiter(cfg BasicLimiterConfig, defaults DefaultLimitConfig) (*BasicLimiter, error) {
	limiter := new(BasicLimiter)
	var err error

	limiter.SystemLimits, err = cfg.System.toLimit(defaults.SystemBaseLimit, defaults.SystemMemory)
	if err != nil {
		return nil, fmt.Errorf("invalid system limit: %w", err)
	}

	limiter.TransientLimits, err = cfg.Transient.toLimit(defaults.TransientBaseLimit, defaults.TransientMemory)
	if err != nil {
		return nil, fmt.Errorf("invalid transient limit: %w", err)
	}

	limiter.DefaultServiceLimits, err = cfg.ServiceDefault.toLimit(defaults.ServiceBaseLimit, defaults.ServiceMemory)
	if err != nil {
		return nil, fmt.Errorf("invlaid default service limit: %w", err)
	}

	limiter.DefaultServicePeerLimits, err = cfg.ServicePeerDefault.toLimit(defaults.ServicePeerBaseLimit, defaults.ServicePeerMemory)
	if err != nil {
		return nil, fmt.Errorf("invlaid default service peer limit: %w", err)
	}

	if len(cfg.Service) > 0 {
		limiter.ServiceLimits = make(map[string]Limit, len(cfg.Service))
		for svc, cfgLimit := range cfg.Service {
			limiter.ServiceLimits[svc], err = cfgLimit.toLimit(defaults.ServiceBaseLimit, defaults.ServiceMemory)
			if err != nil {
				return nil, fmt.Errorf("invalid service limit for %s: %w", svc, err)
			}
		}
	}

	if len(cfg.ServicePeer) > 0 {
		limiter.ServicePeerLimits = make(map[string]Limit, len(cfg.ServicePeer))
		for svc, cfgLimit := range cfg.ServicePeer {
			limiter.ServicePeerLimits[svc], err = cfgLimit.toLimit(defaults.ServicePeerBaseLimit, defaults.ServicePeerMemory)
			if err != nil {
				return nil, fmt.Errorf("invalid service peer limit for %s: %w", svc, err)
			}
		}
	}

	limiter.DefaultProtocolLimits, err = cfg.ProtocolDefault.toLimit(defaults.ProtocolBaseLimit, defaults.ProtocolMemory)
	if err != nil {
		return nil, fmt.Errorf("invlaid default protocol limit: %w", err)
	}

	limiter.DefaultProtocolPeerLimits, err = cfg.ProtocolPeerDefault.toLimit(defaults.ProtocolPeerBaseLimit, defaults.ProtocolPeerMemory)
	if err != nil {
		return nil, fmt.Errorf("invlaid default protocol peer limit: %w", err)
	}

	if len(cfg.Protocol) > 0 {
		limiter.ProtocolLimits = make(map[protocol.ID]Limit, len(cfg.Protocol))
		for p, cfgLimit := range cfg.Protocol {
			limiter.ProtocolLimits[protocol.ID(p)], err = cfgLimit.toLimit(defaults.ProtocolBaseLimit, defaults.ProtocolMemory)
			if err != nil {
				return nil, fmt.Errorf("invalid service limit for %s: %w", p, err)
			}
		}
	}

	if len(cfg.ProtocolPeer) > 0 {
		limiter.ProtocolPeerLimits = make(map[protocol.ID]Limit, len(cfg.ProtocolPeer))
		for p, cfgLimit := range cfg.ProtocolPeer {
			limiter.ProtocolPeerLimits[protocol.ID(p)], err = cfgLimit.toLimit(defaults.ProtocolPeerBaseLimit, defaults.ProtocolPeerMemory)
			if err != nil {
				return nil, fmt.Errorf("invalid service peer limit for %s: %w", p, err)
			}
		}
	}

	limiter.DefaultPeerLimits, err = cfg.PeerDefault.toLimit(defaults.PeerBaseLimit, defaults.PeerMemory)
	if err != nil {
		return nil, fmt.Errorf("invalid peer limit: %w", err)
	}

	if len(cfg.Peer) > 0 {
		limiter.PeerLimits = make(map[peer.ID]Limit, len(cfg.Peer))
		for p, cfgLimit := range cfg.Peer {
			pid, err := peer.Decode(p)
			if err != nil {
				return nil, fmt.Errorf("invalid peer ID %s: %w", p, err)
			}
			limiter.PeerLimits[pid], err = cfgLimit.toLimit(defaults.PeerBaseLimit, defaults.PeerMemory)
			if err != nil {
				return nil, fmt.Errorf("invalid peer limit for %s: %w", p, err)
			}
		}
	}

	limiter.ConnLimits, err = cfg.Conn.toLimitFixed(defaults.ConnBaseLimit, defaults.ConnMemory)
	if err != nil {
		return nil, fmt.Errorf("invalid conn limit: %w", err)
	}

	limiter.StreamLimits, err = cfg.Stream.toLimitFixed(defaults.StreamBaseLimit, defaults.StreamMemory)
	if err != nil {
		return nil, fmt.Errorf("invalid stream limit: %w", err)
	}

	return limiter, nil
}
<|endoftext|>


<commit_before>/*
Copyright (c) 2014-2016 VMware, Inc. All Rights Reserved.

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

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

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

package vm

import (
	"context"
	"flag"
	"fmt"
	"strconv"

	"github.com/vmware/govmomi/govc/cli"
	"github.com/vmware/govmomi/govc/flags"
	"github.com/vmware/govmomi/object"
	"github.com/vmware/govmomi/vim25/types"
)

type hidKey struct {
	Code         int32
	ShiftPressed bool
}

var hidCharacterMap = map[string]hidKey{
	"a": {0x04, false},
	"b": {0x05, false},
	"c": {0x06, false},
	"d": {0x07, false},
	"e": {0x08, false},
	"f": {0x09, false},
	"g": {0x0a, false},
	"h": {0x0b, false},
	"i": {0x0c, false},
	"j": {0x0d, false},
	"k": {0x0e, false},
	"l": {0x0f, false},
	"m": {0x10, false},
	"n": {0x11, false},
	"o": {0x12, false},
	"p": {0x13, false},
	"q": {0x14, false},
	"r": {0x15, false},
	"s": {0x16, false},
	"t": {0x17, false},
	"u": {0x18, false},
	"v": {0x19, false},
	"w": {0x1a, false},
	"x": {0x1b, false},
	"y": {0x1c, false},
	"z": {0x1d, false},
	"1": {0x1e, false},
	"2": {0x1f, false},
	"3": {0x20, false},
	"4": {0x21, false},
	"5": {0x22, false},
	"6": {0x23, false},
	"7": {0x24, false},
	"8": {0x25, false},
	"9": {0x26, false},
	"0": {0x27, false},
	"A": {0x04, true},
	"B": {0x05, true},
	"C": {0x06, true},
	"D": {0x07, true},
	"E": {0x08, true},
	"F": {0x09, true},
	"G": {0x0a, true},
	"H": {0x0b, true},
	"I": {0x0c, true},
	"J": {0x0d, true},
	"K": {0x0e, true},
	"L": {0x0f, true},
	"M": {0x10, true},
	"N": {0x11, true},
	"O": {0x12, true},
	"P": {0x13, true},
	"Q": {0x14, true},
	"R": {0x15, true},
	"S": {0x16, true},
	"T": {0x17, true},
	"U": {0x18, true},
	"V": {0x19, true},
	"W": {0x1a, true},
	"X": {0x1b, true},
	"Y": {0x1c, true},
	"Z": {0x1d, true},
	"!": {0x1e, true},
	"@": {0x1f, true},
	"#": {0x20, true},
	"$": {0x21, true},
	"%": {0x22, true},
	"^": {0x23, true},
	"&": {0x24, true},
	"*": {0x25, true},
	"(": {0x26, true},
	")": {0x27, true},
	"-": {0x2d, false},
	"_": {0x2d, true},
	"=": {0x2e, false},
	"+": {0x2e, true},
	"[": {0x2f, false},
	"{": {0x2f, true},
	"]": {0x30, false},
	"}": {0x30, true},
	`\`: {0x31, false},
	"|": {0x31, true},
	";": {0x33, false},
	":": {0x33, true},
	"'": {0x34, false},
	`"`: {0x34, true},
	"`": {0x35, false},
	"~": {0x35, true},
	",": {0x36, false},
	"<": {0x36, true},
	".": {0x37, false},
	">": {0x37, true},
	"/": {0x38, false},
	"?": {0x38, true},
}

type keystrokes struct {
	*flags.VirtualMachineFlag

	UsbHidCodeValue int32
	UsbHidCode      string
	UsbHidString    string
	LeftControl     bool
	LeftShift       bool
	LeftAlt         bool
	LeftGui         bool
	RightControl    bool
	RightShift      bool
	RightAlt        bool
	RightGui        bool
}

func init() {
	cli.Register("vm.keystrokes", &keystrokes{})
}

func (cmd *keystrokes) Register(ctx context.Context, f *flag.FlagSet) {
	cmd.VirtualMachineFlag, ctx = flags.NewVirtualMachineFlag(ctx)
	cmd.VirtualMachineFlag.Register(ctx, f)

	f.StringVar(&cmd.UsbHidString, "s", "", "Raw String to Send")
	f.StringVar(&cmd.UsbHidCode, "c", "", "USB HID Code (hex)")
	f.Var(flags.NewInt32(&cmd.UsbHidCodeValue), "r", "Raw USB HID Code Value (int32)")
	f.BoolVar(&cmd.LeftControl, "lc", false, "Enable/Disable Left Control")
	f.BoolVar(&cmd.LeftShift, "ls", false, "Enable/Disable Left Shift")
	f.BoolVar(&cmd.LeftAlt, "la", false, "Enable/Disable Left Alt")
	f.BoolVar(&cmd.LeftGui, "lg", false, "Enable/Disable Left Gui")
	f.BoolVar(&cmd.RightControl, "rc", false, "Enable/Disable Right Control")
	f.BoolVar(&cmd.RightShift, "rs", false, "Enable/Disable Right Shift")
	f.BoolVar(&cmd.RightAlt, "ra", false, "Enable/Disable Right Alt")
	f.BoolVar(&cmd.RightGui, "rg", false, "Enable/Disable Right Gui")
}

func (cmd *keystrokes) Usage() string {
	return "VM"
}

func (cmd *keystrokes) Description() string {
	return `Send Keystrokes to VM.

Examples:
 Default Scenario
  govc vm.keystrokes -vm $vm -s "root" 	# writes 'root' to the console
  govc vm.keystrokes -vm $vm -c 0x15 	# writes an 'r' to the console
  govc vm.keystrokes -vm $vm -r 1376263 # writes an 'r' to the console
  govc vm.keystrokes -vm $vm -c 0x28 	# presses ENTER on the console
  govc vm.keystrokes -vm $vm -c 0x4c -la true -lc true 	# sends CTRL+ALT+DEL to console`
}

func (cmd *keystrokes) Process(ctx context.Context) error {
	if err := cmd.VirtualMachineFlag.Process(ctx); err != nil {
		return err
	}
	return nil
}

func (cmd *keystrokes) Run(ctx context.Context, f *flag.FlagSet) error {
	vm, err := cmd.VirtualMachine()
	if err != nil {
		return err
	}

	if vm == nil {
		return flag.ErrHelp
	}

	err = cmd.processUserInput(ctx, vm)
	if err != nil {
		return err
	}
	return nil
}

func (cmd *keystrokes) processUserInput(ctx context.Context, vm *object.VirtualMachine) error {
	if err := cmd.checkValidInputs(); err != nil {
		return err
	}

	codes, err := cmd.processUsbCode()

	if err != nil {
		return err
	}

	var keyEventArray []types.UsbScanCodeSpecKeyEvent
	for _, code := range codes {
		leftShiftSetting := false
		if code.ShiftPressed || cmd.LeftShift {
			leftShiftSetting = true
		}
		modifiers := types.UsbScanCodeSpecModifierType{
			LeftControl:  &cmd.LeftControl,
			LeftShift:    &leftShiftSetting,
			LeftAlt:      &cmd.LeftAlt,
			LeftGui:      &cmd.LeftGui,
			RightControl: &cmd.RightControl,
			RightShift:   &cmd.RightShift,
			RightAlt:     &cmd.RightAlt,
			RightGui:     &cmd.RightGui,
		}
		keyEvent := types.UsbScanCodeSpecKeyEvent{
			UsbHidCode: code.Code,
			Modifiers:  &modifiers,
		}
		keyEventArray = append(keyEventArray, keyEvent)
	}

	spec := types.UsbScanCodeSpec{
		KeyEvents: keyEventArray,
	}

	_, err = vm.PutUsbScanCodes(ctx, spec)

	return err
}

func (cmd *keystrokes) processUsbCode() ([]hidKey, error) {
	if cmd.rawCodeProvided() {
		return []hidKey{{cmd.UsbHidCodeValue, false}}, nil
	}
	if cmd.hexCodeProvided() {
		s, err := hexStringToHidCode(cmd.UsbHidCode)
		if err != nil {
			return nil, err
		}
		return []hidKey{{s, false}}, nil
	}

	if cmd.stringProvided() {
		var retKeyArray []hidKey
		for _, c := range cmd.UsbHidString {
			lookupValue, ok := hidCharacterMap[string(c)]
			if !ok {
				return nil, fmt.Errorf("invalid Character %s in String: %s", string(c), cmd.UsbHidString)
			}
			lookupValue.Code = intToHidCode(lookupValue.Code)
			retKeyArray = append(retKeyArray, lookupValue)
		}
		return retKeyArray, nil
	}
	return nil, nil
}

func hexStringToHidCode(hex string) (int32, error) {
	s, err := strconv.ParseInt(hex, 0, 32)
	if err != nil {
		return 0, err
	}
	return intToHidCode(int32(s)), nil
}

func intToHidCode(v int32) int32 {
	var s int32 = v << 16
	s |= 7
	return s
}

func (cmd *keystrokes) checkValidInputs() error {
	// poor man's boolean XOR -> A xor B xor C = A'BC' + AB'C' + A'B'C + ABC
	if (!cmd.rawCodeProvided() && cmd.hexCodeProvided() && !cmd.stringProvided()) || // A'BC'
		(cmd.rawCodeProvided() && !cmd.hexCodeProvided() && !cmd.stringProvided()) || // AB'C'
		(!cmd.rawCodeProvided() && !cmd.hexCodeProvided() && cmd.stringProvided()) || // A'B'C
		(cmd.rawCodeProvided() && cmd.hexCodeProvided() && cmd.stringProvided()) { // ABC
		return nil
	}
	return fmt.Errorf("specify only 1 argument")
}

func (cmd keystrokes) rawCodeProvided() bool {
	return cmd.UsbHidCodeValue != 0
}

func (cmd keystrokes) hexCodeProvided() bool {
	return cmd.UsbHidCode != ""
}

func (cmd keystrokes) stringProvided() bool {
	return cmd.UsbHidString != ""
}
<commit_msg>vm.keystrokes -s (Allow spaces)<commit_after>/*
Copyright (c) 2014-2016 VMware, Inc. All Rights Reserved.

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

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

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

package vm

import (
	"context"
	"flag"
	"fmt"
	"strconv"

	"github.com/vmware/govmomi/govc/cli"
	"github.com/vmware/govmomi/govc/flags"
	"github.com/vmware/govmomi/object"
	"github.com/vmware/govmomi/vim25/types"
)

type hidKey struct {
	Code         int32
	ShiftPressed bool
}

var hidCharacterMap = map[string]hidKey{
	"a": {0x04, false},
	"b": {0x05, false},
	"c": {0x06, false},
	"d": {0x07, false},
	"e": {0x08, false},
	"f": {0x09, false},
	"g": {0x0a, false},
	"h": {0x0b, false},
	"i": {0x0c, false},
	"j": {0x0d, false},
	"k": {0x0e, false},
	"l": {0x0f, false},
	"m": {0x10, false},
	"n": {0x11, false},
	"o": {0x12, false},
	"p": {0x13, false},
	"q": {0x14, false},
	"r": {0x15, false},
	"s": {0x16, false},
	"t": {0x17, false},
	"u": {0x18, false},
	"v": {0x19, false},
	"w": {0x1a, false},
	"x": {0x1b, false},
	"y": {0x1c, false},
	"z": {0x1d, false},
	"1": {0x1e, false},
	"2": {0x1f, false},
	"3": {0x20, false},
	"4": {0x21, false},
	"5": {0x22, false},
	"6": {0x23, false},
	"7": {0x24, false},
	"8": {0x25, false},
	"9": {0x26, false},
	"0": {0x27, false},
	"A": {0x04, true},
	"B": {0x05, true},
	"C": {0x06, true},
	"D": {0x07, true},
	"E": {0x08, true},
	"F": {0x09, true},
	"G": {0x0a, true},
	"H": {0x0b, true},
	"I": {0x0c, true},
	"J": {0x0d, true},
	"K": {0x0e, true},
	"L": {0x0f, true},
	"M": {0x10, true},
	"N": {0x11, true},
	"O": {0x12, true},
	"P": {0x13, true},
	"Q": {0x14, true},
	"R": {0x15, true},
	"S": {0x16, true},
	"T": {0x17, true},
	"U": {0x18, true},
	"V": {0x19, true},
	"W": {0x1a, true},
	"X": {0x1b, true},
	"Y": {0x1c, true},
	"Z": {0x1d, true},
	"!": {0x1e, true},
	"@": {0x1f, true},
	"#": {0x20, true},
	"$": {0x21, true},
	"%": {0x22, true},
	"^": {0x23, true},
	"&": {0x24, true},
	"*": {0x25, true},
	"(": {0x26, true},
	")": {0x27, true},
	" ": {0x2c, false},
	"-": {0x2d, false},
	"_": {0x2d, true},
	"=": {0x2e, false},
	"+": {0x2e, true},
	"[": {0x2f, false},
	"{": {0x2f, true},
	"]": {0x30, false},
	"}": {0x30, true},
	`\`: {0x31, false},
	"|": {0x31, true},
	";": {0x33, false},
	":": {0x33, true},
	"'": {0x34, false},
	`"`: {0x34, true},
	"`": {0x35, false},
	"~": {0x35, true},
	",": {0x36, false},
	"<": {0x36, true},
	".": {0x37, false},
	">": {0x37, true},
	"/": {0x38, false},
	"?": {0x38, true},
}

type keystrokes struct {
	*flags.VirtualMachineFlag

	UsbHidCodeValue int32
	UsbHidCode      string
	UsbHidString    string
	LeftControl     bool
	LeftShift       bool
	LeftAlt         bool
	LeftGui         bool
	RightControl    bool
	RightShift      bool
	RightAlt        bool
	RightGui        bool
}

func init() {
	cli.Register("vm.keystrokes", &keystrokes{})
}

func (cmd *keystrokes) Register(ctx context.Context, f *flag.FlagSet) {
	cmd.VirtualMachineFlag, ctx = flags.NewVirtualMachineFlag(ctx)
	cmd.VirtualMachineFlag.Register(ctx, f)

	f.StringVar(&cmd.UsbHidString, "s", "", "Raw String to Send")
	f.StringVar(&cmd.UsbHidCode, "c", "", "USB HID Code (hex)")
	f.Var(flags.NewInt32(&cmd.UsbHidCodeValue), "r", "Raw USB HID Code Value (int32)")
	f.BoolVar(&cmd.LeftControl, "lc", false, "Enable/Disable Left Control")
	f.BoolVar(&cmd.LeftShift, "ls", false, "Enable/Disable Left Shift")
	f.BoolVar(&cmd.LeftAlt, "la", false, "Enable/Disable Left Alt")
	f.BoolVar(&cmd.LeftGui, "lg", false, "Enable/Disable Left Gui")
	f.BoolVar(&cmd.RightControl, "rc", false, "Enable/Disable Right Control")
	f.BoolVar(&cmd.RightShift, "rs", false, "Enable/Disable Right Shift")
	f.BoolVar(&cmd.RightAlt, "ra", false, "Enable/Disable Right Alt")
	f.BoolVar(&cmd.RightGui, "rg", false, "Enable/Disable Right Gui")
}

func (cmd *keystrokes) Usage() string {
	return "VM"
}

func (cmd *keystrokes) Description() string {
	return `Send Keystrokes to VM.

Examples:
 Default Scenario
  govc vm.keystrokes -vm $vm -s "root" 	# writes 'root' to the console
  govc vm.keystrokes -vm $vm -c 0x15 	# writes an 'r' to the console
  govc vm.keystrokes -vm $vm -r 1376263 # writes an 'r' to the console
  govc vm.keystrokes -vm $vm -c 0x28 	# presses ENTER on the console
  govc vm.keystrokes -vm $vm -c 0x4c -la true -lc true 	# sends CTRL+ALT+DEL to console`
}

func (cmd *keystrokes) Process(ctx context.Context) error {
	if err := cmd.VirtualMachineFlag.Process(ctx); err != nil {
		return err
	}
	return nil
}

func (cmd *keystrokes) Run(ctx context.Context, f *flag.FlagSet) error {
	vm, err := cmd.VirtualMachine()
	if err != nil {
		return err
	}

	if vm == nil {
		return flag.ErrHelp
	}

	err = cmd.processUserInput(ctx, vm)
	if err != nil {
		return err
	}
	return nil
}

func (cmd *keystrokes) processUserInput(ctx context.Context, vm *object.VirtualMachine) error {
	if err := cmd.checkValidInputs(); err != nil {
		return err
	}

	codes, err := cmd.processUsbCode()

	if err != nil {
		return err
	}

	var keyEventArray []types.UsbScanCodeSpecKeyEvent
	for _, code := range codes {
		leftShiftSetting := false
		if code.ShiftPressed || cmd.LeftShift {
			leftShiftSetting = true
		}
		modifiers := types.UsbScanCodeSpecModifierType{
			LeftControl:  &cmd.LeftControl,
			LeftShift:    &leftShiftSetting,
			LeftAlt:      &cmd.LeftAlt,
			LeftGui:      &cmd.LeftGui,
			RightControl: &cmd.RightControl,
			RightShift:   &cmd.RightShift,
			RightAlt:     &cmd.RightAlt,
			RightGui:     &cmd.RightGui,
		}
		keyEvent := types.UsbScanCodeSpecKeyEvent{
			UsbHidCode: code.Code,
			Modifiers:  &modifiers,
		}
		keyEventArray = append(keyEventArray, keyEvent)
	}

	spec := types.UsbScanCodeSpec{
		KeyEvents: keyEventArray,
	}

	_, err = vm.PutUsbScanCodes(ctx, spec)

	return err
}

func (cmd *keystrokes) processUsbCode() ([]hidKey, error) {
	if cmd.rawCodeProvided() {
		return []hidKey{{cmd.UsbHidCodeValue, false}}, nil
	}
	if cmd.hexCodeProvided() {
		s, err := hexStringToHidCode(cmd.UsbHidCode)
		if err != nil {
			return nil, err
		}
		return []hidKey{{s, false}}, nil
	}

	if cmd.stringProvided() {
		var retKeyArray []hidKey
		for _, c := range cmd.UsbHidString {
			lookupValue, ok := hidCharacterMap[string(c)]
			if !ok {
				return nil, fmt.Errorf("invalid Character %s in String: %s", string(c), cmd.UsbHidString)
			}
			lookupValue.Code = intToHidCode(lookupValue.Code)
			retKeyArray = append(retKeyArray, lookupValue)
		}
		return retKeyArray, nil
	}
	return nil, nil
}

func hexStringToHidCode(hex string) (int32, error) {
	s, err := strconv.ParseInt(hex, 0, 32)
	if err != nil {
		return 0, err
	}
	return intToHidCode(int32(s)), nil
}

func intToHidCode(v int32) int32 {
	var s int32 = v << 16
	s |= 7
	return s
}

func (cmd *keystrokes) checkValidInputs() error {
	// poor man's boolean XOR -> A xor B xor C = A'BC' + AB'C' + A'B'C + ABC
	if (!cmd.rawCodeProvided() && cmd.hexCodeProvided() && !cmd.stringProvided()) || // A'BC'
		(cmd.rawCodeProvided() && !cmd.hexCodeProvided() && !cmd.stringProvided()) || // AB'C'
		(!cmd.rawCodeProvided() && !cmd.hexCodeProvided() && cmd.stringProvided()) || // A'B'C
		(cmd.rawCodeProvided() && cmd.hexCodeProvided() && cmd.stringProvided()) { // ABC
		return nil
	}
	return fmt.Errorf("specify only 1 argument")
}

func (cmd keystrokes) rawCodeProvided() bool {
	return cmd.UsbHidCodeValue != 0
}

func (cmd keystrokes) hexCodeProvided() bool {
	return cmd.UsbHidCode != ""
}

func (cmd keystrokes) stringProvided() bool {
	return cmd.UsbHidString != ""
}
<|endoftext|>


<commit_before>// Package fetcher implements a simple bridge for HTTP request and response
package fetcher

import (
	"appengine"
	"appengine/urlfetch"
	"bufio"
	"net/http"
)

func init() {
	http.HandleFunc("/", hello) // this will handle all except /http and /connect
	http.HandleFunc("/http", HandleHTTP)
	http.HandleFunc("/connect", HandleConnect)
}

// 1. read the real client request R1 from the body of request R2
// 2. round trip the request R1 by urlfetch
// 3. write response P1 of request R1 as the body of response P2
// 4. send back response P2
func HandleHTTP(w http.ResponseWriter, r *http.Request) {
	// create context for gae request
	ctx := appengine.NewContext(r)
	cli := urlfetch.Client(ctx)
	s := NewSession(ctx, w)

	// read client requst
	creq, err := http.ReadRequest(bufio.NewReader(r.Body))
	if err != nil {
		s.HTTPError("ReadRequest: %s", err.Error())
		return
	}

	// record the real client request
    s.Info("%s %s://%s %s\n", creq.Method, creq.URL.Scheme, creq.URL.Host, creq.Proto)

	// round trip the client request
	// in fact RoundTrip supports both http and https
	resp, err := cli.Transport.RoundTrip(creq)
	if err != nil {
		s.HTTPError("RoundTrip: %s", err.Error())
		return
	}
	defer resp.Body.Close()

	// remote server response
	s.Info("RESPONSE %s %s\n", creq.URL.Host, resp.Status)

	// write response and send to client
	if err := resp.Write(w); err != nil {
		s.HTTPError("Write: %s", err.Error())
		return
	}
}

// unimplemented
func HandleConnect(w http.ResponseWriter, r *http.Request) {
	// TODO: socket?
}

func hello(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Hello, world!"))
}
<commit_msg>format comments<commit_after>// Package fetcher implements a simple bridge for HTTP request and response
package fetcher

import (
	"appengine"
	"appengine/urlfetch"
	"bufio"
	"net/http"
)

func init() {
	http.HandleFunc("/", hello) // this will handle all except /http and /connect
	http.HandleFunc("/http", HandleHTTP)
	http.HandleFunc("/connect", HandleConnect)
}

// Flow:
//  1. read the real client request R1 from the body of request R2
//  2. round trip the request R1 by urlfetch
//  3. write response P1 of request R1 as the body of response P2
//  4. send back response P2
func HandleHTTP(w http.ResponseWriter, r *http.Request) {
	// create context for gae request
	ctx := appengine.NewContext(r)
	cli := urlfetch.Client(ctx)
	s := NewSession(ctx, w)

	// read client requst
	creq, err := http.ReadRequest(bufio.NewReader(r.Body))
	if err != nil {
		s.HTTPError("ReadRequest: %s", err.Error())
		return
	}

	// record the real client request
	s.Info("%s %s://%s %s\n", creq.Method, creq.URL.Scheme, creq.URL.Host, creq.Proto)

	// round trip the client request
	// in fact RoundTrip supports both http and https
	resp, err := cli.Transport.RoundTrip(creq)
	if err != nil {
		s.HTTPError("RoundTrip: %s", err.Error())
		return
	}
	defer resp.Body.Close()

	// remote server response
	s.Info("RESPONSE %s %s\n", creq.URL.Host, resp.Status)

	// write response and send to client
	if err := resp.Write(w); err != nil {
		s.HTTPError("Write: %s", err.Error())
		return
	}
}

// unimplemented
func HandleConnect(w http.ResponseWriter, r *http.Request) {
	// TODO: socket?
}

func hello(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Hello, world!"))
}
<|endoftext|>