<commit_before>/**
 * (c) 2014, Caoimhe Chaos <caoimhechaos@protonmail.com>,
 *	     Ancient Solutions. All rights reserved.
 *
 * Redistribution and use in source  and binary forms, with or without
 * modification, are permitted  provided that the following conditions
 * are met:
 *
 * * Redistributions of  source code  must retain the  above copyright
 *   notice, this list of conditions and the following disclaimer.
 * * Redistributions in binary form must reproduce the above copyright
 *   notice, this  list of conditions and the  following disclaimer in
 *   the  documentation  and/or  other  materials  provided  with  the
 *   distribution.
 * * Neither  the  name  of  Ancient Solutions  nor  the  name  of its
 *   contributors may  be used to endorse or  promote products derived
 *   from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS"  AND ANY EXPRESS  OR IMPLIED WARRANTIES  OF MERCHANTABILITY
 * AND FITNESS  FOR A PARTICULAR  PURPOSE ARE DISCLAIMED. IN  NO EVENT
 * SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
 * INDIRECT, INCIDENTAL, SPECIAL,  EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED  TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE,  DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
 * STRICT  LIABILITY,  OR  TORT  (INCLUDING NEGLIGENCE  OR  OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
 * OF THE POSSIBILITY OF SUCH DAMAGE.
 */

// SMTP handler callback.
package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/mail"
	"reflect"
	"regexp"
	"time"

	"ancient-solutions.com/mailpump"
	"ancient-solutions.com/mailpump/smtpump"
)

type smtpCallback struct {
	smtpump.SmtpReceiver
	maxContentLength int64
}

var features = []string{"ETRN", "8BITMIME", "DSN"}

// String representation of an email regular expression.
var email_re string = "([\\w\\+-\\.]+(?:%[\\w\\+-\\.]+)?@[\\w\\+-\\.]+)"

// RE match to extract the mail address from a MAIL From command.
var from_re *regexp.Regexp = regexp.MustCompile(
	"^[Ff][Rr][Oo][Mm]:\\s*(?:<" + email_re + ">|" + email_re + ")$")

// RE match to extract the mail address from a RCPT To command.
var rcpt_re *regexp.Regexp = regexp.MustCompile(
	"^[Tt][Oo]:\\s*(?:<" + email_re + ">|" + email_re + ")$")

func getConnectionData(conn *smtpump.SmtpConnection) *mailpump.MailMessage {
	var msg *mailpump.MailMessage
	var val reflect.Value
	var ud interface{}
	var ok bool

	ud = conn.GetUserdata()
	val = reflect.ValueOf(ud)
	if !val.IsValid() || val.IsNil() {
		msg = new(mailpump.MailMessage)
		conn.SetUserdata(msg)
		return msg
	}

	msg, ok = ud.(*mailpump.MailMessage)
	if !ok {
		log.Print("Connection userdata is not a MailMessage!")
		return nil
	}

	if msg == nil {
		msg = new(mailpump.MailMessage)
		conn.SetUserdata(msg)
	}

	return msg
}

// Store all available information about the peer in the message structure
// for SPAM analysis.
func (self smtpCallback) ConnectionOpened(
	conn *smtpump.SmtpConnection, peer net.Addr) (
	ret smtpump.SmtpReturnCode) {
	var host string
	var msg *mailpump.MailMessage = getConnectionData(conn)
	var err error

	if msg == nil {
		ret.Code = smtpump.SMTP_LOCALERR
		ret.Message = "Unable to allocate connection structures."
		ret.Terminate = true
		return
	}

	host, _, err = net.SplitHostPort(peer.String())
	if err == nil {
		msg.SmtpPeer = &host
	} else {
		host = peer.String()
		msg.SmtpPeer = &host
	}
	msg.SmtpPeerRevdns, _ = net.LookupAddr(host)
	return
}

// Ignore disconnections.
func (self smtpCallback) ConnectionClosed(conn *smtpump.SmtpConnection) {
}

// Just save the host name and respond.
func (self smtpCallback) Helo(
	conn *smtpump.SmtpConnection, hostname string, esmtp bool) (
	ret smtpump.SmtpReturnCode) {
	var msg *mailpump.MailMessage = getConnectionData(conn)
	var response string = fmt.Sprintf("Hello, %s! Nice to meet you.",
		hostname)
	msg.SmtpHelo = &hostname

	if esmtp {
		var pos int
		var capa string
		conn.Respond(smtpump.SMTP_COMPLETED, true, response)

		for pos, capa = range features {
			conn.Respond(smtpump.SMTP_COMPLETED,
				pos < (len(features)-1), capa)
		}
		return
	}

	ret.Code = smtpump.SMTP_COMPLETED
	ret.Message = response
	return
}

// Ensure HELO has been set, then record From.
func (self smtpCallback) MailFrom(
	conn *smtpump.SmtpConnection, sender string) (
	ret smtpump.SmtpReturnCode) {
	var msg *mailpump.MailMessage = getConnectionData(conn)
	var matches []string
	var addr string

	if msg.SmtpHelo == nil {
		ret.Code = smtpump.SMTP_BAD_SEQUENCE
		ret.Message = "Polite people say Hello first!"
		return
	}

	matches = from_re.FindStringSubmatch(sender)
	if len(matches) == 0 {
		if len(sender) > 0 {
			log.Print("Received unparseable address: ", sender)
		}
		ret.Code = smtpump.SMTP_PARAMETER_NOT_IMPLEMENTED
		ret.Message = "Address not understood, sorry."
		return
	}

	for _, addr = range matches {
		if len(addr) > 0 {
			msg.SmtpFrom = new(string)
			*msg.SmtpFrom = addr
		}
	}
	ret.Code = smtpump.SMTP_COMPLETED
	ret.Message = "Ok."
	return
}

// Ensure HELO and MAIL have been set, then record To.
func (self smtpCallback) RcptTo(
	conn *smtpump.SmtpConnection, recipient string) (
	ret smtpump.SmtpReturnCode) {
	var msg *mailpump.MailMessage = getConnectionData(conn)
	var matches []string
	var addr string
	var realaddr string

	if msg.SmtpHelo == nil {
		ret.Code = smtpump.SMTP_BAD_SEQUENCE
		ret.Message = "Polite people say Hello first!"
		return
	}

	if msg.SmtpFrom == nil {
		ret.Code = smtpump.SMTP_BAD_SEQUENCE
		ret.Message = "Need MAIL command before RCPT."
		return
	}

	matches = rcpt_re.FindStringSubmatch(recipient)
	if len(matches) == 0 {
		if len(recipient) > 0 {
			log.Print("Received unparseable address: ", recipient)
		}
		ret.Code = smtpump.SMTP_PARAMETER_NOT_IMPLEMENTED
		ret.Message = "Address not understood, sorry."
		return
	}

	for _, addr = range matches {
		if len(addr) > 0 {
			realaddr = addr
		}
	}
	msg.SmtpTo = append(msg.SmtpTo, realaddr)
	ret.Code = smtpump.SMTP_COMPLETED
	ret.Message = "Ok."
	return
}

// Read the data following the DATA command, up to the configured limit.
func (self smtpCallback) Data(conn *smtpump.SmtpConnection) (
	ret smtpump.SmtpReturnCode) {
	var msg *mailpump.MailMessage = getConnectionData(conn)
	var hdr string
	var vals []string
	var addrs []*mail.Address
	var addr *mail.Address
	var dotreader io.Reader
	var contentsreader *io.LimitedReader
	var message *mail.Message
	var tm time.Time
	var err error

	if msg.SmtpHelo == nil {
		ret.Code = smtpump.SMTP_BAD_SEQUENCE
		ret.Message = "Polite people say Hello first!"
		return
	}

	if msg.SmtpFrom == nil {
		ret.Code = smtpump.SMTP_BAD_SEQUENCE
		ret.Message = "Need MAIL command before DATA."
		return
	}

	if len(msg.SmtpTo) == 0 {
		ret.Code = smtpump.SMTP_BAD_SEQUENCE
		ret.Message = "Need RCPT command before DATA."
		return
	}

	conn.Respond(smtpump.SMTP_PROCEED, false, "Proceed with message.")

	dotreader = conn.GetDotReader()
	contentsreader = &io.LimitedReader{
		R: dotreader,
		N: self.maxContentLength + 1,
	}
	message, err = mail.ReadMessage(contentsreader)
	if err != nil {
		ret.Code = smtpump.SMTP_LOCALERR
		ret.Message = "Unable to read message: " + err.Error()
		// Consume all remaining output before returning an error.
		ioutil.ReadAll(dotreader)
		return
	}

	// See if we ran out of bytes to our limit
	if contentsreader.N <= 0 {
		ret.Code = smtpump.SMTP_MESSAGE_TOO_BIG
		ret.Message = "Size limit exceeded. Thanks for playing."
		ret.Terminate = true
		return
	}

	msg.Body, err = ioutil.ReadAll(message.Body)
	if err != nil {
		ret.Code = smtpump.SMTP_LOCALERR
		ret.Message = "Unable to parse message: " + err.Error()
		return
	}

	for hdr, vals = range message.Header {
		var header = new(mailpump.MailMessage_MailHeader)
		header.Name = new(string)
		*header.Name = hdr
		header.Value = make([]string, len(vals))
		copy(header.Value, vals)
		msg.Headers = append(msg.Headers, header)
	}

	tm, err = message.Header.Date()
	if err == nil {
		msg.DateHdr = new(int64)
		*msg.DateHdr = tm.Unix()
	}

	addrs, _ = message.Header.AddressList("From")
	if len(addrs) > 0 {
		msg.FromHdr = new(string)
		*msg.FromHdr = addrs[0].String()
	}

	addrs, _ = message.Header.AddressList("To")
	for _, addr = range addrs {
		msg.ToHdr = append(msg.ToHdr, addr.String())
	}

	addrs, _ = message.Header.AddressList("Cc")
	for _, addr = range addrs {
		msg.CcHdrs = append(msg.CcHdrs, addr.String())
	}

	addrs, _ = message.Header.AddressList("Sender")
	if len(addrs) > 0 {
		msg.SenderHdr = new(string)
		*msg.SenderHdr = addrs[0].String()
	}

	hdr = message.Header.Get("Message-Id")
	if len(hdr) <= 0 {
		hdr = message.Header.Get("Message-ID")
	}
	if len(hdr) > 0 {
		msg.MsgidHdr = new(string)
		*msg.MsgidHdr = hdr
	}

	// TODO(caoimhe): send this to some server.
	ret.Code = smtpump.SMTP_NOT_IMPLEMENTED
	ret.Message = "Ok, but this doesn't go anywhere yet."
	return
}

// FIXME: STUB.
func (self smtpCallback) Etrn(conn *smtpump.SmtpConnection, domain string) (
	ret smtpump.SmtpReturnCode) {
	ret.Code = smtpump.SMTP_NOT_IMPLEMENTED
	ret.Message = "Not yet implemented."
	return
}

// Forget all connection related data except HELO and the peer information.
func (self smtpCallback) Reset(conn *smtpump.SmtpConnection) (
	ret smtpump.SmtpReturnCode) {
	var msg = getConnectionData(conn)
	var peer, tlsc, helo string
	var rdns []string
	peer = msg.GetSmtpPeer()
	tlsc = msg.GetSmtpPeerTlsInfo()
	rdns = msg.GetSmtpPeerRevdns()
	helo = msg.GetSmtpHelo()
	msg.Reset()

	msg.SmtpPeer = &peer
	if len(tlsc) > 0 {
		msg.SmtpPeerTlsInfo = &tlsc
	}
	if len(rdns) > 0 {
		msg.SmtpPeerRevdns = make([]string, len(rdns))
		copy(msg.SmtpPeerRevdns, rdns)
	}
	if len(helo) > 0 {
		msg.SmtpHelo = &helo
	}
	ret.Code = smtpump.SMTP_COMPLETED
	ret.Message = "Ok."
	return
}

