<commit_before>package main

import (
	"bytes"
	"encoding/csv"
	"fmt"
	"reflect"
	"time"

	"github.com/jmoiron/sqlx"
	"github.com/nu7hatch/gouuid"
)

//Client is our SQL client
type Client struct {
	db      *sqlx.DB
	history []Query
	host    string
	user    string
}

//Row will hold rows of our SQL table
type Row []interface{}

//Result will hold our SQL query resultset
type Result struct {
	Columns []string `json:"columns"`
	Rows    []Row    `json:"rows"`
}

// Query
type Query struct {
	Timestamp int64  `json:"timestamp"`
	Query     string `json:"query"`
}

//NewClientFromURL will create a new mysql client using the URL provided in parameters
func NewClientFromURL(url string) (string, error) {

	db, err := sqlx.Open("mysql", url)

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

	user, host, _, _ := getConnParametersFromString(url)

	u4, err := uuid.NewV4()
	if err != nil {
		return "", err
	}

	strUuid := u4.String()

	dbClientMap[strUuid] = &Client{db: db, host: host, user: user}

	return strUuid, nil
}

//Close disconnects a existing connection
func (client *Client) Close() error {
	//Clear history
	client.history = nil
	client.host = ""
	client.user = ""
	return client.db.Close()
}

//Test if we have a working connection with the database
func (client *Client) Test() error {
	return client.db.Ping()
}

func (client *Client) recordQuery(query string) {
	saveQuery := Query{
		Timestamp: time.Now().UTC().Unix(),
		Query:     query,
	}
	client.history = append(client.history, saveQuery)
}

//Info of our connected database
func (client *Client) Info() (*Result, error) {
	return client.Query(MySQLInfo)
}

//Databases will list all the databases in the system
func (client *Client) Databases() ([]string, error) {
	res, err := client.Query(MySQLDatabases)

	if err != nil {
		return nil, err
	}

	var tables []string

	for _, row := range res.Rows {
		tables = append(tables, row[0].(string))
	}

	return tables, nil
}

//DatabaseTables will give you list of tables belonging to the database
func (client *Client) DatabaseTables(database string) ([]string, error) {
	res, err := client.Query(fmt.Sprintf(MySQLDatabaseTables, database))

	if err != nil {
		return nil, err
	}

	var tables []string

	for _, row := range res.Rows {
		tables = append(tables, row[0].(string))
	}

	return tables, nil
}

//DatabaseViews will give you list of views belonging to the database
func (client *Client) DatabaseViews(database string) ([]string, error) {
	res, err := client.Query(fmt.Sprintf(MySQLDatabaseViews, database))

	if err != nil {
		return nil, err
	}

	var tables []string

	for _, row := range res.Rows {
		tables = append(tables, row[0].(string))
	}

	return tables, nil
}

//DatabaseProcedures returns a list of all the stored procedures in the database
func (client *Client) DatabaseProcedures(database string) ([]string, error) {
	res, err := client.Query(fmt.Sprintf(MySQLDatabaseProcedures, database))

	if err != nil {
		return nil, err
	}

	var tables []string

	for _, row := range res.Rows {
		tables = append(tables, row[0].(string))
	}

	return tables, nil
}

//DatabaseFunctions returns a list of all the functions in the database
func (client *Client) DatabaseFunctions(database string) ([]string, error) {
	res, err := client.Query(fmt.Sprintf(MySQLDatabaseFunctions, database))

	if err != nil {
		return nil, err
	}

	var tables []string

	for _, row := range res.Rows {
		tables = append(tables, row[0].(string))
	}

	return tables, nil
}

//TableInfo will return info like data used, row count etc.
func (client *Client) TableInfo(table string) (*Result, error) {
	return client.Query(fmt.Sprintf(MySQLTableInfo, table))
}

//TableIndexes returns all the indexes of the table
func (client *Client) TableIndexes(table string) (*Result, error) {
	res, err := client.Query(fmt.Sprintf(MySQLTableIndexs, table))

	if err != nil {
		return nil, err
	}

	return res, err
}

func (client *Client) TableColumns(database string, table string) (*Result, error) {
	res, err := client.Query(fmt.Sprintf(MySQLTableColumns, database, table))

	if err != nil {
		return nil, err
	}

	return res, err
}

//ProcedureParameters returns all the paramaters of a stored procedure
func (client *Client) ProcedureParameters(procedure string, database string) (*Result, error) {
	res, err := client.Query(fmt.Sprintf(MySQLProcedureParameters, procedure, database))

	if err != nil {
		return nil, err
	}

	return res, err
}

//DatabaseCollationCharSet returns all the collation and character sets in db
func (client *Client) DatabaseCollationCharSet() (*Result, error) {
	res, err := client.Query(MySQLAllCollationCharSet)

	if err != nil {
		return nil, err
	}

	return res, err
}

//AlterDatabase let's you set character set & collation of the database
func (client *Client) AlterDatabase(database string, charset string, collation string) (*Result, error) {
	res, err := client.Query(fmt.Sprintf(MySQLDatabaseAlter, database, charset, collation))

	if err != nil {
		return nil, err
	}

	return res, err
}

//DropDatabase will drop the database from the system
func (client *Client) DropDatabase(database string) (*Result, error) {
	res, err := client.Query(fmt.Sprintf(MySQLDatabaseDrop, database))

	if err != nil {
		return nil, err
	}

	return res, err
}

//DropTable will drop the table from selected database
func (client *Client) DropTable(database string, table string) (*Result, error) {
	res, err := client.Query(fmt.Sprintf(MySQLTableDrop, database, table))

	if err != nil {
		return nil, err
	}

	return res, err
}

//TruncateTable will truncate the table
func (client *Client) TruncateTable(database string, table string) (*Result, error) {
	res, err := client.Query(fmt.Sprintf(MySQLTableTruncate, database, table))

	if err != nil {
		return nil, err
	}

	return res, err
}

//ProcedureDefinition will give you the create statement of procedure/function
func (client *Client) ProcedureDefinition(procType string, database string, name string) (*Result, error) {
	res, err := client.Query(fmt.Sprintf(MySQLProcedureDefinition, procType, database, name))

	if err != nil {
		return nil, err
	}

	return res, err
}

func (client *Client) DropProcedure(procType string, database string, name string) (bool, error) {
	_, err := client.Execute(fmt.Sprintf(MySQLProcedureDrop, procType, database, name))

	if err != nil {
		return false, err
	}

	return true, err
}

func (client *Client) ProcedureCreate(procType string, database string, name string, definition string) (bool, error) {
	trans, err := client.db.Begin()

	if err != nil {
		return false, err
	}

	//set this as default database
	_, err = trans.Exec(fmt.Sprintf("use %s;", database))

	//Drop existing procedure
	_, err = trans.Exec(fmt.Sprintf(MySQLProcedureDrop, procType, database, name))

	if err != nil {
		return false, err
	}

	//Create new definition
	//yoIndex := 10
	//if procType == "FUNCTION" {
	//	yoIndex = 9
	//}
	//mehIndex := strings.Index(definition, procType+" `")
	//newDef := splice(definition, mehIndex+yoIndex, 0, "`"+database+"`.")

	_, err = trans.Exec(definition)

	if err != nil {
		return false, err
	}

	return true, nil
}

func (client *Client) ViewDefinition(database string, name string) (*Result, error) {
	res, err := client.Query(fmt.Sprintf(MySQLViewDefinition, database, name))

	if err != nil {
		return nil, err
	}

	return res, err
}

func (client *Client) Search(query string) (*Result, error) {
	//Search in table list
	resTbl, err := client.Query(fmt.Sprintf(MySQLSearchTable, query))

	if err != nil {
		return nil, err
	}

	resProc, err := client.Query(fmt.Sprintf(MySQLSearchProcedure, query))

	if err != nil {
		return nil, err
	}

	resFunc, err := client.Query(fmt.Sprintf(MySQLSearchFunction, query))

	if err != nil {
		return nil, err
	}

	resMerge := Result{
		Columns: resTbl.Columns,
		Rows:    append(resTbl.Rows, resProc.Rows...),
	}

	resMerge.Rows = append(resMerge.Rows, resFunc.Rows...)

	return &resMerge, err
}

//Query will execute the sql query passed as parameter, and return the resultset
func (client *Client) Query(query string) (*Result, error) {
	rows, err := client.db.Queryx(query)

	client.recordQuery(query)

	if err != nil {
		return nil, err
	}

	defer rows.Close()

	cols, err := rows.Columns()

	if err != nil {
		return nil, err
	}

	result := Result{
		Columns: cols,
	}

	for rows.Next() {
		obj, err := rows.SliceScan()

		for i, item := range obj {
			if item == nil {
				obj[i] = nil
			} else {
				t := reflect.TypeOf(item).Kind().String()

				if t == "slice" {
					obj[i] = string(item.([]byte))
				}
			}
		}

		if err == nil {
			result.Rows = append(result.Rows, obj)
		}
	}

	return &result, nil
}

func (client *Client) Execute(query string) (int64, error) {
	res, err := client.db.Exec(query)

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

	rowsAffected, err := res.RowsAffected()
	if err != nil {
		return -1, err
	}

	return rowsAffected, nil
}

//Format the resultset
func (res *Result) Format() []map[string]interface{} {
	var items []map[string]interface{}

	for _, row := range res.Rows {
		item := make(map[string]interface{})

		for i, c := range res.Columns {
			item[c] = row[i]
		}

		items = append(items, item)
	}

	return items
}

//CSV will format the sql resultset as CSV
func (res *Result) CSV() []byte {
	buff := &bytes.Buffer{}
	writer := csv.NewWriter(buff)

	writer.Write(res.Columns)

	for _, row := range res.Rows {
		record := make([]string, len(res.Columns))

		for i, item := range row {
			if item != nil {
				record[i] = fmt.Sprintf("%v", item)
			} else {
				record[i] = ""
			}
		}

		err := writer.Write(record)

		if err != nil {
			fmt.Println(err)
			break
		}
	}

	writer.Flush()
	return buff.Bytes()
}
<commit_msg>keep history with latest on top order<commit_after>package main

import (
	"bytes"
	"encoding/csv"
	"fmt"
	"reflect"
	"time"

	"github.com/jmoiron/sqlx"
	"github.com/nu7hatch/gouuid"
)

//Client is our SQL client
type Client struct {
	db      *sqlx.DB
	history []Query
	host    string
	user    string
}

//Row will hold rows of our SQL table
type Row []interface{}

//Result will hold our SQL query resultset
type Result struct {
	Columns []string `json:"columns"`
	Rows    []Row    `json:"rows"`
}

// Query
type Query struct {
	Timestamp int64  `json:"timestamp"`
	Query     string `json:"query"`
}

//NewClientFromURL will create a new mysql client using the URL provided in parameters
func NewClientFromURL(url string) (string, error) {

	db, err := sqlx.Open("mysql", url)

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

	user, host, _, _ := getConnParametersFromString(url)

	u4, err := uuid.NewV4()
	if err != nil {
		return "", err
	}

	strUuid := u4.String()

	dbClientMap[strUuid] = &Client{db: db, host: host, user: user}

	return strUuid, nil
}

//Close disconnects a existing connection
func (client *Client) Close() error {
	//Clear history
	client.history = nil
	client.host = ""
	client.user = ""
	return client.db.Close()
}

//Test if we have a working connection with the database
func (client *Client) Test() error {
	return client.db.Ping()
}

func (client *Client) recordQuery(query string) {
	saveQuery := Query{
		Timestamp: time.Now().UTC().Unix(),
		Query:     query,
	}
	client.history = prepend(saveQuery, client.history)
}
func prepend(v Query, slice []Query) []Query {
	return append([]Query{v}, slice...)
}

//Info of our connected database
func (client *Client) Info() (*Result, error) {
	return client.Query(MySQLInfo)
}

//Databases will list all the databases in the system
func (client *Client) Databases() ([]string, error) {
	res, err := client.Query(MySQLDatabases)

	if err != nil {
		return nil, err
	}

	var tables []string

	for _, row := range res.Rows {
		tables = append(tables, row[0].(string))
	}

	return tables, nil
}

//DatabaseTables will give you list of tables belonging to the database
func (client *Client) DatabaseTables(database string) ([]string, error) {
	res, err := client.Query(fmt.Sprintf(MySQLDatabaseTables, database))

	if err != nil {
		return nil, err
	}

	var tables []string

	for _, row := range res.Rows {
		tables = append(tables, row[0].(string))
	}

	return tables, nil
}

//DatabaseViews will give you list of views belonging to the database
func (client *Client) DatabaseViews(database string) ([]string, error) {
	res, err := client.Query(fmt.Sprintf(MySQLDatabaseViews, database))

	if err != nil {
		return nil, err
	}

	var tables []string

	for _, row := range res.Rows {
		tables = append(tables, row[0].(string))
	}

	return tables, nil
}

//DatabaseProcedures returns a list of all the stored procedures in the database
func (client *Client) DatabaseProcedures(database string) ([]string, error) {
	res, err := client.Query(fmt.Sprintf(MySQLDatabaseProcedures, database))

	if err != nil {
		return nil, err
	}

	var tables []string

	for _, row := range res.Rows {
		tables = append(tables, row[0].(string))
	}

	return tables, nil
}

