<commit_before>package main

// vim: set noexpandtab :

import (
	"database/sql"
	"errors"
	"fmt"
	//"io"
	"github.com/gokyle/cryptobox/box"
	"github.com/golang/glog"
	"time"
)

const (
	challengeExpire int = 2 //minutes
	tokenExpire     int = 5 //minutes
)

var initSQL string = `
SET timezone TO UTC;
BEGIN;
CREATE TABLE IF NOT EXISTS users (
	id       bigserial NOT NULL UNIQUE PRIMARY KEY,
	name     varchar(255),
	pubKey   bytea,
	creation timestamp,
	admin    boolean
);
CREATE TABLE IF NOT EXISTS groups (
	id        bigserial NOT NULL UNIQUE PRIMARY KEY,
	name      varchar(255),
	admin     bigint[],
	member    bigint[],
	userGroup boolean
);
CREATE TABLE IF NOT EXISTS secret (
	id     bigserial NOT NULL UNIQUE PRIMARY KEY,
	uri    text,
	note   text,
	box    bytea,
	signer bigint
);
CREATE TABLE IF NOT EXISTS secretMap (
	id       bigserial NOT NULL UNIQUE PRIMARY KEY,
	isGroup  boolean,
	ownerId  bigint,
	secretId bigint
);
DROP TABLE IF EXISTS session;
CREATE TABLE session (
	id        bigint,
	token     bytea,
	challenge bytea,
	expire    timestamp
);
COMMIT; `

func dbInit(db *sql.DB) error {
	var err error
	_, err = db.Exec(initSQL)
	if err != nil {
		glog.Errorln("error intializing database")
		return err
	}
	q = map[string]*sql.Stmt{}
	prepare := map[string]string{
		"getUser":          "SELECT name, pubkey, creation, admin FROM users WHERE id = $1",
		"getUserIdByName":  "SELECT id FROM users WHERE name = $1",
		"getGroup":         "SELECT name, admin, member, userGroup FROM groups WHERE id = $1",
		"getGroupIdByName": "SELECT id FROM groups WHERE name = $1",
		"checkToken":       "SELECT EXISTS (SELECT 1 FROM session WHERE id = $1 AND token = $2 AND expire > current_timestamp);",
		"checkChallenge":   "SELECT EXISTS (SELECT 1 FROM session WHERE id = $1 AND challenge = $2 AND expire > current_timestamp);",
	}
	for name, query := range prepare {
		x, err := db.Prepare(query)
		if err != nil {
			glog.Errorf("error preparing: %s\t%s\n", name, query)
			return err
		}
		q[name] = x
	}
	return nil
}

// User bits
type User struct {
	Id       int64
	Name     string
	PubKey   box.PublicKey
	Creation time.Time
	Admin    bool
}

func (u *User) String() string {
	return fmt.Sprintf("%x\t%s", u.Id, u.Name)
}

func (u *User) Save() error {
	var err error
	var r sql.Result
	var exists bool
	var id int64
	db.QueryRow("SELECT EXISTS (SELECT 1 FROM users WHERE id = $1);", u.Id).Scan(&exists)
	if !exists {
		r, err = db.Exec("INSERT INTO users (id, name, pubKey, creation, admin) VALUES ($1, $2, $3, $4, $5);",
			u.Id, u.Name, []byte(u.PubKey), u.Creation, u.Admin)
	} else {
		return errors.New("user already exists")
	}
	if err != nil {
		return err
	}
	id, _ = r.LastInsertId()
	if glog.V(1) {
		glog.Infof("User.Save(): Last Insert Id: %x\n", id)
	}
	return nil
}

// Group
type Group struct {
	Id        int64
	Name      string
	Admin     []int64
	Member    []int64
	UserGroup bool
}

func (g *Group) String() string {
	return fmt.Sprintf("%x\t%s", g.Id, g.Name)
}

func (g *Group) Save() error {
	var err error
	var r sql.Result
	var exists bool
	var id int64
	db.QueryRow("SELECT EXISTS (SELECT 1 FROM groups WHERE id = $1);", g.Id).Scan(&exists)
	if !exists {
		r, err = db.Exec("INSERT INTO groups (id, name, admin[1], member[1], userGroup) VALUES ($1, $2, $3, $4, $5);",
			g.Id, g.Name, g.Admin[0], g.Member[0], g.UserGroup)
	} else {
		return errors.New("group already exists")
	}
	if err != nil {
		return err
	}
	id, _ = r.LastInsertId()
	if glog.V(1) {
		glog.Infof("Group.Save(): Last Insert Id: %x\n", id)
	}
	return nil
}

type SecretMap struct {
	Id       int64
	IsGroup  bool
	OwnerId  int64
	SecretId int64
}

func (sm *SecretMap) String() string {
	return fmt.Sprintf("%x\t%x\t%x", sm.Id, sm.OwnerId, sm.SecretId)
}

// Secret Bits
type Secret struct {
	Id   int64
	Uri  string
	Note string
	//Meta   map[string]string
	Box    []byte
	Signer int64
}

func (s *Secret) String() string {
	return fmt.Sprintf("%x\t%s", s.Id, s.Uri)
}

// Token bits
type Token struct {
	Id   int64
	Data []byte
}

type Challenge struct {
	Id   int64
	Data []byte
}

//
// These are all methods for manipulating the datastore w/r/t the previously
// defined objects.
//

// This function is always used to actually return a User record.
// Other methods look up the id from the relevant bit of information and then call this.
func GetUser(id int64) (*User, error) {
	var err error
	var name string
	var key []byte
	var creation time.Time
	var admin bool
	err = q["getUser"].QueryRow(id).Scan(&name, &key, &creation, &admin)
	if err != nil {
		return nil, err
	}
	k := box.PublicKey(key)
	//return &User{id, name, der.(*box.PublicKey), creation, admin}, nil
	return &User{id, name, k, creation, admin}, nil
}

func GetUserByName(name string) (*User, error) {
	var id int64
	var err error
	err = q["getUserIdByName"].QueryRow(name).Scan(&id)
	if err != nil {
		return nil, err
	}
	return GetUser(id)
}

// The following group methods are the same as their User counterparts
func GetGroup(id int64) (*Group, error) {
	var name string
	var admin []int64
	var member []int64
	var primary bool
	err := q["getGroup"].QueryRow(id).Scan(&name, &admin, &member, &primary)
	if err != nil {
		return nil, err
	}
	return &Group{id, name, admin, member, primary}, nil
}

func GetGroupByName(name string) (*Group, error) {
	var id int64
	var err error
	err = q["getGroupIdByName"].QueryRow(name).Scan(&id)
	if err != nil {
		return nil, err
	}
	return GetGroup(id)
}

func SaveToken(id int64, token []byte) error {
	expire := time.Now().UTC().Add(time.Duration(tokenExpire) * time.Minute)
	tx, err := db.Begin()
	if err != nil {
		tx.Rollback()
		return err
	}
	_, err = tx.Exec("DELETE FROM session WHERE id = $1;", id)
	if err != nil {
		tx.Rollback()
		return err
	}
	_, err = tx.Exec("INSERT INTO session (id, token, expire) VALUES ($1, $2, $3);", id, token, expire)
	if err != nil {
		tx.Rollback()
		return err
	}
	err = tx.Commit()
	if err != nil {
		tx.Rollback()
		return err
	}
	return nil
}

func CheckToken(id int64, token []byte) bool {
	var ok bool
	err := db.QueryRow("SELECT EXISTS (SELECT 1 FROM session WHERE id = $1 AND token = $2);", id, token).Scan(&ok)
	if err != nil {
		return false
	}
	return ok
}

func SaveChallenge(id int64, challenge []byte) error {
	expire := time.Now().UTC().Add(time.Duration(challengeExpire) * time.Minute)
	//_, err = tx.Exec("DELETE FROM session WHERE id = $1 AND token IS NULL;", id)
	//if err != nil {
	//	tx.Rollback()
	//	return err
	//}
	_, err := db.Exec("INSERT INTO session (id, challenge, expire) VALUES ($1, $2, $3);", id, challenge, expire)
	if err != nil {
		return err
	}
	return nil
}

func CheckChallenge(id int64, challenge []byte) bool {
	var ok bool
	err := db.QueryRow("SELECT EXISTS (SELECT 1 FROM session WHERE id = $1 AND challenge = $2);", id, challenge).Scan(&ok)
	if err != nil {
		return false
	}
	return ok
}

//
// Higher order functions
//

func NewUser(name string, admin bool) (error, *User, box.PrivateKey) {
	var err error
	var id int64
	id, err = nextId("user")
	priv, pub, ok := box.GenerateKey()
	if !ok {
		glog.Errorln(err)
		return err, nil, nil
	}
	// cast so that we don't need to need to include cryptobox everywhere
	newUser := &User{Id: id, Name: name, PubKey: pub, Creation: time.Now().UTC(), Admin: admin}
	err = newUser.Save()
	if err != nil {
		glog.Errorln(err)
		return err, nil, nil
	}
	return nil, newUser, priv
}

func NewGroup(name string, primaryUserId int64, isUserGroup bool) (error, *Group) {
	var err error
	var nid int64
	nid, err = nextId("group")
	if err != nil {
		return err, nil
	}
	var newGroup *Group = &Group{Id: nid, Name: name, Admin: []int64{primaryUserId},
		Member: []int64{primaryUserId}, UserGroup: isUserGroup}
	err = newGroup.Save()
	if err != nil {
		glog.Errorln(err)
		return err, nil
	}
	return nil, newGroup
}
<commit_msg>re-structure tables<commit_after>package main

// vim: set noexpandtab :

import (
	"database/sql"
	"errors"
	"fmt"
	//"io"
	"github.com/gokyle/cryptobox/box"
	"github.com/golang/glog"
	"time"
)

const (
	challengeExpire int = 2 //minutes
	tokenExpire     int = 5 //minutes
)

var initSQL string = `
SET timezone TO UTC;
BEGIN;
CREATE TABLE IF NOT EXISTS users (
	id       bigserial NOT NULL UNIQUE PRIMARY KEY,
	name     varchar(255),
	pubKey   bytea,
	creation timestamp,
	admin    boolean
);
CREATE TABLE IF NOT EXISTS groups (
	id        bigserial NOT NULL UNIQUE PRIMARY KEY,
	name      varchar(255),
	-- User group means there should only be one member.
	userGroup boolean
);
CREATE TABLE IF NOT EXISTS users_groups (
	uid bigint REFERENCES users (id),
	gid bigint REFERENCES groups (id),
	-- Is user (uid) and admin of group (gid)
	admin bool
);
CREATE TABLE IF NOT EXISTS secrets (
	id     bigserial NOT NULL UNIQUE PRIMARY KEY,
	rev    bigint NOT NULL,
	uri    text,
	note   text,
	box    bytea,
	signer bigint
);
CREATE TABLE IF NOT EXISTS groups_secrets (
	-- All secrets are owned by a group. All users have a group they're the only member of.
	-- This may make some lookups more expensive, but makes lookups logically simpler
	owner  bigint REFERENCES groups (id),
	secret bigint REFERENCES secrets (id)
);
DROP TABLE IF EXISTS session;
CREATE TABLE session (
	id        bigint,
	token     bytea,
	challenge bytea,
	expire    timestamp
);
CREATE OR REPLACE VIEW users_secrets AS
	SELECT u.id AS uid, s.id AS id, s.rev, s.uri, s.note, s.box, s.signer FROM
		users as u
		LEFT JOIN users_groups AS ug ON (u.id = ug.uid)
		LEFT JOIN groups AS g ON (g.id = ug.gid)
		LEFT JOIN groups_secrets AS gs ON (g.id = gs.owner)
		LEFT JOIN secrets AS s ON (s.id = gs.secret)
WHERE s.box IS NOT NULL;
COMMIT; `

func dbInit(db *sql.DB) error {
	var err error
	_, err = db.Exec(initSQL)
	if err != nil {
		glog.Errorln("error intializing database")
		return err
	}
	q = map[string]*sql.Stmt{}
	prepare := map[string]string{
		"getUser":          "SELECT name, pubkey, creation, admin FROM users WHERE id = $1",
		"getUserIdByName":  "SELECT id FROM users WHERE name = $1",
		"getGroup":         "SELECT name, userGroup FROM groups WHERE id = $1",
		"getGroupIdByName": "SELECT id FROM groups WHERE name = $1",
		"checkToken":       "SELECT EXISTS (SELECT 1 FROM session WHERE id = $1 AND token = $2 AND expire > current_timestamp);",
		"checkChallenge":   "SELECT EXISTS (SELECT 1 FROM session WHERE id = $1 AND challenge = $2 AND expire > current_timestamp);",
	}
	for name, query := range prepare {
		x, err := db.Prepare(query)
		if err != nil {
			glog.Errorf("error preparing: %s\t%s\n", name, query)
			return err
		}
		q[name] = x
	}
	return nil
}

