<commit_before>package sql

import (
	"crypto/sha1"
	"database/sql"
	"encoding/hex"
	"errors"
	"fmt"
	"hash"
	"sync"

	"github.com/lib/pq"

	"github.com/barakmich/glog"
	"github.com/google/cayley/graph"
	"github.com/google/cayley/graph/iterator"
	"github.com/google/cayley/quad"
)

const QuadStoreType = "sql"

func init() {
	graph.RegisterQuadStore(QuadStoreType, graph.QuadStoreRegistration{
		NewFunc:           newQuadStore,
		NewForRequestFunc: nil,
		UpgradeFunc:       nil,
		InitFunc:          createSQLTables,
		IsPersistent:      true,
	})
}

var (
	hashPool = sync.Pool{
		New: func() interface{} { return sha1.New() },
	}
	hashSize = sha1.Size
)

type QuadStore struct {
	db           *sql.DB
	sqlFlavor    string
	size         int64
	lru          *cache
	noSizes      bool
	useEstimates bool
}

func connectSQLTables(addr string, _ graph.Options) (*sql.DB, error) {
	// TODO(barakmich): Parse options for more friendly addr, other SQLs.
	conn, err := sql.Open("postgres", addr)
	if err != nil {
		glog.Errorf("Couldn't open database at %s: %#v", addr, err)
		return nil, err
	}
	// "Open may just validate its arguments without creating a connection to the database."
	// "To verify that the data source name is valid, call Ping."
	// Source: http://golang.org/pkg/database/sql/#Open
	if err := conn.Ping(); err != nil {
		glog.Errorf("Couldn't open database at %s: %#v", addr, err)
		return nil, err
	}
	return conn, nil
}

func createSQLTables(addr string, options graph.Options) error {
	conn, err := connectSQLTables(addr, options)
	if err != nil {
		return err
	}
	tx, err := conn.Begin()
	if err != nil {
		glog.Errorf("Couldn't begin creation transaction: %s", err)
		return err
	}

	quadTable, err := tx.Exec(`
	CREATE TABLE quads (
		subject TEXT NOT NULL,
		predicate TEXT NOT NULL,
		object TEXT NOT NULL,
		label TEXT,
		horizon BIGSERIAL PRIMARY KEY,
		id BIGINT,
		ts timestamp,
		subject_hash TEXT NOT NULL,
		predicate_hash TEXT NOT NULL,
		object_hash TEXT NOT NULL,
		label_hash TEXT,
		UNIQUE(subject_hash, predicate_hash, object_hash, label_hash)
	);`)
	if err != nil {
		glog.Errorf("Cannot create quad table: %v", quadTable)
		return err
	}
	factor, factorOk, err := options.IntKey("db_fill_factor")
	if !factorOk {
		factor = 50
	}
	var index sql.Result

	index, err = tx.Exec(fmt.Sprintf(`
	CREATE INDEX spo_index ON quads (subject_hash) WITH (FILLFACTOR = %d);
	CREATE INDEX pos_index ON quads (predicate_hash) WITH (FILLFACTOR = %d);
	CREATE INDEX osp_index ON quads (object_hash) WITH (FILLFACTOR = %d);
	`, factor, factor, factor))
	if err != nil {
		glog.Errorf("Cannot create indices: %v", index)
		return err
	}
	tx.Commit()
	return nil
}

func newQuadStore(addr string, options graph.Options) (graph.QuadStore, error) {
	var qs QuadStore
	conn, err := connectSQLTables(addr, options)
	if err != nil {
		return nil, err
	}
	localOpt, localOptOk, err := options.BoolKey("local_optimize")
	if err != nil {
		return nil, err
	}
	qs.db = conn
	qs.sqlFlavor = "postgres"
	qs.size = -1
	qs.lru = newCache(1024)

	// Skip size checking by default.
	qs.noSizes = true
	if localOptOk {
		if localOpt {
			qs.noSizes = false
		}
	}
	qs.useEstimates, _, err = options.BoolKey("use_estimates")
	if err != nil {
		return nil, err
	}

	return &qs, nil
}

func hashOf(s string) string {
	h := hashPool.Get().(hash.Hash)
	h.Reset()
	defer hashPool.Put(h)
	key := make([]byte, 0, hashSize)
	h.Write([]byte(s))
	key = h.Sum(key)
	return hex.EncodeToString(key)
}

func (qs *QuadStore) copyFrom(tx *sql.Tx, in []graph.Delta) error {
	stmt, err := tx.Prepare(pq.CopyIn("quads", "subject", "predicate", "object", "label", "id", "ts", "subject_hash", "predicate_hash", "object_hash", "label_hash"))
	if err != nil {
		return err
	}
	for _, d := range in {
		_, err := stmt.Exec(
			d.Quad.Subject,
			d.Quad.Predicate,
			d.Quad.Object,
			d.Quad.Label,
			d.ID.Int(),
			d.Timestamp,
			hashOf(d.Quad.Subject),
			hashOf(d.Quad.Predicate),
			hashOf(d.Quad.Object),
			hashOf(d.Quad.Label),
		)
		if err != nil {
			glog.Errorf("couldn't prepare COPY statement: %v", err)
			return err
		}
	}
	_, err = stmt.Exec()
	if err != nil {
		return err
	}
	return stmt.Close()
}

func (qs *QuadStore) runTxPostgres(tx *sql.Tx, in []graph.Delta, opts graph.IgnoreOpts) error {
	allAdds := true
	for _, d := range in {
		if d.Action != graph.Add {
			allAdds = false
		}
	}
	if allAdds {
		return qs.copyFrom(tx, in)
	}

	insert, err := tx.Prepare(`INSERT INTO quads(subject, predicate, object, label, id, ts, subject_hash, predicate_hash, object_hash, label_hash) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`)
	defer insert.Close()
	if err != nil {
		glog.Errorf("Cannot prepare insert statement: %v", err)
		return err
	}
	for _, d := range in {
		switch d.Action {
		case graph.Add:
			_, err := insert.Exec(
				d.Quad.Subject,
				d.Quad.Predicate,
				d.Quad.Object,
				d.Quad.Label,
				d.ID.Int(),
				d.Timestamp,
				hashOf(d.Quad.Subject),
				hashOf(d.Quad.Predicate),
				hashOf(d.Quad.Object),
				hashOf(d.Quad.Label),
			)
			if err != nil {
				glog.Errorf("couldn't prepare INSERT statement: %v", err)
				return err
			}
		case graph.Delete:
			result, err := tx.Exec(`DELETE FROM quads WHERE subject=$1 and predicate=$2 and object=$3 and label=$4;`,
				d.Quad.Subject, d.Quad.Predicate, d.Quad.Object, d.Quad.Label)
			if err != nil {
				glog.Errorf("couldn't exec DELETE statement: %v", err)
				return err
			}
			affected, err := result.RowsAffected()
			if err != nil {
				glog.Errorf("couldn't get DELETE RowsAffected: %v", err)
				return err
			}
			if affected != 1 && !opts.IgnoreMissing {
				return errors.New("deleting non-existent triple; rolling back")
			}
		default:
			panic("unknown action")
		}
	}
	return nil
}

func (qs *QuadStore) ApplyDeltas(in []graph.Delta, opts graph.IgnoreOpts) error {
	// TODO(barakmich): Support more ignoreOpts? "ON CONFLICT IGNORE"
	tx, err := qs.db.Begin()
	if err != nil {
		glog.Errorf("couldn't begin write transaction: %v", err)
		return err
	}
	switch qs.sqlFlavor {
	case "postgres":
		err = qs.runTxPostgres(tx, in, opts)
		if err != nil {
			tx.Rollback()
			return err
		}
	default:
		panic("no support for flavor: " + qs.sqlFlavor)
	}
	return tx.Commit()
}

func (qs *QuadStore) Quad(val graph.Value) quad.Quad {
	return val.(quad.Quad)
}

func (qs *QuadStore) QuadIterator(d quad.Direction, val graph.Value) graph.Iterator {
	return NewSQLLinkIterator(qs, d, val.(string))
}

func (qs *QuadStore) NodesAllIterator() graph.Iterator {
	return NewAllIterator(qs, "nodes")
}

func (qs *QuadStore) QuadsAllIterator() graph.Iterator {
	return NewAllIterator(qs, "quads")
}

func (qs *QuadStore) ValueOf(s string) graph.Value {
	return s
}

func (qs *QuadStore) NameOf(v graph.Value) string {
	if v == nil {
		glog.V(2).Info("NameOf was nil")
		return ""
	}
	return v.(string)
}

func (qs *QuadStore) Size() int64 {
	// TODO(barakmich): Sync size with writes.
	if qs.size != -1 {
		return qs.size
	}

	query := "SELECT COUNT(*) FROM quads;"
	if qs.useEstimates {
		switch qs.sqlFlavor {
		case "postgres":
			query = "SELECT reltuples::BIGINT AS estimate FROM pg_class WHERE relname='quads';"
		default:
			panic("no estimate support for flavor: " + qs.sqlFlavor)
		}
	}

	c := qs.db.QueryRow(query)
	err := c.Scan(&qs.size)
	if err != nil {
		glog.Errorf("Couldn't execute COUNT: %v", err)
		return 0
	}
	return qs.size
}

func (qs *QuadStore) Horizon() graph.PrimaryKey {
	var horizon int64
	err := qs.db.QueryRow("SELECT horizon FROM quads ORDER BY horizon DESC LIMIT 1;").Scan(&horizon)
	if err != nil {
		glog.Errorf("Couldn't execute horizon: %v", err)
		return graph.NewSequentialKey(0)
	}
	return graph.NewSequentialKey(horizon)
}

func (qs *QuadStore) FixedIterator() graph.FixedIterator {
	return iterator.NewFixed(iterator.Identity)
}

func (qs *QuadStore) Close() {
	qs.db.Close()
}

func (qs *QuadStore) QuadDirection(in graph.Value, d quad.Direction) graph.Value {
	q := in.(quad.Quad)
	return q.Get(d)
}

func (qs *QuadStore) Type() string {
	return QuadStoreType
}

func (qs *QuadStore) sizeForIterator(isAll bool, dir quad.Direction, val string) int64 {
	var err error
	if isAll {
		return qs.Size()
	}
	if qs.noSizes {
		if dir == quad.Predicate {
			return (qs.Size() / 100) + 1
		}
		return (qs.Size() / 1000) + 1
	}
	if val, ok := qs.lru.Get(val + string(dir.Prefix())); ok {
		return val
	}
	var size int64
	glog.V(4).Infoln("sql: getting size for select %s, %s", dir.String(), val)
	err = qs.db.QueryRow(
		fmt.Sprintf("SELECT count(*) FROM quads WHERE %s_hash = $1;", dir.String()), hashOf(val)).Scan(&size)
	if err != nil {
		glog.Errorln("Error getting size from SQL database: %v", err)
		return 0
	}
	qs.lru.Put(val+string(dir.Prefix()), size)
	return size
}
<commit_msg>Fix a deadlock behavior with DELETE transactions<commit_after>package sql

import (
	"crypto/sha1"
	"database/sql"
	"encoding/hex"
	"errors"
	"fmt"
	"hash"
	"sync"

	"github.com/lib/pq"

	"github.com/barakmich/glog"
	"github.com/google/cayley/graph"
	"github.com/google/cayley/graph/iterator"
	"github.com/google/cayley/quad"
)

const QuadStoreType = "sql"

func init() {
	graph.RegisterQuadStore(QuadStoreType, graph.QuadStoreRegistration{
		NewFunc:           newQuadStore,
		NewForRequestFunc: nil,
		UpgradeFunc:       nil,
		InitFunc:          createSQLTables,
		IsPersistent:      true,
	})
}

var (
	hashPool = sync.Pool{
		New: func() interface{} { return sha1.New() },
	}
	hashSize = sha1.Size
)

type QuadStore struct {
	db           *sql.DB
	sqlFlavor    string
	size         int64
	lru          *cache
	noSizes      bool
	useEstimates bool
}

func connectSQLTables(addr string, _ graph.Options) (*sql.DB, error) {
	// TODO(barakmich): Parse options for more friendly addr, other SQLs.
	conn, err := sql.Open("postgres", addr)
	if err != nil {
		glog.Errorf("Couldn't open database at %s: %#v", addr, err)
		return nil, err
	}
	// "Open may just validate its arguments without creating a connection to the database."
	// "To verify that the data source name is valid, call Ping."
	// Source: http://golang.org/pkg/database/sql/#Open
	if err := conn.Ping(); err != nil {
		glog.Errorf("Couldn't open database at %s: %#v", addr, err)
		return nil, err
	}
	return conn, nil
}

