<commit_before>package drivers

import (
	"database/sql"
	"fmt"

	// Side-effect import sql driver
	_ "github.com/lib/pq"
	"github.com/nullbio/sqlboiler/bdb"
)

// PostgresDriver holds the database connection string and a handle
// to the database connection.
type PostgresDriver struct {
	connStr string
	dbConn  *sql.DB
}

// NewPostgresDriver takes the database connection details as parameters and
// returns a pointer to a PostgresDriver object. Note that it is required to
// call PostgresDriver.Open() and PostgresDriver.Close() to open and close
// the database connection once an object has been obtained.
func NewPostgresDriver(user, pass, dbname, host string, port int, sslmode string) *PostgresDriver {
	driver := PostgresDriver{
		connStr: fmt.Sprintf("user=%s password=%s dbname=%s host=%s port=%d sslmode=%s",
			user, pass, dbname, host, port, sslmode),
	}

	return &driver
}

// Open opens the database connection using the connection string
func (p *PostgresDriver) Open() error {
	var err error
	p.dbConn, err = sql.Open("postgres", p.connStr)
	if err != nil {
		return err
	}

	return nil
}

// Close closes the database connection
func (p *PostgresDriver) Close() {
	p.dbConn.Close()
}

// TableNames connects to the postgres database and
// retrieves all table names from the information_schema where the
// table schema is public. It excludes common migration tool tables
// such as gorp_migrations
func (p *PostgresDriver) TableNames() ([]string, error) {
	var names []string

	rows, err := p.dbConn.Query(`
		select table_name from information_schema.tables
		where table_schema = 'public' and table_name not like '%migrations%'
	`)

	if err != nil {
		return nil, err
	}

	defer rows.Close()
	for rows.Next() {
		var name string
		if err := rows.Scan(&name); err != nil {
			return nil, err
		}
		names = append(names, name)
	}

	return names, nil
}

// Columns takes a table name and attempts to retrieve the table information
// from the database information_schema.columns. It retrieves the column names
// and column types and returns those as a []Column after TranslateColumnType()
// converts the SQL types to Go types, for example: "varchar" to "string"
func (p *PostgresDriver) Columns(tableName string) ([]bdb.Column, error) {
	var columns []bdb.Column

	rows, err := p.dbConn.Query(`
		select column_name, data_type, column_default, is_nullable,
			(
				select cast(count(*) as bit) as is_unique
				from information_schema.constraint_column_usage as ccu
		    inner join information_schema.table_constraints tc on ccu.constraint_name = tc.constraint_name
		    where ccu.column_name = c.column_name and tc.constraint_type = 'UNIQUE'
			) as is_unique
		from information_schema.columns as c
		where table_name=$1 and table_schema = 'public';
	`, tableName)

	if err != nil {
		return nil, err
	}

	defer rows.Close()
	for rows.Next() {
		var colName, colType, colDefault, nullable string
		var unique bool
		var defaultPtr *string
		if err := rows.Scan(&colName, &colType, &defaultPtr, &nullable, &unique); err != nil {
			return nil, fmt.Errorf("unable to scan for table %s: %s", tableName, err)
		}

		if defaultPtr == nil {
			colDefault = ""
		} else {
			colDefault = *defaultPtr
		}

		column := bdb.Column{
			Name:     colName,
			DBType:   colType,
			Default:  colDefault,
			Nullable: nullable == "YES",
			Unique:   unique,
		}
		columns = append(columns, column)
	}

	return columns, nil
}

// PrimaryKeyInfo looks up the primary key for a table.
func (p *PostgresDriver) PrimaryKeyInfo(tableName string) (*bdb.PrimaryKey, error) {
	pkey := &bdb.PrimaryKey{}
	var err error

	query := `
	select tc.constraint_name
	from information_schema.table_constraints as tc
	where tc.table_name = $1 and tc.constraint_type = 'PRIMARY KEY' and tc.table_schema = 'public';`

	row := p.dbConn.QueryRow(query, tableName)
	if err = row.Scan(&pkey.Name); err != nil {
		if err == sql.ErrNoRows {
			return nil, nil
		}
		return nil, err
	}

	queryColumns := `
	select kcu.column_name
	from   information_schema.key_column_usage as kcu
	where  constraint_name = $1 and table_schema = 'public';`

	var rows *sql.Rows
	if rows, err = p.dbConn.Query(queryColumns, pkey.Name); err != nil {
		return nil, err
	}

	var columns []string
	for rows.Next() {
		var column string

		err = rows.Scan(&column)
		if err != nil {
			return nil, err
		}

		columns = append(columns, column)
	}

	if err = rows.Err(); err != nil {
		return nil, err
	}

	pkey.Columns = columns

	return pkey, nil
}

// ForeignKeyInfo retrieves the foreign keys for a given table name.
func (p *PostgresDriver) ForeignKeyInfo(tableName string) ([]bdb.ForeignKey, error) {
	var fkeys []bdb.ForeignKey

	query := `
	select
		tc.constraint_name,
		kcu.table_name as source_table,
		kcu.column_name as source_column,
		ccu.table_name as dest_table,
		ccu.column_name as dest_column
	from information_schema.table_constraints as tc
		inner join information_schema.key_column_usage as kcu ON tc.constraint_name = kcu.constraint_name
		inner join information_schema.constraint_column_usage as ccu ON tc.constraint_name = ccu.constraint_name
	where tc.table_name = $1 and tc.constraint_type = 'FOREIGN KEY' and tc.table_schema = 'public';`

	var rows *sql.Rows
	var err error
	if rows, err = p.dbConn.Query(query, tableName); err != nil {
		return nil, err
	}

	for rows.Next() {
		var fkey bdb.ForeignKey
		var sourceTable string

		err = rows.Scan(&fkey.Name, &sourceTable, &fkey.Column, &fkey.ForeignTable, &fkey.ForeignColumn)
		if err != nil {
			return nil, err
		}

		fkeys = append(fkeys, fkey)
	}

	if err = rows.Err(); err != nil {
		return nil, err
	}

	return fkeys, nil
}

// TranslateColumnType converts postgres database types to Go types, for example
// "varchar" to "string" and "bigint" to "int64". It returns this parsed data
// as a Column object.
func (p *PostgresDriver) TranslateColumnType(c bdb.Column) bdb.Column {
	if c.Nullable {
		switch c.DBType {
		case "bigint", "bigserial":
			c.Type = "null.Int64"
		case "integer", "serial":
			c.Type = "null.Int32"
		case "smallint", "smallserial":
			c.Type = "null.Int16"
		case "decimal", "numeric", "double precision", "money":
			c.Type = "null.Float64"
		case "real":
			c.Type = "null.Float32"
		case "bit", "interval", "bit varying", "character", "character varying", "cidr", "inet", "json", "macaddr", "text", "uuid", "xml":
			c.Type = "null.String"
		case "bytea":
			c.Type = "[]byte"
		case "boolean":
			c.Type = "null.Bool"
		case "date", "time", "timestamp without time zone", "timestamp with time zone":
			c.Type = "null.Time"
		default:
			c.Type = "null.String"
		}
	} else {
		switch c.DBType {
		case "bigint", "bigserial":
			c.Type = "int64"
		case "integer", "serial":
			c.Type = "int32"
		case "smallint", "smallserial":
			c.Type = "int16"
		case "decimal", "numeric", "double precision", "money":
			c.Type = "float64"
		case "real":
			c.Type = "float32"
		case "bit", "interval", "bit varying", "character", "character varying", "cidr", "inet", "json", "macaddr", "text", "uuid", "xml":
			c.Type = "string"
		case "bytea":
			c.Type = "[]byte"
		case "boolean":
			c.Type = "bool"
		case "date", "time", "timestamp without time zone", "timestamp with time zone":
			c.Type = "time.Time"
		default:
			c.Type = "string"
		}
	}

	return c
}
<commit_msg>Fix a bug in the column constraint detection<commit_after>package drivers

import (
	"database/sql"
	"fmt"

	// Side-effect import sql driver
	_ "github.com/lib/pq"
	"github.com/nullbio/sqlboiler/bdb"
)

// PostgresDriver holds the database connection string and a handle
// to the database connection.
type PostgresDriver struct {
	connStr string
	dbConn  *sql.DB
}

// NewPostgresDriver takes the database connection details as parameters and
// returns a pointer to a PostgresDriver object. Note that it is required to
// call PostgresDriver.Open() and PostgresDriver.Close() to open and close
// the database connection once an object has been obtained.
func NewPostgresDriver(user, pass, dbname, host string, port int, sslmode string) *PostgresDriver {
	driver := PostgresDriver{
		connStr: fmt.Sprintf("user=%s password=%s dbname=%s host=%s port=%d sslmode=%s",
			user, pass, dbname, host, port, sslmode),
	}

	return &driver
}

// Open opens the database connection using the connection string
func (p *PostgresDriver) Open() error {
	var err error
	p.dbConn, err = sql.Open("postgres", p.connStr)
	if err != nil {
		return err
	}

	return nil
}

// Close closes the database connection
func (p *PostgresDriver) Close() {
	p.dbConn.Close()
}

// TableNames connects to the postgres database and
// retrieves all table names from the information_schema where the
// table schema is public. It excludes common migration tool tables
// such as gorp_migrations
func (p *PostgresDriver) TableNames() ([]string, error) {
	var names []string

	rows, err := p.dbConn.Query(`
		select table_name from information_schema.tables
		where table_schema = 'public' and table_name not like '%migrations%'
	`)

	if err != nil {
		return nil, err
	}

	defer rows.Close()
	for rows.Next() {
		var name string
		if err := rows.Scan(&name); err != nil {
			return nil, err
		}
		names = append(names, name)
	}

	return names, nil
}

// Columns takes a table name and attempts to retrieve the table information
// from the database information_schema.columns. It retrieves the column names
// and column types and returns those as a []Column after TranslateColumnType()
// converts the SQL types to Go types, for example: "varchar" to "string"
func (p *PostgresDriver) Columns(tableName string) ([]bdb.Column, error) {
	var columns []bdb.Column

	rows, err := p.dbConn.Query(`
		select column_name, data_type, column_default, is_nullable,
			(
				select cast(count(*) as bit) as is_unique
				from information_schema.constraint_column_usage as ccu
		    inner join information_schema.table_constraints tc on ccu.constraint_name = tc.constraint_name
		    where ccu.table_name = c.table_name and ccu.column_name = c.column_name and tc.constraint_type = 'UNIQUE'
			) as is_unique
		from information_schema.columns as c
		where table_name=$1 and table_schema = 'public';
	`, tableName)

	if err != nil {
		return nil, err
	}

	defer rows.Close()
	for rows.Next() {
		var colName, colType, colDefault, nullable string
		var unique bool
		var defaultPtr *string
		if err := rows.Scan(&colName, &colType, &defaultPtr, &nullable, &unique); err != nil {
			return nil, fmt.Errorf("unable to scan for table %s: %s", tableName, err)
		}

		if defaultPtr == nil {
			colDefault = ""
		} else {
			colDefault = *defaultPtr
		}

		column := bdb.Column{
			Name:     colName,
			DBType:   colType,
			Default:  colDefault,
			Nullable: nullable == "YES",
			Unique:   unique,
		}
		columns = append(columns, column)
	}

	return columns, nil
}

