<commit_before>package gateway

import (
	"errors"
	"sync"
	"time"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/modules"
)

// rpcID is an 8-byte signature that is added to all RPCs to tell the gatway
// what to do with the RPC.
type rpcID [8]byte

// String returns a string representation of an rpcID. Empty elements of rpcID
// will be encoded as spaces.
func (id rpcID) String() string {
	for i := range id {
		if id[i] == 0 {
			id[i] = ' '
		}
	}
	return string(id[:])
}

// handlerName truncates a string to 8 bytes. If len(name) < 8, the remaining
// bytes are 0. A handlerName is specified at the beginning of each network
// call, indicating which function should handle the connection.
func handlerName(name string) (id rpcID) {
	copy(id[:], name)
	return
}

// managedRPC calls an RPC on the given address. managedRPC cannot be called on
// an address that the Gateway is not connected to.
func (g *Gateway) managedRPC(addr modules.NetAddress, name string, fn modules.RPCFunc) error {
	g.mu.RLock()
	peer, ok := g.peers[addr]
	g.mu.RUnlock()
	if !ok {
		return errors.New("can't call RPC on unconnected peer " + string(addr))
	}

	conn, err := peer.open()
	if err != nil {
		// peer probably disconnected without sending a shutdown signal;
		// disconnect from them
		g.log.Debugf("Could not initiate RPC with %v; disconnecting", addr)
		g.mu.Lock()
		delete(g.peers, addr)
		g.mu.Unlock()
		if derr := peer.sess.Close(); err != nil {
			g.log.Debugf("Could not cleanly disconnect from %v: %v", addr, derr)
		}
		return err
	}
	defer conn.Close()

	// write header
	if err := encoding.WriteObject(conn, handlerName(name)); err != nil {
		return err
	}
	// call fn
	return fn(conn)
}

// RPC calls an RPC on the given address. RPC cannot be called on an address
// that the Gateway is not connected to.
func (g *Gateway) RPC(addr modules.NetAddress, name string, fn modules.RPCFunc) error {
	if err := g.threads.Add(); err != nil {
		return err
	}
	defer g.threads.Done()
	return g.managedRPC(addr, name, fn)
}

// RegisterRPC registers an RPCFunc as a handler for a given identifier. To
// call an RPC, use gateway.RPC, supplying the same identifier given to
// RegisterRPC. Identifiers should always use PascalCase. The first 8
// characters of an identifier should be unique, as the identifier used
// internally is truncated to 8 bytes.
func (g *Gateway) RegisterRPC(name string, fn modules.RPCFunc) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.handlers[handlerName(name)]; ok {
		build.Critical("RPC already registered: " + name)
	}
	g.handlers[handlerName(name)] = fn
}

// UnregisterRPC unregisters an RPC and removes the corresponding RPCFunc from
// g.handlers. Future calls to the RPC by peers will fail.
func (g *Gateway) UnregisterRPC(name string) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.handlers[handlerName(name)]; !ok {
		build.Critical("RPC not registered: " + name)
	}
	delete(g.handlers, handlerName(name))
}

// RegisterConnectCall registers a name and RPCFunc to be called on a peer
// upon connecting.
func (g *Gateway) RegisterConnectCall(name string, fn modules.RPCFunc) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.initRPCs[name]; ok {
		build.Critical("ConnectCall already registered: " + name)
	}
	g.initRPCs[name] = fn
}

// UnregisterConnectCall unregisters an on-connect call and removes the
// corresponding RPCFunc from g.initRPCs. Future connections to peers will not
// trigger the RPC to be called on them.
func (g *Gateway) UnregisterConnectCall(name string) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.initRPCs[name]; !ok {
		build.Critical("ConnectCall not registered: " + name)
	}
	delete(g.initRPCs, name)
}

// threadedListenPeer listens for new streams on a peer connection and serves them via
// threadedHandleConn.
func (g *Gateway) threadedListenPeer(p *peer) {
	// threadedListenPeer registers to the peerTG instead of the primary thread
	// group because peer connections can be lifetime in length, but can also
	// be short-lived. The fact that they can be lifetime means that they can't
	// call threads.Add as they will block calls to threads.Flush. The fact
	// that they can be short-lived means that threads.OnStop is not a good
	// tool for closing out the threads. Instead, they register to peerTG,
	// which is cleanly closed upon gateway shutdown but will not block any
	// calls to threads.Flush()
	if g.peerTG.Add() != nil {
		return
	}
	defer g.peerTG.Done()

	// Spin up a goroutine to listen for a shutdown signal from both the peer
	// and from the gateway. In the event of either, close the muxado session.
	connClosedChan := make(chan struct{})
	peerCloseChan := make(chan struct{})
	go func() {
		// Signal that the muxado session has been successfully closed, and
		// that this goroutine has terminated.
		defer close(connClosedChan)

		// Listen for a stop signal.
		select {
		case <-g.threads.StopChan():
		case <-peerCloseChan:
		}

		// Can't call Disconnect because it could return sync.ErrStopped.
		g.mu.Lock()
		delete(g.peers, p.NetAddress)
		g.mu.Unlock()
		if err := p.sess.Close(); err != nil {
			g.log.Debugf("WARN: error disconnecting from peer %q: %v", p.NetAddress, err)
		}
	}()

	for {
		conn, err := p.accept()
		if err != nil {
			g.log.Debugf("Peer connection with %v closed: %v\n", p.NetAddress, err)
			break
		}
		// Set the default deadline on the conn.
		err = conn.SetDeadline(time.Now().Add(rpcStdDeadline))
		if err != nil {
			g.log.Printf("Peer connection (%v) deadline could not be set: %v\n", p.NetAddress, err)
			continue
		}

		// The handler is responsible for closing the connection, though a
		// default deadline has been set.
		go g.threadedHandleConn(conn)
		if !g.managedSleep(peerRPCDelay) {
			break
		}
	}
	// Signal that the goroutine can shutdown.
	close(peerCloseChan)
	// Wait for confirmation that the goroutine has shut down before returning
	// and releasing the threadgroup registration.
	<-connClosedChan
}

// threadedHandleConn reads header data from a connection, then routes it to the
// appropriate handler for further processing.
func (g *Gateway) threadedHandleConn(conn modules.PeerConn) {
	defer conn.Close()
	if g.threads.Add() != nil {
		return
	}
	defer g.threads.Done()

	var id rpcID
	err := conn.SetDeadline(time.Now().Add(rpcStdDeadline))
	if err != nil {
		return
	}
	if err := encoding.ReadObject(conn, &id, 8); err != nil {
		return
	}
	// call registered handler for this ID
	g.mu.RLock()
	fn, ok := g.handlers[id]
	g.mu.RUnlock()
	if !ok {
		g.log.Debugf("WARN: incoming conn %v requested unknown RPC \"%v\"", conn.RPCAddr(), id)
		return
	}
	g.log.Debugf("INFO: incoming conn %v requested RPC \"%v\"", conn.RPCAddr(), id)

	// call fn
	err = fn(conn)
	// don't log benign errors
	if err == modules.ErrDuplicateTransactionSet || err == modules.ErrBlockKnown {
		err = nil
	}
	if err != nil {
		g.log.Debugf("WARN: incoming RPC \"%v\" from conn %v failed: %v", id, conn.RPCAddr(), err)
	}
}

// Broadcast calls an RPC on all of the specified peers. The calls are run in
// parallel. Broadcasts are restricted to "one-way" RPCs, which simply write an
// object and disconnect. This is why Broadcast takes an interface{} instead of
// an RPCFunc.
func (g *Gateway) Broadcast(name string, obj interface{}, peers []modules.Peer) {
	if g.threads.Add() != nil {
		return
	}
	defer g.threads.Done()

	g.log.Debugf("INFO: broadcasting RPC %q to %v peers", name, len(peers))

	// only encode obj once, instead of using WriteObject
	enc := encoding.Marshal(obj)
	fn := func(conn modules.PeerConn) error {
		return encoding.WritePrefix(conn, enc)
	}

	var wg sync.WaitGroup
	for _, p := range peers {
		wg.Add(1)
		go func(addr modules.NetAddress) {
			defer wg.Done()
			err := g.managedRPC(addr, name, fn)
			if err != nil {
				g.log.Debugf("WARN: broadcasting RPC %q to peer %q failed (attempting again in 10 seconds): %v", name, addr, err)
				// try one more time before giving up
				select {
				case <-time.After(10 * time.Second):
				case <-g.threads.StopChan():
					return
				}
				err := g.managedRPC(addr, name, fn)
				if err != nil {
					g.log.Debugf("WARN: broadcasting RPC %q to peer %q failed twice: %v", name, addr, err)
				}
			}
		}(p.NetAddress)
	}
	wg.Wait()
}
<commit_msg>add deadline when writing RPC header<commit_after>package gateway

import (
	"errors"
	"sync"
	"time"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/modules"
)

// rpcID is an 8-byte signature that is added to all RPCs to tell the gatway
// what to do with the RPC.
type rpcID [8]byte

// String returns a string representation of an rpcID. Empty elements of rpcID
// will be encoded as spaces.
func (id rpcID) String() string {
	for i := range id {
		if id[i] == 0 {
			id[i] = ' '
		}
	}
	return string(id[:])
}

// handlerName truncates a string to 8 bytes. If len(name) < 8, the remaining
// bytes are 0. A handlerName is specified at the beginning of each network
// call, indicating which function should handle the connection.
func handlerName(name string) (id rpcID) {
	copy(id[:], name)
	return
}

// managedRPC calls an RPC on the given address. managedRPC cannot be called on
// an address that the Gateway is not connected to.
func (g *Gateway) managedRPC(addr modules.NetAddress, name string, fn modules.RPCFunc) error {
	g.mu.RLock()
	peer, ok := g.peers[addr]
	g.mu.RUnlock()
	if !ok {
		return errors.New("can't call RPC on unconnected peer " + string(addr))
	}

	conn, err := peer.open()
	if err != nil {
		// peer probably disconnected without sending a shutdown signal;
		// disconnect from them
		g.log.Debugf("Could not initiate RPC with %v; disconnecting", addr)
		g.mu.Lock()
		delete(g.peers, addr)
		g.mu.Unlock()
		if derr := peer.sess.Close(); err != nil {
			g.log.Debugf("Could not cleanly disconnect from %v: %v", addr, derr)
		}
		return err
	}
	defer conn.Close()

	// write header
	conn.SetDeadline(time.Now().Add(10 * time.Second))
	if err := encoding.WriteObject(conn, handlerName(name)); err != nil {
		return err
	}
	conn.SetDeadline(time.Time{})
	// call fn
	return fn(conn)
}

// RPC calls an RPC on the given address. RPC cannot be called on an address
// that the Gateway is not connected to.
func (g *Gateway) RPC(addr modules.NetAddress, name string, fn modules.RPCFunc) error {
	if err := g.threads.Add(); err != nil {
		return err
	}
	defer g.threads.Done()
	return g.managedRPC(addr, name, fn)
}

// RegisterRPC registers an RPCFunc as a handler for a given identifier. To
// call an RPC, use gateway.RPC, supplying the same identifier given to
// RegisterRPC. Identifiers should always use PascalCase. The first 8
// characters of an identifier should be unique, as the identifier used
// internally is truncated to 8 bytes.
func (g *Gateway) RegisterRPC(name string, fn modules.RPCFunc) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.handlers[handlerName(name)]; ok {
		build.Critical("RPC already registered: " + name)
	}
	g.handlers[handlerName(name)] = fn
}

// UnregisterRPC unregisters an RPC and removes the corresponding RPCFunc from
// g.handlers. Future calls to the RPC by peers will fail.
func (g *Gateway) UnregisterRPC(name string) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.handlers[handlerName(name)]; !ok {
		build.Critical("RPC not registered: " + name)
	}
	delete(g.handlers, handlerName(name))
}

// RegisterConnectCall registers a name and RPCFunc to be called on a peer
// upon connecting.
func (g *Gateway) RegisterConnectCall(name string, fn modules.RPCFunc) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.initRPCs[name]; ok {
		build.Critical("ConnectCall already registered: " + name)
	}
	g.initRPCs[name] = fn
}

// UnregisterConnectCall unregisters an on-connect call and removes the
// corresponding RPCFunc from g.initRPCs. Future connections to peers will not
// trigger the RPC to be called on them.
func (g *Gateway) UnregisterConnectCall(name string) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.initRPCs[name]; !ok {
		build.Critical("ConnectCall not registered: " + name)
	}
	delete(g.initRPCs, name)
}