func createSQLTables(addr string, options graph.Options) error {
	conn, err := connectSQLTables(addr, options)
	if err != nil {
		return err
	}
	tx, err := conn.Begin()
	if err != nil {
		glog.Errorf("Couldn't begin creation transaction: %s", err)
		return err
	}

	quadTable, err := tx.Exec(`
	CREATE TABLE quads (
		subject TEXT NOT NULL,
		predicate TEXT NOT NULL,
		object TEXT NOT NULL,
		label TEXT,
		horizon BIGSERIAL PRIMARY KEY,
		id BIGINT,
		ts timestamp,
		subject_hash TEXT NOT NULL,
		predicate_hash TEXT NOT NULL,
		object_hash TEXT NOT NULL,
		label_hash TEXT,
		UNIQUE(subject_hash, predicate_hash, object_hash, label_hash)
	);`)
	if err != nil {
		glog.Errorf("Cannot create quad table: %v", quadTable)
		return err
	}
	factor, factorOk, err := options.IntKey("db_fill_factor")
	if !factorOk {
		factor = 50
	}
	var index sql.Result

	index, err = tx.Exec(fmt.Sprintf(`
	CREATE INDEX spo_index ON quads (subject_hash) WITH (FILLFACTOR = %d);
	CREATE INDEX pos_index ON quads (predicate_hash) WITH (FILLFACTOR = %d);
	CREATE INDEX osp_index ON quads (object_hash) WITH (FILLFACTOR = %d);
	`, factor, factor, factor))
	if err != nil {
		glog.Errorf("Cannot create indices: %v", index)
		return err
	}
	tx.Commit()
	return nil
}

func newQuadStore(addr string, options graph.Options) (graph.QuadStore, error) {
	var qs QuadStore
	conn, err := connectSQLTables(addr, options)
	if err != nil {
		return nil, err
	}
	localOpt, localOptOk, err := options.BoolKey("local_optimize")
	if err != nil {
		return nil, err
	}
	qs.db = conn
	qs.sqlFlavor = "postgres"
	qs.size = -1
	qs.lru = newCache(1024)

	// Skip size checking by default.
	qs.noSizes = true
	if localOptOk {
		if localOpt {
			qs.noSizes = false
		}
	}
	qs.useEstimates, _, err = options.BoolKey("use_estimates")
	if err != nil {
		return nil, err
	}

	return &qs, nil
}

func hashOf(s string) string {
	h := hashPool.Get().(hash.Hash)
	h.Reset()
	defer hashPool.Put(h)
	key := make([]byte, 0, hashSize)
	h.Write([]byte(s))
	key = h.Sum(key)
	return hex.EncodeToString(key)
}

func (qs *QuadStore) copyFrom(tx *sql.Tx, in []graph.Delta) error {
	stmt, err := tx.Prepare(pq.CopyIn("quads", "subject", "predicate", "object", "label", "id", "ts", "subject_hash", "predicate_hash", "object_hash", "label_hash"))
	if err != nil {
		return err
	}
	for _, d := range in {
		_, err := stmt.Exec(
			d.Quad.Subject,
			d.Quad.Predicate,
			d.Quad.Object,
			d.Quad.Label,
			d.ID.Int(),
			d.Timestamp,
			hashOf(d.Quad.Subject),
			hashOf(d.Quad.Predicate),
			hashOf(d.Quad.Object),
			hashOf(d.Quad.Label),
		)
		if err != nil {
			glog.Errorf("couldn't prepare COPY statement: %v", err)
			return err
		}
	}
	_, err = stmt.Exec()
	if err != nil {
		return err
	}
	return stmt.Close()
}

func (qs *QuadStore) runTxPostgres(tx *sql.Tx, in []graph.Delta, opts graph.IgnoreOpts) error {
	allAdds := true
	for _, d := range in {
		if d.Action != graph.Add {
			allAdds = false
		}
	}
	if allAdds {
		return qs.copyFrom(tx, in)
	}

	for _, d := range in {
		switch d.Action {
		case graph.Add:
			_, err := tx.Exec(`INSERT INTO quads(subject, predicate, object, label, id, ts, subject_hash, predicate_hash, object_hash, label_hash) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10);`,
				d.Quad.Subject,
				d.Quad.Predicate,
				d.Quad.Object,
				d.Quad.Label,
				d.ID.Int(),
				d.Timestamp,
				hashOf(d.Quad.Subject),
				hashOf(d.Quad.Predicate),
				hashOf(d.Quad.Object),
				hashOf(d.Quad.Label),
			)
			if err != nil {
				glog.Errorf("couldn't exec INSERT statement: %v", err)
				return err
			}
		case graph.Delete:
			result, err := tx.Exec(`DELETE FROM quads WHERE subject=$1 and predicate=$2 and object=$3 and label=$4;`,
				d.Quad.Subject, d.Quad.Predicate, d.Quad.Object, d.Quad.Label)
			if err != nil {
				glog.Errorf("couldn't exec DELETE statement: %v", err)
				return err
			}
			affected, err := result.RowsAffected()
			if err != nil {
				glog.Errorf("couldn't get DELETE RowsAffected: %v", err)
				return err
			}
			if affected != 1 && !opts.IgnoreMissing {
				return errors.New("deleting non-existent triple; rolling back")
			}
		default:
			panic("unknown action")
		}
	}
	return nil
}

func (qs *QuadStore) ApplyDeltas(in []graph.Delta, opts graph.IgnoreOpts) error {
	// TODO(barakmich): Support more ignoreOpts? "ON CONFLICT IGNORE"
	tx, err := qs.db.Begin()
	if err != nil {
		glog.Errorf("couldn't begin write transaction: %v", err)
		return err
	}
	switch qs.sqlFlavor {
	case "postgres":
		err = qs.runTxPostgres(tx, in, opts)
		if err != nil {
			tx.Rollback()
			return err
		}
	default:
		panic("no support for flavor: " + qs.sqlFlavor)
	}
	return tx.Commit()
}

func (qs *QuadStore) Quad(val graph.Value) quad.Quad {
	return val.(quad.Quad)
}

func (qs *QuadStore) QuadIterator(d quad.Direction, val graph.Value) graph.Iterator {
	return NewSQLLinkIterator(qs, d, val.(string))
}

func (qs *QuadStore) NodesAllIterator() graph.Iterator {
	return NewAllIterator(qs, "nodes")
}

func (qs *QuadStore) QuadsAllIterator() graph.Iterator {
	return NewAllIterator(qs, "quads")
}

func (qs *QuadStore) ValueOf(s string) graph.Value {
	return s
}

func (qs *QuadStore) NameOf(v graph.Value) string {
	if v == nil {
		glog.V(2).Info("NameOf was nil")
		return ""
	}
	return v.(string)
}

func (qs *QuadStore) Size() int64 {
	// TODO(barakmich): Sync size with writes.
	if qs.size != -1 {
		return qs.size
	}

	query := "SELECT COUNT(*) FROM quads;"
	if qs.useEstimates {
		switch qs.sqlFlavor {
		case "postgres":
			query = "SELECT reltuples::BIGINT AS estimate FROM pg_class WHERE relname='quads';"
		default:
			panic("no estimate support for flavor: " + qs.sqlFlavor)
		}
	}

	c := qs.db.QueryRow(query)
	err := c.Scan(&qs.size)
	if err != nil {
		glog.Errorf("Couldn't execute COUNT: %v", err)
		return 0
	}
	return qs.size
}

func (qs *QuadStore) Horizon() graph.PrimaryKey {
	var horizon int64
	err := qs.db.QueryRow("SELECT horizon FROM quads ORDER BY horizon DESC LIMIT 1;").Scan(&horizon)
	if err != nil {
		glog.Errorf("Couldn't execute horizon: %v", err)
		return graph.NewSequentialKey(0)
	}
	return graph.NewSequentialKey(horizon)
}

func (qs *QuadStore) FixedIterator() graph.FixedIterator {
	return iterator.NewFixed(iterator.Identity)
}

func (qs *QuadStore) Close() {
	qs.db.Close()
}

func (qs *QuadStore) QuadDirection(in graph.Value, d quad.Direction) graph.Value {
	q := in.(quad.Quad)
	return q.Get(d)
}

func (qs *QuadStore) Type() string {
	return QuadStoreType
}

func (qs *QuadStore) sizeForIterator(isAll bool, dir quad.Direction, val string) int64 {
	var err error
	if isAll {
		return qs.Size()
	}
	if qs.noSizes {
		if dir == quad.Predicate {
			return (qs.Size() / 100) + 1
		}
		return (qs.Size() / 1000) + 1
	}
	if val, ok := qs.lru.Get(val + string(dir.Prefix())); ok {
		return val
	}
	var size int64
	glog.V(4).Infoln("sql: getting size for select %s, %s", dir.String(), val)
	err = qs.db.QueryRow(
		fmt.Sprintf("SELECT count(*) FROM quads WHERE %s_hash = $1;", dir.String()), hashOf(val)).Scan(&size)
	if err != nil {
		glog.Errorln("Error getting size from SQL database: %v", err)
		return 0
	}
	qs.lru.Put(val+string(dir.Prefix()), size)
	return size
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"github.com/leonelquinteros/gorand"
	"gopkg.in/ldap.v2"
	"log"
	"net/http"
	_ "strings"
	"time"
)

var authOptions *AuthorizationOptions

type Error struct {
	Message string
	Path    string
}

type AuthorizationOptions struct {
	realm          string
	ldapServer     string
	port           int
	username       string
	password       string
	ldapSearchBase string
}

type Credentials struct {
	Username string
	Password string
}

// setup handlers
func init() {
	var logoutHandler = http.HandlerFunc(Logout)
	var loginHandler = http.HandlerFunc(Login)
	http.Handle("/login", loginHandler)
	http.Handle("/logout", logoutHandler)
}

// LdapAuthFunc authenticates a user against an LDAP server
// The Request parameter is probably not necessary.
func LdapAuthFunc(creds *Credentials) bool {
	hostaddr := fmt.Sprintf("%s:%d", authOptions.ldapServer, authOptions.port) // move to config file
	conn, err := ldap.Dial("tcp", hostaddr)
	if err != nil {
		log.Fatal(err)
	}

	defer conn.Close()

	err = conn.StartTLS(&tls.Config{InsecureSkipVerify: true})
	if err != nil {
		log.Fatal(err)
	}

	err = conn.Bind(authOptions.realm+"\\"+creds.Username, creds.Password)
	if err == nil {
		return true
	} else {
		log.Printf("LDAP: %s", err.Error())
		return false
	}
}

func FindUser(username string) (Person, error) {
	var user Person
	if authOptions == nil {
		log.Panicf("Auth options should not be nil")
	}
	hostaddr := fmt.Sprintf("%s:%d", authOptions.ldapServer, authOptions.port)
	conn, err := ldap.Dial("tcp", hostaddr)
	if err != nil {
		log.Fatal(err)
		return user, err
	}

	defer conn.Close()

	err = conn.StartTLS(&tls.Config{InsecureSkipVerify: true})
	if err != nil {
		log.Fatal(err)
		return user, err
	}

	err = conn.Bind(authOptions.realm+"\\"+authOptions.username, authOptions.password)
	if err != nil {
		log.Fatal(err)
		return user, err
	}

	searchRequest := ldap.NewSearchRequest(
		authOptions.ldapSearchBase,
		ldap.ScopeWholeSubtree, ldap.NeverDerefAliases, 0, 0, false,
		fmt.Sprintf("(&(objectClass=organizationalPerson)(sAMAccountName=%s))", username),
		[]string{"dn", "cn", "title", "department", "telephoneNumber", "mobile", "physicalDeliveryOfficeName"},
		nil,
	)
	res, err := conn.Search(searchRequest)
	if err != nil || len(res.Entries) != 1 {
		log.Fatal(err)
	}

	ldapPerson := res.Entries[0]

	user = Person{
		Username:   username,
		Name:       ldapPerson.GetAttributeValue("cn"),
		Department: ldapPerson.GetAttributeValue("department"),
		Telephone:  ldapPerson.GetAttributeValue("telephoneNumber"),
		Mobile:     ldapPerson.GetAttributeValue("mobile"),
		Office:     ldapPerson.GetAttributeValue("physicalDeliveryOfficeName"),
	}
	return user, err
}

func CreateUser(username string) error {
	var user Person
	user, err := FindUser(username)
	if err != nil {
		return err
	}

	err = AddPerson(
		username,
		user.Name,
		user.Department,
		user.Telephone,
		user.Mobile,
		user.Office,
	)
	return err
}

