<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"go/build"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"

	"code.google.com/p/go.tools/go/loader"
	"code.google.com/p/go.tools/godoc/vfs"

	"sourcegraph.com/sourcegraph/srclib-go/gog"
	"sourcegraph.com/sourcegraph/srclib-go/gog/definfo"
	defpkg "sourcegraph.com/sourcegraph/srclib-go/golang_def"
	"sourcegraph.com/sourcegraph/srclib/graph"
	"sourcegraph.com/sourcegraph/srclib/grapher"
	"sourcegraph.com/sourcegraph/srclib/repo"
	"sourcegraph.com/sourcegraph/srclib/unit"
)

func init() {
	_, err := parser.AddCommand("graph",
		"graph a Go package",
		"Graph a Go package, producing all defs, refs, and docs.",
		&graphCmd,
	)
	if err != nil {
		log.Fatal(err)
	}

	// Check that we have the '-i' flag.
	cmd := exec.Command("go", "help", "build")
	o, err := cmd.Output()
	if err != nil {
		log.Fatal(err)
	}
	usage := strings.Split(string(o), "\n")[0] // The usage is on the first line.
	matched, err := regexp.MatchString("-i", usage)
	if err != nil {
		log.Fatal(err)
	}
	if !matched {
		log.Fatal("'go build' does not have the '-i' flag. Please upgrade to go1.3+.")
	}
}

type GraphCmd struct{}

var graphCmd GraphCmd

func (c *GraphCmd) Execute(args []string) error {
	var unit *unit.SourceUnit
	if err := json.NewDecoder(os.Stdin).Decode(&unit); err != nil {
		return err
	}
	if err := os.Stdin.Close(); err != nil {
		return err
	}

	if err := unmarshalTypedConfig(unit.Config); err != nil {
		return err
	}
	if err := config.apply(); err != nil {
		return err
	}

	if os.Getenv("IN_DOCKER_CONTAINER") != "" {
		buildPkg, err := UnitDataAsBuildPackage(unit)
		if err != nil {
			return err
		}

		// Make a new GOPATH.
		buildContext.GOPATH = "/tmp/gopath"

		// Set up GOPATH so it has this repo.
		log.Printf("Setting up a new GOPATH at %s", buildContext.GOPATH)
		dir := filepath.Join(buildContext.GOPATH, "src", string(unit.Repo))
		if err := os.MkdirAll(filepath.Dir(dir), 0700); err != nil {
			return err
		}
		log.Printf("Creating symlink to oldname %q at newname %q.", cwd, dir)
		if err := os.Symlink(cwd, dir); err != nil {
			return err
		}

		log.Printf("Changing directory to %q.", dir)
		if err := os.Chdir(dir); err != nil {
			return err
		}
		dockerCWD = cwd

		if config.GOROOT == "" {
			cwd = dir
		}

		// Get and install deps. (Only deps not in this repo; if we call `go
		// get` on this repo, we will either try to check out a different
		// version or fail with 'stale checkout?' because the .dockerignore
		// doesn't copy the .git dir.)
		var externalDeps []string
		for _, dep := range unit.Dependencies {
			importPath := dep.(string)
			if !strings.HasPrefix(importPath, string(unit.Repo)) && importPath != "C" {
				externalDeps = append(externalDeps, importPath)
			}
		}
		cmd := exec.Command("go", "get", "-d", "-t", "-v", "./"+buildPkg.Dir)
		cmd.Args = append(cmd.Args, externalDeps...)
		cmd.Env = config.env()
		cmd.Stdout, cmd.Stderr = os.Stderr, os.Stderr
		log.Printf("Downloading import dependencies: %v (env vars: %v).", cmd.Args, cmd.Env)
		if err := cmd.Run(); err != nil {
			return err
		}
		log.Printf("Finished downloading dependencies.")
	}

	out, err := Graph(unit)
	if err != nil {
		return err
	}

	// Make paths relative to repo.
	for _, gs := range out.Defs {
		if gs.File == "" {
			log.Printf("no file %+v", gs)
		}
		if gs.File != "" {
			gs.File = relPath(cwd, gs.File)
		}
	}
	for _, gr := range out.Refs {
		if gr.File != "" {
			gr.File = relPath(cwd, gr.File)
		}
	}
	for _, gd := range out.Docs {
		if gd.File != "" {
			gd.File = relPath(cwd, gd.File)
		}
	}

	if err := json.NewEncoder(os.Stdout).Encode(out); err != nil {
		return err
	}
	return nil
}

func relPath(base, path string) string {
	rp, err := filepath.Rel(base, path)
	if err != nil {
		log.Fatalf("Failed to make path %q relative to %q: %s", path, base, err)
	}

	// TODO(sqs): hack
	if strings.HasPrefix(rp, "../../../") && dockerCWD != "" {
		rp, err = filepath.Rel(dockerCWD, path)
		if err != nil {
			log.Fatalf("Failed to make path %q relative to %q: %s", path, cwd, err)
		}
	}

	return rp
}

func Graph(unit *unit.SourceUnit) (*grapher.Output, error) {
	pkg, err := UnitDataAsBuildPackage(unit)
	if err != nil {
		return nil, err
	}

	o, err := doGraph(pkg)
	if err != nil {
		return nil, err
	}

	o2 := grapher.Output{}

	uri := string(unit.Repo)

	for _, gs := range o.Defs {
		d, err := convertGoDef(gs, uri)
		if err != nil {
			return nil, err
		}
		if d != nil {
			o2.Defs = append(o2.Defs, d)
		}
	}
	for _, gr := range o.Refs {
		r, err := convertGoRef(gr, uri)
		if err != nil {
			return nil, err
		}
		if r != nil {
			o2.Refs = append(o2.Refs, r)
		}
	}
	for _, gd := range o.Docs {
		d, err := convertGoDoc(gd, uri)
		if err != nil {
			return nil, err
		}
		if d != nil {
			o2.Docs = append(o2.Docs, d)
		}
	}

	return &o2, nil
}

func convertGoDef(gs *gog.Def, repoURI string) (*graph.Def, error) {
	resolvedTarget, err := ResolveDep(gs.DefKey.PackageImportPath, repoURI)
	if err != nil {
		return nil, err
	}
	path := graph.DefPath(pathOrDot(strings.Join(gs.Path, "/")))
	treePath := treePath(string(path))
	if !treePath.IsValid() {
		return nil, fmt.Errorf("'%s' is not a valid tree-path", treePath)
	}

	def := &graph.Def{
		DefKey: graph.DefKey{
			Unit:     resolvedTarget.ToUnit,
			UnitType: resolvedTarget.ToUnitType,
			Path:     path,
		},
		TreePath: treePath,

		Name: gs.Name,
		Kind: definfo.GeneralKindMap[gs.Kind],

		File:     gs.File,
		DefStart: gs.DeclSpan[0],
		DefEnd:   gs.DeclSpan[1],

		Exported: gs.DefInfo.Exported,
		Test:     strings.HasSuffix(gs.File, "_test.go"),
	}

	d := defpkg.DefData{
		PackageImportPath: gs.DefKey.PackageImportPath,
		DefInfo:           gs.DefInfo,
	}
	def.Data, err = json.Marshal(d)
	if err != nil {
		return nil, err
	}

	if def.File == "" {
		// some cgo defs have empty File; omit them
		return nil, nil
	}

	return def, nil
}

func convertGoRef(gr *gog.Ref, repoURI string) (*graph.Ref, error) {
	resolvedTarget, err := ResolveDep(gr.Def.PackageImportPath, repoURI)
	if err != nil {
		return nil, err
	}
	if resolvedTarget == nil {
		return nil, nil
	}

	return &graph.Ref{
		DefRepo:     uriOrEmpty(resolvedTarget.ToRepoCloneURL),
		DefPath:     graph.DefPath(pathOrDot(strings.Join(gr.Def.Path, "/"))),
		DefUnit:     resolvedTarget.ToUnit,
		DefUnitType: resolvedTarget.ToUnitType,
		Def:         gr.IsDef,
		File:        gr.File,
		Start:       gr.Span[0],
		End:         gr.Span[1],
	}, nil
}

func convertGoDoc(gd *gog.Doc, repoURI string) (*graph.Doc, error) {
	resolvedTarget, err := ResolveDep(gd.PackageImportPath, repoURI)
	if err != nil {
		return nil, err
	}
	return &graph.Doc{
		DefKey: graph.DefKey{
			Path:     graph.DefPath(pathOrDot(strings.Join(gd.Path, "/"))),
			Unit:     resolvedTarget.ToUnit,
			UnitType: resolvedTarget.ToUnitType,
		},
		Format: gd.Format,
		Data:   gd.Data,
		File:   gd.File,
		Start:  gd.Span[0],
		End:    gd.Span[1],
	}, nil
}

func uriOrEmpty(cloneURL string) repo.URI {
	if cloneURL == "" {
		return ""
	}
	return repo.MakeURI(cloneURL)
}

func pathOrDot(path string) string {
	if path == "" {
		return "."
	}
	return path
}

func treePath(path string) graph.TreePath {
	if path == "" || path == "." {
		return graph.TreePath(".")
	}
	return graph.TreePath(fmt.Sprintf("./%s", path))
}

func doGraph(pkg *build.Package) (*gog.Output, error) {
	importPath := pkg.ImportPath

	// If we've overridden GOROOT and we're building a package not in
	// $GOROOT/src/pkg (such as "cmd/go"), then we need to virtualize GOROOT
	// because we can't set GOPATH=GOROOT (go/build ignores GOPATH in that
	// case).
	if config.GOROOT != "" && strings.HasPrefix(importPath, "cmd/") {
		// Unset our custom GOROOT (since we're routing FS ops to it using
		// vfs) and set it as our GOPATH.
		buildContext.GOROOT = build.Default.GOROOT
		buildContext.GOPATH = config.GOROOT

		virtualCWD = build.Default.GOROOT

		ns := vfs.NameSpace{}
		ns.Bind(filepath.Join(buildContext.GOROOT, "src/pkg"), vfs.OS(filepath.Join(config.GOROOT, "src/pkg")), "/", vfs.BindBefore)
		ns.Bind("/", vfs.OS("/"), "/", vfs.BindAfter)
		buildContext.IsDir = func(path string) bool {
			fi, err := ns.Stat(path)
			return err == nil && fi.Mode().IsDir()
		}
		buildContext.HasSubdir = func(root, dir string) (rel string, ok bool) { panic("unexpected") }
		buildContext.OpenFile = func(path string) (io.ReadCloser, error) {
			f, err := ns.Open(path)
			return f, err
		}
		buildContext.ReadDir = ns.ReadDir
	}

	if !loaderConfig.SourceImports {
		imports := map[string]struct{}{}
		for _, imp := range pkg.Imports {
			imports[imp] = struct{}{}
		}
		for _, imp := range pkg.TestImports {
			imports[imp] = struct{}{}
		}
		for _, imp := range pkg.XTestImports {
			imports[imp] = struct{}{}
		}

		for imp, _ := range imports {
			if imp == "C" {
				continue
			}
			cmd := exec.Command("go", "install", "-v", imp)
			cmd.Env = config.env()
			cmd.Stdout, cmd.Stderr = os.Stderr, os.Stderr
			log.Printf("Install %q: %v (env vars: %v)", importPath, cmd.Args, cmd.Env)
			if err := cmd.Run(); err != nil {
				return nil, err
			}
		}
	}

	importUnsafe := importPath == "unsafe"

	// Special-case: if this is a Cgo package, treat the CgoFiles as GoFiles or
	// else the character offsets will be junk.
	//
	// See https://codereview.appspot.com/86140043.
	loaderConfig.Build.CgoEnabled = false
	build.Default = *loaderConfig.Build
	if len(pkg.CgoFiles) > 0 {
		var allGoFiles []string
		allGoFiles = append(allGoFiles, pkg.GoFiles...)
		allGoFiles = append(allGoFiles, pkg.CgoFiles...)
		allGoFiles = append(allGoFiles, pkg.TestGoFiles...)
		for i, f := range allGoFiles {
			allGoFiles[i] = filepath.Join(cwd, pkg.Dir, f)
		}
		if err := loaderConfig.CreateFromFilenames(pkg.ImportPath, allGoFiles...); err != nil {
			return nil, err
		}
	} else {
		// Normal import
		if err := loaderConfig.ImportWithTests(importPath); err != nil {
			return nil, err
		}
	}

	if importUnsafe {
		// Special-case "unsafe" because go/loader does not let you load it
		// directly.
		if loaderConfig.ImportPkgs == nil {
			loaderConfig.ImportPkgs = make(map[string]bool)
		}
		loaderConfig.ImportPkgs["unsafe"] = true
	}

	prog, err := loaderConfig.Load()
	if err != nil {
		return nil, err
	}

	g := gog.New(prog)

	var pkgs []*loader.PackageInfo
	for _, pkg := range prog.Created {
		if strings.HasSuffix(pkg.Pkg.Name(), "_test") {
			// ignore xtest packages
			continue
		}
		pkgs = append(pkgs, pkg)
	}
	for _, pkg := range prog.Imported {
		pkgs = append(pkgs, pkg)
	}

	for _, pkg := range pkgs {
		if err := g.Graph(pkg); err != nil {
			return nil, err
		}
	}

	return &g.Output, nil
}
<commit_msg>error-tolerance (allowErrorsInGraph)<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"go/build"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/golang/gddo/gosrc"

	"code.google.com/p/go.tools/go/loader"
	"code.google.com/p/go.tools/godoc/vfs"

	"sourcegraph.com/sourcegraph/srclib-go/gog"
	"sourcegraph.com/sourcegraph/srclib-go/gog/definfo"
	defpkg "sourcegraph.com/sourcegraph/srclib-go/golang_def"
	"sourcegraph.com/sourcegraph/srclib/graph"
	"sourcegraph.com/sourcegraph/srclib/grapher"
	"sourcegraph.com/sourcegraph/srclib/repo"
	"sourcegraph.com/sourcegraph/srclib/unit"
)