// User bits
type User struct {
	Id       int64
	Name     string
	PubKey   box.PublicKey
	Creation time.Time
	Admin    bool
}

func (u *User) String() string {
	return fmt.Sprintf("%x\t%s", u.Id, u.Name)
}

func (u *User) Save() error {
	var err error
	var id int64

	if u.Id != 0 {
		return errors.New("user Id already populated")
	}

	tx, err := db.Begin()
	var gid int64
	if err != nil {
		return errors.New("unable to begin transaction to create user")
	}
	err = tx.QueryRow("INSERT INTO groups (id, name, userGroup) VALUES (default, $1, TRUE) RETURNING id;", fmt.Sprintf("%s/group", u.Name)).
		Scan(&gid)
	if err != nil {
		tx.Rollback()
		return errors.New(fmt.Sprintf("unable to create group for user: %v", err))
	}
	err = tx.QueryRow("INSERT INTO users (id, name, pubKey, creation, admin) VALUES (default, $1, $2, $3, $4) RETURNING id;",
		u.Name, []byte(u.PubKey), u.Creation, u.Admin).Scan(&id)
	if err != nil {
		tx.Rollback()
		return errors.New(fmt.Sprintf("unable to insert user: %v", err))
	}
	u.Id = id
	_, err = tx.Exec("INSERT INTO users_groups (uid, gid, admin) VALUES ($1, $2, TRUE);", id, gid)
	if err != nil {
		tx.Rollback()
		return errors.New(fmt.Sprintf("unable to associate user and group: %v", err))
	}
	err = tx.Commit()
	if err != nil {
		tx.Rollback()
		return errors.New(fmt.Sprintf("unable to commit transaction to create user: %v", err))
	}

	if glog.V(1) {
		glog.Infof("User.Save(): Last Insert Id: %x\n", id)
	}
	return nil
}

// Group
type Group struct {
	Id        int64
	Name      string
	Admin     []int64
	Member    []int64
	UserGroup bool
}

func (g *Group) String() string {
	return fmt.Sprintf("%x\t%s", g.Id, g.Name)
}

func (g *Group) Save() error {
	var err error
	var id int64

	if g.Id != 0 {
		return errors.New("group Id already populated")
	}

	err = db.QueryRow("INSERT INTO groups (id, name, userGroup) VALUES (default, $1, $2) RETURNING id;", g.Name, g.UserGroup).
		Scan(&id)
	if err != nil {
		return err
	}

	g.Id = id

	if glog.V(1) {
		glog.Infof("Group.Save(): Last Insert Id: %x\n", id)
	}
	return nil
}

type SecretMap struct {
	Id       int64
	IsGroup  bool
	OwnerId  int64
	SecretId int64
}

func (sm *SecretMap) String() string {
	return fmt.Sprintf("%x\t%x\t%x", sm.Id, sm.OwnerId, sm.SecretId)
}

// Secret Bits
type Secret struct {
	Id   int64
	Uri  string
	Note string
	//Meta   map[string]string
	Box    []byte
	Signer int64
}

func (s *Secret) String() string {
	return fmt.Sprintf("%x\t%s", s.Id, s.Uri)
}

// Token bits
type Token struct {
	Id   int64
	Data []byte
}

type Challenge struct {
	Id   int64
	Data []byte
}

//
// These are all methods for manipulating the datastore w/r/t the previously
// defined objects.
//

// This function is always used to actually return a User record.
// Other methods look up the id from the relevant bit of information and then call this.
func GetUser(id int64) (*User, error) {
	var err error
	var name string
	var key []byte
	var creation time.Time
	var admin bool
	err = q["getUser"].QueryRow(id).Scan(&name, &key, &creation, &admin)
	if err != nil {
		return nil, err
	}
	k := box.PublicKey(key)
	//return &User{id, name, der.(*box.PublicKey), creation, admin}, nil
	return &User{id, name, k, creation, admin}, nil
}

func GetUserByName(name string) (*User, error) {
	var id int64
	var err error
	err = q["getUserIdByName"].QueryRow(name).Scan(&id)
	if err != nil {
		return nil, err
	}
	return GetUser(id)
}

// The following group methods are the same as their User counterparts
func GetGroup(id int64) (*Group, error) {
	var name string
	var primary bool
	err := q["getGroup"].QueryRow(id).Scan(&name, &primary)
	if err != nil {
		return nil, err
	}
	return &Group{Id: id, Name: name, UserGroup: primary}, nil
}

func GetGroupByName(name string) (*Group, error) {
	var id int64
	var err error
	err = q["getGroupIdByName"].QueryRow(name).Scan(&id)
	if err != nil {
		return nil, err
	}
	return GetGroup(id)
}

func SaveToken(id int64, token []byte) error {
	expire := time.Now().UTC().Add(time.Duration(tokenExpire) * time.Minute)
	tx, err := db.Begin()
	if err != nil {
		tx.Rollback()
		return err
	}
	_, err = tx.Exec("DELETE FROM session WHERE id = $1;", id)
	if err != nil {
		tx.Rollback()
		return err
	}
	_, err = tx.Exec("INSERT INTO session (id, token, expire) VALUES ($1, $2, $3);", id, token, expire)
	if err != nil {
		tx.Rollback()
		return err
	}
	err = tx.Commit()
	if err != nil {
		tx.Rollback()
		return err
	}
	return nil
}

func CheckToken(id int64, token []byte) bool {
	var ok bool
	err := db.QueryRow("SELECT EXISTS (SELECT 1 FROM session WHERE id = $1 AND token = $2);", id, token).Scan(&ok)
	if err != nil {
		return false
	}
	return ok
}

func SaveChallenge(id int64, challenge []byte) error {
	expire := time.Now().UTC().Add(time.Duration(challengeExpire) * time.Minute)
	//_, err = tx.Exec("DELETE FROM session WHERE id = $1 AND token IS NULL;", id)
	//if err != nil {
	//	tx.Rollback()
	//	return err
	//}
	_, err := db.Exec("INSERT INTO session (id, challenge, expire) VALUES ($1, $2, $3);", id, challenge, expire)
	if err != nil {
		return err
	}
	return nil
}

func CheckChallenge(id int64, challenge []byte) bool {
	var ok bool
	err := db.QueryRow("SELECT EXISTS (SELECT 1 FROM session WHERE id = $1 AND challenge = $2);", id, challenge).Scan(&ok)
	if err != nil {
		return false
	}
	return ok
}

//
// Higher order functions
//

func NewUser(name string, admin bool) (error, *User, box.PrivateKey) {
	var err error
	priv, pub, ok := box.GenerateKey()
	if !ok {
		glog.Errorln(err)
		return err, nil, nil
	}
	// cast so that we don't need to need to include cryptobox everywhere
	newUser := &User{Name: name, PubKey: pub, Creation: time.Now().UTC(), Admin: admin}
	err = newUser.Save()
	if err != nil {
		glog.Errorln(err)
		return err, nil, nil
	}
	return nil, newUser, priv
}

func NewGroup(name string, primaryUserId int64, isUserGroup bool) (error, *Group) {
	var err error
	if err != nil {
		return err, nil
	}
	var newGroup *Group = &Group{Name: name, UserGroup: isUserGroup}
	err = newGroup.Save()
	if err != nil {
		glog.Errorln(err)
		return err, nil
	}
	_, err = db.Exec("INSERT INTO users_groups (uid, gid, admin) VALUES ($1, $2, TRUE);", primaryUserId, newGroup.Id)
	if err != nil {
		glog.Errorln(err)
		return err, nil
	}
	return nil, newGroup
}
<|endoftext|>


<commit_before>package docopt

import (
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"unicode"
)

func errKey(key string) error {
	return fmt.Errorf("no such key: %s", key)
}
func errType(key string) error {
	return fmt.Errorf("key: %s failed type conversion", key)
}
func errStrconv(key string, convErr error) error {
	return fmt.Errorf("key: %s failed type conversion: %s", key, convErr)
}

// Opts is a map of command line options to their values, with some convenience
// methods for value type conversion (bool, float64, int, string). For example,
// to get an option value as an int:
//
// 	opts, _ := docopt.ParseDoc("Usage: sleep <seconds>")
// 	secs, _ := opts.Int("<seconds>")
//
// Additionally, Opts.Bind allows you easily populate a struct's fields with the
// values of each option value. See below for examples.
//
// Lastly, you can still treat Opts as a regular map, and do any type checking
// and conversion that you want to yourself. For example:
//
// 	if s, ok := opts["<binary>"].(string); ok {
// 		if val, err := strconv.ParseUint(s, 2, 64); err != nil { ... }
// 	}
//
// Note that any non-boolean option / flag will have a string value in the
// underlying map.
type Opts map[string]interface{}

func (o Opts) String(key string) (s string, err error) {
	v, ok := o[key]
	if !ok {
		err = errKey(key)
		return
	}
	s, ok = v.(string)
	if !ok {
		err = errType(key)
	}
	return
}

func (o Opts) Bool(key string) (b bool, err error) {
	v, ok := o[key]
	if !ok {
		err = errKey(key)
		return
	}
	b, ok = v.(bool)
	if !ok {
		err = errType(key)
	}
	return
}

func (o Opts) Int(key string) (i int, err error) {
	s, err := o.String(key)
	if err != nil {
		return
	}
	i, err = strconv.Atoi(s)
	if err != nil {
		err = errStrconv(key, err)
	}
	return
}

func (o Opts) Float64(key string) (f float64, err error) {
	s, err := o.String(key)
	if err != nil {
		return
	}
	f, err = strconv.ParseFloat(s, 64)
	if err != nil {
		err = errStrconv(key, err)
	}
	return
}

// Bind populates the fields of a given struct with matching option values.
// Each key in Opts will be mapped to an exported field of the struct pointed
// to by `v`, as follows:
//
// 	Field int `docopt:"-"`        // Field is ignored by Bind
// 	Field int `docopt:"--help"`   // Field mapped from key "--help"
// 	Field int `docopt:"-h"`       // Field mapped from key "-h"
// 	Field int                     // Field mapped from key "--field"
// 	F int                         // F mapped from key "-f"
//
func (o Opts) Bind(v interface{}) error {
	value := reflect.ValueOf(v)
	if value.Kind() != reflect.Ptr {
		return newError("'v' argument is not pointer to struct type")
	}
	for value.Kind() == reflect.Ptr {
		value = value.Elem()
	}
	if value.Kind() != reflect.Struct {
		return newError("'v' argument is not pointer to struct type")
	}
	typ := value.Type()
	indexMap := make(map[string]int)
	for i := 0; i < typ.NumField(); i++ {
		field := typ.Field(i)
		if isUnexportedField(field) || field.Anonymous {
			continue
		}
		tag := field.Tag.Get("docopt")
		if tag == "" {
			key := strings.ToLower(field.Name)
			if len(field.Name) == 1 {
				key = "-" + key
			} else {
				key = "--" + key
			}
			indexMap[key] = i
			continue
		}
		for _, t := range strings.Split(tag, ",") {
			indexMap[t] = i
		}
	}
	for k, v := range o {
		i, ok := indexMap[k]
		if !ok {
			return newError("mapping of '%s' is not found in given struct, or an unexported field", k)
		}
		field := value.Field(i)
		if field.Interface() != reflect.Zero(field.Type()).Interface() {
			continue
		}
		val := reflect.ValueOf(v)
		if !val.Type().AssignableTo(field.Type()) {
			return newError("value of '%s' is not assignable to '%s' field", k, value.Type().Field(i).Name)
		}
		if !field.CanSet() {
			return newError("'%s' field cannot be set", value.Type().Field(i).Name)
		}
		field.Set(val)
	}
	return nil
}

// isUnexportedField returns whether the field is unexported.
// isUnexportedField is to avoid the bug in versions older than Go1.3.
// See following links:
//   https://code.google.com/p/go/issues/detail?id=7247
//   http://golang.org/ref/spec#Exported_identifiers
func isUnexportedField(field reflect.StructField) bool {
	return !(field.PkgPath == "" && unicode.IsUpper(rune(field.Name[0])))
}
<commit_msg>Allow bound struct to omit --help and --version keys<commit_after>package docopt

import (
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"unicode"
)

func errKey(key string) error {
	return fmt.Errorf("no such key: %s", key)
}
func errType(key string) error {
	return fmt.Errorf("key: %s failed type conversion", key)
}
func errStrconv(key string, convErr error) error {
	return fmt.Errorf("key: %s failed type conversion: %s", key, convErr)
}

// Opts is a map of command line options to their values, with some convenience
// methods for value type conversion (bool, float64, int, string). For example,
// to get an option value as an int:
//
// 	opts, _ := docopt.ParseDoc("Usage: sleep <seconds>")
// 	secs, _ := opts.Int("<seconds>")
//
// Additionally, Opts.Bind allows you easily populate a struct's fields with the
// values of each option value. See below for examples.
//
// Lastly, you can still treat Opts as a regular map, and do any type checking
// and conversion that you want to yourself. For example:
//
// 	if s, ok := opts["<binary>"].(string); ok {
// 		if val, err := strconv.ParseUint(s, 2, 64); err != nil { ... }
// 	}
//
// Note that any non-boolean option / flag will have a string value in the
// underlying map.
type Opts map[string]interface{}

