<commit_before>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/dustin/go-humanize"

	"github.intel.com/hpdd/ce-tools/pkg/applog"
	"github.intel.com/hpdd/lustre"
	"github.intel.com/hpdd/lustre/fs"
	"github.intel.com/hpdd/lustre/hsm"
)

func init() {
	hsmStateFlags := strings.Join(hsm.GetStateFlagNames(), ",")

	hsmCommand := cli.Command{
		Name:  "hsm",
		Usage: "HSM-related data movement actions",
		Subcommands: []cli.Command{
			{
				Name:      "archive",
				Usage:     "Initiate HSM archive of specified paths",
				ArgsUsage: "[path [path...]]",
				Flags: []cli.Flag{
					cli.IntFlag{
						Name:  "id, i",
						Usage: "Numeric ID of archive backend",
					},
					cli.BoolFlag{
						Name:  "null, 0",
						Usage: "Null-separated paths are read from stdin (e.g. piped from find -print0)",
					},
				},
				Action: hsmAction,
			},
			{
				Name:      "release",
				Usage:     "Release local data of HSM-archived paths",
				ArgsUsage: "[path [path...]]",
				Flags: []cli.Flag{
					cli.BoolFlag{
						Name:  "null, 0",
						Usage: "Null-separated paths are read from stdin (e.g. piped from find -print0)",
					},
				},
				Action: hsmAction,
			},
			{
				Name:      "restore",
				Usage:     "Explicitly restore local data of HSM-archived paths",
				ArgsUsage: "[path [path...]]",
				Flags: []cli.Flag{
					cli.BoolFlag{
						Name:  "null, 0",
						Usage: "Null-separated paths are read from stdin (e.g. piped from find -print0)",
					},
				},
				Action: hsmAction,
			},
			{
				Name:      "remove",
				Usage:     "Remove HSM-archived data of specified paths (local data is not removed)",
				ArgsUsage: "[path [path...]]",
				Flags: []cli.Flag{
					cli.BoolFlag{
						Name:  "null, 0",
						Usage: "Null-separated paths are read from stdin (e.g. piped from find -print0)",
					},
				},
				Action: hsmAction,
			},
			{
				Name:      "cancel",
				Usage:     "Cancel HSM operations being performed on specified paths",
				ArgsUsage: "[path [path...]]",
				Flags: []cli.Flag{
					cli.BoolFlag{
						Name:  "null, 0",
						Usage: "Null-separated paths are read from stdin (e.g. piped from find -print0)",
					},
				},
				Action: hsmAction,
			},
			{
				Name:      "set",
				Usage:     "Set HSM flags or archive ID for specified paths",
				ArgsUsage: "[path [path...]]",
				Flags: []cli.Flag{
					cli.BoolFlag{
						Name:  "null, 0",
						Usage: "Null-separated paths are read from stdin (e.g. piped from find -print0)",
					},
					cli.IntFlag{
						Name:  "id, i",
						Usage: "Numeric ID of archive backend",
					},
					cli.StringSliceFlag{
						Name:  "flag, f",
						Usage: fmt.Sprintf("HSM flag to set (%s)", hsmStateFlags),
						Value: &cli.StringSlice{},
					},
					cli.StringSliceFlag{
						Name:  "clear, F",
						Usage: fmt.Sprintf("HSM flag to clear (%s)", hsmStateFlags),
						Value: &cli.StringSlice{},
					},
				},
				Action: hsmSetAction,
			},
			{
				Name:      "status",
				Usage:     "Display HSM status for specified paths",
				ArgsUsage: "[path [path...]]",
				Flags: []cli.Flag{
					cli.BoolFlag{
						Name:  "action, a",
						Usage: "Include current HSM action",
					},
					cli.BoolFlag{
						Name:  "hide-path, H",
						Usage: "Hide pathname in output",
					},
					cli.BoolFlag{
						Name:  "long, l",
						Usage: "Show long-form states",
					},
					cli.BoolFlag{
						Name:  "progress, p",
						Usage: "Show copy progress for archive/restore actions",
					},
					cli.BoolFlag{
						Name:  "null, 0",
						Usage: "Null-separated paths are read from stdin (e.g. piped from find -print0)",
					},
				},
				Action: hsmStatusAction,
			},
		},
	}
	commands = append(commands, hsmCommand)
}

func getFilePaths(c *cli.Context) ([]string, error) {
	var paths []string

	if c.Bool("null") {
		reader := bufio.NewReader(os.Stdin)
		path, err := reader.ReadBytes('\000')
		for err == nil {
			paths = append(paths, string(path[:len(path)-1]))
			path, err = reader.ReadBytes('\000')
		}
		if err != io.EOF {
			return nil, err
		}
	} else {
		paths = c.Args()
	}

	return paths, nil
}

func getPathStatus(c *cli.Context, filePath string) (string, error) {
	var buf bytes.Buffer

	s, err := hsm.GetFileStatus(filePath)
	if err != nil {
		return "", err
	}

	if !c.Bool("hide-path") {
		fmt.Fprintf(&buf, "%s ", filePath)
	}
	fmt.Fprintf(&buf, hsm.FileStatusString(s, !c.Bool("long")))

	if s.Exists() && c.Bool("action") {
		a, err := hsm.GetFileAction(filePath)
		if err != nil {
			return "", err
		}
		if a.IsNone() {
			fmt.Fprintf(&buf, " -")
		} else {
			fmt.Fprintf(&buf, " [%s:%s]", a.Action(), a.State())
			if c.Bool("progress") && (a.IsArchive() || a.IsRestore()) {
				st, err := os.Stat(filePath)
				if err != nil {
					applog.Fail(err)
				}
				fmt.Fprintf(&buf, "(%s/%s)",
					humanize.IBytes(a.BytesCopied),
					humanize.IBytes(uint64(st.Size())))
			}
		}
	} else {
		fmt.Fprintf(&buf, " -")
	}

	// TODO: Display xattrs, once we've standardized them?

	return buf.String(), nil
}

func hsmSetAction(c *cli.Context) {
	paths, err := getFilePaths(c)
	if err != nil {
		applog.Fail(err)
	}

	if len(paths) < 1 {
		applog.Fail(fmt.Errorf("HSM set request must be made with at least 1 path"))
	}

	setFlags, err := hsm.GetStatusMask(c.StringSlice("flag"))
	if err != nil {
		applog.Fail(err)
	}
	clearFlags, err := hsm.GetStatusMask(c.StringSlice("clear"))
	if err != nil {
		applog.Fail(err)
	}
	archiveID := uint32(c.Int("id"))

	if setFlags == 0 && clearFlags == 0 && archiveID == 0 {
		applog.Fail(fmt.Errorf("HSM set request made with no flags to set or clear, and no new archive ID supplied"))
	}

	// TODO: Parallelize this?
	for _, path := range paths {
		if err := hsm.SetFileStatus(path, setFlags, clearFlags, archiveID); err != nil {
			applog.Fail(err)
		}
	}
}

func hsmStatusAction(c *cli.Context) {
	paths, err := getFilePaths(c)
	if err != nil {
		applog.Fail(err)
	}

	if len(paths) < 1 {
		applog.Fail(fmt.Errorf("HSM status request must be made with at least 1 path"))
	}

	for _, path := range paths {
		status, err := getPathStatus(c, path)
		if err != nil {
			applog.Fail(err)
		}
		fmt.Println(status)
	}
}

func hsmAction(c *cli.Context) {
	paths, err := getFilePaths(c)
	if err != nil {
		applog.Fail(err)
	}

	if err := submitHsmRequest(c.Command.Name, uint(c.Int("id")), paths...); err != nil {
		applog.Fail(err)
	}
}

func submitHsmRequest(actionName string, archiveID uint, paths ...string) error {
	var fids []*lustre.Fid

	if len(paths) < 1 {
		return fmt.Errorf("HSM %s request must be made with at least 1 path", actionName)
	}

	fsID, err := fs.GetID(paths[0])
	if err != nil {
		return fmt.Errorf("Error getting fs ID from %s: %s", paths[0], err)
	}

	fsRoot, err := fs.MountRoot(paths[0])
	if err != nil {
		return fmt.Errorf("Error getting fs root from %s: %s", paths[0], err)
	}

	// TODO: Occurs to me that it might be better to break up a large
	// batch into multiple batches, each serviced by its own goroutine.
	for _, path := range paths {
		if !strings.HasPrefix(path, string(fsRoot)) {
			return fmt.Errorf("All files in HSM request must be in the same filesystem (%s is not in %s)", path, fsRoot)
		}

		fid, err := fs.LookupFid(path)
		if err != nil {
			return fmt.Errorf("Cannot resolve Fid for %s: %s", path, err)
		}
		fids = append(fids, fid)
	}

	switch actionName {
	case "archive":
		err = hsm.RequestArchive(fsID, archiveID, fids)
	case "release":
		err = hsm.RequestRelease(fsID, archiveID, fids)
	case "restore":
		err = hsm.RequestRestore(fsID, archiveID, fids)
	case "remove":
		err = hsm.RequestRemove(fsID, archiveID, fids)
	case "cancel":
		err = hsm.RequestCancel(fsID, archiveID, fids)
	default:
		err = fmt.Errorf("Unhandled HSM action: %s", actionName)
	}

	return err
}
<commit_msg>Support relative path names<commit_after>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/dustin/go-humanize"

	"github.intel.com/hpdd/ce-tools/pkg/applog"
	"github.intel.com/hpdd/lustre"
	"github.intel.com/hpdd/lustre/fs"
	"github.intel.com/hpdd/lustre/hsm"
)

func init() {
	hsmStateFlags := strings.Join(hsm.GetStateFlagNames(), ",")

	hsmCommand := cli.Command{
		Name:  "hsm",
		Usage: "HSM-related data movement actions",
		Subcommands: []cli.Command{
			{
				Name:      "archive",
				Usage:     "Initiate HSM archive of specified paths",
				ArgsUsage: "[path [path...]]",
				Flags: []cli.Flag{
					cli.IntFlag{
						Name:  "id, i",
						Usage: "Numeric ID of archive backend",
					},
					cli.BoolFlag{
						Name:  "null, 0",
						Usage: "Null-separated paths are read from stdin (e.g. piped from find -print0)",
					},
				},
				Action: hsmAction,
			},
			{
				Name:      "release",
				Usage:     "Release local data of HSM-archived paths",
				ArgsUsage: "[path [path...]]",
				Flags: []cli.Flag{
					cli.BoolFlag{
						Name:  "null, 0",
						Usage: "Null-separated paths are read from stdin (e.g. piped from find -print0)",
					},
				},
				Action: hsmAction,
			},
			{
				Name:      "restore",
				Usage:     "Explicitly restore local data of HSM-archived paths",
				ArgsUsage: "[path [path...]]",
				Flags: []cli.Flag{
					cli.BoolFlag{
						Name:  "null, 0",
						Usage: "Null-separated paths are read from stdin (e.g. piped from find -print0)",
					},
				},
				Action: hsmAction,
			},
			{
				Name:      "remove",
				Usage:     "Remove HSM-archived data of specified paths (local data is not removed)",
				ArgsUsage: "[path [path...]]",
				Flags: []cli.Flag{
					cli.BoolFlag{
						Name:  "null, 0",
						Usage: "Null-separated paths are read from stdin (e.g. piped from find -print0)",
					},
				},
				Action: hsmAction,
			},
			{
				Name:      "cancel",
				Usage:     "Cancel HSM operations being performed on specified paths",
				ArgsUsage: "[path [path...]]",
				Flags: []cli.Flag{
					cli.BoolFlag{
						Name:  "null, 0",
						Usage: "Null-separated paths are read from stdin (e.g. piped from find -print0)",
					},
				},
				Action: hsmAction,
			},
			{
				Name:      "set",
				Usage:     "Set HSM flags or archive ID for specified paths",
				ArgsUsage: "[path [path...]]",
				Flags: []cli.Flag{
					cli.BoolFlag{
						Name:  "null, 0",
						Usage: "Null-separated paths are read from stdin (e.g. piped from find -print0)",
					},
					cli.IntFlag{
						Name:  "id, i",
						Usage: "Numeric ID of archive backend",
					},
					cli.StringSliceFlag{
						Name:  "flag, f",
						Usage: fmt.Sprintf("HSM flag to set (%s)", hsmStateFlags),
						Value: &cli.StringSlice{},
					},
					cli.StringSliceFlag{
						Name:  "clear, F",
						Usage: fmt.Sprintf("HSM flag to clear (%s)", hsmStateFlags),
						Value: &cli.StringSlice{},
					},
				},
				Action: hsmSetAction,
			},
			{
				Name:      "status",
				Usage:     "Display HSM status for specified paths",
				ArgsUsage: "[path [path...]]",
				Flags: []cli.Flag{
					cli.BoolFlag{
						Name:  "action, a",
						Usage: "Include current HSM action",
					},
					cli.BoolFlag{
						Name:  "hide-path, H",
						Usage: "Hide pathname in output",
					},
					cli.BoolFlag{
						Name:  "long, l",
						Usage: "Show long-form states",
					},
					cli.BoolFlag{
						Name:  "progress, p",
						Usage: "Show copy progress for archive/restore actions",
					},
					cli.BoolFlag{
						Name:  "null, 0",
						Usage: "Null-separated paths are read from stdin (e.g. piped from find -print0)",
					},
				},
				Action: hsmStatusAction,
			},
		},
	}
	commands = append(commands, hsmCommand)
}