// PrimaryKeyInfo looks up the primary key for a table.
func (p *PostgresDriver) PrimaryKeyInfo(tableName string) (*bdb.PrimaryKey, error) {
	pkey := &bdb.PrimaryKey{}
	var err error

	query := `
	select tc.constraint_name
	from information_schema.table_constraints as tc
	where tc.table_name = $1 and tc.constraint_type = 'PRIMARY KEY' and tc.table_schema = 'public';`

	row := p.dbConn.QueryRow(query, tableName)
	if err = row.Scan(&pkey.Name); err != nil {
		if err == sql.ErrNoRows {
			return nil, nil
		}
		return nil, err
	}

	queryColumns := `
	select kcu.column_name
	from   information_schema.key_column_usage as kcu
	where  constraint_name = $1 and table_schema = 'public';`

	var rows *sql.Rows
	if rows, err = p.dbConn.Query(queryColumns, pkey.Name); err != nil {
		return nil, err
	}

	var columns []string
	for rows.Next() {
		var column string

		err = rows.Scan(&column)
		if err != nil {
			return nil, err
		}

		columns = append(columns, column)
	}

	if err = rows.Err(); err != nil {
		return nil, err
	}

	pkey.Columns = columns

	return pkey, nil
}

// ForeignKeyInfo retrieves the foreign keys for a given table name.
func (p *PostgresDriver) ForeignKeyInfo(tableName string) ([]bdb.ForeignKey, error) {
	var fkeys []bdb.ForeignKey

	query := `
	select
		tc.constraint_name,
		kcu.table_name as source_table,
		kcu.column_name as source_column,
		ccu.table_name as dest_table,
		ccu.column_name as dest_column
	from information_schema.table_constraints as tc
		inner join information_schema.key_column_usage as kcu ON tc.constraint_name = kcu.constraint_name
		inner join information_schema.constraint_column_usage as ccu ON tc.constraint_name = ccu.constraint_name
	where tc.table_name = $1 and tc.constraint_type = 'FOREIGN KEY' and tc.table_schema = 'public';`

	var rows *sql.Rows
	var err error
	if rows, err = p.dbConn.Query(query, tableName); err != nil {
		return nil, err
	}

	for rows.Next() {
		var fkey bdb.ForeignKey
		var sourceTable string

		err = rows.Scan(&fkey.Name, &sourceTable, &fkey.Column, &fkey.ForeignTable, &fkey.ForeignColumn)
		if err != nil {
			return nil, err
		}

		fkeys = append(fkeys, fkey)
	}

	if err = rows.Err(); err != nil {
		return nil, err
	}

	return fkeys, nil
}

// TranslateColumnType converts postgres database types to Go types, for example
// "varchar" to "string" and "bigint" to "int64". It returns this parsed data
// as a Column object.
func (p *PostgresDriver) TranslateColumnType(c bdb.Column) bdb.Column {
	if c.Nullable {
		switch c.DBType {
		case "bigint", "bigserial":
			c.Type = "null.Int64"
		case "integer", "serial":
			c.Type = "null.Int32"
		case "smallint", "smallserial":
			c.Type = "null.Int16"
		case "decimal", "numeric", "double precision", "money":
			c.Type = "null.Float64"
		case "real":
			c.Type = "null.Float32"
		case "bit", "interval", "bit varying", "character", "character varying", "cidr", "inet", "json", "macaddr", "text", "uuid", "xml":
			c.Type = "null.String"
		case "bytea":
			c.Type = "[]byte"
		case "boolean":
			c.Type = "null.Bool"
		case "date", "time", "timestamp without time zone", "timestamp with time zone":
			c.Type = "null.Time"
		default:
			c.Type = "null.String"
		}
	} else {
		switch c.DBType {
		case "bigint", "bigserial":
			c.Type = "int64"
		case "integer", "serial":
			c.Type = "int32"
		case "smallint", "smallserial":
			c.Type = "int16"
		case "decimal", "numeric", "double precision", "money":
			c.Type = "float64"
		case "real":
			c.Type = "float32"
		case "bit", "interval", "bit varying", "character", "character varying", "cidr", "inet", "json", "macaddr", "text", "uuid", "xml":
			c.Type = "string"
		case "bytea":
			c.Type = "[]byte"
		case "boolean":
			c.Type = "bool"
		case "date", "time", "timestamp without time zone", "timestamp with time zone":
			c.Type = "time.Time"
		default:
			c.Type = "string"
		}
	}

	return c
}
<|endoftext|>


<commit_before>package drivers

import (
	"database/sql"
	"fmt"
	"strings"

	// Side-effect import sql driver
	_ "github.com/lib/pq"
	"github.com/pkg/errors"
	"github.com/vattle/sqlboiler/bdb"
	"github.com/vattle/sqlboiler/strmangle"
)

// PostgresDriver holds the database connection string and a handle
// to the database connection.
type PostgresDriver struct {
	connStr string
	dbConn  *sql.DB
}

// NewPostgresDriver takes the database connection details as parameters and
// returns a pointer to a PostgresDriver object. Note that it is required to
// call PostgresDriver.Open() and PostgresDriver.Close() to open and close
// the database connection once an object has been obtained.
func NewPostgresDriver(user, pass, dbname, host string, port int, sslmode string) *PostgresDriver {
	driver := PostgresDriver{
		connStr: PostgresBuildQueryString(user, pass, dbname, host, port, sslmode),
	}

	return &driver
}

// PostgresBuildQueryString builds a query string.
func PostgresBuildQueryString(user, pass, dbname, host string, port int, sslmode string) string {
	parts := []string{}
	if len(user) != 0 {
		parts = append(parts, fmt.Sprintf("user=%s", user))
	}
	if len(pass) != 0 {
		parts = append(parts, fmt.Sprintf("password=%s", pass))
	}
	if len(dbname) != 0 {
		parts = append(parts, fmt.Sprintf("dbname=%s", dbname))
	}
	if len(host) != 0 {
		parts = append(parts, fmt.Sprintf("host=%s", host))
	}
	if port != 0 {
		parts = append(parts, fmt.Sprintf("port=%d", port))
	}
	if len(sslmode) != 0 {
		parts = append(parts, fmt.Sprintf("sslmode=%s", sslmode))
	}

	return strings.Join(parts, " ")
}

// Open opens the database connection using the connection string
func (p *PostgresDriver) Open() error {
	var err error
	p.dbConn, err = sql.Open("postgres", p.connStr)
	if err != nil {
		return err
	}

	return nil
}

// Close closes the database connection
func (p *PostgresDriver) Close() {
	p.dbConn.Close()
}

// UseLastInsertID returns false for postgres
func (p *PostgresDriver) UseLastInsertID() bool {
	return false
}

// TableNames connects to the postgres database and
// retrieves all table names from the information_schema where the
// table schema is schema. It uses a whitelist and blacklist.
func (p *PostgresDriver) TableNames(schema string, whitelist, blacklist []string) ([]string, error) {
	var names []string

	query := fmt.Sprintf(`select table_name from information_schema.tables where table_schema = $1`)
	args := []interface{}{schema}
	if len(whitelist) > 0 {
		query += fmt.Sprintf(" and table_name in (%s);", strmangle.Placeholders(true, len(whitelist), 2, 1))
		for _, w := range whitelist {
			args = append(args, w)
		}
	} else if len(blacklist) > 0 {
		query += fmt.Sprintf(" and table_name not in (%s);", strmangle.Placeholders(true, len(blacklist), 2, 1))
		for _, b := range blacklist {
			args = append(args, b)
		}
	}

	rows, err := p.dbConn.Query(query, args...)

	if err != nil {
		return nil, err
	}

	defer rows.Close()
	for rows.Next() {
		var name string
		if err := rows.Scan(&name); err != nil {
			return nil, err
		}
		names = append(names, name)
	}

	return names, nil
}

// Columns takes a table name and attempts to retrieve the table information
// from the database information_schema.columns. It retrieves the column names
// and column types and returns those as a []Column after TranslateColumnType()
// converts the SQL types to Go types, for example: "varchar" to "string"
func (p *PostgresDriver) Columns(schema, tableName string) ([]bdb.Column, error) {
	var columns []bdb.Column

	rows, err := p.dbConn.Query(`
		select column_name, c.data_type, e.data_type, column_default, c.udt_name, is_nullable,
			(select exists(
		    select 1
				from information_schema.constraint_column_usage as ccu
		    inner join information_schema.table_constraints tc on ccu.constraint_name = tc.constraint_name
		    where ccu.table_name = c.table_name and ccu.column_name = c.column_name and tc.constraint_type = 'UNIQUE'
			)) OR (select exists(
		    select 1
		    from
		      pg_indexes pgix
		      inner join pg_class pgc on pgix.indexname = pgc.relname and pgc.relkind = 'i'
		      inner join pg_index pgi on pgi.indexrelid = pgc.oid
		      inner join pg_attribute pga on pga.attrelid = pgi.indrelid and pga.attnum = ANY(pgi.indkey)
		    where
		      pgix.schemaname = $1 and pgix.tablename = c.table_name and pga.attname = c.column_name and pgi.indisunique = true
		)) as is_unique
		from information_schema.columns as c LEFT JOIN information_schema.element_types e
			ON ((c.table_catalog, c.table_schema, c.table_name, 'TABLE', c.dtd_identifier)
				= (e.object_catalog, e.object_schema, e.object_name, e.object_type, e.collection_type_identifier))
		where c.table_name=$2 and c.table_schema = $1;
	`, schema, tableName)

	if err != nil {
		return nil, err
	}
	defer rows.Close()

	for rows.Next() {
		var colName, udtName, colType, colDefault, nullable string
		var elementType *string
		var unique bool
		var defaultPtr *string
		if err := rows.Scan(&colName, &colType, &elementType, &defaultPtr, &udtName, &nullable, &unique); err != nil {
			return nil, errors.Wrapf(err, "unable to scan for table %s", tableName)
		}

		if defaultPtr == nil {
			colDefault = ""
		} else {
			colDefault = *defaultPtr
		}

		column := bdb.Column{
			Name:     colName,
			DBType:   colType,
			ArrType:  elementType,
			UDTName:  udtName,
			Default:  colDefault,
			Nullable: nullable == "YES",
			Unique:   unique,
		}
		columns = append(columns, column)
	}

	return columns, nil
}

// PrimaryKeyInfo looks up the primary key for a table.
func (p *PostgresDriver) PrimaryKeyInfo(schema, tableName string) (*bdb.PrimaryKey, error) {
	pkey := &bdb.PrimaryKey{}
	var err error

	query := `
	select tc.constraint_name
	from information_schema.table_constraints as tc
	where tc.table_name = $1 and tc.constraint_type = 'PRIMARY KEY' and tc.table_schema = $2;`

	row := p.dbConn.QueryRow(query, tableName, schema)
	if err = row.Scan(&pkey.Name); err != nil {
		if err == sql.ErrNoRows {
			return nil, nil
		}
		return nil, err
	}

	queryColumns := `
	select kcu.column_name
	from   information_schema.key_column_usage as kcu
	where  constraint_name = $1 and table_schema = $2;`

	var rows *sql.Rows
	if rows, err = p.dbConn.Query(queryColumns, pkey.Name, schema); err != nil {
		return nil, err
	}
	defer rows.Close()

	var columns []string
	for rows.Next() {
		var column string

		err = rows.Scan(&column)
		if err != nil {
			return nil, err
		}

		columns = append(columns, column)
	}

	if err = rows.Err(); err != nil {
		return nil, err
	}

	pkey.Columns = columns

	return pkey, nil
}

