<commit_before>package quic

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"errors"
	"fmt"
	"hash"
	"log"
	"net"
	"sync"
	"time"

	"github.com/lucas-clemente/quic-go/internal/protocol"
	"github.com/lucas-clemente/quic-go/internal/utils"
	"github.com/lucas-clemente/quic-go/internal/wire"
	"github.com/lucas-clemente/quic-go/logging"
)

type statelessResetErr struct {
	token protocol.StatelessResetToken
}

func (e statelessResetErr) Error() string {
	return fmt.Sprintf("received a stateless reset with token %x", e.token)
}

// The packetHandlerMap stores packetHandlers, identified by connection ID.
// It is used:
// * by the server to store sessions
// * when multiplexing outgoing connections to store clients
type packetHandlerMap struct {
	mutex sync.Mutex

	conn      connection
	connIDLen int

	handlers    map[string] /* string(ConnectionID)*/ packetHandler
	resetTokens map[protocol.StatelessResetToken] /* stateless reset token */ packetHandler
	server      unknownPacketHandler

	listening chan struct{} // is closed when listen returns
	closed    bool

	deleteRetiredSessionsAfter time.Duration

	statelessResetEnabled bool
	statelessResetMutex   sync.Mutex
	statelessResetHasher  hash.Hash

	tracer logging.Tracer
	logger utils.Logger
}

var _ packetHandlerManager = &packetHandlerMap{}

func setReceiveBuffer(c net.PacketConn, logger utils.Logger) error {
	conn, ok := c.(interface{ SetReadBuffer(int) error })
	if !ok {
		return errors.New("connection doesn't allow setting of receive buffer size")
	}
	size, err := inspectReadBuffer(c)
	if err != nil {
		return fmt.Errorf("failed to determine receive buffer size: %w", err)
	}
	if size >= protocol.DesiredReceiveBufferSize {
		logger.Debugf("Conn has receive buffer of %d kiB (wanted: at least %d kiB)", size/1024, protocol.DesiredReceiveBufferSize/1024)
	}
	if err := conn.SetReadBuffer(protocol.DesiredReceiveBufferSize); err != nil {
		return fmt.Errorf("failed to increase receive buffer size: %w", err)
	}
	newSize, err := inspectReadBuffer(c)
	if err != nil {
		return fmt.Errorf("failed to determine receive buffer size: %w", err)
	}
	if newSize == size {
		return fmt.Errorf("failed to determine receive buffer size: %w", err)
	}
	if newSize < protocol.DesiredReceiveBufferSize {
		return fmt.Errorf("failed to sufficiently increase receive buffer size (was: %d kiB, wanted: %d kiB, got: %d kiB)", size/1024, protocol.DesiredReceiveBufferSize/1024, newSize/1024)
	}
	logger.Debugf("Increased receive buffer size to %d kiB", newSize/1024)
	return nil
}

// only print warnings about the UPD receive buffer size once
var receiveBufferWarningOnce sync.Once

func newPacketHandlerMap(
	c net.PacketConn,
	connIDLen int,
	statelessResetKey []byte,
	tracer logging.Tracer,
	logger utils.Logger,
) (packetHandlerManager, error) {
	if err := setReceiveBuffer(c, logger); err != nil {
		receiveBufferWarningOnce.Do(func() {
			log.Printf("%s. See https://github.com/lucas-clemente/quic-go/wiki/UDP-Receive-Buffer-Size for details.", err)
		})
	}
	conn, err := wrapConn(c)
	if err != nil {
		return nil, err
	}
	m := &packetHandlerMap{
		conn:                       conn,
		connIDLen:                  connIDLen,
		listening:                  make(chan struct{}),
		handlers:                   make(map[string]packetHandler),
		resetTokens:                make(map[protocol.StatelessResetToken]packetHandler),
		deleteRetiredSessionsAfter: protocol.RetiredConnectionIDDeleteTimeout,
		statelessResetEnabled:      len(statelessResetKey) > 0,
		statelessResetHasher:       hmac.New(sha256.New, statelessResetKey),
		tracer:                     tracer,
		logger:                     logger,
	}
	go m.listen()

	if logger.Debug() {
		go m.logUsage()
	}
	return m, nil
}

func (h *packetHandlerMap) logUsage() {
	ticker := time.NewTicker(2 * time.Second)
	var printedZero bool
	for {
		select {
		case <-h.listening:
			return
		case <-ticker.C:
		}

		h.mutex.Lock()
		numHandlers := len(h.handlers)
		numTokens := len(h.resetTokens)
		h.mutex.Unlock()
		// If the number tracked handlers and tokens is zero, only print it a single time.
		hasZero := numHandlers == 0 && numTokens == 0
		if !hasZero || (hasZero && !printedZero) {
			h.logger.Debugf("Tracking %d connection IDs and %d reset tokens.\n", numHandlers, numTokens)
			printedZero = false
			if hasZero {
				printedZero = true
			}
		}
	}
}

func (h *packetHandlerMap) Add(id protocol.ConnectionID, handler packetHandler) bool /* was added */ {
	h.mutex.Lock()
	defer h.mutex.Unlock()

	if _, ok := h.handlers[string(id)]; ok {
		h.logger.Debugf("Not adding connection ID %s, as it already exists.", id)
		return false
	}
	h.handlers[string(id)] = handler
	h.logger.Debugf("Adding connection ID %s.", id)
	return true
}

func (h *packetHandlerMap) AddWithConnID(clientDestConnID, newConnID protocol.ConnectionID, fn func() packetHandler) bool {
	h.mutex.Lock()
	defer h.mutex.Unlock()

	if _, ok := h.handlers[string(clientDestConnID)]; ok {
		h.logger.Debugf("Not adding connection ID %s for a new session, as it already exists.", clientDestConnID)
		return false
	}

	sess := fn()
	h.handlers[string(clientDestConnID)] = sess
	h.handlers[string(newConnID)] = sess
	h.logger.Debugf("Adding connection IDs %s and %s for a new session.", clientDestConnID, newConnID)
	return true
}

func (h *packetHandlerMap) Remove(id protocol.ConnectionID) {
	h.mutex.Lock()
	delete(h.handlers, string(id))
	h.mutex.Unlock()
	h.logger.Debugf("Removing connection ID %s.", id)
}

func (h *packetHandlerMap) Retire(id protocol.ConnectionID) {
	h.logger.Debugf("Retiring connection ID %s in %s.", id, h.deleteRetiredSessionsAfter)
	time.AfterFunc(h.deleteRetiredSessionsAfter, func() {
		h.mutex.Lock()
		delete(h.handlers, string(id))
		h.mutex.Unlock()
		h.logger.Debugf("Removing connection ID %s after it has been retired.", id)
	})
}

func (h *packetHandlerMap) ReplaceWithClosed(id protocol.ConnectionID, handler packetHandler) {
	h.mutex.Lock()
	h.handlers[string(id)] = handler
	h.mutex.Unlock()
	h.logger.Debugf("Replacing session for connection ID %s with a closed session.", id)

	time.AfterFunc(h.deleteRetiredSessionsAfter, func() {
		h.mutex.Lock()
		handler.shutdown()
		delete(h.handlers, string(id))
		h.mutex.Unlock()
		h.logger.Debugf("Removing connection ID %s for a closed session after it has been retired.", id)
	})
}

func (h *packetHandlerMap) AddResetToken(token protocol.StatelessResetToken, handler packetHandler) {
	h.mutex.Lock()
	h.resetTokens[token] = handler
	h.mutex.Unlock()
}

func (h *packetHandlerMap) RemoveResetToken(token protocol.StatelessResetToken) {
	h.mutex.Lock()
	delete(h.resetTokens, token)
	h.mutex.Unlock()
}

func (h *packetHandlerMap) SetServer(s unknownPacketHandler) {
	h.mutex.Lock()
	h.server = s
	h.mutex.Unlock()
}

func (h *packetHandlerMap) CloseServer() {
	h.mutex.Lock()
	if h.server == nil {
		h.mutex.Unlock()
		return
	}
	h.server = nil
	var wg sync.WaitGroup
	for _, handler := range h.handlers {
		if handler.getPerspective() == protocol.PerspectiveServer {
			wg.Add(1)
			go func(handler packetHandler) {
				// blocks until the CONNECTION_CLOSE has been sent and the run-loop has stopped
				handler.shutdown()
				wg.Done()
			}(handler)
		}
	}
	h.mutex.Unlock()
	wg.Wait()
}

// Destroy closes the underlying connection and waits until listen() has returned.
// It does not close active sessions.
func (h *packetHandlerMap) Destroy() error {
	if err := h.conn.Close(); err != nil {
		return err
	}
	<-h.listening // wait until listening returns
	return nil
}

func (h *packetHandlerMap) close(e error) error {
	h.mutex.Lock()
	if h.closed {
		h.mutex.Unlock()
		return nil
	}

	var wg sync.WaitGroup
	for _, handler := range h.handlers {
		wg.Add(1)
		go func(handler packetHandler) {
			handler.destroy(e)
			wg.Done()
		}(handler)
	}

	if h.server != nil {
		h.server.setCloseError(e)
	}
	h.closed = true
	h.mutex.Unlock()
	wg.Wait()
	return getMultiplexer().RemoveConn(h.conn)
}

func (h *packetHandlerMap) listen() {
	defer close(h.listening)
	for {
		p, err := h.conn.ReadPacket()
		if nerr, ok := err.(net.Error); ok && nerr.Temporary() {
			h.logger.Debugf("Temporary error reading from conn: %w", err)
			continue
		}
		if err != nil {
			h.close(err)
			return
		}
		h.handlePacket(p)
	}
}

func (h *packetHandlerMap) handlePacket(p *receivedPacket) {
	connID, err := wire.ParseConnectionID(p.data, h.connIDLen)
	if err != nil {
		h.logger.Debugf("error parsing connection ID on packet from %s: %s", p.remoteAddr, err)
		if h.tracer != nil {
			h.tracer.DroppedPacket(p.remoteAddr, logging.PacketTypeNotDetermined, p.Size(), logging.PacketDropHeaderParseError)
		}
		p.buffer.MaybeRelease()
		return
	}

	h.mutex.Lock()
	defer h.mutex.Unlock()

	if isStatelessReset := h.maybeHandleStatelessReset(p.data); isStatelessReset {
		return
	}

	if handler, ok := h.handlers[string(connID)]; ok { // existing session
		handler.handlePacket(p)
		return
	}
	if p.data[0]&0x80 == 0 {
		go h.maybeSendStatelessReset(p, connID)
		return
	}
	if h.server == nil { // no server set
		h.logger.Debugf("received a packet with an unexpected connection ID %s", connID)
		return
	}
	h.server.handlePacket(p)
}

func (h *packetHandlerMap) maybeHandleStatelessReset(data []byte) bool {
	// stateless resets are always short header packets
	if data[0]&0x80 != 0 {
		return false
	}
	if len(data) < 17 /* type byte + 16 bytes for the reset token */ {
		return false
	}

	var token protocol.StatelessResetToken
	copy(token[:], data[len(data)-16:])
	if sess, ok := h.resetTokens[token]; ok {
		h.logger.Debugf("Received a stateless reset with token %#x. Closing session.", token)
		go sess.destroy(statelessResetErr{token: token})
		return true
	}
	return false
}

func (h *packetHandlerMap) GetStatelessResetToken(connID protocol.ConnectionID) protocol.StatelessResetToken {
	var token protocol.StatelessResetToken
	if !h.statelessResetEnabled {
		// Return a random stateless reset token.
		// This token will be sent in the server's transport parameters.
		// By using a random token, an off-path attacker won't be able to disrupt the connection.
		rand.Read(token[:])
		return token
	}
	h.statelessResetMutex.Lock()
	h.statelessResetHasher.Write(connID.Bytes())
	copy(token[:], h.statelessResetHasher.Sum(nil))
	h.statelessResetHasher.Reset()
	h.statelessResetMutex.Unlock()
	return token
}

func (h *packetHandlerMap) maybeSendStatelessReset(p *receivedPacket, connID protocol.ConnectionID) {
	defer p.buffer.Release()
	if !h.statelessResetEnabled {
		return
	}
	// Don't send a stateless reset in response to very small packets.
	// This includes packets that could be stateless resets.
	if len(p.data) <= protocol.MinStatelessResetSize {
		return
	}
	token := h.GetStatelessResetToken(connID)
	h.logger.Debugf("Sending stateless reset to %s (connection ID: %s). Token: %#x", p.remoteAddr, connID, token)
	data := make([]byte, protocol.MinStatelessResetSize-16, protocol.MinStatelessResetSize)
	rand.Read(data)
	data[0] = (data[0] & 0x7f) | 0x40
	data = append(data, token[:]...)
	if _, err := h.conn.WriteTo(data, p.remoteAddr); err != nil {
		h.logger.Debugf("Error sending Stateless Reset: %s", err)
	}
}
<commit_msg>improve error message when the read buffer size can't be set<commit_after>package quic

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"errors"
	"fmt"
	"hash"
	"log"
	"net"
	"sync"
	"time"

	"github.com/lucas-clemente/quic-go/internal/protocol"
	"github.com/lucas-clemente/quic-go/internal/utils"
	"github.com/lucas-clemente/quic-go/internal/wire"
	"github.com/lucas-clemente/quic-go/logging"
)

