<commit_before>package rest

import (
	"regexp"
	"strings"
)

const (
	intParam      = "i"
	alphaParam    = "a"
	alphaNumParam = "an"
)

var regParam = regexp.MustCompile(`{(\w+)(:\w+)?}`)
var regMap = map[string]string{
	intParam:      `([0-9]+)`,
	alphaParam:    `([A-Za-z]+)`,
	alphaNumParam: `([0-9A-Za-z]+)`,
}

// Routeable is an interface which allows you to create your own router
// * Get(string) *Route returns the route by route name
// * Match(string) *Route takes a URI and returns a *Route it matches.  If it does not
//   it returns nil
// * SetParamNames sets the route parameter names to a route.  Ex. /api/foo/{id}, id is a param
// * SetParamValues sets the value to match the route parameter.  Ex. /api/foo/i, id is the param, and 1 is the value
// * SetActions set the actions for that given route.  The map key is the method.  Ex. map["GET"]Action
type Routeable interface {
	Get(string) *Route
	Match(string) *Route
	SetParamNames(*Route, []string)
	SetParamValues(*Route, []interface{})
	SetActions(*Route, map[string]Action)
}

type router struct {
	prefix     string
	Routes     map[string]*Route
	routingMap map[*Route]*regexp.Regexp
}

type routeBuilder struct {
	path      string
	actions   map[string]Action
	routeName string
}

var routers map[string]router

func init() {
	routers = make(map[string]router)
	routers["default"] = router{"", make(map[string]*Route), make(map[*Route]*regexp.Regexp)}
}

// NewRouter creates and returns a ready to use router based on a name
func NewRouter(name string) router {
	routers[name] = router{"", make(map[string]*Route), make(map[*Route]*regexp.Regexp)}
	return routers[name]
}

// DefaultRouter returns the router with name "default"
func DefaultRouter() router {
	return routers["default"]
}

// Router returns a router specified by a name
func Router(name string) router {
	return routers[name]
}

func (ro router) Prefix(prefix string) router {
	ro.prefix = prefix
	return ro
}

func NewRoute() *routeBuilder {
	return &routeBuilder{actions: make(map[string]Action)}
}

func (rb *routeBuilder) Named(name string) *routeBuilder {
	rb.routeName = name
	return rb
}

func (rb *routeBuilder) For(path string) *routeBuilder {
	rb.path = path
	return rb
}

func (rb *routeBuilder) With(method string, action Action) *routeBuilder {
	rb.actions[method] = action
	return rb
}

func (rb *routeBuilder) And(method string, action Action) *routeBuilder {
	return rb.With(method, action)
}

func (ro router) RouteMap(rbs ...*routeBuilder) router {
	for _, routeBuilder := range rbs {
		route := &Route{
			Path:    ro.prefix + routeBuilder.path,
			Actions: routeBuilder.actions,
		}

		ro.initRoute(route)

		if routeBuilder.routeName == "" {
			ro.Routes[routeBuilder.path] = route
		}
		ro.Routes[routeBuilder.routeName] = route
	}

	return ro
}

func (ro router) initRoute(route *Route) {
	toSub := regParam.FindAllStringSubmatch(route.Path, -1)

	regString := route.Path

	if len(toSub) > 0 {
		params := make([]string, len(toSub))

		for i, v := range toSub {
			whole, param, pType, regex := v[0], v[1], v[2], `([^/]+)`

			params[i] = param

			if len(pType) > 1 {
				if r, ok := regMap[pType[1:]]; ok {
					regex = r
				}
			}
			regString = strings.Replace(regString, whole, regex, -1)
		}
		ro.SetParamNames(route, params)
	}

	ro.routingMap[route] = regexp.MustCompile(regString + "/?")
}

func (ro router) Get(name string) *Route {
	return ro.Routes[name]
}

func (ro router) Match(test string) *Route {
	for route, regex := range ro.routingMap {
		matches := regex.FindStringSubmatch(test)
		if matches != nil && matches[0] == test {
			values := make([]interface{}, len(matches[1:]))

			for i, m := range matches[1:] {
				values[i] = m
			}
			ro.SetParamValues(route, values)

			return route
		}
	}

	return nil
}

func (ro router) SetParamNames(r *Route, pn []string) {
	r.ParamNames = pn
}

func (ro router) SetParamValues(r *Route, pv []interface{}) {
	r.ParamValues = pv
}

func (ro router) SetActions(r *Route, a map[string]Action) {
	r.Actions = a
}
<commit_msg>make sure to create router is named doesnt exit<commit_after>package rest

import (
	"regexp"
	"strings"
)

const (
	intParam      = "i"
	alphaParam    = "a"
	alphaNumParam = "an"
)

var regParam = regexp.MustCompile(`{(\w+)(:\w+)?}`)
var regMap = map[string]string{
	intParam:      `([0-9]+)`,
	alphaParam:    `([A-Za-z]+)`,
	alphaNumParam: `([0-9A-Za-z]+)`,
}

// Routeable is an interface which allows you to create your own router
// * Get(string) *Route returns the route by route name
// * Match(string) *Route takes a URI and returns a *Route it matches.  If it does not
//   it returns nil
// * SetParamNames sets the route parameter names to a route.  Ex. /api/foo/{id}, id is a param
// * SetParamValues sets the value to match the route parameter.  Ex. /api/foo/i, id is the param, and 1 is the value
// * SetActions set the actions for that given route.  The map key is the method.  Ex. map["GET"]Action
type Routeable interface {
	Get(string) *Route
	Match(string) *Route
	SetParamNames(*Route, []string)
	SetParamValues(*Route, []interface{})
	SetActions(*Route, map[string]Action)
}

type router struct {
	prefix     string
	Routes     map[string]*Route
	routingMap map[*Route]*regexp.Regexp
}

type routeBuilder struct {
	path      string
	actions   map[string]Action
	routeName string
}

var routers map[string]router

func init() {
	routers = make(map[string]router)
	routers["default"] = router{"", make(map[string]*Route), make(map[*Route]*regexp.Regexp)}
}

// NewRouter creates and returns a ready to use router based on a name
func NewRouter(name string) router {
	routers[name] = router{"", make(map[string]*Route), make(map[*Route]*regexp.Regexp)}
	return routers[name]
}

// DefaultRouter returns the router with name "default"
func DefaultRouter() router {
	return routers["default"]
}

// Router returns a router specified by a name
func Router(name string) router {
	// create new router if it doesn't exist
	if _, ok := routers[name]; !ok {
		routers[name] = router{"", make(map[string]*Route), make(map[*Route]*regexp.Regexp)}
	}

	return routers[name]
}

func (ro router) Prefix(prefix string) router {
	ro.prefix = prefix
	return ro
}

func NewRoute() *routeBuilder {
	return &routeBuilder{actions: make(map[string]Action)}
}

func (rb *routeBuilder) Named(name string) *routeBuilder {
	rb.routeName = name
	return rb
}

func (rb *routeBuilder) For(path string) *routeBuilder {
	rb.path = path
	return rb
}

func (rb *routeBuilder) With(method string, action Action) *routeBuilder {
	rb.actions[method] = action
	return rb
}

func (rb *routeBuilder) And(method string, action Action) *routeBuilder {
	return rb.With(method, action)
}

func (ro router) RouteMap(rbs ...*routeBuilder) router {
	for _, routeBuilder := range rbs {
		route := &Route{
			Path:    ro.prefix + routeBuilder.path,
			Actions: routeBuilder.actions,
		}

		ro.initRoute(route)

		if routeBuilder.routeName == "" {
			ro.Routes[routeBuilder.path] = route
		}
		ro.Routes[routeBuilder.routeName] = route
	}

	return ro
}

func (ro router) initRoute(route *Route) {
	toSub := regParam.FindAllStringSubmatch(route.Path, -1)

	regString := route.Path

	if len(toSub) > 0 {
		params := make([]string, len(toSub))

		for i, v := range toSub {
			whole, param, pType, regex := v[0], v[1], v[2], `([^/]+)`

			params[i] = param

			if len(pType) > 1 {
				if r, ok := regMap[pType[1:]]; ok {
					regex = r
				}
			}
			regString = strings.Replace(regString, whole, regex, -1)
		}
		ro.SetParamNames(route, params)
	}

	ro.routingMap[route] = regexp.MustCompile(regString + "/?")
}

func (ro router) Get(name string) *Route {
	return ro.Routes[name]
}

func (ro router) Match(test string) *Route {
	for route, regex := range ro.routingMap {
		matches := regex.FindStringSubmatch(test)
		if matches != nil && matches[0] == test {
			values := make([]interface{}, len(matches[1:]))

			for i, m := range matches[1:] {
				values[i] = m
			}
			ro.SetParamValues(route, values)

			return route
		}
	}

	return nil
}

func (ro router) SetParamNames(r *Route, pn []string) {
	r.ParamNames = pn
}

func (ro router) SetParamValues(r *Route, pv []interface{}) {
	r.ParamValues = pv
}

func (ro router) SetActions(r *Route, a map[string]Action) {
	r.Actions = a
}
<|endoftext|>


<commit_before>package echo

type (
	// Router is the registry of all registered routes for an `Echo` instance for
	// request matching and URL path parameter parsing.
	Router struct {
		tree   *node
		routes []Route
		echo   *Echo
	}
	node struct {
		kind          kind
		label         byte
		prefix        string
		parent        *node
		children      children
		ppath         string
		pnames        []string
		methodHandler *methodHandler
	}
	kind          uint8
	children      []*node
	methodHandler struct {
		connect Handler
		delete  Handler
		get     Handler
		head    Handler
		options Handler
		patch   Handler
		post    Handler
		put     Handler
		trace   Handler
	}
)

const (
	skind kind = iota
	pkind
	akind
)

// NewRouter returns a new Router instance.
func NewRouter(e *Echo) *Router {
	return &Router{
		tree: &node{
			methodHandler: new(methodHandler),
		},
		routes: []Route{},
		echo:   e,
	}
}

// Handle implements `echo.Middleware` which makes router a middleware.
func (r *Router) Handle(next Handler) Handler {
	return HandlerFunc(func(c Context) error {
		method := c.Request().Method()
		path := c.Request().URL().Path()
		r.Find(method, path, c)
		return next.Handle(c)
	})
}

// Priority is super secret.
func (r *Router) Priority() int {
	return 0
}

// Add registers a new route for method and path with matching handler.
func (r *Router) Add(method, path string, h Handler, e *Echo) {
	ppath := path        // Pristine path
	pnames := []string{} // Param names

	for i, l := 0, len(path); i < l; i++ {
		if path[i] == ':' {
			j := i + 1

			r.insert(method, path[:i], nil, skind, "", nil, e)
			for ; i < l && path[i] != '/'; i++ {
			}

			pnames = append(pnames, path[j:i])
			path = path[:j] + path[i:]
			i, l = j, len(path)

			if i == l {
				r.insert(method, path[:i], h, pkind, ppath, pnames, e)
				return
			}
			r.insert(method, path[:i], nil, pkind, ppath, pnames, e)
		} else if path[i] == '*' {
			r.insert(method, path[:i], nil, skind, "", nil, e)
			pnames = append(pnames, "_*")
			r.insert(method, path[:i+1], h, akind, ppath, pnames, e)
			return
		}
	}

	r.insert(method, path, h, skind, ppath, pnames, e)
}

