<commit_before>package chain

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/btcsuite/btcd/chaincfg"
	"github.com/btcsuite/btcd/chaincfg/chainhash"
	"github.com/btcsuite/btcd/rpcclient"
	"github.com/btcsuite/btcd/wire"
	"github.com/lightninglabs/gozmq"
)

// BitcoindConn represents a persistent client connection to a bitcoind node
// that listens for events read from a ZMQ connection.
type BitcoindConn struct {
	started int32 // To be used atomically.
	stopped int32 // To be used atomically.

	// rescanClientCounter is an atomic counter that assigns a unique ID to
	// each new bitcoind rescan client using the current bitcoind
	// connection.
	rescanClientCounter uint64

	// chainParams identifies the current network the bitcoind node is
	// running on.
	chainParams *chaincfg.Params

	// client is the RPC client to the bitcoind node.
	client *rpcclient.Client

	// zmqBlockConn is the ZMQ connection we'll use to read raw block
	// events.
	zmqBlockConn *gozmq.Conn

	// zmqTxConn is the ZMQ connection we'll use to read raw transaction
	// events.
	zmqTxConn *gozmq.Conn

	// rescanClients is the set of active bitcoind rescan clients to which
	// ZMQ event notfications will be sent to.
	rescanClientsMtx sync.Mutex
	rescanClients    map[uint64]*BitcoindClient

	quit chan struct{}
	wg   sync.WaitGroup
}

// NewBitcoindConn creates a client connection to the node described by the host
// string. The ZMQ connections are established immediately to ensure liveness.
// If the remote node does not operate on the same bitcoin network as described
// by the passed chain parameters, the connection will be disconnected.
func NewBitcoindConn(chainParams *chaincfg.Params,
	host, user, pass, zmqBlockHost, zmqTxHost string,
	zmqPollInterval time.Duration) (*BitcoindConn, error) {

	clientCfg := &rpcclient.ConnConfig{
		Host:                 host,
		User:                 user,
		Pass:                 pass,
		DisableAutoReconnect: false,
		DisableConnectOnNew:  true,
		DisableTLS:           true,
		HTTPPostMode:         true,
	}

	client, err := rpcclient.New(clientCfg, nil)
	if err != nil {
		return nil, err
	}

	// Establish two different ZMQ connections to bitcoind to retrieve block
	// and transaction event notifications. We'll use two as a separation of
	// concern to ensure one type of event isn't dropped from the connection
	// queue due to another type of event filling it up.
	zmqBlockConn, err := gozmq.Subscribe(
		zmqBlockHost, []string{"rawblock"}, zmqPollInterval,
	)
	if err != nil {
		return nil, fmt.Errorf("unable to subscribe for zmq block "+
			"events: %v", err)
	}

	zmqTxConn, err := gozmq.Subscribe(
		zmqTxHost, []string{"rawtx"}, zmqPollInterval,
	)
	if err != nil {
		zmqBlockConn.Close()
		return nil, fmt.Errorf("unable to subscribe for zmq tx "+
			"events: %v", err)
	}

	conn := &BitcoindConn{
		chainParams:   chainParams,
		client:        client,
		zmqBlockConn:  zmqBlockConn,
		zmqTxConn:     zmqTxConn,
		rescanClients: make(map[uint64]*BitcoindClient),
		quit:          make(chan struct{}),
	}

	return conn, nil
}

// Start attempts to establish a RPC and ZMQ connection to a bitcoind node. If
// successful, a goroutine is spawned to read events from the ZMQ connection.
// It's possible for this function to fail due to a limited number of connection
// attempts. This is done to prevent waiting forever on the connection to be
// established in the case that the node is down.
func (c *BitcoindConn) Start() error {
	if !atomic.CompareAndSwapInt32(&c.started, 0, 1) {
		return nil
	}

	// Verify that the node is running on the expected network.
	net, err := c.getCurrentNet()
	if err != nil {
		return err
	}
	if net != c.chainParams.Net {
		return fmt.Errorf("expected network %v, got %v",
			c.chainParams.Net, net)
	}

	c.wg.Add(2)
	go c.blockEventHandler()
	go c.txEventHandler()

	return nil
}

// Stop terminates the RPC and ZMQ connection to a bitcoind node and removes any
// active rescan clients.
func (c *BitcoindConn) Stop() {
	if !atomic.CompareAndSwapInt32(&c.stopped, 0, 1) {
		return
	}

	for _, client := range c.rescanClients {
		client.Stop()
	}

	close(c.quit)
	c.client.Shutdown()
	c.zmqBlockConn.Close()
	c.zmqTxConn.Close()

	c.client.WaitForShutdown()
	c.wg.Wait()
}

// blockEventHandler reads raw blocks events from the ZMQ block socket and
// forwards them along to the current rescan clients.
//
// NOTE: This must be run as a goroutine.
func (c *BitcoindConn) blockEventHandler() {
	defer c.wg.Done()

	log.Info("Started listening for bitcoind block notifications via ZMQ "+
		"on", c.zmqBlockConn.RemoteAddr())

	for {
		// Before attempting to read from the ZMQ socket, we'll make
		// sure to check if we've been requested to shut down.
		select {
		case <-c.quit:
			return
		default:
		}

		// Poll an event from the ZMQ socket.
		msgBytes, err := c.zmqBlockConn.Receive()
		if err != nil {
			// EOF should only be returned if the connection was
			// explicitly closed, so we can exit at this point.
			if err == io.EOF {
				return
			}

			// It's possible that the connection to the socket
			// continuously times out, so we'll prevent logging this
			// error to prevent spamming the logs.
			netErr, ok := err.(net.Error)
			if ok && netErr.Timeout() {
				continue
			}

			log.Errorf("Unable to receive ZMQ rawblock message: %v",
				err)
			continue
		}

		// We have an event! We'll now ensure it is a block event,
		// deserialize it, and report it to the different rescan
		// clients.
		eventType := string(msgBytes[0])
		switch eventType {
		case "rawblock":
			block := &wire.MsgBlock{}
			r := bytes.NewReader(msgBytes[1])
			if err := block.Deserialize(r); err != nil {
				log.Errorf("Unable to deserialize block: %v",
					err)
				continue
			}

			c.rescanClientsMtx.Lock()
			for _, client := range c.rescanClients {
				select {
				case client.zmqBlockNtfns <- block:
				case <-client.quit:
				case <-c.quit:
					c.rescanClientsMtx.Unlock()
					return
				}
			}
			c.rescanClientsMtx.Unlock()
		default:
			// It's possible that the message wasn't fully read if
			// bitcoind shuts down, which will produce an unreadable
			// event type. To prevent from logging it, we'll make
			// sure it conforms to the ASCII standard.
			if eventType == "" || !isASCII(eventType) {
				continue
			}

			log.Warnf("Received unexpected event type from "+
				"rawblock subscription: %v", eventType)
		}
	}
}

// txEventHandler reads raw blocks events from the ZMQ block socket and forwards
// them along to the current rescan clients.
//
// NOTE: This must be run as a goroutine.
func (c *BitcoindConn) txEventHandler() {
	defer c.wg.Done()

	log.Info("Started listening for bitcoind transaction notifications "+
		"via ZMQ on", c.zmqTxConn.RemoteAddr())

	for {
		// Before attempting to read from the ZMQ socket, we'll make
		// sure to check if we've been requested to shut down.
		select {
		case <-c.quit:
			return
		default:
		}

		// Poll an event from the ZMQ socket.
		msgBytes, err := c.zmqTxConn.Receive()
		if err != nil {
			// EOF should only be returned if the connection was
			// explicitly closed, so we can exit at this point.
			if err == io.EOF {
				return
			}

			// It's possible that the connection to the socket
			// continuously times out, so we'll prevent logging this
			// error to prevent spamming the logs.
			netErr, ok := err.(net.Error)
			if ok && netErr.Timeout() {
				continue
			}

			log.Errorf("Unable to receive ZMQ rawtx message: %v",
				err)
			continue
		}

		// We have an event! We'll now ensure it is a transaction event,
		// deserialize it, and report it to the different rescan
		// clients.
		eventType := string(msgBytes[0])
		switch eventType {
		case "rawtx":
			tx := &wire.MsgTx{}
			r := bytes.NewReader(msgBytes[1])
			if err := tx.Deserialize(r); err != nil {
				log.Errorf("Unable to deserialize "+
					"transaction: %v", err)
				continue
			}

			c.rescanClientsMtx.Lock()
			for _, client := range c.rescanClients {
				select {
				case client.zmqTxNtfns <- tx:
				case <-client.quit:
				case <-c.quit:
					c.rescanClientsMtx.Unlock()
					return
				}
			}
			c.rescanClientsMtx.Unlock()
		default:
			// It's possible that the message wasn't fully read if
			// bitcoind shuts down, which will produce an unreadable
			// event type. To prevent from logging it, we'll make
			// sure it conforms to the ASCII standard.
			if eventType == "" || !isASCII(eventType) {
				continue
			}

			log.Warnf("Received unexpected event type from rawtx "+
				"subscription: %v", eventType)
		}
	}
}

// getCurrentNet returns the network on which the bitcoind node is running.
func (c *BitcoindConn) getCurrentNet() (wire.BitcoinNet, error) {
	hash, err := c.client.GetBlockHash(0)
	if err != nil {
		return 0, err
	}

	switch *hash {
	case *chaincfg.TestNet3Params.GenesisHash:
		return chaincfg.TestNet3Params.Net, nil
	case *chaincfg.RegressionNetParams.GenesisHash:
		return chaincfg.RegressionNetParams.Net, nil
	case *chaincfg.MainNetParams.GenesisHash:
		return chaincfg.MainNetParams.Net, nil
	default:
		return 0, fmt.Errorf("unknown network with genesis hash %v", hash)
	}
}

// NewBitcoindClient returns a bitcoind client using the current bitcoind
// connection. This allows us to share the same connection using multiple
// clients.
func (c *BitcoindConn) NewBitcoindClient() *BitcoindClient {
	return &BitcoindClient{
		quit: make(chan struct{}),

		id: atomic.AddUint64(&c.rescanClientCounter, 1),

		chainParams: c.chainParams,
		chainConn:   c,

		rescanUpdate:     make(chan interface{}),
		watchedAddresses: make(map[string]struct{}),
		watchedOutPoints: make(map[wire.OutPoint]struct{}),
		watchedTxs:       make(map[chainhash.Hash]struct{}),

		notificationQueue: NewConcurrentQueue(20),
		zmqTxNtfns:        make(chan *wire.MsgTx),
		zmqBlockNtfns:     make(chan *wire.MsgBlock),

		mempool:        make(map[chainhash.Hash]struct{}),
		expiredMempool: make(map[int32]map[chainhash.Hash]struct{}),
	}
}

// AddClient adds a client to the set of active rescan clients of the current
// chain connection. This allows the connection to include the specified client
// in its notification delivery.
//
// NOTE: This function is safe for concurrent access.
func (c *BitcoindConn) AddClient(client *BitcoindClient) {
	c.rescanClientsMtx.Lock()
	defer c.rescanClientsMtx.Unlock()

	c.rescanClients[client.id] = client
}

// RemoveClient removes the client with the given ID from the set of active
// rescan clients. Once removed, the client will no longer receive block and
// transaction notifications from the chain connection.
//
// NOTE: This function is safe for concurrent access.
func (c *BitcoindConn) RemoveClient(id uint64) {
	c.rescanClientsMtx.Lock()
	defer c.rescanClientsMtx.Unlock()

	delete(c.rescanClients, id)
}

// isASCII is a helper method that checks whether all bytes in `data` would be
// printable ASCII characters if interpreted as a string.
func isASCII(s string) bool {
	for _, c := range s {
		if c < 32 || c > 126 {
			return false
		}
	}
	return true
}
<commit_msg>chain: extract ZMQ command strings into constants<commit_after>package chain

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/btcsuite/btcd/chaincfg"
	"github.com/btcsuite/btcd/chaincfg/chainhash"
	"github.com/btcsuite/btcd/rpcclient"
	"github.com/btcsuite/btcd/wire"
	"github.com/lightninglabs/gozmq"
)

const (
	// rawBlockZMQCommand is the command used to receive raw block
	// notifications from bitcoind through ZMQ.
	rawBlockZMQCommand = "rawblock"

	// rawTxZMQCommand is the command used to receive raw transaction
	// notifications from bitcoind through ZMQ.
	rawTxZMQCommand = "rawtx"
)