// ForeignKeyInfo retrieves the foreign keys for a given table name.
func (p *PostgresDriver) ForeignKeyInfo(schema, tableName string) ([]bdb.ForeignKey, error) {
	var fkeys []bdb.ForeignKey

	query := `
	select
		tc.constraint_name,
		kcu.table_name as source_table,
		kcu.column_name as source_column,
		ccu.table_name as dest_table,
		ccu.column_name as dest_column
	from information_schema.table_constraints as tc
		inner join information_schema.key_column_usage as kcu ON tc.constraint_name = kcu.constraint_name and tc.constraint_schema = kcu.constraint_schema
		inner join information_schema.constraint_column_usage as ccu ON tc.constraint_name = ccu.constraint_name and tc.constraint_schema = ccu.constraint_schema
	where tc.table_name = $1 and tc.constraint_type = 'FOREIGN KEY' and tc.table_schema = $2;`

	var rows *sql.Rows
	var err error
	if rows, err = p.dbConn.Query(query, tableName, schema); err != nil {
		return nil, err
	}

	for rows.Next() {
		var fkey bdb.ForeignKey
		var sourceTable string

		fkey.Table = tableName
		err = rows.Scan(&fkey.Name, &sourceTable, &fkey.Column, &fkey.ForeignTable, &fkey.ForeignColumn)
		if err != nil {
			return nil, err
		}

		fkeys = append(fkeys, fkey)
	}

	if err = rows.Err(); err != nil {
		return nil, err
	}

	return fkeys, nil
}

// TranslateColumnType converts postgres database types to Go types, for example
// "varchar" to "string" and "bigint" to "int64". It returns this parsed data
// as a Column object.
func (p *PostgresDriver) TranslateColumnType(c bdb.Column) bdb.Column {
	if c.Nullable {
		switch c.DBType {
		case "bigint", "bigserial":
			c.Type = "null.Int64"
		case "integer", "serial":
			c.Type = "null.Int"
		case "smallint", "smallserial":
			c.Type = "null.Int16"
		case "decimal", "numeric", "double precision":
			c.Type = "null.Float64"
		case "real":
			c.Type = "null.Float32"
		case "bit", "interval", "bit varying", "character", "money", "character varying", "cidr", "inet", "macaddr", "text", "uuid", "xml":
			c.Type = "null.String"
		case "bytea":
			c.Type = "null.Bytes"
		case "json", "jsonb":
			c.Type = "null.JSON"
		case "boolean":
			c.Type = "null.Bool"
		case "date", "time", "timestamp without time zone", "timestamp with time zone":
			c.Type = "null.Time"
		case "ARRAY":
			if c.ArrType == nil {
				panic("unable to get postgres ARRAY underlying type")
			}
			c.Type = getArrayType(c)
			// Make DBType something like ARRAYinteger for parsing with randomize.Struct
			c.DBType = c.DBType + *c.ArrType
		case "USER-DEFINED":
			if c.UDTName == "hstore" {
				c.Type = "types.HStore"
				c.DBType = "hstore"
			} else {
				c.Type = "string"
				fmt.Printf("Warning: Incompatible data type detected: %s", c.UDTName)
			}
		default:
			c.Type = "null.String"
		}
	} else {
		switch c.DBType {
		case "bigint", "bigserial":
			c.Type = "int64"
		case "integer", "serial":
			c.Type = "int"
		case "smallint", "smallserial":
			c.Type = "int16"
		case "decimal", "numeric", "double precision":
			c.Type = "float64"
		case "real":
			c.Type = "float32"
		case "bit", "interval", "uuint", "bit varying", "character", "money", "character varying", "cidr", "inet", "macaddr", "text", "uuid", "xml":
			c.Type = "string"
		case "json", "jsonb":
			c.Type = "types.JSON"
		case "bytea":
			c.Type = "[]byte"
		case "boolean":
			c.Type = "bool"
		case "date", "time", "timestamp without time zone", "timestamp with time zone":
			c.Type = "time.Time"
		case "ARRAY":
			c.Type = getArrayType(c)
			// Make DBType something like ARRAYinteger for parsing with randomize.Struct
			c.DBType = c.DBType + *c.ArrType
		case "USER-DEFINED":
			if c.UDTName == "hstore" {
				c.Type = "types.HStore"
				c.DBType = "hstore"
			} else {
				c.Type = "string"
				fmt.Printf("Warning: Incompatible data type detected: %s", c.UDTName)
			}
		default:
			c.Type = "string"
		}
	}

	return c
}

// getArrayType returns the correct boil.Array type for each database type
func getArrayType(c bdb.Column) string {
	switch *c.ArrType {
	case "bigint", "bigserial", "integer", "serial", "smallint", "smallserial":
		return "types.Int64Array"
	case "bytea":
		return "types.BytesArray"
	case "bit", "interval", "uuint", "bit varying", "character", "money", "character varying", "cidr", "inet", "macaddr", "text", "uuid", "xml":
		return "types.StringArray"
	case "boolean":
		return "types.BoolArray"
	case "decimal", "numeric", "double precision", "real":
		return "types.Float64Array"
	default:
		return "types.StringArray"
	}
}

// RightQuote is the quoting character for the right side of the identifier
func (p *PostgresDriver) RightQuote() byte {
	return '"'
}

// LeftQuote is the quoting character for the left side of the identifier
func (p *PostgresDriver) LeftQuote() byte {
	return '"'
}

// IndexPlaceholders returns true to indicate PSQL supports indexed placeholders
func (p *PostgresDriver) IndexPlaceholders() bool {
	return true
}
<commit_msg>Make warnings for incompat data type readable<commit_after>package drivers

import (
	"database/sql"
	"fmt"
	"strings"

	// Side-effect import sql driver
	_ "github.com/lib/pq"
	"github.com/pkg/errors"
	"github.com/vattle/sqlboiler/bdb"
	"github.com/vattle/sqlboiler/strmangle"
)

// PostgresDriver holds the database connection string and a handle
// to the database connection.
type PostgresDriver struct {
	connStr string
	dbConn  *sql.DB
}

// NewPostgresDriver takes the database connection details as parameters and
// returns a pointer to a PostgresDriver object. Note that it is required to
// call PostgresDriver.Open() and PostgresDriver.Close() to open and close
// the database connection once an object has been obtained.
func NewPostgresDriver(user, pass, dbname, host string, port int, sslmode string) *PostgresDriver {
	driver := PostgresDriver{
		connStr: PostgresBuildQueryString(user, pass, dbname, host, port, sslmode),
	}

	return &driver
}

// PostgresBuildQueryString builds a query string.
func PostgresBuildQueryString(user, pass, dbname, host string, port int, sslmode string) string {
	parts := []string{}
	if len(user) != 0 {
		parts = append(parts, fmt.Sprintf("user=%s", user))
	}
	if len(pass) != 0 {
		parts = append(parts, fmt.Sprintf("password=%s", pass))
	}
	if len(dbname) != 0 {
		parts = append(parts, fmt.Sprintf("dbname=%s", dbname))
	}
	if len(host) != 0 {
		parts = append(parts, fmt.Sprintf("host=%s", host))
	}
	if port != 0 {
		parts = append(parts, fmt.Sprintf("port=%d", port))
	}
	if len(sslmode) != 0 {
		parts = append(parts, fmt.Sprintf("sslmode=%s", sslmode))
	}

	return strings.Join(parts, " ")
}

// Open opens the database connection using the connection string
func (p *PostgresDriver) Open() error {
	var err error
	p.dbConn, err = sql.Open("postgres", p.connStr)
	if err != nil {
		return err
	}

	return nil
}

// Close closes the database connection
func (p *PostgresDriver) Close() {
	p.dbConn.Close()
}

// UseLastInsertID returns false for postgres
func (p *PostgresDriver) UseLastInsertID() bool {
	return false
}

// TableNames connects to the postgres database and
// retrieves all table names from the information_schema where the
// table schema is schema. It uses a whitelist and blacklist.
func (p *PostgresDriver) TableNames(schema string, whitelist, blacklist []string) ([]string, error) {
	var names []string

	query := fmt.Sprintf(`select table_name from information_schema.tables where table_schema = $1`)
	args := []interface{}{schema}
	if len(whitelist) > 0 {
		query += fmt.Sprintf(" and table_name in (%s);", strmangle.Placeholders(true, len(whitelist), 2, 1))
		for _, w := range whitelist {
			args = append(args, w)
		}
	} else if len(blacklist) > 0 {
		query += fmt.Sprintf(" and table_name not in (%s);", strmangle.Placeholders(true, len(blacklist), 2, 1))
		for _, b := range blacklist {
			args = append(args, b)
		}
	}

	rows, err := p.dbConn.Query(query, args...)

	if err != nil {
		return nil, err
	}

	defer rows.Close()
	for rows.Next() {
		var name string
		if err := rows.Scan(&name); err != nil {
			return nil, err
		}
		names = append(names, name)
	}

	return names, nil
}

// Columns takes a table name and attempts to retrieve the table information
// from the database information_schema.columns. It retrieves the column names
// and column types and returns those as a []Column after TranslateColumnType()
// converts the SQL types to Go types, for example: "varchar" to "string"
func (p *PostgresDriver) Columns(schema, tableName string) ([]bdb.Column, error) {
	var columns []bdb.Column

	rows, err := p.dbConn.Query(`
		select column_name, c.data_type, e.data_type, column_default, c.udt_name, is_nullable,
			(select exists(
		    select 1
				from information_schema.constraint_column_usage as ccu
		    inner join information_schema.table_constraints tc on ccu.constraint_name = tc.constraint_name
		    where ccu.table_name = c.table_name and ccu.column_name = c.column_name and tc.constraint_type = 'UNIQUE'
			)) OR (select exists(
		    select 1
		    from
		      pg_indexes pgix
		      inner join pg_class pgc on pgix.indexname = pgc.relname and pgc.relkind = 'i'
		      inner join pg_index pgi on pgi.indexrelid = pgc.oid
		      inner join pg_attribute pga on pga.attrelid = pgi.indrelid and pga.attnum = ANY(pgi.indkey)
		    where
		      pgix.schemaname = $1 and pgix.tablename = c.table_name and pga.attname = c.column_name and pgi.indisunique = true
		)) as is_unique
		from information_schema.columns as c LEFT JOIN information_schema.element_types e
			ON ((c.table_catalog, c.table_schema, c.table_name, 'TABLE', c.dtd_identifier)
				= (e.object_catalog, e.object_schema, e.object_name, e.object_type, e.collection_type_identifier))
		where c.table_name=$2 and c.table_schema = $1;
	`, schema, tableName)

	if err != nil {
		return nil, err
	}
	defer rows.Close()

	for rows.Next() {
		var colName, udtName, colType, colDefault, nullable string
		var elementType *string
		var unique bool
		var defaultPtr *string
		if err := rows.Scan(&colName, &colType, &elementType, &defaultPtr, &udtName, &nullable, &unique); err != nil {
			return nil, errors.Wrapf(err, "unable to scan for table %s", tableName)
		}

		if defaultPtr == nil {
			colDefault = ""
		} else {
			colDefault = *defaultPtr
		}

		column := bdb.Column{
			Name:     colName,
			DBType:   colType,
			ArrType:  elementType,
			UDTName:  udtName,
			Default:  colDefault,
			Nullable: nullable == "YES",
			Unique:   unique,
		}
		columns = append(columns, column)
	}

	return columns, nil
}

// PrimaryKeyInfo looks up the primary key for a table.
func (p *PostgresDriver) PrimaryKeyInfo(schema, tableName string) (*bdb.PrimaryKey, error) {
	pkey := &bdb.PrimaryKey{}
	var err error

	query := `
	select tc.constraint_name
	from information_schema.table_constraints as tc
	where tc.table_name = $1 and tc.constraint_type = 'PRIMARY KEY' and tc.table_schema = $2;`

	row := p.dbConn.QueryRow(query, tableName, schema)
	if err = row.Scan(&pkey.Name); err != nil {
		if err == sql.ErrNoRows {
			return nil, nil
		}
		return nil, err
	}

	queryColumns := `
	select kcu.column_name
	from   information_schema.key_column_usage as kcu
	where  constraint_name = $1 and table_schema = $2;`

	var rows *sql.Rows
	if rows, err = p.dbConn.Query(queryColumns, pkey.Name, schema); err != nil {
		return nil, err
	}
	defer rows.Close()

	var columns []string
	for rows.Next() {
		var column string

		err = rows.Scan(&column)
		if err != nil {
			return nil, err
		}

		columns = append(columns, column)
	}

	if err = rows.Err(); err != nil {
		return nil, err
	}

	pkey.Columns = columns

	return pkey, nil
}