func init() {
	_, err := parser.AddCommand("graph",
		"graph a Go package",
		"Graph a Go package, producing all defs, refs, and docs.",
		&graphCmd,
	)
	if err != nil {
		log.Fatal(err)
	}

	// Check that we have the '-i' flag.
	cmd := exec.Command("go", "help", "build")
	o, err := cmd.Output()
	if err != nil {
		log.Fatal(err)
	}
	usage := strings.Split(string(o), "\n")[0] // The usage is on the first line.
	matched, err := regexp.MatchString("-i", usage)
	if err != nil {
		log.Fatal(err)
	}
	if !matched {
		log.Fatal("'go build' does not have the '-i' flag. Please upgrade to go1.3+.")
	}
}

type GraphCmd struct{}

var graphCmd GraphCmd

func (c *GraphCmd) Execute(args []string) error {
	var unit *unit.SourceUnit
	if err := json.NewDecoder(os.Stdin).Decode(&unit); err != nil {
		return err
	}
	if err := os.Stdin.Close(); err != nil {
		return err
	}

	if err := unmarshalTypedConfig(unit.Config); err != nil {
		return err
	}
	if err := config.apply(); err != nil {
		return err
	}

	if os.Getenv("IN_DOCKER_CONTAINER") != "" {
		buildPkg, err := UnitDataAsBuildPackage(unit)
		if err != nil {
			return err
		}

		// Make a new GOPATH.
		buildContext.GOPATH = "/tmp/gopath"

		// Set up GOPATH so it has this repo.
		log.Printf("Setting up a new GOPATH at %s", buildContext.GOPATH)
		dir := filepath.Join(buildContext.GOPATH, "src", string(unit.Repo))
		if err := os.MkdirAll(filepath.Dir(dir), 0700); err != nil {
			return err
		}
		log.Printf("Creating symlink to oldname %q at newname %q.", cwd, dir)
		if err := os.Symlink(cwd, dir); err != nil {
			return err
		}

		log.Printf("Changing directory to %q.", dir)
		if err := os.Chdir(dir); err != nil {
			return err
		}
		dockerCWD = cwd

		if config.GOROOT == "" {
			cwd = dir
		}

		// Get and install deps. (Only deps not in this repo; if we call `go
		// get` on this repo, we will either try to check out a different
		// version or fail with 'stale checkout?' because the .dockerignore
		// doesn't copy the .git dir.)
		var externalDeps []string
		for _, dep := range unit.Dependencies {
			importPath := dep.(string)
			if !strings.HasPrefix(importPath, string(unit.Repo)) && importPath != "C" {
				externalDeps = append(externalDeps, importPath)
			}
		}
		cmd := exec.Command("go", "get", "-d", "-t", "-v", "./"+buildPkg.Dir)
		cmd.Args = append(cmd.Args, externalDeps...)
		cmd.Env = config.env()
		cmd.Stdout, cmd.Stderr = os.Stderr, os.Stderr
		log.Printf("Downloading import dependencies: %v (env vars: %v).", cmd.Args, cmd.Env)
		if err := cmd.Run(); err != nil {
			return err
		}
		log.Printf("Finished downloading dependencies.")
	}

	out, err := Graph(unit)
	if err != nil {
		return err
	}

	// Make paths relative to repo.
	for _, gs := range out.Defs {
		if gs.File == "" {
			log.Printf("no file %+v", gs)
		}
		if gs.File != "" {
			gs.File = relPath(cwd, gs.File)
		}
	}
	for _, gr := range out.Refs {
		if gr.File != "" {
			gr.File = relPath(cwd, gr.File)
		}
	}
	for _, gd := range out.Docs {
		if gd.File != "" {
			gd.File = relPath(cwd, gd.File)
		}
	}

	if err := json.NewEncoder(os.Stdout).Encode(out); err != nil {
		return err
	}
	return nil
}

func relPath(base, path string) string {
	rp, err := filepath.Rel(base, path)
	if err != nil {
		log.Fatalf("Failed to make path %q relative to %q: %s", path, base, err)
	}

	// TODO(sqs): hack
	if strings.HasPrefix(rp, "../../../") && dockerCWD != "" {
		rp, err = filepath.Rel(dockerCWD, path)
		if err != nil {
			log.Fatalf("Failed to make path %q relative to %q: %s", path, cwd, err)
		}
	}

	return rp
}

func Graph(unit *unit.SourceUnit) (*grapher.Output, error) {
	pkg, err := UnitDataAsBuildPackage(unit)
	if err != nil {
		return nil, err
	}

	o, err := doGraph(pkg)
	if err != nil {
		return nil, err
	}

	o2 := grapher.Output{}

	uri := string(unit.Repo)

	for _, gs := range o.Defs {
		d, err := convertGoDef(gs, uri)
		if err != nil {
			return nil, err
		}
		if d != nil {
			o2.Defs = append(o2.Defs, d)
		}
	}
	for _, gr := range o.Refs {
		r, err := convertGoRef(gr, uri)
		if err != nil {
			return nil, err
		}
		if r != nil {
			o2.Refs = append(o2.Refs, r)
		}
	}
	for _, gd := range o.Docs {
		d, err := convertGoDoc(gd, uri)
		if err != nil {
			return nil, err
		}
		if d != nil {
			o2.Docs = append(o2.Docs, d)
		}
	}

	return &o2, nil
}

func convertGoDef(gs *gog.Def, repoURI string) (*graph.Def, error) {
	resolvedTarget, err := ResolveDep(gs.DefKey.PackageImportPath, repoURI)
	if err != nil {
		return nil, err
	}
	path := graph.DefPath(pathOrDot(strings.Join(gs.Path, "/")))
	treePath := treePath(string(path))
	if !treePath.IsValid() {
		return nil, fmt.Errorf("'%s' is not a valid tree-path", treePath)
	}

	def := &graph.Def{
		DefKey: graph.DefKey{
			Unit:     resolvedTarget.ToUnit,
			UnitType: resolvedTarget.ToUnitType,
			Path:     path,
		},
		TreePath: treePath,

		Name: gs.Name,
		Kind: definfo.GeneralKindMap[gs.Kind],

		File:     gs.File,
		DefStart: gs.DeclSpan[0],
		DefEnd:   gs.DeclSpan[1],

		Exported: gs.DefInfo.Exported,
		Test:     strings.HasSuffix(gs.File, "_test.go"),
	}

	d := defpkg.DefData{
		PackageImportPath: gs.DefKey.PackageImportPath,
		DefInfo:           gs.DefInfo,
	}
	def.Data, err = json.Marshal(d)
	if err != nil {
		return nil, err
	}

	if def.File == "" {
		// some cgo defs have empty File; omit them
		return nil, nil
	}

	return def, nil
}

func convertGoRef(gr *gog.Ref, repoURI string) (*graph.Ref, error) {
	resolvedTarget, err := ResolveDep(gr.Def.PackageImportPath, repoURI)
	if err != nil {
		return nil, err
	}
	if resolvedTarget == nil {
		return nil, nil
	}

	return &graph.Ref{
		DefRepo:     uriOrEmpty(resolvedTarget.ToRepoCloneURL),
		DefPath:     graph.DefPath(pathOrDot(strings.Join(gr.Def.Path, "/"))),
		DefUnit:     resolvedTarget.ToUnit,
		DefUnitType: resolvedTarget.ToUnitType,
		Def:         gr.IsDef,
		File:        gr.File,
		Start:       gr.Span[0],
		End:         gr.Span[1],
	}, nil
}

func convertGoDoc(gd *gog.Doc, repoURI string) (*graph.Doc, error) {
	resolvedTarget, err := ResolveDep(gd.PackageImportPath, repoURI)
	if err != nil {
		return nil, err
	}
	return &graph.Doc{
		DefKey: graph.DefKey{
			Path:     graph.DefPath(pathOrDot(strings.Join(gd.Path, "/"))),
			Unit:     resolvedTarget.ToUnit,
			UnitType: resolvedTarget.ToUnitType,
		},
		Format: gd.Format,
		Data:   gd.Data,
		File:   gd.File,
		Start:  gd.Span[0],
		End:    gd.Span[1],
	}, nil
}

func uriOrEmpty(cloneURL string) repo.URI {
	if cloneURL == "" {
		return ""
	}
	return repo.MakeURI(cloneURL)
}

func pathOrDot(path string) string {
	if path == "" {
		return "."
	}
	return path
}

func treePath(path string) graph.TreePath {
	if path == "" || path == "." {
		return graph.TreePath(".")
	}
	return graph.TreePath(fmt.Sprintf("./%s", path))
}

// allowErrorsInGraph is whether the grapher should continue after
// encountering "reasonably common" errors (such as compile errors).
var allowErrorsInGraph = true

func doGraph(pkg *build.Package) (*gog.Output, error) {
	importPath := pkg.ImportPath

	// If we've overridden GOROOT and we're building a package not in
	// $GOROOT/src/pkg (such as "cmd/go"), then we need to virtualize GOROOT
	// because we can't set GOPATH=GOROOT (go/build ignores GOPATH in that
	// case).
	if config.GOROOT != "" && strings.HasPrefix(importPath, "cmd/") {
		// Unset our custom GOROOT (since we're routing FS ops to it using
		// vfs) and set it as our GOPATH.
		buildContext.GOROOT = build.Default.GOROOT
		buildContext.GOPATH = config.GOROOT

		virtualCWD = build.Default.GOROOT

		ns := vfs.NameSpace{}
		ns.Bind(filepath.Join(buildContext.GOROOT, "src/pkg"), vfs.OS(filepath.Join(config.GOROOT, "src/pkg")), "/", vfs.BindBefore)
		ns.Bind("/", vfs.OS("/"), "/", vfs.BindAfter)
		buildContext.IsDir = func(path string) bool {
			fi, err := ns.Stat(path)
			return err == nil && fi.Mode().IsDir()
		}
		buildContext.HasSubdir = func(root, dir string) (rel string, ok bool) { panic("unexpected") }
		buildContext.OpenFile = func(path string) (io.ReadCloser, error) {
			f, err := ns.Open(path)
			return f, err
		}
		buildContext.ReadDir = ns.ReadDir
	}

	if !loaderConfig.SourceImports {
		imports := map[string]struct{}{}
		for _, imp := range pkg.Imports {
			imports[imp] = struct{}{}
		}
		for _, imp := range pkg.TestImports {
			imports[imp] = struct{}{}
		}
		for _, imp := range pkg.XTestImports {
			imports[imp] = struct{}{}
		}

		for imp, _ := range imports {
			if imp == "C" {
				continue
			}
			if gosrc.IsGoRepoPath(imp) {
				// Optimization: don't bother installing builtin
				// packages because they're already installed. (But if
				// we accidentally install one, it's OK and not going
				// to cause any problems.)
				continue
			}
			cmd := exec.Command("go", "install", "-v", imp)
			cmd.Env = config.env()
			cmd.Stdout, cmd.Stderr = os.Stderr, os.Stderr
			log.Printf("Install %q: %v (env vars: %v)", importPath, cmd.Args, cmd.Env)
			if err := cmd.Run(); err != nil {
				if allowErrorsInGraph {
					log.Printf("Warning: failed to install package %q (command %v, env vars %v): %s. Continuing...", imp, cmd.Args, cmd.Env)
				} else {
					return nil, err
				}
			}
		}
	}

	importUnsafe := importPath == "unsafe"

	// Special-case: if this is a Cgo package, treat the CgoFiles as GoFiles or
	// else the character offsets will be junk.
	//
	// See https://codereview.appspot.com/86140043.
	loaderConfig.Build.CgoEnabled = false
	build.Default = *loaderConfig.Build
	if len(pkg.CgoFiles) > 0 {
		var allGoFiles []string
		allGoFiles = append(allGoFiles, pkg.GoFiles...)
		allGoFiles = append(allGoFiles, pkg.CgoFiles...)
		allGoFiles = append(allGoFiles, pkg.TestGoFiles...)
		for i, f := range allGoFiles {
			allGoFiles[i] = filepath.Join(cwd, pkg.Dir, f)
		}
		if err := loaderConfig.CreateFromFilenames(pkg.ImportPath, allGoFiles...); err != nil {
			return nil, err
		}
	} else {
		// Normal import
		if err := loaderConfig.ImportWithTests(importPath); err != nil {
			return nil, err
		}
	}

	if importUnsafe {
		// Special-case "unsafe" because go/loader does not let you load it
		// directly.
		if loaderConfig.ImportPkgs == nil {
			loaderConfig.ImportPkgs = make(map[string]bool)
		}
		loaderConfig.ImportPkgs["unsafe"] = true
	}

	prog, err := loaderConfig.Load()
	if err != nil {
		return nil, err
	}

	g := gog.New(prog)

	var pkgs []*loader.PackageInfo
	for _, pkg := range prog.Created {
		if strings.HasSuffix(pkg.Pkg.Name(), "_test") {
			// ignore xtest packages
			continue
		}
		pkgs = append(pkgs, pkg)
	}
	for _, pkg := range prog.Imported {
		pkgs = append(pkgs, pkg)
	}

	for _, pkg := range pkgs {
		if err := g.Graph(pkg); err != nil {
			return nil, err
		}
	}

	return &g.Output, nil
}
<|endoftext|>


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

