<commit_before>package noopraven

import (
	"github.com/getsentry/raven-go"
)

type RavenClient interface {
	Capture(packet *raven.Packet, captureTags map[string]string) (eventID string, ch chan error)
	CaptureError(err error, tags map[string]string, interfaces ...raven.Interface) string
	CaptureMessage(message string, tags map[string]string, interfaces ...raven.Interface) string
	CapturePanic(f func(), tags map[string]string, interfaces ...raven.Interface) (err interface{}, errorID string)
	Close()
	ProjectID() string
	Release() string
	SetDSN(dsn string) error
	SetRelease(release string)
	URL() string
}

type NoopRavenClient struct {
}

func (client *NoopRavenClient) Capture(packet *raven.Packet, captureTags map[string]string) (eventID string, ch chan error) {
	return "", nil
}

func (client *NoopRavenClient) CaptureError(err error, tags map[string]string, interfaces ...raven.Interface) string {
	return ""
}

func (client *NoopRavenClient) CaptureMessage(message string, tags map[string]string, interfaces ...raven.Interface) string {
	return ""
}

func (client *NoopRavenClient) CapturePanic(f func(), tags map[string]string, interfaces ...raven.Interface) (err interface{}, errorID string) {
	f()
	return
}

func (client *NoopRavenClient) Close() {
}

func (client *NoopRavenClient) ProjectID() string {
	return ""
}

func (client *NoopRavenClient) Release() string {
	return ""
}

func (client *NoopRavenClient) SetDSN(dsn string) error {
	return nil
}

func (client *NoopRavenClient) SetRelease(release string) {
}

func (client *NoopRavenClient) URL() string {
	return ""
}
<commit_msg>adding CaptureErrorAndWait<commit_after>package noopraven

import (
	"github.com/getsentry/raven-go"
)

type RavenClient interface {
	Capture(packet *raven.Packet, captureTags map[string]string) (eventID string, ch chan error)
	CaptureError(err error, tags map[string]string, interfaces ...raven.Interface) string
	CaptureErrorAndWait(err error, tags map[string]string, interfaces ...raven.Interface) string
	CaptureMessage(message string, tags map[string]string, interfaces ...raven.Interface) string
	CapturePanic(f func(), tags map[string]string, interfaces ...raven.Interface) (err interface{}, errorID string)
	Close()
	ProjectID() string
	Release() string
	SetDSN(dsn string) error
	SetRelease(release string)
	URL() string
}

type NoopRavenClient struct {
}

func (client *NoopRavenClient) Capture(packet *raven.Packet, captureTags map[string]string) (eventID string, ch chan error) {
	return "", nil
}

func (client *NoopRavenClient) CaptureError(err error, tags map[string]string, interfaces ...raven.Interface) string {
	return ""
}

func (client *NoopRavenClient) CaptureErrorAndWait(err error, tags map[string]string, interfaces ...raven.Interface) string {
	return ""
}

func (client *NoopRavenClient) CaptureMessage(message string, tags map[string]string, interfaces ...raven.Interface) string {
	return ""
}

func (client *NoopRavenClient) CapturePanic(f func(), tags map[string]string, interfaces ...raven.Interface) (err interface{}, errorID string) {
	f()
	return
}

func (client *NoopRavenClient) Close() {
}

func (client *NoopRavenClient) ProjectID() string {
	return ""
}

func (client *NoopRavenClient) Release() string {
	return ""
}

func (client *NoopRavenClient) SetDSN(dsn string) error {
	return nil
}

func (client *NoopRavenClient) SetRelease(release string) {
}

func (client *NoopRavenClient) URL() string {
	return ""
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package tchannel

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/uber/tchannel-go/relay"

	"github.com/uber-go/atomic"
)

// _maxRelayTombs is the maximum number of tombs we'll accumulate in a single
// relayItems.
const _maxRelayTombs = 1e4

// _relayTombTTL is the length of time we'll keep a tomb before GC'ing it.
const _relayTombTTL = time.Second

type relayItem struct {
	*time.Timer

	stats       relay.CallStats
	remapID     uint32
	destination *Relayer
	tomb        bool
	span        Span
}

type relayItems struct {
	sync.RWMutex

	logger Logger
	tombs  uint64
	items  map[uint32]relayItem
}

func newRelayItems(logger Logger) *relayItems {
	return &relayItems{
		items:  make(map[uint32]relayItem),
		logger: logger,
	}
}

// Count returns the number of non-tombstone items in the relay.
func (r *relayItems) Count() int {
	r.RLock()
	n := len(r.items) - int(r.tombs)
	r.RUnlock()
	return n
}

// Get checks for a relay item by ID, returning the item and a bool indicating
// whether the item was found.
func (r *relayItems) Get(id uint32) (relayItem, bool) {
	r.RLock()
	item, ok := r.items[id]
	r.RUnlock()

	return item, ok
}

// Add adds a relay item.
func (r *relayItems) Add(id uint32, item relayItem) {
	r.Lock()
	r.items[id] = item
	r.Unlock()
}

// Delete removes a relayItem completely (without leaving a tombstone). It
// returns the deleted item, along with a bool indicating whether we completed a
// relayed call.
func (r *relayItems) Delete(id uint32) (relayItem, bool) {
	r.Lock()
	item, ok := r.items[id]
	if !ok {
		r.Unlock()
		r.logger.WithFields(LogField{"id", id}).Warn("Attempted to delete non-existent relay item.")
		return item, false
	}
	delete(r.items, id)
	if item.tomb {
		r.tombs--
	}
	r.Unlock()

	item.Stop()
	return item, !item.tomb
}

// Entomb sets the tomb bit on a relayItem and schedules a garbage collection. It
// returns the entombed item, along with a bool indicating whether we completed
// a relayed call.
func (r *relayItems) Entomb(id uint32, deleteAfter time.Duration) (relayItem, bool) {
	r.Lock()
	if r.tombs > _maxRelayTombs {
		r.Unlock()
		r.logger.WithFields(LogField{"id", id}).Warn("Too many tombstones, deleting relay item immediately.")
		return r.Delete(id)
	}
	item, ok := r.items[id]
	if !ok {
		r.Unlock()
		r.logger.WithFields(LogField{"id", id}).Warn("Can't find relay item to entomb.")
		return item, false
	}
	if item.tomb {
		r.Unlock()
		r.logger.WithFields(LogField{"id", id}).Warn("Re-entombing a tombstone.")
		return item, false
	}
	r.tombs++
	item.tomb = true
	r.items[id] = item
	r.Unlock()

	// TODO: We should be clearing these out in batches, rather than creating
	// individual timers for each item.
	time.AfterFunc(deleteAfter, func() { r.Delete(id) })
	return item, true
}

type frameType int

const (
	requestFrame  frameType = 0
	responseFrame frameType = 1
)

// A Relayer forwards frames.
type Relayer struct {
	stats relay.Stats
	hosts relay.Hosts

	// outbound is the remapping for requests that originated on this
	// connection, and are outbound towards some other connection.
	// It stores remappings for all request frames read on this connection.
	outbound *relayItems

	// inbound is the remapping for requests that originated on some other
	// connection which was directed to this connection.
	// It stores remappings for all response frames read on this connection.
	inbound *relayItems

	peers   *PeerList
	conn    *Connection
	logger  Logger
	pending atomic.Uint32
}

// NewRelayer constructs a Relayer.
func NewRelayer(ch *Channel, conn *Connection) *Relayer {
	return &Relayer{
		stats:    ch.relayStats,
		hosts:    ch.RelayHosts(),
		outbound: newRelayItems(ch.Logger().WithFields(LogField{"relay", "outbound"})),
		inbound:  newRelayItems(ch.Logger().WithFields(LogField{"relay", "inbound"})),
		peers:    ch.Peers(),
		conn:     conn,
		logger:   conn.log,
	}
}

// Hosts returns the RelayHosts guiding peer selection.
func (r *Relayer) Hosts() relay.Hosts {
	return r.hosts
}

// Relay is called for each frame that is read on the connection.
func (r *Relayer) Relay(f *Frame) error {
	if f.messageType() != messageTypeCallReq {
		return r.handleNonCallReq(f)
	}
	return r.handleCallReq(newLazyCallReq(f))
}

// Receive receives frames intended for this connection.
func (r *Relayer) Receive(f *Frame, fType frameType) {
	id := f.Header.ID

	// If we receive a response frame, we expect to find that ID in our outbound.
	// If we receive a request frame, we expect to find that ID in our inbound.
	items := r.receiverItems(fType)

	item, ok := items.Get(id)
	if !ok {
		r.logger.WithFields(
			LogField{"id", id},
		).Warn("Received a frame without a RelayItem.")
		return
	}

	// call res frames don't include the OK bit, so we can't wait until the last
	// frame of a relayed RPC to determine if the call succeeded.
	if fType == responseFrame {
		// If we've gotten a response frame, we're the originating relayer and
		// should handle stats.
		if succeeded, failMsg := determinesCallSuccess(f); succeeded {
			item.stats.Succeeded()
		} else if len(failMsg) > 0 {
			item.stats.Failed(failMsg)
		}
	}

	// When we write the frame to sendCh, we lose ownership of the frame, and it
	// may be released to the frame pool at any point.
	finished := finishesCall(f)

	// TODO: Add some sort of timeout here to avoid blocking forever on a
	// stalled connection.
	r.conn.sendCh <- f

	if finished {
		items := r.receiverItems(fType)
		r.finishRelayItem(items, id)
	}
}

func (r *Relayer) canHandleNewCall() bool {
	var canHandle bool
	r.conn.withStateRLock(func() error {
		canHandle = r.conn.state == connectionActive
		if canHandle {
			r.pending.Inc()
		}
		return nil
	})
	return canHandle
}

func (r *Relayer) getDestination(f lazyCallReq, cs relay.CallStats) (*Connection, bool, error) {
	if _, ok := r.outbound.Get(f.Header.ID); ok {
		r.logger.WithFields(LogField{"id", f.Header.ID}).Warn("received duplicate callReq")
		cs.Failed("relay-" + ErrCodeProtocol.MetricsKey())
		// TODO: this is a protocol error, kill the connection.
		return nil, false, errors.New("callReq with already active ID")
	}

	// Get the destination
	selectedPeer, err := r.hosts.Get(f)
	if err == nil && selectedPeer.HostPort == "" {
		err = errInvalidPeerForGroup(f.Service())
	}
	if err != nil {
		if _, ok := err.(SystemError); !ok {
			err = NewSystemError(ErrCodeDeclined, err.Error())
		}
		cs.Failed("relay-" + GetSystemErrorCode(err).MetricsKey())
		r.conn.SendSystemError(f.Header.ID, f.Span(), err)
		return nil, false, nil
	}

	cs.SetPeer(selectedPeer)
	peer := r.peers.GetOrAdd(selectedPeer.HostPort)

	// TODO: Should connections use the call timeout? Or a separate timeout?
	remoteConn, err := peer.getConnectionTimeout(f.TTL())
	if err != nil {
		r.logger.WithFields(
			ErrField(err),
			LogField{"selectedPeer", selectedPeer},
		).Warn("Failed to connect to relay host.")
		cs.Failed("relay-connection-failed")
		r.conn.SendSystemError(f.Header.ID, f.Span(), NewWrappedSystemError(ErrCodeNetwork, err))
		return nil, false, nil
	}

	return remoteConn, true, nil
}

func (r *Relayer) handleCallReq(f lazyCallReq) error {
	callStats := r.stats.Begin(f)
	if !r.canHandleNewCall() {
		callStats.Failed("relay-channel-closed")
		callStats.End()
		return ErrChannelClosed
	}

	// Get a remote connection and check whether it can handle this call.
	remoteConn, ok, err := r.getDestination(f, callStats)
	if err == nil && ok {
		if !remoteConn.relay.canHandleNewCall() {
			err = NewSystemError(ErrCodeNetwork, "selected closed connection, retry")
			callStats.Failed("relay-connection-closed")
		}
	}
	if err != nil || !ok {
		// Failed to get a remote connection, or the connection is not in the right
		// state to handle this call. Since we already incremented pending on
		// the current relay, we need to decrement it.
		r.decrementPending()
		callStats.End()
		return err
	}

	destinationID := remoteConn.NextMessageID()
	ttl := f.TTL()
	span := f.Span()
	// The remote side of the relay doesn't need to track stats.
	remoteConn.relay.addRelayItem(false /* isOriginator */, destinationID, f.Header.ID, r, ttl, span, nil)
	relayToDest := r.addRelayItem(true /* isOriginator */, f.Header.ID, destinationID, remoteConn.relay, ttl, span, callStats)

	f.Header.ID = destinationID
	relayToDest.destination.Receive(f.Frame, requestFrame)
	return nil
}