// Update all the database users with attributes
// from LDAP. Accordingly, this takes a set of
// LDAP connection options as a parameter
func UpdateLdap(options AuthorizationOptions) error {
	authOptions = &options
	// get users
	people, err := GetUsers()
	if err != nil {
		log.Fatalf("Failed to get users from the database: %s", err.Error())
		return err
	}
	// for each user, get the LDAP entry
	for _, user := range people {
		updated, err := FindUser(user.Username)
		if err != nil {
			log.Printf("Failed to get user %s from the LDAP Server: %s", user.Username, err.Error())
		}
		// and update the database
		user.Name = updated.Name
		user.Department = updated.Department
		user.Office = updated.Office
		user.Telephone = updated.Telephone
		user.Mobile = updated.Mobile
		if SetPersonDetails(user) != nil {
			log.Printf("Failed to update user %s: %s", user.Username, err)
		}

		fmt.Printf(". ")
	}
	fmt.Printf("\n")
	return nil
}

func AuthorizationMiddleware(options AuthorizationOptions, next http.Handler) http.Handler {
	authOptions = &options
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var session string = ""
		cookies := r.Cookies()
		for _, cookie := range cookies {
			if cookie.Name == "session" {
				session = cookie.Value
				log.Println("Session ", cookie.Value)
				break
			}
		}

		if username, err := ValidateSession(session); err == nil {
			next.ServeHTTP(w, r.WithContext(newContextWithUsername(r.Context(), username)))
		} else {
			log.Println("no session found")
			sessionErr := &Error{Message: "unauthorized", Path: "/login"}
			log.Printf(sessionErr.Message)
			content, _ := json.Marshal(sessionErr)
			http.Error(w, string(content), http.StatusUnauthorized)
			if err := json.NewEncoder(w).Encode(sessionErr); err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
			}
		}
	})
}

func Logout(w http.ResponseWriter, r *http.Request) {
	log.Println("logged out")
	var cookie http.Cookie
	cookies := r.Cookies()
	for _, c := range cookies {
		if c.Name == "session" {
			cookie = *c
			break
		}
	}
	if err := RemoveSession(cookie.Value); err != nil {
		http.Error(w, "", http.StatusInternalServerError)
	}
	cookie.Value = ""
	cookie.Expires = time.Unix(0, 0)
	http.SetCookie(w, &cookie)
	w.Write([]byte("Logged out"))
}

func Login(w http.ResponseWriter, r *http.Request) {
	log.Println("Creating session")
	id, err := gorand.UUIDv4()
	session, err := gorand.MarshalUUID(id)
	if err != nil {
		http.Error(w, "", http.StatusInternalServerError)
	}

	creds := new(Credentials)
	err = json.NewDecoder(r.Body).Decode(creds)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	if LdapAuthFunc(creds) {
		var person *Person

		if person, err = GetPerson(creds.Username); err != nil {
			// create user from ldap store
			if err = CreateUser(creds.Username); err != nil {
				log.Fatalf("Failed to create user: %s", err.Error())
				http.Error(w, err.Error(), http.StatusInternalServerError)
			}
			log.Printf("Created user")
			if person, err = GetPerson(creds.Username); err != nil {
				log.Fatal(err)
				http.Error(w, err.Error(), http.StatusInternalServerError)
			}
		}
		if err = CreateSession(session, person.ID); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		log.Printf("created session %s", session)

		cookie := &http.Cookie{Name: "session", Value: session, HttpOnly: false}
		http.SetCookie(w, cookie)
		w.Write([]byte("success"))

	} else {
		log.Printf("login failed")
		sessionErr := Error{Message: "login failed", Path: "/login"}
		if err := json.NewEncoder(w).Encode(sessionErr); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	}
}

type requestKey int

const requestUsernameKey = 0

func usernameFromContext(ctx context.Context) string {
	return ctx.Value(requestUsernameKey).(string)
}

func newContextWithUsername(ctx context.Context, username string) context.Context {
	return context.WithValue(ctx, requestUsernameKey, username)
}

func unauthorizedHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Access-Control-Allow-Origin", "http://localhost:4050")
	w.Header().Add("Access-Control-Allow-Methods", "GET, PUT, OPTIONS, HEAD")
	w.Header().Add("Access-Control-Allow-Headers", "Content-Type, Authorization, X-Requested-With")
	w.Header().Add("Access-Control-Allow-Credentials", "true")
	if r.Method == "OPTIONS" {
		return
	} else {
		http.Error(w, "", http.StatusUnauthorized)
	}
}
<commit_msg>updated docs, deleted dead code<commit_after>package main

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"github.com/leonelquinteros/gorand"
	"gopkg.in/ldap.v2"
	"log"
	"net/http"
	_ "strings"
	"time"
)

var authOptions *AuthorizationOptions

// An error that can be serialized to
// JSON and returned to the client
type Error struct {
	// Error message
	Message string
	// The http path to the api endpoint
	// to fix the error, eg: /login
	Path string
}

// A set of options necessary to find
// and login to a LDAP server
type AuthorizationOptions struct {
	realm          string
	ldapServer     string
	port           int
	username       string
	password       string
	ldapSearchBase string
}

type Credentials struct {
	Username string
	Password string
}

// setup handlers
func init() {
	var logoutHandler = http.HandlerFunc(Logout)
	var loginHandler = http.HandlerFunc(Login)
	http.Handle("/login", loginHandler)
	http.Handle("/logout", logoutHandler)
}

// LdapAuthFunc authenticates a user against an LDAP server
// The Request parameter is probably not necessary.
func LdapAuthFunc(creds *Credentials) bool {
	hostaddr := fmt.Sprintf("%s:%d", authOptions.ldapServer, authOptions.port) // move to config file
	conn, err := ldap.Dial("tcp", hostaddr)
	if err != nil {
		log.Fatal(err)
	}

	defer conn.Close()

	err = conn.StartTLS(&tls.Config{InsecureSkipVerify: true})
	if err != nil {
		log.Fatal(err)
	}

	err = conn.Bind(authOptions.realm+"\\"+creds.Username, creds.Password)
	if err == nil {
		return true
	} else {
		log.Printf("LDAP: %s", err.Error())
		return false
	}
}

// Find a person in LDAP
func FindUser(username string) (Person, error) {
	var user Person
	if authOptions == nil {
		log.Panicf("Auth options should not be nil")
	}
	hostaddr := fmt.Sprintf("%s:%d", authOptions.ldapServer, authOptions.port)
	conn, err := ldap.Dial("tcp", hostaddr)
	if err != nil {
		log.Fatal(err)
		return user, err
	}

	defer conn.Close()

	err = conn.StartTLS(&tls.Config{InsecureSkipVerify: true})
	if err != nil {
		log.Fatal(err)
		return user, err
	}

	err = conn.Bind(authOptions.realm+"\\"+authOptions.username, authOptions.password)
	if err != nil {
		log.Fatal(err)
		return user, err
	}

	searchRequest := ldap.NewSearchRequest(
		authOptions.ldapSearchBase,
		ldap.ScopeWholeSubtree, ldap.NeverDerefAliases, 0, 0, false,
		fmt.Sprintf("(&(objectClass=organizationalPerson)(sAMAccountName=%s))", username),
		[]string{"dn", "cn", "title", "department", "telephoneNumber", "mobile", "physicalDeliveryOfficeName"},
		nil,
	)
	res, err := conn.Search(searchRequest)
	if err != nil || len(res.Entries) != 1 {
		log.Fatal(err)
	}

	ldapPerson := res.Entries[0]

	user = Person{
		Username:   username,
		Name:       ldapPerson.GetAttributeValue("cn"),
		Department: ldapPerson.GetAttributeValue("department"),
		Telephone:  ldapPerson.GetAttributeValue("telephoneNumber"),
		Mobile:     ldapPerson.GetAttributeValue("mobile"),
		Office:     ldapPerson.GetAttributeValue("physicalDeliveryOfficeName"),
	}
	return user, err
}

// Create a user for a given username. This user must exist
// in LDAP
func CreateUser(username string) error {
	var user Person
	user, err := FindUser(username)
	if err != nil {
		return err
	}

	err = AddPerson(
		username,
		user.Name,
		user.Department,
		user.Telephone,
		user.Mobile,
		user.Office,
	)
	return err
}

// Update all the database users with attributes
// from LDAP. Accordingly, this takes a set of
// LDAP connection options as a parameter
func UpdateLdap(options AuthorizationOptions) error {
	authOptions = &options
	// get users
	people, err := GetUsers()
	if err != nil {
		log.Fatalf("Failed to get users from the database: %s", err.Error())
		return err
	}
	// for each user, get the LDAP entry
	for _, user := range people {
		updated, err := FindUser(user.Username)
		if err != nil {
			log.Printf("Failed to get user %s from the LDAP Server: %s", user.Username, err.Error())
		}
		// and update the database
		user.Name = updated.Name
		user.Department = updated.Department
		user.Office = updated.Office
		user.Telephone = updated.Telephone
		user.Mobile = updated.Mobile
		if SetPersonDetails(user) != nil {
			log.Printf("Failed to update user %s: %s", user.Username, err)
		}

		fmt.Printf(". ")
	}
	fmt.Printf("\n")
	return nil
}

// Handles cookie-based authentication. An incoming
// request will have its session ID read from a cookie, and if
// the session is not valid, returns a JSON-encoded response
// redirecting to the Login api endpoint.
func AuthorizationMiddleware(options AuthorizationOptions, next http.Handler) http.Handler {
	authOptions = &options
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var session string = ""
		cookies := r.Cookies()
		for _, cookie := range cookies {
			if cookie.Name == "session" {
				session = cookie.Value
				log.Println("Session ", cookie.Value)
				break
			}
		}

		if username, err := ValidateSession(session); err == nil {
			next.ServeHTTP(w, r.WithContext(newContextWithUsername(r.Context(), username)))
		} else {
			log.Println("no session found")
			sessionErr := &Error{Message: "unauthorized", Path: "/login"}
			log.Printf(sessionErr.Message)
			content, _ := json.Marshal(sessionErr)
			http.Error(w, string(content), http.StatusUnauthorized)
			if err := json.NewEncoder(w).Encode(sessionErr); err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
			}
		}
	})
}

// Removes a session from the database and sets
// the cookie to be expired.
func Logout(w http.ResponseWriter, r *http.Request) {
	log.Println("logged out")
	var cookie http.Cookie
	cookies := r.Cookies()
	for _, c := range cookies {
		if c.Name == "session" {
			cookie = *c
			break
		}
	}
	if err := RemoveSession(cookie.Value); err != nil {
		http.Error(w, "", http.StatusInternalServerError)
	}
	cookie.Value = ""
	cookie.Expires = time.Unix(0, 0)
	http.SetCookie(w, &cookie)
	w.Write([]byte("Logged out"))
}

// Authenticate a user and create a session in the
// database, returning a cookie with the session ID
func Login(w http.ResponseWriter, r *http.Request) {
	log.Println("Creating session")
	id, err := gorand.UUIDv4()
	session, err := gorand.MarshalUUID(id)
	if err != nil {
		http.Error(w, "", http.StatusInternalServerError)
	}

	creds := new(Credentials)
	err = json.NewDecoder(r.Body).Decode(creds)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	if LdapAuthFunc(creds) {
		var person *Person

		if person, err = GetPerson(creds.Username); err != nil {
			// create user from ldap store
			if err = CreateUser(creds.Username); err != nil {
				log.Fatalf("Failed to create user: %s", err.Error())
				http.Error(w, err.Error(), http.StatusInternalServerError)
			}
			log.Printf("Created user")
			if person, err = GetPerson(creds.Username); err != nil {
				log.Fatal(err)
				http.Error(w, err.Error(), http.StatusInternalServerError)
			}
		}
		if err = CreateSession(session, person.ID); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		log.Printf("created session %s", session)

		cookie := &http.Cookie{Name: "session", Value: session, HttpOnly: false}
		http.SetCookie(w, cookie)
		w.Write([]byte("success"))

	} else {
		log.Printf("login failed")
		sessionErr := Error{Message: "login failed", Path: "/login"}
		if err := json.NewEncoder(w).Encode(sessionErr); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	}
}

const requestUsernameKey = 0

// get the username from a context object
func usernameFromContext(ctx context.Context) string {
	return ctx.Value(requestUsernameKey).(string)
}

// return a new context object with the username in it
func newContextWithUsername(ctx context.Context, username string) context.Context {
	return context.WithValue(ctx, requestUsernameKey, username)
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"errors"
	"fmt"
	log "github.com/Sirupsen/logrus"
	_ "github.com/mattn/go-sqlite3"
)

var conn *sql.DB

var statusCodes map[int]Status

func init() {
}

func ValidateSession(sessionID string) (string, error) {
	stmt, err := conn.Prepare("SELECT username FROM sessions JOIN people ON (person_id = people.id) WHERE sessions.id = ?")
	checkErr(err)
	res, err := stmt.Query(sessionID)
	checkErr(err)
	var username string
	rows := 0
	for res.Next() {
		if err = res.Scan(&username); err != nil {
			return "", err
		}
		rows++
	}
	if rows == 1 {
		return username, nil
	} else {
		return "", errors.New("session not found")
	}
}