// BitcoindConn represents a persistent client connection to a bitcoind node
// that listens for events read from a ZMQ connection.
type BitcoindConn struct {
	started int32 // To be used atomically.
	stopped int32 // To be used atomically.

	// rescanClientCounter is an atomic counter that assigns a unique ID to
	// each new bitcoind rescan client using the current bitcoind
	// connection.
	rescanClientCounter uint64

	// chainParams identifies the current network the bitcoind node is
	// running on.
	chainParams *chaincfg.Params

	// client is the RPC client to the bitcoind node.
	client *rpcclient.Client

	// zmqBlockConn is the ZMQ connection we'll use to read raw block
	// events.
	zmqBlockConn *gozmq.Conn

	// zmqTxConn is the ZMQ connection we'll use to read raw transaction
	// events.
	zmqTxConn *gozmq.Conn

	// rescanClients is the set of active bitcoind rescan clients to which
	// ZMQ event notfications will be sent to.
	rescanClientsMtx sync.Mutex
	rescanClients    map[uint64]*BitcoindClient

	quit chan struct{}
	wg   sync.WaitGroup
}

// NewBitcoindConn creates a client connection to the node described by the host
// string. The ZMQ connections are established immediately to ensure liveness.
// If the remote node does not operate on the same bitcoin network as described
// by the passed chain parameters, the connection will be disconnected.
func NewBitcoindConn(chainParams *chaincfg.Params,
	host, user, pass, zmqBlockHost, zmqTxHost string,
	zmqPollInterval time.Duration) (*BitcoindConn, error) {

	clientCfg := &rpcclient.ConnConfig{
		Host:                 host,
		User:                 user,
		Pass:                 pass,
		DisableAutoReconnect: false,
		DisableConnectOnNew:  true,
		DisableTLS:           true,
		HTTPPostMode:         true,
	}

	client, err := rpcclient.New(clientCfg, nil)
	if err != nil {
		return nil, err
	}

	// Establish two different ZMQ connections to bitcoind to retrieve block
	// and transaction event notifications. We'll use two as a separation of
	// concern to ensure one type of event isn't dropped from the connection
	// queue due to another type of event filling it up.
	zmqBlockConn, err := gozmq.Subscribe(
		zmqBlockHost, []string{rawBlockZMQCommand}, zmqPollInterval,
	)
	if err != nil {
		return nil, fmt.Errorf("unable to subscribe for zmq block "+
			"events: %v", err)
	}

	zmqTxConn, err := gozmq.Subscribe(
		zmqTxHost, []string{rawTxZMQCommand}, zmqPollInterval,
	)
	if err != nil {
		zmqBlockConn.Close()
		return nil, fmt.Errorf("unable to subscribe for zmq tx "+
			"events: %v", err)
	}

	conn := &BitcoindConn{
		chainParams:   chainParams,
		client:        client,
		zmqBlockConn:  zmqBlockConn,
		zmqTxConn:     zmqTxConn,
		rescanClients: make(map[uint64]*BitcoindClient),
		quit:          make(chan struct{}),
	}

	return conn, nil
}

// Start attempts to establish a RPC and ZMQ connection to a bitcoind node. If
// successful, a goroutine is spawned to read events from the ZMQ connection.
// It's possible for this function to fail due to a limited number of connection
// attempts. This is done to prevent waiting forever on the connection to be
// established in the case that the node is down.
func (c *BitcoindConn) Start() error {
	if !atomic.CompareAndSwapInt32(&c.started, 0, 1) {
		return nil
	}

	// Verify that the node is running on the expected network.
	net, err := c.getCurrentNet()
	if err != nil {
		return err
	}
	if net != c.chainParams.Net {
		return fmt.Errorf("expected network %v, got %v",
			c.chainParams.Net, net)
	}

	c.wg.Add(2)
	go c.blockEventHandler()
	go c.txEventHandler()

	return nil
}

// Stop terminates the RPC and ZMQ connection to a bitcoind node and removes any
// active rescan clients.
func (c *BitcoindConn) Stop() {
	if !atomic.CompareAndSwapInt32(&c.stopped, 0, 1) {
		return
	}

	for _, client := range c.rescanClients {
		client.Stop()
	}

	close(c.quit)
	c.client.Shutdown()
	c.zmqBlockConn.Close()
	c.zmqTxConn.Close()

	c.client.WaitForShutdown()
	c.wg.Wait()
}

// blockEventHandler reads raw blocks events from the ZMQ block socket and
// forwards them along to the current rescan clients.
//
// NOTE: This must be run as a goroutine.
func (c *BitcoindConn) blockEventHandler() {
	defer c.wg.Done()

	log.Info("Started listening for bitcoind block notifications via ZMQ "+
		"on", c.zmqBlockConn.RemoteAddr())

	for {
		// Before attempting to read from the ZMQ socket, we'll make
		// sure to check if we've been requested to shut down.
		select {
		case <-c.quit:
			return
		default:
		}

		// Poll an event from the ZMQ socket.
		msgBytes, err := c.zmqBlockConn.Receive()
		if err != nil {
			// EOF should only be returned if the connection was
			// explicitly closed, so we can exit at this point.
			if err == io.EOF {
				return
			}

			// It's possible that the connection to the socket
			// continuously times out, so we'll prevent logging this
			// error to prevent spamming the logs.
			netErr, ok := err.(net.Error)
			if ok && netErr.Timeout() {
				continue
			}

			log.Errorf("Unable to receive ZMQ %v message: %v",
				rawBlockZMQCommand, err)
			continue
		}

		// We have an event! We'll now ensure it is a block event,
		// deserialize it, and report it to the different rescan
		// clients.
		eventType := string(msgBytes[0])
		switch eventType {
		case rawBlockZMQCommand:
			block := &wire.MsgBlock{}
			r := bytes.NewReader(msgBytes[1])
			if err := block.Deserialize(r); err != nil {
				log.Errorf("Unable to deserialize block: %v",
					err)
				continue
			}

			c.rescanClientsMtx.Lock()
			for _, client := range c.rescanClients {
				select {
				case client.zmqBlockNtfns <- block:
				case <-client.quit:
				case <-c.quit:
					c.rescanClientsMtx.Unlock()
					return
				}
			}
			c.rescanClientsMtx.Unlock()
		default:
			// It's possible that the message wasn't fully read if
			// bitcoind shuts down, which will produce an unreadable
			// event type. To prevent from logging it, we'll make
			// sure it conforms to the ASCII standard.
			if eventType == "" || !isASCII(eventType) {
				continue
			}

			log.Warnf("Received unexpected event type from %v "+
				"subscription: %v", rawBlockZMQCommand,
				eventType)
		}
	}
}

// txEventHandler reads raw blocks events from the ZMQ block socket and forwards
// them along to the current rescan clients.
//
// NOTE: This must be run as a goroutine.
func (c *BitcoindConn) txEventHandler() {
	defer c.wg.Done()

	log.Info("Started listening for bitcoind transaction notifications "+
		"via ZMQ on", c.zmqTxConn.RemoteAddr())

	for {
		// Before attempting to read from the ZMQ socket, we'll make
		// sure to check if we've been requested to shut down.
		select {
		case <-c.quit:
			return
		default:
		}

		// Poll an event from the ZMQ socket.
		msgBytes, err := c.zmqTxConn.Receive()
		if err != nil {
			// EOF should only be returned if the connection was
			// explicitly closed, so we can exit at this point.
			if err == io.EOF {
				return
			}

			// It's possible that the connection to the socket
			// continuously times out, so we'll prevent logging this
			// error to prevent spamming the logs.
			netErr, ok := err.(net.Error)
			if ok && netErr.Timeout() {
				continue
			}

			log.Errorf("Unable to receive ZMQ %v message: %v",
				rawTxZMQCommand, err)
			continue
		}

		// We have an event! We'll now ensure it is a transaction event,
		// deserialize it, and report it to the different rescan
		// clients.
		eventType := string(msgBytes[0])
		switch eventType {
		case rawTxZMQCommand:
			tx := &wire.MsgTx{}
			r := bytes.NewReader(msgBytes[1])
			if err := tx.Deserialize(r); err != nil {
				log.Errorf("Unable to deserialize "+
					"transaction: %v", err)
				continue
			}

			c.rescanClientsMtx.Lock()
			for _, client := range c.rescanClients {
				select {
				case client.zmqTxNtfns <- tx:
				case <-client.quit:
				case <-c.quit:
					c.rescanClientsMtx.Unlock()
					return
				}
			}
			c.rescanClientsMtx.Unlock()
		default:
			// It's possible that the message wasn't fully read if
			// bitcoind shuts down, which will produce an unreadable
			// event type. To prevent from logging it, we'll make
			// sure it conforms to the ASCII standard.
			if eventType == "" || !isASCII(eventType) {
				continue
			}

			log.Warnf("Received unexpected event type from %v "+
				"subscription: %v", rawTxZMQCommand, eventType)
		}
	}
}

// getCurrentNet returns the network on which the bitcoind node is running.
func (c *BitcoindConn) getCurrentNet() (wire.BitcoinNet, error) {
	hash, err := c.client.GetBlockHash(0)
	if err != nil {
		return 0, err
	}

	switch *hash {
	case *chaincfg.TestNet3Params.GenesisHash:
		return chaincfg.TestNet3Params.Net, nil
	case *chaincfg.RegressionNetParams.GenesisHash:
		return chaincfg.RegressionNetParams.Net, nil
	case *chaincfg.MainNetParams.GenesisHash:
		return chaincfg.MainNetParams.Net, nil
	default:
		return 0, fmt.Errorf("unknown network with genesis hash %v", hash)
	}
}

// NewBitcoindClient returns a bitcoind client using the current bitcoind
// connection. This allows us to share the same connection using multiple
// clients.
func (c *BitcoindConn) NewBitcoindClient() *BitcoindClient {
	return &BitcoindClient{
		quit: make(chan struct{}),

		id: atomic.AddUint64(&c.rescanClientCounter, 1),

		chainParams: c.chainParams,
		chainConn:   c,

		rescanUpdate:     make(chan interface{}),
		watchedAddresses: make(map[string]struct{}),
		watchedOutPoints: make(map[wire.OutPoint]struct{}),
		watchedTxs:       make(map[chainhash.Hash]struct{}),

		notificationQueue: NewConcurrentQueue(20),
		zmqTxNtfns:        make(chan *wire.MsgTx),
		zmqBlockNtfns:     make(chan *wire.MsgBlock),

		mempool:        make(map[chainhash.Hash]struct{}),
		expiredMempool: make(map[int32]map[chainhash.Hash]struct{}),
	}
}

// AddClient adds a client to the set of active rescan clients of the current
// chain connection. This allows the connection to include the specified client
// in its notification delivery.
//
// NOTE: This function is safe for concurrent access.
func (c *BitcoindConn) AddClient(client *BitcoindClient) {
	c.rescanClientsMtx.Lock()
	defer c.rescanClientsMtx.Unlock()

	c.rescanClients[client.id] = client
}

// RemoveClient removes the client with the given ID from the set of active
// rescan clients. Once removed, the client will no longer receive block and
// transaction notifications from the chain connection.
//
// NOTE: This function is safe for concurrent access.
func (c *BitcoindConn) RemoveClient(id uint64) {
	c.rescanClientsMtx.Lock()
	defer c.rescanClientsMtx.Unlock()

	delete(c.rescanClients, id)
}

// isASCII is a helper method that checks whether all bytes in `data` would be
// printable ASCII characters if interpreted as a string.
func isASCII(s string) bool {
	for _, c := range s {
		if c < 32 || c > 126 {
			return false
		}
	}
	return true
}
<|endoftext|>


<commit_before>package bot

import (
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"

	"github.com/belak/irc"
	"github.com/khades/servbot/commandHandlers"
	"github.com/khades/servbot/ircClient"
	"github.com/khades/servbot/models"
	"github.com/khades/servbot/repos"
)