// Handle all frames except messageTypeCallReq.
func (r *Relayer) handleNonCallReq(f *Frame) error {
	frameType := frameTypeFor(f)

	// If we read a request frame, we need to use the outbound map to decide
	// the destination. Otherwise, we use the inbound map.
	items := r.outbound
	if frameType == responseFrame {
		items = r.inbound
	}

	item, ok := items.Get(f.Header.ID)
	if !ok {
		return errors.New("non-callReq for inactive ID")
	}
	if item.tomb {
		// Call timed out, ignore this frame. (We've already handled stats.)
		// TODO: metrics for late-arriving frames.
		return nil
	}
	originalID := f.Header.ID
	f.Header.ID = item.remapID
	item.destination.Receive(f, frameType)

	if finishesCall(f) {
		r.finishRelayItem(items, originalID)
	}
	return nil
}

// addRelayItem adds a relay item to either outbound or inbound.
func (r *Relayer) addRelayItem(isOriginator bool, id, remapID uint32, destination *Relayer, ttl time.Duration, span Span, cs relay.CallStats) relayItem {
	item := relayItem{
		stats:       cs,
		remapID:     remapID,
		destination: destination,
		span:        span,
	}

	items := r.inbound
	if isOriginator {
		items = r.outbound
	}
	item.Timer = time.AfterFunc(ttl, func() { r.timeoutRelayItem(isOriginator, items, id) })
	items.Add(id, item)
	return item
}

func (r *Relayer) timeoutRelayItem(isOriginator bool, items *relayItems, id uint32) {
	item, ok := items.Entomb(id, _relayTombTTL)
	if !ok {
		return
	}
	if isOriginator {
		r.conn.SendSystemError(id, item.span, ErrTimeout)
		item.stats.Failed("timeout")
		item.stats.End()
	}

	r.decrementPending()
}

func (r *Relayer) finishRelayItem(items *relayItems, id uint32) {
	item, ok := items.Delete(id)
	if !ok {
		return
	}
	if item.stats != nil {
		item.stats.End()
	}
	r.decrementPending()
}

func (r *Relayer) decrementPending() {
	r.pending.Dec()
	r.conn.checkExchanges()
}

func (r *Relayer) canClose() bool {
	if r == nil {
		return true
	}
	return r.countPending() == 0
}

func (r *Relayer) countPending() uint32 {
	return r.pending.Load()
}

func (r *Relayer) receiverItems(fType frameType) *relayItems {
	if fType == requestFrame {
		return r.inbound
	}
	return r.outbound
}

func frameTypeFor(f *Frame) frameType {
	switch t := f.Header.messageType; t {
	case messageTypeCallRes, messageTypeCallResContinue, messageTypeError, messageTypePingRes:
		return responseFrame
	case messageTypeCallReq, messageTypeCallReqContinue, messageTypePingReq:
		return requestFrame
	default:
		panic(fmt.Sprintf("unsupported frame type: %v", t))
	}
}

func errInvalidPeerForGroup(group []byte) error {
	return NewSystemError(ErrCodeDeclined, "invalid peer for %q", group)
}

func determinesCallSuccess(f *Frame) (succeeded bool, failMsg string) {
	switch f.messageType() {
	case messageTypeError:
		msg := newLazyError(f).Code().MetricsKey()
		return false, msg
	case messageTypeCallRes:
		if newLazyCallRes(f).OK() {
			return true, ""
		}
		return false, "application-error"
	default:
		return false, ""
	}
}
<commit_msg>relay: Fix use-after-release on finishesCall(frame)<commit_after>// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package tchannel

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/uber/tchannel-go/relay"

	"github.com/uber-go/atomic"
)

// _maxRelayTombs is the maximum number of tombs we'll accumulate in a single
// relayItems.
const _maxRelayTombs = 1e4

// _relayTombTTL is the length of time we'll keep a tomb before GC'ing it.
const _relayTombTTL = time.Second

type relayItem struct {
	*time.Timer

	stats       relay.CallStats
	remapID     uint32
	destination *Relayer
	tomb        bool
	span        Span
}

type relayItems struct {
	sync.RWMutex

	logger Logger
	tombs  uint64
	items  map[uint32]relayItem
}

func newRelayItems(logger Logger) *relayItems {
	return &relayItems{
		items:  make(map[uint32]relayItem),
		logger: logger,
	}
}

// Count returns the number of non-tombstone items in the relay.
func (r *relayItems) Count() int {
	r.RLock()
	n := len(r.items) - int(r.tombs)
	r.RUnlock()
	return n
}

// Get checks for a relay item by ID, returning the item and a bool indicating
// whether the item was found.
func (r *relayItems) Get(id uint32) (relayItem, bool) {
	r.RLock()
	item, ok := r.items[id]
	r.RUnlock()

	return item, ok
}

// Add adds a relay item.
func (r *relayItems) Add(id uint32, item relayItem) {
	r.Lock()
	r.items[id] = item
	r.Unlock()
}

// Delete removes a relayItem completely (without leaving a tombstone). It
// returns the deleted item, along with a bool indicating whether we completed a
// relayed call.
func (r *relayItems) Delete(id uint32) (relayItem, bool) {
	r.Lock()
	item, ok := r.items[id]
	if !ok {
		r.Unlock()
		r.logger.WithFields(LogField{"id", id}).Warn("Attempted to delete non-existent relay item.")
		return item, false
	}
	delete(r.items, id)
	if item.tomb {
		r.tombs--
	}
	r.Unlock()

	item.Stop()
	return item, !item.tomb
}

// Entomb sets the tomb bit on a relayItem and schedules a garbage collection. It
// returns the entombed item, along with a bool indicating whether we completed
// a relayed call.
func (r *relayItems) Entomb(id uint32, deleteAfter time.Duration) (relayItem, bool) {
	r.Lock()
	if r.tombs > _maxRelayTombs {
		r.Unlock()
		r.logger.WithFields(LogField{"id", id}).Warn("Too many tombstones, deleting relay item immediately.")
		return r.Delete(id)
	}
	item, ok := r.items[id]
	if !ok {
		r.Unlock()
		r.logger.WithFields(LogField{"id", id}).Warn("Can't find relay item to entomb.")
		return item, false
	}
	if item.tomb {
		r.Unlock()
		r.logger.WithFields(LogField{"id", id}).Warn("Re-entombing a tombstone.")
		return item, false
	}
	r.tombs++
	item.tomb = true
	r.items[id] = item
	r.Unlock()

	// TODO: We should be clearing these out in batches, rather than creating
	// individual timers for each item.
	time.AfterFunc(deleteAfter, func() { r.Delete(id) })
	return item, true
}

type frameType int

const (
	requestFrame  frameType = 0
	responseFrame frameType = 1
)

// A Relayer forwards frames.
type Relayer struct {
	stats relay.Stats
	hosts relay.Hosts

	// outbound is the remapping for requests that originated on this
	// connection, and are outbound towards some other connection.
	// It stores remappings for all request frames read on this connection.
	outbound *relayItems

	// inbound is the remapping for requests that originated on some other
	// connection which was directed to this connection.
	// It stores remappings for all response frames read on this connection.
	inbound *relayItems

	peers   *PeerList
	conn    *Connection
	logger  Logger
	pending atomic.Uint32
}

// NewRelayer constructs a Relayer.
func NewRelayer(ch *Channel, conn *Connection) *Relayer {
	return &Relayer{
		stats:    ch.relayStats,
		hosts:    ch.RelayHosts(),
		outbound: newRelayItems(ch.Logger().WithFields(LogField{"relay", "outbound"})),
		inbound:  newRelayItems(ch.Logger().WithFields(LogField{"relay", "inbound"})),
		peers:    ch.Peers(),
		conn:     conn,
		logger:   conn.log,
	}
}

// Hosts returns the RelayHosts guiding peer selection.
func (r *Relayer) Hosts() relay.Hosts {
	return r.hosts
}

// Relay is called for each frame that is read on the connection.
func (r *Relayer) Relay(f *Frame) error {
	if f.messageType() != messageTypeCallReq {
		return r.handleNonCallReq(f)
	}
	return r.handleCallReq(newLazyCallReq(f))
}

// Receive receives frames intended for this connection.
func (r *Relayer) Receive(f *Frame, fType frameType) {
	id := f.Header.ID

	// If we receive a response frame, we expect to find that ID in our outbound.
	// If we receive a request frame, we expect to find that ID in our inbound.
	items := r.receiverItems(fType)

	item, ok := items.Get(id)
	if !ok {
		r.logger.WithFields(
			LogField{"id", id},
		).Warn("Received a frame without a RelayItem.")
		return
	}

	// call res frames don't include the OK bit, so we can't wait until the last
	// frame of a relayed RPC to determine if the call succeeded.
	if fType == responseFrame {
		// If we've gotten a response frame, we're the originating relayer and
		// should handle stats.
		if succeeded, failMsg := determinesCallSuccess(f); succeeded {
			item.stats.Succeeded()
		} else if len(failMsg) > 0 {
			item.stats.Failed(failMsg)
		}
	}

	// When we write the frame to sendCh, we lose ownership of the frame, and it
	// may be released to the frame pool at any point.
	finished := finishesCall(f)

	// TODO: Add some sort of timeout here to avoid blocking forever on a
	// stalled connection.
	r.conn.sendCh <- f

	if finished {
		items := r.receiverItems(fType)
		r.finishRelayItem(items, id)
	}
}

func (r *Relayer) canHandleNewCall() bool {
	var canHandle bool
	r.conn.withStateRLock(func() error {
		canHandle = r.conn.state == connectionActive
		if canHandle {
			r.pending.Inc()
		}
		return nil
	})
	return canHandle
}

func (r *Relayer) getDestination(f lazyCallReq, cs relay.CallStats) (*Connection, bool, error) {
	if _, ok := r.outbound.Get(f.Header.ID); ok {
		r.logger.WithFields(LogField{"id", f.Header.ID}).Warn("received duplicate callReq")
		cs.Failed("relay-" + ErrCodeProtocol.MetricsKey())
		// TODO: this is a protocol error, kill the connection.
		return nil, false, errors.New("callReq with already active ID")
	}

	// Get the destination
	selectedPeer, err := r.hosts.Get(f)
	if err == nil && selectedPeer.HostPort == "" {
		err = errInvalidPeerForGroup(f.Service())
	}
	if err != nil {
		if _, ok := err.(SystemError); !ok {
			err = NewSystemError(ErrCodeDeclined, err.Error())
		}
		cs.Failed("relay-" + GetSystemErrorCode(err).MetricsKey())
		r.conn.SendSystemError(f.Header.ID, f.Span(), err)
		return nil, false, nil
	}

	cs.SetPeer(selectedPeer)
	peer := r.peers.GetOrAdd(selectedPeer.HostPort)

	// TODO: Should connections use the call timeout? Or a separate timeout?
	remoteConn, err := peer.getConnectionTimeout(f.TTL())
	if err != nil {
		r.logger.WithFields(
			ErrField(err),
			LogField{"selectedPeer", selectedPeer},
		).Warn("Failed to connect to relay host.")
		cs.Failed("relay-connection-failed")
		r.conn.SendSystemError(f.Header.ID, f.Span(), NewWrappedSystemError(ErrCodeNetwork, err))
		return nil, false, nil
	}

	return remoteConn, true, nil
}

func (r *Relayer) handleCallReq(f lazyCallReq) error {
	callStats := r.stats.Begin(f)
	if !r.canHandleNewCall() {
		callStats.Failed("relay-channel-closed")
		callStats.End()
		return ErrChannelClosed
	}

	// Get a remote connection and check whether it can handle this call.
	remoteConn, ok, err := r.getDestination(f, callStats)
	if err == nil && ok {
		if !remoteConn.relay.canHandleNewCall() {
			err = NewSystemError(ErrCodeNetwork, "selected closed connection, retry")
			callStats.Failed("relay-connection-closed")
		}
	}
	if err != nil || !ok {
		// Failed to get a remote connection, or the connection is not in the right
		// state to handle this call. Since we already incremented pending on
		// the current relay, we need to decrement it.
		r.decrementPending()
		callStats.End()
		return err
	}

	destinationID := remoteConn.NextMessageID()
	ttl := f.TTL()
	span := f.Span()
	// The remote side of the relay doesn't need to track stats.
	remoteConn.relay.addRelayItem(false /* isOriginator */, destinationID, f.Header.ID, r, ttl, span, nil)
	relayToDest := r.addRelayItem(true /* isOriginator */, f.Header.ID, destinationID, remoteConn.relay, ttl, span, callStats)

	f.Header.ID = destinationID
	relayToDest.destination.Receive(f.Frame, requestFrame)
	return nil
}

