<commit_before>package commands

import (
	"fmt"
	"io"
	"sort"
	"strings"
	"text/tabwriter"
	"unicode"

	cid "gx/ipfs/QmPSQnBKM9g7BaUcZCvswUJVscQ1ipjmwxN5PXCjkp9EQ7/go-cid"
	mhash "gx/ipfs/QmPnFwZ2JXKnXgMw8CdBPxn7FWh6LLdjUjxV1fKHuJnkr8/go-multihash"
	cidutil "gx/ipfs/QmQJSeE3CX4zos9qeaG8EhecEK9zvrTEfTG84J8C5NVRwt/go-cidutil"
	cmdkit "gx/ipfs/QmSP88ryZkHSRn1fnngAaV2Vcn63WUJzAavnRM9CVdU1Ky/go-ipfs-cmdkit"
	verifcid "gx/ipfs/QmVkMRSkXrpjqrroEXWuYBvDBnXCdMMY6gsKicBGVGUqKT/go-verifcid"
	cmds "gx/ipfs/QmXTmUCBtDUrzDYVzASogLiNph7EBuYqEgPL7QoHNMzUnz/go-ipfs-cmds"
	mbase "gx/ipfs/QmekxXDhCxCJRNuzmHreuaT3BsuJcsjcXWNrtV9C8DRHtd/go-multibase"
)

var CidCmd = &cmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "Convert and discover properties of CIDs",
	},
	Subcommands: map[string]*cmds.Command{
		"format": cidFmtCmd,
		"base32": base32Cmd,
		"bases":  basesCmd,
		"codecs": codecsCmd,
		"hashes": hashesCmd,
	},
}

var cidFmtCmd = &cmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "Format and convert a CID in various useful ways.",
		LongDescription: `
Format and converts <cid>'s in various useful ways.

The optional format string is a printf style format string:
` + cidutil.FormatRef,
	},
	Arguments: []cmdkit.Argument{
		cmdkit.StringArg("cid", true, true, "Cids to format."),
	},
	Options: []cmdkit.Option{
		cmdkit.StringOption("f", "Printf style format string.").WithDefault("%s"),
		cmdkit.StringOption("v", "CID version to convert to."),
		cmdkit.StringOption("b", "Multibase to display CID in."),
	},
	Run: func(req *cmds.Request, resp cmds.ResponseEmitter, env cmds.Environment) error {
		fmtStr, _ := req.Options["f"].(string)
		verStr, _ := req.Options["v"].(string)
		baseStr, _ := req.Options["b"].(string)

		opts := cidFormatOpts{}

		if strings.IndexByte(fmtStr, '%') == -1 {
			return fmt.Errorf("invalid format string: %s", fmtStr)
		}
		opts.fmtStr = fmtStr

		switch verStr {
		case "":
			// noop
		case "0":
			opts.verConv = toCidV0
		case "1":
			opts.verConv = toCidV1
		default:
			return fmt.Errorf("invalid cid version: %s\n", verStr)
		}

		if baseStr != "" {
			encoder, err := mbase.EncoderByName(baseStr)
			if err != nil {
				return err
			}
			opts.newBase = encoder.Encoding()
		} else {
			opts.newBase = mbase.Encoding(-1)
		}

		return emitCids(req, resp, opts)
	},
	PostRun: cmds.PostRunMap{
		cmds.CLI: streamRes(func(v interface{}, out io.Writer) nonFatalError {
			r := v.(*CidFormatRes)
			if r.ErrorMsg != "" {
				return nonFatalError(fmt.Sprintf("%s: %s", r.CidStr, r.ErrorMsg))
			}
			fmt.Fprintf(out, "%s\n", r.Formatted)
			return ""
		}),
	},
	Type: CidFormatRes{},
}

type CidFormatRes struct {
	CidStr    string // Original Cid String passed in
	Formatted string // Formated Result
	ErrorMsg  string // Error
}

var base32Cmd = &cmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "Convert CIDs to Base32 CID version 1.",
	},
	Arguments: []cmdkit.Argument{
		cmdkit.StringArg("cid", true, true, "Cids to convert.").EnableStdin(),
	},
	Run: func(req *cmds.Request, resp cmds.ResponseEmitter, env cmds.Environment) error {
		opts := cidFormatOpts{
			fmtStr:  "%s",
			newBase: mbase.Encoding(mbase.Base32),
			verConv: toCidV1,
		}
		return emitCids(req, resp, opts)
	},
	PostRun: cidFmtCmd.PostRun,
	Type:    cidFmtCmd.Type,
}

type cidFormatOpts struct {
	fmtStr  string
	newBase mbase.Encoding
	verConv func(cid cid.Cid) (cid.Cid, error)
}

func emitCids(req *cmds.Request, resp cmds.ResponseEmitter, opts cidFormatOpts) error {
	for _, cidStr := range req.Arguments {
		emit := func(fmtd string, err error) {
			res := &CidFormatRes{CidStr: cidStr, Formatted: fmtd}
			if err != nil {
				res.ErrorMsg = err.Error()
			}
			resp.Emit(res)
		}
		c, err := cid.Decode(cidStr)
		if err != nil {
			emit("", err)
			continue
		}
		base := opts.newBase
		if base == -1 {
			base, _ = cid.ExtractEncoding(cidStr)
		}
		if opts.verConv != nil {
			c, err = opts.verConv(c)
			if err != nil {
				emit("", err)
				continue
			}
		}
		str, err := cidutil.Format(opts.fmtStr, base, c)
		if _, ok := err.(cidutil.FormatStringError); ok {
			// no point in continuing if there is a problem with the format string
			return err
		}
		emit(str, err)
	}
	return nil
}

func toCidV0(c cid.Cid) (cid.Cid, error) {
	if c.Type() != cid.DagProtobuf {
		return cid.Cid{}, fmt.Errorf("can't convert non-protobuf nodes to cidv0")
	}
	return cid.NewCidV0(c.Hash()), nil
}

func toCidV1(c cid.Cid) (cid.Cid, error) {
	return cid.NewCidV1(c.Type(), c.Hash()), nil
}

type CodeAndName struct {
	Code int
	Name string
}

var basesCmd = &cmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "List available multibase encodings.",
	},
	Options: []cmdkit.Option{
		cmdkit.BoolOption("prefix", "also include the single leter prefixes in addition to the code"),
		cmdkit.BoolOption("numeric", "also include numeric codes"),
	},
	Run: func(req *cmds.Request, resp cmds.ResponseEmitter, env cmds.Environment) error {
		var res []CodeAndName
		// use EncodingToStr in case at some point there are multiple names for a given code
		for code, name := range mbase.EncodingToStr {
			res = append(res, CodeAndName{int(code), name})
		}
		cmds.EmitOnce(resp, res)
		return nil
	},
	Encoders: cmds.EncoderMap{
		cmds.Text: cmds.MakeEncoder(func(req *cmds.Request, w0 io.Writer, val0 interface{}) error {
			w := tabwriter.NewWriter(w0, 0, 0, 2, ' ', 0)
			prefixes, _ := req.Options["prefix"].(bool)
			numeric, _ := req.Options["numeric"].(bool)
			val := val0.([]CodeAndName)
			sort.Sort(multibaseSorter{val})
			for _, v := range val {
				if prefixes && v.Code >= 32 && v.Code < 127 {
					fmt.Fprintf(w, "%c\t", v.Code)
				} else if prefixes {
					// don't display non-printable prefixes
					fmt.Fprintf(w, "\t")
				}
				if numeric {
					fmt.Fprintf(w, "%d\t%s\n", v.Code, v.Name)
				} else {
					fmt.Fprintf(w, "%s\n", v.Name)
				}
			}
			w.Flush()
			return nil
		}),
	},
	Type: []CodeAndName{},
}

var codecsCmd = &cmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "List available CID codecs.",
	},
	Options: []cmdkit.Option{
		cmdkit.BoolOption("numeric", "also include numeric codes"),
	},
	Run: func(req *cmds.Request, resp cmds.ResponseEmitter, env cmds.Environment) error {
		var res []CodeAndName
		// use CodecToStr as there are multiple names for a given code
		for code, name := range cid.CodecToStr {
			res = append(res, CodeAndName{int(code), name})
		}
		cmds.EmitOnce(resp, res)
		return nil
	},
	Encoders: cmds.EncoderMap{
		cmds.Text: cmds.MakeEncoder(func(req *cmds.Request, w0 io.Writer, val0 interface{}) error {
			w := tabwriter.NewWriter(w0, 0, 0, 2, ' ', 0)
			numeric, _ := req.Options["numeric"].(bool)
			val := val0.([]CodeAndName)
			sort.Sort(codeAndNameSorter{val})
			for _, v := range val {
				if numeric {
					fmt.Fprintf(w, "%d\t%s\n", v.Code, v.Name)
				} else {
					fmt.Fprintf(w, "%s\n", v.Name)
				}
			}
			w.Flush()
			return nil
		}),
	},
	Type: []CodeAndName{},
}

var hashesCmd = &cmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "List available multihashes.",
	},
	Options: codecsCmd.Options,
	Run: func(req *cmds.Request, resp cmds.ResponseEmitter, env cmds.Environment) error {
		var res []CodeAndName
		// use mhash.Codes in case at some point there are multiple names for a given code
		for code, name := range mhash.Codes {
			if !verifcid.IsGoodHash(code) {
				continue
			}
			res = append(res, CodeAndName{int(code), name})
		}
		cmds.EmitOnce(resp, res)
		return nil
	},
	Encoders: codecsCmd.Encoders,
	Type:     codecsCmd.Type,
}

type multibaseSorter struct {
	data []CodeAndName
}

func (s multibaseSorter) Len() int      { return len(s.data) }
func (s multibaseSorter) Swap(i, j int) { s.data[i], s.data[j] = s.data[j], s.data[i] }

func (s multibaseSorter) Less(i, j int) bool {
	a := unicode.ToLower(rune(s.data[i].Code))
	b := unicode.ToLower(rune(s.data[j].Code))
	if a != b {
		return a < b
	}
	// lowecase letters should come before uppercase
	return s.data[i].Code > s.data[j].Code
}

type codeAndNameSorter struct {
	data []CodeAndName
}

func (s codeAndNameSorter) Len() int           { return len(s.data) }
func (s codeAndNameSorter) Swap(i, j int)      { s.data[i], s.data[j] = s.data[j], s.data[i] }
func (s codeAndNameSorter) Less(i, j int) bool { return s.data[i].Code < s.data[j].Code }
<commit_msg>Right align numbers in "ipfs cid bases|codecs|hashes" output.<commit_after>package commands

import (
	"fmt"
	"io"
	"sort"
	"strings"
	"unicode"

	cid "gx/ipfs/QmPSQnBKM9g7BaUcZCvswUJVscQ1ipjmwxN5PXCjkp9EQ7/go-cid"
	mhash "gx/ipfs/QmPnFwZ2JXKnXgMw8CdBPxn7FWh6LLdjUjxV1fKHuJnkr8/go-multihash"
	cidutil "gx/ipfs/QmQJSeE3CX4zos9qeaG8EhecEK9zvrTEfTG84J8C5NVRwt/go-cidutil"
	cmdkit "gx/ipfs/QmSP88ryZkHSRn1fnngAaV2Vcn63WUJzAavnRM9CVdU1Ky/go-ipfs-cmdkit"
	verifcid "gx/ipfs/QmVkMRSkXrpjqrroEXWuYBvDBnXCdMMY6gsKicBGVGUqKT/go-verifcid"
	cmds "gx/ipfs/QmXTmUCBtDUrzDYVzASogLiNph7EBuYqEgPL7QoHNMzUnz/go-ipfs-cmds"
	mbase "gx/ipfs/QmekxXDhCxCJRNuzmHreuaT3BsuJcsjcXWNrtV9C8DRHtd/go-multibase"
)

var CidCmd = &cmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "Convert and discover properties of CIDs",
	},
	Subcommands: map[string]*cmds.Command{
		"format": cidFmtCmd,
		"base32": base32Cmd,
		"bases":  basesCmd,
		"codecs": codecsCmd,
		"hashes": hashesCmd,
	},
}

