<commit_before>package ccg

//go:generate myccg -u AstDecls.Filter -p ccg -o utils.go slice ast.Decl AstDecls
//go:generate myccg -u AstSpecs.Filter -p ccg -o utils.go slice ast.Spec AstSpecs
//go:generate myccg -p ccg -o utils.go set types.Object ObjectSet NewObjectSet

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/format"
	"go/token"
	"io"
	"strings"

	"golang.org/x/tools/go/loader"
	"golang.org/x/tools/go/types"
	"golang.org/x/tools/imports"
)

var (
	pt = fmt.Printf
)

type Config struct {
	From    string
	Params  map[string]string
	Renames map[string]string
	Writer  io.Writer
	Package string
	Decls   []ast.Decl
	FileSet *token.FileSet
	Uses    []string
}

func Copy(config Config) error {
	// load package
	loadConf := loader.Config{
		Fset: config.FileSet,
	}
	loadConf.Import(config.From)
	program, err := loadConf.Load()
	if err != nil {
		return fmt.Errorf("ccg: load package %v", err)
	}
	info := program.Imported[config.From]

	// remove param declarations
	for _, f := range info.Files {
		f.Decls = AstDecls(f.Decls).Filter(func(decl ast.Decl) (ret bool) {
			switch decl := decl.(type) {
			case *ast.GenDecl:
				switch decl.Tok {
				case token.TYPE:
					decl.Specs = AstSpecs(decl.Specs).Filter(func(spec ast.Spec) bool {
						name := spec.(*ast.TypeSpec).Name.Name
						_, exists := config.Params[name]
						return !exists
					})
					ret = len(decl.Specs) > 0
				case token.VAR:
					decl.Specs = AstSpecs(decl.Specs).Filter(func(sp ast.Spec) bool {
						spec := sp.(*ast.ValueSpec)
						names := []*ast.Ident{}
						values := []ast.Expr{}
						for i, name := range spec.Names {
							if _, exists := config.Params[name.Name]; !exists {
								names = append(names, name)
								if i < len(spec.Values) {
									values = append(values, spec.Values[i])
								}
							}
						}
						spec.Names = names
						if len(values) == 0 {
							spec.Values = nil
						} else {
							spec.Values = values
						}
						return len(spec.Names) > 0
					})
					ret = len(decl.Specs) > 0
				default:
					return true
				}
			default:
				return true
			}
			return
		})
	}

	// collect objects to rename
	renamed := map[string]string{}
	objects := make(map[types.Object]string)
	collectObjects := func(mapping map[string]string) error {
		for from, to := range mapping {
			obj := info.Pkg.Scope().Lookup(from)
			if obj == nil {
				return fmt.Errorf("ccg: name not found %s", from)
			}
			objects[obj] = to
			renamed[to] = from
		}
		return nil
	}
	if err := collectObjects(config.Params); err != nil {
		return err
	}
	if err := collectObjects(config.Renames); err != nil {
		return err
	}

	// rename
	rename := func(defs map[*ast.Ident]types.Object) {
		for id, obj := range defs {
			if to, ok := objects[obj]; ok {
				id.Name = to
			}
		}
	}
	rename(info.Defs)
	rename(info.Uses)

	// collect existing decls
	existingVars := make(map[string]func(expr ast.Expr))
	existingTypes := make(map[string]func(expr ast.Expr))
	existingFuncs := make(map[string]func(fn *ast.FuncDecl))
	var decls []ast.Decl
	for i, decl := range config.Decls {
		decls = append(decls, decl)
		switch decl := decl.(type) {
		case *ast.GenDecl:
			switch decl.Tok {
			case token.VAR:
				for _, spec := range decl.Specs {
					spec := spec.(*ast.ValueSpec)
					for i, name := range spec.Names {
						i := i
						spec := spec
						existingVars[name.Name] = func(expr ast.Expr) {
							spec.Values[i] = expr
						}
					}
				}
			case token.TYPE:
				for i, spec := range decl.Specs {
					spec := spec.(*ast.TypeSpec)
					i := i
					decl := decl
					existingTypes[spec.Name.Name] = func(expr ast.Expr) {
						decl.Specs[i].(*ast.TypeSpec).Type = expr
					}
				}
			}
		case *ast.FuncDecl:
			name := decl.Name.Name
			if decl.Recv != nil {
				name = decl.Recv.List[0].Type.(*ast.Ident).Name + "." + name
			}
			i := i
			existingFuncs[name] = func(fndecl *ast.FuncDecl) {
				decls[i] = fndecl
			}
		}
	}

	// collect output declarations
	var newDecls []ast.Decl
	for _, f := range info.Files {
		for _, decl := range f.Decls {
			switch decl := decl.(type) {
			case *ast.GenDecl:
				switch decl.Tok {
				// var
				case token.VAR:
					newDecl := &ast.GenDecl{
						Tok: token.VAR,
					}
					for _, spec := range decl.Specs {
						spec := spec.(*ast.ValueSpec)
						for i, name := range spec.Names {
							if mutator, ok := existingVars[name.Name]; ok {
								mutator(spec.Values[i])
							} else {
								newDecl.Specs = append(newDecl.Specs, spec)
							}
						}
					}
					if len(newDecl.Specs) > 0 {
						newDecls = append(newDecls, newDecl)
					}
				// type
				case token.TYPE:
					newDecl := &ast.GenDecl{
						Tok: token.TYPE,
					}
					for _, spec := range decl.Specs {
						spec := spec.(*ast.TypeSpec)
						name := spec.Name.Name
						if mutator, ok := existingTypes[name]; ok {
							mutator(spec.Type)
						} else {
							newDecl.Specs = append(newDecl.Specs, spec)
						}
					}
					if len(newDecl.Specs) > 0 {
						newDecls = append(newDecls, newDecl)
					}
				// import or const
				default:
					newDecls = append(newDecls, decl)
				}
			// func
			case *ast.FuncDecl:
				name := decl.Name.Name
				if decl.Recv != nil {
					name = decl.Recv.List[0].Type.(*ast.Ident).Name + "." + name
				}
				if mutator, ok := existingFuncs[name]; ok {
					mutator(decl)
				} else {
					newDecls = append(newDecls, decl)
				}
			}
		}
	}

	// filter by uses
	// get function dependencies
	deps := make(map[types.Object]ObjectSet)
	for _, decl := range newDecls {
		switch decl := decl.(type) {
		case *ast.FuncDecl:
			obj := info.ObjectOf(decl.Name)
			set := NewObjectSet()
			var visitor astVisitor
			visitor = func(node ast.Node) astVisitor {
				switch node := node.(type) {
				case *ast.Ident:
					dep := info.ObjectOf(node)
					set.Add(dep)
				}
				return visitor
			}
			ast.Walk(visitor, decl)
			deps[obj] = set
		}
	}

	// get uses objects
	uses := NewObjectSet()
	for _, use := range config.Uses {
		parts := strings.Split(use, ".")
		switch len(parts) {
		case 2: // method
			var ty types.Object
			if from, ok := renamed[parts[0]]; ok { // renamed type, use original type name
				ty = info.Pkg.Scope().Lookup(from)
			} else {
				ty = info.Pkg.Scope().Lookup(parts[0])
			}
			typeName, ok := ty.(*types.TypeName)
			if !ok {
				return fmt.Errorf("%s is not a type", parts[0])
			}
			obj, _, _ := types.LookupFieldOrMethod(typeName.Type(), true, info.Pkg, parts[1])
			uses.Add(obj)
		case 1: // non-method
			obj := info.Pkg.Scope().Lookup(parts[0])
			uses.Add(obj)
		default:
			return fmt.Errorf("invalid use spec: %s", use)
		}
	}

	// filter
	if len(uses) > 0 {
		// calculate uses closure
		for {
			l := len(uses)
			for use := range uses {
				if deps, ok := deps[use]; ok {
					for dep := range deps {
						uses.Add(dep)
					}
				}
			}
			if len(uses) == l {
				break
			}
		}
		newDecls = AstDecls(newDecls).Filter(func(decl ast.Decl) (ret bool) {
			switch decl := decl.(type) {
			case *ast.FuncDecl:
				obj := info.ObjectOf(decl.Name)
				return uses.In(obj)
			case *ast.GenDecl:
				switch decl.Tok {
				case token.TYPE:
					decl.Specs = AstSpecs(decl.Specs).Filter(func(spec ast.Spec) bool {
						obj := info.ObjectOf(spec.(*ast.TypeSpec).Name)
						return uses.In(obj)
					})
					ret = len(decl.Specs) > 0
				default:
					ret = true
				}
			default:
				ret = true
			}
			return
		})
	}

	// merge new and existing decls
	decls = append(decls, newDecls...)

	// decls tidy ups
	newDecls = newDecls[0:0]
	var importDecls []ast.Decl
	for _, decl := range decls {
		// ensure linebreak between decls
		switch decl := decl.(type) {
		case *ast.FuncDecl:
			if decl.Doc == nil {
				decl.Doc = new(ast.CommentGroup)
			}
		case *ast.GenDecl:
			if decl.Doc == nil {
				decl.Doc = new(ast.CommentGroup)
			}
		}
		// move import decls to beginning
		if decl, ok := decl.(*ast.GenDecl); ok && decl.Tok == token.IMPORT {
			importDecls = append(importDecls, decl)
			continue
		}
		newDecls = append(newDecls, decl)
	}
	decls = append(importDecls, newDecls...)

	// output
	if config.Writer != nil {
		if config.Package != "" { // output complete file
			file := &ast.File{
				Name:  ast.NewIdent(config.Package),
				Decls: decls,
			}
			buf := new(bytes.Buffer)
			err = format.Node(buf, program.Fset, file)
			if err != nil { //NOCOVER
				return fmt.Errorf("ccg: format output %v", err)
			}
			bs, err := imports.Process("", buf.Bytes(), nil)
			if err != nil { //NOCOVER
				return fmt.Errorf("ccg: format output %v", err)
			}
			config.Writer.Write(bs)
		} else { // output decls only
			err = format.Node(config.Writer, program.Fset, decls)
			if err != nil { //NOCOVER
				return fmt.Errorf("ccg: format output %v", err)
			}
		}
	}

	return nil
}

type astVisitor func(ast.Node) astVisitor

func (v astVisitor) Visit(node ast.Node) ast.Visitor {
	return v(node)
}
<commit_msg>handle renamed function correctly in Uses<commit_after>package ccg

//go:generate myccg -u AstDecls.Filter -p ccg -o utils.go slice ast.Decl AstDecls
//go:generate myccg -u AstSpecs.Filter -p ccg -o utils.go slice ast.Spec AstSpecs
//go:generate myccg -u ObjectSet.Add,ObjectSet.In,NewObjectSet -p ccg -o utils.go set types.Object ObjectSet NewObjectSet

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/format"
	"go/token"
	"io"
	"strings"

	"golang.org/x/tools/go/loader"
	"golang.org/x/tools/go/types"
	"golang.org/x/tools/imports"
)

var (
	pt = fmt.Printf
)

type Config struct {
	From    string
	Params  map[string]string
	Renames map[string]string
	Writer  io.Writer
	Package string
	Decls   []ast.Decl
	FileSet *token.FileSet
	Uses    []string
}