package ht

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"path"

	"log"

	"net/http/cookiejar"

	"strings"
	"sync"
)

// A Suite is a collection of tests which are run together.
type Suite struct {
	Name        string
	Description string

	// Test contains the actual tests to execute.
	Tests []*Test

	// OmitChecks allows to omit all checks defined on the main tests.
	OmitChecks bool

	// Setup contain tests to be executed before the execution actual tests.
	// If one or more setup test fail, the main tets won't be executed.
	// Teardown tests are excuted after the main test.
	// Setup and Teardown share the cookie jar with the main tests.
	Setup, Teardown []*Test

	// Variables contains global variables to be used during this
	// execution
	Variables map[string]string

	// KeepCookies determines whether to use a cookie jar to keep
	// cookies between tests.
	KeepCookies bool

	// Log is the logger to be used by tests and checks.
	Log *log.Logger
}

// Init initialises the suite by unrolling repeated test in Tests.
// Setup and Teardown tests are not unrolled.
func (s *Suite) Init() {
	// Unroll repeated test to single instances
	var tests []*Test
	for _, t := range s.Tests {
		if len(t.UnrollWith) == 0 {
			tests = append(tests, t)
		} else {
			nreps := lcmOf(t.UnrollWith)
			unrolled := t.Repeat(nreps, t.UnrollWith)
			// Clear repetitions to prevent re-unrolling during
			// a second Prepare.
			for _, u := range unrolled {
				u.UnrollWith = nil
			}
			tests = append(tests, unrolled...)
		}
	}
	tests = append(s.Setup, tests...)
	tests = append(tests, s.Teardown...)
	s.Tests = tests
}

// Compile prepares all tests in s for execution.
func (s *Suite) Compile() error {
	// Create cookie jar.
	var jar *cookiejar.Jar
	if s.KeepCookies {
		jar, _ = cookiejar.New(nil)
	}

	// Compile all tests and inject jar and logger.
	prepare := func(t *Test, which string, omit bool) error {
		err := t.Compile(s.Variables)
		if err != nil {
			return fmt.Errorf("Suite %q, cannot prepare %s %q: %s",
				s.Name, which, t.Name, err)
		}
		if omit {
			t.checks = nil
		}
		t.Jar = jar
		t.Log = s.Log
		return nil
	}
	for _, t := range s.Setup {
		if err := prepare(t, "setup", false); err != nil { // Cannot omit checks in setup.
			return err
		}
	}
	for _, t := range s.Tests {
		if err := prepare(t, "test", s.OmitChecks); err != nil {
			return err
		}
	}
	for _, t := range s.Teardown {
		if err := prepare(t, "teardown", s.OmitChecks); err != nil {
			return err
		}
	}
	return nil
}

// Execute the setup tests in s. The tests are executed sequentialy,
// execution stops on the first error.
func (s *Suite) ExecuteSetup() Result {
	return s.execute(s.Setup)
}

// ExecuteTeardown runs all teardown tests ignoring all errors.
func (s *Suite) ExecuteTeardown() Result {
	return s.execute(s.Teardown)
}

// Execute all non-setup, non-teardown tests of s sequentialy.
func (s *Suite) ExecuteTests() Result {
	return s.execute(s.Tests)
}

func (s *Suite) execute(tests []*Test) Result {
	if len(tests) == 0 {
		return Result{Status: Pass}
	}
	result := Result{Elements: make([]Result, len(tests))}
	for i, test := range tests {
		result.Elements[i] = test.Run()
	}
	result.Status = CombinedStatus(result.Elements)
	return result
}

// Execute the whole suite sequentially.
func (s *Suite) Execute() Result {
	result := s.ExecuteSetup()
	if result.Status != Pass {
		result.Error = fmt.Errorf("Setup failed")
		return result
	}
	result = s.ExecuteTests()
	if result.Status != Pass {
		result.Error = fmt.Errorf("TODO")
		return result
	}
	s.ExecuteTeardown()
	return result
}

// ExecuteConcurrent executes all non-setup, non-teardown tests concurrently.
// But at most maxConcurrent tests of s are executed concurrently.
func (s *Suite) ExecuteConcurrent(maxConcurrent int) error {
	if maxConcurrent > len(s.Tests) {
		maxConcurrent = len(s.Tests)
	}
	s.Log.Printf("Running %d test concurrently", maxConcurrent)
	res := make(chan string, len(s.Tests))

	c := make(chan *Test, maxConcurrent)
	wg := sync.WaitGroup{}
	wg.Add(maxConcurrent)
	for i := 0; i < maxConcurrent; i++ {
		go func() {
			defer wg.Done()
			for test := range c {
				result := test.Run() // TODO
				if result.Status != Pass {
					res <- test.Name
				}
			}
		}()
	}
	for _, test := range s.Tests {
		c <- test
	}
	close(c)
	wg.Wait()
	close(res)

	var failures []string
	for ft := range res {
		failures = append(failures, ft)
	}
	if len(failures) > 0 {
		return fmt.Errorf("Failes %d of %d test: %s", len(failures),
			len(s.Tests), strings.Join(failures, ", "))
	}
	return nil
}

// jsonSuite is the struct corresponding to the JSON serialisation od a Suite.
type jsonSuite struct {
	Tests []*Test
	Suite struct {
		Name, Description       string
		KeepCookies, OmitChecks bool
		Tests, Setup, Teardow   []string
	}
	Variables map[string]string
}

// loadJsonSuite loads the file with the given filename and decodes into a jsonSuite.
// It will try to find filename in all paths.
func loadJsonSuite(filename string, paths []string) (rs jsonSuite, err error) {
	var all []byte
	for _, p := range paths {
		name := path.Join(p, filename)
		all, err = ioutil.ReadFile(name)
		if err == nil {
			break
		}
	}
	if err != nil {
		return rs, err
	}
	err = json.Unmarshal(all, &rs)
	if err != nil {
		return rs, err
	}
	return rs, nil
}

// LoadSuite reads a suite from filename. Uargh
func LoadSuite(filename string, paths []string) (*Suite, error) {
	loaded := make(map[string][]*Test) // filename --> Tests

	s, err := loadJsonSuite(filename, paths)
	if err != nil {
		return nil, err
	}

	suite := Suite{
		Name:        s.Suite.Name,
		Description: s.Suite.Description,
		KeepCookies: s.Suite.KeepCookies,
		OmitChecks:  s.Suite.OmitChecks,
	}

	findTest := func(name string) *Test {
		var collection []*Test
		if i := strings.LastIndex(name, "@"); i != -1 {
			file := strings.TrimSpace(name[i+1:])
			if t, ok := loaded[file]; ok {
				collection = t
			} else {
				rs, err := loadJsonSuite(file, paths)
				if err != nil {
					return nil // TODO, return error here
				}
				loaded[file] = rs.Tests
				collection = rs.Tests
				name = strings.TrimSpace(name[:i])
			}
		} else {
			collection = s.Tests
		}

		for i := range collection {
			if collection[i].Name == name {
				return collection[i]
			}
		}
		return nil
	}

	for _, name := range s.Suite.Setup {
		tp := findTest(name)
		if tp == nil {
			return nil, fmt.Errorf("no test %q found", name)
		}
		suite.Setup = append(suite.Setup, tp)
	}
	for _, name := range s.Suite.Tests {
		tp := findTest(name)
		if tp == nil {
			return nil, fmt.Errorf("no test %q found", name)
		}
		suite.Tests = append(suite.Tests, tp)
	}
	for _, name := range s.Suite.Teardow {
		tp := findTest(name)
		if tp == nil {
			return nil, fmt.Errorf("no test %q found", name)
		}
		suite.Teardown = append(suite.Teardown, tp)
	}

	suite.Variables = s.Variables
	suite.KeepCookies = s.Suite.KeepCookies
	suite.OmitChecks = s.Suite.OmitChecks

	return &suite, nil
}
<commit_msg>Refactor suite loading to handle single tests.<commit_after>// Copyright 2014 Volker Dobler.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ht

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"path"

	"log"

	"net/http/cookiejar"

	"strings"
	"sync"
)

// A Suite is a collection of tests which are run together.
type Suite struct {
	Name        string
	Description string

	// Test contains the actual tests to execute.
	Tests []*Test

	// OmitChecks allows to omit all checks defined on the main tests.
	OmitChecks bool

	// Setup contain tests to be executed before the execution actual tests.
	// If one or more setup test fail, the main tets won't be executed.
	// Teardown tests are excuted after the main test.
	// Setup and Teardown share the cookie jar with the main tests.
	Setup, Teardown []*Test

	// Variables contains global variables to be used during this
	// execution
	Variables map[string]string

	// KeepCookies determines whether to use a cookie jar to keep
	// cookies between tests.
	KeepCookies bool

	// Log is the logger to be used by tests and checks.
	Log *log.Logger
}

// Init initialises the suite by unrolling repeated test in Tests.
// Setup and Teardown tests are not unrolled.
func (s *Suite) Init() {
	// Unroll repeated test to single instances
	var tests []*Test
	for _, t := range s.Tests {
		if len(t.UnrollWith) == 0 {
			tests = append(tests, t)
		} else {
			nreps := lcmOf(t.UnrollWith)
			unrolled := t.Repeat(nreps, t.UnrollWith)
			// Clear repetitions to prevent re-unrolling during
			// a second Prepare.
			for _, u := range unrolled {
				u.UnrollWith = nil
			}
			tests = append(tests, unrolled...)
		}
	}
	tests = append(s.Setup, tests...)
	tests = append(tests, s.Teardown...)
	s.Tests = tests
}

// Compile prepares all tests in s for execution.
func (s *Suite) Compile() error {
	// Create cookie jar.
	var jar *cookiejar.Jar
	if s.KeepCookies {
		jar, _ = cookiejar.New(nil)
	}

	// Compile all tests and inject jar and logger.
	prepare := func(t *Test, which string, omit bool) error {
		err := t.Compile(s.Variables)
		if err != nil {
			return fmt.Errorf("Suite %q, cannot prepare %s %q: %s",
				s.Name, which, t.Name, err)
		}
		if omit {
			t.checks = nil
		}
		t.Jar = jar
		t.Log = s.Log
		return nil
	}
	for _, t := range s.Setup {
		if err := prepare(t, "setup", false); err != nil { // Cannot omit checks in setup.
			return err
		}
	}
	for _, t := range s.Tests {
		if err := prepare(t, "test", s.OmitChecks); err != nil {
			return err
		}
	}
	for _, t := range s.Teardown {
		if err := prepare(t, "teardown", s.OmitChecks); err != nil {
			return err
		}
	}
	return nil
}

// Execute the setup tests in s. The tests are executed sequentialy,
// execution stops on the first error.
func (s *Suite) ExecuteSetup() Result {
	return s.execute(s.Setup)
}

// ExecuteTeardown runs all teardown tests ignoring all errors.
func (s *Suite) ExecuteTeardown() Result {
	return s.execute(s.Teardown)
}

// Execute all non-setup, non-teardown tests of s sequentialy.
func (s *Suite) ExecuteTests() Result {
	return s.execute(s.Tests)
}

func (s *Suite) execute(tests []*Test) Result {
	if len(tests) == 0 {
		return Result{Status: Pass}
	}
	result := Result{Elements: make([]Result, len(tests))}
	for i, test := range tests {
		result.Elements[i] = test.Run()
	}
	result.Status = CombinedStatus(result.Elements)
	return result
}

// Execute the whole suite sequentially.
func (s *Suite) Execute() Result {
	result := s.ExecuteSetup()
	if result.Status != Pass {
		result.Error = fmt.Errorf("Setup failed")
		return result
	}
	result = s.ExecuteTests()
	if result.Status != Pass {
		result.Error = fmt.Errorf("TODO")
		return result
	}
	s.ExecuteTeardown()
	return result
}