var cidFmtCmd = &cmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "Format and convert a CID in various useful ways.",
		LongDescription: `
Format and converts <cid>'s in various useful ways.

The optional format string is a printf style format string:
` + cidutil.FormatRef,
	},
	Arguments: []cmdkit.Argument{
		cmdkit.StringArg("cid", true, true, "Cids to format."),
	},
	Options: []cmdkit.Option{
		cmdkit.StringOption("f", "Printf style format string.").WithDefault("%s"),
		cmdkit.StringOption("v", "CID version to convert to."),
		cmdkit.StringOption("b", "Multibase to display CID in."),
	},
	Run: func(req *cmds.Request, resp cmds.ResponseEmitter, env cmds.Environment) error {
		fmtStr, _ := req.Options["f"].(string)
		verStr, _ := req.Options["v"].(string)
		baseStr, _ := req.Options["b"].(string)

		opts := cidFormatOpts{}

		if strings.IndexByte(fmtStr, '%') == -1 {
			return fmt.Errorf("invalid format string: %s", fmtStr)
		}
		opts.fmtStr = fmtStr

		switch verStr {
		case "":
			// noop
		case "0":
			opts.verConv = toCidV0
		case "1":
			opts.verConv = toCidV1
		default:
			return fmt.Errorf("invalid cid version: %s\n", verStr)
		}

		if baseStr != "" {
			encoder, err := mbase.EncoderByName(baseStr)
			if err != nil {
				return err
			}
			opts.newBase = encoder.Encoding()
		} else {
			opts.newBase = mbase.Encoding(-1)
		}

		return emitCids(req, resp, opts)
	},
	PostRun: cmds.PostRunMap{
		cmds.CLI: streamRes(func(v interface{}, out io.Writer) nonFatalError {
			r := v.(*CidFormatRes)
			if r.ErrorMsg != "" {
				return nonFatalError(fmt.Sprintf("%s: %s", r.CidStr, r.ErrorMsg))
			}
			fmt.Fprintf(out, "%s\n", r.Formatted)
			return ""
		}),
	},
	Type: CidFormatRes{},
}

type CidFormatRes struct {
	CidStr    string // Original Cid String passed in
	Formatted string // Formated Result
	ErrorMsg  string // Error
}

var base32Cmd = &cmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "Convert CIDs to Base32 CID version 1.",
	},
	Arguments: []cmdkit.Argument{
		cmdkit.StringArg("cid", true, true, "Cids to convert.").EnableStdin(),
	},
	Run: func(req *cmds.Request, resp cmds.ResponseEmitter, env cmds.Environment) error {
		opts := cidFormatOpts{
			fmtStr:  "%s",
			newBase: mbase.Encoding(mbase.Base32),
			verConv: toCidV1,
		}
		return emitCids(req, resp, opts)
	},
	PostRun: cidFmtCmd.PostRun,
	Type:    cidFmtCmd.Type,
}

type cidFormatOpts struct {
	fmtStr  string
	newBase mbase.Encoding
	verConv func(cid cid.Cid) (cid.Cid, error)
}

func emitCids(req *cmds.Request, resp cmds.ResponseEmitter, opts cidFormatOpts) error {
	for _, cidStr := range req.Arguments {
		emit := func(fmtd string, err error) {
			res := &CidFormatRes{CidStr: cidStr, Formatted: fmtd}
			if err != nil {
				res.ErrorMsg = err.Error()
			}
			resp.Emit(res)
		}
		c, err := cid.Decode(cidStr)
		if err != nil {
			emit("", err)
			continue
		}
		base := opts.newBase
		if base == -1 {
			base, _ = cid.ExtractEncoding(cidStr)
		}
		if opts.verConv != nil {
			c, err = opts.verConv(c)
			if err != nil {
				emit("", err)
				continue
			}
		}
		str, err := cidutil.Format(opts.fmtStr, base, c)
		if _, ok := err.(cidutil.FormatStringError); ok {
			// no point in continuing if there is a problem with the format string
			return err
		}
		emit(str, err)
	}
	return nil
}

func toCidV0(c cid.Cid) (cid.Cid, error) {
	if c.Type() != cid.DagProtobuf {
		return cid.Cid{}, fmt.Errorf("can't convert non-protobuf nodes to cidv0")
	}
	return cid.NewCidV0(c.Hash()), nil
}

func toCidV1(c cid.Cid) (cid.Cid, error) {
	return cid.NewCidV1(c.Type(), c.Hash()), nil
}

type CodeAndName struct {
	Code int
	Name string
}

var basesCmd = &cmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "List available multibase encodings.",
	},
	Options: []cmdkit.Option{
		cmdkit.BoolOption("prefix", "also include the single leter prefixes in addition to the code"),
		cmdkit.BoolOption("numeric", "also include numeric codes"),
	},
	Run: func(req *cmds.Request, resp cmds.ResponseEmitter, env cmds.Environment) error {
		var res []CodeAndName
		// use EncodingToStr in case at some point there are multiple names for a given code
		for code, name := range mbase.EncodingToStr {
			res = append(res, CodeAndName{int(code), name})
		}
		cmds.EmitOnce(resp, res)
		return nil
	},
	Encoders: cmds.EncoderMap{
		cmds.Text: cmds.MakeEncoder(func(req *cmds.Request, w io.Writer, val0 interface{}) error {
			prefixes, _ := req.Options["prefix"].(bool)
			numeric, _ := req.Options["numeric"].(bool)
			val := val0.([]CodeAndName)
			sort.Sort(multibaseSorter{val})
			for _, v := range val {
				code := v.Code
				if code < 32 || code >= 127 {
					// don't display non-printable prefixes
					code = ' '
				}
				switch {
				case prefixes && numeric:
					fmt.Fprintf(w, "%c %5d  %s\n", code, v.Code, v.Name)
				case prefixes:
					fmt.Fprintf(w, "%c  %s\n", code, v.Name)
				case numeric:
					fmt.Fprintf(w, "%5d  %s\n", v.Code, v.Name)
				default:
					fmt.Fprintf(w, "%s\n", v.Name)
				}
			}
			return nil
		}),
	},
	Type: []CodeAndName{},
}

var codecsCmd = &cmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "List available CID codecs.",
	},
	Options: []cmdkit.Option{
		cmdkit.BoolOption("numeric", "also include numeric codes"),
	},
	Run: func(req *cmds.Request, resp cmds.ResponseEmitter, env cmds.Environment) error {
		var res []CodeAndName
		// use CodecToStr as there are multiple names for a given code
		for code, name := range cid.CodecToStr {
			res = append(res, CodeAndName{int(code), name})
		}
		cmds.EmitOnce(resp, res)
		return nil
	},
	Encoders: cmds.EncoderMap{
		cmds.Text: cmds.MakeEncoder(func(req *cmds.Request, w io.Writer, val0 interface{}) error {
			numeric, _ := req.Options["numeric"].(bool)
			val := val0.([]CodeAndName)
			sort.Sort(codeAndNameSorter{val})
			for _, v := range val {
				if numeric {
					fmt.Fprintf(w, "%5d  %s\n", v.Code, v.Name)
				} else {
					fmt.Fprintf(w, "%s\n", v.Name)
				}
			}
			return nil
		}),
	},
	Type: []CodeAndName{},
}

var hashesCmd = &cmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "List available multihashes.",
	},
	Options: codecsCmd.Options,
	Run: func(req *cmds.Request, resp cmds.ResponseEmitter, env cmds.Environment) error {
		var res []CodeAndName
		// use mhash.Codes in case at some point there are multiple names for a given code
		for code, name := range mhash.Codes {
			if !verifcid.IsGoodHash(code) {
				continue
			}
			res = append(res, CodeAndName{int(code), name})
		}
		cmds.EmitOnce(resp, res)
		return nil
	},
	Encoders: codecsCmd.Encoders,
	Type:     codecsCmd.Type,
}

type multibaseSorter struct {
	data []CodeAndName
}

func (s multibaseSorter) Len() int      { return len(s.data) }
func (s multibaseSorter) Swap(i, j int) { s.data[i], s.data[j] = s.data[j], s.data[i] }

func (s multibaseSorter) Less(i, j int) bool {
	a := unicode.ToLower(rune(s.data[i].Code))
	b := unicode.ToLower(rune(s.data[j].Code))
	if a != b {
		return a < b
	}
	// lowecase letters should come before uppercase
	return s.data[i].Code > s.data[j].Code
}

type codeAndNameSorter struct {
	data []CodeAndName
}

func (s codeAndNameSorter) Len() int           { return len(s.data) }
func (s codeAndNameSorter) Swap(i, j int)      { s.data[i], s.data[j] = s.data[j], s.data[i] }
func (s codeAndNameSorter) Less(i, j int) bool { return s.data[i].Code < s.data[j].Code }
<|endoftext|>


<commit_before>package auth

import (
	"database/sql"

	"github.com/techjanitor/pram-libs/db"
	e "github.com/techjanitor/pram-libs/errors"
)

// user struct
type User struct {
	Id              uint   `json:"id"`
	Name            string `json:"name"`
	Email           string `json:"email"`
	Group           uint   `json:"group"`
	IsModerator     bool   `json:"moderator"`
	IsAdmin         bool   `json:"admin"`
	IsAuthenticated bool   `json:"-"`
	IsConfirmed     bool   `json:"-"`
	IsLocked        bool   `json:"-"`
	IsBanned        bool   `json:"-"`
}

// get the user info from id
func (u *User) Info() (err error) {

	// this needs an id
	if u.Id == 0 || u.Id == 1 {
		return e.ErrInvalidParam
	}

	// Get Database handle
	dbase, err := db.GetDb()
	if err != nil {
		panic(err)
	}

	// get data from users table
	err = dbase.QueryRow(`SELECT user_group_map.usergroup_id,user_name,user_email,user_confirmed,user_locked,user_banned
    FROM users
    INNER JOIN user_group_map ON user_group_map.user_id = users.user_id
    WHERE users.user_id = ?`, u.Id).Scan(&u.Group, &u.Name, &u.Email, &u.IsConfirmed, &u.IsLocked, &u.IsBanned)
	if err == sql.ErrNoRows {
		return e.ErrUserNotExist
	} else if err != nil {
		return e.ErrInternalError
	}

	// if account is not confirmed
	if !u.IsConfirmed {
		return e.ErrUserNotConfirmed
	}

	// if locked
	if u.IsLocked {
		return e.ErrUserLocked
	}

	// if banned
	if u.IsBanned {
		return e.ErrUserBanned
	}

	// mark authenticated
	u.IsAuthenticated = true

	return

}
<commit_msg>new user group map<commit_after>package auth

import (
	"database/sql"

	"github.com/techjanitor/pram-libs/db"
	e "github.com/techjanitor/pram-libs/errors"
)

// user struct
type User struct {
	Id              uint   `json:"id"`
	Name            string `json:"name"`
	Email           string `json:"email"`
	Group           uint   `json:"group"`
	Avatar          string `json:"avatar"`
	IsModerator     bool   `json:"moderator"`
	IsAdmin         bool   `json:"admin"`
	IsAuthenticated bool   `json:"-"`
	IsConfirmed     bool   `json:"-"`
	IsLocked        bool   `json:"-"`
	IsBanned        bool   `json:"-"`
}

// get the user info from id
func (u *User) Info() (err error) {

	// this needs an id
	if u.Id == 0 || u.Id == 1 {
		return e.ErrInvalidParam
	}

	// Get Database handle
	dbase, err := db.GetDb()
	if err != nil {
		panic(err)
	}

	// get data from users table
	err = dbase.QueryRow(`SELECT user_group_map.usergroup_id,user_name,user_email,user_confirmed,user_locked,user_banned,user_avatar
    FROM users
    INNER JOIN user_group_map ON user_group_map.user_id = users.user_id
    WHERE users.user_id = ?`, u.Id).Scan(&u.Group, &u.Name, &u.Email, &u.IsConfirmed, &u.IsLocked, &u.IsBanned, &u.Avatar)
	if err == sql.ErrNoRows {
		return e.ErrUserNotExist
	} else if err != nil {
		return e.ErrInternalError
	}

	// if account is not confirmed
	if !u.IsConfirmed {
		return e.ErrUserNotConfirmed
	}

	// if locked
	if u.IsLocked {
		return e.ErrUserLocked
	}

	// if banned
	if u.IsBanned {
		return e.ErrUserBanned
	}

	// mark authenticated
	u.IsAuthenticated = true

	return

}
<|endoftext|>


<commit_before>/*
Core contains the main functionality of the Livepeer node.

The logical orgnization of the `core` module is as follows:

livepeernode.go: Main struct definition and code that is common to all node types.
broadcaster.go: Code that is called only when the node is in broadcaster mode.
orchestrator.go: Code that is called only when the node is in orchestrator mode.

*/
package core