func (o Opts) String(key string) (s string, err error) {
	v, ok := o[key]
	if !ok {
		err = errKey(key)
		return
	}
	s, ok = v.(string)
	if !ok {
		err = errType(key)
	}
	return
}

func (o Opts) Bool(key string) (b bool, err error) {
	v, ok := o[key]
	if !ok {
		err = errKey(key)
		return
	}
	b, ok = v.(bool)
	if !ok {
		err = errType(key)
	}
	return
}

func (o Opts) Int(key string) (i int, err error) {
	s, err := o.String(key)
	if err != nil {
		return
	}
	i, err = strconv.Atoi(s)
	if err != nil {
		err = errStrconv(key, err)
	}
	return
}

func (o Opts) Float64(key string) (f float64, err error) {
	s, err := o.String(key)
	if err != nil {
		return
	}
	f, err = strconv.ParseFloat(s, 64)
	if err != nil {
		err = errStrconv(key, err)
	}
	return
}

// Bind populates the fields of a given struct with matching option values.
// Each key in Opts will be mapped to an exported field of the struct pointed
// to by `v`, as follows:
//
// 	Field int `docopt:"-"`        // Field is ignored by Bind
// 	Field int `docopt:"--help"`   // Field mapped from key "--help"
// 	Field int `docopt:"-h"`       // Field mapped from key "-h"
// 	Field int                     // Field mapped from key "--field"
// 	F int                         // F mapped from key "-f"
//
func (o Opts) Bind(v interface{}) error {
	value := reflect.ValueOf(v)
	if value.Kind() != reflect.Ptr {
		return newError("'v' argument is not pointer to struct type")
	}
	for value.Kind() == reflect.Ptr {
		value = value.Elem()
	}
	if value.Kind() != reflect.Struct {
		return newError("'v' argument is not pointer to struct type")
	}
	typ := value.Type()
	indexMap := make(map[string]int)
	for i := 0; i < typ.NumField(); i++ {
		field := typ.Field(i)
		if isUnexportedField(field) || field.Anonymous {
			continue
		}
		tag := field.Tag.Get("docopt")
		if tag == "" {
			key := strings.ToLower(field.Name)
			if len(field.Name) == 1 {
				key = "-" + key
			} else {
				key = "--" + key
			}
			indexMap[key] = i
			continue
		}
		for _, t := range strings.Split(tag, ",") {
			indexMap[t] = i
		}
	}
	for k, v := range o {
		i, ok := indexMap[k]
		if !ok {
			if k == "--help" || k == "--version" {
				continue
			}
			return newError("mapping of '%s' is not found in given struct, or an unexported field", k)
		}
		field := value.Field(i)
		if field.Interface() != reflect.Zero(field.Type()).Interface() {
			continue
		}
		val := reflect.ValueOf(v)
		if !val.Type().AssignableTo(field.Type()) {
			return newError("value of '%s' is not assignable to '%s' field", k, value.Type().Field(i).Name)
		}
		if !field.CanSet() {
			return newError("'%s' field cannot be set", value.Type().Field(i).Name)
		}
		field.Set(val)
	}
	return nil
}

// isUnexportedField returns whether the field is unexported.
// isUnexportedField is to avoid the bug in versions older than Go1.3.
// See following links:
//   https://code.google.com/p/go/issues/detail?id=7247
//   http://golang.org/ref/spec#Exported_identifiers
func isUnexportedField(field reflect.StructField) bool {
	return !(field.PkgPath == "" && unicode.IsUpper(rune(field.Name[0])))
}
<|endoftext|>


<commit_before>// Copyright © 2015, 2016 Authors (see AUTHORS file)
//
// License for use of this code is detailed in the LICENSE file

package opus

/*
// Link opus using pkg-config.
#cgo pkg-config: --static opus
#include <opus/opus.h>

// Access the preprocessor from CGO
const int CONST_APPLICATION_VOIP = OPUS_APPLICATION_VOIP;
const int CONST_APPLICATION_AUDIO = OPUS_APPLICATION_AUDIO;
const int CONST_APPLICATION_RESTRICTED_LOWDELAY = OPUS_APPLICATION_RESTRICTED_LOWDELAY;
*/
import "C"

type Application int

// These variables should be constants, but for interoperability with CGO
// they're var. Don't change them, though!
var (
	// Optimize encoding for VOIP
	APPLICATION_VOIP = Application(C.CONST_APPLICATION_VOIP)
	// Optimize encoding for non-voice signals like music
	APPLICATION_AUDIO = Application(C.CONST_APPLICATION_AUDIO)
	// Optimize encoding for low latency applications
	APPLICATION_RESTRICTED_LOWDELAY = Application(C.CONST_APPLICATION_RESTRICTED_LOWDELAY)
)

const (
	xMAX_BITRATE       = 48000
	xMAX_FRAME_SIZE_MS = 60
	xMAX_FRAME_SIZE    = xMAX_BITRATE * xMAX_FRAME_SIZE_MS / 1000
	// Maximum size of an encoded frame. I actually have no idea, but this
	// looks like it's big enough.
	maxEncodedFrameSize = 10000
)

func Version() string {
	return C.GoString(C.opus_get_version_string())
}
<commit_msg>Don't put --static in pkg-config<commit_after>// Copyright © 2015, 2016 Authors (see AUTHORS file)
//
// License for use of this code is detailed in the LICENSE file

package opus

/*
// Link opus using pkg-config.
#cgo pkg-config: opus
#include <opus/opus.h>

// Access the preprocessor from CGO
const int CONST_APPLICATION_VOIP = OPUS_APPLICATION_VOIP;
const int CONST_APPLICATION_AUDIO = OPUS_APPLICATION_AUDIO;
const int CONST_APPLICATION_RESTRICTED_LOWDELAY = OPUS_APPLICATION_RESTRICTED_LOWDELAY;
*/
import "C"

type Application int

// These variables should be constants, but for interoperability with CGO
// they're var. Don't change them, though!
var (
	// Optimize encoding for VOIP
	APPLICATION_VOIP = Application(C.CONST_APPLICATION_VOIP)
	// Optimize encoding for non-voice signals like music
	APPLICATION_AUDIO = Application(C.CONST_APPLICATION_AUDIO)
	// Optimize encoding for low latency applications
	APPLICATION_RESTRICTED_LOWDELAY = Application(C.CONST_APPLICATION_RESTRICTED_LOWDELAY)
)

const (
	xMAX_BITRATE       = 48000
	xMAX_FRAME_SIZE_MS = 60
	xMAX_FRAME_SIZE    = xMAX_BITRATE * xMAX_FRAME_SIZE_MS / 1000
	// Maximum size of an encoded frame. I actually have no idea, but this
	// looks like it's big enough.
	maxEncodedFrameSize = 10000
)

func Version() string {
	return C.GoString(C.opus_get_version_string())
}
<|endoftext|>


<commit_before>package simra

import (
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"testing"

	"github.com/pankona/gomo-simra/simra/database"
)

type mock struct {
	m map[string]interface{}
}

func (db *mock) Open(dirpath string) error {
	if db.m == nil {
		db.m = map[string]interface{}{}
	}
	return nil
}

func (db *mock) Close() {
	db.m = nil
}

func (db *mock) Put(key string, value interface{}) {
	if db.m == nil {
		log.Fatal("database is not opened yet")
		return
	}
	db.m[key] = value
}

func (db *mock) Get(key string) interface{} {
	if db.m == nil {
		log.Fatal("database is not opened yet")
		return nil
	}
	return db.m[key]
}

func TestMock(t *testing.T) {
	db := OpenDB(&mock{}, filepath.Join("."))
	if db == nil {
		t.Error("failed to open database")
	}
	defer db.Close()

	db.Put("key1", "value1")
	value1 := db.Get("key1")
	if value1 != "value1" {
		t.Error("failed to open database")
	}
}

func TestBolt(t *testing.T) {
	tmpdir, err := ioutil.TempDir(filepath.Join("."), "simradb_")
	if err != nil {
		t.Fatalf("failed to create temporary working directory. err: %s", err.Error())
	}
	defer func() {
		_ = os.RemoveAll(tmpdir)
	}()

	db := OpenDB(&database.Boltdb{}, filepath.Join(tmpdir))
	if db == nil {
		t.Error("failed to open database")
	}
	defer db.Close()

	db.Put("key1", "value1")
	fetched := db.Get("key1")
	if bytes, ok := fetched.([]uint8); ok {
		v := string(bytes)
		if v != "value1" {
			t.Errorf("unexpected fetched value. expected: %s, fetched: %s", "value1", v)
		}
	}
}
<commit_msg>[#125] update database_test.go to check errors<commit_after>package simra

import (
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"testing"

	"github.com/pankona/gomo-simra/simra/database"
)

type mock struct {
	m map[string]interface{}
}

func (db *mock) Open(dirpath string) error {
	if db.m == nil {
		db.m = map[string]interface{}{}
	}
	return nil
}

func (db *mock) Close() {
	db.m = nil
}

func (db *mock) Put(key string, value interface{}) {
	if db.m == nil {
		log.Fatal("database is not opened yet")
		return
	}
	db.m[key] = value
}

func (db *mock) Get(key string) interface{} {
	if db.m == nil {
		log.Fatal("database is not opened yet")
		return nil
	}
	return db.m[key]
}

func TestMock(t *testing.T) {
	db, err := OpenDB(&mock{}, filepath.Join("."))
	if err != nil {
		t.Error("failed to open database. err: %s", err.Error())
		return
	}
	defer db.Close()

	db.Put("key1", "value1")
	value1 := db.Get("key1")
	if value1 != "value1" {
		t.Error("failed to open database")
	}
}

func TestBolt(t *testing.T) {
	tmpdir, err := ioutil.TempDir(filepath.Join("."), "simradb_")
	if err != nil {
		t.Fatalf("failed to create temporary working directory. err: %s", err.Error())
	}
	defer func() {
		_ = os.RemoveAll(tmpdir)
	}()

	db, err := OpenDB(&database.Boltdb{}, filepath.Join(tmpdir))
	if err != nil {
		t.Error("failed to open database. err: %s", err.Error())
		return
	}
	defer db.Close()

	db.Put("key1", "value1")
	fetched := db.Get("key1")
	if bytes, ok := fetched.([]uint8); ok {
		v := string(bytes)
		if v != "value1" {
			t.Errorf("unexpected fetched value. expected: %s, fetched: %s", "value1", v)
		}
	}
}
<|endoftext|>


<commit_before>package rxgo

import (
	"testing"

	"github.com/reactivex/rxgo/handlers"
	"github.com/stretchr/testify/assert"
)

func TestCreateNewSingleObserverWithConstructor(t *testing.T) {
	var got int64
	single := Just(1, 2, 3).Count().Subscribe(
		handlers.NextFunc(func(i interface{}) {
			got = i.(int64)
		}))

	v, err := single.Block()

	assert.Nil(t, err)
	assert.Equal(t, int64(3), got)
	assert.Equal(t, int64(3), v)
	assert.True(t, single.IsDisposed())
}

func TestCreateNewSingleObserverFromSingleObserver(t *testing.T) {
	var got int64
	singleObserver := CheckSingleEventHandler(handlers.NextFunc(func(i interface{}) {
		got = i.(int64)
	}))

	single := Just(1, 2, 3).Count().Subscribe(singleObserver)
	v, err := single.Block()

	assert.Nil(t, err)
	assert.Equal(t, int64(3), got)
	assert.Equal(t, int64(3), v)
	assert.True(t, single.IsDisposed())
}
<commit_msg>Testing coverage<commit_after>package rxgo

import (
	"testing"

	"github.com/reactivex/rxgo/errors"

	"github.com/reactivex/rxgo/handlers"
	"github.com/stretchr/testify/assert"
)

func TestCreateNewSingleObserverWithConstructor(t *testing.T) {
	var got int64
	single := Just(1, 2, 3).Count().Subscribe(
		handlers.NextFunc(func(i interface{}) {
			got = i.(int64)
		}))

	v, err := single.Block()

	assert.Nil(t, err)
	assert.Equal(t, int64(3), got)
	assert.Equal(t, int64(3), v)
	assert.True(t, single.IsDisposed())
}

func TestCreateNewSingleObserverFromSingleObserver(t *testing.T) {
	var got int64
	singleObserver := CheckSingleEventHandler(handlers.NextFunc(func(i interface{}) {
		got = i.(int64)
	}))

	single := Just(1, 2, 3).Count().Subscribe(singleObserver)
	v, err := single.Block()

	assert.Nil(t, err)
	assert.Equal(t, int64(3), got)
	assert.Equal(t, int64(3), v)
	assert.True(t, single.IsDisposed())
}

func TestSingleObserverHandle(t *testing.T) {
	var got int64
	singleObserver := CheckSingleEventHandler(handlers.NextFunc(func(i interface{}) {
		got = 10
	}))
	singleObserver.Handle("")
	assert.Equal(t, int64(10), got)
}

func TestSingleObserverHandleWithError(t *testing.T) {
	var got int64
	singleObserver := CheckSingleEventHandler(handlers.ErrFunc(func(err error) {
		got = 10
	}))
	singleObserver.Handle(errors.New(0))
	assert.Equal(t, int64(10), got)
}
<|endoftext|>


<commit_before>package version

// Version of Functions
var Version = "0.3.224"
<commit_msg>fnserver: 0.3.225 release [skip ci]<commit_after>package version

// Version of Functions
var Version = "0.3.225"
<|endoftext|>


<commit_before>package version

// Version of Functions
var Version = "0.3.554"
<commit_msg>fnserver: 0.3.555 release [skip ci]<commit_after>package version

// Version of Functions
var Version = "0.3.555"
<|endoftext|>


<commit_before>package apihandlers

import (
	"fmt"

	log "github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
)

type subnetData struct {
	subnetID   string
	cidrBlock  string
	subnetWrk  string
	subnetName string
}

func getSubnets() []*ec2.Subnet {
	ec2client := getec2client()
	params := &ec2.DescribeSubnetsInput{
		DryRun: aws.Bool(false),
	}
	log.Debug("Calling DescribeSubnets")
	resp, err := ec2client.DescribeSubnets(params)
	if err != nil {
		fmt.Println(err.Error())
	}
	return resp.Subnets
}

func parseSubnetsData(subnets []*ec2.Subnet) (response []subnetData) {

	resp := make([]subnetData, 0, 4)
	log.Debug("Parsing Subnets Data")
	if subnets != nil {
		for _, subnet := range subnets {
			parsedData := new(subnetData)
			parsedData.subnetID = *subnet.SubnetId
			parsedData.cidrBlock = *subnet.CidrBlock
			tags := subnet.Tags
			if tags != nil {
				parsedData.subnetWrk = getTagValue(tags, "WRK")
				parsedData.subnetName = getTagValue(tags, "Name")
			}
			resp = append(resp, *parsedData)
		}
	}
	return resp
}

func getSubnetIdByTag(tagname, tagvalue string) (subnetId string) {
	ec2client := getec2client()

}
<commit_msg>Adding filters for getsubnetbytag func<commit_after>package apihandlers

import (
	"fmt"

	log "github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
)

type subnetData struct {
	subnetID   string
	cidrBlock  string
	subnetWrk  string
	subnetName string
}

func getSubnets() []*ec2.Subnet {
	ec2client := getec2client()
	params := &ec2.DescribeSubnetsInput{
		DryRun: aws.Bool(false),
	}
	log.Debug("Calling DescribeSubnets")
	resp, err := ec2client.DescribeSubnets(params)
	if err != nil {
		fmt.Println(err.Error())
	}
	return resp.Subnets
}

func parseSubnetsData(subnets []*ec2.Subnet) (response []subnetData) {

	resp := make([]subnetData, 0, 4)
	log.Debug("Parsing Subnets Data")
	if subnets != nil {
		for _, subnet := range subnets {
			parsedData := new(subnetData)
			parsedData.subnetID = *subnet.SubnetId
			parsedData.cidrBlock = *subnet.CidrBlock
			tags := subnet.Tags
			if tags != nil {
				parsedData.subnetWrk = getTagValue(tags, "WRK")
				parsedData.subnetName = getTagValue(tags, "Name")
			}
			resp = append(resp, *parsedData)
		}
	}
	return resp
}

func getSubnetIdByTag(tagname, tagvalue string) (subnetId string) {
	ec2client := getec2client()
	params := &ec2.DescribeSubnetsInput{
		Filters = []*ec2.Filter{
			{ // Required
				Name: aws.String("tag:"+tagname),
				Values: []*string{
					aws.String(tagvalue), // Required
					// More values...
				},
			},
			// More values...
		},
	}

	Filters:
  ec2client.DescribeSubnets(input *ec2.DescribeSubnetsInput)
}
<|endoftext|>


<commit_before>// Copyright 2014 The Serviced 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 elastic

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/control-center/serviced/datastore"
	"github.com/control-center/serviced/logging"
	"github.com/elastic/go-elasticsearch/v7"
	"github.com/elastic/go-elasticsearch/v7/esapi"
	"github.com/zenoss/logri"
	"io/ioutil"
)