// FIXME: STUB.
func (self smtpCallback) Quit(conn *smtpump.SmtpConnection) (
	ret smtpump.SmtpReturnCode) {
	ret.Code = smtpump.SMTP_CLOSING
	ret.Message = "See you later!"
	ret.Terminate = true
	return
}
<commit_msg>Add a description for QUIT, which is already implemented.<commit_after>/**
 * (c) 2014, Caoimhe Chaos <caoimhechaos@protonmail.com>,
 *	     Ancient Solutions. All rights reserved.
 *
 * Redistribution and use in source  and binary forms, with or without
 * modification, are permitted  provided that the following conditions
 * are met:
 *
 * * Redistributions of  source code  must retain the  above copyright
 *   notice, this list of conditions and the following disclaimer.
 * * Redistributions in binary form must reproduce the above copyright
 *   notice, this  list of conditions and the  following disclaimer in
 *   the  documentation  and/or  other  materials  provided  with  the
 *   distribution.
 * * Neither  the  name  of  Ancient Solutions  nor  the  name  of its
 *   contributors may  be used to endorse or  promote products derived
 *   from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS"  AND ANY EXPRESS  OR IMPLIED WARRANTIES  OF MERCHANTABILITY
 * AND FITNESS  FOR A PARTICULAR  PURPOSE ARE DISCLAIMED. IN  NO EVENT
 * SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
 * INDIRECT, INCIDENTAL, SPECIAL,  EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED  TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE,  DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
 * STRICT  LIABILITY,  OR  TORT  (INCLUDING NEGLIGENCE  OR  OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
 * OF THE POSSIBILITY OF SUCH DAMAGE.
 */

// SMTP handler callback.
package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/mail"
	"reflect"
	"regexp"
	"time"

	"ancient-solutions.com/mailpump"
	"ancient-solutions.com/mailpump/smtpump"
)

type smtpCallback struct {
	smtpump.SmtpReceiver
	maxContentLength int64
}

var features = []string{"ETRN", "8BITMIME", "DSN"}

// String representation of an email regular expression.
var email_re string = "([\\w\\+-\\.]+(?:%[\\w\\+-\\.]+)?@[\\w\\+-\\.]+)"

// RE match to extract the mail address from a MAIL From command.
var from_re *regexp.Regexp = regexp.MustCompile(
	"^[Ff][Rr][Oo][Mm]:\\s*(?:<" + email_re + ">|" + email_re + ")$")

// RE match to extract the mail address from a RCPT To command.
var rcpt_re *regexp.Regexp = regexp.MustCompile(
	"^[Tt][Oo]:\\s*(?:<" + email_re + ">|" + email_re + ")$")

func getConnectionData(conn *smtpump.SmtpConnection) *mailpump.MailMessage {
	var msg *mailpump.MailMessage
	var val reflect.Value
	var ud interface{}
	var ok bool

	ud = conn.GetUserdata()
	val = reflect.ValueOf(ud)
	if !val.IsValid() || val.IsNil() {
		msg = new(mailpump.MailMessage)
		conn.SetUserdata(msg)
		return msg
	}

	msg, ok = ud.(*mailpump.MailMessage)
	if !ok {
		log.Print("Connection userdata is not a MailMessage!")
		return nil
	}

	if msg == nil {
		msg = new(mailpump.MailMessage)
		conn.SetUserdata(msg)
	}

	return msg
}

// Store all available information about the peer in the message structure
// for SPAM analysis.
func (self smtpCallback) ConnectionOpened(
	conn *smtpump.SmtpConnection, peer net.Addr) (
	ret smtpump.SmtpReturnCode) {
	var host string
	var msg *mailpump.MailMessage = getConnectionData(conn)
	var err error

	if msg == nil {
		ret.Code = smtpump.SMTP_LOCALERR
		ret.Message = "Unable to allocate connection structures."
		ret.Terminate = true
		return
	}

	host, _, err = net.SplitHostPort(peer.String())
	if err == nil {
		msg.SmtpPeer = &host
	} else {
		host = peer.String()
		msg.SmtpPeer = &host
	}
	msg.SmtpPeerRevdns, _ = net.LookupAddr(host)
	return
}

// Ignore disconnections.
func (self smtpCallback) ConnectionClosed(conn *smtpump.SmtpConnection) {
}

// Just save the host name and respond.
func (self smtpCallback) Helo(
	conn *smtpump.SmtpConnection, hostname string, esmtp bool) (
	ret smtpump.SmtpReturnCode) {
	var msg *mailpump.MailMessage = getConnectionData(conn)
	var response string = fmt.Sprintf("Hello, %s! Nice to meet you.",
		hostname)
	msg.SmtpHelo = &hostname

	if esmtp {
		var pos int
		var capa string
		conn.Respond(smtpump.SMTP_COMPLETED, true, response)

		for pos, capa = range features {
			conn.Respond(smtpump.SMTP_COMPLETED,
				pos < (len(features)-1), capa)
		}
		return
	}

	ret.Code = smtpump.SMTP_COMPLETED
	ret.Message = response
	return
}

// Ensure HELO has been set, then record From.
func (self smtpCallback) MailFrom(
	conn *smtpump.SmtpConnection, sender string) (
	ret smtpump.SmtpReturnCode) {
	var msg *mailpump.MailMessage = getConnectionData(conn)
	var matches []string
	var addr string

	if msg.SmtpHelo == nil {
		ret.Code = smtpump.SMTP_BAD_SEQUENCE
		ret.Message = "Polite people say Hello first!"
		return
	}

	matches = from_re.FindStringSubmatch(sender)
	if len(matches) == 0 {
		if len(sender) > 0 {
			log.Print("Received unparseable address: ", sender)
		}
		ret.Code = smtpump.SMTP_PARAMETER_NOT_IMPLEMENTED
		ret.Message = "Address not understood, sorry."
		return
	}

	for _, addr = range matches {
		if len(addr) > 0 {
			msg.SmtpFrom = new(string)
			*msg.SmtpFrom = addr
		}
	}
	ret.Code = smtpump.SMTP_COMPLETED
	ret.Message = "Ok."
	return
}

// Ensure HELO and MAIL have been set, then record To.
func (self smtpCallback) RcptTo(
	conn *smtpump.SmtpConnection, recipient string) (
	ret smtpump.SmtpReturnCode) {
	var msg *mailpump.MailMessage = getConnectionData(conn)
	var matches []string
	var addr string
	var realaddr string

	if msg.SmtpHelo == nil {
		ret.Code = smtpump.SMTP_BAD_SEQUENCE
		ret.Message = "Polite people say Hello first!"
		return
	}

	if msg.SmtpFrom == nil {
		ret.Code = smtpump.SMTP_BAD_SEQUENCE
		ret.Message = "Need MAIL command before RCPT."
		return
	}

	matches = rcpt_re.FindStringSubmatch(recipient)
	if len(matches) == 0 {
		if len(recipient) > 0 {
			log.Print("Received unparseable address: ", recipient)
		}
		ret.Code = smtpump.SMTP_PARAMETER_NOT_IMPLEMENTED
		ret.Message = "Address not understood, sorry."
		return
	}

	for _, addr = range matches {
		if len(addr) > 0 {
			realaddr = addr
		}
	}
	msg.SmtpTo = append(msg.SmtpTo, realaddr)
	ret.Code = smtpump.SMTP_COMPLETED
	ret.Message = "Ok."
	return
}

// Read the data following the DATA command, up to the configured limit.
func (self smtpCallback) Data(conn *smtpump.SmtpConnection) (
	ret smtpump.SmtpReturnCode) {
	var msg *mailpump.MailMessage = getConnectionData(conn)
	var hdr string
	var vals []string
	var addrs []*mail.Address
	var addr *mail.Address
	var dotreader io.Reader
	var contentsreader *io.LimitedReader
	var message *mail.Message
	var tm time.Time
	var err error

	if msg.SmtpHelo == nil {
		ret.Code = smtpump.SMTP_BAD_SEQUENCE
		ret.Message = "Polite people say Hello first!"
		return
	}

	if msg.SmtpFrom == nil {
		ret.Code = smtpump.SMTP_BAD_SEQUENCE
		ret.Message = "Need MAIL command before DATA."
		return
	}

	if len(msg.SmtpTo) == 0 {
		ret.Code = smtpump.SMTP_BAD_SEQUENCE
		ret.Message = "Need RCPT command before DATA."
		return
	}

	conn.Respond(smtpump.SMTP_PROCEED, false, "Proceed with message.")

	dotreader = conn.GetDotReader()
	contentsreader = &io.LimitedReader{
		R: dotreader,
		N: self.maxContentLength + 1,
	}
	message, err = mail.ReadMessage(contentsreader)
	if err != nil {
		ret.Code = smtpump.SMTP_LOCALERR
		ret.Message = "Unable to read message: " + err.Error()
		// Consume all remaining output before returning an error.
		ioutil.ReadAll(dotreader)
		return
	}

	// See if we ran out of bytes to our limit
	if contentsreader.N <= 0 {
		ret.Code = smtpump.SMTP_MESSAGE_TOO_BIG
		ret.Message = "Size limit exceeded. Thanks for playing."
		ret.Terminate = true
		return
	}

	msg.Body, err = ioutil.ReadAll(message.Body)
	if err != nil {
		ret.Code = smtpump.SMTP_LOCALERR
		ret.Message = "Unable to parse message: " + err.Error()
		return
	}

	for hdr, vals = range message.Header {
		var header = new(mailpump.MailMessage_MailHeader)
		header.Name = new(string)
		*header.Name = hdr
		header.Value = make([]string, len(vals))
		copy(header.Value, vals)
		msg.Headers = append(msg.Headers, header)
	}

	tm, err = message.Header.Date()
	if err == nil {
		msg.DateHdr = new(int64)
		*msg.DateHdr = tm.Unix()
	}

	addrs, _ = message.Header.AddressList("From")
	if len(addrs) > 0 {
		msg.FromHdr = new(string)
		*msg.FromHdr = addrs[0].String()
	}

	addrs, _ = message.Header.AddressList("To")
	for _, addr = range addrs {
		msg.ToHdr = append(msg.ToHdr, addr.String())
	}

	addrs, _ = message.Header.AddressList("Cc")
	for _, addr = range addrs {
		msg.CcHdrs = append(msg.CcHdrs, addr.String())
	}

	addrs, _ = message.Header.AddressList("Sender")
	if len(addrs) > 0 {
		msg.SenderHdr = new(string)
		*msg.SenderHdr = addrs[0].String()
	}

	hdr = message.Header.Get("Message-Id")
	if len(hdr) <= 0 {
		hdr = message.Header.Get("Message-ID")
	}
	if len(hdr) > 0 {
		msg.MsgidHdr = new(string)
		*msg.MsgidHdr = hdr
	}

	// TODO(caoimhe): send this to some server.
	ret.Code = smtpump.SMTP_NOT_IMPLEMENTED
	ret.Message = "Ok, but this doesn't go anywhere yet."
	return
}

// FIXME: STUB.
func (self smtpCallback) Etrn(conn *smtpump.SmtpConnection, domain string) (
	ret smtpump.SmtpReturnCode) {
	ret.Code = smtpump.SMTP_NOT_IMPLEMENTED
	ret.Message = "Not yet implemented."
	return
}

// Forget all connection related data except HELO and the peer information.
func (self smtpCallback) Reset(conn *smtpump.SmtpConnection) (
	ret smtpump.SmtpReturnCode) {
	var msg = getConnectionData(conn)
	var peer, tlsc, helo string
	var rdns []string
	peer = msg.GetSmtpPeer()
	tlsc = msg.GetSmtpPeerTlsInfo()
	rdns = msg.GetSmtpPeerRevdns()
	helo = msg.GetSmtpHelo()
	msg.Reset()

	msg.SmtpPeer = &peer
	if len(tlsc) > 0 {
		msg.SmtpPeerTlsInfo = &tlsc
	}
	if len(rdns) > 0 {
		msg.SmtpPeerRevdns = make([]string, len(rdns))
		copy(msg.SmtpPeerRevdns, rdns)
	}
	if len(helo) > 0 {
		msg.SmtpHelo = &helo
	}
	ret.Code = smtpump.SMTP_COMPLETED
	ret.Message = "Ok."
	return
}

// Close the connection with a friendly message.
func (self smtpCallback) Quit(conn *smtpump.SmtpConnection) (
	ret smtpump.SmtpReturnCode) {
	ret.Code = smtpump.SMTP_CLOSING
	ret.Message = "See you later!"
	ret.Terminate = true
	return
}
<|endoftext|>


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

package smstorage

import (
	"sort"
	"strconv"
	"testing"

	"github.com/fullstorydev/gosolr/solrman/solrmanapi"
)