// ForeignKeyInfo retrieves the foreign keys for a given table name.
func (p *PostgresDriver) ForeignKeyInfo(schema, tableName string) ([]bdb.ForeignKey, error) {
	var fkeys []bdb.ForeignKey

	query := `
	select
		tc.constraint_name,
		kcu.table_name as source_table,
		kcu.column_name as source_column,
		ccu.table_name as dest_table,
		ccu.column_name as dest_column
	from information_schema.table_constraints as tc
		inner join information_schema.key_column_usage as kcu ON tc.constraint_name = kcu.constraint_name and tc.constraint_schema = kcu.constraint_schema
		inner join information_schema.constraint_column_usage as ccu ON tc.constraint_name = ccu.constraint_name and tc.constraint_schema = ccu.constraint_schema
	where tc.table_name = $1 and tc.constraint_type = 'FOREIGN KEY' and tc.table_schema = $2;`

	var rows *sql.Rows
	var err error
	if rows, err = p.dbConn.Query(query, tableName, schema); err != nil {
		return nil, err
	}

	for rows.Next() {
		var fkey bdb.ForeignKey
		var sourceTable string

		fkey.Table = tableName
		err = rows.Scan(&fkey.Name, &sourceTable, &fkey.Column, &fkey.ForeignTable, &fkey.ForeignColumn)
		if err != nil {
			return nil, err
		}

		fkeys = append(fkeys, fkey)
	}

	if err = rows.Err(); err != nil {
		return nil, err
	}

	return fkeys, nil
}

// TranslateColumnType converts postgres database types to Go types, for example
// "varchar" to "string" and "bigint" to "int64". It returns this parsed data
// as a Column object.
func (p *PostgresDriver) TranslateColumnType(c bdb.Column) bdb.Column {
	if c.Nullable {
		switch c.DBType {
		case "bigint", "bigserial":
			c.Type = "null.Int64"
		case "integer", "serial":
			c.Type = "null.Int"
		case "smallint", "smallserial":
			c.Type = "null.Int16"
		case "decimal", "numeric", "double precision":
			c.Type = "null.Float64"
		case "real":
			c.Type = "null.Float32"
		case "bit", "interval", "bit varying", "character", "money", "character varying", "cidr", "inet", "macaddr", "text", "uuid", "xml":
			c.Type = "null.String"
		case "bytea":
			c.Type = "null.Bytes"
		case "json", "jsonb":
			c.Type = "null.JSON"
		case "boolean":
			c.Type = "null.Bool"
		case "date", "time", "timestamp without time zone", "timestamp with time zone":
			c.Type = "null.Time"
		case "ARRAY":
			if c.ArrType == nil {
				panic("unable to get postgres ARRAY underlying type")
			}
			c.Type = getArrayType(c)
			// Make DBType something like ARRAYinteger for parsing with randomize.Struct
			c.DBType = c.DBType + *c.ArrType
		case "USER-DEFINED":
			if c.UDTName == "hstore" {
				c.Type = "types.HStore"
				c.DBType = "hstore"
			} else {
				c.Type = "string"
				fmt.Printf("Warning: Incompatible data type detected: %s\n", c.UDTName)
			}
		default:
			c.Type = "null.String"
		}
	} else {
		switch c.DBType {
		case "bigint", "bigserial":
			c.Type = "int64"
		case "integer", "serial":
			c.Type = "int"
		case "smallint", "smallserial":
			c.Type = "int16"
		case "decimal", "numeric", "double precision":
			c.Type = "float64"
		case "real":
			c.Type = "float32"
		case "bit", "interval", "uuint", "bit varying", "character", "money", "character varying", "cidr", "inet", "macaddr", "text", "uuid", "xml":
			c.Type = "string"
		case "json", "jsonb":
			c.Type = "types.JSON"
		case "bytea":
			c.Type = "[]byte"
		case "boolean":
			c.Type = "bool"
		case "date", "time", "timestamp without time zone", "timestamp with time zone":
			c.Type = "time.Time"
		case "ARRAY":
			c.Type = getArrayType(c)
			// Make DBType something like ARRAYinteger for parsing with randomize.Struct
			c.DBType = c.DBType + *c.ArrType
		case "USER-DEFINED":
			if c.UDTName == "hstore" {
				c.Type = "types.HStore"
				c.DBType = "hstore"
			} else {
				c.Type = "string"
				fmt.Printf("Warning: Incompatible data type detected: %s\n", c.UDTName)
			}
		default:
			c.Type = "string"
		}
	}

	return c
}

// getArrayType returns the correct boil.Array type for each database type
func getArrayType(c bdb.Column) string {
	switch *c.ArrType {
	case "bigint", "bigserial", "integer", "serial", "smallint", "smallserial":
		return "types.Int64Array"
	case "bytea":
		return "types.BytesArray"
	case "bit", "interval", "uuint", "bit varying", "character", "money", "character varying", "cidr", "inet", "macaddr", "text", "uuid", "xml":
		return "types.StringArray"
	case "boolean":
		return "types.BoolArray"
	case "decimal", "numeric", "double precision", "real":
		return "types.Float64Array"
	default:
		return "types.StringArray"
	}
}

// RightQuote is the quoting character for the right side of the identifier
func (p *PostgresDriver) RightQuote() byte {
	return '"'
}

// LeftQuote is the quoting character for the left side of the identifier
func (p *PostgresDriver) LeftQuote() byte {
	return '"'
}

// IndexPlaceholders returns true to indicate PSQL supports indexed placeholders
func (p *PostgresDriver) IndexPlaceholders() bool {
	return true
}
<|endoftext|>


<commit_before>package syscalls

import (
	"encoding/binary"
	"fmt"
	"github.com/lunixbochs/struc"
	"os"
	"syscall"

	"../models"
)

func errno(err error) uint64 {
	if err != nil {
		return ^uint64(err.(syscall.Errno))
	}
	return 0
}

type U models.Usercorn

type Syscall struct {
	Func func(u U, args []uint64) uint64
	Args []int
}

func exit(u U, args []uint64) uint64 {
	syscall.Exit(int(args[0]))
	return 0
}

func read(u U, args []uint64) uint64 {
	tmp := make([]byte, args[2])
	n, _ := syscall.Read(int(args[0]), tmp)
	u.MemWrite(args[1], tmp[:n])
	return uint64(n)
}

func write(u U, args []uint64) uint64 {
	mem, _ := u.MemRead(args[1], args[2])
	n, _ := syscall.Write(int(args[0]), mem)
	return uint64(n)
}

func open(u U, args []uint64) uint64 {
	path, _ := u.MemReadStr(args[0])
	fd, _ := syscall.Open(path, int(args[1]), uint32(args[2]))
	return uint64(fd)
}

func _close(u U, args []uint64) uint64 {
	syscall.Close(int(args[0]))
	return 0
}

func lseek(u U, args []uint64) uint64 {
	off, _ := syscall.Seek(int(args[0]), int64(args[1]), int(args[2]))
	return uint64(off)
}

func mmap(u U, args []uint64) uint64 {
	size := args[1]
	addr, _ := u.Mmap(args[0], size)
	fd, off := int(int32(args[4])), int64(args[5])
	if fd > 0 {
		fd2, _ := syscall.Dup(fd)
		f := os.NewFile(uintptr(fd2), "")
		f.Seek(off, 0)
		tmp := make([]byte, size)
		n, _ := f.Read(tmp)
		u.MemWrite(addr, tmp[:n])
	}
	return uint64(addr)
}

func munmap(u U, args []uint64) uint64 {
	return 0
}

func mprotect(u U, args []uint64) uint64 {
	return 0
}

func brk(u U, args []uint64) uint64 {
	// TODO: return is Linux specific
	ret, _ := u.Brk(args[0])
	return ret
}

func fstat(u U, args []uint64) uint64 {
	var stat syscall.Stat_t
	err := syscall.Fstat(int(args[0]), &stat)
	if err != nil {
		return 1
	}
	err = struc.Pack(u.MemWriter(args[1]), &stat)
	if err != nil {
		panic(err)
	}
	return 0
}

func getcwd(u U, args []uint64) uint64 {
	wd, _ := os.Getwd()
	if uint64(len(wd)) > args[1] {
		wd = wd[:args[1]]
	}
	u.MemWrite(args[0], []byte(wd))
	return 0
}

func access(u U, args []uint64) uint64 {
	// TODO: portability
	path, _ := u.MemReadStr(args[0])
	err := syscall.Access(path, uint32(args[1]))
	return errno(err)
}

func writev(u U, args []uint64) uint64 {
	ptr := u.MemReader(args[1])
	var i uint64
	for i = 0; i < args[2]; i++ {
		// TODO: bits support (via Usercorn.Bits() I think)
		var iovec Iovec64
		// TODO: endian support
		struc.UnpackWithOrder(ptr, &iovec, binary.LittleEndian)
		data, _ := u.MemRead(iovec.Base, iovec.Len)
		syscall.Write(int(args[0]), data)
	}
	return 0
}

type A []int

var syscalls = map[string]Syscall{
	"exit": {exit, A{INT}},
	// "fork": {fork, A{}},
	"read":     {read, A{FD, OBUF, LEN}},
	"write":    {write, A{FD, BUF, LEN}},
	"open":     {open, A{STR, INT, INT}},
	"close":    {_close, A{FD}},
	"lseek":    {lseek, A{FD, OFF, INT}},
	"mmap":     {mmap, A{PTR, LEN, INT, INT, FD, OFF}},
	"munmap":   {munmap, A{PTR, LEN}},
	"mprotect": {mprotect, A{PTR, LEN, INT}},
	"brk":      {brk, A{PTR}},
	"fstat":    {fstat, A{FD, PTR}},
	"getcwd":   {getcwd, A{PTR, LEN}},
	"access":   {access, A{STR, INT}},
	"writev":   {writev, A{FD, PTR, INT}},
}

func Call(u models.Usercorn, name string, getArgs func(n int) ([]uint64, error), strace bool) (uint64, error) {
	s, ok := syscalls[name]
	if !ok {
		panic(fmt.Errorf("Unknown syscall: %s", name))
	}
	args, err := getArgs(len(s.Args))
	if err != nil {
		return 0, err
	}
	if strace {
		Trace(u, name, args)
	}
	ret := s.Func(u, args)
	if strace {
		TraceRet(u, name, args, ret)
	}
	return ret, nil
}
<commit_msg>use human-readable syscall arg names<commit_after>package syscalls

import (
	"encoding/binary"
	"fmt"
	"github.com/lunixbochs/struc"
	"os"
	"syscall"

	"../models"
)

func errno(err error) uint64 {
	if err != nil {
		return ^uint64(err.(syscall.Errno))
	}
	return 0
}

type U models.Usercorn

type Syscall struct {
	Func func(u U, a []uint64) uint64
	Args []int
}

func exit(u U, a []uint64) uint64 {
	code := int(a[0])
	syscall.Exit(code)
	return 0
}

func read(u U, a []uint64) uint64 {
	fd, buf, size := int(a[0]), a[1], a[2]
	tmp := make([]byte, size)
	n, _ := syscall.Read(fd, tmp)
	u.MemWrite(buf, tmp[:n])
	return uint64(n)
}