var chatHandler irc.HandlerFunc = func(client *irc.Client, message *irc.Message) {
	//	log.Println(message.String())
	msgID, found := message.Tags.GetTag("msg-id")
	if found {
		switch msgID {
		case "room_mods":
			{
				commaIndex := strings.Index(message.Params[1], ":")
				if commaIndex != -1 {
					//				log.Printf("Channel %v: got mods list", message.Params[0])
					mods := strings.Split(message.Params[1][commaIndex+2:], ", ")
					repos.PushMods(message.Params[0][1:], mods)
				}
			}
		case "resub":
			{
				msgParamMonths, msgParamMonthsFound := message.Tags.GetTag("msg-param-months")
				user, userFound := message.Tags.GetTag("display-name")
				channel := message.Params[0]
				if msgParamMonthsFound && userFound && channel != "" {
					resubCount, resubCountError := strconv.Atoi(msgParamMonths)
					if resubCountError == nil {
						formedMessage := models.ChatMessage{
							Channel:          channel,
							User:             user,
							IsMod:            false,
							IsSub:            true,
							Date:             time.Now(),
							SubscriptionInfo: &models.SubscriptionInfo{Count: resubCount}}
						repos.LogMessage(formedMessage)
						channelInfo, error := repos.GetChannelInfo(channel)
						if error == nil && channelInfo.SubAlert.Enabled == true {
							messageBody := strings.TrimSpace(fmt.Sprintf("%s %s%s",
								channelInfo.SubAlert.RepeatPrefix,
								strings.Repeat(channelInfo.SubAlert.RepeatBody+" ", formedMessage.SubscriptionInfo.Count),
								channelInfo.SubAlert.RepeatPostfix))
							if messageBody != "" {
								IrcClientInstance.SendPublic(models.OutgoingMessage{
									Body:    messageBody,
									Channel: channel,
									User:    user})
							}
						}
						log.Printf("Channel %v: %v resubbed for %v months\n", formedMessage.Channel, formedMessage.User, formedMessage.SubscriptionInfo.Count)
					}
				}
			}
		}
	}

	if message.User == "twitchnotify" {
		log.Println("Got first sub")
		user := strings.Split(message.Params[1], " ")[0]
		channel := message.Params[0]
		if user != "" && channel != "" {
			formedMessage := models.ChatMessage{
				Channel:          channel,
				User:             user,
				IsMod:            false,
				IsSub:            true,
				IsPrime:          strings.Contains(message.String(), "Twitch Prime"),
				Date:             time.Now(),
				SubscriptionInfo: &models.SubscriptionInfo{Count: 1}}
			repos.LogMessage(formedMessage)
			channelInfo, error := repos.GetChannelInfo(channel)
			if error == nil && channelInfo.SubAlert.Enabled == true && channelInfo.SubAlert.FirstMessage != "" {
				IrcClientInstance.SendPublic(models.OutgoingMessage{
					Body:    channelInfo.SubAlert.FirstMessage,
					Channel: channel,
					User:    user})
			}
			log.Printf("Channel %v: %v subbed\n", formedMessage.Channel, formedMessage.User)
		}
	}
	if message.Command == "CLEARCHAT" {
		banDuration, banDurationFound := message.Tags.GetTag("ban-duration")
		intBanDuration := 0
		if banDurationFound {
			parsedValue, parseError := strconv.Atoi(banDuration)
			if parseError == nil {
				intBanDuration = parsedValue
			}
		}
		banReason, _ := message.Tags.GetTag("ban-reason")
		user := message.Params[1]
		channel := message.Params[0]
		formedMessage := models.ChatMessage{
			Channel: channel,
			User:    user,
			IsMod:   false,
			IsSub:   true,
			Date:    time.Now(),
			BanInfo: &models.BanInfo{Duration: intBanDuration, Reason: banReason}}
		repos.LogMessage(formedMessage)
		//	log.Printf("Channel %v: %v is banned for %v \n", channel, user, intBanDuration)
	}
	if message.Command == "PRIVMSG" {
		formedMessage := models.ChatMessage{
			Channel:     message.Params[0][1:],
			User:        message.User,
			MessageBody: message.Params[1],
			IsMod:       message.Tags["mod"] == "1" || message.User == "khadesru",
			IsSub:       message.Tags["subscriber"] == "1",
			IsPrime:     strings.Contains(message.Tags["badges"].Encode(), "premium/1"),
			Date:        time.Now()}
		repos.LogMessage(formedMessage)
		isCommand, commandBody := formedMessage.IsCommand()
		if isCommand {
			handlerFunction := commandHandlers.Router.Go(commandBody.Command)
			handlerFunction(true, &formedMessage, commandBody, &IrcClientInstance)
		}
	}

	if message.Command == "001" {
		client.Write("CAP REQ twitch.tv/tags")
		client.Write("CAP REQ twitch.tv/membership")
		client.Write("CAP REQ twitch.tv/commands")
		for _, value := range repos.Config.Channels {
			client.Write("JOIN #" + value)
		}
		IrcClientInstance = ircClient.IrcClient{Client: client, Bounces: make(map[string]time.Time), Ready: true}
		IrcClientInstance.SendModsCommand()
		log.Println("Bot is started")
	}
}
<commit_msg>Fixed subalert, removing channel Hashtag<commit_after>package bot

import (
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"

	"github.com/belak/irc"
	"github.com/khades/servbot/commandHandlers"
	"github.com/khades/servbot/ircClient"
	"github.com/khades/servbot/models"
	"github.com/khades/servbot/repos"
)

var chatHandler irc.HandlerFunc = func(client *irc.Client, message *irc.Message) {
	//	log.Println(message.String())
	msgID, found := message.Tags.GetTag("msg-id")
	if found {
		switch msgID {
		case "room_mods":
			{
				commaIndex := strings.Index(message.Params[1], ":")
				if commaIndex != -1 {
					//				log.Printf("Channel %v: got mods list", message.Params[0])
					mods := strings.Split(message.Params[1][commaIndex+2:], ", ")
					repos.PushMods(message.Params[0][1:], mods)
				}
			}
		case "resub":
			{
				msgParamMonths, msgParamMonthsFound := message.Tags.GetTag("msg-param-months")
				user, userFound := message.Tags.GetTag("display-name")
				channel := message.Params[0][1:]
				if msgParamMonthsFound && userFound && channel != "" {
					resubCount, resubCountError := strconv.Atoi(msgParamMonths)
					if resubCountError == nil {
						formedMessage := models.ChatMessage{
							Channel:          channel,
							User:             user,
							IsMod:            false,
							IsSub:            true,
							Date:             time.Now(),
							SubscriptionInfo: &models.SubscriptionInfo{Count: resubCount}}
						repos.LogMessage(formedMessage)
						channelInfo, error := repos.GetChannelInfo(channel)
						if error == nil && channelInfo.SubAlert.Enabled == true {
							messageBody := strings.TrimSpace(fmt.Sprintf("%s %s%s",
								channelInfo.SubAlert.RepeatPrefix,
								strings.Repeat(channelInfo.SubAlert.RepeatBody+" ", formedMessage.SubscriptionInfo.Count),
								channelInfo.SubAlert.RepeatPostfix))
							if messageBody != "" {
								IrcClientInstance.SendPublic(models.OutgoingMessage{
									Body:    messageBody,
									Channel: channel,
									User:    user})
							}
						}
						log.Printf("Channel %v: %v resubbed for %v months\n", formedMessage.Channel, formedMessage.User, formedMessage.SubscriptionInfo.Count)
					}
				}
			}
		}
	}

	if message.User == "twitchnotify" {
		log.Println("Got first sub")
		user := strings.Split(message.Params[1], " ")[0]
		channel := message.Params[0][1:]
		if user != "" && channel != "" {
			formedMessage := models.ChatMessage{
				Channel:          channel,
				User:             user,
				IsMod:            false,
				IsSub:            true,
				IsPrime:          strings.Contains(message.String(), "Twitch Prime"),
				Date:             time.Now(),
				SubscriptionInfo: &models.SubscriptionInfo{Count: 1}}
			repos.LogMessage(formedMessage)
			channelInfo, error := repos.GetChannelInfo(channel)
			if error == nil && channelInfo.SubAlert.Enabled == true && channelInfo.SubAlert.FirstMessage != "" {
				IrcClientInstance.SendPublic(models.OutgoingMessage{
					Body:    channelInfo.SubAlert.FirstMessage,
					Channel: channel,
					User:    user})
			}
			log.Printf("Channel %v: %v subbed\n", formedMessage.Channel, formedMessage.User)
		}
	}
	if message.Command == "CLEARCHAT" {
		banDuration, banDurationFound := message.Tags.GetTag("ban-duration")
		intBanDuration := 0
		if banDurationFound {
			parsedValue, parseError := strconv.Atoi(banDuration)
			if parseError == nil {
				intBanDuration = parsedValue
			}
		}
		banReason, _ := message.Tags.GetTag("ban-reason")
		user := message.Params[1]
		channel := message.Params[0]
		formedMessage := models.ChatMessage{
			Channel: channel,
			User:    user,
			IsMod:   false,
			IsSub:   true,
			Date:    time.Now(),
			BanInfo: &models.BanInfo{Duration: intBanDuration, Reason: banReason}}
		repos.LogMessage(formedMessage)
		//	log.Printf("Channel %v: %v is banned for %v \n", channel, user, intBanDuration)
	}
	if message.Command == "PRIVMSG" {
		formedMessage := models.ChatMessage{
			Channel:     message.Params[0][1:],
			User:        message.User,
			MessageBody: message.Params[1],
			IsMod:       message.Tags["mod"] == "1" || message.User == "khadesru",
			IsSub:       message.Tags["subscriber"] == "1",
			IsPrime:     strings.Contains(message.Tags["badges"].Encode(), "premium/1"),
			Date:        time.Now()}
		repos.LogMessage(formedMessage)
		isCommand, commandBody := formedMessage.IsCommand()
		if isCommand {
			handlerFunction := commandHandlers.Router.Go(commandBody.Command)
			handlerFunction(true, &formedMessage, commandBody, &IrcClientInstance)
		}
	}

	if message.Command == "001" {
		client.Write("CAP REQ twitch.tv/tags")
		client.Write("CAP REQ twitch.tv/membership")
		client.Write("CAP REQ twitch.tv/commands")
		for _, value := range repos.Config.Channels {
			client.Write("JOIN #" + value)
		}
		IrcClientInstance = ircClient.IrcClient{Client: client, Bounces: make(map[string]time.Time), Ready: true}
		IrcClientInstance.SendModsCommand()
		log.Println("Bot is started")
	}
}
<|endoftext|>


<commit_before>package bot

import "github.com/yuin/gopher-lua"
import "github.com/jriddick/geoffrey/msg"

// RegisterMessage will register msg.Message struct to Lua
func RegisterMessage(state *lua.LState) {
	// Create the metatable
	meta := state.NewTypeMetatable("config")
	state.SetGlobal("message", meta)

	// Bind our functions
	state.SetField(meta, "__index", state.NewFunction(messageIndex))
}

// PushMessage will push the given msg.Message to the Lua stack
func PushMessage(msg *msg.Message, state *lua.LState) {
	// Create the userdata
	data := state.NewUserData()
	data.Value = msg

	// Set the metatable
	state.SetMetatable(data, state.GetTypeMetatable("message"))

	// Push the message
	state.Push(data)
}

func checkMessage(state *lua.LState) *msg.Message {
	data := state.CheckUserData(1)
	if v, ok := data.Value.(*msg.Message); ok {
		return v
	}
	state.ArgError(1, "message expected")
	return nil
}

func messageIndex(state *lua.LState) int {
	msg := checkMessage(state)
	key := state.CheckString(2)

	switch key {
	case "Command":
		state.Push(lua.LString(msg.Command))
	case "Params":
		params := state.NewTable()
		for key, param := range msg.Params {
			params.RawSetInt(key, lua.LString(param))
		}
		state.Push(params)
	case "Trailing":
		state.Push(lua.LString(msg.Trailing))
	default:
		state.Push(lua.LNil)
	}

	return 1
}
<commit_msg>added bindings for msg.Tags and msg.Prefix<commit_after>package bot

import "github.com/yuin/gopher-lua"
import "github.com/jriddick/geoffrey/msg"

// RegisterMessage will register msg.Message struct to Lua
func RegisterMessage(state *lua.LState) {
	// Create the metatable
	meta := state.NewTypeMetatable("message")
	state.SetGlobal("message", meta)

	// Bind the message index function
	state.SetField(meta, "__index", state.NewFunction(messageIndex))

	// Create the metatable for prefix struct
	prefix := state.NewTypeMetatable("prefix")
	state.SetGlobal("prefix", prefix)

	// Bind the prefix index function
	state.SetField(prefix, "__index", state.NewFunction(prefixIndex))
}

// PushMessage will push the given msg.Message to the Lua stack
func PushMessage(msg *msg.Message, state *lua.LState) {
	// Create the userdata
	data := state.NewUserData()
	data.Value = msg

	// Set the metatable
	state.SetMetatable(data, state.GetTypeMetatable("message"))

	// Push the message
	state.Push(data)
}

func pushPrefix(prefix *msg.Prefix, state *lua.LState) {
	// Create the userdata
	data := state.NewUserData()
	data.Value = prefix

	// Set the metatable
	state.SetMetatable(data, state.GetTypeMetatable("prefix"))

	// Push the prefix
	state.Push(data)
}

func checkMessage(state *lua.LState) *msg.Message {
	data := state.CheckUserData(1)
	if v, ok := data.Value.(*msg.Message); ok {
		return v
	}
	state.ArgError(1, "message expected")
	return nil
}

func checkPrefix(state *lua.LState) *msg.Prefix {
	data := state.CheckUserData(1)
	if v, ok := data.Value.(*msg.Prefix); ok {
		return v
	}
	state.ArgError(1, "prefix expected")
	return nil
}

func messageIndex(state *lua.LState) int {
	msg := checkMessage(state)
	key := state.CheckString(2)

	switch key {
	case "Command":
		state.Push(lua.LString(msg.Command))
	case "Params":
		params := state.NewTable()
		for key, param := range msg.Params {
			params.RawSetInt(key, lua.LString(param))
		}
		state.Push(params)
	case "Trailing":
		state.Push(lua.LString(msg.Trailing))
	case "Prefix":
		pushPrefix(msg.Prefix, state)
	case "Tags":
		tags := state.NewTable()
		for key, value := range msg.Tags {
			tags.RawSetString(key, lua.LString(value))
		}
		state.Push(tags)
	default:
		state.Push(lua.LNil)
	}

	return 1
}

