<commit_before>/**
 * (C) Copyright 2014, Deft Labs
 */

package deftlabsds

import (
	"fmt"
	"time"
	"labix.org/v2/mgo"
	"deftlabs.com/log"
	"deftlabs.com/kernel"
)

// Create a new Mongo component from a configuration path. The path passed must be in the following format.
//
// mongodb: {
//     configDb: {
//         mongoUrl: "mongodb://localhost:27017/test",
//         safeMode: 0,
//         dialTimeoutInMs: 3000,
//         socketTimeoutInMs: 3000,
//         syncTimeoutInMs: 3000,
//         cursorTimeoutInMs: 30000,
//     }
// }
//
// The configPath for this component would be "mongodb.configDb". The path can be any arbitrary set of nested
// json documents (json path). If the path is incorrect, the Start() method will panic when called by the kernel.
//
// All of the params above must be present or
// the Start method will panic. If the componentId or configPath param is nil or empty,
// this method will panic.
func NewMongoFromConfigPath(componentId, configPath string) *Mongo {

	if len(componentId) == 0 {
		panic("When calling NewMongoFromConfigPath you must pass in a non-empty componentId param")
	}

	if len(configPath) == 0 {
		panic("When calling NewMongoFromConfigPath you must pass in a non-empty configPath param")
	}

	return &Mongo{ componentId : componentId, configPath : configPath }
}

// Create a new Mongo component. This method will panic if either of the params are nil or len == 0.
func NewMongo(componentId, mongoUrl string, safeMode, dialTimeoutInMs, socketTimeoutInMs, syncTimeoutInMs, cursorTimeoutInMs int) *Mongo {

	if len(componentId) == 0 {
		panic("When calling NewMongo you must pass in a non-empty component id")
	}

	if len(mongoUrl) == 0 {
		panic("When calling NewMongo you must pass in a non-empty Mongo url")
	}

	return &Mongo{
		componentId : componentId,
		mongoUrl: mongoUrl,
		safeMode : safeMode,
		dialTimeoutInMs : dialTimeoutInMs,
		socketTimeoutInMs : socketTimeoutInMs,
		syncTimeoutInMs : syncTimeoutInMs,
		cursorTimeoutInMs : cursorTimeoutInMs,
	}
}

type Mongo struct {
	kernel *deftlabskernel.Kernel
	slogger.Logger

	configPath string

	componentId string
	mongoUrl string
	safeMode int
	dialTimeoutInMs int
	socketTimeoutInMs int
	syncTimeoutInMs int
	cursorTimeoutInMs int
	session *mgo.Session
}

// Returns the collection from the session.
func (self *Mongo) Collection(dbName, collectionName string) *mgo.Collection { return self.Db(dbName).C(collectionName) }

// Returns the database from the session.
func (self *Mongo) Db(name string) *mgo.Database { return self.session.DB(name) }

// Returns the session struct.
func (self *Mongo) Session() *mgo.Session { return self.session }

// Returns a clone of the session struct.
func (self *Mongo) SessionClone() *mgo.Session { return self.session.Clone() }

// Returns a copy of the session struct.
func (self *Mongo) SessionCopy() *mgo.Session { return self.session.Clone() }

func (self *Mongo) Start(kernel *deftlabskernel.Kernel) error {
	self.kernel = kernel
	self.Logger = kernel.Logger

	var err error

	// This is a configuration based creation. Load the config data first.
	if len(self.configPath) > 0 {
		self.mongoUrl = self.kernel.Configuration.String(fmt.Sprintf("%s.%s", self.configPath, "mongoUrl"), "")
		self.safeMode = self.kernel.Configuration.Int(fmt.Sprintf("%s.%s", self.configPath, "safeMode"), -1)
		self.dialTimeoutInMs = self.kernel.Configuration.Int(fmt.Sprintf("%s.%s", self.configPath, "dialTimeoutInMs"), -1)
		self.socketTimeoutInMs = self.kernel.Configuration.Int(fmt.Sprintf("%s.%s", self.configPath, "socketTimeoutInMs"), -1)
		self.syncTimeoutInMs = self.kernel.Configuration.Int(fmt.Sprintf("%s.%s", self.configPath, "syncTimeoutInMs"), -1)
		self.cursorTimeoutInMs = self.kernel.Configuration.Int(fmt.Sprintf("%s.%s", self.configPath, "cursorTimeoutInMs"), -1)
	}

	// Validate the params
	if len(self.mongoUrl) == 0 {
		panic(fmt.Sprintf("In Mongo - mongoUrl is not set - componentId: %s", self.componentId))
	}

	if self.dialTimeoutInMs < 0 {
		panic(fmt.Sprintf("In Mongo - dialTimeoutInMs is invalid - value: %d - componentId: %s", self.dialTimeoutInMs, self.componentId))
	}

	if self.socketTimeoutInMs < 0 {
		panic(fmt.Sprintf("In Mongo - socketTimeoutInMs is invalid - value: %d - componentId: %s", self.socketTimeoutInMs, self.componentId))
	}

	if self.syncTimeoutInMs < 0 {
		panic(fmt.Sprintf("In Mongo - syncTimeoutInMs is invalid - value: %d - componentId: %s", self.syncTimeoutInMs, self.componentId))
	}

	if self.cursorTimeoutInMs < 0 {
		panic(fmt.Sprintf("In Mongo - cursorTimeoutInMs is invalid - value: %d - componentId: %s", self.cursorTimeoutInMs, self.componentId))
	}

	if self.safeMode < 0  || self.safeMode > 2 {
		panic(fmt.Sprintf("In Mongo - safeMode is invalid - value: %d - componentId: %s", self.safeMode, self.componentId))
	}

	// Create the session.
	if self.session, err = mgo.DialWithTimeout(self.mongoUrl, time.Duration(self.dialTimeoutInMs) * time.Millisecond); err != nil {
		return slogger.NewStackError("Unable to init Mongo session - component: %s - mongodbUrl: %s", self.componentId, self.mongoUrl)
	}

	// This is annoying, but mgo defines these constants as the restricted "mode" type.
	switch self.safeMode {
		case 0: self.session.SetMode(mgo.Eventual, true)
		case 1: self.session.SetMode(mgo.Monotonic, true)
		case 2: self.session.SetMode(mgo.Strong, true)
	}

	self.session.SetSocketTimeout(time.Duration(self.socketTimeoutInMs) * time.Millisecond)
	self.session.SetSyncTimeout(time.Duration(self.syncTimeoutInMs) * time.Millisecond)

	return nil
}

// Stop the component. This will close the base session.
func (self *Mongo) Stop(kernel *deftlabskernel.Kernel) error {

	if self.session != nil {
		self.session.Close()
	}

	return nil
}

<commit_msg>changed safeMode to mode - different param<commit_after>/**
 * (C) Copyright 2014, Deft Labs
 */

package deftlabsds

import (
	"fmt"
	"time"
	"labix.org/v2/mgo"
	"deftlabs.com/log"
	"deftlabs.com/kernel"
)

// Create a new Mongo component from a configuration path. The path passed must be in the following format.
//
// mongodb: {
//     configDb: {
//         mongoUrl: "mongodb://localhost:27017/test",
//         mode: 0,
//         dialTimeoutInMs: 3000,
//         socketTimeoutInMs: 3000,
//         syncTimeoutInMs: 3000,
//         cursorTimeoutInMs: 30000,
//     }
// }
//
// The configPath for this component would be "mongodb.configDb". The path can be any arbitrary set of nested
// json documents (json path). If the path is incorrect, the Start() method will panic when called by the kernel.
//
// All of the params above must be present or
// the Start method will panic. If the componentId or configPath param is nil or empty,
// this method will panic.
func NewMongoFromConfigPath(componentId, configPath string) *Mongo {

	if len(componentId) == 0 {
		panic("When calling NewMongoFromConfigPath you must pass in a non-empty componentId param")
	}

	if len(configPath) == 0 {
		panic("When calling NewMongoFromConfigPath you must pass in a non-empty configPath param")
	}

	return &Mongo{ componentId : componentId, configPath : configPath }
}

// Create a new Mongo component. This method will panic if either of the params are nil or len == 0.
func NewMongo(componentId, mongoUrl string, mode, dialTimeoutInMs, socketTimeoutInMs, syncTimeoutInMs, cursorTimeoutInMs int) *Mongo {

	if len(componentId) == 0 {
		panic("When calling NewMongo you must pass in a non-empty component id")
	}

	if len(mongoUrl) == 0 {
		panic("When calling NewMongo you must pass in a non-empty Mongo url")
	}

	return &Mongo{
		componentId : componentId,
		mongoUrl: mongoUrl,
		mode : mode,
		dialTimeoutInMs : dialTimeoutInMs,
		socketTimeoutInMs : socketTimeoutInMs,
		syncTimeoutInMs : syncTimeoutInMs,
		cursorTimeoutInMs : cursorTimeoutInMs,
	}
}

type Mongo struct {
	kernel *deftlabskernel.Kernel
	slogger.Logger

	configPath string

	componentId string
	mongoUrl string
	mode int
	dialTimeoutInMs int
	socketTimeoutInMs int
	syncTimeoutInMs int
	cursorTimeoutInMs int
	session *mgo.Session
}

// Returns the collection from the session.
func (self *Mongo) Collection(dbName, collectionName string) *mgo.Collection { return self.Db(dbName).C(collectionName) }

// Returns the database from the session.
func (self *Mongo) Db(name string) *mgo.Database { return self.session.DB(name) }

// Returns the session struct.
func (self *Mongo) Session() *mgo.Session { return self.session }

// Returns a clone of the session struct.
func (self *Mongo) SessionClone() *mgo.Session { return self.session.Clone() }

// Returns a copy of the session struct.
func (self *Mongo) SessionCopy() *mgo.Session { return self.session.Clone() }

func (self *Mongo) Start(kernel *deftlabskernel.Kernel) error {
	self.kernel = kernel
	self.Logger = kernel.Logger

	var err error

	// This is a configuration based creation. Load the config data first.
	if len(self.configPath) > 0 {
		self.mongoUrl = self.kernel.Configuration.String(fmt.Sprintf("%s.%s", self.configPath, "mongoUrl"), "")
		self.mode = self.kernel.Configuration.Int(fmt.Sprintf("%s.%s", self.configPath, "mode"), -1)
		self.dialTimeoutInMs = self.kernel.Configuration.Int(fmt.Sprintf("%s.%s", self.configPath, "dialTimeoutInMs"), -1)
		self.socketTimeoutInMs = self.kernel.Configuration.Int(fmt.Sprintf("%s.%s", self.configPath, "socketTimeoutInMs"), -1)
		self.syncTimeoutInMs = self.kernel.Configuration.Int(fmt.Sprintf("%s.%s", self.configPath, "syncTimeoutInMs"), -1)
		self.cursorTimeoutInMs = self.kernel.Configuration.Int(fmt.Sprintf("%s.%s", self.configPath, "cursorTimeoutInMs"), -1)
	}

	// Validate the params
	if len(self.mongoUrl) == 0 {
		panic(fmt.Sprintf("In Mongo - mongoUrl is not set - componentId: %s", self.componentId))
	}

	if self.dialTimeoutInMs < 0 {
		panic(fmt.Sprintf("In Mongo - dialTimeoutInMs is invalid - value: %d - componentId: %s", self.dialTimeoutInMs, self.componentId))
	}

	if self.socketTimeoutInMs < 0 {
		panic(fmt.Sprintf("In Mongo - socketTimeoutInMs is invalid - value: %d - componentId: %s", self.socketTimeoutInMs, self.componentId))
	}

	if self.syncTimeoutInMs < 0 {
		panic(fmt.Sprintf("In Mongo - syncTimeoutInMs is invalid - value: %d - componentId: %s", self.syncTimeoutInMs, self.componentId))
	}

	if self.cursorTimeoutInMs < 0 {
		panic(fmt.Sprintf("In Mongo - cursorTimeoutInMs is invalid - value: %d - componentId: %s", self.cursorTimeoutInMs, self.componentId))
	}

	if self.mode < 0  || self.mode > 2 {
		panic(fmt.Sprintf("In Mongo - mode is invalid - value: %d - componentId: %s", self.mode, self.componentId))
	}

	// Create the session.
	if self.session, err = mgo.DialWithTimeout(self.mongoUrl, time.Duration(self.dialTimeoutInMs) * time.Millisecond); err != nil {
		return slogger.NewStackError("Unable to init Mongo session - component: %s - mongodbUrl: %s", self.componentId, self.mongoUrl)
	}

	// This is annoying, but mgo defines these constants as the restricted "mode" type.
	switch self.mode {
		case 0: self.session.SetMode(mgo.Eventual, true)
		case 1: self.session.SetMode(mgo.Monotonic, true)
		case 2: self.session.SetMode(mgo.Strong, true)
	}

	self.session.SetSocketTimeout(time.Duration(self.socketTimeoutInMs) * time.Millisecond)
	self.session.SetSyncTimeout(time.Duration(self.syncTimeoutInMs) * time.Millisecond)

	return nil
}