// ExecuteConcurrent executes all non-setup, non-teardown tests concurrently.
// But at most maxConcurrent tests of s are executed concurrently.
func (s *Suite) ExecuteConcurrent(maxConcurrent int) error {
	if maxConcurrent > len(s.Tests) {
		maxConcurrent = len(s.Tests)
	}
	s.Log.Printf("Running %d test concurrently", maxConcurrent)
	res := make(chan string, len(s.Tests))

	c := make(chan *Test, maxConcurrent)
	wg := sync.WaitGroup{}
	wg.Add(maxConcurrent)
	for i := 0; i < maxConcurrent; i++ {
		go func() {
			defer wg.Done()
			for test := range c {
				result := test.Run() // TODO
				if result.Status != Pass {
					res <- test.Name
				}
			}
		}()
	}
	for _, test := range s.Tests {
		c <- test
	}
	close(c)
	wg.Wait()
	close(res)

	var failures []string
	for ft := range res {
		failures = append(failures, ft)
	}
	if len(failures) > 0 {
		return fmt.Errorf("Failes %d of %d test: %s", len(failures),
			len(s.Tests), strings.Join(failures, ", "))
	}
	return nil
}

// jsonSuite is the struct corresponding to the JSON serialisation od a Suite.
type jsonSuite struct {
	Tests []*Test
	Suite struct {
		Name, Description       string
		KeepCookies, OmitChecks bool
		Tests, Setup, Teardow   []string
	}
	Variables map[string]string
}

// loadJsonSuite loads the file with the given filename and decodes into a jsonSuite.
// It will try to find filename in all paths and will report the dir path back in which
// the suite was found.
func loadJsonSuite(filename string, paths []string) (rs jsonSuite, dir string, err error) {
	var all []byte
	for _, p := range paths {
		dir = path.Join(p, filename)
		all, err = ioutil.ReadFile(dir)
		if err == nil {
			break
		}
	}
	dir = path.Dir(dir)
	if err != nil {
		return rs, dir, err
	}
	err = json.Unmarshal(all, &rs)
	if err != nil {
		return rs, dir, err
	}
	return rs, dir, nil
}

// LoadSuite reads a suite from filename while looking for filename in the
// filesystem paths.
//
// References to individual tests are made by name of the test with two
// special ways to load test from their own files or from different suits
// to access "non-local tests":
//     @name-of-file-with-one-test             will load a single test from
//                                             the given file
//     Name of test @ filename-of-other suite  will use test "Name of Test"
//                                             from the suite found in the
//                                             given file
// Bothe filenames are relative to the position of the suite beeing read
// (i.e. pathsis ignored while loading such tests).
func LoadSuite(filename string, paths []string) (*Suite, error) {
	s, dir, err := loadJsonSuite(filename, paths)
	if err != nil {
		return nil, err
	}

	suite := Suite{
		Name:        s.Suite.Name,
		Description: s.Suite.Description,
		KeepCookies: s.Suite.KeepCookies,
		OmitChecks:  s.Suite.OmitChecks,
	}

	for _, name := range s.Suite.Setup {
		tp, err := findTest(name, s.Tests, dir)
		if err != nil {
			return nil, fmt.Errorf("test %q: %s", name, err)
		}
		suite.Setup = append(suite.Setup, tp)
	}
	for _, name := range s.Suite.Tests {
		tp, err := findTest(name, s.Tests, dir)
		if err != nil {
			return nil, fmt.Errorf("test %q: %s", name, err)
		}
		suite.Tests = append(suite.Tests, tp)
	}
	for _, name := range s.Suite.Teardow {
		tp, err := findTest(name, s.Tests, dir)
		if err != nil {
			return nil, fmt.Errorf("test %q: %s", name, err)
		}
		suite.Teardown = append(suite.Teardown, tp)
	}

	suite.Variables = s.Variables
	suite.KeepCookies = s.Suite.KeepCookies
	suite.OmitChecks = s.Suite.OmitChecks

	return &suite, nil
}

// findTest tries to find a test with the given name in the given Suite.
// Non-local tests are tried to be loaded from dir.
func findTest(name string, collection []*Test, dir string) (*Test, error) {
	if strings.HasPrefix(name, "@") {
		name = strings.TrimSpace(name[1:])
		name = path.Join(dir, name)
		println("Local Single Test", name)
		return loadJsonTest(name)
	}

	if i := strings.LastIndex(name, "@"); i != -1 {
		// Load suite only from given dir (possible relative to dir).
		file := strings.TrimSpace(name[i+1:])
		name = strings.TrimSpace(name[:i])
		println("Local References Test", file, name)
		rs, _, err := loadJsonSuite(file, []string{dir})
		if err != nil {
			return nil, err
		}
		collection = rs.Tests
	} else {
		println("Plain Test", name)
	}

	for i := range collection {
		if collection[i].Name == name {
			return collection[i], nil
		}
	}
	return nil, errors.New("not found")
}

// loadJsonTest loads the file with the given filename and decodes into a Test.
func loadJsonTest(filename string) (*Test, error) {
	all, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}
	t := &Test{}
	err = json.Unmarshal(all, t)
	if err != nil {
		return nil, err
	}
	return t, nil
}
<|endoftext|>


<commit_before>package peerstream

import (
	"errors"
	"net"
	"sync"
)

// fd is a (file) descriptor, unix style
type fd uint32

type Swarm struct {
	// active streams.
	streams    map[*Stream]struct{}
	streamLock sync.RWMutex

	// active connections. generate new Streams
	conns    map[*Conn]struct{}
	connLock sync.RWMutex

	// active listeners. generate new Listeners
	listeners    map[*Listener]struct{}
	listenerLock sync.RWMutex

	// these handlers should be accessed with their getter/setter
	// as this pointer may be changed at any time.
	handlerLock   sync.RWMutex  // protects the functions below
	connHandler   ConnHandler   // receives Conns intiated remotely
	streamHandler StreamHandler // receives Streams initiated remotely
	selectConn    SelectConn    // default SelectConn function
}

func NewSwarm() *Swarm {
	return &Swarm{
		streams:       make(map[*Stream]struct{}),
		conns:         make(map[*Conn]struct{}),
		listeners:     make(map[*Listener]struct{}),
		selectConn:    SelectRandomConn,
		streamHandler: NoOpStreamHandler,
		connHandler:   NoOpConnHandler,
	}
}

// SetStreamHandler assigns the stream handler in the swarm.
// The handler assumes responsibility for closing the stream.
// This need not happen at the end of the handler, leaving the
// stream open (to be used and closed later) is fine.
// It is also fine to keep a pointer to the Stream.
// This is a threadsafe (atomic) operation
func (s *Swarm) SetStreamHandler(sh StreamHandler) {
	s.handlerLock.Lock()
	defer s.handlerLock.Unlock()
	s.streamHandler = sh
}

// StreamHandler returns the Swarm's current StreamHandler.
// This is a threadsafe (atomic) operation
func (s *Swarm) StreamHandler() StreamHandler {
	s.handlerLock.RLock()
	defer s.handlerLock.RUnlock()
	if s.streamHandler == nil {
		return NoOpStreamHandler
	}
	return s.streamHandler
}

// SetConnHandler assigns the conn handler in the swarm.
// Unlike the StreamHandler, the ConnHandler has less respon-
// ibility for the Connection. The Swarm is still its client.
// This handler is only a notification.
// This is a threadsafe (atomic) operation
func (s *Swarm) SetConnHandler(ch ConnHandler) {
	s.handlerLock.Lock()
	defer s.handlerLock.Unlock()
	s.connHandler = ch
}

// ConnHandler returns the Swarm's current ConnHandler.
// This is a threadsafe (atomic) operation
func (s *Swarm) ConnHandler() ConnHandler {
	s.handlerLock.RLock()
	defer s.handlerLock.RUnlock()
	if s.connHandler == nil {
		return NoOpConnHandler
	}
	return s.connHandler
}

// SetConnSelect assigns the connection selector in the swarm.
// If cs is nil, will use SelectRandomConn
// This is a threadsafe (atomic) operation
func (s *Swarm) SetSelectConn(cs SelectConn) {
	s.handlerLock.Lock()
	defer s.handlerLock.Unlock()
	s.selectConn = cs
}

// ConnSelect returns the Swarm's current connection selector.
// ConnSelect is used in order to select the best of a set of
// possible connections. The default chooses one at random.
// This is a threadsafe (atomic) operation
func (s *Swarm) SelectConn() SelectConn {
	s.handlerLock.RLock()
	defer s.handlerLock.RUnlock()
	if s.selectConn == nil {
		return SelectRandomConn
	}
	return s.selectConn
}

// Conns returns all the connections associated with this Swarm.
func (s *Swarm) Conns() []*Conn {
	conns := make([]*Conn, 0, len(s.conns))
	for c := range s.conns {
		conns = append(conns, c)
	}
	return conns
}

// Listeners returns all the listeners associated with this Swarm.
func (s *Swarm) Listeners() []*Listener {
	out := make([]*Listener, 0, len(s.listeners))
	for c := range s.listeners {
		out = append(out, c)
	}
	return out
}

// Streams returns all the streams associated with this Swarm.
func (s *Swarm) Streams() []*Stream {
	out := make([]*Stream, 0, len(s.streams))
	for c := range s.streams {
		out = append(out, c)
	}
	return out
}

// AddListener adds net.Listener to the Swarm, and immediately begins
// accepting incoming connections.
func (s *Swarm) AddListener(l net.Listener) (*Listener, error) {
	return s.addListener(l)
}

// AddListenerWithRateLimit adds Listener to the Swarm, and immediately
// begins accepting incoming connections. The rate of connection acceptance
// depends on the RateLimit option
// func (s *Swarm) AddListenerWithRateLimit(net.Listner, RateLimit) // TODO

// AddConn gives the Swarm ownership of net.Conn. The Swarm will open a
// SPDY session and begin listening for Streams.
// Returns the resulting Swarm-associated peerstream.Conn.
// Idempotent: if the Connection has already been added, this is a no-op.
func (s *Swarm) AddConn(netConn net.Conn) (*Conn, error) {
	return s.addConn(netConn, false)
}

// NewStream opens a new Stream on the best available connection,
// as selected by current swarm.SelectConn.
func (s *Swarm) NewStream() (*Stream, error) {
	return s.NewStreamSelectConn(s.SelectConn())
}

func (s *Swarm) newStreamSelectConn(selConn SelectConn, conns []*Conn) (*Stream, error) {
	if selConn == nil {
		return nil, errors.New("nil SelectConn")
	}

	best := selConn(conns)
	if best == nil || !ConnInConns(best, conns) {
		return nil, ErrInvalidConnSelected
	}
	return s.NewStreamWithConn(best)
}

// NewStreamWithSelectConn opens a new Stream on a connection selected
// by selConn.
func (s *Swarm) NewStreamSelectConn(selConn SelectConn) (*Stream, error) {
	if selConn == nil {
		return nil, errors.New("nil SelectConn")
	}

	conns := s.Conns()
	if len(conns) == 0 {
		return nil, ErrNoConnections
	}
	return s.newStreamSelectConn(selConn, conns)
}

// NewStreamWithGroup opens a new Stream on an available connection in
// the given group. Uses the current swarm.SelectConn to pick between
// multiple connections.
func (s *Swarm) NewStreamWithGroup(group Group) (*Stream, error) {
	conns := s.ConnsWithGroup(group)
	return s.newStreamSelectConn(s.SelectConn(), conns)
}

// NewStreamWithNetConn opens a new Stream on given net.Conn.
// Calls s.AddConn(netConn).
func (s *Swarm) NewStreamWithNetConn(netConn net.Conn) (*Stream, error) {
	c, err := s.AddConn(netConn)
	if err != nil {
		return nil, err
	}
	return s.NewStreamWithConn(c)
}

// NewStreamWithConnection opens a new Stream on given connection.
func (s *Swarm) NewStreamWithConn(conn *Conn) (*Stream, error) {
	if conn == nil {
		return nil, errors.New("nil Conn")
	}
	if conn.Swarm() != s {
		return nil, errors.New("connection not associated with swarm")
	}

	s.connLock.RLock()
	if _, found := s.conns[conn]; !found {
		s.connLock.RUnlock()
		return nil, errors.New("connection not associated with swarm")
	}
	s.connLock.RUnlock()
	return s.createStream(conn)
}

// AddConnToGroup assigns given Group to conn
func (s *Swarm) AddConnToGroup(conn *Conn, g Group) {
	conn.groups.Add(g)
}

// ConnsWithGroup returns all the connections with a given Group
func (s *Swarm) ConnsWithGroup(g Group) []*Conn {
	return ConnsWithGroup(g, s.Conns())
}

// StreamsWithGroup returns all the streams with a given Group
func (s *Swarm) StreamsWithGroup(g Group) []*Stream {
	return StreamsWithGroup(g, s.Streams())
}