func write(u U, a []uint64) uint64 {
	fd, buf, size := int(a[0]), a[1], a[2]
	mem, _ := u.MemRead(buf, size)
	n, _ := syscall.Write(fd, mem)
	return uint64(n)
}

func open(u U, a []uint64) uint64 {
	path, _ := u.MemReadStr(a[0])
	mode, flags := int(a[1]), uint32(a[2])
	fd, _ := syscall.Open(path, mode, flags)
	return uint64(fd)
}

func _close(u U, a []uint64) uint64 {
	fd := int(a[0])
	syscall.Close(fd)
	return 0
}

func lseek(u U, a []uint64) uint64 {
	fd, offset, whence := int(a[0]), int64(a[1]), int(a[2])
	off, _ := syscall.Seek(fd, offset, whence)
	return uint64(off)
}

func mmap(u U, a []uint64) uint64 {
	addr_hint, size, prot, flags, fd, off := a[0], a[1], a[2], a[3], int(int32(a[4])), int64(a[5])
	prot, flags = flags, prot // ignore go error
	addr, _ := u.Mmap(addr_hint, size)
	if fd > 0 {
		fd2, _ := syscall.Dup(fd)
		f := os.NewFile(uintptr(fd2), "")
		f.Seek(off, 0)
		tmp := make([]byte, size)
		n, _ := f.Read(tmp)
		u.MemWrite(addr, tmp[:n])
	}
	return uint64(addr)
}

func munmap(u U, a []uint64) uint64 {
	return 0
}

func mprotect(u U, a []uint64) uint64 {
	return 0
}

func brk(u U, a []uint64) uint64 {
	// TODO: return is Linux specific
	addr := a[0]
	ret, _ := u.Brk(addr)
	return ret
}

func fstat(u U, a []uint64) uint64 {
	fd, buf := int(a[0]), a[1]
	var stat syscall.Stat_t
	err := syscall.Fstat(fd, &stat)
	if err != nil {
		return 1
	}
	err = struc.Pack(u.MemWriter(buf), &stat)
	if err != nil {
		panic(err)
	}
	return 0
}

func getcwd(u U, a []uint64) uint64 {
	buf, size := a[0], a[1]
	wd, _ := os.Getwd()
	if uint64(len(wd)) > size {
		wd = wd[:size]
	}
	u.MemWrite(buf, []byte(wd))
	return 0
}

func access(u U, a []uint64) uint64 {
	// TODO: portability
	path, _ := u.MemReadStr(a[0])
	amode := uint32(a[1])
	err := syscall.Access(path, amode)
	return errno(err)
}

func writev(u U, a []uint64) uint64 {
	fd, iov, count := int(a[0]), a[1], a[2]
	ptr := u.MemReader(iov)
	var i uint64
	for i = 0; i < count; i++ {
		// TODO: bits support (via Usercorn.Bits() I think)
		var iovec Iovec64
		// TODO: endian support
		struc.UnpackWithOrder(ptr, &iovec, binary.LittleEndian)
		data, _ := u.MemRead(iovec.Base, iovec.Len)
		syscall.Write(fd, data)
	}
	return 0
}

type A []int

var syscalls = map[string]Syscall{
	"exit": {exit, A{INT}},
	// "fork": {fork, A{}},
	"read":     {read, A{FD, OBUF, LEN}},
	"write":    {write, A{FD, BUF, LEN}},
	"open":     {open, A{STR, INT, INT}},
	"close":    {_close, A{FD}},
	"lseek":    {lseek, A{FD, OFF, INT}},
	"mmap":     {mmap, A{PTR, LEN, INT, INT, FD, OFF}},
	"munmap":   {munmap, A{PTR, LEN}},
	"mprotect": {mprotect, A{PTR, LEN, INT}},
	"brk":      {brk, A{PTR}},
	"fstat":    {fstat, A{FD, PTR}},
	"getcwd":   {getcwd, A{PTR, LEN}},
	"access":   {access, A{STR, INT}},
	"writev":   {writev, A{FD, PTR, INT}},
}

func Call(u models.Usercorn, name string, getArgs func(n int) ([]uint64, error), strace bool) (uint64, error) {
	s, ok := syscalls[name]
	if !ok {
		panic(fmt.Errorf("Unknown syscall: %s", name))
	}
	args, err := getArgs(len(s.Args))
	if err != nil {
		return 0, err
	}
	if strace {
		Trace(u, name, args)
	}
	ret := s.Func(u, args)
	if strace {
		TraceRet(u, name, args, ret)
	}
	return ret, nil
}
<|endoftext|>


<commit_before>package kinetic

import (
	"errors"
	"sync"
	"sync/atomic"

	"github.com/aws/aws-sdk-go/aws"

	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	awsKinesis "github.com/aws/aws-sdk-go/service/kinesis"
	"github.com/aws/aws-sdk-go/service/kinesis/kinesisiface"
)

const (
	atSequenceNumber = iota
	afterSequenceNumber
	trimHorizon
	latest

	statusCreating = iota
	statusDeleting
	statusActive
	staticUpdating

	kinesisWritesPerSec int = 1000
	kinesisReadsPerSec  int = 5
	// Timeout TODO
	Timeout = 60
)

// ErrNilShardIterator is an error for when we get back a nil shard iterator
var ErrNilShardIterator = errors.New("Nil shard iterator")

// ErrNilShardStatus is an error for when we get back a nil stream status
var ErrNilShardStatus = errors.New("Nil stream status")

// Empty is an empty struct.  It is mostly used for counting semaphore purposes
type Empty struct{}

var (
	conf = getConfig()

	// ShardIterTypes are the types of iterators to use within Kinesis
	ShardIterTypes shardIteratorTypes = map[int]string{
		atSequenceNumber:    "AT_SEQUENCE_NUMBER",
		afterSequenceNumber: "AFTER_SEQUENCE_NUMBER",
		trimHorizon:         "TRIM_HORIZON",
		latest:              "LATEST",
	}

	streamStatuses streamStatusTypes = map[int]string{
		statusCreating: "CREATING",
		statusDeleting: "DELETING",
		statusActive:   "ACTIVE",
		staticUpdating: "UPDATING",
	}
)

type msgFn func([]byte, *sync.WaitGroup)
type shardIteratorTypes map[int]string
type streamStatusTypes map[int]string

type kinesis struct {
	stream            string
	shard             string
	endPoint          string
	shardIteratorType string
	shardIterator     string
	sequenceNumber    string
	sequenceNumberMu  sync.Mutex

	client kinesisiface.KinesisAPI

	msgCount int64
	errCount int64
}

func (k *kinesis) init(stream, shard, shardIteratorType, accessKey, secretKey, region string) (*kinesis, error) {

	sess, err := authenticate(accessKey, secretKey)
	conf := aws.NewConfig().WithRegion(region)
	if k.endPoint != "" {
		conf = conf.WithEndpoint(k.endPoint)
	}
	client := awsKinesis.New(sess, conf)

	k = &kinesis{
		stream:            stream,
		shard:             shard,
		shardIteratorType: shardIteratorType,
		client:            client,
	}
	if err != nil {
		return k, err
	}

	err = k.initShardIterator()
	if err != nil {
		return k, err
	}

	return k, nil
}

func (k *kinesis) initShardIterator() error {
	// log.Println(k.sequenceNumber)
	var awsSeqNumber *string
	if k.sequenceNumber != "" {
		awsSeqNumber = aws.String(k.sequenceNumber)
	}
	resp, err := k.client.GetShardIterator(&awsKinesis.GetShardIteratorInput{
		ShardId:                aws.String(k.shard),             // Required
		ShardIteratorType:      aws.String(k.shardIteratorType), // Required
		StreamName:             aws.String(k.stream),            // Required
		StartingSequenceNumber: awsSeqNumber,
	})
	if err != nil {
		return err
	}
	if resp.ShardIterator != nil {
		return k.setShardIterator(*resp.ShardIterator)
	}

	return ErrNilShardIterator
}

func (k *kinesis) setSequenceNumber(sequenceNum string) {
	if sequenceNum == "" || len(sequenceNum) == 0 {
		return
	}

	k.sequenceNumberMu.Lock()
	k.sequenceNumber = sequenceNum
	k.sequenceNumberMu.Unlock()
}

func (k *kinesis) setShardIterator(shardIter string) error {
	if shardIter == "" || len(shardIter) == 0 {
		return errors.New("Attempted to set shard iterator with empty value")
	}

	k.shardIterator = shardIter

	return nil
}

func (k *kinesis) checkActive() (bool, error) {
	status, err := k.client.DescribeStream(&awsKinesis.DescribeStreamInput{
		StreamName: aws.String("StreamName"), // Required
	})
	if err != nil {
		return false, err
	}
	if status.StreamDescription.StreamStatus == nil {
		return false, ErrNilShardStatus
	}
	if streamStatuses[statusActive] == *status.StreamDescription.StreamStatus {
		return true, nil
	}
	return false, nil
}

func (k *kinesis) newClient(endpoint, stream string) (*awsKinesis.Kinesis, error) {
	k.endPoint = endpoint
	conf := &aws.Config{}
	conf = conf.WithCredentials(
		credentials.NewStaticCredentials("BAD_ACCESS_KEY", "BAD_SECRET_KEY", "BAD_TOKEN"),
	).WithEndpoint(endpoint).WithRegion("us-east-1").WithDisableSSL(true).WithMaxRetries(3)

	// fake region
	sess, err := session.NewSessionWithOptions(session.Options{Config: *conf})
	return awsKinesis.New(sess, conf), err
}

func (k *kinesis) refreshClient(accessKey, secretKey, region string) error {
	sess, err := authenticate(accessKey, secretKey)
	conf := aws.NewConfig().WithRegion(region).WithEndpoint(k.endPoint)
	if err != nil {
		return err
	}
	k.client = awsKinesis.New(sess, conf)
	return nil
}

func (k *kinesis) decMsgCount() {
	atomic.AddInt64(&k.msgCount, -1)
}

func (k *kinesis) incMsgCount() {
	atomic.AddInt64(&k.msgCount, 1)
}

func (k *kinesis) getMsgCount() int64 {
	return atomic.LoadInt64(&k.msgCount)
}

func (k *kinesis) decErrCount() {
	atomic.AddInt64(&k.errCount, -1)
}

func (k *kinesis) incErrCount() {
	atomic.AddInt64(&k.errCount, 1)
}

func (k *kinesis) getErrCount() int64 {
	return atomic.LoadInt64(&k.errCount)
}

func getLock(sem chan Empty) {
	sem <- Empty{}
}
<commit_msg>automatically set shardIteratorType to AT_SEQUENCE_NUMBER if a sequence number is stored inside kinesis struct<commit_after>package kinetic

import (
	"errors"
	"sync"
	"sync/atomic"

	"github.com/aws/aws-sdk-go/aws"

	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	awsKinesis "github.com/aws/aws-sdk-go/service/kinesis"
	"github.com/aws/aws-sdk-go/service/kinesis/kinesisiface"
)

const (
	atSequenceNumber = iota
	afterSequenceNumber
	trimHorizon
	latest

	statusCreating = iota
	statusDeleting
	statusActive
	staticUpdating

	kinesisWritesPerSec int = 1000
	kinesisReadsPerSec  int = 5
	// Timeout TODO
	Timeout = 60
)

// ErrNilShardIterator is an error for when we get back a nil shard iterator
var ErrNilShardIterator = errors.New("Nil shard iterator")

// ErrNilShardStatus is an error for when we get back a nil stream status
var ErrNilShardStatus = errors.New("Nil stream status")

// Empty is an empty struct.  It is mostly used for counting semaphore purposes
type Empty struct{}

