<commit_before>/******************************************************************************
pageviewsync Source Code
Copyright (C) 2013 Lumen LLC.

This file is part of the pageviewsync Source Code.

pageviewsync is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

pageviewsync is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with pageviewsync.  If not, see <http://www.gnu.org/licenses/>.
*****************************************************************************/

package worker

import (
	"database/sql"
	"errors"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	lmnHttp "lumenlearning.com/util/http"
	lmnCanvas "lumenlearning.com/util/canvas/api"
	lmnTime "lumenlearning.com/util/time"
	"net/http"
	"sync"
)

type Worker struct {
	UserID string
	DBInfo DBConnectInfo
	APIInfo APIConnectInfo

	WorkGrp chan bool
	WaitGrp *sync.WaitGroup
	Logger chan interface{}
}

type DBConnectInfo struct {
	User string
	Pass string
	Schema string
	Table string
}

func (d *DBConnectInfo) String() string {
	return "User:"+d.User+",Pass:"+d.Pass+",Schema:"+d.Schema+",Table:"+d.Table
}

type APIConnectInfo struct {
	Host string
	Auth string
	Client *http.Client
}

func (a *APIConnectInfo) String() string {
	return "Host:"+a.Host+",Auth:"+a.Auth+",Client:"+fmt.Sprint(a.Client)
}

// Indicate that this goroutine is finished.
// Free up space for another worker to start.
func (w *Worker) Done() {
	w.WaitGrp.Done()
	<- w.WorkGrp
}

func (w *Worker) Log(format string, args... interface{}) {
	w.Logger <- fmt.Sprintf("%v: %v", w.UserID, fmt.Sprintf(format, args...))
}

func (w *Worker) RunPageviewUpdate() error {
	// Find out what we need to do to get this user up-to-date.
	// We'll find out what the last timestamp we have recorded is,
	// if any.
	lastRequestID, lastTimestamp, err := GetUpdateReqs(w.UserID, w.DBInfo)
	if err != nil {
		w.Done()
		w.Log("Terminating. worker.GetUpdateReqs => %v", err.Error())
		return errors.New("worker.GetUpdateReqs => " + err.Error())
	}
	var fullUpdate bool = (lastTimestamp == "" || lastRequestID == "")

	w.Log("lastRequestID=%v, lastRequestID=%v, fullUpdate=%v", lastRequestID, lastTimestamp, fullUpdate)

	// Get the information needed to update the local DB
	url := "https://"+w.APIInfo.Host+"/api/v1/users/"+w.UserID+"/page_views?per_page=1000"

	// Go get as many pages as we need
	var pvs []lmnCanvas.Pageview
	keepGoing := true
	for keepGoing {
		w.Log("Fetching %v", url)
		nextLink, results, err := ParsePage(url, w.APIInfo)
		if err != nil {
			w.Done()
			w.Log("Terminating. worker.ParsePage => %v", err.Error())
			return errors.New("worker.ParsePage => " + err.Error())
		}

		// Update url to be the next link
		url = nextLink

		if !fullUpdate {
			// Find out if we need to keep going
			for _, r := range results {
				ts, _ := r["created_at"]
				id, _ := r["request_id"]

				tISO, err := lmnTime.TimeFromISO8601Full(ts.(string))
				if err != nil {
					w.Done()
					w.Log("Terminated. lmnTime.TimeFromISO8601 => %v", err)
					return errors.New("lmnTime.TimeFromISO8601 => " + err.Error())
				}
				tUnix, err := lmnTime.ISO8601BasicFromTime(tISO)
				if err != nil {
					w.Done()
					w.Log("Terminated. lmnTime.ISO8601BasicFromTime => %v", err)
					return errors.New("lmnTime.ISO8601BasicFromTime => " + err.Error())
				}

				if id.(string) == lastRequestID || tUnix < lastTimestamp {
					keepGoing = false
					break
				} else {
					pvs = append(pvs, r)
				}
			}
		} else {
			// Add everything from this page
			pvs = append(pvs, results...)
		}

		// Is there anything left to grab?
		if nextLink == "" {
			keepGoing = false
			break
		}
	}

	insertCount, err := UpdateDB(&w.DBInfo, &pvs)
	if err != nil {
		w.Done()
		w.Log("Terminated. worker.UpdateDB => %v", err.Error())
		return errors.New("worker.UpdateDB => " + err.Error())
	}

	// If we got to this point, no errors!
	w.Done()
	w.Log("Updated with %v rows.", insertCount)
	return nil
}

func UpdateDB(dbInfo *DBConnectInfo, pvs *[]lmnCanvas.Pageview) (int64, error) {
	//// send the results to the local database ////
	// make the connection
	con, err := sql.Open("mysql", dbInfo.User+":"+dbInfo.Pass+"@unix(/var/run/mysqld/mysqld.sock)"+"/"+dbInfo.Schema)
	defer con.Close()
	if err != nil {
		return 0, errors.New("sql.Open => " + err.Error())
	}

	// turn off autocommit to keep our updates atomic
	// This is unnecessary.
/*	_, err = con.Exec("SET autocommit=0;")
	if err != nil {
		return 0, errors.New("con.Exec => " + err.Error())
 }*/

	// start a transaction
	tx, err := con.Begin()
	if err != nil {
		return 0, errors.New("con.Begin => " + err.Error())
	}

	// Keep track of total pageview records updated
	var insertCount int64

	// convert each lmnCanvas.Pageview to an insert statement
	for _, pv := range *pvs {
		ins := "INSERT INTO "+dbInfo.Table+" ("
		val := " VALUES ("
		i := 0
		for k, v := range pv {
			// Get the string representation of the value
			strVal := fmt.Sprint(v)
			if strVal == "<nil>" {
				strVal = "NULL"
			}

			// build insert and value strings for this item
			insI, valI := BuildInsertAndValues(k, strVal, i)
			ins = ins + insI
			val = val + valI
			i += 1

			// If it's the "created_at" field, we also need to make a copy
			//   in YYYY-mm-dd HH:MM:SS format.  It's NOT "Unix" time format,
			//   but that's what I called it for lack of a better term.
			// Also for the "updated_at" field.
			if k == "created_at" {
				insI, valI, err := GetDateTimeValue("created_at_datetime", strVal, i)
				if err != nil {
					return 0, errors.New("worker.GetDateTimeValue => " + err.Error())
				}
				ins = ins + insI
				val = val + valI
				i += 1
			}
			if k == "updated_at" {
				insI, valI, err := GetDateTimeValue("updated_at_datetime", strVal, i)
				if err != nil {
					return 0, errors.New("worker.GetDateTimeValue => " + err.Error())
				}
				ins = ins + insI
				val = val + valI
				i += 1
			}

		}

		// Finish up the query bits
		ins = ins + ")"
		val = val + ")"

		// add each insert to the transaction
		qu := ins + " " + val
		res, err := tx.Exec(qu)
		if err != nil {
			err1 := err
			errorString := "tx.Exec => "+err1.Error()

			err := tx.Rollback()
			if err != nil {
				errorString += "\n"+err.Error()
			}

			return 0, errors.New(errorString)
		}

		// Record number of rows inserted
		rowsAffected, err := res.(sql.Result).RowsAffected()
		if err != nil {
			return 0, errors.New("res.RowsAffected => " + err.Error())
		}
		insertCount += rowsAffected
	}

	// commit the transaction
	err = tx.Commit()
	if err != nil {
		return 0, errors.New("tx.Commit => " + err.Error())
	}

	return insertCount, nil
}

func GetDateTimeValue(k, strVal string, i int) (string, string, error) {
	t, err := lmnTime.TimeFromISO8601Full(strVal)
	if err != nil {
		return "", "", errors.New("lmnCanvas.TimeFromISO8601 => " + err.Error())
	}

	tu, err := lmnTime.ISO8601BasicFromTime(t)
	if err != nil {
		return "", "", errors.New("lmnCanvas.UnixFromTime => " + err.Error())
	}

	insI, valI := BuildInsertAndValues(k, tu, i)
	return insI, valI, nil
}

func BuildInsertAndValues(k, strVal string, i int) (string, string) {
	ins := ""
	val := ""

	// Change the quote character to handle NULLs, which shouldn't be in quotes.
	qc := "\""
	if strVal == "NULL" {
		qc = ""
	}

	if i > 0 {
		ins = fmt.Sprintf(",`%v`", k)
		val  = fmt.Sprintf(",%v%v%v", qc, strVal, qc)
	} else {
		ins = fmt.Sprintf("`%v`", k)
		val = fmt.Sprintf("%v%v%v", qc, strVal, qc)
	}

	return ins, val
}

// Call the API and parse the results
// Return  the "next" link
// Return a []lmnCanvas.Pageview
func ParsePage(url string, apiInfo APIConnectInfo) (string, []lmnCanvas.Pageview, error) {
	// Call the API and get the response
	resp, body, err := GetResponse(url, apiInfo)
	if err != nil {
		return "", nil, errors.New("worker.GetResponse" + err.Error())
	}

	// Get the Link header, check for rel="next"
	nextLink, err := lmnCanvas.GetNextLink(resp)
	if err != nil {
		return "", nil, errors.New("lmnCanvas.GetNextLink => " + err.Error())
	}

	// Get a map representation of the data returned
	obj, err := lmnCanvas.GetObjFromJSON(body)
	if err != nil {
		return "", nil, errors.New("lmnCanvas.GetObjFromJSON => " + err.Error())
	}

	var results []interface{}

	switch vt := (*obj).(type) {
	case []interface{}:
		results = (*obj).([]interface{})
	default:
		return "", nil, errors.New(fmt.Sprintf("Expecting an array, received %v", vt))
	}

	var pvs []lmnCanvas.Pageview

	for _, v := range results {
		switch v.(type) {
		case map[string]interface{}:
			pageviewResult := v.(map[string]interface{})
			var pv lmnCanvas.Pageview = make(lmnCanvas.Pageview)
			for k, v := range pageviewResult {
				pv[k] = v
			}
			pvs = append(pvs, pv)
		}
	}

	return nextLink, pvs, nil
}

func GetResponse(url string, apiInfo APIConnectInfo) (*http.Response, *[]byte, error) {
	resp, err := lmnCanvas.AuthorizedCall(url, apiInfo.Auth, apiInfo.Client)
	if err != nil {
		return nil, nil, errors.New("lmnCanvas.AuthorizedCall => " + err.Error())
	}

	body, err := lmnHttp.ReadResponseBody(resp)
	if err != nil {
		return nil, nil, errors.New("lmnCanvas.ReadResponse => " + err.Error())
	}

	return resp, body, nil
}

func GetUpdateReqs(userID string, dbInfo DBConnectInfo) (string, string, error) {
	// Connect to the DB and the appropriate schema
	con, err := sql.Open("mysql", dbInfo.User+":"+dbInfo.Pass+"@unix(/var/run/mysqld/mysqld.sock)"+"/"+dbInfo.Schema)
	defer con.Close()
	if err != nil {
		return "", "", errors.New("sql.Open => " + err.Error())
	}

	// Build the query
	query := "SELECT request_id, created_at FROM pageviews WHERE user_id = '"+userID+"' ORDER BY created_at DESC LIMIT 1"

	// Find out the last pageview timestamp for this userID
	row := con.QueryRow(query)
	var request_id string
	var created_at string
	err = row.Scan(&request_id, &created_at)

	if err != nil {
		if err == sql.ErrNoRows {
			request_id = ""
			created_at = ""
		} else {
			return "", "", errors.New("row.Scan => " + err.Error())
		}
	}

	return request_id, created_at, nil
}
<commit_msg>Change the way we handle the lastTimestamp.<commit_after>/******************************************************************************
pageviewsync Source Code
Copyright (C) 2013 Lumen LLC.

This file is part of the pageviewsync Source Code.

pageviewsync is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

pageviewsync is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with pageviewsync.  If not, see <http://www.gnu.org/licenses/>.
*****************************************************************************/

package worker

import (
	"database/sql"
	"errors"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	lmnHttp "lumenlearning.com/util/http"
	lmnCanvas "lumenlearning.com/util/canvas/api"
	lmnTime "lumenlearning.com/util/time"
	"net/http"
	"sync"
)

type Worker struct {
	UserID string
	DBInfo DBConnectInfo
	APIInfo APIConnectInfo

	WorkGrp chan bool
	WaitGrp *sync.WaitGroup
	Logger chan interface{}
}

type DBConnectInfo struct {
	User string
	Pass string
	Schema string
	Table string
}

func (d *DBConnectInfo) String() string {
	return "User:"+d.User+",Pass:"+d.Pass+",Schema:"+d.Schema+",Table:"+d.Table
}

type APIConnectInfo struct {
	Host string
	Auth string
	Client *http.Client
}

func (a *APIConnectInfo) String() string {
	return "Host:"+a.Host+",Auth:"+a.Auth+",Client:"+fmt.Sprint(a.Client)
}

// Indicate that this goroutine is finished.
// Free up space for another worker to start.
func (w *Worker) Done() {
	w.WaitGrp.Done()
	<- w.WorkGrp
}

func (w *Worker) Log(format string, args... interface{}) {
	w.Logger <- fmt.Sprintf("%v: %v", w.UserID, fmt.Sprintf(format, args...))
}

