<commit_before>// Copyright (c) 2012-2014 Jeremy Latt
// Copyright (c) 2016-2017 Daniel Oaks <daniel@danieloaks.net>
// released under the MIT license

package irc

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"strings"

	"github.com/goshuirc/irc-go/ircfmt"
	"github.com/oragono/oragono/irc/history"
	"github.com/oragono/oragono/irc/sno"
	"github.com/oragono/oragono/irc/utils"
)

var (
	// anything added here MUST be casefolded:
	restrictedNicknames = map[string]bool{
		"=scene=":  true, // used for rp commands
		"histserv": true, // TODO(slingamn) this should become a real service
	}
)

// returns whether the change succeeded or failed
func performNickChange(server *Server, client *Client, target *Client, session *Session, newnick string, rb *ResponseBuffer) bool {
	nickname := strings.TrimSpace(newnick)
	cfnick, err := CasefoldName(nickname)
	currentNick := client.Nick()

	if len(nickname) < 1 {
		rb.Add(nil, server.name, ERR_NONICKNAMEGIVEN, currentNick, client.t("No nickname given"))
		return false
	}

	if err != nil || len(nickname) > server.Limits().NickLen || restrictedNicknames[cfnick] {
		rb.Add(nil, server.name, ERR_ERRONEUSNICKNAME, currentNick, nickname, client.t("Erroneous nickname"))
		return false
	}

	if target.Nick() == nickname {
		return true
	}

	hadNick := target.HasNick()
	origNickMask := target.NickMaskString()
	details := target.Details()
	err = client.server.clients.SetNick(target, session, nickname)
	if err == errNicknameInUse {
		rb.Add(nil, server.name, ERR_NICKNAMEINUSE, currentNick, nickname, client.t("Nickname is already in use"))
		return false
	} else if err == errNicknameReserved {
		rb.Add(nil, server.name, ERR_NICKNAMEINUSE, currentNick, nickname, client.t("Nickname is reserved by a different account"))
		return false
	} else if err != nil {
		rb.Add(nil, server.name, ERR_UNKNOWNERROR, currentNick, "NICK", fmt.Sprintf(client.t("Could not set or change nickname: %s"), err.Error()))
		return false
	}

	message := utils.MakeSplitMessage("", true)
	histItem := history.Item{
		Type:        history.Nick,
		Nick:        origNickMask,
		AccountName: details.accountName,
		Message:     message,
	}
	histItem.Params[0] = nickname

	client.server.logger.Debug("nick", fmt.Sprintf("%s changed nickname to %s [%s]", origNickMask, nickname, cfnick))
	if hadNick {
		target.server.snomasks.Send(sno.LocalNicks, fmt.Sprintf(ircfmt.Unescape("$%s$r changed nickname to %s"), details.nick, nickname))
		target.server.whoWas.Append(details.WhoWas)
		rb.AddFromClient(message.Time, message.Msgid, origNickMask, details.accountName, nil, "NICK", nickname)
		for session := range target.Friends() {
			if session != rb.session {
				session.sendFromClientInternal(false, message.Time, message.Msgid, origNickMask, details.accountName, nil, "NICK", nickname)
			}
		}
	}

	for _, channel := range client.Channels() {
		channel.history.Add(histItem)
	}

	target.nickTimer.Touch(rb)

	if target.Registered() {
		client.server.monitorManager.AlertAbout(target, true)
	}
	// else: Run() will attempt registration immediately after this
	return true
}

func (server *Server) RandomlyRename(client *Client) {
	prefix := server.AccountConfig().NickReservation.RenamePrefix
	if prefix == "" {
		prefix = "Guest-"
	}
	buf := make([]byte, 8)
	rand.Read(buf)
	nick := fmt.Sprintf("%s%s", prefix, hex.EncodeToString(buf))
	sessions := client.Sessions()
	if len(sessions) == 0 {
		return
	}
	// XXX arbitrarily pick the first session to receive error messages;
	// all other sessions receive a `NICK` line same as a friend would
	rb := NewResponseBuffer(sessions[0])
	performNickChange(server, client, client, nil, nick, rb)
	rb.Send(false)
	// technically performNickChange can fail to change the nick,
	// but if they're still delinquent, the timer will get them later
}
<commit_msg>fix #360<commit_after>// Copyright (c) 2012-2014 Jeremy Latt
// Copyright (c) 2016-2017 Daniel Oaks <daniel@danieloaks.net>
// released under the MIT license

package irc

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"strings"

	"github.com/goshuirc/irc-go/ircfmt"
	"github.com/oragono/oragono/irc/history"
	"github.com/oragono/oragono/irc/sno"
	"github.com/oragono/oragono/irc/utils"
)

var (
	// anything added here MUST be casefolded:
	restrictedNicknames = map[string]bool{
		"=scene=":  true, // used for rp commands
		"histserv": true, // TODO(slingamn) this should become a real service
	}
)

// returns whether the change succeeded or failed
func performNickChange(server *Server, client *Client, target *Client, session *Session, newnick string, rb *ResponseBuffer) bool {
	nickname := strings.TrimSpace(newnick)
	cfnick, err := CasefoldName(nickname)
	currentNick := client.Nick()

	if len(nickname) < 1 {
		rb.Add(nil, server.name, ERR_NONICKNAMEGIVEN, currentNick, client.t("No nickname given"))
		return false
	}

	if err != nil || len(nickname) > server.Limits().NickLen || restrictedNicknames[cfnick] {
		rb.Add(nil, server.name, ERR_ERRONEUSNICKNAME, currentNick, nickname, client.t("Erroneous nickname"))
		return false
	}

	if target.Nick() == nickname {
		return true
	}

	hadNick := target.HasNick()
	origNickMask := target.NickMaskString()
	details := target.Details()
	err = client.server.clients.SetNick(target, session, nickname)
	if err == errNicknameInUse {
		rb.Add(nil, server.name, ERR_NICKNAMEINUSE, currentNick, nickname, client.t("Nickname is already in use"))
		return false
	} else if err == errNicknameReserved {
		rb.Add(nil, server.name, ERR_NICKNAMEINUSE, currentNick, nickname, client.t("Nickname is reserved by a different account"))
		return false
	} else if err != nil {
		rb.Add(nil, server.name, ERR_UNKNOWNERROR, currentNick, "NICK", fmt.Sprintf(client.t("Could not set or change nickname: %s"), err.Error()))
		return false
	}

	message := utils.MakeSplitMessage("", true)
	histItem := history.Item{
		Type:        history.Nick,
		Nick:        origNickMask,
		AccountName: details.accountName,
		Message:     message,
	}
	histItem.Params[0] = nickname

	client.server.logger.Debug("nick", fmt.Sprintf("%s changed nickname to %s [%s]", origNickMask, nickname, cfnick))
	if hadNick {
		if client == target {
			target.server.snomasks.Send(sno.LocalNicks, fmt.Sprintf(ircfmt.Unescape("$%s$r changed nickname to %s"), details.nick, nickname))
		} else {
			target.server.snomasks.Send(sno.LocalNicks, fmt.Sprintf(ircfmt.Unescape("Operator %s changed nickname of $%s$r to %s"), client.Nick(), details.nick, nickname))
		}
		target.server.whoWas.Append(details.WhoWas)
		rb.AddFromClient(message.Time, message.Msgid, origNickMask, details.accountName, nil, "NICK", nickname)
		for session := range target.Friends() {
			if session != rb.session {
				session.sendFromClientInternal(false, message.Time, message.Msgid, origNickMask, details.accountName, nil, "NICK", nickname)
			}
		}
	}

	for _, channel := range client.Channels() {
		channel.history.Add(histItem)
	}

	target.nickTimer.Touch(rb)

	if target.Registered() {
		client.server.monitorManager.AlertAbout(target, true)
	}
	// else: Run() will attempt registration immediately after this
	return true
}

func (server *Server) RandomlyRename(client *Client) {
	prefix := server.AccountConfig().NickReservation.RenamePrefix
	if prefix == "" {
		prefix = "Guest-"
	}
	buf := make([]byte, 8)
	rand.Read(buf)
	nick := fmt.Sprintf("%s%s", prefix, hex.EncodeToString(buf))
	sessions := client.Sessions()
	if len(sessions) == 0 {
		return
	}
	// XXX arbitrarily pick the first session to receive error messages;
	// all other sessions receive a `NICK` line same as a friend would
	rb := NewResponseBuffer(sessions[0])
	performNickChange(server, client, client, nil, nick, rb)
	rb.Send(false)
	// technically performNickChange can fail to change the nick,
	// but if they're still delinquent, the timer will get them later
}
<|endoftext|>


<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"

	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)
}
<commit_msg>Rewrite to use samkumar/etcdstruct library<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/samkumar/etcdstruct"
	"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
}

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

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

// 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) SetRetrievedRevision(rev int64) {
	tdef.retrievedRevision = rev
}

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

// 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):]
}

// 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 = etcdstruct.RetrieveEtcdStruct(ctx, etcdClient, getEtcdKey(username, accountpath), 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 = etcdstruct.RetrieveEtcdStruct(ctx, etcdClient, getEtcdKey(tag, tagpath), tdef)
	return
}

// 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 etcdstruct.UpsertEtcdStruct(ctx, etcdClient, getEtcdKey(acc.Username, accountpath), acc)
}

