<commit_before>package gorethink

import (
	"encoding/json"
	"errors"
	"reflect"

	"github.com/dancannon/gorethink/encoding"
	p "github.com/dancannon/gorethink/ql2"
)

var (
	errCursorClosed = errors.New("connection closed, cannot read cursor")
)

func newCursor(conn *Connection, cursorType string, token int64, term *Term, opts map[string]interface{}) *Cursor {
	if cursorType == "" {
		cursorType = "Cursor"
	}

	cursor := &Cursor{
		conn:       conn,
		token:      token,
		cursorType: cursorType,
		term:       term,
		opts:       opts,
	}

	return cursor
}

// Cursor is the result of a query. Its cursor starts before the first row
// of the result set. A Cursor is not thread safe and should only be accessed
// by a single goroutine at any given time. Use Next to advance through the
// rows:
//
//     cursor, err := query.Run(session)
//     ...
//     defer cursor.Close()
//
//     var response interface{}
//     for cursor.Next(&response) {
//         ...
//     }
//     err = cursor.Err() // get any error encountered during iteration
//     ...
type Cursor struct {
	releaseConn func(error)

	conn       *Connection
	token      int64
	cursorType string
	term       *Term
	opts       map[string]interface{}

	lastErr   error
	fetching  bool
	closed    int32
	finished  bool
	isAtom    bool
	buffer    queue
	responses queue
	profile   interface{}
}

// Profile returns the information returned from the query profiler.
func (c *Cursor) Profile() interface{} {
	return c.profile
}

// Type returns the cursor type (by default "Cursor")
func (c *Cursor) Type() string {
	return c.cursorType
}

// Err returns nil if no errors happened during iteration, or the actual
// error otherwise.
func (c *Cursor) Err() error {
	return c.lastErr
}

// Close closes the cursor, preventing further enumeration. If the end is
// encountered, the cursor is closed automatically. Close is idempotent.
func (c *Cursor) Close() error {
	var err error

	if c.closed != 0 || !atomic.CompareAndSwapInt32(&c.closed, 0, 1) {
		return nil
	}

	conn := c.conn
	if conn == nil {
		return nil
	}
	if conn.conn == nil {
		return nil
	}

	// Stop any unfinished queries
	if c.closed == 0 && !c.finished {
		q := Query{
			Type:  p.Query_STOP,
			Token: c.token,
		}

		_, _, err = conn.Query(q)
	}

	if c.releaseConn != nil {
		c.releaseConn(err)
	}

	c.conn = nil
	c.buffer.elems = nil
	c.responses.elems = nil

	return err
}

// Next retrieves the next document from the result set, blocking if necessary.
// This method will also automatically retrieve another batch of documents from
// the server when the current one is exhausted, or before that in background
// if possible.
//
// Next returns true if a document was successfully unmarshalled onto result,
// and false at the end of the result set or if an error happened.
// When Next returns false, the Err method should be called to verify if
// there was an error during iteration.
//
// Also note that you are able to reuse the same variable multiple times as
// `Next` zeroes the value before scanning in the result.
func (c *Cursor) Next(dest interface{}) bool {
	if c.closed != 0 {
		return false
	}

	hasMore, err := c.loadNext(dest)
	if c.handleError(err) != nil {
		c.Close()
		return false
	}

	if !hasMore {
		c.Close()
	}

	return hasMore
}

func (c *Cursor) loadNext(dest interface{}) (bool, error) {
	for c.lastErr == nil {
		// Check if response is closed/finished
		if c.buffer.Len() == 0 && c.responses.Len() == 0 && c.closed != 0 {
			return false, errCursorClosed
		}

		if c.buffer.Len() == 0 && c.responses.Len() == 0 && !c.finished {
			err := c.fetchMore()
			if err != nil {
				return false, err
			}
		}

		if c.buffer.Len() == 0 && c.responses.Len() == 0 && c.finished {
			return false, nil
		}

		if c.buffer.Len() == 0 && c.responses.Len() > 0 {
			if response, ok := c.responses.Pop().(json.RawMessage); ok {
				var value interface{}
				err := json.Unmarshal(response, &value)
				if err != nil {
					return false, err
				}

				value, err = recursivelyConvertPseudotype(value, c.opts)
				if err != nil {
					return false, err
				}

				// If response is an ATOM then try and convert to an array
				if data, ok := value.([]interface{}); ok && c.isAtom {
					for _, v := range data {
						c.buffer.Push(v)
					}
				} else if value == nil {
					c.buffer.Push(nil)
				} else {
					c.buffer.Push(value)
				}
			}
		}

		if c.buffer.Len() > 0 {
			data := c.buffer.Pop()

			err := encoding.Decode(dest, data)
			if err != nil {
				return false, err
			}

			return true, nil
		}
	}

	return false, c.lastErr
}

// All retrieves all documents from the result set into the provided slice
// and closes the cursor.
//
// The result argument must necessarily be the address for a slice. The slice
// may be nil or previously allocated.
//
// Also note that you are able to reuse the same variable multiple times as
// `All` zeroes the value before scanning in the result. It also attempts
// to reuse the existing slice without allocating any more space by either
// resizing or returning a selection of the slice if necessary.
func (c *Cursor) All(result interface{}) error {
	resultv := reflect.ValueOf(result)
	if resultv.Kind() != reflect.Ptr || resultv.Elem().Kind() != reflect.Slice {
		panic("result argument must be a slice address")
	}
	slicev := resultv.Elem()
	slicev = slicev.Slice(0, slicev.Cap())
	elemt := slicev.Type().Elem()
	i := 0
	for {
		if slicev.Len() == i {
			elemp := reflect.New(elemt)
			if !c.Next(elemp.Interface()) {
				break
			}
			slicev = reflect.Append(slicev, elemp.Elem())
			slicev = slicev.Slice(0, slicev.Cap())
		} else {
			if !c.Next(slicev.Index(i).Addr().Interface()) {
				break
			}
		}
		i++
	}
	resultv.Elem().Set(slicev.Slice(0, i))

	if err := c.Err(); err != nil {
		c.Close()
		return err
	}

	if err := c.Close(); err != nil {
		return err
	}

	return nil
}

// One retrieves a single document from the result set into the provided
// slice and closes the cursor.
//
// Also note that you are able to reuse the same variable multiple times as
// `One` zeroes the value before scanning in the result.
func (c *Cursor) One(result interface{}) error {
	if c.IsNil() {
		c.Close()
		return ErrEmptyResult
	}

	hasResult := c.Next(result)

	if err := c.Err(); err != nil {
		c.Close()
		return err
	}

	if err := c.Close(); err != nil {
		return err
	}

	if !hasResult {
		return ErrEmptyResult
	}

	return nil
}

// Listen listens for rows from the database and sends the result onto the given
// channel. The type that the row is scanned into is determined by the element
// type of the channel.
//
// Also note that this function returns immediately.
//
//     cursor, err := r.Expr([]int{1,2,3}).Run(session)
//     if err != nil {
//         panic(err)
//     }
//
//     ch := make(chan int)
//     cursor.Listen(ch)
//     <- ch // 1
//     <- ch // 2
//     <- ch // 3
func (c *Cursor) Listen(channel interface{}) {
	go func() {
		channelv := reflect.ValueOf(channel)
		if channelv.Kind() != reflect.Chan {
			panic("input argument must be a channel")
		}
		elemt := channelv.Type().Elem()
		for {
			elemp := reflect.New(elemt)
			if !c.Next(elemp.Interface()) {
				break
			}

			channelv.Send(elemp.Elem())
		}

		c.Close()
		channelv.Close()
	}()
}

// IsNil tests if the current row is nil.
func (c *Cursor) IsNil() bool {
	if c.buffer.Len() > 0 {
		bufferedItem := c.buffer.Peek()
		if bufferedItem == nil {
			return true
		}

		return false
	}

	if c.responses.Len() > 0 {
		response := c.responses.Peek()
		if response == nil {
			return true
		}

		if response, ok := response.(json.RawMessage); ok {
			if string(response) == "null" {
				return true
			}
		}

		return false
	}

	return true
}

// fetchMore fetches more rows from the database.
//
// If wait is true then it will wait for the database to reply otherwise it
// will return after sending the continue query.
func (c *Cursor) fetchMore() error {
	var err error
	if !c.fetching {
		c.fetching = true

		if c.closed != 0 {
			return errCursorClosed
		}

		q := Query{
			Type:  p.Query_CONTINUE,
			Token: c.token,
		}

		_, _, err = c.conn.Query(q)
		c.handleError(err)
	}

	return err
}

// handleError sets the value of lastErr to err if lastErr is not yet set.
func (c *Cursor) handleError(err error) error {
	return c.handleErrorLocked(err)
}

// handleError sets the value of lastErr to err if lastErr is not yet set.
func (c *Cursor) handleErrorLocked(err error) error {
	if c.lastErr == nil {
		c.lastErr = err
	}

	return c.lastErr
}

// extend adds the result of a continue query to the cursor.
func (c *Cursor) extend(response *Response) {
	for _, response := range response.Responses {
		c.responses.Push(response)
	}

	c.finished = response.Type != p.Response_SUCCESS_PARTIAL
	c.fetching = false
	c.isAtom = response.Type == p.Response_SUCCESS_ATOM

	putResponse(response)
}

// Queue structure used for storing responses

type queue struct {
	elems               []interface{}
	nelems, popi, pushi int
}

func (q *queue) Len() int {
	if len(q.elems) == 0 {
		return 0
	}

	return q.nelems
}
func (q *queue) Push(elem interface{}) {
	if q.nelems == len(q.elems) {
		q.expand()
	}
	q.elems[q.pushi] = elem
	q.nelems++
	q.pushi = (q.pushi + 1) % len(q.elems)
}
func (q *queue) Pop() (elem interface{}) {
	if q.nelems == 0 {
		return nil
	}
	elem = q.elems[q.popi]
	q.elems[q.popi] = nil // Help GC.
	q.nelems--
	q.popi = (q.popi + 1) % len(q.elems)
	return elem
}
func (q *queue) Peek() (elem interface{}) {
	if q.nelems == 0 {
		return nil
	}
	return q.elems[q.popi]
}
func (q *queue) expand() {
	curcap := len(q.elems)
	var newcap int
	if curcap == 0 {
		newcap = 8
	} else if curcap < 1024 {
		newcap = curcap * 2
	} else {
		newcap = curcap + (curcap / 4)
	}
	elems := make([]interface{}, newcap)
	if q.popi == 0 {
		copy(elems, q.elems)
		q.pushi = curcap
	} else {
		newpopi := newcap - (curcap - q.popi)
		copy(elems, q.elems[:q.popi])
		copy(elems[newpopi:], q.elems[q.popi:])
		q.popi = newpopi
	}
	for i := range q.elems {
		q.elems[i] = nil // Help GC.
	}
	q.elems = elems
}
<commit_msg>Added missing import<commit_after>package gorethink

import (
	"encoding/json"
	"errors"
	"reflect"
	"sync/atomic"

	"github.com/dancannon/gorethink/encoding"
	p "github.com/dancannon/gorethink/ql2"
)

var (
	errCursorClosed = errors.New("connection closed, cannot read cursor")
)

func newCursor(conn *Connection, cursorType string, token int64, term *Term, opts map[string]interface{}) *Cursor {
	if cursorType == "" {
		cursorType = "Cursor"
	}

	cursor := &Cursor{
		conn:       conn,
		token:      token,
		cursorType: cursorType,
		term:       term,
		opts:       opts,
	}

	return cursor
}

// Cursor is the result of a query. Its cursor starts before the first row
// of the result set. A Cursor is not thread safe and should only be accessed
// by a single goroutine at any given time. Use Next to advance through the
// rows:
//
//     cursor, err := query.Run(session)
//     ...
//     defer cursor.Close()
//
//     var response interface{}
//     for cursor.Next(&response) {
//         ...
//     }
//     err = cursor.Err() // get any error encountered during iteration
//     ...
type Cursor struct {
	releaseConn func(error)

	conn       *Connection
	token      int64
	cursorType string
	term       *Term
	opts       map[string]interface{}

	lastErr   error
	fetching  bool
	closed    int32
	finished  bool
	isAtom    bool
	buffer    queue
	responses queue
	profile   interface{}
}

// Profile returns the information returned from the query profiler.
func (c *Cursor) Profile() interface{} {
	return c.profile
}

// Type returns the cursor type (by default "Cursor")
func (c *Cursor) Type() string {
	return c.cursorType
}

// Err returns nil if no errors happened during iteration, or the actual
// error otherwise.
func (c *Cursor) Err() error {
	return c.lastErr
}

// Close closes the cursor, preventing further enumeration. If the end is
// encountered, the cursor is closed automatically. Close is idempotent.
func (c *Cursor) Close() error {
	var err error

	if c.closed != 0 || !atomic.CompareAndSwapInt32(&c.closed, 0, 1) {
		return nil
	}

	conn := c.conn
	if conn == nil {
		return nil
	}
	if conn.conn == nil {
		return nil
	}

	// Stop any unfinished queries
	if c.closed == 0 && !c.finished {
		q := Query{
			Type:  p.Query_STOP,
			Token: c.token,
		}

		_, _, err = conn.Query(q)
	}

	if c.releaseConn != nil {
		c.releaseConn(err)
	}

	c.conn = nil
	c.buffer.elems = nil
	c.responses.elems = nil

	return err
}

// Next retrieves the next document from the result set, blocking if necessary.
// This method will also automatically retrieve another batch of documents from
// the server when the current one is exhausted, or before that in background
// if possible.
//
// Next returns true if a document was successfully unmarshalled onto result,
// and false at the end of the result set or if an error happened.
// When Next returns false, the Err method should be called to verify if
// there was an error during iteration.
//
// Also note that you are able to reuse the same variable multiple times as
// `Next` zeroes the value before scanning in the result.
func (c *Cursor) Next(dest interface{}) bool {
	if c.closed != 0 {
		return false
	}

	hasMore, err := c.loadNext(dest)
	if c.handleError(err) != nil {
		c.Close()
		return false
	}

	if !hasMore {
		c.Close()
	}

	return hasMore
}