func CreateSession(sessionID string, userID int) error {
	stmt, err := conn.Prepare("INSERT INTO sessions (id, person_id) VALUES (?, ?)")
	checkErr(err)
	stmt.Exec(sessionID, userID)
	checkErr(err)
	return err
}

func RemoveSession(sessionID string) error {
	stmt, err := conn.Prepare("DELETE FROM sessions WHERE id = ?")
	checkErr(err)
	_, err = stmt.Exec(sessionID)
	checkErr(err)
	return err
}

func checkErr(err error) {
	if err != nil {
		log.Error(err.Error())
		log.Error(" Noooooo!\n")
	}
}

func AddPerson(username string, name string, department string, telephone string, mobile string, office string, title string) (*Person, error) {
	stmt, err := conn.Prepare("INSERT INTO people (username, name, status, department, mobile, telephone, office, title) VALUES (?,?,?,?,?,?,?,?)")
	if err != nil {
		log.Fatal(err)
	}

	// there has to be a status code 0 in the db or this will fail
	_, err = stmt.Exec(username, name, 0, department, mobile, telephone, office, title)
	if err != nil {
		return nil, err
	}

	person, err := GetPerson(username)
	log.Infof("Added %s to the db", username)
	return person, err
}

func createDb(dbPath string) {
	db, err := sql.Open("sqlite3", dbPath)
	checkErr(err)
	if err != nil {
		log.Fatal("Could not open the database file")
	}
	conn = db

	rows, err := db.Query("SELECT name FROM sqlite_master WHERE type = 'table'")
	tables := make(map[string]string)
	var table string
	checkErr(err)
	for rows.Next() {
		err = rows.Scan(&table)
		checkErr(err)
		log.Infof("found table %s", table)
		tables[table] = table
	}

	if _, ok := tables["status"]; !ok {
		log.Print("creating status table")
		res, err := db.Exec("CREATE TABLE status (id INTEGER PRIMARY KEY, value TEXT)")
		checkErr(err)
		log.Print("insert status values")
		stmt, err := db.Prepare("INSERT INTO status (value) VALUES (?)")
		checkErr(err)
		res, err = stmt.Exec("In")
		rows, err := res.RowsAffected()
		if rows == 0 {
			panic(err)
		}
		checkErr(err)
		res, err = stmt.Exec("Out")
		checkErr(err)
		res, err = stmt.Exec("In Field")
		checkErr(err)
	}
	log.Print("creating people table")
	if _, ok := tables["people"]; !ok {
		_, err = db.Exec("CREATE TABLE people (id INTEGER PRIMARY KEY, username TEXT UNIQUE, name TEXT NOT NULL, department TEXT null, mobile TEXT not null default '', telephone TEXT not null default '', office TEXT not null default '', title TEXT not null default '', status int REFERENCES status(id), notes TEXT DEFAULT '', last_editor INTEGER NULL REFERENCES people(id), last_edit_time datetime DEFAULT CURRENT_TIMESTAMP)")
		checkErr(err)
	}
	if _, ok := tables["sessions"]; !ok {
		log.Print("creating sessions table")
		stmt, err := db.Prepare("CREATE TABLE sessions (id text PRIMARY KEY, person_id INTEGER REFERENCES people(id), create_time DATETIME DEFAULT CURRENT_TIMESTAMP)")
		_, err = stmt.Exec()
		checkErr(err)
	}
}

func GetUsers() ([]*Person, error) {
	log.Print("GetUsers")
	if conn == nil {
		log.Panic("Database was not open")
	}

	rows, err := conn.Query(`SELECT p.id, p.username, p.name, p.department, p.status, p.notes, p.telephone, p.mobile, p.office, p.title, l.name, p.last_edit_time
		FROM people p
		LEFT JOIN people l ON p.last_editor = l.id
		ORDER BY p.department, p.name`)
	checkErr(err)
	if err != nil {
		return nil, err
	}

	var id int
	var username string
	var name string
	var department sql.NullString
	var notes string
	var status int
	var telephone string
	var mobile string
	var office string
	var people []*Person
	var lastEditor sql.NullString
	var lastEditTime NullTime
	var title string
	statuses, err := StatusCodes()
	if err != nil {
		log.Fatalf("Failed to get status codes from the database: %s", err)
	}

	for rows.Next() {
		err = rows.Scan(
			&id,
			&username,
			&name,
			&department,
			&status,
			&notes,
			&telephone,
			&mobile,
			&office,
			&title,
			&lastEditor,
			&lastEditTime)
		checkErr(err)
		if err != nil {
			return nil, err
		}
		p := &Person{
			ID:         id,
			Name:       name,
			Username:   username,
			Department: department.String,
			Status:     statuses[status],
			Remarks:    notes,
			Telephone:  telephone,
			Mobile:     mobile,
			Office:     office,
			Title:      title,
			LastEditor: "",
		}

		p.Status = statuses[status]

		if lastEditor.Valid {
			p.LastEditor = lastEditor.String
		}
		if lastEditTime.Valid {
			p.LastEditTime = lastEditTime.Time.Local()
		}
		people = append(people, p)
	}
	log.Debugf("Got %d people from the db\n", len(people))
	rows.Close()

	return people, nil
}

func GetPerson(username string) (*Person, error) {
	log.Print("GetPerson")
	if conn == nil {
		log.Panic("Database was not open")
	}
	var person *Person
	var id int
	var uname string
	var name string
	var status int
	var notes string
	var department string
	var office sql.NullString
	var telephone sql.NullString
	var mobile sql.NullString
	var title string
	var lastEditor sql.NullString
	var lastEditTime NullTime

	stmt, err := conn.Prepare(`SELECT p.id, p.username, p.name, p.department, p.status, p.notes, p.telephone, p.mobile, p.office, p.title, l.name as last_editor, p.last_edit_time
	FROM people p left join people l on l.id = p.last_editor WHERE p.username = ?`)
	checkErr(err)
	if err != nil {
		return nil, err
	}
	rows, err := stmt.Query(username)
	checkErr(err)
	if err != nil {
		return nil, err
	}
	statuses, err := StatusCodes()
	if err != nil {
		log.Fatalf("Could not get status codes from the database: %s", err)
	}

	if rows.Next() {
		err = rows.Scan(&id, &uname, &name, &department, &status, &notes, &telephone, &mobile, &office, &title, &lastEditor, &lastEditTime)
		checkErr(err)
		if err != nil {
			return nil, err
		}
		person = &Person{
			ID:         id,
			Name:       name,
			Username:   username,
			Department: department,
			Title:      title,
			Remarks:    notes,
		}
		person.Status = statuses[status]

		if lastEditor.Valid {
			person.LastEditor = lastEditor.String
		}
		if lastEditTime.Valid {
			person.LastEditTime = lastEditTime.Time.Local()
		}
		if telephone.Valid {
			person.Telephone = telephone.String
		}
		if mobile.Valid {
			person.Mobile = mobile.String
		}
		if office.Valid {
			person.Office = office.String
		}
	} else {
		err = fmt.Errorf("No user named %s", username)
		person = nil
	}
	rows.Close()

	return person, err
}

func SetPerson(person *Person, username string) error {
	stmt, err := conn.Prepare("UPDATE people SET status = ?, notes = ?, last_editor = (select id from people where username = ?), last_edit_time = current_timestamp WHERE username = ?")
	checkErr(err)
	res, err := stmt.Exec(person.Status.Code, person.Remarks, username, person.Username)
	checkErr(err)

	rows, err := res.RowsAffected()
	checkErr(err)
	if rows != 1 {
		err = fmt.Errorf("Failed to update user %s", person.Username)
	} else {
		err = nil
	}

	return err
}

// Updates details for a user. This is meant to be used
// internally for attributes that are not editable by
// the user.
func SetPersonDetails(person *Person) error {
	stmt, err := conn.Prepare("UPDATE people SET name = ?, department = ?, telephone = ?, mobile = ?, office = ?, title = ? WHERE username = ?")
	checkErr(err)
	res, err := stmt.Exec(person.Name, person.Department, person.Telephone, person.Mobile, person.Office, person.Title, person.Username)
	checkErr(err)
	rows, err := res.RowsAffected()
	checkErr(err)
	if rows != 1 {
		err = fmt.Errorf("Failed to update user %s", person.Username)
	} else {
		err = nil
	}

	return err
}

func StatusCodes() (map[int]Status, error) {
	if statusCodes == nil {
		statusCodes = make(map[int]Status)
	}
	if len(statusCodes) > 0 {
		return statusCodes, nil
	}
	stmt, err := conn.Prepare("SELECT * FROM status")
	checkErr(err)
	if err != nil {
		statusCodes = make(map[int]Status)
		return statusCodes, err
	}
	rows, err := stmt.Query()
	checkErr(err)
	if err != nil {
		statusCodes = make(map[int]Status)
		return statusCodes, err
	}

	var status Status
	for rows.Next() {
		status = Status{}
		err = rows.Scan(&status.Code, &status.Value)
		statusCodes[status.Code] = status
		if err != nil {
			statusCodes = make(map[int]Status)
			return statusCodes, err
		}
	}
	return statusCodes, nil
}
<commit_msg>don't crash on parallel access to the statuscode cache<commit_after>package main

import (
	"database/sql"
	"errors"
	"fmt"
	log "github.com/Sirupsen/logrus"
	_ "github.com/mattn/go-sqlite3"
	"sync"
)

var conn *sql.DB

var statusCodes map[int]Status
var mutex *sync.Mutex

func init() {
	mutex = &sync.Mutex{}
}

func ValidateSession(sessionID string) (string, error) {
	stmt, err := conn.Prepare("SELECT username FROM sessions JOIN people ON (person_id = people.id) WHERE sessions.id = ?")
	checkErr(err)
	res, err := stmt.Query(sessionID)
	checkErr(err)
	var username string
	rows := 0
	for res.Next() {
		if err = res.Scan(&username); err != nil {
			return "", err
		}
		rows++
	}
	if rows == 1 {
		return username, nil
	} else {
		return "", errors.New("session not found")
	}
}

func CreateSession(sessionID string, userID int) error {
	stmt, err := conn.Prepare("INSERT INTO sessions (id, person_id) VALUES (?, ?)")
	checkErr(err)
	stmt.Exec(sessionID, userID)
	checkErr(err)
	return err
}

func RemoveSession(sessionID string) error {
	stmt, err := conn.Prepare("DELETE FROM sessions WHERE id = ?")
	checkErr(err)
	_, err = stmt.Exec(sessionID)
	checkErr(err)
	return err
}

func checkErr(err error) {
	if err != nil {
		log.Error(err.Error())
		log.Error(" Noooooo!\n")
	}
}

func AddPerson(username string, name string, department string, telephone string, mobile string, office string, title string) (*Person, error) {
	stmt, err := conn.Prepare("INSERT INTO people (username, name, status, department, mobile, telephone, office, title) VALUES (?,?,?,?,?,?,?,?)")
	if err != nil {
		log.Fatal(err)
	}

	// there has to be a status code 0 in the db or this will fail
	_, err = stmt.Exec(username, name, 0, department, mobile, telephone, office, title)
	if err != nil {
		return nil, err
	}

	person, err := GetPerson(username)
	log.Infof("Added %s to the db", username)
	return person, err
}

func createDb(dbPath string) {
	db, err := sql.Open("sqlite3", dbPath)
	checkErr(err)
	if err != nil {
		log.Fatal("Could not open the database file")
	}
	conn = db

	rows, err := db.Query("SELECT name FROM sqlite_master WHERE type = 'table'")
	tables := make(map[string]string)
	var table string
	checkErr(err)
	for rows.Next() {
		err = rows.Scan(&table)
		checkErr(err)
		log.Infof("found table %s", table)
		tables[table] = table
	}

	if _, ok := tables["status"]; !ok {
		log.Print("creating status table")
		res, err := db.Exec("CREATE TABLE status (id INTEGER PRIMARY KEY, value TEXT)")
		checkErr(err)
		log.Print("insert status values")
		stmt, err := db.Prepare("INSERT INTO status (value) VALUES (?)")
		checkErr(err)
		res, err = stmt.Exec("In")
		rows, err := res.RowsAffected()
		if rows == 0 {
			panic(err)
		}
		checkErr(err)
		res, err = stmt.Exec("Out")
		checkErr(err)
		res, err = stmt.Exec("In Field")
		checkErr(err)
	}
	log.Print("creating people table")
	if _, ok := tables["people"]; !ok {
		_, err = db.Exec("CREATE TABLE people (id INTEGER PRIMARY KEY, username TEXT UNIQUE, name TEXT NOT NULL, department TEXT null, mobile TEXT not null default '', telephone TEXT not null default '', office TEXT not null default '', title TEXT not null default '', status int REFERENCES status(id), notes TEXT DEFAULT '', last_editor INTEGER NULL REFERENCES people(id), last_edit_time datetime DEFAULT CURRENT_TIMESTAMP)")
		checkErr(err)
	}
	if _, ok := tables["sessions"]; !ok {
		log.Print("creating sessions table")
		stmt, err := db.Prepare("CREATE TABLE sessions (id text PRIMARY KEY, person_id INTEGER REFERENCES people(id), create_time DATETIME DEFAULT CURRENT_TIMESTAMP)")
		_, err = stmt.Exec()
		checkErr(err)
	}
}