func (r *Router) insert(method, path string, h Handler, t kind, ppath string, pnames []string, e *Echo) {
	// Adjust max param
	l := len(pnames)
	if *e.maxParam < l {
		*e.maxParam = l
	}

	cn := r.tree // Current node as root
	if cn == nil {
		panic("echo ⇛ invalid method")
	}
	search := path

	for {
		sl := len(search)
		pl := len(cn.prefix)
		l := 0

		// LCP
		max := pl
		if sl < max {
			max = sl
		}
		for ; l < max && search[l] == cn.prefix[l]; l++ {
		}

		if l == 0 {
			// At root node
			cn.label = search[0]
			cn.prefix = search
			if h != nil {
				cn.kind = t
				cn.addHandler(method, h)
				cn.ppath = ppath
				cn.pnames = pnames
			}
		} else if l < pl {
			// Split node
			n := newNode(cn.kind, cn.prefix[l:], cn, cn.children, cn.methodHandler, cn.ppath, cn.pnames)

			// Reset parent node
			cn.kind = skind
			cn.label = cn.prefix[0]
			cn.prefix = cn.prefix[:l]
			cn.children = nil
			cn.methodHandler = new(methodHandler)
			cn.ppath = ""
			cn.pnames = nil

			cn.addChild(n)

			if l == sl {
				// At parent node
				cn.kind = t
				cn.addHandler(method, h)
				cn.ppath = ppath
				cn.pnames = pnames
			} else {
				// Create child node
				n = newNode(t, search[l:], cn, nil, new(methodHandler), ppath, pnames)
				n.addHandler(method, h)
				cn.addChild(n)
			}
		} else if l < sl {
			search = search[l:]
			c := cn.findChildWithLabel(search[0])
			if c != nil {
				// Go deeper
				cn = c
				continue
			}
			// Create child node
			n := newNode(t, search, cn, nil, new(methodHandler), ppath, pnames)
			n.addHandler(method, h)
			cn.addChild(n)
		} else {
			// Node already exists
			if h != nil {
				cn.addHandler(method, h)
				cn.ppath = ppath
				cn.pnames = pnames
			}
		}
		return
	}
}

func newNode(t kind, pre string, p *node, c children, mh *methodHandler, ppath string, pnames []string) *node {
	return &node{
		kind:          t,
		label:         pre[0],
		prefix:        pre,
		parent:        p,
		children:      c,
		ppath:         ppath,
		pnames:        pnames,
		methodHandler: mh,
	}
}

func (n *node) addChild(c *node) {
	n.children = append(n.children, c)
}

func (n *node) findChild(l byte, t kind) *node {
	for _, c := range n.children {
		if c.label == l && c.kind == t {
			return c
		}
	}
	return nil
}

func (n *node) findChildWithLabel(l byte) *node {
	for _, c := range n.children {
		if c.label == l {
			return c
		}
	}
	return nil
}

func (n *node) findChildByKind(t kind) *node {
	for _, c := range n.children {
		if c.kind == t {
			return c
		}
	}
	return nil
}

func (n *node) addHandler(method string, h Handler) {
	switch method {
	case GET:
		n.methodHandler.get = h
	case POST:
		n.methodHandler.post = h
	case PUT:
		n.methodHandler.put = h
	case DELETE:
		n.methodHandler.delete = h
	case PATCH:
		n.methodHandler.patch = h
	case OPTIONS:
		n.methodHandler.options = h
	case HEAD:
		n.methodHandler.head = h
	case CONNECT:
		n.methodHandler.connect = h
	case TRACE:
		n.methodHandler.trace = h
	}
}

func (n *node) findHandler(method string) Handler {
	switch method {
	case GET:
		return n.methodHandler.get
	case POST:
		return n.methodHandler.post
	case PUT:
		return n.methodHandler.put
	case DELETE:
		return n.methodHandler.delete
	case PATCH:
		return n.methodHandler.patch
	case OPTIONS:
		return n.methodHandler.options
	case HEAD:
		return n.methodHandler.head
	case CONNECT:
		return n.methodHandler.connect
	case TRACE:
		return n.methodHandler.trace
	default:
		return nil
	}
}

func (n *node) check405() HandlerFunc {
	for _, m := range methods {
		if h := n.findHandler(m); h != nil {
			return methodNotAllowedHandler
		}
	}
	return notFoundHandler
}

// Find lookup a handler registed for method and path. It also parses URL for path
// parameters and load them into context.
//
// For performance:
//
// - Get context from `Echo#GetContext()`
// - Reset it `Context#Reset()`
// - Return it `Echo#PutContext()`.
func (r *Router) Find(method, path string, context Context) {
	ctx := context.Object()
	cn := r.tree // Current node as root

	var (
		search = path
		c      *node  // Child node
		n      int    // Param counter
		nk     kind   // Next kind
		nn     *node  // Next node
		ns     string // Next search
	)

	// Search order static > param > any
	for {
		if search == "" {
			goto End
		}

		pl := 0 // Prefix length
		l := 0  // LCP length

		if cn.label != ':' {
			sl := len(search)
			pl = len(cn.prefix)

			// LCP
			max := pl
			if sl < max {
				max = sl
			}
			for ; l < max && search[l] == cn.prefix[l]; l++ {
			}
		}

		if l == pl {
			// Continue search
			search = search[l:]
		} else {
			cn = nn
			search = ns
			if nk == pkind {
				goto Param
			} else if nk == akind {
				goto Any
			}
			// Not found
			return
		}

		if search == "" {
			goto End
		}

		// Static node
		if c = cn.findChild(search[0], skind); c != nil {
			// Save next
			if cn.label == '/' {
				nk = pkind
				nn = cn
				ns = search
			}
			cn = c
			continue
		}

		// Param node
	Param:
		if c = cn.findChildByKind(pkind); c != nil {
			// Issue #378
			if len(ctx.pvalues) == n {
				continue
			}

			// Save next
			if cn.label == '/' {
				nk = akind
				nn = cn
				ns = search
			}

			cn = c
			i, l := 0, len(search)
			for ; i < l && search[i] != '/'; i++ {
			}
			ctx.pvalues[n] = search[:i]
			n++
			search = search[i:]
			continue
		}

		// Any node
	Any:
		if cn = cn.findChildByKind(akind); cn == nil {
			if nn != nil {
				cn = nn
				nn = nil // Next
				search = ns
				if nk == pkind {
					goto Param
				} else if nk == akind {
					goto Any
				}
			}
			// Not found
			return
		}
		ctx.pvalues[len(cn.pnames)-1] = search
		goto End
	}

End:
	ctx.path = cn.ppath
	ctx.pnames = cn.pnames
	ctx.handler = cn.findHandler(method)

	// NOTE: Slow zone...
	if ctx.handler == nil {
		ctx.handler = cn.check405()

		// Dig further for any, might have an empty value for *, e.g.
		// serving a directory. Issue #207.
		if cn = cn.findChildByKind(akind); cn == nil {
			return
		}
		ctx.pvalues[len(cn.pnames)-1] = ""
		if ctx.handler = cn.findHandler(method); ctx.handler == nil {
			ctx.handler = cn.check405()
		}
	}
	return
}
<commit_msg>Refactored some code<commit_after>package echo

type (
	// Router is the registry of all registered routes for an `Echo` instance for
	// request matching and URL path parameter parsing.
	Router struct {
		tree   *node
		routes []Route
		echo   *Echo
	}
	node struct {
		kind          kind
		label         byte
		prefix        string
		parent        *node
		children      children
		ppath         string
		pnames        []string
		methodHandler *methodHandler
	}
	kind          uint8
	children      []*node
	methodHandler struct {
		connect Handler
		delete  Handler
		get     Handler
		head    Handler
		options Handler
		patch   Handler
		post    Handler
		put     Handler
		trace   Handler
	}
)

const (
	skind kind = iota
	pkind
	akind
)

// NewRouter returns a new Router instance.
func NewRouter(e *Echo) *Router {
	return &Router{
		tree: &node{
			methodHandler: new(methodHandler),
		},
		routes: []Route{},
		echo:   e,
	}
}

// Handle implements `echo.Middleware` which makes router a middleware.
func (r *Router) Handle(next Handler) Handler {
	return HandlerFunc(func(c Context) error {
		method := c.Request().Method()
		path := c.Request().URL().Path()
		r.Find(method, path, c)
		return next.Handle(c)
	})
}

// Priority is super secret.
func (r *Router) Priority() int {
	return 0
}

// Add registers a new route for method and path with matching handler.
func (r *Router) Add(method, path string, h Handler, e *Echo) {
	ppath := path        // Pristine path
	pnames := []string{} // Param names

	for i, l := 0, len(path); i < l; i++ {
		if path[i] == ':' {
			j := i + 1

			r.insert(method, path[:i], nil, skind, "", nil, e)
			for ; i < l && path[i] != '/'; i++ {
			}

			pnames = append(pnames, path[j:i])
			path = path[:j] + path[i:]
			i, l = j, len(path)

			if i == l {
				r.insert(method, path[:i], h, pkind, ppath, pnames, e)
				return
			}
			r.insert(method, path[:i], nil, pkind, ppath, pnames, e)
		} else if path[i] == '*' {
			r.insert(method, path[:i], nil, skind, "", nil, e)
			pnames = append(pnames, "_*")
			r.insert(method, path[:i+1], h, akind, ppath, pnames, e)
			return
		}
	}

	r.insert(method, path, h, skind, ppath, pnames, e)
}

func (r *Router) insert(method, path string, h Handler, t kind, ppath string, pnames []string, e *Echo) {
	// Adjust max param
	l := len(pnames)
	if *e.maxParam < l {
		*e.maxParam = l
	}

	cn := r.tree // Current node as root
	if cn == nil {
		panic("echo ⇛ invalid method")
	}
	search := path

	for {
		sl := len(search)
		pl := len(cn.prefix)
		l := 0

		// LCP
		max := pl
		if sl < max {
			max = sl
		}
		for ; l < max && search[l] == cn.prefix[l]; l++ {
		}

		if l == 0 {
			// At root node
			cn.label = search[0]
			cn.prefix = search
			if h != nil {
				cn.kind = t
				cn.addHandler(method, h)
				cn.ppath = ppath
				cn.pnames = pnames
			}
		} else if l < pl {
			// Split node
			n := newNode(cn.kind, cn.prefix[l:], cn, cn.children, cn.methodHandler, cn.ppath, cn.pnames)

			// Reset parent node
			cn.kind = skind
			cn.label = cn.prefix[0]
			cn.prefix = cn.prefix[:l]
			cn.children = nil
			cn.methodHandler = new(methodHandler)
			cn.ppath = ""
			cn.pnames = nil

			cn.addChild(n)

			if l == sl {
				// At parent node
				cn.kind = t
				cn.addHandler(method, h)
				cn.ppath = ppath
				cn.pnames = pnames
			} else {
				// Create child node
				n = newNode(t, search[l:], cn, nil, new(methodHandler), ppath, pnames)
				n.addHandler(method, h)
				cn.addChild(n)
			}
		} else if l < sl {
			search = search[l:]
			c := cn.findChildWithLabel(search[0])
			if c != nil {
				// Go deeper
				cn = c
				continue
			}
			// Create child node
			n := newNode(t, search, cn, nil, new(methodHandler), ppath, pnames)
			n.addHandler(method, h)
			cn.addChild(n)
		} else {
			// Node already exists
			if h != nil {
				cn.addHandler(method, h)
				cn.ppath = ppath
				cn.pnames = pnames
			}
		}
		return
	}
}

func newNode(t kind, pre string, p *node, c children, mh *methodHandler, ppath string, pnames []string) *node {
	return &node{
		kind:          t,
		label:         pre[0],
		prefix:        pre,
		parent:        p,
		children:      c,
		ppath:         ppath,
		pnames:        pnames,
		methodHandler: mh,
	}
}

func (n *node) addChild(c *node) {
	n.children = append(n.children, c)
}

func (n *node) findChild(l byte, t kind) *node {
	for _, c := range n.children {
		if c.label == l && c.kind == t {
			return c
		}
	}
	return nil
}

func (n *node) findChildWithLabel(l byte) *node {
	for _, c := range n.children {
		if c.label == l {
			return c
		}
	}
	return nil
}

