<commit_before>/*
 * Copyright (c) 2017 Sam Kumar <samkumar@berkeley.edu>
 * Copyright (c) 2017 University of California, Berkeley
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *     * Redistributions of source code must retain the above copyright
 *       notice, this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *     * Neither the name of the University of California, Berkeley nor the
 *       names of its contributors may be used to endorse or promote products
 *       derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNERS OR CONTRIBUTORS BE LIABLE FOR
 * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

// Package accounts implements tools to manage account information for Mr. Plotter.
// Account information is stored in etcd, and so a Version 3 etcd client is needed
// for most of the API functions.
package accounts

import (
    "context"
    "fmt"

    "golang.org/x/crypto/bcrypt"

    "github.com/ugorji/go/codec"
    etcd "github.com/coreos/etcd/clientv3"
)

// The "public" tag is used to specify streams viewable to all users, even those
// who are not logged in.
const PUBLIC_TAG = "public"

// The "all" tag represents permissions to view all streams. It is stored
// implicitly (i.e. not in etcd) and cannot be redefined.
const ALL_TAG = "all"

const accountpath string = "mrplotter/accounts/"
const tagpath string = "mrplotter/tagdefs/"
var etcdprefix string = ""

var mp codec.Handle = &codec.MsgpackHandle{}

// MrPlotterAccount abstracts account information for a single user.
type MrPlotterAccount struct {
    Username string
    Tags map[string]struct{}
    PasswordHash []byte

    retrievedRevision int64
}

// MrPlotterTagDef abstracts a mapping from a Tag to set of collection prefixes.
type MrPlotterTagDef struct {
    Tag string
    PathPrefix map[string]struct{}

    retrievedRevision int64
}

// Can be used to abstract a MrPlotterAccount or MrPlotterTagDef
type MrPlotterEtcdToken interface {
    getEtcdKey() string
    setRetrievedRevision(rev int64)
    getRetrievedRevision() int64

    getTypePath() string
}

func (acc *MrPlotterAccount) getEtcdKey() string {
    return getEtcdKey(acc.Username, accountpath)
}

func (acc *MrPlotterAccount) setRetrievedRevision(rev int64) {
    acc.retrievedRevision = rev
}

func (acc *MrPlotterAccount) getRetrievedRevision() int64 {
    return acc.retrievedRevision
}

func (acc *MrPlotterAccount) getTypePath() string {
    return accountpath
}

// Sets the password for a user.
func (acc *MrPlotterAccount) SetPassword(newPassword []byte) error {
    phash, err := bcrypt.GenerateFromPassword(newPassword, bcrypt.DefaultCost)
    if (err == nil) {
        acc.PasswordHash = phash
    }
    return err
}

// Checks whether the provided password matches the user's password. Returns
// true if the provided password is correct, and false if it is not. If the
// returned error is not nil, then the returned boolean should be ignored.
func (acc *MrPlotterAccount) CheckPassword(password []byte) (bool, error) {
    err := bcrypt.CompareHashAndPassword(acc.PasswordHash, password)
    if err == nil {
        return true, nil
    } else if err == bcrypt.ErrMismatchedHashAndPassword {
        return false, nil
    } else {
        return false, err
    }
}

func (tdef *MrPlotterTagDef) getEtcdKey() string {
    return getEtcdKey(tdef.Tag, tagpath)
}

func (tdef *MrPlotterTagDef) setRetrievedRevision(rev int64) {
    tdef.retrievedRevision = rev
}

func (tdef *MrPlotterTagDef) getRetrievedRevision() int64 {
    return tdef.retrievedRevision
}

func (tdef *MrPlotterTagDef) getTypePath() string {
    return tagpath
}

func encodeToken(tok MrPlotterEtcdToken) ([]byte, error) {
    var encoded []byte

    encoder := codec.NewEncoderBytes(&encoded, mp)
    err := encoder.Encode(tok)
    if err != nil {
        return nil, err
    }

    return encoded, nil
}

func decodeToken(encoded []byte, into interface{}) error {
    decoder := codec.NewDecoderBytes(encoded, mp)
    return decoder.Decode(into)
}

// Sets the prefix added to keys in the etcd database.
// The keys used are of the form <prefix>mrplotter/accounts/<username>.
// The prefix allows multiple configurations for Mr. Plotter to coexist in a
// single etcd database system.
//
// Modifying the prefix while an operation is in progress results in undefined
// behavior. Furthermore, MrPlotterAccount structs returned by RetrieveAccount
// or RetrieveMultipleAccounts must always be used with the same prefix with
// which they were generated.
func SetEtcdKeyPrefix(prefix string) {
    etcdprefix = prefix
}

// Gets the base path for account data in etcd.
func GetAccountEtcdPath() string {
    return fmt.Sprintf("%s%s", etcdprefix, accountpath)
}

// Gets the base path for tag definitions in etcd.
func GetTagEtcdPath() string {
    return fmt.Sprintf("%s%s", etcdprefix, tagpath)
}

func getEtcdKey(name string, typepath string) string {
    return fmt.Sprintf("%s%s%s", etcdprefix, typepath, name)
}

func getNameFromEtcdKey(etcdKey string, typepath string) string {
    return etcdKey[len(etcdprefix) + len(typepath):]
}

func retrieveToken(ctx context.Context, etcdClient *etcd.Client, username string, retrieveInto MrPlotterEtcdToken) error {
    etcdKey := getEtcdKey(username, retrieveInto.getTypePath())
    resp, err := etcdClient.Get(ctx, etcdKey)
    if err != nil {
        return err
    }

    /* No token with that name exists. */
    if len(resp.Kvs) == 0 {
        return nil
    }

    err = decodeToken(resp.Kvs[0].Value, retrieveInto)
    if err != nil {
        return err
    }

    retrieveInto.setRetrievedRevision(resp.Kvs[0].ModRevision)
    return nil
}

// Retrieves the account information for the specified user.
// Setting the "Username" field in the returned struct renders it unsuitable
// for use with the "UpsertAccountAtomically" function.
func RetrieveAccount(ctx context.Context, etcdClient *etcd.Client, username string) (acc *MrPlotterAccount, err error) {
    acc = &MrPlotterAccount{}
    err = retrieveToken(ctx, etcdClient, username, acc)
    return
}

// Retrieves the tag definition for the specified tag.
// Setting the "Tag" field in the returned struct renders it unsuitable
// for use with the "UpsertTagDefAtomically" function.
func RetrieveTagDef(ctx context.Context, etcdClient *etcd.Client, tag string) (tdef *MrPlotterTagDef, err error) {
    tdef = &MrPlotterTagDef{}
    err = retrieveToken(ctx, etcdClient, tag, tdef)
    return
}

func upsertToken(ctx context.Context, etcdClient *etcd.Client, tok MrPlotterEtcdToken) error {
    encoded, err := encodeToken(tok)
    if err != nil {
        return err
    }

    etcdKey := tok.getEtcdKey()
    _, err = etcdClient.Put(ctx, etcdKey, string(encoded))
    return err
}

// Updates the account according to the provided account information, creating
// a new user account if a user with that username does not exist.
func UpsertAccount(ctx context.Context, etcdClient *etcd.Client, acc *MrPlotterAccount) error {
    return upsertToken(ctx, etcdClient, acc)
}

func UpsertTagDef(ctx context.Context, etcdClient *etcd.Client, tdef *MrPlotterTagDef) error {
    return upsertToken(ctx, etcdClient, tdef)
}

// Same as UpsertAccount, but fails if the account was updated meanwhile. This
// allows one to ensure that the read-modify-write operation required to update
// account information can be done atomically. Returns true if the operation
// succeeds, and returns false otherwise.
//
// The rules are as follows: if the MrPlotterAccount struct was created
// directly, the operation fails if the account already exists in etcd. If the
// MrPlotterAccount struct was returned by RetrieveAccount or
// RetrieveMultipleAccounts, the operation fails if the data stored in etcd
// was updated after the account information was retrieved. You should not
// modify the "Username" field of a struct returned by RetrieveAccount or
// RetrieveMultipleAccounts and then use it with this function. Setting the
// "Username" field of a struct to be used with this function is only allowed
// for structs that are created directly.
func upsertTokenAtomically(ctx context.Context, etcdClient *etcd.Client, tok MrPlotterEtcdToken) (bool, error) {
    encoded, err := encodeToken(tok)
    if err != nil {
        return false, err
    }

    etcdKey := tok.getEtcdKey()
    resp, err := etcdClient.Txn(ctx).
        If(etcd.Compare(etcd.ModRevision(etcdKey), "=", tok.getRetrievedRevision())).
        Then(etcd.OpPut(etcdKey, string(encoded))).
        Commit()
    if resp != nil {
        return resp.Succeeded, err
    } else {
        return false, err
    }
}

func UpsertAccountAtomically(ctx context.Context, etcdClient *etcd.Client, acc *MrPlotterAccount) (bool, error) {
    return upsertTokenAtomically(ctx, etcdClient, acc)
}

func UpsertTagDefAtomically(ctx context.Context, etcdClient *etcd.Client, tdef *MrPlotterTagDef) (bool, error) {
    return upsertTokenAtomically(ctx, etcdClient, tdef)
}

// Deletes the account of the user with the provided username.
func DeleteAccount(ctx context.Context, etcdClient *etcd.Client, username string) error {
    etcdKey := getEtcdKey(username, accountpath)
    _, err := etcdClient.Delete(ctx, etcdKey)
    return err
}

// Deletes the tag definition with the provided name
func DeleteTagDef(ctx context.Context, etcdClient *etcd.Client, tag string) error {
    etcdKey := getEtcdKey(tag, tagpath)
    _, err := etcdClient.Delete(ctx, etcdKey)
    return err
}

// Retrieves the account information of all users whose username begins with
// the provided prefix.
// If one entry is in a corrupt state and cannot be decoded, its Tags set will
// be set to nil and decoding will continue.
// Setting the "Username" field in the returned struct renders it unsuitable
// for use with the "UpsertAccountAtomically" function.
func RetrieveMultipleAccounts(ctx context.Context, etcdClient *etcd.Client, usernameprefix string) ([]*MrPlotterAccount, error) {
    etcdKeyPrefix := getEtcdKey(usernameprefix, accountpath)
    resp, err := etcdClient.Get(ctx, etcdKeyPrefix, etcd.WithPrefix())
    if err != nil {
        return nil, err
    }

    accs := make([]*MrPlotterAccount, 0, len(resp.Kvs))
    var acc *MrPlotterAccount
    for _, kv := range resp.Kvs {
        acc = nil
        err := decodeToken(kv.Value, &acc)
        if err != nil {
            acc = &MrPlotterAccount{Username: getNameFromEtcdKey(string(kv.Key), accountpath)}
        }
        acc.setRetrievedRevision(kv.ModRevision)
        accs = append(accs, acc)
    }

    return accs, nil
}

// Retrieves the account information of all users whose username begins with
// the provided prefix.
// If one entry is in a corrupt state and cannot be decoded, its Tags set will
// be set to nil and decoding will continue.
// Setting the "Username" field in the returned struct renders it unsuitable
// for use with the "UpsertAccountAtomically" function.
func RetrieveMultipleTagDefs(ctx context.Context, etcdClient *etcd.Client, tagprefix string) ([]*MrPlotterTagDef, error) {
    etcdKeyPrefix := getEtcdKey(tagprefix, tagpath)
    resp, err := etcdClient.Get(ctx, etcdKeyPrefix, etcd.WithPrefix())
    if err != nil {
        return nil, err
    }

    tdefs := make([]*MrPlotterTagDef, 0, len(resp.Kvs))
    var tdef *MrPlotterTagDef
    for _, kv := range resp.Kvs {
        err := decodeToken(kv.Value, &tdef)
        if err != nil {
            tdef = &MrPlotterTagDef{Tag: getNameFromEtcdKey(string(kv.Key), tagpath)}
        }
        tdef.setRetrievedRevision(kv.ModRevision)
        tdefs = append(tdefs, tdef)
    }

    return tdefs, nil
}

func deleteMultipleTokens(ctx context.Context, etcdClient *etcd.Client, prefix string, typepath string) (int64, error) {
    etcdKeyPrefix := getEtcdKey(prefix, typepath)
    resp, err := etcdClient.Delete(ctx, etcdKeyPrefix, etcd.WithPrefix())
    if resp != nil {
        return resp.Deleted, err
    } else {
        return 0, err
    }
}

// Deletes the accounts of all users whose username begins with the provided
// prefix.
func DeleteMultipleAccounts(ctx context.Context, etcdClient *etcd.Client, usernameprefix string) (int64, error) {
    return deleteMultipleTokens(ctx, etcdClient, usernameprefix, accountpath)
}