import (
	"context"
	"errors"
	"net/url"
	"sync"

	"github.com/livepeer/go-livepeer/common"
	"github.com/livepeer/go-livepeer/eth"
	"github.com/livepeer/go-livepeer/ipfs"
)

var ErrLivepeerNode = errors.New("ErrLivepeerNode")
var ErrTranscode = errors.New("ErrTranscode")
var DefaultJobLength = int64(5760) //Avg 1 day in 15 sec blocks
var LivepeerVersion = "0.2.4-unstable"

type NodeType int

const (
	Broadcaster NodeType = iota
	Transcoder
)

//LivepeerNode handles videos going in and coming out of the Livepeer network.
type LivepeerNode struct {

	// Common fields
	VideoSource     VideoSource
	Eth             eth.LivepeerEthClient
	EthEventMonitor eth.EventMonitor
	EthServices     map[string]eth.EventService
	WorkDir         string
	NodeType        NodeType
	Database        *common.DB

	// Transcoder public fields
	ClaimManagers map[int64]eth.ClaimManager
	SegmentChans  map[int64]SegmentChan
	Ipfs          ipfs.IpfsApi
	ServiceURI    *url.URL

	// Transcoder private fields
	claimMutex   *sync.Mutex
	segmentMutex *sync.Mutex
}

//NewLivepeerNode creates a new Livepeer Node. Eth can be nil.
func NewLivepeerNode(e eth.LivepeerEthClient, wd string, dbh *common.DB) (*LivepeerNode, error) {

	return &LivepeerNode{VideoSource: NewBasicVideoSource(), Eth: e, WorkDir: wd, Database: dbh, EthServices: make(map[string]eth.EventService), ClaimManagers: make(map[int64]eth.ClaimManager), SegmentChans: make(map[int64]SegmentChan), claimMutex: &sync.Mutex{}, segmentMutex: &sync.Mutex{}}, nil
}

func (n *LivepeerNode) StartEthServices() error {
	var err error
	for k, s := range n.EthServices {
		// Skip BlockService until the end
		if k == "BlockService" {
			continue
		}
		err = s.Start(context.Background())
		if err != nil {
			return err
		}
	}

	// Make sure to initialize BlockService last so other services can
	// create filters starting from the last seen block
	if s, ok := n.EthServices["BlockService"]; ok {
		if err := s.Start(context.Background()); err != nil {
			return err
		}
	}

	return nil
}

func (n *LivepeerNode) StopEthServices() error {
	var err error
	for _, s := range n.EthServices {
		err = s.Stop()
		if err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>version bump<commit_after>/*
Core contains the main functionality of the Livepeer node.

The logical orgnization of the `core` module is as follows:

livepeernode.go: Main struct definition and code that is common to all node types.
broadcaster.go: Code that is called only when the node is in broadcaster mode.
orchestrator.go: Code that is called only when the node is in orchestrator mode.

*/
package core

import (
	"context"
	"errors"
	"net/url"
	"sync"

	"github.com/livepeer/go-livepeer/common"
	"github.com/livepeer/go-livepeer/eth"
	"github.com/livepeer/go-livepeer/ipfs"
)

var ErrLivepeerNode = errors.New("ErrLivepeerNode")
var ErrTranscode = errors.New("ErrTranscode")
var DefaultJobLength = int64(5760) //Avg 1 day in 15 sec blocks
var LivepeerVersion = "0.3.0-unstable"

type NodeType int

const (
	Broadcaster NodeType = iota
	Transcoder
)

//LivepeerNode handles videos going in and coming out of the Livepeer network.
type LivepeerNode struct {

	// Common fields
	VideoSource     VideoSource
	Eth             eth.LivepeerEthClient
	EthEventMonitor eth.EventMonitor
	EthServices     map[string]eth.EventService
	WorkDir         string
	NodeType        NodeType
	Database        *common.DB

	// Transcoder public fields
	ClaimManagers map[int64]eth.ClaimManager
	SegmentChans  map[int64]SegmentChan
	Ipfs          ipfs.IpfsApi
	ServiceURI    *url.URL

	// Transcoder private fields
	claimMutex   *sync.Mutex
	segmentMutex *sync.Mutex
}

//NewLivepeerNode creates a new Livepeer Node. Eth can be nil.
func NewLivepeerNode(e eth.LivepeerEthClient, wd string, dbh *common.DB) (*LivepeerNode, error) {

	return &LivepeerNode{VideoSource: NewBasicVideoSource(), Eth: e, WorkDir: wd, Database: dbh, EthServices: make(map[string]eth.EventService), ClaimManagers: make(map[int64]eth.ClaimManager), SegmentChans: make(map[int64]SegmentChan), claimMutex: &sync.Mutex{}, segmentMutex: &sync.Mutex{}}, nil
}

func (n *LivepeerNode) StartEthServices() error {
	var err error
	for k, s := range n.EthServices {
		// Skip BlockService until the end
		if k == "BlockService" {
			continue
		}
		err = s.Start(context.Background())
		if err != nil {
			return err
		}
	}

	// Make sure to initialize BlockService last so other services can
	// create filters starting from the last seen block
	if s, ok := n.EthServices["BlockService"]; ok {
		if err := s.Start(context.Background()); err != nil {
			return err
		}
	}

	return nil
}

func (n *LivepeerNode) StopEthServices() error {
	var err error
	for _, s := range n.EthServices {
		err = s.Stop()
		if err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package network

import (
	"context"
	"io"

	ic "github.com/libp2p/go-libp2p-core/crypto"
	"github.com/libp2p/go-libp2p-core/peer"

	ma "github.com/multiformats/go-multiaddr"
)

// Conn is a connection to a remote peer. It multiplexes streams.
// Usually there is no need to use a Conn directly, but it may
// be useful to get information about the peer on the other side:
//  stream.Conn().RemotePeer()
type Conn interface {
	io.Closer

	ConnSecurity
	ConnMultiaddrs

	// ID returns an identifier that uniquely identifies this Conn within this
	// host, during this run. Connection IDs may repeat across restarts.
	ID() string

	// NewStream constructs a new Stream over this conn.
	NewStream(context.Context) (Stream, error)

	// GetStreams returns all open streams over this conn.
	GetStreams() []Stream

	// Stat stores metadata pertaining to this conn.
	Stat() Stat
}

// ConnSecurity is the interface that one can mix into a connection interface to
// give it the security methods.
type ConnSecurity interface {
	// LocalPeer returns our peer ID
	LocalPeer() peer.ID

	// LocalPrivateKey returns our private key
	LocalPrivateKey() ic.PrivKey

	// RemotePeer returns the peer ID of the remote peer.
	RemotePeer() peer.ID

	// RemotePublicKey returns the public key of the remote peer.
	RemotePublicKey() ic.PubKey
}

// ConnMultiaddrs is an interface mixin for connection types that provide multiaddr
// addresses for the endpoints.
type ConnMultiaddrs interface {
	// LocalMultiaddr returns the local Multiaddr associated
	// with this connection
	LocalMultiaddr() ma.Multiaddr

	// RemoteMultiaddr returns the remote Multiaddr associated
	// with this connection
	RemoteMultiaddr() ma.Multiaddr
}
<commit_msg>abstract Conn Stat interface for threading<commit_after>package network

import (
	"context"
	"io"

	ic "github.com/libp2p/go-libp2p-core/crypto"
	"github.com/libp2p/go-libp2p-core/peer"

	ma "github.com/multiformats/go-multiaddr"
)

// Conn is a connection to a remote peer. It multiplexes streams.
// Usually there is no need to use a Conn directly, but it may
// be useful to get information about the peer on the other side:
//  stream.Conn().RemotePeer()
type Conn interface {
	io.Closer

	ConnSecurity
	ConnMultiaddrs
	ConnStat

	// ID returns an identifier that uniquely identifies this Conn within this
	// host, during this run. Connection IDs may repeat across restarts.
	ID() string

	// NewStream constructs a new Stream over this conn.
	NewStream(context.Context) (Stream, error)

	// GetStreams returns all open streams over this conn.
	GetStreams() []Stream
}

// ConnSecurity is the interface that one can mix into a connection interface to
// give it the security methods.
type ConnSecurity interface {
	// LocalPeer returns our peer ID
	LocalPeer() peer.ID

	// LocalPrivateKey returns our private key
	LocalPrivateKey() ic.PrivKey

	// RemotePeer returns the peer ID of the remote peer.
	RemotePeer() peer.ID

	// RemotePublicKey returns the public key of the remote peer.
	RemotePublicKey() ic.PubKey
}

// ConnMultiaddrs is an interface mixin for connection types that provide multiaddr
// addresses for the endpoints.
type ConnMultiaddrs interface {
	// LocalMultiaddr returns the local Multiaddr associated
	// with this connection
	LocalMultiaddr() ma.Multiaddr

	// RemoteMultiaddr returns the remote Multiaddr associated
	// with this connection
	RemoteMultiaddr() ma.Multiaddr
}

type ConnStat interface {
	// Stat stores metadata pertaining to this conn.
	Stat() Stat
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"time"

	"github.com/franela/goreq"
	"github.com/ncodes/cocoon/core/config"
	"github.com/ncodes/cocoon/core/orderer/cmd"
	logging "github.com/op/go-logging"
)

var log *logging.Logger

func init() {
	config.ConfigureLogger()
	log = logging.MustGetLogger("main")
	goreq.SetConnectTimeout(5 * time.Second)
}

func main() {
	// defer profile.Start(profile.MemProfile).Stop()

	// go func() {
	// 	go func() {
	// 		log.Info(http.ListenAndServe("localhost:6060", nil).Error())
	// 	}()

	// 	for {
	// 		var mem runtime.MemStats
	// 		runtime.ReadMemStats(&mem)
	// 		log.Infof("Alloc: %d, Total Alloc: %d, HeapAlloc: %d, HeapSys: %d", mem.Alloc, mem.TotalAlloc, mem.HeapAlloc, mem.HeapSys)
	// 		time.Sleep(10 * time.Second)
	// 	}
	// }()

	if err := cmd.RootCmd.Execute(); err != nil {
		log.Error(err.Error())
		os.Exit(-1)
	}
}
<commit_msg>disabled grpc api<commit_after>package main

import (
	"os"
	"time"

	"google.golang.org/grpc/grpclog"

	"github.com/franela/goreq"
	"github.com/ncodes/cocoon/core/common"
	"github.com/ncodes/cocoon/core/config"
	"github.com/ncodes/cocoon/core/orderer/cmd"
	logging "github.com/op/go-logging"
)

var log *logging.Logger

func init() {
	config.ConfigureLogger()
	log = logging.MustGetLogger("main")
	goreq.SetConnectTimeout(5 * time.Second)
	if len(os.Getenv("ENABLE_GRPC_LOG")) == 0 {
		gl := common.GLogger{}
		gl.Disable(true, true)
		grpclog.SetLogger(&gl)
	}
}

func main() {
	// defer profile.Start(profile.MemProfile).Stop()

	// go func() {
	// 	go func() {
	// 		log.Info(http.ListenAndServe("localhost:6060", nil).Error())
	// 	}()

	// 	for {
	// 		var mem runtime.MemStats
	// 		runtime.ReadMemStats(&mem)
	// 		log.Infof("Alloc: %d, Total Alloc: %d, HeapAlloc: %d, HeapSys: %d", mem.Alloc, mem.TotalAlloc, mem.HeapAlloc, mem.HeapSys)
	// 		time.Sleep(10 * time.Second)
	// 	}
	// }()

	if err := cmd.RootCmd.Execute(); err != nil {
		log.Error(err.Error())
		os.Exit(-1)
	}
}
<|endoftext|>


<commit_before>package server

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"strconv"

	"github.com/gorilla/mux"
)

type Pattern struct {
	Blocks      []BlockLedger      `json:"blocks"`
	Connections []ConnectionLedger `json:"connections"`
	Groups      []Group            `json:"groups"`
}

type Node interface {
	GetID() int
	GetParent() *Group
	SetParent(*Group)
}

type Group struct {
	Id       int    `json:"id"`
	Label    string `json:"label"`
	Children []int  `json:"children"`
	Parent   *Group `json:"-"`
}

func (g *Group) GetID() int {
	return g.Id
}

func (g *Group) GetParent() *Group {
	return g.Parent
}

func (g *Group) SetParent(group *Group) {
	g.Parent = group
}

func (s *Server) ListGroups() []Group {
	groups := []Group{}
	for _, g := range s.groups {
		groups = append(groups, *g)
	}
	return groups
}

func (s *Server) GroupIndexHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.WriteHeader(http.StatusOK)
	if err := json.NewEncoder(w).Encode(s.ListGroups()); err != nil {
		panic(err)
	}
}

func (s *Server) DetachChild(g Node) error {
	parent := g.GetParent()
	if parent == nil {
		return errors.New("no parent to detach from")
	}

	id := g.GetID()

	child := -1
	for i, v := range parent.Children {
		if v == id {
			child = i
		}
	}

	if child == -1 {
		return errors.New("could not remove child from group: child does not exist")
	}

	parent.Children = append(parent.Children[:child], parent.Children[child+1:]...)

	update := struct {
		Id    int `json:"id"`
		Child int `json:"child"`
	}{
		parent.GetID(), g.GetID(),
	}

	s.websocketBroadcast(Update{Action: DELETE, Type: GROUP_CHILD, Data: update})
	return nil
}

func (s *Server) AddChildToGroup(id int, n Node) error {
	newParent, ok := s.groups[id]
	if !ok {
		return errors.New("group not found")
	}

	nid := n.GetID()
	for _, v := range newParent.Children {
		if v == nid {
			return errors.New("node already child of this group")
		}
	}

	newParent.Children = append(newParent.Children, nid)
	if n.GetParent() != nil {
		err := s.DetachChild(n)
		if err != nil {
			return err
		}
	}

	n.SetParent(newParent)

	update := struct {
		Id    int `json:"id"`
		Child int `json:"child"`
	}{
		id, nid,
	}

	s.websocketBroadcast(Update{Action: UPDATE, Type: GROUP_CHILD, Data: update})
	return nil
}

// CreateGroupHandler responds to a POST request to instantiate a new group and add it to the Server.
// Moves all of the specified children out of the parent's group and into the new group.
func (s *Server) GroupCreateHandler(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read request body"})
		return
	}

	var g struct {
		Group    int    `json:"group"`
		Children []int  `json:"children"`
		Label    string `json:"label"`
	}

	err = json.Unmarshal(body, &g)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read JSON"})
		return
	}

	s.Lock()
	defer s.Unlock()

	newGroup := &Group{
		Children: g.Children,
		Label:    g.Label,
		Id:       s.GetNextID(),
	}

	if newGroup.Children == nil {
		newGroup.Children = []int{}
	}

	for _, c := range newGroup.Children {
		_, okb := s.blocks[c]
		_, okg := s.groups[c]
		if !okb && !okg {
			w.WriteHeader(http.StatusBadRequest)
			writeJSON(w, Error{"could not create group: invalid children"})
			return
		}
	}

	s.groups[newGroup.Id] = newGroup
	s.websocketBroadcast(Update{Action: CREATE, Type: GROUP, Data: newGroup})

	err = s.AddChildToGroup(g.Group, newGroup)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	for _, c := range newGroup.Children {
		if cb, ok := s.blocks[c]; ok {
			err = s.AddChildToGroup(newGroup.Id, cb)
		}
		if cg, ok := s.groups[c]; ok {
			err = s.AddChildToGroup(newGroup.Id, cg)
		}
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			writeJSON(w, Error{err.Error()})
			return
		}
	}

	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) DeleteGroup(id int) error {
	group, ok := s.groups[id]
	if !ok {
		return errors.New("could not find group to delete")
	}

	for _, c := range group.Children {
		if _, ok := s.blocks[c]; ok {
			err := s.DeleteBlock(c)
			if err != nil {
				return err
			}
		} else if _, ok := s.groups[c]; ok {
			err := s.DeleteGroup(c)
			if err != nil {
				return err
			}
		}
	}

	update := struct {
		Id int `json:"id"`
	}{
		id,
	}
	s.DetachChild(group)
	delete(s.groups, id)
	s.websocketBroadcast(Update{Action: DELETE, Type: GROUP, Data: update})
	return nil
}