//DatabaseFunctions returns a list of all the functions in the database
func (client *Client) DatabaseFunctions(database string) ([]string, error) {
	res, err := client.Query(fmt.Sprintf(MySQLDatabaseFunctions, database))

	if err != nil {
		return nil, err
	}

	var tables []string

	for _, row := range res.Rows {
		tables = append(tables, row[0].(string))
	}

	return tables, nil
}

//TableInfo will return info like data used, row count etc.
func (client *Client) TableInfo(table string) (*Result, error) {
	return client.Query(fmt.Sprintf(MySQLTableInfo, table))
}

//TableIndexes returns all the indexes of the table
func (client *Client) TableIndexes(table string) (*Result, error) {
	res, err := client.Query(fmt.Sprintf(MySQLTableIndexs, table))

	if err != nil {
		return nil, err
	}

	return res, err
}

func (client *Client) TableColumns(database string, table string) (*Result, error) {
	res, err := client.Query(fmt.Sprintf(MySQLTableColumns, database, table))

	if err != nil {
		return nil, err
	}

	return res, err
}

//ProcedureParameters returns all the paramaters of a stored procedure
func (client *Client) ProcedureParameters(procedure string, database string) (*Result, error) {
	res, err := client.Query(fmt.Sprintf(MySQLProcedureParameters, procedure, database))

	if err != nil {
		return nil, err
	}

	return res, err
}

//DatabaseCollationCharSet returns all the collation and character sets in db
func (client *Client) DatabaseCollationCharSet() (*Result, error) {
	res, err := client.Query(MySQLAllCollationCharSet)

	if err != nil {
		return nil, err
	}

	return res, err
}

//AlterDatabase let's you set character set & collation of the database
func (client *Client) AlterDatabase(database string, charset string, collation string) (*Result, error) {
	res, err := client.Query(fmt.Sprintf(MySQLDatabaseAlter, database, charset, collation))

	if err != nil {
		return nil, err
	}

	return res, err
}

//DropDatabase will drop the database from the system
func (client *Client) DropDatabase(database string) (*Result, error) {
	res, err := client.Query(fmt.Sprintf(MySQLDatabaseDrop, database))

	if err != nil {
		return nil, err
	}

	return res, err
}

//DropTable will drop the table from selected database
func (client *Client) DropTable(database string, table string) (*Result, error) {
	res, err := client.Query(fmt.Sprintf(MySQLTableDrop, database, table))

	if err != nil {
		return nil, err
	}

	return res, err
}

//TruncateTable will truncate the table
func (client *Client) TruncateTable(database string, table string) (*Result, error) {
	res, err := client.Query(fmt.Sprintf(MySQLTableTruncate, database, table))

	if err != nil {
		return nil, err
	}

	return res, err
}

//ProcedureDefinition will give you the create statement of procedure/function
func (client *Client) ProcedureDefinition(procType string, database string, name string) (*Result, error) {
	res, err := client.Query(fmt.Sprintf(MySQLProcedureDefinition, procType, database, name))

	if err != nil {
		return nil, err
	}

	return res, err
}

func (client *Client) DropProcedure(procType string, database string, name string) (bool, error) {
	_, err := client.Execute(fmt.Sprintf(MySQLProcedureDrop, procType, database, name))

	if err != nil {
		return false, err
	}

	return true, err
}

func (client *Client) ProcedureCreate(procType string, database string, name string, definition string) (bool, error) {
	trans, err := client.db.Begin()

	if err != nil {
		return false, err
	}

	//set this as default database
	_, err = trans.Exec(fmt.Sprintf("use %s;", database))

	//Drop existing procedure
	_, err = trans.Exec(fmt.Sprintf(MySQLProcedureDrop, procType, database, name))

	if err != nil {
		return false, err
	}

	//Create new definition
	//yoIndex := 10
	//if procType == "FUNCTION" {
	//	yoIndex = 9
	//}
	//mehIndex := strings.Index(definition, procType+" `")
	//newDef := splice(definition, mehIndex+yoIndex, 0, "`"+database+"`.")

	_, err = trans.Exec(definition)

	if err != nil {
		return false, err
	}

	return true, nil
}

func (client *Client) ViewDefinition(database string, name string) (*Result, error) {
	res, err := client.Query(fmt.Sprintf(MySQLViewDefinition, database, name))

	if err != nil {
		return nil, err
	}

	return res, err
}

func (client *Client) Search(query string) (*Result, error) {
	//Search in table list
	resTbl, err := client.Query(fmt.Sprintf(MySQLSearchTable, query))

	if err != nil {
		return nil, err
	}

	resProc, err := client.Query(fmt.Sprintf(MySQLSearchProcedure, query))

	if err != nil {
		return nil, err
	}

	resFunc, err := client.Query(fmt.Sprintf(MySQLSearchFunction, query))

	if err != nil {
		return nil, err
	}

	resMerge := Result{
		Columns: resTbl.Columns,
		Rows:    append(resTbl.Rows, resProc.Rows...),
	}

	resMerge.Rows = append(resMerge.Rows, resFunc.Rows...)

	return &resMerge, err
}

//Query will execute the sql query passed as parameter, and return the resultset
func (client *Client) Query(query string) (*Result, error) {
	rows, err := client.db.Queryx(query)

	client.recordQuery(query)

	if err != nil {
		return nil, err
	}

	defer rows.Close()

	cols, err := rows.Columns()

	if err != nil {
		return nil, err
	}

	result := Result{
		Columns: cols,
	}

	for rows.Next() {
		obj, err := rows.SliceScan()

		for i, item := range obj {
			if item == nil {
				obj[i] = nil
			} else {
				t := reflect.TypeOf(item).Kind().String()

				if t == "slice" {
					obj[i] = string(item.([]byte))
				}
			}
		}

		if err == nil {
			result.Rows = append(result.Rows, obj)
		}
	}

	return &result, nil
}

func (client *Client) Execute(query string) (int64, error) {
	res, err := client.db.Exec(query)

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

	rowsAffected, err := res.RowsAffected()
	if err != nil {
		return -1, err
	}

	return rowsAffected, nil
}

//Format the resultset
func (res *Result) Format() []map[string]interface{} {
	var items []map[string]interface{}

	for _, row := range res.Rows {
		item := make(map[string]interface{})

		for i, c := range res.Columns {
			item[c] = row[i]
		}

		items = append(items, item)
	}

	return items
}

//CSV will format the sql resultset as CSV
func (res *Result) CSV() []byte {
	buff := &bytes.Buffer{}
	writer := csv.NewWriter(buff)

	writer.Write(res.Columns)

	for _, row := range res.Rows {
		record := make([]string, len(res.Columns))

		for i, item := range row {
			if item != nil {
				record[i] = fmt.Sprintf("%v", item)
			} else {
				record[i] = ""
			}
		}

		err := writer.Write(record)

		if err != nil {
			fmt.Println(err)
			break
		}
	}

	writer.Flush()
	return buff.Bytes()
}
<|endoftext|>


<commit_before>package air

import (
	"bytes"
	"compress/gzip"
	"encoding/binary"
	"fmt"
	"io/ioutil"
	"mime"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/VictoriaMetrics/fastcache"
	"github.com/cespare/xxhash"
	"github.com/fsnotify/fsnotify"
)

// coffer is a binary asset file manager that uses runtime memory to reduce disk
// I/O pressure.
type coffer struct {
	a       *Air
	once    *sync.Once
	assets  *sync.Map
	cache   *fastcache.Cache
	watcher *fsnotify.Watcher
}

// newCoffer returns a new instance of the `coffer` with the a.
func newCoffer(a *Air) *coffer {
	c := &coffer{
		a:      a,
		once:   &sync.Once{},
		assets: &sync.Map{},
	}

	var err error
	if c.watcher, err = fsnotify.NewWatcher(); err != nil {
		panic(fmt.Errorf(
			"air: failed to build coffer watcher: %v",
			err,
		))
	}

	go func() {
		for {
			select {
			case e := <-c.watcher.Events:
				if a.CofferEnabled {
					a.DEBUG(
						"air: asset file event occurs",
						map[string]interface{}{
							"file":  e.Name,
							"event": e.Op.String(),
						},
					)
				}

				if ai, ok := c.assets.Load(e.Name); ok {
					a := ai.(*asset)
					c.assets.Delete(a.name)
					c.cache.Del(a.digest)
					if a.gzippedDigest != nil {
						c.cache.Del(a.gzippedDigest)
					}
				}
			case err := <-c.watcher.Errors:
				if a.CofferEnabled {
					a.ERROR(
						"air: coffer watcher error",
						map[string]interface{}{
							"error": err.Error(),
						},
					)
				}
			}
		}
	}()

	return c
}

// asset returns an `asset` from the c for the name.
func (c *coffer) asset(name string) (*asset, error) {
	c.once.Do(func() {
		c.cache = fastcache.New(c.a.CofferMaxMemoryBytes)
	})

	if ai, ok := c.assets.Load(name); ok {
		return ai.(*asset), nil
	} else if ar, err := filepath.Abs(c.a.AssetRoot); err != nil {
		return nil, err
	} else if !strings.HasPrefix(name, ar) {
		return nil, nil
	}

	ext := filepath.Ext(name)
	if !stringSliceContainsCIly(c.a.AssetExts, ext) {
		return nil, nil
	}

	fi, err := os.Stat(name)
	if err != nil {
		return nil, err
	}

	b, err := ioutil.ReadFile(name)
	if err != nil {
		return nil, err
	}

	var (
		mt       = mime.TypeByExtension(ext)
		minified bool
		gb       []byte
	)

	if mt != "" {
		mt, _, err := mime.ParseMediaType(mt)
		if err != nil {
			return nil, err
		}

		if c.a.MinifierEnabled &&
			stringSliceContains(c.a.MinifierMIMETypes, mt) {
			if b, err = c.a.minifier.minify(mt, b); err != nil {
				return nil, err
			}

			minified = true
		}

		if c.a.GzipEnabled &&
			stringSliceContains(c.a.GzipMIMETypes, mt) {
			buf := bytes.Buffer{}
			if gw, err := gzip.NewWriterLevel(
				&buf,
				c.a.GzipCompressionLevel,
			); err != nil {
				return nil, err
			} else if _, err = gw.Write(b); err != nil {
				return nil, err
			} else if err = gw.Close(); err != nil {
				return nil, err
			}

			gb = buf.Bytes()
		}
	}

	if err := c.watcher.Add(name); err != nil {
		return nil, err
	}

	a := &asset{
		coffer:   c,
		name:     name,
		mimeType: mt,
		modTime:  fi.ModTime(),
		minified: minified,
		digest:   make([]byte, 8),
	}

	binary.BigEndian.PutUint64(a.digest, xxhash.Sum64(b))
	c.cache.Set(a.digest, b)

	if gb != nil {
		a.gzippedDigest = make([]byte, 8)
		binary.BigEndian.PutUint64(a.gzippedDigest, xxhash.Sum64(gb))
		c.cache.Set(a.gzippedDigest, gb)
	}

	c.assets.Store(name, a)

	return a, nil
}

// asset is a binary asset file.
type asset struct {
	coffer        *coffer
	name          string
	mimeType      string
	modTime       time.Time
	minified      bool
	digest        []byte
	gzippedDigest []byte
}

// content returns the content of the a with the gzipped.
func (a *asset) content(gzipped bool) []byte {
	var c []byte
	if gzipped {
		c = a.coffer.cache.Get(nil, a.gzippedDigest)
	} else {
		c = a.coffer.cache.Get(nil, a.digest)
	}

	if len(c) == 0 {
		a.coffer.assets.Delete(a.name)
		a.coffer.cache.Del(a.digest)
		if a.gzippedDigest != nil {
			a.coffer.cache.Del(a.gzippedDigest)
		}

		return nil
	}

	return c
}
<commit_msg>refactor: refine `coffer#asset()`<commit_after>package air

import (
	"bytes"
	"compress/gzip"
	"encoding/binary"
	"fmt"
	"io/ioutil"
	"mime"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/VictoriaMetrics/fastcache"
	"github.com/aofei/mimesniffer"
	"github.com/cespare/xxhash"
	"github.com/fsnotify/fsnotify"
)

// coffer is a binary asset file manager that uses runtime memory to reduce disk
// I/O pressure.
type coffer struct {
	a       *Air
	once    *sync.Once
	assets  *sync.Map
	cache   *fastcache.Cache
	watcher *fsnotify.Watcher
}

// newCoffer returns a new instance of the `coffer` with the a.
func newCoffer(a *Air) *coffer {
	c := &coffer{
		a:      a,
		once:   &sync.Once{},
		assets: &sync.Map{},
	}

	var err error
	if c.watcher, err = fsnotify.NewWatcher(); err != nil {
		panic(fmt.Errorf(
			"air: failed to build coffer watcher: %v",
			err,
		))
	}

	go func() {
		for {
			select {
			case e := <-c.watcher.Events:
				if a.CofferEnabled {
					a.DEBUG(
						"air: asset file event occurs",
						map[string]interface{}{
							"file":  e.Name,
							"event": e.Op.String(),
						},
					)
				}

				if ai, ok := c.assets.Load(e.Name); ok {
					a := ai.(*asset)
					c.assets.Delete(a.name)
					c.cache.Del(a.digest)
					if a.gzippedDigest != nil {
						c.cache.Del(a.gzippedDigest)
					}
				}
			case err := <-c.watcher.Errors:
				if a.CofferEnabled {
					a.ERROR(
						"air: coffer watcher error",
						map[string]interface{}{
							"error": err.Error(),
						},
					)
				}
			}
		}
	}()

	return c
}