// threadedListenPeer listens for new streams on a peer connection and serves them via
// threadedHandleConn.
func (g *Gateway) threadedListenPeer(p *peer) {
	// threadedListenPeer registers to the peerTG instead of the primary thread
	// group because peer connections can be lifetime in length, but can also
	// be short-lived. The fact that they can be lifetime means that they can't
	// call threads.Add as they will block calls to threads.Flush. The fact
	// that they can be short-lived means that threads.OnStop is not a good
	// tool for closing out the threads. Instead, they register to peerTG,
	// which is cleanly closed upon gateway shutdown but will not block any
	// calls to threads.Flush()
	if g.peerTG.Add() != nil {
		return
	}
	defer g.peerTG.Done()

	// Spin up a goroutine to listen for a shutdown signal from both the peer
	// and from the gateway. In the event of either, close the muxado session.
	connClosedChan := make(chan struct{})
	peerCloseChan := make(chan struct{})
	go func() {
		// Signal that the muxado session has been successfully closed, and
		// that this goroutine has terminated.
		defer close(connClosedChan)

		// Listen for a stop signal.
		select {
		case <-g.threads.StopChan():
		case <-peerCloseChan:
		}

		// Can't call Disconnect because it could return sync.ErrStopped.
		g.mu.Lock()
		delete(g.peers, p.NetAddress)
		g.mu.Unlock()
		if err := p.sess.Close(); err != nil {
			g.log.Debugf("WARN: error disconnecting from peer %q: %v", p.NetAddress, err)
		}
	}()

	for {
		conn, err := p.accept()
		if err != nil {
			g.log.Debugf("Peer connection with %v closed: %v\n", p.NetAddress, err)
			break
		}
		// Set the default deadline on the conn.
		err = conn.SetDeadline(time.Now().Add(rpcStdDeadline))
		if err != nil {
			g.log.Printf("Peer connection (%v) deadline could not be set: %v\n", p.NetAddress, err)
			continue
		}

		// The handler is responsible for closing the connection, though a
		// default deadline has been set.
		go g.threadedHandleConn(conn)
		if !g.managedSleep(peerRPCDelay) {
			break
		}
	}
	// Signal that the goroutine can shutdown.
	close(peerCloseChan)
	// Wait for confirmation that the goroutine has shut down before returning
	// and releasing the threadgroup registration.
	<-connClosedChan
}

// threadedHandleConn reads header data from a connection, then routes it to the
// appropriate handler for further processing.
func (g *Gateway) threadedHandleConn(conn modules.PeerConn) {
	defer conn.Close()
	if g.threads.Add() != nil {
		return
	}
	defer g.threads.Done()

	var id rpcID
	err := conn.SetDeadline(time.Now().Add(rpcStdDeadline))
	if err != nil {
		return
	}
	if err := encoding.ReadObject(conn, &id, 8); err != nil {
		return
	}
	// call registered handler for this ID
	g.mu.RLock()
	fn, ok := g.handlers[id]
	g.mu.RUnlock()
	if !ok {
		g.log.Debugf("WARN: incoming conn %v requested unknown RPC \"%v\"", conn.RPCAddr(), id)
		return
	}
	g.log.Debugf("INFO: incoming conn %v requested RPC \"%v\"", conn.RPCAddr(), id)

	// call fn
	err = fn(conn)
	// don't log benign errors
	if err == modules.ErrDuplicateTransactionSet || err == modules.ErrBlockKnown {
		err = nil
	}
	if err != nil {
		g.log.Debugf("WARN: incoming RPC \"%v\" from conn %v failed: %v", id, conn.RPCAddr(), err)
	}
}

// Broadcast calls an RPC on all of the specified peers. The calls are run in
// parallel. Broadcasts are restricted to "one-way" RPCs, which simply write an
// object and disconnect. This is why Broadcast takes an interface{} instead of
// an RPCFunc.
func (g *Gateway) Broadcast(name string, obj interface{}, peers []modules.Peer) {
	if g.threads.Add() != nil {
		return
	}
	defer g.threads.Done()

	g.log.Debugf("INFO: broadcasting RPC %q to %v peers", name, len(peers))

	// only encode obj once, instead of using WriteObject
	enc := encoding.Marshal(obj)
	fn := func(conn modules.PeerConn) error {
		return encoding.WritePrefix(conn, enc)
	}

	var wg sync.WaitGroup
	for _, p := range peers {
		wg.Add(1)
		go func(addr modules.NetAddress) {
			defer wg.Done()
			err := g.managedRPC(addr, name, fn)
			if err != nil {
				g.log.Debugf("WARN: broadcasting RPC %q to peer %q failed (attempting again in 10 seconds): %v", name, addr, err)
				// try one more time before giving up
				select {
				case <-time.After(10 * time.Second):
				case <-g.threads.StopChan():
					return
				}
				err := g.managedRPC(addr, name, fn)
				if err != nil {
					g.log.Debugf("WARN: broadcasting RPC %q to peer %q failed twice: %v", name, addr, err)
				}
			}
		}(p.NetAddress)
	}
	wg.Wait()
}
<|endoftext|>


<commit_before>package gateway

import (
	"errors"
	"sync"
	"time"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/modules"
)

// rpcID is an 8-byte signature that is added to all RPCs to tell the gatway
// what to do with the RPC.
type rpcID [8]byte

// String returns a string representation of an rpcID. Empty elements of rpcID
// will be encoded as spaces.
func (id rpcID) String() string {
	for i := range id {
		if id[i] == 0 {
			id[i] = ' '
		}
	}
	return string(id[:])
}

// handlerName truncates a string to 8 bytes. If len(name) < 8, the remaining
// bytes are 0. A handlerName is specified at the beginning of each network
// call, indicating which function should handle the connection.
func handlerName(name string) (id rpcID) {
	copy(id[:], name)
	return
}

// managedRPC calls an RPC on the given address. managedRPC cannot be called on
// an address that the Gateway is not connected to.
func (g *Gateway) managedRPC(addr modules.NetAddress, name string, fn modules.RPCFunc) error {
	g.mu.RLock()
	peer, ok := g.peers[addr]
	g.mu.RUnlock()
	if !ok {
		return errors.New("can't call RPC on unconnected peer " + string(addr))
	}

	conn, err := peer.open()
	if err != nil {
		return err
	}
	defer conn.Close()

	// write header
	if err := encoding.WriteObject(conn, handlerName(name)); err != nil {
		return err
	}
	// call fn
	return fn(conn)
}

// RPC calls an RPC on the given address. RPC cannot be called on an address
// that the Gateway is not connected to.
func (g *Gateway) RPC(addr modules.NetAddress, name string, fn modules.RPCFunc) error {
	if err := g.threads.Add(); err != nil {
		return err
	}
	defer g.threads.Done()
	return g.managedRPC(addr, name, fn)
}

// RegisterRPC registers an RPCFunc as a handler for a given identifier. To
// call an RPC, use gateway.RPC, supplying the same identifier given to
// RegisterRPC. Identifiers should always use PascalCase. The first 8
// characters of an identifier should be unique, as the identifier used
// internally is truncated to 8 bytes.
func (g *Gateway) RegisterRPC(name string, fn modules.RPCFunc) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.handlers[handlerName(name)]; ok {
		build.Critical("RPC already registered: " + name)
	}
	g.handlers[handlerName(name)] = fn
}

// UnregisterRPC unregisters an RPC and removes the corresponding RPCFunc from
// g.handlers. Future calls to the RPC by peers will fail.
func (g *Gateway) UnregisterRPC(name string) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.handlers[handlerName(name)]; !ok {
		build.Critical("RPC not registered: " + name)
	}
	delete(g.handlers, handlerName(name))
}

// RegisterConnectCall registers a name and RPCFunc to be called on a peer
// upon connecting.
func (g *Gateway) RegisterConnectCall(name string, fn modules.RPCFunc) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.initRPCs[name]; ok {
		build.Critical("ConnectCall already registered: " + name)
	}
	g.initRPCs[name] = fn
}

// UnregisterConnectCall unregisters an on-connect call and removes the
// corresponding RPCFunc from g.initRPCs. Future connections to peers will not
// trigger the RPC to be called on them.
func (g *Gateway) UnregisterConnectCall(name string) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.initRPCs[name]; !ok {
		build.Critical("ConnectCall not registered: " + name)
	}
	delete(g.initRPCs, name)
}

// threadedListenPeer listens for new streams on a peer connection and serves them via
// threadedHandleConn.
func (g *Gateway) threadedListenPeer(p *peer) {
	// threadedListenPeer registers to the peerTG instead of the primary thread
	// group because peer connections can be lifetime in length, but can also
	// be short-lived. The fact that they can be lifetime means that they can't
	// call threads.Add as they will block calls to threads.Flush. The fact
	// that they can be short-lived means that threads.OnStop is not a good
	// tool for closing out the threads. Instead, they register to peerTG,
	// which is cleanly closed upon gateway shutdown but will not block any
	// calls to threads.Flush()
	if g.peerTG.Add() != nil {
		return
	}
	defer g.peerTG.Done()

	// Spin up a goroutine to listen for a shutdown signal from both the peer
	// and from the gateway. In the event of either, close the muxado session.
	connClosedChan := make(chan struct{})
	peerCloseChan := make(chan struct{})
	go func() {
		// Signal that the muxado session has been successfully closed, and
		// that this goroutine has terminated.
		defer close(connClosedChan)

		// Listen for a stop signal.
		select {
		case <-g.threads.StopChan():
		case <-peerCloseChan:
		}

		// Can't call Disconnect because it could return sync.ErrStopped.
		g.mu.Lock()
		delete(g.peers, p.NetAddress)
		g.mu.Unlock()
		if err := p.sess.Close(); err != nil {
			g.log.Debugf("WARN: error disconnecting from peer %q: %v", p.NetAddress, err)
		}
	}()

	for {
		conn, err := p.accept()
		if err != nil {
			g.log.Debugf("Peer connection with %v closed: %v\n", p.NetAddress, err)
			break
		}
		// Set the default deadline on the conn.
		err = conn.SetDeadline(time.Now().Add(rpcStdDeadline))
		if err != nil {
			g.log.Printf("Peer connection (%v) deadline could not be set: %v\n", p.NetAddress, err)
			continue
		}

		// The handler is responsible for closing the connection, though a
		// default deadline has been set.
		go g.threadedHandleConn(conn)
		if !g.managedSleep(peerRPCDelay) {
			break
		}
	}
	// Signal that the goroutine can shutdown.
	close(peerCloseChan)
	// Wait for confirmation that the goroutine has shut down before returning
	// and releasing the threadgroup registration.
	<-connClosedChan
}

// threadedHandleConn reads header data from a connection, then routes it to the
// appropriate handler for further processing.
func (g *Gateway) threadedHandleConn(conn modules.PeerConn) {
	defer conn.Close()
	if g.threads.Add() != nil {
		return
	}
	defer g.threads.Done()

	var id rpcID
	err := conn.SetDeadline(time.Now().Add(rpcStdDeadline))
	if err != nil {
		return
	}
	if err := encoding.ReadObject(conn, &id, 8); err != nil {
		return
	}
	// call registered handler for this ID
	g.mu.RLock()
	fn, ok := g.handlers[id]
	g.mu.RUnlock()
	if !ok {
		g.log.Debugf("WARN: incoming conn %v requested unknown RPC \"%v\"", conn.RPCAddr(), id)
		return
	}
	g.log.Debugf("INFO: incoming conn %v requested RPC \"%v\"", conn.RPCAddr(), id)

	// call fn
	err = fn(conn)
	// don't log benign errors
	if err == modules.ErrDuplicateTransactionSet || err == modules.ErrBlockKnown {
		err = nil
	}
	if err != nil {
		g.log.Debugf("WARN: incoming RPC \"%v\" from conn %v failed: %v", id, conn.RPCAddr(), err)
	}
}

// Broadcast calls an RPC on all of the specified peers. The calls are run in
// parallel. Broadcasts are restricted to "one-way" RPCs, which simply write an
// object and disconnect. This is why Broadcast takes an interface{} instead of
// an RPCFunc.
func (g *Gateway) Broadcast(name string, obj interface{}, peers []modules.Peer) {
	if g.threads.Add() != nil {
		return
	}
	defer g.threads.Done()

	g.log.Debugf("INFO: broadcasting RPC %q to %v peers", name, len(peers))

	// only encode obj once, instead of using WriteObject
	enc := encoding.Marshal(obj)
	fn := func(conn modules.PeerConn) error {
		return encoding.WritePrefix(conn, enc)
	}

	var wg sync.WaitGroup
	for _, p := range peers {
		wg.Add(1)
		go func(addr modules.NetAddress) {
			defer wg.Done()
			err := g.managedRPC(addr, name, fn)
			if err != nil {
				g.log.Debugf("WARN: broadcasting RPC %q to peer %q failed (attempting again in 10 seconds): %v", name, addr, err)
				// try one more time before giving up
				select {
				case <-time.After(10 * time.Second):
				case <-g.threads.StopChan():
					return
				}
				err := g.managedRPC(addr, name, fn)
				if err != nil {
					g.log.Debugf("WARN: broadcasting RPC %q to peer %q failed twice: %v", name, addr, err)
				}
			}
		}(p.NetAddress)
	}
	wg.Wait()
}
<commit_msg>disconnect from peers if RPC can't be initiated<commit_after>package gateway

import (
	"errors"
	"sync"
	"time"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/modules"
)

// rpcID is an 8-byte signature that is added to all RPCs to tell the gatway
// what to do with the RPC.
type rpcID [8]byte

// String returns a string representation of an rpcID. Empty elements of rpcID
// will be encoded as spaces.
func (id rpcID) String() string {
	for i := range id {
		if id[i] == 0 {
			id[i] = ' '
		}
	}
	return string(id[:])
}

// handlerName truncates a string to 8 bytes. If len(name) < 8, the remaining
// bytes are 0. A handlerName is specified at the beginning of each network
// call, indicating which function should handle the connection.
func handlerName(name string) (id rpcID) {
	copy(id[:], name)
	return
}