func Copy(config Config) error {
	// load package
	loadConf := loader.Config{
		Fset: config.FileSet,
	}
	loadConf.Import(config.From)
	program, err := loadConf.Load()
	if err != nil {
		return fmt.Errorf("ccg: load package %v", err)
	}
	info := program.Imported[config.From]

	// remove param declarations
	for _, f := range info.Files {
		f.Decls = AstDecls(f.Decls).Filter(func(decl ast.Decl) (ret bool) {
			switch decl := decl.(type) {
			case *ast.GenDecl:
				switch decl.Tok {
				case token.TYPE:
					decl.Specs = AstSpecs(decl.Specs).Filter(func(spec ast.Spec) bool {
						name := spec.(*ast.TypeSpec).Name.Name
						_, exists := config.Params[name]
						return !exists
					})
					ret = len(decl.Specs) > 0
				case token.VAR:
					decl.Specs = AstSpecs(decl.Specs).Filter(func(sp ast.Spec) bool {
						spec := sp.(*ast.ValueSpec)
						names := []*ast.Ident{}
						values := []ast.Expr{}
						for i, name := range spec.Names {
							if _, exists := config.Params[name.Name]; !exists {
								names = append(names, name)
								if i < len(spec.Values) {
									values = append(values, spec.Values[i])
								}
							}
						}
						spec.Names = names
						if len(values) == 0 {
							spec.Values = nil
						} else {
							spec.Values = values
						}
						return len(spec.Names) > 0
					})
					ret = len(decl.Specs) > 0
				default:
					return true
				}
			default:
				return true
			}
			return
		})
	}

	// collect objects to rename
	renamed := map[string]string{}
	objects := make(map[types.Object]string)
	collectObjects := func(mapping map[string]string) error {
		for from, to := range mapping {
			obj := info.Pkg.Scope().Lookup(from)
			if obj == nil {
				return fmt.Errorf("ccg: name not found %s", from)
			}
			objects[obj] = to
			renamed[to] = from
		}
		return nil
	}
	if err := collectObjects(config.Params); err != nil {
		return err
	}
	if err := collectObjects(config.Renames); err != nil {
		return err
	}

	// rename
	rename := func(defs map[*ast.Ident]types.Object) {
		for id, obj := range defs {
			if to, ok := objects[obj]; ok {
				id.Name = to
			}
		}
	}
	rename(info.Defs)
	rename(info.Uses)

	// collect existing decls
	existingVars := make(map[string]func(expr ast.Expr))
	existingTypes := make(map[string]func(expr ast.Expr))
	existingFuncs := make(map[string]func(fn *ast.FuncDecl))
	var decls []ast.Decl
	for i, decl := range config.Decls {
		decls = append(decls, decl)
		switch decl := decl.(type) {
		case *ast.GenDecl:
			switch decl.Tok {
			case token.VAR:
				for _, spec := range decl.Specs {
					spec := spec.(*ast.ValueSpec)
					for i, name := range spec.Names {
						i := i
						spec := spec
						existingVars[name.Name] = func(expr ast.Expr) {
							spec.Values[i] = expr
						}
					}
				}
			case token.TYPE:
				for i, spec := range decl.Specs {
					spec := spec.(*ast.TypeSpec)
					i := i
					decl := decl
					existingTypes[spec.Name.Name] = func(expr ast.Expr) {
						decl.Specs[i].(*ast.TypeSpec).Type = expr
					}
				}
			}
		case *ast.FuncDecl:
			name := decl.Name.Name
			if decl.Recv != nil {
				name = decl.Recv.List[0].Type.(*ast.Ident).Name + "." + name
			}
			i := i
			existingFuncs[name] = func(fndecl *ast.FuncDecl) {
				decls[i] = fndecl
			}
		}
	}

	// collect output declarations
	var newDecls []ast.Decl
	for _, f := range info.Files {
		for _, decl := range f.Decls {
			switch decl := decl.(type) {
			case *ast.GenDecl:
				switch decl.Tok {
				// var
				case token.VAR:
					newDecl := &ast.GenDecl{
						Tok: token.VAR,
					}
					for _, spec := range decl.Specs {
						spec := spec.(*ast.ValueSpec)
						for i, name := range spec.Names {
							if mutator, ok := existingVars[name.Name]; ok {
								mutator(spec.Values[i])
							} else {
								newDecl.Specs = append(newDecl.Specs, spec)
							}
						}
					}
					if len(newDecl.Specs) > 0 {
						newDecls = append(newDecls, newDecl)
					}
				// type
				case token.TYPE:
					newDecl := &ast.GenDecl{
						Tok: token.TYPE,
					}
					for _, spec := range decl.Specs {
						spec := spec.(*ast.TypeSpec)
						name := spec.Name.Name
						if mutator, ok := existingTypes[name]; ok {
							mutator(spec.Type)
						} else {
							newDecl.Specs = append(newDecl.Specs, spec)
						}
					}
					if len(newDecl.Specs) > 0 {
						newDecls = append(newDecls, newDecl)
					}
				// import or const
				default:
					newDecls = append(newDecls, decl)
				}
			// func
			case *ast.FuncDecl:
				name := decl.Name.Name
				if decl.Recv != nil {
					name = decl.Recv.List[0].Type.(*ast.Ident).Name + "." + name
				}
				if mutator, ok := existingFuncs[name]; ok {
					mutator(decl)
				} else {
					newDecls = append(newDecls, decl)
				}
			}
		}
	}

	// filter by uses
	// get function dependencies
	deps := make(map[types.Object]ObjectSet)
	for _, decl := range newDecls {
		switch decl := decl.(type) {
		case *ast.FuncDecl:
			obj := info.ObjectOf(decl.Name)
			set := NewObjectSet()
			var visitor astVisitor
			visitor = func(node ast.Node) astVisitor {
				switch node := node.(type) {
				case *ast.Ident:
					dep := info.ObjectOf(node)
					set.Add(dep)
				}
				return visitor
			}
			ast.Walk(visitor, decl)
			deps[obj] = set
		}
	}

	// get uses objects
	uses := NewObjectSet()
	for _, use := range config.Uses {
		parts := strings.Split(use, ".")
		switch len(parts) {
		case 2: // method
			var ty types.Object
			if from, ok := renamed[parts[0]]; ok { // renamed type, use original type name
				ty = info.Pkg.Scope().Lookup(from)
			} else {
				ty = info.Pkg.Scope().Lookup(parts[0])
			}
			typeName, ok := ty.(*types.TypeName)
			if !ok {
				return fmt.Errorf("%s is not a type", parts[0])
			}
			obj, _, _ := types.LookupFieldOrMethod(typeName.Type(), true, info.Pkg, parts[1])
			uses.Add(obj)
		case 1: // non-method
			var obj types.Object
			if from, ok := renamed[parts[0]]; ok { // renamed function
				obj = info.Pkg.Scope().Lookup(from)
			} else {
				obj = info.Pkg.Scope().Lookup(parts[0])
			}
			uses.Add(obj)
		default:
			return fmt.Errorf("invalid use spec: %s", use)
		}
	}

	// filter
	if len(uses) > 0 {
		// calculate uses closure
		for {
			l := len(uses)
			for use := range uses {
				if deps, ok := deps[use]; ok {
					for dep := range deps {
						uses.Add(dep)
					}
				}
			}
			if len(uses) == l {
				break
			}
		}
		newDecls = AstDecls(newDecls).Filter(func(decl ast.Decl) (ret bool) {
			switch decl := decl.(type) {
			case *ast.FuncDecl:
				obj := info.ObjectOf(decl.Name)
				return uses.In(obj)
			case *ast.GenDecl:
				switch decl.Tok {
				case token.TYPE:
					decl.Specs = AstSpecs(decl.Specs).Filter(func(spec ast.Spec) bool {
						obj := info.ObjectOf(spec.(*ast.TypeSpec).Name)
						return uses.In(obj)
					})
					ret = len(decl.Specs) > 0
				default:
					ret = true
				}
			default:
				ret = true
			}
			return
		})
	}

	// merge new and existing decls
	decls = append(decls, newDecls...)

	// decls tidy ups
	newDecls = newDecls[0:0]
	var importDecls []ast.Decl
	for _, decl := range decls {
		// ensure linebreak between decls
		switch decl := decl.(type) {
		case *ast.FuncDecl:
			if decl.Doc == nil {
				decl.Doc = new(ast.CommentGroup)
			}
		case *ast.GenDecl:
			if decl.Doc == nil {
				decl.Doc = new(ast.CommentGroup)
			}
		}
		// move import decls to beginning
		if decl, ok := decl.(*ast.GenDecl); ok && decl.Tok == token.IMPORT {
			importDecls = append(importDecls, decl)
			continue
		}
		newDecls = append(newDecls, decl)
	}
	decls = append(importDecls, newDecls...)

	// output
	if config.Writer != nil {
		if config.Package != "" { // output complete file
			file := &ast.File{
				Name:  ast.NewIdent(config.Package),
				Decls: decls,
			}
			buf := new(bytes.Buffer)
			err = format.Node(buf, program.Fset, file)
			if err != nil { //NOCOVER
				return fmt.Errorf("ccg: format output %v", err)
			}
			bs, err := imports.Process("", buf.Bytes(), nil)
			if err != nil { //NOCOVER
				return fmt.Errorf("ccg: format output %v", err)
			}
			config.Writer.Write(bs)
		} else { // output decls only
			err = format.Node(config.Writer, program.Fset, decls)
			if err != nil { //NOCOVER
				return fmt.Errorf("ccg: format output %v", err)
			}
		}
	}

	return nil
}

type astVisitor func(ast.Node) astVisitor

func (v astVisitor) Visit(node ast.Node) ast.Visitor {
	return v(node)
}
<|endoftext|>


<commit_before>// Based on https://github.com/olebedev/staticbin
package staticbin

import (
	"bytes"
	"log"
	"net/http"
	"path"
	"strings"
	"time"

	"github.com/labstack/echo"
	"github.com/labstack/echo/engine/standard"
)

type Options struct {
	// SkipLogging will disable [Static] log messages when a static file is served.
	SkipLogging bool
	// IndexFile defines which file to serve as index if it exists.
	IndexFile string
	// Path prefix
	Dir string
}

func (o *Options) init() {
	if o.IndexFile == "" {
		o.IndexFile = "index.html"
	}
}

// Static returns a middleware handler that serves static files in the given directory.
func Static(asset func(string) ([]byte, error), options ...Options) echo.MiddlewareFunc {
	if asset == nil {
		panic("asset is nil")
	}

	opt := Options{}
	for _, o := range options {
		opt = o
		break
	}
	opt.init()

	modtime := time.Now()

	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			request := c.Request().(*standard.Request).Request
			if request.Method != "GET" && request.Method != "HEAD" {
				// Request is not correct. Go farther.
				// return echo.NewHTTPError(http.StatusMethodNotAllowed)
				return next(c)
			}

			u := request.URL
			url := u.Path
			if !strings.HasPrefix(url, opt.Dir) {
				// return echo.NewHTTPError(http.StatusUnsupportedMediaType)
				return next(c)
			}
			file := strings.TrimPrefix(
				strings.TrimPrefix(url, opt.Dir),
				"/",
			)
			b, err := asset(file)

			if err != nil {
				// Try to serve the index file.
				b, err = asset(path.Join(file, opt.IndexFile))

				if err != nil {
					// Go farther if the asset could not be found.
					return next(c)
				}
			}

			if !opt.SkipLogging {
				log.Println("[Static] Serving " + url)
			}

			// http.ServeContent(c.Writer, c.Request(), url, modtime, bytes.NewReader(b))
			// c.Abort()

			response := c.Response().(*standard.Response).ResponseWriter
			http.ServeContent(response, request, url, modtime, bytes.NewReader(b))

			return nil
		}
	}
}
<commit_msg>staticbin: 去掉standard的约束<commit_after>// Based on https://github.com/olebedev/staticbin
package staticbin

import (
	"bytes"
	"log"
	"path"
	"strings"
	"time"

	"github.com/labstack/echo"
)

type Options struct {
	// SkipLogging will disable [Static] log messages when a static file is served.
	SkipLogging bool
	// IndexFile defines which file to serve as index if it exists.
	IndexFile string
	// Path prefix
	Dir string
}

func (o *Options) init() {
	if o.IndexFile == "" {
		o.IndexFile = "index.html"
	}
}

// Static returns a middleware handler that serves static files in the given directory.
func Static(asset func(string) ([]byte, error), options ...Options) echo.MiddlewareFunc {
	if asset == nil {
		panic("asset is nil")
	}

	opt := Options{}
	for _, o := range options {
		opt = o
		break
	}
	opt.init()

	modtime := time.Now()

	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			request := c.Request()
			method := request.Method()
			if method != "GET" && method != "HEAD" {
				// Request is not correct. Go farther.
				// return echo.NewHTTPError(http.StatusMethodNotAllowed)
				return next(c)
			}

			url := request.URL().Path()
			if !strings.HasPrefix(url, opt.Dir) {
				// return echo.NewHTTPError(http.StatusUnsupportedMediaType)
				return next(c)
			}
			file := strings.TrimPrefix(
				strings.TrimPrefix(url, opt.Dir),
				"/",
			)
			b, err := asset(file)

			if err != nil {
				// Try to serve the index file.
				b, err = asset(path.Join(file, opt.IndexFile))

				if err != nil {
					// Go farther if the asset could not be found.
					return next(c)
				}
			}

			if !opt.SkipLogging {
				log.Println("[Static] Serving " + url)
			}

			// response := c.Response()
			// http.ServeContent(c.Response().Writer(), request, url, modtime, bytes.NewReader(b))

			return c.ServeContent(bytes.NewReader(b), file, modtime)
		}
	}
}
<|endoftext|>


<commit_before>package gogrs

import (
	"encoding/csv"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
	"time"
)

// DailyData start with stock no, date.
type DailyData struct {
	No   string
	Date time.Time
}

// URL return stock csv url path.
func (d DailyData) URL() string {
	return fmt.Sprintf(TWSECSV, d.Date.Year(), d.Date.Month(), d.Date.Year(), d.Date.Month(), d.No, RandInt())
}

// Round will do sub one month.
func (d *DailyData) Round() {
	year, month, _ := d.Date.Date()
	d.Date = time.Date(year, month-1, 1, 0, 0, 0, 0, time.UTC)
}

// GetData return csv data in array.
func (d DailyData) GetData() ([][]string, error) {
	urlpath := fmt.Sprintf("%s%s", TWSEHOST, d.URL())
	csvFiles, err := http.Get(urlpath)
	if err != nil {
		fmt.Println("[err] >>> ", err)
		return nil, err
	}
	defer csvFiles.Body.Close()
	data, _ := ioutil.ReadAll(csvFiles.Body)
	csvArrayContent := strings.Split(string(data), "\n")
	for i := range csvArrayContent {
		csvArrayContent[i] = strings.TrimSpace(csvArrayContent[i])
	}
	if len(csvArrayContent) > 2 {
		csvReader := csv.NewReader(strings.NewReader(strings.Join(csvArrayContent[2:], "\n")))
		return csvReader.ReadAll()
	}
	return nil, errors.New("Not enough data.")
}
<commit_msg>Return full path in DailyData.URL<commit_after>package gogrs

import (
	"encoding/csv"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
	"time"
)

// DailyData start with stock no, date.
type DailyData struct {
	No   string
	Date time.Time
}

// URL return stock csv url path.
func (d DailyData) URL() string {
	path := fmt.Sprintf(TWSECSV, d.Date.Year(), d.Date.Month(), d.Date.Year(), d.Date.Month(), d.No, RandInt())
	return fmt.Sprintf("%s%s", TWSEHOST, path)
}

// Round will do sub one month.
func (d *DailyData) Round() {
	year, month, _ := d.Date.Date()
	d.Date = time.Date(year, month-1, 1, 0, 0, 0, 0, time.UTC)
}