func getFilePaths(c *cli.Context) ([]string, error) {
	var paths []string

	if c.Bool("null") {
		reader := bufio.NewReader(os.Stdin)
		path, err := reader.ReadBytes('\000')
		for err == nil {
			paths = append(paths, string(path[:len(path)-1]))
			path, err = reader.ReadBytes('\000')
		}
		if err != io.EOF {
			return nil, err
		}
	} else {
		paths = c.Args()
	}

	return paths, nil
}

func getPathStatus(c *cli.Context, filePath string) (string, error) {
	var buf bytes.Buffer

	s, err := hsm.GetFileStatus(filePath)
	if err != nil {
		return "", err
	}

	if !c.Bool("hide-path") {
		fmt.Fprintf(&buf, "%s ", filePath)
	}
	fmt.Fprintf(&buf, hsm.FileStatusString(s, !c.Bool("long")))

	if s.Exists() && c.Bool("action") {
		a, err := hsm.GetFileAction(filePath)
		if err != nil {
			return "", err
		}
		if a.IsNone() {
			fmt.Fprintf(&buf, " -")
		} else {
			fmt.Fprintf(&buf, " [%s:%s]", a.Action(), a.State())
			if c.Bool("progress") && (a.IsArchive() || a.IsRestore()) {
				st, err := os.Stat(filePath)
				if err != nil {
					applog.Fail(err)
				}
				fmt.Fprintf(&buf, "(%s/%s)",
					humanize.IBytes(a.BytesCopied),
					humanize.IBytes(uint64(st.Size())))
			}
		}
	} else {
		fmt.Fprintf(&buf, " -")
	}

	// TODO: Display xattrs, once we've standardized them?

	return buf.String(), nil
}

func hsmSetAction(c *cli.Context) {
	paths, err := getFilePaths(c)
	if err != nil {
		applog.Fail(err)
	}

	if len(paths) < 1 {
		applog.Fail(fmt.Errorf("HSM set request must be made with at least 1 path"))
	}

	setFlags, err := hsm.GetStatusMask(c.StringSlice("flag"))
	if err != nil {
		applog.Fail(err)
	}
	clearFlags, err := hsm.GetStatusMask(c.StringSlice("clear"))
	if err != nil {
		applog.Fail(err)
	}
	archiveID := uint32(c.Int("id"))

	if setFlags == 0 && clearFlags == 0 && archiveID == 0 {
		applog.Fail(fmt.Errorf("HSM set request made with no flags to set or clear, and no new archive ID supplied"))
	}

	// TODO: Parallelize this?
	for _, path := range paths {
		if err := hsm.SetFileStatus(path, setFlags, clearFlags, archiveID); err != nil {
			applog.Fail(err)
		}
	}
}

func hsmStatusAction(c *cli.Context) {
	paths, err := getFilePaths(c)
	if err != nil {
		applog.Fail(err)
	}

	if len(paths) < 1 {
		applog.Fail(fmt.Errorf("HSM status request must be made with at least 1 path"))
	}

	for _, path := range paths {
		status, err := getPathStatus(c, path)
		if err != nil {
			applog.Fail(err)
		}
		fmt.Println(status)
	}
}

func hsmAction(c *cli.Context) {
	paths, err := getFilePaths(c)
	if err != nil {
		applog.Fail(err)
	}

	if err := submitHsmRequest(c.Command.Name, uint(c.Int("id")), paths...); err != nil {
		applog.Fail(err)
	}
}

func submitHsmRequest(actionName string, archiveID uint, paths ...string) error {
	var fids []*lustre.Fid

	if len(paths) < 1 {
		return fmt.Errorf("HSM %s request must be made with at least 1 path", actionName)
	}

	fsID, err := fs.GetID(paths[0])
	if err != nil {
		return fmt.Errorf("Error getting fs ID from %s: %s", paths[0], err)
	}

	fsRoot, err := fs.MountRoot(paths[0])
	if err != nil {
		return fmt.Errorf("Error getting fs root from %s: %s", paths[0], err)
	}

	// TODO: Occurs to me that it might be better to break up a large
	// batch into multiple batches, each serviced by its own goroutine.
	for _, path := range paths {
		absPath, err := filepath.Abs(path)
		if !strings.HasPrefix(absPath, string(fsRoot)) {
			return fmt.Errorf("All files in HSM request must be in the same filesystem (%s is not in %s)", path, fsRoot)
		}

		fid, err := fs.LookupFid(path)
		if err != nil {
			return fmt.Errorf("Cannot resolve Fid for %s: %s", path, err)
		}
		fids = append(fids, fid)
	}

	switch actionName {
	case "archive":
		err = hsm.RequestArchive(fsID, archiveID, fids)
	case "release":
		err = hsm.RequestRelease(fsID, archiveID, fids)
	case "restore":
		err = hsm.RequestRestore(fsID, archiveID, fids)
	case "remove":
		err = hsm.RequestRemove(fsID, archiveID, fids)
	case "cancel":
		err = hsm.RequestCancel(fsID, archiveID, fids)
	default:
		err = fmt.Errorf("Unhandled HSM action: %s", actionName)
	}

	return err
}
<|endoftext|>


<commit_before>package cart

import (
	"sync"
	"net/http"
	"os"
	"time"
	"strings"
	"html/template"
	"github.com/gimke/cart/render"
)

type Engine struct {
	Router
	delims          render.Delims
	routers			map[string]*Router	//saved routers
	engine			*Engine
	pool        	sync.Pool
	tree			*node 	//match trees

	NotFound		HandlerFinal

	FuncMap         template.FuncMap
	Template 		*template.Template

	ForwardedByClientIP		bool
	AppEngine				bool
}

var _ http.Handler = &Engine{}

var server *http.Server

func (e *Engine) allocateContext() *Context {
	return &Context{}
}

func (e *Engine) findRouter(absolutePath string) (*Router, bool) {
	router := e.routers[absolutePath]
	if router == nil {
		return nil, false
	}
	return router, true
}

func (e *Engine) getRouter(absolutePath string) (*Router, bool) {
	router := e.routers[absolutePath]
	find := true
	if router == nil {
		find = false
		router = &Router{
			engine:e,
			Path:absolutePath,
			methods:make([]method,0),
		}
	}
	return router, find
}

func (e *Engine) addRoute(router *Router) {
	if router.Path[0] != '/' {
		panic("Path must begin with '/' in path '" + router.Path + "'")
	}
	if e.tree == nil {
		e.tree = &node{}
	}
	if _, found := e.tree.findCaseInsensitivePath(router.Path, true); !found {
		debugPrint("Add Router %s",router.Path)
		e.routers[router.Path] = router
		e.tree.addRoute(router.Path, router)
	}
}

func (e *Engine) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	c := e.pool.Get().(*Context)
	c.reset(w, req)
	e.serveHTTP(c)
	e.pool.Put(c)
}

func (e *Engine) mixMethods(httpMethod string, r *Router) HandlerCompose {
	//http method find any
	var methods HandlerCompose
	if m, find := r.getMethod("ANY"); find {
		methods = compose(m)
	}
	if m, find := r.getMethod(httpMethod); find {
		if methods != nil {
			methods = compose(methods,m)
		} else {
			methods = compose(m)
		}
	}
	return methods
}

func (e *Engine) serveHTTP(c *Context) {
	path := c.Request.URL.Path
	httpMethod := c.Request.Method

	final404 := func() {
		// 404 error
		// make temp router
		c.Router,_ = e.getRouter(path)
		if c.Response.Size() == -1 && c.Response.Status() == 200 {
			if e.NotFound != nil {
				e.NotFound(c);
			} else {
				c.ErrorHTML(404,
					"404 Not Found",
					"The page "+path+" is not found")
				//c.String(404,"404 Not Found")
			}
		}
	}

	if root := e.tree; root != nil {
		if r, ps, tsr := root.getValue(path); r != nil {
			router := r.(*Router)
			c.Router = router
			c.Params = ps

			//methods
			methods := e.mixMethods(httpMethod, router)
			//middleware

			composed := router.composed
			if composed != nil && methods != nil {
				composed = compose(composed, methods)
			} else if composed == nil && methods != nil {
				composed = methods
			}
			if composed != nil {
				composed(c,final404)()
			} else {
				final404()
			}
			c.Response.WriteHeaderNow()
			return
		} else if httpMethod != "CONNECT" && path != "/" {
			code := 301 // Permanent redirect, request with GET method
			if httpMethod != "GET" {
				code = 307
			}
			if tsr {
				if len(path) > 1 && path[len(path)-1] == '/' {
					c.Request.URL.Path = path[:len(path)-1]
				} else {
					c.Request.URL.Path = path + "/"
				}
				http.Redirect(c.Response, c.Request, c.Request.URL.String(), code)
				return
			}
		}
	}
	//find / middleware
	r, composed := e.mixComposed(path)
	if composed != nil {
		c.Router = r
		composed(c,final404)()
	} else {
		final404()
	}
	c.Response.WriteHeaderNow()
}

func (e *Engine) mixComposed(absolutePath string) (*Router, HandlerCompose) {
	sp := strings.Split(absolutePath,"/")
	for i, _ := range sp {
		//find it's self first ..... last is root path / router
		tempPath := strings.Join(sp[0:len(sp)-i], "/")
		if tempPath == "" {
			tempPath = "/"
		}
		if pr, find := e.findRouter(tempPath); find {
			return pr, pr.composed
		}
		//auto add slash then find
		if tempPath[len(tempPath)-1] != '/' && tempPath != absolutePath && tempPath != "/" {
			tempPath = tempPath+"/"
			if pr, find := e.findRouter(tempPath); find {
				return pr, pr.composed
			}
		}

	}
	return nil, nil
}
/*
init new Engine
 */
func (e *Engine) init() {
	e.Router = Router{
		Path: "/",
	}
	e.Router.engine = e
	e.pool.New = func() interface{} {
		return e.allocateContext()
	}
	e.tree = &node{}
	e.routers = make(map[string]*Router)
}
/*
Run the server
 */
func (e *Engine) Run(addr ...string) (err error) {
	defer func() { debugError(err) }()
	address := resolveAddress(addr)
	debugPrint("PID:%d Listening and serving HTTP on %s\n", os.Getpid(), address)

	server = &http.Server{
		Addr: address,
		Handler: e,
		ReadTimeout: time.Second * 90,
		//ReadHeaderTimeout: time.Second * 90,
		WriteTimeout: time.Second * 90,
		//IdleTimeout: time.Second * 90,
	}
	err = server.ListenAndServe()
	return
}

func (engine *Engine) LoadHTMLGlob(pattern string) {

	templ := template.Must(template.New("").Delims(engine.delims.Left, engine.delims.Right).Funcs(engine.FuncMap).ParseGlob(pattern))
	engine.SetHTMLTemplate(templ)

}

func (engine *Engine) SetHTMLTemplate(templ *template.Template) {
	engine.Template = templ.Funcs(engine.FuncMap)
}

func (engine *Engine) SetFuncMap(funcMap template.FuncMap) {
	engine.FuncMap = funcMap
}<commit_msg>add mem session<commit_after>package cart

import (
	"sync"
	"net/http"
	"os"
	"time"
	"strings"
	"html/template"
	"github.com/gimke/cart/render"
)

type Engine struct {
	Router
	delims          render.Delims
	routers			map[string]*Router	//saved routers
	engine			*Engine
	pool        	sync.Pool
	tree			*node 	//match trees

	NotFound		HandlerFinal

	FuncMap         template.FuncMap
	Template 		*template.Template

	ForwardedByClientIP		bool
	AppEngine				bool
}

var _ http.Handler = &Engine{}

var server *http.Server

func (e *Engine) allocateContext() *Context {
	return &Context{}
}

func (e *Engine) findRouter(absolutePath string) (*Router, bool) {
	router := e.routers[absolutePath]
	if router == nil {
		return nil, false
	}
	return router, true
}

func (e *Engine) getRouter(absolutePath string) (*Router, bool) {
	router := e.routers[absolutePath]
	find := true
	if router == nil {
		find = false
		router = &Router{
			engine:e,
			Path:absolutePath,
			methods:make([]method,0),
		}
	}
	return router, find
}

func (e *Engine) addRoute(router *Router) {
	if router.Path[0] != '/' {
		panic("Path must begin with '/' in path '" + router.Path + "'")
	}
	if e.tree == nil {
		e.tree = &node{}
	}
	if _, found := e.tree.findCaseInsensitivePath(router.Path, true); !found {
		debugPrint("Add Router %s",router.Path)
		e.routers[router.Path] = router
		e.tree.addRoute(router.Path, router)
	}
}

func (e *Engine) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	c := e.pool.Get().(*Context)
	c.reset(w, req)
	e.serveHTTP(c)
	e.pool.Put(c)
}

func (e *Engine) mixMethods(httpMethod string, r *Router) HandlerCompose {
	//http method find any
	var methods HandlerCompose
	if m, find := r.getMethod("ANY"); find {
		methods = compose(m)
	}
	if m, find := r.getMethod(httpMethod); find {
		if methods != nil {
			methods = compose(methods,m)
		} else {
			methods = compose(m)
		}
	}
	return methods
}