// Stop the component. This will close the base session.
func (self *Mongo) Stop(kernel *deftlabskernel.Kernel) error {

	if self.session != nil {
		self.session.Close()
	}

	return nil
}

<|endoftext|>


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

// Package csv reads and writes comma-separated values (CSV) files.
// There are many kinds of CSV files; this package supports the format
// described in RFC 4180.
//
// A csv file contains zero or more records of one or more fields per record.
// Each record is separated by the newline character. The final record may
// optionally be followed by a newline character.
//
//	field1,field2,field3
//
// White space is considered part of a field.
//
// Carriage returns before newline characters are silently removed.
//
// Blank lines are ignored. A line with only whitespace characters (excluding
// the ending newline character) is not considered a blank line.
//
// Fields which start and stop with the quote character " are called
// quoted-fields. The beginning and ending quote are not part of the
// field.
//
// The source:
//
//	normal string,"quoted-field"
//
// results in the fields
//
//	{`normal string`, `quoted-field`}
//
// Within a quoted-field a quote character followed by a second quote
// character is considered a single quote.
//
//	"the ""word"" is true","a ""quoted-field"""
//
// results in
//
//	{`the "word" is true`, `a "quoted-field"`}
//
// Newlines and commas may be included in a quoted-field
//
//	"Multi-line
//	field","comma is ,"
//
// results in
//
//	{`Multi-line
//	field`, `comma is ,`}
package csv

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io"
	"unicode"
)

// A ParseError is returned for parsing errors.
// The first line is 1.  The first column is 0.
type ParseError struct {
	Line   int   // Line where the error occurred
	Column int   // Column (rune index) where the error occurred
	Err    error // The actual error
}

func (e *ParseError) Error() string {
	return fmt.Sprintf("line %d, column %d: %s", e.Line, e.Column, e.Err)
}

// These are the errors that can be returned in ParseError.Error
var (
	ErrTrailingComma = errors.New("extra delimiter at end of line") // no longer used
	ErrBareQuote     = errors.New("bare \" in non-quoted-field")
	ErrQuote         = errors.New("extraneous \" in field")
	ErrFieldCount    = errors.New("wrong number of fields in line")
)

// A Reader reads records from a CSV-encoded file.
//
// As returned by NewReader, a Reader expects input conforming to RFC 4180.
// The exported fields can be changed to customize the details before the
// first call to Read or ReadAll.
//
//
type Reader struct {
	// Comma is the field delimiter.
	// It is set to comma (',') by NewReader.
	Comma rune
	// Comment, if not 0, is the comment character. Lines beginning with the
	// Comment character without preceding whitespace are ignored.
	// With leading whitespace the Comment character becomes part of the
	// field, even if TrimLeadingSpace is true.
	Comment rune
	// FieldsPerRecord is the number of expected fields per record.
	// If FieldsPerRecord is positive, Read requires each record to
	// have the given number of fields. If FieldsPerRecord is 0, Read sets it to
	// the number of fields in the first record, so that future records must
	// have the same field count. If FieldsPerRecord is negative, no check is
	// made and records may have a variable number of fields.
	FieldsPerRecord int
	// If LazyQuotes is true, a quote may appear in an unquoted field and a
	// non-doubled quote may appear in a quoted field.
	LazyQuotes    bool
	TrailingComma bool // ignored; here for backwards compatibility
	// If TrimLeadingSpace is true, leading white space in a field is ignored.
	// This is done even if the field delimiter, Comma, is white space.
	TrimLeadingSpace bool

	line   int
	column int
	r      *bufio.Reader
	field  bytes.Buffer
}

// NewReader returns a new Reader that reads from r.
func NewReader(r io.Reader) *Reader {
	return &Reader{
		Comma: ',',
		r:     bufio.NewReader(r),
	}
}

// error creates a new ParseError based on err.
func (r *Reader) error(err error) error {
	return &ParseError{
		Line:   r.line,
		Column: r.column,
		Err:    err,
	}
}

// Read reads one record from r. The record is a slice of strings with each
// string representing one field.
func (r *Reader) Read() (record []string, err error) {
	for {
		record, err = r.parseRecord()
		if record != nil {
			break
		}
		if err != nil {
			return nil, err
		}
	}

	if r.FieldsPerRecord > 0 {
		if len(record) != r.FieldsPerRecord {
			r.column = 0 // report at start of record
			return record, r.error(ErrFieldCount)
		}
	} else if r.FieldsPerRecord == 0 {
		r.FieldsPerRecord = len(record)
	}
	return record, nil
}

// ReadAll reads all the remaining records from r.
// Each record is a slice of fields.
// A successful call returns err == nil, not err == io.EOF. Because ReadAll is
// defined to read until EOF, it does not treat end of file as an error to be
// reported.
func (r *Reader) ReadAll() (records [][]string, err error) {
	for {
		record, err := r.Read()
		if err == io.EOF {
			return records, nil
		}
		if err != nil {
			return nil, err
		}
		records = append(records, record)
	}
}

// readRune reads one rune from r, folding \r\n to \n and keeping track
// of how far into the line we have read.  r.column will point to the start
// of this rune, not the end of this rune.
func (r *Reader) readRune() (rune, error) {
	r1, _, err := r.r.ReadRune()

	// Handle \r\n here. We make the simplifying assumption that
	// anytime \r is followed by \n that it can be folded to \n.
	// We will not detect files which contain both \r\n and bare \n.
	if r1 == '\r' {
		r1, _, err = r.r.ReadRune()
		if err == nil {
			if r1 != '\n' {
				r.r.UnreadRune()
				r1 = '\r'
			}
		}
	}
	r.column++
	return r1, err
}

// skip reads runes up to and including the rune delim or until error.
func (r *Reader) skip(delim rune) error {
	for {
		r1, err := r.readRune()
		if err != nil {
			return err
		}
		if r1 == delim {
			return nil
		}
	}
}

// parseRecord reads and parses a single csv record from r.
func (r *Reader) parseRecord() (fields []string, err error) {
	// Each record starts on a new line. We increment our line
	// number (lines start at 1, not 0) and set column to -1
	// so as we increment in readRune it points to the character we read.
	r.line++
	r.column = -1

	// Peek at the first rune. If it is an error we are done.
	// If we support comments and it is the comment character
	// then skip to the end of line.

	r1, _, err := r.r.ReadRune()
	if err != nil {
		return nil, err
	}

	if r.Comment != 0 && r1 == r.Comment {
		return nil, r.skip('\n')
	}
	r.r.UnreadRune()

	// At this point we have at least one field.
	for {
		haveField, delim, err := r.parseField()
		if haveField {
			// If FieldsPerRecord is greater than 0 we can assume the final
			// length of fields to be equal to FieldsPerRecord.
			if r.FieldsPerRecord > 0 && fields == nil {
				fields = make([]string, 0, r.FieldsPerRecord)
			}
			fields = append(fields, r.field.String())
		}
		if delim == '\n' || err == io.EOF {
			return fields, err
		} else if err != nil {
			return nil, err
		}
	}
}

// parseField parses the next field in the record. The read field is
// located in r.field. Delim is the first character not part of the field
// (r.Comma or '\n').
func (r *Reader) parseField() (haveField bool, delim rune, err error) {
	r.field.Reset()

	r1, err := r.readRune()
	for err == nil && r.TrimLeadingSpace && r1 != '\n' && unicode.IsSpace(r1) {
		r1, err = r.readRune()
	}

	if err == io.EOF && r.column != 0 {
		return true, 0, err
	}
	if err != nil {
		return false, 0, err
	}

	switch r1 {
	case r.Comma:
		// will check below

	case '\n':
		// We are a trailing empty field or a blank line
		if r.column == 0 {
			return false, r1, nil
		}
		return true, r1, nil

	case '"':
		// quoted field
	Quoted:
		for {
			r1, err = r.readRune()
			if err != nil {
				if err == io.EOF {
					if r.LazyQuotes {
						return true, 0, err
					}
					return false, 0, r.error(ErrQuote)
				}
				return false, 0, err
			}
			switch r1 {
			case '"':
				r1, err = r.readRune()
				if err != nil || r1 == r.Comma {
					break Quoted
				}
				if r1 == '\n' {
					return true, r1, nil
				}
				if r1 != '"' {
					if !r.LazyQuotes {
						r.column--
						return false, 0, r.error(ErrQuote)
					}
					// accept the bare quote
					r.field.WriteRune('"')
				}
			case '\n':
				r.line++
				r.column = -1
			}
			r.field.WriteRune(r1)
		}

	default:
		// unquoted field
		for {
			r.field.WriteRune(r1)
			r1, err = r.readRune()
			if err != nil || r1 == r.Comma {
				break
			}
			if r1 == '\n' {
				return true, r1, nil
			}
			if !r.LazyQuotes && r1 == '"' {
				return false, 0, r.error(ErrBareQuote)
			}
		}
	}

	if err != nil {
		if err == io.EOF {
			return true, 0, err
		}
		return false, 0, err
	}

	return true, r1, nil
}
<commit_msg>encoding/csv: avoid allocations when reading records<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package csv reads and writes comma-separated values (CSV) files.
// There are many kinds of CSV files; this package supports the format
// described in RFC 4180.
//
// A csv file contains zero or more records of one or more fields per record.
// Each record is separated by the newline character. The final record may
// optionally be followed by a newline character.
//
//	field1,field2,field3
//
// White space is considered part of a field.
//
// Carriage returns before newline characters are silently removed.
//
// Blank lines are ignored. A line with only whitespace characters (excluding
// the ending newline character) is not considered a blank line.
//
// Fields which start and stop with the quote character " are called
// quoted-fields. The beginning and ending quote are not part of the
// field.
//
// The source:
//
//	normal string,"quoted-field"
//
// results in the fields
//
//	{`normal string`, `quoted-field`}
//
// Within a quoted-field a quote character followed by a second quote
// character is considered a single quote.
//
//	"the ""word"" is true","a ""quoted-field"""
//
// results in
//
//	{`the "word" is true`, `a "quoted-field"`}
//
// Newlines and commas may be included in a quoted-field
//
//	"Multi-line
//	field","comma is ,"
//
// results in
//
//	{`Multi-line
//	field`, `comma is ,`}
package csv

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io"
	"unicode"
)

// A ParseError is returned for parsing errors.
// The first line is 1.  The first column is 0.
type ParseError struct {
	Line   int   // Line where the error occurred
	Column int   // Column (rune index) where the error occurred
	Err    error // The actual error
}

func (e *ParseError) Error() string {
	return fmt.Sprintf("line %d, column %d: %s", e.Line, e.Column, e.Err)
}

// These are the errors that can be returned in ParseError.Error
var (
	ErrTrailingComma = errors.New("extra delimiter at end of line") // no longer used
	ErrBareQuote     = errors.New("bare \" in non-quoted-field")
	ErrQuote         = errors.New("extraneous \" in field")
	ErrFieldCount    = errors.New("wrong number of fields in line")
)

// A Reader reads records from a CSV-encoded file.
//
// As returned by NewReader, a Reader expects input conforming to RFC 4180.
// The exported fields can be changed to customize the details before the
// first call to Read or ReadAll.
//
//
type Reader struct {
	// Comma is the field delimiter.
	// It is set to comma (',') by NewReader.
	Comma rune
	// Comment, if not 0, is the comment character. Lines beginning with the
	// Comment character without preceding whitespace are ignored.
	// With leading whitespace the Comment character becomes part of the
	// field, even if TrimLeadingSpace is true.
	Comment rune
	// FieldsPerRecord is the number of expected fields per record.
	// If FieldsPerRecord is positive, Read requires each record to
	// have the given number of fields. If FieldsPerRecord is 0, Read sets it to
	// the number of fields in the first record, so that future records must
	// have the same field count. If FieldsPerRecord is negative, no check is
	// made and records may have a variable number of fields.
	FieldsPerRecord int
	// If LazyQuotes is true, a quote may appear in an unquoted field and a
	// non-doubled quote may appear in a quoted field.
	LazyQuotes    bool
	TrailingComma bool // ignored; here for backwards compatibility
	// If TrimLeadingSpace is true, leading white space in a field is ignored.
	// This is done even if the field delimiter, Comma, is white space.
	TrimLeadingSpace bool

	line   int
	column int
	r      *bufio.Reader
	// lineBuffer holds the unescaped fields read by readField, one after another.
	// The fields can be accessed by using the indexes in fieldIndexes.
	// Example: for the row `a,"b","c""d",e` lineBuffer will contain `abc"de` and
	// fieldIndexes will contain the indexes 0, 1, 2, 5.
	lineBuffer bytes.Buffer
	// Indexes of fields inside lineBuffer
	// The i'th field starts at offset fieldIndexes[i] in lineBuffer.
	fieldIndexes []int
}