// GetData return csv data in array.
func (d DailyData) GetData() ([][]string, error) {
	csvFiles, err := http.Get(d.URL())
	if err != nil {
		fmt.Println("[err] >>> ", err)
		return nil, err
	}
	defer csvFiles.Body.Close()
	data, _ := ioutil.ReadAll(csvFiles.Body)
	csvArrayContent := strings.Split(string(data), "\n")
	for i := range csvArrayContent {
		csvArrayContent[i] = strings.TrimSpace(csvArrayContent[i])
	}
	if len(csvArrayContent) > 2 {
		csvReader := csv.NewReader(strings.NewReader(strings.Join(csvArrayContent[2:], "\n")))
		return csvReader.ReadAll()
	}
	return nil, errors.New("Not enough data.")
}
<|endoftext|>


<commit_before>package storage

import (
	_ "github.com/lib/pq"
	"log"
	"database/sql"
	"fmt"
	"time"
)

const TABLE_USERS = `create table if not exists users
(
  id           serial,
  created_date timestamp default now(),
  chat_id      bigint unique     not null,
  access_token text unique     not null
);`

type Postgres struct {
	db *sql.DB
}

type User struct {
	id          int
	createdDate time.Time
	chatId      int
	accessToken string
}

func NewPostgres(url string) *Postgres {
	db, err := sql.Open("postgres", url)
	if err != nil {
		log.Fatal(err)
	}

	_, err = db.Exec(TABLE_USERS)
	if err != nil {
		log.Fatal(err)
	}

	return &Postgres{
		db,
	}
}

func (postgres *Postgres) Add(chatId int64, accessToken string) error {
	result, err := postgres.db.Exec("insert into users (chat_id, access_token) values ($1, $2)", chatId, accessToken)
	if err != nil {
		return err
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return err
	}

	fmt.Printf("User %d created successfully (%d row affected)\n", chatId, rowsAffected)

	return nil
}

func (postgres *Postgres) Get(chatId int64) (string, error) {
	stmt, err := postgres.db.Prepare("select u.access_token from users u where u.chat_id=$1")
	if err != nil {
		return "", err
	}

	defer stmt.Close()

	row := stmt.QueryRow(chatId)

	var accessToken string
	err = row.Scan(&accessToken)
	if err == sql.ErrNoRows {
		return "", nil
	} else {
		return "", err
	}

	return accessToken, nil
}

func (postgres *Postgres) Delete(chatId int64) error {
	result, err := postgres.db.Exec("delete from users where chat_id=$1", chatId)
	if err != nil {
		return err
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return err
	}

	fmt.Printf("User %d deleted successfully (%d row affected)\n", chatId, rowsAffected)

	return nil
}
<commit_msg>fix bug select access token from DB<commit_after>package storage

import (
	_ "github.com/lib/pq"
	"log"
	"database/sql"
	"fmt"
	"time"
)

const TABLE_USERS = `create table if not exists users
(
  id           serial,
  created_date timestamp default now(),
  chat_id      bigint unique     not null,
  access_token text unique     not null
);`

type Postgres struct {
	db *sql.DB
}

type User struct {
	id          int
	createdDate time.Time
	chatId      int
	accessToken string
}

func NewPostgres(url string) *Postgres {
	db, err := sql.Open("postgres", url)
	if err != nil {
		log.Fatal(err)
	}

	_, err = db.Exec(TABLE_USERS)
	if err != nil {
		log.Fatal(err)
	}

	return &Postgres{
		db,
	}
}

func (postgres *Postgres) Add(chatId int64, accessToken string) error {
	result, err := postgres.db.Exec("insert into users (chat_id, access_token) values ($1, $2)", chatId, accessToken)
	if err != nil {
		return err
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return err
	}

	fmt.Printf("User %d created successfully (%d row affected)\n", chatId, rowsAffected)

	return nil
}

func (postgres *Postgres) Get(chatId int64) (string, error) {
	stmt, err := postgres.db.Prepare("select u.access_token from users u where u.chat_id=$1")
	if err != nil {
		return "", err
	}

	defer stmt.Close()

	row := stmt.QueryRow(chatId)

	var accessToken string
	err = row.Scan(&accessToken)
	if err == sql.ErrNoRows {
		return "", nil
	} else if err != nil {
		return "", err
	}

	return accessToken, nil
}

func (postgres *Postgres) Delete(chatId int64) error {
	result, err := postgres.db.Exec("delete from users where chat_id=$1", chatId)
	if err != nil {
		return err
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return err
	}

	fmt.Printf("User %d deleted successfully (%d row affected)\n", chatId, rowsAffected)

	return nil
}
<|endoftext|>


<commit_before>package store

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"
	"time"

	sql "github.com/otoolep/rqlite/db"
)

type mockSnapshotSink struct {
	*os.File
}

func (m *mockSnapshotSink) ID() string {
	return "1"
}

func (m *mockSnapshotSink) Cancel() error {
	return nil
}

func Test_OpenStoreSingleNode(t *testing.T) {
	s := mustNewStore(true)
	defer os.RemoveAll(s.Path())

	if err := s.Open(true); err != nil {
		t.Fatalf("failed to open single-node store: %s", err.Error())
	}
}

func Test_OpenStoreCloseSingleNode(t *testing.T) {
	s := mustNewStore(true)
	defer os.RemoveAll(s.Path())

	if err := s.Open(true); err != nil {
		t.Fatalf("failed to open single-node store: %s", err.Error())
	}
	if err := s.Close(true); err != nil {
		t.Fatalf("failed to close single-node store: %s", err.Error())
	}
}

func Test_SingleNodeExecuteQuery(t *testing.T) {
	s := mustNewStore(true)
	defer os.RemoveAll(s.Path())

	if err := s.Open(true); err != nil {
		t.Fatalf("failed to open single-node store: %s", err.Error())
	}
	defer s.Close(true)
	s.WaitForLeader(10 * time.Second)

	queries := []string{
		`CREATE TABLE foo (id INTEGER NOT NULL PRIMARY KEY, name TEXT)`,
		`INSERT INTO foo(id, name) VALUES(1, "fiona")`,
	}
	_, err := s.Execute(queries, false, false)
	if err != nil {
		t.Fatalf("failed to execute on single node: %s", err.Error())
	}
	r, err := s.Query([]string{`SELECT * FROM foo`}, false, false, None)
	if err != nil {
		t.Fatalf("failed to query single node: %s", err.Error())
	}
	r, err = s.Query([]string{`SELECT * FROM foo`}, false, false, None)
	if err != nil {
		t.Fatalf("failed to query single node: %s", err.Error())
	}
	r, err = s.Query([]string{`SELECT * FROM foo`}, false, false, None)
	if err != nil {
		t.Fatalf("failed to query single node: %s", err.Error())
	}
	if exp, got := `["id","name"]`, asJSON(r[0].Columns); exp != got {
		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
	}
	if exp, got := `[[1,"fiona"]]`, asJSON(r[0].Values); exp != got {
		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
	}
}

func Test_SingleNodeExecuteQueryTx(t *testing.T) {
	s := mustNewStore(true)
	defer os.RemoveAll(s.Path())

	if err := s.Open(true); err != nil {
		t.Fatalf("failed to open single-node store: %s", err.Error())
	}
	defer s.Close(true)
	s.WaitForLeader(10 * time.Second)

	queries := []string{
		`CREATE TABLE foo (id INTEGER NOT NULL PRIMARY KEY, name TEXT)`,
		`INSERT INTO foo(id, name) VALUES(1, "fiona")`,
	}
	_, err := s.Execute(queries, false, true)
	if err != nil {
		t.Fatalf("failed to execute on single node: %s", err.Error())
	}
	r, err := s.Query([]string{`SELECT * FROM foo`}, false, true, None)
	if err != nil {
		t.Fatalf("failed to query single node: %s", err.Error())
	}
	r, err = s.Query([]string{`SELECT * FROM foo`}, false, true, Weak)
	if err != nil {
		t.Fatalf("failed to query single node: %s", err.Error())
	}
	r, err = s.Query([]string{`SELECT * FROM foo`}, false, true, Strong)
	if err != nil {
		t.Fatalf("failed to query single node: %s", err.Error())
	}
	if exp, got := `["id","name"]`, asJSON(r[0].Columns); exp != got {
		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
	}
	if exp, got := `[[1,"fiona"]]`, asJSON(r[0].Values); exp != got {
		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
	}
	_, err = s.Execute(queries, false, true)
	if err != nil {
		t.Fatalf("failed to execute on single node: %s", err.Error())
	}
}

func Test_MultiNodeExecuteQuery(t *testing.T) {
	s0 := mustNewStore(true)
	defer os.RemoveAll(s0.Path())
	if err := s0.Open(true); err != nil {
		t.Fatalf("failed to open node for multi-node test: %s", err.Error())
	}
	defer s0.Close(true)
	s0.WaitForLeader(10 * time.Second)

	s1 := mustNewStore(true)
	defer os.RemoveAll(s1.Path())
	if err := s1.Open(false); err != nil {
		t.Fatalf("failed to open node for multi-node test: %s", err.Error())
	}
	defer s1.Close(true)

	// Join the second node to the first.
	if err := s0.Join(s1.Addr().String()); err != nil {
		t.Fatalf("failed to join to node at %s: %s", s0.Addr().String(), err.Error())
	}

	queries := []string{
		`CREATE TABLE foo (id INTEGER NOT NULL PRIMARY KEY, name TEXT)`,
		`INSERT INTO foo(id, name) VALUES(1, "fiona")`,
	}
	_, err := s0.Execute(queries, false, false)
	if err != nil {
		t.Fatalf("failed to execute on single node: %s", err.Error())
	}
	r, err := s0.Query([]string{`SELECT * FROM foo`}, false, false, None)
	if err != nil {
		t.Fatalf("failed to query single node: %s", err.Error())
	}
	if exp, got := `["id","name"]`, asJSON(r[0].Columns); exp != got {
		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
	}
	if exp, got := `[[1,"fiona"]]`, asJSON(r[0].Values); exp != got {
		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
	}

	// Wait until the 3 log entries have been applied to the follower,
	// and then query.
	if err := s1.WaitForAppliedIndex(3, 5*time.Second); err != nil {
		t.Fatalf("error waiting for follower to apply index: %s:", err.Error())
	}
	r, err = s1.Query([]string{`SELECT * FROM foo`}, false, false, Weak)
	if err == nil {
		t.Fatalf("successfully queried non-leader node")
	}
	r, err = s1.Query([]string{`SELECT * FROM foo`}, false, false, Strong)
	if err == nil {
		t.Fatalf("successfully queried non-leader node")
	}
	r, err = s1.Query([]string{`SELECT * FROM foo`}, false, false, None)
	if err != nil {
		t.Fatalf("failed to query single node: %s", err.Error())
	}
	if exp, got := `["id","name"]`, asJSON(r[0].Columns); exp != got {
		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
	}
	if exp, got := `[[1,"fiona"]]`, asJSON(r[0].Values); exp != got {
		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
	}
}

func Test_SingleNodeSnapshot(t *testing.T) {
	s := mustNewStore(false)
	defer os.RemoveAll(s.Path())

	if err := s.Open(true); err != nil {
		t.Fatalf("failed to open single-node store: %s", err.Error())
	}
	defer s.Close(true)
	s.WaitForLeader(10 * time.Second)

	queries := []string{
		`CREATE TABLE foo (id INTEGER NOT NULL PRIMARY KEY, name TEXT)`,
		`INSERT INTO foo(id, name) VALUES(1, "fiona")`,
	}
	_, err := s.Execute(queries, false, false)
	if err != nil {
		t.Fatalf("failed to execute on single node: %s", err.Error())
	}
	_, err = s.Query([]string{`SELECT * FROM foo`}, false, false, None)
	if err != nil {
		t.Fatalf("failed to query single node: %s", err.Error())
	}

	// Snap the node and write to disk.
	f, err := s.Snapshot()
	if err != nil {
		t.Fatalf("failed to snapshot node: %s", err.Error())
	}

	snapDir := mustTempDir()
	defer os.RemoveAll(snapDir)
	snapFile, err := os.Create(filepath.Join(snapDir, "snapshot"))
	if err != nil {
		t.Fatalf("failed to create snapshot file: %s", err.Error())
	}
	sink := &mockSnapshotSink{snapFile}
	if err := f.Persist(sink); err != nil {
		t.Fatalf("failed to persist snapshot to disk: %s", err.Error())
	}

	// Check restoration.
	snapFile, err = os.Open(filepath.Join(snapDir, "snapshot"))
	if err != nil {
		t.Fatalf("failed to open snapshot file: %s", err.Error())
	}
	if err := s.Restore(snapFile); err != nil {
		t.Fatalf("failed to restore snapshot from disk: %s", err.Error())
	}

	// Ensure database is back in the correct state.
	r, err := s.Query([]string{`SELECT * FROM foo`}, false, false, None)
	if err != nil {
		t.Fatalf("failed to query single node: %s", err.Error())
	}
	if exp, got := `["id","name"]`, asJSON(r[0].Columns); exp != got {
		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
	}
	if exp, got := `[[1,"fiona"]]`, asJSON(r[0].Values); exp != got {
		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
	}
}

func mustNewStore(inmem bool) *Store {
	path := mustTempDir()
	defer os.RemoveAll(path)

	cfg := sql.NewConfig()
	cfg.Memory = inmem
	s := New(cfg, path, "localhost:0")
	if s == nil {
		panic("failed to create new store")
	}
	return s
}

func mustTempDir() string {
	var err error
	path, err := ioutil.TempDir("", "rqlilte-test-")
	if err != nil {
		panic("failed to create temp dir")
	}
	return path
}