// Same as UpsertAccount, but for Tag Definitions
func UpsertTagDef(ctx context.Context, etcdClient *etcd.Client, tdef *MrPlotterTagDef) error {
	return etcdstruct.UpsertEtcdStruct(ctx, etcdClient, getEtcdKey(tdef.Tag, tagpath), 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 UpsertAccountAtomically(ctx context.Context, etcdClient *etcd.Client, acc *MrPlotterAccount) (bool, error) {
	return etcdstruct.UpsertEtcdStructAtomic(ctx, etcdClient, getEtcdKey(acc.Username, accountpath), acc)
}

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

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

// Deletes the tag definition with the provided name
func DeleteTagDef(ctx context.Context, etcdClient *etcd.Client, tag string) error {
	_, err := etcdstruct.DeleteEtcdStructs(ctx, etcdClient, getEtcdKey(tag, tagpath))
	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)
	accs := make([]*MrPlotterAccount, 0, 16)
	err := etcdstruct.RetrieveEtcdStructs(ctx, etcdClient, func(key []byte) etcdstruct.EtcdStruct {
		acc := &MrPlotterAccount{}
		accs = append(accs, acc)
		return acc
	}, func(es etcdstruct.EtcdStruct, key []byte) {
		acc := es.(*MrPlotterAccount)
		acc.Username = getNameFromEtcdKey(string(key), tagpath)
		acc.Tags = nil
	}, etcdKeyPrefix, etcd.WithPrefix())
	if err != nil {
		return nil, err
	}

	return accs, 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 RetrieveMultipleTagDefs(ctx context.Context, etcdClient *etcd.Client, tagprefix string) ([]*MrPlotterTagDef, error) {
	etcdKeyPrefix := getEtcdKey(tagprefix, tagpath)
	tdefs := make([]*MrPlotterTagDef, 0, 16)
	err := etcdstruct.RetrieveEtcdStructs(ctx, etcdClient, func(key []byte) etcdstruct.EtcdStruct {
		tdef := &MrPlotterTagDef{}
		tdefs = append(tdefs, tdef)
		return tdef
	}, func(es etcdstruct.EtcdStruct, key []byte) {
		tdef := es.(*MrPlotterTagDef)
		tdef.Tag = getNameFromEtcdKey(string(key), tagpath)
		tdef.PathPrefix = nil
	}, etcdKeyPrefix, etcd.WithPrefix())
	if err != nil {
		return nil, err
	}

	return tdefs, 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 etcdstruct.DeleteEtcdStructs(ctx, etcdClient, getEtcdKey(usernameprefix, accountpath), etcd.WithPrefix())
}

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


<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// GNU/Linux library calls.

package syscall

import "unsafe"

//sys	Openat(dirfd int, path string, flags int, mode uint32) (fd int, err error)
//openat(dirfd int, path *byte, flags int, mode Mode_t) int

//sys	futimesat(dirfd int, path *byte, times *[2]Timeval) (err error)
//futimesat(dirfd int, path *byte, times *[2]Timeval) int
func Futimesat(dirfd int, path string, tv []Timeval) (err error) {
	if len(tv) != 2 {
		return EINVAL
	}
	return futimesat(dirfd, StringBytePtr(path), (*[2]Timeval)(unsafe.Pointer(&tv[0])))
}

func Futimes(fd int, tv []Timeval) (err error) {
	// Believe it or not, this is the best we can do on GNU/Linux
	// (and is what glibc does).
	return Utimes("/proc/self/fd/"+itoa(fd), tv)
}

//sys	ptrace(request int, pid int, addr uintptr, data uintptr) (err error)
//ptrace(request int, pid Pid_t, addr *byte, data *byte) _C_long

//sysnb raw_ptrace(request int, pid int, addr *byte, data *byte) (err Errno)
//ptrace(request int, pid Pid_t, addr *byte, data *byte) _C_long

func ptracePeek(req int, pid int, addr uintptr, out []byte) (count int, err error) {
	// The peek requests are machine-size oriented, so we wrap it
	// to retrieve arbitrary-length data.

	// The ptrace syscall differs from glibc's ptrace.
	// Peeks returns the word in *data, not as the return value.

	var buf [sizeofPtr]byte

	// Leading edge.  PEEKTEXT/PEEKDATA don't require aligned
	// access (PEEKUSER warns that it might), but if we don't
	// align our reads, we might straddle an unmapped page
	// boundary and not get the bytes leading up to the page
	// boundary.
	n := 0
	if addr%sizeofPtr != 0 {
		err = ptrace(req, pid, addr-addr%sizeofPtr, uintptr(unsafe.Pointer(&buf[0])))
		if err != nil {
			return 0, err
		}
		n += copy(out, buf[addr%sizeofPtr:])
		out = out[n:]
	}

	// Remainder.
	for len(out) > 0 {
		// We use an internal buffer to gaurantee alignment.
		// It's not documented if this is necessary, but we're paranoid.
		err = ptrace(req, pid, addr+uintptr(n), uintptr(unsafe.Pointer(&buf[0])))
		if err != nil {
			return n, err
		}
		copied := copy(out, buf[0:])
		n += copied
		out = out[copied:]
	}

	return n, nil
}

func PtracePeekText(pid int, addr uintptr, out []byte) (count int, err error) {
	return ptracePeek(PTRACE_PEEKTEXT, pid, addr, out)
}

func PtracePeekData(pid int, addr uintptr, out []byte) (count int, err error) {
	return ptracePeek(PTRACE_PEEKDATA, pid, addr, out)
}

func ptracePoke(pokeReq int, peekReq int, pid int, addr uintptr, data []byte) (count int, err error) {
	// As for ptracePeek, we need to align our accesses to deal
	// with the possibility of straddling an invalid page.

	// Leading edge.
	n := 0
	if addr%sizeofPtr != 0 {
		var buf [sizeofPtr]byte
		err = ptrace(peekReq, pid, addr-addr%sizeofPtr, uintptr(unsafe.Pointer(&buf[0])))
		if err != nil {
			return 0, err
		}
		n += copy(buf[addr%sizeofPtr:], data)
		word := *((*uintptr)(unsafe.Pointer(&buf[0])))
		err = ptrace(pokeReq, pid, addr-addr%sizeofPtr, word)
		if err != nil {
			return 0, err
		}
		data = data[n:]
	}

	// Interior.
	for len(data) > int(sizeofPtr) {
		word := *((*uintptr)(unsafe.Pointer(&data[0])))
		err = ptrace(pokeReq, pid, addr+uintptr(n), word)
		if err != nil {
			return n, err
		}
		n += int(sizeofPtr)
		data = data[sizeofPtr:]
	}

	// Trailing edge.
	if len(data) > 0 {
		var buf [sizeofPtr]byte
		err = ptrace(peekReq, pid, addr+uintptr(n), uintptr(unsafe.Pointer(&buf[0])))
		if err != nil {
			return n, err
		}
		copy(buf[0:], data)
		word := *((*uintptr)(unsafe.Pointer(&buf[0])))
		err = ptrace(pokeReq, pid, addr+uintptr(n), word)
		if err != nil {
			return n, err
		}
		n += len(data)
	}

	return n, nil
}

func PtracePokeText(pid int, addr uintptr, data []byte) (count int, err error) {
	return ptracePoke(PTRACE_POKETEXT, PTRACE_PEEKTEXT, pid, addr, data)
}

func PtracePokeData(pid int, addr uintptr, data []byte) (count int, err error) {
	return ptracePoke(PTRACE_POKEDATA, PTRACE_PEEKDATA, pid, addr, data)
}

func PtraceSetOptions(pid int, options int) (err error) {
	return ptrace(PTRACE_SETOPTIONS, pid, 0, uintptr(options))
}

func PtraceGetEventMsg(pid int) (msg uint, err error) {
	var data _C_long
	err = ptrace(PTRACE_GETEVENTMSG, pid, 0, uintptr(unsafe.Pointer(&data)))
	msg = uint(data)
	return
}

func PtraceCont(pid int, signal int) (err error) {
	return ptrace(PTRACE_CONT, pid, 0, uintptr(signal))
}

func PtraceSingleStep(pid int) (err error) { return ptrace(PTRACE_SINGLESTEP, pid, 0, 0) }

func PtraceAttach(pid int) (err error) { return ptrace(PTRACE_ATTACH, pid, 0, 0) }

func PtraceDetach(pid int) (err error) { return ptrace(PTRACE_DETACH, pid, 0, 0) }

// FIXME: mksysinfo needs to produce LINUX_REBOOT_MAGIC[12].

// //sys	reboot(magic1 uint, magic2 uint, cmd int, arg string) (err error)
// //reboot(magic1 uint, magic2 uint, cmd int, arg *byte) int
// func Reboot(cmd int) (err error) {
// 	return reboot(LINUX_REBOOT_MAGIC1, LINUX_REBOOT_MAGIC2, cmd, "")
// }

//sys	Acct(path string) (err error)
//acct(path *byte) int

// FIXME: mksysinfo Timex
// //sys	Adjtimex(buf *Timex) (state int, err error)
// //adjtimex(buf *Timex) int

//sys	Faccessat(dirfd int, path string, mode uint32, flags int) (err error)
//faccessat(dirfd int, pathname *byte, mode int, flags int) int

// FIXME: Only in glibc 2.10 and later.
// //sys	Fallocate(fd int, mode uint32, off int64, len int64) (err error)
// //fallocate(fd int, mode int, offset Offset_t, len Offset_t) int

//sys	Fchmodat(dirfd int, path string, mode uint32, flags int) (err error)
//fchmodat(dirfd int, pathname *byte, mode Mode_t, flags int) int

//sys	Fchownat(dirfd int, path string, uid int, gid int, flags int) (err error)
//fchownat(dirfd int, path *byte, owner Uid_t, group Gid_t, flags int) int

//sys	Flock(fd int, how int) (err error)
//flock(fd int, how int) int

// FIXME: mksysinfo statfs
// //sys	Fstatfs(fd int, buf *Statfs_t) (err error)
// //fstatfs(fd int, buf *Statfs_t) int

// FIXME: Only available as a syscall.
// //sysnb	Gettid() (tid int)
// //gettid() Pid_t

// FIXME: mksysinfo linux_dirent
//    Or just abandon this function.
// //sys	Getdents(fd int, buf []byte) (n int, err error)
// //getdents64(fd int, buf *byte, count uint)

//sys	InotifyAddWatch(fd int, pathname string, mask uint32) (watchdesc int, err error)
//inotify_add_watch(fd int, pathname *byte, mask uint32) int

//sysnb	InotifyInit() (fd int, err error)
//inotify_init() int

// FIXME: Only in glibc 2.9 and later.
// //sysnb	InotifyInit1(flags int) (fd int, err error)
// //inotify_init1(flags int) int

//sysnb	InotifyRmWatch(fd int, watchdesc uint32) (success int, err error)
//inotify_rm_watch(fd int, wd uint32) int

//sys	Klogctl(typ int, buf []byte) (n int, err error)
//klogctl(typ int, bufp *byte, len int) int

//sys	Mkdirat(dirfd int, path string, mode uint32) (err error)
//mkdirat(dirfd int, path *byte, mode Mode_t) int

//sys	Mknodat(dirfd int, path string, mode uint32, dev int) (err error)
//mknodat(dirfd int, path *byte, mode Mode_t, dev _dev_t) int

//sys	PivotRoot(newroot string, putold string) (err error)
//pivot_root(newroot *byte, putold *byte) int

//sys	Renameat(olddirfd int, oldpath string, newdirfd int, newpath string) (err error)
//renameat(olddirfd int, oldpath *byte, newdirfd int, newpath *byte) int

//sys	sendfile(outfd int, infd int, offset *Offset_t, count int) (written int, err error)
//sendfile64(outfd int, infd int, offset *Offset_t, count Size_t) Ssize_t
func Sendfile(outfd int, infd int, offset *int64, count int) (written int, err error) {
	var soff Offset_t
	var psoff *Offset_t
	if offset != nil {
		psoff = &soff
	}
	written, err = sendfile(outfd, infd, psoff, count)
	if offset != nil {
		*offset = int64(soff)
	}
	return
}

//sys	Setfsgid(gid int) (err error)
//setfsgid(gid Gid_t) int

//sys	Setfsuid(uid int) (err error)
//setfsuid(uid Uid_t) int

//sysnb	Setresgid(rgid int, egid int, sgid int) (err error)
//setresgid(rgid Gid_t, egid Gid_t, sgid Gid_t) int

//sysnb	Setresuid(ruid int, eguid int, suid int) (err error)
//setresuid(ruid Uid_t, euid Uid_t, suid Uid_t) int

//sys	splice(rfd int, roff *_loff_t, wfd int, woff *_loff_t, len int, flags int) (n int64, err error)
//splice(rfd int, roff *_loff_t, wfd int, woff *_loff_t, len Size_t, flags uint) Ssize_t
func Splice(rfd int, roff *int64, wfd int, woff *int64, len int, flags int) (n int64, err error) {
	var lroff _loff_t
	var plroff *_loff_t
	if roff != nil {
		plroff = &lroff
	}
	var lwoff _loff_t
	var plwoff *_loff_t
	if woff != nil {
		plwoff = &lwoff
	}
	n, err = splice(rfd, plroff, wfd, plwoff, len, flags)
	if roff != nil {
		*roff = int64(lroff)
	}
	if woff != nil {
		*woff = int64(lwoff)
	}
	return
}

// FIXME: mksysinfo statfs
// //sys	Statfs(path string, buf *Statfs_t) (err error)
// //statfs(path *byte, buf *Statfs_t) int

// FIXME: Only in glibc 2.6 and later.
// //sys	SyncFileRange(fd int, off int64, n int64, flags int) (err error)
// //sync_file_range(fd int, off Offset_t, n Offset_t, flags uint) int

// FIXME: mksysinfo Sysinfo_t
// //sysnb	Sysinfo(info *Sysinfo_t) (err error)
// //sysinfo(info *Sysinfo_t) int

//sys	Tee(rfd int, wfd int, len int, flags int) (n int64, err error)
//tee(rfd int, wfd int, len Size_t, flags uint) Ssize_t

// FIXME: Only available as a syscall.
// //sysnb	Tgkill(tgid int, tid int, sig int) (err error)
// //tgkill(tgid int, tid int, sig int) int

//sys	unlinkat(dirfd int, path string, flags int) (err error)
//unlinkat(dirfd int, path *byte, flags int) int

func Unlinkat(dirfd int, path string) (err error) {
	return unlinkat(dirfd, path, 0)
}

//sys	Unmount(target string, flags int) (err error) = SYS_UMOUNT2
//umount2(target *byte, flags int) int

//sys	Unshare(flags int) (err error)
//unshare(flags int) int

// FIXME: mksysinfo Ustat_t
// //sys	Ustat(dev int, ubuf *Ustat_t) (err error)
// //ustat(dev _dev_t, ubuf *Ustat_t) int
<commit_msg>syscall: Add Gettid on GNU/Linux systems.<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// GNU/Linux library calls.

package syscall

import "unsafe"

//sys	Openat(dirfd int, path string, flags int, mode uint32) (fd int, err error)
//openat(dirfd int, path *byte, flags int, mode Mode_t) int

//sys	futimesat(dirfd int, path *byte, times *[2]Timeval) (err error)
//futimesat(dirfd int, path *byte, times *[2]Timeval) int
func Futimesat(dirfd int, path string, tv []Timeval) (err error) {
	if len(tv) != 2 {
		return EINVAL
	}
	return futimesat(dirfd, StringBytePtr(path), (*[2]Timeval)(unsafe.Pointer(&tv[0])))
}

func Futimes(fd int, tv []Timeval) (err error) {
	// Believe it or not, this is the best we can do on GNU/Linux
	// (and is what glibc does).
	return Utimes("/proc/self/fd/"+itoa(fd), tv)
}

//sys	ptrace(request int, pid int, addr uintptr, data uintptr) (err error)
//ptrace(request int, pid Pid_t, addr *byte, data *byte) _C_long

//sysnb raw_ptrace(request int, pid int, addr *byte, data *byte) (err Errno)
//ptrace(request int, pid Pid_t, addr *byte, data *byte) _C_long

func ptracePeek(req int, pid int, addr uintptr, out []byte) (count int, err error) {
	// The peek requests are machine-size oriented, so we wrap it
	// to retrieve arbitrary-length data.

	// The ptrace syscall differs from glibc's ptrace.
	// Peeks returns the word in *data, not as the return value.

	var buf [sizeofPtr]byte

	// Leading edge.  PEEKTEXT/PEEKDATA don't require aligned
	// access (PEEKUSER warns that it might), but if we don't
	// align our reads, we might straddle an unmapped page
	// boundary and not get the bytes leading up to the page
	// boundary.
	n := 0
	if addr%sizeofPtr != 0 {
		err = ptrace(req, pid, addr-addr%sizeofPtr, uintptr(unsafe.Pointer(&buf[0])))
		if err != nil {
			return 0, err
		}
		n += copy(out, buf[addr%sizeofPtr:])
		out = out[n:]
	}

	// Remainder.
	for len(out) > 0 {
		// We use an internal buffer to gaurantee alignment.
		// It's not documented if this is necessary, but we're paranoid.
		err = ptrace(req, pid, addr+uintptr(n), uintptr(unsafe.Pointer(&buf[0])))
		if err != nil {
			return n, err
		}
		copied := copy(out, buf[0:])
		n += copied
		out = out[copied:]
	}

	return n, nil
}

func PtracePeekText(pid int, addr uintptr, out []byte) (count int, err error) {
	return ptracePeek(PTRACE_PEEKTEXT, pid, addr, out)
}

func PtracePeekData(pid int, addr uintptr, out []byte) (count int, err error) {
	return ptracePeek(PTRACE_PEEKDATA, pid, addr, out)
}

func ptracePoke(pokeReq int, peekReq int, pid int, addr uintptr, data []byte) (count int, err error) {
	// As for ptracePeek, we need to align our accesses to deal
	// with the possibility of straddling an invalid page.

	// Leading edge.
	n := 0
	if addr%sizeofPtr != 0 {
		var buf [sizeofPtr]byte
		err = ptrace(peekReq, pid, addr-addr%sizeofPtr, uintptr(unsafe.Pointer(&buf[0])))
		if err != nil {
			return 0, err
		}
		n += copy(buf[addr%sizeofPtr:], data)
		word := *((*uintptr)(unsafe.Pointer(&buf[0])))
		err = ptrace(pokeReq, pid, addr-addr%sizeofPtr, word)
		if err != nil {
			return 0, err
		}
		data = data[n:]
	}

	// Interior.
	for len(data) > int(sizeofPtr) {
		word := *((*uintptr)(unsafe.Pointer(&data[0])))
		err = ptrace(pokeReq, pid, addr+uintptr(n), word)
		if err != nil {
			return n, err
		}
		n += int(sizeofPtr)
		data = data[sizeofPtr:]
	}

	// Trailing edge.
	if len(data) > 0 {
		var buf [sizeofPtr]byte
		err = ptrace(peekReq, pid, addr+uintptr(n), uintptr(unsafe.Pointer(&buf[0])))
		if err != nil {
			return n, err
		}
		copy(buf[0:], data)
		word := *((*uintptr)(unsafe.Pointer(&buf[0])))
		err = ptrace(pokeReq, pid, addr+uintptr(n), word)
		if err != nil {
			return n, err
		}
		n += len(data)
	}

	return n, nil
}

func PtracePokeText(pid int, addr uintptr, data []byte) (count int, err error) {
	return ptracePoke(PTRACE_POKETEXT, PTRACE_PEEKTEXT, pid, addr, data)
}

func PtracePokeData(pid int, addr uintptr, data []byte) (count int, err error) {
	return ptracePoke(PTRACE_POKEDATA, PTRACE_PEEKDATA, pid, addr, data)
}

func PtraceSetOptions(pid int, options int) (err error) {
	return ptrace(PTRACE_SETOPTIONS, pid, 0, uintptr(options))
}

func PtraceGetEventMsg(pid int) (msg uint, err error) {
	var data _C_long
	err = ptrace(PTRACE_GETEVENTMSG, pid, 0, uintptr(unsafe.Pointer(&data)))
	msg = uint(data)
	return
}

func PtraceCont(pid int, signal int) (err error) {
	return ptrace(PTRACE_CONT, pid, 0, uintptr(signal))
}

func PtraceSingleStep(pid int) (err error) { return ptrace(PTRACE_SINGLESTEP, pid, 0, 0) }

func PtraceAttach(pid int) (err error) { return ptrace(PTRACE_ATTACH, pid, 0, 0) }

func PtraceDetach(pid int) (err error) { return ptrace(PTRACE_DETACH, pid, 0, 0) }

// FIXME: mksysinfo needs to produce LINUX_REBOOT_MAGIC[12].

// //sys	reboot(magic1 uint, magic2 uint, cmd int, arg string) (err error)
// //reboot(magic1 uint, magic2 uint, cmd int, arg *byte) int
// func Reboot(cmd int) (err error) {
// 	return reboot(LINUX_REBOOT_MAGIC1, LINUX_REBOOT_MAGIC2, cmd, "")
// }

//sys	Acct(path string) (err error)
//acct(path *byte) int

// FIXME: mksysinfo Timex
// //sys	Adjtimex(buf *Timex) (state int, err error)
// //adjtimex(buf *Timex) int

//sys	Faccessat(dirfd int, path string, mode uint32, flags int) (err error)
//faccessat(dirfd int, pathname *byte, mode int, flags int) int

// FIXME: Only in glibc 2.10 and later.
// //sys	Fallocate(fd int, mode uint32, off int64, len int64) (err error)
// //fallocate(fd int, mode int, offset Offset_t, len Offset_t) int

//sys	Fchmodat(dirfd int, path string, mode uint32, flags int) (err error)
//fchmodat(dirfd int, pathname *byte, mode Mode_t, flags int) int

//sys	Fchownat(dirfd int, path string, uid int, gid int, flags int) (err error)
//fchownat(dirfd int, path *byte, owner Uid_t, group Gid_t, flags int) int

//sys	Flock(fd int, how int) (err error)
//flock(fd int, how int) int

// FIXME: mksysinfo statfs
// //sys	Fstatfs(fd int, buf *Statfs_t) (err error)
// //fstatfs(fd int, buf *Statfs_t) int

func Gettid() (tid int) {
	r1, _, _ := Syscall(SYS_GETTID, 0, 0, 0)
	return int(r1)
}

// FIXME: mksysinfo linux_dirent
//    Or just abandon this function.
// //sys	Getdents(fd int, buf []byte) (n int, err error)
// //getdents64(fd int, buf *byte, count uint)

//sys	InotifyAddWatch(fd int, pathname string, mask uint32) (watchdesc int, err error)
//inotify_add_watch(fd int, pathname *byte, mask uint32) int

//sysnb	InotifyInit() (fd int, err error)
//inotify_init() int

// FIXME: Only in glibc 2.9 and later.
// //sysnb	InotifyInit1(flags int) (fd int, err error)
// //inotify_init1(flags int) int

//sysnb	InotifyRmWatch(fd int, watchdesc uint32) (success int, err error)
//inotify_rm_watch(fd int, wd uint32) int

//sys	Klogctl(typ int, buf []byte) (n int, err error)
//klogctl(typ int, bufp *byte, len int) int

//sys	Mkdirat(dirfd int, path string, mode uint32) (err error)
//mkdirat(dirfd int, path *byte, mode Mode_t) int

//sys	Mknodat(dirfd int, path string, mode uint32, dev int) (err error)
//mknodat(dirfd int, path *byte, mode Mode_t, dev _dev_t) int

//sys	PivotRoot(newroot string, putold string) (err error)
//pivot_root(newroot *byte, putold *byte) int

//sys	Renameat(olddirfd int, oldpath string, newdirfd int, newpath string) (err error)
//renameat(olddirfd int, oldpath *byte, newdirfd int, newpath *byte) int

//sys	sendfile(outfd int, infd int, offset *Offset_t, count int) (written int, err error)
//sendfile64(outfd int, infd int, offset *Offset_t, count Size_t) Ssize_t
func Sendfile(outfd int, infd int, offset *int64, count int) (written int, err error) {
	var soff Offset_t
	var psoff *Offset_t
	if offset != nil {
		psoff = &soff
	}
	written, err = sendfile(outfd, infd, psoff, count)
	if offset != nil {
		*offset = int64(soff)
	}
	return
}

//sys	Setfsgid(gid int) (err error)
//setfsgid(gid Gid_t) int

//sys	Setfsuid(uid int) (err error)
//setfsuid(uid Uid_t) int

//sysnb	Setresgid(rgid int, egid int, sgid int) (err error)
//setresgid(rgid Gid_t, egid Gid_t, sgid Gid_t) int

//sysnb	Setresuid(ruid int, eguid int, suid int) (err error)
//setresuid(ruid Uid_t, euid Uid_t, suid Uid_t) int

//sys	splice(rfd int, roff *_loff_t, wfd int, woff *_loff_t, len int, flags int) (n int64, err error)
//splice(rfd int, roff *_loff_t, wfd int, woff *_loff_t, len Size_t, flags uint) Ssize_t
func Splice(rfd int, roff *int64, wfd int, woff *int64, len int, flags int) (n int64, err error) {
	var lroff _loff_t
	var plroff *_loff_t
	if roff != nil {
		plroff = &lroff
	}
	var lwoff _loff_t
	var plwoff *_loff_t
	if woff != nil {
		plwoff = &lwoff
	}
	n, err = splice(rfd, plroff, wfd, plwoff, len, flags)
	if roff != nil {
		*roff = int64(lroff)
	}
	if woff != nil {
		*woff = int64(lwoff)
	}
	return
}

// FIXME: mksysinfo statfs
// //sys	Statfs(path string, buf *Statfs_t) (err error)
// //statfs(path *byte, buf *Statfs_t) int

// FIXME: Only in glibc 2.6 and later.
// //sys	SyncFileRange(fd int, off int64, n int64, flags int) (err error)
// //sync_file_range(fd int, off Offset_t, n Offset_t, flags uint) int

// FIXME: mksysinfo Sysinfo_t
// //sysnb	Sysinfo(info *Sysinfo_t) (err error)
// //sysinfo(info *Sysinfo_t) int

//sys	Tee(rfd int, wfd int, len int, flags int) (n int64, err error)
//tee(rfd int, wfd int, len Size_t, flags uint) Ssize_t

// FIXME: Only available as a syscall.
// //sysnb	Tgkill(tgid int, tid int, sig int) (err error)
// //tgkill(tgid int, tid int, sig int) int

//sys	unlinkat(dirfd int, path string, flags int) (err error)
//unlinkat(dirfd int, path *byte, flags int) int

func Unlinkat(dirfd int, path string) (err error) {
	return unlinkat(dirfd, path, 0)
}

//sys	Unmount(target string, flags int) (err error) = SYS_UMOUNT2
//umount2(target *byte, flags int) int

//sys	Unshare(flags int) (err error)
//unshare(flags int) int

// FIXME: mksysinfo Ustat_t
// //sys	Ustat(dev int, ubuf *Ustat_t) (err error)
// //ustat(dev _dev_t, ubuf *Ustat_t) int
<|endoftext|>


<commit_before>package datastore

import (
	"fmt"
	"log"
	"reflect"
	"strings"

	"github.com/docker/libkv"
	"github.com/docker/libkv/store"
	"github.com/docker/libkv/store/boltdb"
	"github.com/docker/libkv/store/consul"
	"github.com/docker/libkv/store/etcd"
	"github.com/docker/libkv/store/zookeeper"
	"github.com/docker/libnetwork/types"
)

//DataStore exported
type DataStore interface {
	// GetObject gets data from datastore and unmarshals to the specified object
	GetObject(key string, o KVObject) error
	// PutObject adds a new Record based on an object into the datastore
	PutObject(kvObject KVObject) error
	// PutObjectAtomic provides an atomic add and update operation for a Record
	PutObjectAtomic(kvObject KVObject) error
	// DeleteObject deletes a record
	DeleteObject(kvObject KVObject) error
	// DeleteObjectAtomic performs an atomic delete operation
	DeleteObjectAtomic(kvObject KVObject) error
	// DeleteTree deletes a record
	DeleteTree(kvObject KVObject) error
	// Watchable returns whether the store is watchable are not
	Watchable() bool
	// Watch for changes on a KVObject
	Watch(kvObject KVObject, stopCh <-chan struct{}) (<-chan KVObject, error)
	// List returns of a list of KVObjects belonging to the parent
	// key. The caller must pass a KVObject of the same type as
	// the objects that need to be listed
	List(string, KVObject) ([]KVObject, error)
	// Scope returns the scope of the store
	Scope() string
	// KVStore returns access to the KV Store
	KVStore() store.Store
	// Close closes the data store
	Close()
}

// ErrKeyModified is raised for an atomic update when the update is working on a stale state
var (
	ErrKeyModified = store.ErrKeyModified
	ErrKeyNotFound = store.ErrKeyNotFound
)

type datastore struct {
	scope string
	store store.Store
	cache *cache
}

// KVObject is  Key/Value interface used by objects to be part of the DataStore
type KVObject interface {
	// Key method lets an object to provide the Key to be used in KV Store
	Key() []string
	// KeyPrefix method lets an object to return immediate parent key that can be used for tree walk
	KeyPrefix() []string
	// Value method lets an object to marshal its content to be stored in the KV store
	Value() []byte
	// SetValue is used by the datastore to set the object's value when loaded from the data store.
	SetValue([]byte) error
	// Index method returns the latest DB Index as seen by the object
	Index() uint64
	// SetIndex method allows the datastore to store the latest DB Index into the object
	SetIndex(uint64)
	// True if the object exists in the datastore, false if it hasn't been stored yet.
	// When SetIndex() is called, the object has been stored.
	Exists() bool
	// DataScope indicates the storage scope of the KV object
	DataScope() string
	// Skip provides a way for a KV Object to avoid persisting it in the KV Store
	Skip() bool
}

// KVConstructor interface defines methods which can construct a KVObject from another.
type KVConstructor interface {
	// New returns a new object which is created based on the
	// source object
	New() KVObject
	// CopyTo deep copies the contents of the implementing object
	// to the passed destination object
	CopyTo(KVObject) error
}

// ScopeCfg represents Datastore configuration.
type ScopeCfg struct {
	Client ScopeClientCfg
}

// ScopeClientCfg represents Datastore Client-only mode configuration
type ScopeClientCfg struct {
	Provider string
	Address  string
	Config   *store.Config
}

const (
	// LocalScope indicates to store the KV object in local datastore such as boltdb
	LocalScope = "local"
	// GlobalScope indicates to store the KV object in global datastore such as consul/etcd/zookeeper
	GlobalScope   = "global"
	defaultPrefix = "/var/lib/docker/network/files"
)

const (
	// NetworkKeyPrefix is the prefix for network key in the kv store
	NetworkKeyPrefix = "network"
	// EndpointKeyPrefix is the prefix for endpoint key in the kv store
	EndpointKeyPrefix = "endpoint"
)

var (
	defaultScopes = makeDefaultScopes()
)

func makeDefaultScopes() map[string]*ScopeCfg {
	def := make(map[string]*ScopeCfg)
	def[LocalScope] = &ScopeCfg{
		Client: ScopeClientCfg{
			Provider: "boltdb",
			Address:  defaultPrefix + "/boltdb.db",
			Config: &store.Config{
				Bucket: "libnetwork",
			},
		},
	}

	return def
}

var rootChain = []string{"docker", "libnetwork"}

func init() {
	consul.Register()
	zookeeper.Register()
	etcd.Register()
	boltdb.Register()
}

// DefaultScopes returns a map of default scopes and it's config for clients to use.
func DefaultScopes(dataDir string) map[string]*ScopeCfg {
	if dataDir != "" {
		defaultScopes[LocalScope].Client.Address = dataDir + "/network/files/boltdb.db"
		return defaultScopes
	}

	defaultScopes[LocalScope].Client.Address = defaultPrefix + "/boltdb.db"
	return defaultScopes
}

// IsValid checks if the scope config has valid configuration.
func (cfg *ScopeCfg) IsValid() bool {
	if cfg == nil ||
		strings.TrimSpace(cfg.Client.Provider) == "" ||
		strings.TrimSpace(cfg.Client.Address) == "" {
		return false
	}

	return true
}

//Key provides convenient method to create a Key
func Key(key ...string) string {
	keychain := append(rootChain, key...)
	str := strings.Join(keychain, "/")
	return str + "/"
}

//ParseKey provides convenient method to unpack the key to complement the Key function
func ParseKey(key string) ([]string, error) {
	chain := strings.Split(strings.Trim(key, "/"), "/")

	// The key must atleast be equal to the rootChain in order to be considered as valid
	if len(chain) <= len(rootChain) || !reflect.DeepEqual(chain[0:len(rootChain)], rootChain) {
		return nil, types.BadRequestErrorf("invalid Key : %s", key)
	}
	return chain[len(rootChain):], nil
}

// newClient used to connect to KV Store
func newClient(scope string, kv string, addrs string, config *store.Config, cached bool) (DataStore, error) {
	if cached && scope != LocalScope {
		return nil, fmt.Errorf("caching supported only for scope %s", LocalScope)
	}

	if config == nil {
		config = &store.Config{}
	}
	store, err := libkv.NewStore(store.Backend(kv), []string{addrs}, config)
	if err != nil {
		return nil, err
	}

	ds := &datastore{scope: scope, store: store}
	if cached {
		ds.cache = newCache(ds)
	}

	return ds, nil
}

// NewDataStore creates a new instance of LibKV data store
func NewDataStore(scope string, cfg *ScopeCfg) (DataStore, error) {
	if cfg == nil || cfg.Client.Provider == "" || cfg.Client.Address == "" {
		c, ok := defaultScopes[scope]
		if !ok || c.Client.Provider == "" || c.Client.Address == "" {
			return nil, fmt.Errorf("unexpected scope %s without configuration passed", scope)
		}

		cfg = c
	}

	var cached bool
	if scope == LocalScope {
		cached = true
	}

	return newClient(scope, cfg.Client.Provider, cfg.Client.Address, cfg.Client.Config, cached)
}

func (ds *datastore) Close() {
	ds.store.Close()
}

func (ds *datastore) Scope() string {
	return ds.scope
}

func (ds *datastore) Watchable() bool {
	return ds.scope != LocalScope
}

func (ds *datastore) Watch(kvObject KVObject, stopCh <-chan struct{}) (<-chan KVObject, error) {
	sCh := make(chan struct{})

	ctor, ok := kvObject.(KVConstructor)
	if !ok {
		return nil, fmt.Errorf("error watching object type %T, object does not implement KVConstructor interface", kvObject)
	}

	kvpCh, err := ds.store.Watch(Key(kvObject.Key()...), sCh)
	if err != nil {
		return nil, err
	}

	kvoCh := make(chan KVObject)

	go func() {
		for {
			select {
			case <-stopCh:
				close(sCh)
				return
			case kvPair := <-kvpCh:
				dstO := ctor.New()

				if err := dstO.SetValue(kvPair.Value); err != nil {
					log.Printf("Could not unmarshal kvpair value = %s", string(kvPair.Value))
					break
				}

				dstO.SetIndex(kvPair.LastIndex)
				kvoCh <- dstO
			}
		}
	}()

	return kvoCh, nil
}

func (ds *datastore) KVStore() store.Store {
	return ds.store
}

// PutObjectAtomic adds a new Record based on an object into the datastore
func (ds *datastore) PutObjectAtomic(kvObject KVObject) error {
	var (
		previous *store.KVPair
		pair     *store.KVPair
		err      error
	)

	if kvObject == nil {
		return types.BadRequestErrorf("invalid KV Object : nil")
	}

	kvObjValue := kvObject.Value()

	if kvObjValue == nil {
		return types.BadRequestErrorf("invalid KV Object with a nil Value for key %s", Key(kvObject.Key()...))
	}

	if kvObject.Skip() {
		goto add_cache
	}

	if kvObject.Exists() {
		previous = &store.KVPair{Key: Key(kvObject.Key()...), LastIndex: kvObject.Index()}
	} else {
		previous = nil
	}

	_, pair, err = ds.store.AtomicPut(Key(kvObject.Key()...), kvObjValue, previous, nil)
	if err != nil {
		return err
	}

	kvObject.SetIndex(pair.LastIndex)

add_cache:
	if ds.cache != nil {
		return ds.cache.add(kvObject)
	}

	return nil
}

// PutObject adds a new Record based on an object into the datastore
func (ds *datastore) PutObject(kvObject KVObject) error {
	if kvObject == nil {
		return types.BadRequestErrorf("invalid KV Object : nil")
	}

	if kvObject.Skip() {
		goto add_cache
	}

	if err := ds.putObjectWithKey(kvObject, kvObject.Key()...); err != nil {
		return err
	}

add_cache:
	if ds.cache != nil {
		return ds.cache.add(kvObject)
	}

	return nil
}

func (ds *datastore) putObjectWithKey(kvObject KVObject, key ...string) error {
	kvObjValue := kvObject.Value()

	if kvObjValue == nil {
		return types.BadRequestErrorf("invalid KV Object with a nil Value for key %s", Key(kvObject.Key()...))
	}
	return ds.store.Put(Key(key...), kvObjValue, nil)
}

// GetObject returns a record matching the key
func (ds *datastore) GetObject(key string, o KVObject) error {
	if ds.cache != nil {
		return ds.cache.get(key, o)
	}

	kvPair, err := ds.store.Get(key)
	if err != nil {
		return err
	}

	if err := o.SetValue(kvPair.Value); err != nil {
		return err
	}

	// Make sure the object has a correct view of the DB index in
	// case we need to modify it and update the DB.
	o.SetIndex(kvPair.LastIndex)
	return nil
}

func (ds *datastore) ensureKey(key string) error {
	exists, err := ds.store.Exists(key)
	if err != nil {
		return err
	}
	if exists {
		return nil
	}
	return ds.store.Put(key, []byte{}, nil)
}

func (ds *datastore) List(key string, kvObject KVObject) ([]KVObject, error) {
	if ds.cache != nil {
		return ds.cache.list(kvObject)
	}

	// Bail out right away if the kvObject does not implement KVConstructor
	ctor, ok := kvObject.(KVConstructor)
	if !ok {
		return nil, fmt.Errorf("error listing objects, object does not implement KVConstructor interface")
	}

	// Make sure the parent key exists
	if err := ds.ensureKey(key); err != nil {
		return nil, err
	}

	kvList, err := ds.store.List(key)
	if err != nil {
		return nil, err
	}

	var kvol []KVObject
	for _, kvPair := range kvList {
		if len(kvPair.Value) == 0 {
			continue
		}

		dstO := ctor.New()
		if err := dstO.SetValue(kvPair.Value); err != nil {
			return nil, err
		}

		// Make sure the object has a correct view of the DB index in
		// case we need to modify it and update the DB.
		dstO.SetIndex(kvPair.LastIndex)

		kvol = append(kvol, dstO)
	}

	return kvol, nil
}

// DeleteObject unconditionally deletes a record from the store
func (ds *datastore) DeleteObject(kvObject KVObject) error {
	// cleaup the cache first
	if ds.cache != nil {
		ds.cache.del(kvObject)
	}

	if kvObject.Skip() {
		return nil
	}

	return ds.store.Delete(Key(kvObject.Key()...))
}

// DeleteObjectAtomic performs atomic delete on a record
func (ds *datastore) DeleteObjectAtomic(kvObject KVObject) error {
	if kvObject == nil {
		return types.BadRequestErrorf("invalid KV Object : nil")
	}

	previous := &store.KVPair{Key: Key(kvObject.Key()...), LastIndex: kvObject.Index()}

	if kvObject.Skip() {
		goto del_cache
	}

	if _, err := ds.store.AtomicDelete(Key(kvObject.Key()...), previous); err != nil {
		return err
	}

del_cache:
	// cleanup the cache only if AtomicDelete went through successfully
	if ds.cache != nil {
		return ds.cache.del(kvObject)
	}

	return nil
}

// DeleteTree unconditionally deletes a record from the store
func (ds *datastore) DeleteTree(kvObject KVObject) error {
	// cleaup the cache first
	if ds.cache != nil {
		ds.cache.del(kvObject)
	}

	if kvObject.Skip() {
		return nil
	}

	return ds.store.DeleteTree(Key(kvObject.KeyPrefix()...))
}
<commit_msg>Synchronize datastore apis<commit_after>package datastore

import (
	"fmt"
	"log"
	"reflect"
	"strings"
	"sync"

	"github.com/docker/libkv"
	"github.com/docker/libkv/store"
	"github.com/docker/libkv/store/boltdb"
	"github.com/docker/libkv/store/consul"
	"github.com/docker/libkv/store/etcd"
	"github.com/docker/libkv/store/zookeeper"
	"github.com/docker/libnetwork/types"
)

//DataStore exported
type DataStore interface {
	// GetObject gets data from datastore and unmarshals to the specified object
	GetObject(key string, o KVObject) error
	// PutObject adds a new Record based on an object into the datastore
	PutObject(kvObject KVObject) error
	// PutObjectAtomic provides an atomic add and update operation for a Record
	PutObjectAtomic(kvObject KVObject) error
	// DeleteObject deletes a record
	DeleteObject(kvObject KVObject) error
	// DeleteObjectAtomic performs an atomic delete operation
	DeleteObjectAtomic(kvObject KVObject) error
	// DeleteTree deletes a record
	DeleteTree(kvObject KVObject) error
	// Watchable returns whether the store is watchable are not
	Watchable() bool
	// Watch for changes on a KVObject
	Watch(kvObject KVObject, stopCh <-chan struct{}) (<-chan KVObject, error)
	// List returns of a list of KVObjects belonging to the parent
	// key. The caller must pass a KVObject of the same type as
	// the objects that need to be listed
	List(string, KVObject) ([]KVObject, error)
	// Scope returns the scope of the store
	Scope() string
	// KVStore returns access to the KV Store
	KVStore() store.Store
	// Close closes the data store
	Close()
}

// ErrKeyModified is raised for an atomic update when the update is working on a stale state
var (
	ErrKeyModified = store.ErrKeyModified
	ErrKeyNotFound = store.ErrKeyNotFound
)

type datastore struct {
	scope string
	store store.Store
	cache *cache
	sync.Mutex
}

// KVObject is  Key/Value interface used by objects to be part of the DataStore
type KVObject interface {
	// Key method lets an object to provide the Key to be used in KV Store
	Key() []string
	// KeyPrefix method lets an object to return immediate parent key that can be used for tree walk
	KeyPrefix() []string
	// Value method lets an object to marshal its content to be stored in the KV store
	Value() []byte
	// SetValue is used by the datastore to set the object's value when loaded from the data store.
	SetValue([]byte) error
	// Index method returns the latest DB Index as seen by the object
	Index() uint64
	// SetIndex method allows the datastore to store the latest DB Index into the object
	SetIndex(uint64)
	// True if the object exists in the datastore, false if it hasn't been stored yet.
	// When SetIndex() is called, the object has been stored.
	Exists() bool
	// DataScope indicates the storage scope of the KV object
	DataScope() string
	// Skip provides a way for a KV Object to avoid persisting it in the KV Store
	Skip() bool
}

// KVConstructor interface defines methods which can construct a KVObject from another.
type KVConstructor interface {
	// New returns a new object which is created based on the
	// source object
	New() KVObject
	// CopyTo deep copies the contents of the implementing object
	// to the passed destination object
	CopyTo(KVObject) error
}

// ScopeCfg represents Datastore configuration.
type ScopeCfg struct {
	Client ScopeClientCfg
}

// ScopeClientCfg represents Datastore Client-only mode configuration
type ScopeClientCfg struct {
	Provider string
	Address  string
	Config   *store.Config
}

const (
	// LocalScope indicates to store the KV object in local datastore such as boltdb
	LocalScope = "local"
	// GlobalScope indicates to store the KV object in global datastore such as consul/etcd/zookeeper
	GlobalScope   = "global"
	defaultPrefix = "/var/lib/docker/network/files"
)

const (
	// NetworkKeyPrefix is the prefix for network key in the kv store
	NetworkKeyPrefix = "network"
	// EndpointKeyPrefix is the prefix for endpoint key in the kv store
	EndpointKeyPrefix = "endpoint"
)

var (
	defaultScopes = makeDefaultScopes()
)

func makeDefaultScopes() map[string]*ScopeCfg {
	def := make(map[string]*ScopeCfg)
	def[LocalScope] = &ScopeCfg{
		Client: ScopeClientCfg{
			Provider: "boltdb",
			Address:  defaultPrefix + "/boltdb.db",
			Config: &store.Config{
				Bucket: "libnetwork",
			},
		},
	}

	return def
}

var rootChain = []string{"docker", "libnetwork"}

func init() {
	consul.Register()
	zookeeper.Register()
	etcd.Register()
	boltdb.Register()
}

// DefaultScopes returns a map of default scopes and it's config for clients to use.
func DefaultScopes(dataDir string) map[string]*ScopeCfg {
	if dataDir != "" {
		defaultScopes[LocalScope].Client.Address = dataDir + "/network/files/boltdb.db"
		return defaultScopes
	}

	defaultScopes[LocalScope].Client.Address = defaultPrefix + "/boltdb.db"
	return defaultScopes
}

// IsValid checks if the scope config has valid configuration.
func (cfg *ScopeCfg) IsValid() bool {
	if cfg == nil ||
		strings.TrimSpace(cfg.Client.Provider) == "" ||
		strings.TrimSpace(cfg.Client.Address) == "" {
		return false
	}

	return true
}

//Key provides convenient method to create a Key
func Key(key ...string) string {
	keychain := append(rootChain, key...)
	str := strings.Join(keychain, "/")
	return str + "/"
}

//ParseKey provides convenient method to unpack the key to complement the Key function
func ParseKey(key string) ([]string, error) {
	chain := strings.Split(strings.Trim(key, "/"), "/")

	// The key must atleast be equal to the rootChain in order to be considered as valid
	if len(chain) <= len(rootChain) || !reflect.DeepEqual(chain[0:len(rootChain)], rootChain) {
		return nil, types.BadRequestErrorf("invalid Key : %s", key)
	}
	return chain[len(rootChain):], nil
}

// newClient used to connect to KV Store
func newClient(scope string, kv string, addrs string, config *store.Config, cached bool) (DataStore, error) {
	if cached && scope != LocalScope {
		return nil, fmt.Errorf("caching supported only for scope %s", LocalScope)
	}

	if config == nil {
		config = &store.Config{}
	}
	store, err := libkv.NewStore(store.Backend(kv), []string{addrs}, config)
	if err != nil {
		return nil, err
	}

	ds := &datastore{scope: scope, store: store}
	if cached {
		ds.cache = newCache(ds)
	}

	return ds, nil
}

// NewDataStore creates a new instance of LibKV data store
func NewDataStore(scope string, cfg *ScopeCfg) (DataStore, error) {
	if cfg == nil || cfg.Client.Provider == "" || cfg.Client.Address == "" {
		c, ok := defaultScopes[scope]
		if !ok || c.Client.Provider == "" || c.Client.Address == "" {
			return nil, fmt.Errorf("unexpected scope %s without configuration passed", scope)
		}

		cfg = c
	}

	var cached bool
	if scope == LocalScope {
		cached = true
	}

	return newClient(scope, cfg.Client.Provider, cfg.Client.Address, cfg.Client.Config, cached)
}

func (ds *datastore) Close() {
	ds.store.Close()
}

func (ds *datastore) Scope() string {
	return ds.scope
}

func (ds *datastore) Watchable() bool {
	return ds.scope != LocalScope
}

func (ds *datastore) Watch(kvObject KVObject, stopCh <-chan struct{}) (<-chan KVObject, error) {
	sCh := make(chan struct{})

	ctor, ok := kvObject.(KVConstructor)
	if !ok {
		return nil, fmt.Errorf("error watching object type %T, object does not implement KVConstructor interface", kvObject)
	}

	kvpCh, err := ds.store.Watch(Key(kvObject.Key()...), sCh)
	if err != nil {
		return nil, err
	}

	kvoCh := make(chan KVObject)

	go func() {
		for {
			select {
			case <-stopCh:
				close(sCh)
				return
			case kvPair := <-kvpCh:
				dstO := ctor.New()

				if err := dstO.SetValue(kvPair.Value); err != nil {
					log.Printf("Could not unmarshal kvpair value = %s", string(kvPair.Value))
					break
				}

				dstO.SetIndex(kvPair.LastIndex)
				kvoCh <- dstO
			}
		}
	}()

	return kvoCh, nil
}

func (ds *datastore) KVStore() store.Store {
	return ds.store
}

// PutObjectAtomic adds a new Record based on an object into the datastore
func (ds *datastore) PutObjectAtomic(kvObject KVObject) error {
	var (
		previous *store.KVPair
		pair     *store.KVPair
		err      error
	)
	ds.Lock()
	defer ds.Unlock()

	if kvObject == nil {
		return types.BadRequestErrorf("invalid KV Object : nil")
	}

	kvObjValue := kvObject.Value()

	if kvObjValue == nil {
		return types.BadRequestErrorf("invalid KV Object with a nil Value for key %s", Key(kvObject.Key()...))
	}

	if kvObject.Skip() {
		goto add_cache
	}

	if kvObject.Exists() {
		previous = &store.KVPair{Key: Key(kvObject.Key()...), LastIndex: kvObject.Index()}
	} else {
		previous = nil
	}

	_, pair, err = ds.store.AtomicPut(Key(kvObject.Key()...), kvObjValue, previous, nil)
	if err != nil {
		return err
	}

	kvObject.SetIndex(pair.LastIndex)

add_cache:
	if ds.cache != nil {
		return ds.cache.add(kvObject)
	}

	return nil
}

// PutObject adds a new Record based on an object into the datastore
func (ds *datastore) PutObject(kvObject KVObject) error {
	ds.Lock()
	defer ds.Unlock()

	if kvObject == nil {
		return types.BadRequestErrorf("invalid KV Object : nil")
	}

	if kvObject.Skip() {
		goto add_cache
	}

	if err := ds.putObjectWithKey(kvObject, kvObject.Key()...); err != nil {
		return err
	}

add_cache:
	if ds.cache != nil {
		return ds.cache.add(kvObject)
	}

	return nil
}

func (ds *datastore) putObjectWithKey(kvObject KVObject, key ...string) error {
	kvObjValue := kvObject.Value()

	if kvObjValue == nil {
		return types.BadRequestErrorf("invalid KV Object with a nil Value for key %s", Key(kvObject.Key()...))
	}
	return ds.store.Put(Key(key...), kvObjValue, nil)
}

// GetObject returns a record matching the key
func (ds *datastore) GetObject(key string, o KVObject) error {
	ds.Lock()
	defer ds.Unlock()

	if ds.cache != nil {
		return ds.cache.get(key, o)
	}

	kvPair, err := ds.store.Get(key)
	if err != nil {
		return err
	}

	if err := o.SetValue(kvPair.Value); err != nil {
		return err
	}

	// Make sure the object has a correct view of the DB index in
	// case we need to modify it and update the DB.
	o.SetIndex(kvPair.LastIndex)
	return nil
}

func (ds *datastore) ensureKey(key string) error {
	exists, err := ds.store.Exists(key)
	if err != nil {
		return err
	}
	if exists {
		return nil
	}
	return ds.store.Put(key, []byte{}, nil)
}

func (ds *datastore) List(key string, kvObject KVObject) ([]KVObject, error) {
	ds.Lock()
	defer ds.Unlock()

	if ds.cache != nil {
		return ds.cache.list(kvObject)
	}

	// Bail out right away if the kvObject does not implement KVConstructor
	ctor, ok := kvObject.(KVConstructor)
	if !ok {
		return nil, fmt.Errorf("error listing objects, object does not implement KVConstructor interface")
	}

	// Make sure the parent key exists
	if err := ds.ensureKey(key); err != nil {
		return nil, err
	}

	kvList, err := ds.store.List(key)
	if err != nil {
		return nil, err
	}

	var kvol []KVObject
	for _, kvPair := range kvList {
		if len(kvPair.Value) == 0 {
			continue
		}

		dstO := ctor.New()
		if err := dstO.SetValue(kvPair.Value); err != nil {
			return nil, err
		}

		// Make sure the object has a correct view of the DB index in
		// case we need to modify it and update the DB.
		dstO.SetIndex(kvPair.LastIndex)

		kvol = append(kvol, dstO)
	}

	return kvol, nil
}

// DeleteObject unconditionally deletes a record from the store
func (ds *datastore) DeleteObject(kvObject KVObject) error {
	ds.Lock()
	defer ds.Unlock()

	// cleaup the cache first
	if ds.cache != nil {
		ds.cache.del(kvObject)
	}

	if kvObject.Skip() {
		return nil
	}

	return ds.store.Delete(Key(kvObject.Key()...))
}

// DeleteObjectAtomic performs atomic delete on a record
func (ds *datastore) DeleteObjectAtomic(kvObject KVObject) error {
	ds.Lock()
	defer ds.Unlock()

	if kvObject == nil {
		return types.BadRequestErrorf("invalid KV Object : nil")
	}

	previous := &store.KVPair{Key: Key(kvObject.Key()...), LastIndex: kvObject.Index()}

	if kvObject.Skip() {
		goto del_cache
	}

	if _, err := ds.store.AtomicDelete(Key(kvObject.Key()...), previous); err != nil {
		return err
	}

del_cache:
	// cleanup the cache only if AtomicDelete went through successfully
	if ds.cache != nil {
		return ds.cache.del(kvObject)
	}

	return nil
}

// DeleteTree unconditionally deletes a record from the store
func (ds *datastore) DeleteTree(kvObject KVObject) error {
	ds.Lock()
	defer ds.Unlock()

	// cleaup the cache first
	if ds.cache != nil {
		ds.cache.del(kvObject)
	}

	if kvObject.Skip() {
		return nil
	}

	return ds.store.DeleteTree(Key(kvObject.KeyPrefix()...))
}
<|endoftext|>


<commit_before>package tasks

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/HouzuoGuo/tiedot/db"
	"github.com/deckarep/gosx-notifier"
	"github.com/fatih/color"
	"github.com/olekukonko/tablewriter"
	"os"
	"strconv"
	"strings"
	"time"
)

const myDBDir = "/Users/dlt/.selfcontrol"

var (
	errDuplicateTimerForTask = errors.New("task already has a running timer")
	tasksCollection          *db.Col
	timersCollection         *db.Col
	DB                       *db.DB
	timers                   map[int][]*taskTimer
)

type task map[string]interface{}

type taskTimer struct {
	TaskID       int
	Timer        *time.Timer
	StartedAt    time.Time
	FinishedAt   time.Time
	Message      string
	Fired        bool
}

func (tt *taskTimer) trigger() {
	pushNotification(tt.Message)
	tt.FinishedAt = time.Now()
	tt.Fired = true
}

func (tt *taskTimer) ellapsedTime() time.Duration {
	var finishedAt time.Time
	if tt.unfinished() {
		finishedAt = time.Now()
	} else {
		finishedAt = tt.FinishedAt
	}
	return finishedAt.Sub(tt.StartedAt)
}

func (timer *taskTimer) unfinished() bool {
	return !timer.Fired
}

func init() {
	d, err := db.OpenDB(myDBDir)
	DB = d
	if err != nil {
		panic(err)
	}
	DB.Create("Tasks")
	DB.Create("Timers")
	tasksCollection = DB.Use("Tasks")
	timersCollection = DB.Use("Timers")
	loadTimers()
}

func loadTimers() {
	timers = make(map[int][]*taskTimer)
	timersCollection.ForEachDoc(func(id int, raw []byte) (willMoveOn bool) {
		var doc map[string]interface{}
		json.Unmarshal(raw, &doc)
		startedAt, err := time.Parse(time.RFC3339, doc["StartedAt"].(string))
		if err != nil {
			panic(err)
		}
		finishedAt, err := time.Parse(time.RFC3339, doc["FinishedAt"].(string))
		if err != nil {
			panic(err)
		}
		timer := &taskTimer{
			TaskID:       id,
			Fired:        doc["Fired"].(bool),
			StartedAt:    startedAt,
			FinishedAt:   finishedAt,
		}
		timers[id] = append(timers[id], timer)
		fmt.Println(timer)
		return true
	})
}

func Save() {
	for taskID, taskTimers := range timers {
		for _, timer := range taskTimers {
			timersDoc := map[string]interface{}{
				"TaskID":       string(taskID),
				"Fired":        timer.Fired,
				"Message": 	timer.Message,
				"StartedAt":    timer.StartedAt,
				"FinishedAt":   timer.FinishedAt,
			}
			_, err := timersCollection.Insert(timersDoc)
			if err != nil {
				panic(err)
			}
		}
	}
}

// Create a new task given its name
func Create(name string) {
	task := map[string]interface{}{
		"name":   name,
		"status": "TODO",
	}
	_, err := tasksCollection.Insert(task)
	if err != nil {
		panic(err)
	}
}

// Print all tasks in a ASCII table
func Print() {
	table := tablewriter.NewWriter(os.Stdout)
	table.SetHeader([]string{"ID", "name", "status", "total time"})
	table.SetAutoFormatHeaders(false)
	for _, row := range createRows() {
		table.Append(row)
	}
	table.Render()
}

// Delete a task with given id
func Delete(id int) {
	if err := tasksCollection.Delete(id); err != nil {
		panic(err)
	}
}

// UpdateFields updates the tasks attributes
func UpdateFields(taskID int, fieldValuePairs []string) (bool, error) {
	task, err := tasksCollection.Read(taskID)
	if err != nil {
		return false, err
	}
	for _, pair := range fieldValuePairs {
		arr := strings.Split(pair, ":")
		field, value := arr[0], arr[1]
		if field == "status" {
			value = strings.ToUpper(value)
		}
		task[field] = value
	}
	if err = tasksCollection.Update(taskID, task); err != nil {
		return false, err
	}
	return true, nil
}

// AddTimerForTask adds a timer for a given task id and duration
func AddTimerForTask(taskID int, d time.Duration) (bool, error) {
	task, err := tasksCollection.Read(taskID)
	if err != nil {
		panic(err)
	}

	if hasRunningTimer(taskID) {
		return false, errDuplicateTimerForTask
	}

	name := task["name"].(string)
	timer := time.NewTimer(d)
	message := "Timer for '" + name + "' finished!"
	tTimer := &taskTimer{
		TaskID:     taskID,
		Timer:      timer,
		Message:    message,
		StartedAt:  time.Now(),
		FinishedAt: time.Time{},
	}

	if timers[taskID] == nil {
		timers[taskID] = make([]*taskTimer, 0)
	}

	timers[taskID] = append(timers[taskID], tTimer)

	go func() {
		<-timer.C
		tTimer.trigger()
		Save()
		Print()
	}()
	return true, nil
}

func hasRunningTimer(taskID int) bool {
	for _, timer := range timers[taskID] {
		if timer.unfinished() {
			return true
		}
	}
	return false
}

func pushNotification(message string) {
	gosxnotifier.NewNotification(message).Push()
}

func createRows() [][]string {
	var rows = [][]string{}
	tasksCollection.ForEachDoc(func(id int, raw []byte) (willMoveOn bool) {
		var doc task
		json.Unmarshal(raw, &doc)
		name := doc["name"].(string)
		status := doc["status"].(string)
		row := []string{
			strconv.Itoa(id),
			name,
			coloredStatus(status),
			totalRunningTime(id).String(),
		}
		rows = append(rows, row)
		return true
	})
	return rows
}

func totalRunningTime(taskID int) time.Duration {
	loadTimers()
	total := time.Duration(0)
	for _, timer := range timers[taskID] {
		total += timer.ellapsedTime()
	}
	return total
}

func coloredStatus(status string) string {
	switch status {
	case "TODO":
		return color.CyanString(status)
	case "DOING":
		return color.YellowString(status)
	case "DONE":
		return color.GreenString(status)
	}
	return ""
}
<commit_msg>Persit timers<commit_after>package tasks

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/HouzuoGuo/tiedot/db"
	"github.com/deckarep/gosx-notifier"
	"github.com/fatih/color"
	"github.com/olekukonko/tablewriter"
	"os"
	"strconv"
	"strings"
	"time"
)