func GetUsers() ([]*Person, error) {
	log.Print("GetUsers")
	if conn == nil {
		log.Panic("Database was not open")
	}

	rows, err := conn.Query(`SELECT p.id, p.username, p.name, p.department, p.status, p.notes, p.telephone, p.mobile, p.office, p.title, l.name, p.last_edit_time
		FROM people p
		LEFT JOIN people l ON p.last_editor = l.id
		ORDER BY p.department, p.name`)
	checkErr(err)
	if err != nil {
		return nil, err
	}

	var id int
	var username string
	var name string
	var department sql.NullString
	var notes string
	var status int
	var telephone string
	var mobile string
	var office string
	var people []*Person
	var lastEditor sql.NullString
	var lastEditTime NullTime
	var title string
	statuses, err := StatusCodes()
	if err != nil {
		log.Fatalf("Failed to get status codes from the database: %s", err)
	}

	for rows.Next() {
		err = rows.Scan(
			&id,
			&username,
			&name,
			&department,
			&status,
			&notes,
			&telephone,
			&mobile,
			&office,
			&title,
			&lastEditor,
			&lastEditTime)
		checkErr(err)
		if err != nil {
			return nil, err
		}
		p := &Person{
			ID:         id,
			Name:       name,
			Username:   username,
			Department: department.String,
			Status:     statuses[status],
			Remarks:    notes,
			Telephone:  telephone,
			Mobile:     mobile,
			Office:     office,
			Title:      title,
			LastEditor: "",
		}

		p.Status = statuses[status]

		if lastEditor.Valid {
			p.LastEditor = lastEditor.String
		}
		if lastEditTime.Valid {
			p.LastEditTime = lastEditTime.Time.Local()
		}
		people = append(people, p)
	}
	log.Debugf("Got %d people from the db\n", len(people))
	rows.Close()

	return people, nil
}

func GetPerson(username string) (*Person, error) {
	log.Print("GetPerson")
	if conn == nil {
		log.Panic("Database was not open")
	}
	var person *Person
	var id int
	var uname string
	var name string
	var status int
	var notes string
	var department string
	var office sql.NullString
	var telephone sql.NullString
	var mobile sql.NullString
	var title string
	var lastEditor sql.NullString
	var lastEditTime NullTime

	stmt, err := conn.Prepare(`SELECT p.id, p.username, p.name, p.department, p.status, p.notes, p.telephone, p.mobile, p.office, p.title, l.name as last_editor, p.last_edit_time
	FROM people p left join people l on l.id = p.last_editor WHERE p.username = ?`)
	checkErr(err)
	if err != nil {
		return nil, err
	}
	rows, err := stmt.Query(username)
	checkErr(err)
	if err != nil {
		return nil, err
	}
	statuses, err := StatusCodes()
	if err != nil {
		log.Fatalf("Could not get status codes from the database: %s", err)
	}

	if rows.Next() {
		err = rows.Scan(&id, &uname, &name, &department, &status, &notes, &telephone, &mobile, &office, &title, &lastEditor, &lastEditTime)
		checkErr(err)
		if err != nil {
			return nil, err
		}
		person = &Person{
			ID:         id,
			Name:       name,
			Username:   username,
			Department: department,
			Title:      title,
			Remarks:    notes,
		}
		person.Status = statuses[status]

		if lastEditor.Valid {
			person.LastEditor = lastEditor.String
		}
		if lastEditTime.Valid {
			person.LastEditTime = lastEditTime.Time.Local()
		}
		if telephone.Valid {
			person.Telephone = telephone.String
		}
		if mobile.Valid {
			person.Mobile = mobile.String
		}
		if office.Valid {
			person.Office = office.String
		}
	} else {
		err = fmt.Errorf("No user named %s", username)
		person = nil
	}
	rows.Close()

	return person, err
}

func SetPerson(person *Person, username string) error {
	stmt, err := conn.Prepare("UPDATE people SET status = ?, notes = ?, last_editor = (select id from people where username = ?), last_edit_time = current_timestamp WHERE username = ?")
	checkErr(err)
	res, err := stmt.Exec(person.Status.Code, person.Remarks, username, person.Username)
	checkErr(err)

	rows, err := res.RowsAffected()
	checkErr(err)
	if rows != 1 {
		err = fmt.Errorf("Failed to update user %s", person.Username)
	} else {
		err = nil
	}

	return err
}

// Updates details for a user. This is meant to be used
// internally for attributes that are not editable by
// the user.
func SetPersonDetails(person *Person) error {
	stmt, err := conn.Prepare("UPDATE people SET name = ?, department = ?, telephone = ?, mobile = ?, office = ?, title = ? WHERE username = ?")
	checkErr(err)
	res, err := stmt.Exec(person.Name, person.Department, person.Telephone, person.Mobile, person.Office, person.Title, person.Username)
	checkErr(err)
	rows, err := res.RowsAffected()
	checkErr(err)
	if rows != 1 {
		err = fmt.Errorf("Failed to update user %s", person.Username)
	} else {
		err = nil
	}

	return err
}

func StatusCodes() (map[int]Status, error) {
	if statusCodes == nil {
		statusCodes = make(map[int]Status)
	}
	mutex.Lock()
	if len(statusCodes) > 0 {
		mutex.Unlock()
		return statusCodes, nil
	}
	stmt, err := conn.Prepare("SELECT * FROM status")
	checkErr(err)
	if err != nil {
		statusCodes = make(map[int]Status)
		mutex.Unlock()
		return statusCodes, err
	}
	rows, err := stmt.Query()
	checkErr(err)
	if err != nil {
		statusCodes = make(map[int]Status)
		mutex.Unlock()
		return statusCodes, err
	}

	var status Status
	for rows.Next() {
		status = Status{}
		err = rows.Scan(&status.Code, &status.Value)
		statusCodes[status.Code] = status
		if err != nil {
			statusCodes = make(map[int]Status)
			mutex.Unlock()
			return statusCodes, err
		}
	}
	mutex.Unlock()
	return statusCodes, nil
}
<|endoftext|>


<commit_before>package calcium

import (
	"fmt"
	"io"
	"io/ioutil"
	"path/filepath"
	"strings"

	engineapi "github.com/docker/engine-api/client"
	enginetypes "github.com/docker/engine-api/types"
	enginecontainer "github.com/docker/engine-api/types/container"
	enginenetwork "github.com/docker/engine-api/types/network"
	"gitlab.ricebook.net/platform/core/types"
	"gitlab.ricebook.net/platform/core/utils"
	"golang.org/x/net/context"
)

// As the name says,
// blocks until the stream is empty, until we meet EOF
func ensureReaderClosed(stream io.ReadCloser) {
	if stream == nil {
		return
	}
	io.Copy(ioutil.Discard, stream)
	stream.Close()
}

// Copies config from container
// And make a new name for container
func makeContainerConfig(info enginetypes.ContainerJSON, image string) (
	*enginecontainer.Config,
	*enginecontainer.HostConfig,
	*enginenetwork.NetworkingConfig,
	string,
	error) {

	// we use `_` to join container name
	// since we don't support `_` in entrypoint, and no `_` is in suffix,
	// the last part will be suffix and second last part will be entrypoint,
	// the rest will be the appname
	parts := strings.Split(trimLeftSlash(info.Name), "_")
	length := len(parts)
	if length < 3 {
		return nil, nil, nil, "", fmt.Errorf("Bad container name format: %q", info.Name)
	}

	entrypoint := parts[length-2]
	appname := strings.Join(parts[:length-2], "_")

	suffix := utils.RandomString(6)
	containerName := strings.Join([]string{appname, entrypoint, suffix}, "_")

	config := info.Config
	config.Image = image

	hostConfig := info.HostConfig
	networkConfig := &enginenetwork.NetworkingConfig{
		EndpointsConfig: info.NetworkSettings.Networks,
	}
	return config, hostConfig, networkConfig, containerName, nil
}

// see https://github.com/docker/docker/issues/6705
// docker's stupid problem
func trimLeftSlash(name string) string {
	return strings.TrimPrefix(name, "/")
}

// make mount paths
// app.yaml支持三种写法.
// app.yaml里可以支持mount_paths的写法, 例如
// mount_paths:
//     - "/var/www/html"
//     - "/data/eggsy"
// 这样的路径会被直接挂载到permdir下面去, 例如上面的路径就是
// /mnt/mfs/permdirs/eggsy/data/eggsy
// /mnt/mfs/permdirs/eggsy/var/www/html
// 而且这些路径是可以读写的.
//
// 或者使用volumes, 参数格式跟docker一样, 例如
// volumes:
//     - "/data/test:/test:ro"
//     - "/data/testx:/testx"
// 说明把宿主机的/data/test映射到容器里的/test, 只读, 同时
// 把宿主机的/data/tests映射到容器里的/testx, 读写.
//
// 或者使用binds, 例如
// binds:
//     "/host/path":
//         bind: "/container/path"
//         ro: true
// 说明把宿主机的/host/path映射到容器里的/container/path, 并且只读
func makeMountPaths(specs types.Specs, config types.Config) ([]string, map[string]struct{}) {
	binds := []string{}
	volumes := make(map[string]struct{})
	permDirHost := filepath.Join(config.PermDir, specs.Appname)

	// mount_paths
	for _, path := range specs.MountPaths {
		hostPath := filepath.Join(permDirHost, path)
		binds = append(binds, fmt.Sprintf("%s:%s:rw", hostPath, path))
		volumes[path] = struct{}{}
	}

	// volumes
	for _, path := range specs.Volumes {
		parts := strings.Split(path, ":")
		if len(parts) == 2 {
			binds = append(binds, fmt.Sprintf("%s:%s:ro", parts[0], parts[1]))
			volumes[parts[1]] = struct{}{}
		} else if len(parts) == 3 {
			binds = append(binds, fmt.Sprintf("%s:%s:%s", parts[0], parts[1], parts[2]))
			volumes[parts[1]] = struct{}{}
		}
	}

	// binds
	var mode string
	for hostPath, bind := range specs.Binds {
		if bind.ReadOnly {
			mode = "ro"
		} else {
			mode = "rw"
		}
		binds = append(binds, fmt.Sprintf("%s:%s:%s", hostPath, bind.InContainerPath, mode))
		volumes[bind.InContainerPath] = struct{}{}
	}

	// /proc/sys
	volumes["/writable-proc/sys"] = struct{}{}
	binds = append(binds, "/proc/sys:/writable-proc/sys:ro")
	return binds, volumes
}

// 跑存在labels里的exec
// 为什么要存labels呢, 因为下线容器的时候根本不知道entrypoint是啥
func runExec(client *engineapi.Client, container enginetypes.ContainerJSON, label string) error {
	cmd, ok := container.Config.Labels[label]
	if !ok || cmd == "" {
		return fmt.Errorf("No %q found in container %q", label, container.ID)
	}

	cmds := utils.MakeCommandLineArgs(cmd)
	execConfig := enginetypes.ExecConfig{User: container.Config.User, Cmd: cmds}
	resp, err := client.ContainerExecCreate(context.Background(), container.ID, execConfig)
	if err != nil {
		return err
	}

	return client.ContainerExecStart(context.Background(), resp.ID, enginetypes.ExecStartCheck{})
}
<commit_msg>should be rw<commit_after>package calcium

import (
	"fmt"
	"io"
	"io/ioutil"
	"path/filepath"
	"strings"

	engineapi "github.com/docker/engine-api/client"
	enginetypes "github.com/docker/engine-api/types"
	enginecontainer "github.com/docker/engine-api/types/container"
	enginenetwork "github.com/docker/engine-api/types/network"
	"gitlab.ricebook.net/platform/core/types"
	"gitlab.ricebook.net/platform/core/utils"
	"golang.org/x/net/context"
)

// As the name says,
// blocks until the stream is empty, until we meet EOF
func ensureReaderClosed(stream io.ReadCloser) {
	if stream == nil {
		return
	}
	io.Copy(ioutil.Discard, stream)
	stream.Close()
}