// asset returns an `asset` from the c for the name.
func (c *coffer) asset(name string) (*asset, error) {
	c.once.Do(func() {
		c.cache = fastcache.New(c.a.CofferMaxMemoryBytes)
	})

	if ai, ok := c.assets.Load(name); ok {
		return ai.(*asset), nil
	} else if ar, err := filepath.Abs(c.a.AssetRoot); err != nil {
		return nil, err
	} else if !strings.HasPrefix(name, ar) {
		return nil, nil
	}

	ext := filepath.Ext(name)
	if !stringSliceContainsCIly(c.a.AssetExts, ext) {
		return nil, nil
	}

	fi, err := os.Stat(name)
	if err != nil {
		return nil, err
	}

	b, err := ioutil.ReadFile(name)
	if err != nil {
		return nil, err
	}

	var (
		mt       = mime.TypeByExtension(ext)
		minified bool
		gb       []byte
	)

	if mt == "" {
		l := len(b)
		if l > 512 {
			l = 512
		}

		mt = mimesniffer.Sniff(b[:l])
	}

	pmt, _, err := mime.ParseMediaType(mt)
	if err != nil {
		return nil, err
	}

	if c.a.MinifierEnabled &&
		stringSliceContains(c.a.MinifierMIMETypes, pmt) {
		if b, err = c.a.minifier.minify(pmt, b); err != nil {
			return nil, err
		}

		minified = true
	}

	if c.a.GzipEnabled && stringSliceContains(c.a.GzipMIMETypes, pmt) {
		buf := bytes.Buffer{}
		if gw, err := gzip.NewWriterLevel(
			&buf,
			c.a.GzipCompressionLevel,
		); err != nil {
			return nil, err
		} else if _, err = gw.Write(b); err != nil {
			return nil, err
		} else if err = gw.Close(); err != nil {
			return nil, err
		}

		gb = buf.Bytes()
	}

	if err := c.watcher.Add(name); err != nil {
		return nil, err
	}

	a := &asset{
		coffer:   c,
		name:     name,
		mimeType: mt,
		modTime:  fi.ModTime(),
		minified: minified,
		digest:   make([]byte, 8),
	}

	binary.BigEndian.PutUint64(a.digest, xxhash.Sum64(b))
	c.cache.Set(a.digest, b)

	if gb != nil {
		a.gzippedDigest = make([]byte, 8)
		binary.BigEndian.PutUint64(a.gzippedDigest, xxhash.Sum64(gb))
		c.cache.Set(a.gzippedDigest, gb)
	}

	c.assets.Store(name, a)

	return a, nil
}

// asset is a binary asset file.
type asset struct {
	coffer        *coffer
	name          string
	mimeType      string
	modTime       time.Time
	minified      bool
	digest        []byte
	gzippedDigest []byte
}

// content returns the content of the a with the gzipped.
func (a *asset) content(gzipped bool) []byte {
	var c []byte
	if gzipped {
		c = a.coffer.cache.Get(nil, a.gzippedDigest)
	} else {
		c = a.coffer.cache.Get(nil, a.digest)
	}

	if len(c) == 0 {
		a.coffer.assets.Delete(a.name)
		a.coffer.cache.Del(a.digest)
		if a.gzippedDigest != nil {
			a.coffer.cache.Del(a.gzippedDigest)
		}

		return nil
	}

	return c
}
<|endoftext|>


<commit_before>package canbus

import (
	"bytes"
	"fmt"
	"io"
	"sync"

	"golang.org/x/net/context"

	"github.com/sbinet/lsst-ccs/fwk"
)

// Cmd is a type of command to send/receive on/from the CAN bus.
type Cmd string

// The command types known to the CAN bus.
const (
	Boot Cmd = "boot"
	Info     = "info"
	Rsdo     = "rsdo"
	Wsdo     = "wsdo"
	Sync     = "sync"
	Quit     = "quit"
)

type Error struct {
	Code int
}

func (err Error) Error() string {
	return fmt.Sprintf("canbus: error code=%v", err.Code)
}

// Command is a command sent/received on the CAN bus
type Command struct {
	Name Cmd
	Data []byte
}

func (cmd Command) bytes() []byte {
	o := make([]byte, 0, len(cmd.Name)+1+len(cmd.Data))
	o = append(o, []byte(cmd.Name)...)
	if len(cmd.Data) > 0 {
		o = append(o, sepComma...)
		o = append(o, cmd.Data...)
	}
	o = append(o, '\r', 0, '\n')
	return o
}

func (cmd Command) String() string {
	return fmt.Sprintf("Command{%s,%s}", cmd.Name, string(cmd.Data))
}

func (cmd Command) Err() error {
	node := 0
	ecode := 0
	_, err := fmt.Fscanf(bytes.NewReader(cmd.Data),
		"%x,%x",
		&node,
		&ecode,
	)
	if err != nil {
		return err
	}
	if ecode == 0 {
		return nil
	}
	return Error{ecode}
}

func newCommand(data []byte) Command {
	data = bytes.TrimSpace(data)
	if !bytes.Contains(data, sepComma) {
		return Command{}
	}

	tokens := bytes.SplitN(data, sepComma, 2)
	cmd := Command{
		Name: Cmd(tokens[0]),
		Data: tokens[1],
	}
	return cmd
}

type Bus interface {
	ADC() *ADC
	DAC() *DAC
	Send(cmd Command) (Command, error)
}

type busImpl struct {
	*fwk.Base
	conn  cwrapper
	quit  chan struct{}
	nodes []int

	adc *ADC
	dac *DAC

	devices []fwk.Device

	mux  sync.Mutex
	send chan Command
	recv chan Command
}

func New(name string, port int, adc *ADC, dac *DAC, devices ...fwk.Device) fwk.Module {
	devs := append([]fwk.Device{adc, dac}, devices...)
	bus := &busImpl{
		Base:    fwk.NewBase(name),
		conn:    newCwrapperImpl(port),
		quit:    make(chan struct{}),
		nodes:   make([]int, 0, 2),
		adc:     adc,
		dac:     dac,
		send:    make(chan Command),
		recv:    make(chan Command),
		devices: devs,
	}
	fwk.System.Register(bus)
	for _, dev := range bus.devices {
		fwk.System.Register(dev)
	}

	return bus
}

func (bus *busImpl) Boot(ctx context.Context) error {
	bus.Infof(">>> boot...\n")
	var err error

	err = bus.Base.Boot(ctx)
	if err != nil {
		bus.Errorf("error booting: %v\n", err)
		return err
	}

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	err = bus.init()
	if err != nil {
		bus.Errorf("error: %v\n", err)
		return err
	}

	bus.Infof(">>> boot... [done]\n")
	return err
}

func (bus *busImpl) Start(ctx context.Context) error {
	var err error
	return err
}

func (bus *busImpl) Stop(ctx context.Context) error {
	var err error
	bus.Infof("stopping...\n")

	return err
}

func (bus *busImpl) Shutdown(ctx context.Context) error {
	var err error
	bus.Infof("shutdown...\n")

	_, err = bus.Send(Command{Quit, nil})
	if err != nil {
		bus.Errorf("error closing canbus: %v\n", err)
	}

	err = bus.Close()
	if err != nil {
		return err
	}

	return err
}

func (bus *busImpl) init() error {
	var err error

	err = bus.conn.init(bus.Base.Logger)
	if err != nil {
		bus.Errorf("error initializing cwrapper: %v\n", err)
		return err
	}

	const bufsz = 1024
	buf := make([]byte, bufsz)

	// consume welcome message
	n, err := bus.conn.Read(buf)
	if err != nil {
		bus.Errorf("error receiving welcome message: %v\n", err)
		return err
	}
	if n <= 0 {
		bus.Errorf("empty welcome message!\n")
		return io.ErrUnexpectedEOF
	}

	if !bytes.HasPrefix(buf[:n], []byte("TestBench ISO-8859-1")) {
		bus.Errorf("unexpected welcome message: %q\n", string(buf[:n]))
		return io.ErrUnexpectedEOF
	}

	// discover nodes
	for len(bus.nodes) < len(bus.devices) {
		buf = buf[:bufsz]
		n, err := bus.conn.Read(buf)
		if err != nil {
			bus.Errorf("error receiving boot message: %v\n", err)
			return err
		}
		if n <= 0 {
			// nothing was read...
			continue
		}
		buf = buf[:n]
		cmd := newCommand(buf)
		switch cmd.Name {
		case Boot:
			id := 0
			_, err := fmt.Fscanf(bytes.NewReader(cmd.Data), "%x", &id)
			if err != nil {
				bus.Errorf("error decoding node id: %v\n", err)
				return err
			}
			bus.Infof("detected node 0x%x\n", id)
			bus.nodes = append(bus.nodes, id)
		default:
			bus.Errorf("unexpected command name: %q (cmd=%v)\n", cmd.Name, cmd)
			return fmt.Errorf("unexpected command %q", cmd.Name)
		}
	}

	type Node struct {
		id       int
		device   int
		vendor   int
		product  int
		revision int
		serial   string
	}

	nodes := make([]Node, len(bus.nodes))
	// fetch infos about nodes
	for _, id := range bus.nodes {
		buf := []byte(fmt.Sprintf("%s,%x\n", Info, id))
		_, err := bus.conn.Write(buf)
		if err != nil {
			bus.Errorf("error sending info message: %v\n", err)
			return err
		}

		buf = make([]byte, bufsz)
		n, err := bus.conn.Read(buf)
		if err != nil {
			bus.Errorf("error receiving info message: %v\n", err)
			return err
		}
		if n <= 0 {
			// nothing was read...
			continue
		}
		buf = buf[:n]
		cmd := newCommand(buf)
		switch cmd.Name {
		case Info:
			var node Node
			_, err = fmt.Fscanf(
				bytes.NewReader(cmd.Data),
				"%x,%x,%x,%x,%x,%s",
				&node.id,
				&node.device,
				&node.vendor,
				&node.product,
				&node.revision,
				&node.serial,
			)
			if err != nil {
				bus.Errorf("error decoding %v: %v\n", cmd, err)
				return err
			}
			bus.Infof("node=%#v\n", node)
			nodes = append(nodes, node)
			//TODO(sbinet): better/more-general handling
			switch node.serial {
			case bus.adc.serial:
				bus.adc.node = node.id
				bus.adc.bus = bus
			case bus.dac.serial:
				bus.dac.node = node.id
				bus.dac.bus = bus
			}

		default:
			err = fmt.Errorf("unexpected command name: %q (cmd: %v)", cmd.Name, cmd)
			bus.Errorf("error: %v\n", err)
			return err
		}
	}

	bus.Infof("adc=%#v\n", bus.adc)
	bus.Infof("dac=%#v\n", bus.dac)

	err = bus.adc.init()
	if err != nil {
		bus.Errorf("error initializing ADC: %v\n", err)
		return err
	}

	err = bus.dac.init()
	if err != nil {
		bus.Errorf("error initializing DAC: %v\n", err)
		return err
	}

	go bus.run()

	return err
}

func (bus *busImpl) Close() error {
	if bus.conn == nil {
		return nil
	}
	bus.Infof("closing tcp server\n")
	return bus.conn.Close()
}

func (bus *busImpl) run() {
	bus.Infof("handle...\n")

	const bufsz = 1024
	pool := sync.Pool{
		New: func() interface{} {
			return make([]byte, bufsz)
		},
	}

loop:
	for {
		select {
		case cmd := <-bus.send:
			n, err := bus.conn.Write(cmd.bytes())
			if err != nil {
				bus.Errorf("error sending command %v: %v\n", cmd, err)
				return
			}

			switch cmd.Name {
			case Quit:
				bus.Infof("received 'quit' request...\n")
				break loop
			}

			// TODO(sbinet) only read back when needed?
			buf := pool.Get().([]byte)
			buf = buf[:bufsz]
			n, err = bus.conn.Read(buf)
			if err != nil {
				bus.Errorf("error receiving message: %v\n", err)
				break loop
			}
			buf = buf[:n]
			cmd = newCommand(buf)
			bus.recv <- cmd
			pool.Put(buf)

		case <-bus.quit:
			bus.Infof("quit...\n")
			break loop
		}
	}

	close(bus.send)
	close(bus.recv)
	bus.Infof("handle... [done]\n")
}

// Send sends a command down the bus and returns its reply
func (bus *busImpl) Send(icmd Command) (Command, error) {
	var err error

	bus.mux.Lock()
	defer bus.mux.Unlock()

	bus.send <- icmd
	switch icmd.Name {
	case Quit:
		ocmd := <-bus.recv
		return ocmd, err
	}
	ocmd := <-bus.recv

	if ocmd.Name != icmd.Name {
		return ocmd, fmt.Errorf("unexpected command: %v", ocmd)
	}

	ecode := 0
	_, err = fmt.Fscanf(bytes.NewReader(ocmd.Data),
		"%x",
		&ecode,
	)
	if err != nil {
		return ocmd, err
	}

	// need to synchronize bus
	// FIXME(sbinet) figure out what exactly happens.
	if ecode == -1 {
		buf := make([]byte, 1024)
		n, err := bus.conn.Read(buf)
		if err != nil {
			bus.Errorf("error receiving message: %v\n", err)
			return ocmd, err
		}
		buf = buf[:n]
		cmd := newCommand(buf)
		return cmd, err
	}

	return ocmd, err
}