type statelessResetErr struct {
	token protocol.StatelessResetToken
}

func (e statelessResetErr) Error() string {
	return fmt.Sprintf("received a stateless reset with token %x", e.token)
}

// The packetHandlerMap stores packetHandlers, identified by connection ID.
// It is used:
// * by the server to store sessions
// * when multiplexing outgoing connections to store clients
type packetHandlerMap struct {
	mutex sync.Mutex

	conn      connection
	connIDLen int

	handlers    map[string] /* string(ConnectionID)*/ packetHandler
	resetTokens map[protocol.StatelessResetToken] /* stateless reset token */ packetHandler
	server      unknownPacketHandler

	listening chan struct{} // is closed when listen returns
	closed    bool

	deleteRetiredSessionsAfter time.Duration

	statelessResetEnabled bool
	statelessResetMutex   sync.Mutex
	statelessResetHasher  hash.Hash

	tracer logging.Tracer
	logger utils.Logger
}

var _ packetHandlerManager = &packetHandlerMap{}

func setReceiveBuffer(c net.PacketConn, logger utils.Logger) error {
	conn, ok := c.(interface{ SetReadBuffer(int) error })
	if !ok {
		return errors.New("connection doesn't allow setting of receive buffer size. Not a *net.UDPConn?")
	}
	size, err := inspectReadBuffer(c)
	if err != nil {
		return fmt.Errorf("failed to determine receive buffer size: %w", err)
	}
	if size >= protocol.DesiredReceiveBufferSize {
		logger.Debugf("Conn has receive buffer of %d kiB (wanted: at least %d kiB)", size/1024, protocol.DesiredReceiveBufferSize/1024)
	}
	if err := conn.SetReadBuffer(protocol.DesiredReceiveBufferSize); err != nil {
		return fmt.Errorf("failed to increase receive buffer size: %w", err)
	}
	newSize, err := inspectReadBuffer(c)
	if err != nil {
		return fmt.Errorf("failed to determine receive buffer size: %w", err)
	}
	if newSize == size {
		return fmt.Errorf("failed to determine receive buffer size: %w", err)
	}
	if newSize < protocol.DesiredReceiveBufferSize {
		return fmt.Errorf("failed to sufficiently increase receive buffer size (was: %d kiB, wanted: %d kiB, got: %d kiB)", size/1024, protocol.DesiredReceiveBufferSize/1024, newSize/1024)
	}
	logger.Debugf("Increased receive buffer size to %d kiB", newSize/1024)
	return nil
}

// only print warnings about the UPD receive buffer size once
var receiveBufferWarningOnce sync.Once

func newPacketHandlerMap(
	c net.PacketConn,
	connIDLen int,
	statelessResetKey []byte,
	tracer logging.Tracer,
	logger utils.Logger,
) (packetHandlerManager, error) {
	if err := setReceiveBuffer(c, logger); err != nil {
		receiveBufferWarningOnce.Do(func() {
			log.Printf("%s. See https://github.com/lucas-clemente/quic-go/wiki/UDP-Receive-Buffer-Size for details.", err)
		})
	}
	conn, err := wrapConn(c)
	if err != nil {
		return nil, err
	}
	m := &packetHandlerMap{
		conn:                       conn,
		connIDLen:                  connIDLen,
		listening:                  make(chan struct{}),
		handlers:                   make(map[string]packetHandler),
		resetTokens:                make(map[protocol.StatelessResetToken]packetHandler),
		deleteRetiredSessionsAfter: protocol.RetiredConnectionIDDeleteTimeout,
		statelessResetEnabled:      len(statelessResetKey) > 0,
		statelessResetHasher:       hmac.New(sha256.New, statelessResetKey),
		tracer:                     tracer,
		logger:                     logger,
	}
	go m.listen()

	if logger.Debug() {
		go m.logUsage()
	}
	return m, nil
}

func (h *packetHandlerMap) logUsage() {
	ticker := time.NewTicker(2 * time.Second)
	var printedZero bool
	for {
		select {
		case <-h.listening:
			return
		case <-ticker.C:
		}

		h.mutex.Lock()
		numHandlers := len(h.handlers)
		numTokens := len(h.resetTokens)
		h.mutex.Unlock()
		// If the number tracked handlers and tokens is zero, only print it a single time.
		hasZero := numHandlers == 0 && numTokens == 0
		if !hasZero || (hasZero && !printedZero) {
			h.logger.Debugf("Tracking %d connection IDs and %d reset tokens.\n", numHandlers, numTokens)
			printedZero = false
			if hasZero {
				printedZero = true
			}
		}
	}
}

func (h *packetHandlerMap) Add(id protocol.ConnectionID, handler packetHandler) bool /* was added */ {
	h.mutex.Lock()
	defer h.mutex.Unlock()

	if _, ok := h.handlers[string(id)]; ok {
		h.logger.Debugf("Not adding connection ID %s, as it already exists.", id)
		return false
	}
	h.handlers[string(id)] = handler
	h.logger.Debugf("Adding connection ID %s.", id)
	return true
}

func (h *packetHandlerMap) AddWithConnID(clientDestConnID, newConnID protocol.ConnectionID, fn func() packetHandler) bool {
	h.mutex.Lock()
	defer h.mutex.Unlock()

	if _, ok := h.handlers[string(clientDestConnID)]; ok {
		h.logger.Debugf("Not adding connection ID %s for a new session, as it already exists.", clientDestConnID)
		return false
	}

	sess := fn()
	h.handlers[string(clientDestConnID)] = sess
	h.handlers[string(newConnID)] = sess
	h.logger.Debugf("Adding connection IDs %s and %s for a new session.", clientDestConnID, newConnID)
	return true
}

func (h *packetHandlerMap) Remove(id protocol.ConnectionID) {
	h.mutex.Lock()
	delete(h.handlers, string(id))
	h.mutex.Unlock()
	h.logger.Debugf("Removing connection ID %s.", id)
}

func (h *packetHandlerMap) Retire(id protocol.ConnectionID) {
	h.logger.Debugf("Retiring connection ID %s in %s.", id, h.deleteRetiredSessionsAfter)
	time.AfterFunc(h.deleteRetiredSessionsAfter, func() {
		h.mutex.Lock()
		delete(h.handlers, string(id))
		h.mutex.Unlock()
		h.logger.Debugf("Removing connection ID %s after it has been retired.", id)
	})
}

func (h *packetHandlerMap) ReplaceWithClosed(id protocol.ConnectionID, handler packetHandler) {
	h.mutex.Lock()
	h.handlers[string(id)] = handler
	h.mutex.Unlock()
	h.logger.Debugf("Replacing session for connection ID %s with a closed session.", id)

	time.AfterFunc(h.deleteRetiredSessionsAfter, func() {
		h.mutex.Lock()
		handler.shutdown()
		delete(h.handlers, string(id))
		h.mutex.Unlock()
		h.logger.Debugf("Removing connection ID %s for a closed session after it has been retired.", id)
	})
}

func (h *packetHandlerMap) AddResetToken(token protocol.StatelessResetToken, handler packetHandler) {
	h.mutex.Lock()
	h.resetTokens[token] = handler
	h.mutex.Unlock()
}

func (h *packetHandlerMap) RemoveResetToken(token protocol.StatelessResetToken) {
	h.mutex.Lock()
	delete(h.resetTokens, token)
	h.mutex.Unlock()
}

func (h *packetHandlerMap) SetServer(s unknownPacketHandler) {
	h.mutex.Lock()
	h.server = s
	h.mutex.Unlock()
}

func (h *packetHandlerMap) CloseServer() {
	h.mutex.Lock()
	if h.server == nil {
		h.mutex.Unlock()
		return
	}
	h.server = nil
	var wg sync.WaitGroup
	for _, handler := range h.handlers {
		if handler.getPerspective() == protocol.PerspectiveServer {
			wg.Add(1)
			go func(handler packetHandler) {
				// blocks until the CONNECTION_CLOSE has been sent and the run-loop has stopped
				handler.shutdown()
				wg.Done()
			}(handler)
		}
	}
	h.mutex.Unlock()
	wg.Wait()
}

// Destroy closes the underlying connection and waits until listen() has returned.
// It does not close active sessions.
func (h *packetHandlerMap) Destroy() error {
	if err := h.conn.Close(); err != nil {
		return err
	}
	<-h.listening // wait until listening returns
	return nil
}

func (h *packetHandlerMap) close(e error) error {
	h.mutex.Lock()
	if h.closed {
		h.mutex.Unlock()
		return nil
	}

	var wg sync.WaitGroup
	for _, handler := range h.handlers {
		wg.Add(1)
		go func(handler packetHandler) {
			handler.destroy(e)
			wg.Done()
		}(handler)
	}

	if h.server != nil {
		h.server.setCloseError(e)
	}
	h.closed = true
	h.mutex.Unlock()
	wg.Wait()
	return getMultiplexer().RemoveConn(h.conn)
}

func (h *packetHandlerMap) listen() {
	defer close(h.listening)
	for {
		p, err := h.conn.ReadPacket()
		if nerr, ok := err.(net.Error); ok && nerr.Temporary() {
			h.logger.Debugf("Temporary error reading from conn: %w", err)
			continue
		}
		if err != nil {
			h.close(err)
			return
		}
		h.handlePacket(p)
	}
}

func (h *packetHandlerMap) handlePacket(p *receivedPacket) {
	connID, err := wire.ParseConnectionID(p.data, h.connIDLen)
	if err != nil {
		h.logger.Debugf("error parsing connection ID on packet from %s: %s", p.remoteAddr, err)
		if h.tracer != nil {
			h.tracer.DroppedPacket(p.remoteAddr, logging.PacketTypeNotDetermined, p.Size(), logging.PacketDropHeaderParseError)
		}
		p.buffer.MaybeRelease()
		return
	}

	h.mutex.Lock()
	defer h.mutex.Unlock()

	if isStatelessReset := h.maybeHandleStatelessReset(p.data); isStatelessReset {
		return
	}

	if handler, ok := h.handlers[string(connID)]; ok { // existing session
		handler.handlePacket(p)
		return
	}
	if p.data[0]&0x80 == 0 {
		go h.maybeSendStatelessReset(p, connID)
		return
	}
	if h.server == nil { // no server set
		h.logger.Debugf("received a packet with an unexpected connection ID %s", connID)
		return
	}
	h.server.handlePacket(p)
}

func (h *packetHandlerMap) maybeHandleStatelessReset(data []byte) bool {
	// stateless resets are always short header packets
	if data[0]&0x80 != 0 {
		return false
	}
	if len(data) < 17 /* type byte + 16 bytes for the reset token */ {
		return false
	}

	var token protocol.StatelessResetToken
	copy(token[:], data[len(data)-16:])
	if sess, ok := h.resetTokens[token]; ok {
		h.logger.Debugf("Received a stateless reset with token %#x. Closing session.", token)
		go sess.destroy(statelessResetErr{token: token})
		return true
	}
	return false
}

func (h *packetHandlerMap) GetStatelessResetToken(connID protocol.ConnectionID) protocol.StatelessResetToken {
	var token protocol.StatelessResetToken
	if !h.statelessResetEnabled {
		// Return a random stateless reset token.
		// This token will be sent in the server's transport parameters.
		// By using a random token, an off-path attacker won't be able to disrupt the connection.
		rand.Read(token[:])
		return token
	}
	h.statelessResetMutex.Lock()
	h.statelessResetHasher.Write(connID.Bytes())
	copy(token[:], h.statelessResetHasher.Sum(nil))
	h.statelessResetHasher.Reset()
	h.statelessResetMutex.Unlock()
	return token
}

func (h *packetHandlerMap) maybeSendStatelessReset(p *receivedPacket, connID protocol.ConnectionID) {
	defer p.buffer.Release()
	if !h.statelessResetEnabled {
		return
	}
	// Don't send a stateless reset in response to very small packets.
	// This includes packets that could be stateless resets.
	if len(p.data) <= protocol.MinStatelessResetSize {
		return
	}
	token := h.GetStatelessResetToken(connID)
	h.logger.Debugf("Sending stateless reset to %s (connection ID: %s). Token: %#x", p.remoteAddr, connID, token)
	data := make([]byte, protocol.MinStatelessResetSize-16, protocol.MinStatelessResetSize)
	rand.Read(data)
	data[0] = (data[0] & 0x7f) | 0x40
	data = append(data, token[:]...)
	if _, err := h.conn.WriteTo(data, p.remoteAddr); err != nil {
		h.logger.Debugf("Error sending Stateless Reset: %s", err)
	}
}
<|endoftext|>