// Close shuts down the Swarm, and it's listeners.
func (s *Swarm) Close() error {
	// shut down TODO
	return nil
}
<commit_msg>forgot locks in getters #2<commit_after>package peerstream

import (
	"errors"
	"net"
	"sync"
)

// fd is a (file) descriptor, unix style
type fd uint32

type Swarm struct {
	// active streams.
	streams    map[*Stream]struct{}
	streamLock sync.RWMutex

	// active connections. generate new Streams
	conns    map[*Conn]struct{}
	connLock sync.RWMutex

	// active listeners. generate new Listeners
	listeners    map[*Listener]struct{}
	listenerLock sync.RWMutex

	// these handlers should be accessed with their getter/setter
	// as this pointer may be changed at any time.
	handlerLock   sync.RWMutex  // protects the functions below
	connHandler   ConnHandler   // receives Conns intiated remotely
	streamHandler StreamHandler // receives Streams initiated remotely
	selectConn    SelectConn    // default SelectConn function
}

func NewSwarm() *Swarm {
	return &Swarm{
		streams:       make(map[*Stream]struct{}),
		conns:         make(map[*Conn]struct{}),
		listeners:     make(map[*Listener]struct{}),
		selectConn:    SelectRandomConn,
		streamHandler: NoOpStreamHandler,
		connHandler:   NoOpConnHandler,
	}
}

// SetStreamHandler assigns the stream handler in the swarm.
// The handler assumes responsibility for closing the stream.
// This need not happen at the end of the handler, leaving the
// stream open (to be used and closed later) is fine.
// It is also fine to keep a pointer to the Stream.
// This is a threadsafe (atomic) operation
func (s *Swarm) SetStreamHandler(sh StreamHandler) {
	s.handlerLock.Lock()
	defer s.handlerLock.Unlock()
	s.streamHandler = sh
}

// StreamHandler returns the Swarm's current StreamHandler.
// This is a threadsafe (atomic) operation
func (s *Swarm) StreamHandler() StreamHandler {
	s.handlerLock.RLock()
	defer s.handlerLock.RUnlock()
	if s.streamHandler == nil {
		return NoOpStreamHandler
	}
	return s.streamHandler
}

// SetConnHandler assigns the conn handler in the swarm.
// Unlike the StreamHandler, the ConnHandler has less respon-
// ibility for the Connection. The Swarm is still its client.
// This handler is only a notification.
// This is a threadsafe (atomic) operation
func (s *Swarm) SetConnHandler(ch ConnHandler) {
	s.handlerLock.Lock()
	defer s.handlerLock.Unlock()
	s.connHandler = ch
}

// ConnHandler returns the Swarm's current ConnHandler.
// This is a threadsafe (atomic) operation
func (s *Swarm) ConnHandler() ConnHandler {
	s.handlerLock.RLock()
	defer s.handlerLock.RUnlock()
	if s.connHandler == nil {
		return NoOpConnHandler
	}
	return s.connHandler
}

// SetConnSelect assigns the connection selector in the swarm.
// If cs is nil, will use SelectRandomConn
// This is a threadsafe (atomic) operation
func (s *Swarm) SetSelectConn(cs SelectConn) {
	s.handlerLock.Lock()
	defer s.handlerLock.Unlock()
	s.selectConn = cs
}

// ConnSelect returns the Swarm's current connection selector.
// ConnSelect is used in order to select the best of a set of
// possible connections. The default chooses one at random.
// This is a threadsafe (atomic) operation
func (s *Swarm) SelectConn() SelectConn {
	s.handlerLock.RLock()
	defer s.handlerLock.RUnlock()
	if s.selectConn == nil {
		return SelectRandomConn
	}
	return s.selectConn
}

// Conns returns all the connections associated with this Swarm.
func (s *Swarm) Conns() []*Conn {
	s.connLock.RLock()
	conns := make([]*Conn, 0, len(s.conns))
	for c := range s.conns {
		conns = append(conns, c)
	}
	s.connLock.RUnlock()
	return conns
}

// Listeners returns all the listeners associated with this Swarm.
func (s *Swarm) Listeners() []*Listener {
	s.listenerLock.RLock()
	out := make([]*Listener, 0, len(s.listeners))
	for c := range s.listeners {
		out = append(out, c)
	}
	s.listenerLock.RUnlock()
	return out
}

// Streams returns all the streams associated with this Swarm.
func (s *Swarm) Streams() []*Stream {
	s.streamLock.RLock()
	out := make([]*Stream, 0, len(s.streams))
	for c := range s.streams {
		out = append(out, c)
	}
	s.streamLock.RUnlock()
	return out
}

// AddListener adds net.Listener to the Swarm, and immediately begins
// accepting incoming connections.
func (s *Swarm) AddListener(l net.Listener) (*Listener, error) {
	return s.addListener(l)
}

// AddListenerWithRateLimit adds Listener to the Swarm, and immediately
// begins accepting incoming connections. The rate of connection acceptance
// depends on the RateLimit option
// func (s *Swarm) AddListenerWithRateLimit(net.Listner, RateLimit) // TODO

// AddConn gives the Swarm ownership of net.Conn. The Swarm will open a
// SPDY session and begin listening for Streams.
// Returns the resulting Swarm-associated peerstream.Conn.
// Idempotent: if the Connection has already been added, this is a no-op.
func (s *Swarm) AddConn(netConn net.Conn) (*Conn, error) {
	return s.addConn(netConn, false)
}

// NewStream opens a new Stream on the best available connection,
// as selected by current swarm.SelectConn.
func (s *Swarm) NewStream() (*Stream, error) {
	return s.NewStreamSelectConn(s.SelectConn())
}

func (s *Swarm) newStreamSelectConn(selConn SelectConn, conns []*Conn) (*Stream, error) {
	if selConn == nil {
		return nil, errors.New("nil SelectConn")
	}

	best := selConn(conns)
	if best == nil || !ConnInConns(best, conns) {
		return nil, ErrInvalidConnSelected
	}
	return s.NewStreamWithConn(best)
}

// NewStreamWithSelectConn opens a new Stream on a connection selected
// by selConn.
func (s *Swarm) NewStreamSelectConn(selConn SelectConn) (*Stream, error) {
	if selConn == nil {
		return nil, errors.New("nil SelectConn")
	}

	conns := s.Conns()
	if len(conns) == 0 {
		return nil, ErrNoConnections
	}
	return s.newStreamSelectConn(selConn, conns)
}

// NewStreamWithGroup opens a new Stream on an available connection in
// the given group. Uses the current swarm.SelectConn to pick between
// multiple connections.
func (s *Swarm) NewStreamWithGroup(group Group) (*Stream, error) {
	conns := s.ConnsWithGroup(group)
	return s.newStreamSelectConn(s.SelectConn(), conns)
}

// NewStreamWithNetConn opens a new Stream on given net.Conn.
// Calls s.AddConn(netConn).
func (s *Swarm) NewStreamWithNetConn(netConn net.Conn) (*Stream, error) {
	c, err := s.AddConn(netConn)
	if err != nil {
		return nil, err
	}
	return s.NewStreamWithConn(c)
}

// NewStreamWithConnection opens a new Stream on given connection.
func (s *Swarm) NewStreamWithConn(conn *Conn) (*Stream, error) {
	if conn == nil {
		return nil, errors.New("nil Conn")
	}
	if conn.Swarm() != s {
		return nil, errors.New("connection not associated with swarm")
	}

	s.connLock.RLock()
	if _, found := s.conns[conn]; !found {
		s.connLock.RUnlock()
		return nil, errors.New("connection not associated with swarm")
	}
	s.connLock.RUnlock()
	return s.createStream(conn)
}

// AddConnToGroup assigns given Group to conn
func (s *Swarm) AddConnToGroup(conn *Conn, g Group) {
	conn.groups.Add(g)
}

// ConnsWithGroup returns all the connections with a given Group
func (s *Swarm) ConnsWithGroup(g Group) []*Conn {
	return ConnsWithGroup(g, s.Conns())
}

// StreamsWithGroup returns all the streams with a given Group
func (s *Swarm) StreamsWithGroup(g Group) []*Stream {
	return StreamsWithGroup(g, s.Streams())
}

// Close shuts down the Swarm, and it's listeners.
func (s *Swarm) Close() error {
	// shut down TODO
	return nil
}
<|endoftext|>


<commit_before>package servers

import (
	"net"
	"strings"
	"sync"
	"time"

	"github.com/miekg/dns"

	"github.com/hawkingrei/g53/cache"
	"github.com/hawkingrei/g53/utils"
)

// NewService creates a new service
func NewService() (s *utils.Service) {
	s = &utils.Service{TTL: -1}
	return
}

// ServiceListProvider represents the entrypoint to get containers
type ServiceListProvider interface {
	AddService(utils.Service)
	RemoveService(utils.Service) error
	SetService(utils.Service, utils.Service) error
	GetService(utils.Service) (utils.Service, error)
	GetAllServices() []utils.Service
}

// DNSServer represents a DNS server
type DNSServer struct {
	config     *utils.Config
	server     *dns.Server
	mux        *dns.ServeMux
	lock       *sync.RWMutex
	publicDns  *cache.Cache
	privateDns *cache.Cache
}

// NewDNSServer create a new DNSServer
func NewDNSServer(c *utils.Config) *DNSServer {
	publicDns, _ := cache.New(10000)
	privateDns, _ := cache.New(100000000)
	s := &DNSServer{
		config:     c,
		lock:       &sync.RWMutex{},
		publicDns:  publicDns,
		privateDns: privateDns,
	}

	logger.Debugf("Handling DNS requests for '%s'.", c.Domain.String())

	s.mux = dns.NewServeMux()
	s.mux.HandleFunc(".", s.handleRequest)
	s.server = &dns.Server{Addr: c.DnsAddr, Net: "udp", Handler: s.mux}

	return s
}

// Start starts the DNSServer
func (s *DNSServer) Start() error {
	logger.Infof("start DNS Server")
	return s.server.ListenAndServe()
}

// Stop stops the DNSServer
func (s *DNSServer) Stop() {
	s.server.Shutdown()
}

func (s *DNSServer) SetService(originalValue utils.Service, modifyValue utils.Service) error {
	return s.privateDns.Set(originalValue, modifyValue)
}

// AddService adds a new container and thus new DNS records
func (s *DNSServer) AddService(service utils.Service) {
	if service.RecordType == "CNAME" || service.RecordType == "A" {
		s.privateDns.Add(service)

		logger.Debugf("Added service: '%s'.", service)
		logger.Debugf("Handling DNS requests for '%s'.", service.Aliases)
		s.mux.HandleFunc(service.Aliases+".", s.handleRequest)
		//for _, alias := range service.Aliases {
		//	logger.Debugf("Handling DNS requests for '%s'.", alias)
		//	s.mux.HandleFunc(alias+".", s.handleRequest)
		//}
	} else {
		logger.Warningf("Service '%s' ignored: No RecordType provided:", service)
	}
}

// RemoveService removes a new container and thus DNS records
func (s *DNSServer) RemoveService(service utils.Service) error {
	if err := s.privateDns.Remove(service); err != nil {
		return err
	}
	s.mux.HandleRemove(service.Aliases + ".")
	logger.Debugf("Removeed service '%s'", service)

	return nil
}

// GetService reads a service from the repository
func (s *DNSServer) GetService(service utils.Service) (utils.Service, error) {
	result, err := s.privateDns.Get(service)
	if err != nil {
		return *new(utils.Service), err
	}
	return utils.EntryToServer(result), err
}

// GetAllServices reads all services from the repository

func (s *DNSServer) GetAllServices() []utils.Service {
	return s.privateDns.List()
}

func (s *DNSServer) handleForward(w dns.ResponseWriter, r *dns.Msg) {
	//r.SetEdns0(4096, true)
	logger.Debugf("Using DNS forwarding for '%s'", r.Question[0].Name)
	logger.Debugf("Forwarding DNS nameservers: %s", s.config.Nameservers.String())

	// Otherwise just forward the request to another server
	c := new(dns.Client)
	c.UDPSize = uint16(4096)

	// look at each Nameserver, stop on success
	for i := range s.config.Nameservers {
		logger.Debugf("Using Nameserver %s", s.config.Nameservers[i])

		in, _, err := c.Exchange(r, s.config.Nameservers[i])
		if err == nil {
			w.WriteMsg(in)
			return
		}

		if i == (len(s.config.Nameservers) - 1) {
			logger.Noticef("DNS fowarding for '%s' failed: no more nameservers to try", err.Error())

			// Send failure reply
			m := new(dns.Msg)
			m.SetReply(r)
			m.Ns = s.createSOA()
			m.SetRcode(r, dns.RcodeRefused) // REFUSED
			w.WriteMsg(m)

		} else {
			logger.Errorf("DNS fowarding for '%s' failed: trying next Nameserver...", err.Error())
		}
	}
}