func (n *node) findChildByKind(t kind) *node {
	for _, c := range n.children {
		if c.kind == t {
			return c
		}
	}
	return nil
}

func (n *node) addHandler(method string, h Handler) {
	switch method {
	case GET:
		n.methodHandler.get = h
	case POST:
		n.methodHandler.post = h
	case PUT:
		n.methodHandler.put = h
	case DELETE:
		n.methodHandler.delete = h
	case PATCH:
		n.methodHandler.patch = h
	case OPTIONS:
		n.methodHandler.options = h
	case HEAD:
		n.methodHandler.head = h
	case CONNECT:
		n.methodHandler.connect = h
	case TRACE:
		n.methodHandler.trace = h
	}
}

func (n *node) findHandler(method string) Handler {
	switch method {
	case GET:
		return n.methodHandler.get
	case POST:
		return n.methodHandler.post
	case PUT:
		return n.methodHandler.put
	case DELETE:
		return n.methodHandler.delete
	case PATCH:
		return n.methodHandler.patch
	case OPTIONS:
		return n.methodHandler.options
	case HEAD:
		return n.methodHandler.head
	case CONNECT:
		return n.methodHandler.connect
	case TRACE:
		return n.methodHandler.trace
	default:
		return nil
	}
}

func (n *node) checkMethodNotAllowed() HandlerFunc {
	for _, m := range methods {
		if h := n.findHandler(m); h != nil {
			return methodNotAllowedHandler
		}
	}
	return notFoundHandler
}

// Find lookup a handler registed for method and path. It also parses URL for path
// parameters and load them into context.
//
// For performance:
//
// - Get context from `Echo#GetContext()`
// - Reset it `Context#Reset()`
// - Return it `Echo#PutContext()`.
func (r *Router) Find(method, path string, context Context) {
	ctx := context.Object()
	cn := r.tree // Current node as root

	var (
		search = path
		c      *node  // Child node
		n      int    // Param counter
		nk     kind   // Next kind
		nn     *node  // Next node
		ns     string // Next search
	)

	// Search order static > param > any
	for {
		if search == "" {
			goto End
		}

		pl := 0 // Prefix length
		l := 0  // LCP length

		if cn.label != ':' {
			sl := len(search)
			pl = len(cn.prefix)

			// LCP
			max := pl
			if sl < max {
				max = sl
			}
			for ; l < max && search[l] == cn.prefix[l]; l++ {
			}
		}

		if l == pl {
			// Continue search
			search = search[l:]
		} else {
			cn = nn
			search = ns
			if nk == pkind {
				goto Param
			} else if nk == akind {
				goto Any
			}
			// Not found
			return
		}

		if search == "" {
			goto End
		}

		// Static node
		if c = cn.findChild(search[0], skind); c != nil {
			// Save next
			if cn.label == '/' {
				nk = pkind
				nn = cn
				ns = search
			}
			cn = c
			continue
		}

		// Param node
	Param:
		if c = cn.findChildByKind(pkind); c != nil {
			// Issue #378
			if len(ctx.pvalues) == n {
				continue
			}

			// Save next
			if cn.label == '/' {
				nk = akind
				nn = cn
				ns = search
			}

			cn = c
			i, l := 0, len(search)
			for ; i < l && search[i] != '/'; i++ {
			}
			ctx.pvalues[n] = search[:i]
			n++
			search = search[i:]
			continue
		}

		// Any node
	Any:
		if cn = cn.findChildByKind(akind); cn == nil {
			if nn != nil {
				cn = nn
				nn = nil // Next
				search = ns
				if nk == pkind {
					goto Param
				} else if nk == akind {
					goto Any
				}
			}
			// Not found
			return
		}
		ctx.pvalues[len(cn.pnames)-1] = search
		goto End
	}

End:
	ctx.path = cn.ppath
	ctx.pnames = cn.pnames
	ctx.handler = cn.findHandler(method)

	// NOTE: Slow zone...
	if ctx.handler == nil {
		ctx.handler = cn.checkMethodNotAllowed()

		// Dig further for any, might have an empty value for *, e.g.
		// serving a directory. Issue #207.
		if cn = cn.findChildByKind(akind); cn == nil {
			return
		}
		ctx.pvalues[len(cn.pnames)-1] = ""
		if ctx.handler = cn.findHandler(method); ctx.handler == nil {
			ctx.handler = cn.checkMethodNotAllowed()
		}
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/gorilla/mux"
	"log"
	"net/http"
	"time"
)

func NewRouter() *mux.Router {
	router := mux.NewRouter().StrictSlash(true)
	for _, route := range routes {
		handler := route.HandlerFunc
		handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

			start := time.Now()
			handler.ServeHTTP(w, r)
			log.Printf(
				"%s\t%s\t%s\t%s",
				r.Method,
				r.RequestURI,
				route,
				time.Since(start),
			)
		})

		router.
			Methods(route.Method).
			Path(route.Pattern).
			Name(route.Name).
			Handler(handler)

	}
	return router
}
<commit_msg>revert to old router<commit_after>package main

import (
	"github.com/gorilla/mux"
	"log"
	"net/http"
	"time"
)

func NewRouter() *mux.Router {
	router := mux.NewRouter().StrictSlash(true)
	for _, route := range routes {
		var handler http.Handler
		handler = route.HandlerFunc
		handler = Logger(handler, route.Name)

		router.
			Methods(route.Method).
			Path(route.Pattern).
			Name(route.Name).
			Handler(handler)

	}
	return router
}

func Logger(inner http.Handler, name string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()

		inner.ServeHTTP(w, r)

		log.Printf(
			"%s\t%s\t%s\t%s",
			r.Method,
			r.RequestURI,
			name,
			time.Since(start),
		)
	})
}
<|endoftext|>


<commit_before>package water

import (
	"fmt"
	"net/http"
	"strings"
	"sync"
)

var (
	_HTTP_METHODS = map[string]int{
		"GET":     0,
		"POST":    1,
		"DELETE":  2,
		"PUT":     3,
		"PATCH":   4,
		"HEAD":    5,
		"OPTIONS": 6,
		"TRACE":   7,
	}
)

func methodIndex(method string) int {
	switch method {
	case "GET":
		return 0
	case "POST":
		return 1
	case "DELETE":
		return 2
	case "PUT":
		return 3
	case "PATCH":
		return 4
	case "HEAD":
		return 5
	case "OPTIONS":
		return 6
	case "TRACE":
		return 7
	default:
		return -1
	}
}

// --- route ---

type route struct {
	method, uri string
	handlers    []Handler
}

// routeStore represents a thread-safe store for route pattern.
// 用于检查route pattern是否重复(冲突)及以后打印
type routeStore struct {
	routeMap   map[string]map[string]*route
	routeSlice []*route

	lock sync.Mutex
}

func newRouteStore() *routeStore {
	rs := &routeStore{
		routeMap:   make(map[string]map[string]*route),
		routeSlice: make([]*route, 0),
	}

	for m := range _HTTP_METHODS {
		rs.routeMap[m] = make(map[string]*route)
	}

	return rs
}

func (rs *routeStore) add(r *route) {
	rs.lock.Lock()
	defer rs.lock.Unlock()

	if rs.routeMap[r.method][r.uri] != nil {
		panic(fmt.Sprintf("double uri : %s[%s]", r.method, r.uri))
	}

	rs.routeMap[r.method][r.uri] = r
	rs.routeSlice = append(rs.routeSlice, r)
}

// --- router ---

// 树(单父多子)
type Router struct {
	method  string // 只有终端节点有
	pattern string

	befores  []interface{}
	handlers []interface{} // 只有终端节点有

	parent *Router
	sub    []*Router
}

func NewRouter() *Router {
	return new(Router)
}

func (r *Router) Group(pattern string, fn func(*Router)) {
	rr := &Router{
		pattern: pattern,
		parent:  r,
	}

	r.sub = append(r.sub, rr)

	fn(rr)
}

func (r *Router) Use(handlers ...interface{}) {
	r.befores = append(r.befores, handlers...)
}

func dumpRoute(r *Router, rs *routeStore) {
	if r.sub == nil {
		rs.add(getRoute(r))
		return
	}

	for _, v := range r.sub {
		dumpRoute(v, rs)
	}
}

func getRoute(r *Router) *route {
	ps := []string{}
	hs := []interface{}{}

	tmp := r
	for {
		ps = append(ps, tmp.pattern)

		if len(tmp.handlers) > 0 {
			hs = append(hs, tmp.handlers...)
		}
		if len(tmp.befores) > 0 {
			hstmp := make([]interface{}, len(tmp.befores))

			copy(hstmp, tmp.befores)
			hstmp = append(hstmp, hs...)
			hs = hstmp
		}

		if tmp.parent == nil {
			break
		}

		tmp = tmp.parent
	}

	return &route{
		method:   r.method,
		uri:      strings.Join(reverseStrings(ps), ""),
		handlers: newHandlers(hs),
	}
}

// r is root router.
func (r *Router) Handler() *water {
	if r.parent != nil {
		panic("sub router not allowed: Handler")
	}

	rs := newRouteStore()

	dumpRoute(r, rs)

	w := newWater()
	w.routeStore = rs
	w.buildTree()

	return w
}

// 此时还无法获取Router.afters,因为Router.afters还未执行到
// RouteStor放入Route后才可获取afters的信息
func (r *Router) handle(method, pattern string, handlers []interface{}) {
	rr := &Router{
		method:   method,
		pattern:  pattern,
		parent:   r,
		handlers: handlers,
	}

	r.sub = append(r.sub, rr)
}

func (r *Router) Any(pattern string, handlers ...interface{}) {
	for m := range _HTTP_METHODS {
		r.handle(m, pattern, handlers)
	}
}

func (r *Router) Get(pattern string, handlers ...interface{}) {
	r.handle(http.MethodGet, pattern, handlers)
}

func (r *Router) Post(pattern string, handlers ...interface{}) {
	r.handle(http.MethodPost, pattern, handlers)
}

func (r *Router) Put(pattern string, handlers ...interface{}) {
	r.handle(http.MethodPut, pattern, handlers)
}

func (r *Router) Patch(pattern string, handlers ...interface{}) {
	r.handle(http.MethodPatch, pattern, handlers)
}

func (r *Router) Delete(pattern string, handlers ...interface{}) {
	r.handle(http.MethodDelete, pattern, handlers)
}

func (r *Router) Options(pattern string, handlers ...interface{}) {
	r.handle(http.MethodOptions, pattern, handlers)
}

func (r *Router) Head(pattern string, handlers ...interface{}) {
	r.handle(http.MethodHead, pattern, handlers)
}

func (r *Router) Trace(pattern string, handlers ...interface{}) {
	r.handle(http.MethodTrace, pattern, handlers)
}

func (r *Router) Classic() {
	if r.parent != nil {
		panic("sub router not allowed : Classic()")
	}

	r.Use(Logger())
	r.Use(Recovery())
}
<commit_msg>u 改进错误检查<commit_after>package water

import (
	"fmt"
	"net/http"
	"strings"
	"sync"
)

var (
	_HTTP_METHODS = map[string]int{
		"GET":     0,
		"POST":    1,
		"DELETE":  2,
		"PUT":     3,
		"PATCH":   4,
		"HEAD":    5,
		"OPTIONS": 6,
		"TRACE":   7,
	}
)

func methodIndex(method string) int {
	switch method {
	case "GET":
		return 0
	case "POST":
		return 1
	case "DELETE":
		return 2
	case "PUT":
		return 3
	case "PATCH":
		return 4
	case "HEAD":
		return 5
	case "OPTIONS":
		return 6
	case "TRACE":
		return 7
	default:
		return -1
	}
}

// --- route ---

type route struct {
	method, uri string
	handlers    []Handler
}