func (s *Server) GroupDeleteHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	ids, ok := vars["id"]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"no ID supplied"})
		return
	}

	id, err := strconv.Atoi(ids)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	s.Lock()
	defer s.Unlock()

	err = s.DeleteGroup(id)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusNoContent)
}
func (s *Server) GroupHandler(w http.ResponseWriter, r *http.Request) {
}
func (s *Server) GroupExportHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	ids, ok := vars["id"]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"no ID supplied"})
		return
	}

	id, err := strconv.Atoi(ids)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	s.Lock()
	defer s.Unlock()
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	writeJSON(w, s.ExportGroup(id))
}

func (s *Server) ExportGroup(id int) Pattern {
	p := Pattern{}
	p.Groups = append(p.Groups, *s.groups[id])
	for _, c := range s.groups[id].Children {
		b, ok := s.blocks[c]
		if !ok {
			g := s.ExportGroup(c)
			p.Blocks = append(p.Blocks, g.Blocks...)
			p.Groups = append(p.Groups, g.Groups...)
			p.Connections = append(p.Connections, g.Connections...)
			continue
		}
		p.Blocks = append(p.Blocks, *b)
	}
	return p
}

func (s *Server) GroupImportHandler(w http.ResponseWriter, r *http.Request) {

}
func (s *Server) GroupModifyLabelHandler(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read request body"})
		return
	}

	vars := mux.Vars(r)
	ids, ok := vars["id"]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"no ID supplied"})
		return
	}

	id, err := strconv.Atoi(ids)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	var l string
	err = json.Unmarshal(body, &l)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not unmarshal value"})
		return
	}

	s.Lock()
	defer s.Unlock()

	g, ok := s.groups[id]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"no block found"})
		return
	}

	g.Label = l

	update := struct {
		Label string `json:"label"`
		Id    int    `json:"id"`
	}{
		l, id,
	}

	s.websocketBroadcast(Update{Action: UPDATE, Type: GROUP, Data: update})

	w.WriteHeader(http.StatusNoContent)
}
func (s *Server) GroupModifyAllChildrenHandler(w http.ResponseWriter, r *http.Request) {
}
func (s *Server) GroupModifyChildHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	ids, ok := vars["id"]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"no ID supplied"})
		return
	}

	id, err := strconv.Atoi(ids)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	childs, ok := vars["node_id"]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"no ID supplied"})
		return
	}

	child, err := strconv.Atoi(childs)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	if id == child {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"cannot add group as member of itself"})
		return
	}

	s.Lock()
	defer s.Unlock()

	var n Node

	if _, ok := s.groups[id]; !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not find id"})
		return
	}

	if b, ok := s.blocks[child]; ok {
		n = b
	}
	if g, ok := s.groups[child]; ok {
		n = g
	}

	if n == nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not find id"})
		return
	}

	err = s.AddChildToGroup(id, n)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) GroupPositionHandler(w http.ResponseWriter, r *http.Request) {
}
<commit_msg>adding error checking to export<commit_after>package server

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"strconv"

	"github.com/gorilla/mux"
)

type Pattern struct {
	Blocks      []BlockLedger      `json:"blocks"`
	Connections []ConnectionLedger `json:"connections"`
	Groups      []Group            `json:"groups"`
}

type Node interface {
	GetID() int
	GetParent() *Group
	SetParent(*Group)
}

type Group struct {
	Id       int    `json:"id"`
	Label    string `json:"label"`
	Children []int  `json:"children"`
	Parent   *Group `json:"-"`
}

func (g *Group) GetID() int {
	return g.Id
}

func (g *Group) GetParent() *Group {
	return g.Parent
}

func (g *Group) SetParent(group *Group) {
	g.Parent = group
}

func (s *Server) ListGroups() []Group {
	groups := []Group{}
	for _, g := range s.groups {
		groups = append(groups, *g)
	}
	return groups
}

func (s *Server) GroupIndexHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.WriteHeader(http.StatusOK)
	if err := json.NewEncoder(w).Encode(s.ListGroups()); err != nil {
		panic(err)
	}
}

func (s *Server) DetachChild(g Node) error {
	parent := g.GetParent()
	if parent == nil {
		return errors.New("no parent to detach from")
	}

	id := g.GetID()

	child := -1
	for i, v := range parent.Children {
		if v == id {
			child = i
		}
	}

	if child == -1 {
		return errors.New("could not remove child from group: child does not exist")
	}

	parent.Children = append(parent.Children[:child], parent.Children[child+1:]...)

	update := struct {
		Id    int `json:"id"`
		Child int `json:"child"`
	}{
		parent.GetID(), g.GetID(),
	}

	s.websocketBroadcast(Update{Action: DELETE, Type: GROUP_CHILD, Data: update})
	return nil
}

func (s *Server) AddChildToGroup(id int, n Node) error {
	newParent, ok := s.groups[id]
	if !ok {
		return errors.New("group not found")
	}

	nid := n.GetID()
	for _, v := range newParent.Children {
		if v == nid {
			return errors.New("node already child of this group")
		}
	}

	newParent.Children = append(newParent.Children, nid)
	if n.GetParent() != nil {
		err := s.DetachChild(n)
		if err != nil {
			return err
		}
	}

	n.SetParent(newParent)

	update := struct {
		Id    int `json:"id"`
		Child int `json:"child"`
	}{
		id, nid,
	}

	s.websocketBroadcast(Update{Action: UPDATE, Type: GROUP_CHILD, Data: update})
	return nil
}

// CreateGroupHandler responds to a POST request to instantiate a new group and add it to the Server.
// Moves all of the specified children out of the parent's group and into the new group.
func (s *Server) GroupCreateHandler(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read request body"})
		return
	}

	var g struct {
		Group    int    `json:"group"`
		Children []int  `json:"children"`
		Label    string `json:"label"`
	}

	err = json.Unmarshal(body, &g)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read JSON"})
		return
	}

	s.Lock()
	defer s.Unlock()

	newGroup := &Group{
		Children: g.Children,
		Label:    g.Label,
		Id:       s.GetNextID(),
	}

	if newGroup.Children == nil {
		newGroup.Children = []int{}
	}

	for _, c := range newGroup.Children {
		_, okb := s.blocks[c]
		_, okg := s.groups[c]
		if !okb && !okg {
			w.WriteHeader(http.StatusBadRequest)
			writeJSON(w, Error{"could not create group: invalid children"})
			return
		}
	}

	s.groups[newGroup.Id] = newGroup
	s.websocketBroadcast(Update{Action: CREATE, Type: GROUP, Data: newGroup})

	err = s.AddChildToGroup(g.Group, newGroup)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	for _, c := range newGroup.Children {
		if cb, ok := s.blocks[c]; ok {
			err = s.AddChildToGroup(newGroup.Id, cb)
		}
		if cg, ok := s.groups[c]; ok {
			err = s.AddChildToGroup(newGroup.Id, cg)
		}
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			writeJSON(w, Error{err.Error()})
			return
		}
	}

	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) DeleteGroup(id int) error {
	group, ok := s.groups[id]
	if !ok {
		return errors.New("could not find group to delete")
	}

	for _, c := range group.Children {
		if _, ok := s.blocks[c]; ok {
			err := s.DeleteBlock(c)
			if err != nil {
				return err
			}
		} else if _, ok := s.groups[c]; ok {
			err := s.DeleteGroup(c)
			if err != nil {
				return err
			}
		}
	}

	update := struct {
		Id int `json:"id"`
	}{
		id,
	}
	s.DetachChild(group)
	delete(s.groups, id)
	s.websocketBroadcast(Update{Action: DELETE, Type: GROUP, Data: update})
	return nil
}

func (s *Server) GroupDeleteHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	ids, ok := vars["id"]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"no ID supplied"})
		return
	}

	id, err := strconv.Atoi(ids)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	s.Lock()
	defer s.Unlock()

	err = s.DeleteGroup(id)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusNoContent)
}
func (s *Server) GroupHandler(w http.ResponseWriter, r *http.Request) {
}
func (s *Server) GroupExportHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	ids, ok := vars["id"]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"no ID supplied"})
		return
	}

	id, err := strconv.Atoi(ids)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	s.Lock()
	defer s.Unlock()
	p, err := s.ExportGroup(id)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	writeJSON(w, p)
}

func (s *Server) ExportGroup(id int) (*Pattern, error) {
	p := &Pattern{}
	g, ok := s.groups[id]
	if !ok {
		return nil, errors.New("could not find group to export")
	}

	p.Groups = append(p.Groups, *g)
	for _, c := range g.Children {
		b, ok := s.blocks[c]
		if !ok {
			g, err := s.ExportGroup(c)
			if err != nil {
				return nil, err
			}

			p.Blocks = append(p.Blocks, g.Blocks...)
			p.Groups = append(p.Groups, g.Groups...)
			p.Connections = append(p.Connections, g.Connections...)
			continue
		}
		p.Blocks = append(p.Blocks, *b)
	}

	return p, nil
}