const myDBDir = "/Users/dlt/.selfcontrol"

var (
	errDuplicateTimerForTask = errors.New("task already has a running timer")
	tasksCollection          *db.Col
	timersCollection         *db.Col
	DB                       *db.DB
	timers                   map[int][]*taskTimer
)

type task map[string]interface{}

type taskTimer struct {
	TaskID       int
	Timer        *time.Timer
	StartedAt    time.Time
	FinishedAt   time.Time
	Message      string
	Fired        bool
}

func (tt *taskTimer) trigger() {
	pushNotification(tt.Message)
	tt.FinishedAt = time.Now()
	tt.Fired = true
}

func (tt *taskTimer) ellapsedTime() time.Duration {
	var finishedAt time.Time
	if tt.unfinished() {
		finishedAt = time.Now()
	} else {
		finishedAt = tt.FinishedAt
	}
	return finishedAt.Sub(tt.StartedAt)
}

func (timer *taskTimer) unfinished() bool {
	return !timer.Fired
}

func init() {
	d, err := db.OpenDB(myDBDir)
	DB = d
	if err != nil {
		panic(err)
	}
	DB.Create("Tasks")
	DB.Create("Timers")
	tasksCollection = DB.Use("Tasks")
	timersCollection = DB.Use("Timers")
	timers = make(map[int][]*taskTimer)
	loadTimers()
}