func prefixIndex(state *lua.LState) int {
	prefix := checkPrefix(state)
	key := state.CheckString(2)

	switch key {
	case "Name":
		state.Push(lua.LString(prefix.Name))
	case "User":
		state.Push(lua.LString(prefix.User))
	case "Host":
		state.Push(lua.LString(prefix.Host))
	default:
		state.Push(lua.LNil)
	}

	return 1
}
<|endoftext|>


<commit_before>// Copyright 2018 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 httpreplay provides an API for recording and replaying traffic
// from HTTP-based Google API clients.
//
// To record:
// 1.  Call NewRecorder to get a Recorder.
// 2.  Use its Client method to obtain an HTTP client to use when making API calls.
// 3.  Close the Recorder when you're done. That will save the
//     log of interactions to the file you provided to NewRecorder.
//
// To replay:
// 1.  Call NewReplayer with the same filename you used to record to get a Replayer.
// 2.  Call its Client method and use the client to make the same API calls.
//     You will get back the recorded responses.
// 3.  Close the Replayer when you're done.
//
// This package is EXPERIMENTAL and is subject to change or removal without notice.
// It requires Go version 1.8 or higher.
package httpreplay

// TODO(jba): add examples.

import (
	"context"
	"net/http"

	"cloud.google.com/go/httpreplay/internal/proxy"
	"google.golang.org/api/option"
	htransport "google.golang.org/api/transport/http"
)

// A Recorder records HTTP interactions.
type Recorder struct {
	proxy *proxy.Proxy
}

// NewRecorder creates a recorder that writes to filename. The file will
// also store initial state that can be retrieved to configure replay.
//
// You must call Close on the Recorder to ensure that all data is written.
func NewRecorder(filename string, initial []byte) (*Recorder, error) {
	p, err := proxy.ForRecording(filename, 0)
	if err != nil {
		return nil, err
	}
	p.Initial = initial
	return &Recorder{proxy: p}, nil
}

// RemoveRequestHeaders will remove request headers matching patterns from the log,
// and skip matching them during replay.
//
// Pattern is taken literally except for *, which matches any sequence of characters.
func (r *Recorder) RemoveRequestHeaders(patterns ...string) {
	r.proxy.RemoveRequestHeaders(patterns)
}

// ClearHeaders will replace the value of request and response headers that match
// any of the patterns with CLEARED, on both recording and replay.
// Use ClearHeaders when the header information is secret or may change from run to
// run, but you still want to verify that the headers are being sent and received.
//
// Pattern is taken literally except for *, which matches any sequence of characters.
func (r *Recorder) ClearHeaders(patterns ...string) {
	r.proxy.ClearHeaders(patterns)
}

// RemoveQueryParams will remove URL query parameters matching patterns from the log,
// and skip matching them during replay.
//
// Pattern is taken literally except for *, which matches any sequence of characters.
func (r *Recorder) RemoveQueryParams(patterns ...string) {
	r.proxy.RemoveQueryParams(patterns)
}

// ClearQueryParams will replace the value of URL query parametrs that match any of
// the patterns with CLEARED, on both recording and replay.
// Use ClearQueryParams when the parameter information is secret or may change from
// run to run, but you still want to verify that it are being sent.
//
// Pattern is taken literally except for *, which matches any sequence of characters.
func (r *Recorder) ClearQueryParams(patterns ...string) {
	r.proxy.ClearQueryParams(patterns)
}

// Client returns an http.Client to be used for recording. Provide authentication options
// like option.WithTokenSource as you normally would, or omit them to use Application Default
// Credentials.
func (r *Recorder) Client(ctx context.Context, opts ...option.ClientOption) (*http.Client, error) {
	return proxyClient(ctx, r.proxy, opts...)
}

func proxyClient(ctx context.Context, p *proxy.Proxy, opts ...option.ClientOption) (*http.Client, error) {
	trans, err := htransport.NewTransport(ctx, p.Transport(), opts...)
	if err != nil {
		return nil, err
	}
	return &http.Client{Transport: trans}, nil
}

// Close closes the Recorder and saves the log file.
func (r *Recorder) Close() error {
	return r.proxy.Close()
}

// A Replayer replays previously recorded HTTP interactions.
type Replayer struct {
	proxy *proxy.Proxy
}

// NewReplayer creates a replayer that reads from filename.
func NewReplayer(filename string) (*Replayer, error) {
	p, err := proxy.ForReplaying(filename, 0)
	if err != nil {
		return nil, err
	}
	return &Replayer{proxy: p}, nil
}

// Client returns an HTTP client for replaying. The client does not need to be
// configured with credentials for authenticating to a server, since it never
// contacts a real backend.
func (r *Replayer) Client(ctx context.Context) (*http.Client, error) {
	return proxyClient(ctx, r.proxy, option.WithoutAuthentication())
}

// Initial returns the initial state saved by the Recorder.
func (r *Replayer) Initial() []byte {
	return r.proxy.Initial
}

// IgnoreHeader will not use h when matching requests.
func (r *Replayer) IgnoreHeader(h string) {
	r.proxy.IgnoreHeader(h)
}

// Close closes the replayer.
func (r *Replayer) Close() error {
	return r.proxy.Close()
}

// DebugHeaders helps to determine whether a header should be ignored.
// When true, if requests have the same method, URL and body but differ
// in a header, the first mismatched header is logged.
func DebugHeaders() {
	proxy.DebugHeaders = true
}

// Supported reports whether httpreplay is supported in the current version of Go.
// For Go 1.8 and above, the answer is true.
func Supported() bool { return true }
<commit_msg>httpreplay: fix godoc formatting<commit_after>// Copyright 2018 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 httpreplay provides an API for recording and replaying traffic
// from HTTP-based Google API clients.
//
// To record:
//  1.  Call NewRecorder to get a Recorder.
//  2.  Use its Client method to obtain an HTTP client to use when making API calls.
//  3.  Close the Recorder when you're done. That will save the log of interactions
//      to the file you provided to NewRecorder.
//
// To replay:
//  1.  Call NewReplayer with the same filename you used to record to get a Replayer.
//  2.  Call its Client method and use the client to make the same API calls.
//      You will get back the recorded responses.
//  3.  Close the Replayer when you're done.
//
// This package is EXPERIMENTAL and is subject to change or removal without notice.
// It requires Go version 1.8 or higher.
package httpreplay

// TODO(jba): add examples.

import (
	"context"
	"net/http"

	"cloud.google.com/go/httpreplay/internal/proxy"
	"google.golang.org/api/option"
	htransport "google.golang.org/api/transport/http"
)

// A Recorder records HTTP interactions.
type Recorder struct {
	proxy *proxy.Proxy
}

// NewRecorder creates a recorder that writes to filename. The file will
// also store initial state that can be retrieved to configure replay.
//
// You must call Close on the Recorder to ensure that all data is written.
func NewRecorder(filename string, initial []byte) (*Recorder, error) {
	p, err := proxy.ForRecording(filename, 0)
	if err != nil {
		return nil, err
	}
	p.Initial = initial
	return &Recorder{proxy: p}, nil
}

// RemoveRequestHeaders will remove request headers matching patterns from the log,
// and skip matching them during replay.
//
// Pattern is taken literally except for *, which matches any sequence of characters.
func (r *Recorder) RemoveRequestHeaders(patterns ...string) {
	r.proxy.RemoveRequestHeaders(patterns)
}

// ClearHeaders will replace the value of request and response headers that match
// any of the patterns with CLEARED, on both recording and replay.
// Use ClearHeaders when the header information is secret or may change from run to
// run, but you still want to verify that the headers are being sent and received.
//
// Pattern is taken literally except for *, which matches any sequence of characters.
func (r *Recorder) ClearHeaders(patterns ...string) {
	r.proxy.ClearHeaders(patterns)
}

// RemoveQueryParams will remove URL query parameters matching patterns from the log,
// and skip matching them during replay.
//
// Pattern is taken literally except for *, which matches any sequence of characters.
func (r *Recorder) RemoveQueryParams(patterns ...string) {
	r.proxy.RemoveQueryParams(patterns)
}

// ClearQueryParams will replace the value of URL query parametrs that match any of
// the patterns with CLEARED, on both recording and replay.
// Use ClearQueryParams when the parameter information is secret or may change from
// run to run, but you still want to verify that it are being sent.
//
// Pattern is taken literally except for *, which matches any sequence of characters.
func (r *Recorder) ClearQueryParams(patterns ...string) {
	r.proxy.ClearQueryParams(patterns)
}

// Client returns an http.Client to be used for recording. Provide authentication options
// like option.WithTokenSource as you normally would, or omit them to use Application Default
// Credentials.
func (r *Recorder) Client(ctx context.Context, opts ...option.ClientOption) (*http.Client, error) {
	return proxyClient(ctx, r.proxy, opts...)
}

func proxyClient(ctx context.Context, p *proxy.Proxy, opts ...option.ClientOption) (*http.Client, error) {
	trans, err := htransport.NewTransport(ctx, p.Transport(), opts...)
	if err != nil {
		return nil, err
	}
	return &http.Client{Transport: trans}, nil
}

// Close closes the Recorder and saves the log file.
func (r *Recorder) Close() error {
	return r.proxy.Close()
}

// A Replayer replays previously recorded HTTP interactions.
type Replayer struct {
	proxy *proxy.Proxy
}

// NewReplayer creates a replayer that reads from filename.
func NewReplayer(filename string) (*Replayer, error) {
	p, err := proxy.ForReplaying(filename, 0)
	if err != nil {
		return nil, err
	}
	return &Replayer{proxy: p}, nil
}

// Client returns an HTTP client for replaying. The client does not need to be
// configured with credentials for authenticating to a server, since it never
// contacts a real backend.
func (r *Replayer) Client(ctx context.Context) (*http.Client, error) {
	return proxyClient(ctx, r.proxy, option.WithoutAuthentication())
}

// Initial returns the initial state saved by the Recorder.
func (r *Replayer) Initial() []byte {
	return r.proxy.Initial
}

// IgnoreHeader will not use h when matching requests.
func (r *Replayer) IgnoreHeader(h string) {
	r.proxy.IgnoreHeader(h)
}

// Close closes the replayer.
func (r *Replayer) Close() error {
	return r.proxy.Close()
}

// DebugHeaders helps to determine whether a header should be ignored.
// When true, if requests have the same method, URL and body but differ
// in a header, the first mismatched header is logged.
func DebugHeaders() {
	proxy.DebugHeaders = true
}

// Supported reports whether httpreplay is supported in the current version of Go.
// For Go 1.8 and above, the answer is true.
func Supported() bool { return true }
<|endoftext|>


<commit_before>package ilium

import "sort"

type Distribution1D struct {
	f, cdf []float32
	intF   float32
}

func (d *Distribution1D) SampleDiscrete(u float32) (i int, p float32) {
	greaterThanR := func(i int) bool { return d.cdf[i+1] > u }
	n := len(d.f)
	i = sort.Search(n, greaterThanR)
	p = d.f[i] / (d.intF * float32(n))
	return
}

func MakeDistribution1D(f []float32) Distribution1D {
	n := len(f)
	cdf := make([]float32, n+1)
	for i := 1; i < n+1; i++ {
		cdf[i] = cdf[i-1] + f[i-1]/float32(n)
	}
	intF := cdf[n]
	if intF == 0 {
		for i := 1; i < n; i++ {
			cdf[i] = float32(i) / float32(n)
		}
	} else {
		for i := 1; i < n; i++ {
			cdf[i] /= intF
		}
	}
	cdf[n] = 1
	return Distribution1D{f, cdf, intF}
}
<commit_msg>Implement Distribution1D.ComputeDiscretePdf()<commit_after>package ilium

import "sort"

type Distribution1D struct {
	f, cdf []float32
	intF   float32
}

func (d *Distribution1D) SampleDiscrete(u float32) (i int, p float32) {
	greaterThanR := func(i int) bool { return d.cdf[i+1] > u }
	n := len(d.f)
	i = sort.Search(n, greaterThanR)
	p = d.f[i] / (d.intF * float32(n))
	return
}

func (d *Distribution1D) ComputeDiscretePdf(i int) float32 {
	return d.cdf[i+1] - d.cdf[i]
}

func MakeDistribution1D(f []float32) Distribution1D {
	n := len(f)
	cdf := make([]float32, n+1)
	for i := 1; i < n+1; i++ {
		cdf[i] = cdf[i-1] + f[i-1]/float32(n)
	}
	intF := cdf[n]
	if intF == 0 {
		for i := 1; i < n; i++ {
			cdf[i] = float32(i) / float32(n)
		}
	} else {
		for i := 1; i < n; i++ {
			cdf[i] /= intF
		}
	}
	cdf[n] = 1
	return Distribution1D{f, cdf, intF}
}
<|endoftext|>