func (w *Worker) RunPageviewUpdate() error {
	// Find out what we need to do to get this user up-to-date.
	// We'll find out what the last timestamp we have recorded is,
	// if any.
	lastRequestID, lastTimestamp, err := GetUpdateReqs(w.UserID, w.DBInfo)
	if err != nil {
		w.Done()
		w.Log("Terminating. worker.GetUpdateReqs => %v", err.Error())
		return errors.New("worker.GetUpdateReqs => " + err.Error())
	}
	var fullUpdate bool = (lastTimestamp == 0 || lastRequestID == "")

	w.Log("lastRequestID=%v, lastRequestID=%v, fullUpdate=%v", lastRequestID, lastTimestamp, fullUpdate)

	// Get the information needed to update the local DB
	url := "https://"+w.APIInfo.Host+"/api/v1/users/"+w.UserID+"/page_views?per_page=1000"

	// Go get as many pages as we need
	var pvs []lmnCanvas.Pageview
	keepGoing := true
	for keepGoing {
		w.Log("Fetching %v", url)
		nextLink, results, err := ParsePage(url, w.APIInfo)
		if err != nil {
			w.Done()
			w.Log("Terminating. worker.ParsePage => %v", err.Error())
			return errors.New("worker.ParsePage => " + err.Error())
		}

		// Update url to be the next link
		url = nextLink

		if !fullUpdate {
			// Find out if we need to keep going
			for _, r := range results {
				ts, _ := r["created_at"]
				id, _ := r["request_id"]

				tISO, err := lmnTime.TimeFromISO8601Full(ts.(string))
				if err != nil {
					w.Done()
					w.Log("Terminated. lmnTime.TimeFromISO8601 => %v", err)
					return errors.New("lmnTime.TimeFromISO8601 => " + err.Error())
				}
				tUnix := tISO.Unix()

				// We know we need to stop if we've hit the lastRequestID
				//   or the lastTimestamp (in the case lastRequestID is blank).
				if id.(string) == lastRequestID || (id.(string) == "" && tUnix <= lastTimestamp) {
					keepGoing = false
					break
				} else {
					pvs = append(pvs, r)
				}
			}
		} else {
			// Add everything from this page
			pvs = append(pvs, results...)
		}

		// Is there anything left to grab?
		if nextLink == "" {
			keepGoing = false
			break
		}
	}

	insertCount, err := w.UpdateDB(&w.DBInfo, &pvs)
	if err != nil {
		w.Done()
		w.Log("Terminated. worker.UpdateDB => %v", err.Error())
		return errors.New("worker.UpdateDB => " + err.Error())
	}

	// If we got to this point, no errors!
	w.Done()
	w.Log("Updated with %v rows.", insertCount)
	return nil
}

func (w *Worker) UpdateDB(dbInfo *DBConnectInfo, pvs *[]lmnCanvas.Pageview) (int64, error) {
	//// send the results to the local database ////
	// make the connection
	con, err := sql.Open("mysql", dbInfo.User+":"+dbInfo.Pass+"@unix(/var/run/mysqld/mysqld.sock)"+"/"+dbInfo.Schema)
	defer con.Close()
	if err != nil {
		return 0, errors.New("sql.Open => " + err.Error())
	}

	// start a transaction
	tx, err := con.Begin()
	if err != nil {
		return 0, errors.New("con.Begin => " + err.Error())
	}

	// Keep track of total pageview records updated
	var insertCount int64

	// convert each lmnCanvas.Pageview to an insert statement
	for _, pv := range *pvs {
		ins := "INSERT INTO "+dbInfo.Table+" ("
		val := " VALUES ("
		i := 0
		for k, v := range pv {
			// Get the string representation of the value
			strVal := fmt.Sprint(v)
			if strVal == "<nil>" {
				strVal = "NULL"
			}

			// build insert and value strings for this item
			insI, valI := BuildInsertAndValues(k, strVal, i)
			ins = ins + insI
			val = val + valI
			i += 1

			// Here we're going to create a few workaround fields.
			// If it's the "created_at" field, we also need to make a copy
			//   in Unix timestamp format. (Seconds since
			//   1970-01-01 00:00:00 UTC). ("created_at_unix")
			// Same thing for the "updated_at" field. ("updated_at_unix")
			// The "user_id_requested" field is something I have to create
			//   to keep track of which user ID this pageview actually belongs
			//   to.  The user_id field sometimes shows 123456, somestimes
			//   10000000123456, and this makes it hard to consistently work
			//   with one specific user in the database.  The
			//   user_id_requested field SHOULD (untested assumption) the
			//   123456 format, avoiding the synonym issue altogether.
			if k == "created_at" {
				insI, valI, err := GetUnixTimestampValue("created_at_unix", strVal, i)
				if err != nil {
					return 0, errors.New("worker.GetUnixTimestampValue => " + err.Error())
				}
				ins = ins + insI
				val = val + valI
				i += 1
			}
			if k == "updated_at" {
				insI, valI, err := GetUnixTimestampValue("updated_at_unix", strVal, i)
				if err != nil {
					return 0, errors.New("worker.GetUnixTimestampValue => " + err.Error())
				}
				ins = ins + insI
				val = val + valI
				i += 1
			}
			if k == "user_id" {
				insI, valI := BuildInsertAndValues("user_id_requested", w.UserID, i)
				ins = ins + insI
				val = val + valI
				i += 1
			}

		}

		// Finish up the query bits
		ins = ins + ")"
		val = val + ")"

		// add each insert to the transaction
		qu := ins + " " + val
		res, err := tx.Exec(qu)
		if err != nil {
			err1 := err
			errorString := "tx.Exec => "+err1.Error()

			err := tx.Rollback()
			if err != nil {
				errorString += "\n"+err.Error()
			}

			return 0, errors.New(errorString)
		}

		// Record number of rows inserted
		rowsAffected, err := res.(sql.Result).RowsAffected()
		if err != nil {
			return 0, errors.New("res.RowsAffected => " + err.Error())
		}
		insertCount += rowsAffected
	}

	// commit the transaction
	err = tx.Commit()
	if err != nil {
		return 0, errors.New("tx.Commit => " + err.Error())
	}

	return insertCount, nil
}

func GetUnixTimestampValue(k, strVal string, i int) (string, string, error) {
	t, err := lmnTime.TimeFromISO8601Full(strVal)
	if err != nil {
		return "", "", errors.New("lmnCanvas.TimeFromISO8601 => " + err.Error())
	}

	tu := t.Unix()

	insI, valI := BuildInsertAndValues(k, fmt.Sprintf("%v", tu), i)
	return insI, valI, nil
}

func GetDateTimeValue(k, strVal string, i int) (string, string, error) {
	t, err := lmnTime.TimeFromISO8601Full(strVal)
	if err != nil {
		return "", "", errors.New("lmnCanvas.TimeFromISO8601 => " + err.Error())
	}

	tu, err := lmnTime.ISO8601BasicFromTime(t)
	if err != nil {
		return "", "", errors.New("lmnCanvas.UnixFromTime => " + err.Error())
	}

	insI, valI := BuildInsertAndValues(k, tu, i)
	return insI, valI, nil
}

func BuildInsertAndValues(k, strVal string, i int) (string, string) {
	ins := ""
	val := ""

	// Change the quote character to handle NULLs, which shouldn't be in quotes.
	qc := "\""
	if strVal == "NULL" {
		qc = ""
	}

	if i > 0 {
		ins = fmt.Sprintf(",`%v`", k)
		val  = fmt.Sprintf(",%v%v%v", qc, strVal, qc)
	} else {
		ins = fmt.Sprintf("`%v`", k)
		val = fmt.Sprintf("%v%v%v", qc, strVal, qc)
	}

	return ins, val
}

// Call the API and parse the results
// Return  the "next" link
// Return a []lmnCanvas.Pageview
func ParsePage(url string, apiInfo APIConnectInfo) (string, []lmnCanvas.Pageview, error) {
	// Call the API and get the response
	resp, body, err := GetResponse(url, apiInfo)
	if err != nil {
		return "", nil, errors.New("worker.GetResponse" + err.Error())
	}

	// Get the Link header, check for rel="next"
	nextLink, err := lmnCanvas.GetNextLink(resp)
	if err != nil {
		return "", nil, errors.New("lmnCanvas.GetNextLink => " + err.Error())
	}

	// Get a map representation of the data returned
	obj, err := lmnCanvas.GetObjFromJSON(body)
	if err != nil {
		return "", nil, errors.New("lmnCanvas.GetObjFromJSON => " + err.Error())
	}

	var results []interface{}

	switch vt := (*obj).(type) {
	case []interface{}:
		results = (*obj).([]interface{})
	default:
		return "", nil, errors.New(fmt.Sprintf("Expecting an array, received %v", vt))
	}

	var pvs []lmnCanvas.Pageview

	for _, v := range results {
		switch v.(type) {
		case map[string]interface{}:
			pageviewResult := v.(map[string]interface{})
			var pv lmnCanvas.Pageview = make(lmnCanvas.Pageview)
			for k, v := range pageviewResult {
				pv[k] = v
			}
			pvs = append(pvs, pv)
		}
	}

	return nextLink, pvs, nil
}

func GetResponse(url string, apiInfo APIConnectInfo) (*http.Response, *[]byte, error) {
	resp, err := lmnCanvas.AuthorizedCall(url, apiInfo.Auth, apiInfo.Client)
	if err != nil {
		return nil, nil, errors.New("lmnCanvas.AuthorizedCall => " + err.Error())
	}

	body, err := lmnHttp.ReadResponseBody(resp)
	if err != nil {
		return nil, nil, errors.New("lmnCanvas.ReadResponse => " + err.Error())
	}

	return resp, body, nil
}

func GetUpdateReqs(userID string, dbInfo DBConnectInfo) (string, int64, error) {
	// Connect to the DB and the appropriate schema
	con, err := sql.Open("mysql", dbInfo.User+":"+dbInfo.Pass+"@unix(/var/run/mysqld/mysqld.sock)"+"/"+dbInfo.Schema)
	defer con.Close()
	if err != nil {
		return "", 0, errors.New("sql.Open => " + err.Error())
	}

	// Build the query
	query := fmt.Sprintf("SELECT request_id, created_at_unix FROM pageviews WHERE user_id_requested = '%v' ORDER BY created_at_unix DESC, pageviews_key ASC LIMIT 1", userID)

	// Find out the last pageview timestamp for this userID
	row := con.QueryRow(query)
	var request_id string
	var created_at int64
	err = row.Scan(&request_id, &created_at)

	if err != nil {
		if err == sql.ErrNoRows {
			request_id = ""
			created_at = 0
		} else {
			return "", 0, errors.New("row.Scan => " + err.Error())
		}
	}

	return request_id, created_at, nil
}
<|endoftext|>


<commit_before>package wsmux

import (
	"bytes"
	"encoding/binary"
	"io/ioutil"
	"net"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

/*
TODO: Add ping and pong handlers
*/

const (
	defaultQueueSize            = 20
	defaultStreamQueueSize      = 20
	defaultKeepAliveInterval    = 30 * time.Second
	defaultStreamAcceptDeadline = 30 * time.Second
)

// Session implements net.Listener. Allows creating and acception multiplexed streams over ws
type Session struct {
	mu       sync.Mutex
	streams  map[uint32]*stream
	streamCh chan *stream
	conn     *websocket.Conn

	sendLock sync.Mutex

	keepAliveInterval    time.Duration
	streamAcceptDeadline time.Duration

	logger Logger

	nextID uint32

	closed chan struct{} // nil channel

	closeConn bool

	remoteCloseCallback func()
}

func (s *Session) send(f frame) error {
	select {
	case <-s.closed:
		return ErrSessionClosed
	default:
	}
	s.sendLock.Lock()
	defer s.sendLock.Unlock()
	err := s.conn.WriteMessage(websocket.BinaryMessage, f.Write())
	if err != nil {
		s.logger.Printf("wrote %v", f)
	}
	return err
}

func newSession(conn *websocket.Conn, server bool, conf Config) *Session {
	s := &Session{}
	s.conn = conn
	s.streams = make(map[uint32]*stream)
	s.streamCh = make(chan *stream, defaultStreamQueueSize)

	s.closed = make(chan struct{})
	s.closeConn = true

	s.remoteCloseCallback = conf.RemoteCloseCallback

	if conf.KeepAliveInterval == 0 {
		s.keepAliveInterval = defaultKeepAliveInterval
	} else {
		s.keepAliveInterval = conf.KeepAliveInterval
	}

	if conf.StreamAcceptDeadline == 0 {
		s.streamAcceptDeadline = defaultStreamAcceptDeadline
	} else {
		s.streamAcceptDeadline = conf.StreamAcceptDeadline
	}

	if server {
		s.nextID = 0
	} else {
		s.nextID = 1
	}

	if conf.Log == nil {
		s.logger = &nilLogger{}
	} else {
		s.logger = conf.Log
	}

	s.conn.SetCloseHandler(s.closeHandler)

	go s.recvLoop()
	return s
}

// Accept is used to accept an incoming stream
func (s *Session) Accept() (net.Conn, error) {

	select {
	case <-s.closed:
		return nil, ErrSessionClosed
	case str := <-s.streamCh:
		return str, nil
	}
}

// Open creates a new stream
func (s *Session) Open() (net.Conn, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	select {
	case <-s.closed:
		return nil, ErrSessionClosed
	default:
	}

	id := s.nextID
	// increment here so that we can wait safely
	s.nextID += 2
	if _, ok := s.streams[id]; ok {
		return nil, ErrDuplicateStream
	}

	str := newStream(id, s)
	s.streams[id] = str

	if err := s.send(newSynFrame(id)); err != nil {
		s.nextID -= 2
		return nil, err
	}
	// unlock mutex and wait
	s.mu.Unlock()

	select {
	case <-str.accepted:
		s.mu.Lock()
		return str, nil
	case <-s.closed:
		s.mu.Lock()
		// state of s.nextID doesn't matter here
		delete(s.streams, id)
		return nil, ErrSessionClosed
	case <-time.After(s.streamAcceptDeadline):
		s.mu.Lock()
		// nextID can be cyclically reused, and previous instance
		// may be in use by a different stream
		delete(s.streams, id)
		return nil, ErrAcceptTimeout
	}

}

// Close closes the current session and underlying connection
func (s *Session) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Check if channel has been closed
	select {
	case <-s.closed:
		return ErrSessionClosed
	default:
	}

	close(s.closed)
	if s.closeConn {
		_ = s.conn.Close()
	}

	for _, v := range s.streams {
		_ = v.Close()
	}
	s.streams = nil

	return nil
}