// managedRPC calls an RPC on the given address. managedRPC cannot be called on
// an address that the Gateway is not connected to.
func (g *Gateway) managedRPC(addr modules.NetAddress, name string, fn modules.RPCFunc) error {
	g.mu.RLock()
	peer, ok := g.peers[addr]
	g.mu.RUnlock()
	if !ok {
		return errors.New("can't call RPC on unconnected peer " + string(addr))
	}

	conn, err := peer.open()
	if err != nil {
		// peer probably disconnected without sending a shutdown signal;
		// disconnect from them
		g.log.Debugf("Could not initiate RPC with %v; disconnecting", addr)
		if derr := g.Disconnect(addr); derr != nil {
			g.log.Debugf("Could not cleanly disconnect from %v: %v", addr, derr)
		}
		return err
	}
	defer conn.Close()

	// write header
	if err := encoding.WriteObject(conn, handlerName(name)); err != nil {
		return err
	}
	// call fn
	return fn(conn)
}

// RPC calls an RPC on the given address. RPC cannot be called on an address
// that the Gateway is not connected to.
func (g *Gateway) RPC(addr modules.NetAddress, name string, fn modules.RPCFunc) error {
	if err := g.threads.Add(); err != nil {
		return err
	}
	defer g.threads.Done()
	return g.managedRPC(addr, name, fn)
}

// RegisterRPC registers an RPCFunc as a handler for a given identifier. To
// call an RPC, use gateway.RPC, supplying the same identifier given to
// RegisterRPC. Identifiers should always use PascalCase. The first 8
// characters of an identifier should be unique, as the identifier used
// internally is truncated to 8 bytes.
func (g *Gateway) RegisterRPC(name string, fn modules.RPCFunc) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.handlers[handlerName(name)]; ok {
		build.Critical("RPC already registered: " + name)
	}
	g.handlers[handlerName(name)] = fn
}

// UnregisterRPC unregisters an RPC and removes the corresponding RPCFunc from
// g.handlers. Future calls to the RPC by peers will fail.
func (g *Gateway) UnregisterRPC(name string) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.handlers[handlerName(name)]; !ok {
		build.Critical("RPC not registered: " + name)
	}
	delete(g.handlers, handlerName(name))
}

// RegisterConnectCall registers a name and RPCFunc to be called on a peer
// upon connecting.
func (g *Gateway) RegisterConnectCall(name string, fn modules.RPCFunc) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.initRPCs[name]; ok {
		build.Critical("ConnectCall already registered: " + name)
	}
	g.initRPCs[name] = fn
}

// UnregisterConnectCall unregisters an on-connect call and removes the
// corresponding RPCFunc from g.initRPCs. Future connections to peers will not
// trigger the RPC to be called on them.
func (g *Gateway) UnregisterConnectCall(name string) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.initRPCs[name]; !ok {
		build.Critical("ConnectCall not registered: " + name)
	}
	delete(g.initRPCs, name)
}

// threadedListenPeer listens for new streams on a peer connection and serves them via
// threadedHandleConn.
func (g *Gateway) threadedListenPeer(p *peer) {
	// threadedListenPeer registers to the peerTG instead of the primary thread
	// group because peer connections can be lifetime in length, but can also
	// be short-lived. The fact that they can be lifetime means that they can't
	// call threads.Add as they will block calls to threads.Flush. The fact
	// that they can be short-lived means that threads.OnStop is not a good
	// tool for closing out the threads. Instead, they register to peerTG,
	// which is cleanly closed upon gateway shutdown but will not block any
	// calls to threads.Flush()
	if g.peerTG.Add() != nil {
		return
	}
	defer g.peerTG.Done()

	// Spin up a goroutine to listen for a shutdown signal from both the peer
	// and from the gateway. In the event of either, close the muxado session.
	connClosedChan := make(chan struct{})
	peerCloseChan := make(chan struct{})
	go func() {
		// Signal that the muxado session has been successfully closed, and
		// that this goroutine has terminated.
		defer close(connClosedChan)

		// Listen for a stop signal.
		select {
		case <-g.threads.StopChan():
		case <-peerCloseChan:
		}

		// Can't call Disconnect because it could return sync.ErrStopped.
		g.mu.Lock()
		delete(g.peers, p.NetAddress)
		g.mu.Unlock()
		if err := p.sess.Close(); err != nil {
			g.log.Debugf("WARN: error disconnecting from peer %q: %v", p.NetAddress, err)
		}
	}()

	for {
		conn, err := p.accept()
		if err != nil {
			g.log.Debugf("Peer connection with %v closed: %v\n", p.NetAddress, err)
			break
		}
		// Set the default deadline on the conn.
		err = conn.SetDeadline(time.Now().Add(rpcStdDeadline))
		if err != nil {
			g.log.Printf("Peer connection (%v) deadline could not be set: %v\n", p.NetAddress, err)
			continue
		}

		// The handler is responsible for closing the connection, though a
		// default deadline has been set.
		go g.threadedHandleConn(conn)
		if !g.managedSleep(peerRPCDelay) {
			break
		}
	}
	// Signal that the goroutine can shutdown.
	close(peerCloseChan)
	// Wait for confirmation that the goroutine has shut down before returning
	// and releasing the threadgroup registration.
	<-connClosedChan
}

// threadedHandleConn reads header data from a connection, then routes it to the
// appropriate handler for further processing.
func (g *Gateway) threadedHandleConn(conn modules.PeerConn) {
	defer conn.Close()
	if g.threads.Add() != nil {
		return
	}
	defer g.threads.Done()

	var id rpcID
	err := conn.SetDeadline(time.Now().Add(rpcStdDeadline))
	if err != nil {
		return
	}
	if err := encoding.ReadObject(conn, &id, 8); err != nil {
		return
	}
	// call registered handler for this ID
	g.mu.RLock()
	fn, ok := g.handlers[id]
	g.mu.RUnlock()
	if !ok {
		g.log.Debugf("WARN: incoming conn %v requested unknown RPC \"%v\"", conn.RPCAddr(), id)
		return
	}
	g.log.Debugf("INFO: incoming conn %v requested RPC \"%v\"", conn.RPCAddr(), id)

	// call fn
	err = fn(conn)
	// don't log benign errors
	if err == modules.ErrDuplicateTransactionSet || err == modules.ErrBlockKnown {
		err = nil
	}
	if err != nil {
		g.log.Debugf("WARN: incoming RPC \"%v\" from conn %v failed: %v", id, conn.RPCAddr(), err)
	}
}

// Broadcast calls an RPC on all of the specified peers. The calls are run in
// parallel. Broadcasts are restricted to "one-way" RPCs, which simply write an
// object and disconnect. This is why Broadcast takes an interface{} instead of
// an RPCFunc.
func (g *Gateway) Broadcast(name string, obj interface{}, peers []modules.Peer) {
	if g.threads.Add() != nil {
		return
	}
	defer g.threads.Done()

	g.log.Debugf("INFO: broadcasting RPC %q to %v peers", name, len(peers))

	// only encode obj once, instead of using WriteObject
	enc := encoding.Marshal(obj)
	fn := func(conn modules.PeerConn) error {
		return encoding.WritePrefix(conn, enc)
	}

	var wg sync.WaitGroup
	for _, p := range peers {
		wg.Add(1)
		go func(addr modules.NetAddress) {
			defer wg.Done()
			err := g.managedRPC(addr, name, fn)
			if err != nil {
				g.log.Debugf("WARN: broadcasting RPC %q to peer %q failed (attempting again in 10 seconds): %v", name, addr, err)
				// try one more time before giving up
				select {
				case <-time.After(10 * time.Second):
				case <-g.threads.StopChan():
					return
				}
				err := g.managedRPC(addr, name, fn)
				if err != nil {
					g.log.Debugf("WARN: broadcasting RPC %q to peer %q failed twice: %v", name, addr, err)
				}
			}
		}(p.NetAddress)
	}
	wg.Wait()
}
<|endoftext|>


<commit_before>package miekg

import (
	"net"
	"strings"
	"time"

	"github.com/miekg/dns"
	"github.com/zmap/zdns"
)

type Answer struct {
	Ttl    uint32 `json:"ttl,omitempty"`
	Type   string `json:"type,omitempty"`
	Name   string `json:"name,omitempty"`
	Answer string `json:"answer,omitempty"`
}

type MXAnswer struct {
	Ttl        uint32 `json:"ttl,omitempty"`
	Type       string `json:"type,omitempty"`
	Name       string `json:"name,omitempty"`
	Answer     string `json:"answer,omitempty"`
	Preference uint16 `json:"preference"`
}

// result to be returned by scan of host
type Result struct {
	Answers     []interface{} `json:"answers"`
	Additional  []interface{} `json:"additionals"`
	Authorities []interface{} `json:"authorities"`
	Protocol    string        `json:"protocol"`
}

type Lookup struct {
	DNSType dns.Type
	Prefix  string
	zdns.BaseLookup
}

type GlobalLookupFactory struct {
}

type RoutineLookupFactory struct {
	Client    *dns.Client
	TCPClient *dns.Client
}

func (s *RoutineLookupFactory) Initialize(t time.Duration) {
	s.Client = new(dns.Client)
	s.Client.Timeout = t
	s.TCPClient = new(dns.Client)
	s.TCPClient.Net = "tcp"
	s.TCPClient.Timeout = t
}

func dotName(name string) string {
	return strings.Join([]string{name, "."}, "")
}

func ParseAnswer(ans dns.RR) interface{} {
	var retv Answer
	if a, ok := ans.(*dns.A); ok {
		retv = Answer{Ttl: a.Hdr.Ttl, Type: dns.Type(a.Hdr.Rrtype).String(), Name: a.Hdr.Name, Answer: a.A.String()}
	} else if aaaa, ok := ans.(*dns.AAAA); ok {
		retv = Answer{Ttl: aaaa.Hdr.Ttl, Type: dns.Type(aaaa.Hdr.Rrtype).String(), Name: aaaa.Hdr.Name, Answer: aaaa.AAAA.String()}
	} else if cname, ok := ans.(*dns.CNAME); ok {
		retv = Answer{Ttl: cname.Hdr.Ttl, Type: dns.Type(cname.Hdr.Rrtype).String(), Name: cname.Hdr.Name, Answer: cname.Target}
	} else if txt, ok := ans.(*dns.TXT); ok {
		retv = Answer{Ttl: txt.Hdr.Ttl, Type: dns.Type(txt.Hdr.Rrtype).String(), Name: txt.Hdr.Name, Answer: strings.Join(txt.Txt, "\n")}
	} else if ns, ok := ans.(*dns.NS); ok {
		retv = Answer{Ttl: ns.Hdr.Ttl, Type: dns.Type(ns.Hdr.Rrtype).String(), Name: ns.Hdr.Name, Answer: ns.Ns}
	} else if ptr, ok := ans.(*dns.PTR); ok {
		retv = Answer{Ttl: ptr.Hdr.Ttl, Type: dns.Type(ptr.Hdr.Rrtype).String(), Name: ptr.Hdr.Name, Answer: ptr.Ptr}
	} else if spf, ok := ans.(*dns.SPF); ok {
		retv = Answer{Ttl: spf.Hdr.Ttl, Type: dns.Type(spf.Hdr.Rrtype).String(), Name: spf.Hdr.Name, Answer: spf.String()}

	} else if mx, ok := ans.(*dns.MX); ok {
		return MXAnswer{
			Ttl:        mx.Hdr.Ttl,
			Type:       dns.Type(mx.Hdr.Rrtype).String(),
			Name:       mx.Hdr.Name,
			Answer:     mx.Mx,
			Preference: mx.Preference,
		}
	} else if caa, ok := ans.(*dns.CAA); ok {
		return struct {
			Ttl   uint32 `json:"ttl"`
			Type  string `json:"type"`
			Tag   string `json:"tag"`
			Value string `json:"value"`
			Flag  uint8  `json:"flag"`
		}{
			Ttl:   caa.Hdr.Ttl,
			Type:  dns.Type(caa.Hdr.Rrtype).String(),
			Tag:   caa.Tag,
			Value: caa.Value,
			Flag:  caa.Flag,
		}
	} else if soa, ok := ans.(*dns.SOA); ok {
		return struct {
			Ttl     uint32 `json:"ttl"`
			Type    string `json:"type"`
			Name    string `json:"name"`
			Ns      string `json:"ns"`
			Mbox    string `json:"mbox"`
			Serial  uint32 `json:"serial"`
			Refresh uint32 `json:"refresh"`
			Retry   uint32 `json:"retry"`
			Expire  uint32 `json:"expire"`
			Minttl  uint32 `json:"min_ttl"`
		}{
			Ttl:     soa.Hdr.Ttl,
			Type:    dns.Type(soa.Hdr.Rrtype).String(),
			Name:    soa.Hdr.Name,
			Ns:      strings.TrimSuffix(soa.Ns, "."),
			Mbox:    soa.Mbox,
			Serial:  soa.Serial,
			Refresh: soa.Refresh,
			Retry:   soa.Retry,
			Expire:  soa.Expire,
			Minttl:  soa.Minttl,
		}
	} else if any, ok := ans.(*dns.ANY); ok {
		return struct {
			Type string
		}{
			Type: dns.Type(any.Hdr.Rrtype).String(),
		}
	}
	retv.Name = strings.TrimSuffix(retv.Name, ".")
	return retv
}