func (c *Cursor) loadNext(dest interface{}) (bool, error) {
	for c.lastErr == nil {
		// Check if response is closed/finished
		if c.buffer.Len() == 0 && c.responses.Len() == 0 && c.closed != 0 {
			return false, errCursorClosed
		}

		if c.buffer.Len() == 0 && c.responses.Len() == 0 && !c.finished {
			err := c.fetchMore()
			if err != nil {
				return false, err
			}
		}

		if c.buffer.Len() == 0 && c.responses.Len() == 0 && c.finished {
			return false, nil
		}

		if c.buffer.Len() == 0 && c.responses.Len() > 0 {
			if response, ok := c.responses.Pop().(json.RawMessage); ok {
				var value interface{}
				err := json.Unmarshal(response, &value)
				if err != nil {
					return false, err
				}

				value, err = recursivelyConvertPseudotype(value, c.opts)
				if err != nil {
					return false, err
				}

				// If response is an ATOM then try and convert to an array
				if data, ok := value.([]interface{}); ok && c.isAtom {
					for _, v := range data {
						c.buffer.Push(v)
					}
				} else if value == nil {
					c.buffer.Push(nil)
				} else {
					c.buffer.Push(value)
				}
			}
		}

		if c.buffer.Len() > 0 {
			data := c.buffer.Pop()

			err := encoding.Decode(dest, data)
			if err != nil {
				return false, err
			}

			return true, nil
		}
	}

	return false, c.lastErr
}

// All retrieves all documents from the result set into the provided slice
// and closes the cursor.
//
// The result argument must necessarily be the address for a slice. The slice
// may be nil or previously allocated.
//
// Also note that you are able to reuse the same variable multiple times as
// `All` zeroes the value before scanning in the result. It also attempts
// to reuse the existing slice without allocating any more space by either
// resizing or returning a selection of the slice if necessary.
func (c *Cursor) All(result interface{}) error {
	resultv := reflect.ValueOf(result)
	if resultv.Kind() != reflect.Ptr || resultv.Elem().Kind() != reflect.Slice {
		panic("result argument must be a slice address")
	}
	slicev := resultv.Elem()
	slicev = slicev.Slice(0, slicev.Cap())
	elemt := slicev.Type().Elem()
	i := 0
	for {
		if slicev.Len() == i {
			elemp := reflect.New(elemt)
			if !c.Next(elemp.Interface()) {
				break
			}
			slicev = reflect.Append(slicev, elemp.Elem())
			slicev = slicev.Slice(0, slicev.Cap())
		} else {
			if !c.Next(slicev.Index(i).Addr().Interface()) {
				break
			}
		}
		i++
	}
	resultv.Elem().Set(slicev.Slice(0, i))

	if err := c.Err(); err != nil {
		c.Close()
		return err
	}

	if err := c.Close(); err != nil {
		return err
	}

	return nil
}

// One retrieves a single document from the result set into the provided
// slice and closes the cursor.
//
// Also note that you are able to reuse the same variable multiple times as
// `One` zeroes the value before scanning in the result.
func (c *Cursor) One(result interface{}) error {
	if c.IsNil() {
		c.Close()
		return ErrEmptyResult
	}

	hasResult := c.Next(result)

	if err := c.Err(); err != nil {
		c.Close()
		return err
	}

	if err := c.Close(); err != nil {
		return err
	}

	if !hasResult {
		return ErrEmptyResult
	}

	return nil
}

// Listen listens for rows from the database and sends the result onto the given
// channel. The type that the row is scanned into is determined by the element
// type of the channel.
//
// Also note that this function returns immediately.
//
//     cursor, err := r.Expr([]int{1,2,3}).Run(session)
//     if err != nil {
//         panic(err)
//     }
//
//     ch := make(chan int)
//     cursor.Listen(ch)
//     <- ch // 1
//     <- ch // 2
//     <- ch // 3
func (c *Cursor) Listen(channel interface{}) {
	go func() {
		channelv := reflect.ValueOf(channel)
		if channelv.Kind() != reflect.Chan {
			panic("input argument must be a channel")
		}
		elemt := channelv.Type().Elem()
		for {
			elemp := reflect.New(elemt)
			if !c.Next(elemp.Interface()) {
				break
			}

			channelv.Send(elemp.Elem())
		}

		c.Close()
		channelv.Close()
	}()
}

// IsNil tests if the current row is nil.
func (c *Cursor) IsNil() bool {
	if c.buffer.Len() > 0 {
		bufferedItem := c.buffer.Peek()
		if bufferedItem == nil {
			return true
		}

		return false
	}

	if c.responses.Len() > 0 {
		response := c.responses.Peek()
		if response == nil {
			return true
		}

		if response, ok := response.(json.RawMessage); ok {
			if string(response) == "null" {
				return true
			}
		}

		return false
	}

	return true
}

// fetchMore fetches more rows from the database.
//
// If wait is true then it will wait for the database to reply otherwise it
// will return after sending the continue query.
func (c *Cursor) fetchMore() error {
	var err error
	if !c.fetching {
		c.fetching = true

		if c.closed != 0 {
			return errCursorClosed
		}

		q := Query{
			Type:  p.Query_CONTINUE,
			Token: c.token,
		}

		_, _, err = c.conn.Query(q)
		c.handleError(err)
	}

	return err
}

// handleError sets the value of lastErr to err if lastErr is not yet set.
func (c *Cursor) handleError(err error) error {
	return c.handleErrorLocked(err)
}

// handleError sets the value of lastErr to err if lastErr is not yet set.
func (c *Cursor) handleErrorLocked(err error) error {
	if c.lastErr == nil {
		c.lastErr = err
	}

	return c.lastErr
}

// extend adds the result of a continue query to the cursor.
func (c *Cursor) extend(response *Response) {
	for _, response := range response.Responses {
		c.responses.Push(response)
	}

	c.finished = response.Type != p.Response_SUCCESS_PARTIAL
	c.fetching = false
	c.isAtom = response.Type == p.Response_SUCCESS_ATOM

	putResponse(response)
}

// Queue structure used for storing responses

type queue struct {
	elems               []interface{}
	nelems, popi, pushi int
}

func (q *queue) Len() int {
	if len(q.elems) == 0 {
		return 0
	}

	return q.nelems
}
func (q *queue) Push(elem interface{}) {
	if q.nelems == len(q.elems) {
		q.expand()
	}
	q.elems[q.pushi] = elem
	q.nelems++
	q.pushi = (q.pushi + 1) % len(q.elems)
}
func (q *queue) Pop() (elem interface{}) {
	if q.nelems == 0 {
		return nil
	}
	elem = q.elems[q.popi]
	q.elems[q.popi] = nil // Help GC.
	q.nelems--
	q.popi = (q.popi + 1) % len(q.elems)
	return elem
}
func (q *queue) Peek() (elem interface{}) {
	if q.nelems == 0 {
		return nil
	}
	return q.elems[q.popi]
}
func (q *queue) expand() {
	curcap := len(q.elems)
	var newcap int
	if curcap == 0 {
		newcap = 8
	} else if curcap < 1024 {
		newcap = curcap * 2
	} else {
		newcap = curcap + (curcap / 4)
	}
	elems := make([]interface{}, newcap)
	if q.popi == 0 {
		copy(elems, q.elems)
		q.pushi = curcap
	} else {
		newpopi := newcap - (curcap - q.popi)
		copy(elems, q.elems[:q.popi])
		copy(elems[newpopi:], q.elems[q.popi:])
		q.popi = newpopi
	}
	for i := range q.elems {
		q.elems[i] = nil // Help GC.
	}
	q.elems = elems
}
<|endoftext|>


<commit_before>// Package store abstracts the persistent storage used by elvish.
package store

import (
	"database/sql"
	"fmt"
	"net/url"
	"sync"

	"github.com/elves/elvish/util"

	_ "github.com/mattn/go-sqlite3" // enable the "sqlite3" SQL driver
)

var Logger = util.GetLogger("[store] ")
var initDB = map[string](func(*sql.DB) error){}

// Store is the permanent storage backend for elvish.
type Store struct {
	db    *sql.DB
	Waits sync.WaitGroup
}

// DefaultDB returns the default database for storage.
func DefaultDB(dbname string) (*sql.DB, error) {
	uri := "file:" + url.QueryEscape(dbname) +
		"?mode=rwc&cache=shared&vfs=unix-dotfile"
	return sql.Open("sqlite3", uri)
}

// NewStore creates a new Store with the default database.
func NewStore(dbname string) (*Store, error) {
	db, err := DefaultDB(dbname)
	if err != nil {
		return nil, err
	}
	return NewStoreDB(db)
}

// NewStoreDB creates a new Store with a custom database. The database must be
// a SQLite database.
func NewStoreDB(db *sql.DB) (*Store, error) {
	Logger.Println("initializing store")
	defer Logger.Println("initialized store")
	st := &Store{db, sync.WaitGroup{}}

	if SchemaUpToDate(db) {
		Logger.Println("DB schema up to date")
	} else {
		for name, fn := range initDB {
			err := fn(db)
			if err != nil {
				return nil, fmt.Errorf("failed to %s: %v", name, err)
			}
		}
	}

	return st, nil
}

func (s *Store) Close() error {
	if s == nil || s.db == nil {
		return nil
	}
	s.Waits.Wait()
	return s.db.Close()
}
<commit_msg>Use only one sqlite connection.<commit_after>// Package store abstracts the persistent storage used by elvish.
package store

import (
	"database/sql"
	"fmt"
	"net/url"
	"sync"

	"github.com/elves/elvish/util"

	_ "github.com/mattn/go-sqlite3" // enable the "sqlite3" SQL driver
)

var Logger = util.GetLogger("[store] ")
var initDB = map[string](func(*sql.DB) error){}

// Store is the permanent storage backend for elvish.
type Store struct {
	db    *sql.DB
	Waits sync.WaitGroup
}

// DefaultDB returns the default database for storage.
func DefaultDB(dbname string) (*sql.DB, error) {
	uri := "file:" + url.QueryEscape(dbname) +
		"?mode=rwc&cache=shared&vfs=unix-dotfile"
	db, err := sql.Open("sqlite3", uri)
	if err == nil {
		db.SetMaxOpenConns(1)
	}
	return db, err
}

// NewStore creates a new Store with the default database.
func NewStore(dbname string) (*Store, error) {
	db, err := DefaultDB(dbname)
	if err != nil {
		return nil, err
	}
	return NewStoreDB(db)
}

// NewStoreDB creates a new Store with a custom database. The database must be
// a SQLite database.
func NewStoreDB(db *sql.DB) (*Store, error) {
	Logger.Println("initializing store")
	defer Logger.Println("initialized store")
	st := &Store{db, sync.WaitGroup{}}

	if SchemaUpToDate(db) {
		Logger.Println("DB schema up to date")
	} else {
		for name, fn := range initDB {
			err := fn(db)
			if err != nil {
				return nil, fmt.Errorf("failed to %s: %v", name, err)
			}
		}
	}

	return st, nil
}

func (s *Store) Close() error {
	if s == nil || s.db == nil {
		return nil
	}
	s.Waits.Wait()
	return s.db.Close()
}
<|endoftext|>


<commit_before>package store

import (
	"database/sql"
	"log"
	"time"

	// mysql driver
	_ "github.com/go-sql-driver/mysql"
	"github.com/kusubooru/shimmie"
)

type datastore struct {
	*sql.DB
}

func (db datastore) SQLDB() *sql.DB {
	return db.DB
}

// Open creates a database connection for the given driver and configuration.
func Open(driver, config string) shimmie.Store {
	db := openDB(driver, config)
	return &datastore{db}
}

// openDB opens a new database connection with the specified driver and
// connection string.
func openDB(driver, config string) *sql.DB {
	db, err := sql.Open(driver, config)
	if err != nil {
		log.Fatalln("database connection failed:", err)
	}
	if driver == "mysql" {
		// per issue https://github.com/go-sql-driver/mysql/issues/257
		db.SetMaxIdleConns(0)
	}
	if err := pingDatabase(db); err != nil {
		log.Fatalln("database ping attempts failed:", err)
	}
	return db
}

// helper function to ping the database with backoff to ensure a connection can
// be established before we proceed.
func pingDatabase(db *sql.DB) (err error) {
	for i := 0; i < 10; i++ {
		err = db.Ping()
		if err == nil {
			return
		}
		log.Print("database ping failed. retry in 1s")
		time.Sleep(time.Second)
	}
	return
}

func (db datastore) Close() error {
	return db.DB.Close()
}
<commit_msg>Add Tx helper function to reduce Transaction logic<commit_after>package store

import (
	"database/sql"
	"fmt"
	"log"
	"time"

	// mysql driver
	_ "github.com/go-sql-driver/mysql"
	"github.com/kusubooru/shimmie"
)

type datastore struct {
	*sql.DB
}

func (db datastore) SQLDB() *sql.DB {
	return db.DB
}

// Open creates a database connection for the given driver and configuration.
func Open(driver, config string) shimmie.Store {
	db := openDB(driver, config)
	return &datastore{db}
}

// openDB opens a new database connection with the specified driver and
// connection string.
func openDB(driver, config string) *sql.DB {
	db, err := sql.Open(driver, config)
	if err != nil {
		log.Fatalln("database connection failed:", err)
	}
	if driver == "mysql" {
		// per issue https://github.com/go-sql-driver/mysql/issues/257
		db.SetMaxIdleConns(0)
	}
	if err := pingDatabase(db); err != nil {
		log.Fatalln("database ping attempts failed:", err)
	}
	return db
}

// helper function to ping the database with backoff to ensure a connection can
// be established before we proceed.
func pingDatabase(db *sql.DB) (err error) {
	for i := 0; i < 10; i++ {
		err = db.Ping()
		if err == nil {
			return
		}
		log.Print("database ping failed. retry in 1s")
		time.Sleep(time.Second)
	}
	return
}