func (s *DNSServer) makeServiceCNAME(n string, service utils.Service) dns.RR {
	rr := new(dns.CNAME)

	var ttl int
	if service.TTL != -1 {
		ttl = service.TTL
	} else {
		ttl = s.config.Ttl
	}

	rr.Hdr = dns.RR_Header{
		Name:   n,
		Rrtype: dns.TypeCNAME,
		Class:  dns.ClassINET,
		Ttl:    uint32(ttl),
	}

	if len(service.Value) != 0 {
		if len(service.Value) > 1 {
			logger.Warningf("Multiple IP address found for container '%s'. Only the first address will be used", service.Aliases)
		}
		rr.Target = service.Value + "."
	} else {
		logger.Errorf("No valid IP address found for container '%s' at makeServiceCNAME", service.Aliases)
	}
	return rr
}

func (s *DNSServer) makeServiceA(n string, service utils.Service) dns.RR {
	rr := new(dns.A)

	var ttl int
	if service.TTL != -1 {
		ttl = service.TTL
	} else {
		ttl = s.config.Ttl
	}

	rr.Hdr = dns.RR_Header{
		Name:   n,
		Rrtype: dns.TypeA,
		Class:  dns.ClassINET,
		Ttl:    uint32(ttl),
	}

	if len(service.Value) != 0 {
		if len(service.Value) > 1 {
			logger.Warningf("Multiple IP address found for container '%s'. Only the first address will be used", service.Aliases)
		}
		rr.A = net.ParseIP(service.Value)
	} else {
		logger.Errorf("No valid IP address found for container '%s' at makeServiceA", service.Aliases)
	}

	return rr
}

func (s *DNSServer) handleRequest(w dns.ResponseWriter, r *dns.Msg) {
	m := new(dns.Msg)
	m.Compress = true
	m.SetReply(r)
	m.RecursionAvailable = true

	// Send empty response for empty requests
	if len(r.Question) == 0 {
		m.Ns = s.createSOA()
		w.WriteMsg(m)
		return
	}

	// respond to SOA requests
	if r.Question[0].Qtype == dns.TypeSOA {
		m.Answer = s.createSOA()
		w.WriteMsg(m)
		return
	}

	m.Answer = make([]dns.RR, 0, 2)
	query := r.Question[0].Name

	// trim off any trailing dot
	if query[len(query)-1] == '.' {
		query = query[:len(query)-1]
	}

	result := s.queryServices(utils.Service{dns.TypeToString[r.Question[0].Qtype], "", 0, strings.ToLower(query)})

	logger.Debugf("DNS record found for query '%s'  '%s'", query, dns.TypeToString[r.Question[0].Qtype])
	for service := range result {
		var rr dns.RR
		switch r.Question[0].Qtype {
		case dns.TypeA:
			rr = s.makeServiceA(r.Question[0].Name, utils.EntryToServer(service))
		case dns.TypeCNAME:
			rr = s.makeServiceCNAME(r.Question[0].Name, utils.EntryToServer(service))
		default:
			// this query type isn't supported, but we do have
			// a record with this name. Per RFC 4074 sec. 3, we
			// immediately return an empty NOERROR reply.
			m.Ns = s.createSOA()
			m.MsgHdr.Authoritative = true
			w.WriteMsg(m)
			return
		}

		m.Answer = append(m.Answer, rr)
	}

	// We didn't find a record corresponding to the query
	if len(m.Answer) == 0 {
		s.handleForward(w, r)
		return
	}
	w.WriteMsg(m)
	return
}

func (s *DNSServer) queryServices(service utils.Service) chan *utils.Entry {
	c := make(chan *utils.Entry, 10)
	go func() {
		result, err := s.privateDns.Get(service)
		if err == nil {
			c <- result
		}
		close(c)
	}()
	return c
}

// TTL is used from config so that not-found result responses are not cached
// for a long time. The other defaults left as is(skydns source) because they
// do not have an use case in this situation.
func (s *DNSServer) createSOA() []dns.RR {
	dom := dns.Fqdn(s.config.Domain.String() + ".")
	soa := &dns.SOA{
		Hdr: dns.RR_Header{
			Name:   dom,
			Rrtype: dns.TypeSOA,
			Class:  dns.ClassINET,
			Ttl:    uint32(s.config.Ttl)},
		Ns:      "g53." + dom,
		Mbox:    "g53.g53." + dom,
		Serial:  uint32(time.Now().Truncate(time.Hour).Unix()),
		Refresh: 28800,
		Retry:   7200,
		Expire:  604800,
		Minttl:  uint32(s.config.Ttl),
	}
	return []dns.RR{soa}
}
<commit_msg>fix: update the validation<commit_after>package servers

import (
	"net"
	"strings"
	"sync"
	"time"
	"github.com/miekg/dns"

	"github.com/hawkingrei/g53/cache"
	"github.com/hawkingrei/g53/utils"
)

// NewService creates a new service
func NewService() (s *utils.Service) {
	s = &utils.Service{TTL: -1}
	return
}

// ServiceListProvider represents the entrypoint to get containers
type ServiceListProvider interface {
	AddService(utils.Service)
	RemoveService(utils.Service) error
	SetService(utils.Service, utils.Service) error
	GetService(utils.Service) (utils.Service, error)
	GetAllServices() []utils.Service
}

// DNSServer represents a DNS server
type DNSServer struct {
	config     *utils.Config
	server     *dns.Server
	mux        *dns.ServeMux
	lock       *sync.RWMutex
	publicDns  *cache.Cache
	privateDns *cache.Cache
}

// NewDNSServer create a new DNSServer
func NewDNSServer(c *utils.Config) *DNSServer {
	publicDns, _ := cache.New(10000)
	privateDns, _ := cache.New(100000000)
	s := &DNSServer{
		config:     c,
		lock:       &sync.RWMutex{},
		publicDns:  publicDns,
		privateDns: privateDns,
	}

	logger.Debugf("Handling DNS requests for '%s'.", c.Domain.String())

	s.mux = dns.NewServeMux()
	s.mux.HandleFunc(".", s.handleRequest)
	s.server = &dns.Server{Addr: c.DnsAddr, Net: "udp", Handler: s.mux}

	return s
}

// Start starts the DNSServer
func (s *DNSServer) Start() error {
	logger.Infof("start DNS Server")
	return s.server.ListenAndServe()
}

// Stop stops the DNSServer
func (s *DNSServer) Stop() {
	s.server.Shutdown()
}

func (s *DNSServer) SetService(originalValue utils.Service, modifyValue utils.Service) error {
	return s.privateDns.Set(originalValue, modifyValue)
}

// AddService adds a new container and thus new DNS records
func (s *DNSServer) AddService(service utils.Service) {
	if service.RecordType == "CNAME" || service.RecordType == "A" {
		if  string(service.Aliases[len(service.Aliases)-1]) != "." {
			service.Aliases = string(service.Aliases) + "."
		}

		if service.RecordType == "CNAME" && string(service.Value[len(service.Value)-1]) != "." {
			service.Value = string(service.Value) + "."
		}

		s.privateDns.Add(service)

		logger.Debugf("Added service: '%s'.", service)
		logger.Debugf("Handling DNS requests for '%s'.", service.Aliases)
		s.mux.HandleFunc(service.Aliases+".", s.handleRequest)
		//for _, alias := range service.Aliases {
		//	logger.Debugf("Handling DNS requests for '%s'.", alias)
		//	s.mux.HandleFunc(alias+".", s.handleRequest)
		//}

	} else {
		logger.Warningf("Service '%s' ignored: No RecordType provided:", service)
	}
}

// RemoveService removes a new container and thus DNS records
func (s *DNSServer) RemoveService(service utils.Service) error {
	if err := s.privateDns.Remove(service); err != nil {
		return err
	}
	s.mux.HandleRemove(service.Aliases + ".")
	logger.Debugf("Removeed service '%s'", service)

	return nil
}

// GetService reads a service from the repository
func (s *DNSServer) GetService(service utils.Service) (utils.Service, error) {
	result, err := s.privateDns.Get(service)
	if err != nil {
		return *new(utils.Service), err
	}
	return utils.EntryToServer(result), err
}

// GetAllServices reads all services from the repository

func (s *DNSServer) GetAllServices() []utils.Service {
	return s.privateDns.List()
}

func (s *DNSServer) handleForward(w dns.ResponseWriter, r *dns.Msg) {
	//r.SetEdns0(4096, true)
	logger.Debugf("Using DNS forwarding for '%s'", r.Question[0].Name)
	logger.Debugf("Forwarding DNS nameservers: %s", s.config.Nameservers.String())

	// Otherwise just forward the request to another server
	c := new(dns.Client)
	c.UDPSize = uint16(4096)

	// look at each Nameserver, stop on success
	for i := range s.config.Nameservers {
		logger.Debugf("Using Nameserver %s", s.config.Nameservers[i])

		in, _, err := c.Exchange(r, s.config.Nameservers[i])
		if err == nil {
			w.WriteMsg(in)
			return
		}

		if i == (len(s.config.Nameservers) - 1) {
			logger.Noticef("DNS fowarding for '%s' failed: no more nameservers to try", err.Error())

			// Send failure reply
			m := new(dns.Msg)
			m.SetReply(r)
			m.Ns = s.createSOA()
			m.SetRcode(r, dns.RcodeRefused) // REFUSED
			w.WriteMsg(m)

		} else {
			logger.Errorf("DNS fowarding for '%s' failed: trying next Nameserver...", err.Error())
		}
	}
}

func (s *DNSServer) makeServiceCNAME(n string, service utils.Service) dns.RR {
	rr := new(dns.CNAME)

	var ttl int
	if service.TTL != -1 {
		ttl = service.TTL
	} else {
		ttl = s.config.Ttl
	}

	rr.Hdr = dns.RR_Header{
		Name:   n,
		Rrtype: dns.TypeCNAME,
		Class:  dns.ClassINET,
		Ttl:    uint32(ttl),
	}

	if len(service.Value) != 0 {
		if len(service.Value) > 1 {
			logger.Warningf("Multiple IP address found for container '%s'. Only the first address will be used", service.Aliases)
		}
		rr.Target = service.Value
	} else {
		logger.Errorf("No valid IP address found for container '%s' at makeServiceCNAME", service.Aliases)
	}
	return rr
}

func (s *DNSServer) makeServiceA(n string, service utils.Service) dns.RR {
	rr := new(dns.A)

	var ttl int
	if service.TTL != -1 {
		ttl = service.TTL
	} else {
		ttl = s.config.Ttl
	}

	rr.Hdr = dns.RR_Header{
		Name:   n,
		Rrtype: dns.TypeA,
		Class:  dns.ClassINET,
		Ttl:    uint32(ttl),
	}

	if len(service.Value) != 0 {
		if len(service.Value) > 1 {
			logger.Warningf("Multiple IP address found for container '%s'. Only the first address will be used", service.Aliases)
		}
		rr.A = net.ParseIP(service.Value)
	} else {
		logger.Errorf("No valid IP address found for container '%s' at makeServiceA", service.Aliases)
	}

	return rr
}

func (s *DNSServer) handleRequest(w dns.ResponseWriter, r *dns.Msg) {
	m := new(dns.Msg)
	m.Compress = true
	m.SetReply(r)
	m.RecursionAvailable = true

	// Send empty response for empty requests
	if len(r.Question) == 0 {
		m.Ns = s.createSOA()
		w.WriteMsg(m)
		return
	}

	// respond to SOA requests
	if r.Question[0].Qtype == dns.TypeSOA {
		m.Answer = s.createSOA()
		w.WriteMsg(m)
		return
	}

	m.Answer = make([]dns.RR, 0, 2)
	query := r.Question[0].Name

	// trim off any trailing dot
	if query[len(query)-1] != '.' {
		query = query + "."
	}
	logger.Debugf("DNS record found for query '%s'  '%s'", query, dns.TypeToString[r.Question[0].Qtype])
	result := s.queryServices(utils.Service{dns.TypeToString[r.Question[0].Qtype], "", 0, strings.ToLower(query)})
	for service := range result {
		var rr dns.RR
		switch r.Question[0].Qtype {
		case dns.TypeA:
			rr = s.makeServiceA(r.Question[0].Name, utils.EntryToServer(service))
		case dns.TypeCNAME:
			rr = s.makeServiceCNAME(r.Question[0].Name, utils.EntryToServer(service))
		default:
			// this query type isn't supported, but we do have
			// a record with this name. Per RFC 4074 sec. 3, we
			// immediately return an empty NOERROR reply.
			m.Ns = s.createSOA()
			m.MsgHdr.Authoritative = true
			w.WriteMsg(m)
			return
		}

		m.Answer = append(m.Answer, rr)
	}

	// We didn't find a record corresponding to the query
	if len(m.Answer) == 0 {
		s.handleForward(w, r)
		return
	}
	w.WriteMsg(m)
	return
}

func (s *DNSServer) queryServices(service utils.Service) chan *utils.Entry {
	c := make(chan *utils.Entry, 10)
	go func() {
		result, err := s.privateDns.Get(service)
		if err == nil {
			c <- result
		} else {
			logger.Debugf(err.Error())
		}
		close(c)
	}()
	return c
}