// Copies config from container
// And make a new name for container
func makeContainerConfig(info enginetypes.ContainerJSON, image string) (
	*enginecontainer.Config,
	*enginecontainer.HostConfig,
	*enginenetwork.NetworkingConfig,
	string,
	error) {

	// we use `_` to join container name
	// since we don't support `_` in entrypoint, and no `_` is in suffix,
	// the last part will be suffix and second last part will be entrypoint,
	// the rest will be the appname
	parts := strings.Split(trimLeftSlash(info.Name), "_")
	length := len(parts)
	if length < 3 {
		return nil, nil, nil, "", fmt.Errorf("Bad container name format: %q", info.Name)
	}

	entrypoint := parts[length-2]
	appname := strings.Join(parts[:length-2], "_")

	suffix := utils.RandomString(6)
	containerName := strings.Join([]string{appname, entrypoint, suffix}, "_")

	config := info.Config
	config.Image = image

	hostConfig := info.HostConfig
	networkConfig := &enginenetwork.NetworkingConfig{
		EndpointsConfig: info.NetworkSettings.Networks,
	}
	return config, hostConfig, networkConfig, containerName, nil
}

// see https://github.com/docker/docker/issues/6705
// docker's stupid problem
func trimLeftSlash(name string) string {
	return strings.TrimPrefix(name, "/")
}

// make mount paths
// app.yaml支持三种写法.
// app.yaml里可以支持mount_paths的写法, 例如
// mount_paths:
//     - "/var/www/html"
//     - "/data/eggsy"
// 这样的路径会被直接挂载到permdir下面去, 例如上面的路径就是
// /mnt/mfs/permdirs/eggsy/data/eggsy
// /mnt/mfs/permdirs/eggsy/var/www/html
// 而且这些路径是可以读写的.
//
// 或者使用volumes, 参数格式跟docker一样, 例如
// volumes:
//     - "/data/test:/test:ro"
//     - "/data/testx:/testx"
// 说明把宿主机的/data/test映射到容器里的/test, 只读, 同时
// 把宿主机的/data/tests映射到容器里的/testx, 读写.
//
// 或者使用binds, 例如
// binds:
//     "/host/path":
//         bind: "/container/path"
//         ro: true
// 说明把宿主机的/host/path映射到容器里的/container/path, 并且只读
func makeMountPaths(specs types.Specs, config types.Config) ([]string, map[string]struct{}) {
	binds := []string{}
	volumes := make(map[string]struct{})
	permDirHost := filepath.Join(config.PermDir, specs.Appname)

	// mount_paths
	for _, path := range specs.MountPaths {
		hostPath := filepath.Join(permDirHost, path)
		binds = append(binds, fmt.Sprintf("%s:%s:rw", hostPath, path))
		volumes[path] = struct{}{}
	}

	// volumes
	for _, path := range specs.Volumes {
		parts := strings.Split(path, ":")
		if len(parts) == 2 {
			binds = append(binds, fmt.Sprintf("%s:%s:ro", parts[0], parts[1]))
			volumes[parts[1]] = struct{}{}
		} else if len(parts) == 3 {
			binds = append(binds, fmt.Sprintf("%s:%s:%s", parts[0], parts[1], parts[2]))
			volumes[parts[1]] = struct{}{}
		}
	}

	// binds
	var mode string
	for hostPath, bind := range specs.Binds {
		if bind.ReadOnly {
			mode = "ro"
		} else {
			mode = "rw"
		}
		binds = append(binds, fmt.Sprintf("%s:%s:%s", hostPath, bind.InContainerPath, mode))
		volumes[bind.InContainerPath] = struct{}{}
	}

	// /proc/sys
	volumes["/writable-proc/sys"] = struct{}{}
	binds = append(binds, "/proc/sys:/writable-proc/sys:rw")
	return binds, volumes
}

// 跑存在labels里的exec
// 为什么要存labels呢, 因为下线容器的时候根本不知道entrypoint是啥
func runExec(client *engineapi.Client, container enginetypes.ContainerJSON, label string) error {
	cmd, ok := container.Config.Labels[label]
	if !ok || cmd == "" {
		return fmt.Errorf("No %q found in container %q", label, container.ID)
	}

	cmds := utils.MakeCommandLineArgs(cmd)
	execConfig := enginetypes.ExecConfig{User: container.Config.User, Cmd: cmds}
	resp, err := client.ContainerExecCreate(context.Background(), container.ID, execConfig)
	if err != nil {
		return err
	}

	return client.ContainerExecStart(context.Background(), resp.ID, enginetypes.ExecStartCheck{})
}
<|endoftext|>


<commit_before><commit_msg>Removing the tool:  asdf.  Introduced for testing, not needed anymore.<commit_after><|endoftext|>


<commit_before>// Copyright 2020 the Exposure Notifications Server authors
//
// 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 main runs the debugger service.
package main

import (
	"context"
	"fmt"
	"os/signal"
	"syscall"

	"github.com/google/exposure-notifications-server/internal/admin"
	"github.com/google/exposure-notifications-server/internal/buildinfo"
	"github.com/google/exposure-notifications-server/internal/setup"
	"github.com/google/exposure-notifications-server/pkg/logging"
	"github.com/google/exposure-notifications-server/pkg/server"
)

func main() {
	ctx, done := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)

	logger := logging.NewLoggerFromEnv().
		With("build_id", buildinfo.BuildID).
		With("build_tag", buildinfo.BuildTag)
	ctx = logging.WithLogger(ctx, logger)

	defer func() {
		done()
		if r := recover(); r != nil {
			logger.Fatalw("application panic", "panic", r)
		}
	}()

	err := realMain(ctx)
	done()

	if err != nil {
		logger.Fatal(err)
	}
	logger.Info("successful shutdown")
}

func realMain(ctx context.Context) error {
	logger := logging.FromContext(ctx)

	var config admin.Config
	env, err := setup.Setup(ctx, &config)
	if err != nil {
		return fmt.Errorf("setup.Setup: %w", err)
	}
	defer env.Close(ctx)

	adminServer, err := admin.NewServer(&config, env)
	if err != nil {
		return fmt.Errorf("export.NewServer: %w", err)
	}

	srv, err := server.New(config.Port)
	if err != nil {
		return fmt.Errorf("server.New: %w", err)
	}
	logger.Infof("listening on :%s", config.Port)

	return srv.ServeHTTPHandler(ctx, adminServer.Routes(ctx))
}
<commit_msg>Change admin server error format (#1562)<commit_after>// Copyright 2020 the Exposure Notifications Server authors
//
// 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 main runs the debugger service.
package main

import (
	"context"
	"fmt"
	"os/signal"
	"syscall"

	"github.com/google/exposure-notifications-server/internal/admin"
	"github.com/google/exposure-notifications-server/internal/buildinfo"
	"github.com/google/exposure-notifications-server/internal/setup"
	"github.com/google/exposure-notifications-server/pkg/logging"
	"github.com/google/exposure-notifications-server/pkg/server"
)

func main() {
	ctx, done := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)

	logger := logging.NewLoggerFromEnv().
		With("build_id", buildinfo.BuildID).
		With("build_tag", buildinfo.BuildTag)
	ctx = logging.WithLogger(ctx, logger)

	defer func() {
		done()
		if r := recover(); r != nil {
			logger.Fatalw("application panic", "panic", r)
		}
	}()

	err := realMain(ctx)
	done()

	if err != nil {
		logger.Fatal(err)
	}
	logger.Info("successful shutdown")
}

func realMain(ctx context.Context) error {
	logger := logging.FromContext(ctx)

	var config admin.Config
	env, err := setup.Setup(ctx, &config)
	if err != nil {
		return fmt.Errorf("setup.Setup: %w", err)
	}
	defer env.Close(ctx)

	adminServer, err := admin.NewServer(&config, env)
	if err != nil {
		return fmt.Errorf("admin.NewServer: %w", err)
	}

	srv, err := server.New(config.Port)
	if err != nil {
		return fmt.Errorf("server.New: %w", err)
	}
	logger.Infof("listening on :%s", config.Port)

	return srv.ServeHTTPHandler(ctx, adminServer.Routes(ctx))
}
<|endoftext|>


<commit_before>// bulk_data_gen generates time series data from pre-specified use cases.
//
// Supported formats:
// InfluxDB bulk load format
// ElasticSearch bulk load format
// Cassandra query format
// Mongo custom format
// OpenTSDB bulk HTTP format
// TimescaleDB SQL INSERT and binary COPY FROM
// Graphite plaintext format
//
// Supported use cases:
// Devops: scale_var is the number of hosts to simulate, with log messages
//         every 10 seconds.
package main

import (
	"bufio"
	"flag"
	"fmt"
	"github.com/influxdata/influxdb-comparisons/bulk_data_gen/common"
	"github.com/influxdata/influxdb-comparisons/bulk_data_gen/dashboard"
	"github.com/influxdata/influxdb-comparisons/bulk_data_gen/devops"
	"github.com/influxdata/influxdb-comparisons/bulk_data_gen/iot"
	"log"
	"os"
	"runtime/pprof"
	"strings"
	"time"
)

// Output data format choices:
var formatChoices = []string{"influx-bulk", "es-bulk", "es-bulk6x", "cassandra", "mongo", "opentsdb", "timescaledb-sql", "timescaledb-copyFrom", "graphite-line", "graphite-pickle"}

// Program option vars:
var (
	daemonUrl string
	dbName    string

	format           string
	useCase          string
	configFile       string
	scaleVar         int64
	scaleVarOffset   int64
	samplingInterval time.Duration

	timestampStartStr string
	timestampEndStr   string

	timestampStart time.Time
	timestampEnd   time.Time

	interleavedGenerationGroupID uint
	interleavedGenerationGroups  uint

	seed  int64
	debug int

	cpuProfile string
)

// Parse args:
func init() {
	flag.StringVar(&format, "format", formatChoices[0], fmt.Sprintf("Format to emit. (choices: %s)", strings.Join(formatChoices, ", ")))

	flag.StringVar(&useCase, "use-case", common.UseCaseChoices[0], fmt.Sprintf("Use case to model. (choices: %s)", strings.Join(common.UseCaseChoices, ", ")))
	flag.Int64Var(&scaleVar, "scale-var", 1, "Scaling variable specific to the use case.")
	flag.Int64Var(&scaleVarOffset, "scale-var-offset", 0, "Scaling variable offset specific to the use case.")
	flag.DurationVar(&samplingInterval, "sampling-interval", devops.EpochDuration, "Simulated sampling interval.")
	flag.StringVar(&configFile, "config-file", "", "Simulator config file in TOML format")

	flag.StringVar(&timestampStartStr, "timestamp-start", common.DefaultDateTimeStart, "Beginning timestamp (RFC3339).")
	flag.StringVar(&timestampEndStr, "timestamp-end", common.DefaultDateTimeEnd, "Ending timestamp (RFC3339).")

	flag.Int64Var(&seed, "seed", 0, "PRNG seed (default, or 0, uses the current timestamp).")
	flag.IntVar(&debug, "debug", 0, "Debug printing (choices: 0, 1, 2) (default 0).")

	flag.UintVar(&interleavedGenerationGroupID, "interleaved-generation-group-id", 0, "Group (0-indexed) to perform round-robin serialization within. Use this to scale up data generation to multiple processes.")
	flag.UintVar(&interleavedGenerationGroups, "interleaved-generation-groups", 1, "The number of round-robin serialization groups. Use this to scale up data generation to multiple processes.")

	flag.StringVar(&cpuProfile, "cpu-profile", "", "Write CPU profile to `file`")

	flag.Parse()

	if !(interleavedGenerationGroupID < interleavedGenerationGroups) {
		log.Fatal("incorrect interleaved groups configuration")
	}

	validFormat := false
	for _, s := range formatChoices {
		if s == format {
			validFormat = true
			break
		}
	}
	if !validFormat {
		log.Fatal("invalid format specifier")
	}

	// the default seed is the current timestamp:
	if seed == 0 {
		seed = int64(time.Now().Nanosecond())
	}
	fmt.Fprintf(os.Stderr, "using random seed %d\n", seed)

	// Parse timestamps:
	var err error
	timestampStart, err = time.Parse(time.RFC3339, timestampStartStr)
	if err != nil {
		log.Fatal(err)
	}
	timestampStart = timestampStart.UTC()
	timestampEnd, err = time.Parse(time.RFC3339, timestampEndStr)
	if err != nil {
		log.Fatal(err)
	}
	timestampEnd = timestampEnd.UTC()

	if samplingInterval <= 0 {
		log.Fatal("Invalid sampling interval")
	}
	devops.EpochDuration = samplingInterval
	log.Printf("Using sampling interval %v\n", devops.EpochDuration)
}