// routeStore represents a thread-safe store for route pattern.
// 用于检查route pattern是否重复(冲突)及以后打印
type routeStore struct {
	routeMap   map[string]map[string]*route
	routeSlice []*route

	lock sync.Mutex
}

func newRouteStore() *routeStore {
	rs := &routeStore{
		routeMap:   make(map[string]map[string]*route),
		routeSlice: make([]*route, 0),
	}

	for m := range _HTTP_METHODS {
		rs.routeMap[m] = make(map[string]*route)
	}

	return rs
}

func (rs *routeStore) add(r *route) {
	rs.lock.Lock()
	defer rs.lock.Unlock()

	if rs.routeMap[r.method][r.uri] != nil {
		panic(fmt.Sprintf("double uri : %s[%s]", r.method, r.uri))
	}

	rs.routeMap[r.method][r.uri] = r
	rs.routeSlice = append(rs.routeSlice, r)
}

// --- router ---

// 树(单父多子)
type Router struct {
	method  string // 只有终端节点有
	pattern string

	befores  []interface{}
	handlers []interface{} // 只有终端节点有

	parent *Router
	sub    []*Router
}

func NewRouter() *Router {
	return new(Router)
}

func (r *Router) Group(pattern string, fn func(*Router)) {
	rr := &Router{
		pattern: pattern,
		parent:  r,
	}

	r.sub = append(r.sub, rr)

	fn(rr)
}

func (r *Router) Use(handlers ...interface{}) {
	r.befores = append(r.befores, handlers...)
}

func dumpRoute(r *Router, rs *routeStore) {
	if r.sub == nil {
		rs.add(getRoute(r))
		return
	}

	for _, v := range r.sub {
		dumpRoute(v, rs)
	}
}

func getRoute(r *Router) *route {
	ps := []string{}
	hs := []interface{}{}

	tmp := r
	for {
		ps = append(ps, strings.TrimSpace(tmp.pattern))

		if len(tmp.handlers) > 0 {
			hs = append(hs, tmp.handlers...)
		}
		if len(tmp.befores) > 0 {
			hstmp := make([]interface{}, len(tmp.befores))

			copy(hstmp, tmp.befores)
			hstmp = append(hstmp, hs...)
			hs = hstmp
		}

		if tmp.parent == nil {
			break
		}

		tmp = tmp.parent
	}

	re := &route{
		method:   r.method,
		uri:      strings.Join(reverseStrings(ps), ""),
		handlers: newHandlers(hs),
	}

	if len(re.handlers) == 0 {
		panic(fmt.Sprintf("handler err : empty handlers in route(%s,%s)", re.method, re.uri))
	}

	return re
}

// r is root router.
func (r *Router) Handler() *water {
	if r.parent != nil {
		panic("sub router not allowed: Handler")
	}

	rs := newRouteStore()

	dumpRoute(r, rs)

	w := newWater()
	w.routeStore = rs
	w.buildTree()

	return w
}

// 此时还无法获取Router.afters,因为Router.afters还未执行到
// RouteStor放入Route后才可获取afters的信息
func (r *Router) handle(method, pattern string, handlers []interface{}) {
	for _, v := range handlers {
		if v == nil {
			panic(fmt.Sprintf("handler err : find nil in route(%s,%s)", method, pattern))
		}
	}

	rr := &Router{
		method:   method,
		pattern:  pattern,
		parent:   r,
		handlers: handlers,
	}

	r.sub = append(r.sub, rr)
}

func (r *Router) Any(pattern string, handlers ...interface{}) {
	for m := range _HTTP_METHODS {
		r.handle(m, pattern, handlers)
	}
}

func (r *Router) Get(pattern string, handlers ...interface{}) {
	r.handle(http.MethodGet, pattern, handlers)
}

func (r *Router) Post(pattern string, handlers ...interface{}) {
	r.handle(http.MethodPost, pattern, handlers)
}

func (r *Router) Put(pattern string, handlers ...interface{}) {
	r.handle(http.MethodPut, pattern, handlers)
}

func (r *Router) Patch(pattern string, handlers ...interface{}) {
	r.handle(http.MethodPatch, pattern, handlers)
}

func (r *Router) Delete(pattern string, handlers ...interface{}) {
	r.handle(http.MethodDelete, pattern, handlers)
}

func (r *Router) Options(pattern string, handlers ...interface{}) {
	r.handle(http.MethodOptions, pattern, handlers)
}

func (r *Router) Head(pattern string, handlers ...interface{}) {
	r.handle(http.MethodHead, pattern, handlers)
}

func (r *Router) Trace(pattern string, handlers ...interface{}) {
	r.handle(http.MethodTrace, pattern, handlers)
}

func (r *Router) Classic() {
	if r.parent != nil {
		panic("sub router not allowed : Classic()")
	}

	r.Use(Logger())
	r.Use(Recovery())
}
<|endoftext|>


<commit_before>// largely cribbed from gorilla/pat

package glue

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

	"github.com/gorilla/mux"
)

// newRouter prepares a new router
func newRouter() *router {
	return &router{}
}

// router is a request router that implements a pat-like API.
//
// pat docs: http://godoc.org/github.com/bmizerany/pat
type router struct {
	mux.Router
	NotFoundHandler Handler
}

// Add registers a pattern with a handler for the given request method.
func (r *router) Add(meth, pat string, h Handler) *mux.Route {

	// if not already an http.Handler, wrap it as a routeHandler
	if _, ok := h.(http.Handler); !ok {
		h = routeHandler{Handler: h}
	}

	return r.NewRoute().PathPrefix(pat).Handler(h.(http.Handler)).Methods(meth)
}

// Delete registers a pattern with a handler for DELETE requests.
func (r *router) Delete(pat string, h Handler) *mux.Route {
	return r.Add("DELETE", pat, h)
}

// Get registers a pattern with a handler for GET requests.
func (r *router) Get(pat string, h Handler) *mux.Route {
	return r.Add("GET", pat, h)
}

// Post registers a pattern with a handler for POST requests.
func (r *router) Post(pat string, h Handler) *mux.Route {
	return r.Add("POST", pat, h)
}

// Put registers a pattern with a handler for PUT requests.
func (r *router) Put(pat string, h Handler) *mux.Route {
	return r.Add("PUT", pat, h)
}

// Handle is a glue.Handler that does route matching and invokes the registered
// glue.Handler for a route.
//
// If a route is not found the NotFoundHandler is invoked.
func (r *router) Handle(w http.ResponseWriter, req *http.Request, c Context) {
	// Clean path to canonical form and redirect.
	if p := cleanPath(req.URL.Path); p != req.URL.Path {
		w.Header().Set("Location", p)
		w.WriteHeader(http.StatusMovedPermanently)
		return
	}
	var match mux.RouteMatch
	var handler http.Handler
	if matched := r.Match(req, &match); matched {
		if rhandler, ok := match.Handler.(routeHandler); ok {
			rhandler.ctx = c
			handler = rhandler
		}
		// @todo coersion could fail
		registerVars(req, match.Vars)
	}

	if handler == nil {
		if r.NotFoundHandler == nil {
			handler = http.NotFoundHandler()
		} else {
			handler = routeHandler{r.NotFoundHandler, c}
		}
	}
	handler.ServeHTTP(w, req)
}

// registerVars adds the matched route variables to the URL query.
func registerVars(r *http.Request, vars map[string]string) {
	parts, i := make([]string, len(vars)), 0
	for key, value := range vars {
		parts[i] = url.QueryEscape(":"+key) + "=" + url.QueryEscape(value)
		i++
	}
	q := strings.Join(parts, "&")
	if r.URL.RawQuery == "" {
		r.URL.RawQuery = q
	} else {
		r.URL.RawQuery += "&" + q
	}
}

// cleanPath returns the canonical path for p, eliminating . and .. elements.
// Borrowed from the net/http package.
func cleanPath(p string) string {
	if p == "" {
		return "/"
	}
	if p[0] != '/' {
		p = "/" + p
	}
	np := path.Clean(p)
	// path.Clean removes trailing slash except for root;
	// put the trailing slash back if necessary.
	if p[len(p)-1] == '/' && np != "/" {
		np += "/"
	}
	return np
}

type routeHandler struct {
	Handler
	ctx Context
}

func (rh routeHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	vals, err := rh.ctx.Call(rh.Handler, rh.ctx.g.Injector)
	if err != nil {
		panic(err)
	}
	// if the handler returned something, write it to the http response
	if len(vals) > 0 {
		_, err := rh.ctx.Call(func(rw http.ResponseWriter, rHandler ResponseHandler) {
			rHandler(rw, vals)
		}, rh.ctx.g.Injector)
		if err != nil {
			panic(err)
		}
	}
}
<commit_msg>Set handler regardless of routeHandler type match<commit_after>// largely cribbed from gorilla/pat

package glue

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

	"github.com/gorilla/mux"
)

// newRouter prepares a new router
func newRouter() *router {
	return &router{}
}

// router is a request router that implements a pat-like API.
//
// pat docs: http://godoc.org/github.com/bmizerany/pat
type router struct {
	mux.Router
	NotFoundHandler Handler
}

// Add registers a pattern with a handler for the given request method.
func (r *router) Add(meth, pat string, h Handler) *mux.Route {

	// if not already an http.Handler, wrap it as a routeHandler
	if _, ok := h.(http.Handler); !ok {
		h = routeHandler{Handler: h}
	}

	return r.NewRoute().PathPrefix(pat).Handler(h.(http.Handler)).Methods(meth)
}

// Delete registers a pattern with a handler for DELETE requests.
func (r *router) Delete(pat string, h Handler) *mux.Route {
	return r.Add("DELETE", pat, h)
}

// Get registers a pattern with a handler for GET requests.
func (r *router) Get(pat string, h Handler) *mux.Route {
	return r.Add("GET", pat, h)
}

// Post registers a pattern with a handler for POST requests.
func (r *router) Post(pat string, h Handler) *mux.Route {
	return r.Add("POST", pat, h)
}

// Put registers a pattern with a handler for PUT requests.
func (r *router) Put(pat string, h Handler) *mux.Route {
	return r.Add("PUT", pat, h)
}

// Handle is a glue.Handler that does route matching and invokes the registered
// glue.Handler for a route.
//
// If a route is not found the NotFoundHandler is invoked.
func (r *router) Handle(w http.ResponseWriter, req *http.Request, c Context) {
	// Clean path to canonical form and redirect.
	if p := cleanPath(req.URL.Path); p != req.URL.Path {
		w.Header().Set("Location", p)
		w.WriteHeader(http.StatusMovedPermanently)
		return
	}
	var match mux.RouteMatch
	var handler http.Handler
	if matched := r.Match(req, &match); matched {
		handler = match.Handler
		if rhandler, ok := match.Handler.(routeHandler); ok {
			rhandler.ctx = c
			handler = rhandler
		}
		registerVars(req, match.Vars)
	}

	if handler == nil {
		if r.NotFoundHandler == nil {
			handler = http.NotFoundHandler()
		} else {
			handler = routeHandler{r.NotFoundHandler, c}
		}
	}
	handler.ServeHTTP(w, req)
}

// registerVars adds the matched route variables to the URL query.
func registerVars(r *http.Request, vars map[string]string) {
	parts, i := make([]string, len(vars)), 0
	for key, value := range vars {
		parts[i] = url.QueryEscape(":"+key) + "=" + url.QueryEscape(value)
		i++
	}
	q := strings.Join(parts, "&")
	if r.URL.RawQuery == "" {
		r.URL.RawQuery = q
	} else {
		r.URL.RawQuery += "&" + q
	}
}