// TTL is used from config so that not-found result responses are not cached
// for a long time. The other defaults left as is(skydns source) because they
// do not have an use case in this situation.
func (s *DNSServer) createSOA() []dns.RR {
	dom := dns.Fqdn(s.config.Domain.String() + ".")
	soa := &dns.SOA{
		Hdr: dns.RR_Header{
			Name:   dom,
			Rrtype: dns.TypeSOA,
			Class:  dns.ClassINET,
			Ttl:    uint32(s.config.Ttl)},
		Ns:      "g53." + dom,
		Mbox:    "g53.g53." + dom,
		Serial:  uint32(time.Now().Truncate(time.Hour).Unix()),
		Refresh: 28800,
		Retry:   7200,
		Expire:  604800,
		Minttl:  uint32(s.config.Ttl),
	}
	return []dns.RR{soa}
}
<|endoftext|>


<commit_before>package user

import (
	"context"
	htmlTemplate "html/template"
	"net/http"
	"strconv"
	"time"
	"xorkevin.dev/governor"
	"xorkevin.dev/governor/service/kvstore"
	"xorkevin.dev/governor/service/mail"
	"xorkevin.dev/governor/service/msgqueue"
	"xorkevin.dev/governor/service/user/apikey"
	"xorkevin.dev/governor/service/user/approval/model"
	"xorkevin.dev/governor/service/user/gate"
	"xorkevin.dev/governor/service/user/model"
	"xorkevin.dev/governor/service/user/role"
	"xorkevin.dev/governor/service/user/session/model"
	"xorkevin.dev/governor/service/user/token"
	"xorkevin.dev/governor/util/rank"
	"xorkevin.dev/hunter2"
)

const (
	authRoutePrefix = "/auth"
)

type (
	// User is a user management service
	User interface {
		GetByID(userid string) (*ResUserGet, error)
	}

	Service interface {
		governor.Service
		User
	}

	service struct {
		users             usermodel.Repo
		sessions          sessionmodel.Repo
		approvals         approvalmodel.Repo
		roles             role.Role
		apikeys           apikey.Apikey
		kvnewuser         kvstore.KVStore
		kvemailchange     kvstore.KVStore
		kvpassreset       kvstore.KVStore
		kvsessions        kvstore.KVStore
		queue             msgqueue.Msgqueue
		mailer            mail.Mail
		gate              gate.Gate
		hasher            *hunter2.Blake2bHasher
		verifier          *hunter2.Verifier
		tokenizer         token.Tokenizer
		logger            governor.Logger
		baseURL           string
		authURL           string
		accessTime        int64
		refreshTime       int64
		refreshCacheTime  int64
		confirmTime       int64
		passwordResetTime int64
		newLoginEmail     bool
		passwordMinSize   int
		userApproval      bool
		emailurlbase      string
		tplemailchange    *htmlTemplate.Template
		tplforgotpass     *htmlTemplate.Template
		tplnewuser        *htmlTemplate.Template
	}

	router struct {
		s service
	}

	// NewUserProps are properties of a newly created user
	NewUserProps struct {
		Userid       string `json:"userid"`
		Username     string `json:"username"`
		Email        string `json:"email"`
		FirstName    string `json:"first_name"`
		LastName     string `json:"last_name"`
		CreationTime int64  `json:"creation_time"`
	}

	// DeleteUserProps are properties of a deleted user
	DeleteUserProps struct {
		Userid string `json:"userid"`
	}
)

const (
	NewUserQueueID    = "gov.user.new"
	DeleteUserQueueID = "gov.user.delete"
)

const (
	time5m     int64 = int64(5 * time.Minute / time.Second)
	time24h    int64 = int64(24 * time.Hour / time.Second)
	time6month int64 = time24h * 365 / 2
)

// New creates a new User
func New(users usermodel.Repo, sessions sessionmodel.Repo, approvals approvalmodel.Repo, roles role.Role, apikeys apikey.Apikey, kv kvstore.KVStore, queue msgqueue.Msgqueue, mailer mail.Mail, tokenizer token.Tokenizer, g gate.Gate) Service {
	hasher := hunter2.NewBlake2bHasher()
	verifier := hunter2.NewVerifier()
	verifier.RegisterHash(hasher)

	return &service{
		users:             users,
		sessions:          sessions,
		approvals:         approvals,
		roles:             roles,
		apikeys:           apikeys,
		kvnewuser:         kv.Subtree("newuser"),
		kvemailchange:     kv.Subtree("emailchange"),
		kvpassreset:       kv.Subtree("passreset"),
		kvsessions:        kv.Subtree("sessions"),
		queue:             queue,
		mailer:            mailer,
		gate:              g,
		hasher:            hasher,
		verifier:          verifier,
		tokenizer:         tokenizer,
		accessTime:        time5m,
		refreshTime:       time6month,
		refreshCacheTime:  time24h,
		confirmTime:       time24h,
		passwordResetTime: time24h,
	}
}

func (s *service) Register(r governor.ConfigRegistrar, jr governor.JobRegistrar) {
	r.SetDefault("accesstime", "5m")
	r.SetDefault("refreshtime", "4380h")
	r.SetDefault("refreshcache", "24h")
	r.SetDefault("confirmtime", "24h")
	r.SetDefault("passwordresettime", "24h")
	r.SetDefault("newloginemail", true)
	r.SetDefault("passwordminsize", 8)
	r.SetDefault("userapproval", false)
	r.SetDefault("email.url.base", "http://localhost:8080")
	r.SetDefault("email.url.emailchange", "/a/security/email/confirm?key={{.Userid}}.{{.Key}}")
	r.SetDefault("email.url.forgotpass", "/x/resetpass?key={{.Userid}}.{{.Key}}")
	r.SetDefault("email.url.newuser", "/x/confirm?email={{.Email}}&key={{.Key}}")
}

func (s *service) router() *router {
	return &router{
		s: *s,
	}
}

func (s *service) Init(ctx context.Context, c governor.Config, r governor.ConfigReader, l governor.Logger, m governor.Router) error {
	s.logger = l
	l = s.logger.WithData(map[string]string{
		"phase": "init",
	})

	s.baseURL = c.BaseURL
	s.authURL = c.BaseURL + r.URL() + authRoutePrefix
	if t, err := time.ParseDuration(r.GetStr("accesstime")); err != nil {
		return governor.NewError("Failed to parse access time", http.StatusBadRequest, err)
	} else {
		s.accessTime = int64(t / time.Second)
	}
	if t, err := time.ParseDuration(r.GetStr("refreshtime")); err != nil {
		return governor.NewError("Failed to parse refresh time", http.StatusBadRequest, err)
	} else {
		s.refreshTime = int64(t / time.Second)
	}
	if t, err := time.ParseDuration(r.GetStr("refreshcache")); err != nil {
		return governor.NewError("Failed to parse refresh cache", http.StatusBadRequest, err)
	} else {
		s.refreshCacheTime = int64(t / time.Second)
	}
	if t, err := time.ParseDuration(r.GetStr("confirmtime")); err != nil {
		return governor.NewError("Failed to parse confirm time", http.StatusBadRequest, err)
	} else {
		s.confirmTime = int64(t / time.Second)
	}
	if t, err := time.ParseDuration(r.GetStr("passwordresettime")); err != nil {
		return governor.NewError("Failed to parse password reset time", http.StatusBadRequest, err)
	} else {
		s.passwordResetTime = int64(t / time.Second)
	}
	s.newLoginEmail = r.GetBool("newloginemail")
	s.passwordMinSize = r.GetInt("passwordminsize")
	s.userApproval = r.GetBool("userapproval")

	s.emailurlbase = r.GetStr("email.url.base")
	if t, err := htmlTemplate.New("email.url.emailchange").Parse(r.GetStr("email.url.emailchange")); err != nil {
		return governor.NewError("Failed to parse email change url template", http.StatusBadRequest, err)
	} else {
		s.tplemailchange = t
	}
	if t, err := htmlTemplate.New("email.url.forgotpass").Parse(r.GetStr("email.url.forgotpass")); err != nil {
		return governor.NewError("Failed to parse forgot pass url template", http.StatusBadRequest, err)
	} else {
		s.tplforgotpass = t
	}
	if t, err := htmlTemplate.New("email.url.newuser").Parse(r.GetStr("email.url.newuser")); err != nil {
		return governor.NewError("Failed to parse new user url template", http.StatusBadRequest, err)
	} else {
		s.tplnewuser = t
	}

	l.Info("loaded config", map[string]string{
		"accesstime (s)":        strconv.FormatInt(s.accessTime, 10),
		"refreshtime (s)":       strconv.FormatInt(s.refreshTime, 10),
		"refreshcache (s)":      strconv.FormatInt(s.refreshCacheTime, 10),
		"confirmtime (s)":       strconv.FormatInt(s.confirmTime, 10),
		"passwordresettime (s)": strconv.FormatInt(s.passwordResetTime, 10),
		"newloginemail":         strconv.FormatBool(s.newLoginEmail),
		"passwordminsize":       strconv.Itoa(s.passwordMinSize),
		"issuer":                r.GetStr("issuer"),
		"userapproval":          strconv.FormatBool(s.userApproval),
		"tplemailchange":        r.GetStr("email.url.emailchange"),
		"tplforgotpass":         r.GetStr("email.url.forgotpass"),
		"tplnewuser":            r.GetStr("email.url.newuser"),
	})

	sr := s.router()
	sr.mountRoute(m.Group("/user"))
	sr.mountAuth(m.Group(authRoutePrefix))
	sr.mountApikey(m.Group("/apikey"))
	l.Info("mounted http routes", nil)
	return nil
}

func (s *service) Setup(req governor.ReqSetup) error {
	l := s.logger.WithData(map[string]string{
		"phase": "setup",
	})

	madmin, err := s.users.New(req.Username, req.Password, req.Email, req.Firstname, req.Lastname)
	if err != nil {
		return err
	}

	if err := s.users.Setup(); err != nil {
		return err
	}
	l.Info("created user table", nil)

	if err := s.sessions.Setup(); err != nil {
		return err
	}
	l.Info("created usersession table", nil)

	if err := s.approvals.Setup(); err != nil {
		return err
	}
	l.Info("created userapprovals table", nil)

	if err := s.users.Insert(madmin); err != nil {
		return err
	}
	if err := s.roles.InsertRoles(madmin.Userid, rank.Admin()); err != nil {
		return err
	}
	l.Info("inserted new setup admin", map[string]string{
		"username": madmin.Username,
		"userid":   madmin.Userid,
	})
	return nil
}

func (s *service) Start(ctx context.Context) error {
	return nil
}

func (s *service) Stop(ctx context.Context) {
}

func (s *service) Health() error {
	return nil
}
<commit_msg>Change email confirm url<commit_after>package user

import (
	"context"
	htmlTemplate "html/template"
	"net/http"
	"strconv"
	"time"
	"xorkevin.dev/governor"
	"xorkevin.dev/governor/service/kvstore"
	"xorkevin.dev/governor/service/mail"
	"xorkevin.dev/governor/service/msgqueue"
	"xorkevin.dev/governor/service/user/apikey"
	"xorkevin.dev/governor/service/user/approval/model"
	"xorkevin.dev/governor/service/user/gate"
	"xorkevin.dev/governor/service/user/model"
	"xorkevin.dev/governor/service/user/role"
	"xorkevin.dev/governor/service/user/session/model"
	"xorkevin.dev/governor/service/user/token"
	"xorkevin.dev/governor/util/rank"
	"xorkevin.dev/hunter2"
)

const (
	authRoutePrefix = "/auth"
)

type (
	// User is a user management service
	User interface {
		GetByID(userid string) (*ResUserGet, error)
	}

	Service interface {
		governor.Service
		User
	}

	service struct {
		users             usermodel.Repo
		sessions          sessionmodel.Repo
		approvals         approvalmodel.Repo
		roles             role.Role
		apikeys           apikey.Apikey
		kvnewuser         kvstore.KVStore
		kvemailchange     kvstore.KVStore
		kvpassreset       kvstore.KVStore
		kvsessions        kvstore.KVStore
		queue             msgqueue.Msgqueue
		mailer            mail.Mail
		gate              gate.Gate
		hasher            *hunter2.Blake2bHasher
		verifier          *hunter2.Verifier
		tokenizer         token.Tokenizer
		logger            governor.Logger
		baseURL           string
		authURL           string
		accessTime        int64
		refreshTime       int64
		refreshCacheTime  int64
		confirmTime       int64
		passwordResetTime int64
		newLoginEmail     bool
		passwordMinSize   int
		userApproval      bool
		emailurlbase      string
		tplemailchange    *htmlTemplate.Template
		tplforgotpass     *htmlTemplate.Template
		tplnewuser        *htmlTemplate.Template
	}

	router struct {
		s service
	}

	// NewUserProps are properties of a newly created user
	NewUserProps struct {
		Userid       string `json:"userid"`
		Username     string `json:"username"`
		Email        string `json:"email"`
		FirstName    string `json:"first_name"`
		LastName     string `json:"last_name"`
		CreationTime int64  `json:"creation_time"`
	}

	// DeleteUserProps are properties of a deleted user
	DeleteUserProps struct {
		Userid string `json:"userid"`
	}
)