func (bus *busImpl) ADC() *ADC {
	return bus.adc
}

func (bus *busImpl) DAC() *DAC {
	return bus.dac
}

var (
	sepComma = []byte(",")
)
<commit_msg>fwk/drivers/canbus: remove pool introducing goroutine cross-talk<commit_after>package canbus

import (
	"bytes"
	"fmt"
	"io"
	"sync"

	"golang.org/x/net/context"

	"github.com/sbinet/lsst-ccs/fwk"
)

// Cmd is a type of command to send/receive on/from the CAN bus.
type Cmd string

// The command types known to the CAN bus.
const (
	Boot Cmd = "boot"
	Info     = "info"
	Rsdo     = "rsdo"
	Wsdo     = "wsdo"
	Sync     = "sync"
	Quit     = "quit"
)

type Error struct {
	Code int
}

func (err Error) Error() string {
	return fmt.Sprintf("canbus: error code=%v", err.Code)
}

// Command is a command sent/received on the CAN bus
type Command struct {
	Name Cmd
	Data []byte
}

func (cmd Command) bytes() []byte {
	o := make([]byte, 0, len(cmd.Name)+1+len(cmd.Data))
	o = append(o, []byte(cmd.Name)...)
	if len(cmd.Data) > 0 {
		o = append(o, sepComma...)
		o = append(o, cmd.Data...)
	}
	o = append(o, '\r', 0, '\n')
	return o
}

func (cmd Command) String() string {
	return fmt.Sprintf("Command{%s,%s}", cmd.Name, string(cmd.Data))
}

func (cmd Command) Err() error {
	node := 0
	ecode := 0
	_, err := fmt.Fscanf(bytes.NewReader(cmd.Data),
		"%x,%x",
		&node,
		&ecode,
	)
	if err != nil {
		return err
	}
	if ecode == 0 {
		return nil
	}
	return Error{ecode}
}

func newCommand(data []byte) Command {
	data = bytes.TrimSpace(data)
	if !bytes.Contains(data, sepComma) {
		return Command{}
	}

	tokens := bytes.SplitN(data, sepComma, 2)
	cmd := Command{
		Name: Cmd(tokens[0]),
		Data: tokens[1],
	}
	return cmd
}

type Bus interface {
	ADC() *ADC
	DAC() *DAC
	Send(cmd Command) (Command, error)
}

type busImpl struct {
	*fwk.Base
	conn  cwrapper
	quit  chan struct{}
	nodes []int

	adc *ADC
	dac *DAC

	devices []fwk.Device

	mux  sync.Mutex
	send chan Command
	recv chan Command
}

func New(name string, port int, adc *ADC, dac *DAC, devices ...fwk.Device) fwk.Module {
	devs := append([]fwk.Device{adc, dac}, devices...)
	bus := &busImpl{
		Base:    fwk.NewBase(name),
		conn:    newCwrapperImpl(port),
		quit:    make(chan struct{}),
		nodes:   make([]int, 0, 2),
		adc:     adc,
		dac:     dac,
		send:    make(chan Command),
		recv:    make(chan Command),
		devices: devs,
	}
	fwk.System.Register(bus)
	for _, dev := range bus.devices {
		fwk.System.Register(dev)
	}

	return bus
}

func (bus *busImpl) Boot(ctx context.Context) error {
	bus.Infof(">>> boot...\n")
	var err error

	err = bus.Base.Boot(ctx)
	if err != nil {
		bus.Errorf("error booting: %v\n", err)
		return err
	}

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	err = bus.init()
	if err != nil {
		bus.Errorf("error: %v\n", err)
		return err
	}

	bus.Infof(">>> boot... [done]\n")
	return err
}

func (bus *busImpl) Start(ctx context.Context) error {
	var err error
	return err
}

func (bus *busImpl) Stop(ctx context.Context) error {
	var err error
	bus.Infof("stopping...\n")

	return err
}

func (bus *busImpl) Shutdown(ctx context.Context) error {
	var err error
	bus.Infof("shutdown...\n")

	_, err = bus.Send(Command{Quit, nil})
	if err != nil {
		bus.Errorf("error closing canbus: %v\n", err)
	}

	err = bus.Close()
	if err != nil {
		return err
	}

	return err
}

func (bus *busImpl) init() error {
	var err error

	err = bus.conn.init(bus.Base.Logger)
	if err != nil {
		bus.Errorf("error initializing cwrapper: %v\n", err)
		return err
	}

	const bufsz = 1024
	buf := make([]byte, bufsz)

	// consume welcome message
	n, err := bus.conn.Read(buf)
	if err != nil {
		bus.Errorf("error receiving welcome message: %v\n", err)
		return err
	}
	if n <= 0 {
		bus.Errorf("empty welcome message!\n")
		return io.ErrUnexpectedEOF
	}

	if !bytes.HasPrefix(buf[:n], []byte("TestBench ISO-8859-1")) {
		bus.Errorf("unexpected welcome message: %q\n", string(buf[:n]))
		return io.ErrUnexpectedEOF
	}

	// discover nodes
	for len(bus.nodes) < len(bus.devices) {
		buf = buf[:bufsz]
		n, err := bus.conn.Read(buf)
		if err != nil {
			bus.Errorf("error receiving boot message: %v\n", err)
			return err
		}
		if n <= 0 {
			// nothing was read...
			continue
		}
		buf = buf[:n]
		cmd := newCommand(buf)
		switch cmd.Name {
		case Boot:
			id := 0
			_, err := fmt.Fscanf(bytes.NewReader(cmd.Data), "%x", &id)
			if err != nil {
				bus.Errorf("error decoding node id: %v\n", err)
				return err
			}
			bus.Infof("detected node 0x%x\n", id)
			bus.nodes = append(bus.nodes, id)
		default:
			bus.Errorf("unexpected command name: %q (cmd=%v)\n", cmd.Name, cmd)
			return fmt.Errorf("unexpected command %q", cmd.Name)
		}
	}

	type Node struct {
		id       int
		device   int
		vendor   int
		product  int
		revision int
		serial   string
	}

	nodes := make([]Node, len(bus.nodes))
	// fetch infos about nodes
	for _, id := range bus.nodes {
		buf := []byte(fmt.Sprintf("%s,%x\n", Info, id))
		_, err := bus.conn.Write(buf)
		if err != nil {
			bus.Errorf("error sending info message: %v\n", err)
			return err
		}

		buf = make([]byte, bufsz)
		n, err := bus.conn.Read(buf)
		if err != nil {
			bus.Errorf("error receiving info message: %v\n", err)
			return err
		}
		if n <= 0 {
			// nothing was read...
			continue
		}
		buf = buf[:n]
		cmd := newCommand(buf)
		switch cmd.Name {
		case Info:
			var node Node
			_, err = fmt.Fscanf(
				bytes.NewReader(cmd.Data),
				"%x,%x,%x,%x,%x,%s",
				&node.id,
				&node.device,
				&node.vendor,
				&node.product,
				&node.revision,
				&node.serial,
			)
			if err != nil {
				bus.Errorf("error decoding %v: %v\n", cmd, err)
				return err
			}
			bus.Infof("node=%#v\n", node)
			nodes = append(nodes, node)
			//TODO(sbinet): better/more-general handling
			switch node.serial {
			case bus.adc.serial:
				bus.adc.node = node.id
				bus.adc.bus = bus
			case bus.dac.serial:
				bus.dac.node = node.id
				bus.dac.bus = bus
			}

		default:
			err = fmt.Errorf("unexpected command name: %q (cmd: %v)", cmd.Name, cmd)
			bus.Errorf("error: %v\n", err)
			return err
		}
	}

	bus.Infof("adc=%#v\n", bus.adc)
	bus.Infof("dac=%#v\n", bus.dac)

	err = bus.adc.init()
	if err != nil {
		bus.Errorf("error initializing ADC: %v\n", err)
		return err
	}

	err = bus.dac.init()
	if err != nil {
		bus.Errorf("error initializing DAC: %v\n", err)
		return err
	}

	go bus.run()

	return err
}

func (bus *busImpl) Close() error {
	if bus.conn == nil {
		return nil
	}
	bus.Infof("closing tcp server\n")
	return bus.conn.Close()
}

func (bus *busImpl) run() {
	bus.Infof("handle...\n")

	const bufsz = 1024

loop:
	for {
		select {
		case cmd := <-bus.send:
			n, err := bus.conn.Write(cmd.bytes())
			if err != nil {
				bus.Errorf("error sending command %v: %v\n", cmd, err)
				return
			}

			switch cmd.Name {
			case Quit:
				bus.Infof("received 'quit' request...\n")
				break loop
			}

			// TODO(sbinet) only read back when needed?
			buf := make([]byte, bufsz)
			n, err = bus.conn.Read(buf)
			if err != nil {
				bus.Errorf("error receiving message: %v\n", err)
				break loop
			}
			buf = buf[:n]
			cmd = newCommand(buf)
			bus.recv <- cmd

		case <-bus.quit:
			bus.Infof("quit...\n")
			break loop
		}
	}

	close(bus.send)
	close(bus.recv)
	bus.Infof("handle... [done]\n")
}

// Send sends a command down the bus and returns its reply
func (bus *busImpl) Send(icmd Command) (Command, error) {
	var err error

	bus.mux.Lock()
	defer bus.mux.Unlock()

	bus.send <- icmd
	switch icmd.Name {
	case Quit:
		ocmd := <-bus.recv
		return ocmd, err
	}
	ocmd := <-bus.recv

	if ocmd.Name != icmd.Name {
		return ocmd, fmt.Errorf("unexpected command: %v", ocmd)
	}

	ecode := 0
	_, err = fmt.Fscanf(bytes.NewReader(ocmd.Data),
		"%x",
		&ecode,
	)
	if err != nil {
		return ocmd, err
	}

	// need to synchronize bus
	// FIXME(sbinet) figure out what exactly happens.
	if ecode == -1 {
		buf := make([]byte, 1024)
		n, err := bus.conn.Read(buf)
		if err != nil {
			bus.Errorf("error receiving message: %v\n", err)
			return ocmd, err
		}
		buf = buf[:n]
		cmd := newCommand(buf)
		return cmd, err
	}

	return ocmd, err
}

func (bus *busImpl) ADC() *ADC {
	return bus.adc
}

func (bus *busImpl) DAC() *DAC {
	return bus.dac
}

var (
	sepComma = []byte(",")
)
<|endoftext|>


<commit_before>// Copyright (C) 2014 Nippon Telegraph and Telephone Corporation.
//
// 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 server
import (
	"github.com/stretchr/testify/assert"
	"github.com/osrg/gobgp/zebra"
	"github.com/osrg/gobgp/table"
	"testing"
	"net"
)

func Test_createPathFromIPRouteMessage(t *testing.T) {
	assert := assert.New(t)

	m := &zebra.Message{}
	h := &zebra.Header{
		Len: zebra.HEADER_SIZE,
		Marker: zebra.HEADER_MARKER,
		Version: zebra.VERSION,
		Command: zebra.IPV4_ROUTE_ADD,
	}

	b := &zebra.IPRouteBody{
		Type: zebra.ROUTE_TYPE(zebra.ROUTE_STATIC),
		Flags: zebra.FLAG(zebra.FLAG_SELECTED),
		Message: zebra.MESSAGE_NEXTHOP | zebra.MESSAGE_DISTANCE | zebra.MESSAGE_METRIC,
		SAFI:         zebra.SAFI(zebra.SAFI_UNICAST),
		Prefix:       net.ParseIP("192.168.100.0"),
		PrefixLength :uint8(24),
		Nexthops :    []net.IP{net.ParseIP("0.0.0.0")},
		Ifindexs :    []uint32{1},
		Distance:     uint8(0),
		Metric  :     uint32(100),
		Api :         zebra.API_TYPE(zebra.IPV4_ROUTE_ADD),
	}

	m.Header = *h
	m.Body = b

	pi := &table.PeerInfo{
		AS:      65000,
		LocalID: net.ParseIP("10.0.0.1"),
	}
	p := createPathFromIPRouteMessage(m, pi)
	assert.NotEqual(nil, p)
	assert.Equal("0.0.0.0", p.GetNexthop().String())
	assert.Equal("192.168.100.0/24",p.GetNlri().String())
	assert.True(p.IsFromZebra)
	assert.False(p.IsWithdraw)

	// withdraw
	h.Command = zebra.IPV4_ROUTE_DELETE
	m.Header = *h
	p = createPathFromIPRouteMessage(m, pi)
	assert.NotEqual(nil, p)
	assert.Equal("0.0.0.0", p.GetNexthop().String())
	assert.Equal("192.168.100.0/24",p.GetNlri().String())
	med, _ := p.GetMed()
	assert.Equal(uint32(100), med)
	assert.True(p.IsFromZebra)
	assert.True(p.IsWithdraw)

	// IPv6
	h.Command = zebra.IPV6_ROUTE_ADD
	b.Prefix = net.ParseIP("2001:db8:0:f101::")
	b.PrefixLength = uint8(64)
	b.Nexthops =  []net.IP{net.ParseIP("::")}
	m.Header = *h
	m.Body = b

	p = createPathFromIPRouteMessage(m, pi)
	assert.NotEqual(nil, p)
	assert.Equal("::", p.GetNexthop().String())
	assert.Equal("2001:db8:0:f101::/64",p.GetNlri().String())
	med, _ = p.GetMed()
	assert.Equal(uint32(100), med)
	assert.True(p.IsFromZebra)
	assert.False(p.IsWithdraw)

	// withdraw
	h.Command = zebra.IPV6_ROUTE_DELETE
	m.Header = *h
	p = createPathFromIPRouteMessage(m, pi)
	assert.NotEqual(nil, p)
	assert.Equal("::", p.GetNexthop().String())
	assert.Equal("2001:db8:0:f101::/64",p.GetNlri().String())
	assert.True(p.IsFromZebra)
	assert.True(p.IsWithdraw)

}<commit_msg>server: go fmt<commit_after>// Copyright (C) 2014 Nippon Telegraph and Telephone Corporation.
//
// 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 server