func (e *Engine) serveHTTP(c *Context) {
	path := c.Request.URL.Path
	httpMethod := c.Request.Method

	final404 := func() {
		// 404 error
		// make temp router
		c.Router,_ = e.getRouter(path)
		if c.Response.Size() == -1 && c.Response.Status() == 200 {
			if e.NotFound != nil {
				e.NotFound(c);
			} else {
				c.ErrorHTML(404,
					"404 Not Found",
					"The page "+path+" is not found")
				//c.String(404,"404 Not Found")
			}
		}
	}

	if root := e.tree; root != nil {
		if r, ps, tsr := root.getValue(path); r != nil {
			router := r.(*Router)
			c.Router = router
			c.Params = ps

			//methods
			methods := e.mixMethods(httpMethod, router)
			//middleware

			composed := router.composed
			if composed != nil && methods != nil {
				composed = compose(composed, methods)
			} else if composed == nil && methods != nil {
				composed = methods
			}
			if composed != nil {
				composed(c,final404)()
			} else {
				final404()
			}
			c.Response.WriteHeaderNow()
			return
		} else if httpMethod != "CONNECT" && path != "/" {
			code := 301 // Permanent redirect, request with GET method
			if httpMethod != "GET" {
				code = 307
			}
			if tsr {
				if len(path) > 1 && path[len(path)-1] == '/' {
					c.Request.URL.Path = path[:len(path)-1]
				} else {
					c.Request.URL.Path = path + "/"
				}
				http.Redirect(c.Response, c.Request, c.Request.URL.String(), code)
				return
			}
		}
	}
	//find / middleware
	r, composed := e.mixComposed(path)
	if composed != nil {
		c.Router = r
		composed(c,final404)()
	} else {
		final404()
	}
	c.Response.WriteHeaderNow()
}

func (e *Engine) mixComposed(absolutePath string) (*Router, HandlerCompose) {
	sp := strings.Split(absolutePath,"/")
	for i, _ := range sp {
		//find it's self first ..... last is root path / router
		tempPath := strings.Join(sp[0:len(sp)-i], "/")
		if tempPath == "" {
			tempPath = "/"
		}
		if pr, find := e.findRouter(tempPath); find {
			return pr, pr.composed
		}
		//auto add slash then find
		if tempPath[len(tempPath)-1] != '/' && tempPath != absolutePath && tempPath != "/" {
			tempPath = tempPath+"/"
			if pr, find := e.findRouter(tempPath); find {
				return pr, pr.composed
			}
		}

	}
	return nil, nil
}
/*
init new Engine
 */
func (e *Engine) init() {
	e.Router = Router{
		Path: "/",
	}
	e.Router.engine = e
	e.pool.New = func() interface{} {
		return e.allocateContext()
	}
	e.tree = &node{}
	e.routers = make(map[string]*Router)
}
/*
Run the server
 */
func (e *Engine) Run(addr ...string) (server *http.Server, err error) {
	defer func() { debugError(err) }()
	address := resolveAddress(addr)
	debugPrint("PID:%d Listening and serving HTTP on %s\n", os.Getpid(), address)

	server = &http.Server{
		Addr: address,
		Handler: e,
		ReadTimeout: time.Second * 90,
		//ReadHeaderTimeout: time.Second * 90,
		WriteTimeout: time.Second * 90,
		//IdleTimeout: time.Second * 90,
	}
	err = server.ListenAndServe()
	return
}

/*
RunTLS
 */
func (e *Engine) RunTLS(addr string, certFile string, keyFile string) (server *http.Server,err error) {
	defer func() { debugError(err) }()
	debugPrint("PID:%d Listening and serving HTTPS on %s\n", os.Getpid(), addr)
	server = &http.Server{
		Addr: addr,
		Handler: e,
		ReadTimeout: time.Second * 90,
		//ReadHeaderTimeout: time.Second * 90,
		WriteTimeout: time.Second * 90,
		//IdleTimeout: time.Second * 90,
	}
	err = server.ListenAndServeTLS(certFile,keyFile)
	return
}

func (engine *Engine) LoadHTMLGlob(pattern string) {

	templ := template.Must(template.New("").Delims(engine.delims.Left, engine.delims.Right).Funcs(engine.FuncMap).ParseGlob(pattern))
	engine.SetHTMLTemplate(templ)

}

func (engine *Engine) SetHTMLTemplate(templ *template.Template) {
	engine.Template = templ.Funcs(engine.FuncMap)
}

func (engine *Engine) SetFuncMap(funcMap template.FuncMap) {
	engine.FuncMap = funcMap
}<|endoftext|>


<commit_before>// This is the main package for the `packer` application.
package main

import (
	"bytes"
	"fmt"
	"github.com/mitchellh/packer/packer"
	"github.com/mitchellh/packer/packer/plugin"
	"io/ioutil"
	"log"
	"os"
	"runtime"
)

func main() {
	if os.Getenv("PACKER_LOG") == "" {
		// If we don't have logging explicitly enabled, then disable it
		log.SetOutput(ioutil.Discard)
	} else {
		// Logging is enabled, make sure it goes to stderr
		log.SetOutput(os.Stderr)
	}

	// If there is no explicit number of Go threads to use, then set it
	if os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
	}

	config, err := loadConfig()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error loading configuration: \n\n%s\n", err)
		os.Exit(1)
	}

	log.Printf("Packer config: %+v", config)

	defer plugin.CleanupClients()

	cacheDir := os.Getenv("PACKER_CACHE_DIR")
	if cacheDir == "" {
		cacheDir = "packer_cache"
	}

	if err := os.MkdirAll(cacheDir, 0755); err != nil {
		fmt.Fprintf(os.Stderr, "Error preparing cache directory: \n\n%s\n", err)
		os.Exit(1)
	}

	log.Printf("Setting cache directory: %s", cacheDir)
	cache := &packer.FileCache{CacheDir: cacheDir}

	envConfig := packer.DefaultEnvironmentConfig()
	envConfig.Cache = cache
	envConfig.Commands = config.CommandNames()
	envConfig.Components.Builder = config.LoadBuilder
	envConfig.Components.Command = config.LoadCommand
	envConfig.Components.Hook = config.LoadHook
	envConfig.Components.PostProcessor = config.LoadPostProcessor
	envConfig.Components.Provisioner = config.LoadProvisioner

	env, err := packer.NewEnvironment(envConfig)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Packer initialization error: \n\n%s\n", err)
		os.Exit(1)
	}

	setupSignalHandlers(env)

	exitCode, err := env.Cli(os.Args[1:])
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error executing CLI: %s\n", err.Error())
		os.Exit(1)
	}

	plugin.CleanupClients()
	os.Exit(exitCode)
}

func loadConfig() (*config, error) {
	var config config
	if err := decodeConfig(bytes.NewBufferString(defaultConfig), &config); err != nil {
		return nil, err
	}

	mustExist := true
	configFilePath := os.Getenv("PACKER_CONFIG")
	if configFilePath == "" {
		var err error
		configFilePath, err = configFile()
		mustExist = false

		if err != nil {
			log.Printf("Error detecing default config file path: %s", err)
		}
	}

	if configFilePath == "" {
		return &config, nil
	}

	log.Printf("Attempting to open config file: %s", configFilePath)
	f, err := os.Open(configFilePath)
	if err != nil {
		if !os.IsNotExist(err) {
			return nil, err
		}

		if mustExist {
			return nil, err
		}

		log.Println("File doesn't exist, but doesn't need to. Ignoring.")
		return &config, nil
	}
	defer f.Close()

	if err := decodeConfig(f, &config); err != nil {
		return nil, err
	}

	return &config, nil
}
<commit_msg>Make sure the cache dir is absolute<commit_after>// This is the main package for the `packer` application.
package main

import (
	"bytes"
	"fmt"
	"github.com/mitchellh/packer/packer"
	"github.com/mitchellh/packer/packer/plugin"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"runtime"
)

func main() {
	if os.Getenv("PACKER_LOG") == "" {
		// If we don't have logging explicitly enabled, then disable it
		log.SetOutput(ioutil.Discard)
	} else {
		// Logging is enabled, make sure it goes to stderr
		log.SetOutput(os.Stderr)
	}

	// If there is no explicit number of Go threads to use, then set it
	if os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
	}

	config, err := loadConfig()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error loading configuration: \n\n%s\n", err)
		os.Exit(1)
	}

	log.Printf("Packer config: %+v", config)

	cacheDir := os.Getenv("PACKER_CACHE_DIR")
	if cacheDir == "" {
		cacheDir = "packer_cache"
	}

	cacheDir, err = filepath.Abs(cacheDir)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error preparing cache directory: \n\n%s\n", err)
		os.Exit(1)
	}

	if err := os.MkdirAll(cacheDir, 0755); err != nil {
		fmt.Fprintf(os.Stderr, "Error preparing cache directory: \n\n%s\n", err)
		os.Exit(1)
	}

	log.Printf("Setting cache directory: %s", cacheDir)
	cache := &packer.FileCache{CacheDir: cacheDir}

	defer plugin.CleanupClients()

	envConfig := packer.DefaultEnvironmentConfig()
	envConfig.Cache = cache
	envConfig.Commands = config.CommandNames()
	envConfig.Components.Builder = config.LoadBuilder
	envConfig.Components.Command = config.LoadCommand
	envConfig.Components.Hook = config.LoadHook
	envConfig.Components.PostProcessor = config.LoadPostProcessor
	envConfig.Components.Provisioner = config.LoadProvisioner

	env, err := packer.NewEnvironment(envConfig)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Packer initialization error: \n\n%s\n", err)
		os.Exit(1)
	}

	setupSignalHandlers(env)

	exitCode, err := env.Cli(os.Args[1:])
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error executing CLI: %s\n", err.Error())
		os.Exit(1)
	}

	plugin.CleanupClients()
	os.Exit(exitCode)
}

func loadConfig() (*config, error) {
	var config config
	if err := decodeConfig(bytes.NewBufferString(defaultConfig), &config); err != nil {
		return nil, err
	}

	mustExist := true
	configFilePath := os.Getenv("PACKER_CONFIG")
	if configFilePath == "" {
		var err error
		configFilePath, err = configFile()
		mustExist = false

		if err != nil {
			log.Printf("Error detecing default config file path: %s", err)
		}
	}

	if configFilePath == "" {
		return &config, nil
	}

	log.Printf("Attempting to open config file: %s", configFilePath)
	f, err := os.Open(configFilePath)
	if err != nil {
		if !os.IsNotExist(err) {
			return nil, err
		}

		if mustExist {
			return nil, err
		}

		log.Println("File doesn't exist, but doesn't need to. Ignoring.")
		return &config, nil
	}
	defer f.Close()

	if err := decodeConfig(f, &config); err != nil {
		return nil, err
	}

	return &config, nil
}
<|endoftext|>


<commit_before>package registry

import (
	"fmt"
	"time"
	"errors"
	"net/http"
	"database/sql"

	"github.com/apex/log"
	"github.com/satori/go.uuid"
	"github.com/lib/pq"
	"github.com/spf13/viper"
	"gopkg.in/gin-gonic/gin.v1"
	"github.com/facebookgo/grace/gracehttp"
)

var ctx = log.WithFields(log.Fields{
	"cmd": "registry",
})

func initDatabase() (*sql.DB, error) {
	db, err := sql.Open("postgres", viper.GetString("database.url"))
	if err != nil {
		ctx.Error("failed to open database")
		return nil, err
	}
	return db, err
}

type UpdateReq struct {
	ProbeCC string `json:"probe_cc"`
	ProbeASN string `json:"probe_asn"`
	Platform string `json:"platform"`

	SoftwareName string `json:"software_name"`
	SoftwareVersion string `json:"software_version"`
	SupportedTests []string `json:"supported_tests"`

	NetworkType string `json:"network_type"`
	AvailableBandwidth string `json:"available_bandwidth"`

	Token string `json:"token"`

	ProbeFamily string `json:"probe_family"`
	ProbeID string `json:"probe_id"`
}

func IsClientRegistered(db *sql.DB, clientID string) (bool, error) {
	var found string
	query := fmt.Sprintf(`SELECT id FROM %s WHERE id = $1`,
				pq.QuoteIdentifier(viper.GetString("database.active-probes-table")))
	err := db.QueryRow(query, clientID).Scan(&found)
	if err != nil {
		if err == sql.ErrNoRows {
			return false, nil
		}
		return false, err
	}
	return true, nil
}