// cleanPath returns the canonical path for p, eliminating . and .. elements.
// Borrowed from the net/http package.
func cleanPath(p string) string {
	if p == "" {
		return "/"
	}
	if p[0] != '/' {
		p = "/" + p
	}
	np := path.Clean(p)
	// path.Clean removes trailing slash except for root;
	// put the trailing slash back if necessary.
	if p[len(p)-1] == '/' && np != "/" {
		np += "/"
	}
	return np
}

type routeHandler struct {
	Handler
	ctx Context
}

func (rh routeHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	vals, err := rh.ctx.Call(rh.Handler, rh.ctx.g.Injector)
	if err != nil {
		panic(err)
	}
	// if the handler returned something, write it to the http response
	if len(vals) > 0 {
		_, err := rh.ctx.Call(func(rw http.ResponseWriter, rHandler ResponseHandler) {
			rHandler(rw, vals)
		}, rh.ctx.g.Injector)
		if err != nil {
			panic(err)
		}
	}
}
<|endoftext|>


<commit_before>package govkbot

import (
	"log"
	"strconv"
	"strings"
	"time"
)

const (
	vkAPIURL        = "https://api.vk.com/method/"
	vkAPIVer        = "5.52"
	messagesCount   = 200
	requestInterval = 400 // 3 requests per second VK limit
)

// VKBot - bot config
type VKBot struct {
	msgRoutes        map[string]func(*Message) string
	actionRoutes     map[string]func(*Message) string
	cmdHandlers      map[string]func(*Message) string
	msgHandlers      map[string]func(*Message) string
	errorHandler     func(*Message, error)
	LastMsg          int
	markedMessages   map[int]*Message
	lastUserMessages map[int]int
	lastChatMessages map[int]int
	autoFriend       bool
}

var bot = newBot()

//API - bot API
var API = newAPI()

// SetDebug - enable/disable debug messages logging
func SetDebug(debug bool) {
	API.DEBUG = debug
}

func newBot() *VKBot {
	return &VKBot{
		msgRoutes:        make(map[string]func(*Message) string),
		actionRoutes:     make(map[string]func(*Message) string),
		markedMessages:   make(map[int]*Message),
		lastUserMessages: make(map[int]int),
		lastChatMessages: make(map[int]int),
	}
}

func newAPI() *VkAPI {
	return &VkAPI{
		Token:           "",
		URL:             vkAPIURL,
		Ver:             vkAPIVer,
		MessagesCount:   messagesCount,
		RequestInterval: requestInterval,
		DEBUG:           false,
		HTTPS:           true,
	}
}

// SetToken - set bot token
func SetToken(token string) {
	API.Token = token
}

// SetAutoFriend - enables mutual auto friending
func SetAutoFriend(af bool) {
	bot.autoFriend = af
}

// SetAPI - setup API config
func SetAPI(token string, url string, ver string) {
	SetToken(token)
	if url != "" {
		API.URL = url
	}
	if ver != "" {
		API.Ver = ver
	}
}

// SetLang - sets VK response language. Default auto. Available: en, ru, ua, be, es, fi, de, it
func SetLang(lang string) {
	API.Lang = lang
}

// HandleMessage - add substr message handler.
// Function must return string to reply or "" (if no reply)
// You can use m.Reply(string), if need more replies in handler
func HandleMessage(command string, handler func(*Message) string) {
	bot.msgRoutes[command] = handler
}

// HandleAction - add action handler.
// Function must return string to reply or "" (if no reply)
// You can use m.Reply(string), if need more replies in handler
func HandleAction(command string, handler func(*Message) string) {
	bot.actionRoutes[command] = handler
}

// HandleError - add error handler
func HandleError(handler func(*Message, error)) {
	bot.errorHandler = handler
}

// GetMessages - request unread messages from VK (more than 200)
func GetMessages() ([]*Message, error) {
	var allMessages []*Message
	lastMsg := bot.LastMsg
	offset := 0
	var err error
	var messages *Messages
	for {
		messages, err = API.GetMessages(API.MessagesCount, offset)
		if len(messages.Items) > 0 {
			if messages.Items[0].ID > lastMsg {
				lastMsg = messages.Items[0].ID
			}
		}
		allMessages = append(allMessages, messages.Items...)
		if bot.LastMsg > 0 {
			if len(messages.Items) > 0 {
				if messages.Items[len(messages.Items)-1].ID <= bot.LastMsg {
					bot.LastMsg = lastMsg
					break
				}
			} else {
				break
			}
			offset += API.MessagesCount
		} else {
			bot.LastMsg = lastMsg
			break
		}
	}
	if offset > 0 {
		API.NotifyAdmin("many messages in interval. offset: " + strconv.Itoa(offset))
	}
	return allMessages, err
}

func sendError(msg *Message, err error) {
	if bot.errorHandler != nil {
		bot.errorHandler(msg, err)
	} else {
		log.Fatalf("VKBot error: %+v\n", err.Error())
	}

}

//RouteAction routes an action
func RouteAction(m *Message) (replies []string, err error) {
	if m.Action != "" {
		debugPrint("route action: %+v\n", m.Action)
		for k, v := range bot.actionRoutes {
			if m.Action == k {
				bot.markedMessages[m.ID] = m
				msg := v(m)
				if msg != "" {
					replies = append(replies, msg)
				}
			}
		}
	}
	return replies, nil
}

// RouteMessage routes single message
func RouteMessage(m *Message) (replies []string, err error) {
	message := strings.TrimSpace(strings.ToLower(m.Body))
	if HasPrefix(message, "/ ") {
		message = "/" + TrimPrefix(message, "/ ")
	}
	if m.Action != "" {
		replies, err = RouteAction(m)
		return replies, err
	}
	marked := false
	for k, v := range bot.msgRoutes {
		if HasPrefix(message, k) {
			msg := v(m)
			if msg != "" {
				marked = true
				_, ok := bot.markedMessages[m.ID]
				if ok {
					delete(bot.markedMessages, m.ID)
				}
				replies = append(replies, msg)
			} else {
				if !marked {
					bot.markedMessages[m.ID] = m
				}
			}
		}
	}
	return replies, nil
}

// RouteMessages routes inbound messages
func RouteMessages(messages []*Message) (result map[*Message][]string) {
	result = make(map[*Message][]string)
	for _, m := range messages {
		//if m.ID <= bot.LastMsg {
		//	break
		//}
		if m.ReadState == 0 {
			replies, err := RouteMessage(m)
			if err != nil {
				sendError(m, err)
			}
			if len(replies) > 0 {
				result[m] = replies
			}
		}
	}
	return result
}

// MainRoute - main router func. Working cycle Listen.
func MainRoute() {
	bot.markedMessages = make(map[int]*Message)
	messages, err := GetMessages()
	if err != nil {
		sendError(nil, err)
	}
	replies := RouteMessages(messages)
	for m, msgs := range replies {
		for _, msg := range msgs {
			if msg != "" {
				_, err = m.Reply(msg)
				if err != nil {
					log.Printf("Error sending message: '%+v'\n", msg)
					sendError(m, err)
				}
			}
		}
	}

	for _, m := range bot.markedMessages {
		m.MarkAsRead()
	}
}

// Listen - start server
func Listen(token string, url string, ver string, adminID int) {
	SetAPI(token, url, ver)
	API.AdminID = adminID
	u, err := API.Me()
	if err != nil {
		sendError(nil, err)
	}
	API.UID = u.ID

	go friendReceiver()

	c := time.Tick(3 * time.Second)
	for range c {
		MainRoute()
	}
}

// CheckFriends checking friend invites and mathes and deletes mutual
func CheckFriends() {
	uids, _ := API.GetFriendRequests(false)
	if len(uids) > 0 {
		for _, uid := range uids {
			API.AddFriend(uid)
			for k, v := range bot.actionRoutes {
				if k == "friend_add" {
					m := Message{Action: "friend_add", UserID: uid}
					v(&m)
				}
			}
		}
	}
	uids, _ = API.GetFriendRequests(true)
	if len(uids) > 0 {
		for _, uid := range uids {
			API.DeleteFriend(uid)
			for k, v := range bot.actionRoutes {
				if k == "friend_delete" {
					m := Message{Action: "friend_delete", UserID: uid}
					v(&m)
				}
			}
		}
	}
}

func friendReceiver() {
	CheckFriends()
	c := time.Tick(30 * time.Second)
	for range c {
		CheckFriends()
	}
}

// NotifyAdmin - notify AdminID by VK
func NotifyAdmin(msg string) error {
	return API.NotifyAdmin(msg)
}
<commit_msg>alert if bot try send wrong message<commit_after>package govkbot

import (
	"log"
	"strconv"
	"strings"
	"time"
)

const (
	vkAPIURL        = "https://api.vk.com/method/"
	vkAPIVer        = "5.52"
	messagesCount   = 200
	requestInterval = 400 // 3 requests per second VK limit
)

// VKBot - bot config
type VKBot struct {
	msgRoutes        map[string]func(*Message) string
	actionRoutes     map[string]func(*Message) string
	cmdHandlers      map[string]func(*Message) string
	msgHandlers      map[string]func(*Message) string
	errorHandler     func(*Message, error)
	LastMsg          int
	markedMessages   map[int]*Message
	lastUserMessages map[int]int
	lastChatMessages map[int]int
	autoFriend       bool
}

var bot = newBot()

//API - bot API
var API = newAPI()

// SetDebug - enable/disable debug messages logging
func SetDebug(debug bool) {
	API.DEBUG = debug
}

func newBot() *VKBot {
	return &VKBot{
		msgRoutes:        make(map[string]func(*Message) string),
		actionRoutes:     make(map[string]func(*Message) string),
		markedMessages:   make(map[int]*Message),
		lastUserMessages: make(map[int]int),
		lastChatMessages: make(map[int]int),
	}
}

func newAPI() *VkAPI {
	return &VkAPI{
		Token:           "",
		URL:             vkAPIURL,
		Ver:             vkAPIVer,
		MessagesCount:   messagesCount,
		RequestInterval: requestInterval,
		DEBUG:           false,
		HTTPS:           true,
	}
}

// SetToken - set bot token
func SetToken(token string) {
	API.Token = token
}

// SetAutoFriend - enables mutual auto friending
func SetAutoFriend(af bool) {
	bot.autoFriend = af
}

// SetAPI - setup API config
func SetAPI(token string, url string, ver string) {
	SetToken(token)
	if url != "" {
		API.URL = url
	}
	if ver != "" {
		API.Ver = ver
	}
}

// SetLang - sets VK response language. Default auto. Available: en, ru, ua, be, es, fi, de, it
func SetLang(lang string) {
	API.Lang = lang
}

// HandleMessage - add substr message handler.
// Function must return string to reply or "" (if no reply)
// You can use m.Reply(string), if need more replies in handler
func HandleMessage(command string, handler func(*Message) string) {
	bot.msgRoutes[command] = handler
}

// HandleAction - add action handler.
// Function must return string to reply or "" (if no reply)
// You can use m.Reply(string), if need more replies in handler
func HandleAction(command string, handler func(*Message) string) {
	bot.actionRoutes[command] = handler
}

// HandleError - add error handler
func HandleError(handler func(*Message, error)) {
	bot.errorHandler = handler
}

// GetMessages - request unread messages from VK (more than 200)
func GetMessages() ([]*Message, error) {
	var allMessages []*Message
	lastMsg := bot.LastMsg
	offset := 0
	var err error
	var messages *Messages
	for {
		messages, err = API.GetMessages(API.MessagesCount, offset)
		if len(messages.Items) > 0 {
			if messages.Items[0].ID > lastMsg {
				lastMsg = messages.Items[0].ID
			}
		}
		allMessages = append(allMessages, messages.Items...)
		if bot.LastMsg > 0 {
			if len(messages.Items) > 0 {
				if messages.Items[len(messages.Items)-1].ID <= bot.LastMsg {
					bot.LastMsg = lastMsg
					break
				}
			} else {
				break
			}
			offset += API.MessagesCount
		} else {
			bot.LastMsg = lastMsg
			break
		}
	}
	if offset > 0 {
		API.NotifyAdmin("many messages in interval. offset: " + strconv.Itoa(offset))
	}
	return allMessages, err
}