// Deletes all tag definitions beginning with the given prefix.
func DeleteMultipleTagDefs(ctx context.Context, etcdClient *etcd.Client, tagprefix string) (int64, error) {
    return deleteMultipleTokens(ctx, etcdClient, tagprefix, tagpath)
}
<commit_msg>Fix bug in accounts.go<commit_after>/*
 * Copyright (c) 2017 Sam Kumar <samkumar@berkeley.edu>
 * Copyright (c) 2017 University of California, Berkeley
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *     * Redistributions of source code must retain the above copyright
 *       notice, this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *     * Neither the name of the University of California, Berkeley nor the
 *       names of its contributors may be used to endorse or promote products
 *       derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNERS OR CONTRIBUTORS BE LIABLE FOR
 * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

// Package accounts implements tools to manage account information for Mr. Plotter.
// Account information is stored in etcd, and so a Version 3 etcd client is needed
// for most of the API functions.
package accounts

import (
	"context"
	"fmt"

	"golang.org/x/crypto/bcrypt"

	etcd "github.com/coreos/etcd/clientv3"
	"github.com/ugorji/go/codec"
)

// The "public" tag is used to specify streams viewable to all users, even those
// who are not logged in.
const PUBLIC_TAG = "public"

// The "all" tag represents permissions to view all streams. It is stored
// implicitly (i.e. not in etcd) and cannot be redefined.
const ALL_TAG = "all"

const accountpath string = "mrplotter/accounts/"
const tagpath string = "mrplotter/tagdefs/"

var etcdprefix string = ""

var mp codec.Handle = &codec.MsgpackHandle{}

// MrPlotterAccount abstracts account information for a single user.
type MrPlotterAccount struct {
	Username     string
	Tags         map[string]struct{}
	PasswordHash []byte

	retrievedRevision int64
}

// MrPlotterTagDef abstracts a mapping from a Tag to set of collection prefixes.
type MrPlotterTagDef struct {
	Tag        string
	PathPrefix map[string]struct{}

	retrievedRevision int64
}

// Can be used to abstract a MrPlotterAccount or MrPlotterTagDef
type MrPlotterEtcdToken interface {
	getEtcdKey() string
	setRetrievedRevision(rev int64)
	getRetrievedRevision() int64

	getTypePath() string
}

func (acc *MrPlotterAccount) getEtcdKey() string {
	return getEtcdKey(acc.Username, accountpath)
}

func (acc *MrPlotterAccount) setRetrievedRevision(rev int64) {
	acc.retrievedRevision = rev
}

func (acc *MrPlotterAccount) getRetrievedRevision() int64 {
	return acc.retrievedRevision
}

func (acc *MrPlotterAccount) getTypePath() string {
	return accountpath
}

// Sets the password for a user.
func (acc *MrPlotterAccount) SetPassword(newPassword []byte) error {
	phash, err := bcrypt.GenerateFromPassword(newPassword, bcrypt.DefaultCost)
	if err == nil {
		acc.PasswordHash = phash
	}
	return err
}

// Checks whether the provided password matches the user's password. Returns
// true if the provided password is correct, and false if it is not. If the
// returned error is not nil, then the returned boolean should be ignored.
func (acc *MrPlotterAccount) CheckPassword(password []byte) (bool, error) {
	err := bcrypt.CompareHashAndPassword(acc.PasswordHash, password)
	if err == nil {
		return true, nil
	} else if err == bcrypt.ErrMismatchedHashAndPassword {
		return false, nil
	} else {
		return false, err
	}
}

func (tdef *MrPlotterTagDef) getEtcdKey() string {
	return getEtcdKey(tdef.Tag, tagpath)
}

func (tdef *MrPlotterTagDef) setRetrievedRevision(rev int64) {
	tdef.retrievedRevision = rev
}

func (tdef *MrPlotterTagDef) getRetrievedRevision() int64 {
	return tdef.retrievedRevision
}

func (tdef *MrPlotterTagDef) getTypePath() string {
	return tagpath
}

func encodeToken(tok MrPlotterEtcdToken) ([]byte, error) {
	var encoded []byte

	encoder := codec.NewEncoderBytes(&encoded, mp)
	err := encoder.Encode(tok)
	if err != nil {
		return nil, err
	}

	return encoded, nil
}

func decodeToken(encoded []byte, into interface{}) error {
	decoder := codec.NewDecoderBytes(encoded, mp)
	return decoder.Decode(into)
}

// Sets the prefix added to keys in the etcd database.
// The keys used are of the form <prefix>mrplotter/accounts/<username>.
// The prefix allows multiple configurations for Mr. Plotter to coexist in a
// single etcd database system.
//
// Modifying the prefix while an operation is in progress results in undefined
// behavior. Furthermore, MrPlotterAccount structs returned by RetrieveAccount
// or RetrieveMultipleAccounts must always be used with the same prefix with
// which they were generated.
func SetEtcdKeyPrefix(prefix string) {
	etcdprefix = prefix
}

// Gets the base path for account data in etcd.
func GetAccountEtcdPath() string {
	return fmt.Sprintf("%s%s", etcdprefix, accountpath)
}

// Gets the base path for tag definitions in etcd.
func GetTagEtcdPath() string {
	return fmt.Sprintf("%s%s", etcdprefix, tagpath)
}

func getEtcdKey(name string, typepath string) string {
	return fmt.Sprintf("%s%s%s", etcdprefix, typepath, name)
}

func getNameFromEtcdKey(etcdKey string, typepath string) string {
	return etcdKey[len(etcdprefix)+len(typepath):]
}

func retrieveToken(ctx context.Context, etcdClient *etcd.Client, username string, retrieveInto MrPlotterEtcdToken) error {
	etcdKey := getEtcdKey(username, retrieveInto.getTypePath())
	resp, err := etcdClient.Get(ctx, etcdKey)
	if err != nil {
		return err
	}

	/* No token with that name exists. */
	if len(resp.Kvs) == 0 {
		return nil
	}

	err = decodeToken(resp.Kvs[0].Value, retrieveInto)
	if err != nil {
		return err
	}

	retrieveInto.setRetrievedRevision(resp.Kvs[0].ModRevision)
	return nil
}

// Retrieves the account information for the specified user.
// Setting the "Username" field in the returned struct renders it unsuitable
// for use with the "UpsertAccountAtomically" function.
func RetrieveAccount(ctx context.Context, etcdClient *etcd.Client, username string) (acc *MrPlotterAccount, err error) {
	acc = &MrPlotterAccount{}
	err = retrieveToken(ctx, etcdClient, username, acc)
	return
}

// Retrieves the tag definition for the specified tag.
// Setting the "Tag" field in the returned struct renders it unsuitable
// for use with the "UpsertTagDefAtomically" function.
func RetrieveTagDef(ctx context.Context, etcdClient *etcd.Client, tag string) (tdef *MrPlotterTagDef, err error) {
	tdef = &MrPlotterTagDef{}
	err = retrieveToken(ctx, etcdClient, tag, tdef)
	return
}

func upsertToken(ctx context.Context, etcdClient *etcd.Client, tok MrPlotterEtcdToken) error {
	encoded, err := encodeToken(tok)
	if err != nil {
		return err
	}

	etcdKey := tok.getEtcdKey()
	_, err = etcdClient.Put(ctx, etcdKey, string(encoded))
	return err
}

// Updates the account according to the provided account information, creating
// a new user account if a user with that username does not exist.
func UpsertAccount(ctx context.Context, etcdClient *etcd.Client, acc *MrPlotterAccount) error {
	return upsertToken(ctx, etcdClient, acc)
}

func UpsertTagDef(ctx context.Context, etcdClient *etcd.Client, tdef *MrPlotterTagDef) error {
	return upsertToken(ctx, etcdClient, tdef)
}

func upsertTokenAtomically(ctx context.Context, etcdClient *etcd.Client, tok MrPlotterEtcdToken) (bool, error) {
	encoded, err := encodeToken(tok)
	if err != nil {
		return false, err
	}

	etcdKey := tok.getEtcdKey()
	resp, err := etcdClient.Txn(ctx).
		If(etcd.Compare(etcd.ModRevision(etcdKey), "=", tok.getRetrievedRevision())).
		Then(etcd.OpPut(etcdKey, string(encoded))).
		Commit()
	if resp != nil {
		return resp.Succeeded, err
	} else {
		return false, err
	}
}

// Same as UpsertAccount, but fails if the account was updated meanwhile. This
// allows one to ensure that the read-modify-write operation required to update
// account information can be done atomically. Returns true if the operation
// succeeds, and returns false otherwise.
//
// The rules are as follows: if the MrPlotterAccount struct was created
// directly, the operation fails if the account already exists in etcd. If the
// MrPlotterAccount struct was returned by RetrieveAccount or
// RetrieveMultipleAccounts, the operation fails if the data stored in etcd
// was updated after the account information was retrieved. You should not
// modify the "Username" field of a struct returned by RetrieveAccount or
// RetrieveMultipleAccounts and then use it with this function. Setting the
// "Username" field of a struct to be used with this function is only allowed
// for structs that are created directly.
func UpsertAccountAtomically(ctx context.Context, etcdClient *etcd.Client, acc *MrPlotterAccount) (bool, error) {
	return upsertTokenAtomically(ctx, etcdClient, acc)
}

// Same as UpsertAccountAtomically, but for Tag Definitions.
func UpsertTagDefAtomically(ctx context.Context, etcdClient *etcd.Client, tdef *MrPlotterTagDef) (bool, error) {
	return upsertTokenAtomically(ctx, etcdClient, tdef)
}

// Deletes the account of the user with the provided username.
func DeleteAccount(ctx context.Context, etcdClient *etcd.Client, username string) error {
	etcdKey := getEtcdKey(username, accountpath)
	_, err := etcdClient.Delete(ctx, etcdKey)
	return err
}

// Deletes the tag definition with the provided name
func DeleteTagDef(ctx context.Context, etcdClient *etcd.Client, tag string) error {
	etcdKey := getEtcdKey(tag, tagpath)
	_, err := etcdClient.Delete(ctx, etcdKey)
	return err
}

// Retrieves the account information of all users whose username begins with
// the provided prefix.
// If one entry is in a corrupt state and cannot be decoded, its Tags set will
// be set to nil and decoding will continue.
// Setting the "Username" field in the returned struct renders it unsuitable
// for use with the "UpsertAccountAtomically" function.
func RetrieveMultipleAccounts(ctx context.Context, etcdClient *etcd.Client, usernameprefix string) ([]*MrPlotterAccount, error) {
	etcdKeyPrefix := getEtcdKey(usernameprefix, accountpath)
	resp, err := etcdClient.Get(ctx, etcdKeyPrefix, etcd.WithPrefix())
	if err != nil {
		return nil, err
	}

	accs := make([]*MrPlotterAccount, 0, len(resp.Kvs))
	var acc *MrPlotterAccount
	for _, kv := range resp.Kvs {
		acc = nil
		err := decodeToken(kv.Value, &acc)
		if err != nil {
			acc = &MrPlotterAccount{Username: getNameFromEtcdKey(string(kv.Key), accountpath)}
		}
		acc.setRetrievedRevision(kv.ModRevision)
		accs = append(accs, acc)
	}

	return accs, nil
}

// Retrieves the account information of all users whose username begins with
// the provided prefix.
// If one entry is in a corrupt state and cannot be decoded, its Tags set will
// be set to nil and decoding will continue.
// Setting the "Username" field in the returned struct renders it unsuitable
// for use with the "UpsertAccountAtomically" function.
func RetrieveMultipleTagDefs(ctx context.Context, etcdClient *etcd.Client, tagprefix string) ([]*MrPlotterTagDef, error) {
	etcdKeyPrefix := getEtcdKey(tagprefix, tagpath)
	resp, err := etcdClient.Get(ctx, etcdKeyPrefix, etcd.WithPrefix())
	if err != nil {
		return nil, err
	}

	tdefs := make([]*MrPlotterTagDef, 0, len(resp.Kvs))
	var tdef *MrPlotterTagDef
	for _, kv := range resp.Kvs {
		tdef = nil
		err := decodeToken(kv.Value, &tdef)
		if err != nil {
			tdef = &MrPlotterTagDef{Tag: getNameFromEtcdKey(string(kv.Key), tagpath)}
		}
		tdef.setRetrievedRevision(kv.ModRevision)
		tdefs = append(tdefs, tdef)
	}

	return tdefs, nil
}

func deleteMultipleTokens(ctx context.Context, etcdClient *etcd.Client, prefix string, typepath string) (int64, error) {
	etcdKeyPrefix := getEtcdKey(prefix, typepath)
	resp, err := etcdClient.Delete(ctx, etcdKeyPrefix, etcd.WithPrefix())
	if resp != nil {
		return resp.Deleted, err
	} else {
		return 0, err
	}
}

// Deletes the accounts of all users whose username begins with the provided
// prefix.
func DeleteMultipleAccounts(ctx context.Context, etcdClient *etcd.Client, usernameprefix string) (int64, error) {
	return deleteMultipleTokens(ctx, etcdClient, usernameprefix, accountpath)
}

// Deletes all tag definitions beginning with the given prefix.
func DeleteMultipleTagDefs(ctx context.Context, etcdClient *etcd.Client, tagprefix string) (int64, error) {
	return deleteMultipleTokens(ctx, etcdClient, tagprefix, tagpath)
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/flowcommerce/tools/executor"
	"github.com/flowcommerce/tools/scala_library"
)

func main() {
	executor := executor.Create("lib-reference-scala-json")
	executor = executor.Add("go get -u github.com/flowcommerce/json-reference/common")
	executor = executor.Add("go run script/generate.go")
	executor = executor.Add("sbt test")
	executor = executor.Add("git commit -am 'autocommit: Regenerate models'")

	executor.Run()

	scala_library.Run()
}
<commit_msg>Add --allow empty<commit_after>package main