func (s *Server) GroupImportHandler(w http.ResponseWriter, r *http.Request) {

}
func (s *Server) GroupModifyLabelHandler(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not read request body"})
		return
	}

	vars := mux.Vars(r)
	ids, ok := vars["id"]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"no ID supplied"})
		return
	}

	id, err := strconv.Atoi(ids)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	var l string
	err = json.Unmarshal(body, &l)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not unmarshal value"})
		return
	}

	s.Lock()
	defer s.Unlock()

	g, ok := s.groups[id]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"no block found"})
		return
	}

	g.Label = l

	update := struct {
		Label string `json:"label"`
		Id    int    `json:"id"`
	}{
		l, id,
	}

	s.websocketBroadcast(Update{Action: UPDATE, Type: GROUP, Data: update})

	w.WriteHeader(http.StatusNoContent)
}
func (s *Server) GroupModifyAllChildrenHandler(w http.ResponseWriter, r *http.Request) {
}
func (s *Server) GroupModifyChildHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	ids, ok := vars["id"]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"no ID supplied"})
		return
	}

	id, err := strconv.Atoi(ids)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	childs, ok := vars["node_id"]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"no ID supplied"})
		return
	}

	child, err := strconv.Atoi(childs)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	if id == child {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"cannot add group as member of itself"})
		return
	}

	s.Lock()
	defer s.Unlock()

	var n Node

	if _, ok := s.groups[id]; !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not find id"})
		return
	}

	if b, ok := s.blocks[child]; ok {
		n = b
	}
	if g, ok := s.groups[child]; ok {
		n = g
	}

	if n == nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not find id"})
		return
	}

	err = s.AddChildToGroup(id, n)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) GroupPositionHandler(w http.ResponseWriter, r *http.Request) {
}
<|endoftext|>


<commit_before>package core

import (
	"fmt"
	"io"
	"pfi/sensorbee/sensorbee/data"
	"sync"
)

// SharedState is a state which nodes in a topology can access. It can be a
// machine learning model, a data structure for aggregation (like a histgram),
// a configuration information for specific Boxes, and so on.
//
// SharedState doesn't have methods to read it's internal data because internal
// data representation heavily depends on each SharedState implementation. The
// easiest way to use SharedState from a component is to obtain the actual
// data type via the type assertion. See examples to learn more about how to
// use it.
//
// If a SharedState also implements Writer interface, it can be updated via
// SharedStateSink. Write method in it writes a tuple to the state. How tuples
// are processed depends on each SharedState. For example, a machine learning
// model might use a tuple as a training data, and another state could compute
// the average of a specific field. Write may return fatal or temporary errors
// as Box.Process does. See the documentation of Box.Process for details.
//
// Write method might be called after Terminate method is called. When it
// occurs, Write should return an error. Also, Write and Terminate can be
// called concurrently.
type SharedState interface {
	// Terminate finalizes the state. The state can no longer be used after
	// this method is called. This method doesn't have to be idempotent.
	//
	// Write or other methods the actual instance has might be called after
	// Terminate method is called. When it occurs, they should return an error.
	// Also, Terminate and them can be called concurrently.
	Terminate(ctx *Context) error
}

// SavableSharedState is a SharedState which can be persisted through Save
// method. Providing forward/backward compatibility of the saved file format
// is the responsibility of the author of the state.
//
// Because the best way of implementing Load method depends on each SharedState,
// it doesn't always have to be provided with Save method.
type SavableSharedState interface {
	SharedState

	// Save writes data of the state to a given writer. Save receives parameters
	// which are used to customize the behavior of the method. Parameters are
	// defined by each component and there's no common definition.
	//
	// Save and other methods can be called concurrently.
	Save(ctx *Context, w io.Writer, params data.Map) error
}

// LoadableSharedState is a SharedState which can be persisted through Save
// and Load method.
type LoadableSharedState interface {
	SavableSharedState

	// Load overwrites the state with save data. Parameters don't have to be
	// same as Save's parameters. They can even be completely different.
	// There MUST NOT be a required parameter. Values of required parameters
	// should be saved with the state itself.
	//
	// Load and other methods including Save can be called concurrently.
	Load(ctx *Context, r io.Reader, params data.Map) error
}

// TODO: Add MixiableSharedState interface

// SharedStateRegistry manages SharedState with names assigned to each state.
type SharedStateRegistry interface {
	// Add adds a state to the registry. It fails if the registry already has
	// a state having the same name. Add also calls SharedState.Init. If it
	// fails Add returns an error and doesn't register the SharedState. The
	// caller doesn't have to call Terminate on failure.
	//
	// Don't add the same instance of SharedState more than once to registries.
	// Otherwise, Init and Terminate methods of the state will be called
	// multiple times.
	Add(name, typeName string, s SharedState) error

	// Get returns a SharedState having the name in the registry. It returns
	// NotExistError if the registry doesn't have the state.
	Get(name string) (SharedState, error)

	// Type returns a type of a SharedState. It returns NotExistError if the
	// registry doesn't have the state.
	Type(name string) (string, error)

	// Replace replaces the previous SharedState instance with a new instance.
	// The previous instance is returned on success if any. The previous state
	// will not be terminated by the registry and the caller must call
	// Terminate. The type name must be same as the previous state's type name.
	//
	// The given SharedState is terminated when it cannot be replaced.
	Replace(name, typeName string, s SharedState) (SharedState, error)

	// List returns a map containing all SharedState the registry has.
	// The map returned from this method can safely be modified.
	List() (map[string]SharedState, error)

	// Remove removes a SharedState the registry has. It automatically
	// terminates the state. If SharedState.Terminate failed, Remove returns an
	// error. However, even if it returns an error, the state is removed from
	// the registry.
	//
	// Remove also returns the removed SharedState if the registry has it. When
	// SharedState.Terminate fails, Remove returns both the removed SharedState
	// and an error. If the registry doesn't have a SharedState having the name,
	// it returns a nil SharedState and NotExistError.
	Remove(name string) (SharedState, error)
}

type defaultSharedStateInfo struct {
	state    SharedState
	typeName string
}

type defaultSharedStateRegistry struct {
	ctx    *Context
	m      sync.RWMutex
	states map[string]*defaultSharedStateInfo
}

// NewDefaultSharedStateRegistry create a default registry of SharedStates.
func NewDefaultSharedStateRegistry(ctx *Context) SharedStateRegistry {
	return &defaultSharedStateRegistry{
		ctx:    ctx,
		states: map[string]*defaultSharedStateInfo{},
	}
}

func (r *defaultSharedStateRegistry) Add(name, typeName string, s SharedState) error {
	err := func() error {
		r.m.Lock()
		defer r.m.Unlock()
		if _, ok := r.states[name]; ok {
			return fmt.Errorf("the registry already has a state '%v'", name)
		}
		r.states[name] = &defaultSharedStateInfo{
			state:    s,
			typeName: typeName,
		}
		return nil
	}()
	if err != nil {
		if err := r.closeSharedState(s); err != nil {
			r.ctx.ErrLog(err).WithField("state_name", name).
				Errorf("Cannot terminate a state which couldn't be added to the registry due to name duplication")
		}
		return err // This is the original error
	}
	return nil
}

func (r *defaultSharedStateRegistry) closeSharedState(s SharedState) (err error) {
	defer func() {
		if e := recover(); e != nil {
			if er, ok := e.(error); ok {
				err = er
			} else {
				err = fmt.Errorf("SharedState.Terminate panicked: %v", e)
			}
		}
	}()
	return s.Terminate(r.ctx)
}

func (r *defaultSharedStateRegistry) Get(name string) (SharedState, error) {
	r.m.RLock()
	defer r.m.RUnlock()
	if s, ok := r.states[name]; ok {
		return s.state, nil
	}
	return nil, NotExistError(fmt.Errorf("state '%v' was not found", name))
}

func (r *defaultSharedStateRegistry) Type(name string) (string, error) {
	r.m.RLock()
	defer r.m.RUnlock()
	if s, ok := r.states[name]; ok {
		return s.typeName, nil
	}
	return "", NotExistError(fmt.Errorf("state '%v' was not found", name))
}

func (r *defaultSharedStateRegistry) Replace(name, typeName string, s SharedState) (SharedState, error) {
	r.m.Lock()
	defer r.m.Unlock()
	prev, ok := r.states[name]
	if ok {
		if prev.typeName != typeName {
			if err := r.closeSharedState(s); err != nil {
				r.ctx.ErrLog(err).WithField("state_name", name).
					WithField("state_type", typeName).WithField("prev_state_type", prev.typeName).
					Errorf("Cannot terminate a state which couldn't be replaced due to a type mismatch")
			}
			return nil, fmt.Errorf("state '%v' has a different type from the previous state's type", name)
		}
	}
	r.states[name] = &defaultSharedStateInfo{
		state:    s,
		typeName: typeName,
	}
	if prev == nil {
		return nil, nil
	}
	return prev.state, nil
}

func (r *defaultSharedStateRegistry) List() (map[string]SharedState, error) {
	r.m.RLock()
	defer r.m.RUnlock()
	m := make(map[string]SharedState, len(r.states))
	for n, s := range r.states {
		m[n] = s.state
	}
	return m, nil
}

func (r *defaultSharedStateRegistry) Remove(name string) (SharedState, error) {
	s := func() SharedState {
		r.m.Lock()
		defer r.m.Unlock()
		if s, ok := r.states[name]; ok {
			delete(r.states, name)
			return s.state
		}
		return nil
	}()
	if s == nil {
		return nil, NotExistError(fmt.Errorf("state '%v' was not found", name))
	}

	if err := s.Terminate(r.ctx); err != nil {
		return s, err
	}
	return s, nil
}

// writer points a shared state. sharedStateSink will point to the same shared state
// even after the state is removed from the context.
type sharedStateSink struct {
	writer Writer
}

// NewSharedStateSink creates a sink that writes to SharedState.
func NewSharedStateSink(ctx *Context, name string) (Sink, error) {
	// Get SharedState by name
	state, err := ctx.SharedStates.Get(name)
	if err != nil {
		return nil, err
	}

	// It fails if the shared state cannot be written
	writer, ok := state.(Writer)
	if !ok {
		return nil, fmt.Errorf("'%v' state cannot be written", name)
	}

	s := &sharedStateSink{
		writer: writer,
	}
	return s, nil
}

func (s *sharedStateSink) Write(ctx *Context, t *Tuple) error {
	return s.writer.Write(ctx, t)
}

func (s *sharedStateSink) Close(ctx *Context) error {
	// SharedState must not be terminated when this sink is closed because
	// the state is still being used by other components.
	return nil
}
<commit_msg>Modify sharedStateSink to refer extract sinks by name in Write().<commit_after>package core

import (
	"fmt"
	"io"
	"pfi/sensorbee/sensorbee/data"
	"sync"
)

// SharedState is a state which nodes in a topology can access. It can be a
// machine learning model, a data structure for aggregation (like a histgram),
// a configuration information for specific Boxes, and so on.
//
// SharedState doesn't have methods to read it's internal data because internal
// data representation heavily depends on each SharedState implementation. The
// easiest way to use SharedState from a component is to obtain the actual
// data type via the type assertion. See examples to learn more about how to
// use it.
//
// If a SharedState also implements Writer interface, it can be updated via
// SharedStateSink. Write method in it writes a tuple to the state. How tuples
// are processed depends on each SharedState. For example, a machine learning
// model might use a tuple as a training data, and another state could compute
// the average of a specific field. Write may return fatal or temporary errors
// as Box.Process does. See the documentation of Box.Process for details.
//
// Write method might be called after Terminate method is called. When it
// occurs, Write should return an error. Also, Write and Terminate can be
// called concurrently.
type SharedState interface {
	// Terminate finalizes the state. The state can no longer be used after
	// this method is called. This method doesn't have to be idempotent.
	//
	// Write or other methods the actual instance has might be called after
	// Terminate method is called. When it occurs, they should return an error.
	// Also, Terminate and them can be called concurrently.
	Terminate(ctx *Context) error
}