func (db datastore) Close() error {
	return db.DB.Close()
}

// Tx allows to perform a function in a transaction. It detects error and panic
// and in that case it rollbacks otherwise it commits the transaction.
func Tx(db *sql.DB, txFunc func(*sql.Tx) error) (err error) {
	tx, err := db.Begin()
	if err != nil {
		return
	}
	defer func() {
		if p := recover(); p != nil {
			switch p := p.(type) {
			case error:
				err = p
			default:
				err = fmt.Errorf("%s", p)
			}
		}
		if err != nil {
			tx.Rollback()
			return
		}
		err = tx.Commit()
	}()
	return txFunc(tx)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"github.com/howeyc/fsnotify"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"syscall"
	"time"
)

// Milliseconds to wait for the next job to begin after a file change
const WorkDelay = 900

// Default pattern to match files which trigger a build
const FilePattern = `(.+\.go|.+\.c)$`

var (
	flag_directory = flag.String("directory", ".", "Directory to watch for changes")
	flag_pattern   = flag.String("pattern", FilePattern, "Pattern of watched files")
	flag_command   = flag.String("command", "", "Command to run and restart after build")
	flag_recursive = flag.Bool("recursive", true, "Watch all dirs. recursively")
	flag_build     = flag.String("build", "go build", "Command to rebuild after changes")
)

// Run `go build` and print the output if something's gone wrong.
func build() bool {
	log.Println("Running build command!")

	args := strings.Split(*flag_build, " ")
	if len(args) == 0 {
		// If the user has specified and empty then we are done.
		return true
	}

	cmd := exec.Command(args[0], args[1:]...)

	cmd.Dir = *flag_directory

	output, err := cmd.CombinedOutput()

	if err == nil {
		log.Println("Build ok.")
	} else {
		log.Println("Error while building:\n", string(output))
	}

	return err == nil
}

func matchesPattern(pattern *regexp.Regexp, file string) bool {
	return pattern.MatchString(file)
}

// Accept build jobs and start building when there are no jobs rushing in.
// The inrush protection is WorkDelay milliseconds long, in this period
// every incoming job will reset the timer.
func builder(jobs <-chan string, buildDone chan<- bool) {
	createThreshold := func() <-chan time.Time {
		return time.After(time.Duration(WorkDelay * time.Millisecond))
	}

	threshold := createThreshold()

	for {
		select {
		case <-jobs:
			threshold = createThreshold()
		case <-threshold:
			if build() {
				buildDone <- true
			}
		}
	}
}

func logger(stdoutChan <-chan io.ReadCloser) {
	dumper := func(pipe io.ReadCloser, prefix string) {
		reader := bufio.NewReader(pipe)

	readloop:
		for {
			line, err := reader.ReadString('\n')

			if err != nil {
				break readloop
			}

			log.Print(prefix, " ", line)
		}
	}

	for {
		pipe := <-stdoutChan

		go dumper(pipe, "stdout:")

		pipe = <-stdoutChan

		go dumper(pipe, "stderr:")
	}
}

// Run the command in the given string and restart it after
// a message was received on the buildDone channel.
func runner(command string, buildDone chan bool) {
	var currentProcess *os.Process

	stdoutChan := make(chan io.ReadCloser)

	go logger(stdoutChan)

	for {
		<-buildDone

		args := strings.Split(command, " ")
		cmd := exec.Command(args[0], args[1:]...)

		if currentProcess != nil {
			err := currentProcess.Kill()

			if err != nil {
				log.Fatal("Could not kill child process. Aborting due to danger of infinite forks.")
			}
		}

		log.Println("Restarting the given command.")

		pipe, err := cmd.StdoutPipe()

		if err != nil {
			log.Fatal("Can't get stdout pipe for command:", err)
		}

		stdoutChan <- pipe

		pipe, err = cmd.StderrPipe()

		if err != nil {
			log.Fatal("Can't get stderr pipe for command:", err)
		}

		stdoutChan <- pipe

		err = cmd.Start()

		if err != nil {
			log.Println("Error while running command:", err)
		}

		currentProcess = cmd.Process
	}
}

func flusher(buildDone <-chan bool) {
	for {
		<-buildDone
	}
}

func main() {
	flag.Parse()

	if *flag_directory == "" {
		fmt.Fprintf(os.Stderr, "-directory=... is required.\n")
		os.Exit(1)
	}

	watcher, err := fsnotify.NewWatcher()

	if err != nil {
		log.Fatal(err)
	}

	defer watcher.Close()

	if *flag_recursive == true {
		err = filepath.Walk(*flag_directory, func(path string, info os.FileInfo, err error) error {
			if err == nil && info.IsDir() {
				return watcher.Watch(path)
			}
			return err
		})

		if err != nil {
			log.Fatal("filepath.Walk():", err)
		}

	} else {
		err := watcher.Watch(*flag_directory)

		if err != nil {
			log.Fatal("watcher.Watch():", err)
		}
	}

	pattern := regexp.MustCompile(*flag_pattern)
	jobs := make(chan string)
	buildDone := make(chan bool)

	go builder(jobs, buildDone)

	if *flag_command != "" {
		go runner(*flag_command, buildDone)
	} else {
		go flusher(buildDone)
	}

	for {
		select {
		case ev := <-watcher.Event:
			if ev.Name != "" && matchesPattern(pattern, ev.Name) {
				jobs <- ev.Name
			}

		case err := <-watcher.Error:
			if v, ok := err.(*os.SyscallError); ok {
				if v.Err == syscall.EINTR {
					continue
				}
				log.Fatal("watcher.Error: SyscallError:", v)
			}
			log.Fatal("watcher.Error:", err)
		}
	}
}
<commit_msg>stdoutChan was not an appropriate name<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"github.com/howeyc/fsnotify"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"syscall"
	"time"
)

// Milliseconds to wait for the next job to begin after a file change
const WorkDelay = 900

// Default pattern to match files which trigger a build
const FilePattern = `(.+\.go|.+\.c)$`

var (
	flag_directory = flag.String("directory", ".", "Directory to watch for changes")
	flag_pattern   = flag.String("pattern", FilePattern, "Pattern of watched files")
	flag_command   = flag.String("command", "", "Command to run and restart after build")
	flag_recursive = flag.Bool("recursive", true, "Watch all dirs. recursively")
	flag_build     = flag.String("build", "go build", "Command to rebuild after changes")
)

// Run `go build` and print the output if something's gone wrong.
func build() bool {
	log.Println("Running build command!")

	args := strings.Split(*flag_build, " ")
	if len(args) == 0 {
		// If the user has specified and empty then we are done.
		return true
	}

	cmd := exec.Command(args[0], args[1:]...)

	cmd.Dir = *flag_directory

	output, err := cmd.CombinedOutput()

	if err == nil {
		log.Println("Build ok.")
	} else {
		log.Println("Error while building:\n", string(output))
	}

	return err == nil
}

func matchesPattern(pattern *regexp.Regexp, file string) bool {
	return pattern.MatchString(file)
}

// Accept build jobs and start building when there are no jobs rushing in.
// The inrush protection is WorkDelay milliseconds long, in this period
// every incoming job will reset the timer.
func builder(jobs <-chan string, buildDone chan<- bool) {
	createThreshold := func() <-chan time.Time {
		return time.After(time.Duration(WorkDelay * time.Millisecond))
	}

	threshold := createThreshold()

	for {
		select {
		case <-jobs:
			threshold = createThreshold()
		case <-threshold:
			if build() {
				buildDone <- true
			}
		}
	}
}

func logger(pipeChan <-chan io.ReadCloser) {
	dumper := func(pipe io.ReadCloser, prefix string) {
		reader := bufio.NewReader(pipe)

	readloop:
		for {
			line, err := reader.ReadString('\n')

			if err != nil {
				break readloop
			}

			log.Print(prefix, " ", line)
		}
	}

	for {
		pipe := <-pipeChan
		go dumper(pipe, "stdout:")

		pipe = <-pipeChan
		go dumper(pipe, "stderr:")
	}
}

// Run the command in the given string and restart it after
// a message was received on the buildDone channel.
func runner(command string, buildDone chan bool) {
	var currentProcess *os.Process

	pipeChan := make(chan io.ReadCloser)

	go logger(pipeChan)

	for {
		<-buildDone

		args := strings.Split(command, " ")
		cmd := exec.Command(args[0], args[1:]...)

		if currentProcess != nil {
			err := currentProcess.Kill()

			if err != nil {
				log.Fatal("Could not kill child process. Aborting due to danger of infinite forks.")
			}
		}

		log.Println("Restarting the given command.")

		pipe, err := cmd.StdoutPipe()

		if err != nil {
			log.Fatal("Can't get stdout pipe for command:", err)
		}

		pipeChan <- pipe

		pipe, err = cmd.StderrPipe()

		if err != nil {
			log.Fatal("Can't get stderr pipe for command:", err)
		}

		pipeChan <- pipe

		err = cmd.Start()

		if err != nil {
			log.Println("Error while running command:", err)
		}

		currentProcess = cmd.Process
	}
}

func flusher(buildDone <-chan bool) {
	for {
		<-buildDone
	}
}

func main() {
	flag.Parse()

	if *flag_directory == "" {
		fmt.Fprintf(os.Stderr, "-directory=... is required.\n")
		os.Exit(1)
	}

	watcher, err := fsnotify.NewWatcher()

	if err != nil {
		log.Fatal(err)
	}

	defer watcher.Close()

	if *flag_recursive == true {
		err = filepath.Walk(*flag_directory, func(path string, info os.FileInfo, err error) error {
			if err == nil && info.IsDir() {
				return watcher.Watch(path)
			}
			return err
		})

		if err != nil {
			log.Fatal("filepath.Walk():", err)
		}

	} else {
		err := watcher.Watch(*flag_directory)

		if err != nil {
			log.Fatal("watcher.Watch():", err)
		}
	}

	pattern := regexp.MustCompile(*flag_pattern)
	jobs := make(chan string)
	buildDone := make(chan bool)

	go builder(jobs, buildDone)

	if *flag_command != "" {
		go runner(*flag_command, buildDone)
	} else {
		go flusher(buildDone)
	}

	for {
		select {
		case ev := <-watcher.Event:
			if ev.Name != "" && matchesPattern(pattern, ev.Name) {
				jobs <- ev.Name
			}

		case err := <-watcher.Error:
			if v, ok := err.(*os.SyscallError); ok {
				if v.Err == syscall.EINTR {
					continue
				}
				log.Fatal("watcher.Error: SyscallError:", v)
			}
			log.Fatal("watcher.Error:", err)
		}
	}
}
<|endoftext|>


<commit_before>// Package store provides a simple distributed key-value store. The keys and
// associated values are changed via distributed consensus, meaning that the
// values are changed only when a majority of nodes in the cluster agree on
// the new value.
//
// Distributed consensus is provided via the Raft algorithm.
package store

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/hashicorp/raft"
	"github.com/hashicorp/raft-boltdb"
)

const (
	retainSnapshotCount = 2
	raftTimeout         = 10 * time.Second
)

type command struct {
	Op    string `json:"op,omitempty"`
	Key   string `json:"key,omitempty"`
	Value string `json:"value,omitempty"`
}

// Store is a simple key-value store, where all changes are made via Raft consensus.
type Store struct {
	RaftDir  string
	RaftBind string

	mu sync.Mutex
	m  map[string]string // The key-value store for the system.

	raft *raft.Raft // The consensus mechanism

	logger *log.Logger
}

// New returns a new Store.
func New() *Store {
	return &Store{
		m:      make(map[string]string),
		logger: log.New(os.Stderr, "[store] ", log.LstdFlags),
	}
}

// Open opens the store. If enableSingle is set, and there are no existing peers,
// then this node becomes the first node, and therefore leader, of the cluster.
func (s *Store) Open(enableSingle bool) error {
	// Setup Raft configuration.
	config := raft.DefaultConfig()

	// Check for any existing peers.
	peers, err := readPeersJSON(filepath.Join(s.RaftDir, "peers.json"))
	if err != nil {
		return err
	}

	// Allow the node to entry single-mode, potentially electing itself, if
	// explicitly enabled and there is only 1 node in the cluster already.
	if enableSingle && len(peers) <= 1 {
		s.logger.Println("enabling single-node mode")
		config.EnableSingleNode = true
		config.DisableBootstrapAfterElect = false
	}

	// Setup Raft communication.
	addr, err := net.ResolveTCPAddr("tcp", s.RaftBind)
	if err != nil {
		return err
	}
	transport, err := raft.NewTCPTransport(s.RaftBind, addr, 3, 10*time.Second, os.Stderr)
	if err != nil {
		return err
	}

	// Create peer storage.
	peerStore := raft.NewJSONPeers(s.RaftDir, transport)

	// Create the snapshot store. This allows the Raft to truncate the log.
	snapshots, err := raft.NewFileSnapshotStore(s.RaftDir, retainSnapshotCount, os.Stderr)
	if err != nil {
		return fmt.Errorf("file snapshot store: %s", err)
	}

	// Create the log store and stable store.
	logStore, err := raftboltdb.NewBoltStore(filepath.Join(s.RaftDir, "raft.db"))
	if err != nil {
		return fmt.Errorf("new bolt store: %s", err)
	}

	// Instantiate the Raft systems.
	ra, err := raft.NewRaft(config, (*fsm)(s), logStore, logStore, snapshots, peerStore, transport)
	if err != nil {
		return fmt.Errorf("new raft: %s", err)
	}
	s.raft = ra
	return nil
}

// Get returns the value for the given key.
func (s *Store) Get(key string) (string, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.m[key], nil
}