func asJSON(v interface{}) string {
	b, err := json.Marshal(v)
	if err != nil {
		panic("failed to JSON marshal value")
	}
	return string(b)
}
<commit_msg>Explicitly test file and in-mem<commit_after>package store

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"
	"time"

	sql "github.com/otoolep/rqlite/db"
)

type mockSnapshotSink struct {
	*os.File
}

func (m *mockSnapshotSink) ID() string {
	return "1"
}

func (m *mockSnapshotSink) Cancel() error {
	return nil
}

func Test_OpenStoreSingleNode(t *testing.T) {
	s := mustNewStore(true)
	defer os.RemoveAll(s.Path())

	if err := s.Open(true); err != nil {
		t.Fatalf("failed to open single-node store: %s", err.Error())
	}
}

func Test_OpenStoreCloseSingleNode(t *testing.T) {
	s := mustNewStore(true)
	defer os.RemoveAll(s.Path())

	if err := s.Open(true); err != nil {
		t.Fatalf("failed to open single-node store: %s", err.Error())
	}
	if err := s.Close(true); err != nil {
		t.Fatalf("failed to close single-node store: %s", err.Error())
	}
}

func Test_SingleNodeInMemExecuteQuery(t *testing.T) {
	s := mustNewStore(true)
	defer os.RemoveAll(s.Path())

	if err := s.Open(true); err != nil {
		t.Fatalf("failed to open single-node store: %s", err.Error())
	}
	defer s.Close(true)
	s.WaitForLeader(10 * time.Second)

	queries := []string{
		`CREATE TABLE foo (id INTEGER NOT NULL PRIMARY KEY, name TEXT)`,
		`INSERT INTO foo(id, name) VALUES(1, "fiona")`,
	}
	_, err := s.Execute(queries, false, false)
	if err != nil {
		t.Fatalf("failed to execute on single node: %s", err.Error())
	}
	r, err := s.Query([]string{`SELECT * FROM foo`}, false, false, None)
	if err != nil {
		t.Fatalf("failed to query single node: %s", err.Error())
	}
	r, err = s.Query([]string{`SELECT * FROM foo`}, false, false, None)
	if err != nil {
		t.Fatalf("failed to query single node: %s", err.Error())
	}
	r, err = s.Query([]string{`SELECT * FROM foo`}, false, false, None)
	if err != nil {
		t.Fatalf("failed to query single node: %s", err.Error())
	}
	if exp, got := `["id","name"]`, asJSON(r[0].Columns); exp != got {
		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
	}
	if exp, got := `[[1,"fiona"]]`, asJSON(r[0].Values); exp != got {
		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
	}
}

func Test_SingleNodeFileExecuteQuery(t *testing.T) {
    s := mustNewStore(false)
    defer os.RemoveAll(s.Path())

    if err := s.Open(true); err != nil {
        t.Fatalf("failed to open single-node store: %s", err.Error())
    }
    defer s.Close(true)
    s.WaitForLeader(10 * time.Second)

    queries := []string{
        `CREATE TABLE foo (id INTEGER NOT NULL PRIMARY KEY, name TEXT)`,
        `INSERT INTO foo(id, name) VALUES(1, "fiona")`,
    }
    _, err := s.Execute(queries, false, false)
    if err != nil {
        t.Fatalf("failed to execute on single node: %s", err.Error())
    }
    r, err := s.Query([]string{`SELECT * FROM foo`}, false, false, None)
    if err != nil {
        t.Fatalf("failed to query single node: %s", err.Error())
    }
    r, err = s.Query([]string{`SELECT * FROM foo`}, false, false, None)
    if err != nil {
        t.Fatalf("failed to query single node: %s", err.Error())
    }
    r, err = s.Query([]string{`SELECT * FROM foo`}, false, false, None)
    if err != nil {
        t.Fatalf("failed to query single node: %s", err.Error())
    }
    if exp, got := `["id","name"]`, asJSON(r[0].Columns); exp != got {
        t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
    }
    if exp, got := `[[1,"fiona"]]`, asJSON(r[0].Values); exp != got {
        t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
    }
}

func Test_SingleNodeExecuteQueryTx(t *testing.T) {
	s := mustNewStore(true)
	defer os.RemoveAll(s.Path())

	if err := s.Open(true); err != nil {
		t.Fatalf("failed to open single-node store: %s", err.Error())
	}
	defer s.Close(true)
	s.WaitForLeader(10 * time.Second)

	queries := []string{
		`CREATE TABLE foo (id INTEGER NOT NULL PRIMARY KEY, name TEXT)`,
		`INSERT INTO foo(id, name) VALUES(1, "fiona")`,
	}
	_, err := s.Execute(queries, false, true)
	if err != nil {
		t.Fatalf("failed to execute on single node: %s", err.Error())
	}
	r, err := s.Query([]string{`SELECT * FROM foo`}, false, true, None)
	if err != nil {
		t.Fatalf("failed to query single node: %s", err.Error())
	}
	r, err = s.Query([]string{`SELECT * FROM foo`}, false, true, Weak)
	if err != nil {
		t.Fatalf("failed to query single node: %s", err.Error())
	}
	r, err = s.Query([]string{`SELECT * FROM foo`}, false, true, Strong)
	if err != nil {
		t.Fatalf("failed to query single node: %s", err.Error())
	}
	if exp, got := `["id","name"]`, asJSON(r[0].Columns); exp != got {
		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
	}
	if exp, got := `[[1,"fiona"]]`, asJSON(r[0].Values); exp != got {
		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
	}
	_, err = s.Execute(queries, false, true)
	if err != nil {
		t.Fatalf("failed to execute on single node: %s", err.Error())
	}
}

func Test_MultiNodeExecuteQuery(t *testing.T) {
	s0 := mustNewStore(true)
	defer os.RemoveAll(s0.Path())
	if err := s0.Open(true); err != nil {
		t.Fatalf("failed to open node for multi-node test: %s", err.Error())
	}
	defer s0.Close(true)
	s0.WaitForLeader(10 * time.Second)

	s1 := mustNewStore(true)
	defer os.RemoveAll(s1.Path())
	if err := s1.Open(false); err != nil {
		t.Fatalf("failed to open node for multi-node test: %s", err.Error())
	}
	defer s1.Close(true)

	// Join the second node to the first.
	if err := s0.Join(s1.Addr().String()); err != nil {
		t.Fatalf("failed to join to node at %s: %s", s0.Addr().String(), err.Error())
	}

	queries := []string{
		`CREATE TABLE foo (id INTEGER NOT NULL PRIMARY KEY, name TEXT)`,
		`INSERT INTO foo(id, name) VALUES(1, "fiona")`,
	}
	_, err := s0.Execute(queries, false, false)
	if err != nil {
		t.Fatalf("failed to execute on single node: %s", err.Error())
	}
	r, err := s0.Query([]string{`SELECT * FROM foo`}, false, false, None)
	if err != nil {
		t.Fatalf("failed to query single node: %s", err.Error())
	}
	if exp, got := `["id","name"]`, asJSON(r[0].Columns); exp != got {
		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
	}
	if exp, got := `[[1,"fiona"]]`, asJSON(r[0].Values); exp != got {
		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
	}

	// Wait until the 3 log entries have been applied to the follower,
	// and then query.
	if err := s1.WaitForAppliedIndex(3, 5*time.Second); err != nil {
		t.Fatalf("error waiting for follower to apply index: %s:", err.Error())
	}
	r, err = s1.Query([]string{`SELECT * FROM foo`}, false, false, Weak)
	if err == nil {
		t.Fatalf("successfully queried non-leader node")
	}
	r, err = s1.Query([]string{`SELECT * FROM foo`}, false, false, Strong)
	if err == nil {
		t.Fatalf("successfully queried non-leader node")
	}
	r, err = s1.Query([]string{`SELECT * FROM foo`}, false, false, None)
	if err != nil {
		t.Fatalf("failed to query single node: %s", err.Error())
	}
	if exp, got := `["id","name"]`, asJSON(r[0].Columns); exp != got {
		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
	}
	if exp, got := `[[1,"fiona"]]`, asJSON(r[0].Values); exp != got {
		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
	}
}

func Test_SingleNodeSnapshot(t *testing.T) {
	s := mustNewStore(false)
	defer os.RemoveAll(s.Path())

	if err := s.Open(true); err != nil {
		t.Fatalf("failed to open single-node store: %s", err.Error())
	}
	defer s.Close(true)
	s.WaitForLeader(10 * time.Second)

	queries := []string{
		`CREATE TABLE foo (id INTEGER NOT NULL PRIMARY KEY, name TEXT)`,
		`INSERT INTO foo(id, name) VALUES(1, "fiona")`,
	}
	_, err := s.Execute(queries, false, false)
	if err != nil {
		t.Fatalf("failed to execute on single node: %s", err.Error())
	}
	_, err = s.Query([]string{`SELECT * FROM foo`}, false, false, None)
	if err != nil {
		t.Fatalf("failed to query single node: %s", err.Error())
	}

	// Snap the node and write to disk.
	f, err := s.Snapshot()
	if err != nil {
		t.Fatalf("failed to snapshot node: %s", err.Error())
	}

	snapDir := mustTempDir()
	defer os.RemoveAll(snapDir)
	snapFile, err := os.Create(filepath.Join(snapDir, "snapshot"))
	if err != nil {
		t.Fatalf("failed to create snapshot file: %s", err.Error())
	}
	sink := &mockSnapshotSink{snapFile}
	if err := f.Persist(sink); err != nil {
		t.Fatalf("failed to persist snapshot to disk: %s", err.Error())
	}

	// Check restoration.
	snapFile, err = os.Open(filepath.Join(snapDir, "snapshot"))
	if err != nil {
		t.Fatalf("failed to open snapshot file: %s", err.Error())
	}
	if err := s.Restore(snapFile); err != nil {
		t.Fatalf("failed to restore snapshot from disk: %s", err.Error())
	}

	// Ensure database is back in the correct state.
	r, err := s.Query([]string{`SELECT * FROM foo`}, false, false, None)
	if err != nil {
		t.Fatalf("failed to query single node: %s", err.Error())
	}
	if exp, got := `["id","name"]`, asJSON(r[0].Columns); exp != got {
		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
	}
	if exp, got := `[[1,"fiona"]]`, asJSON(r[0].Values); exp != got {
		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
	}
}

func mustNewStore(inmem bool) *Store {
	path := mustTempDir()
	defer os.RemoveAll(path)

	cfg := sql.NewConfig()
	cfg.Memory = inmem
	s := New(cfg, path, "localhost:0")
	if s == nil {
		panic("failed to create new store")
	}
	return s
}

func mustTempDir() string {
	var err error
	path, err := ioutil.TempDir("", "rqlilte-test-")
	if err != nil {
		panic("failed to create temp dir")
	}
	return path
}

func asJSON(v interface{}) string {
	b, err := json.Marshal(v)
	if err != nil {
		panic("failed to JSON marshal value")
	}
	return string(b)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/gorilla/websocket"
	"github.com/qlm-iot/qlm/df"
	"github.com/qlm-iot/qlm/mi"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"strings"
)

func httpserverconnector(address string, sendPtr, receivePtr *chan []byte){
	send := *sendPtr
	receive := *receivePtr
	for {
		select {
		case raw_msg := <-send:
			msg := string(raw_msg)
			data := url.Values{}
			data.Set("msg", msg)
			response, err := http.PostForm(address, data)
			if err == nil{
				defer response.Body.Close()
				content, err := ioutil.ReadAll(response.Body)
				if err == nil{
					receive <- content
				}else{
					receive <- []byte(err.Error())
				}
			}else{
				receive <- []byte(err.Error())
			}
		}
	}
}
func wsServerConnector(address string, sendPtr, receivePtr *chan []byte){
	send := *sendPtr
	receive := *receivePtr
	for {
		select {
		case raw_msg := <-send:
			var h http.Header

			conn, _, err := websocket.DefaultDialer.Dial(address, h)
			if err == nil{
				if err := conn.WriteMessage(websocket.BinaryMessage, raw_msg); err != nil {
					receive <- []byte(err.Error())
				}
				_, content, err := conn.ReadMessage()
				if err == nil {
					receive <- content
				}else{
					receive <- []byte(err.Error())
				}
			}else{
				receive <- []byte(err.Error())
			}
		}
	}
}
func createServerConnection(address string, send, receive *chan []byte) bool{
	if strings.HasPrefix(address, "http://"){
		go httpserverconnector(address, send, receive)
	}else if strings.HasPrefix(address, "ws://"){
		go wsServerConnector(address, send, receive)
	}else{
		return false
	}
	return true
}
/*
Usage
cli [--server http://localhost/qlm/] test
cli [--server http://localhost/qlm/] read id name
cli [--server http://localhost/qlm/] write id name value
cli [--server http://localhost/qlm/] order id name interval
cli [--server http://localhost/qlm/] order-get req_id
cli [--server http://localhost/qlm/] order-cancel req_id
*/