func Update(db *sql.DB, clientID string, req UpdateReq) (error) {
	tx, err := db.Begin()
	if err != nil {
		ctx.WithError(err).Error("failed to open transaction")
		return err
	}

	// Write into the updates table
	{
		query := fmt.Sprintf(`INSERT INTO %s (
			id, update_time,
			client_id,
			probe_cc, probe_asn,
			platform, software_name,
			software_version, supported_tests,
			network_type, available_bandwidth,
			token, probe_family,
			probe_id, update_type
		) VALUES (
			$1, $2,
			$3, $4,
			$5, $6,
			$7, $8,
			$9, $10,
			$11, $12,
			$13, $14, $15)`,
			pq.QuoteIdentifier(viper.GetString("database.probe-updates-table")))

		stmt, err := tx.Prepare(query)
		if (err != nil) {
			ctx.WithError(err).Error("failed to prepare update probes query")
			return err
		}

		updateID := uuid.NewV4().String()
		_, err = stmt.Exec(updateID, time.Now().UTC(),
							clientID,
							req.ProbeCC, req.ProbeASN,
							req.Platform, req.SoftwareName,
							req.SoftwareVersion, pq.Array(req.SupportedTests),
							req.NetworkType, req.AvailableBandwidth,
							req.Token, req.ProbeFamily,
							req.ProbeID, "register")
		if (err != nil) {
			ctx.WithError(err).Error("failed to add data to update table, rolling back")
			tx.Rollback()
			return errors.New("error in adding data to update probes")
		}
	}

	// Write into the active probes table
	{
		query := fmt.Sprintf(`UPDATE %s SET
			last_updated = $2,
			probe_cc = $3,
			probe_asn = $4,
			platform = $5,
			software_name = $6,
			software_version = $7,
			supported_tests = $8,
			network_type = $9,
			available_bandwidth = $10,
			token = $11,
			probe_family = $12,
			probe_id = $13
			WHERE id = $1`,
			pq.QuoteIdentifier(viper.GetString("database.active-probes-table")))

		stmt, err := tx.Prepare(query)
		if (err != nil) {
			ctx.WithError(err).Error("failed to prepare update probes query")
			return err
		}
		_, err = stmt.Exec(clientID,
							time.Now().UTC(),
							req.ProbeCC,
							req.ProbeASN,
							req.Platform,
							req.SoftwareName,
							req.SoftwareVersion,
							pq.Array(req.SupportedTests),
							req.NetworkType,
							req.AvailableBandwidth,
							req.Token,
							req.ProbeFamily,
							req.ProbeID)
		if (err != nil) {
			ctx.WithError(err).Error("failed to update active table, rolling back")
			tx.Rollback()
			return errors.New("failed to update active table")
		}
	}

	if err := tx.Commit(); err != nil {
		ctx.WithError(err).Error("failed to commit transaction, rolling back")
		return err
	}

	return nil
}

type RegisterReq struct {
	ProbeCC string `json:"probe_cc" binding:"required"`
	ProbeASN string `json:"probe_asn" binding:"required"`
	Platform string `json:"platform" binding:"required"`

	SoftwareName string `json:"software_name" binding:"required"`
	SoftwareVersion string `json:"software_version" binding:"required"`
	SupportedTests []string `json:"supported_tests"`

	NetworkType string `json:"network_type"`
	AvailableBandwidth string `json:"available_bandwidth"`

	Token string `json:"token"`

	ProbeFamily string `json:"probe_family"`
	ProbeID string `json:"probe_id"`
}

func Register(db *sql.DB, req RegisterReq) (string, error) {
	if ((req.Platform == "ios" || req.Platform == "android") && req.Token == "") {
		return "", errors.New("missing device token")
	}

	tx, err := db.Begin()
	if err != nil {
		ctx.WithError(err).Error("failed to open transaction")
		return "", err
	}

	var clientID = uuid.NewV4().String()

	{
		query := fmt.Sprintf(`INSERT INTO %s (
			id, creation_time,
			last_updated,
			probe_cc, probe_asn,
			platform, software_name,
			software_version, supported_tests,
			network_type, available_bandwidth,
			token, probe_family,
			probe_id
		) VALUES (
			$1, $2,
			$3, $4,
			$5, $6,
			$7, $8,
			$9, $10,
			$11, $12,
			$13, $14)`,
			pq.QuoteIdentifier(viper.GetString("database.active-probes-table")))

		stmt, err := tx.Prepare(query)
		if (err != nil) {
			ctx.WithError(err).Error("failed to prepare active probes query")
			return "", err
		}
		defer stmt.Close()

		_, err = stmt.Exec(clientID, time.Now().UTC(),
							time.Now().UTC(),
							req.ProbeCC, req.ProbeASN,
							req.Platform, req.SoftwareName,
							req.SoftwareVersion, pq.Array(req.SupportedTests),
							req.NetworkType, req.AvailableBandwidth,
							req.Token, req.ProbeFamily,
							req.ProbeID)
		if err != nil {
			tx.Rollback()
			ctx.WithError(err).Error("failed to insert into active probes table, rolling back")
			return "", err
		}
	}

	{
		query := fmt.Sprintf(`INSERT INTO %s (
			id, update_time,
			client_id,
			probe_cc, probe_asn,
			platform, software_name,
			software_version, supported_tests,
			network_type, available_bandwidth,
			token, probe_family,
			probe_id, update_type
		) VALUES (
			$1, $2,
			$3, $4,
			$5, $6,
			$7, $8,
			$9, $10,
			$11, $12,
			$13, $14, $15)`,
			pq.QuoteIdentifier(viper.GetString("database.probe-updates-table")))

		stmt, err := tx.Prepare(query)
		if (err != nil) {
			ctx.WithError(err).Error("failed to prepare update probes query")
			return "", err
		}
		defer stmt.Close()

		updateID := uuid.NewV4().String()
		_, err = stmt.Exec(updateID, time.Now().UTC(),
							clientID,
							req.ProbeCC, req.ProbeASN,
							req.Platform, req.SoftwareName,
							req.SoftwareVersion, pq.Array(req.SupportedTests),
							req.NetworkType, req.AvailableBandwidth,
							req.Token, req.ProbeFamily,
							req.ProbeID, "register")
		if (err != nil) {
			ctx.WithError(err).Error("failed to add data to update table, rolling back")
			tx.Rollback()
			return "", errors.New("error in adding data to update probes")
		}
	}

	if err := tx.Commit(); err != nil {
		ctx.WithError(err).Error("failed to commit transaction, rolling back")
		return "", err
	}

	return clientID, nil
}

func Start() {
	db, err := initDatabase()

	if (err != nil) {
		ctx.WithError(err).Error("failed to connect to DB")
		return
	}
	defer db.Close()

	router := gin.Default()
	router.POST("/api/v1/clients", func(c *gin.Context) {
		var registerReq RegisterReq
		err := c.BindJSON(&registerReq)
		if (err != nil) {
			ctx.WithError(err).Error("invalid request")
			c.JSON(http.StatusBadRequest,
					gin.H{"error": "invalid request"})
			return
		}

		clientID , err := Register(db, registerReq)
		if (err != nil) {
			c.JSON(http.StatusBadRequest,
					gin.H{"error": err.Error()})
			return
		}

		c.JSON(http.StatusOK, gin.H{"client_id": clientID})
		return
	})

	// XXX do we also want to support a PATCH method?
	router.PUT("/api/v1/clients/:client_id", func(c *gin.Context) {
		var updateReq UpdateReq
		clientID := c.Param("client_id")
		err := c.BindJSON(&updateReq)
		if (err != nil) {
			ctx.WithError(err).Error("invalid request")
			c.JSON(http.StatusBadRequest,
					gin.H{"error": "invalid request"})
			return
		}
		isRegistered, err := IsClientRegistered(db, clientID)
		if (err != nil) {
			ctx.WithError(err).Error("failed to learn if client is registered")
			c.JSON(http.StatusBadRequest,
					gin.H{"error": err.Error()})
			return
		}
		if (isRegistered == false) {
			c.JSON(http.StatusNotFound,
					gin.H{"error": "client is not registered"})
			return
		}

		err = Update(db, clientID, updateReq)
		if (err != nil) {
			ctx.WithError(err).Error("failed to update")
			c.JSON(http.StatusBadRequest,
					gin.H{"error": err.Error()})
			return
		}
	})

	Addr := fmt.Sprintf("%s:%d", viper.GetString("api.address"),
								viper.GetInt("api.port"))
	ctx.Infof("starting on %s", Addr)
	s := &http.Server{
		Addr: Addr,
		Handler: router,
	}
	gracehttp.Serve(s)
}
<commit_msg>Add listing of clients<commit_after>package registry

import (
	"fmt"
	"time"
	"errors"
	"net/http"
	"database/sql"

	"github.com/apex/log"
	"github.com/satori/go.uuid"
	"github.com/lib/pq"
	"github.com/spf13/viper"
	"gopkg.in/gin-gonic/gin.v1"
	"github.com/facebookgo/grace/gracehttp"
)

var ctx = log.WithFields(log.Fields{
	"cmd": "registry",
})

func initDatabase() (*sql.DB, error) {
	db, err := sql.Open("postgres", viper.GetString("database.url"))
	if err != nil {
		ctx.Error("failed to open database")
		return nil, err
	}
	return db, err
}

type ClientData struct {
	ProbeCC string `json:"probe_cc" binding:"required"`
	ProbeASN string `json:"probe_asn" binding:"required"`
	Platform string `json:"platform" binding:"required"`

	SoftwareName string `json:"software_name" binding:"required"`
	SoftwareVersion string `json:"software_version" binding:"required"`
	SupportedTests []string `json:"supported_tests"`

	NetworkType string `json:"network_type"`
	AvailableBandwidth string `json:"available_bandwidth"`

	Token string `json:"token"`

	ProbeFamily string `json:"probe_family"`
	ProbeID string `json:"probe_id"`
}

func IsClientRegistered(db *sql.DB, clientID string) (bool, error) {
	var found string
	query := fmt.Sprintf(`SELECT id FROM %s WHERE id = $1`,
				pq.QuoteIdentifier(viper.GetString("database.active-probes-table")))
	err := db.QueryRow(query, clientID).Scan(&found)
	if err != nil {
		if err == sql.ErrNoRows {
			return false, nil
		}
		return false, err
	}
	return true, nil
}

func Update(db *sql.DB, clientID string, req ClientData) (error) {
	tx, err := db.Begin()
	if err != nil {
		ctx.WithError(err).Error("failed to open transaction")
		return err
	}

	// Write into the updates table
	{
		query := fmt.Sprintf(`INSERT INTO %s (
			id, update_time,
			client_id,
			probe_cc, probe_asn,
			platform, software_name,
			software_version, supported_tests,
			network_type, available_bandwidth,
			token, probe_family,
			probe_id, update_type
		) VALUES (
			$1, $2,
			$3, $4,
			$5, $6,
			$7, $8,
			$9, $10,
			$11, $12,
			$13, $14, $15)`,
			pq.QuoteIdentifier(viper.GetString("database.probe-updates-table")))

		stmt, err := tx.Prepare(query)
		if (err != nil) {
			ctx.WithError(err).Error("failed to prepare update probes query")
			return err
		}

		updateID := uuid.NewV4().String()
		_, err = stmt.Exec(updateID, time.Now().UTC(),
							clientID,
							req.ProbeCC, req.ProbeASN,
							req.Platform, req.SoftwareName,
							req.SoftwareVersion, pq.Array(req.SupportedTests),
							req.NetworkType, req.AvailableBandwidth,
							req.Token, req.ProbeFamily,
							req.ProbeID, "register")
		if (err != nil) {
			ctx.WithError(err).Error("failed to add data to update table, rolling back")
			tx.Rollback()
			return errors.New("error in adding data to update probes")
		}
	}

	// Write into the active probes table
	{
		query := fmt.Sprintf(`UPDATE %s SET
			last_updated = $2,
			probe_cc = $3,
			probe_asn = $4,
			platform = $5,
			software_name = $6,
			software_version = $7,
			supported_tests = $8,
			network_type = $9,
			available_bandwidth = $10,
			token = $11,
			probe_family = $12,
			probe_id = $13
			WHERE id = $1`,
			pq.QuoteIdentifier(viper.GetString("database.active-probes-table")))

		stmt, err := tx.Prepare(query)
		if (err != nil) {
			ctx.WithError(err).Error("failed to prepare update probes query")
			return err
		}
		_, err = stmt.Exec(clientID,
							time.Now().UTC(),
							req.ProbeCC,
							req.ProbeASN,
							req.Platform,
							req.SoftwareName,
							req.SoftwareVersion,
							pq.Array(req.SupportedTests),
							req.NetworkType,
							req.AvailableBandwidth,
							req.Token,
							req.ProbeFamily,
							req.ProbeID)
		if (err != nil) {
			ctx.WithError(err).Error("failed to update active table, rolling back")
			tx.Rollback()
			return errors.New("failed to update active table")
		}
	}

	if err := tx.Commit(); err != nil {
		ctx.WithError(err).Error("failed to commit transaction, rolling back")
		return err
	}

	return nil
}

func Register(db *sql.DB, req ClientData) (string, error) {
	if ((req.Platform == "ios" || req.Platform == "android") && req.Token == "") {
		return "", errors.New("missing device token")
	}

	tx, err := db.Begin()
	if err != nil {
		ctx.WithError(err).Error("failed to open transaction")
		return "", err
	}

	var clientID = uuid.NewV4().String()

	{
		query := fmt.Sprintf(`INSERT INTO %s (
			id, creation_time,
			last_updated,
			probe_cc, probe_asn,
			platform, software_name,
			software_version, supported_tests,
			network_type, available_bandwidth,
			token, probe_family,
			probe_id
		) VALUES (
			$1, $2,
			$3, $4,
			$5, $6,
			$7, $8,
			$9, $10,
			$11, $12,
			$13, $14)`,
			pq.QuoteIdentifier(viper.GetString("database.active-probes-table")))

		stmt, err := tx.Prepare(query)
		if (err != nil) {
			ctx.WithError(err).Error("failed to prepare active probes query")
			return "", err
		}
		defer stmt.Close()

		_, err = stmt.Exec(clientID, time.Now().UTC(),
							time.Now().UTC(),
							req.ProbeCC, req.ProbeASN,
							req.Platform, req.SoftwareName,
							req.SoftwareVersion, pq.Array(req.SupportedTests),
							req.NetworkType, req.AvailableBandwidth,
							req.Token, req.ProbeFamily,
							req.ProbeID)
		if err != nil {
			tx.Rollback()
			ctx.WithError(err).Error("failed to insert into active probes table, rolling back")
			return "", err
		}
	}

	{
		query := fmt.Sprintf(`INSERT INTO %s (
			id, update_time,
			client_id,
			probe_cc, probe_asn,
			platform, software_name,
			software_version, supported_tests,
			network_type, available_bandwidth,
			token, probe_family,
			probe_id, update_type
		) VALUES (
			$1, $2,
			$3, $4,
			$5, $6,
			$7, $8,
			$9, $10,
			$11, $12,
			$13, $14, $15)`,
			pq.QuoteIdentifier(viper.GetString("database.probe-updates-table")))

		stmt, err := tx.Prepare(query)
		if (err != nil) {
			ctx.WithError(err).Error("failed to prepare update probes query")
			return "", err
		}
		defer stmt.Close()

		updateID := uuid.NewV4().String()
		_, err = stmt.Exec(updateID, time.Now().UTC(),
							clientID,
							req.ProbeCC, req.ProbeASN,
							req.Platform, req.SoftwareName,
							req.SoftwareVersion, pq.Array(req.SupportedTests),
							req.NetworkType, req.AvailableBandwidth,
							req.Token, req.ProbeFamily,
							req.ProbeID, "register")
		if (err != nil) {
			ctx.WithError(err).Error("failed to add data to update table, rolling back")
			tx.Rollback()
			return "", errors.New("error in adding data to update probes")
		}
	}

	if err := tx.Commit(); err != nil {
		ctx.WithError(err).Error("failed to commit transaction, rolling back")
		return "", err
	}

	return clientID, nil
}