import (
	"github.com/flowcommerce/tools/executor"
	"github.com/flowcommerce/tools/scala_library"
)

func main() {
	executor := executor.Create("lib-reference-scala-json")
	executor = executor.Add("go get -u github.com/flowcommerce/json-reference/common")
	executor = executor.Add("go run script/generate.go")
	executor = executor.Add("sbt test")
	executor = executor.Add("git commit --allow-empty -am 'autocommit: Regenerate models'")

	executor.Run()

	scala_library.Run()
}
<|endoftext|>


<commit_before>package jwt

import (
	. "github.com/WindomZ/go-jwt/macro"
	jwt "github.com/dgrijalva/jwt-go"
	"net/http"
	"strings"
)

// set key and prefix value in http.Request header entries
func SetHTTPHeaderKey(key string) {
	SetHTTPHeaderKeyAndValuePrefix(key, "")
}

// Generate the signing string, and set into http request
func SignRequest(req *http.Request, kid string, arg interface{}, minutes int) error {
	token, err := Sign(kid, arg, minutes)
	if err != nil {
		return err
	}
	req.Header.Set(HEADER_KEY(), HEADER_VALUE(token))
	return nil
}

// Parse http request, validate, and return a token.
func ParseRequest(req *http.Request) (interface{}, error) {
	if req == nil {
		return nil, ErrRequest
	} else if ah := req.Header.Get(HEADER_KEY()); ah != "" {
		if len(HEADER_VALUE_PREFIX()) == 0 {
			return parseToken(jwt.Parse(ah, getJwtHandlerKey))
		} else if strings.HasPrefix(
			strings.ToUpper(ah[:len(HEADER_VALUE_PREFIX())]),
			HEADER_VALUE_PREFIX()) {
			return parseToken(jwt.Parse(ah[(len(HEADER_VALUE_PREFIX())+1):],
				getJwtHandlerKey))
		}
	}
	return nil, jwt.ErrNoTokenInRequest
}
<commit_msg>fix: http request<commit_after>package jwt

import (
	. "github.com/WindomZ/go-jwt/macro"
	jwt "github.com/dgrijalva/jwt-go"
	"net/http"
	"strings"
)

// set key and prefix value in http.Request header entries
func SetHTTPHeaderKey(key string) {
	SetHTTPHeaderKeyAndValuePrefix(key, "")
}

// Generate the signing string, and set into http request
func SignRequest(req *http.Request, kid string, arg interface{}, minutes int) error {
	token, err := Sign(kid, arg, minutes)
	if err != nil {
		return err
	}
	req.Header.Set(HEADER_KEY(), HEADER_VALUE(token))
	return nil
}

// Parse http request, validate, and return a token.
func ParseRequest(req *http.Request) (interface{}, error) {
	if req == nil {
		return nil, ErrRequest
	} else if ah := req.Header.Get(HEADER_KEY()); len(ah) != 0 {
		if len(HEADER_VALUE_PREFIX()) == 0 {
			return parseToken(jwt.Parse(ah, getJwtHandlerKey))
		} else if strings.HasPrefix(ah[:len(HEADER_VALUE_PREFIX())], HEADER_VALUE_PREFIX()) {
			return parseToken(jwt.Parse(ah[(len(HEADER_VALUE_PREFIX())+1):],
				getJwtHandlerKey))
		}
	}
	return nil, jwt.ErrNoTokenInRequest
}
<|endoftext|>


<commit_before>package socks5

import (
	"fmt"
	"io"
	"net"
	"strings"

	"golang.org/x/net/context"
)

const (
	ConnectCommand   = uint8(1)
	BindCommand      = uint8(2)
	AssociateCommand = uint8(3)
	ipv4Address      = uint8(1)
	fqdnAddress      = uint8(3)
	ipv6Address      = uint8(4)
)

const (
	successReply uint8 = iota
	serverFailure
	ruleFailure
	networkUnreachable
	hostUnreachable
	connectionRefused
	ttlExpired
	commandNotSupported
	addrTypeNotSupported
)

var (
	unrecognizedAddrType = fmt.Errorf("Unrecognized address type")
)

// AddressRewriter is used to rewrite a destination transparently
type AddressRewriter interface {
	Rewrite(ctx context.Context, request *Request) (context.Context, *AddrSpec)
}

// AddrSpec is used to return the target AddrSpec
// which may be specified as IPv4, IPv6, or a FQDN
type AddrSpec struct {
	FQDN string
	IP   net.IP
	Port int
}

func (a *AddrSpec) String() string {
	if a.FQDN != "" {
		return fmt.Sprintf("%s (%s):%d", a.FQDN, a.IP, a.Port)
	}
	return fmt.Sprintf("%s:%d", a.IP, a.Port)
}

// A Request represents request received by a server
type Request struct {
	// Protocol version
	Version uint8
	// Requested command
	Command uint8
	// AuthContext provided during negotiation
	AuthContext *AuthContext
	// AddrSpec of the the network that sent the request
	RemoteAddr *AddrSpec
	// AddrSpec of the desired destination
	DestAddr *AddrSpec
	// AddrSpec of the actual destination (might be affected by rewrite)
	realDestAddr *AddrSpec
	bufConn      io.Reader
}

type conn interface {
	Write([]byte) (int, error)
	RemoteAddr() net.Addr
}

// NewRequest creates a new Request from the tcp connection
func NewRequest(bufConn io.Reader) (*Request, error) {
	// Read the version byte
	header := []byte{0, 0, 0}
	if _, err := io.ReadAtLeast(bufConn, header, 3); err != nil {
		return nil, fmt.Errorf("Failed to get command version: %v", err)
	}

	// Ensure we are compatible
	if header[0] != socks5Version {
		return nil, fmt.Errorf("Unsupported command version: %v", header[0])
	}

	// Read in the destination address
	dest, err := readAddrSpec(bufConn)
	if err != nil {
		return nil, err
	}

	request := &Request{
		Version:  socks5Version,
		Command:  header[1],
		DestAddr: dest,
		bufConn:  bufConn,
	}

	return request, nil
}

// handleRequest is used for request processing after authentication
func (s *Server) handleRequest(req *Request, conn conn) error {
	ctx := context.Background()

	// Resolve the address if we have a FQDN
	dest := req.DestAddr
	if dest.FQDN != "" {
		ctx_, addr, err := s.config.Resolver.Resolve(ctx, dest.FQDN)
		if err != nil {
			if err := sendReply(conn, hostUnreachable, nil); err != nil {
				return fmt.Errorf("Failed to send reply: %v", err)
			}
			return fmt.Errorf("Failed to resolve destination '%v': %v", dest.FQDN, err)
		}
		ctx = ctx_
		dest.IP = addr
	}

	// Apply any address rewrites
	req.realDestAddr = req.DestAddr
	if s.config.Rewriter != nil {
		ctx, req.realDestAddr = s.config.Rewriter.Rewrite(ctx, req)
	}

	// Switch on the command
	switch req.Command {
	case ConnectCommand:
		return s.handleConnect(ctx, conn, req)
	case BindCommand:
		return s.handleBind(ctx, conn, req)
	case AssociateCommand:
		return s.handleAssociate(ctx, conn, req)
	default:
		if err := sendReply(conn, commandNotSupported, nil); err != nil {
			return fmt.Errorf("Failed to send reply: %v", err)
		}
		return fmt.Errorf("Unsupported command: %v", req.Command)
	}
}

// handleConnect is used to handle a connect command
func (s *Server) handleConnect(ctx context.Context, conn conn, req *Request) error {
	// Check if this is allowed
	if ctx_, ok := s.config.Rules.Allow(ctx, req); !ok {
		if err := sendReply(conn, ruleFailure, nil); err != nil {
			return fmt.Errorf("Failed to send reply: %v", err)
		}
		return fmt.Errorf("Connect to %v blocked by rules", req.DestAddr)
	} else {
		ctx = ctx_
	}

	// Attempt to connect
	addr := (&net.TCPAddr{IP: req.realDestAddr.IP, Port: req.realDestAddr.Port}).String()
	dial := s.config.Dial
	if dial == nil {
		dial = func(ctx context.Context, net_, addr string) (net.Conn, error) {
			return net.Dial(net_, addr)
		}
	}
	target, err := dial(ctx, "tcp", addr)
	if err != nil {
		msg := err.Error()
		resp := hostUnreachable
		if strings.Contains(msg, "refused") {
			resp = connectionRefused
		} else if strings.Contains(msg, "network is unreachable") {
			resp = networkUnreachable
		}
		if err := sendReply(conn, resp, nil); err != nil {
			return fmt.Errorf("Failed to send reply: %v", err)
		}
		return fmt.Errorf("Connect to %v failed: %v", req.DestAddr, err)
	}
	defer target.Close()

	// Send success
	local := target.LocalAddr().(*net.TCPAddr)
	bind := AddrSpec{IP: local.IP, Port: local.Port}
	if err := sendReply(conn, successReply, &bind); err != nil {
		return fmt.Errorf("Failed to send reply: %v", err)
	}

	// Start proxying
	errCh := make(chan error, 2)
	go proxy(target, req.bufConn, errCh)
	go proxy(conn, target, errCh)

	// Wait
	for i := 0; i < 2; i++ {
		e := <-errCh
		if e != nil {
			// return from this function closes target (and conn).
			return e
		}
	}
	return nil
}

// handleBind is used to handle a connect command
func (s *Server) handleBind(ctx context.Context, conn conn, req *Request) error {
	// Check if this is allowed
	if ctx_, ok := s.config.Rules.Allow(ctx, req); !ok {
		if err := sendReply(conn, ruleFailure, nil); err != nil {
			return fmt.Errorf("Failed to send reply: %v", err)
		}
		return fmt.Errorf("Bind to %v blocked by rules", req.DestAddr)
	} else {
		ctx = ctx_
	}

	// TODO: Support bind
	if err := sendReply(conn, commandNotSupported, nil); err != nil {
		return fmt.Errorf("Failed to send reply: %v", err)
	}
	return nil
}

// handleAssociate is used to handle a connect command
func (s *Server) handleAssociate(ctx context.Context, conn conn, req *Request) error {
	// Check if this is allowed
	if ctx_, ok := s.config.Rules.Allow(ctx, req); !ok {
		if err := sendReply(conn, ruleFailure, nil); err != nil {
			return fmt.Errorf("Failed to send reply: %v", err)
		}
		return fmt.Errorf("Associate to %v blocked by rules", req.DestAddr)
	} else {
		ctx = ctx_
	}

	// TODO: Support associate
	if err := sendReply(conn, commandNotSupported, nil); err != nil {
		return fmt.Errorf("Failed to send reply: %v", err)
	}
	return nil
}

// readAddrSpec is used to read AddrSpec.
// Expects an address type byte, follwed by the address and port
func readAddrSpec(r io.Reader) (*AddrSpec, error) {
	d := &AddrSpec{}

	// Get the address type
	addrType := []byte{0}
	if _, err := r.Read(addrType); err != nil {
		return nil, err
	}

	// Handle on a per type basis
	switch addrType[0] {
	case ipv4Address:
		addr := make([]byte, 4)
		if _, err := io.ReadAtLeast(r, addr, len(addr)); err != nil {
			return nil, err
		}
		d.IP = net.IP(addr)

	case ipv6Address:
		addr := make([]byte, 16)
		if _, err := io.ReadAtLeast(r, addr, len(addr)); err != nil {
			return nil, err
		}
		d.IP = net.IP(addr)

	case fqdnAddress:
		if _, err := r.Read(addrType); err != nil {
			return nil, err
		}
		addrLen := int(addrType[0])
		fqdn := make([]byte, addrLen)
		if _, err := io.ReadAtLeast(r, fqdn, addrLen); err != nil {
			return nil, err
		}
		d.FQDN = string(fqdn)

	default:
		return nil, unrecognizedAddrType
	}

	// Read the port
	port := []byte{0, 0}
	if _, err := io.ReadAtLeast(r, port, 2); err != nil {
		return nil, err
	}
	d.Port = (int(port[0]) << 8) | int(port[1])

	return d, nil
}

// sendReply is used to send a reply message
func sendReply(w io.Writer, resp uint8, addr *AddrSpec) error {
	// Format the address
	var addrType uint8
	var addrBody []byte
	var addrPort uint16
	switch {
	case addr == nil:
		addrType = ipv4Address
		addrBody = []byte{0, 0, 0, 0}
		addrPort = 0

	case addr.FQDN != "":
		addrType = fqdnAddress
		addrBody = append([]byte{byte(len(addr.FQDN))}, addr.FQDN...)
		addrPort = uint16(addr.Port)

	case addr.IP.To4() != nil:
		addrType = ipv4Address
		addrBody = []byte(addr.IP.To4())
		addrPort = uint16(addr.Port)

	case addr.IP.To16() != nil:
		addrType = ipv6Address
		addrBody = []byte(addr.IP.To16())
		addrPort = uint16(addr.Port)

	default:
		return fmt.Errorf("Failed to format address: %v", addr)
	}

	// Format the message
	msg := make([]byte, 6+len(addrBody))
	msg[0] = socks5Version
	msg[1] = resp
	msg[2] = 0 // Reserved
	msg[3] = addrType
	copy(msg[4:], addrBody)
	msg[4+len(addrBody)] = byte(addrPort >> 8)
	msg[4+len(addrBody)+1] = byte(addrPort & 0xff)

	// Send the message
	_, err := w.Write(msg)
	return err
}