func main() {
	if cpuProfile != "" {
		f, err := os.Create(cpuProfile)
		if err != nil {
			log.Fatal("could not create CPU profile: ", err)
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			log.Fatal("could not start CPU profile: ", err)
		}
		defer pprof.StopCPUProfile()
	}

	common.Seed(seed)

	if configFile != "" {
		c, err := common.NewConfig(configFile)
		if err != nil {
			log.Printf("Warning: external config error: %v", err)
		} else {
			common.Config = c
			log.Printf("Using config file %s\n", configFile)
		}
	}

	out := bufio.NewWriterSize(os.Stdout, 4<<20)
	defer out.Flush()

	var sim common.Simulator

	switch useCase {
	case common.UseCaseChoices[0]:
		cfg := &devops.DevopsSimulatorConfig{
			Start: timestampStart,
			End:   timestampEnd,

			HostCount:  scaleVar,
			HostOffset: scaleVarOffset,
		}
		sim = cfg.ToSimulator()
	case common.UseCaseChoices[2]:
		cfg := &dashboard.DashboardSimulatorConfig{
			Start: timestampStart,
			End:   timestampEnd,

			HostCount:  scaleVar,
			HostOffset: scaleVarOffset,
		}
		sim = cfg.ToSimulator()
	case common.UseCaseChoices[1]:
		cfg := &iot.IotSimulatorConfig{
			Start: timestampStart,
			End:   timestampEnd,

			SmartHomeCount:  scaleVar,
			SmartHomeOffset: scaleVarOffset,
		}
		sim = cfg.ToSimulator()
	default:
		panic("unreachable")
	}

	var serializer common.Serializer
	switch format {
	case "influx-bulk":
		serializer = common.NewSerializerInflux()
	case "es-bulk":
		serializer = common.NewSerializerElastic("5x")
	case "es-bulk6x":
		serializer = common.NewSerializerElastic("6x")
	case "cassandra":
		serializer = common.NewSerializerCassandra()
	case "mongo":
		serializer = common.NewSerializerMongo()
	case "opentsdb":
		serializer = common.NewSerializerOpenTSDB()
	case "timescaledb-sql":
		serializer = common.NewSerializerTimescaleSql()
	case "timescaledb-copyFrom":
		serializer = common.NewSerializerTimescaleBin()
	case "graphite-line":
		serializer = common.NewSerializerGraphiteLine()
	default:
		panic("unreachable")
	}

	var currentInterleavedGroup uint = 0

	t := time.Now()
	point := common.MakeUsablePoint()
	n := int64(0)
	for !sim.Finished() {
		sim.Next(point)
		n++
		// in the default case this is always true
		if currentInterleavedGroup == interleavedGenerationGroupID {
			//println("printing")
			err := serializer.SerializePoint(out, point)
			if err != nil {
				log.Fatal(err)
			}

		}
		point.Reset()

		currentInterleavedGroup++
		if currentInterleavedGroup == interleavedGenerationGroups {
			currentInterleavedGroup = 0
		}
	}
	if n != sim.SeenPoints() {
		panic(fmt.Sprintf("Logic error, written %d points, generated %d points", n, sim.SeenPoints()))
	}
	serializer.SerializeSize(out, sim.SeenPoints(), sim.SeenValues())
	err := out.Flush()
	dur := time.Now().Sub(t)
	log.Printf("Written %d points, %d values, took %0f seconds\n", n, sim.SeenValues(), dur.Seconds())
	if err != nil {
		log.Fatal(err.Error())
	}
}
<commit_msg>reverting previous commit<commit_after>// bulk_data_gen generates time series data from pre-specified use cases.
//
// Supported formats:
// InfluxDB bulk load format
// ElasticSearch bulk load format
// Cassandra query format
// Mongo custom format
// OpenTSDB bulk HTTP format
// TimescaleDB SQL INSERT and binary COPY FROM
// Graphite plaintext format
//
// Supported use cases:
// Devops: scale_var is the number of hosts to simulate, with log messages
//         every 10 seconds.
package main

import (
	"bufio"
	"flag"
	"fmt"
	"github.com/influxdata/influxdb-comparisons/bulk_data_gen/common"
	"github.com/influxdata/influxdb-comparisons/bulk_data_gen/dashboard"
	"github.com/influxdata/influxdb-comparisons/bulk_data_gen/devops"
	"github.com/influxdata/influxdb-comparisons/bulk_data_gen/iot"
	"log"
	"os"
	"runtime/pprof"
	"strings"
	"time"
)

// Output data format choices:
var formatChoices = []string{"influx-bulk", "es-bulk", "es-bulk6x", "cassandra", "mongo", "opentsdb", "timescaledb-sql", "timescaledb-copyFrom", "graphite-line", "graphite-pickle"}

// Program option vars:
var (
	daemonUrl string
	dbName    string

	format           string
	useCase          string
	configFile       string
	scaleVar         int64
	scaleVarOffset   int64
	samplingInterval time.Duration

	timestampStartStr string
	timestampEndStr   string

	timestampStart time.Time
	timestampEnd   time.Time

	interleavedGenerationGroupID uint
	interleavedGenerationGroups  uint

	seed  int64
	debug int

	cpuProfile string
)

// Parse args:
func init() {
	flag.StringVar(&format, "format", formatChoices[0], fmt.Sprintf("Format to emit. (choices: %s)", strings.Join(formatChoices, ", ")))

	flag.StringVar(&useCase, "use-case", common.UseCaseChoices[0], fmt.Sprintf("Use case to model. (choices: %s)", strings.Join(common.UseCaseChoices, ", ")))
	flag.Int64Var(&scaleVar, "scale-var", 1, "Scaling variable specific to the use case.")
	flag.Int64Var(&scaleVarOffset, "scale-var-offset", 0, "Scaling variable offset specific to the use case.")
	flag.DurationVar(&samplingInterval, "sampling-interval", devops.EpochDuration, "Simulated sampling interval.")
	flag.StringVar(&configFile, "config-file", "", "Simulator config file in TOML format")

	flag.StringVar(&timestampStartStr, "timestamp-start", common.DefaultDateTimeStart, "Beginning timestamp (RFC3339).")
	flag.StringVar(&timestampEndStr, "timestamp-end", common.DefaultDateTimeEnd, "Ending timestamp (RFC3339).")

	flag.Int64Var(&seed, "seed", 0, "PRNG seed (default, or 0, uses the current timestamp).")
	flag.IntVar(&debug, "debug", 0, "Debug printing (choices: 0, 1, 2) (default 0).")

	flag.UintVar(&interleavedGenerationGroupID, "interleaved-generation-group-id", 0, "Group (0-indexed) to perform round-robin serialization within. Use this to scale up data generation to multiple processes.")
	flag.UintVar(&interleavedGenerationGroups, "interleaved-generation-groups", 1, "The number of round-robin serialization groups. Use this to scale up data generation to multiple processes.")

	flag.StringVar(&cpuProfile, "cpu-profile", "", "Write CPU profile to `file`")

	flag.Parse()

	if !(interleavedGenerationGroupID < interleavedGenerationGroups) {
		log.Fatal("incorrect interleaved groups configuration")
	}

	validFormat := false
	for _, s := range formatChoices {
		if s == format {
			validFormat = true
			break
		}
	}
	if !validFormat {
		log.Fatal("invalid format specifier")
	}

	// the default seed is the current timestamp:
	if seed == 0 {
		seed = int64(time.Now().Nanosecond())
	}
	fmt.Fprintf(os.Stderr, "using random seed %d\n", seed)

	// Parse timestamps:
	var err error
	timestampStart, err = time.Parse(time.RFC3339, timestampStartStr)
	if err != nil {
		log.Fatal(err)
	}
	timestampStart = timestampStart.UTC()
	timestampEnd, err = time.Parse(time.RFC3339, timestampEndStr)
	if err != nil {
		log.Fatal(err)
	}
	timestampEnd = timestampEnd.UTC()

	if samplingInterval <= 0 {
		log.Fatal("Invalid sampling interval")
	}
	devops.EpochDuration = samplingInterval
	log.Printf("Using sampling interval %v\n", devops.EpochDuration)
}

func main() {
	if cpuProfile != "" {
		f, err := os.Create(cpuProfile)
		if err != nil {
			log.Fatal("could not create CPU profile: ", err)
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			log.Fatal("could not start CPU profile: ", err)
		}
		defer pprof.StopCPUProfile()
	}

	common.Seed(seed)

	if configFile != "" {
		c, err := common.NewConfig(configFile)
		if err != nil {
			log.Fatalf("external config error: %v", err)
		}
		common.Config = c
		log.Printf("Using config file %s\n", configFile)
	}

	out := bufio.NewWriterSize(os.Stdout, 4<<20)
	defer out.Flush()

	var sim common.Simulator

	switch useCase {
	case common.UseCaseChoices[0]:
		cfg := &devops.DevopsSimulatorConfig{
			Start: timestampStart,
			End:   timestampEnd,

			HostCount:  scaleVar,
			HostOffset: scaleVarOffset,
		}
		sim = cfg.ToSimulator()
	case common.UseCaseChoices[2]:
		cfg := &dashboard.DashboardSimulatorConfig{
			Start: timestampStart,
			End:   timestampEnd,

			HostCount:  scaleVar,
			HostOffset: scaleVarOffset,
		}
		sim = cfg.ToSimulator()
	case common.UseCaseChoices[1]:
		cfg := &iot.IotSimulatorConfig{
			Start: timestampStart,
			End:   timestampEnd,

			SmartHomeCount:  scaleVar,
			SmartHomeOffset: scaleVarOffset,
		}
		sim = cfg.ToSimulator()
	default:
		panic("unreachable")
	}

	var serializer common.Serializer
	switch format {
	case "influx-bulk":
		serializer = common.NewSerializerInflux()
	case "es-bulk":
		serializer = common.NewSerializerElastic("5x")
	case "es-bulk6x":
		serializer = common.NewSerializerElastic("6x")
	case "cassandra":
		serializer = common.NewSerializerCassandra()
	case "mongo":
		serializer = common.NewSerializerMongo()
	case "opentsdb":
		serializer = common.NewSerializerOpenTSDB()
	case "timescaledb-sql":
		serializer = common.NewSerializerTimescaleSql()
	case "timescaledb-copyFrom":
		serializer = common.NewSerializerTimescaleBin()
	case "graphite-line":
		serializer = common.NewSerializerGraphiteLine()
	default:
		panic("unreachable")
	}

	var currentInterleavedGroup uint = 0

	t := time.Now()
	point := common.MakeUsablePoint()
	n := int64(0)
	for !sim.Finished() {
		sim.Next(point)
		n++
		// in the default case this is always true
		if currentInterleavedGroup == interleavedGenerationGroupID {
			//println("printing")
			err := serializer.SerializePoint(out, point)
			if err != nil {
				log.Fatal(err)
			}

		}
		point.Reset()

		currentInterleavedGroup++
		if currentInterleavedGroup == interleavedGenerationGroups {
			currentInterleavedGroup = 0
		}
	}
	if n != sim.SeenPoints() {
		panic(fmt.Sprintf("Logic error, written %d points, generated %d points", n, sim.SeenPoints()))
	}
	serializer.SerializeSize(out, sim.SeenPoints(), sim.SeenValues())
	err := out.Flush()
	dur := time.Now().Sub(t)
	log.Printf("Written %d points, %d values, took %0f seconds\n", n, sim.SeenValues(), dur.Seconds())
	if err != nil {
		log.Fatal(err.Error())
	}
}
<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
	"os"
	"strings"

	"github.com/mitchellh/go-homedir"
)

var pluginCompletions []func(string) []string

// This file is meant (for now) for autocompletion in command mode, not
// while coding. This helps micro autocomplete commands and then filenames
// for example with `vsplit filename`.

// FileComplete autocompletes filenames
func FileComplete(input string) (string, []string) {
	var sep string = string(os.PathSeparator)
	dirs := strings.Split(input, sep)
	var files []os.FileInfo
	var err error
	if len(dirs) > 1 {
		home, _ := homedir.Dir()

		directories := strings.Join(dirs[:len(dirs)-1], sep)
		if strings.HasPrefix(directories, "~") {
			directories = strings.Replace(directories, "~", home, 1)
		}
		files, err = ioutil.ReadDir(directories)
	} else {
		files, err = ioutil.ReadDir(".")
	}
	var suggestions []string
	if err != nil {
		return "", suggestions
	}
	for _, f := range files {
		name := f.Name()
		if f.IsDir() {
			name += sep
		}
		if strings.HasPrefix(name, dirs[len(dirs)-1]) {
			suggestions = append(suggestions, name)
		}
	}

	var chosen string
	if len(suggestions) == 1 {
		if len(dirs) > 1 {
			chosen = strings.Join(dirs[:len(dirs)-1], sep) + sep + suggestions[0]
		} else {
			chosen = suggestions[0]
		}
	} else {
		if len(dirs) > 1 {
			chosen = strings.Join(dirs[:len(dirs)-1], sep) + sep
		}
	}

	return chosen, suggestions
}