// NewReader returns a new Reader that reads from r.
func NewReader(r io.Reader) *Reader {
	return &Reader{
		Comma: ',',
		r:     bufio.NewReader(r),
	}
}

// error creates a new ParseError based on err.
func (r *Reader) error(err error) error {
	return &ParseError{
		Line:   r.line,
		Column: r.column,
		Err:    err,
	}
}

// Read reads one record from r. The record is a slice of strings with each
// string representing one field.
func (r *Reader) Read() (record []string, err error) {
	for {
		record, err = r.parseRecord()
		if record != nil {
			break
		}
		if err != nil {
			return nil, err
		}
	}

	if r.FieldsPerRecord > 0 {
		if len(record) != r.FieldsPerRecord {
			r.column = 0 // report at start of record
			return record, r.error(ErrFieldCount)
		}
	} else if r.FieldsPerRecord == 0 {
		r.FieldsPerRecord = len(record)
	}
	return record, nil
}

// ReadAll reads all the remaining records from r.
// Each record is a slice of fields.
// A successful call returns err == nil, not err == io.EOF. Because ReadAll is
// defined to read until EOF, it does not treat end of file as an error to be
// reported.
func (r *Reader) ReadAll() (records [][]string, err error) {
	for {
		record, err := r.Read()
		if err == io.EOF {
			return records, nil
		}
		if err != nil {
			return nil, err
		}
		records = append(records, record)
	}
}

// readRune reads one rune from r, folding \r\n to \n and keeping track
// of how far into the line we have read.  r.column will point to the start
// of this rune, not the end of this rune.
func (r *Reader) readRune() (rune, error) {
	r1, _, err := r.r.ReadRune()

	// Handle \r\n here. We make the simplifying assumption that
	// anytime \r is followed by \n that it can be folded to \n.
	// We will not detect files which contain both \r\n and bare \n.
	if r1 == '\r' {
		r1, _, err = r.r.ReadRune()
		if err == nil {
			if r1 != '\n' {
				r.r.UnreadRune()
				r1 = '\r'
			}
		}
	}
	r.column++
	return r1, err
}

// skip reads runes up to and including the rune delim or until error.
func (r *Reader) skip(delim rune) error {
	for {
		r1, err := r.readRune()
		if err != nil {
			return err
		}
		if r1 == delim {
			return nil
		}
	}
}

// parseRecord reads and parses a single csv record from r.
func (r *Reader) parseRecord() (fields []string, err error) {
	// Each record starts on a new line. We increment our line
	// number (lines start at 1, not 0) and set column to -1
	// so as we increment in readRune it points to the character we read.
	r.line++
	r.column = -1

	// Peek at the first rune. If it is an error we are done.
	// If we support comments and it is the comment character
	// then skip to the end of line.

	r1, _, err := r.r.ReadRune()
	if err != nil {
		return nil, err
	}

	if r.Comment != 0 && r1 == r.Comment {
		return nil, r.skip('\n')
	}
	r.r.UnreadRune()

	r.lineBuffer.Reset()
	r.fieldIndexes = r.fieldIndexes[:0]

	// At this point we have at least one field.
	for {
		idx := r.lineBuffer.Len()

		haveField, delim, err := r.parseField()
		if haveField {
			r.fieldIndexes = append(r.fieldIndexes, idx)
		}

		if delim == '\n' || err == io.EOF {
			if len(r.fieldIndexes) == 0 {
				return nil, err
			}
			break
		}

		if err != nil {
			return nil, err
		}
	}

	fieldCount := len(r.fieldIndexes)
	// Using this approach (creating a single string and taking slices of it)
	// means that a single reference to any of the fields will retain the whole
	// string. The risk of a nontrivial space leak caused by this is considered
	// minimal and a tradeoff for better performance through the combined
	// allocations.
	line := r.lineBuffer.String()
	fields = make([]string, fieldCount)

	for i, idx := range r.fieldIndexes {
		if i == fieldCount-1 {
			fields[i] = line[idx:]
		} else {
			fields[i] = line[idx:r.fieldIndexes[i+1]]
		}
	}

	return fields, nil
}

// parseField parses the next field in the record. The read field is
// appended to r.lineBuffer. Delim is the first character not part of the field
// (r.Comma or '\n').
func (r *Reader) parseField() (haveField bool, delim rune, err error) {
	r1, err := r.readRune()
	for err == nil && r.TrimLeadingSpace && r1 != '\n' && unicode.IsSpace(r1) {
		r1, err = r.readRune()
	}

	if err == io.EOF && r.column != 0 {
		return true, 0, err
	}
	if err != nil {
		return false, 0, err
	}

	switch r1 {
	case r.Comma:
		// will check below

	case '\n':
		// We are a trailing empty field or a blank line
		if r.column == 0 {
			return false, r1, nil
		}
		return true, r1, nil

	case '"':
		// quoted field
	Quoted:
		for {
			r1, err = r.readRune()
			if err != nil {
				if err == io.EOF {
					if r.LazyQuotes {
						return true, 0, err
					}
					return false, 0, r.error(ErrQuote)
				}
				return false, 0, err
			}
			switch r1 {
			case '"':
				r1, err = r.readRune()
				if err != nil || r1 == r.Comma {
					break Quoted
				}
				if r1 == '\n' {
					return true, r1, nil
				}
				if r1 != '"' {
					if !r.LazyQuotes {
						r.column--
						return false, 0, r.error(ErrQuote)
					}
					// accept the bare quote
					r.lineBuffer.WriteRune('"')
				}
			case '\n':
				r.line++
				r.column = -1
			}
			r.lineBuffer.WriteRune(r1)
		}

	default:
		// unquoted field
		for {
			r.lineBuffer.WriteRune(r1)
			r1, err = r.readRune()
			if err != nil || r1 == r.Comma {
				break
			}
			if r1 == '\n' {
				return true, r1, nil
			}
			if !r.LazyQuotes && r1 == '"' {
				return false, 0, r.error(ErrBareQuote)
			}
		}
	}

	if err != nil {
		if err == io.EOF {
			return true, 0, err
		}
		return false, 0, err
	}

	return true, r1, nil
}
<|endoftext|>


<commit_before>package spotify

import (
	"github.com/fabiofalci/sconsify/infrastructure"
	"github.com/fabiofalci/sconsify/sconsify"
	sp "github.com/op/go-libspotify/spotify"
	webspotify "github.com/zmb3/spotify"
	"strings"
	"time"
)

func (spotify *Spotify) shutdownSpotify() {
	spotify.session.Logout()
	spotify.initCache()
	spotify.events.ShutdownEngine()
}

func (spotify *Spotify) play(trackUri *sconsify.Track) {

	player := spotify.session.Player()
	if !spotify.paused || spotify.currentTrack != trackUri {
		link, err := spotify.session.ParseLink(trackUri.URI)
		if err != nil {
			return
		}

		track, err := link.Track()
		if err != nil {
			return
		}

		if trackUri.IsPartial() {
			trackUri = sconsify.ToSconsifyTrack(track)
		}

		if !spotify.isTrackAvailable(track) {
			if trackUri.IsFromWebApi() {
				retry := trackUri.RetryLoading()
				if retry < 4 {
					go func() {
						time.Sleep(100 * time.Millisecond)
						spotify.events.Play(trackUri)
					}()
					return
				}
			}
			spotify.events.TrackNotAvailable(trackUri)
			return
		}
		if err := player.Load(track); err != nil {
			return
		}
		spotify.events.NewTrackLoaded(track.Duration())
	}
	player.Play()

	spotify.events.TrackPlaying(trackUri)
	spotify.currentTrack = trackUri
	spotify.paused = false
	return
}

func (spotify *Spotify) isTrackAvailable(track *sp.Track) bool {
	return track.Availability() == sp.TrackAvailabilityAvailable
}

func (spotify *Spotify) search(query string) {
	playlists := sconsify.InitPlaylists()

	query = checkAlias(query)
	name := " " + query

	playlist := sconsify.InitSearchPlaylist(name, name, func(playlist *sconsify.Playlist) {
		options := createWebSpotifyOptions(50, playlist.Tracks())
		if searchResult, err := spotify.getWebApiClient().SearchOpt(query, webspotify.SearchTypeTrack, options); err == nil {
			numberOfTracks := len(searchResult.Tracks.Tracks)
			infrastructure.Debugf("Search '%v' returned %v track(s)", query, numberOfTracks)
			for _, track := range searchResult.Tracks.Tracks {
				webArtist := track.Artists[0]
				artist := sconsify.InitArtist(string(webArtist.URI), webArtist.Name)
				playlist.AddTrack(sconsify.InitWebApiTrack(string(track.URI), artist, track.Name, track.TimeDuration().String()))
				infrastructure.Debugf("\tTrack '%v' (%v)", track.URI, track.Name)
			}
		} else {
			infrastructure.Debugf("Spotify search returning error: %v", err)
		}
	})
	playlist.ExecuteLoad()
	playlists.AddPlaylist(playlist)

	spotify.events.NewPlaylist(playlists)
}

func checkAlias(query string) string {
	if strings.HasPrefix(query, "ar:") {
		return strings.Replace(query, "ar:", "artist:", 1)
	} else if strings.HasPrefix(query, "al:") {
		return strings.Replace(query, "al:", "album:", 1)
	} else if strings.HasPrefix(query, "tr:") {
		return strings.Replace(query, "tr:", "track:", 1)
	}
	return query
}

func (spotify *Spotify) getWebApiClient() *webspotify.Client {
	if spotify.client != nil {
		return spotify.client
	}
	return webspotify.DefaultClient
}

func (spotify *Spotify) pause() {
	if spotify.currentTrack != nil {
		if spotify.paused {
			spotify.play(spotify.currentTrack)
		} else {
			spotify.pauseCurrentTrack()
		}
	}
}

func (spotify *Spotify) pauseCurrentTrack() {
	player := spotify.session.Player()
	player.Pause()
	spotify.events.TrackPaused(spotify.currentTrack)
	spotify.paused = true
}

func (spotify *Spotify) artistAlbums(artist *sconsify.Artist) {
	if simpleAlbumPage, err := spotify.client.GetArtistAlbums(webspotify.ID(artist.GetSpotifyID())); err == nil {
		folder := sconsify.InitFolder(artist.URI, "*"+artist.Name, make([]*sconsify.Playlist, 0))

		if fullTracks, err := spotify.client.GetArtistsTopTracks(webspotify.ID(artist.GetSpotifyID()), "GB"); err == nil {
			tracks := make([]*sconsify.Track, len(fullTracks))
			for i, track := range fullTracks {
				tracks[i] = sconsify.InitWebApiTrack(string(track.URI), artist, track.Name, track.TimeDuration().String())
			}

			folder.AddPlaylist(sconsify.InitPlaylist(artist.URI, " "+artist.Name+" Top Tracks", tracks))
		}

		infrastructure.Debugf("# of albums %v", len(simpleAlbumPage.Albums))
		for _, simpleAlbum := range simpleAlbumPage.Albums {
			infrastructure.Debugf("AlbumsID %v = %v", simpleAlbum.URI, simpleAlbum.Name)
			playlist := sconsify.InitOnDemandPlaylist(string(simpleAlbum.URI), " "+simpleAlbum.Name, true, func(playlist *sconsify.Playlist) {
				infrastructure.Debugf("Album id %v", playlist.ToSpotifyID())
				if simpleTrackPage, err := spotify.client.GetAlbumTracks(webspotify.ID(playlist.ToSpotifyID())); err == nil {
					infrastructure.Debugf("# of tracks %v", len(simpleTrackPage.Tracks))
					for _, track := range simpleTrackPage.Tracks {
						playlist.AddTrack(sconsify.InitWebApiTrack(string(track.URI), artist, track.Name, track.TimeDuration().String()))
					}
				}
			})
			folder.AddPlaylist(playlist)
		}

		spotify.events.ArtistAlbums(folder)
	}
}
<commit_msg>Do not use auth client when searching<commit_after>package spotify