func main() {
	var receive chan []byte
	var send chan []byte

	send = make(chan []byte)
	receive = make(chan []byte)

	var address string
	flag.StringVar(&address, "server", "http://localhost/qlm/", "Server address")

	flag.Parse()

	if !createServerConnection(address, &send, &receive){
		fmt.Println("Unsupported server protocol")
		return
	}

	command := flag.Arg(0)
	switch command {
		case "test": {
			send <- createEmptyReadRequest()
		}
		case "read": {
			id := flag.Arg(1)
			name := flag.Arg(2)
			send <- createReadRequest(id, name)
		}
		case "write":{
			id := flag.Arg(1)
			name := flag.Arg(2)
			value := flag.Arg(3)
			send <- createWriteRequest(id, name, value)
		}
		case "order": {
			id := flag.Arg(1)
			name := flag.Arg(2)
			interval, _ := strconv.ParseFloat(flag.Arg(3), 32)
			send <- createSubscriptionRequest(id, name, interval)
		}
		case "order-get": {
			requestId := flag.Arg(1)
			send <- createReadSubscriptionRequest(requestId)
		}
		case "order-cancel": {
			requestId := flag.Arg(1)
			send <- createCancelSubscriptionRequest(requestId)
		}
		default: {
			fmt.Println("Unknown command")
			return
		}
	}

	msg := <-receive
	fmt.Println(string(msg))
}

func createEmptyReadRequest() []byte{
	ret, _ := mi.Marshal(mi.QlmEnvelope{
		Version: "1.0",
		Ttl:     -1,
		Read:    &mi.ReadRequest{},
	})
	return ret
}

func createQLMMessage(id, name string) string{
	objects := df.Objects{
		Objects: []df.Object{
			df.Object{
				Id:   &df.QLMID{Text: id},
				InfoItems: []df.InfoItem{
					df.InfoItem{
						Name: name,
					},
				},
			},
		},
	}
	data, _ := df.Marshal(objects)
	return (string)(data)
}

func createQLMMessageWithValue(id, name, value string) string{
	objects := df.Objects{
		Objects: []df.Object{
			df.Object{
				Id:   &df.QLMID{Text: id},
				InfoItems: []df.InfoItem{
					df.InfoItem{
						Name: name,
						Values: []df.Value{
							df.Value{
								Text: value,
							},
						},
					},
				},
			},
		},
	}
	data, _ := df.Marshal(objects)
	return (string)(data)
}

func createReadRequest(id, name string) []byte{
	ret, _ := mi.Marshal(mi.QlmEnvelope{
		Version: "1.0",
		Ttl:     -1,
		Read: &mi.ReadRequest{
			MsgFormat:  "QLMdf",
			Message:    &mi.Message{
				Data: createQLMMessage(id, name),
			},
		},
	})
	return ret
}

func createSubscriptionRequest(id, name string, interval float64) []byte{
	ret, _ := mi.Marshal(mi.QlmEnvelope{
		Version: "1.0",
		Ttl:     -1,
		Read: &mi.ReadRequest{
			MsgFormat:  "QLMdf",
			Interval:   interval,
			Message:    &mi.Message{
				Data: createQLMMessage(id, name),
			},
		},
	})
	return ret
}

func createReadSubscriptionRequest(requestId string) []byte{
	ret, _ := mi.Marshal(mi.QlmEnvelope{
		Version: "1.0",
		Ttl:     -1,
		Read: &mi.ReadRequest{
			MsgFormat:  "QLMdf",
			RequestIds: []mi.Id{
				mi.Id{Text: requestId},
			},
		},
	})
	return ret
}

func createCancelSubscriptionRequest(requestId string) []byte{
	ret, _ := mi.Marshal(mi.QlmEnvelope{
		Version: "1.0",
		Ttl:     -1,
		Cancel:  &mi.CancelRequest{
			RequestIds: []mi.Id{
				mi.Id{Text: requestId},
			},
		},
	})
	return ret
}

func createWriteRequest(id, name, value string) []byte{
	ret, _ := mi.Marshal(mi.QlmEnvelope{
		Version: "1.0",
		Ttl:     -1,
		Write:   &mi.WriteRequest{
			MsgFormat:  "QLMdf",
			TargetType: "device",
			Message: &mi.Message{
				Data: createQLMMessageWithValue(id, name, value),
			},
		},
	})
	return ret
}
<commit_msg>Fix indentation problems<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/gorilla/websocket"
	"github.com/qlm-iot/qlm/df"
	"github.com/qlm-iot/qlm/mi"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"strings"
)

func httpserverconnector(address string, sendPtr, receivePtr *chan []byte) {
	send := *sendPtr
	receive := *receivePtr
	for {
		select {
		case raw_msg := <-send:
			msg := string(raw_msg)
			data := url.Values{}
			data.Set("msg", msg)
			response, err := http.PostForm(address, data)
			if err == nil {
				defer response.Body.Close()
				content, err := ioutil.ReadAll(response.Body)
				if err == nil {
					receive <- content
				} else {
					receive <- []byte(err.Error())
				}
			} else {
				receive <- []byte(err.Error())
			}
		}
	}
}
func wsServerConnector(address string, sendPtr, receivePtr *chan []byte) {
	send := *sendPtr
	receive := *receivePtr
	for {
		select {
		case raw_msg := <-send:
			var h http.Header

			conn, _, err := websocket.DefaultDialer.Dial(address, h)
			if err == nil {
				if err := conn.WriteMessage(websocket.BinaryMessage, raw_msg); err != nil {
					receive <- []byte(err.Error())
				}
				_, content, err := conn.ReadMessage()
				if err == nil {
					receive <- content
				} else {
					receive <- []byte(err.Error())
				}
			} else {
				receive <- []byte(err.Error())
			}
		}
	}
}
func createServerConnection(address string, send, receive *chan []byte) bool {
	if strings.HasPrefix(address, "http://") {
		go httpserverconnector(address, send, receive)
	} else if strings.HasPrefix(address, "ws://") {
		go wsServerConnector(address, send, receive)
	} else {
		return false
	}
	return true
}

/*
Usage
cli [--server http://localhost/qlm/] test
cli [--server http://localhost/qlm/] read id name
cli [--server http://localhost/qlm/] write id name value
cli [--server http://localhost/qlm/] order id name interval
cli [--server http://localhost/qlm/] order-get req_id
cli [--server http://localhost/qlm/] order-cancel req_id
*/

func main() {
	var receive chan []byte
	var send chan []byte

	send = make(chan []byte)
	receive = make(chan []byte)

	var address string
	flag.StringVar(&address, "server", "http://localhost/qlm/", "Server address")

	flag.Parse()

	if !createServerConnection(address, &send, &receive) {
		fmt.Println("Unsupported server protocol")
		return
	}

	command := flag.Arg(0)
	switch command {
	case "test":
		{
			send <- createEmptyReadRequest()
		}
	case "read":
		{
			id := flag.Arg(1)
			name := flag.Arg(2)
			send <- createReadRequest(id, name)
		}
	case "write":
		{
			id := flag.Arg(1)
			name := flag.Arg(2)
			value := flag.Arg(3)
			send <- createWriteRequest(id, name, value)
		}
	case "order":
		{
			id := flag.Arg(1)
			name := flag.Arg(2)
			interval, _ := strconv.ParseFloat(flag.Arg(3), 32)
			send <- createSubscriptionRequest(id, name, interval)
		}
	case "order-get":
		{
			requestId := flag.Arg(1)
			send <- createReadSubscriptionRequest(requestId)
		}
	case "order-cancel":
		{
			requestId := flag.Arg(1)
			send <- createCancelSubscriptionRequest(requestId)
		}
	default:
		{
			fmt.Println("Unknown command")
			return
		}
	}

	msg := <-receive
	fmt.Println(string(msg))
}

func createEmptyReadRequest() []byte {
	ret, _ := mi.Marshal(mi.QlmEnvelope{
		Version: "1.0",
		Ttl:     -1,
		Read:    &mi.ReadRequest{},
	})
	return ret
}

func createQLMMessage(id, name string) string {
	objects := df.Objects{
		Objects: []df.Object{
			df.Object{
				Id: &df.QLMID{Text: id},
				InfoItems: []df.InfoItem{
					df.InfoItem{
						Name: name,
					},
				},
			},
		},
	}
	data, _ := df.Marshal(objects)
	return (string)(data)
}

func createQLMMessageWithValue(id, name, value string) string {
	objects := df.Objects{
		Objects: []df.Object{
			df.Object{
				Id: &df.QLMID{Text: id},
				InfoItems: []df.InfoItem{
					df.InfoItem{
						Name: name,
						Values: []df.Value{
							df.Value{
								Text: value,
							},
						},
					},
				},
			},
		},
	}
	data, _ := df.Marshal(objects)
	return (string)(data)
}

func createReadRequest(id, name string) []byte {
	ret, _ := mi.Marshal(mi.QlmEnvelope{
		Version: "1.0",
		Ttl:     -1,
		Read: &mi.ReadRequest{
			MsgFormat: "QLMdf",
			Message: &mi.Message{
				Data: createQLMMessage(id, name),
			},
		},
	})
	return ret
}

func createSubscriptionRequest(id, name string, interval float64) []byte {
	ret, _ := mi.Marshal(mi.QlmEnvelope{
		Version: "1.0",
		Ttl:     -1,
		Read: &mi.ReadRequest{
			MsgFormat: "QLMdf",
			Interval:  interval,
			Message: &mi.Message{
				Data: createQLMMessage(id, name),
			},
		},
	})
	return ret
}

func createReadSubscriptionRequest(requestId string) []byte {
	ret, _ := mi.Marshal(mi.QlmEnvelope{
		Version: "1.0",
		Ttl:     -1,
		Read: &mi.ReadRequest{
			MsgFormat: "QLMdf",
			RequestIds: []mi.Id{
				mi.Id{Text: requestId},
			},
		},
	})
	return ret
}

func createCancelSubscriptionRequest(requestId string) []byte {
	ret, _ := mi.Marshal(mi.QlmEnvelope{
		Version: "1.0",
		Ttl:     -1,
		Cancel: &mi.CancelRequest{
			RequestIds: []mi.Id{
				mi.Id{Text: requestId},
			},
		},
	})
	return ret
}

func createWriteRequest(id, name, value string) []byte {
	ret, _ := mi.Marshal(mi.QlmEnvelope{
		Version: "1.0",
		Ttl:     -1,
		Write: &mi.WriteRequest{
			MsgFormat:  "QLMdf",
			TargetType: "device",
			Message: &mi.Message{
				Data: createQLMMessageWithValue(id, name, value),
			},
		},
	})
	return ret
}
<|endoftext|>


<commit_before>package controllers

import (
	"fmt"
	"net/http"

	"getmelange.com/app/framework"
	"getmelange.com/app/models"
	"getmelange.com/app/packaging"
	"getmelange.com/dap"
	"getmelange.com/router"

	"airdispat.ch/identity"
	"airdispat.ch/routing"

	gdb "github.com/huntaub/go-db"
)

// Profile is a JSONObject specifying a request to create a new identity
// (or profile, I suppose).
type identityRequest struct {
	// Profile Information
	FirstName string `json:"first"`
	LastName  string `json:"last"`
	About     string `json:"about"`
	Password  string `json:"password"`
	// AD Information
	Server  string `json:"server"`
	Tracker string `json:"tracker"`
	Alias   string `json:"alias"`
	// Identity Nickname
	Nickname string `json:"nickname"`
}

// SaveIdentity will create, register, alias, and save a new Identity.
type SaveIdentity struct {
	Tables   map[string]gdb.Table
	Store    *models.Store
	Packager *packaging.Packager
}

// Handle performs the specified functions.
func (i *SaveIdentity) Handle(req *http.Request) framework.View {
	// Decode Body
	profileRequest := &identityRequest{}
	err := DecodeJSONBody(req, &profileRequest)

	if err != nil {
		fmt.Println("Error occured while decoding body:", err)
		return framework.Error500
	}

	// Create Identity
	id, err := identity.CreateIdentity()
	if err != nil {
		fmt.Println("Error occured creating an identity:", err)
		return framework.Error500
	}

	//
	// Server Registration
	//

	// Extract Keys
	server, err := i.Packager.ServerFromId(profileRequest.Server)
	if err != nil {
		fmt.Println("Error occured getting server:", err)
		return &framework.HTTPError{
			ErrorCode: 500,
			Message:   "Couldn't get server for id.",
		}
	}

	id.SetLocation(server.URL)

	// Run Registration
	client := &dap.Client{
		Key:    id,
		Server: server.Key,
	}
	err = client.Register(map[string][]byte{
		"name": []byte(profileRequest.FirstName + " " + profileRequest.LastName),
	})
	if err != nil {
		fmt.Println("Error occurred registering on Server", err)
		return framework.Error500
	}

	//
	// Tracker Registration
	//

	tracker, err := i.Packager.TrackerFromId(profileRequest.Tracker)
	if err != nil {
		fmt.Println("Error occured getting tracker:", err)
		return &framework.HTTPError{
			ErrorCode: 500,
			Message:   "Couldn't get tracker for id.",
		}
	}

	err = (&router.Router{
		Origin: id,
		TrackerList: []string{
			tracker.URL,
		},
	}).Register(id, profileRequest.Alias, map[string]routing.Redirect{
		string(routing.LookupTypeTX): routing.Redirect{
			Alias:       server.Alias,
			Fingerprint: server.Fingerprint,
			Type:        routing.LookupTypeTX,
		},
	})

	if err != nil {
		fmt.Println("Error occurred registering on Tracker", err)
		return framework.Error500
	}

	//
	// Database Registration
	//

	modelID, err := models.CreateIdentityFromDispatch(id, "")
	if err != nil {
		fmt.Println("Error occurred encoding Id", err)
		return framework.Error500
	}

	modelID.Nickname = profileRequest.Nickname

	// Load Server Information
	modelID.Server = server.URL
	modelID.ServerKey = server.EncryptionKey
	modelID.ServerFingerprint = server.Fingerprint
	modelID.ServerAlias = server.Alias

	_, err = i.Tables["identity"].Insert(modelID).Exec(i.Store)
	if err != nil {
		fmt.Println("Error saving Identity", err)
		return framework.Error500
	}

	modelAlias := &models.Alias{
		Identity: gdb.ForeignKey(modelID),
		Location: tracker.URL,
		Username: profileRequest.Alias,
	}

	_, err = i.Tables["alias"].Insert(modelAlias).Exec(i.Store)
	if err != nil {
		fmt.Println("Error saving Alias", err)
		return framework.Error500
	}

	// Save as the current identity.
	err = i.Store.Set("current_identity", id.Address.String())
	if err != nil {
		fmt.Println("Error storing current_identity.", err)
		return framework.Error500
	}

	return &framework.JSONView{
		Content: map[string]interface{}{
			"error": false,
		},
	}
}
<commit_msg>Fixed Issue where New Identity didn't Invalidate Cache<commit_after>package controllers