// Addr used for implementing net.Listener
func (s *Session) Addr() net.Addr {
	return s.conn.LocalAddr()
}

func (s *Session) closeHandler(code int, text string) error {
	s.logger.Printf("ws conn closed: code %d : %s", code, text)
	s.mu.Lock()
	defer s.mu.Unlock()
	s.closeConn = false
	return s.Close()
}

func (s *Session) recvLoop() {
	for {
		select {
		case <-s.closed:
			return
		default:
		}

		t, msgReader, err := s.conn.NextReader()
		if err != nil {
			_ = s.Close()
			break
		}
		if t != websocket.BinaryMessage {
			continue
		}

		h, err := getHeader(msgReader)
		id, msgType := h.id(), h.msg()

		switch msgType {
		//Used for creating a new stream
		case msgSYN:
			s.mu.Lock()
			if _, ok := s.streams[id]; ok {
				s.logger.Printf("received duplicate syn frame for stream %d", id)
				s.mu.Unlock()
				break
			}

			str := newStream(id, s)
			// no point in locking here
			str.accept(DefaultCapacity)

			s.streams[id] = str
			if err := s.send(newAckFrame(id, uint32(DefaultCapacity))); err != nil {
				s.logger.Printf("error accepting stream %d", id)
				_ = s.Close()
				return
			}
			s.asyncPushStream(str)
			s.mu.Unlock()

		//received data
		case msgDAT:
			s.mu.Lock()
			str, ok := s.streams[id]
			s.mu.Unlock()
			if !ok {
				s.logger.Printf("received data frame for unknown stream %d", id)
				break
			}
			b, err := ioutil.ReadAll(msgReader)
			if err != nil {
				s.logger.Print(err)
				break
			}
			str.addToBuffer(b)
			s.logger.Printf("received DAT frame on stream %d: %v", id, bytes.NewBuffer(b))

		//received ack frame
		case msgACK:
			s.mu.Lock()
			str, ok := s.streams[id]
			s.mu.Unlock()
			s.logger.Printf("received ack frame")
			if !ok {
				s.logger.Printf("received ack frame for unknown stream %d", id)
				break
			}

			b := make([]byte, 4)
			_, err := msgReader.Read(b)
			if err != nil {
				s.logger.Print(err)
				break
			}
			read := binary.LittleEndian.Uint32(b)
			select {
			case <-str.accepted:
				s.logger.Printf("received ack frame: id %d: remote read %d bytes", id, read)
				str.updateRemoteCapacity(read)
			default:
				// close str.accepted to accept stream
				s.logger.Printf("accepting stream")
				str.accept(read)
				break
			}

		// received fin frame
		case msgFIN:
			s.mu.Lock()
			str, ok := s.streams[id]
			s.mu.Unlock()
			if !ok {
				s.logger.Printf("received fin frame for unknown stream %d", id)
				break
			}

			err := str.setRemoteClosed()
			if err != nil {
				s.logger.Print(err)
			}
		}

	}
}

func (s *Session) removeStream(id uint32) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, ok := s.streams[id]; ok {
		delete(s.streams, id)
	}
}

func (s *Session) asyncPushStream(str *stream) {
	select {
	case s.streamCh <- str:
	default:
	}
}
<commit_msg>fixes to session: added abort<commit_after>package wsmux

import (
	"bytes"
	"encoding/binary"
	"io/ioutil"
	"net"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

/*
TODO: Add ping and pong handlers
*/

const (
	defaultQueueSize            = 20
	defaultStreamQueueSize      = 20
	defaultKeepAliveInterval    = 30 * time.Second
	defaultStreamAcceptDeadline = 30 * time.Second
)

// Session implements net.Listener. Allows creating and acception multiplexed streams over ws
type Session struct {
	mu        sync.Mutex
	streams   map[uint32]*stream
	streamCh  chan *stream
	conn      *websocket.Conn
	acceptErr error

	sendLock sync.Mutex

	keepAliveInterval    time.Duration
	streamAcceptDeadline time.Duration

	logger Logger

	nextID uint32

	closed chan struct{} // nil channel

	closeConn bool

	remoteCloseCallback func()
}

func (s *Session) send(f frame) error {
	select {
	case <-s.closed:
		return ErrSessionClosed
	default:
	}
	s.sendLock.Lock()
	defer s.sendLock.Unlock()
	err := s.conn.WriteMessage(websocket.BinaryMessage, f.Write())
	if err != nil {
		s.logger.Printf("wrote %v", f)
	}
	return err
}

func newSession(conn *websocket.Conn, server bool, conf Config) *Session {
	s := &Session{
		conn:     conn,
		streams:  make(map[uint32]*stream),
		streamCh: make(chan *stream, defaultStreamQueueSize),

		closed:    make(chan struct{}),
		closeConn: true,

		nextID: 0,

		keepAliveInterval:    defaultKeepAliveInterval,
		streamAcceptDeadline: defaultStreamAcceptDeadline,
		logger:               &nilLogger{},
	}

	// streams opened by server are even numbered
	// streams opened by client are odd numbered
	if !server {
		s.nextID = 1
	}

	if conf.KeepAliveInterval != 0 {
		s.keepAliveInterval = conf.KeepAliveInterval
	}
	if conf.StreamAcceptDeadline != 0 {
		s.streamAcceptDeadline = 0
	}
	if conf.Log != nil {
		s.logger = conf.Log
	}

	s.conn.SetCloseHandler(s.closeHandler)

	go s.recvLoop()
	return s
}

// IsClosed returns true if the session is closed
func (s *Session) IsClosed() bool {
	select {
	case <-s.closed:
		return true
	default:
	}
	return false
}

// Accept is used to accept an incoming stream
func (s *Session) Accept() (net.Conn, error) {

	select {
	case <-s.closed:
		return nil, s.acceptErr
	case str := <-s.streamCh:
		return str, nil
	}
}

// Open creates a new stream
func (s *Session) Open() (net.Conn, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	select {
	case <-s.closed:
		return nil, ErrSessionClosed
	default:
	}

	id := s.nextID
	// increment here so that we can wait safely
	s.nextID += 2
	if _, ok := s.streams[id]; ok {
		return nil, ErrDuplicateStream
	}

	str := newStream(id, s)
	s.streams[id] = str

	if err := s.send(newSynFrame(id)); err != nil {
		s.nextID -= 2
		return nil, err
	}
	// unlock mutex and wait
	s.mu.Unlock()

	select {
	case <-str.accepted:
		s.mu.Lock()
		return str, nil
	case <-s.closed:
		s.mu.Lock()
		// state of s.nextID doesn't matter here
		delete(s.streams, id)
		return nil, ErrSessionClosed
	case <-time.After(s.streamAcceptDeadline):
		s.mu.Lock()
		// nextID can be cyclically reused, and previous instance
		// may be in use by a different stream
		delete(s.streams, id)
		return nil, ErrAcceptTimeout
	}

}

// Close closes the current session and underlying connection
func (s *Session) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Check if channel has been closed
	if s.IsClosed() {
		return nil
	}

	var err error
	if s.closeConn {
		err = s.conn.Close()
	}

	for _, v := range s.streams {
		_ = v.Close()
	}
	s.streams = nil

	close(s.closed)
	return err
}

// Addr used for implementing net.Listener
func (s *Session) Addr() net.Addr {
	return s.conn.LocalAddr()
}

func (s *Session) closeHandler(code int, text string) error {
	s.logger.Printf("ws conn closed: code %d : %s", code, text)
	s.mu.Lock()
	defer s.mu.Unlock()
	s.closeConn = false
	return s.Close()
}

func (s *Session) recvLoop() {
	for {
		select {
		case <-s.closed:
			return
		default:
		}

		t, msgReader, err := s.conn.NextReader()
		if err != nil {
			_ = s.abort(err)
			break
		}
		if t != websocket.BinaryMessage {
			continue
		}

		h, err := getHeader(msgReader)
		id, msgType := h.id(), h.msg()

		switch msgType {
		//Used for creating a new stream
		case msgSYN:
			s.mu.Lock()
			if _, ok := s.streams[id]; ok {
				s.logger.Printf("received duplicate syn frame for stream %d", id)
				s.mu.Unlock()
				break
			}

			str := newStream(id, s)
			// no point in locking here
			str.accept(DefaultCapacity)

			s.streams[id] = str
			if err := s.send(newAckFrame(id, uint32(DefaultCapacity))); err != nil {
				s.logger.Printf("error accepting stream %d", id)
				_ = s.Close()
				return
			}
			s.asyncPushStream(str)
			s.mu.Unlock()

		//received data
		case msgDAT:
			s.mu.Lock()
			str, ok := s.streams[id]
			s.mu.Unlock()
			if !ok {
				s.logger.Printf("received data frame for unknown stream %d", id)
				break
			}
			b, err := ioutil.ReadAll(msgReader)
			if err != nil {
				s.logger.Print(err)
				break
			}
			str.addToBuffer(b)
			s.logger.Printf("received DAT frame on stream %d: %v", id, bytes.NewBuffer(b))

		//received ack frame
		case msgACK:
			s.mu.Lock()
			str, ok := s.streams[id]
			s.mu.Unlock()
			s.logger.Printf("received ack frame")
			if !ok {
				s.logger.Printf("received ack frame for unknown stream %d", id)
				break
			}

			b := make([]byte, 4)
			_, err := msgReader.Read(b)
			if err != nil {
				s.logger.Print(err)
				break
			}
			read := binary.LittleEndian.Uint32(b)
			select {
			case <-str.accepted:
				s.logger.Printf("received ack frame: id %d: remote read %d bytes", id, read)
				str.updateRemoteCapacity(read)
			default:
				// close str.accepted to accept stream
				s.logger.Printf("accepting stream")
				str.accept(read)
				break
			}

		// received fin frame
		case msgFIN:
			s.mu.Lock()
			str, ok := s.streams[id]
			s.mu.Unlock()
			if !ok {
				s.logger.Printf("received fin frame for unknown stream %d", id)
				break
			}

			err := str.setRemoteClosed()
			if err != nil {
				s.logger.Print(err)
			}
		}

	}
}

func (s *Session) removeStream(id uint32) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, ok := s.streams[id]; ok {
		delete(s.streams, id)
	}
}

func (s *Session) asyncPushStream(str *stream) {
	select {
	case s.streamCh <- str:
	default:
	}
}

func (s *Session) abort(e error) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.acceptErr = e

	if s.IsClosed() {
		return nil
	}

	return s.Close()
}
<|endoftext|>


<commit_before>package goenocean

import (
	"bytes"
	"encoding/binary"
	"io"
	"time"

	log "github.com/cihub/seelog"
	"github.com/tarm/goserial"
)

type Encoder interface {
	Encode() []byte
}

func Serial(send chan Encoder, recv chan Packet) {
	c := &serial.Config{Name: "/dev/ttyUSB0", Baud: 57600}
	s, err := serial.OpenPort(c)
	if err != nil {
		log.Critical(err)
	}

	response := make(chan Packet, 100)
	go readPackets(s, func(data []byte) {
		reciever(data, recv, response)
	})

	go sender(s, send, response)
}

func sender(data io.ReadWriter, send chan Encoder, response chan Packet) {
	//TODO change to io.Writer

	for p := range send {
		gotResponse := make(chan struct{})
		go waitForResponse(gotResponse, response)
		_, err := data.Write(p.Encode())
		//Dont send next until we have a response from the last one
		<-gotResponse
		if err != nil {
			log.Critical(err)
		}
	}
}

func waitForResponse(weGotResponse chan struct{}, response chan Packet) {
	select {
	case p := <-response:
		log.Debugf("We got response after send: % x\n", p.Encode())
		if !bytes.Equal(p.Data(), []byte{0}) {
			log.Errorf("We got RESPONSE error after send: % x\n", p.Encode())
		}
		weGotResponse <- struct{}{}
		return
	case <-time.After(time.Second * 2):
		log.Error("We got TIMOUT after send")
		weGotResponse <- struct{}{}
		return
	}

}

func reciever(data []byte, recv chan Packet, resp chan Packet) {
	p, err := Decode(data)
	log.Debugf("%#v\n", p)
	log.Debugf("%#v\n", p.Header())
	log.Debugf("Data: %#v\n", p.Data())
	if err != nil {
		log.Error("Decode failed :", err)
		return
	}
	if p.PacketType() == 2 {
		resp <- p
	}
	recv <- p
}

func readPackets(rd io.ReadWriter, f func([]byte)) {
	//TODO change to io.Reader
	//TODO write tests for readPackets. create rd stub which has a  Read(p []byte) (n int, err error)

	buf := make([]byte, 1)
	var rawPacket []byte
	state := 0
	var length int

	for {
		readLen, err := rd.Read(buf)
		if err != nil {
			if err == io.EOF {
				return
			}
			log.Error("ERROR reading:", err)
			continue
		}

		log.Debugf("% x ", buf)

		if readLen > 0 && buf[0] == 0x55 && state == 0 {
			rawPacket = []byte{}
			state = 1
		}

		rawPacket = append(rawPacket, buf...)

		switch state {
		case 1: //header
			if len(rawPacket) > 5 {
				length = int(binary.BigEndian.Uint16(rawPacket[1:3]))
				state = 2
			}

		case 2: // the rest!
			if len(rawPacket) > 5+length+int(rawPacket[3]) {
				state = 3
			}

		case 3: //data crc
			state = 0
			f(rawPacket)
		}

	}

}
<commit_msg>use constant instead of integer<commit_after>package goenocean