import (
	"github.com/fabiofalci/sconsify/infrastructure"
	"github.com/fabiofalci/sconsify/sconsify"
	sp "github.com/op/go-libspotify/spotify"
	webspotify "github.com/zmb3/spotify"
	"strings"
	"time"
)

func (spotify *Spotify) shutdownSpotify() {
	spotify.session.Logout()
	spotify.session.Close()
	spotify.initCache()
	spotify.events.ShutdownEngine()
}

func (spotify *Spotify) play(trackUri *sconsify.Track) {

	player := spotify.session.Player()
	if !spotify.paused || spotify.currentTrack != trackUri {
		link, err := spotify.session.ParseLink(trackUri.URI)
		if err != nil {
			return
		}

		track, err := link.Track()
		if err != nil {
			return
		}

		if trackUri.IsPartial() {
			trackUri = sconsify.ToSconsifyTrack(track)
		}

		if !spotify.isTrackAvailable(track) {
			if trackUri.IsFromWebApi() {
				retry := trackUri.RetryLoading()
				if retry < 4 {
					go func() {
						time.Sleep(100 * time.Millisecond)
						spotify.events.Play(trackUri)
					}()
					return
				}
			}
			spotify.events.TrackNotAvailable(trackUri)
			return
		}
		if err := player.Load(track); err != nil {
			return
		}
		spotify.events.NewTrackLoaded(track.Duration())
	}
	player.Play()

	spotify.events.TrackPlaying(trackUri)
	spotify.currentTrack = trackUri
	spotify.paused = false
	return
}

func (spotify *Spotify) isTrackAvailable(track *sp.Track) bool {
	return track.Availability() == sp.TrackAvailabilityAvailable
}

func (spotify *Spotify) search(query string) {
	playlists := sconsify.InitPlaylists()

	query = checkAlias(query)
	name := " " + query

	playlist := sconsify.InitSearchPlaylist(name, name, func(playlist *sconsify.Playlist) {
		options := createWebSpotifyOptions(50, playlist.Tracks())
		// TODO use spotify.client while still token valid, then switch to Default one
		if searchResult, err := webspotify.DefaultClient.SearchOpt(query, webspotify.SearchTypeTrack, options); err == nil {
			numberOfTracks := len(searchResult.Tracks.Tracks)
			infrastructure.Debugf("Search '%v' returned %v track(s)", query, numberOfTracks)
			for _, track := range searchResult.Tracks.Tracks {
				webArtist := track.Artists[0]
				artist := sconsify.InitArtist(string(webArtist.URI), webArtist.Name)
				playlist.AddTrack(sconsify.InitWebApiTrack(string(track.URI), artist, track.Name, track.TimeDuration().String()))
				infrastructure.Debugf("\tTrack '%v' (%v)", track.URI, track.Name)
			}
		} else {
			infrastructure.Debugf("Spotify search returning error: %v", err)
		}
	})
	playlist.ExecuteLoad()
	playlists.AddPlaylist(playlist)

	spotify.events.NewPlaylist(playlists)
}

func checkAlias(query string) string {
	if strings.HasPrefix(query, "ar:") {
		return strings.Replace(query, "ar:", "artist:", 1)
	} else if strings.HasPrefix(query, "al:") {
		return strings.Replace(query, "al:", "album:", 1)
	} else if strings.HasPrefix(query, "tr:") {
		return strings.Replace(query, "tr:", "track:", 1)
	}
	return query
}

func (spotify *Spotify) pause() {
	if spotify.currentTrack != nil {
		if spotify.paused {
			spotify.play(spotify.currentTrack)
		} else {
			spotify.pauseCurrentTrack()
		}
	}
}

func (spotify *Spotify) pauseCurrentTrack() {
	player := spotify.session.Player()
	player.Pause()
	spotify.events.TrackPaused(spotify.currentTrack)
	spotify.paused = true
}