// Set sets the value for the given key.
func (s *Store) Set(key, value string) error {
	if s.raft.State() != raft.Leader {
		return fmt.Errorf("not leader")
	}

	c := &command{
		Op:    "set",
		Key:   key,
		Value: value,
	}
	b, err := json.Marshal(c)
	if err != nil {
		return err
	}

	f := s.raft.Apply(b, raftTimeout)
	return f.Error()
}

// Delete deletes the given key.
func (s *Store) Delete(key string) error {
	if s.raft.State() != raft.Leader {
		return fmt.Errorf("not leader")
	}

	c := &command{
		Op:  "delete",
		Key: key,
	}
	b, err := json.Marshal(c)
	if err != nil {
		return err
	}

	f := s.raft.Apply(b, raftTimeout)
	return f.Error()
}

// Join joins a node, located at addr, to this store. The node must be ready to
// respond to Raft communications at that address.
func (s *Store) Join(addr string) error {
	s.logger.Printf("received join request for remote node as %s", addr)

	f := s.raft.AddPeer(addr)
	if f.Error() != nil {
		return f.Error()
	}
	s.logger.Printf("node at %s joined successfully", addr)
	return nil
}

type fsm Store

// Apply applies a Raft log entry to the key-value store.
func (f *fsm) Apply(l *raft.Log) interface{} {
	var c command
	if err := json.Unmarshal(l.Data, &c); err != nil {
		panic(fmt.Sprintf("failed to unmarshal command: %s", err.Error()))
	}

	switch c.Op {
	case "set":
		return f.applySet(c.Key, c.Value)
	case "delete":
		return f.applyDelete(c.Key)
	default:
		panic(fmt.Sprintf("unrecognized command op: %s", c.Op))
	}
}

// Snapshot returns a snapshot of the key-value store.
func (f *fsm) Snapshot() (raft.FSMSnapshot, error) {
	f.mu.Lock()
	defer f.mu.Unlock()

	// Clone the map.
	o := make(map[string]string)
	for k, v := range f.m {
		o[k] = v
	}
	return &fsmSnapshot{store: o}, nil
}

// Restore stores the key-value store to a previous state.
func (f *fsm) Restore(rc io.ReadCloser) error {
	o := make(map[string]string)
	if err := json.NewDecoder(rc).Decode(&o); err != nil {
		return err
	}

	// Set the state from the snapshot, no lock required according to
	// Hashicorp docs.
	f.m = o
	return nil
}

func (f *fsm) applySet(key, value string) interface{} {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.m[key] = value
	return nil
}

func (f *fsm) applyDelete(key string) interface{} {
	f.mu.Lock()
	defer f.mu.Unlock()
	delete(f.m, key)
	return nil
}

type fsmSnapshot struct {
	store map[string]string
}

func (f *fsmSnapshot) Persist(sink raft.SnapshotSink) error {
	err := func() error {
		// Encode data.
		b, err := json.Marshal(f.store)
		if err != nil {
			return err
		}

		// Write data to sink.
		if _, err := sink.Write(b); err != nil {
			return err
		}

		// Close the sink.
		if err := sink.Close(); err != nil {
			return err
		}

		return nil
	}()

	if err != nil {
		sink.Cancel()
		return err
	}

	return nil
}

func (f *fsmSnapshot) Release() {}

func readPeersJSON(path string) ([]string, error) {
	b, err := ioutil.ReadFile(path)
	if err != nil && !os.IsNotExist(err) {
		return nil, err
	}

	if len(b) == 0 {
		return nil, nil
	}

	var peers []string
	dec := json.NewDecoder(bytes.NewReader(b))
	if err := dec.Decode(&peers); err != nil {
		return nil, err
	}

	return peers, nil
}
<commit_msg>remote store.readPeersJSON<commit_after>// Package store provides a simple distributed key-value store. The keys and
// associated values are changed via distributed consensus, meaning that the
// values are changed only when a majority of nodes in the cluster agree on
// the new value.
//
// Distributed consensus is provided via the Raft algorithm.
package store

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/hashicorp/raft"
	"github.com/hashicorp/raft-boltdb"
)

const (
	retainSnapshotCount = 2
	raftTimeout         = 10 * time.Second
)

type command struct {
	Op    string `json:"op,omitempty"`
	Key   string `json:"key,omitempty"`
	Value string `json:"value,omitempty"`
}

// Store is a simple key-value store, where all changes are made via Raft consensus.
type Store struct {
	RaftDir  string
	RaftBind string

	mu sync.Mutex
	m  map[string]string // The key-value store for the system.

	raft *raft.Raft // The consensus mechanism

	logger *log.Logger
}

// New returns a new Store.
func New() *Store {
	return &Store{
		m:      make(map[string]string),
		logger: log.New(os.Stderr, "[store] ", log.LstdFlags),
	}
}

// Open opens the store. If enableSingle is set, and there are no existing peers,
// then this node becomes the first node, and therefore leader, of the cluster.
func (s *Store) Open(enableSingle bool) error {
	// Setup Raft configuration.
	config := raft.DefaultConfig()

	// Setup Raft communication.
	addr, err := net.ResolveTCPAddr("tcp", s.RaftBind)
	if err != nil {
		return err
	}
	transport, err := raft.NewTCPTransport(s.RaftBind, addr, 3, 10*time.Second, os.Stderr)
	if err != nil {
		return err
	}

	// Create peer storage.
	peerStore := raft.NewJSONPeers(s.RaftDir, transport)

	// Check for any existing peers.
	peers, err := peerStore.Peers()
	if err != nil {
		return err
	}

	// Allow the node to entry single-mode, potentially electing itself, if
	// explicitly enabled and there is only 1 node in the cluster already.
	if enableSingle && len(peers) <= 1 {
		s.logger.Println("enabling single-node mode")
		config.EnableSingleNode = true
		config.DisableBootstrapAfterElect = false
	}

	// Create the snapshot store. This allows the Raft to truncate the log.
	snapshots, err := raft.NewFileSnapshotStore(s.RaftDir, retainSnapshotCount, os.Stderr)
	if err != nil {
		return fmt.Errorf("file snapshot store: %s", err)
	}

	// Create the log store and stable store.
	logStore, err := raftboltdb.NewBoltStore(filepath.Join(s.RaftDir, "raft.db"))
	if err != nil {
		return fmt.Errorf("new bolt store: %s", err)
	}

	// Instantiate the Raft systems.
	ra, err := raft.NewRaft(config, (*fsm)(s), logStore, logStore, snapshots, peerStore, transport)
	if err != nil {
		return fmt.Errorf("new raft: %s", err)
	}
	s.raft = ra
	return nil
}

// Get returns the value for the given key.
func (s *Store) Get(key string) (string, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.m[key], nil
}

// Set sets the value for the given key.
func (s *Store) Set(key, value string) error {
	if s.raft.State() != raft.Leader {
		return fmt.Errorf("not leader")
	}

	c := &command{
		Op:    "set",
		Key:   key,
		Value: value,
	}
	b, err := json.Marshal(c)
	if err != nil {
		return err
	}

	f := s.raft.Apply(b, raftTimeout)
	return f.Error()
}

// Delete deletes the given key.
func (s *Store) Delete(key string) error {
	if s.raft.State() != raft.Leader {
		return fmt.Errorf("not leader")
	}

	c := &command{
		Op:  "delete",
		Key: key,
	}
	b, err := json.Marshal(c)
	if err != nil {
		return err
	}

	f := s.raft.Apply(b, raftTimeout)
	return f.Error()
}

// Join joins a node, located at addr, to this store. The node must be ready to
// respond to Raft communications at that address.
func (s *Store) Join(addr string) error {
	s.logger.Printf("received join request for remote node as %s", addr)

	f := s.raft.AddPeer(addr)
	if f.Error() != nil {
		return f.Error()
	}
	s.logger.Printf("node at %s joined successfully", addr)
	return nil
}

type fsm Store

// Apply applies a Raft log entry to the key-value store.
func (f *fsm) Apply(l *raft.Log) interface{} {
	var c command
	if err := json.Unmarshal(l.Data, &c); err != nil {
		panic(fmt.Sprintf("failed to unmarshal command: %s", err.Error()))
	}

	switch c.Op {
	case "set":
		return f.applySet(c.Key, c.Value)
	case "delete":
		return f.applyDelete(c.Key)
	default:
		panic(fmt.Sprintf("unrecognized command op: %s", c.Op))
	}
}

// Snapshot returns a snapshot of the key-value store.
func (f *fsm) Snapshot() (raft.FSMSnapshot, error) {
	f.mu.Lock()
	defer f.mu.Unlock()

	// Clone the map.
	o := make(map[string]string)
	for k, v := range f.m {
		o[k] = v
	}
	return &fsmSnapshot{store: o}, nil
}

// Restore stores the key-value store to a previous state.
func (f *fsm) Restore(rc io.ReadCloser) error {
	o := make(map[string]string)
	if err := json.NewDecoder(rc).Decode(&o); err != nil {
		return err
	}

	// Set the state from the snapshot, no lock required according to
	// Hashicorp docs.
	f.m = o
	return nil
}

func (f *fsm) applySet(key, value string) interface{} {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.m[key] = value
	return nil
}

func (f *fsm) applyDelete(key string) interface{} {
	f.mu.Lock()
	defer f.mu.Unlock()
	delete(f.m, key)
	return nil
}

type fsmSnapshot struct {
	store map[string]string
}

func (f *fsmSnapshot) Persist(sink raft.SnapshotSink) error {
	err := func() error {
		// Encode data.
		b, err := json.Marshal(f.store)
		if err != nil {
			return err
		}

		// Write data to sink.
		if _, err := sink.Write(b); err != nil {
			return err
		}

		// Close the sink.
		if err := sink.Close(); err != nil {
			return err
		}

		return nil
	}()

	if err != nil {
		sink.Cancel()
		return err
	}

	return nil
}

func (f *fsmSnapshot) Release() {}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"

	"github.com/davecgh/go-spew/spew"
	"github.com/gow/dyconf"
)

type daemon struct {
	fileName string
	cm       dyconf.ConfigManager
	log      *log.Logger
}

func (d *daemon) start(logger *log.Logger, fileName string, host string, port string) error {
	d.log = logger
	if fileName == "" {
		return fmt.Errorf("Invalid config file. File path cannot be empty.")
	}
	if !filepath.IsAbs(fileName) {
		return fmt.Errorf("Invalid config file: [%s]. File path must be absolute path.", fileName)
	}
	dir := filepath.Dir(fileName)
	if _, err := os.Stat(dir); os.IsNotExist(err) {
		return fmt.Errorf("Invalid config file: [%s]. The enclosing directory [%s] does not exit.", fileName, dir)
	}

	if _, err := os.Stat(fileName); os.IsNotExist(err) {
		d.log.Printf("[INFO]: [%s] doesn't exist. Going to create a new one.", fileName)
	}

	// Open and close the config file.
	var err error
	d.cm, err = dyconf.NewManager(fileName)
	if err != nil {
		return err
	}
	d.log.Println("Existing map: ", spew.Sdump(d.cm.Map()))
	d.log.Println("Starting the daemon...")
	http.HandleFunc("/config/", d.configServer)
	d.log.Fatal(
		"Failed to start the daemon.",
		http.ListenAndServe(fmt.Sprintf("%s:%s", host, port), nil),
	)
	return nil
}

func (d *daemon) stop() error {
	return d.cm.Close()
}

func (d *daemon) configServer(w http.ResponseWriter, req *http.Request) {
	if req.URL.EscapedPath() != "/config/" {
		w.WriteHeader(http.StatusBadRequest)
		return
	}
	switch req.Method {
	case "PUT":
		d.putConfig(w, req)
	case "GET":
		d.getConfig(w, req)
	default:
		w.WriteHeader(http.StatusMethodNotAllowed)
	}
}

func (d *daemon) putConfig(w http.ResponseWriter, req *http.Request) {
	w.WriteHeader(http.StatusNotImplemented)
}
func (d *daemon) getConfig(w http.ResponseWriter, req *http.Request) {
	keys := req.URL.Query()["key"]
	d.log.Print(keys)
	resp := json.NewEncoder(w)
	m, err := d.cm.Map()
	if err != nil {
		resp.Encode(err)
	}

	type kvpair struct {
		Key string `json:"key"`
		Val []byte `json:"value"`
		Err string `json:"error"`
	}
	var kvPairs []kvpair
	for _, key := range keys {
		pair := kvpair{Key: key}
		val, ok := m[key]
		if !ok {
			pair.Err = "Not Found"
		} else {
			pair.Val = val
		}
		kvPairs = append(kvPairs, pair)
	}
	w.WriteHeader(http.StatusOK)
	resp.Encode(kvPairs)
}

/*
type apiHandler interface {
	handleGet(w http.ResponseWriter, req *http.Request)
	handlePut(w http.ResponseWriter, req *http.Request)
	handlePost(w http.ResponseWriter, req *http.Request)
	handleDelete(w http.ResponseWriter, req *http.Request)
}

type api struct {
	h apiHandler
}

func (a *api) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	switch req.Method {
	case "GET":
		a.h.handleGet(w, req)
	case "PUT":
		a.h.handlePut(w, req)
	case "POST":
		a.h.handlePost(w, req)
	case "DELETE":
		a.h.handleDelete(w, req)
	}
}
*/
<commit_msg>Implemented PUT handler<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"

	"github.com/gow/dyconf"
)

type daemon struct {
	fileName string
	cm       dyconf.ConfigManager
	log      *log.Logger
}