var (
	conf = getConfig()

	// ShardIterTypes are the types of iterators to use within Kinesis
	ShardIterTypes shardIteratorTypes = map[int]string{
		atSequenceNumber:    "AT_SEQUENCE_NUMBER",
		afterSequenceNumber: "AFTER_SEQUENCE_NUMBER",
		trimHorizon:         "TRIM_HORIZON",
		latest:              "LATEST",
	}

	streamStatuses streamStatusTypes = map[int]string{
		statusCreating: "CREATING",
		statusDeleting: "DELETING",
		statusActive:   "ACTIVE",
		staticUpdating: "UPDATING",
	}
)

type msgFn func([]byte, *sync.WaitGroup)
type shardIteratorTypes map[int]string
type streamStatusTypes map[int]string

type kinesis struct {
	stream            string
	shard             string
	endPoint          string
	shardIteratorType string
	shardIterator     string
	sequenceNumber    string
	sequenceNumberMu  sync.Mutex

	client kinesisiface.KinesisAPI

	msgCount int64
	errCount int64
}

func (k *kinesis) init(stream, shard, shardIteratorType, accessKey, secretKey, region string) (*kinesis, error) {

	sess, err := authenticate(accessKey, secretKey)
	conf := aws.NewConfig().WithRegion(region)
	if k.endPoint != "" {
		conf = conf.WithEndpoint(k.endPoint)
	}
	client := awsKinesis.New(sess, conf)

	k = &kinesis{
		stream:            stream,
		shard:             shard,
		shardIteratorType: shardIteratorType,
		client:            client,
	}
	if err != nil {
		return k, err
	}

	err = k.initShardIterator()
	if err != nil {
		return k, err
	}

	return k, nil
}

func (k *kinesis) initShardIterator() error {
	// log.Println(k.sequenceNumber)
	var awsSeqNumber *string
	if k.sequenceNumber != "" {
		awsSeqNumber = aws.String(k.sequenceNumber)
		k.shardIteratorType = ShardIterTypes[atSequenceNumber]
	}
	resp, err := k.client.GetShardIterator(&awsKinesis.GetShardIteratorInput{
		ShardId:                aws.String(k.shard),             // Required
		ShardIteratorType:      aws.String(k.shardIteratorType), // Required
		StreamName:             aws.String(k.stream),            // Required
		StartingSequenceNumber: awsSeqNumber,
	})
	if err != nil {
		return err
	}
	if resp.ShardIterator != nil {
		return k.setShardIterator(*resp.ShardIterator)
	}

	return ErrNilShardIterator
}

func (k *kinesis) setSequenceNumber(sequenceNum string) {
	if sequenceNum == "" || len(sequenceNum) == 0 {
		return
	}

	k.sequenceNumberMu.Lock()
	k.sequenceNumber = sequenceNum
	k.sequenceNumberMu.Unlock()
}

func (k *kinesis) setShardIterator(shardIter string) error {
	if shardIter == "" || len(shardIter) == 0 {
		return errors.New("Attempted to set shard iterator with empty value")
	}

	k.shardIterator = shardIter

	return nil
}

func (k *kinesis) checkActive() (bool, error) {
	status, err := k.client.DescribeStream(&awsKinesis.DescribeStreamInput{
		StreamName: aws.String("StreamName"), // Required
	})
	if err != nil {
		return false, err
	}
	if status.StreamDescription.StreamStatus == nil {
		return false, ErrNilShardStatus
	}
	if streamStatuses[statusActive] == *status.StreamDescription.StreamStatus {
		return true, nil
	}
	return false, nil
}

func (k *kinesis) newClient(endpoint, stream string) (*awsKinesis.Kinesis, error) {
	k.endPoint = endpoint
	conf := &aws.Config{}
	conf = conf.WithCredentials(
		credentials.NewStaticCredentials("BAD_ACCESS_KEY", "BAD_SECRET_KEY", "BAD_TOKEN"),
	).WithEndpoint(endpoint).WithRegion("us-east-1").WithDisableSSL(true).WithMaxRetries(3)

	// fake region
	sess, err := session.NewSessionWithOptions(session.Options{Config: *conf})
	return awsKinesis.New(sess, conf), err
}

func (k *kinesis) refreshClient(accessKey, secretKey, region string) error {
	sess, err := authenticate(accessKey, secretKey)
	conf := aws.NewConfig().WithRegion(region).WithEndpoint(k.endPoint)
	if err != nil {
		return err
	}
	k.client = awsKinesis.New(sess, conf)
	return nil
}

func (k *kinesis) decMsgCount() {
	atomic.AddInt64(&k.msgCount, -1)
}

func (k *kinesis) incMsgCount() {
	atomic.AddInt64(&k.msgCount, 1)
}

func (k *kinesis) getMsgCount() int64 {
	return atomic.LoadInt64(&k.msgCount)
}

func (k *kinesis) decErrCount() {
	atomic.AddInt64(&k.errCount, -1)
}

func (k *kinesis) incErrCount() {
	atomic.AddInt64(&k.errCount, 1)
}

func (k *kinesis) getErrCount() int64 {
	return atomic.LoadInt64(&k.errCount)
}

func getLock(sem chan Empty) {
	sem <- Empty{}
}
<|endoftext|>


<commit_before>package system