import (
	"github.com/osrg/gobgp/table"
	"github.com/osrg/gobgp/zebra"
	"github.com/stretchr/testify/assert"
	"net"
	"testing"
)

func Test_createPathFromIPRouteMessage(t *testing.T) {
	assert := assert.New(t)

	m := &zebra.Message{}
	h := &zebra.Header{
		Len:     zebra.HEADER_SIZE,
		Marker:  zebra.HEADER_MARKER,
		Version: zebra.VERSION,
		Command: zebra.IPV4_ROUTE_ADD,
	}

	b := &zebra.IPRouteBody{
		Type:         zebra.ROUTE_TYPE(zebra.ROUTE_STATIC),
		Flags:        zebra.FLAG(zebra.FLAG_SELECTED),
		Message:      zebra.MESSAGE_NEXTHOP | zebra.MESSAGE_DISTANCE | zebra.MESSAGE_METRIC,
		SAFI:         zebra.SAFI(zebra.SAFI_UNICAST),
		Prefix:       net.ParseIP("192.168.100.0"),
		PrefixLength: uint8(24),
		Nexthops:     []net.IP{net.ParseIP("0.0.0.0")},
		Ifindexs:     []uint32{1},
		Distance:     uint8(0),
		Metric:       uint32(100),
		Api:          zebra.API_TYPE(zebra.IPV4_ROUTE_ADD),
	}

	m.Header = *h
	m.Body = b

	pi := &table.PeerInfo{
		AS:      65000,
		LocalID: net.ParseIP("10.0.0.1"),
	}
	p := createPathFromIPRouteMessage(m, pi)
	assert.NotEqual(nil, p)
	assert.Equal("0.0.0.0", p.GetNexthop().String())
	assert.Equal("192.168.100.0/24", p.GetNlri().String())
	assert.True(p.IsFromZebra)
	assert.False(p.IsWithdraw)

	// withdraw
	h.Command = zebra.IPV4_ROUTE_DELETE
	m.Header = *h
	p = createPathFromIPRouteMessage(m, pi)
	assert.NotEqual(nil, p)
	assert.Equal("0.0.0.0", p.GetNexthop().String())
	assert.Equal("192.168.100.0/24", p.GetNlri().String())
	med, _ := p.GetMed()
	assert.Equal(uint32(100), med)
	assert.True(p.IsFromZebra)
	assert.True(p.IsWithdraw)

	// IPv6
	h.Command = zebra.IPV6_ROUTE_ADD
	b.Prefix = net.ParseIP("2001:db8:0:f101::")
	b.PrefixLength = uint8(64)
	b.Nexthops = []net.IP{net.ParseIP("::")}
	m.Header = *h
	m.Body = b

	p = createPathFromIPRouteMessage(m, pi)
	assert.NotEqual(nil, p)
	assert.Equal("::", p.GetNexthop().String())
	assert.Equal("2001:db8:0:f101::/64", p.GetNlri().String())
	med, _ = p.GetMed()
	assert.Equal(uint32(100), med)
	assert.True(p.IsFromZebra)
	assert.False(p.IsWithdraw)

	// withdraw
	h.Command = zebra.IPV6_ROUTE_DELETE
	m.Header = *h
	p = createPathFromIPRouteMessage(m, pi)
	assert.NotEqual(nil, p)
	assert.Equal("::", p.GetNexthop().String())
	assert.Equal("2001:db8:0:f101::/64", p.GetNlri().String())
	assert.True(p.IsFromZebra)
	assert.True(p.IsWithdraw)

}
<|endoftext|>


<commit_before>package service

import (
	"errors"
	"fmt"
	"strconv"
	"strings"

	"github.com/godbus/dbus"
	"github.com/godbus/dbus/introspect"
	"github.com/muka/go-bluetooth/bluez"
	"github.com/muka/go-bluetooth/bluez/profile"
	"github.com/sirupsen/logrus"
)

//UUIDSuffix fixed 128bit UUID [0000]+[xxxx]+[-0000-1000-8000-00805F9B34FB]
const UUIDSuffix = "-0000-1000-8000-00805F9B34FB"

//NewApplication instantiate a new application service
func NewApplication(config *ApplicationConfig) (*Application, error) {

	if config.ObjectName == "" {
		return nil, errors.New("objectName is required")
	}

	if config.ObjectPath == "" {
		return nil, errors.New("objectPath is required")
	}

	if config.conn == nil {
		conn, err := dbus.SystemBus()
		if err != nil {
			return nil, err
		}
		config.conn = conn
	}

	om, err := NewObjectManager(config.conn)
	if err != nil {
		return nil, err
	}

	// props, err := NewProperties(config.conn)
	// if err != nil {
	// 	return nil, err
	// }

	s := &Application{
		config:        config,
		objectManager: om,
		services:      make(map[dbus.ObjectPath]*GattService1),
	}

	return s, nil
}

//GattWriteCallback A callback we can register to handle write requests
type GattWriteCallback func(app *Application, service_uuid string, charUUID string, value []byte) error

//GattDescriptorWriteCallback A callback we can register to handle descriptor write requests
type GattDescriptorWriteCallback func(app *Application, service_uuid string, charUUID string, descUUID string, value []byte) error

//GattReadCallback A callback we can register to handle read requests
type GattReadCallback func(app *Application, service_uuid string, charUUID string) ([]byte, error)

//GattDescriptorReadCallback A callback we can register to handle descriptor ead requests
type GattDescriptorReadCallback func(app *Application, service_uuid string, charUUID string, descUUID string) ([]byte, error)

// ApplicationConfig configuration for the bluetooth service
type ApplicationConfig struct {
	UUIDSuffix   string
	UUID         string
	conn         *dbus.Conn
	ObjectName   string
	ObjectPath   dbus.ObjectPath
	serviceIndex int
	LocalName    string

	WriteFunc     GattWriteCallback
	ReadFunc      GattReadCallback
	DescWriteFunc GattDescriptorWriteCallback
	DescReadFunc  GattDescriptorReadCallback
}

// Application a bluetooth service exposed by bluez
type Application struct {
	config        *ApplicationConfig
	objectManager *ObjectManager
	services      map[dbus.ObjectPath]*GattService1

	adMgr         *profile.LEAdvertisingManager1
	advertisement *LEAdvertisement1
}

//GetObjectManager return the object manager interface handler
func (app *Application) GetObjectManager() *ObjectManager {
	return app.objectManager
}

//Path return the object path
func (app *Application) Path() dbus.ObjectPath {
	return app.config.ObjectPath
}

//Name return the object name
func (app *Application) Name() string {
	return app.config.ObjectName
}

// GenerateUUID generate a 128bit UUID
func (app *Application) GenerateUUID(uuidVal string) string {
	base := app.config.UUID
	if len(uuidVal) == 8 {
		base = ""
	}
	return base + uuidVal + app.config.UUIDSuffix
}

//CreateService create a new GattService1 instance
func (app *Application) CreateService(props *profile.GattService1Properties, advertisedOptional ...bool) (*GattService1, error) {
	app.config.serviceIndex++
	appPath := string(app.Path())
	if appPath == "/" {
		appPath = ""
	}

	advertise := false
	if len(advertisedOptional) > 0 {
		advertise = advertisedOptional[0]
	}

	path := appPath + "/service" + strconv.Itoa(app.config.serviceIndex)
	c := &GattService1Config{
		app:        app,
		objectPath: dbus.ObjectPath(path),
		ID:         app.config.serviceIndex,
		conn:       app.config.conn,
		advertised: advertise,
	}
	s, err := NewGattService1(c, props)
	return s, err
}

//AddService add service to expose
func (app *Application) AddService(service *GattService1) error {

	app.services[service.Path()] = service

	err := service.Expose()
	if err != nil {
		return err
	}

	err = app.exportTree()
	if err != nil {
		return err
	}

	err = app.GetObjectManager().AddObject(service.Path(), service.Properties())
	if err != nil {
		return err
	}

	return err
}

//RemoveService remove an exposed service
func (app *Application) RemoveService(service *GattService1) error {
	if _, ok := app.services[service.Path()]; ok {

		delete(app.services, service.Path())
		err := app.GetObjectManager().RemoveObject(service.Path())

		//TODO: remove chars + descritptors too
		if err != nil {
			return err
		}

		err = app.exportTree()
		if err != nil {
			return err
		}
	}
	return nil
}

//GetServices return the registered services
func (app *Application) GetServices() map[dbus.ObjectPath]*GattService1 {
	return app.services
}

//expose dbus interfaces
func (app *Application) expose() error {

	conn := app.config.conn
	_, err := conn.RequestName(app.Name(), dbus.NameFlagDoNotQueue&dbus.NameFlagReplaceExisting)
	if err != nil {
		return err
	}

	// / path
	err = conn.Export(app.objectManager, app.Path(), bluez.ObjectManagerInterface)
	if err != nil {
		return err
	}

	err = app.exportTree()
	if err != nil {
		return err
	}

	return nil
}

func (app *Application) exportTree() error {

	childrenNode := make([]introspect.Node, 0)

	for servicePath, service := range app.GetServices() {
		childrenNode = append(childrenNode, introspect.Node{
			Name: string(servicePath)[1:],
		})
		for charPath, char := range service.GetCharacteristics() {
			childrenNode = append(childrenNode, introspect.Node{
				Name: string(charPath)[1:],
			})
			for descPath := range char.GetDescriptors() {
				childrenNode = append(childrenNode, introspect.Node{
					Name: string(descPath)[1:],
				})
			}
		}
	}

	// must include also child nodes
	node := &introspect.Node{
		Interfaces: []introspect.Interface{
			//Introspect
			introspect.IntrospectData,
			//ObjectManager
			bluez.ObjectManagerIntrospectData,
		},
		Children: childrenNode,
	}

	err := app.config.conn.Export(
		introspect.NewIntrospectable(node),
		app.Path(),
		"org.freedesktop.DBus.Introspectable")

	return err
}

// CallbackError error from a callback
type CallbackError struct {
	msg  string
	code int
}

func (e *CallbackError) Error() string {
	return e.msg
}

//NewCallbackError create a new callback error
func NewCallbackError(code int, msg string) *CallbackError {
	result := &CallbackError{msg: msg, code: code}
	return result
}

//CallbackNotRegistered callback not registered
const CallbackNotRegistered = -1

//CallbackFunctionError callback reported an error
const CallbackFunctionError = -2

//HandleRead Handle application read
func (app *Application) HandleRead(srvUUID string, uuid string) ([]byte, *CallbackError) {
	if app.config.ReadFunc == nil {
		b := make([]byte, 0)
		return b, NewCallbackError(CallbackNotRegistered, "No callback registered.")
	}

	var cberr *CallbackError
	b, err := app.config.ReadFunc(app, srvUUID, uuid)
	if err != nil {
		if err, ok := err.(*CallbackError); ok {
			// If a CallbackError is returned, simply pass it through
			cberr = err
		} else {
			cberr = NewCallbackError(CallbackFunctionError, err.Error())
		}
	}

	return b, cberr
}

// HandleWrite handle application write
func (app *Application) HandleWrite(srvUUID string, uuid string, value []byte) *CallbackError {
	if app.config.WriteFunc == nil {
		return NewCallbackError(CallbackNotRegistered, "No callback registered.")
	}

	err := app.config.WriteFunc(app, srvUUID, uuid, value)
	if err != nil {
		if err, ok := err.(*CallbackError); ok {
			// If a CallbackError is returned, simply pass it through
			return err
		}

		return NewCallbackError(CallbackFunctionError, err.Error())
	}

	return nil
}

//HandleDescriptorRead handle descriptor read
func (app *Application) HandleDescriptorRead(srvUUID string, charUUID string, descUUID string) ([]byte, *CallbackError) {
	if app.config.DescReadFunc == nil {
		b := make([]byte, 0)
		return b, NewCallbackError(CallbackNotRegistered, "No callback registered.")
	}

	var cberr *CallbackError
	b, err := app.config.DescReadFunc(app, srvUUID, charUUID, descUUID)
	if err != nil {
		if err, ok := err.(*CallbackError); ok {
			// If a CallbackError is returned, simply pass it through
			cberr = err
		} else {
			cberr = NewCallbackError(CallbackFunctionError, err.Error())
		}
	}

	return b, cberr
}