func loadTimers() {
	timersCollection.ForEachDoc(func(id int, raw []byte) (willMoveOn bool) {
		var doc map[string]interface{}
		json.Unmarshal(raw, &doc)
		startedAt, err := time.Parse(time.RFC3339, doc["StartedAt"].(string))
		if err != nil {
			panic(err)
		}
		finishedAt, err := time.Parse(time.RFC3339, doc["FinishedAt"].(string))
		if err != nil {
			panic(err)
		}
		timer := &taskTimer{
			TaskID:       id,
			Fired:        doc["Fired"].(bool),
			StartedAt:    startedAt,
			FinishedAt:   finishedAt,
		}
		timers[id] = append(timers[id], timer)
		fmt.Println(timer)
		return true
	})
}

// Create a new task given its name
func Create(name string) {
	task := map[string]interface{}{
		"name":   name,
		"status": "TODO",
	}
	_, err := tasksCollection.Insert(task)
	if err != nil {
		panic(err)
	}
}

// Print all tasks in a ASCII table
func Print() {
	table := tablewriter.NewWriter(os.Stdout)
	table.SetHeader([]string{"ID", "name", "status", "total time"})
	table.SetAutoFormatHeaders(false)
	for _, row := range createRows() {
		table.Append(row)
	}
	table.Render()
}