func TranslateMiekgErrorCode(err int) zdns.Status {
	return zdns.Status(dns.RcodeToString[err])
}

func DoLookup(udp *dns.Client, tcp *dns.Client, nameServer string, dnsType uint16, name string) (interface{}, zdns.Status, error) {
	// this is where we do scanning
	res := Result{Answers: []interface{}{}, Authorities: []interface{}{}, Additional: []interface{}{}}

	m := new(dns.Msg)
	m.SetQuestion(dotName(name), dnsType)
	m.RecursionDesired = true

	useTCP := false
	res.Protocol = "udp"
	r, _, err := udp.Exchange(m, nameServer)
	if err == dns.ErrTruncated {
		r, _, err = tcp.Exchange(m, nameServer)
		useTCP = true
		res.Protocol = "tcp"
	}
	if err != nil || r == nil {
		if nerr, ok := err.(net.Error); ok {
			if nerr.Timeout() {
				return nil, zdns.STATUS_TIMEOUT, nil
			} else if nerr.Temporary() {
				return nil, zdns.STATUS_TEMPORARY, err
			}
		}
		return nil, zdns.STATUS_ERROR, err
	}
	if r.Rcode == dns.RcodeBadTrunc && !useTCP {
		r, _, err = tcp.Exchange(m, nameServer)
	}
	if err != nil || r == nil {
		return nil, zdns.STATUS_ERROR, err
	}
	if r.Rcode != dns.RcodeSuccess {
		return nil, TranslateMiekgErrorCode(r.Rcode), nil
	}
	for _, ans := range r.Answer {
		inner := ParseAnswer(ans)
		if inner != nil {
			res.Answers = append(res.Answers, inner)
		}
	}
	for _, ans := range r.Extra {
		inner := ParseAnswer(ans)
		if inner != nil {
			res.Additional = append(res.Additional, inner)
		}
	}
	for _, ans := range r.Ns {
		inner := ParseAnswer(ans)
		if inner != nil {
			res.Authorities = append(res.Authorities, inner)
		}
	}
	return res, zdns.STATUS_NOERROR, nil
}

func DoTxtLookup(udp *dns.Client, tcp *dns.Client, nameServer string, prefix string, name string) (string, zdns.Status, error) {
	res, status, err := DoLookup(udp, tcp, nameServer, dns.TypeTXT, name)
	if status != zdns.STATUS_NOERROR {
		return "", status, err
	}
	if parsedResult, ok := res.(Result); ok {
		for _, a := range parsedResult.Answers {
			ans, _ := a.(Answer)
			if strings.HasPrefix(ans.Answer, prefix) {
				return ans.Answer, zdns.STATUS_NOERROR, err
			}
		}
	}
	return "", zdns.STATUS_NO_RECORD, nil
}
<commit_msg>handle dname<commit_after>package miekg

import (
	"net"
	"strings"
	"time"

	"github.com/miekg/dns"
	"github.com/zmap/zdns"
)

type Answer struct {
	Ttl    uint32 `json:"ttl,omitempty"`
	Type   string `json:"type,omitempty"`
	Name   string `json:"name,omitempty"`
	Answer string `json:"answer,omitempty"`
}

type MXAnswer struct {
	Ttl        uint32 `json:"ttl,omitempty"`
	Type       string `json:"type,omitempty"`
	Name       string `json:"name,omitempty"`
	Answer     string `json:"answer,omitempty"`
	Preference uint16 `json:"preference"`
}

// result to be returned by scan of host
type Result struct {
	Answers     []interface{} `json:"answers"`
	Additional  []interface{} `json:"additionals"`
	Authorities []interface{} `json:"authorities"`
	Protocol    string        `json:"protocol"`
}

type Lookup struct {
	DNSType dns.Type
	Prefix  string
	zdns.BaseLookup
}

type GlobalLookupFactory struct {
}

type RoutineLookupFactory struct {
	Client    *dns.Client
	TCPClient *dns.Client
}

func (s *RoutineLookupFactory) Initialize(t time.Duration) {
	s.Client = new(dns.Client)
	s.Client.Timeout = t
	s.TCPClient = new(dns.Client)
	s.TCPClient.Net = "tcp"
	s.TCPClient.Timeout = t
}

func dotName(name string) string {
	return strings.Join([]string{name, "."}, "")
}

func ParseAnswer(ans dns.RR) interface{} {
	var retv Answer
	if a, ok := ans.(*dns.A); ok {
		retv = Answer{Ttl: a.Hdr.Ttl, Type: dns.Type(a.Hdr.Rrtype).String(), Name: a.Hdr.Name, Answer: a.A.String()}
	} else if aaaa, ok := ans.(*dns.AAAA); ok {
		retv = Answer{Ttl: aaaa.Hdr.Ttl, Type: dns.Type(aaaa.Hdr.Rrtype).String(), Name: aaaa.Hdr.Name, Answer: aaaa.AAAA.String()}
	} else if cname, ok := ans.(*dns.CNAME); ok {
		retv = Answer{Ttl: cname.Hdr.Ttl, Type: dns.Type(cname.Hdr.Rrtype).String(), Name: cname.Hdr.Name, Answer: cname.Target}
	} else if dname, ok := ans.(*dns.DNAME); ok {
		retv = Answer{Ttl: dname.Hdr.Ttl, Type: dns.Type(dname.Hdr.Rrtype).String(), Name: dname.Hdr.Name, Answer: dname.Target}
	} else if txt, ok := ans.(*dns.TXT); ok {
		retv = Answer{Ttl: txt.Hdr.Ttl, Type: dns.Type(txt.Hdr.Rrtype).String(), Name: txt.Hdr.Name, Answer: strings.Join(txt.Txt, "\n")}
	} else if ns, ok := ans.(*dns.NS); ok {
		retv = Answer{Ttl: ns.Hdr.Ttl, Type: dns.Type(ns.Hdr.Rrtype).String(), Name: ns.Hdr.Name, Answer: ns.Ns}
	} else if ptr, ok := ans.(*dns.PTR); ok {
		retv = Answer{Ttl: ptr.Hdr.Ttl, Type: dns.Type(ptr.Hdr.Rrtype).String(), Name: ptr.Hdr.Name, Answer: ptr.Ptr}
	} else if spf, ok := ans.(*dns.SPF); ok {
		retv = Answer{Ttl: spf.Hdr.Ttl, Type: dns.Type(spf.Hdr.Rrtype).String(), Name: spf.Hdr.Name, Answer: spf.String()}

	} else if mx, ok := ans.(*dns.MX); ok {
		return MXAnswer{
			Ttl:        mx.Hdr.Ttl,
			Type:       dns.Type(mx.Hdr.Rrtype).String(),
			Name:       mx.Hdr.Name,
			Answer:     mx.Mx,
			Preference: mx.Preference,
		}
	} else if caa, ok := ans.(*dns.CAA); ok {
		return struct {
			Ttl   uint32 `json:"ttl"`
			Type  string `json:"type"`
			Tag   string `json:"tag"`
			Value string `json:"value"`
			Flag  uint8  `json:"flag"`
		}{
			Ttl:   caa.Hdr.Ttl,
			Type:  dns.Type(caa.Hdr.Rrtype).String(),
			Tag:   caa.Tag,
			Value: caa.Value,
			Flag:  caa.Flag,
		}
	} else if soa, ok := ans.(*dns.SOA); ok {
		return struct {
			Ttl     uint32 `json:"ttl"`
			Type    string `json:"type"`
			Name    string `json:"name"`
			Ns      string `json:"ns"`
			Mbox    string `json:"mbox"`
			Serial  uint32 `json:"serial"`
			Refresh uint32 `json:"refresh"`
			Retry   uint32 `json:"retry"`
			Expire  uint32 `json:"expire"`
			Minttl  uint32 `json:"min_ttl"`
		}{
			Ttl:     soa.Hdr.Ttl,
			Type:    dns.Type(soa.Hdr.Rrtype).String(),
			Name:    soa.Hdr.Name,
			Ns:      strings.TrimSuffix(soa.Ns, "."),
			Mbox:    soa.Mbox,
			Serial:  soa.Serial,
			Refresh: soa.Refresh,
			Retry:   soa.Retry,
			Expire:  soa.Expire,
			Minttl:  soa.Minttl,
		}
	} else if any, ok := ans.(*dns.ANY); ok {
		return struct {
			Type string
		}{
			Type: dns.Type(any.Hdr.Rrtype).String(),
		}
	}
	retv.Name = strings.TrimSuffix(retv.Name, ".")
	return retv
}

func TranslateMiekgErrorCode(err int) zdns.Status {
	return zdns.Status(dns.RcodeToString[err])
}

func DoLookup(udp *dns.Client, tcp *dns.Client, nameServer string, dnsType uint16, name string) (interface{}, zdns.Status, error) {
	// this is where we do scanning
	res := Result{Answers: []interface{}{}, Authorities: []interface{}{}, Additional: []interface{}{}}

	m := new(dns.Msg)
	m.SetQuestion(dotName(name), dnsType)
	m.RecursionDesired = true

	useTCP := false
	res.Protocol = "udp"
	r, _, err := udp.Exchange(m, nameServer)
	if err == dns.ErrTruncated {
		r, _, err = tcp.Exchange(m, nameServer)
		useTCP = true
		res.Protocol = "tcp"
	}
	if err != nil || r == nil {
		if nerr, ok := err.(net.Error); ok {
			if nerr.Timeout() {
				return nil, zdns.STATUS_TIMEOUT, nil
			} else if nerr.Temporary() {
				return nil, zdns.STATUS_TEMPORARY, err
			}
		}
		return nil, zdns.STATUS_ERROR, err
	}
	if r.Rcode == dns.RcodeBadTrunc && !useTCP {
		r, _, err = tcp.Exchange(m, nameServer)
	}
	if err != nil || r == nil {
		return nil, zdns.STATUS_ERROR, err
	}
	if r.Rcode != dns.RcodeSuccess {
		return nil, TranslateMiekgErrorCode(r.Rcode), nil
	}
	for _, ans := range r.Answer {
		inner := ParseAnswer(ans)
		if inner != nil {
			res.Answers = append(res.Answers, inner)
		}
	}
	for _, ans := range r.Extra {
		inner := ParseAnswer(ans)
		if inner != nil {
			res.Additional = append(res.Additional, inner)
		}
	}
	for _, ans := range r.Ns {
		inner := ParseAnswer(ans)
		if inner != nil {
			res.Authorities = append(res.Authorities, inner)
		}
	}
	return res, zdns.STATUS_NOERROR, nil
}

func DoTxtLookup(udp *dns.Client, tcp *dns.Client, nameServer string, prefix string, name string) (string, zdns.Status, error) {
	res, status, err := DoLookup(udp, tcp, nameServer, dns.TypeTXT, name)
	if status != zdns.STATUS_NOERROR {
		return "", status, err
	}
	if parsedResult, ok := res.(Result); ok {
		for _, a := range parsedResult.Answers {
			ans, _ := a.(Answer)
			if strings.HasPrefix(ans.Answer, prefix) {
				return ans.Answer, zdns.STATUS_NOERROR, err
			}
		}
	}
	return "", zdns.STATUS_NO_RECORD, nil
}
<|endoftext|>


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

package proxy

import (
	"net/http"
	"net/url"
	"os"
	"sync"

	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/setting"

	"github.com/gobwas/glob"
)

var (
	once         sync.Once
	hostMatchers []glob.Glob
)

// GetProxyURL returns proxy url
func GetProxyURL() string {
	if !setting.Proxy.Enabled {
		return ""
	}

	if setting.Proxy.ProxyURL == "" {
		if os.Getenv("http_proxy") != "" {
			return os.Getenv("http_proxy")
		}
		return os.Getenv("https_proxy")
	}
	return setting.Proxy.ProxyURL
}

// Match return true if url needs to be proxied
func Match(u string) bool {
	if !setting.Proxy.Enabled {
		return false
	}

	// enforce do once
	Proxy()

	for _, v := range hostMatchers {
		if v.Match(u) {
			return true
		}
	}
	return false
}