// proxy is used to suffle data from src to destination, and sends errors
// down a dedicated channel
func proxy(dst io.Writer, src io.Reader, errCh chan error) {
	_, err := io.Copy(dst, src)
	if tcpConn, ok := dst.(*net.TCPConn); ok {
		tcpConn.CloseWrite()
	}
	errCh <- err
}
<commit_msg>half-close all connections supporting `CloseWrite() error`<commit_after>package socks5

import (
	"fmt"
	"io"
	"net"
	"strings"

	"golang.org/x/net/context"
)

const (
	ConnectCommand   = uint8(1)
	BindCommand      = uint8(2)
	AssociateCommand = uint8(3)
	ipv4Address      = uint8(1)
	fqdnAddress      = uint8(3)
	ipv6Address      = uint8(4)
)

const (
	successReply uint8 = iota
	serverFailure
	ruleFailure
	networkUnreachable
	hostUnreachable
	connectionRefused
	ttlExpired
	commandNotSupported
	addrTypeNotSupported
)

var (
	unrecognizedAddrType = fmt.Errorf("Unrecognized address type")
)

// AddressRewriter is used to rewrite a destination transparently
type AddressRewriter interface {
	Rewrite(ctx context.Context, request *Request) (context.Context, *AddrSpec)
}

// AddrSpec is used to return the target AddrSpec
// which may be specified as IPv4, IPv6, or a FQDN
type AddrSpec struct {
	FQDN string
	IP   net.IP
	Port int
}

func (a *AddrSpec) String() string {
	if a.FQDN != "" {
		return fmt.Sprintf("%s (%s):%d", a.FQDN, a.IP, a.Port)
	}
	return fmt.Sprintf("%s:%d", a.IP, a.Port)
}

// A Request represents request received by a server
type Request struct {
	// Protocol version
	Version uint8
	// Requested command
	Command uint8
	// AuthContext provided during negotiation
	AuthContext *AuthContext
	// AddrSpec of the the network that sent the request
	RemoteAddr *AddrSpec
	// AddrSpec of the desired destination
	DestAddr *AddrSpec
	// AddrSpec of the actual destination (might be affected by rewrite)
	realDestAddr *AddrSpec
	bufConn      io.Reader
}

type conn interface {
	Write([]byte) (int, error)
	RemoteAddr() net.Addr
}

// NewRequest creates a new Request from the tcp connection
func NewRequest(bufConn io.Reader) (*Request, error) {
	// Read the version byte
	header := []byte{0, 0, 0}
	if _, err := io.ReadAtLeast(bufConn, header, 3); err != nil {
		return nil, fmt.Errorf("Failed to get command version: %v", err)
	}

	// Ensure we are compatible
	if header[0] != socks5Version {
		return nil, fmt.Errorf("Unsupported command version: %v", header[0])
	}

	// Read in the destination address
	dest, err := readAddrSpec(bufConn)
	if err != nil {
		return nil, err
	}

	request := &Request{
		Version:  socks5Version,
		Command:  header[1],
		DestAddr: dest,
		bufConn:  bufConn,
	}

	return request, nil
}

// handleRequest is used for request processing after authentication
func (s *Server) handleRequest(req *Request, conn conn) error {
	ctx := context.Background()

	// Resolve the address if we have a FQDN
	dest := req.DestAddr
	if dest.FQDN != "" {
		ctx_, addr, err := s.config.Resolver.Resolve(ctx, dest.FQDN)
		if err != nil {
			if err := sendReply(conn, hostUnreachable, nil); err != nil {
				return fmt.Errorf("Failed to send reply: %v", err)
			}
			return fmt.Errorf("Failed to resolve destination '%v': %v", dest.FQDN, err)
		}
		ctx = ctx_
		dest.IP = addr
	}

	// Apply any address rewrites
	req.realDestAddr = req.DestAddr
	if s.config.Rewriter != nil {
		ctx, req.realDestAddr = s.config.Rewriter.Rewrite(ctx, req)
	}

	// Switch on the command
	switch req.Command {
	case ConnectCommand:
		return s.handleConnect(ctx, conn, req)
	case BindCommand:
		return s.handleBind(ctx, conn, req)
	case AssociateCommand:
		return s.handleAssociate(ctx, conn, req)
	default:
		if err := sendReply(conn, commandNotSupported, nil); err != nil {
			return fmt.Errorf("Failed to send reply: %v", err)
		}
		return fmt.Errorf("Unsupported command: %v", req.Command)
	}
}

// handleConnect is used to handle a connect command
func (s *Server) handleConnect(ctx context.Context, conn conn, req *Request) error {
	// Check if this is allowed
	if ctx_, ok := s.config.Rules.Allow(ctx, req); !ok {
		if err := sendReply(conn, ruleFailure, nil); err != nil {
			return fmt.Errorf("Failed to send reply: %v", err)
		}
		return fmt.Errorf("Connect to %v blocked by rules", req.DestAddr)
	} else {
		ctx = ctx_
	}

	// Attempt to connect
	addr := (&net.TCPAddr{IP: req.realDestAddr.IP, Port: req.realDestAddr.Port}).String()
	dial := s.config.Dial
	if dial == nil {
		dial = func(ctx context.Context, net_, addr string) (net.Conn, error) {
			return net.Dial(net_, addr)
		}
	}
	target, err := dial(ctx, "tcp", addr)
	if err != nil {
		msg := err.Error()
		resp := hostUnreachable
		if strings.Contains(msg, "refused") {
			resp = connectionRefused
		} else if strings.Contains(msg, "network is unreachable") {
			resp = networkUnreachable
		}
		if err := sendReply(conn, resp, nil); err != nil {
			return fmt.Errorf("Failed to send reply: %v", err)
		}
		return fmt.Errorf("Connect to %v failed: %v", req.DestAddr, err)
	}
	defer target.Close()

	// Send success
	local := target.LocalAddr().(*net.TCPAddr)
	bind := AddrSpec{IP: local.IP, Port: local.Port}
	if err := sendReply(conn, successReply, &bind); err != nil {
		return fmt.Errorf("Failed to send reply: %v", err)
	}

	// Start proxying
	errCh := make(chan error, 2)
	go proxy(target, req.bufConn, errCh)
	go proxy(conn, target, errCh)

	// Wait
	for i := 0; i < 2; i++ {
		e := <-errCh
		if e != nil {
			// return from this function closes target (and conn).
			return e
		}
	}
	return nil
}

// handleBind is used to handle a connect command
func (s *Server) handleBind(ctx context.Context, conn conn, req *Request) error {
	// Check if this is allowed
	if ctx_, ok := s.config.Rules.Allow(ctx, req); !ok {
		if err := sendReply(conn, ruleFailure, nil); err != nil {
			return fmt.Errorf("Failed to send reply: %v", err)
		}
		return fmt.Errorf("Bind to %v blocked by rules", req.DestAddr)
	} else {
		ctx = ctx_
	}

	// TODO: Support bind
	if err := sendReply(conn, commandNotSupported, nil); err != nil {
		return fmt.Errorf("Failed to send reply: %v", err)
	}
	return nil
}

// handleAssociate is used to handle a connect command
func (s *Server) handleAssociate(ctx context.Context, conn conn, req *Request) error {
	// Check if this is allowed
	if ctx_, ok := s.config.Rules.Allow(ctx, req); !ok {
		if err := sendReply(conn, ruleFailure, nil); err != nil {
			return fmt.Errorf("Failed to send reply: %v", err)
		}
		return fmt.Errorf("Associate to %v blocked by rules", req.DestAddr)
	} else {
		ctx = ctx_
	}

	// TODO: Support associate
	if err := sendReply(conn, commandNotSupported, nil); err != nil {
		return fmt.Errorf("Failed to send reply: %v", err)
	}
	return nil
}

// readAddrSpec is used to read AddrSpec.
// Expects an address type byte, follwed by the address and port
func readAddrSpec(r io.Reader) (*AddrSpec, error) {
	d := &AddrSpec{}

	// Get the address type
	addrType := []byte{0}
	if _, err := r.Read(addrType); err != nil {
		return nil, err
	}

	// Handle on a per type basis
	switch addrType[0] {
	case ipv4Address:
		addr := make([]byte, 4)
		if _, err := io.ReadAtLeast(r, addr, len(addr)); err != nil {
			return nil, err
		}
		d.IP = net.IP(addr)

	case ipv6Address:
		addr := make([]byte, 16)
		if _, err := io.ReadAtLeast(r, addr, len(addr)); err != nil {
			return nil, err
		}
		d.IP = net.IP(addr)

	case fqdnAddress:
		if _, err := r.Read(addrType); err != nil {
			return nil, err
		}
		addrLen := int(addrType[0])
		fqdn := make([]byte, addrLen)
		if _, err := io.ReadAtLeast(r, fqdn, addrLen); err != nil {
			return nil, err
		}
		d.FQDN = string(fqdn)

	default:
		return nil, unrecognizedAddrType
	}

	// Read the port
	port := []byte{0, 0}
	if _, err := io.ReadAtLeast(r, port, 2); err != nil {
		return nil, err
	}
	d.Port = (int(port[0]) << 8) | int(port[1])

	return d, nil
}

// sendReply is used to send a reply message
func sendReply(w io.Writer, resp uint8, addr *AddrSpec) error {
	// Format the address
	var addrType uint8
	var addrBody []byte
	var addrPort uint16
	switch {
	case addr == nil:
		addrType = ipv4Address
		addrBody = []byte{0, 0, 0, 0}
		addrPort = 0

	case addr.FQDN != "":
		addrType = fqdnAddress
		addrBody = append([]byte{byte(len(addr.FQDN))}, addr.FQDN...)
		addrPort = uint16(addr.Port)

	case addr.IP.To4() != nil:
		addrType = ipv4Address
		addrBody = []byte(addr.IP.To4())
		addrPort = uint16(addr.Port)

	case addr.IP.To16() != nil:
		addrType = ipv6Address
		addrBody = []byte(addr.IP.To16())
		addrPort = uint16(addr.Port)

	default:
		return fmt.Errorf("Failed to format address: %v", addr)
	}

	// Format the message
	msg := make([]byte, 6+len(addrBody))
	msg[0] = socks5Version
	msg[1] = resp
	msg[2] = 0 // Reserved
	msg[3] = addrType
	copy(msg[4:], addrBody)
	msg[4+len(addrBody)] = byte(addrPort >> 8)
	msg[4+len(addrBody)+1] = byte(addrPort & 0xff)

	// Send the message
	_, err := w.Write(msg)
	return err
}

type closeWriter interface {
	CloseWrite() error
}

// proxy is used to suffle data from src to destination, and sends errors
// down a dedicated channel
func proxy(dst io.Writer, src io.Reader, errCh chan error) {
	_, err := io.Copy(dst, src)
	if tcpConn, ok := dst.(closeWriter); ok {
		tcpConn.CloseWrite()
	}
	errCh <- err
}
<|endoftext|>


<commit_before>package hrq

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"net/textproto"
	"net/url"
	"os"
	"strings"
	"time"
)

// DefaultTimeout is seconds of timeout.
var DefaultTimeout = 15

// DefaultContentType is a default content-type of request.
var DefaultContentType = "application/x-www-form-urlencoded"

// File is file for multipart/form.
type File struct {
	ContentType string
	FieldName   string
	Name        string
	File        *os.File
}

// Request inherits http.Request.
type Request struct {
	*http.Request
	Timeout time.Duration
	Data    interface{}
	Files   []*File
}

func (r *Request) isPostOrPut() bool {
	return r.Method == "POST" || r.Method == "PUT"
}

func (r *Request) setBody(values *strings.Reader) {
	body := ioutil.NopCloser(values)
	r.Body = body
	r.ContentLength = int64(values.Len())
	s := *values
	r.GetBody = func() (io.ReadCloser, error) {
		r := s
		return ioutil.NopCloser(&r), nil
	}
}

// AddFile sets file for multipart/form.
func (r *Request) AddFile(contentType, fieldName, fileName string, file *os.File) *Request {
	f := &File{
		ContentType: contentType,
		FieldName:   fieldName,
		Name:        fileName,
		File:        file,
	}
	r.Files = append(r.Files, f)
	return r
}

// SetTimeout sets timeout.
func (r *Request) SetTimeout(timeout int) *Request {
	r.Timeout = time.Duration(timeout) * time.Second
	return r
}