import (
	"bytes"
	"encoding/binary"
	"io"
	"time"

	log "github.com/cihub/seelog"
	"github.com/tarm/goserial"
)

type Encoder interface {
	Encode() []byte
}

func Serial(send chan Encoder, recv chan Packet) {
	c := &serial.Config{Name: "/dev/ttyUSB0", Baud: 57600}
	s, err := serial.OpenPort(c)
	if err != nil {
		log.Critical(err)
	}

	response := make(chan Packet, 100)
	go readPackets(s, func(data []byte) {
		reciever(data, recv, response)
	})

	go sender(s, send, response)
}

func sender(data io.ReadWriter, send chan Encoder, response chan Packet) {
	//TODO change to io.Writer

	for p := range send {
		gotResponse := make(chan struct{})
		go waitForResponse(gotResponse, response)
		_, err := data.Write(p.Encode())
		//Dont send next until we have a response from the last one
		<-gotResponse
		if err != nil {
			log.Critical(err)
		}
	}
}

func waitForResponse(weGotResponse chan struct{}, response chan Packet) {
	select {
	case p := <-response:
		log.Debugf("We got response after send: % x\n", p.Encode())
		if !bytes.Equal(p.Data(), []byte{0}) {
			log.Errorf("We got RESPONSE error after send: % x\n", p.Encode())
		}
		weGotResponse <- struct{}{}
		return
	case <-time.After(time.Second * 2):
		log.Error("We got TIMOUT after send")
		weGotResponse <- struct{}{}
		return
	}

}

func reciever(data []byte, recv chan Packet, resp chan Packet) {
	p, err := Decode(data)
	log.Debugf("%#v\n", p)
	log.Debugf("%#v\n", p.Header())
	log.Debugf("Data: %#v\n", p.Data())
	if err != nil {
		log.Error("Decode failed :", err)
		return
	}
	if p.PacketType() == PacketTypeResponse {
		resp <- p
	}
	recv <- p
}

func readPackets(rd io.ReadWriter, f func([]byte)) {
	//TODO change to io.Reader
	//TODO write tests for readPackets. create rd stub which has a  Read(p []byte) (n int, err error)

	buf := make([]byte, 1)
	var rawPacket []byte
	state := 0
	var length int

	for {
		readLen, err := rd.Read(buf)
		if err != nil {
			if err == io.EOF {
				return
			}
			log.Error("ERROR reading:", err)
			continue
		}

		log.Debugf("% x ", buf)

		if readLen > 0 && buf[0] == 0x55 && state == 0 {
			rawPacket = []byte{}
			state = 1
		}

		rawPacket = append(rawPacket, buf...)

		switch state {
		case 1: //header
			if len(rawPacket) > 5 {
				length = int(binary.BigEndian.Uint16(rawPacket[1:3]))
				state = 2
			}

		case 2: // the rest!
			if len(rawPacket) > 5+length+int(rawPacket[3]) {
				state = 3
			}

		case 3: //data crc
			state = 0
			f(rawPacket)
		}

	}

}
<|endoftext|>


<commit_before>// Copyright 2015 Frustra. All rights reserved.
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.

package bbcode

import (
	"strings"
	"testing"
)

var fullTestInput = `the quick brown [b]fox[/b]:
[url=http://example][img]http://example.png[/img][/url]`

var fullTestOutput = `the quick brown <b>fox</b>:<br><a href="http://example"><img src="http://example.png"></a>`

func TestFullBasic(t *testing.T) {
	c := NewCompiler(false, false)
	input := fullTestInput
	output := fullTestOutput
	for in, out := range basicTests {
		input += in
		output += out
	}
	result := c.Compile(input)
	if result != output {
		t.Errorf("Failed to compile %s.\nExpected: %s, got: %s\n", input, output, result)
	}
}

func BenchmarkFullBasic(b *testing.B) {
	c := NewCompiler(false, false)
	input := fullTestInput
	for in := range basicTests {
		input += in
	}
	for i := 0; i < b.N; i++ {
		c.Compile(input)
	}
}

var basicTests = map[string]string{
	``: ``,
	`[url]http://example.com[/url]`: `<a href="http://example.com">http://example.com</a>`,
	`[img]http://example.com[/img]`: `<img src="http://example.com">`,
	`[img][/img]`:                   `<img src="">`,

	`[url=http://example.com]example[/url]`: `<a href="http://example.com">example</a>`,
	`[img=http://example.com][/img]`:        `<img src="http://example.com">`,

	`[B]bold[/b]`:          `<b>bold</b>`,
	`[i]italic[/i]`:        `<i>italic</i>`,
	`[u]underline[/U]`:     `<u>underline</u>`,
	`[s]strikethrough[/s]`: `<s>strikethrough</s>`,

	`[u][b]something[/b] then [b]something else[/b][/u]`: `<u><b>something</b> then <b>something else</b></u>`,
	`blank[b][/b]`:                                       `blank<b></b>`,

	"test\nnewline\nnewline": `test<br>newline<br>newline`,
	"test\n\nnewline":        `test<br><br>newline`,
	"[b]test[/b]\n\nnewline": `<b>test</b><br><br>newline`,
	"[b]test\nnewline[/b]":   `<b>test<br>newline</b>`,

	"[code][b]some[/b]\n[i]stuff[/i]\n[/quote][/code][b]more[/b]":         "
[b]some[/b]\n[i]stuff[/i]\n[/quote]
more",
	"[quote name=Someguy]hello[/quote]":                                   `<blockquote><cite>Someguy said:</cite>hello</blockquote>`,
	"[center]hello[/center]":                                              `<div style="text-align: center;">hello</div>`,
	"[size=6]hello[/size]":                                                `<span class="size6">hello</span>`,
	"[center][b][color=#00BFFF][size=6]hello[/size][/color][/b][/center]": `<div style="text-align: center;"><b><span style="color: #00BFFF;"><span class="size6">hello</span></span></b></div>`,

	`[not a tag][/not ]`: `[not a tag][/not ]`,
	`[not a tag]`:        `[not a tag]`,
}
var basicMultiArgTests = map[string][]string{
	`[img=http://example.com]alt text[/img]`: []string{`<img`, ` src="http://example.com"`, ` alt="alt text"`, ` title="alt text"`, `>`},
	`[img = foo]bar[/img]`:                   []string{`<img`, ` src="foo"`, ` alt="bar"`, ` title="bar"`, `>`},
}

func TestCompile(t *testing.T) {
	c := NewCompiler(false, false)
	for in, out := range basicTests {
		result := c.Compile(in)
		if result != out {
			t.Errorf("Failed to compile %s.\nExpected: %s, got: %s\n", in, out, result)
		}
	}
	for in, out := range basicMultiArgTests {
		result := c.Compile(in)
		if !strings.HasPrefix(result, out[0]) || !strings.HasSuffix(result, out[len(out)-1]) {
			t.Errorf("Failed to compile %s.\nExpected: %s, got: %s\n", in, out, result)
		}
		for i := 1; i < len(out)-1; i++ {
			if !strings.Contains(result, out[i]) {
				t.Errorf("Failed to compile %s.\nExpected: %s, got: %s\n", in, out, result)
			}
		}
	}
}

var sanitizationTests = map[string]string{
	`<script>`:            `<script>`,
	`[url]<script>[/url]`: `<a href="%3Cscript%3E"><script></a>`,

	`[url=<script>]<script>[/url]`: `<a href="%3Cscript%3E"><script></a>`,

	`[url=http://a.b/z?\]link[/url]`: `<a href="http://a.b/z?\">link</a>`,
}
var sanitizationMultiArgTests = map[string][]string{
	`[img=<script>]<script>[/img]`:        []string{`<img`, ` src="%3Cscript%3E"`, ` alt="<script>"`, ` title="<script>"`, `>`},
	`[img="http://\"a.b/z"]"link"\[/img]`: []string{`<img`, ` src="http://"a.b/z"`, ` alt=""link"\"`, ` title=""link"\"`, `>`},
}

func TestSanitization(t *testing.T) {
	c := NewCompiler(false, false)
	for in, out := range sanitizationTests {
		result := c.Compile(in)
		if result != out {
			t.Errorf("Failed to compile %s.\nExpected: %s, got: %s\n", in, out, result)
		}
	}
	for in, out := range sanitizationMultiArgTests {
		result := c.Compile(in)
		if !strings.HasPrefix(result, out[0]) || !strings.HasSuffix(result, out[len(out)-1]) {
			t.Errorf("Failed to compile %s.\nExpected: %s, got: %s\n", in, out, result)
		}
		for i := 1; i < len(out)-1; i++ {
			if !strings.Contains(result, out[i]) {
				t.Errorf("Failed to compile %s.\nExpected: %s, got: %s\n", in, out, result)
			}
		}
	}
}

func TestFullSanitization(t *testing.T) {
	c := NewCompiler(false, false)
	input := fullTestInput
	output := fullTestOutput
	for in, out := range sanitizationTests {
		input += in
		output += out
	}
	result := c.Compile(input)
	if result != output {
		t.Errorf("Failed to compile %s.\nExpected: %s, got: %s\n", input, output, result)
	}
}

func BenchmarkFullSanitization(b *testing.B) {
	c := NewCompiler(false, false)
	input := fullTestInput
	for in := range sanitizationTests {
		input += in
	}
	for i := 0; i < b.N; i++ {
		c.Compile(input)
	}
}

var brokenTests = map[string]string{
	"[b]":        `[b]`,
	"[b]\n":      `[b]<br>`,
	"[b]hello":   `[b]hello`,
	"[b]hello\n": `[b]hello<br>`,
	"the quick brown [b][i]fox[/b][/i]\n[i]\n[b]hi[/b]][b][url=http://example[img]http://example.png[/img][/url][b]": `the quick brown <b>[i]fox</b>[/i]<br>[i]<br><b>hi</b>][b][url=http://example<img src="http://example.png">[/url][b]`,
	"the quick brown[/b][b]hello[/b]":                                                                                `the quick brown[/b]<b>hello</b>`,
	"the quick brown[/b][/code]":                                                                                `the quick brown[/b][/code]`,
	"[ b][	i]the quick brown[/i][/b=hello]": `[ b]<i>the quick brown</i>[/b=hello]`,
	"[b [herp@#$%]]the quick brown[/b]": `[b [herp@#$%]]the quick brown[/b]`,
	"[b=hello a=hi	q]the quick brown[/b]": `<b>the quick brown</b>`,
	"[b]hi[":     `[b]hi[`,
	"[b hi=derp": `[b hi=derp`,
}

func TestBroken(t *testing.T) {
	c := NewCompiler(false, false)
	for in, out := range brokenTests {
		result := c.Compile(in)
		if result != out {
			t.Errorf("Failed to compile %s.\nExpected: %s, got: %s\n", in, out, result)
		}
	}
}

func BenchmarkFullBroken(b *testing.B) {
	c := NewCompiler(false, false)
	input := fullTestInput
	for in := range brokenTests {
		input += in
	}
	for i := 0; i < b.N; i++ {
		c.Compile(input)
	}
}

var customTests = map[string]string{
	`[img]//foo/bar.png[/img]`: `<img src="//custom.png">`,
	`[url]//foo/bar.png[/url]`: `[url]//foo/bar.png[/url]`,
}

func compileImg(node *BBCodeNode) (*HTMLTag, bool) {
	out, appendExpr := DefaultTagCompilers["img"](node)
	out.Attrs["src"] = "//custom.png"
	return out, appendExpr
}

func TestCompileCustom(t *testing.T) {
	c := NewCompiler(false, false)
	c.SetTag("url", nil)
	c.SetTag("img", compileImg)
	for in, out := range customTests {
		result := c.Compile(in)
		if result != out {
			t.Errorf("Failed to compile %s.\nExpected: %s, got: %s\n", in, out, result)
		}
	}
}
<commit_msg>Add bytes/s processed to benchmarks<commit_after>// Copyright 2015 Frustra. All rights reserved.
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.

package bbcode

import (
	"strings"
	"testing"
)

var fullTestInput = `the quick brown [b]fox[/b]:
[url=http://example][img]http://example.png[/img][/url]`

var fullTestOutput = `the quick brown <b>fox</b>:<br><a href="http://example"><img src="http://example.png"></a>`

func TestFullBasic(t *testing.T) {
	c := NewCompiler(false, false)
	input := fullTestInput
	output := fullTestOutput
	for in, out := range basicTests {
		input += in
		output += out
	}
	result := c.Compile(input)
	if result != output {
		t.Errorf("Failed to compile %s.\nExpected: %s, got: %s\n", input, output, result)
	}
}

func BenchmarkFullBasic(b *testing.B) {
	c := NewCompiler(false, false)
	input := fullTestInput
	for in := range basicTests {
		input += in
	}
	b.SetBytes(int64(len(input)))
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		c.Compile(input)
	}
}