// SavableSharedState is a SharedState which can be persisted through Save
// method. Providing forward/backward compatibility of the saved file format
// is the responsibility of the author of the state.
//
// Because the best way of implementing Load method depends on each SharedState,
// it doesn't always have to be provided with Save method.
type SavableSharedState interface {
	SharedState

	// Save writes data of the state to a given writer. Save receives parameters
	// which are used to customize the behavior of the method. Parameters are
	// defined by each component and there's no common definition.
	//
	// Save and other methods can be called concurrently.
	Save(ctx *Context, w io.Writer, params data.Map) error
}

// LoadableSharedState is a SharedState which can be persisted through Save
// and Load method.
type LoadableSharedState interface {
	SavableSharedState

	// Load overwrites the state with save data. Parameters don't have to be
	// same as Save's parameters. They can even be completely different.
	// There MUST NOT be a required parameter. Values of required parameters
	// should be saved with the state itself.
	//
	// Load and other methods including Save can be called concurrently.
	Load(ctx *Context, r io.Reader, params data.Map) error
}

// TODO: Add MixiableSharedState interface

// SharedStateRegistry manages SharedState with names assigned to each state.
type SharedStateRegistry interface {
	// Add adds a state to the registry. It fails if the registry already has
	// a state having the same name. Add also calls SharedState.Init. If it
	// fails Add returns an error and doesn't register the SharedState. The
	// caller doesn't have to call Terminate on failure.
	//
	// Don't add the same instance of SharedState more than once to registries.
	// Otherwise, Init and Terminate methods of the state will be called
	// multiple times.
	Add(name, typeName string, s SharedState) error

	// Get returns a SharedState having the name in the registry. It returns
	// NotExistError if the registry doesn't have the state.
	Get(name string) (SharedState, error)

	// Type returns a type of a SharedState. It returns NotExistError if the
	// registry doesn't have the state.
	Type(name string) (string, error)

	// Replace replaces the previous SharedState instance with a new instance.
	// The previous instance is returned on success if any. The previous state
	// will not be terminated by the registry and the caller must call
	// Terminate. The type name must be same as the previous state's type name.
	//
	// The given SharedState is terminated when it cannot be replaced.
	Replace(name, typeName string, s SharedState) (SharedState, error)

	// List returns a map containing all SharedState the registry has.
	// The map returned from this method can safely be modified.
	List() (map[string]SharedState, error)

	// Remove removes a SharedState the registry has. It automatically
	// terminates the state. If SharedState.Terminate failed, Remove returns an
	// error. However, even if it returns an error, the state is removed from
	// the registry.
	//
	// Remove also returns the removed SharedState if the registry has it. When
	// SharedState.Terminate fails, Remove returns both the removed SharedState
	// and an error. If the registry doesn't have a SharedState having the name,
	// it returns a nil SharedState and NotExistError.
	Remove(name string) (SharedState, error)
}

type defaultSharedStateInfo struct {
	state    SharedState
	typeName string
}

type defaultSharedStateRegistry struct {
	ctx    *Context
	m      sync.RWMutex
	states map[string]*defaultSharedStateInfo
}

// NewDefaultSharedStateRegistry create a default registry of SharedStates.
func NewDefaultSharedStateRegistry(ctx *Context) SharedStateRegistry {
	return &defaultSharedStateRegistry{
		ctx:    ctx,
		states: map[string]*defaultSharedStateInfo{},
	}
}

func (r *defaultSharedStateRegistry) Add(name, typeName string, s SharedState) error {
	err := func() error {
		r.m.Lock()
		defer r.m.Unlock()
		if _, ok := r.states[name]; ok {
			return fmt.Errorf("the registry already has a state '%v'", name)
		}
		r.states[name] = &defaultSharedStateInfo{
			state:    s,
			typeName: typeName,
		}
		return nil
	}()
	if err != nil {
		if err := r.closeSharedState(s); err != nil {
			r.ctx.ErrLog(err).WithField("state_name", name).
				Errorf("Cannot terminate a state which couldn't be added to the registry due to name duplication")
		}
		return err // This is the original error
	}
	return nil
}

func (r *defaultSharedStateRegistry) closeSharedState(s SharedState) (err error) {
	defer func() {
		if e := recover(); e != nil {
			if er, ok := e.(error); ok {
				err = er
			} else {
				err = fmt.Errorf("SharedState.Terminate panicked: %v", e)
			}
		}
	}()
	return s.Terminate(r.ctx)
}

func (r *defaultSharedStateRegistry) Get(name string) (SharedState, error) {
	r.m.RLock()
	defer r.m.RUnlock()
	if s, ok := r.states[name]; ok {
		return s.state, nil
	}
	return nil, NotExistError(fmt.Errorf("state '%v' was not found", name))
}

func (r *defaultSharedStateRegistry) Type(name string) (string, error) {
	r.m.RLock()
	defer r.m.RUnlock()
	if s, ok := r.states[name]; ok {
		return s.typeName, nil
	}
	return "", NotExistError(fmt.Errorf("state '%v' was not found", name))
}

func (r *defaultSharedStateRegistry) Replace(name, typeName string, s SharedState) (SharedState, error) {
	r.m.Lock()
	defer r.m.Unlock()
	prev, ok := r.states[name]
	if ok {
		if prev.typeName != typeName {
			if err := r.closeSharedState(s); err != nil {
				r.ctx.ErrLog(err).WithField("state_name", name).
					WithField("state_type", typeName).WithField("prev_state_type", prev.typeName).
					Errorf("Cannot terminate a state which couldn't be replaced due to a type mismatch")
			}
			return nil, fmt.Errorf("state '%v' has a different type from the previous state's type", name)
		}
	}
	r.states[name] = &defaultSharedStateInfo{
		state:    s,
		typeName: typeName,
	}
	if prev == nil {
		return nil, nil
	}
	return prev.state, nil
}

func (r *defaultSharedStateRegistry) List() (map[string]SharedState, error) {
	r.m.RLock()
	defer r.m.RUnlock()
	m := make(map[string]SharedState, len(r.states))
	for n, s := range r.states {
		m[n] = s.state
	}
	return m, nil
}

func (r *defaultSharedStateRegistry) Remove(name string) (SharedState, error) {
	s := func() SharedState {
		r.m.Lock()
		defer r.m.Unlock()
		if s, ok := r.states[name]; ok {
			delete(r.states, name)
			return s.state
		}
		return nil
	}()
	if s == nil {
		return nil, NotExistError(fmt.Errorf("state '%v' was not found", name))
	}

	if err := s.Terminate(r.ctx); err != nil {
		return s, err
	}
	return s, nil
}

// writer points a shared state. sharedStateSink will point to the same shared state
// even after the state is removed from the context.
type sharedStateSink struct {
	name string
}

// NewSharedStateSink creates a sink that writes to SharedState.
func NewSharedStateSink(ctx *Context, name string) (Sink, error) {
	// Get SharedState by name
	state, err := ctx.SharedStates.Get(name)
	if err != nil {
		return nil, err
	}

	// It fails if the shared state cannot be written
	_, ok := state.(Writer)
	if !ok {
		return nil, fmt.Errorf("'%v' state cannot be written", name)
	}

	s := &sharedStateSink{
		name: name,
	}
	return s, nil
}

func (s *sharedStateSink) Write(ctx *Context, t *Tuple) error {
	state, err := ctx.SharedStates.Get(s.name)
	if err != nil {
		return err
	}

	// It fails if the shared state cannot be written
	writer, ok := state.(Writer)
	if !ok {
		return fmt.Errorf("'%v' state cannot be written", s.name)
	}

	return writer.Write(ctx, t)
}

func (s *sharedStateSink) Close(ctx *Context) error {
	// SharedState must not be terminated when this sink is closed because
	// the state is still being used by other components.
	return nil
}
<|endoftext|>


<commit_before>package set

import (
	"testing"
)

// BenchmarkAdd checks the performance of the set.Add() method
func BenchmarkAdd(b *testing.B) {
	// Create a new set
	set := New()

	// Run set.Add() b.N times
	for i := 0; i < b.N; i++ {
		set.Add(i)
	}
}

// benchmarkCartesianProduct checks the performance of the set.CartesianProduct() method
func benchmarkCartesianProduct(n int, s *Set, t *Set) {
	// Run set.CartesianProduct() n times
	for i := 0; i < n; i++ {
		s.CartesianProduct(t)
	}
}

// BenchmarkCartesianProductSmall checks the performance of the set.CartesianProduct() method
// over a small data set
func BenchmarkCartesianProductSmall(b *testing.B) {
	benchmarkCartesianProduct(b.N, New(1, 2), New(2, 1))
}

// BenchmarkCartesianProductLarge checks the performance of the set.CartesianProduct() method
// over a large data set
func BenchmarkCartesianProductLarge(b *testing.B) {
	benchmarkCartesianProduct(b.N, New(1, 2, 3, 4, 5, 6, 7, 8, 9), New(9, 8, 7, 6, 5, 4, 3, 2, 1))
}

// benchmarkClone checks the performance of the set.Clone() method
func benchmarkClone(n int, s *Set) {
	// Run set.Clone() n times
	for i := 0; i < n; i++ {
		s.Clone()
	}
}

// BenchmarkCloneSmall checks the performance of the set.Clone() method
// over a small data set
func BenchmarkCloneSmall(b *testing.B) {
	benchmarkClone(b.N, New(1, 2))
}

// BenchmarkCloneLarge checks the performance of the set.Clone() method
// over a large data set
func BenchmarkCloneLarge(b *testing.B) {
	benchmarkClone(b.N, New(1, 2, 3, 4, 5, 6, 7, 8, 9))
}

// benchmarkDifference checks the performance of the set.Difference() method
func benchmarkDifference(n int, s *Set, t *Set) {
	// Run set.Difference() n times
	for i := 0; i < n; i++ {
		s.Difference(t)
	}
}

// BenchmarkDifferenceSmall checks the performance of the set.Difference() method
// over a small data set
func BenchmarkDifferenceSmall(b *testing.B) {
	benchmarkDifference(b.N, New(1, 2), New(2, 1))
}

// BenchmarkDifferenceLarge checks the performance of the set.Difference() method
// over a large data set
func BenchmarkDifferenceLarge(b *testing.B) {
	benchmarkDifference(b.N, New(1, 2, 3, 4, 5, 6, 7, 8, 9), New(9, 8, 7, 6, 5, 4, 3, 2, 1))
}

// benchmarkEnumerate checks the performance of the set.Enumerate() method
func benchmarkEnumerate(n int, s *Set) {
	// Run set.Enumerate() n times
	for i := 0; i < n; i++ {
		s.Enumerate()
	}
}

// BenchmarkEnumerateSmall checks the performance of the set.Enumerate() method
// over a small data set
func BenchmarkEnumerateSmall(b *testing.B) {
	benchmarkEnumerate(b.N, New(1, 2))
}

// BenchmarkEnumerateLarge checks the performance of the set.Enumerate() method
// over a large data set
func BenchmarkEnumerateLarge(b *testing.B) {
	benchmarkEnumerate(b.N, New(1, 2, 3, 4, 5, 6, 7, 8, 9))
}
<commit_msg>Add BenchmarkEqual(), small and large variants<commit_after>package set

import (
	"testing"
)

// BenchmarkAdd checks the performance of the set.Add() method
func BenchmarkAdd(b *testing.B) {
	// Create a new set
	set := New()

	// Run set.Add() b.N times
	for i := 0; i < b.N; i++ {
		set.Add(i)
	}
}

// benchmarkCartesianProduct checks the performance of the set.CartesianProduct() method
func benchmarkCartesianProduct(n int, s *Set, t *Set) {
	// Run set.CartesianProduct() n times
	for i := 0; i < n; i++ {
		s.CartesianProduct(t)
	}
}

// BenchmarkCartesianProductSmall checks the performance of the set.CartesianProduct() method
// over a small data set
func BenchmarkCartesianProductSmall(b *testing.B) {
	benchmarkCartesianProduct(b.N, New(1, 2), New(2, 1))
}

// BenchmarkCartesianProductLarge checks the performance of the set.CartesianProduct() method
// over a large data set
func BenchmarkCartesianProductLarge(b *testing.B) {
	benchmarkCartesianProduct(b.N, New(1, 2, 3, 4, 5, 6, 7, 8, 9), New(9, 8, 7, 6, 5, 4, 3, 2, 1))
}

// benchmarkClone checks the performance of the set.Clone() method
func benchmarkClone(n int, s *Set) {
	// Run set.Clone() n times
	for i := 0; i < n; i++ {
		s.Clone()
	}
}