// Delete a task with given id
func Delete(id int) {
	if err := tasksCollection.Delete(id); err != nil {
		panic(err)
	}
}

// UpdateFields updates the tasks attributes
func UpdateFields(taskID int, fieldValuePairs []string) (bool, error) {
	task, err := tasksCollection.Read(taskID)
	if err != nil {
		return false, err
	}
	for _, pair := range fieldValuePairs {
		arr := strings.Split(pair, ":")
		field, value := arr[0], arr[1]
		if field == "status" {
			value = strings.ToUpper(value)
		}
		task[field] = value
	}
	if err = tasksCollection.Update(taskID, task); err != nil {
		return false, err
	}
	return true, nil
}

// AddTimerForTask adds a timer for a given task id and duration
func AddTimerForTask(taskID int, d time.Duration) (bool, error) {
	task, err := tasksCollection.Read(taskID)
	if err != nil {
		panic(err)
	}

	if hasRunningTimer(taskID) {
		return false, errDuplicateTimerForTask
	}

	name := task["name"].(string)
	timer := time.NewTimer(d)
	message := "Timer for '" + name + "' finished!"
	tTimer := &taskTimer{
		TaskID:     taskID,
		Timer:      timer,
		Message:    message,
		StartedAt:  time.Now(),
		FinishedAt: time.Time{},
	}

	if timers[taskID] == nil {
		timers[taskID] = make([]*taskTimer, 0)
	}

	timers[taskID] = append(timers[taskID], tTimer)

	go func() {
		<-timer.C
		tTimer.trigger()
		Save()
		Print()
	}()
	return true, nil
}