type elasticConnection struct {
	index  string
	client *elasticsearch.Client
}

var (
	plog        = logging.PackageLogger() // the standard package logger
	traceLogger *logri.Logger             // a 'trace' logger for an additional level of debug messages
)

func init() {
	traceLoggerName := plog.Name + ".trace"
	traceLogger = logri.GetLogger(traceLoggerName)
}

func BuildID(id string, key string) string {
	return fmt.Sprintf("%s-%s", id, key)
}

func (ec *elasticConnection) Put(key datastore.Key, msg datastore.JSONMessage) error {
	logger := plog.WithFields(log.Fields{
		"kind": key.Kind(),
		"id":   key.ID(),
	})
	logger.Debug("Put")
	traceLogger.WithField("payload", string(msg.Bytes())).Debug("Put")

	//Because of document type depreciation in ES 7.x we combine id and key
	args := []func(*esapi.IndexRequest){
		ec.client.Index.WithDocumentID(BuildID(key.ID(), key.Kind())),
		ec.client.Index.WithRefresh("true"),
	}

	if (msg.Version()["seqNo"] + msg.Version()["primaryTerm"]) > 0 {
		args = append(args, ec.client.Index.WithIfPrimaryTerm(msg.Version()["primaryTerm"]),
			ec.client.Index.WithIfSeqNo(msg.Version()["seqNo"]))
	}

	res, err := ec.client.Index(ec.index, bytes.NewReader(msg.Bytes()), args...)

	defer res.Body.Close()
	if err != nil {
		logger.WithError(err).Error("Put failed")
		//TODO handle new type of ES error
		//if eserr, iseserror := err.(api.ESError); iseserror && eserr.Code == 409 {
		// Conflict
		//return fmt.Errorf("Your changes conflict with those made by another user. Please reload and try your changes again.")
		//}
		return err
	}

	traceLogger.WithField("response", res).Debug("Put")
	if res.IsError() {
		logger.Info("error while put")
		return fmt.Errorf("non OK response: %v", res)
	}
	return nil
}

func (ec *elasticConnection) Get(key datastore.Key) (datastore.JSONMessage, error) {
	logger := plog.WithFields(log.Fields{
		"kind": key.Kind(),
		"id":   key.ID(),
	})
	logger.Debug("Get")
	traceLogger.Debug("Get")

	//Because of document type depreciation in ES 7.x we combine id and key
	response, err := ec.elasticGet(BuildID(key.ID(), key.Kind()))
	if err != nil {
		logger.WithError(err).Error("Get failed")
		return nil, err
	}
	if !response.Found {
		logger.Debug("Entity not found")
		return nil, datastore.ErrNoSuchEntity{Key: key}
	}
	bytes := response.Source
	msg := datastore.NewJSONMessage(bytes, map[string]int{
		"version":     response.Version,
		"primaryTerm": response.PrimaryTerm,
		"seqNo":       response.SeqNo})
	return msg, nil
}

func (ec *elasticConnection) Delete(key datastore.Key) error {
	logger := plog.WithFields(log.Fields{
		"kind": key.Kind(),
		"id":   key.ID(),
	})
	logger.Debug("Delete")

	//Because of document type depreciation in ES 7.x we combine id and key
	res, err := ec.client.Delete(ec.index,
		BuildID(key.ID(), key.Kind()),
		ec.client.Delete.WithRefresh("true"))

	traceLogger.WithField("response", res).Debug("Delete")
	if err != nil {
		return err
	}
	return nil
}

func (ec *elasticConnection) Query(query esapi.SearchRequest) ([]datastore.JSONMessage, error) {

	resp, err := query.Do(context.Background(), ec.client)
	if err != nil {
		err = fmt.Errorf("error executing query %v", err)
		plog.WithError(err).Error("error executing query")
		return nil, err
	}
	defer resp.Body.Close()

	if resp.IsError() {
		var e map[string]interface{}
		if err := json.NewDecoder(resp.Body).Decode(&e); err != nil {
			plog.WithError(err).Errorf("Error parsing the response body: %s", err)
		} else {
			plog.WithError(err).Errorf("[%s] %s: %s",
				resp.Status(),
				e["error"].(map[string]interface{})["type"],
				e["error"].(map[string]interface{})["reason"],
			)
		}
		return nil, err
	}

	return toJSONMessages(resp)
}

// convert search result of json host to dao.Host array
func toJSONMessages(result *esapi.Response) ([]datastore.JSONMessage, error) {
	// Deserialize the response into a map.
	var r elasticSearchResults
	if err := json.NewDecoder(result.Body).Decode(&r); err != nil {
		plog.WithError(err).Errorf("Error parsing the response body: %s", err)
		return nil, err
	}

	var total = r.Hits.Total.Value
	var msgs = make([]datastore.JSONMessage, total)

	logger := plog.WithField("total", total)
	if total > 0 {
		// Note that while it's possible for queries to span types, the vast majority of CC use cases (all?)
		//   only query a single type at a time, so it's sufficient to get the type of the first hit
		logger.WithField("id-type", r.Hits.Hits[0].Id)
	}
	logger.Debug("Query finished")

	for i, hit := range r.Hits.Hits {
		var data bytes.Buffer
		if err := json.NewEncoder(&data).Encode(hit.Source); err != nil {
			plog.WithError(err).Errorf("Error encoding query: %s", err)
			return nil, err
		}

		msg := datastore.NewJSONMessage(data.Bytes(), map[string]int{
			"version":     hit.Version,
			"primaryTerm": hit.PrimaryTerm,
			"seqNo":       hit.SeqNo,
		})
		msgs[i] = msg
	}
	return msgs, nil
}

//Modified from elastigo to use custom response type
func (ec *elasticConnection) elasticGet(id string) (elasticResponse, error) {
	var (
		retval elasticResponse
		res    *esapi.Response
		err    error
	)

	res, err = ec.client.Get(ec.index, id)

	defer res.Body.Close()

	if err != nil {
		plog.Errorf("Error getting response: %s", err)
		return retval, err
	}

	if res.IsError() {
		var e map[string]interface{}
		if err := json.NewDecoder(res.Body).Decode(&e); err != nil {
			plog.Errorf("Error parsing the response body: %s", err)
			return retval, err
		} else {
			// Print the response status and error information.
			plog.Errorf("%s ID:%s", res.Status(), id)
			return retval, nil
		}
	}
	//body, err := api.DoCommand("GET", url, nil)

	// marshall into json
	var body []byte
	body, err = ioutil.ReadAll(res.Body)

	if err != nil {
		return retval, err
	}

	jsonErr := json.Unmarshal(body, &retval)
	if jsonErr != nil {
		return retval, jsonErr
	}
	return retval, err
}

//Modified from elastigo BaseResponse to accept document source into a json.RawMessage
type elasticResponse struct {
	Ok          bool            `json:"ok"`
	Index       string          `json:"_index,omitempty"`
	Type        string          `json:"_type,omitempty"`
	ID          string          `json:"_id,omitempty"`
	Source      json.RawMessage `json:"_source,omitempty"` // depends on the schema you've defined
	Version     int             `json:"_version,omitempty"`
	Found       bool            `json:"found,omitempty"`
	Exists      bool            `json:"exists,omitempty"`
	PrimaryTerm int             `json:"_primary_term,omitempty"`
	SeqNo       int             `json:"_seq_no,omitempty"`
}

type elasticSearchResults struct {
	Took     int            `json:"took,omitempty"`
	TimedOut bool           `json:"timed_out,omitempty"`
	Shards   map[string]int `json:"_shards,omitempty"`
	Hits     struct {
		Total struct {
			Value    int    `json:"value"`
			Relation string `json:"relation"`
		} `json:"total"`
		MaxScore float64 `json:"max_score"`
		Hits     []struct {
			Index       string                 `json:"_index"`
			Type        string                 `json:"_type"`
			Id          string                 `json:"_id"`
			Version     int                    `json:"_version,omitempty"`
			PrimaryTerm int                    `json:"_primary_term,omitempty"`
			SeqNo       int                    `json:"_seq_no,omitempty"`
			Score       float64                `json:"_score"`
			Source      map[string]interface{} `json:"_source"`
		} `json:"hits,omitempty"`
	} `json:"hits,omitempty"`
}
<commit_msg>ZEN-34210: Cloud Snapshot Causing Serviced to Crash on Master Server<commit_after>// Copyright 2014 The Serviced 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 elastic

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/control-center/serviced/datastore"
	"github.com/control-center/serviced/logging"
	"github.com/elastic/go-elasticsearch/v7"
	"github.com/elastic/go-elasticsearch/v7/esapi"
	"github.com/zenoss/logri"
	"io/ioutil"
)