type ActiveClient struct {
	ClientID			string `json:"client_id"`

	ProbeCC				string `json:"probe_cc"`
	ProbeASN			string `json:"probe_asn"`
	Platform			string `json:"platform"`

	SoftwareName		string `json:"software_name"`
	SoftwareVersion		string `json:"software_version"`
	SupportedTests		string `json:"supported_tests"`

	NetworkType			string `json:"network_type"`
	AvailableBandwidth	string `json:"available_bandwidth"`

	Token				string `json:"token"`

	ProbeFamily			string `json:"probe_family"`
	ProbeID				string `json:"probe_id"`

	LastUpdated			time.Time `json:"last_updated"`
	CreationTime		time.Time `json:"creation_time"`
}

func ListClients(db *sql.DB) ([]ActiveClient, error) {
	var activeClients []ActiveClient
	query := fmt.Sprintf(`SELECT
			id, creation_time,
			last_updated,
			probe_cc, probe_asn,
			platform, software_name,
			software_version, supported_tests,
			network_type, available_bandwidth,
			token, probe_family,
			probe_id FROM %s`,
		pq.QuoteIdentifier(viper.GetString("database.active-probes-table")))

	rows, err := db.Query(query)
	if err != nil {
		ctx.WithError(err).Error("failed to list clients")
		return activeClients, err
	}
	defer rows.Close()
	for rows.Next() {
		var ac ActiveClient
		err := rows.Scan(&ac.ClientID,
						&ac.CreationTime,
						&ac.LastUpdated,
						&ac.ProbeCC,
						&ac.ProbeASN,
						&ac.Platform,
						&ac.SoftwareName,
						&ac.SoftwareVersion,
						&ac.SupportedTests,
						&ac.NetworkType,
						&ac.AvailableBandwidth,
						&ac.Token,
						&ac.ProbeFamily,
						&ac.ProbeID)
		if err != nil {
			ctx.WithError(err).Error("failed to iterate over clients")
			return activeClients, err
		}
		activeClients = append(activeClients, ac)
	}
	return activeClients, nil
}

func Start() {
	db, err := initDatabase()

	if (err != nil) {
		ctx.WithError(err).Error("failed to connect to DB")
		return
	}
	defer db.Close()

	router := gin.Default()
	router.GET("/api/v1/clients", func(c *gin.Context) {
		// XXX add authentication
		clientList, err := ListClients(db)
		if err != nil {
			c.JSON(http.StatusBadRequest,
					gin.H{"error": err.Error()})
			return
		}
		c.JSON(http.StatusOK,
				gin.H{"active_clients": clientList})
	})

	router.POST("/api/v1/clients", func(c *gin.Context) {
		var registerReq ClientData
		err := c.BindJSON(&registerReq)
		if (err != nil) {
			ctx.WithError(err).Error("invalid request")
			c.JSON(http.StatusBadRequest,
					gin.H{"error": "invalid request"})
			return
		}

		clientID , err := Register(db, registerReq)
		if (err != nil) {
			c.JSON(http.StatusBadRequest,
					gin.H{"error": err.Error()})
			return
		}

		c.JSON(http.StatusOK, gin.H{"client_id": clientID})
		return
	})

	// XXX do we also want to support a PATCH method?
	router.PUT("/api/v1/clients/:client_id", func(c *gin.Context) {
		var updateReq ClientData
		clientID := c.Param("client_id")
		err := c.BindJSON(&updateReq)
		if (err != nil) {
			ctx.WithError(err).Error("invalid request")
			c.JSON(http.StatusBadRequest,
					gin.H{"error": "invalid request"})
			return
		}
		isRegistered, err := IsClientRegistered(db, clientID)
		if (err != nil) {
			ctx.WithError(err).Error("failed to learn if client is registered")
			c.JSON(http.StatusBadRequest,
					gin.H{"error": err.Error()})
			return
		}
		if (isRegistered == false) {
			c.JSON(http.StatusNotFound,
					gin.H{"error": "client is not registered"})
			return
		}

		err = Update(db, clientID, updateReq)
		if (err != nil) {
			ctx.WithError(err).Error("failed to update")
			c.JSON(http.StatusBadRequest,
					gin.H{"error": err.Error()})
			return
		}
		c.JSON(http.StatusOK,
				gin.H{"status": "ok"})
	})

	Addr := fmt.Sprintf("%s:%d", viper.GetString("api.address"),
								viper.GetInt("api.port"))
	ctx.Infof("starting on %s", Addr)
	s := &http.Server{
		Addr: Addr,
		Handler: router,
	}
	gracehttp.Serve(s)
}
<|endoftext|>


<commit_before>package tokay

import (
	"fmt"
	"html/template"
	"net/http"
	"os"
	"reflect"
	"runtime"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/night-codes/tokay-render"
	"github.com/valyala/fasthttp"
)

type (
	// Render is interface for engine.Render
	Render interface {
		JSON(*fasthttp.RequestCtx, int, interface{}) error
		JSONP(*fasthttp.RequestCtx, int, string, interface{}) error
		HTML(*fasthttp.RequestCtx, int, string, interface{}, ...string) error
		XML(*fasthttp.RequestCtx, int, interface{}) error
	}

	// Handler is the function for handling HTTP requests.
	Handler func(*Context)

	// Engine manages routes and dispatches HTTP requests to the handlers of the matching routes.
	Engine struct {
		RouterGroup
		// Default render engine
		Render Render
		// AppEngine usage marker
		AppEngine bool
		// Print debug messages to log
		Debug bool

		// Enables automatic redirection if the current route can't be matched but a
		// handler for the path with the trailing slash exists.
		// For example if /foo is requested but a route only exists for /foo/, the
		// client is redirected to /foo/ with http status code 301 for GET requests
		// and 307 for all other request methods.
		RedirectTrailingSlash bool

		pool             sync.Pool
		routes           map[string]*Route
		stores           storesMap
		maxParams        int
		notFound         []Handler
		notFoundHandlers []Handler
	}

	// Config is a struct for specifying configuration options for the tokay.Engine object.
	Config struct {
		// Enables automatic redirection if the current route can't be matched but a handler for the path with the trailing slash exists.
		RedirectTrailingSlash bool
		// Print debug messages to log
		Debug bool
		// Extensions to parse template files from. Defaults to [".html"].
		TemplatesExtensions []string
		// Directories to load templates. Default is ["templates"].
		TemplatesDirs []string
		// Left templates delimiter, defaults to {{.
		LeftTemplateDelimiter string
		// Right templates delimiter, defaults to }}.
		RightTemplateDelimiter string
		// Funcs is a slice of FuncMaps to apply to the template upon compilation. This is useful for helper functions. Defaults to [].
		TemplatesFuncs template.FuncMap
	}
)

var (
	// AppEngine usage marker
	AppEngine bool

	// Methods lists all supported HTTP methods by Engine.
	Methods = []string{
		"HEAD",
		"GET",
		"POST",
		"CONNECT",
		"DELETE",
		"OPTIONS",
		"PATCH",
		"PUT",
		"TRACE",
	}
)

// New creates a new Engine object.
func New(config ...*Config) *Engine {
	var r *render.Render
	var cfgRedirectTrailingSlash bool
	var cfgDebug bool
	if len(config) != 0 && config[0] != nil {
		if len(config[0].TemplatesDirs) != 0 {
			r = render.New(&render.Config{
				Directories: config[0].TemplatesDirs,
				Extensions:  config[0].TemplatesExtensions,
				Delims: render.Delims{
					Left: config[0].LeftTemplateDelimiter,
				},
			})
		}
		cfgRedirectTrailingSlash = config[0].RedirectTrailingSlash
		cfgDebug = config[0].Debug
	} else {
		r = render.New()
	}

	debug.Println("Render object created" + r.Templates.DefinedTemplates())

	engine := &Engine{
		AppEngine:             AppEngine,
		routes:                make(map[string]*Route),
		stores:                *newStoresMap(),
		Render:                r,
		RedirectTrailingSlash: cfgRedirectTrailingSlash,
		Debug: cfgDebug,
	}
	engine.RouterGroup = *newRouteGroup("", engine, make([]Handler, 0))
	engine.NotFound(MethodNotAllowedHandler, NotFoundHandler)
	engine.pool.New = func() interface{} {
		return &Context{
			pvalues: make([]string, engine.maxParams),
			engine:  engine,
		}
	}
	return engine
}

func runmsg(addr string, ec chan error, message string) (err error) {
	if message != "" {
		select {
		case err = <-ec:
			return
		case _ = <-time.Tick(time.Second / 4):
			if strings.Contains(message, "%s") {
				fmt.Printf(message+"\n", addr)
			} else {
				fmt.Println(message)
			}
		}
	}
	err = <-ec
	return
}

// Run attaches the engine to a fasthttp server and starts listening and serving HTTP requests.
// It is a shortcut for fasthttp.ListenAndServe(addr, engine.HandleRequest) Note: this method will block the
// calling goroutine indefinitely unless an error happens.
func (engine *Engine) Run(addr string, message ...string) error {
	ec := make(chan error)
	go func() {
		ec <- fasthttp.ListenAndServe(addr, engine.HandleRequest)
	}()
	return runmsg(addr, ec, append(message, "HTTP server started at %s")[0])
}

// RunTLS attaches the engine to a fasthttp server and starts listening and
// serving HTTPS (secure) requests. It is a shortcut for
// fasthttp.ListenAndServeTLS(addr, certFile, keyFile, engine.HandleRequest)
// Note: this method will block the calling goroutine indefinitely unless an error happens.
func (engine *Engine) RunTLS(addr string, certFile, keyFile string, message ...string) error {
	ec := make(chan error)
	go func() {
		ec <- fasthttp.ListenAndServeTLS(addr, certFile, keyFile, engine.HandleRequest)
	}()
	return runmsg(addr, ec, append(message, "HTTPS server started at %s")[0])
}

// RunUnix attaches the engine to a fasthttp server and starts listening and
// serving HTTP requests through the specified unix socket (ie. a file).
// Note: this method will block the calling goroutine indefinitely unless an error happens.
func (engine *Engine) RunUnix(addr string, mode os.FileMode, message ...string) error {
	ec := make(chan error)
	go func() {
		ec <- fasthttp.ListenAndServeUNIX(addr, mode, engine.HandleRequest)
	}()
	return runmsg(addr, ec, append(message, "Unix server started at %s")[0])
}

// HandleRequest handles the HTTP request.
func (engine *Engine) HandleRequest(ctx *fasthttp.RequestCtx) {
	ws := false
	start := time.Now()
	c := engine.pool.Get().(*Context)
	c.init(ctx)
	c.handlers, c.pnames, ws = engine.find(string(ctx.Method()), string(ctx.Path()), c.pvalues)
	fin := func() {
		c.Next()
		engine.pool.Put(c)
		engine.debug(fmt.Sprintf("%-21s | %d | %9v | %-7s %-25s ", time.Now().Format("2006/01/02 - 15:04:05"), c.Response.StatusCode(), time.Since(start), string(ctx.Method()), string(ctx.Path())))
	}
	if ws {
		c.Websocket(fin)
		return
	}
	fin()
}

// Route returns the named route.
// Nil is returned if the named route cannot be found.
func (engine *Engine) Route(name string) *Route {
	return engine.routes[name]
}

// Use appends the specified handlers to the engine and shares them with all routes.
func (engine *Engine) Use(handlers ...Handler) {
	engine.RouterGroup.Use(handlers...)
	engine.notFoundHandlers = combineHandlers(engine.handlers, engine.notFound)
}

// NotFound specifies the handlers that should be invoked when the engine cannot find any route matching a request.
// Note that the handlers registered via Use will be invoked first in this case.
func (engine *Engine) NotFound(handlers ...Handler) {
	engine.notFound = handlers
	engine.notFoundHandlers = combineHandlers(engine.handlers, engine.notFound)
}