// Handle all frames except messageTypeCallReq.
func (r *Relayer) handleNonCallReq(f *Frame) error {
	frameType := frameTypeFor(f)

	// If we read a request frame, we need to use the outbound map to decide
	// the destination. Otherwise, we use the inbound map.
	items := r.outbound
	if frameType == responseFrame {
		items = r.inbound
	}

	item, ok := items.Get(f.Header.ID)
	if !ok {
		return errors.New("non-callReq for inactive ID")
	}
	if item.tomb {
		// Call timed out, ignore this frame. (We've already handled stats.)
		// TODO: metrics for late-arriving frames.
		return nil
	}
	originalID := f.Header.ID
	f.Header.ID = item.remapID

	// Once we call Receive on the frame, we lose ownership of the frame.
	finished := finishesCall(f)

	item.destination.Receive(f, frameType)

	if finished {
		r.finishRelayItem(items, originalID)
	}
	return nil
}

// addRelayItem adds a relay item to either outbound or inbound.
func (r *Relayer) addRelayItem(isOriginator bool, id, remapID uint32, destination *Relayer, ttl time.Duration, span Span, cs relay.CallStats) relayItem {
	item := relayItem{
		stats:       cs,
		remapID:     remapID,
		destination: destination,
		span:        span,
	}

	items := r.inbound
	if isOriginator {
		items = r.outbound
	}
	item.Timer = time.AfterFunc(ttl, func() { r.timeoutRelayItem(isOriginator, items, id) })
	items.Add(id, item)
	return item
}

func (r *Relayer) timeoutRelayItem(isOriginator bool, items *relayItems, id uint32) {
	item, ok := items.Entomb(id, _relayTombTTL)
	if !ok {
		return
	}
	if isOriginator {
		r.conn.SendSystemError(id, item.span, ErrTimeout)
		item.stats.Failed("timeout")
		item.stats.End()
	}

	r.decrementPending()
}

func (r *Relayer) finishRelayItem(items *relayItems, id uint32) {
	item, ok := items.Delete(id)
	if !ok {
		return
	}
	if item.stats != nil {
		item.stats.End()
	}
	r.decrementPending()
}

func (r *Relayer) decrementPending() {
	r.pending.Dec()
	r.conn.checkExchanges()
}

func (r *Relayer) canClose() bool {
	if r == nil {
		return true
	}
	return r.countPending() == 0
}

func (r *Relayer) countPending() uint32 {
	return r.pending.Load()
}

func (r *Relayer) receiverItems(fType frameType) *relayItems {
	if fType == requestFrame {
		return r.inbound
	}
	return r.outbound
}

func frameTypeFor(f *Frame) frameType {
	switch t := f.Header.messageType; t {
	case messageTypeCallRes, messageTypeCallResContinue, messageTypeError, messageTypePingRes:
		return responseFrame
	case messageTypeCallReq, messageTypeCallReqContinue, messageTypePingReq:
		return requestFrame
	default:
		panic(fmt.Sprintf("unsupported frame type: %v", t))
	}
}

func errInvalidPeerForGroup(group []byte) error {
	return NewSystemError(ErrCodeDeclined, "invalid peer for %q", group)
}

func determinesCallSuccess(f *Frame) (succeeded bool, failMsg string) {
	switch f.messageType() {
	case messageTypeError:
		msg := newLazyError(f).Code().MetricsKey()
		return false, msg
	case messageTypeCallRes:
		if newLazyCallRes(f).OK() {
			return true, ""
		}
		return false, "application-error"
	default:
		return false, ""
	}
}
<|endoftext|>


<commit_before>/*
 * Package contains functions to handle eGauge connections and network communication.
 */
package eGaugeHandlers

import (
	"log"
	"net/http"

	"github.com/seadsystem/Backend/DB/landingzone/database"
)

func HandleRequest(res http.ResponseWriter, req *http.Request, db database.DB) {
	log.Println("Got request from:", req.RemoteAddr)
	if req.Method == "POST" {
		log.Printf("Request from (%s) is of type POST\n", req.RemoteAddr)
		log.Printf("%+v\n", req.PostForm)
	}
}
<commit_msg>Added full request dumping.<commit_after>/*
 * Package contains functions to handle eGauge connections and network communication.
 */
package eGaugeHandlers

import (
	"log"
	"net/http"
	"net/http/httputil"

	"github.com/seadsystem/Backend/DB/landingzone/database"
)

func HandleRequest(res http.ResponseWriter, req *http.Request, db database.DB) {
	log.Println("Got request from:", req.RemoteAddr)
	if req.Method == "POST" {
		log.Printf("Request from (%s) is of type POST\n", req.RemoteAddr)
	}

	data, err := httputil.DumpRequestOut(req, true)
	if err != nil {
		log.Println("Error dumping request:", err)
	} else {
		log.Println(data)
	}
}
<|endoftext|>


<commit_before>package regserv

import (
	"testing"

	"github.com/dgrijalva/jwt-go"
	"github.com/koding/kite/kitekey"
	"github.com/koding/kite/testkeys"
	"github.com/koding/kite/testutil"
)

func TestRegister(t *testing.T) {
	regserv := New(testutil.NewConfig(), testkeys.Public, testkeys.Private)

	key, err := regserv.register("foo", "bar")
	if err != nil {
		t.Errorf(err.Error())
		return
	}

	token, err := jwt.Parse(key, kitekey.GetKontrolKey)
	if err != nil {
		t.Errorf(err.Error())
		return
	}

	if username := token.Claims["sub"].(string); username != "foo" {
		t.Errorf("invalid username: %s", username)
		return
	}

	if hostname := token.Claims["aud"].(string); hostname != "bar" {
		t.Errorf("invalid hostname: %s", hostname)
		return
	}
}
<commit_msg>fix regserv test<commit_after>package regserv

import (
	"testing"

	"github.com/dgrijalva/jwt-go"
	"github.com/koding/kite/kitekey"
	"github.com/koding/kite/testkeys"
	"github.com/koding/kite/testutil"
)

func TestRegister(t *testing.T) {
	regserv := New(testutil.NewConfig(), testkeys.Public, testkeys.Private)

	key, err := regserv.register("foo")
	if err != nil {
		t.Errorf(err.Error())
		return
	}

	token, err := jwt.Parse(key, kitekey.GetKontrolKey)
	if err != nil {
		t.Errorf(err.Error())
		return
	}

	if username := token.Claims["sub"].(string); username != "foo" {
		t.Errorf("invalid username: %s", username)
		return
	}
}
<|endoftext|>