// Send sends request.
// If method is POST and content-type is application/x-www-form-urlencoded,
// the request data is urlencoded.
// If method is POST and content-type is application/json,
// the request data is converted to json string.
func (r *Request) Send() (res *Response, err error) {
	if r.isPostOrPut() && r.Data != nil && r.HeaderValue("Content-Type") != "multipart/form-data" {
		if r.HeaderValue("Content-Type") == "application/x-www-form-urlencoded" {
			data, ok := r.Data.(map[string][]string)
			if !ok {
				err := errors.New("data is not a map[string][]string at Request.Send()")
				return nil, err
			}
			values := strings.NewReader(url.Values(data).Encode())
			r.setBody(values)
		} else if r.HeaderValue("Content-Type") == "application/json" {
			jsonBytes, err := json.Marshal(r.Data)
			if err != nil {
				return nil, err
			}
			values := strings.NewReader(string(jsonBytes))
			r.setBody(values)
		}
	} else if r.isPostOrPut() && r.HeaderValue("Content-Type") == "multipart/form-data" {
		var buffer bytes.Buffer
		writer := multipart.NewWriter(&buffer)
		data, ok := r.Data.(map[string]string)
		if !ok {
			err := errors.New("data is not a map[string]string at Request.Send()")
			return nil, err
		}
		for k, v := range data {
			writer.WriteField(k, v)
		}
		for _, file := range r.Files {
			part := make(textproto.MIMEHeader)
			part.Set("Content-Type", file.ContentType)
			desc := fmt.Sprintf(`form-data; name="%s"; filename="%s"`, file.FieldName, file.Name)
			part.Set("Content-Disposition", desc)
			fileWriter, err := writer.CreatePart(part)
			if err != nil {
				return nil, err
			}
			_, err = io.Copy(fileWriter, file.File)
			if err != nil {
				return nil, err
			}
			defer file.File.Close()
		}
		writer.Close()
		r.SetHeader("Content-Type", writer.FormDataContentType())
		r.ContentLength = int64(buffer.Len())
		b := buffer.Bytes()
		r.GetBody = func() (io.ReadCloser, error) {
			reader := bytes.NewReader(b)
			return ioutil.NopCloser(reader), nil
		}
		reader := bytes.NewReader(b)
		r.Body = ioutil.NopCloser(reader)
	}

	cli := &http.Client{
		Timeout: r.Timeout,
	}
	response, err := cli.Do(r.Request)
	if err != nil {
		return
	}
	res = &Response{
		Response: response,
	}
	return
}

// SetHeader sets a value of request header.
func (r *Request) SetHeader(key, value string) *Request {
	r.Header.Set(key, value)
	return r
}

// HeaderValue returns a value of request header.
func (r *Request) HeaderValue(key string) string {
	return r.Header.Get(key)
}

// DelHeader delete a value of request header by key.
func (r *Request) DelHeader(key string) *Request {
	r.Header.Del(key)
	return r
}

// PutCookie makes a cookie which is setted name and value.
// It adds a cookie to a request.
func (r *Request) PutCookie(name, value string) *Request {
	r.AddCookie(&http.Cookie{Name: name, Value: value})
	return r
}

// NewRequest make a Request.
func NewRequest(method, url string, body io.Reader, timeoutSecond int) (req *Request, err error) {
	request, err := http.NewRequest(method, url, nil)
	if err != nil {
		return
	}
	timeout := time.Duration(timeoutSecond) * time.Second
	req = &Request{
		Request: request,
		Timeout: timeout,
	}
	return
}

// Get make a request whose method is GET.
func Get(url string) (req *Request, err error) {
	req, err = NewRequest("GET", url, nil, DefaultTimeout)
	return
}

func postOrPut(method, url string, data interface{}) (req *Request, err error) {
	req, err = NewRequest(method, url, nil, DefaultTimeout)
	if err != nil {
		return
	}
	req.SetHeader("Content-Type", DefaultContentType)
	req.Files = []*File{}
	req.Data = data
	return
}

// Post make a request whose method is POST.
func Post(url string, data interface{}) (req *Request, err error) {
	req, err = postOrPut("POST", url, data)
	return
}

// Put make a request whose method is PUT.
func Put(url string, data interface{}) (req *Request, err error) {
	req, err = postOrPut("PUT", url, data)
	return
}
<commit_msg>Add methods (#39)<commit_after>package hrq

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"net/textproto"
	"net/url"
	"os"
	"strings"
	"time"
)

// DefaultTimeout is seconds of timeout.
var DefaultTimeout = 15

// DefaultContentType is a default content-type of request.
var DefaultContentType = "application/x-www-form-urlencoded"

// File is file for multipart/form.
type File struct {
	ContentType string
	FieldName   string
	Name        string
	File        *os.File
}

// Request inherits http.Request.
type Request struct {
	*http.Request
	Timeout time.Duration
	Data    interface{}
	Files   []*File
}

func (r *Request) isPostOrPut() bool {
	return r.Method == "POST" || r.Method == "PUT"
}

func (r *Request) setBody(values *strings.Reader) {
	body := ioutil.NopCloser(values)
	r.Body = body
	r.ContentLength = int64(values.Len())
	s := *values
	r.GetBody = func() (io.ReadCloser, error) {
		r := s
		return ioutil.NopCloser(&r), nil
	}
}

// AddFile sets file for multipart/form.
func (r *Request) AddFile(contentType, fieldName, fileName string, file *os.File) *Request {
	f := &File{
		ContentType: contentType,
		FieldName:   fieldName,
		Name:        fileName,
		File:        file,
	}
	r.Files = append(r.Files, f)
	return r
}

// SetTimeout sets timeout.
func (r *Request) SetTimeout(timeout int) *Request {
	r.Timeout = time.Duration(timeout) * time.Second
	return r
}

// Send sends request.
// If method is POST and content-type is application/x-www-form-urlencoded,
// the request data is urlencoded.
// If method is POST and content-type is application/json,
// the request data is converted to json string.
func (r *Request) Send() (res *Response, err error) {
	if r.isPostOrPut() && r.Data != nil && r.HeaderValue("Content-Type") != "multipart/form-data" {
		if r.HeaderValue("Content-Type") == "application/x-www-form-urlencoded" {
			data, ok := r.Data.(map[string][]string)
			if !ok {
				err := errors.New("data is not a map[string][]string at Request.Send()")
				return nil, err
			}
			values := strings.NewReader(url.Values(data).Encode())
			r.setBody(values)
		} else if r.HeaderValue("Content-Type") == "application/json" {
			jsonBytes, err := json.Marshal(r.Data)
			if err != nil {
				return nil, err
			}
			values := strings.NewReader(string(jsonBytes))
			r.setBody(values)
		}
	} else if r.isPostOrPut() && r.HeaderValue("Content-Type") == "multipart/form-data" {
		var buffer bytes.Buffer
		writer := multipart.NewWriter(&buffer)
		data, ok := r.Data.(map[string]string)
		if !ok {
			err := errors.New("data is not a map[string]string at Request.Send()")
			return nil, err
		}
		for k, v := range data {
			writer.WriteField(k, v)
		}
		for _, file := range r.Files {
			part := make(textproto.MIMEHeader)
			part.Set("Content-Type", file.ContentType)
			desc := fmt.Sprintf(`form-data; name="%s"; filename="%s"`, file.FieldName, file.Name)
			part.Set("Content-Disposition", desc)
			fileWriter, err := writer.CreatePart(part)
			if err != nil {
				return nil, err
			}
			_, err = io.Copy(fileWriter, file.File)
			if err != nil {
				return nil, err
			}
			defer file.File.Close()
		}
		writer.Close()
		r.SetHeader("Content-Type", writer.FormDataContentType())
		r.ContentLength = int64(buffer.Len())
		b := buffer.Bytes()
		r.GetBody = func() (io.ReadCloser, error) {
			reader := bytes.NewReader(b)
			return ioutil.NopCloser(reader), nil
		}
		reader := bytes.NewReader(b)
		r.Body = ioutil.NopCloser(reader)
	}

	cli := &http.Client{
		Timeout: r.Timeout,
	}
	response, err := cli.Do(r.Request)
	if err != nil {
		return
	}
	res = &Response{
		Response: response,
	}
	return
}

// SetHeader sets a value of request header.
func (r *Request) SetHeader(key, value string) *Request {
	r.Header.Set(key, value)
	return r
}

// HeaderValue returns a value of request header.
func (r *Request) HeaderValue(key string) string {
	return r.Header.Get(key)
}

// DelHeader delete a value of request header by key.
func (r *Request) DelHeader(key string) *Request {
	r.Header.Del(key)
	return r
}

// PutCookie makes a cookie which is setted name and value.
// It adds a cookie to a request.
func (r *Request) PutCookie(name, value string) *Request {
	r.AddCookie(&http.Cookie{Name: name, Value: value})
	return r
}

// NewRequest make a Request.
func NewRequest(method, url string, body io.Reader, timeoutSecond int) (req *Request, err error) {
	request, err := http.NewRequest(method, url, nil)
	if err != nil {
		return
	}
	timeout := time.Duration(timeoutSecond) * time.Second
	req = &Request{
		Request: request,
		Timeout: timeout,
	}
	return
}

// Get make a request whose method is GET.
func Get(url string) (req *Request, err error) {
	req, err = NewRequest("GET", url, nil, DefaultTimeout)
	return
}

// Delete make a request whose method is DELETE.
func Delete(url string) (req *Request, err error) {
	req, err = NewRequest("DELETE", url, nil, DefaultTimeout)
	return
}

// Head make a request whose method is HEAD.
func Head(url string) (req *Request, err error) {
	req, err = NewRequest("HEAD", url, nil, DefaultTimeout)
	return
}

// Options make a request whose method is OPTIONS.
func Options(url string) (req *Request, err error) {
	req, err = NewRequest("OPTIONS", url, nil, DefaultTimeout)
	return
}

func postOrPut(method, url string, data interface{}) (req *Request, err error) {
	req, err = NewRequest(method, url, nil, DefaultTimeout)
	if err != nil {
		return
	}
	req.SetHeader("Content-Type", DefaultContentType)
	req.Files = []*File{}
	req.Data = data
	return
}

// Post make a request whose method is POST.
func Post(url string, data interface{}) (req *Request, err error) {
	req, err = postOrPut("POST", url, data)
	return
}

// Put make a request whose method is PUT.
func Put(url string, data interface{}) (req *Request, err error) {
	req, err = postOrPut("PUT", url, data)
	return
}
<|endoftext|>


<commit_before>// Copyright 2016 Marcel Gotsch. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package goserv

import (
	"encoding/json"
	"net/http"
)

// A Request represents an HTTP request received by the Server.
//
// It embeds the native http.Request, thus all native fields are still available
// through Request. Every Request has it's own Context providing a key-value store to share
// data between multiple Handlers. In case that the Route handling the Request has parameters, the parameter
// values are extracted from the Request's path and stored in .Params.
type Request struct {
	// Embedded http.Request.
	*http.Request

	// Request specific key-value store to share data between Handlers
	Context *Context

	// Key-value store containing named parameter values extracted from
	// the Request's path. See Route.
	Params Params

	// Sanitized http.Request.URL.Path
	SanitizedPath string
}

// JSON parses the request's body using the encoding/json Decoder. In case
// of a decoding error the error is returned.
//
// Note: The request's body is closed after calling this method.
func (r *Request) JSON(v interface{}) error {
	err := json.NewDecoder(r.Body).Decode(v)
	r.Body.Close()

	if err != nil {
		return err
	}

	return nil
}

func newRequest(r *http.Request) *Request {
	return &Request{r, newContext(), make(Params), SanitizePath(r.URL.Path)}
}
<commit_msg>Rename Request.JSON to .ReadJSON<commit_after>// Copyright 2016 Marcel Gotsch. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package goserv

import (
	"encoding/json"
	"net/http"
)

// A Request represents an HTTP request received by the Server.
//
// It embeds the native http.Request, thus all native fields are still available
// through Request. Every Request has it's own Context providing a key-value store to share
// data between multiple Handlers. In case that the Route handling the Request has parameters, the parameter
// values are extracted from the Request's path and stored in .Params.
type Request struct {
	// Embedded http.Request.
	*http.Request

	// Request specific key-value store to share data between Handlers
	Context *Context

	// Key-value store containing named parameter values extracted from
	// the Request's path. See Route.
	Params Params

	// Sanitized http.Request.URL.Path
	SanitizedPath string
}

// ReadJSON parses the request's body using the encoding/json Decoder. In case
// of a decoding error the error is returned.
//
// Note: The request's body is closed after calling this method.
func (r *Request) ReadJSON(v interface{}) error {
	err := json.NewDecoder(r.Body).Decode(v)
	r.Body.Close()

	if err != nil {
		return err
	}

	return nil
}

func newRequest(r *http.Request) *Request {
	return &Request{r, newContext(), make(Params), SanitizePath(r.URL.Path)}
}
<|endoftext|>


<commit_before>package main

import (
	"net/http"
	"net/url"
	"path"
	"strings"

	"github.com/boltdb/bolt"
)

type Request struct {
	Service string
	Method  string
	Path    string
	Data    string

	Settings  Settings
	NoQueries bool
	NoHeaders bool

	URL url.URL
}

// Prepare the http request.  This will substitute all the parameters,
// addd all the headers and query parameters
func (r *Request) Prepare() (*http.Request, error) {
	// prepare the url
	r.URL = r.Settings.URL()
	params := paramReplacer(r.Settings.Parameters)

	r.URL.Path = path.Join(r.Settings.BasePath.String, params.Replace(r.Path))
	r.Data = params.Replace(r.Data)

	if !r.NoQueries {
		q := r.URL.Query()
		for key, value := range r.Settings.Queries {
			q.Set(params.Replace(key), params.Replace(value))
		}
		r.URL.RawQuery = q.Encode()
	}

	// prepare the request
	req, err := http.NewRequest(
		strings.ToUpper(r.Method),
		r.URL.String(),
		strings.NewReader(r.Data),
	)
	if err != nil {
		return nil, err
	}

	if r.Settings.Username.Valid && r.Settings.Password.Valid {
		req.SetBasicAuth(r.Settings.Username.String, r.Settings.Password.String)
	}

	if !r.NoHeaders {
		for key, value := range r.Settings.Headers {
			req.Header.Set(params.Replace(key), params.Replace(value))
		}
	}

	return req, nil
}