// handleError is the error handler for handling any unhandled errors.
func (engine *Engine) handleError(c *Context, err error) {
	c.Error(err.Error(), http.StatusInternalServerError)
}

func (engine *Engine) add(method, path string, handlers []Handler) {
	for _, h := range handlers {
		engine.debug(fmt.Sprintf("%-7s %-25s -->", method, path), runtime.FuncForPC(reflect.ValueOf(h).Pointer()).Name())
	}
	store := engine.stores.Get(method)
	if store == nil {
		store = newStore()
		engine.stores.Set(method, store)
	}
	if n := store.Add(path, handlers); n > engine.maxParams {
		engine.maxParams = n
	}
}

func (engine *Engine) find(method, path string, pvalues []string) (handlers []Handler, pnames []string, ws bool) {
	var hh interface{}
	if store := engine.stores.Get(method); store != nil {
		if hh, pnames = store.Get(path, pvalues); hh != nil {
			return hh.([]Handler), pnames, false
		}
	}
	if method == "GET" {
		if store := engine.stores.Get("WEBSOCKET"); store != nil {
			if hh, pnames = store.Get(path, pvalues); hh != nil {
				return hh.([]Handler), pnames, true
			}
		}
	}

	return engine.notFoundHandlers, pnames, false
}

func (engine *Engine) findAllowedMethods(path string) map[string]bool {
	methods := make(map[string]bool)
	pvalues := make([]string, engine.maxParams)
	engine.stores.Range(func(m string, store routeStore) {
		if handlers, _ := store.Get(path, pvalues); handlers != nil {
			methods[m] = true
		}
	})
	return methods
}

func (engine *Engine) debug(text ...interface{}) {
	if engine.Debug {
		debug.Println(text...)
	}
}

// NotFoundHandler returns a 404 HTTP error indicating a request has no matching route.
func NotFoundHandler(c *Context) {
	if c.engine.RedirectTrailingSlash && redirectTrailingSlash(c) {
		return
	}
	c.String(http.StatusNotFound, http.StatusText(http.StatusNotFound))
}

// MethodNotAllowedHandler handles the situation when a request has matching route without matching HTTP method.
// In this case, the handler will respond with an Allow HTTP header listing the allowed HTTP methods.
// Otherwise, the handler will do nothing and let the next handler (usually a NotFoundHandler) to handle the problem.
func MethodNotAllowedHandler(c *Context) {
	methods := c.Engine().findAllowedMethods(string(c.Path()))
	if len(methods) == 0 {
		return
	}
	methods["OPTIONS"] = true
	ms := make([]string, len(methods))
	i := 0
	for method := range methods {
		ms[i] = method
		i++
	}
	sort.Strings(ms)
	c.Response.Header.Set("Allow", strings.Join(ms, ", "))
	if string(c.Method()) != "OPTIONS" {
		c.Response.SetStatusCode(http.StatusMethodNotAllowed)
	}
	c.Abort()
	return
}

func redirectTrailingSlash(c *Context) bool {
	path := c.Path()
	statusCode := 301 // Permanent redirect, request with GET method
	if c.Method() != "GET" {
		statusCode = 307
	}
	if len(path) == 0 {
		c.Redirect(statusCode, "/")
		return true
	}

	pathSpl := strings.Split(path, "/")
	d := 1
	if path[len(path)-1] == '/' && len(pathSpl) > 1 {
		d = 2
	}
	hasdot := strings.Index(pathSpl[len(pathSpl)-d], ".") != -1

	if path[len(path)-1] != '/' && !hasdot {
		c.Redirect(statusCode, path+"/")
		return true
	}
	if path[len(path)-1] == '/' && hasdot {
		c.Redirect(statusCode, path[:len(path)-1])
		return true
	}

	return false
}
<commit_msg>Remove render debug message<commit_after>package tokay

import (
	"fmt"
	"html/template"
	"net/http"
	"os"
	"reflect"
	"runtime"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/night-codes/tokay-render"
	"github.com/valyala/fasthttp"
)

type (
	// Render is interface for engine.Render
	Render interface {
		JSON(*fasthttp.RequestCtx, int, interface{}) error
		JSONP(*fasthttp.RequestCtx, int, string, interface{}) error
		HTML(*fasthttp.RequestCtx, int, string, interface{}, ...string) error
		XML(*fasthttp.RequestCtx, int, interface{}) error
	}

	// Handler is the function for handling HTTP requests.
	Handler func(*Context)

	// Engine manages routes and dispatches HTTP requests to the handlers of the matching routes.
	Engine struct {
		RouterGroup
		// Default render engine
		Render Render
		// AppEngine usage marker
		AppEngine bool
		// Print debug messages to log
		Debug bool

		// Enables automatic redirection if the current route can't be matched but a
		// handler for the path with the trailing slash exists.
		// For example if /foo is requested but a route only exists for /foo/, the
		// client is redirected to /foo/ with http status code 301 for GET requests
		// and 307 for all other request methods.
		RedirectTrailingSlash bool

		pool             sync.Pool
		routes           map[string]*Route
		stores           storesMap
		maxParams        int
		notFound         []Handler
		notFoundHandlers []Handler
	}

	// Config is a struct for specifying configuration options for the tokay.Engine object.
	Config struct {
		// Enables automatic redirection if the current route can't be matched but a handler for the path with the trailing slash exists.
		RedirectTrailingSlash bool
		// Print debug messages to log
		Debug bool
		// Extensions to parse template files from. Defaults to [".html"].
		TemplatesExtensions []string
		// Directories to load templates. Default is ["templates"].
		TemplatesDirs []string
		// Left templates delimiter, defaults to {{.
		LeftTemplateDelimiter string
		// Right templates delimiter, defaults to }}.
		RightTemplateDelimiter string
		// Funcs is a slice of FuncMaps to apply to the template upon compilation. This is useful for helper functions. Defaults to [].
		TemplatesFuncs template.FuncMap
	}
)

var (
	// AppEngine usage marker
	AppEngine bool

	// Methods lists all supported HTTP methods by Engine.
	Methods = []string{
		"HEAD",
		"GET",
		"POST",
		"CONNECT",
		"DELETE",
		"OPTIONS",
		"PATCH",
		"PUT",
		"TRACE",
	}
)

// New creates a new Engine object.
func New(config ...*Config) *Engine {
	var r *render.Render
	var cfgRedirectTrailingSlash bool
	var cfgDebug bool
	if len(config) != 0 && config[0] != nil {
		if len(config[0].TemplatesDirs) != 0 {
			r = render.New(&render.Config{
				Directories: config[0].TemplatesDirs,
				Extensions:  config[0].TemplatesExtensions,
				Delims: render.Delims{
					Left: config[0].LeftTemplateDelimiter,
				},
			})
		}
		cfgRedirectTrailingSlash = config[0].RedirectTrailingSlash
		cfgDebug = config[0].Debug
	} else {
		r = render.New()
	}

	engine := &Engine{
		AppEngine:             AppEngine,
		routes:                make(map[string]*Route),
		stores:                *newStoresMap(),
		Render:                r,
		RedirectTrailingSlash: cfgRedirectTrailingSlash,
		Debug: cfgDebug,
	}
	engine.RouterGroup = *newRouteGroup("", engine, make([]Handler, 0))
	engine.NotFound(MethodNotAllowedHandler, NotFoundHandler)
	engine.pool.New = func() interface{} {
		return &Context{
			pvalues: make([]string, engine.maxParams),
			engine:  engine,
		}
	}
	return engine
}

func runmsg(addr string, ec chan error, message string) (err error) {
	if message != "" {
		select {
		case err = <-ec:
			return
		case _ = <-time.Tick(time.Second / 4):
			if strings.Contains(message, "%s") {
				fmt.Printf(message+"\n", addr)
			} else {
				fmt.Println(message)
			}
		}
	}
	err = <-ec
	return
}

// Run attaches the engine to a fasthttp server and starts listening and serving HTTP requests.
// It is a shortcut for fasthttp.ListenAndServe(addr, engine.HandleRequest) Note: this method will block the
// calling goroutine indefinitely unless an error happens.
func (engine *Engine) Run(addr string, message ...string) error {
	ec := make(chan error)
	go func() {
		ec <- fasthttp.ListenAndServe(addr, engine.HandleRequest)
	}()
	return runmsg(addr, ec, append(message, "HTTP server started at %s")[0])
}

// RunTLS attaches the engine to a fasthttp server and starts listening and
// serving HTTPS (secure) requests. It is a shortcut for
// fasthttp.ListenAndServeTLS(addr, certFile, keyFile, engine.HandleRequest)
// Note: this method will block the calling goroutine indefinitely unless an error happens.
func (engine *Engine) RunTLS(addr string, certFile, keyFile string, message ...string) error {
	ec := make(chan error)
	go func() {
		ec <- fasthttp.ListenAndServeTLS(addr, certFile, keyFile, engine.HandleRequest)
	}()
	return runmsg(addr, ec, append(message, "HTTPS server started at %s")[0])
}

// RunUnix attaches the engine to a fasthttp server and starts listening and
// serving HTTP requests through the specified unix socket (ie. a file).
// Note: this method will block the calling goroutine indefinitely unless an error happens.
func (engine *Engine) RunUnix(addr string, mode os.FileMode, message ...string) error {
	ec := make(chan error)
	go func() {
		ec <- fasthttp.ListenAndServeUNIX(addr, mode, engine.HandleRequest)
	}()
	return runmsg(addr, ec, append(message, "Unix server started at %s")[0])
}

// HandleRequest handles the HTTP request.
func (engine *Engine) HandleRequest(ctx *fasthttp.RequestCtx) {
	ws := false
	start := time.Now()
	c := engine.pool.Get().(*Context)
	c.init(ctx)
	c.handlers, c.pnames, ws = engine.find(string(ctx.Method()), string(ctx.Path()), c.pvalues)
	fin := func() {
		c.Next()
		engine.pool.Put(c)
		engine.debug(fmt.Sprintf("%-21s | %d | %9v | %-7s %-25s ", time.Now().Format("2006/01/02 - 15:04:05"), c.Response.StatusCode(), time.Since(start), string(ctx.Method()), string(ctx.Path())))
	}
	if ws {
		c.Websocket(fin)
		return
	}
	fin()
}

// Route returns the named route.
// Nil is returned if the named route cannot be found.
func (engine *Engine) Route(name string) *Route {
	return engine.routes[name]
}

// Use appends the specified handlers to the engine and shares them with all routes.
func (engine *Engine) Use(handlers ...Handler) {
	engine.RouterGroup.Use(handlers...)
	engine.notFoundHandlers = combineHandlers(engine.handlers, engine.notFound)
}

// NotFound specifies the handlers that should be invoked when the engine cannot find any route matching a request.
// Note that the handlers registered via Use will be invoked first in this case.
func (engine *Engine) NotFound(handlers ...Handler) {
	engine.notFound = handlers
	engine.notFoundHandlers = combineHandlers(engine.handlers, engine.notFound)
}

// handleError is the error handler for handling any unhandled errors.
func (engine *Engine) handleError(c *Context, err error) {
	c.Error(err.Error(), http.StatusInternalServerError)
}

func (engine *Engine) add(method, path string, handlers []Handler) {
	for _, h := range handlers {
		engine.debug(fmt.Sprintf("%-7s %-25s -->", method, path), runtime.FuncForPC(reflect.ValueOf(h).Pointer()).Name())
	}
	store := engine.stores.Get(method)
	if store == nil {
		store = newStore()
		engine.stores.Set(method, store)
	}
	if n := store.Add(path, handlers); n > engine.maxParams {
		engine.maxParams = n
	}
}

func (engine *Engine) find(method, path string, pvalues []string) (handlers []Handler, pnames []string, ws bool) {
	var hh interface{}
	if store := engine.stores.Get(method); store != nil {
		if hh, pnames = store.Get(path, pvalues); hh != nil {
			return hh.([]Handler), pnames, false
		}
	}
	if method == "GET" {
		if store := engine.stores.Get("WEBSOCKET"); store != nil {
			if hh, pnames = store.Get(path, pvalues); hh != nil {
				return hh.([]Handler), pnames, true
			}
		}
	}

	return engine.notFoundHandlers, pnames, false
}

func (engine *Engine) findAllowedMethods(path string) map[string]bool {
	methods := make(map[string]bool)
	pvalues := make([]string, engine.maxParams)
	engine.stores.Range(func(m string, store routeStore) {
		if handlers, _ := store.Get(path, pvalues); handlers != nil {
			methods[m] = true
		}
	})
	return methods
}

func (engine *Engine) debug(text ...interface{}) {
	if engine.Debug {
		debug.Println(text...)
	}
}

// NotFoundHandler returns a 404 HTTP error indicating a request has no matching route.
func NotFoundHandler(c *Context) {
	if c.engine.RedirectTrailingSlash && redirectTrailingSlash(c) {
		return
	}
	c.String(http.StatusNotFound, http.StatusText(http.StatusNotFound))
}

// MethodNotAllowedHandler handles the situation when a request has matching route without matching HTTP method.
// In this case, the handler will respond with an Allow HTTP header listing the allowed HTTP methods.
// Otherwise, the handler will do nothing and let the next handler (usually a NotFoundHandler) to handle the problem.
func MethodNotAllowedHandler(c *Context) {
	methods := c.Engine().findAllowedMethods(string(c.Path()))
	if len(methods) == 0 {
		return
	}
	methods["OPTIONS"] = true
	ms := make([]string, len(methods))
	i := 0
	for method := range methods {
		ms[i] = method
		i++
	}
	sort.Strings(ms)
	c.Response.Header.Set("Allow", strings.Join(ms, ", "))
	if string(c.Method()) != "OPTIONS" {
		c.Response.SetStatusCode(http.StatusMethodNotAllowed)
	}
	c.Abort()
	return
}