func testStorage_InProgressOps(t *testing.T, s SolrManStorage) {
	assertOps := func(expectOps ...string) {
		ops, err := s.GetInProgressOps()
		if err != nil {
			t.Errorf("GetInProgressOps failed: %s", err)
			return
		}
		if len(expectOps) != len(ops) {
			t.Errorf("expect len %d != actual len %d", len(expectOps), len(ops))
			return
		}
		for i := range ops {
			if expectOps[i] != ops[i].Key() {
				t.Errorf("expect %s != actual %s", expectOps[i], ops[i].Key())
			}
		}
	}

	assertOps()

	ops := []solrmanapi.OpRecord{
		{Collection: "foo", Shard: "1", StartedMs: 1},
		{Collection: "bar", Shard: "2", StartedMs: 2},
		{Collection: "baz", Shard: "3", StartedMs: 3},
	}

	s.AddInProgressOp(ops[0])
	assertOps("SolrOp:foo:1")

	s.AddInProgressOp(ops[2])
	assertOps("SolrOp:baz:3", "SolrOp:foo:1")

	s.AddInProgressOp(ops[1])
	assertOps("SolrOp:baz:3", "SolrOp:bar:2", "SolrOp:foo:1")

	s.DelInProgressOp(ops[0])
	assertOps("SolrOp:baz:3", "SolrOp:bar:2")

	s.DelInProgressOp(ops[2])
	assertOps("SolrOp:bar:2")

	s.DelInProgressOp(ops[1])
	assertOps()
}