<commit_before>package persister

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"

	whisper "github.com/go-graphite/go-whisper"
	"go.uber.org/zap"

	"github.com/lomik/go-carbon/helper"
	"github.com/lomik/go-carbon/points"
	"github.com/lomik/go-carbon/tags"
	"github.com/lomik/zapwriter"
)

const storeMutexCount = 32768

type StoreFunc func(metric string)

// Whisper write data to *.wsp files
type Whisper struct {
	helper.Stoppable
	recv                    func(chan bool) string
	pop                     func(string) (*points.Points, bool)
	confirm                 func(*points.Points)
	tagsEnabled             bool
	taggedFn                func(string, bool)
	schemas                 WhisperSchemas
	aggregation             *WhisperAggregation
	workersCount            int
	rootPath                string
	created                 uint32 // counter
	throttledCreates        uint32 // counter
	updateOperations        uint32 // counter
	committedPoints         uint32 // counter
	extended                uint32 // counter
	sparse                  bool
	flock                   bool
	compressed              bool
	hashFilenames           bool
	maxUpdatesPerSecond     int
	maxCreatesPerSecond     int
	hardMaxCreatesPerSecond bool
	throttleTicker          *ThrottleTicker
	maxCreatesTicker        *ThrottleTicker
	storeMutex              [storeMutexCount]sync.Mutex
	mockStore               func() (StoreFunc, func())
	logger                  *zap.Logger
	createLogger            *zap.Logger
	// blockThrottleNs        uint64 // sum ns counter
	// blockQueueGetNs        uint64 // sum ns counter
	// blockAvoidConcurrentNs uint64 // sum ns counter
	// blockUpdateManyNs      uint64 // sum ns counter
}

// NewWhisper create instance of Whisper
func NewWhisper(
	rootPath string,
	schemas WhisperSchemas,
	aggregation *WhisperAggregation,
	recv func(chan bool) string,
	pop func(string) (*points.Points, bool),
	confirm func(*points.Points)) *Whisper {

	return &Whisper{
		recv:                recv,
		pop:                 pop,
		confirm:             confirm,
		schemas:             schemas,
		aggregation:         aggregation,
		workersCount:        1,
		rootPath:            rootPath,
		maxUpdatesPerSecond: 0,
		logger:              zapwriter.Logger("persister"),
		createLogger:        zapwriter.Logger("whisper:new"),
	}
}

// SetMaxUpdatesPerSecond enable throttling
func (p *Whisper) SetMaxUpdatesPerSecond(maxUpdatesPerSecond int) {
	p.maxUpdatesPerSecond = maxUpdatesPerSecond
}

// SetMaxCreatesPerSecond enable throttling
func (p *Whisper) SetMaxCreatesPerSecond(maxCreatesPerSecond int) {
	p.maxCreatesPerSecond = maxCreatesPerSecond
}

// SetHardMaxCreatesPerSecond enable throttling
func (p *Whisper) SetHardMaxCreatesPerSecond(hardMaxCreatesPerSecond bool) {
	p.hardMaxCreatesPerSecond = hardMaxCreatesPerSecond
}

// GetMaxUpdatesPerSecond returns current throttling speed
func (p *Whisper) GetMaxUpdatesPerSecond() int {
	return p.maxUpdatesPerSecond
}

// SetWorkers count
func (p *Whisper) SetWorkers(count int) {
	if count >= 1 {
		p.workersCount = count
	} else {
		p.workersCount = 1
	}
}

// SetSparse creation
func (p *Whisper) SetSparse(sparse bool) {
	p.sparse = sparse
}

// SetFLock on create and open
func (p *Whisper) SetFLock(flock bool) {
	p.flock = flock
}

func (p *Whisper) SetCompressed(compressed bool) {
	p.compressed = compressed
}

func (p *Whisper) SetHashFilenames(v bool) {
	p.hashFilenames = v
}

func (p *Whisper) SetMockStore(fn func() (StoreFunc, func())) {
	p.mockStore = fn
}

func (p *Whisper) SetTagsEnabled(v bool) {
	p.tagsEnabled = v
}

func (p *Whisper) SetTaggedFn(fn func(string, bool)) {
	p.taggedFn = fn
}

func fnv32(key string) uint32 {
	hash := uint32(2166136261)
	const prime32 = uint32(16777619)
	for i := 0; i < len(key); i++ {
		hash *= prime32
		hash ^= uint32(key[i])
	}
	return hash
}

func (p *Whisper) updateMany(w *whisper.Whisper, path string, points []*whisper.TimeSeriesPoint) {
	defer func() {
		if r := recover(); r != nil {
			p.logger.Error("UpdateMany panic recovered",
				zap.String("path", path),
				zap.String("traceback", fmt.Sprint(r)),
			)
		}
	}()

	// start = time.Now()
	if err := w.UpdateMany(points); err != nil {
		p.logger.Error("fail to update metric",
			zap.String("path", path),
			zap.Error(err),
		)
	}
	if w.Extended {
		atomic.AddUint32(&p.extended, 1)
		p.logger.Info("cwhisper file has extended", zap.String("path", path))
	}
}

func (p *Whisper) store(metric string) {
	// avoid concurrent store same metric
	// @TODO: may be flock?
	// start := time.Now()
	mutexIndex := fnv32(metric) % storeMutexCount
	p.storeMutex[mutexIndex].Lock()
	// atomic.AddUint64(&p.blockAvoidConcurrentNs, uint64(time.Since(start).Nanoseconds()))
	defer p.storeMutex[mutexIndex].Unlock()

	var path string
	if p.tagsEnabled && strings.IndexByte(metric, ';') >= 0 {
		path = tags.FilePath(p.rootPath, metric, p.hashFilenames) + ".wsp"
	} else {
		path = filepath.Join(p.rootPath, strings.Replace(metric, ".", "/", -1)+".wsp")
	}

	w, err := whisper.OpenWithOptions(path, &whisper.Options{
		FLock:      p.flock,
		Compressed: p.compressed,
	})
	if err != nil {
		// create new whisper if file not exists
		if !os.IsNotExist(err) {
			p.logger.Error("failed to open whisper file", zap.String("path", path), zap.Error(err))
			if pathErr, isPathErr := err.(*os.PathError); isPathErr && pathErr.Err == syscall.ENAMETOOLONG {
				p.pop(metric)
			}
			return
		}

		if t := p.maxCreatesThrottling(); t != throttlingOff {
			if t == throttlingHard {
				p.pop(metric)
			}

			atomic.AddUint32(&p.throttledCreates, 1)
			p.logger.Error("metric creation throttled",
				zap.String("name", metric),
				zap.String("operation", "create"),
				zap.Bool("dropped", t == throttlingHard),
			)

			return
		}

		schema, ok := p.schemas.Match(metric)
		if !ok {
			p.logger.Error("no storage schema defined for metric", zap.String("metric", metric))
			return
		}

		aggr := p.aggregation.match(metric)
		if aggr == nil {
			p.logger.Error("no storage aggregation defined for metric", zap.String("metric", metric))
			return
		}

		if err = os.MkdirAll(filepath.Dir(path), os.ModeDir|os.ModePerm); err != nil {
			p.logger.Error("mkdir failed",
				zap.String("dir", filepath.Dir(path)),
				zap.Error(err),
				zap.String("path", path),
			)
			return
		}

		compressed := p.compressed
		if schema.Compressed != nil {
			compressed = *schema.Compressed
		}
		w, err = whisper.CreateWithOptions(path, schema.Retentions, aggr.aggregationMethod, float32(aggr.xFilesFactor), &whisper.Options{
			Sparse:     p.sparse,
			FLock:      p.flock,
			Compressed: compressed,
		})
		if err != nil {
			p.logger.Error("create new whisper file failed",
				zap.String("path", path),
				zap.Error(err),
				zap.String("retention", schema.RetentionStr),
				zap.String("schema", schema.Name),
				zap.String("aggregation", aggr.name),
				zap.Float64("xFilesFactor", aggr.xFilesFactor),
				zap.String("method", aggr.aggregationMethodStr),
				zap.Bool("compressed", compressed),
			)
			return
		}

		if p.tagsEnabled && p.taggedFn != nil && strings.IndexByte(metric, ';') >= 0 {
			p.taggedFn(metric, true)
		}

		p.createLogger.Debug("created",
			zap.String("path", path),
			zap.String("retention", schema.RetentionStr),
			zap.String("schema", schema.Name),
			zap.String("aggregation", aggr.name),
			zap.Float64("xFilesFactor", aggr.xFilesFactor),
			zap.String("method", aggr.aggregationMethodStr),
			zap.Bool("compressed", compressed),
		)

		atomic.AddUint32(&p.created, 1)
	}

	values, exists := p.pop(metric)
	if !exists {
		return
	}
	if p.confirm != nil {
		defer p.confirm(values)
	}

	points := make([]*whisper.TimeSeriesPoint, len(values.Data))
	for i, r := range values.Data {
		points[i] = &whisper.TimeSeriesPoint{Time: int(r.Timestamp), Value: r.Value}
	}

	atomic.AddUint32(&p.committedPoints, uint32(len(values.Data)))
	atomic.AddUint32(&p.updateOperations, 1)

	// start = time.Now()
	p.updateMany(w, path, points)

	w.Close()
	// atomic.AddUint64(&p.blockUpdateManyNs, uint64(time.Since(start).Nanoseconds()))

	if p.tagsEnabled && p.taggedFn != nil && strings.IndexByte(metric, ';') >= 0 {
		p.taggedFn(metric, false)
	}
}

type throttleMode int

const (
	throttlingOff  throttleMode = iota
	throttlingSoft throttleMode = iota
	throttlingHard throttleMode = iota
)

// For test error messages
func (t throttleMode) String() string {
	if t == throttlingHard {
		return "hard"
	}

	if t == throttlingSoft {
		return "soft"
	}

	return "off"
}

func (p *Whisper) maxCreatesThrottling() throttleMode {
	if p.hardMaxCreatesPerSecond {
		keep, open := <-p.maxCreatesTicker.C
		if !open {
			return throttlingHard
		}

		if keep {
			return throttlingOff
		}

		return throttlingHard
	} else {
		select {
		case <-p.maxCreatesTicker.C:
			return throttlingOff

		default:
			return throttlingSoft
		}
	}
}

func (p *Whisper) worker(exit chan bool) {
	storeFunc := p.store
	var doneCb func()
	if p.mockStore != nil {
		storeFunc, doneCb = p.mockStore()
	}

LOOP:
	for {
		// start := time.Now()
		select {
		case <-p.throttleTicker.C:
			// atomic.AddUint64(&p.blockThrottleNs, uint64(time.Since(start).Nanoseconds()))
			// pass
		case <-exit:
			return
		}

		// start = time.Now()
		metric := p.recv(exit)
		// atomic.AddUint64(&p.blockQueueGetNs, uint64(time.Since(start).Nanoseconds()))
		if metric == "" {
			// exit closed
			break LOOP
		}
		storeFunc(metric)
		if doneCb != nil {
			doneCb()
		}
	}
}

// Stat callback
func (p *Whisper) Stat(send helper.StatCallback) {
	updateOperations := atomic.LoadUint32(&p.updateOperations)
	committedPoints := atomic.LoadUint32(&p.committedPoints)
	atomic.AddUint32(&p.updateOperations, -updateOperations)
	atomic.AddUint32(&p.committedPoints, -committedPoints)

	created := atomic.LoadUint32(&p.created)
	atomic.AddUint32(&p.created, -created)

	extended := atomic.LoadUint32(&p.extended)
	atomic.AddUint32(&p.extended, -extended)

	throttledCreates := atomic.LoadUint32(&p.throttledCreates)
	atomic.AddUint32(&p.throttledCreates, -throttledCreates)

	send("updateOperations", float64(updateOperations))
	send("committedPoints", float64(committedPoints))
	if updateOperations > 0 {
		send("pointsPerUpdate", float64(committedPoints)/float64(updateOperations))
	} else {
		send("pointsPerUpdate", 0.0)
	}

	send("created", float64(created))
	send("throttledCreates", float64(throttledCreates))
	send("maxCreatesPerSecond", float64(p.maxCreatesPerSecond))

	send("maxUpdatesPerSecond", float64(p.maxUpdatesPerSecond))
	send("workers", float64(p.workersCount))
	send("extended", float64(extended))

	// helper.SendAndSubstractUint64("blockThrottleNs", &p.blockThrottleNs, send)
	// helper.SendAndSubstractUint64("blockQueueGetNs", &p.blockQueueGetNs, send)
	// helper.SendAndSubstractUint64("blockAvoidConcurrentNs", &p.blockAvoidConcurrentNs, send)
	// helper.SendAndSubstractUint64("blockUpdateManyNs", &p.blockUpdateManyNs, send)

}