type elasticConnection struct {
	index  string
	client *elasticsearch.Client
}

var (
	plog        = logging.PackageLogger() // the standard package logger
	traceLogger *logri.Logger             // a 'trace' logger for an additional level of debug messages
)

func init() {
	traceLoggerName := plog.Name + ".trace"
	traceLogger = logri.GetLogger(traceLoggerName)
}

func BuildID(id string, key string) string {
	return fmt.Sprintf("%s-%s", id, key)
}

func (ec *elasticConnection) Put(key datastore.Key, msg datastore.JSONMessage) error {
	logger := plog.WithFields(log.Fields{
		"kind": key.Kind(),
		"id":   key.ID(),
	})
	logger.Debug("Put")
	traceLogger.WithField("payload", string(msg.Bytes())).Debug("Put")

	//Because of document type depreciation in ES 7.x we combine id and key
	args := []func(*esapi.IndexRequest){
		ec.client.Index.WithDocumentID(BuildID(key.ID(), key.Kind())),
		ec.client.Index.WithRefresh("true"),
	}

	if (msg.Version()["seqNo"] + msg.Version()["primaryTerm"]) > 0 {
		args = append(args, ec.client.Index.WithIfPrimaryTerm(msg.Version()["primaryTerm"]),
			ec.client.Index.WithIfSeqNo(msg.Version()["seqNo"]))
	}

	res, err := ec.client.Index(ec.index, bytes.NewReader(msg.Bytes()), args...)

	defer res.Body.Close()
	if err != nil {
		logger.WithError(err).Error("Put failed")
		//TODO handle new type of ES error
		//if eserr, iseserror := err.(api.ESError); iseserror && eserr.Code == 409 {
		// Conflict
		//return fmt.Errorf("Your changes conflict with those made by another user. Please reload and try your changes again.")
		//}
		return err
	}

	traceLogger.WithField("response", res).Debug("Put")
	if res.IsError() {
		logger.Info("error while put")
		return fmt.Errorf("non OK response: %v", res)
	}
	return nil
}

func (ec *elasticConnection) Get(key datastore.Key) (datastore.JSONMessage, error) {
	logger := plog.WithFields(log.Fields{
		"kind": key.Kind(),
		"id":   key.ID(),
	})
	logger.Debug("Get")
	traceLogger.Debug("Get")

	//Because of document type depreciation in ES 7.x we combine id and key
	response, err := ec.elasticGet(BuildID(key.ID(), key.Kind()))
	if err != nil {
		logger.WithError(err).Error("Get failed")
		return nil, err
	}
	if !response.Found {
		logger.Debug("Entity not found")
		return nil, datastore.ErrNoSuchEntity{Key: key}
	}
	bytes := response.Source
	msg := datastore.NewJSONMessage(bytes, map[string]int{
		"version":     response.Version,
		"primaryTerm": response.PrimaryTerm,
		"seqNo":       response.SeqNo})
	return msg, nil
}

func (ec *elasticConnection) Delete(key datastore.Key) error {
	logger := plog.WithFields(log.Fields{
		"kind": key.Kind(),
		"id":   key.ID(),
	})
	logger.Debug("Delete")

	//Because of document type depreciation in ES 7.x we combine id and key
	res, err := ec.client.Delete(ec.index,
		BuildID(key.ID(), key.Kind()),
		ec.client.Delete.WithRefresh("true"))

	traceLogger.WithField("response", res).Debug("Delete")
	if err != nil {
		return err
	}
	return nil
}

func (ec *elasticConnection) Query(query esapi.SearchRequest) ([]datastore.JSONMessage, error) {

	resp, err := query.Do(context.Background(), ec.client)
	if err != nil {
		err = fmt.Errorf("error executing query %v", err)
		plog.WithError(err).Error("error executing query")
		return nil, err
	}
	defer resp.Body.Close()

	if resp.IsError() {
		var e map[string]interface{}
		if err := json.NewDecoder(resp.Body).Decode(&e); err != nil {
			plog.WithError(err).Errorf("Error parsing the response body: %s", err)
		} else {
			plog.WithError(err).Errorf("[%s] %s: %s",
				resp.Status(),
				e["error"].(map[string]interface{})["type"],
				e["error"].(map[string]interface{})["reason"],
			)
		}
		return nil, err
	}

	return toJSONMessages(resp)
}

// convert search result of json host to dao.Host array
func toJSONMessages(result *esapi.Response) ([]datastore.JSONMessage, error) {
	// Deserialize the response into a map.
	var r elasticSearchResults
	if err := json.NewDecoder(result.Body).Decode(&r); err != nil {
		plog.WithError(err).Errorf("Error parsing the response body: %s", err)
		return nil, err
	}

	var total = r.Hits.Total.Value
	var msgs = make([]datastore.JSONMessage, total)

	logger := plog.WithField("total", total)
	if total > 0 {
		// Note that while it's possible for queries to span types, the vast majority of CC use cases (all?)
		//   only query a single type at a time, so it's sufficient to get the type of the first hit
		logger.WithField("id-type", r.Hits.Hits[0].Id)
	}
	logger.Debug("Query finished")

	for i, hit := range r.Hits.Hits {
		var data bytes.Buffer
		if err := json.NewEncoder(&data).Encode(hit.Source); err != nil {
			plog.WithError(err).Errorf("Error encoding query: %s", err)
			return nil, err
		}

		msg := datastore.NewJSONMessage(data.Bytes(), map[string]int{
			"version":     hit.Version,
			"primaryTerm": hit.PrimaryTerm,
			"seqNo":       hit.SeqNo,
		})
		msgs[i] = msg
	}
	return msgs, nil
}

//Modified from elastigo to use custom response type
func (ec *elasticConnection) elasticGet(id string) (elasticResponse, error) {
	var (
		retval elasticResponse
		res    *esapi.Response
		err    error
	)

	res, err = ec.client.Get(ec.index, id)

	if err != nil {
		plog.Errorf("Error getting response: %s", err)
		return retval, err
	}

	defer res.Body.Close()

	if res.IsError() {
		var e map[string]interface{}
		if err := json.NewDecoder(res.Body).Decode(&e); err != nil {
			plog.Errorf("Error parsing the response body: %s", err)
			return retval, err
		} else {
			// Print the response status and error information.
			plog.Errorf("%s ID:%s", res.Status(), id)
			return retval, nil
		}
	}
	//body, err := api.DoCommand("GET", url, nil)

	// marshall into json
	var body []byte
	body, err = ioutil.ReadAll(res.Body)

	if err != nil {
		return retval, err
	}

	jsonErr := json.Unmarshal(body, &retval)
	if jsonErr != nil {
		return retval, jsonErr
	}
	return retval, err
}

//Modified from elastigo BaseResponse to accept document source into a json.RawMessage
type elasticResponse struct {
	Ok          bool            `json:"ok"`
	Index       string          `json:"_index,omitempty"`
	Type        string          `json:"_type,omitempty"`
	ID          string          `json:"_id,omitempty"`
	Source      json.RawMessage `json:"_source,omitempty"` // depends on the schema you've defined
	Version     int             `json:"_version,omitempty"`
	Found       bool            `json:"found,omitempty"`
	Exists      bool            `json:"exists,omitempty"`
	PrimaryTerm int             `json:"_primary_term,omitempty"`
	SeqNo       int             `json:"_seq_no,omitempty"`
}

type elasticSearchResults struct {
	Took     int            `json:"took,omitempty"`
	TimedOut bool           `json:"timed_out,omitempty"`
	Shards   map[string]int `json:"_shards,omitempty"`
	Hits     struct {
		Total struct {
			Value    int    `json:"value"`
			Relation string `json:"relation"`
		} `json:"total"`
		MaxScore float64 `json:"max_score"`
		Hits     []struct {
			Index       string                 `json:"_index"`
			Type        string                 `json:"_type"`
			Id          string                 `json:"_id"`
			Version     int                    `json:"_version,omitempty"`
			PrimaryTerm int                    `json:"_primary_term,omitempty"`
			SeqNo       int                    `json:"_seq_no,omitempty"`
			Score       float64                `json:"_score"`
			Source      map[string]interface{} `json:"_source"`
		} `json:"hits,omitempty"`
	} `json:"hits,omitempty"`
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"flag"
	"fmt"
	tu "github.com/nyaxt/otaru/testutils"
	"os"
	"time"

	"github.com/dustin/go-humanize"

	"github.com/nyaxt/otaru/btncrypt"
	"github.com/nyaxt/otaru/logger"
)

var mylog = logger.Registry().Category("otaru-btncrypt-benchmark")

var (
	flagSize = flag.String("size", "100MB", "Test target blob size")
)

func Usage() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "  %s\n", os.Args[0])
	flag.PrintDefaults()
}

func main() {
	logger.Registry().AddOutput(logger.WriterLogger{os.Stderr})
	logger.Registry().AddOutput(logger.HandleCritical(func() { os.Exit(1) }))

	flag.Usage = Usage
	flag.Parse()

	size, err := humanize.ParseBytes(*flagSize)
	if err != nil {
		logger.Criticalf(mylog, "Failed to parse size: %s", *flagSize)
	}
	logger.Infof(mylog, "Target blob size: %d", size)
	tstart := time.Now()
	buf := make([]byte, size)
	if _, err := rand.Read(buf); err != nil {
		logger.Criticalf(mylog, "Failed to generate random seq: %v", err)
	}
	logger.Infof(mylog, "Preparing target took: %v", time.Since(tstart))

	tstart = time.Now()
	envelope, err := btncrypt.Encrypt(tu.TestCipher(), buf)
	elapsed := time.Since(tstart)

	if err != nil {
		logger.Criticalf(mylog, "Failed to encrypt: %v", err)
	}
	logger.Infof(mylog, "Encrypt took %v", elapsed)
	sizeMB := (float64)(size) / (1000 * 1000)
	MBps := sizeMB / elapsed.Seconds()
	Mbps := MBps * 8
	logger.Infof(mylog, "%v MB/s %v Mbps", MBps, Mbps)

	_ = envelope
}
<commit_msg>decrypt bench<commit_after>package main

import (
	"crypto/rand"
	"flag"
	"fmt"
	tu "github.com/nyaxt/otaru/testutils"
	"os"
	"time"

	"github.com/dustin/go-humanize"

	"github.com/nyaxt/otaru/btncrypt"
	"github.com/nyaxt/otaru/logger"
)

var mylog = logger.Registry().Category("otaru-btncrypt-benchmark")

var (
	flagSize = flag.String("size", "100MB", "Test target blob size")
)

func Usage() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "  %s\n", os.Args[0])
	flag.PrintDefaults()
}

func main() {
	logger.Registry().AddOutput(logger.WriterLogger{os.Stderr})
	logger.Registry().AddOutput(logger.HandleCritical(func() { os.Exit(1) }))

	flag.Usage = Usage
	flag.Parse()

	size, err := humanize.ParseBytes(*flagSize)
	if err != nil {
		logger.Criticalf(mylog, "Failed to parse size: %s", *flagSize)
	}
	logger.Infof(mylog, "Target blob size: %d", size)
	tstart := time.Now()
	buf := make([]byte, size)
	if _, err := rand.Read(buf); err != nil {
		logger.Criticalf(mylog, "Failed to generate random seq: %v", err)
	}
	logger.Infof(mylog, "Preparing target took: %v", time.Since(tstart))

	var envelope []byte
	{
		tstart = time.Now()
		envelope, err = btncrypt.Encrypt(tu.TestCipher(), buf)
		elapsed := time.Since(tstart)

		if err != nil {
			logger.Criticalf(mylog, "Failed to encrypt: %v", err)
		}
		logger.Infof(mylog, "Encrypt took %v", elapsed)
		sizeMB := (float64)(size) / (1000 * 1000)
		MBps := sizeMB / elapsed.Seconds()
		Mbps := MBps * 8
		logger.Infof(mylog, "%v MB/s %v Mbps", MBps, Mbps)
	}

	{
		tstart = time.Now()
		_, err = btncrypt.Decrypt(tu.TestCipher(), envelope, len(buf))
		elapsed := time.Since(tstart)

		if err != nil {
			logger.Criticalf(mylog, "Failed to decrypt: %v", err)
		}
		logger.Infof(mylog, "Decrypt took %v", elapsed)
		sizeMB := (float64)(size) / (1000 * 1000)
		MBps := sizeMB / elapsed.Seconds()
		Mbps := MBps * 8
		logger.Infof(mylog, "%v MB/s %v Mbps", MBps, Mbps)
	}
}
<|endoftext|>


<commit_before>package controllers

import (
        "fmt"
		"log"
        "gopkg.in/mgo.v2"
        "gopkg.in/mgo.v2/bson"
)

type App struct {
	*revel.Controller
}

func (c App) Index() revel.Result {
	return c.Render()
}<commit_msg>no message<commit_after>package controllers