func testStorage_CompletedOps(t *testing.T, s SolrManStorage) {
	assertOps := func(count int, expectOps ...string) {
		ops, err := s.GetCompletedOps(count)
		if err != nil {
			t.Errorf("GetInProgressOps failed: %s", err)
			return
		}
		if len(expectOps) != len(ops) {
			t.Errorf("expect len %d != actual len %d", len(expectOps), len(ops))
			return
		}
		for i := range ops {
			if expectOps[i] != ops[i].Key() {
				t.Errorf("expect %s != actual %s", expectOps[i], ops[i].Key())
			}
		}
	}

	assertOps(0)
	assertOps(1)
	assertOps(99)

	ops := []solrmanapi.OpRecord{
		{Collection: "foo", Shard: "1", FinishedMs: 1},
		{Collection: "bar", Shard: "2", FinishedMs: 2},
		{Collection: "baz", Shard: "3", FinishedMs: 3},
	}

	s.AddCompletedOp(ops[0])
	assertOps(99, "SolrOp:foo:1")
	assertOps(1, "SolrOp:foo:1")
	assertOps(0)

	s.AddCompletedOp(ops[1])
	assertOps(99, "SolrOp:bar:2", "SolrOp:foo:1")
	assertOps(2, "SolrOp:bar:2", "SolrOp:foo:1")
	assertOps(1, "SolrOp:bar:2")
	assertOps(0)

	s.AddCompletedOp(ops[2])
	assertOps(99, "SolrOp:baz:3", "SolrOp:bar:2", "SolrOp:foo:1")
	assertOps(3, "SolrOp:baz:3", "SolrOp:bar:2", "SolrOp:foo:1")
	assertOps(2, "SolrOp:baz:3", "SolrOp:bar:2")
	assertOps(1, "SolrOp:baz:3")
	assertOps(0)

	// Now add a bunch and make sure only NumStoredCompletedOps come back.
	bigNum := NumStoredCompletedOps * 2
	var expectOps []string
	for i := 0; i < bigNum; i++ {
		op := solrmanapi.OpRecord{Collection: "max", Shard: strconv.Itoa(i), FinishedMs: int64(i)}
		s.AddCompletedOp(op)
		expectOps = append(expectOps, op.Key())
	}
	expectOps = expectOps[len(expectOps)-NumStoredCompletedOps:]
	sort.Sort(sort.Reverse(sort.StringSlice(expectOps))) // should come back in reverse order
	assertOps(bigNum, expectOps...)
	// A negative count also returns all values.
	assertOps(-1, expectOps...)
}
<commit_msg>Fixing test (#51)<commit_after>// Copyright 2017 FullStory, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package smstorage

import (
	"sort"
	"strconv"
	"testing"

	"github.com/fullstorydev/gosolr/solrman/solrmanapi"
)

func testStorage_InProgressOps(t *testing.T, s SolrManStorage) {
	assertOps := func(expectOps ...string) {
		ops, err := s.GetInProgressOps()
		if err != nil {
			t.Errorf("GetInProgressOps failed: %s", err)
			return
		}
		if len(expectOps) != len(ops) {
			t.Errorf("expect len %d != actual len %d", len(expectOps), len(ops))
			return
		}
		for i := range ops {
			if expectOps[i] != ops[i].Key() {
				t.Errorf("expect %s != actual %s", expectOps[i], ops[i].Key())
			}
		}
	}

	assertOps()

	ops := []solrmanapi.OpRecord{
		{Collection: "foo", Shard: "1", StartedMs: 1, Replica: "core1"},
		{Collection: "bar", Shard: "2", StartedMs: 2, Replica: "core2"},
		{Collection: "baz", Shard: "3", StartedMs: 3, Replica: "core3"},
	}

	s.AddInProgressOp(ops[0])
	assertOps("SolrOp:foo:1:core1")

	s.AddInProgressOp(ops[2])
	assertOps("SolrOp:baz:3:core3", "SolrOp:foo:1:core1")

	s.AddInProgressOp(ops[1])
	assertOps("SolrOp:baz:3:core3", "SolrOp:bar:2:core2", "SolrOp:foo:1:core1")

	s.DelInProgressOp(ops[0])
	assertOps("SolrOp:baz:3:core3", "SolrOp:bar:2:core2")

	s.DelInProgressOp(ops[2])
	assertOps("SolrOp:bar:2:core2")

	s.DelInProgressOp(ops[1])
	assertOps()
}

func testStorage_CompletedOps(t *testing.T, s SolrManStorage) {
	assertOps := func(count int, expectOps ...string) {
		ops, err := s.GetCompletedOps(count)
		if err != nil {
			t.Errorf("GetInProgressOps failed: %s", err)
			return
		}
		if len(expectOps) != len(ops) {
			t.Errorf("expect len %d != actual len %d", len(expectOps), len(ops))
			return
		}
		for i := range ops {
			if expectOps[i] != ops[i].Key() {
				t.Errorf("expect %s != actual %s", expectOps[i], ops[i].Key())
			}
		}
	}

	assertOps(0)
	assertOps(1)
	assertOps(99)

	ops := []solrmanapi.OpRecord{
		{Collection: "foo", Shard: "1", FinishedMs: 1},
		{Collection: "bar", Shard: "2", FinishedMs: 2},
		{Collection: "baz", Shard: "3", FinishedMs: 3},
	}

	s.AddCompletedOp(ops[0])
	assertOps(99, "SolrOp:foo:1:")
	assertOps(1, "SolrOp:foo:1:")
	assertOps(0)

	s.AddCompletedOp(ops[1])
	assertOps(99, "SolrOp:bar:2:", "SolrOp:foo:1:")
	assertOps(2, "SolrOp:bar:2:", "SolrOp:foo:1:")
	assertOps(1, "SolrOp:bar:2:")
	assertOps(0)

	s.AddCompletedOp(ops[2])
	assertOps(99, "SolrOp:baz:3:", "SolrOp:bar:2:", "SolrOp:foo:1:")
	assertOps(3, "SolrOp:baz:3:", "SolrOp:bar:2:", "SolrOp:foo:1:")
	assertOps(2, "SolrOp:baz:3:", "SolrOp:bar:2:")
	assertOps(1, "SolrOp:baz:3:")
	assertOps(0)

	// Now add a bunch and make sure only NumStoredCompletedOps come back.
	bigNum := NumStoredCompletedOps * 2
	var expectOps []string
	for i := 0; i < bigNum; i++ {
		op := solrmanapi.OpRecord{Collection: "max", Shard: strconv.Itoa(i), FinishedMs: int64(i)}
		s.AddCompletedOp(op)
		expectOps = append(expectOps, op.Key())
	}
	expectOps = expectOps[len(expectOps)-NumStoredCompletedOps:]
	sort.Sort(sort.Reverse(sort.StringSlice(expectOps))) // should come back in reverse order
	assertOps(bigNum, expectOps...)
	// A negative count also returns all values.
	assertOps(-1, expectOps...)
}
<|endoftext|>


<commit_before>package rtsengine

import (
	"fmt"
	"image"
	"io"
)

// HumanPlayer implements the IPlayer interface for a human player
type HumanPlayer struct {
	// Structures common to all players.
	BasePlayer

	// Live TCPWire to communicate with UI
	Wire *TCPWire
}

// NewHumanPlayer constructs a HumanPlayer
func NewHumanPlayer(description string, worldLocation image.Point, width int, height int, pool *Pool, pathing *AStarPathing, world *World) *HumanPlayer {
	player := HumanPlayer{}

	player.description = description
	player.GenerateView(worldLocation, width, height)
	player.ItemPool = pool
	player.Pathing = pathing
	player.OurWorld = world

	// Add mechanics

	return &player
}

/////////////////////////////////////////////////////////////////////////
//                           IPlayer interface                         //
/////////////////////////////////////////////////////////////////////////
func (player *HumanPlayer) listen(wire *TCPWire) {
	player.Wire = wire
}

func (player *HumanPlayer) isHuman() bool {
	return true
}

func (player *HumanPlayer) isWireAlive() bool {
	// Best guess
	return player.Wire != nil
}

func (player *HumanPlayer) start() error {

	if !player.isWireAlive() {
		return fmt.Errorf("Failed: This player does not have an active wire connection.")
	}

	go player.listenForWireCommands()
	return nil
}

func (player *HumanPlayer) stop() {

}

/////////////////////////////////////////////////////////////////////////
//                           IPlayer interface                         //
/////////////////////////////////////////////////////////////////////////

func (player *HumanPlayer) listenForWireCommands() {
	var packet WirePacket
	for {

		if err := player.Wire.JSONDecoder.Decode(&packet); err == io.EOF {
			fmt.Println("\n\nEOF was detected. Connection lost.")
			return
		}
		packet.Print()

		switch packet.Command {
		case FullView:
			player.View.Span = player.OurWorld.Grid.Span
			player.View.WorldOrigin = player.OurWorld.Grid.WorldOrigin

		case PartialRefreshPlayerToUI:
			for i := 0; i < player.View.Span.Dx(); i++ {
				for j := 0; j < player.View.Span.Dy(); j++ {

					// Convert View point to world and get acre in world.
					worldPoint := player.View.ToWorldPoint(&image.Point{i, j})
					if !player.OurWorld.In(&worldPoint) {
						continue
					}
					ourAcre := player.OurWorld.Matrix[worldPoint.X][worldPoint.Y]

					if ourAcre.IsOccupiedOrNotGrass() {
						packet.Clear()

						// Use View Coordinates
						packet.CurrentX = i
						packet.CurrentY = j
						packet.LocalTerrain = ourAcre.terrain
						if ourAcre.Occupied() {
							packet.Unit = ourAcre.unit.unitType()
						}

						//packet.Life = ourAcre.unit.
						packet.WorldWidth = player.OurWorld.Grid.Span.Dy()
						packet.WorldHeight = player.OurWorld.Grid.Span.Dx()
						packet.WorldX = 0
						packet.WorldY = 0

						packet.ViewWidth = player.View.Span.Dy()
						packet.ViewHeight = player.View.Span.Dx()
						packet.ViewX = player.View.WorldOrigin.X
						packet.ViewY = player.View.WorldOrigin.Y

						if err := player.Wire.JSONEncoder.Encode(&packet); err == io.EOF {
							fmt.Println("\n\nEOF was detected. Connection lost.")
							return
						}

					}
				}
			}

		} //switch

	} // for ever

}
<commit_msg>Comments.<commit_after>package rtsengine

import (
	"fmt"
	"image"
	"io"
)

// HumanPlayer implements the IPlayer interface for a human player
type HumanPlayer struct {
	// Structures common to all players.
	BasePlayer

	// Live TCPWire to communicate with UI
	Wire *TCPWire
}

// NewHumanPlayer constructs a HumanPlayer
func NewHumanPlayer(description string, worldLocation image.Point, width int, height int, pool *Pool, pathing *AStarPathing, world *World) *HumanPlayer {
	player := HumanPlayer{}

	player.description = description
	player.GenerateView(worldLocation, width, height)
	player.ItemPool = pool
	player.Pathing = pathing
	player.OurWorld = world

	// Add mechanics

	return &player
}

/////////////////////////////////////////////////////////////////////////
//                           IPlayer interface                         //
/////////////////////////////////////////////////////////////////////////
func (player *HumanPlayer) listen(wire *TCPWire) {
	player.Wire = wire
}

func (player *HumanPlayer) isHuman() bool {
	return true
}

func (player *HumanPlayer) isWireAlive() bool {
	// Best guess
	return player.Wire != nil
}

func (player *HumanPlayer) start() error {

	if !player.isWireAlive() {
		return fmt.Errorf("Failed: This player does not have an active wire connection.")
	}

	go player.listenForWireCommands()
	return nil
}

func (player *HumanPlayer) stop() {

}

/////////////////////////////////////////////////////////////////////////
//                           IPlayer interface                         //
/////////////////////////////////////////////////////////////////////////

func (player *HumanPlayer) listenForWireCommands() {
	var packet WirePacket
	for {

		if err := player.Wire.JSONDecoder.Decode(&packet); err == io.EOF {
			fmt.Println("\n\nEOF was detected. Connection lost.")
			return
		}
		packet.Print()

		switch packet.Command {

		// Set the View to equial the entire world. Used for testing.
		case FullView:
			player.View.Span = player.OurWorld.Grid.Span
			player.View.WorldOrigin = player.OurWorld.Grid.WorldOrigin

			// Return all non empty or non grass acres in the view.
		case PartialRefreshPlayerToUI:
			for i := 0; i < player.View.Span.Dx(); i++ {
				for j := 0; j < player.View.Span.Dy(); j++ {

					// Convert View point to world and get acre in world.
					worldPoint := player.View.ToWorldPoint(&image.Point{i, j})
					if !player.OurWorld.In(&worldPoint) {
						continue
					}
					ourAcre := player.OurWorld.Matrix[worldPoint.X][worldPoint.Y]

					if ourAcre.IsOccupiedOrNotGrass() {
						packet.Clear()

						// Use View Coordinates
						packet.CurrentX = i
						packet.CurrentY = j
						packet.LocalTerrain = ourAcre.terrain
						if ourAcre.Occupied() {
							packet.Unit = ourAcre.unit.unitType()
						}

						//packet.Life = ourAcre.unit.
						packet.WorldWidth = player.OurWorld.Grid.Span.Dy()
						packet.WorldHeight = player.OurWorld.Grid.Span.Dx()
						packet.WorldX = 0
						packet.WorldY = 0

						packet.ViewWidth = player.View.Span.Dy()
						packet.ViewHeight = player.View.Span.Dx()
						packet.ViewX = player.View.WorldOrigin.X
						packet.ViewY = player.View.WorldOrigin.Y

						if err := player.Wire.JSONEncoder.Encode(&packet); err == io.EOF {
							fmt.Println("\n\nEOF was detected. Connection lost.")
							return
						}

					}
				}
			}

		} //switch

	} // for ever

}
<|endoftext|>


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

package hammer

import (
	"context"
	"flag"
	"math/rand"
	"strings"
	"testing"
	"time"

	"github.com/google/trillian/monitoring"
	"github.com/google/trillian/storage/testdb"
	"github.com/google/trillian/testonly/integration"

	_ "github.com/google/trillian/merkle/coniks"    // register CONIKS_SHA512_256
	_ "github.com/google/trillian/merkle/maphasher" // register TEST_MAP_HASHER
)

var (
	operations = flag.Uint64("operations", 20, "Number of operations to perform")
	singleTX   = flag.Bool("single_transaction", false, "Experimental: whether to use a single transaction when updating the map")
)

func TestRetryExposesDeadlineError(t *testing.T) {
	testdb.SkipIfNoMySQL(t)
	ctx := context.Background()
	env, err := integration.NewMapEnv(ctx, *singleTX)
	if err != nil {
		t.Fatal(err)
	}
	defer env.Close()

	bias := MapBias{
		Bias: map[MapEntrypointName]int{
			GetLeavesName:    10,
			GetLeavesRevName: 10,
			SetLeavesName:    10,
			GetSMRName:       10,
			GetSMRRevName:    10,
		},
		InvalidChance: map[MapEntrypointName]int{
			GetLeavesName:    10,
			GetLeavesRevName: 10,
			SetLeavesName:    10,
			GetSMRName:       0,
			GetSMRRevName:    10,
		},
	}

	seed := time.Now().UTC().UnixNano() & 0xFFFFFFFF
	cfg := MapConfig{
		MapID:         0, // ephemeral tree
		Client:        env.Map,
		Admin:         env.Admin,
		MetricFactory: monitoring.InertMetricFactory{},
		RandSource:    rand.NewSource(seed),
		EPBias:        bias,
		LeafSize:      1000,
		ExtraSize:     100,
		MinLeaves:     800,
		MaxLeaves:     1200,
		Operations:    *operations,
		NumCheckers:   1,
	}

	ctx, cancel := context.WithTimeout(ctx, time.Nanosecond)
	defer cancel()
	if err, wantErr := HitMap(ctx, cfg), context.DeadlineExceeded; !strings.Contains(err.Error(), wantErr.Error()) {
		t.Fatalf("Got err %q, expected %q", err, wantErr)
	}
}

func TestInProcessMapHammer(t *testing.T) {
	testdb.SkipIfNoMySQL(t)
	ctx := context.Background()
	env, err := integration.NewMapEnv(ctx, *singleTX)
	if err != nil {
		t.Fatal(err)
	}
	defer env.Close()

	bias := MapBias{
		Bias: map[MapEntrypointName]int{
			GetLeavesName:    10,
			GetLeavesRevName: 10,
			SetLeavesName:    10,
			GetSMRName:       10,
			GetSMRRevName:    10,
		},
		InvalidChance: map[MapEntrypointName]int{
			GetLeavesName:    10,
			GetLeavesRevName: 10,
			SetLeavesName:    10,
			GetSMRName:       0,
			GetSMRRevName:    10,
		},
	}

	seed := time.Now().UTC().UnixNano() & 0xFFFFFFFF
	cfg := MapConfig{
		MapID:         0, // ephemeral tree
		Client:        env.Map,
		Admin:         env.Admin,
		MetricFactory: monitoring.InertMetricFactory{},
		RandSource:    rand.NewSource(seed),
		EPBias:        bias,
		LeafSize:      1000,
		ExtraSize:     100,
		MinLeaves:     800,
		MaxLeaves:     1200,
		Operations:    *operations,
		NumCheckers:   1,
	}
	if err := HitMap(ctx, cfg); err != nil {
		t.Fatalf("hammer failure: %v", err)
	}
}
<commit_msg>Reduce the min/max leaves in the hammer test (#1849)<commit_after>// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package hammer

import (
	"context"
	"flag"
	"math/rand"
	"strings"
	"testing"
	"time"

	"github.com/google/trillian/monitoring"
	"github.com/google/trillian/storage/testdb"
	"github.com/google/trillian/testonly/integration"

	_ "github.com/google/trillian/merkle/coniks"    // register CONIKS_SHA512_256
	_ "github.com/google/trillian/merkle/maphasher" // register TEST_MAP_HASHER
)

var (
	operations = flag.Uint64("operations", 20, "Number of operations to perform")
	singleTX   = flag.Bool("single_transaction", false, "Experimental: whether to use a single transaction when updating the map")
)

func TestRetryExposesDeadlineError(t *testing.T) {
	testdb.SkipIfNoMySQL(t)
	ctx := context.Background()
	env, err := integration.NewMapEnv(ctx, *singleTX)
	if err != nil {
		t.Fatal(err)
	}
	defer env.Close()

	bias := MapBias{
		Bias: map[MapEntrypointName]int{
			GetLeavesName:    10,
			GetLeavesRevName: 10,
			SetLeavesName:    10,
			GetSMRName:       10,
			GetSMRRevName:    10,
		},
		InvalidChance: map[MapEntrypointName]int{
			GetLeavesName:    10,
			GetLeavesRevName: 10,
			SetLeavesName:    10,
			GetSMRName:       0,
			GetSMRRevName:    10,
		},
	}

	seed := time.Now().UTC().UnixNano() & 0xFFFFFFFF
	cfg := MapConfig{
		MapID:         0, // ephemeral tree
		Client:        env.Map,
		Admin:         env.Admin,
		MetricFactory: monitoring.InertMetricFactory{},
		RandSource:    rand.NewSource(seed),
		EPBias:        bias,
		LeafSize:      1000,
		ExtraSize:     100,
		// TODO(mhutchinson): Increase these when #1845 is understood & fixed.
		MinLeaves:   100,
		MaxLeaves:   150,
		Operations:  *operations,
		NumCheckers: 1,
	}

	ctx, cancel := context.WithTimeout(ctx, time.Nanosecond)
	defer cancel()
	if err, wantErr := HitMap(ctx, cfg), context.DeadlineExceeded; !strings.Contains(err.Error(), wantErr.Error()) {
		t.Fatalf("Got err %q, expected %q", err, wantErr)
	}
}

func TestInProcessMapHammer(t *testing.T) {
	testdb.SkipIfNoMySQL(t)
	ctx := context.Background()
	env, err := integration.NewMapEnv(ctx, *singleTX)
	if err != nil {
		t.Fatal(err)
	}
	defer env.Close()

	bias := MapBias{
		Bias: map[MapEntrypointName]int{
			GetLeavesName:    10,
			GetLeavesRevName: 10,
			SetLeavesName:    10,
			GetSMRName:       10,
			GetSMRRevName:    10,
		},
		InvalidChance: map[MapEntrypointName]int{
			GetLeavesName:    10,
			GetLeavesRevName: 10,
			SetLeavesName:    10,
			GetSMRName:       0,
			GetSMRRevName:    10,
		},
	}

	seed := time.Now().UTC().UnixNano() & 0xFFFFFFFF
	cfg := MapConfig{
		MapID:         0, // ephemeral tree
		Client:        env.Map,
		Admin:         env.Admin,
		MetricFactory: monitoring.InertMetricFactory{},
		RandSource:    rand.NewSource(seed),
		EPBias:        bias,
		LeafSize:      1000,
		ExtraSize:     100,
		// TODO(mhutchinson): Increase these when #1845 is understood & fixed.
		MinLeaves:   100,
		MaxLeaves:   150,
		Operations:  *operations,
		NumCheckers: 1,
	}
	if err := HitMap(ctx, cfg); err != nil {
		t.Fatalf("hammer failure: %v", err)
	}
}
<|endoftext|>


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

package usb_test

import (
	"bytes"
	"log"
	"os"
	"testing"

	. "github.com/kylelemons/gousb/usb"
	"github.com/kylelemons/gousb/usbid"
)

func TestNoop(t *testing.T) {
	if os.Getenv("TRAVIS") == "true" {
		t.Skip("test known to fail on Travis")
	}
	c := NewContext()
	defer c.Close()
	c.Debug(0)
}

func TestEnum(t *testing.T) {
	if os.Getenv("TRAVIS") == "true" {
		t.Skip("test known to fail on Travis")
	}
	c := NewContext()
	defer c.Close()
	c.Debug(0)

	logDevice := func(t *testing.T, desc *Descriptor) {
		t.Logf("%03d.%03d %s", desc.Bus, desc.Address, usbid.Describe(desc))
		t.Logf("- Protocol: %s", usbid.Classify(desc))

		for _, cfg := range desc.Configs {
			t.Logf("- %s:", cfg)
			for _, alt := range cfg.Interfaces {
				t.Logf("  --------------")
				for _, iface := range alt.Setups {
					t.Logf("  - %s", iface)
					t.Logf("    - %s", usbid.Classify(iface))
					for _, end := range iface.Endpoints {
						t.Logf("    - %s (packet size: %d bytes)", end, end.MaxPacketSize)
					}
				}
			}
			t.Logf("  --------------")
		}
	}

	descs := []*Descriptor{}
	devs, err := c.ListDevices(func(desc *Descriptor) bool {
		logDevice(t, desc)
		descs = append(descs, desc)
		return true
	})
	defer func() {
		for _, d := range devs {
			d.Close()
		}
	}()
	if err != nil {
		t.Fatalf("list: %s", err)
	}

	if got, want := len(devs), len(descs); got != want {
		t.Fatalf("len(devs) = %d, want %d", got, want)
	}

	for i := range devs {
		if got, want := devs[i].Descriptor, descs[i]; got != want {
			t.Errorf("dev[%d].Descriptor = %p, want %p", i, got, want)
		}
	}
}

func TestOpenDeviceWithVidPid(t *testing.T) {
	c := NewContext()
	defer c.Close()
	c.Debug(0)

	// Accept for all device
	devs, err := c.ListDevices(func(desc *Descriptor) bool {
		return true
	})
	defer func() {
		for _, d := range devs {
			d.Close()
		}
	}()

	if err != nil {
		t.Fatalf("list: %s", err)
	}

	for i := range devs {
		vid := devs[i].Vendor
		pid := devs[i].Product
		device, err := c.OpenDeviceWithVidPid((int)(vid), (int)(pid))

		// if the context failed to open device
		if err != nil {
			t.Fail()
		}

		// if opened device was not valid
		if device.Descriptor.Bus != devs[i].Bus ||
			device.Descriptor.Address != devs[i].Address ||
			device.Vendor != devs[i].Vendor ||
			device.Product != devs[i].Product {
			t.Fail()
		}

	}
}

func TestMultipleContexts(t *testing.T) {
	var buf bytes.Buffer
	log.SetOutput(&buf)
	for i := 0; i < 2; i++ {
		ctx := NewContext()
		_, err := ctx.ListDevices(func(desc *Descriptor) bool {
			return false
		})
		if err != nil {
			t.Fatal(err)
		}
		ctx.Close()
	}
	log.SetOutput(os.Stderr)
	if buf.Len() > 0 {
		t.Errorf("Non zero output to log, while testing:  %s", buf.String())
	}
}
<commit_msg>more tests failing on travis<commit_after>// Copyright 2013 Google Inc.  All rights reserved.
// Copyright 2016 the gousb Authors.  All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package usb_test

import (
	"bytes"
	"log"
	"os"
	"testing"

	. "github.com/kylelemons/gousb/usb"
	"github.com/kylelemons/gousb/usbid"
)

func TestNoop(t *testing.T) {
	if os.Getenv("TRAVIS") == "true" {
		t.Skip("test known to fail on Travis")
	}
	c := NewContext()
	defer c.Close()
	c.Debug(0)
}

func TestEnum(t *testing.T) {
	if os.Getenv("TRAVIS") == "true" {
		t.Skip("test known to fail on Travis")
	}
	c := NewContext()
	defer c.Close()
	c.Debug(0)

	logDevice := func(t *testing.T, desc *Descriptor) {
		t.Logf("%03d.%03d %s", desc.Bus, desc.Address, usbid.Describe(desc))
		t.Logf("- Protocol: %s", usbid.Classify(desc))

		for _, cfg := range desc.Configs {
			t.Logf("- %s:", cfg)
			for _, alt := range cfg.Interfaces {
				t.Logf("  --------------")
				for _, iface := range alt.Setups {
					t.Logf("  - %s", iface)
					t.Logf("    - %s", usbid.Classify(iface))
					for _, end := range iface.Endpoints {
						t.Logf("    - %s (packet size: %d bytes)", end, end.MaxPacketSize)
					}
				}
			}
			t.Logf("  --------------")
		}
	}

	descs := []*Descriptor{}
	devs, err := c.ListDevices(func(desc *Descriptor) bool {
		logDevice(t, desc)
		descs = append(descs, desc)
		return true
	})
	defer func() {
		for _, d := range devs {
			d.Close()
		}
	}()
	if err != nil {
		t.Fatalf("list: %s", err)
	}

	if got, want := len(devs), len(descs); got != want {
		t.Fatalf("len(devs) = %d, want %d", got, want)
	}

	for i := range devs {
		if got, want := devs[i].Descriptor, descs[i]; got != want {
			t.Errorf("dev[%d].Descriptor = %p, want %p", i, got, want)
		}
	}
}

func TestOpenDeviceWithVidPid(t *testing.T) {
	if os.Getenv("TRAVIS") == "true" {
		t.Skip("test known to fail on Travis")
	}
	c := NewContext()
	defer c.Close()
	c.Debug(0)

	// Accept for all device
	devs, err := c.ListDevices(func(desc *Descriptor) bool {
		return true
	})
	defer func() {
		for _, d := range devs {
			d.Close()
		}
	}()

	if err != nil {
		t.Fatalf("list: %s", err)
	}

	for i := range devs {
		vid := devs[i].Vendor
		pid := devs[i].Product
		device, err := c.OpenDeviceWithVidPid((int)(vid), (int)(pid))

		// if the context failed to open device
		if err != nil {
			t.Fail()
		}

		// if opened device was not valid
		if device.Descriptor.Bus != devs[i].Bus ||
			device.Descriptor.Address != devs[i].Address ||
			device.Vendor != devs[i].Vendor ||
			device.Product != devs[i].Product {
			t.Fail()
		}

	}
}

func TestMultipleContexts(t *testing.T) {
	if os.Getenv("TRAVIS") == "true" {
		t.Skip("test known to fail on Travis")
	}
	var buf bytes.Buffer
	log.SetOutput(&buf)
	for i := 0; i < 2; i++ {
		ctx := NewContext()
		_, err := ctx.ListDevices(func(desc *Descriptor) bool {
			return false
		})
		if err != nil {
			t.Fatal(err)
		}
		ctx.Close()
	}
	log.SetOutput(os.Stderr)
	if buf.Len() > 0 {
		t.Errorf("Non zero output to log, while testing:  %s", buf.String())
	}
}
<|endoftext|>


<commit_before>// Copyright 2013 M-Lab
//
// 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 rtt provides a resolver for mlab-ns2 using RTT-based metrics.
package rtt

import (
	"errors"
	"net"
	"sort"
	"time"
)

const (
	v4PrefixSize = 22 // defines the size of the IPv4 group prefix
	v6PrefixSize = 56 // defines the size of the IPv6 group prefix
)

var (
	v4PrefixMask        = net.CIDRMask(v4PrefixSize, 8*net.IPv4len)
	v6PrefixMask        = net.CIDRMask(v6PrefixSize, 8*net.IPv6len)
	ErrMergeSiteRTT     = errors.New("SiteRTT cannot be merged, mismatching Site IDs.")
	ErrMergeClientGroup = errors.New("ClientGroups cannot be merged, mismatching ClientGroup Prefixes.")
)

// GetClientGroup returns a *net.IPNet which represents a subnet of prefix
// size v4PrefixSize in the case of IPv4 addresses.
func GetClientGroup(ip net.IP) *net.IPNet {
	if ip.To4() == nil {
		return &net.IPNet{IP: ip.Mask(v6PrefixMask), Mask: v6PrefixMask}
	}
	return &net.IPNet{IP: ip.Mask(v4PrefixMask), Mask: v4PrefixMask}
}

// IsEqualClientGroup checks if two IPs are in the same client group defined
// by prefix sizes defined by v4PrefixSize and v6PrefixSize.
func IsEqualClientGroup(a, b net.IP) bool {
	ipnet := GetClientGroup(a)
	return ipnet.Contains(b)
}

// MergeSiteRTTs merges a new SiteRTT entry into an old SiteRTT entry if the new
// entry has lower or equal RTT, and also reports whether the merge has caused
// any changes.
func MergeSiteRTTs(oldSR, newSR *SiteRTT) (bool, error) {
	if oldSR.SiteID != newSR.SiteID {
		return false, ErrMergeSiteRTT
	}
	if newSR.RTT <= oldSR.RTT {
		oldSR.RTT = newSR.RTT
		oldSR.LastUpdated = newSR.LastUpdated
		return true, nil
	}
	return false, nil
}

// MergeClientGroups merges a new list of SiteRTT with an existing list of
// SiteRTT and sorts it in ascending RTT order. It also reports if the merge has
// caused any changes.
// Note: Used for merging new bigquery data with existing datastore data.
func MergeClientGroups(oldCG, newCG *ClientGroup) (bool, error) {
	oIP, nIP := net.IP(oldCG.Prefix), net.IP(newCG.Prefix)
	if !oIP.Equal(nIP) {
		return false, ErrMergeClientGroup
	}

	// Populate temporary maps to ease merge
	oRTTs := make(map[string]*SiteRTT)
	nRTTs := make(map[string]*SiteRTT)
	for i, s := range oldCG.SiteRTTs {
		oRTTs[s.SiteID] = &oldCG.SiteRTTs[i]
	}
	for i, s := range newCG.SiteRTTs {
		nRTTs[s.SiteID] = &newCG.SiteRTTs[i]
	}

	// Keep SiteRTT with lower RTT
	var os *SiteRTT
	var ok, changed, srChanged bool
	var err error
	for k, ns := range nRTTs {
		os, ok = oRTTs[k]
		if !ok {
			oRTTs[k] = ns
			changed = true
		} else {
			srChanged, err = MergeSiteRTTs(os, ns)
			if err != nil {
				return false, err
			}
			if srChanged {
				changed = true
			}
		}
	}

	// Create new list of SiteRTTs
	oldCG.SiteRTTs = make(SiteRTTs, 0, len(oRTTs))
	for _, s := range oRTTs {
		oldCG.SiteRTTs = append(oldCG.SiteRTTs, *s)
	}
	sort.Sort(oldCG.SiteRTTs)

	return changed, nil
}
<commit_msg>rtt: Remove unused import "time"<commit_after>// Copyright 2013 M-Lab
//
// 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 rtt provides a resolver for mlab-ns2 using RTT-based metrics.
package rtt

import (
	"errors"
	"net"
	"sort"
)

const (
	v4PrefixSize = 22 // defines the size of the IPv4 group prefix
	v6PrefixSize = 56 // defines the size of the IPv6 group prefix
)

var (
	v4PrefixMask        = net.CIDRMask(v4PrefixSize, 8*net.IPv4len)
	v6PrefixMask        = net.CIDRMask(v6PrefixSize, 8*net.IPv6len)
	ErrMergeSiteRTT     = errors.New("SiteRTT cannot be merged, mismatching Site IDs.")
	ErrMergeClientGroup = errors.New("ClientGroups cannot be merged, mismatching ClientGroup Prefixes.")
)

// GetClientGroup returns a *net.IPNet which represents a subnet of prefix
// size v4PrefixSize in the case of IPv4 addresses.
func GetClientGroup(ip net.IP) *net.IPNet {
	if ip.To4() == nil {
		return &net.IPNet{IP: ip.Mask(v6PrefixMask), Mask: v6PrefixMask}
	}
	return &net.IPNet{IP: ip.Mask(v4PrefixMask), Mask: v4PrefixMask}
}

// IsEqualClientGroup checks if two IPs are in the same client group defined
// by prefix sizes defined by v4PrefixSize and v6PrefixSize.
func IsEqualClientGroup(a, b net.IP) bool {
	ipnet := GetClientGroup(a)
	return ipnet.Contains(b)
}

// MergeSiteRTTs merges a new SiteRTT entry into an old SiteRTT entry if the new
// entry has lower or equal RTT, and also reports whether the merge has caused
// any changes.
func MergeSiteRTTs(oldSR, newSR *SiteRTT) (bool, error) {
	if oldSR.SiteID != newSR.SiteID {
		return false, ErrMergeSiteRTT
	}
	if newSR.RTT <= oldSR.RTT {
		oldSR.RTT = newSR.RTT
		oldSR.LastUpdated = newSR.LastUpdated
		return true, nil
	}
	return false, nil
}

// MergeClientGroups merges a new list of SiteRTT with an existing list of
// SiteRTT and sorts it in ascending RTT order. It also reports if the merge has
// caused any changes.
// Note: Used for merging new bigquery data with existing datastore data.
func MergeClientGroups(oldCG, newCG *ClientGroup) (bool, error) {
	oIP, nIP := net.IP(oldCG.Prefix), net.IP(newCG.Prefix)
	if !oIP.Equal(nIP) {
		return false, ErrMergeClientGroup
	}

	// Populate temporary maps to ease merge
	oRTTs := make(map[string]*SiteRTT)
	nRTTs := make(map[string]*SiteRTT)
	for i, s := range oldCG.SiteRTTs {
		oRTTs[s.SiteID] = &oldCG.SiteRTTs[i]
	}
	for i, s := range newCG.SiteRTTs {
		nRTTs[s.SiteID] = &newCG.SiteRTTs[i]
	}

	// Keep SiteRTT with lower RTT
	var os *SiteRTT
	var ok, changed, srChanged bool
	var err error
	for k, ns := range nRTTs {
		os, ok = oRTTs[k]
		if !ok {
			oRTTs[k] = ns
			changed = true
		} else {
			srChanged, err = MergeSiteRTTs(os, ns)
			if err != nil {
				return false, err
			}
			if srChanged {
				changed = true
			}
		}
	}

	// Create new list of SiteRTTs
	oldCG.SiteRTTs = make(SiteRTTs, 0, len(oRTTs))
	for _, s := range oRTTs {
		oldCG.SiteRTTs = append(oldCG.SiteRTTs, *s)
	}
	sort.Sort(oldCG.SiteRTTs)

	return changed, nil
}
<|endoftext|>


<commit_before>package nyb

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	kitty "github.com/ugjka/kittybot"
	"gopkg.in/ugjka/go-tz.v2/tz"
)

const helpMsg = "COMMANDS: '%shny <location>', '%stime <location>', '%snext', '%sprevious', '%sremaining', '%shelp', '%ssource'"

func (bot *Settings) addTriggers() {
	irc := bot.irc

	//Log Notices
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "NOTICE"
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("[NOTICE] " + m.Content)
		},
	})

	//Trigger for !source
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"source")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Reply(m, "https://github.com/ugjka/newyearsbot")
		},
	})

	//Trigger for !help
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"help") ||
				normalize(m.Content) == bot.Prefix+"hny"
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying help...")
			b.Reply(m, fmt.Sprintf(helpMsg, bot.Prefix, bot.Prefix, bot.Prefix, bot.Prefix, bot.Prefix, bot.Prefix, bot.Prefix))
		},
	})

	//Trigger for !next
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"next")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying next...")
			dur := time.Minute * time.Duration(bot.next.Offset*60)
			if timeNow().UTC().Add(dur).After(target) {
				b.Reply(m, fmt.Sprintf("No more next, %d is here AoE", target.Year()))
				return
			}
			hdur := humanDur(target.Sub(timeNow().UTC().Add(dur)))
			b.Reply(m, fmt.Sprintf("Next New Year in %s in %s",
				hdur, bot.next))
		},
	})

	//Trigger for !previous
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"previous")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying previous...")
			dur := time.Minute * time.Duration(bot.previous.Offset*60)
			hdur := humanDur(timeNow().UTC().Add(dur).Sub(target))
			if bot.previous.Offset == -12 {
				hdur = humanDur(timeNow().UTC().Add(dur).Sub(target.AddDate(-1, 0, 0)))
			}
			b.Reply(m, fmt.Sprintf("Previous New Year %s ago in %s",
				hdur, bot.previous))
		},
	})

	//Trigger for !remaining
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"remaining")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying remaining...")
			plural := "s"
			if bot.remaining == 1 {
				plural = ""
			}
			pct := (len(bot.zones) - bot.remaining) * 100 / (len(bot.zones) * 100)
			b.Reply(m, fmt.Sprintf("%d timezone%s remaining. %d%% are in the new year", bot.remaining, plural, pct))
		},
	})

	//Trigger for time in location
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"time ")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying time...")
			result, err := bot.time(normalize(m.Content)[len(bot.Prefix)+len("time")+1:])
			if err == errNoZone || err == errNoPlace {
				b.Warn("Query error: " + err.Error())
				b.Reply(m, err.Error())
				return
			}
			if err != nil {
				b.Warn("Query error: " + err.Error())
				b.Reply(m, "Some error occurred!")
				return
			}
			b.Reply(m, result)
		},
	})

	//Trigger for UTC time
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				normalize(m.Content) == bot.Prefix+"time"
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying time...")
			result := "Time is " + time.Now().UTC().Format("Mon Jan 2 15:04:05 -0700 MST 2006")
			b.Reply(m, result)
		},
	})

	//Trigger for new year in location
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"hny ")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			result, err := bot.newYear(normalize(m.Content)[len(bot.Prefix)+len("hny")+1:])
			if err == errNoZone || err == errNoPlace {
				b.Warn("Query error: " + err.Error())
				b.Reply(m, err.Error())
				return
			}
			if err != nil {
				b.Warn("Query error: " + err.Error())
				b.Reply(m, "Some error occurred!")
				return
			}
			b.Reply(m, result)
		},
	})
}