func redirectTrailingSlash(c *Context) bool {
	path := c.Path()
	statusCode := 301 // Permanent redirect, request with GET method
	if c.Method() != "GET" {
		statusCode = 307
	}
	if len(path) == 0 {
		c.Redirect(statusCode, "/")
		return true
	}

	pathSpl := strings.Split(path, "/")
	d := 1
	if path[len(path)-1] == '/' && len(pathSpl) > 1 {
		d = 2
	}
	hasdot := strings.Index(pathSpl[len(pathSpl)-d], ".") != -1

	if path[len(path)-1] != '/' && !hasdot {
		c.Redirect(statusCode, path+"/")
		return true
	}
	if path[len(path)-1] == '/' && hasdot {
		c.Redirect(statusCode, path[:len(path)-1])
		return true
	}

	return false
}
<|endoftext|>


<commit_before>package GOsu

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
)

// Types
const (
	OSU        = "0"
	TAIKO      = "1"
	CTB        = "2"
	MANIA      = "3"
	BEATMAPSET = "s"
	BEATMAPID  = "b"
	USERID     = "u"
)

var (
	API_URL          string = "https://osu.ppy.sh/api/"
	API_RECENT_PLAYS string = "get_user_recent"
	API_GET_BEATMAPS string = "get_beatmaps"
	API_GET_USER     string = "get_user"
)

type Database struct {
	API_KEY string
}

type Beatmap struct {
	Beatmapset_ID     string
	Beatmap_ID        string
	Approved          string
	Approved_Date     string
	Last_Update       string
	Total_Length      string
	Hit_Length        string
	Version           string
	Artist            string
	Title             string
	Creator           string
	Bpm               string
	Source            string
	Difficulty_Rating string
	Diff_Size         string
	Diff_Overall      string
	Diff_Approach     string
	Diff_Drain        string
	Mode              string
}

type Song struct {
	Beatmap_ID   string
	Score        string
	MaxCombo     string
	Count50      string
	Count100     string
	Count300     string
	CountMiss    string
	CountKatu    string
	CountGeki    string
	Perfect      string
	Enabled_Mods string
	User_ID      string
	Date         string
	Rank         string
}

type User struct {
	User_ID       string
	Username      string
	Count300      string
	Count100      string
	Count50       string
	PlayCount     string
	Ranked_Score  string
	Total_Score   string
	PP_Rank       string
	Level         string
	PP_Raw        string
	Accuracy      string
	Count_Rank_SS string
	Count_Rank_S  string
	Count_Rank_A  string
	Country       string
	Events        []Event
}

type Event struct {
	Display_HTML  string
	Beatmap_ID    string
	Beatmapset_ID string
	Date          string
	EpicFactor    string
}

func (d *Database) SetAPIKey() error {
	tempKey, err := ioutil.ReadFile("./APIKEY.txt")

	// If there is no file, try find the API Key in the Environment Variables.
	if err != nil {
		d.API_KEY = os.Getenv("APIKEY")

		if len(d.API_KEY) <= 1 {
			err = errors.New("API Key: unable to locate API Key in environment variables or in local APIKEY.txt file.")
			return err
		} else {
			err = nil
		}
	} else {
		d.API_KEY = string(tempKey)
	}

	// Trims spaces and trailing newlines from the API key so that the URL
	// to retrieve songs can be built properly.
	d.API_KEY = strings.TrimSpace(d.API_KEY)
	d.API_KEY = strings.Trim(d.API_KEY, "\r\n")

	return err
}

func (d Database) BuildRecentURL(USER_ID string, GAME_TYPE string) string {
	return API_URL + API_RECENT_PLAYS + "?k=" + d.API_KEY + "&u=" + USER_ID + "&m=" + GAME_TYPE
}

func (d Database) BuildBeatmapURL(ID string, TYPE string) string {
	return API_URL + API_GET_BEATMAPS + "?k=" + d.API_KEY + "&" + TYPE + "=" + ID
}

func (d Database) BuildUserURL(USER_ID string, GAME_TYPE string, DAYS string) string {
	return API_URL + API_GET_USER + "?k=" + d.API_KEY + "&u=" + USER_ID + "&m=" + GAME_TYPE + "&event_days=" + DAYS
}

func RetrieveHTML(URL string) ([]byte, error) {
	res, err := http.Get(URL)
	defer res.Body.Close()

	if err != nil {
		return nil, errors.New("HTTP: Could not open a connection to the Osu! API server.")
	}

	html, err := ioutil.ReadAll(res.Body)

	if err != nil {
		return nil, errors.New("HTML: Could not read the HTML page grabbed.")
	}

	return html, err
}

func (d Database) GetUser(USER_ID string, GAME_TYPE string, DAYS string) ([]User, error) {
	var user []User
	url := d.BuildUserURL(USER_ID, GAME_TYPE, DAYS)
	html, err := RetrieveHTML(url)

	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(html, &user)

	if err != nil {
		return nil, errors.New("JSON: Couldn't process HTML into JSON data. You might have the wrong page or a wrong API key. The HTML grabbed at " + url + " will be displayed below:\n" + string(html))
	}

	return user, err
}

func (d Database) GetBeatmaps(ID string, TYPE string) ([]Beatmap, error) {
	var beatmaps []Beatmap
	url := d.BuildBeatmapURL(ID, TYPE)
	html, err := RetrieveHTML(url)

	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(html, &beatmaps)

	if err != nil {
		return nil, errors.New("JSON: Couldn't process HTML into JSON data. You might have the wrong page or a wrong API key. The HTML grabbed at " + url + " will be displayed below:\n" + string(html))
	}

	return beatmaps, err
}

func (d Database) GetRecentPlays(USER_ID string, GAME_TYPE string) ([]Song, error) {
	var songs []Song
	url := d.BuildRecentURL(USER_ID, GAME_TYPE)
	html, err := RetrieveHTML(url)

	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(html, &songs)

	if err != nil {
		return nil, errors.New("JSON: Couldn't process HTML into JSON data. You might have the wrong page or a wrong API key. The HTML grabbed at " + url + " will be displayed below:\n" + string(html))
	}

	return songs, err
}

// ONLY A TEMPORARY FUNCTION.
// Use this function if you are behind a proxy/corporate network and want to work off a local file.
// It will serve as a local HTML file for you to test the website.
func GetLocalPlays(path string) ([]Song, error) {
	var songs []Song

	html, err := ioutil.ReadFile(path)

	if err != nil {
		return nil, errors.New("HTML: Could not read the local HTML page properly.")
	}

	err = json.Unmarshal(html, &songs)

	if err != nil {
		return nil, errors.New("JSON: Couldn't process the local HTML page, most likely due to not being in the right format.")
	}

	return songs, err
}
<commit_msg>Added GetScore function.<commit_after>package GOsu

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
)

// Types
const (
	OSU        = "0"
	TAIKO      = "1"
	CTB        = "2"
	MANIA      = "3"
	BEATMAPSET = "s"
	BEATMAPID  = "b"
	USERID     = "u"
)

var (
	API_URL          string = "https://osu.ppy.sh/api/"
	API_RECENT_PLAYS string = "get_user_recent"
	API_GET_BEATMAPS string = "get_beatmaps"
	API_GET_USER     string = "get_user"
	API_GET_SCORES   string = "get_scores"
)

type Database struct {
	API_KEY string
}

type Beatmap struct {
	Beatmapset_ID     string
	Beatmap_ID        string
	Approved          string
	Approved_Date     string
	Last_Update       string
	Total_Length      string
	Hit_Length        string
	Version           string
	Artist            string
	Title             string
	Creator           string
	Bpm               string
	Source            string
	Difficulty_Rating string
	Diff_Size         string
	Diff_Overall      string
	Diff_Approach     string
	Diff_Drain        string
	Mode              string
}

type Song struct {
	Beatmap_ID   string
	Score        string
	MaxCombo     string
	Count50      string
	Count100     string
	Count300     string
	CountMiss    string
	CountKatu    string
	CountGeki    string
	Perfect      string
	Enabled_Mods string
	User_ID      string
	Date         string
	Rank         string
}

type User struct {
	User_ID       string
	Username      string
	Count300      string
	Count100      string
	Count50       string
	PlayCount     string
	Ranked_Score  string
	Total_Score   string
	PP_Rank       string
	Level         string
	PP_Raw        string
	Accuracy      string
	Count_Rank_SS string
	Count_Rank_S  string
	Count_Rank_A  string
	Country       string
	Events        []Event
}

type Event struct {
	Display_HTML  string
	Beatmap_ID    string
	Beatmapset_ID string
	Date          string
	EpicFactor    string
}

type Score struct {
	Score        string
	Username     string
	MaxCombo     string
	Count50      string
	Count100     string
	Count300     string
	CountMiss    string
	CountKatu    string
	CountGeki    string
	Perfect      string
	Enabled_Mods string
	User_ID      string
	Date         string
	Rank         string
	PP           string
}

func (d *Database) SetAPIKey() error {
	tempKey, err := ioutil.ReadFile("./APIKEY.txt")

	// If there is no file, try find the API Key in the Environment Variables.
	if err != nil {
		d.API_KEY = os.Getenv("APIKEY")

		if len(d.API_KEY) <= 1 {
			err = errors.New("API Key: unable to locate API Key in environment variables or in local APIKEY.txt file.")
			return err
		} else {
			err = nil
		}
	} else {
		d.API_KEY = string(tempKey)
	}

	// Trims spaces and trailing newlines from the API key so that the URL
	// to retrieve songs can be built properly.
	d.API_KEY = strings.TrimSpace(d.API_KEY)
	d.API_KEY = strings.Trim(d.API_KEY, "\r\n")

	return err
}

func (d Database) BuildRecentURL(USER_ID string, GAME_TYPE string) string {
	return API_URL + API_RECENT_PLAYS + "?k=" + d.API_KEY + "&u=" + USER_ID + "&m=" + GAME_TYPE
}

func (d Database) BuildBeatmapURL(ID string, TYPE string) string {
	return API_URL + API_GET_BEATMAPS + "?k=" + d.API_KEY + "&" + TYPE + "=" + ID
}

func (d Database) BuildUserURL(USER_ID string, GAME_TYPE string, DAYS string) string {
	return API_URL + API_GET_USER + "?k=" + d.API_KEY + "&u=" + USER_ID + "&m=" + GAME_TYPE + "&event_days=" + DAYS
}

func (d Database) BuildScoreURL(BEATMAP_ID string, USER_ID string, GAME_TYPE string) string {
	return API_URL + API_GET_SCORES + "?k=" + d.API_KEY + "&b=" + BEATMAP_ID + "&m=" + GAME_TYPE + "&u=" + USER_ID
}

func RetrieveHTML(URL string) ([]byte, error) {
	res, err := http.Get(URL)
	defer res.Body.Close()

	if err != nil {
		return nil, errors.New("HTTP: Could not open a connection to the Osu! API server.")
	}

	html, err := ioutil.ReadAll(res.Body)

	if err != nil {
		return nil, errors.New("HTML: Could not read the HTML page grabbed.")
	}

	return html, err
}

func (d Database) GetUser(USER_ID string, GAME_TYPE string, DAYS string) ([]User, error) {
	var user []User
	url := d.BuildUserURL(USER_ID, GAME_TYPE, DAYS)
	html, err := RetrieveHTML(url)

	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(html, &user)

	if err != nil {
		return nil, errors.New("JSON: Couldn't process HTML into JSON data. You might have the wrong page or a wrong API key. The HTML grabbed at " + url + " will be displayed below:\n" + string(html))
	}

	return user, err
}

func (d Database) GetBeatmaps(ID string, TYPE string) ([]Beatmap, error) {
	var beatmaps []Beatmap
	url := d.BuildBeatmapURL(ID, TYPE)
	html, err := RetrieveHTML(url)

	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(html, &beatmaps)

	if err != nil {
		return nil, errors.New("JSON: Couldn't process HTML into JSON data. You might have the wrong page or a wrong API key. The HTML grabbed at " + url + " will be displayed below:\n" + string(html))
	}

	return beatmaps, err
}

func (d Database) GetRecentPlays(USER_ID string, GAME_TYPE string) ([]Song, error) {
	var songs []Song
	url := d.BuildRecentURL(USER_ID, GAME_TYPE)
	html, err := RetrieveHTML(url)

	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(html, &songs)

	if err != nil {
		return nil, errors.New("JSON: Couldn't process HTML into JSON data. You might have the wrong page or a wrong API key. The HTML grabbed at " + url + " will be displayed below:\n" + string(html))
	}

	return songs, err
}

func (d Database) GetScores(BEATMAP_ID string, USER_ID string, GAME_TYPE string) ([]Score, error) {
	var scores []Score
	url := d.BuildScoreURL(BEATMAP_ID, USER_ID, GAME_TYPE)
	html, err := RetrieveHTML(url)

	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(html, &scores)

	if err != nil {
		return nil, errors.New("JSON: Couldn't process HTML into JSON data. You might have the wrong page or a wrong API key. The HTML grabbed at " + url + " will be displayed below:\n" + string(html))
	}

	return scores, err
}