var basicTests = map[string]string{
	``: ``,
	`[url]http://example.com[/url]`: `<a href="http://example.com">http://example.com</a>`,
	`[img]http://example.com[/img]`: `<img src="http://example.com">`,
	`[img][/img]`:                   `<img src="">`,

	`[url=http://example.com]example[/url]`: `<a href="http://example.com">example</a>`,
	`[img=http://example.com][/img]`:        `<img src="http://example.com">`,

	`[B]bold[/b]`:          `<b>bold</b>`,
	`[i]italic[/i]`:        `<i>italic</i>`,
	`[u]underline[/U]`:     `<u>underline</u>`,
	`[s]strikethrough[/s]`: `<s>strikethrough</s>`,

	`[u][b]something[/b] then [b]something else[/b][/u]`: `<u><b>something</b> then <b>something else</b></u>`,
	`blank[b][/b]`:                                       `blank<b></b>`,

	"test\nnewline\nnewline": `test<br>newline<br>newline`,
	"test\n\nnewline":        `test<br><br>newline`,
	"[b]test[/b]\n\nnewline": `<b>test</b><br><br>newline`,
	"[b]test\nnewline[/b]":   `<b>test<br>newline</b>`,

	"[code][b]some[/b]\n[i]stuff[/i]\n[/quote][/code][b]more[/b]":         "
[b]some[/b]\n[i]stuff[/i]\n[/quote]
more",
	"[quote name=Someguy]hello[/quote]":                                   `<blockquote><cite>Someguy said:</cite>hello</blockquote>`,
	"[center]hello[/center]":                                              `<div style="text-align: center;">hello</div>`,
	"[size=6]hello[/size]":                                                `<span class="size6">hello</span>`,
	"[center][b][color=#00BFFF][size=6]hello[/size][/color][/b][/center]": `<div style="text-align: center;"><b><span style="color: #00BFFF;"><span class="size6">hello</span></span></b></div>`,

	`[not a tag][/not ]`: `[not a tag][/not ]`,
	`[not a tag]`:        `[not a tag]`,
}
var basicMultiArgTests = map[string][]string{
	`[img=http://example.com]alt text[/img]`: []string{`<img`, ` src="http://example.com"`, ` alt="alt text"`, ` title="alt text"`, `>`},
	`[img = foo]bar[/img]`:                   []string{`<img`, ` src="foo"`, ` alt="bar"`, ` title="bar"`, `>`},
}

func TestCompile(t *testing.T) {
	c := NewCompiler(false, false)
	for in, out := range basicTests {
		result := c.Compile(in)
		if result != out {
			t.Errorf("Failed to compile %s.\nExpected: %s, got: %s\n", in, out, result)
		}
	}
	for in, out := range basicMultiArgTests {
		result := c.Compile(in)
		if !strings.HasPrefix(result, out[0]) || !strings.HasSuffix(result, out[len(out)-1]) {
			t.Errorf("Failed to compile %s.\nExpected: %s, got: %s\n", in, out, result)
		}
		for i := 1; i < len(out)-1; i++ {
			if !strings.Contains(result, out[i]) {
				t.Errorf("Failed to compile %s.\nExpected: %s, got: %s\n", in, out, result)
			}
		}
	}
}

var sanitizationTests = map[string]string{
	`<script>`:            `<script>`,
	`[url]<script>[/url]`: `<a href="%3Cscript%3E"><script></a>`,

	`[url=<script>]<script>[/url]`: `<a href="%3Cscript%3E"><script></a>`,

	`[url=http://a.b/z?\]link[/url]`: `<a href="http://a.b/z?\">link</a>`,
}
var sanitizationMultiArgTests = map[string][]string{
	`[img=<script>]<script>[/img]`:        []string{`<img`, ` src="%3Cscript%3E"`, ` alt="<script>"`, ` title="<script>"`, `>`},
	`[img="http://\"a.b/z"]"link"\[/img]`: []string{`<img`, ` src="http://"a.b/z"`, ` alt=""link"\"`, ` title=""link"\"`, `>`},
}

func TestSanitization(t *testing.T) {
	c := NewCompiler(false, false)
	for in, out := range sanitizationTests {
		result := c.Compile(in)
		if result != out {
			t.Errorf("Failed to compile %s.\nExpected: %s, got: %s\n", in, out, result)
		}
	}
	for in, out := range sanitizationMultiArgTests {
		result := c.Compile(in)
		if !strings.HasPrefix(result, out[0]) || !strings.HasSuffix(result, out[len(out)-1]) {
			t.Errorf("Failed to compile %s.\nExpected: %s, got: %s\n", in, out, result)
		}
		for i := 1; i < len(out)-1; i++ {
			if !strings.Contains(result, out[i]) {
				t.Errorf("Failed to compile %s.\nExpected: %s, got: %s\n", in, out, result)
			}
		}
	}
}

func TestFullSanitization(t *testing.T) {
	c := NewCompiler(false, false)
	input := fullTestInput
	output := fullTestOutput
	for in, out := range sanitizationTests {
		input += in
		output += out
	}
	result := c.Compile(input)
	if result != output {
		t.Errorf("Failed to compile %s.\nExpected: %s, got: %s\n", input, output, result)
	}
}

func BenchmarkFullSanitization(b *testing.B) {
	c := NewCompiler(false, false)
	input := fullTestInput
	for in := range sanitizationTests {
		input += in
	}
	b.SetBytes(int64(len(input)))
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		c.Compile(input)
	}
}

var brokenTests = map[string]string{
	"[b]":        `[b]`,
	"[b]\n":      `[b]<br>`,
	"[b]hello":   `[b]hello`,
	"[b]hello\n": `[b]hello<br>`,
	"the quick brown [b][i]fox[/b][/i]\n[i]\n[b]hi[/b]][b][url=http://example[img]http://example.png[/img][/url][b]": `the quick brown <b>[i]fox</b>[/i]<br>[i]<br><b>hi</b>][b][url=http://example<img src="http://example.png">[/url][b]`,
	"the quick brown[/b][b]hello[/b]":                                                                                `the quick brown[/b]<b>hello</b>`,
	"the quick brown[/b][/code]":                                                                                `the quick brown[/b][/code]`,
	"[ b][	i]the quick brown[/i][/b=hello]": `[ b]<i>the quick brown</i>[/b=hello]`,
	"[b [herp@#$%]]the quick brown[/b]": `[b [herp@#$%]]the quick brown[/b]`,
	"[b=hello a=hi	q]the quick brown[/b]": `<b>the quick brown</b>`,
	"[b]hi[":     `[b]hi[`,
	"[b hi=derp": `[b hi=derp`,
}

func TestBroken(t *testing.T) {
	c := NewCompiler(false, false)
	for in, out := range brokenTests {
		result := c.Compile(in)
		if result != out {
			t.Errorf("Failed to compile %s.\nExpected: %s, got: %s\n", in, out, result)
		}
	}
}

func BenchmarkFullBroken(b *testing.B) {
	c := NewCompiler(false, false)
	input := fullTestInput
	for in := range brokenTests {
		input += in
	}
	b.SetBytes(int64(len(input)))
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		c.Compile(input)
	}
}

var customTests = map[string]string{
	`[img]//foo/bar.png[/img]`: `<img src="//custom.png">`,
	`[url]//foo/bar.png[/url]`: `[url]//foo/bar.png[/url]`,
}

func compileImg(node *BBCodeNode) (*HTMLTag, bool) {
	out, appendExpr := DefaultTagCompilers["img"](node)
	out.Attrs["src"] = "//custom.png"
	return out, appendExpr
}

func TestCompileCustom(t *testing.T) {
	c := NewCompiler(false, false)
	c.SetTag("url", nil)
	c.SetTag("img", compileImg)
	for in, out := range customTests {
		result := c.Compile(in)
		if result != out {
			t.Errorf("Failed to compile %s.\nExpected: %s, got: %s\n", in, out, result)
		}
	}
}
<|endoftext|>


<commit_before>package conf

import (
	"fmt"
	"strconv"
	"strings"
)

// conf types
const (
	StringType     = "string"
	IntType        = "int"
	Int64Type      = "int64"
	BoolType       = "bool"
	StringListType = "[]string"
	AliasMapType   = "[string string, string]"
)

// MapConf 基于Map的配置信息
type MapConf map[string]string

type AliasKey struct {
	Key   string
	Alias string
}

func ErrConfMissingKey(key, dataType string) error {
	return fmt.Errorf("MissingKey: The configs must contains %s, dataType must be %s", key, dataType)
}

func ErrConfKeyType(key, dataType string) error {
	return fmt.Errorf("TypeError: The configs must contains %s, dataType must be %s", key, dataType)
}

func ErrMissConfigAliasMap(detailKeys string) error {
	return fmt.Errorf("AliasMapType must use format \"a b\" or \"a\",and split with \",\"")
}

func (conf MapConf) Get(key string) (interface{}, error) {
	value, exist := conf[key]
	if !exist {
		return nil, fmt.Errorf("The configs must contains %s", key)
	}
	return value, nil
}

func (conf MapConf) GetStringOr(key string, deft string) (string, error) {
	ret, err := conf.GetString(key)
	if err != nil {
		return deft, err
	}
	return ret, err
}

func (conf MapConf) GetString(key string) (string, error) {
	value, exist := conf[key]
	if !exist {
		return "", ErrConfMissingKey(key, StringType)
	}
	return value, nil
}

func (conf MapConf) GetIntOr(key string, deft int) (int, error) {
	ret, err := conf.GetInt(key)
	if err != nil {
		return deft, err
	}
	return ret, err
}

func (conf MapConf) GetInt(key string) (int, error) {
	value, exist := conf[key]
	if !exist {
		return 0, ErrConfMissingKey(key, IntType)
	}
	v, err := strconv.Atoi(value)
	if err != nil {
		return 0, ErrConfKeyType(key, IntType)
	}
	return v, nil
}

func (conf MapConf) GetInt64Or(key string, deft int64) (int64, error) {
	ret, err := conf.GetInt64(key)
	if err != nil {
		return deft, err
	}
	return ret, nil
}

func (conf MapConf) GetInt64(key string) (int64, error) {
	value, exist := conf[key]
	if !exist {
		return 0, ErrConfMissingKey(key, Int64Type)
	}
	v, err := strconv.ParseInt(value, 10, 64)
	if err != nil {
		return 0, ErrConfKeyType(key, Int64Type)
	}
	return v, nil
}

func (conf MapConf) GetBoolOr(key string, deft bool) (bool, error) {
	ret, err := conf.GetBool(key)
	if err != nil {
		return deft, err
	}
	return ret, err
}

func (conf MapConf) GetBool(key string) (bool, error) {
	value, exist := conf[key]
	if !exist {
		return false, ErrConfMissingKey(key, BoolType)
	}
	v, err := strconv.ParseBool(value)
	if err != nil {
		return false, ErrConfKeyType(key, BoolType)
	}
	return v, nil
}

func (conf MapConf) GetStringListOr(key string, deft []string) ([]string, error) {
	ret, err := conf.GetStringList(key)
	if err != nil {
		return deft, err
	}
	return ret, err
}

func (conf MapConf) GetStringList(key string) ([]string, error) {
	value, exist := conf[key]
	if !exist {
		return []string{}, ErrConfMissingKey(key, StringListType)
	}
	v := strings.Split(value, ",")
	var newV []string
	for _, i := range v {
		trimI := strings.TrimSpace(i)
		if len(trimI) > 0 {
			newV = append(newV, trimI)
		}
	}
	if len(newV) <= 0 {
		return []string{}, ErrConfKeyType(key, StringListType)
	}
	return newV, nil
}

func (c MapConf) GetAliasList(key string) (aks []AliasKey, err error) {
	ks, err := c.GetStringList(key)
	if err != nil {
		return
	}
	for _, k := range ks {
		parts := strings.Fields(k)
		if len(parts) <= 0 {
			continue
		}
		name := strings.TrimSpace(parts[0])
		alias := name
		if len(parts) >= 2 {
			alias = strings.TrimSpace(parts[1])
		}
		aks = append(aks, AliasKey{Key: name, Alias: alias})
	}
	return
}

func (conf MapConf) GetAliasMapOr(key string, deft map[string]string) (map[string]string, error) {
	ret, err := conf.GetAliasMap(key)
	if err != nil {
		return deft, err
	}
	return ret, err
}

func (conf MapConf) GetAliasMap(key string) (map[string]string, error) {
	value, exist := conf[key]
	if !exist || value == "" {
		return make(map[string]string), ErrConfMissingKey(key, AliasMapType)
	}
	v := strings.Split(value, ",")
	newV := make(map[string]string)
	for _, i := range v {
		trimI := strings.TrimSpace(i)
		if len(trimI) <= 0 {
			continue
		}
		var (
			name  string
			alias string
		)
		splits := strings.Fields(trimI)
		switch len(splits) {
		case 1:
			name, alias = splits[0], splits[0]
		case 2:
			name, alias = splits[0], splits[1]
		default:
			return newV, ErrMissConfigAliasMap(trimI)
		}
		newV[name] = alias
	}
	if len(newV) <= 0 {
		return make(map[string]string), ErrConfKeyType(key, AliasMapType)
	}
	return newV, nil
}
<commit_msg>empty string时为默认值<commit_after>package conf

import (
	"fmt"
	"strconv"
	"strings"
)

// conf types
const (
	StringType     = "string"
	IntType        = "int"
	Int64Type      = "int64"
	BoolType       = "bool"
	StringListType = "[]string"
	AliasMapType   = "[string string, string]"
)

// MapConf 基于Map的配置信息
type MapConf map[string]string

type AliasKey struct {
	Key   string
	Alias string
}

func ErrConfMissingKey(key, dataType string) error {
	return fmt.Errorf("MissingKey: The configs must contains %s, dataType must be %s", key, dataType)
}

func ErrConfKeyType(key, dataType string) error {
	return fmt.Errorf("TypeError: The configs must contains %s, dataType must be %s", key, dataType)
}

func ErrMissConfigAliasMap(detailKeys string) error {
	return fmt.Errorf("AliasMapType must use format \"a b\" or \"a\",and split with \",\"")
}

func (conf MapConf) Get(key string) (interface{}, error) {
	value, exist := conf[key]
	if !exist {
		return nil, fmt.Errorf("The configs must contains %s", key)
	}
	return value, nil
}

func (conf MapConf) GetStringOr(key string, deft string) (string, error) {
	ret, err := conf.GetString(key)
	if err != nil || ret == "" {
		return deft, err
	}
	return ret, err
}