import (
	"fmt"
	"net/http"

	"getmelange.com/app/framework"
	"getmelange.com/app/models"
	"getmelange.com/app/packaging"
	"getmelange.com/dap"
	"getmelange.com/router"

	"airdispat.ch/identity"
	"airdispat.ch/routing"

	gdb "github.com/huntaub/go-db"
)

// Profile is a JSONObject specifying a request to create a new identity
// (or profile, I suppose).
type identityRequest struct {
	// Profile Information
	FirstName string `json:"first"`
	LastName  string `json:"last"`
	About     string `json:"about"`
	Password  string `json:"password"`
	// AD Information
	Server  string `json:"server"`
	Tracker string `json:"tracker"`
	Alias   string `json:"alias"`
	// Identity Nickname
	Nickname string `json:"nickname"`
}

// SaveIdentity will create, register, alias, and save a new Identity.
type SaveIdentity struct {
	Tables   map[string]gdb.Table
	Store    *models.Store
	Packager *packaging.Packager
}

// Handle performs the specified functions.
func (i *SaveIdentity) Handle(req *http.Request) framework.View {
	// Decode Body
	profileRequest := &identityRequest{}
	err := DecodeJSONBody(req, &profileRequest)

	if err != nil {
		fmt.Println("Error occured while decoding body:", err)
		return framework.Error500
	}

	// Create Identity
	id, err := identity.CreateIdentity()
	if err != nil {
		fmt.Println("Error occured creating an identity:", err)
		return framework.Error500
	}

	//
	// Server Registration
	//

	// Extract Keys
	server, err := i.Packager.ServerFromId(profileRequest.Server)
	if err != nil {
		fmt.Println("Error occured getting server:", err)
		return &framework.HTTPError{
			ErrorCode: 500,
			Message:   "Couldn't get server for id.",
		}
	}

	id.SetLocation(server.URL)

	// Run Registration
	client := &dap.Client{
		Key:    id,
		Server: server.Key,
	}
	err = client.Register(map[string][]byte{
		"name": []byte(profileRequest.FirstName + " " + profileRequest.LastName),
	})
	if err != nil {
		fmt.Println("Error occurred registering on Server", err)
		return framework.Error500
	}

	//
	// Tracker Registration
	//

	tracker, err := i.Packager.TrackerFromId(profileRequest.Tracker)
	if err != nil {
		fmt.Println("Error occured getting tracker:", err)
		return &framework.HTTPError{
			ErrorCode: 500,
			Message:   "Couldn't get tracker for id.",
		}
	}

	err = (&router.Router{
		Origin: id,
		TrackerList: []string{
			tracker.URL,
		},
	}).Register(id, profileRequest.Alias, map[string]routing.Redirect{
		string(routing.LookupTypeTX): routing.Redirect{
			Alias:       server.Alias,
			Fingerprint: server.Fingerprint,
			Type:        routing.LookupTypeTX,
		},
	})

	if err != nil {
		fmt.Println("Error occurred registering on Tracker", err)
		return framework.Error500
	}

	//
	// Database Registration
	//

	modelID, err := models.CreateIdentityFromDispatch(id, "")
	if err != nil {
		fmt.Println("Error occurred encoding Id", err)
		return framework.Error500
	}

	modelID.Nickname = profileRequest.Nickname

	// Load Server Information
	modelID.Server = server.URL
	modelID.ServerKey = server.EncryptionKey
	modelID.ServerFingerprint = server.Fingerprint
	modelID.ServerAlias = server.Alias

	_, err = i.Tables["identity"].Insert(modelID).Exec(i.Store)
	if err != nil {
		fmt.Println("Error saving Identity", err)
		return framework.Error500
	}

	modelAlias := &models.Alias{
		Identity: gdb.ForeignKey(modelID),
		Location: tracker.URL,
		Username: profileRequest.Alias,
	}

	_, err = i.Tables["alias"].Insert(modelAlias).Exec(i.Store)
	if err != nil {
		fmt.Println("Error saving Alias", err)
		return framework.Error500
	}

	// Save as the current identity.
	err = i.Store.Set("current_identity", id.Address.String())
	if err != nil {
		fmt.Println("Error storing current_identity.", err)
		return framework.Error500
	}

	models.InvalidateCaches()

	return &framework.JSONView{
		Content: map[string]interface{}{
			"error": false,
		},
	}
}
<|endoftext|>


<commit_before>package pointslicepool

import (
	"sync"

	"github.com/grafana/metrictank/schema"
)

// default size is probably bigger than what most responses need, but it saves [re]allocations
// also it's possible that occasionally more size is needed, causing a realloc of underlying array, and that extra space will stick around until next GC run.
const DefaultPointSliceSize = 2000

type PointSlicePool struct {
	defaultSize int
	p           sync.Pool
}

func New(defaultSize int) *PointSlicePool {
	return &PointSlicePool{
		defaultSize: defaultSize,
		p: sync.Pool{
			New: func() interface{} {
				return make([]schema.Point, 0, defaultSize)
			},
		},
	}
}

func (p *PointSlicePool) Put(s []schema.Point) {
	p.p.Put(s[:0])
}

func (p *PointSlicePool) Get() []schema.Point {
	return p.p.Get().([]schema.Point)
}

// GetMin returns a pointslice that has at least minCap capacity
func (p *PointSlicePool) GetMin(minCap int) []schema.Point {
	candidate := p.Get()
	if cap(candidate) >= minCap {
		return candidate
	}
	p.p.Put(candidate)
	if minCap > p.defaultSize {
		return make([]schema.Point, 0, minCap)
	}
	// even if our caller needs a smaller cap now, we expect they will put it back in the pool
	// so it can later be reused.
	// may as well allocate a size now that we expect will be more useful down the road.
	return make([]schema.Point, 0, p.defaultSize)
}
<commit_msg>avoid possibly allocating a too small, garbage pointslice<commit_after>package pointslicepool

import (
	"sync"

	"github.com/grafana/metrictank/schema"
)

// default size is probably bigger than what most responses need, but it saves [re]allocations
// also it's possible that occasionally more size is needed, causing a realloc of underlying array, and that extra space will stick around until next GC run.
const DefaultPointSliceSize = 2000

type PointSlicePool struct {
	defaultSize int
	p           sync.Pool
}

func New(defaultSize int) *PointSlicePool {
	return &PointSlicePool{
		defaultSize: defaultSize,
		p:           sync.Pool{},
	}
}

func (p *PointSlicePool) Put(s []schema.Point) {
	p.p.Put(s[:0])
}

func (p *PointSlicePool) Get() []schema.Point {
	return p.GetMin(p.defaultSize)
}

// GetMin returns a pointslice that has at least minCap capacity
func (p *PointSlicePool) GetMin(minCap int) []schema.Point {
	candidate, ok := p.p.Get().([]schema.Point)
	if ok {
		if cap(candidate) >= minCap {
			return candidate
		}
		p.p.Put(candidate)
	}
	if minCap > p.defaultSize {
		return make([]schema.Point, 0, minCap)
	}
	// even if our caller needs a smaller cap now, we expect they will put it back in the pool
	// so it can later be reused.
	// may as well allocate a size now that we expect will be more useful down the road.
	return make([]schema.Point, 0, p.defaultSize)
}
<|endoftext|>


<commit_before>package v4

import (
	"github.com/giantswarm/versionbundle"
)