var (
	errNoZone  = errors.New("couldn't get timezone for that location")
	errNoPlace = errors.New("Couldn't find that place")
)

func (bot *Settings) time(location string) (string, error) {
	bot.irc.Info("Querying location: " + location)
	data, err := NominatimFetcher(&bot.Email, &bot.Nominatim, &location)
	if err != nil {
		bot.irc.Warn("Nominatim error: " + err.Error())
		return "", err
	}
	var res NominatimResults
	if err = json.Unmarshal(data, &res); err != nil {
		bot.irc.Warn("Nominatim JSON error: " + err.Error())
		return "", err
	}
	if len(res) == 0 {
		return "", errNoPlace
	}
	p := tz.Point{
		Lat: res[0].Lat,
		Lon: res[0].Lon,
	}
	tzid, err := tz.GetZone(p)
	if err != nil {
		return "", errNoZone
	}
	zone, err := time.LoadLocation(tzid[0])
	if err != nil {
		return "", errNoZone
	}
	address := res[0].DisplayName
	msg := fmt.Sprintf("Time in %s is %s", address, time.Now().In(zone).Format("Mon Jan 2 15:04:05 -0700 MST 2006"))
	return msg, nil
}

func (bot *Settings) newYear(location string) (string, error) {
	bot.irc.Info("Querying location: " + location)
	data, err := NominatimFetcher(&bot.Email, &bot.Nominatim, &location)
	if err != nil {
		bot.irc.Warn("Nominatim error: " + err.Error())
		return "", err
	}
	var res NominatimResults
	if err = json.Unmarshal(data, &res); err != nil {
		bot.irc.Warn("Nominatim JSON error: " + err.Error())
		return "", err
	}
	if len(res) == 0 {
		return "", errNoPlace
	}
	p := tz.Point{
		Lat: res[0].Lat,
		Lon: res[0].Lon,
	}
	tzid, err := tz.GetZone(p)
	if err != nil {
		return "", errNoZone
	}
	zone, err := time.LoadLocation(tzid[0])
	if err != nil {
		return "", errNoZone
	}
	offset := zoneOffset(target, zone)
	address := res[0].DisplayName
	if timeNow().UTC().Add(offset).Before(target) {
		hdur := humanDur(target.Sub(timeNow().UTC().Add(offset)))
		const newYearFutureMsg = "New Year in %s will happen in %s"
		return fmt.Sprintf(newYearFutureMsg, address, hdur), nil
	}
	hdur := humanDur(timeNow().UTC().Add(offset).Sub(target))
	const newYearPastMsg = "New Year in %s happened %s ago"
	return fmt.Sprintf(newYearPastMsg, address, hdur), nil
}
<commit_msg>2020 pull request bugs<commit_after>package nyb

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	kitty "github.com/ugjka/kittybot"
	"gopkg.in/ugjka/go-tz.v2/tz"
)