// CommandComplete autocompletes commands
func CommandComplete(input string) (string, []string) {
	var suggestions []string
	for cmd := range commands {
		if strings.HasPrefix(cmd, input) {
			suggestions = append(suggestions, cmd)
		}
	}

	var chosen string
	if len(suggestions) == 1 {
		chosen = suggestions[0]
	}
	return chosen, suggestions
}

// HelpComplete autocompletes help topics
func HelpComplete(input string) (string, []string) {
	var suggestions []string

	for _, topic := range helpFiles {
		if strings.HasPrefix(topic, input) {

			suggestions = append(suggestions, topic)
		}
	}

	var chosen string
	if len(suggestions) == 1 {
		chosen = suggestions[0]
	}
	return chosen, suggestions
}

func contains(s []string, e string) bool {
	for _, a := range s {
		if a == e {
			return true
		}
	}
	return false
}

// OptionComplete autocompletes options
func OptionComplete(input string) (string, []string) {
	var suggestions []string
	localSettings := DefaultLocalSettings()
	for option := range globalSettings {
		if strings.HasPrefix(option, input) {
			suggestions = append(suggestions, option)
		}
	}
	for option := range localSettings {
		if strings.HasPrefix(option, input) && !contains(suggestions, option) {
			suggestions = append(suggestions, option)
		}
	}

	var chosen string
	if len(suggestions) == 1 {
		chosen = suggestions[0]
	}
	return chosen, suggestions
}

// MakeCompletion registeres a function from a plugin for autocomplete commands
func MakeCompletion(function string) Completion {
	pluginCompletions = append(pluginCompletions, LuaFunctionComplete(function))
	return Completion(-len(pluginCompletions))
}

// PluginComplete autocompletes from plugin function
func PluginComplete(complete Completion, input string) (chosen string, suggestions []string) {
	idx := int(-complete) - 1

	if len(pluginCompletions) <= idx {
		return "", nil
	}
	suggestions = pluginCompletions[idx](input)

	if len(suggestions) == 1 {
		chosen = suggestions[0]
	}
	return
}
<commit_msg>include trailing path delimiter<commit_after>package main

import (
	"io/ioutil"
	"os"
	"strings"

	"github.com/mitchellh/go-homedir"
)

var pluginCompletions []func(string) []string

// This file is meant (for now) for autocompletion in command mode, not
// while coding. This helps micro autocomplete commands and then filenames
// for example with `vsplit filename`.

// FileComplete autocompletes filenames
func FileComplete(input string) (string, []string) {
	var sep string = string(os.PathSeparator)
	dirs := strings.Split(input, sep)

	var files []os.FileInfo
	var err error
	if len(dirs) > 1 {
		home, _ := homedir.Dir()

		directories := strings.Join(dirs[:len(dirs)-1], sep) + sep

		if strings.HasPrefix(directories, "~") {
			directories = strings.Replace(directories, "~", home, 1)
		}
		files, err = ioutil.ReadDir(directories)
	} else {
		files, err = ioutil.ReadDir(".")
	}

	var suggestions []string
	if err != nil {
		return "", suggestions
	}
	for _, f := range files {
		name := f.Name()
		if f.IsDir() {
			name += sep
		}
		if strings.HasPrefix(name, dirs[len(dirs)-1]) {
			suggestions = append(suggestions, name)
		}
	}

	var chosen string
	if len(suggestions) == 1 {
		if len(dirs) > 1 {
			chosen = strings.Join(dirs[:len(dirs)-1], sep) + sep + suggestions[0]
		} else {
			chosen = suggestions[0]
		}
	} else {
		if len(dirs) > 1 {
			chosen = strings.Join(dirs[:len(dirs)-1], sep) + sep
		}
	}

	return chosen, suggestions
}

// CommandComplete autocompletes commands
func CommandComplete(input string) (string, []string) {
	var suggestions []string
	for cmd := range commands {
		if strings.HasPrefix(cmd, input) {
			suggestions = append(suggestions, cmd)
		}
	}

	var chosen string
	if len(suggestions) == 1 {
		chosen = suggestions[0]
	}
	return chosen, suggestions
}

// HelpComplete autocompletes help topics
func HelpComplete(input string) (string, []string) {
	var suggestions []string

	for _, topic := range helpFiles {
		if strings.HasPrefix(topic, input) {

			suggestions = append(suggestions, topic)
		}
	}

	var chosen string
	if len(suggestions) == 1 {
		chosen = suggestions[0]
	}
	return chosen, suggestions
}

func contains(s []string, e string) bool {
	for _, a := range s {
		if a == e {
			return true
		}
	}
	return false
}

// OptionComplete autocompletes options
func OptionComplete(input string) (string, []string) {
	var suggestions []string
	localSettings := DefaultLocalSettings()
	for option := range globalSettings {
		if strings.HasPrefix(option, input) {
			suggestions = append(suggestions, option)
		}
	}
	for option := range localSettings {
		if strings.HasPrefix(option, input) && !contains(suggestions, option) {
			suggestions = append(suggestions, option)
		}
	}

	var chosen string
	if len(suggestions) == 1 {
		chosen = suggestions[0]
	}
	return chosen, suggestions
}

// MakeCompletion registeres a function from a plugin for autocomplete commands
func MakeCompletion(function string) Completion {
	pluginCompletions = append(pluginCompletions, LuaFunctionComplete(function))
	return Completion(-len(pluginCompletions))
}

// PluginComplete autocompletes from plugin function
func PluginComplete(complete Completion, input string) (chosen string, suggestions []string) {
	idx := int(-complete) - 1

	if len(pluginCompletions) <= idx {
		return "", nil
	}
	suggestions = pluginCompletions[idx](input)

	if len(suggestions) == 1 {
		chosen = suggestions[0]
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"os"

	"github.com/coreos/coreos-cloudinit/datasource"
	"github.com/coreos/coreos-cloudinit/initialize"
	"github.com/coreos/coreos-cloudinit/system"
)

const version = "0.5.2"

func main() {
	var printVersion bool
	flag.BoolVar(&printVersion, "version", false, "Print the version and exit")

	var ignoreFailure bool
	flag.BoolVar(&ignoreFailure, "ignore-failure", false, "Exits with 0 status in the event of malformed input from user-data")

	var file string
	flag.StringVar(&file, "from-file", "", "Read user-data from provided file")

	var url string
	flag.StringVar(&url, "from-url", "", "Download user-data from provided url")

	var useProcCmdline bool
	flag.BoolVar(&useProcCmdline, "from-proc-cmdline", false, fmt.Sprintf("Parse %s for '%s=<url>', using the cloud-config served by an HTTP GET to <url>", datasource.ProcCmdlineLocation, datasource.ProcCmdlineCloudConfigFlag))

	var workspace string
	flag.StringVar(&workspace, "workspace", "/var/lib/coreos-cloudinit", "Base directory coreos-cloudinit should use to store data")

	var sshKeyName string
	flag.StringVar(&sshKeyName, "ssh-key-name", initialize.DefaultSSHKeyName, "Add SSH keys to the system with the given name")

	flag.Parse()

	if printVersion == true {
		fmt.Printf("coreos-cloudinit version %s\n", version)
		os.Exit(0)
	}

	if file != "" && url != "" && !useProcCmdline {
		fmt.Println("Provide one of --from-file, --from-url or --from-proc-cmdline")
		os.Exit(1)
	}

	var ds datasource.Datasource
	if file != "" {
		ds = datasource.NewLocalFile(file)
	} else if url != "" {
		ds = datasource.NewMetadataService(url)
	} else if useProcCmdline {
		ds = datasource.NewProcCmdline()
	} else {
		fmt.Println("Provide one of --from-file, --from-url or --from-proc-cmdline")
		os.Exit(1)
	}

	log.Printf("Fetching user-data from datasource of type %q", ds.Type())
	userdataBytes, err := ds.Fetch()
	if err != nil {
		log.Printf("Failed fetching user-data from datasource: %v", err)
		if ignoreFailure {
			os.Exit(0)
		} else {
			os.Exit(1)
		}
	}

	if len(userdataBytes) == 0 {
		log.Printf("No user data to handle, exiting.")
		os.Exit(0)
	}

	env := initialize.NewEnvironment("/", workspace)

	userdata := string(userdataBytes)
	userdata = env.Apply(userdata)

	parsed, err := initialize.ParseUserData(userdata)
	if err != nil {
		log.Printf("Failed parsing user-data: %v", err)
		if ignoreFailure {
			os.Exit(0)
		} else {
			os.Exit(1)
		}
	}

	err = initialize.PrepWorkspace(env.Workspace())
	if err != nil {
		log.Fatalf("Failed preparing workspace: %v", err)
	}

	switch t := parsed.(type) {
	case initialize.CloudConfig:
		err = initialize.Apply(t, env)
	case system.Script:
		var path string
		path, err = initialize.PersistScriptInWorkspace(t, env.Workspace())
		if err == nil {
			var name string
			name, err = system.ExecuteScript(path)
			initialize.PersistUnitNameInWorkspace(name, workspace)
		}
	}

	if err != nil {
		log.Fatalf("Failed resolving user-data: %v", err)
	}
}
<commit_msg>chore(release): Bump version to v0.5.2+git<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"os"

	"github.com/coreos/coreos-cloudinit/datasource"
	"github.com/coreos/coreos-cloudinit/initialize"
	"github.com/coreos/coreos-cloudinit/system"
)

const version = "0.5.2+git"

func main() {
	var printVersion bool
	flag.BoolVar(&printVersion, "version", false, "Print the version and exit")

	var ignoreFailure bool
	flag.BoolVar(&ignoreFailure, "ignore-failure", false, "Exits with 0 status in the event of malformed input from user-data")

	var file string
	flag.StringVar(&file, "from-file", "", "Read user-data from provided file")

	var url string
	flag.StringVar(&url, "from-url", "", "Download user-data from provided url")

	var useProcCmdline bool
	flag.BoolVar(&useProcCmdline, "from-proc-cmdline", false, fmt.Sprintf("Parse %s for '%s=<url>', using the cloud-config served by an HTTP GET to <url>", datasource.ProcCmdlineLocation, datasource.ProcCmdlineCloudConfigFlag))

	var workspace string
	flag.StringVar(&workspace, "workspace", "/var/lib/coreos-cloudinit", "Base directory coreos-cloudinit should use to store data")

	var sshKeyName string
	flag.StringVar(&sshKeyName, "ssh-key-name", initialize.DefaultSSHKeyName, "Add SSH keys to the system with the given name")

	flag.Parse()

	if printVersion == true {
		fmt.Printf("coreos-cloudinit version %s\n", version)
		os.Exit(0)
	}

	if file != "" && url != "" && !useProcCmdline {
		fmt.Println("Provide one of --from-file, --from-url or --from-proc-cmdline")
		os.Exit(1)
	}

	var ds datasource.Datasource
	if file != "" {
		ds = datasource.NewLocalFile(file)
	} else if url != "" {
		ds = datasource.NewMetadataService(url)
	} else if useProcCmdline {
		ds = datasource.NewProcCmdline()
	} else {
		fmt.Println("Provide one of --from-file, --from-url or --from-proc-cmdline")
		os.Exit(1)
	}

	log.Printf("Fetching user-data from datasource of type %q", ds.Type())
	userdataBytes, err := ds.Fetch()
	if err != nil {
		log.Printf("Failed fetching user-data from datasource: %v", err)
		if ignoreFailure {
			os.Exit(0)
		} else {
			os.Exit(1)
		}
	}

	if len(userdataBytes) == 0 {
		log.Printf("No user data to handle, exiting.")
		os.Exit(0)
	}

	env := initialize.NewEnvironment("/", workspace)

	userdata := string(userdataBytes)
	userdata = env.Apply(userdata)

	parsed, err := initialize.ParseUserData(userdata)
	if err != nil {
		log.Printf("Failed parsing user-data: %v", err)
		if ignoreFailure {
			os.Exit(0)
		} else {
			os.Exit(1)
		}
	}

	err = initialize.PrepWorkspace(env.Workspace())
	if err != nil {
		log.Fatalf("Failed preparing workspace: %v", err)
	}

	switch t := parsed.(type) {
	case initialize.CloudConfig:
		err = initialize.Apply(t, env)
	case system.Script:
		var path string
		path, err = initialize.PersistScriptInWorkspace(t, env.Workspace())
		if err == nil {
			var name string
			name, err = system.ExecuteScript(path)
			initialize.PersistUnitNameInWorkspace(name, workspace)
		}
	}

	if err != nil {
		log.Fatalf("Failed resolving user-data: %v", err)
	}
}
<|endoftext|>