import (
        "fmt"
		"log"
        "gopkg.in/mgo.v2"
        "gopkg.in/mgo.v2/bson"
        "github.com/revel/revel"
)

type App struct {
	*revel.Controller
}

func (c App) Index() revel.Result {
	return c.Render()
}<|endoftext|>


<commit_before>package duet

import (
	"bytes"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"time"
)

// Configuration represents package configuration (shared by commands)
type Configuration struct {
	Namespace        string
	PairsFile        string
	EmailLookup      string
	CoAuthoredBy     bool
	Global           bool
	RotateAuthor     bool
	SetGitUserConfig bool
	StaleCutoff      time.Duration
}

// NewConfiguration initializes Configuration from the environment
// Returns an error if it cannot parse the staleness timeout as an integer or
// the global var as a bool
func NewConfiguration() (config *Configuration, err error) {
	config = &Configuration{
		Namespace:   getenvDefault("GIT_DUET_CONFIG_NAMESPACE", "duet.env"),
		EmailLookup: os.Getenv("GIT_DUET_EMAIL_LOOKUP_COMMAND"),
	}

	if config.PairsFile, err = getPairsFile(); err != nil {
		return nil, err
	}

	cutoff, err := strconv.Atoi(getenvDefault("GIT_DUET_SECONDS_AGO_STALE", "1200"))
	if err != nil {
		return nil, err
	}

	if config.Global, err = strconv.ParseBool(getenvDefault("GIT_DUET_GLOBAL", "0")); err != nil {
		return nil, err
	}

	if config.RotateAuthor, err = strconv.ParseBool(getenvDefault("GIT_DUET_ROTATE_AUTHOR", "0")); err != nil {
		return nil, err
	}

	if config.CoAuthoredBy, err = strconv.ParseBool(getenvDefault("GIT_DUET_CO_AUTHORED_BY", "0")); err != nil {
		return nil, err
	}

	defaultSetGitUserConfig := "0"
	if config.CoAuthoredBy {
		defaultSetGitUserConfig = "1"
	}
	if config.SetGitUserConfig, err = strconv.ParseBool(getenvDefault("GIT_DUET_SET_GIT_USER_CONFIG", defaultSetGitUserConfig)); err != nil {
		return nil, err
	}

	config.StaleCutoff = time.Duration(cutoff) * time.Second

	return config, nil
}

func getPairsFile() (value string, err error) {
	authorsFile := ".git-authors"
	defaultAuthorsFile := path.Join(os.Getenv("HOME"), authorsFile)

	if os.Getenv("GIT_DUET_AUTHORS_FILE") != "" {
		return os.Getenv("GIT_DUET_AUTHORS_FILE"), nil
	}

	gitDirectory, err := exec.Command("git", "rev-parse", "--show-toplevel").CombinedOutput()
	if err != nil {
		if !bytes.Contains(gitDirectory, []byte("Not a git repository")) {
			return "", err
		}
		return defaultAuthorsFile, nil
	}

	gitDirectoryAuthors := path.Join(strings.TrimSpace(string(gitDirectory)), authorsFile)
	if _, err := os.Stat(gitDirectoryAuthors); err == nil {
		return gitDirectoryAuthors, nil
	}

	return defaultAuthorsFile, nil
}

func getenvDefault(key, defaultValue string) (value string) {
	value = os.Getenv(key)
	if value == "" {
		value = defaultValue
	}

	return value
}
<commit_msg>Handle "not a git repository" errors<commit_after>package duet

import (
	"bytes"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"time"
)

// Configuration represents package configuration (shared by commands)
type Configuration struct {
	Namespace        string
	PairsFile        string
	EmailLookup      string
	CoAuthoredBy     bool
	Global           bool
	RotateAuthor     bool
	SetGitUserConfig bool
	StaleCutoff      time.Duration
}

// NewConfiguration initializes Configuration from the environment
// Returns an error if it cannot parse the staleness timeout as an integer or
// the global var as a bool
func NewConfiguration() (config *Configuration, err error) {
	config = &Configuration{
		Namespace:   getenvDefault("GIT_DUET_CONFIG_NAMESPACE", "duet.env"),
		EmailLookup: os.Getenv("GIT_DUET_EMAIL_LOOKUP_COMMAND"),
	}

	if config.PairsFile, err = getPairsFile(); err != nil {
		return nil, err
	}

	cutoff, err := strconv.Atoi(getenvDefault("GIT_DUET_SECONDS_AGO_STALE", "1200"))
	if err != nil {
		return nil, err
	}

	if config.Global, err = strconv.ParseBool(getenvDefault("GIT_DUET_GLOBAL", "0")); err != nil {
		return nil, err
	}

	if config.RotateAuthor, err = strconv.ParseBool(getenvDefault("GIT_DUET_ROTATE_AUTHOR", "0")); err != nil {
		return nil, err
	}

	if config.CoAuthoredBy, err = strconv.ParseBool(getenvDefault("GIT_DUET_CO_AUTHORED_BY", "0")); err != nil {
		return nil, err
	}

	defaultSetGitUserConfig := "0"
	if config.CoAuthoredBy {
		defaultSetGitUserConfig = "1"
	}
	if config.SetGitUserConfig, err = strconv.ParseBool(getenvDefault("GIT_DUET_SET_GIT_USER_CONFIG", defaultSetGitUserConfig)); err != nil {
		return nil, err
	}

	config.StaleCutoff = time.Duration(cutoff) * time.Second

	return config, nil
}

func getPairsFile() (value string, err error) {
	authorsFile := ".git-authors"
	defaultAuthorsFile := path.Join(os.Getenv("HOME"), authorsFile)

	if os.Getenv("GIT_DUET_AUTHORS_FILE") != "" {
		return os.Getenv("GIT_DUET_AUTHORS_FILE"), nil
	}

	gitDirectory, err := exec.Command("git", "rev-parse", "--show-toplevel").CombinedOutput()
	if err != nil {
		if bytes.Contains(gitDirectory, []byte("Not a git repository")) ||
			bytes.Contains(gitDirectory, []byte("not a git repository")) {
			return defaultAuthorsFile, nil
		}
		return "", err
	}

	gitDirectoryAuthors := path.Join(strings.TrimSpace(string(gitDirectory)), authorsFile)
	if _, err := os.Stat(gitDirectoryAuthors); err == nil {
		return gitDirectoryAuthors, nil
	}

	return defaultAuthorsFile, nil
}

func getenvDefault(key, defaultValue string) (value string) {
	value = os.Getenv(key)
	if value == "" {
		value = defaultValue
	}

	return value
}
<|endoftext|>


<commit_before>// +build windows

package box

import (
	"syscall"
	"unsafe"
)

var kernel32 = syscall.NewLazyDLL("kernel32")

var hConout syscall.Handle

func init() {
	var err error
	hConout, err = syscall.Open("CONOUT$", syscall.O_RDWR, 0)
	if err != nil {
		panic(err.Error())
	}
}

type coord_t struct {
	X int16
	Y int16
}

type small_rect_t struct {
	Left   int16
	Top    int16
	Right  int16
	Bottom int16
}

type console_screen_buffer_info_t struct {
	Size              coord_t
	CursorPosition    coord_t
	Attributes        uint16
	Window            small_rect_t
	MaximumWindowSize coord_t
}

var getConsoleScreenBufferInfo = kernel32.NewProc("GetConsoleScreenBufferInfo")

type console_handle_t syscall.Handle

func newHandle(handle syscall.Handle) console_handle_t {
	return console_handle_t(handle)
}

func (h console_handle_t) GetScreenBufferInfo() *console_screen_buffer_info_t {
	var csbi console_screen_buffer_info_t
	getConsoleScreenBufferInfo.Call(
		uintptr(h),
		uintptr(unsafe.Pointer(&csbi)))
	return &csbi
}

func GetScreenBufferInfo() *console_screen_buffer_info_t {
	return console_handle_t(hConout).GetScreenBufferInfo()
}

func (this *console_screen_buffer_info_t) ViewSize() (int, int) {
	return int(this.Window.Right-this.Window.Left) + 1,
		int(this.Window.Bottom-this.Window.Top) + 1
}

func (this *console_screen_buffer_info_t) CursorPos() (int, int) {
	return int(this.CursorPosition.X), int(this.CursorPosition.Y)
}

func GetLocate() (int, int) {
	return GetScreenBufferInfo().CursorPos()
}
<commit_msg>Remove unused conio_windows.go<commit_after><|endoftext|>


<commit_before>package main

import (
    "fmt"
    "golang.org/x/net/html"
    "net/http"
    "database/sql"
    _ "github.com/go-sql-driver/mysql"
)

// Helper function to pull the id attribute from a Token
func getId(t html.Token) (ok bool, id string) {
    // Iterate over all of the Token's attributes until we find an "id"
    for _, attr := range t.Attr {
        if attr.Key == "data-wallpaper-id" {
            id = attr.Val
            ok = true
        }
    }
    return
}

// Extract all http** links from a given webpage
func crawl(url string, ch chan string, chFinished chan bool) {
    resp, err := http.Get(url)

    defer func() {
        // Notify that we're done after this function
        chFinished <- true
    }()

    if err != nil {
        fmt.Println("ERROR: Failed to crawl \"" + url + "\"")
        return
    }

    b := resp.Body
    defer b.Close() // close Body when the function returns

    z := html.NewTokenizer(b)

    for {
        tt := z.Next()

        switch {
        case tt == html.ErrorToken:
            // End of the document, we're done
            return
        case tt == html.StartTagToken:
            t := z.Token()

            // Check if the token is an <img> tag
            isImg := t.Data == "figure"
            if !isImg {
                continue
            }

            // Extract the href value, if there is one
            ok, id := getId(t)
            if !ok {
                continue
            }
            ch <- id
        }
    }
}

func checkErr(err error) {
    if err != nil {
        panic(err)
    }
}

func main() {
    url := "https://alpha.wallhaven.cc/search?categories=111&purity=110&sorting=random&order=desc"

    // Database init
    db, err := sql.Open("mysql", "root:root@tcp(localhost:3306)/wallswap")
    if err != nil {
        panic(err.Error())
    }
    defer db.Close()

    // Truncate and prepare table
    db.Query("truncate wallpaper")
    stmt, err := db.Prepare("INSERT wallpaper SET thumb_url=?,url=?")
    checkErr(err)

    // Channels
    chIds := make(chan string)
    chFinished := make(chan bool)

    // Kick off the crawl process (concurrently)
    go crawl(url, chIds, chFinished)

    // Subscribe to both channels
    isFinished := false
    for isFinished == false {
        select {
        case id := <-chIds:
            fmt.Println(" - " + id)
            stmt.Exec(
                "https://alpha.wallhaven.cc/wallpapers/thumb/small/th-" + id + ".jpg",
                "https://wallpapers.wallhaven.cc/wallpapers/full/wallhaven-" + id + ".jpg")
            checkErr(err)

        case <-chFinished:
            isFinished = true
        }
    }
    close(chIds)
}
<commit_msg>[*] Changed categories for crawl - no more anime<commit_after>package main

import (
    "fmt"
    "golang.org/x/net/html"
    "net/http"
    "database/sql"
    _ "github.com/go-sql-driver/mysql"
)

// Helper function to pull the id attribute from a Token
func getId(t html.Token) (ok bool, id string) {
    // Iterate over all of the Token's attributes until we find an "id"
    for _, attr := range t.Attr {
        if attr.Key == "data-wallpaper-id" {
            id = attr.Val
            ok = true
        }
    }
    return
}

// Extract all http** links from a given webpage
func crawl(url string, ch chan string, chFinished chan bool) {
    resp, err := http.Get(url)

    defer func() {
        // Notify that we're done after this function
        chFinished <- true
    }()

    if err != nil {
        fmt.Println("ERROR: Failed to crawl \"" + url + "\"")
        return
    }

    b := resp.Body
    defer b.Close() // close Body when the function returns

    z := html.NewTokenizer(b)

    for {
        tt := z.Next()

        switch {
        case tt == html.ErrorToken:
            // End of the document, we're done
            return
        case tt == html.StartTagToken:
            t := z.Token()

            // Check if the token is an <img> tag
            isImg := t.Data == "figure"
            if !isImg {
                continue
            }

            // Extract the href value, if there is one
            ok, id := getId(t)
            if !ok {
                continue
            }
            ch <- id
        }
    }
}

func checkErr(err error) {
    if err != nil {
        panic(err)
    }
}

func main() {
    url := "https://alpha.wallhaven.cc/search?categories=101&purity=110&sorting=random&order=desc"

    // Database init
    db, err := sql.Open("mysql", "root:root@tcp(localhost:3306)/wallswap")
    if err != nil {
        panic(err.Error())
    }
    defer db.Close()

    // Truncate and prepare table
    db.Query("truncate wallpaper")
    stmt, err := db.Prepare("INSERT wallpaper SET thumb_url=?,url=?")
    checkErr(err)

    // Channels
    chIds := make(chan string)
    chFinished := make(chan bool)

    // Kick off the crawl process (concurrently)
    go crawl(url, chIds, chFinished)

    // Subscribe to both channels
    isFinished := false
    for isFinished == false {
        select {
        case id := <-chIds:
            fmt.Println(" - " + id)
            stmt.Exec(
                "https://alpha.wallhaven.cc/wallpapers/thumb/small/th-" + id + ".jpg",
                "https://wallpapers.wallhaven.cc/wallpapers/full/wallhaven-" + id + ".jpg")
            checkErr(err)

        case <-chFinished:
            isFinished = true
        }
    }
    close(chIds)
}
<|endoftext|>