<commit_before>package handlers

import (
	"github.com/getfider/fider/app/actions"
	"github.com/getfider/fider/app/models/cmd"
	"github.com/getfider/fider/app/models/query"
	"github.com/getfider/fider/app/pkg/bus"
	"github.com/getfider/fider/app/pkg/web"
)

// BillingPage is the billing settings page
func BillingPage() web.HandlerFunc {
	return func(c *web.Context) error {
		if c.Tenant().Billing.StripeCustomerID == "" {
			if err := bus.Dispatch(c, &cmd.CreateBillingCustomer{}); err != nil {
				return c.Failure(err)
			}

			if err := bus.Dispatch(c, &cmd.UpdateTenantBillingSettings{
				Settings: c.Tenant().Billing,
			}); err != nil {
				return c.Failure(err)
			}
		}

		getUpcomingInvoiceQuery := &query.GetUpcomingInvoice{}
		if c.Tenant().Billing.StripeSubscriptionID != "" {
			err := bus.Dispatch(c, getUpcomingInvoiceQuery)
			if err != nil {
				return c.Failure(err)
			}
		}

		paymentInfo := query.GetPaymentInfo{}
		err := bus.Dispatch(c, paymentInfo)
		if err != nil {
			return c.Failure(err)
		}

		listPlansQuery := &query.ListBillingPlans{}
		if paymentInfo.Result != nil {
			listPlansQuery.CountryCode = paymentInfo.Result.AddressCountry
			err = bus.Dispatch(c, listPlansQuery)
			if err != nil {
				println(err.Error())
				return c.Failure(err)
			}
		}

		countUsers := &query.CountUsers{}
		allCountries := &query.GetAllCountries{}
		if err := bus.Dispatch(c, countUsers); err != nil {
			return c.Failure(err)
		}

		return c.Page(web.Props{
			Title:     "Billing · Site Settings",
			ChunkName: "Billing.page",
			Data: web.Map{
				"invoiceDue":      getUpcomingInvoiceQuery.Result,
				"tenantUserCount": countUsers.Result,
				"plans":           listPlansQuery.Result,
				"paymentInfo":     paymentInfo.Result,
				"countries":       allCountries.Result,
			},
		})
	}
}

// UpdatePaymentInfo on stripe based on given input
func UpdatePaymentInfo() web.HandlerFunc {
	return func(c *web.Context) error {
		input := new(actions.CreateEditBillingPaymentInfo)
		if result := c.BindTo(input); !result.Ok {
			return c.HandleValidation(result)
		}

		if err := bus.Dispatch(c, &cmd.UpdatePaymentInfo{Input: input.Model}); err != nil {
			return c.Failure(err)
		}

		return c.Ok(web.Map{})
	}
}

// GetBillingPlans returns a list of plans for given country code
func GetBillingPlans() web.HandlerFunc {
	return func(c *web.Context) error {
		countryCode := c.Param("countryCode")
		listPlansQuery := &query.ListBillingPlans{CountryCode: countryCode}
		err := bus.Dispatch(c, listPlansQuery)
		if err != nil {
			return c.Failure(err)
		}
		return c.Ok(listPlansQuery.Result)
	}
}

// BillingSubscribe subscribes current tenant to given plan on stripe
func BillingSubscribe() web.HandlerFunc {
	return func(c *web.Context) error {
		planID := c.Param("planID")

		paymentInfoQuery := &query.GetPaymentInfo{}
		err := bus.Dispatch(c, paymentInfoQuery)
		if err != nil {
			return c.Failure(err)
		}

		getPlanByIDQuery := &query.GetBillingPlanByID{
			PlanID:      planID,
			CountryCode: paymentInfoQuery.Result.AddressCountry,
		}
		err = bus.Dispatch(c, getPlanByIDQuery)
		if err != nil {
			return c.Failure(err)
		}
		plan := getPlanByIDQuery.Result

		countUsers := &query.CountUsers{}
		err = bus.Dispatch(c, countUsers)
		if err != nil {
			return c.Failure(err)
		}

		if plan.MaxUsers > 0 && countUsers.Result > plan.MaxUsers {
			return c.Unauthorized()
		}

		if err = bus.Dispatch(c, &cmd.CreateBillingSubscription{
			PlanID: plan.ID,
		}); err != nil {
			return c.Failure(err)
		}

		updateBilling := &cmd.UpdateTenantBillingSettings{Settings: c.Tenant().Billing}
		activateTenant := &cmd.ActivateTenant{TenantID: c.Tenant().ID}
		if err := bus.Dispatch(c, updateBilling, activateTenant); err != nil {
			return c.Failure(err)
		}

		return c.Ok(web.Map{})
	}
}

// CancelBillingSubscription cancels current subscription from current tenant
func CancelBillingSubscription() web.HandlerFunc {
	return func(c *web.Context) error {
		err := bus.Dispatch(c, &cmd.CancelBillingSubscription{})
		if err != nil {
			return c.Failure(err)
		}

		if err := bus.Dispatch(c, &cmd.UpdateTenantBillingSettings{
			Settings: c.Tenant().Billing,
		}); err != nil {
			return c.Failure(err)
		}

		return c.Ok(web.Map{})
	}
}
<commit_msg>fix: pointer issue in billing (#901)<commit_after>package handlers

import (
	"github.com/getfider/fider/app/actions"
	"github.com/getfider/fider/app/models/cmd"
	"github.com/getfider/fider/app/models/query"
	"github.com/getfider/fider/app/pkg/bus"
	"github.com/getfider/fider/app/pkg/web"
)

// BillingPage is the billing settings page
func BillingPage() web.HandlerFunc {
	return func(c *web.Context) error {
		if c.Tenant().Billing.StripeCustomerID == "" {
			if err := bus.Dispatch(c, &cmd.CreateBillingCustomer{}); err != nil {
				return c.Failure(err)
			}

			if err := bus.Dispatch(c, &cmd.UpdateTenantBillingSettings{
				Settings: c.Tenant().Billing,
			}); err != nil {
				return c.Failure(err)
			}
		}

		getUpcomingInvoiceQuery := &query.GetUpcomingInvoice{}
		if c.Tenant().Billing.StripeSubscriptionID != "" {
			err := bus.Dispatch(c, getUpcomingInvoiceQuery)
			if err != nil {
				return c.Failure(err)
			}
		}

		paymentInfo := &query.GetPaymentInfo{}
		err := bus.Dispatch(c, paymentInfo)
		if err != nil {
			return c.Failure(err)
		}

		listPlansQuery := &query.ListBillingPlans{}
		if paymentInfo.Result != nil {
			listPlansQuery.CountryCode = paymentInfo.Result.AddressCountry
			err = bus.Dispatch(c, listPlansQuery)
			if err != nil {
				println(err.Error())
				return c.Failure(err)
			}
		}

		countUsers := &query.CountUsers{}
		allCountries := &query.GetAllCountries{}
		if err := bus.Dispatch(c, countUsers); err != nil {
			return c.Failure(err)
		}

		return c.Page(web.Props{
			Title:     "Billing · Site Settings",
			ChunkName: "Billing.page",
			Data: web.Map{
				"invoiceDue":      getUpcomingInvoiceQuery.Result,
				"tenantUserCount": countUsers.Result,
				"plans":           listPlansQuery.Result,
				"paymentInfo":     paymentInfo.Result,
				"countries":       allCountries.Result,
			},
		})
	}
}

// UpdatePaymentInfo on stripe based on given input
func UpdatePaymentInfo() web.HandlerFunc {
	return func(c *web.Context) error {
		input := new(actions.CreateEditBillingPaymentInfo)
		if result := c.BindTo(input); !result.Ok {
			return c.HandleValidation(result)
		}

		if err := bus.Dispatch(c, &cmd.UpdatePaymentInfo{Input: input.Model}); err != nil {
			return c.Failure(err)
		}

		return c.Ok(web.Map{})
	}
}

// GetBillingPlans returns a list of plans for given country code
func GetBillingPlans() web.HandlerFunc {
	return func(c *web.Context) error {
		countryCode := c.Param("countryCode")
		listPlansQuery := &query.ListBillingPlans{CountryCode: countryCode}
		err := bus.Dispatch(c, listPlansQuery)
		if err != nil {
			return c.Failure(err)
		}
		return c.Ok(listPlansQuery.Result)
	}
}

// BillingSubscribe subscribes current tenant to given plan on stripe
func BillingSubscribe() web.HandlerFunc {
	return func(c *web.Context) error {
		planID := c.Param("planID")

		paymentInfoQuery := &query.GetPaymentInfo{}
		err := bus.Dispatch(c, paymentInfoQuery)
		if err != nil {
			return c.Failure(err)
		}

		getPlanByIDQuery := &query.GetBillingPlanByID{
			PlanID:      planID,
			CountryCode: paymentInfoQuery.Result.AddressCountry,
		}
		err = bus.Dispatch(c, getPlanByIDQuery)
		if err != nil {
			return c.Failure(err)
		}
		plan := getPlanByIDQuery.Result

		countUsers := &query.CountUsers{}
		err = bus.Dispatch(c, countUsers)
		if err != nil {
			return c.Failure(err)
		}

		if plan.MaxUsers > 0 && countUsers.Result > plan.MaxUsers {
			return c.Unauthorized()
		}

		if err = bus.Dispatch(c, &cmd.CreateBillingSubscription{
			PlanID: plan.ID,
		}); err != nil {
			return c.Failure(err)
		}

		updateBilling := &cmd.UpdateTenantBillingSettings{Settings: c.Tenant().Billing}
		activateTenant := &cmd.ActivateTenant{TenantID: c.Tenant().ID}
		if err := bus.Dispatch(c, updateBilling, activateTenant); err != nil {
			return c.Failure(err)
		}

		return c.Ok(web.Map{})
	}
}

// CancelBillingSubscription cancels current subscription from current tenant
func CancelBillingSubscription() web.HandlerFunc {
	return func(c *web.Context) error {
		err := bus.Dispatch(c, &cmd.CancelBillingSubscription{})
		if err != nil {
			return c.Failure(err)
		}

		if err := bus.Dispatch(c, &cmd.UpdateTenantBillingSettings{
			Settings: c.Tenant().Billing,
		}); err != nil {
			return c.Failure(err)
		}

		return c.Ok(web.Map{})
	}
}
<|endoftext|>


<commit_before>// Copyright 2018 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
//
//     https://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.

// AUTO-GENERATED CODE. DO NOT EDIT.

// Package kms is an auto-generated package for the
// Google Cloud Key Management Service (KMS) API.
//
//   NOTE: This package is in alpha. It is not stable, and is likely to change.
//
// Manages encryption for your cloud services the same way you do
// on-premises.
// You can generate, use, rotate, and destroy AES256 encryption keys.
//
// Use the client at cloud.google.com/go/kms in preference to this.
package kms // import "cloud.google.com/go/kms/apiv1"

import (
	"golang.org/x/net/context"
	"google.golang.org/grpc/metadata"
)

func insertMetadata(ctx context.Context, mds ...metadata.MD) context.Context {
	out, _ := metadata.FromOutgoingContext(ctx)
	out = out.Copy()
	for _, md := range mds {
		for k, v := range md {
			out[k] = append(out[k], v...)
		}
	}
	return metadata.NewOutgoingContext(ctx, out)
}

// DefaultAuthScopes reports the default set of authentication scopes to use with this package.
func DefaultAuthScopes() []string {
	return []string{
		"https://www.googleapis.com/auth/cloud-platform",
	}
}
<commit_msg>kms: regen gapic; removes alpha notice<commit_after>// Copyright 2018 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
//
//     https://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.

// AUTO-GENERATED CODE. DO NOT EDIT.

// Package kms is an auto-generated package for the
// Google Cloud Key Management Service (KMS) API.

//
// Manages encryption for your cloud services the same way you do
// on-premises.
// You can generate, use, rotate, and destroy AES256 encryption keys.
//
// Use the client at cloud.google.com/go/kms in preference to this.
package kms // import "cloud.google.com/go/kms/apiv1"

import (
	"golang.org/x/net/context"
	"google.golang.org/grpc/metadata"
)

func insertMetadata(ctx context.Context, mds ...metadata.MD) context.Context {
	out, _ := metadata.FromOutgoingContext(ctx)
	out = out.Copy()
	for _, md := range mds {
		for k, v := range md {
			out[k] = append(out[k], v...)
		}
	}
	return metadata.NewOutgoingContext(ctx, out)
}

// DefaultAuthScopes reports the default set of authentication scopes to use with this package.
func DefaultAuthScopes() []string {
	return []string{
		"https://www.googleapis.com/auth/cloud-platform",
	}
}
<|endoftext|>


<commit_before>package main

import "time"

func main() {
	const (
		cap = 4
		max = 6
	)
	ch := make(chan int, cap)
	go worker(ch)
	for i := 0; i < max; i++ {
		ch <- i
		println("add", i)
	}
	close(ch)
	println("end of sending")
	time.Sleep(3e9)
}

func worker(in chan int) {
	for {
		select {
		case ch, ok := <-in:
			if !ok {
				println("chan close")
				return
			}
			time.Sleep(1e6)
			println(ch, len(in), "/", cap(in))
		default:
			println("default")
		}
	}
}
<commit_msg>channels/limit: select on send<commit_after>package main

import "time"

func main() {
	const (
		cap = 4
		max = 8
	)
	ch := make(chan int, cap)
	go worker(ch)
	for i := 0; i < max; i++ {
		select {
		case ch <- i:
			println("send", i)
		default:
			println("block", i)
			ch <- i
		}
	}
	close(ch)
	println("end of sending")
	time.Sleep(3e9)
}

func worker(ch chan int) {
	for v := range ch {
		time.Sleep(1e6)
		println("receive", v)
	}
}
<|endoftext|>


<commit_before>package chanotify

import (
	"errors"
	"sync"
)

// Notifier can effectively notify you about receiving from particular channels.
// It operates with pairs <-chan struct{} <-> string which is notification
// channel and its identificator respectively.
// Notification channel is <-chan struc{}, each send to which is spawn
// notification from Notifier, close doesn't spawn anything and removes channel
// from Notifier.
type Notifier struct {
	c chan interface{}

	m      sync.Mutex // guards doneCh
	doneCh map[interface{}]chan struct{}
	closed bool
}

// New returns a new *Notifier.
func New() *Notifier {
	s := &Notifier{
		c:      make(chan interface{}),
		doneCh: make(map[interface{}]chan struct{}),
	}
	return s
}

// Chan returns channel on which client listen for notifications.
// IDs of notifications is sent to the returned channel.
func (n *Notifier) Chan() <-chan interface{} {
	return n.c
}

// Add adds new notification channel to Notifier.
func (n *Notifier) Add(id interface{}, ch <-chan struct{}) error {
	n.m.Lock()
	defer n.m.Unlock()

	if n.closed {
		return errors.New("notifier closed; cannot add the channel on the notifier")
	}
	if _, ok := n.doneCh[id]; ok {
		return errors.New("cannot register duplicate key")
	}

	done := make(chan struct{})
	n.doneCh[id] = done

	n.startWorker(ch, id, done)
	return nil
}

func (n *Notifier) killWorker(id interface{}, done chan struct{}) {
	n.m.Lock()
	delete(n.doneCh, id)
	n.m.Unlock()
}

func (n *Notifier) startWorker(ch <-chan struct{}, id interface{}, done chan struct{}) {
	go func() {
		for {
			select {
			case _, ok := <-ch:
				if !ok {
					// If the channel is closed, we don't need the goroutine
					// or the done channel mechanism running anymore.
					n.killWorker(id, done)
					return
				}
				n.c <- id
			case <-done:
				// We don't need this goroutine running anymore, return.
				n.killWorker(id, done)
				return
			}
		}
	}()
}

// Close closes the notifier and releases its underlying resources.
func (n *Notifier) Close() {
	n.m.Lock()
	defer n.m.Unlock()
	for _, done := range n.doneCh {
		close(done)
	}
	close(n.c)
	n.closed = true
}
<commit_msg>Clarify chanotify docs about the usage of Add and Close<commit_after>package chanotify

import (
	"errors"
	"sync"
)

// Notifier can effectively notify you about receiving from particular channels.
// It operates with pairs <-chan struct{} <-> string which is notification
// channel and its identificator respectively.
// Notification channel is <-chan struc{}, each send to which is spawn
// notification from Notifier, close doesn't spawn anything and removes channel
// from Notifier.
type Notifier struct {
	c chan interface{}

	m      sync.Mutex // guards doneCh
	doneCh map[interface{}]chan struct{}
	closed bool
}

// New returns a new notifier. A notifier must be closed by
// calling, (*Notifier).Close, once it is no longer in use.
func New() *Notifier {
	s := &Notifier{
		c:      make(chan interface{}),
		doneCh: make(map[interface{}]chan struct{}),
	}
	return s
}

// Chan returns channel on which client listen for notifications.
// IDs of notifications is sent to the returned channel.
func (n *Notifier) Chan() <-chan interface{} {
	return n.c
}

// Add adds new notification channel to the notifier.
// Multiple registrations of the same ID is not allowed.
func (n *Notifier) Add(id interface{}, ch <-chan struct{}) error {
	n.m.Lock()
	defer n.m.Unlock()

	if n.closed {
		return errors.New("notifier closed; cannot add the channel on the notifier")
	}
	if _, ok := n.doneCh[id]; ok {
		return errors.New("cannot register duplicate key")
	}

	done := make(chan struct{})
	n.doneCh[id] = done

	n.startWorker(ch, id, done)
	return nil
}

func (n *Notifier) killWorker(id interface{}, done chan struct{}) {
	n.m.Lock()
	delete(n.doneCh, id)
	n.m.Unlock()
}

func (n *Notifier) startWorker(ch <-chan struct{}, id interface{}, done chan struct{}) {
	go func() {
		for {
			select {
			case _, ok := <-ch:
				if !ok {
					// If the channel is closed, we don't need the goroutine
					// or the done channel mechanism running anymore.
					n.killWorker(id, done)
					return
				}
				n.c <- id
			case <-done:
				// We don't need this goroutine running anymore, return.
				n.killWorker(id, done)
				return
			}
		}
	}()
}

// Close closes the notifier and releases its underlying resources.
func (n *Notifier) Close() {
	n.m.Lock()
	defer n.m.Unlock()
	for _, done := range n.doneCh {
		close(done)
	}
	close(n.c)
	n.closed = true
}
<|endoftext|>


<commit_before>package integration_test

import (
	"os/exec"
	"testing"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
)

func TestIntegration(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "CLI Integration")
}