func (d *daemon) start(logger *log.Logger, fileName string, host string, port string) error {
	d.log = logger
	if fileName == "" {
		return fmt.Errorf("Invalid config file. File path cannot be empty.")
	}
	if !filepath.IsAbs(fileName) {
		return fmt.Errorf("Invalid config file: [%s]. File path must be absolute path.", fileName)
	}
	dir := filepath.Dir(fileName)
	if _, err := os.Stat(dir); os.IsNotExist(err) {
		return fmt.Errorf("Invalid config file: [%s]. The enclosing directory [%s] does not exit.", fileName, dir)
	}

	if _, err := os.Stat(fileName); os.IsNotExist(err) {
		d.log.Printf("[INFO]: [%s] doesn't exist. Going to create a new one.", fileName)
	}

	// Open and close the config file.
	var err error
	d.cm, err = dyconf.NewManager(fileName)
	if err != nil {
		return err
	}
	d.log.Println("Starting the daemon...")
	http.HandleFunc("/config/", d.configServer)
	d.log.Fatal(
		"Failed to start the daemon.",
		http.ListenAndServe(fmt.Sprintf("%s:%s", host, port), nil),
	)
	return nil
}

func (d *daemon) stop() error {
	return d.cm.Close()
}

func (d *daemon) configServer(w http.ResponseWriter, req *http.Request) {
	if req.URL.EscapedPath() != "/config/" {
		w.WriteHeader(http.StatusBadRequest)
		return
	}
	switch req.Method {
	case "PUT":
		d.putConfig(w, req)
	case "GET":
		d.getConfig(w, req)
	default:
		w.WriteHeader(http.StatusMethodNotAllowed)
	}
}

func (d *daemon) putConfig(w http.ResponseWriter, req *http.Request) {

	resp := json.NewEncoder(w)
	key := req.FormValue("key")
	if key == "" {
		w.WriteHeader(http.StatusBadRequest)
		resp.Encode(struct{ Error string }{"Invalid Key"})
		return
	}
	val := req.FormValue("value")

	if err := d.cm.Set(key, []byte(val)); err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		resp.Encode(err)
	}
	w.WriteHeader(http.StatusOK)
}

func (d *daemon) getConfig(w http.ResponseWriter, req *http.Request) {
	keys := req.URL.Query()["key"]
	resp := json.NewEncoder(w)
	m, err := d.cm.Map()
	if err != nil {
		resp.Encode(err)
	}

	type kvpair struct {
		Key string `json:"key"`
		Val []byte `json:"value"`
		Err string `json:"error",omitempty`
	}
	var kvPairs []kvpair
	for _, key := range keys {
		pair := kvpair{Key: key}
		val, ok := m[key]
		if !ok {
			pair.Err = "Not Found"
		} else {
			pair.Val = val
		}
		kvPairs = append(kvPairs, pair)
	}
	w.WriteHeader(http.StatusOK)
	resp.Encode(kvPairs)
}
<|endoftext|>


<commit_before>package gateway

import (
	"bytes"
	"fmt"
	"github.com/42wim/matterbridge/bridge"
	"github.com/42wim/matterbridge/bridge/config"
	log "github.com/Sirupsen/logrus"
	//	"github.com/davecgh/go-spew/spew"
	"crypto/sha1"
	"github.com/hashicorp/golang-lru"
	"github.com/peterhellberg/emojilib"
	"net/http"
	"regexp"
	"strings"
	"time"
)

type Gateway struct {
	*config.Config
	Router         *Router
	MyConfig       *config.Gateway
	Bridges        map[string]*bridge.Bridge
	Channels       map[string]*config.ChannelInfo
	ChannelOptions map[string]config.ChannelOptions
	Message        chan config.Message
	Name           string
	Messages       *lru.Cache
}

type BrMsgID struct {
	br *bridge.Bridge
	ID string
}

func New(cfg config.Gateway, r *Router) *Gateway {
	gw := &Gateway{Channels: make(map[string]*config.ChannelInfo), Message: r.Message,
		Router: r, Bridges: make(map[string]*bridge.Bridge), Config: r.Config}
	cache, _ := lru.New(5000)
	gw.Messages = cache
	gw.AddConfig(&cfg)
	return gw
}

func (gw *Gateway) AddBridge(cfg *config.Bridge) error {
	br := gw.Router.getBridge(cfg.Account)
	if br == nil {
		br = bridge.New(gw.Config, cfg, gw.Message)
	}
	gw.mapChannelsToBridge(br)
	gw.Bridges[cfg.Account] = br
	return nil
}

func (gw *Gateway) AddConfig(cfg *config.Gateway) error {
	gw.Name = cfg.Name
	gw.MyConfig = cfg
	gw.mapChannels()
	for _, br := range append(gw.MyConfig.In, append(gw.MyConfig.InOut, gw.MyConfig.Out...)...) {
		err := gw.AddBridge(&br)
		if err != nil {
			return err
		}
	}
	return nil
}

func (gw *Gateway) mapChannelsToBridge(br *bridge.Bridge) {
	for ID, channel := range gw.Channels {
		if br.Account == channel.Account {
			br.Channels[ID] = *channel
		}
	}
}

func (gw *Gateway) reconnectBridge(br *bridge.Bridge) {
	br.Disconnect()
	time.Sleep(time.Second * 5)
RECONNECT:
	log.Infof("Reconnecting %s", br.Account)
	err := br.Connect()
	if err != nil {
		log.Errorf("Reconnection failed: %s. Trying again in 60 seconds", err)
		time.Sleep(time.Second * 60)
		goto RECONNECT
	}
	br.Joined = make(map[string]bool)
	br.JoinChannels()
}

func (gw *Gateway) mapChannelConfig(cfg []config.Bridge, direction string) {
	for _, br := range cfg {
		if isApi(br.Account) {
			br.Channel = "api"
		}
		ID := br.Channel + br.Account
		if _, ok := gw.Channels[ID]; !ok {
			channel := &config.ChannelInfo{Name: br.Channel, Direction: direction, ID: ID, Options: br.Options, Account: br.Account,
				SameChannel: make(map[string]bool)}
			channel.SameChannel[gw.Name] = br.SameChannel
			gw.Channels[channel.ID] = channel
		} else {
			// if we already have a key and it's not our current direction it means we have a bidirectional inout
			if gw.Channels[ID].Direction != direction {
				gw.Channels[ID].Direction = "inout"
			}
		}
		gw.Channels[ID].SameChannel[gw.Name] = br.SameChannel
	}
}

func (gw *Gateway) mapChannels() error {
	gw.mapChannelConfig(gw.MyConfig.In, "in")
	gw.mapChannelConfig(gw.MyConfig.Out, "out")
	gw.mapChannelConfig(gw.MyConfig.InOut, "inout")
	return nil
}

func (gw *Gateway) getDestChannel(msg *config.Message, dest bridge.Bridge) []config.ChannelInfo {
	var channels []config.ChannelInfo
	// if source channel is in only, do nothing
	for _, channel := range gw.Channels {
		// lookup the channel from the message
		if channel.ID == getChannelID(*msg) {
			// we only have destinations if the original message is from an "in" (sending) channel
			if !strings.Contains(channel.Direction, "in") {
				return channels
			}
			continue
		}
	}
	for _, channel := range gw.Channels {
		if _, ok := gw.Channels[getChannelID(*msg)]; !ok {
			continue
		}

		// do samechannelgateway logic
		if channel.SameChannel[msg.Gateway] {
			if msg.Channel == channel.Name && msg.Account != dest.Account {
				channels = append(channels, *channel)
			}
			continue
		}
		if strings.Contains(channel.Direction, "out") && channel.Account == dest.Account && gw.validGatewayDest(msg, channel) {
			channels = append(channels, *channel)
		}
	}
	return channels
}

func (gw *Gateway) handleMessage(msg config.Message, dest *bridge.Bridge) []*BrMsgID {
	var brMsgIDs []*BrMsgID

	// TODO refactor
	// only slack now, check will have to be done in the different bridges.
	// we need to check if we can't use fallback or text in other bridges
	if msg.Extra != nil {
		if dest.Protocol != "discord" &&
			dest.Protocol != "slack" &&
			dest.Protocol != "mattermost" &&
			dest.Protocol != "telegram" &&
			dest.Protocol != "matrix" {
			if msg.Text == "" {
				return brMsgIDs
			}
		}
	}
	// only relay join/part when configged
	if msg.Event == config.EVENT_JOIN_LEAVE && !gw.Bridges[dest.Account].Config.ShowJoinPart {
		return brMsgIDs
	}
	// broadcast to every out channel (irc QUIT)
	if msg.Channel == "" && msg.Event != config.EVENT_JOIN_LEAVE {
		log.Debug("empty channel")
		return brMsgIDs
	}
	originchannel := msg.Channel
	origmsg := msg
	channels := gw.getDestChannel(&msg, *dest)
	for _, channel := range channels {
		// do not send to ourself
		if channel.ID == getChannelID(origmsg) {
			continue
		}
		log.Debugf("Sending %#v from %s (%s) to %s (%s)", msg, msg.Account, originchannel, dest.Account, channel.Name)
		msg.Channel = channel.Name
		msg.Avatar = gw.modifyAvatar(origmsg, dest)
		msg.Username = gw.modifyUsername(origmsg, dest)
		msg.ID = ""
		if res, ok := gw.Messages.Get(origmsg.ID); ok {
			IDs := res.([]*BrMsgID)
			for _, id := range IDs {
				if dest.Protocol == id.br.Protocol {
					msg.ID = id.ID
				}
			}
		}
		// for api we need originchannel as channel
		if dest.Protocol == "api" {
			msg.Channel = originchannel
		}
		mID, err := dest.Send(msg)
		if err != nil {
			fmt.Println(err)
		}
		// append the message ID (mID) from this bridge (dest) to our brMsgIDs slice
		if mID != "" {
			brMsgIDs = append(brMsgIDs, &BrMsgID{dest, mID})
		}
	}
	return brMsgIDs
}

func (gw *Gateway) ignoreMessage(msg *config.Message) bool {
	// if we don't have the bridge, ignore it
	if _, ok := gw.Bridges[msg.Account]; !ok {
		return true
	}
	if msg.Text == "" {
		// we have an attachment or actual bytes
		if msg.Extra != nil && (msg.Extra["attachments"] != nil || len(msg.Extra["file"]) > 0) {
			return false
		}
		log.Debugf("ignoring empty message %#v from %s", msg, msg.Account)
		return true
	}
	for _, entry := range strings.Fields(gw.Bridges[msg.Account].Config.IgnoreNicks) {
		if msg.Username == entry {
			log.Debugf("ignoring %s from %s", msg.Username, msg.Account)
			return true
		}
	}
	// TODO do not compile regexps everytime
	for _, entry := range strings.Fields(gw.Bridges[msg.Account].Config.IgnoreMessages) {
		if entry != "" {
			re, err := regexp.Compile(entry)
			if err != nil {
				log.Errorf("incorrect regexp %s for %s", entry, msg.Account)
				continue
			}
			if re.MatchString(msg.Text) {
				log.Debugf("matching %s. ignoring %s from %s", entry, msg.Text, msg.Account)
				return true
			}
		}
	}
	return false
}

func (gw *Gateway) modifyUsername(msg config.Message, dest *bridge.Bridge) string {
	br := gw.Bridges[msg.Account]
	msg.Protocol = br.Protocol
	if gw.Config.General.StripNick || dest.Config.StripNick {
		re := regexp.MustCompile("[^a-zA-Z0-9]+")
		msg.Username = re.ReplaceAllString(msg.Username, "")
	}
	nick := dest.Config.RemoteNickFormat
	if nick == "" {
		nick = gw.Config.General.RemoteNickFormat
	}

	// loop to replace nicks
	for _, outer := range br.Config.ReplaceNicks {
		search := outer[0]
		replace := outer[1]
		// TODO move compile to bridge init somewhere
		re, err := regexp.Compile(search)
		if err != nil {
			log.Errorf("regexp in %s failed: %s", msg.Account, err)
			break
		}
		msg.Username = re.ReplaceAllString(msg.Username, replace)
	}

	if len(msg.Username) > 0 {
		// fix utf-8 issue #193
		i := 0
		for index := range msg.Username {
			if i == 1 {
				i = index
				break
			}
			i++
		}
		nick = strings.Replace(nick, "{NOPINGNICK}", msg.Username[:i]+""+msg.Username[i:], -1)
	}
	nick = strings.Replace(nick, "{NICK}", msg.Username, -1)
	nick = strings.Replace(nick, "{BRIDGE}", br.Name, -1)
	nick = strings.Replace(nick, "{PROTOCOL}", br.Protocol, -1)
	return nick
}

func (gw *Gateway) modifyAvatar(msg config.Message, dest *bridge.Bridge) string {
	iconurl := gw.Config.General.IconURL
	if iconurl == "" {
		iconurl = dest.Config.IconURL
	}
	iconurl = strings.Replace(iconurl, "{NICK}", msg.Username, -1)
	if msg.Avatar == "" {
		msg.Avatar = iconurl
	}
	return msg.Avatar
}

func (gw *Gateway) modifyMessage(msg *config.Message) {
	// replace :emoji: to unicode
	msg.Text = emojilib.Replace(msg.Text)
	br := gw.Bridges[msg.Account]
	// loop to replace messages
	for _, outer := range br.Config.ReplaceMessages {
		search := outer[0]
		replace := outer[1]
		// TODO move compile to bridge init somewhere
		re, err := regexp.Compile(search)
		if err != nil {
			log.Errorf("regexp in %s failed: %s", msg.Account, err)
			break
		}
		msg.Text = re.ReplaceAllString(msg.Text, replace)
	}
	msg.Gateway = gw.Name
}

func (gw *Gateway) handleFiles(msg *config.Message) {
	if msg.Extra == nil || gw.Config.General.MediaServerUpload == "" {
		return
	}
	if len(msg.Extra["file"]) > 0 {
		client := &http.Client{
			Timeout: time.Second * 5,
		}
		for i, f := range msg.Extra["file"] {
			fi := f.(config.FileInfo)
			sha1sum := fmt.Sprintf("%x", sha1.Sum(*fi.Data))
			reader := bytes.NewReader(*fi.Data)
			url := gw.Config.General.MediaServerUpload + "/" + sha1sum + "/" + fi.Name
			durl := gw.Config.General.MediaServerDownload + "/" + sha1sum + "/" + fi.Name
			extra := msg.Extra["file"][i].(config.FileInfo)
			extra.URL = durl
			msg.Extra["file"][i] = extra
			req, _ := http.NewRequest("PUT", url, reader)
			req.Header.Set("Content-Type", "binary/octet-stream")
			_, err := client.Do(req)
			if err != nil {
				log.Errorf("mediaserver upload failed: %#v", err)
			}
			log.Debugf("mediaserver download URL = %s", durl)
		}
	}
}