// Start worker
func (p *Whisper) Start() error {
	return p.StartFunc(func() error {
		p.throttleTicker = NewThrottleTicker(p.maxUpdatesPerSecond)
		if p.hardMaxCreatesPerSecond {
			p.maxCreatesTicker = NewHardThrottleTicker(p.maxCreatesPerSecond)
		} else {
			p.maxCreatesTicker = NewThrottleTicker(p.maxCreatesPerSecond)
		}

		for i := 0; i < p.workersCount; i++ {
			p.Go(p.worker)
		}

		return nil
	})
}

func (p *Whisper) Stop() {
	p.StopFunc(func() {
		p.throttleTicker.Stop()
		p.maxCreatesTicker.Stop()
	})
}
<commit_msg>persister: handle empty whisper files caused by edge cases like server reboot<commit_after>package persister

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"

	whisper "github.com/go-graphite/go-whisper"
	"go.uber.org/zap"

	"github.com/lomik/go-carbon/helper"
	"github.com/lomik/go-carbon/points"
	"github.com/lomik/go-carbon/tags"
	"github.com/lomik/zapwriter"
)

const storeMutexCount = 32768

type StoreFunc func(metric string)

// Whisper write data to *.wsp files
type Whisper struct {
	helper.Stoppable
	recv                    func(chan bool) string
	pop                     func(string) (*points.Points, bool)
	confirm                 func(*points.Points)
	tagsEnabled             bool
	taggedFn                func(string, bool)
	schemas                 WhisperSchemas
	aggregation             *WhisperAggregation
	workersCount            int
	rootPath                string
	created                 uint32 // counter
	throttledCreates        uint32 // counter
	updateOperations        uint32 // counter
	committedPoints         uint32 // counter
	extended                uint32 // counter
	sparse                  bool
	flock                   bool
	compressed              bool
	hashFilenames           bool
	maxUpdatesPerSecond     int
	maxCreatesPerSecond     int
	hardMaxCreatesPerSecond bool
	throttleTicker          *ThrottleTicker
	maxCreatesTicker        *ThrottleTicker
	storeMutex              [storeMutexCount]sync.Mutex
	mockStore               func() (StoreFunc, func())
	logger                  *zap.Logger
	createLogger            *zap.Logger
	// blockThrottleNs        uint64 // sum ns counter
	// blockQueueGetNs        uint64 // sum ns counter
	// blockAvoidConcurrentNs uint64 // sum ns counter
	// blockUpdateManyNs      uint64 // sum ns counter
}

// NewWhisper create instance of Whisper
func NewWhisper(
	rootPath string,
	schemas WhisperSchemas,
	aggregation *WhisperAggregation,
	recv func(chan bool) string,
	pop func(string) (*points.Points, bool),
	confirm func(*points.Points)) *Whisper {

	return &Whisper{
		recv:                recv,
		pop:                 pop,
		confirm:             confirm,
		schemas:             schemas,
		aggregation:         aggregation,
		workersCount:        1,
		rootPath:            rootPath,
		maxUpdatesPerSecond: 0,
		logger:              zapwriter.Logger("persister"),
		createLogger:        zapwriter.Logger("whisper:new"),
	}
}

// SetMaxUpdatesPerSecond enable throttling
func (p *Whisper) SetMaxUpdatesPerSecond(maxUpdatesPerSecond int) {
	p.maxUpdatesPerSecond = maxUpdatesPerSecond
}

// SetMaxCreatesPerSecond enable throttling
func (p *Whisper) SetMaxCreatesPerSecond(maxCreatesPerSecond int) {
	p.maxCreatesPerSecond = maxCreatesPerSecond
}

// SetHardMaxCreatesPerSecond enable throttling
func (p *Whisper) SetHardMaxCreatesPerSecond(hardMaxCreatesPerSecond bool) {
	p.hardMaxCreatesPerSecond = hardMaxCreatesPerSecond
}

// GetMaxUpdatesPerSecond returns current throttling speed
func (p *Whisper) GetMaxUpdatesPerSecond() int {
	return p.maxUpdatesPerSecond
}

// SetWorkers count
func (p *Whisper) SetWorkers(count int) {
	if count >= 1 {
		p.workersCount = count
	} else {
		p.workersCount = 1
	}
}

// SetSparse creation
func (p *Whisper) SetSparse(sparse bool) {
	p.sparse = sparse
}

// SetFLock on create and open
func (p *Whisper) SetFLock(flock bool) {
	p.flock = flock
}

func (p *Whisper) SetCompressed(compressed bool) {
	p.compressed = compressed
}

func (p *Whisper) SetHashFilenames(v bool) {
	p.hashFilenames = v
}

func (p *Whisper) SetMockStore(fn func() (StoreFunc, func())) {
	p.mockStore = fn
}

func (p *Whisper) SetTagsEnabled(v bool) {
	p.tagsEnabled = v
}

func (p *Whisper) SetTaggedFn(fn func(string, bool)) {
	p.taggedFn = fn
}

func fnv32(key string) uint32 {
	hash := uint32(2166136261)
	const prime32 = uint32(16777619)
	for i := 0; i < len(key); i++ {
		hash *= prime32
		hash ^= uint32(key[i])
	}
	return hash
}

func (p *Whisper) updateMany(w *whisper.Whisper, path string, points []*whisper.TimeSeriesPoint) {
	defer func() {
		if r := recover(); r != nil {
			p.logger.Error("UpdateMany panic recovered",
				zap.String("path", path),
				zap.String("traceback", fmt.Sprint(r)),
			)
		}
	}()

	// start = time.Now()
	if err := w.UpdateMany(points); err != nil {
		p.logger.Error("fail to update metric",
			zap.String("path", path),
			zap.Error(err),
		)
	}
	if w.Extended {
		atomic.AddUint32(&p.extended, 1)
		p.logger.Info("cwhisper file has extended", zap.String("path", path))
	}
}

func (p *Whisper) store(metric string) {
	// avoid concurrent store same metric
	// @TODO: may be flock?
	// start := time.Now()
	mutexIndex := fnv32(metric) % storeMutexCount
	p.storeMutex[mutexIndex].Lock()
	// atomic.AddUint64(&p.blockAvoidConcurrentNs, uint64(time.Since(start).Nanoseconds()))
	defer p.storeMutex[mutexIndex].Unlock()

	var path string
	if p.tagsEnabled && strings.IndexByte(metric, ';') >= 0 {
		path = tags.FilePath(p.rootPath, metric, p.hashFilenames) + ".wsp"
	} else {
		path = filepath.Join(p.rootPath, strings.Replace(metric, ".", "/", -1)+".wsp")
	}

	w, err := whisper.OpenWithOptions(path, &whisper.Options{
		FLock:      p.flock,
		Compressed: p.compressed,
	})
	if err != nil {
		// create new whisper if file not exists
		if !os.IsNotExist(err) {
			// There are cases that new files are created but data are not saved
			// on disk due to edge issues like server panics/reboot. So it's
			// better to treat empty files as NotExist to work around this
			// problem.
			stat, err2 := os.Stat(path)
			if err2 != nil {
				p.logger.Error("failed to stat whisper file", zap.String("path", path), zap.Error(err2))
			}
			if err2 != nil || stat.Size() > 0 {
				p.logger.Error("failed to open whisper file", zap.String("path", path), zap.Error(err))
				if pathErr, isPathErr := err.(*os.PathError); isPathErr && pathErr.Err == syscall.ENAMETOOLONG {
					p.pop(metric)
				}
				return
			}
			if err := os.Remove(path); err != nil {
				p.logger.Error("failed to delete empty whisper file", zap.String("path", path), zap.Error(err))
				p.pop(metric)
				return
			}
			p.logger.Warn("deleted empty whisper file", zap.String("path", path))
		}

		if t := p.maxCreatesThrottling(); t != throttlingOff {
			if t == throttlingHard {
				p.pop(metric)
			}

			atomic.AddUint32(&p.throttledCreates, 1)
			p.logger.Error("metric creation throttled",
				zap.String("name", metric),
				zap.String("operation", "create"),
				zap.Bool("dropped", t == throttlingHard),
			)

			return
		}

		schema, ok := p.schemas.Match(metric)
		if !ok {
			p.logger.Error("no storage schema defined for metric", zap.String("metric", metric))
			return
		}

		aggr := p.aggregation.match(metric)
		if aggr == nil {
			p.logger.Error("no storage aggregation defined for metric", zap.String("metric", metric))
			return
		}

		if err = os.MkdirAll(filepath.Dir(path), os.ModeDir|os.ModePerm); err != nil {
			p.logger.Error("mkdir failed",
				zap.String("dir", filepath.Dir(path)),
				zap.Error(err),
				zap.String("path", path),
			)
			return
		}

		compressed := p.compressed
		if schema.Compressed != nil {
			compressed = *schema.Compressed
		}
		w, err = whisper.CreateWithOptions(path, schema.Retentions, aggr.aggregationMethod, float32(aggr.xFilesFactor), &whisper.Options{
			Sparse:     p.sparse,
			FLock:      p.flock,
			Compressed: compressed,
		})
		if err != nil {
			p.logger.Error("create new whisper file failed",
				zap.String("path", path),
				zap.Error(err),
				zap.String("retention", schema.RetentionStr),
				zap.String("schema", schema.Name),
				zap.String("aggregation", aggr.name),
				zap.Float64("xFilesFactor", aggr.xFilesFactor),
				zap.String("method", aggr.aggregationMethodStr),
				zap.Bool("compressed", compressed),
			)
			return
		}

		if p.tagsEnabled && p.taggedFn != nil && strings.IndexByte(metric, ';') >= 0 {
			p.taggedFn(metric, true)
		}

		p.createLogger.Debug("created",
			zap.String("path", path),
			zap.String("retention", schema.RetentionStr),
			zap.String("schema", schema.Name),
			zap.String("aggregation", aggr.name),
			zap.Float64("xFilesFactor", aggr.xFilesFactor),
			zap.String("method", aggr.aggregationMethodStr),
			zap.Bool("compressed", compressed),
		)

		atomic.AddUint32(&p.created, 1)
	}

	values, exists := p.pop(metric)
	if !exists {
		return
	}
	if p.confirm != nil {
		defer p.confirm(values)
	}

	points := make([]*whisper.TimeSeriesPoint, len(values.Data))
	for i, r := range values.Data {
		points[i] = &whisper.TimeSeriesPoint{Time: int(r.Timestamp), Value: r.Value}
	}

	atomic.AddUint32(&p.committedPoints, uint32(len(values.Data)))
	atomic.AddUint32(&p.updateOperations, 1)

	// start = time.Now()
	p.updateMany(w, path, points)

	w.Close()
	// atomic.AddUint64(&p.blockUpdateManyNs, uint64(time.Since(start).Nanoseconds()))

	if p.tagsEnabled && p.taggedFn != nil && strings.IndexByte(metric, ';') >= 0 {
		p.taggedFn(metric, false)
	}
}

type throttleMode int

const (
	throttlingOff  throttleMode = iota
	throttlingSoft throttleMode = iota
	throttlingHard throttleMode = iota
)

// For test error messages
func (t throttleMode) String() string {
	if t == throttlingHard {
		return "hard"
	}

	if t == throttlingSoft {
		return "soft"
	}

	return "off"
}

func (p *Whisper) maxCreatesThrottling() throttleMode {
	if p.hardMaxCreatesPerSecond {
		keep, open := <-p.maxCreatesTicker.C
		if !open {
			return throttlingHard
		}

		if keep {
			return throttlingOff
		}

		return throttlingHard
	} else {
		select {
		case <-p.maxCreatesTicker.C:
			return throttlingOff

		default:
			return throttlingSoft
		}
	}
}

func (p *Whisper) worker(exit chan bool) {
	storeFunc := p.store
	var doneCb func()
	if p.mockStore != nil {
		storeFunc, doneCb = p.mockStore()
	}

LOOP:
	for {
		// start := time.Now()
		select {
		case <-p.throttleTicker.C:
			// atomic.AddUint64(&p.blockThrottleNs, uint64(time.Since(start).Nanoseconds()))
			// pass
		case <-exit:
			return
		}

		// start = time.Now()
		metric := p.recv(exit)
		// atomic.AddUint64(&p.blockQueueGetNs, uint64(time.Since(start).Nanoseconds()))
		if metric == "" {
			// exit closed
			break LOOP
		}
		storeFunc(metric)
		if doneCb != nil {
			doneCb()
		}
	}
}