// Proxy returns the system proxy
func Proxy() func(req *http.Request) (*url.URL, error) {
	if !setting.Proxy.Enabled {
		return nil
	}
	if setting.Proxy.ProxyURL == "" {
		return http.ProxyFromEnvironment
	}

	once.Do(func() {
		for _, h := range setting.Proxy.ProxyHosts {
			if g, err := glob.Compile(h); err == nil {
				hostMatchers = append(hostMatchers, g)
			} else {
				log.Error("glob.Compile %s failed: %v", h, err)
			}
		}
	})

	return func(req *http.Request) (*url.URL, error) {
		for _, v := range hostMatchers {
			if v.Match(req.URL.Host) {
				return http.ProxyURL(setting.Proxy.ProxyURLFixed)(req)
			}
		}
		return http.ProxyFromEnvironment(req)
	}
}
<commit_msg>Return nil proxy function if proxy not enabled (#16742)<commit_after>// Copyright 2021 The Gitea Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package proxy

import (
	"net/http"
	"net/url"
	"os"
	"sync"

	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/setting"

	"github.com/gobwas/glob"
)

var (
	once         sync.Once
	hostMatchers []glob.Glob
)

// GetProxyURL returns proxy url
func GetProxyURL() string {
	if !setting.Proxy.Enabled {
		return ""
	}

	if setting.Proxy.ProxyURL == "" {
		if os.Getenv("http_proxy") != "" {
			return os.Getenv("http_proxy")
		}
		return os.Getenv("https_proxy")
	}
	return setting.Proxy.ProxyURL
}

// Match return true if url needs to be proxied
func Match(u string) bool {
	if !setting.Proxy.Enabled {
		return false
	}

	// enforce do once
	Proxy()

	for _, v := range hostMatchers {
		if v.Match(u) {
			return true
		}
	}
	return false
}

// Proxy returns the system proxy
func Proxy() func(req *http.Request) (*url.URL, error) {
	if !setting.Proxy.Enabled {
		return func(req *http.Request) (*url.URL, error) {
			return nil, nil
		}
	}
	if setting.Proxy.ProxyURL == "" {
		return http.ProxyFromEnvironment
	}

	once.Do(func() {
		for _, h := range setting.Proxy.ProxyHosts {
			if g, err := glob.Compile(h); err == nil {
				hostMatchers = append(hostMatchers, g)
			} else {
				log.Error("glob.Compile %s failed: %v", h, err)
			}
		}
	})

	return func(req *http.Request) (*url.URL, error) {
		for _, v := range hostMatchers {
			if v.Match(req.URL.Host) {
				return http.ProxyURL(setting.Proxy.ProxyURLFixed)(req)
			}
		}
		return http.ProxyFromEnvironment(req)
	}
}
<|endoftext|>


<commit_before>package storage

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/containers/storage/pkg/ioutils"
	"github.com/containers/storage/pkg/stringid"
)

var (
	// ErrContainerUnknown indicates that there was no container with the specified name or ID
	ErrContainerUnknown = errors.New("container not known")
)

// A Container is a reference to a read-write layer with a metadata string.
// ID is either one specified at create-time or a randomly-generated value.
// Names is an optional set of user-defined convenience values.
// ImageID is the ID of the image which was used to create the container.
// LayerID is the ID of the read-write layer for the container itself.
// It is assumed that the image's top layer is the parent of the container's
// read-write layer.
type Container struct {
	ID       string   `json:"id"`
	Names    []string `json:"names,omitempty"`
	ImageID  string   `json:"image"`
	LayerID  string   `json:"layer"`
	Metadata string   `json:"metadata,omitempty"`
}

// ContainerStore provides bookkeeping for information about Containers.
//
// Create creates a container that has a specified ID (or a random one) and an
// optional name, based on the specified image, using the specified layer as
// its read-write layer.
//
// SetMetadata replaces the metadata associated with a container with the
// supplied value.
//
// Exists checks if there is a container with the given ID or name.
//
// Get retrieves information about a container given an ID or name.
//
// Delete removes the record of the container.
//
// Wipe removes records of all containers.
//
// Lookup attempts to translate a name to an ID.  Most methods do this
// implicitly.
//
// Containers returns a slice enumerating the known containers.
type ContainerStore interface {
	Store
	Create(id string, names []string, image, layer, metadata string) (*Container, error)
	SetMetadata(id, metadata string) error
	SetNames(id string, names []string) error
	Get(id string) (*Container, error)
	Exists(id string) bool
	Delete(id string) error
	Wipe() error
	Lookup(name string) (string, error)
	Containers() ([]Container, error)
}

type containerStore struct {
	lockfile   Locker
	dir        string
	containers []Container
	byid       map[string]*Container
	byname     map[string]*Container
}

func (r *containerStore) Containers() ([]Container, error) {
	return r.containers, nil
}

func (r *containerStore) Load() error {
	rpath := filepath.Join(r.dir, "containers.json")
	data, err := ioutil.ReadFile(rpath)
	if err != nil && !os.IsNotExist(err) {
		return err
	}
	containers := []Container{}
	ids := make(map[string]*Container)
	names := make(map[string]*Container)
	if err = json.Unmarshal(data, &containers); len(data) == 0 || err == nil {
		for n, container := range containers {
			ids[container.ID] = &containers[n]
			for _, name := range container.Names {
				names[name] = &containers[n]
			}
		}
	}
	r.containers = containers
	r.byid = ids
	r.byname = names
	return nil
}

func (r *containerStore) Save() error {
	rpath := filepath.Join(r.dir, "containers.json")
	jdata, err := json.Marshal(&r.containers)
	if err != nil {
		return err
	}
	return ioutils.AtomicWriteFile(rpath, jdata, 0600)
}

func newContainerStore(dir string) (ContainerStore, error) {
	if err := os.MkdirAll(dir, 0700); err != nil {
		return nil, err
	}
	lockfile, err := GetLockfile(filepath.Join(dir, "containers.lock"))
	if err != nil {
		return nil, err
	}
	lockfile.Lock()
	defer lockfile.Unlock()
	cstore := containerStore{
		lockfile:   lockfile,
		dir:        dir,
		containers: []Container{},
		byid:       make(map[string]*Container),
		byname:     make(map[string]*Container),
	}
	if err := cstore.Load(); err != nil {
		return nil, err
	}
	return &cstore, nil
}

func (r *containerStore) Create(id string, names []string, image, layer, metadata string) (container *Container, err error) {
	if id == "" {
		id = stringid.GenerateRandomID()
	}
	for _, name := range names {
		if _, nameInUse := r.byname[name]; nameInUse {
			return nil, ErrDuplicateName
		}
	}
	if err == nil {
		newContainer := Container{
			ID:       id,
			Names:    names,
			ImageID:  image,
			LayerID:  layer,
			Metadata: metadata,
		}
		r.containers = append(r.containers, newContainer)
		container = &r.containers[len(r.containers)-1]
		r.byid[id] = container
		for _, name := range names {
			r.byname[name] = container
		}
		err = r.Save()
	}
	return container, err
}

func (r *containerStore) SetMetadata(id, metadata string) error {
	if container, ok := r.byname[id]; ok {
		id = container.ID
	}
	if container, ok := r.byid[id]; ok {
		container.Metadata = metadata
		return r.Save()
	}
	return ErrContainerUnknown
}

func (r *containerStore) SetNames(id string, names []string) error {
	if container, ok := r.byname[id]; ok {
		id = container.ID
	}
	if container, ok := r.byid[id]; ok {
		for _, name := range container.Names {
			delete(r.byname, name)
		}
		for _, name := range names {
			r.byname[name] = container
		}
		container.Names = names
		return r.Save()
	}
	return ErrContainerUnknown
}

func (r *containerStore) Delete(id string) error {
	if container, ok := r.byname[id]; ok {
		id = container.ID
	}
	if container, ok := r.byid[id]; ok {
		newContainers := []Container{}
		for _, candidate := range r.containers {
			if candidate.ID != id {
				newContainers = append(newContainers, candidate)
			}
		}
		for _, name := range container.Names {
			delete(r.byname, name)
		}
		r.containers = newContainers
		if err := r.Save(); err != nil {
			return err
		}
	}
	return nil
}

func (r *containerStore) Get(id string) (*Container, error) {
	if c, ok := r.byname[id]; ok {
		return c, nil
	}
	if c, ok := r.byid[id]; ok {
		return c, nil
	}
	return nil, ErrContainerUnknown
}

func (r *containerStore) Lookup(name string) (id string, err error) {
	container, ok := r.byname[name]
	if !ok {
		return "", ErrContainerUnknown
	}
	return container.ID, nil
}

func (r *containerStore) Exists(id string) bool {
	if _, ok := r.byname[id]; ok {
		return true
	}
	if _, ok := r.byid[id]; ok {
		return true
	}
	return false
}

func (r *containerStore) Wipe() error {
	ids := []string{}
	for id := range r.byid {
		ids = append(ids, id)
	}
	for _, id := range ids {
		if err := r.Delete(id); err != nil {
			return err
		}
	}
	return nil
}

func (r *containerStore) Lock() {
	r.lockfile.Lock()
}

func (r *containerStore) Unlock() {
	r.lockfile.Unlock()
}

func (r *containerStore) Touch() error {
	return r.lockfile.Touch()
}

func (r *containerStore) Modified() (bool, error) {
	return r.lockfile.Modified()
}
<commit_msg>Look up containers by their layers<commit_after>package storage

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/containers/storage/pkg/ioutils"
	"github.com/containers/storage/pkg/stringid"
)

var (
	// ErrContainerUnknown indicates that there was no container with the specified name or ID
	ErrContainerUnknown = errors.New("container not known")
)

// A Container is a reference to a read-write layer with a metadata string.
// ID is either one specified at create-time or a randomly-generated value.
// Names is an optional set of user-defined convenience values.
// ImageID is the ID of the image which was used to create the container.
// LayerID is the ID of the read-write layer for the container itself.
// It is assumed that the image's top layer is the parent of the container's
// read-write layer.
type Container struct {
	ID       string   `json:"id"`
	Names    []string `json:"names,omitempty"`
	ImageID  string   `json:"image"`
	LayerID  string   `json:"layer"`
	Metadata string   `json:"metadata,omitempty"`
}

// ContainerStore provides bookkeeping for information about Containers.
//
// Create creates a container that has a specified ID (or a random one) and an
// optional name, based on the specified image, using the specified layer as
// its read-write layer.
//
// SetMetadata replaces the metadata associated with a container with the
// supplied value.
//
// Exists checks if there is a container with the given ID or name.
//
// Get retrieves information about a container given an ID or name.
//
// Delete removes the record of the container.
//
// Wipe removes records of all containers.
//
// Lookup attempts to translate a name to an ID.  Most methods do this
// implicitly.
//
// Containers returns a slice enumerating the known containers.
type ContainerStore interface {
	Store
	Create(id string, names []string, image, layer, metadata string) (*Container, error)
	SetMetadata(id, metadata string) error
	SetNames(id string, names []string) error
	Get(id string) (*Container, error)
	Exists(id string) bool
	Delete(id string) error
	Wipe() error
	Lookup(name string) (string, error)
	Containers() ([]Container, error)
}

type containerStore struct {
	lockfile   Locker
	dir        string
	containers []Container
	byid       map[string]*Container
	bylayer    map[string]*Container
	byname     map[string]*Container
}

func (r *containerStore) Containers() ([]Container, error) {
	return r.containers, nil
}

func (r *containerStore) Load() error {
	rpath := filepath.Join(r.dir, "containers.json")
	data, err := ioutil.ReadFile(rpath)
	if err != nil && !os.IsNotExist(err) {
		return err
	}
	containers := []Container{}
	layers := make(map[string]*Container)
	ids := make(map[string]*Container)
	names := make(map[string]*Container)
	if err = json.Unmarshal(data, &containers); len(data) == 0 || err == nil {
		for n, container := range containers {
			ids[container.ID] = &containers[n]
			layers[container.LayerID] = &containers[n]
			for _, name := range container.Names {
				names[name] = &containers[n]
			}
		}
	}
	r.containers = containers
	r.byid = ids
	r.bylayer = layers
	r.byname = names
	return nil
}

func (r *containerStore) Save() error {
	rpath := filepath.Join(r.dir, "containers.json")
	jdata, err := json.Marshal(&r.containers)
	if err != nil {
		return err
	}
	return ioutils.AtomicWriteFile(rpath, jdata, 0600)
}

func newContainerStore(dir string) (ContainerStore, error) {
	if err := os.MkdirAll(dir, 0700); err != nil {
		return nil, err
	}
	lockfile, err := GetLockfile(filepath.Join(dir, "containers.lock"))
	if err != nil {
		return nil, err
	}
	lockfile.Lock()
	defer lockfile.Unlock()
	cstore := containerStore{
		lockfile:   lockfile,
		dir:        dir,
		containers: []Container{},
		byid:       make(map[string]*Container),
		bylayer:    make(map[string]*Container),
		byname:     make(map[string]*Container),
	}
	if err := cstore.Load(); err != nil {
		return nil, err
	}
	return &cstore, nil
}

func (r *containerStore) Create(id string, names []string, image, layer, metadata string) (container *Container, err error) {
	if id == "" {
		id = stringid.GenerateRandomID()
	}
	for _, name := range names {
		if _, nameInUse := r.byname[name]; nameInUse {
			return nil, ErrDuplicateName
		}
	}
	if err == nil {
		newContainer := Container{
			ID:       id,
			Names:    names,
			ImageID:  image,
			LayerID:  layer,
			Metadata: metadata,
		}
		r.containers = append(r.containers, newContainer)
		container = &r.containers[len(r.containers)-1]
		r.byid[id] = container
		r.bylayer[layer] = container
		for _, name := range names {
			r.byname[name] = container
		}
		err = r.Save()
	}
	return container, err
}

func (r *containerStore) SetMetadata(id, metadata string) error {
	if container, ok := r.byname[id]; ok {
		id = container.ID
	} else if container, ok := r.bylayer[id]; ok {
		id = container.ID
	}
	if container, ok := r.byid[id]; ok {
		container.Metadata = metadata
		return r.Save()
	}
	return ErrContainerUnknown
}