const helpMsg = "COMMANDS: '%shny <location>', '%stime <location>', '%snext', '%sprevious', '%sremaining', '%shelp', '%ssource'"

func (bot *Settings) addTriggers() {
	irc := bot.irc

	//Log Notices
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "NOTICE"
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("[NOTICE] " + m.Content)
		},
	})

	//Trigger for !source
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"source")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Reply(m, "https://github.com/ugjka/newyearsbot")
		},
	})

	//Trigger for !help
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"help") ||
				normalize(m.Content) == bot.Prefix+"hny"
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying help...")
			b.Reply(m, fmt.Sprintf(helpMsg, bot.Prefix, bot.Prefix, bot.Prefix, bot.Prefix, bot.Prefix, bot.Prefix, bot.Prefix))
		},
	})

	//Trigger for !next
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"next")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying next...")
			dur := time.Minute * time.Duration(bot.next.Offset*60)
			if timeNow().UTC().Add(dur).After(target) {
				b.Reply(m, fmt.Sprintf("No more next, %d is here AoE", target.Year()))
				return
			}
			hdur := humanDur(target.Sub(timeNow().UTC().Add(dur)))
			b.Reply(m, fmt.Sprintf("Next New Year in %s in %s",
				hdur, bot.next))
		},
	})

	//Trigger for !previous
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"previous")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying previous...")
			dur := time.Minute * time.Duration(bot.previous.Offset*60)
			hdur := humanDur(timeNow().UTC().Add(dur).Sub(target))
			if bot.previous.Offset == -12 {
				hdur = humanDur(timeNow().UTC().Add(dur).Sub(target.AddDate(-1, 0, 0)))
			}
			b.Reply(m, fmt.Sprintf("Previous New Year %s ago in %s",
				hdur, bot.previous))
		},
	})

	//Trigger for !remaining
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"remaining")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying remaining...")
			plural := "s"
			if bot.remaining == 1 {
				plural = ""
			}
			pct := ((float64(len(bot.zones)) - float64(bot.remaining)) / float64(len(bot.zones)) * 100)
			b.Reply(m, fmt.Sprintf("%d timezone%s remaining. %.2f%% are in the new year", bot.remaining, plural, pct))
		},
	})

	//Trigger for time in location
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"time ")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying time...")
			result, err := bot.time(normalize(m.Content)[len(bot.Prefix)+len("time")+1:])
			if err == errNoZone || err == errNoPlace {
				b.Warn("Query error: " + err.Error())
				b.Reply(m, err.Error())
				return
			}
			if err != nil {
				b.Warn("Query error: " + err.Error())
				b.Reply(m, "Some error occurred!")
				return
			}
			b.Reply(m, result)
		},
	})

	//Trigger for UTC time
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				normalize(m.Content) == bot.Prefix+"time"
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			b.Info("Querying time...")
			result := "Time is " + time.Now().UTC().Format("Mon Jan 2 15:04:05 -0700 MST 2006")
			b.Reply(m, result)
		},
	})

	//Trigger for new year in location
	irc.AddTrigger(kitty.Trigger{
		Condition: func(b *kitty.Bot, m *kitty.Message) bool {
			return m.Command == "PRIVMSG" &&
				strings.HasPrefix(normalize(m.Content), bot.Prefix+"hny ")
		},
		Action: func(b *kitty.Bot, m *kitty.Message) {
			result, err := bot.newYear(normalize(m.Content)[len(bot.Prefix)+len("hny")+1:])
			if err == errNoZone || err == errNoPlace {
				b.Warn("Query error: " + err.Error())
				b.Reply(m, err.Error())
				return
			}
			if err != nil {
				b.Warn("Query error: " + err.Error())
				b.Reply(m, "Some error occurred!")
				return
			}
			b.Reply(m, result)
		},
	})
}

var (
	errNoZone  = errors.New("couldn't get timezone for that location")
	errNoPlace = errors.New("Couldn't find that place")
)

func (bot *Settings) time(location string) (string, error) {
	bot.irc.Info("Querying location: " + location)
	data, err := NominatimFetcher(&bot.Email, &bot.Nominatim, &location)
	if err != nil {
		bot.irc.Warn("Nominatim error: " + err.Error())
		return "", err
	}
	var res NominatimResults
	if err = json.Unmarshal(data, &res); err != nil {
		bot.irc.Warn("Nominatim JSON error: " + err.Error())
		return "", err
	}
	if len(res) == 0 {
		return "", errNoPlace
	}
	p := tz.Point{
		Lat: res[0].Lat,
		Lon: res[0].Lon,
	}
	tzid, err := tz.GetZone(p)
	if err != nil {
		return "", errNoZone
	}
	zone, err := time.LoadLocation(tzid[0])
	if err != nil {
		return "", errNoZone
	}
	address := res[0].DisplayName
	msg := fmt.Sprintf("Time in %s is %s", address, time.Now().In(zone).Format("Mon Jan 2 15:04:05 -0700 MST 2006"))
	return msg, nil
}