const (
	NewUserQueueID    = "gov.user.new"
	DeleteUserQueueID = "gov.user.delete"
)

const (
	time5m     int64 = int64(5 * time.Minute / time.Second)
	time24h    int64 = int64(24 * time.Hour / time.Second)
	time6month int64 = time24h * 365 / 2
)

// New creates a new User
func New(users usermodel.Repo, sessions sessionmodel.Repo, approvals approvalmodel.Repo, roles role.Role, apikeys apikey.Apikey, kv kvstore.KVStore, queue msgqueue.Msgqueue, mailer mail.Mail, tokenizer token.Tokenizer, g gate.Gate) Service {
	hasher := hunter2.NewBlake2bHasher()
	verifier := hunter2.NewVerifier()
	verifier.RegisterHash(hasher)

	return &service{
		users:             users,
		sessions:          sessions,
		approvals:         approvals,
		roles:             roles,
		apikeys:           apikeys,
		kvnewuser:         kv.Subtree("newuser"),
		kvemailchange:     kv.Subtree("emailchange"),
		kvpassreset:       kv.Subtree("passreset"),
		kvsessions:        kv.Subtree("sessions"),
		queue:             queue,
		mailer:            mailer,
		gate:              g,
		hasher:            hasher,
		verifier:          verifier,
		tokenizer:         tokenizer,
		accessTime:        time5m,
		refreshTime:       time6month,
		refreshCacheTime:  time24h,
		confirmTime:       time24h,
		passwordResetTime: time24h,
	}
}

func (s *service) Register(r governor.ConfigRegistrar, jr governor.JobRegistrar) {
	r.SetDefault("accesstime", "5m")
	r.SetDefault("refreshtime", "4380h")
	r.SetDefault("refreshcache", "24h")
	r.SetDefault("confirmtime", "24h")
	r.SetDefault("passwordresettime", "24h")
	r.SetDefault("newloginemail", true)
	r.SetDefault("passwordminsize", 8)
	r.SetDefault("userapproval", false)
	r.SetDefault("email.url.base", "http://localhost:8080")
	r.SetDefault("email.url.emailchange", "/a/confirm/email?key={{.Userid}}.{{.Key}}")
	r.SetDefault("email.url.forgotpass", "/x/resetpass?key={{.Userid}}.{{.Key}}")
	r.SetDefault("email.url.newuser", "/x/confirm?email={{.Email}}&key={{.Key}}")
}

func (s *service) router() *router {
	return &router{
		s: *s,
	}
}

func (s *service) Init(ctx context.Context, c governor.Config, r governor.ConfigReader, l governor.Logger, m governor.Router) error {
	s.logger = l
	l = s.logger.WithData(map[string]string{
		"phase": "init",
	})

	s.baseURL = c.BaseURL
	s.authURL = c.BaseURL + r.URL() + authRoutePrefix
	if t, err := time.ParseDuration(r.GetStr("accesstime")); err != nil {
		return governor.NewError("Failed to parse access time", http.StatusBadRequest, err)
	} else {
		s.accessTime = int64(t / time.Second)
	}
	if t, err := time.ParseDuration(r.GetStr("refreshtime")); err != nil {
		return governor.NewError("Failed to parse refresh time", http.StatusBadRequest, err)
	} else {
		s.refreshTime = int64(t / time.Second)
	}
	if t, err := time.ParseDuration(r.GetStr("refreshcache")); err != nil {
		return governor.NewError("Failed to parse refresh cache", http.StatusBadRequest, err)
	} else {
		s.refreshCacheTime = int64(t / time.Second)
	}
	if t, err := time.ParseDuration(r.GetStr("confirmtime")); err != nil {
		return governor.NewError("Failed to parse confirm time", http.StatusBadRequest, err)
	} else {
		s.confirmTime = int64(t / time.Second)
	}
	if t, err := time.ParseDuration(r.GetStr("passwordresettime")); err != nil {
		return governor.NewError("Failed to parse password reset time", http.StatusBadRequest, err)
	} else {
		s.passwordResetTime = int64(t / time.Second)
	}
	s.newLoginEmail = r.GetBool("newloginemail")
	s.passwordMinSize = r.GetInt("passwordminsize")
	s.userApproval = r.GetBool("userapproval")

	s.emailurlbase = r.GetStr("email.url.base")
	if t, err := htmlTemplate.New("email.url.emailchange").Parse(r.GetStr("email.url.emailchange")); err != nil {
		return governor.NewError("Failed to parse email change url template", http.StatusBadRequest, err)
	} else {
		s.tplemailchange = t
	}
	if t, err := htmlTemplate.New("email.url.forgotpass").Parse(r.GetStr("email.url.forgotpass")); err != nil {
		return governor.NewError("Failed to parse forgot pass url template", http.StatusBadRequest, err)
	} else {
		s.tplforgotpass = t
	}
	if t, err := htmlTemplate.New("email.url.newuser").Parse(r.GetStr("email.url.newuser")); err != nil {
		return governor.NewError("Failed to parse new user url template", http.StatusBadRequest, err)
	} else {
		s.tplnewuser = t
	}

	l.Info("loaded config", map[string]string{
		"accesstime (s)":        strconv.FormatInt(s.accessTime, 10),
		"refreshtime (s)":       strconv.FormatInt(s.refreshTime, 10),
		"refreshcache (s)":      strconv.FormatInt(s.refreshCacheTime, 10),
		"confirmtime (s)":       strconv.FormatInt(s.confirmTime, 10),
		"passwordresettime (s)": strconv.FormatInt(s.passwordResetTime, 10),
		"newloginemail":         strconv.FormatBool(s.newLoginEmail),
		"passwordminsize":       strconv.Itoa(s.passwordMinSize),
		"issuer":                r.GetStr("issuer"),
		"userapproval":          strconv.FormatBool(s.userApproval),
		"tplemailchange":        r.GetStr("email.url.emailchange"),
		"tplforgotpass":         r.GetStr("email.url.forgotpass"),
		"tplnewuser":            r.GetStr("email.url.newuser"),
	})

	sr := s.router()
	sr.mountRoute(m.Group("/user"))
	sr.mountAuth(m.Group(authRoutePrefix))
	sr.mountApikey(m.Group("/apikey"))
	l.Info("mounted http routes", nil)
	return nil
}

func (s *service) Setup(req governor.ReqSetup) error {
	l := s.logger.WithData(map[string]string{
		"phase": "setup",
	})

	madmin, err := s.users.New(req.Username, req.Password, req.Email, req.Firstname, req.Lastname)
	if err != nil {
		return err
	}

	if err := s.users.Setup(); err != nil {
		return err
	}
	l.Info("created user table", nil)

	if err := s.sessions.Setup(); err != nil {
		return err
	}
	l.Info("created usersession table", nil)

	if err := s.approvals.Setup(); err != nil {
		return err
	}
	l.Info("created userapprovals table", nil)

	if err := s.users.Insert(madmin); err != nil {
		return err
	}
	if err := s.roles.InsertRoles(madmin.Userid, rank.Admin()); err != nil {
		return err
	}
	l.Info("inserted new setup admin", map[string]string{
		"username": madmin.Username,
		"userid":   madmin.Userid,
	})
	return nil
}

func (s *service) Start(ctx context.Context) error {
	return nil
}

func (s *service) Stop(ctx context.Context) {
}

func (s *service) Health() error {
	return nil
}
<|endoftext|>


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

package controllers

import (
	"net/http"
	"net/url"
	"reflect"

	"github.com/astaxie/beego"
	log "github.com/golang/glog"
	"github.com/opensds/opensds/pkg/model"
)

const (
	StatusOK       = http.StatusOK
	StatusAccepted = http.StatusAccepted
)

type BasePortal struct {
	beego.Controller
}

func (b *BasePortal) GetParameters() (map[string][]string, error) {
	u, err := url.Parse(b.Ctx.Request.URL.String())
	if err != nil {
		return nil, err
	}
	m, err := url.ParseQuery(u.RawQuery)
	if err != nil {
		return nil, err
	}
	return m, nil
}

// Filter some items in spec that no need to transfer to users.
func (b *BasePortal) outputFilter(resp interface{}, whiteList []string) interface{} {
	v := reflect.ValueOf(resp)
	if v.Kind() == reflect.Slice {
		var s []map[string]interface{}
		for i := 0; i < v.Len(); i++ {
			m := b.doFilter(v.Index(i).Interface(), whiteList)
			s = append(s, m)
		}
		return s
	} else {
		return b.doFilter(resp, whiteList)
	}
}

func (b *BasePortal) doFilter(resp interface{}, whiteList []string) map[string]interface{} {
	v := reflect.ValueOf(resp).Elem()
	m := map[string]interface{}{}
	for _, name := range whiteList {
		field := v.FieldByName(name)
		if field.IsValid() {
			m[name] = field.Interface()
		}
	}
	return m
}

func (b *BasePortal) ErrorHandle(errType int, errMsg string) {
	var errBody []byte

	switch errType {
	case model.ErrorBadRequest:
		errBody = model.ErrorBadRequestStatus(errMsg)
	case model.ErrorUnauthorized:
		errBody = model.ErrorUnauthorizedStatus(errMsg)
	case model.ErrorForbidden:
		errBody = model.ErrorForbiddenStatus(errMsg)
	case model.ErrorNotFound:
		errBody = model.ErrorNotFoundStatus(errMsg)
	case model.ErrorInternalServer:
		errBody = model.ErrorInternalServerStatus(errMsg)
	default:
		errBody = model.ErrorNotImplementedStatus(errMsg)
	}

	b.Ctx.Output.SetStatus(errType)
	b.Ctx.Output.Body(errBody)
	log.Error(errMsg)
}

func (b *BasePortal) SuccessHandle(status int, body []byte) {
	b.Ctx.Output.SetStatus(status)
	if body != nil {
		b.Ctx.Output.Body(body)
	}
}
<commit_msg>Add output content type<commit_after>// Copyright (c) 2019 Huawei Technologies Co., Ltd. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package controllers

import (
	"net/http"
	"net/url"
	"reflect"

	"github.com/astaxie/beego"
	log "github.com/golang/glog"
	"github.com/opensds/opensds/pkg/model"
)

const (
	StatusOK       = http.StatusOK
	StatusAccepted = http.StatusAccepted
)

type BasePortal struct {
	beego.Controller
}

func (b *BasePortal) GetParameters() (map[string][]string, error) {
	u, err := url.Parse(b.Ctx.Request.URL.String())
	if err != nil {
		return nil, err
	}
	m, err := url.ParseQuery(u.RawQuery)
	if err != nil {
		return nil, err
	}
	return m, nil
}

// Filter some items in spec that no need to transfer to users.
func (b *BasePortal) outputFilter(resp interface{}, whiteList []string) interface{} {
	v := reflect.ValueOf(resp)
	if v.Kind() == reflect.Slice {
		var s []map[string]interface{}
		for i := 0; i < v.Len(); i++ {
			m := b.doFilter(v.Index(i).Interface(), whiteList)
			s = append(s, m)
		}
		return s
	} else {
		return b.doFilter(resp, whiteList)
	}
}

func (b *BasePortal) doFilter(resp interface{}, whiteList []string) map[string]interface{} {
	v := reflect.ValueOf(resp).Elem()
	m := map[string]interface{}{}
	for _, name := range whiteList {
		field := v.FieldByName(name)
		if field.IsValid() {
			m[name] = field.Interface()
		}
	}
	return m
}

func (b *BasePortal) ErrorHandle(errType int, errMsg string) {
	var errBody []byte

	switch errType {
	case model.ErrorBadRequest:
		errBody = model.ErrorBadRequestStatus(errMsg)
	case model.ErrorUnauthorized:
		errBody = model.ErrorUnauthorizedStatus(errMsg)
	case model.ErrorForbidden:
		errBody = model.ErrorForbiddenStatus(errMsg)
	case model.ErrorNotFound:
		errBody = model.ErrorNotFoundStatus(errMsg)
	case model.ErrorInternalServer:
		errBody = model.ErrorInternalServerStatus(errMsg)
	default:
		errBody = model.ErrorNotImplementedStatus(errMsg)
	}

	b.Ctx.Output.SetStatus(errType)
	b.Ctx.Output.Header("Content-Type", "application/json; charset=utf-8")
	b.Ctx.Output.Body(errBody)
	log.Error(errMsg)
}

func (b *BasePortal) SuccessHandle(status int, body []byte) {
	b.Ctx.Output.SetStatus(status)
	b.Ctx.Output.Header("Content-Type", "application/json; charset=utf-8")
	if body != nil {
		b.Ctx.Output.Body(body)
	}
}
<|endoftext|>