func (r *containerStore) SetNames(id string, names []string) error {
	if container, ok := r.byname[id]; ok {
		id = container.ID
	} else if container, ok := r.bylayer[id]; ok {
		id = container.ID
	}
	if container, ok := r.byid[id]; ok {
		for _, name := range container.Names {
			delete(r.byname, name)
		}
		for _, name := range names {
			r.byname[name] = container
		}
		container.Names = names
		return r.Save()
	}
	return ErrContainerUnknown
}

func (r *containerStore) Delete(id string) error {
	if container, ok := r.byname[id]; ok {
		id = container.ID
	} else if container, ok := r.bylayer[id]; ok {
		id = container.ID
	}
	if container, ok := r.byid[id]; ok {
		newContainers := []Container{}
		for _, candidate := range r.containers {
			if candidate.ID != id {
				newContainers = append(newContainers, candidate)
			}
		}
		for _, name := range container.Names {
			delete(r.byname, name)
		}
		r.containers = newContainers
		if err := r.Save(); err != nil {
			return err
		}
	}
	return nil
}

func (r *containerStore) Get(id string) (*Container, error) {
	if c, ok := r.byname[id]; ok {
		return c, nil
	} else if c, ok := r.bylayer[id]; ok {
		return c, nil
	}
	if c, ok := r.byid[id]; ok {
		return c, nil
	}
	return nil, ErrContainerUnknown
}

func (r *containerStore) Lookup(name string) (id string, err error) {
	container, ok := r.byname[name]
	if !ok {
		return "", ErrContainerUnknown
	}
	return container.ID, nil
}

func (r *containerStore) Exists(id string) bool {
	if _, ok := r.byname[id]; ok {
		return true
	}
	if _, ok := r.bylayer[id]; ok {
		return true
	}
	if _, ok := r.byid[id]; ok {
		return true
	}
	return false
}

func (r *containerStore) Wipe() error {
	ids := []string{}
	for id := range r.byid {
		ids = append(ids, id)
	}
	for _, id := range ids {
		if err := r.Delete(id); err != nil {
			return err
		}
	}
	return nil
}

func (r *containerStore) Lock() {
	r.lockfile.Lock()
}

func (r *containerStore) Unlock() {
	r.lockfile.Unlock()
}

func (r *containerStore) Touch() error {
	return r.lockfile.Touch()
}

func (r *containerStore) Modified() (bool, error) {
	return r.lockfile.Modified()
}
<|endoftext|>


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

// Package test provides common storage testing code.
package test

import (
	"time"

	"github.com/Percona-Lab/PromHouse/prompb"
	"github.com/prometheus/common/model"
)

func GetData() *prompb.WriteRequest {
	start := model.Now().Add(-6 * time.Second)

	return &prompb.WriteRequest{
		TimeSeries: []*prompb.TimeSeries{
			{
				Labels: []*prompb.Label{
					{Name: "__name__", Value: "http_requests_total"},
					{Name: "code", Value: "200"},
					{Name: "handler", Value: "query"},
				},
				Samples: []*prompb.Sample{
					{Value: 13, TimestampMs: int64(start)},
					{Value: 14, TimestampMs: int64(start.Add(1 * time.Second))},
					{Value: 14, TimestampMs: int64(start.Add(2 * time.Second))},
					{Value: 14, TimestampMs: int64(start.Add(3 * time.Second))},
					{Value: 15, TimestampMs: int64(start.Add(4 * time.Second))},
				},
			},
			{
				Labels: []*prompb.Label{
					{Name: "__name__", Value: "http_requests_total"},
					{Name: "code", Value: "400"},
					{Name: "handler", Value: "query_range"},
				},
				Samples: []*prompb.Sample{
					{Value: 9, TimestampMs: int64(start)},
					{Value: 9, TimestampMs: int64(start.Add(1 * time.Second))},
					{Value: 9, TimestampMs: int64(start.Add(2 * time.Second))},
					{Value: 11, TimestampMs: int64(start.Add(3 * time.Second))},
					{Value: 11, TimestampMs: int64(start.Add(4 * time.Second))},
				},
			},
			{
				Labels: []*prompb.Label{
					{Name: "__name__", Value: "http_requests_total"},
					{Name: "code", Value: "200"},
					{Name: "handler", Value: "prometheus"},
				},
				Samples: []*prompb.Sample{
					{Value: 591, TimestampMs: int64(start)},
					{Value: 592, TimestampMs: int64(start.Add(1 * time.Second))},
					{Value: 593, TimestampMs: int64(start.Add(2 * time.Second))},
					{Value: 594, TimestampMs: int64(start.Add(3 * time.Second))},
					{Value: 595, TimestampMs: int64(start.Add(4 * time.Second))},
				},
			},
		},
	}
}

func MakeMetric(labels []*prompb.Label) model.Metric {
	metric := make(model.Metric, len(labels))
	for _, l := range labels {
		metric[model.LabelName(l.Name)] = model.LabelValue(l.Value)
	}
	return metric
}
<commit_msg>Group imports.<commit_after>// PromHouse
// Copyright (C) 2017 Percona LLC
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

// Package test provides common storage testing code.
package test

import (
	"time"

	"github.com/prometheus/common/model"

	"github.com/Percona-Lab/PromHouse/prompb"
)

func GetData() *prompb.WriteRequest {
	start := model.Now().Add(-6 * time.Second)

	return &prompb.WriteRequest{
		TimeSeries: []*prompb.TimeSeries{
			{
				Labels: []*prompb.Label{
					{Name: "__name__", Value: "http_requests_total"},
					{Name: "code", Value: "200"},
					{Name: "handler", Value: "query"},
				},
				Samples: []*prompb.Sample{
					{Value: 13, TimestampMs: int64(start)},
					{Value: 14, TimestampMs: int64(start.Add(1 * time.Second))},
					{Value: 14, TimestampMs: int64(start.Add(2 * time.Second))},
					{Value: 14, TimestampMs: int64(start.Add(3 * time.Second))},
					{Value: 15, TimestampMs: int64(start.Add(4 * time.Second))},
				},
			},
			{
				Labels: []*prompb.Label{
					{Name: "__name__", Value: "http_requests_total"},
					{Name: "code", Value: "400"},
					{Name: "handler", Value: "query_range"},
				},
				Samples: []*prompb.Sample{
					{Value: 9, TimestampMs: int64(start)},
					{Value: 9, TimestampMs: int64(start.Add(1 * time.Second))},
					{Value: 9, TimestampMs: int64(start.Add(2 * time.Second))},
					{Value: 11, TimestampMs: int64(start.Add(3 * time.Second))},
					{Value: 11, TimestampMs: int64(start.Add(4 * time.Second))},
				},
			},
			{
				Labels: []*prompb.Label{
					{Name: "__name__", Value: "http_requests_total"},
					{Name: "code", Value: "200"},
					{Name: "handler", Value: "prometheus"},
				},
				Samples: []*prompb.Sample{
					{Value: 591, TimestampMs: int64(start)},
					{Value: 592, TimestampMs: int64(start.Add(1 * time.Second))},
					{Value: 593, TimestampMs: int64(start.Add(2 * time.Second))},
					{Value: 594, TimestampMs: int64(start.Add(3 * time.Second))},
					{Value: 595, TimestampMs: int64(start.Add(4 * time.Second))},
				},
			},
		},
	}
}

func MakeMetric(labels []*prompb.Label) model.Metric {
	metric := make(model.Metric, len(labels))
	for _, l := range labels {
		metric[model.LabelName(l.Name)] = model.LabelValue(l.Value)
	}
	return metric
}
<|endoftext|>


<commit_before>package wallet

import (
	"fmt"

	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"
)

const (
	numInitialKeys = 1e6
	maxScanKeys    = 100e6
)

var errMaxKeys = fmt.Errorf("refused to generate more than %v keys from seed", maxScanKeys)

type scannedSiacoinOutput struct {
	id        types.SiacoinOutputID
	value     types.Currency
	seedIndex uint64
}

// A seedScanner scans the blockchain for addresses that belong to a given
// seed.
type seedScanner struct {
	seed             modules.Seed
	keys             map[types.UnlockHash]uint64 // map address to seed index
	siacoinOutputs   map[types.SiacoinOutputID]scannedSiacoinOutput
	largestIndexSeen uint64 // largest index that has appeared in the blockchain
}

func (s *seedScanner) isSeedAddress(uh types.UnlockHash) bool {
	_, exists := s.keys[uh]
	return exists
}

// generateKeys generates n additional keys from the seedScanner's seed.
func (s *seedScanner) generateKeys(n uint64) {
	initialProgress := uint64(len(s.keys))
	for i, k := range generateKeys(s.seed, initialProgress, n) {
		s.keys[k.UnlockConditions.UnlockHash()] = initialProgress + uint64(i)
	}
}

func (s *seedScanner) ProcessConsensusChange(cc modules.ConsensusChange) {
	// update outputs
	for _, block := range cc.AppliedBlocks {
		for i, mp := range block.MinerPayouts {
			// if a seed miner output is found, add it to the output set
			// TODO: these outputs can only be spent after a maturity delay
			if index, exists := s.keys[mp.UnlockHash]; exists {
				id := types.SiacoinOutputID(block.MinerPayoutID(uint64(i)))
				s.siacoinOutputs[id] = scannedSiacoinOutput{
					id:        id,
					value:     mp.Value,
					seedIndex: index,
				}
			}
		}
		for _, txn := range block.Transactions {
			for i, sco := range txn.SiacoinOutputs {
				// if a seed output is found, add it to the output set
				if index, exists := s.keys[sco.UnlockHash]; exists {
					id := txn.SiacoinOutputID(uint64(i))
					s.siacoinOutputs[id] = scannedSiacoinOutput{
						id:        id,
						value:     sco.Value,
						seedIndex: index,
					}
				}
			}
			for _, sci := range txn.SiacoinInputs {
				// if a seed output is spent, remove it from the output set
				if _, exists := s.siacoinOutputs[sci.ParentID]; exists {
					delete(s.siacoinOutputs, sci.ParentID)
				}
			}
		}
	}

	// update largestIndexSeen
	var addrs []types.UnlockHash
	for _, diff := range cc.SiacoinOutputDiffs {
		addrs = append(addrs, diff.SiacoinOutput.UnlockHash)
	}
	for _, diff := range cc.SiafundOutputDiffs {
		addrs = append(addrs, diff.SiafundOutput.UnlockHash)
	}
	for _, block := range cc.AppliedBlocks {
		for _, mp := range block.MinerPayouts {
			addrs = append(addrs, mp.UnlockHash)
		}
		for _, txn := range block.Transactions {
			for _, sci := range txn.SiacoinInputs {
				addrs = append(addrs, sci.UnlockConditions.UnlockHash())
			}
			for _, sco := range txn.SiacoinOutputs {
				addrs = append(addrs, sco.UnlockHash)
			}
			for _, sfi := range txn.SiafundInputs {
				addrs = append(addrs, sfi.UnlockConditions.UnlockHash())
			}
			for _, sfo := range txn.SiafundOutputs {
				addrs = append(addrs, sfo.UnlockHash)
			}
		}
	}
	for _, addr := range addrs {
		index, exists := s.keys[addr]
		if exists && index > s.largestIndexSeen {
			s.largestIndexSeen = index
		}
	}
}

// scan subscribes d to cs and scans the blockchain for addresses that belong
// to d's seed. If scan returns errMaximumKeys, additional keys may need to be
// generated to find all the addresses.
func (s *seedScanner) scan(cs modules.ConsensusSet) error {
	// generate a bunch of keys and scan the blockchain looking for them. If
	// none of the keys are found, we are done; otherwise, generate more keys
	// and try again (bounded by a sane default).
	//
	// NOTE: since scanning is very slow, we aim to only scan once, which
	// means generating many keys.
	var numKeys uint64 = numInitialKeys
	for len(s.keys) < maxScanKeys {
		s.generateKeys(numKeys)
		if err := cs.ConsensusSetSubscribe(s, modules.ConsensusChangeBeginning); err != nil {
			return err
		}
		if s.largestIndexSeen < uint64(len(s.keys))/2 {
			cs.Unsubscribe(s)
			return nil
		}
		// double number of keys generated each iteration, capping so that we
		// do not exceed maxScanKeys
		numKeys *= 2
		if numKeys > maxScanKeys-uint64(len(s.keys)) {
			numKeys = maxScanKeys - uint64(len(s.keys))
		}
	}
	cs.Unsubscribe(s)
	return errMaxKeys
}

func newSeedScanner(seed modules.Seed) *seedScanner {
	return &seedScanner{
		seed:           seed,
		keys:           make(map[types.UnlockHash]uint64),
		siacoinOutputs: make(map[types.SiacoinOutputID]scannedSiacoinOutput),
	}
}
<commit_msg>fix sweeping of miner outputs<commit_after>package wallet

import (
	"fmt"

	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"
)

const (
	numInitialKeys = 1e6
	maxScanKeys    = 100e6
)

var errMaxKeys = fmt.Errorf("refused to generate more than %v keys from seed", maxScanKeys)

type scannedSiacoinOutput struct {
	id        types.SiacoinOutputID
	value     types.Currency
	seedIndex uint64
}

// A seedScanner scans the blockchain for addresses that belong to a given
// seed.
type seedScanner struct {
	seed             modules.Seed
	keys             map[types.UnlockHash]uint64 // map address to seed index
	siacoinOutputs   map[types.SiacoinOutputID]scannedSiacoinOutput
	minerOutputs     map[types.BlockHeight][]scannedSiacoinOutput
	largestIndexSeen uint64 // largest index that has appeared in the blockchain
	blockheight      types.BlockHeight
}