func (conf MapConf) GetString(key string) (string, error) {
	value, exist := conf[key]
	if !exist {
		return "", ErrConfMissingKey(key, StringType)
	}
	return value, nil
}

func (conf MapConf) GetIntOr(key string, deft int) (int, error) {
	ret, err := conf.GetInt(key)
	if err != nil {
		return deft, err
	}
	return ret, err
}

func (conf MapConf) GetInt(key string) (int, error) {
	value, exist := conf[key]
	if !exist {
		return 0, ErrConfMissingKey(key, IntType)
	}
	v, err := strconv.Atoi(value)
	if err != nil {
		return 0, ErrConfKeyType(key, IntType)
	}
	return v, nil
}

func (conf MapConf) GetInt64Or(key string, deft int64) (int64, error) {
	ret, err := conf.GetInt64(key)
	if err != nil {
		return deft, err
	}
	return ret, nil
}

func (conf MapConf) GetInt64(key string) (int64, error) {
	value, exist := conf[key]
	if !exist {
		return 0, ErrConfMissingKey(key, Int64Type)
	}
	v, err := strconv.ParseInt(value, 10, 64)
	if err != nil {
		return 0, ErrConfKeyType(key, Int64Type)
	}
	return v, nil
}

func (conf MapConf) GetBoolOr(key string, deft bool) (bool, error) {
	ret, err := conf.GetBool(key)
	if err != nil {
		return deft, err
	}
	return ret, err
}

func (conf MapConf) GetBool(key string) (bool, error) {
	value, exist := conf[key]
	if !exist {
		return false, ErrConfMissingKey(key, BoolType)
	}
	v, err := strconv.ParseBool(value)
	if err != nil {
		return false, ErrConfKeyType(key, BoolType)
	}
	return v, nil
}

func (conf MapConf) GetStringListOr(key string, deft []string) ([]string, error) {
	ret, err := conf.GetStringList(key)
	if err != nil {
		return deft, err
	}
	return ret, err
}

func (conf MapConf) GetStringList(key string) ([]string, error) {
	value, exist := conf[key]
	if !exist {
		return []string{}, ErrConfMissingKey(key, StringListType)
	}
	v := strings.Split(value, ",")
	var newV []string
	for _, i := range v {
		trimI := strings.TrimSpace(i)
		if len(trimI) > 0 {
			newV = append(newV, trimI)
		}
	}
	if len(newV) <= 0 {
		return []string{}, ErrConfKeyType(key, StringListType)
	}
	return newV, nil
}

func (c MapConf) GetAliasList(key string) (aks []AliasKey, err error) {
	ks, err := c.GetStringList(key)
	if err != nil {
		return
	}
	for _, k := range ks {
		parts := strings.Fields(k)
		if len(parts) <= 0 {
			continue
		}
		name := strings.TrimSpace(parts[0])
		alias := name
		if len(parts) >= 2 {
			alias = strings.TrimSpace(parts[1])
		}
		aks = append(aks, AliasKey{Key: name, Alias: alias})
	}
	return
}

func (conf MapConf) GetAliasMapOr(key string, deft map[string]string) (map[string]string, error) {
	ret, err := conf.GetAliasMap(key)
	if err != nil {
		return deft, err
	}
	return ret, err
}

func (conf MapConf) GetAliasMap(key string) (map[string]string, error) {
	value, exist := conf[key]
	if !exist || value == "" {
		return make(map[string]string), ErrConfMissingKey(key, AliasMapType)
	}
	v := strings.Split(value, ",")
	newV := make(map[string]string)
	for _, i := range v {
		trimI := strings.TrimSpace(i)
		if len(trimI) <= 0 {
			continue
		}
		var (
			name  string
			alias string
		)
		splits := strings.Fields(trimI)
		switch len(splits) {
		case 1:
			name, alias = splits[0], splits[0]
		case 2:
			name, alias = splits[0], splits[1]
		default:
			return newV, ErrMissConfigAliasMap(trimI)
		}
		newV[name] = alias
	}
	if len(newV) <= 0 {
		return make(map[string]string), ErrConfKeyType(key, AliasMapType)
	}
	return newV, nil
}
<|endoftext|>


<commit_before>package irckit

import (
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"
	"unicode"

	"github.com/mattermost/platform/model"
)

type CommandHandler interface {
	handle(u *User, c *Command, args []string, service string)
}

type Command struct {
	handler   func(u *User, toUser *User, args []string, service string)
	minParams int
	maxParams int
	login     bool
}

func logout(u *User, toUser *User, args []string, service string) {
	if u.inprogress {
		u.MsgUser(toUser, "login or logout in progress. Please wait")
		return
	}
	switch service {
	case "mattermost":
		u.logoutFromMattermost()
	case "slack":
		u.logoutFromSlack()
	}
}

func login(u *User, toUser *User, args []string, service string) {
	if u.inprogress {
		u.MsgUser(toUser, "login or logout in progress. Please wait")
		return
	}
	if service == "slack" {
		var err error
		if len(args) != 1 {
			u.MsgUser(toUser, "need LOGIN <token>")
			return
		}
		u.Token = args[len(args)-1]
		if u.sc != nil {
			fmt.Println("login, starting logout")
			err := u.logoutFromSlack()
			if err != nil {
				u.MsgUser(toUser, err.Error())
				return
			}
		}
		if u.mc != nil {
			err := u.logoutFromMattermost()
			if err != nil {
				u.MsgUser(toUser, err.Error())
				return
			}
		}
		u.inprogress = true
		defer func() { u.inprogress = false }()
		u.sc, err = u.loginToSlack()
		if err != nil {
			u.MsgUser(toUser, err.Error())
			return
		}
		u.MsgUser(toUser, "login OK")
		return
	}

	cred := &MmCredentials{}
	datalen := 4
	if u.Cfg.DefaultTeam != "" {
		cred.Team = u.Cfg.DefaultTeam
		datalen--
	}
	if u.Cfg.DefaultServer != "" {
		cred.Server = u.Cfg.DefaultServer
		datalen--
	}
	if len(args) == datalen {
		cred.Pass = args[len(args)-1]
		cred.Login = args[len(args)-2]
		// no default server or team specified
		if cred.Server == "" && cred.Team == "" {
			cred.Server = args[len(args)-4]
		}
		if cred.Team == "" {
			cred.Team = args[len(args)-3]
		}
		if cred.Server == "" {
			cred.Server = args[len(args)-3]
		}

	}

	// incorrect arguments
	if len(args) != datalen {
		// no server or team
		if cred.Team != "" && cred.Server != "" {
			u.MsgUser(toUser, "need LOGIN <login> <pass>")
			return
		}
		// server missing
		if cred.Team != "" {
			u.MsgUser(toUser, "need LOGIN <server> <login> <pass>")
			return
		}
		// team missing
		if cred.Server != "" {
			u.MsgUser(toUser, "need LOGIN <team> <login> <pass>")
			return
		}
		u.MsgUser(toUser, "need LOGIN <server> <team> <login> <pass>")
		return
	}

	if !u.isValidMMServer(cred.Server) {
		u.MsgUser(toUser, "not allowed to connect to "+cred.Server)
		return
	}

	if u.sc != nil {
		fmt.Println("login, starting logout")
		err := u.logoutFromSlack()
		if err != nil {
			u.MsgUser(toUser, err.Error())
			return
		}
	}
	if u.mc != nil {
		err := u.logoutFromMattermost()
		if err != nil {
			u.MsgUser(toUser, err.Error())
			return
		}
	}
	u.Credentials = cred
	var err error
	u.mc, err = u.loginToMattermost()
	if err != nil {
		u.MsgUser(toUser, err.Error())
		return
	}
	u.mc.OnWsConnect = u.addUsersToChannels
	go u.mc.StatusLoop()
	u.MsgUser(toUser, "login OK")

}

func search(u *User, toUser *User, args []string, service string) {
	if service == "slack" {
		u.MsgUser(toUser, "not implemented")
	}
	postlist := u.mc.SearchPosts(strings.Join(args, " "))
	if postlist == nil || len(postlist.Order) == 0 {
		u.MsgUser(toUser, "no results")
		return
	}
	for i := len(postlist.Order) - 1; i >= 0; i-- {
		if postlist.Posts[postlist.Order[i]].DeleteAt > postlist.Posts[postlist.Order[i]].CreateAt {
			continue
		}
		timestamp := time.Unix(postlist.Posts[postlist.Order[i]].CreateAt/1000, 0).Format("January 02, 2006 15:04")
		channelname := u.mc.GetChannelName(postlist.Posts[postlist.Order[i]].ChannelId)
		u.MsgUser(toUser, "#"+channelname+" <"+u.mc.GetUser(postlist.Posts[postlist.Order[i]].UserId).Username+"> "+timestamp)
		u.MsgUser(toUser, strings.Repeat("=", len("#"+channelname+" <"+u.mc.GetUser(postlist.Posts[postlist.Order[i]].UserId).Username+"> "+timestamp)))
		for _, post := range strings.Split(postlist.Posts[postlist.Order[i]].Message, "\n") {
			if post != "" {
				u.MsgUser(toUser, post)
			}
		}
		if len(postlist.Posts[postlist.Order[i]].FileIds) > 0 {
			for _, fname := range u.mc.GetFileLinks(postlist.Posts[postlist.Order[i]].FileIds) {
				u.MsgUser(toUser, "download file - "+fname)
			}
		}
		u.MsgUser(toUser, "")
		u.MsgUser(toUser, "")
	}
}

func searchUsers(u *User, toUser *User, args []string, service string) {
	if service == "slack" {
		u.MsgUser(toUser, "not implemented")
	}
	users, resp := u.mc.Client.SearchUsers(&model.UserSearch{Term: strings.Join(args, " ")})
	if resp.Error != nil {
		u.MsgUser(toUser, fmt.Sprint("Error", resp.Error))
		return
	}
	for _, user := range users {
		u.MsgUser(toUser, fmt.Sprint(user.Nickname, user.FirstName, user.LastName))
	}
}

func scrollback(u *User, toUser *User, args []string, service string) {
	if service == "slack" {
		u.MsgUser(toUser, "not implemented")
	}
	if len(args) != 2 {
		u.MsgUser(toUser, "need SCROLLBACK <channel> <lines>")
		u.MsgUser(toUser, "e.g. SCROLLBACK #bugs 10 (show last 10 lines from #bugs)")
		return
	}
	limit, err := strconv.Atoi(args[1])
	if err != nil {
		u.MsgUser(toUser, "need SCROLLBACK <channel> <lines>")
		u.MsgUser(toUser, "e.g. SCROLLBACK #bugs 10 (show last 10 lines from #bugs)")
		return
	}
	if !strings.Contains(args[0], "#") {
		u.MsgUser(toUser, "need SCROLLBACK <channel> <lines>")
		u.MsgUser(toUser, "e.g. SCROLLBACK #bugs 10 (show last 10 lines from #bugs)")
		return
	}
	args[0] = strings.Replace(args[0], "#", "", -1)
	postlist := u.mc.GetPosts(u.mc.GetChannelId(args[0], u.mc.Team.Id), limit)
	if postlist == nil || len(postlist.Order) == 0 {
		u.MsgUser(toUser, "no results")
		return
	}
	for i := len(postlist.Order) - 1; i >= 0; i-- {
		nick := u.mc.GetUser(postlist.Posts[postlist.Order[i]].UserId).Username
		for _, post := range strings.Split(postlist.Posts[postlist.Order[i]].Message, "\n") {
			if post != "" {
				u.MsgUser(toUser, "<"+nick+"> "+post)
			}
		}
		if len(postlist.Posts[postlist.Order[i]].FileIds) > 0 {
			for _, fname := range u.mc.GetFileLinks(postlist.Posts[postlist.Order[i]].FileIds) {
				u.MsgUser(toUser, "<"+nick+"> download file - "+fname)
			}
		}
	}

}

var cmds = map[string]Command{
	"logout":      {handler: logout, login: true, minParams: 0, maxParams: 0},
	"login":       {handler: login, minParams: 2, maxParams: 4},
	"search":      {handler: search, login: true, minParams: 1, maxParams: -1},
	"searchusers": {handler: searchUsers, login: true, minParams: 1, maxParams: -1},
	"scrollback":  {handler: scrollback, login: true, minParams: 2, maxParams: 2},
}

func (u *User) handleServiceBot(service string, toUser *User, msg string) {

	//func (u *User) handleMMServiceBot(toUser *User, msg string) {
	commands, err := parseCommandString(msg)
	if err != nil {
		u.MsgUser(toUser, fmt.Sprintf("\"%s\" is improperly formatted", msg))
		return
	}
	cmd, ok := cmds[strings.ToLower(commands[0])]
	if !ok {
		keys := make([]string, 0)
		for k := range cmds {
			keys = append(keys, k)
		}
		u.MsgUser(toUser, "possible commands: "+strings.Join(keys, ", "))
		u.MsgUser(toUser, "<command> help for more info")
		return
	}
	if cmd.login {
		switch service {
		case "mattermost":
			if u.mc == nil {
				u.MsgUser(toUser, "You're not logged in. Use LOGIN first.")
				return
			}
		case "slack":
			if u.sc == nil {
				u.MsgUser(toUser, "You're not logged in. Use LOGIN first.")
				return
			}
		}
	}
	/*
		if cmd.minParams > len(commands[1:]) {
			u.MsgUser(toUser, fmt.Sprintf("%s requires at least %v arguments", commands[0], cmd.minParams))
			return
		}
	*/
	if cmd.maxParams > -1 && len(commands[1:]) > cmd.maxParams {
		u.MsgUser(toUser, fmt.Sprintf("%s takes at most %v arguments", commands[0], cmd.maxParams))
		return
	}
	cmd.handler(u, toUser, commands[1:], service)
}