func VersionBundle() versionbundle.Bundle {
	return versionbundle.Bundle{
		Changelogs: []versionbundle.Changelog{
			{
				Component:   "cloudconfig",
				Description: "Kubernetes (hyperkube) updated with version 1.11.1.",
				Kind:        versionbundle.KindChanged,
			},
		},
		Components: []versionbundle.Component{
			{
				Name:    "calico",
				Version: "3.0.8",
			},
			{
				Name:    "containerlinux",
				Version: "1745.7.0",
			},
			{
				Name:    "docker",
				Version: "18.03.1",
			},
			{
				Name:    "etcd",
				Version: "3.3.3",
			},
			{
				Name:    "coredns",
				Version: "1.1.1",
			},
			{
				Name:    "kubernetes",
				Version: "1.11.1",
			},
		},
		Name:    "azure-operator",
		Version: "1.2.0",
	}
}
<commit_msg>Bump v4 WIP to 2.0.0 (#330)<commit_after>package v4

import (
	"github.com/giantswarm/versionbundle"
)

func VersionBundle() versionbundle.Bundle {
	return versionbundle.Bundle{
		Changelogs: []versionbundle.Changelog{
			{
				Component:   "kubernetes",
				Description: "Updated to 1.11.1.",
				Kind:        versionbundle.KindChanged,
			},
		},
		Components: []versionbundle.Component{
			{
				Name:    "calico",
				Version: "3.0.8",
			},
			{
				Name:    "containerlinux",
				Version: "1745.7.0",
			},
			{
				Name:    "docker",
				Version: "18.03.1",
			},
			{
				Name:    "etcd",
				Version: "3.3.3",
			},
			{
				Name:    "coredns",
				Version: "1.1.1",
			},
			{
				Name:    "kubernetes",
				Version: "1.11.1",
			},
		},
		Name:    "azure-operator",
		Version: "2.0.0",
	}
}
<|endoftext|>


<commit_before>package executor

import (
	"io"
	"os/exec"
	"sync"
)

// Cmd is the full configuration for a command-line executable.
type Cmd struct {
	Stdout      Log
	Stderr      Log
	Directory   string
	SetUID      bool
	UID         int
	SetGID      bool
	GID         int
	Arguments   []string
	Environment map[string]string
}

// Command creates a new Cmd with generic settings given a set of command-line
// arguments.
func Command(arguments ...string) *Cmd {
	res := new(Cmd)
	res.Arguments = arguments
	res.Stdout = NullLog
	res.Stderr = NullLog
	return res
}

// Clone creates a copy of a Cmd.
// While it does do a completey copy of the Arguments and Environment fields, it
// cannot copy the Logs.
func (c *Cmd) Clone() *Cmd {
	x := *c
	cpy := &x
	cpy.Arguments = make([]string, len(c.Arguments))
	for i, arg := range c.Arguments {
		cpy.Arguments[i] = arg
	}
	cpy.Environment = map[string]string{}
	for key, val := range c.Environment {
		cpy.Environment[key] = val
	}
	return cpy
}

// ToJob creates a Job based on the current configuration in a Cmd.
// If the receiver is modified after a call to ToJob(), the job will not be
// modified.
func (c *Cmd) ToJob() Job {
	return &cmdJob{sync.Mutex{}, c.Clone(), nil}
}

func (c *Cmd) toExecCmd() (*exec.Cmd, error) {
	task := exec.Command(c.Arguments[0], c.Arguments[1:]...)
	for key, value := range c.Environment {
		task.Env = append(task.Env, key+"="+value)
	}

	// TODO: here, set UID and GID

	task.Dir = c.Directory

	// Create output streams
	var err error
	if task.Stdout, err = c.Stdout.Open(); err != nil {
		return nil, err
	}
	if task.Stderr, err = c.Stderr.Open(); err != nil {
		task.Stdout.(io.Closer).Close()
		return nil, err
	}

	return task, nil
}

type cmdJob struct {
	mutex   sync.Mutex
	command *Cmd
	execCmd *exec.Cmd
}

func (c *cmdJob) Start() error {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	// Make sure the job is not already running.
	if c.execCmd != nil {
		return ErrAlreadyRunning
	}

	// Generate the exec.Cmd
	cmd, err := c.command.toExecCmd()
	if err != nil {
		return err
	}

	// Start the command or return an error
	if err := cmd.Start(); err != nil {
		cmd.Stdout.(io.Closer).Close()
		cmd.Stderr.(io.Closer).Close()
		return err
	}

	c.execCmd = cmd
	return nil
}

func (c *cmdJob) Stop() error {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	if c.execCmd == nil {
		return ErrNotRunning
	}
	c.execCmd.Process.Kill()
	return nil
}

func (c *cmdJob) Wait() error {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	if c.execCmd == nil {
		return ErrNotRunning
	}

	c.mutex.Unlock()
	res := c.execCmd.Wait()
	c.mutex.Lock()

	c.execCmd.Stdout.(io.Closer).Close()
	c.execCmd.Stderr.(io.Closer).Close()
	c.execCmd = nil
	return res
}
<commit_msg>added documentation to Cmd fields<commit_after>package executor

import (
	"io"
	"os/exec"
	"sync"
)

// Cmd is the full configuration for a command-line executable.
type Cmd struct {
	// Stdout stores the standard output configuration.
	Stdout Log

	// Stderr stores the standard error configuration.
	Stderr Log

	// Directory is the working directory for the command
	Directory string

	// SetUID specifies whether or not the UID field should be used.
	SetUID bool

	// UID is the UID to run the command under.
	UID int

	// SetGID specifies whether or not the GID field should be used.
	SetGID bool

	// GID is the GID to run the command under.
	GID int

	// Arguments is the command-line arguments for the command.
	Arguments []string

	// Environment is a mapping of environment variables for the command.
	Environment map[string]string
}

// Command creates a new Cmd with generic settings given a set of command-line
// arguments.
func Command(arguments ...string) *Cmd {
	res := new(Cmd)
	res.Arguments = arguments
	res.Stdout = NullLog
	res.Stderr = NullLog
	return res
}

// Clone creates a copy of a Cmd.
// While it does do a completey copy of the Arguments and Environment fields, it
// cannot copy the Logs.
func (c *Cmd) Clone() *Cmd {
	x := *c
	cpy := &x
	cpy.Arguments = make([]string, len(c.Arguments))
	for i, arg := range c.Arguments {
		cpy.Arguments[i] = arg
	}
	cpy.Environment = map[string]string{}
	for key, val := range c.Environment {
		cpy.Environment[key] = val
	}
	return cpy
}

// ToJob creates a Job based on the current configuration in a Cmd.
// If the receiver is modified after a call to ToJob(), the job will not be
// modified.
func (c *Cmd) ToJob() Job {
	return &cmdJob{sync.Mutex{}, c.Clone(), nil}
}

func (c *Cmd) toExecCmd() (*exec.Cmd, error) {
	task := exec.Command(c.Arguments[0], c.Arguments[1:]...)
	for key, value := range c.Environment {
		task.Env = append(task.Env, key+"="+value)
	}

	// TODO: here, set UID and GID

	task.Dir = c.Directory

	// Create output streams
	var err error
	if task.Stdout, err = c.Stdout.Open(); err != nil {
		return nil, err
	}
	if task.Stderr, err = c.Stderr.Open(); err != nil {
		task.Stdout.(io.Closer).Close()
		return nil, err
	}

	return task, nil
}

type cmdJob struct {
	mutex   sync.Mutex
	command *Cmd
	execCmd *exec.Cmd
}

func (c *cmdJob) Start() error {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	// Make sure the job is not already running.
	if c.execCmd != nil {
		return ErrAlreadyRunning
	}

	// Generate the exec.Cmd
	cmd, err := c.command.toExecCmd()
	if err != nil {
		return err
	}

	// Start the command or return an error
	if err := cmd.Start(); err != nil {
		cmd.Stdout.(io.Closer).Close()
		cmd.Stderr.(io.Closer).Close()
		return err
	}

	c.execCmd = cmd
	return nil
}

func (c *cmdJob) Stop() error {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	if c.execCmd == nil {
		return ErrNotRunning
	}
	c.execCmd.Process.Kill()
	return nil
}

func (c *cmdJob) Wait() error {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	if c.execCmd == nil {
		return ErrNotRunning
	}

	c.mutex.Unlock()
	res := c.execCmd.Wait()
	c.mutex.Lock()

	c.execCmd.Stdout.(io.Closer).Close()
	c.execCmd.Stderr.(io.Closer).Close()
	c.execCmd = nil
	return res
}
<|endoftext|>


<commit_before>package main

import (
	// "bufio"
	"bytes"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"github.com/nu7hatch/gouuid"
	"io/ioutil"
	"net/url"
	"os"
	"os/exec"
	"strings"
	"time"
)

// @author Robin Verlangen

type Cmd struct {
	Command       string   // Commands to execute
	Pending       bool     // Did we dispatch it to the client?
	Id            string   // Unique ID for this command
	ClientId      string   // Client ID on which the command is executed
	TemplateId    string   // Reference to the template id
	Signature     string   // makes this only valid from the server to the client based on the preshared token and this is a signature with the command and id
	Timeout       int      // in seconds
	State         string   // Textual representation of the current state, e.g. finished, failed, etc.
	RequestUserId string   // User ID of the user that initiated this command
	Created       int64    // Unix timestamp created
	BufOutput     []string // Standard output
	BufOutputErr  []string // Error output
}

// Sign the command on the server
func (c *Cmd) Sign(client *RegisteredClient) {
	c.Signature = c.ComputeHmac(client.AuthToken)
}

// Set local state
func (c *Cmd) SetState(state string) {
	// Old state for change detection
	oldState := c.State

	// Update
	c.State = state

	// Run validation
	if oldState == "finished_execution" && c.State == "flushed_logs" {
		c._validate()
	}
}

// Validate the execution of a command, only on the server
func (c *Cmd) _validate() {
	// Only on the server
	if conf.IsServer == false {
		return
	}

	// Get template
	template := server.templateStore.Get(c.TemplateId)
	if template == nil {
		log.Printf("Unable to find template %s for validation of cmd %s", c.TemplateId, c.Id)
		return
	}

	// Iterate and run on templates
	var failedValidation = false
	for _, v := range template.ValidationRules {
		// Select stream
		var stream []string
		if v.OutputStream == 1 {
			stream = c.BufOutput
		} else {
			stream = c.BufOutputErr
		}

		// Match on line
		var matched bool = false
		for _, line := range stream {
			if strings.Contains(line, v.Text) {
				matched = true
				break
			}
		}

		// Did we match?
		if v.MustContain == true && matched == false {
			// Should BE there, but is NOT
			c.SetState("failed_validation")
			failedValidation = true
			break
		} else if v.MustContain == false && matched == true {
			// Should NOT be there, but IS
			c.SetState("failed_validation")
			failedValidation = true
			break
		}
	}

	// Done and passed validation
	if failedValidation == false {
		c.SetState("finished")
		log.Printf("Validation passed for %s", c.Id)
	}
}

// Notify state to server
func (c *Cmd) NotifyServer(state string) {
	// Update local client state
	c.SetState(state)

	// Update server state
	client._req("PUT", fmt.Sprintf("client/%s/cmd/%s/state?state=%s", url.QueryEscape(client.Id), url.QueryEscape(c.Id), url.QueryEscape(state)), nil)
}

// Should we flush the local buffer? After X milliseconds or Y lines
func (c *Cmd) _checkFlushLogs() {
	// At least 10 lines
	if len(c.BufOutput) > 10 || len(c.BufOutputErr) > 10 {
		c._flushLogs()
	}
}

// Write logs to server
func (c *Cmd) _flushLogs() {
	// To JSON
	m := make(map[string][]string)
	m["output"] = c.BufOutput
	m["error"] = c.BufOutputErr
	bytes, je := json.Marshal(m)
	if je != nil {
		log.Printf("Failed to convert logs to JSON: %s", je)
		return
	}

	// Post to server
	uri := fmt.Sprintf("client/%s/cmd/%s/logs", url.QueryEscape(client.Id), url.QueryEscape(c.Id))
	b, e := client._req("PUT", uri, bytes)
	if e != nil || len(b) < 1 {
		log.Printf("Failed log write: %s", e)
	}

	// Clear buffers
	c.BufOutput = make([]string, 0)
	c.BufOutputErr = make([]string, 0)
}

// Log output
func (c *Cmd) LogOutput(line string) {
	// No lock, only one routine can access this

	// Append
	c.BufOutput = append(c.BufOutput, line)

	// Check to flush?
	c._checkFlushLogs()
}

// Log error
func (c *Cmd) LogError(line string) {
	// No lock, only one routine can access this

	// Append
	c.BufOutputErr = append(c.BufOutputErr, line)

	// Check to flush?
	c._checkFlushLogs()
}

// Sign the command
func (c *Cmd) ComputeHmac(token string) string {
	bytes, be := base64.URLEncoding.DecodeString(token)
	if be != nil {
		return ""
	}
	mac := hmac.New(sha256.New, bytes)
	mac.Write([]byte(c.Command))
	mac.Write([]byte(c.Id))
	sum := mac.Sum(nil)
	return base64.URLEncoding.EncodeToString(sum)
}

// Execute command on the client
func (c *Cmd) Execute(client *Client) {
	log.Printf("Executing %s: %s", c.Id, c.Command)

	// Validate HMAC
	c.NotifyServer("validating")
	if client != nil {
		// Compute mac
		expectedMac := c.ComputeHmac(client.AuthToken)
		if expectedMac != c.Signature || len(c.Signature) < 1 {
			c.NotifyServer("invalid_signature")
			log.Printf("ERROR! Invalid command signature, communication between server and client might be tampered with")
			return
		}
	} else {
		log.Printf("Executing insecure command, unable to validate HMAC of %s", c.Id)
	}

	// Start
	c.NotifyServer("starting")

	// File contents
	var fileBytes bytes.Buffer
	fileBytes.WriteString("#!/bin/bash\n")
	fileBytes.WriteString(c.Command)

	// Write tmp file
	tmpFileName := fmt.Sprintf("/tmp/indispenso_%s", c.Id)
	ioutil.WriteFile(tmpFileName, fileBytes.Bytes(), 0644)

	// Remove file once done
	defer os.Remove(tmpFileName)

	// Run file
	cmd := exec.Command("bash", tmpFileName)
	var out bytes.Buffer
	var outerr bytes.Buffer
	cmd.Stdout = &out
	cmd.Stderr = &outerr

	// Consume streams
	// go func() {
	// 	p, pe := cmd.StdoutPipe()
	// 	if pe != nil {
	// 		log.Printf("Pipe error: %s", pe)
	// 		return
	// 	}
	// 	scanner := bufio.NewScanner(p)
	// 	for scanner.Scan() {
	// 		txt := scanner.Text()
	// 		c.LogOutput(txt)
	// 		if debug {
	// 			log.Println(scanner.Text())
	// 		}
	// 	}
	// 	if err := scanner.Err(); err != nil {
	// 		fmt.Fprintln(os.Stderr, "reading standard input:", err)
	// 	}
	// }()
	// go func() {
	// 	p, pe := cmd.StderrPipe()
	// 	if pe != nil {
	// 		log.Printf("Pipe error: %s", pe)
	// 		return
	// 	}
	// 	scanner := bufio.NewScanner(p)
	// 	for scanner.Scan() {
	// 		txt := scanner.Text()
	// 		c.LogError(txt)
	// 		if debug {
	// 			log.Println(scanner.Text())
	// 		}
	// 	}
	// 	if err := scanner.Err(); err != nil {
	// 		fmt.Fprintln(os.Stderr, "reading standard input:", err)
	// 	}
	// }()

	// Start
	err := cmd.Start()
	if err != nil {
		c.NotifyServer("failed_execution")
		log.Printf("Failed to start command: %s", err)
		return
	}
	c.NotifyServer("started_execution")

	// Timeout mechanism
	done := make(chan error, 1)
	go func() {
		done <- cmd.Wait()
	}()
	select {
	case <-time.After(time.Duration(c.Timeout) * time.Second):
		if err := cmd.Process.Kill(); err != nil {
			log.Printf("Failed to kill %s: %s", c.Id, err)
			return
		}
		<-done // allow goroutine to exit
		c.NotifyServer("killed_execution")
		log.Printf("Process %s killed", c.Id)
	case err := <-done:
		if err != nil {
			c.NotifyServer("failed_execution")
			log.Printf("Process %s done with error = %v", c.Id, err)
		} else {
			c.NotifyServer("finished_execution")
			log.Printf("Finished %s", c.Id)
		}
	}

	// Logs
	for _, line := range strings.Split(out.String(), "\n") {
		c.LogOutput(line)
	}
	for _, line := range strings.Split(outerr.String(), "\n") {
		c.LogError(line)
	}
	// Final flush
	c._flushLogs()
	c.NotifyServer("flushed_logs")
}

func newCmd(command string, timeout int) *Cmd {
	// Default timeout if not valid
	if timeout < 1 {
		timeout = DEFAULT_COMMAND_TIMEOUT
	}

	// Id
	id, _ := uuid.NewV4()

	// Create instance
	return &Cmd{
		Id:           id.String(),
		Command:      command,
		Pending:      true,
		Timeout:      timeout,
		Created:      time.Now().Unix(),
		BufOutput:    make([]string, 0),
		BufOutputErr: make([]string, 0),
	}
}
<commit_msg>Fix nil<commit_after>package main

import (
	// "bufio"
	"bytes"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"github.com/nu7hatch/gouuid"
	"io/ioutil"
	"net/url"
	"os"
	"os/exec"
	"strings"
	"time"
)

// @author Robin Verlangen

type Cmd struct {
	Command       string   // Commands to execute
	Pending       bool     // Did we dispatch it to the client?
	Id            string   // Unique ID for this command
	ClientId      string   // Client ID on which the command is executed
	TemplateId    string   // Reference to the template id
	Signature     string   // makes this only valid from the server to the client based on the preshared token and this is a signature with the command and id
	Timeout       int      // in seconds
	State         string   // Textual representation of the current state, e.g. finished, failed, etc.
	RequestUserId string   // User ID of the user that initiated this command
	Created       int64    // Unix timestamp created
	BufOutput     []string // Standard output
	BufOutputErr  []string // Error output
}

// Sign the command on the server
func (c *Cmd) Sign(client *RegisteredClient) {
	c.Signature = c.ComputeHmac(client.AuthToken)
}

// Set local state
func (c *Cmd) SetState(state string) {
	// Old state for change detection
	oldState := c.State

	// Update
	c.State = state

	// Run validation
	if oldState == "finished_execution" && c.State == "flushed_logs" {
		c._validate()
	}
}

// Validate the execution of a command, only on the server
func (c *Cmd) _validate() {
	// Only on the server
	if conf.IsServer == false {
		return
	}

	// Get template
	template := server.templateStore.Get(c.TemplateId)
	if template == nil {
		log.Printf("Unable to find template %s for validation of cmd %s", c.TemplateId, c.Id)
		return
	}

	// Iterate and run on templates
	var failedValidation = false
	for _, v := range template.ValidationRules {
		// Select stream
		var stream []string
		if v.OutputStream == 1 {
			stream = c.BufOutput
		} else {
			stream = c.BufOutputErr
		}

		// Match on line
		var matched bool = false
		for _, line := range stream {
			if strings.Contains(line, v.Text) {
				matched = true
				break
			}
		}

		// Did we match?
		if v.MustContain == true && matched == false {
			// Should BE there, but is NOT
			c.SetState("failed_validation")
			failedValidation = true
			break
		} else if v.MustContain == false && matched == true {
			// Should NOT be there, but IS
			c.SetState("failed_validation")
			failedValidation = true
			break
		}
	}

	// Done and passed validation
	if failedValidation == false {
		c.SetState("finished")
		log.Printf("Validation passed for %s", c.Id)
	}
}

// Notify state to server
func (c *Cmd) NotifyServer(state string) {
	// Update local client state
	c.SetState(state)

	// Update server state, only if this has a signature, else it is local
	if len(c.Signature) > 0 {
		client._req("PUT", fmt.Sprintf("client/%s/cmd/%s/state?state=%s", url.QueryEscape(client.Id), url.QueryEscape(c.Id), url.QueryEscape(state)), nil)
	}
}

// Should we flush the local buffer? After X milliseconds or Y lines
func (c *Cmd) _checkFlushLogs() {
	// At least 10 lines
	if len(c.BufOutput) > 10 || len(c.BufOutputErr) > 10 {
		c._flushLogs()
	}
}

// Write logs to server
func (c *Cmd) _flushLogs() {
	// To JSON
	m := make(map[string][]string)
	m["output"] = c.BufOutput
	m["error"] = c.BufOutputErr
	bytes, je := json.Marshal(m)
	if je != nil {
		log.Printf("Failed to convert logs to JSON: %s", je)
		return
	}

	// Post to server
	uri := fmt.Sprintf("client/%s/cmd/%s/logs", url.QueryEscape(client.Id), url.QueryEscape(c.Id))
	b, e := client._req("PUT", uri, bytes)
	if e != nil || len(b) < 1 {
		log.Printf("Failed log write: %s", e)
	}

	// Clear buffers
	c.BufOutput = make([]string, 0)
	c.BufOutputErr = make([]string, 0)
}

// Log output
func (c *Cmd) LogOutput(line string) {
	// No lock, only one routine can access this

	// Append
	c.BufOutput = append(c.BufOutput, line)

	// Check to flush?
	c._checkFlushLogs()
}

// Log error
func (c *Cmd) LogError(line string) {
	// No lock, only one routine can access this

	// Append
	c.BufOutputErr = append(c.BufOutputErr, line)

	// Check to flush?
	c._checkFlushLogs()
}

// Sign the command
func (c *Cmd) ComputeHmac(token string) string {
	bytes, be := base64.URLEncoding.DecodeString(token)
	if be != nil {
		return ""
	}
	mac := hmac.New(sha256.New, bytes)
	mac.Write([]byte(c.Command))
	mac.Write([]byte(c.Id))
	sum := mac.Sum(nil)
	return base64.URLEncoding.EncodeToString(sum)
}

// Execute command on the client
func (c *Cmd) Execute(client *Client) {
	log.Printf("Executing %s: %s", c.Id, c.Command)

	// Validate HMAC
	c.NotifyServer("validating")
	if client != nil {
		// Compute mac
		expectedMac := c.ComputeHmac(client.AuthToken)
		if expectedMac != c.Signature || len(c.Signature) < 1 {
			c.NotifyServer("invalid_signature")
			log.Printf("ERROR! Invalid command signature, communication between server and client might be tampered with")
			return
		}
	} else {
		log.Printf("Executing insecure command, unable to validate HMAC of %s", c.Id)
	}

	// Start
	c.NotifyServer("starting")

	// File contents
	var fileBytes bytes.Buffer
	fileBytes.WriteString("#!/bin/bash\n")
	fileBytes.WriteString(c.Command)

	// Write tmp file
	tmpFileName := fmt.Sprintf("/tmp/indispenso_%s", c.Id)
	ioutil.WriteFile(tmpFileName, fileBytes.Bytes(), 0644)

	// Remove file once done
	defer os.Remove(tmpFileName)

	// Run file
	cmd := exec.Command("bash", tmpFileName)
	var out bytes.Buffer
	var outerr bytes.Buffer
	cmd.Stdout = &out
	cmd.Stderr = &outerr

	// Consume streams
	// go func() {
	// 	p, pe := cmd.StdoutPipe()
	// 	if pe != nil {
	// 		log.Printf("Pipe error: %s", pe)
	// 		return
	// 	}
	// 	scanner := bufio.NewScanner(p)
	// 	for scanner.Scan() {
	// 		txt := scanner.Text()
	// 		c.LogOutput(txt)
	// 		if debug {
	// 			log.Println(scanner.Text())
	// 		}
	// 	}
	// 	if err := scanner.Err(); err != nil {
	// 		fmt.Fprintln(os.Stderr, "reading standard input:", err)
	// 	}
	// }()
	// go func() {
	// 	p, pe := cmd.StderrPipe()
	// 	if pe != nil {
	// 		log.Printf("Pipe error: %s", pe)
	// 		return
	// 	}
	// 	scanner := bufio.NewScanner(p)
	// 	for scanner.Scan() {
	// 		txt := scanner.Text()
	// 		c.LogError(txt)
	// 		if debug {
	// 			log.Println(scanner.Text())
	// 		}
	// 	}
	// 	if err := scanner.Err(); err != nil {
	// 		fmt.Fprintln(os.Stderr, "reading standard input:", err)
	// 	}
	// }()

	// Start
	err := cmd.Start()
	if err != nil {
		c.NotifyServer("failed_execution")
		log.Printf("Failed to start command: %s", err)
		return
	}
	c.NotifyServer("started_execution")

	// Timeout mechanism
	done := make(chan error, 1)
	go func() {
		done <- cmd.Wait()
	}()
	select {
	case <-time.After(time.Duration(c.Timeout) * time.Second):
		if err := cmd.Process.Kill(); err != nil {
			log.Printf("Failed to kill %s: %s", c.Id, err)
			return
		}
		<-done // allow goroutine to exit
		c.NotifyServer("killed_execution")
		log.Printf("Process %s killed", c.Id)
	case err := <-done:
		if err != nil {
			c.NotifyServer("failed_execution")
			log.Printf("Process %s done with error = %v", c.Id, err)
		} else {
			c.NotifyServer("finished_execution")
			log.Printf("Finished %s", c.Id)
		}
	}

	// Logs
	for _, line := range strings.Split(out.String(), "\n") {
		c.LogOutput(line)
	}
	for _, line := range strings.Split(outerr.String(), "\n") {
		c.LogError(line)
	}
	// Final flush
	c._flushLogs()
	c.NotifyServer("flushed_logs")
}

func newCmd(command string, timeout int) *Cmd {
	// Default timeout if not valid
	if timeout < 1 {
		timeout = DEFAULT_COMMAND_TIMEOUT
	}

	// Id
	id, _ := uuid.NewV4()

	// Create instance
	return &Cmd{
		Id:           id.String(),
		Command:      command,
		Pending:      true,
		Timeout:      timeout,
		Created:      time.Now().Unix(),
		BufOutput:    make([]string, 0),
		BufOutputErr: make([]string, 0),
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"strings"
)

func putCmd(cfg *Config, args []string) {
	if len(args) != 1 {
		PrintUsage("put")
	}

	var (
		err error
		val map[string]interface{}
	)

	// Decoded provided argument other read from stdin.
	if len(args) == 2 {
		err = json.Unmarshal([]byte(args[1]), &val)
	} else {
		err = json.NewDecoder(os.Stdin).Decode(&val)
	}

	defer cfg.Mongo.Close()
	o, err := Put(cfg, args[0], val)

	if err != nil {
		log.Fatal(err)
	}

	if o == nil {
		return
	}

	b, err := json.MarshalIndent(o, "", "  ")

	if err != nil {
		log.Fatal(err)
	}

	fmt.Fprintf(os.Stdout, "%s\n", b)
}

func getCmd(cfg *Config, args []string) {
	var (
		v  int
		ts string
	)

	fs := flag.NewFlagSet("get", flag.ExitOnError)

	fs.IntVar(&v, "version", 0, "Specific revision to get.")
	fs.StringVar(&ts, "time", "", "Returns the object as of the specified time.")

	fs.Parse(args)

	args = fs.Args()

	if len(args) != 1 {
		PrintUsage("get")
	}

	t, err := ParseTimeString(ts)

	if v > 0 && t > 0 {
		fmt.Println("error: version and time are mutually exclusive\n")
		PrintUsage("get")
	}

	defer cfg.Mongo.Close()

	o, err := get(cfg, args[0], true)

	if err != nil {
		log.Fatal(err)
	}

	if o == nil {
		return
	}

	if v > 0 {
		o = o.AtVersion(v)
	} else if t > 0 {
		o = o.AtTime(t)
	}

	if o == nil {
		return
	}

	o.History = nil

	b, err := json.MarshalIndent(o, "", "  ")

	if err != nil {
		log.Fatal(err)
	}

	fmt.Fprintf(os.Stdout, "%s\n", b)
}

func keysCmd(cfg *Config, args []string) {
	defer cfg.Mongo.Close()

	keys, err := Keys(cfg)

	if err != nil {
		log.Fatal(err)
	}

	if len(keys) == 0 {
		return
	}

	fmt.Fprintln(os.Stdout, strings.Join(keys, "\n"))
}

func logCmd(cfg *Config, args []string) {
	if len(args) != 1 {
		PrintUsage("log")
	}

	defer cfg.Mongo.Close()
	l, err := Log(cfg, args[0])

	if err != nil {
		log.Fatal(err)
	}

	if l == nil {
		return
	}

	b, err := json.MarshalIndent(l, "", "  ")

	if err != nil {
		log.Fatal(err)
	}

	fmt.Fprintf(os.Stdout, "%s\n", b)
}

func httpCmd(cfg *Config, args []string) {
	fs := flag.NewFlagSet("http", flag.ExitOnError)

	fs.StringVar(&cfg.HTTP.Host, "host", "localhost", "Host to bind to.")
	fs.IntVar(&cfg.HTTP.Port, "port", 5000, "Port to bind to.")

	fs.Parse(args)

	defer cfg.Mongo.Close()

	runHTTP(cfg)
}
<commit_msg>Put command requires at least one argument<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"strings"
)

func putCmd(cfg *Config, args []string) {
	if len(args) < 1 {
		PrintUsage("put")
	}

	var (
		err error
		val map[string]interface{}
	)

	// Decoded provided argument other read from stdin.
	if len(args) == 2 {
		err = json.Unmarshal([]byte(args[1]), &val)
	} else {
		err = json.NewDecoder(os.Stdin).Decode(&val)
	}

	defer cfg.Mongo.Close()
	o, err := Put(cfg, args[0], val)

	if err != nil {
		log.Fatal(err)
	}

	if o == nil {
		return
	}

	b, err := json.MarshalIndent(o, "", "  ")

	if err != nil {
		log.Fatal(err)
	}

	fmt.Fprintf(os.Stdout, "%s\n", b)
}

func getCmd(cfg *Config, args []string) {
	var (
		v  int
		ts string
	)

	fs := flag.NewFlagSet("get", flag.ExitOnError)

	fs.IntVar(&v, "version", 0, "Specific revision to get.")
	fs.StringVar(&ts, "time", "", "Returns the object as of the specified time.")

	fs.Parse(args)

	args = fs.Args()

	if len(args) != 1 {
		PrintUsage("get")
	}

	t, err := ParseTimeString(ts)

	if v > 0 && t > 0 {
		fmt.Println("error: version and time are mutually exclusive\n")
		PrintUsage("get")
	}

	defer cfg.Mongo.Close()

	o, err := get(cfg, args[0], true)

	if err != nil {
		log.Fatal(err)
	}

	if o == nil {
		return
	}

	if v > 0 {
		o = o.AtVersion(v)
	} else if t > 0 {
		o = o.AtTime(t)
	}

	if o == nil {
		return
	}

	o.History = nil

	b, err := json.MarshalIndent(o, "", "  ")

	if err != nil {
		log.Fatal(err)
	}

	fmt.Fprintf(os.Stdout, "%s\n", b)
}

func keysCmd(cfg *Config, args []string) {
	defer cfg.Mongo.Close()

	keys, err := Keys(cfg)

	if err != nil {
		log.Fatal(err)
	}

	if len(keys) == 0 {
		return
	}

	fmt.Fprintln(os.Stdout, strings.Join(keys, "\n"))
}

func logCmd(cfg *Config, args []string) {
	if len(args) != 1 {
		PrintUsage("log")
	}

	defer cfg.Mongo.Close()
	l, err := Log(cfg, args[0])

	if err != nil {
		log.Fatal(err)
	}

	if l == nil {
		return
	}

	b, err := json.MarshalIndent(l, "", "  ")

	if err != nil {
		log.Fatal(err)
	}

	fmt.Fprintf(os.Stdout, "%s\n", b)
}

func httpCmd(cfg *Config, args []string) {
	fs := flag.NewFlagSet("http", flag.ExitOnError)

	fs.StringVar(&cfg.HTTP.Host, "host", "localhost", "Host to bind to.")
	fs.IntVar(&cfg.HTTP.Port, "port", 5000, "Port to bind to.")

	fs.Parse(args)

	defer cfg.Mongo.Close()

	runHTTP(cfg)
}
<|endoftext|>