// ServiceBucket returns the db bucket for the requests service
func (r *Request) ServiceBucket(tx *bolt.Tx) (*bolt.Bucket, error) {
	if r.Service == "" {
		info := tx.Bucket([]byte("info"))
		current := info.Get([]byte("current"))
		r.Service = string(current)
	}

	sb, err := tx.CreateBucketIfNotExists([]byte("services"))
	if err != nil {
		return nil, err
	}

	b, err := sb.CreateBucketIfNotExists([]byte(r.Service))
	if err != nil {
		return nil, err
	}

	return b, nil
}

// PathBucket returns the bucket for the request path, creates if needed
func (r Request) PathBucket(tx *bolt.Tx) (*bolt.Bucket, error) {
	s, err := r.ServiceBucket(tx)
	if err != nil {
		return nil, err
	}

	pb, err := s.CreateBucketIfNotExists([]byte("paths"))
	if err != nil {
		return nil, err
	}

	b, err := pb.CreateBucketIfNotExists([]byte(r.Path))
	if err != nil {
		return nil, err
	}

	return b, nil
}

// MethodBucket returns the bucket for the request method, creates if needed
func (r Request) MethodBucket(tx *bolt.Tx) (*bolt.Bucket, error) {
	s, err := r.PathBucket(tx)
	if err != nil {
		return nil, err
	}

	b, err := s.CreateBucketIfNotExists([]byte(r.Method))
	if err != nil {
		return nil, err
	}

	return b, nil
}

// LoadSettings from the database
func (r *Request) LoadSettings(tx *bolt.Tx) error {
	sb, pb, mb, err := request.Match(tx)
	if err != nil {
		return err
	}

	// Start with blank settings
	r.Settings = NewSettings()

	// no service based settings
	if sb == nil {
		r.Settings.Merge(settings)
		return nil
	}

	// load service settings as base
	r.Settings = LoadSettings(sb)

	if pb == nil {
		r.Settings.Merge(settings)
		return nil
	}

	r.Settings.Merge(LoadSettings(pb))

	if mb == nil {
		r.Settings.Merge(settings)
		return nil
	}

	r.Settings.Merge(LoadSettings(mb))
	r.Settings.Merge(settings)

	return nil
}

// Match returns the relavant db buckets for all requst settings
func (r Request) Match(tx *bolt.Tx) (service, path, method *bolt.Bucket, err error) {
	service, err = r.ServiceBucket(tx)
	if err != nil {
		return nil, nil, nil, err
	}

	pb := service.Bucket([]byte("paths"))
	if pb == nil {
		return service, nil, nil, nil
	}

	// Match the path, this will eventually be expanded to better match
	// the path.  So specific paths will be matched before generic ones
	// at the moment it requires an exact match
	c := pb.Cursor()
	for key, _ := c.First(); key != nil; key, _ = c.Next() {
		if string(key) == r.Path {
			path = pb.Bucket(key)
			break
		}
	}

	if path == nil {
		return service, path, nil, nil
	}

	method = path.Bucket([]byte(r.Method))

	return service, path, method, nil
}
<commit_msg>Simplify loading nested settings into Request<commit_after>package main

import (
	"net/http"
	"net/url"
	"path"
	"strings"

	"github.com/boltdb/bolt"
)

type Request struct {
	Service string
	Method  string
	Path    string
	Data    string

	Settings  Settings
	NoQueries bool
	NoHeaders bool

	URL url.URL
}

// Prepare the http request.  This will substitute all the parameters,
// addd all the headers and query parameters
func (r *Request) Prepare() (*http.Request, error) {
	// prepare the url
	r.URL = r.Settings.URL()
	params := paramReplacer(r.Settings.Parameters)

	r.URL.Path = path.Join(r.Settings.BasePath.String, params.Replace(r.Path))
	r.Data = params.Replace(r.Data)

	if !r.NoQueries {
		q := r.URL.Query()
		for key, value := range r.Settings.Queries {
			q.Set(params.Replace(key), params.Replace(value))
		}
		r.URL.RawQuery = q.Encode()
	}

	// prepare the request
	req, err := http.NewRequest(
		strings.ToUpper(r.Method),
		r.URL.String(),
		strings.NewReader(r.Data),
	)
	if err != nil {
		return nil, err
	}

	if r.Settings.Username.Valid && r.Settings.Password.Valid {
		req.SetBasicAuth(r.Settings.Username.String, r.Settings.Password.String)
	}

	if !r.NoHeaders {
		for key, value := range r.Settings.Headers {
			req.Header.Set(params.Replace(key), params.Replace(value))
		}
	}

	return req, nil
}

// ServiceBucket returns the db bucket for the requests service
func (r *Request) ServiceBucket(tx *bolt.Tx) (*bolt.Bucket, error) {
	if r.Service == "" {
		info := tx.Bucket([]byte("info"))
		current := info.Get([]byte("current"))
		r.Service = string(current)
	}

	sb, err := tx.CreateBucketIfNotExists([]byte("services"))
	if err != nil {
		return nil, err
	}

	b, err := sb.CreateBucketIfNotExists([]byte(r.Service))
	if err != nil {
		return nil, err
	}

	return b, nil
}

// PathBucket returns the bucket for the request path, creates if needed
func (r Request) PathBucket(tx *bolt.Tx) (*bolt.Bucket, error) {
	s, err := r.ServiceBucket(tx)
	if err != nil {
		return nil, err
	}

	pb, err := s.CreateBucketIfNotExists([]byte("paths"))
	if err != nil {
		return nil, err
	}

	b, err := pb.CreateBucketIfNotExists([]byte(r.Path))
	if err != nil {
		return nil, err
	}

	return b, nil
}

// MethodBucket returns the bucket for the request method, creates if needed
func (r Request) MethodBucket(tx *bolt.Tx) (*bolt.Bucket, error) {
	s, err := r.PathBucket(tx)
	if err != nil {
		return nil, err
	}

	b, err := s.CreateBucketIfNotExists([]byte(r.Method))
	if err != nil {
		return nil, err
	}

	return b, nil
}

// LoadSettings from the database
func (r *Request) LoadSettings(tx *bolt.Tx) error {
	sb, pb, mb, err := request.Match(tx)
	if err != nil {
		return err
	}

	// Start with blank settings
	r.Settings = NewSettings()

	// load service settings
	if sb != nil {
		r.Settings = LoadSettings(sb)
	}

	// load path settings
	if pb != nil {
		r.Settings.Merge(LoadSettings(pb))
	}

	// load method settings
	if mb != nil {
		r.Settings.Merge(LoadSettings(mb))
	}

	// load provided cli flags settings
	r.Settings.Merge(settings)

	return nil
}

// Match returns the relavant db buckets for all requst settings
func (r Request) Match(tx *bolt.Tx) (service, path, method *bolt.Bucket, err error) {
	service, err = r.ServiceBucket(tx)
	if err != nil {
		return nil, nil, nil, err
	}

	pb := service.Bucket([]byte("paths"))
	if pb == nil {
		return service, nil, nil, nil
	}

	// Match the path, this will eventually be expanded to better match
	// the path.  So specific paths will be matched before generic ones
	// at the moment it requires an exact match
	c := pb.Cursor()
	for key, _ := c.First(); key != nil; key, _ = c.Next() {
		if string(key) == r.Path {
			path = pb.Bucket(key)
			break
		}
	}

	if path == nil {
		return service, path, nil, nil
	}

	method = path.Bucket([]byte(r.Method))

	return service, path, method, nil
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"koding/artifact"
	"koding/ec2dynamicdata"
	"koding/kites/common"
	"koding/kites/kloud/pkg/dnsclient"
	"os"

	"github.com/koding/kite"
	"github.com/koding/logging"
	"github.com/koding/multiconfig"
	"github.com/koding/tunnel"
	"github.com/mitchellh/goamz/aws"
)

type registerResult struct {
	VirtualHost string
	Identifier  string
}

type tunnelServer struct {
	Port            int
	Debug           bool
	BaseVirtualHost string `required:"true"`

	// ServerAddr is public Address of the running server. Like an assigned Elastic IP
	ServerAddr string

	HostedZone string
	AccessKey  string
	SecretKey  string

	// internal
	Log    logging.Logger
	Server *tunnel.Server
	Dns    *dnsclient.Route53
}

func main() {
	t := new(tunnelServer)
	m := multiconfig.New()
	m.MustLoad(t)
	m.MustValidate(t)

	auth := aws.Auth{
		AccessKey: t.AccessKey,
		SecretKey: t.SecretKey,
	}

	// get from ec2 meta-data if it's not passed explicitly
	if t.ServerAddr == "" {
		var err error
		t.ServerAddr, err = ec2dynamicdata.GetMetadata(ec2dynamicdata.PublicIPv4)
		if err != nil {
			fmt.Fprintf(os.Stderr, err.Error())
			os.Exit(1)
		}
	}

	var err error
	t.Server, err = tunnel.NewServer(&tunnel.ServerConfig{
		Debug: t.Debug,
	})
	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}

	t.Dns = dnsclient.NewRoute53Client(t.HostedZone, auth)
	t.Log = common.NewLogger("tunnelkite", t.Debug)

	k := kite.New("tunnelserver", "0.0.1")
	k.Config.DisableAuthentication = true
	k.Config.Port = t.Port

	k.HandleFunc("register", t.Register)
	k.HandleHTTPFunc("/healthCheck", artifact.HealthCheckHandler(Name))
	k.HandleHTTPFunc("/version", artifact.VersionHandler())
	k.HandleHTTP("/{rest:.*}", t.Server)

	k.Run()
}

func (t *tunnelServer) Register(r *kite.Request) (interface{}, error) {
	virtualHost := fmt.Sprintf("%s.%s", r.Username, t.BaseVirtualHost)
	identifier := randomID(32)

	domain := r.Username + "." + t.Dns.HostedZone()

	t.Log.Debug("Adding domain '%s' with to IP Address '%s'", domain, t.ServerAddr)
	if err := t.Dns.Upsert(domain, t.ServerAddr); err != nil {
		return nil, err
	}

	t.Log.Debug("Adding host '%s' with identifer '%s'", virtualHost, identifier)
	t.Server.AddHost(virtualHost, identifier)

	// cleanup domain and delete in memory virtualhost
	t.Server.OnDisconnect(identifier, func() error {
		t.Log.Debug("Deleting host '%s' and domain '%s'", virtualHost, domain)
		t.Server.DeleteHost(virtualHost)
		return t.Dns.Delete(domain)
	})

	return registerResult{
		VirtualHost: virtualHost,
		Identifier:  identifier,
	}, nil
}

// randomID generates a random string of the given length
func randomID(length int) string {
	r := make([]byte, length*6/8)
	rand.Read(r)
	return base64.URLEncoding.EncodeToString(r)
}
<commit_msg>Tunnelproxy: added forgotten Name varible<commit_after>package main

import (
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"koding/artifact"
	"koding/ec2dynamicdata"
	"koding/kites/common"
	"koding/kites/kloud/pkg/dnsclient"
	"os"

	"github.com/koding/kite"
	"github.com/koding/logging"
	"github.com/koding/multiconfig"
	"github.com/koding/tunnel"
	"github.com/mitchellh/goamz/aws"
)

const Name = "tunnelserver"

type registerResult struct {
	VirtualHost string
	Identifier  string
}

type tunnelServer struct {
	Port            int
	Debug           bool
	BaseVirtualHost string `required:"true"`

	// ServerAddr is public Address of the running server. Like an assigned Elastic IP
	ServerAddr string

	HostedZone string
	AccessKey  string
	SecretKey  string

	// internal
	Log    logging.Logger
	Server *tunnel.Server
	Dns    *dnsclient.Route53
}

func main() {
	t := new(tunnelServer)
	m := multiconfig.New()
	m.MustLoad(t)
	m.MustValidate(t)

	auth := aws.Auth{
		AccessKey: t.AccessKey,
		SecretKey: t.SecretKey,
	}

	// get from ec2 meta-data if it's not passed explicitly
	if t.ServerAddr == "" {
		var err error
		t.ServerAddr, err = ec2dynamicdata.GetMetadata(ec2dynamicdata.PublicIPv4)
		if err != nil {
			fmt.Fprintf(os.Stderr, err.Error())
			os.Exit(1)
		}
	}

	var err error
	t.Server, err = tunnel.NewServer(&tunnel.ServerConfig{
		Debug: t.Debug,
	})
	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}

	t.Dns = dnsclient.NewRoute53Client(t.HostedZone, auth)
	t.Log = common.NewLogger("tunnelkite", t.Debug)

	k := kite.New(Name, "0.0.1")
	k.Config.DisableAuthentication = true
	k.Config.Port = t.Port

	k.HandleFunc("register", t.Register)
	k.HandleHTTPFunc("/healthCheck", artifact.HealthCheckHandler(Name))
	k.HandleHTTPFunc("/version", artifact.VersionHandler())
	k.HandleHTTP("/{rest:.*}", t.Server)

	k.Run()
}