func parseCommandString(line string) ([]string, error) {
	args := []string{}
	buf := ""
	var escaped, doubleQuoted, singleQuoted bool

	got := false

	for _, r := range line {

		// If the string is escaped
		if escaped {
			buf += string(r)
			escaped = false
			continue
		}

		// If "\"
		if r == '\\' {
			if singleQuoted {
				buf += string(r)
			} else {
				escaped = true
			}
			continue
		}

		// If it is whitespace
		if unicode.IsSpace(r) {
			if singleQuoted || doubleQuoted {
				buf += string(r)
			} else if got {
				args = append(args, buf)
				buf = ""
				got = false
			}
			continue
		}
		// If Quoted
		switch r {
		case '"':
			if !singleQuoted {
				doubleQuoted = !doubleQuoted
				continue
			}
		case '\'':
			if !doubleQuoted {
				singleQuoted = !singleQuoted
				continue
			}
		}
		got = true
		buf += string(r)
	}

	if got {
		args = append(args, buf)
	}
	if escaped || singleQuoted || doubleQuoted {
		return nil, errors.New("invalid command line string")
	}

	return args, nil
}
<commit_msg>Make sure to return for not implemented slack search. Closes 42wim/matterircd#187<commit_after>package irckit

import (
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"
	"unicode"

	"github.com/mattermost/platform/model"
)

type CommandHandler interface {
	handle(u *User, c *Command, args []string, service string)
}

type Command struct {
	handler   func(u *User, toUser *User, args []string, service string)
	minParams int
	maxParams int
	login     bool
}

func logout(u *User, toUser *User, args []string, service string) {
	if u.inprogress {
		u.MsgUser(toUser, "login or logout in progress. Please wait")
		return
	}
	switch service {
	case "mattermost":
		u.logoutFromMattermost()
	case "slack":
		u.logoutFromSlack()
	}
}

func login(u *User, toUser *User, args []string, service string) {
	if u.inprogress {
		u.MsgUser(toUser, "login or logout in progress. Please wait")
		return
	}
	if service == "slack" {
		var err error
		if len(args) != 1 {
			u.MsgUser(toUser, "need LOGIN <token>")
			return
		}
		u.Token = args[len(args)-1]
		if u.sc != nil {
			fmt.Println("login, starting logout")
			err := u.logoutFromSlack()
			if err != nil {
				u.MsgUser(toUser, err.Error())
				return
			}
		}
		if u.mc != nil {
			err := u.logoutFromMattermost()
			if err != nil {
				u.MsgUser(toUser, err.Error())
				return
			}
		}
		u.inprogress = true
		defer func() { u.inprogress = false }()
		u.sc, err = u.loginToSlack()
		if err != nil {
			u.MsgUser(toUser, err.Error())
			return
		}
		u.MsgUser(toUser, "login OK")
		return
	}

	cred := &MmCredentials{}
	datalen := 4
	if u.Cfg.DefaultTeam != "" {
		cred.Team = u.Cfg.DefaultTeam
		datalen--
	}
	if u.Cfg.DefaultServer != "" {
		cred.Server = u.Cfg.DefaultServer
		datalen--
	}
	if len(args) == datalen {
		cred.Pass = args[len(args)-1]
		cred.Login = args[len(args)-2]
		// no default server or team specified
		if cred.Server == "" && cred.Team == "" {
			cred.Server = args[len(args)-4]
		}
		if cred.Team == "" {
			cred.Team = args[len(args)-3]
		}
		if cred.Server == "" {
			cred.Server = args[len(args)-3]
		}

	}

	// incorrect arguments
	if len(args) != datalen {
		// no server or team
		if cred.Team != "" && cred.Server != "" {
			u.MsgUser(toUser, "need LOGIN <login> <pass>")
			return
		}
		// server missing
		if cred.Team != "" {
			u.MsgUser(toUser, "need LOGIN <server> <login> <pass>")
			return
		}
		// team missing
		if cred.Server != "" {
			u.MsgUser(toUser, "need LOGIN <team> <login> <pass>")
			return
		}
		u.MsgUser(toUser, "need LOGIN <server> <team> <login> <pass>")
		return
	}

	if !u.isValidMMServer(cred.Server) {
		u.MsgUser(toUser, "not allowed to connect to "+cred.Server)
		return
	}

	if u.sc != nil {
		fmt.Println("login, starting logout")
		err := u.logoutFromSlack()
		if err != nil {
			u.MsgUser(toUser, err.Error())
			return
		}
	}
	if u.mc != nil {
		err := u.logoutFromMattermost()
		if err != nil {
			u.MsgUser(toUser, err.Error())
			return
		}
	}
	u.Credentials = cred
	var err error
	u.mc, err = u.loginToMattermost()
	if err != nil {
		u.MsgUser(toUser, err.Error())
		return
	}
	u.mc.OnWsConnect = u.addUsersToChannels
	go u.mc.StatusLoop()
	u.MsgUser(toUser, "login OK")

}

func search(u *User, toUser *User, args []string, service string) {
	if service == "slack" {
		u.MsgUser(toUser, "not implemented")
		return
	}
	postlist := u.mc.SearchPosts(strings.Join(args, " "))
	if postlist == nil || len(postlist.Order) == 0 {
		u.MsgUser(toUser, "no results")
		return
	}
	for i := len(postlist.Order) - 1; i >= 0; i-- {
		if postlist.Posts[postlist.Order[i]].DeleteAt > postlist.Posts[postlist.Order[i]].CreateAt {
			continue
		}
		timestamp := time.Unix(postlist.Posts[postlist.Order[i]].CreateAt/1000, 0).Format("January 02, 2006 15:04")
		channelname := u.mc.GetChannelName(postlist.Posts[postlist.Order[i]].ChannelId)
		u.MsgUser(toUser, "#"+channelname+" <"+u.mc.GetUser(postlist.Posts[postlist.Order[i]].UserId).Username+"> "+timestamp)
		u.MsgUser(toUser, strings.Repeat("=", len("#"+channelname+" <"+u.mc.GetUser(postlist.Posts[postlist.Order[i]].UserId).Username+"> "+timestamp)))
		for _, post := range strings.Split(postlist.Posts[postlist.Order[i]].Message, "\n") {
			if post != "" {
				u.MsgUser(toUser, post)
			}
		}
		if len(postlist.Posts[postlist.Order[i]].FileIds) > 0 {
			for _, fname := range u.mc.GetFileLinks(postlist.Posts[postlist.Order[i]].FileIds) {
				u.MsgUser(toUser, "download file - "+fname)
			}
		}
		u.MsgUser(toUser, "")
		u.MsgUser(toUser, "")
	}
}

func searchUsers(u *User, toUser *User, args []string, service string) {
	if service == "slack" {
		u.MsgUser(toUser, "not implemented")
	}
	users, resp := u.mc.Client.SearchUsers(&model.UserSearch{Term: strings.Join(args, " ")})
	if resp.Error != nil {
		u.MsgUser(toUser, fmt.Sprint("Error", resp.Error))
		return
	}
	for _, user := range users {
		u.MsgUser(toUser, fmt.Sprint(user.Nickname, user.FirstName, user.LastName))
	}
}

func scrollback(u *User, toUser *User, args []string, service string) {
	if service == "slack" {
		u.MsgUser(toUser, "not implemented")
	}
	if len(args) != 2 {
		u.MsgUser(toUser, "need SCROLLBACK <channel> <lines>")
		u.MsgUser(toUser, "e.g. SCROLLBACK #bugs 10 (show last 10 lines from #bugs)")
		return
	}
	limit, err := strconv.Atoi(args[1])
	if err != nil {
		u.MsgUser(toUser, "need SCROLLBACK <channel> <lines>")
		u.MsgUser(toUser, "e.g. SCROLLBACK #bugs 10 (show last 10 lines from #bugs)")
		return
	}
	if !strings.Contains(args[0], "#") {
		u.MsgUser(toUser, "need SCROLLBACK <channel> <lines>")
		u.MsgUser(toUser, "e.g. SCROLLBACK #bugs 10 (show last 10 lines from #bugs)")
		return
	}
	args[0] = strings.Replace(args[0], "#", "", -1)
	postlist := u.mc.GetPosts(u.mc.GetChannelId(args[0], u.mc.Team.Id), limit)
	if postlist == nil || len(postlist.Order) == 0 {
		u.MsgUser(toUser, "no results")
		return
	}
	for i := len(postlist.Order) - 1; i >= 0; i-- {
		nick := u.mc.GetUser(postlist.Posts[postlist.Order[i]].UserId).Username
		for _, post := range strings.Split(postlist.Posts[postlist.Order[i]].Message, "\n") {
			if post != "" {
				u.MsgUser(toUser, "<"+nick+"> "+post)
			}
		}
		if len(postlist.Posts[postlist.Order[i]].FileIds) > 0 {
			for _, fname := range u.mc.GetFileLinks(postlist.Posts[postlist.Order[i]].FileIds) {
				u.MsgUser(toUser, "<"+nick+"> download file - "+fname)
			}
		}
	}

}

var cmds = map[string]Command{
	"logout":      {handler: logout, login: true, minParams: 0, maxParams: 0},
	"login":       {handler: login, minParams: 2, maxParams: 4},
	"search":      {handler: search, login: true, minParams: 1, maxParams: -1},
	"searchusers": {handler: searchUsers, login: true, minParams: 1, maxParams: -1},
	"scrollback":  {handler: scrollback, login: true, minParams: 2, maxParams: 2},
}

func (u *User) handleServiceBot(service string, toUser *User, msg string) {

	//func (u *User) handleMMServiceBot(toUser *User, msg string) {
	commands, err := parseCommandString(msg)
	if err != nil {
		u.MsgUser(toUser, fmt.Sprintf("\"%s\" is improperly formatted", msg))
		return
	}
	cmd, ok := cmds[strings.ToLower(commands[0])]
	if !ok {
		keys := make([]string, 0)
		for k := range cmds {
			keys = append(keys, k)
		}
		u.MsgUser(toUser, "possible commands: "+strings.Join(keys, ", "))
		u.MsgUser(toUser, "<command> help for more info")
		return
	}
	if cmd.login {
		switch service {
		case "mattermost":
			if u.mc == nil {
				u.MsgUser(toUser, "You're not logged in. Use LOGIN first.")
				return
			}
		case "slack":
			if u.sc == nil {
				u.MsgUser(toUser, "You're not logged in. Use LOGIN first.")
				return
			}
		}
	}
	/*
		if cmd.minParams > len(commands[1:]) {
			u.MsgUser(toUser, fmt.Sprintf("%s requires at least %v arguments", commands[0], cmd.minParams))
			return
		}
	*/
	if cmd.maxParams > -1 && len(commands[1:]) > cmd.maxParams {
		u.MsgUser(toUser, fmt.Sprintf("%s takes at most %v arguments", commands[0], cmd.maxParams))
		return
	}
	cmd.handler(u, toUser, commands[1:], service)
}

func parseCommandString(line string) ([]string, error) {
	args := []string{}
	buf := ""
	var escaped, doubleQuoted, singleQuoted bool

	got := false

	for _, r := range line {

		// If the string is escaped
		if escaped {
			buf += string(r)
			escaped = false
			continue
		}

		// If "\"
		if r == '\\' {
			if singleQuoted {
				buf += string(r)
			} else {
				escaped = true
			}
			continue
		}

		// If it is whitespace
		if unicode.IsSpace(r) {
			if singleQuoted || doubleQuoted {
				buf += string(r)
			} else if got {
				args = append(args, buf)
				buf = ""
				got = false
			}
			continue
		}
		// If Quoted
		switch r {
		case '"':
			if !singleQuoted {
				doubleQuoted = !doubleQuoted
				continue
			}
		case '\'':
			if !doubleQuoted {
				singleQuoted = !singleQuoted
				continue
			}
		}
		got = true
		buf += string(r)
	}

	if got {
		args = append(args, buf)
	}
	if escaped || singleQuoted || doubleQuoted {
		return nil, errors.New("invalid command line string")
	}

	return args, nil
}
<|endoftext|>


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

package aerospike

import (
	"crypto/tls"
	"io"
	"net"
	"runtime"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	. "github.com/aerospike/aerospike-client-go/logger"
	. "github.com/aerospike/aerospike-client-go/types"
)

// Connection represents a connection with a timeout.
type Connection struct {
	node *Node

	// timeout
	timeout time.Duration

	// duration after which connection is considered idle
	idleTimeout  time.Duration
	idleDeadline time.Time

	// connection object
	conn net.Conn

	// to avoid having a buffer pool and contention
	dataBuffer []byte

	lck sync.Mutex
}

// makes sure that the connection is closed eventually, even if it is not consumed
func connectionFinalizer(c *Connection) {
	c.Close()
}

func errToTimeoutErr(err error) error {
	if err, ok := err.(net.Error); ok && err.Timeout() {
		return NewAerospikeError(TIMEOUT, err.Error())
	}
	return err
}

func shouldClose(err error) bool {
	if err == io.EOF {
		return true
	}

	if err, ok := err.(net.Error); ok && err.Timeout() {
		return true
	}

	return false
}

// NewConnection creates a connection on the network and returns the pointer
// A minimum timeout of 2 seconds will always be applied.
// If the connection is not established in the specified timeout,
// an error will be returned
func NewConnection(address string, timeout time.Duration) (*Connection, error) {
	newConn := &Connection{dataBuffer: make([]byte, 1024)}
	runtime.SetFinalizer(newConn, connectionFinalizer)

	// don't wait indefinitely
	if timeout == 0 {
		timeout = 5 * time.Second
	}

	conn, err := net.DialTimeout("tcp", address, timeout)
	if err != nil {
		Logger.Error("Connection to address `" + address + "` failed to establish with error: " + err.Error())
		return nil, errToTimeoutErr(err)
	}
	newConn.conn = conn

	// set timeout at the last possible moment
	if err := newConn.SetTimeout(timeout); err != nil {
		newConn.Close()
		return nil, err
	}

	return newConn, nil
}