// Stat callback
func (p *Whisper) Stat(send helper.StatCallback) {
	updateOperations := atomic.LoadUint32(&p.updateOperations)
	committedPoints := atomic.LoadUint32(&p.committedPoints)
	atomic.AddUint32(&p.updateOperations, -updateOperations)
	atomic.AddUint32(&p.committedPoints, -committedPoints)

	created := atomic.LoadUint32(&p.created)
	atomic.AddUint32(&p.created, -created)

	extended := atomic.LoadUint32(&p.extended)
	atomic.AddUint32(&p.extended, -extended)

	throttledCreates := atomic.LoadUint32(&p.throttledCreates)
	atomic.AddUint32(&p.throttledCreates, -throttledCreates)

	send("updateOperations", float64(updateOperations))
	send("committedPoints", float64(committedPoints))
	if updateOperations > 0 {
		send("pointsPerUpdate", float64(committedPoints)/float64(updateOperations))
	} else {
		send("pointsPerUpdate", 0.0)
	}

	send("created", float64(created))
	send("throttledCreates", float64(throttledCreates))
	send("maxCreatesPerSecond", float64(p.maxCreatesPerSecond))

	send("maxUpdatesPerSecond", float64(p.maxUpdatesPerSecond))
	send("workers", float64(p.workersCount))
	send("extended", float64(extended))

	// helper.SendAndSubstractUint64("blockThrottleNs", &p.blockThrottleNs, send)
	// helper.SendAndSubstractUint64("blockQueueGetNs", &p.blockQueueGetNs, send)
	// helper.SendAndSubstractUint64("blockAvoidConcurrentNs", &p.blockAvoidConcurrentNs, send)
	// helper.SendAndSubstractUint64("blockUpdateManyNs", &p.blockUpdateManyNs, send)

}

// Start worker
func (p *Whisper) Start() error {
	return p.StartFunc(func() error {
		p.throttleTicker = NewThrottleTicker(p.maxUpdatesPerSecond)
		if p.hardMaxCreatesPerSecond {
			p.maxCreatesTicker = NewHardThrottleTicker(p.maxCreatesPerSecond)
		} else {
			p.maxCreatesTicker = NewThrottleTicker(p.maxCreatesPerSecond)
		}

		for i := 0; i < p.workersCount; i++ {
			p.Go(p.worker)
		}

		return nil
	})
}

func (p *Whisper) Stop() {
	p.StopFunc(func() {
		p.throttleTicker.Stop()
		p.maxCreatesTicker.Stop()
	})
}
<|endoftext|>


<commit_before>package libp2pquic

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"errors"
	"math/big"
	"time"

	"github.com/gogo/protobuf/proto"
	ic "github.com/libp2p/go-libp2p-crypto"
	pb "github.com/libp2p/go-libp2p-crypto/pb"
)

// mint certificate selection is broken.
const hostname = "quic.ipfs"

const certValidityPeriod = 180 * 24 * time.Hour

func generateConfig(privKey ic.PrivKey) (*tls.Config, error) {
	key, hostCert, err := keyToCertificate(privKey)
	if err != nil {
		return nil, err
	}
	// The ephemeral key used just for a couple of connections (or a limited time).
	ephemeralKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, err
	}
	// Sign the ephemeral key using the host key.
	// This is the only time that the host's private key of the peer is needed.
	// Note that this step could be done asynchronously, such that a running node doesn't need access its private key at all.
	certTemplate := &x509.Certificate{
		DNSNames:     []string{hostname},
		SerialNumber: big.NewInt(1),
		NotBefore:    time.Now().Add(-24 * time.Hour),
		NotAfter:     time.Now().Add(certValidityPeriod),
	}
	certDER, err := x509.CreateCertificate(rand.Reader, certTemplate, hostCert, ephemeralKey.Public(), key)
	if err != nil {
		return nil, err
	}
	cert, err := x509.ParseCertificate(certDER)
	if err != nil {
		return nil, err
	}
	return &tls.Config{
		ServerName:         hostname,
		InsecureSkipVerify: true, // This is not insecure here. We will verify the cert chain ourselves.
		ClientAuth:         tls.RequireAnyClientCert,
		Certificates: []tls.Certificate{{
			Certificate: [][]byte{cert.Raw, hostCert.Raw},
			PrivateKey:  ephemeralKey,
		}},
	}, nil
}

func getRemotePubKey(chain []*x509.Certificate) (ic.PubKey, error) {
	if len(chain) != 2 {
		return nil, errors.New("expected 2 certificates in the chain")
	}
	pool := x509.NewCertPool()
	pool.AddCert(chain[1])
	if _, err := chain[0].Verify(x509.VerifyOptions{Roots: pool}); err != nil {
		return nil, err
	}
	remotePubKey, err := x509.MarshalPKIXPublicKey(chain[1].PublicKey)
	if err != nil {
		return nil, err
	}
	return ic.UnmarshalRsaPublicKey(remotePubKey)
}

func keyToCertificate(sk ic.PrivKey) (interface{}, *x509.Certificate, error) {
	sn, err := rand.Int(rand.Reader, big.NewInt(1<<62))
	if err != nil {
		return nil, nil, err
	}
	tmpl := &x509.Certificate{
		SerialNumber: sn,
		NotBefore:    time.Now().Add(-24 * time.Hour),
		NotAfter:     time.Now().Add(certValidityPeriod),
		IsCA:         true,
		BasicConstraintsValid: true,
	}

	var publicKey, privateKey interface{}
	keyBytes, err := sk.Bytes()
	if err != nil {
		return nil, nil, err
	}
	pbmes := new(pb.PrivateKey)
	if err := proto.Unmarshal(keyBytes, pbmes); err != nil {
		return nil, nil, err
	}
	switch pbmes.GetType() {
	case pb.KeyType_RSA:
		k, err := x509.ParsePKCS1PrivateKey(pbmes.GetData())
		if err != nil {
			return nil, nil, err
		}
		publicKey = &k.PublicKey
		privateKey = k
	// TODO: add support for ECDSA
	default:
		return nil, nil, errors.New("unsupported key type for TLS")
	}
	certDER, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, publicKey, privateKey)
	if err != nil {
		return nil, nil, err
	}
	cert, err := x509.ParseCertificate(certDER)
	if err != nil {
		return nil, nil, err
	}
	return privateKey, cert, nil
}
<commit_msg>run Go 1.11 gofmt, use Go 1.10.x and 1.11.x on Travis<commit_after>package libp2pquic

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"errors"
	"math/big"
	"time"

	"github.com/gogo/protobuf/proto"
	ic "github.com/libp2p/go-libp2p-crypto"
	pb "github.com/libp2p/go-libp2p-crypto/pb"
)

// mint certificate selection is broken.
const hostname = "quic.ipfs"

const certValidityPeriod = 180 * 24 * time.Hour

func generateConfig(privKey ic.PrivKey) (*tls.Config, error) {
	key, hostCert, err := keyToCertificate(privKey)
	if err != nil {
		return nil, err
	}
	// The ephemeral key used just for a couple of connections (or a limited time).
	ephemeralKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, err
	}
	// Sign the ephemeral key using the host key.
	// This is the only time that the host's private key of the peer is needed.
	// Note that this step could be done asynchronously, such that a running node doesn't need access its private key at all.
	certTemplate := &x509.Certificate{
		DNSNames:     []string{hostname},
		SerialNumber: big.NewInt(1),
		NotBefore:    time.Now().Add(-24 * time.Hour),
		NotAfter:     time.Now().Add(certValidityPeriod),
	}
	certDER, err := x509.CreateCertificate(rand.Reader, certTemplate, hostCert, ephemeralKey.Public(), key)
	if err != nil {
		return nil, err
	}
	cert, err := x509.ParseCertificate(certDER)
	if err != nil {
		return nil, err
	}
	return &tls.Config{
		ServerName:         hostname,
		InsecureSkipVerify: true, // This is not insecure here. We will verify the cert chain ourselves.
		ClientAuth:         tls.RequireAnyClientCert,
		Certificates: []tls.Certificate{{
			Certificate: [][]byte{cert.Raw, hostCert.Raw},
			PrivateKey:  ephemeralKey,
		}},
	}, nil
}

func getRemotePubKey(chain []*x509.Certificate) (ic.PubKey, error) {
	if len(chain) != 2 {
		return nil, errors.New("expected 2 certificates in the chain")
	}
	pool := x509.NewCertPool()
	pool.AddCert(chain[1])
	if _, err := chain[0].Verify(x509.VerifyOptions{Roots: pool}); err != nil {
		return nil, err
	}
	remotePubKey, err := x509.MarshalPKIXPublicKey(chain[1].PublicKey)
	if err != nil {
		return nil, err
	}
	return ic.UnmarshalRsaPublicKey(remotePubKey)
}

func keyToCertificate(sk ic.PrivKey) (interface{}, *x509.Certificate, error) {
	sn, err := rand.Int(rand.Reader, big.NewInt(1<<62))
	if err != nil {
		return nil, nil, err
	}
	tmpl := &x509.Certificate{
		SerialNumber:          sn,
		NotBefore:             time.Now().Add(-24 * time.Hour),
		NotAfter:              time.Now().Add(certValidityPeriod),
		IsCA:                  true,
		BasicConstraintsValid: true,
	}

	var publicKey, privateKey interface{}
	keyBytes, err := sk.Bytes()
	if err != nil {
		return nil, nil, err
	}
	pbmes := new(pb.PrivateKey)
	if err := proto.Unmarshal(keyBytes, pbmes); err != nil {
		return nil, nil, err
	}
	switch pbmes.GetType() {
	case pb.KeyType_RSA:
		k, err := x509.ParsePKCS1PrivateKey(pbmes.GetData())
		if err != nil {
			return nil, nil, err
		}
		publicKey = &k.PublicKey
		privateKey = k
	// TODO: add support for ECDSA
	default:
		return nil, nil, errors.New("unsupported key type for TLS")
	}
	certDER, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, publicKey, privateKey)
	if err != nil {
		return nil, nil, err
	}
	cert, err := x509.ParseCertificate(certDER)
	if err != nil {
		return nil, nil, err
	}
	return privateKey, cert, nil
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/iam"
	randomprovider "github.com/hashicorp/terraform/builtin/providers/random"
	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAWSPolicyAttachment_basic(t *testing.T) {
	var out iam.ListEntitiesForPolicyOutput

	user1 := fmt.Sprintf("test-user-%d", acctest.RandInt())
	user2 := fmt.Sprintf("test-user-%d", acctest.RandInt())
	user3 := fmt.Sprintf("test-user-%d", acctest.RandInt())

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSPolicyAttachmentDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccAWSPolicyAttachConfig(user1),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSPolicyAttachmentExists("aws_iam_policy_attachment.test-attach", 3, &out),
					testAccCheckAWSPolicyAttachmentAttributes([]string{user1}, []string{"test-role"}, []string{"test-group"}, &out),
				),
			},
			resource.TestStep{
				Config: testAccAWSPolicyAttachConfigUpdate(user1, user2, user3),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSPolicyAttachmentExists("aws_iam_policy_attachment.test-attach", 6, &out),
					testAccCheckAWSPolicyAttachmentAttributes([]string{user3, user3}, []string{"test-role2", "test-role3"}, []string{"test-group2", "test-group3"}, &out),
				),
			},
		},
	})
}

func TestAccAWSPolicyAttachment_paginatedEntities(t *testing.T) {
	var out iam.ListEntitiesForPolicyOutput

	resource.Test(t, resource.TestCase{
		PreCheck: func() { testAccPreCheck(t) },
		Providers: map[string]terraform.ResourceProvider{
			"aws":    testAccProvider,
			"random": randomprovider.Provider(),
		},
		CheckDestroy: testAccCheckAWSPolicyAttachmentDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccAWSPolicyPaginatedAttachConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSPolicyAttachmentExists("aws_iam_policy_attachment.test-paginated-attach", 101, &out),
				),
			},
		},
	})
}

func testAccCheckAWSPolicyAttachmentDestroy(s *terraform.State) error {
	return nil
}

func testAccCheckAWSPolicyAttachmentExists(n string, c int64, out *iam.ListEntitiesForPolicyOutput) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No policy name is set")
		}

		conn := testAccProvider.Meta().(*AWSClient).iamconn
		arn := rs.Primary.Attributes["policy_arn"]

		resp, err := conn.GetPolicy(&iam.GetPolicyInput{
			PolicyArn: aws.String(arn),
		})
		if err != nil {
			return fmt.Errorf("Error: Policy (%s) not found", n)
		}
		if c != *resp.Policy.AttachmentCount {
			return fmt.Errorf("Error: Policy (%s) has wrong number of entities attached on initial creation", n)
		}
		resp2, err := conn.ListEntitiesForPolicy(&iam.ListEntitiesForPolicyInput{
			PolicyArn: aws.String(arn),
		})
		if err != nil {
			return fmt.Errorf("Error: Failed to get entities for Policy (%s)", arn)
		}

		*out = *resp2
		return nil
	}
}

func testAccCheckAWSPolicyAttachmentAttributes(users []string, roles []string, groups []string, out *iam.ListEntitiesForPolicyOutput) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		uc := len(users)
		rc := len(roles)
		gc := len(groups)

		for _, u := range users {
			for _, pu := range out.PolicyUsers {
				if u == *pu.UserName {
					uc--
				}
			}
		}
		for _, r := range roles {
			for _, pr := range out.PolicyRoles {
				if r == *pr.RoleName {
					rc--
				}
			}
		}
		for _, g := range groups {
			for _, pg := range out.PolicyGroups {
				if g == *pg.GroupName {
					gc--
				}
			}
		}
		if uc != 0 || rc != 0 || gc != 0 {
			return fmt.Errorf("Error: Number of attached users, roles, or groups was incorrect:\n expected %d users and found %d\nexpected %d roles and found %d\nexpected %d groups and found %d", len(users), len(users)-uc, len(roles), len(roles)-rc, len(groups), len(groups)-gc)
		}
		return nil
	}
}