func sendError(msg *Message, err error) {
	if bot.errorHandler != nil {
		bot.errorHandler(msg, err)
	} else {
		log.Fatalf("VKBot error: %+v\n", err.Error())
	}

}

//RouteAction routes an action
func RouteAction(m *Message) (replies []string, err error) {
	if m.Action != "" {
		debugPrint("route action: %+v\n", m.Action)
		for k, v := range bot.actionRoutes {
			if m.Action == k {
				bot.markedMessages[m.ID] = m
				msg := v(m)
				if msg != "" {
					replies = append(replies, msg)
				}
			}
		}
	}
	return replies, nil
}

// RouteMessage routes single message
func RouteMessage(m *Message) (replies []string, err error) {
	message := strings.TrimSpace(strings.ToLower(m.Body))
	if HasPrefix(message, "/ ") {
		message = "/" + TrimPrefix(message, "/ ")
	}
	if m.Action != "" {
		replies, err = RouteAction(m)
		return replies, err
	}
	marked := false
	for k, v := range bot.msgRoutes {
		if HasPrefix(message, k) {
			msg := v(m)
			if msg != "" {
				marked = true
				_, ok := bot.markedMessages[m.ID]
				if ok {
					delete(bot.markedMessages, m.ID)
				}
				replies = append(replies, msg)
			} else {
				if !marked {
					bot.markedMessages[m.ID] = m
				}
			}
		}
	}
	return replies, nil
}

// RouteMessages routes inbound messages
func RouteMessages(messages []*Message) (result map[*Message][]string) {
	result = make(map[*Message][]string)
	for _, m := range messages {
		//if m.ID <= bot.LastMsg {
		//	break
		//}
		if m.ReadState == 0 {
			replies, err := RouteMessage(m)
			if err != nil {
				sendError(m, err)
			}
			if len(replies) > 0 {
				result[m] = replies
			}
		}
	}
	return result
}

// MainRoute - main router func. Working cycle Listen.
func MainRoute() {
	bot.markedMessages = make(map[int]*Message)
	messages, err := GetMessages()
	if err != nil {
		sendError(nil, err)
	}
	replies := RouteMessages(messages)
	for m, msgs := range replies {
		for _, msg := range msgs {
			if msg != "" {
				_, err = m.Reply(msg)
				if err != nil {
					log.Printf("Error sending message: '%+v'\n", msg)
					sendError(m, err)
					_, err = m.Reply("Cant send message, maybe wrong/china letters?")
					if err != nil {
						sendError(m, err)
					}
				}
			}
		}
	}

	for _, m := range bot.markedMessages {
		m.MarkAsRead()
	}
}

// Listen - start server
func Listen(token string, url string, ver string, adminID int) {
	SetAPI(token, url, ver)
	API.AdminID = adminID
	u, err := API.Me()
	if err != nil {
		sendError(nil, err)
	}
	API.UID = u.ID

	go friendReceiver()

	c := time.Tick(3 * time.Second)
	for range c {
		MainRoute()
	}
}

// CheckFriends checking friend invites and mathes and deletes mutual
func CheckFriends() {
	uids, _ := API.GetFriendRequests(false)
	if len(uids) > 0 {
		for _, uid := range uids {
			API.AddFriend(uid)
			for k, v := range bot.actionRoutes {
				if k == "friend_add" {
					m := Message{Action: "friend_add", UserID: uid}
					v(&m)
				}
			}
		}
	}
	uids, _ = API.GetFriendRequests(true)
	if len(uids) > 0 {
		for _, uid := range uids {
			API.DeleteFriend(uid)
			for k, v := range bot.actionRoutes {
				if k == "friend_delete" {
					m := Message{Action: "friend_delete", UserID: uid}
					v(&m)
				}
			}
		}
	}
}

func friendReceiver() {
	CheckFriends()
	c := time.Tick(30 * time.Second)
	for range c {
		CheckFriends()
	}
}

// NotifyAdmin - notify AdminID by VK
func NotifyAdmin(msg string) error {
	return API.NotifyAdmin(msg)
}
<|endoftext|>