<commit_before>// Copyright 2015-2017 trivago GmbH
//
// 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 consumer

import (
	"github.com/fsnotify/fsnotify"
	"github.com/sirupsen/logrus"
	"github.com/trivago/gollum/core"
	"github.com/trivago/tgo"
	"github.com/trivago/tgo/tio"
	"github.com/trivago/tgo/tsync"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	fileBufferGrowSize = 1024
	fileOffsetStart    = "oldest"
	fileOffsetEnd      = "newest"
)

type fileState int32

const (
	fileStateOpen = fileState(iota)
	fileStateRead = fileState(iota)
	fileStateDone = fileState(iota)
)

const (
	observeModePoll  = "poll"
	observeModeWatch = "watch"
)

// File consumer plugin
//
// The file consumer allows to read from files while looking for a delimiter
// that marks the end of a message. If the file is part of e.g. a log rotation
// the file consumer can be set to a symbolic link of the latest file and
// (optionally) be told to reopen the file by sending a SIGHUP. A symlink to
// a file will automatically be reopened if the underlying file is changed.
//
// Configuration example
//
// myConsumer:
//   Type: consumer.File
//   File: /var/run/system.log
//   DefaultOffset: newest
//   OffsetFile: ""
//   Delimiter: "\n"
//   ObserveMode: poll
//   PollingDelay: 100
//
// File is a mandatory setting and contains the file to read. The file will be
// read from beginning to end and the reader will stay attached until the
// consumer is stopped. I.e. appends to the attached file will be recognized
// automatically.
//
// DefaultOffset defines where to start reading the file. Valid values are
// "oldest" and "newest". If OffsetFile is defined the DefaultOffset setting
// will be ignored unless the file does not exist.
// By default this is set to "newest".
//
// OffsetFile defines the path to a file that stores the current offset inside
// the given file. If the consumer is restarted that offset is used to continue
// reading. By default this is set to "" which disables the offset file.
//
// Delimiter defines the end of a message inside the file. By default this is
// set to "\n".
//
// ObserveMode defines the mode how to observe the target file. You can decide between `poll` and `watch`.
// By default this is set to `poll`.
// Note: The watch implementation uses [fsnotify/fsnotify](https://github.com/fsnotify/fsnotify) package.
// Please check for rotating files (like log rotation) if your system supports continue watching on moved files.
//
// PollingDelay defines the time duration how long the consumer will wait to check a file for new content
// after hitting the end of file (EOF) in milliseconds (ms).
// By default this time duration is set to "100" milliseconds.
// Note: This settings take only an effect if the consumer is running in `poll` mode!
type File struct {
	core.SimpleConsumer `gollumdoc:"embed_type"`

	delimiter   string `config:"Delimiter" default:"\n"`
	observeMode string `config:"ObserveMode" default:"poll"`

	seeker  seeker
	source  sourceFile
	watcher *watcher
}

func init() {
	core.TypeRegistry.Register(File{})
}

// Configure initializes this consumer with values from a plugin config.
func (cons *File) Configure(conf core.PluginConfigReader) {
	cons.SetRollCallback(cons.onRoll)

	var err error
	cons.source, err = newSourceFile(conf)
	if err != nil {
		cons.Logger.Fatal(err)
	}

	cons.seeker = newSeeker(conf)

	// restore default observer mode for invalid config settings
	if cons.observeMode != observeModePoll && cons.observeMode != observeModeWatch {
		cons.Logger.WithField("observeMode", cons.observeMode).Errorf("Unknown observe mode '%s'", cons.observeMode)
		cons.observeMode = observeModePoll
	}
}

// Enqueue creates a new message
func (cons *File) Enqueue(data []byte) {
	metaData := core.Metadata{}

	dir, file := filepath.Split(cons.source.realFileName)
	metaData.SetValue("file", []byte(file))
	metaData.SetValue("dir", []byte(dir))

	cons.EnqueueWithMetadata(data, metaData)
}

func (cons *File) storeOffset() {
	ioutil.WriteFile(cons.source.offsetFileName, []byte(strconv.FormatInt(cons.seeker.offset, 10)), 0644)
}

func (cons *File) enqueueAndPersist(data []byte) {
	cons.seeker.offset, _ = cons.source.file.Seek(io.SeekStart, io.SeekCurrent)
	cons.Enqueue(data)
	cons.storeOffset()
}

func (cons *File) setState(state fileState) {
	cons.source.state = state
}

func (cons *File) initFile() {
	defer cons.setState(fileStateRead)

	if cons.source.file != nil {
		cons.source.file.Close()
		cons.source.file = nil
		cons.seeker.seek = cons.seeker.onRotate
		cons.seeker.offset = 0
		if cons.source.offsetFileName != "" {
			cons.storeOffset()
		}
	}

	if cons.source.offsetFileName != "" {
		fileContents, err := ioutil.ReadFile(cons.source.offsetFileName)
		if err == nil {
			cons.seeker.seek = 1
			cons.seeker.offset, err = strconv.ParseInt(string(fileContents), 10, 64)
			if err != nil {
				cons.Logger.Error("Error reading offset file: ", err)
			}
		}
	}
}

func (cons *File) close() {
	if cons.source.file != nil {
		cons.source.file.Close()
	}
	cons.setState(fileStateDone)
	cons.WorkerDone()
}

func (cons *File) observe() {
	defer cons.close()

	sendFunction := cons.Enqueue
	if cons.source.offsetFileName != "" {
		sendFunction = cons.enqueueAndPersist
	}

	buffer := tio.NewBufferedReader(fileBufferGrowSize, 0, 0, cons.delimiter)

	cons.Logger.WithField("file", cons.source.realFileName).Debugf("Use observe mode '%s'", cons.observeMode)
	if cons.observeMode == observeModeWatch {
		cons.watcher = newWatcher(cons.Logger, &cons.source, func() { cons.read(buffer, sendFunction, func() {}, func() {}) })
		cons.watcher.Watch(buffer, sendFunction)
	} else {
		cons.poll(buffer, sendFunction)
	}
}

func (cons *File) poll(buffer *tio.BufferedReader, sendFunction func(data []byte)) {
	spin := tsync.NewCustomSpinner(cons.source.pollingDelay)

	for cons.source.state != fileStateDone {
		cons.read(buffer, sendFunction, spin.Yield, spin.Reset)
	}
}

func (cons *File) read(buffer *tio.BufferedReader, sendFunction func(data []byte), onEOF func(), onAfterRead func()) {
	// Initialize the seek state if requested
	// Try to read the remains of the file first
	if cons.source.state == fileStateOpen {
		if cons.source.file != nil {
			buffer.ReadAll(cons.source.file, sendFunction)
		}
		cons.initFile()
		buffer.Reset(uint64(cons.seeker.offset))
	}

	// Try to open the file to read from
	if cons.source.state == fileStateRead && cons.source.file == nil {
		file, err := os.OpenFile(cons.source.realFileName, os.O_RDONLY, 0666)

		switch {
		case err != nil:
			if cons.source.printFileOpenError {
				cons.Logger.Warning("Open failed: ", err)
				cons.source.printFileOpenError = false
			}
			time.Sleep(3 * time.Second)
			return // ### continue, retry ###

		default:
			cons.source.file = file
			cons.seeker.offset, _ = cons.source.file.Seek(cons.seeker.offset, cons.seeker.seek)
			cons.source.printFileOpenError = true
		}
	}

	// Try to read from the file
	if cons.source.state == fileStateRead && cons.source.file != nil {
		err := buffer.ReadAll(cons.source.file, sendFunction)

		switch {
		case err == nil: // ok
			onAfterRead()

		case err == io.EOF:
			if cons.source.isRotated() {
				cons.Logger.Info("Rotation detected")
				cons.onRoll()
			}
			onEOF()

		case cons.source.state == fileStateRead:
			cons.Logger.Error("Reading failed: ", err)
			cons.source.file.Close()
			cons.source.file = nil
		}
	}
}

func (source *sourceFile) isRotated() bool {
	if source.file.Name() != source.realFileName {
		return true
	} else if time.Since(source.lastStatCheck) > time.Second {
		newStat, newStatErr := os.Stat(source.realFileName)
		oldStat, oldStatErr := source.file.Stat()

		source.lastStatCheck = time.Now()

		if newStatErr == nil && oldStatErr == nil && !os.SameFile(newStat, oldStat) {
			return true
		}
	}

	return false
}

func (cons *File) onRoll() {
	cons.setState(fileStateOpen)
}

// Consume listens to stdin.
func (cons *File) Consume(workers *sync.WaitGroup) {
	cons.setState(fileStateOpen)
	defer cons.setState(fileStateDone)

	go tgo.WithRecoverShutdown(func() {
		cons.AddMainWorker(workers)
		cons.observe()
	})

	cons.ControlLoop()
}

// -- sourceFile --

type sourceFile struct {
	fileName       string        `config:"File" default:"/var/run/system.log"`
	offsetFileName string        `config:"OffsetFile"`
	pollingDelay   time.Duration `config:"PollingDelay" default:"100" metric:"ms"`

	file               *os.File
	realFileName       string
	state              fileState
	printFileOpenError bool
	lastStatCheck      time.Time
}

func (source *sourceFile) Configure(conf core.PluginConfigReader) {
	source.realFileName = source.getRealFileName()

	source.printFileOpenError = true
}

func (source *sourceFile) getRealFileName() string {
	baseFileName, err := filepath.EvalSymlinks(source.fileName)
	if err != nil {
		baseFileName = source.fileName
	}

	baseFileName, err = filepath.Abs(baseFileName)
	if err != nil {
		baseFileName = source.fileName
	}

	return baseFileName
}

func newSourceFile(conf core.PluginConfigReader) (sourceFile, error) {
	source := sourceFile{}
	err := conf.Configure(&source)
	if err != nil {
		return source, err
	}

	return source, nil
}

// -- seeker --

type seeker struct {
	seek     int
	onRotate int
	offset   int64
}

func newSeeker(conf core.PluginConfigReader) seeker {
	switch strings.ToLower(conf.GetString("DefaultOffset", fileOffsetEnd)) {
	default:
		fallthrough
	case fileOffsetEnd:
		return seeker{
			seek:     io.SeekEnd,
			onRotate: io.SeekCurrent,
			offset:   io.SeekStart,
		}

	case fileOffsetStart:
		return seeker{
			seek:     io.SeekCurrent,
			onRotate: io.SeekCurrent,
			offset:   io.SeekStart,
		}
	}
}

// -- watcher --

type watcher struct {
	logger logrus.FieldLogger
	source *sourceFile
	read   func()

	done chan int
}

func newWatcher(logger logrus.FieldLogger, source *sourceFile, readFunction func()) *watcher {
	watcher := watcher{}

	watcher.logger = logger
	watcher.source = source
	watcher.read = readFunction

	return &watcher
}

func (w *watcher) close(watcher *fsnotify.Watcher) {
	err := watcher.Close()
	if err != nil {
		w.logger.Fatal(err)
	}
}

func (w *watcher) Watch(buffer *tio.BufferedReader, sendFunction func(data []byte)) {
	// init
	w.done = make(chan int)

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		w.logger.Fatal(err)
	}
	defer w.close(watcher)

	// watcher loop in subroutine
	go w.watchLoop(watcher)

	// busy loop for shutdown
	for w.source.state != fileStateDone {
		time.Sleep(time.Second)
	}

	w.logger.Debug("shutdown file watcher ..")
	close(w.done)
}

func (w *watcher) watchLoop(watcher *fsnotify.Watcher) {
outerLoop:
	for {
		if _, err := os.Stat(w.source.realFileName); os.IsNotExist(err) {
			w.logger.WithField("file", w.source.realFileName).
				Warning("watched file not exists. retry in 3sec ..")
			time.Sleep(3 * time.Second)
			continue outerLoop // retry
		}

		// read current file state before watching
		w.read()

		err := watcher.Add(w.source.realFileName)
		if err != nil {
			w.logger.Error("error during adding watcher: ", err)
			time.Sleep(3 * time.Second)
			continue outerLoop // retry
		}

	innerLoop:
		for {
			select {
			case event := <-watcher.Events:
				if event.Op&fsnotify.Write == fsnotify.Write {
					w.logger.Debug("modified file: ", event.Name)
					w.read()
					break
				}

				if event.Op&fsnotify.Rename == fsnotify.Rename || event.Op&fsnotify.Remove == fsnotify.Remove {
					w.logger.WithField("event", event).Debug("file renamed/removed: ", event.Name)
					break innerLoop
				}
			case err := <-watcher.Errors:
				w.logger.Error("Error during watch loop: ", err)
			case <-w.done:
				return
			}
		}
	}
}
<commit_msg>update pr requests<commit_after>// Copyright 2015-2017 trivago GmbH
//
// 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 consumer