//HandleDescriptorWrite handle descriptor write
func (app *Application) HandleDescriptorWrite(srvUUID string, charUUID string, descUUID string, value []byte) *CallbackError {
	if app.config.DescWriteFunc == nil {
		return NewCallbackError(CallbackNotRegistered, "No callback registered.")
	}

	err := app.config.DescWriteFunc(app, srvUUID, charUUID, descUUID, value)
	if err != nil {
		if err, ok := err.(*CallbackError); ok {
			// If a CallbackError is returned, simply pass it through
			return err
		}

		return NewCallbackError(CallbackFunctionError, err.Error())
	}

	return nil
}

//Run start the application
func (app *Application) Run() error {

	err := app.expose()
	if err != nil {
		return err
	}

	return nil
}

//StartAdvertising advertise information for a service
func (app *Application) StartAdvertising(deviceInterface string) error {

	if app.advertisement != nil || app.adMgr != nil {
		logrus.Debugf("Already advertising on %s", deviceInterface)
		return nil
	}

	path := fmt.Sprintf("/org/bluez/advertisement/%s", deviceInterface)

	logrus.Debugf("Registering service on %s (%s)", deviceInterface, path)

	config := &LEAdvertisement1Config{
		conn:       app.config.conn,
		objectPath: dbus.ObjectPath(path),
	}

	serviceUUIDs := make([]string, 0)

	for _, serv := range app.services {
		if serv.Advertised() {
			serviceUUIDs = append(serviceUUIDs, serv.properties.UUID)
		}
	}

	if len([]byte(strings.Join(serviceUUIDs, ","))) > 31 {
		logrus.Warn("Advertisment limit of 31 bytes may have been exceeded. Consider not exposing all services IDs with `CreateService(props, false)`")
	}

	props := &profile.LEAdvertisement1Properties{
		Type:         "peripheral",
		LocalName:    app.config.LocalName,
		ServiceUUIDs: serviceUUIDs,
		// Appearance:   0,
	}

	var err error

	advertisement, err := NewLEAdvertisement1(config, props)
	if err != nil {
		return fmt.Errorf("NewLEAdvertisement1: %s", err)
	}

	app.advertisement = advertisement

	err = advertisement.Expose()
	if err != nil {
		app.advertisement = nil
		return fmt.Errorf("Expose: %s", err)
	}

	options := make(map[string]interface{})

	app.adMgr = profile.NewLEAdvertisingManager1(deviceInterface)

	err = app.adMgr.RegisterAdvertisement(path, options)
	if err != nil {
		app.advertisement = nil
		app.adMgr = nil
		return fmt.Errorf("RegisterAdvertisement: %s", err)
	}

	adapter := profile.NewAdapter1(deviceInterface)

	err = adapter.SetProperty("Discoverable", true)
	if err != nil {
		return err
	}

	err = adapter.SetProperty("Powered", true)
	if err != nil {
		return err
	}

	return nil
}

//StopAdvertising stop advertising information on a service
func (app *Application) StopAdvertising() error {

	if app.advertisement == nil || app.adMgr == nil {
		// Not advertising
		return nil
	}

	err := app.adMgr.UnregisterAdvertisement(string(app.advertisement.config.objectPath))

	app.advertisement.Release()

	app.advertisement = nil
	app.adMgr = nil

	return err
}
<commit_msg>Don't override err in else clause of type assertion<commit_after>package service

import (
	"errors"
	"fmt"
	"strconv"
	"strings"

	"github.com/godbus/dbus"
	"github.com/godbus/dbus/introspect"
	"github.com/muka/go-bluetooth/bluez"
	"github.com/muka/go-bluetooth/bluez/profile"
	"github.com/sirupsen/logrus"
)

//UUIDSuffix fixed 128bit UUID [0000]+[xxxx]+[-0000-1000-8000-00805F9B34FB]
const UUIDSuffix = "-0000-1000-8000-00805F9B34FB"

//NewApplication instantiate a new application service
func NewApplication(config *ApplicationConfig) (*Application, error) {

	if config.ObjectName == "" {
		return nil, errors.New("objectName is required")
	}

	if config.ObjectPath == "" {
		return nil, errors.New("objectPath is required")
	}

	if config.conn == nil {
		conn, err := dbus.SystemBus()
		if err != nil {
			return nil, err
		}
		config.conn = conn
	}

	om, err := NewObjectManager(config.conn)
	if err != nil {
		return nil, err
	}

	// props, err := NewProperties(config.conn)
	// if err != nil {
	// 	return nil, err
	// }

	s := &Application{
		config:        config,
		objectManager: om,
		services:      make(map[dbus.ObjectPath]*GattService1),
	}

	return s, nil
}

//GattWriteCallback A callback we can register to handle write requests
type GattWriteCallback func(app *Application, service_uuid string, charUUID string, value []byte) error

//GattDescriptorWriteCallback A callback we can register to handle descriptor write requests
type GattDescriptorWriteCallback func(app *Application, service_uuid string, charUUID string, descUUID string, value []byte) error

//GattReadCallback A callback we can register to handle read requests
type GattReadCallback func(app *Application, service_uuid string, charUUID string) ([]byte, error)

//GattDescriptorReadCallback A callback we can register to handle descriptor ead requests
type GattDescriptorReadCallback func(app *Application, service_uuid string, charUUID string, descUUID string) ([]byte, error)

// ApplicationConfig configuration for the bluetooth service
type ApplicationConfig struct {
	UUIDSuffix   string
	UUID         string
	conn         *dbus.Conn
	ObjectName   string
	ObjectPath   dbus.ObjectPath
	serviceIndex int
	LocalName    string

	WriteFunc     GattWriteCallback
	ReadFunc      GattReadCallback
	DescWriteFunc GattDescriptorWriteCallback
	DescReadFunc  GattDescriptorReadCallback
}

// Application a bluetooth service exposed by bluez
type Application struct {
	config        *ApplicationConfig
	objectManager *ObjectManager
	services      map[dbus.ObjectPath]*GattService1

	adMgr         *profile.LEAdvertisingManager1
	advertisement *LEAdvertisement1
}

//GetObjectManager return the object manager interface handler
func (app *Application) GetObjectManager() *ObjectManager {
	return app.objectManager
}

//Path return the object path
func (app *Application) Path() dbus.ObjectPath {
	return app.config.ObjectPath
}

//Name return the object name
func (app *Application) Name() string {
	return app.config.ObjectName
}

// GenerateUUID generate a 128bit UUID
func (app *Application) GenerateUUID(uuidVal string) string {
	base := app.config.UUID
	if len(uuidVal) == 8 {
		base = ""
	}
	return base + uuidVal + app.config.UUIDSuffix
}

//CreateService create a new GattService1 instance
func (app *Application) CreateService(props *profile.GattService1Properties, advertisedOptional ...bool) (*GattService1, error) {
	app.config.serviceIndex++
	appPath := string(app.Path())
	if appPath == "/" {
		appPath = ""
	}

	advertise := false
	if len(advertisedOptional) > 0 {
		advertise = advertisedOptional[0]
	}

	path := appPath + "/service" + strconv.Itoa(app.config.serviceIndex)
	c := &GattService1Config{
		app:        app,
		objectPath: dbus.ObjectPath(path),
		ID:         app.config.serviceIndex,
		conn:       app.config.conn,
		advertised: advertise,
	}
	s, err := NewGattService1(c, props)
	return s, err
}

//AddService add service to expose
func (app *Application) AddService(service *GattService1) error {

	app.services[service.Path()] = service

	err := service.Expose()
	if err != nil {
		return err
	}

	err = app.exportTree()
	if err != nil {
		return err
	}

	err = app.GetObjectManager().AddObject(service.Path(), service.Properties())
	if err != nil {
		return err
	}

	return err
}

//RemoveService remove an exposed service
func (app *Application) RemoveService(service *GattService1) error {
	if _, ok := app.services[service.Path()]; ok {

		delete(app.services, service.Path())
		err := app.GetObjectManager().RemoveObject(service.Path())

		//TODO: remove chars + descritptors too
		if err != nil {
			return err
		}

		err = app.exportTree()
		if err != nil {
			return err
		}
	}
	return nil
}

//GetServices return the registered services
func (app *Application) GetServices() map[dbus.ObjectPath]*GattService1 {
	return app.services
}

//expose dbus interfaces
func (app *Application) expose() error {

	conn := app.config.conn
	_, err := conn.RequestName(app.Name(), dbus.NameFlagDoNotQueue&dbus.NameFlagReplaceExisting)
	if err != nil {
		return err
	}

	// / path
	err = conn.Export(app.objectManager, app.Path(), bluez.ObjectManagerInterface)
	if err != nil {
		return err
	}

	err = app.exportTree()
	if err != nil {
		return err
	}

	return nil
}

func (app *Application) exportTree() error {

	childrenNode := make([]introspect.Node, 0)

	for servicePath, service := range app.GetServices() {
		childrenNode = append(childrenNode, introspect.Node{
			Name: string(servicePath)[1:],
		})
		for charPath, char := range service.GetCharacteristics() {
			childrenNode = append(childrenNode, introspect.Node{
				Name: string(charPath)[1:],
			})
			for descPath := range char.GetDescriptors() {
				childrenNode = append(childrenNode, introspect.Node{
					Name: string(descPath)[1:],
				})
			}
		}
	}

	// must include also child nodes
	node := &introspect.Node{
		Interfaces: []introspect.Interface{
			//Introspect
			introspect.IntrospectData,
			//ObjectManager
			bluez.ObjectManagerIntrospectData,
		},
		Children: childrenNode,
	}

	err := app.config.conn.Export(
		introspect.NewIntrospectable(node),
		app.Path(),
		"org.freedesktop.DBus.Introspectable")

	return err
}

// CallbackError error from a callback
type CallbackError struct {
	msg  string
	code int
}

func (e *CallbackError) Error() string {
	return e.msg
}

//NewCallbackError create a new callback error
func NewCallbackError(code int, msg string) *CallbackError {
	result := &CallbackError{msg: msg, code: code}
	return result
}

//CallbackNotRegistered callback not registered
const CallbackNotRegistered = -1

//CallbackFunctionError callback reported an error
const CallbackFunctionError = -2

//HandleRead Handle application read
func (app *Application) HandleRead(srvUUID string, uuid string) ([]byte, *CallbackError) {
	if app.config.ReadFunc == nil {
		b := make([]byte, 0)
		return b, NewCallbackError(CallbackNotRegistered, "No callback registered.")
	}

	var cberr *CallbackError
	b, err := app.config.ReadFunc(app, srvUUID, uuid)
	if err != nil {
		if callbackErr, ok := err.(*CallbackError); ok {
			// If a CallbackError is returned, simply pass it through
			cberr = callbackErr
		} else {
			cberr = NewCallbackError(CallbackFunctionError, err.Error())
		}
	}

	return b, cberr
}

// HandleWrite handle application write
func (app *Application) HandleWrite(srvUUID string, uuid string, value []byte) *CallbackError {
	if app.config.WriteFunc == nil {
		return NewCallbackError(CallbackNotRegistered, "No callback registered.")
	}

	err := app.config.WriteFunc(app, srvUUID, uuid, value)
	if err != nil {
		if callbackErr, ok := err.(*CallbackError); ok {
			// If a CallbackError is returned, simply pass it through
			return callbackErr
		}

		return NewCallbackError(CallbackFunctionError, err.Error())
	}

	return nil
}

//HandleDescriptorRead handle descriptor read
func (app *Application) HandleDescriptorRead(srvUUID string, charUUID string, descUUID string) ([]byte, *CallbackError) {
	if app.config.DescReadFunc == nil {
		b := make([]byte, 0)
		return b, NewCallbackError(CallbackNotRegistered, "No callback registered.")
	}

	var cberr *CallbackError
	b, err := app.config.DescReadFunc(app, srvUUID, charUUID, descUUID)
	if err != nil {
		if callbackErr, ok := err.(*CallbackError); ok {
			// If a CallbackError is returned, simply pass it through
			cberr = callbackErr
		} else {
			cberr = NewCallbackError(CallbackFunctionError, err.Error())
		}
	}

	return b, cberr
}

//HandleDescriptorWrite handle descriptor write
func (app *Application) HandleDescriptorWrite(srvUUID string, charUUID string, descUUID string, value []byte) *CallbackError {
	if app.config.DescWriteFunc == nil {
		return NewCallbackError(CallbackNotRegistered, "No callback registered.")
	}

	err := app.config.DescWriteFunc(app, srvUUID, charUUID, descUUID, value)
	if err != nil {
		if callbackErr, ok := err.(*CallbackError); ok {
			// If a CallbackError is returned, simply pass it through
			return callbackErr
		}

		return NewCallbackError(CallbackFunctionError, err.Error())
	}

	return nil
}

//Run start the application
func (app *Application) Run() error {

	err := app.expose()
	if err != nil {
		return err
	}

	return nil
}