<commit_before>/*

   Copyright 2013 Niklas Voss

   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 golem

import (
	"errors"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"reflect"
)

// Router handles multiplexing of incoming messenges by typenames/events.
// Initially a router uses heartbeats and the default protocol.
type Router struct {
	// Map of callbacks for event types.
	callbacks map[string]func(*Connection, interface{})
	// Protocol extensions
	extensions map[reflect.Type]reflect.Value
	// Function being called if connection is closed.
	closeFunc func(*Connection)
	// Function called after handshake when a WebSocket connection
	// was succesfully established.
	connectionFunc func(*Connection)
	// Function verifying handshake.
	handshakeFunc func(http.ResponseWriter, *http.Request) bool
	// Active protocol
	protocol Protocol
	// Flag to enable or disable heartbeats
	useHeartbeats bool
	//
	connExtensionConstructor reflect.Value
}

// NewRouter intialises a new instance and returns the pointer.
func NewRouter() *Router {
	// Tries to run hub, if already running nothing will happen.
	hub.run()
	// Returns pointer to instance.
	return &Router{
		callbacks:                make(map[string]func(*Connection, interface{})),
		extensions:               make(map[reflect.Type]reflect.Value),
		closeFunc:                func(*Connection) {}, // Empty placeholder close function.
		connectionFunc:           func(*Connection) {},
		handshakeFunc:            func(http.ResponseWriter, *http.Request) bool { return true }, // Handshake always allowed.
		protocol:                 initialProtocol,
		useHeartbeats:            true,
		connExtensionConstructor: defaultConnectionExtension,
	}
}

// Handler creates a handler function for this router, that can be used with the
// http-package to handle WebSocket-Connections.
func (router *Router) Handler() func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		// Check if method used was GET.
		if r.Method != "GET" {
			http.Error(w, "Method not allowed", 405)
			return
		}
		// Disallow cross-origin connections.
		if r.Header.Get("Origin") != "http://"+r.Host {
			http.Error(w, "Origin not allowed", 403)
			return
		}
		// Check if handshake callback verifies upgrade.
		if !router.handshakeFunc(w, r) {
			http.Error(w, "Authorization failed", 403)
			return
		}
		// Upgrade websocket connection.
		socket, err := websocket.Upgrade(w, r, nil, 1024, 1024)
		// Check if handshake was successful
		if _, ok := err.(websocket.HandshakeError); ok {
			http.Error(w, "Not a websocket handshake", 400)
			return
		} else if err != nil {
			log.Println(err)
			return
		}

		// Create the connection.
		conn := newConnection(socket, router)
		//
		if router.connExtensionConstructor.IsValid() {
			conn.extend(router.connExtensionConstructor.Call([]reflect.Value{reflect.ValueOf(conn)})[0].Interface())
		}

		// Connection established with possible extension, so callback
		router.connectionFunc(conn)

		// And start reading and writing routines.
		conn.run()
	}
}

// The On-function adds callbacks by name of the event, that should be handled.
// For type T the callback would be of type:
//     func (*golem.Connection, *T)
// Type T can be any type. By default golem tries to unmarshal json into the
// specified type. If a custom protocol is used, it will be used instead to process the data.
// If type T is registered to use a protocol extension, it will be used instead.
// If type T is interface{} the interstage data of the active protocol will be directly forwarded!
// (Note: the golem wiki has a whole page about this function)
func (router *Router) On(name string, callback interface{}) {

	callbackValue := reflect.ValueOf(callback)
	callbackType := reflect.TypeOf(callback)
	if router.connExtensionConstructor.IsValid() {
		extType := router.connExtensionConstructor.Type().Out(0)
		if callbackType.In(0) == extType {
			// EXTENSION TYPE

			// NO DATA
			if callbackType.NumIn() == 1 {
				router.callbacks[name] = func(conn *Connection, data interface{}) {
					args := []reflect.Value{reflect.ValueOf(conn.extension)}
					callbackValue.Call(args)
				}
				return
			}

			// INTERFACE
			if callbackType.In(1).Kind() == reflect.Interface {
				router.callbacks[name] = func(conn *Connection, data interface{}) {
					args := []reflect.Value{reflect.ValueOf(conn.extension), reflect.ValueOf(data)}
					callbackValue.Call(args)
				}
				return
			}

			// PROTOCOL EXTENSION
			if parser, ok := router.extensions[callbackType.In(1)]; ok {
				router.callbacks[name] = func(conn *Connection, data interface{}) {
					if result := parser.Call([]reflect.Value{reflect.ValueOf(data)}); result[1].Bool() {
						args := []reflect.Value{reflect.ValueOf(conn.extension), result[0]}
						callbackValue.Call(args)
					}
				}
				return
			}

			// PROTOCOL
			callbackDataElem := callbackType.In(1).Elem()
			router.callbacks[name] = func(conn *Connection, data interface{}) {
				result := reflect.New(callbackDataElem)

				err := router.protocol.Unmarshal(data, result.Interface())
				if err == nil {
					args := []reflect.Value{reflect.ValueOf(conn.extension), result}
					callbackValue.Call(args)
				} else {
					// TODO: Proper debug output!
				}
			}
			return
		}
	} else {
		// DEFAULT TYPE

		// NO DATA
		if reflect.TypeOf(callback).NumIn() == 1 {
			router.callbacks[name] = func(conn *Connection, data interface{}) {
				callback.(func(*Connection))(conn)
			}
			return
		}

		// INTERFACE
		if cb, ok := callback.(func(*Connection, interface{})); ok {
			router.callbacks[name] = cb
			return
		}

		// PROTOCOL EXTENSION
		if parser, ok := router.extensions[callbackType.In(1)]; ok {
			router.callbacks[name] = func(conn *Connection, data interface{}) {
				if result := parser.Call([]reflect.Value{reflect.ValueOf(data)}); result[1].Bool() {
					args := []reflect.Value{reflect.ValueOf(conn), result[0]}
					callbackValue.Call(args)
				}
			}
			return
		}

		// PROTOCOL
		callbackDataElem := callbackType.In(1).Elem()
		router.callbacks[name] = func(conn *Connection, data interface{}) {
			result := reflect.New(callbackDataElem)

			err := router.protocol.Unmarshal(data, result.Interface())
			if err == nil {
				args := []reflect.Value{reflect.ValueOf(conn), result}
				callbackValue.Call(args)
			} else {
				// TODO: Proper debug output!
			}
		}
		return
	}
}

// Unpacks incoming data and forwards it to callback.
func (router *Router) processMessage(conn *Connection, in []byte) {
	if name, data, err := router.protocol.Unpack(in); err == nil {
		if callback, ok := router.callbacks[name]; ok {
			callback(conn, data)
		}
	} // TODO: else error logging?

	defer recover()
}

// OnClose sets the callback, that is called when the connection is closed.
// It accept function of the type func(*Connection) by default or functions
// taking extended connection types if previously registered.
func (router *Router) OnClose(callback interface{}) error { //func(*Connection)) {
	if cb, ok := callback.(func(*Connection)); ok {
		router.closeFunc = cb
	} else {
		if router.connExtensionConstructor.IsValid() {
			callbackValue := reflect.ValueOf(callback)
			extType := router.connExtensionConstructor.Type().Out(0)
			if reflect.TypeOf(callback).In(0) == extType {
				router.closeFunc = func(conn *Connection) {
					callbackValue.Call([]reflect.Value{reflect.ValueOf(conn.extension)})
				}
			} else {
				return errors.New("OnClose cannot accept a callback of the type " + reflect.TypeOf(callback).String() + ".")
			}
		} else {
			return errors.New("OnClose can only accept functions of the type func(*Connection), if no extension is registered.")
		}
	}
	return nil
}

// OnConnection sets the callback, that is called when a websocket connection
// was successfully established, it is therefore called after the handshake.
// It accept function of the type func(*Connection) by default or functions
// taking extended connection types if previously registered.
func (router *Router) OnConnect(callback interface{}) error { //func(*Connection)) {
	if cb, ok := callback.(func(*Connection)); ok {
		router.connectionFunc = cb
	} else {
		if router.connExtensionConstructor.IsValid() {
			callbackValue := reflect.ValueOf(callback)
			extType := router.connExtensionConstructor.Type().Out(0)
			if reflect.TypeOf(callback).In(0) == extType {
				router.connectionFunc = func(conn *Connection) {
					callbackValue.Call([]reflect.Value{reflect.ValueOf(conn.extension)})
				}
			} else {
				return errors.New("OnConnection cannot accept a callback of the type " + reflect.TypeOf(callback).String() + ".")
			}
		} else {
			return errors.New("OnConnection can only accept functions of the type func(*Connection), if no extension is registered.")
		}
	}
	return nil
}

// OnHandshake sets the callback for handshake verfication.
// If the handshake function returns false the request will not be upgraded.
func (router *Router) OnHandshake(callback func(http.ResponseWriter, *http.Request) bool) {
	router.handshakeFunc = callback
}

// The AddProtocolExtension-function allows adding of custom parsers for custom types. For any Type T
// the parser function would look like this:
//      func (interface{}) (T, bool)
// The interface's type is depending on the interstage product of the active protocol, by default
// for the JSON-based protocol it is []byte and therefore the function could be simplified to:
//      func ([]byte) (T, bool)
// Or in general if P is the interstage product:
//      func (*P) (T, bool)
// The boolean return value is necessary to verify if parsing was successful.
// All On-handling function accepting T as input data will now automatically use the custom
// extension. For an example see the example_data.go file in the example repository.
func (router *Router) AddProtocolExtension(extensionFunc interface{}) error {
	extensionValue := reflect.ValueOf(extensionFunc)
	extensionType := extensionValue.Type()

	if extensionType.NumIn() != 1 {
		return errors.New("Cannot add function(" + extensionType.String() + ") as parser: To many arguments!")
	}
	if extensionType.NumOut() != 2 {
		return errors.New("Cannot add function(" + extensionType.String() + ") as parser: Wrong number of return values!")
	}
	if extensionType.Out(1).Kind() != reflect.Bool {
		return errors.New("Cannot add function(" + extensionType.String() + ") as parser: Second return value is not Bool!")
	}

	router.extensions[extensionType.Out(0)] = extensionValue

	return nil
}

// SetConnectionExtension sets the extension for this router. A connection extension is an extended connection structure, that afterwards can be
// used by On-handlers as well (use-cases: add persistent storage to connection, additional methods et cetera).
// The SetConnectionExtension function takes the constructor of the custom format to be able to use and create it on
// connection to the router.
// For type E the constructor needs to fulfil the following requirements:
//     func NewE(conn *Connection) *E
// Afterwards On-handler can us this extended type:
//     router.On(func (extendedConn E, data Datatype) { ... })
// For an example have a look at the example repository and have a look at the 'example_connection_extension.go'.
func (router *Router) SetConnectionExtension(constructor interface{}) {
	router.connExtensionConstructor = reflect.ValueOf(constructor)
}

// SetProtocol sets the protocol of the router to the supplied implementation of the Protocol interface.
func (router *Router) SetProtocol(protocol Protocol) {
	router.protocol = protocol
}

//

// SetHeartbeat activates or deactivates the heartbeat depending on the flag parameter. By default heartbeats are activated.
func (router *Router) SetHeartbeat(flag bool) {
	router.useHeartbeats = flag
}
<commit_msg>Be much better about origin handling, and pass the http.Request object through to the OnConnect callback so that request metadata can be used for connection setup.<commit_after>/*

   Copyright 2013 Niklas Voss

   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 golem

import (
	"errors"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"reflect"
)

// Router handles multiplexing of incoming messenges by typenames/events.
// Initially a router uses heartbeats and the default protocol.
type Router struct {
	// Map of callbacks for event types.
	callbacks map[string]func(*Connection, interface{})
	// Protocol extensions
	extensions map[reflect.Type]reflect.Value
	// Function being called if connection is closed.
	closeFunc func(*Connection)
	// Function called after handshake when a WebSocket connection
	// was succesfully established.
	connectionFunc func(*Connection, *http.Request)
	// Function verifying handshake.
	handshakeFunc func(http.ResponseWriter, *http.Request) bool
	// Active protocol
	protocol Protocol
	// Flag to enable or disable heartbeats
	useHeartbeats bool
	//
	connExtensionConstructor reflect.Value
	// If set, the values the Origin header will be checked against and access is only allowed
	// on a match; otherwise no Origin checking is performed. *This overrides the
	// Access-Control-Allow-Origin header!*
	Origins []string
}

// NewRouter intialises a new instance and returns the pointer.
func NewRouter() *Router {
	// Tries to run hub, if already running nothing will happen.
	hub.run()
	// Returns pointer to instance.
	return &Router{
		callbacks:                make(map[string]func(*Connection, interface{})),
		extensions:               make(map[reflect.Type]reflect.Value),
		closeFunc:                func(*Connection) {}, // Empty placeholder close function.
		connectionFunc:           func(*Connection, *http.Request) {},
		handshakeFunc:            func(http.ResponseWriter, *http.Request) bool { return true }, // Handshake always allowed.
		protocol:                 initialProtocol,
		useHeartbeats:            true,
		connExtensionConstructor: defaultConnectionExtension,
		Origins:                  make([]string, 0),
	}
}

// Handler creates a handler function for this router, that can be used with the
// http-package to handle WebSocket-Connections.
func (router *Router) Handler() func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		// Check if method used was GET.
		if r.Method != "GET" {
			http.Error(w, "Method not allowed", 405)
			return
		}

		// Disallow cross-origin connections.
		if len(router.Origins) > 0 {
			originFound := false
			for _, origin := range router.Origins {
				if r.Header.Get("Origin") == origin {
					originFound = true
					break
				}
			}
			if !originFound {
				http.Error(w, "Origin not allowed", 403)
				return
			}
		} else {
			if len(r.Header.Get("Access-Control-Allow-Origin")) > 0 {
				allowedOrigin := r.Header.Get("Access-Control-Allow-Origin")
				if allowedOrigin != "*" {
					if r.URL.Scheme+"://"+r.Host != allowedOrigin {
						http.Error(w, "Origin not allwed", 403)
						return
					}
				}
			}
		}

		// Check if handshake callback verifies upgrade.
		if !router.handshakeFunc(w, r) {
			http.Error(w, "Authorization failed", 403)
			return
		}

		// Upgrade websocket connection.
		protocols := websocket.Subprotocols(r)
		var responseHeader http.Header = nil
		if len(protocols) > 0 {
			responseHeader = http.Header{"Sec-Websocket-Protocol": {protocols[0]}}
		}
		socket, err := websocket.Upgrade(w, r, responseHeader, 1024, 1024)
		// Check if handshake was successful
		if _, ok := err.(websocket.HandshakeError); ok {
			http.Error(w, "Not a websocket handshake", 400)
			return
		} else if err != nil {
			log.Println(err)
			return
		}

		// Create the connection.
		conn := newConnection(socket, router)
		//
		if router.connExtensionConstructor.IsValid() {
			conn.extend(router.connExtensionConstructor.Call([]reflect.Value{reflect.ValueOf(conn)})[0].Interface())
		}

		// Connection established with possible extension, so callback
		router.connectionFunc(conn, r)

		// And start reading and writing routines.
		conn.run()
	}
}

// The On-function adds callbacks by name of the event, that should be handled.
// For type T the callback would be of type:
//     func (*golem.Connection, *T)
// Type T can be any type. By default golem tries to unmarshal json into the
// specified type. If a custom protocol is used, it will be used instead to process the data.
// If type T is registered to use a protocol extension, it will be used instead.
// If type T is interface{} the interstage data of the active protocol will be directly forwarded!
// (Note: the golem wiki has a whole page about this function)
func (router *Router) On(name string, callback interface{}) {

	callbackValue := reflect.ValueOf(callback)
	callbackType := reflect.TypeOf(callback)
	if router.connExtensionConstructor.IsValid() {
		extType := router.connExtensionConstructor.Type().Out(0)
		if callbackType.In(0) == extType {
			// EXTENSION TYPE

			// NO DATA
			if callbackType.NumIn() == 1 {
				router.callbacks[name] = func(conn *Connection, data interface{}) {
					args := []reflect.Value{reflect.ValueOf(conn.extension)}
					callbackValue.Call(args)
				}
				return
			}

			// INTERFACE
			if callbackType.In(1).Kind() == reflect.Interface {
				router.callbacks[name] = func(conn *Connection, data interface{}) {
					args := []reflect.Value{reflect.ValueOf(conn.extension), reflect.ValueOf(data)}
					callbackValue.Call(args)
				}
				return
			}

			// PROTOCOL EXTENSION
			if parser, ok := router.extensions[callbackType.In(1)]; ok {
				router.callbacks[name] = func(conn *Connection, data interface{}) {
					if result := parser.Call([]reflect.Value{reflect.ValueOf(data)}); result[1].Bool() {
						args := []reflect.Value{reflect.ValueOf(conn.extension), result[0]}
						callbackValue.Call(args)
					}
				}
				return
			}

			// PROTOCOL
			callbackDataElem := callbackType.In(1).Elem()
			router.callbacks[name] = func(conn *Connection, data interface{}) {
				result := reflect.New(callbackDataElem)

				err := router.protocol.Unmarshal(data, result.Interface())
				if err == nil {
					args := []reflect.Value{reflect.ValueOf(conn.extension), result}
					callbackValue.Call(args)
				} else {
					// TODO: Proper debug output!
				}
			}
			return
		}
	} else {
		// DEFAULT TYPE

		// NO DATA
		if reflect.TypeOf(callback).NumIn() == 1 {
			router.callbacks[name] = func(conn *Connection, data interface{}) {
				callback.(func(*Connection))(conn)
			}
			return
		}

		// INTERFACE
		if cb, ok := callback.(func(*Connection, interface{})); ok {
			router.callbacks[name] = cb
			return
		}

		// PROTOCOL EXTENSION
		if parser, ok := router.extensions[callbackType.In(1)]; ok {
			router.callbacks[name] = func(conn *Connection, data interface{}) {
				if result := parser.Call([]reflect.Value{reflect.ValueOf(data)}); result[1].Bool() {
					args := []reflect.Value{reflect.ValueOf(conn), result[0]}
					callbackValue.Call(args)
				}
			}
			return
		}

		// PROTOCOL
		callbackDataElem := callbackType.In(1).Elem()
		router.callbacks[name] = func(conn *Connection, data interface{}) {
			result := reflect.New(callbackDataElem)

			err := router.protocol.Unmarshal(data, result.Interface())
			if err == nil {
				args := []reflect.Value{reflect.ValueOf(conn), result}
				callbackValue.Call(args)
			} else {
				// TODO: Proper debug output!
			}
		}
		return
	}
}

// Unpacks incoming data and forwards it to callback.
func (router *Router) processMessage(conn *Connection, in []byte) {
	if name, data, err := router.protocol.Unpack(in); err == nil {
		if callback, ok := router.callbacks[name]; ok {
			callback(conn, data)
		}
	} // TODO: else error logging?

	defer recover()
}

// OnClose sets the callback, that is called when the connection is closed.
// It accept function of the type func(*Connection) by default or functions
// taking extended connection types if previously registered.
func (router *Router) OnClose(callback interface{}) error { //func(*Connection)) {
	if cb, ok := callback.(func(*Connection)); ok {
		router.closeFunc = cb
	} else {
		if router.connExtensionConstructor.IsValid() {
			callbackValue := reflect.ValueOf(callback)
			extType := router.connExtensionConstructor.Type().Out(0)
			if reflect.TypeOf(callback).In(0) == extType {
				router.closeFunc = func(conn *Connection) {
					callbackValue.Call([]reflect.Value{reflect.ValueOf(conn.extension)})
				}
			} else {
				return errors.New("OnClose cannot accept a callback of the type " + reflect.TypeOf(callback).String() + ".")
			}
		} else {
			return errors.New("OnClose can only accept functions of the type func(*Connection), if no extension is registered.")
		}
	}
	return nil
}

// OnConnection sets the callback, that is called when a websocket connection
// was successfully established, it is therefore called after the handshake.
// It accept function of the type func(*Connection, *http.Request) by default or functions
// taking extended connection types if previously registered.
// The http.Request object can be used for connection metadata information.
func (router *Router) OnConnect(callback interface{}) error { //func(*Connection)) {
	if cb, ok := callback.(func(*Connection, *http.Request)); ok {
		router.connectionFunc = cb
	} else {
		if router.connExtensionConstructor.IsValid() {
			callbackValue := reflect.ValueOf(callback)
			extType := router.connExtensionConstructor.Type().Out(0)
			if reflect.TypeOf(callback).In(0) == extType {
				router.connectionFunc = func(conn *Connection, hr *http.Request) {
					callbackValue.Call([]reflect.Value{reflect.ValueOf(conn.extension), reflect.ValueOf(hr)})
				}
			} else {
				return errors.New("OnConnection cannot accept a callback of the type " + reflect.TypeOf(callback).String() + ".")
			}
		} else {
			return errors.New("OnConnection can only accept functions of the type func(*Connection), if no extension is registered.")
		}
	}
	return nil
}

// OnHandshake sets the callback for handshake verfication.
// If the handshake function returns false the request will not be upgraded.
// The http.Request object will be passed into OnConnect as well.
func (router *Router) OnHandshake(callback func(http.ResponseWriter, *http.Request) bool) {
	router.handshakeFunc = callback
}

// The AddProtocolExtension-function allows adding of custom parsers for custom types. For any Type T
// the parser function would look like this:
//      func (interface{}) (T, bool)
// The interface's type is depending on the interstage product of the active protocol, by default
// for the JSON-based protocol it is []byte and therefore the function could be simplified to:
//      func ([]byte) (T, bool)
// Or in general if P is the interstage product:
//      func (*P) (T, bool)
// The boolean return value is necessary to verify if parsing was successful.
// All On-handling function accepting T as input data will now automatically use the custom
// extension. For an example see the example_data.go file in the example repository.
func (router *Router) AddProtocolExtension(extensionFunc interface{}) error {
	extensionValue := reflect.ValueOf(extensionFunc)
	extensionType := extensionValue.Type()

	if extensionType.NumIn() != 1 {
		return errors.New("Cannot add function(" + extensionType.String() + ") as parser: To many arguments!")
	}
	if extensionType.NumOut() != 2 {
		return errors.New("Cannot add function(" + extensionType.String() + ") as parser: Wrong number of return values!")
	}
	if extensionType.Out(1).Kind() != reflect.Bool {
		return errors.New("Cannot add function(" + extensionType.String() + ") as parser: Second return value is not Bool!")
	}

	router.extensions[extensionType.Out(0)] = extensionValue

	return nil
}

// SetConnectionExtension sets the extension for this router. A connection extension is an extended connection structure, that afterwards can be
// used by On-handlers as well (use-cases: add persistent storage to connection, additional methods et cetera).
// The SetConnectionExtension function takes the constructor of the custom format to be able to use and create it on
// connection to the router.
// For type E the constructor needs to fulfil the following requirements:
//     func NewE(conn *Connection) *E
// Afterwards On-handler can us this extended type:
//     router.On(func (extendedConn E, data Datatype) { ... })
// For an example have a look at the example repository and have a look at the 'example_connection_extension.go'.
func (router *Router) SetConnectionExtension(constructor interface{}) {
	router.connExtensionConstructor = reflect.ValueOf(constructor)
}

// SetProtocol sets the protocol of the router to the supplied implementation of the Protocol interface.
func (router *Router) SetProtocol(protocol Protocol) {
	router.protocol = protocol
}

//

// SetHeartbeat activates or deactivates the heartbeat depending on the flag parameter. By default heartbeats are activated.
func (router *Router) SetHeartbeat(flag bool) {
	router.useHeartbeats = flag
}
<|endoftext|>


<commit_before>package gocli

import (
	"fmt"
	"os"
	"regexp"
	"sort"
	"strings"
)

type Router struct {
	Actions   map[string]*Action
	Separator string
}

func NewRouter(mapping map[string]*Action) *Router {
	router := &Router{}
	for path, action := range mapping {
		router.Register(path, action)
	}
	return router
}

func (cli *Router) Register(path string, action *Action) {
	if cli.Actions == nil {
		cli.Actions = make(map[string]*Action)
	}
	cli.Actions[path] = action
}

func (router *Router) Search(patterns []string) map[string]*Action {
	re := regexp.MustCompile(strings.Join(patterns, ".*/"))
	actions := make(map[string]*Action)
	for key, action := range router.Actions {
		if re.MatchString(key) {
			actions[key] = action
		}
	}
	return actions
}