import (
	"github.com/fsnotify/fsnotify"
	"github.com/sirupsen/logrus"
	"github.com/trivago/gollum/core"
	"github.com/trivago/tgo"
	"github.com/trivago/tgo/tio"
	"github.com/trivago/tgo/tsync"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	fileBufferGrowSize = 1024
	fileOffsetStart    = "oldest"
	fileOffsetEnd      = "newest"
)

type fileState int32

const (
	fileStateOpen = fileState(iota)
	fileStateRead = fileState(iota)
	fileStateDone = fileState(iota)
)

const (
	observeModePoll  = "poll"
	observeModeWatch = "watch"
)

// File consumer plugin
//
// The file consumer allows to read from files while looking for a delimiter
// that marks the end of a message. If the file is part of e.g. a log rotation
// the file consumer can be set to a symbolic link of the latest file and
// (optionally) be told to reopen the file by sending a SIGHUP. A symlink to
// a file will automatically be reopened if the underlying file is changed.
//
// Configuration example
//
// myConsumer:
//   Type: consumer.File
//   File: /var/run/system.log
//   DefaultOffset: newest
//   OffsetFile: ""
//   Delimiter: "\n"
//   ObserveMode: poll
//   PollingDelay: 100
//
// File is a mandatory setting and contains the file to read. The file will be
// read from beginning to end and the reader will stay attached until the
// consumer is stopped. I.e. appends to the attached file will be recognized
// automatically.
//
// DefaultOffset defines where to start reading the file. Valid values are
// "oldest" and "newest". If OffsetFile is defined the DefaultOffset setting
// will be ignored unless the file does not exist.
// By default this is set to "newest".
//
// OffsetFile defines the path to a file that stores the current offset inside
// the given file. If the consumer is restarted that offset is used to continue
// reading. By default this is set to "" which disables the offset file.
//
// Delimiter defines the end of a message inside the file. By default this is
// set to "\n".
//
// ObserveMode defines the mode how to observe the target file. You can decide between `poll` and `watch`.
// By default this is set to `poll`.
// Note: The watch implementation uses [fsnotify/fsnotify](https://github.com/fsnotify/fsnotify) package.
// Please check for rotating files (like log rotation) if your system supports continue watching on moved files.
//
// PollingDelay defines the time duration how long the consumer will wait to check a file for new content
// after hitting the end of file (EOF) in milliseconds (ms).
// By default this time duration is set to "100" milliseconds.
// Note: This settings take only an effect if the consumer is running in `poll` mode!
type File struct {
	core.SimpleConsumer `gollumdoc:"embed_type"`

	delimiter   string `config:"Delimiter" default:"\n"`
	observeMode string `config:"ObserveMode" default:"poll"`

	seeker  seeker
	source  sourceFile
	watcher *watcher
}

func init() {
	core.TypeRegistry.Register(File{})
}

// Configure initializes this consumer with values from a plugin config.
func (cons *File) Configure(conf core.PluginConfigReader) {
	cons.SetRollCallback(cons.onRoll)

	var err error
	cons.source, err = newSourceFile(conf)
	if err != nil {
		cons.Logger.Fatal(err)
	}

	cons.seeker = newSeeker(conf)

	// restore default observer mode for invalid config settings
	if cons.observeMode != observeModePoll && cons.observeMode != observeModeWatch {
		cons.Logger.WithField("observeMode", cons.observeMode).Errorf("Unknown observe mode '%s'", cons.observeMode)
		cons.observeMode = observeModePoll
	}
}

// Enqueue creates a new message
func (cons *File) Enqueue(data []byte) {
	metaData := core.Metadata{}

	dir, file := filepath.Split(cons.source.realFileName)
	metaData.SetValue("file", []byte(file))
	metaData.SetValue("dir", []byte(dir))

	cons.EnqueueWithMetadata(data, metaData)
}

func (cons *File) storeOffset() {
	ioutil.WriteFile(cons.source.offsetFileName, []byte(strconv.FormatInt(cons.seeker.offset, 10)), 0644)
}

func (cons *File) enqueueAndPersist(data []byte) {
	cons.seeker.offset, _ = cons.source.file.Seek(io.SeekStart, io.SeekCurrent)
	cons.Enqueue(data)
	cons.storeOffset()
}

func (cons *File) setState(state fileState) {
	cons.source.state = state
}

func (cons *File) initFile() {
	defer cons.setState(fileStateRead)

	if cons.source.file != nil {
		cons.source.file.Close()
		cons.source.file = nil
		cons.seeker.seek = cons.seeker.onRotate
		cons.seeker.offset = 0
		if cons.source.offsetFileName != "" {
			cons.storeOffset()
		}
	}

	if cons.source.offsetFileName != "" {
		fileContents, err := ioutil.ReadFile(cons.source.offsetFileName)
		if err == nil {
			cons.seeker.seek = 1
			cons.seeker.offset, err = strconv.ParseInt(string(fileContents), 10, 64)
			if err != nil {
				cons.Logger.Error("Error reading offset file: ", err)
			}
		}
	}
}

func (cons *File) close() {
	if cons.source.file != nil {
		cons.source.file.Close()
	}
	cons.setState(fileStateDone)
	cons.WorkerDone()
}

func (cons *File) observe() {
	defer cons.close()

	sendFunction := cons.Enqueue
	if cons.source.offsetFileName != "" {
		sendFunction = cons.enqueueAndPersist
	}

	buffer := tio.NewBufferedReader(fileBufferGrowSize, 0, 0, cons.delimiter)

	cons.Logger.WithField("file", cons.source.realFileName).Debugf("Use observe mode '%s'", cons.observeMode)
	if cons.observeMode == observeModeWatch {
		cons.watcher = newWatcher(cons.Logger, &cons.source, func() { cons.read(buffer, sendFunction, func() {}, func() {}) })
		cons.watcher.Watch(buffer, sendFunction)
	} else {
		cons.poll(buffer, sendFunction)
	}
}

func (cons *File) poll(buffer *tio.BufferedReader, sendFunction func(data []byte)) {
	spin := tsync.NewCustomSpinner(cons.source.pollingDelay)

	for cons.source.state != fileStateDone {
		cons.read(buffer, sendFunction, spin.Yield, spin.Reset)
	}
}

func (cons *File) read(buffer *tio.BufferedReader, sendFunction func(data []byte), onEOF func(), onAfterRead func()) {
	// Initialize the seek state if requested
	// Try to read the remains of the file first
	if cons.source.state == fileStateOpen {
		if cons.source.file != nil {
			buffer.ReadAll(cons.source.file, sendFunction)
		}
		cons.initFile()
		buffer.Reset(uint64(cons.seeker.offset))
	}

	// Try to open the file to read from
	if cons.source.state == fileStateRead && cons.source.file == nil {
		file, err := os.OpenFile(cons.source.realFileName, os.O_RDONLY, 0666)

		switch {
		case err != nil:
			if cons.source.printFileOpenError {
				cons.Logger.Warning("Open failed: ", err)
				cons.source.printFileOpenError = false
			}
			time.Sleep(3 * time.Second)
			return // ### continue, retry ###

		default:
			cons.source.file = file
			cons.seeker.offset, _ = cons.source.file.Seek(cons.seeker.offset, cons.seeker.seek)
			cons.source.printFileOpenError = true
		}
	}

	// Try to read from the file
	if cons.source.state == fileStateRead && cons.source.file != nil {
		err := buffer.ReadAll(cons.source.file, sendFunction)

		switch {
		case err == nil: // ok
			onAfterRead()

		case err == io.EOF:
			if cons.source.isRotated() {
				cons.Logger.Info("Rotation detected")
				cons.onRoll()
			}
			onEOF()

		case cons.source.state == fileStateRead:
			cons.Logger.Error("Reading failed: ", err)
			cons.source.file.Close()
			cons.source.file = nil
		}
	}
}

func (source *sourceFile) isRotated() bool {
	if source.file.Name() != source.realFileName {
		return true
	}

	if time.Since(source.lastStatCheck) > time.Second {
		newStat, newStatErr := os.Stat(source.realFileName)
		oldStat, oldStatErr := source.file.Stat()

		source.lastStatCheck = time.Now()

		if newStatErr == nil && oldStatErr == nil && !os.SameFile(newStat, oldStat) {
			return true
		}
	}

	return false
}

func (cons *File) onRoll() {
	cons.setState(fileStateOpen)
}

// Consume listens to stdin.
func (cons *File) Consume(workers *sync.WaitGroup) {
	cons.setState(fileStateOpen)
	defer cons.setState(fileStateDone)

	go tgo.WithRecoverShutdown(func() {
		cons.AddMainWorker(workers)
		cons.observe()
	})

	cons.ControlLoop()
}

// -- sourceFile --

type sourceFile struct {
	fileName       string        `config:"File" default:"/var/run/system.log"`
	offsetFileName string        `config:"OffsetFile"`
	pollingDelay   time.Duration `config:"PollingDelay" default:"100" metric:"ms"`

	file               *os.File
	realFileName       string
	state              fileState
	printFileOpenError bool
	lastStatCheck      time.Time
}

func (source *sourceFile) Configure(conf core.PluginConfigReader) {
	source.realFileName = source.getRealFileName()

	source.printFileOpenError = true
}

func (source *sourceFile) getRealFileName() string {
	baseFileName, err := filepath.EvalSymlinks(source.fileName)
	if err != nil {
		baseFileName = source.fileName
	}

	baseFileName, err = filepath.Abs(baseFileName)
	if err != nil {
		baseFileName = source.fileName
	}

	return baseFileName
}

func newSourceFile(conf core.PluginConfigReader) (sourceFile, error) {
	source := sourceFile{}
	err := conf.Configure(&source)
	if err != nil {
		return source, err
	}

	return source, nil
}

// -- seeker --

type seeker struct {
	seek     int
	onRotate int
	offset   int64
}

func newSeeker(conf core.PluginConfigReader) seeker {
	switch strings.ToLower(conf.GetString("DefaultOffset", fileOffsetEnd)) {
	default:
		fallthrough
	case fileOffsetEnd:
		return seeker{
			seek:     io.SeekEnd,
			onRotate: io.SeekCurrent,
			offset:   io.SeekStart,
		}

	case fileOffsetStart:
		return seeker{
			seek:     io.SeekCurrent,
			onRotate: io.SeekCurrent,
			offset:   io.SeekStart,
		}
	}
}

// -- watcher --

type watcher struct {
	logger logrus.FieldLogger
	source *sourceFile
	read   func()

	done chan int
}

func newWatcher(logger logrus.FieldLogger, source *sourceFile, readFunction func()) *watcher {
	watcher := watcher{}

	watcher.logger = logger
	watcher.source = source
	watcher.read = readFunction

	return &watcher
}

func (w *watcher) close(watcher *fsnotify.Watcher) {
	err := watcher.Close()
	if err != nil {
		w.logger.Fatal(err)
	}
}

func (w *watcher) Watch(buffer *tio.BufferedReader, sendFunction func(data []byte)) {
	// init
	w.done = make(chan int)

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		w.logger.Fatal(err)
	}
	defer w.close(watcher)

	// watcher loop in subroutine
	go w.watchLoop(watcher)

	// busy loop for shutdown
	for w.source.state != fileStateDone {
		time.Sleep(time.Second)
	}

	w.logger.Debug("shutdown file watcher ..")
	close(w.done)
}

func (w *watcher) watchLoop(watcher *fsnotify.Watcher) {
	for {
		if _, err := os.Stat(w.source.realFileName); os.IsNotExist(err) {
			w.logger.WithField("file", w.source.realFileName).
				Warning("watched file not exists. retry in 3sec ..")
			time.Sleep(3 * time.Second)
			continue // retry
		}

		// read current file state before watching
		w.read()

		err := watcher.Add(w.source.realFileName)
		if err != nil {
			w.logger.Error("error during adding watcher: ", err)
			time.Sleep(3 * time.Second)
			continue // retry
		}

	fileEvent:
		for {
			select {
			case event := <-watcher.Events:
				if event.Op&fsnotify.Write == fsnotify.Write {
					w.logger.Debug("modified file: ", event.Name)
					w.read()
					continue fileEvent //break select
				}

				if event.Op&fsnotify.Rename == fsnotify.Rename || event.Op&fsnotify.Remove == fsnotify.Remove {
					w.logger.WithField("event", event).Debug("file renamed/removed: ", event.Name)
					break fileEvent
				}
			case err := <-watcher.Errors:
				w.logger.Error("Error during watch loop: ", err)
			case <-w.done:
				return
			}
		}
	}
}
<|endoftext|>