func getChannelID(msg config.Message) string {
	return msg.Channel + msg.Account
}

func (gw *Gateway) validGatewayDest(msg *config.Message, channel *config.ChannelInfo) bool {
	return msg.Gateway == gw.Name
}

func isApi(account string) bool {
	return strings.HasPrefix(account, "api.")
}
<commit_msg>Change RemoteNickFormat replacement order. Closes #336<commit_after>package gateway

import (
	"bytes"
	"fmt"
	"github.com/42wim/matterbridge/bridge"
	"github.com/42wim/matterbridge/bridge/config"
	log "github.com/Sirupsen/logrus"
	//	"github.com/davecgh/go-spew/spew"
	"crypto/sha1"
	"github.com/hashicorp/golang-lru"
	"github.com/peterhellberg/emojilib"
	"net/http"
	"regexp"
	"strings"
	"time"
)

type Gateway struct {
	*config.Config
	Router         *Router
	MyConfig       *config.Gateway
	Bridges        map[string]*bridge.Bridge
	Channels       map[string]*config.ChannelInfo
	ChannelOptions map[string]config.ChannelOptions
	Message        chan config.Message
	Name           string
	Messages       *lru.Cache
}

type BrMsgID struct {
	br *bridge.Bridge
	ID string
}

func New(cfg config.Gateway, r *Router) *Gateway {
	gw := &Gateway{Channels: make(map[string]*config.ChannelInfo), Message: r.Message,
		Router: r, Bridges: make(map[string]*bridge.Bridge), Config: r.Config}
	cache, _ := lru.New(5000)
	gw.Messages = cache
	gw.AddConfig(&cfg)
	return gw
}

func (gw *Gateway) AddBridge(cfg *config.Bridge) error {
	br := gw.Router.getBridge(cfg.Account)
	if br == nil {
		br = bridge.New(gw.Config, cfg, gw.Message)
	}
	gw.mapChannelsToBridge(br)
	gw.Bridges[cfg.Account] = br
	return nil
}

func (gw *Gateway) AddConfig(cfg *config.Gateway) error {
	gw.Name = cfg.Name
	gw.MyConfig = cfg
	gw.mapChannels()
	for _, br := range append(gw.MyConfig.In, append(gw.MyConfig.InOut, gw.MyConfig.Out...)...) {
		err := gw.AddBridge(&br)
		if err != nil {
			return err
		}
	}
	return nil
}

func (gw *Gateway) mapChannelsToBridge(br *bridge.Bridge) {
	for ID, channel := range gw.Channels {
		if br.Account == channel.Account {
			br.Channels[ID] = *channel
		}
	}
}

func (gw *Gateway) reconnectBridge(br *bridge.Bridge) {
	br.Disconnect()
	time.Sleep(time.Second * 5)
RECONNECT:
	log.Infof("Reconnecting %s", br.Account)
	err := br.Connect()
	if err != nil {
		log.Errorf("Reconnection failed: %s. Trying again in 60 seconds", err)
		time.Sleep(time.Second * 60)
		goto RECONNECT
	}
	br.Joined = make(map[string]bool)
	br.JoinChannels()
}

func (gw *Gateway) mapChannelConfig(cfg []config.Bridge, direction string) {
	for _, br := range cfg {
		if isApi(br.Account) {
			br.Channel = "api"
		}
		ID := br.Channel + br.Account
		if _, ok := gw.Channels[ID]; !ok {
			channel := &config.ChannelInfo{Name: br.Channel, Direction: direction, ID: ID, Options: br.Options, Account: br.Account,
				SameChannel: make(map[string]bool)}
			channel.SameChannel[gw.Name] = br.SameChannel
			gw.Channels[channel.ID] = channel
		} else {
			// if we already have a key and it's not our current direction it means we have a bidirectional inout
			if gw.Channels[ID].Direction != direction {
				gw.Channels[ID].Direction = "inout"
			}
		}
		gw.Channels[ID].SameChannel[gw.Name] = br.SameChannel
	}
}

func (gw *Gateway) mapChannels() error {
	gw.mapChannelConfig(gw.MyConfig.In, "in")
	gw.mapChannelConfig(gw.MyConfig.Out, "out")
	gw.mapChannelConfig(gw.MyConfig.InOut, "inout")
	return nil
}

func (gw *Gateway) getDestChannel(msg *config.Message, dest bridge.Bridge) []config.ChannelInfo {
	var channels []config.ChannelInfo
	// if source channel is in only, do nothing
	for _, channel := range gw.Channels {
		// lookup the channel from the message
		if channel.ID == getChannelID(*msg) {
			// we only have destinations if the original message is from an "in" (sending) channel
			if !strings.Contains(channel.Direction, "in") {
				return channels
			}
			continue
		}
	}
	for _, channel := range gw.Channels {
		if _, ok := gw.Channels[getChannelID(*msg)]; !ok {
			continue
		}

		// do samechannelgateway logic
		if channel.SameChannel[msg.Gateway] {
			if msg.Channel == channel.Name && msg.Account != dest.Account {
				channels = append(channels, *channel)
			}
			continue
		}
		if strings.Contains(channel.Direction, "out") && channel.Account == dest.Account && gw.validGatewayDest(msg, channel) {
			channels = append(channels, *channel)
		}
	}
	return channels
}

func (gw *Gateway) handleMessage(msg config.Message, dest *bridge.Bridge) []*BrMsgID {
	var brMsgIDs []*BrMsgID

	// TODO refactor
	// only slack now, check will have to be done in the different bridges.
	// we need to check if we can't use fallback or text in other bridges
	if msg.Extra != nil {
		if dest.Protocol != "discord" &&
			dest.Protocol != "slack" &&
			dest.Protocol != "mattermost" &&
			dest.Protocol != "telegram" &&
			dest.Protocol != "matrix" {
			if msg.Text == "" {
				return brMsgIDs
			}
		}
	}
	// only relay join/part when configged
	if msg.Event == config.EVENT_JOIN_LEAVE && !gw.Bridges[dest.Account].Config.ShowJoinPart {
		return brMsgIDs
	}
	// broadcast to every out channel (irc QUIT)
	if msg.Channel == "" && msg.Event != config.EVENT_JOIN_LEAVE {
		log.Debug("empty channel")
		return brMsgIDs
	}
	originchannel := msg.Channel
	origmsg := msg
	channels := gw.getDestChannel(&msg, *dest)
	for _, channel := range channels {
		// do not send to ourself
		if channel.ID == getChannelID(origmsg) {
			continue
		}
		log.Debugf("Sending %#v from %s (%s) to %s (%s)", msg, msg.Account, originchannel, dest.Account, channel.Name)
		msg.Channel = channel.Name
		msg.Avatar = gw.modifyAvatar(origmsg, dest)
		msg.Username = gw.modifyUsername(origmsg, dest)
		msg.ID = ""
		if res, ok := gw.Messages.Get(origmsg.ID); ok {
			IDs := res.([]*BrMsgID)
			for _, id := range IDs {
				if dest.Protocol == id.br.Protocol {
					msg.ID = id.ID
				}
			}
		}
		// for api we need originchannel as channel
		if dest.Protocol == "api" {
			msg.Channel = originchannel
		}
		mID, err := dest.Send(msg)
		if err != nil {
			fmt.Println(err)
		}
		// append the message ID (mID) from this bridge (dest) to our brMsgIDs slice
		if mID != "" {
			brMsgIDs = append(brMsgIDs, &BrMsgID{dest, mID})
		}
	}
	return brMsgIDs
}

func (gw *Gateway) ignoreMessage(msg *config.Message) bool {
	// if we don't have the bridge, ignore it
	if _, ok := gw.Bridges[msg.Account]; !ok {
		return true
	}
	if msg.Text == "" {
		// we have an attachment or actual bytes
		if msg.Extra != nil && (msg.Extra["attachments"] != nil || len(msg.Extra["file"]) > 0) {
			return false
		}
		log.Debugf("ignoring empty message %#v from %s", msg, msg.Account)
		return true
	}
	for _, entry := range strings.Fields(gw.Bridges[msg.Account].Config.IgnoreNicks) {
		if msg.Username == entry {
			log.Debugf("ignoring %s from %s", msg.Username, msg.Account)
			return true
		}
	}
	// TODO do not compile regexps everytime
	for _, entry := range strings.Fields(gw.Bridges[msg.Account].Config.IgnoreMessages) {
		if entry != "" {
			re, err := regexp.Compile(entry)
			if err != nil {
				log.Errorf("incorrect regexp %s for %s", entry, msg.Account)
				continue
			}
			if re.MatchString(msg.Text) {
				log.Debugf("matching %s. ignoring %s from %s", entry, msg.Text, msg.Account)
				return true
			}
		}
	}
	return false
}

func (gw *Gateway) modifyUsername(msg config.Message, dest *bridge.Bridge) string {
	br := gw.Bridges[msg.Account]
	msg.Protocol = br.Protocol
	if gw.Config.General.StripNick || dest.Config.StripNick {
		re := regexp.MustCompile("[^a-zA-Z0-9]+")
		msg.Username = re.ReplaceAllString(msg.Username, "")
	}
	nick := dest.Config.RemoteNickFormat
	if nick == "" {
		nick = gw.Config.General.RemoteNickFormat
	}

	// loop to replace nicks
	for _, outer := range br.Config.ReplaceNicks {
		search := outer[0]
		replace := outer[1]
		// TODO move compile to bridge init somewhere
		re, err := regexp.Compile(search)
		if err != nil {
			log.Errorf("regexp in %s failed: %s", msg.Account, err)
			break
		}
		msg.Username = re.ReplaceAllString(msg.Username, replace)
	}

	if len(msg.Username) > 0 {
		// fix utf-8 issue #193
		i := 0
		for index := range msg.Username {
			if i == 1 {
				i = index
				break
			}
			i++
		}
		nick = strings.Replace(nick, "{NOPINGNICK}", msg.Username[:i]+""+msg.Username[i:], -1)
	}
	nick = strings.Replace(nick, "{BRIDGE}", br.Name, -1)
	nick = strings.Replace(nick, "{PROTOCOL}", br.Protocol, -1)
	nick = strings.Replace(nick, "{NICK}", msg.Username, -1)
	return nick
}

func (gw *Gateway) modifyAvatar(msg config.Message, dest *bridge.Bridge) string {
	iconurl := gw.Config.General.IconURL
	if iconurl == "" {
		iconurl = dest.Config.IconURL
	}
	iconurl = strings.Replace(iconurl, "{NICK}", msg.Username, -1)
	if msg.Avatar == "" {
		msg.Avatar = iconurl
	}
	return msg.Avatar
}

func (gw *Gateway) modifyMessage(msg *config.Message) {
	// replace :emoji: to unicode
	msg.Text = emojilib.Replace(msg.Text)
	br := gw.Bridges[msg.Account]
	// loop to replace messages
	for _, outer := range br.Config.ReplaceMessages {
		search := outer[0]
		replace := outer[1]
		// TODO move compile to bridge init somewhere
		re, err := regexp.Compile(search)
		if err != nil {
			log.Errorf("regexp in %s failed: %s", msg.Account, err)
			break
		}
		msg.Text = re.ReplaceAllString(msg.Text, replace)
	}
	msg.Gateway = gw.Name
}

func (gw *Gateway) handleFiles(msg *config.Message) {
	if msg.Extra == nil || gw.Config.General.MediaServerUpload == "" {
		return
	}
	if len(msg.Extra["file"]) > 0 {
		client := &http.Client{
			Timeout: time.Second * 5,
		}
		for i, f := range msg.Extra["file"] {
			fi := f.(config.FileInfo)
			sha1sum := fmt.Sprintf("%x", sha1.Sum(*fi.Data))
			reader := bytes.NewReader(*fi.Data)
			url := gw.Config.General.MediaServerUpload + "/" + sha1sum + "/" + fi.Name
			durl := gw.Config.General.MediaServerDownload + "/" + sha1sum + "/" + fi.Name
			extra := msg.Extra["file"][i].(config.FileInfo)
			extra.URL = durl
			msg.Extra["file"][i] = extra
			req, _ := http.NewRequest("PUT", url, reader)
			req.Header.Set("Content-Type", "binary/octet-stream")
			_, err := client.Do(req)
			if err != nil {
				log.Errorf("mediaserver upload failed: %#v", err)
			}
			log.Debugf("mediaserver download URL = %s", durl)
		}
	}
}

func getChannelID(msg config.Message) string {
	return msg.Channel + msg.Account
}

func (gw *Gateway) validGatewayDest(msg *config.Message, channel *config.ChannelInfo) bool {
	return msg.Gateway == gw.Name
}

func isApi(account string) bool {
	return strings.HasPrefix(account, "api.")
}
<|endoftext|>


<commit_before>package stream

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"sync"
	"testing"
	"time"
)

var (
	testdata = []byte("hello\nworld\n")
	errFail  = errors.New("fail")
)

type badFs struct {
	readers []File
}
type badFile struct{ name string }

func (r badFile) Name() string                            { return r.name }
func (r badFile) Read(p []byte) (int, error)              { return 0, errFail }
func (r badFile) ReadAt(p []byte, off int64) (int, error) { return 0, errFail }
func (r badFile) Write(p []byte) (int, error)             { return 0, errFail }
func (r badFile) Close() error                            { return errFail }

func (fs badFs) Create(name string) (File, error) { return os.Create(name) }
func (fs badFs) Open(name string) (File, error) {
	if len(fs.readers) > 0 {
		f := fs.readers[len(fs.readers)-1]
		fs.readers = fs.readers[:len(fs.readers)-1]
		return f, nil
	}
	return nil, errFail
}
func (fs badFs) Remove(name string) error { return os.Remove(name) }