func (cli *Router) Usage() string {
	keys := []string{}
	for key := range cli.Actions {
		keys = append(keys, key)
	}
	return cli.UsageForKeys(keys, "")
}

func (cli *Router) UsageForKeys(keys []string, pattern string) string {
	sort.Strings(keys)
	table := NewTable()
	if cli.Separator != "" {
		table.Separator = cli.Separator
	}
	maxParts := 0
	selected := []string{}
	for _, key := range keys {
		partsCount := len(strings.Split(key, "/"))
		if partsCount > maxParts {
			maxParts = partsCount
		}
		selected = append(selected, key)
	}
	for _, key := range selected {
		parts := strings.Split(key, "/")
		action := cli.Actions[key]

		// fill up parts
		for i := (maxParts - len(parts)); i > 0; i-- {
			parts = append(parts, "")
		}

		parts = append(parts, action.Usage, action.Description)
		table.Add(parts...)
		if action.Args != nil {
			usage := action.Args.Usage()
			if usage != "" {
				lines := strings.Split(usage, "\n")
				for _, line := range lines {
					usageParts := []string{}
					for j := 0; j < 3; j++ {
						usageParts = append(usageParts, "")
					}
					current := append(usageParts, line)
					table.Add(current...)
				}
			}
		}
	}
	out := []string{"USAGE"}
	out = append(out, table.String())
	return strings.Join(out, "\n")
}

func AddActionUsage(parts []string, table *Table, action *Action) {
	parts = append(parts, action.Usage, action.Description)
	table.Add(parts...)
	if action.Args != nil {
		usage := action.Args.Usage()
		if usage != "" {
			lines := strings.Split(usage, "\n")
			for _, line := range lines {
				usageParts := []string{}
				for j := 0; j < 3; j++ {
					usageParts = append(usageParts, "")
				}
				current := append(usageParts, line)
				table.Add(current...)
			}
		}
	}
}

func (cli *Router) Handle(raw []string) error {
	for i := len(raw); i > 0; i-- {
		parts := raw[1:i]
		actions := cli.Search(parts)
		switch len(actions) {
		case 0:
			fmt.Println(cli.Usage())
			return nil
		case 1:
			var action *Action
			for _, a := range actions {
				action = a
			}
			args := action.Args
			if args == nil {
				args = &Args{}
			}
			e := args.Parse(raw[i:])
			if e == nil {
				e = action.Handler(args)
			}
			if e != nil {
				table := NewTable()
				AddActionUsage(parts, table, action)
				fmt.Println(table.String())
				os.Exit(0)
			}
			return nil
		default:
			keys := []string{}
			for key, _ := range actions {
				keys = append(keys, key)
			}
			fmt.Println(cli.UsageForKeys(keys, ""))
			return nil

		}
	}
	return nil
}
<commit_msg>use key from mapping for path<commit_after>package gocli

import (
	"fmt"
	"os"
	"regexp"
	"sort"
	"strings"
)

type Router struct {
	Actions   map[string]*Action
	Separator string
}

func NewRouter(mapping map[string]*Action) *Router {
	router := &Router{}
	for path, action := range mapping {
		router.Register(path, action)
	}
	return router
}

func (cli *Router) Register(path string, action *Action) {
	if cli.Actions == nil {
		cli.Actions = make(map[string]*Action)
	}
	cli.Actions[path] = action
}

func (router *Router) Search(patterns []string) map[string]*Action {
	re := regexp.MustCompile(strings.Join(patterns, ".*/"))
	actions := make(map[string]*Action)
	for key, action := range router.Actions {
		if re.MatchString(key) {
			actions[key] = action
		}
	}
	return actions
}

func (cli *Router) Usage() string {
	keys := []string{}
	for key := range cli.Actions {
		keys = append(keys, key)
	}
	return cli.UsageForKeys(keys, "")
}

func (cli *Router) UsageForKeys(keys []string, pattern string) string {
	sort.Strings(keys)
	table := NewTable()
	if cli.Separator != "" {
		table.Separator = cli.Separator
	}
	maxParts := 0
	selected := []string{}
	for _, key := range keys {
		partsCount := len(strings.Split(key, "/"))
		if partsCount > maxParts {
			maxParts = partsCount
		}
		selected = append(selected, key)
	}
	for _, key := range selected {
		parts := strings.Split(key, "/")
		action := cli.Actions[key]

		// fill up parts
		for i := (maxParts - len(parts)); i > 0; i-- {
			parts = append(parts, "")
		}

		parts = append(parts, action.Usage, action.Description)
		table.Add(parts...)
		if action.Args != nil {
			usage := action.Args.Usage()
			if usage != "" {
				lines := strings.Split(usage, "\n")
				for _, line := range lines {
					usageParts := []string{}
					for j := 0; j < 3; j++ {
						usageParts = append(usageParts, "")
					}
					current := append(usageParts, line)
					table.Add(current...)
				}
			}
		}
	}
	out := []string{"USAGE"}
	out = append(out, table.String())
	return strings.Join(out, "\n")
}

func AddActionUsage(parts []string, table *Table, action *Action) {
	parts = append(parts, action.Usage, action.Description)
	table.Add(parts...)
	if action.Args != nil {
		usage := action.Args.Usage()
		if usage != "" {
			lines := strings.Split(usage, "\n")
			for _, line := range lines {
				usageParts := []string{}
				for j := 0; j < 3; j++ {
					usageParts = append(usageParts, "")
				}
				current := append(usageParts, line)
				table.Add(current...)
			}
		}
	}
}

func (cli *Router) Handle(raw []string) error {
	for i := len(raw); i > 0; i-- {
		parts := raw[1:i]
		actions := cli.Search(parts)
		switch len(actions) {
		case 0:
			fmt.Println(cli.Usage())
			return nil
		case 1:
			var action *Action
			for k, a := range actions {
				parts = strings.Split(k, "/")
				action = a
			}
			args := action.Args
			if args == nil {
				args = &Args{}
			}
			e := args.Parse(raw[i:])
			if e == nil {
				e = action.Handler(args)
			}
			if e != nil {
				table := NewTable()
				AddActionUsage(parts, table, action)
				fmt.Println(table.String())
				os.Exit(0)
			}
			return nil
		default:
			keys := []string{}
			for key, _ := range actions {
				keys = append(keys, key)
			}
			fmt.Println(cli.UsageForKeys(keys, ""))
			return nil

		}
	}
	return nil
}
<|endoftext|>