//StartAdvertising advertise information for a service
func (app *Application) StartAdvertising(deviceInterface string) error {

	if app.advertisement != nil || app.adMgr != nil {
		logrus.Debugf("Already advertising on %s", deviceInterface)
		return nil
	}

	path := fmt.Sprintf("/org/bluez/advertisement/%s", deviceInterface)

	logrus.Debugf("Registering service on %s (%s)", deviceInterface, path)

	config := &LEAdvertisement1Config{
		conn:       app.config.conn,
		objectPath: dbus.ObjectPath(path),
	}

	serviceUUIDs := make([]string, 0)

	for _, serv := range app.services {
		if serv.Advertised() {
			serviceUUIDs = append(serviceUUIDs, serv.properties.UUID)
		}
	}

	if len([]byte(strings.Join(serviceUUIDs, ","))) > 31 {
		logrus.Warn("Advertisment limit of 31 bytes may have been exceeded. Consider not exposing all services IDs with `CreateService(props, false)`")
	}

	props := &profile.LEAdvertisement1Properties{
		Type:         "peripheral",
		LocalName:    app.config.LocalName,
		ServiceUUIDs: serviceUUIDs,
		// Appearance:   0,
	}

	var err error

	advertisement, err := NewLEAdvertisement1(config, props)
	if err != nil {
		return fmt.Errorf("NewLEAdvertisement1: %s", err)
	}

	app.advertisement = advertisement

	err = advertisement.Expose()
	if err != nil {
		app.advertisement = nil
		return fmt.Errorf("Expose: %s", err)
	}

	options := make(map[string]interface{})

	app.adMgr = profile.NewLEAdvertisingManager1(deviceInterface)

	err = app.adMgr.RegisterAdvertisement(path, options)
	if err != nil {
		app.advertisement = nil
		app.adMgr = nil
		return fmt.Errorf("RegisterAdvertisement: %s", err)
	}

	adapter := profile.NewAdapter1(deviceInterface)

	err = adapter.SetProperty("Discoverable", true)
	if err != nil {
		return err
	}

	err = adapter.SetProperty("Powered", true)
	if err != nil {
		return err
	}

	return nil
}

//StopAdvertising stop advertising information on a service
func (app *Application) StopAdvertising() error {

	if app.advertisement == nil || app.adMgr == nil {
		// Not advertising
		return nil
	}

	err := app.adMgr.UnregisterAdvertisement(string(app.advertisement.config.objectPath))

	app.advertisement.Release()

	app.advertisement = nil
	app.adMgr = nil

	return err
}
<|endoftext|>


<commit_before>package awsauth

import (
	"bufio"
	"bytes"
	"crypto/hmac"
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

func serviceAndRegion(host string) (string, string) {
	var region, service string
	parts := strings.Split(host, ".")

	service = parts[0]

	if len(parts) >= 4 {
		if parts[1] == "s3" {
			region = parts[0]
			service = parts[1]
		} else {
			region = parts[1]
		}
	} else {
		if strings.HasPrefix(parts[0], "s3-") {
			service = parts[0][:2]
			region = parts[0][3:]
		} else {
			region = "us-east-1" // default. http://docs.aws.amazon.com/general/latest/gr/rande.html
		}
	}

	return service, region
}

func checkKeys() {
	if Keys == nil {
		Keys = &Credentials{
			AccessKeyID:     os.Getenv(envAccessKeyID),
			SecretAccessKey: os.Getenv(envSecretAccessKey),
			SecurityToken:   os.Getenv(envSecurityToken),
		}
	}
	// if accesskey and the secretkey are blank, get the key from the role
	if Keys.AccessKeyID == "" && onEC2() {

		Keys = getIAMRoleCredentials()
	}

	// if the expiration is set and it's less than 5 minutes in the future, get a new key
	if Keys.expired() {
		Keys = getIAMRoleCredentials()
	}
}

func onEC2() bool {

	c, err := net.DialTimeout("tcp", "169.254.169.254:80", time.Second)

	if err != nil {
		return false
	} else {
		c.Close()
		return true
	}
}

func getIAMRoleList() []string {
	// Get a list of the roles that are available to this instance
	url := "http://169.254.169.254/latest/meta-data/iam/security-credentials/"
	client := &http.Client{}
	req, _ := http.NewRequest("GET", url, nil)
	resp, _ := client.Do(req)

	// buf := new(bytes.Buffer)
	// buf.ReadFrom(resp.Body)
	// role := buf.String()

	var roles []string

	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		roles = append(roles, scanner.Text())
	}
	return roles
}

func getIAMRoleCredentials() *Credentials {

	roles := getIAMRoleList()

	if len(roles) < 1 {
		return &Credentials{}
	}

	// Use the first role in the list
	role := roles[0]

	url := "http://169.254.169.254/latest/meta-data/iam/security-credentials/"

	// Create the full URL of the role
	var buffer bytes.Buffer
	buffer.WriteString(url)
	buffer.WriteString(role)
	roleurl := buffer.String()

	// Get the role
	rolereq, _ := http.NewRequest("GET", roleurl, nil)
	roleresp, _ := client.Do(rolereq)
	rolebuf := new(bytes.Buffer)
	rolebuf.ReadFrom(roleresp.Body)

	creds := Credentials{}

	_ = json.Unmarshal(rolebuf.Bytes(), &creds)

	return &creds

}

func augmentRequestQuery(req *http.Request, values url.Values) *http.Request {
	for key, arr := range req.URL.Query() {
		for _, val := range arr {
			values.Set(key, val)
		}
	}

	req.URL.RawQuery = values.Encode()

	return req
}

func hmacSHA256(key []byte, content string) []byte {
	mac := hmac.New(sha256.New, key)
	mac.Write([]byte(content))
	return mac.Sum(nil)
}

func hmacSHA1(key []byte, content string) []byte {
	mac := hmac.New(sha1.New, key)
	mac.Write([]byte(content))
	return mac.Sum(nil)
}

func hashSHA256(content []byte) string {
	h := sha256.New()
	h.Write(content)
	return fmt.Sprintf("%x", h.Sum(nil))
}

func hashMD5(content []byte) string {
	h := md5.New()
	h.Write(content)
	return base64.StdEncoding.EncodeToString(h.Sum(nil))
}

func readAndReplaceBody(req *http.Request) []byte {
	if req.Body == nil {
		return []byte{}
	}
	payload, _ := ioutil.ReadAll(req.Body)
	req.Body = ioutil.NopCloser(bytes.NewReader(payload))
	return payload
}

func concat(delim string, str ...string) string {
	return strings.Join(str, delim)
}

var now = func() time.Time {
	return time.Now().UTC()
}
<commit_msg>adding error handling<commit_after>package awsauth

import (
	"bufio"
	"bytes"
	"crypto/hmac"
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

func serviceAndRegion(host string) (string, string) {
	var region, service string
	parts := strings.Split(host, ".")

	service = parts[0]

	if len(parts) >= 4 {
		if parts[1] == "s3" {
			region = parts[0]
			service = parts[1]
		} else {
			region = parts[1]
		}
	} else {
		if strings.HasPrefix(parts[0], "s3-") {
			service = parts[0][:2]
			region = parts[0][3:]
		} else {
			region = "us-east-1" // default. http://docs.aws.amazon.com/general/latest/gr/rande.html
		}
	}

	return service, region
}

func checkKeys() {
	if Keys == nil {
		Keys = &Credentials{
			AccessKeyID:     os.Getenv(envAccessKeyID),
			SecretAccessKey: os.Getenv(envSecretAccessKey),
			SecurityToken:   os.Getenv(envSecurityToken),
		}
	}
	// if accesskey and the secretkey are blank, get the key from the role
	if Keys.AccessKeyID == "" && onEC2() {

		Keys = getIAMRoleCredentials()
	}

	// if the expiration is set and it's less than 5 minutes in the future, get a new key
	if Keys.expired() && onEC2() {
		Keys = getIAMRoleCredentials()
	}
}

func onEC2() bool {

	c, err := net.DialTimeout("tcp", "169.254.169.254:80", time.Second)

	if err != nil {
		return false
	} else {
		c.Close()
		return true
	}
}

// getIAMRoleList gets a list of the roles that are available to this instance
func getIAMRoleList() []string {

	var roles []string
	url := "http://169.254.169.254/latest/meta-data/iam/security-credentials/"

	client := &http.Client{}

	req, err := http.NewRequest("GET", url, nil)

	if err != nil {
		return roles
	}

	resp, err := client.Do(req)

	if err != nil {
		return roles
	}

	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		roles = append(roles, scanner.Text())
	}
	return roles
}

func getIAMRoleCredentials() *Credentials {

	roles := getIAMRoleList()

	if len(roles) < 1 {
		return &Credentials{}
	}

	// Use the first role in the list
	role := roles[0]

	url := "http://169.254.169.254/latest/meta-data/iam/security-credentials/"

	// Create the full URL of the role
	var buffer bytes.Buffer
	buffer.WriteString(url)
	buffer.WriteString(role)
	roleurl := buffer.String()

	// Get the role
	rolereq, err := http.NewRequest("GET", roleurl, nil)

	if err != nil {
		return &Credentials{}
	}

	roleresp, err := client.Do(rolereq)

	if err != nil {
		return &Credentials{}
	}

	rolebuf := new(bytes.Buffer)
	rolebuf.ReadFrom(roleresp.Body)

	creds := Credentials{}

	err = json.Unmarshal(rolebuf.Bytes(), &creds)

	if err != nil {
		return &Credentials{}
	}

	return &creds

}

func augmentRequestQuery(req *http.Request, values url.Values) *http.Request {
	for key, arr := range req.URL.Query() {
		for _, val := range arr {
			values.Set(key, val)
		}
	}

	req.URL.RawQuery = values.Encode()

	return req
}

func hmacSHA256(key []byte, content string) []byte {
	mac := hmac.New(sha256.New, key)
	mac.Write([]byte(content))
	return mac.Sum(nil)
}

func hmacSHA1(key []byte, content string) []byte {
	mac := hmac.New(sha1.New, key)
	mac.Write([]byte(content))
	return mac.Sum(nil)
}

func hashSHA256(content []byte) string {
	h := sha256.New()
	h.Write(content)
	return fmt.Sprintf("%x", h.Sum(nil))
}

func hashMD5(content []byte) string {
	h := md5.New()
	h.Write(content)
	return base64.StdEncoding.EncodeToString(h.Sum(nil))
}

func readAndReplaceBody(req *http.Request) []byte {
	if req.Body == nil {
		return []byte{}
	}
	payload, _ := ioutil.ReadAll(req.Body)
	req.Body = ioutil.NopCloser(bytes.NewReader(payload))
	return payload
}

func concat(delim string, str ...string) string {
	return strings.Join(str, delim)
}

var now = func() time.Time {
	return time.Now().UTC()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"

	"github.com/BurntSushi/toml"
)

// Version returns the version of grimd
const Version = "0.0.1"

type config struct {
	Sources          []string
	Log              string
	LogLevel         int
	Bind             string
	API              string
	Nullroute        string
	Nullroutev6      string
	Nameservers      []string
	Interval         int
	Timeout          int
	Expire           int
	Maxcount         int
	QuestionCacheCap int
	TTL              uint32
	Blocklist        []string
	Whitelist        []string
}

const defaultConfig = `# list of sources to pull adlists from
sources = [
"http://mirror1.malwaredomains.com/files/justdomains",
"https://raw.githubusercontent.com/StevenBlack/hosts/master/hosts",
"http://sysctl.org/cameleon/hosts",
"https://zeustracker.abuse.ch/blocklist.php?download=domainblocklist",
"https://s3.amazonaws.com/lists.disconnect.me/simple_tracking.txt",
"https://s3.amazonaws.com/lists.disconnect.me/simple_ad.txt",
"http://hosts-file.net/ad_servers.txt",
"https://raw.githubusercontent.com/quidsup/notrack/master/trackers.txt"
]

# location of the log file
log = "grimd.log"

# what kind of information should be logged, 0 = errors and important operations, 1 = dns queries, 2 = debug
loglevel = 0

# address to bind to for the DNS server
bind = "0.0.0.0:53"

# address to bind to for the API server
api = "127.0.0.1:8080"

# ipv4 address to forward blocked queries to
nullroute = "0.0.0.0"

# ipv6 address to forward blocked queries to
nullroutev6 = "0:0:0:0:0:0:0:0"

# nameservers to forward queries to
nameservers = ["8.8.8.8:53", "8.8.4.4:53"]

# concurrency interval for lookups in miliseconds
interval = 200

# query timeout for dns lookups in seconds
timeout = 5

# cache entry lifespan in seconds
expire = 600

# cache capacity, 0 for infinite
maxcount = 0

# question cache capacity, 0 for infinite but not recommended (this is used for storing logs)
questioncachecap = 1000

# manual blocklist entries
blocklist = []

# manual whitelist entries (TODO: change to seperate host formatted file)
whitelist = [
	"getsentry.com",
	"www.getsentry.com"
]
`

// Config is the global configuration
var Config config

// LoadConfig loads the given config file
func LoadConfig(path string) error {
	if _, err := os.Stat(path); os.IsNotExist(err) {
		if err := generateConfig(path); err != nil {
			return err
		}
	}

	if _, err := toml.DecodeFile(path, &Config); err != nil {
		return fmt.Errorf("could not load config: %s", err)
	}

	return nil
}

func generateConfig(path string) error {
	output, err := os.Create(path)
	if err != nil {
		return fmt.Errorf("could not generate config: %s", err)
	}
	defer output.Close()

	r := strings.NewReader(defaultConfig)
	if _, err := io.Copy(output, r); err != nil {
		return fmt.Errorf("could not copy default config: %s", err)
	}

	if abs, err := filepath.Abs(path); err == nil {
		log.Printf("generated default config %s\n", abs)
	}

	return nil
}
<commit_msg>updated config<commit_after>package main