import (
	"net/http"
	"github.com/golang/glog"
	"github.com/zenazn/goji/web"
	"github.com/elcct/defaultproject/models"
	"github.com/gorilla/sessions"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

// Makes sure templates are stored in the context
func (application *Application) ApplyTemplates(c *web.C, h http.Handler) http.Handler {
	fn := func(w http.ResponseWriter, r *http.Request) {
		c.Env["Template"] = application.Template
		h.ServeHTTP(w, r)
	}
	return http.HandlerFunc(fn)
}

// Makes sure controllers can have access to session
func (application *Application) ApplySessions(c *web.C, h http.Handler) http.Handler {
	fn := func(w http.ResponseWriter, r *http.Request) {
		session, _ := application.Store.Get(r, "session")
		c.Env["Session"] = session
		h.ServeHTTP(w, r)
	}
	return http.HandlerFunc(fn)
}

// Makes sure controllers can have access to the database
func (application *Application) ApplyDatabase(c *web.C, h http.Handler) http.Handler {
	fn := func(w http.ResponseWriter, r *http.Request) {
		session := application.DBSession.Clone()
		defer session.Close()
		c.Env["DBSession"] = session
		c.Env["DBName"] = application.Configuration.Database.Database
		h.ServeHTTP(w, r)
	}
	return http.HandlerFunc(fn)
}

func (application *Application) ApplyAuth(c *web.C, h http.Handler) http.Handler {
	fn := func(w http.ResponseWriter, r *http.Request) {
		session := c.Env["Session"].(*sessions.Session)
		if userId, ok := session.Values["User"].(bson.ObjectId); ok {
			dbSession := c.Env["DBSession"].(*mgo.Session)
			database := dbSession.DB(c.Env["DBName"].(string))

			user := new(models.User)
			err := database.C("users").Find(bson.M{"_id": userId}).One(&user)
			if err != nil {
				glog.Warningf("Auth error: %v", err)
				c.Env["User"] = nil
			} else {
				c.Env["User"] = user
			}
		}
		h.ServeHTTP(w, r)
	}
	return http.HandlerFunc(fn)
}

<commit_msg>Fixed memory leak noted in #1. Thnx @resal81<commit_after>package system

import (
	"net/http"
	"github.com/golang/glog"
	"github.com/zenazn/goji/web"
	"github.com/elcct/defaultproject/models"
	"github.com/gorilla/context"
	"github.com/gorilla/sessions"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

// Makes sure templates are stored in the context
func (application *Application) ApplyTemplates(c *web.C, h http.Handler) http.Handler {
	fn := func(w http.ResponseWriter, r *http.Request) {
		c.Env["Template"] = application.Template
		h.ServeHTTP(w, r)
	}
	return http.HandlerFunc(fn)
}

// Makes sure controllers can have access to session
func (application *Application) ApplySessions(c *web.C, h http.Handler) http.Handler {
	fn := func(w http.ResponseWriter, r *http.Request) {
		session, _ := application.Store.Get(r, "session")
		c.Env["Session"] = session
		h.ServeHTTP(w, r)
		context.Clear(r)
	}
	return http.HandlerFunc(fn)
}

// Makes sure controllers can have access to the database
func (application *Application) ApplyDatabase(c *web.C, h http.Handler) http.Handler {
	fn := func(w http.ResponseWriter, r *http.Request) {
		session := application.DBSession.Clone()
		defer session.Close()
		c.Env["DBSession"] = session
		c.Env["DBName"] = application.Configuration.Database.Database
		h.ServeHTTP(w, r)
	}
	return http.HandlerFunc(fn)
}

func (application *Application) ApplyAuth(c *web.C, h http.Handler) http.Handler {
	fn := func(w http.ResponseWriter, r *http.Request) {
		session := c.Env["Session"].(*sessions.Session)
		if userId, ok := session.Values["User"].(bson.ObjectId); ok {
			dbSession := c.Env["DBSession"].(*mgo.Session)
			database := dbSession.DB(c.Env["DBName"].(string))

			user := new(models.User)
			err := database.C("users").Find(bson.M{"_id": userId}).One(&user)
			if err != nil {
				glog.Warningf("Auth error: %v", err)
				c.Env["User"] = nil
			} else {
				c.Env["User"] = user
			}
		}
		h.ServeHTTP(w, r)
	}
	return http.HandlerFunc(fn)
}

<|endoftext|>


<commit_before>// +build linux

package keyring

import (
	"github.com/aulanov/go.dbus"
	"fmt"
	"os"
	"encoding/json"
)

const (
	DBUS_SERVICE_NAME = "org.kde.kwalletd"
	DBUS_PATH = "/modules/kwalletd"
	APPID = "aws-vault"
	FOLDER = "aws-vault"
)

func init() {
	wallet, err := newKwallet()
	if err != nil {
		fmt.Fprintln(os.Stderr, "Failed to connect to kwallet:", err)
		return
	}

	supportedBackends[KWalletBackend] = opener(func(name string) (Keyring, error) {
		if name == "" {
			name = "kdewallet"
		}

		return &kwalletKeyring{
			wallet: *wallet,
			name: name,
		}, nil
	})
}

type kwalletKeyring struct {
	wallet kwalletBinding
	name string
	handle int32
}

func (k *kwalletKeyring) openWallet() error {
	isOpen, err := k.wallet.IsOpen(k.handle)
	if err != nil {
		return err
	}

	if !isOpen {
		handle, err := k.wallet.Open(k.name, 0, APPID)
		if err != nil {
			return err
		}
		k.handle = handle
	}

	return nil
}

func (k *kwalletKeyring) Get(key string) (Item, error) {
	err := k.openWallet()
	if err != nil {
		return Item{}, err
	}

	data, err := k.wallet.ReadEntry(k.handle, FOLDER, key, APPID)
	if err != nil {
		return Item{}, err
	}

	item := Item{}
	err = json.Unmarshal(data, &item)
	if err != nil {
		return Item{}, err
	}

	return item, nil
}

func (k *kwalletKeyring) Set(item Item) error {
	err := k.openWallet()
	if err != nil {
		return err
	}

	data, err := json.Marshal(item)
	if err != nil {
		return err
	}

	err = k.wallet.WriteEntry(k.handle, FOLDER, item.Key, data, APPID)
	if err != nil {
		return err
	}

	return nil
}

func (k *kwalletKeyring) Remove(key string) error {
	err := k.openWallet()
	if err != nil {
		return err
	}

	err = k.wallet.RemoveEntry(k.handle, FOLDER, key, APPID)
	if err != nil {
		return err
	}

	return nil
}

func (k *kwalletKeyring) Keys() ([]string, error) {
	err := k.openWallet()
	if err != nil {
		return []string{}, err
	}

	entries, err := k.wallet.EntryList(k.handle, FOLDER, APPID)
	if err != nil {
		return []string{}, err
	}

	return entries, nil
}

func newKwallet() (*kwalletBinding, error) {
	conn, err := dbus.SessionBus()
	if err != nil {
		return nil, err
	}

	return &kwalletBinding{
		conn.Object(DBUS_SERVICE_NAME, DBUS_PATH),
	}, nil
}

// Dumb Dbus bindings for kwallet bindings with types
type kwalletBinding struct {
	dbus dbus.BusObject
}

// method bool org.kde.KWallet.isOpen(int handle)
func (k *kwalletBinding) IsOpen(handle int32) (bool, error) {
	call := k.dbus.Call("org.kde.KWallet.isOpen", 0, handle)
	if call.Err != nil {
		return false, call.Err
	}

	return call.Body[0].(bool), call.Err
}

// method int org.kde.KWallet.open(QString wallet, qlonglong wId, QString appid)
func (k *kwalletBinding) Open(name string, wID int64, appid string) (int32, error) {
	call := k.dbus.Call("org.kde.KWallet.open", 0, name, wID, appid)
	if call.Err != nil {
		return 0, call.Err
	}

	return call.Body[0].(int32), call.Err
}

// method QStringList org.kde.KWallet.entryList(int handle, QString folder, QString appid)
func (k *kwalletBinding) EntryList(handle int32, folder string, appid string) ([]string, error) {
	call := k.dbus.Call("org.kde.KWallet.entryList", 0, handle, folder, appid)
	if call.Err != nil {
		return []string{}, call.Err
	}

	return call.Body[0].([]string), call.Err
}

// method int org.kde.KWallet.writeEntry(int handle, QString folder, QString key, QByteArray value, QString appid)
func (k *kwalletBinding) WriteEntry(handle int32, folder string, key string, value []byte, appid string) error {
	call := k.dbus.Call("org.kde.KWallet.writeEntry", 0, handle, folder, key, value, appid)
	if call.Err != nil {
		return call.Err
	}

	return call.Err
}

// method int org.kde.KWallet.removeEntry(int handle, QString folder, QString key, QString appid)
func (k *kwalletBinding) RemoveEntry(handle int32, folder string, key string, appid string) error {
	call := k.dbus.Call("org.kde.KWallet.removeEntry", 0, handle, folder, key, appid)
	if call.Err != nil {
		return call.Err
	}

	return call.Err
}

// method QByteArray org.kde.KWallet.readEntry(int handle, QString folder, QString key, QString appid)
func (k *kwalletBinding) ReadEntry(handle int32, folder string, key string, appid string) ([]byte, error) {
	call := k.dbus.Call("org.kde.KWallet.readEntry", 0, handle, folder, key, appid)
	if call.Err != nil {
		return []byte{}, call.Err
	}

	return call.Body[0].([]byte), call.Err
}
<commit_msg>Only try and contact dbus if the backend is selected<commit_after>// +build linux

package keyring

import (
	"encoding/json"

	"github.com/aulanov/go.dbus"
)

const (
	DBUS_SERVICE_NAME = "org.kde.kwalletd"
	DBUS_PATH         = "/modules/kwalletd"
	APPID             = "aws-vault"
	FOLDER            = "aws-vault"
)

func init() {
	supportedBackends[KWalletBackend] = opener(func(name string) (Keyring, error) {
		if name == "" {
			name = "kdewallet"
		}

		wallet, err := newKwallet()
		if err != nil {
			return nil, err
		}

		return &kwalletKeyring{
			wallet: *wallet,
			name:   name,
		}, nil
	})
}

type kwalletKeyring struct {
	wallet kwalletBinding
	name   string
	handle int32
}

func (k *kwalletKeyring) openWallet() error {
	isOpen, err := k.wallet.IsOpen(k.handle)
	if err != nil {
		return err
	}

	if !isOpen {
		handle, err := k.wallet.Open(k.name, 0, APPID)
		if err != nil {
			return err
		}
		k.handle = handle
	}

	return nil
}

func (k *kwalletKeyring) Get(key string) (Item, error) {
	err := k.openWallet()
	if err != nil {
		return Item{}, err
	}

	data, err := k.wallet.ReadEntry(k.handle, FOLDER, key, APPID)
	if err != nil {
		return Item{}, err
	}

	item := Item{}
	err = json.Unmarshal(data, &item)
	if err != nil {
		return Item{}, err
	}

	return item, nil
}

func (k *kwalletKeyring) Set(item Item) error {
	err := k.openWallet()
	if err != nil {
		return err
	}

	data, err := json.Marshal(item)
	if err != nil {
		return err
	}

	err = k.wallet.WriteEntry(k.handle, FOLDER, item.Key, data, APPID)
	if err != nil {
		return err
	}

	return nil
}

func (k *kwalletKeyring) Remove(key string) error {
	err := k.openWallet()
	if err != nil {
		return err
	}

	err = k.wallet.RemoveEntry(k.handle, FOLDER, key, APPID)
	if err != nil {
		return err
	}

	return nil
}

func (k *kwalletKeyring) Keys() ([]string, error) {
	err := k.openWallet()
	if err != nil {
		return []string{}, err
	}

	entries, err := k.wallet.EntryList(k.handle, FOLDER, APPID)
	if err != nil {
		return []string{}, err
	}

	return entries, nil
}

func newKwallet() (*kwalletBinding, error) {
	conn, err := dbus.SessionBus()
	if err != nil {
		return nil, err
	}

	return &kwalletBinding{
		conn.Object(DBUS_SERVICE_NAME, DBUS_PATH),
	}, nil
}

// Dumb Dbus bindings for kwallet bindings with types
type kwalletBinding struct {
	dbus dbus.BusObject
}

// method bool org.kde.KWallet.isOpen(int handle)
func (k *kwalletBinding) IsOpen(handle int32) (bool, error) {
	call := k.dbus.Call("org.kde.KWallet.isOpen", 0, handle)
	if call.Err != nil {
		return false, call.Err
	}

	return call.Body[0].(bool), call.Err
}

// method int org.kde.KWallet.open(QString wallet, qlonglong wId, QString appid)
func (k *kwalletBinding) Open(name string, wID int64, appid string) (int32, error) {
	call := k.dbus.Call("org.kde.KWallet.open", 0, name, wID, appid)
	if call.Err != nil {
		return 0, call.Err
	}

	return call.Body[0].(int32), call.Err
}

// method QStringList org.kde.KWallet.entryList(int handle, QString folder, QString appid)
func (k *kwalletBinding) EntryList(handle int32, folder string, appid string) ([]string, error) {
	call := k.dbus.Call("org.kde.KWallet.entryList", 0, handle, folder, appid)
	if call.Err != nil {
		return []string{}, call.Err
	}

	return call.Body[0].([]string), call.Err
}

// method int org.kde.KWallet.writeEntry(int handle, QString folder, QString key, QByteArray value, QString appid)
func (k *kwalletBinding) WriteEntry(handle int32, folder string, key string, value []byte, appid string) error {
	call := k.dbus.Call("org.kde.KWallet.writeEntry", 0, handle, folder, key, value, appid)
	if call.Err != nil {
		return call.Err
	}

	return call.Err
}

// method int org.kde.KWallet.removeEntry(int handle, QString folder, QString key, QString appid)
func (k *kwalletBinding) RemoveEntry(handle int32, folder string, key string, appid string) error {
	call := k.dbus.Call("org.kde.KWallet.removeEntry", 0, handle, folder, key, appid)
	if call.Err != nil {
		return call.Err
	}

	return call.Err
}

// method QByteArray org.kde.KWallet.readEntry(int handle, QString folder, QString key, QString appid)
func (k *kwalletBinding) ReadEntry(handle int32, folder string, key string, appid string) ([]byte, error) {
	call := k.dbus.Call("org.kde.KWallet.readEntry", 0, handle, folder, key, appid)
	if call.Err != nil {
		return []byte{}, call.Err
	}

	return call.Body[0].([]byte), call.Err
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package engine

import (
	"golang.org/x/net/context"

	"github.com/keybase/client/go/kex2"
	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol"
)

// DeviceAdd is an engine.
type DeviceAdd struct {
	libkb.Contextified
}

// NewDeviceAdd creates a DeviceAdd engine.
func NewDeviceAdd(g *libkb.GlobalContext) *DeviceAdd {
	return &DeviceAdd{
		Contextified: libkb.NewContextified(g),
	}
}

// Name is the unique engine name.
func (e *DeviceAdd) Name() string {
	return "DeviceAdd"
}

// GetPrereqs returns the engine prereqs.
func (e *DeviceAdd) Prereqs() Prereqs {
	return Prereqs{Device: true}
}

// RequiredUIs returns the required UIs.
func (e *DeviceAdd) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{libkb.ProvisionUIKind}
}

// SubConsumers returns the other UI consumers for this engine.
func (e *DeviceAdd) SubConsumers() []libkb.UIConsumer {
	return []libkb.UIConsumer{
		&Kex2Provisioner{},
	}
}

// Run starts the engine.
func (e *DeviceAdd) Run(ctx *Context) (err error) {
	e.G().Log.Debug("+ DeviceAdd.Run()")
	defer func() { e.G().Log.Debug("- DeviceAdd.Run() -> %s", libkb.ErrToOk(err)) }()

	arg := keybase1.ChooseDeviceTypeArg{Kind: keybase1.ChooseType_NEW_DEVICE}
	provisioneeType, err := ctx.ProvisionUI.ChooseDeviceType(context.TODO(), arg)
	if err != nil {
		return err
	}
	e.G().Log.Debug("provisionee device type: %v", provisioneeType)

	// make a new secret:
	secret, err := libkb.NewKex2Secret()
	if err != nil {
		return err
	}
	e.G().Log.Debug("secret phrase received")

	// provisioner needs ppstream, and UI is confusing when it asks for
	// it at the same time as asking for the secret, so get it first
	// before prompting for the kex2 secret:
	pps, err := e.G().LoginState().GetPassphraseStream(ctx.SecretUI)
	if err != nil {
		return err
	}

	// create provisioner engine
	provisioner := NewKex2Provisioner(e.G(), secret.Secret(), pps)

	var canceler func()

	// display secret and prompt for secret from X in a goroutine:
	go func() {
		sb := secret.Secret()
		arg := keybase1.DisplayAndPromptSecretArg{
			Secret:          sb[:],
			Phrase:          secret.Phrase(),
			OtherDeviceType: provisioneeType,
		}
		var contxt context.Context
		contxt, canceler = context.WithCancel(context.Background())
		receivedSecret, err := ctx.ProvisionUI.DisplayAndPromptSecret(contxt, arg)
		if err != nil {
			// XXX ???
			e.G().Log.Warning("DisplayAndPromptSecret error: %s", err)
		} else if receivedSecret.Secret != nil && len(receivedSecret.Secret) > 0 {
			e.G().Log.Debug("received secret, adding to provisioner")
			var ks kex2.Secret
			copy(ks[:], receivedSecret.Secret)
			provisioner.AddSecret(ks)
		} else if len(receivedSecret.Phrase) > 0 {
			e.G().Log.Debug("received secret phrase, adding to provisioner")
			ks, err := libkb.NewKex2SecretFromPhrase(receivedSecret.Phrase)
			if err != nil {
				e.G().Log.Warning("DisplayAndPromptSecret error: %s", err)
			} else {
				provisioner.AddSecret(ks.Secret())
			}
		}
	}()

	defer func() {
		if canceler != nil {
			e.G().Log.Debug("canceling DisplayAndPromptSecret call")
			canceler()
		}
	}()

	// run provisioner
	if err = RunEngine(provisioner, ctx); err != nil {
		if err == kex2.ErrHelloTimeout {
			return libkb.CanceledError{M: "Failed to provision device: are you sure you typed the secret properly?"}
		}
		return err
	}

	return nil
}
<commit_msg>Send UserChanged after device add<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package engine

import (
	"golang.org/x/net/context"

	"github.com/keybase/client/go/kex2"
	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol"
)

// DeviceAdd is an engine.
type DeviceAdd struct {
	libkb.Contextified
}

// NewDeviceAdd creates a DeviceAdd engine.
func NewDeviceAdd(g *libkb.GlobalContext) *DeviceAdd {
	return &DeviceAdd{
		Contextified: libkb.NewContextified(g),
	}
}

// Name is the unique engine name.
func (e *DeviceAdd) Name() string {
	return "DeviceAdd"
}

// GetPrereqs returns the engine prereqs.
func (e *DeviceAdd) Prereqs() Prereqs {
	return Prereqs{Device: true}
}

// RequiredUIs returns the required UIs.
func (e *DeviceAdd) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{libkb.ProvisionUIKind}
}

// SubConsumers returns the other UI consumers for this engine.
func (e *DeviceAdd) SubConsumers() []libkb.UIConsumer {
	return []libkb.UIConsumer{
		&Kex2Provisioner{},
	}
}

// Run starts the engine.
func (e *DeviceAdd) Run(ctx *Context) (err error) {
	e.G().Log.Debug("+ DeviceAdd.Run()")
	defer func() { e.G().Log.Debug("- DeviceAdd.Run() -> %s", libkb.ErrToOk(err)) }()

	arg := keybase1.ChooseDeviceTypeArg{Kind: keybase1.ChooseType_NEW_DEVICE}
	provisioneeType, err := ctx.ProvisionUI.ChooseDeviceType(context.TODO(), arg)
	if err != nil {
		return err
	}
	e.G().Log.Debug("provisionee device type: %v", provisioneeType)

	// make a new secret:
	secret, err := libkb.NewKex2Secret()
	if err != nil {
		return err
	}
	e.G().Log.Debug("secret phrase received")

	// provisioner needs ppstream, and UI is confusing when it asks for
	// it at the same time as asking for the secret, so get it first
	// before prompting for the kex2 secret:
	pps, err := e.G().LoginState().GetPassphraseStream(ctx.SecretUI)
	if err != nil {
		return err
	}

	// create provisioner engine
	provisioner := NewKex2Provisioner(e.G(), secret.Secret(), pps)

	var canceler func()

	// display secret and prompt for secret from X in a goroutine:
	go func() {
		sb := secret.Secret()
		arg := keybase1.DisplayAndPromptSecretArg{
			Secret:          sb[:],
			Phrase:          secret.Phrase(),
			OtherDeviceType: provisioneeType,
		}
		var contxt context.Context
		contxt, canceler = context.WithCancel(context.Background())
		receivedSecret, err := ctx.ProvisionUI.DisplayAndPromptSecret(contxt, arg)
		if err != nil {
			// XXX ???
			e.G().Log.Warning("DisplayAndPromptSecret error: %s", err)
		} else if receivedSecret.Secret != nil && len(receivedSecret.Secret) > 0 {
			e.G().Log.Debug("received secret, adding to provisioner")
			var ks kex2.Secret
			copy(ks[:], receivedSecret.Secret)
			provisioner.AddSecret(ks)
		} else if len(receivedSecret.Phrase) > 0 {
			e.G().Log.Debug("received secret phrase, adding to provisioner")
			ks, err := libkb.NewKex2SecretFromPhrase(receivedSecret.Phrase)
			if err != nil {
				e.G().Log.Warning("DisplayAndPromptSecret error: %s", err)
			} else {
				provisioner.AddSecret(ks.Secret())
			}
		}
	}()

	defer func() {
		if canceler != nil {
			e.G().Log.Debug("canceling DisplayAndPromptSecret call")
			canceler()
		}
	}()

	// run provisioner
	if err = RunEngine(provisioner, ctx); err != nil {
		if err == kex2.ErrHelloTimeout {
			return libkb.CanceledError{M: "Failed to provision device: are you sure you typed the secret properly?"}
		}
		return err
	}

	// provisioning was successful, so the user has changed:
	e.G().NotifyRouter.HandleUserChanged(e.G().Env.GetUID())

	return nil
}
<|endoftext|>


<commit_before>package main

import (
    "log"
    "encoding/json"
    "io/ioutil"
    "os"
    "path/filepath"
    "archive/zip"
    "errors"
    "io"
)

type ModInfoList struct {
    Mods []ModInfo `json:"mods"`
    Destination string `json:"-"`
}
type ModInfo struct {
    Name string `json:"name"`
    Version string `json:"version"`
    Title string `json:"title"`
    Author string `json:"author"`
    FileName string `json:"file_name"`
}

func newModInfoList(destination string) (ModInfoList, error) {
    var err error
    mod_info_list := ModInfoList{
        Destination: destination,
    }

    err = mod_info_list.listInstalledMods()
    if err != nil {
        log.Printf("ModInfoList ... error listing installed Mods: %s", err)
        return mod_info_list, err
    }

    return mod_info_list, nil
}

func (mod_info_list *ModInfoList) listInstalledMods() (error) {
    var err error
    mod_info_list.Mods = nil

    err = filepath.Walk(mod_info_list.Destination, func(path string, info os.FileInfo, err error) error {
        if !info.IsDir() && filepath.Ext(path) == ".zip" {
            zip_file, err := zip.OpenReader(path)
            if err != nil {
                log.Fatalln(err)
                return err
            }

            var mod_info ModInfo
            err = mod_info.getModInfo(&zip_file.Reader)
            if err != nil {
                log.Fatalf("Error in getModInfo: %s", err)
            }

            mod_info.FileName = info.Name()
            mod_info_list.Mods = append(mod_info_list.Mods, mod_info)
        }

        return nil
    })

    if err != nil {
        log.Printf("error while walking over the given dir: %s", err)
        return err
    }

    return nil
}

func (mod_info_list *ModInfoList) deleteMod(mod_name string) (error) {
    var err error

    //search for mod, that should be deleted
    for _, mod := range mod_info_list.Mods {
        if mod.Name == mod_name {
            //delete mod
            err = os.Remove(mod_info_list.Destination + "/" + mod.FileName)
            if err != nil {
                log.Printf("ModInfoList ... error when deleting mod: %s", err)
                return err
            }

            //reload mod-list
            err = mod_info_list.listInstalledMods()
            if err != nil {
                log.Printf("ModInfoList ... error while refreshing installedModList: %s", err)
                return err
            }

            return nil
        }
    }

    log.Printf("the mod-file doesnt exists!")
    return nil
}

func (mod_info *ModInfo) getModInfo(reader *zip.Reader) error {
    for _, single_file := range reader.File {
        if single_file.FileInfo().Name() == "info.json" {
            //interpret info.json
            rc, err := single_file.Open()

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

            byte_array, err := ioutil.ReadAll(rc)
            rc.Close()
            if err != nil {
                log.Fatal(err)
                return err
            }

            err = json.Unmarshal(byte_array, mod_info)
            if err != nil {
                log.Fatalln(err)
                return err
            }

            return nil
        }
    }

    return errors.New("info.json not found in zip-file")
}

func (mod_info_list *ModInfoList) createMod(mod_name string, file_name string, mod_file io.Reader) error {
    var err error

    //save uploaded file
    new_file, err := os.Create(mod_info_list.Destination + "/" + file_name)
    if err != nil {
        log.Printf("error on creating new file - %s: %s", file_name, err)
        return err
    }
    defer new_file.Close()

    //reload the list
    err = new_file.Close()
    if err != nil {
        log.Printf("error on closing new created zip-file: %s", err)
        return err
    }

    err = mod_info_list.listInstalledMods()
    if err != nil {
        log.Printf("error on listing mod-infos: %s", err)
        return err
    }

    return nil
}
<commit_msg>reverted delete file creation<commit_after>package main

import (
    "log"
    "encoding/json"
    "io/ioutil"
    "os"
    "path/filepath"
    "archive/zip"
    "errors"
    "io"
)

type ModInfoList struct {
    Mods []ModInfo `json:"mods"`
    Destination string `json:"-"`
}
type ModInfo struct {
    Name string `json:"name"`
    Version string `json:"version"`
    Title string `json:"title"`
    Author string `json:"author"`
    FileName string `json:"file_name"`
}

func newModInfoList(destination string) (ModInfoList, error) {
    var err error
    mod_info_list := ModInfoList{
        Destination: destination,
    }

    err = mod_info_list.listInstalledMods()
    if err != nil {
        log.Printf("ModInfoList ... error listing installed Mods: %s", err)
        return mod_info_list, err
    }

    return mod_info_list, nil
}

func (mod_info_list *ModInfoList) listInstalledMods() (error) {
    var err error
    mod_info_list.Mods = nil

    err = filepath.Walk(mod_info_list.Destination, func(path string, info os.FileInfo, err error) error {
        if !info.IsDir() && filepath.Ext(path) == ".zip" {
            zip_file, err := zip.OpenReader(path)
            if err != nil {
                log.Fatalln(err)
                return err
            }

            var mod_info ModInfo
            err = mod_info.getModInfo(&zip_file.Reader)
            if err != nil {
                log.Fatalf("Error in getModInfo: %s", err)
            }

            mod_info.FileName = info.Name()
            mod_info_list.Mods = append(mod_info_list.Mods, mod_info)
        }

        return nil
    })

    if err != nil {
        log.Printf("error while walking over the given dir: %s", err)
        return err
    }

    return nil
}

func (mod_info_list *ModInfoList) deleteMod(mod_name string) (error) {
    var err error

    //search for mod, that should be deleted
    for _, mod := range mod_info_list.Mods {
        if mod.Name == mod_name {
            //delete mod
            err = os.Remove(mod_info_list.Destination + "/" + mod.FileName)
            if err != nil {
                log.Printf("ModInfoList ... error when deleting mod: %s", err)
                return err
            }

            //reload mod-list
            err = mod_info_list.listInstalledMods()
            if err != nil {
                log.Printf("ModInfoList ... error while refreshing installedModList: %s", err)
                return err
            }

            return nil
        }
    }

    log.Printf("the mod-file doesnt exists!")
    return nil
}

func (mod_info *ModInfo) getModInfo(reader *zip.Reader) error {
    for _, single_file := range reader.File {
        if single_file.FileInfo().Name() == "info.json" {
            //interpret info.json
            rc, err := single_file.Open()

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

            byte_array, err := ioutil.ReadAll(rc)
            rc.Close()
            if err != nil {
                log.Fatal(err)
                return err
            }

            err = json.Unmarshal(byte_array, mod_info)
            if err != nil {
                log.Fatalln(err)
                return err
            }

            return nil
        }
    }

    return errors.New("info.json not found in zip-file")
}

func (mod_info_list *ModInfoList) createMod(mod_name string, file_name string, mod_file io.Reader) error {
    var err error

    //save uploaded file
    new_file, err := os.Create(mod_info_list.Destination + "/" + file_name)
    if err != nil {
        log.Printf("error on creating new file - %s: %s", file_name, err)
        return err
    }
    defer new_file.Close()

    _, err = io.Copy(new_file, mod_file)
    if err != nil {
        log.Printf("error on copying file to disk: %s", err)
        return err
    }

    //reload the list
    err = new_file.Close()
    if err != nil {
        log.Printf("error on closing new created zip-file: %s", err)
        return err
    }

    err = mod_info_list.listInstalledMods()
    if err != nil {
        log.Printf("error on listing mod-infos: %s", err)
        return err
    }

    return nil
}
<|endoftext|>