var _ = Describe("exit codes", func() {
	It("exits non-zero when an unknown command is invoked", func() {
		result := Cf("some-command-that-should-never-actually-be-a-real-thing-i-can-use")

		Eventually(result).Should(Say("not a registered command"))
		Expect(result).Should(Exit(1))
	})
})

func Cf(args ...string) *Session {
	path, err := Build("github.com/cloudfoundry/cli/main")
	Expect(err).NotTo(HaveOccurred())

	session, err := Start(exec.Command(path, args...), GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())

	return session
}
<commit_msg>Allow integration test to take longer<commit_after>package integration_test

import (
	"os/exec"
	"testing"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
)

func TestIntegration(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "CLI Integration")
}

var _ = Describe("exit codes", func() {
	It("exits non-zero when an unknown command is invoked", func() {
		result := Cf("some-command-that-should-never-actually-be-a-real-thing-i-can-use")

		Eventually(result, 10).Should(Say("not a registered command"))
		Expect(result).Should(Exit(1))
	})
})

func Cf(args ...string) *Session {
	path, err := Build("github.com/cloudfoundry/cli/main")
	Expect(err).NotTo(HaveOccurred())

	session, err := Start(exec.Command(path, args...), GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())

	return session
}
<|endoftext|>


<commit_before>package interfaces

import (

)

<commit_msg>remove duplicate file for web interface<commit_after><|endoftext|>


<commit_before>package ydls

import (
	"fmt"
	"html/template"
	"io"
	"net/http"
	"net/url"
	"strings"
)

type baseURLXHeaders int

const (
	trustXHeaders baseURLXHeaders = iota
	dontTrustXHeaders
)

func baseURLFromRequest(r *http.Request, shouldXHeaders baseURLXHeaders) *url.URL {
	schema := ""
	host := ""
	prefix := ""
	if shouldXHeaders == trustXHeaders {
		schema = r.Header.Get("X-Forwarded-Proto")
		host = r.Header.Get("X-Forwarded-Host")
		prefix = r.Header.Get("X-Forwarded-Prefix")
	}

	if schema == "" {
		schema = "http"
		if r.TLS != nil {
			schema = "https"
		}
	}
	if host == "" {
		host = r.Host
	}

	return &url.URL{
		Scheme: schema,
		Host:   host,
		Path:   prefix,
	}
}

// URL encode with space encoded as "%20"
func urlEncode(s string) string {
	return strings.Replace(url.QueryEscape(s), "+", "%20", -1)
}

// make string safe to use in non-encoded content disposition filename
func safeContentDispositionFilename(s string) string {
	rs := []rune(s)
	for i, r := range rs {
		if r < 0x20 || r > 0x7e || r == '"' || r == '/' || r == '\\' {
			rs[i] = '_'
		}
	}

	return string(rs)
}

// Handler is a http.Handler using ydls
type Handler struct {
	YDLS      YDLS
	IndexTmpl *template.Template
	InfoLog   Printer
	DebugLog  Printer
}