func Save() {
	for taskID, taskTimers := range timers {
	fmt.Println("Saving timers for taskid: " + string(taskID))
		for _, timer := range taskTimers {
			timersDoc := map[string]interface{}{
				"TaskID":       string(taskID),
				"Fired":        timer.Fired,
				"Message": 	timer.Message,
				"StartedAt":    timer.StartedAt,
				"FinishedAt":   timer.FinishedAt,
			}
			_, err := timersCollection.Insert(timersDoc)
			if err != nil {
				panic(err)
			}
		}
	}
}

func hasRunningTimer(taskID int) bool {
	for _, timer := range timers[taskID] {
		if timer.unfinished() {
			return true
		}
	}
	return false
}

func pushNotification(message string) {
	gosxnotifier.NewNotification(message).Push()
}

func createRows() [][]string {
	var rows = [][]string{}
	tasksCollection.ForEachDoc(func(id int, raw []byte) (willMoveOn bool) {
		var doc task
		json.Unmarshal(raw, &doc)
		name := doc["name"].(string)
		status := doc["status"].(string)
		row := []string{
			strconv.Itoa(id),
			name,
			coloredStatus(status),
			totalRunningTime(id).String(),
		}
		rows = append(rows, row)
		return true
	})
	return rows
}

func totalRunningTime(taskID int) time.Duration {
	loadTimers()
	total := time.Duration(0)
	for _, timer := range timers[taskID] {
		total += timer.ellapsedTime()
	}
	return total
}