func TestMemFs(t *testing.T) {
	fs := NewMemFS()
	if _, err := fs.Open("not found"); err != ErrNotFoundInMem {
		t.Error(err)
		t.FailNow()
	}
}

func TestSingletonFs(t *testing.T) {
	var fs singletonFs
	if _, err := fs.Create(""); err == nil {
		t.Errorf("expected unsupported, got %s", err)
	}
}

func TestBadFile(t *testing.T) {
	fs := badFs{readers: make([]File, 0, 1)}
	fs.readers = append(fs.readers, badFile{name: "test"})
	f, err := NewStream("test", fs)
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	defer cleanup(f, t)
	defer f.Close()

	r, err := f.NextReader()
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	defer r.Close()
	if r.Name() != "test" {
		t.Errorf("expected name to to be 'test' got %s", r.Name())
		t.FailNow()
	}
	if _, err := r.ReadAt(nil, 0); err == nil {
		t.Error("expected ReadAt error")
		t.FailNow()
	}
	if _, err := r.Read(nil); err == nil {
		t.Error("expected Read error")
		t.FailNow()
	}
}

func TestBadFs(t *testing.T) {
	f, err := NewStream("test", badFs{})
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	defer cleanup(f, t)
	defer f.Close()

	r, err := f.NextReader()
	if err == nil {
		t.Error("expected open error")
		t.FailNow()
	} else {
		return
	}
	r.Close()
}

func TestStd(t *testing.T) {
	f, err := New(t.Name() + ".txt")
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	if f.Name() != t.Name()+".txt" {
		t.Errorf("expected name to be %s.txt: %s", t.Name(), f.Name())
	}
	testFile(f, t)
}

func TestMem(t *testing.T) {
	f, err := NewStream(t.Name()+".txt", NewMemFS())
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	f.Write(nil)
	testFile(f, t)
}

func TestCancelClosesAll(t *testing.T) {
	f, err := New(t.Name() + ".txt")
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	for i := 0; i < 10; i++ {
		_, err := f.NextReader() // we won't even grab the reader
		if err != nil {
			t.Error(err)
			t.FailNow()
		}
	}
	f.Cancel()    // Closes all the idle Readers
	cleanup(f, t) // this will deadlock if any arn't Closed
}

func TestCloseUnblocksBlockingRead(t *testing.T) {
	f, err := New(t.Name() + ".txt")
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	r, err := f.NextReader()
	if err != nil {
		t.Error(err)
	}
	go func() {
		_, err := ioutil.ReadAll(r)
		if err == nil || err == io.EOF {
			t.Error("expected an error on a blocking Read for a closed Reader")
		}
	}()
	<-time.After(100 * time.Millisecond) // wait for blocking read
	r.Close()

	<-time.After(100 * time.Millisecond) // wait until read has been unblocked
	f.Close()

	cleanup(f, t) // this will deadlock if any arn't Closed
}

func TestCancelBeforeClose(t *testing.T) {
	f, err := New(t.Name() + ".txt")
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	f.Write([]byte("Hello"))
	r, err := f.NextReader() // blocking reader
	if err != nil {
		t.Error("error creating new reader: ", err)
	}

	wg := sync.WaitGroup{}
	wg.Add(1)
	go func() {
		_, err := ioutil.ReadAll(r)
		if err != ErrCanceled {
			t.Error("Read after cancel should return an error")
		}
		wg.Done()
	}()
	<-time.After(100 * time.Millisecond) // give Reader time to block, this tests it unblocks

	// When canceling writer, reader is closed, so writer unblocks and test passes
	f.Cancel()

	// ReadAt after cancel
	_, err = ioutil.ReadAll(io.NewSectionReader(r, 0, 1))
	if err != ErrCanceled {
		t.Error("ReadAt after cancel should return an error")
	}

	// NextReader should fail as well
	_, err = f.NextReader()
	if err != ErrCanceled {
		t.Error("NextReader should be canceled, but got: ", err)
	}

	n, err := f.Write([]byte("world"))
	// Writer is closed as well
	if err == nil {
		t.Error("expected write after canceling to fail")
	}
	if n != 0 {
		t.Error("expected write after canceling to not write anything")
	}
	wg.Wait()
	cleanup(f, t)
}

func TestCancelAfterClose(t *testing.T) {
	f, err := New(t.Name() + ".txt")
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	f.Write([]byte("Hello"))
	f.Close()

	r, _ := f.NextReader()

	// When canceling writer, reader is unblocked, so writer unblocks and test passes
	f.Cancel()

	wg := sync.WaitGroup{}
	wg.Add(1)
	go func() {
		_, err := ioutil.ReadAll(r)
		if err != ErrCanceled {
			t.Error("Read after canceling (even after close) should fail")
		}
		wg.Done()
	}()

	wg.Wait()

	cleanup(f, t)
}

func TestMemStreams(t *testing.T) {
	testFile(NewMemStream(), t)
}

func TestReadAtWait(t *testing.T) {
	f, err := NewStream(t.Name()+".txt", NewMemFS())
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	r, err := f.NextReader()
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	io.WriteString(f, "hello")
	go func() {
		<-time.After(100 * time.Millisecond)
		io.WriteString(f, " world")
		f.Close()
	}()
	data := make([]byte, 11)
	_, err = r.ReadAt(data, 0)
	if err != nil {
		t.Error(err)
	}
	if string(data) != "hello world" {
		t.Error("expected to read 'hello world' got ", string(data))
	}
}

func TestShutdown(t *testing.T) {
	f, err := NewStream(t.Name()+".txt", NewMemFS())
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	r, err := f.NextReader()
	if err != nil {
		t.Error("should return valid reader before remove")
	}
	io.WriteString(f, "Hello")

	f.shutdownWithErr(nil) // noop
	n, err := r.Read(make([]byte, 1))
	if err != nil || n != 1 {
		t.Error("expected successful read, got %s, bytes=%d.", err, n)
	}

	er := errors.New("shutdown")

	go func() {
		_, err := f.NextReader()
		if err != er {
			t.Errorf("expected %s, got %s", er, err)
		}

		f.Close()
		io.Copy(ioutil.Discard, r)
		r.Close()
	}()

	f.shutdownWithErr(er)
}

func TestRemove(t *testing.T) {
	f, err := NewStream(t.Name()+".txt", NewMemFS())
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	r, err := f.NextReader()
	if err != nil {
		t.Error("should return valid reader before remove")
	}
	fmt.Fprintf(f, "Hello")
	go cleanup(f, t)
	<-time.After(100 * time.Millisecond)
	waitForReadToFinish := make(chan struct{})
	go func() {
		b, err := ioutil.ReadAll(r)
		r.Close()
		if err != nil {
			t.Errorf("reader should continue as if nothing happens after remove, got: %s", err.Error())
		}
		if string(b) != "Hello World" {
			t.Errorf("reader should read all written data aven after remove, got: %s", string(b))
		}
		close(waitForReadToFinish)
	}()
	r2, err := f.NextReader()
	switch err {
	case ErrRemoving:
	case nil:
		t.Error("expected error on NextReader()")
		r2.Close()
	default:
		t.Error("expected diff error on NextReader()", err)
	}

	// It should still succeed as we didn't closed stream yet what causes a remove
	fmt.Fprintf(f, " World")

	f.Close()

	n, err := fmt.Fprintf(f, " War")

	if err == nil {
		t.Errorf("Write to stream after closing should fail, has: %v", err)
	}

	if n != 0 {
		t.Error("expected write after closing to not write anything")
	}
	<-waitForReadToFinish
}

func testFile(f *Stream, t *testing.T) {

	for i := 0; i < 10; i++ {
		go testReader(f, t)
	}

	for i := 0; i < 10; i++ {
		f.Write(testdata[:10])
		<-time.After(10 * time.Millisecond)
		f.Write(testdata[10:])
	}

	if err := f.Close(); err != nil {
		t.Error("expected successful close, got ", err)
	}
	testReader(f, t)
	cleanup(f, t)
}

func testReader(f *Stream, t *testing.T) {
	r, err := f.NextReader()
	if err != nil {
		t.Error(err)
		t.FailNow()
	}

	defer func() {
		err = r.Close()
		if err != nil {
			t.Error("expected successful close, got ", err)
		}

		_, err = ioutil.ReadAll(r)
		if err == nil {
			t.Error("Read after Reader.Close() should return an error")
		}

		_, err = ioutil.ReadAll(io.NewSectionReader(r, 0, 1))
		if err == nil {
			t.Error("ReadAt after Reader.Close() should return an error")
		}
	}()

	buf := bytes.NewBuffer(nil)
	sr := io.NewSectionReader(r, 1+int64(len(testdata)*5), 5)
	io.CopyBuffer(buf, sr, make([]byte, 3)) // force multiple reads
	if !bytes.Equal(buf.Bytes(), testdata[1:6]) {
		t.Errorf("unequal %s", buf.Bytes())
		return
	}

	buf.Reset()
	io.CopyBuffer(buf, r, make([]byte, 3)) // force multiple reads
	if !bytes.Equal(buf.Bytes(), bytes.Repeat(testdata, 10)) {
		t.Errorf("unequal %s", buf.Bytes())
		return
	}
	l, closed := r.Size()
	if !closed {
		t.Errorf("expected writer to be closed, but was open")
	}
	if l != int64(len(testdata)*10) {
		t.Errorf("expected size to be %d but got %d", len(testdata)*10, l)
	}
}

func cleanup(f *Stream, t *testing.T) {
	if err := f.Remove(); err != nil && err != ErrUnsupported {
		t.Error("error while removing file: ", err)
	}
}
<commit_msg>fixing t.Error to t.Errorf<commit_after>package stream

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"sync"
	"testing"
	"time"
)

var (
	testdata = []byte("hello\nworld\n")
	errFail  = errors.New("fail")
)

type badFs struct {
	readers []File
}
type badFile struct{ name string }

func (r badFile) Name() string                            { return r.name }
func (r badFile) Read(p []byte) (int, error)              { return 0, errFail }
func (r badFile) ReadAt(p []byte, off int64) (int, error) { return 0, errFail }
func (r badFile) Write(p []byte) (int, error)             { return 0, errFail }
func (r badFile) Close() error                            { return errFail }

func (fs badFs) Create(name string) (File, error) { return os.Create(name) }
func (fs badFs) Open(name string) (File, error) {
	if len(fs.readers) > 0 {
		f := fs.readers[len(fs.readers)-1]
		fs.readers = fs.readers[:len(fs.readers)-1]
		return f, nil
	}
	return nil, errFail
}
func (fs badFs) Remove(name string) error { return os.Remove(name) }

func TestMemFs(t *testing.T) {
	fs := NewMemFS()
	if _, err := fs.Open("not found"); err != ErrNotFoundInMem {
		t.Error(err)
		t.FailNow()
	}
}

func TestSingletonFs(t *testing.T) {
	var fs singletonFs
	if _, err := fs.Create(""); err == nil {
		t.Errorf("expected unsupported, got %s", err)
	}
}

func TestBadFile(t *testing.T) {
	fs := badFs{readers: make([]File, 0, 1)}
	fs.readers = append(fs.readers, badFile{name: "test"})
	f, err := NewStream("test", fs)
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	defer cleanup(f, t)
	defer f.Close()

	r, err := f.NextReader()
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	defer r.Close()
	if r.Name() != "test" {
		t.Errorf("expected name to to be 'test' got %s", r.Name())
		t.FailNow()
	}
	if _, err := r.ReadAt(nil, 0); err == nil {
		t.Error("expected ReadAt error")
		t.FailNow()
	}
	if _, err := r.Read(nil); err == nil {
		t.Error("expected Read error")
		t.FailNow()
	}
}

func TestBadFs(t *testing.T) {
	f, err := NewStream("test", badFs{})
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	defer cleanup(f, t)
	defer f.Close()

	r, err := f.NextReader()
	if err == nil {
		t.Error("expected open error")
		t.FailNow()
	} else {
		return
	}
	r.Close()
}

func TestStd(t *testing.T) {
	f, err := New(t.Name() + ".txt")
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	if f.Name() != t.Name()+".txt" {
		t.Errorf("expected name to be %s.txt: %s", t.Name(), f.Name())
	}
	testFile(f, t)
}

func TestMem(t *testing.T) {
	f, err := NewStream(t.Name()+".txt", NewMemFS())
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	f.Write(nil)
	testFile(f, t)
}

func TestCancelClosesAll(t *testing.T) {
	f, err := New(t.Name() + ".txt")
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	for i := 0; i < 10; i++ {
		_, err := f.NextReader() // we won't even grab the reader
		if err != nil {
			t.Error(err)
			t.FailNow()
		}
	}
	f.Cancel()    // Closes all the idle Readers
	cleanup(f, t) // this will deadlock if any arn't Closed
}

func TestCloseUnblocksBlockingRead(t *testing.T) {
	f, err := New(t.Name() + ".txt")
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	r, err := f.NextReader()
	if err != nil {
		t.Error(err)
	}
	go func() {
		_, err := ioutil.ReadAll(r)
		if err == nil || err == io.EOF {
			t.Error("expected an error on a blocking Read for a closed Reader")
		}
	}()
	<-time.After(100 * time.Millisecond) // wait for blocking read
	r.Close()

	<-time.After(100 * time.Millisecond) // wait until read has been unblocked
	f.Close()

	cleanup(f, t) // this will deadlock if any arn't Closed
}