// ONLY A TEMPORARY FUNCTION.
// Use this function if you are behind a proxy/corporate network and want to work off a local file.
// It will serve as a local HTML file for you to test the website.
func GetLocalPlays(path string) ([]Song, error) {
	var songs []Song

	html, err := ioutil.ReadFile(path)

	if err != nil {
		return nil, errors.New("HTML: Could not read the local HTML page properly.")
	}

	err = json.Unmarshal(html, &songs)

	if err != nil {
		return nil, errors.New("JSON: Couldn't process the local HTML page, most likely due to not being in the right format.")
	}

	return songs, err
}
<|endoftext|>


<commit_before>package collectd

import (
	"fmt"
	"net"
	"net/http"
	"sync/atomic"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/cep21/gohelpers/structdefaults"
	"github.com/cep21/gohelpers/workarounds"
	"github.com/gorilla/mux"
	"github.com/signalfx/metricproxy/config"

	"strings"

	"encoding/json"

	"github.com/signalfx/metricproxy/datapoint"
	"github.com/signalfx/metricproxy/datapoint/dpsink"
	"github.com/signalfx/metricproxy/protocol"
	"github.com/signalfx/metricproxy/reqcounter"
	"github.com/signalfx/metricproxy/stats"
	"github.com/signalfx/metricproxy/web"
	"golang.org/x/net/context"
)

// ListenerServer will listen for collectd datapoint connections
type ListenerServer struct {
	stats.Keeper
	name     string
	listener net.Listener
	server   http.Server
}

var _ protocol.Listener = &ListenerServer{}

// Close the socket currently open for collectd JSON connections
func (streamer *ListenerServer) Close() error {
	return streamer.listener.Close()
}

// JSONDecoder can decode collectd's native JSON datapoint format
type JSONDecoder struct {
	SendTo dpsink.Sink

	TotalErrors    int64
	TotalBlankDims int64
}

const sfxDimQueryParamPrefix string = "sfxdim_"

// ServeHTTPC decodes datapoints for the connection and sends them to the decoder's sink
func (decoder *JSONDecoder) ServeHTTPC(ctx context.Context, rw http.ResponseWriter, req *http.Request) {
	err := decoder.Read(ctx, req)
	if err != nil {
		atomic.AddInt64(&decoder.TotalErrors, 1)
		rw.WriteHeader(http.StatusBadRequest)
		rw.Write([]byte(fmt.Sprintf("Unable to decode json: %s", err.Error())))
		return
	}
	rw.Write([]byte(`"OK"`))
}

func (decoder *JSONDecoder) Read(ctx context.Context, req *http.Request) error {
	defaultDims := decoder.defaultDims(req)
	var d JSONWriteBody
	err := json.NewDecoder(req.Body).Decode(&d)
	if err != nil {
		return err
	}
	dps := make([]*datapoint.Datapoint, 0, len(d)*2)
	for _, f := range d {
		if f.TypeS != nil && f.Time != nil {
			for i := range f.Dsnames {
				if i < len(f.Dstypes) && i < len(f.Values) {
					dps = append(dps, NewDatapoint(f, uint(i), defaultDims))
				}
			}
		}
	}
	return decoder.SendTo.AddDatapoints(ctx, dps)
}

func (decoder *JSONDecoder) defaultDims(req *http.Request) map[string]string {
	params := req.URL.Query()
	defaultDims := make(map[string]string)
	for key := range params {
		if strings.HasPrefix(key, sfxDimQueryParamPrefix) {
			value := params.Get(key)
			if len(value) == 0 {
				atomic.AddInt64(&decoder.TotalBlankDims, 1)
				continue
			}
			key = key[len(sfxDimQueryParamPrefix):]
			defaultDims[key] = value
		}
	}
	return defaultDims
}

// Stats about this decoder, including how many datapoints it decoded
func (decoder *JSONDecoder) Stats(dimensions map[string]string) []*datapoint.Datapoint {
	return []*datapoint.Datapoint{
		datapoint.New("total_blank_dims", dimensions, datapoint.NewIntValue(decoder.TotalBlankDims), datapoint.Counter, time.Now()),
		datapoint.New("invalid_collectd_json", dimensions, datapoint.NewIntValue(decoder.TotalErrors), datapoint.Counter, time.Now()),
	}
}

var defaultCollectdConfig = &config.ListenFrom{
	ListenAddr:      workarounds.GolangDoesnotAllowPointerToStringLiteral("127.0.0.1:8081"),
	TimeoutDuration: workarounds.GolangDoesnotAllowPointerToTimeLiteral(time.Second * 30),
	ListenPath:      workarounds.GolangDoesnotAllowPointerToStringLiteral("/post-collectd"),
	Name:            workarounds.GolangDoesnotAllowPointerToStringLiteral("collectd"),
	JSONEngine:      workarounds.GolangDoesnotAllowPointerToStringLiteral("native"),
}

// ListenerLoader loads a listener for collectd write_http protocol
func ListenerLoader(ctx context.Context, sink dpsink.Sink, listenFrom *config.ListenFrom) (*ListenerServer, error) {
	structdefaults.FillDefaultFrom(listenFrom, defaultCollectdConfig)
	log.WithField("listenFrom", listenFrom).Info("Creating listener using final config")
	return StartListeningCollectDHTTPOnPort(ctx, sink, *listenFrom.ListenAddr, *listenFrom.ListenPath, *listenFrom.TimeoutDuration, *listenFrom.Name)
}

// StartListeningCollectDHTTPOnPort servers http collectd requests
func StartListeningCollectDHTTPOnPort(ctx context.Context, sink dpsink.Sink,
	listenAddr string, listenPath string, clientTimeout time.Duration, name string) (*ListenerServer, error) {

	listener, err := net.Listen("tcp", listenAddr)
	if err != nil {
		return nil, err
	}
	h, st := SetupHandler(ctx, name, sink)

	r := mux.NewRouter()
	r.Path(listenPath).Headers("Content-type", "application/json").Handler(h)

	listenServer := ListenerServer{
		Keeper:   st,
		name:     name,
		listener: listener,
		server: http.Server{
			Handler:      r,
			Addr:         listenAddr,
			ReadTimeout:  clientTimeout,
			WriteTimeout: clientTimeout,
		},
	}

	go listenServer.server.Serve(listener)
	return &listenServer, nil
}

// SetupHandler is shared between signalfx and here to setup listening for collectd connections.
// Will do shared basic setup like configuring request counters
func SetupHandler(ctx context.Context, name string, sink dpsink.Sink) (*web.Handler, stats.Keeper) {
	metricTracking := reqcounter.RequestCounter{}
	counter := &dpsink.Counter{}
	collectdDecoder := JSONDecoder{
		SendTo: dpsink.FromChain(sink, dpsink.NextWrap(counter)),
	}
	h := web.NewHandler(ctx, &collectdDecoder).Add(web.NextHTTP(metricTracking.ServeHTTP))
	st := stats.ToKeeperMany(map[string]string{"listener": name, "type": "collectd"}, &metricTracking, counter, &collectdDecoder)
	return h, st
}
<commit_msg>Check for nil values in collectd listener<commit_after>package collectd

import (
	"fmt"
	"net"
	"net/http"
	"sync/atomic"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/cep21/gohelpers/structdefaults"
	"github.com/cep21/gohelpers/workarounds"
	"github.com/gorilla/mux"
	"github.com/signalfx/metricproxy/config"

	"strings"

	"encoding/json"

	"github.com/signalfx/metricproxy/datapoint"
	"github.com/signalfx/metricproxy/datapoint/dpsink"
	"github.com/signalfx/metricproxy/protocol"
	"github.com/signalfx/metricproxy/reqcounter"
	"github.com/signalfx/metricproxy/stats"
	"github.com/signalfx/metricproxy/web"
	"golang.org/x/net/context"
)

// ListenerServer will listen for collectd datapoint connections
type ListenerServer struct {
	stats.Keeper
	name     string
	listener net.Listener
	server   http.Server
}

var _ protocol.Listener = &ListenerServer{}

// Close the socket currently open for collectd JSON connections
func (streamer *ListenerServer) Close() error {
	return streamer.listener.Close()
}

// JSONDecoder can decode collectd's native JSON datapoint format
type JSONDecoder struct {
	SendTo dpsink.Sink

	TotalErrors    int64
	TotalBlankDims int64
}

const sfxDimQueryParamPrefix string = "sfxdim_"

// ServeHTTPC decodes datapoints for the connection and sends them to the decoder's sink
func (decoder *JSONDecoder) ServeHTTPC(ctx context.Context, rw http.ResponseWriter, req *http.Request) {
	err := decoder.Read(ctx, req)
	if err != nil {
		atomic.AddInt64(&decoder.TotalErrors, 1)
		rw.WriteHeader(http.StatusBadRequest)
		rw.Write([]byte(fmt.Sprintf("Unable to decode json: %s", err.Error())))
		return
	}
	rw.Write([]byte(`"OK"`))
}

func (decoder *JSONDecoder) Read(ctx context.Context, req *http.Request) error {
	defaultDims := decoder.defaultDims(req)
	var d JSONWriteBody
	err := json.NewDecoder(req.Body).Decode(&d)
	if err != nil {
		return err
	}
	dps := make([]*datapoint.Datapoint, 0, len(d)*2)
	for _, f := range d {
		if f.TypeS != nil && f.Time != nil {
			for i := range f.Dsnames {
				if i < len(f.Dstypes) && i < len(f.Values) && f.Values[i] != nil {
					dps = append(dps, NewDatapoint(f, uint(i), defaultDims))
				}
			}
		}
	}
	return decoder.SendTo.AddDatapoints(ctx, dps)
}

func (decoder *JSONDecoder) defaultDims(req *http.Request) map[string]string {
	params := req.URL.Query()
	defaultDims := make(map[string]string)
	for key := range params {
		if strings.HasPrefix(key, sfxDimQueryParamPrefix) {
			value := params.Get(key)
			if len(value) == 0 {
				atomic.AddInt64(&decoder.TotalBlankDims, 1)
				continue
			}
			key = key[len(sfxDimQueryParamPrefix):]
			defaultDims[key] = value
		}
	}
	return defaultDims
}

// Stats about this decoder, including how many datapoints it decoded
func (decoder *JSONDecoder) Stats(dimensions map[string]string) []*datapoint.Datapoint {
	return []*datapoint.Datapoint{
		datapoint.New("total_blank_dims", dimensions, datapoint.NewIntValue(decoder.TotalBlankDims), datapoint.Counter, time.Now()),
		datapoint.New("invalid_collectd_json", dimensions, datapoint.NewIntValue(decoder.TotalErrors), datapoint.Counter, time.Now()),
	}
}

var defaultCollectdConfig = &config.ListenFrom{
	ListenAddr:      workarounds.GolangDoesnotAllowPointerToStringLiteral("127.0.0.1:8081"),
	TimeoutDuration: workarounds.GolangDoesnotAllowPointerToTimeLiteral(time.Second * 30),
	ListenPath:      workarounds.GolangDoesnotAllowPointerToStringLiteral("/post-collectd"),
	Name:            workarounds.GolangDoesnotAllowPointerToStringLiteral("collectd"),
	JSONEngine:      workarounds.GolangDoesnotAllowPointerToStringLiteral("native"),
}

// ListenerLoader loads a listener for collectd write_http protocol
func ListenerLoader(ctx context.Context, sink dpsink.Sink, listenFrom *config.ListenFrom) (*ListenerServer, error) {
	structdefaults.FillDefaultFrom(listenFrom, defaultCollectdConfig)
	log.WithField("listenFrom", listenFrom).Info("Creating listener using final config")
	return StartListeningCollectDHTTPOnPort(ctx, sink, *listenFrom.ListenAddr, *listenFrom.ListenPath, *listenFrom.TimeoutDuration, *listenFrom.Name)
}

// StartListeningCollectDHTTPOnPort servers http collectd requests
func StartListeningCollectDHTTPOnPort(ctx context.Context, sink dpsink.Sink,
	listenAddr string, listenPath string, clientTimeout time.Duration, name string) (*ListenerServer, error) {

	listener, err := net.Listen("tcp", listenAddr)
	if err != nil {
		return nil, err
	}
	h, st := SetupHandler(ctx, name, sink)

	r := mux.NewRouter()
	r.Path(listenPath).Headers("Content-type", "application/json").Handler(h)

	listenServer := ListenerServer{
		Keeper:   st,
		name:     name,
		listener: listener,
		server: http.Server{
			Handler:      r,
			Addr:         listenAddr,
			ReadTimeout:  clientTimeout,
			WriteTimeout: clientTimeout,
		},
	}

	go listenServer.server.Serve(listener)
	return &listenServer, nil
}

// SetupHandler is shared between signalfx and here to setup listening for collectd connections.
// Will do shared basic setup like configuring request counters
func SetupHandler(ctx context.Context, name string, sink dpsink.Sink) (*web.Handler, stats.Keeper) {
	metricTracking := reqcounter.RequestCounter{}
	counter := &dpsink.Counter{}
	collectdDecoder := JSONDecoder{
		SendTo: dpsink.FromChain(sink, dpsink.NextWrap(counter)),
	}
	h := web.NewHandler(ctx, &collectdDecoder).Add(web.NextHTTP(metricTracking.ServeHTTP))
	st := stats.ToKeeperMany(map[string]string{"listener": name, "type": "collectd"}, &metricTracking, counter, &collectdDecoder)
	return h, st
}
<|endoftext|>