func testAccAWSPolicyAttachConfig(u1 string) string {
	return fmt.Sprintf(`
resource "aws_iam_user" "user" {
    name = "%s"
}
resource "aws_iam_role" "role" {
    name = "test-role"
	  assume_role_policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": "sts:AssumeRole",
      "Principal": {
        "Service": "ec2.amazonaws.com"
      },
      "Effect": "Allow",
      "Sid": ""
    }
  ]
}
EOF
}
resource "aws_iam_group" "group" {
    name = "test-group"
}

resource "aws_iam_policy" "policy" {
    name = "test-policy"
    description = "A test policy"
    policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": [
        "iam:ChangePassword"
      ],
      "Resource": "*",
      "Effect": "Allow"
    }
  ]
}
EOF
}

resource "aws_iam_policy_attachment" "test-attach" {
    name = "test-attachment"
    users = ["${aws_iam_user.user.name}"]
    roles = ["${aws_iam_role.role.name}"]
    groups = ["${aws_iam_group.group.name}"]
    policy_arn = "${aws_iam_policy.policy.arn}"
}`, u1)
}

func testAccAWSPolicyAttachConfigUpdate(u1, u2, u3 string) string {
	return fmt.Sprintf(`
resource "aws_iam_user" "user" {
    name = "%s"
}
resource "aws_iam_user" "user2" {
    name = "%s"
}
resource "aws_iam_user" "user3" {
    name = "%s"
}
resource "aws_iam_role" "role" {
    name = "test-role"
	  assume_role_policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": "sts:AssumeRole",
      "Principal": {
        "Service": "ec2.amazonaws.com"
      },
      "Effect": "Allow",
      "Sid": ""
    }
  ]
}
EOF
}

resource "aws_iam_role" "role2" {
    name = "test-role2"
	  assume_role_policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": "sts:AssumeRole",
      "Principal": {
        "Service": "ec2.amazonaws.com"
      },
      "Effect": "Allow",
      "Sid": ""
    }
  ]
}
EOF

}
resource "aws_iam_role" "role3" {
    name = "test-role3"
	  assume_role_policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": "sts:AssumeRole",
      "Principal": {
        "Service": "ec2.amazonaws.com"
      },
      "Effect": "Allow",
      "Sid": ""
    }
  ]
}
EOF

}
resource "aws_iam_group" "group" {
    name = "test-group"
}
resource "aws_iam_group" "group2" {
    name = "test-group2"
}
resource "aws_iam_group" "group3" {
    name = "test-group3"
}

resource "aws_iam_policy" "policy" {
    name = "test-policy"
    description = "A test policy"
    policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": [
        "iam:ChangePassword"
      ],
      "Resource": "*",
      "Effect": "Allow"
    }
  ]
}
EOF
}

resource "aws_iam_policy_attachment" "test-attach" {
    name = "test-attachment"
    users = [
        "${aws_iam_user.user2.name}",
        "${aws_iam_user.user3.name}"
    ]
    roles = [
        "${aws_iam_role.role2.name}",
        "${aws_iam_role.role3.name}"
    ]
    groups = [
        "${aws_iam_group.group2.name}",
        "${aws_iam_group.group3.name}"
    ]
    policy_arn = "${aws_iam_policy.policy.arn}"
}`, u1, u2, u3)
}

const testAccAWSPolicyPaginatedAttachConfig = `
resource "random_id" "user_id" {
  byte_length = 10
}

resource "aws_iam_user" "user" {
    count = 101
    name = "${format("paged-test-user-${random_id.user_id.hex}-%d", count.index + 1)}"
}

resource "aws_iam_policy" "policy" {
    name = "test-policy"
    description = "A test policy"
    policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": [
        "iam:ChangePassword"
      ],
      "Resource": "*",
      "Effect": "Allow"
    }
  ]
}
EOF
}

resource "aws_iam_policy_attachment" "test-paginated-attach" {
    name = "test-attachment"
    users = ["${aws_iam_user.user.*.name}"]
    policy_arn = "${aws_iam_policy.policy.arn}"
}
`
<commit_msg>Removed random provider frm iam policy attachment test<commit_after>package aws

import (
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAWSPolicyAttachment_basic(t *testing.T) {
	var out iam.ListEntitiesForPolicyOutput

	user1 := fmt.Sprintf("test-user-%d", acctest.RandInt())
	user2 := fmt.Sprintf("test-user-%d", acctest.RandInt())
	user3 := fmt.Sprintf("test-user-%d", acctest.RandInt())

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSPolicyAttachmentDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccAWSPolicyAttachConfig(user1),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSPolicyAttachmentExists("aws_iam_policy_attachment.test-attach", 3, &out),
					testAccCheckAWSPolicyAttachmentAttributes([]string{user1}, []string{"test-role"}, []string{"test-group"}, &out),
				),
			},
			resource.TestStep{
				Config: testAccAWSPolicyAttachConfigUpdate(user1, user2, user3),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSPolicyAttachmentExists("aws_iam_policy_attachment.test-attach", 6, &out),
					testAccCheckAWSPolicyAttachmentAttributes([]string{user3, user3}, []string{"test-role2", "test-role3"}, []string{"test-group2", "test-group3"}, &out),
				),
			},
		},
	})
}

func TestAccAWSPolicyAttachment_paginatedEntities(t *testing.T) {
	var out iam.ListEntitiesForPolicyOutput
	rInt := acctest.RandInt()

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSPolicyAttachmentDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccAWSPolicyPaginatedAttachConfig(rInt),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSPolicyAttachmentExists("aws_iam_policy_attachment.test-paginated-attach", 101, &out),
				),
			},
		},
	})
}

func testAccCheckAWSPolicyAttachmentDestroy(s *terraform.State) error {
	return nil
}

func testAccCheckAWSPolicyAttachmentExists(n string, c int64, out *iam.ListEntitiesForPolicyOutput) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No policy name is set")
		}

		conn := testAccProvider.Meta().(*AWSClient).iamconn
		arn := rs.Primary.Attributes["policy_arn"]

		resp, err := conn.GetPolicy(&iam.GetPolicyInput{
			PolicyArn: aws.String(arn),
		})
		if err != nil {
			return fmt.Errorf("Error: Policy (%s) not found", n)
		}
		if c != *resp.Policy.AttachmentCount {
			return fmt.Errorf("Error: Policy (%s) has wrong number of entities attached on initial creation", n)
		}
		resp2, err := conn.ListEntitiesForPolicy(&iam.ListEntitiesForPolicyInput{
			PolicyArn: aws.String(arn),
		})
		if err != nil {
			return fmt.Errorf("Error: Failed to get entities for Policy (%s)", arn)
		}

		*out = *resp2
		return nil
	}
}

func testAccCheckAWSPolicyAttachmentAttributes(users []string, roles []string, groups []string, out *iam.ListEntitiesForPolicyOutput) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		uc := len(users)
		rc := len(roles)
		gc := len(groups)

		for _, u := range users {
			for _, pu := range out.PolicyUsers {
				if u == *pu.UserName {
					uc--
				}
			}
		}
		for _, r := range roles {
			for _, pr := range out.PolicyRoles {
				if r == *pr.RoleName {
					rc--
				}
			}
		}
		for _, g := range groups {
			for _, pg := range out.PolicyGroups {
				if g == *pg.GroupName {
					gc--
				}
			}
		}
		if uc != 0 || rc != 0 || gc != 0 {
			return fmt.Errorf("Error: Number of attached users, roles, or groups was incorrect:\n expected %d users and found %d\nexpected %d roles and found %d\nexpected %d groups and found %d", len(users), len(users)-uc, len(roles), len(roles)-rc, len(groups), len(groups)-gc)
		}
		return nil
	}
}

func testAccAWSPolicyAttachConfig(u1 string) string {
	return fmt.Sprintf(`
resource "aws_iam_user" "user" {
    name = "%s"
}
resource "aws_iam_role" "role" {
    name = "test-role"
	  assume_role_policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": "sts:AssumeRole",
      "Principal": {
        "Service": "ec2.amazonaws.com"
      },
      "Effect": "Allow",
      "Sid": ""
    }
  ]
}
EOF
}
resource "aws_iam_group" "group" {
    name = "test-group"
}

resource "aws_iam_policy" "policy" {
    name = "test-policy"
    description = "A test policy"
    policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": [
        "iam:ChangePassword"
      ],
      "Resource": "*",
      "Effect": "Allow"
    }
  ]
}
EOF
}

resource "aws_iam_policy_attachment" "test-attach" {
    name = "test-attachment"
    users = ["${aws_iam_user.user.name}"]
    roles = ["${aws_iam_role.role.name}"]
    groups = ["${aws_iam_group.group.name}"]
    policy_arn = "${aws_iam_policy.policy.arn}"
}`, u1)
}

func testAccAWSPolicyAttachConfigUpdate(u1, u2, u3 string) string {
	return fmt.Sprintf(`
resource "aws_iam_user" "user" {
    name = "%s"
}
resource "aws_iam_user" "user2" {
    name = "%s"
}
resource "aws_iam_user" "user3" {
    name = "%s"
}
resource "aws_iam_role" "role" {
    name = "test-role"
	  assume_role_policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": "sts:AssumeRole",
      "Principal": {
        "Service": "ec2.amazonaws.com"
      },
      "Effect": "Allow",
      "Sid": ""
    }
  ]
}
EOF
}

resource "aws_iam_role" "role2" {
    name = "test-role2"
	  assume_role_policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": "sts:AssumeRole",
      "Principal": {
        "Service": "ec2.amazonaws.com"
      },
      "Effect": "Allow",
      "Sid": ""
    }
  ]
}
EOF

}
resource "aws_iam_role" "role3" {
    name = "test-role3"
	  assume_role_policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": "sts:AssumeRole",
      "Principal": {
        "Service": "ec2.amazonaws.com"
      },
      "Effect": "Allow",
      "Sid": ""
    }
  ]
}
EOF

}
resource "aws_iam_group" "group" {
    name = "test-group"
}
resource "aws_iam_group" "group2" {
    name = "test-group2"
}
resource "aws_iam_group" "group3" {
    name = "test-group3"
}

resource "aws_iam_policy" "policy" {
    name = "test-policy"
    description = "A test policy"
    policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": [
        "iam:ChangePassword"
      ],
      "Resource": "*",
      "Effect": "Allow"
    }
  ]
}
EOF
}

resource "aws_iam_policy_attachment" "test-attach" {
    name = "test-attachment"
    users = [
        "${aws_iam_user.user2.name}",
        "${aws_iam_user.user3.name}"
    ]
    roles = [
        "${aws_iam_role.role2.name}",
        "${aws_iam_role.role3.name}"
    ]
    groups = [
        "${aws_iam_group.group2.name}",
        "${aws_iam_group.group3.name}"
    ]
    policy_arn = "${aws_iam_policy.policy.arn}"
}`, u1, u2, u3)
}

func testAccAWSPolicyPaginatedAttachConfig(rInt int) string {
	return fmt.Sprintf(`
resource "aws_iam_user" "user" {
	count = 101
	name = "${format("paged-test-user-%d-%%d", count.index + 1)}"
}
resource "aws_iam_policy" "policy" {
	name = "test-policy"
	description = "A test policy"
	policy = <<EOF
{
"Version": "2012-10-17",
"Statement": [
	{
		"Action": [
			"iam:ChangePassword"
		],
		"Resource": "*",
		"Effect": "Allow"
	}
]
}
EOF
}
resource "aws_iam_policy_attachment" "test-paginated-attach" {
	name = "test-attachment"
	users = ["${aws_iam_user.user.*.name}"]
	policy_arn = "${aws_iam_policy.policy.arn}"
}`, rInt)
}
<|endoftext|>


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

package formats

import (
	"encoding/base64"
	"encoding/json"
	"io"
	"sort"
	"strconv"
	"strings"
	"unicode/utf8"

	"golang.org/x/net/context"

	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"

	rdb "go.chromium.org/luci/results"
	resultspb "go.chromium.org/luci/results/proto/v1"
	"go.chromium.org/luci/results/util"
)

// GTestResults represents the structure as described to be generated in
// https://cs.chromium.org/chromium/src/base/test/launcher/test_results_tracker.h?l=83&rcl=96020cfd447cb285acfa1a96c37a67ed22fa2499
// (base::TestResultsTracker::SaveSummaryAsJSON)
//
// Fields not used by Test Results are omitted.
type GTestResults struct {
	AllTests   []string `json:"all_tests"`
	GlobalTags []string `json:"global_tags"`

	// PerIterationData is a vector of run iterations, each mapping test names to a list of test data.
	PerIterationData []map[string][]*GTestRunResult `json:"per_iteration_data"`

	// TestLocations maps test names to their location in code.
	TestLocations map[string]*Location `json:"test_locations"`
}