// BenchmarkCloneSmall checks the performance of the set.Clone() method
// over a small data set
func BenchmarkCloneSmall(b *testing.B) {
	benchmarkClone(b.N, New(1, 2))
}

// BenchmarkCloneLarge checks the performance of the set.Clone() method
// over a large data set
func BenchmarkCloneLarge(b *testing.B) {
	benchmarkClone(b.N, New(1, 2, 3, 4, 5, 6, 7, 8, 9))
}

// benchmarkDifference checks the performance of the set.Difference() method
func benchmarkDifference(n int, s *Set, t *Set) {
	// Run set.Difference() n times
	for i := 0; i < n; i++ {
		s.Difference(t)
	}
}

// BenchmarkDifferenceSmall checks the performance of the set.Difference() method
// over a small data set
func BenchmarkDifferenceSmall(b *testing.B) {
	benchmarkDifference(b.N, New(1, 2), New(2, 1))
}

// BenchmarkDifferenceLarge checks the performance of the set.Difference() method
// over a large data set
func BenchmarkDifferenceLarge(b *testing.B) {
	benchmarkDifference(b.N, New(1, 2, 3, 4, 5, 6, 7, 8, 9), New(9, 8, 7, 6, 5, 4, 3, 2, 1))
}

// benchmarkEnumerate checks the performance of the set.Enumerate() method
func benchmarkEnumerate(n int, s *Set) {
	// Run set.Enumerate() n times
	for i := 0; i < n; i++ {
		s.Enumerate()
	}
}

// BenchmarkEnumerateSmall checks the performance of the set.Enumerate() method
// over a small data set
func BenchmarkEnumerateSmall(b *testing.B) {
	benchmarkEnumerate(b.N, New(1, 2))
}

// BenchmarkEnumerateLarge checks the performance of the set.Enumerate() method
// over a large data set
func BenchmarkEnumerateLarge(b *testing.B) {
	benchmarkEnumerate(b.N, New(1, 2, 3, 4, 5, 6, 7, 8, 9))
}

// benchmarkEqual checks the performance of the set.Equal() method
func benchmarkEqual(n int, s *Set, t *Set) {
	// Run set.Equal() n times
	for i := 0; i < n; i++ {
		s.Equal(t)
	}
}

// BenchmarkEqualSmall checks the performance of the set.Equal() method
// over a small data set
func BenchmarkEqualSmall(b *testing.B) {
	benchmarkEqual(b.N, New(1, 2), New(2, 1))
}

// BenchmarkEqualLarge checks the performance of the set.Equal() method
// over a large data set
func BenchmarkEqualLarge(b *testing.B) {
	benchmarkEqual(b.N, New(1, 2, 3, 4, 5, 6, 7, 8, 9), New(9, 8, 7, 6, 5, 4, 3, 2, 1))
}
<|endoftext|>


<commit_before>package set

import (
	"testing"
)

// BenchmarkAdd checks the performance of the set.Add() method
func BenchmarkAdd(b *testing.B) {
	// Create a new set
	set := New()

	// Run set.Add() b.N times
	for i := 0; i < b.N; i++ {
		set.Add(i)
	}
}
<commit_msg>Add BenchmarkCartesianProduct(), small and large variants<commit_after>package set

import (
	"testing"
)

// BenchmarkAdd checks the performance of the set.Add() method
func BenchmarkAdd(b *testing.B) {
	// Create a new set
	set := New()

	// Run set.Add() b.N times
	for i := 0; i < b.N; i++ {
		set.Add(i)
	}
}

// benchmarkCartesianProduct checks the performance of the set.CartesianProduct() method
func benchmarkCartesianProduct(n int, s *Set, t *Set) {
	// Run set.CartesianProduct() n times
	for i := 0; i < n; i++ {
		s.CartesianProduct(t)
	}
}

// BenchmarkCartesianProductSmall checks the performance of the set.CartesianProduct() method
// over a small data set
func BenchmarkCartesianProductSmall(b *testing.B) {
	benchmarkCartesianProduct(b.N, New(1, 2), New(2, 1))
}

// BenchmarkCartesianProductLarge checks the performance of the set.CartesianProduct() method
// over a large data set
func BenchmarkCartesianProductLarge(b *testing.B) {
	benchmarkCartesianProduct(b.N, New(1, 2, 3, 4, 5, 6, 7, 8, 9), New(9, 8, 7, 6, 5, 4, 3, 2, 1))
}
<|endoftext|>


<commit_before>// Socket package provides a concise, powerful and high-performance TCP socket.
//
// Copyright 2017 HenryLee. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
package socket

import (
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/henrylee2cn/goutil"
	"github.com/henrylee2cn/goutil/errors"
	"github.com/henrylee2cn/teleport/codec"
)

type (
	// Socket is a generic stream-oriented network connection.
	//
	// Multiple goroutines may invoke methods on a Socket simultaneously.
	Socket interface {
		// LocalAddr returns the local network address.
		LocalAddr() net.Addr

		// RemoteAddr returns the remote network address.
		RemoteAddr() net.Addr

		// SetDeadline sets the read and write deadlines associated
		// with the connection. It is equivalent to calling both
		// SetReadDeadline and SetWriteDeadline.
		//
		// A deadline is an absolute time after which I/O operations
		// fail with a timeout (see type Error) instead of
		// blocking. The deadline applies to all future and pending
		// I/O, not just the immediately following call to Read or
		// Write. After a deadline has been exceeded, the connection
		// can be refreshed by setting a deadline in the future.
		//
		// An idle timeout can be implemented by repeatedly extending
		// the deadline after successful Read or Write calls.
		//
		// A zero value for t means I/O operations will not time out.
		SetDeadline(t time.Time) error

		// SetReadDeadline sets the deadline for future Read calls
		// and any currently-blocked Read call.
		// A zero value for t means Read will not time out.
		SetReadDeadline(t time.Time) error

		// SetWriteDeadline sets the deadline for future Write calls
		// and any currently-blocked Write call.
		// Even if write times out, it may return n > 0, indicating that
		// some of the data was successfully written.
		// A zero value for t means Write will not time out.
		SetWriteDeadline(t time.Time) error

		// WritePacket writes header and body to the connection.
		// Note: must be safe for concurrent use by multiple goroutines.
		WritePacket(packet *Packet) error

		// ReadPacket reads header and body from the connection.
		// Note: must be safe for concurrent use by multiple goroutines.
		ReadPacket(packet *Packet) error

		// Read reads data from the connection.
		// Read can be made to time out and return an Error with Timeout() == true
		// after a fixed time limit; see SetDeadline and SetReadDeadline.
		Read(b []byte) (n int, err error)

		// Write writes data to the connection.
		// Write can be made to time out and return an Error with Timeout() == true
		// after a fixed time limit; see SetDeadline and SetWriteDeadline.
		Write(b []byte) (n int, err error)

		// Reset reset net.Conn
		// Reset(net.Conn)

		// Close closes the connection socket.
		// Any blocked Read or Write operations will be unblocked and return errors.
		Close() error

		// Public returns temporary public data of Socket.
		Public() goutil.Map
		// PublicLen returns the length of public data of Socket.
		PublicLen() int
		// Id returns the socket id.
		Id() string
		// SetId sets the socket id.
		SetId(string)
		// Reset reset net.Conn and ProtoFunc.
		Reset(netConn net.Conn, protoFunc ...ProtoFunc)
	}
	socket struct {
		net.Conn
		protocol  Proto
		id        string
		idMutex   sync.RWMutex
		ctxPublic goutil.Map
		mu        sync.RWMutex
		curState  int32
		fromPool  bool
	}
)

const (
	normal      int32 = 0
	activeClose int32 = 1
)

var _ net.Conn = Socket(nil)

// ErrProactivelyCloseSocket proactively close the socket error.
var ErrProactivelyCloseSocket = errors.New("socket is closed proactively")

// GetSocket gets a Socket from pool, and reset it.
func GetSocket(c net.Conn, protoFunc ...ProtoFunc) Socket {
	s := socketPool.Get().(*socket)
	s.Reset(c, protoFunc...)
	return s
}

var socketPool = sync.Pool{
	New: func() interface{} {
		s := newSocket(nil, nil)
		s.fromPool = true
		return s
	},
}

// NewSocket wraps a net.Conn as a Socket.
func NewSocket(c net.Conn, protoFunc ...ProtoFunc) Socket {
	return newSocket(c, protoFunc)
}

func newSocket(c net.Conn, protoFuncs []ProtoFunc) *socket {
	var s = &socket{
		protocol: getProto(protoFuncs, c),
		Conn:     c,
	}
	s.optimize()
	return s
}

// WritePacket writes header and body to the connection.
// WritePacket can be made to time out and return an Error with Timeout() == true
// after a fixed time limit; see SetDeadline and SetWriteDeadline.
// Note:
//  For the byte stream type of body, write directly, do not do any processing;
//  Must be safe for concurrent use by multiple goroutines.
func (s *socket) WritePacket(packet *Packet) error {
	s.mu.RLock()
	protocol := s.protocol
	s.mu.RUnlock()
	if packet.BodyCodec() == codec.NilCodecId {
		packet.SetBodyCodec(defaultBodyCodec.Id())
	}
	err := protocol.Pack(packet)
	if err != nil && s.isActiveClosed() {
		err = ErrProactivelyCloseSocket
	}
	return err
}

// ReadPacket reads header and body from the connection.
// Note:
//  For the byte stream type of body, read directly, do not do any processing;
//  Must be safe for concurrent use by multiple goroutines.
func (s *socket) ReadPacket(packet *Packet) error {
	s.mu.RLock()
	protocol := s.protocol
	s.mu.RUnlock()
	return protocol.Unpack(packet)
}

// Public returns temporary public data of Socket.
func (s *socket) Public() goutil.Map {
	if s.ctxPublic == nil {
		s.ctxPublic = goutil.RwMap()
	}
	return s.ctxPublic
}

// PublicLen returns the length of public data of Socket.
func (s *socket) PublicLen() int {
	if s.ctxPublic == nil {
		return 0
	}
	return s.ctxPublic.Len()
}

// Id returns the socket id.
func (s *socket) Id() string {
	s.idMutex.RLock()
	id := s.id
	if len(id) == 0 {
		id = s.RemoteAddr().String()
	}
	s.idMutex.RUnlock()
	return id
}

// SetId sets the socket id.
func (s *socket) SetId(id string) {
	s.idMutex.Lock()
	s.id = id
	s.idMutex.Unlock()
}

// Reset reset net.Conn and ProtoFunc.
func (s *socket) Reset(netConn net.Conn, protoFunc ...ProtoFunc) {
	atomic.StoreInt32(&s.curState, activeClose)
	if s.Conn != nil {
		s.Conn.Close()
	}
	s.mu.Lock()
	s.Conn = netConn
	s.SetId("")
	s.protocol = getProto(protoFunc, netConn)
	atomic.StoreInt32(&s.curState, normal)
	s.optimize()
	s.mu.Unlock()
}