func TestCancelBeforeClose(t *testing.T) {
	f, err := New(t.Name() + ".txt")
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	f.Write([]byte("Hello"))
	r, err := f.NextReader() // blocking reader
	if err != nil {
		t.Error("error creating new reader: ", err)
	}

	wg := sync.WaitGroup{}
	wg.Add(1)
	go func() {
		_, err := ioutil.ReadAll(r)
		if err != ErrCanceled {
			t.Error("Read after cancel should return an error")
		}
		wg.Done()
	}()
	<-time.After(100 * time.Millisecond) // give Reader time to block, this tests it unblocks

	// When canceling writer, reader is closed, so writer unblocks and test passes
	f.Cancel()

	// ReadAt after cancel
	_, err = ioutil.ReadAll(io.NewSectionReader(r, 0, 1))
	if err != ErrCanceled {
		t.Error("ReadAt after cancel should return an error")
	}

	// NextReader should fail as well
	_, err = f.NextReader()
	if err != ErrCanceled {
		t.Error("NextReader should be canceled, but got: ", err)
	}

	n, err := f.Write([]byte("world"))
	// Writer is closed as well
	if err == nil {
		t.Error("expected write after canceling to fail")
	}
	if n != 0 {
		t.Error("expected write after canceling to not write anything")
	}
	wg.Wait()
	cleanup(f, t)
}

func TestCancelAfterClose(t *testing.T) {
	f, err := New(t.Name() + ".txt")
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	f.Write([]byte("Hello"))
	f.Close()

	r, _ := f.NextReader()

	// When canceling writer, reader is unblocked, so writer unblocks and test passes
	f.Cancel()

	wg := sync.WaitGroup{}
	wg.Add(1)
	go func() {
		_, err := ioutil.ReadAll(r)
		if err != ErrCanceled {
			t.Error("Read after canceling (even after close) should fail")
		}
		wg.Done()
	}()

	wg.Wait()

	cleanup(f, t)
}

func TestMemStreams(t *testing.T) {
	testFile(NewMemStream(), t)
}

func TestReadAtWait(t *testing.T) {
	f, err := NewStream(t.Name()+".txt", NewMemFS())
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	r, err := f.NextReader()
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	io.WriteString(f, "hello")
	go func() {
		<-time.After(100 * time.Millisecond)
		io.WriteString(f, " world")
		f.Close()
	}()
	data := make([]byte, 11)
	_, err = r.ReadAt(data, 0)
	if err != nil {
		t.Error(err)
	}
	if string(data) != "hello world" {
		t.Error("expected to read 'hello world' got ", string(data))
	}
}

func TestShutdown(t *testing.T) {
	f, err := NewStream(t.Name()+".txt", NewMemFS())
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	r, err := f.NextReader()
	if err != nil {
		t.Error("should return valid reader before remove")
	}
	io.WriteString(f, "Hello")

	f.shutdownWithErr(nil) // noop
	n, err := r.Read(make([]byte, 1))
	if err != nil || n != 1 {
		t.Errorf("expected successful read, got %s, bytes=%d.", err, n)
	}

	er := errors.New("shutdown")

	go func() {
		_, err := f.NextReader()
		if err != er {
			t.Errorf("expected %s, got %s", er, err)
		}

		f.Close()
		io.Copy(ioutil.Discard, r)
		r.Close()
	}()

	f.shutdownWithErr(er)
}

func TestRemove(t *testing.T) {
	f, err := NewStream(t.Name()+".txt", NewMemFS())
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	r, err := f.NextReader()
	if err != nil {
		t.Error("should return valid reader before remove")
	}
	fmt.Fprintf(f, "Hello")
	go cleanup(f, t)
	<-time.After(100 * time.Millisecond)
	waitForReadToFinish := make(chan struct{})
	go func() {
		b, err := ioutil.ReadAll(r)
		r.Close()
		if err != nil {
			t.Errorf("reader should continue as if nothing happens after remove, got: %s", err.Error())
		}
		if string(b) != "Hello World" {
			t.Errorf("reader should read all written data aven after remove, got: %s", string(b))
		}
		close(waitForReadToFinish)
	}()
	r2, err := f.NextReader()
	switch err {
	case ErrRemoving:
	case nil:
		t.Error("expected error on NextReader()")
		r2.Close()
	default:
		t.Error("expected diff error on NextReader()", err)
	}

	// It should still succeed as we didn't closed stream yet what causes a remove
	fmt.Fprintf(f, " World")

	f.Close()

	n, err := fmt.Fprintf(f, " War")

	if err == nil {
		t.Errorf("Write to stream after closing should fail, has: %v", err)
	}

	if n != 0 {
		t.Error("expected write after closing to not write anything")
	}
	<-waitForReadToFinish
}

func testFile(f *Stream, t *testing.T) {

	for i := 0; i < 10; i++ {
		go testReader(f, t)
	}

	for i := 0; i < 10; i++ {
		f.Write(testdata[:10])
		<-time.After(10 * time.Millisecond)
		f.Write(testdata[10:])
	}

	if err := f.Close(); err != nil {
		t.Error("expected successful close, got ", err)
	}
	testReader(f, t)
	cleanup(f, t)
}

func testReader(f *Stream, t *testing.T) {
	r, err := f.NextReader()
	if err != nil {
		t.Error(err)
		t.FailNow()
	}

	defer func() {
		err = r.Close()
		if err != nil {
			t.Error("expected successful close, got ", err)
		}

		_, err = ioutil.ReadAll(r)
		if err == nil {
			t.Error("Read after Reader.Close() should return an error")
		}

		_, err = ioutil.ReadAll(io.NewSectionReader(r, 0, 1))
		if err == nil {
			t.Error("ReadAt after Reader.Close() should return an error")
		}
	}()

	buf := bytes.NewBuffer(nil)
	sr := io.NewSectionReader(r, 1+int64(len(testdata)*5), 5)
	io.CopyBuffer(buf, sr, make([]byte, 3)) // force multiple reads
	if !bytes.Equal(buf.Bytes(), testdata[1:6]) {
		t.Errorf("unequal %s", buf.Bytes())
		return
	}

	buf.Reset()
	io.CopyBuffer(buf, r, make([]byte, 3)) // force multiple reads
	if !bytes.Equal(buf.Bytes(), bytes.Repeat(testdata, 10)) {
		t.Errorf("unequal %s", buf.Bytes())
		return
	}
	l, closed := r.Size()
	if !closed {
		t.Errorf("expected writer to be closed, but was open")
	}
	if l != int64(len(testdata)*10) {
		t.Errorf("expected size to be %d but got %d", len(testdata)*10, l)
	}
}

func cleanup(f *Stream, t *testing.T) {
	if err := f.Remove(); err != nil && err != ErrUnsupported {
		t.Error("error while removing file: ", err)
	}
}
<|endoftext|>


<commit_before>package gofakeit

import (
	"fmt"
	"testing"
)

type Basic struct {
	s string
	S string
}

type Nested struct {
	A   string
	B   *Basic
	bar *Basic
}

type BuiltIn struct {
	Uint8   *uint8
	Uint16  *uint16
	Uint32  *uint32
	Uint64  *uint64
	Int     *int
	Int8    *int8
	Int16   *int16
	Int32   *int32
	Int64   *int64
	Float32 *float32
	Float64 *float64
	Bool    *bool
}

type Template struct {
	Number *string `fake:"#"`
	Name   *string `fake:"{person.first}"`
}

func TestStructBasic(t *testing.T) {
	var basic Basic
	Struct(&basic)
	if basic.s != "" {
		t.Error("unexported field is not populated")
	}
	if basic.S == "" {
		t.Error("exported field is populated")
	}
}

func TestStructNested(t *testing.T) {
	var nested Nested
	Struct(&nested)
	if nested.A == "" {
		t.Error("exported string field is populated")
	}
	if nested.B == nil {
		t.Error("exported struct field is populated")
	}
	if nested.B.S == "" {
		t.Error("nested struct string field is not populated")
	}
	if nested.bar != nil {
		t.Error("nested struct bar should be be populated due to unexported variable")
	}
}

func TestStructBuiltInTypes(t *testing.T) {
	var builtIn BuiltIn
	Struct(&builtIn)
	if builtIn.Uint8 == nil {
		t.Error("builtIn Uint8 was not set")
	}
	if builtIn.Uint16 == nil {
		t.Error("builtIn Uint16 was not set")
	}
	if builtIn.Uint32 == nil {
		t.Error("builtIn Uint32 was not set")
	}
	if builtIn.Uint64 == nil {
		t.Error("builtIn Uint64 was not set")
	}
	if builtIn.Int == nil {
		t.Error("builtIn int was not set")
	}
	if builtIn.Int8 == nil {
		t.Error("builtIn int8 was not set")
	}
	if builtIn.Int16 == nil {
		t.Error("builtIn int16 was not set")
	}
	if builtIn.Int32 == nil {
		t.Error("builtIn int32 was not set")
	}
	if builtIn.Int64 == nil {
		t.Error("builtIn int64 was not set")
	}
	if builtIn.Float32 == nil {
		t.Error("builtIn float32 was not set")
	}
	if builtIn.Float64 == nil {
		t.Error("builtIn float64 was not set")
	}
	if builtIn.Bool == nil {
		t.Error("builtIn bool was not set")
	}
}

func TestStructWithTemplate(t *testing.T) {
	var template Template
	Struct(&template)
	if *template.Number == "" {
		t.Error("template number should set to number value")
	}
	if *template.Name == "" {
		t.Error("template number should set to number value")
	}
}

func Example_struct() {
	Seed(11)

	type Foo struct {
		Bar     string
		Int     int
		Pointer *int
		Name    string  `fake:"{person.first}"`
		Skip    *string `fake:"skip"`
	}

	var f Foo
	Struct(&f)

	fmt.Printf("%s\n", f.Bar)
	fmt.Printf("%d\n", f.Int)
	fmt.Printf("%d\n", *f.Pointer)
	fmt.Printf("%v\n", f.Name)
	fmt.Printf("%v\n", f.Skip)

	// Output: brmarxhki
	// -8576773003117070818
	// -7054675846543980602
	// Enrique
	// <nil>
}
<commit_msg>generate - update with direct function call usage<commit_after>package gofakeit

import (
	"fmt"
	"testing"
)

type Basic struct {
	s string
	S string
}

type Nested struct {
	A   string
	B   *Basic
	bar *Basic
}

type BuiltIn struct {
	Uint8   *uint8
	Uint16  *uint16
	Uint32  *uint32
	Uint64  *uint64
	Int     *int
	Int8    *int8
	Int16   *int16
	Int32   *int32
	Int64   *int64
	Float32 *float32
	Float64 *float64
	Bool    *bool
}

type Template struct {
	Number *string `fake:"#"`
	Name   *string `fake:"{firstname}"`
}

func TestStructBasic(t *testing.T) {
	var basic Basic
	Struct(&basic)
	if basic.s != "" {
		t.Error("unexported field is not populated")
	}
	if basic.S == "" {
		t.Error("exported field is populated")
	}
}

func TestStructNested(t *testing.T) {
	var nested Nested
	Struct(&nested)
	if nested.A == "" {
		t.Error("exported string field is populated")
	}
	if nested.B == nil {
		t.Error("exported struct field is populated")
	}
	if nested.B.S == "" {
		t.Error("nested struct string field is not populated")
	}
	if nested.bar != nil {
		t.Error("nested struct bar should be be populated due to unexported variable")
	}
}

func TestStructBuiltInTypes(t *testing.T) {
	var builtIn BuiltIn
	Struct(&builtIn)
	if builtIn.Uint8 == nil {
		t.Error("builtIn Uint8 was not set")
	}
	if builtIn.Uint16 == nil {
		t.Error("builtIn Uint16 was not set")
	}
	if builtIn.Uint32 == nil {
		t.Error("builtIn Uint32 was not set")
	}
	if builtIn.Uint64 == nil {
		t.Error("builtIn Uint64 was not set")
	}
	if builtIn.Int == nil {
		t.Error("builtIn int was not set")
	}
	if builtIn.Int8 == nil {
		t.Error("builtIn int8 was not set")
	}
	if builtIn.Int16 == nil {
		t.Error("builtIn int16 was not set")
	}
	if builtIn.Int32 == nil {
		t.Error("builtIn int32 was not set")
	}
	if builtIn.Int64 == nil {
		t.Error("builtIn int64 was not set")
	}
	if builtIn.Float32 == nil {
		t.Error("builtIn float32 was not set")
	}
	if builtIn.Float64 == nil {
		t.Error("builtIn float64 was not set")
	}
	if builtIn.Bool == nil {
		t.Error("builtIn bool was not set")
	}
}

func TestStructWithTemplate(t *testing.T) {
	var template Template
	Struct(&template)
	if *template.Number == "" {
		t.Error("template number should set to number value")
	}
	if *template.Name == "" {
		t.Error("template number should set to number value")
	}
}

func Example_struct() {
	Seed(11)

	type Foo struct {
		Bar     string
		Int     int
		Pointer *int
		Name    string  `fake:"{firstname}"`
		Skip    *string `fake:"skip"`
	}

	var f Foo
	Struct(&f)

	fmt.Printf("%s\n", f.Bar)
	fmt.Printf("%d\n", f.Int)
	fmt.Printf("%d\n", *f.Pointer)
	fmt.Printf("%v\n", f.Name)
	fmt.Printf("%v\n", f.Skip)

	// Output: brmarxhki
	// -8576773003117070818
	// -7054675846543980602
	// Enrique
	// <nil>
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/gogf/gf/g/net/gscanner"
	"net"
	"time"
)

func main() {
	//gscanner.New().SetTimeout(3*time.Second).ScanIp("192.168.2.1", "192.168.2.255", 80, func(conn net.Conn){
	//    fmt.Println(conn.RemoteAddr())
	//})

	gscanner.New().SetTimeout(3*time.Second).ScanIp("120.76.249.1", "120.76.249.255", 80, func(conn net.Conn) {
		fmt.Println(conn.RemoteAddr())
	})

	//gscanner.New().SetTimeout(6*time.Second).ScanPort("120.76.249.2", func(conn net.Conn){
	//    //fmt.Println("yes")
	//    fmt.Println(conn.RemoteAddr())
	//})
}
<commit_msg>remove example codes for gcanner<commit_after><|endoftext|>