func (s *seedScanner) isSeedAddress(uh types.UnlockHash) bool {
	_, exists := s.keys[uh]
	return exists
}

// generateKeys generates n additional keys from the seedScanner's seed.
func (s *seedScanner) generateKeys(n uint64) {
	initialProgress := uint64(len(s.keys))
	for i, k := range generateKeys(s.seed, initialProgress, n) {
		s.keys[k.UnlockConditions.UnlockHash()] = initialProgress + uint64(i)
	}
}

func (s *seedScanner) ProcessConsensusChange(cc modules.ConsensusChange) {
	s.blockheight -= types.BlockHeight(len(cc.RevertedBlocks))
	// update outputs
	for _, block := range cc.AppliedBlocks {
		s.blockheight++
		// when a miner payout matures, move it from minerOutputs to
		// siacoinOutputs
		if matured, exists := s.minerOutputs[s.blockheight]; exists {
			for _, maturedOutput := range matured {
				s.siacoinOutputs[maturedOutput.id] = maturedOutput
			}
			delete(s.minerOutputs, s.blockheight)
		}
		for i, mp := range block.MinerPayouts {
			// if a seed miner output is found, add it to the delayed output
			// set
			if index, exists := s.keys[mp.UnlockHash]; exists {
				maturityHeight := s.blockheight + types.MaturityDelay
				s.minerOutputs[maturityHeight] = append(s.minerOutputs[maturityHeight], scannedSiacoinOutput{
					id:        types.SiacoinOutputID(block.MinerPayoutID(uint64(i))),
					value:     mp.Value,
					seedIndex: index,
				})
			}
		}
		for _, txn := range block.Transactions {
			for i, sco := range txn.SiacoinOutputs {
				// if a seed output is found, add it to the output set
				if index, exists := s.keys[sco.UnlockHash]; exists {
					id := txn.SiacoinOutputID(uint64(i))
					s.siacoinOutputs[id] = scannedSiacoinOutput{
						id:        id,
						value:     sco.Value,
						seedIndex: index,
					}
				}
			}
			for _, sci := range txn.SiacoinInputs {
				// if a seed output is spent, remove it from the output set
				if _, exists := s.siacoinOutputs[sci.ParentID]; exists {
					delete(s.siacoinOutputs, sci.ParentID)
				}
			}
		}
	}

	// update largestIndexSeen
	var addrs []types.UnlockHash
	for _, diff := range cc.SiacoinOutputDiffs {
		addrs = append(addrs, diff.SiacoinOutput.UnlockHash)
	}
	for _, diff := range cc.SiafundOutputDiffs {
		addrs = append(addrs, diff.SiafundOutput.UnlockHash)
	}
	for _, block := range cc.AppliedBlocks {
		for _, mp := range block.MinerPayouts {
			addrs = append(addrs, mp.UnlockHash)
		}
		for _, txn := range block.Transactions {
			for _, sci := range txn.SiacoinInputs {
				addrs = append(addrs, sci.UnlockConditions.UnlockHash())
			}
			for _, sco := range txn.SiacoinOutputs {
				addrs = append(addrs, sco.UnlockHash)
			}
			for _, sfi := range txn.SiafundInputs {
				addrs = append(addrs, sfi.UnlockConditions.UnlockHash())
			}
			for _, sfo := range txn.SiafundOutputs {
				addrs = append(addrs, sfo.UnlockHash)
			}
		}
	}
	for _, addr := range addrs {
		index, exists := s.keys[addr]
		if exists && index > s.largestIndexSeen {
			s.largestIndexSeen = index
		}
	}
}

// scan subscribes d to cs and scans the blockchain for addresses that belong
// to d's seed. If scan returns errMaximumKeys, additional keys may need to be
// generated to find all the addresses.
func (s *seedScanner) scan(cs modules.ConsensusSet) error {
	// generate a bunch of keys and scan the blockchain looking for them. If
	// none of the keys are found, we are done; otherwise, generate more keys
	// and try again (bounded by a sane default).
	//
	// NOTE: since scanning is very slow, we aim to only scan once, which
	// means generating many keys.
	var numKeys uint64 = numInitialKeys
	for len(s.keys) < maxScanKeys {
		s.generateKeys(numKeys)
		if err := cs.ConsensusSetSubscribe(s, modules.ConsensusChangeBeginning); err != nil {
			return err
		}
		if s.largestIndexSeen < uint64(len(s.keys))/2 {
			cs.Unsubscribe(s)
			return nil
		}
		// double number of keys generated each iteration, capping so that we
		// do not exceed maxScanKeys
		numKeys *= 2
		if numKeys > maxScanKeys-uint64(len(s.keys)) {
			numKeys = maxScanKeys - uint64(len(s.keys))
		}
	}
	cs.Unsubscribe(s)
	return errMaxKeys
}

func newSeedScanner(seed modules.Seed) *seedScanner {
	return &seedScanner{
		seed:           seed,
		keys:           make(map[types.UnlockHash]uint64),
		siacoinOutputs: make(map[types.SiacoinOutputID]scannedSiacoinOutput),
		minerOutputs:   make(map[types.BlockHeight][]scannedSiacoinOutput),
	}
}
<|endoftext|>


<commit_before>package plugin

import (
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"sort"
	"strings"

	"github.com/hashicorp/packer/packer"
	packersdk "github.com/hashicorp/packer/packer-plugin-sdk/packer"
)

// PACKERSPACE is used to represent the spaces that separate args for a command
// without being confused with spaces in the path to the command itself.
const PACKERSPACE = "-PACKERSPACE-"

type Config struct {
	PluginMinPort  int
	PluginMaxPort  int
	builders       packer.MapOfBuilder
	provisioners   packer.MapOfProvisioner
	postProcessors packer.MapOfPostProcessor
}

func (c *Config) GetPlugins() (packer.MapOfBuilder, packer.MapOfProvisioner, packer.MapOfPostProcessor) {
	return c.builders, c.provisioners, c.postProcessors
}

// Discover discovers plugins.
//
// Search the directory of the executable, then the plugins directory, and
// finally the CWD, in that order. Any conflicts will overwrite previously
// found plugins, in that order.
// Hence, the priority order is the reverse of the search order - i.e., the
// CWD has the highest priority.
func (c *Config) Discover() error {
	c.builders = packer.MapOfBuilder{}
	c.postProcessors = packer.MapOfPostProcessor{}
	c.provisioners = packer.MapOfProvisioner{}
	// If we are already inside a plugin process we should not need to
	// discover anything.
	if os.Getenv(MagicCookieKey) == MagicCookieValue {
		return nil
	}

	// Next, look in the same directory as the executable.
	exePath, err := os.Executable()
	if err != nil {
		log.Printf("[ERR] Error loading exe directory: %s", err)
	} else {
		if err := c.discoverExternalComponents(filepath.Dir(exePath)); err != nil {
			return err
		}
	}

	// Next, look in the default plugins directory inside the configdir/.packer.d/plugins.
	dir, err := packer.ConfigDir()
	if err != nil {
		log.Printf("[ERR] Error loading config directory: %s", err)
	} else {
		if err := c.discoverExternalComponents(filepath.Join(dir, "plugins")); err != nil {
			return err
		}
	}

	// Next, look in the CWD.
	if err := c.discoverExternalComponents("."); err != nil {
		return err
	}

	// Check whether there is a custom Plugin directory defined. This gets
	// absolute preference.
	if packerPluginPath := os.Getenv("PACKER_PLUGIN_PATH"); packerPluginPath != "" {
		sep := ":"
		if runtime.GOOS == "windows" {
			// on windows, PATH is semicolon-separated
			sep = ";"
		}
		plugPaths := strings.Split(packerPluginPath, sep)
		for _, plugPath := range plugPaths {
			if err := c.discoverExternalComponents(plugPath); err != nil {
				return err
			}
		}
	}

	return nil
}

func (c *Config) discoverExternalComponents(path string) error {
	var err error

	if !filepath.IsAbs(path) {
		path, err = filepath.Abs(path)
		if err != nil {
			return err
		}
	}
	var externallyUsed []string

	pluginPaths, err := c.discoverSingle(filepath.Join(path, "packer-builder-*"))
	if err != nil {
		return err
	}
	for pluginName, pluginPath := range pluginPaths {
		newPath := pluginPath // this needs to be stored in a new variable for the func below
		c.builders[pluginName] = func() (packersdk.Builder, error) {
			return c.Client(newPath).Builder()
		}
		externallyUsed = append(externallyUsed, pluginName)
	}
	if len(externallyUsed) > 0 {
		sort.Strings(externallyUsed)
		log.Printf("using external builders %v", externallyUsed)
		externallyUsed = nil
	}

	pluginPaths, err = c.discoverSingle(filepath.Join(path, "packer-post-processor-*"))
	if err != nil {
		return err
	}
	for pluginName, pluginPath := range pluginPaths {
		newPath := pluginPath // this needs to be stored in a new variable for the func below
		c.postProcessors[pluginName] = func() (packersdk.PostProcessor, error) {
			return c.Client(newPath).PostProcessor()
		}
		externallyUsed = append(externallyUsed, pluginName)
	}
	if len(externallyUsed) > 0 {
		sort.Strings(externallyUsed)
		log.Printf("using external post-processors %v", externallyUsed)
		externallyUsed = nil
	}

	pluginPaths, err = c.discoverSingle(filepath.Join(path, "packer-provisioner-*"))
	if err != nil {
		return err
	}
	for pluginName, pluginPath := range pluginPaths {
		newPath := pluginPath // this needs to be stored in a new variable for the func below
		c.provisioners[pluginName] = func() (packersdk.Provisioner, error) {
			return c.Client(newPath).Provisioner()
		}
		externallyUsed = append(externallyUsed, pluginName)
	}
	if len(externallyUsed) > 0 {
		sort.Strings(externallyUsed)
		log.Printf("using external provisioners %v", externallyUsed)
	}

	return nil
}

func (c *Config) discoverSingle(glob string) (map[string]string, error) {
	matches, err := filepath.Glob(glob)
	if err != nil {
		return nil, err
	}

	res := make(map[string]string)

	prefix := filepath.Base(glob)
	prefix = prefix[:strings.Index(prefix, "*")]
	for _, match := range matches {
		file := filepath.Base(match)

		// On Windows, ignore any plugins that don't end in .exe.
		// We could do a full PATHEXT parse, but this is probably good enough.
		if runtime.GOOS == "windows" && strings.ToLower(filepath.Ext(file)) != ".exe" {
			log.Printf(
				"[DEBUG] Ignoring plugin match %s, no exe extension",
				match)
			continue
		}

		// If the filename has a ".", trim up to there
		if idx := strings.Index(file, ".exe"); idx >= 0 {
			file = file[:idx]
		}

		// Look for foo-bar-baz. The plugin name is "baz"
		pluginName := file[len(prefix):]
		log.Printf("[DEBUG] Discovered plugin: %s = %s", pluginName, match)
		res[pluginName] = match
	}

	return res, nil
}

func (c *Config) Client(path string) *Client {
	originalPath := path

	// Check for special case using `packer plugin PLUGIN`
	args := []string{}
	if strings.Contains(path, PACKERSPACE) {
		parts := strings.Split(path, PACKERSPACE)
		path = parts[0]
		args = parts[1:]
	}

	// First attempt to find the executable by consulting the PATH.
	path, err := exec.LookPath(path)
	if err != nil {
		// If that doesn't work, look for it in the same directory
		// as the `packer` executable (us).
		log.Printf("Plugin could not be found at %s (%v). Checking same directory as executable.", originalPath, err)
		exePath, err := os.Executable()
		if err != nil {
			log.Printf("Couldn't get current exe path: %s", err)
		} else {
			log.Printf("Current exe path: %s", exePath)
			path = filepath.Join(filepath.Dir(exePath), filepath.Base(originalPath))
		}
	}

	// If everything failed, just use the original path and let the error
	// bubble through.
	if path == "" {
		path = originalPath
	}

	log.Printf("Creating plugin client for path: %s", path)
	var config ClientConfig
	config.Cmd = exec.Command(path, args...)
	config.Managed = true
	config.MinPort = c.PluginMinPort
	config.MaxPort = c.PluginMaxPort
	return NewClient(&config)
}
<commit_msg>don't override builder/prov/pproc<commit_after>package plugin

import (
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"sort"
	"strings"

	"github.com/hashicorp/packer/packer"
	packersdk "github.com/hashicorp/packer/packer-plugin-sdk/packer"
)

// PACKERSPACE is used to represent the spaces that separate args for a command
// without being confused with spaces in the path to the command itself.
const PACKERSPACE = "-PACKERSPACE-"

type Config struct {
	PluginMinPort  int
	PluginMaxPort  int
	builders       packer.MapOfBuilder
	provisioners   packer.MapOfProvisioner
	postProcessors packer.MapOfPostProcessor
}

func (c *Config) GetPlugins() (packer.MapOfBuilder, packer.MapOfProvisioner, packer.MapOfPostProcessor) {
	return c.builders, c.provisioners, c.postProcessors
}