// NewSecureConnection creates a TLS connection on the network and returns the pointer.
// A minimum timeout of 2 seconds will always be applied.
// If the connection is not established in the specified timeout,
// an error will be returned
func NewSecureConnection(policy *ClientPolicy, host *Host) (*Connection, error) {
	address := net.JoinHostPort(host.Name, strconv.Itoa(host.Port))
	conn, err := NewConnection(address, policy.Timeout)
	if err != nil {
		return nil, err
	}

	if policy.TlsConfig == nil {
		return conn, nil
	}

	// Use version dependent clone function to clone the config
	tlsConfig := cloneTlsConfig(policy.TlsConfig)
	tlsConfig.ServerName = host.TLSName

	sconn := tls.Client(conn.conn, tlsConfig)
	if err := sconn.Handshake(); err != nil {
		sconn.Close()
		return nil, err
	}

	if host.TLSName != "" && !tlsConfig.InsecureSkipVerify {
		if err := sconn.VerifyHostname(host.TLSName); err != nil {
			sconn.Close()
			Logger.Error("Connection to address `" + address + "` failed to establish with error: " + err.Error())
			return nil, errToTimeoutErr(err)
		}
	}

	conn.conn = sconn
	return conn, nil
}

// Write writes the slice to the connection buffer.
func (ctn *Connection) Write(buf []byte) (total int, err error) {
	// make sure all bytes are written
	// Don't worry about the loop, timeout has been set elsewhere
	length := len(buf)
	var r int
	for total < length {
		if r, err = ctn.conn.Write(buf[total:]); err != nil {
			break
		}
		total += r
	}

	if err == nil {
		return total, nil
	}

	atomic.AddInt64(&ctn.node.stats.ConnectionsFailed, 1)
	ctn.Close()
	return total, errToTimeoutErr(err)
}

// ReadN reads N bytes from connection buffer to the provided Writer.
func (ctn *Connection) ReadN(buf io.Writer, length int64) (total int64, err error) {
	// if all bytes are not read, retry until successful
	// Don't worry about the loop; we've already set the timeout elsewhere
	total, err = io.CopyN(buf, ctn.conn, length)

	if err == nil && total == length {
		return total, nil
	} else if err != nil {
		if shouldClose(err) {
			ctn.Close()
		}
		return total, errToTimeoutErr(err)
	}

	atomic.AddInt64(&ctn.node.stats.ConnectionsFailed, 1)
	ctn.Close()
	return total, NewAerospikeError(SERVER_ERROR)
}

// Read reads from connection buffer to the provided slice.
func (ctn *Connection) Read(buf []byte, length int) (total int, err error) {
	// if all bytes are not read, retry until successful
	// Don't worry about the loop; we've already set the timeout elsewhere
	var r int
	for total < length {
		r, err = ctn.conn.Read(buf[total:length])
		total += r
		if err != nil {
			break
		}
	}

	if err == nil && total == length {
		return total, nil
	} else if err != nil {
		if shouldClose(err) {
			ctn.Close()
		}
		return total, errToTimeoutErr(err)
	}

	atomic.AddInt64(&ctn.node.stats.ConnectionsFailed, 1)
	ctn.Close()
	return total, NewAerospikeError(SERVER_ERROR)
}

// IsConnected returns true if the connection is not closed yet.
func (ctn *Connection) IsConnected() bool {
	return ctn.conn != nil
}

// SetTimeout sets connection timeout for both read and write operations.
func (ctn *Connection) SetTimeout(timeout time.Duration) error {
	// Set timeout ONLY if there is or has been a timeout
	if timeout > 0 || ctn.timeout != 0 {
		ctn.timeout = timeout

		// important: remove deadline when not needed; connections are pooled
		if ctn.conn != nil {
			var deadline time.Time
			if timeout > 0 {
				deadline = time.Now().Add(timeout)
			}
			if err := ctn.conn.SetDeadline(deadline); err != nil {
				return err
			}
		}
	}

	return nil
}

// Close closes the connection
func (ctn *Connection) Close() {
	ctn.lck.Lock()
	defer ctn.lck.Unlock()

	if ctn != nil && ctn.conn != nil {
		// deregister
		if ctn.node != nil {
			ctn.node.connectionCount.DecrementAndGet()
		}

		if err := ctn.conn.Close(); err != nil {
			Logger.Warn(err.Error())
		}
		ctn.conn = nil
	}
}

// Authenticate will send authentication information to the server.
func (ctn *Connection) Authenticate(user string, password []byte) error {
	// need to authenticate
	if user != "" {
		command := newAdminCommand(ctn.dataBuffer)
		if err := command.authenticate(ctn, user, password); err != nil {
			// Socket not authenticated. Do not put back into pool.
			ctn.Close()
			return err
		}
	}
	return nil
}

// setIdleTimeout sets the idle timeout for the connection.
func (ctn *Connection) setIdleTimeout(timeout time.Duration) {
	ctn.idleTimeout = timeout
}

// isIdle returns true if the connection has reached the idle deadline.
func (ctn *Connection) isIdle() bool {
	return ctn.idleTimeout > 0 && !time.Now().Before(ctn.idleDeadline)
}

// refresh extends the idle deadline of the connection.
func (ctn *Connection) refresh() {
	ctn.idleDeadline = time.Now().Add(ctn.idleTimeout)
}
<commit_msg>Add missing stats counters, use sync.Once instead of sync.Mutex for close method, Added DefaultBufferSize<commit_after>// Copyright 2013-2017 Aerospike, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package aerospike

import (
	"crypto/tls"
	"io"
	"net"
	"runtime"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	. "github.com/aerospike/aerospike-client-go/logger"
	. "github.com/aerospike/aerospike-client-go/types"
)

// DefaultBufferSize specifies the initial size of the connection buffer when it is created.
// If not big enough (as big as the average record), it will be reallocated to size again
// which will be more expensive.
const DefaultBufferSize = 64 * 1024 // 64 KiB

// Connection represents a connection with a timeout.
type Connection struct {
	node *Node

	// timeout
	timeout time.Duration

	// duration after which connection is considered idle
	idleTimeout  time.Duration
	idleDeadline time.Time

	// connection object
	conn net.Conn

	// to avoid having a buffer pool and contention
	dataBuffer []byte

	closer sync.Once
}

// makes sure that the connection is closed eventually, even if it is not consumed
func connectionFinalizer(c *Connection) {
	c.Close()
}

func errToTimeoutErr(err error) error {
	if err, ok := err.(net.Error); ok && err.Timeout() {
		return NewAerospikeError(TIMEOUT, err.Error())
	}
	return err
}

func shouldClose(err error) bool {
	if err == io.EOF {
		return true
	}

	if err, ok := err.(net.Error); ok && err.Timeout() {
		return true
	}

	return false
}

// NewConnection creates a connection on the network and returns the pointer
// A minimum timeout of 2 seconds will always be applied.
// If the connection is not established in the specified timeout,
// an error will be returned
func NewConnection(address string, timeout time.Duration) (*Connection, error) {
	newConn := &Connection{dataBuffer: make([]byte, DefaultBufferSize)}
	runtime.SetFinalizer(newConn, connectionFinalizer)

	// don't wait indefinitely
	if timeout == 0 {
		timeout = 5 * time.Second
	}

	conn, err := net.DialTimeout("tcp", address, timeout)
	if err != nil {
		Logger.Error("Connection to address `" + address + "` failed to establish with error: " + err.Error())
		return nil, errToTimeoutErr(err)
	}
	newConn.conn = conn

	// set timeout at the last possible moment
	if err := newConn.SetTimeout(timeout); err != nil {
		newConn.Close()
		return nil, err
	}

	return newConn, nil
}

// NewSecureConnection creates a TLS connection on the network and returns the pointer.
// A minimum timeout of 2 seconds will always be applied.
// If the connection is not established in the specified timeout,
// an error will be returned
func NewSecureConnection(policy *ClientPolicy, host *Host) (*Connection, error) {
	address := net.JoinHostPort(host.Name, strconv.Itoa(host.Port))
	conn, err := NewConnection(address, policy.Timeout)
	if err != nil {
		return nil, err
	}

	if policy.TlsConfig == nil {
		return conn, nil
	}

	// Use version dependent clone function to clone the config
	tlsConfig := cloneTlsConfig(policy.TlsConfig)
	tlsConfig.ServerName = host.TLSName

	sconn := tls.Client(conn.conn, tlsConfig)
	if err := sconn.Handshake(); err != nil {
		sconn.Close()
		return nil, err
	}

	if host.TLSName != "" && !tlsConfig.InsecureSkipVerify {
		if err := sconn.VerifyHostname(host.TLSName); err != nil {
			sconn.Close()
			Logger.Error("Connection to address `" + address + "` failed to establish with error: " + err.Error())
			return nil, errToTimeoutErr(err)
		}
	}

	conn.conn = sconn
	return conn, nil
}

// Write writes the slice to the connection buffer.
func (ctn *Connection) Write(buf []byte) (total int, err error) {
	// make sure all bytes are written
	// Don't worry about the loop, timeout has been set elsewhere
	length := len(buf)
	var r int
	for total < length {
		if r, err = ctn.conn.Write(buf[total:]); err != nil {
			break
		}
		total += r
	}

	if err == nil {
		return total, nil
	}

	if ctn.node != nil {
		atomic.AddInt64(&ctn.node.stats.ConnectionsFailed, 1)
	}
	ctn.Close()
	return total, errToTimeoutErr(err)
}

// ReadN reads N bytes from connection buffer to the provided Writer.
func (ctn *Connection) ReadN(buf io.Writer, length int64) (total int64, err error) {
	// if all bytes are not read, retry until successful
	// Don't worry about the loop; we've already set the timeout elsewhere
	total, err = io.CopyN(buf, ctn.conn, length)

	if err == nil && total == length {
		return total, nil
	} else if err != nil {
		if ctn.node != nil {
			atomic.AddInt64(&ctn.node.stats.ConnectionsFailed, 1)
		}

		if shouldClose(err) {
			ctn.Close()
		}
		return total, errToTimeoutErr(err)
	}

	if ctn.node != nil {
		atomic.AddInt64(&ctn.node.stats.ConnectionsFailed, 1)
	}
	ctn.Close()
	return total, NewAerospikeError(SERVER_ERROR)
}

// Read reads from connection buffer to the provided slice.
func (ctn *Connection) Read(buf []byte, length int) (total int, err error) {
	// if all bytes are not read, retry until successful
	// Don't worry about the loop; we've already set the timeout elsewhere
	var r int
	for total < length {
		r, err = ctn.conn.Read(buf[total:length])
		total += r
		if err != nil {
			break
		}
	}

	if err == nil && total == length {
		return total, nil
	} else if err != nil {
		if ctn.node != nil {
			atomic.AddInt64(&ctn.node.stats.ConnectionsFailed, 1)
		}
		if shouldClose(err) {
			ctn.Close()
		}
		return total, errToTimeoutErr(err)
	}

	if ctn.node != nil {
		atomic.AddInt64(&ctn.node.stats.ConnectionsFailed, 1)
	}
	ctn.Close()
	return total, NewAerospikeError(SERVER_ERROR)
}

// IsConnected returns true if the connection is not closed yet.
func (ctn *Connection) IsConnected() bool {
	return ctn.conn != nil
}

// SetTimeout sets connection timeout for both read and write operations.
func (ctn *Connection) SetTimeout(timeout time.Duration) error {
	// Set timeout ONLY if there is or has been a timeout set before
	if timeout > 0 || ctn.timeout != 0 {
		ctn.timeout = timeout

		// important: remove deadline when not needed; connections are pooled
		if ctn.conn != nil {
			var deadline time.Time
			if timeout > 0 {
				deadline = time.Now().Add(timeout)
			}
			if err := ctn.conn.SetDeadline(deadline); err != nil {
				if ctn.node != nil {
					atomic.AddInt64(&ctn.node.stats.ConnectionsFailed, 1)
				}
				return err
			}
		}
	}

	return nil
}

// Close closes the connection
func (ctn *Connection) Close() {
	ctn.closer.Do(func() {
		if ctn != nil && ctn.conn != nil {
			// deregister
			if ctn.node != nil {
				ctn.node.connectionCount.DecrementAndGet()
			}

			if err := ctn.conn.Close(); err != nil {
				Logger.Warn(err.Error())
			}
			ctn.conn = nil
		}
	})
}

// Authenticate will send authentication information to the server.
func (ctn *Connection) Authenticate(user string, password []byte) error {
	// need to authenticate
	if user != "" {
		command := newAdminCommand(ctn.dataBuffer)
		if err := command.authenticate(ctn, user, password); err != nil {
			if ctn.node != nil {
				atomic.AddInt64(&ctn.node.stats.ConnectionsFailed, 1)
			}
			// Socket not authenticated. Do not put back into pool.
			ctn.Close()
			return err
		}
	}
	return nil
}

// setIdleTimeout sets the idle timeout for the connection.
func (ctn *Connection) setIdleTimeout(timeout time.Duration) {
	ctn.idleTimeout = timeout
}

// isIdle returns true if the connection has reached the idle deadline.
func (ctn *Connection) isIdle() bool {
	return ctn.idleTimeout > 0 && !time.Now().Before(ctn.idleDeadline)
}

// refresh extends the idle deadline of the connection.
func (ctn *Connection) refresh() {
	ctn.idleDeadline = time.Now().Add(ctn.idleTimeout)
}
<|endoftext|>