// Close closes the connection socket.
// Any blocked Read or Write operations will be unblocked and return errors.
// If it is from 'GetSocket()' function(a pool), return itself to pool.
func (s *socket) Close() error {
	if s.isActiveClosed() {
		return nil
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.isActiveClosed() {
		return nil
	}
	atomic.StoreInt32(&s.curState, activeClose)

	var err error
	if s.Conn != nil {
		err = s.Conn.Close()
	}
	if s.fromPool {
		s.Conn = nil
		s.ctxPublic = nil
		s.protocol = nil
		socketPool.Put(s)
	}
	return err
}

func (s *socket) isActiveClosed() bool {
	return atomic.LoadInt32(&s.curState) == activeClose
}

func (s *socket) optimize() {
	if c, ok := s.Conn.(ifaceSetKeepAlive); ok {
		if changeKeepAlive {
			c.SetKeepAlive(keepAlive)
		}
		if keepAlive && keepAlivePeriod >= 0 {
			c.SetKeepAlivePeriod(keepAlivePeriod)
		}
	}
	if c, ok := s.Conn.(ifaceSetBuffer); ok {
		if readBuffer >= 0 {
			c.SetReadBuffer(readBuffer)
		}
		if writeBuffer >= 0 {
			c.SetWriteBuffer(writeBuffer)
		}
	}
}

type (
	ifaceSetKeepAlive interface {
		// SetKeepAlive sets whether the operating system should send
		// keepalive messages on the connection.
		SetKeepAlive(keepalive bool) error
		// SetKeepAlivePeriod sets period between keep alives.
		SetKeepAlivePeriod(d time.Duration) error
	}
	ifaceSetBuffer interface {
		// SetReadBuffer sets the size of the operating system's
		// receive buffer associated with the connection.
		SetReadBuffer(bytes int) error
		// SetWriteBuffer sets the size of the operating system's
		// transmit buffer associated with the connection.
		SetWriteBuffer(bytes int) error
	}
)

// Connection related system configuration
var (
	writeBuffer     int           = -1
	readBuffer      int           = -1
	changeKeepAlive bool          = false
	keepAlive       bool          = false
	keepAlivePeriod time.Duration = -1
)

// SetKeepAlive sets whether the operating system should send
// keepalive messages on the connection.
// Note: If have not called the function, the system defaults are used.
func SetKeepAlive(keepalive bool) {
	changeKeepAlive = true
	keepAlive = keepalive
}

// SetKeepAlivePeriod sets period between keep alives.
// Note: if d=-1, don't change the system default value.
func SetKeepAlivePeriod(d time.Duration) {
	keepAlivePeriod = d
}

// SetReadBuffer sets the size of the operating system's
// receive buffer associated with the connection.
// Note: if bytes=-1, don't change the system default value.
func SetReadBuffer(bytes int) {
	readBuffer = bytes
	resetFastProtoReadBufioSize()
}

// SetWriteBuffer sets the size of the operating system's
// transmit buffer associated with the connection.
// Note: if bytes=-1, don't change the system default value.
func SetWriteBuffer(bytes int) {
	writeBuffer = bytes
}

var fastProtoReadBufioSize int

func init() {
	resetFastProtoReadBufioSize()
}

func resetFastProtoReadBufioSize() {
	if readBuffer < 0 {
		fastProtoReadBufioSize = 1024 * 4
	} else if readBuffer == 0 {
		fastProtoReadBufioSize = 1024 * 35
	} else {
		fastProtoReadBufioSize = readBuffer / 2
	}
}

func getProto(protoFuncs []ProtoFunc, rw io.ReadWriter) Proto {
	if len(protoFuncs) > 0 {
		return protoFuncs[0](rw)
	} else {
		return defaultProtoFunc(rw)
	}
}
<commit_msg>Fix socket keepAlive default value<commit_after>// Socket package provides a concise, powerful and high-performance TCP socket.
//
// Copyright 2017 HenryLee. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
package socket

import (
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/henrylee2cn/goutil"
	"github.com/henrylee2cn/goutil/errors"
	"github.com/henrylee2cn/teleport/codec"
)

type (
	// Socket is a generic stream-oriented network connection.
	//
	// Multiple goroutines may invoke methods on a Socket simultaneously.
	Socket interface {
		// LocalAddr returns the local network address.
		LocalAddr() net.Addr

		// RemoteAddr returns the remote network address.
		RemoteAddr() net.Addr

		// SetDeadline sets the read and write deadlines associated
		// with the connection. It is equivalent to calling both
		// SetReadDeadline and SetWriteDeadline.
		//
		// A deadline is an absolute time after which I/O operations
		// fail with a timeout (see type Error) instead of
		// blocking. The deadline applies to all future and pending
		// I/O, not just the immediately following call to Read or
		// Write. After a deadline has been exceeded, the connection
		// can be refreshed by setting a deadline in the future.
		//
		// An idle timeout can be implemented by repeatedly extending
		// the deadline after successful Read or Write calls.
		//
		// A zero value for t means I/O operations will not time out.
		SetDeadline(t time.Time) error

		// SetReadDeadline sets the deadline for future Read calls
		// and any currently-blocked Read call.
		// A zero value for t means Read will not time out.
		SetReadDeadline(t time.Time) error

		// SetWriteDeadline sets the deadline for future Write calls
		// and any currently-blocked Write call.
		// Even if write times out, it may return n > 0, indicating that
		// some of the data was successfully written.
		// A zero value for t means Write will not time out.
		SetWriteDeadline(t time.Time) error

		// WritePacket writes header and body to the connection.
		// Note: must be safe for concurrent use by multiple goroutines.
		WritePacket(packet *Packet) error

		// ReadPacket reads header and body from the connection.
		// Note: must be safe for concurrent use by multiple goroutines.
		ReadPacket(packet *Packet) error

		// Read reads data from the connection.
		// Read can be made to time out and return an Error with Timeout() == true
		// after a fixed time limit; see SetDeadline and SetReadDeadline.
		Read(b []byte) (n int, err error)

		// Write writes data to the connection.
		// Write can be made to time out and return an Error with Timeout() == true
		// after a fixed time limit; see SetDeadline and SetWriteDeadline.
		Write(b []byte) (n int, err error)

		// Reset reset net.Conn
		// Reset(net.Conn)

		// Close closes the connection socket.
		// Any blocked Read or Write operations will be unblocked and return errors.
		Close() error

		// Public returns temporary public data of Socket.
		Public() goutil.Map
		// PublicLen returns the length of public data of Socket.
		PublicLen() int
		// Id returns the socket id.
		Id() string
		// SetId sets the socket id.
		SetId(string)
		// Reset reset net.Conn and ProtoFunc.
		Reset(netConn net.Conn, protoFunc ...ProtoFunc)
	}
	socket struct {
		net.Conn
		protocol  Proto
		id        string
		idMutex   sync.RWMutex
		ctxPublic goutil.Map
		mu        sync.RWMutex
		curState  int32
		fromPool  bool
	}
)

const (
	normal      int32 = 0
	activeClose int32 = 1
)

var _ net.Conn = Socket(nil)

// ErrProactivelyCloseSocket proactively close the socket error.
var ErrProactivelyCloseSocket = errors.New("socket is closed proactively")

// GetSocket gets a Socket from pool, and reset it.
func GetSocket(c net.Conn, protoFunc ...ProtoFunc) Socket {
	s := socketPool.Get().(*socket)
	s.Reset(c, protoFunc...)
	return s
}

var socketPool = sync.Pool{
	New: func() interface{} {
		s := newSocket(nil, nil)
		s.fromPool = true
		return s
	},
}

// NewSocket wraps a net.Conn as a Socket.
func NewSocket(c net.Conn, protoFunc ...ProtoFunc) Socket {
	return newSocket(c, protoFunc)
}

func newSocket(c net.Conn, protoFuncs []ProtoFunc) *socket {
	var s = &socket{
		protocol: getProto(protoFuncs, c),
		Conn:     c,
	}
	s.optimize()
	return s
}

// WritePacket writes header and body to the connection.
// WritePacket can be made to time out and return an Error with Timeout() == true
// after a fixed time limit; see SetDeadline and SetWriteDeadline.
// Note:
//  For the byte stream type of body, write directly, do not do any processing;
//  Must be safe for concurrent use by multiple goroutines.
func (s *socket) WritePacket(packet *Packet) error {
	s.mu.RLock()
	protocol := s.protocol
	s.mu.RUnlock()
	if packet.BodyCodec() == codec.NilCodecId {
		packet.SetBodyCodec(defaultBodyCodec.Id())
	}
	err := protocol.Pack(packet)
	if err != nil && s.isActiveClosed() {
		err = ErrProactivelyCloseSocket
	}
	return err
}

// ReadPacket reads header and body from the connection.
// Note:
//  For the byte stream type of body, read directly, do not do any processing;
//  Must be safe for concurrent use by multiple goroutines.
func (s *socket) ReadPacket(packet *Packet) error {
	s.mu.RLock()
	protocol := s.protocol
	s.mu.RUnlock()
	return protocol.Unpack(packet)
}

// Public returns temporary public data of Socket.
func (s *socket) Public() goutil.Map {
	if s.ctxPublic == nil {
		s.ctxPublic = goutil.RwMap()
	}
	return s.ctxPublic
}

// PublicLen returns the length of public data of Socket.
func (s *socket) PublicLen() int {
	if s.ctxPublic == nil {
		return 0
	}
	return s.ctxPublic.Len()
}

// Id returns the socket id.
func (s *socket) Id() string {
	s.idMutex.RLock()
	id := s.id
	if len(id) == 0 {
		id = s.RemoteAddr().String()
	}
	s.idMutex.RUnlock()
	return id
}

// SetId sets the socket id.
func (s *socket) SetId(id string) {
	s.idMutex.Lock()
	s.id = id
	s.idMutex.Unlock()
}

// Reset reset net.Conn and ProtoFunc.
func (s *socket) Reset(netConn net.Conn, protoFunc ...ProtoFunc) {
	atomic.StoreInt32(&s.curState, activeClose)
	if s.Conn != nil {
		s.Conn.Close()
	}
	s.mu.Lock()
	s.Conn = netConn
	s.SetId("")
	s.protocol = getProto(protoFunc, netConn)
	atomic.StoreInt32(&s.curState, normal)
	s.optimize()
	s.mu.Unlock()
}

// Close closes the connection socket.
// Any blocked Read or Write operations will be unblocked and return errors.
// If it is from 'GetSocket()' function(a pool), return itself to pool.
func (s *socket) Close() error {
	if s.isActiveClosed() {
		return nil
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.isActiveClosed() {
		return nil
	}
	atomic.StoreInt32(&s.curState, activeClose)

	var err error
	if s.Conn != nil {
		err = s.Conn.Close()
	}
	if s.fromPool {
		s.Conn = nil
		s.ctxPublic = nil
		s.protocol = nil
		socketPool.Put(s)
	}
	return err
}

func (s *socket) isActiveClosed() bool {
	return atomic.LoadInt32(&s.curState) == activeClose
}

func (s *socket) optimize() {
	if c, ok := s.Conn.(ifaceSetKeepAlive); ok {
		if changeKeepAlive {
			c.SetKeepAlive(keepAlive)
		}
		if keepAlive && keepAlivePeriod >= 0 {
			c.SetKeepAlivePeriod(keepAlivePeriod)
		}
	}
	if c, ok := s.Conn.(ifaceSetBuffer); ok {
		if readBuffer >= 0 {
			c.SetReadBuffer(readBuffer)
		}
		if writeBuffer >= 0 {
			c.SetWriteBuffer(writeBuffer)
		}
	}
}

type (
	ifaceSetKeepAlive interface {
		// SetKeepAlive sets whether the operating system should send
		// keepalive messages on the connection.
		SetKeepAlive(keepalive bool) error
		// SetKeepAlivePeriod sets period between keep alives.
		SetKeepAlivePeriod(d time.Duration) error
	}
	ifaceSetBuffer interface {
		// SetReadBuffer sets the size of the operating system's
		// receive buffer associated with the connection.
		SetReadBuffer(bytes int) error
		// SetWriteBuffer sets the size of the operating system's
		// transmit buffer associated with the connection.
		SetWriteBuffer(bytes int) error
	}
)

// Connection related system configuration
var (
	writeBuffer     int           = -1
	readBuffer      int           = -1
	changeKeepAlive bool          = false
	keepAlive       bool          = true
	keepAlivePeriod time.Duration = -1
)

// SetKeepAlive sets whether the operating system should send
// keepalive messages on the connection.
// Note: If have not called the function, the system defaults are used.
func SetKeepAlive(keepalive bool) {
	changeKeepAlive = true
	keepAlive = keepalive
}

// SetKeepAlivePeriod sets period between keep alives.
// Note: if d=-1, don't change the system default value.
func SetKeepAlivePeriod(d time.Duration) {
	keepAlivePeriod = d
}

// SetReadBuffer sets the size of the operating system's
// receive buffer associated with the connection.
// Note: if bytes=-1, don't change the system default value.
func SetReadBuffer(bytes int) {
	readBuffer = bytes
	resetFastProtoReadBufioSize()
}

// SetWriteBuffer sets the size of the operating system's
// transmit buffer associated with the connection.
// Note: if bytes=-1, don't change the system default value.
func SetWriteBuffer(bytes int) {
	writeBuffer = bytes
}

var fastProtoReadBufioSize int

func init() {
	resetFastProtoReadBufioSize()
}

func resetFastProtoReadBufioSize() {
	if readBuffer < 0 {
		fastProtoReadBufioSize = 1024 * 4
	} else if readBuffer == 0 {
		fastProtoReadBufioSize = 1024 * 35
	} else {
		fastProtoReadBufioSize = readBuffer / 2
	}
}

func getProto(protoFuncs []ProtoFunc, rw io.ReadWriter) Proto {
	if len(protoFuncs) > 0 {
		return protoFuncs[0](rw)
	} else {
		return defaultProtoFunc(rw)
	}
}
<|endoftext|>