func (yh *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	infoLog := yh.InfoLog
	if infoLog == nil {
		infoLog = nopPrinter{}
	}
	debugLog := yh.DebugLog
	if debugLog == nil {
		debugLog = nopPrinter{}
	}

	w.Header().Set("X-Content-Type-Options", "nosniff")
	w.Header().Set("X-XSS-Protection", "1; mode=block")

	debugLog.Printf("%s Request %s %s", r.RemoteAddr, r.Method, r.URL.String())

	if r.Method != http.MethodGet {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	if r.URL.Path == "/" && r.URL.RawQuery == "" {
		if yh.IndexTmpl != nil {
			w.Header().Set("Content-Security-Policy", "default-src 'none'; img-src 'self'; style-src 'unsafe-inline'; form-action 'self'")
			yh.IndexTmpl.Execute(w, yh.YDLS.Config.Formats)
		} else {
			http.Error(w, "Not found", http.StatusNotFound)
		}
		return
	} else if r.URL.Path == "/favicon.ico" {
		http.Error(w, "Not found", http.StatusNotFound)
		return
	}

	var requestOptions RequestOptions
	var requestOptionsErr error
	if r.URL.Query().Get("url") != "" {
		// ?url=url&format=format&codec=&codec=...
		requestOptions, requestOptionsErr = NewRequestOptionsFromQuery(r.URL.Query(), yh.YDLS.Config.Formats)
	} else {
		// /opt+opt.../http://...
		requestOptions, requestOptionsErr = NewRequestOptionsFromPath(r.URL, yh.YDLS.Config.Formats)
	}
	if requestOptionsErr != nil {
		infoLog.Printf("%s Invalid request %s %s (%s)", r.RemoteAddr, r.Method, r.URL.Path, requestOptionsErr.Error())
		http.Error(w, requestOptionsErr.Error(), http.StatusBadRequest)
		return
	}

	downloadOptions := DownloadOptions{
		RequestOptions: requestOptions,
		BaseURL:        baseURLFromRequest(r, trustXHeaders),
		DebugLog:       debugLog,
	}

	formatName := "best"
	if requestOptions.Format != nil {
		formatName = requestOptions.Format.Name
	}
	infoLog.Printf("%s Downloading (%s) %s", r.RemoteAddr, formatName, requestOptions.MediaRawURL)

	dr, err := yh.YDLS.Download(
		r.Context(),
		downloadOptions,
	)
	if err != nil {
		infoLog.Printf("%s Download failed %s %s (%s)", r.RemoteAddr, r.Method, r.URL.Path, err.Error())
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	w.Header().Set("Content-Security-Policy", "default-src 'none'; reflected-xss block")
	w.Header().Set("Content-Type", dr.MIMEType)
	if dr.Filename != "" {
		w.Header().Set("Content-Disposition",
			fmt.Sprintf("attachment; filename*=UTF-8''%s; filename=\"%s\"",
				urlEncode(dr.Filename), safeContentDispositionFilename(dr.Filename)),
		)
	}

	io.Copy(w, dr.Media)
	dr.Media.Close()
	dr.Wait()
}
<commit_msg>Relax Content-Security-Policy for index page<commit_after>package ydls

import (
	"fmt"
	"html/template"
	"io"
	"net/http"
	"net/url"
	"strings"
)

type baseURLXHeaders int

const (
	trustXHeaders baseURLXHeaders = iota
	dontTrustXHeaders
)

func baseURLFromRequest(r *http.Request, shouldXHeaders baseURLXHeaders) *url.URL {
	schema := ""
	host := ""
	prefix := ""
	if shouldXHeaders == trustXHeaders {
		schema = r.Header.Get("X-Forwarded-Proto")
		host = r.Header.Get("X-Forwarded-Host")
		prefix = r.Header.Get("X-Forwarded-Prefix")
	}

	if schema == "" {
		schema = "http"
		if r.TLS != nil {
			schema = "https"
		}
	}
	if host == "" {
		host = r.Host
	}

	return &url.URL{
		Scheme: schema,
		Host:   host,
		Path:   prefix,
	}
}

// URL encode with space encoded as "%20"
func urlEncode(s string) string {
	return strings.Replace(url.QueryEscape(s), "+", "%20", -1)
}

// make string safe to use in non-encoded content disposition filename
func safeContentDispositionFilename(s string) string {
	rs := []rune(s)
	for i, r := range rs {
		if r < 0x20 || r > 0x7e || r == '"' || r == '/' || r == '\\' {
			rs[i] = '_'
		}
	}

	return string(rs)
}

// Handler is a http.Handler using ydls
type Handler struct {
	YDLS      YDLS
	IndexTmpl *template.Template
	InfoLog   Printer
	DebugLog  Printer
}

func (yh *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	infoLog := yh.InfoLog
	if infoLog == nil {
		infoLog = nopPrinter{}
	}
	debugLog := yh.DebugLog
	if debugLog == nil {
		debugLog = nopPrinter{}
	}

	w.Header().Set("X-Content-Type-Options", "nosniff")
	w.Header().Set("X-XSS-Protection", "1; mode=block")

	debugLog.Printf("%s Request %s %s", r.RemoteAddr, r.Method, r.URL.String())

	if r.Method != http.MethodGet {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	if r.URL.Path == "/" && r.URL.RawQuery == "" {
		if yh.IndexTmpl != nil {
			w.Header().Set("Content-Security-Policy", "default-src 'self'; style-src 'unsafe-inline'; form-action 'self'")
			yh.IndexTmpl.Execute(w, yh.YDLS.Config.Formats)
		} else {
			http.Error(w, "Not found", http.StatusNotFound)
		}
		return
	} else if r.URL.Path == "/favicon.ico" {
		http.Error(w, "Not found", http.StatusNotFound)
		return
	}

	var requestOptions RequestOptions
	var requestOptionsErr error
	if r.URL.Query().Get("url") != "" {
		// ?url=url&format=format&codec=&codec=...
		requestOptions, requestOptionsErr = NewRequestOptionsFromQuery(r.URL.Query(), yh.YDLS.Config.Formats)
	} else {
		// /opt+opt.../http://...
		requestOptions, requestOptionsErr = NewRequestOptionsFromPath(r.URL, yh.YDLS.Config.Formats)
	}
	if requestOptionsErr != nil {
		infoLog.Printf("%s Invalid request %s %s (%s)", r.RemoteAddr, r.Method, r.URL.Path, requestOptionsErr.Error())
		http.Error(w, requestOptionsErr.Error(), http.StatusBadRequest)
		return
	}

	downloadOptions := DownloadOptions{
		RequestOptions: requestOptions,
		BaseURL:        baseURLFromRequest(r, trustXHeaders),
		DebugLog:       debugLog,
	}

	formatName := "best"
	if requestOptions.Format != nil {
		formatName = requestOptions.Format.Name
	}
	infoLog.Printf("%s Downloading (%s) %s", r.RemoteAddr, formatName, requestOptions.MediaRawURL)

	dr, err := yh.YDLS.Download(
		r.Context(),
		downloadOptions,
	)
	if err != nil {
		infoLog.Printf("%s Download failed %s %s (%s)", r.RemoteAddr, r.Method, r.URL.Path, err.Error())
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	w.Header().Set("Content-Security-Policy", "default-src 'none'; reflected-xss block")
	w.Header().Set("Content-Type", dr.MIMEType)
	if dr.Filename != "" {
		w.Header().Set("Content-Disposition",
			fmt.Sprintf("attachment; filename*=UTF-8''%s; filename=\"%s\"",
				urlEncode(dr.Filename), safeContentDispositionFilename(dr.Filename)),
		)
	}

	io.Copy(w, dr.Media)
	dr.Media.Close()
	dr.Wait()
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"crypto/md5"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"

	"github.com/jessevdk/go-flags"
	"github.com/mackerelio/checkers"
	"github.com/mattn/go-encoding"
	"github.com/mattn/go-zglob"
	enc "golang.org/x/text/encoding"
)

type logOpts struct {
	LogFile             string  `short:"f" long:"file" value-name:"FILE" description:"Path to log file"`
	Pattern             string  `short:"p" long:"pattern" required:"true" value-name:"PAT" description:"Pattern to search for"`
	Exclude             string  `short:"E" long:"exclude" value-name:"PAT" description:"Pattern to exclude from matching"`
	WarnOver            int64   `short:"w" long:"warning-over" description:"Trigger a warning if matched lines is over a number"`
	CritOver            int64   `short:"c" long:"critical-over" description:"Trigger a critical if matched lines is over a number"`
	WarnLevel           float64 `long:"warning-level" value-name:"N" description:"Warning level if pattern has a group"`
	CritLevel           float64 `long:"critical-level" value-name:"N" description:"Critical level if pattern has a group"`
	ReturnContent       bool    `short:"r" long:"return" description:"Return matched line"`
	FilePattern         string  `short:"F" long:"file-pattern" value-name:"FILE" description:"Check a pattern of files, instead of one file"`
	CaseInsensitive     bool    `short:"i" long:"icase" description:"Run a case insensitive match"`
	StateDir            string  `short:"s" long:"state-dir" value-name:"DIR" description:"Dir to keep state files under"`
	NoState             bool    `long:"no-state" description:"Don't use state file and read whole logs"`
	Encoding            string  `long:"encoding" description:"Encoding of log file"`
	Missing             string  `long:"missing" default:"UNKNOWN" value-name:"(CRITICAL|WARNING|OK|UNKNOWN)" description:"Exit status when log files missing"`
	patternReg          *regexp.Regexp
	excludeReg          *regexp.Regexp
	fileListFromGlob    []string
	fileListFromPattern []string
	origArgs            []string
	decoder             *enc.Decoder
}

func (opts *logOpts) prepare() error {
	if opts.LogFile == "" && opts.FilePattern == "" {
		return fmt.Errorf("No log file specified")
	}

	var err error
	if opts.patternReg, err = regCompileWithCase(opts.Pattern, opts.CaseInsensitive); err != nil {
		return fmt.Errorf("pattern is invalid")
	}

	if opts.Exclude != "" {
		opts.excludeReg, err = regCompileWithCase(opts.Exclude, opts.CaseInsensitive)
		if err != nil {
			return fmt.Errorf("exclude pattern is invalid")
		}
	}

	if opts.LogFile != "" {
		opts.fileListFromGlob, err = zglob.Glob(opts.LogFile)
		// unless --missing specified, we should ignore file not found error
		if err != nil && err != os.ErrNotExist {
			return fmt.Errorf("invalid glob for --file")
		}
	}

	if opts.FilePattern != "" {
		dirStr := filepath.Dir(opts.FilePattern)
		filePat := filepath.Base(opts.FilePattern)
		reg, err := regCompileWithCase(filePat, opts.CaseInsensitive)
		if err != nil {
			return fmt.Errorf("file-pattern is invalid")
		}

		fileInfos, err := ioutil.ReadDir(dirStr)
		if err != nil {
			return fmt.Errorf("cannot read the directory:" + err.Error())
		}

		for _, fileInfo := range fileInfos {
			if fileInfo.IsDir() {
				continue
			}
			fname := fileInfo.Name()
			if reg.MatchString(fname) {
				opts.fileListFromPattern = append(opts.fileListFromPattern, dirStr+string(filepath.Separator)+fileInfo.Name())
			}
		}
	}
	if !validateMissing(opts.Missing) {
		return fmt.Errorf("missing option is invalid")
	}
	return nil
}

func main() {
	ckr := run(os.Args[1:])
	ckr.Name = "LOG"
	ckr.Exit()
}

func regCompileWithCase(ptn string, caseInsensitive bool) (*regexp.Regexp, error) {
	if caseInsensitive {
		ptn = "(?i)" + ptn
	}
	return regexp.Compile(ptn)
}

func validateMissing(missing string) bool {
	switch missing {
	case "CRITICAL", "WARNING", "OK", "UNKNOWN", "":
		return true
	default:
		return false
	}
}

func parseArgs(args []string) (*logOpts, error) {
	origArgs := make([]string, len(args))
	copy(origArgs, args)
	opts := &logOpts{}
	_, err := flags.ParseArgs(opts, args)
	opts.origArgs = origArgs
	if opts.StateDir == "" {
		workdir := os.Getenv("MACKEREL_PLUGIN_WORKDIR")
		if workdir == "" {
			workdir = os.TempDir()
		}
		opts.StateDir = filepath.Join(workdir, "check-log")
	}
	return opts, err
}

func run(args []string) *checkers.Checker {
	opts, err := parseArgs(args)
	if err != nil {
		os.Exit(1)
	}

	err = opts.prepare()
	if err != nil {
		return checkers.Unknown(err.Error())
	}

	warnNum := int64(0)
	critNum := int64(0)
	var missingFiles []string
	errorOverall := ""

	if opts.LogFile != "" && len(opts.fileListFromGlob) == 0 {
		missingFiles = append(missingFiles, opts.LogFile)
	}

	for _, f := range append(opts.fileListFromGlob, opts.fileListFromPattern...) {
		_, err := os.Stat(f)
		if err != nil {
			missingFiles = append(missingFiles, f)
			continue
		}
		w, c, errLines, err := opts.searchLog(f)
		if err != nil {
			return checkers.Unknown(err.Error())
		}
		warnNum += w
		critNum += c
		if opts.ReturnContent {
			errorOverall += errLines
		}
	}

	msg := fmt.Sprintf("%d warnings, %d criticals for pattern /%s/.", warnNum, critNum, opts.Pattern)
	if errorOverall != "" {
		msg += "\n" + errorOverall
	}
	checkSt := checkers.OK
	if len(missingFiles) > 0 {
		switch opts.Missing {
		case "OK":
		case "WARNING":
			checkSt = checkers.WARNING
		case "CRITICAL":
			checkSt = checkers.CRITICAL
		default:
			checkSt = checkers.UNKNOWN
		}
		msg += "\n" + fmt.Sprintf("The following %d files are missing.", len(missingFiles))
		for _, f := range missingFiles {
			msg += "\n" + f
		}
	}
	if warnNum > opts.WarnOver {
		checkSt = checkers.WARNING
	}
	if critNum > opts.CritOver {
		checkSt = checkers.CRITICAL
	}
	return checkers.NewChecker(checkSt, msg)
}

func (opts *logOpts) searchLog(logFile string) (int64, int64, string, error) {
	stateFile := getStateFile(opts.StateDir, logFile, opts.origArgs)
	skipBytes := int64(0)
	if !opts.NoState {
		s, err := getBytesToSkip(stateFile)
		if err != nil {
			return 0, 0, "", err
		}
		skipBytes = s
	}

	f, err := os.Open(logFile)
	if err != nil {
		return 0, 0, "", err
	}
	defer f.Close()

	stat, err := f.Stat()
	if err != nil {
		return 0, 0, "", err
	}

	rotated := false
	if stat.Size() < skipBytes {
		rotated = true
	} else if skipBytes > 0 {
		f.Seek(skipBytes, 0)
	}

	var r io.Reader = f
	if opts.Encoding != "" {
		e := encoding.GetEncoding(opts.Encoding)
		if e == nil {
			return 0, 0, "", fmt.Errorf("unknown encoding:" + opts.Encoding)
		}
		opts.decoder = e.NewDecoder()
	}

	warnNum, critNum, readBytes, errLines, err := opts.searchReader(r)
	if err != nil {
		return warnNum, critNum, errLines, err
	}

	if rotated {
		skipBytes = readBytes
	} else {
		skipBytes += readBytes
	}

	if !opts.NoState {
		err = writeBytesToSkip(stateFile, skipBytes)
		if err != nil {
			log.Printf("writeByteToSkip failed: %s\n", err.Error())
		}
	}
	return warnNum, critNum, errLines, nil
}

func (opts *logOpts) searchReader(rdr io.Reader) (warnNum, critNum, readBytes int64, errLines string, err error) {
	r := bufio.NewReader(rdr)
	for {
		lineBytes, rErr := r.ReadBytes('\n')
		if rErr != nil {
			if rErr != io.EOF {
				err = rErr
			}
			break
		}
		readBytes += int64(len(lineBytes))

		if opts.decoder != nil {
			lineBytes, err = opts.decoder.Bytes(lineBytes)
			if err != nil {
				break
			}
		}
		line := strings.Trim(string(lineBytes), "\r\n")
		if matched, matches := opts.match(line); matched {
			if len(matches) > 1 && (opts.WarnLevel > 0 || opts.CritLevel > 0) {
				level, err := strconv.ParseFloat(matches[1], 64)
				if err != nil {
					warnNum++
					critNum++
					errLines += line + "\n"
				} else {
					levelOver := false
					if level > opts.WarnLevel {
						levelOver = true
						warnNum++
					}
					if level > opts.CritLevel {
						levelOver = true
						critNum++
					}
					if levelOver {
						errLines += line + "\n"
					}
				}
			} else {
				warnNum++
				critNum++
				errLines += line + "\n"
			}
		}
	}
	return
}

func (opts *logOpts) match(line string) (bool, []string) {
	pReg := opts.patternReg
	eReg := opts.excludeReg

	matches := pReg.FindStringSubmatch(line)
	matched := len(matches) > 0 && (eReg == nil || !eReg.MatchString(line))
	return matched, matches
}

var stateRe = regexp.MustCompile(`^([A-Z]):[/\\]`)

func getStateFile(stateDir, f string, args []string) string {
	return filepath.Join(
		stateDir,
		fmt.Sprintf(
			"%s-%x",
			stateRe.ReplaceAllString(f, `$1`+string(filepath.Separator)),
			md5.Sum([]byte(strings.Join(args, " "))),
		),
	)
}

func getBytesToSkip(f string) (int64, error) {
	_, err := os.Stat(f)
	if err != nil {
		return 0, nil
	}
	b, err := ioutil.ReadFile(f)
	if err != nil {
		return 0, err
	}
	i, err := strconv.ParseInt(strings.Trim(string(b), " \r\n"), 10, 64)
	if err != nil {
		log.Printf("failed to getBytesToSkip (ignoring): %s", err)
	}
	return i, nil
}

func writeBytesToSkip(f string, num int64) error {
	err := os.MkdirAll(filepath.Dir(f), 0755)
	if err != nil {
		return err
	}
	return writeFileAtomically(f, []byte(fmt.Sprintf("%d", num)))
}

func writeFileAtomically(f string, contents []byte) error {
	tmpf, err := ioutil.TempFile("", "")
	if err != nil {
		return err
	}
	defer os.Remove(tmpf.Name())
	_, err = tmpf.Write(contents)
	tmpf.Close() // should be called before rename
	if err != nil {
		return err
	}
	return os.Rename(tmpf.Name(), f)
}
<commit_msg>add comment for defer os.Remove because it fails on successful case<commit_after>package main

import (
	"bufio"
	"crypto/md5"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"

	"github.com/jessevdk/go-flags"
	"github.com/mackerelio/checkers"
	"github.com/mattn/go-encoding"
	"github.com/mattn/go-zglob"
	enc "golang.org/x/text/encoding"
)

type logOpts struct {
	LogFile             string  `short:"f" long:"file" value-name:"FILE" description:"Path to log file"`
	Pattern             string  `short:"p" long:"pattern" required:"true" value-name:"PAT" description:"Pattern to search for"`
	Exclude             string  `short:"E" long:"exclude" value-name:"PAT" description:"Pattern to exclude from matching"`
	WarnOver            int64   `short:"w" long:"warning-over" description:"Trigger a warning if matched lines is over a number"`
	CritOver            int64   `short:"c" long:"critical-over" description:"Trigger a critical if matched lines is over a number"`
	WarnLevel           float64 `long:"warning-level" value-name:"N" description:"Warning level if pattern has a group"`
	CritLevel           float64 `long:"critical-level" value-name:"N" description:"Critical level if pattern has a group"`
	ReturnContent       bool    `short:"r" long:"return" description:"Return matched line"`
	FilePattern         string  `short:"F" long:"file-pattern" value-name:"FILE" description:"Check a pattern of files, instead of one file"`
	CaseInsensitive     bool    `short:"i" long:"icase" description:"Run a case insensitive match"`
	StateDir            string  `short:"s" long:"state-dir" value-name:"DIR" description:"Dir to keep state files under"`
	NoState             bool    `long:"no-state" description:"Don't use state file and read whole logs"`
	Encoding            string  `long:"encoding" description:"Encoding of log file"`
	Missing             string  `long:"missing" default:"UNKNOWN" value-name:"(CRITICAL|WARNING|OK|UNKNOWN)" description:"Exit status when log files missing"`
	patternReg          *regexp.Regexp
	excludeReg          *regexp.Regexp
	fileListFromGlob    []string
	fileListFromPattern []string
	origArgs            []string
	decoder             *enc.Decoder
}

func (opts *logOpts) prepare() error {
	if opts.LogFile == "" && opts.FilePattern == "" {
		return fmt.Errorf("No log file specified")
	}

	var err error
	if opts.patternReg, err = regCompileWithCase(opts.Pattern, opts.CaseInsensitive); err != nil {
		return fmt.Errorf("pattern is invalid")
	}

	if opts.Exclude != "" {
		opts.excludeReg, err = regCompileWithCase(opts.Exclude, opts.CaseInsensitive)
		if err != nil {
			return fmt.Errorf("exclude pattern is invalid")
		}
	}

	if opts.LogFile != "" {
		opts.fileListFromGlob, err = zglob.Glob(opts.LogFile)
		// unless --missing specified, we should ignore file not found error
		if err != nil && err != os.ErrNotExist {
			return fmt.Errorf("invalid glob for --file")
		}
	}

	if opts.FilePattern != "" {
		dirStr := filepath.Dir(opts.FilePattern)
		filePat := filepath.Base(opts.FilePattern)
		reg, err := regCompileWithCase(filePat, opts.CaseInsensitive)
		if err != nil {
			return fmt.Errorf("file-pattern is invalid")
		}

		fileInfos, err := ioutil.ReadDir(dirStr)
		if err != nil {
			return fmt.Errorf("cannot read the directory:" + err.Error())
		}

		for _, fileInfo := range fileInfos {
			if fileInfo.IsDir() {
				continue
			}
			fname := fileInfo.Name()
			if reg.MatchString(fname) {
				opts.fileListFromPattern = append(opts.fileListFromPattern, dirStr+string(filepath.Separator)+fileInfo.Name())
			}
		}
	}
	if !validateMissing(opts.Missing) {
		return fmt.Errorf("missing option is invalid")
	}
	return nil
}

func main() {
	ckr := run(os.Args[1:])
	ckr.Name = "LOG"
	ckr.Exit()
}

func regCompileWithCase(ptn string, caseInsensitive bool) (*regexp.Regexp, error) {
	if caseInsensitive {
		ptn = "(?i)" + ptn
	}
	return regexp.Compile(ptn)
}

func validateMissing(missing string) bool {
	switch missing {
	case "CRITICAL", "WARNING", "OK", "UNKNOWN", "":
		return true
	default:
		return false
	}
}

func parseArgs(args []string) (*logOpts, error) {
	origArgs := make([]string, len(args))
	copy(origArgs, args)
	opts := &logOpts{}
	_, err := flags.ParseArgs(opts, args)
	opts.origArgs = origArgs
	if opts.StateDir == "" {
		workdir := os.Getenv("MACKEREL_PLUGIN_WORKDIR")
		if workdir == "" {
			workdir = os.TempDir()
		}
		opts.StateDir = filepath.Join(workdir, "check-log")
	}
	return opts, err
}

func run(args []string) *checkers.Checker {
	opts, err := parseArgs(args)
	if err != nil {
		os.Exit(1)
	}

	err = opts.prepare()
	if err != nil {
		return checkers.Unknown(err.Error())
	}

	warnNum := int64(0)
	critNum := int64(0)
	var missingFiles []string
	errorOverall := ""

	if opts.LogFile != "" && len(opts.fileListFromGlob) == 0 {
		missingFiles = append(missingFiles, opts.LogFile)
	}

	for _, f := range append(opts.fileListFromGlob, opts.fileListFromPattern...) {
		_, err := os.Stat(f)
		if err != nil {
			missingFiles = append(missingFiles, f)
			continue
		}
		w, c, errLines, err := opts.searchLog(f)
		if err != nil {
			return checkers.Unknown(err.Error())
		}
		warnNum += w
		critNum += c
		if opts.ReturnContent {
			errorOverall += errLines
		}
	}

	msg := fmt.Sprintf("%d warnings, %d criticals for pattern /%s/.", warnNum, critNum, opts.Pattern)
	if errorOverall != "" {
		msg += "\n" + errorOverall
	}
	checkSt := checkers.OK
	if len(missingFiles) > 0 {
		switch opts.Missing {
		case "OK":
		case "WARNING":
			checkSt = checkers.WARNING
		case "CRITICAL":
			checkSt = checkers.CRITICAL
		default:
			checkSt = checkers.UNKNOWN
		}
		msg += "\n" + fmt.Sprintf("The following %d files are missing.", len(missingFiles))
		for _, f := range missingFiles {
			msg += "\n" + f
		}
	}
	if warnNum > opts.WarnOver {
		checkSt = checkers.WARNING
	}
	if critNum > opts.CritOver {
		checkSt = checkers.CRITICAL
	}
	return checkers.NewChecker(checkSt, msg)
}

func (opts *logOpts) searchLog(logFile string) (int64, int64, string, error) {
	stateFile := getStateFile(opts.StateDir, logFile, opts.origArgs)
	skipBytes := int64(0)
	if !opts.NoState {
		s, err := getBytesToSkip(stateFile)
		if err != nil {
			return 0, 0, "", err
		}
		skipBytes = s
	}

	f, err := os.Open(logFile)
	if err != nil {
		return 0, 0, "", err
	}
	defer f.Close()

	stat, err := f.Stat()
	if err != nil {
		return 0, 0, "", err
	}

	rotated := false
	if stat.Size() < skipBytes {
		rotated = true
	} else if skipBytes > 0 {
		f.Seek(skipBytes, 0)
	}

	var r io.Reader = f
	if opts.Encoding != "" {
		e := encoding.GetEncoding(opts.Encoding)
		if e == nil {
			return 0, 0, "", fmt.Errorf("unknown encoding:" + opts.Encoding)
		}
		opts.decoder = e.NewDecoder()
	}

	warnNum, critNum, readBytes, errLines, err := opts.searchReader(r)
	if err != nil {
		return warnNum, critNum, errLines, err
	}

	if rotated {
		skipBytes = readBytes
	} else {
		skipBytes += readBytes
	}

	if !opts.NoState {
		err = writeBytesToSkip(stateFile, skipBytes)
		if err != nil {
			log.Printf("writeByteToSkip failed: %s\n", err.Error())
		}
	}
	return warnNum, critNum, errLines, nil
}

func (opts *logOpts) searchReader(rdr io.Reader) (warnNum, critNum, readBytes int64, errLines string, err error) {
	r := bufio.NewReader(rdr)
	for {
		lineBytes, rErr := r.ReadBytes('\n')
		if rErr != nil {
			if rErr != io.EOF {
				err = rErr
			}
			break
		}
		readBytes += int64(len(lineBytes))

		if opts.decoder != nil {
			lineBytes, err = opts.decoder.Bytes(lineBytes)
			if err != nil {
				break
			}
		}
		line := strings.Trim(string(lineBytes), "\r\n")
		if matched, matches := opts.match(line); matched {
			if len(matches) > 1 && (opts.WarnLevel > 0 || opts.CritLevel > 0) {
				level, err := strconv.ParseFloat(matches[1], 64)
				if err != nil {
					warnNum++
					critNum++
					errLines += line + "\n"
				} else {
					levelOver := false
					if level > opts.WarnLevel {
						levelOver = true
						warnNum++
					}
					if level > opts.CritLevel {
						levelOver = true
						critNum++
					}
					if levelOver {
						errLines += line + "\n"
					}
				}
			} else {
				warnNum++
				critNum++
				errLines += line + "\n"
			}
		}
	}
	return
}

func (opts *logOpts) match(line string) (bool, []string) {
	pReg := opts.patternReg
	eReg := opts.excludeReg

	matches := pReg.FindStringSubmatch(line)
	matched := len(matches) > 0 && (eReg == nil || !eReg.MatchString(line))
	return matched, matches
}

var stateRe = regexp.MustCompile(`^([A-Z]):[/\\]`)

func getStateFile(stateDir, f string, args []string) string {
	return filepath.Join(
		stateDir,
		fmt.Sprintf(
			"%s-%x",
			stateRe.ReplaceAllString(f, `$1`+string(filepath.Separator)),
			md5.Sum([]byte(strings.Join(args, " "))),
		),
	)
}

func getBytesToSkip(f string) (int64, error) {
	_, err := os.Stat(f)
	if err != nil {
		return 0, nil
	}
	b, err := ioutil.ReadFile(f)
	if err != nil {
		return 0, err
	}
	i, err := strconv.ParseInt(strings.Trim(string(b), " \r\n"), 10, 64)
	if err != nil {
		log.Printf("failed to getBytesToSkip (ignoring): %s", err)
	}
	return i, nil
}

func writeBytesToSkip(f string, num int64) error {
	err := os.MkdirAll(filepath.Dir(f), 0755)
	if err != nil {
		return err
	}
	return writeFileAtomically(f, []byte(fmt.Sprintf("%d", num)))
}

func writeFileAtomically(f string, contents []byte) error {
	tmpf, err := ioutil.TempFile("", "")
	if err != nil {
		return err
	}
	// os.Remove here works successfully when tmpf.Write fails or os.Rename fails.
	// In successful case, os.Remove fails because the temporary file is already renamed.
	defer os.Remove(tmpf.Name())
	_, err = tmpf.Write(contents)
	tmpf.Close() // should be called before rename
	if err != nil {
		return err
	}
	return os.Rename(tmpf.Name(), f)
}
<|endoftext|>