func (t *tunnelServer) Register(r *kite.Request) (interface{}, error) {
	virtualHost := fmt.Sprintf("%s.%s", r.Username, t.BaseVirtualHost)
	identifier := randomID(32)

	domain := r.Username + "." + t.Dns.HostedZone()

	t.Log.Debug("Adding domain '%s' with to IP Address '%s'", domain, t.ServerAddr)
	if err := t.Dns.Upsert(domain, t.ServerAddr); err != nil {
		return nil, err
	}

	t.Log.Debug("Adding host '%s' with identifer '%s'", virtualHost, identifier)
	t.Server.AddHost(virtualHost, identifier)

	// cleanup domain and delete in memory virtualhost
	t.Server.OnDisconnect(identifier, func() error {
		t.Log.Debug("Deleting host '%s' and domain '%s'", virtualHost, domain)
		t.Server.DeleteHost(virtualHost)
		return t.Dns.Delete(domain)
	})

	return registerResult{
		VirtualHost: virtualHost,
		Identifier:  identifier,
	}, nil
}

// randomID generates a random string of the given length
func randomID(length int) string {
	r := make([]byte, length*6/8)
	rand.Read(r)
	return base64.URLEncoding.EncodeToString(r)
}
<|endoftext|>


<commit_before>// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package testlib

import (
	"fmt"
	"strings"
	"time"

	"github.com/youtube/vitess/go/vt/logutil"
	myproto "github.com/youtube/vitess/go/vt/mysqlctl/proto"
	"github.com/youtube/vitess/go/vt/tabletmanager/tmclient"
	"github.com/youtube/vitess/go/vt/topo"
	"github.com/youtube/vitess/go/vt/wrangler"
	"github.com/youtube/vitess/go/vt/zktopo"
	"golang.org/x/net/context"

	"testing"
)

// TestInitMasterShard is the good scenario test, where everything
// works as planned
func TestInitMasterShard(t *testing.T) {
	ctx := context.Background()
	ts := zktopo.NewTestServer(t, []string{"cell1", "cell2"})
	wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient(), time.Second)

	// Create a master, a couple good slaves
	master := NewFakeTablet(t, wr, "cell1", 0, topo.TYPE_MASTER)
	goodSlave1 := NewFakeTablet(t, wr, "cell1", 1, topo.TYPE_REPLICA)
	goodSlave2 := NewFakeTablet(t, wr, "cell2", 2, topo.TYPE_REPLICA)

	// Master: set a plausible ReplicationPosition to return,
	// and expect to add entry in _vt.reparent_journal
	master.FakeMysqlDaemon.CurrentMasterPosition = myproto.ReplicationPosition{
		GTIDSet: myproto.MariadbGTID{
			Domain:   5,
			Server:   456,
			Sequence: 890,
		},
	}
	master.FakeMysqlDaemon.ReadOnly = true
	master.FakeMysqlDaemon.ExpectedExecuteSuperQueryList = []string{
		"CREATE DATABASE IF NOT EXISTS _vt",
		"SUBCREATE TABLE IF NOT EXISTS _vt.reparent_journal",
		"SUBINSERT INTO _vt.reparent_journal (time_created_ns, action_name, master_alias, replication_position) VALUES",
	}
	master.StartActionLoop(t, wr)
	defer master.StopActionLoop(t)

	// Slave1: expect to be re-parented
	goodSlave1.FakeMysqlDaemon.ReadOnly = true
	goodSlave1.FakeMysqlDaemon.StartReplicationCommandsStatus = &myproto.ReplicationStatus{
		Position:           master.FakeMysqlDaemon.CurrentMasterPosition,
		MasterHost:         master.Tablet.Hostname,
		MasterPort:         master.Tablet.Portmap["mysql"],
		MasterConnectRetry: 10,
	}
	goodSlave1.FakeMysqlDaemon.StartReplicationCommandsResult = []string{"cmd1"}
	goodSlave1.FakeMysqlDaemon.ExpectedExecuteSuperQueryList = goodSlave1.FakeMysqlDaemon.StartReplicationCommandsResult
	goodSlave1.StartActionLoop(t, wr)
	defer goodSlave1.StopActionLoop(t)

	// Slave2: expect to be re-parented
	goodSlave2.FakeMysqlDaemon.ReadOnly = true
	goodSlave2.FakeMysqlDaemon.StartReplicationCommandsStatus = &myproto.ReplicationStatus{
		Position:           master.FakeMysqlDaemon.CurrentMasterPosition,
		MasterHost:         master.Tablet.Hostname,
		MasterPort:         master.Tablet.Portmap["mysql"],
		MasterConnectRetry: 10,
	}
	goodSlave2.FakeMysqlDaemon.StartReplicationCommandsResult = []string{"cmd1", "cmd2"}
	goodSlave2.FakeMysqlDaemon.ExpectedExecuteSuperQueryList = goodSlave2.FakeMysqlDaemon.StartReplicationCommandsResult
	goodSlave2.StartActionLoop(t, wr)
	defer goodSlave2.StopActionLoop(t)

	// run InitShardMaster
	if err := wr.InitShardMaster(ctx, master.Tablet.Keyspace, master.Tablet.Shard, master.Tablet.Alias, false /*force*/, 10*time.Second); err != nil {
		t.Fatalf("InitShardMaster failed: %v", err)
	}

	// check what was run
	if master.FakeMysqlDaemon.ReadOnly {
		t.Errorf("master was not turned read-write")
	}
	si, err := ts.GetShard(master.Tablet.Keyspace, master.Tablet.Shard)
	if err != nil {
		t.Fatalf("GetShard failed: %v", err)
	}
	if si.MasterAlias != master.Tablet.Alias {
		t.Errorf("unexpected shard master alias, got %v expected %v", si.MasterAlias, master.Tablet.Alias)
	}
}

// TestInitMasterShardChecks makes sure the safety checks work
func TestInitMasterShardChecks(t *testing.T) {
	ctx := context.Background()
	ts := zktopo.NewTestServer(t, []string{"cell1", "cell2"})
	wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient(), time.Second)

	master := NewFakeTablet(t, wr, "cell1", 0, topo.TYPE_MASTER)

	// InitShardMaster with an unknown tablet
	if err := wr.InitShardMaster(ctx, master.Tablet.Keyspace, master.Tablet.Shard, topo.TabletAlias{
		Cell: master.Tablet.Alias.Cell,
		Uid:  master.Tablet.Alias.Uid + 1,
	}, false /*force*/, 10*time.Second); err == nil || !strings.Contains(err.Error(), "is not in the shard") {
		t.Errorf("InitShardMaster with unknown alias returned wrong error: %v", err)
	}

	// InitShardMaster with two masters in the shard, no force flag
	// (master2 needs to run InitTablet with -force, as it is the second
	// master in the same shard)
	master2 := NewFakeTablet(t, wr, "cell1", 1, topo.TYPE_MASTER, ForceInitTablet())
	if err := wr.InitShardMaster(ctx, master2.Tablet.Keyspace, master2.Tablet.Shard, master2.Tablet.Alias, false /*force*/, 10*time.Second); err == nil || !strings.Contains(err.Error(), "is not the only master in the shard") {
		t.Errorf("InitShardMaster with two masters returned wrong error: %v", err)
	}

	// InitShardMaster where the new master fails (use force flag
	// as we have 2 masters)
	master.FakeMysqlDaemon.BreakSlavesError = fmt.Errorf("forced test error")
	master.StartActionLoop(t, wr)
	defer master.StopActionLoop(t)
	if err := wr.InitShardMaster(ctx, master.Tablet.Keyspace, master.Tablet.Shard, master.Tablet.Alias, true /*force*/, 10*time.Second); err == nil || !strings.Contains(err.Error(), "forced test error") {
		t.Errorf("InitShardMaster with new master failing BreakSlaves returned wrong error: %v", err)
	}
}

// TestInitMasterShardOneSlaveFails makes sure that if one slave fails to
// proceed, the action completes anyway
func TestInitMasterShardOneSlaveFails(t *testing.T) {
	ctx := context.Background()
	ts := zktopo.NewTestServer(t, []string{"cell1", "cell2"})
	wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient(), time.Second)

	// Create a master, a couple slaves
	master := NewFakeTablet(t, wr, "cell1", 0, topo.TYPE_MASTER)
	goodSlave := NewFakeTablet(t, wr, "cell1", 1, topo.TYPE_REPLICA)
	badSlave := NewFakeTablet(t, wr, "cell2", 2, topo.TYPE_REPLICA)

	// Master: set a plausible ReplicationPosition to return,
	// and expect to add entry in _vt.reparent_journal
	master.FakeMysqlDaemon.CurrentMasterPosition = myproto.ReplicationPosition{
		GTIDSet: myproto.MariadbGTID{
			Domain:   5,
			Server:   456,
			Sequence: 890,
		},
	}
	master.FakeMysqlDaemon.ReadOnly = true
	master.FakeMysqlDaemon.ExpectedExecuteSuperQueryList = []string{
		"CREATE DATABASE IF NOT EXISTS _vt",
		"SUBCREATE TABLE IF NOT EXISTS _vt.reparent_journal",
		"SUBINSERT INTO _vt.reparent_journal (time_created_ns, action_name, master_alias, replication_position) VALUES",
	}
	master.StartActionLoop(t, wr)
	defer master.StopActionLoop(t)

	// goodSlave: expect to be re-parented
	goodSlave.FakeMysqlDaemon.ReadOnly = true
	goodSlave.FakeMysqlDaemon.StartReplicationCommandsStatus = &myproto.ReplicationStatus{
		Position:           master.FakeMysqlDaemon.CurrentMasterPosition,
		MasterHost:         master.Tablet.Hostname,
		MasterPort:         master.Tablet.Portmap["mysql"],
		MasterConnectRetry: 10,
	}
	goodSlave.FakeMysqlDaemon.StartReplicationCommandsResult = []string{"cmd1"}
	goodSlave.FakeMysqlDaemon.ExpectedExecuteSuperQueryList = goodSlave.FakeMysqlDaemon.StartReplicationCommandsResult
	goodSlave.StartActionLoop(t, wr)
	defer goodSlave.StopActionLoop(t)

	// badSlave: insert an error by failing the ReplicationStatus input
	// on purpose
	badSlave.FakeMysqlDaemon.ReadOnly = true
	badSlave.FakeMysqlDaemon.StartReplicationCommandsStatus = &myproto.ReplicationStatus{
		Position:           master.FakeMysqlDaemon.CurrentMasterPosition,
		MasterHost:         "",
		MasterPort:         0,
		MasterConnectRetry: 10,
	}
	badSlave.StartActionLoop(t, wr)
	defer badSlave.StopActionLoop(t)

	// run InitShardMaster
	if err := wr.InitShardMaster(ctx, master.Tablet.Keyspace, master.Tablet.Shard, master.Tablet.Alias, false /*force*/, 10*time.Second); err == nil || !strings.Contains(err.Error(), "wrong status for StartReplicationCommands") {
		t.Errorf("InitShardMaster with one failed slave returned wrong error: %v", err)
	}

	// check what was run: master should still be good
	if master.FakeMysqlDaemon.ReadOnly {
		t.Errorf("master was not turned read-write")
	}
	si, err := ts.GetShard(master.Tablet.Keyspace, master.Tablet.Shard)
	if err != nil {
		t.Fatalf("GetShard failed: %v", err)
	}
	if si.MasterAlias != master.Tablet.Alias {
		t.Errorf("unexpected shard master alias, got %v expected %v", si.MasterAlias, master.Tablet.Alias)
	}
}
<commit_msg>One more error path unit tested.<commit_after>// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package testlib

import (
	"fmt"
	"strings"
	"time"

	"github.com/youtube/vitess/go/vt/logutil"
	myproto "github.com/youtube/vitess/go/vt/mysqlctl/proto"
	"github.com/youtube/vitess/go/vt/tabletmanager/tmclient"
	"github.com/youtube/vitess/go/vt/topo"
	"github.com/youtube/vitess/go/vt/wrangler"
	"github.com/youtube/vitess/go/vt/zktopo"
	"golang.org/x/net/context"

	"testing"
)