import (
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"

	"github.com/BurntSushi/toml"
)

// Version returns the version of grimd
const Version = "0.0.1"

type config struct {
	Sources          []string
	Log              string
	LogLevel         int
	Bind             string
	API              string
	Nullroute        string
	Nullroutev6      string
	Nameservers      []string
	Interval         int
	Timeout          int
	Expire           int
	Maxcount         int
	QuestionCacheCap int
	TTL              uint32
	Blocklist        []string
	Whitelist        []string
}

const defaultConfig = `# list of sources to pull blocklists from
sources = [
"http://mirror1.malwaredomains.com/files/justdomains",
"https://raw.githubusercontent.com/StevenBlack/hosts/master/hosts",
"http://sysctl.org/cameleon/hosts",
"https://zeustracker.abuse.ch/blocklist.php?download=domainblocklist",
"https://s3.amazonaws.com/lists.disconnect.me/simple_tracking.txt",
"https://s3.amazonaws.com/lists.disconnect.me/simple_ad.txt",
"http://hosts-file.net/ad_servers.txt",
"https://raw.githubusercontent.com/quidsup/notrack/master/trackers.txt"
]

# location of the log file
log = "grimd.log"

# what kind of information should be logged, 0 = errors and important operations, 1 = dns queries, 2 = debug
loglevel = 0

# address to bind to for the DNS server
bind = "0.0.0.0:53"

# address to bind to for the API server
api = "127.0.0.1:8080"

# ipv4 address to forward blocked queries to
nullroute = "0.0.0.0"

# ipv6 address to forward blocked queries to
nullroutev6 = "0:0:0:0:0:0:0:0"

# nameservers to forward queries to
nameservers = ["8.8.8.8:53", "8.8.4.4:53"]

# concurrency interval for lookups in miliseconds
interval = 200

# query timeout for dns lookups in seconds
timeout = 5

# cache entry lifespan in seconds
expire = 600

# cache capacity, 0 for infinite
maxcount = 0

# question cache capacity, 0 for infinite but not recommended (this is used for storing logs)
questioncachecap = 1000

# manual blocklist entries
blocklist = []

# manual whitelist entries (TODO: change to seperate host formatted file)
whitelist = [
	"getsentry.com",
	"www.getsentry.com"
]
`

// Config is the global configuration
var Config config

// LoadConfig loads the given config file
func LoadConfig(path string) error {
	if _, err := os.Stat(path); os.IsNotExist(err) {
		if err := generateConfig(path); err != nil {
			return err
		}
	}

	if _, err := toml.DecodeFile(path, &Config); err != nil {
		return fmt.Errorf("could not load config: %s", err)
	}

	return nil
}

func generateConfig(path string) error {
	output, err := os.Create(path)
	if err != nil {
		return fmt.Errorf("could not generate config: %s", err)
	}
	defer output.Close()

	r := strings.NewReader(defaultConfig)
	if _, err := io.Copy(output, r); err != nil {
		return fmt.Errorf("could not copy default config: %s", err)
	}

	if abs, err := filepath.Abs(path); err == nil {
		log.Printf("generated default config %s\n", abs)
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"log"
	"os"
	"time"

	"github.com/hashicorp/go-cleanhttp"

	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
	"github.com/aws/aws-sdk-go/aws/defaults"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"

	alks "github.com/Cox-Automotive/alks-go"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sts"
)

// Version number, to be injected at link time
// to set, add `-ldflags "-X main.versionNumber=1.2.3"` to the go build command
var versionNumber string

// Config stores ALKS configuration and credentials
type Config struct {
	URL           string
	AccessKey     string
	SecretKey     string
	Token         string
	CredsFilename string
	Profile       string
	AssumeRole    assumeRoleDetails
}

type assumeRoleDetails struct {
	RoleARN     string
	SessionName string
	ExternalID  string
	Policy      string
}

func getCredentials(c *Config) *credentials.Credentials {
	// Follow the  same priority as the AWS Terraform Provider
	// https://www.terraform.io/docs/providers/aws/#authentication

	// needed for the EC2MetaData service
	sess := session.Must(session.NewSession())

	providers := []credentials.Provider{
		&credentials.StaticProvider{Value: credentials.Value{
			AccessKeyID:     c.AccessKey,
			SecretAccessKey: c.SecretKey,
			SessionToken:    c.Token,
		}},
		&credentials.EnvProvider{},
		&credentials.SharedCredentialsProvider{
			Filename: c.CredsFilename,
			Profile:  c.Profile,
		},
		&ec2rolecreds.EC2RoleProvider{
			Client: ec2metadata.New(sess),
		},
	}

	// Check for ECS container, for more details see:
	// https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task-iam-roles.html
	if uri := os.Getenv("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI"); len(uri) > 0 {
		client := cleanhttp.DefaultClient()
		client.Timeout = 100 * time.Millisecond
		cfg := &aws.Config{
			HTTPClient: client,
		}

		providers = append(providers, defaults.RemoteCredProvider(*cfg, defaults.Handlers()))
	}

	return credentials.NewChainCredentials(providers)
}

// Client returns a properly configured ALKS client or an appropriate error if initialization fails
func (c *Config) Client() (*alks.Client, error) {
	log.Println("[DEBUG] Validting STS credentials")

	// lookup credentials
	creds := getCredentials(c)
	cp, cpErr := creds.Get()

	// validate we have credentials
	if cpErr != nil {
		return nil, errors.New(`No valid credential sources found for ALKS Provider.
Please see https://github.com/Cox-Automotive/terraform-provider-alks#authentication for more information on
providing credentials for the ALKS Provider`)
	}

	// create a new session to test credentails
	sess, err := session.NewSession(&aws.Config{
		Region:      aws.String("us-east-1"),
		Credentials: creds,
	})

	// validate session
	if err != nil {
		return nil, fmt.Errorf("Error creating session from STS. (%v)", err)
	}

	var stsconn *sts.STS
	// we need to assume another role before creating an ALKS client
	if c.AssumeRole.RoleARN != "" {
		arCreds := stscreds.NewCredentials(sess, c.AssumeRole.RoleARN, func(p *stscreds.AssumeRoleProvider) {
			if c.AssumeRole.SessionName != "" {
				p.RoleSessionName = c.AssumeRole.SessionName
			}

			if c.AssumeRole.ExternalID != "" {
				p.ExternalID = &c.AssumeRole.ExternalID
			}

			if c.AssumeRole.Policy != "" {
				p.Policy = &c.AssumeRole.Policy
			}
		})

		cp, cpErr = arCreds.Get()
		if cpErr != nil {
			return nil, fmt.Errorf("The role %q cannot be assumed. Please verify the role ARN, role policies and your base AWS credentials", c.AssumeRole.RoleARN)
		}

		stsconn = sts.New(sess, &aws.Config{
			Region:      aws.String("us-east-1"),
			Credentials: arCreds,
		})
	} else {
		stsconn = sts.New(sess)
	}

	// make a basic api call to test creds are valid
	_, serr := stsconn.GetCallerIdentity(&sts.GetCallerIdentityInput{})

	// check for valid creds
	if serr != nil {
		return nil, serr
	}

	// got good creds, create alks sts client
	client, err := alks.NewSTSClient(c.URL, cp.AccessKeyID, cp.SecretAccessKey, cp.SessionToken)

	if err != nil {
		return nil, err
	}

	client.SetUserAgent(fmt.Sprintf("alks-terraform-provider-%s", getPluginVersion()))

	log.Println("[INFO] ALKS Client configured")

	return client, nil
}

func getPluginVersion() string {
	if versionNumber != "" {
		return versionNumber
	}

	return "unknown"
}
<commit_msg>Added a check to make sure credentials are correct, otherwise an error will be thrown with an appropriate error response.<commit_after>package main

import (
	"errors"
	"fmt"
	"log"
	"os"
	"os/exec"
	"strings"
	"time"

	"github.com/hashicorp/go-cleanhttp"

	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
	"github.com/aws/aws-sdk-go/aws/defaults"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"

	alks "github.com/Cox-Automotive/alks-go"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sts"
)

// Version number, to be injected at link time
// to set, add `-ldflags "-X main.versionNumber=1.2.3"` to the go build command
var versionNumber string

// Config stores ALKS configuration and credentials
type Config struct {
	URL           string
	AccessKey     string
	SecretKey     string
	Token         string
	CredsFilename string
	Profile       string
	AssumeRole    assumeRoleDetails
}

type assumeRoleDetails struct {
	RoleARN     string
	SessionName string
	ExternalID  string
	Policy      string
}

func getCredentials(c *Config) *credentials.Credentials {
	// Follow the  same priority as the AWS Terraform Provider
	// https://www.terraform.io/docs/providers/aws/#authentication

	// needed for the EC2MetaData service
	sess := session.Must(session.NewSession())

	providers := []credentials.Provider{
		&credentials.StaticProvider{Value: credentials.Value{
			AccessKeyID:     c.AccessKey,
			SecretAccessKey: c.SecretKey,
			SessionToken:    c.Token,
		}},
		&credentials.EnvProvider{},
		&credentials.SharedCredentialsProvider{
			Filename: c.CredsFilename,
			Profile:  c.Profile,
		},
		&ec2rolecreds.EC2RoleProvider{
			Client: ec2metadata.New(sess),
		},
	}

	// Check for ECS container, for more details see:
	// https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task-iam-roles.html
	if uri := os.Getenv("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI"); len(uri) > 0 {
		client := cleanhttp.DefaultClient()
		client.Timeout = 100 * time.Millisecond
		cfg := &aws.Config{
			HTTPClient: client,
		}

		providers = append(providers, defaults.RemoteCredProvider(*cfg, defaults.Handlers()))
	}

	return credentials.NewChainCredentials(providers)
}

// Client returns a properly configured ALKS client or an appropriate error if initialization fails
func (c *Config) Client() (*alks.Client, error) {
	log.Println("[DEBUG] Validting STS credentials")

	// lookup credentials
	creds := getCredentials(c)
	cp, cpErr := creds.Get()

	// validate we have credentials
	if cpErr != nil {
		return nil, errors.New(`No valid credential sources found for ALKS Provider.
Please see https://github.com/Cox-Automotive/terraform-provider-alks#authentication for more information on
providing credentials for the ALKS Provider`)
	}

	// create a new session to test credentails
	sess, err := session.NewSession(&aws.Config{
		Region:      aws.String("us-east-1"),
		Credentials: creds,
	})

	// validate session
	if err != nil {
		return nil, fmt.Errorf("Error creating session from STS. (%v)", err)
	}

	var stsconn *sts.STS
	// we need to assume another role before creating an ALKS client
	if c.AssumeRole.RoleARN != "" {
		arCreds := stscreds.NewCredentials(sess, c.AssumeRole.RoleARN, func(p *stscreds.AssumeRoleProvider) {
			if c.AssumeRole.SessionName != "" {
				p.RoleSessionName = c.AssumeRole.SessionName
			}

			if c.AssumeRole.ExternalID != "" {
				p.ExternalID = &c.AssumeRole.ExternalID
			}

			if c.AssumeRole.Policy != "" {
				p.Policy = &c.AssumeRole.Policy
			}
		})

		cp, cpErr = arCreds.Get()
		if cpErr != nil {
			return nil, fmt.Errorf("The role %q cannot be assumed. Please verify the role ARN, role policies and your base AWS credentials", c.AssumeRole.RoleARN)
		}

		stsconn = sts.New(sess, &aws.Config{
			Region:      aws.String("us-east-1"),
			Credentials: arCreds,
		})
	} else {
		stsconn = sts.New(sess)
	}

	// make a basic api call to test creds are valid
	_, serr := stsconn.GetCallerIdentity(&sts.GetCallerIdentityInput{})

	// check for valid creds
	if serr != nil {
		return nil, serr
	}

	// check if the user is using a assume-role IAM admin session
	if isValidIAM() != true {
		return nil, errors.New("Looks like you are not using ALKS IAM credentials. This will result in errors when creating roles. \n " +
			"Note: If using ALKS CLI to get credentials, be sure to use the '-i' flag. \n Please see https://coxautoinc.sharepoint.com/sites/service-internal-tools-team/SitePages/ALKS-Terraform-Provider---Troubleshooting.aspx for more information.")
	}

	// got good creds, create alks sts client
	client, err := alks.NewSTSClient(c.URL, cp.AccessKeyID, cp.SecretAccessKey, cp.SessionToken)

	if err != nil {
		return nil, err
	}

	client.SetUserAgent(fmt.Sprintf("alks-terraform-provider-%s", getPluginVersion()))

	log.Println("[INFO] ALKS Client configured")

	return client, nil
}

func getPluginVersion() string {
	if versionNumber != "" {
		return versionNumber
	}

	return "unknown"
}

func isValidIAM() bool {
	arg0 := "aws"
	arg1 := "sts"
	arg2 := "get-caller-identity"
	arg3 := "--query"
	arg4 := "[Arn]"
	arg5 := "--output"
	arg6 := "text"

	cmd := exec.Command(arg0, arg1, arg2, arg3, arg4, arg5, arg6)
	role, _ := cmd.Output()

	if strings.Contains(string(role), "assumed-role/Admin/") || strings.Contains(string(role), "assumed-role/IAMAdmin/") {
		return true
	}

	return false
}
<|endoftext|>