// GTestRunResult represents the per_iteration_data as described in
// https://cs.chromium.org/chromium/src/base/test/launcher/test_results_tracker.h?l=83&rcl=96020cfd447cb285acfa1a96c37a67ed22fa2499
// (base::TestResultsTracker::SaveSummaryAsJSON)
//
// Fields not used by Test Results are omitted.
type GTestRunResult struct {
	Status        string `json:"status"`
	ElapsedTimeMs int    `json:"elapsed_time_ms"`

	LosslessSnippet     bool   `json:"losless_snippet"`
	OutputSnippetBase64 string `json:"output_snippet_base64"`
}

// Location describes a code location.
type Location struct {
	File string `json:"file"`
	Line int    `json:"line"`
}

// ConvertFromJSON reads the provided reader into the receiver.
func (r *GTestResults) ConvertFromJSON(ctx context.Context, reader io.Reader) error {
	if err := json.NewDecoder(reader).Decode(r); err != nil {
		return err
	}

	if len(r.AllTests) == 0 {
		return errors.Reason(`missing "all_tests" field in JSON`).Err()
	}

	return nil
}

// ToInvocation converts a GTestResults to a (partial) resultspb.Invocation.
func (r *GTestResults) ToInvocation(ctx context.Context, req *resultspb.DeriveInvocationFromSwarmingRequest) (*resultspb.Invocation, error) {
	// testsToVariants maps GTest base name to the variant IDs present to Results.
	// We store these maps to handle the same test and variants appearing in multiple iterations.
	testsToVariants := map[string]map[string][]*resultspb.TestResult{}
	variantDefs := map[string]*resultspb.VariantDef{}

	// In theory, we can have multiple iterations. This seems rare in practice, so log if we do see
	// more than one to confirm and track.
	if len(r.PerIterationData) > 1 {
		logging.Infof(ctx,
			"Got %d iterations for task %s on %s", len(r.PerIterationData), req.Task.Id, req.Task.Hostname)
	}

	// Assume the invocation is complete for now; if any results are NOTRUN, we'll mark as otherwise.
	incomplete := false

	for i, data := range r.PerIterationData {
		for name, results := range data {
			baseName, params := extractGTestParameters(name)

			// Derive the variant definition from base variant and params.
			variantDefMap := rdb.MergeTestVariantMaps(req.BaseTestVariant.Def, params)
			if err := variantDefMap.Validate(); err != nil {
				return nil, errors.Annotate(err,
					"iteration %d of test %s failed to validate variant def", i, name).Err()
			}
			variantDef := variantDefMap.Proto()
			varID := variantDef.Digest
			if _, ok := variantDefs[varID]; !ok {
				variantDefs[varID] = variantDef
			}

			// Generate individual result protos.
			rpbs := make([]*resultspb.TestResult, len(results))
			for i, result := range results {
				if rpb, err := r.convertRunResult(ctx, name, result); err != nil {
					return nil, errors.Annotate(err,
						"iteration %d of test %s failed to convert run result", i, name).Err()

				} else {
					rpbs[i] = rpb
				}

				// TODO(jchinlee): Verify that it's indeed the case that getting NOTRUN results in the final
				// results indicates the task was incomplete.
				// TODO(jchinlee): Check how unexpected SKIPPED tests should be handled.
				if result.Status == "NOTRUN" {
					incomplete = true
				}
			}

			// Store the processed test result into the correct part of the overall map.
			testPath := req.TestPathPrefix + baseName
			if _, ok := testsToVariants[testPath]; !ok {
				testsToVariants[testPath] = map[string][]*resultspb.TestResult{}
			}
			testsToVariants[testPath][varID] = append(testsToVariants[testPath][varID], rpbs...)
		}
	}

	// Now that we have all the results mapped by test path and variant, populate the invocation.
	inv := &resultspb.Invocation{
		Incomplete:  incomplete,
		VariantDefs: variantDefs,
		Tags:        []*resultspb.StringPair{util.StringPair("test_framework", "gtest")},
	}

	for testPath, variants := range testsToVariants {
		testpb := &resultspb.Invocation_Test{Path: testPath}

		for id, results := range variants {
			testpb.Variants = append(testpb.Variants, &resultspb.Invocation_TestVariant{
				VariantId: id,
				Results:   results,
			})
		}

		inv.Tests = append(inv.Tests, testpb)
	}
	sort.Slice(inv.Tests, func(i, j int) bool { return inv.Tests[i].Path < inv.Tests[j].Path })

	// Populate the tags.
	for _, tag := range r.GlobalTags {
		inv.Tags = append(inv.Tags, util.StringPair("gtest_global_tag", tag))
	}
	SortTagsInPlace(inv.Tags)

	return inv, nil
}

func toGTestStatus(s string) (resultspb.Status, error) {
	switch s {
	case "SUCCESS":
		return resultspb.Status_PASS, nil
	case "FAILURE":
		return resultspb.Status_FAIL, nil
	case "FAILURE_ON_EXIT":
		return resultspb.Status_FAIL, nil
	case "TIMEOUT":
		return resultspb.Status_ABORT, nil
	case "CRASH":
		return resultspb.Status_CRASH, nil
	case "SKIPPED":
		return resultspb.Status_SKIP, nil
	case "EXCESSIVE_OUTPUT":
		return resultspb.Status_FAIL, nil
	case "NOTRUN":
		return resultspb.Status_SKIP, nil
	default:
		// This would only happen if the set of possible GTest result statuses change and resultsdb has
		// not been updated to match.
		return resultspb.Status_STATUS_UNSPECIFIED, errors.Reason("unknown GTest status %q", s).Err()
	}
}

// extractGTestParameters extracts parameters from a test path as a mapping with "param/" keys.
func extractGTestParameters(testPath string) (basePath string, params rdb.VariantDefMap) {
	// TODO(jchinlee): Implement.
	basePath = testPath
	params = rdb.VariantDefMap{}
	return
}