// TestInitMasterShard is the good scenario test, where everything
// works as planned
func TestInitMasterShard(t *testing.T) {
	ctx := context.Background()
	ts := zktopo.NewTestServer(t, []string{"cell1", "cell2"})
	wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient(), time.Second)

	// Create a master, a couple good slaves
	master := NewFakeTablet(t, wr, "cell1", 0, topo.TYPE_MASTER)
	goodSlave1 := NewFakeTablet(t, wr, "cell1", 1, topo.TYPE_REPLICA)
	goodSlave2 := NewFakeTablet(t, wr, "cell2", 2, topo.TYPE_REPLICA)

	// Master: set a plausible ReplicationPosition to return,
	// and expect to add entry in _vt.reparent_journal
	master.FakeMysqlDaemon.CurrentMasterPosition = myproto.ReplicationPosition{
		GTIDSet: myproto.MariadbGTID{
			Domain:   5,
			Server:   456,
			Sequence: 890,
		},
	}
	master.FakeMysqlDaemon.ReadOnly = true
	master.FakeMysqlDaemon.ExpectedExecuteSuperQueryList = []string{
		"CREATE DATABASE IF NOT EXISTS _vt",
		"SUBCREATE TABLE IF NOT EXISTS _vt.reparent_journal",
		"SUBINSERT INTO _vt.reparent_journal (time_created_ns, action_name, master_alias, replication_position) VALUES",
	}
	master.StartActionLoop(t, wr)
	defer master.StopActionLoop(t)

	// Slave1: expect to be re-parented
	goodSlave1.FakeMysqlDaemon.ReadOnly = true
	goodSlave1.FakeMysqlDaemon.StartReplicationCommandsStatus = &myproto.ReplicationStatus{
		Position:           master.FakeMysqlDaemon.CurrentMasterPosition,
		MasterHost:         master.Tablet.Hostname,
		MasterPort:         master.Tablet.Portmap["mysql"],
		MasterConnectRetry: 10,
	}
	goodSlave1.FakeMysqlDaemon.StartReplicationCommandsResult = []string{"cmd1"}
	goodSlave1.FakeMysqlDaemon.ExpectedExecuteSuperQueryList = goodSlave1.FakeMysqlDaemon.StartReplicationCommandsResult
	goodSlave1.StartActionLoop(t, wr)
	defer goodSlave1.StopActionLoop(t)

	// Slave2: expect to be re-parented
	goodSlave2.FakeMysqlDaemon.ReadOnly = true
	goodSlave2.FakeMysqlDaemon.StartReplicationCommandsStatus = &myproto.ReplicationStatus{
		Position:           master.FakeMysqlDaemon.CurrentMasterPosition,
		MasterHost:         master.Tablet.Hostname,
		MasterPort:         master.Tablet.Portmap["mysql"],
		MasterConnectRetry: 10,
	}
	goodSlave2.FakeMysqlDaemon.StartReplicationCommandsResult = []string{"cmd1", "cmd2"}
	goodSlave2.FakeMysqlDaemon.ExpectedExecuteSuperQueryList = goodSlave2.FakeMysqlDaemon.StartReplicationCommandsResult
	goodSlave2.StartActionLoop(t, wr)
	defer goodSlave2.StopActionLoop(t)

	// run InitShardMaster
	if err := wr.InitShardMaster(ctx, master.Tablet.Keyspace, master.Tablet.Shard, master.Tablet.Alias, false /*force*/, 10*time.Second); err != nil {
		t.Fatalf("InitShardMaster failed: %v", err)
	}

	// check what was run
	if master.FakeMysqlDaemon.ReadOnly {
		t.Errorf("master was not turned read-write")
	}
	si, err := ts.GetShard(master.Tablet.Keyspace, master.Tablet.Shard)
	if err != nil {
		t.Fatalf("GetShard failed: %v", err)
	}
	if si.MasterAlias != master.Tablet.Alias {
		t.Errorf("unexpected shard master alias, got %v expected %v", si.MasterAlias, master.Tablet.Alias)
	}
}

// TestInitMasterShardChecks makes sure the safety checks work
func TestInitMasterShardChecks(t *testing.T) {
	ctx := context.Background()
	ts := zktopo.NewTestServer(t, []string{"cell1", "cell2"})
	wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient(), time.Second)

	master := NewFakeTablet(t, wr, "cell1", 0, topo.TYPE_MASTER)

	// InitShardMaster with an unknown tablet
	if err := wr.InitShardMaster(ctx, master.Tablet.Keyspace, master.Tablet.Shard, topo.TabletAlias{
		Cell: master.Tablet.Alias.Cell,
		Uid:  master.Tablet.Alias.Uid + 1,
	}, false /*force*/, 10*time.Second); err == nil || !strings.Contains(err.Error(), "is not in the shard") {
		t.Errorf("InitShardMaster with unknown alias returned wrong error: %v", err)
	}

	// InitShardMaster with two masters in the shard, no force flag
	// (master2 needs to run InitTablet with -force, as it is the second
	// master in the same shard)
	master2 := NewFakeTablet(t, wr, "cell1", 1, topo.TYPE_MASTER, ForceInitTablet())
	if err := wr.InitShardMaster(ctx, master2.Tablet.Keyspace, master2.Tablet.Shard, master2.Tablet.Alias, false /*force*/, 10*time.Second); err == nil || !strings.Contains(err.Error(), "is not the only master in the shard") {
		t.Errorf("InitShardMaster with two masters returned wrong error: %v", err)
	}

	// InitShardMaster where the new master fails (use force flag
	// as we have 2 masters)
	master.FakeMysqlDaemon.BreakSlavesError = fmt.Errorf("forced test error")
	master.StartActionLoop(t, wr)
	defer master.StopActionLoop(t)
	if err := wr.InitShardMaster(ctx, master.Tablet.Keyspace, master.Tablet.Shard, master.Tablet.Alias, true /*force*/, 10*time.Second); err == nil || !strings.Contains(err.Error(), "forced test error") {
		t.Errorf("InitShardMaster with new master failing BreakSlaves returned wrong error: %v", err)
	}
}

// TestInitMasterShardOneSlaveFails makes sure that if one slave fails to
// proceed, the action completes anyway
func TestInitMasterShardOneSlaveFails(t *testing.T) {
	ctx := context.Background()
	ts := zktopo.NewTestServer(t, []string{"cell1", "cell2"})
	wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient(), time.Second)

	// Create a master, a couple slaves
	master := NewFakeTablet(t, wr, "cell1", 0, topo.TYPE_MASTER)
	goodSlave := NewFakeTablet(t, wr, "cell1", 1, topo.TYPE_REPLICA)
	badSlave := NewFakeTablet(t, wr, "cell2", 2, topo.TYPE_REPLICA)

	// Master: set a plausible ReplicationPosition to return,
	// and expect to add entry in _vt.reparent_journal
	master.FakeMysqlDaemon.CurrentMasterPosition = myproto.ReplicationPosition{
		GTIDSet: myproto.MariadbGTID{
			Domain:   5,
			Server:   456,
			Sequence: 890,
		},
	}
	master.FakeMysqlDaemon.ReadOnly = true
	master.FakeMysqlDaemon.ExpectedExecuteSuperQueryList = []string{
		"CREATE DATABASE IF NOT EXISTS _vt",
		"SUBCREATE TABLE IF NOT EXISTS _vt.reparent_journal",
		"SUBINSERT INTO _vt.reparent_journal (time_created_ns, action_name, master_alias, replication_position) VALUES",
	}
	master.StartActionLoop(t, wr)
	defer master.StopActionLoop(t)

	// goodSlave: expect to be re-parented
	goodSlave.FakeMysqlDaemon.ReadOnly = true
	goodSlave.FakeMysqlDaemon.StartReplicationCommandsStatus = &myproto.ReplicationStatus{
		Position:           master.FakeMysqlDaemon.CurrentMasterPosition,
		MasterHost:         master.Tablet.Hostname,
		MasterPort:         master.Tablet.Portmap["mysql"],
		MasterConnectRetry: 10,
	}
	goodSlave.FakeMysqlDaemon.StartReplicationCommandsResult = []string{"cmd1"}
	goodSlave.FakeMysqlDaemon.ExpectedExecuteSuperQueryList = goodSlave.FakeMysqlDaemon.StartReplicationCommandsResult
	goodSlave.StartActionLoop(t, wr)
	defer goodSlave.StopActionLoop(t)

	// badSlave: insert an error by failing the ReplicationStatus input
	// on purpose
	badSlave.FakeMysqlDaemon.ReadOnly = true
	badSlave.FakeMysqlDaemon.StartReplicationCommandsStatus = &myproto.ReplicationStatus{
		Position:           master.FakeMysqlDaemon.CurrentMasterPosition,
		MasterHost:         "",
		MasterPort:         0,
		MasterConnectRetry: 10,
	}
	badSlave.StartActionLoop(t, wr)
	defer badSlave.StopActionLoop(t)

	// also change the master alias in the Shard object, to make sure it
	// is set back.
	si, err := ts.GetShard(master.Tablet.Keyspace, master.Tablet.Shard)
	if err != nil {
		t.Fatalf("GetShard failed: %v", err)
	}
	si.MasterAlias.Uid++
	if err := topo.UpdateShard(ctx, ts, si); err != nil {
		t.Fatalf("UpdateShard failed: %v", err)
	}

	// run InitShardMaster without force, it fails because master is
	// changing.
	if err := wr.InitShardMaster(ctx, master.Tablet.Keyspace, master.Tablet.Shard, master.Tablet.Alias, false /*force*/, 10*time.Second); err == nil || !strings.Contains(err.Error(), "is not the shard master") {
		t.Errorf("InitShardMaster with mismatched new master returned wrong error: %v", err)
	}

	// run InitShardMaster
	if err := wr.InitShardMaster(ctx, master.Tablet.Keyspace, master.Tablet.Shard, master.Tablet.Alias, true /*force*/, 10*time.Second); err == nil || !strings.Contains(err.Error(), "wrong status for StartReplicationCommands") {
		t.Errorf("InitShardMaster with one failed slave returned wrong error: %v", err)
	}

	// check what was run: master should still be good
	if master.FakeMysqlDaemon.ReadOnly {
		t.Errorf("master was not turned read-write")
	}
	si, err = ts.GetShard(master.Tablet.Keyspace, master.Tablet.Shard)
	if err != nil {
		t.Fatalf("GetShard failed: %v", err)
	}
	if si.MasterAlias != master.Tablet.Alias {
		t.Errorf("unexpected shard master alias, got %v expected %v", si.MasterAlias, master.Tablet.Alias)
	}
}
<|endoftext|>


<commit_before>package sdees

import (
	"os"
	"os/exec"
	"sort"
	"strings"
	"time"
)

func ListFiles(gitfolder string) []string {
	defer timeTrack(time.Now(), "Listing files")
	cwd, _ := os.Getwd()
	defer os.Chdir(cwd)
	err := os.Chdir(gitfolder)
	if err != nil {
		logger.Error("ListFiles Cannot chdir into " + gitfolder)
	}

	cmd := exec.Command("git", "ls-tree", "--name-only", "master")
	stdout, err := cmd.Output()
	if err != nil {
		logger.Error("Problem doing ls-tree")
	}
	documents := []string{}
	for _, document := range strings.Split(strings.TrimSpace(string(stdout)), "\n") {
		if document[0] == '.' {
			continue
		}
		documents = append(documents, document)
	}

	sort.Strings(documents)
	return documents
}

func ListFileOfOne(gitfolder string, branch string) (string, bool) {
	defer timeTrack(time.Now(), "Listing files")
	cwd, _ := os.Getwd()
	defer os.Chdir(cwd)
	err := os.Chdir(gitfolder)
	if err != nil {
		logger.Error("Cannot chdir into " + gitfolder)
	}

	cmd := exec.Command("git", "ls-tree", "--name-only", branch)
	stdout, err := cmd.Output()
	if err != nil {
		logger.Error("Problem doing ls-tree")
	}
	docMap := make(map[string]bool)
	for _, document := range strings.Split(strings.TrimSpace(string(stdout)), "\n") {
		if document[0] == '.' {
			continue
		}
		docMap[strings.Replace(document, ".gpg", "", -1)] = strings.Contains(document, ".gpg")
	}

	documents := make([]string, len(docMap))
	i := 0
	for k, _ := range docMap {
		documents[i] = k
		i++
	}
	sort.Strings(documents)
	encrypted := make([]bool, len(documents))
	for i, val := range documents {
		encrypted[i] = docMap[val]
	}
	return documents[0], encrypted[0]
}
<commit_msg>Fixed listfiles<commit_after>package sdees

import (
	"os"
	"os/exec"
	"sort"
	"strings"
	"time"
)

func ListFiles(gitfolder string) []string {
	defer timeTrack(time.Now(), "Listing files")
	cwd, _ := os.Getwd()
	defer os.Chdir(cwd)
	err := os.Chdir(gitfolder)
	if err != nil {
		logger.Error("ListFiles Cannot chdir into " + gitfolder)
	}

	cmd := exec.Command("git", "ls-tree", "--name-only", "master")
	stdout, err := cmd.Output()
	if err != nil {
		logger.Error("Problem doing ls-tree")
	}
	documents := []string{}
	for _, document := range strings.Split(strings.TrimSpace(string(stdout)), "\n") {
		if document[0] == '.' && document != ".key" {
			documents = append(documents, document)
		}
	}

	sort.Strings(documents)
	return documents
}

func ListFileOfOne(gitfolder string, branch string) (string, bool) {
	defer timeTrack(time.Now(), "Listing files")
	cwd, _ := os.Getwd()
	defer os.Chdir(cwd)
	err := os.Chdir(gitfolder)
	if err != nil {
		logger.Error("Cannot chdir into " + gitfolder)
	}

	cmd := exec.Command("git", "ls-tree", "--name-only", branch)
	stdout, err := cmd.Output()
	if err != nil {
		logger.Error("Problem doing ls-tree")
	}
	docMap := make(map[string]bool)
	for _, document := range strings.Split(strings.TrimSpace(string(stdout)), "\n") {
		if document[0] == '.' {
			continue
		}
		docMap[strings.Replace(document, ".gpg", "", -1)] = strings.Contains(document, ".gpg")
	}

	documents := make([]string, len(docMap))
	i := 0
	for k, _ := range docMap {
		documents[i] = k
		i++
	}
	sort.Strings(documents)
	encrypted := make([]bool, len(documents))
	for i, val := range documents {
		encrypted[i] = docMap[val]
	}
	return documents[0], encrypted[0]
}
<|endoftext|>