<commit_before>/*
 * Minio Go Library for Amazon S3 Compatible Cloud Storage
 * Copyright 2015-2017 Minio, Inc.
 *
 * 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 minio

import (
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"
)

// MaxRetry is the maximum number of retries before stopping.
var MaxRetry = 10

// MaxJitter will randomize over the full exponential backoff time
const MaxJitter = 1.0

// NoJitter disables the use of jitter for randomizing the exponential backoff time
const NoJitter = 0.0

// DefaultRetryUnit - default unit multiplicative per retry.
// defaults to 1 second.
const DefaultRetryUnit = time.Second

// DefaultRetryCap - Each retry attempt never waits no longer than
// this maximum time duration.
const DefaultRetryCap = time.Second * 30

// newRetryTimer creates a timer with exponentially increasing
// delays until the maximum retry attempts are reached.
func (c Client) newRetryTimer(maxRetry int, unit time.Duration, cap time.Duration, jitter float64, doneCh chan struct{}) <-chan int {
	attemptCh := make(chan int)

	// computes the exponential backoff duration according to
	// https://www.awsarchitectureblog.com/2015/03/backoff.html
	exponentialBackoffWait := func(attempt int) time.Duration {
		// normalize jitter to the range [0, 1.0]
		if jitter < NoJitter {
			jitter = NoJitter
		}
		if jitter > MaxJitter {
			jitter = MaxJitter
		}

		//sleep = random_between(0, min(cap, base * 2 ** attempt))
		sleep := unit * time.Duration(1<<uint(attempt))
		if sleep > cap {
			sleep = cap
		}
		if jitter != NoJitter {
			sleep -= time.Duration(c.random.Float64() * float64(sleep) * jitter)
		}
		return sleep
	}

	go func() {
		defer close(attemptCh)
		for i := 0; i < maxRetry; i++ {
			select {
			// Attempts start from 1.
			case attemptCh <- i + 1:
			case <-doneCh:
				// Stop the routine.
				return
			}
			time.Sleep(exponentialBackoffWait(i))
		}
	}()
	return attemptCh
}

// isNetErrorRetryable - is network error retryable.
func isNetErrorRetryable(err error) bool {
	if err == nil {
		return false
	}
	switch err.(type) {
	case net.Error:
		switch err.(type) {
		case *net.DNSError, *net.OpError, net.UnknownNetworkError:
			return true
		case *url.Error:
			// For a URL error, where it replies back "connection closed"
			// retry again.
			if strings.Contains(err.Error(), "Connection closed by foreign host") {
				return true
			}
		default:
			if strings.Contains(err.Error(), "net/http: TLS handshake timeout") {
				// If error is - tlsHandshakeTimeoutError, retry.
				return true
			} else if strings.Contains(err.Error(), "i/o timeout") {
				// If error is - tcp timeoutError, retry.
				return true
			} else if strings.Contains(err.Error(), "connection timed out") {
				// If err is a net.Dial timeout, retry.
				return true
			} else if strings.Contains(err.Error(), "net/http: HTTP/1.x transport connection broken") {
				// If error is transport connection broken, retry.
				return true
			}
		}
	}
	return false
}

// List of AWS S3 error codes which are retryable.
var retryableS3Codes = map[string]struct{}{
	"RequestError":          {},
	"RequestTimeout":        {},
	"Throttling":            {},
	"ThrottlingException":   {},
	"RequestLimitExceeded":  {},
	"RequestThrottled":      {},
	"InternalError":         {},
	"ExpiredToken":          {},
	"ExpiredTokenException": {},
	// Add more AWS S3 codes here.
}

// isS3CodeRetryable - is s3 error code retryable.
func isS3CodeRetryable(s3Code string) (ok bool) {
	_, ok = retryableS3Codes[s3Code]
	return ok
}

// List of HTTP status codes which are retryable.
var retryableHTTPStatusCodes = map[int]struct{}{
	429: {}, // http.StatusTooManyRequests is not part of the Go 1.5 library, yet
	http.StatusInternalServerError: {},
	http.StatusBadGateway:          {},
	http.StatusServiceUnavailable:  {},
	// Add more HTTP status codes here.
}

// isHTTPStatusRetryable - is HTTP error code retryable.
func isHTTPStatusRetryable(httpStatusCode int) (ok bool) {
	_, ok = retryableHTTPStatusCodes[httpStatusCode]
	return ok
}
<commit_msg>Added SlowDown in retryableS3Codes #971 (#972)<commit_after>/*
 * Minio Go Library for Amazon S3 Compatible Cloud Storage
 * Copyright 2015-2017 Minio, Inc.
 *
 * 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 minio

import (
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"
)

// MaxRetry is the maximum number of retries before stopping.
var MaxRetry = 10

// MaxJitter will randomize over the full exponential backoff time
const MaxJitter = 1.0

// NoJitter disables the use of jitter for randomizing the exponential backoff time
const NoJitter = 0.0

// DefaultRetryUnit - default unit multiplicative per retry.
// defaults to 1 second.
const DefaultRetryUnit = time.Second

// DefaultRetryCap - Each retry attempt never waits no longer than
// this maximum time duration.
const DefaultRetryCap = time.Second * 30

// newRetryTimer creates a timer with exponentially increasing
// delays until the maximum retry attempts are reached.
func (c Client) newRetryTimer(maxRetry int, unit time.Duration, cap time.Duration, jitter float64, doneCh chan struct{}) <-chan int {
	attemptCh := make(chan int)

	// computes the exponential backoff duration according to
	// https://www.awsarchitectureblog.com/2015/03/backoff.html
	exponentialBackoffWait := func(attempt int) time.Duration {
		// normalize jitter to the range [0, 1.0]
		if jitter < NoJitter {
			jitter = NoJitter
		}
		if jitter > MaxJitter {
			jitter = MaxJitter
		}

		//sleep = random_between(0, min(cap, base * 2 ** attempt))
		sleep := unit * time.Duration(1<<uint(attempt))
		if sleep > cap {
			sleep = cap
		}
		if jitter != NoJitter {
			sleep -= time.Duration(c.random.Float64() * float64(sleep) * jitter)
		}
		return sleep
	}

	go func() {
		defer close(attemptCh)
		for i := 0; i < maxRetry; i++ {
			select {
			// Attempts start from 1.
			case attemptCh <- i + 1:
			case <-doneCh:
				// Stop the routine.
				return
			}
			time.Sleep(exponentialBackoffWait(i))
		}
	}()
	return attemptCh
}

// isNetErrorRetryable - is network error retryable.
func isNetErrorRetryable(err error) bool {
	if err == nil {
		return false
	}
	switch err.(type) {
	case net.Error:
		switch err.(type) {
		case *net.DNSError, *net.OpError, net.UnknownNetworkError:
			return true
		case *url.Error:
			// For a URL error, where it replies back "connection closed"
			// retry again.
			if strings.Contains(err.Error(), "Connection closed by foreign host") {
				return true
			}
		default:
			if strings.Contains(err.Error(), "net/http: TLS handshake timeout") {
				// If error is - tlsHandshakeTimeoutError, retry.
				return true
			} else if strings.Contains(err.Error(), "i/o timeout") {
				// If error is - tcp timeoutError, retry.
				return true
			} else if strings.Contains(err.Error(), "connection timed out") {
				// If err is a net.Dial timeout, retry.
				return true
			} else if strings.Contains(err.Error(), "net/http: HTTP/1.x transport connection broken") {
				// If error is transport connection broken, retry.
				return true
			}
		}
	}
	return false
}

// List of AWS S3 error codes which are retryable.
var retryableS3Codes = map[string]struct{}{
	"RequestError":          {},
	"RequestTimeout":        {},
	"Throttling":            {},
	"ThrottlingException":   {},
	"RequestLimitExceeded":  {},
	"RequestThrottled":      {},
	"InternalError":         {},
	"ExpiredToken":          {},
	"ExpiredTokenException": {},
	"SlowDown":              {},
	// Add more AWS S3 codes here.
}

// isS3CodeRetryable - is s3 error code retryable.
func isS3CodeRetryable(s3Code string) (ok bool) {
	_, ok = retryableS3Codes[s3Code]
	return ok
}

// List of HTTP status codes which are retryable.
var retryableHTTPStatusCodes = map[int]struct{}{
	429: {}, // http.StatusTooManyRequests is not part of the Go 1.5 library, yet
	http.StatusInternalServerError: {},
	http.StatusBadGateway:          {},
	http.StatusServiceUnavailable:  {},
	// Add more HTTP status codes here.
}

// isHTTPStatusRetryable - is HTTP error code retryable.
func isHTTPStatusRetryable(httpStatusCode int) (ok bool) {
	_, ok = retryableHTTPStatusCodes[httpStatusCode]
	return ok
}
<|endoftext|>


<commit_before>package common

import (
	"strconv"
	"strings"
	"text/scanner"

	"github.com/neelance/graphql-go/errors"
)

type Literal interface {
	Value(vars map[string]interface{}) interface{}
	String() string
	Location() errors.Location
}

type BasicLit struct {
	Type rune
	Text string
	Loc  errors.Location
}

func (lit *BasicLit) Value(vars map[string]interface{}) interface{} {
	switch lit.Type {
	case scanner.Int, scanner.Float:
		value, err := strconv.ParseFloat(lit.Text, 64)
		if err != nil {
			panic(err)
		}
		return value

	case scanner.String:
		value, err := strconv.Unquote(lit.Text)
		if err != nil {
			panic(err)
		}
		return value

	case scanner.Ident:
		switch lit.Text {
		case "true":
			return true
		case "false":
			return false
		default:
			return lit.Text
		}

	default:
		panic("invalid literal")
	}
}

func (lit *BasicLit) String() string {
	return lit.Text
}

func (lit *BasicLit) Location() errors.Location {
	return lit.Loc
}

type ListLit struct {
	Entries []Literal
	Loc     errors.Location
}

func (lit *ListLit) Value(vars map[string]interface{}) interface{} {
	entries := make([]interface{}, len(lit.Entries))
	for i, entry := range lit.Entries {
		entries[i] = entry.Value(vars)
	}
	return entries
}

func (lit *ListLit) String() string {
	entries := make([]string, len(lit.Entries))
	for i, entry := range lit.Entries {
		entries[i] = entry.String()
	}
	return "[" + strings.Join(entries, ", ") + "]"
}

func (lit *ListLit) Location() errors.Location {
	return lit.Loc
}

type ObjectLit struct {
	Fields []*ObjectLitField
	Loc    errors.Location
}

type ObjectLitField struct {
	Name  Ident
	Value Literal
}

func (lit *ObjectLit) Value(vars map[string]interface{}) interface{} {
	fields := make(map[string]interface{}, len(lit.Fields))
	for _, f := range lit.Fields {
		fields[f.Name.Name] = f.Value.Value(vars)
	}
	return fields
}

func (lit *ObjectLit) String() string {
	entries := make([]string, 0, len(lit.Fields))
	for _, f := range lit.Fields {
		entries = append(entries, f.Name.Name+": "+f.Value.String())
	}
	return "{" + strings.Join(entries, ", ") + "}"
}

func (lit *ObjectLit) Location() errors.Location {
	return lit.Loc
}

type NullLit struct {
	Loc errors.Location
}

func (lit *NullLit) Value(vars map[string]interface{}) interface{} {
	return nil
}

func (lit *NullLit) String() string {
	return "null"
}

func (lit *NullLit) Location() errors.Location {
	return lit.Loc
}

type Variable struct {
	Name string
	Loc  errors.Location
}

func (v Variable) Value(vars map[string]interface{}) interface{} {
	return vars[v.Name]
}

func (v Variable) String() string {
	return "$" + v.Name
}

func (v *Variable) Location() errors.Location {
	return v.Loc
}

func ParseLiteral(l *Lexer, constOnly bool) Literal {
	loc := l.Location()
	switch l.Peek() {
	case '$':
		if constOnly {
			l.SyntaxError("variable not allowed")
			panic("unreachable")
		}
		l.ConsumeToken('$')
		return &Variable{l.ConsumeIdent(), loc}

	case scanner.Int, scanner.Float, scanner.String, scanner.Ident:
		lit := l.ConsumeLiteral()
		if lit.Type == scanner.Ident && lit.Text == "null" {
			return &NullLit{loc}
		}
		lit.Loc = loc
		return lit
	case '-':
		l.ConsumeToken('-')
		lit := l.ConsumeLiteral()
		lit.Text = "-" + lit.Text
		lit.Loc = loc
		return lit
	case '[':
		l.ConsumeToken('[')
		var list []Literal
		for l.Peek() != ']' {
			list = append(list, ParseLiteral(l, constOnly))
		}
		l.ConsumeToken(']')
		return &ListLit{list, loc}

	case '{':
		l.ConsumeToken('{')
		var fields []*ObjectLitField
		for l.Peek() != '}' {
			name := l.ConsumeIdentWithLoc()
			l.ConsumeToken(':')
			value := ParseLiteral(l, constOnly)
			fields = append(fields, &ObjectLitField{name, value})
		}
		l.ConsumeToken('}')
		return &ObjectLit{fields, loc}

	default:
		l.SyntaxError("invalid value")
		panic("unreachable")
	}
}
<commit_msg>Separate literal arg parsing cases (int, float)<commit_after>package common

import (
	"strconv"
	"strings"
	"text/scanner"

	"github.com/neelance/graphql-go/errors"
)

type Literal interface {
	Value(vars map[string]interface{}) interface{}
	String() string
	Location() errors.Location
}

type BasicLit struct {
	Type rune
	Text string
	Loc  errors.Location
}

func (lit *BasicLit) Value(vars map[string]interface{}) interface{} {
	switch lit.Type {
	case scanner.Int:
		value, err := strconv.ParseInt(lit.Text, 10, 32)
		if err != nil {
			panic(err)
		}
		return int32(value)

	case scanner.Float:
		value, err := strconv.ParseFloat(lit.Text, 64)
		if err != nil {
			panic(err)
		}
		return value

	case scanner.String:
		value, err := strconv.Unquote(lit.Text)
		if err != nil {
			panic(err)
		}
		return value

	case scanner.Ident:
		switch lit.Text {
		case "true":
			return true
		case "false":
			return false
		default:
			return lit.Text
		}

	default:
		panic("invalid literal")
	}
}

func (lit *BasicLit) String() string {
	return lit.Text
}

func (lit *BasicLit) Location() errors.Location {
	return lit.Loc
}

type ListLit struct {
	Entries []Literal
	Loc     errors.Location
}

func (lit *ListLit) Value(vars map[string]interface{}) interface{} {
	entries := make([]interface{}, len(lit.Entries))
	for i, entry := range lit.Entries {
		entries[i] = entry.Value(vars)
	}
	return entries
}

func (lit *ListLit) String() string {
	entries := make([]string, len(lit.Entries))
	for i, entry := range lit.Entries {
		entries[i] = entry.String()
	}
	return "[" + strings.Join(entries, ", ") + "]"
}

func (lit *ListLit) Location() errors.Location {
	return lit.Loc
}

type ObjectLit struct {
	Fields []*ObjectLitField
	Loc    errors.Location
}

type ObjectLitField struct {
	Name  Ident
	Value Literal
}

func (lit *ObjectLit) Value(vars map[string]interface{}) interface{} {
	fields := make(map[string]interface{}, len(lit.Fields))
	for _, f := range lit.Fields {
		fields[f.Name.Name] = f.Value.Value(vars)
	}
	return fields
}

func (lit *ObjectLit) String() string {
	entries := make([]string, 0, len(lit.Fields))
	for _, f := range lit.Fields {
		entries = append(entries, f.Name.Name+": "+f.Value.String())
	}
	return "{" + strings.Join(entries, ", ") + "}"
}

func (lit *ObjectLit) Location() errors.Location {
	return lit.Loc
}

type NullLit struct {
	Loc errors.Location
}

func (lit *NullLit) Value(vars map[string]interface{}) interface{} {
	return nil
}

func (lit *NullLit) String() string {
	return "null"
}

func (lit *NullLit) Location() errors.Location {
	return lit.Loc
}

type Variable struct {
	Name string
	Loc  errors.Location
}

func (v Variable) Value(vars map[string]interface{}) interface{} {
	return vars[v.Name]
}

func (v Variable) String() string {
	return "$" + v.Name
}

func (v *Variable) Location() errors.Location {
	return v.Loc
}

func ParseLiteral(l *Lexer, constOnly bool) Literal {
	loc := l.Location()
	switch l.Peek() {
	case '$':
		if constOnly {
			l.SyntaxError("variable not allowed")
			panic("unreachable")
		}
		l.ConsumeToken('$')
		return &Variable{l.ConsumeIdent(), loc}

	case scanner.Int, scanner.Float, scanner.String, scanner.Ident:
		lit := l.ConsumeLiteral()
		if lit.Type == scanner.Ident && lit.Text == "null" {
			return &NullLit{loc}
		}
		lit.Loc = loc
		return lit
	case '-':
		l.ConsumeToken('-')
		lit := l.ConsumeLiteral()
		lit.Text = "-" + lit.Text
		lit.Loc = loc
		return lit
	case '[':
		l.ConsumeToken('[')
		var list []Literal
		for l.Peek() != ']' {
			list = append(list, ParseLiteral(l, constOnly))
		}
		l.ConsumeToken(']')
		return &ListLit{list, loc}

	case '{':
		l.ConsumeToken('{')
		var fields []*ObjectLitField
		for l.Peek() != '}' {
			name := l.ConsumeIdentWithLoc()
			l.ConsumeToken(':')
			value := ParseLiteral(l, constOnly)
			fields = append(fields, &ObjectLitField{name, value})
		}
		l.ConsumeToken('}')
		return &ObjectLit{fields, loc}

	default:
		l.SyntaxError("invalid value")
		panic("unreachable")
	}
}
<|endoftext|>


<commit_before>package victor

import (
    "regexp"
    "log"
)

type RobotAdapter interface {
    Hear(string, func(*TextMessage))
    Respond(string, func(*TextMessage))
    Run()
}

type Robot struct {
    name      string
    options   map[string]string
    adapter   RobotAdapter
    listeners []*Listener
    users     []*User
}

func NewRobot(adapter string, options map[string]string) RobotAdapter {
    name := "robot"

    if _, exists := options["name"]; exists {
        name = options["name"]
    }

    robot := &Robot{
        name: name,
        options: options,
        listeners: make([]*Listener, 0, 1),
    }

    var aRobot RobotAdapter

    if adapter == "shell" {
        aRobot = &Shell{Robot: robot}
    } else if adapter == "campfire" {
        aRobot = &Campfire{Robot: robot}
    } else {
        panic("Unknown adapter.")
    }

    return aRobot
}

func (self *Robot) Hear(expStr string, callback func(*TextMessage)) {
    exp, _ := regexp.Compile(expStr)
    self.listeners = append(self.listeners, NewListener(exp, callback))
}

func (self *Robot) Respond(expStr string, callback func(*TextMessage)) {
    expWithNameStr := "(" + self.name + "[:,]?)\\s*(?:" + expStr + ")"
    exp, _         := regexp.Compile(expWithNameStr)

    self.listeners = append(self.listeners, NewListener(exp, callback))
}

func (self *Robot) Receive(msg *TextMessage) {
    for _, listener := range self.listeners {
        if listener.Test(msg) {
            log.Printf("Listener /%s/ triggered by '%s'", listener.Exp.String(), msg.Body)
            listener.Callback(msg)
        }
    }
}

func (self *Robot) RememberUser(user *User) {
    self.users = append(self.users, user)
}

func (self *Robot) UserForId(id int) *User {
    for _, user := range self.users {
        if user.Id == id {
            return user
        }
    }

    return nil
}

func (self *Robot) Shutdown() {
    log.Print("See ya!")
}
<commit_msg>Beginning of string only.<commit_after>package victor

import (
    "regexp"
    "log"
)

type RobotAdapter interface {
    Hear(string, func(*TextMessage))
    Respond(string, func(*TextMessage))
    Run()
}

type Robot struct {
    name      string
    options   map[string]string
    adapter   RobotAdapter
    listeners []*Listener
    users     []*User
}

func NewRobot(adapter string, options map[string]string) RobotAdapter {
    name := "robot"

    if _, exists := options["name"]; exists {
        name = options["name"]
    }

    robot := &Robot{
        name: name,
        options: options,
        listeners: make([]*Listener, 0, 1),
    }

    var aRobot RobotAdapter

    if adapter == "shell" {
        aRobot = &Shell{Robot: robot}
    } else if adapter == "campfire" {
        aRobot = &Campfire{Robot: robot}
    } else {
        panic("Unknown adapter.")
    }

    return aRobot
}

func (self *Robot) Hear(expStr string, callback func(*TextMessage)) {
    exp, _ := regexp.Compile(expStr)
    self.listeners = append(self.listeners, NewListener(exp, callback))
}

func (self *Robot) Respond(expStr string, callback func(*TextMessage)) {
    expWithNameStr := "^(" + self.name + "[:,]?)\\s*(?:" + expStr + ")"
    exp, _         := regexp.Compile(expWithNameStr)

    self.listeners = append(self.listeners, NewListener(exp, callback))
}

func (self *Robot) Receive(msg *TextMessage) {
    for _, listener := range self.listeners {
        if listener.Test(msg) {
            log.Printf("Listener /%s/ triggered by '%s'", listener.Exp.String(), msg.Body)
            listener.Callback(msg)
        }
    }
}

func (self *Robot) RememberUser(user *User) {
    self.users = append(self.users, user)
}

func (self *Robot) UserForId(id int) *User {
    for _, user := range self.users {
        if user.Id == id {
            return user
        }
    }

    return nil
}

func (self *Robot) Shutdown() {
    log.Print("See ya!")
}
<|endoftext|>


<commit_before>// Copyright 2014 Hajime Hoshi
//
// 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 graphics

import (
	"strings"

	"github.com/hajimehoshi/ebiten/internal/opengl"
)

type shaderId int

const (
	shaderVertexModelview shaderId = iota
	shaderFragmentTexture
)

func shader(id shaderId) string {
	str := shaders[id]
	if !opengl.GetContext().GlslHighpSupported() {
		str = strings.Replace(str, "highp ", "", -1)
		str = strings.Replace(str, "lowp ", "", -1)
	}
	return str
}

var shaders = map[shaderId]string{
	shaderVertexModelview: `
uniform highp mat4 projection_matrix;
attribute highp vec2 vertex;
attribute highp vec2 tex_coord;
attribute highp vec4 geo_matrix_body;
attribute highp vec2 geo_matrix_translation;
varying highp vec2 vertex_out_tex_coord;

void main(void) {
  vertex_out_tex_coord = tex_coord;
  mat4 geo_matrix = mat4(
    vec4(geo_matrix_body[0], geo_matrix_body[2], 0, 0),
    vec4(geo_matrix_body[1], geo_matrix_body[3], 0, 0),
    vec4(0, 0, 1, 0),
    vec4(geo_matrix_translation, 0, 1)
  );
  gl_Position = projection_matrix * geo_matrix * vec4(vertex, 0, 1);
}
`,
	shaderFragmentTexture: `
uniform lowp sampler2D texture;
uniform lowp mat4 color_matrix;
uniform lowp vec4 color_matrix_translation;
varying highp vec2 vertex_out_tex_coord;

void main(void) {
  lowp vec4 color = texture2D(texture, vertex_out_tex_coord);

  // Un-premultiply alpha
  if (0.0 < color.a) {
    color.rgb /= color.a;
  }
  // Apply the color matrix
  color = (color_matrix * color) + color_matrix_translation;
  color = clamp(color, 0.0, 1.0);
  // Premultiply alpha
  color.rgb *= color.a;

  gl_FragColor = color;
}
`,
}
<commit_msg>graphics: Bug fix: Need to specify default precisions (#390)<commit_after>// Copyright 2014 Hajime Hoshi
//
// 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 graphics

import (
	"strings"

	"github.com/hajimehoshi/ebiten/internal/opengl"
)

type shaderId int

const (
	shaderVertexModelview shaderId = iota
	shaderFragmentTexture
)

func shader(id shaderId) string {
	str := shaders[id]
	if !opengl.GetContext().GlslHighpSupported() {
		str = strings.Replace(str, "highp ", "", -1)
		str = strings.Replace(str, "lowp ", "", -1)
	}
	return str
}

var shaders = map[shaderId]string{
	shaderVertexModelview: `
#ifdef GL_ES
precision highp float;
#endif

uniform highp mat4 projection_matrix;
attribute highp vec2 vertex;
attribute highp vec2 tex_coord;
attribute highp vec4 geo_matrix_body;
attribute highp vec2 geo_matrix_translation;
varying highp vec2 vertex_out_tex_coord;

void main(void) {
  vertex_out_tex_coord = tex_coord;
  mat4 geo_matrix = mat4(
    vec4(geo_matrix_body[0], geo_matrix_body[2], 0, 0),
    vec4(geo_matrix_body[1], geo_matrix_body[3], 0, 0),
    vec4(0, 0, 1, 0),
    vec4(geo_matrix_translation, 0, 1)
  );
  gl_Position = projection_matrix * geo_matrix * vec4(vertex, 0, 1);
}
`,
	shaderFragmentTexture: `
#ifdef GL_ES
precision mediump float;
#endif

uniform lowp sampler2D texture;
uniform lowp mat4 color_matrix;
uniform lowp vec4 color_matrix_translation;
varying highp vec2 vertex_out_tex_coord;

void main(void) {
  lowp vec4 color = texture2D(texture, vertex_out_tex_coord);

  // Un-premultiply alpha
  if (0.0 < color.a) {
    color.rgb /= color.a;
  }
  // Apply the color matrix
  color = (color_matrix * color) + color_matrix_translation;
  color = clamp(color, 0.0, 1.0);
  // Premultiply alpha
  color.rgb *= color.a;

  gl_FragColor = color;
}
`,
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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 publish defines the exposure keys publishing API.
package publish

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"time"

	"go.opencensus.io/trace"

	"github.com/google/exposure-notifications-server/internal/authorizedapp"
	"github.com/google/exposure-notifications-server/internal/jsonutil"
	"github.com/google/exposure-notifications-server/internal/logging"
	"github.com/google/exposure-notifications-server/internal/pb"
	"github.com/google/exposure-notifications-server/internal/publish/database"
	"github.com/google/exposure-notifications-server/internal/publish/model"
	"github.com/google/exposure-notifications-server/internal/revision"
	revisiondb "github.com/google/exposure-notifications-server/internal/revision/database"
	"github.com/google/exposure-notifications-server/internal/serverenv"
	"github.com/google/exposure-notifications-server/internal/verification"
	verifydb "github.com/google/exposure-notifications-server/internal/verification/database"
	verifyapi "github.com/google/exposure-notifications-server/pkg/api/v1alpha1"
	"github.com/google/exposure-notifications-server/pkg/base64util"
	"github.com/mikehelmick/go-chaff"
)

// NewHandler creates the HTTP handler for the TTK publishing API.
func NewHandler(ctx context.Context, config *Config, env *serverenv.ServerEnv) (http.Handler, error) {
	logger := logging.FromContext(ctx)

	if env.Database() == nil {
		return nil, fmt.Errorf("missing database in server environment")
	}
	if env.AuthorizedAppProvider() == nil {
		return nil, fmt.Errorf("missing AuthorizedApp provider in server environment")
	}

	transformer, err := model.NewTransformer(config)
	if err != nil {
		return nil, fmt.Errorf("model.NewTransformer: %w", err)
	}
	logger.Infof("max keys per upload: %v", config.MaxKeysOnPublish)
	logger.Infof("max same start interval keys: %v", config.MaxSameStartIntervalKeys)
	logger.Infof("max interval start age: %v", config.MaxIntervalAge)
	logger.Infof("truncate window: %v", config.TruncateWindow)
	if config.DebugReleaseSameDayKeys() {
		logger.Warnf("SERVER IS IN DEBUG MODE. KEYS MAY BE RELEASED EARLY.")
	}

	verifier, err := verification.New(verifydb.New(env.Database()), &config.Verification)
	if err != nil {
		return nil, fmt.Errorf("verification.New: %w", err)
	}

	aadBytes, err := config.RevisionTokenADDBytes()
	if err != nil {
		return nil, fmt.Errorf("must provide ADD for revision token encryption in REVISION_TOKEN_AAD env variable: %w", err)
	}
	revisionKeyConfig := revisiondb.KMSConfig{
		WrapperKeyID: config.RevisionTokenKeyID,
		KeyManager:   env.GetKeyManager(),
	}
	revisionDB, err := revisiondb.New(env.Database(), &revisionKeyConfig)
	if err != nil {
		return nil, fmt.Errorf("revisiondb.New: %w", err)
	}
	tm, err := revision.New(ctx, revisionDB, config.RevisionKeyCacheDuration, config.RevisionTokenMinLength)
	if err != nil {
		return nil, fmt.Errorf("revision.New: %w", err)
	}

	return &publishHandler{
		serverenv:             env,
		transformer:           transformer,
		config:                config,
		database:              database.New(env.Database()),
		tracker:               chaff.New(),
		tokenManager:          tm,
		tokenAAD:              aadBytes,
		authorizedAppProvider: env.AuthorizedAppProvider(),
		verifier:              verifier,
	}, nil
}

type publishHandler struct {
	config                *Config
	serverenv             *serverenv.ServerEnv
	transformer           *model.Transformer
	database              *database.PublishDB
	tokenManager          *revision.TokenManager
	tracker               *chaff.Tracker
	tokenAAD              []byte
	authorizedAppProvider authorizedapp.Provider
	verifier              *verification.Verifier
}

type response struct {
	status      int
	pubResponse *verifyapi.PublishResponse
	metric      string
	count       int // metricCount
}

func (h *publishHandler) handleRequest(w http.ResponseWriter, r *http.Request) response {
	ctx, span := trace.StartSpan(r.Context(), "(*publish.publishHandler).handleRequest")
	defer span.End()

	logger := logging.FromContext(ctx)
	metrics := h.serverenv.MetricsExporter(ctx)

	var data verifyapi.Publish
	code, err := jsonutil.Unmarshal(w, r, &data)
	if err != nil {
		message := fmt.Sprintf("error unmarshaling API call, code: %v: %v", code, err)
		span.SetStatus(trace.Status{Code: trace.StatusCodeInternal, Message: message})
		return response{
			status:      http.StatusBadRequest,
			pubResponse: &verifyapi.PublishResponse{Error: message},
			metric:      "publish-bad-json", count: 1}
	}

	appConfig, err := h.authorizedAppProvider.AppConfig(ctx, data.AppPackageName)
	if err != nil {
		// Config loaded, but app with that name isn't registered. This can also
		// happen if the app was recently registered but the cache hasn't been
		// refreshed.
		if errors.Is(err, authorizedapp.ErrAppNotFound) {
			message := fmt.Sprintf("unauthorized app: %v", data.AppPackageName)
			span.SetStatus(trace.Status{Code: trace.StatusCodeInternal, Message: message})
			return response{status: http.StatusUnauthorized,
				pubResponse: &verifyapi.PublishResponse{Error: message},
				metric:      "publish-app-not-authorized", count: 1}
		}

		// A higher-level configuration error occurred, likely while trying to read
		// from the database. This is retryable, although won't succeed if the error
		// isn't transient.
		// This message (and logging) will only contain the AppPkgName from the request
		// and no other data from the request.
		message := fmt.Sprintf("no AuthorizedApp, dropping data: %v", err)
		span.SetStatus(trace.Status{Code: trace.StatusCodeInternal, Message: message})
		return response{
			status:      http.StatusUnauthorized,
			pubResponse: &verifyapi.PublishResponse{Error: message},
			metric:      "publish-error-loading-authorizedapp",
			count:       1,
		}
	}

	// Verify the request is from a permitted region.
	if len(data.Regions) == 0 {
		message := "no regions provided on publish request"
		span.SetStatus(trace.Status{Code: trace.StatusCodePermissionDenied, Message: message})
		return response{
			status:      http.StatusBadRequest,
			pubResponse: &verifyapi.PublishResponse{Error: message},
			metric:      "publish-region-not-specified",
			count:       1,
		}
	}
	for _, r := range data.Regions {
		if !appConfig.IsAllowedRegion(r) {
			err := fmt.Errorf("app %v tried to write to unauthorized region %v", appConfig.AppPackageName, r)
			message := fmt.Sprintf("verifying allowed regions: %v", err)
			span.SetStatus(trace.Status{Code: trace.StatusCodePermissionDenied, Message: message})
			return response{
				status:      http.StatusUnauthorized,
				pubResponse: &verifyapi.PublishResponse{Error: message},
				metric:      "publish-region-not-authorized",
				count:       1,
			}
		}
	}

	// Perform health authority certificate verification.
	verifiedClaims, err := h.verifier.VerifyDiagnosisCertificate(ctx, appConfig, &data)
	if err != nil {
		if appConfig.BypassHealthAuthorityVerification {
			logger.Warnf("bypassing health authority certificate verification for app: %v", appConfig.AppPackageName)
			metrics.WriteInt("publish-health-authority-verification-bypassed", true, 1)
		} else {
			message := fmt.Sprintf("unable to validate diagnosis verification: %v", err)
			span.SetStatus(trace.Status{Code: trace.StatusCodeInvalidArgument, Message: message})
			return response{
				status:      http.StatusUnauthorized,
				pubResponse: &verifyapi.PublishResponse{Error: message},
				metric:      "publish-bad-verification", count: 1}
		}
	}

	// Examine the revision token.
	var token *pb.RevisionTokenData
	if len(data.RevisionToken) != 0 {
		encryptedToken, err := base64util.DecodeString(data.RevisionToken)
		if err != nil {
			logger.Errorf("unable to decode revision token, proceeding without: %v", err)
		} else {
			token, err = h.tokenManager.UnmarshalRevisionToken(ctx, encryptedToken, h.tokenAAD)
			if err != nil {
				logger.Errorf("unable to unmarshall / descrypt revision token: %v", err)
				token = nil // just in case.
			}
		}
	}

	batchTime := time.Now()
	exposures, err := h.transformer.TransformPublish(ctx, &data, verifiedClaims, batchTime)
	if err != nil {
		message := fmt.Sprintf("unable to read request data: %v", err)
		logger.Errorf(message)
		span.SetStatus(trace.Status{Code: trace.StatusCodeInvalidArgument, Message: message})
		return response{
			status:      http.StatusBadRequest,
			pubResponse: &verifyapi.PublishResponse{Error: message},
			metric:      "publish-transform-fail", count: 1}
	}

	n, err := h.database.InsertAndReviseExposures(ctx, exposures, token, appConfig.BypassRevisionToken)
	if err != nil {
		message := fmt.Sprintf("error writing exposure record: %v", err)
		logger.Errorf(message)
		span.SetStatus(trace.Status{Code: trace.StatusCodeInternal, Message: message})
		return response{
			status:      http.StatusInternalServerError,
			pubResponse: &verifyapi.PublishResponse{Error: http.StatusText(http.StatusInternalServerError)},
			metric:      "publish-db-write-error", count: 1}
	}

	// Build the new revision token. Union of existing token take + new exposures.
	var keep pb.RevisionTokenData
	if token != nil {
		// put existing tokens that aren't too old back in the token.
		retainInterval := model.IntervalNumber(batchTime.Add(-1 * h.config.MaxIntervalAge))
		for _, rk := range token.RevisableKeys {
			if rk.IntervalNumber+rk.IntervalCount >= retainInterval {
				keep.RevisableKeys = append(keep.RevisableKeys, rk)
			}
		}
	}
	newToken, err := h.tokenManager.MakeRevisionToken(ctx, &keep, exposures, h.tokenAAD)
	if err != nil {
		logger.Errorf("unable to make new revision token: %v", err)
	}

	message := fmt.Sprintf("Inserted %d exposures.", n)
	span.AddAttributes(trace.Int64Attribute("inserted_exposures", int64(n)))
	logger.Info(message)
	return response{
		status: http.StatusOK,
		pubResponse: &verifyapi.PublishResponse{
			RevisionToken:     base64.StdEncoding.EncodeToString(newToken),
			InsertedExposures: n,
		},
		metric: "publish-exposures-written",
		count:  n,
	}
}

// ServeHTTP handles the publish event. It tracks requests and can handle chaff
// requests when provided a request with the X-Chaff header.
func (h *publishHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	h.tracker.HandleTrack(chaff.HeaderDetector("X-Chaff"),
		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			response := h.handleRequest(w, r)

			if response.metric != "" {
				ctx := r.Context()
				metrics := h.serverenv.MetricsExporter(ctx)
				metrics.WriteInt(response.metric, true, response.count)
			}

			data, err := json.Marshal(response.pubResponse)
			if err != nil {
				w.WriteHeader(http.StatusInternalServerError)
				fmt.Fprintf(w, "{\"error\": \"%v\"}", err.Error())
				return
			}
			w.WriteHeader(response.status)
			fmt.Fprintf(w, "%s", data)
		})).ServeHTTP(w, r)
}
<commit_msg>fix inverted bypass switch. (#807)<commit_after>// Copyright 2020 Google LLC
//
// 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 publish defines the exposure keys publishing API.
package publish

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"time"

	"go.opencensus.io/trace"

	"github.com/google/exposure-notifications-server/internal/authorizedapp"
	"github.com/google/exposure-notifications-server/internal/jsonutil"
	"github.com/google/exposure-notifications-server/internal/logging"
	"github.com/google/exposure-notifications-server/internal/pb"
	"github.com/google/exposure-notifications-server/internal/publish/database"
	"github.com/google/exposure-notifications-server/internal/publish/model"
	"github.com/google/exposure-notifications-server/internal/revision"
	revisiondb "github.com/google/exposure-notifications-server/internal/revision/database"
	"github.com/google/exposure-notifications-server/internal/serverenv"
	"github.com/google/exposure-notifications-server/internal/verification"
	verifydb "github.com/google/exposure-notifications-server/internal/verification/database"
	verifyapi "github.com/google/exposure-notifications-server/pkg/api/v1alpha1"
	"github.com/google/exposure-notifications-server/pkg/base64util"
	"github.com/mikehelmick/go-chaff"
)

// NewHandler creates the HTTP handler for the TTK publishing API.
func NewHandler(ctx context.Context, config *Config, env *serverenv.ServerEnv) (http.Handler, error) {
	logger := logging.FromContext(ctx)

	if env.Database() == nil {
		return nil, fmt.Errorf("missing database in server environment")
	}
	if env.AuthorizedAppProvider() == nil {
		return nil, fmt.Errorf("missing AuthorizedApp provider in server environment")
	}

	transformer, err := model.NewTransformer(config)
	if err != nil {
		return nil, fmt.Errorf("model.NewTransformer: %w", err)
	}
	logger.Infof("max keys per upload: %v", config.MaxKeysOnPublish)
	logger.Infof("max same start interval keys: %v", config.MaxSameStartIntervalKeys)
	logger.Infof("max interval start age: %v", config.MaxIntervalAge)
	logger.Infof("truncate window: %v", config.TruncateWindow)
	if config.DebugReleaseSameDayKeys() {
		logger.Warnf("SERVER IS IN DEBUG MODE. KEYS MAY BE RELEASED EARLY.")
	}

	verifier, err := verification.New(verifydb.New(env.Database()), &config.Verification)
	if err != nil {
		return nil, fmt.Errorf("verification.New: %w", err)
	}

	aadBytes, err := config.RevisionTokenADDBytes()
	if err != nil {
		return nil, fmt.Errorf("must provide ADD for revision token encryption in REVISION_TOKEN_AAD env variable: %w", err)
	}
	revisionKeyConfig := revisiondb.KMSConfig{
		WrapperKeyID: config.RevisionTokenKeyID,
		KeyManager:   env.GetKeyManager(),
	}
	revisionDB, err := revisiondb.New(env.Database(), &revisionKeyConfig)
	if err != nil {
		return nil, fmt.Errorf("revisiondb.New: %w", err)
	}
	tm, err := revision.New(ctx, revisionDB, config.RevisionKeyCacheDuration, config.RevisionTokenMinLength)
	if err != nil {
		return nil, fmt.Errorf("revision.New: %w", err)
	}

	return &publishHandler{
		serverenv:             env,
		transformer:           transformer,
		config:                config,
		database:              database.New(env.Database()),
		tracker:               chaff.New(),
		tokenManager:          tm,
		tokenAAD:              aadBytes,
		authorizedAppProvider: env.AuthorizedAppProvider(),
		verifier:              verifier,
	}, nil
}

type publishHandler struct {
	config                *Config
	serverenv             *serverenv.ServerEnv
	transformer           *model.Transformer
	database              *database.PublishDB
	tokenManager          *revision.TokenManager
	tracker               *chaff.Tracker
	tokenAAD              []byte
	authorizedAppProvider authorizedapp.Provider
	verifier              *verification.Verifier
}

type response struct {
	status      int
	pubResponse *verifyapi.PublishResponse
	metric      string
	count       int // metricCount
}

func (h *publishHandler) handleRequest(w http.ResponseWriter, r *http.Request) response {
	ctx, span := trace.StartSpan(r.Context(), "(*publish.publishHandler).handleRequest")
	defer span.End()

	logger := logging.FromContext(ctx)
	metrics := h.serverenv.MetricsExporter(ctx)

	var data verifyapi.Publish
	code, err := jsonutil.Unmarshal(w, r, &data)
	if err != nil {
		message := fmt.Sprintf("error unmarshaling API call, code: %v: %v", code, err)
		span.SetStatus(trace.Status{Code: trace.StatusCodeInternal, Message: message})
		return response{
			status:      http.StatusBadRequest,
			pubResponse: &verifyapi.PublishResponse{Error: message},
			metric:      "publish-bad-json", count: 1}
	}

	appConfig, err := h.authorizedAppProvider.AppConfig(ctx, data.AppPackageName)
	if err != nil {
		// Config loaded, but app with that name isn't registered. This can also
		// happen if the app was recently registered but the cache hasn't been
		// refreshed.
		if errors.Is(err, authorizedapp.ErrAppNotFound) {
			message := fmt.Sprintf("unauthorized app: %v", data.AppPackageName)
			span.SetStatus(trace.Status{Code: trace.StatusCodeInternal, Message: message})
			return response{status: http.StatusUnauthorized,
				pubResponse: &verifyapi.PublishResponse{Error: message},
				metric:      "publish-app-not-authorized", count: 1}
		}

		// A higher-level configuration error occurred, likely while trying to read
		// from the database. This is retryable, although won't succeed if the error
		// isn't transient.
		// This message (and logging) will only contain the AppPkgName from the request
		// and no other data from the request.
		message := fmt.Sprintf("no AuthorizedApp, dropping data: %v", err)
		span.SetStatus(trace.Status{Code: trace.StatusCodeInternal, Message: message})
		return response{
			status:      http.StatusUnauthorized,
			pubResponse: &verifyapi.PublishResponse{Error: message},
			metric:      "publish-error-loading-authorizedapp",
			count:       1,
		}
	}

	// Verify the request is from a permitted region.
	if len(data.Regions) == 0 {
		message := "no regions provided on publish request"
		span.SetStatus(trace.Status{Code: trace.StatusCodePermissionDenied, Message: message})
		return response{
			status:      http.StatusBadRequest,
			pubResponse: &verifyapi.PublishResponse{Error: message},
			metric:      "publish-region-not-specified",
			count:       1,
		}
	}
	for _, r := range data.Regions {
		if !appConfig.IsAllowedRegion(r) {
			err := fmt.Errorf("app %v tried to write to unauthorized region %v", appConfig.AppPackageName, r)
			message := fmt.Sprintf("verifying allowed regions: %v", err)
			span.SetStatus(trace.Status{Code: trace.StatusCodePermissionDenied, Message: message})
			return response{
				status:      http.StatusUnauthorized,
				pubResponse: &verifyapi.PublishResponse{Error: message},
				metric:      "publish-region-not-authorized",
				count:       1,
			}
		}
	}

	// Perform health authority certificate verification.
	verifiedClaims, err := h.verifier.VerifyDiagnosisCertificate(ctx, appConfig, &data)
	if err != nil {
		if appConfig.BypassHealthAuthorityVerification {
			logger.Warnf("bypassing health authority certificate verification for app: %v", appConfig.AppPackageName)
			metrics.WriteInt("publish-health-authority-verification-bypassed", true, 1)
		} else {
			message := fmt.Sprintf("unable to validate diagnosis verification: %v", err)
			span.SetStatus(trace.Status{Code: trace.StatusCodeInvalidArgument, Message: message})
			return response{
				status:      http.StatusUnauthorized,
				pubResponse: &verifyapi.PublishResponse{Error: message},
				metric:      "publish-bad-verification", count: 1}
		}
	}

	// Examine the revision token.
	var token *pb.RevisionTokenData
	if len(data.RevisionToken) != 0 {
		encryptedToken, err := base64util.DecodeString(data.RevisionToken)
		if err != nil {
			logger.Errorf("unable to decode revision token, proceeding without: %v", err)
		} else {
			token, err = h.tokenManager.UnmarshalRevisionToken(ctx, encryptedToken, h.tokenAAD)
			if err != nil {
				logger.Errorf("unable to unmarshall / descrypt revision token: %v", err)
				token = nil // just in case.
			}
		}
	}

	batchTime := time.Now()
	exposures, err := h.transformer.TransformPublish(ctx, &data, verifiedClaims, batchTime)
	if err != nil {
		message := fmt.Sprintf("unable to read request data: %v", err)
		logger.Errorf(message)
		span.SetStatus(trace.Status{Code: trace.StatusCodeInvalidArgument, Message: message})
		return response{
			status:      http.StatusBadRequest,
			pubResponse: &verifyapi.PublishResponse{Error: message},
			metric:      "publish-transform-fail", count: 1}
	}

	n, err := h.database.InsertAndReviseExposures(ctx, exposures, token, !appConfig.BypassRevisionToken)
	if err != nil {
		message := fmt.Sprintf("error writing exposure record: %v", err)
		logger.Errorf(message)
		span.SetStatus(trace.Status{Code: trace.StatusCodeInternal, Message: message})
		return response{
			status:      http.StatusInternalServerError,
			pubResponse: &verifyapi.PublishResponse{Error: http.StatusText(http.StatusInternalServerError)},
			metric:      "publish-db-write-error", count: 1}
	}

	// Build the new revision token. Union of existing token take + new exposures.
	var keep pb.RevisionTokenData
	if token != nil {
		// put existing tokens that aren't too old back in the token.
		retainInterval := model.IntervalNumber(batchTime.Add(-1 * h.config.MaxIntervalAge))
		for _, rk := range token.RevisableKeys {
			if rk.IntervalNumber+rk.IntervalCount >= retainInterval {
				keep.RevisableKeys = append(keep.RevisableKeys, rk)
			}
		}
	}
	newToken, err := h.tokenManager.MakeRevisionToken(ctx, &keep, exposures, h.tokenAAD)
	if err != nil {
		logger.Errorf("unable to make new revision token: %v", err)
	}

	message := fmt.Sprintf("Inserted %d exposures.", n)
	span.AddAttributes(trace.Int64Attribute("inserted_exposures", int64(n)))
	logger.Info(message)
	return response{
		status: http.StatusOK,
		pubResponse: &verifyapi.PublishResponse{
			RevisionToken:     base64.StdEncoding.EncodeToString(newToken),
			InsertedExposures: n,
		},
		metric: "publish-exposures-written",
		count:  n,
	}
}

// ServeHTTP handles the publish event. It tracks requests and can handle chaff
// requests when provided a request with the X-Chaff header.
func (h *publishHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	h.tracker.HandleTrack(chaff.HeaderDetector("X-Chaff"),
		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			response := h.handleRequest(w, r)

			if response.metric != "" {
				ctx := r.Context()
				metrics := h.serverenv.MetricsExporter(ctx)
				metrics.WriteInt(response.metric, true, response.count)
			}

			data, err := json.Marshal(response.pubResponse)
			if err != nil {
				w.WriteHeader(http.StatusInternalServerError)
				fmt.Fprintf(w, "{\"error\": \"%v\"}", err.Error())
				return
			}
			w.WriteHeader(response.status)
			fmt.Fprintf(w, "%s", data)
		})).ServeHTTP(w, r)
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package runtime

// mtail programs may be created, updated, and deleted while mtail is running, and they will be
// reloaded without having to restart the mtail process -- mtail will handle these on a HUP signal.

import (
	"bytes"
	"context"
	"crypto/sha256"
	"expvar"
	"io"
	"io/ioutil"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/golang/glog"
	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus"

	"github.com/google/mtail/internal/logline"
	"github.com/google/mtail/internal/metrics"
	"github.com/google/mtail/internal/runtime/compiler"
	"github.com/google/mtail/internal/runtime/vm"
)

var (
	// LineCount counts the number of lines received by the program loader.
	LineCount = expvar.NewInt("lines_total")
	// ProgLoads counts the number of program load events.
	ProgLoads = expvar.NewMap("prog_loads_total")
	// ProgLoadErrors counts the number of program load errors.
	ProgLoadErrors = expvar.NewMap("prog_load_errors_total")
)

const (
	fileExt = ".mtail"
)

// LoadAllPrograms loads all programs in a directory and starts watching the
// directory for filesystem changes.  Any compile errors are stored for later retrieival.
// This function returns an error if an internal error occurs.
func (r *Runtime) LoadAllPrograms() error {
	if r.programPath == "" {
		glog.V(2).Info("Programpath is empty, loading nothing")
		return nil
	}
	s, err := os.Stat(r.programPath)
	if err != nil {
		return errors.Wrapf(err, "failed to stat %q", r.programPath)
	}
	switch {
	case s.IsDir():
		fis, rerr := ioutil.ReadDir(r.programPath)
		if rerr != nil {
			return errors.Wrapf(rerr, "Failed to list programs in %q", r.programPath)
		}

		for _, fi := range fis {
			if fi.IsDir() {
				continue
			}
			err = r.LoadProgram(filepath.Join(r.programPath, fi.Name()))
			if err != nil {
				if r.errorsAbort {
					return err
				}
				glog.Warning(err)
			}
		}
	default:
		err = r.LoadProgram(r.programPath)
		if err != nil {
			if r.errorsAbort {
				return err
			}
			glog.Warning(err)
		}
	}
	return nil
}

// LoadProgram loads or reloads a program from the full pathname programPath.  The name of
// the program is the basename of the file.
func (r *Runtime) LoadProgram(programPath string) error {
	name := filepath.Base(programPath)
	if strings.HasPrefix(name, ".") {
		glog.V(2).Infof("Skipping %s because it is a hidden file.", programPath)
		return nil
	}
	if filepath.Ext(name) != fileExt {
		glog.V(2).Infof("Skipping %s due to file extension.", programPath)
		return nil
	}
	f, err := os.OpenFile(programPath, os.O_RDONLY, 0600)
	if err != nil {
		ProgLoadErrors.Add(name, 1)
		return errors.Wrapf(err, "Failed to read program %q", programPath)
	}
	defer func() {
		if err := f.Close(); err != nil {
			glog.Warning(err)
		}
	}()
	r.programErrorMu.Lock()
	defer r.programErrorMu.Unlock()
	r.programErrors[name] = r.CompileAndRun(name, f)
	if r.programErrors[name] != nil {
		if r.errorsAbort {
			return r.programErrors[name]
		}
		glog.Infof("Compile errors for %s:\n%s", name, r.programErrors[name])
	}
	return nil
}

// CompileAndRun compiles a program read from the input, starting execution if
// it succeeds.  If an existing virtual machine of the same name already
// exists, the previous virtual machine is terminated and the new loaded over
// it.  If the new program fails to compile, any existing virtual machine with
// the same name remains running.
func (r *Runtime) CompileAndRun(name string, input io.Reader) error {
	glog.V(2).Infof("CompileAndRun %s", name)
	var buf bytes.Buffer
	tee := io.TeeReader(input, &buf)
	hasher := sha256.New()
	if _, err := io.Copy(hasher, tee); err != nil {
		ProgLoadErrors.Add(name, 1)
		return errors.Wrapf(err, "hashing failed for %q", name)
	}
	contentHash := hasher.Sum(nil)
	r.handleMu.RLock()
	vh, ok := r.handles[name]
	r.handleMu.RUnlock()
	if ok && bytes.Equal(vh.contentHash, contentHash) {
		glog.V(1).Infof("contents match, not recompiling %q", name)
		return nil
	}
	obj, errs := r.c.Compile(name, &buf)
	if errs != nil {
		ProgLoadErrors.Add(name, 1)
		return errors.Errorf("compile failed for %s:\n%s", name, errs)
	}
	if obj == nil {
		ProgLoadErrors.Add(name, 1)
		return errors.Errorf("Internal error: Compilation failed for %s: No program returned, but no errors.", name)
	}
	v := vm.New(name, obj, r.syslogUseCurrentYear, r.overrideLocation, r.logRuntimeErrors)

	if r.dumpBytecode {
		glog.Info("Dumping program objects and bytecode\n", v.DumpByteCode())
	}

	// Load the metrics from the compilation into the global metric storage for export.
	for _, m := range v.Metrics {
		if !m.Hidden {
			if r.omitMetricSource {
				m.Source = ""
			}
			err := r.ms.Add(m)
			if err != nil {
				return err
			}
		}
	}

	ProgLoads.Add(name, 1)
	glog.Infof("Loaded program %s", name)

	if r.compileOnly {
		return nil
	}

	r.handleMu.Lock()
	defer r.handleMu.Unlock()
	// Terminates the existing vm.
	if handle, ok := r.handles[name]; ok {
		close(handle.lines)
	}
	lines := make(chan *logline.LogLine)
	r.handles[name] = &vmHandle{contentHash: contentHash, vm: v, lines: lines}
	r.wg.Add(1)
	go v.Run(lines, &r.wg)
	return nil
}

type vmHandle struct {
	contentHash []byte
	vm          *vm.VM
	lines       chan *logline.LogLine
}

// Runtime handles the lifecycle of programs and virtual machines, by watching
// the configured program source directory, compiling changes to programs, and
// managing the virtual machines.
type Runtime struct {
	ctx context.Context // a cancellable context
	wg  sync.WaitGroup  // used to await vm shutdown

	ms  *metrics.Store        // pointer to metrics.Store to pass to compiler
	reg prometheus.Registerer // plce to reg metrics

	cOpts []compiler.Option // options for constructing `c`
	c     *compiler.Compiler

	programPath string // Path that contains mtail programs.

	handleMu sync.RWMutex         // guards accesses to handles
	handles  map[string]*vmHandle // map of program names to virtual machines

	programErrorMu sync.RWMutex     // guards access to programErrors
	programErrors  map[string]error // errors from the last compile attempt of the program

	overrideLocation     *time.Location // Instructs the vm to override the timezone with the specified zone.
	compileOnly          bool           // Only compile programs and report errors, do not load VMs.
	errorsAbort          bool           // Compiler errors abort the loader.
	dumpBytecode         bool           // Instructs the loader to dump to stdout the compiled program after compilation.
	syslogUseCurrentYear bool           // Instructs the VM to overwrite zero years with the current year in a strptime instruction.
	omitMetricSource     bool
	logRuntimeErrors     bool // Instruct the VM to emit runtime errors to the log.

	signalQuit chan struct{} // When closed stops the signal handler goroutine.
}

// New creates a new program loader that reads programs from programPath.
func New(lines <-chan *logline.LogLine, wg *sync.WaitGroup, programPath string, store *metrics.Store, options ...Option) (*Runtime, error) {
	if store == nil {
		return nil, errors.New("loader needs a store")
	}
	r := &Runtime{
		ms:            store,
		programPath:   programPath,
		handles:       make(map[string]*vmHandle),
		programErrors: make(map[string]error),
		signalQuit:    make(chan struct{}),
	}
	initDone := make(chan struct{})
	defer close(initDone)
	var err error
	if err = r.SetOption(options...); err != nil {
		return nil, err
	}
	if r.c, err = compiler.New(r.cOpts...); err != nil {
		return nil, err
	}
	// Defer shutdown handling to avoid a race on r.wg.
	wg.Add(1)
	defer func() {
		go func() {
			defer wg.Done()
			<-initDone
			r.wg.Wait()
		}()
	}()
	// This goroutine is the main consumer/producer loop.
	r.wg.Add(1)
	go func() {
		defer r.wg.Done() // signal to owner we're done
		<-initDone
		for line := range lines {
			LineCount.Add(1)
			r.handleMu.RLock()
			for prog := range r.handles {
				r.handles[prog].lines <- line
			}
			r.handleMu.RUnlock()
		}
		glog.Info("END OF LINE")
		close(r.signalQuit)
		r.handleMu.Lock()
		for prog := range r.handles {
			close(r.handles[prog].lines)
			delete(r.handles, prog)
		}
		r.handleMu.Unlock()
	}()
	if r.programPath == "" {
		glog.Info("No program path specified, no programs will be loaded.")
		return r, nil
	}

	// Create one goroutine that handles reload signals.
	r.wg.Add(1)
	go func() {
		defer r.wg.Done()
		<-initDone
		if r.programPath == "" {
			glog.Info("no program reload on SIGHUP without programPath")
			return
		}
		n := make(chan os.Signal, 1)
		signal.Notify(n, syscall.SIGHUP)
		defer signal.Stop(n)
		for {
			select {
			case <-r.signalQuit:
				return
			case <-n:
				if err := r.LoadAllPrograms(); err != nil {
					glog.Info(err)
				}
			}
		}
	}()
	// Guarantee all existing programmes get loaded before we leave.
	if err := r.LoadAllPrograms(); err != nil {
		return nil, err
	}
	return r, nil
}

// SetOption takes one or more option functions and applies them in order to Runtime.
func (r *Runtime) SetOption(options ...Option) error {
	for _, option := range options {
		if err := option(r); err != nil {
			return err
		}
	}
	return nil
}

// UnloadProgram removes the named program, any currently running VM goroutine.
func (r *Runtime) UnloadProgram(pathname string) {
	name := filepath.Base(pathname)
	r.handleMu.Lock()
	defer r.handleMu.Unlock()
	if _, ok := r.handles[name]; ok {
		delete(r.handles, name)
	}
}
<commit_msg>Turn on optimiser pass during development.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package runtime

// mtail programs may be created, updated, and deleted while mtail is running, and they will be
// reloaded without having to restart the mtail process -- mtail will handle these on a HUP signal.

import (
	"bytes"
	"context"
	"crypto/sha256"
	"expvar"
	"io"
	"io/ioutil"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/golang/glog"
	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus"

	"github.com/google/mtail/internal/logline"
	"github.com/google/mtail/internal/metrics"
	"github.com/google/mtail/internal/runtime/compiler"
	"github.com/google/mtail/internal/runtime/vm"
)

var (
	// LineCount counts the number of lines received by the program loader.
	LineCount = expvar.NewInt("lines_total")
	// ProgLoads counts the number of program load events.
	ProgLoads = expvar.NewMap("prog_loads_total")
	// ProgLoadErrors counts the number of program load errors.
	ProgLoadErrors = expvar.NewMap("prog_load_errors_total")
)

const (
	fileExt = ".mtail"
)

// LoadAllPrograms loads all programs in a directory and starts watching the
// directory for filesystem changes.  Any compile errors are stored for later retrieival.
// This function returns an error if an internal error occurs.
func (r *Runtime) LoadAllPrograms() error {
	if r.programPath == "" {
		glog.V(2).Info("Programpath is empty, loading nothing")
		return nil
	}
	s, err := os.Stat(r.programPath)
	if err != nil {
		return errors.Wrapf(err, "failed to stat %q", r.programPath)
	}
	switch {
	case s.IsDir():
		fis, rerr := ioutil.ReadDir(r.programPath)
		if rerr != nil {
			return errors.Wrapf(rerr, "Failed to list programs in %q", r.programPath)
		}

		for _, fi := range fis {
			if fi.IsDir() {
				continue
			}
			err = r.LoadProgram(filepath.Join(r.programPath, fi.Name()))
			if err != nil {
				if r.errorsAbort {
					return err
				}
				glog.Warning(err)
			}
		}
	default:
		err = r.LoadProgram(r.programPath)
		if err != nil {
			if r.errorsAbort {
				return err
			}
			glog.Warning(err)
		}
	}
	return nil
}

// LoadProgram loads or reloads a program from the full pathname programPath.  The name of
// the program is the basename of the file.
func (r *Runtime) LoadProgram(programPath string) error {
	name := filepath.Base(programPath)
	if strings.HasPrefix(name, ".") {
		glog.V(2).Infof("Skipping %s because it is a hidden file.", programPath)
		return nil
	}
	if filepath.Ext(name) != fileExt {
		glog.V(2).Infof("Skipping %s due to file extension.", programPath)
		return nil
	}
	f, err := os.OpenFile(programPath, os.O_RDONLY, 0600)
	if err != nil {
		ProgLoadErrors.Add(name, 1)
		return errors.Wrapf(err, "Failed to read program %q", programPath)
	}
	defer func() {
		if err := f.Close(); err != nil {
			glog.Warning(err)
		}
	}()
	r.programErrorMu.Lock()
	defer r.programErrorMu.Unlock()
	r.programErrors[name] = r.CompileAndRun(name, f)
	if r.programErrors[name] != nil {
		if r.errorsAbort {
			return r.programErrors[name]
		}
		glog.Infof("Compile errors for %s:\n%s", name, r.programErrors[name])
	}
	return nil
}

// CompileAndRun compiles a program read from the input, starting execution if
// it succeeds.  If an existing virtual machine of the same name already
// exists, the previous virtual machine is terminated and the new loaded over
// it.  If the new program fails to compile, any existing virtual machine with
// the same name remains running.
func (r *Runtime) CompileAndRun(name string, input io.Reader) error {
	glog.V(2).Infof("CompileAndRun %s", name)
	var buf bytes.Buffer
	tee := io.TeeReader(input, &buf)
	hasher := sha256.New()
	if _, err := io.Copy(hasher, tee); err != nil {
		ProgLoadErrors.Add(name, 1)
		return errors.Wrapf(err, "hashing failed for %q", name)
	}
	contentHash := hasher.Sum(nil)
	r.handleMu.RLock()
	vh, ok := r.handles[name]
	r.handleMu.RUnlock()
	if ok && bytes.Equal(vh.contentHash, contentHash) {
		glog.V(1).Infof("contents match, not recompiling %q", name)
		return nil
	}
	obj, errs := r.c.Compile(name, &buf)
	if errs != nil {
		ProgLoadErrors.Add(name, 1)
		return errors.Errorf("compile failed for %s:\n%s", name, errs)
	}
	if obj == nil {
		ProgLoadErrors.Add(name, 1)
		return errors.Errorf("Internal error: Compilation failed for %s: No program returned, but no errors.", name)
	}
	v := vm.New(name, obj, r.syslogUseCurrentYear, r.overrideLocation, r.logRuntimeErrors)

	if r.dumpBytecode {
		glog.Info("Dumping program objects and bytecode\n", v.DumpByteCode())
	}

	// Load the metrics from the compilation into the global metric storage for export.
	for _, m := range v.Metrics {
		if !m.Hidden {
			if r.omitMetricSource {
				m.Source = ""
			}
			err := r.ms.Add(m)
			if err != nil {
				return err
			}
		}
	}

	ProgLoads.Add(name, 1)
	glog.Infof("Loaded program %s", name)

	if r.compileOnly {
		return nil
	}

	r.handleMu.Lock()
	defer r.handleMu.Unlock()
	// Terminates the existing vm.
	if handle, ok := r.handles[name]; ok {
		close(handle.lines)
	}
	lines := make(chan *logline.LogLine)
	r.handles[name] = &vmHandle{contentHash: contentHash, vm: v, lines: lines}
	r.wg.Add(1)
	go v.Run(lines, &r.wg)
	return nil
}

type vmHandle struct {
	contentHash []byte
	vm          *vm.VM
	lines       chan *logline.LogLine
}

// Runtime handles the lifecycle of programs and virtual machines, by watching
// the configured program source directory, compiling changes to programs, and
// managing the virtual machines.
type Runtime struct {
	ctx context.Context // a cancellable context
	wg  sync.WaitGroup  // used to await vm shutdown

	ms  *metrics.Store        // pointer to metrics.Store to pass to compiler
	reg prometheus.Registerer // plce to reg metrics

	cOpts []compiler.Option // options for constructing `c`
	c     *compiler.Compiler

	programPath string // Path that contains mtail programs.

	handleMu sync.RWMutex         // guards accesses to handles
	handles  map[string]*vmHandle // map of program names to virtual machines

	programErrorMu sync.RWMutex     // guards access to programErrors
	programErrors  map[string]error // errors from the last compile attempt of the program

	overrideLocation     *time.Location // Instructs the vm to override the timezone with the specified zone.
	compileOnly          bool           // Only compile programs and report errors, do not load VMs.
	errorsAbort          bool           // Compiler errors abort the loader.
	dumpBytecode         bool           // Instructs the loader to dump to stdout the compiled program after compilation.
	syslogUseCurrentYear bool           // Instructs the VM to overwrite zero years with the current year in a strptime instruction.
	omitMetricSource     bool
	logRuntimeErrors     bool // Instruct the VM to emit runtime errors to the log.

	signalQuit chan struct{} // When closed stops the signal handler goroutine.
}

// New creates a new program loader that reads programs from programPath.
func New(lines <-chan *logline.LogLine, wg *sync.WaitGroup, programPath string, store *metrics.Store, options ...Option) (*Runtime, error) {
	if store == nil {
		return nil, errors.New("loader needs a store")
	}
	r := &Runtime{
		ms:            store,
		programPath:   programPath,
		handles:       make(map[string]*vmHandle),
		programErrors: make(map[string]error),
		signalQuit:    make(chan struct{}),
	}
	r.cOpts = append(r.cOpts, compiler.Optimise())
	initDone := make(chan struct{})
	defer close(initDone)
	var err error
	if err = r.SetOption(options...); err != nil {
		return nil, err
	}
	if r.c, err = compiler.New(r.cOpts...); err != nil {
		return nil, err
	}
	// Defer shutdown handling to avoid a race on r.wg.
	wg.Add(1)
	defer func() {
		go func() {
			defer wg.Done()
			<-initDone
			r.wg.Wait()
		}()
	}()
	// This goroutine is the main consumer/producer loop.
	r.wg.Add(1)
	go func() {
		defer r.wg.Done() // signal to owner we're done
		<-initDone
		for line := range lines {
			LineCount.Add(1)
			r.handleMu.RLock()
			for prog := range r.handles {
				r.handles[prog].lines <- line
			}
			r.handleMu.RUnlock()
		}
		glog.Info("END OF LINE")
		close(r.signalQuit)
		r.handleMu.Lock()
		for prog := range r.handles {
			close(r.handles[prog].lines)
			delete(r.handles, prog)
		}
		r.handleMu.Unlock()
	}()
	if r.programPath == "" {
		glog.Info("No program path specified, no programs will be loaded.")
		return r, nil
	}

	// Create one goroutine that handles reload signals.
	r.wg.Add(1)
	go func() {
		defer r.wg.Done()
		<-initDone
		if r.programPath == "" {
			glog.Info("no program reload on SIGHUP without programPath")
			return
		}
		n := make(chan os.Signal, 1)
		signal.Notify(n, syscall.SIGHUP)
		defer signal.Stop(n)
		for {
			select {
			case <-r.signalQuit:
				return
			case <-n:
				if err := r.LoadAllPrograms(); err != nil {
					glog.Info(err)
				}
			}
		}
	}()
	// Guarantee all existing programmes get loaded before we leave.
	if err := r.LoadAllPrograms(); err != nil {
		return nil, err
	}
	return r, nil
}

// SetOption takes one or more option functions and applies them in order to Runtime.
func (r *Runtime) SetOption(options ...Option) error {
	for _, option := range options {
		if err := option(r); err != nil {
			return err
		}
	}
	return nil
}

// UnloadProgram removes the named program, any currently running VM goroutine.
func (r *Runtime) UnloadProgram(pathname string) {
	name := filepath.Base(pathname)
	r.handleMu.Lock()
	defer r.handleMu.Unlock()
	if _, ok := r.handles[name]; ok {
		delete(r.handles, name)
	}
}
<|endoftext|>