func (r *GTestResults) convertRunResult(ctx context.Context, name string, result *GTestRunResult) (*resultspb.TestResult, error) {
	status, err := toGTestStatus(result.Status)
	if err != nil {
		return nil, err
	}

	rpb := &resultspb.TestResult{
		Status:   status,
		Tags:     []*resultspb.StringPair{util.StringPair("gtest_status", result.Status)},
		Duration: secondsToDuration(1e-6 * float64(result.ElapsedTimeMs)),
	}

	// Write the summary.
	if result.OutputSnippetBase64 != "" {
		outputBytes, err := base64.StdEncoding.DecodeString(result.OutputSnippetBase64)
		if err != nil {
			// Log the error, but we shouldn't fail to convert an entire invocation just because we can't
			// convert a summary.
			logging.Errorf(ctx, "Failed to convert OutputSnippetBase64 %q", result.OutputSnippetBase64)
		} else {
			// TODO(jchinlee): Escape Markdown.
			rpb.SummaryMarkdown = strings.ToValidUTF8(string(outputBytes), string(utf8.RuneError))
		}
	}

	// Store the correct output snippet.
	rpb.Tags = append(
		rpb.Tags,
		util.StringPair("lossless_snippet", strconv.FormatBool(result.LosslessSnippet)),
	)

	// Store the test code location.
	if loc, ok := r.TestLocations[name]; ok {
		rpb.Tags = append(rpb.Tags,
			util.StringPair("file", loc.File),
			util.StringPair("line", strconv.Itoa(loc.Line)),
		)
	}

	SortTagsInPlace(rpb.Tags)

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

package formats

import (
	"encoding/base64"
	"encoding/json"
	"io"
	"sort"
	"strconv"
	"strings"
	"unicode/utf8"

	"golang.org/x/net/context"

	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"

	rdb "go.chromium.org/luci/results"
	resultspb "go.chromium.org/luci/results/proto/v1"
	"go.chromium.org/luci/results/util"
)

// GTestResults represents the structure as described to be generated in
// https://cs.chromium.org/chromium/src/base/test/launcher/test_results_tracker.h?l=83&rcl=96020cfd447cb285acfa1a96c37a67ed22fa2499
// (base::TestResultsTracker::SaveSummaryAsJSON)
//
// Fields not used by Test Results are omitted.
type GTestResults struct {
	AllTests   []string `json:"all_tests"`
	GlobalTags []string `json:"global_tags"`

	// PerIterationData is a vector of run iterations, each mapping test names to a list of test data.
	PerIterationData []map[string][]*GTestRunResult `json:"per_iteration_data"`

	// TestLocations maps test names to their location in code.
	TestLocations map[string]*Location `json:"test_locations"`
}

// GTestRunResult represents the per_iteration_data as described in
// https://cs.chromium.org/chromium/src/base/test/launcher/test_results_tracker.h?l=83&rcl=96020cfd447cb285acfa1a96c37a67ed22fa2499
// (base::TestResultsTracker::SaveSummaryAsJSON)
//
// Fields not used by Test Results are omitted.
type GTestRunResult struct {
	Status        string `json:"status"`
	ElapsedTimeMs int    `json:"elapsed_time_ms"`

	LosslessSnippet     bool   `json:"losless_snippet"`
	OutputSnippetBase64 string `json:"output_snippet_base64"`
}

// Location describes a code location.
type Location struct {
	File string `json:"file"`
	Line int    `json:"line"`
}

// ConvertFromJSON reads the provided reader into the receiver.
//
// The receiver is cleared and its fields overwritten.
func (r *GTestResults) ConvertFromJSON(ctx context.Context, reader io.Reader) error {
	*r = GTestResults{}
	if err := json.NewDecoder(reader).Decode(r); err != nil {
		return err
	}

	if len(r.AllTests) == 0 {
		return errors.Reason(`missing "all_tests" field in JSON`).Err()
	}

	return nil
}

// ToInvocation converts a GTestResults to a (partial) resultspb.Invocation.
func (r *GTestResults) ToInvocation(ctx context.Context, req *resultspb.DeriveInvocationFromSwarmingRequest) (*resultspb.Invocation, error) {
	// testsToVariants maps GTest base name to the variant IDs present to Results.
	// We store these maps to handle the same test and variants appearing in multiple iterations.
	testsToVariants := map[string]map[string][]*resultspb.TestResult{}
	variantDefs := map[string]*resultspb.VariantDef{}

	// In theory, we can have multiple iterations. This seems rare in practice, so log if we do see
	// more than one to confirm and track.
	if len(r.PerIterationData) > 1 {
		logging.Infof(ctx,
			"Got %d iterations for task %s on %s", len(r.PerIterationData), req.Task.Id, req.Task.Hostname)
	}

	// Assume the invocation is complete for now; if any results are NOTRUN, we'll mark as otherwise.
	incomplete := false

	for i, data := range r.PerIterationData {
		for name, results := range data {
			baseName, params := extractGTestParameters(name)

			// Derive the variant definition from base variant and params.
			variantDefMap := rdb.MergeTestVariantMaps(req.BaseTestVariant.Def, params)
			if err := variantDefMap.Validate(); err != nil {
				return nil, errors.Annotate(err,
					"iteration %d of test %s failed to validate variant def", i, name).Err()
			}
			variantDef := variantDefMap.Proto()
			varID := variantDef.Digest
			if _, ok := variantDefs[varID]; !ok {
				variantDefs[varID] = variantDef
			}

			// Generate individual result protos.
			rpbs := make([]*resultspb.TestResult, len(results))
			for i, result := range results {
				if rpb, err := r.convertRunResult(ctx, name, result); err != nil {
					return nil, errors.Annotate(err,
						"iteration %d of test %s failed to convert run result", i, name).Err()

				} else {
					rpbs[i] = rpb
				}

				// TODO(jchinlee): Verify that it's indeed the case that getting NOTRUN results in the final
				// results indicates the task was incomplete.
				// TODO(jchinlee): Check how unexpected SKIPPED tests should be handled.
				if result.Status == "NOTRUN" {
					incomplete = true
				}
			}

			// Store the processed test result into the correct part of the overall map.
			testPath := req.TestPathPrefix + baseName
			if _, ok := testsToVariants[testPath]; !ok {
				testsToVariants[testPath] = map[string][]*resultspb.TestResult{}
			}
			testsToVariants[testPath][varID] = append(testsToVariants[testPath][varID], rpbs...)
		}
	}

	// Now that we have all the results mapped by test path and variant, populate the invocation.
	inv := &resultspb.Invocation{
		Incomplete:  incomplete,
		VariantDefs: variantDefs,
		Tags:        []*resultspb.StringPair{util.StringPair("test_framework", "gtest")},
	}

	for testPath, variants := range testsToVariants {
		testpb := &resultspb.Invocation_Test{Path: testPath}

		for id, results := range variants {
			testpb.Variants = append(testpb.Variants, &resultspb.Invocation_TestVariant{
				VariantId: id,
				Results:   results,
			})
		}

		inv.Tests = append(inv.Tests, testpb)
	}
	sort.Slice(inv.Tests, func(i, j int) bool { return inv.Tests[i].Path < inv.Tests[j].Path })

	// Populate the tags.
	for _, tag := range r.GlobalTags {
		inv.Tags = append(inv.Tags, util.StringPair("gtest_global_tag", tag))
	}
	SortTagsInPlace(inv.Tags)

	return inv, nil
}

func toGTestStatus(s string) (resultspb.Status, error) {
	switch s {
	case "SUCCESS":
		return resultspb.Status_PASS, nil
	case "FAILURE":
		return resultspb.Status_FAIL, nil
	case "FAILURE_ON_EXIT":
		return resultspb.Status_FAIL, nil
	case "TIMEOUT":
		return resultspb.Status_ABORT, nil
	case "CRASH":
		return resultspb.Status_CRASH, nil
	case "SKIPPED":
		return resultspb.Status_SKIP, nil
	case "EXCESSIVE_OUTPUT":
		return resultspb.Status_FAIL, nil
	case "NOTRUN":
		return resultspb.Status_SKIP, nil
	default:
		// This would only happen if the set of possible GTest result statuses change and resultsdb has
		// not been updated to match.
		return resultspb.Status_STATUS_UNSPECIFIED, errors.Reason("unknown GTest status %q", s).Err()
	}
}

// extractGTestParameters extracts parameters from a test path as a mapping with "param/" keys.
func extractGTestParameters(testPath string) (basePath string, params rdb.VariantDefMap) {
	// TODO(jchinlee): Implement.
	basePath = testPath
	params = rdb.VariantDefMap{}
	return
}

func (r *GTestResults) convertRunResult(ctx context.Context, name string, result *GTestRunResult) (*resultspb.TestResult, error) {
	status, err := toGTestStatus(result.Status)
	if err != nil {
		return nil, err
	}

	rpb := &resultspb.TestResult{
		Status:   status,
		Tags:     []*resultspb.StringPair{util.StringPair("gtest_status", result.Status)},
		Duration: secondsToDuration(1e-6 * float64(result.ElapsedTimeMs)),
	}

	// Write the summary.
	if result.OutputSnippetBase64 != "" {
		outputBytes, err := base64.StdEncoding.DecodeString(result.OutputSnippetBase64)
		if err != nil {
			// Log the error, but we shouldn't fail to convert an entire invocation just because we can't
			// convert a summary.
			logging.Errorf(ctx, "Failed to convert OutputSnippetBase64 %q", result.OutputSnippetBase64)
		} else {
			// TODO(jchinlee): Escape Markdown.
			rpb.SummaryMarkdown = strings.ToValidUTF8(string(outputBytes), string(utf8.RuneError))
		}
	}

	// Store the correct output snippet.
	rpb.Tags = append(
		rpb.Tags,
		util.StringPair("lossless_snippet", strconv.FormatBool(result.LosslessSnippet)),
	)

	// Store the test code location.
	if loc, ok := r.TestLocations[name]; ok {
		rpb.Tags = append(rpb.Tags,
			util.StringPair("file", loc.File),
			util.StringPair("line", strconv.Itoa(loc.Line)),
		)
	}

	SortTagsInPlace(rpb.Tags)

	return rpb, nil
}
<|endoftext|>


<commit_before>package controllers

import (
	"database/sql"
	// _ "github.com/go-sql-driver/mysql"
	"github.com/robfig/revel"
	"github.com/robfig/revel/modules/db/app"
	"math/rand"
	"runtime"
	"sort"
	"sync"
)

type MessageStruct struct {
	Message string `json:"message"`
}

type World struct {
	Id           uint16 `json:"id"`
	RandomNumber uint16 `json:"randomNumber"`
}

type Fortune struct {
	Id      uint16 `json:"id"`
	Message string `json:"message"`
}

const (
	WorldSelect        = "SELECT id,randomNumber FROM World where id=?"
	FortuneSelect      = "SELECT id,message FROM Fortune"
	WorldUpdate        = "UPDATE World SET randomNumber = ? where id = ?"
	WorldRowCount      = 10000
	MaxConnectionCount = 100
)

var (
	worldStatement   *sql.Stmt
	fortuneStatement *sql.Stmt
	updateStatement  *sql.Stmt
)

func init() {
	revel.OnAppStart(func() {
		var err error
		runtime.GOMAXPROCS(runtime.NumCPU())
		db.DbPlugin{}.OnAppStart()
		db.Db.SetMaxIdleConns(MaxConnectionCount)
		if worldStatement, err = db.Db.Prepare(WorldSelect); err != nil {
			revel.ERROR.Fatalln(err)
		}
		if fortuneStatement, err = db.Db.Prepare(FortuneSelect); err != nil {
			revel.ERROR.Fatalln(err)
		}
		if updateStatement, err = db.Db.Prepare(WorldUpdate); err != nil {
			revel.ERROR.Fatalln(err)
		}
	})
}

type App struct {
	*revel.Controller
}

func (c App) Json() revel.Result {
	c.Response.ContentType = "application/json"
	return c.RenderJson(MessageStruct{"Hello, world"})
}

func (c App) Db(queries int) revel.Result {
	if queries <= 1 {
		rowNum := rand.Intn(WorldRowCount) + 1
		var w World
		worldStatement.QueryRow(rowNum).Scan(&w.Id, &w.RandomNumber)
		return c.RenderJson(w)
	}

	ww := make([]World, queries)
	var wg sync.WaitGroup
	wg.Add(queries)
	for i := 0; i < queries; i++ {
		go func(i int) {
			rowNum := rand.Intn(WorldRowCount) + 1
			err := worldStatement.QueryRow(rowNum).Scan(&ww[i].Id, &ww[i].RandomNumber)
			if err != nil {
				revel.ERROR.Fatalf("Error scanning world row: %v", err)
			}
			wg.Done()
		}(i)
	}
	wg.Wait()
	return c.RenderJson(ww)
}

func (c App) Update(queries int) revel.Result {
	if queries <= 1 {
		rowNum := rand.Intn(WorldRowCount) + 1
		var w World
		worldStatement.QueryRow(rowNum).Scan(&w.Id, &w.RandomNumber)
		w.RandomNumber = uint16(rand.Intn(WorldRowCount) + 1)
		updateStatement.Exec(w.RandomNumber, w.Id)
		return c.RenderJson(&w)
	}

	var (
		ww = make([]World, queries)
		wg sync.WaitGroup
	)
	wg.Add(queries)
	for i := 0; i < queries; i++ {
		go func(i int) {
			rowNum := rand.Intn(WorldRowCount) + 1
			err := worldStatement.QueryRow(rowNum).Scan(&ww[i].Id, &ww[i].RandomNumber)
			if err != nil {
				revel.ERROR.Fatalf("Error scanning world row: %v", err)
			}
			ww[i].RandomNumber = uint16(rand.Intn(WorldRowCount) + 1)
			updateStatement.Exec(ww[i].RandomNumber, ww[i].Id)
			wg.Done()
		}(i)
	}
	wg.Wait()
	return c.RenderJson(ww)
}

func (c App) Fortune() revel.Result {
	fortunes := make([]*Fortune, 0, 16)

	rows, err := fortuneStatement.Query()
	if err != nil {
		revel.ERROR.Fatalf("Error preparing statement: %v", err)
	}

	i := 0
	var fortune *Fortune
	for rows.Next() {
		fortune = new(Fortune)
		if err = rows.Scan(&fortune.Id, &fortune.Message); err != nil {
			revel.ERROR.Fatalf("Error scanning fortune row: %v", err)
		}
		fortunes = append(fortunes, fortune)
		i++
	}
	fortunes = append(fortunes, &Fortune{Message: "Additional fortune added at request time."})

	sort.Sort(ByMessage{fortunes})
	return c.Render(fortunes)
}

type Fortunes []*Fortune

func (s Fortunes) Len() int      { return len(s) }
func (s Fortunes) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

type ByMessage struct{ Fortunes }

func (s ByMessage) Less(i, j int) bool { return s.Fortunes[i].Message < s.Fortunes[j].Message }
<commit_msg>remove temp variable<commit_after>package controllers

import (
	"database/sql"
	"github.com/robfig/revel"
	"github.com/robfig/revel/modules/db/app"
	"math/rand"
	"runtime"
	"sort"
	"sync"
)

type MessageStruct struct {
	Message string `json:"message"`
}

type World struct {
	Id           uint16 `json:"id"`
	RandomNumber uint16 `json:"randomNumber"`
}

type Fortune struct {
	Id      uint16 `json:"id"`
	Message string `json:"message"`
}

const (
	WorldSelect        = "SELECT id,randomNumber FROM World where id=?"
	FortuneSelect      = "SELECT id,message FROM Fortune"
	WorldUpdate        = "UPDATE World SET randomNumber = ? where id = ?"
	WorldRowCount      = 10000
	MaxConnectionCount = 100
)

var (
	worldStatement   *sql.Stmt
	fortuneStatement *sql.Stmt
	updateStatement  *sql.Stmt
)

func init() {
	revel.OnAppStart(func() {
		var err error
		runtime.GOMAXPROCS(runtime.NumCPU())
		db.DbPlugin{}.OnAppStart()
		db.Db.SetMaxIdleConns(MaxConnectionCount)
		if worldStatement, err = db.Db.Prepare(WorldSelect); err != nil {
			revel.ERROR.Fatalln(err)
		}
		if fortuneStatement, err = db.Db.Prepare(FortuneSelect); err != nil {
			revel.ERROR.Fatalln(err)
		}
		if updateStatement, err = db.Db.Prepare(WorldUpdate); err != nil {
			revel.ERROR.Fatalln(err)
		}
	})
}

type App struct {
	*revel.Controller
}

func (c App) Json() revel.Result {
	c.Response.ContentType = "application/json"
	return c.RenderJson(MessageStruct{"Hello, world"})
}

func (c App) Db(queries int) revel.Result {
	if queries <= 1 {
		var w World
		worldStatement.QueryRow(rand.Intn(WorldRowCount)+1).Scan(&w.Id, &w.RandomNumber)
		return c.RenderJson(w)
	}

	ww := make([]World, queries)
	var wg sync.WaitGroup
	wg.Add(queries)
	for i := 0; i < queries; i++ {
		go func(i int) {
			err := worldStatement.QueryRow(rand.Intn(WorldRowCount)+1).
				Scan(&ww[i].Id, &ww[i].RandomNumber)
			if err != nil {
				revel.ERROR.Fatalf("Error scanning world row: %v", err)
			}
			wg.Done()
		}(i)
	}
	wg.Wait()
	return c.RenderJson(ww)
}

func (c App) Update(queries int) revel.Result {
	if queries <= 1 {
		var w World
		worldStatement.QueryRow(rand.Intn(WorldRowCount)+1).Scan(&w.Id, &w.RandomNumber)
		w.RandomNumber = uint16(rand.Intn(WorldRowCount) + 1)
		updateStatement.Exec(w.RandomNumber, w.Id)
		return c.RenderJson(&w)
	}

	var (
		ww = make([]World, queries)
		wg sync.WaitGroup
	)
	wg.Add(queries)
	for i := 0; i < queries; i++ {
		go func(i int) {
			err := worldStatement.QueryRow(rand.Intn(WorldRowCount)+1).
				Scan(&ww[i].Id, &ww[i].RandomNumber)
			if err != nil {
				revel.ERROR.Fatalf("Error scanning world row: %v", err)
			}
			ww[i].RandomNumber = uint16(rand.Intn(WorldRowCount) + 1)
			updateStatement.Exec(ww[i].RandomNumber, ww[i].Id)
			wg.Done()
		}(i)
	}
	wg.Wait()
	return c.RenderJson(ww)
}

func (c App) Fortune() revel.Result {
	fortunes := make([]*Fortune, 0, 16)

	rows, err := fortuneStatement.Query()
	if err != nil {
		revel.ERROR.Fatalf("Error preparing statement: %v", err)
	}

	i := 0
	var fortune *Fortune
	for rows.Next() {
		fortune = new(Fortune)
		if err = rows.Scan(&fortune.Id, &fortune.Message); err != nil {
			revel.ERROR.Fatalf("Error scanning fortune row: %v", err)
		}
		fortunes = append(fortunes, fortune)
		i++
	}
	fortunes = append(fortunes, &Fortune{Message: "Additional fortune added at request time."})

	sort.Sort(ByMessage{fortunes})
	return c.Render(fortunes)
}

type Fortunes []*Fortune

func (s Fortunes) Len() int      { return len(s) }
func (s Fortunes) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

type ByMessage struct{ Fortunes }

func (s ByMessage) Less(i, j int) bool { return s.Fortunes[i].Message < s.Fortunes[j].Message }
<|endoftext|>