// Discover discovers plugins.
//
// Search the directory of the executable, then the plugins directory, and
// finally the CWD, in that order. Any conflicts will overwrite previously
// found plugins, in that order.
// Hence, the priority order is the reverse of the search order - i.e., the
// CWD has the highest priority.
func (c *Config) Discover() error {
	if c.builders == nil {
		c.builders = packer.MapOfBuilder{}
	}
	if c.provisioners == nil {
		c.provisioners = packer.MapOfProvisioner{}
	}
	if c.postProcessors == nil {
		c.postProcessors = packer.MapOfPostProcessor{}
	}

	// If we are already inside a plugin process we should not need to
	// discover anything.
	if os.Getenv(MagicCookieKey) == MagicCookieValue {
		return nil
	}

	// Next, look in the same directory as the executable.
	exePath, err := os.Executable()
	if err != nil {
		log.Printf("[ERR] Error loading exe directory: %s", err)
	} else {
		if err := c.discoverExternalComponents(filepath.Dir(exePath)); err != nil {
			return err
		}
	}

	// Next, look in the default plugins directory inside the configdir/.packer.d/plugins.
	dir, err := packer.ConfigDir()
	if err != nil {
		log.Printf("[ERR] Error loading config directory: %s", err)
	} else {
		if err := c.discoverExternalComponents(filepath.Join(dir, "plugins")); err != nil {
			return err
		}
	}

	// Next, look in the CWD.
	if err := c.discoverExternalComponents("."); err != nil {
		return err
	}

	// Check whether there is a custom Plugin directory defined. This gets
	// absolute preference.
	if packerPluginPath := os.Getenv("PACKER_PLUGIN_PATH"); packerPluginPath != "" {
		sep := ":"
		if runtime.GOOS == "windows" {
			// on windows, PATH is semicolon-separated
			sep = ";"
		}
		plugPaths := strings.Split(packerPluginPath, sep)
		for _, plugPath := range plugPaths {
			if err := c.discoverExternalComponents(plugPath); err != nil {
				return err
			}
		}
	}

	return nil
}

func (c *Config) discoverExternalComponents(path string) error {
	var err error

	if !filepath.IsAbs(path) {
		path, err = filepath.Abs(path)
		if err != nil {
			return err
		}
	}
	var externallyUsed []string

	pluginPaths, err := c.discoverSingle(filepath.Join(path, "packer-builder-*"))
	if err != nil {
		return err
	}
	for pluginName, pluginPath := range pluginPaths {
		newPath := pluginPath // this needs to be stored in a new variable for the func below
		c.builders[pluginName] = func() (packersdk.Builder, error) {
			return c.Client(newPath).Builder()
		}
		externallyUsed = append(externallyUsed, pluginName)
	}
	if len(externallyUsed) > 0 {
		sort.Strings(externallyUsed)
		log.Printf("using external builders %v", externallyUsed)
		externallyUsed = nil
	}

	pluginPaths, err = c.discoverSingle(filepath.Join(path, "packer-post-processor-*"))
	if err != nil {
		return err
	}
	for pluginName, pluginPath := range pluginPaths {
		newPath := pluginPath // this needs to be stored in a new variable for the func below
		c.postProcessors[pluginName] = func() (packersdk.PostProcessor, error) {
			return c.Client(newPath).PostProcessor()
		}
		externallyUsed = append(externallyUsed, pluginName)
	}
	if len(externallyUsed) > 0 {
		sort.Strings(externallyUsed)
		log.Printf("using external post-processors %v", externallyUsed)
		externallyUsed = nil
	}

	pluginPaths, err = c.discoverSingle(filepath.Join(path, "packer-provisioner-*"))
	if err != nil {
		return err
	}
	for pluginName, pluginPath := range pluginPaths {
		newPath := pluginPath // this needs to be stored in a new variable for the func below
		c.provisioners[pluginName] = func() (packersdk.Provisioner, error) {
			return c.Client(newPath).Provisioner()
		}
		externallyUsed = append(externallyUsed, pluginName)
	}
	if len(externallyUsed) > 0 {
		sort.Strings(externallyUsed)
		log.Printf("using external provisioners %v", externallyUsed)
	}

	return nil
}

func (c *Config) discoverSingle(glob string) (map[string]string, error) {
	matches, err := filepath.Glob(glob)
	if err != nil {
		return nil, err
	}

	res := make(map[string]string)

	prefix := filepath.Base(glob)
	prefix = prefix[:strings.Index(prefix, "*")]
	for _, match := range matches {
		file := filepath.Base(match)

		// On Windows, ignore any plugins that don't end in .exe.
		// We could do a full PATHEXT parse, but this is probably good enough.
		if runtime.GOOS == "windows" && strings.ToLower(filepath.Ext(file)) != ".exe" {
			log.Printf(
				"[DEBUG] Ignoring plugin match %s, no exe extension",
				match)
			continue
		}

		// If the filename has a ".", trim up to there
		if idx := strings.Index(file, ".exe"); idx >= 0 {
			file = file[:idx]
		}

		// Look for foo-bar-baz. The plugin name is "baz"
		pluginName := file[len(prefix):]
		log.Printf("[DEBUG] Discovered plugin: %s = %s", pluginName, match)
		res[pluginName] = match
	}

	return res, nil
}

func (c *Config) Client(path string) *Client {
	originalPath := path

	// Check for special case using `packer plugin PLUGIN`
	args := []string{}
	if strings.Contains(path, PACKERSPACE) {
		parts := strings.Split(path, PACKERSPACE)
		path = parts[0]
		args = parts[1:]
	}

	// First attempt to find the executable by consulting the PATH.
	path, err := exec.LookPath(path)
	if err != nil {
		// If that doesn't work, look for it in the same directory
		// as the `packer` executable (us).
		log.Printf("Plugin could not be found at %s (%v). Checking same directory as executable.", originalPath, err)
		exePath, err := os.Executable()
		if err != nil {
			log.Printf("Couldn't get current exe path: %s", err)
		} else {
			log.Printf("Current exe path: %s", exePath)
			path = filepath.Join(filepath.Dir(exePath), filepath.Base(originalPath))
		}
	}

	// If everything failed, just use the original path and let the error
	// bubble through.
	if path == "" {
		path = originalPath
	}

	log.Printf("Creating plugin client for path: %s", path)
	var config ClientConfig
	config.Cmd = exec.Command(path, args...)
	config.Managed = true
	config.MinPort = c.PluginMinPort
	config.MaxPort = c.PluginMaxPort
	return NewClient(&config)
}
<|endoftext|>


<commit_before>package thankyouplugin

import (
	"strings"

	"github.com/iopred/bruxism"
)

// ThankyouCommand returns bot statistics.
func ThankyouCommand(bot *bruxism.Bot, service bruxism.Service, message bruxism.Message, command string, parts []string) {
	out := "Septapus currently has no patrons.\nBecome a patron at: <https://www.patreon.com/iopred>"

	if service.SupportsMultiline() {
		service.SendMessage(message.Channel(), out)
	} else {
		service.SendMessage(message.Channel(), strings.Join(strings.Split(out, "\n"), ", "))
	}
}

// StatsHelp is the help for the stats command.
var StatsHelp = bruxism.NewCommandHelp("", "Thank you to all my patrons.")
<commit_msg>:ok_hand:<commit_after>package thankyouplugin

import (
	"strings"

	"github.com/iopred/bruxism"
)

// ThankyouCommand returns bot statistics.
func ThankyouCommand(bot *bruxism.Bot, service bruxism.Service, message bruxism.Message, command string, parts []string) {
	out := "Septapus currently has no patrons.\nBecome a patron at: <https://www.patreon.com/iopred>"

	if service.SupportsMultiline() {
		service.SendMessage(message.Channel(), out)
	} else {
		service.SendMessage(message.Channel(), strings.Join(strings.Split(out, "\n"), ", "))
	}
}

// ThankyouHelp is the help for the stats command.
var ThankyouHelp = bruxism.NewCommandHelp("", "Thank you to all my patrons.")
<|endoftext|>


<commit_before>package ec2
import (
	"bytes"
	"io"
	"io/ioutil"
	"fmt"
	"launchpad.net/goamz/s3"
	"launchpad.net/goyaml"
)

const stateFile = "provider-state"

type bootstrapState struct {
	ZookeeperInstances []string	`yaml:"zookeeper-instances"`
}

func (e *environ) saveState(state *bootstrapState) error {
	data, err := goyaml.Marshal(state)
	if err != nil {
		return err
	}
	return e.PutFile(stateFile, bytes.NewBuffer(data), int64(len(data)))
}

func (e *environ) loadState() (*bootstrapState, error) {
	r, err := e.GetFile(stateFile)
	if err != nil {
		return nil, fmt.Errorf("cannot read %q: %v", stateFile, err)
	}
	defer r.Close()
	data, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, fmt.Errorf("error reading %q: %v", stateFile, err)
	}
	var state bootstrapState
	err = goyaml.Unmarshal(data, &state)
	if err != nil {
		return nil, fmt.Errorf("error unmarshalling %q: %v", stateFile, err)
	}
	return &state, nil
}

func (e *environ) deleteState() error {
	// TODO delete the bucket contents and the bucket itself.
	err := e.bucket().Del(stateFile)
	// If we can't delete the object because the bucket doesn't
	// exist, then we don't care.
	if err, _ := err.(*s3.Error); err != nil && err.StatusCode == 404 {
		return nil
	}
	if err != nil {
		return fmt.Errorf("cannot delete provider state: %v", err)
	}
	return nil
}

// makeBucket makes the environent's control bucket, the
// place where bootstrap information and deployed charms
// are stored. To avoid two round trips on every PUT operation,
// we do this only once for each environ.
func (e *environ) makeBucket() error {
	e.checkBucket.Do(func() {
		// try to make the bucket - PutBucket will succeed if the
		// bucket already exists.
		e.checkBucketError = e.bucket().PutBucket(s3.Private)
	})
	return e.checkBucketError
}

func (e *environ) PutFile(file string, r io.Reader, length int64) error {
	if err := e.makeBucket(); err != nil {
		return fmt.Errorf("cannot make S3 control bucket: %v", err)
	}
	err := e.bucket().PutReader(file, r, length, "binary/octet-stream", s3.Private)
	if err != nil {
		return fmt.Errorf("cannot write file %q to control bucket: %v", file, err)
	}
	return nil
}

func (e *environ) GetFile(file string) (io.ReadCloser, error) {
	return e.bucket().GetReader(file)
}

func (e *environ) RemoveFile(file string) error {
	return e.bucket().Del(file)
}

func (e *environ) bucket() *s3.Bucket {
	return e.s3.Bucket(e.config.bucket)
}
<commit_msg>go fmt<commit_after>package ec2

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"launchpad.net/goamz/s3"
	"launchpad.net/goyaml"
)

const stateFile = "provider-state"

type bootstrapState struct {
	ZookeeperInstances []string `yaml:"zookeeper-instances"`
}

func (e *environ) saveState(state *bootstrapState) error {
	data, err := goyaml.Marshal(state)
	if err != nil {
		return err
	}
	return e.PutFile(stateFile, bytes.NewBuffer(data), int64(len(data)))
}

func (e *environ) loadState() (*bootstrapState, error) {
	r, err := e.GetFile(stateFile)
	if err != nil {
		return nil, fmt.Errorf("cannot read %q: %v", stateFile, err)
	}
	defer r.Close()
	data, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, fmt.Errorf("error reading %q: %v", stateFile, err)
	}
	var state bootstrapState
	err = goyaml.Unmarshal(data, &state)
	if err != nil {
		return nil, fmt.Errorf("error unmarshalling %q: %v", stateFile, err)
	}
	return &state, nil
}

func (e *environ) deleteState() error {
	// TODO delete the bucket contents and the bucket itself.
	err := e.bucket().Del(stateFile)
	// If we can't delete the object because the bucket doesn't
	// exist, then we don't care.
	if err, _ := err.(*s3.Error); err != nil && err.StatusCode == 404 {
		return nil
	}
	if err != nil {
		return fmt.Errorf("cannot delete provider state: %v", err)
	}
	return nil
}

// makeBucket makes the environent's control bucket, the
// place where bootstrap information and deployed charms
// are stored. To avoid two round trips on every PUT operation,
// we do this only once for each environ.
func (e *environ) makeBucket() error {
	e.checkBucket.Do(func() {
		// try to make the bucket - PutBucket will succeed if the
		// bucket already exists.
		e.checkBucketError = e.bucket().PutBucket(s3.Private)
	})
	return e.checkBucketError
}

func (e *environ) PutFile(file string, r io.Reader, length int64) error {
	if err := e.makeBucket(); err != nil {
		return fmt.Errorf("cannot make S3 control bucket: %v", err)
	}
	err := e.bucket().PutReader(file, r, length, "binary/octet-stream", s3.Private)
	if err != nil {
		return fmt.Errorf("cannot write file %q to control bucket: %v", file, err)
	}
	return nil
}

func (e *environ) GetFile(file string) (io.ReadCloser, error) {
	return e.bucket().GetReader(file)
}

func (e *environ) RemoveFile(file string) error {
	return e.bucket().Del(file)
}

func (e *environ) bucket() *s3.Bucket {
	return e.s3.Bucket(e.config.bucket)
}
<|endoftext|>