func (spotify *Spotify) artistAlbums(artist *sconsify.Artist) {
	if simpleAlbumPage, err := spotify.client.GetArtistAlbums(webspotify.ID(artist.GetSpotifyID())); err == nil {
		folder := sconsify.InitFolder(artist.URI, "*"+artist.Name, make([]*sconsify.Playlist, 0))

		if fullTracks, err := spotify.client.GetArtistsTopTracks(webspotify.ID(artist.GetSpotifyID()), "GB"); err == nil {
			tracks := make([]*sconsify.Track, len(fullTracks))
			for i, track := range fullTracks {
				tracks[i] = sconsify.InitWebApiTrack(string(track.URI), artist, track.Name, track.TimeDuration().String())
			}

			folder.AddPlaylist(sconsify.InitPlaylist(artist.URI, " "+artist.Name+" Top Tracks", tracks))
		}

		infrastructure.Debugf("# of albums %v", len(simpleAlbumPage.Albums))
		for _, simpleAlbum := range simpleAlbumPage.Albums {
			infrastructure.Debugf("AlbumsID %v = %v", simpleAlbum.URI, simpleAlbum.Name)
			playlist := sconsify.InitOnDemandPlaylist(string(simpleAlbum.URI), " "+simpleAlbum.Name, true, func(playlist *sconsify.Playlist) {
				infrastructure.Debugf("Album id %v", playlist.ToSpotifyID())
				if simpleTrackPage, err := spotify.client.GetAlbumTracks(webspotify.ID(playlist.ToSpotifyID())); err == nil {
					infrastructure.Debugf("# of tracks %v", len(simpleTrackPage.Tracks))
					for _, track := range simpleTrackPage.Tracks {
						playlist.AddTrack(sconsify.InitWebApiTrack(string(track.URI), artist, track.Name, track.TimeDuration().String()))
					}
				}
			})
			folder.AddPlaylist(playlist)
		}

		spotify.events.ArtistAlbums(folder)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"reflect"
	"sort"
	"testing"
	"time"
)

func init() {
	// run against config template
	loadConfig("meowkov.conf.template")
}

func TestProcessInput(t *testing.T) {
	input := "1 2 3 4 5 6"
	expWords := []string{"1", "2", "3", "4", "5", "6", stop}
	expSeeds := [][]string{
		{"1", "2", "3"},
		{"2", "3", "4"},
		{"3", "4", "5"},
		{"4", "5", "6"},
		{"5", "6", stop},
	}
	words, seeds := processInput(input, false)
	if !reflect.DeepEqual(words, expWords) {
		t.Error("processInput words do not match expected value")
	}
	if !reflect.DeepEqual(seeds, expSeeds) {
		t.Error("processInput seeds do not match expected value")
	}

}

func TestParseInput(t *testing.T) {
	test := func(input string, expected []string) {
		words := parseInput(input)
		if !reflect.DeepEqual(words, expected) {
			t.Error("parseInput words " + dump(words) + " do not match expected " + dump(expected))
		}
	}

	// plain message
	input := "1 2 3"
	expectedWords := []string{"1", "2", "3", stop}
	test(input, expectedWords)

	// remove mentions present at the beginning
	input = config.BotName + ": 1 2 3"
	test(input, expectedWords)
	input = config.BotName + ", 1 2 3"
	test(input, expectedWords)

	// remove BotName if used as mention at the beginning
	input = config.BotName + ": look: 2 3"
	expectedWords = []string{"look:", "2", "3", stop}
	test(input, expectedWords)
	input = config.BotName + ", look: 2 3"
	test(input, expectedWords)

	// do not remove BotName if in the middle
	input = "1 " + config.BotName + " 2 3"
	expectedWords = []string{"1", config.BotName, "2", "3", stop}
	test(input, expectedWords)

	// lowercase input with exception of URLs
	input = "PlAy PiAno https://yt.aergia.eu/#v=T0rs3R4E1Sk&t=23;30"
	expectedWords = []string{"play", "piano", "https://yt.aergia.eu/#v=T0rs3R4E1Sk&t=23;30", stop}
	test(input, expectedWords)
}

func TestNormalizeWord(t *testing.T) {
	test := func(input string, expected string) {
		normalized := normalizeWord(input)
		if normalized != expected {
			t.Error("normalizeWord result >" + normalized + "< does not match expected >" + expected + "<")
		}
	}

	// strip spaces and lowercase input
	test(" CaSe ", "case")

	// strip spaces but no not lowercase URL
	test("  https://yt.aergia.eu/#v=T0rs3R4E1Sk&t=23;3 ", "https://yt.aergia.eu/#v=T0rs3R4E1Sk&t=23;3")

	// remove " from beginning and/or end
	test(" \"foo", "foo")
	test(" foo\" ", "foo")
	test(" \"foo\" ", "foo")
	test(" f\"oo ", "f\"oo")

	// remove ' from beginning and/or end
	test(" 'foo", "foo")
	test(" foo' ", "foo")
	test(" 'foo' ", "foo")
	test(" f'oo ", "f'oo")

	// remove ( and ) from beginning and/or end
	test(" (foo)", "foo")
	test(" (foo ", "foo")
	test(" foo) ", "foo")
	test(" f(oo ", "f(oo")

	// remove [ and ] from beginning and/or end
	test(" [foo]", "foo")
	test(" [foo ", "foo")
	test(" foo] ", "foo")
	test(" f[oo ", "f[oo")

	// remove ? and ! from end only
	test(" foo? ", "foo")
	test(" foo! ", "foo")
	test(" foo!?!?!? ", "foo")
	test(" foo!?bar ", "foo!?bar")
}

func TestGetRedisServer(t *testing.T) {
	config.RedisServer = "foo:1234"
	if getRedisServer() != "foo:1234" {
		t.Error("redis address should be loaded from config")
	}
	os.Setenv("REDIS_PORT_1234_TCP_ADDR", "bar")
	if getRedisServer() != "bar:1234" {
		t.Error("redis address should come from ENV when run in docker")
	}
}

func TestIsEmpty(t *testing.T) {
	problem := !isEmpty("") || !isEmpty(stop)
	if problem {
		t.Error("Empty string should be empty ;-)")
	}
}

func TestIsChainEmpty(t *testing.T) {
	problem := !isChainEmpty([]string{stop}) || !isChainEmpty([]string{}) || !isChainEmpty([]string{""})
	if problem {
		t.Error("Empty slice should be empty ;-)")
	}
}

func TestCalculateChattiness(t *testing.T) {
	privateQuery := false
	chattiness := calculateChattiness("foo bar one two", "nickname", privateQuery)
	if chattiness != config.DefaultChattiness {
		t.Error("calculateChattiness should return DefaultChattiness if bot's nickname is not mentioned")
	}
	chattiness = calculateChattiness("foo bar nickname one two", "nickname", privateQuery)
	if chattiness != always {
		t.Error("calculateChattiness should return 1.0 if nickname is mentioned")
	}
	privateQuery = true
	chattiness = calculateChattiness("foo bar one two", "nickname", privateQuery)
	if chattiness != always {
		t.Error("calculateChattiness should return 1.0 if input is from a private query")
	}

}

func TestInputSource(t *testing.T) {
	rawChannelMsg := ":foo!~bar@unaffiliated/foobar PRIVMSG #test :foo"
	rawPrivateMsg := ":foo!~bar@unaffiliated/foobar PRIVMSG meowkov :foo"
	ownNick := "meowkov"

	source, priv := inputSource(rawChannelMsg, ownNick)
	if source != "#test" || priv {
		t.Error("inputSource should return channel")
	}
	source, priv = inputSource(rawPrivateMsg, ownNick)
	if source != "foo" || !priv {
		t.Error("inputSource should return private query")
	}

}

func TestCreateSeeds(t *testing.T) {
	input := []string{"1", "2", "3", "4", "5", "6"}
	expected := [][]string{
		{"1", "2", "3"},
		{"2", "3", "4"},
		{"3", "4", "5"},
		{"4", "5", "6"},
	}
	output := createSeeds(input)
	if !reflect.DeepEqual(output, expected) {
		t.Error("createSeeds returns incorrect chain groups")
	}
}

func TestiAppendTransliterations(t *testing.T) {
	test := func(input [][]string, expected [][]string) {
		output := chainTransliterations(input)
		if !reflect.DeepEqual(output, expected) {
			t.Error("chainTransliterations returns incorrect chain groups")
		}
	}

	test([][]string{
		{"2", "3", "4"},
		{"3", "ź", "5"},
		{"4", "5", "żółć"},
	}, [][]string{
		{"3", "z", "5"},
		{"4", "5", "zolc"},
	})

	test([][]string{
		{"2", "3", "4"},
		{"3", "4", "5"},
	}, [][]string{})
}

func TestContains(t *testing.T) {
	items := []string{"1", "2", "3"}
	test := func(items []string, item string, expected bool) {
		if contains(items, item) != expected {
			t.Error("contains(" + dump(items) + "," + item + ") should return " + fmt.Sprint(expected))
		}
	}
	test(items, "1", true)
	test(items, "2", true)
	test(items, "3", true)
	test(items, "A", false)

}

func TestMutateChain(t *testing.T) {
	input := []string{"1", "2"}
	word := "A"
	expected := []string{"A", "1", "A", "2", "A"}
	output := mutateChain(word, input)
	if !reflect.DeepEqual(output, expected) {
		t.Error("mutateChain should return " + dump(expected))
	}
}

func TestRandomSmiley(t *testing.T) {
	if !contains(config.Smileys, randomSmiley()) {
		t.Error("randomSmiley should return random item from the list in config file")
	}
}

func TestRemoveBlacklistedWords(t *testing.T) {
	blacklistOrig := config.Blacklist
	dontEndOrig := config.DontEndWith

	config.Blacklist = []string{"2"}
	config.DontEndWith = []string{"5", "6"}
	input := []string{"1", "2", "3", "4", "5", "6"}
	expected := []string{"1", "3", "4"}
	output := removeBlacklistedWords(input)
	if !reflect.DeepEqual(output, expected) {
		t.Error("removeBlacklistedWords should return " + dump(expected))
	}

	config.Blacklist = blacklistOrig
	config.DontEndWith = dontEndOrig
}

func TestMedian(t *testing.T) {
	input := []int{6, 2, 3, 4, 5, 1}
	expected := 3
	output := median(input)
	if output != expected {
		t.Error("median should return " + fmt.Sprint(expected) + " but got " + fmt.Sprint(output))
	}
}

func TestNormalizeResponseChains(t *testing.T) {
	input := make(uniqueTexts)
	input["1"] = struct{}{}
	input["22"] = struct{}{}
	input["333"] = struct{}{}
	input["4444"] = struct{}{}
	input["55555"] = struct{}{}
	input["666666"] = struct{}{}
	expected := []string{"333", "4444", "55555", "666666"}
	output := normalizeResponseChains(input)
	sort.Strings(output)
	if !reflect.DeepEqual(output, expected) {
		t.Error("normalizeResponseChains should return " + dump(expected) + " but got " + dump(output))
	}
}

func TestDump(t *testing.T) {
	input := []string{"1", "2", "3"}
	expected := `["1", "2", "3"]`
	output := dump(input)
	if expected != output {
		t.Error("dump should return " + expected + " but got " + output)
	}
}

func TestTypingDelay(t *testing.T) {
	start := time.Now()
	typingDelay("fooo bar", time.Unix(start.Unix()-1, 0))
	end := time.Now()
	if end.Sub(start) > 1*time.Second {
		t.Error("typingDelay should occur if response took long time to generate")
	}
}
<commit_msg>go vet: first letter after 'Test' must not be lowercase<commit_after>package main

import (
	"fmt"
	"os"
	"reflect"
	"sort"
	"testing"
	"time"
)

func init() {
	// run against config template
	loadConfig("meowkov.conf.template")
}

func TestProcessInput(t *testing.T) {
	input := "1 2 3 4 5 6"
	expWords := []string{"1", "2", "3", "4", "5", "6", stop}
	expSeeds := [][]string{
		{"1", "2", "3"},
		{"2", "3", "4"},
		{"3", "4", "5"},
		{"4", "5", "6"},
		{"5", "6", stop},
	}
	words, seeds := processInput(input, false)
	if !reflect.DeepEqual(words, expWords) {
		t.Error("processInput words do not match expected value")
	}
	if !reflect.DeepEqual(seeds, expSeeds) {
		t.Error("processInput seeds do not match expected value")
	}

}

func TestParseInput(t *testing.T) {
	test := func(input string, expected []string) {
		words := parseInput(input)
		if !reflect.DeepEqual(words, expected) {
			t.Error("parseInput words " + dump(words) + " do not match expected " + dump(expected))
		}
	}

	// plain message
	input := "1 2 3"
	expectedWords := []string{"1", "2", "3", stop}
	test(input, expectedWords)

	// remove mentions present at the beginning
	input = config.BotName + ": 1 2 3"
	test(input, expectedWords)
	input = config.BotName + ", 1 2 3"
	test(input, expectedWords)

	// remove BotName if used as mention at the beginning
	input = config.BotName + ": look: 2 3"
	expectedWords = []string{"look:", "2", "3", stop}
	test(input, expectedWords)
	input = config.BotName + ", look: 2 3"
	test(input, expectedWords)

	// do not remove BotName if in the middle
	input = "1 " + config.BotName + " 2 3"
	expectedWords = []string{"1", config.BotName, "2", "3", stop}
	test(input, expectedWords)

	// lowercase input with exception of URLs
	input = "PlAy PiAno https://yt.aergia.eu/#v=T0rs3R4E1Sk&t=23;30"
	expectedWords = []string{"play", "piano", "https://yt.aergia.eu/#v=T0rs3R4E1Sk&t=23;30", stop}
	test(input, expectedWords)
}

func TestNormalizeWord(t *testing.T) {
	test := func(input string, expected string) {
		normalized := normalizeWord(input)
		if normalized != expected {
			t.Error("normalizeWord result >" + normalized + "< does not match expected >" + expected + "<")
		}
	}

	// strip spaces and lowercase input
	test(" CaSe ", "case")

	// strip spaces but no not lowercase URL
	test("  https://yt.aergia.eu/#v=T0rs3R4E1Sk&t=23;3 ", "https://yt.aergia.eu/#v=T0rs3R4E1Sk&t=23;3")

	// remove " from beginning and/or end
	test(" \"foo", "foo")
	test(" foo\" ", "foo")
	test(" \"foo\" ", "foo")
	test(" f\"oo ", "f\"oo")

	// remove ' from beginning and/or end
	test(" 'foo", "foo")
	test(" foo' ", "foo")
	test(" 'foo' ", "foo")
	test(" f'oo ", "f'oo")

	// remove ( and ) from beginning and/or end
	test(" (foo)", "foo")
	test(" (foo ", "foo")
	test(" foo) ", "foo")
	test(" f(oo ", "f(oo")

	// remove [ and ] from beginning and/or end
	test(" [foo]", "foo")
	test(" [foo ", "foo")
	test(" foo] ", "foo")
	test(" f[oo ", "f[oo")

	// remove ? and ! from end only
	test(" foo? ", "foo")
	test(" foo! ", "foo")
	test(" foo!?!?!? ", "foo")
	test(" foo!?bar ", "foo!?bar")
}

func TestGetRedisServer(t *testing.T) {
	config.RedisServer = "foo:1234"
	if getRedisServer() != "foo:1234" {
		t.Error("redis address should be loaded from config")
	}
	os.Setenv("REDIS_PORT_1234_TCP_ADDR", "bar")
	if getRedisServer() != "bar:1234" {
		t.Error("redis address should come from ENV when run in docker")
	}
}

func TestIsEmpty(t *testing.T) {
	problem := !isEmpty("") || !isEmpty(stop)
	if problem {
		t.Error("Empty string should be empty ;-)")
	}
}

func TestIsChainEmpty(t *testing.T) {
	problem := !isChainEmpty([]string{stop}) || !isChainEmpty([]string{}) || !isChainEmpty([]string{""})
	if problem {
		t.Error("Empty slice should be empty ;-)")
	}
}

func TestCalculateChattiness(t *testing.T) {
	privateQuery := false
	chattiness := calculateChattiness("foo bar one two", "nickname", privateQuery)
	if chattiness != config.DefaultChattiness {
		t.Error("calculateChattiness should return DefaultChattiness if bot's nickname is not mentioned")
	}
	chattiness = calculateChattiness("foo bar nickname one two", "nickname", privateQuery)
	if chattiness != always {
		t.Error("calculateChattiness should return 1.0 if nickname is mentioned")
	}
	privateQuery = true
	chattiness = calculateChattiness("foo bar one two", "nickname", privateQuery)
	if chattiness != always {
		t.Error("calculateChattiness should return 1.0 if input is from a private query")
	}

}

func TestInputSource(t *testing.T) {
	rawChannelMsg := ":foo!~bar@unaffiliated/foobar PRIVMSG #test :foo"
	rawPrivateMsg := ":foo!~bar@unaffiliated/foobar PRIVMSG meowkov :foo"
	ownNick := "meowkov"

	source, priv := inputSource(rawChannelMsg, ownNick)
	if source != "#test" || priv {
		t.Error("inputSource should return channel")
	}
	source, priv = inputSource(rawPrivateMsg, ownNick)
	if source != "foo" || !priv {
		t.Error("inputSource should return private query")
	}

}

func TestCreateSeeds(t *testing.T) {
	input := []string{"1", "2", "3", "4", "5", "6"}
	expected := [][]string{
		{"1", "2", "3"},
		{"2", "3", "4"},
		{"3", "4", "5"},
		{"4", "5", "6"},
	}
	output := createSeeds(input)
	if !reflect.DeepEqual(output, expected) {
		t.Error("createSeeds returns incorrect chain groups")
	}
}

func TestAppendTransliterations(t *testing.T) {
	test := func(input [][]string, expected [][]string) {
		output := chainTransliterations(input)
		if !reflect.DeepEqual(output, expected) {
			t.Error("chainTransliterations returns incorrect chain groups")
		}
	}

	test([][]string{
		{"2", "3", "4"},
		{"3", "ź", "5"},
		{"4", "5", "żółć"},
	}, [][]string{
		{"3", "z", "5"},
		{"4", "5", "zolc"},
	})

	test([][]string{
		{"2", "3", "4"},
		{"3", "4", "5"},
	}, [][]string{})
}

func TestContains(t *testing.T) {
	items := []string{"1", "2", "3"}
	test := func(items []string, item string, expected bool) {
		if contains(items, item) != expected {
			t.Error("contains(" + dump(items) + "," + item + ") should return " + fmt.Sprint(expected))
		}
	}
	test(items, "1", true)
	test(items, "2", true)
	test(items, "3", true)
	test(items, "A", false)

}

func TestMutateChain(t *testing.T) {
	input := []string{"1", "2"}
	word := "A"
	expected := []string{"A", "1", "A", "2", "A"}
	output := mutateChain(word, input)
	if !reflect.DeepEqual(output, expected) {
		t.Error("mutateChain should return " + dump(expected))
	}
}

func TestRandomSmiley(t *testing.T) {
	if !contains(config.Smileys, randomSmiley()) {
		t.Error("randomSmiley should return random item from the list in config file")
	}
}

func TestRemoveBlacklistedWords(t *testing.T) {
	blacklistOrig := config.Blacklist
	dontEndOrig := config.DontEndWith

	config.Blacklist = []string{"2"}
	config.DontEndWith = []string{"5", "6"}
	input := []string{"1", "2", "3", "4", "5", "6"}
	expected := []string{"1", "3", "4"}
	output := removeBlacklistedWords(input)
	if !reflect.DeepEqual(output, expected) {
		t.Error("removeBlacklistedWords should return " + dump(expected))
	}

	config.Blacklist = blacklistOrig
	config.DontEndWith = dontEndOrig
}

func TestMedian(t *testing.T) {
	input := []int{6, 2, 3, 4, 5, 1}
	expected := 3
	output := median(input)
	if output != expected {
		t.Error("median should return " + fmt.Sprint(expected) + " but got " + fmt.Sprint(output))
	}
}

func TestNormalizeResponseChains(t *testing.T) {
	input := make(uniqueTexts)
	input["1"] = struct{}{}
	input["22"] = struct{}{}
	input["333"] = struct{}{}
	input["4444"] = struct{}{}
	input["55555"] = struct{}{}
	input["666666"] = struct{}{}
	expected := []string{"333", "4444", "55555", "666666"}
	output := normalizeResponseChains(input)
	sort.Strings(output)
	if !reflect.DeepEqual(output, expected) {
		t.Error("normalizeResponseChains should return " + dump(expected) + " but got " + dump(output))
	}
}

func TestDump(t *testing.T) {
	input := []string{"1", "2", "3"}
	expected := `["1", "2", "3"]`
	output := dump(input)
	if expected != output {
		t.Error("dump should return " + expected + " but got " + output)
	}
}

func TestTypingDelay(t *testing.T) {
	start := time.Now()
	typingDelay("fooo bar", time.Unix(start.Unix()-1, 0))
	end := time.Now()
	if end.Sub(start) > 1*time.Second {
		t.Error("typingDelay should occur if response took long time to generate")
	}
}
<|endoftext|>


<commit_before>package graph

import (
	"fmt"
	"geo"
	"mm"
	"path"
	"sort"
)

type OverlayGraphFile struct {
	*GraphFile
	Cluster          []uint16      // cluster id -> vertex indices
	VertexIndices    []int         // vertex indices -> cluster id
	Matrices         [][][]float32 // transport mode -> metric -> (cluster id, i, j) -> weight
	ClusterEdgeCount int           // combined boundary edge count of the clusters
	EdgeCounts       []int         // cluster id -> id of first edge inside the cluster
}

// I/O

func computeVertexIndices(g *OverlayGraphFile) {
	g.VertexIndices = make([]int, g.VertexCount())
	for i := 0; i < g.ClusterCount(); i++ {
		for j := g.Cluster[i]; j < g.Cluster[i+1]; j++ {
			g.VertexIndices[j] = i
		}
	}
}

func computeEdgeCounts(g *OverlayGraphFile) {
	g.EdgeCounts = make([]int, g.ClusterCount()+1)
	g.EdgeCounts[0] = g.GraphFile.EdgeCount()
	for i := 0; i < g.ClusterCount(); i++ {
		g.EdgeCounts[i+1] = g.EdgeCounts[i] + g.ClusterSize(i)*g.ClusterSize(i)
	}
}

func loadAllMatrices(g *OverlayGraphFile, base string) error {
	g.Matrices = make([][][]float32, TransportMax)
	for t := 0; t < int(TransportMax); t++ {
		g.Matrices[t] = make([][]float32, MetricMax)
		for m := 0; m < int(MetricMax); m++ {
			var matrixFile []float32
			fileName := fmt.Sprintf("matrices.trans%d.metric%d.ftf", t+1, m+1)
			err := mm.Open(path.Join(base, fileName), &matrixFile)
			if err != nil {
				return err
			}
			g.Matrices[t][m] = matrixFile
		}
	}
	return nil
}

func OpenOverlay(base string, loadMatrices, ignoreErrors bool) (*OverlayGraphFile, error) {
	overlayBaseDir := path.Join(base, "/overlay")
	g, err := OpenGraphFile(overlayBaseDir, ignoreErrors)
	if err != nil && !ignoreErrors {
		return nil, err
	}

	overlay := &OverlayGraphFile{GraphFile: g}
	files := []struct {
		name string
		p    interface{}
	}{
		{"partitions.ftf", &overlay.Cluster},
	}

	for _, file := range files {
		err = mm.Open(path.Join(overlayBaseDir, file.name), file.p)
		if err != nil && !ignoreErrors {
			return nil, err
		}
	}

	computeVertexIndices(overlay)
	computeEdgeCounts(overlay)
	if loadMatrices {
		err = loadAllMatrices(overlay, base)
		if err != nil && !ignoreErrors {
			return nil, err
		}
	}

	for i := 0; i < overlay.ClusterCount(); i++ {
		overlay.ClusterEdgeCount += overlay.ClusterSize(i) * overlay.ClusterSize(i)
	}

	return overlay, nil
}

func CloseOverlay(overlay *OverlayGraphFile) error {
	err := CloseGraphFile(overlay.GraphFile)
	if err != nil {
		return err
	}

	files := []interface{}{
		&overlay.Cluster,
	}

	for _, p := range files {
		err = mm.Close(p)
		if err != nil {
			return err
		}
	}

	return nil
}

// Graph Interface

func (g *OverlayGraphFile) EdgeCount() int {
	// Count edges and matrices...
	return g.GraphFile.EdgeCount() + g.ClusterEdgeCount
}

func (g *OverlayGraphFile) VertexEdges(v Vertex, forward bool, t Transport, buf []Edge) []Edge {
	// Add the cut edges
	result := g.GraphFile.VertexEdges(v, forward, t, buf)
	// Add the precomputed edges.
	cluster, indexInCluster := g.VertexCluster(v)
	clusterStart := g.EdgeCounts[cluster]
	clusterSize := g.ClusterSize(cluster)
	if forward {
		// out edges
		outEdgesStart := clusterStart + int(indexInCluster)*clusterSize
		for i := 0; i < clusterSize; i++ {
			result = append(result, Edge(outEdgesStart+i))
		}
	} else {
		// in edges
		for i := 0; i < clusterSize; i++ {
			result = append(result, Edge(i*clusterSize+int(indexInCluster)))
		}
	}
	return result
}

func (g *OverlayGraphFile) IsCutEdge(e Edge) bool {
	return int(e) < g.GraphFile.EdgeCount()
}

func (g *OverlayGraphFile) EdgeOpposite(e Edge, v Vertex) Vertex {
	if g.IsCutEdge(e) {
		return g.GraphFile.EdgeOpposite(e, v)
	}
	// binary search for cluster id
	cluster := sort.Search(g.ClusterCount(), func(i int) bool { return int(e) < g.EdgeCounts[i+1] })

	e = e - Edge(g.EdgeCounts[cluster])
	vCheck := int(e) / g.ClusterSize(cluster)
	if int(v) != vCheck {
		panic("index of v is not as expected")
	}
	u := int(e) % g.ClusterSize(cluster)
	return Vertex(u)
}

func (g *OverlayGraphFile) EdgeSteps(e Edge, from Vertex) []geo.Coordinate {
	// Return nil unless the edge is a cross partition edge.
	// In this case, defer to the normal Graph interface.
	if g.IsCutEdge(e) {
		return g.GraphFile.EdgeSteps(e, from)
	}
	return nil
}

func (g *OverlayGraphFile) EdgeWeight(e Edge, t Transport, m Metric) float64 {
	// Return the normal weight if e is a cross partition edge,
	// otherwise return the precomputed weight for t and m.
	if g.IsCutEdge(e) {
		return g.GraphFile.EdgeWeight(e, t, m)
	}
	edgeIndex := int(e) - g.GraphFile.EdgeCount()
	return float64(g.Matrices[t][m][edgeIndex])
}

// Overlay Interface

func (g *OverlayGraphFile) ClusterCount() int {
	return len(g.Cluster) - 1
}

// actually: boundary vertex count
func (g *OverlayGraphFile) ClusterSize(i int) int {
	// cluster id -> number of vertices
	return int(g.Cluster[i+1] - g.Cluster[i])
}

func (g *OverlayGraphFile) VertexCluster(v Vertex) (int, Vertex) {
	// overlay vertex id -> cluster id, cluster vertex id
	i := g.VertexIndices[v]
	return i, v - Vertex(g.Cluster[i])
}

func (g *OverlayGraphFile) ClusterVertex(i int, v Vertex) Vertex {
	// cluster id, cluster vertex id -> overlay vertex id
	return Vertex(g.Cluster[i]) + v
}
<commit_msg>Fix VertexEdges<commit_after>package graph

import (
	"fmt"
	"geo"
	"mm"
	"path"
	"sort"
)

type OverlayGraphFile struct {
	*GraphFile
	Cluster          []uint16      // cluster id -> vertex indices
	VertexIndices    []int         // vertex indices -> cluster id
	Matrices         [][][]float32 // transport mode -> metric -> (cluster id, i, j) -> weight
	ClusterEdgeCount int           // combined boundary edge count of the clusters
	EdgeCounts       []int         // cluster id -> id of first edge inside the cluster
}

// I/O

func computeVertexIndices(g *OverlayGraphFile) {
	g.VertexIndices = make([]int, g.VertexCount())
	for i := 0; i < g.ClusterCount(); i++ {
		for j := g.Cluster[i]; j < g.Cluster[i+1]; j++ {
			g.VertexIndices[j] = i
		}
	}
}

func computeEdgeCounts(g *OverlayGraphFile) {
	g.EdgeCounts = make([]int, g.ClusterCount()+1)
	g.EdgeCounts[0] = g.GraphFile.EdgeCount()
	for i := 0; i < g.ClusterCount(); i++ {
		g.EdgeCounts[i+1] = g.EdgeCounts[i] + g.ClusterSize(i)*g.ClusterSize(i)
	}
}

func loadAllMatrices(g *OverlayGraphFile, base string) error {
	g.Matrices = make([][][]float32, TransportMax)
	for t := 0; t < int(TransportMax); t++ {
		g.Matrices[t] = make([][]float32, MetricMax)
		for m := 0; m < int(MetricMax); m++ {
			var matrixFile []float32
			fileName := fmt.Sprintf("matrices.trans%d.metric%d.ftf", t+1, m+1)
			err := mm.Open(path.Join(base, fileName), &matrixFile)
			if err != nil {
				return err
			}
			g.Matrices[t][m] = matrixFile
		}
	}
	return nil
}

func OpenOverlay(base string, loadMatrices, ignoreErrors bool) (*OverlayGraphFile, error) {
	overlayBaseDir := path.Join(base, "/overlay")
	g, err := OpenGraphFile(overlayBaseDir, ignoreErrors)
	if err != nil && !ignoreErrors {
		return nil, err
	}

	overlay := &OverlayGraphFile{GraphFile: g}
	files := []struct {
		name string
		p    interface{}
	}{
		{"partitions.ftf", &overlay.Cluster},
	}

	for _, file := range files {
		err = mm.Open(path.Join(overlayBaseDir, file.name), file.p)
		if err != nil && !ignoreErrors {
			return nil, err
		}
	}

	computeVertexIndices(overlay)
	computeEdgeCounts(overlay)
	if loadMatrices {
		err = loadAllMatrices(overlay, base)
		if err != nil && !ignoreErrors {
			return nil, err
		}
	}

	for i := 0; i < overlay.ClusterCount(); i++ {
		overlay.ClusterEdgeCount += overlay.ClusterSize(i) * overlay.ClusterSize(i)
	}

	return overlay, nil
}

func CloseOverlay(overlay *OverlayGraphFile) error {
	err := CloseGraphFile(overlay.GraphFile)
	if err != nil {
		return err
	}

	files := []interface{}{
		&overlay.Cluster,
	}

	for _, p := range files {
		err = mm.Close(p)
		if err != nil {
			return err
		}
	}

	return nil
}

// Graph Interface

func (g *OverlayGraphFile) EdgeCount() int {
	// Count edges and matrices...
	return g.GraphFile.EdgeCount() + g.ClusterEdgeCount
}

func (g *OverlayGraphFile) VertexEdges(v Vertex, forward bool, t Transport, buf []Edge) []Edge {
	// Add the cut edges
	result := g.GraphFile.VertexEdges(v, forward, t, buf)
	// Add the precomputed edges.
	cluster, indexInCluster := g.VertexCluster(v)
	clusterStart := g.EdgeCounts[cluster]
	clusterSize := g.ClusterSize(cluster)
	if forward {
		// out edges
		outEdgesStart := clusterStart + int(indexInCluster)*clusterSize
		for i := 0; i < clusterSize; i++ {
			result = append(result, Edge(outEdgesStart+i))
		}
	} else {
		// in edges
		inEdgesStart := clusterStart + int(indexInCluster)
		for i := 0; i < clusterSize; i++ {
			result = append(result, Edge(inEdgesStart + i*clusterSize))
		}
	}
	return result
}

func (g *OverlayGraphFile) IsCutEdge(e Edge) bool {
	return int(e) < g.GraphFile.EdgeCount()
}

func (g *OverlayGraphFile) EdgeOpposite(e Edge, v Vertex) Vertex {
	if g.IsCutEdge(e) {
		return g.GraphFile.EdgeOpposite(e, v)
	}
	// binary search for cluster id
	cluster := sort.Search(g.ClusterCount(), func(i int) bool { return int(e) < g.EdgeCounts[i+1] })

	e = e - Edge(g.EdgeCounts[cluster])
	vCheck := int(e) / g.ClusterSize(cluster)
	if int(v) != vCheck {
		panic("index of v is not as expected")
	}
	u := int(e) % g.ClusterSize(cluster)
	return Vertex(u)
}

func (g *OverlayGraphFile) EdgeSteps(e Edge, from Vertex) []geo.Coordinate {
	// Return nil unless the edge is a cross partition edge.
	// In this case, defer to the normal Graph interface.
	if g.IsCutEdge(e) {
		return g.GraphFile.EdgeSteps(e, from)
	}
	return nil
}

func (g *OverlayGraphFile) EdgeWeight(e Edge, t Transport, m Metric) float64 {
	// Return the normal weight if e is a cross partition edge,
	// otherwise return the precomputed weight for t and m.
	if g.IsCutEdge(e) {
		return g.GraphFile.EdgeWeight(e, t, m)
	}
	edgeIndex := int(e) - g.GraphFile.EdgeCount()
	return float64(g.Matrices[t][m][edgeIndex])
}

// Overlay Interface

func (g *OverlayGraphFile) ClusterCount() int {
	return len(g.Cluster) - 1
}

// actually: boundary vertex count
func (g *OverlayGraphFile) ClusterSize(i int) int {
	// cluster id -> number of vertices
	return int(g.Cluster[i+1] - g.Cluster[i])
}

func (g *OverlayGraphFile) VertexCluster(v Vertex) (int, Vertex) {
	// overlay vertex id -> cluster id, cluster vertex id
	i := g.VertexIndices[v]
	return i, v - Vertex(g.Cluster[i])
}

func (g *OverlayGraphFile) ClusterVertex(i int, v Vertex) Vertex {
	// cluster id, cluster vertex id -> overlay vertex id
	return Vertex(g.Cluster[i]) + v
}
<|endoftext|>


<commit_before>package sarama

import (
	"runtime"
	"testing"
	"time"
)

var (
	emptyMessage = []byte{
		167, 236, 104, 3, // CRC
		0x00,                   // magic version byte
		0x00,                   // attribute flags
		0xFF, 0xFF, 0xFF, 0xFF, // key
		0xFF, 0xFF, 0xFF, 0xFF} // value

	emptyV1Message = []byte{
		204, 47, 121, 217, // CRC
		0x01,                                           // magic version byte
		0x00,                                           // attribute flags
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // timestamp
		0xFF, 0xFF, 0xFF, 0xFF, // key
		0xFF, 0xFF, 0xFF, 0xFF} // value

	emptyV2Message = []byte{
		167, 236, 104, 3, // CRC
		0x02,                   // magic version byte
		0x00,                   // attribute flags
		0xFF, 0xFF, 0xFF, 0xFF, // key
		0xFF, 0xFF, 0xFF, 0xFF} // value

	emptyGzipMessage = []byte{
		97, 79, 149, 90, //CRC
		0x00,                   // magic version byte
		0x01,                   // attribute flags
		0xFF, 0xFF, 0xFF, 0xFF, // key
		// value
		0x00, 0x00, 0x00, 0x17,
		0x1f, 0x8b,
		0x08,
		0, 0, 9, 110, 136, 0, 255, 1, 0, 0, 255, 255, 0, 0, 0, 0, 0, 0, 0, 0}

	emptyGzipMessage18 = []byte{
		132, 99, 80, 148, //CRC
		0x00,                   // magic version byte
		0x01,                   // attribute flags
		0xFF, 0xFF, 0xFF, 0xFF, // key
		// value
		0x00, 0x00, 0x00, 0x17,
		0x1f, 0x8b,
		0x08,
		0, 0, 0, 0, 0, 0, 255, 1, 0, 0, 255, 255, 0, 0, 0, 0, 0, 0, 0, 0}

	emptyLZ4Message = []byte{
		132, 219, 238, 101, // CRC
		0x01,                          // version byte
		0x03,                          // attribute flags: lz4
		0, 0, 1, 88, 141, 205, 89, 56, // timestamp
		0xFF, 0xFF, 0xFF, 0xFF, // key
		0x00, 0x00, 0x00, 0x0f, // len
		0x04, 0x22, 0x4D, 0x18, // LZ4 magic number
		100,                  // LZ4 flags: version 01, block indepedant, content checksum
		112, 185, 0, 0, 0, 0, // LZ4 data
		5, 93, 204, 2, // LZ4 checksum
	}

	emptyBulkSnappyMessage = []byte{
		180, 47, 53, 209, //CRC
		0x00,                   // magic version byte
		0x02,                   // attribute flags
		0xFF, 0xFF, 0xFF, 0xFF, // key
		0, 0, 0, 42,
		130, 83, 78, 65, 80, 80, 89, 0, // SNAPPY magic
		0, 0, 0, 1, // min version
		0, 0, 0, 1, // default version
		0, 0, 0, 22, 52, 0, 0, 25, 1, 16, 14, 227, 138, 104, 118, 25, 15, 13, 1, 8, 1, 0, 0, 62, 26, 0}

	emptyBulkGzipMessage = []byte{
		139, 160, 63, 141, //CRC
		0x00,                   // magic version byte
		0x01,                   // attribute flags
		0xFF, 0xFF, 0xFF, 0xFF, // key
		0x00, 0x00, 0x00, 0x27, // len
		0x1f, 0x8b, // Gzip Magic
		0x08, // deflate compressed
		0, 0, 0, 0, 0, 0, 0, 99, 96, 128, 3, 190, 202, 112, 143, 7, 12, 12, 255, 129, 0, 33, 200, 192, 136, 41, 3, 0, 199, 226, 155, 70, 52, 0, 0, 0}

	emptyBulkLZ4Message = []byte{
		246, 12, 188, 129, // CRC
		0x01,                                  // Version
		0x03,                                  // attribute flags (LZ4)
		255, 255, 249, 209, 212, 181, 73, 201, // timestamp
		0xFF, 0xFF, 0xFF, 0xFF, // key
		0x00, 0x00, 0x00, 0x47, // len
		0x04, 0x22, 0x4D, 0x18, // magic number lz4
		100, // lz4 flags 01100100
		// version: 01, block indep: 1, block checksum: 0, content size: 0, content checksum: 1, reserved: 00
		112, 185, 52, 0, 0, 128, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 14, 121, 87, 72, 224, 0, 0, 255, 255, 255, 255, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 14, 121, 87, 72, 224, 0, 0, 255, 255, 255, 255, 0, 0, 0, 0, 0, 0, 0, 0,
		71, 129, 23, 111, // LZ4 checksum
	}
)

func TestMessageEncoding(t *testing.T) {
	message := Message{}
	testEncodable(t, "empty", &message, emptyMessage)

	message.Value = []byte{}
	message.Codec = CompressionGZIP
	if runtime.Version() == "go1.8" || runtime.Version() == "go1.8.1" {
		testEncodable(t, "empty gzip", &message, emptyGzipMessage18)
	} else {
		testEncodable(t, "empty gzip", &message, emptyGzipMessage)
	}

	message.Value = []byte{}
	message.Codec = CompressionLZ4
	message.Timestamp = time.Unix(1479847795, 0)
	message.Version = 1
	testEncodable(t, "empty lz4", &message, emptyLZ4Message)
}

func TestMessageDecoding(t *testing.T) {
	message := Message{}
	testDecodable(t, "empty", &message, emptyMessage)
	if message.Codec != CompressionNone {
		t.Error("Decoding produced compression codec where there was none.")
	}
	if message.Key != nil {
		t.Error("Decoding produced key where there was none.")
	}
	if message.Value != nil {
		t.Error("Decoding produced value where there was none.")
	}
	if message.Set != nil {
		t.Error("Decoding produced set where there was none.")
	}

	testDecodable(t, "empty gzip", &message, emptyGzipMessage)
	if message.Codec != CompressionGZIP {
		t.Error("Decoding produced incorrect compression codec (was gzip).")
	}
	if message.Key != nil {
		t.Error("Decoding produced key where there was none.")
	}
	if message.Value == nil || len(message.Value) != 0 {
		t.Error("Decoding produced nil or content-ful value where there was an empty array.")
	}
}

func TestMessageDecodingBulkSnappy(t *testing.T) {
	message := Message{}
	testDecodable(t, "bulk snappy", &message, emptyBulkSnappyMessage)
	if message.Codec != CompressionSnappy {
		t.Errorf("Decoding produced codec %d, but expected %d.", message.Codec, CompressionSnappy)
	}
	if message.Key != nil {
		t.Errorf("Decoding produced key %+v, but none was expected.", message.Key)
	}
	if message.Set == nil {
		t.Error("Decoding produced no set, but one was expected.")
	} else if len(message.Set.Messages) != 2 {
		t.Errorf("Decoding produced a set with %d messages, but 2 were expected.", len(message.Set.Messages))
	}
}

func TestMessageDecodingBulkGzip(t *testing.T) {
	message := Message{}
	testDecodable(t, "bulk gzip", &message, emptyBulkGzipMessage)
	if message.Codec != CompressionGZIP {
		t.Errorf("Decoding produced codec %d, but expected %d.", message.Codec, CompressionGZIP)
	}
	if message.Key != nil {
		t.Errorf("Decoding produced key %+v, but none was expected.", message.Key)
	}
	if message.Set == nil {
		t.Error("Decoding produced no set, but one was expected.")
	} else if len(message.Set.Messages) != 2 {
		t.Errorf("Decoding produced a set with %d messages, but 2 were expected.", len(message.Set.Messages))
	}
}

func TestMessageDecodingBulkLZ4(t *testing.T) {
	message := Message{}
	testDecodable(t, "bulk lz4", &message, emptyBulkLZ4Message)
	if message.Codec != CompressionLZ4 {
		t.Errorf("Decoding produced codec %d, but expected %d.", message.Codec, CompressionLZ4)
	}
	if message.Key != nil {
		t.Errorf("Decoding produced key %+v, but none was expected.", message.Key)
	}
	if message.Set == nil {
		t.Error("Decoding produced no set, but one was expected.")
	} else if len(message.Set.Messages) != 2 {
		t.Errorf("Decoding produced a set with %d messages, but 2 were expected.", len(message.Set.Messages))
	}
}

func TestMessageDecodingVersion1(t *testing.T) {
	message := Message{Version: 1}
	testDecodable(t, "decoding empty v1 message", &message, emptyV1Message)
}

func TestMessageDecodingUnknownVersions(t *testing.T) {
	message := Message{Version: 2}
	err := decode(emptyV2Message, &message)
	if err == nil {
		t.Error("Decoding did not produce an error for an unknown magic byte")
	}
	if err.Error() != "kafka: error decoding packet: unknown magic byte (2)" {
		t.Error("Decoding an unknown magic byte produced an unknown error ", err)
	}
}
<commit_msg>Fix gzip test for all possible 1.8 point releases<commit_after>package sarama

import (
	"runtime"
	"strings"
	"testing"
	"time"
)

var (
	emptyMessage = []byte{
		167, 236, 104, 3, // CRC
		0x00,                   // magic version byte
		0x00,                   // attribute flags
		0xFF, 0xFF, 0xFF, 0xFF, // key
		0xFF, 0xFF, 0xFF, 0xFF} // value

	emptyV1Message = []byte{
		204, 47, 121, 217, // CRC
		0x01,                                           // magic version byte
		0x00,                                           // attribute flags
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // timestamp
		0xFF, 0xFF, 0xFF, 0xFF, // key
		0xFF, 0xFF, 0xFF, 0xFF} // value

	emptyV2Message = []byte{
		167, 236, 104, 3, // CRC
		0x02,                   // magic version byte
		0x00,                   // attribute flags
		0xFF, 0xFF, 0xFF, 0xFF, // key
		0xFF, 0xFF, 0xFF, 0xFF} // value

	emptyGzipMessage = []byte{
		97, 79, 149, 90, //CRC
		0x00,                   // magic version byte
		0x01,                   // attribute flags
		0xFF, 0xFF, 0xFF, 0xFF, // key
		// value
		0x00, 0x00, 0x00, 0x17,
		0x1f, 0x8b,
		0x08,
		0, 0, 9, 110, 136, 0, 255, 1, 0, 0, 255, 255, 0, 0, 0, 0, 0, 0, 0, 0}

	emptyGzipMessage18 = []byte{
		132, 99, 80, 148, //CRC
		0x00,                   // magic version byte
		0x01,                   // attribute flags
		0xFF, 0xFF, 0xFF, 0xFF, // key
		// value
		0x00, 0x00, 0x00, 0x17,
		0x1f, 0x8b,
		0x08,
		0, 0, 0, 0, 0, 0, 255, 1, 0, 0, 255, 255, 0, 0, 0, 0, 0, 0, 0, 0}

	emptyLZ4Message = []byte{
		132, 219, 238, 101, // CRC
		0x01,                          // version byte
		0x03,                          // attribute flags: lz4
		0, 0, 1, 88, 141, 205, 89, 56, // timestamp
		0xFF, 0xFF, 0xFF, 0xFF, // key
		0x00, 0x00, 0x00, 0x0f, // len
		0x04, 0x22, 0x4D, 0x18, // LZ4 magic number
		100,                  // LZ4 flags: version 01, block indepedant, content checksum
		112, 185, 0, 0, 0, 0, // LZ4 data
		5, 93, 204, 2, // LZ4 checksum
	}

	emptyBulkSnappyMessage = []byte{
		180, 47, 53, 209, //CRC
		0x00,                   // magic version byte
		0x02,                   // attribute flags
		0xFF, 0xFF, 0xFF, 0xFF, // key
		0, 0, 0, 42,
		130, 83, 78, 65, 80, 80, 89, 0, // SNAPPY magic
		0, 0, 0, 1, // min version
		0, 0, 0, 1, // default version
		0, 0, 0, 22, 52, 0, 0, 25, 1, 16, 14, 227, 138, 104, 118, 25, 15, 13, 1, 8, 1, 0, 0, 62, 26, 0}

	emptyBulkGzipMessage = []byte{
		139, 160, 63, 141, //CRC
		0x00,                   // magic version byte
		0x01,                   // attribute flags
		0xFF, 0xFF, 0xFF, 0xFF, // key
		0x00, 0x00, 0x00, 0x27, // len
		0x1f, 0x8b, // Gzip Magic
		0x08, // deflate compressed
		0, 0, 0, 0, 0, 0, 0, 99, 96, 128, 3, 190, 202, 112, 143, 7, 12, 12, 255, 129, 0, 33, 200, 192, 136, 41, 3, 0, 199, 226, 155, 70, 52, 0, 0, 0}

	emptyBulkLZ4Message = []byte{
		246, 12, 188, 129, // CRC
		0x01,                                  // Version
		0x03,                                  // attribute flags (LZ4)
		255, 255, 249, 209, 212, 181, 73, 201, // timestamp
		0xFF, 0xFF, 0xFF, 0xFF, // key
		0x00, 0x00, 0x00, 0x47, // len
		0x04, 0x22, 0x4D, 0x18, // magic number lz4
		100, // lz4 flags 01100100
		// version: 01, block indep: 1, block checksum: 0, content size: 0, content checksum: 1, reserved: 00
		112, 185, 52, 0, 0, 128, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 14, 121, 87, 72, 224, 0, 0, 255, 255, 255, 255, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 14, 121, 87, 72, 224, 0, 0, 255, 255, 255, 255, 0, 0, 0, 0, 0, 0, 0, 0,
		71, 129, 23, 111, // LZ4 checksum
	}
)

func TestMessageEncoding(t *testing.T) {
	message := Message{}
	testEncodable(t, "empty", &message, emptyMessage)

	message.Value = []byte{}
	message.Codec = CompressionGZIP
	if runtime.Version() == "go1.8" || strings.HasPrefix(runtime.Version(), "go1.8.") {
		testEncodable(t, "empty gzip", &message, emptyGzipMessage18)
	} else {
		testEncodable(t, "empty gzip", &message, emptyGzipMessage)
	}

	message.Value = []byte{}
	message.Codec = CompressionLZ4
	message.Timestamp = time.Unix(1479847795, 0)
	message.Version = 1
	testEncodable(t, "empty lz4", &message, emptyLZ4Message)
}

func TestMessageDecoding(t *testing.T) {
	message := Message{}
	testDecodable(t, "empty", &message, emptyMessage)
	if message.Codec != CompressionNone {
		t.Error("Decoding produced compression codec where there was none.")
	}
	if message.Key != nil {
		t.Error("Decoding produced key where there was none.")
	}
	if message.Value != nil {
		t.Error("Decoding produced value where there was none.")
	}
	if message.Set != nil {
		t.Error("Decoding produced set where there was none.")
	}

	testDecodable(t, "empty gzip", &message, emptyGzipMessage)
	if message.Codec != CompressionGZIP {
		t.Error("Decoding produced incorrect compression codec (was gzip).")
	}
	if message.Key != nil {
		t.Error("Decoding produced key where there was none.")
	}
	if message.Value == nil || len(message.Value) != 0 {
		t.Error("Decoding produced nil or content-ful value where there was an empty array.")
	}
}

func TestMessageDecodingBulkSnappy(t *testing.T) {
	message := Message{}
	testDecodable(t, "bulk snappy", &message, emptyBulkSnappyMessage)
	if message.Codec != CompressionSnappy {
		t.Errorf("Decoding produced codec %d, but expected %d.", message.Codec, CompressionSnappy)
	}
	if message.Key != nil {
		t.Errorf("Decoding produced key %+v, but none was expected.", message.Key)
	}
	if message.Set == nil {
		t.Error("Decoding produced no set, but one was expected.")
	} else if len(message.Set.Messages) != 2 {
		t.Errorf("Decoding produced a set with %d messages, but 2 were expected.", len(message.Set.Messages))
	}
}

func TestMessageDecodingBulkGzip(t *testing.T) {
	message := Message{}
	testDecodable(t, "bulk gzip", &message, emptyBulkGzipMessage)
	if message.Codec != CompressionGZIP {
		t.Errorf("Decoding produced codec %d, but expected %d.", message.Codec, CompressionGZIP)
	}
	if message.Key != nil {
		t.Errorf("Decoding produced key %+v, but none was expected.", message.Key)
	}
	if message.Set == nil {
		t.Error("Decoding produced no set, but one was expected.")
	} else if len(message.Set.Messages) != 2 {
		t.Errorf("Decoding produced a set with %d messages, but 2 were expected.", len(message.Set.Messages))
	}
}

func TestMessageDecodingBulkLZ4(t *testing.T) {
	message := Message{}
	testDecodable(t, "bulk lz4", &message, emptyBulkLZ4Message)
	if message.Codec != CompressionLZ4 {
		t.Errorf("Decoding produced codec %d, but expected %d.", message.Codec, CompressionLZ4)
	}
	if message.Key != nil {
		t.Errorf("Decoding produced key %+v, but none was expected.", message.Key)
	}
	if message.Set == nil {
		t.Error("Decoding produced no set, but one was expected.")
	} else if len(message.Set.Messages) != 2 {
		t.Errorf("Decoding produced a set with %d messages, but 2 were expected.", len(message.Set.Messages))
	}
}

func TestMessageDecodingVersion1(t *testing.T) {
	message := Message{Version: 1}
	testDecodable(t, "decoding empty v1 message", &message, emptyV1Message)
}

func TestMessageDecodingUnknownVersions(t *testing.T) {
	message := Message{Version: 2}
	err := decode(emptyV2Message, &message)
	if err == nil {
		t.Error("Decoding did not produce an error for an unknown magic byte")
	}
	if err.Error() != "kafka: error decoding packet: unknown magic byte (2)" {
		t.Error("Decoding an unknown magic byte produced an unknown error ", err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 Google 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 awskms provides integration with the AWS Cloud KMS.
package awskms

import (
	"encoding/hex"
	"errors"
	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/kms"
	"github.com/aws/aws-sdk-go/service/kms/kmsiface"
)

// AWSAEAD represents a AWS KMS service to a particular URI.
type AWSAEAD struct {
	keyURI string
	kms    kmsiface.KMSAPI
}

// newAWSAEAD returns a new AWS KMS service.
// keyURI must have the following format: 'arn:<partition>:kms:<region>:[:path]'.
// See http://docs.aws.amazon.com/general/latest/gr/aws-arns-and-namespaces.html.
func newAWSAEAD(keyURI string, kms kmsiface.KMSAPI) *AWSAEAD {
	return &AWSAEAD{
		keyURI: keyURI,
		kms:    kms,
	}
}

// Encrypt AEAD encrypts the plaintext data and uses addtionaldata from authentication.
func (a *AWSAEAD) Encrypt(plaintext, additionalData []byte) ([]byte, error) {
	ad := hex.EncodeToString(additionalData)
	req := &kms.EncryptInput{
		KeyId:             aws.String(a.keyURI),
		Plaintext:         plaintext,
		EncryptionContext: map[string]*string{"additionalData": &ad},
	}
	if ad == "" {
		req = &kms.EncryptInput{
			KeyId:     aws.String(a.keyURI),
			Plaintext: plaintext,
		}
	}
	resp, err := a.kms.Encrypt(req)
	if err != nil {
		return nil, err
	}

	return resp.CiphertextBlob, nil
}

// Decrypt AEAD decrypts the data and verified the additional data.
//
// Returns an error if the KeyId field in the response does not match the KeyURI
// provided when creating the client. If we don't do this, the possibility exists
// for the ciphertext to be replaced by one under a key we don't control/expect,
// but do have decrypt permissions on.
//
// This check is disabled if AWESAEAD.keyURI is not in key ARN format.
//
// See https://docs.aws.amazon.com/kms/latest/developerguide/concepts.html#key-id.
func (a *AWSAEAD) Decrypt(ciphertext, additionalData []byte) ([]byte, error) {
	ad := hex.EncodeToString(additionalData)
	req := &kms.DecryptInput{
		KeyId:             aws.String(a.keyURI),
		CiphertextBlob:    ciphertext,
		EncryptionContext: map[string]*string{"additionalData": &ad},
	}
	if ad == "" {
		req = &kms.DecryptInput{
			CiphertextBlob: ciphertext,
		}
	}
	resp, err := a.kms.Decrypt(req)
	if err != nil {
		return nil, err
	}
  if isKeyArnFormat(a.keyURI) && strings.Compare(*resp.KeyId, a.keyURI) != 0 {
		return nil, errors.New("decryption failed: wrong key id")
	}
	return resp.Plaintext, nil
}

// isKeyArnFormat returns true if the keyURI is the KMS Key ARN format; false otherwise.
func isKeyArnFormat(keyURI string) bool {
	tokens := strings.Split(keyURI, ":")
	return len(tokens) == 6 && strings.HasPrefix(tokens[5], "key/")
}
<commit_msg>Update aws_kms_aead.go<commit_after>// Copyright 2017 Google 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 awskms provides integration with the AWS Cloud KMS.
package awskms

import (
	"encoding/hex"
	"errors"
	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/kms"
	"github.com/aws/aws-sdk-go/service/kms/kmsiface"
)

// AWSAEAD represents a AWS KMS service to a particular URI.
type AWSAEAD struct {
	keyURI string
	kms    kmsiface.KMSAPI
}

// newAWSAEAD returns a new AWS KMS service.
// keyURI must have the following format: 'arn:<partition>:kms:<region>:[:path]'.
// See http://docs.aws.amazon.com/general/latest/gr/aws-arns-and-namespaces.html.
func newAWSAEAD(keyURI string, kms kmsiface.KMSAPI) *AWSAEAD {
	return &AWSAEAD{
		keyURI: keyURI,
		kms:    kms,
	}
}

// Encrypt AEAD encrypts the plaintext data and uses addtionaldata from authentication.
func (a *AWSAEAD) Encrypt(plaintext, additionalData []byte) ([]byte, error) {
	ad := hex.EncodeToString(additionalData)
	req := &kms.EncryptInput{
		KeyId:             aws.String(a.keyURI),
		Plaintext:         plaintext,
		EncryptionContext: map[string]*string{"additionalData": &ad},
	}
	if ad == "" {
		req = &kms.EncryptInput{
			KeyId:     aws.String(a.keyURI),
			Plaintext: plaintext,
		}
	}
	resp, err := a.kms.Encrypt(req)
	if err != nil {
		return nil, err
	}

	return resp.CiphertextBlob, nil
}

// Decrypt AEAD decrypts the data and verified the additional data.
//
// Returns an error if the KeyId field in the response does not match the KeyURI
// provided when creating the client. If we don't do this, the possibility exists
// for the ciphertext to be replaced by one under a key we don't control/expect,
// but do have decrypt permissions on.
//
// This check is disabled if AWESAEAD.keyURI is not in key ARN format.
//
// See https://docs.aws.amazon.com/kms/latest/developerguide/concepts.html#key-id.
func (a *AWSAEAD) Decrypt(ciphertext, additionalData []byte) ([]byte, error) {
	ad := hex.EncodeToString(additionalData)
	req := &kms.DecryptInput{
		KeyId:             aws.String(a.keyURI),
		CiphertextBlob:    ciphertext,
		EncryptionContext: map[string]*string{"additionalData": &ad},
	}
	if ad == "" {
		req = &kms.DecryptInput{
			CiphertextBlob: ciphertext,
		}
	}
	resp, err := a.kms.Decrypt(req)
	if err != nil {
		return nil, err
	}
	if isKeyArnFormat(a.keyURI) && strings.Compare(*resp.KeyId, a.keyURI) != 0 {
		return nil, errors.New("decryption failed: wrong key id")
	}
	return resp.Plaintext, nil
}

// isKeyArnFormat returns true if the keyURI is the KMS Key ARN format; false otherwise.
func isKeyArnFormat(keyURI string) bool {
	tokens := strings.Split(keyURI, ":")
	return len(tokens) == 6 && strings.HasPrefix(tokens[5], "key/")
}
<|endoftext|>