func (bot *Settings) newYear(location string) (string, error) {
	bot.irc.Info("Querying location: " + location)
	data, err := NominatimFetcher(&bot.Email, &bot.Nominatim, &location)
	if err != nil {
		bot.irc.Warn("Nominatim error: " + err.Error())
		return "", err
	}
	var res NominatimResults
	if err = json.Unmarshal(data, &res); err != nil {
		bot.irc.Warn("Nominatim JSON error: " + err.Error())
		return "", err
	}
	if len(res) == 0 {
		return "", errNoPlace
	}
	p := tz.Point{
		Lat: res[0].Lat,
		Lon: res[0].Lon,
	}
	tzid, err := tz.GetZone(p)
	if err != nil {
		return "", errNoZone
	}
	zone, err := time.LoadLocation(tzid[0])
	if err != nil {
		return "", errNoZone
	}
	offset := zoneOffset(target, zone)
	address := res[0].DisplayName
	if timeNow().UTC().Add(offset).Before(target) {
		hdur := humanDur(target.Sub(timeNow().UTC().Add(offset)))
		const newYearFutureMsg = "New Year in %s will happen in %s"
		return fmt.Sprintf(newYearFutureMsg, address, hdur), nil
	}
	hdur := humanDur(timeNow().UTC().Add(offset).Sub(target))
	const newYearPastMsg = "New Year in %s happened %s ago"
	return fmt.Sprintf(newYearPastMsg, address, hdur), nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2017, 2019, Oracle and/or its affiliates. All rights reserved.

package oci

import (
	"log"
)

const Version = "3.70.0"

func PrintVersion() {
	log.Printf("[INFO] terraform-provider-oci %s\n", Version)
}
<commit_msg>Finalize changelog and release for version v3.71.0<commit_after>// Copyright (c) 2017, 2019, Oracle and/or its affiliates. All rights reserved.

package oci

import (
	"log"
)

const Version = "3.71.0"

func PrintVersion() {
	log.Printf("[INFO] terraform-provider-oci %s\n", Version)
}
<|endoftext|>


<commit_before>package template

import (
	"bytes"
	"errors"
	"fmt"
	"html/template"
	"io/ioutil"
	"os"
	"os/exec"
	"os/signal"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/HeavyHorst/memkv"
	"github.com/HeavyHorst/remco/backends"
	"github.com/HeavyHorst/remco/template/fileutil"
	"github.com/cloudflare/cfssl/log"
	"github.com/flosch/pongo2"
	flag "github.com/spf13/pflag"
)

type SrcDst struct {
	Src       string
	Dst       string
	Mode      string
	stageFile *os.File
	fileMode  os.FileMode
}

type StoreConfig struct {
	backends.StoreClient
	Name     string
	Onetime  bool
	Watch    bool
	Prefix   string
	Interval int
	Keys     []string
	store    *memkv.Store
}

// Resource is the representation of a parsed template resource.
type Resource struct {
	storeClients []StoreConfig
	ReloadCmd    string
	CheckCmd     string
	Uid          int
	Gid          int
	funcMap      map[string]interface{}
	store        memkv.Store
	syncOnly     bool
	sources      []*SrcDst
}

var ErrEmptySrc = errors.New("empty src template")

// NewResource creates a Resource.
func NewResource(storeClients []StoreConfig, sources []*SrcDst, reloadCmd, checkCmd string) (*Resource, error) {
	if len(storeClients) == 0 {
		return nil, errors.New("A valid StoreClient is required.")
	}

	for _, v := range sources {
		if v.Src == "" {
			return nil, ErrEmptySrc
		}
	}

	// TODO implement flags for these values
	syncOnly := false
	uid := os.Geteuid()
	gid := os.Getegid()

	tr := &Resource{
		storeClients: storeClients,
		ReloadCmd:    reloadCmd,
		CheckCmd:     checkCmd,
		store:        memkv.New(),
		funcMap:      newFuncMap(),
		Uid:          uid,
		Gid:          gid,
		syncOnly:     syncOnly,
		sources:      sources,
	}

	// initialize the inidividual backend memkv Stores
	for i := range tr.storeClients {
		store := memkv.New()
		tr.storeClients[i].store = &store
	}

	addFuncs(tr.funcMap, tr.store.FuncMap)

	return tr, nil
}

// NewResourceFromFlags creates a Resource from the provided commandline flags.
func NewResourceFromFlags(s backends.Store, flags *flag.FlagSet, watch bool) (*Resource, error) {
	src, _ := flags.GetString("src")
	dst, _ := flags.GetString("dst")
	keys, _ := flags.GetStringSlice("keys")
	fileMode, _ := flags.GetString("fileMode")
	prefix, _ := flags.GetString("prefix")
	reloadCmd, _ := flags.GetString("reload_cmd")
	checkCmd, _ := flags.GetString("check_cmd")
	onetime, _ := flags.GetBool("onetime")
	interval, _ := flags.GetInt("interval")

	sd := &SrcDst{
		Src:  src,
		Dst:  dst,
		Mode: fileMode,
	}

	b := StoreConfig{
		StoreClient: s.Client,
		Name:        s.Name,
		Onetime:     onetime,
		Prefix:      prefix,
		Watch:       watch,
		Interval:    interval,
		Keys:        keys,
	}

	return NewResource([]StoreConfig{b}, []*SrcDst{sd}, reloadCmd, checkCmd)
}

// setVars sets the Vars for template resource.
func (t *Resource) setVars(storeClient StoreConfig) error {
	var err error
	log.Debug("Retrieving keys from store: " + storeClient.Name)
	log.Debug("Key prefix set to " + storeClient.Prefix)

	result, err := storeClient.GetValues(appendPrefix(storeClient.Prefix, storeClient.Keys))
	if err != nil {
		return err
	}

	storeClient.store.Purge()

	for k, res := range result {
		storeClient.store.Set(path.Join("/", strings.TrimPrefix(k, storeClient.Prefix)), res)
	}

	//merge all stores
	t.store.Purge()
	for _, v := range t.storeClients {
		for _, kv := range v.store.GetAllKVs() {
			t.store.Set(kv.Key, kv.Value)
		}
	}

	return nil
}

// createStageFile stages the src configuration file by processing the src
// template and setting the desired owner, group, and mode. It also sets the
// StageFile for the template resource.
// It returns an error if any.
func (t *Resource) createStageFileAndSync() error {
	for _, s := range t.sources {
		log.Debug("Using source template " + s.Src)

		if !fileutil.IsFileExist(s.Src) {
			return errors.New("Missing template: " + s.Src)
		}

		log.Debug("Compiling source template(pongo2) " + s.Src)
		tmpl, err := pongo2.FromFile(s.Src)
		if err != nil {
			return fmt.Errorf("Unable to process template %s, %s", s.Src, err)
		}

		// create TempFile in Dest directory to avoid cross-filesystem issues
		temp, err := ioutil.TempFile(filepath.Dir(s.Dst), "."+filepath.Base(s.Dst))
		if err != nil {
			return err
		}

		if err = tmpl.ExecuteWriter(t.funcMap, temp); err != nil {
			temp.Close()
			os.Remove(temp.Name())
			return err
		}

		defer temp.Close()

		// Set the owner, group, and mode on the stage file now to make it easier to
		// compare against the destination configuration file later.
		os.Chmod(temp.Name(), s.fileMode)
		os.Chown(temp.Name(), t.Uid, t.Gid)
		s.stageFile = temp

		if err = t.sync(s); err != nil {
			return err
		}
	}
	return nil
}

// sync compares the staged and dest config files and attempts to sync them
// if they differ. sync will run a config check command if set before
// overwriting the target config file. Finally, sync will run a reload command
// if set to have the application or service pick up the changes.
// It returns an error if any.
func (t *Resource) sync(s *SrcDst) error {
	staged := s.stageFile.Name()
	defer os.Remove(staged)

	log.Debug("Comparing candidate config to " + s.Dst)
	ok, err := fileutil.SameFile(staged, s.Dst)
	if err != nil {
		log.Error(err.Error())
	}

	if !ok {
		log.Info("Target config " + s.Dst + " out of sync")
		if !t.syncOnly && t.CheckCmd != "" {
			if err := t.check(staged); err != nil {
				return errors.New("Config check failed: " + err.Error())
			}
		}
		log.Debug("Overwriting target config " + s.Dst)
		err := os.Rename(staged, s.Dst)
		if err != nil {
			if strings.Contains(err.Error(), "device or resource busy") {
				log.Debug("Rename failed - target is likely a mount. Trying to write instead")
				// try to open the file and write to it
				var contents []byte
				var rerr error
				contents, rerr = ioutil.ReadFile(staged)
				if rerr != nil {
					return rerr
				}
				err := ioutil.WriteFile(s.Dst, contents, s.fileMode)
				// make sure owner and group match the temp file, in case the file was created with WriteFile
				os.Chown(s.Dst, t.Uid, t.Gid)
				if err != nil {
					return err
				}
			} else {
				return err
			}
		}
		if !t.syncOnly && t.ReloadCmd != "" {
			if err := t.reload(); err != nil {
				return err
			}
		}
		log.Info("Target config " + s.Dst + " has been updated")
	} else {
		log.Debug("Target config " + s.Dst + " in sync")
	}
	return nil
}

// setFileMode sets the FileMode.
func (t *Resource) setFileMode() error {
	for _, s := range t.sources {
		if s.Mode == "" {
			if !fileutil.IsFileExist(s.Dst) {
				s.fileMode = 0644
			} else {
				fi, err := os.Stat(s.Dst)
				if err != nil {
					return err
				}
				s.fileMode = fi.Mode()
			}
		} else {
			mode, err := strconv.ParseUint(s.Mode, 0, 32)
			if err != nil {
				return err
			}
			s.fileMode = os.FileMode(mode)
		}
	}
	return nil
}

// Process is a convenience function that wraps calls to the three main tasks
// required to keep local configuration files in sync. First we gather vars
// from the store, then we stage a candidate configuration file, and finally sync
// things up.
// It returns an error if any.
func (t *Resource) process(storeClient StoreConfig) error {
	if err := t.setFileMode(); err != nil {
		return err
	}
	if err := t.setVars(storeClient); err != nil {
		return err
	}
	if err := t.createStageFileAndSync(); err != nil {
		return err
	}
	return nil
}

// check executes the check command to validate the staged config file. The
// command is modified so that any references to src template are substituted
// with a string representing the full path of the staged file. This allows the
// check to be run on the staged file before overwriting the destination config
// file.
// It returns nil if the check command returns 0 and there are no other errors.
func (t *Resource) check(stageFile string) error {
	var cmdBuffer bytes.Buffer
	data := make(map[string]string)
	data["src"] = stageFile
	tmpl, err := template.New("checkcmd").Parse(t.CheckCmd)
	if err != nil {
		return err
	}
	if err := tmpl.Execute(&cmdBuffer, data); err != nil {
		return err
	}
	log.Debug("Running " + cmdBuffer.String())
	c := exec.Command("/bin/sh", "-c", cmdBuffer.String())
	output, err := c.CombinedOutput()
	if err != nil {
		log.Error(fmt.Sprintf("%q", string(output)))
		return err
	}
	log.Debug(fmt.Sprintf("%q", string(output)))
	return nil
}

// reload executes the reload command.
// It returns nil if the reload command returns 0.
func (t *Resource) reload() error {
	log.Debug("Running " + t.ReloadCmd)
	c := exec.Command("/bin/sh", "-c", t.ReloadCmd)
	output, err := c.CombinedOutput()
	if err != nil {
		log.Error(fmt.Sprintf("%q", string(output)))
		return err
	}
	log.Debug(fmt.Sprintf("%q", string(output)))
	return nil
}

func (s StoreConfig) watch(stopChan chan bool, processChan chan StoreConfig) {
	var lastIndex uint64
	keysPrefix := appendPrefix(s.Prefix, s.Keys)

	for {
		select {
		case <-stopChan:
			return
		default:
			index, err := s.WatchPrefix(s.Prefix, keysPrefix, lastIndex, stopChan)
			if err != nil {
				log.Error(err)
				// Prevent backend errors from consuming all resources.
				time.Sleep(time.Second * 2)
				continue
			}
			processChan <- s
			lastIndex = index
		}
	}
}

func (s StoreConfig) interval(stopChan chan bool, processChan chan StoreConfig) {
	if s.Onetime {
		processChan <- s
		return
	}
	for {
		select {
		case <-stopChan:
			return
		case <-time.After(time.Duration(s.Interval) * time.Second):
			processChan <- s
		}
	}
}

func (t *Resource) Monitor() {
	wg := &sync.WaitGroup{}
	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
	stopChan := make(chan bool)
	processChan := make(chan StoreConfig)

	defer close(processChan)

	//load all KV-pairs the first time
	if err := t.setFileMode(); err != nil {
		log.Error(err)
	}
	for _, storeClient := range t.storeClients {
		if err := t.setVars(storeClient); err != nil {
			log.Error(err)
		}
	}
	if err := t.createStageFileAndSync(); err != nil {
		log.Error(err)
	}

	for _, sc := range t.storeClients {
		wg.Add(1)
		if sc.Watch {
			go func(s StoreConfig) {
				defer wg.Done()
				//t.watch(s, stopChan, processChan)
				s.watch(stopChan, processChan)
			}(sc)
		} else {
			go func(s StoreConfig) {
				defer wg.Done()
				s.interval(stopChan, processChan)
			}(sc)
		}
	}

	go func() {
		// If there is no go routine left - quit
		wg.Wait()
		signalChan <- syscall.SIGINT
	}()

	for {
		select {
		case storeClient := <-processChan:
			if err := t.process(storeClient); err != nil {
				log.Error(err)
			}
		case s := <-signalChan:
			log.Info(fmt.Sprintf("Captured %v. Exiting...", s))
			close(stopChan)
			// drain processChan
			go func() {
				for range processChan {
				}
			}()
			wg.Wait()
			return
		}
	}
}
<commit_msg>run only once not twice in onetime mode<commit_after>package template

import (
	"bytes"
	"errors"
	"fmt"
	"html/template"
	"io/ioutil"
	"os"
	"os/exec"
	"os/signal"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/HeavyHorst/memkv"
	"github.com/HeavyHorst/remco/backends"
	"github.com/HeavyHorst/remco/template/fileutil"
	"github.com/cloudflare/cfssl/log"
	"github.com/flosch/pongo2"
	flag "github.com/spf13/pflag"
)

type SrcDst struct {
	Src       string
	Dst       string
	Mode      string
	stageFile *os.File
	fileMode  os.FileMode
}

type StoreConfig struct {
	backends.StoreClient
	Name     string
	Onetime  bool
	Watch    bool
	Prefix   string
	Interval int
	Keys     []string
	store    *memkv.Store
}

// Resource is the representation of a parsed template resource.
type Resource struct {
	storeClients []StoreConfig
	ReloadCmd    string
	CheckCmd     string
	Uid          int
	Gid          int
	funcMap      map[string]interface{}
	store        memkv.Store
	syncOnly     bool
	sources      []*SrcDst
}

var ErrEmptySrc = errors.New("empty src template")

// NewResource creates a Resource.
func NewResource(storeClients []StoreConfig, sources []*SrcDst, reloadCmd, checkCmd string) (*Resource, error) {
	if len(storeClients) == 0 {
		return nil, errors.New("A valid StoreClient is required.")
	}

	for _, v := range sources {
		if v.Src == "" {
			return nil, ErrEmptySrc
		}
	}

	// TODO implement flags for these values
	syncOnly := false
	uid := os.Geteuid()
	gid := os.Getegid()

	tr := &Resource{
		storeClients: storeClients,
		ReloadCmd:    reloadCmd,
		CheckCmd:     checkCmd,
		store:        memkv.New(),
		funcMap:      newFuncMap(),
		Uid:          uid,
		Gid:          gid,
		syncOnly:     syncOnly,
		sources:      sources,
	}

	// initialize the inidividual backend memkv Stores
	for i := range tr.storeClients {
		store := memkv.New()
		tr.storeClients[i].store = &store
	}

	addFuncs(tr.funcMap, tr.store.FuncMap)

	return tr, nil
}

// NewResourceFromFlags creates a Resource from the provided commandline flags.
func NewResourceFromFlags(s backends.Store, flags *flag.FlagSet, watch bool) (*Resource, error) {
	src, _ := flags.GetString("src")
	dst, _ := flags.GetString("dst")
	keys, _ := flags.GetStringSlice("keys")
	fileMode, _ := flags.GetString("fileMode")
	prefix, _ := flags.GetString("prefix")
	reloadCmd, _ := flags.GetString("reload_cmd")
	checkCmd, _ := flags.GetString("check_cmd")
	onetime, _ := flags.GetBool("onetime")
	interval, _ := flags.GetInt("interval")

	sd := &SrcDst{
		Src:  src,
		Dst:  dst,
		Mode: fileMode,
	}

	b := StoreConfig{
		StoreClient: s.Client,
		Name:        s.Name,
		Onetime:     onetime,
		Prefix:      prefix,
		Watch:       watch,
		Interval:    interval,
		Keys:        keys,
	}

	return NewResource([]StoreConfig{b}, []*SrcDst{sd}, reloadCmd, checkCmd)
}

// setVars sets the Vars for template resource.
func (t *Resource) setVars(storeClient StoreConfig) error {
	var err error
	log.Debug("Retrieving keys from store: " + storeClient.Name)
	log.Debug("Key prefix set to " + storeClient.Prefix)

	result, err := storeClient.GetValues(appendPrefix(storeClient.Prefix, storeClient.Keys))
	if err != nil {
		return err
	}

	storeClient.store.Purge()

	for k, res := range result {
		storeClient.store.Set(path.Join("/", strings.TrimPrefix(k, storeClient.Prefix)), res)
	}

	//merge all stores
	t.store.Purge()
	for _, v := range t.storeClients {
		for _, kv := range v.store.GetAllKVs() {
			t.store.Set(kv.Key, kv.Value)
		}
	}

	return nil
}

// createStageFile stages the src configuration file by processing the src
// template and setting the desired owner, group, and mode. It also sets the
// StageFile for the template resource.
// It returns an error if any.
func (t *Resource) createStageFileAndSync() error {
	for _, s := range t.sources {
		log.Debug("Using source template " + s.Src)

		if !fileutil.IsFileExist(s.Src) {
			return errors.New("Missing template: " + s.Src)
		}

		log.Debug("Compiling source template(pongo2) " + s.Src)
		tmpl, err := pongo2.FromFile(s.Src)
		if err != nil {
			return fmt.Errorf("Unable to process template %s, %s", s.Src, err)
		}

		// create TempFile in Dest directory to avoid cross-filesystem issues
		temp, err := ioutil.TempFile(filepath.Dir(s.Dst), "."+filepath.Base(s.Dst))
		if err != nil {
			return err
		}

		if err = tmpl.ExecuteWriter(t.funcMap, temp); err != nil {
			temp.Close()
			os.Remove(temp.Name())
			return err
		}

		defer temp.Close()

		// Set the owner, group, and mode on the stage file now to make it easier to
		// compare against the destination configuration file later.
		os.Chmod(temp.Name(), s.fileMode)
		os.Chown(temp.Name(), t.Uid, t.Gid)
		s.stageFile = temp

		if err = t.sync(s); err != nil {
			return err
		}
	}
	return nil
}

// sync compares the staged and dest config files and attempts to sync them
// if they differ. sync will run a config check command if set before
// overwriting the target config file. Finally, sync will run a reload command
// if set to have the application or service pick up the changes.
// It returns an error if any.
func (t *Resource) sync(s *SrcDst) error {
	staged := s.stageFile.Name()
	defer os.Remove(staged)

	log.Debug("Comparing candidate config to " + s.Dst)
	ok, err := fileutil.SameFile(staged, s.Dst)
	if err != nil {
		log.Error(err.Error())
	}

	if !ok {
		log.Info("Target config " + s.Dst + " out of sync")
		if !t.syncOnly && t.CheckCmd != "" {
			if err := t.check(staged); err != nil {
				return errors.New("Config check failed: " + err.Error())
			}
		}
		log.Debug("Overwriting target config " + s.Dst)
		err := os.Rename(staged, s.Dst)
		if err != nil {
			if strings.Contains(err.Error(), "device or resource busy") {
				log.Debug("Rename failed - target is likely a mount. Trying to write instead")
				// try to open the file and write to it
				var contents []byte
				var rerr error
				contents, rerr = ioutil.ReadFile(staged)
				if rerr != nil {
					return rerr
				}
				err := ioutil.WriteFile(s.Dst, contents, s.fileMode)
				// make sure owner and group match the temp file, in case the file was created with WriteFile
				os.Chown(s.Dst, t.Uid, t.Gid)
				if err != nil {
					return err
				}
			} else {
				return err
			}
		}
		if !t.syncOnly && t.ReloadCmd != "" {
			if err := t.reload(); err != nil {
				return err
			}
		}
		log.Info("Target config " + s.Dst + " has been updated")
	} else {
		log.Debug("Target config " + s.Dst + " in sync")
	}
	return nil
}

// setFileMode sets the FileMode.
func (t *Resource) setFileMode() error {
	for _, s := range t.sources {
		if s.Mode == "" {
			if !fileutil.IsFileExist(s.Dst) {
				s.fileMode = 0644
			} else {
				fi, err := os.Stat(s.Dst)
				if err != nil {
					return err
				}
				s.fileMode = fi.Mode()
			}
		} else {
			mode, err := strconv.ParseUint(s.Mode, 0, 32)
			if err != nil {
				return err
			}
			s.fileMode = os.FileMode(mode)
		}
	}
	return nil
}

// Process is a convenience function that wraps calls to the three main tasks
// required to keep local configuration files in sync. First we gather vars
// from the store, then we stage a candidate configuration file, and finally sync
// things up.
// It returns an error if any.
func (t *Resource) process(storeClient StoreConfig) error {
	if err := t.setFileMode(); err != nil {
		return err
	}
	if err := t.setVars(storeClient); err != nil {
		return err
	}
	if err := t.createStageFileAndSync(); err != nil {
		return err
	}
	return nil
}

// check executes the check command to validate the staged config file. The
// command is modified so that any references to src template are substituted
// with a string representing the full path of the staged file. This allows the
// check to be run on the staged file before overwriting the destination config
// file.
// It returns nil if the check command returns 0 and there are no other errors.
func (t *Resource) check(stageFile string) error {
	var cmdBuffer bytes.Buffer
	data := make(map[string]string)
	data["src"] = stageFile
	tmpl, err := template.New("checkcmd").Parse(t.CheckCmd)
	if err != nil {
		return err
	}
	if err := tmpl.Execute(&cmdBuffer, data); err != nil {
		return err
	}
	log.Debug("Running " + cmdBuffer.String())
	c := exec.Command("/bin/sh", "-c", cmdBuffer.String())
	output, err := c.CombinedOutput()
	if err != nil {
		log.Error(fmt.Sprintf("%q", string(output)))
		return err
	}
	log.Debug(fmt.Sprintf("%q", string(output)))
	return nil
}

// reload executes the reload command.
// It returns nil if the reload command returns 0.
func (t *Resource) reload() error {
	log.Debug("Running " + t.ReloadCmd)
	c := exec.Command("/bin/sh", "-c", t.ReloadCmd)
	output, err := c.CombinedOutput()
	if err != nil {
		log.Error(fmt.Sprintf("%q", string(output)))
		return err
	}
	log.Debug(fmt.Sprintf("%q", string(output)))
	return nil
}

func (s StoreConfig) watch(stopChan chan bool, processChan chan StoreConfig) {
	var lastIndex uint64
	keysPrefix := appendPrefix(s.Prefix, s.Keys)

	for {
		select {
		case <-stopChan:
			return
		default:
			index, err := s.WatchPrefix(s.Prefix, keysPrefix, lastIndex, stopChan)
			if err != nil {
				log.Error(err)
				// Prevent backend errors from consuming all resources.
				time.Sleep(time.Second * 2)
				continue
			}
			processChan <- s
			lastIndex = index
		}
	}
}

func (s StoreConfig) interval(stopChan chan bool, processChan chan StoreConfig) {
	if s.Onetime {
		return
	}
	for {
		select {
		case <-stopChan:
			return
		case <-time.After(time.Duration(s.Interval) * time.Second):
			processChan <- s
		}
	}
}

func (t *Resource) Monitor() {
	wg := &sync.WaitGroup{}
	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
	stopChan := make(chan bool)
	processChan := make(chan StoreConfig)

	defer close(processChan)

	//load all KV-pairs the first time
	if err := t.setFileMode(); err != nil {
		log.Error(err)
	}
	for _, storeClient := range t.storeClients {
		if err := t.setVars(storeClient); err != nil {
			log.Error(err)
		}
	}
	if err := t.createStageFileAndSync(); err != nil {
		log.Error(err)
	}

	for _, sc := range t.storeClients {
		wg.Add(1)
		if sc.Watch {
			go func(s StoreConfig) {
				defer wg.Done()
				//t.watch(s, stopChan, processChan)
				s.watch(stopChan, processChan)
			}(sc)
		} else {
			go func(s StoreConfig) {
				defer wg.Done()
				s.interval(stopChan, processChan)
			}(sc)
		}
	}

	go func() {
		// If there is no goroutine left - quit
		wg.Wait()
		signalChan <- syscall.SIGINT
	}()

	for {
		select {
		case storeClient := <-processChan:
			if err := t.process(storeClient); err != nil {
				log.Error(err)
			}
		case s := <-signalChan:
			log.Info(fmt.Sprintf("Captured %v. Exiting...", s))
			close(stopChan)
			// drain processChan
			go func() {
				for range processChan {
				}
			}()
			wg.Wait()
			return
		}
	}
}
<|endoftext|>