func coloredStatus(status string) string {
	switch status {
	case "TODO":
		return color.CyanString(status)
	case "DOING":
		return color.YellowString(status)
	case "DONE":
		return color.GreenString(status)
	}
	return ""
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/sagemaker"
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation"
)

func resourceAwsSagemakerCodeRepository() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsSagemakerCodeRepositoryCreate,
		Read:   resourceAwsSagemakerCodeRepositoryRead,
		Update: resourceAwsSagemakerCodeRepositoryUpdate,
		Delete: resourceAwsSagemakerCodeRepositoryDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"arn": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"code_repository_name": {
				Type:         schema.TypeString,
				Required:     true,
				ForceNew:     true,
				ValidateFunc: validation.StringLenBetween(1, 63),
			},
			"git_config": {
				Type:     schema.TypeList,
				Required: true,
				MaxItems: 1,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"repository_url": {
							Type:     schema.TypeString,
							ForceNew: true,
							Required: true,
						},
						"branch": {
							Type:         schema.TypeString,
							ForceNew:     true,
							Optional:     true,
							ValidateFunc: validation.StringLenBetween(1, 1024),
						},
						"secret_arn": {
							Type:         schema.TypeString,
							Optional:     true,
							ValidateFunc: validateArn,
						},
					},
				},
			},
		},
	}
}

func resourceAwsSagemakerCodeRepositoryCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).sagemakerconn

	name := d.Get("code_repository_name").(string)

	createOpts := &sagemaker.CreateCodeRepositoryInput{
		CodeRepositoryName: aws.String(name),
		GitConfig:          expandSagemakerCodeRepositoryGitConfig(d.Get("git_config").([]interface{})),
	}

	log.Printf("[DEBUG] sagemaker code repository create config: %#v", *createOpts)
	_, err := conn.CreateCodeRepository(createOpts)
	if err != nil {
		return fmt.Errorf("error creating SageMaker code repository: %w", err)
	}

	d.SetId(name)

	return resourceAwsSagemakerCodeRepositoryRead(d, meta)
}

func resourceAwsSagemakerCodeRepositoryRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).sagemakerconn

	input := &sagemaker.DescribeCodeRepositoryInput{
		CodeRepositoryName: aws.String(d.Id()),
	}

	codeRepository, err := conn.DescribeCodeRepository(input)
	if err != nil {
		if isAWSErr(err, "ValidationException", "Cannot find CodeRepository") {
			d.SetId("")
			log.Printf("[WARN] Unable to find SageMaker code repository (%s); removing from state", d.Id())
			return nil
		}
		return fmt.Errorf("error reading SageMaker code repository (%s): %w", d.Id(), err)

	}

	d.Set("code_repository_name", codeRepository.CodeRepositoryName)
	d.Set("arn", codeRepository.CodeRepositoryArn)

	if err := d.Set("git_config", flattenSagemakerCodeRepositoryGitConfig(codeRepository.GitConfig)); err != nil {
		return fmt.Errorf("error setting git_config for sagemaker code repository (%s): %w", d.Id(), err)
	}

	return nil
}

func resourceAwsSagemakerCodeRepositoryUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).sagemakerconn

	input := &sagemaker.UpdateCodeRepositoryInput{
		CodeRepositoryName: aws.String(d.Id()),
		GitConfig:          expandSagemakerCodeRepositoryUpdateGitConfig(d.Get("git_config").([]interface{})),
	}

	log.Printf("[DEBUG] sagemaker code repository update config: %#v", *input)
	_, err := conn.UpdateCodeRepository(input)
	if err != nil {
		return fmt.Errorf("error updating SageMaker code repository: %w", err)
	}

	return resourceAwsSagemakerCodeRepositoryRead(d, meta)
}

func resourceAwsSagemakerCodeRepositoryDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).sagemakerconn

	input := &sagemaker.DeleteCodeRepositoryInput{
		CodeRepositoryName: aws.String(d.Id()),
	}

	if _, err := conn.DeleteCodeRepository(input); err != nil {
		if isAWSErr(err, "ValidationException", "Cannot find CodeRepository") {
			return nil
		}
		return fmt.Errorf("error deleting SageMaker code repository (%s): %w", d.Id(), err)
	}

	return nil
}

func expandSagemakerCodeRepositoryGitConfig(l []interface{}) *sagemaker.GitConfig {
	if len(l) == 0 || l[0] == nil {
		return nil
	}

	m := l[0].(map[string]interface{})

	config := &sagemaker.GitConfig{
		RepositoryUrl: aws.String(m["repository_url"].(string)),
	}

	if v, ok := m["branch"].(string); ok && v != "" {
		config.Branch = aws.String(v)
	}

	if v, ok := m["secret_arn"].(string); ok && v != "" {
		config.SecretArn = aws.String(v)
	}

	return config
}

func flattenSagemakerCodeRepositoryGitConfig(config *sagemaker.GitConfig) []map[string]interface{} {
	if config == nil {
		return []map[string]interface{}{}
	}

	m := map[string]interface{}{
		"repository_url": aws.StringValue(config.RepositoryUrl),
	}

	if config.Branch != nil {
		m["branch"] = aws.StringValue(config.Branch)
	}

	if config.SecretArn != nil {
		m["secret_arn"] = aws.StringValue(config.SecretArn)
	}

	return []map[string]interface{}{m}
}

func expandSagemakerCodeRepositoryUpdateGitConfig(l []interface{}) *sagemaker.GitConfigForUpdate {
	if len(l) == 0 || l[0] == nil {
		return nil
	}

	m := l[0].(map[string]interface{})

	config := &sagemaker.GitConfigForUpdate{
		SecretArn: aws.String(m["secret_arn"].(string)),
	}

	return config
}
<commit_msg>add validation to name<commit_after>package aws

import (
	"fmt"
	"log"
	"regexp"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/sagemaker"
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation"
)

func resourceAwsSagemakerCodeRepository() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsSagemakerCodeRepositoryCreate,
		Read:   resourceAwsSagemakerCodeRepositoryRead,
		Update: resourceAwsSagemakerCodeRepositoryUpdate,
		Delete: resourceAwsSagemakerCodeRepositoryDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"arn": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"code_repository_name": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
				ValidateFunc: validation.All(
					validation.StringLenBetween(1, 63),
					validation.StringMatch(regexp.MustCompile(`^[a-zA-Z0-9](-*[a-zA-Z0-9])*$`), "Valid characters are a-z, A-Z, 0-9, and - (hyphen)."),
				),
			},
			"git_config": {
				Type:     schema.TypeList,
				Required: true,
				MaxItems: 1,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"repository_url": {
							Type:     schema.TypeString,
							ForceNew: true,
							Required: true,
						},
						"branch": {
							Type:         schema.TypeString,
							ForceNew:     true,
							Optional:     true,
							ValidateFunc: validation.StringLenBetween(1, 1024),
						},
						"secret_arn": {
							Type:         schema.TypeString,
							Optional:     true,
							ValidateFunc: validateArn,
						},
					},
				},
			},
		},
	}
}

func resourceAwsSagemakerCodeRepositoryCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).sagemakerconn

	name := d.Get("code_repository_name").(string)

	createOpts := &sagemaker.CreateCodeRepositoryInput{
		CodeRepositoryName: aws.String(name),
		GitConfig:          expandSagemakerCodeRepositoryGitConfig(d.Get("git_config").([]interface{})),
	}

	log.Printf("[DEBUG] sagemaker code repository create config: %#v", *createOpts)
	_, err := conn.CreateCodeRepository(createOpts)
	if err != nil {
		return fmt.Errorf("error creating SageMaker code repository: %w", err)
	}

	d.SetId(name)

	return resourceAwsSagemakerCodeRepositoryRead(d, meta)
}

func resourceAwsSagemakerCodeRepositoryRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).sagemakerconn

	input := &sagemaker.DescribeCodeRepositoryInput{
		CodeRepositoryName: aws.String(d.Id()),
	}

	codeRepository, err := conn.DescribeCodeRepository(input)
	if err != nil {
		if isAWSErr(err, "ValidationException", "Cannot find CodeRepository") {
			d.SetId("")
			log.Printf("[WARN] Unable to find SageMaker code repository (%s); removing from state", d.Id())
			return nil
		}
		return fmt.Errorf("error reading SageMaker code repository (%s): %w", d.Id(), err)

	}

	d.Set("code_repository_name", codeRepository.CodeRepositoryName)
	d.Set("arn", codeRepository.CodeRepositoryArn)

	if err := d.Set("git_config", flattenSagemakerCodeRepositoryGitConfig(codeRepository.GitConfig)); err != nil {
		return fmt.Errorf("error setting git_config for sagemaker code repository (%s): %w", d.Id(), err)
	}

	return nil
}

func resourceAwsSagemakerCodeRepositoryUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).sagemakerconn

	input := &sagemaker.UpdateCodeRepositoryInput{
		CodeRepositoryName: aws.String(d.Id()),
		GitConfig:          expandSagemakerCodeRepositoryUpdateGitConfig(d.Get("git_config").([]interface{})),
	}

	log.Printf("[DEBUG] sagemaker code repository update config: %#v", *input)
	_, err := conn.UpdateCodeRepository(input)
	if err != nil {
		return fmt.Errorf("error updating SageMaker code repository: %w", err)
	}

	return resourceAwsSagemakerCodeRepositoryRead(d, meta)
}

func resourceAwsSagemakerCodeRepositoryDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).sagemakerconn

	input := &sagemaker.DeleteCodeRepositoryInput{
		CodeRepositoryName: aws.String(d.Id()),
	}

	if _, err := conn.DeleteCodeRepository(input); err != nil {
		if isAWSErr(err, "ValidationException", "Cannot find CodeRepository") {
			return nil
		}
		return fmt.Errorf("error deleting SageMaker code repository (%s): %w", d.Id(), err)
	}

	return nil
}

func expandSagemakerCodeRepositoryGitConfig(l []interface{}) *sagemaker.GitConfig {
	if len(l) == 0 || l[0] == nil {
		return nil
	}

	m := l[0].(map[string]interface{})

	config := &sagemaker.GitConfig{
		RepositoryUrl: aws.String(m["repository_url"].(string)),
	}

	if v, ok := m["branch"].(string); ok && v != "" {
		config.Branch = aws.String(v)
	}

	if v, ok := m["secret_arn"].(string); ok && v != "" {
		config.SecretArn = aws.String(v)
	}

	return config
}

func flattenSagemakerCodeRepositoryGitConfig(config *sagemaker.GitConfig) []map[string]interface{} {
	if config == nil {
		return []map[string]interface{}{}
	}

	m := map[string]interface{}{
		"repository_url": aws.StringValue(config.RepositoryUrl),
	}

	if config.Branch != nil {
		m["branch"] = aws.StringValue(config.Branch)
	}

	if config.SecretArn != nil {
		m["secret_arn"] = aws.StringValue(config.SecretArn)
	}

	return []map[string]interface{}{m}
}

func expandSagemakerCodeRepositoryUpdateGitConfig(l []interface{}) *sagemaker.GitConfigForUpdate {
	if len(l) == 0 || l[0] == nil {
		return nil
	}

	m := l[0].(map[string]interface{})

	config := &sagemaker.GitConfigForUpdate{
		SecretArn: aws.String(m["secret_arn"].(string)),
	}

	return config
}
<|endoftext|>